[Date Prev][Date Next] [Thread Prev][Thread Next] [Date Index] [Thread Index]

Bug#1106702: [preauth] unblock: scifem/0.6.0-1



Package: release.debian.org
Severity: normal
X-Debbugs-Cc: scifem@packages.debian.org
Control: affects -1 + src:scifem
User: release.debian.org@packages.debian.org
Usertags: unblock

Please unblock package scifem

[ Reason ]

scifem made a recent release upstream after we entered the stable
release freeze.

The updates would be valuable to users
- updates blocked solver API
- updates read_mri_data_to_function
- fixes xdmf and adios2 file handling
and it could be good for long term users to make this release available
to them.

[ Impact ]

Without the unblock, users will have the previous release without
these fixes. With the blocked solver updates it means user's projects
using the package will not match upstream's. The long-term stable user
experience will be better if this release is allowed into trixie.

[ Tests ]

autopkgtest runs upstream tests. All tests are passing on all
architectures using scifem/0.6.0-1 from experimental
(loong64 has not yet launched its test)

[ Risks ]

scifem is a leaf project.  The package affects no other package, and
is used only by users who deliberately intend to use it.
(The fenicsx virtual package checks only that it installs and loads
in python).  Since tests are passing, there is negligible risk to trixie.

[ Checklist ]
  [x ] all changes are documented in the d/changelog
  [x ] I reviewed all changes and I approve them
  [x ] attach debdiff against the package in testing

[ Other info ]
scifem 0.6.0 is in experimental.  This unblock request is a request
for permission to upload to unstable for inclusion in trixie.

unblock scifem/0.6.0-1
diff -Nru scifem-0.5.0/debian/changelog scifem-0.6.0/debian/changelog
--- scifem-0.5.0/debian/changelog	2025-04-04 21:20:13.000000000 +0200
+++ scifem-0.6.0/debian/changelog	2025-05-27 12:30:36.000000000 +0200
@@ -1,3 +1,12 @@
+scifem (0.6.0-1) experimental; urgency=medium
+
+  * New upstream release
+    - updates blocked solver API
+    - updates read_mri_data_to_function
+    - fixes xdmf and adios2 file handling
+
+ -- Drew Parsons <dparsons@debian.org>  Tue, 27 May 2025 12:30:36 +0200
+
 scifem (0.5.0-4) unstable; urgency=medium
 
   * mpich (on 32-bit systems) sends warning messages to stderr,
diff -Nru scifem-0.5.0/examples/real_function_space.py scifem-0.6.0/examples/real_function_space.py
--- scifem-0.5.0/examples/real_function_space.py	2025-03-21 20:54:41.000000000 +0100
+++ scifem-0.6.0/examples/real_function_space.py	2025-05-23 21:28:34.000000000 +0200
@@ -59,6 +59,7 @@
 
 import numpy as np
 from scifem import create_real_functionspace, assemble_scalar
+from scifem.petsc import apply_lifting_and_set_bc, zero_petsc_vector
 import ufl
 
 M = 20
@@ -106,8 +107,10 @@
 L0 = ufl.inner(f, du) * ufl.dx + ufl.inner(g, du) * ufl.ds
 L1 = ufl.inner(zero, dl) * ufl.dx
 
-a = dolfinx.fem.form([[a00, a01], [a10, None]])
-L = dolfinx.fem.form([L0, L1])
+a = [[a00, a01], [a10, None]]
+L = [L0, L1]
+a_compiled = dolfinx.fem.form(a)
+L_compiled = dolfinx.fem.form(L)
 
 # Note that we have defined the variational form in a block form, and
 # that we have not included $h$ in the variational form. We will enforce this
@@ -115,29 +118,64 @@
 
 # We can now assemble the matrix and vector
 
-A = dolfinx.fem.petsc.assemble_matrix_block(a)
+try:
+    A = dolfinx.fem.petsc.assemble_matrix_block(a_compiled)
+except AttributeError:
+    A = dolfinx.fem.petsc.assemble_matrix(a_compiled, kind="mpi")
 A.assemble()
-b = dolfinx.fem.petsc.assemble_vector_block(L, a, bcs=[])
+
+
+# On the main branch of DOLFINx, the `assemble_vector` function for blocked spaces has been rewritten to reflect how
+# it works for standard assembly and `nest` assembly. This means that lifting is applied manually.
+# In this case, with no Dirichlet BC, we could skip those steps.
+# However, for clarity we include them here.
+
+bcs = []
+main_assembly = False
+try:
+    b = dolfinx.fem.petsc.assemble_vector_block(L_compiled, a_compiled, bcs=bcs)
+except AttributeError:
+    main_assembly = True
+    b = dolfinx.fem.petsc.assemble_vector(L_compiled, kind="mpi")
+    apply_lifting_and_set_bc(b, a_compiled, bcs=bcs)
 
 # Next, we modify the second part of the block to contain `h`
-# We start by enforcing the multiplier constraint $h$ by modifying the right hand side vector
+# We start by enforcing the multiplier constraint $h$ by modifying the right hand side vector.
+# On the main branch, this is greatly simplified
 
-if dolfinx.__version__ == "0.8.0":
-    maps = [(V.dofmap.index_map, V.dofmap.index_map_bs), (R.dofmap.index_map, R.dofmap.index_map_bs)]
-elif Version(dolfinx.__version__) >= Version("0.9.0.0"):
-    maps = [(Wi.dofmap.index_map, Wi.dofmap.index_map_bs) for Wi in W.ufl_sub_spaces()]
 
-b_local = get_local_vectors(b, maps)
-b_local[1][:] = h
-scatter_local_vectors(
-        b,
-        b_local,
-        maps,
-    )
-b.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
+uh = dolfinx.fem.Function(V, name="u")
 
-# We can now solve the linear system
+if main_assembly:
+    # We start by inserting the value in the real space
+    rh = dolfinx.fem.Function(R)
+    rh.x.array[0] = h
+    # Next we need to add this value to the existing right hand side vector.
+    # Therefore we create assign 0s to the primal space
+    b_real_space = b.duplicate()
+    uh.x.array[:] = 0
+    # Transfer the data to `b_real_space`
+    dolfinx.fem.petsc.assign([uh, rh], b_real_space)
+    # And accumulate the values in the right hand side vector
+    b.axpy(1, b_real_space)
+    # We destroy the temporary work vector after usage
+    b_real_space.destroy()
+else:
+    if Version(dolfinx.__version__) < Version("0.9.0"):
+        maps = [(V.dofmap.index_map, V.dofmap.index_map_bs), (R.dofmap.index_map, R.dofmap.index_map_bs)]
+    else:
+        maps = [(Wi.dofmap.index_map, Wi.dofmap.index_map_bs) for Wi in W.ufl_sub_spaces()]
+
+    b_local = get_local_vectors(b, maps)
+    b_local[1][:] = h
+    scatter_local_vectors(
+            b,
+            b_local,
+            maps,
+        )
+    b.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
 
+# We can now solve the linear system
 ksp = PETSc.KSP().create(mesh.comm)
 ksp.setOperators(A)
 ksp.setType("preonly")
@@ -145,17 +183,30 @@
 pc.setType("lu")
 pc.setFactorSolverType("mumps")
 
-xh = dolfinx.fem.petsc.create_vector_block(L)
+if main_assembly:
+    xh = b.duplicate()
+else:
+    xh = dolfinx.fem.petsc.create_vector_block(L_compiled)
+
 ksp.solve(b, xh)
 xh.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
 
 # Finally, we extract the solution u from the blocked system and compute the error
 
 uh = dolfinx.fem.Function(V, name="u")
-x_local = get_local_vectors(xh, maps)
-uh.x.array[: len(x_local[0])] = x_local[0]
-uh.x.scatter_forward()
-
+if main_assembly:
+    dolfinx.fem.petsc.assign(xh, [uh, rh])
+else:
+    x_local = get_local_vectors(xh, maps)
+    uh.x.array[: len(x_local[0])] = x_local[0]
+    uh.x.scatter_forward()
+
+# We destroy all PETSc objects
+
+b.destroy()
+xh.destroy()
+A.destroy()
+ksp.destroy()
 
 diff = uh - u_exact(x)
 error = dolfinx.fem.form(ufl.inner(diff, diff) * ufl.dx)
@@ -168,7 +219,7 @@
 
 import pyvista
 
-pyvista.start_xvfb()
+pyvista.start_xvfb(1.0)
 grid = pyvista.UnstructuredGrid(*vtk_mesh)
 grid.point_data["u"] = uh.x.array.real
 
diff -Nru scifem-0.5.0/.pre-commit-config.yaml scifem-0.6.0/.pre-commit-config.yaml
--- scifem-0.5.0/.pre-commit-config.yaml	2025-03-21 20:54:41.000000000 +0100
+++ scifem-0.6.0/.pre-commit-config.yaml	2025-05-23 21:28:34.000000000 +0200
@@ -15,7 +15,7 @@
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: 'v0.11.0'
+    rev: 'v0.11.6'
     hooks:
        # Run the linter.
       - id: ruff
diff -Nru scifem-0.5.0/pyproject.toml scifem-0.6.0/pyproject.toml
--- scifem-0.5.0/pyproject.toml	2025-03-21 20:54:41.000000000 +0100
+++ scifem-0.6.0/pyproject.toml	2025-05-23 21:28:34.000000000 +0200
@@ -4,7 +4,7 @@
 
 [project]
 name = "scifem"
-version = "0.5.0"
+version = "0.6.0"
 description = "Scientific tools for finite element methods"
 readme = "README.md"
 requires-python = ">=3.10"
@@ -26,6 +26,7 @@
 test = ["pytest", "petsc4py", "h5py", "scifem[biomed]"]
 biomed = ["nibabel"]
 all = ["scifem[docs,dev,test,audio2,h5py,biomed]"]
+petsc = ["petsc4py"]
 
 [tool.pytest.ini_options]
 testpaths = ["tests"]
