Ude implementation: UDESolver to evaluate UDE models + OptaxBackend to find optimal solution

.github/workflows/python-test.yml

@@ -125,6 +125,7 @@ jobs:
         python -m pip install flake8 pytest
         pip install .[pyabc,pymoo,interactive,numpyro]
         pip install -e case_studies/lotka_volterra_case_study
+        pip install -e case_studies/lotka_volterra_UDE_case_study
 
     - name: Lint with flake8
       if: env.full-test == 'true' && needs.decide-to-test.outputs.changes == 'true' && needs.decide-to-test.outputs.tagged_commit == 'false' && github.event_name == 'pull_request'

case_studies/lotka_volterra_UDE_case_study/.gitignore

@@ -0,0 +1,4 @@
+__pycache__
+results
+*.code-workspace
+*.egg-info

case_studies/lotka_volterra_UDE_case_study/.pre-commit-config.yaml

@@ -0,0 +1,15 @@
+repos:
+- repo: https://github.com/pre-commit/pre-commit-hooks
+  rev: v2.3.0
+  hooks:
+  - id: check-yaml
+  - id: check-toml
+
+- repo: local
+  hooks:
+  - id: pytest-check
+    name: pytest-check
+    entry: test.sh
+    language: script
+    pass_filenames: false
+    always_run: true

case_studies/lotka_volterra_UDE_case_study/__init__.py

@@ -0,0 +1 @@
+__version__ = "1.0.0"

case_studies/lotka_volterra_UDE_case_study/lotka_volterra_UDE_case_study/__init__.py

@@ -0,0 +1,7 @@
+from . import data
+from . import mod
+from . import plot
+from . import prob
+from . import sim
+
+__version__ = "1.0.0"

case_studies/lotka_volterra_UDE_case_study/lotka_volterra_UDE_case_study/mod.py

@@ -0,0 +1,91 @@
+import equinox as eqx
+import jax.nn as jnn
+import jax.numpy as jnp
+import jax
+from pymob.utils.UDE import UDEBase
+
+class Func(UDEBase):
+
+    mlp_depth: int = 3
+    mlp_width: int = 3
+    mlp_in_size: int = 2
+    mlp_out_size: int = 2
+
+    alpha: jax.Array
+    delta: jax.Array
+
+    def __init__(self, params, weights=None, bias=None, *, key, **kwargs):
+        self.init_MLP(weights, bias, key=key)
+        self.init_params(params)
+
+    def __call__(self, t, y):
+        """
+        Returns the growth rates of predator and prey depending on their current state.
+
+        Parameters
+        ----------
+        t : scalar
+            Just here to fulfill the requirements by diffeqsolve(). Has no effect and
+            can be set to None.
+        y : jax.ArrayImpl
+            Array containing two values: the current abundance of prey and predator,
+            respectively.
+
+        Returns:
+        --------
+        jax.ArrayImpl
+            An array containing the growth rates of prey and predators, respectively.
+        """
+
+        params = self.preprocess_params()
+
+        prey, predator = y
+
+        dprey_dt_ode = params["alpha"] * prey 
+        dpredator_dt_ode = - params["delta"] * predator
+        dprey_dt_nn, dpredator_dt_nn = self.mlp(y) * jnp.array([jnp.tanh(prey).astype(float), jnp.tanh(predator).astype(float)])
+
+        dprey_dt = dprey_dt_ode + dprey_dt_nn
+        dpredator_dt = dpredator_dt_ode + dpredator_dt_nn
+
+        return jnp.array([dprey_dt.astype(float),dpredator_dt.astype(float)])
+
+class Func1D(UDEBase):
+
+    mlp_depth: int = 3
+    mlp_width: int = 3
+    mlp_in_size: int = 1
+    mlp_out_size: int = 1
+
+    r: jax.Array
+
+    def __init__(self, params, weights=None, bias=None, *, key, **kwargs):
+        self.init_MLP(weights, bias, key=key)
+        self.init_params(params)
+
+    def __call__(self, t, y):
+        """
+        Returns the growth rates of predator and prey depending on their current state.
+
+        Parameters
+        ----------
+        t : scalar
+            Just here to fulfill the requirements by diffeqsolve(). Has no effect and
+            can be set to None.
+        y : jax.ArrayImpl
+            Array containing two values: the current abundance of prey and predator,
+            respectively.
+
+        Returns:
+        --------
+        jax.ArrayImpl
+            An array containing the growth rates of prey and predators, respectively.
+        """
+
+        params = self.preprocess_params()
+
+        X = y
+
+        dX_dt = params["r"] * X + self.mlp(y)
+
+        return jnp.array(dX_dt.astype(float))

