Compare commits

...
31 Commits
Author SHA1 Message Date
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
Veselin Dobrev 1ed3b48c2e In class MultiVector, remove the need for Memory flag synchronizations
in some cases. This required changes in the internals of the class.

Added some new methods in class MultiVector.
2026-02-26 09:57:21 -08:00
Veselin Dobrev fbd9189e7b Restrist with 'enable_if' the variadic template MultiVector ctor and
MakeRef method to be considered only when the arg types are convertible
to (Vector &).
2026-02-25 19:17:07 -08:00
Veselin Dobrev 1dd889cb16 Add support for constructing and re-constructing MultiVectors to reference
multiple Vectors given as arguments.
2026-02-25 17:44:31 -08:00
Veselin Dobrev 2e8fbd661a Fix a warning in a miniapp. 2026-02-25 14:56:28 -08:00
Veselin Dobrev 6e424dba6e Draft implementation of an array-of-Vectors class where each Vector generally
has a different size and is allocated independently.

The tentative name for the new class is MultiVector.

In class Operator, added new virtual methods Mult() and GetGradient() that
use MultiVectors.
2026-02-25 13:51:43 -08:00
15 changed files with 3051 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.
+3
View File
@@ -224,6 +224,9 @@ public:
/** @see GetGradient(const Vector &) */
Operator &GetGradient(const Vector &x, bool finalize) const;
/// Suppress a warning about hiding overloaded virtual function.
using Operator::GetGradient;
/// Update the NonlinearForm to propagate updates of the associated FE space.
/** After calling this method, the essential boundary conditions need to be
set again. */
+2
View File
@@ -27,6 +27,7 @@ list(APPEND SRCS
handle.cpp
matrix.cpp
mma.cpp
multivector.cpp
ode.cpp
operator.cpp
ordering.cpp
@@ -63,6 +64,7 @@ list(APPEND HDRS
linalg.hpp
matrix.hpp
mma.hpp
multivector.hpp
ode.hpp
operator.hpp
ordering.hpp
+1
View File
@@ -15,6 +15,7 @@
// Linear algebra header file
#include "vector.hpp"
#include "multivector.hpp"
#include "operator.hpp"
#include "matrix.hpp"
#include "sparsemat.hpp"
+60
View File
@@ -0,0 +1,60 @@
// 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 "multivector.hpp"
namespace mfem
{
MultiVector::MultiVector(const Array<int> &vector_sizes)
{
SetSizes(vector_sizes);
}
MultiVector::MultiVector(const Array<int> &vector_sizes, MemoryType mt)
{
SetSizes(vector_sizes, mt);
}
MultiVector::MultiVector(Vector &base, const Array<int> &vector_sizes)
{
MakeRef(base, vector_sizes);
}
void MultiVector::SetSizes(const Array<int> &vector_sizes)
{
blocks.resize(vector_sizes.Size());
for (int i = 0; i < vector_sizes.Size(); i++)
{
operator[](i).SetSize(vector_sizes[i]);
}
}
void MultiVector::SetSizes(const Array<int> &vector_sizes, MemoryType mt)
{
blocks.resize(vector_sizes.Size());
for (int i = 0; i < vector_sizes.Size(); i++)
{
operator[](i).SetSize(vector_sizes[i], mt);
}
}
void MultiVector::MakeRef(Vector &base, const Array<int> &vector_sizes)
{
blocks.resize(vector_sizes.Size());
for (int offset = 0, i = 0; i < vector_sizes.Size(); i++)
{
blocks[i].emplace<0>(base, offset, vector_sizes[i]);
offset += vector_sizes[i];
}
}
} // namespace mfem
+198
View File
@@ -0,0 +1,198 @@
// 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_MULTIVECTOR_HPP
#define MFEM_MULTIVECTOR_HPP
#include "../general/array.hpp"
#include "vector.hpp"
#include <vector>
#include <array>
#include <variant>
namespace mfem
{
/// Class representing an array of Vectors with generally different sizes.
/** This class is similar to BlockVector with the following two main
differences:
- the data for the individual Vector blocks does not need to be part of one
big contiguous memory allocation;
- this class does not inherit from class Vector (as a consequence of the
first bullet).
Internally, each Vector block is represented as either:
- (default) a Vector object constructed and owned by this class; this
object, in turn, as any Vector object, can own its Memory allocation or
refer to a sub-Memory of another Memory object; or
- a pointer to an externally allocated Vector or classes derived from
Vector. */
class MultiVector
{
private:
std::vector<std::variant<Vector,Vector*>> blocks;
public:
/// Create an empty MultiVector with zero blocks.
MultiVector() = default;
/** @brief Create a MultiVector with @a num_blocks blocks. The individual
Vector blocks are default initialized, i.e. they all have size zero. */
MultiVector(int num_blocks)
: blocks(num_blocks) { }
/** @brief Construct a MultiVector with number of blocks and individual block
Vector sizes given by @a vector_sizes.
@note The memory of the individual Vector blocks is NOT initialized. */
MultiVector(const Array<int> &vector_sizes);
/** @brief Construct a MultiVector with number of blocks and individual block
Vector sizes given by @a vector_sizes. All Vector blocks use the
MemoryType @a mt.
@note The memory of the individual Vector blocks is NOT initialized. */
MultiVector(const Array<int> &vector_sizes, MemoryType mt);
/** @brief Construct a MultiVector referencing data within a given monolithic
Vector @a base.
With this constructor, the Memory flags of @a base and of the individual
Vector blocks may need to be explicitly synchronized when data is moved
between host and device. */
MultiVector(Vector &base, const Array<int> &vector_sizes);
/** @brief Construct a MultiVector referencing multiple Vectors given as
arguments.
The VectorTypes reference arguments are expected to be static_cast-able
to (Vector &) which is the case if the types are derived from Vector,
e.g. HypreParVector, GridFunction, etc.
With this constructor, operations on individual Vector blocks are
performed directly on the objects @a vs. In particular, there is no need
to synchronize the Memory flags of @a vs and the ones of the individual
Vector blocks when data is moved between host and device. */
template <typename... VectorTypes,
std::enable_if_t<
std::conjunction_v<
std::is_convertible<VectorTypes&,Vector&>...>, bool> = true>
MultiVector(VectorTypes &...vs) { MakeRef(vs...); }
/// Return the number of Vectors in the MultiVector.
int NumBlocks() const { return blocks.size(); }
/** @brief Set the number of Vectors in the MultiVector. Existing Vector
blocks will remain unmodified. New Vector blocks will be default
initialized, i.e. they all have size zero. */
void SetNumBlocks(int num_blocks) { blocks.resize(num_blocks); }
/// Read-write access to the i-th Vector.
inline Vector &operator[](int i);
/// Read-only access to the i-th Vector.
inline const Vector &operator[](int i) const;
/** @brief Update the MultiVector according to the given @a vector_sizes.
This method can be used to add or remove blocks. The individual Vector
sizes are updated using the method Vector::SetSize(int). */
void SetSizes(const Array<int> &vector_sizes);
/** @brief Update the MultiVector according to the given @a vector_sizes and
MemoryType @a mt.
This method can be used to add or remove blocks. The individual Vector
sizes and MemoryType are updated using the method
Vector::SetSize(int, MemoryType). */
void SetSizes(const Array<int> &vector_sizes, MemoryType mt);
/** @brief Update the MultiVector to reference data within a given monolithic
Vector @a base.
After calling this method, the Memory flags of @a base and of the
individual Vector blocks may need to be explicitly synchronized when data
is moved between host and device.*/
void MakeRef(Vector &base, const Array<int> &vector_sizes);
/** @brief Update the @a i-th MultiVector block to reference data within the
given monolithic Vector @a base at the given @a offset and with the given
@a size.
After calling this method, the Memory flags of @a base and of the @a i-th
Vector block may need to be explicitly synchronized when data is moved
between host and device.*/
inline void MakeRef(int i, Vector &base, int offset, int size)
{
blocks[i].emplace<0>(base, offset, size);
}
/** @brief Update the MultiVector to reference multiple Vectors given as
arguments.
The VectorTypes reference arguments are expected to be static_cast-able
to (Vector &) which is the case if the types are derived from Vector,
e.g. HypreParVector, GridFunction, etc.
After calling this method, operations on individual Vector blocks are
performed directly on the objects @a vs. In particular, there is no need
to synchronize the Memory flags of @a vs and the ones of the individual
Vector blocks when data is moved between host and device. */
template <typename... VectorTypes,
std::enable_if_t<
std::conjunction_v<
std::is_convertible<VectorTypes&,Vector&>...>, bool> = true>
inline void MakeRef(VectorTypes &...vs);
/** @brief Update the @a i-th MultiVector block to reference the given
Vector @a v.
After calling this method, operations on the @a i-th Vector block are
performed directly on the Vector @a v. In particular, there is no need
to synchronize the Memory flags of @a v and the ones of the @a i-th
Vector blocks when data is moved between host and device. */
inline void MakeRef(int i, Vector &v) { blocks[i] = &v; }
};
// Inline and template methods
inline Vector &MultiVector::operator[](int i)
{
auto &bi = blocks[i];
return (bi.index() == 0) ? std::get<0>(bi) : *std::get<1>(bi);
}
inline const Vector &MultiVector::operator[](int i) const
{
auto &bi = blocks[i];
return (bi.index() == 0) ? std::get<0>(bi) : *std::get<1>(bi);
}
template <typename... VectorTypes,
std::enable_if_t<
std::conjunction_v<
std::is_convertible<VectorTypes&,Vector&>...>, bool>>
inline void MultiVector::MakeRef(VectorTypes &...vs)
{
blocks.resize(sizeof...(vs));
if constexpr (sizeof...(vs) > 0)
{
const std::array vs_p{&static_cast<Vector&>(vs)...};
for (std::size_t i = 0; i < sizeof...(vs); i++)
{
blocks[i] = vs_p[i];
}
}
}
} // namespace mfem
#endif // MFEM_MULTIVECTOR_HPP
+10
View File
@@ -111,6 +111,16 @@ void Operator::ArrayAddMultTranspose(const Array<const Vector *> &X,
}
}
void Operator::Mult(const MultiVector &, MultiVector &)
{
MFEM_ABORT("this method is not overriden for this class!");
}
Operator &Operator::GetGradient(const MultiVector &) const
{
MFEM_ABORT("this method is not overriden for this class!");
}
void Operator::FormLinearSystem(const Array<int> &ess_tdof_list,
Vector &x, Vector &b,
Operator* &Aout, Vector &X, Vector &B,
+21
View File
@@ -13,6 +13,7 @@
#define MFEM_OPERATOR
#include "vector.hpp"
#include "multivector.hpp"
namespace mfem
{
@@ -129,6 +130,16 @@ public:
virtual void ArrayAddMultTranspose(const Array<const Vector *> &X,
Array<Vector *> &Y, const real_t a = 1.0) const;
/** @brief Operator application, y = A(x), where the input @a x and the
output @a y are MultiVector objects, i.e. they generally use
non-contiguous memory representation.
The typical use case for this method are block operators like
DifferentiableOperator.
The base class implementation for the method is to generate an error. */
virtual void Mult(const MultiVector &x, MultiVector &y);
/** @brief Evaluate the gradient operator at the point @a x. The default
behavior in class Operator is to generate an error. */
virtual Operator &GetGradient(const Vector &x) const
@@ -137,6 +148,16 @@ public:
return const_cast<Operator &>(*this);
}
/** @brief Evaluate the gradient operator at the point @a x. The input @a x
is provided as a MultiVector, i.e. it generally uses non-contiguous
memory representation.
The typical use case for this method are block operators like
DifferentiableOperator.
The base class implementation for the method is to generate an error. */
virtual Operator &GetGradient(const MultiVector &x) const;
/** @brief Computes the diagonal entries into @a diag. Typically, this
operation only makes sense for linear Operator%s. In some cases, only an
approximation of the diagonal is computed. */
+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)
@@ -80,6 +80,8 @@ public:
// limitations
void MultRT_2D(const Vector &x, Vector &y, Mode mode) const;
void MultRT_3D(const Vector &x, Vector &y, Mode mode) const;
// suppress warning about hiding overloaded virtual function:
using Operator::Mult;
};
} // namespace mfem
+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()
+895
View File
@@ -0,0 +1,895 @@
#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));
const_cast<DiffusionCoefficient*>(this)->Mult(xmv, ymv);
}
void Mult(const MultiVector &x, MultiVector &y) 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
}
}
// TODO: Possibly delete and only support MultiVector version of GradientMult
void GradientMult(const Vector &x, const Vector &dx, Vector &dy) const override
{
BlockVector xb(x.GetData(), InputOffsets());
BlockVector dxb(dx.GetData(), InputOffsets());
BlockVector dyb(dy.GetData(), OutputOffsets());
MultiVector xmv(1), dxmv(1), dymv(1);
xmv.MakeRef(0, xb.GetBlock(0));
dxmv.MakeRef(0, dxb.GetBlock(0));
dymv.MakeRef(0, dyb.GetBlock(0));
const_cast<DiffusionCoefficient*>(this)->GradientMult(xmv, dxmv, dymv);
}
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));
const_cast<ProductGridFunctions*>(this)->Mult(xmv, ymv);
}
void Mult(const MultiVector &x, MultiVector &y) 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);
}
// TODO: Possibly delete and only support MultiVector version of GradientMult
void GradientMult(const Vector &x, const Vector &dx, Vector &dy) const override
{
BlockVector xb(x.GetData(), InputOffsets());
BlockVector dxb(dx.GetData(), InputOffsets());
BlockVector dyb(dy.GetData(), OutputOffsets());
MultiVector xmv(x_gf.size()), dxmv(x_gf.size()), dymv(1);
for (size_t i = 0; i < x_gf.size(); i++)
{
xmv.MakeRef(i, xb.GetBlock(i));
dxmv.MakeRef(i, dxb.GetBlock(i));
}
dymv.MakeRef(0, dyb.GetBlock(0));
const_cast<ProductGridFunctions*>(this)->GradientMult(xmv, dxmv, dymv);
}
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));
const_cast<DiffusionOperator*>(this)->Mult(xmv, ymv);
}
void Mult(const MultiVector &x, MultiVector &y) 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");
}
// TODO: Possibly delete and only support MultiVector version of GradientMult
void GradientMult(const Vector &x, const Vector &dx, Vector &dy) const override
{
BlockVector xb(x.GetData(), InputOffsets());
BlockVector dxb(dx.GetData(), InputOffsets());
BlockVector dyb(dy.GetData(), OutputOffsets());
Vector &Tadj = dxb.GetBlock(0);
Vector &kadj = dxb.GetBlock(1);
Vector &yadj = dyb.GetBlock(0);
Vector &tdofs = xb.GetBlock(0);
Vector &kdofs = xb.GetBlock(1);
dfdT_mat->Mult(Tadj, yadj);
dfdk_mat->AddMult(kadj, yadj);
}
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;
}
+873
View File
@@ -0,0 +1,873 @@
// 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));
}
const_cast<DAGraph*>(this)->Mult(xmv, ymv);
}
void DAGraph::Mult(const MultiVector &x, MultiVector &y)
{
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];
}
}
if(input_type == InputType::VECTOR)
{
x_node.SetSize(MaxWidth());
y_node.SetSize(MaxHeight());
for (auto node : nodes)
{
x_node.SetSize(node->Width());
y_node.SetSize(node->Height());
// Assemble input fields into a single vector for the node
auto node_inputs = node->InputFields();
auto ioffsets = node->InputOffsets();
for (int i=0; i < node_inputs.Size(); i++)
{
auto in_field = node_inputs[i];
int idx = index_map.Get(in_field->ID());
x_node.SetVector(y[idx],ioffsets[i]);
}
node->Mult(x_node, y_node);
// Disassemble output vector back
auto node_outputs = node->OutputFields();
BlockVector ynb(y_node.GetData(), node->OutputOffsets());
for (int i=0; i < node_outputs.Size(); i++)
{
auto out_field = node_outputs[i];
int idx = index_map.Get(out_field->ID());
y[idx] = ynb.GetBlock(i);
}
}
}
else if(input_type == InputType::MULTIVECTOR)
{
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->Mult(xmv_node, ymv_node);
}
}
else if(input_type == InputType::NONE)
{
Vector x_unused, y_unused;
for (auto node : nodes)
{
node->Mult(x_unused, y_unused);
}
}
else
{
MFEM_ABORT("DAGraph::Execute() not implemented for input type: "
<< static_cast<int>(input_type));
}
}
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));
}
const_cast<GraphGradient*>(this)->Mult(xmv, ymv); // Forward mode: compute JVP, y = J(z) * x
}
void GraphGradient::Mult(const MultiVector &x, MultiVector &y)
{
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 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->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));
}
const_cast<GraphGradient*>(this)->MultTranspose(xmv, ymv); // Reverse mode: compute VJP, y = J(z)^T * x
}
void GraphGradient::MultTranspose(const MultiVector &x, MultiVector &y)
{
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 member function?
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 in_type = graph->GetInputType();
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];
}
}
if(in_type == InputType::VECTOR)
{
x0.SetSize(graph->MaxWidth());
dx.SetSize(graph->MaxWidth());
dy.SetSize(graph->MaxHeight());
auto nodes = graph->Nodes();
for (auto node : nodes)
{
x0.SetSize(node->Width());
dx.SetSize(node->Width());
dy.SetSize(node->Height());
// Assemble input fields into a single vector for the node
auto node_inputs = node->InputFields();
auto ioffsets = node->InputOffsets();
for(int i=0; i < node_inputs.Size(); i++)
{
auto in_field = node_inputs[i];
MFEM_ASSERT(index_map.Has(in_field->ID()), "Input field ID not found in index_map");
int idx = index_map.Get(in_field->ID());
x0.SetVector(xlin[idx], ioffsets[i]);
dx.SetVector(y[idx], ioffsets[i]);
}
node->GradientMult(x0, dx, dy); // Compute JVP for the node
// Disassemble output vector back
auto node_outputs = node->OutputFields();
BlockVector ynb(dy.GetData(), node->OutputOffsets());
for(int i=0; i < node_outputs.Size(); i++)
{
auto out_field = node_outputs[i];
MFEM_ASSERT(index_map.Has(out_field->ID()), "Output field ID not found in index_map");
int idx = index_map.Get(out_field->ID());
y[idx] = ynb.GetBlock(i);
}
}
}
else if(in_type == InputType::MULTIVECTOR)
{
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
}
}
else if(in_type == InputType::NONE)
{
Vector x_unused, dx_unused, dy_unused;
auto nodes = graph->Nodes();
for (auto node : nodes)
{
node->GradientMult(x_unused, dx_unused, dy_unused);
}
}
else
{
MFEM_ABORT("GraphGradient::Forward() not implemented for input type: "
<< static_cast<int>(in_type));
}
}
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 in_type = graph->GetInputType();
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];
}
}
if(in_type == InputType::VECTOR)
{
x0.SetSize(graph->MaxWidth());
dx.SetSize(graph->MaxHeight());
dy.SetSize(graph->MaxWidth());
for (int i=nnodes-1; i >= 0; i--)
{
auto node = graph->GetNode(i);
x0.SetSize(node->Width());
dx.SetSize(node->Height());
dy.SetSize(node->Width());
auto node_inputs = node->InputFields();
auto ioffsets = node->InputOffsets();
for(int i=0; i < node_inputs.Size(); i++)
{
auto in_field = node_inputs[i];
MFEM_ASSERT(index_map.Has(in_field->ID()), "Input field ID not found in index_map");
int idx = index_map.Get(in_field->ID());
x0.SetVector(xlin[idx], ioffsets[i]);
}
auto node_outputs = node->OutputFields();
auto ooffsets = node->OutputOffsets();
for(int i=0; i < node_outputs.Size(); i++)
{
auto out_field = node_outputs[i];
MFEM_ASSERT(index_map.Has(out_field->ID()), "Output field ID not found in index_map");
int idx = index_map.Get(out_field->ID());
dx.SetVector(y[idx], ooffsets[i]);
}
node->GradientMultTranspose(x0, dx, dy); // Compute JVP for the node
BlockVector dynb(dy.GetData(), node->InputOffsets());
for(int i=0; i < node_inputs.Size(); i++)
{
int idx = index_map.Get(node_inputs[i]->ID());
y[idx] = dynb.GetBlock(i);
}
}
}
else if(in_type == InputType::MULTIVECTOR)
{
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
}
}
else if(in_type == InputType::NONE)
{
Vector x_unused, dx_unused, dy_unused;
for (int i=nnodes-1; i >= 0; i--)
{
auto node = graph->GetNode(i);
node->GradientMultTranspose(x_unused, dx_unused, dy_unused); // Compute VJP for the node
}
}
else
{
MFEM_ABORT("GraphGradient::Reverse() not implemented for input type: "
<< static_cast<int>(in_type));
}
}
Operator& GraphGradient::GetGradient(const Vector &x) const
{
// Used to build Jacobian matrix
MFEM_ABORT("GraphGradient::GetGradient() not implemented");
}
} // namespace mfem
+838
View File
@@ -0,0 +1,838 @@
// 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 Get the field associated with the given name, or nullptr if not found
Field* GetField(const std::string &field_name) const
{
return fields.Get(field_name);
}
/// @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 field_collection.GetField(f); }
FieldCollection::FieldMap Fields() const { return field_collection.Fields(); }
Field* Fields(const std::string &f) const { return field_collection.GetField(f); }
Array<Field*>& InputFields() const { return field_collection.InputFields(); }
Array<Field*>& OutputFields() const { return field_collection.OutputFields(); }
Field* InputField(int i) const { return field_collection.InputField(i); }
Field* OutputField(int i) const { return field_collection.OutputField(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)), ...);
}
virtual void Save (std::ostream &out) const
{
out << "\"Node-" << id << "\" : " << std::endl;
out << "{\n";
out << "\"Name\": \"" << name << "\",\n";
field_collection.Save(out);
out << "}";
}
virtual void Mult(const Vector &x, Vector &y) const override
{
MFEM_ABORT("GraphNode::Mult() not implemented");
}
virtual void Mult(const MultiVector &x, MultiVector &y) override
{
MFEM_ABORT("GraphNode::Mult(MultiVector) not implemented");
}
using Operator::GetGradient;
// TODO: Possibly remove this and only support MultiVector version of GradientMult
virtual void GradientMult(const Vector &x, const Vector &dx, Vector &dy) const
{
MFEM_ABORT("GraphNode::GradientMult() not implemented");
GetGradient(x).Mult(dx, dy);
}
virtual void GradientMult(const MultiVector &x, const MultiVector &dx, MultiVector &dy) const
{
MFEM_ABORT("GraphNode::GradientMult() not implemented");
GetGradient(x).Mult(dx, dy);
}
// TODO: Possibly remove this and only support MultiVector version of GradientMultTranspose
virtual void GradientMultTranspose(const Vector &x, const Vector &dx, Vector &dy) const
{
MFEM_ABORT("GraphNode::GradientMultTranspose() not implemented");
GetGradient(x).MultTranspose(dx, dy);
}
virtual void GradientMultTranspose(const MultiVector &x, const MultiVector &dx, MultiVector &dy) const
{
MFEM_ABORT("GraphNode::GradientMultTranspose() not implemented");
// GetGradient(x).MultTranspose(dx, dy); // Not yet implemented
}
/// @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 ~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
};
enum InputType
{
VECTOR, ///< Asemble the input blockvector from individual fields
MULTIVECTOR, ///< Asemble the multivector from individual fields
NONE ///< No input
};
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
InputType input_type = InputType::MULTIVECTOR; ///< Input type for the graph
mutable Vector x_node, y_node; ///< Temporary vectors for evaluating nodes
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;
}
}
void SetInputType(InputType type) { input_type = type; }
InputType GetInputType() const { return input_type; }
/**
@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 Mult(const MultiVector &x, MultiVector &y) 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
{
public:
using InputType = DAGraph::InputType;
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 Vector x0, dx, dy;
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 Mult(const MultiVector &x, MultiVector &y) override;
void MultTranspose(const Vector &x, Vector &y) const override;
void MultTranspose(const MultiVector &x, MultiVector &y);
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