diff --git a/PARCtorch/demos/ShearFlow.ipynb b/PARCtorch/demos/ShearFlow.ipynb new file mode 100644 index 0000000..36b09ec --- /dev/null +++ b/PARCtorch/demos/ShearFlow.ipynb @@ -0,0 +1,615 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "
\n", + "

Physics Aware Recurrent Convolutional Neural Network (PARC): Shear-Flow Demo

\n", + "
\n", + " \"Image\n", + " \"Image\n", + " \"Image\n", + "
\n", + "
\n", + "

A customizable framework to embed physics in Deep Learning. PARC's formulation is inspired by Advection-Diffusion-Reaction processes and uses an Inductive Bias approach to constrain the Neural Network.

\n", + "\n", + "\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Why PARC?\n", + "

PARC brings together deep learning with first prinicples physics - its recurrent convolutional core is able to learn complex spatiotemporal patterns, while built-in biases (conservation laws & advection-diffusion operators) ensure every prediction is physically plausible. PARC is constructed in such a manner that is does not need to \"re-learn\" fundemental dynamics - allowing faster training with far less data.

\n", + "\n", + "

PINN (Physics-Informed Neural Networks) exists as an adjacent model to PARC. While PINN may preform decently on generics, it struggles with advection or chaotic/non-linear systems - net leading to higher computational cost. In such situations PARC is deemed the superior model with it being more scalable, efficient, and accurate than other physics-based models.

\n", + "\n", + "### Internal General PARC PDE\n", + "

Below is the general form of the partial differential equation that PARCv2 is learning - and its inital boundaries. In the case of 2D periodic incompressible Shear-Flow - we can describe the certain vairables to represent the nessearcy points in which PARC tries to model after the equation.

\n", + "\n", + "$$\n", + "\\frac{\\partial X}{\\partial t} = f(X,\\mu) + \\epsilon\n", + "$$\n", + "\n", + "- $X$ is the fields of interest - Temperature, Pressure, Reynolds, Velocities (U & V)\n", + "- $μ$ is the microstructure\n", + "\n", + "## Goal\n", + "

The goal of this notebook is to walk you through the use cases of PARC via 2D Periodic Incompressible Shear Flow. At the end of this notebook you will be presented with visualization in accordance to the equations representing conservation of momentum for fluids and fully describing fluid motion - predicting velocity fields and tracer distributions based on current states

\n", + "\n", + "### Shear Flow Modeling\n", + "

Lets have a brief overview of what Shear Flow is - it's a type of fluid motion in which adjacent layers move parallel to each other at different velocities. When each layer slides past its neighboring one, there is a velocity gradient perpendicular to the main flow direction - this gradeint is what produces a shear stress within the fluid (defined by Newtons Law of Viscosity). An example of this can be seen with an aeroplane - as the plane moves forward the air closest to the wings \"stick\" to the wings surface, while air farther away moves at near free-stream velocity.

\n", + "\n", + "

Mathematically, the primary equations govenrning our simulation of Shear Flow are:

\n", + "\n", + "$$\n", + "\\frac{\\partial u}{\\partial t} + \\nabla p - \\nu \\Delta u = -u\\cdot \\nabla u\n", + "$$\n", + "\n", + "$$\n", + "\\frac{\\partial s}{\\partial t} - D\\Delta s = -u \\cdot \\nabla s\n", + "$$\n", + "\n", + "- $u = (u_x, u_y)$ is the velocity field (with horizontal and vertical components)\n", + "- $p$ is pressure ($\\int p\\ = 0$ enforcing gauge)\n", + "- $s$ is a passive tracer (scalar field)\n", + "- $v = \\frac{1}{Re}$ is the kinematic viscocity - Reynolds Number (parameter one)\n", + "- $D = \\frac{v}{Sc}$ is the tracer diffusivity - Schmidts Number (parameter two)\n", + "\n", + "\n", + "### Setting Up\n", + "The notebook will guide you through from start to finish in preparing, training, and modeling a physics-based equation. The notebook will cover the following:\n", + "\n", + "- loading and preparing data for Shear Flow\n", + "- Using PARCv2 Model to learn and predict time evolution of $u(x,t)$\n", + "- Evaluating the model's preformance and compare predicted resutls to ground truth\n", + "\n", + "

