R_code.qmd

---
title: "R Code Test"
author: "Abdullah Mahmood"
date: today
format: html
editor: source
jupyter: main
---

```{python}
from cmdstanpy import CmdStanModel
import bridgestan as bs
import numpy as np
import os
import json
from scipy.special import expit
import scipy.stats as stats
import pandas as pd

import logging
import warnings
warnings.simplefilter(action='ignore', category=FutureWarning)
warnings.filterwarnings( "ignore", module = "plotnine/..*" )

import cmdstanpy as csp
csp.utils.get_logger().setLevel(logging.ERROR)

def inline_plot(plot_func, *args, **kwargs):
    plt.clf()  
    plot_func(*args, **kwargs)
    plt.show()
    plt.close()
```

```{python}
class Stan(CmdStanModel):
    def __init__(self, stan_file: str, stan_code: str, force_compile=False):
        """Load or compile a Stan model"""
        stan_src = f"{stan_file}.stan"
        exe_file = stan_file
        
        # Check for the compiled executable
        if not os.path.isfile(exe_file) or force_compile:
            with open(stan_src, 'w') as f:
                f.write(stan_code)
            super().__init__(stan_file=stan_src, force_compile=True, cpp_options={'STAN_THREADS': 'true', 'parallel_chains': 4})
        else:
            super().__init__(stan_file=stan_src, exe_file=exe_file)

class BridgeStan(bs.StanModel):
    def __init__(self, stan_file: str, data: dict, force_compile=False):
        """Load or compile a BridgeStan shared object"""
        stan_so = f"{stan_file}_model.so"
        make_args = ['BRIDGESTAN_AD_HESSIAN=true', 'STAN_THREADS=true']
        data = json.dumps(data)
        if not os.path.isfile(stan_so) or force_compile:  # If the shared object does not exist, compile it
            super().__init__(f"{stan_file}.stan", data, make_args=make_args)
        else:
            super().__init__(stan_so, data, make_args=make_args, warn=False)

class StanQuap(object):
    """
    Description:
    Find mode of posterior distribution for arbitrary fixed effect models and 
    then produce an approximation of the full posterior using the quadratic    
    curvature at the mode.
  
    This command provides a convenient interface for finding quadratic approximations 
    of posterior distributions for models defined in Stan. This procedure is equivalent 
    to penalized maximum likelihood estimation and the use of a Hessian for profiling, 
    and therefore can be used to define many common regularization procedures. The point 
    estimates returned correspond to a maximum a posterior, or MAP, estimate. However the 
    intention is that users will use `extract_samples` and `laplace_sample` and other methods to work 
    with the full posterior.
    """
    def __init__(self,
                 stan_file: str, 
                 stan_code: str, 
                 data: dict, 
                 algorithm = 'Newton',
                 jacobian: bool = False,
                 force_compile = False,
                 **kwargs):
        self.train_data = data
        self.stan_model = Stan(stan_file, stan_code, force_compile)
        self.bs_model = BridgeStan(stan_file, self.train_data, force_compile)
        self.opt_model = self.stan_model.optimize(
                              data=self.train_data,
                              algorithm=algorithm,
                              jacobian=jacobian,
                              **kwargs
                        )
        self.param_names = self.bs_model.param_names()
        self.opt_params = {param: self.opt_model.stan_variable(param) for param in self.param_names}
        self.params_unc = self.bs_model.param_unconstrain(
                              np.array(list(self.opt_params.values()))
                        )
        self.jacobian = jacobian

    def log_density_hessian(self):
        log_dens, gradient, hessian = self.bs_model.log_density_hessian(
            self.params_unc, 
            jacobian=self.jacobian
        )
        return log_dens, gradient, hessian
    
    def vcov_matrix(self, param_types=None, eps=1e-6):
        _, _, hessian_unc = self.log_density_hessian()
        vcov_unc = np.linalg.inv(-hessian_unc)
        cov_matrix = self.transform_vcov(vcov_unc, param_types, eps)
        return cov_matrix
    
    def laplace_sample(self, data: dict = None, draws: int = 100_000):
        if data is None:
            data = self.train_data
        return self.stan_model.laplace_sample(data=data, 
                                              mode=self.opt_model, 
                                              draws=draws, 
                                              jacobian=self.jacobian)
      
    def extract_samples(self, n: int = 100_000, dict_out: bool = True):
        laplace_obj = self.laplace_sample(draws=n)
        if dict_out:
            stan_var_dict = laplace_obj.stan_variables()
            return {param: stan_var_dict[param] for param in self.param_names}
        return laplace_obj.draws()
    
    def sim(self, data: dict = None, n = 1000, dict_out: bool = True):
        if data is None:
            data = self.train_data      
        laplace_obj = self.laplace_sample(data=data, draws=n)
        if dict_out:
            stan_var_dict = laplace_obj.stan_variables()
            return {param: stan_var_dict[param] for param in stan_var_dict if param not in self.param_names}
        return laplace_obj.draws()

    def compute_jacobian_analytical(self, param_types):
        """
        Analytical computation of the Jacobian matrix for transforming 
        variance-covariance matrix from unconstrained to constrained space.
        """
        dim = len(self.params_unc)
        J = np.zeros((dim, dim))  # Initialize Jacobian matrix
        
        for i in range(dim):
            if param_types[i] == 'uncons':  # Unconstrained (Identity transformation)
                J[i, i] = 1
            elif param_types[i] == 'pos_real':  # Positive real (Exp transformation)
                J[i, i] = np.exp(self.params_unc[i])
            elif param_types[i] == 'prob':  # Probability (Logit transformation)
                x = 1 / (1 + np.exp(-self.params_unc[i]))  # Sigmoid function
                J[i, i] = x * (1 - x)
            else:
                raise ValueError(f"Unknown parameter type: {param_types[i]}")
      
        return J      

    def compute_jacobian_numerical(self, eps=1e-6):
        """
         Analytical computation of the Jacobian matrix for transforming 
         variance-covariance matrix from unconstrained to constrained space.
        """
        dim = len(self.params_unc)
        J = np.zeros((dim, dim))  # Full Jacobian matrix
    
        # Compute Jacobian numerically for each parameter
        for i in range(dim):
            perturbed = self.params_unc.copy()
            # Perturb parameter i
            perturbed[i] += eps
            constrained_plus = np.array(self.bs_model.param_constrain(perturbed))
            perturbed[i] -= 2 * eps
            constrained_minus = np.array(self.bs_model.param_constrain(perturbed))
            # Compute numerical derivative
            J[:, i] = (constrained_plus - constrained_minus) / (2 * eps)
    
        return J

    def transform_vcov(self, vcov_unc, param_types=None, eps=1e-6):
        """
        Transform the variance-covariance matrix from the unconstrained space to the constrained space.
        Args:
        - vcov_unc (np.array): variance-covariance matrix in the unconstrained space.
        - param_types (list) [Required for analytical solution]: List of strings specifying the type of each parameter. 
          Options: 'uncons' (unconstrained), 'pos_real' (positive real), 'prob' (0 to 1).
        - eps (float) [Required for numerical solution]: Small perturbation for numerical differentiation.
        Returns:
        - vcov_con (np.array): variance-covariance matrix in the constrained space.
        """
        if param_types is None:
            J = self.compute_jacobian_numerical(eps)
        else:
            J = self.compute_jacobian_analytical(param_types)
        vcov_con = J.T @ vcov_unc @ J
        return vcov_con

    def precis(self, param_types=None, prob=0.89, eps=1e-6):
        vcov_mat = self.vcov_matrix(param_types, eps)
        pos_mu = np.array(list(self.opt_params.values()))
        pos_sigma = np.sqrt(np.diag(vcov_mat))
        plo = (1-prob)/2
        phi = 1 - plo
        lo = pos_mu + pos_sigma * stats.norm.ppf(plo)
        hi = pos_mu + pos_sigma * stats.norm.ppf(phi)
        res = pd.DataFrame({
          'Parameter': list(self.opt_params.keys()),
          'Mean': pos_mu,
          'StDev': pos_sigma,
          f'{plo:.1%}': lo,
          f'{phi:.1%}': hi})
        return res.set_index('Parameter')
```


