# ---
# 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
# ---

# In this notebook we parse deepmod data into a nice dataframe.

# +
# DeepMoD stuff
import torch
from deepymod import DeepMoD
from deepymod.model.func_approx import NN
from deepymod.model.library import Library1D
from deepymod.model.constraint import LeastSquares
from deepymod.model.sparse_estimators import Threshold
from deepymod.analysis import load_tensorboard
from deepymod.data import Dataset
from deepymod.data.burgers import BurgersDelta

from natsort import natsorted
from os import listdir, path
import pandas as pd
import numpy as np

import seaborn as sns
import matplotlib.pyplot as plt


# -

def correct_eq(found_coeffs):
    # Correct coeffs for burgers
    correct_coeffs = np.zeros((12, 1))
    correct_coeffs[[2, 5]] = 1.0

    n_active_terms_incorrect = np.sum(found_coeffs[correct_coeffs != 0.0] == 0)
    n_inactive_terms_incorrect = np.sum(found_coeffs[correct_coeffs == 0.0] != 0)
    if n_active_terms_incorrect + n_inactive_terms_incorrect > 0:
        correct = False
    else:
        correct = True
    return correct


# +
# Getting deepmod data
data_folder = 'runs/shifted_noise20/'
identifier = 'shiftednoise20_grid_'

v=0.25
dataset = Dataset(BurgersDelta, A=1, v=v)
df = pd.DataFrame()
df['model'] = [key for key in natsorted(listdir(data_folder)) if (key[-3:]!= '.pt') if key.find(identifier) != -1 if key.find('nt') == -1]


# +
# Getting deepmod data
data_folder = 'runs/revision_normal_grid/'
identifier = 'normal_grid'


v=0.25
dataset = Dataset(BurgersDelta, A=1, v=v)
df = pd.DataFrame()
df['model'] = [key for key in natsorted(listdir(data_folder)) if (key[-3:]!= '.pt') if key.find(identifier) != -1 if key.find('nt') == -1]


# +
# Getting deepmod data
data_folder = 'runs/revision_grid_noise_20/'
identifier = 'normal_grid'

v=0.25
dataset = Dataset(BurgersDelta, A=1, v=v)
df = pd.DataFrame()
df['model'] = [key for key in natsorted(listdir(data_folder)) if (key[-3:]!= '.pt') if key.find(identifier) != -1 if key.find('nt') == -1]


# +
# Getting deepmod data
data_folder = 'runs/normal_grid_noise_20/'
identifier = 'normal_grid'

v=0.25
dataset = Dataset(BurgersDelta, A=1, v=v)
df = pd.DataFrame()
df['model'] = [key for key in natsorted(listdir(data_folder)) if (key[-3:]!= '.pt') if key.find(identifier) != -1 if key.find('nt') == -1]


# +
# Getting deepmod data
data_folder = 'runs/rev_random_20/'
identifier = 'rev_random_grid_20_'


v=0.25
dataset = Dataset(BurgersDelta, A=1, v=v)
df = pd.DataFrame()
df['model'] = [key for key in natsorted(listdir(data_folder)) if (key[-3:]!= '.pt') if key.find(identifier) != -1 if key.find('nt') == -1]

# -

#df['n_x'] = [int(row.model.split('_')[2][:]) for idx, row in df.iterrows()]
df['n_x'] = [int(row.model[17:]) for idx, row in df.iterrows()]
df['n_t'] = 100

8.2/40/np.sqrt(4*0.25*0.1)

8.2/40

df

# Adding grids
df['x'] = [np.linspace(-4, 4, row.n_x) for idx, row in df.iterrows()]
df['t'] = [np.linspace(0.1, 1.1, row.n_t) for idx, row in df.iterrows()]
df['t_grid'] = [np.meshgrid(row.t, row.x, indexing='ij')[0] for idx, row in df.iterrows()]
df['x_grid'] = [np.meshgrid(row.t, row.x, indexing='ij')[1] for idx, row in df.iterrows()]

# +
# Calculating true derivatives
df['u_t'] = [dataset.time_deriv(row.x_grid, row.t_grid).reshape(row.x_grid.shape) for idx, row in df.iterrows()]
df['u_x'] = [dataset.library(row.x_grid.reshape(-1, 1), row.t_grid.reshape(-1, 1), poly_order=2, deriv_order=3)[:, 1].reshape(row.x_grid.shape) for idx, row in df.iterrows()]
df['u_xx'] = [dataset.library(row.x_grid.reshape(-1, 1), row.t_grid.reshape(-1, 1), poly_order=2, deriv_order=3)[:, 2].reshape(row.x_grid.shape) for idx, row in df.iterrows()]
df['u_xxx'] = [dataset.library(row.x_grid.reshape(-1, 1), row.t_grid.reshape(-1, 1), poly_order=2, deriv_order=3)[:, 3].reshape(row.x_grid.shape) for idx, row in df.iterrows()]

# Calculating normalizing properties
df['l'] = [np.sqrt(4 * v * row.t)[:, None] for idx, row in df.iterrows()]
df['dz'] = [(np.ones_like(row.t)[:, None] * np.diff(row.x)[0] / row.l) for idx, row in df.iterrows()]
df['u0'] = [np.sqrt(v / (np.pi * row.t))[:, None] for idx, row in df.iterrows()]
# -

[print(data_folder + row.model) for idx, row in df.iterrows()]

# Loading data
deepmod_data = [load_tensorboard(path.join(data_folder+row.model)) for idx, row in df.iterrows()]


coeff_keys = natsorted([key for key in deepmod_data[0].keys() if key[:9] == 'estimator'])