This serves as a guide to training a PARC model for 2D Periodic Incompressible Shear-Flow. Here are the inital steps to take before you begin training your PARC model!

\n", + "\n", + "Download & Prepare Data:\n", + "- Extact the data and ensure that it placed in the following directory: PARCtorch/PARCtorch/data \n", + "- If needed update the paths in the notebook accordlingly (for train_dir, test_dir)\n", + "\n", + "Install PARCtorch:\n", + "- Ensure PARCtorch is installed in your Python Environemnt - view instructions on installation at: https://github.com/baeklab/PARCtorch\\n\",\n", + "- Note: Ensure to be in the same environemnt/kernel when continuing in this notebook.\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Define Filesystem Paths for Data\n", + "

To begin with, we will import some necessary Python modules and utility functions. In this cell we are defining the filesystem paths for our Shear Flow training and testing data - these paths are then used late for computing normalization and building DataLoaders

" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import sys\n", + "import os\n", + "from pathlib import Path\n", + "from PARCtorch.data.normalization import compute_min_max\n", + "\n", + "# Add the root directory (PARCtorch) to the system path\n", + "sys.path.append(os.path.abspath(os.path.join(os.getcwd(), \"..\")))\n", + "\n", + "# Define data directories\n", + "train_dir = Path(\"/scratch/wex8ke/shear-flow/data/train\")\n", + "test_dir = Path(\"/scratch/wex8ke/shear-flow/data/test\")\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Construct Data Loader for Training\n", + "

We next create a DataLoader for training. The DataLoader in PyTorch is a crucial utility that facilitates efficient data handling for training and evaluating machine learning models. It abstracts the process of fetching, batching, and shuffling data, ensuring that the model is fed with properly formatted inputs in an optimal way. Specifically, it helps with:

\n", + "\n", + "\n", + "\n", + "

We are now instantiating the WellDatasetInterface using min_val & max_val, future_steps count - all wrapped in a PyTorch DataLoader

" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "from PARCtorch.data.dataset import WellDatasetInterface\n", + "from the_well.data.datasets import WellDataset\n", + "from torch.utils.data import DataLoader\n", + "from torchvision.transforms import Compose, RandomCrop, CenterCrop, Resize\n", + "import PARCtorch.data.dataset as ds_mod\n", + "import torch\n", + "\n", + "# override the __del__ so it never tries to iterate that missing attr\n", + "ds_mod.WellDatasetInterface.__del__ = lambda self: None\n", + "\n", + "# Now import the utilities\n", + "from PARCtorch.data.dataset import (\n", + " WellDatasetInterface,\n", + " custom_collate_fn,\n", + " InitialConditionDataset,\n", + " initial_condition_collate_fn,\n", + ")\n", + "from PARCtorch.utilities.viz import (\n", + " visualize_channels,\n", + " save_gifs_with_ground_truth,\n", + ")\n", + "\n", + "orig_h, orig_w = 256, 512\n", + "spatial_percent = 0.0625\n", + "\n", + "new_h = int(orig_h * spatial_percent) # 16\n", + "new_w = int(orig_w * spatial_percent) # 32\n", + "\n", + "\n", + "# spatial transformation\n", + "patch_size = (64, 64)\n", + "train_spatial_transform = Compose([\n", + " RandomCrop(patch_size),\n", + " Resize((new_h, new_w)),\n", + "])\n", + "gt_spatial_transform = Compose([\n", + " CenterCrop(patch_size),\n", + " Resize((new_h, new_w)),\n", + "])\n", + "\n", + "future_steps = 8\n", + "batch_size = 2\n", + "\n", + "# Initialize the dataset\n", + "ds = WellDatasetInterface(\n", + " future_steps=future_steps,\n", + " min_val=torch.tensor([0.03, 0.3217, 0.0541, 0.0, 0.0]),\n", + " max_val=torch.tensor([3.16, 196.3853, 2.2342, 1.5395, 1.5395]),\n", + " well_dataset_args={\n", + " \"well_base_path\": \"/scratch/wex8ke\",\n", + " \"well_dataset_name\": \"shear_flow\",\n", + " \"well_split_name\": \"train\",\n", + " \"exclude_filters\": [\"shear_flow_Reynolds_5e5_Schmidt_1e0.hdf5\"]\n", + " }\n", + ")\n", + "\n", + "ds.transform = train_spatial_transform\n", + "\n", + "# Create DataLoader for training dataset\n", + "train_loader = DataLoader(\n", + " ds,\n", + " batch_size=batch_size,\n", + " shuffle=False,\n", + " num_workers=1,\n", + " collate_fn=custom_collate_fn,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Visualize the Data - Check it was Loaded Properly\n", + "

