# Hierarchical Predator Prey modelling The Lotka-Volterra predator-prey model is the archetypical model for dynamical systems, depicting the fluctuating population development of the dynamical system. It is simple enough to fit parameters and estimate their uncertainty in a single replicate. But what if there was some environmental fluctuation we wanted ```python import numpy as np import arviz as az import xarray as xr import matplotlib.pyplot as plt import preliz as pz import jax jax.config.update("jax_enable_x64", True) from pymob import Config from pymob.inference.scipy_backend import ScipyBackend from pymob.sim.parameters import Param, RandomVariable from pymob.sim.config import Modelparameters from pymob.solvers.diffrax import JaxSolver from pymob.inference.analysis import plot_pair ``` WARNING (pytensor.tensor.blas): Using NumPy C-API based implementation for BLAS functions. ```python # import case study and simulation config = Config("../scenarios/lotka_volterra_hierarchical_hyperpriors/settings.cfg") config.case_study.package = "../.." config.case_study.scenario = "lotka_volterra_hierarchical_vaying_y0" config.import_casestudy_modules(reset_path=True) from sim import HierarchicalSimulation sim = HierarchicalSimulation(config) sim.setup() ``` Inserted './../..' in PATH at index=0 Inserted './../../lotka_volterra_case_study' in PATH at index=0 MinMaxScaler(variable=rabbits, min=0.0, max=1329.0) MinMaxScaler(variable=wolves, min=0.0, max=1019.0) Results directory exists at '/home/flo-schu/projects/pymob/case_studies/lotka_volterra_case_study/results/lotka_volterra_hierarchical_vaying_y0'. Scenario directory exists at '/home/flo-schu/projects/pymob/case_studies/lotka_volterra_case_study/scenarios/lotka_volterra_hierarchical_vaying_y0'. ## Investigate the structure of $y_0$ For simulating our artificial data (`hierarchical_model.ipynb`), we assumed some initial values of $y_0$. The $y_0$ values were generated from a uniform distribution between 2 and 15 for wolves and a uniform distribution between 35 and 70. Then, after simulating the observations, a poisson noise model was added on top of the deterministic simulation. So far, we have assumed that the noisy observation at $t=0$ are the true initial values for the simulation. To demonstrate this effect. We look at two trajectories that have different starting values ```python # expand time coordinates and constrain index coordinates for demonstration purposes sim.coordinates["time"] = np.linspace(0,10,100) # sim.coordinates["id"] = np.arange(0, 3) sim.dispatch_constructor() # TODO: Only partially replace the y0 values (like in theta) e = sim.dispatch( theta={"alpha": 1, "beta": 0.02}, y0={"rabbits": [50], "wolves":np.arange(1,121)} ) e() fig, (ax1, ax2) = plt.subplots(2,1) for i in sim.coordinates["id"]: e.results.sel(id=i).wolves.plot(ax=ax1, label=f"id={i}") e.results.sel(id=i).rabbits.plot(ax=ax2, label=f"id={i}") # plt.legend() e.results ```

<xarray.Dataset>
Dimensions:               (id: 120, time: 100)
Coordinates:
  * id                    (id) int32 0 1 2 3 4 5 6 ... 114 115 116 117 118 119
  * time                  (time) float64 0.0 0.101 0.202 ... 9.798 9.899 10.0
    rabbit_species        (id) object 'Cottontail' 'Cottontail' ... 'Jackrabbit'
    experiment            (id) object '2010' '2010' '2010' ... '2012' '2012'
    rabbit_species_index  (id) int64 0 0 0 0 0 0 0 0 0 0 ... 1 1 1 1 1 1 1 1 1 1
    experiment_index      (id) int64 0 0 0 0 0 0 0 0 0 0 ... 2 2 2 2 2 2 2 2 2 2
Data variables:
    rabbits               (id, time) float64 50.0 55.2 60.94 ... 27.61 29.72
    wolves                (id, time) float64 1.0 1.023 1.052 ... 13.67 13.65

