pychnosz 1.1.4__cp311-cp311-win_amd64.whl

pychnosz/core/equilibrate.py

@@ -0,0 +1,813 @@
+"""
+Equilibrate module for calculating equilibrium activities of species.
+
+This module provides Python equivalents of the R functions in equilibrate.R:
+- equilibrate(): Calculate equilibrium activities from chemical affinities
+- equil.boltzmann(): Boltzmann distribution method
+- equil.reaction(): Reaction-based equilibration method
+- balance(): Determine balancing coefficients
+- Supporting utilities for species equilibration
+
+Author: CHNOSZ Python port
+"""
+
+import numpy as np
+import pandas as pd
+from typing import Union, List, Optional, Dict, Any, Tuple
+import warnings
+from scipy.optimize import brentq
+
+from .thermo import thermo
+from .info import info
+
+
+def equilibrate(aout: Dict[str, Any],
+                balance: Optional[Union[str, int, List[float]]] = None,
+                loga_balance: Optional[Union[float, List[float]]] = None,
+                ispecies: Optional[Union[List[int], List[bool]]] = None,
+                normalize: Union[bool, List[bool]] = False,
+                as_residue: bool = False,
+                method: Optional[Union[str, List[str]]] = None,
+                tol: float = np.finfo(float).eps ** 0.25,
+                messages: bool = True) -> Dict[str, Any]:
+    """
+    Calculate equilibrium activities of species from affinities.
+
+    This function calculates the equilibrium activities of species in
+    (metastable) equilibrium from the affinities of their formation reactions
+    from basis species at given activities.
+
+    Parameters
+    ----------
+    aout : dict
+        Output from affinity() containing chemical affinities
+    balance : str, int, or list of float, optional
+        Balancing method:
+        - None: Autoselect using which_balance()
+        - str: Name of basis species to balance on
+        - "length": Balance on protein length (for proteins)
+        - "volume": Balance on standard-state volume
+        - 1: Balance on one mole of species (formula units)
+        - list: User-defined balancing coefficients
+    loga_balance : float or list of float, optional
+        Logarithm of total activity of the balancing basis species
+        If None, calculated from species initial activities and n.balance
+    ispecies : list of int or list of bool, optional
+        Indices or boolean mask of species to include in equilibration
+        Default: all species except those with state "cr" (crystalline)
+    normalize : bool or list of bool, default False
+        Normalize formulas by balancing coefficients?
+    as_residue : bool, default False
+        Use residue basis for proteins?
+    method : str or list of str, optional
+        Equilibration method:
+        - "boltzmann": Boltzmann distribution (for n.balance = 1)
+        - "reaction": Reaction-based equilibration (general method)
+        If None, chooses "boltzmann" if all n.balance == 1, else "reaction"
+    tol : float, default np.finfo(float).eps**0.25
+        Tolerance for root-finding in reaction method
+    messages : bool, default True
+        Whether to print informational messages
+
+    Returns
+    -------
+    dict
+        Dictionary containing all aout contents plus:
+        - balance : str or list, Balancing description
+        - m_balance : list, Molar formula divisors
+        - n_balance : list, Balancing coefficients
+        - loga_balance : float or array, Log activity of balanced quantity
+        - Astar : list of arrays, Normalized affinities
+        - loga_equil : list of arrays, Equilibrium log activities
+
+    Examples
+    --------
+    >>> import pychnosz
+    >>> pychnosz.basis("CHNOS")
+    >>> pychnosz.basis("NH3", -2)
+    >>> pychnosz.species(["alanine", "glycine", "serine"])
+    >>> a = pychnosz.affinity(NH3=[-80, 60], T=55, P=2000)
+    >>> e = pychnosz.equilibrate(a, balance="CO2")
+
+    Notes
+    -----
+    This is a 1:1 replica of the R CHNOSZ equilibrate() function.
+    - Handles both Boltzmann and reaction-based equilibration
+    - Supports normalization and residue basis for proteins
+    - Properly handles crystalline species via predominance diagrams
+    - Implements identical balancing logic to R version
+    """
+
+    # Handle mosaic output (not implemented yet, but keep structure)
+    if aout.get('fun') == 'mosaic':
+        raise NotImplementedError("mosaic equilibration not yet implemented")
+
+    # Number of possible species
+    # affinity() returns values as a dict with ispecies as keys
+    if isinstance(aout['values'], dict):
+        # Convert dict to list ordered by species dataframe
+        values_list = []
+        for i in range(len(aout['species'])):
+            species_idx = aout['species']['ispecies'].iloc[i]
+            if species_idx in aout['values']:
+                values_list.append(aout['values'][species_idx])
+            else:
+                # Species not in values dict - use NaN array
+                values_list.append(np.array([np.nan]))
+        aout['values'] = values_list
+
+    nspecies = len(aout['values'])
+
+    # Get the balancing coefficients
+    bout = _balance(aout, balance, messages)
+    n_balance_orig = bout['n_balance'].copy()
+    n_balance = bout['n_balance'].copy()
+    balance = bout['balance']
+
+    # If solids (cr) species are present, find them on a predominance diagram
+    iscr = [('cr' in str(state)) for state in aout['species']['state']]
+    ncr = sum(iscr)
+
+    # Set default ispecies to exclude cr species (matching R default)
+    if ispecies is None:
+        ispecies = [not is_cr for is_cr in iscr]
+
+    if ncr > 0:
+        # Import diagram here to avoid circular imports
+        from .diagram import diagram
+        dout = diagram(aout, balance=balance, normalize=normalize,
+                      as_residue=as_residue, plot_it=False, limit_water=False, messages=messages)
+
+    if ncr == nspecies:
+        # We get here if there are only solids
+        m_balance = None
+        Astar = None
+        loga_equil = []
+        for i in range(len(aout['values'])):
+            la = np.array(aout['values'][i], copy=True)
+            la[:] = np.nan
+            loga_equil.append(la)
+    else:
+        # We get here if there are any aqueous species
+        # Take selected species in 'ispecies'
+        if len(ispecies) == 0:
+            raise ValueError("the length of ispecies is zero")
+
+        # Convert boolean to indices if needed
+        if isinstance(ispecies, list) and len(ispecies) > 0:
+            if isinstance(ispecies[0], bool):
+                ispecies = [i for i, x in enumerate(ispecies) if x]
+
+        # Take out species that have NA affinities
+        ina = [all(np.isnan(np.array(x).flatten())) for x in aout['values']]
+        ispecies = [i for i in ispecies if not ina[i]]
+
+        if len(ispecies) == 0:
+            raise ValueError("all species have NA affinities")
+
+        if ispecies != list(range(nspecies)):
+            if messages:
+                print(f"equilibrate: using {len(ispecies)} of {nspecies} species")
+            aout_species_df = aout['species']
+            aout['species'] = aout_species_df.iloc[ispecies].reset_index(drop=True)
+            aout['values'] = [aout['values'][i] for i in ispecies]
+            n_balance = [n_balance[i] for i in ispecies]
+
+        # Number of species that are left
+        nspecies = len(aout['values'])
+
+        # Say what the balancing coefficients are
+        if len(n_balance) < 100:
+            if messages:
+                print(f"equilibrate: n.balance is {', '.join(map(str, n_balance))}")
+
+        # Logarithm of total activity of the balancing basis species
+        if loga_balance is None:
+            # Sum up the activities, then take absolute value
+            # in case n.balance is negative
+            logact = np.array([aout['species']['logact'].iloc[i] for i in range(len(aout['species']))])
+            sumact = abs(sum(10**logact * n_balance))
+            loga_balance = np.log10(sumact)
+
+        # Make loga.balance the same length as the values of affinity
+        if isinstance(loga_balance, (int, float)):
+            loga_balance = float(loga_balance)
+        else:
+            loga_balance = np.array(loga_balance).flatten()
+
+        nvalues = len(np.array(aout['values'][0]).flatten())
+
+        if isinstance(loga_balance, float) or len(np.atleast_1d(loga_balance)) == 1:
+            # We have a constant loga.balance
+            if isinstance(loga_balance, np.ndarray):
+                loga_balance = float(loga_balance[0])
+            if messages:
+                print(f"equilibrate: loga.balance is {loga_balance}")
+            loga_balance = np.full(nvalues, loga_balance)
+        else:
+            # We are using a variable loga.balance (supplied by the user)
+            if len(loga_balance) != nvalues:
+                raise ValueError(f"length of loga.balance ({len(loga_balance)}) doesn't match "
+                               f"the affinity values ({nvalues})")
+            if messages:
+                print(f"equilibrate: loga.balance has same length as affinity values ({len(loga_balance)})")
+
+        # Normalize the molar formula by the balance coefficients
+        m_balance = n_balance.copy()
+        isprotein = ['_' in str(name) for name in aout['species']['name']]
+
+        # Handle normalize parameter
+        if isinstance(normalize, bool):
+            normalize = [normalize] * nspecies
+        elif not isinstance(normalize, list):
+            normalize = list(normalize)
+
+        if any(normalize) or as_residue:
+            if any(n < 0 for n in n_balance):
+                raise ValueError("one or more negative balancing coefficients prohibit using normalized molar formulas")
+
+            for i in range(nspecies):
+                if normalize[i] or as_residue:
+                    n_balance[i] = 1
+
+            if as_residue:
+                if messages:
+                    print("equilibrate: using 'as.residue' for molar formulas")
+            else:
+                if messages:
+                    print("equilibrate: using 'normalize' for molar formulas")
+
+            # Set the formula divisor (m.balance) to 1 for species whose formulas are *not* normalized
+            m_balance = [m_balance[i] if (normalize[i] or as_residue) else 1
+                        for i in range(nspecies)]
+        else:
+            m_balance = [1] * nspecies
+
+        # Astar: the affinities/2.303RT of formation reactions with
+        # formed species in their standard-state activities
+        Astar = []
+        for i in range(nspecies):
+            # 'starve' the affinity of the activity of the species,
+            # and normalize the value by the molar ratio
+            logact_i = aout['species']['logact'].iloc[i]
+            astar_i = (np.array(aout['values'][i]) + logact_i) / m_balance[i]
+            Astar.append(astar_i)
+
+        # Choose a method and compute the equilibrium activities of species
+        if method is None:
+            if all(n == 1 for n in n_balance):
+                method = ["boltzmann"]
+            else:
+                method = ["reaction"]
+        elif isinstance(method, str):
+            method = [method]
+
+        if messages:
+            print(f"equilibrate: using {method[0]} method")
+
+        if method[0] == "boltzmann":
+            loga_equil = equil_boltzmann(Astar, n_balance, loga_balance)
+        elif method[0] == "reaction":
+            loga_equil = equil_reaction(Astar, n_balance, loga_balance, tol)
+        else:
+            raise ValueError(f"unknown method: {method[0]}")
+
+        # If we normalized the formulas, get back to activities of species
+        if any(normalize) and not as_residue:
+            loga_equil = [loga_equil[i] - np.log10(m_balance[i])
+                         for i in range(nspecies)]
+
+    # Process cr species
+    if ncr > 0:
+        # cr species were excluded from equilibrium calculation,
+        # so get values back to original lengths
+        norig = len(dout['values'])
+        n_balance = n_balance_orig
+
+        # Ensure ispecies is in index form (not boolean)
+        # When ncr == nspecies, ispecies was never converted from boolean to indices
+        if isinstance(ispecies, list) and len(ispecies) > 0:
+            if isinstance(ispecies[0], bool):
+                ispecies = [i for i, x in enumerate(ispecies) if x]
+
+        # Match indices back to original
+        imatch = [None] * norig
+        for j, orig_idx in enumerate(range(norig)):
+            if orig_idx in ispecies:
+                imatch[orig_idx] = ispecies.index(orig_idx)
+
+        # Handle None values (when ncr == nspecies, these are set to None)
+        # In R, indexing NULL returns NULL, so we need to check for None in Python
+        if m_balance is not None:
+            m_balance = [m_balance[imatch[i]] if imatch[i] is not None else None
+                        for i in range(norig)]
+        if Astar is not None:
+            Astar = [Astar[imatch[i]] if imatch[i] is not None else None
+                    for i in range(norig)]
+
+        # Get a template from first loga_equil to determine shape
+        loga_equil1 = loga_equil[0]
+        loga_equil_orig = [None] * norig
+
+        for i in range(norig):
+            if imatch[i] is not None:
+                loga_equil_orig[i] = loga_equil[imatch[i]]
+
+        # Replace None loga_equil with -999 for cr-only species (will be set to 0 where predominant)
+        # Use np.full with shape, not full_like, to avoid inheriting NaN values
+        ina = [i for i in range(norig) if imatch[i] is None]
+        for i in ina:
+            loga_equil_orig[i] = np.full(loga_equil1.shape, -999.0)
+        loga_equil = loga_equil_orig
+        aout['species'] = dout['species']
+        aout['values'] = dout['values']
+
+        # Find the grid points where any cr species is predominant
+        icr = [i for i in range(len(dout['species']))
+               if 'cr' in str(dout['species']['state'].iloc[i])]
+
+        # predominant uses 1-based R indexing (1, 2, 3, ...), convert to 0-based for Python
+        predominant = dout['predominant']
+        iscr_mask = np.zeros_like(predominant, dtype=bool)
+        for icr_idx in icr:
+            # Compare with icr_idx + 1 because predominant is 1-based
+            iscr_mask |= (predominant == icr_idx + 1)
+
+        # At those grid points, make the aqueous species' activities practically zero
+        for i in range(norig):
+            if i not in icr:
+                loga_equil[i] = np.array(loga_equil[i], copy=True)
+                loga_equil[i][iscr_mask] = -999
+
+        # At the grid points where cr species predominate, set their loga_equil to 0 (standard state)
+        for i in icr:
+            # Compare with i + 1 because predominant is 1-based
+            ispredom = (predominant == i + 1)
+            loga_equil[i] = np.array(loga_equil[i], copy=True)
+            # Set to standard state activity (logact, typically 0) where predominant
+            loga_equil[i][ispredom] = dout['species']['logact'].iloc[i]
+
+    # Put together the output
+    out = aout.copy()
+    out['fun'] = 'equilibrate'  # Mark this as equilibrate output
+    out['balance'] = balance
+    out['m_balance'] = m_balance
+    out['n_balance'] = n_balance
+    out['loga_balance'] = loga_balance
+    out['Astar'] = Astar
+    out['loga_equil'] = loga_equil
+
+    return out
+
+
+def equil_boltzmann(Astar: List[np.ndarray],
+                   n_balance: List[float],
+                   loga_balance: np.ndarray) -> List[np.ndarray]:
+    """
+    Calculate equilibrium activities using Boltzmann distribution.
+
+    This method works using the Boltzmann distribution:
+    A/At = e^(Astar/n.balance) / sum(e^(Astar/n.balance))
+
+    where A is activity of the ith residue and At is total activity of residues.
+
+    Advantages:
+    - Loops over species only - much faster than equil.reaction
+    - No root finding - those games might fail at times
+
+    Disadvantage:
+    - Only works for per-residue reactions (n.balance = 1)
+    - Can create NaN logacts if the Astars are huge/small
+
+    Parameters
+    ----------
+    Astar : list of ndarray
+        Normalized affinities for each species
+    n_balance : list of float
+        Balancing coefficients (must all be 1)
+    loga_balance : ndarray
+        Log activity of the balanced quantity
+
+    Returns
+    -------
+    list of ndarray
+        Equilibrium log activities for each species
+    """
+
+    if not all(n == 1 for n in n_balance):
+        raise ValueError("won't run equil.boltzmann for balance != 1")
+
+    # Initialize output object
+    A = [np.array(a, copy=True) for a in Astar]
+
+    # Remember the dimensions of elements of Astar
+    Astardim = Astar[0].shape if Astar[0].ndim > 0 else (len(Astar[0]),)
+
+    # First loop: make vectors
+    A = [a.flatten() for a in A]
+    loga_balance_vec = loga_balance.flatten()
+
+    # Second loop: get the exponentiated Astars (numerators)
+    # Need to convert /2.303RT to /RT
+    A = [np.exp(np.log(10) * Astar[i].flatten() / n_balance[i])
+         for i in range(len(A))]
+
+    # Third loop: accumulate the denominator
+    # Initialize variable to hold the sum
+    At = np.zeros_like(A[0])
+    for i in range(len(A)):
+        At = At + A[i] * n_balance[i]
+
+    # Fourth loop: calculate log abundances
+    A = [loga_balance_vec + np.log10(A[i] / At) for i in range(len(A))]
+
+    # Fifth loop: restore dimensions
+    A = [a.reshape(Astardim) for a in A]
+
+    return A
+
+
+def equil_reaction(Astar: List[np.ndarray],
+                  n_balance: List[float],
+                  loga_balance: np.ndarray,
+                  tol: float = np.finfo(float).eps ** 0.25) -> List[np.ndarray]:
+    """
+    Calculate equilibrium activities using reaction-based method.
+
+    To turn the affinities/RT (A) of formation reactions into
+    logactivities of species (logact(things)) at metastable equilibrium.
+
+    For any reaction stuff = thing,
+      A = logK - logQ
+        = logK - logact(thing) + logact(stuff)
+    given Astar = A + logact(thing),
+    given Abar = A / n.balance,
+      logact(thing) = Astar - Abar * n.balance  [2]
+
+    where n.balance is the number of the balanced quantity
+    (conserved component) in each species.
+
+    Equilibrium values of logact(thing) satisfy:
+    1) Abar is equal for all species
+    2) log10(sum of (10^logact(thing) * n.balance)) = loga.balance  [1]
+
+    Because of the logarithms, we can't solve the equations directly.
+    Instead, use root-finding to compute Abar satisfying [1].
+
+    Parameters
+    ----------
+    Astar : list of ndarray
+        Normalized affinities for each species
+    n_balance : list of float
+        Balancing coefficients
+    loga_balance : ndarray
+        Log activity of the balanced quantity
+    tol : float
+        Tolerance for root-finding
+
+    Returns
+    -------
+    list of ndarray
+        Equilibrium log activities for each species
+    """
+
+    # We can't run on one species
+    if len(Astar) == 1:
+        raise ValueError("at least two species needed for reaction-based equilibration")
+
+    # Remember the dimensions and names
+    Adim = Astar[0].shape if Astar[0].ndim > 0 else None
+
+    # Make a matrix out of the list of Astar
+    Astar_array = np.array([a.flatten() for a in Astar]).T
+
+    if len(loga_balance) != Astar_array.shape[0]:
+        raise ValueError("length of loga.balance must be equal to the number of conditions for affinity()")
+
+    # Function definitions:
+    def logafun(logact):
+        """Calculate log of activity of balanced quantity from logact(thing) of all species [1]"""
+        # Use log-sum-exp trick for numerical stability
+        # log10(sum(10^x_i * n_i)) = log10(sum(n_i * 10^x_i))
+        # = max(x) + log10(sum(n_i * 10^(x_i - max(x))))
+        # This prevents overflow when x_i values are very large or very small
+
+        logact = np.asarray(logact)
+        n_balance_arr = np.asarray(n_balance)
+
+        # Find maximum for numerical stability
+        max_logact = np.max(logact)
+
+        # Compute sum in log space with shifted values
+        # sum(n_i * 10^x_i) = 10^max(x) * sum(n_i * 10^(x_i - max(x)))
+        shifted = logact - max_logact
+        sum_shifted = np.sum(n_balance_arr * 10**shifted)
+
+        # Convert back: log10(10^max(x) * sum(...)) = max(x) + log10(sum(...))
+        return max_logact + np.log10(sum_shifted)
+
+    def logactfun(Abar, i):
+        """Calculate logact(thing) from Abar for the ith condition [2]"""
+        return Astar_array[i, :] - Abar * np.array(n_balance)
+
+    def logadiff(Abar, i):
+        """Calculate difference between logafun and loga.balance for the ith condition"""
+        return loga_balance[i] - logafun(logactfun(Abar, i))
+
+    def Abarrange(i):
+        """Calculate a range of Abar that gives negative and positive values of logadiff for the ith condition"""
+        # Starting guess of Abar (min/max) from range of Astar / n.balance
+        Abar_range = [
+            np.min(Astar_array[i, :] / n_balance),
+            np.max(Astar_array[i, :] / n_balance)
+        ]
+
+        # diff(Abar.range) can't be 0 (dlogadiff.dAbar becomes NaN)
+        if Abar_range[1] - Abar_range[0] == 0:
+            Abar_range[0] -= 0.1
+            Abar_range[1] += 0.1
+
+        # The range of logadiff
+        logadiff_min = logadiff(Abar_range[0], i)
+        logadiff_max = logadiff(Abar_range[1], i)
+
+        # We're out of luck if they're both infinite
+        if np.isinf(logadiff_min) and np.isinf(logadiff_max):
+            raise ValueError("FIXME: there are no initial guesses for Abar that give "
+                           "finite values of the differences in logarithm of activity "
+                           "of the conserved component")
+
+        # If one of them is infinite we might have a chance
+        if np.isinf(logadiff_min):
+            # Decrease the Abar range by increasing the minimum
+            Abar_range[0] = Abar_range[0] + 0.99 * (Abar_range[1] - Abar_range[0])
+            logadiff_min = logadiff(Abar_range[0], i)
+            if np.isinf(logadiff_min):
+                raise ValueError("FIXME: the second initial guess for Abar.min failed")
+
+        if np.isinf(logadiff_max):
+            # Decrease the Abar range by decreasing the maximum
+            Abar_range[1] = Abar_range[1] - 0.99 * (Abar_range[1] - Abar_range[0])
+            logadiff_max = logadiff(Abar_range[1], i)
+            if np.isinf(logadiff_max):
+                raise ValueError("FIXME: the second initial guess for Abar.max failed")
+
+        iter_count = 0
+        while logadiff_min > 0 or logadiff_max < 0:
+            # The change of logadiff with Abar
+            # It's a weighted mean of the n.balance
+            dlogadiff_dAbar = (logadiff_max - logadiff_min) / (Abar_range[1] - Abar_range[0])
+
+            # Change Abar to center logadiff (min/max) on zero
+            logadiff_mean = (logadiff_min + logadiff_max) / 2
+            Abar_range[0] -= logadiff_mean / dlogadiff_dAbar
+            Abar_range[1] -= logadiff_mean / dlogadiff_dAbar
+
+            # One iteration is enough for the examples in the package
+            # but there might be a case where the range of logadiff doesn't cross zero
+            logadiff_min = logadiff(Abar_range[0], i)
+            logadiff_max = logadiff(Abar_range[1], i)
+            iter_count += 1
+
+            if iter_count > 5:
+                raise ValueError("FIXME: we seem to be stuck! This function (Abarrange() in "
+                               "equil.reaction()) can't find a range of Abar such that the differences "
+                               "in logarithm of activity of the conserved component cross zero")
+
+        return Abar_range
+
+    def Abarfun(i):
+        """Calculate an equilibrium Abar for the ith condition"""
+        # Get limits of Abar where logadiff brackets zero
+        Abar_range = Abarrange(i)
+
+        # Now for the real thing: brentq (Python's uniroot)!
+        Abar = brentq(logadiff, Abar_range[0], Abar_range[1], args=(i,), xtol=tol)
+        return Abar
+
+    # Calculate the logact(thing) for each condition
+    logact = []
+    for i in range(Astar_array.shape[0]):
+        # Get the equilibrium Abar for each condition
+        Abar = Abarfun(i)
+        logact.append(logactfun(Abar, i))
+
+    # Restore the dimensions
+    logact = np.array(logact)
+
+    # Convert back to list of arrays with original dimensions
+    result = []
+    for i in range(logact.shape[1]):
+        thisla = logact[:, i]
+        if Adim is not None:
+            thisla = thisla.reshape(Adim)
+        result.append(thisla)
+
+    return result
+
+
+def _balance(aout: Dict[str, Any],
+            balance: Optional[Union[str, int, List[float]]] = None,
+            messages: bool = True) -> Dict[str, Any]:
+    """
+    Return balancing coefficients and description.
+
+    Generate n.balance from user-given or automatically identified basis species.
+
+    Parameters
+    ----------
+    aout : dict
+        Output from affinity()
+    balance : str, int, or list of float, optional
+        Balance specification:
+        - None: autoselect using which_balance
+        - name of basis species: balanced on this basis species
+        - "length": balanced on sequence length of proteins
+        - "volume": standard-state volume listed in thermo()$OBIGT
+        - 1: balanced on one mole of species (formula units)
+        - numeric vector: user-defined n.balance
+
+    Returns
+    -------
+    dict
+        Dictionary with keys:
+        - n_balance : list, Balancing coefficients
+        - balance : str or list, Balancing description
+    """
+
+    # The index of the basis species that might be balanced
+    ibalance = None
+
+    # Deal with proteins
+    isprotein = ['_' in str(name) for name in aout['species']['name']]
+    if balance is None and all(isprotein):
+        balance = "length"
+
+    # Try to automatically find a balance
+    if balance is None:
+        ibalance = which_balance(aout['species'])
+        # No shared basis species and balance not specified by user - an error
+        if ibalance is None or len(ibalance) == 0:
+            raise ValueError("no basis species is present in all formation reactions")
+
+    # Change "1" to 1 (numeric)
+    if balance == "1":
+        balance = 1
+
+    if isinstance(balance, (int, float, list, np.ndarray)):
+        # A numeric vector
+        if isinstance(balance, (int, float)):
+            balance = [balance]
+        n_balance = list(balance) * (len(aout['values']) // len(balance) + 1)
+        n_balance = n_balance[:len(aout['values'])]
+
+        msgtxt = f"balance: on supplied numeric argument ({','.join(map(str, balance))})"
+        if balance == [1]:
+            msgtxt = f"{msgtxt} [1 means balance on formula units]"
+        if messages:
+            print(msgtxt)
+    else:
+        # "length" for balancing on protein length
+        if balance == "length":
+            if not all(isprotein):
+                raise ValueError("'length' was the requested balance, but some species are not proteins")
+            n_balance = [protein_length(name) for name in aout['species']['name']]
+            if messages:
+                print("balance: on protein length")
+        elif balance == "volume":
+            ispecies_list = aout['species']['ispecies'].tolist()
+            volumes = info(ispecies_list, check_it=False, messages=messages)['V']
+            n_balance = volumes.tolist()
+            if messages:
+                print("balance: on volume")
+        else:
+            # Is the balance the name of a basis species?
+            if ibalance is None or len(ibalance) == 0:
+                # Get basis rownames
+                basis_names = list(aout['basis'].index)
+                try:
+                    ibalance = [basis_names.index(balance)]
+                except ValueError:
+                    raise ValueError(f"basis species ({balance}) not available to balance reactions")
+
+            # The name of the basis species (need this if we got ibalance from which_balance, above)
+            balance = list(aout['species'].columns)[ibalance[0]]
+            if messages:
+                print(f"balance: on moles of {balance} in formation reactions")
+
+            # The balancing coefficients
+            n_balance = aout['species'].iloc[:, ibalance[0]].tolist()
+
+            # We check that all formation reactions contain this basis species
+            if any(n == 0 for n in n_balance):
+                raise ValueError(f"some species have no {balance} in the formation reaction")
+
+    return {'n_balance': n_balance, 'balance': balance}
+
+
+def which_balance(species: pd.DataFrame) -> List[int]:
+    """
+    Return column(s) of species that all have non-zero values.
+
+    Find the first basis species that is present in all species of interest.
+    It can be used to balance the system.
+
+    Parameters
+    ----------
+    species : pd.DataFrame
+        Species dataframe from affinity output
+
+    Returns
+    -------
+    list of int
+        Indices of basis species columns that have non-zero values for all species
+    """
+
+    # Number of basis species columns (exclude the last 4 metadata columns)
+    nbasis = len(species.columns) - 4
+
+    ib = []
+    for i in range(nbasis):
+        coeff = species.iloc[:, i]
+        # Check if all coefficients are non-zero
+        if all(c != 0 for c in coeff):
+            ib.append(i)
+            break  # R version returns first match
+
+    return ib
+
+
+def protein_length(name: Union[str, List[str]]) -> Union[int, List[int]]:
+    """
+    Get protein sequence length.
+
+    Parameters
+    ----------
+    name : str or list of str
+        Protein name(s) (with underscore separator)
+
+    Returns
+    -------
+    int or list of int
+        Sequence length(s)
+    """
+
+    if isinstance(name, str):
+        # Single protein
+        if '_' not in name:
+            raise ValueError(f"protein name '{name}' does not contain underscore")
+        # For now, return a placeholder - would need actual protein database
+        # In R this would look up the actual sequence length
+        return 100  # Placeholder
+    else:
+        # Multiple proteins
+        return [protein_length(n) for n in name]
+
+
+def moles(eout: Dict[str, Any]) -> Dict[str, np.ndarray]:
+    """
+    Calculate total moles of elements from equilibrate output.
+
+    Parameters
+    ----------
+    eout : dict
+        Output from equilibrate()
+
+    Returns
+    -------
+    dict
+        Dictionary with element names as keys and mole arrays as values
+    """
+
+    # Exponentiate loga.equil to get activities
+    act = [10**np.array(x) for x in eout['loga_equil']]
+
+    # Initialize list for moles of basis species
+    nbasis_list = [act[0] * 0 for _ in range(len(eout['basis']))]
+
+    # Loop over species
+    for i in range(len(eout['species'])):
+        # Loop over basis species
+        for j in range(len(eout['basis'])):
+            # The coefficient of this basis species in the formation reaction of this species
+            n = eout['species'].iloc[i, j]
+            # Accumulate the number of moles of basis species
+            nbasis_list[j] = nbasis_list[j] + act[i] * n
+
+    # Initialize list for moles of elements (same as number of basis species)
+    nelem = [act[0] * 0 for _ in range(len(eout['basis']))]
+
+    # Loop over basis species
+    for i in range(len(eout['basis'])):
+        # Loop over elements
+        for j in range(len(eout['basis'])):
+            # The coefficient of this element in the formula of this basis species
+            n = eout['basis'].iloc[i, j]
+            # Accumulate the number of moles of elements
+            nelem[j] = nelem[j] + nbasis_list[i] * n
+
+    # Add element names
+    element_names = list(eout['basis'].columns)[:len(eout['basis'])]
+    result = {element_names[i]: nelem[i] for i in range(len(nelem))}
+
+    return result