<jupyter_start><jupyter_text>Total Correation with Iterative Gaussianization<jupyter_code>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

# library functions

# library functions
from rbig_jax.data import get_classic
from rbig_jax.plots import plot_info_loss, plot_joint

# 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>Data<jupyter_code># %%wandb
# get data
data = get_classic(100_000)

# plot data
plot_joint(data[:1_000], "blue", "Original Data", kind='kde', logger=False)<jupyter_output><empty_output><jupyter_text>InitializationFirst we need to initialize our RBIG method with some parameters. For example, which marginal gaussianization method, what precision do we want, and what rotation method do we want.* `method` - the method used to marginally transform the data. We have histogram or kde currently implemented* `support_extension` - this is the amount we want to extend the support for our marginal distributions. Higher values indicate higher* `precision` - we use interpolation because it is much faster. So the CDF, PDF and PPF are interpolation schemes. This value will allow one to have more "precision" for the estimated quantities but it will be slower.  Total Correlation  Toy Data<jupyter_code>from sklearn.utils import check_random_state

#Parameters
n_samples = 10_000
d_dimensions = 30

seed = 123

rng = check_random_state(seed)

# Generate random normal data
data_original = rng.randn(n_samples, d_dimensions)

# Generate random Data
A = rng.rand(d_dimensions, d_dimensions)

data = data_original @ A

# covariance matrix
C = A.T @ A
vv = onp.diag(C)

# closed Form
tc_original = onp.log(onp.sqrt(vv)).sum() - 0.5 * onp.log(onp.linalg.det(C))

print(f"TC: {tc_original:.4f}")<jupyter_output>TC: 35.2632<jupyter_text>RBIG Estimator<jupyter_code>from rbig_jax.information.entropy import histogram_entropy
from rbig_jax.transforms.histogram import histogram_transform
from rbig_jax.information.total_corr import rbig_total_correlation<jupyter_output><empty_output><jupyter_text>Marginal Uniformization<jupyter_code>from rbig_jax.transforms.histogram import get_hist_params

support_extension = 10
alpha = 1e-5
precision = 100
nbins = int(np.sqrt(data.shape[0]))
return_params = False

uniformize_transform = jax.partial(
    get_hist_params,
    nbins=nbins,
    support_extension=support_extension,
    precision=precision,
    alpha=alpha,
    return_params=return_params
)<jupyter_output>/home/emmanuel/.conda/envs/jax_py38/lib/python3.8/site-packages/jax/lib/xla_bridge.py:130: UserWarning: No GPU/TPU found, falling back to CPU.
  warnings.warn('No GPU/TPU found, falling back to CPU.')<jupyter_text>Univariate Entropy<jupyter_code># define marginal entropy function
entropy_f = jax.partial(histogram_entropy, nbins=1_000, base=2)<jupyter_output><empty_output><jupyter_text>TC w. Scan Function<jupyter_code>%%time
n_iterations = 20
p = 0.1

X_trans, loss = rbig_total_correlation(
    np.array(data), 
    marginal_uni=uniformize_transform, 
    uni_entropy=entropy_f, 
    n_iterations=n_iterations,
    p=p
)
%%time
total_corr = np.sum(loss) * np.log(2)
%%time

print(f"Estimated TC: {total_corr:.4f}")
plot_info_loss(loss, n_layers=len(loss))<jupyter_output><empty_output><jupyter_text>While Loop<jupyter_code>init_state = MaxLoss(layer=0, max_layers=n_layers), X_g, params

def body(state):
    loss_state, x = state

    x, _ = fit_transform_func_jitted(x)
    
    params.append(layer_params)  
    state = MaxLoss(layer=loss_state.layer + 1, max_layers=loss_state.max_layers), x
    
    return state
    
final_state = jax.lax.while_loop(max_condition, body, init_state)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.09270406934745967, 'fraction_numerical': 0.0149289670117987, 'mean_word_length': 5.946488294314381, 'pattern_counts': {'":': 0, '<': 31, '<?xml version=': 0, '>': 31, 'https://': 0, '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': '8', 'dir': 'jupyter-structured-clean-dedup', 'id': '547806', 'license': 'MIT', 'n_tokens_mistral': 1499, 'n_tokens_neox': 1421, 'n_words': 352, 'path': 'notebooks/info_theory/2.1_total_corr.ipynb', 'repo_name': 'alexhepburn/rbig_jax', 'size': '4153'}