case_studies/lotka_volterra_UDE_case_study/lotka_volterra_UDE_case_study/plot.py

@@ -0,0 +1 @@
+from lotka_volterra_case_study.plot import *

case_studies/lotka_volterra_UDE_case_study/lotka_volterra_UDE_case_study/prob.py

@@ -0,0 +1 @@
+from lotka_volterra_case_study.prob import *

case_studies/lotka_volterra_UDE_case_study/lotka_volterra_UDE_case_study/sim.py

@@ -0,0 +1 @@
+from lotka_volterra_case_study.sim import Simulation_v2

case_studies/lotka_volterra_UDE_case_study/pyproject.toml

@@ -0,0 +1,72 @@
+
+[build-system]
+requires = ["setuptools"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "lotka_volterra_UDE_case_study"
+version = "1.0.0"
+authors = [
+  { name="Florian Schunck", email="[email protected]" },
+]
+description = "Lotka Volterra Predator-Prey case study"
+readme = "README.md"
+requires-python = ">=3.10"
+dependencies=[
+    "pymob[numpyro] >= 0.5.0a19",
+    "preliz",
+]
+license = {file = "LICENSE"}
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Programming Language :: Python :: 3",
+    "Natural Language :: English",
+    "License :: OSI Approved :: GNU General Public License v3 (GPLv3)",
+    "Operating System :: OS Independent",
+    "Topic :: Scientific/Engineering :: Bio-Informatics",
+]
+
+[project.urls]
+"Homepage" = "https://github.com/flo-schu/lotka_volterra_case_study"
+"Issue Tracker" = "https://github.com/flo-schu/lotka_volterra_case_study/issues"
+
+[project.optional-dependencies]
+dev = [
+    "pytest >= 7.3", 
+    "bumpver", 
+    "pre-commit",
+    "ipykernel",
+    "ipywidgets"
+]
+
+[tool.setuptools.packages.find]
+include = ["lotka_volterra_UDE_case_study*"]
+
+[tool.bumpver]
+current_version = "1.0.0"
+version_pattern = "MAJOR.MINOR.PATCH[PYTAGNUM]"
+commit_message = "bump version {old_version} -> {new_version}"
+tag_message = "{new_version}"
+tag_scope = "default"
+pre_commit_hook = ""
+post_commit_hook = ""
+commit = true
+tag = true
+push = true
+
+[tool.bumpver.file_patterns]
+"pyproject.toml" = [
+    'current_version = "{version}"',
+    'version = "{version}"'
+]
+"lotka_volterra_case_study/__init__.py" = [
+    '__version__ = "{version}"'
+]
+"README.md" = [
+    'git clone [email protected]:flo-schu/lotka_volterra_case_study/{version}'
+]
+
+[tool.pytest.ini_options]
+markers = [
+    "slow='mark test as slow.'"
+]

case_studies/lotka_volterra_UDE_case_study/scenarios/InfererTest/settings.cfg

