Model Building and Core Constructs

class Model:
    """
    Main model container class for PyMC3 Bayesian models.
    
    The Model class serves as a context manager that collects random variables,
    deterministic variables, and potential functions to define a complete
    Bayesian model for inference.
    """
    
    def __init__(self, name='', model=None):
        """
        Create a new Model.
        
        Parameters:
        - name: str, model name for identification
        - model: Model, parent model for nested models
        """
    
    def Var(self, name, dist, data=None, total_size=None, dims=None):
        """
        Create a random variable within the model.
        
        Parameters:
        - name: str, variable name
        - dist: Distribution, probability distribution
        - data: array-like, observed data (makes variable observed)
        - total_size: int, total size for minibatch inference
        - dims: tuple, dimension names for ArviZ
        
        Returns:
        - TensorVariable: model variable
        """
    
    def set_data(self, new_data, model=None):
        """
        Update shared data variables in the model.
        
        Parameters:
        - new_data: dict, mapping variable names to new data
        - model: Model, model context (uses current if None)
        """
    
    @property
    def logp(self):
        """Combined log-probability of all random variables."""
    
    @property
    def logpt(self):
        """Theano tensor for log-probability computation."""
    
    @property  
    def free_RVs(self):
        """List of free (unobserved) random variables."""
    
    @property
    def observed_RVs(self):
        """List of observed random variables."""
    
    @property
    def deterministics(self):
        """List of deterministic variables."""
    
    @property
    def potentials(self):
        """List of potential functions."""
    
    def compile_logp(self, vars=None, jacobian=True):
        """
        Compile log-probability function.
        
        Parameters:
        - vars: list, variables to include (all free RVs if None)
        - jacobian: bool, include Jacobian of transformations
        
        Returns:
        - function: compiled log-probability function
        """
    
    def compile_dlogp(self, vars=None, jacobian=True):
        """
        Compile gradient of log-probability function.
        
        Parameters:
        - vars: list, variables to include (all free RVs if None)  
        - jacobian: bool, include Jacobian of transformations
        
        Returns:
        - function: compiled gradient function
        """
    
    def check_test_point(self, test_point=None):
        """
        Check that model can be evaluated at test point.
        
        Parameters:
        - test_point: dict, point to test (uses model.initial_point if None)
        
        Returns:
        - dict: evaluation results and diagnostics
        """

def modelcontext(model=None):
    """
    Get current model context.
    
    Parameters:
    - model: Model, specific model to return (current context if None)
    
    Returns:
    - Model: current or specified model
    
    Raises:
    - TypeError: if no model is found and none provided
    """

def set_data(new_data, model=None):
    """
    Set new values for shared data variables.
    
    Parameters:
    - new_data: dict, mapping from variable names to new data arrays
    - model: Model, model containing the variables (current context if None)
    
    Example:
    with pm.Model() as model:
        x_data = pm.Data('x_data', np.array([1, 2, 3]))
        y = pm.Normal('y', mu=x_data, sigma=1)
        
    # Later update the data
    pm.set_data({'x_data': np.array([4, 5, 6])}, model=model)
    """

def Point(*args, **kwargs):
    """
    Create a point dictionary for model evaluation.
    
    Parameters:
    - args: values in order of model.free_RVs
    - kwargs: mapping from variable names to values
    
    Returns:
    - dict: point dictionary mapping variable names to values
    """

class FreeRV:
    """
    Free (unobserved) random variable with prior distribution.
    
    Represents parameters or latent variables that will be inferred
    from data through MCMC or variational inference.
    """
    
    @property
    def distribution(self):
        """Associated probability distribution."""
    
    @property
    def transformed(self):
        """Transformed version for unconstrained sampling."""
    
    @property
    def tag(self):
        """Variable metadata and properties."""

class ObservedRV:
    """
    Observed random variable representing data or likelihood.
    
    Links data to the model through a probability distribution,
    defining the likelihood component of Bayesian inference.
    """
    
    @property
    def distribution(self):
        """Associated probability distribution."""
    
    @property
    def observations(self):
        """Observed data values."""
    
    @property
    def missing_values(self):
        """Locations of missing data (if any)."""

