Support for real forms with complex arithmetic

[jsagg931]

Describe new/missing feature

Dear community,

Within the current version of dolfinx it is not possible to assemble a complex PETSc matrix from a real dolfinx.fem.form. This is actually, a drawback for several applications, for instance for linearization of DG forms which require of Riemmann/numerical fluxes. In this applications one usually requires of the linear form and the jacobian obtained with AD as the bilinear form. The DG numerical fluxes are often defined with ufl.conditional, and this is fine when solving for a steady-state via a Newton iteration of time-stepping in real arithmetic. However, there are applications, in complex arithmetics, where the Jacobian matrix is evaluated at a real solution but applied on complex vectors. In such a case, the ufl.conditional is well defined if evaluated on a real dolfinx.fem.Function then converted to a Float64_Form and assemble a complex Petsc matrix.
The final user should use create_mat_real_to_complex (see below).

I propose the following interface, which modifies cpp/dolfinx/fem/assembler.h, cpp/dolfinx/fem/assembler_matrix_impl.h, python/dolfinx/wrapper/petsc.cpp and python/dolfinx/fem/petsc.py. Simpler modifications could have been made by playing with the template of assemble_matrix, possibly adding a new option to the template that allows a different type than PetscScalar. However, I opted for this option because it is less invasive, it does not modify existing code, and I wait for your comments.

Please, also let me know if this is the right way to do the issue (at first I wanted to attach the files but it did not allow me).

Suggested user interface

In cpp/dolfinx/fem/assembler.h to include

/// @brief Assemble bilinear real form into a complex matrix. Matrix must already be
/// initialised. Does not zero or finalise the matrix.
/// @param[in] mat_add The function for adding values into the matrix
/// @param[in] a The bilinear form to assemble
/// @param[in] constants Constants that appear in `a`
/// @param[in] coefficients Coefficients that appear in `a`
/// @param[in] dof_marker0 Boundary condition markers for the rows. If
/// bc[i] is true then rows i in A will be zeroed. The index i is a
/// local index.
/// @param[in] dof_marker1 Boundary condition markers for the columns.
/// If bc[i] is true then rows i in A will be zeroed. The index i is a
/// local index.
template <dolfinx::scalar T, std::floating_point U>
void assemble_matrix_cast(
    la::MatSet<PetscScalar> auto mat_add, const Form<T, U>& a,
    std::span<const T> constants,
    const std::map<std::pair<IntegralType, int>,
                   std::pair<std::span<const T>, int>>& coefficients,
    std::span<const std::int8_t> dof_marker0,
    std::span<const std::int8_t> dof_marker1)

{
  using mdspanx3_t
      = md::mdspan<const scalar_value_t<T>,
                   md::extents<std::size_t, md::dynamic_extent, 3>>;

  std::shared_ptr<const mesh::Mesh<U>> mesh = a.mesh();
  assert(mesh);
  std::span x = mesh->geometry().x();
  if constexpr (std::is_same_v<U, scalar_value_t<T>>)
  {
    impl::assemble_matrix_cast(mat_add, a, mdspanx3_t(x.data(), x.size() / 3, 3),
                               constants, coefficients, dof_marker0, dof_marker1);
  }
  else
  {
    std::vector<scalar_value_t<T>> _x(x.begin(), x.end());
    impl::assemble_matrix_cast(mat_add, a, mdspanx3_t(_x.data(), _x.size() / 3, 3),
                               constants, coefficients, dof_marker0, dof_marker1);
  }

}

/// Assemble bilinear form into a matrix
/// @param[in] mat_add The function for adding values into the matrix
/// @param[in] a The bilinear from to assemble
/// @param[in] constants Constants that appear in `a`
/// @param[in] coefficients Coefficients that appear in `a`
/// @param[in] bcs Boundary conditions to apply. For boundary condition
///  dofs the row and column are zeroed. The diagonal  entry is not set.

