<jupyter_start><jupyter_text><jupyter_code># if run on colab
!pip install --quiet git+https://github.com/deepmind/dm-haiku
# need to install nsec...
%pylab inline 
%load_ext autoreload
%autoreload 2

import numpy as onp
import jax
import jax.numpy as jnp
from jax.experimental import optix
import haiku as hk
import tensorflow_probability as tfpp; tfp = tfpp.experimental.substrates.jax
from functools import partial

from nsec.models.dae.dae import DAE
from nsec.utils import display_score_two_moons
from nsec.datasets.two_moons import get_two_moons

tfd = tfp.distributions
tfb = tfp.bijectors
two_moons = get_two_moons(0.05)

# This code is taken from make_moons in sklearn
resolution = 1024
outer_circ_x = jnp.cos(jnp.linspace(0, jnp.pi, resolution))
outer_circ_y = jnp.sin(jnp.linspace(0, jnp.pi, resolution))
inner_circ_x = 1 - jnp.cos(jnp.linspace(0, jnp.pi, resolution))
inner_circ_y = 1 - jnp.sin(jnp.linspace(0, jnp.pi, resolution)) - .5

X = jnp.append(outer_circ_x, inner_circ_x)
Y = jnp.append(outer_circ_y, inner_circ_y)
coords = jnp.vstack([X,Y])
plot(outer_circ_x,outer_circ_y)
plot(inner_circ_x,inner_circ_y)
sigma=0.02

# We approximate the 2 moons with a large mixture
distribution = tfd.MixtureSameFamily(
    mixture_distribution=tfd.Categorical(
        probs=np.ones(2*resolution)/resolution/2),
    components_distribution=tfd.MultivariateNormalDiag(
        loc=coords.T, scale_identity_multiplier=sigma) 
)

coords.T.shape
rng_key = jax.random.PRNGKey(seed=0)
samps = distribution.sample(10000, seed=rng_key)

# Plotting samples
hist2d(samps[:,0], samps[:,1],100);
# But now we can also compute the gradients of log p :-)
score = jax.vmap(jax.grad(distribution.log_prob))
score(np.array([0.,0.]).reshape([1,2]))
figure(dpi=100)
X = np.arange(-1.5, 2.5, 0.1)
Y = np.arange(-1, 1.5, 0.1)
points = stack(meshgrid(X, Y), axis=-1).reshape((-1, 2))
g = score(points).reshape([len(Y), len(X),2])

quiver(X, Y, g[:,:,0], g[:,:,1])
X = np.arange(-1.5, 2.5, 0.01)
Y = np.arange(-1., 1.5, 0.01)
points = stack(meshgrid(X, Y), axis=-1).reshape((-1, 2))
g = score(points).reshape([len(Y), len(X),2])

imshow(sqrt(g[:,:,0]**2 + g[:,:,1]**2))<jupyter_output><empty_output><jupyter_text>**Implement a denoiser** ```class DAE(hk.Module):    def __init__(self):        super(Denoiser, self).__init__()        def __call__(self, x, is_training=False):         Encoder        net = hk.Linear(128)(x)        net = hk.BatchNorm(True, True, 0.9)(net, is_training)        net = jax.nn.leaky_relu(net)        net = hk.Linear(128)(net)        net = hk.BatchNorm(True, True, 0.9)(net, is_training)        net = jax.nn.leaky_relu(net)        net = hk.Linear(2)(net)                 Decoder        net = hk.Linear(128)(net)        net = hk.BatchNorm(True, True, 0.9)(net, is_training)        net = jax.nn.leaky_relu(net)        net = hk.Linear(128)(net)        net = hk.BatchNorm(True, True, 0.9)(net, is_training)        net = jax.nn.leaky_relu(net)        net = hk.Linear(2)(net)                return net```<jupyter_code>def forward(x, is_training=False):
    denoiser = DAE()
    return denoiser(x, is_training)
             
model = hk.transform_with_state(forward)
batch_size = 512
step_size = 0.05

def get_batch(rng_key):
    dist = tfd.MixtureSameFamily(
        mixture_distribution=tfd.Categorical(
            probs=np.ones(2*resolution)/resolution/2),
        components_distribution=tfd.MultivariateNormalDiag(
            loc=coords.T, scale_identity_multiplier=sigma)
    )
    
    y = distribution.sample(batch_size, seed=rng_key)
    x = y + step_size * onp.random.randn(batch_size, 2)
    
    # x is a noisy sample, y is a sample from the distribution
    return {'x':x, 'y':y}
optimizer = optix.adam(1e-3)
rng_seq = hk.PRNGSequence(42)
@jax.jit
def loss_fn(params, state, rng_key, batch):
    res, state = model.apply(params, state, rng_key, batch['x'], is_training=True)
    ypred = batch['x'] + res
    loss = jnp.mean((ypred - batch['y'])**2)
    return loss, state
@jax.jit
def update(params, state, rng_key, opt_state, batch):
    (loss, state), grads = jax.value_and_grad(loss_fn, has_aux=True)(params, state, rng_key, batch)
    updates, new_opt_state = optimizer.update(grads, opt_state)
    new_params = optix.apply_updates(params, updates)
    return loss, new_params, state, new_opt_state
params, state = model.init(next(rng_seq), jnp.zeros((1, 2)), is_training=True)
opt_state = optimizer.init(params)
losses = []
for step in range(5000):
    batch = get_batch(next(rng_seq))
    loss, params, state, opt_state = update(params, state, next(rng_seq), opt_state, batch)
    losses.append(loss)
    if step%100==0:
        print(step, loss)
semilogy(np.array(losses[:]), label='loss')
legend()
X = np.arange(-1.5, 2.5, 0.1)
Y = np.arange(-1, 1.5, 0.1)
points = stack(meshgrid(X, Y), axis=-1).reshape((-1, 2))

score = partial(model.apply, params, state, next(rng_seq))
res, state = score(points, is_training=False)
g = res.reshape([len(Y), len(X),2])
figure(figsize=(14,7))
quiver(X, Y, g[:,:,0], g[:,:,1])
display_score_two_moons(score, two_moons, is_amortized=False, is_reg=True)
X = np.arange(-1.5, 2.5, 0.01)
Y = np.arange(-1, 1.5, 0.01)
points = stack(meshgrid(X, Y), axis=-1).reshape((-1, 2))
res, state = score(points, is_training=False)
g = res.reshape([len(Y), len(X),2])
imshow(np.sqrt(g[:,:,0]**2 + g[:,:,1]**2))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1613083070061327, 'fraction_numerical': 0.0306634454562349, 'mean_word_length': 4.225242718446602, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 9, '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': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '526208', 'license': 'MIT', 'n_tokens_mistral': 2222, 'n_tokens_neox': 2145, 'n_words': 426, 'path': 'notebooks/Denoiser_haiku.ipynb', 'repo_name': 'b-remy/score-estimation-comparison', 'size': '5381'}