Add old state (#50)

Author: dschwen

include/tensor_buffers/TensorBuffer.h

@@ -11,7 +11,7 @@
 #include "TensorBufferBase.h"
 
 /**
- * Tensor wrapper arbitrary tensor value dimensions
+ * Tensor wrapper for arbitrary tensor value dimensions
  */
 template <typename T>
 class TensorBuffer : public TensorBufferBase

include/tensor_computes/NEML2TensorCompute.h

@@ -32,7 +32,14 @@ class NEML2TensorCompute : public TensorOperatorBase
 #ifdef NEML2_ENABLED
   neml2::Model & _model;
 
-  std::vector<std::pair<const torch::Tensor *, neml2::LabeledAxisAccessor>> _input_mapping;
+  std::vector<std::tuple<const torch::Tensor *, neml2::TensorType, neml2::LabeledAxisAccessor>>
+      _input_mapping;
+  std::vector<std::tuple<const std::vector<torch::Tensor> *,
+                         const torch::Tensor *,
+                         neml2::TensorType,
+                         neml2::LabeledAxisAccessor>>
+      _old_input_mapping;
+
   std::vector<std::pair<neml2::LabeledAxisAccessor, torch::Tensor *>> _output_mapping;
 #endif
 };

include/tensor_computes/TensorOperatorBase.h

@@ -59,6 +59,13 @@ class TensorOperatorBase : public MooseObject,
   template <typename T = torch::Tensor>
   T & getOutputBufferByName(const TensorOutputBufferName & buffer_name);
 
+  template <typename T = torch::Tensor>
+  const std::vector<T> & getBufferOld(const std::string & param, unsigned int max_states);
+
+  template <typename T = torch::Tensor>
+  const std::vector<T> & getBufferOldByName(const std::string & buffer_name,
+                                            unsigned int max_states);
+
   std::set<std::string> _requested_buffers;
   std::set<std::string> _supplied_buffers;
 
@@ -110,3 +117,17 @@ TensorOperatorBase::getOutputBufferByName(const TensorOutputBufferName & buffer_
   _supplied_buffers.insert(buffer_name);
   return _tensor_problem.getBuffer<T>(buffer_name);
 }
+
+template <typename T>
+const std::vector<T> &
+TensorOperatorBase::getBufferOld(const std::string & param, unsigned int max_states)
+{
+  return getBufferOldByName<T>(getParam<TensorInputBufferName>(param), max_states);
+}
+
+template <typename T>
+const std::vector<T> &
+TensorOperatorBase::getBufferOldByName(const std::string & buffer_name, unsigned int max_states)
+{
+  return _tensor_problem.getBufferOld<T>(buffer_name, max_states);
+}

src/tensor_computes/NEML2TensorCompute.C

@@ -67,14 +67,56 @@ NEML2TensorCompute::NEML2TensorCompute(const InputParameters & params)
   NEML2Utils::assertNEML2Enabled();
 
 #ifdef NEML2_ENABLED
-  for (const auto & [swift_input_name, neml2_input_name] :
-       getParam<TensorInputBufferName, std::string>("swift_inputs", "neml2_inputs"))
+  const auto inputs = getParam<TensorInputBufferName, std::string>("swift_inputs", "neml2_inputs");
+  std::map<neml2::LabeledAxisAccessor, TensorInputBufferName> lookup_swift_name;
+  const auto model_inputs = _model.consumed_items();
+
+  // current inputs
+  for (const auto & [swift_input_name, neml2_input_name] : inputs)
   {
-    const auto * input_buffer = &getInputBufferByName<>(swift_input_name);
-    _input_mapping.emplace_back(
-        input_buffer, neml2::LabeledAxisAccessor(NEML2Utils::parseVariableName(neml2_input_name)));
+    const auto neml2_input =
+        neml2::LabeledAxisAccessor(NEML2Utils::parseVariableName(neml2_input_name));
+
+    // populate reverse lookup map
+    if (lookup_swift_name.find(neml2_input) != lookup_swift_name.end())
+      mooseError("Repeated NEML2 input ", neml2_input_name);
+    lookup_swift_name[neml2_input] = swift_input_name;
+
+    // the user should only specify current neml2 axis
+    if (!neml2_input.is_state() && !neml2_input.is_force())
+      mooseError("Specify only current forces or states as inputs. Old forces and states are "
+                 "automatically coupled when needed.");
+
+    // add input if the model requires it
+    if (model_inputs.count(neml2_input))
+    {
+      const auto * input_buffer = &getInputBufferByName<>(swift_input_name);
+      const auto type = _model.input_variable(neml2_input).type();
+      _input_mapping.emplace_back(input_buffer, type, neml2_input);
+    }
   }
 
+  // old state inputs
+  for (const auto & neml2_input : model_inputs)
+    if (neml2_input.is_old_state())
+    {
+      // check if we couple the current state
+      auto it = lookup_swift_name.find(neml2_input.current());
+      if (it == lookup_swift_name.end())
+        mooseError("The model requires ",
+                   neml2_input,
+                   " but no tensor buffer is assigned to ",
+                   neml2_input.current(),
+                   ".");
+      const auto & swift_input_name = it->second;
+
+      const auto * old_states = &getBufferOldByName<>(swift_input_name, 1);
+      // we also get the current state here just to step zero, when no old state exists!
+      const auto * input_buffer = &getInputBufferByName<>(swift_input_name);
+      const auto type = _model.input_variable(neml2_input).type();
+      _old_input_mapping.emplace_back(old_states, input_buffer, type, neml2_input);
+    }
+
   for (const auto & [neml2_output_name, swift_output_name] :
        getParam<std::string, TensorInputBufferName>("neml2_outputs", "swift_outputs"))
   {
@@ -99,20 +141,30 @@ NEML2TensorCompute::computeBuffer()
 {
 #ifdef NEML2_ENABLED
   neml2::ValueMap in;
-  for (const auto & [tensor_ptr, label] : _input_mapping)
+  auto insert_tensor = [&in, this](const auto & tensor, auto type, const auto & label)
   {
     // convert tensors on the fly at runtime
-    auto sizes = tensor_ptr->sizes();
-    mooseInfoRepeated(name(), " sizes size ", sizes.size(), " is ", Moose::stringify(sizes));
-    if (sizes.size() == _dim)
-      in[label] = neml2::Scalar(*tensor_ptr);
-    else if (sizes.size() == _dim + 1)
-      in[label] = neml2::Vec(*tensor_ptr, _domain.getShape());
-    else if (sizes.size() == _dim + 3)
-      in[label] = neml2::R2(*tensor_ptr, _domain.getShape());
+    auto sizes = tensor.sizes();
+    if (sizes.size() == _dim && type == neml2::TensorType::kScalar)
+      in[label] = neml2::Scalar(tensor);
+    else if (sizes.size() == _dim + 1 && type == neml2::TensorType::kVec)
+      in[label] = neml2::Vec(tensor, _domain.getShape());
+    else if (sizes.size() == _dim + 3 && type == neml2::TensorType::kR2)
+      in[label] = neml2::R2(tensor, _domain.getShape());
     else
-      mooseError("Unsupported tensor dimension");
-  }
+      mooseError("Unsupported/mismatching tensor dimension");
+  };
+
+  // insert current state
+  for (const auto & [current_state, type, label] : _input_mapping)
+    insert_tensor(*current_state, type, label);
+
+  // insert old state
+  for (const auto & [old_states, current_state, type, label] : _old_input_mapping)
+    if (old_states->empty())
+      insert_tensor(*current_state, type, label);
+    else
+      insert_tensor((*old_states)[0], type, label);
 
   auto out = _model.value(in);