Formatting changes

Author: saikrishnanc-nv

examples/structural_mechanics/crash/tests/test_vtp_reader.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: Copyright (c) 2023 - 2025 NVIDIA CORPORATION & AFFILIATES.
 # SPDX-FileCopyrightText: All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
@@ -22,6 +22,7 @@
 
 # Import functions from vtp_reader
 import sys
+
 sys.path.insert(0, str(Path(__file__).parent.parent))
 from vtp_reader import (
     load_vtp_file,
@@ -34,136 +35,164 @@
 def simple_vtp_file():
     """Create a simple VTP file for testing."""
     # Create a simple quad mesh (2x2 grid)
-    points = np.array([
-        [0, 0, 0],
-        [1, 0, 0],
-        [0, 1, 0],
-        [1, 1, 0],
-    ], dtype=np.uint8)
-    
+    points = np.array(
+        [
+            [0, 0, 0],
+            [1, 0, 0],
+            [0, 1, 0],
+            [1, 1, 0],
+        ],
+        dtype=np.uint8,
+    )
+
     # Single quad cell
     faces = np.array([4, 0, 1, 3, 2])  # quad with 4 vertices
-    
+
     mesh = pv.PolyData(points, faces, force_float=False)
-    
+
     # Add displacement fields for 3 timesteps
-    mesh.point_data['displacement_t0.000'] = np.array([
-        [0, 0, 0],
-        [0, 0, 0],
-        [0, 0, 0],
-        [0, 0, 0],
-    ], dtype=np.uint8)
-    
-    mesh.point_data['displacement_t0.005'] = np.array([
-        [1, 0, 0],
-        [1, 0, 0],
-        [1, 0, 0],
-        [1, 0, 0],
-    ], dtype=np.uint8)
-    
-    mesh.point_data['displacement_t0.010'] = np.array([
-        [2, 0, 0],
-        [2, 0, 0],
-        [2, 0, 0],
-        [2, 0, 0],
-    ], dtype=np.uint8)
-    
+    mesh.point_data["displacement_t0.000"] = np.array(
+        [
+            [0, 0, 0],
+            [0, 0, 0],
+            [0, 0, 0],
+            [0, 0, 0],
+        ],
+        dtype=np.uint8,
+    )
+
+    mesh.point_data["displacement_t0.005"] = np.array(
+        [
+            [1, 0, 0],
+            [1, 0, 0],
+            [1, 0, 0],
+            [1, 0, 0],
+        ],
+        dtype=np.uint8,
+    )
+
+    mesh.point_data["displacement_t0.010"] = np.array(
+        [
+            [2, 0, 0],
+            [2, 0, 0],
+            [2, 0, 0],
+            [2, 0, 0],
+        ],
+        dtype=np.uint8,
+    )
+
     # Add thickness as additional point data
-    mesh.point_data['thickness'] = np.array([1, 1, 1, 1], dtype=np.uint8)
-    
+    mesh.point_data["thickness"] = np.array([1, 1, 1, 1], dtype=np.uint8)
+
     # Save to temporary file
-    with tempfile.NamedTemporaryFile(suffix='.vtp', delete=False) as f:
+    with tempfile.NamedTemporaryFile(suffix=".vtp", delete=False) as f:
         temp_path = f.name
-    
+
     mesh.save(temp_path)
     yield temp_path
-    
+
     # Cleanup
     Path(temp_path).unlink(missing_ok=True)
 
 
 def test_load_vtp_file_basic(simple_vtp_file):
     """Test basic VTP file loading."""
     pos_raw, mesh_connectivity, point_data_dict = load_vtp_file(simple_vtp_file)
-    
+
     # Check positions shape: (timesteps, nodes, 3)
     assert pos_raw.shape == (3, 4, 3), f"Expected shape (3, 4, 3), got {pos_raw.shape}"
-    
+
     # Check mesh connectivity
     assert len(mesh_connectivity) == 1, f"Expected 1 cell, got {len(mesh_connectivity)}"
-    assert len(mesh_connectivity[0]) == 4, f"Expected quad with 4 vertices, got {len(mesh_connectivity[0])}"
-    
+    assert len(mesh_connectivity[0]) == 4, (
+        f"Expected quad with 4 vertices, got {len(mesh_connectivity[0])}"
+    )
+
     # Check point data dict contains thickness
-    assert 'thickness' in point_data_dict, "Thickness not found in point_data_dict"
-    assert point_data_dict['thickness'].shape == (4,), f"Expected thickness shape (4,), got {point_data_dict['thickness'].shape}"
+    assert "thickness" in point_data_dict, "Thickness not found in point_data_dict"
+    assert point_data_dict["thickness"].shape == (4,), (
+        f"Expected thickness shape (4,), got {point_data_dict['thickness'].shape}"
+    )
 
 
 def test_load_vtp_file_displacements(simple_vtp_file):
     """Test that displacements are correctly applied."""
     pos_raw, _, _ = load_vtp_file(simple_vtp_file)
-    
+
     # First timestep should be reference coords (displacement = 0)
-    expected_t0 = np.array([
-        [0, 0, 0],
-        [1, 0, 0],
-        [0, 1, 0],
-        [1, 1, 0],
-    ])
+    expected_t0 = np.array(
+        [
+            [0, 0, 0],
+            [1, 0, 0],
+            [0, 1, 0],
+            [1, 1, 0],
+        ]
+    )
     np.testing.assert_array_almost_equal(pos_raw[0], expected_t0, decimal=5)
-    
+
     # Second timestep should include displacement
     expected_t1 = expected_t0 + np.array([[1, 0, 0]] * 4)
     np.testing.assert_array_almost_equal(pos_raw[1], expected_t1, decimal=5)
-    
+
     # Third timestep
     expected_t2 = expected_t0 + np.array([[2, 0, 0]] * 4)
     np.testing.assert_array_almost_equal(pos_raw[2], expected_t2, decimal=5)
 
 
 def test_extract_mesh_connectivity():
     """Test mesh connectivity extraction from PolyData."""
-    points = np.array([
-        [0, 0, 0],
-        [1, 0, 0],
-        [1, 1, 0],
-        [0, 1, 0],
-    ])
-    
+    points = np.array(
+        [
+            [0, 0, 0],
+            [1, 0, 0],
+            [1, 1, 0],
+            [0, 1, 0],
+        ]
+    )
+
     # Create a single quad
     faces = np.array([4, 0, 1, 2, 3])
     poly = pv.PolyData(points, faces, force_float=False)
-    
+
     connectivity = extract_mesh_connectivity_from_polydata(poly)
-    
+
     assert len(connectivity) == 1, f"Expected 1 cell, got {len(connectivity)}"
     assert len(connectivity[0]) == 4, f"Expected 4 vertices, got {len(connectivity[0])}"
-    assert connectivity[0] == [0, 1, 2, 3], f"Expected [0, 1, 2, 3], got {connectivity[0]}"
+    assert connectivity[0] == [0, 1, 2, 3], (
+        f"Expected [0, 1, 2, 3], got {connectivity[0]}"
+    )
 
 
 def test_build_edges_from_mesh_connectivity():
     """Test edge building from mesh connectivity."""
     # Single quad: should produce 4 edges
     mesh_connectivity = [[0, 1, 2, 3]]
     edges = build_edges_from_mesh_connectivity(mesh_connectivity)
-    
+
     expected_edges = {(0, 1), (1, 2), (2, 3), (0, 3)}
     assert edges == expected_edges, f"Expected {expected_edges}, got {edges}"
 
 
 def test_point_data_extraction(simple_vtp_file):
     """Test that non-displacement point data is extracted correctly."""
     _, _, point_data_dict = load_vtp_file(simple_vtp_file)
-    
+
     # Should have thickness
-    assert 'thickness' in point_data_dict, "Thickness not in point_data_dict"
-    
+    assert "thickness" in point_data_dict, "Thickness not in point_data_dict"
+
     # Should NOT have displacement fields
-    assert 'displacement_t0.000' not in point_data_dict, "Displacement fields should not be in point_data_dict"
-    assert 'displacement_t0.005' not in point_data_dict, "Displacement fields should not be in point_data_dict"
-    
+    assert "displacement_t0.000" not in point_data_dict, (
+        "Displacement fields should not be in point_data_dict"
+    )
+    assert "displacement_t0.005" not in point_data_dict, (
+        "Displacement fields should not be in point_data_dict"
+    )
+
     # Check thickness values
     expected_thickness = np.array([1, 1, 1, 1], dtype=np.uint8)
-    np.testing.assert_array_almost_equal(point_data_dict['thickness'], expected_thickness, decimal=5)
+    np.testing.assert_array_almost_equal(
+        point_data_dict["thickness"], expected_thickness, decimal=5
+    )
 
 
 def test_missing_displacement_fields():
@@ -172,12 +201,12 @@ def test_missing_displacement_fields():
     points = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0]])
     faces = np.array([3, 0, 1, 2])
     mesh = pv.PolyData(points, faces, force_float=False)