We are now going to visualize our data to verify if the correct data is loading and formatting. In the cell below we are iterating over the previously defined DataLoader and fetching one batch of data. We can now plot its six channels: Velocity U, Velocity V, Tracer, Preassure, Rayleigh, and Schmidt. We also verify that the following tensors match our expectation - in terms of shape and numerical range:

\n", + "\n", + "\n", + "\n", + "

Once confirmed that the training batches are correct, we can now assemble our PARC model.

" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Fetch a batch and visualize\n", + "for batch in train_loader:\n", + " ic, t0, t1, target = batch\n", + " channel_names = [\n", + " \"Pressure (P)\",\n", + " \"Reynolds (R)\",\n", + " \"Velocity U\",\n", + " \"Velocity V\",\n", + " ]\n", + " custom_cmaps = [\"seismic\", \"seismic\", \"seismic\", \"seismic\"]\n", + "\n", + " visualize_channels(\n", + " ic,\n", + " t0,\n", + " t1,\n", + " target,\n", + " channel_names=channel_names,\n", + " channel_cmaps=custom_cmaps,\n", + " )\n", + " break # Visualize one batch for now" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "

Build Your PARC Model

\n", + "

In this section, we are constructing a PARCv2 model, which is designed to handle spatiotemporal data, such as fluid dynamics simulations. The model leverages various components, including differentiators and integrators, to solve physical equations like 2D periodic incompressible Shear-Flow.

\n", + "\n", + "

Key Components:

\n", + "\n", + "\n", + "

The model is then wrapped into the PARCv2 class, which combines the differentiator, integrator, and loss function (L1Loss). Finally, an Adam optimizer is initialized to train the model by adjusting its parameters to minimize the error between predictions and ground truth data.

\n", + "\n", + "

This setup allows the model to learn how to predict future states in complex physical systems by embedding domain-specific knowledge.

\n" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "from PARCtorch.PARCv2 import PARCv2\n", + "from PARCtorch.differentiator.differentiator import ADRDifferentiator\n", + "from PARCtorch.differentiator.finitedifference import FiniteDifference\n", + "from PARCtorch.integrator.integrator import Integrator\n", + "from PARCtorch.integrator.heun import Heun\n", + "from PARCtorch.utilities.unet import UNet\n", + "\n", + "from torch.optim import Adam" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "# Defining 2D Imcompressible Shear Flow\n", + "PARCv2.check = lambda self: True # temp - for no attr named check in PARCv2\n", + "PARCv2.check_msg = lambda self, chk: None # temp - has no attribute named check_msg\n", + "\n", + "n_fe_features = 64 # Number of features extracted by the UNet\n", + "\n", + "# Initialize the UNet architecture for feature extraction\n", + "unet_sf = UNet(\n", + " [64, 64* 2, 64 * 4], \n", + " 6, \n", + " n_fe_features,\n", + " 3,\n", + " \"zeros\",\n", + " up_block_use_concat=[False, True], \n", + " skip_connection_indices=[0], \n", + ")\n", + "\n", + "# Initialize finite difference method for numerical differentiation\n", + "right_diff = FiniteDifference(padding_mode=\"replicate\").cuda() # Use replication padding to handle boundary conditions\n", + "\n", + "# Initialize Heun's method for numerical integration\n", + "heun_int = Heun().cuda() \n", + "\n", + "# Create the Differentiator, responsible for calculating advection and diffusion\n", + "sf_diff = ADRDifferentiator(\n", + " 4, \n", + " n_fe_features, \n", + " [0, 1, 2, 3, 4, 5], \n", + " [0, 1, 2, 3, 4, 5], \n", + " unet_sf, \n", + " \"constant\", \n", + " right_diff, \n", + " False \n", + ").cuda()\n", + "\n", + "# Create the Integrator, responsible for solving the Poisson equation and performing integration\n", + "sf_int = Integrator(\n", + " True, \n", + " [], \n", + " heun_int, \n", + " [None, None, None, None], \n", + " \"constant\", \n", + " right_diff,\n", + " 3, \n", + " 64\n", + ").cuda()\n", + "\n", + "# Define the loss function (L1 Loss is typically used for regression tasks)\n", + "criterion = torch.nn.L1Loss().cuda()\n", + "\n", + "# Initialize the PARCv2 model with the differentiator, integrator, and loss function\n", + "model = PARCv2(\n", + " sf_diff, \n", + " sf_int, \n", + " criterion\n", + ").cuda()\n", + "\n", + "# Set up the optimizer (Adam is a popular choice for adaptive learning rate optimization)\n", + "optimizer = Adam(\n", + " model.parameters(), \n", + " lr=1e-5\n", + ") " + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Train the Model \n", + "

