Thermal_domino_ag

physicsnemo/datapipes/cae/domino_datapipe.py

@@ -1048,43 +1048,45 @@ def compute_scaling_factors(cfg: DictConfig, input_path: str, use_cache: bool) -
                 gpu_output=True,
             )
 
-            # Calculate mean
+            # Calculate mean and std
             if cfg.model.normalization == "mean_std_scaling":
+            # Collect all data from the dataset in a single loop.
+                all_vol_fields = []
                 for j in range(len(fm_dict)):
-                    print("On iteration {j}")
                     d_dict = fm_dict[j]
                     vol_fields = d_dict["volume_fields"]
 
                     if vol_fields is not None:
-                        if j == 0:
-                            vol_fields_sum = np.mean(vol_fields, 0)
-                        else:
-                            vol_fields_sum += np.mean(vol_fields, 0)
-                    else:
-                        vol_fields_sum = 0.0
+                        all_vol_fields.append(vol_fields)
 
-                vol_fields_mean = vol_fields_sum / len(fm_dict)
+                if all_vol_fields:
+                    # Concatenate all tensors into one large tensor.
+                    full_vol_tensor = torch.cat(all_vol_fields, dim=0)
 
-                for j in range(len(fm_dict)):
-                    print("On iteration {j} again")
-                    d_dict = fm_dict[j]
-                    vol_fields = d_dict["volume_fields"]
+                    if full_vol_tensor.device.type == "cuda":
+                        try:
+                            xp = cp
 
-                    if vol_fields is not None:
-                        if j == 0:
-                            vol_fields_sum_square = np.mean(
-                                (vol_fields - vol_fields_mean) ** 2.0, 0
-                            )
-                        else:
-                            vol_fields_sum_square += np.mean(
-                                (vol_fields - vol_fields_mean) ** 2.0, 0
-                            )
+                            full_vol_array = cp.from_dlpack(full_vol_tensor)
+                        except ImportError:
+                            xp = np
+                            full_vol_array = full_vol_tensor.cpu().numpy()
                     else:
-                        vol_fields_sum_square = 0.0
+                        xp = np
+                        full_vol_array = full_vol_tensor.cpu().numpy()
 
-                vol_fields_std = np.sqrt(vol_fields_sum_square / len(fm_dict))
+                    # Compute mean and std deviation
+                    vol_fields_mean = xp.mean(full_vol_array, axis=0)
+                    vol_fields_std = xp.std(full_vol_array, axis=0)
 
+                else:
+                    vol_fields_mean = 0.0
+                    vol_fields_std = 0.0
+
+                # Store the final scaling factors
                 vol_scaling_factors = [vol_fields_mean, vol_fields_std]
+                print("vol_scaling_factors: ", vol_scaling_factors)
+
 
             if cfg.model.normalization == "min_max_scaling":
                 for j in range(len(fm_dict)):
@@ -1160,43 +1162,52 @@ def compute_scaling_factors(cfg: DictConfig, input_path: str, use_cache: bool) -
                 compute_scaling_factors=True,
             )
 
-            # Calculate mean
+            # Calculate mean and std for surface fields
             if cfg.model.normalization == "mean_std_scaling":
+
+                # Collect all data from the dataset in a single loop.
+                all_surf_fields = []
                 for j in range(len(fm_dict)):
-                    print(f"Mean std scaling on iteration {j}")
                     d_dict = fm_dict[j]
-                    surf_fields = d_dict["surface_fields"].cpu().numpy()
+                    surf_fields = d_dict.get("surface_fields")
 
                     if surf_fields is not None:
-                        if j == 0:
-                            surf_fields_sum = np.mean(surf_fields, 0)
-                        else:
-                            surf_fields_sum += np.mean(surf_fields, 0)
+                        all_surf_fields.append(surf_fields)
+
+                if all_surf_fields:
+                    # Concatenate all tensors into one large tensor.
+                    full_surf_tensor = torch.cat(all_surf_fields, dim=0)
+
+                    if full_surf_tensor.device.type == "cuda":
+                        try:
+                            xp = cp
+                            full_surf_array = cp.from_dlpack(full_surf_tensor)
+                        except (ImportError, NameError):
+                            xp = np
+                            full_surf_array = full_surf_tensor.cpu().numpy()
                     else:
-                        surf_fields_sum = 0.0
+                        xp = np
+                        full_surf_array = full_surf_tensor.cpu().numpy()
 
-                surf_fields_mean = surf_fields_sum / len(fm_dict)
+                    # Compute mean and std deviation
+                    surf_fields_mean = xp.mean(full_surf_array, axis=0)
+                    surf_fields_std = xp.std(full_surf_array, axis=0)
 