class MultiObservedRV:
    """
    Container for multiple observed random variables.
    
    Used when multiple related observations share parameters
    or when vectorizing likelihood computations.
    """

def Deterministic(name, var, model=None, dims=None):
    """
    Create a deterministic variable from a Theano expression.
    
    Deterministic variables are functions of other model variables
    and are tracked for posterior analysis without being sampled directly.
    
    Parameters:
    - name: str, variable name
    - var: TensorVariable, Theano expression defining the transformation
    - model: Model, model to add variable to (current context if None)
    - dims: tuple, dimension names for ArviZ integration
    
    Returns:
    - TensorVariable: deterministic variable
    
    Example:
    with pm.Model() as model:
        mu = pm.Normal('mu', 0, 1)
        sigma = pm.HalfNormal('sigma', 1)
        
        # Deterministic transformation
        precision = pm.Deterministic('precision', 1 / sigma**2)
        
        y = pm.Normal('y', mu=mu, tau=precision, observed=data)
    """

def Potential(name, var, model=None):
    """
    Add a potential (log-likelihood) term to the model.
    
    Potentials allow adding arbitrary log-probability terms that are not
    associated with specific random variables, useful for custom likelihoods,
    constraints, or incorporating external information.
    
    Parameters:
    - name: str, potential name for identification
    - var: TensorVariable, log-probability expression to add to model
    - model: Model, model to add potential to (current context if None)
    
    Returns:
    - TensorVariable: the potential term
    
    Example:
    with pm.Model() as model:
        theta = pm.Uniform('theta', 0, 1, shape=3)
        
        # Add constraint that parameters sum to 1
        constraint = pm.Potential('constraint', 
                                 tt.switch(tt.abs_(theta.sum() - 1) < 1e-6, 
                                          0, 
                                          -np.inf))
    """

class Factor:
    """
    Base class for model factors (components contributing to log-probability).
    
    Factors include random variables, deterministic variables, and potentials.
    They provide a unified interface for model introspection and manipulation.
    """
    
    @property
    def logp(self):
        """Log-probability contribution of this factor."""
    
    @property  
    def logpt(self):
        """Theano tensor for log-probability computation."""

class Data:
    """
    Shared data container for observations that can be updated.
    
    Data objects allow updating observed values without recompiling
    the model, enabling out-of-sample prediction and minibatch inference.
    """
    
    def __init__(self, name, value, dims=None, export_index_as_dims=False):
        """
        Create a shared data variable.
        
        Parameters:
        - name: str, variable name
        - value: array-like, initial data values
        - dims: tuple, dimension names for ArviZ
        - export_index_as_dims: bool, use integer indices as dimension names
        
        Returns:
        - TensorSharedVariable: shared data tensor
        """
    
    def set_value(self, new_value):
        """Update data values."""

class Minibatch:
    """
    Minibatch container for stochastic variational inference.
    
    Automatically handles data subsampling and scaling for minibatch
    training with variational inference methods.
    """
    
    def __init__(self, data, batch_size=None, dtype=None, broadcastable=None, 
                 name=None, random_seed=None, update_shared_f=None, 
                 in_memory_slices=None):
        """
        Create minibatch data container.
        
        Parameters:
        - data: array-like, full dataset
        - batch_size: int, size of minibatches
        - dtype: data type for batches
        - broadcastable: tuple, broadcastable dimensions
        - name: str, variable name
        - random_seed: int, random seed for sampling
        - update_shared_f: function, custom update function
        - in_memory_slices: tuple, slices to keep in memory
        """

class GeneratorAdapter:
    """
    Adapter for using Python generators as data sources.
    
    Enables streaming data processing and infinite data generators
    for online learning scenarios.
    """

