Module pychnosz.core.equilibrium

Chemical equilibrium solver for CHNOSZ.

This module implements equilibrium calculations including: - Activity coefficient models - Chemical speciation - Equilibrium constants - Activity-concentration relationships

Functions

def affinity(species: str | List[str] | None = None,
property: str = 'A',
T: float = 298.15,
P: float = 1.0) ‑> pandas.core.frame.DataFrame

Expand source code

def affinity(species: Optional[Union[str, List[str]]] = None,
            property: str = 'A', T: float = 298.15, P: float = 1.0) -> pd.DataFrame:
    """
    Calculate chemical affinity for formation reactions.
    
    Parameters
    ----------
    species : str, list, or None
        Species to calculate affinity for. If None, uses current thermo species.
    property : str, default 'A'  
        Property to calculate ('A' for affinity, 'logK' for log K, 'logQ' for log Q)
    T : float, default 298.15
        Temperature in Kelvin
    P : float, default 1.0
        Pressure in bar
        
    Returns
    -------
    DataFrame
        Affinities and related properties
    """
    
    # This would interface with the basis system to calculate formation reaction affinities
    # For now, return a placeholder
    
    if species is None:
        if thermo.species is None:
            raise ValueError("No species defined. Use species() function first.")
        species_list = thermo.species['name'].tolist()
    elif isinstance(species, str):
        species_list = [species]
    else:
        species_list = species
    
    results = []
    for sp in species_list:
        # Calculate formation reaction from basis species
        # This requires implementing the basis system and formation reactions
        result = {
            'species': sp,
            'T': T,
            'P': P,
            'A': 0.0,  # Placeholder - would calculate actual affinity
            'logK': 0.0,  # Placeholder
            'logQ': 0.0   # Placeholder
        }
        results.append(result)
    
    return pd.DataFrame(results)

Calculate chemical affinity for formation reactions.

Parameters

species : str, list, or None

Species to calculate affinity for. If None, uses current thermo species.

property : str, default 'A'

Property to calculate ('A' for affinity, 'logK' for log K, 'logQ' for log Q)

T : float, default 298.15

Temperature in Kelvin

P : float, default 1.0

Pressure in bar

Returns

DataFrame

Affinities and related properties

def equilibrate(aout: pandas.core.frame.DataFrame | None = None,
balance: str | int | None = None,
normalize: bool = False,
as_residue: bool = False) ‑> Dict[str, Any]

Expand source code

def equilibrate(aout: Optional[pd.DataFrame] = None, 
               balance: Optional[Union[str, int]] = None,
               normalize: bool = False,
               as_residue: bool = False) -> Dict[str, Any]:
    """
    Find chemical equilibrium using an optimization approach.
    
    Parameters
    ----------
    aout : DataFrame, optional
        Output from affinity() function
    balance : str or int, optional
        Balanced chemical component
    normalize : bool, default False
        Normalize activities to sum to 1
    as_residue : bool, default False
        Return residue of optimization
        
    Returns
    -------
    dict
        Equilibrium results
    """
    
    # Placeholder implementation
    # This would implement the equilibrium optimization using affinity calculations
    
    if aout is None:
        raise ValueError("affinity output required")
    
    # For now, return equal distribution
    n_species = len(aout)
    activities = np.ones(n_species) / n_species
    
    result = {
        'activities': activities,
        'residual': 0.0,
        'converged': True
    }
    
    return result

Find chemical equilibrium using an optimization approach.

Parameters

aout : DataFrame, optional

Output from affinity() function

balance : str or int, optional

Balanced chemical component

normalize : bool, default False

Normalize activities to sum to 1

as_residue : bool, default False

Return residue of optimization

Returns

dict

Equilibrium results

def solubility(species: str | List[str],
mineral: str,
T: float = 298.15,
P: float = 1.0,
pH_range: Tuple[float, float] | None = None) ‑> Dict[str, Any]

Expand source code

