{"metadata":{"kaggle":{"accelerator":"none","dataSources":[{"sourceId":56537,"databundleVersionId":8877088,"sourceType":"competition"},{"sourceId":188251279,"sourceType":"kernelVersion"},{"sourceId":190226151,"sourceType":"kernelVersion"},{"sourceId":82363,"sourceType":"modelInstanceVersion","modelInstanceId":68826,"modelId":93974}],"dockerImageVersionId":30732,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false},"kernelspec":{"name":"python3","display_name":"Python 3","language":"python"},"language_info":{"name":"python","version":"3.10.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Intro\n\n#### My goal here is to simulate whether it will rain in 20, 40, 60....X minutes from now, **without using timeseries-like models** (as was described in the intention of the competition)\n\n### Takeaways\n\n* Experienced stability issues with Tanh vs LeakyRelu or GELU activation functions.\n* **Once a prediction pushes your next timestep off of the data manifold, it can cause instability. If it pushes the prediction to the wrong position on the data manifold it will cause a poorly calibrated prediction.** I'm almost certain this can be helped with more data but I think there are some interesting tricks that can be used to help alleviate this that I did not get a chance to try.\n\n### Possible Improvements / Missed Opportunities\n\n- I wish that I had used this information\nhttps://leap-stc.github.io/ClimSim/demo_notebooks/water_conservation.html\n\n- It would have been interesting to predict some of the other scalar values for which I could not use the dx/dt Euler's approach (I actually probably could have done this for every value but that would require using time series information in training itself). **Is this against the spirit of model learning physics?**\n\n- Considering the TSNE graphs I imagine some sort of adversarial method would help with keeping the simulated points \"data-like\". This, of course, would not lend the simulator to perform well once out of distribution. Perhaps this can be alleviated by producing even more training data with extreme weather.\n\n- KL regularization: I tried it but I don't think it was well calibrated and gave up on the approach (see last few cells for more info).","metadata":{}},{"cell_type":"markdown","source":"## NOTE\n\nRetraining the models may yield very different results because they were trained on a different subset of the data than just the first N_ROWS of the test set.","metadata":{}},{"cell_type":"code","source":"import gc\nimport os\nimport random\nimport time\nimport torch\nimport datetime\nimport numpy as np\nimport pandas as pd\nimport polars as pl\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.optim.lr_scheduler import StepLR","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-07-29T05:08:07.83496Z","iopub.execute_input":"2024-07-29T05:08:07.83567Z","iopub.status.idle":"2024-07-29T05:08:07.847507Z","shell.execute_reply.started":"2024-07-29T05:08:07.83561Z","shell.execute_reply":"2024-07-29T05:08:07.845672Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"! pip install opentsne","metadata":{"execution":{"iopub.status.busy":"2024-07-29T05:08:07.85093Z","iopub.execute_input":"2024-07-29T05:08:07.851579Z","iopub.status.idle":"2024-07-29T05:08:27.19333Z","shell.execute_reply.started":"2024-07-29T05:08:07.851508Z","shell.execute_reply":"2024-07-29T05:08:27.190984Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Actually Predicting the Weather\n\nI will use the example of wind speed and air temperature. \n\n`state_u` and `state_v` are in m/s. `state_t` is in K.\n\nFor each of these we have actual target variable values of `ptend_t`, `ptend_u`, `ptend_v`.\n\nThese are in: `m/s^2` and `K/s`. \n\n$$v_t = v_0 + at$$\n\nWhere:\n- $v_t$ is the velocity at time $t$\n- $v_0$ is the initial velocity\n- $a$ is the constant acceleration\n- $t$ is the time elapsed\n\n### What procedure will we use?\n\nWe will perform a fairly simple procedure that is similar to **Euler's Method**, or other **Numerical Integration** methods that do not include multiple averaged timesteps. Basically we will predict the dx/dt at each timestep and add dX/dt * dt to X_t. We then repeat the procedure until we have arrived at the desired timestep. \n\n### Can we back out actual time information from only this dataset?\n\nApparently from discussions we can.","metadata":{}},{"cell_type":"code","source":"device = \"cuda\" if torch.cuda.is_available() else \"cpu\"","metadata":{"execution":{"iopub.status.busy":"2024-07-29T05:08:27.195838Z","iopub.execute_input":"2024-07-29T05:08:27.196441Z","iopub.status.idle":"2024-07-29T05:08:27.206178Z","shell.execute_reply.started":"2024-07-29T05:08:27.196396Z","shell.execute_reply":"2024-07-29T05:08:27.204082Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def set_seed(seed=42):\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = False\n    \n# Usage\nset_seed(42)  # or any other seed value","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-07-29T05:08:27.208288Z","iopub.execute_input":"2024-07-29T05:08:27.208774Z","iopub.status.idle":"2024-07-29T05:08:27.23237Z","shell.execute_reply.started":"2024-07-29T05:08:27.208737Z","shell.execute_reply":"2024-07-29T05:08:27.228717Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from pathlib import Path","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-07-29T05:08:27.237989Z","iopub.execute_input":"2024-07-29T05:08:27.238661Z","iopub.status.idle":"2024-07-29T05:08:27.256399Z","shell.execute_reply.started":"2024-07-29T05:08:27.23861Z","shell.execute_reply":"2024-07-29T05:08:27.253631Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**NOTE**: Even if we don't retrain the model we need this for normalizations. Of course the normalizations used here and the model being pre-trained would actually be different so it would be best to train on the same set and where the normalizations were derived.","metadata":{}},{"cell_type":"code","source":"LEAP_PATH = Path(\"/kaggle/input/leap-atmospheric-physics-ai-climsim/\")\nDATA_PATH = Path(\"/kaggle/input/save-files/\")\nN_ROWS = 500_000","metadata":{"execution":{"iopub.status.busy":"2024-07-29T05:08:27.25837Z","iopub.execute_input":"2024-07-29T05:08:27.259054Z","iopub.status.idle":"2024-07-29T05:08:27.274487Z","shell.execute_reply.started":"2024-07-29T05:08:27.258997Z","shell.execute_reply":"2024-07-29T05:08:27.272644Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Make sure the train data is periodic\n\nAs <blank> pointed out the test data repeats every 384. Just want to double check with the training data.","metadata":{}},{"cell_type":"code","source":"weights = pd.read_csv(LEAP_PATH / \"./sample_submission.csv\", nrows=1)\ndel weights['sample_id']\nif 'Unnamed: 0' in weights.columns:\n    del weights['Unnamed: 0']\nweights = weights.T\nweights = weights.to_dict()[0]","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2024-07-29T05:08:27.276718Z","iopub.execute_input":"2024-07-29T05:08:27.277282Z","iopub.status.idle":"2024-07-29T05:08:27.334958Z","shell.execute_reply.started":"2024-07-29T05:08:27.277234Z","shell.execute_reply":"2024-07-29T05:08:27.332727Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Unshuffled","metadata":{}},{"cell_type":"code","source":"import polars as pl\n\ntrain = (\n    pl.read_csv(LEAP_PATH / \"./train.csv\", n_rows=N_ROWS // 4)\n    .drop(\"sample_id\")\n    .with_row_index(\"sample_id\")\n    .with_columns(pl.col(\"sample_id\").mod(384).alias(\"grid_id\"))\n)","metadata":{"execution":{"iopub.status.busy":"2024-07-29T05:08:27.337387Z","iopub.execute_input":"2024-07-29T05:08:27.33821Z","iopub.status.idle":"2024-07-29T05:08:36.343125Z","shell.execute_reply.started":"2024-07-29T05:08:27.338156Z","shell.execute_reply":"2024-07-29T05:08:36.341493Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train = train.to_pandas()","metadata":{"execution":{"iopub.status.busy":"2024-07-29T05:08:36.345391Z","iopub.execute_input":"2024-07-29T05:08:36.34597Z","iopub.status.idle":"2024-07-29T05:08:38.33698Z","shell.execute_reply.started":"2024-07-29T05:08:36.34592Z","shell.execute_reply":"2024-07-29T05:08:38.3356Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import pandas as pd\nimport matplotlib.pyplot as plt\nimport matplotlib.colors as mcolors\n\n# Assuming you have a pandas dataframe named 'df' with the required columns\n# and the values are already in a numerical format\n\n# Create a new figure and axis\nfig, ax = plt.subplots(figsize=(10, 6))\n\nstate_t_df = train[['grid_id'] + [c for c in train.columns if 'pbuf_LHFLX' in c]]\n\nstate_t_df = state_t_df[state_t_df['grid_id'].isin(set(range(0,100)))]\n\n# Get a list of unique gridIds\ngrid_ids = state_t_df['grid_id'].unique()\n\n# Create a scatter plot for each gridId\n# for grid_id in grid_ids[0:100]:\n#     # Filter the dataframe for the current gridId\n#     grid_data = state_t_df[state_t_df['grid_id'] == grid_id]\n\n# Get the unique grid_id values\nunique_grid_ids = state_t_df['grid_id'].unique()\n\nnum_colors = len(unique_grid_ids)\ncmap = plt.cm.get_cmap('viridis', num_colors)\ncustom_cmap = mcolors.ListedColormap(cmap.colors)\n\n# Map grid_id values to colors\ngrid_id_to_color = {grid_id: i for i, grid_id in enumerate(unique_grid_ids)}\ncolors = [grid_id_to_color[grid_id] for grid_id in state_t_df['grid_id']]\n    \n# Plot the data for the current gridId\nax.scatter(\n    state_t_df.index, \n    state_t_df[state_t_df.columns[1]], \n    c=colors,\n    cmap=custom_cmap,\n    #label=f'grid_id={grid_id}', \n    s=10, alpha=0.5\n)\n\n# Set axis labels and title\nax.set_xlabel('Index')\nax.set_ylabel('Value')\nax.set_title('2D Value Plot')\n\n\n# Show the plot\nplt.show()","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-07-29T05:08:38.339741Z","iopub.execute_input":"2024-07-29T05:08:38.340259Z","iopub.status.idle":"2024-07-29T05:08:39.6484Z","shell.execute_reply.started":"2024-07-29T05:08:38.340215Z","shell.execute_reply":"2024-07-29T05:08:39.646732Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Shuffled","metadata":{}},{"cell_type":"code","source":"train2 = (\n    pl.from_pandas(train)\n    .sample(fraction=1, shuffle=True)\n    .with_columns(pl.col(\"sample_id\").mod(384).alias(\"grid_id\"))\n)\ntrain2 = train2.to_pandas()","metadata":{"execution":{"iopub.status.busy":"2024-07-29T05:08:39.650301Z","iopub.execute_input":"2024-07-29T05:08:39.650819Z","iopub.status.idle":"2024-07-29T05:08:42.350095Z","shell.execute_reply.started":"2024-07-29T05:08:39.650774Z","shell.execute_reply":"2024-07-29T05:08:42.348881Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import pandas as pd\nimport matplotlib.pyplot as plt\nimport matplotlib.colors as mcolors\n\n# Assuming you have a pandas dataframe named 'df' with the required columns\n# and the values are already in a numerical format\n\n# Create a new figure and axis\nfig, ax = plt.subplots(figsize=(10, 6))\n\nstate_t_df = train2[['grid_id'] + [c for c in train2.columns if 'state_t' in c]]\n\nstate_t_df = state_t_df[state_t_df['grid_id'].isin(set(range(0,100)))]\n\n# Get a list of unique gridIds\ngrid_ids = state_t_df['grid_id'].unique()\n\n# Create a scatter plot for each gridId\n# for grid_id in grid_ids[0:100]:\n#     # Filter the dataframe for the current gridId\n#     grid_data = state_t_df[state_t_df['grid_id'] == grid_id]\n\n# Get the unique grid_id values\nunique_grid_ids = state_t_df['grid_id'].unique()\n\nnum_colors = len(unique_grid_ids)\ncmap = plt.cm.get_cmap('viridis', num_colors)\ncustom_cmap = mcolors.ListedColormap(cmap.colors)\n\n# Map grid_id values to colors\ngrid_id_to_color = {grid_id: i for i, grid_id in enumerate(unique_grid_ids)}\ncolors = [grid_id_to_color[grid_id] for grid_id in state_t_df['grid_id']]\n    \n# Plot the data for the current gridId\nax.scatter(\n    state_t_df.index, \n    state_t_df[state_t_df.columns[1]], \n    c=colors,\n    cmap=custom_cmap,\n    #label=f'grid_id={grid_id}', \n    s=10, alpha=0.5\n)\n\n# Set axis labels and title\nax.set_xlabel('Index')\nax.set_ylabel('Value')\nax.set_title('2D Value Plot')\n\n# Add a colorbar\n# cbar = plt.colorbar()\n# cbar.set_label('Value')\n\n# Add a legend\n# ax.legend()\n\n# Show the plot\nplt.show()","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-07-29T05:08:42.352188Z","iopub.execute_input":"2024-07-29T05:08:42.352599Z","iopub.status.idle":"2024-07-29T05:08:43.454235Z","shell.execute_reply.started":"2024-07-29T05:08:42.352565Z","shell.execute_reply":"2024-07-29T05:08:43.453078Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Summary\n\n#### Confirmed that the train data is periodic. ","metadata":{}},{"cell_type":"markdown","source":"## Prediction Code","metadata":{}},{"cell_type":"code","source":"import torch\n\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-07-29T05:08:43.456052Z","iopub.execute_input":"2024-07-29T05:08:43.456543Z","iopub.status.idle":"2024-07-29T05:08:43.463096Z","shell.execute_reply.started":"2024-07-29T05:08:43.456505Z","shell.execute_reply":"2024-07-29T05:08:43.461622Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"FEAT_COLS = train.columns[1:557]\nTARGET_COLS = [c for c in train.columns[557:] if c != \"grid_id\"]","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-07-29T05:08:43.469247Z","iopub.execute_input":"2024-07-29T05:08:43.469862Z","iopub.status.idle":"2024-07-29T05:08:43.477349Z","shell.execute_reply.started":"2024-07-29T05:08:43.469827Z","shell.execute_reply":"2024-07-29T05:08:43.475637Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Utils","metadata":{}},{"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nfrom sklearn.preprocessing import MinMaxScaler\n\nMIN_STD = 1e-11  # Define a minimum standard deviation to prevent division by zero\n\ndef normalize_dataframe(df, log_columns=None, scaling_method='standard'):\n    \"\"\"\n    Normalize a pandas DataFrame, apply log to specified columns,\n    and store mean, std, and log transformation information for later use.\n    \n    Parameters:\n    df (pd.DataFrame): DataFrame to be normalized.\n    log_columns (list): List of column names to apply log transformation.\n    scaling_method (str): Method to scale the data, either 'standard' or 'minmax'.\n    \n    Returns:\n    pd.DataFrame: Normalized DataFrame.\n    dict: Dictionary containing the mean, std, and log info for each column.\n    \"\"\"\n    mean_std_log = {}\n    normalized_df = df.copy()  # Make a copy to avoid modifying the original DataFrame\n    \n    if log_columns is None:\n        log_columns = []\n    \n    for column in df.columns:\n        # Apply log transformation if the column is in log_columns\n        if column in log_columns:\n            normalized_df[column] = np.log1p(df[column])\n            mean_std_log[column] = {'log': True}\n        else:\n            mean_std_log[column] = {'log': False}\n        \n        mean = normalized_df[column].mean()\n        std = normalized_df[column].std()\n        std = np.maximum(std, MIN_STD)\n        mean_std_log[column].update({'mean': mean, 'std': std})\n        \n        if scaling_method == 'minmax':\n            # Apply min-max scaling\n            scaler = MinMaxScaler()\n            normalized_df[column] = scaler.fit_transform(normalized_df[[column]])\n            # Store min and max for later use\n            min_val = normalized_df[column].min()\n            max_val = normalized_df[column].max()\n            mean_std_log[column].update({'min': min_val, 'max': max_val})\n        else:\n            # Normalize the column using standard scaling\n            if std != 0:\n                normalized_df[column] = (normalized_df[column] - mean) / std\n            else:\n                normalized_df[column] = normalized_df[column] - mean\n    \n    return normalized_df, mean_std_log\n\ndef denormalize_dataframe(normalized_df, mean_std_log):\n    \"\"\"\n    Denormalize a pandas DataFrame using stored mean, std, log, min, and max values.\n    \n    Parameters:\n    normalized_df (pd.DataFrame): Normalized DataFrame to be denormalized.\n    mean_std_log (dict): Dictionary containing the mean, std, log, min, and max info for each column.\n    \n    Returns:\n    pd.DataFrame: Denormalized DataFrame.\n    \"\"\"\n    denormalized_df = normalized_df.copy()\n    \n    for column in normalized_df.columns:\n        if 'min' in mean_std_log[column] and 'max' in mean_std_log[column]:\n            # Reverse min-max scaling\n            min_val = mean_std_log[column]['min']\n            max_val = mean_std_log[column]['max']\n            denormalized_df[column] = normalized_df[column] * (max_val - min_val) + min_val\n        else:\n            mean = mean_std_log[column]['mean']\n            std = mean_std_log[column]['std']\n            \n            # Denormalize the column, handle case where std was zero\n            if std != 0:\n                denormalized_df[column] = (normalized_df[column] * std) + mean\n            else:\n                denormalized_df[column] = normalized_df[column] + mean\n        \n        # Reverse log transformation if applied\n        if mean_std_log[column].get('log', False):\n            denormalized_df[column] = np.expm1(denormalized_df[column])  # expm1 to reverse log1p\n    \n    return denormalized_df\n\n# Example usage:\n# Create a sample DataFrame\ndata = {'feature1': [1, 2, 3, 4, 5], 'q000_feature': [1, 2, 3, 4, 5], 'log_feature': [1, 10, 100, 1000, 10000]}\ndf = pd.DataFrame(data)\n\n# Specify columns for log transform\nlog_columns = ['log_feature']\n\n# Normalize the DataFrame\nnormalized_df, mean_std_log = normalize_dataframe(df, log_columns=log_columns)\nprint(\"Normalized DataFrame:\")\nprint(normalized_df)\n\n# Denormalize the DataFrame\ndenormalized_df = denormalize_dataframe(normalized_df, mean_std_log)\nprint(\"\\nDenormalized DataFrame:\")\nprint(denormalized_df)","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2024-07-29T05:08:43.479114Z","iopub.execute_input":"2024-07-29T05:08:43.479501Z","iopub.status.idle":"2024-07-29T05:08:43.518071Z","shell.execute_reply.started":"2024-07-29T05:08:43.479469Z","shell.execute_reply":"2024-07-29T05:08:43.516607Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def prepare_dataframe(train_set, weights, scaler_type=\"standard\"):\n    for target in weights:\n        train_set[target] = train_set[target] * weights[target]\n    df, mean_std_ = normalize_dataframe(train_set, scaler_type)\n    normalizations = mean_std_\n    print('final shape')\n    print(df.shape)\n    return df, mean_std_","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-07-29T05:08:43.519878Z","iopub.execute_input":"2024-07-29T05:08:43.520295Z","iopub.status.idle":"2024-07-29T05:08:43.528001Z","shell.execute_reply.started":"2024-07-29T05:08:43.520261Z","shell.execute_reply":"2024-07-29T05:08:43.526778Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def dataframe_to_numpy(rows, sequence_columns, scalar_columns, sequence_targets, scalar_target_columns):\n    # make sure to call this after standard scaling bc it is in float32\n    inputs_to_stack = []\n    for col_set in sequence_columns:\n        np_rows = rows[col_set].to_numpy(np.float32)\n        inputs_to_stack.append(\n            torch.tensor(np_rows, dtype=torch.float32).unsqueeze(-1))\n    targets_to_stack = []\n    for col_set in sequence_targets:\n        np_rows = rows[col_set].to_numpy(np.float32)\n        targets_to_stack.append(\n            torch.tensor(np_rows, dtype=torch.float32).unsqueeze(-1))\n    sequence_inputs_stacked = torch.cat(inputs_to_stack, axis=-1)\n    sequence_targets_stacked = torch.cat(targets_to_stack, axis=-1)\n\n    if scalar_columns is not None:\n        np_rows = rows[scalar_columns].astype(np.float32).to_numpy(dtype=np.float32)\n        scalars = torch.tensor(np_rows, dtype=torch.float32)\n    else:\n        scalars = None\n\n    if scalar_target_columns is not None:\n        np_rows = rows[scalar_target_columns].astype(np.float32).to_numpy(dtype=np.float32)\n        scalar_targets = torch.tensor(\n            np_rows, dtype=torch.float32)\n    else:\n        scalar_targets = None\n    return sequence_inputs_stacked, scalars, sequence_targets_stacked, scalar_targets","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-07-29T05:08:43.52972Z","iopub.execute_input":"2024-07-29T05:08:43.530145Z","iopub.status.idle":"2024-07-29T05:08:43.547815Z","shell.execute_reply.started":"2024-07-29T05:08:43.530104Z","shell.execute_reply":"2024-07-29T05:08:43.546156Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Models","metadata":{}},{"cell_type":"code","source":"class AttentionBlock(nn.Module):\n    def __init__(self, features):\n        super(AttentionBlock, self).__init__()\n        self.attention = nn.Sequential(\n            nn.Conv1d(features, features // 8, kernel_size=1),\n            nn.BatchNorm1d(features // 8),\n            nn.ReLU(),\n            nn.Conv1d(features // 8, 1, kernel_size=1),\n            nn.Sigmoid()\n        )\n\n    def forward(self, x):\n        attention_weights = self.attention(x)\n        return x * attention_weights\n\nclass UNet1D(nn.Module):\n    def __init__(self, in_channels, out_channels, input_length=60, init_features=60, scalar_dim=0, output_dim=1):\n        super(UNet1D, self).__init__()\n    \n        self.seq_dim = in_channels\n        self.seq_len = input_length\n        self.scalar_dim = scalar_dim\n        features = init_features\n        self.encoder1 = self._block(in_channels + scalar_dim, features, name=\"enc1\")\n        self.pool1 = nn.MaxPool1d(kernel_size=2, stride=2)\n        self.encoder2 = self._block(features, features * 2, name=\"enc2\")\n        self.pool2 = nn.MaxPool1d(kernel_size=2, stride=2)\n        \n        self.bottleneck = self._block(features * 2, features * 4, name=\"bottleneck\")\n        self.bottleneck2 = self._block(features * 4, features * 4, name=\"bottleneck\")\n        \n        self.upconv2 = nn.ConvTranspose1d(features * 4, features * 2, kernel_size=2, stride=2)\n        self.decoder2 = self._block(features * 4, features * 2, name=\"dec2\")  # Updated to match concatenation\n        self.upconv1 = nn.ConvTranspose1d(features * 2, features, kernel_size=2, stride=2)\n        self.decoder1 = self._block(features * 2, features, name=\"dec1\")  # Updated to match concatenation\n        \n        self.conv = nn.Conv1d(in_channels=features, out_channels=out_channels, kernel_size=1)\n        self.linear = nn.Linear(input_length * out_channels, input_length * out_channels)\n        self.final = nn.Linear(input_length * out_channels, output_dim)\n        self.tanh = nn.Tanh()\n        self.gru = nn.GELU()\n        self.dropout = nn.Dropout(0.05)\n    \n    def forward(self, x):\n        batch_size_ = x.shape[0]\n        x_seq = x[:, 0:-self.scalar_dim]\n        scalars = x[:, -self.scalar_dim:]\n        x_seq = x_seq.reshape([batch_size_, self.seq_dim, self.seq_len])\n        scalar_expanded = scalars.unsqueeze(2).expand(-1, -1, x_seq.size(2))\n        # Concatenate along feature dimension\n        x = torch.cat([x_seq, scalar_expanded], dim=1)\n        enc1 = self.encoder1(x)\n        enc2 = self.encoder2(self.pool1(enc1))\n        bottleneck = self.bottleneck(self.pool2(enc2))\n        bottleneck = self.bottleneck2(bottleneck)\n        dec2 = self.upconv2(bottleneck)\n        dec2 = torch.cat((dec2, enc2), dim=1)\n        dec2 = self.decoder2(dec2)\n        dec1 = self.upconv1(dec2)\n        dec1 = torch.cat((dec1, enc1), dim=1)\n        dec1 = self.decoder1(dec1)\n        final = self.conv(dec1)\n        final = self.tanh(\n            self.linear(final.reshape(batch_size_, -1))\n        )\n        final = self.dropout(final)\n        x = self.final(final)\n        return x\n    \n    @staticmethod\n    def _block(in_channels, features, name):\n        return nn.Sequential(\n            nn.Conv1d(in_channels, features, kernel_size=3, padding=1, stride=1, bias=False),\n            nn.BatchNorm1d(features),\n            nn.LeakyReLU(),\n            nn.Conv1d(features, features, kernel_size=3, padding=1, stride=1, bias=False),\n            nn.BatchNorm1d(features),\n            AttentionBlock(features),\n            nn.Tanh(),\n        )\n    \nin_channels = 2\nout_channels = 2\ninit_features = 8\ninput_length = 60\n# Instantiate the model with masking\nunet_model_test = UNet1D(in_channels, out_channels, input_length, init_features, scalar_dim=10, output_dim=12).to(device)\n# Example input (batch_size, in_channels, sequence_length)\nexample_input = torch.randn(32, 130).to(device)\n\n# Forward pass\noutput = unet_model_test(example_input)\nprint(output.shape)  # Should be (32, out_channels, 60)","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2024-07-29T05:08:43.550348Z","iopub.execute_input":"2024-07-29T05:08:43.551222Z","iopub.status.idle":"2024-07-29T05:08:43.613252Z","shell.execute_reply.started":"2024-07-29T05:08:43.551181Z","shell.execute_reply":"2024-07-29T05:08:43.611895Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from torch.nn.utils.parametrizations import weight_norm\n\nclass Permute(nn.Module):\n    def __init__(self, *dims):\n        super(Permute, self).__init__()\n        self.dims = dims\n\n    def forward(self, x):\n        return x.permute(*self.dims)\n\nclass LSTMConvnet(nn.Module):\n    def __init__(self, seq_dim, seq_len, scalar_dim, hidden_dim, output_dim, activation=\"tanh\"):\n        super(LSTMConvnet, self).__init__()\n        \n        input_dim = seq_dim * seq_len + scalar_dim * 2\n        input_dim_conv = 30 * 16 + scalar_dim * 2\n\n        self.scalar_dim = scalar_dim\n        self.seq_dim = seq_dim\n        self.seq_len = seq_len\n\n        bidirectional = True\n\n        lstm_hidden_size = 64\n        self.lstm_hidden_size = lstm_hidden_size\n        \n        # Define the layers with weight normalization\n        self.fc1 = weight_norm(nn.Linear(lstm_hidden_size * 2 * 60, output_dim))\n        self.fc_skip = weight_norm(nn.Linear(scalar_dim + seq_dim * seq_len, output_dim))\n        self.fc2 = weight_norm(nn.Linear(output_dim, output_dim))\n\n        self.lstm_expand = nn.Sequential(\n            nn.Conv1d(\n                seq_dim + scalar_dim, lstm_hidden_size, kernel_size=5, padding='same'\n            ),\n            nn.BatchNorm1d(lstm_hidden_size),\n            nn.Tanh(),\n            Permute(0, 2, 1),\n            nn.LSTM(\n                lstm_hidden_size, lstm_hidden_size, 1, batch_first=True, bidirectional=bidirectional\n            )\n        )\n        self.lstm2_shape = lstm_hidden_size * 2 * (1 + int(bidirectional)) * 60\n        \n        if activation == \"tanh\":\n            self.activation = nn.Tanh()\n        elif activation == \"gelu\":\n            self.activation = nn.GELU()\n        else:\n            self.activation = nn.LeakyReLU()\n\n    def lstm_to_output(self, x):\n        batch_size, seq_len, _ = x.shape\n        x_out = x.reshape(batch_size, -1)\n        return x_out\n    \n    def forward(self, x):\n        batch_size_ = x.shape[0]\n        seq_flattened = x[:, 0:-self.scalar_dim]\n        scalars = x[:, -self.scalar_dim:]\n        x_seq = seq_flattened.reshape([batch_size_, self.seq_dim, self.seq_len])\n        scalar_expanded = scalars.unsqueeze(2).expand(-1, -1, x_seq.size(2))\n        # Concatenate along feature dimension\n        combined_input = torch.cat([x_seq, scalar_expanded], dim=1)\n        lstm_expand, _ = self.lstm_expand(combined_input)\n        # lstm_contract, _ = self.lstm_contract(lstm_expand)\n        x_out = self.lstm_to_output(lstm_expand)\n        x_out = self.activation(self.fc1(x_out))\n        x = x_out + self.activation(self.fc_skip(x))\n        x = self.fc2(x)\n        return x\n    \nseq_dim = 4\nseq_len = 60  # Change according to your input dimension\nscalar_dim = 8\nhidden_dim = 512  # Change according to your hidden layer size\noutput_dim = seq_dim * seq_len + scalar_dim   # Change according to your output dimension\n\nmodel = LSTMConvnet(seq_dim, seq_len, scalar_dim, hidden_dim, output_dim).to(device)\n\n# Example input\nx = torch.randn(256, seq_dim * seq_len + scalar_dim).to(device)  # Batch size of 256\nprint(f'x: {x.shape}')\noutput = model(x)\nprint(output.shape)","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2024-07-29T05:08:43.61648Z","iopub.execute_input":"2024-07-29T05:08:43.617088Z","iopub.status.idle":"2024-07-29T05:08:43.806107Z","shell.execute_reply.started":"2024-07-29T05:08:43.617039Z","shell.execute_reply":"2024-07-29T05:08:43.804289Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.nn.utils.parametrizations import weight_norm\n\nclass ConvNet(nn.Module):\n    def __init__(self, seq_dim, seq_len, scalar_dim, hidden_dim, output_dim):\n        super(ConvNet, self).__init__()\n        \n        input_dim = seq_dim * seq_len + scalar_dim * 2\n        input_dim_conv = 30 * 16 + scalar_dim * 2\n\n        self.scalar_dim = scalar_dim\n        self.seq_len = seq_len\n        self.seq_dim = seq_dim\n\n        bidirectional = True\n\n        lstm_hidden_size = 16\n        \n        # Define the layers with weight normalization\n        self.scalar_to_hidden = weight_norm(nn.Linear(scalar_dim, lstm_hidden_size))\n        self.fc1 = weight_norm(nn.Linear(15 * 60 + 128, hidden_dim))\n        self.fc2 = weight_norm(nn.Linear(hidden_dim + 128, seq_dim * seq_len + scalar_dim))\n        self.fc3 = weight_norm(nn.Linear(seq_dim * seq_len + scalar_dim, seq_dim * seq_len + scalar_dim))\n        self.fc_out = nn.Linear(seq_dim * seq_len + scalar_dim, output_dim)\n        self.conv = nn.Conv1d(seq_dim + scalar_dim, 30, kernel_size=3, stride=2, padding=1)\n        self.conv2 = nn.Conv1d(30, 60, kernel_size=3, stride=2, padding=1)\n        # GeLU activation function\n        self.bottleneck = nn.Sequential(\n            nn.Linear(seq_dim * seq_len + scalar_dim, 256), \n            nn.SiLU(),\n            nn.Linear(256, 128),\n            nn.Tanh(),\n        )\n        self.gelu = nn.GELU()\n        self.Tanh = nn.Tanh()\n        self.sigmoid = nn.Sigmoid()\n        self.binary_layer = nn.Linear(seq_dim * seq_len + scalar_dim, 128)\n        self.dropout = nn.Dropout(0.05)\n    \n    def forward(self, x):\n        batch_size_ = x.shape[0]\n        x_seq = x[:, 0:-self.scalar_dim]\n        scalars = x[:, -self.scalar_dim:]        \n        x_seq = x_seq.reshape([batch_size_, self.seq_dim, self.seq_len])\n        scalar_expanded = scalars.unsqueeze(2).expand(-1, -1, x_seq.size(2))\n        # Concatenate along feature dimension\n        combined_input = torch.cat([x_seq, scalar_expanded], dim=1)\n        batch_size_, c_, seq_len_ = x_seq.size()\n        seq_conv = self.conv(combined_input)\n        conv_bottleneck = self.conv2(seq_conv)\n        x_bottleneck = self.bottleneck(x)\n        \n        conv_bottleneck = conv_bottleneck.reshape(batch_size_, -1)\n        x_all = torch.cat([conv_bottleneck, x_bottleneck], -1)\n        x1 = self.gelu(self.fc1(x_all))\n        x_binary = self.sigmoid(self.binary_layer(x))\n        x2 = self.gelu(\n            self.fc2(\n                torch.cat([x1, x_binary], dim=1)\n            )\n        )\n        x3 = self.Tanh(self.fc3(x2 + x))\n        x3 = self.dropout(x3)\n        x = self.fc_out(x3)\n        return x\n    \nseq_dim = 4\nseq_len = 60  # Change according to your input dimension\nscalar_dim = 8\nhidden_dim = 512  # Change according to your hidden layer size\noutput_dim = seq_dim * seq_len + scalar_dim   # Change according to your output dimension\n\nmodel = ConvNet(seq_dim, seq_len, scalar_dim, hidden_dim, output_dim).to(device)\n\n# Example input\nx = torch.randn(64, seq_dim * seq_len + scalar_dim).to(device)  # Batch size of 64\nscalar = torch.randn(64, scalar_dim)\noutput = model(x)\nprint(output.shape)\n","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-07-29T05:08:43.808423Z","iopub.execute_input":"2024-07-29T05:08:43.808993Z","iopub.status.idle":"2024-07-29T05:08:43.862286Z","shell.execute_reply.started":"2024-07-29T05:08:43.808957Z","shell.execute_reply":"2024-07-29T05:08:43.860942Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Select Subset of Columns\n\nSo **we only have certain columns where we actually have dx/dt.**\nWe can update all of those target columns but the other columns won't be updating. \n\nSo we will be feeding the model a combination of updated temperatures, wind speeds, and humidities...while keeping other columns constant.\n\n**I am also including scalar columns that don't change**.\n","metadata":{}},{"cell_type":"code","source":"sequence_cols_dxt_dt = [\n    'state_t', \n    'state_u', \n    'state_v', \n    'state_q0001', \n    'state_q0002', \n    'state_q0003',\n]\nsequence_target_cols_dxt_dt = [\n    'ptend_t', \n    'ptend_q0001',\n    'ptend_q0002', \n    'ptend_q0003',\n    'ptend_u', \n    'ptend_v', \n]\n\nsequence_col_all_dx_dt = []\nsequence_target_col_all_dx_dt = []\nfor col in sequence_cols_dxt_dt:\n    sequence_col_all_dx_dt = sequence_col_all_dx_dt + [c for c in FEAT_COLS if col in c]\nfor col in sequence_target_cols_dxt_dt:\n    sequence_target_col_all_dx_dt = sequence_target_col_all_dx_dt + [c for c in TARGET_COLS if col in c]","metadata":{"_kg_hide-input":false,"execution":{"iopub.status.busy":"2024-07-29T05:08:43.864316Z","iopub.execute_input":"2024-07-29T05:08:43.864914Z","iopub.status.idle":"2024-07-29T05:08:43.87422Z","shell.execute_reply.started":"2024-07-29T05:08:43.864872Z","shell.execute_reply":"2024-07-29T05:08:43.872646Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import pandas as pd\nimport torch\nfrom torch.utils.data import Dataset\nimport os\nimport random\nfrom sklearn.preprocessing import StandardScaler\n\nsequence_cols = sequence_cols_dxt_dt\nsequence_target_cols = sequence_target_cols_dxt_dt\nsequence_col_sets = []\nsequence_target_col_sets = []\nfor col in sequence_cols:\n    sequence_col_sets.append([c for c in FEAT_COLS if col in c])\nfor col in sequence_target_cols:\n    sequence_target_col_sets.append([c for c in TARGET_COLS if col in c])\n\nsequence_cols_all = [\n    x\n    for cols in sequence_col_sets\n    for x in cols\n]\nsequence_target_cols_all = [\n    x\n    for cols in sequence_target_col_sets\n    for x in cols\n]\n\n# this is lazy but from a quick investigation these look to be constant\nscalar_columns = ['cam_in_LANDFRAC', 'cam_in_OCNFRAC']\nscalar_targets = [\n     'cam_out_PRECC',\n     'cam_out_PRECSC',\n     'cam_out_FLWDS',\n]","metadata":{"_kg_hide-input":false,"execution":{"iopub.status.busy":"2024-07-29T05:08:43.87625Z","iopub.execute_input":"2024-07-29T05:08:43.876727Z","iopub.status.idle":"2024-07-29T05:08:43.889074Z","shell.execute_reply.started":"2024-07-29T05:08:43.876689Z","shell.execute_reply":"2024-07-29T05:08:43.887856Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Dataloader","metadata":{}},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader\nimport pandas as pd\n\nclass CustomDatasetMultiCol(Dataset):\n    def __init__(self, data, x_features, y_features):\n        \"\"\"\n        Initializes the LeapDataset.\n\n        Parameters\n        ----------\n        file_path : str\n            Path to the parquet file containing the dataset.\n        x_features : list of str\n            List of feature names for the input data.\n        y_features : list of str\n            List of feature names for the target data.\n        y_weights : torch.Tensor\n            Weights to be applied to the target data.\n        \"\"\"\n        \n        super().__init__()\n        \n        self.x = data[x_features]\n\n        self.x = self.x.to_numpy()\n        self.x = torch.from_numpy(self.x)\n        \n        self.y = data[y_features]\n        self.y = self.y.to_numpy()\n        self.y = torch.from_numpy(self.y)\n        \n        \n    def __getitem__(self, idx):\n        x = self.x[idx]\n        y = self.y[idx]\n        \n        x = x.to(torch.float32)\n        y = y.to(torch.float32)\n        \n        return x, y\n    \n    def __len__(self):\n        return len(self.y)","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-07-29T05:08:43.890504Z","iopub.execute_input":"2024-07-29T05:08:43.89102Z","iopub.status.idle":"2024-07-29T05:08:43.909138Z","shell.execute_reply.started":"2024-07-29T05:08:43.890974Z","shell.execute_reply":"2024-07-29T05:08:43.907971Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Load Training Data","metadata":{}},{"cell_type":"code","source":"train = (\n    pl.read_csv(LEAP_PATH / \"./train.csv\", n_rows=N_ROWS)\n    .drop(\"sample_id\")\n    .with_row_index(\"sample_id\")\n    .with_columns(pl.col(\"sample_id\").mod(384).alias(\"grid_id\"))\n)","metadata":{"execution":{"iopub.status.busy":"2024-07-29T05:08:43.910906Z","iopub.execute_input":"2024-07-29T05:08:43.911329Z","iopub.status.idle":"2024-07-29T05:09:09.803739Z","shell.execute_reply.started":"2024-07-29T05:08:43.911266Z","shell.execute_reply":"2024-07-29T05:09:09.802192Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for target in weights:\n    train = train.with_columns(pl.col(target).mul(weights[target]))\ntrain = train.to_pandas()\ntrain = train.drop(columns=[\"sample_id\", \"grid_id\"])\ntrain, mean_std = normalize_dataframe(train, scaling_method='standard', log_columns=[-3])","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-07-29T05:09:09.805765Z","iopub.execute_input":"2024-07-29T05:09:09.80626Z","iopub.status.idle":"2024-07-29T05:09:34.858067Z","shell.execute_reply.started":"2024-07-29T05:09:09.806214Z","shell.execute_reply":"2024-07-29T05:09:34.856904Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_sequence_indices_dx_dt = train.columns.get_indexer(sequence_col_all_dx_dt)\ntarget_sequence_indices_dx_dt = train.columns.get_indexer(sequence_target_col_all_dx_dt)\n\ntrain_only = train.iloc[0:N_ROWS - 10_000]\ndf_holdout = train.iloc[N_ROWS - 10_000:]","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-07-29T05:09:34.859413Z","iopub.execute_input":"2024-07-29T05:09:34.859799Z","iopub.status.idle":"2024-07-29T05:09:34.879879Z","shell.execute_reply.started":"2024-07-29T05:09:34.859767Z","shell.execute_reply":"2024-07-29T05:09:34.878047Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# NOTE: the order here matters for sequence vs scalars\ndataset = CustomDatasetMultiCol(train_only, sequence_cols_all + scalar_columns, sequence_target_cols_all + scalar_targets)","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-07-29T05:09:34.882356Z","iopub.execute_input":"2024-07-29T05:09:34.882862Z","iopub.status.idle":"2024-07-29T05:09:38.892952Z","shell.execute_reply.started":"2024-07-29T05:09:34.882825Z","shell.execute_reply":"2024-07-29T05:09:38.89157Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"holdout_dataset = CustomDatasetMultiCol(df_holdout, sequence_cols_all + scalar_columns, sequence_target_cols_all + scalar_targets)","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-07-29T05:09:38.895105Z","iopub.execute_input":"2024-07-29T05:09:38.895695Z","iopub.status.idle":"2024-07-29T05:09:38.979563Z","shell.execute_reply.started":"2024-07-29T05:09:38.895648Z","shell.execute_reply":"2024-07-29T05:09:38.97805Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from torch.utils.data import DataLoader, Sampler\nbatch_size = 2000\ndataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)","metadata":{"execution":{"iopub.status.busy":"2024-07-29T05:09:38.981719Z","iopub.execute_input":"2024-07-29T05:09:38.982282Z","iopub.status.idle":"2024-07-29T05:09:38.999995Z","shell.execute_reply.started":"2024-07-29T05:09:38.982234Z","shell.execute_reply":"2024-07-29T05:09:38.998489Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"holdout_dataloader = DataLoader(holdout_dataset, batch_size=4000, shuffle=True)","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-07-29T05:09:39.001942Z","iopub.execute_input":"2024-07-29T05:09:39.002328Z","iopub.status.idle":"2024-07-29T05:09:39.018335Z","shell.execute_reply.started":"2024-07-29T05:09:39.002297Z","shell.execute_reply":"2024-07-29T05:09:39.016948Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"device = torch.device(\"cpu\")\nif torch.cuda.is_available():\n    device = \"cuda\"","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-07-29T05:09:39.020159Z","iopub.execute_input":"2024-07-29T05:09:39.020594Z","iopub.status.idle":"2024-07-29T05:09:39.031596Z","shell.execute_reply.started":"2024-07-29T05:09:39.020557Z","shell.execute_reply":"2024-07-29T05:09:39.030141Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Train","metadata":{}},{"cell_type":"code","source":"def train_model(model_type=\"LSTM\", num_epochs=40):\n    \n    seq_len = 60\n    hidden_dim = 800\n    seq_dim = len(sequence_cols_dxt_dt)\n    scalar_dim = len(scalar_columns)\n    init_features = 48\n    n_sequence_outputs = len(sequence_cols_all)\n    n_scalar_outputs = len(scalar_targets)\n    scalar_dim = len(scalar_columns)\n    output_dim = n_scalar_outputs + n_sequence_outputs\n\n    print(f'Model with: {len(sequence_col_sets)} sequence cols, {len(scalar_columns)} scalar cols')\n\n    print(f'Predicting: {len(sequence_target_col_sets)} sequence cols, {len(scalar_targets)} scalar cols')\n\n    print(scalar_dim)\n\n    if model_type.lower() == \"conv\":\n        model = ConvNet(seq_dim, seq_len, scalar_dim, hidden_dim, output_dim)\n    elif model_type.lower() == \"lstm\":\n        model = LSTMConvnet(seq_dim, seq_len, scalar_dim, hidden_dim, output_dim)\n    elif model_type == \"unet\":\n        model = UNet1D(seq_dim, seq_dim + scalar_dim, seq_len, init_features, scalar_dim=scalar_dim, output_dim=output_dim)\n        \n    model = model.to(device)\n\n    print(f'Device: {device}')\n    \n    total_params = sum(p.numel() for p in model.parameters())\n    print(f\"Number of parameters: {total_params}\")\n\n    criterion = nn.MSELoss()\n    optimizer = optim.Adam(model.parameters(), lr=0.0015)\n    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, gamma=0.96, step_size=1)\n\n    print(f'n epochs {num_epochs}')\n    print(next(model.parameters()).device)\n    \n    start = time.time()\n\n    for epoch in range(num_epochs):\n        model.train()\n        moving_avg_loss = []\n        st = time.time()\n        epoch_losses = []\n        i = 0\n        for (\n            input_data, target_data\n        ) in dataloader:\n            optimizer.zero_grad()\n            batch_dim, blah = input_data.size()\n            input_data = input_data.to(device)\n            outputs = model(input_data)  # Forward pass\n            loss = criterion(\n                outputs.to(device), \n                target_data.to(device)#.permute([0, 2, 1])\n            )  # Compute loss\n            loss.backward()  # Backward pass\n            optimizer.step()  # Update weights\n            moving_avg_loss.append(loss.item())\n            i += 1\n        if (epoch + 1) % 2 == 0:\n            moving_avg_loss_val = np.mean(moving_avg_loss)\n            t = time.time() - start\n            print(f'Epoch {epoch+1}/{num_epochs}, Loss ({len(moving_avg_loss)}): {moving_avg_loss_val} {t} s')\n        scheduler.step()\n\n    return model","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-07-29T05:09:39.033134Z","iopub.execute_input":"2024-07-29T05:09:39.033596Z","iopub.status.idle":"2024-07-29T05:09:39.055328Z","shell.execute_reply.started":"2024-07-29T05:09:39.033527Z","shell.execute_reply":"2024-07-29T05:09:39.053123Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"seq_len = 60\nhidden_dim = 800\nseq_dim = len(sequence_cols_dxt_dt)\nscalar_dim = len(scalar_columns)\ninit_features = 48\nn_sequence_outputs = len(sequence_cols_all)\nn_scalar_outputs = len(scalar_targets)\nscalar_dim = len(scalar_columns)\noutput_dim = n_scalar_outputs + n_sequence_outputs\n\nN_EPOCHS = 40\n\nimport time\n\nretrain = True\n\n# this will use some base model I trained a while ago\n# MODEL_PATH = Path('/kaggle/input/rain-forecast-lg/pytorch/default/2')\n\n# this will use the most recently trained model\nMODEL_PATH = Path('/kaggle/input/rain-simulation')\n\nstart = time.time()\n\nmodel_type = \"lstm\"\nif False:\n    print('training')\n    lstm = train_model(model_type, N_EPOCHS) # just takes so long to train\n    torch.save(lstm.state_dict(), f\"model_weights_sim_{model_type}.pt\")\n    print('saved')\nelse:\n    print('loading lstm')\n    lstm = LSTMConvnet(seq_dim, seq_len, scalar_dim, hidden_dim, output_dim)\n    lstm_weights = torch.load(MODEL_PATH / f\"model_weights_sim_{model_type}.pt\", map_location=device)\n    lstm.load_state_dict(lstm_weights)\n    torch.save(lstm.state_dict(), f\"model_weights_sim_{model_type}.pt\")\nprint(f'{time.time() - start} s')\n\n\nmodel_type = \"conv\"\nif False:\n    convnet = train_model(model_type, N_EPOCHS)\n    torch.save(convnet.state_dict(), f\"model_weights_sim_{model_type}.pt\")\n    print('saved')\nelse:\n    print('loading convnet')\n    convnet_weights = torch.load(MODEL_PATH / f\"model_weights_sim_{model_type}.pt\", map_location=device)\n    convnet = ConvNet(seq_dim, seq_len, scalar_dim, hidden_dim, output_dim)\n    convnet.load_state_dict(convnet_weights)\n    torch.save(convnet.state_dict(), f\"model_weights_sim_{model_type}.pt\")\nprint(f'{time.time() - start} s')\n\nmodel_type = \"unet\"\nif retrain:\n    unet = train_model(model_type, N_EPOCHS)\n    torch.save(unet.state_dict(), f\"model_weights_sim_{model_type}.pt\")\n    print('saved')\nelse:\n    print('loading unet')\n    unet_weights = torch.load(MODEL_PATH / f\"model_weights_sim_{model_type}.pt\", map_location=device)\n    unet = UNet1D(seq_dim, seq_dim + scalar_dim, seq_len, init_features, scalar_dim=scalar_dim, output_dim=output_dim)\n    unet.load_state_dict(unet_weights)\n    torch.save(unet.state_dict(), f\"model_weights_sim_{model_type}.pt\")\nprint(f'{time.time() - start} s')","metadata":{"execution":{"iopub.status.busy":"2024-07-29T05:09:39.057978Z","iopub.execute_input":"2024-07-29T05:09:39.058494Z","iopub.status.idle":"2024-07-29T05:09:39.453583Z","shell.execute_reply.started":"2024-07-29T05:09:39.058441Z","shell.execute_reply":"2024-07-29T05:09:39.451906Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Assess Model\n\nAs you will be able to see below, the performance of the model isn't bad (although not amazing). But you can see decent correlation and R2.","metadata":{}},{"cell_type":"code","source":"lstm.to(device)\nunet.to(device)\nlstm.eval()\nunet.eval()\nconvnet.to(device)\nconvnet.eval()","metadata":{"_kg_hide-output":true,"execution":{"iopub.status.busy":"2024-07-29T05:09:39.455283Z","iopub.execute_input":"2024-07-29T05:09:39.455817Z","iopub.status.idle":"2024-07-29T05:09:39.473801Z","shell.execute_reply.started":"2024-07-29T05:09:39.455764Z","shell.execute_reply":"2024-07-29T05:09:39.472229Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def make_predictions(model, dataloader):\n    all_preds = []\n    all_targets = []\n    i = 0\n    for (\n        input_data, target_data\n    ) in dataloader:\n        with torch.no_grad():\n            batch_dim, blah = input_data.size()\n            input_data = input_data.to(device)\n            outputs = model(input_data)  # Forward pass\n            print(outputs.shape)\n            print(target_data.shape)\n            all_preds.append(outputs.detach().cpu().numpy())\n            all_targets.append(target_data.detach().cpu().numpy())\n            i += 1\n\n    all_preds = np.vstack(all_preds)\n    all_targets = np.vstack(all_targets)\n    print(all_preds.shape)\n    print(all_targets.shape)\n    return all_preds, all_targets","metadata":{"scrolled":true,"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-07-29T05:09:39.475259Z","iopub.execute_input":"2024-07-29T05:09:39.475643Z","iopub.status.idle":"2024-07-29T05:09:39.4853Z","shell.execute_reply.started":"2024-07-29T05:09:39.475611Z","shell.execute_reply":"2024-07-29T05:09:39.483843Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"unet_preds, unet_targets = make_predictions(unet, holdout_dataloader)\nlstm_preds, lstm_targets = make_predictions(lstm, holdout_dataloader)\nconv_preds, conv_targets = make_predictions(convnet, holdout_dataloader)","metadata":{"execution":{"iopub.status.busy":"2024-07-29T05:09:39.487096Z","iopub.execute_input":"2024-07-29T05:09:39.487533Z","iopub.status.idle":"2024-07-29T05:09:50.51116Z","shell.execute_reply.started":"2024-07-29T05:09:39.487497Z","shell.execute_reply":"2024-07-29T05:09:50.509809Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import pandas as pd\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nfrom sklearn.metrics import r2_score\nfrom scipy.stats import pearsonr\n\n\ndef generate_scatter_plots_with_metrics(all_preds, all_targets, offset=256):\n    # Assuming all_other_outputs_np and all_other_targets_np are already defined\n\n    # Create a figure with 6 subplots\n    fig, axes = plt.subplots(3, 4, figsize=(9, 8))\n\n    # Flatten the axes array for easy iteration\n    axes = axes.flatten()\n\n    scaler = \"standard\"\n\n    # Loop through indices 2 to 6\n    for i in range(offset, offset + 12):  # Adjust range to match 2 to 6 inclusively\n        df_scatter = pd.DataFrame(data={'pred': all_preds[:, i], 'tgt': all_targets[:, i]})\n        r2_sklearn = r2_score(all_preds[:, i], all_targets[:, i], multioutput='raw_values')\n        pearson = pearsonr(all_preds[:, i].squeeze(), all_targets[:, i].squeeze())\n        print(pearson.statistic)\n        sns.scatterplot(data=df_scatter, x='pred', y='tgt', ax=axes[i - offset])\n        # fig, axes = plt.subplots(1, 1, figsize=(10, 8))\n        # Adding distribution plots\n        axes[i - offset].set_title(\n            f'{i}, col: {df_holdout[sequence_target_col_all_dx_dt + scalar_targets].columns[i]}, \\n R2: {r2_sklearn} \\n Pearson: {np.round(pearson.statistic, 4)}', fontsize=9\n        )\n        if scaler == 'standard':\n            axes[i - offset].set_xlim(-4.5, 4.5)\n            axes[i - offset].set_ylim(-4.5, 4.5)\n        else:\n            axes[i - offset].set_xlim(-0.1, 1.1)\n            axes[i - offset].set_ylim(-0.1, 1.1)\n\n    # Adjust layout\n    plt.tight_layout()\n    plt.show()\n","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-07-29T05:09:50.513522Z","iopub.execute_input":"2024-07-29T05:09:50.51413Z","iopub.status.idle":"2024-07-29T05:09:50.918992Z","shell.execute_reply.started":"2024-07-29T05:09:50.514071Z","shell.execute_reply":"2024-07-29T05:09:50.917497Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Sanity Check","metadata":{}},{"cell_type":"code","source":"generate_scatter_plots_with_metrics(lstm_preds, lstm_targets, 351)","metadata":{"execution":{"iopub.status.busy":"2024-07-29T05:09:50.92076Z","iopub.execute_input":"2024-07-29T05:09:50.921154Z","iopub.status.idle":"2024-07-29T05:09:54.678207Z","shell.execute_reply.started":"2024-07-29T05:09:50.921122Z","shell.execute_reply":"2024-07-29T05:09:54.676774Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"generate_scatter_plots_with_metrics(unet_preds, unet_targets, 351)","metadata":{"execution":{"iopub.status.busy":"2024-07-29T05:09:54.679815Z","iopub.execute_input":"2024-07-29T05:09:54.680238Z","iopub.status.idle":"2024-07-29T05:09:58.650545Z","shell.execute_reply.started":"2024-07-29T05:09:54.680203Z","shell.execute_reply":"2024-07-29T05:09:58.646631Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"generate_scatter_plots_with_metrics(lstm_preds, lstm_targets, 351)","metadata":{"execution":{"iopub.status.busy":"2024-07-29T05:09:58.655914Z","iopub.execute_input":"2024-07-29T05:09:58.657186Z","iopub.status.idle":"2024-07-29T05:10:08.670397Z","shell.execute_reply.started":"2024-07-29T05:09:58.657023Z","shell.execute_reply":"2024-07-29T05:10:08.667173Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"all_r2s_unet = []\nfor i in range(unet_preds.shape[1]):\n    r2_sklearn = r2_score(unet_preds[:, i], unet_targets[:, i], multioutput='raw_values')\n    if r2_sklearn[0] == 1:\n        r2_sklearn = [0]\n    all_r2s_unet.append(r2_sklearn[0])\n\nall_r2s = []\nfor i in range(lstm_preds.shape[1]):\n    r2_sklearn = r2_score(lstm_preds[:, i], lstm_targets[:, i], multioutput='raw_values')\n    if r2_sklearn[0] == 1:\n        r2_sklearn = [0]\n    all_r2s.append(r2_sklearn[0])\n    \nall_r2s_conv = []\nfor i in range(conv_preds.shape[1]):\n    r2_sklearn = r2_score(conv_preds[:, i], conv_targets[:, i], multioutput='raw_values')\n    if r2_sklearn[0] == 1:\n        r2_sklearn = [0]\n    all_r2s_conv.append(r2_sklearn[0])","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-07-29T05:10:08.695008Z","iopub.execute_input":"2024-07-29T05:10:08.695477Z","iopub.status.idle":"2024-07-29T05:10:10.439003Z","shell.execute_reply.started":"2024-07-29T05:10:08.695444Z","shell.execute_reply":"2024-07-29T05:10:10.43776Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"fig, axes = plt.subplots(2, 2, figsize=(12, 8))\n\n# First subplot\naxes[0][0].bar(x=range(0, 363), height=all_r2s)\naxes[0][0].set_ylim(-10, 1)\naxes[0][0].set_title('All R2s LSTM')\naxes[0][0].set_xlabel('Index')\naxes[0][0].set_ylabel('R2 Value')\n\n# Second subplot\naxes[0][1].bar(x=range(0, 363), height=all_r2s_unet)\naxes[0][1].set_ylim(-10, 1)\naxes[0][1].set_title('All R2s UNet')\naxes[0][1].set_xlabel('Index')\naxes[0][1].set_ylabel('R2 Value')\n\n# Third subplot\naxes[1][0].bar(x=range(0, 363), height=all_r2s_conv)\naxes[1][0].set_ylim(-10, 1)\naxes[1][0].set_title('All R2s Conv')\naxes[1][0].set_xlabel('Index')\naxes[1][0].set_ylabel('R2 Value')\n\n# Adjust layout\nplt.tight_layout()\nplt.show()","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-07-29T05:10:10.442168Z","iopub.execute_input":"2024-07-29T05:10:10.442601Z","iopub.status.idle":"2024-07-29T05:10:17.320842Z","shell.execute_reply.started":"2024-07-29T05:10:10.442535Z","shell.execute_reply":"2024-07-29T05:10:17.319285Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Run the Simulation","metadata":{}},{"cell_type":"code","source":"# just get more training data\nsimulation_set = (\n    pl.read_csv(LEAP_PATH / \"train.csv\", n_rows=160_000)\n    .drop(\"sample_id\")\n    .with_row_index(\"sample_id\")\n    .with_columns(pl.col(\"sample_id\").mod(384).alias(\"grid_id\"))\n)","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-07-29T05:10:17.322592Z","iopub.execute_input":"2024-07-29T05:10:17.323038Z","iopub.status.idle":"2024-07-29T05:10:26.652053Z","shell.execute_reply.started":"2024-07-29T05:10:17.323001Z","shell.execute_reply":"2024-07-29T05:10:26.650646Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for target in weights:\n    simulation_set = simulation_set.with_columns(pl.col(target).mul(weights[target]))\nsimulation_set = simulation_set.to_pandas()\nsimulation_set_no_ids = simulation_set.drop(columns=[\"sample_id\", \"grid_id\"])\nprepared_dataframe, mean_std = normalize_dataframe(simulation_set_no_ids, scaling_method='standard', log_columns=[-3])","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-07-29T05:10:26.653851Z","iopub.execute_input":"2024-07-29T05:10:26.654263Z","iopub.status.idle":"2024-07-29T05:10:34.531059Z","shell.execute_reply.started":"2024-07-29T05:10:26.654228Z","shell.execute_reply":"2024-07-29T05:10:34.529606Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"simulation_set = simulation_set[[\"sample_id\", \"grid_id\"] + list(simulation_set_no_ids.columns)]","metadata":{"execution":{"iopub.status.busy":"2024-07-29T05:10:34.533353Z","iopub.execute_input":"2024-07-29T05:10:34.533792Z","iopub.status.idle":"2024-07-29T05:10:35.010322Z","shell.execute_reply.started":"2024-07-29T05:10:34.533755Z","shell.execute_reply":"2024-07-29T05:10:35.008917Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Setup Targets\n\n#### So I want to take 10 different grid ID locations.","metadata":{}},{"cell_type":"code","source":"seq_data, scalar_data, seq_target_data, scalar_target_data = dataframe_to_numpy(\n    prepared_dataframe, \n    sequence_columns=sequence_col_sets, \n    sequence_targets=sequence_target_col_sets,\n    scalar_columns=scalar_columns,\n    scalar_target_columns=scalar_targets\n)","metadata":{"execution":{"iopub.status.busy":"2024-07-29T05:10:35.012486Z","iopub.execute_input":"2024-07-29T05:10:35.012939Z","iopub.status.idle":"2024-07-29T05:10:36.988749Z","shell.execute_reply.started":"2024-07-29T05:10:35.012904Z","shell.execute_reply":"2024-07-29T05:10:36.987345Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"timesteps_ahead = 12\ntarget_rows = np.array(list([range(150, 160)])) + 384 * timesteps_ahead\nfollow_up_rows = []\ntimesteps_fwd = 80\nn_skip = timesteps_fwd\ntarget_rows = target_rows.squeeze()\nfor target_row in target_rows:\n    furs = []\n    for i in range(0, timesteps_fwd):\n        furs = furs + [target_row + 384 * i]\n    follow_up_rows.append(furs)\n\nseq_data_row_batch, \\\n    scalar_data_row_batch, \\\n    seq_target_data_row_batch, \\\n    scalar_target_data_row_batch = (\n        seq_data[target_rows], scalar_data[target_rows], seq_target_data[target_rows], scalar_target_data[target_rows]\n    )\n","metadata":{"execution":{"iopub.status.busy":"2024-07-29T05:10:36.990336Z","iopub.execute_input":"2024-07-29T05:10:36.990762Z","iopub.status.idle":"2024-07-29T05:10:37.008402Z","shell.execute_reply.started":"2024-07-29T05:10:36.990724Z","shell.execute_reply":"2024-07-29T05:10:37.00698Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## For faster Normalization/Standardization in the simulation Loop (need numpy arrays)","metadata":{}},{"cell_type":"markdown","source":"### Numpy only Normalization\n\nThis is faster.","metadata":{"_kg_hide-input":true}},{"cell_type":"code","source":"import numpy as np\n\ndef standardize_data(values, log_columns=None):\n    \"\"\"\n    Standardize the data using z-score scaling and return the standardized data along with mean and standard deviation.\n    Apply log transform to specified columns before standardization.\n    \n    Parameters:\n    values (np.array): Array of values to be standardized.\n    log_columns (list): List of column indices to apply log transformation.\n    \n    Returns:\n    np.array: Standardized data.\n    np.array: Array of mean values for each feature.\n    np.array: Array of standard deviation values for each feature.\n    dict: Information about log transforms.\n    \"\"\"\n    if log_columns is None:\n        log_columns = []\n    \n    transformed_values = values.copy()\n    \n    for col in log_columns:\n        transformed_values[:, col] = np.log1p(np.max(values[:, col], 0))\n    \n    mean_values = transformed_values.mean(axis=0)\n    std_values = transformed_values.std(axis=0)\n    std_values[std_values == 0] = 1  # Avoid division by zero\n    \n    standardized_values = (transformed_values - mean_values) / std_values\n    \n    return standardized_values, mean_values, std_values\n\ndef standardize_data_with_given_params(values, mean_values, std_values, log_cols=None):\n    \"\"\"\n    Standardize the data using given mean and standard deviation values.\n    Apply log transform to specified columns before standardization.\n    \n    Parameters:\n    values (np.array): Array of values to be standardized.\n    mean_values (np.array): Array of mean values for each feature.\n    std_values (np.array): Array of standard deviation values for each feature.\n    log_info (dict): Information about which columns to log transform.\n    \n    Returns:\n    np.array: Standardized data.\n    \"\"\"\n\n    if log_cols is None:\n        log_cols = []\n    transformed_values = values.copy()\n    \n    for col in log_cols:\n        transformed_values[:, col] = np.log1p(np.max(values[:, col], 0))\n\n    standardized_values = (transformed_values - mean_values) / std_values\n    return standardized_values\n\ndef destandardize_data(standardized_values, mean_values, std_values, log_cols=None):\n    \"\"\"\n    De-standardize the data using z-score scaling and reverse log transform if applied.\n    \n    Parameters:\n    standardized_values (np.array): Array of standardized values.\n    mean_values (np.array): Array of mean values for each feature.\n    std_values (np.array): Array of standard deviation values for each feature.\n    log_info (dict): Information about which columns were log transformed.\n    \n    Returns:\n    np.array: De-standardized data.\n    \"\"\"\n    if log_cols is None:\n        log_cols = []\n\n    destandardized_values = standardized_values * std_values + mean_values\n    \n    for col in log_cols:\n        destandardized_values[:, col] = np.expm1(destandardized_values[:, col])\n    \n    return destandardized_values\n\n# Example usage\ndata = np.array([\n    [1., 10., 100.],\n    [2., 20., 200.],\n    [3., 30., 300.]\n])\n\n# Specify columns for log transform (0-based index)\nlog_columns = [2]  # Apply log transform to the third column\n\n# Standardize the data\nstandardized_data, mean_values, std_values = standardize_data(data, log_columns)\nprint(\"\\nStandardized Data:\")\nprint(standardized_data)\nprint(\"\\nMean Values:\")\nprint(mean_values)\nprint(\"\\nStandard Deviation Values:\")\nprint(std_values)\n\n# De-standardize the data\ndestandardized_data = destandardize_data(standardized_data, mean_values, std_values, log_columns)\nprint(\"\\nDe-standardized Data:\")\nprint(destandardized_data)\n\n# Standardize new data with given parameters\nnew_data = np.array([\n    [4, 40, 400],\n    [5, 50, 500]\n])\nstandardized_new_data = standardize_data_with_given_params(new_data, mean_values, std_values, log_columns)\nprint(\"\\nStandardized New Data:\")\nprint(standardized_new_data)","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2024-07-29T05:10:37.010789Z","iopub.execute_input":"2024-07-29T05:10:37.011227Z","iopub.status.idle":"2024-07-29T05:10:37.037733Z","shell.execute_reply.started":"2024-07-29T05:10:37.011191Z","shell.execute_reply":"2024-07-29T05:10:37.036359Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"means_inputs = []\nstds_inputs = []\nfor col in prepared_dataframe[sequence_cols_all + scalar_columns].columns:\n    means_inputs.append(mean_std[col]['mean'])\n    stds_inputs.append(mean_std[col]['std'])\n    \nmeans_targets = []\nstds_targets = []\nfor col in prepared_dataframe[sequence_target_cols_all + scalar_targets].columns:\n    means_targets.append(mean_std[col]['mean'])\n    stds_targets.append(mean_std[col]['std'])\n\nmeans_inp = np.array(means_inputs)\nstds_inp = np.array(stds_inputs)\n\nmeans_tgt = np.array(means_targets)\nstds_tgt = np.array(stds_targets)","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-07-29T05:10:37.039437Z","iopub.execute_input":"2024-07-29T05:10:37.039899Z","iopub.status.idle":"2024-07-29T05:10:38.068495Z","shell.execute_reply.started":"2024-07-29T05:10:37.039847Z","shell.execute_reply":"2024-07-29T05:10:38.067059Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"initial_values = prepared_dataframe.loc[target_rows, sequence_cols_all + scalar_columns]\ntarget_df = prepared_dataframe.loc[\n    [row for row_set in follow_up_rows for row in row_set], sequence_target_col_all_dx_dt + scalar_targets\n]\ninitial_targets = prepared_dataframe.loc[target_rows, sequence_target_cols_all + scalar_targets]\ncurrent_values = initial_values.to_numpy()\nprint(current_values.shape)","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2024-07-29T05:10:38.070682Z","iopub.execute_input":"2024-07-29T05:10:38.07149Z","iopub.status.idle":"2024-07-29T05:10:39.593971Z","shell.execute_reply.started":"2024-07-29T05:10:38.071435Z","shell.execute_reply":"2024-07-29T05:10:39.592566Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Initial Prediction For Calibration\n\nSo here I just start with the initial target rows. Given the scatterplots shown earlier, the predictions should at least be reasonable. ","metadata":{}},{"cell_type":"code","source":"seq_data_rows_ = torch.tensor(\n    current_values[:, 0:len(sequence_cols_all)], dtype=torch.float32\n)\nscalar_data_rows_ = torch.tensor(current_values[:, len(sequence_cols_all):], dtype=torch.float32)\nseq_data_rows_ = seq_data_rows_.to(device)\nscalar_data_rows_ = scalar_data_rows_.to(device)\noutputs = lstm(torch.cat([seq_data_rows_, scalar_data_rows_], dim=1))","metadata":{"execution":{"iopub.status.busy":"2024-07-29T05:10:39.595443Z","iopub.execute_input":"2024-07-29T05:10:39.595957Z","iopub.status.idle":"2024-07-29T05:10:39.619623Z","shell.execute_reply.started":"2024-07-29T05:10:39.595912Z","shell.execute_reply":"2024-07-29T05:10:39.618414Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pd.DataFrame(data=outputs.detach().cpu()[:, -7:].numpy())","metadata":{"_kg_hide-output":true,"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-07-29T05:10:39.621537Z","iopub.execute_input":"2024-07-29T05:10:39.622102Z","iopub.status.idle":"2024-07-29T05:10:39.646016Z","shell.execute_reply.started":"2024-07-29T05:10:39.622051Z","shell.execute_reply":"2024-07-29T05:10:39.644611Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"errors = (outputs.detach().cpu()[:, -7:].numpy() - target_df.iloc[::n_skip, -7:].to_numpy())\nprint(errors.sum(axis=1))","metadata":{"execution":{"iopub.status.busy":"2024-07-29T05:10:39.647664Z","iopub.execute_input":"2024-07-29T05:10:39.648103Z","iopub.status.idle":"2024-07-29T05:10:39.658001Z","shell.execute_reply.started":"2024-07-29T05:10:39.648069Z","shell.execute_reply":"2024-07-29T05:10:39.656632Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Assessing Model Predictions","metadata":{}},{"cell_type":"code","source":"sns.set_theme(rc={'figure.figsize':(12.7,8.27)})\nerrors = outputs.detach().cpu()[:, 0:65].numpy() - target_df.iloc[::n_skip, 0:65].to_numpy()\nsns.barplot(errors)\nplt.ylim(-20, 20)\nplt.figure(figsize=(4,4))","metadata":{"execution":{"iopub.status.busy":"2024-07-29T05:10:39.659802Z","iopub.execute_input":"2024-07-29T05:10:39.66086Z","iopub.status.idle":"2024-07-29T05:10:43.199236Z","shell.execute_reply.started":"2024-07-29T05:10:39.66081Z","shell.execute_reply":"2024-07-29T05:10:43.197488Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"sns.set_theme(rc={'figure.figsize':(12.7,8.27)})\nerrors = outputs.detach().cpu()[:, -65:].numpy() - target_df.iloc[::n_skip, -65:].to_numpy()\nsns.barplot(errors)\nplt.ylim(-20, 20)","metadata":{"execution":{"iopub.status.busy":"2024-07-29T05:10:43.200932Z","iopub.execute_input":"2024-07-29T05:10:43.20133Z","iopub.status.idle":"2024-07-29T05:10:46.880008Z","shell.execute_reply.started":"2024-07-29T05:10:43.201296Z","shell.execute_reply":"2024-07-29T05:10:46.878542Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Simulation Loop","metadata":{}},{"cell_type":"code","source":"def simulation_loop(initial_targets, initial_values, model, number_of_minutes=80, dt=300, r2s_=None):\n    \"\"\"\n    initial_targets: actual targets\n    current_values: \n    \"\"\"\n    current_values = initial_values.to_numpy()\n\n    # confusing this here is that outputs from 0:360 are dx/dt where 360+ are in the same\n    # units as the current values list, so really im only adding this current_values to outputs list\n    # in order to get the last 3 entries\n    outputs_list = [initial_targets]\n    real_outputs_list = [initial_targets]\n    current_values_list = [current_values]\n    real_value_list = [current_values]\n    \n    if not r2s_:\n        r2s_ = all_r2s\n\n    # we are going to assume that the changes are constant for 5 minutes at a time\n    number_of_seconds = number_of_minutes * 60\n    print(number_of_seconds)\n    n_iters = int(number_of_seconds / dt)\n\n    print(f'n_iters {n_iters}')\n\n    torch.cuda.empty_cache()\n\n    outputs_normed = []\n    rain_values = []\n    real_outputs_normed_list = []\n\n    CLIP_VAL = 6\n\n    if not r2s_:\n        r2s_ = np.ones_like(means_targets)\n\n    dont_override = np.array(r2s_) > 0\n    override = np.array(r2s_) < 0\n    override = override * 0.01\n    dont_override = dont_override + override\n    \n    print(f'dont override shape: {dont_override.shape}')\n\n    break_early = 0\n    timestep_real = 1\n\n    for i in range(0, n_iters):\n        with torch.no_grad():\n            seq_data_rows_ = torch.tensor(\n                current_values[:, 0:len(sequence_cols_all)], dtype=torch.float32\n            )\n            scalar_data_rows_ = torch.tensor(current_values[:, len(sequence_cols_all):], dtype=torch.float32)\n            seq_data_rows_ = seq_data_rows_.to(device)\n            scalar_data_rows_ = scalar_data_rows_.to(device)\n            data_rows = torch.cat([seq_data_rows_, scalar_data_rows_], axis=1)\n            outputs = model(data_rows)\n            outputs_as_np = outputs.detach().cpu().numpy().squeeze() * dont_override\n            # do not allow outputs more than N standard deviations from the mean\n            outputs_as_np = outputs_as_np.clip(-CLIP_VAL, CLIP_VAL)\n            outputs_as_np_denorm = destandardize_data(outputs_as_np, means_tgt, stds_tgt, [-3])\n            # create zeros to offset the updatable ones\n            dX_dt = np.zeros_like(current_values)\n            # get only the updatable columns (ones with levels)\n            n_indices_to_keep = len(target_sequence_indices_dx_dt)\n            \n            # overwrite the zeros there\n            dX_dt[:, list(range(n_indices_to_keep))] = outputs_as_np_denorm[:, 0:n_indices_to_keep]\n            current_values_original_units = destandardize_data(current_values, means_inp, stds_inp, [-3])\n            current_values_original_units = current_values_original_units + dX_dt * dt\n            current_values = standardize_data_with_given_params(current_values_original_units, means_inp, stds_inp, [-3])\n            \n            # every 20 minutes we add to the outputs list\n            # for actual comparison I want the current_values\n            # for debugging I want tht dx/dt or outputs_list\n            if break_early and i > break_early:\n                print(break_early)\n                break\n            if (i + 1) % (20 * 60 / dt) == 0 or i < break_early:\n                target_row_np = np.array(target_rows) + (384 * timestep_real)\n                if i < 4:\n                    print(f'iter: {i + 1}')\n                    print(f'n clipped {np.sum(outputs_as_np == CLIP_VAL, axis=1) + np.sum(outputs_as_np == -CLIP_VAL, axis=1)}')\n                if (i + 1) % (20 * 60 / dt) == 0:\n                    if i < 4:\n                        print('target_row_np')\n                        print(target_row_np)\n                    outputs_normed.append(outputs_as_np)\n                    outputs_list.append(outputs_as_np_denorm)\n                    real_outputs = simulation_set.loc[target_row_np, sequence_target_cols_all + scalar_targets]\n                    normed_real = prepared_dataframe.loc[target_row_np, sequence_target_cols_all + scalar_targets]\n                    real_outputs_normed_list.append(normed_real)\n                    real_outputs_list.append(real_outputs)\n                    current_values_list.append(current_values)\n                    real_value_list.append(simulation_set.loc[target_row_np, sequence_cols_all + scalar_columns])\n                    timestep_real += 1\n    return outputs_normed, outputs_list, real_outputs_normed_list, real_outputs_list, real_value_list, current_values_list","metadata":{"execution":{"iopub.status.busy":"2024-07-29T05:10:46.882224Z","iopub.execute_input":"2024-07-29T05:10:46.882766Z","iopub.status.idle":"2024-07-29T05:10:46.907069Z","shell.execute_reply.started":"2024-07-29T05:10:46.882721Z","shell.execute_reply":"2024-07-29T05:10:46.90566Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Are these simulations stable?\n\n#### Certain architectures DO provide stability (at least for a little while). \n\nThese were some important changes that helped provide stability:\n\n- Use Tanh instead of LeakyRelu or GELU\n\n**RELU functions were absolutely horrible at predicting OOD but were *better* at lowering MSE.**\n\n\n*Also if you look at the scatterplot's orange points, you can see that some of the *true* steps are actually outside the training manifold as well.*\n\n<hr>\n\n### N-Clipped\n\nThis shows how many predicted values were clipped at 5 (standard deviations).\n\n### Interesting Find\n\nUsing a GELU or LeakyRELU activation function causes this model to perform really poorly out of distribution, to the tune of\n\n```\nn clipped [ 0  0 10  0  0  0  1  0  0  0]\niter: 1\n20723510.873216335\nn clipped [320 332 351 305 321 328 317 293 307 319]\niter: 2\n-19483236.164180506\nn clipped [363 362 359 362 360 360 359 361 362 361]\niter: 3\n-57178096.18944801\nn clipped [363 363 361 361 362 362 359 362 363 363]\n```","metadata":{}},{"cell_type":"code","source":"dt = 300\nnumber_of_minutes = 60 * 10","metadata":{"execution":{"iopub.status.busy":"2024-07-29T05:10:46.908771Z","iopub.execute_input":"2024-07-29T05:10:46.909227Z","iopub.status.idle":"2024-07-29T05:10:46.927756Z","shell.execute_reply.started":"2024-07-29T05:10:46.909191Z","shell.execute_reply":"2024-07-29T05:10:46.926318Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"outputs_normed, outputs_list, real_outputs_normed_list, real_outputs_list, real_value_list, simulated_values_list = simulation_loop(\n    initial_targets, initial_values, lstm, number_of_minutes=number_of_minutes, dt=5 * 60\n)\noutputs_normed_20_min, outputs_list_20_min, real_outputs_normed_list_20_min, real_outputs_list_20_min, \\\nreal_value_list_20_min, simulated_values_list_20_min = simulation_loop(\n    initial_targets, initial_values, lstm, number_of_minutes=number_of_minutes, dt=20 * 60\n)","metadata":{"_kg_hide-output":true,"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-07-29T05:16:16.954543Z","iopub.execute_input":"2024-07-29T05:16:16.955212Z","iopub.status.idle":"2024-07-29T05:16:51.999836Z","shell.execute_reply.started":"2024-07-29T05:16:16.955167Z","shell.execute_reply":"2024-07-29T05:16:51.997889Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"outputs_normed_conv_20_min, outputs_list_conv_20_min, real_outputs_normed_list_conv_20_min, real_outputs_list_conv_20_min, \\\nreal_value_list_conv_20_min, simulated_values_list_conv_20_min = simulation_loop(\n    initial_targets, initial_values, convnet, number_of_minutes=number_of_minutes, dt=20 * 60\n)","metadata":{"execution":{"iopub.status.busy":"2024-07-29T05:18:40.958838Z","iopub.execute_input":"2024-07-29T05:18:40.959378Z","iopub.status.idle":"2024-07-29T05:18:57.448252Z","shell.execute_reply.started":"2024-07-29T05:18:40.95934Z","shell.execute_reply":"2024-07-29T05:18:57.447004Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"outputs_normed_unet, \\\noutputs_list_unet, real_outputs_normed_list_unet, \\\nreal_outputs_list_unet, real_value_list_unet, \\\nsimulated_values_list_unet = simulation_loop(\n    initial_targets, initial_values, unet, number_of_minutes=number_of_minutes, dt=dt\n)\n\noutputs_normed_unet_20_min, \\\noutputs_list_unet_20_min, real_outputs_normed_list_unet_20_min, \\\nreal_outputs_list_unet_20_min, real_value_list_unet_20_min, \\\nsimulated_values_list_unet_20_min = simulation_loop(\n    initial_targets, initial_values, unet, number_of_minutes=number_of_minutes, dt=60 * 20\n)","metadata":{"_kg_hide-output":true,"execution":{"iopub.status.busy":"2024-07-29T05:12:11.5245Z","iopub.execute_input":"2024-07-29T05:12:11.524968Z","iopub.status.idle":"2024-07-29T05:13:05.554974Z","shell.execute_reply.started":"2024-07-29T05:12:11.52493Z","shell.execute_reply":"2024-07-29T05:13:05.552701Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import numpy as np\nimport matplotlib.pyplot as plt\nimport matplotlib.ticker as mtick\n\nfrom openTSNE import TSNE\nfrom sklearn.decomposition import PCA\nfrom sklearn.preprocessing import StandardScaler\nfrom matplotlib.pyplot import figure","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-07-29T05:13:11.836026Z","iopub.execute_input":"2024-07-29T05:13:11.838734Z","iopub.status.idle":"2024-07-29T05:13:12.207754Z","shell.execute_reply.started":"2024-07-29T05:13:11.838645Z","shell.execute_reply":"2024-07-29T05:13:12.204903Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Plotting Utils","metadata":{}},{"cell_type":"code","source":"def fit_and_plot_tsne(concated):\n    tsne = TSNE(\n        n_components=2,\n        perplexity=40,\n        metric=\"euclidean\",\n        n_jobs=12,\n        random_state=42,\n        verbose=True,\n    )\n    X_embedded = tsne.fit(concated)\n    \n    figure(figsize=(12, 9), dpi=80)\n    \n    y = [\"grey\"]*15_000 + [\"limegreen\"]*10 + [\"red\"]*10 + [\"cornflowerblue\"]*10 + [\"royalblue\"]*10 + [\"orange\"]*10\n    scatter = plt.scatter(X_embedded[:, 0], X_embedded[:, 1], c=y, alpha=0.4)\n    plt.colorbar(scatter)\n    plt.title('t-SNE visualization of digits dataset')\n    plt.xlabel('t-SNE 1')\n    plt.ylabel('t-SNE 2')\n    plt.show()\n\n    scaler = StandardScaler()\n    scaled_data = scaler.fit_transform(concated)\n    \n    pca = PCA(n_components=2)\n    pca_result = pca.fit_transform(scaled_data)\n    \n    plt.figure(figsize=(8, 6))\n    plt.scatter(\n        pca_result[:, 0], pca_result[:, 1], \n        alpha=0.7, c=[\"grey\"]*15_000 + [\"green\"]*10 + [\"red\"]*10 + [\"cornflowerblue\"]*10 + [\"royalblue\"]*10 + [\"orange\"]*10)\n    plt.title('PCA of All Columns')\n    plt.xlabel('First Principal Component')\n    plt.ylabel('Second Principal Component')\n    plt.show()","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-07-29T05:13:13.674072Z","iopub.execute_input":"2024-07-29T05:13:13.67514Z","iopub.status.idle":"2024-07-29T05:13:13.691222Z","shell.execute_reply.started":"2024-07-29T05:13:13.675089Z","shell.execute_reply":"2024-07-29T05:13:13.689163Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import numpy as np\nimport matplotlib.pyplot as plt\nimport matplotlib.ticker as mtick\nimport pandas as pd\n\ndef plot_of_errors(data1, data2, plot_type=\"line\", time_mult=1):\n    \"\"\"\n    Plots a comparison of two NxJxK arrays or lists of lists, keeping the N in order.\n    Shows both sets of actual values with dashed lines for predicted values.\n    Adds a bold baseline of the minimum real value across all timesteps.\n    \n    Parameters:\n    data1: list of lists or numpy array of shape (N, J, K) - Predicted values\n    data2: list of lists or numpy array of shape (N, J, K) - Real values\n    plot_type: string, either \"line\" or \"bar\"\n    \"\"\"\n    # Convert to numpy arrays if they are lists\n    if isinstance(data1, list):\n        data1 = np.array(data1)\n    if isinstance(data2, list):\n        data2 = np.array(data2)\n    \n    # Check if the dimensions match\n    if data1.shape != data2.shape:\n        raise ValueError(\"The input arrays must have the same shape\")\n    \n    # Number of N and J values\n    N, J = data1.shape\n    \n    # Clip predicted values\n    data1 = np.clip(data1, -1.5, 5)\n    \n    # Calculate errors\n    errors = data2 - data1\n    \n    print(\"Forecasted\")\n    display(pd.DataFrame(data1)[0:5])\n    print(\"Real\")\n    display(pd.DataFrame(data2)[0:5])\n    print(\"MSE\")\n    display(pd.DataFrame(np.mean(np.power(errors, 2), axis=0)))\n    \n    # Create a figure and axis\n    fig, axes = plt.subplots(2, 1, figsize=(10, 10))\n    \n    # Calculate the minimum real value across all timesteps\n    baseline = np.min(data2)\n    \n    # Define color cycle\n    colors = plt.cm.tab10(np.linspace(0, 1, J))\n    \n    # Plot for each J group in two subplots\n    x = np.arange(1, N+1) * 20 * time_mult\n    \n    for i, ax in enumerate(axes):\n        cols = range(i*4, (i+1)*4)\n        for j in cols:\n            if j >= J:\n                continue\n            color = colors[j % 10]\n            if plot_type == \"line\":\n                ax.plot(x, data1[:, j], color=color, linestyle='--', marker='o', label=f'N={j}')\n                ax.plot(x, data2[:, j], color=color, linestyle='-', marker='s')\n            else:  # bar plot\n                bar_width = 0.35\n                ax.bar(x - bar_width/2, data1[:, j], bar_width, color=color, alpha=0.5, label=f'N={j}')\n                ax.bar(x + bar_width/2, data2[:, j], bar_width, color=color, alpha=0.9)\n        \n        # Add baseline\n        ax.axhline(y=baseline, color='grey', linestyle='-', linewidth=3, label='Rain vs Dry Line', alpha=0.4)\n        \n        # Add labels and legend\n        ax.set_ylabel('Values')\n        ax.set_title(f'Comparison of Predicted and Real Values (Columns {i*4} to {(i+1)*4-1})')\n        ax.set_xlabel('Timestep')\n        ax.legend()\n        ax.yaxis.set_major_formatter(mtick.FormatStrFormatter('%.3e'))\n    \n    # Adjust layout and show plot\n    plt.tight_layout()\n    plt.show()\n\n# Example usage with dummy data\n# data1 = np.random.rand(100, 8, 5)\n# data2 = np.random.rand(100, 8, 5)\n# plot_of_errors(data1, data2, plot_type=\"line\", time_mult=1)\n","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-07-29T05:13:24.838916Z","iopub.execute_input":"2024-07-29T05:13:24.840148Z","iopub.status.idle":"2024-07-29T05:13:24.861999Z","shell.execute_reply.started":"2024-07-29T05:13:24.840102Z","shell.execute_reply":"2024-07-29T05:13:24.860403Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"sampled_df = prepared_dataframe.loc[:, sequence_cols_all + scalar_columns].sample(15_000)\nprint(sampled_df.shape)\nprint(sampled_df.isna().sum().sum())\nsimulated_values = np.concatenate(simulated_values_list[0:4])\nprint(np.isnan(simulated_values).sum().sum())","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2024-07-29T05:13:26.968044Z","iopub.execute_input":"2024-07-29T05:13:26.968559Z","iopub.status.idle":"2024-07-29T05:13:27.738059Z","shell.execute_reply.started":"2024-07-29T05:13:26.968508Z","shell.execute_reply":"2024-07-29T05:13:27.736246Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Plot of Errors\n\n- Dashed = predicted\n- Solid = actual\n\nYou can mostly see that the initial **20 minute** prediction is the most accurate (just 1 timestep).\n","metadata":{}},{"cell_type":"markdown","source":"## LSTM dt = 5 min","metadata":{}},{"cell_type":"code","source":"plot_of_errors(np.array(outputs_normed)[:, :8, -3], np.array(real_outputs_normed_list)[:, :8, -3], \"line\")","metadata":{"execution":{"iopub.status.busy":"2024-07-29T05:13:36.830341Z","iopub.execute_input":"2024-07-29T05:13:36.830874Z","iopub.status.idle":"2024-07-29T05:13:38.198292Z","shell.execute_reply.started":"2024-07-29T05:13:36.830834Z","shell.execute_reply":"2024-07-29T05:13:38.196862Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## LSTM dt = 20 min","metadata":{}},{"cell_type":"code","source":"plot_of_errors(np.array(outputs_normed_20_min)[:, :8, -3], np.array(real_outputs_normed_list)[:, :8, -3], \"line\")","metadata":{"execution":{"iopub.status.busy":"2024-07-29T05:17:52.276309Z","iopub.execute_input":"2024-07-29T05:17:52.276919Z","iopub.status.idle":"2024-07-29T05:17:53.613091Z","shell.execute_reply.started":"2024-07-29T05:17:52.276874Z","shell.execute_reply":"2024-07-29T05:17:53.611658Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Convnet dt = 20 minutes","metadata":{}},{"cell_type":"code","source":"plot_of_errors(np.array(outputs_normed_conv_20_min)[:, :8, -3], np.array(real_outputs_normed_list)[:, :8, -3], \"line\")","metadata":{"execution":{"iopub.status.busy":"2024-07-29T05:18:57.450633Z","iopub.execute_input":"2024-07-29T05:18:57.451064Z","iopub.status.idle":"2024-07-29T05:18:58.760189Z","shell.execute_reply.started":"2024-07-29T05:18:57.45103Z","shell.execute_reply":"2024-07-29T05:18:58.758808Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## UNet dt = 300 Seconds","metadata":{}},{"cell_type":"code","source":"plot_of_errors(np.array(outputs_normed_unet)[:, :8, -3], np.array(real_outputs_normed_list)[:, :8, -3], \"line\")","metadata":{"execution":{"iopub.status.busy":"2024-07-29T05:14:20.939962Z","iopub.execute_input":"2024-07-29T05:14:20.940566Z","iopub.status.idle":"2024-07-29T05:14:22.305169Z","shell.execute_reply.started":"2024-07-29T05:14:20.940495Z","shell.execute_reply":"2024-07-29T05:14:22.303138Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## UNet dt = 20 minutes","metadata":{}},{"cell_type":"code","source":"plot_of_errors(np.array(outputs_normed_unet_20_min)[:, :8, -3], np.array(real_outputs_normed_list_unet_20_min)[:, :8, -3], \"line\")","metadata":{"execution":{"iopub.status.busy":"2024-07-29T05:13:40.293262Z","iopub.execute_input":"2024-07-29T05:13:40.29375Z","iopub.status.idle":"2024-07-29T05:13:41.580726Z","shell.execute_reply.started":"2024-07-29T05:13:40.293713Z","shell.execute_reply":"2024-07-29T05:13:41.578896Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Visualizing Where the Predictions Are on TSNE\n\n#### The basic idea here is that the small (or large) errors in the dx/dt steps will push us so far off our starting point that our subsequent predictions will fare worse and worse.\n\n- **Green**: Initial Values\n- **Red**: After first prediction (1 second)\n- **Blue**: After 2 predictions (red val + dx * 30 sec)\n- **Orange**: After 1 \"prediction\" -- using actual dx/dt values","metadata":{}},{"cell_type":"code","source":"initial_values.index","metadata":{"execution":{"iopub.status.busy":"2024-07-29T05:11:24.624216Z","iopub.status.idle":"2024-07-29T05:11:24.624654Z","shell.execute_reply.started":"2024-07-29T05:11:24.624439Z","shell.execute_reply":"2024-07-29T05:11:24.624455Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"initial_values","metadata":{"execution":{"iopub.status.busy":"2024-07-29T05:11:24.627097Z","iopub.status.idle":"2024-07-29T05:11:24.627754Z","shell.execute_reply.started":"2024-07-29T05:11:24.627425Z","shell.execute_reply":"2024-07-29T05:11:24.62745Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def process_results(simulated_values_list, n_skip=20, dt=300):\n    \"\"\"\n    Produces an array of: sampled_values, simulated\n    \"\"\"\n    initial_values_np = initial_values.to_numpy()\n    dX_dt = np.zeros_like(current_values)\n    # get only the updatable columns\n    n_indices_to_keep = len(target_sequence_indices_dx_dt)\n    # overwrite the zeros there\n    actual_dx_dt_initial = destandardize_data(target_df.iloc[::n_skip, :].to_numpy(), means_tgt, stds_tgt, [-3])\n    dX_dt[:, list(range(n_indices_to_keep))] = actual_dx_dt_initial[:, 0:n_indices_to_keep]\n    # ok so i have the actual values\n    actual_values_original_units = destandardize_data(initial_values_np, means_inp, stds_inp)\n    actual_values_original_units_next_timestep = actual_values_original_units + dX_dt * dt\n    actual_current_values = simulation_set.loc[(initial_values.index).tolist(), sequence_col_all_dx_dt + scalar_columns + sequence_target_col_all_dx_dt]\n    actual_next_values = simulation_set.loc[(initial_values.index + 384).tolist(), sequence_col_all_dx_dt + scalar_columns + sequence_target_col_all_dx_dt]\n    \n    real_standardized = standardize_data_with_given_params(\n        actual_values_original_units_next_timestep, means_inp, stds_inp\n    )\n    sampled_df = prepared_dataframe.loc[:, sequence_cols_all + scalar_columns].sample(15_000)\n    #print(sampled_df.shape)\n    #print(sampled_df.isna().sum().sum())\n    simulated_values = np.concatenate(simulated_values_list[0:4])\n    #print(np.isnan(simulated_values).sum().sum())\n    \n    concated = np.concatenate(\n        [sampled_df, simulated_values, real_standardized]\n    ).astype(np.float32) + .000000001\n    #print(f'concatenated {concated.shape} {concated.dtype}')\n    return concated","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-07-29T05:11:24.629485Z","iopub.status.idle":"2024-07-29T05:11:24.629942Z","shell.execute_reply.started":"2024-07-29T05:11:24.629744Z","shell.execute_reply":"2024-07-29T05:11:24.629761Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"concated_unet_20 = process_results(simulated_values_list_unet_20_min, n_skip, dt=60*60)\nconcated_lstm = process_results(simulated_values_list, n_skip, dt=300)\nconcated_unet = process_results(simulated_values_list_unet, n_skip, dt=300)","metadata":{"execution":{"iopub.status.busy":"2024-07-29T05:11:24.631937Z","iopub.status.idle":"2024-07-29T05:11:24.632591Z","shell.execute_reply.started":"2024-07-29T05:11:24.632253Z","shell.execute_reply":"2024-07-29T05:11:24.632278Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"fit_and_plot_tsne(concated_lstm)","metadata":{"execution":{"iopub.status.busy":"2024-07-29T05:11:24.634124Z","iopub.status.idle":"2024-07-29T05:11:24.634916Z","shell.execute_reply.started":"2024-07-29T05:11:24.634435Z","shell.execute_reply":"2024-07-29T05:11:24.634459Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"fit_and_plot_tsne(concated_unet)","metadata":{"execution":{"iopub.status.busy":"2024-07-29T05:11:24.637276Z","iopub.status.idle":"2024-07-29T05:11:24.637743Z","shell.execute_reply.started":"2024-07-29T05:11:24.637503Z","shell.execute_reply":"2024-07-29T05:11:24.637519Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"fit_and_plot_tsne(concated_unet_20)","metadata":{"execution":{"iopub.status.busy":"2024-07-29T05:11:24.63932Z","iopub.status.idle":"2024-07-29T05:11:24.639779Z","shell.execute_reply.started":"2024-07-29T05:11:24.639539Z","shell.execute_reply":"2024-07-29T05:11:24.639578Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Can We Improve This (Stable for longer)?\n\n*The real way to do this would be to use physics constraints, but I don't know if there is much I can add in that regard.*\n\n\n### The basic idea here is to force the data back onto the manifold via training on a KL objective. \n\n- Can we fine tune the model to produce data that looks like a typical weather pattern?\n- We will use a pre-trained autoencoder to compare the predicted columns and the actual columns.\n- **I kind of didn't train the autoencoder that much, so I don't know how much I trust it.**\n\nThis is similar to: https://proceedings.neurips.cc/paper/2020/file/8e2c381d4dd04f1c55093f22c59c3a08-Paper.pdf\n\n### I tried this but nothing worth publishing.","metadata":{}}]}