```{python}
d = pd.read_csv("data/Howell1.csv", sep=';')
d2 = d[d['age'] >= 18]

m4_3_pred = '''
data {
    int<lower=1> N;
    vector[N] height;
    vector[N] weight;
    real xbar;
    
    int<lower=1> N_tilde;
    vector[N_tilde] weight_tilde;
}
parameters {
    real a;
    real<lower=0> b;
    real<lower=0, upper=50> sigma;
}
model {
    vector[N] mu;
    mu = a + b * (weight - xbar);
    
    // Likelihood Function
    height ~ normal(mu, sigma);
    
    // Priors
    a ~ normal(178, 20);
    b ~ lognormal(0, 1);
    sigma ~ uniform(0, 50);
}
generated quantities {
    vector[N_tilde] y_tilde; 
    for (i in 1:N_tilde) {
        y_tilde[i] = normal_rng(a + b * (weight_tilde[i] - xbar), sigma);
    }
}
'''
weight_seq = np.arange(25, 71)

data = {'N': len(d2), 
        'xbar': d2['weight'].mean(),
        'height': d2['height'].tolist(), 
        'weight': d2['weight'].tolist(),
        'N_tilde': len(weight_seq),
        'weight_tilde': weight_seq.tolist()}

m4_3_pred_model = StanQuap('stan_models/m4_3_pred', m4_3_pred, data, algorithm = 'Newton')
m4_3_pred_model.precis().round(2)
m4_3_pred_model.vcov_matrix().round(3)

m4_3_pred_model.sim(data=data, n=100)
```

