diff --git a/indica/converters/abstractconverter.py b/indica/converters/abstractconverter.py index 7c79f1c1..869ee725 100644 --- a/indica/converters/abstractconverter.py +++ b/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( diff --git a/indica/examples/example_diagnostic_models.py b/indica/examples/example_diagnostic_models.py index 8d74e5aa..c0b40f48 100644 --- a/indica/examples/example_diagnostic_models.py +++ b/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) diff --git a/indica/examples/example_plasma.py b/indica/examples/example_plasma.py new file mode 100644 index 00000000..3e604110 --- /dev/null +++ b/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 diff --git a/indica/examples/example_transforms.py b/indica/examples/example_transforms.py new file mode 100644 index 00000000..b5f36ee8 --- /dev/null +++ b/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 diff --git a/indica/models/bolometer_camera.py b/indica/models/bolometer_camera.py index a63f7760..ae4ae836 100644 --- a/indica/models/bolometer_camera.py +++ b/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() diff --git a/indica/models/charge_exchange.py b/indica/models/charge_exchange.py index e1fbfca8..54d0e7f2 100644 --- a/indica/models/charge_exchange.py +++ b/indica/models/charge_exchange.py @@ -6,10 +6,10 @@ from indica.converters import TransectCoordinates 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 +from indica.utilities import set_plot_rcparams class ChargeExchange(DiagnosticModel): @@ -115,96 +115,40 @@ def __call__( return self.bckc + def plot(self, nplot: int = 1): + set_plot_rcparams("profiles") -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 example_run( - pulse: int = None, - diagnostic_name: str = "cxrs", - plasma=None, - plot=False, -): - # TODO: LOS sometimes crossing bad EFIT reconstruction - - if plasma is None: - from indica.equilibrium import fake_equilibrium - - 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) - - # Create new interferometers diagnostics - transect_transform = pi_transform_example(5) - transect_transform.set_equilibrium(plasma.equilibrium) - model = ChargeExchange( - diagnostic_name, - ) - model.set_transform(transect_transform) - model.set_plasma(plasma) - - bckc = model() - - if plot: - it = int(len(plasma.t) / 2) - tplot = plasma.t[it].values - - cols_time = cm.gnuplot2(np.linspace(0.1, 0.75, len(plasma.t), dtype=float)) + cols_time = cm.gnuplot2(np.linspace(0.1, 0.75, len(self.t), dtype=float)) - model.transect_transform.plot(tplot) + self.transect_transform.plot() # Plot back-calculated profiles plt.figure() - for i, t in enumerate(plasma.t.values): - plasma.toroidal_rotation.sel(t=t, element=model.element).plot( + for i, t in enumerate(self.t): + if i % nplot: + continue + Vtor = self.bckc["vtor"].sel(t=t, method="nearest") + rho = Vtor.transform.rho.sel(t=t, method="nearest") + plt.scatter( + rho, + Vtor, color=cols_time[i], - label=f"t={t:1.2f} s", + marker="o", alpha=0.7, + label=f"t={t:1.2f} s", ) - Vtor = bckc["vtor"].sel(t=t, method="nearest") - rho = Vtor.transform.rho.sel(t=t, method="nearest") - plt.scatter(rho, Vtor, color=cols_time[i], marker="o", alpha=0.7) plt.xlabel("Channel") plt.ylabel("Measured toroidal rotation (rad/s)") plt.legend() plt.figure() - for i, t in enumerate(plasma.t.values): - plasma.ion_temperature.sel(t=t, element=model.element).plot( - color=cols_time[i], - label=f"t={t:1.2f} s", - alpha=0.7, - ) - Ti = bckc["ti"].sel(t=t, method="nearest") + for i, t in enumerate(self.t): + if i % nplot: + continue + Ti = self.bckc["ti"].sel(t=t, method="nearest") rho = Ti.transform.rho.sel(t=t, method="nearest") plt.scatter(rho, Ti, color=cols_time[i], marker="o", alpha=0.7) plt.xlabel("Channel") plt.ylabel("Measured ion temperature (eV)") plt.legend() plt.show() - - return plasma, model, bckc - - -if __name__ == "__main__": - plt.ioff() - example_run(plot=True) - plt.show() diff --git a/indica/models/diode_filters.py b/indica/models/diode_filters.py index 66d5e140..2ab61056 100644 --- a/indica/models/diode_filters.py +++ b/indica/models/diode_filters.py @@ -8,12 +8,13 @@ from indica.converters import LineOfSightTransform from indica.models.abstractdiagnostic import DiagnosticModel -from indica.models.plasma import example_plasma from indica.numpy_typing import LabeledArray import indica.physics as ph from indica.readers.available_quantities import AVAILABLE_QUANTITIES from indica.utilities import assign_datatype from indica.utilities import format_coord +from indica.utilities import set_axis_sci +from indica.utilities import set_plot_rcparams class BremsstrahlungDiode(DiagnosticModel): @@ -226,95 +227,53 @@ def __call__( self._build_bckc_dictionary() return self.bckc + def plot(self, nplot: int = 1): + set_plot_rcparams("profiles") + if len(self.bckc) == 0: + print("No model results to plot") + return -def example_geometry(nchannels: int = 12): + # Line-of-sight information + self.los_transform.plot(np.mean(self.t)) - 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 - - return origin, direction - - -def example_run( - pulse: int = None, nchannels: int = 12, plasma=None, plot: bool = False -): - if plasma is None: - from indica.equilibrium import fake_equilibrium - - 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) - - # Create new interferometers diagnostics - diagnostic_name = "diode_brems" - origin, direction = example_geometry(nchannels=nchannels) - 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, - ) - los_transform.set_equilibrium(plasma.equilibrium) - model = BremsstrahlungDiode( - diagnostic_name, - ) - model.set_transform(los_transform) - model.set_plasma(plasma) - bckc = model() - - if plot: - it = int(len(plasma.t) / 2) - tplot = plasma.t[it].values - - model.los_transform.plot(tplot) - - # Plot back-calculated values + # Back-calculated profiles + cols_time = cm.gnuplot2(np.linspace(0.1, 0.75, len(self.t), dtype=float)) plt.figure() - cols_chan = cm.gnuplot2( - np.linspace(0.1, 0.75, len(model.los_transform.x1), dtype=float) - ) - for chan in model.los_transform.x1: - bckc["brightness"].sel(channel=chan).plot( - label=f"CH{chan}", color=cols_chan[chan] - ) - plt.xlabel("Time (s)") - plt.ylabel("Bremsstrahlung LOS-integrals (W/m^2)") + for i, t in enumerate(np.array(self.t)): + if i % nplot: + continue + + _brightness = self.bckc["brightness"].sel(t=t, method="nearest") + if "beamlet" in _brightness.dims: + plt.fill_between( + _brightness.channel, + _brightness.max("beamlet"), + _brightness.min("beamlet"), + color=cols_time[i], + alpha=0.5, + ) + brightness = _brightness.mean("beamlet") + else: + brightness = _brightness + brightness.plot(label=f"t={t:1.2f} s", color=cols_time[i]) + set_axis_sci() + plt.title(self.name.upper()) + plt.xlabel("Channel") + plt.ylabel("Measured brightness (W/m^2)") plt.legend() - # Plot the profiles - cols_time = cm.gnuplot2(np.linspace(0.1, 0.75, len(plasma.t), dtype=float)) + # Local emissivity profiles plt.figure() - for i, t in enumerate(plasma.t.values): + for i, t in enumerate(np.array(self.t)): + if i % nplot: + continue plt.plot( - model.emissivity.rho_poloidal, - model.emissivity.sel(t=t), + self.emissivity.rho_poloidal, + self.emissivity.sel(t=t), color=cols_time[i], label=f"t={t:1.2f} s", ) + set_axis_sci() plt.xlabel("rho") - plt.ylabel("Bremsstrahlung emission (W/m^3)") + plt.ylabel("Local radiated power (W/m^3)") plt.legend() - - return plasma, model, bckc - - -if __name__ == "__main__": - plt.ioff() - example_run(plot=True) - plt.show() diff --git a/indica/models/equilibrium_reconstruction.py b/indica/models/equilibrium_reconstruction.py index 6e4ee5e1..63a6186a 100644 --- a/indica/models/equilibrium_reconstruction.py +++ b/indica/models/equilibrium_reconstruction.py @@ -1,8 +1,6 @@ -import matplotlib.pyplot as plt import xarray as xr from indica.models.abstractdiagnostic import DiagnosticModel -from indica.models.plasma import example_plasma from indica.readers.available_quantities import AVAILABLE_QUANTITIES from indica.utilities import assign_datatype from indica.utilities import check_time_present @@ -62,37 +60,3 @@ def __call__( self.wp = self.plasma.wp.interp(t=t) self._build_bckc_dictionary() return self.bckc - - -def example_run( - diagnostic_name: str = "efit", - plasma=None, - plot=False, - t=None, -): - if plasma is None: - from indica.equilibrium import fake_equilibrium - - plasma = example_plasma() - 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) - - model = EquilibriumReconstruction(diagnostic_name) - model.set_plasma(plasma) - bckc = model() - - if plot: - bckc["wp"].plot() - - return plasma, model, bckc - - -if __name__ == "__main__": - example_run(plot=True) - plt.show(block=True) diff --git a/indica/models/helike_spectroscopy.py b/indica/models/helike_spectroscopy.py index 9f099268..88087144 100644 --- a/indica/models/helike_spectroscopy.py +++ b/indica/models/helike_spectroscopy.py @@ -5,7 +5,6 @@ from xarray import DataArray from indica.converters import LineOfSightTransform -from indica.defaults.load_defaults import load_default_objects from indica.models.abstractdiagnostic import DiagnosticModel from indica.numpy_typing import LabeledArray import indica.physics as ph @@ -561,57 +560,3 @@ def plot(self): # plt.legend() plt.show(block=True) - - -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 example_run(plasma=None, plot=False, moment_analysis: bool = False, **kwargs): - - if plasma is None: - plasma = load_default_objects("st40", "plasma") - equilibrium = load_default_objects("st40", "equilibrium") - plasma.set_equilibrium(equilibrium) - - diagnostic_name = "xrcs" - los_transform = helike_transform_example(3) - los_transform.set_equilibrium(plasma.equilibrium) - model = HelikeSpectrometer( - diagnostic_name, - window_masks=[], - ) - model.set_transform(los_transform) - model.set_plasma(plasma) - - bckc = model(moment_analysis=moment_analysis, **kwargs) - - # Plot spectra - if plot: - model.plot() - - return plasma, model, bckc - - -if __name__ == "__main__": - example_run(plot=True, moment_analysis=True) diff --git a/indica/models/interferometry.py b/indica/models/interferometry.py index 60b76802..ddd9680b 100644 --- a/indica/models/interferometry.py +++ b/indica/models/interferometry.py @@ -1,14 +1,13 @@ -from matplotlib import cm import matplotlib.pylab as plt -import numpy as np import xarray as xr from xarray import DataArray from indica.converters import LineOfSightTransform 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 set_axis_sci +from indica.utilities import set_plot_rcparams class Interferometry(DiagnosticModel): @@ -88,90 +87,31 @@ def __call__( self._build_bckc_dictionary() return self.bckc + def plot(self, nplot: int = 1): + set_plot_rcparams("profiles") + if len(self.bckc) == 0: + print("No model results to plot") + return -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 example_run(pulse: int = None, plasma=None, plot=False): - if plasma is None: - from indica.equilibrium import fake_equilibrium - - 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) - - # Create new interferometers diagnostics - diagnostic_name = "smmh1" - model = Interferometry( - diagnostic_name, - ) - los_transform = smmh1_transform_example(3) - los_transform.set_equilibrium(plasma.equilibrium) - model.set_transform(los_transform) - model.set_plasma(plasma) - - bckc = model() - - if plot: - it = int(len(plasma.t) / 2) - tplot = plasma.t[it].values + # Line-of-sight information + self.los_transform.plot() - model.los_transform.plot(tplot) - - # Plot back-calculated values + # Back-calculated profiles plt.figure() - cols_chan = cm.gnuplot2( - np.linspace(0.1, 0.75, len(model.los_transform.x1), dtype=float) - ) - for chan in model.los_transform.x1: - bckc["ne"].sel(channel=chan).plot(label=f"CH{chan}", color=cols_chan[chan]) - plt.xlabel("Time (s)") - plt.ylabel("Ne LOS-integrals (m^-2)") - plt.legend() - - # Plot the profiles - cols_time = cm.gnuplot2(np.linspace(0.1, 0.75, len(plasma.t), dtype=float)) - plt.figure() - for i, t in enumerate(plasma.t.values): - plt.plot( - model.Ne.rho_poloidal, - model.Ne.sel(t=t), - color=cols_time[i], - label=f"t={t:1.2f} s", + _value = self.bckc["ne"] + if "beamlet" in _value.dims: + plt.fill_between( + _value.t, + _value.max("beamlet"), + _value.min("beamlet"), + alpha=0.5, ) - plt.xlabel("rho") - plt.ylabel("Ne (m^-3)") + value = _value.mean("beamlet") + else: + value = _value + value.plot() + set_axis_sci() + plt.title(self.name.upper()) + plt.xlabel("Time (s)") + plt.ylabel("Measured LOS-integrated density (m^-2)") plt.legend() - plt.show(block=True) - - return plasma, model, bckc - - -if __name__ == "__main__": - plt.ioff() - example_run(plot=True) - plt.show() diff --git a/indica/models/plasma.py b/indica/models/plasma.py index 7367f2e2..f488aea5 100644 --- a/indica/models/plasma.py +++ b/indica/models/plasma.py @@ -1,7 +1,6 @@ from copy import deepcopy from functools import lru_cache import hashlib -from pathlib import Path import pickle from typing import Callable from typing import List @@ -923,86 +922,3 @@ def __call__(self, parameters: dict, t: float = None): else: datatype = profile_datatype getattr(plasma, datatype).loc[dict(t=t)] = profiler_output - - -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" - ) - - # TODO: swap all profiles to new version! - - 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 - - -if __name__ == "__main__": - example_plasma() diff --git a/indica/models/thomson_scattering.py b/indica/models/thomson_scattering.py index 4fa63e5d..b117d92c 100644 --- a/indica/models/thomson_scattering.py +++ b/indica/models/thomson_scattering.py @@ -6,7 +6,6 @@ from indica.converters import TransectCoordinates 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 @@ -159,95 +158,3 @@ def plot(self): plt.xlabel("Channel") plt.ylabel("Measured electron temperature (eV)") plt.legend() - - -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 - - -def example_run( - pulse: int = None, - diagnostic_name: str = "ts", - plasma=None, - plot=False, -): - - # TODO: LOS sometime crossing bad EFIT reconstruction and crashing... - - if plasma is None: - from indica.equilibrium import fake_equilibrium - - 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) - - transect_transform = ts_transform_example(11) - transect_transform.set_equilibrium(plasma.equilibrium) - model = ThomsonScattering( - diagnostic_name, - ) - model.set_transform(transect_transform) - model.set_plasma(plasma) - - bckc = model() - - if plot: - it = int(len(plasma.t) / 2) - tplot = plasma.t[it].values - - cols_time = cm.gnuplot2(np.linspace(0.1, 0.75, len(plasma.t), dtype=float)) - - model.transect_transform.plot(tplot) - - # Plot back-calculated profiles - plt.figure() - for i, t in enumerate(plasma.t.values): - plasma.electron_density.sel(t=t).plot( - color=cols_time[i], - label=f"t={t:1.2f} s", - alpha=0.7, - ) - Ne = bckc["ne"].sel(t=t, method="nearest") - rho = Ne.transform.rho.sel(t=t, method="nearest") - plt.scatter(rho, Ne, color=cols_time[i], marker="o", alpha=0.7) - plt.xlabel("Channel") - plt.ylabel("Measured electron density (m^-3)") - plt.legend() - - plt.figure() - for i, t in enumerate(plasma.t.values): - plasma.electron_temperature.sel(t=t).plot( - color=cols_time[i], - label=f"t={t:1.2f} s", - alpha=0.7, - ) - Te = bckc["te"].sel(t=t, method="nearest") - rho = Te.transform.rho.sel(t=t, method="nearest") - plt.scatter(rho, Te, color=cols_time[i], marker="o", alpha=0.7) - plt.xlabel("Channel") - plt.ylabel("Measured electron temperature (eV)") - plt.legend() - - return plasma, model, bckc - - -if __name__ == "__main__": - plt.ioff() - example_run(plot=True) - plt.show() diff --git a/indica/workflows/bayes_workflow.py b/indica/workflows/bayes_workflow.py index 50b59db1..22e0d95f 100644 --- a/indica/workflows/bayes_workflow.py +++ b/indica/workflows/bayes_workflow.py @@ -24,17 +24,17 @@ from indica.bayesblackbox import ln_prior from indica.defaults.load_defaults import load_default_objects from indica.equilibrium import Equilibrium +from indica.examples.example_transforms import helike_transform_example +from indica.examples.example_transforms import pi_transform_example +from indica.examples.example_transforms import smmh1_transform_example +from indica.examples.example_transforms import ts_transform_example from indica.models.charge_exchange import ChargeExchange -from indica.models.charge_exchange import pi_transform_example from indica.models.equilibrium_reconstruction import EquilibriumReconstruction -from indica.models.helike_spectroscopy import helike_transform_example from indica.models.helike_spectroscopy import HelikeSpectrometer from indica.models.interferometry import Interferometry -from indica.models.interferometry import smmh1_transform_example from indica.models.plasma import Plasma from indica.models.plasma import PlasmaProfiles from indica.models.thomson_scattering import ThomsonScattering -from indica.models.thomson_scattering import ts_transform_example from indica.operators.gpr_fit import gpr_fit_ts from indica.operators.gpr_fit import post_process_ts from indica.readers.read_st40 import ReadST40 diff --git a/tests/integration/models/test_diagnostic_models.py b/tests/integration/models/test_diagnostic_models.py index 6723c948..60f5143d 100644 --- a/tests/integration/models/test_diagnostic_models.py +++ b/tests/integration/models/test_diagnostic_models.py @@ -1,28 +1,12 @@ from typing import Callable -from typing import Dict - -import numpy as np -import pytest from indica.defaults.load_defaults import load_default_objects -from indica.models.bolometer_camera import example_run as bolo -from indica.models.charge_exchange import example_run as cxrs -from indica.models.diode_filters import example_run as diodes -from indica.models.equilibrium_reconstruction import example_run as equil_recon -from indica.models.helike_spectroscopy import example_run as helike -from indica.models.interferometry import example_run as interf -from indica.models.thomson_scattering import example_run as ts - - -EXAMPLES: Dict[str, Callable] = { - "bolometer_camera": bolo, - "diode_filters": diodes, - "interferometry": interf, - "helike_spectroscopy": helike, - "thomson_scattering": ts, - "charge_exchange": cxrs, - "equilibrium_reconstruction": equil_recon, -} +from indica.models import Bolometer +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 class TestModels: @@ -30,99 +14,42 @@ class TestModels: def setup_class(self): machine = "st40" - equilibrium = load_default_objects(machine, "equilibrium") - _plasma = load_default_objects(machine, "plasma") - _plasma.set_equilibrium(equilibrium) - - self.plasma = _plasma - nt = np.size(self.plasma.t) - tstart = self.plasma.tstart - tend = self.plasma.tend - dt = self.plasma.dt - it = int(nt / 2.0) - self.time_single_pass = float(self.plasma.t[it].values) - self.time_single_fail = float(np.max(self.plasma.equilibrium.rho.t) + 1.0) - self.time_interp = np.linspace(tstart + dt, tend - dt, num=int(nt / 3)) - - def _test_timepoint_pass(self, model_name: str, **kwargs): - """Test that model can be called for single time-point""" - _, model, bckc = EXAMPLES[model_name](plasma=self.plasma, **kwargs) - model(t=self.time_single_pass) - - def _test_timepoint_fail(self, model_name: str, **kwargs): - """Test that model can be called for single time-point""" - _, model, bckc = EXAMPLES[model_name](plasma=self.plasma, **kwargs) - with pytest.raises(Exception): - model(t=self.time_single_fail) - - def _test_time_interpolation(self, model_name: str, **kwargs): - """Test that model correctly interpolates data on new axis""" - _, model, bckc = EXAMPLES[model_name](plasma=self.plasma, **kwargs) - bckc = model(t=self.time_interp, **kwargs) - - for quantity, value in bckc.items(): - if "t" in value.dims: - assert np.array_equal(value.t.values, self.time_interp) - - def test_cxrs_timepoint_fail(self): - self._test_timepoint_fail("charge_exchange") - - def test_cxrs_timepoint_pass(self): - self._test_timepoint_pass("charge_exchange") - - def test_cxrs_interpolation(self): - self._test_time_interpolation("charge_exchange") - - def test_ts_timepoint_fail(self): - self._test_timepoint_fail("thomson_scattering") - - def test_ts_timepoint_pass(self): - self._test_timepoint_pass("thomson_scattering") - - def test_ts_interpolation(self): - self._test_time_interpolation("thomson_scattering") - - def test_bolo_timepoint_fail(self): - self._test_timepoint_fail("bolometer_camera") - - def test_bolo_timepoint_pass(self): - self._test_timepoint_pass("bolometer_camera") - - def test_bolo_interpolation(self): - self._test_time_interpolation("bolometer_camera") - - def test_interf_timepoint_fail(self): - self._test_timepoint_fail("interferometry") - - def test_interf_timepoint_pass(self): - self._test_timepoint_pass("interferometry") + self.machine = machine + self.transforms = load_default_objects(machine, "geometry") + self.equilibrium = load_default_objects(machine, "equilibrium") + self.plasma = load_default_objects(machine, "plasma") + self.plasma.set_equilibrium(self.equilibrium) - def test_interf_interpolation(self): - self._test_time_interpolation("interferometry") + def run_model(self, instrument: str, model: Callable): + """ + Make sure model runs without errors + """ - def test_diodes_timepoint_fail(self): - self._test_timepoint_fail("diode_filters") + model = model(instrument) + if instrument in self.transforms.keys(): + transform = self.transforms[instrument] + transform.set_equilibrium(self.equilibrium) + model.set_transform(transform) + model.set_plasma(self.plasma) - def test_diodes_timepoint_pass(self): - self._test_timepoint_pass("diode_filters") + bckc = model(sum_beamlets=False) - # def test_diodes_interpolation(self): - # self._test_time_interpolation("diode_filters") + assert type(bckc) == dict - def test_helike_timepoint_fail(self): - self._test_timepoint_fail("helike_spectroscopy") + def test_interferometer(self): + self.run_model("smmh", Interferometry) - def test_helike_timepoint_pass(self): - self._test_timepoint_pass("helike_spectroscopy") + def test_helike_spectroscopy(self): + self.run_model("xrcs", HelikeSpectrometer) - # def test_helike_interpolation(self): - # self._test_time_interpolation("helike_spectroscopy") + def test_thomson_scattering(self): + self.run_model("ts", ThomsonScattering) - def test_equil_recon_timepoint_fail(self): - self._test_timepoint_fail("equilibrium_reconstruction") + def test_charge_exchange(self): + self.run_model("cxff_pi", ChargeExchange) - def test_equil_recon_timepoint_pass(self): - self._test_timepoint_pass("equilibrium_reconstruction") + def test_equilibrium(self): + self.run_model("efit", EquilibriumReconstruction) - def test_equil_recon_interpolation(self): - self._test_time_interpolation("equilibrium_reconstruction") + def test_bolometer(self): + self.run_model("blom_xy1", Bolometer) diff --git a/tests/integration/models/test_helike_spectroscopy.py b/tests/integration/models/test_helike_spectroscopy.py index d74a0049..896f6f2e 100644 --- a/tests/integration/models/test_helike_spectroscopy.py +++ b/tests/integration/models/test_helike_spectroscopy.py @@ -1,6 +1,6 @@ from indica.defaults.load_defaults import load_default_objects +from indica.examples.example_transforms import helike_transform_example import indica.models.helike_spectroscopy as helike -from indica.models.helike_spectroscopy import helike_transform_example class TestHelike: diff --git a/tests/integration/models/test_plasma.py b/tests/integration/models/test_plasma.py index cb02e62f..6315e8a4 100644 --- a/tests/integration/models/test_plasma.py +++ b/tests/integration/models/test_plasma.py @@ -1,6 +1,6 @@ -import indica.models.plasma as plasma +from indica.examples.example_plasma import example_plasma def test_example_plasma(): - _ = plasma.example_plasma() + _ = example_plasma()