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

423 lines
11 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"
#include "../fem/dfem/doperator.hpp"
#include "../fem/dfem/backends/local_qf/prelude.hpp"
#include "linalg/tensor_arrays.hpp"
#ifdef MFEM_USE_MPI
using namespace mfem;
using namespace mfem::future;
using mfem::future::tensor;
#ifdef MFEM_USE_ENZYME
using dscalar_t = real_t;
#else
using mfem::future::dual;
using dscalar_t = dual<real_t, real_t>;
#endif
constexpr int DIM = 2;
class DummyParameterSpace : public ParameterSpace
{
public:
class Bimpl : public Operator
{
virtual void Mult(const Vector &x, Vector &y) const
{
for (int i = 0; i < y.Size(); i++)
{
y(i) = x(0);
}
}
};
class Btimpl : public Operator
{
virtual void Mult(const Vector &x, Vector &y) const
{
y(0) = x(0);
}
};
DummyParameterSpace() : ParameterSpace(1) {}
virtual int GetTrueVSize() const override
{
return 1;
}
virtual int GetVSize() const override
{
return 1;
}
virtual const Operator* GetB() const override
{
if (!B)
{
B.reset(new Bimpl());
}
return B.get();
}
virtual const Operator* GetBt() const override
{
if (!Bt)
{
Bt.reset(new Btimpl());
}
return Bt.get();
}
};
struct massqf
{
inline MFEM_HOST_DEVICE
void operator()(
tensor_array<const real_t> &u,
tensor_array<const real_t, DIM, DIM> &J,
tensor_array<const real_t> &w,
tensor_array<real_t> &out1,
tensor_array<real_t> &out2) const
{
for (size_t q = 0; q < u.size(); q++)
{
const auto v = u(q) * det(J(q)) * w(q);
out1(q) = v;
out2(q) = v;
}
}
};
struct mass_diffusion_qdata_qf
{
inline MFEM_HOST_DEVICE
void operator()(
tensor_array<const real_t> &u,
tensor_array<const real_t, DIM> &dudxi,
tensor_array<const real_t, DIM, DIM> &J,
tensor_array<const real_t, DIM, DIM> &qdata,
tensor_array<const real_t> &w,
tensor_array<const real_t> &dummy_parameter,
tensor_array<real_t> &out1,
tensor_array<real_t, DIM> &out2,
tensor_array<real_t, DIM, DIM> &out3) const
{
for (size_t q = 0; q < u.size(); q++)
{
const auto invJq = inv(J(q));
const auto detJq = det(J(q));
out1(q) = u(q) * detJq * w(q);
// out2(q) = (dudxi(q) * invJq) * transpose(invJq) * (detJq * w(q));
out3(q) = J(q);
}
jit_bounds(dudxi, J, w, out2, u.size());
}
// XXX: Attribute instrumentation does not work due to ABI differences that
// change the argument number.
//__attribute__((annotate("jit", 5)))
void jit_bounds(
tensor_array<const real_t, DIM> &dudxi,
tensor_array<const real_t, DIM, DIM> &J,
tensor_array<const real_t> &w,
tensor_array<real_t, DIM> &out,
size_t NQ) const
{
for (size_t q = 0; q < NQ; q++)
{
const auto invJq = inv(J(q));
const auto detJq = det(J(q));
out(q) = (dudxi(q) * invJq) * transpose(invJq) * (detJq * w(q));
}
}
};
struct massqflocal
{
inline MFEM_HOST_DEVICE
void operator()(
const tensor<real_t> &u,
const tensor<real_t, DIM, DIM> &J,
const tensor<real_t> &w,
tensor<real_t> &out1,
tensor<real_t> &out2) const
{
const auto v = u * det(J) * w;
out1 = v;
out2 = v;
}
};
TEST_CASE("dFEM Multiple Outputs", "[Parallel][dFEM]")
{
const bool all_tests = launch_all_non_regression_tests;
const auto p = !all_tests ? 2 : GENERATE(1, 2, 3);
const char *filename = "../../data/inline-quad.mesh";
CAPTURE(filename, DIM, p);
Mesh smesh(filename);
MFEM_ASSERT(smesh.Dimension() == DIM, "DIM and mesh dimension have to match");
ParMesh pmesh(MPI_COMM_WORLD, smesh);
pmesh.EnsureNodes();
auto* nodes = static_cast<ParGridFunction*>(pmesh.GetNodes());
smesh.Clear();
H1_FECollection fec(p, DIM);
ParFiniteElementSpace fes(&pmesh, &fec);
const auto *ir = &IntRules.Get(pmesh.GetTypicalElementGeometry(), 2 * p);
ParGridFunction x(&fes), y(&fes), z(&fes);
ConstantCoefficient one(1.0);
Array<int> all_domain_attr;
if (pmesh.attributes.Size() > 0)
{
all_domain_attr.SetSize(pmesh.attributes.Max());
all_domain_attr = 1;
}
// {
// Array<int> inoffsets(3);
// inoffsets[0] = 0;
// inoffsets[1] = fes.GetTrueVSize();
// inoffsets[2] = nodes->ParFESpace()->GetTrueVSize();
// inoffsets.PartialSum();
// BlockVector X(inoffsets);
// X.GetBlock(0).Randomize(1);
// X.GetBlock(1) = *nodes;
// x.SetFromTrueDofs(X.GetBlock(0));
// Array<int> outoffsets(2);
// outoffsets[0] = 0;
// outoffsets[1] = fes.GetTrueVSize();
// outoffsets.PartialSum();
// BlockVector Z(outoffsets);
// ParBilinearForm blf(&fes);
// blf.AddDomainIntegrator(new MassIntegrator(one, ir));
// blf.SetAssemblyLevel(AssemblyLevel::PARTIAL);
// blf.Assemble();
// blf.Mult(x, y);
// Vector Y(fes.GetTrueVSize());
// fes.GetProlongationMatrix()->MultTranspose(y, Y);
// static constexpr int U = 0, COORDINATES = 1, V = 2;
// const std::vector<FieldDescriptor> in
// {
// {U, &fes},
// {COORDINATES, nodes->ParFESpace()}
// };
// const std::vector<FieldDescriptor> out // test spaces?
// {
// {V, &fes},
// };
// DifferentiableOperator dop(in, out, pmesh);
// auto derivatives = std::integer_sequence<size_t, U> {};
// auto mass_qfunc = massqf{};
// dop.AddDomainIntegrator(mass_qfunc,
// tuple{ Value<U>{}, Gradient<COORDINATES>{}, Weight{} },
// tuple{ Value<V>{}, Value<V>{} },
// *ir, all_domain_attr, derivatives);
// fes.GetRestrictionMatrix()->Mult(x, X.GetBlock(0));
// dop.Mult(X, Z);
// Vector Y0(Y);
// Y0 *= 2.0;
// Y0 -= Z.GetBlock(0);
// real_t norm_g, norm_l = Y0.Normlinf();
// MPI_Allreduce(&norm_l, &norm_g, 1, MPI_DOUBLE, MPI_MAX, pmesh.GetComm());
// REQUIRE(norm_g == MFEM_Approx(0.0));
// MPI_Barrier(MPI_COMM_WORLD);
// auto ddop = dop.GetDerivative(U, X);
// ddop->Mult(X.GetBlock(0), Z);
// Y0 = Y;
// Y0 *= 2.0;
// Y0 -= Z.GetBlock(0);
// norm_l = Y0.Normlinf();
// MPI_Allreduce(&norm_l, &norm_g, 1, MPI_DOUBLE, MPI_MAX, pmesh.GetComm());
// REQUIRE(norm_g == MFEM_Approx(0.0));
// MPI_Barrier(MPI_COMM_WORLD);
// }
{
QuadratureSpace qs(pmesh, *ir);
QuadratureFunction qdata(qs, DIM*DIM);
DummyParameterSpace dps;
ParameterFunction dpf(dps);
dpf = 9.12345;
auto coef_func = [](const Vector &coords)
{
return coords[0] * coords[1] * (DIM == 3 ? coords[2] : 1.0);
};
FunctionCoefficient coef(coef_func);
x.ProjectCoefficient(coef);
Vector xtvec, ytvec, ytvecmfem;
x.GetTrueDofs(xtvec);
ytvec.SetSize(xtvec.Size());
ytvecmfem.SetSize(xtvec.Size());
Vector nodestvec;
nodes->GetTrueDofs(nodestvec);
qdata = 123.0;
Vector yqdata(qdata.Size());
MultiVector X{xtvec, nodestvec, qdata, dpf};
MultiVector Z{ytvec, yqdata};
ParBilinearForm blf(&fes);
blf.AddDomainIntegrator(new MassIntegrator(ir));
blf.AddDomainIntegrator(new DiffusionIntegrator(ir));
blf.SetAssemblyLevel(AssemblyLevel::PARTIAL);
blf.Assemble();
blf.Mult(x, y);
fes.GetProlongationMatrix()->MultTranspose(y, ytvecmfem);
static constexpr int U = 0, COORDINATES = 1, V = 2, S = 3, L = 4;
const std::vector<FieldDescriptor> in
{
{U, &fes},
{COORDINATES, nodes->ParFESpace()},
{S, &qdata},
{L, &dps}
};
const std::vector<FieldDescriptor> out
{
{V, &fes},
{S, &qdata}
};
{
DifferentiableOperator dop(in, out, pmesh);
dop.SetQLayouts({{Value<U>{}, {1, 0}}}, {});
auto derivatives = std::integer_sequence<size_t, U> {};
auto mass_diffusion_qfunc = mass_diffusion_qdata_qf{};
dop.AddDomainIntegrator(
mass_diffusion_qfunc,
tuple{Value<U>{}, Gradient<U>{}, Gradient<COORDINATES>{}, Identity<S>{}, Weight{}, Value<L>{}},
tuple{Value<V>{}, Gradient<V>{}, Identity<S>{}},
*ir, all_domain_attr, derivatives);
fes.GetRestrictionMatrix()->Mult(x, xtvec);
dop.Mult(X, Z);
Vector Y0(ytvecmfem);
Y0 -= Z[0];
real_t norm_l = Y0.Normlinf();
real_t norm_g = norm_l;
MPI_Allreduce(&norm_l, &norm_g, 1, MPI_DOUBLE, MPI_MAX, pmesh.GetComm());
REQUIRE(norm_g == MFEM_Approx(0.0));
MPI_Barrier(MPI_COMM_WORLD);
auto ddop = dop.GetDerivative(U, X);
ddop->Mult(X[0], Z);
Y0 = ytvecmfem;
Y0 -= Z[0];
norm_l = Y0.Normlinf();
norm_g = norm_l;
MPI_Allreduce(&norm_l, &norm_g, 1, MPI_DOUBLE, MPI_MAX, pmesh.GetComm());
REQUIRE(norm_g == MFEM_Approx(0.0));
MPI_Barrier(MPI_COMM_WORLD);
}
{
static constexpr int W = 0;
ParBilinearForm blf(&fes);
blf.AddDomainIntegrator(new MassIntegrator(ir));
blf.SetAssemblyLevel(AssemblyLevel::PARTIAL);
blf.Assemble();
blf.Mult(x, y);
fes.GetProlongationMatrix()->MultTranspose(y, ytvecmfem);
const std::vector<FieldDescriptor> in
{
{U, &fes},
{COORDINATES, nodes->ParFESpace()},
};
const std::vector<FieldDescriptor> out
{
{V, &fes},
{W, &fes},
};
DifferentiableOperator dop(in, out, pmesh);
auto mass_qfunclocal = massqflocal{};
dop.AddDomainIntegrator<LocalQFBackend>(
mass_qfunclocal,
tuple{Value<U>{}, Gradient<COORDINATES>{}, Weight{}},
tuple{Value<V>{}, Value<W>{}},
*ir, all_domain_attr);
Vector nodestv;
nodes->GetTrueDofs(nodestv);
fes.GetRestrictionMatrix()->Mult(x, xtvec);
Vector ztvec(xtvec.Size());
Vector zztvec(xtvec.Size());
MultiVector X{xtvec, nodestv};
MultiVector Z{ztvec, zztvec};
dop.Mult(X, Z);
Vector Y0(ytvecmfem);
Y0 -= Z[0];
real_t norm_l = Y0.Normlinf();
real_t norm_g = norm_l;
MPI_Allreduce(&norm_l, &norm_g, 1, MPI_DOUBLE, MPI_MAX, pmesh.GetComm());
REQUIRE(norm_g == MFEM_Approx(0.0));
Vector Y1(ytvecmfem);
Y1 -= Z[1];
norm_l = Y1.Normlinf();
norm_g = norm_l;
MPI_Allreduce(&norm_l, &norm_g, 1, MPI_DOUBLE, MPI_MAX, pmesh.GetComm());
REQUIRE(norm_g == MFEM_Approx(0.0));
MPI_Barrier(MPI_COMM_WORLD);
}
}
}
#endif // MFEM_USE_MPI