backsolve pedestal peaking in jch profiles and provide pedestal peaking outputs so that inputs and outputs are consistent regardless of which profiles are selected

Author: jlogan-cfs

cfspopcon/formulas/atomic_data/atomic_data.py

@@ -1,6 +1,7 @@
 """Module defining the AtomicData class, used for interfacing with radas files."""
 
 import warnings
+from collections.abc import Callable
 from pathlib import Path
 from typing import Union
 
@@ -14,6 +15,19 @@
 from .coeff_interpolator import CoeffInterpolator
 
 
+class _LazyValueDict(dict):
+    """Dictionary which populates missing values on first access."""
+
+    def __init__(self, loader: Callable):
+        super().__init__()
+        self._loader = loader
+
+    def __missing__(self, key):
+        value = self._loader(key)
+        self[key] = value
+        return value
+
+
 class AtomicData:
     """A class to manage atomic data for various species, providing facilities for accessing datasets directly or by constructing interpolators for different datasets (i.e. radiated power curves).
 
@@ -28,62 +42,116 @@ class AtomicData:
     """
 
     def __init__(self, atomic_data_directory: Path = Path() / "radas_dir") -> None:
-        """Initializes the AtomicData object by loading atomic data from the specified directory.
+        """Initializes the AtomicData object by indexing the atomic data files in the specified directory.
 
         Parameters:
         - atomic_data_directory (Path): The path to the directory containing atomic data files.
         """
         self.atomic_data_directory = atomic_data_directory
-        self.datasets = self.read_atomic_data(atomic_data_directory)  # Load atomic data into the datasets attribute
-        self.available_species = list(self.datasets.keys())  # List available species based on the loaded datasets
-
-        # Initialize dictionaries to hold interpolators for different data types and conditions
-        self.coronal_Lz_interpolators: dict[AtomicSpecies, CoeffInterpolator] = dict()
-        self.coronal_Z_interpolators: dict[AtomicSpecies, CoeffInterpolator] = dict()
-        self.noncoronal_Lz_interpolators: dict[tuple[AtomicSpecies, float], CoeffInterpolator] = dict()
-        self.noncoronal_Z_interpolators: dict[tuple[AtomicSpecies, float], CoeffInterpolator] = dict()
-
-        self.species_ne_tau: dict[AtomicSpecies, xr.DataArray] = dict()
+        self.atomic_data_files = self.find_atomic_data_files(atomic_data_directory)
+        self.datasets: dict[AtomicSpecies, xr.Dataset] = _LazyValueDict(self._load_dataset)
+        self.available_species = list(self.atomic_data_files.keys())
+
+        # Build interpolators only when a particular species/table is requested.
+        self.coronal_Lz_interpolators: dict[AtomicSpecies, CoeffInterpolator] = _LazyValueDict(self._load_coronal_Lz_interpolator)
+        self.coronal_Z_interpolators: dict[AtomicSpecies, CoeffInterpolator] = _LazyValueDict(self._load_coronal_Z_interpolator)
+        self.noncoronal_Lz_interpolators: dict[tuple[AtomicSpecies, float], CoeffInterpolator] = _LazyValueDict(
+            self._load_noncoronal_Lz_interpolator
+        )
+        self.noncoronal_Z_interpolators: dict[tuple[AtomicSpecies, float], CoeffInterpolator] = _LazyValueDict(
+            self._load_noncoronal_Z_interpolator
+        )
+
+        self.species_ne_tau: dict[AtomicSpecies, xr.DataArray] = _LazyValueDict(self._load_species_ne_tau)
         self.ne_tau_units = ureg.m**-3 * ureg.s
 
-        self.radas_version: str = ""
+        self._radas_version: str = ""
+        self._radas_version_checked_species: set[AtomicSpecies] = set()
 
+    @property
+    def radas_version(self) -> str:
+        """Return the RADAS version after checking the available datasets."""
         for species in self.available_species:
-            dataset = self[species]
-
-            ref = dict(
-                reference_electron_density=dataset.reference_electron_density,
-                reference_electron_temp=dataset.reference_electron_temp,
-            )
-
-            self.coronal_Lz_interpolators[species] = CoeffInterpolator(dataset.coronal_Lz, **ref)
-            self.coronal_Z_interpolators[species] = CoeffInterpolator(dataset.coronal_mean_charge_state, **ref)
-
-            self.species_ne_tau[species] = dataset["ne_tau"].pint.to(self.ne_tau_units).pint.dequantify()
-
-            for ne_tau, dataset_at_single_ne_tau in dataset.groupby("dim_ne_tau"):
-                subds = dataset_at_single_ne_tau.squeeze(dim="dim_ne_tau")
-                self.noncoronal_Lz_interpolators[(species, ne_tau)] = CoeffInterpolator(subds.equilibrium_Lz, **ref)
-                self.noncoronal_Z_interpolators[(species, ne_tau)] = CoeffInterpolator(subds.equilibrium_mean_charge_state, **ref)
-
-                self._check_radas_version(getattr(dataset, "radas_version", "UNDEFINED"))
+            _ = self[species]
+        return self._radas_version
+
+    def _load_dataset(self, species: AtomicSpecies) -> xr.Dataset:
+        dataset = xr.open_dataset(self.atomic_data_files[species]).pint.quantify()
+        if species not in self._radas_version_checked_species:
+            self._check_radas_version(getattr(dataset, "radas_version", "UNDEFINED"))
+            self._radas_version_checked_species.add(species)
+        return dataset
+
+    def _get_reference_values(self, species: AtomicSpecies) -> dict[str, Quantity]:
+        dataset = self[species]
+        return dict(
+            reference_electron_density=dataset.reference_electron_density,
+            reference_electron_temp=dataset.reference_electron_temp,
+        )
+
+    def _load_coronal_Lz_interpolator(self, species: AtomicSpecies) -> CoeffInterpolator:
+        dataset = self[species]
+        return CoeffInterpolator(dataset.coronal_Lz, **self._get_reference_values(species))
+
+    def _load_coronal_Z_interpolator(self, species: AtomicSpecies) -> CoeffInterpolator:
+        dataset = self[species]
+        return CoeffInterpolator(dataset.coronal_mean_charge_state, **self._get_reference_values(species))
+
+    def _load_species_ne_tau(self, species: AtomicSpecies) -> xr.DataArray:
+        return self[species]["ne_tau"].pint.to(self.ne_tau_units).pint.dequantify()
+
+    def _load_noncoronal_Lz_interpolator(self, key: tuple[AtomicSpecies, float]) -> CoeffInterpolator:
+        species, ne_tau = key
+        dataset = self[species].sel(dim_ne_tau=ne_tau)
+        if "dim_ne_tau" in dataset.dims:
+            dataset = dataset.squeeze(dim="dim_ne_tau")
+        return CoeffInterpolator(dataset.equilibrium_Lz, **self._get_reference_values(species))
+
+    def _load_noncoronal_Z_interpolator(self, key: tuple[AtomicSpecies, float]) -> CoeffInterpolator:
+        species, ne_tau = key
+        dataset = self[species].sel(dim_ne_tau=ne_tau)
+        if "dim_ne_tau" in dataset.dims:
+            dataset = dataset.squeeze(dim="dim_ne_tau")
+        return CoeffInterpolator(dataset.equilibrium_mean_charge_state, **self._get_reference_values(species))
 
     def _check_radas_version(self, test_version: str) -> None:
         """Checks that the provided test_version matches radas_version (if set).
 
         If radas_version is not set, sets radas_version = test_version.
         If a mismatch is found, sets radas_version = UNDEFINED.
         """
-        if self.radas_version == "":
-            self.radas_version = test_version
-        elif self.radas_version != test_version:
+        if self._radas_version == "":
+            self._radas_version = test_version
+        elif self._radas_version != test_version:
             warnings.warn(
-                f"Found multiple radas radas versions ({self.radas_version} != {test_version}) in the requested atomic data. Will set radas_version = UNDEFINED.",
+                f"Found multiple radas radas versions ({self._radas_version} != {test_version}) in the requested atomic data. Will set radas_version = UNDEFINED.",
                 stacklevel=2,
             )
 
     @staticmethod
-    def read_atomic_data(atomic_data_directory: Path = Path() / "radas_dir") -> dict[AtomicSpecies, xr.Dataset]:
+    def find_atomic_data_files(atomic_data_directory: Path = Path() / "radas_dir") -> dict[AtomicSpecies, Path]:
+        """Return the available RADAS netCDF files indexed by species."""
+        if not atomic_data_directory.exists():
+            raise FileNotFoundError(f"atomic_data_directory ({atomic_data_directory.absolute()}) does not exist.")
+
+        if not (atomic_data_directory / "output").exists():
+            raise FileNotFoundError(
+                f"atomic_data_directory ({atomic_data_directory}) does not contain a subfolder called 'output'. Make sure you have executed `poetry run radas` before calling this function."
+            )
+
+        atomic_data_files: dict[AtomicSpecies, Path] = dict()
+        for file in (atomic_data_directory / "output").iterdir():
+            if file.suffix == ".nc":
+                species = file.stem
+                try:
+                    atomic_data_files[AtomicSpecies[species.capitalize()]] = file
+                except KeyError:
+                    print(f"No AtomicSpecies found corresponding to {species}")
+
+        return atomic_data_files
+
+    @classmethod
+    def read_atomic_data(cls, atomic_data_directory: Path = Path() / "radas_dir") -> dict[AtomicSpecies, xr.Dataset]:
         """Reads atomic data from netCDF files located in the specified directory.
 
         This function scans a directory for netCDF files (.nc), each representing atomic data for a different species.