template <dolfinx::scalar T, std::floating_point U>
void assemble_matrix_cast(
  auto mat_add, const Form<T, U>& a, std::span<const T> constants,
  const std::map<std::pair<IntegralType, int>,
                 std::pair<std::span<const T>, int>>& coefficients,
  const std::vector<std::reference_wrapper<const DirichletBC<T, U>>>& bcs)
{

  // Index maps for dof ranges
  // NOTE: For mixed-topology meshes, there will be multiple DOF maps,
  // but the index maps are the same.
  auto map0 = a.function_spaces().at(0)->dofmaps(0)->index_map;
  auto map1 = a.function_spaces().at(1)->dofmaps(0)->index_map;
  auto bs0 = a.function_spaces().at(0)->dofmaps(0)->index_map_bs();
  auto bs1 = a.function_spaces().at(1)->dofmaps(0)->index_map_bs();

  // Build dof markers
  std::vector<std::int8_t> dof_marker0, dof_marker1;
  assert(map0);
  std::int32_t dim0 = bs0 * (map0->size_local() + map0->num_ghosts());
  assert(map1);
  std::int32_t dim1 = bs1 * (map1->size_local() + map1->num_ghosts());
  for (std::size_t k = 0; k < bcs.size(); ++k)
  {
    assert(bcs[k].get().function_space());
    if (a.function_spaces().at(0)->contains(*bcs[k].get().function_space()))
    {
      dof_marker0.resize(dim0, false);
      bcs[k].get().mark_dofs(dof_marker0);
    }

    if (a.function_spaces().at(1)->contains(*bcs[k].get().function_space()))
    {
      dof_marker1.resize(dim1, false);
      bcs[k].get().mark_dofs(dof_marker1);
    }
  }

  // Assemble
  assemble_matrix_cast(mat_add, a, constants, coefficients, dof_marker0,
                  dof_marker1);
  }
} // namespace dolfinx::fem

In cpp/dolfinx/fem/assembler_matrix_impl.h

// It shall be invoked with T = double
template <dolfinx::scalar T>
void assemble_cells_matrix_cast(
    la::MatSet<PetscScalar> auto mat_set, mdspan2_t x_dofmap,
    md::mdspan<const scalar_value_t<T>,
               md::extents<std::size_t, md::dynamic_extent, 3>>
        x,
    std::span<const std::int32_t> cells,
    std::tuple<mdspan2_t, int, std::span<const std::int32_t>> dofmap0,
    fem::DofTransformKernel<T> auto P0,
    std::tuple<mdspan2_t, int, std::span<const std::int32_t>> dofmap1,
    fem::DofTransformKernel<T> auto P1T, std::span<const std::int8_t> bc0,
    std::span<const std::int8_t> bc1, FEkernel<T> auto kernel,
    md::mdspan<const T, md::dextents<std::size_t, 2>> coeffs,
    std::span<const T> constants, std::span<const std::uint32_t> cell_info0,
    std::span<const std::uint32_t> cell_info1)
{
  if (cells.empty())
    return;

  const auto [dmap0, bs0, cells0] = dofmap0;
  const auto [dmap1, bs1, cells1] = dofmap1;

  // Iterate over active cells
  const int num_dofs0 = dmap0.extent(1);
  const int num_dofs1 = dmap1.extent(1);
  const int ndim0 = bs0 * num_dofs0;
  const int ndim1 = bs1 * num_dofs1;
  std::vector<T> Ae(ndim0 * ndim1);
  std::vector<scalar_value_t<T>> cdofs(3 * x_dofmap.extent(1));

  // Iterate over active cells
  assert(cells0.size() == cells.size());
  assert(cells1.size() == cells.size());
  for (std::size_t c = 0; c < cells.size(); ++c)
  {
    // Cell index in integration domain mesh (c), test function mesh
    // (c0) and trial function mesh (c1)
    std::int32_t cell = cells[c];
    std::int32_t cell0 = cells0[c];
    std::int32_t cell1 = cells1[c];

    // Get cell coordinates/geometry
    auto x_dofs = md::submdspan(x_dofmap, cell, md::full_extent);
    for (std::size_t i = 0; i < x_dofs.size(); ++i)
      std::copy_n(&x(x_dofs[i], 0), 3, std::next(cdofs.begin(), 3 * i));

    // Tabulate tensor
    std::ranges::fill(Ae, 0);
    kernel(Ae.data(), &coeffs(c, 0), constants.data(), cdofs.data(), nullptr,
           nullptr, nullptr);

    // Compute A = P_0 \tilde{A} P_1^T (dof transformation)
    P0(Ae, cell_info0, cell0, ndim1);  // B = P0 \tilde{A}
    P1T(Ae, cell_info1, cell1, ndim0); // A =  B P1_T

    // Zero rows/columns for essential bcs
    std::span dofs0(dmap0.data_handle() + cell0 * num_dofs0, num_dofs0);
    std::span dofs1(dmap1.data_handle() + cell1 * num_dofs1, num_dofs1);

    if (!bc0.empty())
    {
      for (int i = 0; i < num_dofs0; ++i)
      {
        for (int k = 0; k < bs0; ++k)
        {
          if (bc0[bs0 * dofs0[i] + k])
          {
            // Zero row bs0 * i + k
            const int row = bs0 * i + k;
            std::fill_n(std::next(Ae.begin(), ndim1 * row), ndim1, 0);
          }
        }
      }
    }

    if (!bc1.empty())
    {
      for (int j = 0; j < num_dofs1; ++j)
      {
        for (int k = 0; k < bs1; ++k)
        {
          if (bc1[bs1 * dofs1[j] + k])
          {
            // Zero column bs1 * j + k
            const int col = bs1 * j + k;
            for (int row = 0; row < ndim0; ++row)
              Ae[row * ndim1 + col] = 0;
          }
        }
      }
    }

    if constexpr (std::is_same_v<PetscScalar, scalar_value_t<T>>)
    {
      mat_set(dofs0, dofs1, Ae);
    }
    else if constexpr (dolfinx::is_complex<PetscScalar>)
    {
      std::vector<PetscScalar> Ae_Petsc(Ae.size());
      std::transform(Ae.begin(), Ae.end(), Ae_Petsc.begin(),
                      [](T val) { return PetscScalar(val, 0.0); });
      mat_set(dofs0, dofs1, Ae_Petsc);
    }
    else
    {
      std::vector<PetscScalar> Ae_Petsc(Ae.size());
      std::transform(Ae.begin(), Ae.end(), Ae_Petsc.begin(),
                      [](T val) { return static_cast<PetscScalar>(val); });
      mat_set(dofs0, dofs1, Ae_Petsc);
    }
  }
}