@@ -125,7 +126,7 @@
 sign_tags = false
 tag_name = "v{new_version}"
 tag_message = "Bump version: {current_version} → {new_version}"
-current_version = "0.5.0"
+current_version = "0.6.0"
 
 
 [[tool.bumpversion.files]]
diff -Nru scifem-0.5.0/src/scifem/biomedical.py scifem-0.6.0/src/scifem/biomedical.py
--- scifem-0.5.0/src/scifem/biomedical.py	2025-03-21 20:54:41.000000000 +0100
+++ scifem-0.6.0/src/scifem/biomedical.py	2025-05-23 21:28:34.000000000 +0200
@@ -3,20 +3,26 @@
 from pathlib import Path
 import numpy as np
 import numpy.typing as npt
+import warnings
 
 
 def apply_mri_transform(
-    path: Path, coordinates: npt.NDArray[np.floating], use_tkr: bool = False
+    path: Path,
+    coordinates: npt.NDArray[np.floating],
+    vox2ras_transform: npt.NDArray[np.floating] | None = None,
+    use_tkr: bool = False,
 ) -> npt.NDArray[np.inexact]:
     """Given a path to a set of MRI voxel data, return the data evaluated at the given coordinates.
 
     Args:
         path: Path to the MRI data, should be a file format supported by nibabel.
         coordinates: Coordinates to evaluate the MRI data at.
+        vox2ras_transform: Optional transformation matrix to convert from voxel to ras coordinates.
+            If None, use the transformation matrix from the given MRI data.
         use_tkr: If true, use the old freesurfer `tkregister`, see:
             https://www.mail-archive.com/freesurfer@nmr.mgh.harvard.edu/msg69541.html
             for more details. Else use the standard VOX2RAS transform (equivalent to attaching
-            an affine map to a nibabel image.
+            an affine map to a nibabel image.)
 
     Returns:
         The data evaluated at these points.
@@ -28,21 +34,43 @@
         raise ImportError("This function requires the nibabel package to be installed")
     image = nibabel.load(path)
     data = image.get_fdata()
+    # Check if the image is in the correct orientation
+    orientation = nibabel.aff2axcodes(image.affine)
+    if orientation != ("R", "A", "S"):
+        warnings.warn(f"Image orientation is {orientation}, expected (R, A, S). ")
+
+    # Define mgh_header depending on MRI data types (i.e. nifti or mgz)
+    if isinstance(image, nibabel.freesurfer.mghformat.MGHImage):
+        mgh_header = image.header
+    elif isinstance(image, nibabel.nifti1.Nifti1Image):
+        mgh = nibabel.MGHImage(image.dataobj, image.affine)
+        mgh_header = mgh.header
+    else:
+        raise ValueError(
+            f"Unsupported image type: {type(image)} - Supported types are mgz and nifti"
+        )
+
     # This depends on how one uses FreeSurfer
-    if use_tkr:
-        vox2ras = image.header.get_vox2ras_tkr()
+    if vox2ras_transform is not None:
+        vox2ras = vox2ras_transform
     else:
-        # VOX to ras explanation: https://surfer.nmr.mgh.harvard.edu/ftp/articles/vox2ras.pdf
-        vox2ras = image.header.get_vox2ras()
+        if use_tkr:
+            vox2ras = mgh_header.get_vox2ras_tkr()
+        else:
+            # VOX to ras explanation: https://surfer.nmr.mgh.harvard.edu/ftp/articles/vox2ras.pdf
+            vox2ras = mgh_header.get_vox2ras()
+    # Check shape
+    if vox2ras.shape != (4, 4):
+        raise ValueError(f"vox2ras transform must be a 4x4 matrix, got shape {vox2ras.shape}")
+
     ras2vox = np.linalg.inv(vox2ras)
-    ijk_vectorized = naff.apply_affine(ras2vox, coordinates)
 
-    # Round indices to nearest integer
+    ijk_vectorized = naff.apply_affine(ras2vox, coordinates)
     # The file standard assumes that the voxel coordinates refer to the center of each voxel.
     # https://brainder.org/2012/09/23/the-nifti-file-format/
     ijk_rounded = np.rint(ijk_vectorized).astype("int")
-
     assert np.all(ijk_rounded >= 0)
+
     return data[ijk_rounded[:, 0], ijk_rounded[:, 1], ijk_rounded[:, 2]]
 
 
@@ -51,6 +79,8 @@
     mesh: dolfinx.mesh.Mesh,
     edim: int,
     entities: npt.NDArray[np.int32] | None = None,
+    vox2ras_transform: npt.NDArray[np.floating] | None = None,
+    use_tkr: bool = False,
 ) -> dolfinx.mesh.MeshTags:
     """Read in MRI data over a set of entities in the mesh and attach it to a mesh tag.
 
@@ -71,7 +101,7 @@
         entities = np.arange(num_entities, dtype=np.int32)
 
     midpoints = dolfinx.mesh.compute_midpoints(mesh, edim, entities)
-    data = apply_mri_transform(Path(mri_data_path), midpoints)
+    data = apply_mri_transform(Path(mri_data_path), midpoints, vox2ras_transform, use_tkr=use_tkr)
     et = dolfinx.mesh.meshtags(mesh, edim, entities, data.astype(np.int32))
     return et
 
@@ -80,8 +110,10 @@
     mri_data_path: str | Path,
     mesh: dolfinx.mesh.Mesh,
     cells: npt.NDArray[np.int32] | None = None,
+    vox2ras_transform: npt.NDArray[np.floating] | None = None,
     degree: int = 0,
     dtype: np.inexact = dolfinx.default_scalar_type,
