diff --git a/.github/workflows/build_docs.yml b/.github/workflows/build_docs.yml index 7397ec1..947bb56 100644 --- a/.github/workflows/build_docs.yml +++ b/.github/workflows/build_docs.yml @@ -23,7 +23,7 @@ jobs: DEB_PYTHON_INSTALL_LAYOUT: deb_system runs-on: ubuntu-22.04 - container: dolfinx/dolfinx:nightly + container: ghcr.io/fenics/dolfinx/dolfinx:nightly steps: # This action sets the current path to the root of your github repo - uses: actions/checkout@v4 diff --git a/.github/workflows/check_formatting.yml b/.github/workflows/check_formatting.yml index 4d0fde2..f41f293 100644 --- a/.github/workflows/check_formatting.yml +++ b/.github/workflows/check_formatting.yml @@ -10,7 +10,7 @@ jobs: DEB_PYTHON_INSTALL_LAYOUT: deb_system runs-on: ubuntu-22.04 - container: dolfinx/dolfinx:nightly + container: ghcr.io/fenics/dolfinx/dolfinx:nightly steps: # This action sets the current path to the root of your github repo - uses: actions/checkout@v4 diff --git a/.github/workflows/create_legacy_checkpoint.yml b/.github/workflows/create_legacy_checkpoint.yml new file mode 100644 index 0000000..dbcd899 --- /dev/null +++ b/.github/workflows/create_legacy_checkpoint.yml @@ -0,0 +1,27 @@ +name: Generate adios4dolfinx legacy data + +on: + workflow_call: + +env: + data_dir: "legacy_checkpoint" + adios4dolfinx_version: "0.7.1" + +jobs: + create-adios-data: + runs-on: "ubuntu-22.04" + container: ghcr.io/fenics/dolfinx/dolfinx:v0.7.3 + + steps: + - uses: actions/checkout@v4 + + - name: Install legacy version of adios4dolfinx + - run: python3 -m pip install adios4dolfinx==${adios4dolfinx_version} + + - name: Create datasets + run: python3 ./tests/create_legacy_checkpoint.py --output-dir=$data_dir + + - uses: actions/upload-artifact@v4 + with: + name: ${{ env.data_dir }} + path: ./${{ env.data_dir }} diff --git a/.github/workflows/create_legacy_data.yml b/.github/workflows/create_legacy_data.yml index 8c829d7..79f923e 100644 --- a/.github/workflows/create_legacy_data.yml +++ b/.github/workflows/create_legacy_data.yml @@ -1,4 +1,4 @@ -name: Generate legacy data +name: Generate data from Legacy DOLFIN on: workflow_call: @@ -8,7 +8,7 @@ env: jobs: - create-data: + create-dolfin-data: runs-on: "ubuntu-22.04" container: ghcr.io/scientificcomputing/fenics:2023-11-15 steps: diff --git a/.github/workflows/test_package.yml b/.github/workflows/test_package.yml index 718313e..f24225c 100644 --- a/.github/workflows/test_package.yml +++ b/.github/workflows/test_package.yml @@ -21,12 +21,15 @@ jobs: create-datasets: uses: ./.github/workflows/create_legacy_data.yml + create-legacy-datasets: + uses: ./.github/workflows/create_legacy_checkpoint.yml + check-formatting: uses: ./.github/workflows/check_formatting.yml test-code: runs-on: "ubuntu-22.04" - needs: [create-datasets, check-formatting] + needs: [create-datasets, create-legacy-datasets, check-formatting] container: ghcr.io/fenics/dolfinx/dolfinx:nightly env: DEB_PYTHON_INSTALL_LAYOUT: deb_system @@ -43,7 +46,13 @@ jobs: with: name: legacy path: ./legacy - + + - name: Download legacy data + uses: actions/download-artifact@v4 + with: + name: legacy_checkpoint + path: ./legacy_checkpoint + - name: Install package run: python3 -m pip install .[test] diff --git a/README.md b/README.md index 292f8bf..0964466 100644 --- a/README.md +++ b/README.md @@ -37,6 +37,18 @@ For scalability, the code uses [MPI Neighbourhood collectives](https://www.mpi-f - Reading and writing meshtags associated to meshes `adios4dolfinx.read/write_meshtags` - Reading checkpoints for any element (serial and parallel, one checkpoint per file). Use `adios4dolfinx.read/write_function`. +> [!IMPORTANT] +> For a checkpoint to be valid, you first have to store the mesh with `write_mesh`, then use `write_function` to append to the checkpoint file. + +> [!IMPORTANT] +> A checkpoint file supports multiple functions and multiple time steps, as long as the functions are associated with the same mesh + +> [!IMPORTANT] +> Only one mesh per file is allowed + +> [!WARNING] +> If you are using checkpoints written with `adios4dolfinx` prior to time-dependent/multi-function support, please use the `legacy=True` flag for reading in the checkpoint + ## Legacy DOLFIN Only checkpoints for `Lagrange` or `DG` functions are supported from legacy DOLFIN diff --git a/src/adios4dolfinx/adios2_helpers.py b/src/adios4dolfinx/adios2_helpers.py index 01589cd..8df8ec4 100644 --- a/src/adios4dolfinx/adios2_helpers.py +++ b/src/adios4dolfinx/adios2_helpers.py @@ -128,7 +128,7 @@ def read_dofmap( break infile.EndStep() if dofmap_offsets not in io.AvailableVariables().keys(): - raise KeyError(f"Dof offsets not found at '{dofmap_offsets}'") + raise KeyError(f"Dof offsets not found at '{dofmap_offsets}' in {filename}") # Get global shape of dofmap-offset, and read in data with an overlap d_offsets = io.InquireVariable(dofmap_offsets) @@ -147,7 +147,7 @@ def read_dofmap( # Assuming dofmap is saved in stame step # Get the relevant part of the dofmap if dofmap not in io.AvailableVariables().keys(): - raise KeyError(f"Dof offsets not found at {dofmap}") + raise KeyError(f"Dof offsets not found at {dofmap} in {filename}") cell_dofs = io.InquireVariable(dofmap) cell_dofs.SetSelection([[in_offsets[0]], [in_offsets[-1] - in_offsets[0]]]) in_dofmap = np.empty( @@ -166,9 +166,10 @@ def read_dofmap( def read_array( - adios: adios2.ADIOS, - filename: pathlib.Path, array_name: str, engine: str, comm: MPI.Intracomm -) -> Tuple[npt.NDArray[valid_function_types], int]: + adios: adios2.ADIOS, + filename: pathlib.Path, array_name: str, engine: str, comm: MPI.Intracomm, + time: float = 0., time_name: str = "", + legacy: bool = False) -> Tuple[npt.NDArray[valid_function_types], int]: """ Read an array from file, return the global starting position of the local array @@ -178,6 +179,8 @@ def read_array( array_name: Name of array in file engine: Name of engine to use to read file comm: MPI communicator used for reading the data + time_name: Name of time variable for modern checkpoints + legacy: If True ignore time_name and read the first available step Returns: Local part of array and its global starting position """ @@ -185,13 +188,36 @@ def read_array( io.SetEngine(engine) infile = io.Open(str(filename), adios2.Mode.Read) - for i in range(infile.Steps()): - infile.BeginStep() - if array_name in io.AvailableVariables().keys(): - break - infile.EndStep() - if array_name not in io.AvailableVariables().keys(): - raise KeyError(f"No array found at {array_name}") + # Get time-stamp from first available step + if legacy: + for i in range(infile.Steps()): + infile.BeginStep() + if array_name in io.AvailableVariables().keys(): + break + infile.EndStep() + if array_name not in io.AvailableVariables().keys(): + raise KeyError(f"No array found at {array_name}") + else: + for i in range(infile.Steps()): + infile.BeginStep() + if time_name in io.AvailableVariables().keys(): + arr = io.InquireVariable(time_name) + time_shape = arr.Shape() + arr.SetSelection([[0], [time_shape[0]]]) + times = np.empty(time_shape[0], dtype=adios_to_numpy_dtype[arr.Type()]) + infile.Get(arr, times, adios2.Mode.Sync) + if times[0] == time: + break + if i == infile.Steps() - 1: + raise KeyError(f"No data associated with {time_name}={time} found in {filename}") + + infile.EndStep() + + if time_name not in io.AvailableVariables().keys(): + raise KeyError(f"No data associated with {time_name}={time} found in {filename}") + + if array_name not in io.AvailableVariables().keys(): + raise KeyError(f"No array found at {time=} for {array_name}") arr = io.InquireVariable(array_name) arr_shape = arr.Shape() diff --git a/src/adios4dolfinx/checkpointing.py b/src/adios4dolfinx/checkpointing.py index f01e8e0..f190f75 100644 --- a/src/adios4dolfinx/checkpointing.py +++ b/src/adios4dolfinx/checkpointing.py @@ -108,12 +108,8 @@ def write_mesh(mesh: dolfinx.mesh.Mesh, filename: Path, engine: str = "BP4"): # Write basix properties cmap = mesh.geometry.cmap - io.DefineAttribute( - "Degree", np.array([cmap.degree], dtype=np.int32) - ) - io.DefineAttribute( - "LagrangeVariant", np.array([cmap.variant], dtype=np.int32) - ) + io.DefineAttribute("Degree", np.array([cmap.degree], dtype=np.int32)) + io.DefineAttribute("LagrangeVariant", np.array([cmap.variant], dtype=np.int32)) # Write topology g_imap = mesh.geometry.index_map() @@ -126,9 +122,9 @@ def write_mesh(mesh: dolfinx.mesh.Mesh, filename: Path, engine: str = "BP4"): dofs_out = np.zeros((num_cells_local, num_dofs_per_cell), dtype=np.int64) assert g_dmap.shape[1] == num_dofs_per_cell - dofs_out[:, :] = np.asarray(g_imap.local_to_global( - g_dmap[:num_cells_local, :].reshape(-1) - )).reshape(dofs_out.shape) + dofs_out[:, :] = np.asarray( + g_imap.local_to_global(g_dmap[:num_cells_local, :].reshape(-1)) + ).reshape(dofs_out.shape) dvar = io.DefineVariable( "Topology", @@ -309,7 +305,8 @@ def read_meshtags(filename: str, mesh: dolfinx.mesh.Mesh, meshtag_name: str, return mt -def read_function(u: dolfinx.fem.Function, filename: Path, engine: str = "BP4"): +def read_function(u: dolfinx.fem.Function, filename: Path, engine: str = "BP4", time: float = 0., + legacy: bool = False): """ Read checkpoint from file and fill it into `u`. @@ -317,10 +314,12 @@ def read_function(u: dolfinx.fem.Function, filename: Path, engine: str = "BP4"): u: Function to fill filename: Path to checkpoint engine: ADIOS engine type used for reading + legacy: If checkpoint is from prior to time-dependent writing set to True """ mesh = u.function_space.mesh comm = mesh.comm adios = adios2.ADIOS(comm) + name = u.name # ----------------------Step 1--------------------------------- # Compute index of input cells and get cell permutation num_owned_cells = mesh.topology.index_map(mesh.topology.dim).size_local @@ -341,8 +340,12 @@ def read_function(u: dolfinx.fem.Function, filename: Path, engine: str = "BP4"): # -------------------Step 3----------------------------------- # Read dofmap from file and compute dof owners - dofmap_path = "Dofmap" - xdofmap_path = "XDofmap" + if legacy: + dofmap_path = "Dofmap" + xdofmap_path = "XDofmap" + else: + dofmap_path = f"{name}_dofmap" + xdofmap_path = f"{name}_XDofmap" input_dofmap = read_dofmap( adios, comm, filename, dofmap_path, xdofmap_path, num_cells_global, engine ) @@ -355,8 +358,13 @@ def read_function(u: dolfinx.fem.Function, filename: Path, engine: str = "BP4"): # --------------------Step 4----------------------------------- # Read array from file and communicate them to input dofmap process - array_path = "Values" - input_array, starting_pos = read_array(adios, filename, array_path, engine, comm) + if legacy: + array_path = "Values" + else: + array_path = f"{name}_values" + time_name = f"{name}_time" + input_array, starting_pos = read_array(adios, filename, array_path, engine, comm, time, time_name, + legacy=legacy) recv_array = send_dofs_and_recv_values( input_dofmap.array, dof_owner, comm, input_array, starting_pos ) @@ -438,21 +446,21 @@ def read_mesh( # Get mesh cell type if "CellType" not in io.AvailableAttributes().keys(): - raise KeyError("Mesh cell type not found at CellType") + raise KeyError(f"Mesh cell type not found at CellType in {file}") celltype = io.InquireAttribute("CellType") cell_type = celltype.DataString()[0] # Get basix info if "LagrangeVariant" not in io.AvailableAttributes().keys(): - raise KeyError("Mesh LagrangeVariant not found") + raise KeyError(f"Mesh LagrangeVariant not found in {file}") lvar = io.InquireAttribute("LagrangeVariant").Data()[0] if "Degree" not in io.AvailableAttributes().keys(): - raise KeyError("Mesh degree not found") + raise KeyError(f"Mesh degree not found in {file}") degree = io.InquireAttribute("Degree").Data()[0] # Get mesh geometry if "Points" not in io.AvailableVariables().keys(): - raise KeyError("Mesh coordinates not found at Points") + raise KeyError(f"Mesh coordinates not found at Points in {file}") geometry = io.InquireVariable("Points") x_shape = geometry.Shape() geometry_range = compute_local_range(comm, x_shape[0]) @@ -460,12 +468,13 @@ def read_mesh( [[geometry_range[0], 0], [geometry_range[1] - geometry_range[0], x_shape[1]]] ) mesh_geometry = np.empty( - (geometry_range[1] - geometry_range[0], x_shape[1]), dtype=adios_to_numpy_dtype[geometry.Type()] + (geometry_range[1] - geometry_range[0], x_shape[1]), + dtype=adios_to_numpy_dtype[geometry.Type()], ) infile.Get(geometry, mesh_geometry, adios2.Mode.Deferred) # Get mesh topology (distributed) if "Topology" not in io.AvailableVariables().keys(): - raise KeyError("Mesh topology not found at Topology'") + raise KeyError("Mesh topology not found at Topology in {file}") topology = io.InquireVariable("Topology") shape = topology.Shape() local_range = compute_local_range(comm, shape[0]) @@ -489,7 +498,7 @@ def read_mesh( basix.LagrangeVariant(int(lvar)), shape=(mesh_geometry.shape[1],), gdim=mesh_geometry.shape[1], - dtype=mesh_geometry.dtype + dtype=mesh_geometry.dtype, ) domain = ufl.Mesh(element) partitioner = dolfinx.cpp.mesh.create_cell_partitioner(ghost_mode) @@ -503,6 +512,7 @@ def write_function( filename: Path, engine: str = "BP4", mode: adios2.Mode = adios2.Mode.Append, + time: float = 0.0 ): """ Write function checkpoint to file. @@ -510,8 +520,9 @@ def write_function( Args: u: Function to write to file filename: Path to write to - egine: ADIOS2 engine + engine: ADIOS2 engine mode: Write or append. + time: Time-stamp for simulation """ dofmap = u.function_space.dofmap values = u.x.array @@ -521,14 +532,43 @@ def write_function( adios = adios2.ADIOS(comm) io = adios.DeclareIO("FunctionWriter") io.SetEngine(engine) - outfile = io.Open(str(filename), adios2.Mode.Append) + + # If mode is append, check if we have written the function to file before + name = u.name + + if mode == adios2.Mode.Append: + # First open the file in read-mode to check if the function has been written before + read_file = io.Open(str(filename), adios2.Mode.Read) + io.SetEngine(engine) + first_write = True + for _ in range(read_file.Steps()): + read_file.BeginStep() + if name in io.AvailableAttributes(): + first_write = False + break + read_file.EndStep() + read_file.Close() + outfile = io.Open(str(filename), mode) + io.DefineAttribute(name, name) + outfile.BeginStep() + + # Add time step to file + t_arr = np.array([time], dtype=np.float64) + time_var = io.DefineVariable( + f"{name}_time", + t_arr, + shape=[1], + start=[0], + count=[1 if mesh.comm.rank == 0 else 0], + ) + outfile.Put(time_var, t_arr) + # Write local part of vector num_dofs_local = dofmap.index_map.size_local * dofmap.index_map_bs num_dofs_global = dofmap.index_map.size_global * dofmap.index_map_bs local_start = dofmap.index_map.local_range[0] * dofmap.index_map_bs - outfile.BeginStep() val_var = io.DefineVariable( - "Values", + f"{name}_values", np.zeros(num_dofs_local, dtype=u.dtype), shape=[num_dofs_global], start=[local_start], @@ -536,7 +576,14 @@ def write_function( ) outfile.Put(val_var, values[:num_dofs_local]) - # Create global, unrolled dofmap + if not first_write: + outfile.PerformPuts() + outfile.EndStep() + outfile.Close() + assert adios.RemoveIO("FunctionWriter") + return + + # Convert local dofmap into global_dofmap dmap = dofmap.list num_dofs_per_cell = dmap.shape[1] dofmap_bs = dofmap.bs @@ -559,7 +606,7 @@ def write_function( # Get offsets of dofmap dofmap_imap = dolfinx.common.IndexMap(mesh.comm, num_dofs_local_dmap) dofmap_var = io.DefineVariable( - "Dofmap", + f"{name}_dofmap", np.zeros(num_dofs_local_dmap, dtype=np.int64), shape=[dofmap_imap.size_global], start=[dofmap_imap.local_range[0]], @@ -571,7 +618,7 @@ def write_function( cell_start = mesh.topology.index_map(mesh.topology.dim).local_range[0] local_dofmap_offsets += dofmap_imap.local_range[0] xdofmap_var = io.DefineVariable( - "XDofmap", + f"{name}_XDofmap", np.zeros(num_cells_local + 1, dtype=np.int64), shape=[num_cells_global + 1], start=[cell_start], diff --git a/src/adios4dolfinx/legacy_readers.py b/src/adios4dolfinx/legacy_readers.py index 9e5f5f8..0cb19b7 100644 --- a/src/adios4dolfinx/legacy_readers.py +++ b/src/adios4dolfinx/legacy_readers.py @@ -406,7 +406,8 @@ def read_function_from_legacy_h5( # Read input data adios = adios2.ADIOS(comm) - local_array, starting_pos = read_array(adios, filename, f"/{group}/{vector_group}", "HDF5", comm) + local_array, starting_pos = read_array(adios, filename, f"/{group}/{vector_group}", "HDF5", + comm, legacy=True) # Send global dof indices to correct input process, and receive value of given dof local_values = send_dofs_and_recv_values( diff --git a/tests/create_legacy_checkpoint.py b/tests/create_legacy_checkpoint.py new file mode 100644 index 0000000..25ac67b --- /dev/null +++ b/tests/create_legacy_checkpoint.py @@ -0,0 +1,68 @@ +# Copyright (C) 2024 Jørgen Schartum Dokken +# +# This file is part of adios4dolfinx +# +# SPDX-License-Identifier: MIT + + +""" +Functions to create checkpoints with adios4dolfinx v0.7.x +""" + +import argparse +import pathlib + +import numpy as np +from mpi4py import MPI +import adios4dolfinx +import dolfinx +from importlib.metadata import version + +a4d_version = version("adios4dolfinx") +assert a4d_version < "0.8.0", f"Creating a legacy checkpoint requires adios4dolfinx < 0.8.0, you have {a4d_version}." + + +def f(x): + values = np.zeros((2, x.shape[1]), dtype=np.float64) + values[0] = x[0] + values[1] = -x[1] + return values + + +def write_checkpoint(filename, mesh, el, f): + V = dolfinx.fem.FunctionSpace(mesh, el) + uh = dolfinx.fem.Function(V, dtype=np.float64) + uh.interpolate(f) + + adios4dolfinx.write_mesh(V.mesh, filename) + adios4dolfinx.write_function(uh, filename) + + +def verify_checkpoint(filename, el, f): + mesh = adios4dolfinx.read_mesh(MPI.COMM_WORLD, filename, "BP4", dolfinx.mesh.GhostMode.shared_facet) + V = dolfinx.fem.FunctionSpace(mesh, el) + uh = dolfinx.fem.Function(V, dtype=np.float64) + adios4dolfinx.read_function(uh, filename) + + u_ex = dolfinx.fem.Function(V, dtype=np.float64) + u_ex.interpolate(f) + + np.testing.assert_allclose(u_ex.x.array, uh.x.array, atol=1e-15) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + formatter_class=argparse.ArgumentDefaultsHelpFormatter + ) + parser.add_argument("--output-dir", type=str, default="legacy_checkpoint", dest="dir") + + inputs = parser.parse_args() + path = pathlib.Path(inputs.dir) + path.mkdir(exist_ok=True, parents=True) + filename = path / "adios4dolfinx_checkpoint.bp" + + mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 10, 10) + el = ("N1curl", 3) + write_checkpoint(filename, mesh, el, f) + MPI.COMM_WORLD.Barrier() + verify_checkpoint(filename, el, f) diff --git a/tests/create_legacy_data.py b/tests/create_legacy_data.py index 931ded2..7008592 100644 --- a/tests/create_legacy_data.py +++ b/tests/create_legacy_data.py @@ -10,7 +10,6 @@ """ import argparse -import os import pathlib import dolfin @@ -135,8 +134,7 @@ def verify_xdmf( inputs = parser.parse_args() path = pathlib.Path(inputs.dir) - if not os.path.exists(path): - os.mkdir(path) + path.mkdir(exist_ok=True, parents=True) h5_filename = path / f"{inputs.name}.h5" xdmf_filename = path / f"{inputs.name}_checkpoint.xdmf" diff --git a/tests/test_checkpointing.py b/tests/test_checkpointing.py index 8b54154..3e916e9 100644 --- a/tests/test_checkpointing.py +++ b/tests/test_checkpointing.py @@ -7,7 +7,7 @@ import pytest from mpi4py import MPI -from .test_utils import get_dtype, read_function, write_function +from .test_utils import read_function, write_function, get_dtype, read_function_time_dep, write_function_time_dep dtypes = [np.float64, np.float32] # Mesh geometry dtypes write_comm = [MPI.COMM_SELF, MPI.COMM_WORLD] # Communicators for creating mesh @@ -86,3 +86,73 @@ def f(x): MPI.COMM_WORLD.Barrier() read_function(read_comm, el, f, hash, f_dtype) + + +@pytest.mark.parametrize("complex", [True, False]) +@pytest.mark.parametrize("family", ["Lagrange", "DG"]) +@pytest.mark.parametrize("degree", [1, 4]) +@pytest.mark.parametrize("read_comm", [MPI.COMM_SELF, MPI.COMM_WORLD]) +def test_read_write_P_2D_time(read_comm, family, degree, complex, mesh_2D): + mesh = mesh_2D + f_dtype = get_dtype(mesh.geometry.x.dtype, complex) + + el = basix.ufl.element(family, + mesh.ufl_cell().cellname(), + degree, + basix.LagrangeVariant.gll_warped, + gdim=mesh.geometry.dim, + shape=(mesh.geometry.dim, ), + dtype=mesh.geometry.x.dtype) + + def f0(x): + values = np.empty((2, x.shape[1]), dtype=f_dtype) + values[0] = np.full(x.shape[1], np.pi) + x[0] + x[1] * 1j + values[1] = x[0] + 3j * x[1] + return values + + def f1(x): + values = np.empty((2, x.shape[1]), dtype=f_dtype) + values[0] = 2*np.full(x.shape[1], np.pi) + x[0] + x[1] * 1j + values[1] = -x[0] + 3j * x[1] + 2*x[1] + return values + + t0 = 0.8 + t1 = 0.6 + hash = write_function_time_dep(mesh, el, f0, f1, t0, t1, f_dtype) + MPI.COMM_WORLD.Barrier() + read_function_time_dep(read_comm, el, f0, f1, t0, t1, hash, f_dtype) + + +@pytest.mark.parametrize("complex", [True, False]) +@pytest.mark.parametrize("family", ["Lagrange", "DG"]) +@pytest.mark.parametrize("degree", [1, 4]) +@pytest.mark.parametrize("read_comm", [MPI.COMM_SELF, MPI.COMM_WORLD]) +def test_read_write_P_3D_time(read_comm, family, degree, complex, mesh_3D): + mesh = mesh_3D + f_dtype = get_dtype(mesh.geometry.x.dtype, complex) + el = basix.ufl.element(family, + mesh.ufl_cell().cellname(), + degree, + basix.LagrangeVariant.gll_warped, + gdim=mesh.geometry.dim, + shape=(mesh.geometry.dim, )) + + def f(x): + values = np.empty((3, x.shape[1]), dtype=f_dtype) + values[0] = np.pi + x[0] + 2j*x[2] + values[1] = x[1] + 2 * x[0] + values[2] = 1j*x[1] + np.cos(x[2]) + return values + + def g(x): + values = np.empty((3, x.shape[1]), dtype=f_dtype) + values[0] = x[0] + np.pi * 2j*x[2] + values[1] = 1j*x[2] + 2 * x[0] + values[2] = x[0] + 1j*np.cos(x[1]) + return values + + t0 = 0.1 + t1 = 1.3 + hash = write_function_time_dep(mesh, el, g, f, t0, t1, f_dtype) + MPI.COMM_WORLD.Barrier() + read_function_time_dep(read_comm, el, g, f, t0, t1, hash, f_dtype) diff --git a/tests/test_checkpointing_vector.py b/tests/test_checkpointing_vector.py index f1402fd..408bdf9 100644 --- a/tests/test_checkpointing_vector.py +++ b/tests/test_checkpointing_vector.py @@ -140,3 +140,36 @@ def f(x): hash = write_function(mesh, el, f, dtype=f_dtype) MPI.COMM_WORLD.Barrier() read_function(read_comm, el, f, hash, dtype=f_dtype) + + +@pytest.mark.parametrize("complex", [True, False]) +@pytest.mark.parametrize("family", ["RTCF"]) +@pytest.mark.parametrize("degree", [1, 2, 3]) +@pytest.mark.parametrize("read_comm", [MPI.COMM_SELF, MPI.COMM_WORLD]) +def test_read_write_multiple(read_comm, family, degree, complex, non_simplex_mesh_2D): + mesh = non_simplex_mesh_2D + f_dtype = get_dtype(mesh.geometry.x.dtype, complex) + el = basix.ufl.element(family, + mesh.ufl_cell().cellname(), + degree, + gdim=mesh.geometry.dim) + + def f(x): + values = np.empty((2, x.shape[1]), dtype=f_dtype) + values[0] = np.full(x.shape[1], np.pi) + 2j*x[2] + values[1] = x[1] + 2 * x[0] + 2j*np.cos(x[2]) + return values + + def g(x): + values = np.empty((2, x.shape[1]), dtype=f_dtype) + values[0] = (1+1j) * x[0] + values[1] = (3-3j) * x[1] + return values + + hash_f = write_function(mesh, el, f, dtype=f_dtype, name="f", append=False) + hash_g = write_function(mesh, el, g, dtype=f_dtype, name="g", append=True) + assert hash_f == hash_g + + MPI.COMM_WORLD.Barrier() + read_function(read_comm, el, f, hash_f, dtype=f_dtype, name="f") + read_function(read_comm, el, g, hash_g, dtype=f_dtype, name="g") diff --git a/tests/test_legacy_readers.py b/tests/test_legacy_readers.py index c102148..a73152f 100644 --- a/tests/test_legacy_readers.py +++ b/tests/test_legacy_readers.py @@ -11,16 +11,19 @@ import dolfinx import numpy as np +import pytest import ufl from dolfinx.fem.petsc import LinearProblem -from adios4dolfinx import (read_function_from_legacy_h5, - read_mesh_from_legacy_h5) +from adios4dolfinx import (read_function, read_function_from_legacy_h5, + read_mesh, read_mesh_from_legacy_h5) def test_legacy_mesh(): comm = MPI.COMM_WORLD path = (pathlib.Path("legacy") / "mesh.h5").absolute() + if not path.exists(): + pytest.skip(f"{path} does not exist") mesh = read_mesh_from_legacy_h5(comm=comm, filename=path, group="/mesh") assert mesh.topology.dim == 3 volume = mesh.comm.allreduce( @@ -42,6 +45,8 @@ def test_legacy_mesh(): def test_read_legacy_mesh_from_checkpoint(): comm = MPI.COMM_WORLD filename = (pathlib.Path("legacy") / "mesh_checkpoint.h5").absolute() + if not filename.exists(): + pytest.skip(f"{filename} does not exist") mesh = read_mesh_from_legacy_h5(comm=comm, filename=filename, group="/Mesh/mesh") assert mesh.topology.dim == 3 volume = mesh.comm.allreduce( @@ -63,6 +68,8 @@ def test_read_legacy_mesh_from_checkpoint(): def test_legacy_function(): comm = MPI.COMM_WORLD path = (pathlib.Path("legacy") / "mesh.h5").absolute() + if not path.exists(): + pytest.skip(f"{path} does not exist") mesh = read_mesh_from_legacy_h5(comm, path, "/mesh") V = dolfinx.fem.functionspace(mesh, ("DG", 2)) u = ufl.TrialFunction(V) @@ -80,7 +87,7 @@ def test_legacy_function(): u_in = dolfinx.fem.Function(V) read_function_from_legacy_h5(mesh.comm, path, u_in, group="v") - assert np.allclose(uh.x.array, u_in.x.array) + np.testing.assert_allclose(uh.x.array, u_in.x.array, atol=1e-14) W = dolfinx.fem.functionspace(mesh, ("DG", 2, (mesh.geometry.dim, ))) wh = dolfinx.fem.Function(W) @@ -89,12 +96,14 @@ def test_legacy_function(): read_function_from_legacy_h5(mesh.comm, path, w_in, group="w") - assert np.allclose(wh.x.array, w_in.x.array) + np.testing.assert_allclose(wh.x.array, w_in.x.array, atol=1e-14) def test_read_legacy_function_from_checkpoint(): comm = MPI.COMM_WORLD path = (pathlib.Path("legacy") / "mesh_checkpoint.h5").absolute() + if not path.exists(): + pytest.skip(f"{path} does not exist") mesh = read_mesh_from_legacy_h5(comm, path, "/Mesh/mesh") V = dolfinx.fem.functionspace(mesh, ("DG", 2)) @@ -126,8 +135,33 @@ def test_read_legacy_function_from_checkpoint(): w_in = dolfinx.fem.Function(W) read_function_from_legacy_h5(mesh.comm, path, w_in, group="w", step=0) - assert np.allclose(wh.x.array, w_in.x.array) + np.testing.assert_allclose(wh.x.array, w_in.x.array, atol=1e-14) wh.interpolate(lambda x: np.vstack((x[0], 0*x[0], x[1]))) read_function_from_legacy_h5(mesh.comm, path, w_in, group="w", step=1) - assert np.allclose(wh.x.array, w_in.x.array) + np.testing.assert_allclose(wh.x.array, w_in.x.array, atol=1e-14) + + +def test_adios4dolfinx_legacy(): + comm = MPI.COMM_WORLD + path = (pathlib.Path("legacy_checkpoint") / "adios4dolfinx_checkpoint.bp").absolute() + if not path.exists(): + pytest.skip(f"{path} does not exist") + + el = ("N1curl", 3) + mesh = read_mesh(comm, path, "BP4", dolfinx.mesh.GhostMode.shared_facet) + + def f(x): + values = np.zeros((2, x.shape[1]), dtype=np.float64) + values[0] = x[0] + values[1] = -x[1] + return values + + V = dolfinx.fem.functionspace(mesh, el) + u = dolfinx.fem.Function(V) + read_function(u, path, engine="BP4", legacy=True) + + u_ex = dolfinx.fem.Function(V) + u_ex.interpolate(f) + + np.testing.assert_allclose(u.x.array, u_ex.x.array, atol=1e-14) diff --git a/tests/test_utils.py b/tests/test_utils.py index 6501cbd..ee4210a 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -9,43 +9,51 @@ import adios4dolfinx -def write_function(mesh, el, f, dtype) -> str: +def write_function(mesh, el, f, dtype, name="uh", append: bool = False) -> str: V = dolfinx.fem.functionspace(mesh, el) uh = dolfinx.fem.Function(V, dtype=dtype) uh.interpolate(f) + uh.name = name el_hash = ( V.element.signature() .replace(" ", "") .replace(",", "") .replace("(", "") .replace(")", "") + .replace("[", "") + .replace("]", "") ) - filename = pathlib.Path(f"output/mesh{el_hash}_{dtype}.bp") - if mesh.comm.size != 1: - adios4dolfinx.write_mesh(mesh, filename) - adios4dolfinx.write_function(uh, filename) + file_hash = f"{el_hash}_{np.dtype(dtype).name}" + filename = pathlib.Path(f"output/mesh_{file_hash}.bp") + if mesh.comm.size != 1: + if not append: + adios4dolfinx.write_mesh(mesh, filename) + adios4dolfinx.write_function(uh, filename, time=0.0) else: if MPI.COMM_WORLD.rank == 0: - adios4dolfinx.write_mesh(mesh, filename) - adios4dolfinx.write_function(uh, filename) - return f"{el_hash}_{dtype}" + if not append: + adios4dolfinx.write_mesh(mesh, filename) + adios4dolfinx.write_function(uh, filename, time=0.0) + return file_hash -def read_function(comm, el, f, hash, dtype): - filename = f"output/mesh{hash}.bp" + +def read_function(comm, el, f, hash, dtype, name="uh"): + filename = f"output/mesh_{hash}.bp" engine = "BP4" mesh = adios4dolfinx.read_mesh( comm, filename, engine, dolfinx.mesh.GhostMode.shared_facet ) V = dolfinx.fem.functionspace(mesh, el) v = dolfinx.fem.Function(V, dtype=dtype) + v.name = name adios4dolfinx.read_function(v, filename, engine) v_ex = dolfinx.fem.Function(V, dtype=dtype) v_ex.interpolate(f) res = np.finfo(dtype).resolution - assert np.allclose(v.x.array, v_ex.x.array, atol=10*res, rtol=10*res) + assert np.allclose(v.x.array, v_ex.x.array, atol=10 * res, rtol=10 * res) def get_dtype(in_dtype: np.dtype, complex: bool): @@ -63,3 +71,58 @@ def get_dtype(in_dtype: np.dtype, complex: bool): else: raise ValueError("Unsuported dtype") return dtype + + +def write_function_time_dep(mesh, el, f0, f1, t0, t1, dtype) -> str: + V = dolfinx.fem.functionspace(mesh, el) + uh = dolfinx.fem.Function(V, dtype=dtype) + uh.interpolate(f0) + el_hash = ( + V.element.signature() + .replace(" ", "") + .replace(",", "") + .replace("(", "") + .replace(")", "") + .replace("[", "") + .replace("]", "") + ) + file_hash = f"{el_hash}_{np.dtype(dtype).name}" + filename = pathlib.Path(f"output/mesh_{file_hash}.bp") + if mesh.comm.size != 1: + adios4dolfinx.write_mesh(mesh, filename) + adios4dolfinx.write_function(uh, filename, time=t0) + uh.interpolate(f1) + adios4dolfinx.write_function(uh, filename, time=t1) + + else: + if MPI.COMM_WORLD.rank == 0: + adios4dolfinx.write_mesh(mesh, filename) + adios4dolfinx.write_function(uh, filename, time=t0) + uh.interpolate(f1) + adios4dolfinx.write_function(uh, filename, time=t1) + + return file_hash + + +def read_function_time_dep(comm, el, f0, f1, t0, t1, hash, dtype): + filename = f"output/mesh_{hash}.bp" + engine = "BP4" + mesh = adios4dolfinx.read_mesh( + comm, filename, engine, dolfinx.mesh.GhostMode.shared_facet + ) + V = dolfinx.fem.functionspace(mesh, el) + v = dolfinx.fem.Function(V, dtype=dtype) + + adios4dolfinx.read_function(v, filename, engine, time=t1) + v_ex = dolfinx.fem.Function(V, dtype=dtype) + v_ex.interpolate(f1) + + res = np.finfo(dtype).resolution + assert np.allclose(v.x.array, v_ex.x.array, atol=10 * res, rtol=10 * res) + + adios4dolfinx.read_function(v, filename, engine, time=t0) + v_ex = dolfinx.fem.Function(V, dtype=dtype) + v_ex.interpolate(f0) + + res = np.finfo(dtype).resolution + assert np.allclose(v.x.array, v_ex.x.array, atol=10 * res, rtol=10 * res)