@@ -99,30 +167,7 @@ def read_atomic_data(atomic_data_directory: Path = Path() / "radas_dir") -> dict
         Raises:
         - FileNotFoundError: If the atomic_data_directory or its 'output' subdirectory does not exist.
         """
-        # Ensure the atomic data directory and its 'output' subdirectory exist
-        if not atomic_data_directory.exists():
-            raise FileNotFoundError(f"atomic_data_directory ({atomic_data_directory.absolute()}) does not exist.")
-
-        if not (atomic_data_directory / "output").exists():
-            raise FileNotFoundError(
-                f"atomic_data_directory ({atomic_data_directory}) does not contain a subfolder called 'output'. Make sure you have executed `poetry run radas` before calling this function."
-            )
-
-        atomic_data = dict()
-        # Iterate through each netCDF file in the 'output' directory
-        for file in (atomic_data_directory / "output").iterdir():
-            if file.suffix == ".nc":  # Check if the file is a netCDF file
-                species = file.stem  # Extract the species name from the file name
-                try:
-                    # Attempt to map the file name to an AtomicSpecies enum
-                    species_enum = AtomicSpecies[species.capitalize()]
-                    # Read the netCDF file into an xarray Dataset and quantify it with pint
-                    atomic_data[species_enum] = xr.open_dataset(file).pint.quantify()
-                except KeyError:
-                    # If no matching AtomicSpecies enum is found, print a warning
-                    print(f"No AtomicSpecies found corresponding to {species}")
-
-        return atomic_data
+        return {species: xr.open_dataset(file).pint.quantify() for species, file in cls.find_atomic_data_files(atomic_data_directory).items()}
 
     @staticmethod
     def key_to_enum(species: Union[str, AtomicSpecies]) -> AtomicSpecies: