<jupyter_start><jupyter_text>Import Packages<jupyter_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_no_spacetime import FNO1dComplexNoSpacetime
%load_ext autoreload
%autoreload 1
%aimport plotting_utils<jupyter_output><empty_output><jupyter_text>Load Data and Models<jupyter_code>DATA_DIR = '/local/meliao/projects/fourier_neural_operator/data/2021-08-14_NLS_data_files'
MODEL_DIR = '/local/meliao/projects/fourier_neural_operator/experiments/20_investigate_frequency_response/models'
PLOTS_DIR = '/local/meliao/projects/fourier_neural_operator/experiments/20_investigate_frequency_response/plots/Compare_N_X_datasets'
RESULTS_DIR = '/local/meliao/projects/fourier_neural_operator/experiments/20_investigate_frequency_response/results'


model_fp_dd = {'Dataset 0': os.path.join(MODEL_DIR, 'dset_00_time_1_ep_1000'),
              'Dataset 1': os.path.join(MODEL_DIR, 'dset_01_time_1_ep_1000'),
              'Dataset 2': os.path.join(MODEL_DIR, 'dset_02_time_1_ep_1000'),
              'Dataset 3': os.path.join(MODEL_DIR, 'dset_03_time_1_ep_1000'),
              'Dataset 4': os.path.join(MODEL_DIR, 'dset_04_time_1_ep_1000'),
              }

whole_dset_fp_dd = {'Dataset 0': os.path.join(MODEL_DIR, 'whole_dset_00_ep_1000'),
              'Dataset 1': os.path.join(MODEL_DIR, 'whole_dset_01_ep_1000'),
              'Dataset 2': os.path.join(MODEL_DIR, 'whole_dset_02_ep_1000'),
              'Dataset 3': os.path.join(MODEL_DIR, 'whole_dset_03_ep_1000'),
            #   'Dataset 4': os.path.join(MODEL_DIR, 'whole_dset_04_ep_1000'),
              }


model_dd = {k: torch.load(v, map_location='cpu') for k,v in model_fp_dd.items()}

whole_model_dd = {k: torch.load(v, map_location='cpu') for k,v in whole_dset_fp_dd.items()}

if not os.path.isdir(PLOTS_DIR):
    os.mkdir(PLOTS_DIR)
data_fp_dd = {'Dataset 0': os.path.join(DATA_DIR, '00_test.mat'),
                'Dataset 1': os.path.join(DATA_DIR, '01_test.mat'),
                'Dataset 2': os.path.join(DATA_DIR, '02_test.mat'),
                'Dataset 3': os.path.join(DATA_DIR, '03_test.mat'),
                'Dataset 4': os.path.join(DATA_DIR, '04_test.mat'),
                'Dataset 5': os.path.join(DATA_DIR, '05_test.mat')}
data_dd = {k: sio.loadmat(v) for k,v in data_fp_dd.items()}

dataset_dd = {k: OneStepDataSet(v['output'], v['t'], v['x']) for k,v in data_dd.items()}<jupyter_output><empty_output><jupyter_text>Prediction Differences Between Similar ICs<jupyter_code>def prepare_input(X):
    # 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))

    # FNO code appends the spatial grid to the input as below:
    x_grid = torch.linspace(-np.pi, np.pi, s).view(-1,1)
    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
with torch.no_grad():
    preds_dd = {}
    errors_dd = {}
    for k in model_dd.keys():
        model_k = model_dd[k]
        dset_k = dataset_dd[k]
        input = prepare_input(dset_k.X[:,0])
        target = dset_k.X[:,1]
        preds_k = model_k(input)
        preds_dd[k] = preds_k
        errors_dd[k] = l2_normalized_error(preds_k, target)
        print("Finished with model ", k)
        

    whole_dset_preds_dd = {}
    whole_dset_errors_dd = {}
    for k in whole_model_dd.keys():
        model_k = whole_model_dd[k]
        dset_k = dataset_dd[k]
        input = prepare_input(dset_k.X[:,0])
        target = dset_k.X[:,1]
        preds_k = model_k(input)
        whole_dset_preds_dd[k] = preds_k
        whole_dset_errors_dd[k] = l2_normalized_error(preds_k, target)
        print("Finished with model ", k)<jupyter_output>Finished with model  Dataset 0
