Add all required modules and requirements.txt

LSTM_model.py

@@ -0,0 +1,249 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+# In[1]:
+
+
+import numpy as np
+import torch
+import torch.nn as nn
+import time
+import math
+import torch
+
+num_time_steps = 500
+x = np.linspace(0.0,1.0,num=128)
+dx = 1.0/np.shape(x)[0]
+tsteps = np.linspace(0.0,2.0,num=num_time_steps)
+dt = 2.0/np.shape(tsteps)[0]
+
+class AE_Encoder(nn.Module):
+    def __init__(self, input_dim, latent_dim=2, feats=[512, 256, 128, 64, 32]):
+        super(AE_Encoder, self).__init__()
+        self.latent_dim = latent_dim
+        self._net = nn.Sequential(
+            nn.Linear(input_dim, feats[0]),
+            nn.GELU(),
+            nn.Linear(feats[0], feats[1]),
+            nn.GELU(),
+            nn.Linear(feats[1], feats[2]),
+            nn.GELU(),
+            nn.Linear(feats[2], feats[3]),
+            nn.GELU(),
+            nn.Linear(feats[3], feats[4]),
+            nn.GELU(),
+            nn.Linear(feats[4], latent_dim)
+        )
+
+    def forward(self, x):
+      Z = self._net(x)
+      return Z
+
+
+class AE_Decoder(nn.Module):
+    def __init__(self, latent_dim, output_dim, feats=[32, 64, 128, 256, 512]):
+        super(AE_Decoder, self).__init__()
+        self.output_dim = output_dim
+        self._net = nn.Sequential(
+            nn.Linear(latent_dim, feats[0]),
+            nn.GELU(),
+            nn.Linear(feats[0], feats[1]),
+            nn.GELU(),
+            nn.Linear(feats[1], feats[2]),
+            nn.GELU(),
+            nn.Linear(feats[2], feats[3]),
+            nn.GELU(),
+            nn.Linear(feats[3], feats[4]),
+            nn.GELU(),
+            nn.Linear(feats[4], output_dim),
+        )
+
+    def forward(self, x):
+      y = self._net(x)
+      return y
+
+
+class AE_Model(nn.Module):
+    def __init__(self, encoder, decoder):
+        super(AE_Model, self).__init__()
+        self.encoder = encoder
+        self.decoder = decoder # decoder for x(t)
+
+    def forward(self, x):
+        z = self.encoder(x)
+        # Reconstruction
+        x_hat = self.decoder(z)  # Reconstruction of x(t)
+
+        return x_hat
+        
+        
+class PytorchLSTM(nn.Module):
+    def __init__(self, input_dim=3, hidden_dim=40, output_dim=2):
+        super().__init__()
+        # First LSTM: simulates return_sequences=True
+        self.lstm1 = nn.LSTM(input_dim, hidden_dim, batch_first=True)
+        # Second LSTM: simulates return_sequences=False
+        self.lstm2 = nn.LSTM(hidden_dim, hidden_dim, batch_first=True)
+        # Dense layer
+        self.fc = nn.Linear(hidden_dim, output_dim)
+
+    def forward(self, x):
+        """
+        x shape: [batch_size, time_window, input_dim]
+        """
+        # LSTM1 (return_sequences=True)
+        out1, (h1, c1) = self.lstm1(x)
+        # out1 shape: [batch_size, time_window, hidden_dim]
+
+        # LSTM2 (return_sequences=False -> we only use the last time step)
+        out2, (h2, c2) = self.lstm2(out1)
+        # out2 shape: [batch_size, time_window, hidden_dim]
+        # Last timestep (since we didn't set return_sequences=True)
+        # is effectively out2[:, -1, :], but PyTorch LSTM always returns full seq unless you slice.
+
+        last_timestep = out2[:, -1, :]  # shape: [batch_size, hidden_dim]
+
+        # Dense -> 2 outputs
+        output = self.fc(last_timestep)  # shape: [batch_size, 2]
+        return output
+        
+def measure_lstm_prediction_time(
+    decoder,
+    lstm_model,
+    lstm_testing_data,
+    sim_num,
+    final_time,
+    time_window=40
+):
+    """
+    Predicts up to `final_time` in a walk-forward manner for simulation `sim_num`,
+    measures the elapsed time, and returns the final predicted latent + the true latent.
+
+    Parameters
+    ----------
+    decoder : torch.nn.Module
+    	The trained weights of the decoder
+    model : torch.nn.Module
+        Trained PyTorch LSTM model. We'll set model.eval() inside.
+    lstm_testing_data : np.ndarray
+        Shape (num_test_snapshots, num_time_steps, 3).
+        The last dimension typically holds (2 latents + 1 param) or similar.
+    sim_num : int
+        Which simulation index to use (e.g., 0 for the first).
+    final_time : int
+        The final timestep index you want to predict up to (>= time_window).
+        For example, if time_window=10 and final_time=20, we will predict from t=10..19.
+    time_window : int
+        Size of the rolling window (default=40).
+
+    Returns
+    -------
+    float
+        Elapsed time (seconds) for performing the predictions from t=time_window up to t=final_time.
+    np.ndarray
+        The final predicted latent at time=final_time (shape (2,)).
+    np.ndarray
+        The true latent at time=final_time (shape (2,)).
+    """
+
+    # Basic shape info
+    num_time_steps = lstm_testing_data.shape[1]
+    if final_time > num_time_steps:
+        raise ValueError(
+            f"final_time={final_time} exceeds available time steps={num_time_steps}."
+        )
+    if final_time < time_window:
+        raise ValueError(
+            f"final_time={final_time} is less than time_window={time_window}, no prediction needed."
+        )
+
+    # Initialize the rolling window with first `time_window` steps
+    input_seq = np.zeros((1, time_window, 3), dtype=np.float32)
+    input_seq[0, :, :] = lstm_testing_data[sim_num, 0:time_window, :]
+
+    lstm_model.eval()  # inference mode
+
+    final_pred = None  # store the final predicted latent
+    start_time = time.time()
+
+    with torch.no_grad():
+        # Predict from t=time_window to t=final_time-1
+        # so that at the end of the loop we've generated a prediction for index final_time.
+        # If you want the model's prediction at final_time itself, we do a loop up to final_time.
+        for t in range(time_window, final_time):
+            inp_tensor = torch.from_numpy(input_seq).float()  # shape [1, 10, 3]
+            pred = lstm_model(inp_tensor)  # shape [1, 2]
+            pred_np = pred.numpy()[0, :]  # shape (2,)
+
+            # Shift the rolling window
+            temp = input_seq[0, 1:time_window, :].copy()
+            input_seq[0, 0:time_window - 1, :] = temp
+            input_seq[0, time_window - 1, 0:2] = pred_np
+
+            # Keep track of the last prediction
+            final_pred = pred_np
+            
+    x_hat_tau_pred = decoder(torch.tensor(final_pred, dtype = torch.float32))
+    
+    end_time = time.time()
+    
+    elapsed = end_time - start_time
+
+    # final_pred is the LSTM's predicted latent for step `final_time`.
+    # The *true* latent at that time is:
+    final_true = lstm_testing_data[sim_num, final_time, 0:2]  # shape (2,)
+    
+    return elapsed, final_pred, final_true
+    
+    
+def collect_snapshots(Rnum):
+    snapshot_matrix = np.zeros(shape=(np.shape(x)[0],np.shape(tsteps)[0]))
+
+    trange = np.arange(np.shape(tsteps)[0])
+    for t in trange:
+        snapshot_matrix[:,t] = exact_solution(Rnum,tsteps[t])[:]
+
+    return snapshot_matrix
+
+def collect_multiparam_snapshots_train():
+    rnum_vals = np.arange(900,2900,100)
+    
+    rsnap = 0
+    for rnum_val in rnum_vals:
+        snapshots_temp = np.transpose(collect_snapshots(rnum_val))
+        
+        if rsnap == 0:
+            all_snapshots = snapshots_temp
+        else:
+            
+            all_snapshots = np.concatenate((all_snapshots,snapshots_temp),axis=0)
+            
+        rsnap = rsnap + 1    
+    return all_snapshots, rnum_vals/1000
+
+def collect_multiparam_snapshots_test():
+    rnum_vals = np.arange(1050,2850,200)
+    
+    rsnap = 0
+    for rnum_val in rnum_vals:
+        snapshots_temp = np.transpose(collect_snapshots(rnum_val))
+        
+        if rsnap == 0:
+            all_snapshots = snapshots_temp
+        else:
+            
+            all_snapshots = np.concatenate((all_snapshots,snapshots_temp),axis=0)
+            
+        rsnap = rsnap + 1    
+    return all_snapshots, rnum_vals/1000
+    
+    
+
+    return elapsed, final_pred, final_true
+
+def exact_solution(Rnum,t):
+    x = np.linspace(0.0,1.0,num=128)
+    t0 = np.exp(Rnum/8.0)
+    return (x/(t+1))/(1.0+np.sqrt((t+1)/t0)*np.exp(Rnum*(x*x)/(4.0*t+4)))
+    
+        

