Improve root-finding convergence with bisection and log-scale formulation

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(

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)