tessl/pypi-pymc

Probabilistic Programming in Python: Bayesian Modeling and Probabilistic Machine Learning with PyTensor

—

Pending

Overview

Eval results

Files

PyMC Ordinary Differential Equations (ODE)

Name: tessl/pypi-pymc
Author: tessl

PyMC provides comprehensive support for probabilistic modeling with ordinary differential equations. This enables Bayesian parameter estimation for dynamic systems, time series modeling with mechanistic models, and uncertainty quantification in scientific models described by ODEs.

Core ODE Interface

The ODE module integrates differential equation solvers with PyMC's probabilistic programming framework:

from pymc import ode
import numpy as np

# Main ODE class
class ode.DifferentialEquation:
    """
    Solve ordinary differential equations as part of PyMC models.
    
    Parameters:
    - func (callable): Function defining the ODE system dy/dt = func(y, t, theta)
    - times (array): Time points at which to solve the ODE
    - n_states (int): Number of state variables in the ODE system
    - n_theta (int): Number of parameters in the ODE system
    - t0 (float): Initial time point
    """
    
    def __init__(self, func, times, n_states, n_theta, t0=0):
        pass
    
    def __call__(self, theta, y0):
        """
        Solve the ODE system.
        
        Parameters:
        - theta: Parameter tensor for the ODE system
        - y0: Initial conditions tensor
        
        Returns:
        - Solution tensor with shape (n_times, n_states)
        """
        pass

Capabilities

Basic ODE Integration

Solve systems of ordinary differential equations within probabilistic models:

import pymc as pm
from pymc import ode
import numpy as np

# Define ODE system
def simple_exponential(y, t, theta):
    """Simple exponential growth: dy/dt = theta[0] * y"""
    return theta[0] * y

# Solve in model context
with pm.Model() as exponential_model:
    # Parameter prior
    growth_rate = pm.Normal('growth_rate', mu=0.1, sigma=0.05)
    
    # Initial condition
    y0 = pm.Normal('y0', mu=1.0, sigma=0.1)
    
    # Define ODE
    ode_solution = ode.DifferentialEquation(
        func=simple_exponential,
        times=np.linspace(0, 10, 50),
        n_states=1,
        n_theta=1,
        t0=0
    )
    
    # Solve ODE
    solution = ode_solution([growth_rate], [y0])
    
    # Likelihood for observed data
    sigma = pm.HalfNormal('sigma', sigma=0.1)
    obs = pm.Normal('obs', mu=solution, sigma=sigma, observed=observed_data)

Multi-Dimensional ODE Systems

Handle complex dynamical systems with multiple interacting components:

def lotka_volterra(y, t, theta):
    """
    Lotka-Volterra predator-prey model.
    
    y[0]: prey population
    y[1]: predator population
    theta[0]: prey growth rate (alpha)
    theta[1]: predation rate (beta) 
    theta[2]: predator efficiency (gamma)
    theta[3]: predator death rate (delta)
    """
    prey, predator = y
    alpha, beta, gamma, delta = theta
    
    dprey_dt = alpha * prey - beta * prey * predator
    dpredator_dt = gamma * prey * predator - delta * predator
    
    return [dprey_dt, dpredator_dt]

# Model with multiple state variables
with pm.Model() as lotka_volterra_model:
    # Population dynamics parameters
    alpha = pm.LogNormal('alpha', mu=np.log(1.0), sigma=0.2)  # Prey growth
    beta = pm.LogNormal('beta', mu=np.log(0.5), sigma=0.2)    # Predation
    gamma = pm.LogNormal('gamma', mu=np.log(0.3), sigma=0.2)  # Predator efficiency
    delta = pm.LogNormal('delta', mu=np.log(0.8), sigma=0.2)  # Predator death
    
    # Initial populations
    prey_0 = pm.LogNormal('prey_0', mu=np.log(10), sigma=0.1)
    predator_0 = pm.LogNormal('predator_0', mu=np.log(5), sigma=0.1)
    
    # Solve the ODE system
    ode_solution = ode.DifferentialEquation(
        func=lotka_volterra,
        times=observation_times,
        n_states=2,
        n_theta=4,
        t0=0
    )
    
    solution = ode_solution([alpha, beta, gamma, delta], [prey_0, predator_0])
    
    # Separate observed populations
    prey_solution = solution[:, 0]
    predator_solution = solution[:, 1]
    
    # Observation models
    prey_sigma = pm.HalfNormal('prey_sigma', sigma=1.0)
    predator_sigma = pm.HalfNormal('predator_sigma', sigma=1.0)
    
    prey_obs = pm.Normal('prey_obs', mu=prey_solution, sigma=prey_sigma, 
                        observed=prey_data)
    predator_obs = pm.Normal('predator_obs', mu=predator_solution, 
                           sigma=predator_sigma, observed=predator_data)