+    use_tkr: bool = False,
 ) -> dolfinx.fem.Function:
     """Read in MRI data over a set of cells in the mesh and attach it to an appropriate function.
 
@@ -89,10 +121,17 @@
         mri_data_path: Path to MRI data
         mesh: The mesh to attach the MRI data to.
         cells: Subset of cells to evaluate the MRI data at. If None, all cells are used.
+        vox2ras_transform: Optional transformation matrix to convert from voxel to ras coordinates.
+            If None, use the transformation matrix from the given MRI data.
         degree: Degree of the (quadrature) function space to attach the MRI data to. Defaults to 0.
             If degree is 0, use a DG-0 space instead of a quadrature space, to simplify
             post-processing.
         dtype: Data type used for input data. Can be used for rounding data.
+        use_tkr: If true, use the old freesurfer `tkregister`, see:
+            https://www.mail-archive.com/freesurfer@nmr.mgh.harvard.edu/msg69541.html
+            for more details. Else use the standard VOX2RAS transform (equivalent to attaching
+            an affine map to a nibabel image.)
+
 
     Raises:
         ValueError: If degree is negative.
@@ -117,7 +156,9 @@
     dof_positions = Vh.tabulate_dof_coordinates()
     dofmap = Vh.dofmap.list
     dofmap_pos = dofmap[cells].flatten()
-    data = apply_mri_transform(Path(mri_data_path), dof_positions[dofmap_pos])
+    data = apply_mri_transform(
+        Path(mri_data_path), dof_positions[dofmap_pos], vox2ras_transform, use_tkr=use_tkr
+    )
 
     v = dolfinx.fem.Function(Vh)
     v.x.array[dofmap_pos] = data.astype(dtype)
diff -Nru scifem-0.5.0/src/scifem/petsc.py scifem-0.6.0/src/scifem/petsc.py
--- scifem-0.5.0/src/scifem/petsc.py	1970-01-01 01:00:00.000000000 +0100
+++ scifem-0.6.0/src/scifem/petsc.py	2025-05-23 21:28:34.000000000 +0200
@@ -0,0 +1,106 @@
+import dolfinx
+import typing
+
+__all__ = [
+    "zero_petsc_vector",
+    "ghost_update",
+    "apply_lifting_and_set_bc",
+]
+
+try:
+    from petsc4py import PETSc
+
+    def zero_petsc_vector(b):
+        """Zero a PETSc vector, including ghosts"""
+
+        if b.getType() == PETSc.Vec.Type.NEST:
+            for b_sub in b.getNestSubVecs():
+                with b_sub.localForm() as b_local:
+                    b_local.set(0.0)
+        else:
+            with b.localForm() as b_loc:
+                b_loc.set(0)
+
+    def ghost_update(x, insert_mode, scatter_mode):
+        """Ghost update a vector"""
+        if x.getType() == PETSc.Vec.Type.NEST:
+            for x_sub in x.getNestSubVecs():
+                x_sub.ghostUpdate(addv=insert_mode, mode=scatter_mode)
+        else:
+            x.ghostUpdate(addv=insert_mode, mode=scatter_mode)
+
+except ImportError:
+
+    def zero_petsc_vector(_b):
+        """Zero a PETSc vector, including ghosts"""
+        raise RuntimeError("petsc4py is not available. Cannot zero vector.")
+
+    def ghost_update(_x, _insert_mode, _scatter_mode):
+        """Ghost update a vector"""
+        raise RuntimeError("petsc4py is not available. Cannot ghost update vector.")
+
+
+if dolfinx.has_petsc4py:
+    from petsc4py import PETSc
+    import dolfinx.fem.petsc
+
+    def apply_lifting_and_set_bc(
+        b: PETSc.Vec,
+        a: typing.Union[
+            typing.Iterable[dolfinx.fem.Form], typing.Iterable[typing.Iterable[dolfinx.fem.Form]]
+        ],
+        bcs: typing.Union[
+            typing.Iterable[dolfinx.fem.DirichletBC],
+            typing.Iterable[typing.Iterable[dolfinx.fem.DirichletBC]],
+        ],
+        x: typing.Optional[PETSc.Vec] = None,
+        alpha: float = 1.0,
+    ):
+        """Apply lifting to a vector and set boundary conditions.
+
+        Convenience function to apply lifting and set boundary conditions for multiple matrix types.
+        This modifies the vector b such that
+
+        .. math::
+            b_{free} = b - alpha a[i] (u_bc[i] - x[i])
+            b_{bc} = u_bc[i]
+
+        where :math:`b_{free}` is the free part of the vector, :math:`b_{bc}` is the part that has
+        boundary conditions applied.
+
+        Args:
+            b: The vector to apply lifting to.
+            a: Sequence fo forms to apply lifting from. If the system is blocked or nested,
+                his is a nested list of forms.
+            bcs: The boundary conditions to apply. If the form is blocked or nested, this is a list,
+                while if it is a single form, this is a nested list.
+            x: Vector to subtract from the boundary conditions. Usually used in a Newton iteration.
+            alpha: The scaling factor for the boundary conditions.
+        """
+        if hasattr(dolfinx.fem.petsc, "create_vector_nest"):
+            raise NotImplementedError(
+                "This function is only implemented for later versions of DOLFINx."
+            )
+
+        try:
+            bcs0 = dolfinx.fem.bcs_by_block(dolfinx.fem.extract_function_spaces(a, 0), bcs)
+            bcs1 = dolfinx.fem.bcs_by_block(dolfinx.fem.extract_function_spaces(a, 1), bcs)
+        except AssertionError:
+            bcs0 = bcs
+            bcs1 = bcs
+
+        dolfinx.fem.petsc.apply_lifting(b, a, bcs=bcs1, x0=x, alpha=alpha)
+        ghost_update(b, PETSc.InsertMode.ADD_VALUES, PETSc.ScatterMode.REVERSE)
+        try:
+            dolfinx.fem.petsc.set_bc(b, bcs0, x0=x, alpha=alpha)
+        except AttributeError:
+            for _bcs in bcs0:
+                dolfinx.fem.petsc.set_bc(b, _bcs, x0=x, alpha=alpha)
+        ghost_update(b, PETSc.InsertMode.INSERT_VALUES, PETSc.ScatterMode.FORWARD)
+
+else:
+
+    def apply_lifting_and_set_bc(_b, _a, _bcs, _x, _alpha):  # type: ignore
+        raise RuntimeError(
+            "petsc4py is not available. Cannot apply lifting and set boundary conditions."
+        )
diff -Nru scifem-0.5.0/src/scifem/solvers.py scifem-0.6.0/src/scifem/solvers.py
--- scifem-0.5.0/src/scifem/solvers.py	2025-03-21 20:54:41.000000000 +0100
+++ scifem-0.6.0/src/scifem/solvers.py	2025-05-23 21:28:34.000000000 +0200
@@ -5,435 +5,547 @@
 from packaging.version import parse as _v
 
 import numpy as np
-from petsc4py import PETSc
 import ufl
 import dolfinx
 
-
 __all__ = ["NewtonSolver", "BlockedNewtonSolver"]
 
-logger = logging.getLogger(__name__)
+if dolfinx.has_petsc4py and dolfinx.has_petsc:
+    logger = logging.getLogger(__name__)
+    from petsc4py import PETSc
+    import dolfinx.fem.petsc
+
+    # assemble_vector_block(scale=...) is renamed assemble_vector_block(alpha=...)
+    # in 0.9
+    _alpha_kw: str = "alpha"
+    if _v(dolfinx.__version__) < _v("0.9"):
+        _alpha_kw = "scale"
+
+    class NewtonSolver:
+        max_iterations: int
+        _bcs: list[dolfinx.fem.DirichletBC]
+        A: PETSc.Mat
+        b: PETSc.Vec
+        _J: dolfinx.fem.Form | ufl.form.Form
+        _F: dolfinx.fem.Form | ufl.form.Form
+        _dx: PETSc.Vec
+        _error_on_convergence: bool
+        _pre_solve_callback: Callable[["NewtonSolver"], None] | None
+        _post_solve_callback: Callable[["NewtonSolver"], None] | None
+
+        def __init__(
+            self,
+            F: list[dolfinx.fem.Form],
+            J: list[list[dolfinx.fem.Form]],
+            w: list[dolfinx.fem.Function],
+            bcs: list[dolfinx.fem.DirichletBC] | None = None,
+            max_iterations: int = 5,
+            petsc_options: dict[str, str | float | int | None] | None = None,
+            error_on_nonconvergence: bool = True,
+        ):
+            """
+            Create a Newton solver for a block nonlinear problem ``F(u) = 0``.
+            Solved as ``J(u) du = -F(u)``, where ``J`` is the Jacobian of ``F``.
+
+            Args:
+                F: List of forms defining the residual
+                J: List of lists of forms defining the Jacobian
+                w: List of functions representing the solution
+                bcs: List of Dirichlet boundary conditions
+                max_iterations: Maximum number of iterations in Newton solver
+                petsc_options: PETSc options for Krylov subspace solver.
+                error_on_nonconvergence: Raise an error if the linear solver
+                    does not converge or Newton solver doesn't converge
+            """
+            # Compile forms if not compiled. Will throw error it requires entity maps
+            self._F = dolfinx.fem.form(F)
+            self._J = dolfinx.fem.form(J)
+
+            # Store solution and accessible/modifiable properties
+            self._pre_solve_callback = None
+            self._post_solve_callback = None
+            self.max_iterations = max_iterations
+            self._error_on_convergence = error_on_nonconvergence
+            self.bcs = [] if bcs is None else bcs
+            self.w = w
+
+            # Create PETSc objects for block assembly
+            try:
+                self.b = dolfinx.fem.petsc.create_vector_block(self._F)
+                self.A = dolfinx.fem.petsc.create_matrix_block(self._J)
+                self.dx = dolfinx.fem.petsc.create_vector_block(self._F)
+                self.x = dolfinx.fem.petsc.create_vector_block(self._F)
+            except AttributeError:
+                self.b = dolfinx.fem.petsc.create_vector(self._F, kind="mpi")
+                self.A = dolfinx.fem.petsc.create_matrix(self._J, kind="mpi")
+                self.dx = dolfinx.fem.petsc.create_vector(self._F, kind="mpi")
+                self.x = dolfinx.fem.petsc.create_vector(self._F, kind="mpi")
 
-# assemble_vector_block(scale=...) is renamed assemble_vector_block(alpha=...)
-# in 0.9
-_alpha_kw: str = "alpha"
-if _v(dolfinx.__version__) < _v("0.9"):
-    _alpha_kw = "scale"
-
-
-class NewtonSolver:
-    max_iterations: int
-    _bcs: list[dolfinx.fem.DirichletBC]
-    A: PETSc.Mat
-    b: PETSc.Vec
-    _J: dolfinx.fem.Form | ufl.form.Form
-    _F: dolfinx.fem.Form | ufl.form.Form
-    _dx: PETSc.Vec
-    _error_on_convergence: bool
-    _pre_solve_callback: Callable[["NewtonSolver"], None] | None
-    _post_solve_callback: Callable[["NewtonSolver"], None] | None
-
-    def __init__(
-        self,
-        F: list[dolfinx.fem.Form],
-        J: list[list[dolfinx.fem.Form]],
-        w: list[dolfinx.fem.Function],
-        bcs: list[dolfinx.fem.DirichletBC] | None = None,
-        max_iterations: int = 5,
-        petsc_options: dict[str, str | float | int | None] | None = None,
-        error_on_nonconvergence: bool = True,
-    ):
-        """
-        Create a Newton solver for a block nonlinear problem ``F(u) = 0``.
-        Solved as ``J(u) du = -F(u)``, where ``J`` is the Jacobian of ``F``.
-
-        Args:
-            F: List of forms defining the residual
-            J: List of lists of forms defining the Jacobian
-            w: List of functions representing the solution
-            bcs: List of Dirichlet boundary conditions
-            max_iterations: Maximum number of iterations in Newton solver
-            petsc_options: PETSc options for Krylov subspace solver.
-            error_on_nonconvergence: Raise an error if the linear solver
-                does not converge or Newton solver doesn't converge
-        """
-        # Compile forms if not compiled. Will throw error it requires entity maps
-        self._F = dolfinx.fem.form(F)
-        self._J = dolfinx.fem.form(J)
-
-        # Store solution and accessible/modifiable properties
-        self._pre_solve_callback = None
-        self._post_solve_callback = None
-        self.max_iterations = max_iterations
-        self._error_on_convergence = error_on_nonconvergence
-        self.bcs = [] if bcs is None else bcs
-        self.w = w
-
-        # Create PETSc objects for block assembly
-        self.b = dolfinx.fem.petsc.create_vector_block(self._F)
-        self.A = dolfinx.fem.petsc.create_matrix_block(self._J)
-        self.dx = dolfinx.fem.petsc.create_vector_block(self._F)
-        self.x = dolfinx.fem.petsc.create_vector_block(self._F)
-
-        # Set PETSc options
-        opts = PETSc.Options()
-        if petsc_options is not None:
-            for k, v in petsc_options.items():
-                opts[k] = v
-
-        # Define KSP solver
-        self._solver = PETSc.KSP().create(self.b.getComm().tompi4py())
-        self._solver.setOperators(self.A)
-        self._solver.setFromOptions()
-
-        # Set matrix and vector PETSc options
-        self.A.setFromOptions()
-        self.b.setFromOptions()
-
-    def set_pre_solve_callback(self, callback: Callable[["NewtonSolver"], None]):
-        """Set a callback function that is called before each Newton iteration."""
-        self._pre_solve_callback = callback
-
-    def set_post_solve_callback(self, callback: Callable[["NewtonSolver"], None]):
-        """Set a callback function that is called after each Newton iteration."""
-        self._post_solve_callback = callback
-
-    def solve(self, atol=1e-6, rtol=1e-8, beta=1.0) -> int:
-        """Solve the nonlinear problem using Newton's method.
-
-        Args:
-            atol: Absolute tolerance for the update.
-            rtol: Relative tolerance for the update.
-            beta: Damping parameter for the update.
-
-        Returns:
-            The number of Newton iterations used to converge.
-
-        Note:
-            The tolerance is on the 0-norm of the update.
-        """
-        i = 1
+                self._maps = [
+                    (
+                        form.function_spaces[0].dofmaps(0).index_map,
+                        form.function_spaces[0].dofmaps(0).index_map_bs,
+                    )
+                    for form in self.F
+                ]
 
-        while i <= self.max_iterations:
-            if self._pre_solve_callback is not None:
-                self._pre_solve_callback(self)
+            # Set PETSc options
+            opts = PETSc.Options()
+            if petsc_options is not None:
+                for k, v in petsc_options.items():
+                    opts[k] = v
+
+            # Define KSP solver
+            self._solver = PETSc.KSP().create(self.b.getComm().tompi4py())
+            self._solver.setOperators(self.A)
+            self._solver.setFromOptions()
+
+            # Set matrix and vector PETSc options
+            self.A.setFromOptions()
+            self.b.setFromOptions()
+
+        def set_pre_solve_callback(self, callback: Callable[["NewtonSolver"], None]):
+            """Set a callback function that is called before each Newton iteration."""
+            self._pre_solve_callback = callback
+
+        def set_post_solve_callback(self, callback: Callable[["NewtonSolver"], None]):
+            """Set a callback function that is called after each Newton iteration."""
+            self._post_solve_callback = callback
+
+        def solve(self, atol=1e-6, rtol=1e-8, beta=1.0) -> int:
+            """Solve the nonlinear problem using Newton's method.
+
+            Args:
+                atol: Absolute tolerance for the update.
+                rtol: Relative tolerance for the update.
+                beta: Damping parameter for the update.
+
+            Returns:
+                The number of Newton iterations used to converge.
+
+            Note:
+                The tolerance is on the 0-norm of the update.
+            """
+            i = 1
+
+            while i <= self.max_iterations:
+                if self._pre_solve_callback is not None:
+                    self._pre_solve_callback(self)
+
+                # Pack constants and coefficients
+                try:
+                    constants_L = dolfinx.fem.pack_constants(self._F)
+                    coeffs_L = dolfinx.fem.pack_coefficients(self._F)
+                    constants_a = [
+                        dolfinx.fem.pack_constants(form)
+                        if form is not None
+                        else np.array([], dtype=self.x.array.dtype)
+                        for form in self._J
+                    ]
+                    coeffs_a = [
+                        {} if form is None else dolfinx.fem.pack_coefficients(form)
+                        for form in self._J
+                    ]
+                except AttributeError:
+                    # NOTE: DOLFINx 0.9 compatibility
+                    # Pack constants and coefficients
+                    constants_L = [
+                        form and dolfinx.cpp.fem.pack_constants(form._cpp_object)
+                        for form in self._F
+                    ]
+                    coeffs_L = [
+                        dolfinx.cpp.fem.pack_coefficients(form._cpp_object) for form in self._F
+                    ]
+                    constants_a = [
+                        [
+                            dolfinx.cpp.fem.pack_constants(form._cpp_object)
+                            if form is not None
+                            else np.array([], dtype=self.x.array.dtype)
+                            for form in forms
+                        ]
+                        for forms in self._J
+                    ]
+                    coeffs_a = [
+                        [
+                            {}
+                            if form is None
+                            else dolfinx.cpp.fem.pack_coefficients(form._cpp_object)
+                            for form in forms
+                        ]
+                        for forms in self._J
+                    ]
+                # Scatter previous solution `w` to `self.x`, the blocked version used for lifting
+                dolfinx.cpp.la.petsc.scatter_local_vectors(
+                    self.x,
+                    [si.x.petsc_vec.array_r for si in self.w],
+                    [
+                        (
+                            si.function_space.dofmap.index_map,
+                            si.function_space.dofmap.index_map_bs,
+                        )
+                        for si in self.w
+                    ],
+                )
+                self.x.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
 
-            # Pack constants and coefficients
-            constants_L = [
-                form and dolfinx.cpp.fem.pack_constants(form._cpp_object) for form in self._F
-            ]
-            coeffs_L = [dolfinx.cpp.fem.pack_coefficients(form._cpp_object) for form in self._F]
-            constants_a = [
-                [
-                    dolfinx.cpp.fem.pack_constants(form._cpp_object)
-                    if form is not None
-                    else np.array([], dtype=PETSc.ScalarType)
-                    for form in forms
-                ]
-                for forms in self._J
-            ]
-            coeffs_a = [
-                [
-                    {} if form is None else dolfinx.cpp.fem.pack_coefficients(form._cpp_object)
-                    for form in forms
+                # Assemble F(u_{i-1}) - J(u_D - u_{i-1}) and set du|_bc= u_D - u_{i-1}
+                with self.b.localForm() as b_local:
+                    b_local.set(0.0)
+                try:
+                    dolfinx.fem.petsc.assemble_vector_block(
+                        self.b,
+                        self._F,
+                        self._J,
+                        bcs=self.bcs,
+                        x0=self.x,
+                        coeffs_a=coeffs_a,
+                        constants_a=constants_a,
+                        coeffs_L=coeffs_L,
+                        constants_L=constants_L,
+                        # dolfinx 0.8 compatibility
+                        # this is called 'scale' in 0.8, 'alpha' in 0.9
+                        **{_alpha_kw: -1.0},
+                    )
+                except AttributeError:
+                    dolfinx.fem.petsc.assemble_vector(
+                        self.b, self._F, coeffs=coeffs_L, constants=constants_L
+                    )
+                    bcs1 = dolfinx.fem.bcs_by_block(
+                        dolfinx.fem.extract_function_spaces(self._J, 1), self.bcs
+                    )
+                    dolfinx.fem.petsc.apply_lifting(
+                        self.b,
+                        self._J,
+                        bcs=bcs1,
+                        x0=self.x,
+                        coeffs=coeffs_a,
+                        constants=constants_a,
+                        alpha=-1.0,
+                    )
+                    self.b.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
+                    bcs0 = dolfinx.fem.bcs_by_block(
+                        dolfinx.fem.extract_function_spaces(self._F), self.bcs
+                    )
+                    dolfinx.fem.petsc.set_bc(self.b, bcs0, x0=self.x, alpha=-1.0)
+                    self.b.ghostUpdate(PETSc.InsertMode.INSERT_VALUES, PETSc.ScatterMode.FORWARD)
+
+                # Assemble Jacobian
+                self.A.zeroEntries()
+                try:
+                    dolfinx.fem.petsc.assemble_matrix_block(
+                        self.A, self._J, bcs=self.bcs, constants=constants_a, coeffs=coeffs_a
+                    )
+                except AttributeError:
+                    dolfinx.fem.petsc.assemble_matrix(
+                        self.A, self._J, self.bcs, coeffs=coeffs_a, constants=constants_a
+                    )
+                self.A.assemble()
+
+                self._solver.solve(self.b, self.dx)
+                if self._error_on_convergence:
+                    if (status := self._solver.getConvergedReason()) <= 0:
+                        raise RuntimeError(f"Linear solver did not converge, got reason: {status}")
+
+                # Update solution
+                offset_start = 0
+                for s in self.w:
+                    num_sub_dofs = (
+                        s.function_space.dofmap.index_map.size_local
+                        * s.function_space.dofmap.index_map_bs
+                    )
+                    s.x.petsc_vec.array_w[:num_sub_dofs] -= (
+                        beta * self.dx.array_r[offset_start : offset_start + num_sub_dofs]
+                    )
+                    s.x.petsc_vec.ghostUpdate(
+                        addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD
+                    )
+                    offset_start += num_sub_dofs
+
+                if self._post_solve_callback is not None:
+                    self._post_solve_callback(self)
+
+                # Compute norm of update
+                residual = self.dx.norm(PETSc.NormType.NORM_2)
+                if i == 1:
+                    self.residual_0 = residual
+                relative_residual = residual / max(self.residual_0, atol)
+
+                logger.info(
+                    f"Newton iteration {i}"
+                    f": r (abs) = {residual} (tol={atol}), "
+                    f"r (rel) = {relative_residual} (tol={rtol})"
+                )
+                if relative_residual < rtol or residual < atol:
+                    return i
+                i += 1
+
+            if self._error_on_convergence:
+                raise RuntimeError("Newton solver did not converge")
+            else:
+                return self.max_iterations
+
+        @property
+        def F(self):
+            """The list of residuals where each entry is a ``dolfinx.fem.Form``."""
+            return self._F
+
+        @property
+        def J(self):
+            """
+            The Jacobian blocks represented as lists of lists where each entry
+            is a ``dolfinx.fem.Form``.
+            """
+            return self._J
+
+        def __del__(self):
+            """Clean up the solver by destroying PETSc objects."""
+            self.A.destroy()
+            self.b.destroy()
+            self.dx.destroy()
+            self._solver.destroy()
+            self.x.destroy()
+
+    class BlockedNewtonSolver(dolfinx.cpp.nls.petsc.NewtonSolver):
+        def __init__(
+            self,
+            F: list[ufl.form.Form],
+            u: list[dolfinx.fem.Function],
+            bcs: list[dolfinx.fem.DirichletBC] = [],
+            J: list[list[ufl.form.Form]] | None = None,
+            form_compiler_options: dict | None = None,
+            jit_options: dict | None = None,
+            petsc_options: dict | None = None,
+            entity_maps: dict | None = None,
+        ):
+            """Initialize solver for solving a non-linear problem using Newton's method.
+            Args:
+                F: List of PDE residuals [F_0(u, v_0), F_1(u, v_1), ...]
+                u: List of unknown functions u=[u_0, u_1, ...]
+                bcs: List of Dirichlet boundary conditions
+                J: UFL representation of the Jacobian (Optional)
+                    Note:
+                        If not provided, the Jacobian will be computed using the
+                        assumption that the test functions come from a ``ufl.MixedFunctionSpace``
+                form_compiler_options: Options used in FFCx
+                    compilation of this form. Run ``ffcx --help`` at the
+                    command line to see all available options.
+                jit_options: Options used in CFFI JIT compilation of C
+                    code generated by FFCx. See ``python/dolfinx/jit.py``
+                    for all available options. Takes priority over all
+                    other option values.
+                petsc_options:
+                    Options passed to the PETSc Krylov solver.
+                entity_maps: Maps used to map entities between different meshes.
+                    Only needed if the forms have not been compiled a priori,
+                    and has coefficients, test, or trial functions that are defined on
+                    different meshes.
+            """
+            # Initialize base class
+            super().__init__(u[0].function_space.mesh.comm)
+
+            # Set PETSc options for Krylov solver
+            prefix = self.krylov_solver.getOptionsPrefix()
+            if prefix is None:
+                prefix = ""
+            if petsc_options is not None:
+                # Set PETSc options
+                opts = PETSc.Options()
+                opts.prefixPush(prefix)
+                for k, v in petsc_options.items():
+                    opts[k] = v
+                opts.prefixPop()
+                self.krylov_solver.setFromOptions()
+            self._F = dolfinx.fem.form(
+                F,
+                form_compiler_options=form_compiler_options,
+                jit_options=jit_options,
+                entity_maps=entity_maps,
+            )
+
+            # Create the Jacobian matrix, dF/du
+            if J is None:
+                if _v(dolfinx.__version__) < _v("0.9"):
+                    raise RuntimeError(
+                        "Automatic computation of Jacobian for blocked problem is only"
+                        + "supported in DOLFINx 0.9 and later"
+                    )
+                du = ufl.TrialFunctions(ufl.MixedFunctionSpace(*[ui.function_space for ui in u]))
+                J = ufl.extract_blocks(
+                    sum(ufl.derivative(sum(F), u[i], du[i]) for i in range(len(u)))
+                )
+            self._a = dolfinx.fem.form(
+                J,
+                form_compiler_options=form_compiler_options,
+                jit_options=jit_options,
+                entity_maps=entity_maps,
+            )
+
+            self._bcs = bcs
+            self._u = u
+            self._pre_solve_callback: Callable[["BlockedNewtonSolver"], None] | None = None
+            self._post_solve_callback: Callable[["BlockedNewtonSolver"], None] | None = None
+
+            # Create structures for holding arrays and matrix
+            try:
+                self._b = dolfinx.fem.petsc.create_vector_block(self._F)
+                self._J = dolfinx.fem.petsc.create_matrix_block(self._a)
+                self._dx = dolfinx.fem.petsc.create_vector_block(self._F)
+                self._x = dolfinx.fem.petsc.create_vector_block(self._F)
+            except AttributeError:
+                self._b = dolfinx.fem.petsc.create_vector(self._F, kind="mpi")
+                self._J = dolfinx.fem.petsc.create_matrix(self._a, kind="mpi")
+                self._dx = dolfinx.fem.petsc.create_vector(self._F, kind="mpi")
+                self._x = dolfinx.fem.petsc.create_vector(self._F, kind="mpi")
+
+                self._maps = [
+                    (
+                        form.function_spaces[0].dofmaps(0).index_map,
+                        form.function_spaces[0].dofmaps(0).index_map_bs,
+                    )
+                    for form in self._F
                 ]
-                for forms in self._J
-            ]
 
-            # Scatter previous solution `w` to `self.x`, the blocked version used for lifting
+            self._J.setOptionsPrefix(prefix)
+            self._J.setFromOptions()
+
+            self.setJ(self._assemble_jacobian, self._J)
+            self.setF(self._assemble_residual, self._b)
+            self.set_form(self._pre_newton_iteration)
+            self.set_update(self._update_function)
+
+        def set_pre_solve_callback(self, callback: Callable[["BlockedNewtonSolver"], None]):
+            """Set a callback function that is called before each Newton iteration."""
+            self._pre_solve_callback = callback
+
+        def set_post_solve_callback(self, callback: Callable[["BlockedNewtonSolver"], None]):
+            """Set a callback function that is called after each Newton iteration."""
+            self._post_solve_callback = callback
+
+        @property
+        def L(self) -> list[dolfinx.fem.Form]:
+            """Compiled linear form (the residual form)"""
+            return self._F
+
+        @property
+        def a(self) -> list[list[dolfinx.fem.Form]]:
+            """Compiled bilinear form (the Jacobian form)"""
+            return self._a
+
+        @property
+        def u(self):
+            return self._u
+
+        def __del__(self):
+            self._J.destroy()
+            self._b.destroy()
+            self._dx.destroy()
+            self._x.destroy()
+
+        def _pre_newton_iteration(self, x: PETSc.Vec) -> None:
+            """Function called before the residual or Jacobian is
+            computed.
+            Args:
+            x: The vector containing the latest solution
+            """
+            if self._pre_solve_callback is not None:
+                self._pre_solve_callback(self)
+            # Scatter previous solution `u=[u_0, ..., u_N]` to `x`; the
+            # blocked version used for lifting
             dolfinx.cpp.la.petsc.scatter_local_vectors(
-                self.x,
-                [si.x.petsc_vec.array_r for si in self.w],
+                x,
+                [ui.x.petsc_vec.array_r for ui in self._u],
                 [
                     (
-                        si.function_space.dofmap.index_map,
-                        si.function_space.dofmap.index_map_bs,
+                        ui.function_space.dofmap.index_map,
+                        ui.function_space.dofmap.index_map_bs,
                     )
-                    for si in self.w
+                    for ui in self._u
                 ],
             )
-            self.x.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
+            x.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
 
+        def _assemble_residual(self, x: PETSc.Vec, b: PETSc.Vec) -> None:
+            """Assemble the residual F into the vector b.
+            Args:
+                x: The vector containing the latest solution
+                b: Vector to assemble the residual into
+            """
             # Assemble F(u_{i-1}) - J(u_D - u_{i-1}) and set du|_bc= u_D - u_{i-1}
-            with self.b.localForm() as b_local:
+            with b.localForm() as b_local:
                 b_local.set(0.0)
-            dolfinx.fem.petsc.assemble_vector_block(
-                self.b,
-                self._F,
-                self._J,
-                bcs=self.bcs,
-                x0=self.x,
-                coeffs_a=coeffs_a,
-                constants_a=constants_a,
-                coeffs_L=coeffs_L,
-                constants_L=constants_L,
-                # dolfinx 0.8 compatibility
-                # this is called 'scale' in 0.8, 'alpha' in 0.9
-                **{_alpha_kw: -1.0},
-            )
-            self.b.ghostUpdate(PETSc.InsertMode.INSERT_VALUES, PETSc.ScatterMode.FORWARD)
+            try:
+                dolfinx.fem.petsc.assemble_vector_block(
+                    b,
+                    self._F,
+                    self._a,
+                    bcs=self._bcs,
+                    x0=x,
+                    # dolfinx 0.8 compatibility
+                    # this is called 'scale' in 0.8, 'alpha' in 0.9
+                    **{_alpha_kw: -1.0},
+                )
+            except AttributeError:
+                dolfinx.fem.petsc.assemble_vector(b, self._F)
+                bcs1 = dolfinx.fem.bcs_by_block(
+                    dolfinx.fem.extract_function_spaces(self._a, 1), self._bcs
+                )
+                dolfinx.fem.petsc.apply_lifting(b, self._a, bcs=bcs1, x0=x, alpha=-1.0)
+                b.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
+                bcs0 = dolfinx.fem.bcs_by_block(
+                    dolfinx.fem.extract_function_spaces(self._F), self._bcs
+                )
+                dolfinx.fem.petsc.set_bc(b, bcs0, x0=x, alpha=-1.0)
+            b.ghostUpdate(PETSc.InsertMode.INSERT_VALUES, PETSc.ScatterMode.FORWARD)
 
+        def _assemble_jacobian(self, x: PETSc.Vec, A: PETSc.Mat) -> None:
+            """Assemble the Jacobian matrix.
+            Args:
+                x: The vector containing the latest solution
+            """
             # Assemble Jacobian
-            self.A.zeroEntries()
-            dolfinx.fem.petsc.assemble_matrix_block(
-                self.A, self._J, bcs=self.bcs, constants=constants_a, coeffs=coeffs_a
-            )
-            self.A.assemble()
-
-            self._solver.solve(self.b, self.dx)
-            if self._error_on_convergence:
-                if (status := self._solver.getConvergedReason()) <= 0:
-                    raise RuntimeError(f"Linear solver did not converge, got reason: {status}")
+            A.zeroEntries()
+            try:
+                dolfinx.fem.petsc.assemble_matrix_block(A, self._a, bcs=self._bcs)
+            except AttributeError:
+                dolfinx.fem.petsc.assemble_matrix(A, self._a, self._bcs)
+            A.assemble()
 
+        def _update_function(self, solver, dx: PETSc.Vec, x: PETSc.Vec):
+            if self._post_solve_callback is not None:
+                self._post_solve_callback(self)
             # Update solution
             offset_start = 0
-            for s in self.w:
-                num_sub_dofs = (
-                    s.function_space.dofmap.index_map.size_local
-                    * s.function_space.dofmap.index_map_bs
-                )
-                s.x.petsc_vec.array_w[:num_sub_dofs] -= (
-                    beta * self.dx.array_r[offset_start : offset_start + num_sub_dofs]
+            for ui in self._u:
+                Vi = ui.function_space
+                num_sub_dofs = Vi.dofmap.index_map.size_local * Vi.dofmap.index_map_bs
+                ui.x.petsc_vec.array_w[:num_sub_dofs] -= (
+                    self.relaxation_parameter
+                    * dx.array_r[offset_start : offset_start + num_sub_dofs]
                 )
-                s.x.petsc_vec.ghostUpdate(
+                ui.x.petsc_vec.ghostUpdate(
                     addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD
                 )
                 offset_start += num_sub_dofs
 
-            if self._post_solve_callback is not None:
-                self._post_solve_callback(self)
-
-            # Compute norm of update
-            residual = self.dx.norm(PETSc.NormType.NORM_2)
-            if i == 1:
-                self.residual_0 = residual
-            relative_residual = residual / max(self.residual_0, atol)
-
-            logger.info(
-                f"Newton iteration {i}"
-                f": r (abs) = {residual} (tol={atol}), "
-                f"r (rel) = {relative_residual} (tol={rtol})"
+        def solve(self):
+            """Solve non-linear problem into function. Returns the number
+            of iterations and if the solver converged."""
+            n, converged = super().solve(self._x)
+            self._x.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
+            return n, converged
+
+else:
+
+    class NewtonSolver:  # type: ignore[no-redef]
+        def __init__(self, *args, **kwargs):
+            raise RuntimeError(
+                "NewtonSolver is not available in this version of DOLFINx. "
+                "Please install a version of DOLFINx with PETSc4py support."
             )
-            if relative_residual < rtol or residual < atol:
-                return i
-            i += 1
-
-        if self._error_on_convergence:
-            raise RuntimeError("Newton solver did not converge")
-        else:
-            return self.max_iterations
-
-    @property
-    def F(self):
-        """The list of residuals where each entry is a ``dolfinx.fem.Form``."""
-        return self._F
-
-    @property
-    def J(self):
-        """
-        The Jacobian blocks represented as lists of lists where each entry
-        is a ``dolfinx.fem.Form``.
-        """
-        return self._J
-
-    def __del__(self):
-        """Clean up the solver by destroying PETSc objects."""
-        self.A.destroy()
-        self.b.destroy()
-        self.dx.destroy()
-        self._solver.destroy()
-        self.x.destroy()
-
-
-class BlockedNewtonSolver(dolfinx.cpp.nls.petsc.NewtonSolver):
-    def __init__(
-        self,
-        F: list[ufl.form.Form],
-        u: list[dolfinx.fem.Function],
-        bcs: list[dolfinx.fem.DirichletBC] = [],
-        J: list[list[ufl.form.Form]] | None = None,
-        form_compiler_options: dict | None = None,
-        jit_options: dict | None = None,
-        petsc_options: dict | None = None,
-        entity_maps: dict | None = None,
-    ):
-        """Initialize solver for solving a non-linear problem using Newton's method.
-        Args:
-            F: List of PDE residuals [F_0(u, v_0), F_1(u, v_1), ...]
-            u: List of unknown functions u=[u_0, u_1, ...]
-            bcs: List of Dirichlet boundary conditions
-            J: UFL representation of the Jacobian (Optional)
-                Note:
-                    If not provided, the Jacobian will be computed using the
-                    assumption that the test functions come from a ``ufl.MixedFunctionSpace``
-            form_compiler_options: Options used in FFCx
-                compilation of this form. Run ``ffcx --help`` at the
-                command line to see all available options.
-            jit_options: Options used in CFFI JIT compilation of C
-                code generated by FFCx. See ``python/dolfinx/jit.py``
-                for all available options. Takes priority over all
-                other option values.
-            petsc_options:
-                Options passed to the PETSc Krylov solver.
-            entity_maps: Maps used to map entities between different meshes.
-                Only needed if the forms have not been compiled a priori,
-                and has coefficients, test, or trial functions that are defined on different meshes.
-        """
-        # Initialize base class
-        super().__init__(u[0].function_space.mesh.comm)
-
-        # Set PETSc options for Krylov solver
-        prefix = self.krylov_solver.getOptionsPrefix()
-        if prefix is None:
-            prefix = ""
-        if petsc_options is not None:
-            # Set PETSc options
-            opts = PETSc.Options()
-            opts.prefixPush(prefix)
-            for k, v in petsc_options.items():
-                opts[k] = v
-            opts.prefixPop()
-            self.krylov_solver.setFromOptions()
-        self._F = dolfinx.fem.form(
-            F,
-            form_compiler_options=form_compiler_options,
-            jit_options=jit_options,
-            entity_maps=entity_maps,
-        )
-
-        # Create the Jacobian matrix, dF/du
-        if J is None:
-            if _v(dolfinx.__version__) < _v("0.9"):
-                raise RuntimeError(
-                    "Automatic computation of Jacobian for blocked problem is only"
-                    + "supported in DOLFINx 0.9 and later"
-                )
-            du = ufl.TrialFunctions(ufl.MixedFunctionSpace(*[ui.function_space for ui in u]))
-            J = ufl.extract_blocks(sum(ufl.derivative(sum(F), u[i], du[i]) for i in range(len(u))))
-        self._a = dolfinx.fem.form(
-            J,
-            form_compiler_options=form_compiler_options,
-            jit_options=jit_options,
-            entity_maps=entity_maps,
-        )
-
-        self._bcs = bcs
-        self._u = u
-        self._pre_solve_callback: Callable[["BlockedNewtonSolver"], None] | None = None
-        self._post_solve_callback: Callable[["BlockedNewtonSolver"], None] | None = None
-
-        # Create structures for holding arrays and matrix
-        self._b = dolfinx.fem.petsc.create_vector_block(self._F)
-        self._J = dolfinx.fem.petsc.create_matrix_block(self._a)
-        self._dx = dolfinx.fem.petsc.create_vector_block(self._F)
-        self._x = dolfinx.fem.petsc.create_vector_block(self._F)
-        self._J.setOptionsPrefix(prefix)
-        self._J.setFromOptions()
-
-        self.setJ(self._assemble_jacobian, self._J)
-        self.setF(self._assemble_residual, self._b)
-        self.set_form(self._pre_newton_iteration)
-        self.set_update(self._update_function)
-
-    def set_pre_solve_callback(self, callback: Callable[["BlockedNewtonSolver"], None]):
-        """Set a callback function that is called before each Newton iteration."""
-        self._pre_solve_callback = callback
-
-    def set_post_solve_callback(self, callback: Callable[["BlockedNewtonSolver"], None]):
-        """Set a callback function that is called after each Newton iteration."""
-        self._post_solve_callback = callback
-
-    @property
-    def L(self) -> list[dolfinx.fem.Form]:
-        """Compiled linear form (the residual form)"""
-        return self._F
-
-    @property
-    def a(self) -> list[list[dolfinx.fem.Form]]:
-        """Compiled bilinear form (the Jacobian form)"""
-        return self._a
-
-    @property
-    def u(self):
-        return self._u
-
-    def __del__(self):
-        self._J.destroy()
-        self._b.destroy()
-        self._dx.destroy()
-        self._x.destroy()
-
-    def _pre_newton_iteration(self, x: PETSc.Vec) -> None:
-        """Function called before the residual or Jacobian is
-        computed.
-        Args:
-           x: The vector containing the latest solution
-        """
-        if self._pre_solve_callback is not None:
-            self._pre_solve_callback(self)
-        # Scatter previous solution `u=[u_0, ..., u_N]` to `x`; the
-        # blocked version used for lifting
-        dolfinx.cpp.la.petsc.scatter_local_vectors(
-            x,
-            [ui.x.petsc_vec.array_r for ui in self._u],
-            [
-                (
-                    ui.function_space.dofmap.index_map,
-                    ui.function_space.dofmap.index_map_bs,
-                )
-                for ui in self._u
-            ],
-        )
-        x.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
-
-    def _assemble_residual(self, x: PETSc.Vec, b: PETSc.Vec) -> None:
-        """Assemble the residual F into the vector b.
-        Args:
-            x: The vector containing the latest solution
-            b: Vector to assemble the residual into
-        """
-        # Assemble F(u_{i-1}) - J(u_D - u_{i-1}) and set du|_bc= u_D - u_{i-1}
-        with b.localForm() as b_local:
-            b_local.set(0.0)
-        dolfinx.fem.petsc.assemble_vector_block(
-            b,
-            self._F,
-            self._a,
-            bcs=self._bcs,
-            x0=x,
-            # dolfinx 0.8 compatibility
-            # this is called 'scale' in 0.8, 'alpha' in 0.9
-            **{_alpha_kw: -1.0},
-        )
-        b.ghostUpdate(PETSc.InsertMode.INSERT_VALUES, PETSc.ScatterMode.FORWARD)
-
-    def _assemble_jacobian(self, x: PETSc.Vec, A: PETSc.Mat) -> None:
-        """Assemble the Jacobian matrix.
-        Args:
-            x: The vector containing the latest solution
-        """
-        # Assemble Jacobian
-        A.zeroEntries()
-        dolfinx.fem.petsc.assemble_matrix_block(A, self._a, bcs=self._bcs)
-        A.assemble()
-
-    def _update_function(self, solver, dx: PETSc.Vec, x: PETSc.Vec):
-        if self._post_solve_callback is not None:
-            self._post_solve_callback(self)
-        # Update solution
-        offset_start = 0
-        for ui in self._u:
-            Vi = ui.function_space
-            num_sub_dofs = Vi.dofmap.index_map.size_local * Vi.dofmap.index_map_bs
-            ui.x.petsc_vec.array_w[:num_sub_dofs] -= (
-                self.relaxation_parameter * dx.array_r[offset_start : offset_start + num_sub_dofs]
-            )
-            ui.x.petsc_vec.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
-            offset_start += num_sub_dofs
 
-    def solve(self):
-        """Solve non-linear problem into function. Returns the number
-        of iterations and if the solver converged."""
-        n, converged = super().solve(self._x)
-        self._x.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
-        return n, converged
+    class BlockedNewtonSolver:  # type: ignore[no-redef]
+        def __init__(self, *args, **kwargs):
+            raise RuntimeError(
+                "BlockedNewtonSolver is not available in this version of DOLFINx. "
+                "Please install a version of DOLFINx with PETSc4py support."
+            )
diff -Nru scifem-0.5.0/src/scifem/xdmf.py scifem-0.6.0/src/scifem/xdmf.py
--- scifem-0.5.0/src/scifem/xdmf.py	2025-03-21 20:54:41.000000000 +0100
+++ scifem-0.6.0/src/scifem/xdmf.py	2025-05-23 21:28:34.000000000 +0200
@@ -517,8 +517,12 @@
 
     def _close_adios(self) -> None:
         logger.debug("Closing ADIOS2 file")
-        self._outfile.Close()
-        assert self._adios.RemoveIO("Point cloud writer")
+        try:
+            self._outfile.Close()
+            assert self._adios.RemoveIO("Point cloud writer")
+        except ValueError:
+            # File is allready closed
+            logger.debug("ADIOS2 file already closed")
 
     def close(self) -> None:
         """Close the XDMF file."""
@@ -596,6 +600,9 @@
     def __exit__(self, exc_type, exc_value, traceback):
         self.close()
 
+    def __del__(self):
+        self.close()
+
     def write(self, time: float) -> None:
         """Write the point cloud at a given time.
 
