ids_toplevel import and geometry_directory variable name

Author: mikesndrs

torax/_src/geometry/pydantic_model.py

@@ -19,6 +19,9 @@
 import inspect
 from typing import Annotated, Any, Literal, TypeAlias, TypeVar
 import pydantic
+
+from imas.ids_toplevel import IDSToplevel
+
 from torax._src.geometry import circular_geometry
 from torax._src.geometry import geometry
 from torax._src.geometry import geometry_provider
@@ -257,7 +260,7 @@ class IMASConfig(torax_pydantic.BaseModelFrozen):
     hires_factor: Only used when the initial condition ``psi`` is from plasma
       current. Sets up a higher resolution mesh with ``nrho_hires = nrho *
       hi_res_fac``, used for ``j`` to ``psi`` conversions.
-    geometry_dir: Optionally overrides the default geometry directory.
+    geometry_directory: Optionally overrides the default geometry directory.
     equilibrium_object: Either directly the equilbrium IDS containing the relevant data,
       or the path to the IMAS netCDF file containing the equilibrium.
     Ip_from_parameters: Toggles whether total plasma current is read from the
@@ -268,17 +271,17 @@ class IMASConfig(torax_pydantic.BaseModelFrozen):
   geometry_type: Annotated[Literal['imas'], TIME_INVARIANT] = 'imas'
   n_rho: Annotated[pydantic.PositiveInt, TIME_INVARIANT] = 25
   hires_factor: pydantic.PositiveInt = 4
-  geometry_dir: Annotated[str | None, TIME_INVARIANT] = None
-  equilibrium_object: str | Any = 'ITERhybrid_COCOS17_IDS_ddv4.nc'
-  Ip_from_parameters: Annotated[bool, TIME_INVARIANT] = False
+  geometry_directory: Annotated[str | None, TIME_INVARIANT] = None
+  equilibrium_object: str | IDSToplevel = 'ITERhybrid_COCOS17_IDS_ddv4.nc'
+  Ip_from_parameters: Annotated[bool, TIME_INVARIANT] = True
 
   @pydantic.model_validator(mode='after')
   def _validate_model(self) -> typing_extensions.Self:
     if isinstance(self.equilibrium_object, str) and self.equilibrium_object[-3:] == '.h5':
       raise ValueError(
-          "If you are using hdf5 backend, the path to the data must point the directory containing the equilibrium.h5 and master.h5 files. As the function concatenates the str for geometry_dir and equilibrium_object to give the path, your \
-          equilibrium_object must be either the repository containing these files or an empty string '' (if your geometry_dir is already this specific repository).\n \
-        In any case, make sure geometry_dir + equilibrium_object gives the path to this directory and not to the .h5 file."
+          "If you are using hdf5 backend, the path to the data must point the directory containing the equilibrium.h5 and master.h5 files. As the function concatenates the str for geometry_directory and equilibrium_object to give the path, your \
+          equilibrium_object must be either the repository containing these files or an empty string '' (if your geometry_directory is already this specific repository).\n \
+        In any case, make sure geometry_directory + equilibrium_object gives the path to this directory and not to the .h5 file."
       )
     return self
 

torax/_src/geometry/standard_geometry.py