Pharmacokinetic/Pharmacodynamic Models

Model drug concentration and effect over time:

def pk_two_compartment(y, t, theta):
    """
    Two-compartment pharmacokinetic model.
    
    y[0]: central compartment concentration
    y[1]: peripheral compartment concentration
    theta[0]: elimination rate (ke)
    theta[1]: distribution rate central->peripheral (k12)
    theta[2]: distribution rate peripheral->central (k21)
    """
    central, peripheral = y
    ke, k12, k21 = theta
    
    dcentral_dt = -ke * central - k12 * central + k21 * peripheral
    dperipheral_dt = k12 * central - k21 * peripheral
    
    return [dcentral_dt, dperipheral_dt]

# Pharmacokinetic model with dosing
with pm.Model() as pk_model:
    # PK parameters
    ke = pm.LogNormal('ke', mu=np.log(0.1), sigma=0.3)    # Elimination rate
    k12 = pm.LogNormal('k12', mu=np.log(0.05), sigma=0.3) # Central to peripheral
    k21 = pm.LogNormal('k21', mu=np.log(0.03), sigma=0.3) # Peripheral to central
    
    # Initial concentrations (after IV dose)
    dose = 100  # mg
    volume_central = 10  # L
    initial_central = dose / volume_central
    initial_peripheral = 0.0
    
    # Solve ODE
    ode_solution = ode.DifferentialEquation(
        func=pk_two_compartment,
        times=sample_times,
        n_states=2,
        n_theta=3,
        t0=0
    )
    
    concentrations = ode_solution([ke, k12, k21], 
                                 [initial_central, initial_peripheral])
    
    # Observable is central compartment concentration
    central_conc = concentrations[:, 0]
    
    # Proportional error model
    sigma_prop = pm.HalfNormal('sigma_prop', sigma=0.1)
    observed_conc = pm.Normal('observed_conc', 
                             mu=central_conc, 
                             sigma=sigma_prop * central_conc,
                             observed=concentration_data)

Epidemiological Models

Model disease spread using compartmental models:

def sir_model(y, t, theta):
    """
    SIR epidemiological model.
    
    y[0]: susceptible population (S)
    y[1]: infectious population (I) 
    y[2]: recovered population (R)
    theta[0]: transmission rate (beta)
    theta[1]: recovery rate (gamma)
    """
    S, I, R = y
    beta, gamma = theta
    N = S + I + R  # Total population
    
    dS_dt = -beta * S * I / N
    dI_dt = beta * S * I / N - gamma * I
    dR_dt = gamma * I
    
    return [dS_dt, dI_dt, dR_dt]

# Epidemiological model  
with pm.Model() as sir_epidemic:
    # Disease parameters
    beta = pm.LogNormal('beta', mu=np.log(0.3), sigma=0.2)   # Transmission rate
    gamma = pm.LogNormal('gamma', mu=np.log(0.1), sigma=0.2) # Recovery rate
    
    # Initial conditions
    N = 1000000  # Total population
    I0 = 10      # Initial infected
    S0 = N - I0  # Initial susceptible
    R0 = 0       # Initial recovered
    
    # Solve epidemic dynamics
    ode_solution = ode.DifferentialEquation(
        func=sir_model,
        times=np.arange(0, 365),  # One year daily
        n_states=3,
        n_theta=2,
        t0=0
    )
    
    solution = ode_solution([beta, gamma], [S0, I0, R0])
    
    # Extract infectious population over time
    I_t = solution[:, 1]
    
    # Observation model for reported cases
    reporting_rate = pm.Beta('reporting_rate', alpha=2, beta=8)
    expected_reports = I_t * reporting_rate
    
    # Negative binomial for overdispersed count data
    alpha = pm.HalfNormal('alpha', sigma=10)
    reported_cases = pm.NegativeBinomial('reported_cases',
                                       mu=expected_reports,
                                       alpha=alpha,
                                       observed=case_data)

