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

# # Load Packages and Code

import os
import matplotlib.pyplot as plt
import scipy.io as sio
import torch
import numpy as np
import pandas as pd
import logging
import re

from train_models import FNO1dComplex, SpectralConv1d, OneStepDataSet
from train_models_with_rescaling import FNO1dComplexTime, TimeScalingDataSet

# + tags=[]
# %load_ext autoreload
# %autoreload 1
# %aimport plotting_utils
# -

# # Data Paths

P = '/local/meliao/projects/fourier_neural_operator/'
DATA_DIR = os.path.join(P, 'data')
MODEL_DIR = os.path.join(P, 'experiments/28_systematic_FNO_dist/models')
PLOTS_DIR = os.path.join(P, 'experiments/28_systematic_FNO_dist/plots/')
RESULTS_DIR = os.path.join(P, 'experiments/28_systematic_FNO_dist/results')

if not os.path.isdir(PLOTS_DIR):
    os.mkdir(PLOTS_DIR)

# +
DSETS = ['FNO-R',]
TRAINING = ['WR', 'TD', 'standard']
model_lst = []
model_pattern = os.path.join(MODEL_DIR, 'dset_{}_train_{}_freq_16_ep_1000')


for dset_k in DSETS:
    for training_k in TRAINING:
        dd = {'dset': dset_k, 'train': training_k}
        dd['model'] = torch.load(model_pattern.format(dset_k, training_k), map_location='cpu')
        model_lst.append(dd)

# +
data_fp = os.path.join(DATA_DIR, '2021-09-10_NLS_data_19/NLS_data_seed_251.mat')
data = sio.loadmat(data_fp)

for k in ['output', 't', 'x']:
    print(k, data[k].shape)

time = np.linspace(0, 20, 20001).reshape(1, 20001)
print('time', time.shape)

# +
START_T_IDX = 0
END_T_IDX = 3001


output = data['output'][:, START_T_IDX:END_T_IDX]
t = time[:, START_T_IDX:END_T_IDX]
x_grid = data['x']
# -

dset = TimeScalingDataSet(output, t, x_grid, do_rescaling=False)


# +
# d_2 = sio.loadmat(os.path.join(DATA_DIR, '2021-08-31_NLS_data_16_test.mat'))
# for k in ['output', 't', 'x']:
#     print(k, d_2[k].shape)

# +
# dset = TimeScalingDataSet(d_2['output'], d_2['t'], d_2['x'])
# -

# # Make Predictions

def prepare_input(X, x_grid=None):
    # X has shape (nbatch, 1, grid_size)
    s = X.shape[-1]
    n_batches = X.shape[0]

    # Convert to tensor
    X_input = torch.view_as_real(torch.tensor(X, dtype=torch.cfloat))

    if x_grid is None:
        # FNO code appends the spatial grid to the input as below:
        x_grid = np.linspace(-np.pi, np.pi, s+1)
        x_grid = x_grid[:s]
        x_grid = torch.tensor(x_grid, dtype=torch.float).view(-1,1)
    # print(x_grid.shape)
    # print(X_input.shape)
    X_input = torch.cat((X_input, x_grid.repeat(n_batches, 1, 1)), axis=2)

    return X_input
def l2_normalized_error(pred, actual):
    """Short summary.

    Parameters
    ----------
    pred : type
        Description of parameter `pred`.
    actual : type
        Description of parameter `actual`.

    Returns
    -------
    types
        Description of returned object.

    """
    errors = pred - actual
    error_norms = torch.linalg.norm(torch.tensor(errors), dim=-1, ord=2)
    actual_norms = torch.linalg.norm(torch.tensor(actual), dim=-1, ord=2)
    normalized_errors = torch.divide(error_norms, actual_norms)
    return normalized_errors


# +
def make_lin_interp_predictions(model, dset, n_obs_per_unit_time=1000):
    """
    """
    preds = torch.zeros_like(dset.X)
    errors = torch.zeros((dset.X.shape[0], dset.X.shape[1]))
    n_integer_time_points = int(np.ceil(dset.t.flatten().shape[0] / n_obs_per_unit_time))
    
    model_input = prepare_input(dset.X[:, 0])
    
    pred = dset.X[:,0]
    future_pred = model(prepare_input(pred))
    for i, t_val in enumerate(dset.t[1:], 1):
        # When i is a multiple of n_obs_per_unit_time, we use the model to predict
        if i % n_obs_per_unit_time == 0:
            pred = future_pred
            p = pred
            model_input = prepare_input(pred)
            # future_pred has the time at a future integer step and we will use it in
            # subsequent loops
            future_pred = model(model_input)
            
            # p stores the final prediction that will be entered into preds array
        
        else:
            # Here we do the linear interpolation between pred and future_pred
            t_remainder = np.modf(t_val)[0]
            p = (1 - t_remainder) * pred + t_remainder * future_pred
        
        preds[:, i] = p
        errors[:, i] = l2_normalized_error(p, dset.X[:, i])
    return preds, errors
    
    
