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             : #include "NEML2BatchIndexGenerator.h"
      11             : #include "NEML2Utils.h"
      12             : 
      13             : registerMooseObject("MooseApp", NEML2BatchIndexGenerator);
      14             : 
      15             : InputParameters
      16       14310 : NEML2BatchIndexGenerator::validParams()
      17             : {
      18       14310 :   auto params = ElementUserObject::validParams();
      19       14310 :   params.addClassDescription("Generates the element to batch index map for MOOSEToNEML2 gatherers, "
      20             :                              "NEML2ToMOOSE retrievers, and the NEML2 executor");
      21             : 
      22             :   // Since we use the NEML2 model to evaluate the residual AND the Jacobian at the same time, we
      23             :   // want to execute this user object only at execute_on = LINEAR (i.e. during residual evaluation).
      24             :   // The NONLINEAR exec flag below is for computing Jacobian during automatic scaling.
      25       14310 :   ExecFlagEnum execute_options = MooseUtils::getDefaultExecFlagEnum();
      26       57240 :   execute_options = {EXEC_INITIAL, EXEC_LINEAR, EXEC_NONLINEAR};
      27       14310 :   params.set<ExecFlagEnum>("execute_on") = execute_options;
      28             : 
      29       28620 :   return params;
      30       28620 : }
      31             : 
      32          23 : NEML2BatchIndexGenerator::NEML2BatchIndexGenerator(const InputParameters & params)
      33          23 :   : ElementUserObject(params), _outdated(true)
      34             : {
      35          23 : }
      36             : 
      37             : void
      38           8 : NEML2BatchIndexGenerator::meshChanged()
      39             : {
      40           8 :   _outdated = true;
      41           8 : }
      42             : 
      43             : void
      44        3233 : NEML2BatchIndexGenerator::initialize()
      45             : {
      46        3233 :   if (!NEML2Utils::shouldCompute(_fe_problem))
      47        1360 :     return;
      48             : 
      49        1873 :   if (!_outdated)
      50        1641 :     return;
      51             : 
      52         232 :   _elem_to_batch_index.clear();
      53         232 :   _elem_to_batch_index_cache = {libMesh::invalid_uint, 0};
      54         232 :   _batch_index = 0;
      55             : }
      56             : 
      57             : void
      58       25030 : NEML2BatchIndexGenerator::execute()
      59             : {
      60       25030 :   if (!NEML2Utils::shouldCompute(_fe_problem))
      61        9500 :     return;
      62             : 
      63       15530 :   if (!_outdated)
      64       13155 :     return;
      65             : 
      66        2375 :   _elem_to_batch_index[_current_elem->id()] = _batch_index;
      67        2375 :   _batch_index += _qrule->n_points();
      68             : }
      69             : 
      70             : void
      71         352 : NEML2BatchIndexGenerator::threadJoin(const UserObject & uo)
      72             : {
      73         352 :   if (!NEML2Utils::shouldCompute(_fe_problem))
      74         150 :     return;
      75             : 
      76         202 :   if (!_outdated)
      77         201 :     return;
      78             : 
      79           1 :   const auto & m2n = static_cast<const NEML2BatchIndexGenerator &>(uo);
      80             : 
      81             :   // append and renumber maps
      82           6 :   for (const auto & [elem_id, batch_index] : m2n._elem_to_batch_index)
      83           5 :     _elem_to_batch_index[elem_id] = _batch_index + batch_index;
      84             : 
      85           1 :   _batch_index += m2n._batch_index;
      86             : }
      87             : 
      88             : void
      89        2881 : NEML2BatchIndexGenerator::finalize()
      90             : {
      91        2881 :   _outdated = false;
      92        2881 : }
      93             : 
      94             : std::size_t
      95       67640 : NEML2BatchIndexGenerator::getBatchIndex(dof_id_type elem_id) const
      96             : {
      97             :   // return cached map lookup if applicable
      98       67640 :   if (_elem_to_batch_index_cache.first == elem_id)
      99       33347 :     return _elem_to_batch_index_cache.second;
     100             : 
     101             :   // else, search the map
     102       34293 :   const auto it = _elem_to_batch_index.find(elem_id);
     103       34293 :   if (it == _elem_to_batch_index.end())
     104           0 :     mooseError("No batch index found for element id ", elem_id);
     105       34293 :   _elem_to_batch_index_cache = *it;
     106       34293 :   return it->second;
     107             : }