diff -Nru scifem-0.5.0/tests/test_assembly.py scifem-0.6.0/tests/test_assembly.py
--- scifem-0.5.0/tests/test_assembly.py	2025-03-21 20:54:41.000000000 +0100
+++ scifem-0.6.0/tests/test_assembly.py	2025-05-23 21:28:34.000000000 +0200
@@ -1,8 +1,19 @@
 from mpi4py import MPI
 
 from scifem import assemble_scalar, norm
-from dolfinx.mesh import create_unit_square
-from dolfinx.fem import Function, functionspace, Constant
+import scifem.petsc
+from dolfinx.mesh import create_unit_square, exterior_facet_indices
+import dolfinx
+import basix.ufl
+from dolfinx.fem import (
+    Function,
+    functionspace,
+    Constant,
+    form,
+    locate_dofs_topological,
+    dirichletbc,
+)
+from dolfinx import default_scalar_type
 import ufl
 import pytest
 import numpy as np
@@ -100,3 +111,95 @@
     reference = assemble_scalar(ref_form)
     tol = 50 * np.finfo(dtype).eps
     assert np.isclose(result, reference, atol=tol)
+
+
+@pytest.mark.skipif(not dolfinx.has_petsc4py, reason="Requires DOLFINX with PETSc4py")
+@pytest.mark.skipif(
+    hasattr(dolfinx.fem.petsc, "create_matrix_nest"),
+    reason="Requires latest version of DOLFINx PETSc API",
+)
+@pytest.mark.parametrize("kind", [None, "mpi", "nest"])
+@pytest.mark.parametrize(
+    "alpha", [dolfinx.default_scalar_type(3.0), dolfinx.default_scalar_type(-2.0)]
+)
+def test_lifting_helper(kind, alpha):
+    from petsc4py import PETSc
+    import dolfinx.fem.petsc
+
+    mesh = create_unit_square(MPI.COMM_WORLD, 12, 15)
+    el0 = basix.ufl.element("Lagrange", mesh.basix_cell(), 1)
+    el1 = basix.ufl.element("Lagrange", mesh.basix_cell(), 2)
+
+    if kind is None:
+        el = basix.ufl.mixed_element([el0, el1])
+        W = functionspace(mesh, el)
+        V, _ = W.sub(0).collapse()
+        Q, _ = W.sub(1).collapse()
+    else:
+        V = functionspace(mesh, el0)
+        Q = functionspace(mesh, el1)
+        W = ufl.MixedFunctionSpace(V, Q)
+
+    k = Function(V)
+    k.interpolate(lambda x: 3 + x[0] * x[1])
+    u, p = ufl.TrialFunctions(W)
+    v, q = ufl.TestFunctions(W)
+
+    if kind is None:
+        L = [
+            ufl.inner(dolfinx.fem.Constant(mesh, dolfinx.default_scalar_type(0.0)), v) * ufl.dx,
+            ufl.inner(dolfinx.fem.Constant(mesh, dolfinx.default_scalar_type(0.0)), q) * ufl.dx,
+        ]
+    else:
+        L = [ufl.ZeroBaseForm((v,)), ufl.ZeroBaseForm((q,))]
+    a = [
+        [k * ufl.inner(u, v) * ufl.dx, None],
+        [ufl.inner(u, q) * ufl.dx, k * ufl.inner(p, q) * ufl.dx],
+    ]
+    if kind is None:
+        # Flatten a and L
+        L = sum(L)
+        a = [sum([aij if aij is not None else 0 for ai in a for aij in ai])]
+
+    L_compiled = form(L)
+    a_compiled = form(a)
+
+    mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim)
+    bndry_facets = exterior_facet_indices(mesh.topology)
+
+    bc_space = W.sub(1) if kind is None else Q
+
+    bndry_dofs = locate_dofs_topological(bc_space, mesh.topology.dim - 1, bndry_facets)
+    bcs = [dirichletbc(default_scalar_type(2.0), bndry_dofs, bc_space)]
+    if kind is None:
+        lifting_bcs = [bcs]
+    else:
+        lifting_bcs = bcs
+
+    b = dolfinx.fem.petsc.create_vector(L_compiled, kind=kind)
+    scifem.petsc.zero_petsc_vector(b)
+    scifem.petsc.apply_lifting_and_set_bc(b, a_compiled, lifting_bcs, alpha=alpha)
+
+    # Create reference multiplication of alpha Ag
+    if kind is None:
+        a_compiled = a_compiled[0]
+    A = dolfinx.fem.petsc.create_matrix(a_compiled, kind=kind)
+    dolfinx.fem.petsc.assemble_matrix(A, a_compiled)
+    A.assemble()
+
+    g_ = dolfinx.fem.petsc.create_vector(L_compiled, kind=kind)
+    scifem.petsc.zero_petsc_vector(g_)
+
+    if kind is not None:
+        bcs0 = dolfinx.fem.bcs_by_block(dolfinx.fem.extract_function_spaces(a_compiled, 0), bcs)
+    else:
+        bcs0 = bcs
+    dolfinx.fem.petsc.set_bc(g_, bcs0)
+
+    b_ref = dolfinx.fem.petsc.create_vector(L_compiled, kind=kind)
+    A.mult(g_, b_ref)
+    b_ref.scale(-alpha)
+    dolfinx.fem.petsc.set_bc(b_ref, bcs0, alpha=alpha)
+    scifem.petsc.ghost_update(b_ref, PETSc.InsertMode.INSERT_VALUES, PETSc.ScatterMode.FORWARD)
+    with b.localForm() as b_local, b_ref.localForm() as b_ref_local:
+        np.testing.assert_allclose(b_local.array, b_ref_local.array)
diff -Nru scifem-0.5.0/tests/test_blocked_newton_solver.py scifem-0.6.0/tests/test_blocked_newton_solver.py
--- scifem-0.5.0/tests/test_blocked_newton_solver.py	2025-03-21 20:54:41.000000000 +0100
+++ scifem-0.6.0/tests/test_blocked_newton_solver.py	2025-05-23 21:28:34.000000000 +0200
@@ -2,31 +2,40 @@
 
 import numpy as np
 import dolfinx
+import dolfinx.nls.petsc
 from petsc4py import PETSc
-from scifem import NewtonSolver, assemble_scalar, BlockedNewtonSolver
+from scifem import assemble_scalar, BlockedNewtonSolver, NewtonSolver
+import basix.ufl
 import ufl
 import pytest
-from packaging.version import parse as _v
 
 
 @pytest.mark.parametrize("factor", [1, -1])
-def test_NewtonSolver(factor):
+@pytest.mark.parametrize("auto_split", [True, False])
+def test_NewtonSolver(factor, auto_split):
+    def left_boundary(x):
+        return np.isclose(x[0], 0.0)
+
     dtype = PETSc.RealType
     ftype = PETSc.ScalarType
     mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 10, 12, dtype=dtype)
-    V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1))
-    Q = dolfinx.fem.functionspace(mesh, ("Lagrange", 2))
+    mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim)
+    boundary_facets = dolfinx.mesh.locate_entities_boundary(mesh, 1, left_boundary)
+
+    el_0 = basix.ufl.element("Lagrange", mesh.basix_cell(), 1)
+    el_1 = basix.ufl.element("Lagrange", mesh.basix_cell(), 2)
+    V = dolfinx.fem.functionspace(mesh, el_0)
+    Q = dolfinx.fem.functionspace(mesh, el_1)
 