config_adv_dif.py

@@ -0,0 +1,25 @@
+from dataclasses import dataclass
+import json
+
+@dataclass
+class Config:
+    # default values. DO NOT TOUCH
+    name: str = 'FlexiPropagator_2D'
+    latent_dim: int = 3
+    batch_size: int = 64
+    lr: float = 3e-4
+    num_epochs: int = 25
+    num_time_steps: int = 500
+
+    gamma: float = 3.25
+    beta: float = 1e-3
+
+    val_every: float = 0.25
+    plot_train_every: float = 0.01
+
+    save_dir: str = 'checkpoints'
+
+def load_config(path):
+    with open(path, 'r') as f:
+        config = json.load(f)
+    return Config(**config)

config_burgers.py

@@ -0,0 +1,31 @@
+from dataclasses import dataclass
+import json
+
+@dataclass
+class Config:
+    # default values. DO NOT TOUCH
+    name: str = 'FlexiPropagator'
+    latent_dim: int = 2
+    input_dim: int = 128
+    batch_size: int = 128
+    lr: float = 3e-4
+    num_epochs: int = 200
+    n_samples_train: int = 8_00_000
+
+    num_time_steps: int = 500
+
+    tau_left_fraction: float = 0.35
+    tau_right_fraction: float = 0.85
+
+    gamma: float = 3.25
+    beta: float = 1e-4
+
+    val_every: float = 0.25
+    plot_train_every: float = 0.01
+
+    save_dir: str = 'checkpoints'
+
+def load_config(path):
+    with open(path, 'r') as f:
+        config = json.load(f)
+    return Config(**config)

data_adv_dif.py