-    
-    with tempfile.NamedTemporaryFile(suffix='.vtp', delete=False) as f:
+
+    with tempfile.NamedTemporaryFile(suffix=".vtp", delete=False) as f:
         temp_path = f.name
-    
+
     mesh.save(temp_path)
-    
+
     try:
         with pytest.raises(ValueError, match="No displacement fields found"):
             load_vtp_file(temp_path)
@@ -189,5 +218,5 @@ def test_empty_mesh_connectivity():
     """Test edge building with empty connectivity."""
     mesh_connectivity = []
     edges = build_edges_from_mesh_connectivity(mesh_connectivity)
-    
+
     assert len(edges) == 0, f"Expected 0 edges for empty connectivity, got {len(edges)}"

examples/structural_mechanics/crash/vtp_reader.py

@@ -105,13 +105,13 @@ def natural_key(name):
 
     pos_raw = np.stack(pos_list, axis=0)
     mesh_connectivity = extract_mesh_connectivity_from_polydata(poly)
-    
+
     # Extract all other point data fields (not displacement fields)
     point_data_dict = {}
     for name in poly.point_data.keys():
         if not name.startswith("displacement_"):
             point_data_dict[name] = np.asarray(poly.point_data[name])
-    
+
     return pos_raw, mesh_connectivity, point_data_dict
 
 
@@ -208,7 +208,7 @@ def process_vtp_data(data_dir, num_samples=2, write_vtp=False, logger=None):
             write_vtp=write_vtp,
             logger=logger,
         )
-        
+
         # Create record with coords and all other point data fields
         record = {"coords": mesh_pos_all}
         record.update(point_data_dict)  # Add thickness and any other fields