Add old state and derivative access to NEML2 compute

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);
+}

include/tensor_computes/TimeTensorCompute.h

@@ -0,0 +1,24 @@
+/**********************************************************************/
+/*                    DO NOT MODIFY THIS HEADER                       */
+/*             Swift, a Fourier spectral solver for MOOSE             */
+/*                                                                    */
+/*            Copyright 2024 Battelle Energy Alliance, LLC            */
+/*                        ALL RIGHTS RESERVED                         */
+/**********************************************************************/
+
+#pragma once
+
+#include "TensorOperator.h"
+
+/**
+ * Current simulation time
+ */
+class TimeTensorCompute : public TensorOperator<>
+{
+public:
+  static InputParameters validParams();
+
+  TimeTensorCompute(const InputParameters & parameters);
+
+  virtual void computeBuffer() override;
+};

src/tensor_computes/NEML2TensorCompute.C

@@ -8,6 +8,7 @@
 
 #pragma once
 
+#include "Conversion.h"
 #include "NEML2TensorCompute.h"
 #include "NEML2Utils.h"
 
@@ -67,14 +68,57 @@ 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();
+  mooseInfo("Inputs: ", Moose::stringify(model_inputs));
+
+  // 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() || neml2_input.is_old_force())
+    {
+      // 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 +143,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);
 

src/tensor_computes/TimeTensorCompute.C

@@ -0,0 +1,31 @@
+/**********************************************************************/
+/*                    DO NOT MODIFY THIS HEADER                       */
+/*             Swift, a Fourier spectral solver for MOOSE             */
+/*                                                                    */
+/*            Copyright 2024 Battelle Energy Alliance, LLC            */
+/*                        ALL RIGHTS RESERVED                         */
+/**********************************************************************/
+
+#include "SwiftUtils.h"
+#include "TimeTensorCompute.h"
+
+registerMooseObject("SwiftApp", TimeTensorCompute);
+
+InputParameters
+TimeTensorCompute::validParams()
+{
+  InputParameters params = TensorOperator<>::validParams();
+  params.addClassDescription("Scalar tensor with the current simulation time.");
+  return params;
+}
+
+TimeTensorCompute::TimeTensorCompute(const InputParameters & parameters)
+  : TensorOperator<>(parameters)
+{
+}
+
+void
+TimeTensorCompute::computeBuffer()
+{
+  _u = torch::tensor({_time}, MooseTensor::floatTensorOptions()).expand(_domain.getShape());
+}

test/tests/neml2/neml2_input.i

@@ -8,5 +8,12 @@
     from_var = 'forces/A forces/B'
     to_var = 'state/C'
   []
+
+  [rate]
+    type = ScalarVariableRate
+    time = forces/t
+    variable = forces/A
+    rate = state/dAdt
+  []
 []
 

test/tests/neml2/old_state.i

@@ -0,0 +1,54 @@
+[Domain]
+  dim = 2
+  nx = 2
+  ny = 2
+  xmax = 1
+  ymax = 1
+  mesh_mode = DUMMY
+[]
+
+[Problem]
+  type = TensorProblem
+[]
+
+[TensorComputes]
+  [Initialize]
+    [A]
+      type = ConstantTensor
+      buffer = A
+      real = 2
+    []
+  []
+  [Solve]
+    [time]
+      type = TimeTensorCompute
+      buffer = time
+    []
+    [dAdt]
+      type = NEML2TensorCompute
+      neml2_input_file = neml2_input.i
+      neml2_model = rate
+      swift_inputs = 'A time'
+      neml2_inputs = 'forces/A forces/t'
+      neml2_outputs = 'state/dAdt'
+      swift_outputs = 'dAdt'
+    []
+  []
+[]
+
+[Postprocessors]
+  [dAdt]
+    type = TensorAveragePostprocessor
+    buffer = dAdt
+    execute_on = 'TIMESTEP_END'
+  []
+[]
+
+[Executioner]
+  type = Transient
+  num_steps = 5
+[]
+
+[Outputs]
+  csv = true
+[]

test/tests/neml2/tests

@@ -7,4 +7,12 @@
     design = NEML2TensorCompute.md
     requirement = 'The system shall be able to use NEML2 models as tensor computes'
   []
+  [old_state]
+    type = CSVDiff
+    input = old_state.i
+    csvdiff = old_state_out.csv
+    issues = '#50'
+    design = NEML2TensorCompute.md
+    requirement = 'The system shall be able to use old state in NEML2 models'
+  []
 []