add support for stress and plastic strain

examples/structural_mechanics/crash/crash_data_processors.py

@@ -72,6 +72,26 @@ def load_d3plot_data(data_path: str):
     return coords, pos_raw, mesh_connectivity, part_ids, actual_part_ids
 
 
+def load_d3plot_element_fields(data_path: str):
+    """
+    Load optional element-level fields from a d3plot file.
+    Returns (element_shell_stress, element_shell_effective_plastic_strain),
+    each possibly None if absent in the dataset.
+    """
+    dp = D3plot(data_path)
+    element_shell_stress = (
+        dp.arrays[ArrayType.element_shell_stress]
+        if ArrayType.element_shell_stress in dp.arrays
+        else None
+    )  # expected shape: (T, E, 2, 6)
+    element_shell_effective_plastic_strain = (
+        dp.arrays[ArrayType.element_shell_effective_plastic_strain]
+        if ArrayType.element_shell_effective_plastic_strain in dp.arrays
+        else None
+    )  # expected shape: (T, E, 2)
+    return element_shell_stress, element_shell_effective_plastic_strain
+
+
 def find_k_file(run_dir: Path) -> Optional[Path]:
     """Find .k file in run directory.
 
@@ -228,3 +248,73 @@ def compute_node_thickness(
             node_thickness[i] /= node_thickness_count[i]
 
     return node_thickness
+
+
+def reduce_shell_layers_scalar(x: np.ndarray) -> np.ndarray:
+    """
+    Average over through-thickness layers (axis=2) for scalar shell fields.
+    Input shape: (T, E, 2) -> Output shape: (T, E)
+    """
+    return np.nanmean(x, axis=2)
+
+
+def reduce_shell_layers_stress_voigt(sig: np.ndarray) -> np.ndarray:
+    """
+    Average over through-thickness layers (axis=2) for stress in Voigt form.
+    Input shape: (T, E, 2, 6) -> Output shape: (T, E, 6)
+    """
+    return np.nanmean(sig, axis=2)
+
+
+def von_mises_from_voigt(sig: np.ndarray) -> np.ndarray:
+    """
+    Compute von Mises stress from Voigt components [sx, sy, sz, txy, tyz, tzx].
+    Input shape: (T, E, 6) -> Output shape: (T, E)
+    """
+    sx = sig[..., 0]
+    sy = sig[..., 1]
+    sz = sig[..., 2]
+    txy = sig[..., 3]
+    tyz = sig[..., 4]
+    tzx = sig[..., 5]
+    j2 = 0.5 * ((sx - sy) ** 2 + (sy - sz) ** 2 + (sz - sx) ** 2) + 3.0 * (
+        txy**2 + tyz**2 + tzx**2
+    )
+    # Numerically safe clamp
+    return np.sqrt(np.maximum(j2 * 2.0 / 3.0, 0.0))
+
+
+def compute_element_and_node_fields(
+    mesh_connectivity: np.ndarray,
+    element_shell_stress: Optional[np.ndarray] = None,
+    element_shell_effective_plastic_strain: Optional[np.ndarray] = None,
+    compute_von_mises: bool = True,
+) -> dict:
+    """
+    Convenience function to compute reduced element fields.
+    Inputs can be None; outputs will carry None accordingly.
+    Returns a dict with key 'element' containing arrays or None.
+    """
+    out = {"element": {}}
+
+    # Effective plastic strain (scalar)
+    if element_shell_effective_plastic_strain is not None:
+        eps_elem = reduce_shell_layers_scalar(element_shell_effective_plastic_strain)  # (T, E)
+        out["element"]["effective_plastic_strain"] = eps_elem
+    else:
+        out["element"]["effective_plastic_strain"] = None
+
+    # Stress
+    if element_shell_stress is not None:
+        stress_elem_voigt = reduce_shell_layers_stress_voigt(element_shell_stress)  # (T, E, 6)
+        out["element"]["stress_voigt"] = stress_elem_voigt
+        if compute_von_mises:
+            vm_elem = von_mises_from_voigt(stress_elem_voigt)  # (T, E)
+            out["element"]["stress_vm"] = vm_elem
+        else:
+            out["element"]["stress_vm"] = None
+    else:
+        out["element"]["stress_voigt"] = None
+        out["element"]["stress_vm"] = None
+
+    return out

examples/structural_mechanics/crash/data_sources.py

@@ -26,6 +26,7 @@
     compute_node_thickness,
     find_k_file,
     load_d3plot_data,
+    load_d3plot_element_fields,
     parse_k_file,
 )
 from schemas import CrashExtractedDataInMemory, CrashMetadata
@@ -77,6 +78,11 @@ def read_file(self, run_id: str) -> CrashExtractedDataInMemory:
             load_d3plot_data(str(d3plot_path))
         )
 
+        # Optional element-level fields (stress, effective plastic strain)
+        element_shell_stress, element_shell_effective_plastic_strain = (
+            load_d3plot_element_fields(str(d3plot_path))
+        )
+
         # Parse .k file for thickness
         k_file_path = find_k_file(run_dir=run_dir)
         node_thickness = np.zeros(len(coords))
@@ -101,6 +107,8 @@ def read_file(self, run_id: str) -> CrashExtractedDataInMemory:
             pos_raw=pos_raw,
             mesh_connectivity=mesh_connectivity,
             node_thickness=node_thickness,
+            element_shell_stress=element_shell_stress,
+            element_shell_effective_plastic_strain=element_shell_effective_plastic_strain,
         )
 
     def _get_output_path(self, filename: str) -> Path:
@@ -241,6 +249,17 @@ def _write_impl_temp_file(
 
             mesh.point_data[field_name] = displacement
 
+            # Optional: add per-timestep cell scalar fields
+            n_cells = int(mesh.n_cells)
+            if getattr(data, "filtered_element_effective_plastic_strain", None) is not None:
+                eps_elem_t = data.filtered_element_effective_plastic_strain[t, :]
+                if eps_elem_t.shape[0] == n_cells:
+                    mesh.cell_data[f"cell_effective_plastic_strain_{time_str}"] = eps_elem_t
+            if getattr(data, "filtered_element_stress_vm", None) is not None:
+                vm_elem_t = data.filtered_element_stress_vm[t, :]
+                if vm_elem_t.shape[0] == n_cells:
+                    mesh.cell_data[f"cell_stress_vm_{time_str}"] = vm_elem_t
+
         # Save the single VTP file
         mesh.save(output_path)
 
@@ -472,6 +491,27 @@ def _write_impl_temp_file(
             compressors=(self.compressor,),
         )
 
+        # Optionally write processed element/node fields if available
+        def _maybe_write(name: str, arr: np.ndarray) -> None:
+            if arr is None:
+                return
+            arr32 = arr.astype(np.float32)
+            chunks = self._calculate_chunks(arr32)
+            root.create_array(
+                name=name,
+                data=arr32,
+                chunks=chunks,
+                compressors=(self.compressor,),
+            )
+
+        _maybe_write("element_stress_voigt", data.filtered_element_stress_voigt)
+        _maybe_write("element_stress_vm", data.filtered_element_stress_vm)
+        _maybe_write(
+            "element_effective_plastic_strain",
+            data.filtered_element_effective_plastic_strain,
+        )
+        # Node-level stress/strain intentionally not written (cell-only per request)
+
         # Add some statistics as metadata
         root.attrs["thickness_min"] = float(np.min(data.filtered_node_thickness))
         root.attrs["thickness_max"] = float(np.max(data.filtered_node_thickness))

examples/structural_mechanics/crash/data_transformations.py

@@ -19,7 +19,11 @@
 from typing import Callable, Optional
 
 import numpy as np
-from crash_data_processors import build_edges_from_mesh_connectivity, compute_node_type
+from crash_data_processors import (
+    build_edges_from_mesh_connectivity,
+    compute_element_and_node_fields,
+    compute_node_type,
+)
 from schemas import CrashExtractedDataInMemory
 
 from physicsnemo_curator.etl.data_transformations import DataTransformation
@@ -79,10 +83,12 @@ def transform(
 
         # Step 3: Remap mesh connectivity
         filtered_mesh_connectivity = []
-        for cell in data.mesh_connectivity:
+        kept_element_indices = []
+        for e_idx, cell in enumerate(data.mesh_connectivity):
             filtered_cell = [node_map[n] for n in cell if n in node_map]
             if len(filtered_cell) >= 3:
                 filtered_mesh_connectivity.append(filtered_cell)
+                kept_element_indices.append(e_idx)
 
         # Step 4: Compact to contiguous indices (reuse your logic)
         used = np.unique(
@@ -103,6 +109,38 @@ def transform(
             ]
             num_kept = filtered_pos_raw.shape[1]
 
+        # Optional: filter element fields to kept elements
+        filtered_element_stress_voigt = None
+        filtered_element_stress_vm = None
+        filtered_element_effective_plastic_strain = None
+
+        has_epsp = data.element_shell_effective_plastic_strain is not None
+        has_stress = data.element_shell_stress is not None
+        if has_epsp or has_stress:
+            epsp_kept = (
+                data.element_shell_effective_plastic_strain[:, kept_element_indices, :]
+                if has_epsp
+                else None
+            )
+            stress_kept = (
+                data.element_shell_stress[:, kept_element_indices, :, :]
+                if has_stress
+                else None
+            )
+            # Compute reduced element/node fields on filtered mesh
+            fields = compute_element_and_node_fields(
+                filtered_mesh_connectivity,
+                element_shell_stress=stress_kept,
+                element_shell_effective_plastic_strain=epsp_kept,
+                compute_von_mises=True,
+            )
+            # Unpack
+            filtered_element_effective_plastic_strain = fields["element"][
+                "effective_plastic_strain"
+            ]
+            filtered_element_stress_voigt = fields["element"]["stress_voigt"]
+            filtered_element_stress_vm = fields["element"]["stress_vm"]
+
         # Step 6: Build edges
         edges = build_edges_from_mesh_connectivity(filtered_mesh_connectivity)
 
@@ -115,11 +153,16 @@ def transform(
         data.pos_raw = None
         data.mesh_connectivity = None
         data.node_thickness = None
+        data.element_shell_stress = None
+        data.element_shell_effective_plastic_strain = None
 
         return CrashExtractedDataInMemory(
             metadata=data.metadata,
             filtered_pos_raw=filtered_pos_raw,
             filtered_mesh_connectivity=filtered_mesh_connectivity,
             filtered_node_thickness=filtered_node_thickness,
             edges=edges,
+            filtered_element_stress_voigt=filtered_element_stress_voigt,
+            filtered_element_stress_vm=filtered_element_stress_vm,
+            filtered_element_effective_plastic_strain=filtered_element_effective_plastic_strain,
         )

examples/structural_mechanics/crash/schemas.py

@@ -37,6 +37,7 @@ class CrashExtractedDataInMemory:
 
     Version history:
     - 1.0: Initial version with expected data fields.
+    - 1.1: Added optional element stress/strain raw fields and processed element fields.
     """
 
     # Metadata
@@ -46,9 +47,16 @@ class CrashExtractedDataInMemory:
     pos_raw: np.ndarray = None
     mesh_connectivity: np.ndarray = None
     node_thickness: np.ndarray = None
+    element_shell_stress: np.ndarray = None  # (T, E, 2, 6) or None
+    element_shell_effective_plastic_strain: np.ndarray = None  # (T, E, 2) or None
 
     # Processed data
     filtered_pos_raw: np.ndarray = None
     filtered_mesh_connectivity: np.ndarray = None
     filtered_node_thickness: np.ndarray = None
     edges: np.ndarray = None
+
+    # Processed optional fields
+    filtered_element_stress_voigt: np.ndarray = None  # (T, E, 6)
+    filtered_element_stress_vm: np.ndarray = None  # (T, E)
+    filtered_element_effective_plastic_strain: np.ndarray = None  # (T, E)