Disease spread programs

import pymc as pm


def model(data):
    """Infected count against time: the exponential-approach curve and prior of BoxingGym's
    disease-spread environment, with a Binomial likelihood over a population of 50.
    """
    with pm.Model() as result:
        t = pm.Data("t", data["Time"])
        y = pm.Data("y", data["Infected_Count"])
        theta = pm.TruncatedNormal("theta", mu=1, sigma=1, lower=0, upper=2)
        pm.Binomial("Infected_Count", n=50, p=1-pm.math.exp(-theta*t), observed=y)
    return result

import pymc as pm


def model(data):
    """Floor model: ignores time and fits one infection probability."""
    with pm.Model() as result:
        y = pm.Data("y", data["Infected_Count"])
        p = pm.Beta("p", alpha=1, beta=1)
        pm.Binomial("Infected_Count", n=50, p=p, observed=y)
    return result

import pymc as pm
import numpy as np

def model(data):
    """
    Construct a PyMC model for the death_process task.
    
    Parameters:
    data (dict): A dictionary containing 'Time' and 'Infected_Count' arrays.
    
    Returns:
    pm.Model: The PyMC model.
    """
    times = np.array(data['Time'])
    infected_counts = np.array(data['Infected_Count'])
    N = 50  # Total population size
    
    with pm.Model() as pmodel:
        # Define parameters for the model
        # We assume a linear relationship between time and the logit of the probability of infection
        # p(t) = sigmoid(alpha + beta * t)
        # Then the number of infected follows a Binomial(N, p(t))
        
        # Priors for the linear coefficients
        alpha = pm.Normal('alpha', mu=0, sigma=10)
        beta = pm.Normal('beta', mu=0, sigma=10)
        
        # Calculate the probability of infection at each time point
        logit_p = alpha + beta * times
        p = pm.math.sigmoid(logit_p)
        
        # Observed likelihood
        infected_observed = pm.Binomial('infected_observed', n=N, p=p, observed=infected_counts)
        
    return pmodel

import pymc as pm
import numpy as np

def model(data):
    """
    Defines a PyMC model for predicting infected counts in a population of 50.
    
    Args:
    data (dict): A dictionary containing 'Time' and 'Infected_Count' arrays.
    
    Returns:
    pm.Model: The defined PyMC model.
    """
    time = np.array(data['Time'])
    infected_counts = np.array(data['Infected_Count'])
    N = 50  # Total population size
    
    with pm.Model() as m:
        # Priors for the logistic regression coefficients
        # alpha: intercept, beta: slope
        alpha = pm.Normal('alpha', mu=0, sigma=5)
        beta = pm.Normal('beta', mu=0, sigma=5)
        
        # Compute the probability of infection at each time point
        # Using the logistic function: p = 1 / (1 + exp(-(alpha + beta * t)))
        logit_p = alpha + beta * time
        p = pm.math.invlogit(logit_p)
        
        # Observed likelihood: Binomial distribution
        # The number of infected individuals is binomially distributed with N trials and probability p
        infected_observed = pm.Binomial(
            'infected_observed',
            n=N,
            p=p,
            observed=infected_counts
        )
        
    return m

import pymc as pm
import numpy as np

def model(data):
    time = np.array(data['Time'])
    infected_counts = np.array(data['Infected_Count'])
    
    with pm.Model() as model:
        # Priors for the logistic curve parameters
        # beta: growth rate, should be positive
        beta = pm.Gamma('beta', alpha=2, beta=1)
        
        # t_mid: midpoint of the logistic curve, centered around 1.0 based on data
        t_mid = pm.Normal('t_mid', mu=1.0, sigma=0.5)
        
        # Total population is fixed at 50
        K = 50
        
        # Deterministic logistic function
        # I(t) = K / (1 + exp(-beta * (t - t_mid)))
        expected_infected = pm.Deterministic(
            'expected_infected',
            K / (1 + pm.math.exp(-beta * (time - t_mid)))
        )
        
        # Likelihood
        # Using Gaussian likelihood for simplicity, as counts are relatively high
        # and the Gaussian approximation to Binomial is reasonable.
        # Alternatively, Binomial likelihood: Binomial(n=50, p=expected_infected/50)
        # Let's use Binomial for better statistical grounding.
        p_infected = expected_infected / K
        # Clip p to [0, 1] to avoid numerical issues
        p_infected = pm.math.clip(p_infected, 0, 1)
        
        likelihood = pm.Binomial(
            'likelihood',
            n=50,
            p=p_infected,
            observed=infected_counts
        )
        
    return model