@@ -0,0 +1,133 @@
+[case-study]
+name = lotka_volterra_UDE_case_study
+pymob_version = 0.6.4
+scenario = InfererTest
+package = case_studies
+modules = sim mod prob data plot
+simulation = Simulation
+observations = UDE_obs_inferer_test.nc
+logging = DEBUG
+
+[simulation]
+y0 = 
+x_in = 
+input_files = 
+n_ode_states = 2
+batch_dimension = batch_id
+x_dimension = time
+modeltype = deterministic
+seed = 1
+
+[data-structure]
+prey = dimensions=['batch_id','time'] min=nan max=nan observed=True
+predator = dimensions=['batch_id','time'] min=0.11841753125190735 max=5.719013214111328 observed=True
+
+[solverbase]
+x_dim = time
+exclude_kwargs_model = t time x_in y x Y X
+exclude_kwargs_postprocessing = t time interpolation results
+
+[jax-solver]
+diffrax_solver = Dopri5
+rtol = 1e-06
+atol = 1e-07
+pcoeff = 0.0
+icoeff = 1.0
+dcoeff = 0.0
+max_steps = 100000
+throw_exception = True
+
+[inference]
+eps = 1e-08
+objective_function = total_average
+n_objectives = 1
+objective_names = 
+extra_vars = 
+n_predictions = 100
+
+[model-parameters]
+
+[error-model]
+
+[multiprocessing]
+cores = 1
+
+[inference.pyabc]
+sampler = SingleCoreSampler
+population_size = 100
+minimum_epsilon = 0.0
+min_eps_diff = 0.0
+max_nr_populations = 1000
+database_path = C:\Users\Markus\AppData\Local\Temp/pyabc.db
+
+[inference.pyabc.redis]
+password = nopassword
+port = 1111
+eval.n_predictions = 50
+eval.history_id = -1
+eval.model_id = 0
+
+[inference.pymoo]
+algortihm = UNSGA3
+population_size = 100
+max_nr_populations = 1000
+ftol = 1e-05
+xtol = 1e-07
+cvtol = 1e-07
+verbose = True
+
+[inference.numpyro]
+gaussian_base_distribution = False
+kernel = nuts
+init_strategy = init_to_uniform
+chains = 1
+draws = 2000
+warmup = 1000
+thinning = 1
+nuts_draws = 2000
+nuts_step_size = 0.8
+nuts_max_tree_depth = 10
+nuts_target_accept_prob = 0.8
+nuts_dense_mass = True
+nuts_adapt_step_size = True
+nuts_adapt_mass_matrix = True
+svi_iterations = 10000
+svi_learning_rate = 0.0001
+
+[inference.optax]
+UDE_parameters = alpha = value=1.3 dims=[] hyper=False free=False  delta = value=1.8 dims=[] prior=uniform(loc=1.0,scale=2.0) hyper=False free=True  
+MLP_weight_dist = normal()
+MLP_bias_dist = normal()
+length_strategy = 0.1 1
+steps_strategy = 1000 1000
+lr_strategy = 0.003 0.003
+clip_strategy = 0.1 0.1
+batch_size = 32
+data_split = 0.8
+multiple_runs_target = 3
+multiple_runs_limit = 5
+multiple_runs_plot = 5
+
+[report]
+debug_report = False
+pandoc_output_format = html
+model = True
+parameters = True
+parameters_format = pandas
+diagnostics = True
+diagnostics_with_batch_dim_vars = False
+diagnostics_exclude_vars = 
+goodness_of_fit = True
+goodness_of_fit_use_predictions = True
+goodness_of_fit_nrmse_mode = range
+table_parameter_estimates = True
+table_parameter_estimates_format = csv
+table_parameter_estimates_significant_figures = 3
+table_parameter_estimates_error_metric = sd
+table_parameter_estimates_parameters_as_rows = True
+table_parameter_estimates_with_batch_dim_vars = False
+table_parameter_estimates_exclude_vars = 
+table_parameter_estimates_override_names = 
+plot_trace = True
+plot_parameter_pairs = True
+