Checkpointing with mixed function spaces

I/O
, ,
1

Hi,
I am trying to modify the original checkpointing for mixed function spaces.
The code is

# # Writing a function checkpoint
# In the previous sections, we have gone in to quite some detail as to how
# to store meshes with adios4dolfinx.
# This section will explain how to store functions, and how to read them back in.

# We start by creating a mesh and an appropriate function

import logging
from pathlib import Path

from mpi4py import MPI

import dolfinx
import ipyparallel as ipp
from basix.ufl import element, mixed_element
import adios4dolfinx

assert MPI.COMM_WORLD.size == 1, "This example should only be run with 1 MPI process"

mesh = dolfinx.mesh.create_unit_square(
    MPI.COMM_WORLD, nx=10, ny=10, cell_type=dolfinx.cpp.mesh.CellType.quadrilateral
)

# Next, we create a function, and interpolate a polynomial function into the function space
V_el = element("DG", mesh.basix_cell(),2, shape=(mesh.geometry.dim, ))
Q_el = element("DG", mesh.basix_cell(), 1)
W_el= mixed_element([V_el, Q_el])
W = dolfinx.fem.functionspace(mesh, W_el)

def f(x):
    return x[0],x[1]
def g(x):
    return -(x[0] ** 2)+ x[1] - 2 * x[0]

u = dolfinx.fem.Function(W)

(u.sub(0)).interpolate(f)
(u.sub(1)).interpolate(g)

# For the checkpointing, we start by storing the mesh to file

filename = Path("function_checkpoint.bp")
adios4dolfinx.write_mesh(filename, mesh)

# Next, we store the function to file, and associate it with a name.
# Note that we can also associate a time stamp with it,
u.sub(0).name="velocity"
u.sub(1).name="pressure"
adios4dolfinx.write_function(filename, u.sub(0),time=0.3, name="velocity")
adios4dolfinx.write_function(filename, u.sub(1),time=0.3, name="pressure")

def read_function0(filename: Path,timestamp: float):
    from mpi4py import MPI

    import dolfinx
    import numpy as np

    import adios4dolfinx
    from basix.ufl import element, mixed_element
    in_mesh = adios4dolfinx.read_mesh(filename, MPI.COMM_WORLD)
    V_el = element("DG", in_mesh.basix_cell(),2, shape=(in_mesh.geometry.dim, ))
    Q_el = element("DG", in_mesh.basix_cell(), 1)
    W_el= mixed_element([V_el, Q_el])
    W = dolfinx.fem.functionspace(in_mesh, W_el)
  #  W = dolfinx.fem.functionspace(in_mesh, (el, degree))
    u_ref = dolfinx.fem.Function(W)
    u_ref.sub(0).interpolate(f)
    u_ref.sub(1).interpolate(g)
    u_in = dolfinx.fem.Function(W)
    u_in.sub(0).name="velocity"
    u_in.sub(1).name="pressure"

    adios4dolfinx.read_function(filename, u_in.sub(0), time=timestamp, name="velocity")
    adios4dolfinx.read_function(filename, u_in.sub(1), time=timestamp, name="pressure")
    return u_in
u_comb=read_function0(filename, 0.3)

I am getting following error message:
Traceback (most recent call last):
File “/home/david/fenics/adios4dolfinx/docs_new/checkpointing_mixed.py”, line 78, in
u_comb=read_function0(filename, 0.3)
File “/home/david/fenics/adios4dolfinx/docs_new/checkpointing_mixed.py”, line 74, in read_function0
adios4dolfinx.read_function(filename, u_in.sub(0), time=timestamp, name=“velocity”)
File “/home/david/.local/lib/python3.10/site-packages/adios4dolfinx/checkpointing.py”, line 412, in read_function
input_cells = V.mesh.topology.original_cell_index[local_cells]
IndexError: index 102 is out of bounds for axis 0 with size 100

Any help to solve this is appreciated.
Thank You