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

# %load_ext autoreload
# %autoreload 2

# cd ..

# +
from sklearn.datasets import load_breast_cancer, load_digits
from sklearn.preprocessing import MinMaxScaler

from sklearn.model_selection import StratifiedShuffleSplit

import pymc3 as pm

from modules.neural_networks import BayesianMLP
from modules.viz_utils import visulize_bernoulli_post, visulize_categorical_post
# -

# # Binary - Breast Cancer 

# +
X, y = load_breast_cancer(
    return_X_y=True
)

for tr_i, ts_i in StratifiedShuffleSplit(n_splits=1).split(X, y):
    
    X_tr, X_ts = X[tr_i], X[ts_i]
    y_tr, y_ts = y[tr_i], y[ts_i]
    
    scaler = MinMaxScaler()
    scaler.fit(X_tr)
    X_tr = scaler.transform(X_tr)
    X_ts = scaler.transform(X_ts)
# -

bernoulli_perceptron = BayesianMLP(
    X=X_tr, 
    y=y_tr, 
    shape_out=1, 
    likelyhood_model='bernoulli_lk',
    prior=pm.Laplace,
    layers=(15, 5, 3), 
    activation='tanh',
    mu=0.,
    b=1.
)
bernoulli_perceptron.show_graph()

bernoulli_perceptron.fit(n=500000)
bernoulli_traces = bernoulli_perceptron.predict(
    X_ts, 
    y_ts, 
    ['p'], 
    samples=3000
)

visulize_bernoulli_post(
    traces=bernoulli_traces, 
    y=y_ts
)

# # Multiclass - MNIST

# +
X, y = load_digits(return_X_y=True)

for tr_i, ts_i in StratifiedShuffleSplit(n_splits=1).split(X, y):
    
    X_tr, X_ts = X[tr_i], X[ts_i]
    y_tr, y_ts = y[tr_i], y[ts_i]

    scaler = MinMaxScaler()
    scaler.fit(X_tr)

    X_tr = scaler.transform(X_tr)
    X_ts = scaler.transform(X_ts)
# -

categorical_perceptron = BayesianMLP(
    X=X_tr, 
    y=y_tr, 
    shape_out=10, 
    likelyhood_model='categorical_lk',
    layers=(64, 32, 16), 
    activation='tanh',
    prior=pm.Laplace,
    mu=0.,
    b=1.,
    batch_size=64
)
categorical_perceptron.show_graph()

categorical_perceptron.fit(n=100000)
categorical_traces = categorical_perceptron.predict(
    X_ts, 
    y_ts,
    ['p'],
    samples=3000
)

visulize_categorical_post(
    X=X_ts, 
    p=categorical_traces['p'],
    index=83
)

StarcoderdataJupyterScriptsDedupFiltered