-    backward_compat = _v(dolfinx.__version__) < _v("0.9")
-    if backward_compat:
+    if auto_split:
+        W = ufl.MixedFunctionSpace(V, Q)
+        v, q = ufl.TestFunctions(W)
+        du, dp = ufl.TrialFunctions(W)
+    else:
         v = ufl.TestFunction(V)
         q = ufl.TestFunction(Q)
         du = ufl.TrialFunction(V)
         dp = ufl.TrialFunction(Q)
-    else:
-        W = ufl.MixedFunctionSpace(V, Q)
-        v, q = ufl.TestFunctions(W)
-        du, dp = ufl.TrialFunctions(W)
 
     # Set nonzero initial guess
     # Choose initial guess acording to the input factor
@@ -41,22 +50,36 @@
     c1 = dolfinx.fem.Constant(mesh, ftype(0.82))
     F0 = ufl.inner(c0 * u**2, v) * ufl.dx - ufl.inner(u_expr, v) * ufl.dx
     F1 = ufl.inner(c1 * p**2, q) * ufl.dx - ufl.inner(p_expr, q) * ufl.dx
-    if backward_compat:
+    if auto_split:
+        F_ = F0 + F1
+        F = list(ufl.extract_blocks(F_))
+        J = ufl.extract_blocks(ufl.derivative(F_, u, du) + ufl.derivative(F_, p, dp))
+    else:
         F = [F0, F1]
         J = [
             [ufl.derivative(F0, u, du), ufl.derivative(F0, p, dp)],
             [ufl.derivative(F1, u, du), ufl.derivative(F1, p, dp)],
         ]
