```python
from sbmltoodejax.utils import load_biomodel
# load and simulate model
model, _, _, _ = load_biomodel(10)
n_secs = 150*60
n_steps = int(n_secs / model.deltaT)
ys, ws, ts = model(n_steps)
# plot time-course simulation
...
```
```python
# Load Model
model, y0, w0, c = load_model(model_idx)
# Optax pipeline
@jit
def loss_fn(c, model):
"""loss function"""
ys, ws, ts = model(n_steps, y0, w0, c)
loss = jnp.sqrt(jnp.square(ys[y_indexes["Ca_Cyt"]]
-target_Ca_Cyt).sum())
return loss
@jit
def make_step(c, model, opt_state):
"""update function"""
loss, grads = value_and_grad(loss_fn)(c, model)
updates, opt_state = optim.update(grads, opt_state)
c = optax.apply_updates(c, updates)
return loss, c, opt_state
n_optim_steps = 1000
optim = optax.adam(1e-3)
opt_state = optim.init(c)
train_loss = []
for optim_step_idx in range(n_optim_steps):
loss, c, opt_state = make_step(c, model, opt_state)
train_loss.append(loss)
# Plot train loss
plt.plot(train_loss)
plt.ylabel("training loss")
plt.xlabel("train steps")
plt.show()
```
We can see that for short reaction times (here <100 secs with $\Delta T=0.1$, see Log Scale), SBMLtoODEjax simulation takes longer than the original SBMLtoODEpy
library because when calling *ModelStep* for the first time, it takes some time to generate the compiled trace.
However, the advantage of SBMLtoODEjax becomes clear when considering longer rollouts where we obtain huge speed-ups with respect to original SBMLtoODEpy library (see Linear Scale).
This is because the original SBMLtoODEpy python code uses for-loops, hence have linear increase of compute time, whereas the scanned JIT-compiled step function executes much faster.
You can check our [Benchmarking](tutorials/benchmark.ipynb) tutorial for more details on the comparison.
### Automatic vectorization
`vmap` is another core function transformations in JAX which enables efficient (and seamless) vectorization of functions.
Here, it is particularly useful for running simulation in parallel with batched computations such as batch of initial conditions.
As shown below, doing this in JAX is very simple as one simply need to use `vmap` transformation to vectorize the model function and then call the batched model in the exact same way:
```python
# vector of initial conditions
batched_y0 = ...
# batch model
batched_model = vmap(model, in_axes=(None, 0), out_axes=(0, 0, None))
# run simulation in batch mode
batched_ys, batched_ws, ts = batched_model(n_steps, batched_y0)
```
Below we compare the average simulation time of model rollouts for 1) the SBMLtoODEpy-generated models and for loop computations over the inputs (pink), 2) the SBMLtoODEpy-generated models and pooling over the inputs (orange)
and 3) the SBMLtoODEjax library with vectorized computations (blue). We show results for different batch sizes (x-axis), where all runs have been done on the same machine and for
a rollout length of 10 seconds with $\Delta T=0.1$.
Again, similar conclusions can be drawn where SBMLtoODEjax is less efficient for small batch sizes (and small rollout lengths), but becomes very advantageous
for larger batch sizes.
You can check our [Benchmarking](tutorials/benchmark.ipynb) tutorial for more details on the comparison.
### Automatic differentiation
Finally, `grad` is another core function transformations in JAX which enables automatic differentiation
and allows seamless integration of SBMLtoODEjax models with [Optax](https://optax.readthedocs.io/en/latest/) pipelines, a gradient processing and optimization library for JAX.
Whereas Optax has typically been used to optimize parameters of neural networks, we can use it in the very same way to optimize parameters of our SBMLtoODEjax models
as shown below:
👉 Those limitations should be resolvable and will hopefully be fixed in future versions.
Please consider contributing and check our [Contribution Guidelines](contributing.md) to do so!