# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# This extension serves a narrow use case for writing, compiling, error-checking and validating stan models within jupyter notebook. I am not a comfortable R user and STAN currently leans towards R. I wrote this simple script as I wanted an uninterruped workflow for data manipulation, stan model creation and analysis  all in jupyter notebook. Once the codemirror support is merged in, sysntax highlighting for stan code could also be enabled. Since  few people were asking me for this, I have packaged this as a extension.
#
# Tested to work on Linux/Mac. Python versions 2.7/3.6

# ## Installation

# `pip install git+https://github.com/Arvinds-ds/stanmagic.git`

# ## Testing %%stan magic

import numpy as np
from collections import OrderedDict
import pystan

# ### Load stanmagic extension

# %load_ext stanmagic 

# #### Optional: Using stanc compiler instead of pystan.stanc compiler
#
# To get cmdstan installed:-
#
# 1. Downlad latest cmdstan-xxxx.tar.gz https://github.com/stan-dev/cmdstan/releases
# 2. extract to <cmdstan_path>
# 3. `cd <cmdstan_path>`
# 4. `make`
# 3. `make build`
# 4. export PATH=$PATH:<cmdstan_path>/bin/stanc
#
#
# If the compiler is not in your path or you don't want to edit your PATH, you may have to pass --stanc [compiler_path] to %%stan. See below

# ### Generate test data

X = np.random.randn(100,3)
beta = np.array([0.1,0.2,0.3])
alpha = 4
sigma = 1.7
y = np.dot(X,beta) + alpha + np.random.randn(100)*sigma
N=100
K=3

data = OrderedDict({'X':X, 'y': y, 'N':N, 'K': K})

# ### 1. %%stan -f [stan_file_name]
#
# Saves the cell code to a file specified in [stan_file_name]. The file name can also be accessed in __`_stan_model.model_file`__ generated in local namespace

# +
# %%stan -f test.stan
data {
  
  int<lower=1> N;
  int<lower=1> K;
  
  matrix[N,K] X;
  vector[N] y;
}

parameters{
  real alpha;
  ordered[K] beta;
  real<lower=0> sigma;
  
}

model {
  
  alpha ~ normal(0,10);
  beta[1] ~ normal(.1,.1);
  beta[2] ~ normal(.2,.1);
  beta[3] ~ normal(.3,.1);
  sigma ~ exponential(1);
  for (n in 1:N) {
    y[n] ~ normal(X[n]*beta + alpha, sigma);
  } 
  
}

generated quantities {
  vector[N] y_rep;
  for (n in 1:N) {
    y_rep[n] = normal_rng(X[n]*beta + alpha, sigma);
  }
}
# -

_stan_model

model = pystan.StanModel(file=_stan_model.model_file)

model.sampling(data=data)

# ### 2. %%stan -v model_test

# Saves the cell code to a code string. The `-v` specifies the object name where the output object is saved.
# The model code can  be accessed in __`model_test.model_code`__ generated in local namespace

# +
# %%stan -v model_test

data {
  
  int<lower=1> N;
  int<lower=1> K;
  
  matrix[N,K] X;
  vector[N] y;
}

parameters{
  real alpha;
  ordered[K] beta;
  real<lower=0> sigma;
  
}

model {
  
  alpha ~ normal(0,10);
  beta[1] ~ normal(.1,.1);
  beta[2] ~ normal(.2,.1);
  beta[3] ~ normal(.3,.1);
  sigma ~ exponential(1);
  for (n in 1:N) {
    y[n] ~ normal(X[n]*beta + alpha, sigma);
  } 
  
}

generated quantities {
  vector[N] y_rep;
  for (n in 1:N) {
    y_rep[n] = normal_rng(X[n]*beta + alpha, sigma);
  }
}
# -

model_test

model = pystan.StanModel(model_code=model_test.model_code)

# ### 3. %%stan -f test3.stan -v test3_model

# +
# %%stan -f test3.stan -v test3_model

data {
  
  int<lower=1> N;
  int<lower=1> K;
  
  matrix[N,K] X;
  vector[N] y;
}

parameters{
  real alpha;
  ordered[K] beta;
  real<lower=0> sigma;
  
}

model {
  
  alpha ~ normal(0,10);
  beta[1] ~ normal(.1,.1);
  beta[2] ~ normal(.2,.1);
  beta[3] ~ normal(.3,.1);
  sigma ~ exponential(1);
  for (n in 1:N) {
    y[n] ~ normal(X[n]*beta + alpha, sigma);
  } 
  
}

generated quantities {
  vector[N] y_rep;
  for (n in 1:N) {
    y_rep[n] = normal_rng(X[n]*beta + alpha, sigma);
  }
}
# -

test3_model

model = pystan.StanModel(file=test3_model.model_file)