import numpy as np
import pymc as pm
import theano.tensor as tt


def model(data):
    """
    Build a PyMC model for the death_process / infection spread problem.
    
    The data consists of Time and Infected_Count observations from a population of 50.
    We need to model the number of infected individuals as a function of time.
    
    A common model for disease spread is the logistic growth model (SIR-like), where:
    I(t) = N / (1 + ((N - I0) / I0) * exp(-r * t))
    
    However, since we're observing noisy counts at specific times, we can use a simpler
    parametric model or a stochastic process model.
    
    Given the summary statistics:
    - Time ranges from ~0 to ~2
    - Infected_Count ranges from 1 to 47
    - Mean infected count is ~30
    
    Let's use a logistic growth model with parameters:
    - r: growth rate
    - I0: initial infected count
    
    The logistic function: I(t) = N / (1 + ((N - I0) / I0) * exp(-r * t))
    
    We'll place priors on r and I0, and use a likelihood that accounts for the discrete
    nature of the counts (e.g., Normal or Poisson approximation, or just Normal for simplicity
    since counts are relatively large).
    """
    
    N = 50  # Total population
    
    times = data['Time']
    infected_counts = data['Infected_Count']
    
    with pm.Model() as pm_model:
        # Priors on parameters
        # Growth rate r: positive, reasonable range
        r = pm.HalfNormal('r', sigma=5)
        
        # Initial infected count I0: between 1 and N
        # Using a Beta distribution scaled to [1, N]
        # Or directly use a discrete uniform or normal truncated
        I0 = pm.Beta('I0', alpha=2, beta=5, transform=pm.distributions.transforms.ordered)
        # Scale I0 to be between 1 and N-1
        I0_scaled = I0 * (N - 2) + 1
        
        # Deterministic: predicted infected count at each time point
        # Logistic growth: I(t) = N / (1 + ((N - I0) / I0) * exp(-r * t))
        # Handle edge cases where I0 is very small
        # Using a safer formulation
        
        # Let's use a different parameterization to avoid numerical issues
        # I(t) = N * sigmoid(r * (t - t0))
        # where t0 is the inflection point
        # sigmoid(x) = 1 / (1 + exp(-x))
        
        # Alternative: use r and K (carrying capacity, which is N)
        # I(t) = N / (1 + exp(-r * (t - t_mid)))
        # where t_mid is the time at which I(t) = N/2
        
        t_mid = pm.Normal('t_mid', mu=1.0, sigma=0.5)
        
        # Predicted infected count using logistic function
        I_pred = N / (1 + tt.exp(-r * (times - t_mid)))
        
        # Likelihood: observed counts are approximately normal around the predicted values
        # with some noise. Since counts are integers, we could use a Poisson or Negative Binomial,
        # but for simplicity and given the moderate counts, a Normal likelihood with a small
        # standard deviation works well.
        
        sigma = pm.HalfNormal('sigma', sigma=5)
        
        # Observed likelihood
        infected_obs = pm.Normal('infected_obs', mu=I_pred, sigma=sigma, observed=infected_counts)
    
    return pm_model

import numpy as np
import pymc as pm