def solubility(species: Union[str, List[str]], 
              mineral: str,
              T: float = 298.15, P: float = 1.0,
              pH_range: Optional[Tuple[float, float]] = None) -> Dict[str, Any]:
    """
    Calculate mineral solubility in aqueous solution.
    
    Parameters
    ----------  
    species : str or list
        Aqueous species in equilibrium with mineral
    mineral : str
        Mineral name
    T : float, default 298.15
        Temperature in Kelvin
    P : float, default 1.0  
        Pressure in bar
    pH_range : tuple, optional
        pH range for calculation (min_pH, max_pH)
        
    Returns
    -------
    dict
        Solubility results
    """
    
    if pH_range is None:
        pH_range = (0, 14)
    
    pH_values = np.linspace(pH_range[0], pH_range[1], 100)
    
    # Calculate dissolution reaction
    # This requires implementing mineral dissolution reactions
    
    results = {
        'pH': pH_values,
        'solubility': np.zeros_like(pH_values),  # Placeholder
        'species_distribution': {}  # Placeholder
    }
    
    return results

Calculate mineral solubility in aqueous solution.

Parameters

species : str or list Aqueous species in equilibrium with mineral mineral : str Mineral name T : float, default 298.15 Temperature in Kelvin P : float, default 1.0
Pressure in bar pH_range : tuple, optional pH range for calculation (min_pH, max_pH)

Returns

dict

Solubility results

Classes

class EquilibriumSolver

Expand source code

