From 8a095b49dbed49d3630e72d8579c8a6a39131fc5 Mon Sep 17 00:00:00 2001 From: Aaron Meyer Date: Wed, 26 Aug 2026 20:52:41 -0700 Subject: [PATCH] Improve root-finding convergence with bisection and log-scale formulation --- valentbind/model.py | 81 +++++++++++++++++++++--------- valentbind/test/test_edge_cases.py | 33 ++++++++++++ 2 files changed, 89 insertions(+), 25 deletions(-) diff --git a/valentbind/model.py b/valentbind/model.py index 09ce2d7..1ae2af2 100644 --- a/valentbind/model.py +++ b/valentbind/model.py @@ -26,8 +26,8 @@ def Req_polyfc( potential that is consistent with the mass balance for the total receptor and free ligand concentrations. - :param Phisum: Current guess for the binding potential (a length-1 - array, since the model reduces to a single scalar unknown). + :param Phisum: Current guess for the binding potential (a scalar array, + since the model reduces to a single scalar unknown). :param args: Tuple of ``(Rtot, L0, KxStar, f, A)`` where ``Rtot`` is the total receptor abundance per receptor type, ``L0`` is the total ligand complex concentration, ``KxStar`` is the detailed-balance @@ -42,17 +42,18 @@ def Req_polyfc( def Req_polyc( - Req: jax.Array, + log_Req: jax.Array, args: tuple[jax.Array, float, float, jax.Array, jax.Array, jax.Array], ) -> jax.Array: """ - Mass balance residual for the heterogeneous-complex (polyc) binding model. + Mass balance residual in log-space for the heterogeneous-complex (polyc) + binding model. This is the root-finding target passed to the solver in :func:`polyc`; - it is zero when ``Req`` is the vector of free-receptor abundances + it is zero when ``log_Req`` is the natural logarithm of free-receptor abundances consistent with the mass balance for every receptor type. - :param Req: Current guess for the free receptor abundance per receptor + :param log_Req: Current guess for the log of free receptor abundance per receptor type. :param args: Tuple of ``(Rtot, L0, KxStar, Cplx, Ctheta, Kav)`` where ``Rtot`` is the total receptor abundance per receptor type, ``L0`` @@ -61,10 +62,11 @@ def Req_polyc( monomer composition of each ligand complex, ``Ctheta`` is the relative abundance of each complex, and ``Kav`` is the monomer ligand/receptor affinity matrix. - :return: The residual ``Req + Rbound - Rtot``, which the solver drives - to zero. + :return: The log-ratio residual ``log(Req + Rbound) - log(Rtot)``, which the solver + drives to zero. """ Rtot, L0, KxStar, Cplx, Ctheta, Kav = args + Req = jnp.exp(log_Req) Psi = Req * Kav * KxStar Psirs = Psi.sum(axis=1).reshape(-1, 1) + 1 Psinorm = Psi / Psirs @@ -79,7 +81,7 @@ def Req_polyc( axis=0, ) ) - return Req + Rbound - Rtot + return jnp.log(Req + Rbound) - jnp.log(Rtot) def commonChecks( @@ -159,12 +161,18 @@ def polyfc( A = jnp.dot(LigC.T, Kav) - # Find Phisum by fixed point iteration - solver = opt.LevenbergMarquardt(rtol=1e-9, atol=1e-9) + # Find Phisum by guaranteed bracketed bisection + solver = opt.Bisection(rtol=1e-12, atol=1e-12) + upper = jnp.maximum(jnp.dot(A * KxStar, Rtot.T), 1e-12) result = opt.root_find( - Req_polyfc, solver, y0=jnp.zeros(1), args=(Rtot, L0, KxStar, f, A), throw=True + Req_polyfc, + solver, + y0=jnp.array(0.0), + args=(Rtot, L0, KxStar, f, A), + options=dict(lower=jnp.array(0.0), upper=upper), + throw=True, ) - Phisum = result.value[0] + Phisum = result.value Lbound = L0 / KxStar * ((1 + Phisum) ** f - 1) Rbound = L0 / KxStar * f * Phisum * (1 + Phisum) ** (f - 1) @@ -182,24 +190,47 @@ def polyfc( return Lbound, Rbound, vieq, Rmulti_n -def Req_solve(func: Callable[..., jax.Array], Rtot: jax.Array, *args) -> jax.Array: +def Req_solve( + func: Callable[..., jax.Array], + Rtot: jax.Array, + L0: float, + KxStar: float, + Cplx: jax.Array, + Ctheta: jax.Array, + Kav: jax.Array, +) -> jax.Array: """ - Run Levenberg-Marquardt root finding to calculate the free receptor vector. - - :param func: Residual function to find the root of; called as - ``func(Req, (Rtot, *args))``. - :param Rtot: Total abundance of each receptor type on the cell; also - used as the shape template for the initial guess (zeros). - :param args: Additional positional arguments forwarded to ``func`` - after ``Rtot``. + Run Levenberg-Marquardt root finding in log-space to calculate the free + receptor vector. + + Initializes from an analytical 1:1 Langmuir binding approximation to ensure + rapid and robust convergence. + + :param func: Residual function to find the root of in log space; called as + ``func(log_Req, (Rtot, L0, KxStar, Cplx, Ctheta, Kav))``. + :param Rtot: Total abundance of each receptor type on the cell. + :param L0: Total ligand complex concentration. + :param KxStar: Detailed-balance corrected cross-linking constant. + :param Cplx: Monomer ligand composition of each complex. + :param Ctheta: Relative abundance of each complex. + :param Kav: Matrix of monomer ligand/receptor affinities. :return: The free receptor abundance vector ``Req`` that zeroes ``func``. """ - solver = opt.LevenbergMarquardt(rtol=1e-9, atol=1e-9) + L_monomer = jnp.dot(Ctheta, Cplx) * L0 + A_eff = jnp.dot(L_monomer, Kav) + Req_init = Rtot / (1.0 + A_eff) + log_Req_0 = jnp.log(jnp.maximum(Req_init, 1e-30)) + + solver = opt.LevenbergMarquardt(rtol=1e-10, atol=1e-10) result = opt.root_find( - func, solver, y0=jnp.zeros_like(Rtot), args=(Rtot, *args), throw=True + func, + solver, + y0=log_Req_0, + args=(Rtot, L0, KxStar, Cplx, Ctheta, Kav), + throw=True, ) - return result.value + return jnp.exp(result.value) def polyc( diff --git a/valentbind/test/test_edge_cases.py b/valentbind/test/test_edge_cases.py index 3bc6d42..6a5546a 100644 --- a/valentbind/test/test_edge_cases.py +++ b/valentbind/test/test_edge_cases.py @@ -56,3 +56,36 @@ def test_polyfc_Rbound_increases_with_Rtot() -> None: _, Rbound_high, _, _ = polyfc(L0, KxStar, f, np.array([1e5, 1e5]), LigC, Kav) assert float(Rbound_high) > float(Rbound_low) + + +def test_polyc_extreme_regime_convergence() -> None: + """Test convergence on steep multivalent binding with mixed complexes.""" + from ..model import polyc + + L0 = 1e-14 + KxStar = 1e-6 + Rtot = np.array([8e4, 7e4]) + Cplx = [[2, 3, 2, 1], [1, 1, 1, 1], [2, 2, 1, 3], [0, 1, 3, 0]] + Ctheta = [0.25, 0.25, 0.25, 0.25] + Kav = [[8e3, 9e5], [6e2, 1e3], [6e5, 6e3], [5e4, 9e1]] + + Lbound, Rbound, Lfbnd = polyc(L0, KxStar, Rtot, Cplx, Ctheta, Kav) + assert np.all(Lbound > 0.0) + assert np.all(Rbound > 0.0) + assert np.all(np.sum(Rbound, axis=0) <= Rtot * (1.0 + 1e-6)) + + +def test_polyfc_high_valency() -> None: + """Test polyfc convergence with high valency (f=12) and strong binding.""" + L0 = 1e-10 + KxStar = 1e-11 + f = 12 + Rtot = np.array([1e6]) + LigC = [1.0] + Kav = [[1e8]] + + Lbound, Rbound, vieq, Rmulti_n = polyfc(L0, KxStar, f, Rtot, LigC, Kav) + assert float(Lbound) > 0.0 + assert float(Rbound) > 0.0 + assert len(vieq) == 12 + np.testing.assert_allclose(float(Lbound), float(np.sum(vieq)), rtol=1e-6)