Add on_input_mesh to write_mesh and write_function

docs/backends/opencarp.rst

@@ -9,9 +9,9 @@ Writes and reads linear meshes in the ASCII mesh format of the cardiac simulator
 and surface (``.surf``) files. The backend needs no extra dependencies.
 
 The mesh, its tags and its fiber and sheet directions are written with
-:func:`io4dolfinx.write_mesh_input_order` (or :func:`io4dolfinx.write_mesh`),
-:func:`io4dolfinx.write_meshtags` and :func:`io4dolfinx.write_cell_data`, as shown in the module
-documentation below. Rank 0 writes all files. openCARP expects coordinates in µm, so the factor
+:func:`io4dolfinx.write_mesh`, :func:`io4dolfinx.write_meshtags` and
+:func:`io4dolfinx.write_cell_data`, with the same ``on_input_mesh`` in every call, as shown in the
+module documentation below. Rank 0 writes all files. openCARP expects coordinates in µm, so the factor
 converting the mesh coordinates to µm must always be given as ``backend_args["scale"]``.
 
 With ``markers={"ENDO": (1, 2), "EPI": (2, 2)}``, a heart mesh gives ``heart.pts``,

docs/original_checkpoint.py

@@ -10,7 +10,10 @@
 # An optional way of store an N-to-M checkpoint is to store the function data in the same
 # ordering as the mesh. The write operation will be more expensive, as it requires data
 # communication to ensure contiguous data being written to the checkpoint.
-# The method is exposed as {py:func}`io4dolfinx.write_function_on_input_mesh`.
+# This is done with {py:func}`io4dolfinx.write_function` with `on_input_mesh=True`,
+# and the mesh can be written in the same ordering with {py:func}`io4dolfinx.write_mesh`
+# with `on_input_mesh=True`. (`write_function_on_input_mesh` and `write_mesh_input_order`
+# do the same, but are deprecated.)
 # Below we will demonstrate this method.
 
 # +
