Time dependency and multiple functions per file (#49)

* Start adding functionality for time dependency.

* Add tests for time-dependent/multi-function checkpointing

* Add notes

* Remove newline

* Fix kwarg

* Add Dolfin legacy compatible interface.
Add adios4dolfinx backwards compatibility

* Linting

* Add note to readme

* Change to github images

* Add legacy checkpoint test

* Try fixing pipelines

* Add new workflow to test-code

.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

.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

.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 }}

.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:

.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]
 

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

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,20 +179,45 @@ 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
     """
     io = adios.DeclareIO("ArrayReader")
     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()