From 635e35c9d3b52d30fda4d65e38401c6aec433e92 Mon Sep 17 00:00:00 2001 From: gtrevisan Date: Tue, 28 Apr 2026 16:59:26 -0400 Subject: [PATCH 1/9] offer data getters from params itself --- disruption_py/core/physics_method/params.py | 39 +++++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/disruption_py/core/physics_method/params.py b/disruption_py/core/physics_method/params.py index 1b7f33f2..d102d5e6 100644 --- a/disruption_py/core/physics_method/params.py +++ b/disruption_py/core/physics_method/params.py @@ -69,3 +69,42 @@ def to_coords(self) -> Dict[str, Tuple[str, np.ndarray]]: }, dim="idx", ) + + def get_data(self, *args, **kwargs) -> np.ndarray: + """ + Get data using the data connection. + + Parameters + ---------- + *args + Positional arguments for `data_conn.get_data`. + **kwargs + Keyword arguments for `data_conn.get_data`. + """ + return self.data_conn.get_data(*args, **kwargs) + + def get_data_with_dims(self, *args, **kwargs) -> np.ndarray: + """ + Get data with dimensions using the data connection. + + Parameters + ---------- + *args + Positional arguments for `data_conn.get_data_with_dims`. + **kwargs + Keyword arguments for `data_conn.get_data_with_dims`. + """ + return self.data_conn.get_data_with_dims(*args, **kwargs) + + def get_dims(self, *args, **kwargs) -> np.ndarray: + """ + Get dimensions using the data connection. + + Parameters + ---------- + *args + Positional arguments for `data_conn.get_dims`. + **kwargs + Keyword arguments for `data_conn.get_dims`. + """ + return self.data_conn.get_dims(*args, **kwargs) From 07a18b2c6a338207a97dcf0f5563a0924e2972b7 Mon Sep 17 00:00:00 2001 From: gtrevisan Date: Tue, 28 Apr 2026 17:03:04 -0400 Subject: [PATCH 2/9] remove params.data_conn shortcuts --- disruption_py/machine/mast/efit.py | 5 ++- disruption_py/machine/mast/physics.py | 51 ++++++++++++--------------- 2 files changed, 24 insertions(+), 32 deletions(-) diff --git a/disruption_py/machine/mast/efit.py b/disruption_py/machine/mast/efit.py index 7f51004a..3bb00789 100644 --- a/disruption_py/machine/mast/efit.py +++ b/disruption_py/machine/mast/efit.py @@ -50,13 +50,12 @@ def get_efit_parameters(params: PhysicsMethodParams): dict A dictionary containing the retrieved EFIT parameters. """ - conn = params.data_conn - eq_time = conn.get_data("equilibrium/time") + eq_time = params.data_conn.get_data("equilibrium/time") times = params.times outputs = {} for key, prop in MastEfitMethods.efit_properties.items(): - signal = conn.get_data(f"equilibrium/{prop}") + signal = params.data_conn.get_data(f"equilibrium/{prop}") item = MastUtilMethods.interpolate_1d(eq_time, signal, times) outputs[key] = item diff --git a/disruption_py/machine/mast/physics.py b/disruption_py/machine/mast/physics.py index 11f772ce..d8c3216b 100644 --- a/disruption_py/machine/mast/physics.py +++ b/disruption_py/machine/mast/physics.py @@ -40,11 +40,10 @@ def get_ip_parameters(params: PhysicsMethodParams): A dictionary containing plasma current (`ip`), its time derivative (`dip_dt`), programmed plasma current (`ip_prog`), and its time derivative (`dipprog_dt`). """ - conn = params.data_conn - ip = conn.get_data("summary/ip") - ip_prog = conn.get_data("pulse_schedule/i_plasma") - ip_prog_time = conn.get_data("pulse_schedule/time") - magtime = conn.get_data("summary/time") + ip = params.data_conn.get_data("summary/ip") + ip_prog = params.data_conn.get_data("pulse_schedule/i_plasma") + ip_prog_time = params.data_conn.get_data("pulse_schedule/time") + magtime = params.data_conn.get_data("summary/time") dip_dt = np.gradient(ip, magtime) dipprog_dt = np.gradient(ip_prog, ip_prog_time) @@ -82,11 +81,10 @@ def get_power(params: PhysicsMethodParams): A dictionary containing neutral beam injection power (`power_nbi`) and radiated power (`power_radiated`). """ - conn = params.data_conn - power_nbi = conn.get_data("summary/power_nbi") - power_radiated = conn.get_data("summary/power_radiated") - base_time = conn.get_data("summary/time") + power_nbi = params.data_conn.get_data("summary/power_nbi") + power_radiated = params.data_conn.get_data("summary/power_radiated") + base_time = params.data_conn.get_data("summary/time") times = params.times power_nbi = MastUtilMethods.interpolate_1d(base_time, power_nbi, times) @@ -114,12 +112,11 @@ def get_gas(params: PhysicsMethodParams): A dictionary containing total injected gas (`total_injected`), inboard total gas (`inboard_total`), and outboard total gas (`outboard_total`). """ - conn = params.data_conn - total_injected = conn.get_data("gas_injection/total_injected") - inboard_total = conn.get_data("gas_injection/inboard_total") - outboard_total = conn.get_data("gas_injection/outboard_total") - base_time = conn.get_data("gas_injection/time") + total_injected = params.data_conn.get_data("gas_injection/total_injected") + inboard_total = params.data_conn.get_data("gas_injection/inboard_total") + outboard_total = params.data_conn.get_data("gas_injection/outboard_total") + base_time = params.data_conn.get_data("gas_injection/time") times = params.times total_injected = MastUtilMethods.interpolate_1d( @@ -155,11 +152,10 @@ def get_ts_parameters(params: PhysicsMethodParams): core electron density (`n_e_core`). """ times = params.times - conn = params.data_conn - t_e_core = conn.get_data("thomson_scattering/t_e_core") - n_e_core = conn.get_data("thomson_scattering/n_e_core") - base_time = conn.get_data("thomson_scattering/time") + t_e_core = params.data_conn.get_data("thomson_scattering/t_e_core") + n_e_core = params.data_conn.get_data("thomson_scattering/n_e_core") + base_time = params.data_conn.get_data("thomson_scattering/time") t_e_core = MastUtilMethods.interpolate_1d(base_time, t_e_core, times) n_e_core = MastUtilMethods.interpolate_1d(base_time, n_e_core, times) @@ -198,13 +194,12 @@ def get_densities(params: PhysicsMethodParams): """ - conn = params.data_conn - n_e = conn.get_data("summary/line_average_n_e") - t_n = conn.get_data("summary/time") - ip = conn.get_data("summary/ip") - t_ip = conn.get_data("summary/time") - a_minor = conn.get_data("equilibrium/minor_radius") - t_a = conn.get_data("equilibrium/time") + n_e = params.data_conn.get_data("summary/line_average_n_e") + t_n = params.data_conn.get_data("summary/time") + ip = params.data_conn.get_data("summary/ip") + t_ip = params.data_conn.get_data("summary/time") + a_minor = params.data_conn.get_data("equilibrium/minor_radius") + t_a = params.data_conn.get_data("equilibrium/time") return MastPhysicsMethods._get_densities( params.times, n_e, t_n, ip, t_ip, a_minor, t_a @@ -273,8 +268,7 @@ def get_sxr(params: PhysicsMethodParams): A dictionary containing SXR data (`sxr_data`) and corresponding time points (`sxr_time`). """ - conn = params.data_conn - hcam = conn.get_data("soft_x_rays/horizontal_cam_upper", return_xarray=True) + hcam = params.data_conn.get_data("soft_x_rays/horizontal_cam_upper", return_xarray=True) if hcam is not None: hcam = hcam.isel(horizontal_cam_upper_channel=7) @@ -309,9 +303,8 @@ def get_dalpha(params: PhysicsMethodParams): dict A dictionary containing D-alpha signal data (`dalpha`). """ - conn = params.data_conn - dalpha = conn.get_data( + dalpha = params.data_conn.get_data( "spectrometer_visible/filter_spectrometer_dalpha_voltage", return_xarray=True, ) From 51f0dde5c39f99e4f54dedb9edc42fade2e710be Mon Sep 17 00:00:00 2001 From: gtrevisan Date: Tue, 28 Apr 2026 17:03:54 -0400 Subject: [PATCH 3/9] call data getters from params --- disruption_py/machine/cmod/efit.py | 6 +- disruption_py/machine/cmod/physics.py | 132 ++++++++++++------------- disruption_py/machine/cmod/thomson.py | 34 +++---- disruption_py/machine/d3d/efit.py | 8 +- disruption_py/machine/d3d/physics.py | 112 ++++++++++----------- disruption_py/machine/d3d/util.py | 2 +- disruption_py/machine/east/efit.py | 8 +- disruption_py/machine/east/physics.py | 84 ++++++++-------- disruption_py/machine/hbtep/physics.py | 50 +++++----- disruption_py/machine/mast/efit.py | 4 +- disruption_py/machine/mast/physics.py | 44 ++++----- disruption_py/settings/time_setting.py | 22 ++--- docs/examples/custom_physics_method.py | 6 +- docs/examples/custom_time_setting.py | 2 +- 14 files changed, 257 insertions(+), 257 deletions(-) diff --git a/disruption_py/machine/cmod/efit.py b/disruption_py/machine/cmod/efit.py index 8bc856c9..349d2b61 100644 --- a/disruption_py/machine/cmod/efit.py +++ b/disruption_py/machine/cmod/efit.py @@ -69,7 +69,7 @@ def get_efit_parameters(params: PhysicsMethodParams): dict A dictionary containing the retrieved EFIT parameters. """ - efit_time = params.data_conn.get_data( + efit_time = params.get_data( r"\efit_aeqdsk:time", tree_name="_efit_tree" ) # [s] efit_data = {} @@ -77,7 +77,7 @@ def get_efit_parameters(params: PhysicsMethodParams): # Get data from each of the columns in efit_cols one at a time for param, path in CmodEfitMethods.efit_cols.items(): try: - efit_data[param] = params.data_conn.get_data( + efit_data[param] = params.get_data( path=path, tree_name="_efit_tree", ) @@ -125,7 +125,7 @@ def efit_check(params: PhysicsMethodParams): ] _n = values[2].data() valid_indices = np.nonzero(_n) - (times,) = params.data_conn.get_dims( + (times,) = params.get_dims( r"\efit_aeqdsk:lflag", tree_name="analysis" ) return valid_indices, times[valid_indices] diff --git a/disruption_py/machine/cmod/physics.py b/disruption_py/machine/cmod/physics.py index 23d5cbaa..1f2b87a6 100644 --- a/disruption_py/machine/cmod/physics.py +++ b/disruption_py/machine/cmod/physics.py @@ -61,7 +61,7 @@ def _get_active_wire_segments(params: PhysicsMethodParams): for node_path in children_paths: node_path = node_path.strip() if node_path.split(".")[-1].startswith("SEG_"): - is_on = params.data_conn.get_data( + is_on = params.get_data( 'getnci($, "STATE")', arguments=f"{node_path}:SEG_NUM" ) # 0 represents node being on, 1 represents node being off @@ -70,7 +70,7 @@ def _get_active_wire_segments(params: PhysicsMethodParams): active_segments.append( ( node_path, - params.data_conn.get_data( + params.get_data( f"{node_path}:start_time", tree_name="pcs" ), ) @@ -226,17 +226,17 @@ def get_ip_parameters(params: PhysicsMethodParams): for node_path, start in active_segments: # Ip wire can be one of 16 but is normally no. 16 for wire_index in range(16, 0, -1): - wire_node_name = params.data_conn.get_data( + wire_node_name = params.get_data( f"{node_path}:P_{wire_index:02d}:name", tree_name="pcs" ) if wire_node_name == "IP": try: - pid_gains = params.data_conn.get_data( + pid_gains = params.get_data( f"{node_path}:P_{wire_index:02d}:pid_gains", tree_name="pcs", ) if np.any(pid_gains): - signal, sigtime = params.data_conn.get_data_with_dims( + signal, sigtime = params.get_data_with_dims( f"{node_path}:P_{wire_index:02d}", tree_name="pcs" ) ip_prog_temp = interp1( @@ -257,7 +257,7 @@ def get_ip_parameters(params: PhysicsMethodParams): params.logger.warning(repr(e)) params.logger.opt(exception=True).debug(e) break # Break out of wire_index loop - ip, magtime = params.data_conn.get_data_with_dims( + ip, magtime = params.get_data_with_dims( r"\ip", tree_name="magnetics" ) # [A], [s] return CmodPhysicsMethods._get_ip_parameters( @@ -380,17 +380,17 @@ def get_z_parameters(params: PhysicsMethodParams): for node_path, start in active_wire_segments: for wire_index in range(1, 17): - wire_node_name = params.data_conn.get_data( + wire_node_name = params.get_data( f"{node_path}:P_{wire_index:02d}:name", tree_name="pcs" ) if wire_node_name == "ZCUR": try: - pid_gains = params.data_conn.get_data( + pid_gains = params.get_data( f"{node_path}:P_{wire_index:02d}:pid_gains", tree_name="pcs", ) if np.any(pid_gains): - signal, sigtime = params.data_conn.get_data_with_dims( + signal, sigtime = params.get_data_with_dims( f"{node_path}:P_{wire_index:02d}", tree_name="pcs" ) end = sigtime[ @@ -421,7 +421,7 @@ def get_z_parameters(params: PhysicsMethodParams): raise CalculationError("Data source error: No ZCUR wire was found") # Read in A_OUT, which is a 16xN matrix of the errors for *all* 16 wires for # *all* of the segments. Note that DPCS time is usually taken at 10kHz. - wire_errors, dpcstime = params.data_conn.get_data_with_dims( + wire_errors, dpcstime = params.get_data_with_dims( r"\top.hardware.dpcs.signals:a_out", tree_name="hybrid", dim_nums=[1] ) # The value of Z_error we read is not in the units we want. It must be *divided* @@ -438,7 +438,7 @@ def get_z_parameters(params: PhysicsMethodParams): else: end = active_wire_segments[i + 1][1] # DPCS refers to PCS so we need to open the common ancestor tree, HYBRID - z_factor = params.data_conn.get_data( + z_factor = params.get_data( rf"\dpcs::top.seg_{i + 1:02d}:p_{z_wire_index:02d}:predictor:factor", tree_name="hybrid", ) @@ -452,16 +452,16 @@ def get_z_parameters(params: PhysicsMethodParams): # before 2015. # TODO: Try to fix this if params.shot_id > 1150101000: - ip_without_factor = params.data_conn.get_data( + ip_without_factor = params.get_data( r"\top.hardware.dpcs.signals.a_in:input_056", tree_name="hybrid" ) - ip_factor = params.data_conn.get_data( + ip_factor = params.get_data( r"\top.dpcs_config.inputs:input_056:p_to_v_expr", tree_name="hybrid", ) ip = ip_without_factor * ip_factor # [A] else: - ip, ip_time = params.data_conn.get_data_with_dims( + ip, ip_time = params.get_data_with_dims( r"\ip", tree_name="magnetics" ) # [A], [s] ip = interp1(ip_time, ip, dpcstime) @@ -542,7 +542,7 @@ def get_ohmic_parameters(params: PhysicsMethodParams): - pull requests: #[367](https://github.com/MIT-PSFC/disruption-py/pull/367) """ try: - v_loop, v_loop_time = params.data_conn.get_data_with_dims( + v_loop, v_loop_time = params.get_data_with_dims( r"\top.mflux:v0", tree_name="analysis" ) # [V], [s] # Apply 6-point boxcar smoothing to raw vloop signal. @@ -552,17 +552,17 @@ def get_ohmic_parameters(params: PhysicsMethodParams): params.logger.verbose( r"v_loop: Failed to get \top.mflux:v0 data. Use \efit_aeqdsk:vloopt instead." ) - v_loop, v_loop_time = params.data_conn.get_data_with_dims( + v_loop, v_loop_time = params.get_data_with_dims( r"\efit_aeqdsk:vloopt", tree_name="_efit_tree" ) # [V], [s] if len(v_loop_time) <= 1: raise CalculationError("No data for v_loop_time") - li, efittime = params.data_conn.get_data_with_dims( + li, efittime = params.get_data_with_dims( r"\efit_aeqdsk:ali", tree_name="_efit_tree" ) # [dimensionless], [s] ip_parameters = CmodPhysicsMethods.get_ip_parameters(params=params) - r0 = params.data_conn.get_data( + r0 = params.get_data( r"\efit_aeqdsk:rmagx/100", tree_name="_efit_tree" ) # [m] @@ -739,7 +739,7 @@ def get_power(params: PhysicsMethodParams): p = f"p_{val}" t = f"t_{val}" try: - kwa[p], kwa[t] = params.data_conn.get_data_with_dims( + kwa[p], kwa[t] = params.get_data_with_dims( node, tree_name=tree ) except (mdsExceptions.TreeFOPENR, mdsExceptions.TreeNNF): @@ -751,7 +751,7 @@ def get_power(params: PhysicsMethodParams): except mdsExceptions.TreeException: kwa["p_ohm"] = np.full(len(params.times), np.nan) # Plasma magnetic energy, and respective time base - kwa["wmhd"], kwa["efit_time"] = params.data_conn.get_data_with_dims( + kwa["wmhd"], kwa["efit_time"] = params.get_data_with_dims( r"\efit_aeqdsk:wplasm", tree_name="_efit_tree" ) # [J], [s] return CmodPhysicsMethods._get_power(params.times, **kwa) @@ -807,13 +807,13 @@ def get_kappa_area(params: PhysicsMethodParams): - original source: [get_kappa_area.m](https://github.com/MIT-PSFC/disruption-py/ blob/matlab/CMOD/matlab-core/get_kappa_area.m) """ - aminor = params.data_conn.get_data( + aminor = params.get_data( r"\efit_aeqdsk:aout/100", tree_name="_efit_tree" ) # [m] - area = params.data_conn.get_data( + area = params.get_data( r"\efit_aeqdsk:areao/1e4", tree_name="_efit_tree" ) # [m^2] - times = params.data_conn.get_data( + times = params.get_data( r"\efit_aeqdsk:time", tree_name="_efit_tree" ) # [s] @@ -891,11 +891,11 @@ def get_n_equal_1_amplitude(params: PhysicsMethodParams): a = [] path = r"\mag_bp_coils." - bp_node_names = params.data_conn.get_data( + bp_node_names = params.get_data( f"{path}nodename", tree_name="magnetics" ) - phi = params.data_conn.get_data(f"{path}phi", tree_name="magnetics") # [degree] - btor_pickup_coeffs = params.data_conn.get_data( + phi = params.get_data(f"{path}phi", tree_name="magnetics") # [degree] + btor_pickup_coeffs = params.get_data( f"{path}btor_pickup", tree_name="magnetics" ) # [dimensionless] _, bp13_indices, _ = np.intersect1d( @@ -903,7 +903,7 @@ def get_n_equal_1_amplitude(params: PhysicsMethodParams): ) bp13_phi = phi[bp13_indices] + 360 # INFO bp13_btor_pickup_coeffs = btor_pickup_coeffs[bp13_indices] - btor, t_mag = params.data_conn.get_data_with_dims( + btor, t_mag = params.get_data_with_dims( r"\btor", tree_name="magnetics" ) # [T], [s] # Toroidal power supply takes time to turn on, from ~ -1.8 and should be @@ -918,7 +918,7 @@ def get_n_equal_1_amplitude(params: PhysicsMethodParams): for i, bp13_name in enumerate(bp13_names): try: - signal = params.data_conn.get_data( + signal = params.get_data( path + bp13_name, tree_name="magnetics" ) # [T] # Sensor not available, skip @@ -1069,17 +1069,17 @@ def get_densities(params: PhysicsMethodParams): matlab/CMOD/matlab-core/get_densities.m) """ # Line-integrated density - n_e, t_n = params.data_conn.get_data_with_dims( + n_e, t_n = params.get_data_with_dims( r".tci.results:nl_04", tree_name="electrons" ) # [m^-3], [s] # Divide by chord length of ~0.6m to get line averaged density. # For future refernce, chord length is stored in # .01*\analysis::efit_aeqdsk:rco2v[3,*] n_e = np.squeeze(n_e) / 0.6 - ip, t_ip = params.data_conn.get_data_with_dims( + ip, t_ip = params.get_data_with_dims( r"\ip", tree_name="magnetics" ) # [A], [s] - a_minor, t_a = params.data_conn.get_data_with_dims( + a_minor, t_a = params.get_data_with_dims( r"\efit_aeqdsk:aout/100", tree_name="_efit_tree" ) # [m], [s] @@ -1129,7 +1129,7 @@ def get_efc_current(params: PhysicsMethodParams): - original source: [get_efc_current.m](https://github.com/MIT-PSFC/disruption-py/ blob/matlab/CMOD/matlab-core/get_efc_current.m) """ - iefc, t_iefc = params.data_conn.get_data_with_dims( + iefc, t_iefc = params.get_data_with_dims( r"\efc:u_bus_r_cur", tree_name="engineering" ) # [A], [s] return CmodPhysicsMethods._get_efc_current(params.times, iefc, t_iefc) @@ -1243,13 +1243,13 @@ def get_ts_parameters(params: PhysicsMethodParams): # dependent dimensions being time and z (vertical coordinate) node_path = ".yag_new.results.profiles" - ts_data, ts_time = params.data_conn.get_data_with_dims( + ts_data, ts_time = params.get_data_with_dims( f"{node_path}:te_rz", tree_name="electrons" ) # [keV], [s] - ts_z = params.data_conn.get_data( + ts_z = params.get_data( f"{node_path}:z_sorted", tree_name="electrons" ) # [m] - ts_error = params.data_conn.get_data( + ts_error = params.get_data( f"{node_path}:te_err", tree_name="electrons" ) # [keV] @@ -1404,38 +1404,38 @@ def get_peaking_factors(params: PhysicsMethodParams): raise CalculationError("Shot is on blacklist") # Fetch data # Get EFIT geometry data - z0 = params.data_conn.get_data( + z0 = params.get_data( r"\efit_aeqdsk:zmagx/100", tree_name="_efit_tree" ) # [m] - kappa = params.data_conn.get_data( + kappa = params.get_data( r"\efit_aeqdsk:kappa", tree_name="_efit_tree" ) # [dimensionless] - aminor, efit_time = params.data_conn.get_data_with_dims( + aminor, efit_time = params.get_data_with_dims( r"\efit_aeqdsk:aout/100", tree_name="_efit_tree" ) # [m], [s] bminor = aminor * kappa # Get Te data and TS time basis node_ext = ".yag_new.results.profiles" - ts_te_core, ts_time = params.data_conn.get_data_with_dims( + ts_te_core, ts_time = params.get_data_with_dims( f"{node_ext}:te_rz", tree_name="electrons" ) # [keV], [s] ts_te_core = ts_te_core * 1000 # [keV] -> [eV] - ts_te_edge = params.data_conn.get_data(r"\ts_te") # [eV] + ts_te_edge = params.get_data(r"\ts_te") # [eV] ts_te = np.concatenate((ts_te_core, ts_te_edge)) * 11600 # [eV] -> [K] # Get ne data - ts_ne_core = params.data_conn.get_data( + ts_ne_core = params.get_data( f"{node_ext}:ne_rz", tree_name="electrons" ) # [m^-3] - ts_ne_edge = params.data_conn.get_data(r"\ts_ne") # [m^-3] + ts_ne_edge = params.get_data(r"\ts_ne") # [m^-3] ts_ne = np.concatenate((ts_ne_core, ts_ne_edge)) # Get TS chord positions - ts_z_core = params.data_conn.get_data( + ts_z_core = params.get_data( f"{node_ext}:z_sorted", tree_name="electrons" ) # [m] - ts_z_edge = params.data_conn.get_data(r"\fiber_z", tree_name="electrons") # [m] + ts_z_edge = params.get_data(r"\fiber_z", tree_name="electrons") # [m] ts_z = np.concatenate((ts_z_core, ts_z_edge)) # Make sure that there are equal numbers of edge position and edge temperature points if len(ts_z_edge) != ts_te_edge.shape[0]: @@ -1793,19 +1793,19 @@ def get_te_profile_params_ece(params: PhysicsMethodParams): """ # Get magnetic axis data from EFIT - r0 = params.data_conn.get_data( + r0 = params.get_data( r"\efit_aeqdsk:rmagx/100", tree_name="_efit_tree" ) # [m] - aminor, efit_time = params.data_conn.get_data_with_dims( + aminor, efit_time = params.get_data_with_dims( r"\efit_aeqdsk:aout/100", tree_name="_efit_tree" ) # [m], [s] # Btor and LH Power used for filtering okay time slices - btor, t_mag = params.data_conn.get_data_with_dims( + btor, t_mag = params.get_data_with_dims( r"\btor", tree_name="magnetics" ) # [T], [s] try: - lh_power, lh_time = params.data_conn.get_data_with_dims( + lh_power, lh_time = params.get_data_with_dims( ".results:netpow", tree_name="lh" ) # [kW], [s] except mdsExceptions.MdsException: @@ -1817,15 +1817,15 @@ def get_te_profile_params_ece(params: PhysicsMethodParams): # Read in Te profile measurements from GPC2 (19 channels) node_path = ".gpc_2.results" - gpc2_te_data, gpc2_te_time = params.data_conn.get_data_with_dims( + gpc2_te_data, gpc2_te_time = params.get_data_with_dims( f"{node_path}:gpc2_te", tree_name="electrons" ) # [keV], [s] - gpc2_rad_data, gpc2_rad_time = params.data_conn.get_data_with_dims( + gpc2_rad_data, gpc2_rad_time = params.get_data_with_dims( f"{node_path}:radii", tree_name="electrons" ) # [m], [s] # Te0 from GPC2 useful to check for outliers in the GPC channels # which be caused by some artifact or systematic error - gpc2_te0 = params.data_conn.get_data(r"\gpc2_te0", tree_name="electrons") + gpc2_te0 = params.get_data(r"\gpc2_te0", tree_name="electrons") # Line average density [m^-3] to check for cutoffs densities = CmodPhysicsMethods.get_densities(params) n_e = densities["n_e"] @@ -1878,27 +1878,27 @@ def get_prad_peaking(params: PhysicsMethodParams): """ prad_peaking = np.full(len(params.times), np.nan) nan_output = {"prad_peaking": prad_peaking} - r0 = params.data_conn.get_data( + r0 = params.get_data( r"\efit_aeqdsk:rmagx/100", tree_name="_efit_tree" ) # [m] - z0 = params.data_conn.get_data( + z0 = params.get_data( r"\efit_aeqdsk:zmagx/100", tree_name="_efit_tree" ) # [m] - aminor, efit_time = params.data_conn.get_data_with_dims( + aminor, efit_time = params.get_data_with_dims( r"\efit_aeqdsk:aout/100", tree_name="_efit_tree" ) # [m], [s] got_axa = False try: - bright_axa, t_axa, r_axa = params.data_conn.get_data_with_dims( + bright_axa, t_axa, r_axa = params.get_data_with_dims( r"\TOP.BOLOMETER.RESULTS.DIODE.AXA:BRIGHT", tree_name="spectroscopy", dim_nums=[1, 0], ) # [W/m^2], [s], [m] - z_axa = params.data_conn.get_data( + z_axa = params.get_data( r"\TOP.BOLOMETER.DIODE_CALIB.AXA:Z_O", tree_name="spectroscopy", ) # [m] - good_axa = params.data_conn.get_data( + good_axa = params.get_data( r"\TOP.BOLOMETER.DIODE_CALIB.AXA:GOOD", tree_name="spectroscopy", ) # [index] @@ -1907,16 +1907,16 @@ def get_prad_peaking(params: PhysicsMethodParams): params.logger.debug("Failed to get AXA data") got_axj = False try: - bright_axj, t_axj, r_axj = params.data_conn.get_data_with_dims( + bright_axj, t_axj, r_axj = params.get_data_with_dims( r"\TOP.BOLOMETER.RESULTS.DIODE.AXJ:BRIGHT", tree_name="spectroscopy", dim_nums=[1, 0], ) # [W/m^2], [s], [m] - z_axj = params.data_conn.get_data( + z_axj = params.get_data( r"\TOP.BOLOMETER.DIODE_CALIB.AXJ:Z_O", tree_name="spectroscopy", ) # [m] - good_axj = params.data_conn.get_data( + good_axj = params.get_data( r"\TOP.BOLOMETER.DIODE_CALIB.AXJ:GOOD", tree_name="spectroscopy", ) # [index] @@ -1993,7 +1993,7 @@ def get_sxr_data(params: PhysicsMethodParams): blob/matlab/CMOD/matlab-core/get_sxr_data.m) """ - sxr, t_sxr = params.data_conn.get_data_with_dims( + sxr, t_sxr = params.get_data_with_dims( r"\top.brightnesses.array_1:chord_16", tree_name="xtomo", ) # [W/m^2], [s] @@ -2046,16 +2046,16 @@ def get_beta_normalized(params: PhysicsMethodParams): """ # Get signals from EFIT tree - beta_t, efittime = params.data_conn.get_data_with_dims( + beta_t, efittime = params.get_data_with_dims( r"\efit_aeqdsk:betat", tree_name="_efit_tree" ) # [%], [s] - ip = params.data_conn.get_data( + ip = params.get_data( r"\efit_aeqdsk:cpasma/1e6", tree_name="_efit_tree" ) # [MA] - aminor = params.data_conn.get_data( + aminor = params.get_data( r"\efit_aeqdsk:aout/100", tree_name="_efit_tree" ) # [m] - btor = params.data_conn.get_data( + btor = params.get_data( r"\efit_aeqdsk:btaxp", tree_name="_efit_tree" ) # [T] @@ -2106,7 +2106,7 @@ def get_surface_voltage(params: PhysicsMethodParams): are zeros for every shot. """ # Get signals from EFIT tree - sibdry, efit_time = params.data_conn.get_data_with_dims( + sibdry, efit_time = params.get_data_with_dims( r"\efit_aeqdsk:sibdry", tree_name="_efit_tree" ) # [V*s/rad], [s] diff --git a/disruption_py/machine/cmod/thomson.py b/disruption_py/machine/cmod/thomson.py index 6ecbcec9..76deac3c 100644 --- a/disruption_py/machine/cmod/thomson.py +++ b/disruption_py/machine/cmod/thomson.py @@ -41,10 +41,10 @@ def compare_ts_tci(params: PhysicsMethodParams, nlnum=4): nl_tci2 = [1e32] ts_time1 = [1e32] ts_time2 = [1e32] - (tci_time,) = params.data_conn.get_dims( + (tci_time,) = params.get_dims( ".YAG_NEW.RESULTS.PROFILES:NE_RZ", tree_name="electrons" ) - tci, tci_t = params.data_conn.get_data_with_dims( + tci, tci_t = params.get_data_with_dims( f".TCI.RESULTS:NL_{nlnum:02d}", tree_name="electrons" ) nlts, nlts_t = CmodThomsonDensityMeasure._integrate_ts_tci(params, nlnum) @@ -83,12 +83,12 @@ def _parse_yags(params: PhysicsMethodParams): tuple Tuple containing the number of YAGs and their indices. """ - nyag1 = params.data_conn.get_data(r"\knobs:pulses_q", tree_name="electrons") - nyag2 = params.data_conn.get_data(r"\knobs:pulses_q_2", tree_name="electrons") + nyag1 = params.get_data(r"\knobs:pulses_q", tree_name="electrons") + nyag2 = params.get_data(r"\knobs:pulses_q_2", tree_name="electrons") indices1 = -1 indices2 = -1 - dark = params.data_conn.get_data(r"\n_dark_prior", tree_name="electrons") - ntotal = params.data_conn.get_data(r"\n_total", tree_name="electrons") + dark = params.get_data(r"\n_dark_prior", tree_name="electrons") + ntotal = params.get_data(r"\n_total", tree_name="electrons") nt = ntotal - dark if nyag1 == 0: if nyag2 != 0: @@ -186,39 +186,39 @@ def _map_ts2tci(params: PhysicsMethodParams, nlnum): n_e_sig = [1e32] flag = 1 valid_indices, efit_times = CmodEfitMethods.efit_check(params) - ip = params.data_conn.get_data(r"\ip", "cmod") + ip = params.get_data(r"\ip", "cmod") if np.mean(ip) > 0: flag = 0 - efit_times = params.data_conn.get_data( + efit_times = params.get_data( r"\efit_aeqdsk:time", tree_name="_efit_tree" ) t1 = np.amin(efit_times) t2 = np.amax(efit_times) - psia, psia_t = params.data_conn.get_data_with_dims( + psia, psia_t = params.get_data_with_dims( r"\efit_aeqdsk:SIBDRY", tree_name="_efit_tree" ) psi_0 = params.data_conn.get(r"\efit_aeqdsk:SIMAGX", tree_name="_efit_tree") - nets_core, nets_core_t = params.data_conn.get_data_with_dims( + nets_core, nets_core_t = params.get_data_with_dims( ".YAG_NEW.RESULTS.PROFILES:NE_RZ", tree_name="electrons" ) - nets_core_err = params.data_conn.get_data( + nets_core_err = params.get_data( ".YAG_NEW.RESULTS.PROFILES:NE_ERR", tree_name="electrons" ) - zts_core = params.data_conn.get_data( + zts_core = params.get_data( ".YAG_NEW.RESULTS.PROFILES:Z_SORTED", tree_name="electrons" ) mts_core = len(zts_core) - zts_edge = params.data_conn.get_data(r"\fiber_z") + zts_edge = params.get_data(r"\fiber_z") mts_edge = len(zts_edge) try: - nets_edge = params.data_conn.get_data(r"\ts_ne") - nets_edge_err = params.data_conn.get_data(r"\ts_ne_err") + nets_edge = params.get_data(r"\ts_ne") + nets_edge_err = params.get_data(r"\ts_ne_err") except mdsExceptions.MdsException: nets_edge = np.zeros((len(nets_core[:, 1]), mts_edge)) nets_edge_err = nets_edge + 1e20 mts = mts_core + mts_edge rts = params.data_conn.get(".YAG.RESULTS.PARAM:R") + np.zeros((1, mts)) - rtci = params.data_conn.get_data(".tci.results:rad") + rtci = params.get_data(".tci.results:rad") nts = len(nets_core_t) zts = np.zeros((1, mts)) zts[:, :mts_core] = zts_core @@ -298,7 +298,7 @@ def _efit_rz2psi(params: PhysicsMethodParams, r, z, t, tree="analysis"): r = r.flatten() z = z.flatten() psi = np.full((len(r), len(t)), np.nan) - psirz, rgrid, zgrid, times = params.data_conn.get_data_with_dims( + psirz, rgrid, zgrid, times = params.get_data_with_dims( r"\efit_geqdsk:psirz", tree_name=tree, dim_nums=[0, 1, 2] ) rgrid, zgrid = np.meshgrid(rgrid, zgrid) diff --git a/disruption_py/machine/d3d/efit.py b/disruption_py/machine/d3d/efit.py index 78cbc25b..9b256be0 100644 --- a/disruption_py/machine/d3d/efit.py +++ b/disruption_py/machine/d3d/efit.py @@ -61,11 +61,11 @@ def get_efit_parameters(params: PhysicsMethodParams): A dictionary containing the EFIT parameters and their derivatives. """ efit_data = { - k: params.data_conn.get_data(v, tree_name="_efit_tree") + k: params.get_data(v, tree_name="_efit_tree") for k, v in D3DEfitMethods.efit_cols.items() } efit_time = ( - params.data_conn.get_data(r"\efit_a_eqdsk:atime", tree_name="_efit_tree") + params.get_data(r"\efit_a_eqdsk:atime", tree_name="_efit_tree") / 1.0e3 ) # [ms] -> [s] @@ -108,11 +108,11 @@ def get_rt_efit_parameters(params: PhysicsMethodParams): A dictionary containing the real-time EFIT parameters. """ efit_data = { - k: params.data_conn.get_data(v, tree_name="efitrt1") + k: params.get_data(v, tree_name="efitrt1") for k, v in D3DEfitMethods.rt_efit_cols.items() } efit_time = ( - params.data_conn.get_data(r"\efit_a_eqdsk:atime", tree_name="efitrt1") + params.get_data(r"\efit_a_eqdsk:atime", tree_name="efitrt1") / 1.0e3 ) # [ms] -> [s] # EFIT reconstructions are sometimes invalid, particularly when very close diff --git a/disruption_py/machine/d3d/physics.py b/disruption_py/machine/d3d/physics.py index 20d306f3..9ea7406b 100644 --- a/disruption_py/machine/d3d/physics.py +++ b/disruption_py/machine/d3d/physics.py @@ -81,7 +81,7 @@ def get_h98(params: PhysicsMethodParams): "h98": [np.nan], } try: - h_98, t_h_98 = params.data_conn.get_data_with_dims( + h_98, t_h_98 = params.get_data_with_dims( r"\H_THH98Y2", tree_name="transport" ) t_h_98 /= 1e3 # [ms] -> [s] @@ -119,7 +119,7 @@ def get_h_alpha(params: PhysicsMethodParams): "h_alpha": [np.nan], } try: - h_alpha, t_h_alpha = params.data_conn.get_data_with_dims( + h_alpha, t_h_alpha = params.get_data_with_dims( r"\fs04", tree_name="d3d" ) t_h_alpha /= 1e3 # [ms] -> [s] @@ -166,7 +166,7 @@ def get_power_parameters(params: PhysicsMethodParams): """ # Get neutral beam injected power try: - p_nbi, t_nbi = params.data_conn.get_data_with_dims( + p_nbi, t_nbi = params.get_data_with_dims( r"\top.nb:pinj", tree_name="d3d" ) t_nbi /= 1e3 # [ms] -> [s] @@ -191,7 +191,7 @@ def get_power_parameters(params: PhysicsMethodParams): # Get electron cyclotron heating (ECH) power. It's point data, so it's not # stored in an MDSplus tree try: - p_ech, t_ech = params.data_conn.get_data_with_dims( + p_ech, t_ech = params.get_data_with_dims( r"\top.ech.total:echpwrc", tree_name="rf" ) t_ech /= 1e3 # [ms] -> [s] @@ -234,7 +234,7 @@ def get_power_parameters(params: PhysicsMethodParams): smoothing_window = 0.010 # [s] try: - bol_prm, _ = params.data_conn.get_data_with_dims( + bol_prm, _ = params.get_data_with_dims( r"\bol_prm", tree_name="bolom" ) except mdsExceptions.MdsException as e: @@ -245,11 +245,11 @@ def get_power_parameters(params: PhysicsMethodParams): bol_channels = upper_channels + lower_channels bol_signals = [] for i in range(48): - bol_signal = params.data_conn.get_data( + bol_signal = params.get_data( rf"\top.raw:{bol_channels[i]}", tree_name="bolom" ) bol_signals.append(bol_signal) - bol_time = params.data_conn.get_dims( + bol_time = params.get_dims( rf"\top.raw:{bol_channels[0]}", tree_name="bolom" )[0] bol_time /= 1e3 # [ms] -> [s] @@ -320,14 +320,14 @@ def get_ohmic_parameters(params: PhysicsMethodParams): - issues: #[229](https://github.com/MIT-PSFC/disruption-py/issues/229) """ # Get edge loop voltage and smooth it a bit with a median filter - v_loop, t_v_loop = params.data_conn.get_data_with_dims( + v_loop, t_v_loop = params.get_data_with_dims( f'ptdata("vloopb", {params.shot_id})' ) t_v_loop /= 1e3 # [ms] -> [s] v_loop = scipy.signal.medfilt(v_loop, 11) v_loop = interp1(t_v_loop, v_loop, params.times, "linear") # Get plasma current - ip, t_ip = params.data_conn.get_data_with_dims( + ip, t_ip = params.get_data_with_dims( f"ptdata('ip', {params.shot_id})" ) t_ip /= 1e3 # [ms] -> [s] @@ -349,19 +349,19 @@ def get_ohmic_parameters(params: PhysicsMethodParams): ends_type=1, slew_rate=0, ) - li, t_li = params.data_conn.get_data_with_dims( + li, t_li = params.get_data_with_dims( r"\efit_a_eqdsk:li", tree_name="_efit_tree" ) t_li /= 1e3 # Use chisq to determine which time slices are invalid - chisq = params.data_conn.get_data( + chisq = params.get_data( r"\efit_a_eqdsk:chisq", tree_name="_efit_tree" ) # Filter out invalid indices of efit reconstruction (invalid_indices,) = np.where(chisq > 50) li[invalid_indices] = np.nan - r_0, t_r0 = params.data_conn.get_data_with_dims( + r_0, t_r0 = params.get_data_with_dims( r"\top.results.geqdsk:rmaxis", tree_name="_efit_tree" ) # [m], [ms] t_r0 /= 1e3 # [ms] -> [s] @@ -410,7 +410,7 @@ def get_density_parameters(params: PhysicsMethodParams): - pull requests: #[249](https://github.com/MIT-PSFC/disruption-py/pull/249) """ try: - ne, t_ne = params.data_conn.get_data_with_dims( + ne, t_ne = params.get_data_with_dims( r"\density", tree_name="_efit_tree" ) except mdsExceptions.MdsException: @@ -425,8 +425,8 @@ def get_density_parameters(params: PhysicsMethodParams): # - r"\efit01:density" gives ne = array([4.06199e19]), t_ne = array([1800]) # - "\d3d:denv2" gives actual density data if not np.isfinite(ne).any() or len(ne) < 2: - ne, t_ne = params.data_conn.get_data_with_dims(r"\denv2", tree_name="d3d") - tree_name = params.data_conn.get_tree_name_of_nickname("_efit_tree") + ne, t_ne = params.get_data_with_dims(r"\denv2", tree_name="d3d") + tree_name = params.get_tree_name_of_nickname("_efit_tree") params.logger.verbose( rf"density: data from \{tree_name}:density is either empty or invalid." r" Use \d3d:denv2 instead." @@ -451,13 +451,13 @@ def get_density_parameters(params: PhysicsMethodParams): bounds_error=False, ) try: - ip, t_ip = params.data_conn.get_data_with_dims( + ip, t_ip = params.get_data_with_dims( f"ptdata('ip', {params.shot_id})" ) # [A], [ms] t_ip = t_ip / 1.0e3 # [ms] -> [s] ipsign = np.sign(np.sum(ip)) ip = interp1(t_ip, ip * ipsign, params.times, "linear") # positive definite - a_minor, t_a = params.data_conn.get_data_with_dims( + a_minor, t_a = params.get_data_with_dims( r"\efit_a_eqdsk:aminor", tree_name="_efit_tree" ) # [m], [ms] t_a = t_a / 1.0e3 # [ms] -> [s] @@ -514,7 +514,7 @@ def get_rt_density_parameters(params: PhysicsMethodParams): /disruption-py/blob/matlab/DIII-D/get_density_parameters_RT.m) - pull requests: #[251](https://github.com/MIT-PSFC/disruption-py/pull/251) """ - ne_rt, t_ne_rt = params.data_conn.get_data_with_dims( + ne_rt, t_ne_rt = params.get_data_with_dims( f"ptdata('dssdenest', {params.shot_id})" ) # [10^19 m^-3] t_ne_rt = t_ne_rt / 1.0e3 # [ms] to [s] @@ -526,12 +526,12 @@ def get_rt_density_parameters(params: PhysicsMethodParams): # Get real time ip to calculate the Greenwald density try: - ip_rt, t_ip_rt = params.data_conn.get_data_with_dims( + ip_rt, t_ip_rt = params.get_data_with_dims( f"ptdata('ipsip', {params.shot_id})" ) # [MA], [ms] t_ip_rt = t_ip_rt / 1.0e3 # [ms] to [s] except mdsExceptions.MdsException: - ip_rt, t_ip_rt = params.data_conn.get_data_with_dims( + ip_rt, t_ip_rt = params.get_data_with_dims( f"ptdata('ipspr15v', {params.shot_id})" ) # [volts; 2 V/MA], [ms] t_ip_rt = t_ip_rt / 1.0e3 # [ms] to [s] @@ -548,7 +548,7 @@ def get_rt_density_parameters(params: PhysicsMethodParams): # For the real-time (RT) signals, read from the EFITRT1 tree try: - a_minor_rt, t_a_rt = params.data_conn.get_data_with_dims( + a_minor_rt, t_a_rt = params.get_data_with_dims( r"\efit_a_eqdsk:aminor", tree_name="efitrt1" ) # [m], [ms] t_a_rt = t_a_rt / 1.0e3 # [ms] -> [s] @@ -619,7 +619,7 @@ def get_ip_parameters(params: PhysicsMethodParams): ip_error = np.full(len(params.times), np.nan) # Get measured plasma current parameters try: - ip, t_ip = params.data_conn.get_data_with_dims( + ip, t_ip = params.get_data_with_dims( f"ptdata('ip', {params.shot_id})" ) # [A], [ms] t_ip = t_ip / 1.0e3 # [ms] -> [s] @@ -631,7 +631,7 @@ def get_ip_parameters(params: PhysicsMethodParams): params.logger.opt(exception=True).debug(e) # Get programmed plasma current parameters try: - ip_prog, t_ip_prog = params.data_conn.get_data_with_dims( + ip_prog, t_ip_prog = params.get_data_with_dims( f"ptdata('iptipp', {params.shot_id})" ) # [A], [ms] t_ip_prog = t_ip_prog / 1.0e3 # [ms] -> [s] @@ -653,7 +653,7 @@ def get_ip_parameters(params: PhysicsMethodParams): # Anything else: not in normal Ip feedback mode. In this case, the # 'ip_prog' signal is irrelevant, and therefore 'ip_error' is not defined. try: - ipimode, t_ipimode = params.data_conn.get_data_with_dims( + ipimode, t_ipimode = params.get_data_with_dims( f"ptdata('ipimode', {params.shot_id})" ) t_ipimode = t_ipimode / 1.0e3 # [ms] -> [s] @@ -676,7 +676,7 @@ def get_ip_parameters(params: PhysicsMethodParams): # PCS feedback control of Ip is not being applied. Therefore the # 'ip_error' parameter is undefined for these times. try: - epsoff, t_epsoff = params.data_conn.get_data_with_dims( + epsoff, t_epsoff = params.get_data_with_dims( f"ptdata('epsoff', {params.shot_id})" ) t_epsoff = t_epsoff / 1.0e3 # [ms] -> [s] @@ -743,7 +743,7 @@ def get_rt_ip_parameters(params: PhysicsMethodParams): # Get measured plasma current parameters # TODO: Why open d3d and not the rt efit tree? try: - ip_rt, t_ip_rt = params.data_conn.get_data_with_dims( + ip_rt, t_ip_rt = params.get_data_with_dims( f"ptdata('ipsip', {params.shot_id})" ) # [MA], [ms] t_ip_rt = t_ip_rt / 1.0e3 # [ms] -> [s] @@ -758,7 +758,7 @@ def get_rt_ip_parameters(params: PhysicsMethodParams): params.logger.opt(exception=True).debug(e) # Get programmed plasma current parameters try: - ip_prog_rt, t_ip_prog_rt = params.data_conn.get_data_with_dims( + ip_prog_rt, t_ip_prog_rt = params.get_data_with_dims( f"ptdata('ipsiptargt', {params.shot_id})" ) # [MA], [ms] t_ip_prog_rt = t_ip_prog_rt / 1.0e3 # [ms] -> [s] @@ -774,7 +774,7 @@ def get_rt_ip_parameters(params: PhysicsMethodParams): ) params.logger.opt(exception=True).debug(e) try: - ip_error_rt, t_ip_error_rt = params.data_conn.get_data_with_dims( + ip_error_rt, t_ip_error_rt = params.get_data_with_dims( f"ptdata('ipeecoil', {params.shot_id})" ) # [MA], [ms] t_ip_error_rt = t_ip_error_rt / 1.0e3 # [ms] to [s] @@ -795,7 +795,7 @@ def get_rt_ip_parameters(params: PhysicsMethodParams): # Anything else: not in normal Ip feedback mode. In this case, the # 'ip_prog' signal is irrelevant, and therefore 'ip_error' is not defined. try: - ipimode, t_ipimode = params.data_conn.get_data_with_dims( + ipimode, t_ipimode = params.get_data_with_dims( f"ptdata('ipimode', {params.shot_id})" ) t_ipimode = t_ipimode / 1.0e3 # [ms] -> [s] @@ -814,7 +814,7 @@ def get_rt_ip_parameters(params: PhysicsMethodParams): # PCS feedback control of Ip is not being applied. Therefore the # 'ip_error' parameter is undefined for these times. try: - epsoff, t_epsoff = params.data_conn.get_data_with_dims( + epsoff, t_epsoff = params.get_data_with_dims( f"ptdata('epsoff', {params.shot_id})" ) t_epsoff = t_epsoff / 1.0e3 # [ms] -> [s] @@ -877,7 +877,7 @@ def get_z_parameters(params: PhysicsMethodParams): """ nominal_flattop_radius = 0.59 # Get z_cur - z_cur, t_z_cur = params.data_conn.get_data_with_dims( + z_cur, t_z_cur = params.get_data_with_dims( f"ptdata('vpszp', {params.shot_id})" ) t_z_cur = t_z_cur / 1.0e3 # [ms] -> [s] @@ -885,11 +885,11 @@ def get_z_parameters(params: PhysicsMethodParams): z_cur = interp1(t_z_cur, z_cur, params.times, "linear") # Compute z_cur_norm try: - a_minor, t_a = params.data_conn.get_data_with_dims( + a_minor, t_a = params.get_data_with_dims( r"\efit_a_eqdsk:aminor", tree_name="_efit_tree" ) # [m], [ms] t_a = t_a / 1.0e3 # [ms] -> [s] - chisq = params.data_conn.get_data( + chisq = params.get_data( r"\efit_a_eqdsk:chisq", tree_name="_efit_tree" ) (invalid_indices,) = np.where(chisq > 50) @@ -926,13 +926,13 @@ def get_n1rms_parameters(params: PhysicsMethodParams): - pull requests: #[257](https://github.com/MIT-PSFC/disruption-py/pull/257) """ # Get n1rms signal from d3d tree - n1rms, t_n1rms = params.data_conn.get_data_with_dims(r"\n1rms", tree_name="d3d") + n1rms, t_n1rms = params.get_data_with_dims(r"\n1rms", tree_name="d3d") n1rms *= 1.0e-4 # Gauss -> Tesla t_n1rms /= 1e3 # [ms] -> [s] n1rms = interp1(t_n1rms, n1rms, params.times) # Calculate n1rms_norm try: - b_tor, t_b_tor = params.data_conn.get_data_with_dims( + b_tor, t_b_tor = params.get_data_with_dims( f"ptdata('bt', {params.shot_id})" ) t_b_tor /= 1e3 # [ms] -> [s] @@ -1037,12 +1037,12 @@ def get_peaking_factors(params: PhysicsMethodParams): # Get precomputed rad_cva & rad_xdiv data stored in ptdata tree calculate_prad_pf = False try: - rad_cva, t_rad_cva = params.data_conn.get_data_with_dims( + rad_cva, t_rad_cva = params.get_data_with_dims( f"ptdata('dpsrrdcva', {params.shot_id})" ) # [], [ms] t_rad_cva /= 1e3 # [ms] -> [s] rad_cva = interp1(t_rad_cva, rad_cva, params.times) - rad_xdiv, t_rad_xdiv = params.data_conn.get_data_with_dims( + rad_xdiv, t_rad_xdiv = params.get_data_with_dims( f"ptdata('dpsrrdxdiv', {params.shot_id})" ) # [], [ms] t_rad_xdiv /= 1e3 # [ms] -> [s] @@ -1280,7 +1280,7 @@ def get_zeff_parameters(params: PhysicsMethodParams): A dictionary containing `z_eff` """ # Get Zeff - zeff, t_zeff = params.data_conn.get_data_with_dims( + zeff, t_zeff = params.get_data_with_dims( r"\top.spectroscopy.vb.zeff:zeff", tree_name="d3d" ) t_zeff = t_zeff / 1.0e3 # [ms] -> [s] @@ -1327,14 +1327,14 @@ def get_kappa_area(params: PhysicsMethodParams): /blob/matlab/DIII-D/get_kappa_area.m) - pull requests: #[256](https://github.com/MIT-PSFC/disruption-py/pull/256) """ - a_minor = params.data_conn.get_data( + a_minor = params.get_data( r"\efit_a_eqdsk:aminor", tree_name="_efit_tree" ) - area = params.data_conn.get_data(r"\efit_a_eqdsk:area", tree_name="_efit_tree") - chisq = params.data_conn.get_data( + area = params.get_data(r"\efit_a_eqdsk:area", tree_name="_efit_tree") + chisq = params.get_data( r"\efit_a_eqdsk:chisq", tree_name="_efit_tree" ) - t = params.data_conn.get_data(r"\efit_a_eqdsk:atime", tree_name="_efit_tree") + t = params.get_data(r"\efit_a_eqdsk:atime", tree_name="_efit_tree") t /= 1e3 # [ms] -> [s] kappa_area = area / (np.pi * a_minor**2) invalid_indices = np.where(chisq > 50) @@ -1369,16 +1369,16 @@ def get_shape_parameters(params: PhysicsMethodParams): - pull requests: #[258](https://github.com/MIT-PSFC/disruption-py/pull/258) """ # Get efit_time - efit_time = params.data_conn.get_data( + efit_time = params.get_data( r"\efit_a_eqdsk:atime", tree_name="_efit_tree" ) efit_time /= 1e3 # [ms] -> [s] # Compute triangularity try: - tritop = params.data_conn.get_data( + tritop = params.get_data( r"\efit_a_eqdsk:tritop", tree_name="_efit_tree" ) # meters - tribot = params.data_conn.get_data( + tribot = params.get_data( r"\efit_a_eqdsk:tribot", tree_name="_efit_tree" ) # meters delta = (tritop + tribot) / 2.0 @@ -1388,10 +1388,10 @@ def get_shape_parameters(params: PhysicsMethodParams): delta = None # Compute squareness try: - sqfod = params.data_conn.get_data( + sqfod = params.get_data( r"\efit_a_eqdsk:sqfod", tree_name="_efit_tree" ) - sqfou = params.data_conn.get_data( + sqfou = params.get_data( r"\efit_a_eqdsk:sqfou", tree_name="_efit_tree" ) squareness = (sqfod + sqfou) / 2.0 @@ -1401,7 +1401,7 @@ def get_shape_parameters(params: PhysicsMethodParams): squareness = None # Get aminor try: - aminor = params.data_conn.get_data( + aminor = params.get_data( r"\efit_a_eqdsk:aminor", tree_name="_efit_tree" ) except mdsExceptions.MdsException as e: @@ -1410,7 +1410,7 @@ def get_shape_parameters(params: PhysicsMethodParams): aminor = None # Check chisq for invalid indices try: - chisq = params.data_conn.get_data( + chisq = params.get_data( r"\efit_a_eqdsk:chisq", tree_name="_efit_tree" ) invalid_indices = np.where(chisq > 50) @@ -1493,7 +1493,7 @@ def _get_ne_te( lasers[laser] = {} sub_tree = f"{mds_path}{laser}" try: - (t_sub_tree,) = params.data_conn.get_dims( + (t_sub_tree,) = params.get_dims( f"{sub_tree}:temp", tree_name="electrons" ) # lasers[laser]['time'] gets overwritten in the loop later @@ -1517,7 +1517,7 @@ def _get_ne_te( } for node, name in child_nodes.items(): try: - lasers[laser][node] = params.data_conn.get_data( + lasers[laser][node] = params.get_data( f"{sub_tree}:{name}", tree_name="electrons" ) except mdsExceptions.MdsException as e: @@ -1594,7 +1594,7 @@ def _get_p_rad(params: PhysicsMethodParams, fan="custom", smoothing_window=50): return False # Get bolometry data - bol_prm, _ = params.data_conn.get_data_with_dims(r"\bol_prm", tree_name="bolom") + bol_prm, _ = params.get_data_with_dims(r"\bol_prm", tree_name="bolom") upper_channels = [f"bol_u{i+1:02d}_v" for i in range(24)] lower_channels = [f"bol_l{i+1:02d}_v" for i in range(24)] bol_channels = upper_channels + lower_channels @@ -1603,7 +1603,7 @@ def _get_p_rad(params: PhysicsMethodParams, fan="custom", smoothing_window=50): [] ) # TODO: Decide whether to actually use all bol_times instead of just first one for i in range(48): - bol_signal, bol_time = params.data_conn.get_data_with_dims( + bol_signal, bol_time = params.get_data_with_dims( rf"\top.raw:{bol_channels[i]}", tree_name="bolom" ) bol_time /= 1e3 # [ms] -> [s] @@ -1618,7 +1618,7 @@ def _get_p_rad(params: PhysicsMethodParams, fan="custom", smoothing_window=50): smoothing_window, ) b_struct = matlab_power(a_struct) - r_major_axis, efit_time = params.data_conn.get_data_with_dims( + r_major_axis, efit_time = params.get_data_with_dims( r"\top.results.geqdsk:rmaxis", tree_name="_efit_tree" ) efit_time /= 1e3 # [ms] -> [s] @@ -1709,13 +1709,13 @@ def _get_efit_dict(params: PhysicsMethodParams): """ path = r"\top.results.geqdsk:" nodes = ["z", "r", "rhovn", "psirz", "zmaxis", "ssimag", "ssibry"] - (efit_dict_time,) = params.data_conn.get_dims( + (efit_dict_time,) = params.get_dims( f"{path}psirz", tree_name="_efit_tree", dim_nums=[2] ) efit_dict = {"time": efit_dict_time / 1e3} # [ms] -> [s] for node in nodes: try: - efit_dict[node] = params.data_conn.get_data( + efit_dict[node] = params.get_data( f"{path}{node}", tree_name="_efit_tree" ) except mdsExceptions.MdsException as e: diff --git a/disruption_py/machine/d3d/util.py b/disruption_py/machine/d3d/util.py index 9ba031c2..266b3fd9 100644 --- a/disruption_py/machine/d3d/util.py +++ b/disruption_py/machine/d3d/util.py @@ -33,7 +33,7 @@ def get_polarity(params: PhysicsMethodParams): polarity value, -1 or 1. """ polarity = np.unique( - params.data_conn.get_data(f"ptdata('iptdirect', {params.shot_id})") + params.get_data(f"ptdata('iptdirect', {params.shot_id})") ) if len(polarity) > 1: params.logger.info( diff --git a/disruption_py/machine/east/efit.py b/disruption_py/machine/east/efit.py index 857658cb..56b1de1a 100644 --- a/disruption_py/machine/east/efit.py +++ b/disruption_py/machine/east/efit.py @@ -75,10 +75,10 @@ def get_efit_parameters(params: PhysicsMethodParams): # pylint: disable=duplicate-code efit_data = { - k: params.data_conn.get_data(v, tree_name="_efit_tree") + k: params.get_data(v, tree_name="_efit_tree") for k, v in EastEfitMethods.efit_cols.items() } - efit_time = params.data_conn.get_data( + efit_time = params.get_data( r"\efit_aeqdsk:atime", tree_name="_efit_tree" ) # TODO: [unit?] @@ -122,7 +122,7 @@ def get_pefit_parameters(params: PhysicsMethodParams): """ # pylint: disable=duplicate-code - efit_time = params.data_conn.get_data( + efit_time = params.get_data( r"\efit_a_eqdsk:atime", tree_name="pefit_east" ) # TODO: [unit?] @@ -130,7 +130,7 @@ def get_pefit_parameters(params: PhysicsMethodParams): efit_time, unique_indices = np.unique(efit_time, return_index=True) efit_data = { - k: params.data_conn.get_data(v, tree_name="pefit_east")[unique_indices] + k: params.get_data(v, tree_name="pefit_east")[unique_indices] for k, v in EastEfitMethods.pefit_cols.items() } diff --git a/disruption_py/machine/east/physics.py b/disruption_py/machine/east/physics.py index 35434dd7..b58ec7bf 100644 --- a/disruption_py/machine/east/physics.py +++ b/disruption_py/machine/east/physics.py @@ -116,12 +116,12 @@ def get_ip_parameters(params: PhysicsMethodParams): # Start with \lmtipref (standard Ip programming), which is defined for the # entire shot, and defines the timebase for the programmed Ip signal. - ip_prog, ip_prog_time = params.data_conn.get_data_with_dims( + ip_prog, ip_prog_time = params.get_data_with_dims( r"\lmtipref*1e6", tree_name="pcs_east" ) # [A], [s] # Now check to see if there is a transition to isoflux control - sytps1, time_sytps1 = params.data_conn.get_data_with_dims( + sytps1, time_sytps1 = params.get_data_with_dims( r"\sytps1", tree_name="pcs_east" ) (sytps1_indices,) = np.where(sytps1 > 0) @@ -132,7 +132,7 @@ def get_ip_parameters(params: PhysicsMethodParams): if len(sytps1_indices) > 0 and time_sytps1[0] <= ip_prog_time[-1]: try: - ietip, ietip_time = params.data_conn.get_data_with_dims( + ietip, ietip_time = params.get_data_with_dims( r"\ietip", tree_name="pcs_east" ) ietip = interp1(ietip_time, ietip, ip_prog_time) @@ -140,7 +140,7 @@ def get_ip_parameters(params: PhysicsMethodParams): ietip = np.full(len(ip_prog_time), np.nan) try: - idtip, idtip_time = params.data_conn.get_data_with_dims( + idtip, idtip_time = params.get_data_with_dims( r"\idtip", tree_name="pcs_east" ) idtip = interp1(idtip_time, idtip, ip_prog_time) @@ -148,7 +148,7 @@ def get_ip_parameters(params: PhysicsMethodParams): idtip = np.full(len(ip_prog_time), np.nan) try: - istip, istip_time = params.data_conn.get_data_with_dims( + istip, istip_time = params.get_data_with_dims( r"\istip", tree_name="pcs_east" ) istip = interp1(istip_time, istip, ip_prog_time) @@ -156,7 +156,7 @@ def get_ip_parameters(params: PhysicsMethodParams): istip = np.full(len(ip_prog_time), np.nan) try: - iutip, iutip_time = params.data_conn.get_data_with_dims( + iutip, iutip_time = params.get_data_with_dims( r"\iutip", tree_name="pcs_east" ) iutip = interp1(iutip_time, iutip, ip_prog_time) @@ -242,14 +242,14 @@ def get_v_loop(params: PhysicsMethodParams): # Get "\vp1_s" signal from the EAST tree. (This signal is a sub-sampled # version of "vp1".) try: - v_loop, v_loop_time = params.data_conn.get_data_with_dims( + v_loop, v_loop_time = params.get_data_with_dims( r"\vp1_s", tree_name="east" ) except mdsExceptions.MdsException: params.logger.verbose( r"v_loop: Failed to get \vp1_s data. Use \pcvloop from pcs_east instead." ) - v_loop, v_loop_time = params.data_conn.get_data_with_dims( + v_loop, v_loop_time = params.get_data_with_dims( r"\pcvloop", tree_name="pcs_east" ) # [V] @@ -301,7 +301,7 @@ def get_z_error(params: PhysicsMethodParams): blob/matlab/EAST/get_Z_error.m) """ # Read in the calculated zcur from EFIT - zcur, zcur_time = params.data_conn.get_data_with_dims( + zcur, zcur_time = params.get_data_with_dims( r"\efit_aeqdsk:zcur", tree_name="_efit_tree" ) # [A], [s] # Deal with rare bug @@ -310,15 +310,15 @@ def get_z_error(params: PhysicsMethodParams): # Read in aminor from EFIT # TODO: use \aminor or \aout? -- MATLAB: \aminor - aminor = params.data_conn.get_data(r"\aminor", tree_name="_efit_tree") # [m] + aminor = params.get_data(r"\aminor", tree_name="_efit_tree") # [m] aminor = aminor[unique_indices] # Next, get the programmed/requested/target Z from PCS, and the # calculated Z-centroid from PCS - z_prog, z_prog_time = params.data_conn.get_data_with_dims( + z_prog, z_prog_time = params.get_data_with_dims( r"\lmtzref/100", tree_name="pcs_east" ) # [m], [s] (Node says 'm' but it's wrong) - zcur_lmsz, lmsz_time = params.data_conn.get_data_with_dims( + zcur_lmsz, lmsz_time = params.get_data_with_dims( r"\lmsz", tree_name="pcs_east" ) # [m], [s] @@ -390,7 +390,7 @@ def get_density_parameters(params: PhysicsMethodParams): dn_dt = [np.nan] # Get the density and calculate dn_dt - ne, netime = params.data_conn.get_data_with_dims( + ne, netime = params.get_data_with_dims( r"\dfsdev*1e19", tree_name="pcs_east" ) # [m^-3], [s] dn_dt = np.gradient(ne, netime) # [m^-3/s] @@ -401,7 +401,7 @@ def get_density_parameters(params: PhysicsMethodParams): # Calculate Greenwald density # TODO: use \aminor or \aout? -- MATLAB: \aout - aminor, efittime = params.data_conn.get_data_with_dims( + aminor, efittime = params.get_data_with_dims( r"\efit_aeqdsk:aout", tree_name="_efit_tree" ) # [m], [s] aminor = interp1(efittime, aminor, params.times) @@ -425,7 +425,7 @@ def _get_raw_axuv_data(params: PhysicsMethodParams): these functions different. """ # Get XUV data - (xuvtime,) = params.data_conn.get_dims(r"\pxuv1", tree_name="east_1") # [s] + (xuvtime,) = params.get_dims(r"\pxuv1", tree_name="east_1") # [s] # There are 64 AXUV chords, arranged in 4 arrays of 16 channels each xuv = np.full((len(xuvtime), 64), np.nan) @@ -436,7 +436,7 @@ def _get_raw_axuv_data(params: PhysicsMethodParams): for iarray in range(4): for ichan in range(16): ichord = 16 * iarray + ichan - signal = params.data_conn.get_data( + signal = params.get_data( r"\pxuv" + str(ichord + 1), tree_name="east_1" ) # Subtract baseline @@ -586,7 +586,7 @@ def get_heating_power(nodes, tree): """ heating_power = np.zeros(params.times.shape) for node in nodes: - power_node, time_node = params.data_conn.get_data_with_dims( + power_node, time_node = params.get_data_with_dims( node, tree_name=tree ) heating_power += interp1( @@ -622,7 +622,7 @@ def get_heating_power(nodes, tree): # Get ECRH power try: - p_ecrh, ecrh_time = params.data_conn.get_data_with_dims( + p_ecrh, ecrh_time = params.get_data_with_dims( r"\pecrh1i*1e3", tree_name="analysis" ) # [W], [s] (baseline_indices,) = np.where(ecrh_time < 0) @@ -683,7 +683,7 @@ def get_heating_power(nodes, tree): # Get Wmhd, calculate dWmhd_dt, and calculate p_loss try: - wmhd, efittime = params.data_conn.get_data_with_dims( + wmhd, efittime = params.get_data_with_dims( r"\efit_aeqdsk:wmhd", tree_name="_efit_tree" ) # [W], [s] dwmhd_dt = np.gradient(wmhd, efittime) @@ -742,14 +742,14 @@ def get_p_ohm(params: PhysicsMethodParams): """ # Get raw signals try: - vloop, vloop_time = params.data_conn.get_data_with_dims( + vloop, vloop_time = params.get_data_with_dims( r"\pcvloop", tree_name="pcs_east" ) # [V] - li, li_time = params.data_conn.get_data_with_dims( + li, li_time = params.get_data_with_dims( r"\efit_aeqdsk:li", tree_name="_efit_tree" ) # [H] # Fetch raw ip signal to calculate dip_dt and apply smoothing - ip, ip_time = params.data_conn.get_data_with_dims( + ip, ip_time = params.get_data_with_dims( r"\pcrl01", tree_name="pcs_east" ) # [A] except mdsExceptions.MdsException: @@ -830,19 +830,19 @@ def get_n_equal_1_data(params: PhysicsMethodParams): # Get the rmp coil currents # Translated from get_rmp_and_saddle_signals.m - (rmptime,) = params.data_conn.get_dims(r"\irmpu1", tree_name="east") + (rmptime,) = params.get_dims(r"\irmpu1", tree_name="east") rmp = np.full((len(rmptime), 16), np.nan) for i in range(8): # Get irmpu1 to irmpu8 - signal = params.data_conn.get_data(rf"\irmpu{i+1}", tree_name="east") + signal = params.get_data(rf"\irmpu{i+1}", tree_name="east") if len(signal) == len(rmptime): rmp[:, i] = signal # Get irmpl1 to irmpl8 - signal = params.data_conn.get_data(rf"\irmpl{i+1}", tree_name="east") + signal = params.get_data(rf"\irmpl{i+1}", tree_name="east") if len(signal) == len(rmptime): rmp[:, i + 8] = signal # Get saddle coil signals - (saddletime,) = params.data_conn.get_dims(r"\sad_pa", tree_name="east") + (saddletime,) = params.get_dims(r"\sad_pa", tree_name="east") saddle = np.full((len(saddletime), 8), np.nan) saddle_nodes = [ r"\sad_pa", @@ -856,11 +856,11 @@ def get_n_equal_1_data(params: PhysicsMethodParams): ] for i, node in enumerate(saddle_nodes[:7]): try: - saddle[:, i] = params.data_conn.get_data(node, tree_name="east") + saddle[:, i] = params.get_data(node, tree_name="east") except mdsExceptions.MdsException: saddle[:, i] = 0 - sad_lo = params.data_conn.get_data(r"\sad_lo", tree_name="east") - sad_lm = params.data_conn.get_data(r"\sad_lm", tree_name="east") + sad_lo = params.get_data(r"\sad_lo", tree_name="east") + sad_lm = params.get_data(r"\sad_lm", tree_name="east") saddle[:, 7] = sad_lo - sad_lm # Calculate RMP n=1 Fourier component amplitude and phase (on the timebase @@ -942,7 +942,7 @@ def get_btor(params: PhysicsMethodParams): tree = "pcs_east" else: tree = "eng_tree" - itf, btor_time = params.data_conn.get_data_with_dims(r"\it", tree_name=tree) + itf, btor_time = params.get_data_with_dims(r"\it", tree_name=tree) btor = ( (4 * np.pi * 1e-7) * itf * (16 * 130) / (2 * np.pi * 1.8) ) # about 4,327 amps/tesla @@ -1071,7 +1071,7 @@ def get_pcs_parameters(params: PhysicsMethodParams): } for name, node in signals.items(): try: - signal, timearray = params.data_conn.get_data_with_dims( + signal, timearray = params.get_data_with_dims( node, tree_name="pcs_east" ) signal = interp1( @@ -1083,7 +1083,7 @@ def get_pcs_parameters(params: PhysicsMethodParams): # Get q95_rt try: - q95_rt, q95_rt_time = params.data_conn.get_data_with_dims( + q95_rt, q95_rt_time = params.get_data_with_dims( r"\q95", tree_name="pefitrt_east" ) # Deal with bug @@ -1126,7 +1126,7 @@ def get_pcs_power(params: PhysicsMethodParams): """ # Get p_rad_rt try: - p_rad_rt, timearray = params.data_conn.get_data_with_dims( + p_rad_rt, timearray = params.get_data_with_dims( r"\pcprad", tree_name="pcs_east" ) p_rad_rt = interp1( @@ -1146,7 +1146,7 @@ def get_pcs_power(params: PhysicsMethodParams): try: nbi_signals = {} for name, node in nbi_nodes.items(): - signal, timearray = params.data_conn.get_data_with_dims( + signal, timearray = params.get_data_with_dims( node, tree_name="pefitrt_east" ) signal = interp1( @@ -1170,7 +1170,7 @@ def get_pcs_power(params: PhysicsMethodParams): try: lh_signals = {} for name, node in lh_nodes.items(): - signal, timearray = params.data_conn.get_data_with_dims( + signal, timearray = params.get_data_with_dims( node, tree_name="pefitrt_east" ) signal = interp1( @@ -1314,7 +1314,7 @@ def get_mirnov_std(params: PhysicsMethodParams): # Get the Mirnov signal from \cmp1t (5 MHz) time_window = 0.001 - bp_dot, bp_dot_time = params.data_conn.get_data_with_dims( + bp_dot, bp_dot_time = params.get_data_with_dims( r"\cmp1t", tree_name="east" ) # [T/s], [s] for i, time in enumerate(params.times): @@ -1367,7 +1367,7 @@ def get_n1rms_n2rms(params: PhysicsMethodParams): n1rms_normalized = [np.nan] n2rms_normalized = [np.nan] - (mirtime,) = params.data_conn.get_dims(r"\mitab2", tree_name="east") + (mirtime,) = params.get_dims(r"\mitab2", tree_name="east") mir = np.full((len(mirtime), 16), np.nan) mir_nodes = [ r"\mitab2", @@ -1389,7 +1389,7 @@ def get_n1rms_n2rms(params: PhysicsMethodParams): ] for i, node in enumerate(mir_nodes): try: - mir[:, i] = params.data_conn.get_data(node, tree_name="east") + mir[:, i] = params.get_data(node, tree_name="east") except mdsExceptions.MdsException: continue @@ -1447,7 +1447,7 @@ def get_h98(params: PhysicsMethodParams): """ h98_y2 = [np.nan] - h98_y2, h98_y2_time = params.data_conn.get_data_with_dims( + h98_y2, h98_y2_time = params.get_data_with_dims( r"\h98_mhd", tree_name="energy_east" ) @@ -1466,7 +1466,7 @@ def _get_efit_gaps(params: PhysicsMethodParams, tree: str = "_efit_tree"): lower_gap = [np.nan] # Get plasma boundary data - data, efittime = params.data_conn.get_data_with_dims( + data, efittime = params.get_data_with_dims( r"\top.results.geqdsk:bdry", tree_name=tree ) # Convert the order of indices to MATLAB order @@ -1475,10 +1475,10 @@ def _get_efit_gaps(params: PhysicsMethodParams, tree: str = "_efit_tree"): xcoords, ycoords = data # Get first wall geometry data - xfirstwall = params.data_conn.get_data( + xfirstwall = params.get_data( r"\top.results.geqdsk:xlim", tree_name=tree ) - yfirstwall = params.data_conn.get_data( + yfirstwall = params.get_data( r"\top.results.geqdsk:ylim", tree_name=tree ) seed = np.ones((len(xcoords), 1)) diff --git a/disruption_py/machine/hbtep/physics.py b/disruption_py/machine/hbtep/physics.py index 930298e2..de2b2f50 100644 --- a/disruption_py/machine/hbtep/physics.py +++ b/disruption_py/machine/hbtep/physics.py @@ -31,7 +31,7 @@ def get_ip(params: PhysicsMethodParams): """ Get the plasma current """ - ip, t_ip = params.data_conn.get_data_with_dims( + ip, t_ip = params.get_data_with_dims( r"\top.sensors.rogowskis:ip", tree_name="hbtep2" ) # [A], [s] ip = interp1(t_ip, ip, params.times, "linear") @@ -43,10 +43,10 @@ def get_cap_bank_parameters(params: PhysicsMethodParams): """ Get the ohmic heating and vertical field coil currents """ - i_vfc, t_vfc = params.data_conn.get_data_with_dims( + i_vfc, t_vfc = params.get_data_with_dims( r"\top.sensors.vf_current", tree_name="hbtep2" ) # [A], [s] - i_ohc, t_ohc = params.data_conn.get_data_with_dims( + i_ohc, t_ohc = params.get_data_with_dims( r"\top.sensors.oh_current", tree_name="hbtep2" ) # [A], [s] i_vfc = interp1(t_vfc, i_vfc, params.times, "linear") @@ -69,21 +69,21 @@ def get_plasma_radii(params: PhysicsMethodParams): oh_pickup = 7.0723416e-08 # get vf and oh data - i_vfc, t_vfc = params.data_conn.get_data_with_dims( + i_vfc, t_vfc = params.get_data_with_dims( r"\top.sensors.vf_current", tree_name="hbtep2" ) # [A], [s] - i_ohc, t_ohc = params.data_conn.get_data_with_dims( + i_ohc, t_ohc = params.get_data_with_dims( r"\top.sensors.oh_current", tree_name="hbtep2" ) # [A], [s] # get plasma current - ip, t_ip = params.data_conn.get_data_with_dims( + ip, t_ip = params.get_data_with_dims( r"\top.sensors.rogowskis:ip", tree_name="hbtep2" ) # [A], [s] ip *= 1212.3 * 1e-9 # ip gain # get cosine Rogowski data - cos1_raw, t_cos1_raw = params.data_conn.get_data_with_dims( + cos1_raw, t_cos1_raw = params.get_data_with_dims( r"\top.sensors.rogowskis:cos_1:raw", tree_name="hbtep2" ) # [A/s], [s] # calculate and subtract offset @@ -137,7 +137,7 @@ def get_btor(params: PhysicsMethodParams): """ Calculate B_tor from the TF probe data """ - btor, t_btor = params.data_conn.get_data_with_dims( + btor, t_btor = params.get_data_with_dims( r"\top.sensors.tf_probe", tree_name="hbtep2" ) # [T], [s] btor = interp1(t_btor, btor, params.times, "linear") @@ -168,7 +168,7 @@ def get_v_loop(params: PhysicsMethodParams): """ Get v_loop measurement """ - v_loop, t_v_loop = params.data_conn.get_data_with_dims( + v_loop, t_v_loop = params.get_data_with_dims( r"\top.sensors.loop_voltage", tree_name="hbtep2" ) # [V], [s] v_loop = interp1(t_v_loop, v_loop, params.times, "linear") @@ -180,7 +180,7 @@ def get_h_alpha(params: PhysicsMethodParams): """ Get D alpha line emission from spectrometer """ - h_alpha, t_h_alpha = params.data_conn.get_data_with_dims( + h_alpha, t_h_alpha = params.get_data_with_dims( r"\top.sensors.spectrometer", tree_name="hbtep2" ) # [arb], [s] h_alpha = interp1(t_h_alpha, h_alpha, params.times, "linear") @@ -192,7 +192,7 @@ def get_sxr_midplane(params: PhysicsMethodParams): """ Get the soft x-ray midplane sensor data """ - sxr, t_sxr = params.data_conn.get_data_with_dims( + sxr, t_sxr = params.get_data_with_dims( r"\top.devices.north_rack:cpci:input_74", tree_name="hbtep2" ) # [arb], [s] # offset subtraction @@ -503,7 +503,7 @@ def _get_ta_data(params: PhysicsMethodParams): ta_pol_data_filt = [] ta_pol_time = [] for sensor_name in output["ta_pol_names"]: - sensor_data_raw, sensor_time = params.data_conn.get_data_with_dims( + sensor_data_raw, sensor_time = params.get_data_with_dims( r"\top.sensors.magnetic:" + sensor_name, tree_name="hbtep2" ) # [T], [s] fs_ta_pol = round(1 / (sensor_time[1] - sensor_time[0])) @@ -517,7 +517,7 @@ def _get_ta_data(params: PhysicsMethodParams): ta_rad_data_filt = [] ta_rad_time = [] for sensor_name in output["ta_rad_names"]: - sensor_data_raw, sensor_time = params.data_conn.get_data_with_dims( + sensor_data_raw, sensor_time = params.get_data_with_dims( r"\top.sensors.magnetic:" + sensor_name, tree_name="hbtep2" ) # [T], [s] fs_ta_rad = round(1 / (sensor_time[1] - sensor_time[0])) @@ -632,7 +632,7 @@ def _get_pa_data(params: PhysicsMethodParams): pa1_data_filt = [] pa1_time = [] for sensor_name in output["pa1_names"]: - sensor_data_raw, sensor_time = params.data_conn.get_data_with_dims( + sensor_data_raw, sensor_time = params.get_data_with_dims( r"\top.sensors.magnetic:" + sensor_name, tree_name="hbtep2" ) fs_pa1 = round(1 / (sensor_time[1] - sensor_time[0])) @@ -646,7 +646,7 @@ def _get_pa_data(params: PhysicsMethodParams): pa2_data_filt = [] pa2_time = [] for sensor_name in output["pa2_names"]: - sensor_data_raw, sensor_time = params.data_conn.get_data_with_dims( + sensor_data_raw, sensor_time = params.get_data_with_dims( r"\top.sensors.magnetic:" + sensor_name, tree_name="hbtep2" ) fs_pa2 = round(1 / (sensor_time[1] - sensor_time[0])) @@ -679,7 +679,7 @@ def _get_sxr_data(params: PhysicsMethodParams): output["sensor_num"] = [1, 2, 3, 4, 6, 9, 11, 12, 14, 16] # channels = output['sensor_num'] + 75 - time = params.data_conn.get_dims( + time = params.get_dims( r"\top.sensors.sxr_fan.channel_01:raw", tree_name="hbtep2" ) output["time"] = time[0] @@ -691,25 +691,25 @@ def _get_sxr_data(params: PhysicsMethodParams): midplane = [] for n in output["sensor_num"]: data.append( - params.data_conn.get_data( + params.get_data( rf"\top.sensors.sxr_fan.channel_{n:02d}:raw", tree_name="hbtep2", ) ) r.append( - params.data_conn.get_data( + params.get_data( rf"\top.sensors.sxr_fan.channel_{n:02d}:r", tree_name="hbtep2", ) ) z.append( - params.data_conn.get_data( + params.get_data( rf"\top.sensors.sxr_fan.channel_{n:02d}:z", tree_name="hbtep2", ) ) midplane.append( - params.data_conn.get_data( + params.get_data( rf"\top.sensors.sxr_fan.channel_{n:02d}:midplane", tree_name="hbtep2", ) @@ -736,7 +736,7 @@ def _get_euv_data(params: PhysicsMethodParams): output = {} output["detectors"] = [0, 25, 90, 270] # poloidal location of the 4 fan arrays - time = params.data_conn.get_dims( + time = params.get_dims( r"\top.sensors.euv.pol.det000.channel_01:raw", tree_name="hbtep2" ) output["time"] = time[0] @@ -748,12 +748,12 @@ def _get_euv_data(params: PhysicsMethodParams): rf"\top.sensors.euv.pol.det{detector:03d}.channel_{channel:02d}" ) data.append( - params.data_conn.get_data(address + ":raw", tree_name="hbtep2") + params.get_data(address + ":raw", tree_name="hbtep2") ) - r.append(params.data_conn.get_data(address + ":r", tree_name="hbtep2")) - z.append(params.data_conn.get_data(address + ":z", tree_name="hbtep2")) + r.append(params.get_data(address + ":r", tree_name="hbtep2")) + z.append(params.get_data(address + ":z", tree_name="hbtep2")) gain.append( - params.data_conn.get_data(address + ":gain", tree_name="hbtep2") + params.get_data(address + ":gain", tree_name="hbtep2") ) output[f"euv_{detector:03d}_data"] = data output[f"euv_{detector:03d}_r"] = r diff --git a/disruption_py/machine/mast/efit.py b/disruption_py/machine/mast/efit.py index 3bb00789..1e6eba37 100644 --- a/disruption_py/machine/mast/efit.py +++ b/disruption_py/machine/mast/efit.py @@ -50,12 +50,12 @@ def get_efit_parameters(params: PhysicsMethodParams): dict A dictionary containing the retrieved EFIT parameters. """ - eq_time = params.data_conn.get_data("equilibrium/time") + eq_time = params.get_data("equilibrium/time") times = params.times outputs = {} for key, prop in MastEfitMethods.efit_properties.items(): - signal = params.data_conn.get_data(f"equilibrium/{prop}") + signal = params.get_data(f"equilibrium/{prop}") item = MastUtilMethods.interpolate_1d(eq_time, signal, times) outputs[key] = item diff --git a/disruption_py/machine/mast/physics.py b/disruption_py/machine/mast/physics.py index d8c3216b..99daf64e 100644 --- a/disruption_py/machine/mast/physics.py +++ b/disruption_py/machine/mast/physics.py @@ -40,10 +40,10 @@ def get_ip_parameters(params: PhysicsMethodParams): A dictionary containing plasma current (`ip`), its time derivative (`dip_dt`), programmed plasma current (`ip_prog`), and its time derivative (`dipprog_dt`). """ - ip = params.data_conn.get_data("summary/ip") - ip_prog = params.data_conn.get_data("pulse_schedule/i_plasma") - ip_prog_time = params.data_conn.get_data("pulse_schedule/time") - magtime = params.data_conn.get_data("summary/time") + ip = params.get_data("summary/ip") + ip_prog = params.get_data("pulse_schedule/i_plasma") + ip_prog_time = params.get_data("pulse_schedule/time") + magtime = params.get_data("summary/time") dip_dt = np.gradient(ip, magtime) dipprog_dt = np.gradient(ip_prog, ip_prog_time) @@ -82,9 +82,9 @@ def get_power(params: PhysicsMethodParams): radiated power (`power_radiated`). """ - power_nbi = params.data_conn.get_data("summary/power_nbi") - power_radiated = params.data_conn.get_data("summary/power_radiated") - base_time = params.data_conn.get_data("summary/time") + power_nbi = params.get_data("summary/power_nbi") + power_radiated = params.get_data("summary/power_radiated") + base_time = params.get_data("summary/time") times = params.times power_nbi = MastUtilMethods.interpolate_1d(base_time, power_nbi, times) @@ -113,10 +113,10 @@ def get_gas(params: PhysicsMethodParams): inboard total gas (`inboard_total`), and outboard total gas (`outboard_total`). """ - total_injected = params.data_conn.get_data("gas_injection/total_injected") - inboard_total = params.data_conn.get_data("gas_injection/inboard_total") - outboard_total = params.data_conn.get_data("gas_injection/outboard_total") - base_time = params.data_conn.get_data("gas_injection/time") + total_injected = params.get_data("gas_injection/total_injected") + inboard_total = params.get_data("gas_injection/inboard_total") + outboard_total = params.get_data("gas_injection/outboard_total") + base_time = params.get_data("gas_injection/time") times = params.times total_injected = MastUtilMethods.interpolate_1d( @@ -153,9 +153,9 @@ def get_ts_parameters(params: PhysicsMethodParams): """ times = params.times - t_e_core = params.data_conn.get_data("thomson_scattering/t_e_core") - n_e_core = params.data_conn.get_data("thomson_scattering/n_e_core") - base_time = params.data_conn.get_data("thomson_scattering/time") + t_e_core = params.get_data("thomson_scattering/t_e_core") + n_e_core = params.get_data("thomson_scattering/n_e_core") + base_time = params.get_data("thomson_scattering/time") t_e_core = MastUtilMethods.interpolate_1d(base_time, t_e_core, times) n_e_core = MastUtilMethods.interpolate_1d(base_time, n_e_core, times) @@ -194,12 +194,12 @@ def get_densities(params: PhysicsMethodParams): """ - n_e = params.data_conn.get_data("summary/line_average_n_e") - t_n = params.data_conn.get_data("summary/time") - ip = params.data_conn.get_data("summary/ip") - t_ip = params.data_conn.get_data("summary/time") - a_minor = params.data_conn.get_data("equilibrium/minor_radius") - t_a = params.data_conn.get_data("equilibrium/time") + n_e = params.get_data("summary/line_average_n_e") + t_n = params.get_data("summary/time") + ip = params.get_data("summary/ip") + t_ip = params.get_data("summary/time") + a_minor = params.get_data("equilibrium/minor_radius") + t_a = params.get_data("equilibrium/time") return MastPhysicsMethods._get_densities( params.times, n_e, t_n, ip, t_ip, a_minor, t_a @@ -268,7 +268,7 @@ def get_sxr(params: PhysicsMethodParams): A dictionary containing SXR data (`sxr_data`) and corresponding time points (`sxr_time`). """ - hcam = params.data_conn.get_data("soft_x_rays/horizontal_cam_upper", return_xarray=True) + hcam = params.get_data("soft_x_rays/horizontal_cam_upper", return_xarray=True) if hcam is not None: hcam = hcam.isel(horizontal_cam_upper_channel=7) @@ -304,7 +304,7 @@ def get_dalpha(params: PhysicsMethodParams): A dictionary containing D-alpha signal data (`dalpha`). """ - dalpha = params.data_conn.get_data( + dalpha = params.get_data( "spectrometer_visible/filter_spectrometer_dalpha_voltage", return_xarray=True, ) diff --git a/disruption_py/settings/time_setting.py b/disruption_py/settings/time_setting.py index ffb209d6..0d6b1132 100644 --- a/disruption_py/settings/time_setting.py +++ b/disruption_py/settings/time_setting.py @@ -258,16 +258,16 @@ def _get_times(self, params: TimeSettingParams) -> np.ndarray: np.ndarray Array of times in the timebase. """ - (efit_time,) = params.data_conn.get_dims( + (efit_time,) = params.get_dims( r"\efit_aeqdsk:ali", tree_name="_efit_tree" ) - efit_time_unit = params.data_conn.get_data( + efit_time_unit = params.get_data( r"units_of(dim_of(\efit_aeqdsk:ali))", tree_name="_efit_tree" ) if efit_time_unit not in {"s", "ms", "us"}: params.logger.verbose( "Failed to get the time units of EFIT tree '{tree}', assuming seconds.", - tree=params.data_conn.get_tree_name_of_nickname("_efit_tree"), + tree=params.get_tree_name_of_nickname("_efit_tree"), ) return _postprocess(times=efit_time, units=efit_time_unit) @@ -338,7 +338,7 @@ def d3d_times(self, params: TimeSettingParams) -> np.ndarray: np.ndarray Array of times in the timebase. """ - raw_ip, ip_time = params.data_conn.get_data_with_dims( + raw_ip, ip_time = params.get_data_with_dims( f"ptdata('ip', {params.shot_id})" ) ip_time = ip_time / 1.0e3 @@ -372,7 +372,7 @@ def hbtep_times(self, params: TimeSettingParams) -> np.ndarray: This will be replaced once a get_disruption_time method is implemented for HBT-EP """ - t_ip = params.data_conn.get_dims( + t_ip = params.get_dims( r"\top.sensors.rogowskis:ip", tree_name="hbtep2" ) # [s] t_ip = t_ip[0] @@ -568,7 +568,7 @@ def cmod_times(self, params: TimeSettingParams) -> np.ndarray: np.ndarray Array of times in the timebase. """ - (ip_time,) = params.data_conn.get_dims(r"\ip", tree_name="magnetics") + (ip_time,) = params.get_dims(r"\ip", tree_name="magnetics") return ip_time def d3d_times(self, params: TimeSettingParams) -> np.ndarray: @@ -585,7 +585,7 @@ def d3d_times(self, params: TimeSettingParams) -> np.ndarray: np.ndarray Array of times in the timebase. """ - (ip_time,) = params.data_conn.get_dims(f"ptdata('ip', {params.shot_id})") + (ip_time,) = params.get_dims(f"ptdata('ip', {params.shot_id})") ip_time /= 1e3 # [ms] -> [s] return ip_time @@ -603,7 +603,7 @@ def east_times(self, params: TimeSettingParams) -> np.ndarray: np.ndarray Array of times in the timebase. """ - (ip_time,) = params.data_conn.get_dims(r"\pcrl01", tree_name="pcs_east") + (ip_time,) = params.get_dims(r"\pcrl01", tree_name="pcs_east") # For shots before year 2014, the PCRL01 timebase needs to be shifted # by 17.0 ms if params.shot_id < 44432: @@ -624,7 +624,7 @@ def hbtep_times(self, params: TimeSettingParams) -> np.ndarray: np.ndarray Array of times in the timebase. """ - (ip_time,) = params.data_conn.get_dims( + (ip_time,) = params.get_dims( r"\top.sensors.rogowskis:ip", tree_name="hbtep2" ) return ip_time @@ -695,7 +695,7 @@ def _get_times(self, params: TimeSettingParams) -> np.ndarray: Array of times in the timebase. """ try: - (signal_time,) = params.data_conn.get_dims( + (signal_time,) = params.get_dims( self.signal_path, tree_name=self.tree_name ) except mdsExceptions.MdsException: @@ -704,7 +704,7 @@ def _get_times(self, params: TimeSettingParams) -> np.ndarray: signal_path=self.signal_path, ) raise - signal_unit = params.data_conn.get_data( + signal_unit = params.get_data( f"units_of(dim_of({self.signal_path}))", tree_name=self.tree_name ) if ( diff --git a/docs/examples/custom_physics_method.py b/docs/examples/custom_physics_method.py index 70c817db..5e8fd69f 100644 --- a/docs/examples/custom_physics_method.py +++ b/docs/examples/custom_physics_method.py @@ -43,9 +43,9 @@ def decorated_physics_method(params: PhysicsMethodParams) -> dict: # pylint: disable=duplicate-code @physics_method(columns=["custom_kappa_area"], tokamak=Tokamak.CMOD) def get_custom_kappa_area(params: PhysicsMethodParams): - aminor = params.data_conn.get_data(r"\efit_aeqdsk:aminor", tree_name="_efit_tree") - area = params.data_conn.get_data(r"\efit_aeqdsk:area", tree_name="_efit_tree") - times = params.data_conn.get_data(r"\efit_aeqdsk:time", tree_name="_efit_tree") + aminor = params.get_data(r"\efit_aeqdsk:aminor", tree_name="_efit_tree") + area = params.get_data(r"\efit_aeqdsk:area", tree_name="_efit_tree") + times = params.get_data(r"\efit_aeqdsk:time", tree_name="_efit_tree") # Ensure aminor and area are not 0 or less than 0 aminor[aminor <= 0] = 0.001 diff --git a/docs/examples/custom_time_setting.py b/docs/examples/custom_time_setting.py index 3131d4ed..41f7ed72 100644 --- a/docs/examples/custom_time_setting.py +++ b/docs/examples/custom_time_setting.py @@ -11,7 +11,7 @@ class PRadTime(TimeSetting): def _get_times(self, params: TimeSettingParams): """Return prad times""" - (time_array,) = params.data_conn.get_dims( + (time_array,) = params.get_dims( r"\twopi_diode", tree_name="spectroscopy" ) time_array = time_array[time_array > 0] From 586ffbb148e4d0d83904354f95c08cc54561195f Mon Sep 17 00:00:00 2001 From: gtrevisan Date: Tue, 28 Apr 2026 17:04:10 -0400 Subject: [PATCH 4/9] relint all files --- disruption_py/machine/cmod/efit.py | 8 +--- disruption_py/machine/cmod/physics.py | 56 ++++++---------------- disruption_py/machine/cmod/thomson.py | 4 +- disruption_py/machine/d3d/efit.py | 6 +-- disruption_py/machine/d3d/physics.py | 64 +++++++------------------- disruption_py/machine/d3d/util.py | 4 +- disruption_py/machine/east/physics.py | 20 ++------ disruption_py/machine/hbtep/physics.py | 8 +--- disruption_py/settings/time_setting.py | 20 ++------ docs/examples/custom_time_setting.py | 4 +- 10 files changed, 49 insertions(+), 145 deletions(-) diff --git a/disruption_py/machine/cmod/efit.py b/disruption_py/machine/cmod/efit.py index 349d2b61..47c2da26 100644 --- a/disruption_py/machine/cmod/efit.py +++ b/disruption_py/machine/cmod/efit.py @@ -69,9 +69,7 @@ def get_efit_parameters(params: PhysicsMethodParams): dict A dictionary containing the retrieved EFIT parameters. """ - efit_time = params.get_data( - r"\efit_aeqdsk:time", tree_name="_efit_tree" - ) # [s] + efit_time = params.get_data(r"\efit_aeqdsk:time", tree_name="_efit_tree") # [s] efit_data = {} # Get data from each of the columns in efit_cols one at a time @@ -125,7 +123,5 @@ def efit_check(params: PhysicsMethodParams): ] _n = values[2].data() valid_indices = np.nonzero(_n) - (times,) = params.get_dims( - r"\efit_aeqdsk:lflag", tree_name="analysis" - ) + (times,) = params.get_dims(r"\efit_aeqdsk:lflag", tree_name="analysis") return valid_indices, times[valid_indices] diff --git a/disruption_py/machine/cmod/physics.py b/disruption_py/machine/cmod/physics.py index 1f2b87a6..05e5863f 100644 --- a/disruption_py/machine/cmod/physics.py +++ b/disruption_py/machine/cmod/physics.py @@ -70,9 +70,7 @@ def _get_active_wire_segments(params: PhysicsMethodParams): active_segments.append( ( node_path, - params.get_data( - f"{node_path}:start_time", tree_name="pcs" - ), + params.get_data(f"{node_path}:start_time", tree_name="pcs"), ) ) @@ -562,9 +560,7 @@ def get_ohmic_parameters(params: PhysicsMethodParams): r"\efit_aeqdsk:ali", tree_name="_efit_tree" ) # [dimensionless], [s] ip_parameters = CmodPhysicsMethods.get_ip_parameters(params=params) - r0 = params.get_data( - r"\efit_aeqdsk:rmagx/100", tree_name="_efit_tree" - ) # [m] + r0 = params.get_data(r"\efit_aeqdsk:rmagx/100", tree_name="_efit_tree") # [m] return CmodPhysicsMethods._get_ohmic_parameters( params.times, @@ -739,9 +735,7 @@ def get_power(params: PhysicsMethodParams): p = f"p_{val}" t = f"t_{val}" try: - kwa[p], kwa[t] = params.get_data_with_dims( - node, tree_name=tree - ) + kwa[p], kwa[t] = params.get_data_with_dims(node, tree_name=tree) except (mdsExceptions.TreeFOPENR, mdsExceptions.TreeNNF): kwa[p], kwa[t] = None, None # Ohmic power @@ -813,9 +807,7 @@ def get_kappa_area(params: PhysicsMethodParams): area = params.get_data( r"\efit_aeqdsk:areao/1e4", tree_name="_efit_tree" ) # [m^2] - times = params.get_data( - r"\efit_aeqdsk:time", tree_name="_efit_tree" - ) # [s] + times = params.get_data(r"\efit_aeqdsk:time", tree_name="_efit_tree") # [s] aminor[aminor <= 0] = 0.001 # make sure aminor is not 0 or less than 0 # make sure area is not 0 or less than 0 @@ -891,9 +883,7 @@ def get_n_equal_1_amplitude(params: PhysicsMethodParams): a = [] path = r"\mag_bp_coils." - bp_node_names = params.get_data( - f"{path}nodename", tree_name="magnetics" - ) + bp_node_names = params.get_data(f"{path}nodename", tree_name="magnetics") phi = params.get_data(f"{path}phi", tree_name="magnetics") # [degree] btor_pickup_coeffs = params.get_data( f"{path}btor_pickup", tree_name="magnetics" @@ -918,9 +908,7 @@ def get_n_equal_1_amplitude(params: PhysicsMethodParams): for i, bp13_name in enumerate(bp13_names): try: - signal = params.get_data( - path + bp13_name, tree_name="magnetics" - ) # [T] + signal = params.get_data(path + bp13_name, tree_name="magnetics") # [T] # Sensor not available, skip except mdsExceptions.MdsException: continue @@ -1076,9 +1064,7 @@ def get_densities(params: PhysicsMethodParams): # For future refernce, chord length is stored in # .01*\analysis::efit_aeqdsk:rco2v[3,*] n_e = np.squeeze(n_e) / 0.6 - ip, t_ip = params.get_data_with_dims( - r"\ip", tree_name="magnetics" - ) # [A], [s] + ip, t_ip = params.get_data_with_dims(r"\ip", tree_name="magnetics") # [A], [s] a_minor, t_a = params.get_data_with_dims( r"\efit_aeqdsk:aout/100", tree_name="_efit_tree" ) # [m], [s] @@ -1246,9 +1232,7 @@ def get_ts_parameters(params: PhysicsMethodParams): ts_data, ts_time = params.get_data_with_dims( f"{node_path}:te_rz", tree_name="electrons" ) # [keV], [s] - ts_z = params.get_data( - f"{node_path}:z_sorted", tree_name="electrons" - ) # [m] + ts_z = params.get_data(f"{node_path}:z_sorted", tree_name="electrons") # [m] ts_error = params.get_data( f"{node_path}:te_err", tree_name="electrons" ) # [keV] @@ -1404,9 +1388,7 @@ def get_peaking_factors(params: PhysicsMethodParams): raise CalculationError("Shot is on blacklist") # Fetch data # Get EFIT geometry data - z0 = params.get_data( - r"\efit_aeqdsk:zmagx/100", tree_name="_efit_tree" - ) # [m] + z0 = params.get_data(r"\efit_aeqdsk:zmagx/100", tree_name="_efit_tree") # [m] kappa = params.get_data( r"\efit_aeqdsk:kappa", tree_name="_efit_tree" ) # [dimensionless] @@ -1793,9 +1775,7 @@ def get_te_profile_params_ece(params: PhysicsMethodParams): """ # Get magnetic axis data from EFIT - r0 = params.get_data( - r"\efit_aeqdsk:rmagx/100", tree_name="_efit_tree" - ) # [m] + r0 = params.get_data(r"\efit_aeqdsk:rmagx/100", tree_name="_efit_tree") # [m] aminor, efit_time = params.get_data_with_dims( r"\efit_aeqdsk:aout/100", tree_name="_efit_tree" ) # [m], [s] @@ -1878,12 +1858,8 @@ def get_prad_peaking(params: PhysicsMethodParams): """ prad_peaking = np.full(len(params.times), np.nan) nan_output = {"prad_peaking": prad_peaking} - r0 = params.get_data( - r"\efit_aeqdsk:rmagx/100", tree_name="_efit_tree" - ) # [m] - z0 = params.get_data( - r"\efit_aeqdsk:zmagx/100", tree_name="_efit_tree" - ) # [m] + r0 = params.get_data(r"\efit_aeqdsk:rmagx/100", tree_name="_efit_tree") # [m] + z0 = params.get_data(r"\efit_aeqdsk:zmagx/100", tree_name="_efit_tree") # [m] aminor, efit_time = params.get_data_with_dims( r"\efit_aeqdsk:aout/100", tree_name="_efit_tree" ) # [m], [s] @@ -2049,15 +2025,11 @@ def get_beta_normalized(params: PhysicsMethodParams): beta_t, efittime = params.get_data_with_dims( r"\efit_aeqdsk:betat", tree_name="_efit_tree" ) # [%], [s] - ip = params.get_data( - r"\efit_aeqdsk:cpasma/1e6", tree_name="_efit_tree" - ) # [MA] + ip = params.get_data(r"\efit_aeqdsk:cpasma/1e6", tree_name="_efit_tree") # [MA] aminor = params.get_data( r"\efit_aeqdsk:aout/100", tree_name="_efit_tree" ) # [m] - btor = params.get_data( - r"\efit_aeqdsk:btaxp", tree_name="_efit_tree" - ) # [T] + btor = params.get_data(r"\efit_aeqdsk:btaxp", tree_name="_efit_tree") # [T] # Calculate beta_n with np.errstate(divide="ignore", invalid="ignore"): diff --git a/disruption_py/machine/cmod/thomson.py b/disruption_py/machine/cmod/thomson.py index 76deac3c..5f300d09 100644 --- a/disruption_py/machine/cmod/thomson.py +++ b/disruption_py/machine/cmod/thomson.py @@ -189,9 +189,7 @@ def _map_ts2tci(params: PhysicsMethodParams, nlnum): ip = params.get_data(r"\ip", "cmod") if np.mean(ip) > 0: flag = 0 - efit_times = params.get_data( - r"\efit_aeqdsk:time", tree_name="_efit_tree" - ) + efit_times = params.get_data(r"\efit_aeqdsk:time", tree_name="_efit_tree") t1 = np.amin(efit_times) t2 = np.amax(efit_times) psia, psia_t = params.get_data_with_dims( diff --git a/disruption_py/machine/d3d/efit.py b/disruption_py/machine/d3d/efit.py index 9b256be0..585bfc3b 100644 --- a/disruption_py/machine/d3d/efit.py +++ b/disruption_py/machine/d3d/efit.py @@ -65,8 +65,7 @@ def get_efit_parameters(params: PhysicsMethodParams): for k, v in D3DEfitMethods.efit_cols.items() } efit_time = ( - params.get_data(r"\efit_a_eqdsk:atime", tree_name="_efit_tree") - / 1.0e3 + params.get_data(r"\efit_a_eqdsk:atime", tree_name="_efit_tree") / 1.0e3 ) # [ms] -> [s] # EFIT reconstructions are sometimes invalid, particularly when very close @@ -112,8 +111,7 @@ def get_rt_efit_parameters(params: PhysicsMethodParams): for k, v in D3DEfitMethods.rt_efit_cols.items() } efit_time = ( - params.get_data(r"\efit_a_eqdsk:atime", tree_name="efitrt1") - / 1.0e3 + params.get_data(r"\efit_a_eqdsk:atime", tree_name="efitrt1") / 1.0e3 ) # [ms] -> [s] # EFIT reconstructions are sometimes invalid, particularly when very close # to a disruption. There are a number of EFIT parameters that can indicate diff --git a/disruption_py/machine/d3d/physics.py b/disruption_py/machine/d3d/physics.py index 9ea7406b..0328badb 100644 --- a/disruption_py/machine/d3d/physics.py +++ b/disruption_py/machine/d3d/physics.py @@ -119,9 +119,7 @@ def get_h_alpha(params: PhysicsMethodParams): "h_alpha": [np.nan], } try: - h_alpha, t_h_alpha = params.get_data_with_dims( - r"\fs04", tree_name="d3d" - ) + h_alpha, t_h_alpha = params.get_data_with_dims(r"\fs04", tree_name="d3d") t_h_alpha /= 1e3 # [ms] -> [s] h_alpha = interp1(t_h_alpha, h_alpha, params.times, "linear") output["h_alpha"] = h_alpha @@ -166,9 +164,7 @@ def get_power_parameters(params: PhysicsMethodParams): """ # Get neutral beam injected power try: - p_nbi, t_nbi = params.get_data_with_dims( - r"\top.nb:pinj", tree_name="d3d" - ) + p_nbi, t_nbi = params.get_data_with_dims(r"\top.nb:pinj", tree_name="d3d") t_nbi /= 1e3 # [ms] -> [s] p_nbi = p_nbi * 1e3 # [KW] -> [W] if len(t_nbi) > 2: @@ -234,9 +230,7 @@ def get_power_parameters(params: PhysicsMethodParams): smoothing_window = 0.010 # [s] try: - bol_prm, _ = params.get_data_with_dims( - r"\bol_prm", tree_name="bolom" - ) + bol_prm, _ = params.get_data_with_dims(r"\bol_prm", tree_name="bolom") except mdsExceptions.MdsException as e: params.logger.warning("Failed to open bolom tree.") params.logger.opt(exception=True).debug(e) @@ -249,9 +243,7 @@ def get_power_parameters(params: PhysicsMethodParams): rf"\top.raw:{bol_channels[i]}", tree_name="bolom" ) bol_signals.append(bol_signal) - bol_time = params.get_dims( - rf"\top.raw:{bol_channels[0]}", tree_name="bolom" - )[0] + bol_time = params.get_dims(rf"\top.raw:{bol_channels[0]}", tree_name="bolom")[0] bol_time /= 1e3 # [ms] -> [s] a_struct = matlab_get_bolo( shot_id=params.shot_id, @@ -327,9 +319,7 @@ def get_ohmic_parameters(params: PhysicsMethodParams): v_loop = scipy.signal.medfilt(v_loop, 11) v_loop = interp1(t_v_loop, v_loop, params.times, "linear") # Get plasma current - ip, t_ip = params.get_data_with_dims( - f"ptdata('ip', {params.shot_id})" - ) + ip, t_ip = params.get_data_with_dims(f"ptdata('ip', {params.shot_id})") t_ip /= 1e3 # [ms] -> [s] # Alessandro Pau (JET & AUG) has given Cristina a robust routine that @@ -354,9 +344,7 @@ def get_ohmic_parameters(params: PhysicsMethodParams): ) t_li /= 1e3 # Use chisq to determine which time slices are invalid - chisq = params.get_data( - r"\efit_a_eqdsk:chisq", tree_name="_efit_tree" - ) + chisq = params.get_data(r"\efit_a_eqdsk:chisq", tree_name="_efit_tree") # Filter out invalid indices of efit reconstruction (invalid_indices,) = np.where(chisq > 50) li[invalid_indices] = np.nan @@ -410,9 +398,7 @@ def get_density_parameters(params: PhysicsMethodParams): - pull requests: #[249](https://github.com/MIT-PSFC/disruption-py/pull/249) """ try: - ne, t_ne = params.get_data_with_dims( - r"\density", tree_name="_efit_tree" - ) + ne, t_ne = params.get_data_with_dims(r"\density", tree_name="_efit_tree") except mdsExceptions.MdsException: ne = [np.nan] t_ne = [np.nan] @@ -877,9 +863,7 @@ def get_z_parameters(params: PhysicsMethodParams): """ nominal_flattop_radius = 0.59 # Get z_cur - z_cur, t_z_cur = params.get_data_with_dims( - f"ptdata('vpszp', {params.shot_id})" - ) + z_cur, t_z_cur = params.get_data_with_dims(f"ptdata('vpszp', {params.shot_id})") t_z_cur = t_z_cur / 1.0e3 # [ms] -> [s] z_cur = z_cur / 1.0e2 # [cm] -> [m] z_cur = interp1(t_z_cur, z_cur, params.times, "linear") @@ -889,9 +873,7 @@ def get_z_parameters(params: PhysicsMethodParams): r"\efit_a_eqdsk:aminor", tree_name="_efit_tree" ) # [m], [ms] t_a = t_a / 1.0e3 # [ms] -> [s] - chisq = params.get_data( - r"\efit_a_eqdsk:chisq", tree_name="_efit_tree" - ) + chisq = params.get_data(r"\efit_a_eqdsk:chisq", tree_name="_efit_tree") (invalid_indices,) = np.where(chisq > 50) a_minor[invalid_indices] = np.nan a_minor = interp1(t_a, a_minor, params.times, "linear") @@ -1327,13 +1309,9 @@ def get_kappa_area(params: PhysicsMethodParams): /blob/matlab/DIII-D/get_kappa_area.m) - pull requests: #[256](https://github.com/MIT-PSFC/disruption-py/pull/256) """ - a_minor = params.get_data( - r"\efit_a_eqdsk:aminor", tree_name="_efit_tree" - ) + a_minor = params.get_data(r"\efit_a_eqdsk:aminor", tree_name="_efit_tree") area = params.get_data(r"\efit_a_eqdsk:area", tree_name="_efit_tree") - chisq = params.get_data( - r"\efit_a_eqdsk:chisq", tree_name="_efit_tree" - ) + chisq = params.get_data(r"\efit_a_eqdsk:chisq", tree_name="_efit_tree") t = params.get_data(r"\efit_a_eqdsk:atime", tree_name="_efit_tree") t /= 1e3 # [ms] -> [s] kappa_area = area / (np.pi * a_minor**2) @@ -1369,9 +1347,7 @@ def get_shape_parameters(params: PhysicsMethodParams): - pull requests: #[258](https://github.com/MIT-PSFC/disruption-py/pull/258) """ # Get efit_time - efit_time = params.get_data( - r"\efit_a_eqdsk:atime", tree_name="_efit_tree" - ) + efit_time = params.get_data(r"\efit_a_eqdsk:atime", tree_name="_efit_tree") efit_time /= 1e3 # [ms] -> [s] # Compute triangularity try: @@ -1388,12 +1364,8 @@ def get_shape_parameters(params: PhysicsMethodParams): delta = None # Compute squareness try: - sqfod = params.get_data( - r"\efit_a_eqdsk:sqfod", tree_name="_efit_tree" - ) - sqfou = params.get_data( - r"\efit_a_eqdsk:sqfou", tree_name="_efit_tree" - ) + sqfod = params.get_data(r"\efit_a_eqdsk:sqfod", tree_name="_efit_tree") + sqfou = params.get_data(r"\efit_a_eqdsk:sqfou", tree_name="_efit_tree") squareness = (sqfod + sqfou) / 2.0 except mdsExceptions.MdsException as e: params.logger.warning("Failed to obtain squareness signals") @@ -1401,18 +1373,14 @@ def get_shape_parameters(params: PhysicsMethodParams): squareness = None # Get aminor try: - aminor = params.get_data( - r"\efit_a_eqdsk:aminor", tree_name="_efit_tree" - ) + aminor = params.get_data(r"\efit_a_eqdsk:aminor", tree_name="_efit_tree") except mdsExceptions.MdsException as e: params.logger.warning("Failed to obtain aminor signals") params.logger.opt(exception=True).debug(e) aminor = None # Check chisq for invalid indices try: - chisq = params.get_data( - r"\efit_a_eqdsk:chisq", tree_name="_efit_tree" - ) + chisq = params.get_data(r"\efit_a_eqdsk:chisq", tree_name="_efit_tree") invalid_indices = np.where(chisq > 50) except mdsExceptions.MdsException as e: params.logger.warning( diff --git a/disruption_py/machine/d3d/util.py b/disruption_py/machine/d3d/util.py index 266b3fd9..59e67ea0 100644 --- a/disruption_py/machine/d3d/util.py +++ b/disruption_py/machine/d3d/util.py @@ -32,9 +32,7 @@ def get_polarity(params: PhysicsMethodParams): ------- polarity value, -1 or 1. """ - polarity = np.unique( - params.get_data(f"ptdata('iptdirect', {params.shot_id})") - ) + polarity = np.unique(params.get_data(f"ptdata('iptdirect', {params.shot_id})")) if len(polarity) > 1: params.logger.info( "Polarity of Ip target is not constant. Using value at first timestep.", diff --git a/disruption_py/machine/east/physics.py b/disruption_py/machine/east/physics.py index b58ec7bf..b60dd88f 100644 --- a/disruption_py/machine/east/physics.py +++ b/disruption_py/machine/east/physics.py @@ -242,9 +242,7 @@ def get_v_loop(params: PhysicsMethodParams): # Get "\vp1_s" signal from the EAST tree. (This signal is a sub-sampled # version of "vp1".) try: - v_loop, v_loop_time = params.get_data_with_dims( - r"\vp1_s", tree_name="east" - ) + v_loop, v_loop_time = params.get_data_with_dims(r"\vp1_s", tree_name="east") except mdsExceptions.MdsException: params.logger.verbose( r"v_loop: Failed to get \vp1_s data. Use \pcvloop from pcs_east instead." @@ -436,9 +434,7 @@ def _get_raw_axuv_data(params: PhysicsMethodParams): for iarray in range(4): for ichan in range(16): ichord = 16 * iarray + ichan - signal = params.get_data( - r"\pxuv" + str(ichord + 1), tree_name="east_1" - ) + signal = params.get_data(r"\pxuv" + str(ichord + 1), tree_name="east_1") # Subtract baseline signal = signal - np.mean(signal[:100]) # TODO: change this to causal smoothing @@ -586,9 +582,7 @@ def get_heating_power(nodes, tree): """ heating_power = np.zeros(params.times.shape) for node in nodes: - power_node, time_node = params.get_data_with_dims( - node, tree_name=tree - ) + power_node, time_node = params.get_data_with_dims(node, tree_name=tree) heating_power += interp1( time_node, power_node, @@ -1475,12 +1469,8 @@ def _get_efit_gaps(params: PhysicsMethodParams, tree: str = "_efit_tree"): xcoords, ycoords = data # Get first wall geometry data - xfirstwall = params.get_data( - r"\top.results.geqdsk:xlim", tree_name=tree - ) - yfirstwall = params.get_data( - r"\top.results.geqdsk:ylim", tree_name=tree - ) + xfirstwall = params.get_data(r"\top.results.geqdsk:xlim", tree_name=tree) + yfirstwall = params.get_data(r"\top.results.geqdsk:ylim", tree_name=tree) seed = np.ones((len(xcoords), 1)) xfirstwall = np.reshape(xfirstwall, (-1, 1)) yfirstwall = np.reshape(yfirstwall, (-1, 1)) diff --git a/disruption_py/machine/hbtep/physics.py b/disruption_py/machine/hbtep/physics.py index de2b2f50..0f4fce3f 100644 --- a/disruption_py/machine/hbtep/physics.py +++ b/disruption_py/machine/hbtep/physics.py @@ -747,14 +747,10 @@ def _get_euv_data(params: PhysicsMethodParams): address = ( rf"\top.sensors.euv.pol.det{detector:03d}.channel_{channel:02d}" ) - data.append( - params.get_data(address + ":raw", tree_name="hbtep2") - ) + data.append(params.get_data(address + ":raw", tree_name="hbtep2")) r.append(params.get_data(address + ":r", tree_name="hbtep2")) z.append(params.get_data(address + ":z", tree_name="hbtep2")) - gain.append( - params.get_data(address + ":gain", tree_name="hbtep2") - ) + gain.append(params.get_data(address + ":gain", tree_name="hbtep2")) output[f"euv_{detector:03d}_data"] = data output[f"euv_{detector:03d}_r"] = r output[f"euv_{detector:03d}_z"] = z diff --git a/disruption_py/settings/time_setting.py b/disruption_py/settings/time_setting.py index 0d6b1132..a1cc4165 100644 --- a/disruption_py/settings/time_setting.py +++ b/disruption_py/settings/time_setting.py @@ -258,9 +258,7 @@ def _get_times(self, params: TimeSettingParams) -> np.ndarray: np.ndarray Array of times in the timebase. """ - (efit_time,) = params.get_dims( - r"\efit_aeqdsk:ali", tree_name="_efit_tree" - ) + (efit_time,) = params.get_dims(r"\efit_aeqdsk:ali", tree_name="_efit_tree") efit_time_unit = params.get_data( r"units_of(dim_of(\efit_aeqdsk:ali))", tree_name="_efit_tree" ) @@ -338,9 +336,7 @@ def d3d_times(self, params: TimeSettingParams) -> np.ndarray: np.ndarray Array of times in the timebase. """ - raw_ip, ip_time = params.get_data_with_dims( - f"ptdata('ip', {params.shot_id})" - ) + raw_ip, ip_time = params.get_data_with_dims(f"ptdata('ip', {params.shot_id})") ip_time = ip_time / 1.0e3 baseline = np.mean(raw_ip[:10]) ip = raw_ip - baseline @@ -372,9 +368,7 @@ def hbtep_times(self, params: TimeSettingParams) -> np.ndarray: This will be replaced once a get_disruption_time method is implemented for HBT-EP """ - t_ip = params.get_dims( - r"\top.sensors.rogowskis:ip", tree_name="hbtep2" - ) # [s] + t_ip = params.get_dims(r"\top.sensors.rogowskis:ip", tree_name="hbtep2") # [s] t_ip = t_ip[0] t_ip = t_ip[(t_ip >= 0) & (t_ip <= 12e-3)] steps = round(10e-6 / (t_ip[1] - t_ip[0])) @@ -624,9 +618,7 @@ def hbtep_times(self, params: TimeSettingParams) -> np.ndarray: np.ndarray Array of times in the timebase. """ - (ip_time,) = params.get_dims( - r"\top.sensors.rogowskis:ip", tree_name="hbtep2" - ) + (ip_time,) = params.get_dims(r"\top.sensors.rogowskis:ip", tree_name="hbtep2") return ip_time def mast_times(self, params: TimeSettingParams) -> np.ndarray: @@ -695,9 +687,7 @@ def _get_times(self, params: TimeSettingParams) -> np.ndarray: Array of times in the timebase. """ try: - (signal_time,) = params.get_dims( - self.signal_path, tree_name=self.tree_name - ) + (signal_time,) = params.get_dims(self.signal_path, tree_name=self.tree_name) except mdsExceptions.MdsException: params.logger.error( "Failed to set up timebase for signal {signal_path}", diff --git a/docs/examples/custom_time_setting.py b/docs/examples/custom_time_setting.py index 41f7ed72..98d67bf1 100644 --- a/docs/examples/custom_time_setting.py +++ b/docs/examples/custom_time_setting.py @@ -11,9 +11,7 @@ class PRadTime(TimeSetting): def _get_times(self, params: TimeSettingParams): """Return prad times""" - (time_array,) = params.get_dims( - r"\twopi_diode", tree_name="spectroscopy" - ) + (time_array,) = params.get_dims(r"\twopi_diode", tree_name="spectroscopy") time_array = time_array[time_array > 0] return time_array From 83ed9d4f1b98b1a9c6884b1cca42329b11105a23 Mon Sep 17 00:00:00 2001 From: gtrevisan Date: Wed, 29 Apr 2026 08:42:46 -0400 Subject: [PATCH 5/9] offer data getters from TimeSettingParams, too --- disruption_py/settings/time_setting.py | 39 ++++++++++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/disruption_py/settings/time_setting.py b/disruption_py/settings/time_setting.py index a1cc4165..87ea3f8e 100644 --- a/disruption_py/settings/time_setting.py +++ b/disruption_py/settings/time_setting.py @@ -64,6 +64,45 @@ def disrupted(self) -> bool: """ return self.disruption_time is not None + def get_data(self, *args, **kwargs) -> np.ndarray: + """ + Get data using the data connection. + + Parameters + ---------- + *args + Positional arguments for `data_conn.get_data`. + **kwargs + Keyword arguments for `data_conn.get_data`. + """ + return self.data_conn.get_data(*args, **kwargs) + + def get_data_with_dims(self, *args, **kwargs) -> np.ndarray: + """ + Get data with dimensions using the data connection. + + Parameters + ---------- + *args + Positional arguments for `data_conn.get_data_with_dims`. + **kwargs + Keyword arguments for `data_conn.get_data_with_dims`. + """ + return self.data_conn.get_data_with_dims(*args, **kwargs) + + def get_dims(self, *args, **kwargs) -> np.ndarray: + """ + Get dimensions using the data connection. + + Parameters + ---------- + *args + Positional arguments for `data_conn.get_dims`. + **kwargs + Keyword arguments for `data_conn.get_dims`. + """ + return self.data_conn.get_dims(*args, **kwargs) + TimeSettingType = Union[ "TimeSetting", str, Dict[Tokamak, "TimeSettingType"], List["TimeSettingType"] From 860452c9a83951535e78814d9dd37dcba46f2e88 Mon Sep 17 00:00:00 2001 From: gtrevisan Date: Wed, 29 Apr 2026 08:43:30 -0400 Subject: [PATCH 6/9] revert changes to get_tree_name_of_nickname calls --- disruption_py/machine/d3d/physics.py | 2 +- disruption_py/settings/time_setting.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/disruption_py/machine/d3d/physics.py b/disruption_py/machine/d3d/physics.py index 0328badb..049df15e 100644 --- a/disruption_py/machine/d3d/physics.py +++ b/disruption_py/machine/d3d/physics.py @@ -412,7 +412,7 @@ def get_density_parameters(params: PhysicsMethodParams): # - "\d3d:denv2" gives actual density data if not np.isfinite(ne).any() or len(ne) < 2: ne, t_ne = params.get_data_with_dims(r"\denv2", tree_name="d3d") - tree_name = params.get_tree_name_of_nickname("_efit_tree") + tree_name = params.data_conn.get_tree_name_of_nickname("_efit_tree") params.logger.verbose( rf"density: data from \{tree_name}:density is either empty or invalid." r" Use \d3d:denv2 instead." diff --git a/disruption_py/settings/time_setting.py b/disruption_py/settings/time_setting.py index 87ea3f8e..4dd98fa9 100644 --- a/disruption_py/settings/time_setting.py +++ b/disruption_py/settings/time_setting.py @@ -304,7 +304,7 @@ def _get_times(self, params: TimeSettingParams) -> np.ndarray: if efit_time_unit not in {"s", "ms", "us"}: params.logger.verbose( "Failed to get the time units of EFIT tree '{tree}', assuming seconds.", - tree=params.get_tree_name_of_nickname("_efit_tree"), + tree=params.data_conn.get_tree_name_of_nickname("_efit_tree"), ) return _postprocess(times=efit_time, units=efit_time_unit) From abe28248a658b6cbfeba3026a5e10ef7aa25f649 Mon Sep 17 00:00:00 2001 From: gtrevisan Date: Wed, 29 Apr 2026 08:57:15 -0400 Subject: [PATCH 7/9] use base class type hints for params getters --- disruption_py/core/physics_method/params.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/disruption_py/core/physics_method/params.py b/disruption_py/core/physics_method/params.py index d102d5e6..596706cb 100644 --- a/disruption_py/core/physics_method/params.py +++ b/disruption_py/core/physics_method/params.py @@ -83,7 +83,7 @@ def get_data(self, *args, **kwargs) -> np.ndarray: """ return self.data_conn.get_data(*args, **kwargs) - def get_data_with_dims(self, *args, **kwargs) -> np.ndarray: + def get_data_with_dims(self, *args, **kwargs) -> Tuple: """ Get data with dimensions using the data connection. @@ -96,7 +96,7 @@ def get_data_with_dims(self, *args, **kwargs) -> np.ndarray: """ return self.data_conn.get_data_with_dims(*args, **kwargs) - def get_dims(self, *args, **kwargs) -> np.ndarray: + def get_dims(self, *args, **kwargs) -> Tuple: """ Get dimensions using the data connection. From d0de1e5679ff783e829db6bacdade09f6d34fab9 Mon Sep 17 00:00:00 2001 From: gtrevisan Date: Wed, 29 Apr 2026 09:11:50 -0400 Subject: [PATCH 8/9] fixup! use base class type hints for params getters --- disruption_py/settings/time_setting.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/disruption_py/settings/time_setting.py b/disruption_py/settings/time_setting.py index 4dd98fa9..379c5086 100644 --- a/disruption_py/settings/time_setting.py +++ b/disruption_py/settings/time_setting.py @@ -77,7 +77,7 @@ def get_data(self, *args, **kwargs) -> np.ndarray: """ return self.data_conn.get_data(*args, **kwargs) - def get_data_with_dims(self, *args, **kwargs) -> np.ndarray: + def get_data_with_dims(self, *args, **kwargs) -> Tuple: """ Get data with dimensions using the data connection. @@ -90,7 +90,7 @@ def get_data_with_dims(self, *args, **kwargs) -> np.ndarray: """ return self.data_conn.get_data_with_dims(*args, **kwargs) - def get_dims(self, *args, **kwargs) -> np.ndarray: + def get_dims(self, *args, **kwargs) -> Tuple: """ Get dimensions using the data connection. From 2c03197ade65857d566068455b4366fa8f650fea Mon Sep 17 00:00:00 2001 From: gtrevisan Date: Wed, 29 Apr 2026 12:48:16 -0400 Subject: [PATCH 9/9] shorten docstrings for params getters --- disruption_py/core/physics_method/params.py | 33 ++------------------- disruption_py/settings/time_setting.py | 33 ++------------------- 2 files changed, 6 insertions(+), 60 deletions(-) diff --git a/disruption_py/core/physics_method/params.py b/disruption_py/core/physics_method/params.py index 596706cb..0d88b5ac 100644 --- a/disruption_py/core/physics_method/params.py +++ b/disruption_py/core/physics_method/params.py @@ -71,40 +71,13 @@ def to_coords(self) -> Dict[str, Tuple[str, np.ndarray]]: ) def get_data(self, *args, **kwargs) -> np.ndarray: - """ - Get data using the data connection. - - Parameters - ---------- - *args - Positional arguments for `data_conn.get_data`. - **kwargs - Keyword arguments for `data_conn.get_data`. - """ + """Get data using the data connection.""" return self.data_conn.get_data(*args, **kwargs) def get_data_with_dims(self, *args, **kwargs) -> Tuple: - """ - Get data with dimensions using the data connection. - - Parameters - ---------- - *args - Positional arguments for `data_conn.get_data_with_dims`. - **kwargs - Keyword arguments for `data_conn.get_data_with_dims`. - """ + """Get data with dimensions using the data connection.""" return self.data_conn.get_data_with_dims(*args, **kwargs) def get_dims(self, *args, **kwargs) -> Tuple: - """ - Get dimensions using the data connection. - - Parameters - ---------- - *args - Positional arguments for `data_conn.get_dims`. - **kwargs - Keyword arguments for `data_conn.get_dims`. - """ + """Get dimensions using the data connection.""" return self.data_conn.get_dims(*args, **kwargs) diff --git a/disruption_py/settings/time_setting.py b/disruption_py/settings/time_setting.py index 379c5086..ee00898b 100644 --- a/disruption_py/settings/time_setting.py +++ b/disruption_py/settings/time_setting.py @@ -65,42 +65,15 @@ def disrupted(self) -> bool: return self.disruption_time is not None def get_data(self, *args, **kwargs) -> np.ndarray: - """ - Get data using the data connection. - - Parameters - ---------- - *args - Positional arguments for `data_conn.get_data`. - **kwargs - Keyword arguments for `data_conn.get_data`. - """ + """Get data using the data connection.""" return self.data_conn.get_data(*args, **kwargs) def get_data_with_dims(self, *args, **kwargs) -> Tuple: - """ - Get data with dimensions using the data connection. - - Parameters - ---------- - *args - Positional arguments for `data_conn.get_data_with_dims`. - **kwargs - Keyword arguments for `data_conn.get_data_with_dims`. - """ + """Get data with dimensions using the data connection.""" return self.data_conn.get_data_with_dims(*args, **kwargs) def get_dims(self, *args, **kwargs) -> Tuple: - """ - Get dimensions using the data connection. - - Parameters - ---------- - *args - Positional arguments for `data_conn.get_dims`. - **kwargs - Keyword arguments for `data_conn.get_dims`. - """ + """Get dimensions using the data connection.""" return self.data_conn.get_dims(*args, **kwargs)