<jupyter_start><jupyter_text>Total Correation with Iterative Gaussianization<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(root))
except ModuleNotFoundError:
    import os

    os.system("pip install chex")
    os.system("pip install git+https://github.com/IPL-UV/rbig_jax.git#egg=rbig_jax")
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()))

# jax packages
import jax
import jax.numpy as np
from jax.config import config
import chex
config.update("jax_enable_x64", True)

import numpy as onp
from functools import partial

# logging
import tqdm

# 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>Toy Data<jupyter_code>import sys, os
pysim_root = "/home/emmanuel/code/pysim"
sys.path.append(str(pysim_root))
from pysim.data.information.gaussian import (
    generate_gaussian_data,
    generate_gaussian_rotation_data,
)
from pysim.data.information.linear import generate_linear_entropy_data
from pysim.data.information.studentt import generate_studentt_data
from pysim.information.entropy import marginal_entropy
from pysim.information.histogram import hist_entropy

from typing import NamedTuple


def get_tc_datasets(n_samples, n_features, seed, dataset="gaussian", **kwargs):

    if dataset == "gaussian":

        res = generate_gaussian_data(
            n_samples=n_samples,
            n_features=n_features,
            seed=seed,
            n_base_samples=500_000,
        )

    elif dataset == "gaussian_rotation":

        res = generate_gaussian_rotation_data(
            n_samples=n_samples,
            n_features=n_features,
            seed=seed,
            n_base_samples=500_000,
        )

    elif dataset == "linear_rotation":

        f = partial(marginal_entropy, estimator=hist_entropy, bins="sqrt")

        res = generate_linear_entropy_data(
            n_samples=n_samples,
            n_features=n_features,
            seed=seed,
            marg_h_estimator=f,
            estimator_name="histogram",
            n_base_samples=500_000,
        )
    elif dataset == "studentt":

        res = generate_studentt_data(
            n_samples=n_samples,
            n_features=n_features,
            seed=seed,
            n_base_samples=500_000,
            df=kwargs.get("df", 3.0),
        )

    elif dataset == "cauchy":

        res = generate_studentt_data(
            n_samples=n_samples,
            n_features=n_features,
            seed=seed,
            n_base_samples=500_000,
            df=1.0,
        )

    else:
        raise ValueError(f"Unrecognized dataset.")

    return res
# generate data
res_tuple = get_tc_datasets(
    n_samples=100_000,
    n_features=10,
    seed=123,
    dataset="linear_rotation",
)
fig = corner.corner(res_tuple.X[:, :10], hist_factor=2)<jupyter_output>WARNING:root:Too few points to create valid contours
WARNING:root:Too few points to create valid contours
WARNING:root:Too few points to create valid contours
WARNING:root:Too few points to create valid contours
WARNING:root:Too few points to create valid contours
WARNING:root:Too few points to create valid contours
WARNING:root:Too few points to create valid contours
WARNING:root:Too few points to create valid contours
WARNING:root:Too few points to create valid contours
WARNING:root:Too few points to create valid contours<jupyter_text>RBIG Approximation<jupyter_code>from rbig_jax.models import RBIG
import jax.numpy as jnp

def default_rbig_params(data):
    
    X = jnp.array(data, dtype=np.float64)
    
    params = {
        "X": X,
        # marginal PDF params
        "support_extension": 20, 
        "precision": 100, 
        "eps": 1e-5, 
        "method": "histogram", 
        # Histogram specific params
        "nbins": int(jnp.sqrt(X.shape[0])),
        "alpha": 1e-5,
        # KDE specific params
        "bw": "scott",
        # info reduction loss parameters
        "max_layers" : 250,
        "zero_tolerance" : 10,
        "p" : 0.25,
        # training loop params
        "verbose" : True,
        "n_layers_remove" : 50,
        "interval" : 10,
        # everything (makes it fast...)
        "jitted" : True,
    }

    
    return params
X_g, rbig_model = RBIG(
    **default_rbig_params(res_tuple.X)
)
print(f"True Value: {res_tuple.TC:.4f}")
print(f"RBIG Approx: {rbig_model.info_loss.sum():.4f}")<jupyter_output>True Value: 8.7336
RBIG Approx: 8.6011<jupyter_text>RBIG Model<jupyter_code>from rbig_jax.information.entropy import init_marginal_histogram_entropy
import jax.numpy as jnp
X = jnp.array(data, dtype=np.float64)
nbins = int(jnp.sqrt(X.shape[0]))
H_x = init_marginal_histogram_entropy(nbins=nbins)(X).sum()
from rbig_jax.models import RBIG
import jax.numpy as jnp

X = jnp.array(data, dtype=np.float64)

# a lot of parameters...
support_extension = 20
precision = 100
eps = 1e-5
method = "kde"
# Histogram specific params
nbins = int(jnp.sqrt(X.shape[0]))
alpha = 1e-5
# KDE specific params
bw = "scott"

# info reduction loss parameters
max_layers = 250
zero_tolerance = 10
p = 0.25

# initialize the max layers parameters
verbose = True
n_layers_remove = 50
interval = 10

# everything (makes it fast...)
jitted = True
X_g, rbig_model = RBIG(
    X=X,
    support_extension=support_extension,
    method=method,
    precision=precision,
    alpha=alpha,
    nbins=nbins,
    bw=bw,
    eps=eps,
    max_layers=max_layers,
    zero_tolerance=zero_tolerance,
    p=p,
    verbose=verbose,
    n_layers_remove=n_layers_remove,
    interval=interval,
    jitted=jitted,
)
fig = corner.corner(X_g, hist_factor=2)
rbig_model.info_loss.sum()
fig, ax = plt.subplots()
ax.plot(jnp.cumsum(rbig_model.info_loss))
ax.set(xlabel="Iterations", ylabel="$\Delta$ Info. Reduction")
plt.show()
%%time


X = np.array(data, dtype=np.float32)

tc = rbig_total_correlation(
    X=X.block_until_ready(),
    zero_tolerance=zero_tolerance,
    return_all=False
)
print(f"RBIG TC: {tc:.5f}")<jupyter_output>RBIG TC: 34.75812
CPU times: user 16.8 s, sys: 7.33 s, total: 24.1 s
Wall time: 13.1 s
{'fraction_non_alphanumeric': 0.10160012428149759, 'fraction_numerical': 0.02159391020661799, 'mean_word_length': 3.3529411764705883, 'pattern_counts': {'":': 13, '<': 14, '<?xml version=': 0, '>': 14, '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': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '390803', 'license': 'MIT', 'n_tokens_mistral': 2330, 'n_tokens_neox': 2171, 'n_words': 500, 'path': 'research_notebook/content/notes/info_theory/experiments/.ipynb_checkpoints/rbigjax_tc-checkpoint.ipynb', 'repo_name': 'jejjohnson/research_notebook', 'size': '6437'}