Module `pychnosz.core.balance`

Reaction balancing utilities.

This module provides functions to balance chemical reactions using basis species, without calculating thermodynamic properties. This is more efficient than subcrt() when only the reaction stoichiometry is needed.

Functions

Expand source code

def balance_reaction(species: Union[str, List[str], int, List[int]],
                    coeff: Union[int, float, List[Union[int, float]]],
                    state: Optional[Union[str, List[str]]] = None,
                    basis: Optional[pd.DataFrame] = None,
                    messages: bool = False) -> Optional[Tuple[List, List]]:
    """
    Balance a chemical reaction using basis species.

    This function checks if a reaction is balanced and, if not, attempts to
    balance it by adding basis species. Unlike subcrt(), this function only
    performs the balancing calculation without computing thermodynamic properties,
    making it much more efficient for reaction generation.

    Parameters
    ----------
    species : str, int, list of str, or list of int
        Species names or indices in the reaction
    coeff : int, float, or list
        Stoichiometric coefficients for the species
    state : str, list of str, or None
        Physical states for species (optional)
    basis : pd.DataFrame, optional
        Basis species definition to use. If None, uses global basis from thermo()
    messages : bool
        Whether to print informational messages

    Returns
    -------
    tuple or None
        If reaction is balanced or can be balanced:
            (balanced_species, balanced_coeffs) where both are lists
        If reaction cannot be balanced:
            None

    Examples
    --------
    >>> import chnosz
    >>> pychnosz.reset()
    >>> pychnosz.basis(['H2O', 'H+', 'Fe+2'])
    >>> # Balance reaction for Fe+3
    >>> species, coeffs = balance_reaction('Fe+3', [-1])
    >>> print(f"Species: {species}")
    >>> print(f"Coefficients: {coeffs}")
    """

    # Convert inputs to lists
    if not isinstance(species, list):
        species = [species]
    if not isinstance(coeff, list):
        coeff = [coeff]
    if state is not None and not isinstance(state, list):
        state = [state]

    # Validate lengths
    if len(species) != len(coeff):
        raise ValueError("Length of species and coeff must match")

    # Get basis definition
    thermo_sys = thermo()
    if basis is None:
        if hasattr(thermo_sys, 'basis') and thermo_sys.basis is not None:
            basis = thermo_sys.basis
        else:
            raise RuntimeError("Basis species not defined. Call pychnosz.basis() first.")

    # Look up species indices
    ispecies = []
    for i, sp in enumerate(species):
        if isinstance(sp, (int, np.integer)):
            ispecies.append(int(sp))
        else:
            sp_state = state[i] if state and i < len(state) else None
            sp_idx = info(sp, sp_state, messages=messages)
            if sp_idx is None or (isinstance(sp_idx, float) and np.isnan(sp_idx)):
                raise ValueError(f"Species not found: {sp}")
            ispecies.append(sp_idx)

    # Calculate mass balance
    try:
        mass_balance = makeup(ispecies, coeff, sum_formulas=True)

        # Check if balanced
        tolerance = 1e-6
        unbalanced_elements = {elem: val for elem, val in mass_balance.items()
                             if abs(val) > tolerance}

        if not unbalanced_elements:
            # Already balanced
            if messages:
                print("Reaction is already balanced")
            return (species, coeff)

        # Reaction is unbalanced - try to balance using basis species
        missing_composition = {elem: -val for elem, val in unbalanced_elements.items()}

        if messages:
            print("Reaction is not balanced; missing composition:")
            elem_names = list(missing_composition.keys())
            elem_values = list(missing_composition.values())
            print(" ".join(elem_names))
            print(" ".join([f"{val:.4f}" for val in elem_values]))

        # Get basis element columns
        basis_elements = [col for col in basis.columns
                        if col not in ['ispecies', 'logact', 'state']]

        # Check if all missing elements are in basis
        missing_elements = set(missing_composition.keys())
        if not missing_elements.issubset(set(basis_elements)):
            if messages:
                print(f"Cannot balance: elements {missing_elements - set(basis_elements)} not in basis")
            return None

        # Calculate coefficients for missing composition from basis species
        missing_matrix = np.zeros((1, len(basis_elements)))
        for i, elem in enumerate(basis_elements):
            missing_matrix[0, i] = missing_composition.get(elem, 0)

        # Get basis matrix
        basis_matrix = basis[basis_elements].values.T  # Transpose: (elements × basis_species)

        try:
            # Try to find simple integer solutions first
            basis_coeffs = _find_simple_integer_solution(
                basis_matrix.T,
                missing_matrix.flatten(),
                basis['ispecies'].tolist(),
                missing_composition
            )

            if basis_coeffs is None:
                # Fall back to linear algebra solution
                basis_coeffs = np.linalg.solve(basis_matrix, missing_matrix.T).flatten()

                # Apply zapsmall equivalent (digits=7)
                basis_coeffs = np.around(basis_coeffs, decimals=7)

                # Clean up very small numbers
                basis_coeffs[np.abs(basis_coeffs) < 1e-7] = 0

            # Get non-zero coefficients and corresponding basis species
            nonzero_indices = np.abs(basis_coeffs) > 1e-6
            if not np.any(nonzero_indices):
                if messages:
                    print("No basis species needed to balance (coefficients are zero)")
                return (species, coeff)

            # Get basis species info
            basis_indices = basis['ispecies'].values[nonzero_indices]
            basis_coeffs_nz = basis_coeffs[nonzero_indices]

            # Create new species list and coefficients
            new_species = list(species) + [int(idx) for idx in basis_indices]
            new_coeff = list(coeff) + list(basis_coeffs_nz)

            if messages:
                print("Balanced reaction by adding basis species:")
                for sp_idx, cf in zip(basis_indices, basis_coeffs_nz):
                    sp_name = thermo_sys.obigt.loc[int(sp_idx)]['name']
                    print(f"  {cf:.4f} {sp_name}")

            # CRITICAL: Consolidate duplicate species by summing coefficients
            # This prevents infinite recursion and matches subcrt's behavior
            consolidated_species = []
            consolidated_coeffs = []

            # Convert all species to indices for consolidation
            species_indices = []
            for sp in new_species:
                if isinstance(sp, (int, np.integer)):
                    species_indices.append(int(sp))
                else:
                    sp_idx = info(sp, None, messages=False)
                    if sp_idx is None or (isinstance(sp_idx, float) and np.isnan(sp_idx)):
                        # Keep as string if not found
                        species_indices.append(sp)
                    else:
                        species_indices.append(sp_idx)

            # Group by species index and sum coefficients
            species_coeff_map = {}
            for sp_idx, coeff in zip(species_indices, new_coeff):
                if sp_idx in species_coeff_map:
                    species_coeff_map[sp_idx] += coeff
                else:
                    species_coeff_map[sp_idx] = coeff

            # Remove species with zero coefficient (cancelled out)
            for sp_idx, coeff in species_coeff_map.items():
                if abs(coeff) > tolerance:
                    consolidated_species.append(sp_idx)
                    consolidated_coeffs.append(coeff)

            # Now check if consolidated reaction is balanced
            # If not, recursively balance again
            try:
                final_mass_balance = makeup(consolidated_species, consolidated_coeffs, sum_formulas=True)
                final_unbalanced = {elem: val for elem, val in final_mass_balance.items()
                                   if abs(val) > tolerance}

                if final_unbalanced:
                    # Still unbalanced after consolidation - recursively balance
                    if messages:
                        print(f"After consolidation, reaction still unbalanced: {final_unbalanced}")
                        print(f"Attempting recursive balance...")
                    return balance_reaction(consolidated_species, consolidated_coeffs, state=None,
                                           basis=basis, messages=messages)
                else:
                    # Balanced! Return consolidated result
                    if messages:
                        print(f"Reaction balanced after consolidation")
                    return (consolidated_species, consolidated_coeffs)

            except Exception as e:
                # If check fails, return consolidated result anyway
                if messages:
                    print(f"Could not verify final balance: {e}")
                return (consolidated_species, consolidated_coeffs)

        except np.linalg.LinAlgError:
            if messages:
                print("Cannot balance: singular basis matrix")
            return None

    except Exception as e:
        if messages:
            print(f"Error checking reaction balance: {e}")
            import traceback
            traceback.print_exc()
        return None

