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

# +
import logging
import importlib
importlib.reload(logging) # see https://stackoverflow.com/a/21475297/1469195
log = logging.getLogger()
log.setLevel('INFO')
import sys

logging.basicConfig(format='%(asctime)s %(levelname)s : %(message)s',
                     level=logging.INFO, stream=sys.stdout)

# +
# %%capture
import os
import site
os.sys.path.insert(0, '/home/schirrmr/code/reversible/')
os.sys.path.insert(0, '/home/schirrmr/braindecode/code/braindecode/')
os.sys.path.insert(0, '/home/schirrmr/code/explaining/reversible//')


# %load_ext autoreload
# %autoreload 2
import numpy as np
import logging
log = logging.getLogger()
log.setLevel('INFO')
import sys
logging.basicConfig(format='%(asctime)s %(levelname)s : %(message)s',
                     level=logging.INFO, stream=sys.stdout)
import matplotlib
from matplotlib import pyplot as plt
from matplotlib import cm
# %matplotlib inline
# %config InlineBackend.figure_format = 'png'
matplotlib.rcParams['figure.figsize'] = (12.0, 1.0)
matplotlib.rcParams['font.size'] = 14
import seaborn
seaborn.set_style('darkgrid')

from reversible2.sliced import sliced_from_samples
from numpy.random import RandomState

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
import numpy as np
import copy
import math

import itertools
import torch as th
from braindecode.torch_ext.util import np_to_var, var_to_np
from reversible2.splitter import SubsampleSplitter

from reversible2.view_as import ViewAs
from reversible2.invert import invert

from reversible2.affine import AdditiveBlock
from reversible2.plot import display_text, display_close
from reversible2.bhno import load_file, create_inputs
# -

import sklearn.datasets
X,y  = sklearn.datasets.make_moons(200, shuffle=False, noise=1e-4)
plt.figure(figsize=(4,4))
plt.scatter(X[:100,0], X[:100,1])
plt.scatter(X[100:,0], X[100:,1])
train_inputs_a = np_to_var(X[0:100:2], dtype=np.float32)
train_inputs_b = np_to_var(X[1:100:2], dtype=np.float32)
cuda = False

# ### let's try less data

import sklearn.datasets
X,y  = sklearn.datasets.make_moons(20, shuffle=False, noise=1e-4)
plt.figure(figsize=(4,4))
plt.scatter(X[:10,0], X[:10,1])
plt.scatter(X[10:,0], X[10:,1])
train_inputs_a = np_to_var(X[0:10:2], dtype=np.float32)
train_inputs_b = np_to_var(X[1:10:2], dtype=np.float32)
cuda = False

plt.figure(figsize=(4,4))
plt.scatter(var_to_np(train_inputs_a)[:,0], var_to_np(train_inputs_a)[:,1])
plt.scatter(var_to_np(train_inputs_b)[:,0], var_to_np(train_inputs_b)[:,1])


# +
from reversible2.distribution import TwoClassDist

from reversible2.blocks import dense_add_block, conv_add_block_3x3
from reversible2.rfft import RFFT, Interleave
from reversible2.util import set_random_seeds
from torch.nn import ConstantPad2d
import torch as th
from reversible2.splitter import SubsampleSplitter


set_random_seeds(2019011641, cuda)
feature_model_a = nn.Sequential(
    dense_add_block(2,200),
    dense_add_block(2,200),
    dense_add_block(2,200),
    dense_add_block(2,200),
)
feature_model_b = nn.Sequential(
    dense_add_block(2,200),
    dense_add_block(2,200),
    dense_add_block(2,200),
    dense_add_block(2,200),
)
if cuda:
    feature_model_a.cuda()
if cuda:
    feature_model_b.cuda()
from reversible2.ot_exact import ot_euclidean_loss_for_samples
class_dist_a = TwoClassDist(2,0, [0,1])
if cuda:
    class_dist_a.cuda()
class_dist_b = TwoClassDist(2,0, [0,1])
if cuda:
    class_dist_b.cuda()

optim_model_a = th.optim.Adam(feature_model_a.parameters())
optim_model_b = th.optim.Adam(feature_model_b.parameters())
optim_dist_a = th.optim.Adam(class_dist_a.parameters(), lr=1e-2)
optim_dist_b = th.optim.Adam(class_dist_b.parameters(), lr=1e-2)

