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/tensors/jit.h"
#include "neml2/equation_systems/SparseTensorList.h"
 
namespace neml2
{
// Foward declarations
class Model;
class LabeledAxis;
class VariableBase;
template <typename T>
class Variable;
template <typename T>
struct TensorName;

ValueMap bind(const std::vector<VariableName> &, const std::vector<Tensor> &);
 
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);

  LabeledAxis & input_axis();
  const LabeledAxis & input_axis() const;

  LabeledAxis & output_axis();
  const LabeledAxis & output_axis() const;
 
  using VariableStorage = std::map<VariableName, std::unique_ptr<VariableBase>>;
  using DerivSparsity = std::vector<std::pair<VariableBase *, const VariableBase *>>;
  using SecDerivSparsity =
      std::vector<std::tuple<VariableBase *, const VariableBase *, const 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_undefined_input();

  void cache_derivative_sparsity();
  void cache_second_derivative_sparsity();
  const std::optional<DerivSparsity> & derivative_sparsity() const;
  const std::optional<SecDerivSparsity> & second_derivative_sparsity() const;

  void assign_input(const ValueMap & vals);
  void assign_input(const std::vector<VariableName> &, const SparseTensorList &);
  void assign_output(const ValueMap & vals);
  void assign_output(const std::vector<VariableName> &, const SparseTensorList &);
  void assign_output_derivatives(const DerivMap & derivs);
  void assign_output_derivatives(const std::vector<VariableName> &,
                                 const std::vector<VariableName> &,
                                 const SparseTensorList &);

  ValueMap collect_input() const;
  SparseTensorList collect_input(const std::vector<VariableName> &) const;
  ValueMap collect_output() const;
  SparseTensorList collect_output(const std::vector<VariableName> &) const;
  DerivMap collect_output_derivatives() const;
  SparseTensorList collect_output_derivatives(const std::vector<VariableName> &,
                                              const std::vector<VariableName> &) 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, bool allow_duplicate = false);

  template <typename T>
  const Variable<T> & declare_input_variable(const VariableName & name,
                                             bool allow_duplicate = false);

  template <typename T>
  Variable<T> & declare_output_variable(const char * name);

  template <typename T>
  Variable<T> & declare_output_variable(const VariableName & name);

  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,
                                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;

  std::map<VariableName, Size> _input_intmd_dims;

  std::map<VariableName, Size> _output_intmd_dims;

  std::optional<DerivSparsity> _deriv_sparsity = std::nullopt;

  std::optional<SecDerivSparsity> _secderiv_sparsity = std::nullopt;

  std::optional<DerivSparsity> _deriv_sparsity_nl_sys = std::nullopt;

  std::optional<SecDerivSparsity> _secderiv_sparsity_nl_sys = std::nullopt;
};
} // namespace neml2