def model(data):
    """
    Logistic growth model for disease spread.
    
    Assumptions:
    - Disease spread follows logistic growth
    - Total population is 50 (carrying capacity)
    - Observed counts follow a Binomial distribution
    - Priors are weakly informative to allow flexibility
    """
    # Extract data
    time = data['Time']
    infected = data['Infected_Count']
    N = 50  # Total population
    
    # Create PyMC model
    with pm.Model() as model:
        # Define feature container for time
        time_data = pm.Data('Time', time, dims='obs_id')
        
        # Priors for logistic growth parameters
        # r: growth rate (positive)
        # t0: inflection point (centered around 1, since that's where we see ~25 infected)
        r = pm.HalfNormal('r', sigma=5)
        t0 = pm.Normal('t0', mu=1.0, sigma=0.5)
        
        # Logistic function: I(t) = N / (1 + exp(-r*(t - t0)))
        # Compute the expected number of infected
        expected_infected = N / (1 + pm.math.exp(-r * (time_data - t0)))
        
        # Convert to probability for Binomial distribution
        p = expected_infected / N
        
        # Clip to avoid numerical issues
        p = pm.math.minimum(pm.math.maximum(p, 1e-10), 1 - 1e-10)
        
        # Observed likelihood - Binomial
        y_obs = pm.Binomial('y_obs', n=N, p=p, observed=infected, dims='obs_id')
    
    return model


def gen_model(observed_data):
    built_model = model({column: observed_data[column].to_numpy() for column in observed_data.columns})
    with built_model:
        trace = pm.sample(200, tune=200, target_accept=0.90, chains=1, cores=1, random_seed=42, idata_kwargs={"log_likelihood": True})
        posterior_predictive = pm.sample_posterior_predictive(trace, random_seed=314, return_inferencedata=False)
    return built_model, posterior_predictive, trace

import numpy as np
import pymc as pm


def model(data):
    time = data['Time']
    infected = data['Infected_Count']
    N = 50  # total population
    
    with pm.Model() as model:
        # Feature container for time
        time_data = pm.Data('Time', time, dims='obs_id')
        
        # Logistic growth parameters
        r = pm.HalfNormal('r', sigma=5)
        t0 = pm.Normal('t0', mu=1.0, sigma=0.5)
        
        # Deterministic logistic curve
        mu_infected = pm.Deterministic(
            'mu_infected',
            N / (1 + pm.math.exp(-r * (time_data - t0))),
            dims='obs_id'
        )
        
        # Overdispersion parameter
        sigma = pm.HalfNormal('sigma', sigma=5)
        
        # Negative Binomial likelihood
        y_obs = pm.NegativeBinomial(
            'y_obs',
            mu=mu_infected,
            sigma=sigma,
            observed=infected,
            dims='obs_id'
        )
    
    return model


def gen_model(observed_data):
    built_model = model({column: observed_data[column].to_numpy() for column in observed_data.columns})
    with built_model:
        trace = pm.sample(200, tune=200, target_accept=0.90, chains=1, cores=1, random_seed=42, idata_kwargs={"log_likelihood": True})
        posterior_predictive = pm.sample_posterior_predictive(trace, random_seed=314, return_inferencedata=False)
    return built_model, posterior_predictive, trace

import numpy as np
import pymc as pm
import pandas as pd

def model(data):
    """
    Logistic growth model for disease spread with Negative Binomial likelihood.
    
    Assumptions:
    - Disease spread follows logistic growth dynamics.
    - Total population is fixed at 50.
    - Observed counts are subject to overdispersion relative to the logistic mean.
    - Negative Binomial likelihood is used to model the count data, allowing for 
      extra variance beyond the Binomial constraint.
    """
    # Extract data
    time = data['Time']
    infected = data['Infected_Count']
    N = 50  # Total population
    
    # Create PyMC model
    with pm.Model() as model:
        # Define feature container for time
        time_data = pm.Data('Time', time, dims='obs_id')
        
        # Priors for logistic growth parameters
        # r: growth rate (positive). HalfNormal with sigma=5 allows for rapid growth.
        r = pm.HalfNormal('r', sigma=5)
        # t0: inflection point. Centered at 1.0 with reasonable uncertainty.
        t0 = pm.Normal('t0', mu=1.0, sigma=0.5)
        
        # Logistic function: I(t) = N / (1 + exp(-r*(t - t0)))
        # Compute the expected number of infected
        expected_infected = N / (1 + pm.math.exp(-r * (time_data - t0)))
        
        # Clip expected infected to ensure it stays within valid range for probability calculation
        # though NegativeBinomial takes mu directly, we might want to ensure positivity.
        # The logistic function naturally stays between 0 and N.
        
        # Overdispersion parameter for Negative Binomial
        # alpha is the dispersion parameter. Smaller alpha means more dispersion.
        # Using a HalfNormal prior for alpha.
        alpha = pm.HalfNormal('alpha', sigma=5)
        
        # Observed likelihood: Negative Binomial
        # pm.NegativeBinomial('y_obs', mu=expected_infected, alpha=alpha, observed=infected, dims='obs_id')
        # Note: In some PyMC versions, the parameterization might differ. 
        # Standard parameterization is mu (mean) and alpha (dispersion).
        # Variance = mu + mu^2 / alpha
        y_obs = pm.NegativeBinomial('y_obs', mu=expected_infected, alpha=alpha, observed=infected, dims='obs_id')
    
    return model


