NEML2 submodule update

- Update neml2 submodule
- Remove dependency on LabeledTensor
- Update all neml2 models
- Update all moose-neml2 input files
- Attempted to speedup data transfer between moose and neml2

ref idaholab#29579

modules/combined/test/tests/optimization/invOpt_elasticity_modular/grad.i

@@ -69,7 +69,6 @@
   verbose = true
   device = 'cpu'
   [forward]
-    enable_AD = true
     model = 'forward_elasticity_model'
 
     moose_input_types = 'MATERIAL'

modules/solid_mechanics/include/neml2/actions/NEML2Action.h

@@ -15,6 +15,8 @@
 
 #include "Action.h"
 
+class NEML2ActionCommon;
+
 /**
  * Action to set up NEML2 objects.
  */
@@ -28,6 +30,8 @@ class NEML2Action : public Action
   virtual void act() override;
 
 protected:
+  const NEML2ActionCommon & getCommonAction() const;
+
 #ifdef NEML2_ENABLED
 
   enum class MOOSEIOType
@@ -132,4 +136,8 @@ class NEML2Action : public Action
 
   /// Material property additional outputs
   std::map<MaterialPropertyName, std::vector<OutputName>> _export_output_targets;
+
+private:
+  /// Get the maximum length of all MOOSE names (for printing purposes)
+  std::size_t getMaximumMOOSEName() const;
 };

modules/solid_mechanics/include/neml2/actions/NEML2ActionCommon.h

@@ -47,13 +47,19 @@ class NEML2ActionCommon : public Action
   /// Parameters that can be specified EITHER under the common area OR under sub-blocks
   static InputParameters commonParams();
 
+  /// Parameters that can ONLY be specified under the common area
   static InputParameters validParams();
 
   NEML2ActionCommon(const InputParameters &);
 
   virtual void act() override;
 
+  const FileName & fname() const { return _fname; }
+
 protected:
   /// Name of the NEML2 input file
   const FileName _fname;
+
+  /// List of cli-args
+  const std::vector<std::string> _cli_args;
 };

modules/solid_mechanics/include/neml2/interfaces/MOOSEToNEML2.h

@@ -13,7 +13,6 @@
 #include "InputParameters.h"
 
 #ifdef NEML2_ENABLED
-#include "neml2/tensors/LabeledVector.h"
 #include "neml2/models/Model.h"
 #endif
 
@@ -70,7 +69,7 @@ class MOOSEToNEML2
   virtual neml2::Tensor gatheredData() const = 0;
 
   /// Insert the gathered data into the NEML2 material model
-  void insertInto(neml2::Model &) const;
+  void insertInto(neml2::ValueMap &, std::map<std::string, neml2::Tensor> &) const;
 
 protected:
   /// Whether we should insert into NEML2 input variable or NEML2 model parameter

modules/solid_mechanics/include/neml2/interfaces/NEML2ModelInterface.h

@@ -9,15 +9,15 @@
 
 #pragma once
 
+#include "NEML2Utils.h"
+
 #ifdef NEML2_ENABLED
 #include "neml2/models/Model.h"
 #include "neml2/misc/parser_utils.h"
 #include "Material.h"
 #include "UserObject.h"
 #endif
 
-#include "NEML2Utils.h"
-
 /**
  * Interface class to provide common input parameters, members, and methods for MOOSEObjects that
  * use NEML2 models.
@@ -35,29 +35,17 @@ class NEML2ModelInterface : public T
 
 protected:
   /**
-   * Validate the NEML2 material model. This method should throw a moose error with the first
-   * encountered problem. Note that the developer is responsible for calling this method at the
-   * appropriate times, for example, at initialSetup().
+   * Validate the NEML2 material model. Note that the developer is responsible for calling this
+   * method at the appropriate times, for example, at initialSetup().
    */
   virtual void validateModel() const;
 
-  /// Initialize the model with a batch shape
-  void initModel(neml2::TensorShapeRef batch_shape);
-
   /// Get the NEML2 model
   neml2::Model & model() const { return _model; }
 
   /// Get the target compute device
   const torch::Device & device() const { return _device; }
 
