pH equilibria optimization leads to numerical instability

Question

I'm creating a pH equilibrium calculator. Starting with compositions of acids / bases the calculator should come up with the equilibrated concentrations of each component, based on physical constants. I formulated this as a set of equations which need to be solved simultaneously. One key constraint is H_plus*OH_minus must be 10^-14. I included this in the equations.

Fsolve. This works but is sensitive to initial conditions and I can't put constraints (I need concentrations to be >=0):

import numpy as np
from scipy.optimize import fsolve


def equations(variables, Ka, Kw, initial_concentrations):
    HA, A_minus, H_plus, OH_minus, ion_Na = variables
    HA_init, A_minus_init, H_plus_init, OH_minus_init, ion_Na_init = initial_concentrations
    # Define the equations
    eq1 = HA + A_minus - HA_init - A_minus_init # Mass balance on acetate
    eq2 = H_plus + ion_Na - OH_minus - A_minus  # charge balance
    eq3 = H_plus * OH_minus - Kw                # Kw for water
    eq4 = Ka * HA - A_minus * H_plus            # HA acid dissociation
    eq5 = ion_Na - ion_Na_init                  # inert ion
    return [eq1, eq2, eq3, eq4, eq5]

def equilibrate(initial_concentrations, initial_guess):
    # User-supplied equilibrium constant (Ka) and initial concentration of acetic acid (HA_initial)
    Ka = 1.8e-5             # mol/L, acetic acid dissociation constant
    Kw = 1e-14              # mol/L, ionic product of water

    # Use fsolve from SciPy to find the equilibrium concentrations
    solution = fsolve(equations, initial_guess, args=(Ka, Kw, initial_concentrations), xtol=1e-18)

    return solution


initial_concentrations = [0.0, 0.1, 1e-7, 1e-7, 0.1]
initial_guess = initial_concentrations

# Extract the concentrations
HA, A_minus, H_plus, OH_minus, ion_Na = equilibrate(initial_concentrations, initial_guess)

# Calculate pH
pH = -np.log10(H_plus)
pKw = -np.log10(H_plus) -np.log10(OH_minus)

# Display the results
print(f"HA: {HA:.8f}")
print(f"A-: {A_minus:.8f}")
print(f"H+: {H_plus:.8f}")
print(f"OH-: {OH_minus:.8f}")
print(f"pH: {pH:.2f}")
print(f"pKw: {pKw:.2f}") # THIS IS THE CONDIRMATORY CHECK - Should be 14.

Nelder-Mead. This converges despite the tight tolerances I set, but gives meaningless values (pKw):

import numpy as np
from scipy.optimize import minimize

def equations(variables, Ka, Kw, initial_concentrations):
    HA, A_minus, H_plus, OH_minus, ion_Na = variables
    HA_init, A_minus_init, H_plus_init, OH_minus_init, ion_Na_init = initial_concentrations
    # Define the equations
    eq1 = HA + A_minus - HA_init - A_minus_init # Mass balance on acetate
    eq2 = H_plus + ion_Na - OH_minus - A_minus  # Charge balance
    eq3 = H_plus * OH_minus - Kw                # Kw for water
    eq4 = Ka * HA - A_minus * H_plus            # HA acid dissociation
    eq5 = ion_Na - ion_Na_init                  # Inert ion
    return eq1**2 + eq2**2 + eq3**2 + eq4**2 + eq5**2  # Use squared error for the minimization

def equilibrate(initial_concentrations, initial_guess, Ka, Kw):
    # Define bounds for variables (concentrations should be >= 0)
    bounds = [[0, 1], [0, 1], [0, 1], [0, 1], [0, 1]]

    # Use minimize from SciPy with BFGS method to find the equilibrium concentrations
    result = minimize(equations, initial_guess, args=(Ka, Kw, initial_concentrations),
                      bounds=bounds,method='Nelder-Mead', options={'disp': True, 'xatol':1E-40, 'fatol':1E-40, 'maxfev':1000000})

    return result.x

# User-supplied equilibrium constant (Ka) and initial concentration of acetic acid (HA_initial)
Ka = 1.8e-5             # mol/L, acetic acid dissociation constant
Kw = 1e-14              # mol/L, ionic product of water

initial_concentrations = [0.00, 0.1, 1e-7, 1e-7, 0.1]

initial_guess = initial_concentrations

# Extract the concentrations
HA, A_minus, H_plus, OH_minus, ion_Na = equilibrate(initial_concentrations, initial_guess, Ka, Kw)

# Calculate pH
pH = -np.log10(H_plus)

pKw = -np.log10(H_plus) -np.log10(OH_minus)

# Display the results
print(f"HA: {HA:.8f}")
print(f"A-: {A_minus:.8f}")
print(f"H+: {H_plus:.8f}")
print(f"OH-: {OH_minus:.8f}")
print(f"pH: {pH:.2f}")
print(f"pKw: {pKw:.2f}") # THIS IS THE CONDIRMATORY CHECK - Should be 14.

The confirmatory check value is all over the place instead of being 14.

