257 lines
8.2 KiB
C++
257 lines
8.2 KiB
C++
// Copyright (c) 2010-2025, Lawrence Livermore National Security, LLC. Produced
|
|
// at the Lawrence Livermore National Laboratory. All Rights reserved. See files
|
|
// LICENSE and NOTICE for details. LLNL-CODE-806117.
|
|
//
|
|
// This file is part of the MFEM library. For more information and source code
|
|
// availability visit https://mfem.org.
|
|
//
|
|
// MFEM is free software; you can redistribute it and/or modify it under the
|
|
// terms of the BSD-3 license. We welcome feedback and contributions, see file
|
|
// CONTRIBUTING.md for details.
|
|
|
|
#include "doperator.hpp"
|
|
|
|
#include <algorithm>
|
|
|
|
#ifdef MFEM_USE_MPI
|
|
|
|
using namespace mfem;
|
|
using namespace mfem::future;
|
|
|
|
DifferentiableOperator::DifferentiableOperator(
|
|
const std::vector<FieldDescriptor> &infds,
|
|
const std::vector<FieldDescriptor> &outfds,
|
|
const ParMesh &mesh) :
|
|
Operator(),
|
|
mesh(mesh),
|
|
infds(infds),
|
|
outfds(outfds)
|
|
{
|
|
unionfds.clear();
|
|
unionfds.insert(unionfds.end(), infds.begin(), infds.end());
|
|
unionfds.insert(unionfds.end(), outfds.begin(), outfds.end());
|
|
std::sort(unionfds.begin(), unionfds.end());
|
|
auto last = std::unique(unionfds.begin(), unionfds.end());
|
|
unionfds.erase(last, unionfds.end());
|
|
|
|
infields_l.resize(infds.size());
|
|
for (size_t i = 0; i < infds.size(); i++)
|
|
{
|
|
infields_l[i] = new Vector(GetVSize(infds[i]));
|
|
}
|
|
|
|
infields_e.resize(infds.size());
|
|
}
|
|
|
|
void DifferentiableOperator::SetMultLevel(MultLevel level)
|
|
{
|
|
mult_level = level;
|
|
}
|
|
|
|
void DifferentiableOperator::Mult(const Vector &x, Vector &y) const
|
|
{
|
|
MFEM_ASSERT(!action_callbacks.empty(),
|
|
"no integrators have been set");
|
|
|
|
MFEM_ASSERT(dynamic_cast<const BlockVector*>(&x),
|
|
"x needs to be a BlockVector");
|
|
|
|
MFEM_ASSERT(dynamic_cast<const BlockVector*>(&y),
|
|
"y needs to be a BlockVector");
|
|
|
|
const auto &bx = static_cast<const BlockVector &>(x);
|
|
auto &by = static_cast<BlockVector &>(y);
|
|
|
|
Mult(bx, by);
|
|
}
|
|
|
|
void DifferentiableOperator::DisableTensorProductStructure(bool disable)
|
|
{
|
|
use_tensor_product_structure = !disable;
|
|
}
|
|
|
|
std::shared_ptr<DerivativeOperator> DifferentiableOperator::GetDerivative(
|
|
size_t derivative_id, const Vector &x)
|
|
{
|
|
MFEM_ASSERT(derivative_action_callbacks.find(derivative_id) !=
|
|
derivative_action_callbacks.end(),
|
|
"no derivative action has been found for ID " << derivative_id);
|
|
|
|
const size_t dfidx = FindIdx(derivative_id, infds);
|
|
|
|
// Get transpose callbacks
|
|
std::vector<derivative_action_t> transpose_callbacks;
|
|
auto it_transpose = daction_transpose_callbacks.find(derivative_id);
|
|
if (it_transpose != daction_transpose_callbacks.end())
|
|
{
|
|
transpose_callbacks = it_transpose->second;
|
|
}
|
|
|
|
// Get assemble callbacks
|
|
std::vector<assemble_derivative_sparsematrix_callback_t> assemble_sparse_cbs;
|
|
auto it_sparse = assemble_derivative_sparsematrix_callbacks.find(derivative_id);
|
|
if (it_sparse != assemble_derivative_sparsematrix_callbacks.end())
|
|
{
|
|
assemble_sparse_cbs = it_sparse->second;
|
|
}
|
|
|
|
std::vector<assemble_derivative_hypreparmatrix_callback_t> assemble_hypre_cbs;
|
|
auto it_hypre = assemble_derivative_hypreparmatrix_callbacks.find(
|
|
derivative_id);
|
|
if (it_hypre != assemble_derivative_hypreparmatrix_callbacks.end())
|
|
{
|
|
assemble_hypre_cbs = it_hypre->second;
|
|
}
|
|
|
|
std::vector<assemble_diagonal_callback_t> assemble_diag_cbs;
|
|
auto it_diag = assemble_diagonal_callbacks.find(derivative_id);
|
|
if (it_diag != assemble_diagonal_callbacks.end())
|
|
{
|
|
assemble_diag_cbs = it_diag->second;
|
|
}
|
|
|
|
// If setup callbacks are available, run them to populate the cache
|
|
if (derivative_setup_callbacks.find(derivative_id) !=
|
|
derivative_setup_callbacks.end())
|
|
{
|
|
// Get the direction field from x
|
|
MFEM_ASSERT(dynamic_cast<const BlockVector*>(&x),
|
|
"x needs to be a BlockVector");
|
|
const auto &bx = static_cast<const BlockVector &>(x);
|
|
|
|
// Prolong to L-space
|
|
prolongation(infds, bx, infields_l);
|
|
|
|
// Restrict to element space
|
|
restriction<Entity::Element>(infds, infields_l, infields_e);
|
|
|
|
// Run all setup callbacks to populate the cache
|
|
for (const auto &setup_callback : derivative_setup_callbacks[derivative_id])
|
|
{
|
|
setup_callback(infields_e);
|
|
}
|
|
|
|
return std::make_shared<DerivativeOperator>(
|
|
height,
|
|
GetTrueVSize(infds[dfidx]),
|
|
derivative_apply_callbacks[derivative_id],
|
|
transpose_callbacks,
|
|
infds[dfidx],
|
|
x,
|
|
infds,
|
|
outfds,
|
|
assemble_sparse_cbs,
|
|
assemble_hypre_cbs,
|
|
assemble_diag_cbs);
|
|
}
|
|
else
|
|
{
|
|
// Fallback to old behavior (direct action callbacks)
|
|
return std::make_shared<DerivativeOperator>(
|
|
height,
|
|
GetTrueVSize(infds[dfidx]),
|
|
derivative_action_callbacks[derivative_id],
|
|
transpose_callbacks,
|
|
infds[dfidx],
|
|
x,
|
|
infds,
|
|
outfds,
|
|
assemble_sparse_cbs,
|
|
assemble_hypre_cbs,
|
|
assemble_diag_cbs);
|
|
}
|
|
}
|
|
|
|
std::shared_ptr<DerivativeOperator> DifferentiableOperator::GetDerivative(
|
|
size_t derivative_id, const MultiVector &x)
|
|
{
|
|
MFEM_ASSERT(derivative_action_callbacks.find(derivative_id) !=
|
|
derivative_action_callbacks.end(),
|
|
"no derivative action has been found for ID " << derivative_id);
|
|
|
|
const size_t dfidx = FindIdx(derivative_id, infds);
|
|
|
|
// Get transpose callbacks
|
|
std::vector<derivative_action_t> transpose_callbacks;
|
|
auto it_transpose = daction_transpose_callbacks.find(derivative_id);
|
|
if (it_transpose != daction_transpose_callbacks.end())
|
|
{
|
|
transpose_callbacks = it_transpose->second;
|
|
}
|
|
|
|
// Get assemble callbacks
|
|
std::vector<assemble_derivative_sparsematrix_callback_t> assemble_sparse_cbs;
|
|
auto it_sparse = assemble_derivative_sparsematrix_callbacks.find(derivative_id);
|
|
if (it_sparse != assemble_derivative_sparsematrix_callbacks.end())
|
|
{
|
|
assemble_sparse_cbs = it_sparse->second;
|
|
}
|
|
|
|
std::vector<assemble_derivative_hypreparmatrix_callback_t> assemble_hypre_cbs;
|
|
auto it_hypre = assemble_derivative_hypreparmatrix_callbacks.find(
|
|
derivative_id);
|
|
if (it_hypre != assemble_derivative_hypreparmatrix_callbacks.end())
|
|
{
|
|
assemble_hypre_cbs = it_hypre->second;
|
|
}
|
|
|
|
std::vector<assemble_diagonal_callback_t> assemble_diag_cbs;
|
|
auto it_diag = assemble_diagonal_callbacks.find(derivative_id);
|
|
if (it_diag != assemble_diagonal_callbacks.end())
|
|
{
|
|
assemble_diag_cbs = it_diag->second;
|
|
}
|
|
|
|
// If setup callbacks are available, run them to populate the cache
|
|
if (derivative_setup_callbacks.find(derivative_id) !=
|
|
derivative_setup_callbacks.end())
|
|
{
|
|
// Prolong to L-space
|
|
prolongation(infds, x, infields_l);
|
|
|
|
// Get the direction field
|
|
Vector direction_l(GetVSize(infds[dfidx]));
|
|
prolongation(infds[dfidx], x[dfidx], direction_l);
|
|
|
|
// Restrict to element space
|
|
restriction<Entity::Element>(infds, infields_l, infields_e);
|
|
|
|
// Run all setup callbacks to populate the cache
|
|
for (const auto &setup_callback : derivative_setup_callbacks[derivative_id])
|
|
{
|
|
setup_callback(infields_e);
|
|
}
|
|
|
|
return std::make_shared<DerivativeOperator>(
|
|
height,
|
|
GetTrueVSize(infds[dfidx]),
|
|
derivative_apply_callbacks[derivative_id],
|
|
transpose_callbacks,
|
|
infds[dfidx],
|
|
x,
|
|
infds,
|
|
outfds,
|
|
assemble_sparse_cbs,
|
|
assemble_hypre_cbs,
|
|
assemble_diag_cbs);
|
|
}
|
|
else
|
|
{
|
|
// Fallback to old behavior
|
|
return std::make_shared<DerivativeOperator>(
|
|
height,
|
|
GetTrueVSize(infds[dfidx]),
|
|
derivative_action_callbacks[derivative_id],
|
|
transpose_callbacks,
|
|
infds[dfidx],
|
|
x,
|
|
infds,
|
|
outfds,
|
|
assemble_sparse_cbs,
|
|
assemble_hypre_cbs,
|
|
assemble_diag_cbs);
|
|
}
|
|
}
|
|
|
|
#endif // MFEM_USE_MPI
|