diff --git a/.gitignore b/.gitignore index 00ca225b..389146ff 100644 --- a/.gitignore +++ b/.gitignore @@ -63,3 +63,16 @@ Temporary Items # IDE .vscode +# build folder +build + +# egg +pcntoolkit.egg-info + +# Demo folder +HBR_demo + +ssh_error_log.txt +HBR_demo/* +dist/pcntoolkit-0.27-py3.11.egg + diff --git a/CHANGES b/CHANGES index c41977f2..93e0fa36 100644 --- a/CHANGES +++ b/CHANGES @@ -73,3 +73,8 @@ version 0.27 - Fixed a translation problem between the previous naming convention for HBR models (only Gaussian models) and the current naming (also SHASH models) - Minor updates to fix synchronisation problems in PCNportal (related to the HBR updates above) - Added configuration files for containerisation with Docker + +version 0.28 +- Updated to PyMC5 (including migrating back-end to PyTensor +- Added support for longitudinal normative modelling with BLR (see Buckova-Rehak et al 2023) +- Changed default optimiser for trend surface models (for scalability) diff --git a/build/lib/pcntoolkit/configs.py b/build/lib/pcntoolkit/configs.py new file mode 100644 index 00000000..98b56f17 --- /dev/null +++ b/build/lib/pcntoolkit/configs.py @@ -0,0 +1,9 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Created on Mon Dec 7 12:51:07 2020 + +@author: seykia +""" + +PICKLE_PROTOCOL = 4 diff --git a/dist/pcntoolkit-0.27-py3.11.egg b/dist/pcntoolkit-0.27-py3.11.egg new file mode 100644 index 00000000..7ba241b2 Binary files /dev/null and b/dist/pcntoolkit-0.27-py3.11.egg differ diff --git a/dist/pcntoolkit-0.28-py3.10.egg b/dist/pcntoolkit-0.28-py3.10.egg new file mode 100644 index 00000000..d9d9b9f2 Binary files /dev/null and b/dist/pcntoolkit-0.28-py3.10.egg differ diff --git a/doc/build/html/_modules/utils.html b/doc/build/html/_modules/utils.html index dd5365f3..2bd7c2bc 100644 --- a/doc/build/html/_modules/utils.html +++ b/doc/build/html/_modules/utils.html @@ -1,181 +1,198 @@ - - - + + + - - - - - utils — Predictive Clinical Neuroscience Toolkit 0.20 documentation - - - - - - - - - - - - - - - - - - + + + + + utils — Predictive Clinical Neuroscience Toolkit 0.20 documentation + + + + + + + + + + + + + + + + + + - - - - - - - - + + + + + + + + - + - -
- - + +
+ + + + + +
+ +
+ + + + + + + + + + + + + + + + + +
+ + + + +
+
+
+
+ +

Source code for utils

+
+ 
 from __future__ import print_function
 
 import os
@@ -191,7 +208,7 @@ 
Source code for utils
 import bspline
 from bspline import splinelab
 from sklearn.datasets import make_regression
-import pymc3 as pm
+import pymc as pm
 
 # -----------------
 # Utility functions
@@ -947,45 +964,49 @@ 
Source code for utils
     plt.xticks(rotation=90, fontsize=7)
     plt.tight_layout()
     plt.show()

-

+
+
-
- -
-
- +
-
+
+ +
+

+ © Copyright 2020, Andre F. Marquand -

-
+

+ + Built with Sphinx using a theme provided by Read the Docs. + + + + + + + + + + + + + - - - - - - - - + \ No newline at end of file diff --git a/doc/build/html/_sources/pages/HBR_NormativeModel_FCONdata_Tutorial.rst.txt b/doc/build/html/_sources/pages/HBR_NormativeModel_FCONdata_Tutorial.rst.txt index 4db24f06..d4933e38 100644 --- a/doc/build/html/_sources/pages/HBR_NormativeModel_FCONdata_Tutorial.rst.txt +++ b/doc/build/html/_sources/pages/HBR_NormativeModel_FCONdata_Tutorial.rst.txt @@ -24,7 +24,6 @@ Step 0: Install necessary libraries & grab data files .. code:: ipython3 - ! pip uninstall -y Theano-PyMC # conflicts with Theano on some environments ! pip install pcntoolkit==0.26 diff --git a/doc/build/html/_sources/pages/normative_modelling_walkthrough.rst.txt b/doc/build/html/_sources/pages/normative_modelling_walkthrough.rst.txt index 06eb5017..f63ec8fa 100644 --- a/doc/build/html/_sources/pages/normative_modelling_walkthrough.rst.txt +++ b/doc/build/html/_sources/pages/normative_modelling_walkthrough.rst.txt @@ -46,7 +46,6 @@ bad if you cannot complete everything during the tutorial. .. code:: ipython3 #install normative modeling - ! pip uninstall -y Theano-PyMC # conflicts with Theano on some environments ! pip install pcntoolkit==0.26 diff --git a/doc/build/html/pages/HBR_NormativeModel_FCONdata_Tutorial.html b/doc/build/html/pages/HBR_NormativeModel_FCONdata_Tutorial.html index ec9e668b..2e1a829a 100644 --- a/doc/build/html/pages/HBR_NormativeModel_FCONdata_Tutorial.html +++ b/doc/build/html/pages/HBR_NormativeModel_FCONdata_Tutorial.html @@ -207,7 +207,7 @@

Hierarchical Bayesian Regressionhttps://colab.research.google.com/assets/colab-badge.svg

Step 0: Install necessary libraries & grab data files

-
! pip uninstall -y Theano-PyMC  # conflicts with Theano on some environments
+
! 
 ! pip install pcntoolkit==0.26
 

 

diff --git a/doc/build/html/pages/normative_modelling_walkthrough.html b/doc/build/html/pages/normative_modelling_walkthrough.html
index 0ed663c0..d2b3d562 100644
--- a/doc/build/html/pages/normative_modelling_walkthrough.html
+++ b/doc/build/html/pages/normative_modelling_walkthrough.html
@@ -1,39 +1,40 @@
-
-
 
 
-  
+
+ 
+
 
   
   
 
   
-  
+
   GPR tutorial — Predictive Clinical Neuroscience Toolkit 0.20 documentation
-  
 
-  
-  
-  
-  
 
-  
+
+
+
+
+
+
   
-  
-    
-      
-        
-        
-        
-        
-        
-        
-    
-    
-
-    
-
-  
+
+
+  
+  
+  
+  
+  
+  
+  
+
+  
+
+
+
+
   
   
   
@@ -43,118 +44,142 @@
   
   
   
-    
-    
-    
-     
+  
+  
+  
+  
 
 
 
 
-   
+
   

-    
+
     
 
     

 
-      
+
       
 
 
       

-        
+
         

-        
-          
 
 
 
@@ -170,74 +195,91 @@
 
 
 
-

 
-  
 
-  
-  

-

+          

+
+            
+
+
+            

+          

           

-           

-            
-  

-
Gaussian Process Regression

-
Created by

-
Mariam Zabihi @m_zabihi

-
Saige Rutherford @being_saige

-
Thomas Wolfers @ThomasWolfers
-_______________________________________________________________________________

-https://colab.research.google.com/assets/colab-badge.svg
-

-
Background Story

-
Morten and Ingrid are concerned about the health of their father,
-Nordan. He recently turned 65 years. A few months ago he could not find
-his way home. Together, they visit a neurologist/psychiatrist to conduct
-a number of cognitive tests. However, those tests were inconclusive.
-While Nordan has a relatively low IQ it could not explain his trouble
-returning home.

-
Recently, the family heard about a new screening technique called
-normative modeling with which one can place individuals in reference to
-a population norm on for instance measures such as brain volume. Nordan
-would like to undertake this procedure to better know what is going on
-and to potentially find targets for treatment. Therefore, the family
-booked an appointment with you, the normative modeling specialist. To
-find out what is going on you compare Nordan’s hyppocampus to the norm
-and to a group of persons with Dementia disorders, who have a similar
-IQ, age as well as the same sex as Nordan.

-
Do your best to get as far as you can. However, you do not need to feel
-bad if you cannot complete everything during the tutorial.

-

-

-
Task 0: Load data and install the pcntoolkit

-
#install normative modeling
-! pip uninstall -y Theano-PyMC  # conflicts with Theano on some environments
+            

+
+              

+                
Gaussian Process Regression

+                
Created by

+                
Mariam Zabihi @m_zabihi

+                
Saige Rutherford @being_saige

+                
Thomas Wolfers @ThomasWolfers
+                  _______________________________________________________________________________

+                https://colab.research.google.com/assets/colab-badge.svg
+                

+                  
Background Story

+                  
Morten and Ingrid are concerned about the health of their father,
+                    Nordan. He recently turned 65 years. A few months ago he could not find
+                    his way home. Together, they visit a neurologist/psychiatrist to conduct
+                    a number of cognitive tests. However, those tests were inconclusive.
+                    While Nordan has a relatively low IQ it could not explain his trouble
+                    returning home.

+                  
Recently, the family heard about a new screening technique called
+                    normative modeling with which one can place individuals in reference to
+                    a population norm on for instance measures such as brain volume. Nordan
+                    would like to undertake this procedure to better know what is going on
+                    and to potentially find targets for treatment. Therefore, the family
+                    booked an appointment with you, the normative modeling specialist. To
+                    find out what is going on you compare Nordan’s hyppocampus to the norm
+                    and to a group of persons with Dementia disorders, who have a similar
+                    IQ, age as well as the same sex as Nordan.

+                  
Do your best to get as far as you can. However, you do not need to feel
+                    bad if you cannot complete everything during the tutorial.

+                

+                

+                  
Task 0: Load data and install the pcntoolkit

+                  

+                    

+                      
#install normative modeling
+!
 ! pip install pcntoolkit==0.26
-

-

-
Option 1: Connect your Google Drive account, and load data from
-Google Drive. Having Google Drive connected will allow you to save any
-files created back to your Drive folder. This step will require you to
-download the csv files from
-Github
-to your computer, and then make a folder in your Google Drive account
-and upload the csv files to this folder.

-
from google.colab import drive
+

+                    

+                  

+                  
Option 1: Connect your Google Drive account, and load data from
+                    Google Drive. Having Google Drive connected will allow you to save any
+                    files created back to your Drive folder. This step will require you to
+                    download the csv files from
+                    Github
+                    to your computer, and then make a folder in your Google Drive account
+                    and upload the csv files to this folder.
+                  

+                  

+                    

+                      
from google.colab import drive
 drive.mount('/content/drive')
 
 #change dir to data on your google drive
@@ -245,19 +287,25 @@ 
Task 0: Load data and install the pcntoolkitos.chdir('drive/My Drive/name-of-folder-where-you-uploaded-csv-files-from-Github/') #Change this path to match the path to your data in Google Drive
 
 # code by T. Wolfers
-

-

-
Option 2: Import the files directly from Github, and skip adding
-them to Google Drive.

-
!wget -nc https://raw.githubusercontent.com/predictive-clinical-neuroscience/PCNtoolkit-demo/master/tutorials/CPC_2020/data/camcan_demographics.csv
+

+                    

+                  

+                  
Option 2: Import the files directly from Github, and skip adding
+                    them to Google Drive.

+                  

+                    

+                      
!wget -nc https://raw.githubusercontent.com/predictive-clinical-neuroscience/PCNtoolkit-demo/master/tutorials/CPC_2020/data/camcan_demographics.csv
 !wget -nc https://raw.githubusercontent.com/predictive-clinical-neuroscience/PCNtoolkit-demo/master/tutorials/CPC_2020/data/camcan_demographics_nordan.csv
 !wget -nc https://raw.githubusercontent.com/predictive-clinical-neuroscience/PCNtoolkit-demo/master/tutorials/CPC_2020/data/camcan_features.csv
 !wget -nc https://raw.githubusercontent.com/predictive-clinical-neuroscience/PCNtoolkit-demo/master/tutorials/CPC_2020/data/camcan_features_nordan.csv
 
 # code by S. Rutherford
-

-

-
--2022-02-17 15:03:58--  https://raw.githubusercontent.com/predictive-clinical-neuroscience/PCNtoolkit-demo/master/tutorials/CPC_2020/data/camcan_demographics.csv
+

+                    

+                  

+                  

+                    

+                      
--2022-02-17 15:03:58--  https://raw.githubusercontent.com/predictive-clinical-neuroscience/PCNtoolkit-demo/master/tutorials/CPC_2020/data/camcan_demographics.csv
 Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.110.133, 185.199.108.133, 185.199.111.133, ...
 Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.110.133|:443... connected.
 HTTP request sent, awaiting response... 200 OK
@@ -300,23 +348,29 @@ 
Task 0: Load data and install the pcntoolkit
-
TASK 1: Format input data

-
You have four files. The features and demographics file for the
-normsample and two files of the same name for Nordan your test sample.
-As one of your coworkers has done the preporcessing and quality control
-there are more subjects in the demographics file than in the features
-file of the norm sample. Please select the overlap of participants
-between those two files.

-
Question for your understanding:

-
    
-
  1. Why do we have to select the overlap between participants in terms of
-featrues and demographics?
    2. 
-

-
import pandas as pd
+

+                    

+                  

+                

+                

+                  
TASK 1: Format input data

+                  
You have four files. The features and demographics file for the
+                    normsample and two files of the same name for Nordan your test sample.
+                    As one of your coworkers has done the preporcessing and quality control
+                    there are more subjects in the demographics file than in the features
+                    file of the norm sample. Please select the overlap of participants
+                    between those two files.

+                  
Question for your understanding:

+                  
    
+                    
  1. 
+                      
    Why do we have to select the overlap between participants in terms of
+                        featrues and demographics?
    
+                    
    2. 
+                  

+                  

+                    

+                      
import pandas as pd
 
 # read in the files.
 norm_demographics = pd.read_csv('camcan_demographics.csv',
@@ -340,9 +394,12 @@ 
TASK 1: Format input dataprint(norm_demographics_features)
 
 # code by T. Wolfers
-

-

-
             age sex_name  sex  IQ_random
+

+                    

+                  

+                  

+                    

+                      
             age sex_name  sex  IQ_random
 paricipants
 CC110033      24     MALE    1         73
 CC110037      18     MALE    1        103
@@ -386,34 +443,42 @@ 
TASK 1: Format input data
-
TASK 2: Prepare the covariate_normsample and testresponse_normsample file.

-
As mentioned in the introductory presentation those files need a
-specific format and the entries need to be seperated by spaces. Use
-whatever method you know to prepare those files based on the data
-provided in TASK 1. Save those files in .txt format in your drive. Also
-get rid of the column names and participant IDs.

-
Given that we only have limited time in this practical we have to make a
-selection for the features based on your prior knowledge. With the
-information in mind that Nordan does not remember his way home, which
-subfield of the hyppocampus is probably a good target for the
-investigations? Select a maximum of four hyppocampal regions as
-features.

-
NOTE: Normative modeling is a screening tool we just make this selection
-due to time constraints, in reality we build these models on millions of
-putative biomarkers that are not restricted to brain imaging.

-
Qestions for your understanding:

-
    
-
  1. What is the requirement for the features in terms of variable
-properties (e.g. dicotomous or continous)? 3) What is the requirement
-for the covariates in terms of these properties? 4) What are the
-requirements for both together? 5) How does this depent on the
-algorithm used?
    2. 
-

-
# perpare covariate_normsample for sex and age
+

+                    

+                  

+                

+                

+                  
TASK 2: Prepare the covariate_normsample and testresponse_normsample file.

+                  
As mentioned in the introductory presentation those files need a
+                    specific format and the entries need to be seperated by spaces. Use
+                    whatever method you know to prepare those files based on the data
+                    provided in TASK 1. Save those files in .txt format in your drive. Also
+                    get rid of the column names and participant IDs.

+                  
Given that we only have limited time in this practical we have to make a
+                    selection for the features based on your prior knowledge. With the
+                    information in mind that Nordan does not remember his way home, which
+                    subfield of the hyppocampus is probably a good target for the
+                    investigations? Select a maximum of four hyppocampal regions as
+                    features.

+                  
NOTE: Normative modeling is a screening tool we just make this selection
+                    due to time constraints, in reality we build these models on millions of
+                    putative biomarkers that are not restricted to brain imaging.

+                  
Qestions for your understanding:

+                  
    
+                    
  1. 
+                      
    What is the requirement for the features in terms of variable
+                        properties (e.g. dicotomous or continous)? 3) What is the requirement
+                        for the covariates in terms of these properties? 4) What are the
+                        requirements for both together? 5) How does this depent on the
+                        algorithm used?
    
+                    
    2. 
+                  

+                  

+                    

+                      
# perpare covariate_normsample for sex and age
 covariate_normsample = norm_demographics_features[['sex',
                                                    'age']]
 
@@ -434,24 +499,30 @@ 
TASK 2: Prepare the covariate_normsample and testresponse_n
                            index = False)
 
 # code by T. Wolfers
-

-

-

-

-
TASK 3: Estimate normative model

-
Once you have prepared and saved all the necessary files. Look at the
-pcntoolkit for running normative modeling. Select an appropritate method
-set up the toolkit and run your analyses using 2-fold cross validation
-in the normsample. Change the output suffix from estimate to ’_2fold’.

-
HINT: You primarily need the estimate function.

-
SUGGESTION: While this process is running you can go to the next TASK 4,
-you will have no doubt when it is correctly running.

-
Question for your understaning:

-
    
-
  1. What does cvfolds mean and why do we use it? 7) What is the output of
-the estimate function and what does it mean?
    2. 
-

-
import pcntoolkit as pcn
+

+                    

+                  

+                

+                

+                  
TASK 3: Estimate normative model

+                  
Once you have prepared and saved all the necessary files. Look at the
+                    pcntoolkit for running normative modeling. Select an appropritate method
+                    set up the toolkit and run your analyses using 2-fold cross validation
+                    in the normsample. Change the output suffix from estimate to ’_2fold’.

+                  
HINT: You primarily need the estimate function.

+                  
SUGGESTION: While this process is running you can go to the next TASK 4,
+                    you will have no doubt when it is correctly running.

+                  
Question for your understaning:

+                  
    
+                    
  1. 
+                      
    What does cvfolds mean and why do we use it? 7) What is the output of
+                        the estimate function and what does it mean?
    
+                    
    2. 
+                  

+                  

+                    

+                      
import pcntoolkit as pcn
 
 # run normative modeling using 2-fold cross-validation
 
@@ -462,9 +533,12 @@ 
TASK 3: Estimate normative modeloutputsuffix = '_2fold')
 
 # code by T. Wolfers
-

-

-
Processing data in features_normsample.txt
+

+                    

+                  

+                  

+                    

+                      
Processing data in features_normsample.txt
 Estimating model  1 of 4
 Warning: Estimation of posterior distribution failed
 Warning: Estimation of posterior distribution failed
@@ -554,23 +628,30 @@ 
TASK 3: Estimate normative model

-

-

-

-
TASK 4: Estimate the forward model of the normative model

-
In order to visulize the normative trajectories you first need to run
-the forward model. To this end you need to set up an appropriate
-covariate_forwardmodel file that covers the age range appropriately for
-both sexes. Save this file as .txt . Then you can input the files you
-made in TASK 1 as well as the file you made now and run the forward
-model using the appropriate specifications.

-
Question for your understaning:

-
    
-
  1. What is yhat and ys2? 9) Why does the output of the forward model
-does not inlcude the Z-scores?
    2. 
-

-
# create covariate_forwardmodel.txt file
+

+                    

+                  

+                

+                

+                  
TASK 4: Estimate the forward model of the normative model

+                  
In order to visulize the normative trajectories you first need to run
+                    the forward model. To this end you need to set up an appropriate
+                    covariate_forwardmodel file that covers the age range appropriately for
+                    both sexes. Save this file as .txt . Then you can input the files you
+                    made in TASK 1 as well as the file you made now and run the forward
+                    model using the appropriate specifications.

+                  
Question for your understaning:

+                  
    
+                    
  1. 
+                      
    What is yhat and ys2? 9) Why does the output of the forward model
+                        does not inlcude the Z-scores?
    
+                    
    2. 
+                  

+                  

+                    

+                      
# create covariate_forwardmodel.txt file
 covariate_forwardmodel = {'sex': [0, 0, 0, 0, 0, 0, 0,
                                   1, 1, 1, 1, 1, 1, 1],
                           'age': [20, 30, 40, 50, 60, 70, 80,
@@ -591,9 +672,12 @@ 
TASK 4: Estimate the forward model of the normative modeloutputsuffix = '_forward')
 
 # code by T. Wolfers
-

-

-
Processing data in features_normsample.txt
+

+                    

+                  

+                  

+                    

+                      
Processing data in features_normsample.txt
 Estimating model  1 of 4
 Warning: Estimation of posterior distribution failed
 Warning: Estimation of posterior distribution failed
@@ -652,23 +736,29 @@ 
TASK 4: Estimate the forward model of the normative model

-

-

-

-
TASK 5: Visualize forward model

-
Visualize the forward model of the normative model similar to the figure
-below.

-

-1-s2.0-S245190221830329X-gr2.jpg
-

-
1-s2.0-S245190221830329X-gr2.jpg

-

-

-
HINT: First create a function that calculates the confidence intervals
-and then plot yhat, y2 of the forward model. Finally, plot the data of
-individual participants.

-
import numpy as np
+

+                    

+                  

+                

+                

+                  
TASK 5: Visualize forward model

+                  
Visualize the forward model of the normative model similar to the figure
+                    below.

+                  

+                    1-s2.0-S245190221830329X-gr2.jpg
+                    

+                      
1-s2.0-S245190221830329X-gr2.jpg

+                    

+                  

+                  
HINT: First create a function that calculates the confidence intervals
+                    and then plot yhat, y2 of the forward model. Finally, plot the data of
+                    individual participants.

+                  

+                    

+                      
import numpy as np
 import matplotlib.pyplot as plt
 
 # confidence interval calculation at x_forward
@@ -726,20 +816,34 @@ 
TASK 5: Visualize forward modelplt.close()
 
 # code by M. Zabihi
-

-

-../_images/normative_modelling_walkthrough_22_0.png
-../_images/normative_modelling_walkthrough_22_1.png
-../_images/normative_modelling_walkthrough_22_2.png
-../_images/normative_modelling_walkthrough_22_3.png
-../_images/normative_modelling_walkthrough_22_4.png
-../_images/normative_modelling_walkthrough_22_5.png
-../_images/normative_modelling_walkthrough_22_6.png
-../_images/normative_modelling_walkthrough_22_7.png
-

-

-
TASK 6: Apply the normative model to Nordan’s data and the dementia patients.

-
# read in Nordan's as well as the patient's demographics and features
+

+                    

+                  

+                  ../_images/normative_modelling_walkthrough_22_0.png
+                  ../_images/normative_modelling_walkthrough_22_1.png
+                  ../_images/normative_modelling_walkthrough_22_2.png
+                  ../_images/normative_modelling_walkthrough_22_3.png
+                  ../_images/normative_modelling_walkthrough_22_4.png
+                  ../_images/normative_modelling_walkthrough_22_5.png
+                  ../_images/normative_modelling_walkthrough_22_6.png
+                  ../_images/normative_modelling_walkthrough_22_7.png
+                

+                

+                  
TASK 6: Apply the normative model to Nordan’s data and the dementia patients.

+                  

+                    

+                      
# read in Nordan's as well as the patient's demographics and features
 demographics_nordan = pd.read_csv('camcan_demographics_nordan.csv',
                                        sep= ",",
                                        index_col = 0)
@@ -776,9 +880,12 @@ 
TASK 6: Apply the normative model to Nordan’s data and th
                        outputsuffix = '_nordan')
 
 # code by T. Wolfers
-

-

-
Processing data in features_normsample.txt
+

+                    

+                  

+                  

+                    

+                      
Processing data in features_normsample.txt
 Estimating model  1 of 4
 Warning: Estimation of posterior distribution failed
 Warning: Estimation of posterior distribution failed
@@ -838,32 +945,45 @@ 
TASK 6: Apply the normative model to Nordan’s data and th
          Gradient evaluations: 104
 Evaluating the model ...
 Writing outputs ...
-

-

-

-

-
TASK 7: In which hyppocampal subfield(s) does Nordan deviate extremely?

-
No coding necessary just create a presentation which includes
-recommendations to Nordan and his family. Use i) |Z| > 3.6 ii) |Z| >
-1.96 as definitions for extreme normative deviations.

-

-

-
TASK 8 (OPTIONAL): Implement a function that calculates percentage change.

-
Percentage change = \frac{x1 - x2}{|x2|}*100

-
# function that calculates percentage change
+

+                    

+                  

+                

+                

+                  
TASK 7: In which hyppocampal subfield(s) does Nordan deviate extremely?

+                  
No coding necessary just create a presentation which includes
+                    recommendations to Nordan and his family. Use i) |Z| > 3.6 ii) |Z| >
+                    1.96 as definitions for extreme normative deviations.

+                

+                

+                  
TASK 8 (OPTIONAL): Implement a function that calculates percentage change.

+                  
Percentage change = \frac{x1 - x2}{|x2|}*100

+                  

+                    

+                      
# function that calculates percentage change
 def calculate_percentage_change(x1, x2):
   percentage_change = ((x1 - x2) / abs(x2)) * 100
   return percentage_change
 
 # code by T. Wolfers
-

-

-

-

-
TASK 9 (OPTIONAL): Visualize percent change

-
Plot the prercentage change in Yhat of the forward model in reference to
-age 20. Do that for both sexes seperately.

-
import matplotlib.pyplot as plt
+

+                    

+                  

+                

+                

+                  
TASK 9 (OPTIONAL): Visualize percent change

+                  
Plot the prercentage change in Yhat of the forward model in reference to
+                    age 20. Do that for both sexes seperately.

+                  

+                    

+                      
import matplotlib.pyplot as plt
 
 forward_yhat = pd.read_csv('yhat_forward.txt', sep = ' ', header=None)
 
@@ -907,42 +1027,52 @@ 
TASK 9 (OPTIONAL): Visualize percent changeplt.ylim(-20, 2)
 
 # code by T. Wolfers
-

-

-
(-20.0, 2.0)
-

-

-../_images/normative_modelling_walkthrough_32_1.png
-

-

-
-
-           

-           
+

+                    

+                  

+                  

+                    

+                      
(-20.0, 2.0)
+

+                    

+                  

+                  ../_images/normative_modelling_walkthrough_32_1.png
+                

+              

+
+
+            

+
           

           
+            
+
+
+            

+
+            

+              

+                © Copyright 2020, Andre F. Marquand
+
+              

+            

+            Built with Sphinx using a theme provided by Read the Docs.
+
+          
 
         

       

@@ -950,19 +1080,20 @@ 
TASK 9 (OPTIONAL): Visualize percent change
-      jQuery(function () {
-          SphinxRtdTheme.Navigation.enable(true);
-      });
+    jQuery(function () {
+      SphinxRtdTheme.Navigation.enable(true);
+    });
   
 
-  
-  
-    
-   
+
+
+
+
 
 
+
 
\ No newline at end of file
diff --git a/doc/requirements.txt b/doc/requirements.txt
index 54e34c46..9f871205 100644
--- a/doc/requirements.txt
+++ b/doc/requirements.txt
@@ -12,4 +12,4 @@ numpy>=1.19.5
 scipy>=1.3.2
 pandas>=0.25.3
 torch>=1.1.0 
-pymc3>=3.11.2
+pymc>=4
diff --git a/doc/source/pages/HBR_NormativeModel_FCONdata_Tutorial.rst b/doc/source/pages/HBR_NormativeModel_FCONdata_Tutorial.rst
index 4db24f06..d4933e38 100644
--- a/doc/source/pages/HBR_NormativeModel_FCONdata_Tutorial.rst
+++ b/doc/source/pages/HBR_NormativeModel_FCONdata_Tutorial.rst
@@ -24,7 +24,6 @@ Step 0: Install necessary libraries & grab data files
 
 .. code:: ipython3
 
-    ! pip uninstall -y Theano-PyMC  # conflicts with Theano on some environments
     ! pip install pcntoolkit==0.26
 
 
diff --git a/doc/source/pages/normative_modelling_walkthrough.rst b/doc/source/pages/normative_modelling_walkthrough.rst
index 06eb5017..f63ec8fa 100644
--- a/doc/source/pages/normative_modelling_walkthrough.rst
+++ b/doc/source/pages/normative_modelling_walkthrough.rst
@@ -46,7 +46,6 @@ bad if you cannot complete everything during the tutorial.
 .. code:: ipython3
 
     #install normative modeling
-    ! pip uninstall -y Theano-PyMC  # conflicts with Theano on some environments
     ! pip install pcntoolkit==0.26
 
 
diff --git a/pcntoolkit/model/SHASH.py b/pcntoolkit/model/SHASH.py
index 5ab91e51..1312020e 100644
--- a/pcntoolkit/model/SHASH.py
+++ b/pcntoolkit/model/SHASH.py
@@ -1,36 +1,70 @@
-import theano.tensor
-from pymc3.distributions import Continuous, draw_values, generate_samples
-import theano.tensor as tt
+from typing import Union, List, Optional
+import pymc as pm
+from pymc import floatX
+from pymc.distributions import Continuous
+
+import pytensor as pt
+import pytensor.tensor as ptt
+from pytensor.graph.op import Op
+from pytensor.graph import Apply
+from pytensor.gradient import grad_not_implemented
+from pytensor.tensor.random.basic import normal
+from pytensor.tensor.random.op import RandomVariable
+
+
 import numpy as np
-from pymc3.distributions.dist_math import bound
 import scipy.special as spp
-from theano import as_op
-from theano.gof.fg import NullType
-from theano.gof.op import Op
-from theano.gof.graph import Apply
-from theano.gradient import grad_not_implemented
+import matplotlib.pyplot as plt
+
 
 """
 @author: Stijn de Boer (AuguB)
-
 See: Jones et al. (2009), Sinh-Arcsinh distributions.
 """
 
 
+def numpy_P(q):
+    """
+    The P function as given in Jones et al.
+    :param q:
+    :return:
+    """
+    frac = np.exp(1.0 / 4.0) / np.power(8.0 * np.pi, 1.0 / 2.0)
+    K1 = numpy_K((q + 1) / 2, 1.0 / 4.0)
+    K2 = numpy_K((q - 1) / 2, 1.0 / 4.0)
+    a = (K1 + K2) * frac
+    return a
+
+
+def numpy_K(p, x):
+    """
+    Computes the values of spp.kv(p,x) for only the unique values of p
+    """
+
+    ps, idxs = np.unique(p, return_inverse=True)
+    return spp.kv(ps, x)[idxs].reshape(p.shape)
+
+
 class K(Op):
     """
-    Modified Bessel function of the second kind, theano implementation
+    Modified Bessel function of the second kind, pytensor implementation
     """
+
+    __props__ = ()
+
     def make_node(self, p, x):
-        p = theano.tensor.as_tensor_variable(p, 'floatX')
-        x = theano.tensor.as_tensor_variable(x, 'floatX')
-        return Apply(self, [p,x], [p.type()])
+        p = pt.tensor.as_tensor_variable(p)
+        x = pt.tensor.as_tensor_variable(x)
+        return Apply(self, [p, x], [p.type()])
 
-    def perform(self, node, inputs, output_storage, params=None):
+    def perform(self, node, inputs_storage, output_storage):
         # Doing this on the unique values avoids doing A LOT OF double work, apparently scipy doesn't do this by itself
-        unique_inputs, inverse_indices = np.unique(inputs[0], return_inverse=True)
-        unique_outputs = spp.kv(unique_inputs, inputs[1])
-        outputs = unique_outputs[inverse_indices].reshape(inputs[0].shape)
+
+        unique_inputs, inverse_indices = np.unique(
+            inputs_storage[0], return_inverse=True
+        )
+        unique_outputs = spp.kv(unique_inputs, inputs_storage[1])
+        outputs = unique_outputs[inverse_indices].reshape(inputs_storage[0].shape)
         output_storage[0][0] = outputs
 
     def grad(self, inputs, output_grads):
@@ -38,234 +72,274 @@ def grad(self, inputs, output_grads):
         dp = 1e-10
         p = inputs[0]
         x = inputs[1]
-        grad = (self(p+dp,x) - self(p, x))/dp
-        return [output_grads[0]*grad, grad_not_implemented(0,1,2,3)]
+        grad = (self(p + dp, x) - self(p, x)) / dp
+        return [output_grads[0] * grad, grad_not_implemented(0, 1, 2, 3)]
+
+
+def S(x, epsilon, delta):
+    """
+    :param epsilon:
+    :param delta:
+    :param x:
+    :return: The sinharcsinh transformation of x
+    """
+    return np.sinh(np.arcsinh(x) * delta - epsilon)
+
+
+def S_inv(x, epsilon, delta):
+    return np.sinh((np.arcsinh(x) + epsilon) / delta)
+
+
+def C(x, epsilon, delta):
+    """
+    :param epsilon:
+    :param delta:
+    :param x:
+    :return: the cosharcsinh transformation of x
+    Be aware that this is sqrt(1+S(x)^2), so you may save some compute if you can re-use the result from S.
+    """
+    return np.cosh(np.arcsinh(x) * delta - epsilon)
+
+
+def P(q):
+    """
+    The P function as given in Jones et al.
+    :param q:
+    :return:
+    """
+    frac = np.exp(1.0 / 4.0) / np.power(8.0 * np.pi, 1.0 / 2.0)
+    K1 = K()((q + 1) / 2, 1.0 / 4.0)
+    K2 = K()((q - 1) / 2, 1.0 / 4.0)
+    a = (K1 + K2) * frac
+    return a
+
+
+def m(epsilon, delta, r):
+    """
+    :param epsilon:
+    :param delta:
+    :param r:
+    :return:  The r'th uncentered moment of the SHASH distribution parameterized by epsilon and delta. Given by Jones et al.
+    The first four moments are given in closed form.
+    """
+    if r == 1:
+        return np.sinh(epsilon / delta) * P(1 / delta)
+    elif r == 2:
+        return (np.cosh(2 * epsilon / delta) * P(2 / delta) - 1) / 2
+    elif r == 3:
+        return (
+            np.sinh(3 * epsilon / delta) * P(3 / delta)
+            - 3 * np.sinh(epsilon / delta) * P(1 / delta)
+        ) / 4
+    elif r == 4:
+        return (
+            np.cosh(4 * epsilon / delta) * P(4 / delta)
+            - 4 * np.cosh(2 * epsilon / delta) * P(2 / delta)
+            + 3
+        ) / 8
+    # else:
+    #     frac1 = ptt.as_tensor_variable(1 / pm.power(2, r))
+    #     acc = ptt.as_tensor_variable(0)
+    #     for i in range(r + 1):
+    #         combs = spp.comb(r, i)
+    #         flip = pm.power(-1, i)
+    #         ex = np.exp((r - 2 * i) * epsilon / delta)
+    #         p = P((r - 2 * i) / delta)
+    #         acc += combs * flip * ex * p
+    #     return frac1 * acc
+
+
+class SHASH(RandomVariable):
+    name = "shash"
+    ndim_supp = 0
+    ndims_params = [0, 0]
+    dtype = "floatX"
+    _print_name = ("SHASH", "\\operatorname{SHASH}")
+
+    @classmethod
+    def rng_fn(cls, rng, epsilon, delta, size=None) -> np.ndarray:
+        return np.sinh(
+            (np.arcsinh(rng.normal(loc=0, scale=1, size=size)) + epsilon) / delta
+        )
+
+
+shash = SHASH()
+
 
 class SHASH(Continuous):
+    rv_op = shash
     """
     SHASH described by Jones et al., based on a standard normal
     All SHASH subclasses inherit from this
     """
-    def __init__(self, epsilon, delta, **kwargs):
-        super().__init__(**kwargs)
-        self.epsilon = tt.as_tensor_variable(epsilon)
-        self.delta = tt.as_tensor_variable(delta)
-        self.K = K()
 
+    @classmethod
+    def dist(cls, epsilon, delta, **kwargs):
+        epsilon = ptt.as_tensor_variable(floatX(epsilon))
+        delta = ptt.as_tensor_variable(floatX(delta))
+        return super().dist([epsilon, delta], **kwargs)
+
+    def logp(value, epsilon, delta):
+        this_S = S(value, epsilon, delta)
+        this_S_sqr = ptt.sqr(this_S)
+        this_C_sqr = 1 + this_S_sqr
+        frac1 = -ptt.log(ptt.constant(2 * np.pi)) / 2
+        frac2 = (
+            ptt.log(delta) + ptt.log(this_C_sqr) / 2 - ptt.log(1 + ptt.sqr(value)) / 2
+        )
+        exp = -this_S_sqr / 2
+        return frac1 + frac2 + exp
 
-    def random(self, point=None, size=None):
-        epsilon, delta = draw_values([self.epsilon, self.delta],
-                                     point=point, size=size)
 
-        def _random(epsilon, delta, size=None):
-            samples_transformed = np.sinh((np.arcsinh(np.random.randn(*size)) + epsilon) / delta)
-            return samples_transformed
+class SHASHoRV(RandomVariable):
+    name = "shasho"
+    ndim_supp = 0
+    ndims_params = [0, 0, 0, 0]
+    dtype = "floatX"
+    _print_name = ("SHASHo", "\\operatorname{SHASHo}")
 
-        return generate_samples(_random, epsilon=epsilon, delta=delta, dist_shape=self.shape, size=size)
+    @classmethod
+    def rng_fn(cls, rng, mu, sigma, epsilon, delta, size=None) -> np.ndarray:
+        s = rng.normal(size=size)
+        return np.sinh((np.arcsinh(s) + epsilon) / delta) * sigma + mu
 
-    def logp(self, value):
-        epsilon = self.epsilon
-        delta = self.delta + tt.np.finfo(np.float32).eps
 
-        this_S = self.S(value)
-        this_S_sqr = tt.sqr(this_S)
-        this_C_sqr = 1+this_S_sqr
-        frac1 = -tt.log(tt.constant(2 * tt.np.pi))/2
-        frac2 = tt.log(delta) + tt.log(this_C_sqr)/2 - tt.log(1 + tt.sqr(value)) / 2
-        exp = -this_S_sqr / 2
+shasho = SHASHoRV()
+
 
-        return bound(frac1 + frac2 + exp, delta > 0)
-
-    def S(self, x):
-        """
-
-        :param epsilon:
-        :param delta:
-        :param x:
-        :return: The sinharcsinh transformation of x
-        """
-        return tt.sinh(tt.arcsinh(x) * self.delta - self.epsilon)
-
-    def S_inv(self, x):
-        return tt.sinh((tt.arcsinh(x) + self.epsilon) / self.delta)
-
-    def C(self, x):
-        """
-        :param epsilon:
-        :param delta:
-        :param x:
-        :return: the cosharcsinh transformation of x
-        Be aware that this is sqrt(1+S(x)^2), so you may save some compute if you can re-use the result from S.
-        """
-        return tt.cosh(tt.arcsinh(x) * self.delta - self.epsilon)
-
-    def P(self, q):
-        """
-        The P function as given in Jones et al.
-        :param q:
-        :return:
-        """
-        frac = np.exp(1 / 4) / np.power(8 * np.pi, 1 / 2)
-        K1 = self.K((q+1)/2,1/4)
-        K2 = self.K((q-1)/2,1/4)
-        a = (K1 + K2) * frac
-        return a
-
-    def m(self, r):
-        """
-        :param epsilon:
-        :param delta:
-        :param r:
-        :return:  The r'th uncentered moment of the SHASH distribution parameterized by epsilon and delta. Given by Jones et al.
-        """
-        frac1 = tt.as_tensor_variable(1 / np.power(2, r))
-        acc = tt.as_tensor_variable(0)
-        for i in range(r + 1):
-            combs = spp.comb(r, i)
-            flip = np.power(-1, i)
-            ex = np.exp((r - 2 * i) * self.epsilon / self.delta)
-            # This is the reason we can not sample delta/kurtosis using NUTS; the gradient of P is unknown to pymc3
-            # TODO write a class that inherits theano.Op and do the gradient in there :)
-            p = self.P((r - 2 * i) / self.delta)
-            acc += combs * flip * ex * p
-        return frac1 * acc
-
-class SHASHo(SHASH):
+class SHASHo(Continuous):
+    rv_op = shasho
     """
     This is the shash where the location and scale parameters have simply been applied as an linear transformation
     directly on the original shash.
     """
 
-    def __init__(self, mu, sigma, epsilon, delta, **kwargs):
-        super().__init__(epsilon, delta, **kwargs)
-        self.mu = tt.as_tensor_variable(mu)
-        self.sigma = tt.as_tensor_variable(sigma)
-
-    def random(self, point=None, size=None):
-        mu, sigma, epsilon, delta = draw_values([self.mu, self.sigma, self.epsilon, self.delta],
-                                                point=point, size=size)
-
-        def _random(mu, sigma, epsilon, delta, size=None):
-            samples_transformed = np.sinh((np.arcsinh(np.random.randn(*size)) + epsilon) / delta) * sigma + mu
-            return samples_transformed
-
-        return generate_samples(_random, mu=mu, sigma=sigma, epsilon=epsilon, delta=delta,
-                                dist_shape=self.shape,
-                                size=size)
-
-    def logp(self, value):
-        mu = self.mu
-        sigma = self.sigma + tt.np.finfo(np.float32).eps
-        epsilon = self.epsilon
-        delta = self.delta + tt.np.finfo(np.float32).eps
-
-        value_transformed = (value - mu) / sigma
-
-        this_S = self.S( value_transformed)
-        this_S_sqr = tt.sqr(this_S)
-        this_C_sqr = 1+this_S_sqr
-        frac1 = -tt.log(tt.constant(2 * tt.np.pi))/2
-        frac2 = tt.log(delta) + tt.log(this_C_sqr)/2 - tt.log(
-            1 + tt.sqr(value_transformed)) / 2
+    @classmethod
+    def dist(cls, mu, sigma, epsilon, delta, **kwargs):
+        mu = ptt.as_tensor_variable(floatX(mu))
+        sigma = ptt.as_tensor_variable(floatX(sigma))
+        epsilon = ptt.as_tensor_variable(floatX(epsilon))
+        delta = ptt.as_tensor_variable(floatX(delta))
+        return super().dist([mu, sigma, epsilon, delta], **kwargs)
+
+    def logp(value, mu, sigma, epsilon, delta):
+        remapped_value = (value - mu) / sigma
+        this_S = S(remapped_value, epsilon, delta)
+        this_S_sqr = ptt.sqr(this_S)
+        this_C_sqr = 1 + this_S_sqr
+        frac1 = -ptt.log(ptt.constant(2 * np.pi)) / 2
+        frac2 = (
+            ptt.log(delta)
+            + ptt.log(this_C_sqr) / 2
+            - ptt.log(1 + ptt.sqr(remapped_value)) / 2
+        )
         exp = -this_S_sqr / 2
-        change_of_variable = -tt.log(sigma)
+        return frac1 + frac2 + exp - ptt.log(sigma)
 
-        return bound(frac1 + frac2 + exp + change_of_variable, sigma > 0, delta > 0)
 
+class SHASHo2RV(RandomVariable):
+    name = "shasho2"
+    ndim_supp = 0
+    ndims_params = [0, 0, 0, 0]
+    dtype = "floatX"
+    _print_name = ("SHASHo2", "\\operatorname{SHASHo2}")
 
-class SHASHo2(SHASH):
-    """
-    This is the shash where we apply the reparameterization provided in section 4.3 in Jones et al.
-    """
-
-    def __init__(self, mu, sigma, epsilon, delta, **kwargs):
-        super().__init__(epsilon, delta, **kwargs)
-        self.mu = tt.as_tensor_variable(mu)
-        self.sigma = tt.as_tensor_variable(sigma)
-
-    def random(self, point=None, size=None):
-        mu, sigma, epsilon, delta = draw_values(
-            [self.mu, self.sigma, self.epsilon, self.delta],
-            point=point, size=size)
+    @classmethod
+    def rng_fn(cls, rng, mu, sigma, epsilon, delta, size=None) -> np.ndarray:
+        s = rng.normal(size=size)
         sigma_d = sigma / delta
+        return np.sinh((np.arcsinh(s) + epsilon) / delta) * sigma_d + mu
 
-        def _random(mu, sigma, epsilon, delta, size=None):
-            samples_transformed = np.sinh(
-                (np.arcsinh(np.random.randn(*size)) + epsilon) / delta) * sigma_d + mu
-            return samples_transformed
 
-        return generate_samples(_random, mu=mu, sigma=sigma_d, epsilon=epsilon, delta=delta,
-                                dist_shape=self.shape,
-                                size=size)
+shasho2 = SHASHo2RV()
 
-    def logp(self, value):
-        mu = self.mu
-        sigma = self.sigma + tt.np.finfo(np.float32).eps
-        epsilon = self.epsilon
-        delta = self.delta + tt.np.finfo(np.float32).eps
-        sigma_d = sigma / delta
 
+class SHASHo2(Continuous):
+    rv_op = shasho2
+    """
+    This is the shash where we apply the reparameterization provided in section 4.3 in Jones et al.
+    """
 
-        # Here a double change of variables is applied
-        value_transformed = ((value - mu) / sigma_d)
+    @classmethod
+    def dist(cls, mu, sigma, epsilon, delta, **kwargs):
+        mu = ptt.as_tensor_variable(floatX(mu))
+        sigma = ptt.as_tensor_variable(floatX(sigma))
+        epsilon = ptt.as_tensor_variable(floatX(epsilon))
+        delta = ptt.as_tensor_variable(floatX(delta))
+        return super().dist([mu, sigma, epsilon, delta], **kwargs)
 