class EquilibriumSolver:
    """
    Chemical equilibrium solver for aqueous systems.
    
    This class implements various equilibrium calculation methods:
    - Activity coefficient corrections (Debye-Hückel, B-dot, Pitzer)
    - Chemical speciation calculations
    - Reaction equilibrium constants
    - Mass balance constraints
    """
    
    def __init__(self):
        """Initialize the equilibrium solver."""
        self.activity_models = {
            'ideal': self._activity_ideal,
            'debye_huckel': self._activity_debye_huckel,
            'bdot': self._activity_bdot,
            'pitzer': self._activity_pitzer
        }
        
        # Default parameters
        self.ionic_strength_limit = 3.0  # mol/kg
        self.max_iterations = 100
        self.tolerance = 1e-8
    
    def calculate_logK(self, reaction: Dict[str, float],
                      T: Union[float, np.ndarray] = 298.15,
                      P: Union[float, np.ndarray] = 1.0) -> np.ndarray:
        """
        Calculate equilibrium constant for a reaction.
        
        Parameters
        ----------
        reaction : dict
            Reaction dictionary with species names as keys and 
            stoichiometric coefficients as values (negative for reactants)
        T : float or array, default 298.15
            Temperature in Kelvin
        P : float or array, default 1.0
            Pressure in bar
            
        Returns
        -------
        array
            log K values at given T and P
        """
        
        # Get species names and coefficients
        species_names = list(reaction.keys())
        coefficients = list(reaction.values())
        
        # Calculate standard properties
        result = subcrt(species_names, coefficients, T=T, P=P, show=False)

        if result.out is not None and 'logK' in result.out.columns:
            return result.out['logK'].values
        else:
            raise ValueError("Could not calculate reaction properties")
    
    def calculate_speciation(self, total_concentrations: Dict[str, float],
                           reactions: Dict[str, Dict[str, float]],
                           T: float = 298.15, P: float = 1.0,
                           pH: Optional[float] = None,
                           ionic_strength: Optional[float] = None,
                           activity_model: str = 'debye_huckel') -> Dict[str, Any]:
        """
        Calculate chemical speciation for an aqueous system.
        
        Parameters
        ----------
        total_concentrations : dict
            Total concentrations of components (mol/kg)
        reactions : dict
            Formation reactions for each species
        T : float, default 298.15
            Temperature in Kelvin
        P : float, default 1.0  
            Pressure in bar
        pH : float, optional
            pH constraint (if provided)
        ionic_strength : float, optional
            Ionic strength (if known, otherwise calculated)
        activity_model : str, default 'debye_huckel'
            Activity coefficient model to use
            
        Returns
        -------
        dict
            Speciation results with concentrations, activities, and properties
        """
        
        # Get equilibrium constants for all reactions
        logK_values = {}
        for species, reaction in reactions.items():
            try:
                logK = self.calculate_logK(reaction, T, P)
                logK_values[species] = logK[0] if hasattr(logK, '__len__') else logK
            except Exception as e:
                warnings.warn(f"Could not calculate logK for {species}: {e}")
                logK_values[species] = 0.0
        
        # Initial guess for species concentrations
        species_names = list(reactions.keys())
        basis_species = set()
        for reaction in reactions.values():
            basis_species.update(reaction.keys())
        basis_species = list(basis_species)
        
        # Create initial guess (equal distribution)
        n_species = len(species_names)
        n_basis = len(basis_species)
        
        if n_species == 0:
            return {'concentrations': {}, 'activities': {}, 'ionic_strength': 0.0}
        
        # Initial concentrations (log scale for stability)
        x0 = np.ones(n_species + n_basis) * (-6.0)  # log concentrations
        
        if pH is not None:
            # Find H+ index and set pH constraint
            if 'H+' in basis_species:
                h_idx = basis_species.index('H+')
                x0[n_species + h_idx] = -pH
        
        # Solve equilibrium system
        try:
            solution = fsolve(self._equilibrium_equations, x0, 
                            args=(species_names, basis_species, reactions, 
                                  logK_values, total_concentrations, pH,
                                  T, P, activity_model))
            
            if not np.allclose(self._equilibrium_equations(solution, species_names, basis_species,
                                                          reactions, logK_values, 
                                                          total_concentrations, pH,
                                                          T, P, activity_model), 0, atol=1e-6):
                warnings.warn("Equilibrium solution may not have converged")
            
        except Exception as e:
            warnings.warn(f"Equilibrium calculation failed: {e}")
            solution = x0  # Use initial guess
        
        # Extract results
        log_species_conc = solution[:n_species]
        log_basis_conc = solution[n_species:]
        
        species_conc = 10**log_species_conc
        basis_conc = 10**log_basis_conc
        
        # Calculate ionic strength
        ionic_str = self._calculate_ionic_strength(species_names, species_conc, 
                                                  basis_species, basis_conc)
        
        # Calculate activity coefficients
        gamma_species = {}
        gamma_basis = {}
        
        for i, species in enumerate(species_names):
            gamma_species[species] = self._get_activity_coefficient(
                species, ionic_str, T, P, activity_model)
        
        for i, species in enumerate(basis_species):
            gamma_basis[species] = self._get_activity_coefficient(
                species, ionic_str, T, P, activity_model)
        
        # Calculate activities
        activities_species = {species: conc * gamma_species[species] 
                            for species, conc in zip(species_names, species_conc)}
        activities_basis = {species: conc * gamma_basis[species]
                           for species, conc in zip(basis_species, basis_conc)}
        
        return {
            'concentrations': {**dict(zip(species_names, species_conc)),
                             **dict(zip(basis_species, basis_conc))},
            'activities': {**activities_species, **activities_basis},
            'activity_coefficients': {**gamma_species, **gamma_basis},
            'ionic_strength': ionic_str,
            'pH': -np.log10(basis_conc[basis_species.index('H+')]) if 'H+' in basis_species else None
        }
    
    def _equilibrium_equations(self, x: np.ndarray, species_names: List[str],
                              basis_species: List[str], reactions: Dict[str, Dict[str, float]],
                              logK_values: Dict[str, float], 
                              total_concentrations: Dict[str, float],
                              pH: Optional[float], T: float, P: float,
                              activity_model: str) -> np.ndarray:
        """
        Equilibrium equations to solve for speciation.
        
        Returns array of residuals for:
        1. Mass balance equations
        2. Equilibrium constant equations  
        3. Charge balance (if applicable)
        4. pH constraint (if provided)
        """
        
        n_species = len(species_names)
        n_basis = len(basis_species)
        
        # Extract log concentrations
        log_species_conc = x[:n_species] 
        log_basis_conc = x[n_species:]
        
        species_conc = 10**log_species_conc
        basis_conc = 10**log_basis_conc
        
        # Calculate ionic strength for activity coefficients
        ionic_str = self._calculate_ionic_strength(species_names, species_conc,
                                                  basis_species, basis_conc)
        
        equations = []
        
        # 1. Equilibrium constant equations
        for i, species in enumerate(species_names):
            if species in reactions:
                reaction = reactions[species]
                logK = logK_values[species]
                
                # log K = log(activity of products) - log(activity of reactants)
                log_activity_ratio = 0
                
                for reactant, coeff in reaction.items():
                    if reactant in basis_species:
                        idx = basis_species.index(reactant)
                        gamma = self._get_activity_coefficient(reactant, ionic_str, T, P, activity_model)
                        log_activity_ratio += coeff * (log_basis_conc[idx] + np.log10(gamma))
                
                # Activity of species being formed
                gamma_species = self._get_activity_coefficient(species, ionic_str, T, P, activity_model)
                log_species_activity = log_species_conc[i] + np.log10(gamma_species)
                
                # Equilibrium equation: logK - log_activity_ratio + log_species_activity = 0
                equations.append(logK - log_activity_ratio + log_species_activity)
        
        # 2. Mass balance equations
        for component, total_conc in total_concentrations.items():
            if total_conc <= 0:
                continue
                
            mass_balance = 0
            
            # Contribution from basis species
            if component in basis_species:
                idx = basis_species.index(component)
                mass_balance += basis_conc[idx]
            
            # Contributions from formed species
            for i, species in enumerate(species_names):
                if species in reactions and component in reactions[species]:
                    coeff = abs(reactions[species][component])  # Take absolute value for mass balance
                    mass_balance += coeff * species_conc[i]
            
            # Mass balance equation: (calculated - total) / total = 0
            equations.append((mass_balance - total_conc) / total_conc)
        
        # 3. pH constraint
        if pH is not None and 'H+' in basis_species:
            h_idx = basis_species.index('H+')
            equations.append(log_basis_conc[h_idx] + pH)  # log[H+] + pH = 0
        
        return np.array(equations)
    
    def _calculate_ionic_strength(self, species_names: List[str], species_conc: np.ndarray,
                                 basis_species: List[str], basis_conc: np.ndarray) -> float:
        """Calculate ionic strength of the solution."""
        
        ionic_strength = 0
        
        # Contributions from basis species (assume they have charges)
        for i, species in enumerate(basis_species):
            charge = self._get_species_charge(species)
            ionic_strength += 0.5 * basis_conc[i] * charge**2
        
        # Contributions from formed species  
        for i, species in enumerate(species_names):
            charge = self._get_species_charge(species)
            ionic_strength += 0.5 * species_conc[i] * charge**2
        
        return ionic_strength
    
    def _get_species_charge(self, species: str) -> int:
        """Extract charge from species name (simplified)."""
        
        if '+' in species:
            charge_str = species.split('+')[-1]
            try:
                return int(charge_str) if charge_str.isdigit() else 1
            except:
                return 1
        elif '-' in species:
            charge_str = species.split('-')[-1]
            try:
                return -int(charge_str) if charge_str.isdigit() else -1
            except:
                return -1
        else:
            return 0
    
    def _get_activity_coefficient(self, species: str, ionic_strength: float,
                                 T: float, P: float, model: str) -> float:
        """Calculate activity coefficient for a species."""
        
        if model in self.activity_models:
            return self.activity_models[model](species, ionic_strength, T, P)
        else:
            warnings.warn(f"Unknown activity model: {model}, using ideal")
            return 1.0
    
    def _activity_ideal(self, species: str, I: float, T: float, P: float) -> float:
        """Ideal activity coefficient (γ = 1)."""
        return 1.0
    
    def _activity_debye_huckel(self, species: str, I: float, T: float, P: float) -> float:
        """Debye-Hückel activity coefficient."""
        
        charge = self._get_species_charge(species)
        
        if charge == 0:
            return 1.0  # Neutral species
        
        # Debye-Hückel parameters (approximate)
        A = 0.509  # at 25°C, valid for higher T too approximately
        
        # Extended Debye-Hückel equation
        if I <= 0.1:
            log_gamma = -A * charge**2 * np.sqrt(I) / (1 + np.sqrt(I))
        else:
            # For higher ionic strengths, use extended form
            B = 0.328  # approximate
            log_gamma = -A * charge**2 * np.sqrt(I) / (1 + B * np.sqrt(I))
        
        return 10**log_gamma
    
    def _activity_bdot(self, species: str, I: float, T: float, P: float) -> float:
        """B-dot activity coefficient model."""
        
        charge = self._get_species_charge(species)
        
        if charge == 0:
            return 1.0
        
        # B-dot equation parameters (simplified)
        A = 0.509  # Debye-Hückel A parameter
        B = 0.328  # B parameter
        bdot = 0.041  # B-dot parameter (approximate)
        
        sqrt_I = np.sqrt(I)
        log_gamma = -A * charge**2 * sqrt_I / (1 + B * sqrt_I) + bdot * I
        
        return 10**log_gamma
    
    def _activity_pitzer(self, species: str, I: float, T: float, P: float) -> float:
        """Pitzer activity coefficient model (simplified)."""
        
        # For now, fall back to B-dot model
        return self._activity_bdot(species, I, T, P)