# +
mix_log_stds_a = th.zeros(len(train_inputs_a), 2, requires_grad=True)
mix_log_stds_b = th.zeros(len(train_inputs_b), 2, requires_grad=True)
optim_log_stds_a = th.optim.Adam([mix_log_stds_a], lr=1e-2)
optim_log_stds_b = th.optim.Adam([mix_log_stds_b], lr=1e-2)

from reversible2.gaussian import get_gaussian_log_probs

def get_mix_gauss_samples(n_samples, mix_means, mix_stds):
    samples = th.randn(100, *mix_means.shape)
    samples = ((samples * mix_stds.unsqueeze(0)) + mix_means.unsqueeze(0))
    return samples


# +
n_epochs = 1000
n_samples = 200
for i_epoch in range(n_epochs):
    for model, dist, mix_log_stds, optim_stds, optim_model, optim_dist, other_model, other_dist, train_inputs in ((
        feature_model_a, class_dist_a, mix_log_stds_a, optim_log_stds_a, 
        optim_model_a, optim_dist_a, feature_model_b, class_dist_b, train_inputs_a),
        (
        feature_model_b, class_dist_b, mix_log_stds_b, optim_log_stds_b, 
            optim_model_b, optim_dist_b, feature_model_a, class_dist_a, train_inputs_b)):

        with th.no_grad():
            outs = model(train_inputs).detach()
        
        mixed = get_mix_gauss_samples(n_samples, outs, th.exp(mix_log_stds))
        mixed = mixed.view(-1, mixed.shape[-1])
        translated = other_model(invert(model, mixed))
        log_probs = other_dist.get_total_log_prob(0, translated)
        nll_loss = -th.mean(log_probs)
        optim_stds.zero_grad()
        nll_loss.backward()
        optim_stds.step()
        
        outs = model(train_inputs)
        mixed = get_mix_gauss_samples(n_samples, outs, th.exp(mix_log_stds))
        
        own_log_prob = dist.get_total_log_prob(0, mixed)
        nll_loss = -th.mean(own_log_prob)
        optim_model.zero_grad()
        optim_dist.zero_grad()
        #optim_stds.zero_grad()
        nll_loss.backward()
        optim_model.step()
        optim_dist.step()
        #optim_stds.step()
        

    if i_epoch % (n_epochs // 20) == 0:

        fig, axes = plt.subplots(1,3, figsize=(16,5))
        radians = np.linspace(0,2*np.pi,24)
        circle_points = np.stack([np.cos(radians), np.sin(radians)], axis=-1)
        circle_th = np_to_var(circle_points, device=train_inputs.device, dtype=np.float32)
        for i_model, (model, dist, inputs, mix_log_stds, other_model) in enumerate((
            (feature_model_a, class_dist_a, train_inputs_a, mix_log_stds_a, feature_model_b),
            (feature_model_b, class_dist_b, train_inputs_b, mix_log_stds_b, feature_model_a),)):
            outs = model(train_inputs)
            for ax in axes[1:]:
                ax.scatter(var_to_np(outs)[:,0], var_to_np(outs)[:,1], s=30, label="Real data {:d}".format(i_model),
                                   c=[seaborn.color_palette()[i_model * 2]])

            stds = th.exp(mix_log_stds)
            circles_per_point = outs.unsqueeze(1) + (circle_th.unsqueeze(0) * stds.unsqueeze(1))

            for c in var_to_np(circles_per_point):
                axes[1].plot(c[:,0], c[:,1],color=seaborn.color_palette()[i_model * 2],
                            alpha=0.5, lw=0.5)

            translated_circles = other_model(invert(model, circles_per_point.view(-1, circles_per_point.shape[-1]))).view(
                *circles_per_point.shape)
            for c in var_to_np(translated_circles):
                axes[2].plot(c[:,0], c[:,1],color=seaborn.color_palette()[i_model * 2],
                            alpha=0.5, lw=0.5)

            in_circles = invert(model, circles_per_point.view(-1, circles_per_point.shape[-1])).view(
                *circles_per_point.shape)
            for c in var_to_np(in_circles):
                axes[0].plot(c[:,0], c[:,1],color=seaborn.color_palette()[i_model * 2],
                            alpha=0.5, lw=0.5)


            out_circle = (circle_th * dist.get_mean_std(0)[1]) + dist.get_mean_std(0)[0]
            axes[1].plot(var_to_np(out_circle)[:,0], var_to_np(out_circle)[:,1],color=seaborn.color_palette()[i_model * 2],
                        alpha=0.5, lw=0.5)
            in_circle = invert(model, out_circle)
            axes[0].plot(var_to_np(in_circle)[:,0], var_to_np(in_circle)[:,1],color=seaborn.color_palette()[i_model * 2],
                        alpha=1, lw=3,)


            axes[0].scatter(var_to_np(inputs)[:,0], var_to_np(inputs)[:,1], s=30, label="Real data {:d}".format(i_model),
                                   c=[seaborn.color_palette()[i_model * 2]])
            axes[0].set_title("Input space")
            axes[0].axis('equal')
            axes[1].set_title("Output space")
            axes[1].axis('equal')
            axes[2].set_title("Output space swapped")
            axes[2].axis('equal')

        display_close(fig)
        for model, dist, train_inputs in [(feature_model_a, class_dist_a, train_inputs_a),
                                            (feature_model_b, class_dist_b, train_inputs_b)]:
                    fig,axes = plt.subplots(1,2, figsize=(10,4))
                    rng = RandomState(201904114)
                    outs = model(train_inputs)
                    other_X  = sklearn.datasets.make_moons(200, shuffle=False, noise=1e-4)[0][:100]
                    other_ins = np_to_var(other_X, dtype=np.float32)
                    other_outs = model(other_ins)

                    axes[0].plot(var_to_np(other_outs[:,0]), var_to_np(other_outs[:,1]), label="All Outputs",
                            color=seaborn.color_palette()[1])
                    axes[0].scatter(var_to_np(outs[:,0]), var_to_np(outs[:,1]), s=30, c=[seaborn.color_palette()[0]],
                               label="Actual data outputs")
                    axes[0].axis('equal')
                    axes[0].set_title("Output space")
                    plt.axis('equal')
                    samples = dist.get_samples(0, 100)
                    inverted = invert(model, samples)
                    axes[1].scatter(var_to_np(inverted)[:,0], var_to_np(inverted)[:,1], s=30, label="Fake/Unknown Samples",
                               c=[seaborn.color_palette()[1]])
                    axes[1].scatter(var_to_np(train_inputs)[:,0], var_to_np(train_inputs)[:,1], s=30, label="Real data",
                               c=[seaborn.color_palette()[0]])
                    axes[1].legend(bbox_to_anchor=(1,1,0,0))
                    axes[1].set_title("Input space")
                    axes[1].axis('equal')
                    display_close(fig)




# -

out_circle.shape

fig, axes = plt.subplots(1,3, figsize=(16,5))
radians = np.linspace(0,2*np.pi,24)
circle_points = np.stack([np.cos(radians), np.sin(radians)], axis=-1)
circle_th = np_to_var(circle_points, device=train_inputs.device, dtype=np.float32)
for i_model, (model, dist, inputs, mix_log_stds, other_model) in enumerate((
    (feature_model_a, class_dist_a, train_inputs_a, mix_log_stds_a, feature_model_b),
    (feature_model_b, class_dist_b, train_inputs_b, mix_log_stds_b, feature_model_a),)):
    outs = model(train_inputs)
    for ax in axes[1:]:
        ax.scatter(var_to_np(outs)[:,0], var_to_np(outs)[:,1], s=30, label="Real data {:d}".format(i_model),
                           c=[seaborn.color_palette()[i_model * 2]])

    stds = th.exp(mix_log_stds)
    circles_per_point = outs.unsqueeze(1) + (circle_th.unsqueeze(0) * stds.unsqueeze(1))

    for c in var_to_np(circles_per_point):
        axes[1].plot(c[:,0], c[:,1],color=seaborn.color_palette()[i_model * 2],
                    alpha=0.5, lw=0.5)

    translated_circles = other_model(invert(model, circles_per_point.view(-1, circles_per_point.shape[-1]))).view(
        *circles_per_point.shape)
    for c in var_to_np(translated_circles):
        axes[2].plot(c[:,0], c[:,1],color=seaborn.color_palette()[i_model * 2],
                    alpha=0.5, lw=0.5)

    in_circles = invert(model, circles_per_point.view(-1, circles_per_point.shape[-1])).view(
        *circles_per_point.shape)
    for c in var_to_np(in_circles):
        axes[0].plot(c[:,0], c[:,1],color=seaborn.color_palette()[i_model * 2],
                    alpha=0.5, lw=0.5)
        
        
    out_circle = (circle_th * dist.get_mean_std(0)[1]) + dist.get_mean_std(0)[0]
    axes[1].plot(var_to_np(out_circle)[:,0], var_to_np(out_circle)[:,1],color=seaborn.color_palette()[i_model * 2],
                alpha=0.5, lw=0.5)
    in_circle = invert(model, out_circle)
    axes[0].plot(var_to_np(in_circle)[:,0], var_to_np(in_circle)[:,1],color=seaborn.color_palette()[i_model * 2],
                alpha=0.5, lw=1,)
    

    axes[0].scatter(var_to_np(inputs)[:,0], var_to_np(inputs)[:,1], s=30, label="Real data {:d}".format(i_model),
                           c=[seaborn.color_palette()[i_model * 2]])
    axes[0].set_title("Input space")
    axes[0].axis('equal')
    axes[1].set_title("Output space")
    axes[1].axis('equal')
    axes[2].set_title("Output space swapped")
    axes[2].axis('equal')

# ### old more data

# +
from reversible2.distribution import TwoClassDist

from reversible2.blocks import dense_add_block, conv_add_block_3x3
from reversible2.rfft import RFFT, Interleave
from reversible2.util import set_random_seeds
from torch.nn import ConstantPad2d
import torch as th
from reversible2.splitter import SubsampleSplitter


set_random_seeds(2019011641, cuda)
feature_model_a = nn.Sequential(
    dense_add_block(2,200),
    dense_add_block(2,200),
    dense_add_block(2,200),
    dense_add_block(2,200),
)
feature_model_b = nn.Sequential(
    dense_add_block(2,200),
    dense_add_block(2,200),
    dense_add_block(2,200),
    dense_add_block(2,200),
)
if cuda:
    feature_model_a.cuda()
if cuda:
    feature_model_b.cuda()
from reversible2.ot_exact import ot_euclidean_loss_for_samples
class_dist_a = TwoClassDist(2,0, [0,1])
if cuda:
    class_dist_a.cuda()
class_dist_b = TwoClassDist(2,0, [0,1])
if cuda:
    class_dist_b.cuda()

optim_model_a = th.optim.Adam(feature_model_a.parameters())
optim_model_b = th.optim.Adam(feature_model_b.parameters())
optim_dist_a = th.optim.Adam(class_dist_a.parameters(), lr=1e-2)
optim_dist_b = th.optim.Adam(class_dist_b.parameters(), lr=1e-2)

# +
mix_log_stds_a = th.zeros(len(train_inputs_a), 2, requires_grad=True)
mix_log_stds_b = th.zeros(len(train_inputs_b), 2, requires_grad=True)
optim_log_stds_a = th.optim.Adam([mix_log_stds_a], lr=1e-2)
optim_log_stds_b = th.optim.Adam([mix_log_stds_b], lr=1e-2)

from reversible2.gaussian import get_gaussian_log_probs

def get_mix_gauss_samples(n_samples, mix_means, mix_stds):
    samples = th.randn(100, *mix_means.shape)
    samples = ((samples * mix_stds.unsqueeze(0)) + mix_means.unsqueeze(0))
    return samples


# +
n_epochs = 1000
n_samples = 200
for i_epoch in range(n_epochs):
    for model, dist, mix_log_stds, optim_stds, optim_model, optim_dist, other_model, other_dist, train_inputs in ((
        feature_model_a, class_dist_a, mix_log_stds_a, optim_log_stds_a, 
        optim_model_a, optim_dist_a, feature_model_b, class_dist_b, train_inputs_a),
        (
        feature_model_b, class_dist_b, mix_log_stds_b, optim_log_stds_b, 
            optim_model_b, optim_dist_b, feature_model_a, class_dist_a, train_inputs_b)):

        with th.no_grad():
            outs = model(train_inputs).detach()
        
        mixed = get_mix_gauss_samples(n_samples, outs, th.exp(mix_log_stds))
        mixed = mixed.view(-1, mixed.shape[-1])
        translated = other_model(invert(model, mixed))
        log_probs = other_dist.get_total_log_prob(0, translated)
        nll_loss = -th.mean(log_probs)
        optim_stds.zero_grad()
        nll_loss.backward()
        optim_stds.step()
        
        outs = model(train_inputs)
        mixed = get_mix_gauss_samples(n_samples, outs, th.exp(mix_log_stds))
        
        own_log_prob = dist.get_total_log_prob(0, mixed)
        nll_loss = -th.mean(own_log_prob)
        optim_model.zero_grad()
        optim_dist.zero_grad()
        #optim_stds.zero_grad()
        nll_loss.backward()
        optim_model.step()
        optim_dist.step()
        #optim_stds.step()
        

    if i_epoch % (n_epochs // 20) == 0:

        fig, axes = plt.subplots(1,3, figsize=(16,5))
        radians = np.linspace(0,2*np.pi,24)
        circle_points = np.stack([np.cos(radians), np.sin(radians)], axis=-1)
        circle_th = np_to_var(circle_points, device=train_inputs.device, dtype=np.float32)

        for i_model, (model, dist, inputs, mix_log_stds, other_model) in enumerate((
            (feature_model_a, class_dist_a, train_inputs_a, mix_log_stds_a, feature_model_b),
            (feature_model_b, class_dist_b, train_inputs_b, mix_log_stds_b, feature_model_a),)):
            outs = model(train_inputs)
            for ax in axes[1:]:
                ax.scatter(var_to_np(outs)[:,0], var_to_np(outs)[:,1], s=30, label="Real data {:d}".format(i_model),
                                   c=[seaborn.color_palette()[i_model * 2]])

            stds = th.exp(mix_log_stds)
            circles_per_point = outs.unsqueeze(1) + (circle_th.unsqueeze(0) * stds.unsqueeze(1))

            for c in var_to_np(circles_per_point):
                axes[1].plot(c[:,0], c[:,1],color=seaborn.color_palette()[i_model * 2],
                            alpha=0.5, lw=0.5)

            translated_circles = other_model(invert(model, circles_per_point.view(-1, circles_per_point.shape[-1]))).view(
                *circles_per_point.shape)
            for c in var_to_np(translated_circles):
                axes[2].plot(c[:,0], c[:,1],color=seaborn.color_palette()[i_model * 2],
                            alpha=0.5, lw=0.5)

            in_circles = invert(model, circles_per_point.view(-1, circles_per_point.shape[-1])).view(
                *circles_per_point.shape)
            for c in var_to_np(in_circles):
                axes[0].plot(c[:,0], c[:,1],color=seaborn.color_palette()[i_model * 2],
                            alpha=0.5, lw=0.5)

            axes[0].scatter(var_to_np(inputs)[:,0], var_to_np(inputs)[:,1], s=30, label="Real data {:d}".format(i_model),
                                   c=[seaborn.color_palette()[i_model * 2]])
            samples = dist.get_samples(0, 100)
            inverted = invert(model, samples)
            axes[0].scatter(var_to_np(inverted)[:,0], var_to_np(inverted)[:,1], s=30, label="Fake data {:d}".format(i_model),
                                   c=[seaborn.color_palette()[i_model * 2 + 1]])
            for ax in axes[1:]:
                ax.scatter(var_to_np(samples)[:,0], var_to_np(samples)[:,1], s=30, label="Fake data {:d}".format(i_model),
                                       c=[seaborn.color_palette()[i_model * 2 + 1]])
            axes[0].set_title("Input space")
            axes[0].axis('equal')
            axes[1].set_title("Output space")
            axes[1].axis('equal')
            axes[2].set_title("Output space swapped")
            axes[2].axis('equal')

        display_close(fig)
        for model, dist, train_inputs in [(feature_model_a, class_dist_a, train_inputs_a),
                                            (feature_model_b, class_dist_b, train_inputs_b)]:
                    fig,axes = plt.subplots(1,2, figsize=(10,4))
                    rng = RandomState(201904114)
                    outs = model(train_inputs)
                    other_X  = sklearn.datasets.make_moons(200, shuffle=False, noise=1e-4)[0][:100]
                    other_ins = np_to_var(other_X, dtype=np.float32)
                    other_outs = model(other_ins)

                    axes[0].plot(var_to_np(other_outs[:,0]), var_to_np(other_outs[:,1]), label="All Outputs",
                            color=seaborn.color_palette()[1])
                    axes[0].scatter(var_to_np(outs[:,0]), var_to_np(outs[:,1]), s=30, c=[seaborn.color_palette()[0]],
                               label="Actual data outputs")
                    axes[0].axis('equal')
                    axes[0].set_title("Output space")
                    plt.axis('equal')
                    samples = dist.get_samples(0, 100)
                    inverted = invert(model, samples)
                    axes[1].scatter(var_to_np(inverted)[:,0], var_to_np(inverted)[:,1], s=30, label="Fake/Unknown Samples",
                               c=[seaborn.color_palette()[1]])
                    axes[1].scatter(var_to_np(train_inputs)[:,0], var_to_np(train_inputs)[:,1], s=30, label="Real data",
                               c=[seaborn.color_palette()[0]])
                    axes[1].legend(bbox_to_anchor=(1,1,0,0))
                    axes[1].set_title("Input space")
                    axes[1].axis('equal')
                    display_close(fig)





# +
n_epochs = 2001
n_samples = 200
for i_epoch in range(n_epochs):
    for model, dist, mix_log_stds, optim_stds, optim_model, optim_dist, other_model, other_dist, train_inputs in ((
        feature_model_a, class_dist_a, mix_log_stds_a, optim_log_stds_a, 
        optim_model_a, optim_dist_a, feature_model_b, class_dist_b, train_inputs_a),
        (
        feature_model_b, class_dist_b, mix_log_stds_b, optim_log_stds_b, 
            optim_model_b, optim_dist_b, feature_model_a, class_dist_a, train_inputs_b)):

        with th.no_grad():
            outs = model(train_inputs).detach()
        
        mixed = get_mix_gauss_samples(n_samples, outs, th.exp(mix_log_stds))
        mixed = mixed.view(-1, mixed.shape[-1])
        translated = other_model(invert(model, mixed))
        log_probs = other_dist.get_total_log_prob(0, translated)
        nll_loss = -th.mean(log_probs)
        optim_stds.zero_grad()
        nll_loss.backward()
        optim_stds.step()
        
        outs = model(train_inputs)
        mixed = get_mix_gauss_samples(n_samples, outs, th.exp(mix_log_stds))
        
        own_log_prob = dist.get_total_log_prob(0, mixed)
        nll_loss = -th.mean(own_log_prob)
        optim_model.zero_grad()
        optim_dist.zero_grad()
        #optim_stds.zero_grad()
        nll_loss.backward()
        optim_model.step()
        optim_dist.step()
        #optim_stds.step()
        

    if i_epoch % (n_epochs // 20) == 0:
        display_text("Epoch {:d} of {:d}".format(i_epoch, n_epochs))
        display_text("Std 0: {:.1E} {:.1E}".format(*class_dist_a.get_mean_std(0)[1].detach().numpy()))
        display_text("Std 1: {:.1E} {:.1E}".format(*class_dist_b.get_mean_std(0)[1].detach().numpy()))

        fig, axes = plt.subplots(1,3, figsize=(16,5))
        radians = np.linspace(0,2*np.pi,24)
        circle_points = np.stack([np.cos(radians), np.sin(radians)], axis=-1)
        circle_th = np_to_var(circle_points, device=train_inputs.device, dtype=np.float32)

        for i_model, (model, dist, inputs, mix_log_stds, other_model) in enumerate((
            (feature_model_a, class_dist_a, train_inputs_a, mix_log_stds_a, feature_model_b),
            (feature_model_b, class_dist_b, train_inputs_b, mix_log_stds_b, feature_model_a),)):
            outs = model(train_inputs)
            for ax in axes[1:]:
                ax.scatter(var_to_np(outs)[:,0], var_to_np(outs)[:,1], s=30, label="Real data {:d}".format(i_model),
                                   c=[seaborn.color_palette()[i_model * 2]])

            stds = th.exp(mix_log_stds)
            circles_per_point = outs.unsqueeze(1) + (circle_th.unsqueeze(0) * stds.unsqueeze(1))

            for c in var_to_np(circles_per_point):
                axes[1].plot(c[:,0], c[:,1],color=seaborn.color_palette()[i_model * 2],
                            alpha=0.5, lw=0.5)

            translated_circles = other_model(invert(model, circles_per_point.view(-1, circles_per_point.shape[-1]))).view(
                *circles_per_point.shape)
            for c in var_to_np(translated_circles):
                axes[2].plot(c[:,0], c[:,1],color=seaborn.color_palette()[i_model * 2],
                            alpha=0.5, lw=0.5)

            in_circles = invert(model, circles_per_point.view(-1, circles_per_point.shape[-1])).view(
                *circles_per_point.shape)
            for c in var_to_np(in_circles):
                axes[0].plot(c[:,0], c[:,1],color=seaborn.color_palette()[i_model * 2],
                            alpha=0.5, lw=0.5)

            axes[0].scatter(var_to_np(inputs)[:,0], var_to_np(inputs)[:,1], s=30, label="Real data {:d}".format(i_model),
                                   c=[seaborn.color_palette()[i_model * 2]])
            samples = dist.get_samples(0, 100)
            inverted = invert(model, samples)
            axes[0].scatter(var_to_np(inverted)[:,0], var_to_np(inverted)[:,1], s=30, label="Fake data {:d}".format(i_model),
                                   c=[seaborn.color_palette()[i_model * 2 + 1]])
            for ax in axes[1:]:
                ax.scatter(var_to_np(samples)[:,0], var_to_np(samples)[:,1], s=30, label="Fake data {:d}".format(i_model),
                                       c=[seaborn.color_palette()[i_model * 2 + 1]])
            axes[0].set_title("Input space")
            axes[0].axis('equal')
            axes[1].set_title("Output space")
            axes[1].axis('equal')
            axes[2].set_title("Output space swapped")
            axes[2].axis('equal')

        display_close(fig)
        for model, dist, train_inputs in [(feature_model_a, class_dist_a, train_inputs_a),
                                            (feature_model_b, class_dist_b, train_inputs_b)]:
                    fig,axes = plt.subplots(1,2, figsize=(10,4))
                    rng = RandomState(201904114)
                    outs = model(train_inputs)
                    other_X  = sklearn.datasets.make_moons(200, shuffle=False, noise=1e-4)[0][:100]
                    other_ins = np_to_var(other_X, dtype=np.float32)
                    other_outs = model(other_ins)

                    axes[0].plot(var_to_np(other_outs[:,0]), var_to_np(other_outs[:,1]), label="All Outputs",
                            color=seaborn.color_palette()[1])
                    axes[0].scatter(var_to_np(outs[:,0]), var_to_np(outs[:,1]), s=30, c=[seaborn.color_palette()[0]],
                               label="Actual data outputs")
                    axes[0].axis('equal')
                    axes[0].set_title("Output space")
                    plt.axis('equal')
                    samples = dist.get_samples(0, 100)
                    inverted = invert(model, samples)
                    axes[1].scatter(var_to_np(inverted)[:,0], var_to_np(inverted)[:,1], s=30, label="Fake/Unknown Samples",
                               c=[seaborn.color_palette()[1]])
                    axes[1].scatter(var_to_np(train_inputs)[:,0], var_to_np(train_inputs)[:,1], s=30, label="Real data",
                               c=[seaborn.color_palette()[0]])
                    axes[1].legend(bbox_to_anchor=(1,1,0,0))
                    axes[1].set_title("Input space")
                    axes[1].axis('equal')
                    display_close(fig)





# +

display(class_dist_b.get_mean_std(0)[1].detach().numpy())
# -

class_dist_b.get_mean_std(0)

# +
display_text("Epoch {:d} of {:d}".format(i_epoch, n_epochs))
display_text("Std 0: {:.1E} {:.1E}".format(*class_dist_a.get_mean_std(0)[1].detach().numpy()))
display_text("Std 1: {:.1E} {:.1E}".format(*class_dist_b.get_mean_std(0)[1].detach().numpy()))

fig, axes = plt.subplots(1,3, figsize=(16,5))
radians = np.linspace(0,2*np.pi,24)
circle_points = np.stack([np.cos(radians), np.sin(radians)], axis=-1)
circle_th = np_to_var(circle_points, device=train_inputs.device, dtype=np.float32)

for i_model, (model, dist, inputs, mix_log_stds, other_model) in enumerate((
    (feature_model_a, class_dist_a, train_inputs_a, mix_log_stds_a, feature_model_b),
    (feature_model_b, class_dist_b, train_inputs_b, mix_log_stds_b, feature_model_a),)):
    outs = model(train_inputs)
    for ax in axes[1:]:
        ax.scatter(var_to_np(outs)[:,0], var_to_np(outs)[:,1], s=30, label="Real data {:d}".format(i_model),
                           c=[seaborn.color_palette()[i_model * 2]])

    stds = th.exp(mix_log_stds)
    circles_per_point = outs.unsqueeze(1) + (circle_th.unsqueeze(0) * stds.unsqueeze(1))

    for c in var_to_np(circles_per_point):
        axes[1].plot(c[:,0], c[:,1],color=seaborn.color_palette()[i_model * 2],
                    alpha=0.5, lw=0.5)

    translated_circles = other_model(invert(model, circles_per_point.view(-1, circles_per_point.shape[-1]))).view(
        *circles_per_point.shape)
    for c in var_to_np(translated_circles):
        axes[2].plot(c[:,0], c[:,1],color=seaborn.color_palette()[i_model * 2],
                    alpha=0.5, lw=0.5)

    in_circles = invert(model, circles_per_point.view(-1, circles_per_point.shape[-1])).view(
        *circles_per_point.shape)
    for c in var_to_np(in_circles):
        axes[0].plot(c[:,0], c[:,1],color=seaborn.color_palette()[i_model * 2],
                    alpha=0.5, lw=0.5)

    axes[0].scatter(var_to_np(inputs)[:,0], var_to_np(inputs)[:,1], s=30, label="Real data {:d}".format(i_model),
                           c=[seaborn.color_palette()[i_model * 2]])
    samples = dist.get_samples(0, 100)
    inverted = invert(model, samples)
    axes[0].scatter(var_to_np(inverted)[:,0], var_to_np(inverted)[:,1], s=30, label="Fake data {:d}".format(i_model),
                           c=[seaborn.color_palette()[i_model * 2 + 1]])
    
    axes[0].set_title("Input space")
    axes[0].axis('equal')
    axes[1].set_title("Output space")
    axes[1].axis('equal')
    axes[2].set_title("Output space swapped")
    axes[2].axis('equal')

display_close(fig)
for model, dist, train_inputs in [(feature_model_a, class_dist_a, train_inputs_a),
                                    (feature_model_b, class_dist_b, train_inputs_b)]:
            fig,axes = plt.subplots(1,2, figsize=(10,4))
            rng = RandomState(201904114)
            outs = model(train_inputs)
            other_X  = sklearn.datasets.make_moons(200, shuffle=False, noise=1e-4)[0][:100]
            other_ins = np_to_var(other_X, dtype=np.float32)
            other_outs = model(other_ins)

            axes[0].plot(var_to_np(other_outs[:,0]), var_to_np(other_outs[:,1]), label="All Outputs",
                    color=seaborn.color_palette()[1])
            axes[0].scatter(var_to_np(outs[:,0]), var_to_np(outs[:,1]), s=30, c=[seaborn.color_palette()[0]],
                       label="Actual data outputs")
            axes[0].axis('equal')
            axes[0].set_title("Output space")
            plt.axis('equal')
            samples = dist.get_samples(0, 100)
            inverted = invert(model, samples)
            axes[1].scatter(var_to_np(inverted)[:,0], var_to_np(inverted)[:,1], s=30, label="Fake/Unknown Samples",
                       c=[seaborn.color_palette()[1]])
            axes[1].scatter(var_to_np(train_inputs)[:,0], var_to_np(train_inputs)[:,1], s=30, label="Real data",
                       c=[seaborn.color_palette()[0]])
            axes[1].legend(bbox_to_anchor=(1,1,0,0))
            axes[1].set_title("Input space")
            axes[1].axis('equal')
            display_close(fig)




StarcoderdataJupyterScriptsDedupFiltered