Generalize the encode_parameter in DoMINO to any arbitrary no. of global parameters input (#903)

* changes based on updated main branch

* update to model.py and end to end testing

* changes to sharded parts of the code

* Update README

* Update inference_on_stl.py to comply with new method

* minor refactor

* update

* Tested training

* remove hardcoded stuff from inference_on_stl.py

* Removed comments from model.py

* Remove air_density and stream_velocity from domino_sharded

* Remove comments from domino_datapipe

* Removed names and make paths generic

* make encode_parameters false

* Update and remove comments

* Update README

* Update README, remove redundant text

* Update model.py to remove air_density and stream_velocity

* Update inference_on_stl.py to be consistent with main

* Update README.md to be compliant with main

* Update tests

* changes based on CI

* small cleaning config.yaml

* Update changelog

* fixing doctest issue

---------

Co-authored-by: Peter Sharpe <peterdsharpe@gmail.com>
Co-authored-by: Rishikesh Ranade <dr.rranade@gmail.com>
Co-authored-by: RishikeshRanade <65631160+RishikeshRanade@users.noreply.github.com>

Author: 3 people

.gitignore

@@ -153,6 +153,7 @@ cython_debug/
 
 # VsCode
 .vscode/
+.cursor/
 
 # VIM
 *.swp

CHANGELOG.md

@@ -14,6 +14,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Safe API to override `__init__`'s arguments saved in checkpoint file with
   `Module.from_checkpoint("chkpt.mdlus", models_args)`.
 - PyTorch Geometric MeshGraphNet backend.
+- Functionality in DoMINO to take arbitrary number of `scalar` or `vector`
+  global parameters and encode them using `class ParameterModel`
 
 ### Changed
 

examples/cfd/external_aerodynamics/domino/README.md

@@ -241,6 +241,52 @@ The DoMINO model can be evaluated directly on unknown STLs using the pre-trained
 4. The surface predictions are carried out on the STL surface. The drag and lift
  accuracy will depend on the resolution of the STL.
 
+### Incorporating multiple global simulation parameters for training/inference
+
+DoMINO supports incorporating multiple global simulation parameters (such as inlet
+velocity, air density, etc.) that can vary across different simulations.
+
+1. Define global parameters in the `variables.global_parameters` section of
+   `conf/config.yaml`. Each parameter must specify its type (`vector` or `scalar`)
+   and reference values for non-dimensionalization.
+
+2. For `vector` type parameters:
+   - If values are single-direction vectors (e.g., [30, 0, 0]), define reference as [30]
+   - If values are two-direction vectors (e.g., [30, 30, 0]), define reference as [30, 30]
+
+3. Enable parameter encoding in the model configuration by setting
+   `model.encode_parameters: true`. This will:
+   - Create a dedicated parameter encoding network (`ParameterModel`)
+   - Non-dimensionalize parameters using reference values from `config.yaml`
+   - Integrate parameter encodings into both surface and volume predictions
+
+4. Ensure your simulation data includes global parameter values. The DoMINO
+   datapipe expects these parameters in the pre-processed `.npy`/`.npz` files:
+   - Examine `openfoam_datapipe.py` and `process_data.py` for examples of how global
+     parameter values are incorporated for external aerodynamics
+   - For the automotive example, `air_density` and `inlet_velocity` remain constant
+     across simulations
+   - Adapt these files for your specific case to correctly calculate
+     `global_params_values` and `global_params_reference` during data preprocessing
+
+5. During training, the model automatically handles global parameter encoding when
+   `model.encode_parameters: true` is set
+   - You may need to adapt `train.py` if you plan to use global parameters in loss
+     functions or de-non-dimensionalization
+
+6. During testing with `test.py`, define `global_params_values` for each test sample:
+   - Global parameters must match those defined in `config.yaml`
+   - For each parameter (e.g., "inlet_velocity", "air_density"), provide appropriate
+     values for each simulation
+   - See the `main()` function in `test.py` for implementation examples
+   - If using global parameters for de-non-dimensionalization, modify `test_step()`
+
+7. When inferencing on unseen geometries with `inference_on_stl.py`:
+   - Define `global_params_values` and `global_params_reference` in both
+     `compute_solution_in_volume()` and `compute_solution_on_surface()` methods
+   - Adjust these parameters based on your specific use case and parameters defined
+     in `config.yaml`
+
 ## Extending DoMINO to a custom dataset
 
 This repository includes examples of **DoMINO** training on the DrivAerML dataset.

examples/cfd/external_aerodynamics/domino/src/conf/config.yaml

@@ -59,6 +59,13 @@ variables:
       UMeanTrim: vector
       pMeanTrim: scalar
       nutMeanTrim: scalar
+  global_parameters:
+    inlet_velocity:
+      type: vector
+      reference: [30.0] # vector [30, 0, 0] should be specified as [30], while [30, 30, 0] should be [30, 30].
+    air_density:
+      type: scalar
+      reference: 1.226
 
 # ┌───────────────────────────────────────────┐
 # │          Training Data Configs            │
@@ -146,7 +153,6 @@ model:
     base_layer: 512
   parameter_model:
     base_layer: 512
-    scaling_params: [30.0, 1.226] # [inlet_velocity, air_density]
     fourier_features: false
     num_modes: 5
 
@@ -163,7 +169,7 @@ train: # Training configurable parameters
     shuffle: true
     drop_last: false
   checkpoint_dir: /user/models/ # Use only for retraining
-  
+
 # ┌───────────────────────────────────────────┐
 # │          Validation Configs               │
 # └───────────────────────────────────────────┘  

examples/cfd/external_aerodynamics/domino/src/inference_on_stl.py

@@ -151,7 +151,6 @@ def _bvh_query_distance(
     sdf_hit_point: wp.array(dtype=wp.vec3f),
     sdf_hit_point_id: wp.array(dtype=wp.int32),
 ):
-
     """
     Computes the signed distance from each point in the given array `points`
     to the mesh represented by `mesh`,within the maximum distance `max_dist`,
@@ -690,7 +689,11 @@ def __init__(
         self.air_density = torch.full((1, 1), 1.205, dtype=torch.float32).to(
             self.device
         )
-        self.num_vol_vars, self.num_surf_vars = self.get_num_variables()
+        (
+            self.num_vol_vars,
+            self.num_surf_vars,
+            self.num_global_features,
+        ) = self.get_num_variables()
         self.model = None
         self.grid_resolution = torch.tensor(self.cfg.model.interp_res).to(self.device)
         self.vol_factors = None
@@ -755,28 +758,25 @@ def load_bounding_box(self):
             self.bounding_box_surface_min_max = [c_min, c_max]
 
     def load_volume_scaling_factors(self):
-        vol_factors = np.array(
-            [
-                [2.2642279, 2.2397292, 1.8689916, 0.7547227],
-                [-1.2899836, -2.2787743, -1.866153, -2.7116761],
-            ],
-            dtype=np.float32,
+        scaling_param_path = self.cfg.eval.scaling_param_path
+        vol_factors_path = os.path.join(
+            scaling_param_path, "volume_scaling_factors.npy"
         )
 
+        vol_factors = np.load(vol_factors_path, allow_pickle=True)
         vol_factors = torch.from_numpy(vol_factors).to(self.device)
 
         return vol_factors
 
     def load_surface_scaling_factors(self):
-        surf_factors = np.array(
-            [
-                [0.8215038, 0.01063187, 0.01514608, 0.01327803],
-                [-2.1505525, -0.01865184, -0.01514422, -0.0121509],
-            ],
-            dtype=np.float32,
+        scaling_param_path = self.cfg.eval.scaling_param_path
+        surf_factors_path = os.path.join(
+            scaling_param_path, "surface_scaling_factors.npy"
         )
 
+        surf_factors = np.load(surf_factors_path, allow_pickle=True)
         surf_factors = torch.from_numpy(surf_factors).to(self.device)
+
         return surf_factors
 
     def read_stl(self):
@@ -842,14 +842,28 @@ def get_num_variables(self):
                 num_surf_vars += 3
             else:
                 num_surf_vars += 1
-        return num_vol_vars, num_surf_vars
+
+        num_global_features = 0
+        global_params_names = list(cfg.variables.global_parameters.keys())
+        for param in global_params_names:
+            if cfg.variables.global_parameters[param].type == "vector":
+                num_global_features += len(
+                    cfg.variables.global_parameters[param].reference
+                )
+            elif cfg.variables.global_parameters[param].type == "scalar":
+                num_global_features += 1
+            else:
+                raise ValueError(f"Unknown global parameter type")
+
+        return num_vol_vars, num_surf_vars, num_global_features
 
     def initialize_model(self, model_path):
         model = (
             DoMINO(
                 input_features=3,
                 output_features_vol=self.num_vol_vars,
                 output_features_surf=self.num_surf_vars,
+                global_features=self.num_global_features,
                 model_parameters=self.cfg.model,
             )
             .to(self.device)
@@ -1358,6 +1372,21 @@ def compute_solution_on_surface(
         inlet_velocity,
         air_density,
     ):
+        """
+        Global parameters: For this particular case, the model was trained on single velocity/density values
+        across all simulations. Hence, global_params_values and global_params_reference are the same.
+        """
+        global_params_values = torch.cat(
+            (inlet_velocity, air_density), axis=1
+        )  # (1, 2)
+        global_params_values = torch.unsqueeze(global_params_values, -1)  # (1, 2, 1)
+
+        global_params_reference = torch.cat(
+            (inlet_velocity, air_density), axis=1
+        )  # (1, 2)
+        global_params_reference = torch.unsqueeze(
+            global_params_reference, -1
+        )  # (1, 2, 1)
 
         if self.dist.world_size == 1:
             geo_encoding_local = model.geo_encoding_local(
@@ -1383,8 +1412,8 @@ def compute_solution_on_surface(
                 surface_neighbors_normals,
                 surface_areas,
                 surface_neighbors_areas,
-                inlet_velocity,
-                air_density,
+                global_params_values,
+                global_params_reference,
             )
         else:
             pos_encoding = model.module.position_encoder(
@@ -1399,8 +1428,8 @@ def compute_solution_on_surface(
                 surface_neighbors_normals,
                 surface_areas,
                 surface_neighbors_areas,
-                inlet_velocity,
-                air_density,
+                global_params_values,
+                global_params_reference,
             )
 
         return tpredictions_batch
@@ -1420,6 +1449,19 @@ def compute_solution_in_volume(
         air_density,
     ):
 
+        ## Global parameters
+        global_params_values = torch.cat(
+            (inlet_velocity, air_density), axis=1
+        )  # (1, 2)
+        global_params_values = torch.unsqueeze(global_params_values, -1)  # (1, 2, 1)
+
+        global_params_reference = torch.cat(
+            (inlet_velocity, air_density), axis=1
+        )  # (1, 2)
+        global_params_reference = torch.unsqueeze(
+            global_params_reference, -1
+        )  # (1, 2, 1)
+
         if self.dist.world_size == 1:
             geo_encoding_local = model.geo_encoding_local(
                 geo_encoding, volume_mesh_centers, p_grid, mode="volume"
@@ -1441,8 +1483,8 @@ def compute_solution_in_volume(
                 volume_mesh_centers,
                 geo_encoding_local,
                 pos_encoding,
-                inlet_velocity,
-                air_density,
+                global_params_values,
+                global_params_reference,
                 num_sample_points=self.stencil_size,
                 eval_mode="volume",
             )
@@ -1454,8 +1496,8 @@ def compute_solution_in_volume(
                 volume_mesh_centers,
                 geo_encoding_local,
                 pos_encoding,
-                inlet_velocity,
-                air_density,
+                global_params_values,
+                global_params_reference,
                 num_sample_points=self.stencil_size,
                 eval_mode="volume",
             )
@@ -1481,8 +1523,8 @@ def compute_solution_in_volume(
 
     domino = dominoInference(cfg, dist, False)
     domino.initialize_model(
-        model_path="/lustre/rranade/modulus_dev/modulus_forked/modulus/examples/cfd/external_aerodynamics/domino/outputs/AWS_Dataset/19/models/DoMINO.0.201.pt"
-    )
+        model_path="/lustre/models/DoMINO.0.7.pt"
+    )  ## Replace the model path with location of the trained model
 
     for count, dirname in enumerate(dirnames_per_gpu):
         # print(f"Processing file {dirname}")
@@ -1525,9 +1567,8 @@ def compute_solution_in_volume(
             "Lift:",
             out_dict["lift_force"],
         )
-        vtp_path = f"/lustre/rranade/modulus_dev/modulus_forked/modulus/examples/cfd/external_aerodynamics/domino/pred_{dirname}_4.vtp"
+        vtp_path = f"/lustre/snidhan/physicsnemo-work/domino-global-param-runs/stl-results/pred_{dirname}_4.vtp"
         domino.mesh_stl.save(vtp_path)
-        # vtp_path = f"/lustre/rranade/modulus_dev/modulus_demo/modulus_rishi/modulus/examples/cfd/external_aerodynamics/domino_gtc_demo/sensitivity_pred_{dirname}.vtp"
         reader = vtk.vtkXMLPolyDataReader()
         reader.SetFileName(f"{vtp_path}")
         reader.Update()