template <dolfinx::scalar T>
void assemble_exterior_facets_cast(
    la::MatSet<PetscScalar> auto mat_set, mdspan2_t x_dofmap,
    md::mdspan<const scalar_value_t<T>,
               md::extents<std::size_t, md::dynamic_extent, 3>>
        x,
    md::mdspan<const std::int32_t,
               std::extents<std::size_t, md::dynamic_extent, 2>>
        facets,
    std::tuple<mdspan2_t, int,
               md::mdspan<const std::int32_t,
                          std::extents<std::size_t, md::dynamic_extent, 2>>>
        dofmap0,
    fem::DofTransformKernel<T> auto P0,
    std::tuple<mdspan2_t, int,
               md::mdspan<const std::int32_t,
                          std::extents<std::size_t, md::dynamic_extent, 2>>>
        dofmap1,
    fem::DofTransformKernel<T> auto P1T, std::span<const std::int8_t> bc0,
    std::span<const std::int8_t> bc1, FEkernel<T> auto kernel,
    md::mdspan<const T, md::dextents<std::size_t, 2>> coeffs,
    std::span<const T> constants, std::span<const std::uint32_t> cell_info0,
    std::span<const std::uint32_t> cell_info1,
    md::mdspan<const std::uint8_t, md::dextents<std::size_t, 2>> perms)
{
  if (facets.empty())
    return;

  const auto [dmap0, bs0, facets0] = dofmap0;
  const auto [dmap1, bs1, facets1] = dofmap1;

  // Data structures used in assembly
  std::vector<scalar_value_t<T>> cdofs(3 * x_dofmap.extent(1));
  const int num_dofs0 = dmap0.extent(1);
  const int num_dofs1 = dmap1.extent(1);
  const int ndim0 = bs0 * num_dofs0;
  const int ndim1 = bs1 * num_dofs1;
  std::vector<T> Ae(ndim0 * ndim1);
  assert(facets0.size() == facets.size());
  assert(facets1.size() == facets.size());
  for (std::size_t f = 0; f < facets.extent(0); ++f)
  {
    // Cell in the integration domain, local facet index relative to the
    // integration domain cell, and cells in the test and trial function
    // meshes
    std::int32_t cell = facets(f, 0);
    std::int32_t local_facet = facets(f, 1);
    std::int32_t cell0 = facets0(f, 0);
    std::int32_t cell1 = facets1(f, 0);

    // Get cell coordinates/geometry
    auto x_dofs = md::submdspan(x_dofmap, cell, md::full_extent);
    for (std::size_t i = 0; i < x_dofs.size(); ++i)
      std::copy_n(&x(x_dofs[i], 0), 3, std::next(cdofs.begin(), 3 * i));

    // Permutations
    std::uint8_t perm = perms.empty() ? 0 : perms(cell, local_facet);

    // Tabulate tensor
    std::ranges::fill(Ae, 0);
    kernel(Ae.data(), &coeffs(f, 0), constants.data(), cdofs.data(),
           &local_facet, &perm, nullptr);
    P0(Ae, cell_info0, cell0, ndim1);
    P1T(Ae, cell_info1, cell1, ndim0);

    // Zero rows/columns for essential bcs
    std::span dofs0(dmap0.data_handle() + cell0 * num_dofs0, num_dofs0);
    std::span dofs1(dmap1.data_handle() + cell1 * num_dofs1, num_dofs1);
    if (!bc0.empty())
    {
      for (int i = 0; i < num_dofs0; ++i)
      {
        for (int k = 0; k < bs0; ++k)
        {
          if (bc0[bs0 * dofs0[i] + k])
          {
            // Zero row bs0 * i + k
            const int row = bs0 * i + k;
            std::fill_n(std::next(Ae.begin(), ndim1 * row), ndim1, 0);
          }
        }
      }
    }
    if (!bc1.empty())
    {
      for (int j = 0; j < num_dofs1; ++j)
      {
        for (int k = 0; k < bs1; ++k)
        {
          if (bc1[bs1 * dofs1[j] + k])
          {
            // Zero column bs1 * j + k
            const int col = bs1 * j + k;
            for (int row = 0; row < ndim0; ++row)
              Ae[row * ndim1 + col] = 0;
          }
        }
      }
    }
 

    if constexpr (std::is_same_v<PetscScalar, scalar_value_t<T>>)
    {
      mat_set(dofs0, dofs1, Ae);
    }
    else if constexpr (dolfinx::is_complex<PetscScalar>)
    {
      std::vector<PetscScalar> Ae_Petsc(Ae.size());
      std::transform(Ae.begin(), Ae.end(), Ae_Petsc.begin(),
                      [](T val) { return PetscScalar(val, 0.0); });
      mat_set(dofs0, dofs1, Ae_Petsc);
    }
    else
    {
      std::vector<PetscScalar> Ae_Petsc(Ae.size());
      std::transform(Ae.begin(), Ae.end(), Ae_Petsc.begin(),
                      [](T val) { return static_cast<PetscScalar>(val); });
      mat_set(dofs0, dofs1, Ae_Petsc);
    }
  }
}