@@ -0,0 +1,261 @@
+import os
+import pickle
+import numpy as np
+import matplotlib.pyplot as plt
+import torch
+from dataclasses import dataclass, asdict
+import json
+
+
+# Rnum = 1000
+num_time_steps = 500
+
+def get_dt(num_time_steps):
+    return 2.0/num_time_steps
+
+dt = get_dt(num_time_steps)
+
+
+def exact_solution(alpha, t, L=2.0, Nx=128, Ny=128, c=1.0):
+    nu = 1.0 / alpha
+    x_vals = np.linspace(-L, L, Nx)
+    y_vals = np.linspace(-L, L, Ny)
+    X, Y = np.meshgrid(x_vals, y_vals)
+    if t <= 0:
+        return np.zeros_like(X)
+    rx = X - c * t
+    ry = Y
+    r2 = rx**2 + ry**2
+    denominator = 4.0 * nu * t
+    amplitude = 1.0 / (4.0 * np.pi * nu * t)
+    U = amplitude * np.exp(-r2 / denominator)
+    return U
+
+
+class AdvectionDiffussionDataset:
+    def __init__(self, 
+                X: np.ndarray = None,
+                X_tau: np.ndarray = None,
+                t_values: np.ndarray = None,
+                tau_values: np.ndarray = None,
+                alpha_values: np.ndarray = None):
+        self.X = X
+        self.X_tau = X_tau
+        self.t_values = t_values
+        self.tau_values = tau_values
+        self.alpha_values = alpha_values
+
+    def append(self, other):
+        self.X = np.concatenate([self.X, other.X]) if self.X is not None else other.X
+        self.X_tau = np.concatenate([self.X_tau, other.X_tau]) if self.X_tau is not None else other.X_tau
+        self.t_values = np.concatenate([self.t_values, other.t_values]) if self.t_values is not None else other.t_values
+        self.tau_values = np.concatenate([self.tau_values, other.tau_values]) if self.tau_values is not None else other.tau_values
+        self.alpha_values = np.concatenate([self.alpha_values, other.alpha_values]) if self.alpha_values is not None else other.alpha_values
+
+@dataclass
+class IntervalSplit:
+    interpolation: tuple
+    extrapolation_left: tuple
+    extrapolation_right: tuple
+
+def prepare_adv_diff_dataset(alpha_range=(0.01, 10), tau_range=(150, 400), dt=dt, n_samples=500):
+    X = []
+    X_tau = []
+    t_values = []
+    tau_values = []
+    alpha_values = []
+    TRANGE = (0.01, 2.0)
+    while len(X) < n_samples:
+        # sample alpha uniformly
+        alpha = np.random.uniform(*alpha_range)
+        t = np.random.uniform(*TRANGE)
+        x_t = exact_solution(alpha, t)
+        tau = np.random.randint(*tau_range)
+        x_tau = exact_solution(alpha, t+(tau*dt))
+        
+        X.append(x_t)
+        X_tau.append(x_tau)
+        t_values.append(t)
+        tau_values.append(tau)
+        alpha_values.append(alpha)
+
+    X = np.array(X)
+    X_tau = np.array(X_tau)
+    t_values = np.array(t_values)
+    tau_values = np.array(tau_values)
+    alpha_values = np.array(alpha_values)
+    dataset = AdvectionDiffussionDataset(X, X_tau, t_values, tau_values, alpha_values)
+    return dataset
+   
+
+def train_test_split_range(interval, interpolation_span=0.1, extrapolation_left_span=0.1, extrapolation_right_span=0.1):
+    """
+    Split the range into train and test ranges
+    We have three test folds:
+    1. Interpolation fold: Re and tau values are within the training (min, max) range but not in the training set
+        We sample an interval of length x_interpolation_span% randomly from the total range
+    2. Extrapolation fold: Re and tau values are outside the training (min, max) range
+        We sample two intervals of length x_extrapolation_right_span% and x_extrapolation_left_span% from the total range
+    3. Validation fold: Re and tau values are randomly sampled from the total set
+
+    Overall interval looks like:
+    Extrapolation_left_test | normal | Interpolation_test | normal | Extrapolation_right_test
+    (min, extrapolation_left) | (extraplation_left, interpolation_min) | (interpolation_min, interpolation_max) | (interpolation_max, extrapolation_right) | (extrapolation_right, max)
+    and
+    train, val = split(normal, val_split)
+    """
+    r_min, r_max = interval
+    length = (r_max-r_min)
+    extra_left_length = extrapolation_left_span * length
+    extra_right_length = extrapolation_right_span * length
+    inter_length = interpolation_span * length
+
+    extrapolation_left = (r_min, r_min + extra_left_length)
+    extrapolation_right = (r_max - extra_right_length, r_max)
+
+    interpolation_min = np.random.uniform(extrapolation_left[1], extrapolation_right[0] - inter_length)
+    interpolation = (interpolation_min, interpolation_min + inter_length)
+
+    train_ranges = [(extrapolation_left[1], interpolation[0]), (interpolation[1], extrapolation_right[0])]
+    return IntervalSplit(interpolation, extrapolation_left, extrapolation_right), train_ranges
+
+def get_train_ranges(interval_split):
+    return [
+        (interval_split.extrapolation_left[1], interval_split.interpolation[0]),
+        (interval_split.interpolation[1], interval_split.extrapolation_right[0])
+    ]
+
+def get_train_val_test_folds(alpha_range, tau_range,
+                             alpha_interpolation_span=0.10,
+                             alpha_extrapolation_left_span=0.10,
+                             alpha_extrapolation_right_span=0.10,
+                             tau_interpolation_span=0.10,
+                             tau_extrapolation_left_span=0.10,
+                             tau_extrapolation_right_span=0.10,
+                             n_samples_train=500,
+                             n_samples_val=200):
+    """
+    Generate train (4 sub-regions) and val (left extrp, interp, right extrp
+    for alpha x left extrp, interp, right extrp for tau) datasets.
+    
+    Returns:
+        dataset_train  : AdvectionDiffussionDataset
+        dataset_val    : AdvectionDiffussionDataset
+        alpha_interval_split: IntervalSplit
+        tau_interval_split  : IntervalSplit
+    """
+
+    # ---------------------------------------------------------------------
+    # 1) Split alpha into 4 regions: left extrp, interp, right extrp, train
+    # 2) Split tau   into 4 regions: left extrp, interp, right extrp, train
+    # ---------------------------------------------------------------------
+    alpha_interval_split, alpha_train_ranges = train_test_split_range(
+        alpha_range,
+        alpha_interpolation_span,
+        alpha_extrapolation_left_span,
+        alpha_extrapolation_right_span
+    )
+    tau_interval_split, tau_train_ranges = train_test_split_range(
+        tau_range,
+        tau_interpolation_span,
+        tau_extrapolation_left_span,
+        tau_extrapolation_right_span
+    )
+
+    # alpha_train_ranges and tau_train_ranges each have 2 intervals:
+    #   alpha_train_ranges = [ (a1_lo, a1_hi), (a2_lo, a2_hi) ]
+    #   tau_train_ranges   = [ (t1_lo, t1_hi), (t2_lo, t2_hi) ]
+    #
+    # Meanwhile, alpha_interval_split has:
+    #   alpha_interval_split.extrapolation_left  = (a_left_lo, a_left_hi)
+    #   alpha_interval_split.interpolation       = (a_int_lo, a_int_hi)
+    #   alpha_interval_split.extrapolation_right = (a_right_lo, a_right_hi)
+    # and similarly for tau_interval_split.
+
+    # -------------------------------------------------------------
+    # 3) Build the TRAIN dataset from the Cartesian product
+    #    of alpha_train_ranges x tau_train_ranges => 4 combos
+    # -------------------------------------------------------------
+    dataset_train = AdvectionDiffussionDataset()
+    for alpha_subrange in alpha_train_ranges:  # 2 intervals
+        for tau_subrange in tau_train_ranges:  # 2 intervals
+            subset = prepare_adv_diff_dataset(
+                alpha_range=alpha_subrange,
+                tau_range=tau_subrange,
+                n_samples=n_samples_train
+            )
+            dataset_train.append(subset)
+
+    # -------------------------------------------------------------
+    # 4) Build the VAL dataset from the leftover intervals:
+    #    alpha in { left extrp, interp, right extrp }
+    #  x tau   in { left extrp, interp, right extrp } => up to 9 combos
+    # -------------------------------------------------------------
+    alpha_val_intervals = [
+        alpha_interval_split.extrapolation_left,
+        alpha_interval_split.interpolation,
+        alpha_interval_split.extrapolation_right
+    ]
+    tau_val_intervals = [
+        tau_interval_split.extrapolation_left,
+        tau_interval_split.interpolation,
+        tau_interval_split.extrapolation_right
+    ]
+
+    dataset_val = AdvectionDiffussionDataset()
+
+    for a_val_range in alpha_val_intervals:
+        for t_val_range in tau_val_intervals:
+            subset_val = prepare_adv_diff_dataset(
+                alpha_range=a_val_range,
+                tau_range=t_val_range,
+                n_samples=n_samples_val
+            )
+            dataset_val.append(subset_val)
+
+    return dataset_train, dataset_val, alpha_interval_split, tau_interval_split
+
+
+def plot_sample(dataset, i):
+    """
+    Plot a sample pair from the dataset.
+    """
+    X = dataset.X
+    X_tau = dataset.X_tau
+    t_values = dataset.t_values
+    tau_values = dataset.tau_values
+    alpha_values = dataset.alpha_values
+
+    print("Shape of X:", X.shape)
+
+    fig, axs = plt.subplots(1, 2, figsize=(12, 5))
+    im1 = axs[0].imshow(X[i], extent=[0, 1, 0, 1], origin='lower', cmap='hot')
+    axs[0].set_title(f'Initial State (t: {t_values[i]})')
+    plt.colorbar(im1, ax=axs[0])
+
+    im2 = axs[1].imshow(X_tau[i], extent=[0, 1, 0, 1], origin='lower', cmap='hot')
+    axs[1].set_title(f'Shifted State (t + tau): {t_values[i]+tau_values[i]*dt}')
+    plt.colorbar(im2, ax=axs[1])
+
+    fig.suptitle(f'Tau: {tau_values[i]}, Alpha: {alpha_values[i]:.4f}')
+    plt.show()
+    
+    
+def load_from_path(path):
+    dataset_train_path = os.path.join(path, 'dataset_train.pkl')
+    dataset_val_path = os.path.join(path, 'dataset_val.pkl')
+    alpha_interval_split_path = os.path.join(path, 'alpha_interval_split.json')
+    tau_interval_split_path = os.path.join(path, 'tau_interval_split.json')
+
+    with open(dataset_train_path, 'rb') as f:
+        dataset_train = pickle.load(f)
+    with open(dataset_val_path, 'rb') as f:
+        dataset_val = pickle.load(f)
+    with open(alpha_interval_split_path, 'r') as f:
+        alpha_interval_split = json.load(f)
+        alpha_interval_split = IntervalSplit(**alpha_interval_split)
+    with open(tau_interval_split_path, 'r') as f:
+        tau_interval_split = json.load(f)
+        tau_interval_split = IntervalSplit(**tau_interval_split)
+
+    return dataset_train, dataset_val, alpha_interval_split, tau_interval_split

