Files
mfem/tests/unit/dfem/test_diffusion.cpp
T

693 lines
22 KiB
C++

// Copyright (c) 2010-2025, Lawrence Livermore National Security, LLC. Produced
// at the Lawrence Livermore National Laboratory. All Rights reserved. See files
// LICENSE and NOTICE for details. LLNL-CODE-806117.
//
// This file is part of the MFEM library. For more information and source code
// availability visit https://mfem.org.
//
// MFEM is free software; you can redistribute it and/or modify it under the
// terms of the BSD-3 license. We welcome feedback and contributions, see file
// CONTRIBUTING.md for details.
#include "../unit_tests.hpp"
#include "mfem.hpp"
#ifdef MFEM_USE_MPI
#include "../linalg/test_same_matrices.hpp"
#include "../../../fem/dfem/doperator.hpp"
#include "../../../fem/dfem/backends/local_qf/prelude.hpp"
#include "../../../linalg/tensor_arrays.hpp"
using namespace mfem;
using namespace mfem::future;
#ifdef MFEM_USE_ENZYME
using dscalar_t = real_t;
#else
using dscalar_t = dual<real_t, real_t>;
#endif
// ────────────────────────────────────────────────────────────────────────────
template <int DIM>
struct Diffusion
{
using dvecd_t = tensor<dscalar_t, DIM>;
using matd_t = tensor<real_t, DIM, DIM>;
struct MFApply
{
MFEM_HOST_DEVICE inline auto operator()(
const dvecd_t &dudxi,
const matd_t &J,
const real_t &w,
dvecd_t &dvdxi) const
{
const auto invJ = inv(J);
const auto invJt = transpose(invJ);
dvdxi = (dudxi * invJ) * invJt * det(J) * w;
}
};
struct PASetup
{
MFEM_HOST_DEVICE inline auto operator()(
const matd_t &J,
const real_t &w,
matd_t &qdata) const
{
qdata = inv(J) * transpose(inv(J)) * det(J) * w;
}
};
struct PAApply
{
MFEM_HOST_DEVICE inline auto operator()(
const dvecd_t &dudxi,
const matd_t &qdata,
dvecd_t &dvdxi) const
{
dvdxi = qdata * dudxi;
};
};
};
// ────────────────────────────────────────────────────────────────────────────
// Global-QF diffusion. Used to exercise the GlobalQFBackend derivative cache
// with residual_size_on_qp = DIM*DIM > 1; the other GlobalQFBackend tests are
// all scalar Value-in/Value-out, where residual_size_on_qp == 1 and the cache
// index layout is degenerate.
template <int DIM>
struct GlobalDiffusion
{
struct MFApply
{
void operator()(tensor_array<const dscalar_t, DIM> &dudxi,
tensor_array<const real_t, DIM, DIM> &J,
tensor_array<const real_t> &w,
tensor_array<dscalar_t, DIM> &dvdxi) const
{
mfem::forall(w.size(), [=] MFEM_HOST_DEVICE(int q)
{
const auto invJ = inv(J(q));
dvdxi(q) = (dudxi(q) * invJ) * transpose(invJ) * det(J(q)) * w(q);
});
}
};
};
// ────────────────────────────────────────────────────────────────────────────
template <int DIM> struct VectorDiffusion
{
using dmatd_t = tensor<dscalar_t, DIM, DIM>;
using matd_t = tensor<real_t, DIM, DIM>;
struct MFApply
{
MFEM_HOST_DEVICE inline auto operator()(
const dmatd_t &dudxi,
const matd_t &J,
const real_t &w,
dmatd_t &dvdxi) const
{
const auto invJ = inv(J);
const auto invJt = transpose(invJ);
dvdxi = (dudxi * invJ) * invJt * det(J) * w;
}
};
};
// ────────────────────────────────────────────────────────────────────────────
template <int DIM, typename QFBackend = LocalQFBackend>
void diffusion(const char *filename, int p)
{
CAPTURE(filename, DIM, p);
Mesh smesh(filename);
ParMesh pmesh(MPI_COMM_WORLD, smesh);
MFEM_VERIFY(pmesh.Dimension() == DIM, "Mesh dimension mismatch");
pmesh.EnsureNodes();
auto *nodes = static_cast<ParGridFunction *>(pmesh.GetNodes());
smesh.Clear();
p = std::max(p, pmesh.GetNodalFESpace()->GetMaxElementOrder());
Array<int> all_domain_attr;
if (pmesh.attributes.Size() > 0)
{
all_domain_attr.SetSize(pmesh.attributes.Max());
all_domain_attr = 1;
}
ParFiniteElementSpace *mfes = nodes->ParFESpace();
const auto *ir = &IntRules.Get(pmesh.GetTypicalElementGeometry(), 2 * p);
H1_FECollection fec(p, DIM);
static constexpr int U = 0, Coords = 1;
SECTION("Scalar")
{
ParFiniteElementSpace pfes(&pmesh, &fec);
ParGridFunction x(&pfes), y(&pfes), z(&pfes);
Vector xtvec(pfes.GetTrueVSize()), ytvec(pfes.GetTrueVSize()),
ztvec(pfes.GetTrueVSize());
xtvec.Randomize(1);
x.SetFromTrueDofs(xtvec);
ParBilinearForm blf_fa(&pfes);
blf_fa.AddDomainIntegrator(new DiffusionIntegrator(ir));
blf_fa.SetAssemblyLevel(AssemblyLevel::FULL);
blf_fa.Assemble();
blf_fa.Finalize();
const auto in_fds = std::vector
{
FieldDescriptor{ U, &pfes },
FieldDescriptor{ Coords, mfes }
};
const auto out_fds = std::vector{ FieldDescriptor{ U, &pfes } };
SECTION("Scalar Action")
{
DifferentiableOperator dop_mf(in_fds, out_fds, pmesh);
typename Diffusion<DIM>::MFApply mf_apply_qf;
dop_mf.AddDomainIntegrator<QFBackend>(
mf_apply_qf,
Inputs<Gradient<U>, Gradient<Coords>, Weight> {},
Outputs<Gradient<U>> {},
*ir, all_domain_attr);
Vector nodestv;
nodes->GetTrueDofs(nodestv);
pfes.GetRestrictionMatrix()->Mult(x, xtvec);
MultiVector X{xtvec, nodestv};
MultiVector Z{ztvec};
dop_mf.Mult(X, Z);
blf_fa.Mult(x, y);
pfes.GetProlongationMatrix()->MultTranspose(y, ytvec);
ytvec -= ztvec;
real_t norm_global = 0.0;
real_t norm_local = ytvec.Normlinf();
MPI_Allreduce(&norm_local, &norm_global, 1, MPI_DOUBLE, MPI_MAX,
pmesh.GetComm());
REQUIRE(norm_global == MFEM_Approx(0.0));
MPI_Barrier(MPI_COMM_WORLD);
}
SECTION("Scalar Action Partial Assembly")
{
static constexpr int QData = 2;
QuadratureSpace qspace(pmesh, *ir);
VectorQuadratureSpace qspace_vec(qspace, DIM * DIM);
QuadratureFunction qd(qspace_vec);
DifferentiableOperator setupPAData(
{
{Coords, mfes}
},
{
{QData, &qspace_vec}
}, pmesh);
typename Diffusion<DIM>::PASetup pa_setup_qf;
setupPAData.AddDomainIntegrator<QFBackend>(
pa_setup_qf,
Inputs<Gradient<Coords>, Weight> {},
Outputs<Identity<QData>> {},
*ir, all_domain_attr);
{
Vector nodestv;
nodes->GetTrueDofs(nodestv);
MultiVector X{nodestv};
MultiVector Y{qd};
setupPAData.Mult(X, Y);
}
DifferentiableOperator applyPAData(
{
{U, &pfes}, {QData, &qspace_vec}
},
{
{U, &pfes}
}, pmesh);
typename Diffusion<DIM>::PAApply pa_apply_qf;
applyPAData.AddDomainIntegrator<QFBackend>(
pa_apply_qf,
Inputs<Gradient<U>, Identity<QData>> {},
Outputs<Gradient<U>> {},
*ir, all_domain_attr);
{
pfes.GetRestrictionMatrix()->Mult(x, xtvec);
MultiVector X{xtvec, qd};
MultiVector Z{ztvec};
applyPAData.Mult(X, Z);
}
blf_fa.Mult(x, y);
pfes.GetProlongationMatrix()->MultTranspose(y, ytvec);
ytvec -= ztvec;
real_t norm_global = 0.0;
real_t norm_local = ytvec.Normlinf();
MPI_Allreduce(&norm_local, &norm_global, 1, MPI_DOUBLE, MPI_MAX,
pmesh.GetComm());
REQUIRE(norm_global == MFEM_Approx(0.0));
MPI_Barrier(MPI_COMM_WORLD);
}
SECTION("Scalar Action Linearized")
{
DifferentiableOperator dop_mf(in_fds, out_fds, pmesh);
typename Diffusion<DIM>::MFApply mf_apply_qf;
dop_mf.AddDomainIntegrator<QFBackend>(
mf_apply_qf,
Inputs<Gradient<U>, Gradient<Coords>, Weight> {},
Outputs<Gradient<U>> {},
*ir, all_domain_attr,
Derivatives<U> {});
pfes.GetRestrictionMatrix()->Mult(x, xtvec);
Vector nodestv;
nodes->GetTrueDofs(nodestv);
MultiVector X{xtvec, nodestv};
MultiVector Z{ztvec};
auto ddop = dop_mf.GetDerivative(U, X);
// Randomize again s.t. the PA setup like cache can't
// trivially succeed by caching one direction only.
xtvec.Randomize(567);
x.SetFromTrueDofs(xtvec);
Vector dztvec(ztvec.Size());
MultiVector DZ{dztvec};
ddop->Mult(X[0], DZ);
blf_fa.Mult(x, y);
pfes.GetProlongationMatrix()->MultTranspose(y, ytvec);
ytvec -= dztvec;
real_t norm_global = 0.0;
real_t norm_local = ytvec.Normlinf();
MPI_Allreduce(&norm_local, &norm_global, 1, MPI_DOUBLE, MPI_MAX,
pmesh.GetComm());
REQUIRE(norm_global == MFEM_Approx(0.0));
MPI_Barrier(MPI_COMM_WORLD);
}
SECTION("Scalar SparseMatrix")
{
DifferentiableOperator dop_mf(in_fds, out_fds, pmesh);
typename Diffusion<DIM>::MFApply mf_apply_qf;
dop_mf.AddDomainIntegrator<QFBackend>(
mf_apply_qf,
Inputs<Gradient<U>, Gradient<Coords>, Weight> {},
Outputs<Gradient<U>> {},
*ir, all_domain_attr,
Derivatives<U> {});
pfes.GetRestrictionMatrix()->Mult(x, xtvec);
Vector nodestv;
nodes->GetTrueDofs(nodestv);
MultiVector X{xtvec, nodestv};
auto dRdU = dop_mf.GetDerivative(U, X);
SparseMatrix *A = nullptr;
dRdU->Assemble(A);
TestSameMatrices(*A, blf_fa.SpMat());
delete A;
MPI_Barrier(MPI_COMM_WORLD);
}
SECTION("Scalar Assemble Diagonal")
{
DifferentiableOperator dop_mf(in_fds, out_fds, pmesh);
typename Diffusion<DIM>::MFApply mf_apply_qf;
dop_mf.AddDomainIntegrator<QFBackend>(
mf_apply_qf,
Inputs<Gradient<U>, Gradient<Coords>, Weight> {},
Outputs<Gradient<U>> {},
*ir, all_domain_attr,
Derivatives<U> {});
pfes.GetRestrictionMatrix()->Mult(x, xtvec);
Vector nodestv;
nodes->GetTrueDofs(nodestv);
MultiVector X{xtvec, nodestv};
auto dRdU = dop_mf.GetDerivative(U, X);
Vector dfem_diagonal(pfes.GetTrueVSize());
dRdU->AssembleDiagonal(dfem_diagonal);
Vector mfem_diagonal(pfes.GetTrueVSize());
blf_fa.AssembleDiagonal(mfem_diagonal);
dfem_diagonal -= mfem_diagonal;
real_t norm_global = 0.0;
real_t norm_local = dfem_diagonal.Normlinf();
MPI_Allreduce(&norm_local, &norm_global, 1, MPI_DOUBLE, MPI_MAX,
pmesh.GetComm());
REQUIRE(norm_global == MFEM_Approx(0.0));
MPI_Barrier(MPI_COMM_WORLD);
}
}
SECTION("Vector")
{
ParFiniteElementSpace vpfes(&pmesh, &fec, DIM);
ParGridFunction vx(&vpfes), vy(&vpfes);
Vector vX(vpfes.GetTrueVSize()), vY(vpfes.GetTrueVSize()),
vZ(vpfes.GetTrueVSize());
SECTION("Vector Diffusion Action")
{
vX.Randomize(1);
vx.SetFromTrueDofs(vX);
DifferentiableOperator dop_mf(
{
{U, &vpfes},
{Coords, mfes},
},
{
{U, &vpfes}
}, pmesh);
typename VectorDiffusion<DIM>::MFApply mf_vector_diffusion_qf;
dop_mf.AddDomainIntegrator<QFBackend>(
mf_vector_diffusion_qf,
Inputs<Gradient<U>, Gradient<Coords>, Weight> {},
Outputs<Gradient<U>> {},
*ir, all_domain_attr);
Vector nodestv;
nodes->GetTrueDofs(nodestv);
MultiVector X{vX, nodestv};
MultiVector Z{vZ};
dop_mf.Mult(X, Z);
ParBilinearForm vblf_fa(&vpfes);
vblf_fa.AddDomainIntegrator(new VectorDiffusionIntegrator(ir));
vblf_fa.SetAssemblyLevel(AssemblyLevel::LEGACYFULL);
vblf_fa.Assemble();
vblf_fa.Finalize();
vblf_fa.Mult(vx, vy);
vpfes.GetProlongationMatrix()->MultTranspose(vy, vY);
vY -= vZ;
real_t norm_global = 0.0;
real_t norm_local = vY.Normlinf();
MPI_Allreduce(&norm_local, &norm_global, 1, MPI_DOUBLE, MPI_MAX,
pmesh.GetComm());
REQUIRE(norm_global == MFEM_Approx(0.0));
MPI_Barrier(MPI_COMM_WORLD);
}
SECTION("Vector Diffusion Action Linearized")
{
vX.Randomize(1);
vx.SetFromTrueDofs(vX);
DifferentiableOperator dop_mf(
{
{U, &vpfes},
{Coords, mfes},
},
{
{U, &vpfes}
}, pmesh);
typename VectorDiffusion<DIM>::MFApply mf_vector_diffusion_qf;
dop_mf.AddDomainIntegrator<QFBackend>(
mf_vector_diffusion_qf,
Inputs<Gradient<U>, Gradient<Coords>, Weight> {},
Outputs<Gradient<U>> {},
*ir, all_domain_attr,
Derivatives<U> {});
Vector nodestv;
nodes->GetTrueDofs(nodestv);
MultiVector X{vX, nodestv};
const auto ddop = dop_mf.GetDerivative(U, X);
MultiVector Z{vZ};
ddop->Mult(vX, Z);
ParBilinearForm vblf_fa(&vpfes);
vblf_fa.AddDomainIntegrator(new VectorDiffusionIntegrator(ir));
vblf_fa.SetAssemblyLevel(AssemblyLevel::LEGACYFULL);
vblf_fa.Assemble();
vblf_fa.Finalize();
vblf_fa.Mult(vx, vy);
vpfes.GetProlongationMatrix()->MultTranspose(vy, vY);
vY -= vZ;
real_t norm_global = 0.0;
real_t norm_local = vY.Normlinf();
MPI_Allreduce(&norm_local, &norm_global, 1, MPI_DOUBLE, MPI_MAX,
pmesh.GetComm());
REQUIRE(norm_global == MFEM_Approx(0.0));
MPI_Barrier(MPI_COMM_WORLD);
}
SECTION("Vector SparseMatrix")
{
vX.Randomize(1);
vx.SetFromTrueDofs(vX);
DifferentiableOperator dop_mf(
{
{U, &vpfes},
{Coords, mfes},
},
{
{U, &vpfes}
}, pmesh);
typename VectorDiffusion<DIM>::MFApply mf_vector_diffusion_qf;
dop_mf.AddDomainIntegrator<QFBackend>(
mf_vector_diffusion_qf,
Inputs<Gradient<U>, Gradient<Coords>, Weight> {},
Outputs<Gradient<U>> {},
*ir, all_domain_attr,
Derivatives<U> {});
Vector nodestv;
nodes->GetTrueDofs(nodestv);
MultiVector X{vX, nodestv};
const auto ddop = dop_mf.GetDerivative(U, X);
MultiVector Z{vZ};
ddop->Mult(vX, Z);
ParBilinearForm vblf_fa(&vpfes);
vblf_fa.AddDomainIntegrator(new VectorDiffusionIntegrator(ir));
vblf_fa.SetAssemblyLevel(AssemblyLevel::LEGACYFULL);
vblf_fa.Assemble();
vblf_fa.Finalize();
SparseMatrix *A = nullptr;
ddop->Assemble(A);
TestSameMatrices(*A, vblf_fa.SpMat());
delete A;
MPI_Barrier(MPI_COMM_WORLD);
}
}
}
// ────────────────────────────────────────────────────────────────────────────
TEST_CASE("dFEM Diffusion 2D", "[Parallel][dFEM][GPU]")
{
const auto p = GenAll({1}, {2, 3});
const auto meshs = { "../../data/inline-quad.mesh" };
const auto extra = { "../../data/star.mesh",
"../../data/star-q3.mesh",
"../../data/rt-2d-q3.mesh",
"../../data/periodic-square.mesh"
};
diffusion<2>(GenAll(meshs, extra), p);
}
// ────────────────────────────────────────────────────────────────────────────
// The GlobalQFBackend writes the derivative qp_cache, but Assemble and
// AssembleDiagonal are served by the LocalQF implementations (see
// GlobalQFBackend::MakeDerivativeAssemble*), so writer and readers live in
// different backends and must agree on the cache layout.
template <int DIM>
void diffusion_globalqf(const char *filename, int p)
{
CAPTURE(filename, DIM, p);
Mesh smesh(filename);
ParMesh pmesh(MPI_COMM_WORLD, smesh);
MFEM_VERIFY(pmesh.Dimension() == DIM, "Mesh dimension mismatch");
pmesh.EnsureNodes();
auto *nodes = static_cast<ParGridFunction *>(pmesh.GetNodes());
smesh.Clear();
p = std::max(p, pmesh.GetNodalFESpace()->GetMaxElementOrder());
Array<int> all_domain_attr;
if (pmesh.attributes.Size() > 0)
{
all_domain_attr.SetSize(pmesh.attributes.Max());
all_domain_attr = 1;
}
ParFiniteElementSpace *mfes = nodes->ParFESpace();
const auto *ir = &IntRules.Get(pmesh.GetTypicalElementGeometry(), 2 * p);
H1_FECollection fec(p, DIM);
ParFiniteElementSpace pfes(&pmesh, &fec);
static constexpr int U = 0, Coords = 1;
ParGridFunction x(&pfes), y(&pfes);
Vector xtvec(pfes.GetTrueVSize()), ytvec(pfes.GetTrueVSize()),
ztvec(pfes.GetTrueVSize());
xtvec.Randomize(1);
x.SetFromTrueDofs(xtvec);
ParBilinearForm blf_fa(&pfes);
blf_fa.AddDomainIntegrator(new DiffusionIntegrator(ir));
blf_fa.SetAssemblyLevel(AssemblyLevel::FULL);
blf_fa.Assemble();
blf_fa.Finalize();
const auto in_fds = std::vector
{
FieldDescriptor{ U, &pfes },
FieldDescriptor{ Coords, mfes }
};
const auto out_fds = std::vector{ FieldDescriptor{ U, &pfes } };
DifferentiableOperator dop(in_fds, out_fds, pmesh);
typename GlobalDiffusion<DIM>::MFApply global_qfn;
dop.AddDomainIntegrator<GlobalQFBackend>(
global_qfn,
Inputs<Gradient<U>, Gradient<Coords>, Weight> {},
Outputs<Gradient<U>> {},
*ir, all_domain_attr,
Derivatives<U> {});
Vector nodestv;
nodes->GetTrueDofs(nodestv);
pfes.GetRestrictionMatrix()->Mult(x, xtvec);
MultiVector X{ xtvec, nodestv };
auto dRdU = dop.GetDerivative(U, X);
const auto max_error = [&](const Vector &v)
{
real_t norm_global = 0.0, norm_local = v.Normlinf();
MPI_Allreduce(&norm_local, &norm_global, 1, MPI_DOUBLE, MPI_MAX,
pmesh.GetComm());
return norm_global;
};
// GlobalQF setup writer -> GlobalQF apply reader
{
MultiVector Z{ ztvec };
dRdU->Mult(xtvec, Z);
blf_fa.Mult(x, y);
pfes.GetProlongationMatrix()->MultTranspose(y, ytvec);
ytvec -= ztvec;
REQUIRE(max_error(ytvec) == MFEM_Approx(0.0));
}
// GlobalQF setup writer -> GlobalQF apply-transpose reader
{
Vector wtvec(pfes.GetTrueVSize());
wtvec.Randomize(0x9e3779b9);
MultiVector W{ wtvec }, Z{ ztvec };
dRdU->MultTranspose(W, Z);
ParGridFunction w(&pfes);
w.SetFromTrueDofs(wtvec);
blf_fa.MultTranspose(w, y);
pfes.GetProlongationMatrix()->MultTranspose(y, ytvec);
ytvec -= ztvec;
REQUIRE(max_error(ytvec) == MFEM_Approx(0.0));
}
// GlobalQF setup writer -> LocalQF assemble reader
{
SparseMatrix *A = nullptr;
dRdU->Assemble(A);
TestSameMatrices(*A, blf_fa.SpMat());
delete A;
}
// GlobalQF setup writer -> LocalQF assemble-diagonal reader
{
Vector dfem_diag(pfes.GetTrueVSize()), mfem_diag(pfes.GetTrueVSize());
dRdU->AssembleDiagonal(dfem_diag);
blf_fa.AssembleDiagonal(mfem_diag);
dfem_diag -= mfem_diag;
REQUIRE(max_error(dfem_diag) == MFEM_Approx(0.0));
}
MPI_Barrier(MPI_COMM_WORLD);
}
// ────────────────────────────────────────────────────────────────────────────
TEST_CASE("dFEM Diffusion GlobalQF cache 2D", "[Parallel][dFEM]")
{
if constexpr (!mfem_use_gpu)
{
const auto p = GenAll({1}, {2, 3});
const auto meshs = { "../../data/inline-quad.mesh" };
const auto extra = { "../../data/star.mesh" };
diffusion_globalqf<2>(GenAll(meshs, extra), p);
}
}
// ────────────────────────────────────────────────────────────────────────────
TEST_CASE("dFEM Diffusion GlobalQF cache 3D", "[Parallel][dFEM]")
{
if constexpr (!mfem_use_gpu)
{
const auto p = GenAll({1}, {2, 3});
const auto meshs = { "../../data/inline-hex.mesh" };
const auto extra = { "../../data/fichera.mesh" };
diffusion_globalqf<3>(GenAll(meshs, extra), p);
}
}
// ────────────────────────────────────────────────────────────────────────────
TEST_CASE("dFEM Diffusion 3D", "[Parallel][dFEM][GPU]")
{
const auto p = GenAll({1}, {2, 3});
const auto meshs = { "../../data/inline-hex.mesh" };
const auto extra = { "../../data/fichera.mesh",
"../../data/fichera-q3.mesh",
"../../data/toroid-hex.mesh",
"../../data/periodic-cube.mesh"
};
diffusion<3>(GenAll(meshs, extra), p);
}
#endif // MFEM_USE_MPI