-        this_S = self.S(value_transformed)
-        this_S_sqr = tt.sqr(this_S)
-        this_C = tt.sqrt(1+this_S_sqr)
-        frac1 = -tt.log(tt.sqrt(tt.constant(2 * tt.np.pi)))
-        frac2 = tt.log(delta) + tt.log(this_C) - tt.log(
-            1 + tt.sqr(value_transformed)) / 2
+    def logp(value, mu, sigma, epsilon, delta):
+        sigma_d = sigma / delta
+        remapped_value = (value - mu) / sigma_d
+        this_S = S(remapped_value, epsilon, delta)
+        this_S_sqr = ptt.sqr(this_S)
+        this_C_sqr = 1 + this_S_sqr
+        frac1 = -ptt.log(ptt.constant(2 * np.pi)) / 2
+        frac2 = (
+            ptt.log(delta)
+            + ptt.log(this_C_sqr) / 2
+            - ptt.log(1 + ptt.sqr(remapped_value)) / 2
+        )
         exp = -this_S_sqr / 2
-        change_of_variable = -tt.log(sigma_d)
-
-        # the change of variables is accounted for in the density by division and multiplication (adding and subtracting for logp)
-        return bound(frac1 + frac2 + exp + change_of_variable, delta > 0, sigma > 0)
-
-class SHASHb(SHASH):
+        return frac1 + frac2 + exp - ptt.log(sigma_d)
+
+
+class SHASHbRV(RandomVariable):
+    name = "shashb"
+    ndim_supp = 0
+    ndims_params = [0, 0, 0, 0]
+    dtype = "floatX"
+    _print_name = ("SHASHo2", "\\operatorname{SHASHo2}")
+
+    @classmethod
+    def rng_fn(
+        cls,
+        rng: np.random.RandomState,
+        mu: Union[np.ndarray, float],
+        sigma: Union[np.ndarray, float],
+        epsilon: Union[np.ndarray, float],
+        delta: Union[np.ndarray, float],
+        size: Optional[Union[List[int], int]],
+    ) -> np.ndarray:
+        s = rng.normal(size=size)
+        mean = np.sinh(epsilon / delta) * numpy_P(1 / delta)
+        var = ((np.cosh(2 * epsilon / delta) * numpy_P(2 / delta) - 1) / 2) - mean**2
+        out = (
+            (np.sinh((np.arcsinh(s) + epsilon) / delta) - mean) / np.sqrt(var)
+        ) * sigma + mu
+        return out
+
+
+shashb = SHASHbRV()
+
+
+class SHASHb(Continuous):
+    rv_op = shashb
     """
     This is the shash where the location and scale parameters been applied as an linear transformation on the shash
     distribution which was corrected for mean and variance.
     """
 
-    def __init__(self, mu, sigma, epsilon, delta, **kwargs):
-        super().__init__(epsilon, delta, **kwargs)
-        self.mu = tt.as_tensor_variable(mu)
-        self.sigma = tt.as_tensor_variable(sigma)
-
-    def random(self, point=None, size=None):
-        mu, sigma, epsilon, delta = draw_values(
-            [self.mu, self.sigma, self.epsilon, self.delta],
-            point=point, size=size)
-        mean = (tt.sinh(epsilon/delta)*self.P(1/delta)).eval()
-        var =  ((tt.cosh(2*epsilon/delta)*self.P(2/delta)-1)/2).eval() - mean**2
-
-        def _random(mean, var, mu, sigma, epsilon, delta, size=None):
-            samples_transformed = ((np.sinh(
-                (np.arcsinh(np.random.randn(*size)) + epsilon) / delta) - mean) / np.sqrt(var)) * sigma + mu
-            return samples_transformed
-
-        return generate_samples(_random, mean=mean, var=var, mu=mu, sigma=sigma, epsilon=epsilon, delta=delta,
-                                dist_shape=self.shape,
-                                size=size)
-
-    def logp(self, value):
-        mu = self.mu
-        sigma = self.sigma + tt.np.finfo(np.float32).eps
-        epsilon = self.epsilon
-        delta = self.delta + tt.np.finfo(np.float32).eps
-        mean = tt.sinh(epsilon/delta)*self.P(1/delta)
-        var = (tt.cosh(2*epsilon/delta)*self.P(2/delta)-1)/2 - tt.sqr(mean)
-
-        # Here a double change of variables is applied
-        value_transformed = ((value - mu) / sigma) * tt.sqrt(var) + mean
-
-        this_S = self.S(value_transformed)
-        this_S_sqr = tt.sqr(this_S)
-        this_C_sqr = 1+this_S_sqr
-        frac1 = -tt.log(tt.constant(2 * tt.np.pi))/2
-        frac2 = tt.log(delta) + tt.log(this_C_sqr)/2 - tt.log(1 + tt.sqr(value_transformed)) / 2
+    @classmethod
+    def dist(cls, mu, sigma, epsilon, delta, **kwargs):
+        mu = ptt.as_tensor_variable(floatX(mu))
+        sigma = ptt.as_tensor_variable(floatX(sigma))
+        epsilon = ptt.as_tensor_variable(floatX(epsilon))
+        delta = ptt.as_tensor_variable(floatX(delta))
+        return super().dist([mu, sigma, epsilon, delta], **kwargs)
+
+    def logp(value, mu, sigma, epsilon, delta):
+        mean = m(epsilon, delta, 1)
+        var = m(epsilon, delta, 2)
+        remapped_value = ((value - mu) / sigma) * np.sqrt(var) + mean
+        this_S = S(remapped_value, epsilon, delta)
+        this_S_sqr = np.square(this_S)
+        this_C_sqr = 1 + this_S_sqr
+        frac1 = -np.log(2 * np.pi) / 2
+        frac2 = (
+            np.log(delta)
+            + np.log(this_C_sqr) / 2
+            - np.log(1 + np.square(remapped_value)) / 2
+        )
         exp = -this_S_sqr / 2
-        change_of_variable = tt.log(var)/2 - tt.log(sigma)
-
-        # the change of variables is accounted for in the density by division and multiplication (addition and subtraction in the log domain)
-        return bound(frac1 + frac2 + exp + change_of_variable, delta > 0, sigma > 0, var > 0)
+        return frac1 + frac2 + exp + np.log(var) / 2 - np.log(sigma)
diff --git a/pcntoolkit/model/hbr.py b/pcntoolkit/model/hbr.py
index 98c71037..9a608d52 100644
--- a/pcntoolkit/model/hbr.py
+++ b/pcntoolkit/model/hbr.py
@@ -9,45 +9,49 @@
 
 from __future__ import print_function
 from __future__ import division
+from collections import OrderedDict
+
 from ast import Param
 from tkinter.font import names
 
 import numpy as np
-import pymc3 as pm
-import theano
+import pymc as pm
+import pytensor
+import arviz as az
+import xarray
+
 from itertools import product
 from functools import reduce
 
-import theano
-from pymc3 import Metropolis, NUTS, Slice, HamiltonianMC
+from pymc import Metropolis, NUTS, Slice, HamiltonianMC
 from scipy import stats
 import bspline
 from bspline import splinelab
 
-from model.SHASH import SHASHo2, SHASHb, SHASHo
 from util.utils import create_poly_basis
 from util.utils import expand_all
 from pcntoolkit.util.utils import cartesian_product
+from pcntoolkit.model.SHASH import *
 
-from theano import printing, function
 
 def bspline_fit(X, order, nknots):
     feature_num = X.shape[1]
     bsp_basis = []
 
     for i in range(feature_num):
-        minx = np.min(X[:,i])
-        maxx = np.max(X[:,i])
-        delta = maxx-minx
+        minx = np.min(X[:, i])
+        maxx = np.max(X[:, i])
+        delta = maxx - minx
         # Expand range by 20% (10% on both sides)
-        splinemin = minx-0.1*delta
-        splinemax = maxx+0.1*delta
+        splinemin = minx - 0.1 * delta
+        splinemax = maxx + 0.1 * delta
         knots = np.linspace(splinemin, splinemax, nknots)
         k = splinelab.augknt(knots, order)
         bsp_basis.append(bspline.Bspline(k, order))
 
     return bsp_basis
 
+
 def bspline_transform(X, bsp_basis):
     if type(bsp_basis) != list:
         temp = []
@@ -62,8 +66,9 @@ def bspline_transform(X, bsp_basis):
 
     return X_transformed
 
+
 def create_poly_basis(X, order):
-    """ compute a polynomial basis expansion of the specified order"""
+    """compute a polynomial basis expansion of the specified order"""
 
     if len(X.shape) == 1:
         X = X[:, np.newaxis]
@@ -71,19 +76,13 @@ def create_poly_basis(X, order):
     Phi = np.zeros((X.shape[0], D * order))
     colid = np.arange(0, D)
     for d in range(1, order + 1):
-        Phi[:, colid] = X ** d
+        Phi[:, colid] = X**d
         colid += D
-
     return Phi
 
 
-def from_posterior(param, samples, distribution=None, half=False, freedom=1):
-    if len(samples.shape) > 1:
-        shape = samples.shape[1:]
-    else:
-        shape = None
-
-    if (distribution is None):
+def from_posterior(param, samples, shape,  distribution=None, half=False, freedom=1):
+    if distribution is None:
         smin, smax = np.min(samples), np.max(samples)
         width = smax - smin
         x = np.linspace(smin, smax, 1000)
@@ -98,114 +97,223 @@ def from_posterior(param, samples, distribution=None, half=False, freedom=1):
             return pm.distributions.Interpolated(param, x, y)
         else:
             return pm.distributions.Interpolated(param, x, y, shape=shape)
-    elif (distribution == 'normal'):
+    elif distribution == "normal":
         temp = stats.norm.fit(samples)
         if shape is None:
             return pm.Normal(param, mu=temp[0], sigma=freedom * temp[1])
         else:
             return pm.Normal(param, mu=temp[0], sigma=freedom * temp[1], shape=shape)
-    elif (distribution == 'hnormal'):
+    elif distribution == "hnormal":
         temp = stats.halfnorm.fit(samples)
         if shape is None:
             return pm.HalfNormal(param, sigma=freedom * temp[1])
         else:
             return pm.HalfNormal(param, sigma=freedom * temp[1], shape=shape)
-    elif (distribution == 'hcauchy'):
+    elif distribution == "hcauchy":
         temp = stats.halfcauchy.fit(samples)
         if shape is None:
             return pm.HalfCauchy(param, freedom * temp[1])
         else:
             return pm.HalfCauchy(param, freedom * temp[1], shape=shape)
-    elif (distribution == 'uniform'):
+    elif distribution == "uniform":
         upper_bound = np.percentile(samples, 95)
         lower_bound = np.percentile(samples, 5)
         r = np.abs(upper_bound - lower_bound)
         if shape is None:
-            return pm.Uniform(param, lower=lower_bound - freedom * r,
-                              upper=upper_bound + freedom * r)
+            return pm.Uniform(
+                param, lower=lower_bound - freedom * r, upper=upper_bound + freedom * r
+            )
         else:
-            return pm.Uniform(param, lower=lower_bound - freedom * r,
-                              upper=upper_bound + freedom * r, shape=shape)
-    elif (distribution == 'huniform'):
+            return pm.Uniform(
+                param,
+                lower=lower_bound - freedom * r,
+                upper=upper_bound + freedom * r,
+                shape=shape,
+            )
+    elif distribution == "huniform":
         upper_bound = np.percentile(samples, 95)
         lower_bound = np.percentile(samples, 5)
         r = np.abs(upper_bound - lower_bound)
         if shape is None:
             return pm.Uniform(param, lower=0, upper=upper_bound + freedom * r)
         else:
-            return pm.Uniform(param, lower=0, upper=upper_bound + freedom * r, shape=shape)
+            return pm.Uniform(
+                param, lower=0, upper=upper_bound + freedom * r, shape=shape
+            )
 
-    elif (distribution == 'gamma'):
+    elif distribution == "gamma":
         alpha_fit, loc_fit, invbeta_fit = stats.gamma.fit(samples)
         if shape is None:
-            return pm.Gamma(param, alpha=freedom * alpha_fit, beta=freedom / invbeta_fit)
+            return pm.Gamma(
+                param, alpha=freedom * alpha_fit, beta=freedom / invbeta_fit
+            )
         else:
-            return pm.Gamma(param, alpha=freedom * alpha_fit, beta=freedom / invbeta_fit, shape=shape)
-
-    elif (distribution == 'igamma'):
+            return pm.Gamma(
+                param,
+                alpha=freedom * alpha_fit,
+                beta=freedom / invbeta_fit,
+                shape=shape,
+            )
+
+    elif distribution == "igamma":
         alpha_fit, loc_fit, beta_fit = stats.gamma.fit(samples)
         if shape is None:
-            return pm.InverseGamma(param, alpha=freedom * alpha_fit, beta=freedom * beta_fit)
+            return pm.InverseGamma(
+                param, alpha=freedom * alpha_fit, beta=freedom * beta_fit
+            )
         else:
-            return pm.InverseGamma(param, alpha=freedom * alpha_fit, beta=freedom * beta_fit, shape=shape)
+            return pm.InverseGamma(
+                param, alpha=freedom * alpha_fit, beta=freedom * beta_fit, shape=shape
+            )
 
 
-def hbr(X, y, batch_effects, batch_effects_size, configs, trace=None):
+def hbr(X, y, batch_effects, configs, idata=None):
     """
     :param X: [N×P] array of clinical covariates
     :param y: [N×1] array of neuroimaging measures
     :param batch_effects: [N×M] array of batch effects
     :param batch_effects_size: [b1, b2,...,bM] List of counts of unique values of batch effects
     :param configs:
-    :param trace:
-    :param return_shared_variables: If true, returns references to the shared variables. The values of the shared variables can be set manually, allowing running the same model on different data without re-compiling it. 
+    :param idata:
+    :param return_shared_variables: If true, returns references to the shared variables. The values of the shared variables can be set manually, allowing running the same model on different data without re-compiling it.
     :return:
     """
-    X = theano.shared(X)
-    X = theano.tensor.cast(X,'floatX')
-    y = theano.shared(y)
-    y = theano.tensor.cast(y,'floatX')
 
+    # Make a param builder that will make the correct calls
+    pb = ParamBuilder(X, y, batch_effects, idata, configs)
+
+    def get_sample_dims(var):
+        if configs[f'random_{var}']:
+            return 'datapoints'
+        elif configs[f'random_slope_{var}']:
+            return 'datapoints'
+        elif configs[f'random_intercept_{var}']:
+            return 'datapoints'
+        elif configs[f'linear_{var}']:
+            return 'datapoints'
+        return None
+
+    with pm.Model(coords=pb.coords) as model:
+        model.add_coord("datapoints", np.arange(X.shape[0]), mutable=True)
+        X = pm.MutableData("X", X, dims=("datapoints", "basis_functions"))
+        pb.X = X
+        y = pm.MutableData("y", np.squeeze(y), dims="datapoints")
+        pb.model = model
+        pb.batch_effect_indices = tuple(
+            [
+                pm.Data(
+                    pb.batch_effect_dim_names[i],
+                    pb.batch_effect_indices[i],
+                    mutable=True,
+                    dims="datapoints",
+                )
+                for i in range(len(pb.batch_effect_indices))
+            ]
+        )
 
-    with pm.Model() as model:
-
-        # Make a param builder that will make the correct calls
-        pb = ParamBuilder(model, X, y, batch_effects, trace, configs)
-
-        if configs['likelihood'] == 'Normal':
-            mu = pb.make_param("mu", mu_slope_mu_params = (0.,10.), 
-                               sigma_slope_mu_params = (5.,), 
-                               mu_intercept_mu_params=(0.,10.), 
-                               sigma_intercept_mu_params = (5.,)).get_samples(pb)
-            sigma = pb.make_param("sigma", mu_sigma_params = (10., 5.),
-                                  sigma_sigma_params = (5.,)).get_samples(pb)
-            sigma_plus = pm.math.log(1+pm.math.exp(sigma))
-            y_like = pm.Normal('y_like',mu=mu, sigma=sigma_plus, observed=y)
-
-        elif configs['likelihood'] in ['SHASHb','SHASHo','SHASHo2']:
+        if configs["likelihood"] == "Normal":
+            
+            mu = pm.Deterministic(
+                "mu_samples",
+                pb.make_param(
+                    "mu",
+                    mu_slope_mu_params=(0.0, 3.0),
+                    sigma_slope_mu_params=(3.0,),
+                    mu_intercept_mu_params=(0.0, 3.0),
+                    sigma_intercept_mu_params=(3.0,),
+                ).get_samples(pb),
+                dims=get_sample_dims('mu'),
+            )
+            sigma = pm.Deterministic(
+                "sigma_samples",
+                pb.make_param(
+                    "sigma", mu_sigma_params=(0.0, 2.0), sigma_sigma_params=(2.0,)
+                ).get_samples(pb),
+                dims=get_sample_dims('sigma'),
+            )
+            sigma_plus = pm.Deterministic(
+                "sigma_plus_samples", pm.math.log(1 + pm.math.exp(sigma/3))*3, dims=get_sample_dims('sigma')
+            )
+            y_like = pm.Normal(
+                "y_like", mu, sigma=sigma_plus, observed=y, dims="datapoints"
+            )
+
+        elif configs["likelihood"] in ["SHASHb", "SHASHo", "SHASHo2"]:
             """
             Comment 1
-            The current parameterizations are tuned towards standardized in- and output data. 
+            The current parameterizations are tuned towards standardized in- and output data.
             It is possible to adjust the priors through the XXX_dist and XXX_params kwargs, like here we do with epsilon_params.
-            Supported distributions are listed in the Prior class. 
-
+            Supported distributions are listed in the Prior class.
             Comment 2
-            Any mapping that is applied here after sampling should also be applied in util.hbr_utils.forward in order for the functions there to properly work. 
+            Any mapping that is applied here after sampling should also be applied in util.hbr_utils.forward in order for the functions there to properly work.
             For example, the softplus applied to sigma here is also applied in util.hbr_utils.forward
             """
-            SHASH_map = {'SHASHb':SHASHb,'SHASHo':SHASHo,'SHASHo2':SHASHo2}
-
-            mu =            pb.make_param("mu",         slope_mu_params = (0.,3.), mu_intercept_mu_params=(0.,1.), sigma_intercept_mu_params = (1.,)).get_samples(pb)
-            sigma =         pb.make_param("sigma",      sigma_params = (1.,2.),    slope_sigma_params=(0.,1.),     intercept_sigma_params = (1., 1.)).get_samples(pb)
-            sigma_plus =    pm.math.log(1+pm.math.exp(sigma))
-            epsilon =       pb.make_param("epsilon",    epsilon_params = (0.,1.),  slope_epsilon_params=(0.,1.), intercept_epsilon_params=(0.,1)).get_samples(pb)
-            delta =         pb.make_param("delta",      delta_params=(1.5,2.),     slope_delta_params=(0.,1),   intercept_delta_params=(2., 1)).get_samples(pb)
-            delta_plus =    pm.math.log(1+pm.math.exp(delta)) + 0.3
-            y_like = SHASH_map[configs['likelihood']]('y_like', mu=mu, sigma=sigma_plus, epsilon=epsilon, delta=delta_plus, observed = y)
-
-    return model
-
-
+            SHASH_map = {"SHASHb": SHASHb, "SHASHo": SHASHo, "SHASHo2": SHASHo2}
+
+            mu = pm.Deterministic(
+                "mu_samples",
+                pb.make_param(
+                    "mu",
+                    slope_mu_params=(0.0, 2.0),
+                    mu_slope_mu_params=(0.0, 2.0),
+                    sigma_slope_mu_params=(2.0,),
+                    mu_intercept_mu_params=(0.0, 2.0),
+                    sigma_intercept_mu_params=(2.0,),
+                ).get_samples(pb),
+                dims=get_sample_dims('mu'),
+            )
+            sigma = pm.Deterministic(
+                "sigma_samples",
+                pb.make_param(
+                    "sigma",
+                    sigma_params=(1.0, 1.0),
+                    sigma_dist="normal",
+                    slope_sigma_params=(0.0, 1.0),
+                    intercept_sigma_params=(1.0, 1.0),
+                ).get_samples(pb),
+                dims=get_sample_dims('sigma'),
+            )
+            sigma_plus = pm.Deterministic(
+                "sigma_plus_samples", np.log(1 + np.exp(sigma)), dims=get_sample_dims('sigma')
+            )
+            epsilon = pm.Deterministic(
+                "epsilon_samples",
+                pb.make_param(
+                    "epsilon",
+                    epsilon_params=(0.0, 1.0),
+                    slope_epsilon_params=(0.0, 0.2),
+                    intercept_epsilon_params=(0.0, 0.2),
+                ).get_samples(pb),
+                dims=get_sample_dims('epsilon'),
+            )
+            delta = pm.Deterministic(
+                "delta_samples",
+                pb.make_param(
+                    "delta",
+                    delta_params=(1.0, 1.0),
+                    delta_dist="normal",
+                    slope_delta_params=(0.0, 0.2),
+                    intercept_delta_params=(1.0, 0.3),
+                ).get_samples(pb),
+                dims=get_sample_dims('delta'),
+            )
+            delta_plus = pm.Deterministic(
+                "delta_plus_samples",
+                np.log(1 + np.exp(delta * 10)) / 10 + 0.3,
+                dims=get_sample_dims('delta'),
+            )
+            y_like = SHASH_map[configs["likelihood"]](
+                "y_like",
+                mu=mu,
+                sigma=sigma_plus,
+                epsilon=epsilon,
+                delta=delta_plus,
+                observed=y,
+                dims="datapoints",
+            )
+        return model
+
+  
 
 class HBR:
 
@@ -214,7 +322,7 @@ class HBR:
     Basic usage::
 
         model = HBR(configs)
-        trace = model.estimate(X, y, batch_effects)
+        idata = model.estimate(X, y, batch_effects)
         ys,s2 = model.predict(X, batch_effects)
 
     where the variables are
@@ -230,245 +338,282 @@ class HBR:
     Written by S.M. Kia
     """
 
-    def get_step_methods(self, m):
-        """
-        This can be used to assign default step functions. However, the nuts initialization keyword doesnt work together with this... so better not use it. 
-
-        STEP_METHODS = (
-            NUTS,
-            HamiltonianMC,
-            Metropolis,
-            BinaryMetropolis,
-            BinaryGibbsMetropolis,
-            Slice,
-            CategoricalGibbsMetropolis,
-        )
-        :param m: a PyMC3 model
-        :return:
-        """
-        samplermap = {'NUTS': NUTS, 'MH': Metropolis, 'Slice': Slice, 'HMC': HamiltonianMC}
-        fallbacks = [Metropolis]         # We are using MH as a fallback method here
-        if self.configs['sampler'] == 'NUTS':
-            step_kwargs = {'nuts': {'target_accept': self.configs['target_accept']}}
-        else:
-            step_kwargs = None
-        return pm.sampling.assign_step_methods(m, methods=[samplermap[self.configs['sampler']]] + fallbacks,
-                                               step_kwargs=step_kwargs)
-
     def __init__(self, configs):
         self.bsp = None
-        self.model_type = configs['type']
+        self.model_type = configs["type"]
         self.configs = configs
 
     def get_modeler(self):
-        return {'nn': nn_hbr}.get(self.model_type, hbr)
-        
+        return hbr
+
     def transform_X(self, X):
-        if self.model_type == 'polynomial':
-            Phi = create_poly_basis(X, self.configs['order'])
-        elif self.model_type == 'bspline':
+        if self.model_type == "polynomial":
+            Phi = create_poly_basis(X, self.configs["order"])
+        elif self.model_type == "bspline":
             if self.bsp is None:
-                self.bsp = bspline_fit(X, self.configs['order'], self.configs['nknots'])
+                self.bsp = bspline_fit(X, self.configs["order"], self.configs["nknots"])
             bspline = bspline_transform(X, self.bsp)
-            Phi = np.concatenate((X, bspline), axis = 1)
+            Phi = np.concatenate((X, bspline), axis=1)
         else:
             Phi = X
         return Phi
 
-
-    def find_map(self, X, y, batch_effects,method='L-BFGS-B'):
-        """ Function to estimate the model """
+    def find_map(self, X, y, batch_effects, method="L-BFGS-B"):
+        """Function to estimate the model"""
         X, y, batch_effects = expand_all(X, y, batch_effects)
-
-        self.batch_effects_num = batch_effects.shape[1]
-        self.batch_effects_size = []
-        for i in range(self.batch_effects_num):
-            self.batch_effects_size.append(len(np.unique(batch_effects[:, i])))
-
         X = self.transform_X(X)
         modeler = self.get_modeler()
-        with modeler(X, y, batch_effects, self.batch_effects_size, self.configs) as m:
+        with modeler(X, y, batch_effects, self.configs) as m:
             self.MAP = pm.find_MAP(method=method)
         return self.MAP
 
-    def estimate(self, X, y, batch_effects):
-
-        """ Function to estimate the model """
+    def estimate(self, X, y, batch_effects, **kwargs):
+        """Function to estimate the model"""
         X, y, batch_effects = expand_all(X, y, batch_effects)
-
-        self.batch_effects_num = batch_effects.shape[1]
-        self.batch_effects_size = []
-        for i in range(self.batch_effects_num):
-            self.batch_effects_size.append(len(np.unique(batch_effects[:, i])))
-
         X = self.transform_X(X)
         modeler = self.get_modeler()
-        with modeler(X, y, batch_effects, self.batch_effects_size, self.configs) as m:
-            self.trace = pm.sample(draws=self.configs['n_samples'],
-                                   tune=self.configs['n_tuning'],
-                                   chains=self.configs['n_chains'],
-                                   init=self.configs['init'], n_init=500000,
-                                   cores=self.configs['cores'])
-        return self.trace
-
-    def predict(self, X, batch_effects, pred='single'):
-        """ Function to make predictions from the model """
+        if hasattr(self, 'idata'):
+            del self.idata
+        with modeler(X, y, batch_effects, self.configs) as m:
+            self.idata = pm.sample(
+                draws=self.configs["n_samples"],
+                tune=self.configs["n_tuning"],
+                chains=self.configs["n_chains"],
+                init=self.configs["init"],
+                n_init=500000,
+                cores=self.configs["cores"],
+            )
+        self.vars_to_sample = ['y_like']
+        if self.configs['remove_datapoints_from_posterior']:
+            chain = self.idata.posterior.coords['chain'].data
+            draw = self.idata.posterior.coords['draw'].data
+            for j in self.idata.posterior.variables.mapping.keys():
+                if j.endswith('_samples'):
+                    dummy_array = xarray.DataArray(data = np.zeros((len(chain), len(draw), 1)), coords = {'chain':chain, 'draw':draw,'empty':np.array([0])}, name=j)
+                    self.idata.posterior[j] = dummy_array
+                    self.vars_to_sample.append(j)
+        return self.idata
+
+    def predict(
+        self, X, batch_effects, batch_effects_maps, pred="single", var_names=None, **kwargs
+    ):
+        """Function to make predictions from the model
+        Args:
+            X: Covariates
+            batch_effects: batch effects corresponding to X
+            all_batch_effects: combinations of all batch effects that were present the training data
+        """
         X, batch_effects = expand_all(X, batch_effects)
 
-        samples = self.configs['n_samples']
+        samples = self.configs["n_samples"]
         y = np.zeros([X.shape[0], 1])
+        X = self.transform_X(X)
+        modeler = self.get_modeler()
 
-        if pred == 'single':
-            X = self.transform_X(X)
-            modeler = self.get_modeler()
-            with modeler(X, y, batch_effects, self.batch_effects_size, self.configs):
-                ppc = pm.sample_posterior_predictive(self.trace, samples=samples, progressbar=True)
-            pred_mean = ppc['y_like'].mean(axis=0)
-            pred_var = ppc['y_like'].var(axis=0)
+        # Make an array with occurences of all the values in be_train, but with the same size as be_test
+        truncated_batch_effects_train = np.stack(
+            [
+                np.resize(np.array(list(batch_effects_maps[i].keys())), X.shape[0])
+                for i in range(batch_effects.shape[1])
+            ],
+            axis=1,
+        )
+
+        # See if a list of var_names is provided, set to self.vars_to_sample otherwise
+        if (var_names is None) or (var_names == ['y_like']):
+            var_names = self.vars_to_sample
+
+        n_samples = X.shape[0]
+        with modeler(X, y, truncated_batch_effects_train, self.configs) as model:
+            # For each batch effect dim
+            for i in range(batch_effects.shape[1]):
+                # Make a map that maps batch effect values to their index
+                valmap = batch_effects_maps[i]
+                # Compute those indices for the test data
+                indices = list(map(lambda x: valmap[x], batch_effects[:, i]))
+                # Those indices need to be used by the model
+                pm.set_data({f"batch_effect_{i}": indices})
+
+            self.idata = pm.sample_posterior_predictive(
+                trace=self.idata,
+                extend_inferencedata=True,
+                progressbar=True,
+                var_names=var_names,
+            )
+        pred_mean = self.idata.posterior_predictive["y_like"].to_numpy().mean(axis=(0, 1))
+        pred_var = self.idata.posterior_predictive["y_like"].to_numpy().var(axis=(0, 1))
 
         return pred_mean, pred_var
 
     def estimate_on_new_site(self, X, y, batch_effects):
-        """ Function to adapt the model """
+        """Function to adapt the model"""
         X, y, batch_effects = expand_all(X, y, batch_effects)
-
-        self.batch_effects_num = batch_effects.shape[1]
-        self.batch_effects_size = []
-        for i in range(self.batch_effects_num):
-            self.batch_effects_size.append(len(np.unique(batch_effects[:, i])))
-
         X = self.transform_X(X)
         modeler = self.get_modeler()
-        with modeler(X, y, batch_effects, self.batch_effects_size,
-                     self.configs, trace=self.trace) as m:
-            self.trace = pm.sample(self.configs['n_samples'],
-                                   tune=self.configs['n_tuning'],
-                                   chains=self.configs['n_chains'],
-                                   target_accept=self.configs['target_accept'],
-                                   init=self.configs['init'], n_init=50000,
-                                   cores=self.configs['cores'])
-        return self.trace
+        with modeler(X, y, batch_effects, self.configs, idata=self.idata) as m:
+            self.idata = pm.sample(
+                self.configs["n_samples"],
+                tune=self.configs["n_tuning"],
+                chains=self.configs["n_chains"],
+                target_accept=self.configs["target_accept"],
+                init=self.configs["init"],
+                n_init=50000,
+                cores=self.configs["cores"],
+            )
+        return self.idata
 
     def predict_on_new_site(self, X, batch_effects):
-        """ Function to make predictions from the model """
+        """Function to make predictions from the model"""
         X, batch_effects = expand_all(X, batch_effects)
-
-        samples = self.configs['n_samples']
+        samples = self.configs["n_samples"]
         y = np.zeros([X.shape[0], 1])
-
         X = self.transform_X(X)
         modeler = self.get_modeler()
-        with modeler(X, y, batch_effects, self.batch_effects_size, self.configs, trace=self.trace):
-            ppc = pm.sample_posterior_predictive(self.trace, samples=samples, progressbar=True)
-        pred_mean = ppc['y_like'].mean(axis=0)
-        pred_var = ppc['y_like'].var(axis=0)
+        with modeler(X, y, batch_effects, self.configs, idata=self.idata):
+            self.idata = pm.sample_posterior_predictive(
+                self.idata, extend_inferencedata=True, progressbar=True
+            )
+        pred_mean = self.idata.posterior_predictive["y_like"].mean(axis=(0, 1))
+        pred_var = self.idata.posterior_predictive["y_like"].var(axis=(0, 1))
 
         return pred_mean, pred_var
 
     def generate(self, X, batch_effects, samples):
-        """ Function to generate samples from posterior predictive distribution """
+        """Function to generate samples from posterior predictive distribution"""
         X, batch_effects = expand_all(X, batch_effects)
 
         y = np.zeros([X.shape[0], 1])
 
         X = self.transform_X(X)
         modeler = self.get_modeler()
-        with modeler(X, y, batch_effects, self.batch_effects_size, self.configs):
-            ppc = pm.sample_posterior_predictive(self.trace, samples=samples, progressbar=True)
-
-        generated_samples = np.reshape(ppc['y_like'].squeeze().T, [X.shape[0] * samples, 1])
+        with modeler(X, y, batch_effects, self.configs):
+            ppc = pm.sample_posterior_predictive(self.idata, progressbar=True)
+        generated_samples = np.reshape(
+            ppc.posterior_predictive["y_like"].squeeze().T, [X.shape[0] * samples, 1]
+        )
         X = np.repeat(X, samples)
         if len(X.shape) == 1:
             X = np.expand_dims(X, axis=1)
         batch_effects = np.repeat(batch_effects, samples, axis=0)
         if len(batch_effects.shape) == 1:
             batch_effects = np.expand_dims(batch_effects, axis=1)
-
         return X, batch_effects, generated_samples
 
-    def sample_prior_predictive(self, X, batch_effects, samples, trace=None):
-        """ Function to sample from prior predictive distribution """
-
-        if len(X.shape) == 1:
-            X = np.expand_dims(X, axis=1)
-        if len(batch_effects.shape) == 1:
-            batch_effects = np.expand_dims(batch_effects, axis=1)
-
-        self.batch_effects_num = batch_effects.shape[1]
-        self.batch_effects_size = []
-        for i in range(self.batch_effects_num):
-            self.batch_effects_size.append(len(np.unique(batch_effects[:, i])))
-
-        y = np.zeros([X.shape[0], 1])
+    def sample_prior_predictive(self, X, batch_effects, samples, y = None, idata=None):
+        """Function to sample from prior predictive distribution"""
+        if y is None:
+            y = np.zeros([X.shape[0], 1])
+        X, y, batch_effects = expand_all(X, y, batch_effects)
 
-        if self.model_type == 'linear':
-            with hbr(X, y, batch_effects, self.batch_effects_size, self.configs,
-                     trace):
-                ppc = pm.sample_prior_predictive(samples=samples)
-        return ppc
+        X = self.transform_X(X)
+        modeler = self.get_modeler()
+        with modeler(X, y, batch_effects, self.configs, idata):
+            self.idata = pm.sample_prior_predictive(samples=samples)
+        return self.idata
 
     def get_model(self, X, y, batch_effects):
         X, y, batch_effects = expand_all(X, y, batch_effects)
-
-        self.batch_effects_num = batch_effects.shape[1]
-        self.batch_effects_size = []
-        for i in range(self.batch_effects_num):
-            self.batch_effects_size.append(len(np.unique(batch_effects[:, i])))
         modeler = self.get_modeler()
         X = self.transform_X(X)
-        return modeler(X, y, batch_effects, self.batch_effects_size, self.configs, self.trace)
-
-    def create_dummy_inputs(self, covariate_ranges = [[0.1,0.9,0.01]]):
+        idata = self.idata if hasattr(self, "idata") else None
+        return modeler(X, y, batch_effects, self.configs, idata=idata)
 
+    def create_dummy_inputs(self, covariate_ranges=[[0.1, 0.9, 0.01]]):
         arrays = []
         for i in range(len(covariate_ranges)):
-            arrays.append(np.arange(covariate_ranges[i][0],covariate_ranges[i][1], covariate_ranges[i][2]))
+            arrays.append(
+                np.arange(
+                    covariate_ranges[i][0],
+                    covariate_ranges[i][1],
+                    covariate_ranges[i][2],
+                )
+            )
         X = cartesian_product(arrays)
         X_dummy = np.concatenate([X for i in range(np.prod(self.batch_effects_size))])
-
         arrays = []
         for i in range(self.batch_effects_num):
             arrays.append(np.arange(0, self.batch_effects_size[i]))
         batch_effects = cartesian_product(arrays)
         batch_effects_dummy = np.repeat(batch_effects, X.shape[0], axis=0)
-
-        return X_dummy, batch_effects_dummy 
+        return X_dummy, batch_effects_dummy
+
+    def Rhats(self, var_names, thin = 1, resolution = 100):
+        """Get Rhat of posterior samples as function of sampling iteration"""
+        idata = self.idata
+        testvars = az.extract(idata, group='posterior', var_names=var_names, combined=False)
+        rhat_dict={}
+        for var_name in var_names:
+            var = np.stack(testvars[var_name].to_numpy())[:,::thin]     
+            var = var.reshape((var.shape[0], var.shape[1], -1))
+            vardim = var.shape[2]
+            interval_skip=var.shape[1]//resolution
+            rhats_var = np.zeros((resolution, vardim))
+            for v in range(vardim):
+                for j in range(resolution):
+                    rhats_var[j,v] = az.rhat(var[:,:j*interval_skip,v])
+            rhat_dict[var_name] = rhats_var
+        return rhat_dict
 
 class Prior:
     """
-    A wrapper class for a PyMC3 distribution. 
-    - creates a fitted distribution from the trace, if one is present
+    A wrapper class for a PyMC distribution.
+    - creates a fitted distribution from the idata, if one is present
     - overloads the __getitem__ function with something that switches between indexing or not, based on the shape
     """
-    def __init__(self, name, dist, params, pb, shape=(1,)) -> None:
+
+    def __init__(self, name, dist, params, pb, has_random_effect=False) -> None:
         self.dist = None
         self.name = name
-        self.shape = shape
-        self.has_random_effect = True if len(shape)>1 else False
-        self.distmap = {'normal': pm.Normal,
-                   'hnormal': pm.HalfNormal,
-                   'gamma': pm.Gamma,
-                   'uniform': pm.Uniform,
-                   'igamma': pm.InverseGamma,
-                   'hcauchy': pm.HalfCauchy}
+        self.has_random_effect = has_random_effect
+        self.distmap = {
+            "normal": pm.Normal,
+            "hnormal": pm.HalfNormal,
+            "gamma": pm.Gamma,
+            "uniform": pm.Uniform,
+            "igamma": pm.InverseGamma,
+            "hcauchy": pm.HalfCauchy,
+            "hstudt": pm.HalfStudentT,
+            "studt": pm.StudentT,
+        }
         self.make_dist(dist, params, pb)
- 
+
     def make_dist(self, dist, params, pb):
-        """This creates a pymc3 distribution. If there is a trace, the distribution is fitted to the trace. If there isn't a trace, the prior is parameterized by the values in (params)"""
+        """This creates a pymc distribution. If there is a idata, the distribution is fitted to the idata. If there isn't a idata, the prior is parameterized by the values in (params)"""
         with pb.model as m:
-            if (pb.trace is not None) and (not self.has_random_effect):
-                int_dist = from_posterior(param=self.name,
-                                            samples=pb.trace[self.name],
-                                            distribution=dist,
-                                            freedom=pb.configs['freedom'])
-                self.dist = int_dist.reshape(self.shape)
+            if pb.idata is not None:
+                # Get samples
+                samples = az.extract(pb.idata, var_names=self.name)
+                # Define mapping to new shape
+                def get_new_dim_size(tup):
+                    oldsize, name = tup
+                    if name.startswith('batch_effect_'):
+                        ind = pb.batch_effect_dim_names.index(name)
+                        return len(np.unique(pb.batch_effect_indices[ind].container.data))
+                    return oldsize
+                
+                new_shape = list(map(get_new_dim_size, zip(samples.shape,samples.dims)))
+                if len(new_shape) == 1:
+                    new_shape = None
+                else:
+                    new_shape = new_shape[:-1]
+                self.dist = from_posterior(
+                    param=self.name,
+                    samples=samples.to_numpy(),
+                    shape = new_shape,
+                    distribution=dist,
+                    freedom=pb.configs["freedom"],
+                )
+
             else:
-                shape_prod = np.product(np.array(self.shape))
-                print(self.name)
-                print(f"dist={dist}")
-                print(f"params={params}")
-                int_dist = self.distmap[dist](self.name, *params, shape=shape_prod)
-                self.dist = int_dist.reshape(self.shape)
+                dims = []
+                if self.has_random_effect:
+                    dims = dims + pb.batch_effect_dim_names
+                if self.name.startswith("slope") or self.name.startswith("offset_slope"):
+                    dims = dims + ["basis_functions"]
+                if dims == []:
+                    self.dist = self.distmap[dist](self.name, *params)
+                else:
+                    self.dist = self.distmap[dist](self.name, *params, dims=dims)
 
     def __getitem__(self, idx):
         """The idx here is the index of the batch-effect. If the prior does not model batch effects, this should return the same value for each index"""
@@ -481,166 +626,202 @@ def __getitem__(self, idx):
 
 class ParamBuilder:
     """
-    A class that simplifies the construction of parameterizations. 
-    It has a lot of attributes necessary for creating the model, including the data, but it is never saved with the model. 
+    A class that simplifies the construction of parameterizations.
+    It has a lot of attributes necessary for creating the model, including the data, but it is never saved with the model.
     It also contains a lot of decision logic for creating the parameterizations.
     """
 
-    def __init__(self, model, X, y, batch_effects, trace, configs):
+    def __init__(self, X, y, batch_effects, idata, configs):
         """
 
         :param model: model to attach all the distributions to
         :param X: Covariates
         :param y: IDPs
-        :param batch_effects: I guess this speaks for itself
-        :param trace:  idem
+        :param batch_effects: array of batch effects
+        :param idata:  idem
         :param configs: idem
         """
-        self.model = model
+        self.model = None  # Needs to be set later, because coords need to be passed at construction of Model
         self.X = X
+        self.n_basis_functions = X.shape[1]
         self.y = y
-        self.batch_effects = batch_effects
-        self.trace = trace
+        self.batch_effects = batch_effects.astype(np.int16)
+        self.idata: az.InferenceData = idata
         self.configs = configs
 
-        self.feature_num = X.shape[1].eval().item()
-        self.y_shape = y.shape.eval()
-        self.n_ys = y.shape[0].eval().item()
+        self.y_shape = y.shape
+        self.n_ys = y.shape[0]
         self.batch_effects_num = batch_effects.shape[1]
 
-        self.batch_effects_size = []
-        self.all_idx = []
+        self.batch_effect_dim_names = []
+        self.batch_effect_indices = []
+        self.coords = OrderedDict()
+        self.coords["basis_functions"] = np.arange(self.n_basis_functions)
+
         for i in range(self.batch_effects_num):
-            # Count the unique values for each batch effect
-            self.batch_effects_size.append(len(np.unique(self.batch_effects[:, i])))
-            # Store the unique values for each batch effect
-            self.all_idx.append(np.int16(np.unique(self.batch_effects[:, i])))
-        # Make a cartesian product of all the unique values of each batch effect
-        self.be_idx = list(product(*self.all_idx))
-
-        # Make tuples of batch effects ID's and indices of datapoints with that specific combination of batch effects
-        self.be_idx_tups = []
-        for be in self.be_idx:
-            a = []
-            for i, b in enumerate(be):
-                a.append(self.batch_effects[:, i] == b)
-            idx = reduce(np.logical_and, a).nonzero()
-            if idx[0].shape[0] != 0:
-                self.be_idx_tups.append((be, idx))
-
-    def make_param(self, name, dim = (1,), **kwargs):
-        if self.configs.get(f'linear_{name}', False):
-            # First make a slope and intercept, and use those to make a linear parameterization 
-            slope_parameterization = self.make_param(f'slope_{name}', dim=[self.feature_num], **kwargs)
-            intercept_parameterization = self.make_param(f'intercept_{name}', **kwargs)
-            return LinearParameterization(name=name, dim=dim, 
-                                    slope_parameterization=slope_parameterization, 
-                                    intercept_parameterization=intercept_parameterization,
-                                    pb=self, 
-                                    **kwargs)
-        
-        elif self.configs.get(f'random_{name}', False):
-            if self.configs.get(f'centered_{name}', True):
-                return CentralRandomFixedParameterization(name=name, pb=self, dim=dim, **kwargs)
+            batch_effect_dim_name = f"batch_effect_{i}"
+            self.batch_effect_dim_names.append(batch_effect_dim_name)
+            this_be_values, this_be_indices = np.unique(
+                self.batch_effects[:, i], return_inverse=True
+            )
+            self.coords[batch_effect_dim_name] = this_be_values
+            self.batch_effect_indices.append(this_be_indices)
+
+    def make_param(self, name, **kwargs):
+        if self.configs.get(f"linear_{name}", False):
+            # First make a slope and intercept, and use those to make a linear parameterization
+            slope_parameterization = self.make_param(f"slope_{name}", **kwargs)
+            intercept_parameterization = self.make_param(f"intercept_{name}", **kwargs)
+            return LinearParameterization(
+                name=name,
+                slope_parameterization=slope_parameterization,
+                intercept_parameterization=intercept_parameterization,
+                **kwargs,
+            )
+
+        elif self.configs.get(f"random_{name}", False):
+            if self.configs.get(f"centered_{name}", True):
+                return CentralRandomFixedParameterization(name=name, pb=self, **kwargs)
             else:
-                return NonCentralRandomFixedParameterization(name=name, pb=self, dim=dim, **kwargs)
+                return NonCentralRandomFixedParameterization(
+                    name=name, pb=self, **kwargs
+                )
         else:
-            return FixedParameterization(name=name, dim=dim, pb=self,**kwargs)
+            return FixedParameterization(name=name, pb=self, **kwargs)
 
 
 class Parameterization:
     """
     This is the top-level parameterization class from which all the other parameterizations inherit.
     """
-    def __init__(self, name, dim):
+
+    def __init__(self, name):
         self.name = name
-        self.dim = dim
-        print(name, type(self))
+        # print(name, type(self))
 
     def get_samples(self, pb):
-
-        with pb.model:
-            samples = theano.tensor.zeros([pb.n_ys, *self.dim])
-            for be, idx in pb.be_idx_tups:
-                samples = theano.tensor.set_subtensor(samples[idx], self.dist[be])
-        return samples
+        pass
 
 
 class FixedParameterization(Parameterization):
     """
     A parameterization that takes a single value for all input. It does not depend on anything except its hyperparameters
     """
-    def __init__(self, name, dim, pb:ParamBuilder, **kwargs):
-        super().__init__(name, dim)
-        dist = kwargs.get(f'{name}_dist','normal')
-        params = kwargs.get(f'{name}_params',(0.,1.))
-        self.dist = Prior(name, dist, params, pb, shape = dim)
+
+    def __init__(self, name, pb: ParamBuilder, **kwargs):
+        super().__init__(name)
+        dist = kwargs.get(f"{name}_dist", "normal")
+        params = kwargs.get(f"{name}_params", (0.0, 1.0))
+        self.dist = Prior(name, dist, params, pb)
+
+    def get_samples(self, pb):
+        with pb.model:
+            return self.dist[0]
 
 
 class CentralRandomFixedParameterization(Parameterization):
     """
     A parameterization that is fixed for each batch effect. This is sampled in a central fashion;
-    the values are sampled from normal distribution with a group mean and group variance 
+    the values are sampled from normal distribution with a group mean and group variance
     """
-    def __init__(self, name, dim, pb:ParamBuilder, **kwargs):
-        super().__init__(name, dim)
+
+    def __init__(self, name, pb: ParamBuilder, **kwargs):
+        super().__init__(name)
 
         # Normal distribution is default for mean
-        mu_dist = kwargs.get(f'mu_{name}_dist','normal')
-        mu_params = kwargs.get(f'mu_{name}_params',(0.,1.))
-        mu_prior = Prior(f'mu_{name}', mu_dist, mu_params, pb, shape = dim)
+        mu_dist = kwargs.get(f"mu_{name}_dist", "normal")
+        mu_params = kwargs.get(f"mu_{name}_params", (0.0, 1.0))
+        mu_prior = Prior(f"mu_{name}", mu_dist, mu_params, pb)
+
+        # HalfNormal is default for sigma
+        sigma_dist = kwargs.get(f"sigma_{name}_dist", "hnormal")
+        sigma_params = kwargs.get(f"sigma_{name}_params", (1.0,))
+        sigma_prior = Prior(f"sigma_{name}", sigma_dist, sigma_params, pb)
+
+        dims = (
+            [*pb.batch_effect_dim_names, "basis_functions"]
+            if self.name.startswith("slope")
+            else pb.batch_effect_dim_names
+        )
+        self.dist = pm.Normal(
+            name=name,
+            mu=mu_prior.dist,
+            sigma=sigma_prior.dist,
+            dims=dims,
+        )
 
-        # HalfCauchy is default for sigma
-        sigma_dist = kwargs.get(f'sigma_{name}_dist','hcauchy')
-        sigma_params = kwargs.get(f'sigma_{name}_params',(1.,))
-        sigma_prior = Prior(f'sigma_{name}',sigma_dist, sigma_params, pb, shape = [*pb.batch_effects_size, *dim])
+    def get_samples(self, pb: ParamBuilder):
+        with pb.model:
+            samples = self.dist[pb.batch_effect_indices]
+            return samples
 
-        self.dist = pm.Normal(name=name, mu=mu_prior.dist, sigma=sigma_prior.dist, shape = [*pb.batch_effects_size, *dim])
-    
 
 class NonCentralRandomFixedParameterization(Parameterization):
     """
     A parameterization that is fixed for each batch effect. This is sampled in a non-central fashion;
-    the values are a sum of a group mean and noise values scaled with a group scaling factor 
+    the values are a sum of a group mean and noise values scaled with a group scaling factor
     """
-    def __init__(self, name,dim,  pb:ParamBuilder, **kwargs):
-        super().__init__(name, dim)
+
+    def __init__(self, name, pb: ParamBuilder, **kwargs):
+        super().__init__(name)
 
         # Normal distribution is default for mean
-        mu_dist = kwargs.get(f'mu_{name}_dist','normal')
-        mu_params = kwargs.get(f'mu_{name}_params',(0.,1.))
-        mu_prior = Prior(f'mu_{name}', mu_dist, mu_params, pb, shape = dim)
+        mu_dist = kwargs.get(f"mu_{name}_dist", "normal")
+        mu_params = kwargs.get(f"mu_{name}_params", (0.0, 1.0))
+        mu_prior = Prior(f"mu_{name}", mu_dist, mu_params, pb)
 
-        # HalfCauchy is default for sigma
-        sigma_dist = kwargs.get(f'sigma_{name}_dist','hcauchy')
-        sigma_params = kwargs.get(f'sigma_{name}_params',(1.,))
-        sigma_prior = Prior(f'sigma_{name}',sigma_dist, sigma_params, pb, shape = dim)
+        # HalfNormal is default for sigma
+        sigma_dist = kwargs.get(f"sigma_{name}_dist", "hnormal")
+        sigma_params = kwargs.get(f"sigma_{name}_params", (1.0,))
+        sigma_prior = Prior(f"sigma_{name}", sigma_dist, sigma_params, pb)
 
         # Normal is default for offset
-        offset_dist = kwargs.get(f'offset_{name}_dist','normal')
-        offset_params = kwargs.get(f'offset_{name}_params',(0.,1.))
-        offset_prior = Prior(f'offset_{name}',offset_dist, offset_params, pb, shape = [*pb.batch_effects_size, *dim])
+        offset_dist = kwargs.get(f"offset_{name}_dist", "normal")
+        offset_params = kwargs.get(f"offset_{name}_params", (0.0, 1.0))
+        offset_prior = Prior(
+            f"offset_{name}", offset_dist, offset_params, pb, has_random_effect=True
+        )
+        dims = (
+            [*pb.batch_effect_dim_names, "basis_functions"]
+            if self.name.startswith("slope")
+            else pb.batch_effect_dim_names
+        )
+        self.dist = pm.Deterministic(
+            name=name,
+            var=mu_prior.dist + sigma_prior.dist * offset_prior.dist,
+            dims=dims,
+        )
 
-        self.dist = pm.Deterministic(name=name, var=mu_prior.dist+sigma_prior.dist*offset_prior.dist)
+    def get_samples(self, pb: ParamBuilder):
+        with pb.model:
+            samples = self.dist[pb.batch_effect_indices]
+            return samples
 
 
 class LinearParameterization(Parameterization):
     """
-    A parameterization that can model a linear dependence on X. 
+    A parameterization that can model a linear dependence on X.
     """
-    def __init__(self, name, dim, slope_parameterization, intercept_parameterization, pb, **kwargs):
-        super().__init__( name, dim)
+
+    def __init__(
+        self, name, slope_parameterization, intercept_parameterization, **kwargs
+    ):
+        super().__init__(name)
         self.slope_parameterization = slope_parameterization
         self.intercept_parameterization = intercept_parameterization
 
-    def get_samples(self, pb:ParamBuilder):
+    def get_samples(self, pb):
         with pb.model:
-            samples = theano.tensor.zeros([pb.n_ys, *self.dim])
-            for be, idx in pb.be_idx_tups:
-                dot = theano.tensor.dot(pb.X[idx,:], self.slope_parameterization.dist[be]).T
-                intercept = self.intercept_parameterization.dist[be]
-                samples = theano.tensor.set_subtensor(samples[idx,:],dot+intercept)
-        return samples
+            intc = self.intercept_parameterization.get_samples(pb)
+            slope_samples = self.slope_parameterization.get_samples(pb)
+            if pb.configs[f"random_slope_{self.name}"]:
+                slope = pb.X * slope_samples
+                slope = slope.sum(axis=-1)
+            else:
+                slope = pb.X @ self.slope_parameterization.get_samples(pb)
+
+            samples = pm.math.flatten(intc) + pm.math.flatten(slope)
+            return samples
 
 
 def get_design_matrix(X, nm, basis="linear"):
@@ -653,10 +834,9 @@ def get_design_matrix(X, nm, basis="linear"):
     return Phi
 
 
-
-def nn_hbr(X, y, batch_effects, batch_effects_size, configs, trace=None):
-    n_hidden = configs['nn_hidden_neuron_num']
-    n_layers = configs['nn_hidden_layers_num']
+def nn_hbr(X, y, batch_effects, batch_effects_size, configs, idata=None):
+    n_hidden = configs["nn_hidden_neuron_num"]
+    n_layers = configs["nn_hidden_layers_num"]
     feature_num = X.shape[1]
     batch_effects_num = batch_effects.shape[1]
     all_idx = []
@@ -665,14 +845,18 @@ def nn_hbr(X, y, batch_effects, batch_effects_size, configs, trace=None):
     be_idx = list(product(*all_idx))
 
     # Initialize random weights between each layer for the mu:
-    init_1 = pm.floatX(np.random.randn(feature_num, n_hidden) * np.sqrt(1 / feature_num))
+    init_1 = pm.floatX(
+        np.random.randn(feature_num, n_hidden) * np.sqrt(1 / feature_num)
+    )
     init_out = pm.floatX(np.random.randn(n_hidden) * np.sqrt(1 / n_hidden))
 
     std_init_1 = pm.floatX(np.random.rand(feature_num, n_hidden))
     std_init_out = pm.floatX(np.random.rand(n_hidden))
 
     # And initialize random weights between each layer for sigma_noise:
-    init_1_noise = pm.floatX(np.random.randn(feature_num, n_hidden) * np.sqrt(1 / feature_num))
+    init_1_noise = pm.floatX(
+        np.random.randn(feature_num, n_hidden) * np.sqrt(1 / feature_num)
+    )
     init_out_noise = pm.floatX(np.random.randn(n_hidden) * np.sqrt(1 / n_hidden))
 
     std_init_1_noise = pm.floatX(np.random.rand(feature_num, n_hidden))
@@ -682,93 +866,137 @@ def nn_hbr(X, y, batch_effects, batch_effects_size, configs, trace=None):
     if n_layers == 2:
         init_2 = pm.floatX(np.random.randn(n_hidden, n_hidden) * np.sqrt(1 / n_hidden))
         std_init_2 = pm.floatX(np.random.rand(n_hidden, n_hidden))
-        init_2_noise = pm.floatX(np.random.randn(n_hidden, n_hidden) * np.sqrt(1 / n_hidden))
+        init_2_noise = pm.floatX(
+            np.random.randn(n_hidden, n_hidden) * np.sqrt(1 / n_hidden)
+        )
         std_init_2_noise = pm.floatX(np.random.rand(n_hidden, n_hidden))
 
     with pm.Model() as model:
-
-        X = pm.Data('X', X)
-        y = pm.Data('y', y)
-
-        if trace is not None:  # Used when estimating/predicting on a new site
-            weights_in_1_grp = from_posterior('w_in_1_grp', trace['w_in_1_grp'],
-                                              distribution='normal', freedom=configs['freedom'])
-
-            weights_in_1_grp_sd = from_posterior('w_in_1_grp_sd', trace['w_in_1_grp_sd'],
-                                                 distribution='hcauchy', freedom=configs['freedom'])
+        X = pm.Data("X", X)
+        y = pm.Data("y", y)
+
+        if idata is not None:  # Used when estimating/predicting on a new site
+            weights_in_1_grp = from_posterior(
+                "w_in_1_grp",
+                idata["w_in_1_grp"],
+                distribution="normal",
+                freedom=configs["freedom"],
+            )
+
+            weights_in_1_grp_sd = from_posterior(
+                "w_in_1_grp_sd",
+                idata["w_in_1_grp_sd"],
+                distribution="hcauchy",
+                freedom=configs["freedom"],
+            )
 
             if n_layers == 2:
-                weights_1_2_grp = from_posterior('w_1_2_grp', trace['w_1_2_grp'],
-                                                 distribution='normal', freedom=configs['freedom'])
-
-                weights_1_2_grp_sd = from_posterior('w_1_2_grp_sd', trace['w_1_2_grp_sd'],
-                                                    distribution='hcauchy', freedom=configs['freedom'])
-
-            weights_2_out_grp = from_posterior('w_2_out_grp', trace['w_2_out_grp'],
-                                               distribution='normal', freedom=configs['freedom'])
-
-            weights_2_out_grp_sd = from_posterior('w_2_out_grp_sd', trace['w_2_out_grp_sd'],
-                                                  distribution='hcauchy', freedom=configs['freedom'])
-
-            mu_prior_intercept = from_posterior('mu_prior_intercept', trace['mu_prior_intercept'],
-                                                distribution='normal', freedom=configs['freedom'])
-            sigma_prior_intercept = from_posterior('sigma_prior_intercept', trace['sigma_prior_intercept'],
-                                                   distribution='hcauchy', freedom=configs['freedom'])
+                weights_1_2_grp = from_posterior(
+                    "w_1_2_grp",
+                    idata["w_1_2_grp"],
+                    distribution="normal",
+                    freedom=configs["freedom"],
+                )
+
+                weights_1_2_grp_sd = from_posterior(
+                    "w_1_2_grp_sd",
+                    idata["w_1_2_grp_sd"],
+                    distribution="hcauchy",
+                    freedom=configs["freedom"],
+                )
+
+            weights_2_out_grp = from_posterior(
+                "w_2_out_grp",
+                idata["w_2_out_grp"],
+                distribution="normal",
+                freedom=configs["freedom"],
+            )
+
+            weights_2_out_grp_sd = from_posterior(
+                "w_2_out_grp_sd",
+                idata["w_2_out_grp_sd"],
+                distribution="hcauchy",
+                freedom=configs["freedom"],
+            )
+
+            mu_prior_intercept = from_posterior(
+                "mu_prior_intercept",
+                idata["mu_prior_intercept"],
+                distribution="normal",
+                freedom=configs["freedom"],
+            )
+            sigma_prior_intercept = from_posterior(
+                "sigma_prior_intercept",
+                idata["sigma_prior_intercept"],
+                distribution="hcauchy",
+                freedom=configs["freedom"],
+            )
 
         else:
             # Group the mean distribution for input to the hidden layer:
-            weights_in_1_grp = pm.Normal('w_in_1_grp', 0, sd=1,
-                                         shape=(feature_num, n_hidden), testval=init_1)
+            weights_in_1_grp = pm.Normal(
+                "w_in_1_grp", 0, sd=1, shape=(feature_num, n_hidden), testval=init_1
+            )
 
             # Group standard deviation:
-            weights_in_1_grp_sd = pm.HalfCauchy('w_in_1_grp_sd', 1.,
-                                                shape=(feature_num, n_hidden), testval=std_init_1)
+            weights_in_1_grp_sd = pm.HalfCauchy(
+                "w_in_1_grp_sd", 1.0, shape=(feature_num, n_hidden), testval=std_init_1
+            )
 
             if n_layers == 2:
                 # Group the mean distribution for hidden layer 1 to hidden layer 2:
-                weights_1_2_grp = pm.Normal('w_1_2_grp', 0, sd=1,
-                                            shape=(n_hidden, n_hidden), testval=init_2)
+                weights_1_2_grp = pm.Normal(
+                    "w_1_2_grp", 0, sd=1, shape=(n_hidden, n_hidden), testval=init_2
+                )
 
                 # Group standard deviation:
-                weights_1_2_grp_sd = pm.HalfCauchy('w_1_2_grp_sd', 1.,
-                                                   shape=(n_hidden, n_hidden), testval=std_init_2)
+                weights_1_2_grp_sd = pm.HalfCauchy(
+                    "w_1_2_grp_sd", 1.0, shape=(n_hidden, n_hidden), testval=std_init_2
+                )
 
             # Group the mean distribution for hidden to output:
-            weights_2_out_grp = pm.Normal('w_2_out_grp', 0, sd=1,
-                                          shape=(n_hidden,), testval=init_out)
+            weights_2_out_grp = pm.Normal(
+                "w_2_out_grp", 0, sd=1, shape=(n_hidden,), testval=init_out
+            )
 
             # Group standard deviation:
-            weights_2_out_grp_sd = pm.HalfCauchy('w_2_out_grp_sd', 1.,
-                                                 shape=(n_hidden,), testval=std_init_out)
+            weights_2_out_grp_sd = pm.HalfCauchy(
+                "w_2_out_grp_sd", 1.0, shape=(n_hidden,), testval=std_init_out
+            )
 
             # mu_prior_intercept = pm.Uniform('mu_prior_intercept', lower=-100, upper=100)
-            mu_prior_intercept = pm.Normal('mu_prior_intercept', mu=0., sigma=1e3)
-            sigma_prior_intercept = pm.HalfCauchy('sigma_prior_intercept', 5)
+            mu_prior_intercept = pm.Normal("mu_prior_intercept", mu=0.0, sigma=1e3)
+            sigma_prior_intercept = pm.HalfCauchy("sigma_prior_intercept", 5)
 
         # Now create separate weights for each group, by doing
         # weights * group_sd + group_mean, we make sure the new weights are
         # coming from the (group_mean, group_sd) distribution.
-        weights_in_1_raw = pm.Normal('w_in_1', 0, sd=1,
-                                     shape=(batch_effects_size + [feature_num, n_hidden]))
+        weights_in_1_raw = pm.Normal(
+            "w_in_1", 0, sd=1, shape=(batch_effects_size + [feature_num, n_hidden])
+        )
         weights_in_1 = weights_in_1_raw * weights_in_1_grp_sd + weights_in_1_grp
 
         if n_layers == 2:
-            weights_1_2_raw = pm.Normal('w_1_2', 0, sd=1,
-                                        shape=(batch_effects_size + [n_hidden, n_hidden]))
+            weights_1_2_raw = pm.Normal(
+                "w_1_2", 0, sd=1, shape=(batch_effects_size + [n_hidden, n_hidden])
+            )
             weights_1_2 = weights_1_2_raw * weights_1_2_grp_sd + weights_1_2_grp
 
-        weights_2_out_raw = pm.Normal('w_2_out', 0, sd=1,
-                                      shape=(batch_effects_size + [n_hidden]))
+        weights_2_out_raw = pm.Normal(
+            "w_2_out", 0, sd=1, shape=(batch_effects_size + [n_hidden])
+        )
         weights_2_out = weights_2_out_raw * weights_2_out_grp_sd + weights_2_out_grp
 
-        intercepts_offset = pm.Normal('intercepts_offset', mu=0, sd=1,
-                                      shape=(batch_effects_size))
+        intercepts_offset = pm.Normal(
+            "intercepts_offset", mu=0, sd=1, shape=(batch_effects_size)
+        )
 
-        intercepts = pm.Deterministic('intercepts', intercepts_offset +
-                                      mu_prior_intercept * sigma_prior_intercept)
+        intercepts = pm.Deterministic(
+            "intercepts", intercepts_offset + mu_prior_intercept * sigma_prior_intercept
+        )
 
         # Build the neural network and estimate y_hat:
-        y_hat = theano.tensor.zeros(y.shape)
+        y_hat = pytensor.tensor.zeros(y.shape)
         for be in be_idx:
             # Find the indices corresponding to 'group be':
             a = []
@@ -776,86 +1004,147 @@ def nn_hbr(X, y, batch_effects, batch_effects_size, configs, trace=None):
                 a.append(batch_effects[:, i] == b)
             idx = reduce(np.logical_and, a).nonzero()
             if idx[0].shape[0] != 0:
-                act_1 = pm.math.tanh(theano.tensor.dot(X[idx, :], weights_in_1[be]))
+                act_1 = pm.math.tanh(pytensor.tensor.dot(X[idx, :], weights_in_1[be]))
                 if n_layers == 2:
-                    act_2 = pm.math.tanh(theano.tensor.dot(act_1, weights_1_2[be]))
-                    y_hat = theano.tensor.set_subtensor(y_hat[idx, 0],
-                                                        intercepts[be] + theano.tensor.dot(act_2, weights_2_out[be]))
+                    act_2 = pm.math.tanh(pytensor.tensor.dot(act_1, weights_1_2[be]))
+                    y_hat = pytensor.tensor.set_subtensor(
+                        y_hat[idx, 0],
+                        intercepts[be] + pytensor.tensor.dot(act_2, weights_2_out[be]),
+                    )
                 else:
-                    y_hat = theano.tensor.set_subtensor(y_hat[idx, 0],
-                                                        intercepts[be] + theano.tensor.dot(act_1, weights_2_out[be]))
+                    y_hat = pytensor.tensor.set_subtensor(
+                        y_hat[idx, 0],
+                        intercepts[be] + pytensor.tensor.dot(act_1, weights_2_out[be]),
+                    )
 
         # If we want to estimate varying noise terms across groups:
-        if configs['random_noise']:
-            if configs['hetero_noise']:
-                if trace is not None:  # # Used when estimating/predicting on a new site
-                    weights_in_1_grp_noise = from_posterior('w_in_1_grp_noise',
-                                                            trace['w_in_1_grp_noise'],
-                                                            distribution='normal', freedom=configs['freedom'])
-
-                    weights_in_1_grp_sd_noise = from_posterior('w_in_1_grp_sd_noise',
-                                                               trace['w_in_1_grp_sd_noise'],
-                                                               distribution='hcauchy', freedom=configs['freedom'])
+        if configs["random_noise"]:
+            if configs["hetero_noise"]:
+                if idata is not None:  # # Used when estimating/predicting on a new site
+                    weights_in_1_grp_noise = from_posterior(
+                        "w_in_1_grp_noise",
+                        idata["w_in_1_grp_noise"],
+                        distribution="normal",
+                        freedom=configs["freedom"],
+                    )
+
+                    weights_in_1_grp_sd_noise = from_posterior(
+                        "w_in_1_grp_sd_noise",
+                        idata["w_in_1_grp_sd_noise"],
+                        distribution="hcauchy",
+                        freedom=configs["freedom"],
+                    )
 
                     if n_layers == 2:
-                        weights_1_2_grp_noise = from_posterior('w_1_2_grp_noise',
-                                                               trace['w_1_2_grp_noise'],
-                                                               distribution='normal', freedom=configs['freedom'])
-
-                        weights_1_2_grp_sd_noise = from_posterior('w_1_2_grp_sd_noise',
-                                                                  trace['w_1_2_grp_sd_noise'],
-                                                                  distribution='hcauchy', freedom=configs['freedom'])
-
-                    weights_2_out_grp_noise = from_posterior('w_2_out_grp_noise',
-                                                             trace['w_2_out_grp_noise'],
-                                                             distribution='normal', freedom=configs['freedom'])
-
-                    weights_2_out_grp_sd_noise = from_posterior('w_2_out_grp_sd_noise',
-                                                                trace['w_2_out_grp_sd_noise'],
-                                                                distribution='hcauchy', freedom=configs['freedom'])
+                        weights_1_2_grp_noise = from_posterior(
+                            "w_1_2_grp_noise",
+                            idata["w_1_2_grp_noise"],
+                            distribution="normal",
+                            freedom=configs["freedom"],
+                        )
+
+                        weights_1_2_grp_sd_noise = from_posterior(
+                            "w_1_2_grp_sd_noise",
+                            idata["w_1_2_grp_sd_noise"],
+                            distribution="hcauchy",
+                            freedom=configs["freedom"],
+                        )
+
+                    weights_2_out_grp_noise = from_posterior(
+                        "w_2_out_grp_noise",
+                        idata["w_2_out_grp_noise"],
+                        distribution="normal",
+                        freedom=configs["freedom"],
+                    )
+
+                    weights_2_out_grp_sd_noise = from_posterior(
+                        "w_2_out_grp_sd_noise",
+                        idata["w_2_out_grp_sd_noise"],
+                        distribution="hcauchy",
+                        freedom=configs["freedom"],
+                    )
 
                 else:
                     # The input layer to the first hidden layer:
-                    weights_in_1_grp_noise = pm.Normal('w_in_1_grp_noise', 0, sd=1,
-                                                       shape=(feature_num, n_hidden),
-                                                       testval=init_1_noise)
-                    weights_in_1_grp_sd_noise = pm.HalfCauchy('w_in_1_grp_sd_noise', 1,
-                                                              shape=(feature_num, n_hidden),
-                                                              testval=std_init_1_noise)
+                    weights_in_1_grp_noise = pm.Normal(
+                        "w_in_1_grp_noise",
+                        0,
+                        sd=1,
+                        shape=(feature_num, n_hidden),
+                        testval=init_1_noise,
+                    )
+                    weights_in_1_grp_sd_noise = pm.HalfCauchy(
+                        "w_in_1_grp_sd_noise",
+                        1,
+                        shape=(feature_num, n_hidden),
+                        testval=std_init_1_noise,
+                    )
 
                     # The first hidden layer to second hidden layer:
                     if n_layers == 2:
-                        weights_1_2_grp_noise = pm.Normal('w_1_2_grp_noise', 0, sd=1,
-                                                          shape=(n_hidden, n_hidden),
-                                                          testval=init_2_noise)
-                        weights_1_2_grp_sd_noise = pm.HalfCauchy('w_1_2_grp_sd_noise', 1,
-                                                                 shape=(n_hidden, n_hidden),
-                                                                 testval=std_init_2_noise)
+                        weights_1_2_grp_noise = pm.Normal(
+                            "w_1_2_grp_noise",
+                            0,
+                            sd=1,
+                            shape=(n_hidden, n_hidden),
+                            testval=init_2_noise,
+                        )
+                        weights_1_2_grp_sd_noise = pm.HalfCauchy(
+                            "w_1_2_grp_sd_noise",
+                            1,
+                            shape=(n_hidden, n_hidden),
+                            testval=std_init_2_noise,
+                        )
 
                     # The second hidden layer to output layer:
-                    weights_2_out_grp_noise = pm.Normal('w_2_out_grp_noise', 0, sd=1,
-                                                        shape=(n_hidden,),
-                                                        testval=init_out_noise)
-                    weights_2_out_grp_sd_noise = pm.HalfCauchy('w_2_out_grp_sd_noise', 1,
-                                                               shape=(n_hidden,),
-                                                               testval=std_init_out_noise)
+                    weights_2_out_grp_noise = pm.Normal(
+                        "w_2_out_grp_noise",
+                        0,
+                        sd=1,
+                        shape=(n_hidden,),
+                        testval=init_out_noise,
+                    )
+                    weights_2_out_grp_sd_noise = pm.HalfCauchy(
+                        "w_2_out_grp_sd_noise",
+                        1,
+                        shape=(n_hidden,),
+                        testval=std_init_out_noise,
+                    )
 
                     # mu_prior_intercept_noise = pm.HalfNormal('mu_prior_intercept_noise', sigma=1e3)
                     # sigma_prior_intercept_noise = pm.HalfCauchy('sigma_prior_intercept_noise', 5)
 
                 # Now create separate weights for each group:
-                weights_in_1_raw_noise = pm.Normal('w_in_1_noise', 0, sd=1,
-                                                   shape=(batch_effects_size + [feature_num, n_hidden]))
-                weights_in_1_noise = weights_in_1_raw_noise * weights_in_1_grp_sd_noise + weights_in_1_grp_noise
+                weights_in_1_raw_noise = pm.Normal(
+                    "w_in_1_noise",
+                    0,
+                    sd=1,
+                    shape=(batch_effects_size + [feature_num, n_hidden]),
+                )
+                weights_in_1_noise = (
+                    weights_in_1_raw_noise * weights_in_1_grp_sd_noise
+                    + weights_in_1_grp_noise
+                )
 
                 if n_layers == 2:
-                    weights_1_2_raw_noise = pm.Normal('w_1_2_noise', 0, sd=1,
-                                                      shape=(batch_effects_size + [n_hidden, n_hidden]))
-                    weights_1_2_noise = weights_1_2_raw_noise * weights_1_2_grp_sd_noise + weights_1_2_grp_noise
-
-                weights_2_out_raw_noise = pm.Normal('w_2_out_noise', 0, sd=1,
-                                                    shape=(batch_effects_size + [n_hidden]))
-                weights_2_out_noise = weights_2_out_raw_noise * weights_2_out_grp_sd_noise + weights_2_out_grp_noise
+                    weights_1_2_raw_noise = pm.Normal(
+                        "w_1_2_noise",
+                        0,
+                        sd=1,
+                        shape=(batch_effects_size + [n_hidden, n_hidden]),
+                    )
+                    weights_1_2_noise = (
+                        weights_1_2_raw_noise * weights_1_2_grp_sd_noise
+                        + weights_1_2_grp_noise
+                    )
+
+                weights_2_out_raw_noise = pm.Normal(
+                    "w_2_out_noise", 0, sd=1, shape=(batch_effects_size + [n_hidden])
+                )
+                weights_2_out_noise = (
+                    weights_2_out_raw_noise * weights_2_out_grp_sd_noise
+                    + weights_2_out_grp_noise
+                )
 
                 # intercepts_offset_noise = pm.Normal('intercepts_offset_noise', mu=0, sd=1,
                 #                          shape=(batch_effects_size))
@@ -864,65 +1153,98 @@ def nn_hbr(X, y, batch_effects, batch_effects_size, configs, trace=None):
                 #                      intercepts_offset_noise * sigma_prior_intercept_noise)
 
                 # Build the neural network and estimate the sigma_y:
-                sigma_y = theano.tensor.zeros(y.shape)
+                sigma_y = pytensor.tensor.zeros(y.shape)
                 for be in be_idx:
                     a = []
                     for i, b in enumerate(be):
                         a.append(batch_effects[:, i] == b)
                     idx = reduce(np.logical_and, a).nonzero()
                     if idx[0].shape[0] != 0:
-                        act_1_noise = pm.math.sigmoid(theano.tensor.dot(X[idx, :], weights_in_1_noise[be]))
+                        act_1_noise = pm.math.sigmoid(
+                            pytensor.tensor.dot(X[idx, :], weights_in_1_noise[be])
+                        )
                         if n_layers == 2:
-                            act_2_noise = pm.math.sigmoid(theano.tensor.dot(act_1_noise, weights_1_2_noise[be]))
-                            temp = pm.math.log1pexp(theano.tensor.dot(act_2_noise, weights_2_out_noise[be])) + 1e-5
+                            act_2_noise = pm.math.sigmoid(
+                                pytensor.tensor.dot(act_1_noise, weights_1_2_noise[be])
+                            )
+                            temp = (
+                                pm.math.log1pexp(
+                                    pytensor.tensor.dot(
+                                        act_2_noise, weights_2_out_noise[be]
+                                    )
+                                )
+                                + 1e-5
+                            )
                         else:
-                            temp = pm.math.log1pexp(theano.tensor.dot(act_1_noise, weights_2_out_noise[be])) + 1e-5
-                        sigma_y = theano.tensor.set_subtensor(sigma_y[idx, 0], temp)
+                            temp = (
+                                pm.math.log1pexp(
+                                    pytensor.tensor.dot(
+                                        act_1_noise, weights_2_out_noise[be]
+                                    )
+                                )
+                                + 1e-5
+                            )
+                        sigma_y = pytensor.tensor.set_subtensor(sigma_y[idx, 0], temp)
 
             else:  # homoscedastic noise:
-                if trace is not None:  # Used for transferring the priors
-                    upper_bound = np.percentile(trace['sigma_noise'], 95)
-                    sigma_noise = pm.Uniform('sigma_noise', lower=0, upper=2 * upper_bound, shape=(batch_effects_size))
+                if idata is not None:  # Used for transferring the priors
+                    upper_bound = np.percentile(idata["sigma_noise"], 95)
+                    sigma_noise = pm.Uniform(
+                        "sigma_noise",
+                        lower=0,
+                        upper=2 * upper_bound,
+                        shape=(batch_effects_size),
+                    )
                 else:
-                    sigma_noise = pm.Uniform('sigma_noise', lower=0, upper=100, shape=(batch_effects_size))
+                    sigma_noise = pm.Uniform(
+                        "sigma_noise", lower=0, upper=100, shape=(batch_effects_size)
+                    )
 
-                sigma_y = theano.tensor.zeros(y.shape)
+                sigma_y = pytensor.tensor.zeros(y.shape)
                 for be in be_idx:
                     a = []
                     for i, b in enumerate(be):
                         a.append(batch_effects[:, i] == b)
                     idx = reduce(np.logical_and, a).nonzero()
                     if idx[0].shape[0] != 0:
-                        sigma_y = theano.tensor.set_subtensor(sigma_y[idx, 0], sigma_noise[be])
+                        sigma_y = pytensor.tensor.set_subtensor(
+                            sigma_y[idx, 0], sigma_noise[be]
+                        )
 
         else:  # do not allow for random noise terms across groups:
-            if trace is not None:  # Used for transferring the priors
-                upper_bound = np.percentile(trace['sigma_noise'], 95)
-                sigma_noise = pm.Uniform('sigma_noise', lower=0, upper=2 * upper_bound)
+            if idata is not None:  # Used for transferring the priors
+                upper_bound = np.percentile(idata["sigma_noise"], 95)
+                sigma_noise = pm.Uniform("sigma_noise", lower=0, upper=2 * upper_bound)
             else:
-                sigma_noise = pm.Uniform('sigma_noise', lower=0, upper=100)
-            sigma_y = theano.tensor.zeros(y.shape)
+                sigma_noise = pm.Uniform("sigma_noise", lower=0, upper=100)
+            sigma_y = pytensor.tensor.zeros(y.shape)
             for be in be_idx:
                 a = []
                 for i, b in enumerate(be):
                     a.append(batch_effects[:, i] == b)
                 idx = reduce(np.logical_and, a).nonzero()
                 if idx[0].shape[0] != 0:
-                    sigma_y = theano.tensor.set_subtensor(sigma_y[idx, 0], sigma_noise)
-
-        if configs['skewed_likelihood']:
-            skewness = pm.Uniform('skewness', lower=-10, upper=10, shape=(batch_effects_size))
-            alpha = theano.tensor.zeros(y.shape)
+                    sigma_y = pytensor.tensor.set_subtensor(
+                        sigma_y[idx, 0], sigma_noise
+                    )
+
+        if configs["skewed_likelihood"]:
+            skewness = pm.Uniform(
+                "skewness", lower=-10, upper=10, shape=(batch_effects_size)
+            )
+            alpha = pytensor.tensor.zeros(y.shape)
             for be in be_idx:
                 a = []
                 for i, b in enumerate(be):
                     a.append(batch_effects[:, i] == b)
                 idx = reduce(np.logical_and, a).nonzero()
                 if idx[0].shape[0] != 0:
-                    alpha = theano.tensor.set_subtensor(alpha[idx, 0], skewness[be])
+                    alpha = pytensor.tensor.set_subtensor(alpha[idx, 0], skewness[be])
         else:
             alpha = 0  # symmetrical normal distribution
 
-        y_like = pm.SkewNormal('y_like', mu=y_hat, sigma=sigma_y, alpha=alpha, observed=y)
+        y_like = pm.SkewNormal(
+            "y_like", mu=y_hat, sigma=sigma_y, alpha=alpha, observed=y
+        )
 
     return model
