8000 Converge ionic strength by secant and warm-start equilibrium searches (#61) by rafatahmed · Pull Request #62 · rafatahmed/Flahax · GitHub
[go: up one dir, main page]

Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

## Unreleased

- Speed up equilibrium solves 3-5x and their pH and dose searches 10-25x in pure Python (#61), found by the #40 comparison with PHREEQC. Each solve averaged its ionic-strength update and needed about 35 outer passes to reach |dI| < 1e-12; ionic strength now takes a safeguarded secant step (the average remains the fallback) and water activity plain substitution. `solve(..., start=result)` warm-starts from a nearby state (free-ion activities, ionic strength, water, precipitated phases) and falls back to the unchanged cold start if it fails; the nitric-target, calculated-pH, acid-selection and measured-water searches warm-start from their previous solve. Converged states are unchanged within solver tolerance (species within 7e-10 relative on every PHREEQC fixture; calculated pH within its 1e-7 bisection width). The mixed-product nitric target drops from about 1080 to 73 ms (PHREEQC: 4.5 ms), calculated-pH irrigation EC from 0.9-2.6 s to 74-244 ms per recipe, the test suite from 133 to 40 s. The shipped phosphate conductivity coefficients are unchanged; refitting them now reproduces them to 7.4e-9 relative instead of bit-exactly, so the reproduction test bounds coefficients at 1e-7 relative and percentages at 1e-5 points (owner decision).
- `recommend()` is 2.8-4.0x faster in pure Python with identical recipes (#59). NNLS was 97% of its time: about 44 solves per call, each rebuilding a Householder factorisation over every row. Simplification trials now warm-start NNLS from the current solution, passive-set least squares use Cholesky on a lazily cached Gram matrix (Householder fallback and refinement when ill-conditioned), dual weights come from the cached Gram entries, and the miss diagnosis computes each drop-one fit once. Over the 52 benchmark cases (best of three, same machine): 248 -> 89 ms with the default ammonium limit, 297 -> 95 ms with a Cl ceiling, 146 -> 36 ms without ceilings. Every recipe is unchanged (largest final concentration change 3e-11 mg/L, identical product sets). Zero runtime dependencies are kept.
- Fix the equilibrium solver failing to converge between pH 4.2 and 5.05 for chelate-rich mixed solutions (#52). From the cold start (every free activity at total/1000), Newton crawled toward strongly complexed free ions (EDTA, DTPA and EDDHA with Fe and Cu) until its iteration limit; an acid target of pH 5.0 for the mixed-product fixture failed. A failed solve now restarts from a continued-fraction mass-balance estimate (Wigley 1977). A pH 4-9 scan of the three mixed fixtures (1,503 solves) goes from 215 failures to none; converging solves are unchanged because the restart runs only after a failure.
- The NNLS solver raises `InputError` when it is still not optimal at its iteration cap, instead of silently returning a non-optimal fit (#55). No benchmark fit reaches the cap.
Expand Down
2 changes: 1 addition & 1 deletion docs/site/models/fertilizer-equilibrium-model.md
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ SI_phase = sum(phase_power * log10(a_basis)) - log_K_phase

Pivoted elimination and damped line search solve all balances together. An active set enforces nonnegative solids, SI = 0 for present phases and SI <= 0 for eligible absent phases. Each phase change resolves the whole mixture, not disconnected macro/trace kernels. Zero-total species are omitted exactly; stalled solves raise `DeliveryError('nonconvergent', ...)`.

Ionic strength and dilute water activity iterate to 1e-12. Component Newton residual is below 2e-10 log10 units; tests require elemental residual below 1e-9 relative + 1e-15 mol/kgw. Davies uses A = 0.509; neutral log10(gamma) = 0.1 I; water activity is approximated by `1 - sum(m)/55.5084`. PHREEQC retains its native ion-specific gamma parameters. Fixture tolerances explicitly bound this difference; the activity models are not claimed identical at every composition.
Ionic strength and dilute water activity iterate to 1e-12 (ionic strength by a safeguarded secant step on I(speciation at I) - I, water activity by substitution). Dose and pH searches warm-start each solve from the previous one; the start changes only the iteration path. Component Newton residual is below 2e-10 log10 units; tests require elemental residual below 1e-9 relative + 1e-15 mol/kgw. Davies uses A = 0.509; neutral log10(gamma) = 0.1 I; water activity is approximated by `1 - sum(m)/55.5084`. PHREEQC retains its native ion-specific gamma parameters. Fixture tolerances explicitly bound this difference; the activity models are not claimed identical at every composition.

Only 25 C and I <= 0.1 mol/kgw are supported. Negative/nonfinite/unknown totals, unsupported phases, pH outside [0,14], and out-of-range temperature/ionic strength fail explicitly. Numerical pH bounds are not experimental validation of every constant over that entire interval. No Pitzer extrapolation, gas exchange, urea hydrolysis, redox kinetics, sorption or precipitation kinetics is inferred.

Expand Down
7 changes: 5 additions & 2 deletions src/flahax/acid_selection.py
Original file line number Diff line number Diff line change
Expand Up @@ -109,10 +109,13 @@ def plan_acid_options(totals, initial_ph, target_ph, *, products, water_mass_kg,
raise DeliveryError('wrong_reagent_direction', 'acid selection cannot satisfy a base requirement')
basis, molar_mass, nutrients = ACIDS[formula]

def at(dose):
last = [initial]

def at(dose, last=last):
amended = dict(totals)
amended[basis] = amended.get(basis, 0) + dose
return solve(amended, target_ph, phases=phases)
last[0] = solve(amended, target_ph, phases=phases, start=last[0]) # warm start, #61
return last[0]

dose = 0.
target = initial if target_ph == initial_ph else at(0.)
Expand Down
47 changes: 41 additions & 6 deletions src/flahax/aqueous_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -101,7 +101,18 @@ def lgamma(z):
break
return logs

def solve(totals, ph, *, phases=(), ionic_strength=None, temperature_c=25.):
def solve(totals, ph, *, phases=(), ionic_strength=None, temperature_c=25., start=None):
"""Equilibrium at fixed pH. ``start`` (an AqueousResult of a nearby state, such as
the previous step of a dose or pH search) only seeds the iteration; the converged
state is the same within solver tolerance, and a failed warm start falls back to
the cold start below."""
if start is not None:
try:
return _solve(totals, ph, phases=phases, ionic_strength=ionic_strength,
temperature_c=temperature_c, start=start)
except DeliveryError as exc:
if exc.code != 'nonconvergent':
raise
try:
return _solve(totals, ph, phases=phases, ionic_strength=ionic_strength, temperature_c=temperature_c)
except DeliveryError as exc:
Expand All @@ -112,7 +123,7 @@ def solve(totals, ph, *, phases=(), ionic_strength=None, temperature_c=25.):
return _solve(totals, ph, phases=phases, ionic_strength=ionic_strength, temperature_c=temperature_c,
presolve=True)

def _solve(totals, ph, *, phases=(), ionic_strength=None, temperature_c=25., presolve=False):
def _solve(totals, ph, *, phases=(), ionic_strength=None, temperature_c=25., presolve=False, start=None):
if temperature_c != 25:
raise DeliveryError('temperature_out_of_range','reaction dataset is restricted to 25 C')
if not math.isfinite(ph) or not 0 <= ph <= 14:
Expand Down Expand Up @@ -142,7 +153,15 @@ def _solve(totals, ph, *, phases=(), ionic_strength=None, temperature_c=25., pre
active=[]
solids={}
water=1.
if start is not None:
logs=[math.log10(start.activities[b]) if start.activities.get(b,0)>0 else v for b,v in zip(bases,logs)]
if ionic_strength is None:
strength=min(.1,max(1e-8,start.ionic_strength))
water=start.water_activity
active=[p for p,v in start.precipitated.items() if v>0 and p in candidates]
solids={p:start.precipitated[p] for p in active}
count=0
previous=None
def lgamma(z):
root=math.sqrt(strength)
return -.509*z*z*(root/(1+root)-.3*strength) if z else .1*strength
Expand Down Expand Up @@ -187,7 +206,9 @@ def evaluate(x, jacobian=False):
count+=1
residual,amounts,matrix=evaluate(x,True)
norm=max(map(abs,residual),default=0.)
if norm<2e-10:
# A warm start can begin inside the tolerance with the previous state's
# residual; one Newton step restores the cold path's quadratic finish (#61).
if norm<2e-10 and (iteration or start is None or norm<1e-14):
break
step=linear(matrix,[-r for r in residual])
alpha=min(1.,4/max([abs(v) for v in step[:n]]+[4.]))
Expand Down Expand Up @@ -234,8 +255,19 @@ def evaluate(x, jacobian=False):
{p:max(0,v) for p,v in solids.items()},residuals,count,
sum(a*r['charge'] for a,r in zip(amounts,records)),water)
if ionic_strength is None:
strength=(strength+updated)/2
water=(water+next_water)/2
# Secant step on g(I) = I(speciation at I) - I, which converges superlinearly
# whether ion pairing makes g's slope positive or negative; the plain average
# (linear, ~35 passes to 1e-12, #61) remains the fallback.
gap=updated-strength
step=(strength+updated)/2
if previous is not None and gap!=previous[1]:
secant=strength-gap*(strength-previous[0])/(gap-previous[1])
if 0<secant<=.1:
step=secant
previous=(strength,gap)
strength=step
# Water activity depends only weakly on itself (through H2O-bearing species).
water=next_water
raise DeliveryError('nonconvergent','ionic strength/water iteration limit')

@dataclass(frozen=True)
Expand All @@ -258,8 +290,11 @@ def nitric_target(totals, initial_ph, target_ph, *, phases=()):
if target_ph>initial_ph:
raise DeliveryError('wrong_reagent_direction','target requires base')
lo,hi=0.,.001
last=[initial]
def at(dose):
return solve(dict(totals)|{'NO3-':totals.get('NO3-',0.)+dose},target_ph,phases=phases)
# Warm-start each dose from the previous solve (#61); only the iteration path changes.
last[0]=solve(dict(totals)|{'NO3-':totals.get('NO3-',0.)+dose},target_ph,phases=phases,start=last[0])
return last[0]
target=at(0.)
if target.charge_balance < initial.charge_balance-1e-12:
raise DeliveryError('wrong_reagent_direction','target requires base')
Expand Down
7 changes: 6 additions & 1 deletion src/flahax/calculated_ph.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,8 +34,13 @@ def calculate_equilibrium_ph(totals, *, complete_analysis, carbon_boundary,
if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value) or value < 0:
raise DeliveryError('invalid_totals', 'analytical totals must be finite nonnegative numbers')

last = []

def at(ph):
return solve(dict(totals), ph, phases=phases, temperature_c=temperature_c)
# Warm-start from the previous pH step (#61); the bracket sequence is unchanged.
last[:] = [solve(dict(totals), ph, phases=phases, temperature_c=temperature_c,
start=last[0] if last else None)]
return last[0]

# Start near neutral to avoid unnecessarily evaluating extreme acidic/basic
# states outside the Davies range. No out-of-domain solve is extrapolated.
8BD0 Expand Down
9 changes: 6 additions & 3 deletions src/flahax/measured_water.py
Original file line number Diff line number Diff line change
Expand Up @@ -70,9 +70,12 @@ def _imbalance_percent(state):

def _carbon_for_alkalinity(totals, ph, alkalinity, phases):
"""Total inorganic carbon whose speciation at pH gives the alkalinity."""
last = [solve({k: v for k, v in totals.items() if k != 'CO3-2'}, ph, phases=phases)]

def alk(carbon):
return carbonate_alkalinity(solve(totals | {'CO3-2': carbon}, ph, phases=phases))
base = carbonate_alkalinity(solve({k: v for k, v in totals.items() if k != 'CO3-2'}, ph, phases=phases))
last[0] = solve(totals | {'CO3-2': carbon}, ph, phases=phases, start=last[0]) # warm start, #61
return carbonate_alkalinity(last[0])
base = carbonate_alkalinity(last[0])
if alkalinity <= base:
raise DeliveryError('alkalinity_out_of_range',
'measured alkalinity is not above the carbon-free alkalinity at the measured pH')
Expand Down Expand Up @@ -153,7 +156,7 @@ def balance_water_analysis(totals, *, measured_ph, alkalinity_eq_per_kg=None, ba
raise DeliveryError('balance_ion_exhausted', f'{balance_ion} cannot absorb the charge imbalance')
if carbon_source == 'alkalinity':
working['CO3-2'] = _carbon_for_alkalinity(working, measured_ph, alkalinity_eq_per_kg, phases)
state = solve(working, measured_ph, phases=phases)
state = solve(working, measured_ph, phases=phases, start=state)
else:
raise DeliveryError('nonconvergent', 'measured-pH charge balance did not close')
return BalancedWaterAnalysis(dict(working), float(measured_ph), state, balance_ion, analysed,
Expand Down
11 changes: 9 additions & 2 deletions tests/test_irrigation_ec.py
Original file line number Diff line number Diff line change
Expand Up @@ -287,13 +287,20 @@ def test_all_shipped_formulas_estimated(self):


def _assert_close_tree(case, got, want, path='supplementary'):
"""Equal structure; numbers within 1e-9, everything else equal."""
"""Equal structure; numbers within 1e-7 relative (1e-9 absolute floor), everything else equal.

The fit speciates through calculated pH, which returns the midpoint of a bisection
bracket up to 1e-7 pH wide; any change in the solver's iteration path (#61) moves
the refit in the ninth digit, far below the ~1e-4 precision of Table III."""
if isinstance(want, dict):
case.assertEqual(set(got), set(want), path)
for key in want:
_assert_close_tree(case, got[key], want[key], f'{path}.{key}')
elif isinstance(want, (int, float)) and not isinstance(want, bool):
case.assertAlmostEqual(got, want, delta=1e-9, msg=path)
# Percentages (fit residuals, gaps to other lambda0 sources) inherit the 1e-7
# relative bound of the conductivities they compare: 1e-5 percentage points.
floor = 1e-5 if path.endswith('percent') else 1e-9
case.assertAlmostEqual(got, want, delta=max(floor, 1e-7 * abs(want)), msg=path)
else:
case.assertEqual(got, want, path)

Expand Down
63 changes: 63 additions & 0 deletions tests/test_solver_convergence.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,63 @@
"""#61: ionic-strength convergence and warm starts change the iteration path, not the state."""
import json
import unittest
from pathlib import Path

from flahax.aqueous_model import nitric_target, solve

FIXTURES = Path(__file__).resolve().parent / 'fixtures' / 'phreeqc'


def fixture(name):
return json.loads((FIXTURES / name / 'expected.json').read_text())


class WarmStartTests(unittest.TestCase):
def assert_same_state(self, a, b, relative=1e-9):
self.assertEqual(set(a.species), set(b.species))
for name, amount in a.species.items():
if amount > 1e-20:
self.assertLess(abs(b.species[name] - amount), relative * amount, name)
self.assertAlmostEqual(a.ionic_strength, b.ionic_strength, delta=1e-11) # both stop at |dI| < 1e-12
self.assertEqual({p for p, v in a.precipitated.items() if v > 0},
{p for p, v in b.precipitated.items() if v > 0})

def test_warm_start_from_a_neighbouring_ph_reaches_the_cold_state(self):
for name in ('products/mixed_products', 'phases/mixed_product_phases', 'phases/gypsum'):
with self.subTest(name):
data = fixture(name)
ph = data['selected_rows'][-1]['pH']
phases = data.get('phases', ())
neighbour = solve(data['runtime_totals'], ph + .3, phases=phases)
cold = solve(data['runtime_totals'], ph, phases=phases)
warm = solve(data['runtime_totals'], ph, phases=phases, start=neighbour)
self.assert_same_state(cold, warm)

def test_warm_start_from_the_solution_takes_one_newton_step(self):
data = fixture('products/mixed_products')
ph = data['selected_rows'][-1]['pH']
cold = solve(data['runtime_totals'], ph)
warm = solve(data['runtime_totals'], ph, start=cold)
self.assertLessEqual(warm.iterations, 2)
self.assertLess(warm.iterations, cold.iterations)
self.assertLessEqual(abs(warm.charge_balance - cold.charge_balance), 1e-15)

def test_nitric_target_dose_is_unchanged(self):
data = fixture('products/mixed_target_ph_nitric')
plan = nitric_target(data['runtime_totals'], data['initial_ph'], data['selected_rows'][-1]['pH'])
tolerance = data['tolerances']
self.assertAlmostEqual(plan.nitric_acid_molal, data['nitric_acid_molal'],
delta=tolerance['acid_absolute'] + tolerance['acid_relative'] * data['nitric_acid_molal'])


class IonicStrengthConvergenceTests(unittest.TestCase):
def test_strong_electrolyte_converges_in_few_newton_iterations(self):
# Gypsum-saturated solution near I = 0.055: the averaged update needed ~47 outer
# passes (102 Newton iterations); the secant step needs a fraction of that.
data = fixture('phases/gypsum')
result = solve(data['runtime_totals'], data['selected_rows'][-1]['pH'], phases=data['phases'])
self.assertLess(result.iterations, 60)


if __name__ == '__main__':
unittest.main()
Loading
0