init: modeling Ahmed Body dataset with DoMINO using physicsnemo

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

Dockerfile

@@ -0,0 +1,57 @@
+FROM nvcr.io/nvidia/physicsnemo/physicsnemo:25.06
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    wget \
+    git \
+    xvfb \
+    && rm -rf /var/lib/apt/lists/*
+
+# Upgrade pip
+RUN python -m pip install --upgrade pip
+
+# Create non-root user
+RUN useradd -m -u 1001 user
+USER user
+ENV HOME=/home/user
+RUN mkdir $HOME/.cache $HOME/.config $HOME/.local && chmod -R 777 $HOME
+ENV PATH=/home/user/.local/bin:$PATH
+
+# Set working directory
+WORKDIR $HOME/app
+
+# Create data directory with proper permissions
+USER root
+RUN mkdir /domino-demo && chown user:user /domino-demo && chmod 777 /domino-demo
+# Add NGC CLI
+RUN wget https://ngc.nvidia.com/downloads/ngccli_linux.zip
+RUN unzip ngccli_linux.zip
+RUN chmod u+x ngc-cli/ngc
+RUN --mount=target=/root/on_startup.sh,source=on_startup.sh,readwrite \
+	bash /root/on_startup.sh
+USER user
+
+# Install Python dependencies
+RUN pip install --no-cache-dir numpy pyvista vtk matplotlib tqdm numpy-stl torchinfo
+
+COPY --chown=user login.html $HOME/app/
+COPY --chown=user on_startup.sh $HOME/app/
+COPY --chown=user README.md $HOME/app/
+COPY --chown=user start_server.sh $HOME/app/
+COPY --chown=user domino-data-preprocessing.ipynb /domino-demo/
+COPY --chown=user domino-training-test.ipynb /domino-demo/
+COPY --chown=user introduction.ipynb /domino-demo/
+COPY --chown=user outputs /domino-demo/outputs
+COPY --chown=user README.md /domino-demo/
+
+RUN chmod -R 777 /domino-demo/
+
+RUN chmod +x start_server.sh
+
+COPY --chown=user login.html /usr/local/lib/python3.12/dist-packages/jupyter_server/templates/login.html
+
+# Expose port
+EXPOSE 7860
+
+# Start the server
+CMD ["./start_server.sh"]

README.md

@@ -0,0 +1,17 @@
+---
+title: DoMINO with Ahmed Body Dataset - Multi-Scale Neural Operator for CFD
+emoji: 🟢
+colorFrom: green
+colorTo: gray
+sdk: docker
+app_port: 7860
+pinned: false
+tags:
+  - physics
+  - cfd
+  - machine-learning
+  - neural-operators
+  - fluid-dynamics
+  - scientific-computing
+---
+

domino-data-preprocessing.ipynb

@@ -0,0 +1,676 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Physics NeMo External Aerodynamics DLI\n",
+    "\n",
+    "## Notebook 1 - Preprocessing Ahmed body *surface* dataset\n",
+    "\n",
+    "### Introduction\n",
+    "\n",
+    "For educational purposes, it's important to use lightweight datasets that are easy to store and manage, especially for users who may not have access to high-performance computing resources. One such dataset is the **Ahmed body surface data**, which includes 3D surface geometry, pressure and wall shear stress data for variations in the Ahmed body geometry and inlet Reynolds number. This dataset is a great choice because it is relatively small in size, yet provides valuable information about aerodynamic simulations. It’s ideal for teaching and experimentation, as it won’t demand excessive storage or computational power. *Note that this dataset was created by the NVIDIA PhysicsNeMo development team and differs from other similar datasets hosted on cloud platforms like AWS.*\n",
+    "\n",
+    "In this notebook, we will walk through the preprocessing steps required to prepare the **Ahmed body surface dataset** for training with the **DoMINO model**, to predict surface quantities like pressure and wall shear stress. The DoMINO model requires 3D surface geometry in **STL format**. The **STL (Stereolithography)** format is a widely used file format for representing 3D surface geometry in computer-aided design (CAD) applications. It describes the surface of a 3D object using a collection of triangular facets, making it a common format for 3D printing and computational geometry. So, as the first step, we’ll extract the 3D surface geometry from the **VTP files**. These files are commonly used in the **VTK (Visualization Toolkit)** format, which stores surface data as **PolyData**—a structure that represents points, lines, and polygons on the surface.\n",
+    "\n",
+    "To make the dataset more suitable for machine learning, we will convert the **VTP (VTK PolyData)** format into **NPY (NumPy)** format. This conversion makes the data easier to work with in machine learning workflows, as NumPy arrays are optimized for numerical operations, making computations faster and more efficient. After converting the data into NPY format, it can be stored on disk, where it will be readily accessible for training the model and further analysis.\n",
+    "\n",
+    "Key aspects of this training:\n",
+    "- Understanding the Ahmed body geometry and its aerodynamic characteristics\n",
+    "- Processing CFD mesh data for deep learning applications\n",
+    "\n",
+    "## Table of Contents\n",
+    "- [Step 1: Define Experiment Parameters and Dependencies](#step-1-define-experiment-parameters-and-dependencies)\n",
+    "  - [Loading Required Libraries](#loading-required-libraries)\n",
+    "  - [Dependencies](#dependencies)\n",
+    "  - [Experiment Parameters and Variables](#experiment-parameters-and-variables)\n",
+    "- [Step 2: Convert VTK to STL Files](#step-2-convert-vtk-to-stl-files)\n",
+    "  - [Understanding the Conversion Process](#understanding-the-conversion-process)\n",
+    "  - [Key Components and Libraries](#key-components-and-libraries)\n",
+    "  - [Important Considerations](#important-considerations)\n",
+    "  - [Implementation Overview](#implementation-overview)\n",
+    "- [Step 3: Visualizing STL Meshes](#Step-3:-Visualizing-STL-Meshes)\n",
+    "- [Step 4: Convert CFD Results to NPY Format](#Step-4:-Convert-CFD-Results-to-NPY-Format)   "
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### **Step 1: Define Experiment Parameters and Dependencies**\n",
+    "\n",
+    "The first step in training the DoMINO model on the Ahmed body surface dataset is to set up our experiment environment and define the necessary parameters. This includes specifying paths to our data, configuring training settings, and ensuring all required libraries are available.\n",
+    "\n",
+    "Key components we need to set up:\n",
+    "- Data paths for training and validation sets\n",
+    "- Model hyperparameters and training configurations\n",
+    "- Visualization settings for results\n",
+    "- Required Python libraries for mesh processing and deep learning\n",
+    "\n",
+    "### Loading Required Libraries\n",
+    "\n",
+    "Before we proceed with the experiment setup, let's first import all the necessary libraries. These libraries will be used for:\n",
+    "- Mesh processing and visualization (`vtk`, `pyvista`)\n",
+    "- Data handling and file operations (`pathlib`, `concurrent.futures`)\n",
+    "- Progress tracking and visualization (`tqdm`, `matplotlib`)\n",
+    "- PyTorch provides data primitives: `torch.utils.data.Dataset` that allow you to use pre-loaded datasets as well as your own data. `Dataset` stores the samples and their corresponding labels.\n",
+    "- Important utilities for data processing and training, testing DoMINO (`physicsnemo.utils.domino.utils`)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import random\n",
+    "from concurrent.futures import ProcessPoolExecutor\n",
+    "from pathlib import Path\n",
+    "from typing import Union\n",
+    "\n",
+    "import numpy as np\n",
+    "import pyvista as pv\n",
+    "import vtk\n",
+    "from stl import mesh\n",
+    "from tqdm import tqdm\n",
+    "\n",
+    "from physicsnemo.utils.domino.utils import *\n",
+    "from torch.utils.data import Dataset"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Experiment Parameters and Variables\n",
+    "\n",
+    "In this section, we set up all the essential parameters and variables required for the Ahmed body experiment. \n",
+    "**Before proceeding, navigate to the data directory and extract the `ahmed_body_dataset.zip` archive. This file contains several sample `.vtp` files needed to run the scripts in this notebook and others.**"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Directory and Path Configuration\n",
+    "DATA_DIR = Path(\"/data/physicsnemo_ahmed_body_dataset_vv1/dataset\")  # Root directory for dataset\n",
+    "\n",
+    "# Physical Variables\n",
+    "VOLUME_VARS = [\"p\"]  # Volume variables to predict (pressure)\n",
+    "SURFACE_VARS = [\"p\", \"wallShearStress\"]  # Surface variables to predict\n",
+    "AIR_DENSITY = 1.205  # Air density in kg/m³"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### **Step 2: Convert VTK to STL Files**\n",
+    "\n",
+    "The second step in our workflow involves converting the CFD simulation data from VTK format to STL format.\n",
+    "\n",
+    "#### Understanding the Conversion Process\n",
+    "\n",
+    "The conversion from VTK to STL involves several key steps:\n",
+    "1. Reading the VTK PolyData file using specialized readers\n",
+    "2. Extracting the surface geometry and mesh data\n",
+    "3. Converting the data while preserving topology and surface properties\n",
+    "4. Saving the result in binary STL format\n",
+    "\n",
+    "#### Key Components and Libraries\n",
+    "\n",
+    "We'll use the following libraries for this conversion:\n",
+    "\n",
+    "1. **VTK (Visualization Toolkit)**\n",
+    "   - `vtk`: Reads VTK PolyData files (.vtp)\n",
+    "   - `vtkSTLWriter`: Writes data in STL format\n",
+    "   - `vtkPolyData`: Manages surface mesh data structures\n",
+    "\n",
+    "2. **File System Operations**\n",
+    "   - `os.path`: Handles file paths and directory operations\n",
+    "   - `pathlib.Path`: Provides modern path handling capabilities\n",
+    "\n",
+    "#### Important Considerations\n",
+    "\n",
+    "During the conversion process, we need to ensure:\n",
+    "- Surface normal vectors are preserved correctly\n",
+    "- Mesh quality and topology are maintained\n",
+    "- The output is compatible with the DoMINO model's requirements\n",
+    "- Memory is managed efficiently for large datasets\n",
+    "\n",
+    "#### Implementation Overview\n",
+    "\n",
+    "The conversion is implemented through two main functions:\n",
+    "\n",
+    "1. **Environment Setup**\n",
+    "```python\n",
+    "def setup_environment(data_dir: str):\n",
+    "    \"\"\"Sets up the working directory and returns relevant paths.\"\"\"\n",
+    "    # Returns paths for dataset, info files, STL files, and surface data\n",
+    "```\n",
+    "\n",
+    "2. **VTK to STL Conversion**\n",
+    "```python\n",
+    "def convert_vtk_to_stl(vtk_filename: str, stl_filename: str):\n",
+    "    \"\"\"Converts a single .vtp file to .stl format.\"\"\"\n",
+    "    # Uses vtkXMLPolyDataReader and vtkSTLWriter for conversion\n",
+    "```\n",
+    "\n",
+    "Let's proceed with implementing these functions and performing the conversion:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def setup_environment(data_dir: str):\n",
+    "    \"\"\"Sets up the working directory and returns relevant paths.\"\"\"\n",
+    "    print(\"=== Environment Setup ===\")\n",
+    "    print(f\"Current data directory: {data_dir}\")\n",
+    "\n",
+    "    dataset_paths = {split: os.path.join(data_dir, split) for split in [\"train\", \"validation\", \"test\"]}\n",
+    "    info_paths = {k: os.path.join(data_dir, f\"{k}_info\") for k in dataset_paths}\n",
+    "    stl_paths = {k: os.path.join(data_dir, f\"{k}_stl_files\") for k in dataset_paths}\n",
+    "    surface_paths = {k: os.path.join(data_dir, f\"{k}_prepared_surface_data\") for k in dataset_paths}\n",
+    "\n",
+    "    return dataset_paths, info_paths, stl_paths, surface_paths\n",
+    "\n",
+    "\n",
+    "def convert_vtk_to_stl(vtk_filename: str, stl_filename: str):\n",
+    "    \"\"\"Converts a single .vtp file to .stl format.\"\"\"\n",
+    "    reader = vtk.vtkXMLPolyDataReader()\n",
+    "    reader.SetFileName(vtk_filename)\n",
+    "    reader.Update()\n",
+    "\n",
+    "    if not reader.GetOutput():\n",
+    "        print(f\"[ERROR] Failed to read {vtk_filename}\")\n",
+    "        return\n",
+    "\n",
+    "    writer = vtk.vtkSTLWriter()\n",
+    "    writer.SetFileName(stl_filename)\n",
+    "    writer.SetInputConnection(reader.GetOutputPort())\n",
+    "    writer.Write()\n",
+    "\n",
+    "    del reader, writer  # Free memory\n",
+    "\n",
+    "\n",
+    "def process_file(vtp_file: str, output_path: str):\n",
+    "    \"\"\"Processes a single .vtp file and saves it as .stl.\"\"\"\n",
+    "    output_file = os.path.join(output_path, os.path.basename(vtp_file).replace(\".vtp\", \".stl\"))\n",
+    "    convert_vtk_to_stl(vtp_file, output_file)\n",
+    "\n",
+    "\n",
+    "def convert_vtp_to_stl_batch(dataset_paths: dict, stl_paths: dict):\n",
+    "    \"\"\"Processes all .vtp files in dataset_paths and saves them in output_paths.\"\"\"\n",
+    "    print(\"\\n=== Starting Conversion Process ===\")\n",
+    "\n",
+    "    for path in stl_paths.values(): os.makedirs(path, exist_ok=True)\n",
+    "\n",
+    "    for key, dataset_path in dataset_paths.items():\n",
+    "        stl_path = stl_paths[key]\n",
+    "        vtp_files = [os.path.join(dataset_path, f) for f in os.listdir(dataset_path) if f.endswith('.vtp')]\n",
+    "\n",
+    "        if not vtp_files:\n",
+    "            print(f\"[WARNING] No .vtp files found in {dataset_path}\")\n",
+    "            continue\n",
+    "\n",
+    "        print(f\"\\nProcessing {len(vtp_files)} files from {dataset_path} → {stl_path}...\")\n",
+    "\n",
+    "        with ProcessPoolExecutor() as executor:\n",
+    "            list(tqdm(executor.map(process_file, vtp_files, [stl_path] * len(vtp_files)),\n",
+    "                      total=len(vtp_files), desc=f\"Converting {key}\", dynamic_ncols=True))\n",
+    "\n",
+    "    print(\"=== All Conversions Completed Successfully ===\")\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Lets convert the files:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dataset_paths, info_paths, stl_paths, surface_paths = setup_environment(DATA_DIR)\n",
+    "convert_vtp_to_stl_batch(dataset_paths, stl_paths)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### **Step 3: Visualizing STL Meshes**\n",
+    "\n",
+    "The third step in our workflow focuses on visualizing the converted STL meshes to verify the success of our conversion process. This step is crucial because:\n",
+    "\n",
+    "#### Understanding the Visualization Process\n",
+    "\n",
+    "The visualization of STL meshes involves several key aspects:\n",
+    "1. Loading the STL files using appropriate visualization libraries\n",
+    "2. Setting up the visualization environment with proper parameters\n",
+    "3. Rendering the mesh with appropriate colors and properties\n",
+    "4. Providing interactive controls for inspection\n",
+    "\n",
+    "#### Key Components and Libraries\n",
+    "\n",
+    "We'll use the following libraries for visualization:\n",
+    "\n",
+    "1. **PyVista**\n",
+    "   - `pv.read()`: Loads STL files\n",
+    "   - `pv.Plotter()`: Creates interactive visualization windows\n",
+    "   - `pv.wrap()`: Converts VTK objects to PyVista meshes\n",
+    "\n",
+    "2. **Matplotlib**\n",
+    "   - For static 2D visualizations if needed\n",
+    "   - For saving visualization outputs\n",
+    "\n",
+    "#### Important Visualization Parameters\n",
+    "\n",
+    "During the visualization process, we need to consider:\n",
+    "- Mesh surface properties (color, opacity)\n",
+    "- Camera position and orientation\n",
+    "- Lighting conditions\n",
+    "- Interactive controls for rotation and zoom\n",
+    "- Quality of the rendered output\n",
+    "\n",
+    "#### Implementation Overview\n",
+    "\n",
+    "The visualization is implemented through several key functions:\n",
+    "\n",
+    "1. **Mesh Loading**\n",
+    "```python\n",
+    "def load_stl_mesh(file_path: str):\n",
+    "    \"\"\"Loads an STL file and returns a PyVista mesh object.\"\"\"\n",
+    "    # Uses PyVista to read and process the STL file\n",
+    "```\n",
+    "\n",
+    "2. **Interactive Visualization**\n",
+    "```python\n",
+    "def plot_stl_comparison(mesh1, mesh2, title1=\"Case 3\", title2=\"Case 9\", volume1=None, volume2=None):\n",
+    "    \"\"\"Create a multi-view comparison visualization of two STL meshes.\"\"\"\n",
+    "    # Sets up the plotter and displays the meshes\n",
+    "```\n",
+    "\n",
+    "Let's proceed with implementing these functions and visualizing our converted meshes:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def load_stl(stl_path: str):\n",
+    "    \"\"\"Load an STL file and return its PyVista mesh.\"\"\"\n",
+    "    stl_file = Path(stl_path)\n",
+    "    if not stl_file.exists():\n",
+    "        print(f\"[ERROR] STL file not found: {stl_path}\")\n",
+    "        return None\n",
+    "\n",
+    "    try:\n",
+    "        mesh = pv.read(str(stl_file))\n",
+    "        return mesh\n",
+    "    except Exception as e:\n",
+    "        print(f\"[ERROR] Failed to load STL file: {e}\")\n",
+    "        return None\n",
+    "\n",
+    "def load_stl(stl_path: str):\n",
+    "    \"\"\"Load an STL file and return its PyVista mesh.\"\"\"\n",
+    "    stl_file = Path(stl_path)\n",
+    "    if not stl_file.exists():\n",
+    "        print(f\"[ERROR] STL file not found: {stl_path}\")\n",
+    "        return None\n",
+    "\n",
+    "    try:\n",
+    "        mesh = pv.read(str(stl_file))\n",
+    "        return mesh\n",
+    "    except Exception as e:\n",
+    "        print(f\"[ERROR] Failed to load STL file: {e}\")\n",
+    "        return None\n",
+    "\n",
+    "def plot_stl_comparison(mesh1, mesh2, title1=\"Case 3\", title2=\"Case 9\"):\n",
+    "    \"\"\"Create a multi-view comparison visualization of two STL meshes.\"\"\"\n",
+    "    # Start virtual frame buffer\n",
+    "    pv.start_xvfb()\n",
+    "\n",
+    "    # Create plotter with off-screen rendering\n",
+    "    pl = pv.Plotter(shape=(2, 4), window_size=[1600, 800], off_screen=True)\n",
+    "\n",
+    "\n",
+    "    # Display parameters for each case with brighter colors\n",
+    "    params_case1 = {\n",
+    "        \"show_edges\": True,\n",
+    "        \"opacity\": 1.0,\n",
+    "        \"edge_color\": 'black',\n",
+    "        \"line_width\": 0.5,\n",
+    "        \"color\": [0.6, 0.8, 1.0],\n",
+    "        \"smooth_shading\": True,\n",
+    "        \"specular\": 0.4,\n",
+    "        \"specular_power\": 10,\n",
+    "        \"diffuse\": 0.9,\n",
+    "        \"ambient\": 1.0,\n",
+    "    }\n",
+    "\n",
+    "    params_case2 = {\n",
+    "        \"show_edges\": True,\n",
+    "        \"opacity\": 1.0,\n",
+    "        \"edge_color\": 'black',\n",
+    "        \"line_width\": 0.5,\n",
+    "        \"color\": [0.7, 1.0, 0.7],\n",
+    "        \"smooth_shading\": True,\n",
+    "        \"specular\": 0.4,\n",
+    "        \"specular_power\": 10,\n",
+    "        \"diffuse\": 0.9,\n",
+    "        \"ambient\": 1.0,\n",
+    "    }\n",
+    "\n",
+    "    def setup_view(pl, mesh, title, view_type, params):\n",
+    "        \"\"\"Helper function to set up consistent views\"\"\"\n",
+    "        pl.add_mesh(mesh, **params)\n",
+    "        pl.add_text(title, position=\"upper_edge\", font_size=10, color='black')\n",
+    "        pl.add_axes(line_width=2)\n",
+    "\n",
+    "        if view_type == \"xy\":\n",
+    "            pl.view_xy()\n",
+    "        elif view_type == \"xz\":\n",
+    "            pl.view_xz()\n",
+    "        elif view_type == \"yz\":\n",
+    "            pl.view_yz()\n",
+    "        else:  # isometric\n",
+    "            pl.view_isometric()\n",
+    "\n",
+    "        pl.reset_camera()\n",
+    "        pl.camera.zoom(0.85)\n",
+    "        # Add directional lighting\n",
+    "\n",
+    "    # Plot views for both meshes\n",
+    "    views = [(\"xy\", \"Top\"), (\"xz\", \"Front\"), (\"yz\", \"Side\"), (\"isometric\", \"Isometric\")]\n",
+    "\n",
+    "    # First mesh (top row)\n",
+    "    for col, (view_type, view_name) in enumerate(views):\n",
+    "        pl.subplot(0, col)\n",
+    "        setup_view(pl, mesh1, f\"{title1} - {view_name}\", view_type, params_case1)\n",
+    "\n",
+    "    # Second mesh (bottom row)\n",
+    "    for col, (view_type, view_name) in enumerate(views):\n",
+    "        pl.subplot(1, col)\n",
+    "        setup_view(pl, mesh2, f\"{title2} - {view_name}\", view_type, params_case2)\n",
+    "\n",
+    "    pl.background_color = '#f0f0f0'\n",
+    "\n",
+    "    # Display in notebook\n",
+    "    return pl.show(jupyter_backend='static')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Lets plot two geometries:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Define STL file paths\n",
+    "STL_FILE_1 = DATA_DIR / \"train_stl_files/case102.stl\"\n",
+    "STL_FILE_2 = DATA_DIR / \"train_stl_files/case116.stl\"\n",
+    "\n",
+    "# Load STL files\n",
+    "mesh1 = load_stl(STL_FILE_1)\n",
+    "mesh2 = load_stl(STL_FILE_2)\n",
+    "\n",
+    "# Print comparison\n",
+    "print(\"\\n=== STL File Comparison ===\\n\")\n",
+    "print(f\"File: {STL_FILE_1.name} vs {STL_FILE_2.name}\")\n",
+    "print(f\"Number of Faces: {mesh1.n_faces} vs {mesh2.n_faces}\")\n",
+    "print(f\"Surface Area: {mesh1.area:.2f} vs {mesh2.area:.2f}\")\n",
+    "print(f\"Volume: {mesh1.volume:.3f} vs {mesh2.volume:.3f} m³\")\n",
+    "\n",
+    "# Create visualization with volumes included\n",
+    "plot_stl_comparison(mesh1, mesh2, title1=\"Case 102\", title2=\"Case 116\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### **Step 4: Convert CFD Results to NPY Format**\n",
+    "\n",
+    "The fourth step in our workflow focuses on converting CFD simulation results into NumPy (.npy) format for efficient training. As previously mentioned, the advantages of using the NPY format include:\n",
+    "- NPY format provides faster data loading during training\n",
+    "- Enables efficient memory management for large datasets\n",
+    "- Facilitates parallel processing of simulation data\n",
+    "- Optimizes data access patterns for deep learning frameworks\n",
+    "\n",
+    "#### Understanding the Conversion Process\n",
+    "\n",
+    "The conversion process involves several key aspects:\n",
+    "1. Reading CFD simulation data from VTK files\n",
+    "2. Extracting surface and volume variables\n",
+    "3. Processing mesh geometry and physical quantities\n",
+    "4. Normalizing data using appropriate scaling factors\n",
+    "5. Saving processed data in NPY format\n",
+    "\n",
+    "#### Key Components and Libraries\n",
+    "\n",
+    "We'll use the following libraries for the conversion:\n",
+    "\n",
+    "1. **VTK and PyVista**\n",
+    "   - For reading CFD simulation data\n",
+    "   - For processing mesh geometry and surface properties\n",
+    "   - For computing cell sizes and normals\n",
+    "\n",
+    "2. **NumPy**\n",
+    "   - For efficient array operations\n",
+    "   - For saving data in NPY format\n",
+    "\n",
+    "3. **Concurrent Processing**\n",
+    "   - For parallel processing of multiple files\n",
+    "   - For improved conversion speed\n",
+    "\n",
+    "#### Important Data Processing Parameters\n",
+    "\n",
+    "During the conversion process, we need to consider:\n",
+    "- Surface variables (pressure, wall shear stress)\n",
+    "- Mean variables (mean wall shear stress, mean pressure)\n",
+    "- Data normalization using inlet velocity and density. **The inlet velocity will be read from the info files**\n",
+    "  ```python\n",
+    "          with open(info_path, \"r\") as file:\n",
+    "            velocity = next(float(line.split(\":\")[1].strip()) for line in file if \"Velocity\" in line)\n",
+    "  ```\n",
+    "- Mesh geometry preservation\n",
+    "- Memory efficiency for large datasets\n",
+    "\n",
+    "#### Implementation Overview\n",
+    "\n",
+    "The conversion is implemented through several key components:\n",
+    "\n",
+    "1. **Dataset Class**\n",
+    "```python\n",
+    "class OpenFoamAhmedBodySurfaceDataset(Dataset):\n",
+    "    \"\"\"Datapipe for converting OpenFOAM dataset to npy.\"\"\"\n",
+    "    # Handles data loading and processing\n",
+    "```\n",
+    "\n",
+    "2. **File Processing Functions**\n",
+    "```python\n",
+    "def process_file(fname: str, fm_data, output_path: str):\n",
+    "    \"\"\"Processes a single surface data file.\"\"\"\n",
+    "    # Converts individual files to NPY format\n",
+    "```\n",
+    "\n",
+    "3. **Batch Processing**\n",
+    "```python\n",
+    "def process_surface_data_batch(dataset_paths, info_paths, stl_paths, surface_paths):\n",
+    "    \"\"\"Processes all surface data files in parallel.\"\"\"\n",
+    "    # Handles parallel processing of multiple files\n",
+    "```\n",
+    "\n",
+    "Let's proceed with implementing these components and converting our CFD results:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class OpenFoamAhmedBodySurfaceDataset(Dataset):\n",
+    "    \"\"\"Datapipe for converting OpenFOAM dataset to npy.\"\"\"\n",
+    "\n",
+    "    def __init__(self, data_path: Union[str, Path], info_path: Union[str, Path], stl_path: Union[str, Path], surface_variables=None, volume_variables=None, device: int = 0):\n",
+    "        self.data_path = Path(data_path).expanduser()\n",
+    "        self.stl_path = Path(stl_path).expanduser()\n",
+    "        self.info_path = Path(info_path).expanduser()\n",
+    "        assert self.data_path.exists(), f\"Path {self.data_path} does not exist\"\n",
+    "\n",
+    "        self.filenames = get_filenames(self.data_path)\n",
+    "        random.shuffle(self.filenames)\n",
+    "        self.surface_variables = surface_variables or [\"p\", \"wallShearStress\"]\n",
+    "        self.volume_variables = volume_variables or [\"UMean\", \"pMean\"]\n",
+    "        self.device = device\n",
+    "\n",
+    "    def __len__(self):\n",
+    "        return len(self.filenames)\n",
+    "\n",
+    "    def __getitem__(self, idx):\n",
+    "        cfd_filename = self.filenames[idx]\n",
+    "        car_dir = self.data_path / cfd_filename\n",
+    "\n",
+    "        stl_path = self.stl_path / f\"{car_dir.stem}.stl\"\n",
+    "        info_path = self.info_path / f\"{car_dir.stem}_info.txt\"\n",
+    "\n",
+    "        with open(info_path, \"r\") as file:\n",
+    "            velocity = next(float(line.split(\":\")[1].strip()) for line in file if \"Velocity\" in line)\n",
+    "\n",
+    "        mesh_stl = pv.get_reader(stl_path).read()\n",
+    "        stl_faces = mesh_stl.faces.reshape(-1, 4)[:, 1:]\n",
+    "        stl_sizes = np.array(mesh_stl.compute_cell_sizes(length=False, area=True, volume=False).cell_data[\"Area\"])\n",
+    "\n",
+    "        reader = vtk.vtkXMLPolyDataReader()\n",
+    "        reader.SetFileName(str(car_dir))\n",
+    "        reader.Update()\n",
+    "        polydata = reader.GetOutput()\n",
+    "\n",
+    "        celldata = get_node_to_elem(polydata).GetCellData()\n",
+    "        surface_fields = np.concatenate(get_fields(celldata, self.surface_variables), axis=-1) / (AIR_DENSITY * velocity**2)\n",
+    "\n",
+    "        mesh = pv.PolyData(polydata)\n",
+    "        surface_sizes = np.array(mesh.compute_cell_sizes(length=False, area=True, volume=False).cell_data[\"Area\"])\n",
+    "        surface_normals = mesh.cell_normals / np.linalg.norm(mesh.cell_normals, axis=1)[:, np.newaxis]\n",
+    "\n",
+    "        return {\n",
+    "            \"stl_coordinates\": mesh_stl.points.astype(np.float32),\n",
+    "            \"stl_centers\": mesh_stl.cell_centers().points.astype(np.float32),\n",
+    "            \"stl_faces\": stl_faces.flatten().astype(np.float32),\n",
+    "            \"stl_areas\": stl_sizes.astype(np.float32),\n",
+    "            \"surface_mesh_centers\": mesh.cell_centers().points.astype(np.float32),\n",
+    "            \"surface_normals\": surface_normals.astype(np.float32),\n",
+    "            \"surface_areas\": surface_sizes.astype(np.float32),\n",
+    "            \"volume_fields\": None,\n",
+    "            \"volume_mesh_centers\": None,\n",
+    "            \"surface_fields\": surface_fields.astype(np.float32),\n",
+    "            \"filename\": cfd_filename,\n",
+    "            \"stream_velocity\": velocity,\n",
+    "            \"air_density\": AIR_DENSITY,\n",
+    "        }\n",
+    "\n",
+    "def process_file(fname: str, fm_data, output_path: str):\n",
+    "    \"\"\"Processes a single surface data file.\"\"\"\n",
+    "    full_path, output_file = os.path.join(fm_data.data_path, fname), os.path.join(output_path, f\"{fname}.npy\")\n",
+    "    if os.path.exists(output_file) or not os.path.exists(full_path) or os.path.getsize(full_path) == 0:\n",
+    "        return\n",
+    "    np.save(output_file, fm_data[fm_data.filenames.index(fname)])\n",
+    "\n",
+    "def process_surface_data_batch(dataset_paths: dict, info_paths: dict, stl_paths: dict, surface_paths: dict):\n",
+    "    \"\"\"Processes all surface data files in dataset_paths and saves them in output_paths.\"\"\"\n",
+    "\n",
+    "    for path in surface_paths.values(): os.makedirs(path, exist_ok=True)\n",
+    "\n",
+    "    print(\"=== Starting Processing ===\")\n",
+    "    for key, dataset_path in dataset_paths.items():\n",
+    "        surface_path = surface_paths[key]\n",
+    "        os.makedirs(surface_path, exist_ok=True)\n",
+    "        fm_data = OpenFoamAhmedBodySurfaceDataset(dataset_path, info_paths[key], stl_paths[key], VOLUME_VARS, SURFACE_VARS)\n",
+    "        file_list = [fname for fname in fm_data.filenames if fname.endswith(\".vtp\")]\n",
+    "\n",
+    "        print(f\"\\nProcessing {len(file_list)} files from {dataset_path} → {surface_path}...\")\n",
+    "\n",
+    "        with ProcessPoolExecutor() as executor:\n",
+    "            list(tqdm(executor.map(process_file, file_list, [fm_data]*len(file_list), [surface_path]*len(file_list)),\n",
+    "                      total=len(file_list), desc=f\"Processing {key}\", dynamic_ncols=True))\n",
+    "\n",
+    "    print(\"=== All Processing Completed Successfully ===\")\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Lets convert the files:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "process_surface_data_batch(dataset_paths, info_paths, stl_paths, surface_paths)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The files are converted to NPY format and saved in the surface_paths directory. We can now use this data to train a model, which is covered in the next notebook - [domino-training-test.ipynb](domino-training-test.ipynb)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "os._exit(00)"
+   ]
+  }
+ ],
+ "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.12.3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

domino-training-test.ipynb

@@ -0,0 +1,1343 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Physics NeMo External Aerodynamics DLI\n",
+    "\n",
+    "## Notebook 3 - Training DoMINO Model on the Ahmed body surface dataset and Running Inference\n",
+    "\n",
+    "In this notebook, we will first provide a detailed explanation of the DoMINO architecture, which is a multi-scale, iterative neural operator designed for modeling large-scale engineering simulations. We will break down the key components of DoMINO, including its use of local geometry representations, multi-scale point convolution kernels, and its efficient handling of complex geometries. Afterward, we will train the model using the **Ahmed body surface dataset**, a widely used dataset in automotive aerodynamics simulations. *As indicated in the previous notebook this dataset was created by the NVIDIA Physics NeMo development team and differs from other similar datasets hosted on cloud platforms like AWS.*\n",
+    "\n",
+    "*The DoMINO model is capable of training both volume fields (such as velocity and pressure) and surface fields (including pressure and wall shear stress). However, for the sake of simplicity and educational purposes, this notebook will *focus solely on training the surface fields* using the Ahmed body surface dataset.*\n",
+    "\n",
+    "## Guid line\n",
+    "- Before starting training, ensure GPU memory is cleared by running:\n",
+    "```python\n",
+    "import os\n",
+    "os._exit(00)\n",
+    "```\n",
+    "- The data preprocessing notebook should be run beforehand as a prerequisite to prepare the data in NPY format.\n",
+    "- The amount of GPU memory allocated during training can be controlled through the following configuration parameters:\n",
+    "```yaml\n",
+    "  interp_res=GRID_RESOLUTION # Resolution of the latent space; coarser resolutions reduce memory usage.\n",
+    "  NUM_SURFACE_NEIGHBORS = 7  # Number of neighboring surface points used to compute solution variables such as pressure. \n",
+    "  SURFACE_POINTS_SAMPLE = 8192  # Number of surface points sampled per epoch; fewer points reduce memory consumption\n",
+    "```\n",
+    "    To adjust GPU memory usage, modify the above configuration values in the \"Experiment Parameters and Variables\" section.\n",
+    "\n",
+    "\n",
+    "# Table of Contents\n",
+    "- [DoMINO Architecture](#Domino-Architecture)\n",
+    "  - [Geometric Understanding: Global and Local Perspectives](#Geometric-Understanding:-Global-and-Local-Perspectives)\n",
+    "  - [The Prediction Engine: Basis Functions and the Aggregation Network](#The-Prediction-Engine:-Basis-Functions-and-the-Aggregation-Network)\n",
+    "- [Training Process](#training)\n",
+    "  - [Step 1: Define Experiment Parameters and Dependencies](#step-1-define-experiment-parameters-and-dependencies)\n",
+    "    - [Loading Required Libraries](#loading-required-libraries)\n",
+    "    - [Dependencies](#dependencies)\n",
+    "    - [Experiment Parameters and Variables](#experiment-parameters-and-variables)\n",
+    "  - [Step 2: Train the DoMINO Model](#step-2-train-the-domino-model)\n",
+    "    - [Understanding the Training Process](#understanding-the-training-process)\n",
+    "    - [Key Components and Libraries](#key-components-and-libraries)\n",
+    "    - [Important Training Parameters](#important-training-parameters)\n",
+    "    - [Implementation Overview](#implementation-overview)\n",
+    "- [Load Model Checkpoint & Run Inference](#Load-Model-Checkpoint-&-Run-Inference)\n",
+    "- [Visualizing the predicted results](#Visualizing-the-predicted-results)\n",
+    "\n",
+    "\n",
+    "## DoMINO Architecture\n",
+    "\n",
+    "DoMINO, which stands for Decomposable Multiscale Iterative Neural Operator, is a novel machine learning model designed to address key challenges in accelerating simulations, such as accuracy, scalability, and generalization to new geometries—particularly in the context of automotive aerodynamics. As a neural operator, DoMINO is capable of predicting point-wise volume and surface fields and is inherently scalable to large domains. Its decomposable architecture is central to its performance: it learns local geometric representations within sub-regions of the domain, enabling greater accuracy by focusing on areas with detailed physical features. The model also employs a multiscale approach using learnable point kernels, allowing it to capture both fine and coarse geometric patterns directly from STL files.\n",
+    "\n",
+    "DoMINO operates iteratively to enable long-range interactions, progressively propagating the learned local geometric features across the entire computational domain. Separately, it builds dynamic, point-based computational stencils—drawing inspiration from traditional numerical methods—which it uses to learn non-linear basis functions tailored to the geometry. These basis functions, combined with localized geometry encodings, allow the model to predict volume and surface solution fields at specified points. A major advantage of DoMINO is that it only requires STL geometry as input during inference; it does not depend on mesh generation or the density and structure of input point clouds. The process begins by encoding global geometry onto a fixed grid using a combination of learnable point convolutions, CNNs, and dense layers. From this, the model constructs localized subdomains around evaluation points to extract rich geometric encodings, which are then used in conjunction with the learned basis functions for final prediction.\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### **Geometric Understanding: Global and Local Perspectives**\n",
+    "DoMINO's ability to accurately predict complex flow phenomena relies heavily on its sophisticated understanding of geometry, processed at both global and local scales.\n",
+    "\n",
+    "\n",
+    "- **A.** **Global Geometry Representation**:\n",
+    "The Global Geometry Representation refers to the overall shape and structure of the entire object or domain that you are modeling. This representation captures all the geometric details across the entire computational domain.\\\n",
+    "Step-by-Step Explanation of Global Geometry Representation:\n",
+    "\n",
+    "    - **Step 1**: Construct Bounding Boxes\n",
+    "    \t- A tight-fitting surface bounding box is created around the STL (3D geometry) to hold the geometry.\n",
+    "        - A computational domain bounding box is also defined, which is larger than the surface bounding box to encompass the whole computational domain.\n",
+    "        - Both bounding boxes can be specified in  ```conf.yaml```\n",
+    "    - **Step 2**: Project STL Vertices onto Structured Grid\n",
+    "    \t- The geometric features of the point cloud, such as spatial coordinates, are projected onto an N-dimensional structured grid of resolution m×m×m×f, which is overlaid on the surface bounding box using **learnable point convolution kernels**.\n",
+    "    \t- The learnable point convolution kernels are created using **differentiable ball query layers**. This means that the method:\n",
+    "    \t- Uses a \"ball\" (a sphere in 3D space) around each point to query or find its neighbors.\n",
+    "    \t- The ball query layer is \"differentiable,\" meaning it can be included in the neural network and updated via back propagation (i.e., during training, the network can learn how to adjust the kernels to improve performance).\n",
+    "    \t- The radius of the ball (radius of influence) defines how far around each point we look for neighboring points to include in the convolution. This defines, in fact, how far the geometry can affect the grid. A range of point convolutional kernel sizes can be learned by specifying several radii. Moreover, different kernels are learned to represent information on the surface bounding box and computational domain bounding box. This enables multi-scale learning of geometry encoding by representing both short- and long-range interactions of the surface and flow fields.  The radii of influence are defined as **list** in the ```conf.yaml``` file:\n",
+    "            ```yaml\n",
+    "            volume_radii: [0.1, 0.5]\n",
+    "            surface_radii: [0.05]\n",
+    "             ```\n",
+    "            These radius are used in the DoMINO model (```physicsnemo/models/domino/model.py```) to compute two **BQWarp** accordingly:\n",
+    "          \n",
+    "\n",
+    "     \n",
+    "            ```python\n",
+    "            class GeometryRep(nn.Module):\n",
+    "            \n",
+    "                def __init__(self, input_features: int, radii, model_parameters=None):\n",
+    "                    \"\"\"\n",
+    "                    Initialize the GeometryRep module.\n",
+    "            \n",
+    "                    Args:\n",
+    "                        input_features: Number of input feature dimensions\n",
+    "                        model_parameters: Configuration parameters for the model\n",
+    "                    \"\"\"\n",
+    "                    super().__init__()\n",
+    "                    geometry_rep = model_parameters.geometry_rep\n",
+    "                    self.geo_encoding_type = model_parameters.geometry_encoding_type\n",
+    "            \n",
+    "                    self.bq_warp = nn.ModuleList()\n",
+    "                    self.geo_processors = nn.ModuleList()\n",
+    "                    for j, p in enumerate(radii):\n",
+    "                        self.bq_warp.append(\n",
+    "                            BQWarp(\n",
+    "                                grid_resolution=model_parameters.interp_res,\n",
+    "                                radius=radii[j],\n",
+    "                            )\n",
+    "                        )\n",
+    "                        self.geo_processors.append(\n",
+    "                            GeoProcessor(\n",
+    "                                input_filters=geometry_rep.geo_conv.base_neurons_out,\n",
+    "                                model_parameters=geometry_rep.geo_processor,\n",
+    "                            )\n",
+    "                        )\n",
+    "            ```\n",
+    "    - **Step 3**: Use Multi-Resolution Approach for Detailed and Coarse Features\n",
+    "    \t- The grid resolution in the bounding box determines the level of detail of the geometry: \n",
+    "    \t- Finer resolution captures more detailed features of the geometry.\n",
+    "    \t- Coarser resolution captures larger, broader features.\n",
+    "    \t- A multi-resolution approach is adopted, meaning multiple grids at different resolutions (levels) are maintained to capture both fine and coarse features of the geometry. The number of resolution levels is a parameter that can be adjusted in conf.yaml file as \\\n",
+    "         ```yaml  GRID_RESOLUTION = [128, 64, 48]  # Resolution of the interpolation grid  ```\n",
+    "        - **Currently, the DoMINO model allows specification of a single resolution but this configuration will be provided in a future release.**\n",
+    "\n",
+    "    - **Step 4**: Propagate Geometry Features into the Computational Domain\n",
+    "        - The computational domain is much larger than the surface bounding box, so the geometry information needs to be extended.\n",
+    "    \t- Geometry features are propagated into the computational domain using two methods: \n",
+    "    \t- As explained in **step 2** Multi-scale **point convolution kernels** project the **geometry information** onto the surface bounding box (**see below the left figure**).\n",
+    "    \t- **Features** from the surface grid of the bounding box (i.e., Gs) are propagated into the computational domain grid (i.e., Gc) using **CNN blocks** that contain convolution, pooling, and unpooling layers (**see below the left figure**).\n",
+    "          As you see in the code snippet above first `BQWarp` is calculated and then  CNN blocks using `GeoProcessor`:\n",
+    "          ```python\n",
+    "          class GeoProcessor(nn.Module):\n",
+    "          ```\n",
+    "        - The **CNN blocks are iterated** for a specified number of steps to refine the geometry representation. Currently, the DoMINO model is configured to run a single iteration. An option to change this will be provided in ```conf.yaml``` in a future release.\n",
+    "\n",
+    "    - **Step 5**: Calculate Signed Distance Function (SDF) and its Gradients\n",
+    "    \t- Additionally, the Signed Distance Function (SDF) and its gradient components are calculated on the computational domain grid.\n",
+    "    \t- These SDF and gradient values are added to the learned features, providing additional information about the topology of the geometry (i.e., the geometry's shape, distances to surfaces, etc.).\n",
+    "\n",
+    "    - **Step 6**: Final Global Geometry Representation\n",
+    "    \t- The final geometry representation of the STL is formed by combining the learned features from the structured grids at different resolutions in both the bounding box and the computational domain.\n",
+    " \n",
+    "  \n",
+    "     \n",
+    "    **Once the computational domain is created, the next step would be local geometry representation.**\n",
+    "\n",
+    "\n",
+    "- **B.**  **Local geometry representation**\n",
+    "The Local Geometry Representation focuses on the geometry in the immediate vicinity of a sampled point p (the points in simulation mesh). The idea is to understand how the geometry behaves around a specific point and its neighbors, which can be important for accurate predictions. While the Global Geometry Representation gives the big picture, the Local Geometry Representation zooms in on a small region of interest around each sampled point. The key difference is that local geometry represents a smaller, more detailed portion of the global geometry, typically focusing on the small-scale features close to a point. For each sampled point p, neighboring points are sampled randomly around them to form a computational stencil of points similar to finite volume and element methods. The local geometry representation is learned by drawing a subregion around the computational stencil of\n",
+    "p + 1 points. The size of the subregion are defined as **list** in the ```conf.yaml``` file:\n",
+    "\n",
+    "  ```yaml\n",
+    "    geometry_local.volume_radii: [0.05, 0.1]\n",
+    "    geometry_local.surface_radii: [0.05]\n",
+    "  ```\n",
+    "Similar to Global geometry representation a point convolution kernel is used here to extract the local features in the subregion from the global geometry representation on the computational domain.\n",
+    "The **BQWarp** are computed for Local geometry representation in ```physicsnemo/models/domino/model.py``` in ```class DoMINO(nn.Module): ```:\n",
+    "\n",
+    "```python\n",
+    "        for ct, j in enumerate(self.surface_radius):\n",
+    "            if self.geo_encoding_type == \"both\":\n",
+    "                total_neighbors_in_radius = self.surface_neighbors_in_radius[ct] * (\n",
+    "                    len(model_parameters.geometry_rep.geo_conv.surface_radii) + 1\n",
+    "                )\n",
+    "            elif self.geo_encoding_type == \"stl\":\n",
+    "                total_neighbors_in_radius = self.surface_neighbors_in_radius[ct] * (\n",
+    "                    len(model_parameters.geometry_rep.geo_conv.surface_radii)\n",
+    "                )\n",
+    "            elif self.geo_encoding_type == \"sdf\":\n",
+    "                total_neighbors_in_radius = self.surface_neighbors_in_radius[ct]\n",
+    "\n",
+    "            self.surface_bq_warp.append(\n",
+    "                BQWarp(\n",
+    "                    grid_resolution=model_parameters.interp_res,\n",
+    "                    radius=self.surface_radius[ct],\n",
+    "                    neighbors_in_radius=self.surface_neighbors_in_radius[ct],\n",
+    "                )\n",
+    "            )\n",
+    "            self.surface_local_point_conv.append(\n",
+    "                LocalPointConv(\n",
+    "                    input_features=total_neighbors_in_radius,\n",
+    "                    base_layer=512,\n",
+    "                    output_features=self.surface_neighbors_in_radius[ct],\n",
+    "                )\n",
+    "            )\n",
+    "```\n",
+    "\n",
+    "Note that **BQWarp** for **Global geometry representation** are computed in ``` class GeometryRep(nn.Module): ```\n",
+    "\n",
+    "**How Does the Multi-Resolution Global Geometry Affect the Local Geometry Representation?**\n",
+    " - Coarse resolution: At the coarse resolution, you get a broad view of the object. This can give information about the general shape and large-scale features of the geometry (e.g., the overall shape of the object, major boundaries, etc.). When local geometry is extracted from the coarse resolution, the features are relatively less detailed, and it might capture larger, more general features of the object.\n",
+    " - Fine resolution: At the fine resolution, you get a detailed view of the geometry, capturing small features such as intricate surface details, small holes, or sharp edges. The local geometry representation derived from the fine resolution will be more detailed and capture smaller variations in the geometry near each sampled point.\n",
+    "\n",
+    "**Thus, the global multi-resolution geometry allows the local geometry to be learned at different levels of detail, depending on the resolution of the grid that is used to represent the geometry.** \n",
+    "\n",
+    "\n",
+    "### **The Prediction Engine: Basis Functions and the Aggregation Network**\n",
+    "Once DoMINO has processed the local geometry around each point of interest, the next stage involves taking these learned geometric features, along with other relevant information, and predicting the actual flow field values. This is accomplished by an \"Aggregation Network,\" which has a structure inspired by DeepONet.\n",
+    "\n",
+    "\n",
+    "- **A.** **Adding Spatial Context: The Role of Positional Encoding**\n",
+    "Neural networks, particularly simpler architectures like Multi-Layer Perceptrons (MLPs), often process sets of input features without an inherent understanding of their absolute or relative spatial positions. If a network receives only a list of coordinates, it might not easily distinguish whether a particular point is \"at the front\" or \"at the back\" of an object in a global sense, or how two distant points are spatially related if they are not immediate neighbors. To overcome this, Positional Encoding (PE) techniques are employed. PE injects information about the position of points into the feature vectors that the network processes (*see the right panel in the figure below*). In DoMINO, \"Position encoding is calculated between the coordinates of the sampled point and the center of mass of the geometry STL.\n",
+    "\n",
+    "- **B.** **Basis Function Neural Network (Latent Vector):**\n",
+    "    - Separately, DoMINO builds dynamic, point-based computational stencils—drawing inspiration from traditional numerical methods—which it uses to learn non-linear basis functions tailored to the geometry. These basis functions, combined with localized geometry encodings, allow the model to predict volume and surface solution fields at specified points (*see the right panel in the figure below*).\\\n",
+    "      What happens here:    \n",
+    "        - The input features (coordinates, SDF, normal vectors, etc. and their fourier features) for each point in the stencil are fed into the Basis Function Neural Network. This is a fully connected neural network that processes these features.\n",
+    "        - The network then computes a latent vector for each point in the stencil. A latent vector is a compressed mathematical representation that encodes the important information about each point’s geometry and position.\n",
+    "        - Purpose: The latent vector captures the essential characteristics of each point’s geometry in a compact form, which will be used in later steps for predicting the solution at that point.\n",
+    "\n",
+    "- **C.** **Concatenating the Latent Vector with the Local Geometry Encoding:**\n",
+    "    - After calculating the latent vector for each point, this vector is concatenated with the local geometry encoding (which includes the previously computed information from the surrounding points and the global geometry) and positional encoding.\n",
+    "    - Why this is done: Concatenating these two representations allows the network to use both the specific local features of each point and the broader context of the surrounding geometry to make predictions.\n",
+    "\n",
+    "- **D.** **Passing Through Additional Neural Network Layers (Solution Prediction):**\n",
+    "    - The combined information (latent vector + positional encoding + local geometry encoding) is passed through another set of fully connected layers (a new neural network).\n",
+    "    - What happens here: These layers process the combined information and predict a solution vector for each point in the stencil. The solution vector could represent various physical quantities such as temperature, pressure, or other simulation results at the sampled point.\n",
+    "    - Purpose: This step produces the predicted solution at each point, based on the local and global geometry.\n",
+    "\n",
+    "- **E.** **Aggregation Network:**\n",
+    "\n",
+    "    The core of the prediction engine in DoMINO is the **Aggregation network**. This is described as **a fully connected neural network with a DeepONet like structure.**  This network takes the processed local geometric features and the basis functions (derived from the Basis Function Neural Network and local geometry encodings) and combines them to compute the final solution field values.   \n",
+    "    To understand DoMINO's aggregation network, it's helpful to recall the DeepONet architecture. DeepONet is specifically designed for learning operators, mathematical mappings between function spaces. This makes it highly suitable for problems in physics described by PDEs. A DeepONet typically consists of two main sub-networks:   \n",
+    "\n",
+    "    - Branch Network: Processes the input function (e.g., PDE coefficients).    \n",
+    "    - Trunk Network: Processes the coordinates where the output function is evaluated.  The outputs of these two networks are then combined to produce the prediction.\n",
+    "\n",
+    "    In DoMINO, this DeepONet structure is adapted as follows:\n",
+    "\n",
+    "    - **Branch Net**: This part of the aggregation network takes the Local geometry rep as input. These are the η features, which encapsulate the detailed local geometric information around point i and its neighbors j.\n",
+    "    - **Trunk Net**: This part processes the basis functions. As described above, these \"basis functions\" are the output of the Basis Function Neural Network after concatenation with local geometry encodings and further MLP layers. They represent a learned, rich description of the query point i's properties, its positional encoding, and its context within its local stencil.    \n",
+    "    The DeepONet-like architecture allows DoMINO to effectively learn how different local geometric environments (processed by the branch net) influence the flow solution at various specific locations (whose context, encoded in the basis functions, is processed by the trunk net).\n",
+    "    The aggregation network, with its DeepONet-like structure, \"computes the solution field on the sampled point, i and its neighbors j.\" This means that for a given point of interest i, the network doesn't just predict the solution at i in isolation. Instead, it leverages the local stencil of points (point i itself and its defined neighbors j) to make predictions across this local cloud.\n",
+    "\n",
+    "    After the aggregation network produces these individual predictions for point i and its neighbors j, the solutions are then averaged using an inverse distance weighted interpolation (**IDW**) scheme.\n",
+    "    In DoMINO, the \"*known values*\"  $u_k$ are the predictions made by the aggregation network at point $i$ and its neighbors $j$. The IDW scheme then blends these predictions to yield the final solution at point $i$, giving more credence to the predictions made at or very near $i$. \n",
+    "    This IDW step can be seen as a form of learned solution refinement or a consensus mechanism. By predicting solutions across a local stencil and then averaging them with IDW, the model can produce a more robust and spatially consistent output at point $i$, potentially smoothing out minor errors or noise from individual network predictions within that stencil. \n",
+    "    This enforces a degree of local coherence in the predicted flow field, which is a desirable characteristic for physical simulations. IDW is often applied within a defined \"search neighborhood\" ; in DoMINO, this neighborhood is implicitly the set of points $j$ (and $i$ itself) for which the aggregation network computes initial solutions.   \n",
+    "\n",
+    "<div style=\"display: flex; justify-content: center; gap: 10px;\">\n",
+    "  <figure style=\"text-align: center;\">\n",
+    "    <img src=\"https://raw.githubusercontent.com/openhackathons-org/End-to-End-AI-for-Science/main/workspace/python/jupyter_notebook/DoMINO/images/global_geo_rep.png\" style=\"width: 100%; height: auto;\">\n",
+    "    <figcaption>Computation and surface Bounding box representation.</figcaption>\n",
+    "  </figure>\n",
+    "  <figure style=\"text-align: center;\">\n",
+    "    <img src=\"https://raw.githubusercontent.com/openhackathons-org/End-to-End-AI-for-Science/main/workspace/python/jupyter_notebook/DoMINO/images/aggregation_net.png\" style=\"width: 55%; height: auto;\">\n",
+    "    <figcaption>Aggregation network is a fully connected neural network with a DeepONet like structure, where Local geometry rep is branch net and basis functions are trunk net.</figcaption>\n",
+    "  </figure>\n",
+    "</div>\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## **Training**\n",
+    "### **Step 1: Define Experiment Parameters and Dependencies**\n",
+    "\n",
+    "The first step in training the DoMINO model on the Ahmed body dataset is to set up our experiment environment and define the necessary parameters. This includes specifying paths to our data, configuring training settings, and ensuring all required libraries are available.\n",
+    "\n",
+    "Key components we need to set up:\n",
+    "- Data paths for training and validation sets\n",
+    "- Model hyperparameters and training configurations\n",
+    "- Visualization settings for results\n",
+    "- Required Python libraries for mesh processing and deep learning\n",
+    "\n",
+    "#### Loading Required Libraries\n",
+    "\n",
+    "Before we proceed with the experiment setup, let's first import all the necessary libraries. These libraries will be used for:\n",
+    "- Deep learning and numerical computations (torch, numpy)\n",
+    "- Progress tracking and visualization (tqdm, matplotlib)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import time\n",
+    "import os\n",
+    "import re\n",
+    "import torch\n",
+    "import torchinfo\n",
+    "\n",
+    "\n",
+    "\n",
+    "import pyvista as pv\n",
+    "from tqdm import tqdm\n",
+    "from pathlib import Path\n",
+    "from types import SimpleNamespace\n",
+    "import matplotlib.pyplot as plt\n",
+    "from pathlib import Path\n",
+    "import apex\n",
+    "import numpy as np\n",
+    "import hydra\n",
+    "from hydra.utils import to_absolute_path\n",
+    "from omegaconf import DictConfig, OmegaConf\n",
+    "\n",
+    "from torch.cuda.amp import GradScaler, autocast\n",
+    "from torch.nn.parallel import DistributedDataParallel\n",
+    "from torch.utils.data import DataLoader\n",
+    "from torch.utils.data.distributed import DistributedSampler\n",
+    "from torch.utils.tensorboard import SummaryWriter\n",
+    "\n",
+    "from physicsnemo.distributed import DistributedManager\n",
+    "from physicsnemo.launch.utils import load_checkpoint, save_checkpoint\n",
+    "from physicsnemo.utils.sdf import signed_distance_field\n",
+    "\n",
+    "from physicsnemo.datapipes.cae.domino_datapipe import DoMINODataPipe\n",
+    "from physicsnemo.models.domino.model import DoMINO\n",
+    "from physicsnemo.utils.domino.utils import *"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Experiment Parameters and Variables\n",
+    "\n",
+    "In this section, we define all the necessary parameters and variables for our Ahmed body experiment. These parameters control various aspects of the training process, data processing, and model configuration.\n",
+    "\n",
+    "These parameters are carefully chosen based on:\n",
+    "- The physical dimensions of the Ahmed body\n",
+    "- The computational requirements of the DoMINO model\n",
+    "- The desired resolution for accurate flow prediction\n",
+    "- The available computational resources\n",
+    "- The specific requirements of the aerodynamic analysis\n",
+    "- `GEOMETRY_REP` contains the hyperparameters for the global geometry representation network.  \n",
+    "- `GEOMETRY_LOCAL` contains the hyperparameters for the local geometry representation.  \n",
+    "- As described in the theoretical section, the point convolution kernel relies on two additional factors: the radius of influence and the number of points included in the kernel. The radii of influence are specified as `volume_radii` and `surface_radii` within both `GEOMETRY_REP` and `GEOMETRY_LOCAL`.  \n",
+    "- The number of points within the kernel is defined as `volume_neighbors_in_radius=[128, 128]` and `surface_neighbors_in_radius=[128]` for the local geometry representation. For the global geometry representation, these values are not explicitly set in the `config.yaml` file, so the default value of `10` is used.  \n",
+    "- The **bounding box parameters** play a crucial role, as they define the computational domain for both volume and surface meshes, ensuring that all relevant flow features around the Ahmed body are accurately captured.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Directory and Path Configuration\n",
+    "EXPERIMENT_TAG = 4  # Unique identifier for this experiment run\n",
+    "PROJECT_NAME = \"ahmed_body_dataset\"  # Name of the project\n",
+    "OUTPUT_DIR = Path(f\"./outputs/{PROJECT_NAME}/{EXPERIMENT_TAG}\")  # Directory for experiment outputs\n",
+    "DATA_DIR = Path(\"/data/physicsnemo_ahmed_body_dataset_vv1/dataset\")  # Root directory for dataset\n",
+    "\n",
+    "CHECKPOINT_DIR = OUTPUT_DIR / \"models\"  # Directory for saving model checkpoints\n",
+    "SAVE_PATH = DATA_DIR / \"mesh_predictions_surf_final1\" # path to save prediction results\n",
+    "\n",
+    "# Ensure directories exist\n",
+    "\n",
+    "os.makedirs(CHECKPOINT_DIR, exist_ok=True)\n",
+    "\n",
+    "# Physical Variables\n",
+    "VOLUME_VARS = [\"p\"]  # Volume variables to predict (pressure)\n",
+    "SURFACE_VARS = [\"p\", \"wallShearStress\"]  # Surface variables to predict\n",
+    "MODEL_TYPE = \"surface\"  # Type of model (surface-only prediction)\n",
+    "AIR_DENSITY = 1.205  # Air density in kg/m³\n",
+    "\n",
+    "# Training Hyperparameters\n",
+    "NUM_EPOCHS = 3  # Number of training epochs\n",
+    "LR = 0.001  # Learning rate\n",
+    "BATCH_SIZE = 1  # Batch size for training\n",
+    "GRID_RESOLUTION = [128, 64, 48]  # Resolution of the interpolation grid\n",
+    "SURFACE_POINTS_SAMPLE = 8192\n",
+    "GEOMETRY_REP=SimpleNamespace(\n",
+    "    geo_conv=SimpleNamespace(base_neurons=32, base_neurons_out=1, volume_radii=[0.1, 0.5], surface_radii=[0.05], hops=1),\n",
+    "    geo_processor=SimpleNamespace(base_filters=8),\n",
+    "    geo_processor_sdf=SimpleNamespace(base_filters=8)\n",
+    ") # Hyperparameters for global geometry representation network\n",
+    "\n",
+    "GEOMETRY_LOCAL=SimpleNamespace(volume_neighbors_in_radius=[128, 128], surface_neighbors_in_radius=[128], volume_radii=[0.05, 0.1], surface_radii=[0.05], base_layer=512) # Hyperparameters for local geometry extraction \n",
+    "\n",
+    "NUM_SURFACE_NEIGHBORS = 7  # Number of neighbors for surface operations\n",
+    "NORMALIZATION = \"min_max_scaling\"  # Data normalization method\n",
+    "INTEGRAL_LOSS_SCALING = 0  # Scaling factor for integral loss\n",
+    "GEOMETRY_ENCODING_TYPE= \"both\" # geometry encoder type, sdf, stl, both\n",
+    "NUM_SURF_VARS = 4  # Number of surface variables to predict, 3 for vectore (wallShearStress) and 1 for scalar (p)\n",
+    "CHECKPOINT_INTERVAL = 1  # Save checkpoint every N epochs\n",
+    "\n",
+    "\n",
+    "# Dataset Paths\n",
+    "DATA_PATHS = {\n",
+    "    \"train\": f\"{DATA_DIR}/train_prepared_surface_data/\",\n",
+    "    \"val\": f\"{DATA_DIR}/validation_prepared_surface_data/\",\n",
+    "    \"test\": f\"{DATA_DIR}/test\"\n",
+    "}\n",
+    "\n",
+    "# Model and Scaling Factor Paths\n",
+    "MODEL_SAVE_DIR = \"./outputs/ahmed_body_dataset/4/models\"\n",
+    "SURF_SAVE_PATH = './outputs/ahmed_body_dataset/surface_scaling_factors.npy'\n",
+    "\n",
+    "# Bounding Box Configuration for Volume and Surface Meshes\n",
+    "BOUNDING_BOX = SimpleNamespace(\n",
+    "    max=[0.5, 0.6, 0.6],  # Maximum coordinates for volume mesh\n",
+    "    min=[-2.5, -0.5, -0.5] # Minimum coordinates for volume mesh\n",
+    ")\n",
+    "BOUNDING_BOX_SURF = SimpleNamespace(\n",
+    "    max=[0.01, 0.6, 0.4],  # Maximum coordinates for surface mesh\n",
+    "    min=[-1.6, -0.01, -0.01] # Minimum coordinates for surface mesh\n",
+    ")\n",
+    "\n",
+    "# Set cuDNN benchmark mode\n",
+    "torch.backends.cudnn.benchmark = True\n",
+    "torch.backends.cudnn.deterministic = False"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### **Step 2: Train the DoMINO Model**\n",
+    "\n",
+    "The fifth step in our workflow focuses on training the DoMINO model on our processed CFD data. This step is crucial because:\n",
+    "- It enables the model to learn complex fluid dynamics patterns\n",
+    "- Provides a foundation for accurate flow field predictions\n",
+    "- Allows for efficient inference on new geometries\n",
+    "- Supports distributed training for improved performance\n",
+    "\n",
+    "#### Understanding the Training Process\n",
+    "\n",
+    "The training process involves several key components:\n",
+    "1. Setting up distributed training environment\n",
+    "2. Creating and configuring datasets and dataloaders\n",
+    "3. Initializing the DoMINO model architecture\n",
+    "4. Implementing training and validation loops\n",
+    "5. Managing model checkpoints and metrics\n",
+    "\n",
+    "#### Key Components and Libraries\n",
+    "\n",
+    "We'll use the following for training:\n",
+    "\n",
+    "- **PyTorch**\n",
+    "   - `torch.distributed`: For distributed training\n",
+    "   - `torch.cuda`: For GPU acceleration\n",
+    "   - `torch.optim`: For optimization algorithms\n",
+    "\n",
+    "- **Data Management**\n",
+    "   - Custom dataset classes for CFD data\n",
+    "   - Distributed samplers for efficient data loading\n",
+    "   - Distributed samplers for efficient data loading\n",
+    "\n",
+    "#### Important Training Parameters\n",
+    "\n",
+    "During the training process, we need to consider:\n",
+    "- Batch size and learning rate\n",
+    "- Number of epochs and validation frequency\n",
+    "- Model architecture parameters\n",
+    "- Loss function configuration\n",
+    "- Checkpointing strategy\n",
+    "\n",
+    "#### Implementation Overview\n",
+    "\n",
+    "The training is implemented through several key components:\n",
+    "\n",
+    "1. **Model Creation**\n",
+    "Creates the DoMINO model, applies configuration, and wraps it with DistributedDataParallel if training is distributed.\n",
+    "DoMINO model is initialized with full configuration. the model is moved to the given device (usually a CUDA GPU).\n",
+    "If training in distributed mode, wrap the model with DistributedDataParallel to sync gradients across processes.\n",
+    "   \n",
+    "```python\n",
+    "def create_model(device, rank, world_size):\n",
+    "    \"\"\"Create and configure DoMINO model.\"\"\"\n",
+    "    # Initializes model with specified parameters\n",
+    "```\n",
+    "2. **Create DoMINO dataset**\n",
+    "Constructs and returns a DoMINODataPipe object for a specified phase (`train` or `val`), fully configured for loading data.\n",
+    "Loads path-specific and global configs (grid resolution, surface vars, encoding, etc.).\n",
+    "\n",
+    "```python\n",
+    "def create_dataset(phase):\n",
+    "    \"\"\"\n",
+    "    Create DoMINO dataset for specified phase (train/val).\n",
+    "    \n",
+    "    Args:\n",
+    "        phase (str): Dataset phase ('train' or 'val')\n",
+    "    \n",
+    "    Returns:\n",
+    "        DoMINODataPipe: Configured dataset\n",
+    "    \"\"\"\n",
+    "```\n",
+    "\n",
+    "3. **Training Loop**\n",
+    "    Main training loop for running multiple epochs. Handles distributed settings, optimizer/scaler setup, and calls `run_epoch()` for training and evaluation.\n",
+    "    It gets dataloaders and samplers for train/val datasets. The optimizer (FusedAdam) and gradient scaler will be set up.\n",
+    "   \n",
+    "    For each epoch:\n",
+    "   - Update distributed sampler epoch (important for shuffling).\n",
+    "   - Call run_epoch() to train and validate.\n",
+    "   - Update best validation loss.\n",
+    "   \n",
+    "```python\n",
+    "def train(model, device, rank, world_size):\n",
+    "    \"\"\"Orchestrates the training process.\"\"\"\n",
+    "    # Handles training loop, validation, and checkpointing\n",
+    "```\n",
+    "\n",
+    "4. **Runs one training epoch run_epoch()**\n",
+    "Runs one training epoch, performs forward and backward passes, computes losses, and evaluates on validation data. Supports mixed precision and distributed training.\n",
+    "   \n",
+    "Step-by-step:\n",
+    "- Training phase:\n",
+    "    - Set model to training mode.\n",
+    "    - Use tqdm progress bar if rank is 0 (main process).\n",
+    "    - For each batch:\n",
+    "       - Move it to the correct device.\n",
+    "       - Run model and compute predictions.\n",
+    "       - Compute masked MSE loss using mse_loss_fn.\n",
+    "       - Apply gradient scaling for mixed-precision training.\n",
+    "       - Step the optimizer and clear gradients.\n",
+    "       - Log training loss (on rank 0).\n",
+    "- Validation phase:\n",
+    "    - Run inference without gradient tracking.\n",
+    "    - Average validation loss over the val loader.\n",
+    "\n",
+    "- Checkpointing:\n",
+    "   - Save best model if current val loss is better.\n",
+    "   - Save periodic checkpoint if epoch meets interval.\n",
+    "   - Return validation loss for tracking best model.\n",
+    "\n",
+    "Let's proceed with implementing these components and training our model:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def mse_loss_fn(output, target, padded_value=-10):\n",
+    "    \"\"\"\n",
+    "    Compute masked MSE loss, ignoring padded values.\n",
+    "    \n",
+    "    Args:\n",
+    "        output (torch.Tensor): Model predictions\n",
+    "        target (torch.Tensor): Ground truth values\n",
+    "        padded_value (float): Value used for padding (default: -10)\n",
+    "    \n",
+    "    Returns:\n",
+    "        torch.Tensor: Mean squared error loss\n",
+    "    \"\"\"\n",
+    "    # Move target to same device as output\n",
+    "    target = target.to(output.device)\n",
+    "    # Create mask for non-padded values\n",
+    "    mask = torch.abs(target - padded_value) > 1e-3\n",
+    "    # Compute masked loss\n",
+    "    masked_loss = torch.sum(((output - target) ** 2) * mask) / torch.sum(mask)\n",
+    "    return masked_loss.mean()\n",
+    "\n",
+    "def create_dataset(phase):\n",
+    "    \"\"\"\n",
+    "    Create DoMINO dataset for specified phase (train/val).\n",
+    "    \n",
+    "    Args:\n",
+    "        phase (str): Dataset phase ('train' or 'val')\n",
+    "    \n",
+    "    Returns:\n",
+    "        DoMINODataPipe: Configured dataset\n",
+    "    \"\"\"\n",
+    "    return DoMINODataPipe(\n",
+    "        DATA_PATHS[phase],\n",
+    "        phase=phase,\n",
+    "        grid_resolution=GRID_RESOLUTION,\n",
+    "        surface_variables=SURFACE_VARS,\n",
+    "        normalize_coordinates=True,\n",
+    "        sampling=True,\n",
+    "        sample_in_bbox=True,\n",
+    "        volume_points_sample=8192,\n",
+    "        surface_points_sample=SURFACE_POINTS_SAMPLE,\n",
+    "        geom_points_sample=60000,\n",
+    "        positional_encoding=False,\n",
+    "        surface_factors=np.load(SURF_SAVE_PATH),\n",
+    "        scaling_type=NORMALIZATION,\n",
+    "        model_type=MODEL_TYPE,\n",
+    "        bounding_box_dims=BOUNDING_BOX,\n",
+    "        bounding_box_dims_surf=BOUNDING_BOX_SURF,\n",
+    "        num_surface_neighbors=NUM_SURFACE_NEIGHBORS,\n",
+    "        gpu_preprocessing=False\n",
+    "        )\n",
+    "\n",
+    "\n",
+    "\n",
+    "def create_dataloaders(rank, world_size):\n",
+    "    \"\"\"\n",
+    "    Create train and validation dataloaders with distributed sampling.\n",
+    "    \n",
+    "    Args:\n",
+    "        rank (int): Process rank\n",
+    "        world_size (int): Total number of processes\n",
+    "    \n",
+    "    Returns:\n",
+    "        tuple: (train_loader, val_loader, train_sampler, val_sampler)\n",
+    "    \"\"\"\n",
+    "    # Create datasets\n",
+    "    train_dataset, val_dataset = create_dataset(\"train\"), create_dataset(\"val\")\n",
+    "    \n",
+    "    # Configure distributed samplers if needed\n",
+    "    train_sampler = DistributedSampler(train_dataset, num_replicas=world_size, rank=rank) if world_size > 1 else None\n",
+    "    val_sampler = DistributedSampler(val_dataset, num_replicas=world_size, rank=rank) if world_size > 1 else None\n",
+    "    \n",
+    "    # Create dataloaders\n",
+    "    return (\n",
+    "        DataLoader(train_dataset, batch_size=BATCH_SIZE, sampler=train_sampler, shuffle=train_sampler is None),\n",
+    "        DataLoader(val_dataset, batch_size=BATCH_SIZE, sampler=val_sampler, shuffle=False),\n",
+    "        train_sampler, val_sampler\n",
+    "    )\n",
+    "\n",
+    "def create_model(device, rank, world_size):\n",
+    "    \"\"\"\n",
+    "    Create and configure DoMINO model with distributed training support.\n",
+    "    \n",
+    "    Args:\n",
+    "        device (torch.device): Computation device\n",
+    "        rank (int): Process rank\n",
+    "        world_size (int): Total number of processes\n",
+    "    \n",
+    "    Returns:\n",
+    "        DoMINO: Configured model (wrapped in DistributedDataParallel if world_size > 1)\n",
+    "    \"\"\"\n",
+    "\n",
+    "  \n",
+    "    # Initialize model with configuration\n",
+    "    model = DoMINO(\n",
+    "        input_features=3,\n",
+    "        output_features_vol=None,\n",
+    "        output_features_surf=NUM_SURF_VARS,\n",
+    "        model_parameters=SimpleNamespace(\n",
+    "            interp_res=GRID_RESOLUTION,\n",
+    "            surface_neighbors=NUM_SURFACE_NEIGHBORS,\n",
+    "            use_surface_normals=True,\n",
+    "            use_surface_area=True,\n",
+    "            encode_parameters=True,\n",
+    "            positional_encoding=False,\n",
+    "            integral_loss_scaling_factor=INTEGRAL_LOSS_SCALING,\n",
+    "            normalization=NORMALIZATION,\n",
+    "            use_sdf_in_basis_func=True,\n",
+    "            geometry_encoding_type= GEOMETRY_ENCODING_TYPE, # geometry encoder type, sdf, stl, both\n",
+    "            geometry_rep=GEOMETRY_REP,\n",
+    "            nn_basis_functions=SimpleNamespace(base_layer=512, fourier_features=False, num_modes=5),\n",
+    "            parameter_model=SimpleNamespace(base_layer=512, scaling_params=[30.0, 1.226], fourier_features=False, num_modes=5),\n",
+    "            position_encoder=SimpleNamespace(base_neurons=512),\n",
+    "            geometry_local=GEOMETRY_LOCAL,\n",
+    "            aggregation_model=SimpleNamespace(base_layer=512),\n",
+    "            model_type=MODEL_TYPE\n",
+    "        ),\n",
+    "    ).to(device)\n",
+    "    \n",
+    "    # Wrap model for distributed training if needed\n",
+    "    if world_size > 1:\n",
+    "        model = DistributedDataParallel(\n",
+    "            model, \n",
+    "            device_ids=[rank], \n",
+    "            output_device=rank, \n",
+    "            find_unused_parameters=True\n",
+    "        )\n",
+    "    \n",
+    "    return model\n",
+    "\n",
+    "def run_epoch(train_loader, val_loader, model, optimizer, scaler, device, epoch, best_vloss, rank, world_size):\n",
+    "    \"\"\"\n",
+    "    Run one training epoch with validation.\n",
+    "    \n",
+    "    Args:\n",
+    "        train_loader (DataLoader): Training data loader\n",
+    "        val_loader (DataLoader): Validation data loader\n",
+    "        model (DoMINO): Model to train\n",
+    "        optimizer (torch.optim.Optimizer): Optimizer\n",
+    "        scaler (GradScaler): Gradient scaler for mixed precision\n",
+    "        device (torch.device): Computation device\n",
+    "        epoch (int): Current epoch number\n",
+    "        best_vloss (float): Best validation loss so far\n",
+    "        rank (int): Process rank\n",
+    "        world_size (int): Total number of processes\n",
+    "    \n",
+    "    Returns:\n",
+    "        float: Validation loss for this epoch\n",
+    "    \"\"\"\n",
+    "    # Training phase\n",
+    "    model.train()\n",
+    "    train_loss = 0.0\n",
+    "    pbar = tqdm(train_loader, desc=f\"Epoch {epoch+1}/{NUM_EPOCHS}\") if rank == 0 else train_loader\n",
+    "    \n",
+    "    for batch in pbar:\n",
+    "        # Move batch to device\n",
+    "        batch = dict_to_device(batch, device)\n",
+    "        \n",
+    "        # Forward pass with mixed precision\n",
+    "        with autocast():\n",
+    "            _, pred_surf = model(batch)\n",
+    "            loss = mse_loss_fn(pred_surf, batch[\"surface_fields\"])\n",
+    "        \n",
+    "        # Backward pass with gradient scaling\n",
+    "        scaler.scale(loss).backward()\n",
+    "        scaler.step(optimizer)\n",
+    "        scaler.update()\n",
+    "        optimizer.zero_grad()\n",
+    "        \n",
+    "        # Update loss tracking\n",
+    "        train_loss += loss.item()\n",
+    "        if rank == 0:\n",
+    "            pbar.set_postfix({\n",
+    "                \"train_loss\": f\"{train_loss/(pbar.n+1):.5e}\", \n",
+    "                \"lr\": f\"{optimizer.param_groups[0]['lr']:.2e}\"\n",
+    "            })\n",
+    "    \n",
+    "    # Compute average training loss\n",
+    "    avg_train_loss = train_loss / len(train_loader)\n",
+    "    \n",
+    "    # Validation phase\n",
+    "    model.eval()\n",
+    "    with torch.no_grad():\n",
+    "        val_loss = sum(\n",
+    "            mse_loss_fn(model(dict_to_device(batch, device))[1], batch[\"surface_fields\"].to(device)).item() \n",
+    "            for batch in val_loader\n",
+    "        ) / len(val_loader)\n",
+    "    \n",
+    "    # Handle distributed training metrics\n",
+    "    if world_size > 1:\n",
+    "        avg_train_loss, val_loss = [torch.tensor(v, device=device) for v in [avg_train_loss, val_loss]]\n",
+    "        torch.distributed.all_reduce(avg_train_loss, op=torch.distributed.ReduceOp.SUM)\n",
+    "        torch.distributed.all_reduce(val_loss, op=torch.distributed.ReduceOp.SUM)\n",
+    "        avg_train_loss, val_loss = avg_train_loss.item() / world_size, val_loss.item() / world_size\n",
+    "    \n",
+    "    # Save checkpoints on main process\n",
+    "    if rank == 0:\n",
+    "        if val_loss < best_vloss:\n",
+    "            save_checkpoint(\n",
+    "                os.path.join(MODEL_SAVE_DIR, \"best_model\"), \n",
+    "                models=model,\n",
+    "                optimizer=optimizer,\n",
+    "                scaler=scaler\n",
+    "            )\n",
+    "\n",
+    "        if (epoch + 1) % CHECKPOINT_INTERVAL == 0:\n",
+    "            save_checkpoint(\n",
+    "                MODEL_SAVE_DIR, \n",
+    "                models=model, \n",
+    "                optimizer=optimizer, \n",
+    "                scaler=scaler, \n",
+    "                epoch=epoch\n",
+    "            )\n",
+    "    \n",
+    "    return val_loss\n",
+    "\n",
+    "def train(model, device, rank, world_size):\n",
+    "    \"\"\"\n",
+    "    Function that orchestrates the training process.\n",
+    "    Handles distributed training setup, model creation and training loop.\n",
+    "    \"\"\"\n",
+    "\n",
+    "\n",
+    "    # Create output directory on main process\n",
+    "    os.makedirs(MODEL_SAVE_DIR, exist_ok=True) if rank == 0 else None\n",
+    "    \n",
+    "    # Set up data\n",
+    "    train_loader, val_loader, train_sampler, val_sampler = create_dataloaders(rank, world_size)\n",
+    "\n",
+    "    optimizer = apex.optimizers.FusedAdam(model.parameters(), lr=0.001)\n",
+    "    \n",
+    "    # Initialize learning rate scheduler and gradient scaler\n",
+    "    #scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[1, 2], gamma=0.5)\n",
+    "    scaler = GradScaler()\n",
+    "    \n",
+    "    # Training loop\n",
+    "    best_vloss = float('inf')\n",
+    "    for epoch in range(NUM_EPOCHS):\n",
+    "        if world_size > 1:\n",
+    "            train_sampler.set_epoch(epoch)\n",
+    "        best_vloss = min(\n",
+    "            best_vloss, \n",
+    "            run_epoch(\n",
+    "                train_loader, val_loader, model, optimizer, \n",
+    "                scaler, device, epoch, best_vloss, rank, world_size\n",
+    "            )\n",
+    "        )\n",
+    "        #scheduler.step()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Lets run the train for few epochs:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Initialize distributed training\n",
+    "os.environ[\"RANK\"] = f\"0\"\n",
+    "os.environ[\"WORLD_SIZE\"] = f\"1\"\n",
+    "os.environ[\"MASTER_ADDR\"] = \"localhost\"\n",
+    "os.environ[\"MASTER_PORT\"] = str(12355)\n",
+    "os.environ[\"LOCAL_RANK\"] = f\"0\"\n",
+    "\n",
+    "\n",
+    "DistributedManager.initialize()\n",
+    "dist = DistributedManager()\n",
+    "device=dist.device\n",
+    "rank=dist.rank\n",
+    "world_size=dist.world_size\n",
+    "print(device)\n",
+    "# Set up model\n",
+    "model = create_model(device, rank, world_size)\n",
+    "# Run training\n",
+    "train(model, device, rank, world_size)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## **Load Model Checkpoint & Run Inference**\n",
+    "\n",
+    "The sixth step in our workflow focuses on evaluating our trained DoMINO model by loading the best checkpoint and running inference on sample cases. \n",
+    "To run the inference, the script needs several key **inputs**. It requires the 3D shape of the object defined in an STL file and a corresponding surface mesh provided as a VTP file, which also contain results from a traditional simulation for comparison purposes. \n",
+    "Crucially, it needs the pre-trained DoMINO AI model loaded from a checkpoint file. Additionally, basic flow conditions like air speed (`STREAM_VELOCITY`) and density (`AIR_DENSITY`), along with specific scaling factors saved from the training phase (used to convert model outputs to physical values), must be provided.\n",
+    "\n",
+    "\n",
+    "As its main **output**, the script generates new VTP files for each tested geometry. These files include the original surface mesh data but are augmented with new data fields representing the AI model's predictions for aerodynamic quantities such as surface pressure and wall shear stress. Furthermore, the script calculates aerodynamic forces based on these predictions and prints a comparison against forces derived from reference data directly to the console.\n",
+    "The code snippet below takes geometry files (STL) and corresponding simulation setup data (partially from VTP files and config parameters), preprocesses them, feeds them into the model to predict aerodynamic quantities (like surface pressure and shear stress), and saves these predictions back into VTP files for analysis and visualization.\n",
+    "\n",
+    "#### Understanding the Testing Process\n",
+    "\n",
+    "The testing process involves several key components:\n",
+    "1. Loading the best model checkpoint\n",
+    "2. Preparing test data\n",
+    "3. Running inference on test cases\n",
+    "4. Analyzing prediction results\n",
+    "5. Comparing with ground truth values\n",
+    "\n",
+    "#### Key Components and Libraries\n",
+    "\n",
+    "We'll use the following libraries for testing:\n",
+    "\n",
+    "1. **PyTorch**\n",
+    "   - `torch.load()`: For loading model checkpoints\n",
+    "   - `model.load_state_dict()`: For restoring model weights\n",
+    "   - `torch.no_grad()`: For efficient inference\n",
+    "\n",
+    "2. **Custom Testing Functions**\n",
+    "   - `test_step()`: For running inference on test cases\n",
+    "   - Data processing utilities for test data preparation\n",
+    "\n",
+    "#### Implementation Overview\n",
+    "\n",
+    "The testing is implemented through several key components:\n",
+    "\n",
+    "1. **Function test_step**\n",
+    "    Within the **test_step** function, several key operations execute in sequence. Initially, torch.no_grad() is used to disable gradient tracking in PyTorch, optimizing performance by saving memory and computation time as gradients are unnecessary during inference. \n",
+    "Next, the necessary data is prepared by extracting inputs like air density, stream velocity, geometry coordinates, bounding box grid information (surf_grid), and the Signed Distance Field (SDF) from the data_dict; the SDF is particularly important as it helps the model understand the position of points relative to the geometry surface. \n",
+    "Following this, a global geometry encoding is generated using model.geo_rep_surface, which takes normalized geometry points, the grid, and the SDF to create a comprehensive representation of the overall shape. \n",
+    "Surface-specific data, including mesh points, normals, areas, and neighbor details (found via methods like KDTree during preprocessing), are then extracted. \n",
+    "To refine the focus, model.geo_encoding_local_surface extracts relevant local geometric features from the global encoding specifically for the surface points where predictions are needed. \n",
+    "Positional awareness is added using model.position_encoder to encode the relative location of surface points. \n",
+    "The core prediction then occurs via model.calculate_solution_with_neighbors, combining local geometry, positional encoding, surface point details, neighbor information, and flow conditions to estimate the target surface fields like pressure coefficient or wall shear stress. Since the model output is normalized, a final un-normalization step converts these predictions back into physical units using the provided surf_factors, stream velocity, and air density. The function concludes by returning the predicted surface fields (prediction_surf), as this specific code path concentrates only on surface predictions.\n",
+    "\n",
+    "\n",
+    "\n",
+    "```python\n",
+    "def test_step(model, data_dict, surf_factors, device):\n",
+    "    \"\"\"\n",
+    "    Executes the core inference logic for a single test case using the trained DoMINO model.\n",
+    "    \n",
+    "    Args:\n",
+    "        model (DoMINO): The trained model\n",
+    "        data_dict: A dictionary containing all necessary input data for this specific test case (geometry, mesh points, flow conditions, etc.), already preprocessed and formatted.\n",
+    "        surf_factors: Scaling factors used during training to normalize the target surface data. Needed here to un-normalize the model's predictions back to physical values.\n",
+    "        device: The computational device (CPU or GPU) to run the calculations on.\n",
+    "        \n",
+    "    Returns:\n",
+    "        tuple: (prediction_vol, prediction_surf) - Model predictions for volume and surface\n",
+    "    \"\"\"\n",
+    "```\n",
+    "\n",
+    "\n",
+    "2. **Function: test**\n",
+    "```python\n",
+    "def test(model, test_dataloader, device):\n",
+    "    \"\"\"\n",
+    "    Run testing on the model using the provided test dataloader.\n",
+    "    \n",
+    "    Args:\n",
+    "        model (DoMINO): The trained model\n",
+    "        test_dataloader (DataLoader): DataLoader containing test data\n",
+    "        device (torch.device): Device to run inference on\n",
+    "        \n",
+    "    Returns:\n",
+    "        list: List of tuples containing (prediction_vol, prediction_surf) for each test case\n",
+    "    \"\"\"\n",
+    "```\n",
+    "\n",
+    "On the other hand, the test function is a higher-level function that organizes and controls the overall testing process. It begins by checking if surface scaling factors have been pre-computed and stored in a .npy file. If the file exists, it loads these factors; if not, it defaults to None. \n",
+    "\n",
+    "**The function then loads a pre-trained model from a checkpoint file (DoMINO.0.....pt) and loads its state into the model**. \n",
+    "```python\n",
+    "    # Load the best model checkpoint\n",
+    "    best_checkpoint = torch.load(CHECKPOINT_DIR / \"best_model/DoMINO.0.401.pt\")\n",
+    "    model.load_state_dict(best_checkpoint)  # Load the model state\n",
+    "    print(\"Model loaded\")\n",
+    "```\n",
+    "\n",
+    "After the model is loaded, it creates a directory for saving predictions if it doesn't already exist. The dirname parameter is used to extract a tag, which helps identify the current test case.\n",
+    "\n",
+    "Next, the function proceeds to load the necessary input files. It reads an STL file that contains the 3D geometry of the surface and extracts relevant data like vertices, faces, and areas. The bounding box dimensions are calculated, and the surface’s center of mass is computed. Then, it prepares a grid (surf_grid) and calculates the signed distance function (SDF) over this grid using the surface geometry, which helps in understanding the geometry’s proximity to the grid points. The function then reads the VTP file, which holds additional surface-related data such as pressure and shear force values.\n",
+    "\n",
+    "The surface fields are then prepared by interpolating the surface mesh data and its corresponding attributes. These fields are normalized to fit within the bounding box dimensions. The data dictionary is assembled, containing all the relevant inputs needed for the model’s prediction. This dictionary includes things like normalized surface coordinates, surface areas, and field values such as stream velocity and air density. The dictionary is converted to PyTorch tensors, making it compatible with the model.\n",
+    "\n",
+    "The test_step function is then called with this prepared data to compute the model's predictions. After the predictions are generated, the function compares the predicted surface forces (pressure and shear stress) with the true values from the surface fields. It calculates the predicted forces and prints out the comparison between the predicted and true values. The predicted surface fields are then converted to VTK format and saved to a file. Finally, the function finishes by returning, completing the testing process. This function provides a complete pipeline for testing a trained model on surface data, generating predictions, and saving them for further analysis.\n",
+    "\n",
+    "\n",
+    "Let's proceed with loading our trained model and running the tests:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def test_step(model, data_dict, surf_factors, device):\n",
+    "    \"\"\"\n",
+    "    Run a single test step on the model.\n",
+    "    \n",
+    "    Args:\n",
+    "        model (DoMINO): The trained model\n",
+    "        data_dict (dict): Dictionary containing test data\n",
+    "        device (torch.device): Device to run inference on\n",
+    "        \n",
+    "    Returns:\n",
+    "        tuple: (prediction_vol, prediction_surf) - Model predictions for volume and surface\n",
+    "    \"\"\"\n",
+    "    \n",
+    "    avg_tloss_vol = 0.0  # Placeholder for average volume loss (not currently used)\n",
+    "    avg_tloss_surf = 0.0  # Placeholder for average surface loss (not currently used)\n",
+    "\n",
+    "    with torch.no_grad():  # Disable gradient computation to save memory and computation during inference\n",
+    "        # Move input data to the specified device (CPU or GPU)\n",
+    "        data_dict = dict_to_device(data_dict, device)\n",
+    "\n",
+    "        # Extract non-dimensionalization factors (important for scaling the inputs)\n",
+    "        air_density = data_dict[\"air_density\"]\n",
+    "        stream_velocity = data_dict[\"stream_velocity\"]\n",
+    "        length_scale = data_dict[\"length_scale\"]\n",
+    "\n",
+    "        # Extract geometry coordinates (nodes of the surface)\n",
+    "        geo_centers = data_dict[\"geometry_coordinates\"]\n",
+    "\n",
+    "        # Extract bounding box grid and signed distance function (SDF) grid for the surface\n",
+    "        s_grid = data_dict[\"surf_grid\"]\n",
+    "        sdf_surf_grid = data_dict[\"sdf_surf_grid\"]\n",
+    "\n",
+    "        # Extract scaling factors for surface (used for un-normalization)\n",
+    "        surf_max = data_dict[\"surface_min_max\"][:, 1]\n",
+    "        surf_min = data_dict[\"surface_min_max\"][:, 0]\n",
+    "\n",
+    "        # Normalize geometry coordinates to fit within a bounding box [-1, 1]\n",
+    "        geo_centers_surf = (\n",
+    "            2.0 * (geo_centers - surf_min) / (surf_max - surf_min) - 1\n",
+    "        )\n",
+    "\n",
+    "        # Generate geometric representation of the surface\n",
+    "        encoding_g_surf = model.geo_rep_surface(\n",
+    "            geo_centers_surf, s_grid, sdf_surf_grid\n",
+    "        )\n",
+    "\n",
+    "        prediction_vol = None  # Volume prediction is not computed in this function\n",
+    "\n",
+    "        # Extract information about the surface: mesh centers, normals, areas, and neighbors\n",
+    "        surface_mesh_centers = data_dict[\"surface_mesh_centers\"]\n",
+    "        surface_normals = data_dict[\"surface_normals\"]\n",
+    "        surface_areas = data_dict[\"surface_areas\"]\n",
+    "\n",
+    "        surface_mesh_neighbors = data_dict[\"surface_mesh_neighbors\"]\n",
+    "        surface_neighbors_normals = data_dict[\"surface_neighbors_normals\"]\n",
+    "        surface_neighbors_areas = data_dict[\"surface_neighbors_areas\"]\n",
+    "\n",
+    "        surface_areas = torch.unsqueeze(surface_areas, -1)  # Add extra dimension\n",
+    "        surface_neighbors_areas = torch.unsqueeze(surface_neighbors_areas, -1)  # Add extra dimension\n",
+    "        pos_surface_center_of_mass = data_dict[\"pos_surface_center_of_mass\"]\n",
+    "        num_points = surface_mesh_centers.shape[1]  # Number of surface points\n",
+    "        \n",
+    "        # Extract target surface fields (for comparison later)\n",
+    "        target_surf = data_dict[\"surface_fields\"]\n",
+    "        prediction_surf = np.zeros_like(target_surf.cpu().numpy())  # Initialize prediction array\n",
+    "\n",
+    "        start_time = time.time()  # Record the start time for performance measurement\n",
+    "\n",
+    "        # Generate local geometric encoding for each surface point\n",
+    "        geo_encoding_local = model.geo_encoding_local(\n",
+    "            0.5 * encoding_g_surf, surface_mesh_centers, s_grid, mode=\"surface\"\n",
+    "        )\n",
+    "\n",
+    "\n",
+    "        # Position encoding based on the center of mass of the surface\n",
+    "        pos_encoding = pos_surface_center_of_mass\n",
+    "        pos_encoding = model.position_encoder(pos_encoding, eval_mode=\"surface\")\n",
+    "\n",
+    "        # Perform the model prediction using neighbors and other surface data\n",
+    "        tpredictions = (\n",
+    "            model.calculate_solution_with_neighbors(\n",
+    "                surface_mesh_centers,\n",
+    "                geo_encoding_local,\n",
+    "                pos_encoding,\n",
+    "                surface_mesh_neighbors,\n",
+    "                surface_normals,\n",
+    "                surface_neighbors_normals,\n",
+    "                surface_areas,\n",
+    "                surface_neighbors_areas,\n",
+    "                stream_velocity,\n",
+    "                air_density,\n",
+    "            )\n",
+    "        )\n",
+    "\n",
+    "        # Convert model predictions to numpy arrays for further processing\n",
+    "        prediction_surf = tpredictions.cpu().numpy()\n",
+    "\n",
+    "        # Unnormalize the surface predictions and scale them using physical quantities\n",
+    "        prediction_surf = (\n",
+    "            unnormalize(prediction_surf, surf_factors[0], surf_factors[1])\n",
+    "            * stream_velocity[0, 0].cpu().numpy() ** 2.0\n",
+    "            * air_density[0, 0].cpu().numpy()\n",
+    "        )\n",
+    "\n",
+    "    return prediction_vol, prediction_surf  # Return volume and surface predictions\n",
+    "\n",
+    "def test(filepath, dirname, CKPT_NUMBER: int=0):\n",
+    "    \"\"\"\n",
+    "    High-level function to manage the testing pipeline, including data preparation, model loading, and prediction saving.\n",
+    "    \n",
+    "    Args:\n",
+    "        filepath (str): Path to the test data directory\n",
+    "        dirname (str): Directory name for the test case\n",
+    "    \n",
+    "    Returns:\n",
+    "        None\n",
+    "    \"\"\"\n",
+    "    # Define names of surface variables to be predicted\n",
+    "    surface_variable_names = SURFACE_VARS\n",
+    "    \n",
+    "    # Check if surface scaling factors are available\n",
+    "    surf_save_path = os.path.join(\n",
+    "        \"outputs\", PROJECT_NAME , \"surface_scaling_factors.npy\"\n",
+    "    )\n",
+    "    if os.path.exists(surf_save_path):\n",
+    "        surf_factors = np.load(surf_save_path)  # Load scaling factors if available\n",
+    "    else:\n",
+    "        surf_factors = None  # If not available, set to None\n",
+    "    \n",
+    "    # Load the best model checkpoint\n",
+    "    best_checkpoint = torch.load(CHECKPOINT_DIR / f\"best_model/DoMINO.0.{CKPT_NUMBER}.pt\")\n",
+    "    model.load_state_dict(best_checkpoint)  # Load the model state\n",
+    "    print(\"Model loaded\")\n",
+    "    \n",
+    "    # Set the path to save predictions\n",
+    "    pred_save_path = SAVE_PATH\n",
+    "    create_directory(pred_save_path)  # Create the output directory if it doesn't exist\n",
+    "    \n",
+    "    # Extract test case identifier from the directory name\n",
+    "    tag = int(re.findall(r\"(\\w+?)(\\d+)\", dirname)[0][1])\n",
+    "    vtp_path = filepath  # Path to the VTP file with surface data\n",
+    "    \n",
+    "    # Prepare the path to save predicted results\n",
+    "    vtp_pred_save_path = os.path.join(\n",
+    "        pred_save_path, f\"boundary_{tag}_predicted.vtp\"\n",
+    "    )\n",
+    "    \n",
+    "    # Load the STL file for the geometry\n",
+    "    path_stl = Path(filepath)\n",
+    "    stl_path = path_stl.parent.parent.joinpath(\"test_stl_files\", path_stl.stem + \".stl\")\n",
+    "    print(\"stl_path::\", stl_path)\n",
+    "    print(\"filepath::\", filepath)\n",
+    "\n",
+    "\n",
+    "    info_path = path_stl.with_name(path_stl.stem + \".txt\").parent.parent.joinpath(\"test_info\", path_stl.stem + \"_info.txt\")\n",
+    "    print(\"info_path:::\", info_path)\n",
+    "\n",
+    "    with open(info_path, \"r\") as file:\n",
+    "         for line in file:\n",
+    "             #print(\"line::\",line)\n",
+    "              if \"Velocity\" in line:\n",
+    "                  velocity = float(line.split(\":\")[1].strip())\n",
+    "                  print(f\"Velocity: {velocity}\")\n",
+    "\n",
+    "    STREAM_VELOCITY = velocity\n",
+    "    \n",
+    "    # Read and process the STL file\n",
+    "    reader = pv.get_reader(stl_path)\n",
+    "    mesh_stl = reader.read()\n",
+    "    stl_vertices = mesh_stl.points\n",
+    "    stl_faces = np.array(mesh_stl.faces).reshape((-1, 4))[:, 1:]  # Extract triangular faces\n",
+    "    mesh_indices_flattened = stl_faces.flatten()\n",
+    "    length_scale = np.amax(np.amax(stl_vertices, 0) - np.amin(stl_vertices, 0))  # Compute scale of the geometry\n",
+    "    stl_sizes = mesh_stl.compute_cell_sizes(length=False, area=True, volume=False)\n",
+    "    stl_sizes = np.array(stl_sizes.cell_data[\"Area\"], dtype=np.float32)\n",
+    "    stl_centers = np.array(mesh_stl.cell_centers().points, dtype=np.float32)\n",
+    "    \n",
+    "    # Calculate the center of mass of the surface\n",
+    "    center_of_mass = calculate_center_of_mass(stl_centers, stl_sizes)\n",
+    "    \n",
+    "    # Extract bounding box dimensions for the surface\n",
+    "    bounding_box_dims_surf = []\n",
+    "    bounding_box_dims_surf.append(np.asarray(BOUNDING_BOX_SURF.max))\n",
+    "    bounding_box_dims_surf.append(np.asarray(BOUNDING_BOX_SURF.min))\n",
+    "    s_max = np.float32(bounding_box_dims_surf[0])\n",
+    "    s_min = np.float32(bounding_box_dims_surf[1])\n",
+    "    \n",
+    "    # Create a 3D grid for the surface\n",
+    "    nx, ny, nz = GRID_RESOLUTION\n",
+    "    surf_grid = create_grid(s_max, s_min, [nx, ny, nz])\n",
+    "    surf_grid_reshaped = surf_grid.reshape(nx * ny * nz, 3)\n",
+    "    \n",
+    "    # Compute the Signed Distance Field (SDF) on the surface grid\n",
+    "    sdf_surf_grid = (\n",
+    "        signed_distance_field(\n",
+    "            stl_vertices,\n",
+    "            mesh_indices_flattened,\n",
+    "            surf_grid_reshaped,\n",
+    "            use_sign_winding_number=True,\n",
+    "        )\n",
+    "        .reshape(nx, ny, nz)\n",
+    "    )\n",
+    "    surf_grid = np.float32(surf_grid)\n",
+    "    sdf_surf_grid = np.float32(sdf_surf_grid)\n",
+    "    surf_grid_max_min = np.float32(np.asarray([s_min, s_max]))\n",
+    "    \n",
+    "    # Read the VTP file containing surface data\n",
+    "    reader = vtk.vtkXMLPolyDataReader()\n",
+    "    reader.SetFileName(vtp_path)\n",
+    "    reader.Update()\n",
+    "    polydata_surf = reader.GetOutput()\n",
+    "    celldata_all = get_node_to_elem(polydata_surf)\n",
+    "    celldata = celldata_all.GetCellData()\n",
+    "    surface_fields = get_fields(celldata, surface_variable_names)\n",
+    "    surface_fields = np.concatenate(surface_fields, axis=-1)\n",
+    "    mesh = pv.PolyData(polydata_surf)\n",
+    "    \n",
+    "    # Extract surface mesh coordinates, neighbors, and normals\n",
+    "    surface_coordinates = np.array(mesh.cell_centers().points, dtype=np.float32)\n",
+    "    interp_func = KDTree(surface_coordinates)\n",
+    "    dd, ii = interp_func.query(surface_coordinates, k=NUM_SURFACE_NEIGHBORS)\n",
+    "    surface_neighbors = surface_coordinates[ii]\n",
+    "    surface_neighbors = surface_neighbors[:, 1:]\n",
+    "    surface_normals = np.array(mesh.cell_normals, dtype=np.float32)\n",
+    "    surface_sizes = mesh.compute_cell_sizes(length=False, area=True, volume=False)\n",
+    "    surface_sizes = np.array(surface_sizes.cell_data[\"Area\"], dtype=np.float32)\n",
+    "    \n",
+    "    # Normalize the surface normals and neighbors\n",
+    "    surface_normals = (\n",
+    "        surface_normals / np.linalg.norm(surface_normals, axis=1)[:, np.newaxis]\n",
+    "    )\n",
+    "    surface_neighbors_normals = surface_normals[ii]\n",
+    "    surface_neighbors_normals = surface_neighbors_normals[:, 1:]\n",
+    "    surface_neighbors_sizes = surface_sizes[ii]\n",
+    "    surface_neighbors_sizes = surface_neighbors_sizes[:, 1:]\n",
+    "    \n",
+    "    # Calculate the grid resolution and normalize the surface data\n",
+    "    dx, dy, dz = (\n",
+    "        (s_max[0] - s_min[0]) / nx,\n",
+    "        (s_max[1] - s_min[1]) / ny,\n",
+    "        (s_max[2] - s_min[2]) / nz,\n",
+    "    )\n",
+    "    pos_surface_center_of_mass = surface_coordinates - center_of_mass\n",
+    "    surface_coordinates = normalize(surface_coordinates, s_max, s_min)\n",
+    "    surface_neighbors = normalize(surface_neighbors, s_max, s_min)\n",
+    "    surf_grid = normalize(surf_grid, s_max, s_min)\n",
+    "    \n",
+    "    # Prepare the data dictionary for model input\n",
+    "    geom_centers = np.float32(stl_vertices)\n",
+    "    data_dict = {\n",
+    "        \"pos_surface_center_of_mass\": np.float32(pos_surface_center_of_mass),\n",
+    "        \"geometry_coordinates\": np.float32(geom_centers),\n",
+    "        \"surf_grid\": np.float32(surf_grid),\n",
+    "        \"sdf_surf_grid\": np.float32(sdf_surf_grid),\n",
+    "        \"surface_mesh_centers\": np.float32(surface_coordinates),\n",
+    "        \"surface_mesh_neighbors\": np.float32(surface_neighbors),\n",
+    "        \"surface_normals\": np.float32(surface_normals),\n",
+    "        \"surface_neighbors_normals\": np.float32(surface_neighbors_normals),\n",
+    "        \"surface_areas\": np.float32(surface_sizes),\n",
+    "        \"surface_neighbors_areas\": np.float32(surface_neighbors_sizes),\n",
+    "        \"surface_fields\": np.float32(surface_fields),\n",
+    "        \"surface_min_max\": np.float32(surf_grid_max_min),\n",
+    "        \"length_scale\": np.array(length_scale, dtype=np.float32),\n",
+    "        \"stream_velocity\": np.expand_dims(\n",
+    "            np.array(STREAM_VELOCITY, dtype=np.float32), axis=-1\n",
+    "        ),\n",
+    "        \"air_density\": np.expand_dims(\n",
+    "            np.array(AIR_DENSITY, dtype=np.float32), axis=-1\n",
+    "        ),\n",
+    "    }\n",
+    "    \n",
+    "    # Convert data dictionary to PyTorch tensors\n",
+    "    data_dict = {\n",
+    "        key: torch.from_numpy(np.expand_dims(np.float32(value), 0))\n",
+    "        for key, value in data_dict.items()\n",
+    "    }\n",
+    "    \n",
+    "    # Perform a test step to get the predictions\n",
+    "    prediction_vol, prediction_surf = test_step(\n",
+    "        model, data_dict, surf_factors, device\n",
+    "    )\n",
+    "    \n",
+    "    # Process the predicted and true surface values to compute forces\n",
+    "    surface_sizes = np.expand_dims(surface_sizes, -1)\n",
+    "    pres_x_pred = np.sum(\n",
+    "        prediction_surf[0, :, 0] * surface_normals[:, 0] * surface_sizes[:, 0]\n",
+    "    )\n",
+    "    shear_x_pred = np.sum(prediction_surf[0, :, 1] * surface_sizes[:, 0])\n",
+    "    pres_x_true = np.sum(\n",
+    "        surface_fields[:, 0] * surface_normals[:, 0] * surface_sizes[:, 0]\n",
+    "    )\n",
+    "    shear_x_true = np.sum(surface_fields[:, 1] * surface_sizes[:, 0])\n",
+    "    force_x_pred = np.sum(\n",
+    "        prediction_surf[0, :, 0] * surface_normals[:, 0] * surface_sizes[:, 0]\n",
+    "        - prediction_surf[0, :, 1] * surface_sizes[:, 0]\n",
+    "    )\n",
+    "    force_x_true = np.sum(\n",
+    "        surface_fields[:, 0] * surface_normals[:, 0] * surface_sizes[:, 0]\n",
+    "        - surface_fields[:, 1] * surface_sizes[:, 0]\n",
+    "    )\n",
+    "    \n",
+    "    # Print the computed forces for comparison\n",
+    "    print(dirname, force_x_pred, force_x_true)\n",
+    "    \n",
+    "    # Convert predictions to VTK format and save the results\n",
+    "    surfParam_vtk = numpy_support.numpy_to_vtk(prediction_surf[0, :, 0:1])\n",
+    "    surfParam_vtk.SetName(f\"{surface_variable_names[0]}Pred\")\n",
+    "    celldata_all.GetCellData().AddArray(surfParam_vtk)\n",
+    "    surfParam_vtk = numpy_support.numpy_to_vtk(prediction_surf[0, :, 1:])\n",
+    "    surfParam_vtk.SetName(f\"{surface_variable_names[1]}Pred\")\n",
+    "    celldata_all.GetCellData().AddArray(surfParam_vtk)\n",
+    "    write_to_vtp(celldata_all, vtp_pred_save_path)  # Save to VTP file\n",
+    "    \n",
+    "    return  # End of the test function"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "input_path = DATA_PATHS[\"test\"]\n",
+    "\n",
+    "CHECKPOINT_NUMBER = 1\n",
+    "\n",
+    "dirnames = get_filenames(input_path)\n",
+    "\n",
+    "for count, dirname in enumerate(dirnames):\n",
+    "    print(f\"Processing file {dirname}\")\n",
+    "    filepath = os.path.join(input_path, dirname)\n",
+    "    print(\"filepath::\",filepath)\n",
+    "    test(filepath, dirname, CKPT_NUMBER=CHECKPOINT_NUMBER)\n",
+    "\n",
+    "folder = Path(SAVE_PATH)\n",
+    "predcited_files = list(folder.glob(\"*.vtp\"))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## **Visualizing the predicted results**\n",
+    "You can Visualize the predicted surface pressure using either PyVista or ParaView. In the following, use  we `pyvista` and display both the predicted and ground truth pressure values, which are stored in .vtp files located in the SAVE_PATH directory."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "##### import pyvista as pv\n",
+    "\n",
+    "# Start virtual framebuffer for off-screen rendering (useful in Jupyter/containers)\n",
+    "pv.start_xvfb()\n",
+    "\n",
+    "# Read the VTP mesh\n",
+    "mesh = pv.read(f\"{DATA_DIR}/mesh_predictions_surf_final1/boundary_119_predicted.vtp\")\n",
+    "print(\"Available cell data keys:\", mesh.cell_data.keys())\n",
+    "\n",
+    "# Create a Plotter with 2 vertical subplots\n",
+    "plotter = pv.Plotter(shape=(2, 1), window_size=[1600, 800], off_screen=True)\n",
+    "\n",
+    "# Plot 'p' (ground truth or reference)\n",
+    "plotter.subplot(0, 0)\n",
+    "plotter.add_text(\"Pressure (p)\", font_size=12)\n",
+    "plotter.add_mesh(mesh, scalars=\"p\", show_edges=False)\n",
+    "\n",
+    "# Plot 'pPred' (predicted pressure)\n",
+    "plotter.subplot(1, 0)\n",
+    "plotter.add_text(\"Predicted Pressure (pPred)\", font_size=12)\n",
+    "plotter.add_mesh(mesh, scalars=\"pPred\", show_edges=False)\n",
+    "\n",
+    "# Show both subplots\n",
+    "plotter.show(jupyter_backend='static')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Conclusion\n",
+    "This notebook successfully demonstrates the complete workflow for training and deploying the DoMINO (Decomposable Multiscale Iterative Neural Operator) model on automotive aerodynamics data. Through comprehensive exploration of DoMINO's sophisticated architecture—which combines global and local geometric representations with a DeepONet-inspired aggregation network—we achieved successful training using distributed training capabilities with mixed precision optimization. The model learned to predict surface pressure and wall shear stress fields from the Ahmed body dataset, demonstrating the effectiveness of its multi-scale geometric encoding approach that processes STL geometries directly without requiring mesh generation. Post-training evaluation revealed the model's capability to generalize to new geometries, generating predictions that can be directly compared against traditional CFD simulations, while the force calculations and aerodynamic coefficient predictions provide quantitative validation of the model's accuracy. The successful implementation offers significant potential for rapid aerodynamic analysis during vehicle design iterations, real-time flow field predictions for design optimization, and reduced computational costs compared to traditional CFD simulations. By combining geometric understanding with deep learning, DoMINO represents a promising step toward more efficient and accessible computational fluid dynamics tools that could revolutionize how we approach aerodynamic simulations in automotive engineering."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "os._exit(00)"
+   ]
+  }
+ ],
+ "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.12.3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

introduction.ipynb

@@ -0,0 +1,208 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "528a57dc",
+   "metadata": {},
+   "source": [
+    "## Introduction\n",
+    "\n",
+    "Participants will work with the **Ahmed body** benchmark — a canonical problem in automotive aerodynamics — to build a surrogate model. The course blends classical CFD data handling with modern AI-based modeling. \n",
+    "For the sake of education here we will use only **Ahmed Body Surface data** for training and not volume data. "
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d61bf370",
+   "metadata": {},
+   "source": [
+    "## Course Modules\n",
+    "\n",
+    "### 0. [Dataset download](#download-ahmed-body-surface-dataset)\n",
+    "\n",
+    "### 1. [Preprocessing Ahmed body dataset for training with DoMINO](domino-data-preprocessing.ipynb)\n",
+    "- Loading and understanding VTK simulation data\n",
+    "- Data normalization techniques for ML\n",
+    "- Extracting key physical quantities\n",
+    "- Creating standardized datasets for training\n",
+    "\n",
+    "### 2. [Training and Inferencing DoMINO Model](domino-training-test.ipynb)\n",
+    "- Understanding DoMINO architecture and physics-informed learning\n",
+    "- Training process\n",
+    "- Load Model Checkpoint & Run Inference\n",
+    "- Visualizing the predicted results\n",
+    "\n",
+    "### Hardware Requirements\n",
+    "- NVIDIA GPU with CUDA support\n",
+    "- 4GB+ RAM\n",
+    "- 100GB+ disk space\n",
+    "\n",
+    "### Software Requirements\n",
+    "- Docker with NVIDIA Container Toolkit\n",
+    "- NVIDIA Drivers (version 545 or later)\n",
+    "- SLURM workload manager"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "8643a4d3",
+   "metadata": {},
+   "source": [
+    "## Quick Start Guide\n",
+    "\n",
+    "### Course Environment\n",
+    "This course is set up to use the the PhysicsNeMo docker environment from NVIDIA NGC, and all other additional dependencies are already installed. If you wish to set up the environment outside of this space, the following steps can be used, but note, they are not necessary for this HuggingFace Space.\n",
+    "\n",
+    "### Required Setup Steps\n",
+    "The only required step is to pull the dataset from NVIDIA NGC.\n",
+    "\n",
+    "#### Download Ahmed Body surface dataset\n",
+    "\n",
+    "The complete Ahmed body surface dataset is hosted on NGC and accessible from the following link:\n",
+    "\n",
+    "https://catalog.ngc.nvidia.com/orgs/nvidia/teams/physicsnemo/resources/physicsnemo_ahmed_body_dataset\n",
+    "\n",
+    "The data can be downloaded using the NVIDIA NGC cli. Please note that the dataset contains VTP files, but training DoMINO and X-MeshGraphNet also requires STL files. Therefore, in the `domino-data-preprocessing.ipynb` notebook, STL files are extracted from the available VTP data.\n",
+    "\n",
+    "The data should be available in this space already at \"/data/physicsnemo_ahmed_body_dataset_vv1\". If it is not, we can download it again:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "99acc658",
+   "metadata": {
+    "vscode": {
+     "languageId": "plaintext"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "!if [ ! -d \"/data/physicsnemo_ahmed_body_dataset_vv1\" ]; then \\\n",
+    "    echo \"Dataset not found. Downloading...\"; \\\n",
+    "    ~/app/ngc-cli/ngc registry resource download-version \"nvidia/physicsnemo/physicsnemo_ahmed_body_dataset:v1\" --dest /data; \\\n",
+    "else \\\n",
+    "    echo \"Dataset already exists at /data/physicsnemo_ahmed_body_dataset_vv1\"; \\\n",
+    "fi"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "93554844",
+   "metadata": {},
+   "source": [
+    "### Optional Steps for Standalone Environment Setup\n",
+    "\n",
+    "The next set of steps are only necessary if you would like to rebuild the environment outside of HuggingFace Spaces.\n",
+    "\n",
+    "Once the dataset is downloaded, you man procede to [domino-data-preprocessing.ipynb](domino-data-preprocessing.ipynb)!"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0086ae4a",
+   "metadata": {},
+   "source": [
+    "#### Setting your environment to pull the PhysicsNeMo container from NGC:\n",
+    "\n",
+    "Please refer to the following link for instructions on setting up your environment to pull the PhysicsNeMo container from NGC:\n",
+    "https://docs.nvidia.com/launchpad/ai/base-command-coe/latest/bc-coe-docker-basics-step-02.html\n",
+    "\n",
+    "\n",
+    "#### Running PhysicsNeMo container using docker command\n",
+    "\n",
+    "Pull the PhysicsNeMo container with the following command:\n",
+    "```bash\n",
+    "docker pull nvcr.io/nvidia/physicsnemo/physicsnemo:25.06\n",
+    " ```\n",
+    "\n",
+    "To launch the PhysicsNeMo container using docker, use the following command:\n",
+    "\n",
+    "```bash\n",
+    "docker run --gpus 1  --shm-size=2g -p 7008:7008 --ulimit memlock=-1 --ulimit stack=67108864 --runtime nvidia -v <path_on_host>:/workspace -it --rm $(docker images | grep 25.03| awk '{print $3}')\n",
+    " \n",
+    "```\n",
+    "\n",
+    "Make sure to replace <path_on_host> with the absolute path to the directory on the host system that contains your Jupyter notebooks and Ahmed Body surface data. This path will be mounted as `/workspace` inside the container, providing access to your data and scripts during the session\n",
+    "\n",
+    "#### Start Jupyter Lab\n",
+    "\n",
+    "From the terminal inside the container run the following command to start Jupyter Lab in the background:\n",
+    "\n",
+    "```bash\n",
+    "nohup python3 -m jupyter lab --ip=0.0.0.0 --port=7008 --allow-root  --no-browser --NotebookApp.token='' --notebook-dir='/workspace/' --NotebookApp.allow_origin='*' > /dev/null 2>&1 &\n",
+    "```\n",
+    "\n",
+    "Then from your labtop start a SSH tunnel using the host which your job is running and the port which you assigned above `--port=1234`: \n",
+    "\n",
+    "```bash\n",
+    "ssh -L 3030:eos0311:7008 eos\n",
+    "```\n",
+    "Access Jupyter Lab using `http://localhost:1234` in your browser. \n",
+    "\n",
+    "#### Dependencies\n",
+    "\n",
+    "Please install the following in the container.\n",
+    "```bash\n",
+    "pip install numpy pyvista vtk matplotlib tqdm numpy-stl torchinfo\n",
+    "apt install  xvfb\n",
+    "```"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "cc765d6a",
+   "metadata": {},
+   "source": [
+    "## Troubleshooting\n",
+    "\n",
+    "### Common Issues\n",
+    "\n",
+    "1. **Memory Issues**\n",
+    "   - Monitor GPU memory: `nvidia-smi`\n",
+    "   - Check system memory: `free -h`\n",
+    "\n",
+    "2. **GPU Support**\n",
+    "\n",
+    "   - Run the following command to verify your container runtime supports NVIDIA GPUs:\n",
+    "\n",
+    "```bash\n",
+    "docker run --rm --runtime=nvidia --gpus all ubuntu nvidia-smi\n",
+    "```\n",
+    "\n",
+    "Expected output should show your GPU information, for example:\n",
+    "```\n",
+    "+-----------------------------------------------------------------------------------------+\n",
+    "| NVIDIA-SMI 560.35.03              Driver Version: 560.35.03      CUDA Version: 12.6     |\n",
+    "|-----------------------------------------+------------------------+----------------------+\n",
+    "| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |\n",
+    "| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |\n",
+    "|                                         |                        |               MIG M. |\n",
+    "|=========================================+========================+======================|\n",
+    "|   0  NVIDIA A100-PCIE-40GB          Off |   00000000:41:00.0 Off |                    0 |\n",
+    "| N/A   37C    P0             35W /  250W |   40423MiB /  40960MiB |      0%      Default |\n",
+    "|                                         |                        |             Disabled |\n",
+    "+-----------------------------------------+------------------------+----------------------+\n",
+    "```\n",
+    "\n",
+    "## Additional Resources\n",
+    "\n",
+    "- [NVIDIA Physics NeMo Documentation](https://docs.nvidia.com/deeplearning/nemo-physics/user-guide)\n",
+    "- [PyVista Documentation](https://docs.pyvista.org/) for 3D visualization\n",
+    "- [Ahmed Body Benchmark](https://www.cfd-online.com/Wiki/Ahmed_body) for background\n",
+    "- [Physics-Informed Neural Networks](https://www.sciencedirect.com/science/article/abs/pii/S0021999118307125) paper\n",
+    "- [Graph Neural Networks](https://arxiv.org/abs/2106.10943) for scientific computing\n",
+    "- [Neural Operators](https://arxiv.org/abs/2108.08481) for PDEs\n",
+    "- [DGL Documentation](https://www.dgl.ai/) for graph neural networks\n",
+    "- [Triton Inference Server](https://github.com/triton-inference-server/server) for deployment\n"
+   ]
+  }
+ ],
+ "metadata": {
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

login.html

@@ -0,0 +1,68 @@
+{% extends "page.html" %}
+
+
+{% block stylesheet %}
+{% endblock %}
+
+{% block site %}
+
+<div id="jupyter-main-app" class="container">
+    
+    <img src="https://huggingface.co/front/assets/huggingface_logo-noborder.svg" alt="Hugging Face Logo">
+    <h4>Welcome to JupyterLab</h4>
+    
+    <h5>The default token is <span style="color:orange;">huggingface</span></h5>
+  
+    {% if login_available %}
+    {# login_available means password-login is allowed. Show the form. #}
+    <div class="row">
+        <div class="navbar col-sm-8">
+            <div class="navbar-inner">
+                <div class="container">
+                    <div class="center-nav">
+                        <form action="{{base_url}}login?next={{next}}" method="post" class="navbar-form pull-left">
+                            {{ xsrf_form_html() | safe }}
+                            {% if token_available %}
+                            <label for="password_input"><strong>{% trans %}Jupyter token <span title="This is the secret you set up when deploying your JupyterLab space">ⓘ</span> {% endtrans
+                                    %}</strong></label>
+                            {% else %}
+                            <label for="password_input"><strong>{% trans %}Jupyter password:{% endtrans %}</strong></label>
+                            {% endif %}
+                            <input type="password" name="password" id="password_input" class="form-control">
+                            <button type="submit" class="btn btn-default" id="login_submit">{% trans %}Log in{% endtrans
+                                %}</button>
+                        </form>
+                    </div>
+                </div>
+            </div>
+        </div>
+    </div>
+    {% else %}
+    <p>{% trans %}No login available, you shouldn't be seeing this page.{% endtrans %}</p>
+    {% endif %}
+
+      <h5>If you don't have the credentials for this Jupyter space, <a target="_blank" href="https://huggingface.co/spaces/SpacesExamples/jupyterlab?duplicate=true">create your own.</a></h5>
+    <br>
+  
+    <p>This template was created by <a href="https://twitter.com/camenduru" target="_blank" >camenduru</a> and <a href="https://huggingface.co/nateraw" target="_blank" >nateraw</a>, with contributions of <a href="https://huggingface.co/osanseviero" target="_blank" >osanseviero</a> and <a href="https://huggingface.co/azzr" target="_blank" >azzr</a> </p>
+    {% if message %}
+    <div class="row">
+        {% for key in message %}
+        <div class="message {{key}}">
+            {{message[key]}}
+        </div>
+        {% endfor %}
+    </div>
+    {% endif %}
+    {% if token_available %}
+    {% block token_message %}
+    
+    {% endblock token_message %}
+    {% endif %}
+</div>
+
+{% endblock %}
+
+
+{% block script %}
+{% endblock %}

on_startup.sh

@@ -0,0 +1,5 @@
+#!/bin/bash
+# Write some commands here that will run on root user before startup.
+# For example, to clone transformers and install it in dev mode:
+# git clone https://github.com/huggingface/transformers.git
+# cd transformers && pip install -e ".[dev]"

requirements.txt

@@ -0,0 +1,38 @@
+# Core dependencies
+torch>=2.0.0
+numpy>=1.21.0
+scipy>=1.10.1
+
+# VTK and visualization
+vtk>=9.2.6
+pyvista>=0.34.0
+matplotlib==3.6.2
+
+# NVIDIA PhysicsNeMo with all features
+nvidia-physicsnemo[all]  # Includes VTK, PyVista, and all other optional features
+
+# Jupyter environment
+jupyterlab>=4.0.0
+notebook>=7.0.0
+ipywidgets>=8.0.0
+jupyter-archive 
+jupyter-resource-usage
+jupyterlab-nvdashboard
+jupyterlab-widgets>=3.0.0
+widgetsnbextension>=4.0.0
+
+# Additional project dependencies
+apex
+boto3==1.34.69
+opencv-python==4.8.0.74
+numpy-stl>=2.16.3
+torchdata>=0.7.0
+zarr[jupyter]
+httpx>=0.24.0
+trimesh>=3.9.0
+python-multipart==0.0.20
+hydra-core==1.3.2
+warp-lang==1.5.1
+loguru
+# Triton Inference Server client
+tritonclient[all]>=2.39.0

start_server.sh

@@ -0,0 +1,19 @@
+#!/bin/bash
+JUPYTER_TOKEN="${JUPYTER_TOKEN:=huggingface}"
+
+NOTEBOOK_DIR="/domino-demo"
+
+jupyter labextension disable "@jupyterlab/apputils-extension:announcements"
+
+jupyter-lab \
+    --ip 0.0.0.0 \
+    --port 7860 \
+    --no-browser \
+    --allow-root \
+    --ServerApp.token="$JUPYTER_TOKEN" \
+    --ServerApp.tornado_settings="{'headers': {'Content-Security-Policy': 'frame-ancestors *'}}" \
+    --ServerApp.cookie_options="{'SameSite': 'None', 'Secure': True}" \
+    --ServerApp.disable_check_xsrf=True \
+    --LabApp.news_url=None \
+    --LabApp.check_for_updates_class="jupyterlab.NeverCheckForUpdate" \
+    --notebook-dir=$NOTEBOOK_DIR