-    else:
-        F_ = F0 + F1
-        F = list(ufl.extract_blocks(F_))
-        J = ufl.extract_blocks(ufl.derivative(F_, u, du) + ufl.derivative(F_, p, dp))
-    petsc_options = {"ksp_type": "preonly", "pc_type": "lu", "pc_factor_mat_solver_type": "mumps"}
-    solver = NewtonSolver(F, J, [u, p], max_iterations=25, petsc_options=petsc_options)
-    solver.solve()
 
+    # Reference solution
     u_ex = factor * ufl.sqrt(u_expr / c0)
     p_ex = factor * ufl.sqrt(p_expr / c1)
+
+    # Create BC on second space
+    dofs_q = dolfinx.fem.locate_dofs_topological(Q, mesh.topology.dim - 1, boundary_facets)
+    p_bc = dolfinx.fem.Function(Q, dtype=ftype)
+    try:
+        ip = Q.element.interpolation_points()
+
+    except TypeError:
+        ip = Q.element.interpolation_points
+    p_bc.interpolate(dolfinx.fem.Expression(p_ex, ip))
+    bc_q = dolfinx.fem.dirichletbc(p_bc, dofs_q)
+
+    petsc_options = {"ksp_type": "preonly", "pc_type": "lu", "pc_factor_mat_solver_type": "mumps"}
+    solver = NewtonSolver(F, J, [u, p], bcs=[bc_q], max_iterations=25, petsc_options=petsc_options)
+    solver.solve()
+
     err_u = ufl.inner(u_ex - u, u_ex - u) * ufl.dx
     err_p = ufl.inner(p_ex - p, p_ex - p) * ufl.dx
     tol = np.finfo(dtype).eps * 1.0e3
