Source code for uclchem.advanced.advanced_heating

"""Heating and cooling mechanism configuration for UCLCHEM.

This module provides class-based interface for accessing and modifying the Fortran
heating module that controls UCLCHEM's heating and cooling mechanisms during execution.

**Thread Safety Warning:**
HeatingSettings modifies global Fortran module state and is **NOT thread-safe**.
Do not use with multiprocessing, multithreading, or concurrent model runs.
Settings should only be modified during initialization, before running models.

Note: Changes made through HeatingSettings affect the global Fortran state and persist
across model runs in the same Python session.

"""

import importlib.util
import logging
import os
import warnings
from pathlib import Path

import numpy as np
import uclchemwrap
from uclchemwrap import f2py_constants as f2py_constants_module
from uclchemwrap import heating as heating_module

[docs] logger = logging.getLogger(__name__)
[docs] class HeatingSettings: """Class-based wrapper for UCLCHEM heating and cooling configuration. This class provides access to the heating and cooling mechanism controls, solver parameters, and computed values from heating.f90. Attributes ---------- PHOTOELECTRIC : dict Photoelectric heating implementations - 'BAKES': Bakes & Tielens method (index 1) - 'WEINGARTNER': Weingartner method (index 2) (Only one can be enabled at a time) H2_FORMATION : int Index for H2 formation heating (3) H2_PHOTODISSOCIATION : int Index for H2 photodissociation heating (4) H2_FUV_PUMPING : int Index for H2 FUV pumping heating (5) CARBON_IONIZATION : int Index for carbon ionization heating (6) COSMIC_RAY : int Index for cosmic ray heating (7) TURBULENT : int Index for turbulent heating (8) GAS_GRAIN_COLLISIONS : int Index for gas-grain collisional heating/cooling (9) ATOMIC_LINE_COOLING : int Index for atomic line cooling H2_COLLISIONALLY_INDUCED : int Index for H2 collisionally induced emission COMPTON_COOLING : int Index for Compton cooling CONTINUUM_EMISSION : int Index for continuum emission cooling MOLECULAR_LINE_COOLING : int Index for molecular line cooling DUST_TEMP_HOCUK : int Hocuk et al. 2017 dust temperature method DUST_TEMP_HOLLENBACH : int Hollenbach 1991 dust temperature method COOLANT_WARM : int Warm restart mode - initialize LTE, rescale on density change (default) COOLANT_FORCE_LTE : int Always reset to LTE before SE iteration (original behavior) COOLANT_FORCE_GROUND : int Always reset to ground state before SE iteration Examples -------- >>> from uclchem.advanced import HeatingSettings >>> settings = HeatingSettings() >>> settings.print_configuration() # doctest: +SKIP ... >>> # Switch photoelectric method (auto-disables Bakes when enabling Weingartner) >>> settings.set_heating_mechanism(settings.PHOTOELECTRIC['WEINGARTNER'], True) >>> # Disable H2 formation >>> settings.set_heating_mechanism(settings.H2_FORMATION, False) """ # Heating mechanism indices (matching heating.f90) # Mechanisms with multiple implementations are dicts; only one can be enabled at a time # Fortran Indexed:
[docs] PHOTOELECTRIC = { "BAKES": 1, # Bakes & Tielens photoelectric heating "WEINGARTNER": 2, # Weingartner photoelectric heating }
[docs] H2_FORMATION = 3
[docs] H2_PHOTODISSOCIATION = 4
[docs] H2_FUV_PUMPING = 5
[docs] CARBON_IONIZATION = 6
[docs] COSMIC_RAY = 7
[docs] TURBULENT = 8
[docs] GAS_GRAIN_COLLISIONS = 9
# Cooling mechanism indices (matching heating.f90)
[docs] ATOMIC_LINE_COOLING = 1
[docs] H2_COLLISIONALLY_INDUCED = 2
[docs] COMPTON_COOLING = 3
[docs] CONTINUUM_EMISSION = 4
[docs] MOLECULAR_LINE_COOLING = 5
# Dust-gas coupling methods
[docs] DUST_TEMP_HOCUK = 1 # Hocuk et al. 2017 parametric formulation
[docs] DUST_TEMP_HOLLENBACH = 2 # Hollenbach 1991 detailed balance method
# Coolant population restart modes
[docs] COOLANT_WARM = 0 # Initialize to LTE on first call, then rescale on density change (default, most efficient)
[docs] COOLANT_FORCE_LTE = 1 # Always reset to LTE before SE iteration (original behavior)
[docs] COOLANT_FORCE_GROUND = 2 # Always reset to ground state before SE iteration
def __init__(self): """Initialize the HeatingSettings wrapper. This connects to the Fortran heating module and provides access to its configuration parameters. Captures the current state as the default configuration for reset_to_defaults(). """ self._heating_module = heating_module self._f2py_constants_module = f2py_constants_module self._uclchemwrap = uclchemwrap # Build a mapping of mechanism IDs to their groups (for mutual exclusion) self._heating_groups = {} for attr_name in dir(self): if attr_name.isupper() and not attr_name.startswith("_"): attr = getattr(self, attr_name) if isinstance(attr, dict): # This is a group with multiple implementations for impl_id in attr.values(): self._heating_groups[impl_id] = list(attr.values()) # Backup initial state for reset_to_defaults() self._default_heating_modules = self._heating_module.heating_modules.copy() self._default_cooling_modules = self._heating_module.cooling_modules.copy() self._default_dust_gas_coupling_method = ( self._heating_module.dust_gas_coupling_method ) self._default_line_solver_attempts = self._heating_module.line_solver_attempts self._default_pahabund = self._heating_module.pahabund self._default_coolant_data_dir = self._f2py_constants_module.coolantdatadir self._default_coolant_active = self._f2py_constants_module.coolant_active.copy() # Check if coolant restart mode functions are available self._coolant_functions_available = hasattr( self._uclchemwrap, "get_coolant_restart_mode_wrap" ) if self._coolant_functions_available: self._default_coolant_restart_mode = ( self._uclchemwrap.get_coolant_restart_mode_wrap() ) else: # Functions not yet exposed, use default value self._default_coolant_restart_mode = 0 # WARM mode
[docs] def set_heating_mechanism(self, mechanism_id: int, enabled: bool = True) -> None: """Enable or disable a specific heating mechanism. For mechanisms with multiple implementations (like PHOTOELECTRIC), enabling one automatically disables others in the same group. Parameters ---------- mechanism_id : int Index of the heating mechanism (1-9). Use class constants (e.g., self.PHOTOELECTRIC['BAKES'], self.H2_FORMATION) enabled : bool True to enable, False to disable. Default: True Raises ------ ValueError If mechanism_id is not between 1 and the number of heating mechanisms. Examples -------- >>> settings = HeatingSettings() >>> # Enable Weingartner photoelectric (auto-disables Bakes) >>> settings.set_heating_mechanism(settings.PHOTOELECTRIC['WEINGARTNER'], True) >>> # Or disable H2 formation >>> settings.set_heating_mechanism(settings.H2_FORMATION, False) """ if not 1 <= mechanism_id <= self._heating_module.nheating: msg = f"mechanism_id must be between 1 and {self._heating_module.nheating}" raise ValueError(msg) # If this mechanism is part of a mutually exclusive group and we're enabling it, # disable all others in the group first if enabled and mechanism_id in self._heating_groups: group = self._heating_groups[mechanism_id] for other_id in group: if other_id != mechanism_id: self._heating_module.heating_modules[other_id - 1] = False # Now set the requested mechanism self._heating_module.heating_modules[mechanism_id - 1] = enabled
[docs] def set_coolant_active(self, coolant_name: str, enabled: bool = True) -> None: """Enable or disable a specific line coolant species. This controls individual coolant species within the molecular line cooling mechanism (cooling_modules index 5). When a coolant is disabled, its level population calculation is skipped entirely and its cooling contribution is set to zero. Parameters ---------- coolant_name : str Name of the coolant species (e.g. "CO", "C+", "p-H2"). Use get_coolant_active() to see available names. enabled : bool True to enable, False to disable. Default: True Raises ------ ValueError If coolant_name is not found in the network. Examples -------- >>> settings = HeatingSettings() >>> settings.set_coolant_active("CO", False) >>> settings.set_coolant_active("p-H2", False) """ names = [ str(np.char.decode(name)).strip() for name in self._f2py_constants_module.coolantnames ] if coolant_name not in names: msg = f"Coolant '{coolant_name}' not found. Available coolants: {names}" raise ValueError(msg) idx = names.index(coolant_name) self._f2py_constants_module.coolant_active[idx] = enabled
[docs] def get_coolant_active(self) -> dict[str, bool]: """Get the active state of all line coolant species. Returns ------- dict[str, bool] Dictionary mapping coolant names to their enabled state Examples -------- >>> settings = HeatingSettings() >>> state = settings.get_coolant_active() >>> print(state) {'H': True, 'C+': True, 'O': True, ...} """ names = [ str(np.char.decode(name)).strip() for name in self._f2py_constants_module.coolantnames ] return { name: bool(self._f2py_constants_module.coolant_active[i]) for i, name in enumerate(names) }
[docs] def set_cooling_mechanism(self, mechanism_id: int, enabled: bool = True) -> None: """Enable or disable a specific cooling mechanism. Parameters ---------- mechanism_id : int Index of the cooling mechanism (1-5). Use class constants (e.g., self.ATOMIC_LINE_COOLING) enabled : bool True to enable, False to disable. Default: True Raises ------ ValueError If mechanism_id is not between 1 and the number of cooling mechanisms. Examples -------- >>> settings = HeatingSettings() >>> settings.set_cooling_mechanism(settings.MOLECULAR_LINE_COOLING, False) """ if not 1 <= mechanism_id <= self._heating_module.ncooling: msg = f"mechanism_id must be between 1 and {self._heating_module.ncooling}" raise ValueError(msg) self._heating_module.cooling_modules[mechanism_id - 1] = enabled
[docs] def get_heating_modules(self) -> dict[str, bool]: """Obtain the state (on/off) of all heating mechanisms. Returns ------- dict[str, bool] Dictionary mapping mechanism names to their enabled state Examples -------- >>> settings = HeatingSettings() >>> state = settings.get_heating_modules() >>> print(state['H2Formation']) False """ labels = [ str(np.char.decode(label)).strip() for label in self._heating_module.heatinglabels ] return { label: bool(self._heating_module.heating_modules[i]) for i, label in enumerate(labels) }
[docs] def get_cooling_modules(self) -> dict[str, bool]: """Obtain the state (on/off) of all cooling mechanisms. Returns ------- dict[str, bool] Dictionary mapping mechanism names to their enabled state Examples -------- >>> settings = HeatingSettings() >>> state = settings.get_cooling_modules() >>> print(state['AtomicLineEmission']) True """ labels = [ str(np.char.decode(label)).strip() for label in self._heating_module.coolinglabels ] return { label: bool(self._heating_module.cooling_modules[i]) for i, label in enumerate(labels) }
[docs] def set_dust_gas_coupling_method(self, method: int) -> None: """Set the dust-gas temperature coupling method. Parameters ---------- method : int Coupling method to use: - 1 (DUST_TEMP_HOCUK): Hocuk et al. 2017 - 2 (DUST_TEMP_HOLLENBACH): Hollenbach 1991 Raises ------ ValueError If method is not one of `[1, 2]`. Examples -------- >>> settings = HeatingSettings() >>> settings.set_dust_gas_coupling_method(settings.DUST_TEMP_HOLLENBACH) """ if method not in {1, 2}: msg = "method must be 1 (Hocuk) or 2 (Hollenbach)" raise ValueError(msg) self._heating_module.dust_gas_coupling_method = method
[docs] def get_dust_gas_coupling_method(self) -> int: """Get the current dust-gas temperature coupling method. Returns ------- int Current method (1=Hocuk, 2=Hollenbach) """ return self._heating_module.dust_gas_coupling_method
[docs] def set_line_solver_attempts(self, attempts: int) -> None: """Set the number of line cooling solver iterations. The line cooling is computed multiple times and the median is used. Parameters ---------- attempts : int Number of solver attempts (odd number recommended). Default is 5. Recommended range: 3-11. Raises ------ ValueError If attempts is less than 1. Examples -------- >>> settings = HeatingSettings() >>> settings.set_line_solver_attempts(7) """ if attempts < 1: msg = "attempts must be at least 1" raise ValueError(msg) if attempts % 2 == 0: warnings.warn( "Even number of attempts may not produce proper median. " "Consider using an odd number.", stacklevel=2, ) self._heating_module.line_solver_attempts = attempts
[docs] def get_line_solver_attempts(self) -> int: """Get the current number of line cooling solver attempts. Returns ------- int Number of solver attempts """ return self._heating_module.line_solver_attempts
[docs] def set_pah_abundance(self, abundance: float) -> None: """Set the PAH (Polycyclic Aromatic Hydrocarbon) abundance. Parameters ---------- abundance : float PAH abundance relative to hydrogen. Default: 6e-7 Raises ------ ValueError if abundance is smaller than 0 Examples -------- >>> settings = HeatingSettings() >>> settings.set_pah_abundance(1e-6) """ if abundance < 0: msg = "abundance must be non-negative" raise ValueError(msg) self._heating_module.pahabund = abundance
[docs] def get_pah_abundance(self) -> float: """Get the current PAH abundance. Returns ------- float PAH abundance """ return self._heating_module.pahabund
[docs] def set_coolant_directory(self, directory: str | Path) -> None: """Set the directory containing collisional rate data files. This directory must contain the LAMDA-format collisional rate files (e.g., co.dat, o-h2.dat, etc.) used for molecular line cooling calculations. Parameters ---------- directory : str | Path Path to directory containing LAMDA-format rate files. Must end with '/'. Max 255 characters. Raises ------ ValueError If path is too long or doesn't end with '/' FileNotFoundError If directory doesn't exist Examples -------- >>> settings = HeatingSettings() >>> settings.set_coolant_directory("/custom/rates/") # doctest: +SKIP """ # Normalize to str for Fortran compatibility dir_str = str(directory) if not dir_str.endswith("/"): dir_str += "/" dir_path = Path(dir_str) if not dir_path.exists(): msg = f"Directory not found: {dir_str}" raise FileNotFoundError(msg) if not dir_path.is_dir(): msg = f"Not a directory: {dir_str}" raise ValueError(msg) if len(dir_str) > 255: # noqa: PLR2004 msg = f"Path too long (max 255): {dir_str}" raise ValueError(msg) # Pad and set current = self._f2py_constants_module.coolantdatadir max_len = int(current.dtype.itemsize) padded = dir_str.ljust(max_len) self._f2py_constants_module.coolantdatadir = padded
[docs] def get_coolant_directory(self) -> str: """Get the current collisional rate data directory. Returns ------- str Directory path as string (stripped of trailing spaces) """ return str(np.char.decode(self._f2py_constants_module.coolantdatadir)).strip()
# TODO: refactor once Fortran is exposed
[docs] def set_coolant_restart_mode(self, mode: int) -> None: """Set the coolant population restart mode. Parameters ---------- mode : int Restart mode (0=WARM, 1=FORCE_LTE, 2=FORCE_GROUND) Raises ------ ValueError If mode is not one of `[0, 1, 2]`. AssertionError If the mode could not be set in the Fortran module. Examples -------- >>> settings = HeatingSettings() >>> settings.set_coolant_restart_mode(settings.COOLANT_WARM) """ if mode not in {0, 1, 2}: msg = f"mode must be 0, 1, or 2, got {mode}" raise ValueError(msg) self._uclchemwrap.uclchemwrap.set_coolant_restart_mode_wrap(mode) if self.get_coolant_restart_mode() != mode: msg = "Failed to set coolant restart mode" raise AssertionError(msg)
# TODO: refactor once Fortran is exposed
[docs] def get_coolant_restart_mode(self) -> int: """Get the current coolant population restart mode. Returns ------- int Current restart mode (0=WARM, 1=FORCE_LTE, 2=FORCE_GROUND) """ return self._uclchemwrap.uclchemwrap.get_coolant_restart_mode_wrap()
[docs] def reset_to_defaults(self) -> None: """Reset all heating and cooling mechanisms to their initial values. Restores the configuration that was present when this HeatingSettings instance was first initialized. This allows reverting any changes made during the session back to the starting state. Examples -------- >>> settings = HeatingSettings() >>> settings.set_heating_mechanism(settings.H2_FORMATION, False) >>> settings.reset_to_defaults() # Restores original state """ # Restore backed-up initial state self._heating_module.heating_modules[:] = self._default_heating_modules self._heating_module.cooling_modules[:] = self._default_cooling_modules self._heating_module.dust_gas_coupling_method = ( self._default_dust_gas_coupling_method ) self._heating_module.line_solver_attempts = self._default_line_solver_attempts self._heating_module.pahabund = self._default_pahabund self._f2py_constants_module.coolantdatadir = self._default_coolant_data_dir self._f2py_constants_module.coolant_active[:] = self._default_coolant_active if self._coolant_functions_available: self._uclchemwrap.set_coolant_restart_mode_wrap( self._default_coolant_restart_mode )
[docs] def print_configuration(self) -> None: """Print the current heating and cooling configuration. Examples -------- >>> settings = HeatingSettings() >>> settings.print_configuration() # doctest: +SKIP ... """ print("=" * 60) print("UCLCHEM Heating & Cooling Configuration") print("=" * 60) print("\nHeating Mechanisms:") print("-" * 40) heating_state = self.get_heating_modules() for name, enabled in heating_state.items(): status = "ENABLED " if enabled else "DISABLED" print(f" {name:30s} : {status}") print("\nCooling Mechanisms:") print("-" * 40) cooling_state = self.get_cooling_modules() for name, enabled in cooling_state.items(): status = "ENABLED " if enabled else "DISABLED" print(f" {name:30s} : {status}") print("\nLine Coolants (individual species):") print("-" * 40) coolant_state = self.get_coolant_active() for name, enabled in coolant_state.items(): status = "ENABLED " if enabled else "DISABLED" print(f" {name:30s} : {status}") print("\nSolver Parameters:") print("-" * 40) method = self.get_dust_gas_coupling_method() method_name = "Hocuk et al. 2017" if method == 1 else "Hollenbach 1991" print(f" Dust-Gas Coupling Method : {method} ({method_name})") print(f" Line Solver Attempts : {self.get_line_solver_attempts()}") print(f" PAH Abundance : {self.get_pah_abundance():.2e}") print(f" Coolant Data Directory : {self.get_coolant_directory()}") restart_mode = self.get_coolant_restart_mode() restart_names = {0: "WARM", 1: "FORCE_LTE", 2: "FORCE_GROUND"} restart_name = restart_names.get(restart_mode, "UNKNOWN") print(f" Coolant Restart Mode : {restart_mode} ({restart_name})") print("=" * 60)
[docs] def initialize_coolant_directory() -> str: """Locate and return the collisional rate data directory. This function searches for coolant data files in the following order: 1. UCLCHEM_COOLANT_DATA environment variable (if set) 2. Installed package data via importlib.resources (normal installation) 3. Development mode: Makerates/data/collisional_rates/ (relative to project root) Returns ------- str Absolute path to the coolant data directory (with trailing slash) Raises ------ RuntimeError If the Fortran heating module is not available (not compiled) FileNotFoundError If coolant data directory cannot be found in any location Examples -------- >>> from uclchem.advanced.advanced_heating import initialize_coolant_directory >>> coolant_dir = initialize_coolant_directory() >>> print(f"Coolant data at: {coolant_dir}") Coolant data at: ... """ # Check if heating module is available if importlib.util.find_spec("uclchemwrap") is None: msg = ( "UCLCHEM heating module not available. " "The Fortran extension may not be compiled. " "Install UCLCHEM with: pip install ." ) raise RuntimeError(msg) # Priority 1: Environment variable env_dir = os.environ.get("UCLCHEM_COOLANT_DATA") if env_dir: env_path = Path(env_dir) if env_path.is_dir() and list(env_path.glob("*.dat")): coolant_dir = str(env_path.resolve()) if not coolant_dir.endswith("/"): coolant_dir += "/" logger.info("Using coolant data from UCLCHEM_COOLANT_DATA: %s", coolant_dir) return coolant_dir else: logger.warning( "UCLCHEM_COOLANT_DATA set to %s, but directory not found or empty. " "Searching other locations...", env_dir, ) # Priority 2: Installed package data (importlib.resources for Python 3.9+) try: # Try new API first (Python 3.9+) try: from importlib.resources import files # noqa: PLC0415 resource = files("uclchem") / "data" / "collisional_rates" # Convert Traversable to Path so we can call .glob() and .resolve() package_data_path = Path(str(resource)) except (ImportError, TypeError): # Fallback to older API (Python 3.7-3.8) from importlib.resources import path as resource_path # noqa: PLC0415 with resource_path("uclchem.data", "collisional_rates") as p: package_data_path = Path(p) if package_data_path.is_dir() and list(package_data_path.glob("*.dat")): coolant_dir = str(package_data_path.resolve()) if not coolant_dir.endswith("/"): coolant_dir += "/" logger.debug("Using installed coolant data: %s", coolant_dir) return coolant_dir except (ImportError, FileNotFoundError, AttributeError) as e: logger.debug("Installed package data not found: %s", e) # Priority 3: Development mode - search for Makerates/data/collisional_rates/ try: from uclchem.utils import UCLCHEM_ROOT_DIR # noqa: PLC0415 # Try relative to UCLCHEM_ROOT_DIR (src/uclchem/) candidates = [ UCLCHEM_ROOT_DIR.parent.parent / "Makerates" / "data" / "collisional_rates", # from src/uclchem to project root Path.cwd() / "Makerates" / "data" / "collisional_rates", # from current working directory Path.cwd().parent / "Makerates" / "data" / "collisional_rates", # one level up ] for candidate in candidates: if candidate.is_dir() and list(candidate.glob("*.dat")): coolant_dir = str(candidate.resolve()) if not coolant_dir.endswith("/"): coolant_dir += "/" logger.info("Using development mode coolant data: %s", coolant_dir) return coolant_dir except Exception as e: logger.debug("Development mode search failed: %s", e) # Not found in any location msg = ( "Could not locate coolant data files (.dat files for collisional rates). " "Searched:\n" " 1. UCLCHEM_COOLANT_DATA environment variable\n" " 2. Installed package data (uclchem/data/collisional_rates/)\n" " 3. Development mode (Makerates/data/collisional_rates/)\n" "\n" "To fix:\n" " - For installed package: Run 'python makerates.py' then 'pip install .'\n" " - For development: Ensure Makerates/data/collisional_rates/*.dat files exist\n" " - Or set UCLCHEM_COOLANT_DATA=/path/to/coolant/data/" ) raise FileNotFoundError(msg)
[docs] def auto_initialize_coolant_directory() -> bool: """Automatically initialize the coolant data directory for the Fortran module. This is a convenience wrapper around initialize_coolant_directory() that: - Attempts to locate coolant data files - Sets the coolant directory in the Fortran module if found - Logs warnings instead of raising exceptions if initialization fails This function is called automatically when the uclchem module is imported. Returns ------- bool True if initialization succeeded, False if it failed Examples -------- >>> from uclchem.advanced.advanced_heating import auto_initialize_coolant_directory >>> if auto_initialize_coolant_directory(): ... print("Coolant data initialized successfully") ... Coolant data initialized successfully """ try: coolant_dir = initialize_coolant_directory() settings = HeatingSettings() settings.set_coolant_directory(coolant_dir) logger.debug("Auto-initialized coolant directory: %s", coolant_dir) return True except RuntimeError as e: # Heating module not available - this is expected for makerates-only builds logger.debug("Coolant initialization skipped: %s", e) return False except FileNotFoundError as e: # Could not find coolant data - warn user logger.warning( "Could not auto-initialize coolant data directory: %s\n" "Heating/cooling calculations may fail. " "Run 'python makerates.py' and reinstall if needed.", e, ) return False except Exception as e: # Unexpected error - warn but don't crash logger.warning( "Unexpected error during coolant initialization: %s" "\nEnabling heating and cooling might cause errors at runtime.", e, ) return False