```{python}
test_stan = '''
data {
    int<lower=0> W;
    int<lower=0> L;
}
parameters {
    real<lower=0, upper=1> p;
}
model {
    p ~ uniform(0, 1);
    W ~ binomial(W + L, p);
}
'''

data = {'W': 24, 'L': 36-24}
model = StanQuap('stan_models/test', test_stan, data)
model.precis().round(7)
model.precis(param_types=['prob'])
```

```{python}
res_df = model.precis().round(7)
res_df

laplace_draws = model.extract_samples()['p']
sigma_draw = laplace_draws.std()
sigma_draw

laplace_draws.mean()

import seaborn as sns
import matplotlib.pyplot as plt
import scipy.stats as stats

def dens_plot():
  x = np.linspace(0,1,100)
  y = stats.norm.pdf(x, model.opt_params['p'], sigma_draw)
  plt.plot(x, y, label='Laplace Draws')
  y = stats.norm.pdf(x, res_df['Mean'][0], res_df['StDev'][0])
  plt.plot(x, y, label='MAP')
  plt.plot(x, stats.beta.pdf(x, 24 + 1, 36 - 24 + 1), label='True Posterior')
  plt.legend()

inline_plot(dens_plot)
```

```{python}
d = pd.read_csv("data/Howell1.csv", sep=';')
d2 = d[d['age'] >= 18]

m4_1 = '''
data {
    int<lower=0> N;
    vector[N] height;
}
parameters {
    real mu;
    real<lower=0, upper=50> sigma;
}
model {
    height ~ normal(mu, sigma);
    mu ~ normal(178, 20);
    sigma ~ uniform(0, 50);
}
'''
data = {'N': len(d2), 'height': d2['height'].tolist()}

model = StanQuap('stan_models/m4_1', m4_1, data)

model.precis().round(2)

model.vcov_matrix()

np.diag(model.vcov_matrix())

def cov2cor( A ):
    """
    covariance matrix to correlation matrix.
    """
    d = np.sqrt(A.diagonal()) # Compute standard deviations
    A = ((A.T/d).T)/d # # Normalize each element by sqrt(C_ii * C_jj)
    return A

cov2cor(model.vcov_matrix())

def cov2cor(c: np.ndarray) -> np.ndarray:
    """
    Return a correlation matrix given a covariance matrix.
    : c = covariance matrix
    """
    D = np.zeros(c.shape)
    np.fill_diagonal(D, np.sqrt(np.diag(c)))
    invD = np.linalg.inv(D)
    return invD @ c @ invD

cov2cor(model.vcov_matrix())
```

```{python}
m4_2 = '''
data {
    int<lower=0> N;
    vector[N] height;
}
parameters {
    real mu;
    real<lower=0, upper=50> sigma;
}
model {
    height ~ normal(mu, sigma);
    mu ~ normal(178, 0.1);
    sigma ~ uniform(0, 50);
}
'''
data = {'N': len(d2), 'height': d2['height'].tolist()}

model = StanQuap('stan_models/m4_2', m4_2, data)

model.precis().round(2)
```


