Update DoMINO to enable multi-batch tests to pass.

Deliberately looping an index_select over the batch index, rather
than something vectorized, because we don't want to expand that op
and it's domain parallelized.

Author: coreyjadams

physicsnemo/models/domino/encodings.py

@@ -159,10 +159,11 @@ def forward(
         for j in range(encoding_g.shape[1]):
             geo_encoding = rearrange(encoding_g[:, j], "b nx ny nz -> b 1 (nx ny nz)")
 
-            geo_encoding_sampled = torch.index_select(
-                geo_encoding, 2, mapping.flatten()
-            )
-            geo_encoding_sampled = torch.reshape(geo_encoding_sampled, mask.shape)
+            sampled = []
+            for b in range(batch_size):
+                s = torch.index_select(geo_encoding[b], 1, mapping[b].flatten())
+                sampled.append(s.reshape(1, *mask.shape[1:]))
+            geo_encoding_sampled = torch.cat(sampled, dim=0)
             geo_encoding_sampled = geo_encoding_sampled * mask
 
             encoding_g_inner.append(geo_encoding_sampled)

physicsnemo/models/domino/geometry_rep.py

@@ -195,7 +195,7 @@ def forward(
             self.grid_resolution[1],
             self.grid_resolution[2],
         )
-        grid = grid.reshape(1, nx * ny * nz, 3, 1)
+        grid = grid.reshape(grid.shape[0], nx * ny * nz, 3, 1)
 
         # Rearrange input to flatten spatial and neighbor dimensions
         x = rearrange(

physicsnemo/models/domino/model.py

@@ -574,8 +574,8 @@ def forward(
 
             # Normalize geometry coordinates based on computational domain
             if "volume_min_max" in data_dict.keys():
-                vol_max = data_dict["volume_min_max"][:, 1]
-                vol_min = data_dict["volume_min_max"][:, 0]
+                vol_max = data_dict["volume_min_max"][:, 1].unsqueeze(1)
+                vol_min = data_dict["volume_min_max"][:, 0].unsqueeze(1)
                 geo_centers_vol = (
                     2.0 * (geo_centers - vol_min) / (vol_max - vol_min) - 1
                 )
@@ -618,8 +618,8 @@ def forward(
             # Represent geometry on bounding box
             # Normalize geometry coordinates based on surface bounding box
             if "surface_min_max" in data_dict.keys():
-                surf_max = data_dict["surface_min_max"][:, 1]
-                surf_min = data_dict["surface_min_max"][:, 0]
+                surf_max = data_dict["surface_min_max"][:, 1].unsqueeze(1)
+                surf_min = data_dict["surface_min_max"][:, 0].unsqueeze(1)
                 geo_centers_surf = (
                     2.0 * (geo_centers - surf_min) / (surf_max - surf_min) - 1
                 )