Finished with model  Dataset 1
Finished with model  Dataset 2
Finished with model  Dataset 3
Finished with model  Dataset 4
Finished with model  Dataset 0
Finished with model  Dataset 1
Finished with model  Dataset 2
Finished with model  Dataset 3<jupyter_text>Prediction plots for ICs freq [1, ..., 5]<jupyter_code>for i in range(1):
    pred_dd = {'FNO predictions': preds_dd['Dataset 0'][i].numpy()}
    x_vals_dd = {'FNO predictions': np.linspace(-np.pi, np.pi, 1024)}
    soln_x_vals = np.linspace(-np.pi, np.pi, 1024)
    soln = dataset_dd['Dataset 0'].X[i,1].numpy()
    plotting_utils.quick_prediction_plot(pred_dd, soln, x_vals_dd=x_vals_dd, soln_x_vals=soln_x_vals)
def plot_average_freq_mistakes(preds_dd, solns_dd, show_n_modes=50, fp=None, abs_val_errors=True):
    """
    solns has different ICs along axis 0
    so does the values in preds_dd
    """

    fig, ax = plt.subplots()
    for k in preds_dd.keys():
        preds_dft = np.abs(np.fft.fft(preds_dd[k], axis=-1))
        soln_dft = np.abs(np.fft.fft(solns_dd[k], axis=-1))
        if abs_val_errors:
            diffs = np.abs(preds_dft - soln_dft)
        else:
            diffs = preds_dft - soln_dft
        means = np.mean(diffs, axis=0)
        stds = np.std(diffs, axis=0)
        if abs_val_errors:
            plt_low = np.clip(means - stds, a_min=0., a_max=None)
        else:
            plt_low = means - stds
        plt_high = means + stds

        ax.plot(means[:show_n_modes], label=k)
        ax.fill_between(np.arange(show_n_modes), plt_low[:show_n_modes], plt_high[:show_n_modes], alpha=0.2)
    ax.hlines(0, xmin=0, xmax=show_n_modes, linestyles='dashed', color='black', alpha=0.5)

    ax.legend()
    if fp is not None:
        plt.savefig(fp)
    else:
        plt.show()
    plt.close(fig)




dd_for_plt = {'ICs freq [1, ..., 5]': preds_dd['Dataset 0'].numpy(),
                'ICs freq [6, ..., 10]': preds_dd['Dataset 1'].numpy(),
                'ICs freq [11, ..., 15]': preds_dd['Dataset 2'].numpy(), }
soln_dd_for_plt = {'ICs freq [1, ..., 5]': dataset_dd['Dataset 0'].X[:,1].numpy(),
                    'ICs freq [6, ..., 10]': dataset_dd['Dataset 1'].X[:,1].numpy(),
                    'ICs freq [11, ..., 15]': dataset_dd['Dataset 2'].X[:,1].numpy()}
plot_average_freq_mistakes(dd_for_plt, soln_dd_for_plt, abs_val_errors=False)<jupyter_output><empty_output><jupyter_text>Prediction plots for ICs freq [6, ..., 10]<jupyter_code>i = 0
pred_dd = {'FNO predictions': preds_dd['Dataset 1'][i].numpy()}
x_vals_dd = {'FNO predictions': np.linspace(-np.pi, np.pi, 1024)}
soln_x_vals = np.linspace(-np.pi, np.pi, 1024)
soln = dataset_dd['Dataset 1'].X[i,1].numpy()
plotting_utils.quick_prediction_plot(pred_dd, soln, x_vals_dd=x_vals_dd, soln_x_vals=soln_x_vals)

def quick_boxplot(errors_dd, names_dd=None, ref_hline=None, fp=None, title=None):
    error_lst = []
    key_lst = []
    for k, errors in errors_dd.items():
        error_lst.append(errors)
        key_lst.append(k)

    if names_dd is not None:
        key_lst = [names_dd[k] for k in key_lst]

    
    fig, ax = plt.subplots()

    ax.set_yscale('log')
    ax.set_ylabel('L2 Normalized Error')
    ax.set_xlabel('FNO Model')
    ax.set_title(title)
    ax.set_xticklabels(labels=key_lst, rotation=45, ha='right')
    if ref_hline is not None:
        ax.hlines(ref_hline, xmin=0.5, xmax=len(key_lst)+ 0.5, linestyles='dashed')
    fig.patch.set_facecolor('white')
    
    ax.boxplot(error_lst)
    fig.tight_layout()
    if fp is not None:
        plt.savefig(fp)
    else:
        plt.show()
    plt.close(fig)
names_dd = {'Dataset 0': 'ICs freq [1, ..., 5]',
            'Dataset 1': 'ICs freq [6, ..., 10]',
            'Dataset 2': 'ICs freq [11, ..., 15]',
            'Dataset 3': 'ICs freq [16, ..., 20]'}
t = 'FNO test errors on different datasets'
fp=os.path.join(PLOTS_DIR, 'time_1_test_errors.png')
quick_boxplot(errors_dd, names_dd = names_dd, title=t, fp=fp)
quick_boxplot(whole_dset_errors_dd, names_dd=names_dd, title='FNO test errors on different datasets')
def double_boxplot(errors_dd_1, errors_dd_2, ref_hline=None, fp=None, title=None):
    FUDGE_X_PARAM = 0.125
    WIDTH_PARAM = 0.25
    error_lst_1 = []
    error_lst_2 = []
    key_lst = []
    for k in errors_dd_1.keys():
        error_lst_1.append(errors_dd_1[k].numpy())
        error_lst_2.append(errors_dd_2[k].numpy())
        key_lst.append(k)
    
    fig, ax = plt.subplots()

    x_ticks = np.arange(len(error_lst_1))
    x_data_1 = x_ticks - FUDGE_X_PARAM
    print(x_data_1)
    x_data_2 = x_ticks + FUDGE_X_PARAM

    bplot1 = ax.boxplot(error_lst_1, widths=WIDTH_PARAM, positions=x_data_1)
    bplot2 = ax.boxplot(error_lst_2, widths=WIDTH_PARAM, positions=x_data_2)

    ax.set_yscale('log')
    ax.set_ylabel('L2 Normalized Error')
    ax.set_xlabel('FNO Model')
    ax.set_title(title)
    ax.set_xticks(ticks=x_ticks)
    ax.set_xticklabels(labels=key_lst, rotation=45, ha='right')
    if ref_hline is not None:
        ax.hlines(ref_hline, xmin=0.5, xmax=len(key_lst)+ 0.5, linestyles='dashed')
    fig.patch.set_facecolor('white')
    

    fig.tight_layout()
    if fp is not None:
        plt.savefig(fp)
    else:
        plt.show()
    plt.close(fig)
# errors_dd.pop('Dataset 4')
double_boxplot(errors_dd, whole_dset_errors_dd)
def make_composed_predictions(model, dset):
    """
    """
    # print(ones_input.shape)
    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]
        predictions_i = model(inputs_i)
        preds[:, t_idx] = predictions_i
        inputs_i = prepare_input(predictions_i)
        errors_i = l2_normalized_error(predictions_i, dset.X[:,t_idx])
        errors[:,t_idx] = errors_i

    return preds, errors
comp_pred_dd = {}
comp_error_dd = {}

with torch.no_grad():
    for k in model_dd.keys():
        model_k = model_dd[k]
        dset_k = dataset_dd[k]
        preds_k, errors_k = make_composed_predictions(model_k, dset_k)
        comp_pred_dd[k] = preds_k
        comp_error_dd[k] = errors_k
        print("Finished with model ", k)

whole_dset_comp_pred_dd = {}
whole_dset_comp_error_dd = {}

with torch.no_grad():
    for k in whole_model_dd.keys():
        model_k = whole_model_dd[k]
        dset_k = dataset_dd[k]
        preds_k, errors_k = make_composed_predictions(model_k, dset_k)
        whole_dset_comp_pred_dd[k] = preds_k
        whole_dset_comp_error_dd[k] = errors_k
        print("Finished with model ", k)