template <dolfinx::scalar T>
void assemble_interior_facets_cast(
    la::MatSet<PetscScalar> auto mat_set,  mdspan2_t x_dofmap,
    md::mdspan<const scalar_value_t<T>,
               md::extents<std::size_t, md::dynamic_extent, 3>>
        x,
    md::mdspan<const std::int32_t,
               std::extents<std::size_t, md::dynamic_extent, 2, 2>>
        facets,
    std::tuple<const DofMap&, int,
               md::mdspan<const std::int32_t,
                          std::extents<std::size_t, md::dynamic_extent, 2, 2>>>
        dofmap0,
    fem::DofTransformKernel<T> auto P0,
    std::tuple<const DofMap&, int,
               md::mdspan<const std::int32_t,
                          std::extents<std::size_t, md::dynamic_extent, 2, 2>>>
        dofmap1,
    fem::DofTransformKernel<T> auto P1T, std::span<const std::int8_t> bc0,
    std::span<const std::int8_t> bc1, FEkernel<T> auto kernel,
    md::mdspan<const T, md::extents<std::size_t, md::dynamic_extent, 2,
                                    md::dynamic_extent>>
        coeffs,
    std::span<const T> constants, std::span<const std::uint32_t> cell_info0,
    std::span<const std::uint32_t> cell_info1,
    md::mdspan<const std::uint8_t, md::dextents<std::size_t, 2>> perms)
{
  if (facets.empty())
    return;

  const auto [dmap0, bs0, facets0] = dofmap0;
  const auto [dmap1, bs1, facets1] = dofmap1;

  // Data structures used in assembly
  using X = scalar_value_t<T>;
  std::vector<X> cdofs(2 * x_dofmap.extent(1) * 3);
  std::span<X> cdofs0(cdofs.data(), x_dofmap.extent(1) * 3);
  std::span<X> cdofs1(cdofs.data() + x_dofmap.extent(1) * 3,
                      x_dofmap.extent(1) * 3);

  const std::size_t dmap0_size = dmap0.map().extent(1);
  const std::size_t dmap1_size = dmap1.map().extent(1);
  const int num_rows = bs0 * 2 * dmap0_size;
  const int num_cols = bs1 * 2 * dmap1_size;

  // Temporaries for joint dofmaps
  std::vector<T> Ae(num_rows * num_cols), be(num_rows);
  std::vector<std::int32_t> dmapjoint0(2 * dmap0_size);
  std::vector<std::int32_t> dmapjoint1(2 * dmap1_size);
  assert(facets0.size() == facets.size());
  assert(facets1.size() == facets.size());
  for (std::size_t f = 0; f < facets.extent(0); ++f)
  {
    // Cells in integration domain,  test function domain and trial
    // function domain
    std::array cells{facets(f, 0, 0), facets(f, 1, 0)};
    std::array cells0{facets0(f, 0, 0), facets0(f, 1, 0)};
    std::array cells1{facets1(f, 0, 0), facets1(f, 1, 0)};

    // Local facets indices
    std::array local_facet{facets(f, 0, 1), facets(f, 1, 1)};

    // Get cell geometry
    auto x_dofs0 = md::submdspan(x_dofmap, cells[0], md::full_extent);
    for (std::size_t i = 0; i < x_dofs0.size(); ++i)
      std::copy_n(&x(x_dofs0[i], 0), 3, std::next(cdofs0.begin(), 3 * i));
    auto x_dofs1 = md::submdspan(x_dofmap, cells[1], md::full_extent);
    for (std::size_t i = 0; i < x_dofs1.size(); ++i)
      std::copy_n(&x(x_dofs1[i], 0), 3, std::next(cdofs1.begin(), 3 * i));

    // Get dof maps for cells and pack
    // When integrating over interfaces between two domains, the test function
    // might only be defined on one side, so we check which cells exist in the
    // test function domain
    std::span<const std::int32_t> dmap0_cell0
        = cells0[0] >= 0 ? dmap0.cell_dofs(cells0[0])
                         : std::span<const std::int32_t>();
    std::span<const std::int32_t> dmap0_cell1
        = cells0[1] >= 0 ? dmap0.cell_dofs(cells0[1])
                         : std::span<const std::int32_t>();

    std::ranges::copy(dmap0_cell0, dmapjoint0.begin());
    std::ranges::copy(dmap0_cell1, std::next(dmapjoint0.begin(), dmap0_size));

    // Check which cells exist in the trial function domain
    std::span<const std::int32_t> dmap1_cell0
        = cells1[0] >= 0 ? dmap1.cell_dofs(cells1[0])
                         : std::span<const std::int32_t>();
    std::span<const std::int32_t> dmap1_cell1
        = cells1[1] >= 0 ? dmap1.cell_dofs(cells1[1])
                         : std::span<const std::int32_t>();

    std::ranges::copy(dmap1_cell0, dmapjoint1.begin());
    std::ranges::copy(dmap1_cell1, std::next(dmapjoint1.begin(), dmap1_size));

    // Tabulate tensor
    std::ranges::fill(Ae, 0);
    std::array perm = perms.empty()
                          ? std::array<std::uint8_t, 2>{0, 0}
                          : std::array{perms(cells[0], local_facet[0]),
                                       perms(cells[1], local_facet[1])};
    kernel(Ae.data(), &coeffs(f, 0, 0), constants.data(), cdofs.data(),
           local_facet.data(), perm.data(), nullptr);

    // Local element layout is a 2x2 block matrix with structure
    //
    //   cell0cell0  |  cell0cell1
    //   cell1cell0  |  cell1cell1
    //
    // where each block is element tensor of size (dmap0, dmap1).

    // Only apply transformation when cells exist
    if (cells0[0] >= 0)
      P0(Ae, cell_info0, cells0[0], num_cols);
    if (cells0[1] >= 0)
    {
      std::span sub_Ae0(Ae.data() + bs0 * dmap0_size * num_cols,
                        bs0 * dmap0_size * num_cols);

      P0(sub_Ae0, cell_info0, cells0[1], num_cols);
    }
    if (cells1[0] >= 0)
      P1T(Ae, cell_info1, cells1[0], num_rows);

    if (cells1[1] >= 0)
    {
      for (int row = 0; row < num_rows; ++row)
      {
        // DOFs for dmap1 and cell1 are not stored contiguously in the
        // block matrix, so each row needs a separate span access
        std::span sub_Ae1(Ae.data() + row * num_cols + bs1 * dmap1_size,
                          bs1 * dmap1_size);
        P1T(sub_Ae1, cell_info1, cells1[1], 1);
      }
    }

    // Zero rows/columns for essential bcs
    if (!bc0.empty())
    {
      for (std::size_t i = 0; i < dmapjoint0.size(); ++i)
      {
        for (int k = 0; k < bs0; ++k)
        {
          if (bc0[bs0 * dmapjoint0[i] + k])
          {
            // Zero row bs0 * i + k
            std::fill_n(std::next(Ae.begin(), num_cols * (bs0 * i + k)),
                        num_cols, 0);
          }
        }
      }
    }
    if (!bc1.empty())
    {
      for (std::size_t j = 0; j < dmapjoint1.size(); ++j)
      {
        for (int k = 0; k < bs1; ++k)
        {
          if (bc1[bs1 * dmapjoint1[j] + k])
          {
            // Zero column bs1 * j + k
            for (int m = 0; m < num_rows; ++m)
              Ae[m * num_cols + bs1 * j + k] = 0;
          }
        }
      }
    }


    if constexpr (std::is_same_v<PetscScalar, scalar_value_t<T>>)
    {
      mat_set(dmapjoint0, dmapjoint1, Ae);
    }
    else if constexpr (dolfinx::is_complex<PetscScalar>)
    {
      std::vector<PetscScalar> Ae_Petsc(Ae.size());
      std::transform(Ae.begin(), Ae.end(), Ae_Petsc.begin(),
                      [](T val) { return PetscScalar(val, 0.0); });
        mat_set(dmapjoint0, dmapjoint1, Ae_Petsc);
    }
    else
    {
      std::vector<PetscScalar> Ae_Petsc(Ae.size());
      std::transform(Ae.begin(), Ae.end(), Ae_Petsc.begin(),
                      [](T val) { return static_cast<PetscScalar>(val); });
        mat_set(dmapjoint0, dmapjoint1, Ae_Petsc);
    }


  }
}


