From bb8a7762db6a0ba144d4b0847fd2154775860029 Mon Sep 17 00:00:00 2001 From: Marco Sertoli Date: Mon, 18 Sep 2023 15:37:00 +0100 Subject: [PATCH] Adding capability of recomputing rho-poloidal if fbnd or faxs are corrupted. --- indica/equilibrium.py | 24 ++++++++++++++++++------ 1 file changed, 18 insertions(+), 6 deletions(-) diff --git a/indica/equilibrium.py b/indica/equilibrium.py index 239505b2..04020285 100644 --- a/indica/equilibrium.py +++ b/indica/equilibrium.py @@ -65,12 +65,29 @@ def __init__( self.faxs = equilibrium_data["faxs"] self.fbnd = equilibrium_data["fbnd"] self.ftor = equilibrium_data["ftor"] + self.rmag = equilibrium_data["rmag"] + self.rbnd = equilibrium_data["rbnd"] + self.zmag = equilibrium_data["zmag"] + self.zbnd = equilibrium_data["zbnd"] + self.zx = self.zbnd.min("arbitrary_index") self.rhotor = np.sqrt( (self.ftor - self.ftor.sel(rho_poloidal=0.0)) / (self.ftor.sel(rho_poloidal=1.0) - self.ftor.sel(rho_poloidal=0.0)) ) self.psi = equilibrium_data["psi"] - self.rho = np.sqrt((self.psi - self.faxs) / (self.fbnd - self.faxs)) + + # Including workaround in case faxs or fbnd had messy data + rho: DataArray = np.sqrt((self.psi - self.faxs) / (self.fbnd - self.faxs)) + if np.any(np.isnan(rho.interp(R=self.rmag, z=self.zmag))): + self.faxs = self.psi.interp(R=self.rmag, z=self.zmag).drop(["R", "z"]) + self.fbnd = self.psi.interp(R=self.rbnd, z=self.zbnd).mean( + "arbitrary_index" + ) + rho = np.sqrt((self.psi - self.faxs) / (self.fbnd - self.faxs)) + if np.any(np.isnan(rho)): + rho = xr.where(rho > 0, rho, 0.0) + + self.rho = rho self.t = self.rho.t if "vjac" in equilibrium_data and "ajac" in equilibrium_data: self.psin = equilibrium_data["psin"] @@ -85,11 +102,6 @@ def __init__( if "rmji" and "rmjo" in equilibrium_data: self.rmji = equilibrium_data["rmji"] self.rmjo = equilibrium_data["rmjo"] - self.rmag = equilibrium_data["rmag"] - self.rbnd = equilibrium_data["rbnd"] - self.zmag = equilibrium_data["zmag"] - self.zbnd = equilibrium_data["zbnd"] - self.zx = self.zbnd.min("arbitrary_index") self.R_offset = R_shift self.z_offset = z_shift