Chemical equilibrium solver for aqueous systems.

This class implements various equilibrium calculation methods: - Activity coefficient corrections (Debye-Hückel, B-dot, Pitzer) - Chemical speciation calculations - Reaction equilibrium constants - Mass balance constraints

Initialize the equilibrium solver.

Methods

def calculate_logK(self,
reaction: Dict[str, float],
T: float | numpy.ndarray = 298.15,
P: float | numpy.ndarray = 1.0) ‑> numpy.ndarray

Expand source code

def calculate_logK(self, reaction: Dict[str, float],
                  T: Union[float, np.ndarray] = 298.15,
                  P: Union[float, np.ndarray] = 1.0) -> np.ndarray:
    """
    Calculate equilibrium constant for a reaction.
    
    Parameters
    ----------
    reaction : dict
        Reaction dictionary with species names as keys and 
        stoichiometric coefficients as values (negative for reactants)
    T : float or array, default 298.15
        Temperature in Kelvin
    P : float or array, default 1.0
        Pressure in bar
        
    Returns
    -------
    array
        log K values at given T and P
    """
    
    # Get species names and coefficients
    species_names = list(reaction.keys())
    coefficients = list(reaction.values())
    
    # Calculate standard properties
    result = subcrt(species_names, coefficients, T=T, P=P, show=False)

    if result.out is not None and 'logK' in result.out.columns:
        return result.out['logK'].values
    else:
        raise ValueError("Could not calculate reaction properties")

