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             : #pragma once
      11             : 
      12             : #include "ElementUserObject.h"
      13             : 
      14             : #include <map>
      15             : 
      16             : /**
      17             :  * NEML2BatchIndexGenerator iterates over the mesh and generates a map from element ID to batch
      18             :  * index which is used by NEML2ModelExecutor for transfer data between MOOSE and NEML2.
      19             :  */
      20             : class NEML2BatchIndexGenerator : public ElementUserObject
      21             : {
      22             : public:
      23             :   static InputParameters validParams();
      24             : 
      25             :   NEML2BatchIndexGenerator(const InputParameters & params);
      26             : 
      27             :   void initialize() override;
      28             :   void execute() override;
      29             :   void threadJoin(const UserObject &) override;
      30             :   void finalize() override;
      31             : 
      32             :   void meshChanged() override;
      33             : 
      34             :   /// Get the current batch index (in almost all cases this is the total batch size)
      35        1649 :   std::size_t getBatchIndex() const { return _batch_index; }
      36             : 
      37             :   /// Get the batch index for the given element ID
      38             :   std::size_t getBatchIndex(dof_id_type elem_id) const;
      39             : 
      40             : protected:
      41             :   /// Whether the batch index map is outdated
      42             :   bool _outdated;
      43             : 
      44             :   /// Highest current batch index
      45             :   std::size_t _batch_index;
      46             : 
      47             :   /// Map from element IDs to batch indices
      48             :   std::map<dof_id_type, std::size_t> _elem_to_batch_index;
      49             : 
      50             :   /// cache the index for the current element
      51             :   mutable std::pair<dof_id_type, std::size_t> _elem_to_batch_index_cache;
      52             : };