def get_data(filename):
    """
    Load data from file (supports various formats).
    
    Parameters:
    - filename: str, path to data file
    
    Returns:
    - array: loaded data
    """

def align_minibatches(*minibatch_tensors):
    """
    Align multiple minibatch tensors for synchronized sampling.
    
    Parameters:
    - minibatch_tensors: MiniBatch objects to align
    
    Returns:
    - list: aligned minibatch tensors
    """

def fn(outs, ins, mode=None, **kwargs):
    """
    Compile Theano function from model variables.
    
    Parameters:
    - outs: list, output variables
    - ins: list, input variables  
    - mode: str, Theano compilation mode
    - kwargs: additional arguments for theano.function
    
    Returns:
    - function: compiled Theano function
    """

def fastfn(outs, ins, mode=None, **kwargs):
    """
    Compile fast Theano function with optimizations.
    
    Parameters:
    - outs: list, output variables
    - ins: list, input variables
    - mode: str, Theano compilation mode ('FAST_COMPILE' if None)
    - kwargs: additional arguments for theano.function
    
    Returns:
    - function: compiled optimized function
    """

class FastPointFunc:
    """
    Fast function wrapper for point evaluation.
    
    Provides efficient evaluation of model functions at specific
    parameter values without full Theano overhead.
    """

class LoosePointFunc:
    """
    Loose function wrapper for flexible point evaluation.
    
    More flexible than FastPointFunc but potentially slower,
    handles shape mismatches and missing variables gracefully.
    """

class ValueGradFunction:
    """
    Combined value and gradient function for optimization.
    
    Efficiently computes both function values and gradients
    in a single pass, useful for MAP estimation and optimization.
    """

def model_to_graphviz(model, var_names=None, formatting='plain', 
                      save=None, figsize=(5, 5), dpi=300):
    """
    Convert PyMC3 model to GraphViz representation.
    
    Creates a directed graph showing relationships between model variables,
    useful for understanding model structure and debugging complex models.
    
    Parameters:
    - model: Model, PyMC3 model to visualize
    - var_names: list, variables to include (all if None)
    - formatting: str, node formatting style ('plain', 'fancy')
    - save: str, filename to save graph (displays if None)
    - figsize: tuple, figure size in inches
    - dpi: int, resolution for saved figure
    
    Returns:
    - graphviz.Digraph: GraphViz graph object
    """

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

# Basic linear regression model
with pm.Model() as basic_model:
    # Priors
    alpha = pm.Normal('alpha', mu=0, sigma=10)
    beta = pm.Normal('beta', mu=0, sigma=10)  
    sigma = pm.HalfNormal('sigma', sigma=1)
    
    # Linear predictor
    mu = alpha + beta * x_data
    
    # Likelihood
    y_obs = pm.Normal('y_obs', mu=mu, sigma=sigma, observed=y_data)
    
    # Deterministic quantities for posterior analysis
    R2 = pm.Deterministic('R2', 1 - tt.var(y_data - mu) / tt.var(y_data))

# Hierarchical model with group-level effects
with pm.Model() as hierarchical_model:
    # Hyperpriors
    mu_alpha = pm.Normal('mu_alpha', mu=0, sigma=10)
    sigma_alpha = pm.HalfNormal('sigma_alpha', sigma=5)
    
    mu_beta = pm.Normal('mu_beta', mu=0, sigma=10)
    sigma_beta = pm.HalfNormal('sigma_beta', sigma=5)
    
    # Group-level parameters
    alpha = pm.Normal('alpha', mu=mu_alpha, sigma=sigma_alpha, shape=n_groups)
    beta = pm.Normal('beta', mu=mu_beta, sigma=sigma_beta, shape=n_groups)
    
    # Observation-level parameters
    sigma = pm.HalfNormal('sigma', sigma=1)
    
    # Expected values
    mu = alpha[group_idx] + beta[group_idx] * x_data
    
    # Likelihood
    y_obs = pm.Normal('y_obs', mu=mu, sigma=sigma, observed=y_data)