Calculate equilibrium constant for a reaction.

Parameters

reaction : dict

Reaction dictionary with species names as keys and stoichiometric coefficients as values (negative for reactants)

T : float or array, default 298.15

Temperature in Kelvin

P : float or array, default 1.0

Pressure in bar

Returns

array

log K values at given T and P

def calculate_speciation(self,
total_concentrations: Dict[str, float],
reactions: Dict[str, Dict[str, float]],
T: float = 298.15,
P: float = 1.0,
pH: float | None = None,
ionic_strength: float | None = None,
activity_model: str = 'debye_huckel') ‑> Dict[str, Any]

Expand source code

def calculate_speciation(self, total_concentrations: Dict[str, float],
                       reactions: Dict[str, Dict[str, float]],
                       T: float = 298.15, P: float = 1.0,
                       pH: Optional[float] = None,
                       ionic_strength: Optional[float] = None,
                       activity_model: str = 'debye_huckel') -> Dict[str, Any]:
    """
    Calculate chemical speciation for an aqueous system.
    
    Parameters
    ----------
    total_concentrations : dict
        Total concentrations of components (mol/kg)
    reactions : dict
        Formation reactions for each species
    T : float, default 298.15
        Temperature in Kelvin
    P : float, default 1.0  
        Pressure in bar
    pH : float, optional
        pH constraint (if provided)
    ionic_strength : float, optional
        Ionic strength (if known, otherwise calculated)
    activity_model : str, default 'debye_huckel'
        Activity coefficient model to use
        
    Returns
    -------
    dict
        Speciation results with concentrations, activities, and properties
    """
    
    # Get equilibrium constants for all reactions
    logK_values = {}
    for species, reaction in reactions.items():
        try:
            logK = self.calculate_logK(reaction, T, P)
            logK_values[species] = logK[0] if hasattr(logK, '__len__') else logK
        except Exception as e:
            warnings.warn(f"Could not calculate logK for {species}: {e}")
            logK_values[species] = 0.0
    
    # Initial guess for species concentrations
    species_names = list(reactions.keys())
    basis_species = set()
    for reaction in reactions.values():
        basis_species.update(reaction.keys())
    basis_species = list(basis_species)
    
    # Create initial guess (equal distribution)
    n_species = len(species_names)
    n_basis = len(basis_species)
    
    if n_species == 0:
        return {'concentrations': {}, 'activities': {}, 'ionic_strength': 0.0}
    
    # Initial concentrations (log scale for stability)
    x0 = np.ones(n_species + n_basis) * (-6.0)  # log concentrations
    
    if pH is not None:
        # Find H+ index and set pH constraint
        if 'H+' in basis_species:
            h_idx = basis_species.index('H+')
            x0[n_species + h_idx] = -pH
    
    # Solve equilibrium system
    try:
        solution = fsolve(self._equilibrium_equations, x0, 
                        args=(species_names, basis_species, reactions, 
                              logK_values, total_concentrations, pH,
                              T, P, activity_model))
        
        if not np.allclose(self._equilibrium_equations(solution, species_names, basis_species,
                                                      reactions, logK_values, 
                                                      total_concentrations, pH,
                                                      T, P, activity_model), 0, atol=1e-6):
            warnings.warn("Equilibrium solution may not have converged")
        
    except Exception as e:
        warnings.warn(f"Equilibrium calculation failed: {e}")
        solution = x0  # Use initial guess
    
    # Extract results
    log_species_conc = solution[:n_species]
    log_basis_conc = solution[n_species:]
    
    species_conc = 10**log_species_conc
    basis_conc = 10**log_basis_conc
    
    # Calculate ionic strength
    ionic_str = self._calculate_ionic_strength(species_names, species_conc, 
                                              basis_species, basis_conc)
    
    # Calculate activity coefficients
    gamma_species = {}
    gamma_basis = {}
    
    for i, species in enumerate(species_names):
        gamma_species[species] = self._get_activity_coefficient(
            species, ionic_str, T, P, activity_model)
    
    for i, species in enumerate(basis_species):
        gamma_basis[species] = self._get_activity_coefficient(
            species, ionic_str, T, P, activity_model)
    
    # Calculate activities
    activities_species = {species: conc * gamma_species[species] 
                        for species, conc in zip(species_names, species_conc)}
    activities_basis = {species: conc * gamma_basis[species]
                       for species, conc in zip(basis_species, basis_conc)}
    
    return {
        'concentrations': {**dict(zip(species_names, species_conc)),
                         **dict(zip(basis_species, basis_conc))},
        'activities': {**activities_species, **activities_basis},
        'activity_coefficients': {**gamma_species, **gamma_basis},
        'ionic_strength': ionic_str,
        'pH': -np.log10(basis_conc[basis_species.index('H+')]) if 'H+' in basis_species else None
    }

Calculate chemical speciation for an aqueous system.

Parameters

total_concentrations : dict

Total concentrations of components (mol/kg)

reactions : dict

Formation reactions for each species

T : float, default 298.15

Temperature in Kelvin

P : float, default 1.0

Pressure in bar

pH : float, optional

pH constraint (if provided)

ionic_strength : float, optional

Ionic strength (if known, otherwise calculated)

activity_model : str, default 'debye_huckel'

Activity coefficient model to use

Returns

dict

Speciation results with concentrations, activities, and properties