def transferMetadata(cls, ref_array: xr.DataArray,
new_array: xr.DataArray):
new_attrs = {
key: value
for key, value in ref_array.attrs.items()
if key not in new_array.attrs
}
return new_array.assign_attrs(new_attrs)
def add_experiment_metadata_to_data_array(self, data_array: xr.DataArray):
params = {}
params.update(self.config["pipeline"])
params.update(self.config["redox"])
params.update(self.config["registration"])
to_remove = ["trimmed_regions", "untrimmed_regions"]
for k in to_remove:
del params[k]
return data_array.assign_attrs(**params)
def _insert_units_if_necessary(da: xr.DataArray) -> xr.DataArray:
if "units" not in da.attrs:
return da.assign_attrs(units="")
else:
return da
def interactive_offset_choice(
guess: float, T_e: DataArray, flux: DataArray, offset_at_time: Optional[DataArray]
) -> Tuple[float, bool]:
"""Plots electron temperature against flux surface when using best
guess for location of magnetic axis. User then has a chance to accept this
value or specify their own alternative. The function can be called again
with the new value, if necessary.
Because user input is prompted on the command-line, this function is not
suitable for use within a GUI.
Parameters
----------
guess
Estimate for appropriate offset of the magnetic axis location.
T_e
Electron temperature data (from HRTS on JET).
flux
Normalised flux surface for each channel at each times.
offset_at_time
The optimal R-offset for each time. If present can be used for an
additional plot.
Returns
-------
offset : float
The offset along the major radius to apply to equilibrium data accepted
or suggested by the user.
accept : bool
Whether the user accepted the offset value or whether, instead, this is
a new suggestion. If the latter the function can be called again with
the new guess.
"""
print(
f"Displaying electron temperatures as a function of normalised magnetic "
f"flux\nfor the specified R-offset of {guess}m. Once finished inspecting "
"the plot(s),\nclose them."
)
etemps = Dataset(
{
"T_e": (T_e / 1e3).assign_attrs(
long_name=r"Electron Temperature $T_{\rm e}$", units="keV"
),
"rho_poloidal": flux.assign_attrs(
long_name=r"Poloidal Magnetic Flux $\rho_{\rm pol}$"
),
}
).isel(t=slice(None, None, 10))
etemps.plot.scatter(x="rho_poloidal", y="T_e")
plt.tight_layout()
if offset_at_time is not None:
plt.figure()
offset_at_time = cast(
DataArray,
offset_at_time.assign_attrs(long_name="Optimal R-offset", units="m"),
)
offset_at_time.plot()
plt.xlabel("Time [s]")
plt.tight_layout()
plt.show()
accept = ask_user(f"Use R-offset of {guess}m?")
if not accept:
while True:
try:
new_guess = float(input("Specify a new R-offset: "))
break
except ValueError:
print("Please enter a valid floating-point number.")
continue
else:
new_guess = guess
return new_guess, accept