# Model with custom potential for non-standard likelihood
with pm.Model() as custom_model:
    # Parameters
    theta = pm.Beta('theta', alpha=1, beta=1)
    
    # Custom log-likelihood function
    def logp_custom(obs, theta):
        # Example: zero-truncated Poisson
        return tt.sum(obs * tt.log(theta) - theta - tt.gammaln(obs + 1) 
                     - tt.log(1 - tt.exp(-theta)))
    
    # Add custom likelihood as potential
    likelihood = pm.Potential('likelihood', 
                             logp_custom(observed_data, theta))

# Model using shared data for prediction
x_shared = pm.Data('x_shared', x_train)
y_shared = pm.Data('y_shared', y_train)

with pm.Model() as prediction_model:
    # Model parameters
    alpha = pm.Normal('alpha', mu=0, sigma=10)
    beta = pm.Normal('beta', mu=0, sigma=10)
    sigma = pm.HalfNormal('sigma', sigma=1)
    
    # Linear predictor using shared data
    mu = alpha + beta * x_shared
    
    # Likelihood
    y_obs = pm.Normal('y_obs', mu=mu, sigma=sigma, observed=y_shared)
    
    # Sample posterior
    trace = pm.sample(1000, tune=1000)

# Update data for out-of-sample prediction
pm.set_data({'x_shared': x_test}, model=prediction_model)

# Generate posterior predictions
with prediction_model:
    post_pred = pm.sample_posterior_predictive(trace)

# Model with multiple deterministic quantities
with pm.Model() as transform_model:
    # Raw parameters
    raw_effects = pm.Normal('raw_effects', mu=0, sigma=1, shape=5)
    raw_scale = pm.Normal('raw_scale', mu=0, sigma=1)
    
    # Transformed parameters
    effects = pm.Deterministic('effects', raw_effects * 0.5)
    scale = pm.Deterministic('scale', tt.exp(raw_scale))
    
    # Derived quantities
    total_effect = pm.Deterministic('total_effect', tt.sum(effects))
    effect_variance = pm.Deterministic('effect_variance', tt.var(effects))
    
    # Model prediction
    prediction = pm.Deterministic('prediction', 
                                 tt.dot(x_data, effects))
    
    # Likelihood
    y_obs = pm.Normal('y_obs', mu=prediction, sigma=scale, 
                     observed=y_data)

# Model with minibatch training for large datasets
batch_size = 100
n_data = len(large_dataset)

# Create minibatch containers
x_minibatch = pm.Minibatch(x_large, batch_size=batch_size)
y_minibatch = pm.Minibatch(y_large, batch_size=batch_size)

with pm.Model() as minibatch_model:
    # Model parameters
    weights = pm.Normal('weights', mu=0, sigma=1, shape=n_features)
    intercept = pm.Normal('intercept', mu=0, sigma=10)
    
    # Scale factor for minibatch
    scale_factor = n_data // batch_size
    
    # Linear predictor
    prediction = intercept + tt.dot(x_minibatch, weights)
    
    # Scaled likelihood for minibatch
    y_obs = pm.Normal('y_obs', mu=prediction, sigma=1, 
                     observed=y_minibatch, 
                     total_size=n_data)

# Model checking and diagnostics
with pm.Model() as diagnostic_model:
    # Model specification...
    theta = pm.Beta('theta', alpha=2, beta=2)
    y_obs = pm.Binomial('y_obs', n=10, p=theta, observed=data)
    
    # Check model evaluation at test point
    test_point = diagnostic_model.test_point
    print("Test point:", test_point)
    
    # Check log-probability evaluation
    logp_check = diagnostic_model.check_test_point()
    print("Log-probability check:", logp_check)
    
    # Compile log-probability function
    logp_fn = diagnostic_model.compile_logp()
    
    # Compile gradient function  
    dlogp_fn = diagnostic_model.compile_dlogp()
    
    # Test function evaluation
    logp_val = logp_fn(test_point)
    grad_val = dlogp_fn(test_point)