![png](hierarchical_model_varying_y0_files/hierarchical_model_varying_y0_4_1.png) In this *mild* case, only the starting population of the rabbits vary (56, 44), while wolves are identical. Despite, we see quite some differences in the dynamic, although the model parameters are the same. We need to esimtate the true $y_0$ values to remove this bias. Assuming that $y_0$ is not known, means we also have to define a prior for the starting values and draw realizations of the starting population from a distribution. This gives us two approaches: 1. We know nothing about the true initial population. This would result in a Uniform prior over the entire span of the data and then add some more, because the true value could lie above or below the range (in our case it will lie only above). 2. We know the observed $y_0$ value and use this as a mean for a prior distribution and assume the error of this prior is the same for each initial value accross all experiments. This can of course become arbitrarily complex, where we could assume that the error on the initial value is different from year to year or species to species, but saying the error on the prior distribution for y0 is always the same seems to be a good first approximation (and we know it's true.) In order to not make our lives harder for an artificial problem, lets take a look at the distributions of the starting values. ```python y0 = sim.parse_input("y0", reference_data=sim.observations, drop_dims=["time"]) unif_wolves = pz.Uniform() pois_wolves = pz.Poisson() lnorm_wolves = pz.LogNormal() gamma_wolves = pz.Gamma() _, ax = pz.mle([pois_wolves, unif_wolves, lnorm_wolves, gamma_wolves], y0["wolves"], plot=4) unif_rabbits = pz.Uniform() pois_rabbits = pz.Poisson() lnorm_rabbits = pz.LogNormal() gamma_rabbits = pz.Gamma() _, ax = pz.mle([pois_rabbits, unif_rabbits, lnorm_wolves, gamma_rabbits], y0["rabbits"], plot=4) ``` ![png](hierarchical_model_varying_y0_files/hierarchical_model_varying_y0_6_0.png) ## Fitting the initial values ```python sim.config.jaxsolver.diffrax_solver = "Dopri5" sim.config.jaxsolver.atol = 1e-12 sim.config.jaxsolver.rtol = 1e-10 ``` ```python wolves_y0 = Param(value=8, dims=("id",), prior="lognorm(scale=4,s=0.6)") rabbits_y0 = Param(value=60, dims=("id",), prior="lognorm(scale=53,s=0.2)") sim.config.model_parameters.wolves_y0 = wolves_y0 sim.config.model_parameters.rabbits_y0 = rabbits_y0 sim.config.model_parameters.beta.prior = "lognorm(scale=0.02,s=2)" sim.config.model_parameters.alpha_species_mu.prior = "halfnorm(scale=5)" sim.config.model_parameters.alpha_species_sigma.prior = "halfnorm(scale=1)" sim.config.model_parameters.alpha_species.prior = "lognorm(scale=[alpha_species_mu],s=alpha_species_sigma)" ``` ```python sim.reset_coordinate("time") sim.config.inference_numpyro.kernel = "svi" sim.dispatch_constructor() sim.set_inferer("numpyro") sim.config.inference.n_predictions = 50 sim.prior_predictive_checks() sim.inferer.prior ``` /home/flo-schu/projects/pymob/pymob/sim/plot.py:155: UserWarning: There were 2 NaN or Inf values in the idata group 'prior_predictive'. See Simulation.inf_preds for a mask with the coordinates. warnings.warn( /home/flo-schu/projects/pymob/pymob/sim/plot.py:155: UserWarning: There were 2 NaN or Inf values in the idata group 'prior_predictive'. See Simulation.inf_preds for a mask with the coordinates. warnings.warn( {'alpha_species_mu': HalfNormalTrans(scale=5, dims=('rabbit_species=2',), obs=None), 'alpha_species_sigma': HalfNormalTrans(scale=1, dims=(), obs=None), 'alpha_species': LogNormalTrans(loc=[alpha_species_mu], scale=alpha_species_sigma, dims=('experiment=3', 'rabbit_species=2'), obs=None), 'alpha_sigma': HalfNormalTrans(scale=1, dims=(), obs=None), 'alpha': LogNormalTrans(scale=alpha_sigma, loc=alpha_species[experiment_index, rabbit_species_index], dims=('id=120',), obs=None), 'beta': LogNormalTrans(loc=0.02, scale=2, dims=(), obs=None), 'wolves_y0': LogNormalTrans(loc=4, scale=0.6, dims=('id=120',), obs=None), 'rabbits_y0': LogNormalTrans(loc=53, scale=0.2, dims=('id=120',), obs=None)} ![png](hierarchical_model_varying_y0_files/hierarchical_model_varying_y0_10_2.png) ```python if False: # sim.inferer.run() sim.inferer.store_results(f"{sim.output_path}/numpyro_svi_posterior.nc") else: sim.inferer.load_results("numpyro_svi_posterior.nc") ``` ```python sim.inferer.idata.posterior.beta.mean(("chain", "draw")) ```

<xarray.DataArray 'beta' ()>
array(0.01748016)

```python sim.posterior_predictive_checks() ``` ![png](hierarchical_model_varying_y0_files/hierarchical_model_varying_y0_13_0.png) ```python az.hdi(sim.inferer.idata.posterior["beta"], hdi_prob=0.95) ```

<xarray.Dataset>
Dimensions:  (hdi: 2)
Coordinates:
  * hdi      (hdi) <U6 'lower' 'higher'
Data variables:
    beta     (hdi) float64 0.01709 0.01786

```python fig, ax1 = plt.subplots(1, 1, figsize=(4,20)) az.plot_forest( data=[sim.inferer.idata.posterior], var_names=["beta", "alpha_species_mu", "alpha_species", "alpha"], ax=ax1, combined=True, hdi_prob=0.95, textsize=8 ) ax1.vlines(0.017648710084435453,*ax1.get_ylim(), color="black") ax1.vlines(1,*ax1.get_ylim(), color="black") ax1.vlines(3,*ax1.get_ylim(), color="black") ``` ![png](hierarchical_model_varying_y0_files/hierarchical_model_varying_y0_15_1.png) It seems we are nailing it already. The parameter estimates provided by SVI contain the true values in their estimate. Yay 🎉 We see that the alpha values of the *Jackrabbit* species vary more than the *Cottontail* alphas. This is caused by using lognormal priors for generating the alpha values for the IDs. The underestimation of the alpha_species_mu posterior parameter esimtate of the Jackrabbit species could originate from stochasticity in the data generation (drawing of alpha values). Parameter estimation was also successfully achieved from pretty uninformative distributions. This also is a success The downside is that NUTS takes a long time. The only thing up next is using our initially observed values as prior means for the initial values ```python rabbits_y0_mu = str(sim.model_parameters["y0"]["rabbits"].values.tolist()).replace(" ", "") wolves_y0_mu = str(sim.model_parameters["y0"]["wolves"].values.tolist()).replace(" ", "") sim.config.model_parameters.wolves_y0.prior = f"lognorm(scale={wolves_y0_mu},s=0.5)" sim.config.model_parameters.rabbits_y0.prior = f"lognorm(scale={rabbits_y0_mu},s=0.5)" sim.set_inferer("numpyro") sim.prior_predictive_checks() sim.inferer.prior ``` /home/flo-schu/projects/pymob/pymob/sim/plot.py:155: UserWarning: There were 3 NaN or Inf values in the idata group 'prior_predictive'. See Simulation.inf_preds for a mask with the coordinates. warnings.warn( /home/flo-schu/projects/pymob/pymob/sim/plot.py:155: UserWarning: There were 3 NaN or Inf values in the idata group 'prior_predictive'. See Simulation.inf_preds for a mask with the coordinates. warnings.warn( ![png](hierarchical_model_varying_y0_files/hierarchical_model_varying_y0_17_1.png) ```python sim.inferer.run() ``` Trace Shapes: Param Sites: Sample Sites: alpha_species_mu_normal_base dist 2 | value 2 | alpha_species_sigma_normal_base dist | value | alpha_species_normal_base dist 3 2 | value 3 2 | alpha_sigma_normal_base dist | value | alpha_normal_base dist 120 | value 120 | beta_normal_base dist | value | wolves_y0_normal_base dist 120 | value 120 | rabbits_y0_normal_base dist 120 | value 120 | rabbits_obs dist 120 12 | value 120 12 | wolves_obs dist 120 12 | value 120 12 | 100%|██████████| 2000/2000 [02:47<00:00, 11.96it/s, init loss: 72264840585.0052, avg. loss [1901-2000]: 4630.1432] arviz - WARNING - Shape validation failed: input_shape: (1, 2000), minimum_shape: (chains=2, draws=4) mean sd hdi_3% hdi_97% mcse_mean mcse_sd ess_bulk \ alpha[0] 1.010 0.036 0.944 1.079 0.001 0.001 1924.0 alpha[1] 1.112 0.037 1.038 1.175 0.001 0.001 1887.0 alpha[2] 1.093 0.038 1.019 1.163 0.001 0.001 2012.0 alpha[3] 1.088 0.036 1.017 1.151 0.001 0.001 1524.0 alpha[4] 1.048 0.033 0.989 1.110 0.001 0.001 2010.0 ... ... ... ... ... ... ... ... wolves_y0[115] 14.685 1.258 12.317 16.990 0.029 0.021 1891.0 wolves_y0[116] 13.725 1.302 11.381 16.144 0.029 0.021 1984.0 wolves_y0[117] 11.830 0.881 10.236 13.474 0.020 0.014 1948.0 wolves_y0[118] 3.966 0.260 3.486 4.430 0.006 0.004 1842.0 wolves_y0[119] 7.461 0.675 6.157 8.634 0.016 0.011 1823.0 ess_tail r_hat alpha[0] 2003.0 NaN alpha[1] 1923.0 NaN alpha[2] 1850.0 NaN alpha[3] 1811.0 NaN alpha[4] 2004.0 NaN ... ... ... wolves_y0[115] 1865.0 NaN wolves_y0[116] 2035.0 NaN wolves_y0[117] 2003.0 NaN wolves_y0[118] 1774.0 NaN wolves_y0[119] 1954.0 NaN [371 rows x 9 columns] ![png](hierarchical_model_varying_y0_files/hierarchical_model_varying_y0_18_3.png) ```python fig, ax1 = plt.subplots(1, 1, figsize=(4,20)) az.plot_forest( data=[sim.inferer.idata.posterior], var_names=["beta", "alpha_species_mu", "alpha_species", "alpha"], ax=ax1, combined=True, hdi_prob=0.95, textsize=8 ) ax1.vlines(0.017648710084435453,*ax1.get_ylim(), color="black") ax1.vlines(1,*ax1.get_ylim(), color="black") ax1.vlines(3,*ax1.get_ylim(), color="black") ``` ![png](hierarchical_model_varying_y0_files/hierarchical_model_varying_y0_19_1.png) The relevant population parameters alpha[Cottontail], alpha[Jackrabbit] and beta[Wolves] are identified with good precision and uncertainty. Using the prior information for the starting values is a good idea. ```python sim.posterior_predictive_checks() ``` ![png](hierarchical_model_varying_y0_files/hierarchical_model_varying_y0_21_0.png) ```python sim.config.case_study.scenario = "lotka_volterra_hierarchical_final" sim.config.create_directory("scenario", force=True) sim.config.save(force=True) ``` Scenario directory created at '/home/flo-schu/projects/pymob/case_studies/lotka_volterra_case_study/scenarios/lotka_volterra_hierarchical_final'. TODO: + Deterministic distribution (identity?) to only broadcast alphas without random drawing a parameter Done: + something is not working with how y0 is handled. + check if priors are correctly specified. Look at the end of the notebook and save priors there? + calibrate without passing y0 and see if it works + Enable only partial y0 passing.