Source code for GUIBRUSHR.Retrieval.ModelCalculation.Classes.SpeciesMolec

"""
Module containing the SpeciesMolec class for handling molecular species data.

This module defines the SpeciesMolec class which manages various molecular
species properties including opacity, isotopes, condensed molecules, and
composition calculations for atmospheric modeling.
"""
import re

from numpy import array, append

from GUIBRUSHR.General_Constants.FunctionsAndConstants.Constant_Variables import ConstantVariables




class SpeciesMolec:
    """
    Class for managing molecular species data and composition.
    
    This class handles the initialization and management of various molecular
    species properties including continuum opacity, line species, isotopes,
    condensed molecules, and Rayleigh scattering species. It also manages
    the composition calculation based on initial parameters or manual input.
    """
    


    def __init__(
        self,
        continum_opacity,
        line_species,
        line_species_isotope,
        line_species_complete_name_hr,
        line_species_complete_name_lr,
        list_condensed_molecules,
        rayleigh_species,
        mass_vector,
        start_molecs,
        initial_params=None,
        manual_model_composition=None,
    ):
        """
        Initialize the SpeciesMolec instance.
        
        Args:
            continum_opacity: Continuum opacity data for the species
            line_species: Line species data
            line_species_isotope: Isotope information for line species
            line_species_complete_name_hr: Complete HR names of line species
            line_species_complete_name_lr: Complete LR names of line species
            list_condensed_molecules: List of condensed molecular species
            rayleigh_species: Rayleigh scattering species data
            mass_vector: Vector containing mass information
            start_molecs: Starting index for molecular parameters
            initial_params: Initial parameter array (default: None)
            manual_model_composition: Manual composition override (default: None)
        """
        # Handle default mutable argument
        if initial_params is None:
            initial_params = []
        
        # Store input parameters as instance attributes
        self.continum_opacity = continum_opacity
        self.line_species = line_species
        self.line_species_isotope = line_species_isotope
        self.line_species_complete_name_hr = line_species_complete_name_hr
        self.line_species_complete_name_lr = line_species_complete_name_lr
        self.list_condensed_molecules = list_condensed_molecules
        self.rayleigh_species = rayleigh_species
        self.mass_vector = mass_vector
        self.start_molecs = start_molecs
        
        # Configure molecular parameter handling
        # Skip H2, He (and potentially H-, H, e- based on comments)
        self.skip_molec_param_array = 2
        
        # Check if H- (hydrogen minus) is included in continuum opacity
        self.include_h_m = "H-" in continum_opacity
        
        # Initialize composition based on available data
        self._initialize_composition(initial_params, manual_model_composition)

    
    def _initialize_composition(self, initial_params, manual_model_composition):
        """
        Initialize the molecular composition.
        
        This method sets up the composition dictionary either from initial
        parameters or from a manually provided composition.
        
        Args:
            initial_params: Array of initial parameters
            manual_model_composition: Manual composition override
        """
        if manual_model_composition is None:
            # Extract molecular and condensate arrays from initial parameters
            array_molecules = array(
                initial_params[self.start_molecs:ConstantVariables.start_elements]
            )
            array_condensate = array(
                initial_params[ConstantVariables.start_condensed_molecs:]
            )
            
            # Create composition dictionary from parameter objects
            # Note: This won't be defined for manual_model scenarios
            self.composition = {
                param.name: param.get_starting_value()
                for param in append(array_molecules, array_condensate)
                if param is not None
            }
        else:
            # Use manually provided composition
            self.composition = manual_model_composition