1. Why does the Nelder-Mead algorithm do so poorly? Is it because the absolute values are close to the error? I thought I could resolve this by tightening tolerances.
2. How to impose positive constraints in fsolve?
3. Is there a better way to do this?

Answer 1

Don't force Nelder-Mead unless you have a very good reason, and I don't see one here.

ion_Na is not actually a degree of freedom, so either pin it to its initial value via bounds, or don't include it in the variables at all.

Include (and test) the Jacobian of the cost function.

Minimize, explicit variables

from typing import NamedTuple

import numpy as np
from scipy.optimize import check_grad, minimize, Bounds, NonlinearConstraint


class Concentrations(NamedTuple):
    HA: float
    A_minus: float
    H_plus: float
    OH_minus: float
    ion_Na: float

    @property
    def pH(self) -> float:
        return -np.log10(self.H_plus)

    @property
    def pKw(self) -> float:
        return -np.log10(self.H_plus) - np.log10(self.OH_minus)


def cost(variables: np.ndarray, init: Concentrations, Ka: float, Kw: float) -> tuple[
    float,  # cost
    tuple[float, ...],  # jacobian
]:
    conc = Concentrations(*variables)

    acetate_mass_balance = conc.HA + conc.A_minus - init.HA - init.A_minus
    charge_balance = conc.H_plus + conc.ion_Na - conc.OH_minus - conc.A_minus
    kw_water = conc.H_plus * conc.OH_minus - Kw
    ha_acid_dissociation = Ka * conc.HA - conc.A_minus * conc.H_plus
    inert_ion = conc.ion_Na - init.ion_Na

    cost = (
        acetate_mass_balance**2 + charge_balance**2 + kw_water**2 + ha_acid_dissociation**2 + inert_ion**2
    )

    jac = (
        2*acetate_mass_balance + 2*ha_acid_dissociation*Ka,
        2*acetate_mass_balance - 2*charge_balance - 2*ha_acid_dissociation*conc.H_plus,
        2*charge_balance + 2*kw_water*conc.OH_minus - 2*ha_acid_dissociation*conc.A_minus,
        -2*charge_balance + 2*kw_water*conc.H_plus,
        2*charge_balance + 2*inert_ion,
    )

    return cost, jac


def conc_pKw(variables: np.ndarray) -> float:
    return Concentrations(*variables).pKw


def jac_pKw(variables: np.ndarray) -> tuple[float, ...]:
    conc = Concentrations(*variables)
    num = -1 / np.log(10)
    return (0, 0, num/conc.H_plus, num/conc.OH_minus, 0)


def test_jacobians() -> None:
    x0 = Concentrations(2, 3, 5, 7, 11)
    error = check_grad(
        lambda v, *args: cost(v, *args)[0],
        lambda v, *args: cost(v, *args)[1],
        x0,
        Concentrations(2.1, 3.1, 5.1, 7.1, 11.1),
        0.5, 0.7,
    )
    assert error < 1e-5

    error = check_grad(func=conc_pKw, grad=jac_pKw, x0=x0)
    assert error < 1e-8


def main() -> None:
    test_jacobians()

    Ka = 1.8e-5
    Kw = 1e-14
    x0 = Concentrations(HA=0, A_minus=0.1, H_plus=1e-7, OH_minus=1e-7, ion_Na=0.1)

    result = minimize(
        fun=cost, jac=True, args=(x0, Ka, Kw), x0=x0, tol=1e-32,
        # bounds=Bounds(lb=0),
        bounds=Bounds(
            lb=(0, 0, 0, 0, x0.ion_Na),
            ub=(1, 1, 1, 1, x0.ion_Na),
        ),
        constraints=NonlinearConstraint(
            fun=conc_pKw, jac=jac_pKw,
            lb=14, ub=14,
        )
    )
    assert result.success, result.message

    print(f'{result.message} in {result.nit} iters, err={result.fun:.1e}')
    conc = Concentrations(*result.x)
    print(conc)
    print(f'pKw = {conc.pKw:.2f}')
    print(f'pH = {conc.pH:.2f}')


if __name__ == '__main__':
    main()

Optimization terminated successfully in 33 iters, err=1.9e-34
Concentrations(HA=7.452611421901198e-06, A_minus=0.0999925473885781, H_plus=1.34157003817923e-09, OH_minus=7.453952991952563e-06, ion_Na=0.1)
pKw = 14.00
pH = 8.87

Minimize, implicit variables

There is a different approach where you set the concentration variables implicitly based on your constraints for pKw and ion_Na; it is much faster but probably incorrect:

from typing import NamedTuple

import numpy as np
from scipy.optimize import check_grad, minimize, Bounds


class Concentrations(NamedTuple):
    HA: float
    A_minus: float
    H_plus: float

    @property
    def OH_minus(self) -> float:
        """
        For fixed pKw:
        -np.log10(self.H_plus) - np.log10(self.OH_minus) == 14
        np.log10(Hplus * OHminus) == -14
        Hplus * OHminus == 1e-14
        """
        return 1e-14 / self.H_plus

    @property
    def pH(self) -> float:
        return -np.log10(self.H_plus)

    @property
    def pKw(self) -> float:
        return -np.log10(self.H_plus) - np.log10(self.OH_minus)


