From ef5763c57b320b392942f3b941fc7bc3ebe8de34 Mon Sep 17 00:00:00 2001 From: Tyagi Date: Tue, 3 Mar 2026 16:59:48 +1100 Subject: [PATCH] docs(beginner): add create_xdmf guide for PETSc HDF5->XDMF workflow explain vertex/cell compatibility mapping, KDTree remap, and parallel write behavior document ParaView tensor requirements (2D/3D -> 3x3, 9 components) and PyVista post-processing alternative --- docs/beginner/create_xdmf.md | 324 +++++++++++++++ docs/beginner/index.md | 10 +- .../discretisation/discretisation_mesh.py | 371 ++++++++++++++++-- 3 files changed, 675 insertions(+), 30 deletions(-) create mode 100644 docs/beginner/create_xdmf.md diff --git a/docs/beginner/create_xdmf.md b/docs/beginner/create_xdmf.md new file mode 100644 index 000000000..f5f141bc4 --- /dev/null +++ b/docs/beginner/create_xdmf.md @@ -0,0 +1,324 @@ +--- +title: "Create XDMF from PETSc HDF5 Output" +--- + +# Create XDMF from PETSc HDF5 Output + +This guide explains how to generate ParaView-ready XDMF/HDF5 outputs with `mesh.write_timestep(create_xdmf=True)`, why this is needed with newer PETSc HDF5 layouts, and how tensors are converted for correct ParaView interpretation. + +This note explains the output-format issue seen with newer PETSc versions and the fix now available in `Mesh.write_timestep()`. + +## 1. What new PETSc writes to HDF5 + +With newer PETSc HDF5 viewer behavior, variable files commonly contain only: + +- `/fields/` +- `/fields/coordinates` + +Instead of relying on a fixed file example, you can verify this directly: + +```python +import h5py +import underworld3 as uw + +mesh = uw.meshing.StructuredQuadBox(elementRes=(4, 4), minCoords=(0, 0), maxCoords=(1, 1)) +v = uw.discretisation.MeshVariable("V", mesh, mesh.dim, degree=1, continuous=True) +p = uw.discretisation.MeshVariable("P", mesh, 1, degree=0, continuous=False) + +with uw.synchronised_array_update(): + v.array[:, 0, :] = 1.0 + p.array[:, 0, 0] = 2.0 + +# Case A: HDF5 only (no XDMF compatibility groups) +mesh.write_timestep( + "check_a", + index=0, + outputPath=".", + meshVars=[v, p], + create_xdmf=False, +) + +# Case B: create XDMF + compatibility groups +mesh.write_timestep( + "check_b", + index=0, + outputPath=".", + meshVars=[v, p], + create_xdmf=True, +) + +for fname in ["check_a.mesh.V.00000.h5", "check_b.mesh.V.00000.h5"]: + print(f"\n{fname}") + with h5py.File(fname, "r") as h5: + def _show(name, obj): + if isinstance(obj, h5py.Dataset): + print(f" {name} {obj.shape}") + else: + print(f" {name}/") + h5.visititems(_show) +``` + +Older workflows also had explicit compatibility groups such as: + +- `/vertex_fields/...` +- `/cell_fields/...` + +which many existing XDMF templates expected. + +## 2. Why `/fields` alone is not sufficient for robust XDMF + +XDMF needs a clear **data center**: + +- node-centered data must match mesh vertices (`Nvertices`) +- cell-centered data must match mesh elements (`Ncells`) + +But `/fields` can store high-order DOF layouts that do not equal vertex count or cell count (for example packed element-point coordinates/values). +If XDMF points `Center="Node"` data to such arrays, ParaView reads mismatched lengths and fails or misinterprets arrays. + +So, `/fields` is useful raw data, but it is not always directly visualization-ready for a mesh topology/geometry pair. + +## 3. What changed in `write_timestep()` and why + +`Mesh.write_timestep()` now supports: + +```python +mesh.write_timestep(..., create_xdmf=True) +``` + +When `create_xdmf=True`: + +1. It writes compatibility datasets per variable: + - node-like vars -> `/vertex_fields/coordinates` and `/vertex_fields/_` + - cell-like vars -> `/cell_fields/_` +2. It writes XDMF that references these compatibility datasets. + +This gives XDMF arrays that match mesh vertex/cell counts, which ParaView expects. + +Implementation notes: + +- continuous, non-degree-0 variables are treated as node-like +- discontinuous or degree-0 variables are treated as cell-like +- high-order `/fields` layouts are unpacked/remapped to vertex/cell-compatible arrays + +### How vertex vs cell field is chosen + +For each mesh variable during `write_timestep(..., create_xdmf=True)`: + +- **Cell field** if: + - variable is discontinuous (`continuous == False`), or + - variable polynomial degree is 0 (`degree == 0`) +- **Vertex field** otherwise. + +In code terms, this is equivalent to: + +```python +is_cell = (not var.continuous) or (var.degree == 0) +``` + +Then: + +- `is_cell == True` -> write `/cell_fields/_` and XDMF `Center="Cell"` +- `is_cell == False` -> write `/vertex_fields/_` and XDMF `Center="Node"` + +### How `/fields` data is remapped (KDTree) + +When a variable is vertex-centered but `/fields` does not already match mesh vertex count: + +1. Source coordinates/values are read from `/fields/coordinates` and `/fields/`. +2. Packed high-order layouts are unpacked into point-wise coordinate/value rows. +3. A nearest-neighbor search is done with `uw.kdtree.KDTree` to map source points to mesh vertices. +4. The mapped values are written to `/vertex_fields/_`. + +For cell-centered variables, values are written into `/cell_fields/_` using row slices. + +### Parallel execution details + +This remapping/writing path is parallel-aware: + +- work is partitioned by MPI rank (each rank owns a slice of vertices / rows) +- each rank computes mapping for its local slice +- output datasets are written by per-rank slices (no overlapping writes) +- with MPI-enabled `h5py`, HDF5 writes use parallel I/O; otherwise a serial fallback is used + +So the compatibility-field generation is computed and written in parallel when MPI-HDF5 is available. + +## 4. Tensor representation for ParaView (2D and 3D) + +ParaView expects `Attribute Type="Tensor"` data as 9 components per point/cell (flattened `3x3`). + +ParaView expects tensors in a `3x3` format, even for 2D problems. + +### The key idea + +A 2D tensor + +```text +[ s_xx s_xy ] +[ s_yx s_yy ] +``` + +must be embedded into a 3D tensor + +```text +[ s_xx s_xy 0 ] +[ s_yx s_yy 0 ] +[ 0 0 0 ] +``` + +ParaView works internally with 3D tensors, so writing only 2D tensor components is not enough for full tensor-aware filters. + +### Expected tensor format in XDMF / HDF5 + +- XDMF attribute: `Type="Tensor"` +- dataset shape: `(N, 9)` +- component order (row-major): + - `[Txx, Txy, Txz, Tyx, Tyy, Tyz, Tzx, Tzy, Tzz]` + +### 2D tensors + +For 2D tensors, values are embedded into 3D by zero-filling Z terms. + +Example (full 2D tensor with 4 comps): + +- input components interpreted as in-plane terms +- output `9` components with `Txz, Tyz, Tzx, Tzy, Tzz = 0` + +Example (2D symmetric with 3 comps): + +- `Txy == Tyx` +- output still written as `9` components. + +### 3D tensors + +- full tensors: `(N, 9)` passed through +- symmetric tensors with 6 comps are expanded into `(N, 9)` with mirrored off-diagonals + +### What we changed in Underworld3 + +When `mesh.write_timestep(..., create_xdmf=True)` is used: + +1. Node/cell compatibility arrays are written to `/vertex_fields` and `/cell_fields`. +2. Tensor-like mesh variables are converted to ParaView-ready `9` components before writing: + - 2D full tensor (`4` comps) -> embedded `3x3` (`9` comps) + - 2D symmetric (`3` comps) -> embedded `3x3` (`9` comps, mirrored shear terms) + - 3D full (`9` comps) -> unchanged + - 3D symmetric (`6` comps) -> expanded to full `3x3` (`9` comps) +3. XDMF attributes for these variables are emitted as `Type="Tensor"` and reference the `9`-component datasets. + +This is why tensor variables now appear in HDF5/XDMF as 9-component arrays for ParaView compatibility. + +## 5. Post-processing without XDMF or extra HDF5 groups + +If you want to avoid writing `/vertex_fields` and `/cell_fields` (to reduce file size and write overhead), you can: + +- save only raw `/fields` data (`create_xdmf=False`) +- do mapping in Python during post-processing +- attach mapped arrays directly to a `pyvista` mesh in memory + +This avoids storing additional redundant datasets in HDF5. + +### 5.1 Write raw HDF5 only + +```python +mesh.write_timestep( + "output", + index=0, + outputPath=".", + meshVars=[v, p, stress], + create_xdmf=False, # no xdmf, no /vertex_fields or /cell_fields +) +``` + +### 5.2 Load mesh + raw field data and map in PyVista + +```python +import h5py +import numpy as np +import pyvista as pv +from scipy.spatial import cKDTree + + +def _flatten_fields_layout(field_values, field_coords, dim): + """Handle packed high-order field layout into point-wise rows.""" + if field_values.ndim == 1: + field_values = field_values.reshape(-1, 1) + + if field_coords.shape[1] == dim: + return field_values, field_coords + + if field_coords.shape[1] % dim != 0: + raise RuntimeError(f"Cannot unpack coords shape {field_coords.shape} for dim={dim}") + + dof_per_row = field_coords.shape[1] // dim + coords = field_coords.reshape(-1, dim) + + if field_values.shape[1] == dof_per_row: + values = field_values.reshape(-1, 1) + elif field_values.shape[1] % dof_per_row == 0: + ncomp = field_values.shape[1] // dof_per_row + values = field_values.reshape(field_values.shape[0], dof_per_row, ncomp).reshape(-1, ncomp) + else: + raise RuntimeError( + f"Cannot unpack values shape {field_values.shape} with dof_per_row={dof_per_row}" + ) + + return values, coords + + +def load_h5_field_to_pvmesh(pvmesh, mesh_h5, var_h5, field_name, location="auto"): + """ + Attach raw /fields data to pvmesh as point_data or cell_data. + location: 'point', 'cell', or 'auto' + """ + with h5py.File(mesh_h5, "r") as mh: + verts = mh["geometry/vertices"][()] + nverts = verts.shape[0] + ncells = pvmesh.n_cells + dim = verts.shape[1] + + with h5py.File(var_h5, "r") as vh: + values = vh["fields"][field_name][()] + coords = vh["fields"]["coordinates"][()] + + if values.ndim == 1: + values = values.reshape(-1, 1) + + # Direct matches: already node or cell sized + if location in ("auto", "point") and values.shape[0] == nverts: + pvmesh.point_data[field_name] = values + return "point" + if location in ("auto", "cell") and values.shape[0] == ncells: + pvmesh.cell_data[field_name] = values + return "cell" + + # Packed/high-order case: map by nearest coordinate + unpacked_values, unpacked_coords = _flatten_fields_layout(values, coords, dim) + tree = cKDTree(unpacked_coords) + _, idx = tree.query(verts, k=1) + pvmesh.point_data[field_name] = unpacked_values[idx, :] + return "point" +``` + +### 5.3 Example usage + +```python +mesh_h5 = "output.mesh.00000.h5" +vel_h5 = "output.mesh.Velocity.00000.h5" +prs_h5 = "output.mesh.Pressure.00000.h5" + +# Build pvmesh from the mesh file (reader choice depends on your workflow) +pvmesh = pv.read(mesh_h5) + +load_h5_field_to_pvmesh(pvmesh, mesh_h5, vel_h5, "Velocity", location="auto") +load_h5_field_to_pvmesh(pvmesh, mesh_h5, prs_h5, "Pressure", location="auto") + +# Continue in-memory analysis/plotting with pyvista +plotter = pv.Plotter() +plotter.add_mesh(pvmesh, scalars="Pressure") +plotter.show() +``` + +Notes: + +- This workflow is ideal for Python-based analysis. +- For direct ParaView file-open workflows, `create_xdmf=True` remains the simpler option. diff --git a/docs/beginner/index.md b/docs/beginner/index.md index 64b36bce1..b4b6785dc 100644 --- a/docs/beginner/index.md +++ b/docs/beginner/index.md @@ -29,7 +29,12 @@ Configure your models for notebooks and command-line execution. **[→ Parameters Guide](parameters.md)** -### 4. Interactive Tutorials +### 4. XDMF / HDF5 Compatibility +Understand the PETSc output-format change and how to write ParaView-ready XDMF/HDF5. + +**[→ XDMF / HDF5 Compatibility Guide](create_xdmf.md)** + +### 5. Interactive Tutorials Work through hands-on notebooks covering all core concepts. **[→ Start Tutorials](tutorials/Notebook_Index.ipynb)** @@ -88,6 +93,7 @@ Once comfortable with the basics, explore: installation quickstart parameters +create_xdmf tutorials/Notebook_Index tutorials/1-Meshes tutorials/2-Variables @@ -104,4 +110,4 @@ tutorials/12-Units_System tutorials/13-Scaling-problems-with-physical-units tutorials/14-Timestepping-with-physical-units tutorials/15-Thermal-convection-with-units -``` \ No newline at end of file +``` diff --git a/src/underworld3/discretisation/discretisation_mesh.py b/src/underworld3/discretisation/discretisation_mesh.py index 5fb4f5279..cc2df2c0f 100644 --- a/src/underworld3/discretisation/discretisation_mesh.py +++ b/src/underworld3/discretisation/discretisation_mesh.py @@ -1750,13 +1750,30 @@ def write_timestep( meshVars: Optional[list] = [], swarmVars: Optional[list] = [], meshUpdates: bool = False, + create_xdmf: bool = False, + vertex_cell_chunk_size: Optional[int] = None, ): """ - Write the selected mesh, variables and swarm variables (as proxies) for later visualisation. - An xdmf file is generated and the overall package can then be read by paraview or pyvista. - Vertex values (on the mesh points) are stored for all variables regardless of their interpolation order - """ + Write mesh and selected variables for visualisation output. + + This writes: + - one mesh HDF5 file (shared/static or per-step, depending on ``meshUpdates``) + - one HDF5 file per mesh variable + - optional proxy files for swarm variables + - optional XDMF file linking all output files + + When ``create_xdmf=True``, variable files also include + compatibility groups: + - ``/vertex_fields`` for node-like variables + - ``/cell_fields`` for cell-like variables + + and ``checkpoint_xdmf()`` is called on rank 0. + + ``vertex_cell_chunk_size`` controls source-point chunking for + node-like remapping. Default is automatic. Set a positive integer + to manually override the auto-chosen chunk size. + """ options = PETSc.Options() options.setValue("viewer_hdf5_sp_output", True) options.setValue("viewer_hdf5_collective", False) @@ -1789,19 +1806,30 @@ def write_timestep( self.write(mesh_file) else: - self.write(output_base_name + f".mesh.{index:05}.h5") + mesh_file = output_base_name + f".mesh.{index:05}.h5" + self.write(mesh_file) if meshVars is not None: for var in meshVars: save_location = output_base_name + f".mesh.{var.clean_name}.{index:05}.h5" var.write(save_location) + if create_xdmf: + is_cell = (not getattr(var, "continuous")) or (getattr(var, "degree") == 0) + _write_vertex_cell_groups( + mesh_h5=mesh_file, + var_h5=save_location, + var_name=var.clean_name, + var_obj=var, + is_cell=is_cell, + vertex_cell_chunk_size=vertex_cell_chunk_size, + ) if swarmVars is not None: for svar in swarmVars: save_location = output_base_name + f".proxy.{svar.clean_name}.{index:05}.h5" svar.write_proxy(save_location) - if uw.mpi.rank == 0: + if create_xdmf and uw.mpi.rank == 0: checkpoint_xdmf( output_base_name, meshUpdates, @@ -3127,6 +3155,274 @@ def mesh_update_callback(array, change_context): ## Simplified to allow us to decide how we want to checkpoint +def _write_vertex_cell_groups( + mesh_h5: str, + var_h5: str, + var_name: str, + var_obj=None, + is_cell: bool = False, + vertex_cell_chunk_size: Optional[int] = None, +): + """ + Write vertex/cell compatibility groups into a variable file. + + Source layout: + - ``/fields/`` + - ``/fields/coordinates`` (for node-style mapping) + + Output layout: + - cell-like variables -> ``/cell_fields/_`` + - node-like variables -> ``/vertex_fields/coordinates`` and + ``/vertex_fields/_`` + + Tensor-like node variables are written as flattened 3x3 tensors + (N, 9) in ``/vertex_fields/_`` for ParaView + compatibility. + """ + + import h5py + import numpy as np + import underworld3 as uw + + comm = uw.mpi.comm + rank = uw.mpi.rank + size = uw.mpi.size + + def _chunk_bounds(total_count, proc_rank, proc_size): + base = total_count // proc_size + rem = total_count % proc_size + start = proc_rank * base + min(proc_rank, rem) + stop = start + base + (1 if proc_rank < rem else 0) + return start, stop + + def _auto_chunk_size( + total_source_rows: int, + coordinate_columns: int, + field_columns: int, + dtype_itemsize: int, + ) -> int: + # Approximate per-rank transient memory for source coordinates + values. + target_bytes_per_rank = 192 * 1024 * 1024 + bytes_per_source_row = max( + dtype_itemsize * max(1, coordinate_columns + field_columns), 8 + ) + chunk_from_memory = max(10_000, target_bytes_per_rank // bytes_per_source_row) + + # Keep enough chunks to spread work and reduce tail effects as rank count grows. + chunk_from_ranks = max( + 10_000, (total_source_rows + (size * 8) - 1) // max(size * 8, 1) + ) + + return int(min(total_source_rows, max(10_000, min(chunk_from_memory, chunk_from_ranks)))) + + def _flatten_source_points(source_coordinate_chunk, source_field_chunk, dim): + if source_coordinate_chunk.ndim != 2: + raise RuntimeError( + f"Unsupported coordinate rank for {var_name}: {source_coordinate_chunk.ndim}" + ) + + if source_coordinate_chunk.shape[1] == dim: + coords = source_coordinate_chunk + values = source_field_chunk + if values.ndim == 1: + values = values.reshape(-1, 1) + return coords, values + + if source_coordinate_chunk.shape[1] % dim != 0: + raise RuntimeError( + f"Coordinate dimension mismatch: mesh={dim}, field={source_coordinate_chunk.shape[1]}" + ) + + dof_per_row = source_coordinate_chunk.shape[1] // dim + coords = source_coordinate_chunk.reshape(-1, dim) + values = source_field_chunk + + if values.ndim == 1: + values = values.reshape(-1, 1) + values = np.repeat(values, dof_per_row, axis=0) + return coords, values + + if values.shape[1] == dof_per_row: + values = values.reshape(-1, 1) + return coords, values + + if values.shape[1] % dof_per_row == 0: + ncomp_eff = values.shape[1] // dof_per_row + values = values.reshape(values.shape[0], dof_per_row, ncomp_eff).reshape(-1, ncomp_eff) + return coords, values + + raise RuntimeError( + f"Cannot unpack field layout for {var_name}: " + f"values={values.shape}, coords={source_coordinate_chunk.shape}, dim={dim}" + ) + + is_tensor = False + if var_obj is not None and hasattr(var_obj, "vtype"): + is_tensor = var_obj.vtype in ( + uw.VarType.TENSOR, + uw.VarType.SYM_TENSOR, + uw.VarType.MATRIX, + ) + + parallel_h5 = bool(getattr(h5py.get_config(), "mpi", False)) + if parallel_h5: + open_kwargs = {"driver": "mpio", "comm": comm} + else: + # Fallback: single-writer path when MPI-HDF5 is unavailable. + if rank != 0: + return + open_kwargs = {} + + with h5py.File(mesh_h5, "r", **open_kwargs) as mesh_file: + mesh_vertices_dataset = mesh_file["geometry/vertices"] + num_vertices = mesh_vertices_dataset.shape[0] + dim = mesh_vertices_dataset.shape[1] + local_vertex_start, local_vertex_stop = _chunk_bounds(num_vertices, rank, size) + local_vertices = mesh_vertices_dataset[local_vertex_start:local_vertex_stop, :] + + with h5py.File(var_h5, "r+", **open_kwargs) as variable_file: + if "fields" not in variable_file: + raise RuntimeError(f"Missing group '/fields' in {var_h5}") + if var_name not in variable_file["fields"]: + raise RuntimeError(f"Missing dataset '/fields/{var_name}' in {var_h5}") + + field_dataset = variable_file["fields"][var_name] + field_dataset_dtype = field_dataset.dtype + output_dataset_name = f"{var_name}_{var_name}" + + if is_cell: + num_rows = field_dataset.shape[0] + local_row_start, local_row_stop = _chunk_bounds(num_rows, rank, size) + local_field_chunk = field_dataset[local_row_start:local_row_stop] + if local_field_chunk.ndim == 1: + local_field_chunk = local_field_chunk.reshape(-1, 1) + output_components = local_field_chunk.shape[1] + + cell_group = variable_file.require_group("cell_fields") + if rank == 0 and output_dataset_name in cell_group: + del cell_group[output_dataset_name] + if parallel_h5: + comm.Barrier() + output_dataset = cell_group.require_dataset( + output_dataset_name, shape=(num_rows, output_components), dtype=field_dataset_dtype + ) + output_dataset[local_row_start:local_row_stop, :] = local_field_chunk + if parallel_h5: + comm.Barrier() + return + + if "coordinates" not in variable_file["fields"]: + raise RuntimeError(f"Missing dataset '/fields/coordinates' in {var_h5}") + + coordinate_dataset = variable_file["fields"]["coordinates"] + num_source_rows = coordinate_dataset.shape[0] + coordinate_columns = coordinate_dataset.shape[1] + if field_dataset.ndim == 1: + field_columns = 1 + else: + field_columns = field_dataset.shape[1] + + if vertex_cell_chunk_size is None or vertex_cell_chunk_size <= 0: + effective_chunk_size = _auto_chunk_size( + total_source_rows=num_source_rows, + coordinate_columns=coordinate_columns, + field_columns=field_columns, + dtype_itemsize=field_dataset_dtype.itemsize, + ) + else: + effective_chunk_size = int(vertex_cell_chunk_size) + + local_vertex_count = max(local_vertex_stop - local_vertex_start, 0) + local_best_distance_sq = np.full(local_vertex_count, np.inf, dtype=np.float64) + local_best_values = None + + for source_start in range(0, num_source_rows, effective_chunk_size): + source_stop = min(source_start + effective_chunk_size, num_source_rows) + source_coordinates = coordinate_dataset[source_start:source_stop, :] + source_values = field_dataset[source_start:source_stop] + source_coordinates, source_values = _flatten_source_points( + source_coordinates, source_values, dim + ) + + if local_best_values is None: + local_best_values = np.zeros( + (local_vertex_count, source_values.shape[1]), dtype=field_dataset_dtype + ) + + source_tree = uw.kdtree.KDTree(source_coordinates) + nearest_distance, nearest_index = source_tree.query( + local_vertices, k=1, sqr_dists=False + ) + nearest_distance = np.asarray(nearest_distance).reshape(-1) + nearest_index = np.asarray(nearest_index).reshape(-1) + nearest_distance_sq = nearest_distance * nearest_distance + improved = nearest_distance_sq < local_best_distance_sq + + if np.any(improved): + local_best_distance_sq[improved] = nearest_distance_sq[improved] + local_best_values[improved, :] = source_values[nearest_index[improved], :] + + if local_best_values is None: + # Degenerate case: empty source + local_best_values = np.zeros((local_vertex_count, 1), dtype=field_dataset_dtype) + + if is_tensor: + local_component_count = local_best_values.shape[1] + tensor_values = np.zeros((local_vertex_count, 9), dtype=field_dataset_dtype) + + if dim == 2 and local_component_count == 4: + tensor_values[:, 0] = local_best_values[:, 0] + tensor_values[:, 4] = local_best_values[:, 1] + tensor_values[:, 1] = local_best_values[:, 2] + tensor_values[:, 3] = local_best_values[:, 3] + elif dim == 2 and local_component_count == 3: + tensor_values[:, 0] = local_best_values[:, 0] + tensor_values[:, 4] = local_best_values[:, 1] + tensor_values[:, 1] = local_best_values[:, 2] + tensor_values[:, 3] = local_best_values[:, 2] + elif dim == 3 and local_component_count == 9: + tensor_values[:, :] = local_best_values[:, :] + elif dim == 3 and local_component_count == 6: + tensor_values[:, 0] = local_best_values[:, 0] + tensor_values[:, 4] = local_best_values[:, 1] + tensor_values[:, 8] = local_best_values[:, 2] + tensor_values[:, 1] = local_best_values[:, 3] + tensor_values[:, 3] = local_best_values[:, 3] + tensor_values[:, 5] = local_best_values[:, 4] + tensor_values[:, 7] = local_best_values[:, 4] + tensor_values[:, 2] = local_best_values[:, 5] + tensor_values[:, 6] = local_best_values[:, 5] + else: + tensor_values = None + + if tensor_values is not None: + local_best_values = tensor_values + + output_components = local_best_values.shape[1] + vertex_group = variable_file.require_group("vertex_fields") + if rank == 0: + if "coordinates" in vertex_group: + del vertex_group["coordinates"] + if output_dataset_name in vertex_group: + del vertex_group[output_dataset_name] + if parallel_h5: + comm.Barrier() + + coordinates_output = vertex_group.require_dataset( + "coordinates", shape=(num_vertices, dim), dtype=local_vertices.dtype + ) + values_output = vertex_group.require_dataset( + output_dataset_name, + shape=(num_vertices, output_components), + dtype=field_dataset_dtype, + ) + coordinates_output[local_vertex_start:local_vertex_stop, :] = local_vertices + values_output[local_vertex_start:local_vertex_stop, :] = local_best_values + + if parallel_h5: + comm.Barrier() + + def checkpoint_xdmf( filename: str, meshUpdates: bool = True, @@ -3240,34 +3536,53 @@ def checkpoint_xdmf( ## The mesh Var attributes - def get_cell_field_size(h5_filename, mesh_var): - try: - with h5py.File(h5_filename, "r") as f: - size = f[f"cell_fields/{mesh_var.clean_name}_{mesh_var.clean_name}"].shape[0] - return size - except: - with h5py.File(h5_filename, "r") as f: - size = f[f"fields/{mesh_var.clean_name}"].shape[0] - return size + def get_field_info(h5_filename, mesh_var, center): + """ + Return (num_items, num_components, dataset_path) for a mesh variable. + Prefers vertex/cell compatibility groups, falls back to /fields layout. + """ + compat_name = f"{mesh_var.clean_name}_{mesh_var.clean_name}" + candidates = [] + + if center == "Cell": + candidates = [f"cell_fields/{compat_name}", f"fields/{mesh_var.clean_name}"] + else: + candidates = [f"vertex_fields/{compat_name}", f"fields/{mesh_var.clean_name}"] + + with h5py.File(h5_filename, "r") as f: + for path in candidates: + if path in f: + shp = f[path].shape + if len(shp) == 1: + return shp[0], 1, path + return shp[0], shp[1], path + + raise RuntimeError( + f"Could not locate data for variable '{mesh_var.clean_name}' in {h5_filename}" + ) attributes = "" for var in meshVars: var_filename = filename + f".mesh.{var.clean_name}.{index:05}.h5" - if var.num_components == 1: - variable_type = "Scalar" - else: - variable_type = "Vector" - # Determine if data is stored on nodes (vertex_fields) or cells (cell_fields) if not getattr(var, "continuous") or getattr(var, "degree") == 0: center = "Cell" - numItems = get_cell_field_size(var_filename, var) - field_group = "cell_fields" else: center = "Node" - numItems = numVertices - field_group = "vertex_fields" + numItems, numComponents, dataset_path = get_field_info(var_filename, var, center) + + # Use variable type when available, but reflect actual stored component count. + if hasattr(var, "vtype") and var.vtype in ( + uw.VarType.TENSOR, + uw.VarType.SYM_TENSOR, + uw.VarType.MATRIX, + ): + variable_type = "Tensor" + elif numComponents == 1: + variable_type = "Scalar" + else: + variable_type = "Vector" var_attribute = f""" 0 0 0 1 1 1 - 1 {numItems} {var.num_components} + 1 {numItems} {numComponents} - &{var.clean_name+"_Data"};:/{field_group}/{var.clean_name+"_"+var.clean_name} + &{var.clean_name+"_Data"};:/{dataset_path}