From c3f52d7748a61a84cfba6c6681907bca8ad10029 Mon Sep 17 00:00:00 2001 From: Vinicius Port Date: Wed, 5 Aug 2026 20:15:24 -0300 Subject: [PATCH 01/11] deps: Add diffrax --- pyproject.toml | 1 + uv.lock | 87 ++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 88 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index ccf8a355..07ffd0fd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -53,6 +53,7 @@ dependencies = [ [project.optional-dependencies] cli = ["rich>=13,<16"] fast = [ + "diffrax>=0.7,<0.8", "jax>=0.4", "jaxlib>=0.4", ] diff --git a/uv.lock b/uv.lock index 1436246c..b6a63122 100644 --- a/uv.lock +++ b/uv.lock @@ -623,6 +623,24 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/07/6c/aa3f2f849e01cb6a001cd8554a88d4c77c5c1a31c95bdf1cf9301e6d9ef4/defusedxml-0.7.1-py2.py3-none-any.whl", hash = "sha256:a352e7e428770286cc899e2542b6cdaedb2b4953ff269a210103ec58f6198a61", size = 25604, upload-time = "2021-03-08T10:59:24.45Z" }, ] +[[package]] +name = "diffrax" +version = "0.7.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "equinox" }, + { name = "jax" }, + { name = "jaxtyping" }, + { name = "lineax" }, + { name = "optimistix" }, + { name = "typing-extensions" }, + { name = "wadler-lindig" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/5a/ea/a22853377a840d70d1d174d43183f12990c66eecf34ba7e76eec781249d0/diffrax-0.7.2.tar.gz", hash = "sha256:9adc70fe90dba83f7dd839845839b2531a3dd736a3944fe1aa26598507cfb48e", size = 155159, upload-time = "2026-02-18T01:14:04.964Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/6c/48/c42f657f3d6b4d9d9674c9bc48e763d024511ea4cc1d6e0c68c7b4380c43/diffrax-0.7.2-py3-none-any.whl", hash = "sha256:1beee2f991cd049962fe3c89da28f933753aa66d80bef857533276c9f23301ed", size = 199673, upload-time = "2026-02-18T01:14:03.38Z" }, +] + [[package]] name = "dill" version = "0.4.1" @@ -632,6 +650,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/1e/77/dc8c558f7593132cf8fefec57c4f60c83b16941c574ac5f619abb3ae7933/dill-0.4.1-py3-none-any.whl", hash = "sha256:1e1ce33e978ae97fcfcff5638477032b801c46c7c65cf717f95fbc2248f79a9d", size = 120019, upload-time = "2026-01-19T02:36:55.663Z" }, ] +[[package]] +name = "equinox" +version = "0.13.8" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "jax" }, + { name = "jaxtyping" }, + { name = "typing-extensions" }, + { name = "wadler-lindig" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/19/ff/522336d2f8264f2ad97119710b76e2cddf66145d03a1e89899175d26b192/equinox-0.13.8.tar.gz", hash = "sha256:dd075050018e2dd02e252e9d29d3060f7e67f085622d8d27a8e89e24bb8523db", size = 145257, upload-time = "2026-05-05T10:03:43.258Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ae/d6/69a76c8ccdef14af687c497040292a46e59fc7a0ab24724b60e50ca61030/equinox-0.13.8-py3-none-any.whl", hash = "sha256:ca004348533cc30a63ebe8823d7dd4bb626dce17743d40bbddb89b402ef2a240", size = 185813, upload-time = "2026-05-05T10:03:41.673Z" }, +] + [[package]] name = "executing" version = "2.2.1" @@ -923,6 +956,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c9/6a/6dff1e6e3f9d918bc777e087091bdefbd7d33328c1d1b152429c6cdcf723/jaxlib-0.9.2-cp314-cp314t-manylinux_2_27_x86_64.whl", hash = "sha256:bbe59bdef668ff5fd998c6d88e8df9a32ab95bec0dea3d2b5f7a11b86a9a6788", size = 83425685, upload-time = "2026-03-18T23:28:07.906Z" }, ] +[[package]] +name = "jaxtyping" +version = "0.3.11" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "wadler-lindig" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/91/c1/091b8852bd7cbf50bd655543c8506033cf4029300c67f8c176c1286879a9/jaxtyping-0.3.11.tar.gz", hash = "sha256:b09c14acf6686feb9e0df5b0d8c6e7c5b6f8d36bf059ee54cd522a186c2ef050", size = 46489, upload-time = "2026-06-13T18:35:23.167Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8b/38/c66bbdc5047f4776c2bd3e47e5295a350e3fa44d5b8942105e71c2a876a0/jaxtyping-0.3.11-py3-none-any.whl", hash = "sha256:8a4bedc4e3f963fa82df41bd13c7ebc2bad925601eb48614c65798f21329d4e3", size = 56593, upload-time = "2026-06-13T18:35:22.01Z" }, +] + [[package]] name = "jedi" version = "0.19.2" @@ -1320,6 +1365,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/82/3d/14ce75ef66813643812f3093ab17e46d3a206942ce7376d31ec2d36229e7/lark-1.3.1-py3-none-any.whl", hash = "sha256:c629b661023a014c37da873b4ff58a817398d12635d3bbb2c5a03be7fe5d1e12", size = 113151, upload-time = "2025-10-27T18:25:54.882Z" }, ] +[[package]] +name = "lineax" +version = "0.1.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "equinox" }, + { name = "jax" }, + { name = "jaxtyping" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/35/d6/4e28416a6fe58dd6bc7565b1ffa330f4d0ba7d74212642b1b734c511299e/lineax-0.1.0.tar.gz", hash = "sha256:5f1a8f060142af2cdbf7d66b99e8d3071c3aa734b677df6339df4b4c4c0554d2", size = 50209, upload-time = "2026-01-27T21:17:26.652Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/80/0c/2ed47112fc1958a0a81c9b015d4e1861953a1ec3a17b081c0180a25ce82c/lineax-0.1.0-py3-none-any.whl", hash = "sha256:f00911c6b07d427c4835db46856970c8348bc82a035b51f4386ad09382af957a", size = 74600, upload-time = "2026-01-27T21:17:25.33Z" }, +] + [[package]] name = "markdown-it-py" version = "4.0.0" @@ -1770,6 +1830,22 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/23/cd/066e86230ae37ed0be70aae89aabf03ca8d9f39c8aea0dec8029455b5540/opt_einsum-3.4.0-py3-none-any.whl", hash = "sha256:69bb92469f86a1565195ece4ac0323943e83477171b91d24c35afe028a90d7cd", size = 71932, upload-time = "2024-09-26T14:33:23.039Z" }, ] +[[package]] +name = "optimistix" +version = "0.1.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "equinox" }, + { name = "jax" }, + { name = "jaxtyping" }, + { name = "lineax" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/5c/0b/3bdc0e698cb0d264e5dca0b5bb171342a950a43fa6e046d8c57189298422/optimistix-0.1.0.tar.gz", hash = "sha256:f05c9104748e87e1dc10a0b4a2be94e427f98c3eab0b1d41ea24a593d76be03d", size = 70164, upload-time = "2026-02-16T13:35:43.991Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/63/53/6a81a6ebb1739f19c32002139719d4ee021a349d5bdbe066910feed327a1/optimistix-0.1.0-py3-none-any.whl", hash = "sha256:c8edda7bf7fe48839c93fcd72bce750c950ed9f7033158303150b7ab395add81", size = 101740, upload-time = "2026-02-16T13:35:42.971Z" }, +] + [[package]] name = "overreact" version = "1.2.0" @@ -1785,6 +1861,7 @@ cli = [ { name = "rich" }, ] fast = [ + { name = "diffrax" }, { name = "jax" }, { name = "jaxlib" }, ] @@ -1814,6 +1891,7 @@ dev = [ [package.metadata] requires-dist = [ { name = "cclib", specifier = ">=1,<2" }, + { name = "diffrax", marker = "extra == 'fast'", specifier = ">=0.7,<0.8" }, { name = "importlib", specifier = ">=1.0.4,<2" }, { name = "jax", marker = "extra == 'fast'", specifier = ">=0.4" }, { name = "jaxlib", marker = "extra == 'fast'", specifier = ">=0.4" }, @@ -2907,6 +2985,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/7f/3e/5db95bcf282c52709639744ca2a8b149baccf648e39c8cc87553df9eae0c/urllib3-2.7.0-py3-none-any.whl", hash = "sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897", size = 131087, upload-time = "2026-05-07T16:13:17.151Z" }, ] +[[package]] +name = "wadler-lindig" +version = "0.1.7" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/1e/67/cbae4bf7683a64755c2c1778c418fea96d00e34395bb91743f08bd951571/wadler_lindig-0.1.7.tar.gz", hash = "sha256:81d14d3fe77d441acf3ebd7f4aefac20c74128bf460e84b512806dccf7b2cd55", size = 15842, upload-time = "2025-06-18T07:00:42.843Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8d/96/04e7b441807b26b794da5b11e59ed7f83b2cf8af202bd7eba8ad2fa6046e/wadler_lindig-0.1.7-py3-none-any.whl", hash = "sha256:e3ec83835570fd0a9509f969162aeb9c65618f998b1f42918cfc8d45122fe953", size = 20516, upload-time = "2025-06-18T07:00:41.684Z" }, +] + [[package]] name = "wcwidth" version = "0.6.0" From 5f4198d87217ba7cf74ac88cff251b11cf30009c Mon Sep 17 00:00:00 2001 From: Vinicius Port Date: Thu, 6 Aug 2026 19:17:10 -0300 Subject: [PATCH 02/11] feat: add an initial draft for diffrax integration Initial idea for integrating diffrax, currently dispatches to a _get_y_diffrax function if the selected method is one of the provided by this integration --- overreact/simulate.py | 112 +++++++++++++++++++++++++++++++++--------- 1 file changed, 88 insertions(+), 24 deletions(-) diff --git a/overreact/simulate.py b/overreact/simulate.py index 1cfb2b89..161cdd34 100644 --- a/overreact/simulate.py +++ b/overreact/simulate.py @@ -34,6 +34,8 @@ # TODO(schneiderfelipe): this should probably be exposed to the user and use the actual simulation temperature. EF = np.exp(1.25 * constants.kcal / (constants.R * 298.15)) +_DIFFRAX_METHODS = ("Kvaerno3", "Kvaerno4", "Kvaerno5") + logger = logging.getLogger(__name__) @@ -58,7 +60,7 @@ def get_y( t_span=None, method="RK23", max_step=np.inf, - first_step=np.finfo(np.float64).eps, + first_step=None, rtol=1e-3, atol=1e-6, max_time=1 * 60 * 60, @@ -89,9 +91,9 @@ def get_y( Maximum step to be performed by the integrator. Defaults to half the total time span. first_step : float, optional - First step size. - Defaults to half the maximum step, or `np.finfo(np.float64).eps`, - whichever is smallest. + First step size. If not given, Diffrax chooses one automatically, + while the SciPy backend uses `np.finfo(np.float64).eps` for backwards + compatibility. rtol, atol : array-like, optional See `scipy.integrate.solve_ivp` for details. max_time : float, optional @@ -160,28 +162,44 @@ def get_y( max_step = np.min([max_step, (t_span[1] - t_span[0]) / 2.0]) logger.warning(f"max step = {max_step} s") - first_step = np.min([first_step, max_step / 2.0]) + if first_step is not None: + first_step = np.min([first_step, max_step / 2.0]) logger.warning(f"first step = {first_step} s") - jac = None - if hasattr(dydt, "jac"): - jac = dydt.jac # noqa: F841 - - logger.warning(f"@t = \x1b[94m{0:10.3f} \x1b[ms\x1b[K") - res = solve_ivp( - dydt, - t_span, - y0, - method=method, - dense_output=True, - max_step=max_step, - first_step=first_step, - rtol=rtol, - atol=atol, - # jac=jac, # noqa: ERA001 - ) - logger.warning(res) - y = res.sol + if method in _DIFFRAX_METHODS: + y = _get_y_diffrax( + dydt, + y0, + t_span, + method, + max_step, + first_step, + rtol, + atol, + ) + else: + if first_step is None: + first_step = np.finfo(np.float64).eps + + jac = None + if hasattr(dydt, "jac"): + jac = dydt.jac # noqa: F841 + + logger.warning(f"@t = \x1b[94m{0:10.3f} \x1b[ms\x1b[K") + res = solve_ivp( + dydt, + t_span, + y0, + method=method, + dense_output=True, + max_step=max_step, + first_step=first_step, + rtol=rtol, + atol=atol, + # jac=jac, # noqa: ERA001 + ) + logger.warning(res) + y = res.sol def r(t): # TODO(schneiderfelipe): this is probably not the best way to @@ -194,6 +212,52 @@ def r(t): return y, r +def _get_y_diffrax(dydt, y0, t_span, method, max_step, first_step, rtol, atol): + """Solve an initial value problem with a Diffrax stiff solver.""" + try: + import diffrax + import jax + import jax.numpy as jnp + except ImportError as exc: + msg = ( + f"the {method} solver requires Diffrax; " + 'install it with `pip install "overreact[fast]"`' + ) + raise ImportError(msg) from exc + + solver = { + "Kvaerno3": diffrax.Kvaerno3, + "Kvaerno4": diffrax.Kvaerno4, + "Kvaerno5": diffrax.Kvaerno5, + }[method]() + term = diffrax.ODETerm(lambda t, y, _args: dydt(t, y)) + stepsize_controller = diffrax.PIDController( + rtol=rtol, + atol=atol, + dtmax=max_step, + ) + solution = diffrax.diffeqsolve( + term, + solver, + t0=t_span[0], + t1=t_span[1], + dt0=first_step, + y0=jnp.asarray(y0), + saveat=diffrax.SaveAt(dense=True), + stepsize_controller=stepsize_controller, + ) + + def y(t): + if np.ndim(t) == 0: + return np.asarray(solution.evaluate(t)) + values = jax.vmap(solution.evaluate)(jnp.asarray(t)) + return np.asarray(values).T + + y.t_min = float(solution.t0) + y.t_max = float(solution.t1) + return y + + def get_dydt(scheme, k, ef=EF): """Generate a rate function that models a reaction scheme. From 6d8bc697b9f869797f21881e1506bc12fff3ee75 Mon Sep 17 00:00:00 2001 From: Vinicius Port Date: Thu, 6 Aug 2026 19:35:56 -0300 Subject: [PATCH 03/11] feat: add new methods to the CLI, update documentation. --- overreact/_cli.py | 17 +++++++++++++++-- overreact/simulate.py | 19 +++++++++++++------ 2 files changed, 28 insertions(+), 8 deletions(-) diff --git a/overreact/_cli.py b/overreact/_cli.py index 112cb54b..0dd78849 100644 --- a/overreact/_cli.py +++ b/overreact/_cli.py @@ -852,8 +852,21 @@ def main(arguments=None): ) parser.add_argument( "--method", - help="integrator used in solving the ODE system of the microkinetic simulation", - choices=["RK23", "DOP853", "RK45", "LSODA", "BDF", "Radau"], + help=( + "integrator used in solving the ODE system of the microkinetic " + "simulation (Kvaerno methods require overreact[fast])" + ), + choices=[ + "RK23", + "DOP853", + "RK45", + "LSODA", + "BDF", + "Radau", + "Kvaerno3", + "Kvaerno4", + "Kvaerno5", + ], default="RK23", ) parser.add_argument( diff --git a/overreact/simulate.py b/overreact/simulate.py index 161cdd34..a9cd293e 100644 --- a/overreact/simulate.py +++ b/overreact/simulate.py @@ -69,8 +69,8 @@ def get_y( This function provides two functions that calculate the concentrations and the rates of formation at any point in time for any compound. It does that - by solving an initial value problem (IVP) through scipy's ``solve_ivp`` - under the hood. + by solving an initial value problem (IVP) through scipy's ``solve_ivp`` or + Diffrax's ``diffeqsolve`` under the hood. Parameters ---------- @@ -84,9 +84,12 @@ def get_y( is chosen based on the system at hand (the method of choice works for any zeroth-, first- or second-order reactions). method : str, optional - Integration method to use. See `scipy.integrate.solve_ivp` for details. - Kinetics problems are very often stiff and, as such, "RK23" and "RK45" may be - unsuited. "LSODA", "BDF", and "Radau" are worth a try if things go bad. + Integration method to use. All existing methods are + provided by `scipy.integrate.solve_ivp`, except for "Kvaerno3", + "Kvaerno4", and "Kvaerno5", which use Diffrax instead. + Kinetics problems are very often stiff and, as such, + "RK23" and "RK45" may be unsuited. "LSODA", "BDF", "Radau", and the + Kvaerno methods are worth trying for stiff systems. max_step : float, optional Maximum step to be performed by the integrator. Defaults to half the total time span. @@ -104,8 +107,12 @@ def get_y( ------- y, r : callable Concentrations and reaction rates as functions of time. The y object - is an OdeSolution and stores attributes t_min and t_max. + stores attributes t_min and t_max. + Notes + ----- + Diffrax's implicit Kvaerno solvers use adaptive step sizes controlled by + ``rtol`` and ``atol``. Examples -------- From e4e1c7dfbed9419f354ccf3b1c8a41a269556cfd Mon Sep 17 00:00:00 2001 From: Vinicius Port Date: Thu, 6 Aug 2026 22:20:52 -0300 Subject: [PATCH 04/11] tests: test new diffrax methods --- tests/test_simulate.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_simulate.py b/tests/test_simulate.py index d71bea19..6683524b 100644 --- a/tests/test_simulate.py +++ b/tests/test_simulate.py @@ -118,7 +118,7 @@ def test_get_y_conservation_in_equilibria() -> None: for sub0 in [0.01, 0.02] for keq in [1.0, 10.0, 100.0] for kcat in [1e-1, 1e1, 1e10, 1e11, 1e13] - for method in ("RK23", "LSODA", "Radau", "BDF") + for method in ("RK23", "LSODA", "Radau", "BDF", "Kvaerno5") ], ) def test_simple_michaelis_menten( @@ -166,7 +166,7 @@ def test_simple_michaelis_menten( for sub0 in [0.01, 0.02] for keq in [1.0, 10.0, 100.0] for kcat in [1e-1, 1e1, 1e10, 1e11, 1e13] - for method in ("RK23", "LSODA", "Radau", "BDF") + for method in ("RK23", "LSODA", "Radau", "BDF", "Kvaerno5") ], ) def test_consuming_michaelis_menten( From f76af8795089f9a7230ddfa3294ddb4746e1c7b5 Mon Sep 17 00:00:00 2001 From: Vinicius Port Date: Thu, 6 Aug 2026 22:23:41 -0300 Subject: [PATCH 05/11] fix(diffrax): no more NaN jacobians --- overreact/simulate.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/overreact/simulate.py b/overreact/simulate.py index a9cd293e..067076b5 100644 --- a/overreact/simulate.py +++ b/overreact/simulate.py @@ -112,7 +112,7 @@ def get_y( Notes ----- Diffrax's implicit Kvaerno solvers use adaptive step sizes controlled by - ``rtol`` and ``atol``. + ``rtol`` and ``atol``. Examples -------- @@ -331,7 +331,10 @@ def get_dydt(scheme, k, ef=EF): k_adj = _adjust_k(scheme, k, ef=ef) def _dydt(_t, y): - r = k_adj * jnp.prod(jnp.power(y, M), axis=1) + # Avoid differentiating 0**0 for compounds that do not participate in + # a reaction, this causes NaN fileed jacobians in diffrax otherwise. + bases = jnp.where(M == 0, 1.0, y) + r = k_adj * jnp.prod(jnp.power(bases, M), axis=1) return jnp.dot(A, r) if _found_jax: From 73b61700987fdb99e7b40fb2320f971b21de97f5 Mon Sep 17 00:00:00 2001 From: Vinicius Port Date: Thu, 27 Aug 2026 16:24:27 -0300 Subject: [PATCH 06/11] fix: create methods from overreact.simulate Uses the constants defined in overreact.simulate to sync the possible choices with the simulate module. Improves the error message to also print the available scipy solvers that don't require jax. A supported_methods constant also is exposed by the module. --- overreact/_cli.py | 12 +----------- overreact/simulate.py | 16 ++++++++++++++-- 2 files changed, 15 insertions(+), 13 deletions(-) diff --git a/overreact/_cli.py b/overreact/_cli.py index 0dd78849..c2f73e6b 100644 --- a/overreact/_cli.py +++ b/overreact/_cli.py @@ -856,17 +856,7 @@ def main(arguments=None): "integrator used in solving the ODE system of the microkinetic " "simulation (Kvaerno methods require overreact[fast])" ), - choices=[ - "RK23", - "DOP853", - "RK45", - "LSODA", - "BDF", - "Radau", - "Kvaerno3", - "Kvaerno4", - "Kvaerno5", - ], + choices=rx.simulate.SUPPORTED_SOLVERS, default="RK23", ) parser.add_argument( diff --git a/overreact/simulate.py b/overreact/simulate.py index 067076b5..6ac2f572 100644 --- a/overreact/simulate.py +++ b/overreact/simulate.py @@ -34,7 +34,17 @@ # TODO(schneiderfelipe): this should probably be exposed to the user and use the actual simulation temperature. EF = np.exp(1.25 * constants.kcal / (constants.R * 298.15)) -_DIFFRAX_METHODS = ("Kvaerno3", "Kvaerno4", "Kvaerno5") +_SCIPY_SOLVERS = ( + "RK23", + "DOP853", + "RK45", + "LSODA", + "BDF", + "Radau", +) +_DIFFRAX_SOLVERS = ("Kvaerno3", "Kvaerno4", "Kvaerno5") + +SUPPORTED_SOLVERS = _SCIPY_SOLVERS + _DIFFRAX_SOLVERS logger = logging.getLogger(__name__) @@ -173,7 +183,7 @@ def get_y( first_step = np.min([first_step, max_step / 2.0]) logger.warning(f"first step = {first_step} s") - if method in _DIFFRAX_METHODS: + if method in _DIFFRAX_SOLVERS: y = _get_y_diffrax( dydt, y0, @@ -229,6 +239,8 @@ def _get_y_diffrax(dydt, y0, t_span, method, max_step, first_step, rtol, atol): msg = ( f"the {method} solver requires Diffrax; " 'install it with `pip install "overreact[fast]"`' + "or choose one of the following solvers: " + f"{', '.join(_SCIPY_SOLVERS)}" ) raise ImportError(msg) from exc From e6148a656666c0cf8a0bb1d056a07d6a3847d345 Mon Sep 17 00:00:00 2001 From: Vinicius Port Date: Thu, 27 Aug 2026 17:22:16 -0300 Subject: [PATCH 07/11] Clear the first step mess --- overreact/simulate.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/overreact/simulate.py b/overreact/simulate.py index 6ac2f572..fa536d48 100644 --- a/overreact/simulate.py +++ b/overreact/simulate.py @@ -179,10 +179,6 @@ def get_y( max_step = np.min([max_step, (t_span[1] - t_span[0]) / 2.0]) logger.warning(f"max step = {max_step} s") - if first_step is not None: - first_step = np.min([first_step, max_step / 2.0]) - logger.warning(f"first step = {first_step} s") - if method in _DIFFRAX_SOLVERS: y = _get_y_diffrax( dydt, @@ -196,7 +192,8 @@ def get_y( ) else: if first_step is None: - first_step = np.finfo(np.float64).eps + first_step = np.min([np.finfo(np.float64).eps, max_step / 2.0]) + logger.warning(f"first step = {first_step} s") jac = None if hasattr(dydt, "jac"): @@ -244,6 +241,11 @@ def _get_y_diffrax(dydt, y0, t_span, method, max_step, first_step, rtol, atol): ) raise ImportError(msg) from exc + if first_step is not None: + logger.warning(f"first step = {first_step} s") + else: + logger.warning("no first step given, diffrax will choose automatically") + solver = { "Kvaerno3": diffrax.Kvaerno3, "Kvaerno4": diffrax.Kvaerno4, From 0045781140fd17a0804e05fe0a8c0cd49c189fb3 Mon Sep 17 00:00:00 2001 From: Vinicius Port Date: Thu, 27 Aug 2026 17:22:50 -0300 Subject: [PATCH 08/11] Add defaults for get_y_diffrax --- overreact/simulate.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/overreact/simulate.py b/overreact/simulate.py index fa536d48..50fa9629 100644 --- a/overreact/simulate.py +++ b/overreact/simulate.py @@ -226,7 +226,16 @@ def r(t): return y, r -def _get_y_diffrax(dydt, y0, t_span, method, max_step, first_step, rtol, atol): +def _get_y_diffrax( + dydt, + y0, + t_span, + method="Kvaerno3", + max_step=np.inf, + first_step=None, + rtol=1e-3, + atol=1e-6, +): """Solve an initial value problem with a Diffrax stiff solver.""" try: import diffrax From 45293ffe17fb2ad1c243ff2d9d138a3d2becaec1 Mon Sep 17 00:00:00 2001 From: Vinicius Port Date: Thu, 27 Aug 2026 17:25:35 -0300 Subject: [PATCH 09/11] Allow jax in scipy's solve_ivp. After avoiding differentiating 0**0 we basically solved the issue where we couldn't use jax in solve_ivp. --- overreact/simulate.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/overreact/simulate.py b/overreact/simulate.py index 50fa9629..26f0ae2a 100644 --- a/overreact/simulate.py +++ b/overreact/simulate.py @@ -197,7 +197,7 @@ def get_y( jac = None if hasattr(dydt, "jac"): - jac = dydt.jac # noqa: F841 + jac = dydt.jac logger.warning(f"@t = \x1b[94m{0:10.3f} \x1b[ms\x1b[K") res = solve_ivp( @@ -210,7 +210,7 @@ def get_y( first_step=first_step, rtol=rtol, atol=atol, - # jac=jac, # noqa: ERA001 + jac=jac, ) logger.warning(res) y = res.sol @@ -355,7 +355,7 @@ def get_dydt(scheme, k, ef=EF): def _dydt(_t, y): # Avoid differentiating 0**0 for compounds that do not participate in - # a reaction, this causes NaN fileed jacobians in diffrax otherwise. + # a reaction, this causes NaN filled jacobians in jax otherwise. bases = jnp.where(M == 0, 1.0, y) r = k_adj * jnp.prod(jnp.power(bases, M), axis=1) return jnp.dot(A, r) From 378c54dd82c5b8c8dae69b6b95372429197ad8a5 Mon Sep 17 00:00:00 2001 From: Vinicius Port Date: Thu, 27 Aug 2026 17:37:56 -0300 Subject: [PATCH 10/11] Remove jax and jaxlib from pyproject.toml --- pyproject.toml | 2 -- 1 file changed, 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 07ffd0fd..09918891 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -54,8 +54,6 @@ dependencies = [ cli = ["rich>=13,<16"] fast = [ "diffrax>=0.7,<0.8", - "jax>=0.4", - "jaxlib>=0.4", ] solvents = ["thermo>=0.2"] From c787b008e9257e23908655e45f5632dd4be08f7c Mon Sep 17 00:00:00 2001 From: Vinicius Port Date: Mon, 31 Aug 2026 16:27:30 -0300 Subject: [PATCH 11/11] Update uv.lock --- uv.lock | 20 ++++++++------------ 1 file changed, 8 insertions(+), 12 deletions(-) diff --git a/uv.lock b/uv.lock index b6a63122..7feb1b01 100644 --- a/uv.lock +++ b/uv.lock @@ -1862,8 +1862,6 @@ cli = [ ] fast = [ { name = "diffrax" }, - { name = "jax" }, - { name = "jaxlib" }, ] solvents = [ { name = "thermo" }, @@ -1893,9 +1891,7 @@ requires-dist = [ { name = "cclib", specifier = ">=1,<2" }, { name = "diffrax", marker = "extra == 'fast'", specifier = ">=0.7,<0.8" }, { name = "importlib", specifier = ">=1.0.4,<2" }, - { name = "jax", marker = "extra == 'fast'", specifier = ">=0.4" }, - { name = "jaxlib", marker = "extra == 'fast'", specifier = ">=0.4" }, - { name = "rich", marker = "extra == 'cli'", specifier = ">=13,<15" }, + { name = "rich", marker = "extra == 'cli'", specifier = ">=13,<16" }, { name = "scipy", specifier = ">=1.10,<2" }, { name = "thermo", marker = "extra == 'solvents'", specifier = ">=0.2" }, ] @@ -1905,19 +1901,19 @@ provides-extras = ["cli", "fast", "solvents"] dev = [ { name = "debugpy", specifier = ">=1,<2" }, { name = "flynt", specifier = ">=0.77,<1.1" }, - { name = "ipython", specifier = ">=8,<9" }, + { name = "ipython", specifier = ">=8,<10" }, { name = "jupyter", specifier = ">=1.0.0,<2" }, { name = "matplotlib", specifier = ">=3,<4" }, { name = "mypy", specifier = ">=0.991,<2.4" }, - { name = "pdoc", specifier = ">=12,<15" }, + { name = "pdoc", specifier = ">=12,<17" }, { name = "perflint", specifier = ">=0.7.1,<0.9.0" }, { name = "pytest", specifier = ">=7.2,<10.0" }, - { name = "pytest-cov", specifier = ">=4,<6" }, - { name = "rich", specifier = ">=13,<15" }, - { name = "ruff", specifier = ">=0.0.210,<0.16.3" }, + { name = "pytest-cov", specifier = ">=4,<8" }, + { name = "rich", specifier = ">=13,<16" }, + { name = "ruff", specifier = ">=0.0.210,<0.16.5" }, { name = "seaborn", specifier = ">=0.12,<0.14" }, { name = "thermo", specifier = ">=0.2" }, - { name = "types-setuptools", specifier = ">=65,<84" }, + { name = "types-setuptools", specifier = ">=65,<85" }, ] [[package]] @@ -2059,7 +2055,7 @@ name = "pexpect" version = "4.9.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "ptyprocess" }, + { name = "ptyprocess", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/42/92/cc564bf6381ff43ce1f4d06852fc19a2f11d180f23dc32d9588bee2f149d/pexpect-4.9.0.tar.gz", hash = "sha256:ee7d41123f3c9911050ea2c2dac107568dc43b2d3b0c7557a33212c398ead30f", size = 166450, upload-time = "2023-11-25T09:07:26.339Z" } wheels = [