From 06c5d7947b463c28dd5c8f858550c284626a8d3b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Anna=20St=C3=B6riko?= Date: Thu, 23 Jul 2026 11:57:52 +0200 Subject: [PATCH] Add additional example notebooks to the documentation --- docs/_quarto.yml | 3 + docs/examples/index.qmd | 11 +- docs/examples/mixed-reactor.qmd | 101 +++++++++ docs/examples/pathogen-model.qmd | 310 ++++++++++++++++++++++++++ docs/examples/transport-model.qmd | 2 +- docs/examples/transport_with_pymc.qmd | 302 +++++++++++++++++++++++++ notebooks/pathogen-model.ipynb | 1 + 7 files changed, 724 insertions(+), 6 deletions(-) create mode 100644 docs/examples/mixed-reactor.qmd create mode 100644 docs/examples/pathogen-model.qmd create mode 100644 docs/examples/transport_with_pymc.qmd diff --git a/docs/_quarto.yml b/docs/_quarto.yml index 265385d..f97aa0d 100644 --- a/docs/_quarto.yml +++ b/docs/_quarto.yml @@ -49,6 +49,9 @@ website: contents: - examples/index.qmd - examples/transport-model.qmd + - examples/pathogen-model.qmd + - examples/transport_with_pymc.qmd + - examples/mixed-reactor.qmd - text: "---" - href: api/index.qmd diff --git a/docs/examples/index.qmd b/docs/examples/index.qmd index 45078f0..73eba91 100644 --- a/docs/examples/index.qmd +++ b/docs/examples/index.qmd @@ -1,9 +1,10 @@ --- -title: "Examples" +title: "Overview" --- -This section collects worked examples [TODO: finish this introduction with a brief description of each example]. +This section collects worked examples: -- [Reactive transport model](transport-model.qmd) - demonstrates how to set up and run a simple reactive transport model with Reactix. -- [TODO: Example 2] -- [TODO: Example 3] +- [Reactive transport model](transport-model.qmd) -- demonstrates how to set up and run a simple reactive transport model with Reactix. +- [Pathogen transport model](pathogen-model.qmd) -- a more complex example showing reactive transport of pathogens with attachment and detachment. +- [PyMC integration](transport_with_pymc.qmd) -- demonstrates integration of Reactix models with PyMC for Bayesian parameter estimation +- [Mixed-reactor model](mixed-reactor.qmd) -- show how to set up a model for a mixed system (without transport) diff --git a/docs/examples/mixed-reactor.qmd b/docs/examples/mixed-reactor.qmd new file mode 100644 index 0000000..5088a4c --- /dev/null +++ b/docs/examples/mixed-reactor.qmd @@ -0,0 +1,101 @@ +--- +title: Mixed reactor example +jupyter: python3 +--- + + +This notebook shows how to use Reactix to simulate an mixed reactor system. + +The system is described by the following ordinary differential equation: + +$$\frac{dc}{dt} = \frac{Q}{V}\left(c - c_{\mathrm{in}}\right) + r$$ + +where $c$ is the concentration, $Q$ is the flow rate of the reactor, $V$ its volume, $c_{\mathrm{in}}$ the inflowing concentration, and $r$ the reaction rates. Initially the concentration in the system is $c_0$. + +In this example, we will consider a first order decay reaction: + +$$r = - k c$$ + +However, more complicated reaction kinetics can easily be implemented with Reactix. + + +```{python} +import jax +import jax.numpy as jnp +import matplotlib.pyplot as plt +import numpy as np + +from reactix import MixedSystem, declare_species, make_solver, KineticReaction, reaction + +jax.config.update("jax_enable_x64", True) +``` + +```{python} +Species = declare_species(["tracer"]) +``` + +## Define the reactions + +```{python} +@reaction +class FirstOrderDecay(KineticReaction): + decay_coefficient: jax.Array + + def rate(self, time, state, system): + return self.decay_coefficient * state.tracer + + def stoichiometry(self, time, state, system): + return { + "tracer": -1, + } + +k = 0.02 +reactions = [FirstOrderDecay(decay_coefficient=k)] +``` + +## Set up the system + +```{python} +c_in = jnp.array(1.0) +Q = jnp.array(0.5) +V = jnp.array(10.0) +reactor_system = MixedSystem.build( + reactions=reactions, + discharge=Q, + volume=V, + inflow_concentration=Species(tracer=c_in), +) +``` + +## Run the solver + +```{python} +t_points = jnp.linspace(0, 200, 123) +solve = make_solver(t_max=200, t_points=t_points, rtol=1e-6, atol=1e-18) +``` + +```{python} +# Create an array of zeros with the length of the number of cells +val0 = jnp.zeros(()) + +# Set up the initial conditions (all zeros in this case) +state = Species(tracer=val0) +``` + +```{python} +solution = solve(state, reactor_system) +``` + +## Compare the solution against an analytical solution + +```{python} +c_0 = val0 +c_steady = Q / (Q + k * V) * c_in +c_analytical = c_steady + (c_0 - c_steady) * np.exp(-(Q / V + k) * t_points) + +plt.plot(t_points, solution.ys.tracer, label="Numerical Solution") +plt.plot(t_points, c_analytical, label="Analytical Solution", linestyle="--") +plt.legend() +``` + + diff --git a/docs/examples/pathogen-model.qmd b/docs/examples/pathogen-model.qmd new file mode 100644 index 0000000..7b1f844 --- /dev/null +++ b/docs/examples/pathogen-model.qmd @@ -0,0 +1,310 @@ +--- +title: Pathogen transport model +jupyter: python3 +--- + +```{python} +import jax +import jax.numpy as jnp +import matplotlib.pyplot as plt + +import reactix as rx + +jax.config.update("jax_enable_x64", True) +``` + + +This notebook demonstrates how to set up a reactive-transport model that simulates transport of a pathogen (modeled as a colloid) in porous media. The model couples advective-dispersive transport with attachment, detachment, and decay processes. + +## Model equations + +The transport of mobile pathogens and attached pathogens is governed by the following system of coupled partial differential equations, assuming constant coefficients. + +Mobile phase pathogen (aqueous concentration, mol/L of water): + +$$\frac{\partial C_m}{\partial t} + \frac{\rho_b}{\theta}\frac{\partial S}{\partial t} = -v\frac{\partial C_m}{\partial x} + D\frac{\partial^2 C_m}{\partial x^2} - k_a C_m + k_d \frac{S}{\rho_b/\theta} - \lambda_m C_m$$ + +Solid phase pathogen (attached concentration, mol/kg of solids): + +$$\frac{\partial S}{\partial t} = k_a \frac{\theta}{\rho_b} C_m - k_d S - \lambda_s S$$ + +Where: + +- $C_m$ = mobile pathogen concentration (mol/L of water) +- $S$ = attached pathogen concentration (mol/kg of solids) +- $\theta$ = porosity (fraction) +- $\rho_b$ = bulk density of porous medium (kg/m³) +- $v$ = pore water velocity (m/d) +- $D$ = hydrodynamic dispersion coefficient (m²/d) +- $k_a$ = attachment rate coefficient (1/d) +- $k_d$ = detachment rate coefficient (1/d) +- $\lambda_m$ = decay rate for mobile pathogens (1/d) +- $\lambda_s$ = decay rate for attached pathogens (1/d) + +## Simulated species + +The following code set up three species that will be used in the simulation – mobile and attached pathogens, and a conservative tracer for comparison. The second line indicates for each species if transport calculations are carried out – for the attached pathogens this is deactivated. + + +```{python} +Species = rx.declare_species(["tracer", "mobile_pathogen", "attached_pathogen"]) +species_is_mobile = Species(tracer=True, mobile_pathogen=True, attached_pathogen=False) +``` + +## System-wide parameters + +Converting between solid-phase and aqueous-phase concentrations in the reaction terms requires knowing the porosity and bulk density of the porous medium. +These are both system-wide properties rather than being specific to one reaction. +While the porosity is a native property of a `System` in Reactix (because it is also required for transport calculations), the bulk density is, by default, not stored in the `System` + +However, users can define a custom class that holds system-wide parameters and pass it to the `System`. Parameter `x` can be accessed as `system.parameters.x` where `system` is an `System` object. + +The custom parameter class needs to be defined using the `user_system_parameters` decorator. + +```{python} +@rx.user_system_parameters +class SystemParameters: + # Provide parameter names as attributes of the class + solid_density: jax.Array + + # Parameters can also be calculated from other system parameters + def bulk_density(self, system): + return (1 - system.porosity) * self.solid_density +``` + +```{python} +# Create a SystemParameters instance to store the parameter values +# Solid density typical for quartz sand: 2.65 g/cm³ +system_parameters = SystemParameters(solid_density=rx.SpatiallyConst(jnp.array(2.65))) # g/cm³ +``` + +## Reaction kinetics and stoichiometry + +Each reaction term is defined in a separate class, providing a custom rate expression and the stoichiometric coefficients. Note how we can use the same variable name for the decay parameter of the decay reaction of attached and mobile pathogens because they are contained by different classes. + +```{python} +@rx.reaction +class MobilePathogenDecay(rx.KineticReaction): + decay_coefficient: jax.Array + + def rate(self, time, state, system): + return self.decay_coefficient * state.mobile_pathogen + + def stoichiometry(self, time, state, system): + return { + "mobile_pathogen": -1, + } + +@rx.reaction +class AttachedPathogenDecay(rx.KineticReaction): + decay_coefficient: jax.Array + + def rate(self, time, state, system): + return self.decay_coefficient * state.attached_pathogen + + def stoichiometry(self, time, state, system): + return { + "attached_pathogen": -1, + } + +@rx.reaction +class Attachment(rx.KineticReaction): + attachment_coefficient: jax.Array + + def rate(self, time, state, system): + return self.attachment_coefficient * state.mobile_pathogen + + def stoichiometry(self, time, state, system): + bulk_density = system.parameters.bulk_density(system) + return { + "mobile_pathogen": -1, + "attached_pathogen": system.porosity / bulk_density + } + +@rx.reaction +class Detachment(rx.KineticReaction): + detachment_coefficient: jax.Array + + def rate(self, time, state, system): + return self.detachment_coefficient * state.attached_pathogen + + def stoichiometry(self, time, state, system): + bulk_density = system.parameters.bulk_density(system) + return { + "mobile_pathogen": bulk_density / system.porosity, + "attached_pathogen": -1, + } +``` + +After defining the form of the reaction rates in the reaction classes, specific parameter values are provided when creating an instance of each reaction class. + +```{python} +reactions = [ + Attachment(attachment_coefficient=1), # 1/d + Detachment(detachment_coefficient=1e-3), # 1/d + MobilePathogenDecay(decay_coefficient=5e-2), # 1/d + AttachedPathogenDecay(decay_coefficient=5e-2), # 1/d +] +``` + +## Discretization and transport parameters + +The simulation is set up on a uniform grid, and parameters and options for the dispersion and advection terms are provided. + +```{python} +n_cells = 200 +column_length = 10.0 # meters +interface_areas = jnp.ones(n_cells + 1) # Uniform cross-sectional area +cells = rx.Cells.equally_spaced(column_length, n_cells, interface_area=interface_areas) + + +dispersion = rx.Dispersion.build( + cells=cells, + dispersivity=jnp.array(0.05), # m + pore_diffusion=Species( + tracer=jnp.array(1e-9 * 3600 * 24), # m²/d + mobile_pathogen=jnp.array(1e-9 * 3600 * 24), + attached_pathogen=jnp.array(1e-9 * 3600 * 24), + ), +) +advection = rx.Advection.build( + limiter_type="minmod", +) +``` + +## Boundary conditions + +The model simulates a pulse injection experiment. A rectangular pulse of pathogen-bearing water is injected from the left boundary, followed by a flushing period with pathogen-free water. + +In total, we need to set up four boundary conditions – one per side for the tracer and the mobile pathogens. Since the attached pathogens are not transported, no boundary condition is needed for them. + +```{python} +t_stop_injection = 20 +bcs = [ + rx.FixedConcentrationBoundary( + boundary="left", + species_selector=lambda s: getattr(s, "tracer"), + fixed_concentration=lambda t: jnp.select([t < t_stop_injection], [jnp.array(10.0)], default=jnp.array(0.0)), + ), + rx.FixedConcentrationBoundary( + boundary="right", + species_selector=lambda s: getattr(s, "tracer"), + fixed_concentration=lambda t: jnp.array(0.0), + ), + rx.FixedConcentrationBoundary( + boundary="left", + species_selector=lambda s: getattr(s, "mobile_pathogen"), + fixed_concentration=lambda t: jnp.select([t < t_stop_injection], [jnp.array(10.0)], default=jnp.array(0.0)), + ), + rx.FixedConcentrationBoundary( + boundary="right", + species_selector=lambda s: getattr(s, "mobile_pathogen"), + fixed_concentration=lambda t: jnp.array(0.0), + ) +] +``` + +## System set-up + +All previously defined components of the model are now passed to a `System` object to set up the model. + +```{python} + +porosity= jnp.ones(n_cells) * 0.3 +system = rx.TransportSystem.build( + porosity=porosity, + # velocity=lambda t: jnp.array(1 / 365) * jnp.sin(np.pi * 2 * 1 / 5000 * t), + discharge=lambda t: jnp.array(1), + cells=cells, + advection=advection, + dispersion=dispersion, + species_is_mobile=species_is_mobile, + bcs=bcs, + reactions=reactions, + parameters=system_parameters +) +``` + +## Initial conditions + +Lastly, all concentrations are set to zero throughout the domain as the initial value. + +```{python} +zeros = jnp.zeros(cells.n_cells) + +initial_state = Species( + tracer=zeros, + mobile_pathogen=zeros, + attached_pathogen=zeros +) +``` + +## Provide solver options and run the simulation + +```{python} +t_points = jnp.linspace(0, 100, 300) +solver = rx.make_solver(t_max=100, t_points=t_points, rtol=1e-6, atol=1e-8) +solution = solver(initial_state, system) +#%timeit solution = solver(state, system) +``` + +## Simulation results + +This section presents the results of the reactive-transport simulation, showing how the three species (conservative tracer, mobile pathogen, and attached pathogen) evolve in space and time under the imposed pulse-injection boundary conditions. + +### Tracer concentration profiles at several times + +The first plot shows how the conservative tracer moves quickly through the entire domain. + +```{python} +plt.plot(cells.centers[:], solution.ys.tracer.T[:,:120:3]); +plt.xlabel("Distance (m)") +plt.ylabel("Tracer Concentration") +``` + +### Mobile pathogen concentration profiles + +In contrast to the conservative tracer, the mobile pathogens are subject to attachment and filtration on the porous medium surfaces. The plot below shows significant attenuation in the first several meters of the domain, where the high attachment coefficient leads to rapid removal from the mobile phase. + +```{python} +plt.plot(cells.centers[:], solution.ys.mobile_pathogen.T[:, 0:30]); +plt.xlabel("Distance (m)") +plt.ylabel("Mobile Pathogen Concentration") +#plt.yscale("log") +#plt.ylim(bottom=1e-3) +``` + +### Breakthrough curve of the mobile pathogens + +The breakthrough curve (BTC) plots the concentration of mobile pathogens at a fixed location (5 m into the column) as a function of time. +After the stop of the injection, it shows the typical tailing behaviour caused by detachment of pathogens from the solid phase. + +```{python} +plt.plot(solution.ts, solution.ys.mobile_pathogen.T[100,:]); +plt.xlabel("Time") +plt.ylabel("Mobile Pathogen Concentration") +plt.yscale("log") +plt.ylim(bottom=1e-4) +``` + +```{python} +%matplotlib widget +from matplotlib import animation, collections + +collections.Collection() +fig, ax = plt.subplots() + +artists = [] +for data in zip(solution.ys.tracer, solution.ys.mobile_pathogen, solution.ys.attached_pathogen): + containers = [ax.plot(cells.centers, y, color=f"C{i}") for i, y in enumerate(data)] + artist = [] + for container in containers: + artist.extend(container) + artists.append(artist) + + +ani = animation.ArtistAnimation(fig=fig, artists=artists, interval=40) +plt.show() +``` + + diff --git a/docs/examples/transport-model.qmd b/docs/examples/transport-model.qmd index e258b51..7fd9c11 100644 --- a/docs/examples/transport-model.qmd +++ b/docs/examples/transport-model.qmd @@ -1,5 +1,5 @@ --- -title: Conservative tracer and first-order decay +title: Conservative transport and first-order decay jupyter: python3 execute: enabled: true diff --git a/docs/examples/transport_with_pymc.qmd b/docs/examples/transport_with_pymc.qmd new file mode 100644 index 0000000..30ad5eb --- /dev/null +++ b/docs/examples/transport_with_pymc.qmd @@ -0,0 +1,302 @@ +--- +title: Parameter estimation with PyMC +jupyter: python3 +--- + + +This notebook demonstrates how Reactix can be coupled with PyMC for Bayesian inference of reaction parameter values. + +[PyMC](https://www.pymc.io/) is a library for Bayesian modeling and parameter inference. It can be used to define the prior distributions and likelihood, and offers a range of modern samplers, among them the No-U-turn sampler (NUTS). Using this sampler requires the computation of gradients of the posterior with respect to parameters. Since Reactix is built on JAX, this can be achieved through automatic differentiation. All this is done under the hood by PyMC, so the user only needs to define the model and run the sampler. + +This notebooks shows a toy example with a simple first-order decay reaction coupled to advective-dispersive transport. +Based on an artificial concentration measurement, we will estimate the decay constant. + +## Package imports + + +```{python} +# Ensure that that the ODE solver return NaN values when it fails instead of raising an error. +# This is important to prevent the NUTS sampler from crashing when the ODE solver fails to integrate the system of equations, +# but rather registering a divergence for the given parameter combination. +%env EQX_ON_ERROR=nan +``` + +```{python} +import pytensor +pytensor.config.exception_verbosity = "high" +``` + +```{python} +import jax +import jax.numpy as jnp +import numpy as np +import matplotlib.pyplot as plt +import dataclasses +import pandas as pd +import xarray as xr + +from reactix import ( + Advection, + Cells, + Dispersion, + FixedConcentrationBoundary, + TransportSystem, + make_solver, + declare_species, + KineticReaction, + reaction, +) +import pytensor.tensor as pt +import pymc as pm +import nutpie +``` + +```{python} +# Enable 64-bit floating-point precision in JAX +# See here: https://docs.kidger.site/diffrax/usage/how-to-choose-a-solver/#stiff-problems +jax.config.update("jax_enable_x64", True) +``` + +## Define the reactive-transport model + +In this example, we will use a simple reactive transport model with a reactive compound that is degraded with first-order decay kinetics. + +### Reaction kinetics + +We define a custom first-order decay reaction with Reactix: + +```{python} +@reaction +class FirstOrderDecay(KineticReaction): + decay_coefficient: jax.Array + + def rate(self, time, state, system): + return self.decay_coefficient * state.reactive_tracer + + def stoichiometry(self, time, state, system): + return { + "reactive_tracer": -1, + } + +Species = declare_species(["tracer", "reactive_tracer"]) +``` + +### Transport system + +We set up a system with constant flow velocity and uniform discretization. At the inflow boundary, a constant-concentration boundary is applied with a time-invariant inflow concentration. + +To keep the code in the PyMC model definition concise, we define a function for setting up the reactive-transport system that takes all the parameters that we want to estimate as inputs. In this case, we want to estimate the first-order decay constant. +Note that `make_system` takes an argument `decay_coefficient` which is then passed as a value to the `FirstOrderDecay` reaction object. + +```{python} +def make_system(decay_coefficient): + # Define the set of species and reactions + species_is_mobile = Species(tracer=True, reactive_tracer=True) + n_cells = 200 + reactions = [ + FirstOrderDecay( + decay_coefficient=decay_coefficient + ) + ] + # Define the cell geometry + interface_areas = jnp.ones(n_cells + 1) + cells = Cells.equally_spaced(10, n_cells, interface_area=interface_areas) + + # Set transport process options + dispersion = Dispersion.build( + cells=cells, + dispersivity=jnp.array(0.1), + pore_diffusion=Species( + tracer=jnp.array(1e-9 * 3600 * 24), + reactive_tracer=jnp.array(1e-9 * 3600 * 24), + ), + ) + advection = Advection.build( + limiter_type="upwind", + ) + # Define the boundary conditions + bcs = [ + FixedConcentrationBoundary( + boundary="left", + species_selector=lambda s: getattr(s, "tracer"), + fixed_concentration=lambda t: jnp.array(10.0), + ), + FixedConcentrationBoundary( + boundary="right", + species_selector=lambda s: getattr(s, "tracer"), + fixed_concentration=lambda t: jnp.array(3.0), + ), + FixedConcentrationBoundary( + boundary="left", + species_selector=lambda s: getattr(s, "reactive_tracer"), + fixed_concentration=lambda t: jnp.array(1.0), + ), + FixedConcentrationBoundary( + boundary="right", + species_selector=lambda s: getattr(s, "reactive_tracer"), + fixed_concentration=lambda t: jnp.array(3.0), + ) + ] + + # Set up the system + porosity= jnp.ones(n_cells) * 0.3 + #porosity = porosity.at[100:].set(0.1) + return TransportSystem.build( + porosity=porosity, + # velocity=lambda t: jnp.array(1 / 365) * jnp.sin(np.pi * 2 * 1 / 5000 * t), + discharge=lambda t: jnp.array(1 / 365) * 0.3, + cells=cells, + advection=advection, + dispersion=dispersion, + species_is_mobile=species_is_mobile, + bcs=bcs, + reactions=reactions + ) +``` + +## Define the PyMC model + +### Measurement point + +Below, we create an concentration data point that will be used for parameter estimation. +We set a concentration of 0.5 at the 40th cell at the 20th time point. + +```{python} +idx_time = np.array([20,]) +idx_space = np.array([40,]) + +measurement = xr.DataArray( + data=[0.5], + name="concentration_measurement", + dims="measurement_nr", + coords={"measurement_nr": pd.RangeIndex(len(idx_time))} +) +``` + +### Passing coordinates and defining prior distributions + +PyMC return an `InferenceData` object with named arrays that have dimensions and coordinates (based on Xarray). +In order to save the time as a coordinate for the simulated concentrations, we need to create a dictionary with all the coodinates. + +```{python} +t_points = jnp.linspace(0, 8000, 123) + +coords = { + "time_dense": np.array(t_points), + **measurement.coords +} +``` + +We can then pass these coordinates to the PyMC model. The first step in the model is to define prior distributions for all parameters that we would like to estimate. + +Here, we chose a half-normal distribution with a standard deviation of 0.01 for the decay coefficient and pass it to the function that builds the reactive-transport system. + +Finally, we also add the cell center locations as a coordinate to the PyMC `model`. This can only be done after creating the `system`, when the cell geometry is known. + +```{python} + +with pm.Model(coords=coords) as model: + decay_coefficient = pm.HalfNormal("k_dec", 0.01) + system = make_system(decay_coefficient) + model.add_coord("x", np.array(system.cells.centers)) + +``` + +### Setting up the solver + +PyMC is built around the PyTensor library, whereas Reactix uses JAX. In order to be able to use JAX together with PyMC, the JAX code needs to be wrapped with the `pytensor.wrap_jax` decorator. + +```{python} + +with model: + # Initial conditions: all species start at zero concentration + zeros = jnp.zeros(system.cells.n_cells) + initial_state = Species( + tracer=zeros, + reactive_tracer=zeros + ) + + + # Wrap the solver function with PyTensor to make it compatible with PyMC + @pytensor.wrap_jax + def solve_pt(y0, system): + solve_fn_dense = make_solver(t_max=8000, t_points=t_points, rtol=1e-6, atol=1e-6) + return solve_fn_dense(y0, system).ys + + # Solve the system of ODEs using the initial state and the system definition + solution = solve_pt(initial_state, system) +``` + +### Saving outputs and defining the likelihood + +```{python} + +with model: + # Save all concentrations as deterministic variables in the PyMC model for later analysis + fields = dataclasses.fields(solution) + for field in fields: + pm.Deterministic( + field.name, + getattr(solution, field.name), + dims=("time_dense", "x"), + ) + + # Define a normal likelihood for the observations, using the model's predicted concentration at the specified time and space indices as the mean + pm.Normal("made_up_point", mu=solution.reactive_tracer[idx_time, idx_space], sigma=0.1, observed=measurement.values, dims=("measurement_nr")) +``` + +## Sample the prior distribution + +```{python} +with model: + prior = pm.sample_prior_predictive(draws=50) +``` + +### Prior predictive checks + +The following plots show several draws from the prior of the simulated concentrations. This is a good check to assess if the chosen prior distributions for the parameters are reasonable. The measured concentrations is plotted alongside as a black dot. + +#### Concentration profile at the time of measurement + +```{python} +prior.prior.reactive_tracer.isel(chain=0, time_dense=idx_time[0]).plot.line(x="x", hue="draw", add_legend=False); +plt.scatter(system.cells.centers[idx_space], measurement.values, color="k", label="measurement") +``` + +#### Breakthrough curve at the measurement location + +```{python} +prior.prior.reactive_tracer.isel(chain=0, x=idx_space[0]).plot.line(x="time_dense", hue="draw", add_legend=False); +plt.scatter(t_points[idx_time], prior.observed_data.made_up_point, color="k", label="measurement") +``` + +## Sample the posterior distribution + +Finally, we can sample the posterior distribution. We can use a state-of-the-art HMC sampler since gradient calculations are provided through automatic differentiation by JAX under the hood. Here we use the Nutpie sampler for PyMC. Note that you need to specify that JAX should be used as a backend for the posterior and gradient evaluation. + +Note that sampling can take a while since thousands of model runs are needed. In order to be able to see intermediate results while the sampler is still running, we let the sampler run in the background by using nutpie's option `blocking=False`. + +```{python} +compiled = nutpie.compile_pymc_model(model, backend="jax", gradient_backend="jax") +sampler = nutpie.sample(compiled, chains=2, blocking=False) + +``` + +After a while, we can use `sampler.inspect()` to obtain the intermediate results of the sampler and inspect them. + +```{python} +# Inspect intermediate results of the sampler +trace = sampler.inspect() +``` + +```{python} +# Plot the evolution of the decay coefficient during the warmup phase of the sampler +trace.warmup_posterior.k_dec.plot.line(x="draw") +``` + +```{python} +# Plot concentration profiles in the warmup posterior +trace.warmup_posterior.reactive_tracer.isel(chain=0, time_dense=20, draw=slice(20, None)).plot.line(x="x", hue="draw", add_legend=False); +``` + + diff --git a/notebooks/pathogen-model.ipynb b/notebooks/pathogen-model.ipynb index 927b948..2044de0 100644 --- a/notebooks/pathogen-model.ipynb +++ b/notebooks/pathogen-model.ipynb @@ -38,6 +38,7 @@ "$$\\frac{\\partial S}{\\partial t} = k_a \\frac{\\theta}{\\rho_b} C_m - k_d S - \\lambda_s S$$\n", "\n", "Where:\n", + "\n", "- $C_m$ = mobile pathogen concentration (mol/L of water)\n", "- $S$ = attached pathogen concentration (mol/kg of solids)\n", "- $\\theta$ = porosity (fraction)\n",