Now, we will be training the PARC model using the training data and the paramaters we have setup above. Here we iterate over the dataset given by train_loader batch-by-batch and in each iteration we:

\n", + "\n", + "" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "from PARCtorch.train import train_model\n", + "\n", + "# Example usage:\n", + "train_model(\n", + " model,\n", + " train_loader,\n", + " criterion,\n", + " optimizer,\n", + " num_epochs=10,\n", + " save_dir='/scratch/wex8ke/shear_flow',\n", + " app=\"ns\",\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Load the Model\n", + "

Post training, we can now load the previously trained model weights into our PARC model for evalutation. Here, we locate the model weights and load them into the exisiting PARC model architecture in order to prep for visualization.

" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "from PARCtorch.utilities.load import load_model_weights\n", + "\n", + "model_loc = '/scratch/wex8ke/shear_flow/model.pth'\n", + "\n", + "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n", + "model_weights_path = (model_loc)\n", + "model = load_model_weights(model, model_weights_path, device)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Sequence DataLoader\n", + "

We now create a sequence DataLoader designed to test the dataset - providing the initial conditions as t=0 for generating predictiond over multiple timesteps. Using DataLoader we wrap the dataset for efficent batching and processing.

\n", + "\n", + "

We also create the Initial Condition Dataset, which loades the initial state for the physical system, the previously computed normalization parameters, and prepares the dataset to be fed into the model. We also " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from torch.utils.data import DataLoader\n", + "from PARCtorch.data.dataset import WellDatasetInterface, initial_condition_collate_fn\n", + "\n", + "min_val_tensor = torch.tensor([0.03, 0.3217, 0.0541, 0.0, 0.0, 0.0])\n", + "max_val_tensor = torch.tensor([3.16, 196.3853, 2.2342, 1.5395, 1.5395, 1.0]) \n", + "\n", + "test_dir = \"/scratch/wex8ke/shear-flow/data/test\"\n", + "future_steps = 5\n", + "batch_size = 2\n", + "\n", + "gt_spatial_transform = Compose([\n", + " CenterCrop(patch_size),\n", + " Resize((128, 128)),\n", + "])\n", + "\n", + "seq_dataset = WellDatasetInterface(\n", + " well_dataset_args = {\n", + " \"path\": test_dir\n", + " },\n", + " future_steps = future_steps,\n", + " min_val = min_val_tensor,\n", + " max_val = max_val_tensor,\n", + " delta_t = 2.0,\n", + " add_constant_scalars= True, \n", + ")\n", + "seq_dataset.transform = gt_spatial_transform\n", + "\n", + "seq_loader = DataLoader(\n", + " seq_dataset,\n", + " batch_size = batch_size,\n", + " shuffle = False,\n", + " num_workers = 1,\n", + " pin_memory = True,\n", + " collate_fn = initial_condition_collate_fn,\n", + ")\n", + "\n", + "\n", + "print(f\"seq_loader ready: {len(seq_loader)} batches\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Ground Truth Loader\n", + "

