Marcosertoli/equilibrium not nan

indica/converters/transect.py

@@ -69,35 +69,17 @@ def __init__(
         self._machine_dims = machine_dimensions
         self.x = xr.DataArray(
             x_positions,
-            coords={
-                self.x1_name: self.x1.data
-                if isinstance(self.x1, xr.DataArray)
-                else self.x1
-            },
-            dims=self.x1_name,
+            coords={self.x1_name: self.x1.data},
         )
         self.y = xr.DataArray(
             y_positions,
-            coords={
-                self.x1_name: self.x1.data
-                if isinstance(self.x1, xr.DataArray)
-                else self.x1
-            },
-            dims=self.x1.name,
+            coords={self.x1_name: self.x1.data},
         )
         self.z = xr.DataArray(
             z_positions,
-            coords={
-                self.x1_name: self.x1.data
-                if isinstance(self.x1, xr.DataArray)
-                else self.x1
-            },
-            dims=self.x1_name,
+            coords={self.x1_name: self.x1.data},
         )
 
-        # self.x: DataArray = DataArray(x_positions, coords=[(self.x1_name, self.x1)])
-        # self.y: DataArray = DataArray(y_positions, coords=[(self.x1_name, self.x1)])
-        # self.z: DataArray = DataArray(z_positions, coords=[(self.x1_name, self.x1)])
         self.R: DataArray = np.sqrt(self.x**2 + self.y**2)
         self.rho: DataArray
         self.theta: DataArray

indica/equilibrium.py

@@ -64,6 +64,20 @@ def __init__(
         self.rbnd = equilibrium_data["rbnd"]
         self.zmag = equilibrium_data["zmag"]
         self.zbnd = equilibrium_data["zbnd"]
+        if "rmji" and "rmjo" in equilibrium_data:
+            self.rmji = equilibrium_data["rmji"]
+            self.rmjo = equilibrium_data["rmjo"]
+        if "vjac" in equilibrium_data and "ajac" in equilibrium_data:
+            psin = equilibrium_data["vjac"].rho_poloidal ** 2
+            dpsin = psin[1] - psin[0]
+            self.volume = (equilibrium_data["vjac"] * dpsin).cumsum("rho_poloidal")
+            self.area = (equilibrium_data["ajac"] * dpsin).cumsum("rho_poloidal")
+        elif "volume" in equilibrium_data and "area" in equilibrium_data:
+            self.volume = equilibrium_data["volume"]
+            self.area = equilibrium_data["area"]
+        else:
+            raise ValueError("No volume or area information")
+
         self.zx = self.zbnd.min("arbitrary_index")
         self.rhotor = np.sqrt(
             (self.ftor - self.ftor.sel(rho_poloidal=0.0))
@@ -72,7 +86,9 @@ def __init__(
         self.psi = equilibrium_data["psi"]
         self.rho = np.sqrt((self.psi - self.faxs) / (self.fbnd - self.faxs))
 
-        # Shift equilibrium
+        # TODO: shift of equilibrium is a bad idea, but useful...
+        #   - psi (R, z) is restricted to new limits to avoid NaNs
+        #   - volume is not changed so is inconsistent!
         if isinstance(R_shift, float):
             R_offset = xr.full_like(self.t, R_shift)
         else:
@@ -81,45 +97,29 @@ def __init__(
             z_offset = xr.full_like(self.t, z_shift)
         else:
             z_offset = z_shift.interp(t=self.t, kwargs={"fill_value": 0})
-        self.R_offset = R_offset
-        self.z_offset = z_offset
-
-        R_new = self.psi.R + self.R_offset
-        z_new = self.psi.z + self.z_offset
-        self.psi = self.psi.interp(R=R_new, z=z_new)
-        self.rho = self.rho.interp(
-            R=R_new, z=z_new, kwargs={"fill_value": self.rho.max()}
-        )
-        self.rmag -= self.R_offset
-        self.zmag -= self.z_offset
-        self.rbnd -= self.R_offset
-        self.zbnd -= self.z_offset
-        self.zx -= self.z_offset
-
-        # Including workaround in case faxs or fbnd had messy data
-        if np.any(np.isnan(self.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"
-            )
+
+        self.R_offset = xr.where(np.isfinite(R_offset), R_offset, 0)
+        self.z_offset = xr.where(np.isfinite(z_offset), z_offset, 0)
+        if np.any(np.abs(R_offset) > 0) or np.any(np.abs(z_offset) > 0):
+            R_new = self.psi.R + self.R_offset
+            z_new = self.psi.z + self.z_offset
+            R_range = slice(R_new.min("R").max("t"), R_new.max("R").min("t"))
+            z_range = slice(z_new.min("z").max("t"), z_new.max("z").min("t"))
+            self.psi = self.psi.interp(R=R_new, z=z_new).sel(R=R_range, z=z_range)
+            self.rho = self.rho.interp(R=R_new, z=z_new).sel(R=R_range, z=z_range)
+            self.rmag -= self.R_offset
+            self.zmag -= self.z_offset
+            self.rbnd -= self.R_offset
+            self.zbnd -= self.z_offset
+            self.zx -= self.z_offset
+            if hasattr(self, "rmji"):
+                self.rmji -= self.R_offset
+            if hasattr(self, "rmjo"):
+                self.rmjo -= self.R_offset
+
         if np.any(np.isnan(self.rho)):
             self.rho = xr.where(self.rho > 0, self.rho, 0.0)
-
         self.t = self.rho.t
-        if "vjac" in equilibrium_data and "ajac" in equilibrium_data:
-            psin = equilibrium_data["vjac"].rho_poloidal ** 2
-            dpsin = psin[1] - psin[0]
-            self.volume = (equilibrium_data["vjac"] * dpsin).cumsum("rho_poloidal")
-            self.area = (equilibrium_data["ajac"] * dpsin).cumsum("rho_poloidal")
-        elif "volume" in equilibrium_data and "area" in equilibrium_data:
-            self.volume = equilibrium_data["volume"]
-            self.area = equilibrium_data["area"]
-        else:
-            raise ValueError("No volume or area information")
-        if "rmji" and "rmjo" in equilibrium_data:
-            self.rmji = equilibrium_data["rmji"] - self.R_offset
-            self.rmjo = equilibrium_data["rmjo"] - self.R_offset
-
         self.Rmin = min(self.rho.coords["R"])
         self.Rmax = max(self.rho.coords["R"])
         self.zmin = min(self.rho.coords["z"])