diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index 016c7c66f8..6f6e84e35d 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -249,3 +249,5 @@ endif() if(MFEM_USE_MOONOLITH) add_subdirectory(moonolith) endif() + +add_subdirectory(dfem) diff --git a/examples/dfem/CMakeLists.txt b/examples/dfem/CMakeLists.txt new file mode 100644 index 0000000000..ae2e70bcb5 --- /dev/null +++ b/examples/dfem/CMakeLists.txt @@ -0,0 +1,117 @@ +# 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(DFEM_EXAMPLES_SRCS) + +if (MFEM_USE_MPI) + list(APPEND DFEM_EXAMPLES_SRCS + diffusion3d.cpp + plasticity.cpp + laghos.cpp + ) +endif() + +# Include the source directory where mfem.hpp and mfem-performance.hpp are. +include_directories(BEFORE ${PROJECT_BINARY_DIR}) + +# Add "test_dfem" target, see below. +add_custom_target(test_dfem + ${CMAKE_CTEST_COMMAND} -R dfem USES_TERMINAL) + +# Add one executable per cpp file, adding "dfem_" as prefix so the CMake +# target is unique from those in the non-dFEM examples. Also sets +# "test_dfem" as a target that depends on the given dFEM examples. +set(PFX dfem_) +add_mfem_examples(DFEM_EXAMPLES_SRCS ${PFX} "" test_dfem) + +# Remove "dfem_" prefix from exectuable name for consistency with GNU build +# system. +foreach(SRC_FILE ${DFEM_EXAMPLES_SRCS}) + get_filename_component(SRC_FILENAME ${SRC_FILE} NAME) + string(REPLACE ".cpp" "" TARGET_NAME "${PFX}${SRC_FILENAME}") + string(REPLACE ${PFX} "" EXE_NAME ${TARGET_NAME}) + set_target_properties(${TARGET_NAME} PROPERTIES OUTPUT_NAME ${EXE_NAME}) +endforeach() + +# Testing. +# The dFEM tests can be run separately using the target "test_dfem" +# which builds the examples and runs: +# ctest -R dfem + +if (MFEM_ENABLE_TESTING) + # Command line options for the tests. + # Example 9: test CVODE with CV_ADAMS (non-stiff implicit) time stepping + # set(EX9_COMMON_OPTS -m ../../data/periodic-hexagon.mesh -p 0 -s 7) + # set(EX9_TEST_OPTS ${EX9_COMMON_OPTS} -r 2 -dt 0.0018 -vs 25) + # set(EX9P_TEST_OPTS ${EX9_COMMON_OPTS} -rp 1 -dt 0.0009 -vs 50) + # Example 10: test CVODE with CV_BDF (stiff implicit) time stepping + # set(EX10_COMMON_OPTS -m ../../data/beam-quad.mesh -o 2 -s 5 -dt 0.15 -tf 6 -vs 10) + # set(EX10_TEST_OPTS ${EX10_COMMON_OPTS} -r 2) + # set(EX10P_TEST_OPTS ${EX10_COMMON_OPTS} -rp 1) + # Example 16: test ARKODE with implicit time stepping using mass form + # set(EX16_COMMON_OPTS -s 15) + # set(EX16_TEST_OPTS ${EX16_COMMON_OPTS}) + # set(EX16P_TEST_OPTS ${EX16_COMMON_OPTS}) + + # Add the tests: one test per source file. + foreach(SRC_FILE ${DFEM_EXAMPLES_SRCS}) + get_filename_component(SRC_FILENAME ${SRC_FILE} NAME) + string(REPLACE ".cpp" "" TEST_NAME ${SRC_FILENAME}) + string(TOUPPER ${TEST_NAME} UP_TEST_NAME) + set(TEST_NAME ${PFX}${TEST_NAME}) + + set(THIS_TEST_OPTIONS "-no-vis") + list(APPEND THIS_TEST_OPTIONS ${${UP_TEST_NAME}_TEST_OPTS}) + # message(STATUS "Test ${TEST_NAME} options: ${THIS_TEST_OPTIONS}") + + if (NOT (${TEST_NAME} MATCHES ".*p$")) + add_test(NAME ${TEST_NAME}_ser + COMMAND ${TEST_NAME} ${THIS_TEST_OPTIONS}) + else() + add_test(NAME ${TEST_NAME}_np=${MFEM_MPI_NP} + COMMAND ${MPIEXEC} ${MPIEXEC_NUMPROC_FLAG} ${MFEM_MPI_NP} + ${MPIEXEC_PREFLAGS} + $ ${THIS_TEST_OPTIONS} + ${MPIEXEC_POSTFLAGS}) + endif() + endforeach() + + # Add CUDA/HIP tests. + set(DEVICE_EXAMPLES + # parallel examples with device support: + # ex9p + ) + set(MFEM_TEST_DEVICE) + if (MFEM_USE_CUDA) + set(MFEM_TEST_DEVICE "cuda") + elseif (MFEM_USE_HIP) + set(MFEM_TEST_DEVICE "hip") + endif() + if (MFEM_TEST_DEVICE) + foreach(TEST_NAME ${DEVICE_EXAMPLES}) + string(TOUPPER ${TEST_NAME} UP_TEST_NAME) + + set(THIS_TEST_OPTIONS "-no-vis" "-d" "${MFEM_TEST_DEVICE}") + list(APPEND THIS_TEST_OPTIONS ${${UP_TEST_NAME}_TEST_OPTS}) + + if (NOT (${TEST_NAME} MATCHES ".*p$")) + add_test(NAME ${PFX}${TEST_NAME}_${MFEM_TEST_DEVICE}_ser + COMMAND ${PFX}${TEST_NAME} ${THIS_TEST_OPTIONS}) + else() + add_test(NAME ${PFX}${TEST_NAME}_${MFEM_TEST_DEVICE}_np=${MFEM_MPI_NP} + COMMAND ${MPIEXEC} ${MPIEXEC_NUMPROC_FLAG} ${MFEM_MPI_NP} + ${MPIEXEC_PREFLAGS} + $ ${THIS_TEST_OPTIONS} + ${MPIEXEC_POSTFLAGS}) + endif() + endforeach() + endif(MFEM_TEST_DEVICE) +endif(MFEM_ENABLE_TESTING) diff --git a/examples/dfem/diffusion3d.cpp b/examples/dfem/diffusion3d.cpp new file mode 100644 index 0000000000..03677061b0 --- /dev/null +++ b/examples/dfem/diffusion3d.cpp @@ -0,0 +1,428 @@ +#include + +#include "fem/qinterp/det.cpp" +#include "fem/qinterp/grad.hpp" // IWYU pragma: keep + +#include +#include +#include + +#include +#include +#include +#include + +#include "fem/dfem/kernels_regs.hpp" + +using namespace mfem; + +using mfem::future::tuple; +using mfem::future::tensor; + +using future::DifferentiableOperator; +using future::ParametricSpace; +using future::ParametricFunction; +using future::FieldDescriptor; +using future::Gradient; +using future::Weight; +using future::None; + +#undef NVTX_COLOR +#define NVTX_COLOR nvtx::kAquamarine +#include "general/nvtx.hpp" + +static int gD1D = 0, gQ1D = 0; + +/////////////////////////////////////////////////////////////////////////////// +struct StiffnessIntegrator : public BilinearFormIntegrator +{ + const FiniteElementSpace *fes; + const real_t *B, *G, *DX; + int ne, d1d, q1d; + Vector J0, dx; + +public: + StiffnessIntegrator() + { + dbg(); + StiffnessKernels::Specialization<2, 3>::Add(); + StiffnessKernels::Specialization<3, 5>::Add(); + StiffnessKernels::Specialization<4, 8>::Add(); + StiffnessKernels::Specialization<5, 10>::Add(); + StiffnessKernels::Specialization<7, 15>::Add(); + } + + void AssemblePA(const FiniteElementSpace &fespace) override + { + fes = &fespace; + auto *mesh = fes->GetMesh(); + const int DIM = mesh->Dimension(); + ne = mesh->GetNE(); + const auto p = fes->GetFE(0)->GetOrder(); + const auto q = 2 * p + mesh->GetElementTransformation(0)->OrderW(); + const auto type = mesh->GetElementBaseGeometry(0); + const IntegrationRule &ir = IntRules.Get(type, q); + const int NQPT = ir.GetNPoints(); + d1d = p + 1; + q1d = IntRules.Get(Geometry::SEGMENT, ir.GetOrder()).GetNPoints(); + MFEM_VERIFY(d1d == gD1D, "D1D mismatch: " << d1d << " != " << gD1D); + MFEM_VERIFY(q1d == gQ1D, "Q1D mismatch: " << q1d << " != " << gQ1D); + MFEM_VERIFY(NQPT == q1d * q1d * q1d, ""); + const DofToQuad *maps = &fes->GetFE(0)->GetDofToQuad(ir, DofToQuad::TENSOR); + const GridFunction *nodes = (mesh->EnsureNodes(), mesh->GetNodes()); + const FiniteElementSpace *nfes = nodes->FESpace(); + const int nVDIM = nfes->GetVDim(); + dx.SetSize(nVDIM * DIM * NQPT * ne, Device::GetDeviceMemoryType()); + J0.SetSize(nVDIM * DIM * NQPT * ne, Device::GetDeviceMemoryType()); + dx.UseDevice(true), J0.UseDevice(true); + B = maps->B.Read(), G = maps->G.Read(), DX = dx.Read(); + + const Operator *NR = nfes->GetElementRestriction( + ElementDofOrdering::LEXICOGRAPHIC); + const QuadratureInterpolator *nqi = nfes->GetQuadratureInterpolator(ir); + nqi->SetOutputLayout(QVectorLayout::byVDIM); + const int nd = nfes->GetFE(0)->GetDof(); + Vector xe(nVDIM * nd * ne, Device::GetDeviceMemoryType()); + NR->Mult(*nodes, (xe.UseDevice(true), xe)); + nqi->Derivatives(xe, J0); + + const auto w_r = ir.GetWeights().Read(); + const auto W = Reshape(w_r, q1d, q1d, q1d); + const auto J = Reshape(J0.Read(), 3, 3, q1d, q1d, q1d, ne); + auto DX_w = Reshape(dx.Write(), 3, 3, q1d, q1d, q1d, ne); + mfem::forall_3D(ne, q1d, q1d, q1d, [=] MFEM_HOST_DEVICE(int e) + { + MFEM_FOREACH_THREAD1(qz, z, q1d) + { + MFEM_FOREACH_THREAD1(qy, y, q1d) + { + MFEM_FOREACH_THREAD1(qx, q, q1d) + { + const real_t w = W(qx, qy, qz); + const real_t *Jtr = &J(0, 0, qx, qy, qz, e); + const real_t detJ = kernels::Det<3>(Jtr); + const real_t wd = w * detJ; + real_t Jrt[9], A[9], D[9] = + { + wd, 0.0, 0.0, + 0.0, wd, 0.0, + 0.0, 0.0, wd + }; + kernels::CalcInverse<3>(Jtr, Jrt); + kernels::MultABt(3, 3, 3, D, Jrt, A); + kernels::Mult(3, 3, 3, A, Jrt, &DX_w(0, 0, qx, qy, qz, e)); + } + } + } + MFEM_SYNC_THREAD; + }); + } + + template + static void StiffnessMult(const int NE, + const real_t *b, const real_t *g, + const real_t *dx, + const real_t *xe, real_t *ye, + const int d1d, const int q1d) + { + const int D1D = T_D1D ? T_D1D : d1d; + const int Q1D = T_Q1D ? T_Q1D : q1d; + + constexpr int DIM = 3, VDIM = 1; + const auto XE = Reshape(xe, D1D, D1D, D1D, VDIM, NE); + const auto DX = Reshape(dx, 3, 3, Q1D, Q1D, Q1D, NE); + auto YE = Reshape(ye, D1D, D1D, D1D, VDIM, NE); + + mfem::forall_2D(NE, Q1D, Q1D, [=] MFEM_HOST_DEVICE(int e) + { + MFEM_SHARED real_t smem[MQ1][MQ1]; + MFEM_SHARED real_t sB[MD1][MQ1], sG[MD1][MQ1]; + regs5d_t r0, r1; + + LoadMatrix(D1D, Q1D, b, sB); + LoadMatrix(D1D, Q1D, g, sG); + + LoadDofs3d(e, D1D, XE, r0); + Grad3d(D1D, Q1D, smem, sB, sG, r0, r1); + + for (int qz = 0; qz < Q1D; qz++) + { + MFEM_FOREACH_THREAD1(qy, y, Q1D) + { + MFEM_FOREACH_THREAD1(qx, x, Q1D) + { + real_t v[3], u[3] = { r1[0][0][qz][qy][qx], + r1[0][1][qz][qy][qx], + r1[0][2][qz][qy][qx] + }; + const real_t *dx = &DX(0, 0, qx, qy, qz, e); + kernels::Mult(3, 3, dx, u, v); + r0[0][0][qz][qy][qx] = v[0]; + r0[0][1][qz][qy][qx] = v[1]; + r0[0][2][qz][qy][qx] = v[2]; + } + } + } + GradTranspose3d(D1D, Q1D, smem, sB, sG, r0, r1); + WriteDofs3d(e, D1D, r1, YE); + }); + } + + using StiffnessKernelType = decltype(&StiffnessMult<1,1>); + MFEM_REGISTER_KERNELS(StiffnessKernels, StiffnessKernelType, (int, int)); + + void AddMultPA(const Vector &x, Vector &y) const override + { + StiffnessKernels::Run(d1d, q1d, ne, B, G, DX, x.Read(), y.ReadWrite(), + d1d, q1d); + } +}; + +template +StiffnessIntegrator::StiffnessKernelType +StiffnessIntegrator::StiffnessKernels::Kernel() +{ + return StiffnessMult; +} + +StiffnessIntegrator::StiffnessKernelType +StiffnessIntegrator::StiffnessKernels::Fallback(int d1d, int q1d) +{ + dbg("\x1b[33mFallback d1d:{} q1d:{}", d1d, q1d); + return StiffnessMult; +} + +/////////////////////////////////////////////////////////////////////////////// +void AddKernelSpecializations() +{ + dbg(); + using Det = QuadratureInterpolator::DetKernels; + Det::Specialization<3, 3, 2, 2>::Add(); + Det::Specialization<3, 3, 4, 4>::Add(); + + using Grad = QuadratureInterpolator::GradKernels; + Grad::Specialization<3, QVectorLayout::byVDIM, false, 3, 2, 3>::Add(); + Grad::Specialization<3, QVectorLayout::byVDIM, false, 3, 3, 5>::Add(); + Grad::Specialization<3, QVectorLayout::byNODES, false, 3, 4, 5>::Add(); + Grad::Specialization<3, QVectorLayout::byVDIM, false, 3, 4, 8>::Add(); +} + +/////////////////////////////////////////////////////////////////////////////// +int main(int argc, char* argv[]) +{ + constexpr int DIM = 3; + + Mpi::Init(); + AddKernelSpecializations(); + + const char* device_config = "cpu"; + int version = 0; + int order = 1; + int refinements = 1; + bool visualization = true; + + OptionsParser args(argc, argv); + args.AddOption(&version, "-v", "--version", ""); + args.AddOption(&order, "-o", "--order", ""); + args.AddOption(&refinements, "-r", "--refinements", ""); + args.AddOption(&device_config, "-d", "--device", + "Device configuration string, see Device::Configure()."); + args.AddOption(&visualization, "-vis", "--visualization", "-no-vis", + "--no-visualization", "Enable or disable GLVis visualization."); + args.ParseCheck(); + + Device device(device_config); + if (Mpi::Root() == 0) + { + device.Print(); + } + + out << std::setprecision(8); + + Mesh smesh = Mesh::MakeCartesian3D(4, 4, 4, Element::HEXAHEDRON); + smesh.EnsureNodes(); + MFEM_ASSERT(smesh.Dimension() == DIM, "incorrect mesh dimension"); + + for (int i = 0; i < refinements; i++) + { + smesh.UniformRefinement(); + } + + ParMesh pmesh(MPI_COMM_WORLD, smesh); + pmesh.SetCurvature(order); + smesh.Clear(); + + out << "#el: " << pmesh.GetNE() << "\n"; + + auto* nodes = static_cast(pmesh.GetNodes()); + ParFiniteElementSpace& mfes = *nodes->ParFESpace(); + + H1_FECollection fec(order, DIM); + ParFiniteElementSpace fes(&pmesh, &fec); + + const auto p = fes.GetFE(0)->GetOrder(); + const auto q = 2 * p + pmesh.GetElementTransformation(0)->OrderW(); + const auto type = pmesh.GetElementBaseGeometry(0); + const IntegrationRule &ir = IntRules.Get(type, q); + gD1D = p + 1; + gQ1D = IntRules.Get(Geometry::SEGMENT, ir.GetOrder()).GetNPoints(); + dbg("D1D: {}, Q1D: {}", gD1D, gQ1D); + + const int NE = pmesh.GetNE(); + const int NQPT = ir.GetNPoints(); + + ParGridFunction x(&fes), y(&fes); + + Array ess_tdof_list, ess_bdr(pmesh.bdr_attributes.Max()); + ess_bdr = 1; + fes.GetEssentialTrueDofs(ess_bdr, ess_tdof_list); + + ConstantCoefficient one(1.0); + + ParLinearForm b(&fes); + b.AddDomainIntegrator(new DomainLFIntegrator(one)); + b.UseFastAssembly(true); + b.Assemble(); + + std::unique_ptr a; + std::unique_ptr ∂op; + + + const int elem_size = DIM * DIM * NQPT; + const int total_size = elem_size * NE; + dbg("DIM: {}, local_size: {}, elem_size: {}, total_size: {}", + DIM, DIM * DIM, elem_size, total_size); + ParametricSpace qdata_space(DIM, DIM * DIM, elem_size, total_size, gD1D, gQ1D); + ParametricFunction qd(qdata_space); + + if (version < 2) + { + a = std::make_unique(&fes); + a->SetAssemblyLevel(AssemblyLevel::PARTIAL); + if (version == 0) { a->AddDomainIntegrator(new DiffusionIntegrator(&ir)); } + if (version == 1) { a->AddDomainIntegrator(new StiffnessIntegrator()); } + a->Assemble(); + if (version == 0) + { + BilinearFormIntegrator *bfi = a->GetDBFI()->operator[](0); + auto *di = dynamic_cast(bfi); + assert(di); + const int d1d = di->dofs1D, q1d = di->quad1D; + dbg("\x1b[33md1d: {} q1d: {}", d1d, q1d); + MFEM_VERIFY(d1d == gD1D, "D1D mismatch: " << d1d << " != " << gD1D); + MFEM_VERIFY(q1d == gQ1D, "Q1D mismatch: " << q1d << " != " << gQ1D); + } + } + else if (version == 2) // MF ∂fem + { + constexpr int U = 0, Ξ = 1; + auto solutions = std::vector{FieldDescriptor{U, &fes}}; + auto parameters = std::vector{FieldDescriptor{Ξ, &mfes}}; + auto diffusion_mf_kernel = + [] MFEM_HOST_DEVICE (const tensor& ∇u, + const tensor& J, + const real_t& w) + { + auto invJ = inv(J); + return tuple{((∇u * invJ)) * transpose(invJ) * det(J) * w}; + }; + ∂op = std::make_unique(solutions, parameters, pmesh); + ∂op->SetParameters({nodes}); + ∂op->AddDomainIntegrator(diffusion_mf_kernel, + tuple{Gradient{}, Gradient<Ξ>{}, Weight{}}, + tuple{Gradient{}}, + ir, ess_bdr); + } + else if (version == 3) // PA ∂fem + { + constexpr int U = 0, Ξ = 1, Q = 2; + FieldDescriptor u_fd{U, &fes}, Ξ_fd{Ξ, &mfes}, q_fd{Q, &qd.space}; + auto w = Weight{}; + auto q = None {}; + auto u = None {}; + auto ∇u = Gradient {}; + auto ∇Ξ = Gradient<Ξ> {}; + auto u_sol = std::vector{u_fd}, + q_param = std::vector{q_fd}, + Ξ_q_params = std::vector{Ξ_fd, q_fd}; + tuple u_J_w = {u, ∇Ξ, w}; + tuple ∇u_q = {∇u, q}; + + auto setup = + [] MFEM_HOST_DEVICE(const real_t &u, + const tensor &J, + const real_t &w) + { + return tuple{inv(J) * transpose(inv(J)) * det(J) * w}; + }; + DifferentiableOperator ∂Setup(u_sol, Ξ_q_params, pmesh); + ∂Setup.SetParameters({nodes, &qd}); + ∂Setup.AddDomainIntegrator(setup, u_J_w, tuple{q}, ir, ess_bdr); + ∂Setup.Mult(Vector{fes.GetTrueVSize()}, qd); + + auto apply = + [] MFEM_HOST_DEVICE(const tensor &∇u, + const tensor &q) + { + return tuple{q * ∇u}; + }; + ∂op = std::make_unique(u_sol, q_param, pmesh); + ∂op->SetParameters({ &qd }); + ∂op->AddDomainIntegrator(apply, ∇u_q, tuple{∇u}, ir, ess_bdr); + } + else { MFEM_ABORT("Invalid version"); } + + OperatorHandle A; + Vector B, X; + if (version >= 2) + { + Operator *A_ptr; + ∂op->FormLinearSystem(ess_tdof_list, x, b, A_ptr, X, B); + A.Reset(A_ptr); + } + else + { + a->FormLinearSystem(ess_tdof_list, x, b, A, X, B); + } + + const real_t rtol = 0.0; + const int max_it = 32, print_lvl = -1; + CGSolver cg(MPI_COMM_WORLD); + cg.SetOperator(*A); + cg.iterative_mode = false; + if constexpr (true) // check + { + cg.SetPrintLevel(1); + cg.SetMaxIter(100); + cg.SetRelTol(1e-8); + cg.SetAbsTol(0.0); + cg.Mult(B, X); + MFEM_VERIFY(cg.GetConverged(), "CG solver did not converge."); + MFEM_DEVICE_SYNC; + mfem::out << "✅" << std::endl; + } + cg.SetAbsTol(0.0); + cg.SetRelTol(rtol); + cg.SetMaxIter(max_it); + cg.SetPrintLevel(print_lvl); + + if (visualization) + { + if (version >= 2) + { + ∂op->RecoverFEMSolution(X, b, x); + } + else + { + a->RecoverFEMSolution(X, b, x); + } + int visport = 19916; + char vishost[] = "localhost"; + socketstream sol_sock(vishost, visport); + sol_sock.precision(8); + sol_sock << "parallel " << Mpi::WorldSize() << " " << Mpi::WorldRank() << "\n"; + sol_sock << "solution\n" << pmesh << x << std::flush; + } + + return EXIT_SUCCESS; +} \ No newline at end of file diff --git a/examples/dfem/laghos.cpp b/examples/dfem/laghos.cpp new file mode 100644 index 0000000000..494bfa733b --- /dev/null +++ b/examples/dfem/laghos.cpp @@ -0,0 +1,1945 @@ +// 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 + +// TODO: Do we want this to be included from mfem.hpp automatically now? +#include +#include + +#include +#include + +using namespace mfem; + +using mfem::future::tuple; +using mfem::future::tensor; + +using future::DifferentiableOperator; +using future::DerivativeOperator; +using future::ParametricFunction; +using future::ParametricSpace; +using future::FieldDescriptor; +using future::Gradient; +using future::Value; +using future::Weight; +using future::None; + +constexpr int VELOCITY = 0; +constexpr int DENSITY0 = 1; +constexpr int COORDINATES0 = 2; +constexpr int COORDINATES = 3; +constexpr int MATERIAL = 4; +constexpr int SPECIFIC_INTERNAL_ENERGY = 5; +constexpr int ELEMENT_SIZE0 = 6; +constexpr int ORDER_VEL = 7; +constexpr int DT_EST = 8; +constexpr int STRESS_TENSOR = 9; + +constexpr int DIMENSION = 2; + +int problem = 0; + +void threshold(Vector &v) +{ + for (int i = 0; i < v.Size(); i++) + { + if (abs(v(i)) <= 1e-12) + { + v(i) = 0.0; + } + } +} + +MFEM_HOST_DEVICE inline +real_t taylor_source(const Vector &x) +{ + return 3.0 / 8.0 * M_PI * ( cos(3.0*M_PI*x(0)) * cos(M_PI*x(1)) - + cos(M_PI*x(0)) * cos(3.0*M_PI*x(1)) ); +}; + +// Smooth transition between 0 and 1 for x in [-eps, eps]. +MFEM_HOST_DEVICE inline +real_t smooth_step_01(real_t x, real_t eps) +{ + const real_t y = (x + eps) / (2.0 * eps); + if (y < 0.0) { return 0.0; } + if (y > 1.0) { return 1.0; } + return (3.0 - 2.0 * y) * y * y; +} + +MFEM_HOST_DEVICE inline +void ComputeMaterialProperties(const real_t &gamma, const real_t &rho, + const real_t &E, real_t &p, real_t &cs) +{ + p = (gamma - 1.0) * rho * E; + cs = sqrt(gamma * (gamma - 1.0) * E); +} + +using vecd = tensor; +using matd = tensor; + +template +MFEM_HOST_DEVICE inline +tuple qdata_setup( + const matd &dvdxi, + const real_t &rho0, + const matd &J0, + const matd &J, + const real_t &gamma, + const real_t &E, + const real_t &h0, + const real_t &order_v, + const real_t &w, + const real_t &cfl, + const bool &use_viscosity) +{ + constexpr real_t eps = 1e-12; + constexpr real_t vorticity_coeff = 1.0; + real_t p, cs; + real_t detJ = det(J); + matd invJ = inv(J); + matd stress{{{0.0}}}; + const real_t rho = rho0 * det(J0) / detJ; + const real_t Ez = fmax(0.0, E); + real_t visc_coeff = 0.0; + real_t dt_est = std::numeric_limits::infinity(); + + ComputeMaterialProperties(gamma, rho, Ez, p, cs); + + for (int d = 0; d < DIMENSION; d++) + { + stress(d, d) = -p; + } + + if (use_viscosity) + { + auto symdvdx = sym(dvdxi * invJ); + auto [eigvals, eigvecs] = eig(symdvdx); + vecd compr_dir = get_col(eigvecs, 0); + auto ph_dir = (J * inv(J0)) * compr_dir; + const real_t h = h0 * norm(ph_dir) / norm(compr_dir); + // Measure of maximal compression. + const real_t mu = eigvals(0); + visc_coeff = 2.0 * rho * h * h * fabs(mu); + visc_coeff += 0.5 * rho * h * cs * vorticity_coeff * + (1.0 - smooth_step_01(mu - 2.0 * eps, eps)); + stress += visc_coeff * symdvdx; + } + + if constexpr (compute_dtest) + { + if (detJ < 0.0) + { + // This will force repetition of the step with smaller dt. + dt_est = 0.0; + } + else + { + const real_t h_min = calcsv(J, DIMENSION-1) / static_cast(order_v); + const real_t idt = cs / h_min + 2.5 * visc_coeff / rho / h_min / h_min; + + if (idt > 0.0) + { + dt_est = cfl / idt; + } + else + { + dt_est = std::numeric_limits::infinity(); + } + } + } + + matd stressJiT = stress * transpose(invJ) * detJ * w; + return tuple{stressJiT, dt_est}; +} + +struct TimeStepEstimateQFunction +{ + TimeStepEstimateQFunction(const real_t *external_data) : + external_data(external_data) {} + + MFEM_HOST_DEVICE inline + auto operator()( + const matd &dvdxi, + const real_t &rho0, + const matd &J0, + const matd &J, + const real_t &gamma, + const real_t &E, + const real_t &h0, + const real_t &order_v, + const real_t &w) const + { + real_t dt_est = get<1>( + qdata_setup(dvdxi, rho0, J0, J, gamma, E, h0, order_v, w, + external_data[0], static_cast(external_data[2]))); + return tuple{dt_est}; + } + + const real_t *external_data; +}; + +struct UpdateQuadratureDataQFunction +{ + UpdateQuadratureDataQFunction(const real_t *external_data) : + external_data(external_data) {} + + MFEM_HOST_DEVICE inline + auto operator()( + const matd &dvdxi, + const real_t &rho0, + const matd &J0, + const matd &J, + const real_t &gamma, + const real_t &E, + const real_t &h0, + const real_t &order_v, + const real_t &w) const + { + matd stressJiT = get<0>( + qdata_setup(dvdxi, rho0, J0, J, gamma, E, h0, order_v, w, + external_data[0], static_cast(external_data[2]))); + return tuple{stressJiT}; + } + + const real_t *external_data; +}; + +class MomentumQFunction +{ +public: + MomentumQFunction(const real_t *external_data) : + external_data(external_data) {} + + MFEM_HOST_DEVICE inline + auto operator()( + const matd &dvdxi, + const real_t &rho0, + const matd &J0, + const matd &J, + const real_t &gamma, + const real_t &E, + const real_t &h0, + const real_t &order_v, + const real_t &w) const + { + auto stressJiT = get<0>( + qdata_setup(dvdxi, rho0, J0, J, gamma, E, h0, order_v, w, external_data[0], + static_cast(external_data[2]))); + + // out << gamma << " " << rho << " " << Ez << " " << p << " " << cs << "\n"; + // out << stressJiT << "\n"; + // TODO-bug: investigate transpose of matrices in return types + // return tuple{transpose(stressJiT)}; + return tuple{stressJiT}; + } + + const real_t *external_data; +}; + +class MomentumPAQFunction +{ +public: + MomentumPAQFunction() = default; + + MFEM_HOST_DEVICE inline + auto operator()( + const matd &stressJiT) const + { + return tuple{stressJiT}; + } +}; + +class EnergyConservationQFunction +{ +public: + EnergyConservationQFunction(const real_t *external_data) : + external_data(external_data) {} + + MFEM_HOST_DEVICE inline + auto operator()( + const matd &dvdxi, + const real_t &rho0, + const matd &J0, + const matd &J, + const real_t &gamma, + const real_t &E, + const real_t &h0, + const real_t &order_v, + const real_t &w) const + { + auto stressJiT = get<0>( + qdata_setup(dvdxi, rho0, J0, J, gamma, E, h0, order_v, w, external_data[0], + static_cast(external_data[2]))); + return tuple{ddot(stressJiT, dvdxi)}; + } + + const real_t *external_data; +}; + +class EnergyConservationPAQFunction +{ +public: + EnergyConservationPAQFunction() = default; + + MFEM_HOST_DEVICE inline + auto operator()( + const matd &dvdxi, + const matd &stressJiT) const + { + return tuple{ddot(stressJiT, dvdxi)}; + } +}; + +class TotalInternalEnergyQFunction +{ +public: + TotalInternalEnergyQFunction() = default; + + MFEM_HOST_DEVICE inline + auto operator() ( + const real_t &E, + const real_t &rho0, + const matd &J0, + const real_t &w) const + { + return tuple{rho0 * E * det(J0) * w}; + } +}; + +class TotalKineticEnergyQFunction +{ +public: + TotalKineticEnergyQFunction() = default; + + MFEM_HOST_DEVICE inline + auto operator() ( + const vecd &v, + const real_t &rho0, + const matd &J0, + const real_t &w) const + { + return tuple{rho0 * 0.5 * v * v * det(J0) * w}; + } +}; + +class DensityQFunction +{ +public: + DensityQFunction() = default; + + MFEM_HOST_DEVICE inline + auto operator() ( + const real_t &rho0, + const matd &J0, + const real_t &w) const + { + return tuple{rho0 * det(J0) * w}; + } +}; + +struct QuadratureData +{ + static constexpr int aux_dim = 1; + QuadratureData(const ParMesh &mesh, const IntegrationRule &ir) : + StressSpace(mesh.Dimension(), mesh.Dimension()*mesh.Dimension(), + ir.GetNPoints(), + mesh.Dimension()*mesh.Dimension()*ir.GetNPoints()*mesh.GetNE()), + stressp(StressSpace), + R(mesh.Dimension(), + aux_dim, + ir.GetNPoints(), + aux_dim*ir.GetNPoints()*mesh.GetNE()), + h0(R), + order_v(R), + dt_est(R) + { + h0.UseDevice(true); + order_v.UseDevice(true); + dt_est.UseDevice(true); + stressp.UseDevice(true); + } + + ParametricSpace StressSpace; + ParametricFunction stressp; + + ParametricSpace R; + ParametricFunction h0, order_v, dt_est; +}; + +class MassPAOperator : public Operator +{ +public: + MassPAOperator(ParFiniteElementSpace &pfes, + const IntegrationRule &ir, + Coefficient &Q) : + Operator(pfes.GetTrueVSize()), + comm(pfes.GetParMesh()->GetComm()), + dim(pfes.GetMesh()->Dimension()), + NE(pfes.GetMesh()->GetNE()), + vsize(pfes.GetVSize()), + pabf(&pfes), + ess_tdofs_count(0), + ess_tdofs(0) + { + if (dim > 1) + { + pabf.SetAssemblyLevel(AssemblyLevel::PARTIAL); + } + pabf.AddDomainIntegrator(new mfem::MassIntegrator(Q, &ir)); + pabf.Assemble(); + pabf.FormSystemMatrix(mfem::Array(), mass); + } + + void SetEssentialTrueDofs(Array &dofs) + { + ess_tdofs_count = dofs.Size(); + if (ess_tdofs.Size() == 0) + { + int ess_tdofs_sz; + MPI_Allreduce(&ess_tdofs_count,&ess_tdofs_sz, 1, MPI_INT, MPI_SUM, comm); + MFEM_ASSERT(ess_tdofs_sz > 0, "ess_tdofs_sz should be positive!"); + ess_tdofs.SetSize(ess_tdofs_sz); + } + if (ess_tdofs_count == 0) { return; } + ess_tdofs = dofs; + } + + void EliminateRHS(Vector &b) const + { + if (ess_tdofs_count > 0) { b.SetSubVector(ess_tdofs, 0.0); } + } + + void Mult(const Vector &x, Vector &y) const override + { + mass->Mult(x, y); + if (ess_tdofs_count > 0) { y.SetSubVector(ess_tdofs, 0.0); } + } + + void FullAddMult(const Vector &x, Vector &y) const + { + mass->AddMult(x, y); + } + + const ParBilinearForm &GetBF() const { return pabf; } + + const MPI_Comm comm; + const int dim, NE, vsize; + ParBilinearForm pabf; + int ess_tdofs_count; + Array ess_tdofs; + OperatorPtr mass; +}; + +class LagrangianHydroJacobianOperator : public Operator +{ +public: + LagrangianHydroJacobianOperator(real_t h, int H1tsize, int L2tsize) : + Operator(2*H1tsize + L2tsize), h(h), H1tsize(H1tsize), L2tsize(L2tsize) {} + + void Mult(const Vector &k, Vector &y) const override + { + jvp(k, y); + } + + template + void Setup(hydro_t &hydro, + std::shared_ptr dRvdx, + std::shared_ptr dRvdv, + std::shared_ptr dRvde, + std::shared_ptr dRedx, + std::shared_ptr dRedv, + std::shared_ptr dRede) + { + w.SetSize(this->height); + z.SetSize(this->height); + + jvp = [dRvdx, dRvdv, dRvde, dRedx, dRedv, dRede, this, &hydro] + (const Vector &u, Vector &y) + { + w = u; + Vector wx, wv, we; + wx.MakeRef(w, 0, H1tsize); + wv.MakeRef(w, H1tsize, H1tsize); + we.MakeRef(w, 2*H1tsize, L2tsize); + + Vector zx, zv, ze; + zx.MakeRef(z, 0, H1tsize); + zv.MakeRef(z, H1tsize, H1tsize); + ze.MakeRef(z, 2*H1tsize, L2tsize); + + Vector yx, yv, ye; + yx.MakeRef(y, 0, H1tsize); + yv.MakeRef(y, H1tsize, H1tsize); + ye.MakeRef(y, 2*H1tsize, L2tsize); + + // position + yx = wv; + yx *= -h; + yx += wx; + + // velocity + // wv.SetSubVector(hydro.ess_tdof, 0.0); + dRvdx->Mult(wx, zv); + zv *= h; + yv = zv; + dRvdv->Mult(wv, zv); + zv *= h; + yv += zv; + // hydro.Mv.TrueAddMult(wv, yv); + Vector wvc, yvc; + for (int c = 0; c < hydro.H1.GetMesh()->Dimension(); c++) + { + wvc.MakeRef(wv, c*hydro.H1c.GetVSize(), hydro.H1c.GetVSize()); + yvc.MakeRef(yv, c*hydro.H1c.GetVSize(), hydro.H1c.GetVSize()); + hydro.Mv->FullAddMult(wvc, yvc); + yvc.SyncAliasMemory(yv); + } + yv.SyncAliasMemory(y); + + dRvde->Mult(we, zv); + zv *= h; + yv += zv; + yv.SetSubVector(hydro.ess_tdof, 0.0); + // for (int i = 0; i < hydro.ess_tdof.Size(); i++) + // { + // // yv(hydro.ess_tdof[i]) = uv(hydro.ess_tdof[i]); + // yv(hydro.ess_tdof[i]) = 0.0; + // } + // yv = 0.0; + + // energy + // [ wx ] + // [ dRe/dx dRe/dv dRe/de ] [ wv ] + // [ we ] + // + + dRedx->Mult(wx, ze); + ze *= -h; + ye = ze; + + dRedv->Mult(wv, ze); + ze *= -h; + ye += ze; + + dRede->Mult(we, ze); + ze *= -h; + // hydro.Me.TrueAddMult(we, ze); + hydro.Me->FullAddMult(we, ze); + + ye += ze; + + yx.SyncAliasMemory(y); + yv.SyncAliasMemory(y); + ye.SyncAliasMemory(y); + }; + } + + virtual MemoryClass GetMemoryClass() const override + { + return Device::GetDeviceMemoryClass(); + } + + real_t h; + std::function jvp; + const int H1tsize; + const int L2tsize; + Vector w, z; +}; + +template +class LagrangianHydroResidualOperator : public Operator +{ +public: + LagrangianHydroResidualOperator(hydro_t &hydro, const real_t dt, + const Vector &x, bool fd_gradient) : + Operator(2*hydro.H1.GetTrueVSize()+hydro.L2.GetTrueVSize()), + hydro(hydro), + dt(dt), + x(x), + u(x.Size()), + H1tsize(hydro.H1.GetTrueVSize()), + L2tsize(hydro.L2.GetTrueVSize()), + fd_gradient(fd_gradient) {} + + void Mult(const Vector &k, Vector &R) const override + { + hydro.UpdateMesh(u); + + u = k; + u *= dt; + u += x; + + hydro.mesh_nodes.SyncMemory(u); + + auto kptr = const_cast(&k); + Vector kx, kv, ke; + kx.MakeRef(*kptr, 0, H1tsize); + kv.MakeRef(*kptr, H1tsize, H1tsize); + ke.MakeRef(*kptr, 2*H1tsize, L2tsize); + + Vector ux, uv, ue; + ux.MakeRef(u, 0, H1tsize); + uv.MakeRef(u, H1tsize, H1tsize); + ue.MakeRef(u, 2*H1tsize, L2tsize); + + Vector Rx, Rv, Re; + Rx.MakeRef(R, 0, H1tsize); + Rv.MakeRef(R, H1tsize, H1tsize); + Re.MakeRef(R, 2*H1tsize, L2tsize); + + Rx = kx; + Rx -= uv; + + hydro.momentum_mf->SetParameters({&hydro.rho0, &hydro.x0, &ux, &hydro.material, &ue, &hydro.qdata->h0, &hydro.qdata->order_v}); + hydro.momentum_mf->Mult(uv, Rv); + + // hydro.Mv.TrueAddMult(kv, Rv); + Vector kvc, Rvc; + for (int c = 0; c < hydro.H1.GetMesh()->Dimension(); c++) + { + kvc.MakeRef(kv, c*hydro.H1c.GetVSize(), hydro.H1c.GetVSize()); + Rvc.MakeRef(Rv, c*hydro.H1c.GetVSize(), hydro.H1c.GetVSize()); + hydro.Mv->FullAddMult(kvc, Rvc); + Rvc.SyncAliasMemory(Rv); + } + Rv.SyncAliasMemory(R); + + Rv.SetSubVector(hydro.ess_tdof, 0.0); + // Rv = 0.0; + + hydro.energy_conservation_mf->SetParameters({&uv, &hydro.rho0, &hydro.x0, &ux, &hydro.material, &hydro.qdata->h0, &hydro.qdata->order_v}); + hydro.energy_conservation_mf->Mult(ue, Re); + + Re.Neg(); + + if (problem == 0) + { + LinearForm e_source(&hydro.L2); + hydro.L2.GetMesh()->DeleteGeometricFactors(); + FunctionCoefficient coeff(taylor_source); + DomainLFIntegrator *d = new DomainLFIntegrator(coeff, &hydro.ir); + e_source.AddDomainIntegrator(d); + e_source.UseFastAssembly(true); + e_source.Assemble(); + + Re -= e_source; + } + + // hydro.Me.TrueAddMult(ke, Re); + hydro.Me->FullAddMult(ke, Re); + + Rx.SyncAliasMemory(R); + Rv.SyncAliasMemory(R); + Re.SyncAliasMemory(R); + } + + Operator& GetGradient(const Vector &k) const override + { + jacobian.reset(new LagrangianHydroJacobianOperator(dt, H1tsize, L2tsize)); + + u = k; + u *= dt; + u += x; + + auto kptr = const_cast(&k); + Vector kx, kv, ke; + kx.MakeRef(*kptr, 0, H1tsize); + kv.MakeRef(*kptr, H1tsize, H1tsize); + ke.MakeRef(*kptr, 2*H1tsize, L2tsize); + + Vector ux, uv, ue; + ux.MakeRef(u, 0, H1tsize); + uv.MakeRef(u, H1tsize, H1tsize); + ue.MakeRef(u, 2*H1tsize, L2tsize); + + if (fd_gradient) + { + fd_jacobian.reset(new future::FDJacobian(*this, k)); + return *fd_jacobian; + } + else + { + auto dRvdx = hydro.momentum_mf->GetDerivative(COORDINATES, {&uv}, + {&hydro.rho0, &hydro.x0, &ux, &hydro.material, &ue, &hydro.qdata->h0, &hydro.qdata->order_v}); + + auto dRvdv = hydro.momentum_mf->GetDerivative(VELOCITY, {&uv}, + {&hydro.rho0, &hydro.x0, &ux, &hydro.material, &ue, &hydro.qdata->h0, &hydro.qdata->order_v}); + + auto dRvde = hydro.momentum_mf->GetDerivative(SPECIFIC_INTERNAL_ENERGY, {&uv}, + {&hydro.rho0, &hydro.x0, &ux, &hydro.material, &ue, &hydro.qdata->h0, &hydro.qdata->order_v}); + + auto dRedx = hydro.energy_conservation_mf->GetDerivative(COORDINATES, {&ue}, + {&uv, &hydro.rho0, &hydro.x0, &ux, &hydro.material, &hydro.qdata->h0, &hydro.qdata->order_v}); + + auto dRedv = hydro.energy_conservation_mf->GetDerivative(VELOCITY, {&ue}, + {&uv, &hydro.rho0, &hydro.x0, &ux, &hydro.material, &hydro.qdata->h0, &hydro.qdata->order_v}); + + auto dRede = hydro.energy_conservation_mf->GetDerivative( + SPECIFIC_INTERNAL_ENERGY, {&ue}, + {&uv, &hydro.rho0, &hydro.x0, &ux, &hydro.material, &hydro.qdata->h0, &hydro.qdata->order_v}); + + jacobian->Setup(hydro, dRvdx, dRvdv, dRvde, dRedx, dRedv, dRede); + return *jacobian; + } + } + + hydro_t &hydro; + const real_t dt; + const Vector &x; + mutable Vector u; + const int H1tsize; + const int L2tsize; + mutable std::shared_ptr fd_jacobian; + mutable std::shared_ptr jacobian; + bool fd_gradient; +}; + +class LagrangianHydroOperator : public TimeDependentOperator +{ +public: + LagrangianHydroOperator( + ParFiniteElementSpace &H1, + ParFiniteElementSpace &L2, + const Array &ess_tdof, + const IntegrationRule &ir, + FunctionCoefficient &rho0_coeff, + ParGridFunction &x0_gf, + ParGridFunction &rho0_gf, + ParGridFunction &material_gf, + std::shared_ptr update_qdata, + std::shared_ptr dtest_mf, + std::shared_ptr momentum_mf, + std::shared_ptr momentum_pa, + std::shared_ptr energy_conservation_mf, + std::shared_ptr energy_conservation_pa, + std::shared_ptr total_internal_energy_mf, + std::shared_ptr total_kinetic_energy_mf, + std::shared_ptr density_mf, + std::shared_ptr qdata, + bool fd_gradient, + const int nonlinear_maximum_iterations, + const real_t nonlinear_relative_tolerance) : + TimeDependentOperator(2*H1.GetVSize()+L2.GetVSize()), + H1(H1), + L2(L2), + H1c(H1.GetParMesh(), H1.FEColl(), 1), + ess_tdof(ess_tdof), + ir(ir), + x0(x0_gf), + rho0(rho0_gf), + material(material_gf), + update_qdata(update_qdata), + dtest_mf(dtest_mf), + momentum_mf(momentum_mf), + momentum_pa(momentum_pa), + energy_conservation_mf(energy_conservation_mf), + energy_conservation_pa(energy_conservation_pa), + total_internal_energy_mf(total_internal_energy_mf), + total_kinetic_energy_mf(total_kinetic_energy_mf), + density_mf(density_mf), + qdata(qdata), + mesh_nodes(&H1), + rhsvc(&H1c), + dvc(&H1c), + rho0_coeff(rho0_coeff), + RHSv(H1.GetTrueVSize()), + rhsv(H1.GetVSize()), + X(2*H1.GetTrueVSize()+L2.GetTrueVSize()), + Xv(H1.GetTrueVSize()), + Xvc(H1c.GetTrueVSize()), + Xe(L2.GetTrueVSize()), + K(2*H1.GetTrueVSize()+L2.GetTrueVSize()), + B(H1c.GetTrueVSize()), + RHSe(L2.GetTrueVSize()), + rhse(L2.GetVSize()), + nl2dofs(L2.GetFE(0)->GetDof()), + fd_gradient(fd_gradient), + nonlinear_maximum_iterations(nonlinear_maximum_iterations), + nonlinear_relative_tolerance(nonlinear_relative_tolerance) + { + Mv = new MassPAOperator(H1c, ir, rho0_coeff); + Array empty_tdofs; + Mv_Jprec = new OperatorJacobiSmoother(Mv->GetBF(), empty_tdofs); + + Me = new MassPAOperator(L2, ir, rho0_coeff); + + // Inside the above constructors for mass, there is reordering of the mesh + // nodes which is performed on the host. Since the mesh nodes are a + // subvector, so we need to sync with the rest of the base vector (which + // is assumed to be in the memory space used by the mfem::Device). + H1.GetParMesh()->GetNodes()->ReadWrite(); + // Attributes 1/2/3 correspond to fixed-x/y/z boundaries, i.e., + // we must enforce v_x/y/z = 0 for the velocity components. + const int bdr_attr_max = H1.GetMesh()->bdr_attributes.Max(); + Array ess_bdr(bdr_attr_max); + for (int c = 0; c < H1.GetMesh()->Dimension(); c++) + { + ess_bdr = 0; + ess_bdr[c] = 1; + H1c.GetEssentialTrueDofs(ess_bdr, c_tdofs[c]); + c_tdofs[c].Read(); + } + } + + void Mult(const Vector &S, Vector &dSdt) const override + { + UpdateMesh(S); + UpdateQuadratureData(S); + + auto sptr = const_cast(&S); + const int H1vsize = H1.GetVSize(); + + ParGridFunction x, v, e; + x.MakeRef(&H1, *sptr, 0); + v.MakeRef(&H1, *sptr, H1vsize); + e.MakeRef(&L2, *sptr, 2*H1vsize); + + ParGridFunction dx, dv, de; + dx.MakeRef(&H1, dSdt, 0); + dv.MakeRef(&H1, dSdt, H1vsize); + de.MakeRef(&L2, dSdt, 2*H1vsize); + + // solve position + dx = v; + + // solve velocity + { + dv = 0.0; + + // momentum_mf->SetParameters({&rho0, &x0, &x, &material, &e, &qdata->h0, &qdata->order_v}); + momentum_pa->SetParameters({&qdata->stressp}); + H1.GetRestrictionMatrix()->Mult(v, Xv); + // momentum_mf->Mult(Xv, RHSv); + momentum_pa->Mult(Xv, RHSv); + RHSv.Neg(); + H1.GetRestrictionMatrix()->MultTranspose(RHSv, rhsv); + + // solve for each velocity component + const int size = H1c.GetVSize(); + const Operator *Pconf = H1c.GetProlongationMatrix(); + for (int c = 0; c < H1.GetMesh()->Dimension(); c++) + { + dvc.MakeRef(&H1c, dSdt, H1vsize + c*size); + rhsvc.MakeRef(&H1c, rhsv, c*size); + if (Pconf) + { + Pconf->MultTranspose(rhsvc, B); + } + else + { + B = rhsvc; + } + + CGSolver cg(H1c.GetParMesh()->GetComm()); + cg.SetPreconditioner(*Mv_Jprec); + cg.SetOperator(*Mv); + cg.SetRelTol(1e-8); + cg.SetAbsTol(0.0); + cg.SetMaxIter(300); + cg.SetPrintLevel(-1); + + H1c.GetRestrictionMatrix()->Mult(dvc, Xvc); + Mv->SetEssentialTrueDofs(c_tdofs[c]); + Mv->EliminateRHS(B); + cg.Mult(B, Xvc); + if (Pconf) + { + Pconf->Mult(Xvc, dvc); + } + else + { + dvc = Xvc; + } + dvc.GetMemory().SyncAlias(dSdt.GetMemory(), dvc.Size()); + } + } + + // solve energy + { + de = 0.0; + + // energy_conservation_mf->SetParameters({&v, &rho0, &x0, &x, &material, &qdata->h0, &qdata->order_v}); + energy_conservation_pa->SetParameters({&v, &qdata->stressp}); + L2.GetRestrictionMatrix()->Mult(e, Xe); + // energy_conservation_mf->Mult(Xe, RHSe); + energy_conservation_pa->Mult(Xe, RHSe); + L2.GetRestrictionMatrix()->MultTranspose(RHSe, rhse); + + if (problem == 0) + { + LinearForm e_source(&L2); + L2.GetMesh()->DeleteGeometricFactors(); + FunctionCoefficient coeff(taylor_source); + DomainLFIntegrator *d = new DomainLFIntegrator(coeff, &ir); + e_source.AddDomainIntegrator(d); + e_source.UseFastAssembly(true); + e_source.Assemble(); + rhse += e_source; + } + + CGSolver cg(L2.GetParMesh()->GetComm()); + cg.SetOperator(*Me); + cg.iterative_mode = false; + cg.SetRelTol(1e-8); + cg.SetAbsTol(0.0); + cg.SetMaxIter(300); + cg.SetPrintLevel(-1); + cg.Mult(rhse, de); + de.GetMemory().SyncAlias(dSdt.GetMemory(), de.Size()); + } + } + + void ImplicitSolve(const real_t dt, const Vector &x, Vector &k) override + { + auto xptr = const_cast(&x); + + Vector xx, xv, xe; + xx.MakeRef(*xptr, 0, H1.GetVSize()); + xv.MakeRef(*xptr, H1.GetVSize(), H1.GetVSize()); + xe.MakeRef(*xptr, 2*H1.GetVSize(), L2.GetVSize()); + + Vector Xx, Xv, Xe; + Xx.MakeRef(X, 0, H1.GetTrueVSize()); + Xv.MakeRef(X, H1.GetTrueVSize(), H1.GetTrueVSize()); + Xe.MakeRef(X, 2*H1.GetTrueVSize(), L2.GetTrueVSize()); + + H1.GetRestrictionMatrix()->Mult(xx, Xx); + H1.GetRestrictionMatrix()->Mult(xv, Xv); + L2.GetRestrictionMatrix()->Mult(xe, Xe); + + Xx.SyncAliasMemory(X); + Xv.SyncAliasMemory(X); + Xe.SyncAliasMemory(X); + + auto residual = LagrangianHydroResidualOperator(*this, dt, X, fd_gradient); + + GMRESSolver gmres(MPI_COMM_WORLD); + gmres.SetMaxIter(500); + gmres.SetKDim(500); + gmres.SetRelTol(1e-8); + gmres.SetAbsTol(1e-12); + gmres.SetPrintLevel(IterativeSolver::PrintLevel().None()); + + NewtonSolver newton(MPI_COMM_WORLD); + newton.SetPrintLevel(IterativeSolver::PrintLevel().None()); + newton.SetOperator(residual); + newton.SetSolver(gmres); + newton.SetAdaptiveLinRtol(); + newton.SetMaxIter(nonlinear_maximum_iterations); + newton.SetRelTol(nonlinear_relative_tolerance); + newton.SetAbsTol(1e-12); + + Vector zero; + K = X; + newton.Mult(zero, K); + + Vector Kx, Kv, Ke; + Kx.MakeRef(K, 0, H1.GetTrueVSize()); + Kv.MakeRef(K, H1.GetTrueVSize(), H1.GetTrueVSize()); + Ke.MakeRef(K, 2*H1.GetTrueVSize(), L2.GetTrueVSize()); + + Vector kx, kv, ke; + kx.MakeRef(k, 0, H1.GetVSize()); + kv.MakeRef(k, H1.GetVSize(), H1.GetVSize()); + ke.MakeRef(k, 2*H1.GetVSize(), L2.GetVSize()); + + H1.GetProlongationMatrix()->Mult(Kx, kx); + H1.GetProlongationMatrix()->Mult(Kv, kv); + L2.GetProlongationMatrix()->Mult(Ke, ke); + // kx.SyncAliasMemory(k); + // kv.SyncAliasMemory(k); + // ke.SyncAliasMemory(k); + } + + void UpdateMesh(const Vector &S) const + { + Vector* sptr = const_cast(&S); + mesh_nodes.MakeRef(&H1, *sptr, 0); + H1.GetParMesh()->NewNodes(mesh_nodes, false); + } + + real_t GetTimeStepEstimate(const Vector &S) + { + UpdateMesh(S); + + auto sptr = const_cast(&S); + const int H1vsize = H1.GetVSize(); + ParGridFunction x, v, e; + x.MakeRef(&H1, *sptr, 0); + v.MakeRef(&H1, *sptr, H1vsize); + e.MakeRef(&L2, *sptr, 2*H1vsize); + dtest_mf->SetParameters({&v, &rho0, &x0, &x, &material, &e, &qdata->h0, &qdata->order_v}); + auto &dt_est = qdata->dt_est; + dtest_mf->Mult(dt_est, dt_est); + + real_t dt_est_local = std::numeric_limits::infinity(); + for (int i = 0; i < dt_est.Size(); i++) + { + if (dt_est(i) == 0.0) + { + return 0.0; + } + dt_est_local = fmin(dt_est_local, dt_est(i)); + } + + real_t dt_est_global; + MPI_Allreduce(&dt_est_local, &dt_est_global, 1, MPI_DOUBLE, MPI_MIN, + L2.GetComm()); + + return dt_est_global; + } + + real_t InternalEnergy(ParGridFunction &e) + { + const auto mt = Device::GetDeviceMemoryType(); + Vector E(L2.GetTrueVSize(), mt), Y(L2.GetTrueVSize(), mt); + total_internal_energy_mf->SetParameters({&rho0, &x0}); + L2.GetRestrictionMatrix()->Mult(e, E); + total_internal_energy_mf->Mult(E, Y); + const real_t ie_local = Y.Sum(); + real_t ie_global = 0.0; + MPI_Allreduce(&ie_local, &ie_global, 1, MPI_DOUBLE, MPI_SUM, + L2.GetParMesh()->GetComm()); + return ie_global; + } + + real_t KineticEnergy(ParGridFunction &v) + { + const auto mt = Device::GetDeviceMemoryType(); + Vector V(H1.GetTrueVSize(), mt), Y(L2.GetTrueVSize(), mt); + total_kinetic_energy_mf->SetParameters({&rho0, &x0}); + H1.GetRestrictionMatrix()->Mult(v, V); + total_kinetic_energy_mf->Mult(V, Y); + const real_t ke_local = Y.Sum(); + real_t ke_global = 0.0; + MPI_Allreduce(&ke_local, &ke_global, 1, MPI_DOUBLE, MPI_SUM, + H1.GetParMesh()->GetComm()); + return ke_global; + } + + void ComputeDensity(ParGridFunction &rho) + { + rho.SetSpace(&L2); + + ParGridFunction rhs_l(&L2); + + Vector rho0_t(L2.GetTrueVSize()), + rho_t(L2.GetTrueVSize()), + rhs(L2.GetTrueVSize()); + + const int l2dofs_cnt = L2.GetFE(0)->GetDof(); + DenseMatrix Mrho(l2dofs_cnt); + DenseMatrixInverse inv(&Mrho); + Vector rhs_e(l2dofs_cnt), rho_z(l2dofs_cnt); + Array dofs(l2dofs_cnt); + MassIntegrator mi(&ir); + + density_mf->SetParameters({&x0}); + L2.GetProlongationMatrix()->MultTranspose(rho0, rho0_t); + density_mf->Mult(rho0_t, rhs); + L2.GetProlongationMatrix()->Mult(rhs, rhs_l); + + for (int e = 0; e < L2.GetParMesh()->GetNE(); e++) + { + const FiniteElement &fe = *L2.GetFE(e); + ElementTransformation &eltr = *L2.GetElementTransformation(e); + L2.GetElementDofs(e, dofs); + mi.AssembleElementMatrix(fe, eltr, Mrho); + inv.Factor(); + rhs_l.GetElementDofValues(e, rhs_e); + inv.Mult(rhs_e, rho_z); + rho.SetSubVector(dofs, rho_z); + } + } + + void UpdateQuadratureData(const Vector &S) const + { + auto sptr = const_cast(&S); + const int H1vsize = H1.GetVSize(); + ParGridFunction x, v, e; + x.MakeRef(&H1, *sptr, 0); + v.MakeRef(&H1, *sptr, H1vsize); + e.MakeRef(&L2, *sptr, 2*H1vsize); + update_qdata->SetParameters({&v, &rho0, &x0, &x, &material, &e, &qdata->h0, &qdata->order_v}); + update_qdata->Mult(qdata->stressp, qdata->stressp); + } + + virtual MemoryClass GetMemoryClass() const override + { + return Device::GetDeviceMemoryClass(); + } + + ParFiniteElementSpace &H1; + ParFiniteElementSpace &L2; + mutable ParFiniteElementSpace H1c; + const Array &ess_tdof; + mutable Array c_tdofs[3]; + const IntegrationRule &ir; + ParGridFunction &x0; + ParGridFunction &rho0; + ParGridFunction &material; + std::shared_ptr update_qdata; + std::shared_ptr dtest_mf; + std::shared_ptr momentum_mf; + std::shared_ptr momentum_pa; + std::shared_ptr energy_conservation_mf; + std::shared_ptr energy_conservation_pa; + std::shared_ptr total_internal_energy_mf; + std::shared_ptr total_kinetic_energy_mf; + std::shared_ptr density_mf; + std::shared_ptr qdata; + mutable ParGridFunction mesh_nodes, rhsvc, dvc; + mutable MassPAOperator *Mv = nullptr, *Me = nullptr; + mutable FunctionCoefficient rho0_coeff; + OperatorJacobiSmoother *Mv_Jprec = nullptr; + mutable Vector RHSv, rhsv, X, Xx, Xv, Xvc, Xe, K, Kx, Kv, Ke, B, RHSe, rhse; + const int nl2dofs; + bool fd_gradient; + const int nonlinear_maximum_iterations; + const real_t nonlinear_relative_tolerance; +}; + +static auto CreateLagrangianHydroOperator( + ParFiniteElementSpace &H1, + ParFiniteElementSpace &L2, + const Array &ess_tdof, + FunctionCoefficient &rho0_coeff, + ParGridFunction &x0_gf, + ParGridFunction &rho0_gf, + ParGridFunction &material_gf, + Vector &external_data, + const IntegrationRule &ir, + bool fd_gradient, + const int nonlinear_maximum_iterations, + const real_t nonlinear_relative_tolerance) +{ + const int order_v = H1.GetOrder(0); + ParMesh &mesh = *H1.GetParMesh(); + + auto qdata = std::make_shared(mesh, ir); + + int ne_loc = mesh.GetNE(), ne_global = 0; + real_t vol_loc = 0.0, vol_global = 0.0; + for (int e = 0; e < mesh.GetNE(); e++) + { + vol_loc += mesh.GetElementVolume(e); + } + MPI_Allreduce(&vol_loc, &vol_global, 1, MPI_DOUBLE, MPI_SUM, mesh.GetComm()); + MPI_Allreduce(&ne_loc, &ne_global, 1, MPI_INT, MPI_SUM, mesh.GetComm()); + + switch (mesh.GetElementBaseGeometry(0)) + { + case Geometry::SEGMENT: qdata->h0 = vol_global / ne_global; break; + case Geometry::SQUARE: qdata->h0 = sqrt(vol_global / ne_global); break; + case Geometry::TRIANGLE: qdata->h0 = sqrt(2.0 * vol_global / ne_global); break; + case Geometry::CUBE: qdata->h0 = pow(vol_global / ne_global, 1./3.); break; + case Geometry::TETRAHEDRON: qdata->h0 = pow(6.0 * vol_global / ne_global, + 1./3.); break; + default: MFEM_ABORT("Unknown zone type!"); + } + qdata->h0 /= (double) H1.GetOrder(0); + + // const real_t h0 = sqrt(vol_global / ne_global) / + // static_cast(H1.GetOrder(0)); + + qdata->order_v = order_v; + qdata->dt_est = std::numeric_limits::infinity(); + + // external_data(2) = qdata->h0; + const auto d_external_data = external_data.Read(); + + Array all_domain_attr(mesh.attributes.Max()); + all_domain_attr = 1; + + std::shared_ptr dt_est; + { + tuple dt_est_kernel_ao = + { + Gradient{}, + Value{}, + Gradient{}, + Gradient{}, + Value{}, + Value{}, + None{}, + None{}, + Weight{} + }; + + tuple dt_est_kernel_oo = {None{}}; + + std::vector dt_est_solutions = + { + FieldDescriptor{DT_EST, &qdata->R} + }; + + std::vector dt_est_parameters = + { + FieldDescriptor{VELOCITY, &H1}, + FieldDescriptor{DENSITY0, &L2}, + FieldDescriptor{COORDINATES0, &H1}, + FieldDescriptor{COORDINATES, &H1}, + FieldDescriptor{MATERIAL, material_gf.ParFESpace()}, + FieldDescriptor{SPECIFIC_INTERNAL_ENERGY, &L2}, + FieldDescriptor{ELEMENT_SIZE0, &qdata->R}, + FieldDescriptor{ORDER_VEL, &qdata->R} + }; + + dt_est = std::make_shared( + dt_est_solutions, dt_est_parameters, mesh); + TimeStepEstimateQFunction dt_est_qf(d_external_data); + dt_est->AddDomainIntegrator(dt_est_qf, dt_est_kernel_ao, + dt_est_kernel_oo, + ir, + all_domain_attr); + } + + std::shared_ptr update_qdata; + { + tuple update_qdata_kernel_ao = + { + Gradient{}, + Value{}, + Gradient{}, + Gradient{}, + Value{}, + Value{}, + None{}, + None{}, + Weight{} + }; + + tuple update_qdata_kernel_oo = {None{}}; + + std::vector update_qdata_solutions = + { + {STRESS_TENSOR, &qdata->StressSpace} + }; + + std::vector update_qdata_parameters = + { + {VELOCITY, &H1}, + {DENSITY0, &L2}, + {COORDINATES0, &H1}, + {COORDINATES, &H1}, + {MATERIAL, material_gf.ParFESpace()}, + {SPECIFIC_INTERNAL_ENERGY, &L2}, + {ELEMENT_SIZE0, &qdata->R}, + {ORDER_VEL, &qdata->R} + }; + + update_qdata = std::make_shared( + update_qdata_solutions, update_qdata_parameters, mesh); + UpdateQuadratureDataQFunction update_qdata_qf(d_external_data); + update_qdata->AddDomainIntegrator(update_qdata_qf, update_qdata_kernel_ao, + update_qdata_kernel_oo, + ir, + all_domain_attr); + } + + // Create momentum operator + std::shared_ptr momentum_mf; + { + tuple momentum_mf_kernel_ao = + { + Gradient{}, + Value{}, + Gradient{}, + Gradient{}, + Value{}, + Value{}, + None{}, + None{}, + Weight{} + }; + + tuple momentum_mf_kernel_oo = {Gradient{}}; + + // * det(J) * weights + // + + std::vector momentum_mf_solutions = + { + FieldDescriptor{VELOCITY, &H1} + }; + + std::vector momentum_mf_parameters = + { + FieldDescriptor{DENSITY0, &L2}, + FieldDescriptor{COORDINATES0, &H1}, + FieldDescriptor{COORDINATES, &H1}, + FieldDescriptor{MATERIAL, material_gf.ParFESpace()}, + FieldDescriptor{SPECIFIC_INTERNAL_ENERGY, &L2}, + FieldDescriptor{ELEMENT_SIZE0, &qdata->R}, + FieldDescriptor{ORDER_VEL, &qdata->R} + }; + + momentum_mf = std::make_shared( + momentum_mf_solutions, momentum_mf_parameters, mesh); + + MomentumQFunction momentum_qf(d_external_data); + auto derivatives = + std::integer_sequence {}; + momentum_mf->AddDomainIntegrator(momentum_qf, momentum_mf_kernel_ao, + momentum_mf_kernel_oo, ir, all_domain_attr, derivatives); + } + + std::shared_ptr momentum_pa; + { + tuple momentum_pa_kernel_ao = {None{}}; + tuple momentum_pa_kernel_oo = {Gradient{}}; + + std::vector momentum_pa_solutions = {{VELOCITY, &H1}}; + std::vector momentum_pa_parameters = {{STRESS_TENSOR, &qdata->StressSpace}}; + + momentum_pa = std::make_shared( + momentum_pa_solutions, momentum_pa_parameters, mesh); + + MomentumPAQFunction momentum_pa_qf; + momentum_pa->AddDomainIntegrator(momentum_pa_qf, momentum_pa_kernel_ao, + momentum_pa_kernel_oo, ir, all_domain_attr); + } + + // Create energy conservation operator + std::shared_ptr energy_conservation_mf; + { + tuple energy_conservation_mf_kernel_ao = + { + Gradient{}, + Value{}, + Gradient{}, + Gradient{}, + Value{}, + Value{}, + None{}, + None{}, + Weight{} + }; + + tuple energy_conservation_mf_kernel_oo = {Value{}}; + + // * det(J) * w + // + + std::vector energy_conservation_mf_solutions = + { + FieldDescriptor{SPECIFIC_INTERNAL_ENERGY, &L2} + }; + + std::vector energy_conservation_mf_parameters = + { + FieldDescriptor{VELOCITY, &H1}, + FieldDescriptor{DENSITY0, &L2}, + FieldDescriptor{COORDINATES0, &H1}, + FieldDescriptor{COORDINATES, &H1}, + FieldDescriptor{MATERIAL, material_gf.ParFESpace()}, + FieldDescriptor{ELEMENT_SIZE0, &qdata->R}, + FieldDescriptor{ORDER_VEL, &qdata->R} + }; + + energy_conservation_mf = + std::make_shared( + energy_conservation_mf_solutions, energy_conservation_mf_parameters, mesh); + + EnergyConservationQFunction energy_conservation_qf(d_external_data); + auto derivatives = + std::integer_sequence {}; + energy_conservation_mf->AddDomainIntegrator( + energy_conservation_qf, energy_conservation_mf_kernel_ao, + energy_conservation_mf_kernel_oo, ir, all_domain_attr, derivatives); + } + + std::shared_ptr energy_conservation_pa; + { + tuple energy_conservation_pa_kernel_ao = {Gradient{}, None{}}; + tuple energy_conservation_pa_kernel_oo = {Value{}}; + + std::vector energy_conservation_pa_solutions = + { + {SPECIFIC_INTERNAL_ENERGY, &L2} + }; + + std::vector energy_conservation_pa_parameters = + { + {VELOCITY, &H1}, + {STRESS_TENSOR, &qdata->StressSpace} + }; + + energy_conservation_pa = + std::make_shared( + energy_conservation_pa_solutions, energy_conservation_pa_parameters, mesh); + + EnergyConservationPAQFunction energy_conservation_qf; + energy_conservation_pa->AddDomainIntegrator( + energy_conservation_qf, energy_conservation_pa_kernel_ao, + energy_conservation_pa_kernel_oo, ir, all_domain_attr); + } + + // Create total internal energy operator + std::shared_ptr total_internal_energy_mf; + { + tuple total_internal_energy_kernel_ao = + { + Value{}, + Value{}, + Gradient{}, + Weight{} + }; + + tuple total_internal_energy_kernel_oo = {Value{}}; + + std::vector total_internal_energy_solutions = + { + FieldDescriptor{SPECIFIC_INTERNAL_ENERGY, &L2} + }; + + std::vector total_internal_energy_parameters = + { + FieldDescriptor{DENSITY0, &L2}, + FieldDescriptor{COORDINATES0, &H1} + }; + + total_internal_energy_mf = + std::make_shared( + total_internal_energy_solutions, + total_internal_energy_parameters, + mesh); + + TotalInternalEnergyQFunction total_internal_energy_qf; + total_internal_energy_mf->AddDomainIntegrator( + total_internal_energy_qf, total_internal_energy_kernel_ao, + total_internal_energy_kernel_oo, ir, all_domain_attr); + } + + // Create total kinetic energy operator + std::shared_ptr total_kinetic_energy_mf; + { + tuple total_kinetic_energy_kernel_ao = + { + Value{}, + Value{}, + Gradient{}, + Weight{} + }; + + tuple total_kinetic_energy_kernel_oo = {Value{}}; + + std::vector total_kinetic_energy_solutions = + { + FieldDescriptor{VELOCITY, &H1} + }; + + std::vector total_kinetic_energy_parameters = + { + FieldDescriptor{DENSITY0, &L2}, + FieldDescriptor{COORDINATES0, &H1} + }; + + total_kinetic_energy_mf = + std::make_shared( + total_kinetic_energy_solutions, + total_kinetic_energy_parameters, mesh); + TotalKineticEnergyQFunction total_kinetic_energy_qf; + total_kinetic_energy_mf->AddDomainIntegrator( + total_kinetic_energy_qf, total_kinetic_energy_kernel_ao, + total_kinetic_energy_kernel_oo, ir, all_domain_attr); + } + + // Create density operator + std::shared_ptr density_mf; + { + tuple density_kernel_ao = + { + Value{}, + Gradient{}, + Weight{} + }; + + tuple density_kernel_oo = {Value{}}; + + std::vector density_solutions = + { + FieldDescriptor{DENSITY0, &L2} + }; + + std::vector density_parameters = + { + FieldDescriptor{COORDINATES0, &H1} + }; + + density_mf = std::make_shared( + density_solutions, density_parameters, mesh); + + DensityQFunction density_qf; + density_mf->AddDomainIntegrator(density_qf, density_kernel_ao, + density_kernel_oo, ir, all_domain_attr); + } + + return LagrangianHydroOperator( + H1, + L2, + ess_tdof, + ir, + rho0_coeff, + x0_gf, + rho0_gf, + material_gf, + update_qdata, + dt_est, + momentum_mf, + momentum_pa, + energy_conservation_mf, + energy_conservation_pa, + total_internal_energy_mf, + total_kinetic_energy_mf, + density_mf, + qdata, + fd_gradient, + nonlinear_maximum_iterations, + nonlinear_relative_tolerance); +} + +int main(int argc, char *argv[]) +{ + Mpi::Init(); + Hypre::Init(); + + const char *device_config = "cpu"; + + const char *mesh_file = + "/Users/andrej1/repos/Laghos/data/rectangle01_quad.mesh"; + + int refinements = 0; + int order_v = 2; + int order_e = 1; + int order_q = -1; + real_t t_final = 0.0; + real_t blast_energy = 0.25; + real_t blast_position[] = {0.0, 0.0, 0.0}; + int ode_solver_type = 4; + bool fd_gradient = false; + bool use_viscosity = false; + real_t cfl = 0.5; + real_t nonlinear_relative_tolerance = 1e-5; + int nonlinear_maximum_iterations = 10; + + OptionsParser args(argc, argv); + args.AddOption(&mesh_file, "-m", "--mesh", + "Mesh file to use."); + args.AddOption(&refinements, "-rs", "--ref", ""); + args.AddOption(&order_v, "-ov", "--ov", ""); + args.AddOption(&order_e, "-oe", "--oe", ""); + args.AddOption(&order_q, "-oq", "--oq", ""); + args.AddOption(&t_final, "-tf", "--tf", ""); + args.AddOption(&problem, "-p", "--p", ""); + args.AddOption(&cfl, "-cfl", "--cfl", ""); + args.AddOption(&device_config, "-d", "--device", + "Device configuration string, see Device::Configure()."); + args.AddOption(&use_viscosity, "-av", "--av", "-no-av", "--no-av", ""); + args.AddOption(&fd_gradient, "-fd", "--fd", "-no-fd", "--no-fd", ""); + args.AddOption(&ode_solver_type, "-s", "--ode-solver", + "ODE solver: 1 - Forward Euler,\n\t" + " 2 - RK2 SSP, 3 - RK3 SSP, 4 - RK4, 6 - RK6,\n\t" + " 7 - RK2Avg." + " 11 - Backward Euler" + " 12 - Implicit Midpoint" + " 13 - SDIRK33Solver"); + args.AddOption(&nonlinear_maximum_iterations, "-nmi", "--nmi", + "Maximum number of nonlinear iterations."); + args.AddOption(&nonlinear_relative_tolerance, "-nrt", "--nrt", + "Nonlinear relative tolerance."); + args.ParseCheck(); + + Device device(device_config); + if (Mpi::Root()) { device.Print(); } + + Mesh serial_mesh = Mesh(mesh_file, true, true); + + if (problem == 0 || problem == 1) + { + serial_mesh = Mesh(Mesh::MakeCartesian2D(1, 1, Element::QUADRILATERAL, + true)); + + const int NBE = serial_mesh.GetNBE(); + for (int b = 0; b < NBE; b++) + { + Element *bel = serial_mesh.GetBdrElement(b); + const int attr = (b < NBE/2) ? 2 : 1; + bel->SetAttribute(attr); + } + } + + if (problem == 2) + { + serial_mesh = Mesh(Mesh::MakeCartesian1D(2)); + serial_mesh.GetBdrElement(0)->SetAttribute(1); + serial_mesh.GetBdrElement(1)->SetAttribute(1); + } + + for (int i = 0; i < refinements; i++) + { + serial_mesh.UniformRefinement(); + } + + // serial_mesh.EnsureNCMesh(); + // serial_mesh.RandomRefinement(0.1); + + ParMesh mesh = ParMesh(MPI_COMM_WORLD, serial_mesh); + const int dim = mesh.Dimension(); + + MFEM_ASSERT(dim == DIMENSION, "mesh dimension inconsistency"); + + // Define the parallel finite element spaces. We use: + // - H1 (Gauss-Lobatto, continuous) for position and velocity. + // - L2 (Bernstein, discontinuous) for specific internal energy. + H1_FECollection H1FEC(order_v, dim); + ParFiniteElementSpace H1FESpace(&mesh, &H1FEC, dim); + L2_FECollection L2FEC(order_e, dim, BasisType::Positive); + ParFiniteElementSpace L2FESpace(&mesh, &L2FEC); + + const auto global_ne = mesh.GetGlobalNE(); + const auto global_h1tsize = H1FESpace.GlobalTrueVSize(); + const auto global_l2tsize = L2FESpace.GlobalTrueVSize(); + + if (Mpi::Root()) + { + out << "num el: " << global_ne << "\n"; + out << "num kinematic dofs: " << global_h1tsize << "\n"; + out << "num thermodynamic dofs: " << global_l2tsize << "\n"; + } + + Array ess_tdof, ess_vdofs; + { + Array ess_bdr(mesh.bdr_attributes.Max()), dofs_marker, dofs_list; + for (int d = 0; d < mesh.Dimension(); d++) + { + // Attributes 1/2/3 correspond to fixed-x/y/z boundaries, + // i.e., we must enforce v_x/y/z = 0 for the velocity components. + ess_bdr = 0; ess_bdr[d] = 1; + H1FESpace.GetEssentialTrueDofs(ess_bdr, dofs_list, d); + ess_tdof.Append(dofs_list); + H1FESpace.GetEssentialVDofs(ess_bdr, dofs_marker, d); + FiniteElementSpace::MarkerToList(dofs_marker, dofs_list); + ess_vdofs.Append(dofs_list); + } + } + + // The monolithic BlockVector stores unknown fields as: + // - 0 -> position + // - 1 -> velocity + // - 2 -> specific internal energy + const int Vsize_l2 = L2FESpace.GetVSize(); + const int Vsize_h1 = H1FESpace.GetVSize(); + Array offset(4); + offset[0] = 0; + offset[1] = offset[0] + Vsize_h1; + offset[2] = offset[1] + Vsize_h1; + offset[3] = offset[2] + Vsize_l2; + BlockVector S(offset, Device::GetDeviceMemoryType()); + + ParGridFunction x_gf, v_gf, e_gf; + x_gf.MakeRef(&H1FESpace, S, offset[0]); + v_gf.MakeRef(&H1FESpace, S, offset[1]); + e_gf.MakeRef(&L2FESpace, S, offset[2]); + + mesh.SetNodalGridFunction(&x_gf); + x_gf.SyncAliasMemory(S); + + ParGridFunction x0_gf = x_gf; + + auto v0 = [](const Vector &x, Vector &v) + { + switch (problem) + { + case 0: + v(0) = sin(M_PI*x(0)) * cos(M_PI*x(1)); + v(1) = -cos(M_PI*x(0)) * sin(M_PI*x(1)); + if (x.Size() == 3) + { + v(0) *= cos(M_PI*x(2)); + v(1) *= cos(M_PI*x(2)); + v(2) = 0.0; + } + break; + case 1: v = 0.0; break; + case 2: v = 0.0; break; + case 3: v = 0.0; break; + default: MFEM_ABORT("error"); + } + }; + + VectorFunctionCoefficient v_coeff(dim, v0); + v_gf.ProjectCoefficient(v_coeff); + for (int i = 0; i < ess_vdofs.Size(); i++) + { + v_gf(ess_vdofs[i]) = 0.0; + } + v_gf.SyncAliasMemory(S); + + auto rho0 = [&dim](const Vector &x) + { + switch (problem) + { + case 0: return 1.0; + case 1: return 1.0; + case 2: return (x(0) < 0.5) ? 1.0 : 0.1; + case 3: return (dim == 2) ? (x(0) > 1.0 && x(1) > 1.5) ? 0.125 : 1.0 + : x(0) > 1.0 && ((x(1) < 1.5 && x(2) < 1.5) || + (x(1) > 1.5 && x(2) > 1.5)) ? 0.125 : 1.0; + default: MFEM_ABORT("error"); + } + }; + + ParGridFunction rho0_gf(&L2FESpace); + FunctionCoefficient rho0_coeff(rho0); + L2_FECollection l2_fec(order_e, mesh.Dimension()); + ParFiniteElementSpace l2_fes(&mesh, &l2_fec); + ParGridFunction l2_rho0_gf(&l2_fes), l2_e(&l2_fes); + l2_rho0_gf.ProjectCoefficient(rho0_coeff); + rho0_gf.ProjectGridFunction(l2_rho0_gf); + + auto gamma_func = [](const Vector &x) + { + switch (problem) + { + case 0: return 5.0 / 3.0; + case 1: return 1.4; + case 2: return 1.4; + case 3: return (x(0) > 1.0 && x(1) <= 1.5) ? 1.4 : 1.5; + default: MFEM_ABORT("error"); + } + }; + + auto e0 = [&rho0, &gamma_func](const Vector &x) + { + switch (problem) + { + case 0: + { + const real_t denom = 2.0 / 3.0; // (5/3 - 1) * density. + real_t val; + if (x.Size() == 2) + { + val = 1.0 + (cos(2*M_PI*x(0)) + cos(2*M_PI*x(1))) / 4.0; + } + else + { + val = 100.0 + ((cos(2*M_PI*x(2)) + 2) * + (cos(2*M_PI*x(0)) + cos(2*M_PI*x(1))) - 2) / 16.0; + } + return val/denom; + } + case 1: return 0.0; // This case in initialized in main(). + case 2: return (x(0) < 0.5) ? 1.0 / rho0(x) / (gamma_func(x) - 1.0) + : 0.1 / rho0(x) / (gamma_func(x) - 1.0); + case 3: return (x(0) > 1.0) ? 0.1 / rho0(x) / (gamma_func(x) - 1.0) + : 1.0 / rho0(x) / (gamma_func(x) - 1.0); + default: MFEM_ABORT("error"); + } + }; + + if (problem == 1) + { + DeltaCoefficient e_coeff(blast_position[0], blast_position[1], + blast_position[2], blast_energy); + l2_e.ProjectCoefficient(e_coeff); + } + else + { + FunctionCoefficient e_coeff(e0); + l2_e.ProjectCoefficient(e_coeff); + } + + e_gf.ProjectGridFunction(l2_e); + e_gf.SyncAliasMemory(S); + + L2_FECollection material_fec(0, dim); + ParFiniteElementSpace L2CFESpace(&mesh, &material_fec); + ParGridFunction material_gf(&L2CFESpace); + FunctionCoefficient material_coeff(gamma_func); + material_gf.ProjectCoefficient(material_coeff); + + ParGridFunction rho_gf(&L2FESpace); + + IntegrationRule ir = IntRules.Get(mesh.GetElementBaseGeometry(0), + 3 * H1FESpace.GetOrder(0) + L2FESpace.GetOrder(0) - 1); + + if (Mpi::Root()) + { + out << "num qp: " << ir.GetNPoints() << "\n"; + } + + // Create external data vector + // Layout is [cfl, order_velocity, use_viscosity, h0] + Vector external_data(4); + external_data[0] = cfl; + external_data[1] = order_v; + external_data[2] = use_viscosity; + + auto hydro = CreateLagrangianHydroOperator(H1FESpace, + L2FESpace, + ess_tdof, + rho0_coeff, + x0_gf, + rho0_gf, + material_gf, + external_data, + ir, + fd_gradient, + nonlinear_maximum_iterations, + nonlinear_relative_tolerance); + + ODESolver *ode_solver = NULL; + switch (ode_solver_type) + { + case 1: ode_solver = new ForwardEulerSolver; break; + case 2: ode_solver = new RK2Solver(0.5); break; + case 3: ode_solver = new RK3SSPSolver; break; + case 4: ode_solver = new RK4Solver; break; + case 6: ode_solver = new RK6Solver; break; + case 11: ode_solver = new BackwardEulerSolver; break; + case 12: ode_solver = new ImplicitMidpointSolver; break; + case 13: ode_solver = new SDIRK33Solver; break; + default: + out << "Unknown ODE solver type: " << ode_solver_type << '\n'; + return -1; + } + ode_solver->Init(hydro); + + hydro.ComputeDensity(rho_gf); + const real_t energy_init = hydro.InternalEnergy(e_gf) + + hydro.KineticEnergy(v_gf); + + if (Mpi::Root()) + { + out << "energy initial: " << energy_init << "\n"; + } + + out << "IE " << hydro.InternalEnergy(e_gf) << "\n" + << "KE "<< hydro.KineticEnergy(v_gf) << "\n"; + + real_t t = 0.0; + real_t dt = hydro.GetTimeStepEstimate(S); + out << "time step estimate: " << dt << "\n"; + real_t t_old; + bool last_step = false; + [[maybe_unused]] int steps = 0; + BlockVector S_old(S); + + ParGridFunction verr_gf(v_gf); + verr_gf.ProjectCoefficient(v_coeff); + v_gf.SyncAliasMemory(S); + v_gf.HostRead(); + verr_gf.HostReadWrite(); + for (int i = 0; i < verr_gf.Size(); i++) + { + verr_gf(i) = abs(verr_gf(i) - std::as_const(v_gf)(i)); + } + + ParaViewDataCollection paraview_dc("dfem", &mesh); + paraview_dc.SetPrefixPath("ParaView"); + paraview_dc.SetLevelsOfDetail(order_v); + paraview_dc.SetDataFormat(VTKFormat::BINARY); + paraview_dc.SetHighOrderOutput(true); + paraview_dc.SetCycle(0); + paraview_dc.SetTime(0.0); + paraview_dc.RegisterField("velocity", &v_gf); + paraview_dc.RegisterField("density", &rho_gf); + paraview_dc.RegisterField("specific_internal_energy", &e_gf); + paraview_dc.RegisterField("material", &material_gf); + // paraview_dc.RegisterField("velocity_error", &verr_gf); + + paraview_dc.SetCycle(0); + paraview_dc.SetTime(0); + paraview_dc.Save(); + + for (int ti = 1; !last_step; ti++) + { + if (t + dt >= t_final) + { + dt = t_final - t; + last_step = true; + } + S_old = S; + t_old = t; + + // S is the vector of dofs, t is the current time, and dt is the time step + // to advance. + ode_solver->Step(S, t, dt); + steps++; + + // Adaptive time step control. + const real_t dt_est = hydro.GetTimeStepEstimate(S); + if (dt_est < dt) + { + // Repeat (solve again) with a decreased time step - decrease of the + // time estimate suggests appearance of oscillations. + dt *= 0.85; + if (dt < std::numeric_limits::epsilon()) + { MFEM_ABORT("The time step crashed!"); } + t = t_old; + S = S_old; + if (Mpi::Root()) { out << "Repeating step " << ti << std::endl; } + ti--; continue; + } + else if (dt_est > 1.25 * dt) { dt *= 1.02; } + + x_gf.SyncAliasMemory(S); + v_gf.SyncAliasMemory(S); + e_gf.SyncAliasMemory(S); + + // Make sure that the mesh corresponds to the new solution state. This is + // needed, because some time integrators use different S-type vectors + // and the oper object might have redirected the mesh positions to those. + mesh.NewNodes(x_gf, false); + + // out << "x_gf outer loop\n"; + // print_vector(x_gf); + + if (Mpi::Root()) + { + out << "step " << std::setw(5) << ti + << ",\tt = " << std::setw(5) << std::setprecision(4) << t + << ",\tdt = " << std::setw(5) << std::setprecision(6) << dt; + out << std::endl; + } + + // verr_gf.ProjectCoefficient(v_coeff); + // for (int i = 0; i < verr_gf.Size(); i++) + // { + // verr_gf(i) = abs(verr_gf(i) - v_gf(i)); + // } + + hydro.ComputeDensity(rho_gf); + + paraview_dc.SetCycle(ti); + paraview_dc.SetTime(t); + paraview_dc.Save(); + } + + const real_t energy_final = hydro.InternalEnergy(e_gf) + + hydro.KineticEnergy(v_gf); + const real_t v_err_max = v_gf.ComputeMaxError(v_coeff); + const real_t v_err_l1 = v_gf.ComputeL1Error(v_coeff); + const real_t v_err_l2 = v_gf.ComputeL2Error(v_coeff); + + if (Mpi::Root()) + { + out << std::scientific << std::setprecision(2) + << "Energy diff: " << fabs(energy_init - energy_final) << std::endl + << "L_inf error: " << v_err_max << std::endl + << "L_1 error: " << v_err_l1 << std::endl + << "L_2 error: " << v_err_l2 << std::endl; + } + + return 0; +} diff --git a/examples/dfem/plasticity.cpp b/examples/dfem/plasticity.cpp new file mode 100644 index 0000000000..4ec536ae30 --- /dev/null +++ b/examples/dfem/plasticity.cpp @@ -0,0 +1,597 @@ +// 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 + +// TODO: Do we want this to be included from mfem.hpp automatically now? +#include +#include + +#include + +using namespace mfem; + +using mfem::future::tuple; +using mfem::future::tensor; + +using future::DifferentiableOperator; +using future::ParametricSpace; +using future::ParametricFunction; +using future::FieldDescriptor; +using future::Gradient; +using future::Weight; +using future::None; + +constexpr int DIMENSION = 2; + +template +MFEM_HOST_DEVICE inline +tensor tensor_to_3D(const tensor& A) +{ + tensor A3D{}; + for (int i = 0; i < dim; i++) + { + for (int j = 0; j < dim; j++) + { + A3D[i][j] = A[i][j]; + } + } + return A3D; +} + +template +struct InternalStateQFunction +{ + InternalStateQFunction() = default; + + MFEM_HOST_DEVICE inline + auto operator()( + const tensor &dudxi, + const tensor &J, + const tensor &internal_state, + const double &w) const + { + auto invJ = inv(J); + auto dudX = dudxi * invJ; + auto dudX3D = tensor_to_3D(dudX); + //auto internal_state_new = get<1>(material(dudX3D, internal_state)); + auto [stress, internal_state_new] = material(dudX3D, internal_state); + // real_t vm = sqrt(1.5)*norm(dev(stress)); + // out << vm << " " << internal_state_new[9] << std::endl; + return tuple{internal_state_new}; + } + + Material material; +}; + +template +struct MomentumRefStateQFunction +{ + MomentumRefStateQFunction() = default; + + MFEM_HOST_DEVICE inline + auto operator()( + const tensor &dudxi, + const tensor &J, + const tensor &internal_state, + const double &w) const + { + auto invJ = inv(J); + auto dudX = dudxi * invJ; + auto dudX3D = tensor_to_3D(dudX); + auto [P3D, Qnew] = material(dudX3D, internal_state); + auto P = future::make_tensor([&P3D](int i, int j) { return P3D[i][j]; }); + auto JxW = det(J) * w * transpose(invJ); + return tuple{P * JxW}; + } + + Material material; +}; + + +struct J2SmallStrain +{ + static constexpr int dim = 3; ///< spatial dimension + static constexpr int n_internal_states = 10; + static constexpr double tol = + 1e-10; ///< relative tolerance on residual mag to judge convergence of return map + + real_t E; ///< Young's modulus + real_t nu; ///< Poisson's ratio + real_t sigma_y; ///< Yield strength + real_t Hi; ///< Isotropic hardening modulus + real_t density; ///< Mass density + + /// @brief variables required to characterize the hysteresis response + struct InternalState + { + tensor plastic_strain; ///< plastic strain + double accumulated_plastic_strain; ///< uniaxial equivalent plastic strain + }; + + MFEM_HOST_DEVICE inline + InternalState unpack_internal_state(const tensor & + packed_state) const + { + // we could use type punning here to avoid copies + auto plastic_strain = future::make_tensor( + [&packed_state](int i, int j) { return packed_state[dim*i + j]; }); + real_t accumulated_plastic_strain = packed_state[n_internal_states - 1]; + return {plastic_strain, accumulated_plastic_strain}; + } + + MFEM_HOST_DEVICE inline + tensor pack_internal_state(const + tensor & plastic_strain, + real_t accumulated_plastic_strain) const + { + tensor packed_state{}; + for (int i = 0, ij = 0; i < dim; i++) + { + for (int j = 0; j < dim; j++, ij++) + { + packed_state[ij] = plastic_strain[i][j]; + } + } + packed_state[n_internal_states - 1] = accumulated_plastic_strain; + return packed_state; + } + + MFEM_HOST_DEVICE inline + tuple, tensor> + operator()(const tensor & dudX, + const tensor & internal_state) const + { + auto I = future::Identity(); + const real_t K = E / (3.0 * (1.0 - 2.0 * nu)); + const real_t G = 0.5 * E / (1.0 + nu); + + auto [plastic_strain, accumulated_plastic_strain] = unpack_internal_state( + internal_state); + + // (i) elastic predictor + auto el_strain = sym(dudX) - plastic_strain; + auto p = K * tr(el_strain); + auto s = 2.0 * G * dev(el_strain); + auto q = sqrt(1.5) * norm(s); + [[maybe_unused]] real_t delta_eqps = 0.0; + + [[maybe_unused]] auto flow_strength = [this](real_t eqps) { return this->sigma_y + this->Hi*eqps; }; + + // (ii) admissibility + if (q - (sigma_y + Hi*accumulated_plastic_strain) > tol*sigma_y) + { + // (iii) return mapping + real_t delta_eqps = (q - sigma_y - Hi*accumulated_plastic_strain)/(3*G + Hi); + auto Np = 1.5 * s / q; + s -= 2.0 * G * delta_eqps * Np; + plastic_strain += delta_eqps * Np; + accumulated_plastic_strain += delta_eqps; + } + auto stress = s + p * I; + auto internal_state_new = pack_internal_state(plastic_strain, + accumulated_plastic_strain); + return {stress, internal_state_new}; + } +}; + +class ElasticityOperator : public Operator +{ + static constexpr int Displacement = 0; + static constexpr int Coordinates = 1; + static constexpr int InternalState = 2; + +public: + class ElasticityJacobianOperator : public Operator + { + public: + ElasticityJacobianOperator(const ElasticityOperator *elasticity, + const Vector &x) : + Operator(elasticity->Height()), + elasticity(elasticity), + z(elasticity->Height()) + { + ParGridFunction u(&elasticity->displacement_fes); + u.SetFromTrueDofs(x); + auto mesh_nodes = static_cast + (elasticity->displacement_fes.GetParMesh()->GetNodes()); + momentum_du = elasticity->momentum->GetDerivative(Displacement, {&u}, {mesh_nodes, &elasticity->internal_state}); + } + + void Mult(const Vector &x, Vector &y) const override + { + z = x; + z.SetSubVector(elasticity->displacement_ess_tdof, 0.0); + + momentum_du->Mult(z, y); + + for (int i = 0; i < elasticity->displacement_ess_tdof.Size(); i++) + { + y[elasticity->displacement_ess_tdof[i]] = + x[elasticity->displacement_ess_tdof[i]]; + } + } + + const ElasticityOperator *elasticity; + std::shared_ptr momentum_du; + mutable Vector z; + }; + + template + ElasticityOperator(ParFiniteElementSpace &displacement_fes, + Array &vel_ess_tdofs, + const IntegrationRule &displacement_ir, + ParametricFunction &internal_state, + Material material) : + Operator(displacement_fes.GetTrueVSize()), + density(1.0e3), + body_force(displacement_fes.GetTrueVSize()), + displacement_ess_tdof(vel_ess_tdofs), + displacement_fes(displacement_fes), + displacement_ir(displacement_ir), + internal_state(internal_state) + { + auto mesh = displacement_fes.GetParMesh(); + mesh_nodes = static_cast(mesh->GetNodes()); + ParFiniteElementSpace& mesh_fes = *mesh_nodes->ParFESpace(); + + { + auto solutions = std::vector + { + FieldDescriptor{Displacement, &displacement_fes}, + }; + + auto parameters = std::vector + { + FieldDescriptor{Coordinates, &mesh_fes}, + FieldDescriptor{InternalState, &internal_state.space} + }; + + momentum = + std::make_shared(solutions, parameters, *mesh); + momentum->DisableTensorProductStructure(); + + tuple inputs{Gradient{}, Gradient{}, None{}, Weight{}}; + tuple outputs{Gradient{}}; + + auto momentum_qf = MomentumRefStateQFunction {.material = material}; + auto derivatives = std::integer_sequence {}; + Array solid_domain_attr(mesh->attributes.Max()); + solid_domain_attr[0] = 1; + momentum->AddDomainIntegrator( + momentum_qf, inputs, outputs, displacement_ir, solid_domain_attr, derivatives); + } + + { + Vector g(DIMENSION); + g = 0.0; + + ParLinearForm body_force_lf(&displacement_fes); + body_force_coef = new VectorConstantCoefficient(g); + auto integ = new VectorDomainLFIntegrator(*body_force_coef); + integ->SetIntRule(&displacement_ir); + body_force_lf.AddDomainIntegrator(integ); + body_force_lf.Assemble(); + body_force_lf.ParallelAssemble(body_force); + } + } + + void Mult(const Vector &displacement, Vector &r) const override + { + momentum->SetParameters({mesh_nodes, &internal_state}); + momentum->Mult(displacement, r); + r -= body_force; + r.SetSubVector(displacement_ess_tdof, 0.0); + } + + void Reaction(const Vector &displacement, Vector &r) const + { + momentum->SetParameters({mesh_nodes, &internal_state}); + momentum->Mult(displacement, r); + r -= body_force; + r.Neg(); + } + + Operator &GetGradient(const Vector &x) const override + { + jacobian_operator = std::make_shared(this, x); + return *jacobian_operator; + + // fd_jacobian = std::make_shared(*this, x); + // return *fd_jacobian; + } + + real_t density; + std::shared_ptr momentum; + mutable std::shared_ptr A; + VectorConstantCoefficient *body_force_coef = nullptr; + Vector body_force; + + ParGridFunction *mesh_nodes; + + const Array displacement_ess_tdof; + + ParFiniteElementSpace &displacement_fes; + IntegrationRule displacement_ir; + + ParametricFunction& internal_state; + + mutable std::shared_ptr jacobian_operator; + mutable std::shared_ptr fd_jacobian; +}; + + +class InternalStateUpdater : public Operator +{ +public: + + static constexpr int Displacement = 0; + static constexpr int Coordinates = 1; + static constexpr int InternalState = 2; + + template + InternalStateUpdater(ParFiniteElementSpace &displacement_fes, + const IntegrationRule &displacement_ir, + ParametricFunction &internal_state, + Material material) : + Operator(displacement_fes.GetTrueVSize()), + displacement_fes(displacement_fes), + displacement_ir(displacement_ir), + internal_state(internal_state) + { + auto mesh = displacement_fes.GetParMesh(); + mesh_nodes = static_cast(mesh->GetNodes()); + ParFiniteElementSpace& mesh_fes = *mesh_nodes->ParFESpace(); + + auto solutions = std::vector + { + FieldDescriptor{Displacement, &displacement_fes} + }; + + auto parameters = std::vector + { + FieldDescriptor{Coordinates, &mesh_fes}, + FieldDescriptor{InternalState, &internal_state.space} + }; + + op = std::make_shared(solutions, parameters, *mesh); + op->DisableTensorProductStructure(); + + tuple inputs{Gradient{}, Gradient{}, None{}, Weight{}}; + tuple outputs{None{}}; + + auto qfunction = InternalStateQFunction {.material = material}; + // just a placeholder for now. We want vjps wrt both displacement and old internal state eventually + auto derivatives = std::integer_sequence {}; + Array solid_domain_attr(mesh->attributes.Max()); + solid_domain_attr[0] = 1; + op->AddDomainIntegrator( + qfunction, inputs, outputs, displacement_ir, solid_domain_attr, derivatives); + } + + void Mult(const Vector &displacement, Vector& internal_state_new) const override + { + op->SetParameters({mesh_nodes, &internal_state}); + op->Mult(displacement, internal_state_new); + } + + void VjpDisplacement(ParGridFunction &u, Vector& internal_state_old, + Vector& internal_state_new_bar, Vector& displacement_bar) const + { + // u, internal_state_old, internal_state_new_bar should be const + out << "Sizes " << "u " << u.Size() << ", qold " << internal_state_old.Size() << + ", qbar " << internal_state_new_bar.Size() << ", ubar " << + displacement_bar.Size() << std::endl; + auto grad_op = op->GetDerivative(Displacement, {&u}, {mesh_nodes, &internal_state_old}); + out << "grad_op " << grad_op->Height() << " " << grad_op->Width() << std::endl; + out << "grad_op^T " << grad_op->Width() << " " << grad_op->Height() << + std::endl; + grad_op->MultTranspose(internal_state_new_bar, displacement_bar); + } + + ParGridFunction *mesh_nodes; + ParFiniteElementSpace &displacement_fes; + std::shared_ptr op; + IntegrationRule displacement_ir; + ParametricFunction& internal_state; +}; + + +int main(int argc, char* argv[]) +{ + constexpr int dim = 2; + + Mpi::Init(); + + const char* device_config = "cpu"; + int polynomial_order = 1; + int ir_order = 2; + int refinements = 0; + int nonlinear_solver_type = 0; + + OptionsParser args(argc, argv); + args.AddOption(&polynomial_order, "-o", "--order", ""); + args.AddOption(&refinements, "-r", "--refinements", ""); + args.AddOption(&ir_order, "-iro", "--integration-rule-order", ""); + args.AddOption(&device_config, "-d", "--device", + "Device configuration string, see Device::Configure()."); + args.AddOption(&nonlinear_solver_type, "-nls", "--nonlinear-solver", ""); + args.ParseCheck(); + + Device device(device_config); + if (Mpi::Root() == 0) + { + device.Print(); + } + + out << std::setprecision(8); + + Mesh mesh_serial = Mesh::MakeCartesian2D(1, 1, Element::QUADRILATERAL, + false, 1.0, 0.1); + mesh_serial.EnsureNodes(); + auto mesh_beam = ParMesh(MPI_COMM_WORLD, mesh_serial); + + out << "#el: " << mesh_beam.GetNE() << "\n"; + + H1_FECollection displacement_fec(polynomial_order, dim); + ParFiniteElementSpace displacement_fes(&mesh_beam, &displacement_fec, dim); + + HYPRE_BigInt global_size = displacement_fes.GlobalTrueVSize(); + if (Mpi::Root()) + { + out << "Number of unknowns: " << global_size << "\n"; + } + + const IntegrationRule &displacement_ir = + IntRules.Get(displacement_fes.GetFE(0)->GetGeomType(), + 2 * ir_order + displacement_fes.GetFE(0)->GetOrder()); + + constexpr int n_internal_state_variables = 10; + ParametricSpace internal_state_space(dim, n_internal_state_variables, + displacement_ir.GetNPoints(), + n_internal_state_variables*displacement_ir.GetNPoints()*mesh_beam.GetNE()); + + ParametricFunction internal_state(internal_state_space); + internal_state = 0.0; + ParametricFunction internal_state_old(internal_state_space); + internal_state_old = 0.0; + + Array bdr_attr_is_ess(mesh_beam.bdr_attributes.Max()); + Array displacement_ess_tdof; + Array bc_tdof; + + bdr_attr_is_ess = 0; + bdr_attr_is_ess[0] = 1; + displacement_fes.GetEssentialTrueDofs(bdr_attr_is_ess, bc_tdof, 1); + for (auto td : bc_tdof) { displacement_ess_tdof.Append(td); }; + + bdr_attr_is_ess = 0; + bdr_attr_is_ess[3] = 1; + displacement_fes.GetEssentialTrueDofs(bdr_attr_is_ess, bc_tdof, 0); + for (auto td : bc_tdof) { displacement_ess_tdof.Append(td); }; + + bdr_attr_is_ess = 0; + bdr_attr_is_ess[1] = 1; + displacement_fes.GetEssentialTrueDofs(bdr_attr_is_ess, bc_tdof, 0); + for (auto td : bc_tdof) { displacement_ess_tdof.Append(td); }; + + ParGridFunction u(&displacement_fes); + u = 0.0; + + using Material = J2SmallStrain; // StVenantKirchhoff + Material material{.E = 1000.0, .nu = 0.25, .sigma_y = 0.53333, .Hi = 40.0, .density = 1.0}; + // Material material{.mu = 0.5e6, .nu = 0.4}; + + ElasticityOperator elasticity(displacement_fes, displacement_ess_tdof, + displacement_ir, internal_state, material); + + CGSolver solver(MPI_COMM_WORLD); + solver.SetAbsTol(0.0); + solver.SetRelTol(1e-10); + solver.SetMaxIter(1000); + solver.SetPrintLevel(2); + + std::shared_ptr nonlinear_solver; + if (nonlinear_solver_type == 0) + { + nonlinear_solver = std::make_shared(MPI_COMM_WORLD); + } + // else if (nonlinear_solver_type == 1) + // { + // nonlinear_solver = std::make_shared(MPI_COMM_WORLD, KIN_LINESEARCH); + // } + else + { + MFEM_ABORT("invalid nonlinear solver type"); + } + nonlinear_solver->SetOperator(elasticity); + nonlinear_solver->SetRelTol(1e-9); + nonlinear_solver->SetMaxIter(25); + nonlinear_solver->SetSolver(solver); + nonlinear_solver->SetPrintLevel(1); + + // variables for output + QuadratureSpace output_internal_state_space(mesh_beam, displacement_ir); + QuadratureFunction output_internal_state(&output_internal_state_space, + internal_state.GetData(), material.n_internal_states); + Vector r(displacement_fes.GetTrueVSize()); + ParGridFunction reaction(&displacement_fes); + Vector end_forces_x(bc_tdof.Size()); + + ParaViewDataCollection dc("dfem_plasticity", &mesh_beam); + dc.SetHighOrderOutput(true); + dc.SetLevelsOfDetail(1); + dc.RegisterField("displacement", &u); + dc.RegisterField("reaction", &reaction); + dc.RegisterQField("internal_state", &output_internal_state); + dc.SetCycle(0); + dc.Save(); + + InternalStateUpdater internal_state_update(displacement_fes, displacement_ir, + internal_state, material); + //Vector q(internal_state_space.GetTotalSize()); + + auto applied_displacement = [](double t) { return 1.2e-2*t; }; + + real_t time = 0.0; + std::ofstream history_file("history_output.csv"); + history_file << applied_displacement(time) << " " << 0.0 << std::endl; + + Vector zero, x(displacement_fes.GetTrueVSize()); + + constexpr int max_cycles = 3; + const real_t dt = 1.0/(max_cycles - 1); + for (int cycle = 1; cycle < max_cycles; cycle++) + { + time += dt; + out << "-------------------------------------------" << std::endl; + out << "TIME STEP " << cycle << std::endl; + out << "t = " << time << std::endl; + + real_t ubc = applied_displacement(time); + u.SetSubVector(bc_tdof, ubc); + + u.GetTrueDofs(x); + nonlinear_solver->Mult(zero, x); + u.SetFromTrueDofs(x); + + // update internal variables + internal_state_old.Set(1.0, internal_state); + internal_state_update.Mult(u, internal_state); + + // Compute reactions + elasticity.Reaction(x, r); + reaction.SetFromTrueDofs(r); + reaction.GetSubVector(bc_tdof, end_forces_x); + real_t force = -end_forces_x.Sum(); + out << "u = " << applied_displacement(time) << ", Force = " << force << + std::endl; + history_file << applied_displacement(time) << " " << force << std::endl; + + output_internal_state = internal_state; + + dc.SetCycle(cycle); + dc.SetTime(time); + dc.Save(); + } + + // try to use the derivative to see if it works + ParametricFunction internal_state_bar(internal_state_space); + internal_state_bar = 1.0; + //ParGridFunction u_bar(displacement_fes); + Vector u_bar(displacement_fes.GetTrueVSize()); + internal_state_update.VjpDisplacement(u, internal_state_old, internal_state_bar, + u_bar); + + future::pretty_print(u_bar); + + history_file.close(); + return 0; +} diff --git a/fem/dfem/doperator.hpp b/fem/dfem/doperator.hpp index 1590258871..c84deaed22 100644 --- a/fem/dfem/doperator.hpp +++ b/fem/dfem/doperator.hpp @@ -23,6 +23,10 @@ #include "qf_derivative_dual.hpp" #include "integrate.hpp" +#undef NVTX_COLOR +#define NVTX_COLOR nvtx::kPurple +#include "general/nvtx.hpp" + namespace mfem::future { @@ -232,10 +236,13 @@ public: /// solutions_t. The result is a T-dof vector. void Mult(const Vector &solutions_t, Vector &result_t) const override { + // dbg(); MFEM_ASSERT(!action_callbacks.empty(), "no integrators have been set"); + // dbg("prolongation"); prolongation(solutions, solutions_t, solutions_l); for (auto &action : action_callbacks) { + // dbg("action"); action(solutions_l, parameters_l, residual_l); } prolongation_transpose(residual_l, result_t); @@ -478,6 +485,7 @@ void DifferentiableOperator::AddDomainIntegrator( auto &output_e_size = output_e_sz; output_restriction_transpose = output_rt; + residual_e.UseDevice(true); residual_e.SetSize(output_e_size); // The explicit captures are necessary to avoid dependency on @@ -502,6 +510,13 @@ void DifferentiableOperator::AddDomainIntegrator( [[maybe_unused]] const int num_elements = GetNumEntities(mesh); const int num_entities = GetNumEntities(mesh); const int num_qp = integration_rule.GetNPoints(); + dbg("num_qp:{}", num_qp); + dbg("num_entities:{}", num_entities); + dbg("dimension:{}", dimension); + dbg("num_fields:{}", num_fields); + dbg("num_inputs:{}", num_inputs); + dbg("num_outputs:{}", num_outputs); + dbg("num_elements:{}", num_entities); if constexpr (is_one_fop::value) { @@ -514,9 +529,12 @@ void DifferentiableOperator::AddDomainIntegrator( residual_l.SetSize(residual_lsize); height = GetTrueVSize(fields[test_space_field_idx]); } + dbg("residual_lsize:{}", residual_l.Size()); + dbg("height:{}", height); // TODO: Is this a hack? width = GetTrueVSize(fields[0]); + dbg("width:{}", width); std::vector dtq; for (const auto &field : fields) @@ -527,10 +545,12 @@ void DifferentiableOperator::AddDomainIntegrator( doftoquad_mode)); } const int q1d = (int)floor(std::pow(num_qp, 1.0/dimension) + 0.5); + dbg("q1d:{}", q1d); const int residual_size_on_qp = GetSizeOnQP(output_fop, fields[test_space_field_idx]); + dbg("residual_size_on_qp:{}", residual_size_on_qp); auto input_dtq_maps = create_dtq_maps(inputs, dtq, input_to_field); auto output_dtq_maps = create_dtq_maps(outputs, dtq, output_to_field); @@ -580,6 +600,8 @@ void DifferentiableOperator::AddDomainIntegrator( { restriction_cb(sol, par, fields_e); + // MFEM_GPU_CHECK(hipGetLastError()); + // dbg("residual_e = 0.0"); residual_e = 0.0; auto ye = Reshape(residual_e.ReadWrite(), test_vdim, num_test_dof, num_entities); @@ -610,6 +632,13 @@ void DifferentiableOperator::AddDomainIntegrator( auto fhat = Reshape(&residual_shmem(0, 0), test_vdim, test_op_dim, num_qp); auto y = Reshape(&ye(0, 0, e), num_test_dof, test_vdim); + + const DofToQuadMap &dtq_0 = output_dtq_shmem[0]; + // dbg("dtq_0.B.GetShape()[0]:{}", dtq_0.B.GetShape()[0]); + // dbg("dtq_0.B.GetShape()[1]:{}", dtq_0.B.GetShape()[1]); + // dbg("dtq_0.B.GetShape()[2]:{}", dtq_0.B.GetShape()[2]); + assert(dtq_0.B.GetShape()[0] == q1d); + map_quadrature_data_to_fields( y, fhat, output_fop, output_dtq_shmem[0], scratch_shmem, dimension, use_sum_factorization);