jabsim 0.1.0__tar.gz

jabsim-0.1.0/LICENCE

@@ -0,0 +1,21 @@
+MIT LICENCE
+
+Copyright © 2026 Kirill Sechkar
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the “Software”), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

jabsim-0.1.0/PKG-INFO

@@ -0,0 +1,132 @@
+Metadata-Version: 2.4
+Name: jabsim
+Version: 0.1.0
+Summary: Simple ODE simulation with JAX for biological systems
+Author-email: Kirill Sechkar <useful.instructive@proton.me>
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENCE
+Requires-Dist: jax>=0.4
+Requires-Dist: numpy>=1.24
+Requires-Dist: scipy>=1.10
+Provides-Extra: test
+Requires-Dist: pytest>=7; extra == "test"
+Dynamic: license-file
+
+# jabsim
+
+A jax-based package for simulating ODE models of biological systems,
+where **all variables are non-negative**. This enforcement of non-negativity, which neither `scipy.solve_ivp` nor `diffrax`
+solvers can give you, is why you may need jabsim.
+This is achieved by clamping the state variables to zero whenever they are negative 
+
+jabsim is powered by the [jax](https://github.com/jax-ml/jax) package for high-performance computing and parallelisation.
+This means that jabsim simulations can be jit-compiled and parallelised on a GPU or TPU using `jax.vmap` or shard mapping.
+
+## How to use jabsim
+
+1. Make an ODE function which calculates the derivative $\frac{dx}{dt}$ from the arguments `t, x, par` in this order.
+   - What do the arguments stand for?     
+     - `t` is the time at this point in the simulation
+     - `x` is the state vector at this time point
+     - `args` is a tuple of extra arguments passed on to the ODE function.
+   - **IMPORTANT**: if you don't know jax, there are a few differences to keep in mind:
+     - If you normally use `numpy` functions in your ODE, import `jax.numpy` as `jnp` and use it instead of `np`. jax with `jax.numpy` will get automatically installed as a dependency of jabsim when you install it.
+     - Be careful using loops and if-statements. If you're an amateur programmer, just avoid doing all that. Otherwise, have a look at the [JAX documentation](https://docs.jax.dev/en/latest/index.html) .
+2. Import `jabsim` and call `jabsim.sim` to simulate. The arguments are as follows:
+   - `par`: a list, array or dict of model parameters as in your ODE function
+   - `model_ode`: the ODE function you created
+   - `x0`: the initial state vector as a 1D `np.array` or `jnp.array`
+   - `tf`: tuple, array or list. The ODE will be simulated 
+   - `savetimestep`: interval between the time points at which the trajectory is saved
+   - `simulator`: string specifyingthe simulation method to use
+        - `"euler"`: Euler simulator.
+        - `"rk4"`: Runge-Kutta 4th order simulator. Slower per ODE integration step but more accurate, hence allowing larger steps for the same accuracy.
+   - `ode_steps_in_savetimestep`: number of ODE integration steps within one timestep
+        - e.g. if `savetimestep=0.5` hours and `ode_step_in_savetimesteps=100`, there will be 100 integration steps per 0.5 hour, so the ODE integration step size will be 0.5/100=0.05 hours.
+        - high `ode_steps_in_savetimestep` number increases accuracy but increases runtimes
+   - `return_numpy`: if `True` (by default, it is) the output will be in the `np.array` format, otherwise it will be `jnp.array`.
+3. Running `jabsim.sim()` will return the arrays `ts` and `xs` as `np.array` or `jnp.array`, as well as a boolean value `success`.
+   - `ts`: array of timepoints between `tf[0]` and `tf[1]` with `savetimestep` hours, seconds or whatever units you are using between each two consecutive point
+   - `xs`: system trajectory saved as an array at the time points in `ts` - axis 0 for time, axis 1 for entries in the state vector (i.e. `xs.shape[0]=len(ts)`).
+   - `success`: boolean value; `True` if no entry in `xs` is `nan` or `inf`, `False` otherwise.
+
+## Notes
+- In practice, 500 steps per hour (e.g. `savetimestep=0.5, ode_steps_in_savetimestep=250` or `savetimestep=1.0, ode_steps_in_savetimestep=500`) works well for the RK4 solver. For the Euler solver, 1e4 steps per hours is reasonably good.
+- For benchmarking, you can also set `simulator="scipy"` to simulate your ODE with `scipy.solve_ivp` (but without any of the delicious jax features of the solvers above). In that case, don't use the arguments `ode_steps_in_savetimestep` and `savetimestep`. Instead, you can *optionally* specify:
+  - `solver`: string describing any solver which may be used with `scipy.solve_ivp`. By default, we have `solver="LSODA"`.
+  - `tols`: dictionary of relative and absolute tolerances for the scipy solver. By default, `tols={'rtol': 1e-6, 'atol': 1e-9}`.
+  - `dt0`: starting integration step size. By default, `dt0=0.1`.
+- If you want to make use of jax parallelisation, make sure to set `return_numpy=False` so that the solver would operate with `jnp.array` objects only.
+
+## Example
+
+Let us integrate a simple one-dimensional ODE $\frac{dx}{dt} = a x^2$. For the initial condition $x_0=1$ and $a=0.4$, 
+this has the analytical solution $x = \frac{1}{1-0.4t}$. This means we can verify that for  `savetimestep=0.5`, 
+`jabsim.sim()` produces `ts=np.array([0, 0.5, 1.0])` and `xs=np.array([1.0, 1.25, 1.66666667])`. 
+All entries in `xs` are finite, hence `success=True`.
+
+```python
+# import jabsim
+import jabsim
+
+# import jax.numpy for numpy operations
+import jax.numpy as jnp
+
+# our model ODE function returning a list of one element
+def model_ode(t, x, args):
+    # unpack args - get the dictiory of parameters
+    par, = args
+    
+    # use jnp to square x 
+    # (here you could just as well use x[0]**2, we just want to make a point)
+    x_squared = jnp.square(x[0])
+    
+    # return dx/dt as a list - with one entry for a one-dimensional ODE
+    return [par['a'] * x_squared]
+
+# our dictionary of paramneters
+par = {'a': 0.4}
+
+ts, xs, success = jabsim.sim(
+    model_ode=model_ode,
+    args=(par,),
+    x0=jnp.array([1.0]),
+    tf=(0.0, 1.0),
+    savetimestep=0.5,
+    simulator='rk4',
+    ode_steps_in_savetimestep=10,
+)
+
+# print the timne
+print(ts)
+print(xs)
+print(success)
+```
+
+# Citation
+If you find this package useful in your work, please cite the paper below:
+code for its Showcase 2 served as jabsim's direct ideological precursor.
+```bibtex
+@article{Gallup2024,
+	author = {Gallup, Olivia and Sechkar, Kirill and Towers, Sebastian and Steel, Harrison},
+	title = {Computational Synthetic Biology Enabled through JAX: A Showcase},
+	journal = {ACS Synth. Biol.},
+	volume = {13},
+	number = {9},
+	pages = {3046},
+	year = {2024},
+	doi = {10.1021/acssynbio.4c00307}
+}
+```
+
+The original JAX package should be cited as:
+```bibtex
+@software{jax2018github,
+  author = {James Bradbury and Roy Frostig and Peter Hawkins and Matthew James Johnson and Yash Katariya and Chris Leary and Dougal Maclaurin and George Necula and Adam Paszke and Jake Vander{P}las and Skye Wanderman-{M}ilne and Qiao Zhang},
+  title = {{JAX}: composable transformations of {P}ython+{N}um{P}y programs},
+  url = {http://github.com/jax-ml/jax},
+  version = {0.3.13},
+  year = {2018},
+}
+```

jabsim-0.1.0/README.md

@@ -0,0 +1,117 @@
+# jabsim
+
+A jax-based package for simulating ODE models of biological systems,
+where **all variables are non-negative**. This enforcement of non-negativity, which neither `scipy.solve_ivp` nor `diffrax`
+solvers can give you, is why you may need jabsim.
+This is achieved by clamping the state variables to zero whenever they are negative 
+
+jabsim is powered by the [jax](https://github.com/jax-ml/jax) package for high-performance computing and parallelisation.
+This means that jabsim simulations can be jit-compiled and parallelised on a GPU or TPU using `jax.vmap` or shard mapping.
+
+## How to use jabsim
+
+1. Make an ODE function which calculates the derivative $\frac{dx}{dt}$ from the arguments `t, x, par` in this order.
+   - What do the arguments stand for?     
+     - `t` is the time at this point in the simulation
+     - `x` is the state vector at this time point
+     - `args` is a tuple of extra arguments passed on to the ODE function.
+   - **IMPORTANT**: if you don't know jax, there are a few differences to keep in mind:
+     - If you normally use `numpy` functions in your ODE, import `jax.numpy` as `jnp` and use it instead of `np`. jax with `jax.numpy` will get automatically installed as a dependency of jabsim when you install it.
+     - Be careful using loops and if-statements. If you're an amateur programmer, just avoid doing all that. Otherwise, have a look at the [JAX documentation](https://docs.jax.dev/en/latest/index.html) .
+2. Import `jabsim` and call `jabsim.sim` to simulate. The arguments are as follows:
+   - `par`: a list, array or dict of model parameters as in your ODE function
+   - `model_ode`: the ODE function you created
+   - `x0`: the initial state vector as a 1D `np.array` or `jnp.array`
+   - `tf`: tuple, array or list. The ODE will be simulated 
+   - `savetimestep`: interval between the time points at which the trajectory is saved
+   - `simulator`: string specifyingthe simulation method to use
+        - `"euler"`: Euler simulator.
+        - `"rk4"`: Runge-Kutta 4th order simulator. Slower per ODE integration step but more accurate, hence allowing larger steps for the same accuracy.
+   - `ode_steps_in_savetimestep`: number of ODE integration steps within one timestep
+        - e.g. if `savetimestep=0.5` hours and `ode_step_in_savetimesteps=100`, there will be 100 integration steps per 0.5 hour, so the ODE integration step size will be 0.5/100=0.05 hours.
+        - high `ode_steps_in_savetimestep` number increases accuracy but increases runtimes
+   - `return_numpy`: if `True` (by default, it is) the output will be in the `np.array` format, otherwise it will be `jnp.array`.
+3. Running `jabsim.sim()` will return the arrays `ts` and `xs` as `np.array` or `jnp.array`, as well as a boolean value `success`.
+   - `ts`: array of timepoints between `tf[0]` and `tf[1]` with `savetimestep` hours, seconds or whatever units you are using between each two consecutive point
+   - `xs`: system trajectory saved as an array at the time points in `ts` - axis 0 for time, axis 1 for entries in the state vector (i.e. `xs.shape[0]=len(ts)`).
+   - `success`: boolean value; `True` if no entry in `xs` is `nan` or `inf`, `False` otherwise.
+
+## Notes
+- In practice, 500 steps per hour (e.g. `savetimestep=0.5, ode_steps_in_savetimestep=250` or `savetimestep=1.0, ode_steps_in_savetimestep=500`) works well for the RK4 solver. For the Euler solver, 1e4 steps per hours is reasonably good.
+- For benchmarking, you can also set `simulator="scipy"` to simulate your ODE with `scipy.solve_ivp` (but without any of the delicious jax features of the solvers above). In that case, don't use the arguments `ode_steps_in_savetimestep` and `savetimestep`. Instead, you can *optionally* specify:
+  - `solver`: string describing any solver which may be used with `scipy.solve_ivp`. By default, we have `solver="LSODA"`.
+  - `tols`: dictionary of relative and absolute tolerances for the scipy solver. By default, `tols={'rtol': 1e-6, 'atol': 1e-9}`.
+  - `dt0`: starting integration step size. By default, `dt0=0.1`.
+- If you want to make use of jax parallelisation, make sure to set `return_numpy=False` so that the solver would operate with `jnp.array` objects only.
+
+## Example
+
+Let us integrate a simple one-dimensional ODE $\frac{dx}{dt} = a x^2$. For the initial condition $x_0=1$ and $a=0.4$, 
+this has the analytical solution $x = \frac{1}{1-0.4t}$. This means we can verify that for  `savetimestep=0.5`, 
+`jabsim.sim()` produces `ts=np.array([0, 0.5, 1.0])` and `xs=np.array([1.0, 1.25, 1.66666667])`. 
+All entries in `xs` are finite, hence `success=True`.
+
+```python
+# import jabsim
+import jabsim
+
+# import jax.numpy for numpy operations
+import jax.numpy as jnp
+
+# our model ODE function returning a list of one element
+def model_ode(t, x, args):
+    # unpack args - get the dictiory of parameters
+    par, = args
+    
+    # use jnp to square x 
+    # (here you could just as well use x[0]**2, we just want to make a point)
+    x_squared = jnp.square(x[0])
+    
+    # return dx/dt as a list - with one entry for a one-dimensional ODE
+    return [par['a'] * x_squared]
+
+# our dictionary of paramneters
+par = {'a': 0.4}
+
+ts, xs, success = jabsim.sim(
+    model_ode=model_ode,
+    args=(par,),
+    x0=jnp.array([1.0]),
+    tf=(0.0, 1.0),
+    savetimestep=0.5,
+    simulator='rk4',
+    ode_steps_in_savetimestep=10,
+)
+
+# print the timne
+print(ts)
+print(xs)
+print(success)
+```
+
+# Citation
+If you find this package useful in your work, please cite the paper below:
+code for its Showcase 2 served as jabsim's direct ideological precursor.
+```bibtex
+@article{Gallup2024,
+	author = {Gallup, Olivia and Sechkar, Kirill and Towers, Sebastian and Steel, Harrison},
+	title = {Computational Synthetic Biology Enabled through JAX: A Showcase},
+	journal = {ACS Synth. Biol.},
+	volume = {13},
+	number = {9},
+	pages = {3046},
+	year = {2024},
+	doi = {10.1021/acssynbio.4c00307}
+}
+```
+
+The original JAX package should be cited as:
+```bibtex
+@software{jax2018github,
+  author = {James Bradbury and Roy Frostig and Peter Hawkins and Matthew James Johnson and Yash Katariya and Chris Leary and Dougal Maclaurin and George Necula and Adam Paszke and Jake Vander{P}las and Skye Wanderman-{M}ilne and Qiao Zhang},
+  title = {{JAX}: composable transformations of {P}ython+{N}um{P}y programs},
+  url = {http://github.com/jax-ml/jax},
+  version = {0.3.13},
+  year = {2018},
+}
+```

jabsim-0.1.0/__init__.py

@@ -0,0 +1,13 @@
+"""Simple JAX-based simulation of biosystem ODEs."""
+
+# import simulator functions
+from .simulators import (
+    loopy_euler,
+    loopy_rk4,
+    sim,
+)
+
+__all__ = [
+    'sim'
+]
+

jabsim-0.1.0/jabsim.egg-info/PKG-INFO

@@ -0,0 +1,132 @@
+Metadata-Version: 2.4
+Name: jabsim
+Version: 0.1.0
+Summary: Simple ODE simulation with JAX for biological systems
+Author-email: Kirill Sechkar <useful.instructive@proton.me>
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENCE
+Requires-Dist: jax>=0.4
+Requires-Dist: numpy>=1.24
+Requires-Dist: scipy>=1.10
+Provides-Extra: test
+Requires-Dist: pytest>=7; extra == "test"
+Dynamic: license-file
+
+# jabsim
+
+A jax-based package for simulating ODE models of biological systems,
+where **all variables are non-negative**. This enforcement of non-negativity, which neither `scipy.solve_ivp` nor `diffrax`
+solvers can give you, is why you may need jabsim.
+This is achieved by clamping the state variables to zero whenever they are negative 
+
+jabsim is powered by the [jax](https://github.com/jax-ml/jax) package for high-performance computing and parallelisation.
+This means that jabsim simulations can be jit-compiled and parallelised on a GPU or TPU using `jax.vmap` or shard mapping.
+
+## How to use jabsim
+
+1. Make an ODE function which calculates the derivative $\frac{dx}{dt}$ from the arguments `t, x, par` in this order.
+   - What do the arguments stand for?     
+     - `t` is the time at this point in the simulation
+     - `x` is the state vector at this time point
+     - `args` is a tuple of extra arguments passed on to the ODE function.
+   - **IMPORTANT**: if you don't know jax, there are a few differences to keep in mind:
+     - If you normally use `numpy` functions in your ODE, import `jax.numpy` as `jnp` and use it instead of `np`. jax with `jax.numpy` will get automatically installed as a dependency of jabsim when you install it.
+     - Be careful using loops and if-statements. If you're an amateur programmer, just avoid doing all that. Otherwise, have a look at the [JAX documentation](https://docs.jax.dev/en/latest/index.html) .
+2. Import `jabsim` and call `jabsim.sim` to simulate. The arguments are as follows:
+   - `par`: a list, array or dict of model parameters as in your ODE function
+   - `model_ode`: the ODE function you created
+   - `x0`: the initial state vector as a 1D `np.array` or `jnp.array`
+   - `tf`: tuple, array or list. The ODE will be simulated 
+   - `savetimestep`: interval between the time points at which the trajectory is saved
+   - `simulator`: string specifyingthe simulation method to use
+        - `"euler"`: Euler simulator.
+        - `"rk4"`: Runge-Kutta 4th order simulator. Slower per ODE integration step but more accurate, hence allowing larger steps for the same accuracy.
+   - `ode_steps_in_savetimestep`: number of ODE integration steps within one timestep
+        - e.g. if `savetimestep=0.5` hours and `ode_step_in_savetimesteps=100`, there will be 100 integration steps per 0.5 hour, so the ODE integration step size will be 0.5/100=0.05 hours.
+        - high `ode_steps_in_savetimestep` number increases accuracy but increases runtimes
+   - `return_numpy`: if `True` (by default, it is) the output will be in the `np.array` format, otherwise it will be `jnp.array`.
+3. Running `jabsim.sim()` will return the arrays `ts` and `xs` as `np.array` or `jnp.array`, as well as a boolean value `success`.
+   - `ts`: array of timepoints between `tf[0]` and `tf[1]` with `savetimestep` hours, seconds or whatever units you are using between each two consecutive point
+   - `xs`: system trajectory saved as an array at the time points in `ts` - axis 0 for time, axis 1 for entries in the state vector (i.e. `xs.shape[0]=len(ts)`).
+   - `success`: boolean value; `True` if no entry in `xs` is `nan` or `inf`, `False` otherwise.
+
+## Notes
+- In practice, 500 steps per hour (e.g. `savetimestep=0.5, ode_steps_in_savetimestep=250` or `savetimestep=1.0, ode_steps_in_savetimestep=500`) works well for the RK4 solver. For the Euler solver, 1e4 steps per hours is reasonably good.
+- For benchmarking, you can also set `simulator="scipy"` to simulate your ODE with `scipy.solve_ivp` (but without any of the delicious jax features of the solvers above). In that case, don't use the arguments `ode_steps_in_savetimestep` and `savetimestep`. Instead, you can *optionally* specify:
+  - `solver`: string describing any solver which may be used with `scipy.solve_ivp`. By default, we have `solver="LSODA"`.
+  - `tols`: dictionary of relative and absolute tolerances for the scipy solver. By default, `tols={'rtol': 1e-6, 'atol': 1e-9}`.
+  - `dt0`: starting integration step size. By default, `dt0=0.1`.
+- If you want to make use of jax parallelisation, make sure to set `return_numpy=False` so that the solver would operate with `jnp.array` objects only.
+
+## Example
+
+Let us integrate a simple one-dimensional ODE $\frac{dx}{dt} = a x^2$. For the initial condition $x_0=1$ and $a=0.4$, 
+this has the analytical solution $x = \frac{1}{1-0.4t}$. This means we can verify that for  `savetimestep=0.5`, 
+`jabsim.sim()` produces `ts=np.array([0, 0.5, 1.0])` and `xs=np.array([1.0, 1.25, 1.66666667])`. 
+All entries in `xs` are finite, hence `success=True`.
+
+```python
+# import jabsim
+import jabsim
+
+# import jax.numpy for numpy operations
+import jax.numpy as jnp
+
+# our model ODE function returning a list of one element
+def model_ode(t, x, args):
+    # unpack args - get the dictiory of parameters
+    par, = args
+    
+    # use jnp to square x 
+    # (here you could just as well use x[0]**2, we just want to make a point)
+    x_squared = jnp.square(x[0])
+    
+    # return dx/dt as a list - with one entry for a one-dimensional ODE
+    return [par['a'] * x_squared]
+
+# our dictionary of paramneters
+par = {'a': 0.4}
+
+ts, xs, success = jabsim.sim(
+    model_ode=model_ode,
+    args=(par,),
+    x0=jnp.array([1.0]),
+    tf=(0.0, 1.0),
+    savetimestep=0.5,
+    simulator='rk4',
+    ode_steps_in_savetimestep=10,
+)
+
+# print the timne
+print(ts)
+print(xs)
+print(success)
+```
+
+# Citation
+If you find this package useful in your work, please cite the paper below:
+code for its Showcase 2 served as jabsim's direct ideological precursor.
+```bibtex
+@article{Gallup2024,
+	author = {Gallup, Olivia and Sechkar, Kirill and Towers, Sebastian and Steel, Harrison},
+	title = {Computational Synthetic Biology Enabled through JAX: A Showcase},
+	journal = {ACS Synth. Biol.},
+	volume = {13},
+	number = {9},
+	pages = {3046},
+	year = {2024},
+	doi = {10.1021/acssynbio.4c00307}
+}
+```
+
+The original JAX package should be cited as:
+```bibtex
+@software{jax2018github,
+  author = {James Bradbury and Roy Frostig and Peter Hawkins and Matthew James Johnson and Yash Katariya and Chris Leary and Dougal Maclaurin and George Necula and Adam Paszke and Jake Vander{P}las and Skye Wanderman-{M}ilne and Qiao Zhang},
+  title = {{JAX}: composable transformations of {P}ython+{N}um{P}y programs},
+  url = {http://github.com/jax-ml/jax},
+  version = {0.3.13},
+  year = {2018},
+}
+```

jabsim-0.1.0/jabsim.egg-info/SOURCES.txt

@@ -0,0 +1,13 @@
+LICENCE
+README.md
+__init__.py
+pyproject.toml
+simulators.py
+./__init__.py
+./simulators.py
+jabsim.egg-info/PKG-INFO
+jabsim.egg-info/SOURCES.txt
+jabsim.egg-info/dependency_links.txt
+jabsim.egg-info/requires.txt
+jabsim.egg-info/top_level.txt
+tests/test_simulators.py

jabsim-0.1.0/jabsim.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

jabsim-0.1.0/jabsim.egg-info/requires.txt

@@ -0,0 +1,6 @@
+jax>=0.4
+numpy>=1.24
+scipy>=1.10
+
+[test]
+pytest>=7

jabsim-0.1.0/jabsim.egg-info/top_level.txt

@@ -0,0 +1 @@
+jabsim

jabsim-0.1.0/pyproject.toml

@@ -0,0 +1,23 @@
+[build-system]
+requires = ["setuptools>=68"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "jabsim"
+version = "0.1.0"
+description = "Simple ODE simulation with JAX for biological systems"
+readme = "README.md"
+authors = [{name = "Kirill Sechkar", email = "useful.instructive@proton.me"}]
+requires-python = ">=3.10"
+dependencies = [
+    "jax>=0.4",
+    "numpy>=1.24",
+    "scipy>=1.10",
+]
+
+[project.optional-dependencies]
+test = ["pytest>=7"]
+
+[tool.setuptools]
+packages = ["jabsim"]
+package-dir = {jabsim = "."}

jabsim-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

jabsim-0.1.0/simulators.py

@@ -0,0 +1,224 @@
+# JABSIM/SIMULATORS.PY - functions for ODE simulation
+
+# PACKAGE IMPORTS
+# import configured jax
+import jax
+# import everything else
+import numpy as np
+import scipy.integrate
+import jax.numpy as jnp
+
+
+# SOLVERS --------------------------------------------------------------------------------------------------------------
+
+# Euler solver using jax.lax.fori_loop
+def loopy_euler(model_ode,
+                save_cntr,
+                sim_state_record,
+                ode_step,
+                ode_steps_in_savetimestep,
+                args):
+    """
+        Use Euler integration to get the simulator state to be recorded at the current saving point into
+        sim_state_record - to be used with jax.lax.fori_loop
+
+        Args:
+            model_ode: ODE function handle
+            save_cntr: current saving point
+            sim_state_record: dictionary containing a record of simulation time points and state vector values
+            ode_step: simulation time step duration
+            ode_steps_in_savetimestep: number of ODE integration steps between each recording
+            args: additional arguments to pass to the ODE function
+
+        Returns:
+            new_sim_state_record: record of simulation time points and state vector values - now with the right x at xs[save_cntr,:]
+    """
+
+    # EULER STEP FUNCTION (with non-negativity constraints)
+    def euler_step(step_cntr, t_x):
+        return {
+            # entries updated over the course of the Euler step
+            't': t_x['t'] + ode_step,
+            'x': jnp.maximum(t_x['x'] + ode_step * model_ode(t_x['t'], t_x['x'], args),0),
+        }
+    # GET THE PRESENT TIME POINT AND STATE
+    last_t_x = {'t': sim_state_record['ts'][save_cntr-1], 'x': sim_state_record['xs'][save_cntr-1, :]}
+    this_t_x = jax.lax.fori_loop(0, ode_steps_in_savetimestep, euler_step, last_t_x)
+
+    # RETURN UPDATED SIMULATOR STATE
+    new_sim_state_record = {'ts': sim_state_record['ts'], 'xs': sim_state_record['xs'].at[save_cntr, :].set(this_t_x['x'])}
+    return new_sim_state_record
+
+# Fourth-order Runge-Kutta solver using jax.lax.fori_loop
+def loopy_rk4(model_ode,
+              save_cntr,
+              sim_state_record,
+              ode_step,
+              ode_steps_in_savetimestep,
+              args):
+    """
+        Use fourth-order Runge-Kutta integration to get the simulator state to be recorded at the current saving point
+        into sim_state_record - to be used with jax.lax.fori_loop
+
+        Args:
+            model_ode: ODE function handle
+            save_cntr: current saving point
+            sim_state_record: dictionary containing a record of simulation time points and state vector values
+            ode_step: simulation time step duration
+            ode_steps_in_savetimestep: number of ODE integration steps between each recording
+            args: additional arguments to pass to the ODE function
+
+        Returns:
+            new_sim_state_record: record of simulation time points and state vector values - now with the right x at xs[save_cntr,:]
+    """
+
+    # FOURTH-ORDER RUNGE-KUTTA STEP FUNCTION (with non-negativity constraints)
+    def rk4_step(step_cntr, t_x):
+        k1 = model_ode(t_x['t'], t_x['x'], args)
+        k2 = model_ode(t_x['t'] + ode_step / 2, jnp.maximum(t_x['x'] + ode_step * k1 / 2, 0), args)
+        k3 = model_ode(t_x['t'] + ode_step / 2, jnp.maximum(t_x['x'] + ode_step * k2 / 2, 0), args)
+        k4 = model_ode(t_x['t'] + ode_step, jnp.maximum(t_x['x'] + ode_step * k3, 0), args)
+        return {
+            # entries updated over the course of the Runge-Kutta step
+            't': t_x['t'] + ode_step,
+            'x': jnp.maximum(t_x['x'] + ode_step * (k1 + 2 * k2 + 2 * k3 + k4) / 6, 0),
+        }
+
+    # GET THE PRESENT TIME POINT AND STATE
+    last_t_x = {'t': sim_state_record['ts'][save_cntr-1], 'x': sim_state_record['xs'][save_cntr-1, :]}
+    this_t_x = jax.lax.fori_loop(0, ode_steps_in_savetimestep, rk4_step, last_t_x)
+
+    # RETURN UPDATED SIMULATOR STATE
+    new_sim_state_record = {'ts': sim_state_record['ts'], 'xs': sim_state_record['xs'].at[save_cntr, :].set(this_t_x['x'])}
+    return new_sim_state_record
+
+
+# SIMULATOR FUNCTION ---------------------------------------------------------------------------------------------------
+# simulator function
+def sim(model_ode, args, x0, tf, savetimestep, simulator='rk4', return_numpy=True, **kwargs):
+    """
+        Simulate an ODE model.
+
+        Args:
+            model_ode: ODE function handle
+            args: extra arguments for the ODE function
+            x0: initial condition
+            tf: time span of the simulation
+            savetimestep: saving the siulation every savetimestep hours
+            simulator: simulation method
+            return_numpy: True if returning numpy arrays, False is jax.numpy arrays
+            **kwargs: additional arguments to pass to the simulation method - specific to the simulator
+
+        Returns:
+            xs: system state at each time point specfied
+    """
+
+    if (simulator == 'euler'):
+        # define custom Euler parameters if not specified
+        ## number of ODE integration steps between each record point
+        ode_steps_in_savetimestep = kwargs.get('ode_steps_in_savetimestep', 1e4)
+
+        # define the time points at which we save the solution
+        ts = jnp.concatenate((jnp.arange(tf[0], tf[1], savetimestep), jnp.array([tf[1]])), dtype=jnp.float64)
+
+        # calculate the time step for the Euler integration
+        ode_step = savetimestep / ode_steps_in_savetimestep
+
+        # make the model ode function return a jnp.array
+        model_ode_jnp = lambda t, x, args: jnp.array(model_ode(t, x, args))
+
+        # make the retrieval of next x a lambda-function for jax.lax.scanning
+        loop_step = lambda save_cntr, sim_state_record: loopy_euler(model_ode_jnp,
+                                                                    save_cntr,
+                                                                    sim_state_record,  # simulator state
+                                                                    ode_step,  # simulation time step
+                                                                    int(ode_steps_in_savetimestep), # number of ODE integration steps between each recording
+                                                                    args)
+
+        # initalise the simulator state: (t, x) - x initialised with initial conditions
+        sim_state_record = {'ts': ts, 'xs': jnp.tile(jnp.array(x0), (ts.shape[0], 1))}
+
+        # simulate
+        sim_state_rec_final = jax.lax.fori_loop(1, ts.shape[0], loop_step, sim_state_record)
+        xs = sim_state_rec_final['xs']
+
+        # check for simulation success (i.e. no nans or infs in x)
+        success = not bool(jnp.any(jnp.isnan(xs)) or jnp.any(jnp.isinf(xs)))
+
+        # return numpy or jax.numpy arrays
+        if(return_numpy):
+            return np.array(ts), np.array(xs), success
+        else:
+            return ts, xs, success
+
+    elif (simulator == 'rk4'):
+        # define custom fourth-order Runge-Kutta parameters if not specified
+        ## number of ODE integration steps between each record point
+        ode_steps_in_savetimestep = kwargs.get('ode_steps_in_savetimestep', 1e4)
+
+        # define the time points at which we save the solution
+        ts = jnp.concatenate((jnp.arange(tf[0], tf[1], savetimestep), jnp.array([tf[1]])), dtype=jnp.float64)
+
+        # calculate the time step for the fourth-order Runge-Kutta integration
+        ode_step = savetimestep / ode_steps_in_savetimestep
+
+        # make the model ode function return a jnp.array
+        model_ode_jnp = lambda t, x, args: jnp.array(model_ode(t, x, args))
+
+        # make the retrieval of next x a lambda-function for jax.lax.fori_loop
+        loop_step = lambda save_cntr, sim_state_record: loopy_rk4(model_ode_jnp,
+                                                                  save_cntr,
+                                                                  sim_state_record,  # simulator state
+                                                                  ode_step,  # simulation time step
+                                                                  int(ode_steps_in_savetimestep), # number of ODE integration steps between each recording
+                                                                  args)
+
+        # initalise the simulator state: (t, x) - x initialised with initial conditions
+        sim_state_record = {'ts': ts, 'xs': jnp.tile(jnp.array(x0), (ts.shape[0], 1))}
+
+        # simulate
+        sim_state_rec_final = jax.lax.fori_loop(1, ts.shape[0], loop_step, sim_state_record)
+        xs = sim_state_rec_final['xs']
+
+        # check for simulation success (i.e. no nans or infs in x)
+        success = not bool(jnp.any(jnp.isnan(xs)) or jnp.any(jnp.isinf(xs)))
+
+        # return numpy or jax.numpy arrays
+        if (return_numpy):
+            return np.array(ts), np.array(xs), success
+        else:
+            return ts, xs, success
+
+    elif (simulator == 'scipy'):
+        # define ODE integration term
+        term = lambda t, x: model_ode(t, x, args)
+
+        # define the time points at which we save the solution
+        ts = np.concatenate((np.arange(tf[0], tf[1], savetimestep), np.array([tf[1]])))
+
+        # define diffrax parameters if not specified
+        ## ODE solver
+        solver = kwargs.get('solver', 'LSODA')
+        ## ODE integration tolerances to specify the step size controller
+        tols = kwargs.get('tols', {'rtol': 1e-6, 'atol': 1e-9})
+        ## initial time step
+        dt0 = kwargs.get('dt0', 0.1)
+
+        # solve the ODE
+        result = scipy.integrate.solve_ivp(term,
+                                           t_span=(tf[0], tf[-1]),
+                                           y0=x0,
+                                           t_eval=ts,
+                                           method=solver,
+                                           rtol=tols['rtol'], atol=tols['atol'],
+                                           first_step=dt0)
+        xs = (result.y).T
+
+        # return numpy or jax.numpy arrays
+        if (return_numpy):
+            return ts, xs, result.success
+        else:
+            return jnp.array(ts), jnp.array(xs), result.success
+
+    else:
+        raise ValueError("Unknown simulator: {}".format(simulator))

jabsim-0.1.0/tests/test_simulators.py

@@ -0,0 +1,29 @@
+import numpy as np
+import jax.numpy as jnp
+import pytest
+
+import jax.numpy as jnp
+from jabsim import sim
+
+
+def decay(t, x, args):
+    par, = args
+    for i in range(0,len(x)):
+        x0_sq = jnp.square(x[i])
+    return [par['rate'] * x0_sq]
+
+
+@pytest.mark.parametrize('simulator', ['scipy', 'euler', 'rk4'])
+def test_simulates_exponential_decay(simulator):
+    kwargs = {'ode_steps_in_savetimestep': 100} if simulator != 'scipy' else {}
+    ts, xs, success = sim(decay, ({'rate': 0.5},), np.array([1.0]), (0.0, 0.5), 0.05,
+                 simulator=simulator, **kwargs)
+
+    np.testing.assert_allclose(ts, np.linspace(0.0, 0.5, 11))
+    np.testing.assert_allclose(np.asarray(xs[:, 0]), 1/(1-0.5*np.asarray(ts)), rtol=2e-2, atol=1e-4)
+    assert(success)
+
+
+def test_rejects_unknown_simulator():
+    with pytest.raises(ValueError, match='Unknown simulator: unknown'):
+        sim({}, decay, np.array([1.0]), (0.0, 1.0), 0.25, simulator='unknown')