Advanced Features

Time-Varying Parameters

Handle parameters that change over time:

def time_varying_ode(y, t, theta_func):
    """ODE with time-varying parameters."""
    theta_t = theta_func(t)  # Parameters as function of time
    return theta_t[0] * y - theta_t[1] * y**2

# Implementation with interpolated parameters
with pm.Model() as time_varying_model:
    # Time-varying growth rate
    growth_nodes = pm.Normal('growth_nodes', mu=0.1, sigma=0.05, shape=5)
    
    # Create interpolation function (conceptual - actual implementation varies)
    def theta_interpolated(t):
        # Interpolate growth_nodes over time
        return [interpolate_growth(t, growth_nodes), 0.01]
    
    # Use in ODE solving...

Stochastic Differential Equations

While PyMC's ODE module focuses on deterministic ODEs, stochastic elements can be incorporated through observation models and parameter uncertainty.

Integration with Experimental Design

Use ODE models to optimize experimental design:

# Design optimal sampling times for parameter estimation
with pm.Model() as design_model:
    # Parameters
    theta = pm.LogNormal('theta', mu=0, sigma=1, shape=n_params)
    
    # Candidate sampling times
    sample_times_candidates = pm.Uniform('sample_times', 
                                       lower=0, upper=T_max, 
                                       shape=n_samples)
    
    # Solve ODE at candidate times
    ode_solution = ode.DifferentialEquation(
        func=model_equations,
        times=sample_times_candidates,
        n_states=n_states,
        n_theta=n_params,
        t0=0
    )
    
    predictions = ode_solution(theta, y0)
    
    # Expected information criterion for design optimization
    # (Implementation depends on specific design criterion)

Performance Considerations

Solver Selection

PyMC's ODE module uses efficient numerical solvers. For optimal performance:

# Install sunode for high-performance solving
# pip install sunode

# Configure solver settings (conceptual)
ode_solution = ode.DifferentialEquation(
    func=model_equations,
    times=times,
    n_states=n_states,
    n_theta=n_theta,
    solver='dopri5',    # Runge-Kutta solver
    rtol=1e-6,         # Relative tolerance
    atol=1e-8          # Absolute tolerance
)

Computational Efficiency

Use compiled functions when possible
Minimize the number of time points when not all are needed for observations
Consider gradient-based sampling methods (NUTS) for efficient exploration
Profile your models to identify computational bottlenecks

Common Patterns

Parameter Transformations

# Log-transform positive parameters for unconstrained sampling
log_params = pm.Normal('log_params', mu=0, sigma=1, shape=n_params)
params = pm.math.exp(log_params)

# Use transformed parameters in ODE
solution = ode_solution(params, y0)

Missing Data and Irregular Observations

# Handle missing observations
mask = ~np.isnan(observed_data)
solution_observed = solution[mask]
data_observed = observed_data[mask]

obs = pm.Normal('obs', mu=solution_observed, sigma=sigma, observed=data_observed)

Model Comparison

# Compare different ODE structures
with pm.Model() as model_1:
    # Simple exponential growth
    pass

with pm.Model() as model_2:
    # Logistic growth with carrying capacity
    pass

# Use information criteria for model selection
loo_1 = pm.loo(trace_1, model_1)
loo_2 = pm.loo(trace_2, model_2)
pm.compare({'exponential': loo_1, 'logistic': loo_2})

PyMC's ODE integration provides a powerful framework for mechanistic modeling in a Bayesian context, enabling principled uncertainty quantification for dynamic systems across scientific domains.

Install with Tessl CLI