-                for j in range(len(fm_dict)):
-                    print(f"Mean std scaling on iteration {j} again")
-                    d_dict = fm_dict[j]
-                    surf_fields = d_dict["surface_fields"]
+                else:
+                    surf_fields_mean = 0.0
+                    surf_fields_std = 0.0
 
-                    if surf_fields is not None:
-                        if j == 0:
-                            surf_fields_sum_square = np.mean(
-                                (surf_fields - surf_fields_mean) ** 2.0, 0
-                            )
-                        else:
-                            surf_fields_sum_square += np.mean(
-                                (surf_fields - surf_fields_mean) ** 2.0, 0
-                            )
-                    else:
-                        surf_fields_sum_square = 0.0
+                # Store the final scaling factors.
+                surf_scaling_factors = [surf_fields_mean, surf_fields_std]
+                print("surf_scaling_factors:",surf_scaling_factors)
 
-                surf_fields_std = np.sqrt(surf_fields_sum_square / len(fm_dict))
+                # Save the final scaling factors.
+                if xp and xp.__name__ == 'cupy':
+                    surf_scaling_factors_np = [arr.get() for arr in surf_scaling_factors]
+                    np.save(surf_save_path, surf_scaling_factors_np)
+                else:
+                    np.save(surf_save_path, surf_scaling_factors)
 
-                surf_scaling_factors = [surf_fields_mean, surf_fields_std]
 
             if cfg.model.normalization == "min_max_scaling":
                 for j in range(len(fm_dict)):

physicsnemo/utils/domino/utils.py

@@ -504,37 +504,63 @@ def get_vertices(polydata: "vtk.vtkPolyData") -> np.ndarray:
         for each vertex.
 
     """
+    # Use GetPoints() to get the geometric coordinates.
     vtk_points = polydata.GetPoints()
-    vertices = numpy_support.vtk_to_numpy(vtk_points.GetData())
+
+    if vtk_points and vtk_points.GetNumberOfPoints() > 0:
+
+        vertices = numpy_support.vtk_to_numpy(vtk_points.GetData())
+    else:
+        print("\nWarning: No points found in the polydata object.")
+        vertices = np.empty((0, 3))
     return vertices
 
 
+
+
 def get_volume_data(
     polydata: "vtk.vtkPolyData", variable_names: list[str]
 ) -> tuple[np.ndarray, list[np.ndarray]]:
-    """Extract vertices and field data from 3D volumetric mesh.
+    """
+    Extracts vertices and field data from a 3D volumetric mesh.
+    Ensures that the final field data is associated with points. If data is only
+    found on cells, it will be interpolated to the points.
+    """
+    point_data = polydata.GetPointData()
+    cell_data = polydata.GetCellData()
+    fields = None
 
-    This function extracts both geometric information (vertex coordinates)
-    and field data from a 3D volumetric mesh. It's commonly used for
-    processing finite element analysis results.
+    # First, check if the data already exists on the points.
+    if point_data.HasArray(variable_names[0]):
+        vertices = get_vertices(polydata)
+        print("Using existing PointData.")
+        fields = get_fields(point_data, variable_names)
 
-    Args:
-        polydata: VTK polydata representing a 3D volumetric mesh.
-        variable_names: List of field variable names to extract.
+    # If not on points, check if it's on the cells and convert it.
+    elif cell_data.HasArray(variable_names[0]):
+        print("Data found on cells. Converting CellData to PointData...")
 
-    Returns:
-        Tuple containing:
-        - Vertex coordinates as NumPy array of shape (n_vertices, 3)
-        - List of field arrays, one per variable
+        # Create the conversion filter.
+        c2p_filter = vtk.vtkCellDataToPointData()
 
-    """
-    vertices = get_vertices(polydata)
-    point_data = polydata.GetPointData()
-    fields = get_fields(point_data, variable_names)
+        # Set the input mesh for the filter.
+        c2p_filter.SetInputData(polydata)
+        c2p_filter.Update()
+        processed_polydata = c2p_filter.GetOutput()
+
+        vertices = get_vertices(processed_polydata)
+
+        # Extract the fields from the new point data.
+        fields = get_fields(processed_polydata.GetPointData(), variable_names)
+
+    else:
+        print(f"\nWarning: Could not find variables '{variable_names}' in PointData or CellData.")
+        return vertices, []
 
     return vertices, fields
 
 
+
 def get_surface_data(
     polydata: "vtk.vtkPolyData", variable_names: list[str]
 ) -> tuple[np.ndarray, list[np.ndarray], list[tuple[int, int]]]: