musica 0.14.4__cp310-cp310-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl

musica/tuvx/radiator_map.py

@@ -0,0 +1,173 @@
+# Copyright (C) 2023-2026 University Corporation for Atmospheric Research
+# SPDX-License-Identifier: Apache-2.0
+"""
+TUV-x RadiatorMap class.
+
+This module provides a class for managing collections of TUV-x radiators.
+The RadiatorMap class allows dictionary-style access to radiators using their names as keys.
+
+Note: TUV-x is only available on macOS and Linux platforms.
+"""
+
+from typing import Iterator
+from .. import backend
+from .radiator import Radiator
+
+_backend = backend.get_backend()
+
+RadiatorMap = _backend._tuvx._RadiatorMap if backend.tuvx_available() else None
+
+if backend.tuvx_available():
+    original_init = RadiatorMap.__init__
+
+    def __init__(self, **kwargs):
+        """Initialize a RadiatorMap instance.
+
+        Args:
+            **kwargs: Additional arguments passed to the C++ constructor
+        """
+        original_init(self, **kwargs)
+
+    RadiatorMap.__init__ = __init__
+
+    def __str__(self):
+        """User-friendly string representation."""
+        return f"RadiatorMap(num_radiators={len(self)})"
+
+    RadiatorMap.__str__ = __str__
+
+    def __repr__(self):
+        """Detailed string representation for debugging."""
+        radiator_details = []
+        for i in range(len(self)):
+            radiator = self.get_radiator_by_index(i)
+            radiator_details.append(f"{radiator.name}")
+        return f"RadiatorMap(radiators={radiator_details})"
+
+    RadiatorMap.__repr__ = __repr__
+
+    def __len__(self):
+        """Return the number of radiators in the map."""
+        return self.get_number_of_radiators()
+
+    RadiatorMap.__len__ = __len__
+
+    def __bool__(self):
+        """Return True if the map has any radiators."""
+        return len(self) > 0
+
+    RadiatorMap.__bool__ = __bool__
+
+    def __getitem__(self, key) -> Radiator:
+        """Get a radiator by name.
+
+        Args:
+            key: Name of the radiator to retrieve
+
+        Returns:
+            Radiator instance corresponding to the given name
+
+        Raises:
+            KeyError: If no radiator with the given name exists
+        """
+        if not isinstance(key, str):
+            raise TypeError("Radiator name must be a string.")
+        try:
+            radiator = self.get_radiator(key)
+            if radiator is None:
+                raise KeyError(f"Radiator with name '{key}' not found.")
+            return radiator
+        except Exception as e:
+            raise KeyError(f"No radiator found with name='{key}'") from e
+
+    RadiatorMap.__getitem__ = __getitem__
+
+    def __setitem__(self, key: str, radiator: Radiator):
+        """Set a radiator in the map by name.
+
+        Args:
+            key: Name of the radiator to set
+            radiator: Radiator instance to associate with the given name
+        """
+        if not isinstance(key, str):
+            raise TypeError("Radiator name must be a string.")
+        if not isinstance(radiator, Radiator):
+            raise TypeError("Value must be a Radiator instance.")
+        if radiator.name != key:
+            raise ValueError("Radiator name does not match the key.")
+        self.add_radiator(radiator)
+
+    RadiatorMap.__setitem__ = __setitem__
+
+    def __iter__(self) -> Iterator[Radiator]:
+        """Iterator over the radiators in the map.
+
+        Yields:
+            Radiator instances in the map
+        """
+        for i in range(len(self)):
+            yield self.get_radiator_by_index(i).name
+
+    RadiatorMap.__iter__ = __iter__
+
+    def __contains__(self, key: str) -> bool:
+        """Check if a radiator with the given name exists in the map.
+
+        Args:
+            key: Name of the radiator to check
+
+        Returns:
+            True if a radiator with the given name exists, False otherwise
+        """
+        if not isinstance(key, str):
+            return False
+        try:
+            radiator = self.get_radiator(str(key))
+            return radiator is not None
+        except (ValueError, KeyError):
+            return False
+
+    RadiatorMap.__contains__ = __contains__
+
+    def clear(self):
+        """Remove all radiators from the map."""
+        num_radiators = len(self)
+        for _ in range(num_radiators):
+            self.remove_radiator_by_index(0)
+
+    RadiatorMap.clear = clear
+
+    def items(self):
+        """Iterator over (name, radiator) pairs in the map.
+
+        Yields:
+            Tuples of (radiator_name, Radiator instance)
+        """
+        for i in range(len(self)):
+            radiator = self.get_radiator_by_index(i)
+            yield (radiator.name, radiator)
+
+    RadiatorMap.items = items
+
+    def keys(self):
+        """Iterator over radiator names in the map.
+
+        Yields:
+            Radiator names as strings
+        """
+        for i in range(len(self)):
+            radiator = self.get_radiator_by_index(i)
+            yield radiator.name
+
+    RadiatorMap.keys = keys
+
+    def values(self):
+        """Iterator over Radiator instances in the map.
+
+        Yields:
+            Radiator instances
+        """
+        for i in range(len(self)):
+            yield self.get_radiator_by_index(i)
+
+    RadiatorMap.values = values

