Implementing an ODE model in Pymob#

In this tutorial, we will implement a simple ODE model, create simulation results and infer an unknown parameter from artificially generated data. It is recommended to work through this notebook after the introductiory tutorial where something very similar is done for a linear regression model.

After setting up the simulation manually (Chapter 1), we will save our settings and create a new simulation from those settings (Chapter 2).

Chapter 1: Setting up the model πŸ‘©β€πŸ’»#

πŸ‘‰ Let’s begin with setting up a Pymob simulation for an ODE model. This will follow roughly the same procedure as the introductory tutorial. We do, however, need to make some tweaks to allow for the needs of an ODE model.

# First, import the necessary python packages
import numpy as np
import matplotlib.pyplot as plt
import xarray as xr
from scipy.integrate import solve_ivp

# Import the pymob modules
from pymob.simulation import SimulationBase
from pymob.solvers.diffrax import JaxSolver
from pymob.sim.config import Param, DataVariable

1.1 Creating the sim object 🧩#

πŸ‘‰ As an example for a relatively simple ODE model, we will use the well-known Lotka-Volterra model describing a predator-prey relationship.

πŸ‘‰ The equations for this model look like this (\(X\) and \(Y\) denote prey and predator, respectively):

\(\frac{dX}{dt} = \alpha X - \beta X Y\)

\(\frac{dY}{dt} = \gamma X Y - \delta Y\)

\(\newline \alpha, \beta, \gamma, \delta > 0\)

πŸ‘‰ In the following cell, we will define our model. To work with our solver (we will later use pymob.solvers.diffrax.JaxSolver which calls diffrax.diffeqsolve), our Python function needs to have a signature of the form fun(t, y, *args) where t represents the current time within the system, y represents the current system state and *args is a placeholder for all model parameters.

πŸ‘‰ Note that the argument t is not used inside the function as the derivatives generated by the Lotka Volterra model are independent from time. It still needs to be included in the signature to satisfy the needs of the solver.

def lotkavolterra(t, y, alpha, beta, gamma, delta):
    X, Y = y
    dXdt = alpha * X - beta * X * Y
    dYdt = gamma * X * Y - delta * Y
    return dXdt, dYdt

πŸ‘‰ We can then create our simulation object and assign the model and the solver to it:

# Initialize the simulation object
sim = SimulationBase()

# Configure the case study
sim.config.case_study.name = "ODEtutorial"
sim.config.case_study.scenario = "lotkavolterra"

# Add the model to the simulation
sim.model = lotkavolterra