@@ -64,15 +87,15 @@
     assert assemble_scalar(err_p) < tol
 
     # Check consistency with other Newton solver
-    if backward_compat:
-        blocked_solver = dolfinx.nls.NewtonSolver(F, J, [u, p], petsc_options=petsc_options)
-    else:
-        blocked_solver = BlockedNewtonSolver(F, [u, p], J=None, petsc_options=petsc_options)
+    blocked_solver = BlockedNewtonSolver(
+        F, [u, p], bcs=[bc_q], J=None if auto_split else J, petsc_options=petsc_options
+    )
 
     dolfinx.log.set_log_level(dolfinx.log.LogLevel.ERROR)
     u.x.array[:] = factor * 0.1
     p.x.array[:] = factor * 0.02
     blocked_solver.convergence_criterion = "incremental"
+    dolfinx.log.set_log_level(dolfinx.log.LogLevel.INFO)
     blocked_solver.solve()
 
     err_u = ufl.inner(u_ex - u, u_ex - u) * ufl.dx
diff -Nru scifem-0.5.0/tests/test_read_mri_data.py scifem-0.6.0/tests/test_read_mri_data.py
--- scifem-0.5.0/tests/test_read_mri_data.py	2025-03-21 20:54:41.000000000 +0100
+++ scifem-0.6.0/tests/test_read_mri_data.py	2025-05-23 21:28:34.000000000 +0200
@@ -13,7 +13,12 @@
 @pytest.mark.parametrize("Nz", [9, 17])
 @pytest.mark.parametrize("theta", [np.pi / 3, 0, -np.pi])
 @pytest.mark.parametrize("translation", [np.array([0, 0, 0]), np.array([2.1, 1.3, 0.4])])
-def test_read_mri_data_to_function(degree, M, Nx, Ny, Nz, theta, translation, tmp_path):
+@pytest.mark.parametrize("mri_data_format", ["nifti", "mgh"])
+@pytest.mark.parametrize("external_affine", [True, False])
+@pytest.mark.parametrize("use_tkr", [False])
+def test_read_mri_data_to_function(
+    degree, M, Nx, Ny, Nz, theta, translation, tmp_path, mri_data_format, external_affine, use_tkr
+):
     nibabel = pytest.importorskip("nibabel")
 
     # Generate rotation and scaling matrix equivalent to a unit cube
@@ -25,14 +30,25 @@
     # Generate the affine mapping for nibabel
     A = np.append(np.dot(rotation_matrix_3D, scale_matrix), (translation).reshape(3, 1), axis=1)
     A = np.vstack([A, [0, 0, 0, 1]])
+    Id = np.identity(4)
 
     # Write transformed data to file
-    data = np.arange(1, M**3 + 1, dtype=np.float64).reshape(M, M, M)
-    image = nibabel.nifti1.Nifti1Image(data, affine=A)
-
     path = MPI.COMM_WORLD.bcast(tmp_path, root=0)
+    if mri_data_format == "nifti":
+        data = np.arange(1, M**3 + 1, dtype=np.float64).reshape(M, M, M)
+        image = nibabel.nifti1.Nifti1Image(data, affine=Id if external_affine else A)
+    elif mri_data_format == "mgh":
+        data = np.arange(1, M**3 + 1, dtype=np.int32).reshape(M, M, M)
+        image = nibabel.freesurfer.mghformat.MGHImage(data, affine=Id if external_affine else A)
+
+    # Reorient the image to RAS
+    orig_ornt = nibabel.io_orientation(image.affine)
+    targ_ornt = nibabel.orientations.axcodes2ornt("RAS")
+    transform = nibabel.orientations.ornt_transform(orig_ornt, targ_ornt)
+    image = image.as_reoriented(transform)
     filename = path.with_suffix(".mgz")
 
+    # Save the image to file
     if MPI.COMM_WORLD.rank == 0:
         nibabel.save(image, filename)
     MPI.COMM_WORLD.Barrier()
@@ -69,7 +85,14 @@
     mesh.geometry.x[:] = np.dot(rotation_matrix_3D, shifted_coords.T).T + translation
 
     # # Read data to function
-    func = read_mri_data_to_function(filename, mesh, degree=degree, dtype=np.float64)
+    if external_affine:
+        func = read_mri_data_to_function(
+            filename, mesh, degree=degree, dtype=np.float64, vox2ras_transform=A, use_tkr=use_tkr
+        )
+    else:
+        func = read_mri_data_to_function(
+            filename, mesh, degree=degree, dtype=np.float64, use_tkr=use_tkr
+        )
 
     np.testing.assert_allclose(func.x.array, reference_data)
 
@@ -78,5 +101,17 @@
         num_cells_local = cell_map.size_local + cell_map.num_ghosts
         # Pick every second cell
         cells = np.arange(num_cells_local, dtype=np.int32)[::2]
-        tag = read_mri_data_to_tag(filename, mesh, edim=mesh.topology.dim, entities=cells)
+        if external_affine:
+            tag = read_mri_data_to_tag(
+                filename,
+                mesh,
+                edim=mesh.topology.dim,
+                entities=cells,
+                vox2ras_transform=A,
+                use_tkr=use_tkr,
+            )
+        else:
+            tag = read_mri_data_to_tag(
+                filename, mesh, edim=mesh.topology.dim, entities=cells, use_tkr=use_tkr
+            )
         np.testing.assert_allclose(tag.values, reference_data[::2])
diff -Nru scifem-0.5.0/tests/test_real_functionspace.py scifem-0.6.0/tests/test_real_functionspace.py
--- scifem-0.5.0/tests/test_real_functionspace.py	2025-03-21 20:54:41.000000000 +0100
+++ scifem-0.6.0/tests/test_real_functionspace.py	2025-05-23 21:28:34.000000000 +0200
@@ -141,12 +141,28 @@
     a = dolfinx.fem.form([[a00, a01], [a10, None]], dtype=stype)
     L = dolfinx.fem.form([L0, L1], dtype=stype)
 
-    A = dolfinx.fem.petsc.assemble_matrix_block(a)
+    new_assemble_mode = False
+    try:
+        A = dolfinx.fem.petsc.assemble_matrix_block(a)
+    except AttributeError:
+        new_assemble_mode = True
+        A = dolfinx.fem.petsc.assemble_matrix(a, kind="mpi")
     A.assemble()
-    b = dolfinx.fem.petsc.create_vector_block(L)
+
+    if new_assemble_mode:
+        b = dolfinx.fem.petsc.create_vector(L, kind="mpi")
+    else:
+        b = dolfinx.fem.petsc.create_vector_block(L)
+
     with b.localForm() as loc:
         loc.set(0)
-    dolfinx.fem.petsc.assemble_vector_block(b, L, a, bcs=[])
+
+    if new_assemble_mode:
+        dolfinx.fem.petsc.assemble_vector(b, L)
+        b.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
+        b.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
+    else:
+        dolfinx.fem.petsc.assemble_vector_block(b, L, a, bcs=[])
 
     ksp = PETSc.KSP().create(mesh.comm)
     ksp.setOperators(A)
@@ -155,16 +171,33 @@
     pc.setType("lu")
     pc.setFactorSolverType("mumps")
 
-    x = dolfinx.fem.petsc.create_vector_block(L)
+    if new_assemble_mode:
+        x = dolfinx.fem.petsc.create_vector(L, kind="mpi")
+    else:
+        x = dolfinx.fem.petsc.create_vector_block(L)
+
     ksp.solve(b, x)
     x.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
+
     uh = dolfinx.fem.Function(V)
-    x_local = dolfinx.cpp.la.petsc.get_local_vectors(
-        x,
-        [(V.dofmap.index_map, V.dofmap.index_map_bs), (R.dofmap.index_map, R.dofmap.index_map_bs)],
-    )
-    uh.x.array[: len(x_local[0])] = x_local[0]
-    uh.x.scatter_forward()
+    rh = dolfinx.fem.Function(R)
+    if new_assemble_mode:
+        dolfinx.fem.petsc.assign(x, [uh, rh])
+    else:
+        x_local = dolfinx.cpp.la.petsc.get_local_vectors(
+            x,
+            [
+                (V.dofmap.index_map, V.dofmap.index_map_bs),
+                (R.dofmap.index_map, R.dofmap.index_map_bs),
+            ],
+        )
+        uh.x.array[: len(x_local[0])] = x_local[0]
+        uh.x.scatter_forward()
+
+    b.destroy()
+    x.destroy()
+    A.destroy()
+    ksp.destroy()
 
     error = dolfinx.fem.form(ufl.inner(u_ex - uh, u_ex - uh) * dx, dtype=stype)
 

Reply to: