Marcosertoli/separate examples (#342)

* New files with example models and transforms - moved from model files

* Adapting all files to new example folder/file structure

* Fixed transform plotting for both transect and los

* Fixed linting

indica/converters/abstractconverter.py

@@ -603,22 +603,24 @@ def plot_geometry(
     colors: ArrayLike,
     marker: str = "o",
 ):
+
+    if "LineOfSight" in trans_name:
+        marker = None
+    if hasattr(abscissa, "beamlet"):
+        beamlets = abscissa.beamlet.values
+    else:
+        beamlets = [
+            0,
+        ]
     channels = abscissa.channel.values
-    beamlets = abscissa.beamlet.values
     for channel, beamlet in itertools.product(channels, beamlets):
-        if "LineOfSight" in trans_name:
-            plt.plot(
-                abscissa.sel(channel=channel, beamlet=beamlet),
-                ordinate.sel(channel=channel, beamlet=beamlet),
-                color=colors[channel],
-            )
-        elif "Transect" in trans_name:
-            plt.scatter(
-                abscissa.sel(channel=channel, beamlet=beamlet),
-                ordinate.sel(channel=channel, beamlet=beamlet),
-                color=colors[channel],
-                marker=marker,
-            )
+        col = colors[channel]
+        x = abscissa.sel(channel=channel)
+        y = ordinate.sel(channel=channel)
+        if hasattr(abscissa, "beamlet"):
+            x = x.sel(beamlet=beamlet)
+            y = y.sel(beamlet=beamlet)
+        plt.plot(x, y, color=col, marker=marker)
 
 
 def find_wall_intersections(

indica/examples/example_diagnostic_models.py

@@ -2,9 +2,11 @@
 
 import matplotlib.pylab as plt
 
-from indica.defaults.read_write_defaults import load_default_objects
+from indica.defaults.load_defaults import load_default_objects
 from indica.models import Bolometer
+from indica.models import BremsstrahlungDiode
 from indica.models import ChargeExchange
+from indica.models import EquilibriumReconstruction
 from indica.models import HelikeSpectrometer
 from indica.models import Interferometry
 from indica.models import ThomsonScattering
@@ -88,5 +90,23 @@ def example_interferometer(
     return run_example_diagnostic_model(machine, instrument, _model, plot=plot)
 
 
+def example_diode_filter(
+    plot=False,
+):
+    machine = "st40"
+    instrument = "brems"
+    _model = BremsstrahlungDiode
+    return run_example_diagnostic_model(machine, instrument, _model, plot=plot)
+
+
+def example_equilibrium(
+    plot=False,
+):
+    machine = "st40"
+    instrument = "efit"
+    _model = EquilibriumReconstruction
+    return run_example_diagnostic_model(machine, instrument, _model, plot=plot)
+
+
 if __name__ == "__main__":
     example_helike_spectroscopy(plot=True)

indica/examples/example_plasma.py

@@ -0,0 +1,85 @@
+from pathlib import Path
+import pickle
+from typing import Tuple
+
+import numpy as np
+
+from indica.models import Plasma
+from indica.models.plasma import PlasmaProfiles
+
+
+def example_plasma(
+    machine: str = "st40",
+    pulse: int = None,
+    tstart=0.02,
+    tend=0.1,
+    dt=0.01,
+    main_ion="h",
+    impurities: Tuple[str, ...] = ("c", "ar", "he"),
+    load_from_pkl: bool = True,
+    **kwargs,
+):
+    default_plasma_file = (
+        f"{Path(__file__).parent.parent}/data/{machine}_default_plasma_phantom.pkl"
+    )
+
+    if load_from_pkl and pulse is not None:
+        try:
+            print(f"\n Loading phantom plasma class from {default_plasma_file}. \n")
+            return pickle.load(open(default_plasma_file, "rb"))
+        except FileNotFoundError:
+            print(
+                f"\n\n No phantom plasma class file {default_plasma_file}. \n"
+                f" Building it and saving to file. \n\n"
+            )
+
+    plasma = Plasma(
+        tstart=tstart,
+        tend=tend,
+        dt=dt,
+        main_ion=main_ion,
+        impurities=impurities,
+        **kwargs,
+    )
+    plasma.build_atomic_data()
+
+    update_profiles = PlasmaProfiles(plasma)
+
+    # Assign profiles to time-points
+    nt = len(plasma.t)
+    ne_peaking = np.linspace(1, 2, nt)
+    te_peaking = np.linspace(1, 2, nt)
+    _y0 = update_profiles.profilers["toroidal_rotation"].y0
+    vrot0 = np.linspace(
+        _y0 * 1.1,
+        _y0 * 2.5,
+        nt,
+    )
+    vrot_peaking = np.linspace(1, 2, nt)
+
+    _y0 = update_profiles.profilers["ion_temperature"].y0
+    ti0 = np.linspace(_y0 * 1.1, _y0 * 2.5, nt)
+
+    _y0 = update_profiles.profilers[f"impurity_density:{impurities[0]}"].y0
+    nimp_y0 = _y0 * 5 * np.linspace(1, 8, nt)
+    nimp_peaking = np.linspace(1, 5, nt)
+    nimp_wcenter = np.linspace(0.4, 0.1, nt)
+    for i, t in enumerate(plasma.t):
+        parameters = {
+            "electron_temperature.peaking": te_peaking[i],
+            "ion_temperature.peaking": te_peaking[i],
+            "ion_temperature.y0": ti0[i],
+            "toroidal_rotation.peaking": vrot_peaking[i],
+            "toroidal_rotation.y0": vrot0[i],
+            "electron_density.peaking": ne_peaking[i],
+            "impurity_density:ar.peaking": nimp_peaking[i],
+            "impurity_density:ar.y0": nimp_y0[i],
+            "impurity_density:ar.wcenter": nimp_wcenter[i],
+        }
+        update_profiles(parameters, t=t)
+
+    if load_from_pkl and pulse is not None:
+        print(f"\n Saving phantom plasma class in {default_plasma_file} \n")
+        pickle.dump(plasma, open(default_plasma_file, "wb"))
+
+    return plasma

indica/examples/example_transforms.py

@@ -0,0 +1,78 @@
+import numpy as np
+
+from indica.converters import LineOfSightTransform
+from indica.converters import TransectCoordinates
+
+
+def pi_transform_example(nchannels: int):
+    x_positions = np.linspace(0.2, 0.8, nchannels)
+    y_positions = np.linspace(0.0, 0.0, nchannels)
+    z_positions = np.linspace(0.0, 0.0, nchannels)
+
+    transect_transform = TransectCoordinates(
+        x_positions,
+        y_positions,
+        z_positions,
+        "pi",
+        machine_dimensions=((0.15, 0.95), (-0.7, 0.7)),
+    )
+    return transect_transform
+
+
+def helike_transform_example(nchannels):
+    los_end = np.full((nchannels, 3), 0.0)
+    los_end[:, 0] = 0.17
+    los_end[:, 1] = 0.0
+    los_end[:, 2] = np.linspace(0.2, -0.5, nchannels)
+    los_start = np.array([[0.9, 0, 0]] * los_end.shape[0])
+    los_start[:, 2] = -0.1
+    origin = los_start
+    direction = los_end - los_start
+
+    los_transform = LineOfSightTransform(
+        origin[0:nchannels, 0],
+        origin[0:nchannels, 1],
+        origin[0:nchannels, 2],
+        direction[0:nchannels, 0],
+        direction[0:nchannels, 1],
+        direction[0:nchannels, 2],
+        name="xrcs",
+        machine_dimensions=((0.15, 0.95), (-0.7, 0.7)),
+        passes=1,
+    )
+    return los_transform
+
+
+def smmh1_transform_example(nchannels):
+    los_start = np.array([[0.8, 0, 0]]) * np.ones((nchannels, 3))
+    los_start[:, 2] = np.linspace(0, -0.2, nchannels)
+    los_end = np.array([[0.17, 0, 0]]) * np.ones((nchannels, 3))
+    los_end[:, 2] = np.linspace(0, -0.2, nchannels)
+    origin = los_start
+    direction = los_end - los_start
+    los_transform = LineOfSightTransform(
+        origin[:, 0],
+        origin[:, 1],
+        origin[:, 2],
+        direction[:, 0],
+        direction[:, 1],
+        direction[:, 2],
+        name="smmh1",
+        machine_dimensions=((0.15, 0.95), (-0.7, 0.7)),
+        passes=2,
+    )
+    return los_transform
+
+
+def ts_transform_example(nchannels):
+    x_positions = np.linspace(0.2, 0.8, nchannels)
+    y_positions = np.linspace(0.0, 0.0, nchannels)
+    z_positions = np.linspace(0.0, 0.0, nchannels)
+    transform = TransectCoordinates(
+        x_positions,
+        y_positions,
+        z_positions,
+        "ts",
+        machine_dimensions=((0.15, 0.95), (-0.7, 0.7)),
+    )
+    return transform

indica/models/bolometer_camera.py

@@ -5,9 +5,7 @@
 from xarray import DataArray
 
 from indica.converters import LineOfSightTransform
-from indica.equilibrium import fake_equilibrium
 from indica.models.abstractdiagnostic import DiagnosticModel
-from indica.models.plasma import example_plasma
 from indica.numpy_typing import LabeledArray
 from indica.readers.available_quantities import AVAILABLE_QUANTITIES
 from indica.utilities import assign_datatype
@@ -167,78 +165,3 @@ def plot(self, nplot: int = 1):
         plt.xlabel("rho")
         plt.ylabel("Local radiated power (W/m^3)")
         plt.legend()
-
-
-def example_run(
-    pulse: int = None,
-    diagnostic_name: str = "bolo_xy",
-    origin: LabeledArray = None,
-    direction: LabeledArray = None,
-    plasma=None,
-    plot=False,
-    tplot=None,
-    nchannels: int = 11,
-):
-
-    if plasma is None:
-        plasma = example_plasma(pulse=pulse)
-        machine_dims = plasma.machine_dimensions
-        equilibrium = fake_equilibrium(
-            tstart=plasma.tstart,
-            tend=plasma.tend,
-            dt=plasma.dt / 2.0,
-            machine_dims=machine_dims,
-        )
-        plasma.set_equilibrium(equilibrium)
-
-    # return plasma
-    # Create new interferometers diagnostics
-    if origin is None or direction is None:
-        los_end = np.full((nchannels, 3), 0.0)
-        los_end[:, 0] = 0.0
-        los_end[:, 1] = np.linspace(-0.2, -1, nchannels)
-        los_end[:, 2] = 0.0
-        los_start = np.array([[1.5, 0, 0]] * los_end.shape[0])
-        origin = los_start
-        direction = los_end - los_start
-
-        # los_end = np.full((nchannels, 3), 0.0)
-        # los_end[:, 0] = 0.17
-        # los_end[:, 1] = 0.0
-        # los_end[:, 2] = np.linspace(0.6, -0.6, nchannels)
-        # los_start = np.array([[1.0, 0, 0]] * los_end.shape[0])
-        # origin = los_start
-        # direction = los_end - los_start
-
-    los_transform = LineOfSightTransform(
-        origin[:, 0],
-        origin[:, 1],
-        origin[:, 2],
-        direction[:, 0],
-        direction[:, 1],
-        direction[:, 2],
-        name=diagnostic_name,
-        machine_dimensions=plasma.machine_dimensions,
-        passes=1,
-        beamlets=16,
-        spot_width=0.03,
-    )
-    los_transform.set_equilibrium(plasma.equilibrium)
-    model = Bolometer(
-        diagnostic_name,
-    )
-    model.set_transform(los_transform)
-    model.set_plasma(plasma)
-
-    bckc = model(sum_beamlets=False)
-
-    if plot:
-        model.plot()
-
-    return plasma, model, bckc
-
-
-if __name__ == "__main__":
-    plt.ioff()
-    example_run(plot=True)
-    plt.show()