def cost(variables: np.ndarray, init: Concentrations, Ka: float, Kw: float, ion_Na: float) -> tuple[
    float,  # cost
    tuple[float, ...],  # jacobian
]:
    conc = Concentrations(*variables)

    acetate_mass_balance = conc.HA + conc.A_minus - init.HA - init.A_minus
    charge_balance = conc.H_plus + ion_Na - conc.OH_minus - conc.A_minus
    kw_water = conc.H_plus * conc.OH_minus - Kw
    ha_acid_dissociation = Ka * conc.HA - conc.A_minus * conc.H_plus

    cost = (
        acetate_mass_balance**2 + charge_balance**2 + kw_water**2 + ha_acid_dissociation**2
    )

    jac = (
        2*acetate_mass_balance + 2*ha_acid_dissociation*Ka,
        2*acetate_mass_balance - 2*charge_balance - 2*ha_acid_dissociation*conc.H_plus,
        2*charge_balance + 2*kw_water*conc.OH_minus - 2*ha_acid_dissociation*conc.A_minus,
    )

    return cost, jac


def test_jacobians() -> None:
    x0 = Concentrations(2, 3, 5)
    error = check_grad(
        lambda v, *args: cost(v, *args)[0],
        lambda v, *args: cost(v, *args)[1],
        x0,
        Concentrations(2.1, 3.2, 5.3),
        0.5, 0.7, 0.1,
    )
    assert error < 1e-5


def main() -> None:
    test_jacobians()

    Ka = 1.8e-5
    Kw = 1e-14
    ion_Na = 0.1
    x0 = Concentrations(HA=0, A_minus=0.1, H_plus=1e-7)

    result = minimize(
        fun=cost, jac=True, x0=x0, tol=1e-48,
        args=(x0, Ka, Kw, ion_Na),
        bounds=Bounds(lb=1e-24, ub=1),
    )
    assert result.success, result.message

    print(f'{result.message} in {result.nit} iters, err={result.fun:.1e}')
    conc = Concentrations(*result.x)
    print(conc)
    print(f'OH- = {conc.OH_minus}')
    print(f'pKw = {conc.pKw:.2f}')
    print(f'pH = {conc.pH:.2f}')


if __name__ == '__main__':
    main()

The fundamental problem is that your system operates at wildly different orders of magnitude. Normalization scaling, and/or calculations in the logarithmic domain, will probably make this better-behaved.

Minimize, constraints only

This approach is certainly the most accurate, but somewhat slower, and also - instead of giving you an approximate solution if the inputs are malformed and would produce negative concentrations - will simply reject the problem as infeasible.

import numpy as np
from scipy.optimize import minimize, Bounds, LinearConstraint, NonlinearConstraint


def noop(x: np.ndarray) -> tuple[float, tuple[float, ...]]:
    return 0, (0, 0, 0, 0)


def water_Kw(x: np.ndarray) -> float:
    HA, Aminus, Hplus, OHminus = x
    return Hplus*OHminus


def water_Kw_jac(x: np.ndarray) -> tuple[float, ...]:
    HA, Aminus, Hplus, OHminus = x
    return 0, 0, OHminus, Hplus


def main() -> None:
    Ka = 1.8e-5
    Kw = 1e-14
    HA_init = 0
    A_minus_init = 0.1
    ion_Na = 0.1

    def HA_acid_dissociation(x: np.ndarray) -> float:
        HA, Aminus, Hplus, OHminus = x
        return Ka*HA - Aminus*Hplus

    def HA_jac(x: np.ndarray) -> tuple[float, ...]:
        HA, Aminus, Hplus, OHminus = x
        return Ka, -Hplus, -Aminus, 0

    result = minimize(
        fun=noop, tol=1e-24, jac=True,
        #   HA    A-    H+   OH-
        x0=( 0, 0.1, 1e-7, 1e-7),
        bounds=Bounds(lb=0, ub=1),
        constraints=(
            LinearConstraint(
                A=(# HA  A-  H+  OH-
                    ( 1,  1,  0,  0),  # Acetate mass balance
                    ( 0,  1, -1,  1),  # Charge balance
                ),
                lb=(HA_init + A_minus_init, ion_Na),
                ub=(HA_init + A_minus_init, ion_Na),
            ),
            NonlinearConstraint(fun=water_Kw, jac=water_Kw_jac, lb=Kw, ub=Kw),
            NonlinearConstraint(fun=HA_acid_dissociation, jac=HA_jac, lb=0, ub=0),
        ),
    )
    assert result.success, result.message

    print(f'{result.message} in {result.nit} iters')
    print('Concentrations:', result.x)
    print('water_kW error:', water_Kw(result.x) - 1e-14)
    print(f'Acid dissoc error:', HA_acid_dissociation(result.x))


if __name__ == '__main__':
    main()

Optimization terminated successfully in 51 iters
Concentrations: [7.45261142e-06 9.99925474e-02 1.34157004e-09 7.45395299e-06]
water_kW error: 0.0
Acid dissoc error: 0.0