Balance a chemical reaction using basis species.

This function checks if a reaction is balanced and, if not, attempts to balance it by adding basis species. Unlike subcrt(), this function only performs the balancing calculation without computing thermodynamic properties, making it much more efficient for reaction generation.

Parameters

species : str, int, list of str, or list of int: Species names or indices in the reaction
coeff : int, float, or list: Stoichiometric coefficients for the species
state : str, list of str, or None: Physical states for species (optional)
basis : pd.DataFrame, optional: Basis species definition to use. If None, uses global basis from thermo()
messages : bool: Whether to print informational messages

Returns

tuple or None: If reaction is balanced or can be balanced: (balanced_species, balanced_coeffs) where both are lists If reaction cannot be balanced: None

Examples

>>> import chnosz
>>> pychnosz.reset()
>>> pychnosz.basis(['H2O', 'H+', 'Fe+2'])
>>> # Balance reaction for Fe+3
>>> species, coeffs = balance_reaction('Fe+3', [-1])
>>> print(f"Species: {species}")
>>> print(f"Coefficients: {coeffs}")

def format_reaction(species: List[str | int], coeffs: List[float]) ‑> str

Expand source code

def format_reaction(species: List[Union[str, int]], coeffs: List[float]) -> str:
    """
    Format a reaction as a string for EQ3/6 input.

    Parameters
    ----------
    species : list
        Species names or indices
    coeffs : list
        Stoichiometric coefficients

    Returns
    -------
    str
        Formatted reaction string like "-1.0000 Fe+3 1.0000 Fe+2 0.2500 O2(g)"
    """
    thermo_sys = thermo()
    parts = []

    for sp, coeff in zip(species, coeffs):
        # Get species name if we have an index
        if isinstance(sp, (int, np.integer)):
            sp_name = thermo_sys.obigt.loc[int(sp)]['name']
        else:
            sp_name = sp

        # Replace 'water' with 'H2O' for EQ3 compatibility
        if sp_name == 'water':
            sp_name = 'H2O'

        parts.append(f"{coeff:.4f}")
        parts.append(sp_name)

    return " ".join(parts)

Format a reaction as a string for EQ3/6 input.

Parameters

species : list: Species names or indices
coeffs : list: Stoichiometric coefficients

Returns

str: Formatted reaction string like "-1.0000 Fe+3 1.0000 Fe+2 0.2500 O2(g)"