<jupyter_start><jupyter_text>Gaussianization Flows<jupyter_code>#@title Install Packages
# %%capture
try:
    import sys, os
    from pyprojroot import here

    # spyder up to find the root
    root = here(project_files=[".here"])

    # append to path
    sys.path.append(str(here()))
except ModuleNotFoundError:
    %%capture
    import os
    os.system("pip install objax wandb chex")
    os.system("pip install git+https://github.com/IPL-UV/rbig_jax.git#egg=rbig_jax")
# jax packages
import jax
import jax.numpy as np
from jax.config import config
# import chex
config.update("jax_enable_x64", False)

import numpy as onp
from functools import partial
import objax
import chex

# library functions
from rbig_jax.data import get_classic
from rbig_jax.plots import plot_joint, plot_joint_prob, plot_info_loss
from rbig_jax.transforms.mixture import MixtureGaussianCDF
from rbig_jax.transforms.logit import Logit
from rbig_jax.transforms.inversecdf import InverseGaussCDF
from rbig_jax.transforms.linear import HouseHolder
from rbig_jax.transforms.base import CompositeTransform
from rbig_jax.models.gaussflow import GaussianizationFlow


# logging
import tqdm
import wandb

# plot methods
import matplotlib.pyplot as plt
import seaborn as sns
sns.reset_defaults()
sns.set_context(context="talk", font_scale=0.7)

%matplotlib inline
%load_ext autoreload
%autoreload 2<jupyter_output><empty_output><jupyter_text>Demo Data<jupyter_code># %%wandb
# get data

n_samples = 10_000
n_features = 2


data = get_classic(n_samples)

# plot data
sns.jointplot(data[:, 0], data[:, 1], s=5, color='blue')<jupyter_output><empty_output><jupyter_text>Model<jupyter_code># model hyperparameters
n_components = 20
n_reflections = 2
generator = objax.random.Generator(123)
learn_temperature = False
n_features = data.shape[1]

# initialize model
bijections = CompositeTransform([
    # Layer I
    MixtureGaussianCDF(n_features=n_features, n_components=n_components),
    Logit(learn_temperature=learn_temperature),
    HouseHolder(n_features=n_features, n_reflections=n_reflections, generator=generator),
    # Layer II
    MixtureGaussianCDF(n_features=n_features, n_components=n_components),
    Logit(learn_temperature=learn_temperature),
    HouseHolder(n_features=n_features, n_reflections=n_reflections, generator=generator),
    # Layer III
    MixtureGaussianCDF(n_features=n_features, n_components=n_components),
    Logit(learn_temperature=learn_temperature),
    HouseHolder(n_features=n_features, n_reflections=n_reflections, generator=generator),
    # Layer IV
    MixtureGaussianCDF(n_features=n_features, n_components=n_components),
    Logit(learn_temperature=learn_temperature),
    HouseHolder(n_features=n_features, n_reflections=n_reflections, generator=generator),
])

# initialize base distribution
base_dist = jax.scipy.stats.norm

# initialize Model
gf_model = GaussianizationFlow(n_features=n_features, bijections=bijections, base_dist=base_dist)
from pprint import pprint

# pprint(gf_model.vars())<jupyter_output><empty_output><jupyter_text>Loss Function<jupyter_code># vectorized the model to allow for batches
# model_vectorized = objax.Vectorize(model, vc=model.vars())

# model_jitted = objax.Jit(model, vc=model.vars())


@objax.Function.with_vars(gf_model.vars())
def nll_loss(x):
    return gf_model.score(x)
nll_loss(data)<jupyter_output><empty_output><jupyter_text>Optimizer<jupyter_code># define the optimizer
opt = objax.optimizer.Adam(gf_model.vars())

# get grad values
gv = objax.GradValues(nll_loss, gf_model.vars())
lr = 0.01
epochs = 500
batch_size = 128

@objax.Function.with_vars(gf_model.vars() + opt.vars())
def train_op(x):
    g, v = gv(x)  # returns gradients, loss
    opt(lr, g)
    return v

# This line is optional: it is compiling the code to make it faster.
train_op = objax.Jit(train_op)<jupyter_output><empty_output><jupyter_text>Testing<jupyter_code>import itertools
from jax import device_put
import tqdm
from jax import random


# initialize parameters
key = random.PRNGKey(123)
itercount = itertools.count()
permute_rng, rng = random.split(key)

train_data = np.array(data)

losses = list()

pbar = tqdm.trange(epochs)

with pbar:
    for i in pbar:

        # batch processing
        permute_rng, rng = random.split(rng)

        # randomly shuffle the data
        train_data = random.permutation(permute_rng, train_data)

        # Train
        avg_loss = []

        for batch_index in range(0, n_samples, batch_size):
            # compute loss
            loss = float(train_op( train_data[batch_index:batch_index+batch_size])[0])
            # append batch
            avg_loss.append(loss)
        # average loss
        batch_loss = np.mean(np.stack(avg_loss))

        # Log losses
        losses.append(batch_loss)
        pbar.set_postfix({"loss": f"{batch_loss:.4f}"})<jupyter_output>100%|██████████| 500/500 [03:34<00:00,  2.33it/s, loss=1.4646]<jupyter_text>Losses<jupyter_code>plt.plot(losses)<jupyter_output><empty_output><jupyter_text>Forward Transformation<jupyter_code># forward transformation
z = gf_model.transform(np.array(data))
sns.jointplot(z[:, 0], z[:, 1], s=5, color='red')<jupyter_output><empty_output><jupyter_text>Inverse Transformation<jupyter_code># forward transformation
X_approx = gf_model.inverse_transform(z)
sns.jointplot(X_approx[:, 0], X_approx[:, 1], s=5, color='green')<jupyter_output><empty_output><jupyter_text>Application I - Generating Samples<jupyter_code>%%time
# generate samples in the latent domain
n_gen_samples = 10_000

# inverse transformation
X_samples = gf_model.sample(n_gen_samples)
sns.jointplot(X_samples[:, 0], X_samples[:, 1], s=5, color='purple')<jupyter_output><empty_output><jupyter_text>Probability Density Estimation<jupyter_code>%%time
# log probability
X_log_prob = gf_model.score_samples(np.array(data))
cmap = "Reds"
probs = X_log_prob
# probs = np.clip(probs, 0.0, 1.0)
probs = np.clip(probs, None, 0.0)
title = "Log Probability"

fig, ax = plt.subplots()
h = ax.scatter(data[:, 0], data[:, 1], s=1, c=probs, cmap=cmap, vmax=0.0)
plt.xlabel("X")
plt.ylabel("Y")
plt.colorbar(h)
ax.set_title(title)
plt.tight_layout()<jupyter_output><empty_output><jupyter_text>Probability<jupyter_code>ngrid = 1_000
buffer = 0.1
xline = np.linspace(data[:,0].min()-buffer, data[:,0].max()+buffer, ngrid)
yline = np.linspace(data[:,1].min()-buffer, data[:,1].max()+buffer, ngrid)
xgrid, ygrid = np.meshgrid(xline, yline)
xyinput = np.concatenate([xgrid.reshape(-1, 1), ygrid.reshape(-1, 1)], axis=1)
%%time
# log probability
log_prob = gf_model.score_samples(np.array(xyinput))
cmap = "Reds"
probs = np.exp(log_prob)
# probs = np.clip(probs, 0.0, 1.0)
title = "Probability"

fig, ax = plt.subplots()
h = ax.scatter(xyinput[:, 0], xyinput[:, 1], s=1, c=probs,  cmap=cmap, vmin=0.0, vmax=1.0)
plt.xlabel("X")
plt.ylabel("Y")
plt.colorbar(h)
ax.set_title(title)
plt.tight_layout()<jupyter_output><empty_output><jupyter_text>Negative Log-Likelihood<jupyter_code>nll = gf_model.score(data)

print("Negative Log-Likelihood:", nll)<jupyter_output>Negative Log-Likelihood: 1.4600139
{'fraction_non_alphanumeric': 0.12272277926846491, 'fraction_numerical': 0.016946758932354186, 'mean_word_length': 5.855759922555663, 'pattern_counts': {'":': 1, '<': 52, '<?xml version=': 0, '>': 51, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '13', 'dir': 'jupyter-structured-clean-dedup', 'id': '68375', 'license': 'MIT', 'n_tokens_mistral': 2734, 'n_tokens_neox': 2602, 'n_words': 489, 'path': 'notebooks/parametric/1.0_Demo_GaussFlows.ipynb', 'repo_name': 'alexhepburn/rbig_jax', 'size': '7081'}