VariableStore.h

// Copyright 2024, UChicago Argonne, LLC
// All Rights Reserved
// Software Name: NEML2 -- the New Engineering material Model Library, version 2
// By: Argonne National Laboratory
// OPEN SOURCE LICENSE (MIT)
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.
 
#pragma once
 
#include <memory>
 
#include "neml2/models/map_types_fwd.h"
#include "neml2/base/LabeledAxisAccessor.h"
#include "neml2/misc/types.h"
#include "neml2/jit/types.h"
 
namespace neml2
{
// Foward declarations
class Model;
class LabeledAxis;
class VariableBase;
template <typename T>
class Variable;
template <typename T>
struct TensorName;
 
class VariableStore
{
public:
  VariableStore(Model * object);
 
  VariableStore(const VariableStore &) = delete;
  VariableStore(VariableStore &&) = delete;
  VariableStore & operator=(const VariableStore &) = delete;
  VariableStore & operator=(VariableStore &&) = delete;
  virtual ~VariableStore() = default;
 
  LabeledAxis & declare_axis(const std::string & name);
 
  virtual void setup_layout();
 
  LabeledAxis & input_axis() { return _input_axis; }
  const LabeledAxis & input_axis() const { return _input_axis; }
 
  LabeledAxis & output_axis() { return _output_axis; }
  const LabeledAxis & output_axis() const { return _output_axis; }
 
  using VariableStorage = std::map<VariableName, std::unique_ptr<VariableBase>>;
 
  VariableStorage & input_variables() { return _input_variables; }
  const VariableStorage & input_variables() const { return _input_variables; }
  VariableStorage & output_variables() { return _output_variables; }
  const VariableStorage & output_variables() const { return _output_variables; }
 
  VariableBase & input_variable(const VariableName &);
  const VariableBase & input_variable(const VariableName &) const;
  VariableBase & output_variable(const VariableName &);
  const VariableBase & output_variable(const VariableName &) const;
 
  const TensorOptions & variable_options() const { return _options; }
 
  virtual void clear_input();
  virtual void clear_output();
  virtual void clear_derivatives();
 
  virtual void zero_input();
  virtual void zero_output();
 
  void assign_input(const ValueMap & vals);
  void assign_input(ValueMap && vals);
  void assign_output(const ValueMap & vals);
  void assign_output_derivatives(const DerivMap & derivs);
 
  ValueMap collect_input() const;
  ValueMap collect_output() const;
  DerivMap collect_output_derivatives() const;
  SecDerivMap collect_output_second_derivatives() const;
 
protected:
  virtual void send_variables_to(const TensorOptions & options);
 
  template <typename T>
  const Variable<T> & declare_input_variable(const char * name,
                                             TensorShapeRef list_shape = {},
                                             bool allow_duplicate = false);
 
  template <typename T>
  const Variable<T> & declare_input_variable(const VariableName & name,
                                             TensorShapeRef list_shape = {},
                                             bool allow_duplicate = false);
 
  template <typename T>
  Variable<T> & declare_output_variable(const char * name, TensorShapeRef list_shape = {});
 
  template <typename T>
  Variable<T> & declare_output_variable(const VariableName & name, TensorShapeRef list_shape = {});
 
  const VariableBase * clone_input_variable(const VariableBase & var,
                                            const VariableName & new_name = {});
 
  VariableBase * clone_output_variable(const VariableBase & var,
                                       const VariableName & new_name = {});
 
  void assign_input_stack(jit::Stack & stack);
 
  void assign_output_stack(jit::Stack & stack, bool out, bool dout, bool d2out);
 
  jit::Stack collect_input_stack() const;
 
  jit::Stack collect_output_stack(bool out, bool dout, bool d2out) const;
 
  // TensorName resolution may require declare_input_variable
  template <typename T>
  friend const T & resolve_tensor_name(const TensorName<T> &, Model *, const std::string &);
 
private:
  // Create a variable
  template <typename T>
  Variable<T> * create_variable(VariableStorage & variables,
                                const VariableName & name,
                                TensorShapeRef list_shape,
                                bool allow_duplicate = false);
 
  Model * _object;
 
  std::map<std::string, std::unique_ptr<LabeledAxis>> _axes;
 
  LabeledAxis & _input_axis;
 
  LabeledAxis & _output_axis;
 
  VariableStorage _input_variables;
 
  VariableStorage _output_variables;
 
  TensorOptions _options;
};
} // namespace neml2