Compare commits

...
30 Commits
Author SHA1 Message Date
Sohail Reddy a3137af040 Reorded member functions and removed redundant code 2026-08-13 15:22:34 -07:00
Sohail Reddy 489f87acff Updated with MultiVector changes 2026-08-12 09:11:22 -07:00
Sohail Reddy 6ada45f19a Merge branch 'multi-vector-dev' into multiapp-rm-vec 2026-08-11 20:44:04 -07:00
Sohail Reddy 1370cb4f36 Removed support for monolithic vectors in multiapps... functionality can be implemented by overloading existing methods in Operator class 2026-08-11 20:01:31 -07:00
Sohail Reddy c1ea8aefb7 Added support for MultiVector and cleaned up. 2026-08-11 13:40:38 -07:00
Sohail Reddy 04444dc37e Merge branch 'multi-vector-dev' into multiapp-io 2026-08-03 15:39:46 -07:00
Sohail Reddy 6f204dde67 Store intermediate fields for jacobian ccompuations 2026-08-03 11:26:36 -07:00
Sohail Reddy f04f98752b Merge branch 'stdmap-generic' into multiapp-io 2026-07-27 18:48:58 -07:00
Sohail Reddy 9b31004ce5 Changed GenericMap to GenericFieldMap 2026-07-27 18:24:51 -07:00
Sohail Reddy 04571c5ddd Using new GenericMap inplace of NamedFieldMap 2026-07-27 16:50:31 -07:00
Sohail Reddy e448b71831 Merge branch 'stdmap-generic' into multiapp-io 2026-07-27 15:56:49 -07:00
Sohail Reddy 48e415e17e Added a generic version of the NamedFieldsMap 2026-07-27 15:56:35 -07:00
Sohail Reddy 8b249e8004 Added Mult(Vector,Vector) support for app call 2026-07-27 14:49:30 -07:00
Sohail Reddy 1585e7d9f1 Updated GraphGradient 2026-07-21 20:16:35 -07:00
Sohail Reddy ae4a2ee9a6 Added easier method to define input and output size 2026-07-21 19:43:58 -07:00
Sohail Reddy 14e8c6ce45 Removed DataNodes for input and output fields to the graph 2026-07-21 19:36:59 -07:00
Sohail Reddy 05f5967267 fixed typo 2026-07-21 15:11:03 -07:00
Sohail Reddy d9e0018e98 Changed function naming JVP/VJP to GradientMult/GradientMultTranspose and added topological sorting of the graph nodes 2026-07-21 14:19:15 -07:00
Sohail Reddy f08dba7b84 Clean up and added operator() to execute nodes 2026-07-15 20:18:53 -07:00
Sohail Reddy 7e75c7f6fb Changed Field type from Source/Target to Output/Input 2026-07-15 19:31:26 -07:00
Sohail Reddy a1962fa492 Changed field collection name 2026-07-15 19:00:42 -07:00
Sohail Reddy ce7fefb6eb Bug fix for indexing into input and output Field vectors 2026-07-15 18:50:48 -07:00
Sohail Reddy 503b286f98 Removed support for FieldEdge and improved input/output field indexing 2026-07-15 18:48:15 -07:00
Sohail Reddy cd0616fae0 Improved support for input and output fields andcleaned up old infrastructure 2026-07-15 16:38:28 -07:00
Sohail Reddy 1ff144a358 Cleaned up field and edge fetching and adding inputs/outputs 2026-07-14 16:23:54 -07:00
Sohail Reddy 2c02d1b3cb Added initial support for Input and Output Fields to Nodes 2026-07-14 13:17:24 -07:00
Sohail Reddy f77dc8d7e9 Removed support for Vectors in FieldEdge and removed unused features 2026-07-13 11:28:28 -07:00
Sohail Reddy c23b850b77 Added FDJacobian::GetGradient to update finite difference Jacobian 2026-07-10 01:14:02 -07:00
Sohail Reddy ebda52d76c Merge branch 'gslib-uint-fix' into multiapp-mult 2026-07-09 20:05:30 -07:00
Sohail Reddy c8ab3cbf69 Branched off master 2026-07-09 15:34:24 -07:00
7 changed files with 2501 additions and 32 deletions
+79 -30
View File
@@ -25,21 +25,35 @@
namespace mfem
{
/// Lightweight adaptor over an std::map from strings to pointer to T
template<typename T>
class NamedFieldsMap
/// Lightweight adaptor over an std::map from type K to type to V
template<typename K, typename V,
typename = typename std::enable_if<std::is_default_constructible<V>::value>::type>
class GenericFieldMap
{
private:
static constexpr bool ValueIsPointer = std::is_pointer<V>::value;
public:
typedef std::map<std::string, T*> MapType;
typedef std::map<K, V> MapType;
typedef typename MapType::iterator iterator;
typedef typename MapType::const_iterator const_iterator;
/// Register field @a field with name @a fname
/** Replace existing field associated with @a fname (and optionally
delete associated pointer if @a own_data is true) */
void Register(const std::string& fname, T* field, bool own_data)
/// Register field @a field with name @a key
/// Only enabled if the template parameter V is not a pointer
template<typename = std::enable_if<!ValueIsPointer, bool>>
void Register(const K& key, V field)
{
T*& ref = field_map[fname];
field_map[key] = field;
}
/// Register field @a field with name @a key
/** Replace existing field associated with @a key (and optionally
delete associated pointer if @a own_data is true).
Only enabled if the template parameter V is a pointer*/
template<typename = std::enable_if<ValueIsPointer, bool>>
void Register(const K& key, V field, bool own_data)
{
V& ref = field_map[key];
if (own_data)
{
delete ref; // if newly allocated -> ref is null -> OK
@@ -47,23 +61,40 @@ public:
ref = field;
}
/// Unregister association between field @a field and name @a fname
/** Optionally delete associated pointer if @a own_data is true */
void Deregister(const std::string& fname, bool own_data)
/// Unregister association between field @a field and name @a key
/// Only enabled if the template parameter V is not a pointer
template<typename = std::enable_if<!ValueIsPointer, bool>>
void Deregister(const K& key)
{
iterator it = field_map.find(fname);
iterator it = field_map.find(key);
if ( it != field_map.end() )
{
field_map.erase(it);
}
}
/// Unregister association between field @a field and name @a key
/** Optionally delete associated pointer if @a own_data is true.
Only enabled if the template parameter V is a pointer */
template<typename = std::enable_if<ValueIsPointer, bool>>
void Deregister(const K& key, bool own_data)
{
iterator it = field_map.find(key);
if ( it != field_map.end() )
{
if (own_data)
{
delete it->second;
it->second = nullptr;
}
field_map.erase(it);
}
}
/// Clear all associations between names and fields
/** Delete associated pointers when @a own_data is true */
/** Delete associated pointers when @a own_data is true.
Only enabled if the template parameter V is a pointer */
template<typename = std::enable_if<ValueIsPointer, bool>>
void DeleteData(bool own_data)
{
for (iterator it = field_map.begin(); it != field_map.end(); ++it)
@@ -76,22 +107,37 @@ public:
}
}
/// Predicate to check if a field is associated with name @a fname
bool Has(const std::string& fname) const
/// Predicate to check if a field is associated with name @a key
bool Has(const K& key) const
{
return field_map.find(fname) != field_map.end();
return field_map.find(key) != field_map.end();
}
/// Get a pointer to the field associated with name @a fname
/** @return Pointer to field associated with @a fname or NULL */
T* Get(const std::string& fname) const
/// Get a pointer to the field associated with name @a key
/** @return Field associated with @a key or NULL,
if value is pointer and key not found */
V Get(const K& key) const
{
const_iterator it = field_map.find(fname);
return it != field_map.end() ? it->second : NULL;
const_iterator it = field_map.find(key);
if (it != field_map.end())
{
return it->second;
}
else
{
if constexpr (ValueIsPointer)
{
return nullptr;
}
else
{
return V(); // Return default-constructed value for non-pointer types
}
}
}
/// Returns a const reference to the underlying map
const MapType& GetMap() const { return field_map; }
const MapType &GetMap() const { return field_map; }
/// Returns the number of registered fields
int NumFields() const { return field_map.size(); }
@@ -106,21 +152,24 @@ public:
/// Returns an end const iterator to the registered fields
const_iterator end() const { return field_map.end(); }
/// Returns an iterator to the field @a fname
iterator find(const std::string& fname)
{ return field_map.find(fname); }
/// Returns an iterator to the field @a key
iterator find(const K& key)
{ return field_map.find(key); }
/// Returns a const iterator to the field @a fname
const_iterator find(const std::string& fname) const
{ return field_map.find(fname); }
/// Returns a const iterator to the field @a key
const_iterator find(const K& key) const
{ return field_map.find(key); }
/// Clears the map of registered fields without reclaiming memory
/// Clears the map of registered fields
void clear() { field_map.clear(); }
protected:
MapType field_map;
};
/// Lightweight adaptor over an std::map from strings to pointer to T
template<typename T>
using NamedFieldsMap = GenericFieldMap<std::string, T*>;
/** A class for collecting finite element data that is part of the same
simulation. Currently, this class groups together grid functions (fields),
+16 -2
View File
@@ -671,6 +671,20 @@ public:
MPI_COMM_WORLD);
}
Operator& GetGradient(const Vector &x0) const override
{
x = x0;
f.UseDevice(x.UseDevice());
xpev.UseDevice(x.UseDevice());
op.Mult(x, f);
const real_t xnorm_local = x.Norml2();
MPI_Allreduce(&xnorm_local, &xnorm, 1, MPITypeMap<real_t>::mpi_type, MPI_SUM,
MPI_COMM_WORLD);
return const_cast<FDJacobian&>(*this);
}
void Mult(const Vector &v, Vector &y) const override
{
// See [1] for choice of eps.
@@ -725,11 +739,11 @@ public:
private:
const Operator &op;
Vector x, f;
mutable Vector x, f;
mutable Vector xpev;
real_t lambda = 1.0e-6;
real_t fixed_eps;
real_t xnorm;
mutable real_t xnorm;
};
/// @brief Find the index of a field descriptor in a vector of field descriptors.
+1
View File
@@ -22,6 +22,7 @@ add_subdirectory(common)
add_subdirectory(contact)
add_subdirectory(dfem)
add_subdirectory(diag-smoothers)
add_subdirectory(multiapp)
add_subdirectory(dpg)
add_subdirectory(electromagnetics)
add_subdirectory(fluids/navier)
+52
View File
@@ -0,0 +1,52 @@
# 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.
set(MESH_FILES
backward-facing-step.msh
channel-cylinder.msh
)
# Add a target to copy the mesh files from the source directory; used by sample
# runs.
set(SRC_MESH_FILES)
foreach(MESH_FILE ${MESH_FILES})
list(APPEND SRC_MESH_FILES ${CMAKE_CURRENT_SOURCE_DIR}/${MESH_FILE})
endforeach()
add_custom_command(OUTPUT data_is_copied
COMMAND ${CMAKE_COMMAND} -E copy_if_different ${SRC_MESH_FILES} .
COMMAND ${CMAKE_COMMAND} -E touch data_is_copied
COMMENT "Copying multiapp miniapps data files ...")
add_custom_target(copy_miniapps_multiapp_data DEPENDS data_is_copied)
list(APPEND MULTIAPP_COMMON_SOURCES
multiapp.cpp)
list(APPEND MULTIAPP_COMMON_HEADERS
multiapp.hpp)
set(MULTIAPP_COMMON_FILES
EXTRA_SOURCES ${MULTIAPP_COMMON_SOURCES}
EXTRA_HEADERS ${MULTIAPP_COMMON_HEADERS})
# Parallel apps.
if (MFEM_USE_MPI)
add_mfem_miniapp(coupled-diffusion
MAIN coupled-diffusion.cpp
${MFEM_MINIAPPS_COMMON_HEADERS} ${MULTIAPP_COMMON_FILES}
LIBRARIES mfem-common)
add_dependencies(coupled-diffusion copy_miniapps_multiapp_data)
# Add parallel tests.
# if (MFEM_ENABLE_TESTING)
# endif()
endif()
+846
View File
@@ -0,0 +1,846 @@
#include "mfem.hpp"
#include "multiapp.hpp"
#include <fstream>
using namespace mfem;
using namespace std;
struct CaseContext
{
int ser_ref = 1; // Serial mesh refinement
int order = 3; // Finite element order
bool visualization = true;// Visualization on/off
int grad_mode = 1; // Gradient mode for the coupled operator - 0: finite difference,
// 1: back/forward propagation
bool coupled = true; // Coupled (true) vs. uncoupled (false) solves
int nl_iter = 50; // Maximum number of nonlinear iterations
int lin_iter = 2000; // Maximum number of linear iterations
#if defined(MFEM_USE_DOUBLE)
real_t tol_nsolve = 1e-4;
real_t tol_lsolve = 1e-6;
#elif defined(MFEM_USE_SINGLE)
real_t tol_nsolve = 1e-3;
real_t tol_lsolve = 1e-3;
#else
#error "Only single and double precision are supported!"
real_t tol_nsolve = 0;
real_t tol_lsolve = 0;
#endif
} ctx;
void SetSolverParameters(IterativeSolver *solver, real_t rtol, real_t atol , int max_it,
int print_level, bool iterative_mode);
/// A functional diffusion coefficient (i.e., k(T))
class FunctionalCoefficient : public Coefficient
{
public:
enum Mode { FUNC = 0, GRAD = 1};
protected:
ParGridFunction *T_gf = nullptr;
real_t kref = 1.0;
real_t a0 = 0.0, a1 = 0.0, a2 = 0.0;
int findex = 0;
Mode mode = Mode::FUNC; // otherwise, grad
public:
FunctionalCoefficient(ParGridFunction *T_gf, real_t kref):
T_gf(T_gf), kref(kref) { }
FunctionalCoefficient(ParGridFunction *T_gf, real_t kref, real_t a0):
T_gf(T_gf), kref(kref), a0(a0) { findex = 1; }
FunctionalCoefficient(ParGridFunction *T_gf, real_t kref,
real_t a0, real_t a1, real_t a2): T_gf(T_gf),
kref(kref), a0(a0), a1(a1), a2(a2) { findex = 2; }
real_t Exponential(real_t x, bool eval_f) const
{
real_t f = kref*exp(a0*x);
return (eval_f ? f : a0*f);
}
real_t Polynomial(real_t x, bool eval_f) const
{
return (eval_f ? kref*(a0 + a1*x + a2*x*x) : kref*(a1 + 2*a2*x));
}
void SetMode(Mode mode) { this->mode = mode; }
Mode GetMode() const { return mode; }
void UpdateGridFunction(ParGridFunction *gf) { T_gf = gf; }
real_t Eval(real_t x, bool eval_f) const
{
switch (findex)
{
case 1:
return Exponential(x, eval_f);
case 2:
return Polynomial(x, eval_f);
default:
return kref;
}
}
real_t Eval(ElementTransformation &Tr,
const IntegrationPoint &ip) override
{
real_t T = T_gf ? T_gf->GetValue(Tr, ip) : 0.0;
bool eval_f = (mode == Mode::FUNC);
return Eval(T, eval_f);
}
};
/// A coefficient defined by the product of grid functions, e.g. k(T) = prod_i x_i
class GridFunctionProductCoefficient : public Coefficient
{
protected:
std::vector<ParGridFunction*> &x;
public:
GridFunctionProductCoefficient(std::vector<ParGridFunction*> &x) : x(x) { }
real_t Eval(ElementTransformation &Tr, const IntegrationPoint &ip) override
{
real_t prod = 1.0;
for(size_t i = 0; i < x.size(); i++)
{
real_t val = x[i]->GetValue(Tr, ip);
prod *= val;
}
return prod;
}
};
class CoefficientIntegrator : public NonlinearFormIntegrator
{
protected:
FunctionalCoefficient *func = nullptr;
Vector shape;
public:
CoefficientIntegrator(FunctionalCoefficient *func) : func(func) { }
void SetCoefficient(FunctionalCoefficient *f) { func = f; }
void AssembleElementVector(const FiniteElement &el,
ElementTransformation &Tr,
const Vector &elfun, Vector &elvect)
{
int dof = el.GetDof();
shape.SetSize(dof);
elvect.SetSize(dof);
elvect = 0.0;
const IntegrationRule *ir = &el.GetNodes();
for (int i = 0; i < ir->GetNPoints(); i++)
{
const IntegrationPoint &ip = ir->IntPoint(i);
el.CalcShape(ip, shape);
Tr.SetIntPoint(&ip);
real_t x = elfun * shape; // Evaluate the function at the integration point
real_t fval = func->Eval(x, true);
for (int j = 0; j < dof; j++)
{
elvect(j) += fval * shape(j);
}
}
}
void AssembleElementGrad(const FiniteElement &el, ElementTransformation &Tr,
const Vector &elfun, DenseMatrix &elmat)
{
int dof = el.GetDof();
shape.SetSize(dof);
elmat.SetSize(dof);
elmat = 0.0;
const IntegrationRule *ir = &el.GetNodes();
for (int i = 0; i < ir->GetNPoints(); i++)
{
const IntegrationPoint &ip = ir->IntPoint(i);
el.CalcShape(ip, shape);
Tr.SetIntPoint(&ip);
real_t x = elfun * shape; // Evaluate the function at the integration point
real_t dfdx = func->Eval(x, false); // Evaluate the derivative of the function at the integration point
for (int j = 0; j < dof; j++)
{
elmat(j,j) += dfdx * shape(j); // Diagonal contribution to the Jacobian
}
}
}
};
class NonlinearDiffusionIntegrator : public NonlinearFormIntegrator
{
protected:
Coefficient *k;
Coefficient *dk;
Vector u, vec, shape;
DenseMatrix dshape, dshapedxt, adjJ;
public:
NonlinearDiffusionIntegrator(Coefficient *kappa, Coefficient *dkappa) :
k(kappa), dk(dkappa) { }
virtual void AssembleElementVector(const FiniteElement &el,
ElementTransformation &Tr,
const Vector &elfun, Vector &elvec)
{
int dim = el.GetDim();
int dof = el.GetDof();
real_t w;
elvec.SetSize(dof);
elvec = 0.0;
const IntegrationRule *ir = IntRule ? IntRule : &IntRules.Get(el.GetGeomType(), 2*el.GetOrder());
u.SetSize(dim);
vec.SetSize(dim);
dshape.SetSize(dof, dim);
adjJ.SetSize(dim, dim);
for (int i = 0; i < ir->GetNPoints(); i++)
{
const IntegrationPoint &ip = ir->IntPoint(i);
el.CalcDShape(ip, dshape);
Tr.SetIntPoint(&ip);
CalcAdjugate(Tr.Jacobian(), adjJ);
w = ip.weight / Tr.Weight();
dshape.MultTranspose(elfun, u);
adjJ.MultTranspose(u, vec);
if(k)
{
w *= k->Eval(Tr, ip);
}
vec *= w;
adjJ.Mult(vec, u);
dshape.AddMult(u, elvec);
}
}
void AssembleElementGrad(const FiniteElement &el, ElementTransformation &Tr,
const Vector &elfun, DenseMatrix &elmat)
{
int dim = el.GetDim();
int dof = el.GetDof();
real_t w, k0 = 0.0, dk0 = 0.0;
elmat.SetSize(dof);
elmat = 0.0;
const IntegrationRule *ir = IntRule ? IntRule : &IntRules.Get(el.GetGeomType(), 2*el.GetOrder());
u.SetSize(dim);
shape.SetSize(dof);
vec.SetSize(dof);
dshape.SetSize(dof, dim);
dshapedxt.SetSize(dof, dim);
// f = grad(psi) * k(u) * grad(T)
// df/dT = grad(psi) ( k(u0) * grad(T) + k'(u0) * grad(u0) * T )
for (int i = 0; i < ir->GetNPoints(); i++)
{
const IntegrationPoint &ip = ir->IntPoint(i);
el.CalcShape(ip, shape);
el.CalcDShape(ip, dshape);
Tr.SetIntPoint(&ip);
w = ip.weight / Tr.Weight();
Mult(dshape, Tr.AdjugateJacobian(), dshapedxt);
k0 = k ? k->Eval(Tr, ip) : 0.0;
dk0 = dk ? dk->Eval(Tr, ip) : 0.0;
if(k0 != 0.0) // grad(psi) * k(u0) * grad(T)
{
real_t kdT = w*k0;
AddMult_a_AAt(kdT, dshapedxt, elmat);
}
if(dk0 != 0.0) // grad(psi) * (k'(T0) * grad(T0)) * T
{
dk0 = w*dk->Eval(Tr, ip);
dshapedxt.MultTranspose(elfun, u); // grad(T0) in physical space
u *= dk0; // k'(T0) * grad(T0)
dshapedxt.Mult(u, vec); // grad(psi) * k'(T0) * grad(T0)
AddMultVWt(vec, shape, elmat); // grad(psi) * k'(T0) * grad(T0) * T
}
}
}
};
/// An application that takes an input field T, and computes an output field k(T)
// represented by the FunctionalCoefficient class.
class DiffusionCoefficient : public GraphNode
{
public:
using Mode = FunctionalCoefficient::Mode;
protected:
ParFiniteElementSpace &fes;
mutable ParGridFunction T, k;
mutable FunctionalCoefficient *kc;
// mutable Vector tdof, kdof, dk_dof, dT_dof;
mutable Mode mode = Mode::FUNC;
mutable ParNonlinearForm Nform;
mutable Operator *J = nullptr; // Jacobian for the nonlinear form
CoefficientIntegrator *coeff_integrator = nullptr;
public:
DiffusionCoefficient(ParFiniteElementSpace &fes) :
GraphNode(fes.GetTrueVSize()), fes(fes), T(&fes), k(&fes),
kc(new FunctionalCoefficient(&T, 1.0, 5.0e-2)),
Nform(&fes),
coeff_integrator(new CoefficientIntegrator(kc))
{
k = 0.0;
T = 0.0;
k.ProjectCoefficient(*kc);
// Testing with the nonlinear form framework to compute k(T) and dk/dT
Nform.AddDomainIntegrator(coeff_integrator); // Transfer ownership
Nform.SetGradientType(Operator::Type::Hypre_ParCSR);
Nform.Setup();
SetInputOffsets(Array<int>({0, fes.GetTrueVSize()}));
SetOutputOffsets(Array<int>({0, fes.GetTrueVSize()}));
}
void SetMode(Mode mode) { this->mode = mode; }
FunctionalCoefficient* GetCoefficient() { return kc; }
void SetCoefficient(FunctionalCoefficient *fc)
{
if(kc) delete kc;
kc = fc;
kc->SetMode(mode);
kc->UpdateGridFunction(&T);
coeff_integrator->SetCoefficient(kc);
}
void Mult(const Vector &x, Vector &y) const override
{
BlockVector xb(x.GetData(), InputOffsets());
BlockVector yb(y.GetData(), OutputOffsets());
MultiVector xmv(1), ymv(1);
xmv.MakeRef(0, xb.GetBlock(0));
ymv.MakeRef(0, yb.GetBlock(0));
MultMV(xmv, ymv);
}
void MultMV(const MultiVector &x, MultiVector &y) const override
{
const Vector &tdof = x[0];
Vector &kdof = y[0];
Nform.Mult(tdof, kdof);
if(exec_mode == GraphNode::GRADIENT_MODE)
{
J = &Nform.GetGradient(tdof); // Store jacobian for JVP
}
else
{
J = nullptr; // Clear the Jacobian if not in gradient mode
}
}
void GradientMult(const MultiVector &x, const MultiVector &dx, MultiVector &dy) const override
{
const Vector &tdof = x[0];
const Vector &xadj = dx[0];
Vector &yadj = dy[0];
if(J)
{
J->Mult(xadj, yadj);
}
else
{
J = &Nform.GetGradient(tdof); // Store jacobian for JVP
J->Mult(xadj, yadj);
}
}
~DiffusionCoefficient() override
{
if(kc) delete kc;
}
};
/// An application that takes n input fields x_i, and computes an output
/// field prod(x) := y = prod_i x_i.
/// Also provides the derivative dy/dx_i = prod_{j!=i} x_j * dx_i/dx for i = 0,...,n-1.
class ProductGridFunctions : public GraphNode
{
protected:
ParFiniteElementSpace &fes;
mutable std::vector<ParGridFunction*> x_gf;
mutable Vector dfdx;
mutable ParGridFunction y_gf;
mutable GridFunctionProductCoefficient prod_coeff;
public:
ProductGridFunctions(ParFiniteElementSpace &fes, int n) :
// GraphNode(fes.GetTrueVSize()),
GraphNode(fes.GetTrueVSize(), fes.GetTrueVSize() * n),
fes(fes), x_gf(n),
y_gf(&fes), prod_coeff(x_gf)
{
Array<int> offsets(n+1);
offsets[0] = 0;
for (int i = 0; i < n; i++)
{
x_gf[i] = new ParGridFunction(&fes);
*x_gf[i] = 0.0;
offsets[i+1] = offsets[i] + fes.GetTrueVSize();
}
y_gf = 0.0;
y_gf.ProjectCoefficient(prod_coeff);
SetInputOffsets(offsets);
SetOutputOffsets(Array<int>({0, fes.GetTrueVSize()}));
}
void Mult(const Vector &x, Vector &y) const override
{
BlockVector xb(x.GetData(), InputOffsets());
BlockVector yb(y.GetData(), OutputOffsets());
MultiVector xmv(x_gf.size()), ymv(1);
for (size_t i = 0; i < x_gf.size(); i++)
{
xmv.MakeRef(i, xb.GetBlock(i));
}
ymv.MakeRef(0, yb.GetBlock(0));
MultMV(xmv, ymv);
}
void MultMV(const MultiVector &x, MultiVector &y) const override
{
for (size_t i = 0; i < x_gf.size(); i++)
{
const Vector &x_dof = x[i];
x_gf[i]->SetFromTrueDofs(x_dof);
}
Field *out_field = OutputField(0);
Vector &y_dof = y[0];
y_gf.ProjectCoefficient(prod_coeff);
y_gf.GetTrueDofs(y_dof);
}
void GradientMult(const MultiVector &x, const MultiVector &dx, MultiVector &dy) const override
{
// Jacobian vector product for y = prod_i x_i is:
// dy/dx = sum_i (prod_{j!=i} x_j * dx_i/dx)
for (size_t i = 0; i < x_gf.size(); i++)
{
const Vector &x_dof = x[i];
x_gf[i]->SetFromTrueDofs(x_dof); // Set all x_i
}
Vector &jvp = dy[0];
jvp = 0.0;
for (size_t i = 0; i < x_gf.size(); i++)
{
const Vector &x_dof = x[i];
const Vector &dx_dof = dx[i]; // Get dx_i/dx
x_gf[i]->SetFromTrueDofs(dx_dof); // Set x_i = dx_i/dx for i-th term in the product
y_gf.ProjectCoefficient(prod_coeff); // Recompute product with x_i replaced by dx_i/dx
y_gf.GetTrueDofs(dfdx); // Get prod_{j!=i} x_j * dx_i/dx for i-th term
jvp += dfdx; // Accumulate contribution from i-th term
x_gf[i]->SetFromTrueDofs(x_dof); // reset to original value for next iteration
}
}
~ProductGridFunctions() override
{
for (size_t i = 0; i < x_gf.size(); i++)
{
if(x_gf[i]) delete x_gf[i];
}
}
};
/// An application that represents the nonlinear diffusion operator: f(T) = -Div(k(u) grad(T))
/// with input field T and k, and output field f(T).
class DiffusionOperator : public GraphNode
{
public:
// Mesh and finite element space
ParMesh &mesh;
ParFiniteElementSpace &fes;
/// Essential dof array.
Array<int> ess_tdofs;
/// Grid functions for the temperature and heat flux
mutable ParGridFunction T, k, dk;
mutable GridFunctionCoefficient k_gfc, dk_gfc;
mutable ParNonlinearForm Nform;
mutable ParLinearForm bform;
mutable Vector b;
ConstantCoefficient zero_coeff, one_coeff;
mutable FunctionalCoefficient *kc = nullptr;
mutable HypreParMatrix *dfdk_mat = nullptr, *dfdT_mat = nullptr;
public:
DiffusionOperator(ParFiniteElementSpace &fes_) :
// GraphNode(fes_.GetTrueVSize()),
GraphNode(fes_.GetTrueVSize(),2*fes_.GetTrueVSize()),
mesh(*fes_.GetParMesh()), fes(fes_),
T(&fes), k(&fes), dk(&fes),
k_gfc(&k), dk_gfc(&dk),
Nform(&fes), bform(&fes),
zero_coeff(0.0), one_coeff(1.0)
{
fes.GetBoundaryTrueDofs(ess_tdofs);
T = 0.0;
k = 0.0;
dk = 0.0;
bform.AddDomainIntegrator(new DomainLFIntegrator(one_coeff));
Nform.AddDomainIntegrator(new NonlinearDiffusionIntegrator(&k_gfc, &dk_gfc));
Nform.SetGradientType(Operator::Type::Hypre_ParCSR);
b.SetSize(fes.GetTrueVSize()); b = 0.0;
Assemble();
SetInputOffsets(Array<int>({0, fes.GetTrueVSize(), 2*fes.GetTrueVSize()}));
SetOutputOffsets(Array<int>({0, fes.GetTrueVSize()}));
}
void SetCoefficient(FunctionalCoefficient *fc) { kc = fc; }
void Assemble()
{
AssembleLinearForms();
AssembleBilinearForms();
AssembleNonlinearForms();
}
void AssembleBilinearForms()
{}
void AssembleNonlinearForms()
{
Nform.SetEssentialTrueDofs(ess_tdofs);
Nform.Setup();
}
void AssembleLinearForms()
{
bform.Assemble();
bform.ParallelAssemble(b);
}
void Mult(const Vector &x, Vector &y) const override
{
BlockVector xb(x.GetData(), InputOffsets());
BlockVector yb(y.GetData(), OutputOffsets());
MultiVector xmv(2), ymv(1);
xmv.MakeRef(0, xb.GetBlock(0));
xmv.MakeRef(1, xb.GetBlock(1));
ymv.MakeRef(0, yb.GetBlock(0));
MultMV(xmv, ymv);
}
void MultMV(const MultiVector &x, MultiVector &y) const override
{
const Vector &tdofs = x[0];
const Vector &kdofs = x[1];
Vector &fdofs = y[0];
k.SetFromTrueDofs(kdofs); // update for use in k_gfc
if(exec_mode == GraphNode::GRADIENT_MODE)
{
if(dfdT_mat) delete dfdT_mat;
if(dfdk_mat) delete dfdk_mat;
dk = 0.0;
k.SetFromTrueDofs(kdofs);
Operator* grad = &Nform.GetGradient(tdofs);
dfdT_mat = new HypreParMatrix(dynamic_cast<const HypreParMatrix&>(*grad)); // deep copy
dk = 1.0;
k = 0.0;
grad = &Nform.GetGradient(tdofs);
dfdk_mat = new HypreParMatrix(dynamic_cast<const HypreParMatrix&>(*grad)); // deep copy
}
else
{
if(dfdT_mat) { delete dfdT_mat; dfdT_mat = nullptr; }
if(dfdk_mat) { delete dfdk_mat; dfdk_mat = nullptr; }
}
Nform.Mult(tdofs, fdofs);
fdofs.SetSubVector(ess_tdofs, 0.0);
}
// Exact block jacobian [df/dT, df/dk]
Operator& GetGradient(const Vector &x) const override
{
MFEM_ABORT("GetGradient not implemented for DiffusionOperator");
}
void GradientMult(const MultiVector &x, const MultiVector &dx, MultiVector &dy) const override
{
const Vector &Tadj = dx[0];
const Vector &kadj = dx[1];
Vector &yadj = dy[0];
const Vector &tdofs = x[0];
const Vector &kdofs = x[1];
dfdT_mat->Mult(Tadj, yadj);
dfdk_mat->AddMult(kadj, yadj);
}
/// @brief Destroy the DiffusionOperator object
~DiffusionOperator() override
{
if(dfdT_mat) delete dfdT_mat;
if(dfdk_mat) delete dfdk_mat;
}
};
int main(int argc, char *argv[])
{
Mpi::Init();
Hypre::Init();
using GradMode = DAGraph::GradMode;
OptionsParser args(argc, argv);
args.AddOption(&ctx.order, "-o", "--order",
"Finite element order (polynomial degree).");
args.AddOption(&ctx.visualization, "-vis", "--visualization", "-no-vis",
"--no-visualization",
"Enable or disable GLVis visualization.");
args.AddOption(&ctx.ser_ref, "-rs", "--serial-refine",
"Number of times to refine the mesh in serial.");
args.AddOption(&ctx.grad_mode, "-gm", "--grad-mode",
"Gradient mode for the coupled operator (0: exact, 1: finite difference, 2: algorithmic differentiation)");
args.AddOption(&ctx.coupled, "-cp", "--coupled", "-ucp", "--uncoupled",
"Coupled (true) vs. uncoupled (false) solves.");
args.ParseCheck();
int order = ctx.order;
std::string mesh_file = "../../data/star.mesh";
Mesh *serial_mesh = new Mesh(mesh_file);
int dim = serial_mesh->Dimension();
for (int i = 0; i < ctx.ser_ref; ++i) { serial_mesh->UniformRefinement(); }
serial_mesh->SetCurvature(order, false, dim, Ordering::byNODES);
ParMesh pmesh(MPI_COMM_WORLD, *serial_mesh);
delete serial_mesh;
pmesh.UniformRefinement();
// Finite element spaces
H1_FECollection fec(order, dim);
ParFiniteElementSpace fes(&pmesh, &fec);
// Build all operator nodes
DiffusionCoefficient diff_coeff_1(fes);
diff_coeff_1.SetName("k(T1)");
diff_coeff_1.SetCoefficient(new FunctionalCoefficient(nullptr, 1.0, 3.5e-2));
// diff_coeff_1.SetCoefficient(new FunctionalCoefficient(nullptr, 1.0, 1.0, 0.1, 0.0));
DiffusionCoefficient diff_coeff_2(fes);
diff_coeff_2.SetName("k(T2)");
diff_coeff_2.SetCoefficient(new FunctionalCoefficient(nullptr, 1.0, 1.0, 2.0, 0.0));
// diff_coeff_2.SetCoefficient(new FunctionalCoefficient(nullptr, 1.5, 2.5e-2));
ProductGridFunctions prod_coeff(fes, 2);
prod_coeff.SetName("k(T1,T2)");
DiffusionOperator diff_op1(fes);
diff_op1.SetName("Div(k(T1,T2) grad(T1))");
diff_op1.SetCoefficient(diff_coeff_1.GetCoefficient());
DiffusionOperator diff_op2(fes);
diff_op2.SetName("Div(k(T1,T2) grad(T2))");
diff_op2.SetCoefficient(diff_coeff_2.GetCoefficient());
// Build the DAG in any order, and then sort it to ensure the correct execution order
DAGraph dag(5);
dag.AddOperator(&diff_coeff_1);
dag.AddOperator(&diff_op1, fes.GetTrueVSize());
dag.AddOperator(&diff_op2, fes.GetTrueVSize());
dag.AddOperator(&diff_coeff_2);
dag.AddOperator(&prod_coeff);
Vector k1vec(fes.GetTrueVSize()); k1vec = 0.0;
Vector k2vec(fes.GetTrueVSize()); k2vec = 0.0;
Vector kpvec(fes.GetTrueVSize()); kpvec = 0.0;
Vector k1adj(fes.GetTrueVSize()); k1adj = 0.0;
Vector k2adj(fes.GetTrueVSize()); k2adj = 0.0;
// Vector kpadj(fes.GetTrueVSize()); kpadj = 0.0;
// Input fields get data from 'x' in DAGraph::Mult(x, y)
Field T1_field(nullptr, nullptr);
Field T2_field(nullptr, nullptr);
// Write space for data and adjoint only needed
// for the intermediate fields k1, k2, and k_prod
Field k1_field(&k1vec, &k1adj);
Field k2_field(&k2vec, &k2adj);
Field kp_field(&kpvec, &kpvec); // can use same space for data & adjoint
// Output fields get data from 'y' in DAGraph::Mult(x, y)
Field f1_field(nullptr, nullptr);
Field f2_field(nullptr, nullptr);
// Add input and output to the DAG
int sz = fes.GetTrueVSize();
dag.AddInput(&T1_field, sz);
dag.AddInput(&T2_field, sz);
dag.AddOutput(&f1_field, sz);
dag.AddOutput(&f2_field, sz);
// Form connections between the nodes in the DAG
diff_coeff_1.AddInput(&T1_field);
diff_coeff_1.AddOutput(&k1_field);
diff_coeff_2.AddInput(&T2_field);
diff_coeff_2.AddOutput(&k2_field);
prod_coeff.AddInputs(&k1_field, &k2_field);
prod_coeff.AddOutput(&kp_field);
diff_op1.AddInput(&T1_field);
diff_op1.AddOutput(&f1_field);
diff_op2.AddInput(&T2_field);
diff_op2.AddOutput(&f2_field);
if(ctx.coupled)
{
diff_op1.AddInput(&kp_field); // kp_field
diff_op2.AddInput(prod_coeff.OutputField(0)); // Can also use kp_field directly
}
else
{
diff_op1.AddInput(&k1_field); // Can also use diff_coeff_1.OutputField(0)
diff_op2.AddInput(&k2_field); // Can also use diff_coeff_2.OutputField(0)
}
// Assemble DAG: topological sort, validate nodes, etc.
dag.Assemble();
std::string output_prefix = ctx.coupled ? "Coupled_Diffusion" : "Uncoupled_Diffusion";
if(Mpi::Root())
{
std::ofstream fout(output_prefix+"-dag.txt");
fout << "{\n";
dag.Save(fout);
fout << "}\n";
fout << std::flush;
fout.close();
}
// Set initial guess and boundary conditions for T1 and T2
Array<int> ess_tdofs;
fes.GetBoundaryTrueDofs(ess_tdofs);
int T1_idx = 0;
int T2_idx = 1;
BlockVector xb(dag.InputOffsets());
BlockVector yb(dag.OutputOffsets());
xb.GetBlock(T1_idx).Randomize();
xb.GetBlock(T2_idx).Randomize();
xb.GetBlock(T1_idx).SetSubVector(ess_tdofs, 0.0);
xb.GetBlock(T2_idx).SetSubVector(ess_tdofs, 0.0);
// Build the nonlinear solver and linear solver for the DAG
NewtonSolver newton_solver(pmesh.GetComm());
GMRESSolver linear_solver(pmesh.GetComm());
linear_solver.SetKDim(500);
SetSolverParameters(&newton_solver, ctx.tol_nsolve, 0.0, ctx.nl_iter, 1, true);
SetSolverParameters(&linear_solver, ctx.tol_lsolve, 0.0, ctx.lin_iter, 1, false);
newton_solver.SetPreconditioner(linear_solver);
linear_solver.SetPrintLevel(1);
// Set the gradient mode for the DAG and solve the coupled system
GradMode gm = static_cast<GradMode>(ctx.grad_mode);
dag.SetGradientMode(gm);
newton_solver.SetOperator(dag);
newton_solver.Mult(xb, yb);
ParaViewDataCollection *pv = nullptr;
if (ctx.visualization)
{
std::string pv_prefix;
switch (ctx.grad_mode)
{
case 0: pv_prefix = "FD"; break;
case 1: pv_prefix = "MF"; break;
default: pv_prefix = "Unknown"; break;
}
pv = new ParaViewDataCollection(output_prefix+"-"+pv_prefix, &pmesh);
pv->SetLevelsOfDetail(order);
pv->SetDataFormat(VTKFormat::BINARY);
pv->SetHighOrderOutput(true);
ParGridFunction T1_gf(&fes);
ParGridFunction T2_gf(&fes);
T1_gf.SetFromTrueDofs(yb.GetBlock(T1_idx));
T2_gf.SetFromTrueDofs(yb.GetBlock(T2_idx));
pv->RegisterField("T1", &T1_gf);
pv->RegisterField("T2", &T2_gf);
pv->Save();
delete pv;
}
std::cout << "Finished solving the coupled diffusion problem." << std::endl;
return 0;
}
void SetSolverParameters(IterativeSolver *solver, real_t rtol, real_t atol,
int max_it, int print_level, bool iterative_mode)
{
solver->SetRelTol(rtol);
solver->SetAbsTol(atol);
solver->SetMaxIter(max_it);
solver->SetPrintLevel(print_level);
solver->iterative_mode = iterative_mode;
}
+706
View File
@@ -0,0 +1,706 @@
// 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 "multiapp.hpp"
namespace mfem
{
DAGraph::~DAGraph()
{
for(int i=0; i < nnodes; i++)
{
if(node_owned[i] && nodes[i]) delete nodes[i];
}
if(grad) delete grad;
}
void DAGraph::Assemble()
{
// Sort graph nodes topologically to ensure correct execution order
// Ordering is not unique, hence, id->index maps are needed
TopologicalSort();
// Collect all fields from the nodes into the field map
CollectFieldMaps();
// Compute depth of the graph nodes
ComputeDepth();
// Validate each node
for (auto &node : nodes)
{
ValidateNode(*node);
}
// Update width and height of the DAG from offsets
// Check that the input and output offsets are consistent
ValidateOffsets();
width = input_offsets.Last();
height = output_offsets.Last();
// Delete any existing gradient operator as node ordering may have changed
if (grad) delete grad;
assembled = true;
}
void DAGraph::ValidateOffsets()
{
// Check that the input and output offsets are consistent
// with the number of inputs and outputs
if(InputFields().Size() > 1)
{
MFEM_ASSERT(input_offsets.Size() == InputFields().Size() + 1,
"Input offsets size inconsistent with number of input fields");
}
else
{
input_offsets = Array<int>({0, nodes[0]->Width()});
}
if(OutputFields().Size() > 1)
{
MFEM_ASSERT(output_offsets.Size() == OutputFields().Size() + 1,
"Output offsets size inconsistent with number of output fields");
}
else
{
output_offsets = Array<int>({0, nodes.Last()->Height()});
}
}
void DAGraph::ValidateNode(GraphNode &node)
{
// Validate that the node's input and output fields are consistent with the graph's field map
auto inputs = node.InputFields();
auto outputs = node.OutputFields();
// Check offsets match width and height of the node
MFEM_ASSERT(node.InputOffsets().Last() == node.Width(),
"Node ID: " << node.ID() << " input offsets do not match node width.");
MFEM_ASSERT(node.OutputOffsets().Last() == node.Height(),
"Node ID: " << node.ID() << " output offsets do not match node height.");
// Check number of input and output fields match the offsets
MFEM_ASSERT(node.InputOffsets().Size() == inputs.Size() + 1,
"Node input offsets size inconsistent with number of input fields");
MFEM_ASSERT(node.OutputOffsets().Size() == outputs.Size() + 1,
"Node output offsets size inconsistent with number of output fields");
// Check that all input and output fields are registered in the graph's field map
for(auto input_field : inputs)
{
MFEM_ASSERT(fid_to_index.Has(input_field->ID()),
"Input field ID " << input_field->ID() << " not found in graph's field map");
}
for(auto output_field : outputs)
{
MFEM_ASSERT(fid_to_index.Has(output_field->ID()),
"Output field ID " << output_field->ID() << " not found in graph's field map");
}
}
void DAGraph::TopologicalSort()
{
Array<int> sorted_indices;
sorted_indices.Reserve(nnodes);
Array<bool> visited(nnodes);
visited = false; // Initialize all nodes as unvisited
// Perform a depth-first search to sort the nodes topologically
std::function<void(int)> DepthFirstSearch = [&](int node_index)
{
if(visited[node_index]) return;
visited[node_index] = true;
auto node = nodes[node_index];
// Visit all nodes that this node depends on
for(auto input_field : node->InputFields())
{
for(int j=0; j < nnodes; j++)
{
auto other_node = nodes[j];
if(other_node == node) continue;
for(auto output_field : other_node->OutputFields())
{
if(input_field->ID() == output_field->ID()) // Compare by unique ID
{
DepthFirstSearch(j);
}
}
}
}
sorted_indices.push_back(node_index);
};
for(int i=0; i < nnodes; i++)
{
DepthFirstSearch(i);
}
nodes.Permute(sorted_indices);
node_owned.Permute(sorted_indices);
// Update the node indices after sorting
for(int i=0; i < nnodes; i++)
{
nodes[i]->SetNodeIndex(i);
}
sorted = true;
}
void DAGraph::ComputeDepth()
{
// Compute depth of ordered nodes
node_depth.SetSize(nnodes);
node_depth = 0;
for(int i=0; i < nnodes; i++)
{
int max_depth = 0;
auto node = nodes[i];
for(auto input_field : node->InputFields())
{
for(int j=0; j < i; j++)
{
auto other_node = nodes[j];
if(other_node == node) continue;
for(auto output_field : other_node->OutputFields())
{
if(input_field->ID() == output_field->ID()) // Compare by unique ID
{
max_depth = std::max(max_depth, node_depth[j] + 1);
}
}
}
}
node_depth[i] = max_depth;
}
}
void DAGraph::CollectFieldMaps()
{
MFEM_ASSERT(sorted, "DAGraph must be topologically sorted before collecting fields");
fid_to_index.clear();
fid_to_field.clear();
int nfields = 0;
for (auto f : InputFields())
{
fid_to_index.Register(f->ID(), nfields++);
fid_to_field.Register(f->ID(), f);
}
for (auto &node : nodes)
{
for (auto f : node->OutputFields())
{
if (!fid_to_index.Has(f->ID()))
{
fid_to_index.Register(f->ID(), nfields++);
}
if (!fid_to_field.Has(f->ID()))
{
fid_to_field.Register(f->ID(), f);
}
}
}
// TODO: Possibly add all intermediate fields from nodes to the graph's FieldCollection
}
void DAGraph::Mult(const Vector &x, Vector &y) const
{
MFEM_ASSERT(width == x.Size(), "Input vector size (" << x.Size()
<< ") must match matrix width (" << width << ")");
MFEM_ASSERT(height == y.Size(), "Output vector size (" << y.Size()
<< ") must match matrix height (" << height << ")");
auto inputs = InputFields();
auto outputs = OutputFields();
BlockVector xb(x.GetData(), input_offsets);
BlockVector yb(y.GetData(), output_offsets);
MultiVector xmv(inputs.Size()), ymv(outputs.Size());
// Set the data pointers of the input and output fields
// of the graph to point to the corresponding blocks of
// the input and output vectors
for(int i=0; i < inputs.Size(); i++)
{
xmv.MakeRef(i, xb.GetBlock(i));
}
for(int i=0; i < outputs.Size(); i++)
{
ymv.MakeRef(i, yb.GetBlock(i));
}
MultMV(xmv, ymv);
}
void DAGraph::MultMV(const MultiVector &x, MultiVector &y) const
{
auto inputs = InputFields();
auto outputs = OutputFields();
MFEM_ASSERT(inputs.Size() == x.NumBlocks(), "Number of input blocks (" << x.NumBlocks()
<< ") must match number of input fields (" << inputs.Size() << ")");
MFEM_ASSERT(outputs.Size() == y.NumBlocks(), "Number of output blocks (" << y.NumBlocks()
<< ") must match number of output fields (" << outputs.Size() << ")");
for(int i=0; i < inputs.Size(); i++)
{
inputs[i]->SetData(const_cast<Vector*>(&x[i]));
}
for (int i=0; i < outputs.Size(); i++)
{
outputs[i]->SetData(&y[i]);
}
auto index_map = GetFieldIdToIndexMap();
auto fld_map = GetFieldIdToFieldMap();
int nfields = index_map.NumFields();
MultiVector ymv(nfields); // TODO: Should this be a member function?
// Assemble the multivector from the individual fields based on their IDs
// This multivector contains all input, output, and intermediate fields in the graph
for (auto const& [id, idx] : index_map)
{
if (fld_map.Has(id))
{
auto field = fld_map.Get(id);
ymv.MakeRef(idx, *field->Data());
}
else
{
MFEM_ABORT("Field ID " << id << " not found in field map");
}
}
Execute(x, ymv);
for(auto &f : inputs)
{
f->SetData(nullptr);
}
for(auto &f : outputs)
{
f->SetData(nullptr);
}
}
void DAGraph::Execute(const MultiVector &x, MultiVector &y) const
{
MFEM_ASSERT(assembled, "DAGraph must be assembled before calling Execute()");
MFEM_ASSERT(x.NumBlocks() == InputFields().Size(),
"Number of input blocks (" << x.NumBlocks()
<< ") must match number of input fields (" << InputFields().Size() << ")");
auto index_map = GetFieldIdToIndexMap();
MFEM_ASSERT(y.NumBlocks() == index_map.NumFields(),
"Number of output blocks (" << y.NumBlocks()
<< ") must match number of fields (" << index_map.NumFields() << ")");
auto inputs = InputFields();
for(int i=0; i < inputs.Size(); i++)
{
int idx = index_map.Get(inputs[i]->ID());
if(&y[idx] != &x[i]) // copy data, if address is different
{
y[idx] = x[i];
}
}
for (auto node : nodes)
{
auto node_inputs = node->InputFields();
auto node_outputs = node->OutputFields();
xmv_node.SetNumBlocks(node_inputs.Size());
ymv_node.SetNumBlocks(node_outputs.Size());
for (int i=0; i < node_inputs.Size(); i++)
{
int idx = index_map.Get(node_inputs[i]->ID());
xmv_node.MakeRef(i, y[idx]);
}
for (int i=0; i < node_outputs.Size(); i++)
{
int idx = index_map.Get(node_outputs[i]->ID());
ymv_node.MakeRef(i, y[idx]);
}
node->MultMV(xmv_node, ymv_node);
}
}
Operator& DAGraph::GetGradient(const Vector &x) const
{
// TODO: Should/could be removed
if(grad_mode == GradMode::FINITE_DIFF)
{
if(!grad)
{
grad = new future::FDJacobian(*this, x, 1e-6);
}
else
{
grad->GetGradient(x); // Update the FDJacobian with new point x
}
return *grad;
}
MFEM_ASSERT(static_cast<int>(grad_mode) < static_cast<int>(GradMode::NONE),
"DAGraph::GetGradient() called with invalid grad_mode: "
<< static_cast<int>(grad_mode));
if(!grad)
{
grad = new GraphGradient(const_cast<DAGraph&>(*this));
}
if(grad_mode == GradMode::ASSEMBLED)
{
return grad->GetGradient(x); // Assemble the Jacobian matrix
}
else // GradMode::MATRIX_FREE
{
dynamic_cast<GraphGradient*>(grad)->Update(x); // Update the GraphGradient with new point x
}
return *grad;
}
GraphGradient::GraphGradient(DAGraph &dag) : Operator(dag.Height(), dag.Width()),
graph(&dag)
{
MFEM_ASSERT(graph->IsAssembled(), "GraphGradient requires an assembled DAGraph.");
MFEM_ASSERT(graph->IsSorted(), "GraphGradient requires a topologically sorted DAGraph.");
auto index_map = graph->GetFieldIdToIndexMap();
auto field_map = graph->GetFieldIdToFieldMap();
MFEM_ASSERT(index_map.NumFields() == field_map.NumFields(),
"Mismatch in number of fields between index_map and field_map");
x_work.DeleteAll(); // Clear any existing pointers
x_work.SetSize(index_map.NumFields());
x_work = nullptr; // Initialize all pointers to nullptr
xlin.SetNumBlocks(index_map.NumFields());
for (auto const& [id, idx] : index_map)
{
MFEM_ASSERT(idx >= 0 && idx < x_work.Size(), "Index out of bounds for field ID: " << id);
MFEM_ASSERT(field_map.Has(id), "Field ID not found in field_map: " << id);
if(x_work[idx] == nullptr)
{
x_work[idx] = new Vector(); // Allocate a new Vector for this field
}
xlin.MakeRef(idx, *x_work[idx]); // Make xlin refer to the allocated Vector
}
}
void GraphGradient::Update(const Vector &x)
{
MFEM_ASSERT(graph != nullptr, "GraphGradient operator requires a non-null DAGraph pointer.");
auto set_exec_mode = [&](DAGraph::ExecutionMode mode)
{
for (auto &node : graph->Nodes())
{
node->SetExecutionMode(mode);
}
};
auto inputs = graph->InputFields();
BlockVector xb(x.GetData(), graph->InputOffsets());
MultiVector xmv(inputs.Size());
for(int i=0; i < inputs.Size(); i++)
{
xmv.MakeRef(i, xb.GetBlock(i));
}
set_exec_mode(DAGraph::ExecutionMode::GRADIENT_MODE);
graph->Execute(xmv, xlin); // Forward pass to populate fields for gradient computations
set_exec_mode(DAGraph::ExecutionMode::DEFAULT_MODE); // Reset execution mode for forward pass
}
void GraphGradient::Mult(const Vector &x, Vector &y) const
{
MFEM_ASSERT(x.Size() == graph->Width(), "Input vector size (" << x.Size()
<< ") must match graph width (" << graph->Width() << ")");
MFEM_ASSERT(y.Size() == graph->Height(), "Output vector size (" << y.Size()
<< ") must match graph height (" << graph->Height() << ")");
auto in_offsets = graph->InputOffsets();
auto out_offsets = graph->OutputOffsets();
auto inputs = graph->InputFields();
auto outputs = graph->OutputFields();
BlockVector xb(x.GetData(), in_offsets);
BlockVector yb(y.GetData(), out_offsets);
MultiVector xmv(inputs.Size()), ymv(outputs.Size());
for(int i=0; i < inputs.Size(); i++)
{
xmv.MakeRef(i, xb.GetBlock(i));
}
for(int i=0; i < outputs.Size(); i++)
{
ymv.MakeRef(i, yb.GetBlock(i));
}
MultMV(xmv, ymv); // Forward mode: compute JVP, y = J(z) * x
}
void GraphGradient::MultMV(const MultiVector &x, MultiVector &y) const
{
auto inputs = graph->InputFields();
auto outputs = graph->OutputFields();
MFEM_ASSERT(inputs.Size() == x.NumBlocks(), "Number of input blocks (" << x.NumBlocks()
<< ") must match number of input fields (" << inputs.Size() << ")");
MFEM_ASSERT(outputs.Size() == y.NumBlocks(), "Number of output blocks (" << y.NumBlocks()
<< ") must match number of output fields (" << outputs.Size() << ")");
for(int i=0; i < inputs.Size(); i++)
{
inputs[i]->SetAdjoint(const_cast<Vector*>(&x[i]));
}
for (int i=0; i < outputs.Size(); i++)
{
outputs[i]->SetAdjoint(&y[i]);
}
auto index_map = graph->GetFieldIdToIndexMap();
auto fld_map = graph->GetFieldIdToFieldMap();
int nfields = index_map.NumFields();
MultiVector ymv(nfields); // TODO: Should this be a class member?
// Assemble the multivector from the individual fields based on their IDs
// This multivector contains all input, output, and intermediate fields in the graph
for (auto const& [id, idx] : index_map)
{
if (fld_map.Has(id))
{
auto field = fld_map.Get(id);
ymv.MakeRef(idx, *field->Adjoint());
}
else
{
MFEM_ABORT("Field ID " << id << " not found in field map");
}
}
Forward(x, ymv); // Forward mode: compute JVP, y = J(z) * x
for (auto &f : inputs)
{
f->SetAdjoint(nullptr);
}
for (auto &f : outputs)
{
f->SetAdjoint(nullptr);
}
}
void GraphGradient::MultTranspose(const Vector &x, Vector &y) const
{
MFEM_ASSERT(x.Size() == graph->Height(), "Input vector size (" << x.Size()
<< ") must match graph height (" << graph->Height() << ")");
MFEM_ASSERT(y.Size() == graph->Width(), "Output vector size (" << y.Size()
<< ") must match graph width (" << graph->Width() << ")");
auto in_offsets = graph->InputOffsets();
auto out_offsets = graph->OutputOffsets();
auto inputs = graph->InputFields();
auto outputs = graph->OutputFields();
BlockVector xb(x.GetData(), out_offsets);
BlockVector yb(y.GetData(), in_offsets);
MultiVector xmv(outputs.Size()), ymv(inputs.Size());
for(int i=0; i < inputs.Size(); i++)
{
xmv.MakeRef(i, xb.GetBlock(i));
}
for(int i=0; i < outputs.Size(); i++)
{
ymv.MakeRef(i, yb.GetBlock(i));
}
MultTransposeMV(xmv, ymv); // Reverse mode: compute VJP, y = J(z)^T * x
}
void GraphGradient::MultTransposeMV(const MultiVector &x, MultiVector &y) const
{
auto inputs = graph->InputFields();
auto outputs = graph->OutputFields();
MFEM_ASSERT(outputs.Size() == x.NumBlocks(), "Number of input blocks (" << x.NumBlocks()
<< ") must match number of output fields (" << outputs.Size() << ")");
MFEM_ASSERT(inputs.Size() == y.NumBlocks(), "Number of output blocks (" << y.NumBlocks()
<< ") must match number of input fields (" << inputs.Size() << ")");
for(int i=0; i < outputs.Size(); i++)
{
outputs[i]->SetAdjoint(const_cast<Vector*>(&x[i]));
}
for (int i=0; i < inputs.Size(); i++)
{
inputs[i]->SetAdjoint(&y[i]);
}
auto index_map = graph->GetFieldIdToIndexMap();
auto fld_map = graph->GetFieldIdToFieldMap();
int nfields = index_map.NumFields();
MultiVector ymv(nfields); // TODO: Should this be a class member?
for(auto const& [id, idx] : index_map)
{
if (fld_map.Has(id))
{
auto field = fld_map.Get(id);
ymv.MakeRef(idx, *field->Adjoint());
}
else
{
MFEM_ABORT("Field ID " << id << " not found in field map");
}
}
Reverse(x, ymv); // Reverse mode: compute VJP, y = J(z)^T * x
for (auto &f : outputs)
{
f->SetAdjoint(nullptr);
}
for (auto &f : inputs)
{
f->SetAdjoint(nullptr);
}
}
void GraphGradient::Forward(const MultiVector &x, MultiVector &y) const
{
MFEM_ASSERT(x.NumBlocks() == graph->InputFields().Size(),
"Number of input blocks (" << x.NumBlocks()
<< ") must match number of input fields (" << graph->InputFields().Size() << ")");
auto index_map = graph->GetFieldIdToIndexMap();
auto field_map = graph->GetFieldIdToFieldMap();
MFEM_ASSERT(y.NumBlocks() == index_map.NumFields(),
"Number of output blocks (" << y.NumBlocks()
<< ") must match number of fields (" << index_map.NumFields() << ")");
auto inputs = graph->InputFields();
for(int i=0; i < inputs.Size(); i++)
{
int idx = index_map.Get(inputs[i]->ID());
if(&y[idx] != &x[i]) // copy data, if address is different
{
y[idx] = x[i];
}
}
auto nodes = graph->Nodes();
for (auto node : nodes)
{
auto node_inputs = node->InputFields();
auto node_outputs = node->OutputFields();
x0_mv.SetNumBlocks(node_inputs.Size());
dx_mv.SetNumBlocks(node_inputs.Size());
dy_mv.SetNumBlocks(node_outputs.Size());
for(int i=0; i < node_inputs.Size(); i++)
{
int idx = index_map.Get(node_inputs[i]->ID());
x0_mv.MakeRef(i, xlin[idx]);
dx_mv.MakeRef(i, y[idx]);
}
for(int i=0; i < node_outputs.Size(); i++)
{
int idx = index_map.Get(node_outputs[i]->ID());
dy_mv.MakeRef(i, y[idx]);
}
node->GradientMult(x0_mv, dx_mv, dy_mv); // Compute JVP for the node
}
}
void GraphGradient::Reverse(const MultiVector &x, MultiVector &y) const
{
MFEM_ASSERT(x.NumBlocks() == graph->OutputFields().Size(),
"Number of input blocks (" << x.NumBlocks()
<< ") must match number of output fields (" << graph->OutputFields().Size() << ")");
auto index_map = graph->GetFieldIdToIndexMap();
auto field_map = graph->GetFieldIdToFieldMap();
int nnodes = graph->Size();
MFEM_ASSERT(y.NumBlocks() == index_map.NumFields(),
"Number of output blocks (" << y.NumBlocks()
<< ") must match number of fields (" << index_map.NumFields() << ")");
auto outputs = graph->OutputFields();
for(int i=0; i < outputs.Size(); i++)
{
int idx = index_map.Get(outputs[i]->ID());
if(&y[idx] != &x[i]) // copy data, if address is different
{
y[idx] = x[i];
}
}
for (int i=nnodes-1; i >= 0; i--)
{
auto node = graph->GetNode(i);
auto node_inputs = node->InputFields();
auto node_outputs = node->OutputFields();
x0_mv.SetNumBlocks(node_inputs.Size());
dx_mv.SetNumBlocks(node_outputs.Size());
dy_mv.SetNumBlocks(node_inputs.Size());
for(int i=0; i < node_inputs.Size(); i++)
{
int idx = index_map.Get(node_inputs[i]->ID());
x0_mv.MakeRef(i, xlin[idx]);
dy_mv.MakeRef(i, y[idx]);
}
for(int i=0; i < node_outputs.Size(); i++)
{
int idx = index_map.Get(node_outputs[i]->ID());
dx_mv.MakeRef(i, y[idx]);
}
node->GradientMultTranspose(x0_mv, dx_mv, dy_mv); // Compute JVP for the node
}
}
Operator& GraphGradient::GetGradient(const Vector &x) const
{
// Used to build Jacobian matrix
MFEM_ABORT("GraphGradient::GetGradient() not implemented");
}
} // namespace mfem
+801
View File
@@ -0,0 +1,801 @@
// 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.
#ifndef MFEM_MULTIAPP_HPP
#define MFEM_MULTIAPP_HPP
#include "mfem.hpp"
namespace mfem
{
/// Forward declarations needed below
class Field;
class FieldCollection;
class GraphNode;
class DAGraph;
class GraphGradient;
/// @brief Base class for storing data (Vector) and distinguishing
/// fields variables
class Field
{
public:
enum Type ///< Not used for now, but could be used to distinguish between input/output fields
{
INPUT , ///< Input field
OUTPUT, ///< Output field
DEFAULT ///< Any field
};
friend class GraphNode;
private:
Type type = Type::DEFAULT;
inline static int next_id = 0;
protected:
Vector *data = nullptr;
Vector *adjoint = nullptr; // For storing derivative info
int id = -1; // initialized to invalid id
std::string name; // Optional name for the field
Operator *oper = nullptr; // Operator that outputs this field
int GetValidID(int id_, int lb=0, int ub = std::numeric_limits<int>::max())
{
return (id_ >= lb && id_ <= ub) ? id_ : next_id++;
}
public:
///@brief Constructor for a Field of type Type with optional ID
Field(Vector *field, Vector *adjoint, Type type, int id_ = -1) :
type(type), data(field), adjoint(adjoint), id(GetValidID(id_)),
name("Field_" + std::to_string(id)) { }
///@brief Constructor for a Field of Default type with optional ID
Field(Vector *field, Vector *adjoint, int id_ = -1) :
Field(field, adjoint, Type::DEFAULT, id_) { }
///@brief Constructor for an input field
Field(Vector *field, int id_ = -1) :
Field(field, nullptr, Type::DEFAULT, id_) { }
///@brief Constructor for a Field of type Type
Field(Vector *field, Type type, int id_ = -1) :
Field(field, nullptr, type, id_) { }
///@brief Get the stored internally stored data pointer
Vector* Data() const { return data; }
Vector* Adjoint() const { return adjoint; }
Operator* GetOperator() const { return oper; }
///@brief Set the internally stored data pointer
virtual void SetData(Vector *field) { data = field; }
virtual void SetAdjoint(Vector *adj) { adjoint = adj; }
virtual void SetOperator(Operator *op) { oper = op; }
virtual void GetData(Vector &field) const { field = *data; }
virtual void GetAdjoint(Vector &adj) const { adj = *adjoint; }
std::string Name() const { return name; }
void SetName(const std::string &n) { name = n; }
int ID() const { return id; }
void SetID(int i)
{
MFEM_ASSERT(i >= 0, "ID must be non-negative.");
id = i;
}
bool IsInput() const {return (type == Type::INPUT);}
bool IsOutput() const {return (type == Type::OUTPUT);}
bool IsDefault() const {return (type == Type::DEFAULT);}
virtual ~Field() = default;
protected:
///@brief Set the type of the field (prevents changing type of input/output fields)
void SetType(Type t)
{
type = t;
}
};
/// @brief A collection of Fields, each identified by a name
class FieldCollection
{
public:
using FieldMap = GenericFieldMap<std::string, Field*>;
using IndexMap = GenericFieldMap<std::string, int>;
private:
std::string name; /// Name of the collection
Operator *oper = nullptr; /// Operator associated with this collection (not owned)
FieldMap fields; /// Map from field name to Field pointer
IndexMap index_map; /// Map from field name to index in input/output vectors
Array<Field*> input_fields; // Input fields for this node
Array<Field*> output_fields; // Output fields for this node
public:
FieldCollection() = default;
/// @brief Constructor with collection name and optional associated operator
FieldCollection(std::string collection_name, Operator *op = nullptr):
name(collection_name), oper(op) {}
/// @brief Constructor with associated operator and default collection name
FieldCollection(Operator *op) : name("FieldCollection"), oper(op) {}
/// @brief Get the number of fields in the collection
int Size() const { return fields.NumFields(); }
/// @brief Set the name of the collection
void SetName(const std::string &collection_name) { name = collection_name;}
/// @brief Get the name of the collection
std::string Name() const { return name; }
/// @brief Set the operator associated with this collection
void SetOperator(Operator *op){ oper = op; }
/// @brief Get the operator associated with this collection
const Operator* GetOperator() const { return oper; }
/// @brief Add a field to the collection with a given name and ownership flag
void AddField(const std::string &field_name, Field *field, bool own = false)
{
if(fields.Has(field_name))
{
MFEM_WARNING("FieldCollection::AddField: Field with name "
<< field_name << " already exists. Replacing existing field.");
}
fields.Register(field_name, field, own);
}
void AddInput(const std::string &field_name,
Field *field, bool own = false)
{
bool has_field = fields.Has(field_name);
bool has_index = index_map.Has(field_name);
if(has_field && has_index)
{
int i = index_map.Get(field_name);
input_fields[i] = field;
}
else
{
input_fields.push_back(field);
index_map.Register(field_name, input_fields.Size() - 1);
}
AddField(field_name, field, own);
}
void AddOutput(const std::string &field_name,
Field *field, bool own = false)
{
bool has_field = fields.Has(field_name);
bool has_index = index_map.Has(field_name);
if(has_field && has_index)
{
int i = index_map.Get(field_name);
output_fields[i] = field;
}
else
{
output_fields.push_back(field);
index_map.Register(field_name, output_fields.Size() - 1);
}
AddField(field_name, field, own);
if(field->GetOperator() == nullptr)
{
field->SetOperator(oper);
}
}
Array<Field*>& InputFields() { return input_fields; }
Array<Field*>& OutputFields() { return output_fields; }
Field* InputField(int i) const { return input_fields[i]; }
Field *InputField(const std::string &field_name) const
{
bool has_index = index_map.Has(field_name);
if(!has_index)
{
MFEM_WARNING("FieldCollection::InputField: Field with name "
<< field_name << " does not exist in the collection.");
return nullptr;
}
int index = index_map.Get(field_name);
MFEM_VERIFY(index >= 0 && index < input_fields.Size(),
"FieldCollection::InputField: Invalid index for field name: "
<< field_name << ".");
return input_fields[index];
}
Field* OutputField(int i) const { return output_fields[i]; }
Field *OutputField(const std::string &field_name) const
{
bool has_index = index_map.Has(field_name);
if(!has_index)
{
MFEM_WARNING("FieldCollection::OutputField: Field with name "
<< field_name << " does not exist in the collection.");
return nullptr;
}
int index = index_map.Get(field_name);
MFEM_VERIFY(index >= 0 && index < output_fields.Size(),
"FieldCollection::OutputField: Invalid index for field name: "
<< field_name << ".");
return output_fields[index];
}
FieldMap &Fields() { return fields; }
FieldMap Fields() const { return fields; }
virtual void Save (std::ostream &out) const
{
out << "\"Fields\":\n";
out << "{\n";
for (auto f = fields.begin(); f != fields.end(); ++f)
{
std::string f_name = f->first;
Field *f_obj = f->second;
// out << " " << f_name << ": ID " << f_obj->ID() << ",\n";
// out << f_obj->ID() << ": " << f_name << ",\n";
out << '\"' << f_obj->ID() << "\": \"" << f_name << "\"";
if(f != std::prev(fields.end())) out << ",";
out << "\n";
}
out << "},\n";
out << "\"Inputs\":\n";
out << "{\n";
for (int i = 0; i < input_fields.Size(); ++i)
{
Field *f_obj = input_fields[i];
out << '\"' << f_obj->ID() << "\": \"" << f_obj->Name() << "\"";
if(i != input_fields.Size() - 1) out << ",";
out << "\n";
}
out << "},\n";
out << "\"Outputs\":\n";
out << "{\n";
for (int i = 0; i < output_fields.Size(); ++i)
{
Field *f_obj = output_fields[i];
out << '\"' << f_obj->ID() << "\": \"" << f_obj->Name() << "\"";
if(i != output_fields.Size() - 1) out << ",";
out << "\n";
}
out << "}\n";
}
Field* HasField(const Field &field) const
{
for (auto f = fields.begin(); f != fields.end(); ++f)
{
if(f->second == &field)
{
return f->second;
}
}
return nullptr;
}
Field* HasField(const std::string &field_name) const
{
return fields.Get(field_name);
}
Field* HasField(const int id) const
{
for (auto f = fields.begin(); f != fields.end(); ++f)
{
if(f->second->ID() == id)
{
return f->second;
}
}
return nullptr;
}
~FieldCollection(){}
};
class GraphNode : public Operator
{
public:
enum ExecutionMode
{
GRADIENT_MODE, ///< Node is being executed as part of a gradient evaluation
DEFAULT_MODE ///< Node is being executed as default, operator evaluation
};
private:
inline static int next_id = 0;
protected:
int id = -1;
int node_index = -1;
mutable ExecutionMode exec_mode = DEFAULT_MODE;
std::string name;
mutable FieldCollection field_collection; ///< Collection of fields associated with this node
// Offsets to be used for operation on BlockVector
Array<int> input_offsets; ///< Offsets for input fields
Array<int> output_offsets; ///< Offsets for output fields
int GetValidID(int id_, int lb=0, int ub = std::numeric_limits<int>::max())
{
return (id_ >= lb && id_ <= ub) ? id_ : next_id++;
}
public:
GraphNode(int h, int w) : Operator(h,w), id(GetValidID(-1)),
name("Node_" + std::to_string(id)),
field_collection(this) { }
GraphNode(int s = 0) : GraphNode(s, s) { }
void SetNodeIndex(int index){ node_index = index; }
int GetNodeIndex() const { return node_index; }
void SetExecutionMode(ExecutionMode mode) { exec_mode = mode; }
ExecutionMode GetExecutionMode() const { return exec_mode; }
void SetName(const std::string &name_) { name = name_; }
std::string Name() const { return name; }
void SetID(int id_) { id = id_; }
int ID() const { return id; }
FieldCollection::FieldMap& Fields() { return field_collection.Fields(); }
Field* Fields(const std::string &f) { return Fields().Get(f); }
FieldCollection::FieldMap Fields() const { return field_collection.Fields(); }
Field* Fields(const std::string &f) const { return Fields().Get(f); }
Array<Field*>& InputFields() const { return field_collection.InputFields(); }
Array<Field*>& OutputFields() const { return field_collection.OutputFields(); }
Field* InputField(int i) const { return InputFields()[i]; }
Field* OutputField(int i) const { return OutputFields()[i]; }
virtual void AddInput(const std::string &field_name,
Field *field, bool own = false)
{ field_collection.AddInput(field_name, field, own); }
virtual void AddInput(Field *field, bool own = false)
{ AddInput(field->Name(), field, own); }
template<bool OwnInputs = false,
typename... Args,
bool AreFields = std::conjunction<std::is_base_of<Field, std::remove_pointer_t<Args>> ...>::value,
typename std::enable_if<AreFields, bool>::type = true >
void AddInputs(Args... args)
{
((AddInput(std::forward<Args>(args), OwnInputs)), ...);
}
virtual void AddOutput(const std::string &field_name,
Field *field, bool own = false)
{ field_collection.AddOutput(field_name, field, own); }
virtual void AddOutput(Field *field, bool own = false)
{ AddOutput(field->Name(), field, own); }
template<bool OwnOutputs = false,
typename... Args,
bool AreFields = std::conjunction<std::is_base_of<Field, std::remove_pointer_t<Args>> ...>::value,
typename std::enable_if<AreFields, bool>::type = true >
void AddOutputs(Args... args)
{
((AddOutput(std::forward<Args>(args), OwnOutputs)), ...);
}
/// @brief Return the input offsets for block starts.
Array<int>& InputOffsets() { return input_offsets; }
/// @brief Read only access to the input offsets for block starts.
const Array<int>& InputOffsets() const { return input_offsets; }
void SetInputOffsets(const Array<int> &offsets) { input_offsets = offsets; }
/// @brief Return the output offsets for block starts.
Array<int>& OutputOffsets() { return output_offsets; }
/// @brief Read only access to the output offsets for block starts.
const Array<int>& OutputOffsets() const { return output_offsets; }
void SetOutputOffsets(const Array<int> &offsets) { output_offsets = offsets; }
virtual void MultMV(const MultiVector &x, MultiVector &y) const
{
MFEM_ABORT("This method is not overridden for this class!");
// Include a default implementation that calls Mult() with Vector arguments
}
virtual void MultTransposeMV(const MultiVector &x, MultiVector &y) const
{
MFEM_ABORT("This method is not overridden for this class!");
// Include a default implementation that calls MultTranspose() with Vector arguments
}
virtual void GradientMult(const MultiVector &x, const MultiVector &dx, MultiVector &dy) const
{
MFEM_ABORT("This method is not overridden for this class!");
GetGradientMV(x).MultMV(dx, dy);
}
virtual void GradientMultTranspose(const MultiVector &x, const MultiVector &dx, MultiVector &dy) const
{
MFEM_ABORT("This method is not overridden for this class!");
GetGradientMV(x).MultTransposeMV(dx, dy); // Not yet implemented
}
virtual void Save (std::ostream &out) const
{
out << "\"Node-" << id << "\" : " << std::endl;
out << "{\n";
out << "\"Name\": \"" << name << "\",\n";
field_collection.Save(out);
out << "}";
}
virtual ~GraphNode() = default;
};
/**
@brief An abstract, type-erased class to define the interface for
operators, not inherited from @a GraphNode. It performs SFINAE
checks for stored operator's member functions and override the Mult
to call the stored object's functions.
*/
template <typename OpType>
class AbstractOperator : public GraphNode
{
protected:
/// Define a template class 'check' to test for the existence of member functions
template <typename C>
class CheckMember{
private:
/// @brief A type trait to check if the erased class has the function Mult
/// with the needed signatures.
template<class T>
using Mult = decltype(std::declval<T&>().Mult(std::declval<const Vector&>(),
std::declval<Vector&>()));
template<class T>
using MultPtr = decltype(std::declval<T&>().Mult(std::declval<const int>(),
std::declval<const real_t*>(),
std::declval<const int>(),
std::declval<real_t*>()));
// ---------------------------------------------------------------------
template <typename T, template<typename> typename Func, typename R>
static constexpr auto Check(T*) -> typename std::is_same< Func<T>, R>::type;
template <typename, template<typename> typename, typename >
static constexpr std::false_type Check(...);
// --- Check for the existence of the member functions
typedef decltype(Check<C,Mult,void>(0)) Has_Mult;
typedef decltype(Check<C,MultPtr,void>(0)) Has_MultPtr;
public:
static constexpr bool HasMult = Has_Mult::value;
static constexpr bool HasMultPtr = Has_MultPtr::value;
};
OpType *op; ///< Pointer to the operator
public:
constexpr bool HasExecute(){return CheckMember<OpType>::HasStep;}
constexpr bool HasMult(){return CheckMember<OpType>::HasMult;}
/// @brief Constructor for the type-erased AbstractOperator class
AbstractOperator(OpType *op_, int h, int w) : GraphNode(h,w), op(op_)
{ }
/// @brief Constructor for the type-erased AbstractOperator class.
AbstractOperator(OpType *op_, int s = 0) : AbstractOperator(op_,s,s) {}
/**
@brief Perform Mult operation with the stored operator, if it exists.
*/
void Mult(const Vector &x, Vector &y) const override
{
if constexpr (CheckMember<OpType>::HasMult)
{
op->Mult(x,y);
}
else if constexpr (CheckMember<OpType>::HasMultPtr)
{
op->Mult(x.Size(), x.GetData(), y.Size(), y.GetData());
}
else
{
MFEM_ABORT("The AbstractOperator does not have the function, "
"Mult(const Vector&, Vector&) or "
"Mult(int, double*, int, double*).");
}
}
};
/**
@brief A class to store and coupled multiple operators together.
*/
class DAGraph : public GraphNode
{
public:
using IntToIntMap = GenericFieldMap<int, int>;
using IntToFieldMap = GenericFieldMap<int, Field*>;
enum class GradMode
{
FINITE_DIFF = 0, ///< Finite difference Jacobian
MATRIX_FREE = 1, ///< Matrix-free Jacobian
ASSEMBLED = 2, ///< Assembled Jacobian
NONE = 3 ///< Not implemented
};
protected:
Array<GraphNode*> nodes; ///< Vector of individual operators
Array<bool> node_owned; ///< Whether the operators are owned
Array<int> node_depth; ///< Depth of each operator in the graph
int max_width = 0; ///< Largest operator width
int max_height = 0; ///< Largest operator height
int nnodes = 0; ///< The number of nodes
bool sorted = false; ///< True if the nodes are topologically sorted
bool assembled = false; ///< True if the graph is assembled
GradMode grad_mode = GradMode::MATRIX_FREE; ///< Gradient mode for the graph
mutable Operator *grad = nullptr; ///< Gradient operator
mutable MultiVector xmv_node, ymv_node; ///< Temporary multivectors for evaluating nodes
IntToFieldMap fid_to_field; ///< Map from Field ID to Field pointer
IntToIntMap fid_to_index; ///< Map from ID to index in an array; needed since ordering is not unique
friend class GraphGradient;
public:
/**
@brief Construct a new CoupledOperator object.
@param nop Total number of operators to couple
*/
DAGraph(const int nop) : GraphNode()
{
nodes.Reserve(nop);
node_owned.Reserve(nop);
}
/**
@brief Construct a new CoupledOperator object for an
abstract non/mfem operator.
*/
template <class OpType>
DAGraph(const OpType &op) : DAGraph(1)
{
AddOperator(op);
}
/**
@brief Add an operator to the list of coupled operator and
return pointer to it. Not owned unless it's not derived from GraphNode.
*/
template <class OpType>
GraphNode* AddOperator(OpType *op_, int h, int w)
{
// Add operator to list of operators
if constexpr(std::is_base_of<GraphNode, OpType>::value)
{
nodes.push_back(op_);
node_owned.Append(false);
}
else
{
nodes.push_back(new AbstractOperator<OpType>(op_,h,w));
node_owned.Append(true);
}
nnodes++;
// Update size of the coupled operator and the block offsets
GraphNode* op = nodes.Last();
op->SetNodeIndex(nnodes-1); // Set the index of the operator
int ht = op->Height();
int wt = op->Width();
max_width = std::max(max_width, wt);
max_height = std::max(max_height, ht);
sorted = false;
return op;
}
/// @brief Add an operator to the list of coupled operator and return pointer to it.
template <class OpType>
GraphNode* AddOperator(OpType *op_, int s = 0) { return AddOperator(op_,s,s);}
/// @brief Get the number of coupled operators
int Size(){return nnodes;}
/// @brief Get the size of the largest operator
int MaxWidth() const {return max_width;}
int MaxHeight() const {return max_height;}
IntToIntMap &GetFieldIdToIndexMap() { return fid_to_index; }
IntToIntMap GetFieldIdToIndexMap() const { return fid_to_index; }
IntToFieldMap &GetFieldIdToFieldMap() { return fid_to_field; }
IntToFieldMap GetFieldIdToFieldMap() const { return fid_to_field; }
/// @brief Get the operator at index @a i
GraphNode* GetNode(const int i)
{
MFEM_ASSERT(i >= 0 && i < nnodes,
"index [" << i << "] is out of range [0," << nnodes << ")");
return nodes[i];
}
Array<GraphNode*>& Nodes() { return nodes; }
/// @brief Specify whether the operator at index @a i is owned.
void OwnNode(const int i, bool own = true)
{
MFEM_ASSERT(i >= 0 && i < nnodes,
"index [" << i << "] is out of range [0," << nnodes << ")");
node_owned[i] = own;
}
void Assemble();
bool IsAssembled() const { return assembled; }
void TopologicalSort();
bool IsSorted() const { return sorted; }
void ComputeDepth();
void ValidateOffsets();
void ValidateNode(GraphNode &node);
void CollectFieldMaps();
using GraphNode::AddInput;
void AddInput(Field *field, int sz, bool own = false)
{
if(input_offsets.Size() == 0)
{ // First entry
input_offsets.Append(0);
}
input_offsets.Append(input_offsets.Last() + sz);
AddInput(field, own);
}
using GraphNode::AddOutput;
void AddOutput(Field *field, int sz, bool own = false)
{
if(output_offsets.Size() == 0)
{ // First entry
output_offsets.Append(0);
}
output_offsets.Append(output_offsets.Last() + sz);
AddOutput(field, own);
}
/// @brief Set the gradient mode for the coupled operator
void SetGradientMode(GradMode mode)
{
if(mode != grad_mode)
{
if(grad) { delete grad; grad = nullptr; }
grad_mode = mode;
}
}
/**
@brief Apply the operator to the vector @a x
and return the result in @a y.
*/
virtual void Mult(const Vector &x, Vector &y) const override;
virtual void MultMV(const MultiVector &x, MultiVector &y) const override;
virtual void Execute(const MultiVector &x, MultiVector &y) const;
virtual void Save (std::ostream &out) const
{
out << "\"DAGraph\":\n";
out << "{\n";
// out << "\"nodes\" : " << nnodes << ",\n";
out << "\"Nodes\":\n";
out << "{\n";
for (int i = 0; i < nodes.Size(); i++)
{
nodes[i]->Save(out);
if(i != nodes.Size()-1) out << ",";
out << "\n";
}
out << "},\n"; // End of Nodes
field_collection.Save(out);
out << "}\n";
}
Operator& GetGradient(const Vector &x) const override;
/// @brief Destroy the Coupled Application object
~DAGraph();
};
class GraphGradient : public Operator
{
protected:
mutable DAGraph *graph = nullptr; ///< Pointer to the DAGraph for which this is the gradient operator
Array<Vector*> x_work; ///< Array to store linearization point (intermediate fields)
mutable MultiVector xlin;
mutable MultiVector x0_mv, dx_mv, dy_mv;
public:
GraphGradient(DAGraph &dag);
void Update(const Vector &x);
void Mult(const Vector &x, Vector &y) const override;
void MultMV(const MultiVector &x, MultiVector &y) const override;
void MultTranspose(const Vector &x, Vector &y) const override;
void MultTransposeMV(const MultiVector &x, MultiVector &y) const override;
Operator &GetGradient(const Vector &x) const override;
void Forward(const MultiVector &x, MultiVector &y) const;
void Reverse(const MultiVector &x, MultiVector &y) const;
~GraphGradient()
{
for (auto &v : x_work)
{
if(v) { delete v; v = nullptr; }
}
x_work.DeleteAll();
}
};
} //mfem namespace
#endif