data_burgers.py

@@ -0,0 +1,243 @@
+
+import os
+import pickle
+import numpy as np
+import matplotlib.pyplot as plt
+import torch
+from dataclasses import dataclass, asdict
+import json
+
+
+# Rnum = 1000
+num_time_steps = 500
+# # x = np.linspace(0.0,1.0,num=128)
+# dx = 1.0/np.shape(x)[0]
+# TSTEPS = np.linspace(0.0,2.0,num=num_time_steps)
+# dt = 2.0/np.shape(TSTEPS)[0]
+
+def get_dt(num_time_steps):
+    return 2.0/num_time_steps
+
+dt = get_dt(num_time_steps)
+
+
+def exact_solution(Rnum,t):
+    x = np.linspace(0.0,1.0,num=128)
+    t0 = np.exp(Rnum/8.0)
+    return (x/(t+1))/(1.0+np.sqrt((t+1)/t0)*np.exp(Rnum*(x*x)/(4.0*t+4)))
+
+
+class ReDataset:
+    def __init__(self, 
+                X: np.ndarray = None,
+                X_tau: np.ndarray = None,
+                t_values: np.ndarray = None,
+                tau_values: np.ndarray = None,
+                Re_values: np.ndarray = None):
+        self.X = X
+        self.X_tau = X_tau
+        self.t_values = t_values
+        self.tau_values = tau_values
+        self.Re_values = Re_values
+
+    def append(self, other):
+        self.X = np.concatenate([self.X, other.X]) if self.X is not None else other.X
+        self.X_tau = np.concatenate([self.X_tau, other.X_tau]) if self.X_tau is not None else other.X_tau
+        self.t_values = np.concatenate([self.t_values, other.t_values]) if self.t_values is not None else other.t_values
+        self.tau_values = np.concatenate([self.tau_values, other.tau_values]) if self.tau_values is not None else other.tau_values
+        self.Re_values = np.concatenate([self.Re_values, other.Re_values]) if self.Re_values is not None else other.Re_values
+
+@dataclass
+class IntervalSplit:
+    interpolation: tuple
+    extrapolation_left: tuple
+    extrapolation_right: tuple
+
+def get_time_shifts(snapshots, tau_range=(100, 500), n_samples=100):
+    X = []
+    X_tau = []
+    tau_values = []
+    while len(X) < n_samples:
+        tau = np.random.randint(*tau_range)
+        i = np.random.randint(0, len(snapshots)-tau)
+        X.append(snapshots[i])
+        X_tau.append(snapshots[i+tau])
+        tau_values.append(tau)
+    X = np.array(X)
+    X_tau = np.array(X_tau)
+    tau_values = np.array(tau_values)
+    return X, X_tau, tau_values
+
+def prepare_Re_dataset(Re_range=(100, 2000), tau_range=(500, 1900), dt=dt, n_samples=5000):
+    X = []
+    X_tau = []
+    t_values = []
+    tau_values = []
+    Re_values = []
+    TRANGE = (0,2)
+    while len(X) < n_samples:
+        # sample Re log uniformly
+        logRe = np.random.uniform(np.log(Re_range[0]), np.log(Re_range[1]))
+        Re = np.exp(logRe).round().astype(int)
+        t = np.random.uniform(*TRANGE)
+        x_t = exact_solution(Re, t)
+        # print('tau_range', tau_range)
+        tau = np.random.randint(*tau_range)
+        x_tau = exact_solution(Re, t+(tau*dt))
+        
+        X.append(x_t)
+        X_tau.append(x_tau)
+        t_values.append(t)
+        tau_values.append(tau)
+        Re_values.append(Re)
+
+    X = np.array(X)
+    X_tau = np.array(X_tau)
+    t_values = np.array(t_values)
+    tau_values = np.array(tau_values)
+    Re_values = np.array(Re_values)
+    # return X, X_tau, tau_values, Re_values
+    dataset = ReDataset(X, X_tau, t_values, tau_values, Re_values)
+    return dataset
+
+
+def train_test_split_range(interval, interpolation_span=0.1, extrapolation_left_span=0.1, extrapolation_right_span=0.1):
+    """
+    Split the range into train and test ranges
+    We have three test folds:
+    1. Interpolation fold: Re and tau values are within the training (min, max) range but not in the training set
+        We sample an interval of length x_interpolation_span% randomly from the total range
+    2. Extrapolation fold: Re and tau values are outside the training (min, max) range
+        We sample two intervals of length x_extrapolation_right_span% and x_extrapolation_left_span% from the total range
+    3. Validation fold: Re and tau values are randomly sampled from the total set
+
+    Overall interval looks like:
+    Extrapolation_left_test | normal | Interpolation_test | normal | Extrapolation_right_test
+    (min, extrapolation_left) | (extraplation_left, interpolation_min) | (interpolation_min, interpolation_max) | (interpolation_max, extrapolation_right) | (extrapolation_right, max)
+    and
+    train, val = split(normal, val_split)
+    """
+    r_min, r_max = interval
+    length = (r_max-r_min)
+    extra_left_length = extrapolation_left_span * length
+    extra_right_length = extrapolation_right_span * length
+    inter_length = interpolation_span * length
+
+    extrapolation_left = (r_min, r_min + extra_left_length)
+    extrapolation_right = (r_max - extra_right_length, r_max)
+
+    interpolation_min = np.random.uniform(extrapolation_left[1], extrapolation_right[0] - inter_length)
+    interpolation = (interpolation_min, interpolation_min + inter_length)
+
+    train_ranges = [(extrapolation_left[1], interpolation[0]), (interpolation[1], extrapolation_right[0])]
+    return IntervalSplit(interpolation, extrapolation_left, extrapolation_right), train_ranges
+
+def get_train_ranges(interval_split):
+    return [
+        (interval_split.extrapolation_left[1], interval_split.interpolation[0]),
+        (interval_split.interpolation[1], interval_split.extrapolation_right[0])
+    ]
+
+# def get_dataset_from_ranges(train_ranges):
+#     dataset = ReDataset()
+#     for re_train_range, tau_train_range in zip(Re_train_ranges, tau_train_ranges):
+#         train_dataset = prepare_Re_dataset(Re_range=re_train_range, tau_range=tau_train_range, n_samples=n_samples_train)
+#         dataset.append(train_dataset)
+    
+#     return dataset
+
+
+def get_train_val_test_folds(Re_range, tau_range,
+                        re_interpolation_span=0.10,
+                        re_extrapolation_left_span=0.1,
+                        re_extrapolation_right_span=0.10,
+                        tau_interpolation_span=0.10,
+                        tau_extrapolation_left_span=0.1,
+                        tau_extrapolation_right_span=0.10,
+                        n_samples_train=1000000,
+                        val_split=0.2):
+    
+    Re_interval_split, Re_train_ranges = train_test_split_range(Re_range, re_interpolation_span, re_extrapolation_left_span, re_extrapolation_right_span)
+    tau_interval_split, tau_train_ranges = train_test_split_range(tau_range, tau_interpolation_span, tau_extrapolation_left_span, tau_extrapolation_right_span)
+
+    # print(Re_interval_split, Re_train_ranges)
+    # print(tau_interval_split, tau_train_ranges)
+    # prepare train dataset
+    dataset = ReDataset()
+    for re_train_range, tau_train_range in zip(Re_train_ranges, tau_train_ranges):
+        train_dataset = prepare_Re_dataset(Re_range=re_train_range, tau_range=tau_train_range, n_samples=n_samples_train)
+        dataset.append(train_dataset)
+
+    inds = np.arange(len(dataset.X))
+    np.random.shuffle(inds)
+    train_inds = inds[:int(len(inds)*(1-val_split))]
+    val_inds = inds[int(len(inds)*(1-val_split)):]
+    dataset_train = ReDataset(dataset.X[train_inds], dataset.X_tau[train_inds], dataset.t_values[train_inds], dataset.tau_values[train_inds], dataset.Re_values[train_inds])
+    dataset_val = ReDataset(dataset.X[val_inds], dataset.X_tau[val_inds],dataset.t_values[val_inds], dataset.tau_values[val_inds], dataset.Re_values[val_inds])
+    return dataset_train, dataset_val, Re_interval_split, tau_interval_split
+
+def plot_sample(dataset, i):
+    X = dataset.X
+    X_tau = dataset.X_tau
+    Tau = dataset.tau_values
+    Re_total = dataset.Re_values
+    plt.plot(X[i], label = "Initial State")
+    plt.plot(X_tau[i], label = "Mapped State")
+    plt.title(f'Tau: {Tau[i]}, Re: {Re_total[i]}')
+    plt.legend()
+    plt.show()
+
+def save_to_path(path, dataset_train, dataset_val, Re_interval_split, tau_interval_split):
+    if not os.path.exists(path):
+        os.makedirs(path)
+    # save dataset_train, dataset_val, Re_interval_split, tau_interval_split to pkl files
+    dataset_train_path = os.path.join(path, 'dataset_train.pkl')
+    dataset_val_path = os.path.join(path, 'dataset_val.pkl')
+    Re_interval_split_path = os.path.join(path, 'Re_interval_split.json')
+    tau_interval_split_path = os.path.join(path, 'tau_interval_split.json')
+
+    with open(dataset_train_path, 'wb') as f:
+        pickle.dump(dataset_train, f)
+    with open(dataset_val_path, 'wb') as f:
+        pickle.dump(dataset_val, f)
+
+    with open(Re_interval_split_path, 'w') as f:
+        json.dump(asdict(Re_interval_split), f)
+    with open(tau_interval_split_path, 'w') as f:
+        json.dump(asdict(tau_interval_split), f)
+
+def load_from_path(path):
+    dataset_train_path = os.path.join(path, 'dataset_train.pkl')
+    dataset_val_path = os.path.join(path, 'dataset_val.pkl')
+    Re_interval_split_path = os.path.join(path, 'Re_interval_split.json')
+    tau_interval_split_path = os.path.join(path, 'tau_interval_split.json')
+
+    with open(dataset_train_path, 'rb') as f:
+        dataset_train = pickle.load(f)
+    with open(dataset_val_path, 'rb') as f:
+        dataset_val = pickle.load(f)
+    with open(Re_interval_split_path, 'r') as f:
+        Re_interval_split = json.load(f)
+        Re_interval_split = IntervalSplit(**Re_interval_split)
+    with open(tau_interval_split_path, 'r') as f:
+        tau_interval_split = json.load(f)
+        tau_interval_split = IntervalSplit(**tau_interval_split)
+
+    return dataset_train, dataset_val, Re_interval_split, tau_interval_split
+
+
+def main():
+    #Re_range = (100, 3000)
+    #num_time_steps = 500
+    #tau_range = (175, 425)
+    #dataset_train, dataset_val, Re_interval_split, tau_interval_split = get_train_val_test_folds(Re_range, tau_range)
+    #save_to_path('data', dataset_train, dataset_val, Re_interval_split, tau_interval_split)
+    
+    
+
+    load_from_path('data')
+
+
+if __name__ == '__main__':
+    main()
+

model_adv_dif.py

@@ -0,0 +1,190 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+# In[1]:
+
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+import math
+import torch
+
+import os
+import torch
+import torch.nn as nn
+import numpy as np
+import pickle
+from dataclasses import dataclass, asdict
+import json
+from torch.utils.data import DataLoader
+
+
+# Normalization Layer for Conv2D
+class Norm(nn.Module):
+    def __init__(self, num_channels, num_groups=4):
+        super(Norm, self).__init__()
+        self.norm = nn.GroupNorm(num_groups, num_channels)
+
+    def forward(self, x):
+        return self.norm(x)
+
+# Encoder using Conv2D
+class Encoder(nn.Module):
+    def __init__(self, latent_dim=3):
+        super(Encoder, self).__init__()
+        self.conv_layers = nn.Sequential(
+            # Input: (batch_size, 1, 256, 256)
+            nn.Conv2d(1, 32, kernel_size=2, stride=2, padding=0),  # (batch_size, 64, 128, 128)
+            nn.GELU(),
+            Norm(32),
+            nn.Conv2d(32, 64, kernel_size=2, stride=2, padding=0),  # (batch_size, 128, 64, 64)
+            nn.GELU(),
+            Norm(64),
+            nn.Conv2d(64, 128, kernel_size=2, stride=2, padding=0),  # (batch_size, 256, 32, 32)
+            nn.GELU(),
+            Norm(128),
+            nn.Conv2d(128, 256, kernel_size=2, stride=2, padding=0),  # (batch_size, 512, 16, 16)
+            nn.GELU(),
+            Norm(256),
+            nn.Conv2d(256, 512, kernel_size=2, stride=2, padding=0),  # (batch_size, 512, 8, 8)
+            nn.GELU(),
+            Norm(512),
+        )
+        self.flatten = nn.Flatten()
+        self.fc_mean = nn.Linear(512 * 4 * 4, latent_dim)
+        self.fc_log_var = nn.Linear(512 * 4 * 4, latent_dim)
+
+    def forward(self, x):
+        x = self.conv_layers(x)
+        x = self.flatten(x)
+        mean = self.fc_mean(x)
+        log_var = self.fc_log_var(x)
+        return mean, log_var
+        
+        
+        
+class Decoder(nn.Module):
+    def __init__(self, latent_dim=3):
+        super(Decoder, self).__init__()
+        # Fully connected layer to transform the latent vector back to the shape (batch_size, 512, 8, 8)
+        self.fc = nn.Linear(latent_dim, 512 * 4 * 4)
+
+        self.deconv_layers = nn.Sequential(
+            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
+            nn.Conv2d(512, 256, kernel_size=1),
+            nn.GELU(),
+            Norm(256),
+
+
+            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
+            nn.Conv2d(256, 128, kernel_size=1),
+            nn.GELU(),
+            Norm(128),
+
+            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
+            nn.Conv2d(128, 64, kernel_size=1),
+            nn.GELU(),
+            Norm(64),
+
+            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
+            nn.Conv2d(64, 32, kernel_size=1),
+            nn.GELU(),
+            Norm(32),
+
+            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
+            nn.Conv2d(32, 1, kernel_size=1),
+            nn.ReLU()
+        )
+
+    def forward(self, z):
+        # Transform the latent vector to match the shape of the feature maps
+        x = self.fc(z)
+        x = x.view(-1, 512, 4, 4)  # Reshape to (batch_size, 512, 4, 4)
+        x = self.deconv_layers(x)
+        return x
+        
+        
+class Propagator_concat(nn.Module): 
+    """
+    Takes in (z(t), tau, alpha) and outputs z(t+tau)
+    """
+    def __init__(self, latent_dim, feats=[16, 32, 64, 32, 16]):
+        """
+        Initialize the propagator network.
+        Input : (z(t), tau)
+        Output: z(t+tau)
+        """
+        super(Propagator_concat, self).__init__()
+
+        self._net = nn.Sequential(
+            nn.Linear(latent_dim + 2, feats[0]),  # 1 is for tau; more params will increase this
+            nn.GELU(),
+            nn.Linear(feats[0], feats[1]),
+            nn.GELU(),
+            nn.Linear(feats[1], feats[2]),
+            nn.GELU(),
+            nn.Linear(feats[2], feats[3]),
+            nn.GELU(),
+            nn.Linear(feats[3], feats[4]),
+            nn.GELU(),
+            nn.Linear(feats[4], latent_dim),
+        )
+
+    def forward(self, z, tau, alpha):
+        """
+        Forward pass of the propagator.
+        Concatenates latent vector z with tau and processes through the network.
+        """
+        zproj = z.squeeze(1)  # Adjust z dimensions if necessary
+        z_ = torch.cat((zproj, tau, alpha), dim=1)  # Concatenate z and tau along the last dimension
+        z_tau = self._net(z_)
+        return z_tau, z_
+        
+        
+        
+        
+class Model(nn.Module):
+    def __init__(self, encoder, decoder, propagator):
+        super(Model, self).__init__()
+        self.encoder = encoder
+        self.decoder = decoder # decoder for x(t)
+        self.propagator = propagator  # used to time march z(t) to z(t+tau)
+
+    def reparameterization(self, mean, var):
+        epsilon = torch.randn_like(var)
+        z = mean + var * epsilon
+        return z
+
+    def forward(self, x, tau, alpha):
+        mean, log_var = self.encoder(x)
+        z = self.reparameterization(mean, torch.exp(0.5 * log_var))
+
+        # Update small fcnn to get z(t+tau) from z(t)
+        z_tau, z_ = self.propagator(z, tau, alpha)
+
+        # Reconstruction
+        x_hat = self.decoder(z)  # Reconstruction of x(t)
+        x_hat_tau = self.decoder(z_tau)
+
+        return x_hat, x_hat_tau, mean, log_var, z_tau, z_
+        
+def loss_function(x, x_tau, x_hat, x_hat_tau, mean, log_var):
+    """
+    Compute the VAE loss components.
+    :param x: Original input
+    :param x_tau: Future input (ground truth)
+    :param x_hat: Reconstructed x(t)
+    :param x_hat_tau: Predicted x(t+tau)
+    :param mean: Mean of the latent distribution
+    :param log_var: Log variance of the latent distribution
+    :return: reconstruction_loss1, reconstruction_loss2, KLD
+    """
+    reconstruction_loss1 = nn.MSELoss()(x, x_hat)  # Reconstruction loss for x(t)
+    reconstruction_loss2 = nn.MSELoss()(x_tau, x_hat_tau)  # Prediction loss for x(t+tau)
+    
+    # Kullback-Leibler Divergence
+    KLD = torch.mean(-0.5 * torch.sum(1 + log_var - mean.pow(2) - log_var.exp(), dim=1))  # Updated dim
+    
+    return reconstruction_loss1, reconstruction_loss2, KLD

model_io_adv_dif.py

@@ -0,0 +1,22 @@
+import torch
+from dataclasses import dataclass, asdict
+from data_adv_dif import IntervalSplit
+from config_adv_dif import Config
+
+
+def save_model(path, model, tau_interval_split, alpha_interval_split, config):
+    torch.save({
+        'model_state_dict': model.state_dict(),
+        'alpha_interval_split': asdict(alpha_interval_split),
+        'tau_interval_split': asdict(tau_interval_split),
+        'config': asdict(config),
+    }, path)
+
+
+def load_model(path, model):
+    checkpoint = torch.load(path)
+    model.load_state_dict(checkpoint['model_state_dict'])
+    alpha_interval_split = IntervalSplit(**checkpoint['alpha_interval_split'])
+    tau_interval_split = IntervalSplit(**checkpoint['tau_interval_split'])
+    config = Config(**checkpoint['config'])
+    return model, alpha_interval_split, tau_interval_split, config

model_io_burgers.py

@@ -0,0 +1,22 @@
+import torch
+from dataclasses import dataclass, asdict
+from data_burgers import IntervalSplit
+from config_burgers import Config
+
+
+def save_model(path, model, tau_interval_split, re_interval_split, config):
+    torch.save({
+        'model_state_dict': model.state_dict(),
+        're_interval_split': asdict(re_interval_split),
+        'tau_interval_split': asdict(tau_interval_split),
+        'config': asdict(config),
+    }, path)
+
+
+def load_model(path, model):
+    checkpoint = torch.load(path)
+    model.load_state_dict(checkpoint['model_state_dict'])
+    re_interval_split = IntervalSplit(**checkpoint['re_interval_split'])
+    tau_interval_split = IntervalSplit(**checkpoint['tau_interval_split'])
+    config = Config(**checkpoint['config'])
+    return model, re_interval_split, tau_interval_split, config

model_v2.py

@@ -0,0 +1,380 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+# In[1]:
+
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+import math
+import torch
+
+
+# In[2]:
+
+
+def positionalencoding1d(d_model, length):
+    """
+    :param d_model: dimension of the model
+    :param length: length of positions
+    :return: length*d_model position matrix
+    """
+    if d_model % 2 != 0:
+        raise ValueError("Cannot use sin/cos positional encoding with "
+                         "odd dim (got dim={:d})".format(d_model))
+    pe = torch.zeros(length, d_model)
+    position = torch.arange(0, length).unsqueeze(1)
+    div_term = torch.exp((torch.arange(0, d_model, 2, dtype=torch.float) *
+                         -(math.log(10000.0) / d_model)))
+    pe[:, 0::2] = torch.sin(position.float() * div_term)
+    pe[:, 1::2] = torch.cos(position.float() * div_term)
+
+    return pe
+
+
+# In[3]:
+
+
+class Norm(nn.Module):
+  def __init__(self, num_channels, num_groups=4):
+    super(Norm, self).__init__()
+    self.norm = nn.GroupNorm(num_groups, num_channels)
+
+  def forward(self, x):
+    return self.norm(x.permute(0,2,1)).permute(0,2,1)
+
+
+# In[4]:
+
+class Norm_new(nn.Module):
+    def __init__(self, num_channels, num_groups=4):
+        super(Norm_new, self).__init__()
+        self.norm = nn.GroupNorm(num_groups, num_channels)
+
+    def forward(self, x):
+        if x.dim() == 2:
+            # Reshape to (batch_size, num_channels, 1)
+            x = x.unsqueeze(-1)
+            x = self.norm(x)
+            # Reshape back to (batch_size, num_channels)
+            x = x.squeeze(-1)
+        else:
+            x = self.norm(x.permute(0, 2, 1)).permute(0, 2, 1)
+        return x
+
+
+class Encoder(nn.Module):
+    def __init__(self, input_dim, latent_dim=2, feats=[512, 256, 128, 64, 32]):
+        super(Encoder, self).__init__()
+        self.latent_dim = latent_dim
+        self._net = nn.Sequential(
+            nn.Linear(input_dim, feats[0]),
+            nn.GELU(),
+            Norm_new(feats[0]),
+            nn.Linear(feats[0], feats[1]),
+            nn.GELU(),
+            Norm_new(feats[1]),
+            nn.Linear(feats[1], feats[2]),
+            nn.GELU(),
+            Norm_new(feats[2]),
+            nn.Linear(feats[2], feats[3]),
+            nn.GELU(),
+            Norm_new(feats[3]),
+            nn.Linear(feats[3], feats[4]),
+            nn.GELU(),
+            Norm_new(feats[4]),
+            nn.Linear(feats[4], 2 * latent_dim)
+        )
+
+    def forward(self, x):
+      Z = self._net(x)
+      mean, log_var = torch.split(Z, self.latent_dim, dim=-1)
+      return mean, log_var
+
+
+# In[5]:
+
+
+class Decoder(nn.Module):
+    def __init__(self, latent_dim, output_dim, feats=[32, 64, 128, 256, 512]):
+        super(Decoder, self).__init__()
+        self.output_dim = output_dim
+        self._net = nn.Sequential(
+            nn.Linear(latent_dim, feats[0]),
+            nn.GELU(),
+            Norm_new(feats[0]),
+            nn.Linear(feats[0], feats[1]),
+            nn.GELU(),
+            Norm_new(feats[1]),
+            nn.Linear(feats[1], feats[2]),
+            nn.GELU(),
+            Norm_new(feats[2]),
+            nn.Linear(feats[2], feats[3]),
+            nn.GELU(),
+            Norm_new(feats[3]),
+            nn.Linear(feats[3], feats[4]),
+            nn.GELU(),
+            Norm_new(feats[4]),
+            nn.Linear(feats[4], output_dim),
+            nn.Tanh()
+        )
+
+    def forward(self, x):
+      y = self._net(x)
+      return y
+
+
+# In[6]:
+
+
+class Propagator(nn.Module): #taken in (z(t), tau) and outputs z(t+tau)  [2, 5, 10, 2]
+  def __init__(self, latent_dim, feats=[16, 32], max_tau=10000, encoding_dim=64):
+    
+    """
+    Input : (z(t), tau)
+    Output: z(t+tau)
+    """
+    self.max_tau = max_tau
+    super(Propagator, self).__init__()
+    self.register_buffer('encodings',  positionalencoding1d(encoding_dim, max_tau))  # shape: max_tau, 64
+
+    self.projector = nn.Sequential(
+        nn.Linear(latent_dim, encoding_dim),
+        nn.ReLU(),
+        Norm(encoding_dim),
+        nn.Linear(encoding_dim, encoding_dim),
+    )
+
+    self._net = nn.Sequential(
+            nn.Linear(encoding_dim, feats[0]),
+            nn.ReLU(),
+            Norm(feats[0]),
+            nn.Linear(feats[0], feats[1]),
+            nn.ReLU(),
+            Norm(feats[1]),
+            nn.Linear(feats[1], latent_dim),
+            )
+
+
+  def forward(self, z, tau):
+    zproj = self.projector(z)
+    enc = self.encodings[tau.long()]
+    # z: 2
+    # enc: 64
+    # [z1, z2, enc1, enc2, ..., enc64]
+    z = zproj + enc
+
+    z_tau = self._net(z)
+    return z_tau
+
+
+# Doing this for the embedding for Re
+class Propagator_encoding(nn.Module): #taken in (z(t), tau) and outputs z(t+tau)  [2, 5, 10, 2]
+  def __init__(self, latent_dim, feats=[16, 32], max_tau=10000, encoding_dim=64, max_re = 5000):
+    
+    """
+    Input : (z(t), tau, re)
+    Output: z(t+tau)
+    """
+    self.max_tau = max_tau
+    self.max_re = max_re
+    super(Propagator_encoding, self).__init__()
+    self.register_buffer('tau_encodings',  positionalencoding1d(encoding_dim, max_tau))  # shape: max_tau, 64
+    self.register_buffer('re_encodings',  positionalencoding1d(encoding_dim, max_re))  # shape: max_re, 64
+    
+    self.projector = nn.Sequential(
+        nn.Linear(latent_dim, encoding_dim),
+        nn.ReLU(),
+        Norm(encoding_dim),
+        nn.Linear(encoding_dim, encoding_dim),
+    )
+
+    self._net = nn.Sequential(
+            nn.Linear(encoding_dim, feats[0]),
+            nn.ReLU(),
+            Norm(feats[0]),
+            nn.Linear(feats[0], feats[1]),
+            nn.ReLU(),
+            Norm(feats[1]),
+            nn.Linear(feats[1], latent_dim),
+            )
+
+
+  def forward(self, z, tau, re):
+    zproj = self.projector(z)
+    tau_enc = self.tau_encodings[tau.long()]
+    re_enc = self.re_encodings[re.long()]
+    # z: 2
+    # enc: 64
+    # [z1, z2, enc1, enc2, ..., enc64]
+    z = zproj + tau_enc + re_enc
+    #print("shape after enc addition: ", z.shape)
+    z_tau = self._net(z)
+    #print("shape z_tau: ", z_tau.shape)
+    return z_tau
+
+
+
+class Propagator_concat(nn.Module): #taken in (z(t), tau) and outputs z(t+tau)  [2, 5, 10, 2]
+  def __init__(self, latent_dim, feats = [16, 32]):
+    
+    """
+    Input : (z(t), tau, re)
+    Output: z(t+tau)
+    """
+    super(Propagator_concat, self).__init__()
+    
+    self._net = nn.Sequential(
+            nn.Linear(latent_dim + 2, feats[0]),
+            nn.ReLU(),
+            #Norm(feats[1]),
+            nn.Linear(feats[0], feats[1]),
+            nn.ReLU(),
+            #Norm(feats[2]),
+            nn.Linear(feats[1], latent_dim),
+            )
+    
+  def forward(self, z, tau, re):
+    zproj = z.squeeze(1)
+    z_ = torch.cat((zproj, tau, re), dim = 1)   
+    z_tau = self._net(z_)
+    z_tau = z_tau[:, None, :]
+    
+    return z_tau
+
+
+
+class Propagator_concat_one_step(nn.Module): #taken in (z(t), Re) and outputs z(t+tau)  [2, 5, 10, 2]
+  def __init__(self, latent_dim, feats = [16, 32]):
+    
+    """
+    Input : (z(t), re)
+    Output: z(t+1*dt)
+    """
+    super(Propagator_concat_one_step, self).__init__()
+    
+    self._net = nn.Sequential(
+            nn.Linear(latent_dim + 1, feats[0]),
+            nn.ReLU(),
+            #Norm(feats[1]),
+            nn.Linear(feats[0], feats[1]),
+            nn.Tanh(),
+            #Norm(feats[2]),
+            nn.Linear(feats[1], latent_dim),
+            )
+    
+  def forward(self, z, re):
+    #zproj = z.squeeze(1)
+    zproj = z
+    z_ = torch.cat((zproj, re), dim = 1)   
+    z_tau = self._net(z_)
+    #z_tau = z_tau[:, None, :]
+    
+    return z_tau
+
+
+
+class Model(nn.Module):
+    def __init__(self, encoder, decoder, propagator):
+        super(Model, self).__init__()
+        self.encoder = encoder
+        self.decoder = decoder # decoder for x(t)
+        self.propagator = propagator  # used to time march z(t) to z(t+tau)
+
+    def reparameterization(self, mean, var):
+        epsilon = torch.randn_like(var)
+        z = mean + var * epsilon
+        return z
+
+    def forward(self, x, tau, re):
+        mean, log_var = self.encoder(x)
+        z = self.reparameterization(mean, torch.exp(0.5 * log_var))
+
+        # Update small fcnn to get z(t+tau) from z(t)
+        z_tau = self.propagator(z, tau, re)
+
+        # Reconstruction
+        x_hat = self.decoder(z)  # Reconstruction of x(t)
+        x_hat_tau = self.decoder(z_tau)
+
+        return x_hat, x_hat_tau, mean, log_var, z_tau
+
+    
+class Model_One_Step(nn.Module): # Only takes in X and Re as the parameter and not the tau as tau = 1
+    def __init__(self, encoder, decoder, propagator):
+        super(Model_One_Step, self).__init__()
+        self.encoder = encoder
+        self.decoder = decoder # decoder for x(t)
+        self.propagator = propagator  # used to time march z(t) to z(t+tau)
+
+    def reparameterization(self, mean, var):
+        epsilon = torch.randn_like(var)
+        z = mean + var * epsilon
+        return z
+
+    def forward(self, x, re):
+        mean, log_var = self.encoder(x)
+        z = self.reparameterization(mean, torch.exp(0.5 * log_var))
+
+        # Update small fcnn to get z(t+1*dt) from z(t) -- We will use the Propagator_concat_one_step here!
+        z_tau = self.propagator(z, re)
+
+        # Reconstruction
+        x_hat = self.decoder(z)  # Reconstruction of x(t)
+        x_hat_tau = self.decoder(z_tau)
+
+        return x_hat, x_hat_tau, mean, log_var, z_tau
+    
+    
+
+class Model_reproduce(nn.Module):
+    def __init__(self, encoder, decoder):
+        super(Model_reproduce, self).__init__()
+        self.encoder = encoder
+        self.decoder = decoder # decoder for x(t)
+
+    def reparameterization(self, mean, var):
+        epsilon = torch.randn_like(var)
+        z = mean + var * epsilon
+        return z
+
+    def forward(self, x):
+        mean, log_var = self.encoder(x)
+        z = self.reparameterization(mean, torch.exp(0.5 * log_var))
+
+        # Reconstruction
+        x_hat = self.decoder(z)  # Reconstruction of x(t)
+
+        return x_hat, mean, log_var
+    
+
+# Define loss function
+def loss_function_reproduce(x, x_hat, mean, log_var):
+    reconstruction_loss1 = nn.MSELoss()(x, x_hat) 
+    
+    KLD = torch.mean(-0.5 * torch.sum(1 + log_var - mean.pow(2) - log_var.exp(), dim=2))
+    return reconstruction_loss1, KLD
+
+
+# Define loss function
+def loss_function(x, x_tau, x_hat, x_hat_tau, mean, log_var):
+    reconstruction_loss1 = nn.MSELoss()(x, x_hat) 
+    reconstruction_loss2 = nn.MSELoss()(x_tau, x_hat_tau)
+    
+    KLD = torch.mean(-0.5 * torch.sum(1 + log_var - mean.pow(2) - log_var.exp(), dim=2))
+    return reconstruction_loss1, reconstruction_loss2, KLD
+
+
+def count_parameters(model):
+    total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    return total_params
+
+
+
+
+
+
+
+

requirements.txt

@@ -0,0 +1,7 @@
+torch
+numpy
+gradio
+matplotlib
+dataclasses
+json5
+