def gen_model(observed_data):
    built_model = model({column: observed_data[column].to_numpy() for column in observed_data.columns})
    with built_model:
        trace = pm.sample(200, tune=200, target_accept=0.90, chains=1, cores=1, random_seed=42, idata_kwargs={"log_likelihood": True})
        posterior_predictive = pm.sample_posterior_predictive(trace, random_seed=314, return_inferencedata=False)
    return built_model, posterior_predictive, trace

import numpy as np
import pymc as pm
import pandas as pd

def model(data):
    """
    Logistic growth model for disease spread with Negative Binomial likelihood.
    
    Assumptions:
    - Disease spread follows a logistic growth mean trend.
    - Total population is 50 (carrying capacity).
    - Observed counts exhibit overdispersion relative to Binomial, modeled via Negative Binomial.
    - Priors are weakly informative.
    """
    # Extract data
    time = data['Time']
    infected = data['Infected_Count']
    N = 50  # Total population
    
    # Create PyMC model
    with pm.Model() as model:
        # Define feature container for time
        time_data = pm.Data('Time', time, dims='obs_id')
        
        # Priors for logistic growth parameters
        # r: growth rate (positive, steep rise observed)
        r = pm.HalfNormal('r', sigma=10)
        # t0: inflection point (centered around 1.0, sigma=1.0 to be flexible)
        t0 = pm.Normal('t0', mu=1.0, sigma=1.0)
        
        # Logistic function: I(t) = N / (1 + exp(-r*(t - t0)))
        # Compute the expected number of infected
        expected_infected = N / (1 + pm.math.exp(-r * (time_data - t0)))
        
        # Clip expected infected to be within [1, N-1] to avoid log issues or zero/negative means
        # We want mean > 0 for NegativeBinomial. 
        # Also, NegativeBinomial parametrization in PyMC uses mu and alpha (or n, p).
        # Using mu and alpha parametrization: mu is mean, alpha is dispersion.
        # Var = mu + mu^2 / alpha
        
        # Clip mean to be strictly positive and <= N
        mu = pm.math.minimum(pm.math.maximum(expected_infected, 0.5), N - 0.5)
        
        # Dispersion parameter for Negative Binomial
        # High alpha means low dispersion (closer to Poisson/Binomial)
        # Low alpha means high dispersion
        alpha = pm.HalfNormal('alpha', sigma=5)
        
        # Observed likelihood - Negative Binomial
        # mu is the mean, alpha is the dispersion parameter
        y_obs = pm.NegativeBinomial('y_obs', mu=mu, alpha=alpha, observed=infected, dims='obs_id')
    
    return model

def gen_model(observed_data):
    built_model = model({column: observed_data[column].to_numpy() for column in observed_data.columns})
    with built_model:
        trace = pm.sample(200, tune=200, target_accept=0.90, chains=1, cores=1, random_seed=42, idata_kwargs={"log_likelihood": True})
        posterior_predictive = pm.sample_posterior_predictive(trace, random_seed=314, return_inferencedata=False)
    return built_model, posterior_predictive, trace