# Define a solver
sim.solver = JaxSolver
/home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/pymob/sim/config/base.py:397: UserWarning: Case study 'unnamed_case_study' could not be imported. Install the case study with `pip install unnamed_case_study`.
  warnings.warn(

1.2 Generating artificial data πŸ“ˆ#

πŸ‘‰ Now we generate some artificial data that we will later use as our observations. To do this, we generate a time series of the Lotka-Volterra model with parameters \(\alpha = 0.7, \beta = 0.1, \gamma = 0.1, \delta = 0.9\) from the initial condition \(X = 10, Y = 5\) using solve_ivp (we could also use diffrax.diffeqsolve here, that would make no difference). This is done for 101 steps with \(\Delta t = 0.5\).

πŸ‘‰ We then add some noise to the data and make sure that predator and prey abundances in our data are always positive as negative abundances would never be measured in reality.

πŸ‘‰ After running the code, you can take a look at our artificial data and recognize the characteristic periodic oscillations produced by the Lotka-Volterra model.

# Generate Lotka Volterra time series
sol = solve_ivp(lotkavolterra, (0, 50), np.array([10,5]), "LSODA", np.linspace(0,50,101), args=[0.7,0.1,0.1,0.9])

# Add "random" noise (example is made reproducible by setting a fixed seed)
rng = np.random.default_rng(seed=1)
noise = rng.normal(0, 0.5, (2,101))
y_obs = sol.y + noise
y_obs = np.greater(y_obs, np.zeros(y_obs.shape)) * y_obs

# Save the evaluated time points
t = sol.t

# Plot the generated data
fig, ax = plt.subplots(figsize=(5, 4))
ax.plot(t, y_obs.transpose(), label='Datapoints')
ax.set(xlabel='t [-]', ylabel='y_obs [-]', title ='Artificial Data')
plt.tight_layout()

png

1.3 Adding data to the sim object 🀝#

πŸ‘‰ Let’s prepare our observations. As seen in the introductory tutorial, Pymob uses xArray datasets. Because our model has two state variables, the dataset containing our artificial data also needs to have two data variables. It also needs to include the time points we generated the data for as a coordinate axis. This can be achieved like this (or probably in an easier way):

# Create an xArray dataset containing the artificial data
data_obs_1 = xr.DataArray(y_obs[0], coords={"time": t}).to_dataset(name="prey")
data_obs_2 = xr.DataArray(y_obs[1], coords={"time": t}).to_dataset(name="predator")
data_obs = xr.merge([data_obs_1, data_obs_2])

# Look at the structure of the generated datatset
data_obs
<xarray.Dataset>
Dimensions:   (time: 101)
Coordinates:

  • time (time) float64 0.0 0.5 1.0 1.5 2.0 … 48.0 48.5 49.0 49.5 50.0 Data variables: prey (time) float64 10.17 11.36 11.85 11.33 … 11.08 11.16 12.37 11.56 predator (time) float64 5.431 5.33 6.397 7.604 … 5.544 5.436 7.871 9.127

πŸ‘‰ As our next step, we add our artificial data to the model. As you can see in the cell output, Pymob automatically detects the two data variables and the time axis and creates two pymob.sim.config.DataVariable objects within the simulation’s pymob.sim.config.DataStructure instance. That’s why it’s so important to prepare the data in the way we did above!

# Add our dataset to the simulation
sim.observations = data_obs

# Take a look at the layout of the data
sim.config.data_structure
MinMaxScaler(variable=prey, min=5.844172888098378, max=12.525948698266157)
MinMaxScaler(variable=predator, min=4.053933700151413, max=10.925258075625722)


/home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/pymob/simulation.py:366: UserWarning: `sim.config.data_structure.prey = Datavariable(dimensions=['time'] unit=None min=5.844172888098378 max=12.525948698266157 observed=True dimensions_evaluator=None)` has been assumed from `sim.observations`. If the order of the dimensions should be different, specify `sim.config.data_structure.prey = DataVariable(dimensions=[...], ...)` manually.
  warnings.warn(
/home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/pymob/simulation.py:366: UserWarning: `sim.config.data_structure.predator = Datavariable(dimensions=['time'] unit=None min=4.053933700151413 max=10.925258075625722 observed=True dimensions_evaluator=None)` has been assumed from `sim.observations`. If the order of the dimensions should be different, specify `sim.config.data_structure.predator = DataVariable(dimensions=[...], ...)` manually.
  warnings.warn(





Datastructure(indices=[], prey=DataVariable(dimensions=['time'], unit=None, min=5.844172888098378, max=12.525948698266157, observed=True, dimensions_evaluator=None), predator=DataVariable(dimensions=['time'], unit=None, min=4.053933700151413, max=10.925258075625722, observed=True, dimensions_evaluator=None))

πŸ‘‰ Because the results of ODE models strongly depend on their initial conditions, our simulation object need to know those. The correct place to put this information is model_parameters["y0"].

πŸ‘‰ The initial conditions also have to be an xArray dataset with two data variables (but without the time coordinate). We can do this manually like before by creating a xArray.Dataset object from our initial conditions:

# Create an xArray dataset
y0_obs_1 = xr.DataArray(10).to_dataset(name="prey")
y0_obs_2 = xr.DataArray(5).to_dataset(name="predator")
y0_obs = xr.merge([y0_obs_1, y0_obs_2])

# Add the initial condition to the simulation
sim.model_parameters["y0"] = y0_obs

Using parse_input()

Otherwise we can use which extracts all the necessary information from the configuration. This is, however, only possible after we give add this information to the configuration. This might seem unnecessary at the moment but you will later see why it makes sense in certain situations.

# Pass the initial condition to the simulation
#
# Note: The input needs to be a list containing a separate string for every state variable.
#       Those strings must have the format "variableName=initialValue" (without any spaces!).
sim.config.simulation.y0 = ["prey=10", "predator=5"]

# Let parse_input() create an xArray dataset
#
# Note: The input variable drop_dims makes sure that the dataset only contains a single value
#       instead of a full time series filled with the same value over and over again.
y0_obs = sim.parse_input("y0", drop_dims=['time'])

# Add the initial condition to the simulation
sim.model_parameters["y0"] = y0_obs

1.4 Setting parameters and running the model πŸ‘Ÿ#

πŸ‘‰ The next step is defining the parameters of the system, similarly as in the introductiory tutorial. In this case, we want to have three fixed parameters (\(\alpha = 0.7, \beta = 0.1, \gamma = 0.1\)) and a single free parameter (\(\delta\)). You will soon see why we made that choice.

# Parameterize the model
sim.config.model_parameters.alpha = Param(value=0.7, free=False)
sim.config.model_parameters.beta = Param(value=0.1, free=False)
sim.config.model_parameters.gamma = Param(value=0.1, free=False)
sim.config.model_parameters.delta = Param(value=0.9, free=True)

# Make sure the model parameters are available to the model
sim.model_parameters["parameters"] = sim.config.model_parameters.value_dict

# Look at the parameter values passed to the model
sim.model_parameters["parameters"]
{'alpha': 0.7, 'beta': 0.1, 'gamma': 0.1, 'delta': 0.9}

πŸ‘‰ We do not need to define model_parameters["x_in"] as we don’t wave any input data in this case. If we wanted to make the growth rates in our model depend on weather conditions and use a corresponding dataset, model_parameters["x_in"] would be the place to include our external data.

πŸ‘‰ Instead, we follow the same routine as in the introductory tutorial, let Pymob initialize the simulation and look at the resulting time series (with \(\delta = 0.9\)):

# Put everything in place for running the simulation
sim.dispatch_constructor()

# Create an evaluator, run the simulation and obtain the results
evaluator = sim.dispatch(theta={"delta":0.9})
evaluator()
data_res = evaluator.results

# Plot the results
fig, ax = plt.subplots(figsize=(5, 4))
ax.plot(data_obs.time, data_obs.prey, ls="-", color="tab:blue", alpha=.5, label ="observation data")
ax.plot(data_obs.time, data_obs.predator, ls="-", color="tab:blue", alpha=.5, label ="observation data")
ax.plot(data_res.time, data_res.prey, color="black", label ="result")
ax.plot(data_res.time, data_res.predator, color="black", label ="result")
ax.legend()
/home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/pymob/simulation.py:723: UserWarning: The number of ODE states was not specified in the config file [simulation] > 'n_ode_states = <n>'. Extracted the return arguments ['dXdt', 'dYdt'] from the source code. Setting 'n_ode_states=2.
  warnings.warn(





<matplotlib.legend.Legend at 0x72007baae2d0>

png

1.5 Finding out the value of \(\delta\) πŸ”Ž#

πŸ‘‰ Now let’s see which value for \(\delta\) best fits our data. To do that, we use the inferer in the same way as in the introductory tutorial. We do, however, need to apply our error model to both of our state variables. Also, we changed the prior for \(\delta\) to a uniform distribution from 0.5 to 1.5 because that’s a better guess.

Caution

The following code will throw an error. This is not your fault, just look at the error message and continue with the next markdown cell.

from jaxlib.xla_extension import XlaRuntimeError
# Add parameters to use in our error model
sim.config.model_parameters.sigma_prey = Param(free=True , prior="lognorm(scale=1,s=1)", min=0, max=1)
sim.config.model_parameters.sigma_predator = Param(free=True , prior="lognorm(scale=1,s=1)", min=0, max=1)

# Define the error model for both state variables
sim.config.error_model.prey = "normal(loc=prey,scale=sigma_prey)"
sim.config.error_model.predator = "normal(loc=predator,scale=sigma_predator)"

# Choose a prior distribution for delta
sim.config.model_parameters.delta.prior = "uniform(loc=0.5,scale=1)"

try:

    # Create the inferer (NumPyro backend, NUTS kernel) and let it do its work
    sim.set_inferer("numpyro")
    sim.inferer.config.inference_numpyro.kernel = "nuts"
    sim.inferer.run()

    # Plot the results
    sim.config.simulation.x_dimension = "time"
    sim.posterior_predictive_checks(pred_hdi_style={"alpha": 0.1})

except XlaRuntimeError as e:

    # Print the error message
    print("An error occurred:", type(e).__name__, ":", e)
Jax 64 bit mode: False
Absolute tolerance: 1e-07


/home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/pymob/inference/numpyro_backend.py:565: UserWarning: Model is not rendered, because the graphviz executable is not found. Try search for 'graphviz executables not found' and the used OS. This should be an easy fix :-)
  warnings.warn(


      Trace Shapes:      
       Param Sites:      
      Sample Sites:      
         delta dist     |
              value     |
    sigma_prey dist     |
              value     |
sigma_predator dist     |
              value     |
      prey_obs dist 101 |
              value 101 |
  predator_obs dist 101 |
              value 101 |
An error occurred: XlaRuntimeError : INTERNAL: Generated function failed: CpuCallback error: _EquinoxRuntimeError: The maximum number of solver steps was reached. Try increasing `max_steps`.


--------------------
An error occurred during the runtime of your JAX program! Unfortunately you do not appear to be using `equinox.filter_jit` (perhaps you are using `jax.jit` instead?) and so further information about the error cannot be displayed. (Probably you are seeing a very large but uninformative error message right now.) Please wrap your program with `equinox.filter_jit`.
--------------------


At:
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/equinox/_errors.py(89): raises
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/jax/_src/callback.py(258): _flat_callback
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/jax/_src/callback.py(52): pure_callback_impl
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/jax/_src/callback.py(188): _callback
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/jax/_src/interpreters/mlir.py(2327): _wrapped_callback
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/jax/_src/interpreters/pxla.py(1145): __call__
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/jax/_src/profiler.py(334): wrapper
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/jax/_src/pjit.py(1178): _pjit_call_impl_python
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/jax/_src/pjit.py(1222): call_impl_cache_miss
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/jax/_src/pjit.py(1238): _pjit_call_impl
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/jax/_src/core.py(893): process_primitive
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/jax/_src/core.py(405): bind_with_trace
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/jax/_src/core.py(2682): bind
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/jax/_src/pjit.py(166): _python_pjit_helper
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/jax/_src/pjit.py(255): cache_miss
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/jax/_src/traceback_util.py(177): reraise_with_filtered_traceback
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/pymob/solvers/base.py(83): __call__
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/pymob/sim/evaluator.py(351): __call__
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/pymob/inference/numpyro_backend.py(274): evaluator
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/pymob/inference/numpyro_backend.py(498): model
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/numpyro/primitives.py(105): __call__
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/numpyro/primitives.py(105): __call__
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/numpyro/primitives.py(105): __call__
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/numpyro/handlers.py(171): get_trace
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/numpyro/infer/util.py(450): _get_model_transforms
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/numpyro/infer/util.py(656): initialize_model
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/numpyro/infer/hmc.py(657): _init_state
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/numpyro/infer/hmc.py(713): init
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/numpyro/infer/mcmc.py(416): _single_chain_mcmc
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/numpyro/infer/mcmc.py(634): run
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/pymob/inference/numpyro_backend.py(712): run_mcmc
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/pymob/inference/numpyro_backend.py(579): run
  /tmp/ipykernel_770/119426844.py(17): <module>
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/IPython/core/interactiveshell.py(3548): run_code
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/IPython/core/interactiveshell.py(3488): run_ast_nodes
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/IPython/core/interactiveshell.py(3306): run_cell_async
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/IPython/core/async_helpers.py(129): _pseudo_sync_runner
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/IPython/core/interactiveshell.py(3101): _run_cell
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/IPython/core/interactiveshell.py(3046): run_cell
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/ipykernel/zmqshell.py(663): run_cell
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/ipykernel/ipkernel.py(464): do_execute
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/ipykernel/kernelbase.py(834): execute_request
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/ipykernel/ipkernel.py(372): execute_request
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/ipykernel/kernelbase.py(478): dispatch_shell
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/ipykernel/kernelbase.py(621): shell_main
  /home/docs/.asdf/installs/python/3.11.12/lib/python3.11/asyncio/events.py(84): _run
  /home/docs/.asdf/installs/python/3.11.12/lib/python3.11/asyncio/base_events.py(1936): _run_once
  /home/docs/.asdf/installs/python/3.11.12/lib/python3.11/asyncio/base_events.py(608): run_forever
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/tornado/platform/asyncio.py(211): start
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/ipykernel/kernelapp.py(758): start
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/traitlets/config/application.py(1075): launch_instance
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/ipykernel_launcher.py(18): <module>
  <frozen runpy>(88): _run_code
  <frozen runpy>(198): _run_module_as_main

πŸ‘‰ What you see is an error that originated during runtime. The error message should tell you:

_EquinoxRuntimeError: The maximum number of solver steps was reached. Try increasing 'max_steps'.

πŸ‘‰ This means that our solver has to deal with a very difficult problem. To accomodate that, it needs to be very precise and work with extremely small time steps which causes it to exceed the maximum number of steps it is allowed to take.

πŸ‘‰ We can solve this in two ways:

  1. Increase max_steps: The simplest work to deal with this problem. It might not always work, though, because with very extreme model dynamics, even a high number of steps can be exceeded.

  2. Set throw_exception to False: With this setting, exceeding the maximum number of steps will not result in an error but return inf values as the result. In that case, the loss would also be infinite and the corresponding value of \(\delta\) would simply be rejected. That means that difficult problems are being thrown out and we make our decision about \(\delta\) based on the remaining runs. In many cases, extreme model behavior resulting in max_steps being exceeded will not fit the data anyway and rejecting the corresponding parameter value is justified. But to make such an assumption, you should know your system very well and check whether the assumption is valid.

πŸ‘‰ We will first try option 1:

# Increase max_steps
sim.config.jaxsolver.max_steps = 100000000

# Put everything in place (needs to be run again because we changed an important setting)
sim.dispatch_constructor()

try:

    # Try running the inferer again
    sim.inferer.run()

    # Plot the results
    sim.config.simulation.x_dimension = "time"
    sim.posterior_predictive_checks(pred_hdi_style={"alpha": 0.1})

except XlaRuntimeError as e:

    # Print the error message
    print("An error occurred:", type(e).__name__, ":", e)
/home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/pymob/inference/numpyro_backend.py:565: UserWarning: Model is not rendered, because the graphviz executable is not found. Try search for 'graphviz executables not found' and the used OS. This should be an easy fix :-)
  warnings.warn(


      Trace Shapes:      
       Param Sites:      
      Sample Sites:      
         delta dist     |
              value     |
    sigma_prey dist     |
              value     |
sigma_predator dist     |
              value     |
      prey_obs dist 101 |
              value 101 |
  predator_obs dist 101 |
              value 101 |


An error occurred: XlaRuntimeError : INTERNAL: Generated function failed: CpuCallback error: _EquinoxRuntimeError: The maximum number of solver steps was reached. Try increasing `max_steps`.


--------------------
An error occurred during the runtime of your JAX program! Unfortunately you do not appear to be using `equinox.filter_jit` (perhaps you are using `jax.jit` instead?) and so further information about the error cannot be displayed. (Probably you are seeing a very large but uninformative error message right now.) Please wrap your program with `equinox.filter_jit`.
--------------------


At:
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/equinox/_errors.py(89): raises
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/jax/_src/callback.py(258): _flat_callback
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/jax/_src/callback.py(52): pure_callback_impl
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/jax/_src/callback.py(188): _callback
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/jax/_src/interpreters/mlir.py(2327): _wrapped_callback
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/jax/_src/interpreters/pxla.py(1145): __call__
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/jax/_src/profiler.py(334): wrapper
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/jax/_src/pjit.py(1178): _pjit_call_impl_python
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/jax/_src/pjit.py(1222): call_impl_cache_miss
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/jax/_src/pjit.py(1238): _pjit_call_impl
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/jax/_src/core.py(893): process_primitive
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/jax/_src/core.py(405): bind_with_trace
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/jax/_src/core.py(2682): bind
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/jax/_src/pjit.py(166): _python_pjit_helper
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/jax/_src/pjit.py(255): cache_miss
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/jax/_src/traceback_util.py(177): reraise_with_filtered_traceback
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/pymob/solvers/base.py(83): __call__
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/pymob/sim/evaluator.py(351): __call__
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/pymob/inference/numpyro_backend.py(274): evaluator
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/pymob/inference/numpyro_backend.py(498): model
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/numpyro/primitives.py(105): __call__
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/numpyro/primitives.py(105): __call__
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/numpyro/primitives.py(105): __call__
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/numpyro/handlers.py(171): get_trace
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/numpyro/infer/util.py(450): _get_model_transforms
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/numpyro/infer/util.py(656): initialize_model
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/numpyro/infer/hmc.py(657): _init_state
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/numpyro/infer/hmc.py(713): init
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/numpyro/infer/mcmc.py(416): _single_chain_mcmc
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/numpyro/infer/mcmc.py(634): run
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/pymob/inference/numpyro_backend.py(712): run_mcmc
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/pymob/inference/numpyro_backend.py(579): run
  /tmp/ipykernel_770/2085724305.py(10): <module>
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/IPython/core/interactiveshell.py(3548): run_code
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/IPython/core/interactiveshell.py(3488): run_ast_nodes
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/IPython/core/interactiveshell.py(3306): run_cell_async
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/IPython/core/async_helpers.py(129): _pseudo_sync_runner
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/IPython/core/interactiveshell.py(3101): _run_cell
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/IPython/core/interactiveshell.py(3046): run_cell
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/ipykernel/zmqshell.py(663): run_cell
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/ipykernel/ipkernel.py(464): do_execute
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/ipykernel/kernelbase.py(834): execute_request
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/ipykernel/ipkernel.py(372): execute_request
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/ipykernel/kernelbase.py(478): dispatch_shell
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/ipykernel/kernelbase.py(621): shell_main
  /home/docs/.asdf/installs/python/3.11.12/lib/python3.11/asyncio/events.py(84): _run
  /home/docs/.asdf/installs/python/3.11.12/lib/python3.11/asyncio/base_events.py(1936): _run_once
  /home/docs/.asdf/installs/python/3.11.12/lib/python3.11/asyncio/base_events.py(608): run_forever
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/tornado/platform/asyncio.py(211): start
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/ipykernel/kernelapp.py(758): start
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/traitlets/config/application.py(1075): launch_instance
  /home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/ipykernel_launcher.py(18): <module>
  <frozen runpy>(88): _run_code
  <frozen runpy>(198): _run_module_as_main

πŸ‘‰ Even with max_steps set to 100.000.000 (the default value is 4096), we still get a runtime error, it just needs a little longer to appear. That means that we probably have an extremely sensitive numerical problem for some of our prior values, exceeding even an unreasonable amount of solver steps. So let’s try option 2:

# Decrease max_steps to a reasonable value and set throw_exception to False
sim.config.jaxsolver.max_steps = 10000
sim.config.jaxsolver.throw_exception = False

# Put everything in place (needs to be run again because we changed an important setting)
sim.dispatch_constructor()

try:

    # Try running the inferer again
    sim.inferer.run()

    # Plot the results
    sim.config.simulation.x_dimension = "time"
    sim.posterior_predictive_checks(pred_hdi_style={"alpha": 0.1})

except XlaRuntimeError as e:

    # Print the error message
    print("An error occurred:", type(e).__name__, ":", e)
/home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/pymob/inference/numpyro_backend.py:565: UserWarning: Model is not rendered, because the graphviz executable is not found. Try search for 'graphviz executables not found' and the used OS. This should be an easy fix :-)
  warnings.warn(


      Trace Shapes:      
       Param Sites:      
      Sample Sites:      
         delta dist     |
              value     |
    sigma_prey dist     |
              value     |
sigma_predator dist     |
              value     |
      prey_obs dist 101 |
              value 101 |
  predator_obs dist 101 |
              value 101 |

0%| | 0/3000 [00:00<?, ?it/s]

warmup: 0%| | 1/3000 [00:04<3:44:38, 4.49s/it, 1 steps of size 1.87e+00. acc. prob=0.00]

warmup: 1%|▏ | 38/3000 [00:04<04:15, 11.61it/s, 1 steps of size 2.44e-04. acc. prob=0.69]

warmup: 2%|▏ | 60/3000 [00:04<02:28, 19.84it/s, 127 steps of size 1.98e-03. acc. prob=0.74]

warmup: 3%|β–Ž | 78/3000 [00:04<01:44, 27.83it/s, 3 steps of size 1.78e-03. acc. prob=0.75]

warmup: 3%|β–Ž | 94/3000 [00:05<01:19, 36.33it/s, 43 steps of size 2.08e-03. acc. prob=0.75]

warmup: 4%|β–Ž | 109/3000 [00:05<01:08, 42.16it/s, 31 steps of size 2.05e-01. acc. prob=0.76]

warmup: 4%|▍ | 130/3000 [00:05<00:48, 59.33it/s, 63 steps of size 8.87e-02. acc. prob=0.77]

warmup: 5%|β–Œ | 154/3000 [00:05<00:34, 82.58it/s, 63 steps of size 5.77e-02. acc. prob=0.77]

warmup: 6%|β–Œ | 176/3000 [00:05<00:27, 103.15it/s, 31 steps of size 1.67e-01. acc. prob=0.77]

warmup: 7%|β–‹ | 203/3000 [00:05<00:21, 132.39it/s, 27 steps of size 1.65e-01. acc. prob=0.78]

warmup: 8%|β–Š | 230/3000 [00:05<00:17, 160.18it/s, 15 steps of size 2.18e-01. acc. prob=0.78]

warmup: 9%|β–‰ | 263/3000 [00:05<00:13, 197.32it/s, 15 steps of size 6.02e-01. acc. prob=0.78]

warmup: 10%|β–‰ | 290/3000 [00:05<00:12, 214.43it/s, 3 steps of size 6.16e-02. acc. prob=0.78]

warmup: 11%|β–ˆ | 325/3000 [00:06<00:10, 246.22it/s, 63 steps of size 1.45e-01. acc. prob=0.78]

warmup: 12%|β–ˆβ– | 360/3000 [00:06<00:09, 271.87it/s, 15 steps of size 2.98e-01. acc. prob=0.78]

warmup: 13%|β–ˆβ–Ž | 397/3000 [00:06<00:08, 298.24it/s, 15 steps of size 2.29e-01. acc. prob=0.78]

warmup: 14%|β–ˆβ– | 432/3000 [00:06<00:08, 311.37it/s, 31 steps of size 2.30e-01. acc. prob=0.79]

warmup: 16%|β–ˆβ–Œ | 470/3000 [00:06<00:07, 329.65it/s, 15 steps of size 2.27e-01. acc. prob=0.78]

warmup: 17%|β–ˆβ–‹ | 508/3000 [00:06<00:07, 342.93it/s, 3 steps of size 2.47e-01. acc. prob=0.79]

warmup: 18%|β–ˆβ–Š | 548/3000 [00:06<00:06, 358.23it/s, 7 steps of size 6.67e-01. acc. prob=0.79]

warmup: 20%|β–ˆβ–‰ | 591/3000 [00:06<00:06, 376.79it/s, 7 steps of size 5.90e-01. acc. prob=0.79]

warmup: 21%|β–ˆβ–ˆ | 637/3000 [00:06<00:05, 400.58it/s, 3 steps of size 2.04e-01. acc. prob=0.79]

warmup: 23%|β–ˆβ–ˆβ–Ž | 680/3000 [00:07<00:05, 408.28it/s, 15 steps of size 4.07e-01. acc. prob=0.79]

warmup: 24%|β–ˆβ–ˆβ– | 724/3000 [00:07<00:05, 416.50it/s, 3 steps of size 5.07e-01. acc. prob=0.79]

warmup: 26%|β–ˆβ–ˆβ–Œ | 771/3000 [00:07<00:05, 429.79it/s, 7 steps of size 4.87e-01. acc. prob=0.79]

warmup: 27%|β–ˆβ–ˆβ–‹ | 815/3000 [00:07<00:05, 432.18it/s, 7 steps of size 3.62e-01. acc. prob=0.79]

warmup: 29%|β–ˆβ–ˆβ–Š | 860/3000 [00:07<00:04, 434.33it/s, 7 steps of size 3.21e-01. acc. prob=0.79]

warmup: 30%|β–ˆβ–ˆβ–ˆ | 907/3000 [00:07<00:04, 442.13it/s, 7 steps of size 3.93e-01. acc. prob=0.79]

warmup: 32%|β–ˆβ–ˆβ–ˆβ– | 952/3000 [00:07<00:04, 442.09it/s, 1 steps of size 6.92e-01. acc. prob=0.79]

warmup: 33%|β–ˆβ–ˆβ–ˆβ–Ž | 997/3000 [00:07<00:04, 423.18it/s, 7 steps of size 1.04e+00. acc. prob=0.79]

sample: 35%|β–ˆβ–ˆβ–ˆβ– | 1043/3000 [00:07<00:04, 432.06it/s, 7 steps of size 4.78e-01. acc. prob=0.87]

sample: 36%|β–ˆβ–ˆβ–ˆβ–‹ | 1089/3000 [00:07<00:04, 439.79it/s, 3 steps of size 4.78e-01. acc. prob=0.89]

sample: 38%|β–ˆβ–ˆβ–ˆβ–Š | 1135/3000 [00:08<00:04, 445.37it/s, 7 steps of size 4.78e-01. acc. prob=0.90]

sample: 39%|β–ˆβ–ˆβ–ˆβ–‰ | 1180/3000 [00:08<00:04, 445.88it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 41%|β–ˆβ–ˆβ–ˆβ–ˆ | 1226/3000 [00:08<00:03, 448.73it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 42%|β–ˆβ–ˆβ–ˆβ–ˆβ– | 1272/3000 [00:08<00:03, 451.28it/s, 15 steps of size 4.78e-01. acc. prob=0.91]

sample: 44%|β–ˆβ–ˆβ–ˆβ–ˆβ– | 1319/3000 [00:08<00:03, 454.70it/s, 3 steps of size 4.78e-01. acc. prob=0.91]

sample: 46%|β–ˆβ–ˆβ–ˆβ–ˆβ–Œ | 1365/3000 [00:08<00:03, 453.75it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 47%|β–ˆβ–ˆβ–ˆβ–ˆβ–‹ | 1412/3000 [00:08<00:03, 456.06it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 49%|β–ˆβ–ˆβ–ˆβ–ˆβ–Š | 1458/3000 [00:08<00:03, 454.43it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 50%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆ | 1505/3000 [00:08<00:03, 457.17it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 52%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ– | 1551/3000 [00:08<00:03, 454.04it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 53%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Ž | 1597/3000 [00:09<00:03, 454.26it/s, 3 steps of size 4.78e-01. acc. prob=0.91]

sample: 55%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ– | 1644/3000 [00:09<00:02, 456.43it/s, 3 steps of size 4.78e-01. acc. prob=0.91]

sample: 56%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‹ | 1690/3000 [00:09<00:02, 453.08it/s, 3 steps of size 4.78e-01. acc. prob=0.91]

sample: 58%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Š | 1736/3000 [00:09<00:02, 452.45it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 59%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‰ | 1782/3000 [00:09<00:02, 452.45it/s, 3 steps of size 4.78e-01. acc. prob=0.91]

sample: 61%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ | 1830/3000 [00:09<00:02, 459.70it/s, 3 steps of size 4.78e-01. acc. prob=0.91]

sample: 63%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Ž | 1876/3000 [00:09<00:02, 454.25it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 64%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ– | 1924/3000 [00:09<00:02, 459.96it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 66%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Œ | 1971/3000 [00:09<00:02, 460.15it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 67%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‹ | 2018/3000 [00:09<00:02, 460.92it/s, 3 steps of size 4.78e-01. acc. prob=0.91]

sample: 69%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‰ | 2065/3000 [00:10<00:02, 460.17it/s, 3 steps of size 4.78e-01. acc. prob=0.91]

sample: 70%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ | 2112/3000 [00:10<00:01, 459.57it/s, 3 steps of size 4.78e-01. acc. prob=0.91]

sample: 72%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ– | 2158/3000 [00:10<00:01, 458.23it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 74%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Ž | 2205/3000 [00:10<00:01, 459.51it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 75%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Œ | 2251/3000 [00:10<00:01, 453.38it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 77%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‹ | 2297/3000 [00:10<00:01, 452.63it/s, 3 steps of size 4.78e-01. acc. prob=0.91]

sample: 78%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Š | 2343/3000 [00:10<00:01, 453.82it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 80%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‰ | 2391/3000 [00:10<00:01, 458.93it/s, 3 steps of size 4.78e-01. acc. prob=0.91]

sample: 81%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ | 2437/3000 [00:10<00:01, 456.41it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 83%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Ž | 2483/3000 [00:10<00:01, 451.88it/s, 11 steps of size 4.78e-01. acc. prob=0.91]

sample: 84%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ– | 2530/3000 [00:11<00:01, 455.98it/s, 3 steps of size 4.78e-01. acc. prob=0.91]

sample: 86%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Œ | 2578/3000 [00:11<00:00, 461.68it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 88%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Š | 2625/3000 [00:11<00:00, 458.52it/s, 5 steps of size 4.78e-01. acc. prob=0.91]

sample: 89%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‰ | 2672/3000 [00:11<00:00, 461.10it/s, 1 steps of size 4.78e-01. acc. prob=0.91]

sample: 91%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ | 2720/3000 [00:11<00:00, 464.07it/s, 3 steps of size 4.78e-01. acc. prob=0.91]

sample: 92%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–| 2767/3000 [00:11<00:00, 460.36it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 94%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–| 2814/3000 [00:11<00:00, 458.26it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 95%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Œ| 2860/3000 [00:11<00:00, 458.27it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 97%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‹| 2906/3000 [00:11<00:00, 456.14it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 98%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Š| 2952/3000 [00:12<00:00, 457.16it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 100%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‰| 2998/3000 [00:12<00:00, 455.83it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 100%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ| 3000/3000 [00:12<00:00, 247.55it/s, 3 steps of size 4.78e-01. acc. prob=0.91]

                      mean       std    median      5.0%     95.0%     n_eff     r_hat
           delta      0.90      0.00      0.90      0.89      0.90   2363.86      1.00
  sigma_predator      0.52      0.04      0.52      0.46      0.58   1303.58      1.00
      sigma_prey      0.43      0.03      0.43      0.38      0.48   1137.19      1.00

Number of divergences: 0

png

πŸ‘‰ This worked, so now we can have a look at the results:

# Report the results
sim.report()
[WARNING] Could not fetch resource 'probability_model.png': replacing image with description
[WARNING] This document format requires a nonempty <title> element.
  Defaulting to 'report' as the title.
  To specify a title, use 'title' in metadata or --metadata title="...".

Chapter 2: Saving and retrieving a simulation πŸ’Ύ#

πŸ‘‰ In this chapter, we will save our Pymob simulation and create a new simulation from it. You will see that this makes the process much shorter than above.

πŸ‘‰ Let’s start by saving our configuration and observations.

Caution

The observations have to be saved before the configuration. Otherwise the configuration doesn’t save the location the observations were saved in which causes problems down the line.

# Set the data paths we want to save to and create the necessary folders if they don't exist yet
import os
sim.config.create_directory("scenario", force=True)
sim.config.create_directory("results", force=True)
os.makedirs(sim.data_path, exist_ok=True)

# Save our configuration and observations
sim.save_observations(force=True)
sim.config.save(force=True)
Scenario directory created at '/home/docs/checkouts/readthedocs.org/user_builds/pymob/checkouts/latest/docs/source/user_guide/case_studies/ODEtutorial/scenarios/lotkavolterra'.
Results directory exists at '/home/docs/checkouts/readthedocs.org/user_builds/pymob/checkouts/latest/docs/source/user_guide/case_studies/ODEtutorial/results/lotkavolterra'.

2.1 Creating a new sim file from a saved configuration πŸ†•#

πŸ‘‰ In the next part we try to generate a new simulation object from the configuration file we just created. To do this, we first have to make an additional import:

from pymob import Config

πŸ‘‰ After we’ve done that, we can now create a pymob.config.Config object from our file. This can then be passed to the constructor of pymob.SimulationBase.

# Load configuration to a Config instance
config = Config("case_studies/ODEtutorial/scenarios/lotkavolterra/settings.cfg")

# Create a new simulation from the configuration
sim2 = SimulationBase(config)
/home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/pymob/sim/config/base.py:397: UserWarning: Case study 'ODEtutorial' could not be imported. Install the case study with `pip install ODEtutorial`.
  warnings.warn(

πŸ‘‰ Essentially, passing the pymob.config.Config file to the pymob.SimulationBase constructor just copies it to config.

config == sim2.config
True

πŸ‘‰ Now that our simulation knows about its configuration, we can call the pymob.sim.SimulationBase.initialize() function which prepares all of our data for us. It fetches the observation data from the specified location and handles the initial condition as well as external inputs (which we don’t have here). That means that a well-prepared config file can save a lot of work!

πŸ‘‰ We do, however, still need to specify some additional features of the pymob.sim.SimulationBase object. That includes the model, its parameters and the solver.

Subclassing SimulationBase

By subclassing pymob.SimulationBase and writing a customized initialize() function that also includes these tasks, this can be avoided (see the last three cells of this notebook). But for now, we will keep it simple and do it manually.

# Add data and initial conditions to the simulation
sim2.initialize(config)

# Add model, model parameters, and solver to the simulation
sim2.model = lotkavolterra
sim2.model_parameters["parameters"] = sim2.config.model_parameters.value_dict
sim2.solver = JaxSolver
MinMaxScaler(variable=prey, min=5.844172888098378, max=12.525948698266157)
MinMaxScaler(variable=predator, min=4.053933700151413, max=10.925258075625722)


/home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/pymob/simulation.py:1564: UserWarning: Using default initialize method, (load observations, define 'y0', define 'x_in'). This may be insufficient for more complex simulations.
  warnings.warn(

2.2 Running the model and parameter inference πŸ‘ŸπŸ”#

πŸ‘‰ As before, we want to create an evaluator for running the system. This is essentially the same code as above, let’s see how it goes:

# Put everything in place for running the simulation
sim2.dispatch_constructor()

try:

    # Create an evaluator, run the simulation and obtain the results
    evaluator2 = sim2.dispatch(theta={"delta":0.9})
    evaluator2()

    # Plot the results
    fig, ax = plt.subplots(figsize=(5, 4))
    data_res2 = evaluator2.results
    ax.plot(data_obs.time, data_obs.prey, ls="-", color="tab:blue", alpha=.5, label ="observation data")
    ax.plot(data_obs.time, data_obs.predator, ls="-", color="tab:blue", alpha=.5, label ="observation data")
    ax.plot(data_res2.time, data_res2.prey, color="black", label ="result")
    ax.plot(data_res2.time, data_res2.predator, color="black", label ="result")
    ax.legend()

except ValueError as e:

    # Print the error message
    print("An error occurred:", type(e).__name__, ":", e)

png

πŸ‘‰ If you chose to ignore the note about in the beginning of this notebook and added the initial conditions manually, you should see the following error message now:

ValueError: vmap in_axes must be an int, None, or a tuple of entries corresponding to the positional arguments passed to the function, but got len(in_axes)=6, len(args)=4

πŸ‘‰ The reason for this is that our model takes four parameters (\(\alpha, \beta, \gamma, \delta\)) along with two initial conditions (for prey and predator, respectively) but we only gave it the model parameters. If we had chosen the formulation before, we would have run the following line of code:

sim.config.simulation.y0 = ["prey=10", "predator=5"]

πŸ‘‰ In this case, the function pymob.SimulationBase.initialize() above would have run and added the initial condition \(X = 10, Y = 5\) to sim2. But in our case, because the initial condition has never been defined in the configuration, this doesn’t happen and we get this error. If you ran into this problem, run the following cell which sets the initial conditions manually. Otherwise, just scroll past it.

y0_obs_1 = xr.DataArray(10).to_dataset(name="prey")
y0_obs_2 = xr.DataArray(5).to_dataset(name="predator")
y0_obs = xr.merge([y0_obs_1, y0_obs_2])

sim2.model_parameters["y0"] = y0_obs

# Put everything in place for running the simulation
sim2.dispatch_constructor()

try:

    # Create an evaluator, run the simulation and obtain the results
    evaluator2 = sim2.dispatch(theta={"delta":0.9})
    evaluator2()

    # Plot the results
    fig, ax = plt.subplots(figsize=(5, 4))
    data_res2 = evaluator2.results
    ax.plot(data_obs.time, data_obs.prey, ls="-", color="tab:blue", alpha=.5, label ="observation data")
    ax.plot(data_obs.time, data_obs.predator, ls="-", color="tab:blue", alpha=.5, label ="observation data")
    ax.plot(data_res2.time, data_res2.prey, color="black", label ="result")
    ax.plot(data_res2.time, data_res2.predator, color="black", label ="result")
    ax.legend()

except ValueError as e:

    # Print the error message
    print("An error occurred:", type(e).__name__, ":", e)

png

πŸ‘‰ Now let’s start the parameter inference again. The result should be the same as before.

# Create the inferer (NumPyro backend, NUTS kernel) and let it do its work
sim2.set_inferer("numpyro")
sim2.inferer.config.inference_numpyro.kernel = "nuts"
sim2.inferer.run()

# Plot the results
sim2.config.simulation.x_dimension = "time"
sim2.posterior_predictive_checks(pred_hdi_style={"alpha": 0.1})
Jax 64 bit mode: False
Absolute tolerance: 1e-07


/home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/pymob/inference/numpyro_backend.py:565: UserWarning: Model is not rendered, because the graphviz executable is not found. Try search for 'graphviz executables not found' and the used OS. This should be an easy fix :-)
  warnings.warn(


      Trace Shapes:      
       Param Sites:      
      Sample Sites:      
         delta dist     |
              value     |
    sigma_prey dist     |
              value     |
sigma_predator dist     |
              value     |
      prey_obs dist 101 |
              value 101 |
  predator_obs dist 101 |
              value 101 |

0%| | 0/3000 [00:00<?, ?it/s]

warmup: 0%| | 1/3000 [00:04<3:46:52, 4.54s/it, 1 steps of size 1.87e+00. acc. prob=0.00]

warmup: 1%|▏ | 38/3000 [00:04<04:17, 11.50it/s, 1 steps of size 2.44e-04. acc. prob=0.69]

warmup: 2%|▏ | 60/3000 [00:04<02:29, 19.66it/s, 127 steps of size 1.98e-03. acc. prob=0.74]

warmup: 3%|β–Ž | 78/3000 [00:04<01:45, 27.60it/s, 3 steps of size 1.78e-03. acc. prob=0.75]

warmup: 3%|β–Ž | 94/3000 [00:05<01:20, 36.03it/s, 43 steps of size 2.08e-03. acc. prob=0.75]

warmup: 4%|β–Ž | 109/3000 [00:05<01:09, 41.85it/s, 31 steps of size 2.05e-01. acc. prob=0.76]

warmup: 4%|▍ | 130/3000 [00:05<00:48, 58.85it/s, 63 steps of size 8.87e-02. acc. prob=0.77]

warmup: 5%|β–Œ | 154/3000 [00:05<00:34, 81.87it/s, 63 steps of size 5.77e-02. acc. prob=0.77]

warmup: 6%|β–Œ | 176/3000 [00:05<00:27, 102.45it/s, 31 steps of size 1.67e-01. acc. prob=0.77]

warmup: 7%|β–‹ | 203/3000 [00:05<00:21, 131.52it/s, 27 steps of size 1.65e-01. acc. prob=0.78]

warmup: 8%|β–Š | 230/3000 [00:05<00:17, 158.99it/s, 15 steps of size 2.18e-01. acc. prob=0.78]

warmup: 9%|β–Š | 262/3000 [00:05<00:14, 194.86it/s, 31 steps of size 3.40e-01. acc. prob=0.78]

warmup: 10%|β–‰ | 288/3000 [00:06<00:12, 210.07it/s, 15 steps of size 2.83e-01. acc. prob=0.78]

warmup: 11%|β–ˆ | 323/3000 [00:06<00:10, 245.63it/s, 15 steps of size 4.89e-01. acc. prob=0.78]

warmup: 12%|β–ˆβ– | 354/3000 [00:06<00:10, 261.38it/s, 15 steps of size 3.30e-01. acc. prob=0.78]

warmup: 13%|β–ˆβ–Ž | 390/3000 [00:06<00:09, 288.17it/s, 11 steps of size 4.12e-01. acc. prob=0.78]

warmup: 14%|β–ˆβ– | 426/3000 [00:06<00:08, 306.36it/s, 31 steps of size 1.88e-01. acc. prob=0.78]

warmup: 15%|β–ˆβ–Œ | 462/3000 [00:06<00:07, 321.40it/s, 15 steps of size 5.17e-01. acc. prob=0.79]

warmup: 17%|β–ˆβ–‹ | 500/3000 [00:06<00:07, 337.63it/s, 15 steps of size 2.32e-01. acc. prob=0.79]

warmup: 18%|β–ˆβ–Š | 539/3000 [00:06<00:06, 352.41it/s, 15 steps of size 3.07e-01. acc. prob=0.79]

warmup: 19%|β–ˆβ–‰ | 581/3000 [00:06<00:06, 370.18it/s, 3 steps of size 5.44e-01. acc. prob=0.79]

warmup: 21%|β–ˆβ–ˆ | 625/3000 [00:06<00:06, 390.59it/s, 3 steps of size 3.74e-01. acc. prob=0.79]

warmup: 22%|β–ˆβ–ˆβ– | 669/3000 [00:07<00:05, 403.61it/s, 7 steps of size 3.96e-01. acc. prob=0.79]

warmup: 24%|β–ˆβ–ˆβ–Ž | 711/3000 [00:07<00:05, 406.59it/s, 7 steps of size 3.43e-01. acc. prob=0.79]

warmup: 25%|β–ˆβ–ˆβ–Œ | 757/3000 [00:07<00:05, 422.19it/s, 3 steps of size 2.47e-01. acc. prob=0.79]

warmup: 27%|β–ˆβ–ˆβ–‹ | 801/3000 [00:07<00:05, 426.44it/s, 1 steps of size 2.34e-01. acc. prob=0.79]

warmup: 28%|β–ˆβ–ˆβ–Š | 844/3000 [00:07<00:05, 426.73it/s, 7 steps of size 4.33e-01. acc. prob=0.79]

warmup: 30%|β–ˆβ–ˆβ–‰ | 889/3000 [00:07<00:04, 431.56it/s, 3 steps of size 6.88e-01. acc. prob=0.79]

warmup: 31%|β–ˆβ–ˆβ–ˆ | 934/3000 [00:07<00:04, 436.30it/s, 7 steps of size 3.88e-01. acc. prob=0.79]

warmup: 33%|β–ˆβ–ˆβ–ˆβ–Ž | 978/3000 [00:07<00:04, 420.48it/s, 7 steps of size 8.69e-01. acc. prob=0.79]

sample: 34%|β–ˆβ–ˆβ–ˆβ– | 1021/3000 [00:07<00:04, 422.16it/s, 3 steps of size 4.78e-01. acc. prob=0.82]

sample: 36%|β–ˆβ–ˆβ–ˆβ–Œ | 1066/3000 [00:07<00:04, 427.83it/s, 7 steps of size 4.78e-01. acc. prob=0.88]

sample: 37%|β–ˆβ–ˆβ–ˆβ–‹ | 1113/3000 [00:08<00:04, 437.20it/s, 7 steps of size 4.78e-01. acc. prob=0.90]

sample: 39%|β–ˆβ–ˆβ–ˆβ–Š | 1158/3000 [00:08<00:04, 439.45it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 40%|β–ˆβ–ˆβ–ˆβ–ˆ | 1202/3000 [00:08<00:04, 436.54it/s, 3 steps of size 4.78e-01. acc. prob=0.91]

sample: 42%|β–ˆβ–ˆβ–ˆβ–ˆβ– | 1248/3000 [00:08<00:03, 440.94it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 43%|β–ˆβ–ˆβ–ˆβ–ˆβ–Ž | 1293/3000 [00:08<00:03, 443.32it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 45%|β–ˆβ–ˆβ–ˆβ–ˆβ– | 1339/3000 [00:08<00:03, 447.54it/s, 3 steps of size 4.78e-01. acc. prob=0.91]

sample: 46%|β–ˆβ–ˆβ–ˆβ–ˆβ–Œ | 1385/3000 [00:08<00:03, 448.27it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 48%|β–ˆβ–ˆβ–ˆβ–ˆβ–Š | 1432/3000 [00:08<00:03, 453.15it/s, 3 steps of size 4.78e-01. acc. prob=0.91]

sample: 49%|β–ˆβ–ˆβ–ˆβ–ˆβ–‰ | 1478/3000 [00:08<00:03, 449.24it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 51%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆ | 1524/3000 [00:08<00:03, 449.67it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 52%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ– | 1569/3000 [00:09<00:03, 447.75it/s, 5 steps of size 4.78e-01. acc. prob=0.91]

sample: 54%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ– | 1615/3000 [00:09<00:03, 449.82it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 55%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Œ | 1661/3000 [00:09<00:02, 451.77it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 57%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‹ | 1707/3000 [00:09<00:02, 444.67it/s, 3 steps of size 4.78e-01. acc. prob=0.91]

sample: 58%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Š | 1753/3000 [00:09<00:02, 446.75it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 60%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‰ | 1798/3000 [00:09<00:02, 447.69it/s, 3 steps of size 4.78e-01. acc. prob=0.91]

sample: 61%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ– | 1844/3000 [00:09<00:02, 451.20it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 63%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Ž | 1890/3000 [00:09<00:02, 450.19it/s, 3 steps of size 4.78e-01. acc. prob=0.91]

sample: 65%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ– | 1936/3000 [00:09<00:02, 451.98it/s, 3 steps of size 4.78e-01. acc. prob=0.91]

sample: 66%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Œ | 1982/3000 [00:10<00:02, 453.62it/s, 3 steps of size 4.78e-01. acc. prob=0.91]

sample: 68%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Š | 2028/3000 [00:10<00:02, 453.89it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 69%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‰ | 2074/3000 [00:10<00:02, 454.70it/s, 3 steps of size 4.78e-01. acc. prob=0.91]

sample: 71%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ | 2120/3000 [00:10<00:01, 450.89it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 72%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ– | 2166/3000 [00:10<00:01, 449.82it/s, 1 steps of size 4.78e-01. acc. prob=0.91]

sample: 74%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Ž | 2212/3000 [00:10<00:01, 449.95it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 75%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Œ | 2258/3000 [00:10<00:01, 447.14it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 77%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‹ | 2304/3000 [00:10<00:01, 448.94it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 78%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Š | 2350/3000 [00:10<00:01, 451.14it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 80%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‰ | 2397/3000 [00:10<00:01, 455.29it/s, 3 steps of size 4.78e-01. acc. prob=0.91]

sample: 81%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ– | 2443/3000 [00:11<00:01, 451.58it/s, 3 steps of size 4.78e-01. acc. prob=0.91]

sample: 83%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Ž | 2489/3000 [00:11<00:01, 445.60it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 85%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ– | 2536/3000 [00:11<00:01, 450.02it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 86%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Œ | 2584/3000 [00:11<00:00, 455.72it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 88%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Š | 2630/3000 [00:11<00:00, 451.99it/s, 3 steps of size 4.78e-01. acc. prob=0.91]

sample: 89%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‰ | 2676/3000 [00:11<00:00, 453.56it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 91%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ | 2723/3000 [00:11<00:00, 455.79it/s, 3 steps of size 4.78e-01. acc. prob=0.91]

sample: 92%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–| 2769/3000 [00:11<00:00, 452.22it/s, 3 steps of size 4.78e-01. acc. prob=0.91]

sample: 94%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–| 2815/3000 [00:11<00:00, 450.93it/s, 3 steps of size 4.78e-01. acc. prob=0.91]

sample: 95%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Œ| 2861/3000 [00:11<00:00, 449.88it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 97%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‹| 2906/3000 [00:12<00:00, 448.04it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 98%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Š| 2952/3000 [00:12<00:00, 449.01it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 100%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‰| 2997/3000 [00:12<00:00, 446.76it/s, 3 steps of size 4.78e-01. acc. prob=0.91]

sample: 100%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ| 3000/3000 [00:12<00:00, 244.53it/s, 3 steps of size 4.78e-01. acc. prob=0.91]

                      mean       std    median      5.0%     95.0%     n_eff     r_hat
           delta      0.90      0.00      0.90      0.89      0.90   2363.86      1.00
  sigma_predator      0.52      0.04      0.52      0.46      0.58   1303.58      1.00
      sigma_prey      0.43      0.03      0.43      0.38      0.48   1137.19      1.00

Number of divergences: 0

png

2.3 Summary#

πŸ‘‰ Creating the simulation from a pre-saved configuration saved us the following steps:

  • Adding data to the simulation

  • If done right: Adding initial conditions to the simulation

  • Creating the Lotka-Volterra parameters

  • Specifying the error model along with the corresponding parameters

  • Telling the evaluator not to throw exceptions if max_steps is exceeded

  • Chossing a prior for parameter inference

πŸ‘‰ We still had to:

  • Define a model

  • Pass parameter values to the simulation

  • Specify the solver

πŸ‘‰ By subclassing pymob.SimulationBase, even those last steps can be avoided. This will be explained in detail in another tutorial as it mostly makes sense in the context of case studies. But in this case, it is pretty straightforward, so here’s a little sneak peek:

# Define the simulation class
class LotkaVolterraSim(SimulationBase):
    model = lotkavolterra
    solver = JaxSolver
    def initialize(self, input=None):
        super().initialize(input)
        self.model_parameters["parameters"] = self.config.model_parameters.value_dict
        self.dispatch_constructor()
    
# Create and initialize simulation (no further steps necessary)
sim3 = LotkaVolterraSim("case_studies/ODEtutorial/scenarios/lotkavolterra/settings.cfg")
sim3.initialize()
MinMaxScaler(variable=prey, min=5.844172888098378, max=12.525948698266157)
MinMaxScaler(variable=predator, min=4.053933700151413, max=10.925258075625722)


/home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/pymob/sim/config/base.py:397: UserWarning: Case study 'ODEtutorial' could not be imported. Install the case study with `pip install ODEtutorial`.
  warnings.warn(
/home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/pymob/simulation.py:1564: UserWarning: Using default initialize method, (load observations, define 'y0', define 'x_in'). This may be insufficient for more complex simulations.
  warnings.warn(
# Put everything in place for running the simulation
sim3.dispatch_constructor()

try:

    # Create an evaluator, run the simulation and obtain the results
    evaluator3 = sim3.dispatch(theta={"delta":0.9})
    evaluator3()

    # Plot the results
    fig, ax = plt.subplots(figsize=(5, 4))
    data_res3 = evaluator3.results
    ax.plot(data_obs.time, data_obs.prey, ls="-", color="tab:blue", alpha=.5, label ="observation data")
    ax.plot(data_obs.time, data_obs.predator, ls="-", color="tab:blue", alpha=.5, label ="observation data")
    ax.plot(data_res3.time, data_res3.prey, color="black", label ="result")
    ax.plot(data_res3.time, data_res3.predator, color="black", label ="result")
    ax.legend()

except ValueError as e:

    # Print the error message
    print("An error occurred:", type(e).__name__, ":", e)

png

# Create the inferer (NumPyro backend, NUTS kernel) and let it do its work
sim3.set_inferer("numpyro")
sim3.inferer.config.inference_numpyro.kernel = "nuts"
sim3.inferer.run()

# Plot the results
sim3.config.simulation.x_dimension = "time"
sim3.posterior_predictive_checks(pred_hdi_style={"alpha": 0.1})
Jax 64 bit mode: False
Absolute tolerance: 1e-07


/home/docs/checkouts/readthedocs.org/user_builds/pymob/envs/latest/lib/python3.11/site-packages/pymob/inference/numpyro_backend.py:565: UserWarning: Model is not rendered, because the graphviz executable is not found. Try search for 'graphviz executables not found' and the used OS. This should be an easy fix :-)
  warnings.warn(


      Trace Shapes:      
       Param Sites:      
      Sample Sites:      
         delta dist     |
              value     |
    sigma_prey dist     |
              value     |
sigma_predator dist     |
              value     |
      prey_obs dist 101 |
              value 101 |
  predator_obs dist 101 |
              value 101 |

0%| | 0/3000 [00:00<?, ?it/s]

warmup: 0%| | 1/3000 [00:04<3:48:41, 4.58s/it, 1 steps of size 1.87e+00. acc. prob=0.00]

warmup: 1%| | 37/3000 [00:04<04:26, 11.11it/s, 1 steps of size 4.43e-04. acc. prob=0.70]

warmup: 2%|▏ | 59/3000 [00:04<02:32, 19.33it/s, 255 steps of size 1.25e-03. acc. prob=0.73]

warmup: 3%|β–Ž | 77/3000 [00:05<01:48, 27.05it/s, 63 steps of size 4.27e-03. acc. prob=0.75]

warmup: 3%|β–Ž | 92/3000 [00:05<01:23, 34.96it/s, 23 steps of size 9.46e-04. acc. prob=0.75]

warmup: 4%|β–Ž | 106/3000 [00:05<01:11, 40.33it/s, 31 steps of size 2.96e-01. acc. prob=0.76]

warmup: 4%|▍ | 129/3000 [00:05<00:48, 59.19it/s, 127 steps of size 6.20e-02. acc. prob=0.77]

warmup: 5%|β–Œ | 153/3000 [00:05<00:34, 82.10it/s, 3 steps of size 4.38e-02. acc. prob=0.77]

warmup: 6%|β–Œ | 172/3000 [00:05<00:28, 98.57it/s, 19 steps of size 1.60e-01. acc. prob=0.77]

warmup: 7%|β–‹ | 199/3000 [00:05<00:21, 128.71it/s, 35 steps of size 1.78e-01. acc. prob=0.78]

warmup: 8%|β–Š | 226/3000 [00:05<00:17, 156.49it/s, 63 steps of size 1.13e-01. acc. prob=0.78]

warmup: 9%|β–Š | 258/3000 [00:05<00:14, 192.22it/s, 7 steps of size 6.43e-01. acc. prob=0.78]

warmup: 10%|β–‰ | 285/3000 [00:06<00:12, 210.59it/s, 15 steps of size 5.37e-01. acc. prob=0.78]

warmup: 11%|β–ˆ | 320/3000 [00:06<00:10, 245.93it/s, 7 steps of size 1.56e-01. acc. prob=0.78]

warmup: 12%|β–ˆβ– | 352/3000 [00:06<00:10, 264.08it/s, 7 steps of size 1.74e-01. acc. prob=0.78]

warmup: 13%|β–ˆβ–Ž | 388/3000 [00:06<00:09, 288.50it/s, 15 steps of size 3.20e-01. acc. prob=0.78]

warmup: 14%|β–ˆβ– | 425/3000 [00:06<00:08, 309.72it/s, 7 steps of size 1.47e-01. acc. prob=0.78]

warmup: 15%|β–ˆβ–Œ | 461/3000 [00:06<00:07, 321.53it/s, 15 steps of size 2.82e-01. acc. prob=0.78]

warmup: 16%|β–ˆβ–‹ | 495/3000 [00:06<00:07, 325.56it/s, 7 steps of size 2.89e-01. acc. prob=0.79]

warmup: 18%|β–ˆβ–Š | 535/3000 [00:06<00:07, 346.52it/s, 7 steps of size 4.85e-01. acc. prob=0.79]

warmup: 19%|β–ˆβ–‰ | 574/3000 [00:06<00:06, 357.40it/s, 3 steps of size 3.31e-01. acc. prob=0.79]

warmup: 21%|β–ˆβ–ˆ | 618/3000 [00:06<00:06, 380.90it/s, 7 steps of size 5.66e-01. acc. prob=0.79]

warmup: 22%|β–ˆβ–ˆβ– | 662/3000 [00:07<00:05, 397.63it/s, 7 steps of size 3.56e-01. acc. prob=0.79]

warmup: 23%|β–ˆβ–ˆβ–Ž | 703/3000 [00:07<00:05, 398.93it/s, 7 steps of size 4.10e-01. acc. prob=0.79]

warmup: 25%|β–ˆβ–ˆβ–Œ | 750/3000 [00:07<00:05, 417.22it/s, 7 steps of size 4.02e-01. acc. prob=0.79]

warmup: 27%|β–ˆβ–ˆβ–‹ | 796/3000 [00:07<00:05, 426.47it/s, 15 steps of size 3.14e-01. acc. prob=0.79]

warmup: 28%|β–ˆβ–ˆβ–Š | 839/3000 [00:07<00:05, 427.21it/s, 3 steps of size 4.62e-01. acc. prob=0.79]

warmup: 29%|β–ˆβ–ˆβ–‰ | 882/3000 [00:07<00:04, 427.50it/s, 1 steps of size 3.86e-01. acc. prob=0.79]

warmup: 31%|β–ˆβ–ˆβ–ˆ | 929/3000 [00:07<00:04, 437.93it/s, 7 steps of size 5.46e-01. acc. prob=0.79]

warmup: 32%|β–ˆβ–ˆβ–ˆβ– | 973/3000 [00:07<00:04, 425.62it/s, 3 steps of size 1.12e+00. acc. prob=0.79]

sample: 34%|β–ˆβ–ˆβ–ˆβ– | 1016/3000 [00:07<00:04, 421.79it/s, 1 steps of size 4.78e-01. acc. prob=0.81]

sample: 35%|β–ˆβ–ˆβ–ˆβ–Œ | 1060/3000 [00:08<00:04, 426.38it/s, 7 steps of size 4.78e-01. acc. prob=0.88]

sample: 37%|β–ˆβ–ˆβ–ˆβ–‹ | 1107/3000 [00:08<00:04, 438.16it/s, 7 steps of size 4.78e-01. acc. prob=0.89]

sample: 38%|β–ˆβ–ˆβ–ˆβ–Š | 1152/3000 [00:08<00:04, 441.52it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 40%|β–ˆβ–ˆβ–ˆβ–‰ | 1197/3000 [00:08<00:04, 439.73it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 41%|β–ˆβ–ˆβ–ˆβ–ˆβ– | 1242/3000 [00:08<00:03, 442.25it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 43%|β–ˆβ–ˆβ–ˆβ–ˆβ–Ž | 1288/3000 [00:08<00:03, 446.99it/s, 3 steps of size 4.78e-01. acc. prob=0.91]

sample: 44%|β–ˆβ–ˆβ–ˆβ–ˆβ– | 1334/3000 [00:08<00:03, 448.27it/s, 15 steps of size 4.78e-01. acc. prob=0.91]

sample: 46%|β–ˆβ–ˆβ–ˆβ–ˆβ–Œ | 1380/3000 [00:08<00:03, 449.36it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 48%|β–ˆβ–ˆβ–ˆβ–ˆβ–Š | 1426/3000 [00:08<00:03, 450.89it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 49%|β–ˆβ–ˆβ–ˆβ–ˆβ–‰ | 1472/3000 [00:08<00:03, 447.37it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 51%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆ | 1517/3000 [00:09<00:03, 447.53it/s, 3 steps of size 4.78e-01. acc. prob=0.91]

sample: 52%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ– | 1562/3000 [00:09<00:03, 444.75it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 54%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Ž | 1608/3000 [00:09<00:03, 447.60it/s, 3 steps of size 4.78e-01. acc. prob=0.91]

sample: 55%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Œ | 1654/3000 [00:09<00:02, 449.51it/s, 3 steps of size 4.78e-01. acc. prob=0.91]

sample: 57%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‹ | 1699/3000 [00:09<00:02, 444.08it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 58%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Š | 1745/3000 [00:09<00:02, 447.80it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 60%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‰ | 1790/3000 [00:09<00:02, 446.73it/s, 3 steps of size 4.78e-01. acc. prob=0.91]

sample: 61%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ | 1837/3000 [00:09<00:02, 453.52it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 63%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Ž | 1883/3000 [00:09<00:02, 447.69it/s, 3 steps of size 4.78e-01. acc. prob=0.91]

sample: 64%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ– | 1930/3000 [00:09<00:02, 452.20it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 66%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Œ | 1976/3000 [00:10<00:02, 454.45it/s, 3 steps of size 4.78e-01. acc. prob=0.91]

sample: 67%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‹ | 2023/3000 [00:10<00:02, 456.23it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 69%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‰ | 2069/3000 [00:10<00:02, 454.49it/s, 3 steps of size 4.78e-01. acc. prob=0.91]

sample: 70%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ | 2115/3000 [00:10<00:01, 453.16it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 72%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ– | 2161/3000 [00:10<00:01, 453.63it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 74%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Ž | 2207/3000 [00:10<00:01, 455.18it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 75%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Œ | 2253/3000 [00:10<00:01, 451.31it/s, 3 steps of size 4.78e-01. acc. prob=0.91]

sample: 77%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‹ | 2299/3000 [00:10<00:01, 448.64it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 78%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Š | 2344/3000 [00:10<00:01, 448.17it/s, 1 steps of size 4.78e-01. acc. prob=0.91]

sample: 80%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‰ | 2390/3000 [00:10<00:01, 450.23it/s, 3 steps of size 4.78e-01. acc. prob=0.91]

sample: 81%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ | 2436/3000 [00:11<00:01, 449.01it/s, 3 steps of size 4.78e-01. acc. prob=0.91]

sample: 83%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Ž | 2481/3000 [00:11<00:01, 444.93it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 84%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ– | 2527/3000 [00:11<00:01, 447.31it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 86%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Œ | 2574/3000 [00:11<00:00, 452.43it/s, 3 steps of size 4.78e-01. acc. prob=0.91]

sample: 87%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‹ | 2620/3000 [00:11<00:00, 450.39it/s, 3 steps of size 4.78e-01. acc. prob=0.91]

sample: 89%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‰ | 2667/3000 [00:11<00:00, 453.67it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 90%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ | 2714/3000 [00:11<00:00, 457.68it/s, 3 steps of size 4.78e-01. acc. prob=0.91]

sample: 92%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–| 2760/3000 [00:11<00:00, 454.29it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 94%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Ž| 2806/3000 [00:11<00:00, 451.84it/s, 3 steps of size 4.78e-01. acc. prob=0.91]

sample: 95%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Œ| 2852/3000 [00:11<00:00, 450.60it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 97%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‹| 2898/3000 [00:12<00:00, 448.87it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 98%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Š| 2944/3000 [00:12<00:00, 450.03it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 100%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‰| 2990/3000 [00:12<00:00, 448.32it/s, 7 steps of size 4.78e-01. acc. prob=0.91]

sample: 100%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ| 3000/3000 [00:12<00:00, 243.61it/s, 3 steps of size 4.78e-01. acc. prob=0.91]

                      mean       std    median      5.0%     95.0%     n_eff     r_hat
           delta      0.90      0.00      0.90      0.89      0.90   2363.86      1.00
  sigma_predator      0.52      0.04      0.52      0.46      0.58   1303.58      1.00
      sigma_prey      0.43      0.03      0.43      0.38      0.48   1137.19      1.00

Number of divergences: 0

png