Pathogen transport model

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.

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.

@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

# 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.

@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.

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.

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.

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.

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.

zeros = jnp.zeros(cells.n_cells)

initial_state = Species(
    tracer=zeros,
    mobile_pathogen=zeros,
    attached_pathogen=zeros
)

Provide solver options and run the simulation

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.

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.

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.

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)

%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()

--- 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() ```