@@ -933,7 +933,7 @@ def calculate_area(x, z):
   def from_IMAS(
       cls,
       equilibrium_object: str | Any,
-      geometry_dir: str | None,
+      geometry_directory: str | None,
       Ip_from_parameters: bool,
       n_rho: int,
       hires_factor: int,
@@ -942,11 +942,11 @@ def from_IMAS(
 
     Args:
       equilibrium_object: Either directly the equilbrium IDS containing the relevant data, or the name of the IMAS netCDF file containing the equilibrium.
-      geometry_dir: Directory where to find the scenario file ontaining the parameters of the Data entry to read.
+      geometry_directory: Directory where to find the scenario file ontaining the parameters of the Data entry to read.
         If None, then it defaults to another dir. See implementation.
       Ip_from_parameters: If True, the Ip is taken from the parameters and the
-        values in the Geometry are resacled to match the new Ip.
-      n_rho: Radial grid points (num cells)
+        values in the Geometry are rescaled to match the new Ip.
+      n_rho: Radial grid points (num cells).
       hires_factor: Grid refinement factor for poloidal flux <--> plasma current
         calculations.
 
@@ -956,7 +956,7 @@ def from_IMAS(
     """
     inputs = imas_equilibrium.geometry_from_IMAS(
         equilibrium_object=equilibrium_object,
-        geometry_dir=geometry_dir,
+        geometry_directory=geometry_directory,
         Ip_from_parameters=Ip_from_parameters,
         n_rho=n_rho,
         hires_factor=hires_factor,

torax/imas_tools/equilibrium.py

@@ -49,7 +49,7 @@ def write_ids_equilibrium_into_config(
 
 def geometry_from_IMAS(
     equilibrium_object: str | IDSToplevel,
-    geometry_dir: str | None = None,
+    geometry_directory: str | None = None,
     Ip_from_parameters: bool = False,
     n_rho: int = 25,
     hires_factor: int = 4,
@@ -59,7 +59,7 @@ def geometry_from_IMAS(
       equilibrium_object: Either directly the equilbrium IDS containing the
         relevant data, or the name of the IMAS netCDF file containing the
         equilibrium.
-      geometry_dir: Directory where to find the equilibrium object.
+      geometry_directory: Directory where to find the equilibrium object.
         If None, it defaults to another dir. See `load_geo_data`
         implementation.
       Ip_from_parameters: If True, the Ip is taken from the parameters and the
@@ -74,7 +74,7 @@ def geometry_from_IMAS(
     # If the equilibrium_object is the file name, load the ids from the netCDF.
     if isinstance(equilibrium_object, str):
         equilibrium = geometry_loader.load_geo_data(
-            geometry_dir,
+            geometry_directory,
             equilibrium_object,
             geometry_loader.GeometrySource.IMAS,
         )
@@ -132,6 +132,11 @@ def geometry_from_IMAS(
     # -> Ip_profile = integrate(y = spr * jtor, x= rhon, initial = 0.0)
     jtor = -1 * IMAS_data.profiles_1d.j_phi
     rhon = IMAS_data.profiles_1d.rho_tor_norm
+    if len(rhon) == 0:
+        if B_0 is None or len(IMAS_data.profiles_1d.phi) == 0:
+            raise ValueError("rho_tor_norm not provided and cannot be calculated from given equilibrium IDS")
+        rho_tor = np.sqrt(IMAS_data.profiles_1d.phi / (np.pi * B_0))
+        rhon = rho_tor / rho_tor[-1]
     vpr = 4 * np.pi * Phi[-1] * rhon / (F * flux_surf_avg_1_over_R2)
     spr = vpr / (2 * np.pi * R_major)
     Ip_profile_unscaled = scipy.integrate.cumulative_trapezoid(y=spr * jtor, x=rhon, initial=0.0)    # this Ip_profile by integration results in a discrepancy between this term and the total ip from IDS

torax/imas_tools/tests/equilibrium_test.py

@@ -56,9 +56,9 @@ def test_save_geometry_to_IMAS(
             atol = self.atol
         # Input equilibrium reading
         config_module = self._get_config_module(config_name)
-        geometry_dir = "torax/data/third_party/geo"
+        geometry_directory = "torax/data/third_party/geo"
         path = os.path.join(
-            geometry_dir, config_module.CONFIG["geometry"]["equilibrium_object"]
+            geometry_directory, config_module.CONFIG["geometry"]["equilibrium_object"]
         )
         equilibrium_in = imas_util.load_IMAS_data(path, "equilibrium")
         # Build TORAXSimState object and write output to equilibrium IDS.

torax/imas_tools/util.py

@@ -48,26 +48,14 @@ def save_netCDF(
         date_str = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
         file_name = "IDS_file_" + date_str
     filepath = os.path.join(directory_path, file_name) + ".nc"
-    with imas.DBEntry(filepath, "w") as netcdf_entry:
-        netcdf_entry.put(IDS)
+    with imas.DBEntry(uri=filepath, mode="w") as netcdf_entry:
+        netcdf_entry.put(ids=IDS)
 
 
 def load_IMAS_data(uri: str, ids_name: str) -> IDSToplevel:
     """
     Loads a full IDS for a given uri or path_name and a given ids_name.
     """
-    db = imas.DBEntry(uri, "r")
-    ids = db.get(ids_name)
+    db = imas.DBEntry(uri=uri, mode="r")
+    ids = db.get(ids_name=ids_name)
     return ids
-
-# todo check if we can copy from geometry without weird dependency loops
-def face_to_cell(face):
-    """Infers cell values corresponding to a vector of face values.
-    Args:
-      face: An array containing face values.
-
-    Returns:
-      cell: An array containing cell values.
-    """
-
-    return 0.5 * (face[:-1] + face[1:])