# ### 3. %%stan -f [stan_file_name] --save_only
#
# Saves the cell code to a file specified in [stan_file_name].  Skips compile step

# +
# %%stan -f test1.stan --save_only
data {
  
  int<lower=1> N;
  int<lower=1> K;
  
  matrix[N,K] X;
  vector[N] y;
}

parameters{
  real alpha;
  ordered[K] beta;
  real<lower=0> sigma;
  
}

model {
  
  alpha ~ normal(0,10);
  beta[1] ~ normal(.1,.1);
  beta[2] ~ normal(.2,.1);
  beta[3] ~ normal(.3,.1);
  sigma ~ exponential(1);
  for (n in 1:N) {
    y[n] ~ normal(X[n]*beta + alpha, sigma);
  } 
  
}

generated quantities {
  vector[N] y_rep;
  for (n in 1:N) {
    y_rep[n] = normal_rng(X[n]*beta + alpha, sigma);
  }
}
# -

model = pystan.StanModel(file='test1.stan')

# ### 4. %%stan 
#
# Saves the cell code to a code string. The code string can be accessed via __`_stan_model.model_code`__

# +
# %%stan
data {
  
  int<lower=1> N;
  int<lower=1> K;
  
  matrix[N,K] X;
  vector[N] y;
}

parameters{
  real alpha;
  ordered[K] beta;
  real<lower=0> sigma;
  
}

model {
  
  alpha ~ normal(0,10);
  beta[1] ~ normal(.1,.1);
  beta[2] ~ normal(.2,.1);
  beta[3] ~ normal(.3,.1);
  sigma ~ exponential(1);
  for (n in 1:N) {
    y[n] ~ normal(X[n]*beta + alpha, sigma);
  } 
  
}

generated quantities {
  vector[N] y_rep;
  for (n in 1:N) {
    y_rep[n] = normal_rng(X[n]*beta + alpha, sigma);
  }
}
# -

_stan_model.model_code

model = pystan.StanModel(model_code=_stan_model.model_code)

model.sampling(data=data)

# ### 5. %%stan --stanc pystan

# +
# %%stan --stanc pystan
data {
  
  int<lower=1> N;
  int<lower=1> K;
  
  matrix[N,K] X;
  vector[N] y;
}

parameters{
  real alpha;
  ordered[K] beta;
  real<lower=0> sigma;
  
}

model {
  
  alpha ~ normal(0,10);
  beta[1] ~ normal(.1,.1);
  beta[2] ~ normal(.2,.1);
  beta[3] ~ normal(.3,.1);
  sigma ~ exponential(1);
  for (n in 1:N) {
    y[n] ~ normal(X[n]*beta + alpha, sigma);
  } 
  
}

generated quantities {
  vector[N] y_rep;
  for (n in 1:N) {
    y_rep[n] = normal_rng(X[n]*beta + alpha, sigma);
  }
}
# -

model = pystan.StanModel(model_code=_stan_model.model_code)

# ### 6. %%stan --stanc "~/Downloads/cmdstan-2.16.0/bin/stanc"

# +
# %%stan --stanc "~/Downloads/cmdstan-2.16.0/bin/stanc"
data {
  
  int<lower=1> N;
  int<lower=1> K;
  
  matrix[N,K] X;
  vector[N] y;
}

parameters{
  real alpha;
  ordered[K] beta;
  real<lower=0> sigma;
  
}

model {
  
  alpha ~ normal(0,10);
  beta[1] ~ normal(.1,.1);
  beta[2] ~ normal(.2,.1);
  beta[3] ~ normal(.3,.1);
  sigma ~ exponential(1);
  for (n in 1:N) {
    y[n] ~ normal(X[n]*beta + alpha, sigma);
  } 
  
}

generated quantities {
  vector[N] y_rep;
  for (n in 1:N) {
    y_rep[n] = normal_rng(X[n]*beta + alpha, sigma);
  }
}
# -

model = pystan.StanModel(model_code=_stan_model.model_code)

# ### 7. %%stan -f [stan_file_name] -o [cpp_file_name]
# Saves the cell code to a file specified in [stan_file_name] and outputs the compiled cpp file to the file name specified by [cpp_file_name]
#

# ### 8. %% stan -f  [stan_file_name] --allow_undefined
# passes the --allow_undefined argument to stanc compiler

# ### 9.%%stan -f  [stan_file_name] --stanc [stanc_compiler_path]
# Saves the cell code to a file specified in [stan_file_name] and compiles using the stan compiler specified in [stanc_compiler]. By default, it uses stanc compiler in your path. If your path does not have the stanc compiler, use this option (e.g %%stan binom.stan --stanc "~/cmdstan-2.16.0/bin/stanc")

assert True



StarcoderdataJupyterScriptsDedupFiltered