@@ -72,12 +75,13 @@ def write_function(
     u = dolfinx.fem.Function(V)
     u.interpolate(f)
 
-    io4dolfinx.write_function_on_input_mesh(
+    io4dolfinx.write_function(
         function_filename.with_suffix(".bp"),
         u,
         mode=io4dolfinx.FileMode.write,
         time=0.0,
         name="Output",
+        on_input_mesh=True,
     )
     print(
         f"{mesh.comm.rank + 1}/{mesh.comm.size} Function written to ",

docs/quickstart.md

@@ -123,7 +123,7 @@ Sometimes you want to save a solution that corresponds exactly to the input mesh
 This is useful for visualization or post-processing on the original geometry.
 
 ```python
-io4dolfinx.write_function_on_input_mesh("solution_on_input.bp", u)
+io4dolfinx.write_function("solution_on_input.bp", u, on_input_mesh=True)
 ```
 
 See the [Checkpoint on input mesh](./original_checkpoint.py) for more details and examples.

docs/writing_functions_checkpoint.py

@@ -61,7 +61,7 @@ def f(x):
 # before reading in the function.
 # We **cannot** use {py:func}`io4dolfinx.read_function` to read function data to the
 # original `mesh` object.
-# To do this, we need to use {py:func}`io4dolfinx.write_function_on_input_mesh`, see
+# To do this, we need to use {py:func}`io4dolfinx.write_function` with `on_input_mesh=True`, see
 # [Writing function on input mesh checkpoint](./original_checkpoint).
 # for more details.
 # ```

src/io4dolfinx/__init__.py

@@ -19,11 +19,12 @@
     write_attributes,
     write_cell_data,
     write_function,
+    write_function_on_input_mesh,
     write_mesh,
+    write_mesh_input_order,
     write_meshtags,
     write_point_data,
 )
-from .original_checkpoint import write_function_on_input_mesh, write_mesh_input_order
 from .readers import (
     read_cell_data,
     read_function_from_legacy_h5,

src/io4dolfinx/backends/opencarp/backend.py

@@ -30,7 +30,9 @@
 .. code-block:: python
 
     args = {"scale": 1000.0, "markers": {"ENDO": (1, 2), "EPI": (2, 2)}}  # mesh in mm
-    io4dolfinx.write_mesh_input_order("heart", mesh, backend="opencarp", backend_args=args)
+    io4dolfinx.write_mesh(
+        "heart", mesh, backend="opencarp", backend_args=args, on_input_mesh=True
+    )
     for tags in (cell_tags, facet_tags):
         io4dolfinx.write_meshtags(
             "heart", mesh, tags, backend="opencarp", backend_args=args, on_input_mesh=True
@@ -50,14 +52,14 @@
   <io4dolfinx.FileMode.append>` adds the sheets. Fields are interpolated into DG-0,
   so fields in a quadrature space must first be generated in a P-1 or DG-0 space.
 
-The ordering of nodes and elements must be the same in every step:
+The ordering of nodes and elements must be the same in every step, so pass the same
+``on_input_mesh`` to every function:
 
-- With :func:`io4dolfinx.write_mesh_input_order` and ``on_input_mesh=True`` (as
-  above), the files use the ordering of the input mesh and are the same whatever the
-  number of processes. Use this for reproducible files.
-- With :func:`io4dolfinx.write_mesh` and ``on_input_mesh=False`` (the default of
-  :func:`io4dolfinx.write_meshtags` and :func:`io4dolfinx.write_cell_data`), the files
-  use the DOLFINx ordering, which depends on the partitioning.
+- With ``on_input_mesh=True`` (as above), the files use the ordering of the input
+  mesh and are the same whatever the number of processes. Use this for reproducible
+  files.
+- With ``on_input_mesh=False`` (the default), the files use the DOLFINx ordering,
+  which depends on the partitioning.
 
 Mixing the two gives directions and markers on the wrong elements and nodes.
 
@@ -658,8 +660,7 @@ def write_mesh(
     Args:
         filename: Base name of the files, or the path of one of them
         comm: MPI communicator used in storage
-        mesh: Mesh data, from :func:`io4dolfinx.write_mesh_input_order` or
-            :func:`io4dolfinx.write_mesh`
+        mesh: Mesh data, from :func:`io4dolfinx.write_mesh`
         backend_args: Arguments to backend, ``"scale"`` is required
         mode: Only :attr:`FileMode.write` is supported
         time: Not used

src/io4dolfinx/checkpointing.py

@@ -8,6 +8,7 @@
 
 import logging
 import typing
+import warnings
 from pathlib import Path
 from typing import Any, Callable
 
@@ -27,6 +28,7 @@
     send_dofs_and_recv_values,
     send_to_input_order,
 )
+from .original_checkpoint import create_function_data_on_original_mesh, create_original_mesh_data
 from .readers import create_geometry_function_space
 from .structures import ArrayData, FunctionData, MeshTagsData
 from .utils import (
@@ -569,6 +571,7 @@ def write_mesh(
     backend_args: dict[str, Any] | None = None,
     backend: str | None = None,
     mesh_name: str | None = None,
+    on_input_mesh: bool = False,
 ):
     """
     Write a mesh to file.
@@ -581,11 +584,26 @@ def write_mesh(
         mesh_name: Name of the mesh in the file to act on. Defaults to the
             mesh written without an explicit name, so existing checkpoints and
             call sites are unaffected.
+        on_input_mesh: If True, the mesh is written with the node and cell ordering of
+            the input mesh, so the file does not depend on the number of processes.
+            Cannot be combined with ``store_partition_info``.
     """
     backend_args = _with_mesh_name(backend_args, mesh_name)
     logger.debug(f"Writing mesh to {filename}")
-    logger.debug(f"Preparing mesh data for storage storing partition info: {store_partition_info}")
-    mesh_data = prepare_meshdata_for_storage(mesh=mesh, store_partition_info=store_partition_info)
+    if on_input_mesh:
+        if store_partition_info:
+            raise ValueError(
+                "Cannot store partition info for a mesh written in the input ordering."
+            )
+        logger.debug("Preparing mesh data in the input ordering")
+        mesh_data = create_original_mesh_data(mesh)
+    else:
+        logger.debug(
+            f"Preparing mesh data for storage storing partition info: {store_partition_info}"
+        )
+        mesh_data = prepare_meshdata_for_storage(
+            mesh=mesh, store_partition_info=store_partition_info
+        )
     logger.debug(f"Write mesh using {backend} backend, with arguments {backend_args}")
     logger.debug(f"Mode {mode} and time {time}")
     _internal_mesh_writer(
@@ -608,6 +626,7 @@ def write_function(
     backend_args: dict[str, Any] | None = None,
     backend: str | None = None,
     mesh_name: str | None = None,
+    on_input_mesh: bool = False,
 ):
     """
     Write function checkpoint to file.
@@ -623,11 +642,25 @@ def write_function(
         mesh_name: Name of the mesh in the file to act on. Defaults to the
             mesh written without an explicit name, so existing checkpoints and
             call sites are unaffected.
+        on_input_mesh: If True, the function is written against the input mesh, to be
+            read with a mesh written with ``on_input_mesh=True`` (or read from the
+            input mesh file).
     """
     backend_args = _with_mesh_name(backend_args, mesh_name)
     n = u.name if name is None else name
     logger.debug(f"Writing function checkpoint to {filename} for function {n} at time {time}")
     logger.debug(f"Using {backend} backend with arguments {backend_args}")
+    if on_input_mesh:
+        _internal_function_writer(
+            Path(filename),
+            u.function_space.mesh.comm,
+            create_function_data_on_original_mesh(u, name),
+            time,
+            backend_args=backend_args,
+            backend=backend,
+            mode=mode,
+        )
+        return
     dofmap = u.function_space.dofmap
     values = u.x.array
     mesh = u.function_space.mesh
@@ -689,6 +722,91 @@ def write_function(
     )
 
 
+def write_function_on_input_mesh(
+    filename: Path | str,
+    u: dolfinx.fem.Function,
+    time: float = 0.0,
+    name: str | None = None,
+    mode: FileMode = FileMode.append,
+    backend_args: dict[str, Any] | None = None,
+    backend: str | None = None,
+):
+    """
+    Write function checkpoint (to be read with the input mesh).
+
+    Deprecated: use :func:`write_function` with ``on_input_mesh=True``.
+
+    Note:
+        Requires backend to implement :meth:`io4dolfinx.backends.IOBackend.write_function`.
+
+    Args:
+        filename: The filename to write to
+        u: The function to checkpoint
+        time: Time-stamp associated with function at current write step
+        mode: The mode to use (write or append)
+        name: Name of function. If None, the name of the function is used.
+        backend_args: Arguments to backend
+        backend: Choice of backend module
+    """
+    warnings.warn(
+        "write_function_on_input_mesh is deprecated and will be removed in a future release. "
+        + "Use write_function with on_input_mesh=True instead.",
+        category=DeprecationWarning,
+        stacklevel=2,
+    )
+    write_function(
+        filename,
+        u,
+        time=time,
+        mode=mode,
+        name=name,
+        backend_args=backend_args,
+        backend=backend,
+        on_input_mesh=True,
+    )
+
+
+def write_mesh_input_order(
+    filename: Path | str,
+    mesh: dolfinx.mesh.Mesh,
+    time: float = 0.0,
+    mode: FileMode = FileMode.write,
+    backend: str | None = None,
+    backend_args: dict[str, Any] | None = None,
+):
+    """
+    Write mesh to checkpoint file in original input ordering.
+
+    Deprecated: use :func:`write_mesh` with ``on_input_mesh=True``.
+
+    Note:
+        Requires backend to implement :meth:`io4dolfinx.backends.IOBackend.write_mesh`.
+
+    Args:
+        filename: The filename to write to
+        mesh: Mesh to checkpoint
+        time: Time-stamp associated with function at current write step
+        mode: The mode to use (write or append)
+        backend_args: Arguments to backend
+        backend: Choice of backend module
+    """
+    warnings.warn(
+        "write_mesh_input_order is deprecated and will be removed in a future release. "
+        + "Use write_mesh with on_input_mesh=True instead.",
+        category=DeprecationWarning,
+        stacklevel=2,
+    )
+    write_mesh(
+        Path(filename),
+        mesh,
+        mode=mode,
+        time=time,
+        backend_args=backend_args,
+        backend=backend,
+        on_input_mesh=True,
+    )
+
+
 def read_function_names(
     filename: Path | str,
     comm: MPI.Comm,
@@ -740,7 +858,7 @@ def write_point_data(
         mesh_name: Name of the mesh in the file to write the data against.
             Defaults to the mesh written without an explicit name.
         on_input_mesh: If True, the node values are written in the node ordering of the
-            input mesh, matching a mesh written with :func:`write_mesh_input_order`.
+            input mesh, matching a mesh written with :func:`write_mesh` with ``on_input_mesh=True``.
     """
     backend_args = _with_mesh_name(backend_args, mesh_name)
     logger.debug(f"Writing point data to {filename} for function {u.name} at time {time}")
@@ -799,7 +917,7 @@ def write_cell_data(
         mesh_name: Name of the mesh in the file to write the data against.
             Defaults to the mesh written without an explicit name.
         on_input_mesh: If True, the cell values are written in the cell ordering of the
-            input mesh, matching a mesh written with :func:`write_mesh_input_order`.
+            input mesh, matching a mesh written with :func:`write_mesh` with ``on_input_mesh=True``.
     """
     backend_args = _with_mesh_name(backend_args, mesh_name)
     logger.debug(f"Writing cell data to {filename} for function {u.name} at time {time}")