def make_rescaled_predictions_on_the_fly(model, dset):
    """
    assumes dset is of type TimeScalingDataSet but has too large a time grid
    to do the whole rescaling operation at once.
    
    So dset has methods _rescale() which operates on one batch at a time
    """
    preds = torch.zeros_like(dset.X)
    errors = torch.zeros((dset.X.shape[0], dset.X.shape[1]))
    preds[:, 0] = dset.X[:, 0]

    for t_idx in range(1, dset.n_tsteps+1):
        time = dset.t[t_idx]
        
        rescaled_ICs = torch.zeros_like(dset.X[:, t_idx])
        for b_idx in range(dset.n_batches):
            rescaled_ICs[b_idx] = dset._rescale(dset.X[b_idx, 0], time)
        
        if t_idx % 1000 == 0:
            print("Finished with t_idx ", t_idx)
        
        rescaled_ICs = prepare_input(rescaled_ICs)
        time_i = time.repeat(rescaled_ICs.shape[0], 1,1)
        predictions_i = model(rescaled_ICs, time_i)
        inv_root_t = 1 / torch.sqrt(time)
        root_t = torch.sqrt(time)
        predictions_i = inv_root_t * predictions_i
        preds[:, t_idx] = predictions_i
        errors_i = l2_normalized_error(predictions_i, dset.X[:,t_idx])
        errors[:,t_idx] = errors_i
    return preds, errors


def make_TD_predictions(model, dset):
    """
    """
    preds = torch.zeros_like(dset.X)
    errors = torch.zeros((dset.X.shape[0], dset.X.shape[1]))
    preds[:, 0] = dset.X[:, 0]
    inputs_i = prepare_input(dset.X[:, 0])
    for t_idx in range(1, dset.n_tsteps+1):
        time = dset.t[t_idx].repeat(inputs_i.shape[0], 1, 1)
        predictions_i = model(inputs_i, time)
        preds[:, t_idx] = predictions_i
        errors_i = l2_normalized_error(predictions_i, dset.X[:,t_idx])
        errors[:,t_idx] = errors_i
    return preds, errors


# + tags=[]
with torch.no_grad():
    preds_dd = {}
    errors_dd = {}
    for model_dd in model_lst:
        preds_k = model_dd['dset'] + '_' + model_dd['train'] 
        dset_to_use = dset
        model = model_dd['model']
        if model_dd['train'] == 'WR':
#             pass
            preds_dd[preds_k], errors_dd[preds_k] = make_rescaled_predictions_on_the_fly(model, dset_to_use)
        elif model_dd['train'] == 'TD':
#             pass
            preds_dd[preds_k], errors_dd[preds_k] = make_TD_predictions(model, dset_to_use)
        else:
            preds_dd[preds_k], errors_dd[preds_k] = make_lin_interp_predictions(model, dset_to_use, n_obs_per_unit_time=1000)
        print("Finished with model", preds_k)
# -

# # Plot Predictions and Errors

# +
time_idx_start = 1001
errors_dd_for_plt = {'TD': errors_dd['FNO-R_TD'].numpy()[:, time_idx_start:],
                    'WR': errors_dd['FNO-R_WR'].numpy()[:, time_idx_start:],
                    'standard': errors_dd['FNO-R_standard'].numpy()[:, time_idx_start:]}

names_dd = {'TD': 'Time-Dependent FNO',
           'WR': 'FNO Trained with Rescaling',
           'standard': "Standard FNO with linear interpolation"}
fp = os.path.join(PLOTS_DIR, 'intermediate_times.png')
plotting_utils.plot_time_errors(errors_dd_for_plt,
                                special_x_ticks=False,
                                t_grid=dset.t.flatten()[time_idx_start:],
                                x_label='Time',
#                                 y_label='$L_2$-Normalized Errors (log scale)', 
                                title='Test performance on non-integer time points',
                                names_dd=names_dd, 
                                log_scale=False, fp=fp)
# -

for i in range(5):
    plt.plot(errors_dd_for_plt['TD'][i], label=i)
plt.legend()
plt.yscale('log')
plt.show()

# + tags=[]
time_idx_start = 1001
errors_dd_for_plt = {'standard': errors_dd['FNO-R_standard'].numpy()}

names_dd = {'TD': 'Time-Dependent FNO',
           'WR': 'FNO Trained with Rescaling',
           'standard': "Standard FNO with linear interpolation"}
fp = os.path.join(PLOTS_DIR, 'test_lin_interp_intermediate_times.png')
plotting_utils.plot_time_errors(errors_dd_for_plt,
                                special_x_ticks=False,
                                t_grid=dset.t.flatten(),
                                x_label='Time',
#                                 y_label='$L_2$-Normalized Errors (log scale)', 
                                title='Test performance on non-integer time points',
                                names_dd=names_dd, 
                                log_scale=False)
# -

for i in range(5):
    plt.plot(errors_dd['FNO-R_standard'].numpy()[i], label=i)
plt.legend()
plt.grid()
plt.show()



StarcoderdataJupyterScriptsDedupFiltered