-  /**
-   * @brief Convert a raw string to a neml2::VariableName
-   *
-   * @param raw_str
-   * @return neml2::VariableName
-   */
-  neml2::VariableName getNEML2VariableName(const std::string & raw_str) const;
-
 private:
   /// The NEML2 material model
   neml2::Model & _model;
@@ -86,11 +74,6 @@ NEML2ModelInterface<T>::validParams()
       ":<device-index> optionally specifies a device index. For example, device='cpu' sets the "
       "target compute device to be CPU, and device='cuda:1' sets the target compute device to be "
       "CUDA with device ID 1.");
-  params.addParam<bool>("enable_AD",
-                        false,
-                        "Set to true to enable PyTorch AD. When set to false (default), no "
-                        "function graph or gradient is computed, which speeds up model "
-                        "evaluation.");
   return params;
 }
 
@@ -109,7 +92,7 @@ template <class T>
 template <typename... P>
 NEML2ModelInterface<T>::NEML2ModelInterface(const InputParameters & params, P &&... args)
   : T(params, args...),
-    _model(neml2::get_model(params.get<std::string>("model"), params.get<bool>("enable_AD"))),
+    _model(neml2::get_model(params.get<std::string>("model"))),
     _device(params.get<std::string>("device"))
 {
 }
@@ -121,18 +104,4 @@ NEML2ModelInterface<T>::validateModel() const
   neml2::diagnose(_model);
 }
 
-template <class T>
-void
-NEML2ModelInterface<T>::initModel(neml2::TensorShapeRef batch_shape)
-{
-  _model.reinit(batch_shape, /*deriv_order=*/1, _device, /*dtype=*/torch::kFloat64);
-}
-
-template <class T>
-neml2::VariableName
-NEML2ModelInterface<T>::getNEML2VariableName(const std::string & raw_str) const
-{
-  return neml2::utils::parse<neml2::VariableName>(raw_str);
-}
-
 #endif // NEML2_ENABLED

modules/solid_mechanics/include/neml2/materials/NEML2ToMOOSEMaterialProperty.h

@@ -43,8 +43,8 @@ class NEML2ToMOOSEMaterialProperty : public Material
   /// Initial condition
   const MaterialProperty<T> * _prop0;
 
-  /// labled view to the requested output (or its derivative)
-  const neml2::Tensor & _output_view;
+  /// Referece to the requested output (or its derivative) value
+  const neml2::Tensor & _value;
 #endif
 };
 
@@ -54,4 +54,3 @@ class NEML2ToMOOSEMaterialProperty : public Material
 DefineNEML2ToMOOSEMaterialPropertyAlias(Real, Real);
 DefineNEML2ToMOOSEMaterialPropertyAlias(SymmetricRankTwoTensor, SymmetricRankTwoTensor);
 DefineNEML2ToMOOSEMaterialPropertyAlias(SymmetricRankFourTensor, SymmetricRankFourTensor);
-DefineNEML2ToMOOSEMaterialPropertyAlias(std::vector<Real>, StdVector);

modules/solid_mechanics/include/neml2/userobjects/MOOSEMaterialPropertyToNEML2.h

@@ -18,7 +18,7 @@
  * Gather a MOOSE material property for insertion into the NEML2 model.
  */
 template <typename T, unsigned int state>
