Line data    Source code

       1             : //* This file is part of the MOOSE framework
       2             : //* https://mooseframework.inl.gov
       3             : //*
       4             : //* All rights reserved, see COPYRIGHT for full restrictions
       5             : //* https://github.com/idaholab/moose/blob/master/COPYRIGHT
       6             : //*
       7             : //* Licensed under LGPL 2.1, please see LICENSE for details
       8             : //* https://www.gnu.org/licenses/lgpl-2.1.html
       9             : 
      10             : #ifdef NEML2_ENABLED
      11             : 
      12             : // libmesh includes
      13             : #include "libmesh/petsc_vector.h"
      14             : 
      15             : // Torch includes
      16             : #include <ATen/ops/from_blob.h>
      17             : 
      18             : // NEML2 includes
      19             : #include "neml2/tensors/functions/discretization/scatter.h"
      20             : #include "neml2/tensors/functions/discretization/interpolate.h"
      21             : 
      22             : // MOOSE includes
      23             : #include "NEML2FEInterpolation.h"
      24             : 
      25             : using namespace libMesh;
      26             : 
      27             : registerMooseObject("MooseApp", NEML2FEInterpolation);
      28             : 
      29             : InputParameters
      30           2 : NEML2FEInterpolation::validParams()
      31             : {
      32           2 :   InputParameters params = ElementUserObject::validParams();
      33             : 
      34           4 :   params.addClassDescription(
      35             :       "This user object provides an interface to NEML2 for finite element "
      36             :       "interpolation of variables and their gradients. It gathers the shape "
      37             :       "functions and DOF maps for each variable in the assembly and provides "
      38             :       "them as NEML2 tensors for use in NEML2 models.");
      39             : 
      40           6 :   params.addRequiredParam<UserObjectName>(
      41             :       "assembly", "The NEML2Assembly object to use to provide assembly information");
      42             : 
      43           2 :   ExecFlagEnum execute_options = MooseUtils::getDefaultExecFlagEnum();
      44           6 :   execute_options = {EXEC_INITIAL, EXEC_LINEAR};
      45           4 :   params.set<ExecFlagEnum>("execute_on") = execute_options;
      46           2 :   params.suppressParameter<ExecFlagEnum>("execute_on");
      47             : 
      48           4 :   return params;
      49           4 : }
      50             : 
      51           0 : NEML2FEInterpolation::NEML2FEInterpolation(const InputParameters & parameters)
      52           0 :   : ElementUserObject(parameters), _neml2_assembly(getUserObject<NEML2Assembly>("assembly"))
      53             : {
      54           0 : }
      55             : 
      56             : const neml2::Tensor &
      57           0 : NEML2FEInterpolation::getValue(const std::string & var_name)
      58             : {
      59           0 :   auto [it, success] = _vars.emplace(var_name, neml2::Tensor());
      60             : 
      61           0 :   if (success)
      62             :   {
      63           0 :     const auto * var = getMOOSEVariable(var_name);
      64           0 :     _phis.emplace(var->feType(), &var->phi());
      65           0 :     _moose_vars.emplace(var_name, var);
      66             :   }
      67             : 
      68           0 :   return it->second;
      69             : }
      70             : 
      71             : const neml2::Tensor &
      72           0 : NEML2FEInterpolation::getGradient(const std::string & var_name)
      73             : {
      74           0 :   auto [it, success] = _grad_vars.emplace(var_name, neml2::Tensor());
      75             : 
      76           0 :   if (success)
      77             :   {
      78           0 :     const auto * var = getMOOSEVariable(var_name);
      79           0 :     _grad_phis.emplace(var->feType(), &var->gradPhi());
      80           0 :     _moose_vars.emplace(var_name, var);
      81             :   }
      82             : 
      83           0 :   return it->second;
      84             : }
      85             : 
      86             : const neml2::Tensor &
      87           0 : NEML2FEInterpolation::getPhi(const std::string & var_name)
      88             : {
      89           0 :   const auto * var = getMOOSEVariable(var_name);
      90             : 
      91           0 :   const auto it = _neml2_phi.find(var->feType());
      92           0 :   if (it != _neml2_phi.end())
      93           0 :     return it->second;
      94             : 
      95           0 :   _phis.emplace(var->feType(), &var->phi());
      96           0 :   auto [it2, success] = _neml2_phi.emplace(var->feType(), neml2::Tensor());
      97           0 :   return it2->second;
      98             : }
      99             : 
     100             : const neml2::Tensor &
     101           0 : NEML2FEInterpolation::getPhiGradient(const std::string & var_name)
     102             : {
     103           0 :   const auto * var = getMOOSEVariable(var_name);
     104             : 
     105           0 :   const auto it = _neml2_grad_phi.find(var->feType());
     106           0 :   if (it != _neml2_grad_phi.end())
     107           0 :     return it->second;
     108             : 
     109           0 :   _grad_phis.emplace(var->feType(), &var->gradPhi());
     110           0 :   auto [it2, success] = _neml2_grad_phi.emplace(var->feType(), neml2::Tensor());
     111           0 :   return it2->second;
     112             : }
     113             : 
     114             : const neml2::Tensor &
     115           0 : NEML2FEInterpolation::getDofMap(const std::string & var_name)
     116             : {
     117           0 :   return _neml2_dof_map[var_name];
     118             : }
     119             : 
     120             : const std::vector<dof_id_type> &
     121           0 : NEML2FEInterpolation::getGlobalDofMap(const std::string & var_name)
     122             : {
     123           0 :   return _moose_dof_map_global[var_name];
     124             : }
     125             : 
     126             : int64_t
     127           0 : NEML2FEInterpolation::local_ndof() const
     128             : {
     129           0 :   return _local_ndof;
     130             : }
     131             : 
     132             : const MooseVariableFE<Real> *
     133           0 : NEML2FEInterpolation::getMOOSEVariable(const std::string & var_name) const
     134             : {
     135           0 :   const auto * var = &_fe_problem.getVariable(
     136             :       0, var_name, Moose::VarKindType::VAR_SOLVER, Moose::VarFieldType::VAR_FIELD_STANDARD);
     137           0 :   const auto * var_fe = dynamic_cast<const MooseVariableFE<Real> *>(var);
     138             : 
     139           0 :   if (!var_fe)
     140           0 :     mooseError("NEML2FEInterpolation only supports variables of type MooseVariableFE<Real>");
     141             : 
     142           0 :   if (var_fe->scalingFactor() != 1)
     143           0 :     mooseError("Scaling factors other than unity are not yet supported");
     144             : 
     145             :   // check domain restrictions for compatibility
     146           0 :   if (!var_fe->hasBlocks(blockIDs()))
     147           0 :     mooseError("The variable '",
     148           0 :                var_fe->name(),
     149             :                "' must be defined on all blocks '",
     150           0 :                name(),
     151             :                "' is defined on.");
     152             : 
     153           0 :   return var_fe;
     154             : }
     155             : 
     156             : void
     157           0 : NEML2FEInterpolation::initialSetup()
     158             : {
     159           0 :   _petsc_solution = dynamic_cast<const PetscVector<Real> *>(_sys.currentSolution());
     160             : 
     161             :   // check if the solution vector is of a supported type
     162           0 :   if (!_petsc_solution)
     163           0 :     mooseError("Only solution vectors of type PetscVector are currently supported");
     164             : 
     165           0 :   if (_tid != 0)
     166           0 :     syncWithMainThread();
     167           0 : }
     168             : 
     169             : void
     170           0 : NEML2FEInterpolation::invalidateFEMContext()
     171             : {
     172           0 :   _fem_context_up_to_date = false;
     173           0 :   invalidateInterpolations();
     174           0 : }
     175             : 
     176             : void
     177           0 : NEML2FEInterpolation::invalidateInterpolations()
     178             : {
     179           0 :   _interp_up_to_date = false;
     180           0 : }
     181             : 
     182             : void
     183           0 : NEML2FEInterpolation::meshChanged()
     184             : {
     185           0 :   invalidateFEMContext();
     186           0 : }
     187             : 
     188             : void
     189           0 : NEML2FEInterpolation::initialize()
     190             : {
     191           0 :   if (_fem_context_up_to_date)
     192           0 :     return;
     193             : 
     194           0 :   _ndofe.clear();
     195           0 :   _moose_dof_map.clear();
     196           0 :   _moose_dof_map_global.clear();
     197           0 :   _moose_phi.clear();
     198           0 :   _moose_grad_phi.clear();
     199           0 :   _local_ndof = 0;
     200             : }
     201             : 
     202             : void
     203           0 : NEML2FEInterpolation::syncWithMainThread()
     204             : {
     205           0 :   auto & main_uo = _fe_problem.getUserObject<NEML2FEInterpolation>(name(), /*tid=*/0);
     206           0 :   for (const auto & [var_name, var] : main_uo._moose_vars)
     207             :   {
     208           0 :     _moose_vars.emplace(var_name, getMOOSEVariable(var_name));
     209           0 :     if (main_uo._phis.count(var->feType()))
     210           0 :       getPhi(var_name);
     211           0 :     if (main_uo._grad_phis.count(var->feType()))
     212           0 :       getPhiGradient(var_name);
     213             :   }
     214           0 : }
     215             : 
     216             : void
     217           0 : NEML2FEInterpolation::threadJoin(const UserObject & y)
     218             : {
     219           0 :   const auto & other = static_cast<const NEML2FEInterpolation &>(y);
     220             :   mooseAssert(_fem_context_up_to_date == other._fem_context_up_to_date,
     221             :               "NEML2FEInterpolation becomes out of sync with other thread");
     222             : 
     223           0 :   if (_fem_context_up_to_date)
     224           0 :     return;
     225             : 
     226           0 :   auto merge_map_vecs = [](auto & map1, const auto & map2)
     227             :   {
     228           0 :     for (const auto & [key, map2_val] : map2)
     229             :     {
     230           0 :       auto & map1_val = map1[key];
     231           0 :       map1_val.insert(map1_val.end(), map2_val.begin(), map2_val.end());
     232             :     }
     233           0 :   };
     234             : 
     235           0 :   merge_map_vecs(_moose_dof_map, other._moose_dof_map);
     236           0 :   merge_map_vecs(_moose_phi, other._moose_phi);
     237           0 :   merge_map_vecs(_moose_grad_phi, other._moose_grad_phi);
     238             : }
     239             : 
     240             : void
     241           0 : NEML2FEInterpolation::execute()
     242             : {
     243           0 :   if (_fem_context_up_to_date)
     244           0 :     return;
     245             : 
     246             :   // DOF indices
     247           0 :   const auto & nl_dof_map = _sys.dofMap();
     248           0 :   for (const auto & [var_name, var] : _moose_vars)
     249             :   {
     250           0 :     nl_dof_map.dof_indices(_current_elem, _dof_indices, var->number());
     251           0 :     auto [it, success] = _ndofe.emplace(var->feType(), _dof_indices.size());
     252           0 :     if (!success && std::size_t(it->second) != _dof_indices.size())
     253           0 :       mooseError("DOF map size mismatch for variable ",
     254             :                  var_name,
     255             :                  ", got ",
     256           0 :                  it->second,
     257             :                  " and ",
     258           0 :                  _dof_indices.size());
     259           0 :     auto & moose_dof_map = _moose_dof_map[var_name];
     260           0 :     auto & moose_dof_map_global = _moose_dof_map_global[var_name];
     261           0 :     for (auto dof : _dof_indices)
     262             :     {
     263           0 :       moose_dof_map.push_back(_petsc_solution->map_global_to_local_index(dof));
     264           0 :       moose_dof_map_global.push_back(dof);
     265             :     }
     266             :   }
     267             : 
     268             :   // shape function values
     269           0 :   for (const auto & [fetype, phi] : _phis)
     270             :   {
     271           0 :     auto & moose_phi = _moose_phi[fetype];
     272           0 :     for (auto i : index_range(*phi))
     273           0 :       for (auto qp : index_range(_q_point))
     274           0 :         moose_phi.push_back((*phi)[i][qp]);
     275             :   }
     276             : 
     277             :   // shape function gradients
     278           0 :   for (const auto & [fetype, grad_phi] : _grad_phis)
     279             :   {
     280           0 :     auto & moose_grad_phi = _moose_grad_phi[fetype];
     281           0 :     for (auto i : index_range(*grad_phi))
     282           0 :       for (auto qp : index_range(_q_point))
     283           0 :         for (auto j : make_range(3))
     284           0 :           moose_grad_phi.push_back((*grad_phi)[i][qp](j));
     285             :   }
     286             : }
     287             : 
     288             : void
     289           0 : NEML2FEInterpolation::finalize()
     290             : {
     291           0 :   TIME_SECTION("finalize", 1, "Updating FEM context and interpolations for NEML2");
     292             : 
     293           0 :   if (!_fem_context_up_to_date)
     294           0 :     updateFEMContext();
     295             : 
     296           0 :   if (!_interp_up_to_date)
     297           0 :     updateInterpolations();
     298           0 : }
     299             : 
     300             : void
     301           0 : NEML2FEInterpolation::updateFEMContext()
     302             : {
     303           0 :   TIME_SECTION("updateFEMContext", 2, "Updating FEM context for NEML2");
     304             : 
     305           0 :   updateDofMap();
     306           0 :   updatePhi();
     307           0 :   updateGradPhi();
     308             : 
     309             :   // done
     310           0 :   _fem_context_up_to_date = true;
     311           0 : }
     312             : 
     313             : void
     314           0 : NEML2FEInterpolation::updateDofMap()
     315             : {
     316           0 :   auto device = _app.getLibtorchDevice();
     317           0 :   auto nelem = _neml2_assembly.numElem();
     318             : 
     319           0 :   for (auto & [var_name, moose_dof_map] : _moose_dof_map)
     320             :   {
     321           0 :     auto var = _moose_vars.at(var_name);
     322           0 :     auto ndofe = _ndofe.at(var->feType());
     323             : 
     324             :     // sanity check on sizes
     325           0 :     if (moose_dof_map.size() != std::size_t(nelem * ndofe))
     326           0 :       mooseError(
     327           0 :           "dof map size mismatch, expected ", nelem * ndofe, " but got ", moose_dof_map.size());
     328             : 
     329             :     // convert to neml2 tensor
     330           0 :     _neml2_dof_map[var_name] =
     331           0 :         neml2::Tensor(at::from_blob(moose_dof_map.data(), {nelem, ndofe}, torch::kInt64), 2)
     332           0 :             .to(device);
     333             : 
     334           0 :     _local_ndof = std::max(_local_ndof, _neml2_dof_map[var_name].max().item<int64_t>() + 1);
     335             :   }
     336           0 : }
     337             : 
     338             : void
     339           0 : NEML2FEInterpolation::updatePhi()
     340             : {
     341           0 :   auto device = _app.getLibtorchDevice();
     342           0 :   auto nelem = _neml2_assembly.numElem();
     343           0 :   auto nqp = _neml2_assembly.numQP();
     344             : 
     345           0 :   for (auto & [fetype, moose_phi] : _moose_phi)
     346             :   {
     347           0 :     auto ndofe = _ndofe.at(fetype);
     348             : 
     349             :     // sanity check on sizes
     350           0 :     if (moose_phi.size() != std::size_t(nelem * ndofe * nqp))
     351           0 :       mooseError("shape function size mismatch, expected ",
     352           0 :                  nelem * ndofe * nqp,
     353             :                  " but got ",
     354           0 :                  moose_phi.size());
     355           0 :     _neml2_phi[fetype] =
     356           0 :         neml2::Tensor(at::from_blob(moose_phi.data(), {nelem, ndofe, nqp}, torch::kFloat64), 3)
     357           0 :             .to(device);
     358             :   }
     359           0 : }
     360             : 
     361             : void
     362           0 : NEML2FEInterpolation::updateGradPhi()
     363             : {
     364           0 :   auto device = _app.getLibtorchDevice();
     365           0 :   auto nelem = _neml2_assembly.numElem();
     366           0 :   auto nqp = _neml2_assembly.numQP();
     367             : 
     368           0 :   for (auto & [fetype, moose_grad_phi] : _moose_grad_phi)
     369             :   {
     370           0 :     auto ndofe = _ndofe.at(fetype);
     371             : 
     372             :     // sanity check on sizes
     373           0 :     if (moose_grad_phi.size() != std::size_t(nelem * ndofe * nqp * 3))
     374           0 :       mooseError("shape function gradient size mismatch, expected ",
     375           0 :                  nelem * ndofe * nqp * 3,
     376             :                  " but got ",
     377           0 :                  moose_grad_phi.size());
     378           0 :     _neml2_grad_phi[fetype] =
     379           0 :         neml2::Tensor(at::from_blob(moose_grad_phi.data(), {nelem, ndofe, nqp, 3}, torch::kFloat64),
     380             :                       3)
     381           0 :             .to(device);
     382             :   }
     383           0 : }
     384             : 
     385             : void
     386           0 : NEML2FEInterpolation::updateInterpolations()
     387             : {
     388           0 :   TIME_SECTION("updateInterpolations", 2, "Updating FEM interpolations for NEML2");
     389             : 
     390             :   // convert the local solution vector to neml2 tensor
     391           0 :   auto sol = at::from_blob(const_cast<Real *>(_petsc_solution->get_array_read()),
     392           0 :                            {local_ndof()},
     393             :                            torch::kFloat64)
     394           0 :                  .to(_app.getLibtorchDevice());
     395             : 
     396             :   // interpolate variable values
     397           0 :   for (auto & [var_name, val] : _vars)
     398             :   {
     399           0 :     const auto & dof_map = _neml2_dof_map[var_name];
     400           0 :     const auto fetype = _moose_vars[var_name]->feType();
     401           0 :     const auto & phi = _neml2_phi[fetype];
     402           0 :     auto sol_scattered = neml2::discretization::scatter(sol, dof_map);
     403           0 :     val = neml2::discretization::interpolate(sol_scattered, phi);
     404           0 :   }
     405             : 
     406             :   // interpolate variable gradients
     407           0 :   for (auto & [var_name, val] : _grad_vars)
     408             :   {
     409           0 :     const auto & dof_map = _neml2_dof_map[var_name];
     410           0 :     const auto fetype = _moose_vars[var_name]->feType();
     411           0 :     const auto & grad_phi = _neml2_grad_phi[fetype];
     412           0 :     auto sol_scattered = neml2::discretization::scatter(sol, dof_map);
     413           0 :     val = neml2::discretization::interpolate(sol_scattered, grad_phi);
     414           0 :   }
     415             : 
     416             :   // close solution and residual vector access
     417           0 :   const_cast<PetscVector<Real> *>(_petsc_solution)->restore_array();
     418             : 
     419             :   // done
     420           0 :   _interp_up_to_date = true;
     421           0 : }
     422             : 
     423             : #endif