error_dd_for_plt = {k: v.numpy()[:, :10] for k,v in comp_error_dd.items()}
error_dd_for_plt.pop('Dataset 2')
plotting_utils.plot_time_errors(error_dd_for_plt)
fp_pattern_dd = {'ICs freq [1, ..., 5]': os.path.join(RESULTS_DIR, "dset_00_time_1_{}.txt"),
                'ICs freq [6, ..., 10]': os.path.join(RESULTS_DIR, "dset_01_time_1_{}.txt"),
                'ICs freq [11, ..., 15]': os.path.join(RESULTS_DIR, "dset_02_time_1_{}.txt"),
                'ICs freq [16, ..., 20]': os.path.join(RESULTS_DIR, "dset_03_time_1_{}.txt")}
train_df_dd = {k: pd.read_table(v.format('train')) for k,v in fp_pattern_dd.items()}
test_df_dd = {k: pd.read_table(v.format('test')) for k,v in fp_pattern_dd.items()}
for k in fp_pattern_dd.keys():
    t = "Train/test data: " + k
    plotting_utils.make_train_test_plot(train_df_dd[k], test_df_dd[k], title=t, log_scale=True)
with torch.no_grad():
    ones_vec = torch.tensor(1., dtype=torch.float).repeat((input1.shape[0], 1,1))
    output1_dd = {model_k: model(input1) for model_k, model in model_dd.items()}
    no_W_output1_dd = {model_k: model(input1) for model_k, model in no_W_model_dd.items()}
    output2_dd = {model_k: model(input2) for model_k, model in model_dd.items()}
    no_W_output2_dd = {model_k: model(input2) for model_k, model in no_W_model_dd.items()}

def plot_two_solutions(preds_dd, ic_dd, soln_dd, x_vals, fp=None, title=None, alpha=0.7, show_n_modes=50):
    fig, ax = plt.subplots(2,3)
    fig.set_size_inches(14, 7.5)
    ax[0,0].set_title("$Re(u)$", size=20)
    ax[0,1].set_title("$Im(u)$", size=20)
    ax[0,2].set_title("$Abs(DFT(u))$", size=20)

    ax[0,0].set_ylabel("ICs", size=20)
    ax[1,0].set_ylabel("Predictions", size=20)
    for k,v in ic_dd.items():
        ax[0,0].plot(x_vals, np.real(v), alpha=alpha, label=k)
        ax[0,1].plot(x_vals, np.imag(v), alpha=alpha, label=k)
        v_dft = np.fft.fft(v)
        v_dft_abs = np.abs(v_dft)
        ax[0,2].plot(v_dft_abs[:show_n_modes],  label=k)
        # ax[2].plot(v_dft_abs[:show_n_modes], alpha=.3, label=k)
        ax[0,2].set_ylabel("Abs(DFT(u))", size=13)
        ax[0,2].set_xlabel("Frequency", size=13)
    for k,v in preds_dd.items():
        ax[1,0].plot(x_vals, np.real(v), alpha=alpha, label=k)
        ax[1,1].plot(x_vals, np.imag(v), alpha=alpha, label=k)
        v_dft = np.fft.fft(v)
        v_dft_abs = np.abs(v_dft)
        ax[1,2].plot(v_dft_abs[:show_n_modes],  label=k)
        # ax[2].plot(v_dft_abs[:show_n_modes], alpha=.3, label=k)
    for k,v in soln_dd.items():
        ax[1,0].plot(x_vals, np.real(v), '--', alpha=alpha, label=k)
        ax[1,1].plot(x_vals, np.imag(v), '--', alpha=alpha, label=k)
        v_dft = np.fft.fft(v)
        v_dft_abs = np.abs(v_dft)
        ax[1,2].plot(v_dft_abs[:show_n_modes], '--', alpha=alpha, label=k)
    ax[1,0].legend()
    ax[0,0].legend()
    ax[1,2].set_ylabel("Abs(DFT(u))", size=13)
    ax[1,2].set_xlabel("Frequency", size=13)
    ax[0,2].legend()
    ax[1,2].legend()
    fig.suptitle(title, size=20)
    fig.patch.set_facecolor('white')
    fig.tight_layout()
    if fp is not None:
        plt.savefig(fp)
    else:
        plt.show()
    plt.close(fig)

def plot_two_solutions_only_DFT(preds_dd, ic_dd, soln_dd, x_vals, fp=None, title=None, alpha=0.7, show_n_modes=50):
    fig, ax = plt.subplots(1,2)
    fig.set_size_inches(14, 7.5)
    ax[0].set_title("$|\\mathcal{F}(u(0, x))|$", size=20)
    ax[1].set_title("$|\\mathcal{F}(u(1, x))|$", size=20)

    # ax[0,0].set_ylabel("ICs", size=20)
    # ax[1,0].set_ylabel("Predictions", size=20)
    for k,v in ic_dd.items():
        # ax[0,0].plot(x_vals, np.real(v), alpha=alpha, label=k)
        # ax[0,1].plot(x_vals, np.imag(v), alpha=alpha, label=k)
        v_dft = np.fft.fft(v)
        v_dft_abs = np.abs(v_dft)
        ax[0].plot(v_dft_abs[:show_n_modes],  label=k)
        # ax[2].plot(v_dft_abs[:show_n_modes], alpha=.3, label=k)
    ax[0].set_ylabel("Initial Conditions", size=13)
    ax[0].set_xlabel("Frequency", size=13)
    for k,v in preds_dd.items():
        # ax[1,0].plot(x_vals, np.real(v), alpha=alpha, label=k)
        # ax[1,1].plot(x_vals, np.imag(v), alpha=alpha, label=k)
        v_dft = np.fft.fft(v)
        v_dft_abs = np.abs(v_dft)
        ax[1].plot(v_dft_abs[:show_n_modes],  label=k)
        # ax[2].plot(v_dft_abs[:show_n_modes], alpha=.3, label=k)
    ax[1].set_ylabel("Predictions/Solutions", size=13)
    ax[1].set_xlabel("Frequency", size=13)
    for k,v in soln_dd.items():
        # ax[1,0].plot(x_vals, np.real(v), '--', alpha=alpha, label=k)
        # ax[1,1].plot(x_vals, np.imag(v), '--', alpha=alpha, label=k)
        v_dft = np.fft.fft(v)
        v_dft_abs = np.abs(v_dft)
        ax[1].plot(v_dft_abs[:show_n_modes], '--', alpha=alpha, label=k)
    ax[1].legend()
    ax[0].legend()
    # ax[1,2].set_ylabel("Abs(DFT(u))", size=13)
    # ax[1,2].set_xlabel("Frequency", size=13)
    # ax[0,2].legend()
    # ax[1,2].legend()
    fig.suptitle(title, size=20)
    fig.patch.set_facecolor('white')
    fig.tight_layout()
    if fp is not None:
        plt.savefig(fp)
    else:
        plt.show()
    plt.close(fig)

for k in model_dd.keys():
    for i in range(5):
        preds_dd = {'Preds_1': output1_dd[k].numpy()[i],
                    'Preds_2': output2_dd[k].numpy()[i]}
        ic_dd = {'IC_1': dset1.X[i,0].numpy(),
                'IC_2': dset2.X[i,0]}
        soln_dd = {'Soln_1': dset1.X[i,1].numpy(),
                    'Soln_2': dset2.X[i,1].numpy()}
        solns = dset1.X[i, 1].numpy()
        title = 'Test case ' + str(i) + ', model trained on ' + k
        fp_i = os.path.join(PLOTS_DIR, 'compare_predictions_model_{}_test_case_{}.png'.format(model_name_dd[k], i))

        plot_two_solutions_only_DFT(preds_dd, ic_dd, soln_dd, np.linspace(-np.pi, np.pi, 1024), title=title, fp=fp_i)

for k in no_W_model_dd.keys():
    for i in range(5):
        preds_dd = {'Preds_1': no_W_output1_dd[k].numpy()[i],
                    'Preds_2': no_W_output2_dd[k].numpy()[i]}
        ic_dd = {'IC_1': dset1.X[i,0].numpy(),
                'IC_2': dset2.X[i,0].numpy()}
        soln_dd = {'Soln_1': dset1.X[i,1].numpy(),
                    'Soln_2': dset2.X[i,1].numpy()}
        # solns = dset1.X[i, 1].numpy()
        title = 'Test case ' + str(i) + ', No W channel, model trained on ' + k
        fp_i = os.path.join(PLOTS_DIR, 'no_W_compare_predictions_model_{}_test_case_{}.png'.format(model_name_dd[k], i))

        plot_two_solutions_only_DFT(preds_dd, ic_dd, soln_dd, np.linspace(-np.pi, np.pi, 1024), title=title, fp=fp_i)