# Checking if correct
df['correct'] = np.stack([correct_eq(data.tail(1)[coeff_keys].to_numpy().T) for data in deepmod_data], axis=-1)
df['coeffs'] = [data.tail(1)[coeff_keys].to_numpy().T for data in deepmod_data]

# Loading final MSE error
df['test_error'] = [data.tail(1)['remaining_MSE_test'].item() for data in deepmod_data]

# +
# Loading different models and getting stuff out
network = NN(2, [30, 30, 30, 30], 1)
library = Library1D(poly_order=2, diff_order=3) # Library function
estimator = Threshold(0.2) # Sparse estimator 
constraint = LeastSquares() # How to constrain
model = DeepMoD(network, library, estimator, constraint) # Putting it all in the model

dt = []
dx = []
d2x = []
d3x = []
trained_model = [i for i in listdir(data_folder) if i[-4:] == 'l.pt']
# -

for idx, row in df.iterrows():
    trained_model = path.join(data_folder,df.model[idx]+'model.pt')
    model.load_state_dict(torch.load(trained_model))
    X = torch.tensor(np.concatenate((row.t_grid.reshape(-1, 1), row.x_grid.reshape(-1, 1)), axis=1), dtype=torch.float32)
    prediction, time_deriv, theta = model(X)
    time_deriv = time_deriv[0].cpu().detach().numpy()
    theta = theta[0].cpu().detach().numpy()
    
    dt.append(time_deriv.reshape(row.t_grid.shape))
    dx.append(theta[:, 1].reshape(row.t_grid.shape))
    d2x.append(theta[:, 2].reshape(row.t_grid.shape))
    d3x.append(theta[:, 3].reshape(row.t_grid.shape))

df['u_t_deepmod'] = dt
df['u_x_deepmod'] = dx
df['u_xx_deepmod'] = d2x
df['u_xxx_deepmod'] = d3x

# +
# Calculating errors
df['u_t_error'] = [np.mean(np.abs(row.u_t - row.u_t_deepmod) * (row.l**0 / row.u0), axis=1) for idx, row in df.iterrows()]
df['u_x_error'] = [np.mean(np.abs(row.u_x - row.u_x_deepmod) * (row.l**1 / row.u0), axis=1) for idx, row in df.iterrows()]
df['u_xx_error'] = [np.mean(np.abs(row.u_xx - row.u_xx_deepmod) * (row.l**2 / row.u0), axis=1) for idx, row in df.iterrows()]
df['u_xxx_error'] = [np.mean(np.abs(row.u_xxx - row.u_xxx_deepmod) * (row.l**3 / row.u0), axis=1) for idx, row in df.iterrows()]

# Making some composite errors
df['full_error'] = [(np.mean(np.abs((row.u_t - row.u_t_deepmod) / np.linalg.norm(row.u_t, axis=1, keepdims=True)) , axis=1) 
                    + np.mean(np.abs((row.u_x - row.u_x_deepmod) / np.linalg.norm(row.u_x, axis=1, keepdims=True)) , axis=1)
                    + np.mean(np.abs((row.u_xx - row.u_xx_deepmod) / np.linalg.norm(row.u_xx, axis=1, keepdims=True)) , axis=1)
                    + np.mean(np.abs((row.u_xxx - row.u_xxx_deepmod) / np.linalg.norm(row.u_xxx, axis=1, keepdims=True)) , axis=1)) 
                    for idx, row in df.iterrows()]
df['deriv_error'] = [(np.mean(np.abs((row.u_x - row.u_x_deepmod) / np.linalg.norm(row.u_x, axis=1, keepdims=True)) , axis=1)
                    + np.mean(np.abs((row.u_xx - row.u_xx_deepmod) / np.linalg.norm(row.u_xx, axis=1, keepdims=True)) , axis=1)
                    + np.mean(np.abs((row.u_xxx - row.u_xxx_deepmod) / np.linalg.norm(row.u_xxx, axis=1, keepdims=True)) , axis=1)) 
                    for idx, row in df.iterrows()]
# -

df.to_pickle('revision_random_20.pd')

np.mean(df.deriv_error[0])

# # Deepmod 2D

# +
# Getting deepmod data
data_folder = 'runs/'
identifier = 'noiseless'

v=0.25
dataset = Dataset(BurgersDelta, A=1, v=v)
df = pd.DataFrame()
#df['model'] = [key for key in natsorted(listdir(data_folder)) if key[-3:]!= '.pt' and key.find(identifier) != -1 if key.find('nt') != -1]
df['model'] = [key for key in natsorted(listdir(data_folder)) if key[-3:]!= '.pt' and key.find(identifier) == -1 if key.find('nt') != -1]
# -

# Adding grid data
df['n_x'] = [int(row.model.split('_')[3]) for idx, row in df.iterrows()]
df['n_t'] = [int(row.model.split('_')[5]) for idx, row in df.iterrows()]
df['run'] = [int(row.model.split('_')[7]) for idx, row in df.iterrows()]

df

# Loading data
deepmod_data = [load_tensorboard(path.join(data_folder, row.model)) for idx, row in df.iterrows()]
coeff_keys = natsorted([key for key in deepmod_data[0].keys() if key[:9] == 'estimator'])

# Checking if correct
df['correct'] = np.stack([correct_eq(data.tail(1)[coeff_keys].to_numpy().T) for data in deepmod_data], axis=-1)
df['coeffs'] = [data.tail(1)[coeff_keys].to_numpy().T for data in deepmod_data]

df.to_pickle('deepmod_2d_noisy.pd')

ax = sns.heatmap(df.pivot(index='n_t', columns='n_x', values='correct'))
ax.invert_yaxis()

ax = sns.heatmap(df.pivot(index='n_t', columns='n_x', values='correct'))
ax.invert_yaxis()



StarcoderdataJupyterScriptsDedupFiltered