musica/tuvx/tuvx.py

@@ -0,0 +1,283 @@
+"""
+TUV-x photolysis calculator Python interface.
+
+This module provides a simplified Python interface to the TUV-x photolysis calculator.
+It allows users to create a TUV-x instance from a JSON configuration file and
+calculate photolysis rates and heating rates.
+
+Note: TUV-x is only available on macOS and Linux platforms.
+"""
+
+import os
+import json
+import tempfile
+from typing import Dict, Optional, Tuple
+import numpy as np
+from .. import backend
+
+_backend = backend.get_backend()
+
+# Import the GridMap class from the backend
+if backend.tuvx_available():
+    from .grid_map import GridMap
+    from .profile import Profile
+    from .profile_map import ProfileMap
+    from .radiator import Radiator
+    from .radiator_map import RadiatorMap
+else:
+    GridMap = None
+    Profile = None
+    ProfileMap = None
+    Radiator = None
+    RadiatorMap = None
+
+
+class TUVX:
+    """
+    A Python interface to the TUV-x photolysis calculator.
+
+    This class provides a simplified interface that only requires a JSON/YAML configuration
+    to set up and run TUV-x calculations. All parameters (solar zenith angle,
+    earth-sun distance, atmospheric profiles, etc.) are specified in the JSON/YAML configuration.
+
+    The configuration can be in a file or provided as a string. Exactly one of `config_path` or
+    `config_string` must be provided.
+    """
+
+    def __init__(self,
+                 grid_map: GridMap,
+                 profile_map: ProfileMap,
+                 radiator_map: RadiatorMap,
+                 config_path: Optional[str] = None,
+                 config_string: Optional[str] = None):
+        """
+        Initialize a TUV-x instance from a configuration file.
+
+        Args:
+            grid_map: GridMap instance containing grid definitions (height, wavelength)
+            profile_map: ProfileMap instance containing atmospheric profiles (temperature, species concentrations, surface albedo, ET flux)
+            radiator_map: RadiatorMap instance containing radiator definitions (optically active species)
+            config_path: Path to the JSON configuration file
+            config_string: JSON configuration as a string
+
+        Raises:
+            FileNotFoundError: If the configuration file doesn't exist
+            ValueError: If TUV-x initialization fails or TUVX is not available
+        """
+        if not backend.tuvx_available():
+            raise ValueError("TUV-x backend is not available.")
+
+        if (config_path is None and config_string is None) or \
+           (config_path is not None and config_string is not None):
+            raise ValueError(
+                "Exactly one of config_path or config_string must be provided.")
+
+        if config_path is not None and not os.path.exists(config_path):
+            raise FileNotFoundError(
+                f"Configuration file not found: {config_path}")
+
+        if config_path is not None:
+            self._tuvx_instance = _backend._tuvx._create_tuvx_from_file(
+                config_path, grid_map, profile_map, radiator_map)
+            self._config_path = config_path
+            self._config_string = None
+        elif config_string is not None:
+            self._tuvx_instance = _backend._tuvx._create_tuvx_from_string(
+                config_string, grid_map, profile_map, radiator_map)
+            self._config_path = None
+            self._config_string = config_string
+
+        # Cache the names for efficiency
+        self._photolysis_names = None
+        self._heating_names = None
+        self._dose_names = None
+
+    def __del__(self):
+        """Clean up the TUV-x instance."""
+        if hasattr(self, '_tuvx_instance') and self._tuvx_instance is not None:
+            _backend._tuvx._delete_tuvx(self._tuvx_instance)
+
+    @property
+    def photolysis_rate_names(self) -> dict[str, int]:
+        """
+        Get the names of photolysis rates.
+
+        Returns:
+            Dictionary mapping photolysis rate names to their indices in the output arrays
+        """
+        if self._photolysis_names is None:
+            self._photolysis_names = _backend._tuvx._get_photolysis_rate_constants_ordering(
+                self._tuvx_instance)
+        return self._photolysis_names
+
+    @property
+    def heating_rate_names(self) -> dict[str, int]:
+        """
+        Get the names of heating rates.
+
+        Returns:
+            Dictionary mapping heating rate names to their indices in the output arrays
+        """
+        if self._heating_names is None:
+            self._heating_names = _backend._tuvx._get_heating_rates_ordering(
+                self._tuvx_instance)
+        return self._heating_names
+
+    @property
+    def dose_rate_names(self) -> dict[str, int]:
+        """
+        Get the names of dose rates.
+
+        Returns:
+            Dictionary mapping dose rate names to their indices in the output arrays
+        """
+        if self._dose_names is None:
+            self._dose_names = _backend._tuvx._get_dose_rates_ordering(
+                self._tuvx_instance)
+        return self._dose_names
+
+    def run(self, sza: float, earth_sun_distance: float) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+        """
+        Run the TUV-x photolysis calculator.
+
+        All parameters (solar zenith angle, Earth-Sun distance, atmospheric profiles,
+        etc.) are read from the JSON configuration file.
+
+        Args:
+            sza: Solar zenith angle in radians
+            earth_sun_distance: Earth-Sun distance in astronomical units (AU)
+
+        Returns:
+            Tuple of (photolysis_rate_constants, heating_rates) as numpy arrays
+            - photolysis_rate_constants: Shape (n_layers, n_reactions) [s^-1]
+            - heating_rates: Shape (n_layers, n_heating_rates) [K s^-1]
+            - dose_rates: Shape (n_layers, n_dose_rates) [W m^-2]
+        """
+        photolysis_rates, heating_rates, dose_rates = _backend._tuvx._run_tuvx(
+            self._tuvx_instance, sza, earth_sun_distance)
+
+        return photolysis_rates, heating_rates, dose_rates
+
+    def get_photolysis_rate_constant(
+        self,
+        reaction_name: str,
+        photolysis_rates: np.ndarray
+    ) -> np.ndarray:
+        """
+        Extract photolysis rate constants for a specific reaction.
+
+        Args:
+            reaction_name: Name of the photolysis reaction
+            photolysis_rates: Output from run() method
+
+        Returns:
+            1D array of photolysis rate constants for all layers [s^-1]
+
+        Raises:
+            KeyError: If reaction_name is not found
+        """
+        names = self.photolysis_rate_names
+        if reaction_name not in names:
+            raise KeyError(
+                f"Reaction '{reaction_name}' not found. "
+                f"Available reactions: {names}"
+            )
+
+        reaction_index = names[reaction_name]
+        return photolysis_rates[:, reaction_index]
+
+    def get_heating_rate(
+        self,
+        rate_name: str,
+        heating_rates: np.ndarray
+    ) -> np.ndarray:
+        """
+        Extract heating rates for a specific rate type.
+
+        Args:
+            rate_name: Name of the heating rate
+            heating_rates: Output from run() method
+
+        Returns:
+            1D array of heating rates for all layers [K s^-1]
+
+        Raises:
+            KeyError: If rate_name is not found
+        """
+        names = self.heating_rate_names
+        if rate_name not in names:
+            raise KeyError(
+                f"Heating rate '{rate_name}' not found. "
+                f"Available rates: {names}"
+            )
+
+        rate_index = names[rate_name]
+        return heating_rates[:, rate_index]
+
+    def get_dose_rate(
+        self,
+        rate_name: str,
+        dose_rates: np.ndarray
+    ) -> np.ndarray:
+        """
+        Extract dose rates for a specific rate type.
+
+        Args:
+            rate_name: Name of the dose rate
+            dose_rates: Output from run() method
+
+        Returns:
+            1D array of dose rates for all layers [W m^-2]
+
+        Raises:
+            KeyError: If rate_name is not found
+        """
+        names = self.dose_rate_names
+        if rate_name not in names:
+            raise KeyError(
+                f"Dose rate '{rate_name}' not found. "
+                f"Available rates: {names}"
+            )
+
+        rate_index = names[rate_name]
+        return dose_rates[:, rate_index]
+
+    @staticmethod
+    def create_config_from_dict(config_dict: Dict) -> 'TUVX':
+        """
+        Create a TUVX instance from a configuration dictionary.
+
+        Args:
+            config_dict: Configuration dictionary
+
+        Returns:
+            TUVX instance initialized with the configuration
+
+        Raises:
+            ValueError: If TUV-x backend is not available
+            FileNotFoundError: If required data files are not found
+        """
+        with tempfile.NamedTemporaryFile(
+                mode='w', suffix='.json', delete=True) as temp_file:
+            json.dump(config_dict, temp_file, indent=2)
+            temp_file.flush()  # Ensure all data is written to disk
+            return TUVX(temp_file.name)
+
+    @staticmethod
+    def create_config_from_json_string(json_string: str) -> 'TUVX':
+        """
+        Create a TUVX instance from a JSON configuration string.
+
+        Args:
+            json_string: JSON configuration as string
+
+        Returns:
+            TUVX instance initialized with the configuration
+
+        Raises:
+            json.JSONDecodeError: If json_string is not valid JSON
+            ValueError: If TUV-x backend is not available
+            FileNotFoundError: If required data files are not found
+        """
+        config_dict = json.loads(json_string)
+        return TUVX.create_config_from_dict(config_dict)