def make_rescaled_predictions(model, dset):
    """
    """
    # print(ones_input.shape)
    preds = torch.zeros_like(dset.X)
    x_vals = torch.zeros((dset.X.shape[1], dset.X.shape[2]))
    errors = torch.zeros((dset.X.shape[0], dset.X.shape[1]))
    # print(x_vals.shape)
    x_vals[0] = dset.x_grid.reshape((1,-1))
    preds[:, 0] = dset.X[:, 0]

    for t_idx in range(1, dset.n_tsteps+1):
        time = dset.t[t_idx]
        rescaled_ICs = prepare_input(dset.rescaled_ICs[:,t_idx])
        # print(rescaled_ICs.shape)
        # x_vals[t_idx] = rescaled_ICs[0, :,2]
        predictions_i = model(rescaled_ICs)
        # inv_root_t = 1 / torch.sqrt(time)
        root_t = torch.sqrt(time)
        predictions_i = 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
#         print("Finished predictions at ", t_idx, inv_root_t)
    return preds, errors
# def make_composed_predictions(model, dset):
#     """
#     """
#     ones_input = torch.tensor(1.).repeat(dset.n_batches, 1,1)
#     # print(ones_input.shape)
#     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]
#         # rescaled_ICs = dset.make_x_train_rescaled_batched(dset.X[:, 0], time)
#         predictions_i = model(inputs_i, ones_input)
#         preds[:, t_idx] = predictions_i
#         inputs_i = prepare_input(predictions_i)
#         errors_i = l2_normalized_error(predictions_i, dset.X[:,t_idx])
#         errors[:,t_idx] = errors_i

# #         print("Finished predictions at ", t_idx)
#     return preds, errors

preds_dd = {}
errors_dd = {}
with torch.no_grad():
    for k, model in model_dd.items():
        preds_i, errors_i = make_rescaled_predictions(model, scaling_dset)
        preds_dd[k] = preds_i
        errors_dd[k] = errors_i
        print("Finished with ", k)
#     preds_composed, errors_composed = make_composed_predictions(model, time_dset)
#     preds_rescaled, x_vals_rescaled, errors_rescaled = make_rescaled_predictions(model, scaling_dset)
errors_dd_i = {k: np.delete(v.numpy(), [59], axis=0) for k,v in errors_dd.items()}
fp_time_errors = os.path.join(PLOTS_DIR, 'scaling_time_errors.png')
plotting_utils.plot_time_errors(errors_dd_i, title='Time-Rescaling Preds with FNO trained on different ICs') #, fp=fp_time_errors)
test_cases_for_plot = list(range(3))
for test_case in test_cases_for_plot:
    solns = scaling_dset.X.numpy()[test_case]
    for k,v in preds_dd.items():
        fp_i = os.path.join(PLOTS_DIR, 'model_{}_test_case_{}.png'.format(model_name_dd[k], test_case))
        print("Working on model {}, case {}".format(model_name_dd[k], test_case))
        preds_dd_i = {k: v.numpy()[test_case]}
        plotting_utils.plot_one_testcase_panels(preds_dd_i, solns, plot_errors=True, show_n_timesteps=10, fp=fp_i)
#     break

pred_arr = preds_dd['Mixed ICs']
print(pred_arr.shape)

plt.plot(np.real(pred_arr[0,2,:].numpy()))
train_pattern = os.path.join(RESULTS_DIR, '{}_train_FNO_train.txt')
test_pattern = os.path.join(RESULTS_DIR, '{}_train_FNO_test.txt')
for k,v in model_name_dd.items():
    train_fp_i = train_pattern.format(v)
    test_fp_i = test_pattern.format(v)
    train_df = pd.read_table(train_fp_i)
    test_df = pd.read_table(test_fp_i)

    title_i = 'Training set: ' + k
    fp_i = os.path.join(PLOTS_DIR, 'train_test_{}.png'.format(v))

    plotting_utils.make_train_test_plot(train_df, test_df, log_scale=True, title=title_i, fp=fp_i)


DATA_DIR = '/local/meliao/projects/fourier_neural_operator/data/'
NEW_PLOTS_DIR = '/local/meliao/projects/fourier_neural_operator/experiments/18_train_with_rescaling/plots/mixed_IC_model'
if not os.path.isdir(NEW_PLOTS_DIR):
    os.mkdir(NEW_PLOTS_DIR)
test_dset_fp_dd = {'ICs freq [1, ..., 5]': os.path.join(DATA_DIR, '2021-06-24_NLS_data_04_test.mat'),
              'ICs freq [6, ..., 10]': os.path.join(DATA_DIR, '2021-07-22_NLS_data_06_test.mat'),
              'ICs freq [11, ..., 15]': os.path.join(DATA_DIR, '2021-08-04_NLS_data_09_test.mat'),
              'ICs freq [16, ..., 20]': os.path.join(DATA_DIR, '2021-08-04_NLS_data_10_test.mat'),
              'Mixed ICs': os.path.join(DATA_DIR, '2021-08-08_NLS_mixed_IC_data_test.mat'),
              }
test_data_dd = {k: sio.loadmat(v) for k,v in test_dset_fp_dd.items()}

test_dset_dd = {k: TimeScalingDataSet(v['output'], v['t'], v['x']) for k,v in test_data_dd.items()}
preds_dd = {}
errors_dd = {}
mixed_model = model_dd['Mixed ICs']
with torch.no_grad():
    for k, dset in test_dset_dd.items():
        preds_i, _ , errors_i = make_rescaled_predictions(mixed_model, dset)
        preds_dd[k] = preds_i
        errors_dd[k] = errors_i
        print("Finished with ", k)
#     preds_composed, errors_composed = make_composed_predictions(model, time_dset)
#     preds_rescaled, x_vals_rescaled, errors_rescaled = make_rescaled_predictions(model, scaling_dset)
errors_dd_i = {k: v.numpy() for k,v in errors_dd.items()}
t = 'Model trained on Mixed ICs and tested on different datasets'
fp = os.path.join(NEW_PLOTS_DIR, 'mixed_ICs_time_errors.png')
plotting_utils.plot_time_errors(errors_dd_i, title=t, fp=fp)
test_cases_for_plot = list(range(3))
for test_case in test_cases_for_plot:
    for k, dset in test_dset_dd.items():
        solns = dset.X.numpy()[test_case]
        preds_dd_i = {k: preds_dd[k].numpy()[test_case]}
        fp_i = os.path.join(NEW_PLOTS_DIR, 'panels_dset_{}_test_case_{}.png'.format(model_name_dd[k], test_case))
        plotting_utils.plot_one_testcase_panels(preds_dd_i, solns, show_n_timesteps=10, fp=fp_i)
        print("Finished dset {} and test case {}".format(model_name_dd[k], test_case))
    # break<jupyter_output>Finished dset 00 and test case 0
Finished dset 01 and test case 0
Finished dset 02 and test case 0
Finished dset 03 and test case 0
Finished dset 04 and test case 0
Finished dset 00 and test case 1
Finished dset 01 and test case 1
Finished dset 02 and test case 1
Finished dset 03 and test case 1
Finished dset 04 and test case 1
Finished dset 00 and test case 2
Finished dset 01 and test case 2
Finished dset 02 and test case 2
Finished dset 03 and test case 2
Finished dset 04 and test case 2
{'fraction_non_alphanumeric': 0.15135000208012647, 'fraction_numerical': 0.02745766942630112, 'mean_word_length': 3.435873777449714, 'pattern_counts': {'":': 0, '<': 19, '<?xml version=': 0, '>': 19, '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': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '512699', 'license': 'MIT', 'n_tokens_mistral': 10122, 'n_tokens_neox': 9712, 'n_words': 1678, 'path': 'experiments/20_investigate_frequency_response/Compare_n_x_grid.ipynb', 'repo_name': 'meliao/fourier_neural_operator', 'size': '24037'}