diff --git a/mfem b/mfem index fb7466c81..8bfe86c4c 160000 --- a/mfem +++ b/mfem @@ -1 +1 @@ -Subproject commit fb7466c81925eb9a7b5614fa54d5a8bd9395ee36 +Subproject commit 8bfe86c4c682552173bc5df901fe9bdd811a58d7 diff --git a/src/smith/physics/dfem_weak_form.hpp b/src/smith/physics/dfem_weak_form.hpp index c28ea7029..d004e9221 100644 --- a/src/smith/physics/dfem_weak_form.hpp +++ b/src/smith/physics/dfem_weak_form.hpp @@ -247,39 +247,31 @@ class DfemWeakForm : public WeakForm { } /// @overload - void vjp(double /*time*/, double dt, ConstFieldPtr /*shape_disp*/, const std::vector& /*fields*/, - const std::vector& /*quad_fields*/, const std::vector& /*v_fields*/, - DualFieldPtr /*vjp_shape_disp_sensitivity*/, const std::vector& /*vjp_sensitivities*/, + void vjp(double /*time*/, double dt, ConstFieldPtr /*shape_disp*/, const std::vector& fields, + const std::vector& /*quad_fields*/, const std::vector& v_fields, + DualFieldPtr /*vjp_shape_disp_sensitivity*/, const std::vector& vjp_sensitivities, const std::vector& /*vjp_quad_field_sensitivities*/) const override { - SLIC_ERROR_ROOT("DfemWeakForm does not support vjp calculations"); - - // SLIC_ERROR_IF(vjp_sensitivities.size() != fields.size(), - // "Invalid number of field sensitivities relative to the number of fields"); - // SLIC_ERROR_IF(v_fields.size() != 1, "FunctionalResidual nonlinear systems only supports 1 output residual"); + SLIC_ERROR_IF(vjp_sensitivities.size() != fields.size(), + "Invalid number of field sensitivities relative to the number of fields"); + SLIC_ERROR_IF(v_fields.size() != 1, "FunctionalResidual nonlinear systems only supports 1 output residual"); dt_ = dt; - // TODO (EBC): add in a future PR... - // std::vector test_par_gf({&v_fields[0]->gridFunction()}); - // std::vector field_par_gf = getLVectors(fields); - // // field_par_gf.push_back(&v_fields[0]->gridFunction()); - - // for (size_t input_col = 0; input_col < fields.size(); ++input_col) { - // if (vjp_sensitivities[input_col] != nullptr) { - // auto deriv_op = v_dot_weak_form_residual_.GetDerivative(input_col, test_par_gf, field_par_gf); - // // do this entry by entry until assembly is supported - // mfem::Vector direction(vjp_sensitivities[input_col]->Size()); - // direction = 0.0; - // for (int i = 0; i < vjp_sensitivities[input_col]->Size(); ++i) { - // direction[i] = 1.0; - // mfem::Vector value(1); - // deriv_op->Mult(direction, value); - // (*vjp_sensitivities[input_col])[i] += value[0]; - // direction[i] = 0.0; - // } - // } - // } + std::vector test_par_gf({&v_fields[0]->gridFunction()}); + std::vector field_par_gf = getLVectors(fields); + // field_par_gf.push_back(&v_fields[0]->gridFunction()); + + for (size_t input_col = 0; input_col < fields.size(); ++input_col) { + if (vjp_sensitivities[input_col] != nullptr) { + auto deriv_op = v_dot_weak_form_residual_.GetDerivative(input_col, test_par_gf, field_par_gf); + mfem::Vector vec_jac_mfem_vector(vjp_sensitivities[input_col]->Size()); + vec_jac_mfem_vector = 0.0; + // NOTE: this is happening entry by entry inside this (temporary) dfem function + deriv_op->Assemble(vec_jac_mfem_vector); + (*vjp_sensitivities[input_col]) += vec_jac_mfem_vector; + } + } } protected: diff --git a/src/smith/physics/tests/CMakeLists.txt b/src/smith/physics/tests/CMakeLists.txt index 8f848d148..55f0205e6 100644 --- a/src/smith/physics/tests/CMakeLists.txt +++ b/src/smith/physics/tests/CMakeLists.txt @@ -52,7 +52,8 @@ endif() if (SMITH_USE_DFEM) list(APPEND physics_serial_test_sources - test_dfem_explicit_dynamics.cpp) + test_dfem_explicit_dynamics.cpp + test_dfem_weak_form.cpp) endif() smith_add_tests(SOURCES ${physics_serial_test_sources} diff --git a/src/smith/physics/tests/test_dfem_weak_form.cpp b/src/smith/physics/tests/test_dfem_weak_form.cpp new file mode 100644 index 000000000..850e3d09a --- /dev/null +++ b/src/smith/physics/tests/test_dfem_weak_form.cpp @@ -0,0 +1,290 @@ +// Copyright (c) Lawrence Livermore National Security, LLC and +// other Smith Project Developers. See the top-level LICENSE file for +// details. +// +// SPDX-License-Identifier: (BSD-3-Clause) + +#include +#include +#include +#include +#include + +#include "gtest/gtest.h" +#include "mpi.h" +#include "mfem.hpp" + +#include "smith/infrastructure/application_manager.hpp" +#include "smith/physics/mesh.hpp" +#include "smith/physics/state/state_manager.hpp" +#include "smith/physics/tests/physics_test_utils.hpp" +#include "smith/physics/dfem_weak_form.hpp" +#include "smith/physics/functional_weak_form.hpp" +#include "smith/numerics/functional/finite_element.hpp" // for H1 +#include "smith/numerics/functional/functional.hpp" +#include "smith/numerics/functional/tensor.hpp" +#include "smith/numerics/functional/tuple.hpp" +#include "smith/physics/common.hpp" +#include "smith/physics/field_types.hpp" +#include "smith/physics/state/finite_element_dual.hpp" +#include "smith/physics/state/finite_element_state.hpp" + +auto element_shape = mfem::Element::QUADRILATERAL; + +struct WeakFormsFixture : public testing::Test { + static constexpr int dim = 2; + static constexpr int disp_order = 1; + + using VectorSpace = smith::H1; + using DensitySpace = smith::L2; + + static constexpr int NUM_FUNCTIONAL_STATES = 2; + static constexpr int NUM_DFEM_STATES = 3; + + enum STATE + { + DISP, + VELO, + COORD + }; + + enum PAR + { + DENSITY + }; + + void SetUp() + { + MPI_Barrier(MPI_COMM_WORLD); + smith::StateManager::initialize(datastore, "solid_dynamics"); + + double length = 0.5; + double width = 2.0; + mesh = std::make_shared(mfem::Mesh::MakeCartesian2D(6, 20, element_shape, true, length, width), + "this_mesh_name", 0, 0); + + smith::FiniteElementState disp = smith::StateManager::newState(VectorSpace{}, "displacement", mesh->tag()); + smith::FiniteElementState velo = smith::StateManager::newState(VectorSpace{}, "velocity", mesh->tag()); + smith::FiniteElementState density = smith::StateManager::newState(DensitySpace{}, "density", mesh->tag()); + + shape_disp = std::make_unique(mesh->newShapeDisplacement()); + shape_disp_dual = std::make_unique(mesh->newShapeDisplacementDual()); + + states = {disp, velo}; + params = {density}; + + for (auto s : states) { + state_duals.push_back(smith::FiniteElementDual(s.space(), s.name() + "_dual")); + } + for (auto p : params) { + param_duals.push_back(smith::FiniteElementDual(p.space(), p.name() + "_dual")); + } + + state_tangents = states; + param_tangents = params; + + std::string physics_name = "fake_physics"; + + using TrialSpace = VectorSpace; + using FunctionalWeakFormT = + smith::FunctionalWeakForm>; + std::vector f_inputs{&states[DISP].space(), &states[VELO].space(), + ¶ms[DENSITY].space()}; + auto f_weak_form = std::make_shared(physics_name, mesh, states[DISP].space(), f_inputs); + + std::vector d_inputs{ + &states[DISP].space(), &states[VELO].space(), + static_cast(mesh->mfemParMesh().GetNodalFESpace())}; + auto d_weak_form = std::make_shared(physics_name, mesh, states[DISP].space(), d_inputs); + mfem::Array solid_attrib({1}); + constexpr int ir_order = 2; + const mfem::IntegrationRule& displacement_ir = mfem::IntRules.Get(disp.space().GetFE(0)->GetGeomType(), ir_order); + // none of the q functions are dependent on the parameter field + constexpr auto deriv_indices = std::make_index_sequence(); + + // body force 1: displacement + f_weak_form->addBodySource(smith::DependsOn<0>{}, mesh->entireBodyName(), + [](double /*t*/, auto /*x*/, auto u) { return u; }); + mfem::future::tuple, mfem::future::Value, mfem::future::Gradient, + mfem::future::Weight> + qfn1_inputs{}; + mfem::future::tuple> qfn1_outputs{}; + d_weak_form->addBodyIntegral( + solid_attrib, + [=] SMITH_HOST_DEVICE(const mfem::future::tensor& u, + const mfem::future::tensor&, + const mfem::future::tensor& dX_dxi, mfem::real_t weight) { + auto J = mfem::future::det(dX_dxi) * weight; + return mfem::future::tuple{-u * J}; + }, + qfn1_inputs, qfn1_outputs, displacement_ir, deriv_indices); + + // body force 2: position + f_weak_form->addBodySource(mesh->entireBodyName(), [](double /*t*/, auto x) { return 0.5 * x; }); + mfem::future::tuple, mfem::future::Value, mfem::future::Value, + mfem::future::Gradient, mfem::future::Weight> + qfn2_inputs{}; + mfem::future::tuple> qfn2_outputs{}; + d_weak_form->addBodyIntegral( + solid_attrib, + [=] SMITH_HOST_DEVICE(const mfem::future::tensor&, + const mfem::future::tensor&, + const mfem::future::tensor& X, + const mfem::future::tensor& dX_dxi, mfem::real_t weight) { + auto J = mfem::future::det(dX_dxi) * weight; + return mfem::future::tuple{-0.5 * X * J}; + }, + qfn2_inputs, qfn2_outputs, displacement_ir, deriv_indices); + + // initialize fields for testing + + for (auto& s : state_tangents) { + pseudoRand(s); + } + for (auto& p : param_tangents) { + pseudoRand(p); + } + + state_duals[DISP] = 1.0; + state_duals[VELO] = 2.0; + param_duals[DENSITY] = 3.0; + + states[DISP].setFromFieldFunction([](smith::tensor x) { + auto u = 0.1 * x; + return u; + }); + + states[VELO].setFromFieldFunction([](smith::tensor x) { + auto u = -0.01 * x; + return u; + }); + + params[DENSITY] = 1.2; + + // functional_weak_form and dfem_weak_form are abstract WeakForm classes to ensure usage only through WeakForm + // interface + functional_weak_form = f_weak_form; + dfem_weak_form = d_weak_form; + } + + const double time = 0.0; + const double dt = 1.0; + + std::string velo_name = "solid_velocity"; + + axom::sidre::DataStore datastore; + std::shared_ptr mesh; + std::shared_ptr functional_weak_form; + std::shared_ptr dfem_weak_form; + + std::unique_ptr shape_disp; + std::unique_ptr shape_disp_dual; + + std::vector states; + std::vector params; + + std::vector state_duals; + std::vector param_duals; + + std::vector state_tangents; + std::vector param_tangents; +}; + +TEST_F(WeakFormsFixture, VjpConsistency) +{ + // initialize the displacement and acceleration to a non-trivial field + auto f_input_fields = getConstFieldPointers(states, params); + smith::FiniteElementState coords( + *static_cast(mesh->mfemParMesh().GetNodalFESpace()), "coords"); + coords.setFromGridFunction(*static_cast(mesh->mfemParMesh().GetNodes())); + std::vector d_input_fields = {&states[DISP], &states[VELO], &coords}; + + smith::FiniteElementDual f_res_vector(states[DISP].space(), "functional_residual"); + f_res_vector = functional_weak_form->residual(time, dt, shape_disp.get(), f_input_fields); + smith::FiniteElementDual d_res_vector(states[DISP].space(), "dfem_residual"); + d_res_vector = dfem_weak_form->residual(time, dt, shape_disp.get(), d_input_fields); + EXPECT_NEAR(f_res_vector.Norml2(), d_res_vector.Norml2(), 1.e-12); + + // test vjp + smith::FiniteElementState v(f_res_vector.space(), "v"); + pseudoRand(v); + auto f_field_vjps = getFieldPointers(state_duals, param_duals); + + smith::FiniteElementDual coords_dual(coords.space(), coords.name() + "_dual"); + coords_dual = 0.0; + std::vector field_vjps_copy = {state_duals[DISP], state_duals[VELO], + std::move(coords_dual)}; + auto d_field_vjps = getFieldPointers(field_vjps_copy); + + dfem_weak_form->vjp(time, dt, shape_disp.get(), d_input_fields, {}, getConstFieldPointers(v), shape_disp_dual.get(), + d_field_vjps, {}); + functional_weak_form->vjp(time, dt, shape_disp.get(), f_input_fields, {}, getConstFieldPointers(v), + shape_disp_dual.get(), f_field_vjps, {}); + + for (size_t i = 0; i < NUM_FUNCTIONAL_STATES; ++i) { + EXPECT_NEAR(f_field_vjps[i]->Norml2(), d_field_vjps[i]->Norml2(), 1e-12) + << " " << d_field_vjps[i]->name() << std::endl; + } +} + +TEST_F(WeakFormsFixture, JvpConsistency) +{ + // initialize the displacement and acceleration to a non-trivial field + auto input_fields = getConstFieldPointers(states, params); + + smith::FiniteElementDual res_vector(states[DISP].space(), "residual"); + res_vector = functional_weak_form->residual(time, dt, shape_disp.get(), input_fields); + ASSERT_NE(0.0, res_vector.Norml2()); + + auto jacobianWeights = [&](size_t i) { + std::vector tangents(input_fields.size()); + tangents[i] = 1.0; + return tangents; + }; + + auto selectStates = [&](size_t i) { + auto field_tangents = getConstFieldPointers(state_tangents, param_tangents); + for (size_t j = 0; j < field_tangents.size(); ++j) { + if (i != j) { + field_tangents[j] = nullptr; + } + } + return field_tangents; + }; + + smith::FiniteElementDual jvp_slow(states[DISP].space(), "jvp_slow"); + smith::FiniteElementDual jvp(states[DISP].space(), "jvp"); + jvp = 4.0; // set to some value to test jvp resets these values + auto jvps = getFieldPointers(jvp); + + auto field_tangents = getConstFieldPointers(state_tangents, param_tangents); + + for (size_t i = 0; i < input_fields.size(); ++i) { + auto J = functional_weak_form->jacobian(time, dt, shape_disp.get(), input_fields, jacobianWeights(i)); + J->Mult(*field_tangents[i], jvp_slow); + functional_weak_form->jvp(time, dt, shape_disp.get(), input_fields, {}, nullptr, selectStates(i), {}, jvps); + EXPECT_NEAR(jvp_slow.Norml2(), jvp.Norml2(), 1e-12); + } + + // test jacobians in weighted combinations + { + field_tangents[NUM_FUNCTIONAL_STATES] = nullptr; + + double velo_factor = 0.2; + std::vector jacobian_weights = {1.0, velo_factor, 0.0}; + + auto J = functional_weak_form->jacobian(time, dt, shape_disp.get(), input_fields, jacobian_weights); + J->Mult(*field_tangents[DISP], jvp_slow); + + state_tangents[VELO] *= velo_factor; + functional_weak_form->jvp(time, dt, shape_disp.get(), input_fields, {}, nullptr, field_tangents, {}, jvps); + EXPECT_NEAR(jvp_slow.Norml2(), jvp.Norml2(), 1e-12); + } +} + +int main(int argc, char* argv[]) +{ + ::testing::InitGoogleTest(&argc, argv); + smith::ApplicationManager applicationManager(argc, argv); + return RUN_ALL_TESTS(); +}