/// The matrix A must already be initialised. The matrix may be a proxy,
/// i.e. a view into a larger matrix, and assembly is performed using
/// local indices. Rows (bc0) and columns (bc1) with Dirichlet
/// conditions are zeroed. Markers (bc0 and bc1) can be empty if no bcs
/// are applied. Matrix is not finalised.
template <dolfinx::scalar T, std::floating_point U>
void assemble_matrix_cast(
    la::MatSet<PetscScalar> auto mat_set, const Form<T, U>& a,
    md::mdspan<const scalar_value_t<T>,
               md::extents<std::size_t, md::dynamic_extent, 3>>
        x,
    std::span<const T> constants,
    const std::map<std::pair<IntegralType, int>,
                   std::pair<std::span<const T>, int>>& coefficients,
    std::span<const std::int8_t> bc0, std::span<const std::int8_t> bc1)
{
  // Integration domain mesh
  std::shared_ptr<const mesh::Mesh<U>> mesh = a.mesh();
  assert(mesh);

  // Test function mesh
  auto mesh0 = a.function_spaces().at(0)->mesh();
  assert(mesh0);

  // Trial function mesh
  auto mesh1 = a.function_spaces().at(1)->mesh();
  assert(mesh1);
  

  // TODO: Mixed topology with exterior and interior facet integrals.
  //
  // NOTE: Can't just loop over cell types for interior facet integrals
  // because we have a kernel per combination of comparable cell types,
  // rather than one per cell type. Also, we need the dofmaps for two
  // different cell types at the same time.
  const int num_cell_types = mesh->topology()->cell_types().size();
  for (int cell_type_idx = 0; cell_type_idx < num_cell_types; ++cell_type_idx)
  {
    // Geometry dofmap and data
    mdspan2_t x_dofmap = mesh->geometry().dofmap(cell_type_idx);

    // Get dofmap data
    std::shared_ptr<const fem::DofMap> dofmap0
        = a.function_spaces().at(0)->dofmaps(cell_type_idx);
    std::shared_ptr<const fem::DofMap> dofmap1
        = a.function_spaces().at(1)->dofmaps(cell_type_idx);
    assert(dofmap0);
    assert(dofmap1);
    auto dofs0 = dofmap0->map();
    const int bs0 = dofmap0->bs();
    auto dofs1 = dofmap1->map();
    const int bs1 = dofmap1->bs();

    auto element0 = a.function_spaces().at(0)->elements(cell_type_idx);
    assert(element0);
    auto element1 = a.function_spaces().at(1)->elements(cell_type_idx);
    assert(element1);
    fem::DofTransformKernel<T> auto P0
        = element0->template dof_transformation_fn<T>(doftransform::standard);
    fem::DofTransformKernel<T> auto P1T
        = element1->template dof_transformation_right_fn<T>(
            doftransform::transpose);

    std::span<const std::uint32_t> cell_info0;
    std::span<const std::uint32_t> cell_info1;
    if (element0->needs_dof_transformations()
        or element1->needs_dof_transformations()
        or a.needs_facet_permutations())
    {
      mesh0->topology_mutable()->create_entity_permutations();
      mesh1->topology_mutable()->create_entity_permutations();
      cell_info0 = std::span(mesh0->topology()->get_cell_permutation_info());
      cell_info1 = std::span(mesh1->topology()->get_cell_permutation_info());
    }

    for (int i = 0; i < a.num_integrals(IntegralType::cell, cell_type_idx); ++i)
    {
      auto fn = a.kernel(IntegralType::cell, i, cell_type_idx);
      assert(fn);
      std::span cells = a.domain(IntegralType::cell, i, cell_type_idx);
      std::span cells0 = a.domain_arg(IntegralType::cell, 0, i, cell_type_idx);
      std::span cells1 = a.domain_arg(IntegralType::cell, 1, i, cell_type_idx);
      auto& [coeffs, cstride] = coefficients.at({IntegralType::cell, i});
      assert(cells.size() * cstride == coeffs.size());
      impl::assemble_cells_matrix_cast(
          mat_set, x_dofmap, x, cells, {dofs0, bs0, cells0}, P0,
          {dofs1, bs1, cells1}, P1T, bc0, bc1, fn,
          md::mdspan(coeffs.data(), cells.size(), cstride), constants,
          cell_info0, cell_info1);
    }

    md::mdspan<const std::uint8_t, md::dextents<std::size_t, 2>> perms;
    if (a.needs_facet_permutations())
    {
      mesh::CellType cell_type = mesh->topology()->cell_types()[cell_type_idx];
      int num_facets_per_cell
          = mesh::cell_num_entities(cell_type, mesh->topology()->dim() - 1);
      mesh->topology_mutable()->create_entity_permutations();
      const std::vector<std::uint8_t>& p
          = mesh->topology()->get_facet_permutations();
      perms = md::mdspan(p.data(), p.size() / num_facets_per_cell,
                         num_facets_per_cell);
    }

    for (int i = 0;
         i < a.num_integrals(IntegralType::exterior_facet, cell_type_idx); ++i)
    {
      if (num_cell_types > 1)
      {
        throw std::runtime_error("Exterior facet integrals with mixed "
                                 "topology aren't supported yet");
      }

      using mdspanx2_t
          = md::mdspan<const std::int32_t,
                       md::extents<std::size_t, md::dynamic_extent, 2>>;

      auto fn = a.kernel(IntegralType::exterior_facet, i, 0);
      assert(fn);
      auto& [coeffs, cstride]
          = coefficients.at({IntegralType::exterior_facet, i});

      std::span f = a.domain(IntegralType::exterior_facet, i, 0);
      mdspanx2_t facets(f.data(), f.size() / 2, 2);
      std::span f0 = a.domain_arg(IntegralType::exterior_facet, 0, i, 0);
      mdspanx2_t facets0(f0.data(), f0.size() / 2, 2);
      std::span f1 = a.domain_arg(IntegralType::exterior_facet, 1, i, 0);
      mdspanx2_t facets1(f1.data(), f1.size() / 2, 2);
      assert((facets.size() / 2) * cstride == coeffs.size());
      impl::assemble_exterior_facets_cast(
          mat_set, x_dofmap, x, facets, {dofs0, bs0, facets0}, P0,
          {dofs1, bs1, facets1}, P1T, bc0, bc1, fn,
          md::mdspan(coeffs.data(), facets.extent(0), cstride), constants,
          cell_info0, cell_info1, perms);
    }

    for (int i = 0;
         i < a.num_integrals(IntegralType::interior_facet, cell_type_idx); ++i)
    {
      if (num_cell_types > 1)
      {
        throw std::runtime_error("Interior facet integrals with mixed "
                                 "topology aren't supported yet");
      }

      using mdspanx22_t
          = md::mdspan<const std::int32_t,
                       md::extents<std::size_t, md::dynamic_extent, 2, 2>>;
      using mdspanx2x_t
          = md::mdspan<const T, md::extents<std::size_t, md::dynamic_extent, 2,
                                            md::dynamic_extent>>;

      auto fn = a.kernel(IntegralType::interior_facet, i, 0);
      assert(fn);
      auto& [coeffs, cstride]
          = coefficients.at({IntegralType::interior_facet, i});

      std::span facets = a.domain(IntegralType::interior_facet, i, 0);
      std::span facets0 = a.domain_arg(IntegralType::interior_facet, 0, i, 0);
      std::span facets1 = a.domain_arg(IntegralType::interior_facet, 1, i, 0);
      assert((facets.size() / 4) * 2 * cstride == coeffs.size());
      impl::assemble_interior_facets_cast(
          mat_set, x_dofmap, x,
          mdspanx22_t(facets.data(), facets.size() / 4, 2, 2),
          {*dofmap0, bs0,
           mdspanx22_t(facets0.data(), facets0.size() / 4, 2, 2)},
          P0,
          {*dofmap1, bs1,
           mdspanx22_t(facets1.data(), facets1.size() / 4, 2, 2)},
          P1T, bc0, bc1, fn,
          mdspanx2x_t(coeffs.data(), facets.size() / 4, 2, cstride), constants,
          cell_info0, cell_info1, perms);
    }
  }
}

In python/dolfinx/wrappers/petsc.cpp

m.def(
        "assemble_matrix_real_to_complex",
        [](Mat A, const dolfinx::fem::Form<PetscReal, PetscReal>& a,
          nb::ndarray<const PetscReal, nb::ndim<1>, nb::c_contig> constants,
          const std::map<std::pair<dolfinx::fem::IntegralType, int>,
          nb::ndarray<const PetscReal, nb::ndim<2>,
          nb::c_contig>>& coefficients,
          const std::vector<const dolfinx::fem::DirichletBC<PetscReal, PetscReal>*>& bcs,
          bool unrolled)
          {
            std::vector<std::reference_wrapper<
                const dolfinx::fem::DirichletBC<PetscReal, PetscReal>>>
                _bcs;
            for (auto bc : bcs)
            {
              assert(bc);
              _bcs.push_back(*bc);
            }   
            if (unrolled)
            {
              auto set_fn = dolfinx::la::petsc::Matrix::set_block_expand_fn(
                  A, a.function_spaces()[0]->dofmap()->bs(),
                  a.function_spaces()[1]->dofmap()->bs(), ADD_VALUES);
              dolfinx::fem::assemble_matrix_cast(
                  set_fn, a, std::span(constants.data(), constants.size()),
                  dolfinx_wrappers::py_to_cpp_coeffs(coefficients), _bcs);
            }
            else 
            {
              dolfinx::fem::assemble_matrix_cast(
                  dolfinx::la::petsc::Matrix::set_block_fn(A, ADD_VALUES),
                  a,
                  std::span(constants.data(), constants.size()),
                  dolfinx_wrappers::py_to_cpp_coeffs(coefficients),
                  _bcs);
            }
        },
        nb::arg("A"), nb::arg("a"), nb::arg("constants"), nb::arg("coeffs"),
        nb::arg("bcs"), nb::arg("unrolled") = false,
        "Assemble bilinear form into an existing PETSc matrix");

