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 <optional>
 
#include "neml2/misc/types.h"
#include "neml2/tensors/jit.h"
#include "neml2/equation_systems/SparseMatrix.h"
#include "neml2/equation_systems/SparseVector.h"
 
namespace neml2
{
// Foward declarations
class Model;
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;
 
  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
  set_output_derivative_filter(const std::vector<std::pair<VariableName, VariableName>> & derivs);

  const std::optional<std::vector<std::pair<VariableName, VariableName>>> &
  requested_output_derivatives() const
  {
    return _requested_derivs;
  }

  void assign_input(const ValueMap &, bool allow_nonexistent = false);
  void assign_input(const SparseVector &, bool allow_nonexistent = false);
  void assign_output(const SparseVector &);
  void assign_output_derivatives(const SparseMatrix &);

  SparseVector collect_input(const AxisLayout &) const;
  ValueMap collect_output() const;
  SparseVector collect_output(const AxisLayout &) const;
  DerivMap collect_output_derivatives() const;
  SparseMatrix collect_output_derivatives(const AxisLayout &, const AxisLayout &) const;
  SecDerivMap collect_output_second_derivatives() const;
 
protected:
  void set_output_derivative_filter_nl_sys(
      const std::vector<std::pair<VariableName, VariableName>> & derivs);

  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,
                                            std::optional<VariableName> new_name = std::nullopt);

  VariableBase * clone_output_variable(const VariableBase & var,
                                       std::optional<VariableName> new_name = std::nullopt);

  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;

  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;

  std::optional<std::vector<std::pair<VariableName, VariableName>>> _requested_derivs =
      std::nullopt;

  std::optional<std::vector<std::pair<VariableName, VariableName>>> _requested_derivs_nl_sys =
      std::nullopt;
};
} // namespace neml2