```{python}
m4_3 = '''
data {
    int<lower=0> N;
    vector[N] height;
    vector[N] weight;
    real xbar;
}
parameters {
    real a;
    real<lower=0> b;
    real<lower=0, upper=50> sigma;
}
model {
    vector[N] mu;
    mu = a + b * (weight - xbar);
    
    // Likelihood Function
    height ~ normal(mu, sigma);
    
    // Priors
    a ~ normal(178, 20);
    b ~ lognormal(0, 1);
    sigma ~ uniform(0, 50);
}
'''

data = {'N': len(d2), 
        'xbar': d2['weight'].mean(),
        'height': d2['height'].tolist(), 
        'weight': d2['weight'].tolist()}

model = StanQuap('stan_models/m4_3', m4_3, data, algorithm = 'Newton')
model.precis().round(2)
model.vcov_matrix().round(3)
```


```{r}
library(rethinking)
data(Howell1)
d <- Howell1
d2 <- d[ d$age >= 18 , ]
```

```{r}
flist <- alist(
height ~ dnorm( mu , sigma ) ,
mu ~ dnorm( 178 , 20 ) ,
sigma ~ dunif( 0 , 50 )
)

m4.1 <- quap( flist , data=d2 )

precis( m4.1 )
```

```{r}
# load data again, since it's a long way back
library(rethinking)
data(Howell1); d <- Howell1; d2 <- d[ d$age >= 18 , ]

# define the average weight, x-bar
xbar <- mean(d2$weight)

# fit model
m4.3 <- quap(
  alist(
    height ~ dnorm( mu , sigma ) ,
    mu <- a + b*( weight - xbar ) ,
    a ~ dnorm( 178 , 20 ) ,
    b ~ dlnorm( 0 , 1 ) ,
    sigma ~ dunif( 0 , 50 )
  ) , data=d2 )

precis( m4.3 )

coef(m4.3)
```

```{r}
globe.qa <- quap(
    alist(
        W ~ dbinom( W+L ,p) , # binomial likelihood
        p ~ dunif(0,1) # uniform prior
    ),
    data=list(W=24,L=36-24))

# display summary of quadratic approximation
round(precis(globe.qa), 5)
(vcov(globe.qa))^(1/2)
vcov(globe.qa)
```

```{r}
print(globe.qa)
```

```{r}

# analytical calculation
W <- 24
L <- 36-24
curve( dbeta( x , W+1 , L+1 ) , from=0 , to=1 )
# quadratic approximation
curve( dnorm( x , 0.67 , 0.07857 ) , lty=2 , add=TRUE )
```



```{r}
library(rethinking)
data(Howell1)
d <- Howell1
d2 <- d[ d$age >= 18 , ]
plot( d2$height ~ d2$weight )
```


```{r}
# load data again, since it's a long way back
library(rethinking)
data(Howell1); d <- Howell1; d2 <- d[ d$age >= 18 , ]
# define the average weight, x-bar
xbar <- mean(d2$weight)
# fit model

m4.3 <- quap(
alist(
height ~ dnorm( mu , sigma ) ,
mu <- a + b*( weight - xbar ) ,
a ~ dnorm( 178 , 20 ) ,
b ~ dlnorm( 0 , 1 ) ,
sigma ~ dunif( 0 , 50 )
) , data=d2 )
```

```{r}
precis( m4.3 )
pairs(m4.3)
```

```{r}
# define sequence of weights to compute predictions for
# these values will be on the horizontal axis
weight.seq <- seq( from=25 , to=70 , by=1 )
# use link to compute mu
# for each sample from posterior
# and for each weight in weight.seq
mu <- link( m4.3 , data=data.frame(weight=weight.seq) )
str(mu)
```

```{r}
# use type="n" to hide raw data
plot( height ~ weight , d2 , type="n" )
# loop over samples and plot each mu value
for ( i in 1:100 )
points( weight.seq , mu[i,] , pch=16 , col=col.alpha(rangi2,0.1) )
```

```{r}
sim.height <- sim( m4.3 , data=list(weight=weight.seq) )
sim.mean <- apply( sim.height , 2 , mean )
sim.PI <- apply( sim.height , 2 , PI , prob=0.89 )

sim.mean
sim.PI
```