-class MOOSEMaterialPropertyToNEML2 : public MOOSEToNEML2Batched
+class MOOSEMaterialPropertyToNEML2 : public MOOSEToNEML2Batched<T>
 {
 public:
   static InputParameters validParams();
@@ -27,7 +27,7 @@ class MOOSEMaterialPropertyToNEML2 : public MOOSEToNEML2Batched
 
 #ifdef NEML2_ENABLED
 protected:
-  torch::Tensor convertQpMOOSEData() const override;
+  const MooseArray<T> & elemMOOSEData() const override { return _mat_prop.get(); }
 
   /// MOOSE material property to read data from
   const MaterialProperty<T> & _mat_prop;

modules/solid_mechanics/include/neml2/userobjects/MOOSEToNEML2Batched.h

@@ -12,6 +12,20 @@
 #include "MOOSEToNEML2.h"
 #include "ElementUserObject.h"
 
+/**
+ * @brief Generic gatherer for collecting "batched" MOOSE data for NEML2
+ *
+ * It is generic in the sense that it can be used for most MOOSE data types that take the form of
+ * MooseArray<T>.
+ *
+ * It is not so generic in the sense that the collected data is always a std::vector of
+ * MooseArray<T>, where the vector size is generally the number of elements this ElementUserObject
+ * operates on, and the MooseArray<T> size is generally the number of quadrature points in each
+ * element.
+ *
+ * @tparam T Type of the underlying MOOSE data, e.g., Real, SymmetricRankTwoTensor, etc.
+ */
+template <typename T>
 class MOOSEToNEML2Batched : public MOOSEToNEML2, public ElementUserObject
 {
 public:
@@ -32,17 +46,82 @@ class MOOSEToNEML2Batched : public MOOSEToNEML2, public ElementUserObject
 
   neml2::Tensor gatheredData() const override;
 
-  // The number of gathered data items (for setting the model batch size)
-  std::size_t size() const { return _buffer.size(); }
+  // The number of elements
+  std::size_t nElem() const { return _buffer.size(); }
+
+  /// The number of quadrature points in each element
+  // Note that this assumes that all elements have the same number of quadrature points
+  std::size_t nQP() const { return _nqp; }
 
 protected:
-  /// Convert the underlying MOOSE data to a torch::Tensor
-  virtual torch::Tensor convertQpMOOSEData() const = 0;
+  /// MOOSE data for the current element
+  virtual const MooseArray<T> & elemMOOSEData() const = 0;
 
   /// Intermediate data buffer, filled during the element loop
-  std::vector<torch::Tensor> _buffer;
+  std::vector<MooseArray<T>> _buffer;
 
-  /// Current element's quadrature point indexing
-  unsigned int _qp;
+  /// Number of quadrature points
+  std::size_t _nqp;
 #endif
 };
+
+template <typename T>
+InputParameters
+MOOSEToNEML2Batched<T>::validParams()
+{
+  auto params = MOOSEToNEML2::validParams();
+  params += ElementUserObject::validParams();
+
+  // Since we use the NEML2 model to evaluate the residual AND the Jacobian at the same time, we
+  // want to execute this user object only at execute_on = LINEAR (i.e. during residual evaluation).
+  // The NONLINEAR exec flag below is for computing Jacobian during automatic scaling.
+  ExecFlagEnum execute_options = MooseUtils::getDefaultExecFlagEnum();
+  execute_options = {EXEC_INITIAL, EXEC_LINEAR, EXEC_NONLINEAR};
+  params.set<ExecFlagEnum>("execute_on") = execute_options;
+
+  return params;
+}
+
+template <typename T>
+MOOSEToNEML2Batched<T>::MOOSEToNEML2Batched(const InputParameters & params)
+  : MOOSEToNEML2(params), ElementUserObject(params)
+{
+}
+
+#ifdef NEML2_ENABLED
+template <typename T>
+void
+MOOSEToNEML2Batched<T>::initialize()
+{
+  _buffer.clear();
+}
+
+template <typename T>
+void
+MOOSEToNEML2Batched<T>::execute()
+{
+  if (_buffer.empty())
+    _nqp = elemMOOSEData().size();
+  else
+    mooseAssert(_nqp == elemMOOSEData().size(),
+                "Number of quadrature points must be the same for all elements");
+  _buffer.push_back(this->elemMOOSEData());
+}
+
+template <typename T>
+void
+MOOSEToNEML2Batched<T>::threadJoin(const UserObject & uo)
+{
+  // append vectors
+  const auto & m2n = static_cast<const MOOSEToNEML2Batched<T> &>(uo);
+  _buffer.insert(_buffer.end(), m2n._buffer.begin(), m2n._buffer.end());
+  mooseAssert(_nqp == m2n._nqp, "Number of quadrature points must be the same for all elements");
+}
+
+template <typename T>
+neml2::Tensor
+MOOSEToNEML2Batched<T>::gatheredData() const
+{
+  return NEML2Utils::from_blob(_buffer);
+}
+#endif

modules/solid_mechanics/include/neml2/userobjects/MOOSEVariableToNEML2.h

@@ -15,7 +15,7 @@
  * Gather a MOOSE variable for insertion into the NEML2 model.
  */
 template <unsigned int state>
-class MOOSEVariableToNEML2Templ : public MOOSEToNEML2Batched
+class MOOSEVariableToNEML2Templ : public MOOSEToNEML2Batched<Real>
 {
 public:
   static InputParameters validParams();
@@ -24,7 +24,7 @@ class MOOSEVariableToNEML2Templ : public MOOSEToNEML2Batched
 
 #ifdef NEML2_ENABLED
 protected:
-  torch::Tensor convertQpMOOSEData() const override;
+  const MooseArray<Real> & elemMOOSEData() const override { return _moose_variable; }
 
   /// Coupled MOOSE variable to read data from
   const VariableValue & _moose_variable;

modules/solid_mechanics/include/neml2/userobjects/NEML2BatchIndexGenerator.h

@@ -37,10 +37,16 @@ class NEML2BatchIndexGenerator : public ElementUserObject
   /// Get the batch index for the given element ID
   std::size_t getBatchIndex(dof_id_type elem_id) const;
 
+  /// Number of quadrature points in each element
+  std::size_t nQp() const { return _n_qp; }
+
 protected:
   /// Whether the batch index map is outdated
   bool _outdated;
 
+  /// Number of elements in each element
+  std::size_t _n_qp;
+
   /// Highest current batch index
   std::size_t _batch_index;
 

modules/solid_mechanics/include/neml2/userobjects/NEML2ModelExecutor.h

@@ -13,12 +13,6 @@
 #include "GeneralUserObject.h"
 #include "NEML2BatchIndexGenerator.h"
 
-#ifdef NEML2_ENABLED
-#include "neml2/tensors/LabeledVector.h"
-#include "neml2/tensors/LabeledMatrix.h"
-#include "neml2/models/Model.h"
-#endif
-
 class MOOSEToNEML2;
 
 /**
@@ -92,14 +86,17 @@ class NEML2ModelExecutor : public NEML2ModelInterface<GeneralUserObject>
   /// flag that indicates if output data has been fully computed
   bool _output_ready;
 
-  /// The input vector of the material model
-  const neml2::LabeledVector & _in;
+  /// The model parameters to update (gathered from MOOSE)
+  std::map<std::string, neml2::Tensor> _model_params;
+
+  /// The input variables of the material model
+  neml2::ValueMap _in;
 
-  /// The output vector of the material model
-  const neml2::LabeledVector & _out;
+  /// The output variables of the material model
+  neml2::ValueMap _out;
 
-  /// The derivative of the output vector w.r.t. the input vector
-  const neml2::LabeledMatrix & _dout_din;
+  /// The derivative of the output variables w.r.t. the input variables
+  neml2::DerivMap _dout_din;
 
   // set of variables to skip
   std::set<neml2::VariableName> _skip_vars;
@@ -114,14 +111,13 @@ class NEML2ModelExecutor : public NEML2ModelInterface<GeneralUserObject>
   std::vector<const MOOSEToNEML2 *> _gatherers;
 
   /// set of output variables that were retrieved (by other objects)
-  mutable std::map<neml2::VariableName, neml2::Tensor> _retrieved_outputs;
+  mutable neml2::ValueMap _retrieved_outputs;
 
   /// set of derivatives that were retrieved (by other objects)
-  mutable std::map<std::pair<neml2::VariableName, neml2::VariableName>, neml2::Tensor>
-      _retrieved_derivatives;
+  mutable neml2::DerivMap _retrieved_derivatives;
 
   /// set of parameter derivatives that were retrieved (by other objects)
-  mutable std::map<std::pair<neml2::VariableName, std::string>, neml2::Tensor>
+  mutable std::map<neml2::VariableName, std::map<std::string, neml2::Tensor>>
       _retrieved_parameter_derivatives;
 #endif
 };