Compare commits
31
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
c1ea8aefb7 | ||
|
|
04444dc37e | ||
|
|
6f204dde67 | ||
|
|
f04f98752b | ||
|
|
9b31004ce5 | ||
|
|
04571c5ddd | ||
|
|
e448b71831 | ||
|
|
48e415e17e | ||
|
|
8b249e8004 | ||
|
|
1585e7d9f1 | ||
|
|
ae4a2ee9a6 | ||
|
|
14e8c6ce45 | ||
|
|
05f5967267 | ||
|
|
d9e0018e98 | ||
|
|
f08dba7b84 | ||
|
|
7e75c7f6fb | ||
|
|
a1962fa492 | ||
|
|
ce7fefb6eb | ||
|
|
503b286f98 | ||
|
|
cd0616fae0 | ||
|
|
1ff144a358 | ||
|
|
2c02d1b3cb | ||
|
|
f77dc8d7e9 | ||
|
|
c23b850b77 | ||
|
|
ebda52d76c | ||
|
|
c8ab3cbf69 | ||
|
|
1ed3b48c2e | ||
|
|
fbd9189e7b | ||
|
|
1dd889cb16 | ||
|
|
2e8fbd661a | ||
|
|
6e424dba6e |
+79
-30
@@ -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
@@ -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.
|
||||
|
||||
@@ -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. */
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -15,6 +15,7 @@
|
||||
// Linear algebra header file
|
||||
|
||||
#include "vector.hpp"
|
||||
#include "multivector.hpp"
|
||||
#include "operator.hpp"
|
||||
#include "matrix.hpp"
|
||||
#include "sparsemat.hpp"
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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,
|
||||
|
||||
@@ -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. */
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
@@ -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;
|
||||
}
|
||||
@@ -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
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user