Add significantly more coverage of the domino datapipe to catch more errors (#1102)

* Add significantly more coverage of the domino datapipe to catch more errors.

* update changelog

* Update tests based on PR review

* Rework dataset tests for domino by using synthetic datasets.

* Remove length scale

Author: coreyjadams

CHANGELOG.md

@@ -13,6 +13,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Added lead-time aware training support to the StormCast example.
 - Add a device aware kNN method to physicsnemo.utils.neighbors. Works with CPU or GPU
   by dispatching to the proper optimized library, and torch.compile compatible.
+- Added additional testing of the DoMINO datapipe.
 
 ### Changed
 

physicsnemo/datapipes/cae/domino_datapipe.py

@@ -522,8 +522,6 @@ def preprocess_combined(self, data_dict):
         if mesh_indices_flattened.dtype != xp.int32:
             mesh_indices_flattened = mesh_indices_flattened.astype(xp.int32)
 
-        length_scale = xp.amax(xp.amax(stl_vertices, 0) - xp.amin(stl_vertices, 0))
-
         center_of_mass = calculate_center_of_mass(stl_centers, stl_sizes)
 
         if self.config.bounding_box_dims_surf is None:
@@ -570,7 +568,6 @@ def preprocess_combined(self, data_dict):
         surf_grid_max_min = xp.stack([s_min, s_max])
 
         return_dict = {
-            "length_scale": length_scale,
             "surf_grid": surf_grid,
             "sdf_surf_grid": sdf_surf_grid,
             "surface_min_max": surf_grid_max_min,
@@ -651,7 +648,7 @@ def preprocess_surface(self, data_dict, core_dict, center_of_mass, s_min, s_max)
                     (s_max[2] - s_min[2]) / nz,
                 )
                 pos_normals_com_surface = calculate_normal_positional_encoding(
-                    surface_coordinates, center_of_mass, cell_length=[dx, dy, dz]
+                    surface_coordinates, center_of_mass, cell_dimensions=[dx, dy, dz]
                 )
             else:
                 pos_normals_com_surface = surface_coordinates - xp.asarray(
@@ -744,7 +741,13 @@ def preprocess_surface(self, data_dict, core_dict, center_of_mass, s_min, s_max)
 
             else:
                 # We are *not* sampling, kNN on ALL points:
-                ii = knn.kneighbors(surface_coordinates, return_distance=False)
+                if self.array_provider == cp:
+                    ii = knn.kneighbors(surface_coordinates, return_distance=False)
+                else:
+                    _, ii = interp_func.query(
+                        surface_coordinates,
+                        k=self.config.num_surface_neighbors,
+                    )
 
                 # Construct the neighbors arrays:
                 surface_neighbors = surface_coordinates[ii][:, 1:]
@@ -892,10 +895,10 @@ def preprocess_volume(
                 pos_normals_closest_vol = calculate_normal_positional_encoding(
                     volume_coordinates,
                     sdf_node_closest_point,
-                    cell_length=[dx, dy, dz],
+                    cell_dimensions=[dx, dy, dz],
                 )
                 pos_normals_com_vol = calculate_normal_positional_encoding(
-                    volume_coordinates, center_of_mass, cell_length=[dx, dy, dz]
+                    volume_coordinates, center_of_mass, cell_dimensions=[dx, dy, dz]
                 )
             else:
                 pos_normals_closest_vol = volume_coordinates - sdf_node_closest_point