# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: 'Python 3.9.5 64-bit (''base'': conda)'
#     name: python3
# ---

# # PCM for tuning regularization parameters in Ridge regression 
#

# Import necessary libraries
import PcmPy as pcm
import numpy as np
import pickle
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from numpy import exp, sqrt
from sklearn.linear_model import Ridge, LinearRegression
from sklearn.pipeline import Pipeline
from sklearn.model_selection import GridSearchCV

# +
# Make the training data: 
N = 100 # Number of observations 
Q = 10  # Number of random effects regressors   
P = 10  # Number of variables 
Z = np.random.normal(0,1,(N,Q)) # Make random design matrix 
U = np.random.normal(0,1,(Q,P))*0.5 # Make random effects  
Y = Z @ U + np.random.normal(0,1,(N,P)) # Generate training data  

# Make testing data: 
Zt = np.random.normal(0,1,(N,Q))
Yt = Zt @ U + np.random.normal(0,1,(N,P))
# -

# Build the datasets from the Data and condition vectors 
comp = np.array([0,0,0,0,0,0,1,1,1,1])
M1 = pcm.regression.RidgeDiag(comp, fit_intercept = True)
M1.optimize_regularization(Z,Y)
print('Estimated thetas:', M1.theta_.round(2))
print('Regularisation:', (exp(M1.theta_[-1])/exp(M1.theta_[:-1])).round(2))

# Now you can fit the model 
M1.fit(Z,Y)
Yp = M1.predict(Zt)
R2 = 1- np.sum((Yt-Yp)**2)/np.sum((Yt)**2)
print('r2 :', R2.round(3))

Yp = M1.predict(Zt)

StarcoderdataJupyterScriptsDedupFiltered