In python/dolfinx/fem/petsc.py to include

def assemble_matrix_mat_real_to_complex(
    A: PETSc.Mat,
    a: Form,
    bcs: list[DirichletBC] = [],
    diagonal: float = 1.0,
    constants=None,
    coeffs=None,
) -> PETSc.Mat:
    """Assemble bilinear form into a matrix.

    The returned matrix is not finalised, i.e. ghost values are not
    accumulated.
    """
    constants = _pack_constants(a._cpp_object) if constants is None else constants
    coeffs = _pack_coefficients(a._cpp_object) if coeffs is None else coeffs
    _bcs = [bc._cpp_object for bc in bcs]
    _cpp.fem.petsc.assemble_matrix_real_to_complex(A, a._cpp_object, constants, coeffs, _bcs)
    if a.function_spaces[0] is a.function_spaces[1]:
        A.assemblyBegin(PETSc.Mat.AssemblyType.FLUSH)
        A.assemblyEnd(PETSc.Mat.AssemblyType.FLUSH)
        _cpp.fem.petsc.insert_diagonal(A, a.function_spaces[0], _bcs, diagonal)
    return A

def create_mat_real_to_complex(a_form: ufl.Form, comm = MPI.COMM_WORLD) -> PETSc.Mat:
        A_form = dolfinx.fem.form(a_form, dtype=np.float64)
        sp = dolfinx.fem.create_sparsity_pattern(A_form)
        sp.finalize()
        import dolfinx.cpp as _cpp
        petsc_matrix = _cpp.la.petsc.create_matrix(comm, sp, "mpiaij")
        petsc_matrix = dolfinx.fem.petsc.assemble_matrix_mat_real_to_complex(petsc_matrix, A_form)
        petsc_matrix.assemble()
        return petsc_matrix

[garth-wells]

I think this should be pretty straightforward to support, and more generally for matrices and vectors with different float types.

[garth-wells]

A first step towards this would be removing all dynamic memory allocations inside low-level assembler functions (started in #3746, but with other changes would need to begin again).

In removing allocation, we pass buffers into the low-level functions. We can pass an extra buffer for the type conversion.