diff --git a/pcntoolkit/normative.py b/pcntoolkit/normative.py
index d4a8cbed..275f8e99 100755
--- a/pcntoolkit/normative.py
+++ b/pcntoolkit/normative.py
@@ -217,7 +217,7 @@ def evaluate(Y, Yhat, S2=None, mY=None, sY=None, nlZ=None, nm=None, Xz_tr=None,
     
     return results
 
-def save_results(respfile, Yhat, S2, maskvol, Z=None, outputsuffix=None, 
+def save_results(respfile, Yhat, S2, maskvol, Z=None, Y=None, outputsuffix=None, 
                  results=None, save_path=''):
     
     print("Writing outputs ...")
@@ -244,7 +244,9 @@ def save_results(respfile, Yhat, S2, maskvol, Z=None, outputsuffix=None,
     if Z is not None:
         fileio.save(Z, os.path.join(save_path, 'Z' + ext), example=exfile, 
                     mask=maskvol)
-
+    if Y is not None:
+        fileio.save(Y, os.path.join(save_path, 'Y' + ext), example=exfile, 
+                    mask=maskvol)
     if results is not None:        
         for metric in list(results.keys()):
             if (metric == 'NLL' or metric == 'BIC') and file_ext == '.nii.gz':
@@ -670,12 +672,14 @@ def predict(covfile, respfile, maskfile=None, **kwargs):
     :param job_id: batch id
     :param fold: which cross-validation fold to use (default = 0)
     :param fold: list of model IDs to predict (if not specified all are computed)
+    :param return_y: return the (transformed) response variable (default = False)
 
     All outputs are written to disk in the same format as the input. These are:
 
     :outputs: * Yhat - predictive mean
               * S2 - predictive variance
               * Z - Z scores
+              * Y - response variable (if return_y is True)
     '''
     
     
@@ -689,6 +693,7 @@ def predict(covfile, respfile, maskfile=None, **kwargs):
     alg = kwargs.pop('alg')
     fold = kwargs.pop('fold',0)
     models = kwargs.pop('models', None)
+    return_y = kwargs.pop('return_y', False)
     
     if alg == 'gpr':
         raise(ValueError, "gpr is not supported with predict()")
@@ -816,10 +821,15 @@ def predict(covfile, respfile, maskfile=None, **kwargs):
                            metrics = ['Rho', 'RMSE', 'SMSE', 'EXPV'])
         
         print("Evaluations Writing outputs ...")
-        save_results(respfile, Yhat, S2, maskvol, Z=Z, 
-                     outputsuffix=outputsuffix, results=results)
         
-        return (Yhat, S2, Z)
+        if return_y:
+            save_results(respfile, Yhat, S2, maskvol, Z=Z, Y=Y,
+                     outputsuffix=outputsuffix, results=results)
+            return (Yhat, S2, Z, Y)
+        else:
+            save_results(respfile, Yhat, S2, maskvol, Z=Z, 
+                     outputsuffix=outputsuffix, results=results)
+            return (Yhat, S2, Z)
 
     
 def transfer(covfile, respfile, testcov=None, testresp=None, maskfile=None, 
@@ -1437,4 +1447,4 @@ def main(*args):
 
 # For running from the command line:
 if __name__ == "__main__":
-    main(sys.argv[1:])
+    main(sys.argv[1:])
\ No newline at end of file
diff --git a/pcntoolkit/normative_model/norm_hbr.py b/pcntoolkit/normative_model/norm_hbr.py
index d769163e..f3b145ac 100644
--- a/pcntoolkit/normative_model/norm_hbr.py
+++ b/pcntoolkit/normative_model/norm_hbr.py
@@ -9,12 +9,18 @@
 
 from __future__ import print_function
 from __future__ import division
+from sys import exit
 
+from itertools import product
 
 import os
 import warnings
 import sys
+
+import xarray
+import arviz as az
 import numpy as np
+from scipy import special as spp
 from ast import literal_eval as make_tuple
 
 try:
@@ -34,19 +40,19 @@
 
 
 class NormHBR(NormBase):
-        
-    """ HBR multi-batch normative modelling class. By default, this function 
-    estimates a linear model with random intercept, random slope, and random 
+
+    """HBR multi-batch normative modelling class. By default, this function
+    estimates a linear model with random intercept, random slope, and random
     homoscedastic noise.
-    
+
     :param X: [N×P] array of clinical covariates
     :param y: [N×1] array of neuroimaging measures
     :param trbefile: the address to the batch effects file for the training set.
         the batch effect array should be a [N×M] array where M is the number of
         the type of batch effects. For example when the site and gender is modeled
-        as batch effects M=2. Each column in the batch effect array contains the 
+        as batch effects M=2. Each column in the batch effect array contains the
         batch ID (starting from 0) for each sample. If not specified (default=None)
-        then all samples assumed to be from the same batch (i.e., the batch effect 
+        then all samples assumed to be from the same batch (i.e., the batch effect
                                                             is not modelled).
     :param tsbefile: Similar to trbefile for the test set.
     :param model_type: Specifies the type of the model from 'linear', 'plynomial',
@@ -59,146 +65,179 @@ class NormHBR(NormBase):
         parametrized on a linear function (heteroscedastic noise) or it is fixed for
         each batch (homoscedastic noise).
     :param linear_epsilon: Boolean (default='False') to decide the parametrization
-        of epsilon for the SHASH likelihood that controls its skewness. 
-        If True, epsilon is  parametrized on a linear function 
+        of epsilon for the SHASH likelihood that controls its skewness.
+        If True, epsilon is  parametrized on a linear function
         (thus changes with covariates) otherwise it is fixed for each batch.
     :param linear_delta: Boolean (default='False') to decide the parametrization
-        of delta for the SHASH likelihood that controls its kurtosis. 
-        If True, delta is  parametrized on a linear function 
+        of delta for the SHASH likelihood that controls its kurtosis.
+        If True, delta is  parametrized on a linear function
         (thus changes with covariates) otherwise it is fixed for each batch.
-    :param random_intercept_{parameter}: if parameters mu (default='True'), 
+    :param random_intercept_{parameter}: if parameters mu (default='True'),
         sigma (default='False'), epsilon (default='False'), and delta (default='False')
-        are parametrized on a linear function, then this boolean decides 
+        are parametrized on a linear function, then this boolean decides
         whether the intercept can vary across batches.
-    :param random_slope_{parameter}: if parameters mu (default='True'), 
+    :param random_slope_{parameter}: if parameters mu (default='True'),
         sigma (default='False'), epsilon (default='False'), and delta (default='False')
-        are parametrized on a linear function, then this boolean decides 
+        are parametrized on a linear function, then this boolean decides
         whether the slope can vary across batches.
-    :param centered_intercept_{parameter}: if parameters mu (default='False'), 
+    :param centered_intercept_{parameter}: if parameters mu (default='False'),
         sigma (default='False'), epsilon (default='False'), and delta (default='False')
-        are parametrized on a linear function, then this boolean decides 
+        are parametrized on a linear function, then this boolean decides
         whether the parameters of intercept are estimated in a centered or
         non-centered manner (default). While centered estimation runs faster
         it may cause some problems for the sampler (the funnel of hell).
-    :param centered_slope_{parameter}: if parameters mu (default='False'), 
+    :param centered_slope_{parameter}: if parameters mu (default='False'),
         sigma (default='False'), epsilon (default='False'), and delta (default='False')
-        are parametrized on a linear function, then this boolean decides 
+        are parametrized on a linear function, then this boolean decides
         whether the parameters of slope are estimated in a centered or
         non-centered manner (default). While centered estimation runs faster
         it may cause some problems for the sampler (the funnel of hell).
     :param sampler: specifies the type of PyMC sampler (Defauls is 'NUTS').
-    :param n_samples: The number of samples to draw (Default is '1000'). Please 
-        note that this parameter must be specified in a string fromat ('1000' and 
-                                                                       not 1000).  
-    :param n_tuning: String that specifies the number of iterations to adjust 
+    :param n_samples: The number of samples to draw (Default is '1000'). Please
+        note that this parameter must be specified in a string fromat ('1000' and
+                                                                       not 1000).
+    :param n_tuning: String that specifies the number of iterations to adjust
         the samplers's step sizes, scalings or similar (defauls is '500').
-    :param n_chains: String that specifies the number of chains to sample. Defauls 
-        is '1' for faster estimation, but note that sampling independent chains 
+    :param n_chains: String that specifies the number of chains to sample. Defauls
+        is '1' for faster estimation, but note that sampling independent chains
         is important for some convergence checks.
     :param cores: String that specifies the number of chains to run in parallel.
         (defauls is '1').
     :param Initialization method to use for auto-assigned NUTS samplers. The
-        defauls is 'jitter+adapt_diag' that starts with a identity mass matrix 
+        defauls is 'jitter+adapt_diag' that starts with a identity mass matrix
         and then adapt a diagonal based on the variance of the tuning samples
         while adding a uniform jitter in [-1, 1] to the starting point in each chain.
-    :param target_accept: String that of a float in [0, 1] that regulates the 
+    :param target_accept: String that of a float in [0, 1] that regulates the
         step size such that we approximate this acceptance rate. The defauls is '0.8'
         but higher values like 0.9 or 0.95 often work better for problematic posteriors.
     :param order: String that defines the order of bspline or polynomial model.
         The defauls is '3'.
     :param nknots: String that defines the numbers of knots for the bspline model.
-        The defauls is '5'. Higher values increase the model complexity with negative 
+        The defauls is '5'. Higher values increase the model complexity with negative
         effect on the spped of estimations.
     :param nn_hidden_layers_num: String the specifies the number of hidden layers
         in neural network model. It can be either '1' or '2'. The default is set to '2'.
-    :param nn_hidden_neuron_num: String that specifies the number of neurons in 
+    :param nn_hidden_neuron_num: String that specifies the number of neurons in
         the hidden layers. The defauls is set to  '2'.
-        
+
     Written by S.de Boer and S.M. Kia
-    
+
     """
 
     def __init__(self, **kwargs):
-
         self.configs = dict()
         # inputs
-        self.configs['trbefile'] = kwargs.pop('trbefile', None)
-        self.configs['tsbefile'] = kwargs.pop('tsbefile', None)
+        self.configs["trbefile"] = kwargs.get("trbefile", None)
+        self.configs["tsbefile"] = kwargs.get("tsbefile", None)
         # Model settings
-        self.configs['type'] = kwargs.pop('model_type', 'linear')
-        self.configs['random_noise'] = kwargs.pop('random_noise', 'True') == 'True'
-        self.configs['likelihood'] = kwargs.pop('likelihood', 'Normal')
+        self.configs["type"] = kwargs.get("model_type", "linear")
+        self.configs["random_noise"] = kwargs.get("random_noise", "True") == "True"
+        self.configs["likelihood"] = kwargs.get("likelihood", "Normal")
         # sampler settings
-        self.configs['n_samples'] = int(kwargs.pop('n_samples', '1000'))
-        self.configs['n_tuning'] = int(kwargs.pop('n_tuning', '500'))
-        self.configs['n_chains'] = int(kwargs.pop('n_chains', '1'))
-        self.configs['sampler'] = kwargs.pop('sampler', 'NUTS')
-        self.configs['target_accept'] = float(kwargs.pop('target_accept', '0.8'))
-        self.configs['init'] = kwargs.pop('init', 'jitter+adapt_diag')
-        self.configs['cores'] = int(kwargs.pop('cores', '1'))
+        self.configs["n_samples"] = int(kwargs.get("n_samples", "1000"))
+        self.configs["n_tuning"] = int(kwargs.get("n_tuning", "500"))
+        self.configs["n_chains"] = int(kwargs.get("n_chains", "1"))
+        self.configs["sampler"] = kwargs.get("sampler", "NUTS")
+        self.configs["target_accept"] = float(kwargs.get("target_accept", "0.8"))
+        self.configs["init"] = kwargs.get("init", "jitter+adapt_diag")
+        self.configs["cores"] = int(kwargs.get("cores", "1"))
+        self.configs["remove_datapoints_from_posterior"] = kwargs.get("remove_datapoints_from_posterior","True") == "True"
         # model transfer setting
-        self.configs['freedom'] = int(kwargs.pop('freedom', '1'))
-        self.configs['transferred'] = False
+        self.configs["freedom"] = int(kwargs.get("freedom", "1"))
+        self.configs["transferred"] = False
         # deprecated settings
-        self.configs['skewed_likelihood'] = kwargs.pop('skewed_likelihood', 'False') == 'True'
+        self.configs["skewed_likelihood"] = (
+            kwargs.get("skewed_likelihood", "False") == "True"
+        )
         # misc
-        self.configs['pred_type'] = kwargs.pop('pred_type', 'single')
-
-        if self.configs['type'] == 'bspline':
-            self.configs['order'] = int(kwargs.pop('order', '3'))
-            self.configs['nknots'] = int(kwargs.pop('nknots', '5'))
-        elif self.configs['type'] == 'polynomial':
-            self.configs['order'] = int(kwargs.pop('order', '3'))
-        elif self.configs['type'] == 'nn':
-            self.configs['nn_hidden_neuron_num'] = int(kwargs.pop('nn_hidden_neuron_num', '2'))
-            self.configs['nn_hidden_layers_num'] = int(kwargs.pop('nn_hidden_layers_num', '2'))
-            if self.configs['nn_hidden_layers_num'] > 2:
-                raise ValueError("Using " + str(self.configs['nn_hidden_layers_num']) \
-                                 + " layers was not implemented. The number of " \
-                                 + " layers has to be less than 3.")
-        elif self.configs['type'] == 'linear':
+        self.configs["pred_type"] = kwargs.get("pred_type", "single")
+
+        if self.configs["type"] == "bspline":
+            self.configs["order"] = int(kwargs.get("order", "3"))
+            self.configs["nknots"] = int(kwargs.get("nknots", "5"))
+        elif self.configs["type"] == "polynomial":
+            self.configs["order"] = int(kwargs.get("order", "3"))
+        elif self.configs["type"] == "nn":
+            self.configs["nn_hidden_neuron_num"] = int(
+                kwargs.get("nn_hidden_neuron_num", "2")
+            )
+            self.configs["nn_hidden_layers_num"] = int(
+                kwargs.get("nn_hidden_layers_num", "2")
+            )
+            if self.configs["nn_hidden_layers_num"] > 2:
+                raise ValueError(
+                    "Using "
+                    + str(self.configs["nn_hidden_layers_num"])
+                    + " layers was not implemented. The number of "
+                    + " layers has to be less than 3."
+                )
+        elif self.configs["type"] == "linear":
             pass
         else:
-            raise ValueError("Unknown model type, please specify from 'linear', \
-                             'polynomial', 'bspline', or 'nn'.")
-
-        if self.configs['type'] in ['bspline', 'polynomial', 'linear']:
+            raise ValueError(
+                "Unknown model type, please specify from 'linear', \
+                             'polynomial', 'bspline', or 'nn'."
+            )
 
-            for p in ['mu', 'sigma', 'epsilon', 'delta']:
-                self.configs[f'linear_{p}'] = kwargs.pop(f'linear_{p}', 'False') == 'True'
+        if self.configs["type"] in ["bspline", "polynomial", "linear"]:
+            for p in ["mu", "sigma", "epsilon", "delta"]:
+                self.configs[f"linear_{p}"] = (
+                    kwargs.get(f"linear_{p}", "False") == "True"
+                )
 
                 ######## Deprecations (remove in later version)
-                if f'{p}_linear' in kwargs.keys():
-                    print(f'The keyword \'{p}_linear\' is deprecated. It is now automatically replaced with \'linear_{p}\'')
-                    self.configs[f'linear_{p}'] = kwargs.pop(f'{p}_linear', 'False') == 'True'
-                ##### End Deprecations 
-
-                for c in ['centered','random']:
-                    self.configs[f'{c}_{p}'] = kwargs.pop(f'{c}_{p}', 'False') == 'True'
-                    for sp in ['slope','intercept']:
-                        self.configs[f'{c}_{sp}_{p}'] = kwargs.pop(f'{c}_{sp}_{p}', 'False') == 'True'
+                if f"{p}_linear" in kwargs.keys():
+                    print(
+                        f"The keyword '{p}_linear' is deprecated. It is now automatically replaced with 'linear_{p}'"
+                    )
+                    self.configs[f"linear_{p}"] = (
+                        kwargs.get(f"{p}_linear", "False") == "True"
+                    )
+                ##### End Deprecations
+
+                for c in ["centered", "random"]:
+                    self.configs[f"{c}_{p}"] = kwargs.get(f"{c}_{p}", "False") == "True"
+                    for sp in ["slope", "intercept"]:
+                        self.configs[f"{c}_{sp}_{p}"] = (
+                            kwargs.get(f"{c}_{sp}_{p}", "False") == "True"
+                        )
 
             ######## Deprecations (remove in later version)
-            if self.configs['linear_sigma']:
-                if 'random_noise' in kwargs.keys():
-                    print("The keyword \'random_noise\' is deprecated. It is now automatically replaced with \'random_intercept_sigma\', because sigma is linear")
-                    self.configs['random_intercept_sigma'] = kwargs.pop('random_noise','True') == 'True'
-            elif 'random_noise' in kwargs.keys():
-                print("The keyword \'random_noise\' is deprecated. It is now automatically replaced with \'random_sigma\', because sigma is fixed")
-                self.configs['random_sigma'] = kwargs.pop('random_noise','True') == 'True'
-            if 'random_slope' in kwargs.keys():
-                print("The keyword \'random_slope\' is deprecated. It is now automatically replaced with \'random_intercept_mu\'")
-                self.configs['random_slope_mu'] = kwargs.pop('random_slope','True') == 'True'
-            ##### End Deprecations 
-
+            if self.configs["linear_sigma"]:
+                if "random_noise" in kwargs.keys():
+                    print(
+                        "The keyword 'random_noise' is deprecated. It is now automatically replaced with 'random_intercept_sigma', because sigma is linear"
+                    )
+                    self.configs["random_intercept_sigma"] = (
+                        kwargs.get("random_noise", "True") == "True"
+                    )
+            elif "random_noise" in kwargs.keys():
+                print(
+                    "The keyword 'random_noise' is deprecated. It is now automatically replaced with 'random_sigma', because sigma is fixed"
+                )
+                self.configs["random_sigma"] = (
+                    kwargs.get("random_noise", "True") == "True"
+                )
+            if "random_slope" in kwargs.keys():
+                print(
+                    "The keyword 'random_slope' is deprecated. It is now automatically replaced with 'random_intercept_mu'"
+                )
+                self.configs["random_slope_mu"] = (
+                    kwargs.get("random_slope", "True") == "True"
+                )
+            ##### End Deprecations
 
         ## Default parameters
-        self.configs['linear_mu'] = kwargs.pop('linear_mu','True') == 'True'
-        self.configs['random_mu'] = kwargs.pop('random_mu','True') == 'True'
-        self.configs['random_intercept_mu'] = kwargs.pop('random_intercept_mu','True') == 'True'
-        self.configs['random_slope_mu'] = kwargs.pop('random_slope_mu','True') == 'True'
-        self.configs['random_sigma'] = kwargs.pop('random_sigma','True') == 'True'
-        self.configs['centered_sigma'] = kwargs.pop('centered_sigma','True') == 'True'
+        self.configs["linear_mu"] = kwargs.get("linear_mu", "True") == "True"
+        self.configs["random_mu"] = kwargs.get("random_mu", "True") == "True"
+        self.configs["random_intercept_mu"] = (
+            kwargs.get("random_intercept_mu", "True") == "True"
+        )
+        self.configs["random_slope_mu"] = (
+            kwargs.get("random_slope_mu", "True") == "True"
+        )
+        self.configs["random_sigma"] = kwargs.get("random_sigma", "True") == "True"
+        self.configs["centered_sigma"] = kwargs.get("centered_sigma", "True") == "True"
         ## End default parameters
 
         self.hbr = HBR(self.configs)
@@ -212,116 +251,374 @@ def neg_log_lik(self):
         return -1
 
     def estimate(self, X, y, **kwargs):
-
-        trbefile = kwargs.pop('trbefile', None)
+        trbefile = kwargs.get("trbefile", None)
         if trbefile is not None:
             batch_effects_train = fileio.load(trbefile)
         else:
-            print('Could not find batch-effects file! Initilizing all as zeros ...')
+            print("Could not find batch-effects file! Initilizing all as zeros ...")
             batch_effects_train = np.zeros([X.shape[0], 1])
 
+        self.batch_effects_maps = [
+            {v: i for i, v in enumerate(np.unique(batch_effects_train[:, j]))}
+            for j in range(batch_effects_train.shape[1])
+        ]
+
         self.hbr.estimate(X, y, batch_effects_train)
 
         return self
 
     def predict(self, Xs, X=None, Y=None, **kwargs):
-
-        tsbefile = kwargs.pop('tsbefile', None)
+        tsbefile = kwargs.get("tsbefile", None)
         if tsbefile is not None:
             batch_effects_test = fileio.load(tsbefile)
         else:
-            print('Could not find batch-effects file! Initilizing all as zeros ...')
+            print("Could not find batch-effects file! Initilizing all as zeros ...")
             batch_effects_test = np.zeros([Xs.shape[0], 1])
 
-        pred_type = self.configs['pred_type']
+        pred_type = self.configs["pred_type"]
 
-        if self.configs['transferred'] == False:
-            yhat, s2 = self.hbr.predict(Xs, batch_effects_test, pred=pred_type)
+        if self.configs["transferred"] == False:
+            yhat, s2 = self.hbr.predict(
+                X=Xs,
+                batch_effects=batch_effects_test,
+                batch_effects_maps=self.batch_effects_maps,
+                pred=pred_type,
+                **kwargs,
+            )
         else:
-            raise ValueError("This is a transferred model. Please use predict_on_new_sites function.")
+            raise ValueError(
+                "This is a transferred model. Please use predict_on_new_sites function."
+            )
 
         return yhat.squeeze(), s2.squeeze()
 
+    
+
     def estimate_on_new_sites(self, X, y, batch_effects):
         self.hbr.estimate_on_new_site(X, y, batch_effects)
-        self.configs['transferred'] = True
+        self.configs["transferred"] = True
         return self
 
     def predict_on_new_sites(self, X, batch_effects):
-
         yhat, s2 = self.hbr.predict_on_new_site(X, batch_effects)
         return yhat, s2
 
-    def extend(self, X, y, batch_effects, X_dummy_ranges=[[0.1, 0.9, 0.01]],
-               merge_batch_dim=0, samples=10, informative_prior=False):
-
+    def extend(
+        self,
+        X,
+        y,
+        batch_effects,
+        X_dummy_ranges=[[0.1, 0.9, 0.01]],
+        merge_batch_dim=0,
+        samples=10,
+        informative_prior=False,
+    ):
         X_dummy, batch_effects_dummy = self.hbr.create_dummy_inputs(X_dummy_ranges)
 
-        X_dummy, batch_effects_dummy, Y_dummy = self.hbr.generate(X_dummy,
-                                                                  batch_effects_dummy, samples)
+        X_dummy, batch_effects_dummy, Y_dummy = self.hbr.generate(
+            X_dummy, batch_effects_dummy, samples
+        )
 
-        batch_effects[:, merge_batch_dim] = batch_effects[:, merge_batch_dim] + \
-                                            np.max(batch_effects_dummy[:, merge_batch_dim]) + 1
+        batch_effects[:, merge_batch_dim] = (
+            batch_effects[:, merge_batch_dim]
+            + np.max(batch_effects_dummy[:, merge_batch_dim])
+            + 1
+        )
 
         if informative_prior:
-            self.hbr.adapt(np.concatenate((X_dummy, X)),
-                           np.concatenate((Y_dummy, y)),
-                           np.concatenate((batch_effects_dummy, batch_effects)))
+            self.hbr.adapt(
+                np.concatenate((X_dummy, X)),
+                np.concatenate((Y_dummy, y)),
+                np.concatenate((batch_effects_dummy, batch_effects)),
+            )
         else:
-            self.hbr.estimate(np.concatenate((X_dummy, X)),
-                              np.concatenate((Y_dummy, y)),
-                              np.concatenate((batch_effects_dummy, batch_effects)))
+            self.hbr.estimate(
+                np.concatenate((X_dummy, X)),
+                np.concatenate((Y_dummy, y)),
+                np.concatenate((batch_effects_dummy, batch_effects)),
+            )
 
         return self
 
-    def tune(self, X, y, batch_effects, X_dummy_ranges=[[0.1, 0.9, 0.01]],
-             merge_batch_dim=0, samples=10, informative_prior=False):
-
+    def tune(
+        self,
+        X,
+        y,
+        batch_effects,
+        X_dummy_ranges=[[0.1, 0.9, 0.01]],
+        merge_batch_dim=0,
+        samples=10,
+        informative_prior=False,
+    ):
         tune_ids = list(np.unique(batch_effects[:, merge_batch_dim]))
 
         X_dummy, batch_effects_dummy = self.hbr.create_dummy_inputs(X_dummy_ranges)
 
         for idx in tune_ids:
             X_dummy = X_dummy[batch_effects_dummy[:, merge_batch_dim] != idx, :]
-            batch_effects_dummy = batch_effects_dummy[batch_effects_dummy[:, merge_batch_dim] != idx, :]
+            batch_effects_dummy = batch_effects_dummy[
+                batch_effects_dummy[:, merge_batch_dim] != idx, :
+            ]
 
-        X_dummy, batch_effects_dummy, Y_dummy = self.hbr.generate(X_dummy,
-                                                                  batch_effects_dummy, samples)
+        X_dummy, batch_effects_dummy, Y_dummy = self.hbr.generate(
+            X_dummy, batch_effects_dummy, samples
+        )
 
         if informative_prior:
-            self.hbr.adapt(np.concatenate((X_dummy, X)),
-                           np.concatenate((Y_dummy, y)),
-                           np.concatenate((batch_effects_dummy, batch_effects)))
+            self.hbr.adapt(
+                np.concatenate((X_dummy, X)),
+                np.concatenate((Y_dummy, y)),
+                np.concatenate((batch_effects_dummy, batch_effects)),
+            )
         else:
-            self.hbr.estimate(np.concatenate((X_dummy, X)),
-                              np.concatenate((Y_dummy, y)),
-                              np.concatenate((batch_effects_dummy, batch_effects)))
+            self.hbr.estimate(
+                np.concatenate((X_dummy, X)),
+                np.concatenate((Y_dummy, y)),
+                np.concatenate((batch_effects_dummy, batch_effects)),
+            )
 
         return self
 
-    def merge(self, nm, X_dummy_ranges=[[0.1, 0.9, 0.01]], merge_batch_dim=0,
-              samples=10):
-
+    def merge(
+        self, nm, X_dummy_ranges=[[0.1, 0.9, 0.01]], merge_batch_dim=0, samples=10
+    ):
         X_dummy1, batch_effects_dummy1 = self.hbr.create_dummy_inputs(X_dummy_ranges)
         X_dummy2, batch_effects_dummy2 = nm.hbr.create_dummy_inputs(X_dummy_ranges)
 
-        X_dummy1, batch_effects_dummy1, Y_dummy1 = self.hbr.generate(X_dummy1,
-                                                                     batch_effects_dummy1, samples)
-        X_dummy2, batch_effects_dummy2, Y_dummy2 = nm.hbr.generate(X_dummy2,
-                                                                   batch_effects_dummy2, samples)
-
-        batch_effects_dummy2[:, merge_batch_dim] = batch_effects_dummy2[:, merge_batch_dim] + \
-                                                   np.max(batch_effects_dummy1[:, merge_batch_dim]) + 1
-
-        self.hbr.estimate(np.concatenate((X_dummy1, X_dummy2)),
-                          np.concatenate((Y_dummy1, Y_dummy2)),
-                          np.concatenate((batch_effects_dummy1,
-                                          batch_effects_dummy2)))
+        X_dummy1, batch_effects_dummy1, Y_dummy1 = self.hbr.generate(
+            X_dummy1, batch_effects_dummy1, samples
+        )
+        X_dummy2, batch_effects_dummy2, Y_dummy2 = nm.hbr.generate(
+            X_dummy2, batch_effects_dummy2, samples
+        )
+
+        batch_effects_dummy2[:, merge_batch_dim] = (
+            batch_effects_dummy2[:, merge_batch_dim]
+            + np.max(batch_effects_dummy1[:, merge_batch_dim])
+            + 1
+        )
+
+        self.hbr.estimate(
+            np.concatenate((X_dummy1, X_dummy2)),
+            np.concatenate((Y_dummy1, Y_dummy2)),
+            np.concatenate((batch_effects_dummy1, batch_effects_dummy2)),
+        )
 
         return self
 
     def generate(self, X, batch_effects, samples=10):
-
-        X, batch_effects, generated_samples = self.hbr.generate(X, batch_effects,
-                                                                samples)
+        X, batch_effects, generated_samples = self.hbr.generate(
+            X, batch_effects, samples
+        )
         return X, batch_effects, generated_samples
+        
+    def get_mcmc_quantiles(self, X, batch_effects=None, z_scores=None):
+
+        """
+        Computes quantiles of an estimated normative model.
+
+        Args:
+            X ([N*p]ndarray): covariates for which the quantiles are computed (must be scaled if scaler is set)
+            batch_effects (ndarray): the batch effects corresponding to X
+            z_scores (ndarray): Use this to determine which quantiles will be computed. The resulting quantiles will have the z-scores given in this list. 
+        """
+        # Set batch effects to zero if none are provided
+        if batch_effects is None:
+            batch_effects = batch_effects_test = np.zeros([X.shape[0], 1])
+            
+
+        # Set the z_scores for which the quantiles are computed
+        if z_scores is None:
+            z_scores = np.arange(-3, 4)
+        likelihood=self.configs['likelihood']
+
+        # Determine the variables to predict
+        if self.configs["likelihood"] == "Normal":
+            var_names = ["mu_samples", "sigma_plus_samples"]
+        elif self.configs["likelihood"].startswith("SHASH"):
+            var_names = [
+                "mu_samples",
+                "sigma_plus_samples",
+                "epsilon_samples",
+                "delta_plus_samples",
+            ]
+        else:
+            exit("Unknown likelihood: " + self.configs["likelihood"])
+
+        # Delete the posterior predictive if it already exists
+        if 'posterior_predictive' in self.hbr.idata.groups():
+            del self.hbr.idata.posterior_predictive
+            
+        # Do a forward to get the posterior predictive in the idata
+        self.hbr.predict(
+            X=X,
+            batch_effects=batch_effects,
+            batch_effects_maps=self.batch_effects_maps,
+            pred="single",
+            var_names=var_names+["y_like"],
+        )
+
+        # Extract the relevant samples from the idata
+        post_pred = az.extract(
+            self.hbr.idata, "posterior_predictive", var_names=var_names
+        )
+
+        # Separate the samples into a list so that they can be unpacked 
+        array_of_vars = list(map(lambda x: post_pred[x], var_names))
+
+        # Create an array to hold the quantiles
+        len_synth_data, n_mcmc_samples = post_pred["mu_samples"].shape
+        quantiles = np.zeros((z_scores.shape[0], len_synth_data, n_mcmc_samples))
+
+        # Compute the quantile iteratively for each z-score
+        for i, j in enumerate(z_scores):
+            zs = np.full((len_synth_data, n_mcmc_samples), j, dtype=float)
+            quantiles[i] = xarray.apply_ufunc(
+                quantile,
+                *array_of_vars,
+                kwargs={"zs": zs, "likelihood":  self.configs['likelihood']},
+            )
+        return quantiles.mean(axis=-1)
+    
+
+    def get_mcmc_zscores(self, X, y, batch_effects=None):
+
+        """
+        Computes zscores of data given an estimated model
+
+        Args:
+            X ([N*p]ndarray): covariates
+            y ([N*1]ndarray): response variables
+            batch_effects (ndarray): the batch effects corresponding to X
+        """
+        # Set batch effects to zero if none are provided
+        print(self.configs['likelihood'])
+        if batch_effects is None:
+            batch_effects = batch_effects_test = np.zeros([X.shape[0], 1])
+            
+        # Determine the variables to predict
+        if self.configs["likelihood"] == "Normal":
+            var_names = ["mu_samples", "sigma_plus_samples"]
+        elif self.configs["likelihood"].startswith("SHASH"):
+            var_names = [
+                "mu_samples",
+                "sigma_plus_samples",
+                "epsilon_samples",
+                "delta_plus_samples",
+            ]
+        else:
+            exit("Unknown likelihood: " + self.configs["likelihood"])
+
+        # Delete the posterior predictive if it already exists
+        if 'posterior_predictive' in self.hbr.idata.groups():
+            del self.hbr.idata.posterior_predictive
+            
+        # Do a forward to get the posterior predictive in the idata
+        self.hbr.predict(
+            X=X,
+            batch_effects=batch_effects,
+            batch_effects_maps=self.batch_effects_maps,
+            pred="single",
+            var_names=var_names+["y_like"],
+        )
+
+        # Extract the relevant samples from the idata
+        post_pred = az.extract(
+            self.hbr.idata, "posterior_predictive", var_names=var_names
+        )
+
+        # Separate the samples into a list so that they can be unpacked 
+        array_of_vars = list(map(lambda x: post_pred[x], var_names))
+
+        # Create an array to hold the quantiles
+        len_data, n_mcmc_samples = post_pred["mu_samples"].shape
+
+        # Compute the quantile iteratively for each z-score
+        z_scores =  xarray.apply_ufunc(
+            z_score,
+            *array_of_vars,
+            kwargs={"y": y, "likelihood": self.configs['likelihood']},
+        )
+        return z_scores.mean(axis=-1)
+    
+    
+
+def S_inv(x, e, d):
+    return np.sinh((np.arcsinh(x) + e) / d)
+
+def K(p, x):
+    """
+    Computes the values of spp.kv(p,x) for only the unique values of p
+    """
+
+    ps, idxs = np.unique(p, return_inverse=True)
+    return spp.kv(ps, x)[idxs].reshape(p.shape)
+
+def P(q):
+    """
+    The P function as given in Jones et al.
+    :param q:
+    :return:
+    """
+    frac = np.exp(1 / 4) / np.sqrt(8 * np.pi)
+    K1 = K((q + 1) / 2, 1 / 4)
+    K2 = K((q - 1) / 2, 1 / 4)
+    a = (K1 + K2) * frac
+    return a
+
+def m(epsilon, delta, r):
+    """
+    The r'th uncentered moment. Given by Jones et al.
+    """
+    frac1 = 1 / np.power(2, r)
+    acc = 0
+    for i in range(r + 1):
+        combs = spp.comb(r, i)
+        flip = np.power(-1, i)
+        ex = np.exp((r - 2 * i) * epsilon / delta)
+        p = P((r - 2 * i) / delta)
+        acc += combs * flip * ex * p
+    return frac1 * acc
+
+def quantile( mu, sigma, epsilon=None, delta=None, zs=0, likelihood = "Normal"):
+    """Get the zs'th quantiles given likelihood parameters"""
+    if likelihood.startswith('SHASH'):
+        if likelihood == "SHASHo":
+            quantiles = S_inv(zs,epsilon,delta)*sigma + mu
+        elif likelihood == "SHASHo2":
+            sigma_d = sigma/delta
+            quantiles = S_inv(zs,epsilon,delta)*sigma_d + mu
+        elif likelihood == "SHASHb":
+            true_mu = m(epsilon, delta, 1)
+            true_sigma = np.sqrt((m(epsilon, delta, 2) - true_mu ** 2))
+            SHASH_c = ((S_inv(zs,epsilon,delta)-true_mu)/true_sigma)
+            quantiles = SHASH_c *sigma + mu
+    elif likelihood == 'Normal':
+        quantiles = zs*sigma + mu
+    else:
+        exit("Unsupported likelihood")
+    return quantiles
+
+    
+def z_score(mu, sigma, epsilon=None, delta=None, y=None, likelihood = "Normal"):
+    """Get the z-scores of Y, given likelihood parameters"""
+    if likelihood.startswith('SHASH'):
+        if likelihood == "SHASHo":
+            SHASH = (y-mu)/sigma
+            Z = np.sinh(np.arcsinh(SHASH)*delta - epsilon)
+        elif likelihood == "SHASHo2":
+            sigma_d = sigma/delta
+            SHASH = (y-mu)/sigma_d
+            Z = np.sinh(np.arcsinh(SHASH)*delta - epsilon)
+        elif likelihood == "SHASHb":
+            true_mu = m(epsilon, delta, 1)
+            true_sigma = np.sqrt((m(epsilon, delta, 2) - true_mu ** 2))
+            SHASH_c = ((y-mu)/sigma)
+            SHASH = SHASH_c * true_sigma + true_mu
+            Z = np.sinh(np.arcsinh(SHASH) * delta - epsilon)
+    elif likelihood == 'Normal':
+        Z = (y-mu)/sigma
+    else:
+        exit("Unsupported likelihood")
+    return Z
+
diff --git a/pcntoolkit/normative_parallel.py b/pcntoolkit/normative_parallel.py
index ad9d7c35..268e9461 100755
--- a/pcntoolkit/normative_parallel.py
+++ b/pcntoolkit/normative_parallel.py
@@ -916,7 +916,7 @@ def bashwrap_nm(processing_dir,
                              job_call + ["\n"])
 
     # changes permissoins for bash.sh file
-    os.chmod(processing_dir + job_name, 0o700)
+    os.chmod(processing_dir + job_name, 0o770)
 
 
 def qsub_nm(job_path,
@@ -1121,7 +1121,7 @@ def sbatchwrap_nm(processing_dir,
                              job_call + ["\n"] + [sbatch_exit])
 
     # changes permissoins for bash.sh file
-    os.chmod(processing_dir + job_name, 0o700)
+    os.chmod(processing_dir + job_name, 0o770)
 
 def sbatch_nm(job_path,
               log_path):
diff --git a/pcntoolkit/trendsurf.py b/pcntoolkit/trendsurf.py
index d7e03732..f009ad29 100644
--- a/pcntoolkit/trendsurf.py
+++ b/pcntoolkit/trendsurf.py
@@ -134,7 +134,7 @@ def get_args(*args):
 
 
 def estimate(filename, maskfile, basis, ard=False, outputall=False,
-             saveoutput=True):
+             saveoutput=True, **kwargs):
     """ Estimate a trend surface model
 
     This will estimate a trend surface model, independently for each subject.
@@ -166,7 +166,10 @@ def estimate(filename, maskfile, basis, ard=False, outputall=False,
               * explainedvar - explained variance
               * rmse - standardised mean squared error
     """
-
+    
+    # parse arguments
+    optim = kwargs.get('optimizer', 'powell')
+    
     # load data
     print("Processing data in", filename)
     Y, X, mask = load_data(filename, maskfile)
@@ -204,7 +207,7 @@ def estimate(filename, maskfile, basis, ard=False, outputall=False,
     for i in range(0, N):
         print("Estimating model ", i+1, "of", N)
         breg = BLR()
-        hyp[i, :] = breg.estimate(hyp0, Phi, Yz[:, i])
+        hyp[i, :] = breg.estimate(hyp0, Phi, Yz[:, i], optimizer=optim)
         m[i, :] = breg.m
         nlZ[i] = breg.nlZ
 
diff --git a/pcntoolkit/util/hbr_utils.py b/pcntoolkit/util/hbr_utils.py
index db7a5ccd..89213100 100644
--- a/pcntoolkit/util/hbr_utils.py
+++ b/pcntoolkit/util/hbr_utils.py
@@ -7,7 +7,7 @@
 import pickle
 import matplotlib.pyplot as plt
 import pandas as pd
-import pymc3 as pm
+import pymc as pm
 from pcntoolkit.model.SHASH import *
 from pcntoolkit.model.hbr import bspline_transform
 
@@ -59,7 +59,7 @@ def z_score(Y, params, likelihood = "Normal"):
             SHASH_c = ((Y-mu)/sigma)
             SHASH = SHASH_c * true_sigma + true_mu
             Z = np.sinh(np.arcsinh(SHASH) * delta - epsilon)
-    elif likelihood == 'Normal':
+    if likelihood == 'Normal':
         Z = (Y-params['mu'])/params['sigma']
     else:
         exit("Unsupported likelihood")
@@ -113,7 +113,7 @@ def quantile(zs, params, likelihood = "Normal"):
             true_sigma = np.sqrt((m(epsilon, delta, 2) - true_mu ** 2))
             SHASH_c = ((S_inv(zs,epsilon,delta)-true_mu)/true_sigma)
             quantiles = SHASH_c *sigma + mu
-    elif likelihood == 'Normal':
+    if likelihood == 'Normal':
         quantiles = zs*params['sigma'] + params['mu']
     else:
         exit("Unsupported likelihood")
@@ -157,8 +157,8 @@ def forward(X, Z, model, sample):
 
     if likelihood == 'Normal':
         parameter_list = ['mu','sigma']
-    elif likelihood in ['SHASHb','SHASHo','SHASHo2']:
-        parameter_list = ['mu','sigma','epsilon','delta']
+    # elif likelihood in ['SHASHb','SHASHo','SHASHo2']:
+    #     parameter_list = ['mu','sigma','epsilon','delta']
     else:
         exit("Unsupported likelihood")
 
diff --git a/pcntoolkit/util/utils.py b/pcntoolkit/util/utils.py
index e81c8158..3d22bb57 100644
--- a/pcntoolkit/util/utils.py
+++ b/pcntoolkit/util/utils.py
@@ -14,7 +14,7 @@
 import bspline
 from bspline import splinelab
 from sklearn.datasets import make_regression
-import pymc3 as pm
+import pymc as pm
 from io import StringIO
 import subprocess
 import re
@@ -1322,16 +1322,16 @@ def get_package_versions():
     versions['Python'] = platform.python_version()
     
     try: 
-        import theano
-        versions['Theano'] = theano.__version__
+        import pytensor
+        versions['pytensor'] = pytensor.__version__
     except:
-        versions['Theano'] = ''
+        versions['pytensor'] = ''
         
     try: 
-        import pymc3
-        versions['PyMC3'] = pymc3.__version__
+        import pymc
+        versions['PyMC'] = pymc.__version__
     except:
-        versions['PyMC3'] = ''
+        versions['PyMC'] = ''
         
     try: 
         import pcntoolkit
@@ -1450,7 +1450,7 @@ def yes_or_no(question):
 
 
 
-#====== This is stuff used for the SHASH distributions, but using numpy (not pymc or theano) ===
+#====== This is stuff used for the SHASH distributions, but using numpy (not pymc or pytensor) ===
 
 def K(p, x):
     return np.array(spp.kv(p, x))
diff --git a/setup.py b/setup.py
index 74ce24e3..811f7cb1 100644
--- a/setup.py
+++ b/setup.py
@@ -1,7 +1,7 @@
 from setuptools import setup, find_packages
 
 setup(name='pcntoolkit',
-      version='0.27',
+      version='0.28',
       description='Predictive Clinical Neuroscience toolkit',
       url='http://github.com/amarquand/PCNtoolkit',
       author='Andre Marquand',
@@ -15,14 +15,12 @@
           'scikit-learn', 
           'bspline',
           'matplotlib',
-          'numpy>=1.19.5,<1.23',
+          'numpy',
           'scipy>=1.3.2',
           'pandas>=0.25.3',
           'torch>=1.1.0', 
           'sphinx-tabs',
-          'pymc3>=3.8,<=3.9.3',
-          'theano==1.0.5',
-          'arviz==0.11.0'
+          'pymc>=5.1.0',
+          'arviz==0.13.0'
       ],
-      #python_requires='<3.10',
       zip_safe=False)
diff --git a/tests/testHBR.py b/tests/testHBR.py
index 6b98ed9f..0a65a2ed 100644
--- a/tests/testHBR.py
+++ b/tests/testHBR.py
@@ -1,4 +1,4 @@
-#!/usr/bin/env python3
+#!/usr/bin/env python
 # -*- coding: utf-8 -*-
 """
 Created on Mon Jul 29 13:26:35 2019
@@ -16,16 +16,20 @@
 import matplotlib.pyplot as plt
 from pcntoolkit.normative import estimate
 from warnings import filterwarnings
+from pcntoolkit.util.utils import scaler
+import xarray
+
 filterwarnings('ignore')
 
+np.random.seed(10)
 
 ########################### Experiment Settings ###############################
 
 
-working_dir = '/home/preclineu/seykia/temp/tests/'  # Specift a working directory
+working_dir = '/home/stijn/temp/'  # Specifyexit() a working directory
                                                     # to save data and results.
 
-simulation_method = 'linear' # 'non-linear'
+simulation_method = 'linear'
 n_features = 1      # The number of input features of X
 n_grps = 2          # Number of batches in data 
 n_samples = 500     # Number of samples in each group (use a list for different
@@ -39,17 +43,16 @@
 X_train, Y_train, grp_id_train, X_test, Y_test, grp_id_test, coef = \
     simulate_data(simulation_method, n_samples, n_features, n_grps, 
                   working_dir=working_dir, plot=True)
-    
 
 ################################# Methods Tests ###############################
     
     
 for model_type in model_types:
     
-    nm = norm_init(X_train, Y_train, alg='hbr', model_type=model_type)
+    nm = norm_init(X_train, Y_train, alg='hbr',likelihood='Normal', model_type=model_type,n_samples=100,n_tuning=10)
     nm.estimate(X_train, Y_train, trbefile=working_dir+'trbefile.pkl')
     yhat, ys2 = nm.predict(X_test, tsbefile=working_dir+'tsbefile.pkl')
-    
+
     for i in range(n_features):
         sorted_idx = X_test[:,i].argsort(axis=0).squeeze()
         temp_X = X_test[sorted_idx,i]
@@ -62,7 +65,7 @@
         for j in range(n_grps):
             plt.scatter(temp_X[temp_be==j,], temp_Y[temp_be==j,], 
                         label='Group' + str(j))
-            plt.plot(temp_X[temp_be==j,], temp_yhat[[temp_be==j,]])
+            plt.plot(temp_X[temp_be==j,], temp_yhat[temp_be==j,])
             plt.fill_between(temp_X[temp_be==j,], temp_yhat[temp_be==j,] - 
                              1.96 * np.sqrt(temp_s2[temp_be==j,]), 
                              temp_yhat[temp_be==j,] + 
@@ -70,6 +73,7 @@
                              color='gray', alpha=0.2)
         plt.title('Model %s, Feature %d' %(model_type, i))
         plt.legend()
+        plt.show()
 
 
 ############################## Normative Modelling Test #######################
diff --git a/tests/testHBR_transfer.py b/tests/testHBR_transfer.py
new file mode 100644
index 00000000..93c3d7f2
--- /dev/null
+++ b/tests/testHBR_transfer.py
@@ -0,0 +1,132 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+Created on Mon Jul 29 13:26:35 2019
+
+@author: seykia
+
+This script tests HBR models with default configs on toy data.
+
+"""
+
+import os
+import numpy as np
+from pcntoolkit.normative_model.norm_utils import norm_init
+from pcntoolkit.util.utils import simulate_data
+import matplotlib.pyplot as plt
+from pcntoolkit.normative import estimate
+from warnings import filterwarnings
+from pcntoolkit.util.utils import scaler
+import xarray
+
+filterwarnings('ignore')
+
+np.random.seed(10)
+
+########################### Experiment Settings ###############################
+
+
+working_dir = '/home/stijn/temp/'  # Specifyexit() a working directory
+                                                    # to save data and results.
+
+simulation_method = 'linear'
+n_features = 1      # The number of input features of X
+n_grps = 5          # Number of batches in data 
+n_transfer_groups = 2 # number of batches in transfer data
+n_samples = 500     # Number of samples in each group (use a list for different
+                    # sample numbers across different batches)
+n_transfer_samples = 100
+
+model_types = ['linear']  # models to try
+
+############################## Data Simulation ################################
+
+
+X_train, Y_train, grp_id_train, X_test, Y_test, grp_id_test, coef = \
+    simulate_data(simulation_method, n_samples, n_features, n_grps, 
+                  working_dir=working_dir, plot=True)
+
+X_train_transfer, Y_train_transfer, grp_id_train_transfer, X_test_transfer, Y_test_transfer, grp_id_test_transfer, coef= simulate_data(simulation_method, n_transfer_samples, n_features = n_features, n_grps=n_transfer_groups, plot=True)
+
+################################# Methods Tests ###############################
+    
+    
+for model_type in model_types:
+    nm = norm_init(X_train, Y_train, alg='hbr',likelihood='Normal', model_type=model_type,n_chains = 4,cores=4,n_samples=100,n_tuning=50, freedom = 5, nknots=8, target_accept="0.99")
+
+    print("Now Estimating on original train data ==============================================")
+    nm.estimate(X_train, Y_train, trbefile=working_dir+'trbefile.pkl')
+    print("Now Predicting on original test data ==============================================")
+    yhat, ys2 = nm.predict(X_test, tsbefile=working_dir+'tsbefile.pkl')
+
+    for i in range(n_features):
+        sorted_idx = X_test[:,i].argsort(axis=0).squeeze()
+        temp_X = X_test[sorted_idx,i]
+        temp_Y = Y_test[sorted_idx,]
+        temp_be = grp_id_test[sorted_idx,:].squeeze()
+        temp_yhat = yhat[sorted_idx,]
+        temp_s2 = ys2[sorted_idx,]
+        
+        plt.figure()
+        for j in range(n_grps):
+            plt.scatter(temp_X[temp_be==j,], temp_Y[temp_be==j,], 
+                        label='Group' + str(j))
+            plt.plot(temp_X[temp_be==j,], temp_yhat[temp_be==j,])
+            plt.fill_between(temp_X[temp_be==j,], temp_yhat[temp_be==j,] - 
+                             1.96 * np.sqrt(temp_s2[temp_be==j,]), 
+                             temp_yhat[temp_be==j,] + 
+                             1.96 * np.sqrt(temp_s2[temp_be==j,]),
+                             color='gray', alpha=0.2)
+        plt.title('Model %s, Feature %d' %(model_type, i))
+        plt.legend()
+        plt.show()
+
+    print("Now Estimating on transfer train data ==============================================")
+    nm.estimate_on_new_sites(X_train_transfer, Y_train_transfer, grp_id_train_transfer)
+    print("Now Estimating on transfer test data ==============================================")
+    yhat, s2 = nm.predict_on_new_sites(X_test_transfer, grp_id_test_transfer)
+
+    for i in range(n_features):
+        sorted_idx = X_test_transfer[:,i].argsort(axis=0).squeeze()
+        temp_X = X_test_transfer[sorted_idx,i]
+        temp_Y = Y_test_transfer[sorted_idx,]
+        temp_be = grp_id_test_transfer[sorted_idx,:].squeeze()
+        temp_yhat = yhat[sorted_idx,]
+        temp_s2 = ys2[sorted_idx,]
+        
+        for j in range(n_transfer_groups):
+            plt.scatter(temp_X[temp_be==j,], temp_Y[temp_be==j,], 
+                        label='Group' + str(j))
+            plt.plot(temp_X[temp_be==j,], temp_yhat[temp_be==j,])
+            plt.fill_between(temp_X[temp_be==j,], temp_yhat[temp_be==j,] - 
+                             1.96 * np.sqrt(temp_s2[temp_be==j,]), 
+                             temp_yhat[temp_be==j,] + 
+                             1.96 * np.sqrt(temp_s2[temp_be==j,]),
+                             color='gray', alpha=0.2)
+        plt.title('Transfer model %s, Feature %d' %(model_type, i))
+        plt.legend()
+        plt.show()
+
+
+############################## Normative Modelling Test #######################
+        
+
+model_type = model_types[0]
+
+covfile = working_dir + 'X_train.pkl'
+respfile = working_dir + 'Y_train.pkl'
+testcov = working_dir + 'X_test.pkl'
+testresp = working_dir + 'Y_test.pkl'
+trbefile = working_dir + 'trbefile.pkl'
+tsbefile = working_dir + 'tsbefile.pkl'
+
+os.chdir(working_dir)
+
+estimate(covfile, respfile, testcov=testcov, testresp=testresp, trbefile=trbefile, 
+         tsbefile=tsbefile, alg='hbr', outputsuffix='_' + model_type, 
+         inscaler='None', outscaler='None', model_type=model_type, 
+         savemodel='True', saveoutput='True')
+
+
+###############################################################################
+
diff --git a/tests/test_hbr_pymc.py b/tests/test_hbr_pymc.py
new file mode 100644
index 00000000..49935864
--- /dev/null
+++ b/tests/test_hbr_pymc.py
@@ -0,0 +1,101 @@
+import os
+import pandas as pd
+import pcntoolkit as ptk
+import pymc as pm
+import numpy as np
+import pickle
+from matplotlib import pyplot as plt
+import arviz as az
+processing_dir = "HBR_demo/"    # replace with a path to your working directory
+if not os.path.isdir(processing_dir):
+    os.makedirs(processing_dir)
+os.chdir(processing_dir)
+processing_dir = os.getcwd()
+
+
+def main():
+    fcon_tr = pd.read_csv('https://raw.githubusercontent.com/predictive-clinical-neuroscience/PCNtoolkit-demo/main/data/fcon1000_tr.csv')
+    fcon_te = pd.read_csv('https://raw.githubusercontent.com/predictive-clinical-neuroscience/PCNtoolkit-demo/main/data/fcon1000_te.csv')
+
+    idps = ['rh_MeanThickness_thickness']
+    covs = ['age']
+    batch_effects = ['sitenum','sex']
+
+    X_train = fcon_tr[covs].to_numpy(dtype=float)
+    Y_train = fcon_tr[idps].to_numpy(dtype=float)
+    batch_effects_train = fcon_tr[batch_effects].to_numpy(dtype=int)
+
+
+    X_test = fcon_te[covs].to_numpy(dtype=float)
+    Y_test = fcon_te[idps].to_numpy(dtype=float)
+    batch_effects_test = fcon_te[batch_effects].to_numpy(dtype=int)
+
+    print(X_test.shape, Y_test.shape, batch_effects_test.shape)
+
+    with open('X_train.pkl', 'wb') as file:
+        pickle.dump(pd.DataFrame(X_train), file)
+    with open('Y_train.pkl', 'wb') as file:
+        pickle.dump(pd.DataFrame(Y_train), file) 
+    with open('trbefile.pkl', 'wb') as file:
+        pickle.dump(pd.DataFrame(batch_effects_train), file) 
+    with open('X_test.pkl', 'wb') as file:
+        pickle.dump(pd.DataFrame(X_test), file)
+    with open('Y_test.pkl', 'wb') as file:
+        pickle.dump(pd.DataFrame(Y_test), file) 
+    with open('tsbefile.pkl', 'wb') as file:
+        pickle.dump(pd.DataFrame(batch_effects_test), file) 
+
+    # a simple function to quickly load pickle files    
+    def ldpkl(filename: str): 
+        with open(filename, 'rb') as f:
+            return pickle.load(f)
+
+    respfile = os.path.join(processing_dir, 'Y_train.pkl') 
+    covfile = os.path.join(processing_dir, 'X_train.pkl') 
+
+    testrespfile_path = os.path.join(processing_dir, 'Y_test.pkl')
+    testcovfile_path = os.path.join(processing_dir, 'X_test.pkl') 
+
+    trbefile = os.path.join(processing_dir, 'trbefile.pkl')
+    tsbefile = os.path.join(processing_dir, 'tsbefile.pkl')
+
+    output_path = os.path.join(processing_dir, 'Models/')    
+    log_dir = os.path.join(processing_dir, 'log/')           
+    if not os.path.isdir(output_path):
+        os.mkdir(output_path)
+    if not os.path.isdir(log_dir):
+        os.mkdir(log_dir)
+
+
+    outputsuffix = '_estimate' 
+    nm = ptk.normative.estimate(covfile=covfile, 
+                    respfile=respfile,
+                    trbefile=trbefile,
+                    testcov=testcovfile_path,
+                    testresp=testrespfile_path,
+                    tsbefile=tsbefile,
+                    alg='hbr', 
+                    likelihood='Normal',
+                    # model_type='bspline',
+                   linear_mu='False',
+                    random_intercept_mu = 'False',
+                    random_slope_mu = 'False',
+                    random_intercept_sigma='False',
+                    random_sigma='False',
+                    random_mu='False',
+                    log_path=log_dir, 
+                    binary='True',
+                    n_samples=17,
+                    n_tuning=13,
+                    n_chains=1,
+                    cores=1,
+                    target_accept=0.99,
+                    init='adapt_diag',
+                    inscaler='standardize',
+                    outscaler='standardize',
+                    output_path=output_path, 
+                    outputsuffix=outputsuffix, 
+                    savemodel=True)
+    
+if __name__=="__main__":
+    main()
\ No newline at end of file