Here, we set up the Ground Truth Loader to provide a comparison of the actual future states of the system against the models predictions. Similar set up to previous DataLoader setups in this notebook.

" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "test_dir = \"/standard/sds_baek_energetic/data/physics/the_well/datasets/shear_flow/data/test\"\n", + "\n", + "gt_dataset = WellDatasetInterface(\n", + " well_dataset_args = {\"path\": test_dir},\n", + " future_steps = future_steps,\n", + " min_val = min_val_tensor,\n", + " max_val = max_val_tensor,\n", + " delta_t = 2.0,\n", + " add_constant_scalars= True,\n", + ")\n", + "\n", + "gt_loader = DataLoader(\n", + " gt_dataset,\n", + " batch_size = batch_size,\n", + " shuffle = False,\n", + " num_workers = 1,\n", + " pin_memory = True,\n", + " collate_fn = custom_collate_fn,\n", + ")\n", + "print(f\"gt_loader built → {len(gt_loader)} batches\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Visualize the Results\n", + "

We now finially visualize our results. Here, we run inferences using thr trained PARC model's predictions against the ground truth. Visualization gifs of the predictions and ground truth are outputted - just runs on the first batch for demo.

" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "import os\n", + "import numpy as np\n", + "import matplotlib.pyplot as plt\n", + "import imageio\n", + "\n", + "filename_prefix = \"comparison_sample0\"\n", + "\n", + "preds_np = predictions.detach().cpu().numpy()\n", + "gts_np = ground_truth.detach().cpu().numpy()\n", + "\n", + "T, B, C, H, W = preds_np.shape\n", + "assert gts_np.shape == (T, B, C, H, W)\n", + "assert len(channels) == C\n", + "assert len(cmaps) == C\n", + "\n", + "out_dir = \"/scratch/wex8ke/shear-flow\"\n", + "os.makedirs(out_dir, exist_ok=True)\n", + "\n", + "for c in range(C):\n", + " frames = []\n", + " for t in range(T):\n", + " fig, ax = plt.subplots(figsize=(W/100, H/100), dpi=100)\n", + "\n", + " # Use batch_idx = 0\n", + " pred_img = preds_np[t, 0, c, :, :]\n", + " gt_img = gts_np [t, 0, c, :, :]\n", + "\n", + " im = ax.imshow(pred_img, cmap=cmaps[c], vmin=gt_img.min(), vmax=gt_img.max())\n", + " ax.set_title(f\"{channels[c]}: t={t} (pred)\")\n", + " ax.axis(\"off\")\n", + " fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)\n", + "\n", + " fig.canvas.draw()\n", + " buf, (width, height) = fig.canvas.print_to_buffer()\n", + " rgba = np.frombuffer(buf, dtype=np.uint8).reshape((height, width, 4))\n", + " rgb = rgba[..., :3]\n", + " frames.append(rgb)\n", + " plt.close(fig)\n", + "\n", + " out_path = os.path.join(out_dir, f\"{filename_prefix}_ch{c}.gif\")\n", + " imageio.mimsave(out_path, frames, fps=max(1, int(1/0.2)))\n", + " print(f\"Saved GIF: {out_path}\")\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import pickle\n", + "import matplotlib.pyplot as plt\n", + " \n", + " \n", + "with open( \"/scratch/wex8ke/shear_flow/training_losses.pkl\", 'rb' ) as f: \n", + " history = pickle.load(f)\n", + " \n", + "plt.figure(figsize=(10,6))\n", + "plt.title('Reduced Resolution (64 x 32) Model Loss Curve')\n", + "plt.xlabel('Epoch')\n", + "plt.ylabel('Loss')\n", + "plt.plot( history)\n", + "plt.savefig('normalization_loss_curve.png')\n", + " " + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.6" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +}