Model.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 "neml2/base/DependencyDefinition.h"
#include "neml2/base/DiagnosticsInterface.h"
 
#include "neml2/models/Data.h"
#include "neml2/models/ParameterStore.h"
#include "neml2/models/VariableStore.h"
#include "neml2/solvers/NonlinearSystem.h"
 
#include "neml2/tensors/LabeledVector.h"
#include "neml2/tensors/LabeledMatrix.h"
#include "neml2/tensors/LabeledTensor3D.h"
 
namespace neml2
{
class Model : public std::enable_shared_from_this<Model>,
              public Data,
              public ParameterStore,
              public VariableStore,
              public NonlinearSystem,
              public DependencyDefinition<VariableName>,
              public DiagnosticsInterface
{
public:
  static OptionSet expected_options();
 
  Model(const OptionSet & options);
 
  virtual void diagnose(std::vector<Diagnosis> &) const override;
 
  void diagnose_nl_sys(std::vector<Diagnosis> & diagnoses) const;
 
  virtual bool is_nonlinear_system() const { return _nonlinear_system; }
 
  std::tuple<const Tensor &, const Tensor &, const Tensor &, const Tensor &, const Tensor &>
  get_system_matrices() const;
 
  virtual void reinit(const Tensor & tensor, int deriv_order);
 
  virtual void reinit(TensorShapeRef batch_shape = {},
                      int deriv_order = 0,
                      const torch::Device & device = default_device(),
                      const torch::Dtype & dtype = default_dtype());
 
  Size batch_dim() const { return _batch_sizes.size(); }
  TensorShapeRef batch_sizes() const { return _batch_sizes; }
  const torch::TensorOptions & options() const { return _options; }
  torch::Dtype scalar_type() const { return _options.dtype().toScalarType(); }
  torch::Device device() const { return _options.device(); }
 
  const std::vector<Model *> & registered_models() const { return _registered_models; }
  Model * registered_model(const std::string & name) const;
 
  virtual std::set<VariableName> consumed_items() const override;
  virtual std::set<VariableName> provided_items() const override;
 
  virtual bool is_AD_enabled() const override { return _enable_AD; }
  bool using_AD_1st_derivative() const { return _AD_1st_deriv; }
  bool using_AD_2nd_derivative() const { return _AD_2nd_deriv; }
  void use_AD_derivatives(bool first = true, bool second = true);
 
  void prepare();
  virtual void value();
  virtual void value_and_dvalue();
  virtual void dvalue();
  virtual void value_and_dvalue_and_d2value();
  virtual void d2value();
  virtual void dvalue_and_d2value();
  virtual LabeledVector value(const LabeledVector & in);
  virtual std::tuple<LabeledVector, LabeledMatrix> value_and_dvalue(const LabeledVector & in);
  virtual LabeledMatrix dvalue(const LabeledVector & in);
  virtual std::tuple<LabeledVector, LabeledMatrix, LabeledTensor3D>
  value_and_dvalue_and_d2value(const LabeledVector & in);
  virtual LabeledTensor3D d2value(const LabeledVector & in);
  virtual std::tuple<LabeledMatrix, LabeledTensor3D> dvalue_and_d2value(const LabeledVector & in);
 
  friend class ParameterStore;
 
  friend class ComposedModel;
 
protected:
  virtual void setup() override;
 
  virtual void check_AD_limitation() const;
  void check_inplace_dbg();
  virtual void check_input(const LabeledVector & in) const;
 
 
  virtual void reinit(bool in, bool out);
  using VariableStore::cache;
  virtual void cache(TensorShapeRef batch_shape,
                     int deriv_order,
                     const torch::Device & device,
                     const torch::Dtype & dtype);
  bool requires_grad() const { return _deriv_order >= 1; }
  bool requires_2nd_grad() const { return _deriv_order >= 2; }
  using VariableStore::allocate_variables;
  virtual void allocate_variables(bool in, bool out);
  virtual void setup_input_views(VariableStore * host = nullptr) override;
  virtual void setup_submodel_input_views(VariableStore * host);
  using VariableStore::setup_output_views;
  virtual void setup_output_views();
  virtual void setup_submodel_output_views();
  virtual void setup_nonlinear_system();
 
  using VariableStore::zero;
  virtual void zero();
  virtual void set_input(const LabeledVector & in);
  virtual void set_solution(const Tensor & x) override;
  virtual void set_value(bool out, bool dout_din, bool d2out_din2) = 0;
  virtual void assemble(bool residual, bool Jacobian) override;
 
  virtual LabeledVector get_output();
  virtual LabeledMatrix get_doutput_dinput();
  virtual LabeledTensor3D get_d2output_dinput2();
 
  template <typename T, typename = typename std::enable_if_t<std::is_base_of_v<Model, T>>>
  T & register_model(const std::string & name,
                     int extra_deriv_order = 0,
                     bool nonlinear = false,
                     bool merge_input = true)
  {
    OptionSet extra_opts;
    extra_opts.set<NEML2Object *>("_host") = host();
    extra_opts.set<int>("_extra_derivative_order") = extra_deriv_order;
    extra_opts.set<bool>("_nonlinear_system") = nonlinear;
    extra_opts.set<bool>("_enable_AD") = input_options().get<bool>("_enable_AD");
 
    auto model = Factory::get_object_ptr<Model>("Models", name, extra_opts);
 
    if (merge_input)
      for (auto && [name, var] : model->input_variables())
        declare_input_variable(var.base_storage(), var.type(), name);
 
    _registered_models.push_back(model.get());
    return *(std::dynamic_pointer_cast<T>(model));
  }
 
  std::vector<Model *> _registered_models;
 
private:
  void input_requires_grad_(bool req = true);
  void extract_derivatives(bool retain_graph, bool create_graph, bool allow_unused);
  void extract_second_derivatives(bool retain_graph, bool create_graph, bool allow_unused);
 
  TensorShape _batch_sizes;
 
  torch::TensorOptions _options;
 
  bool _nonlinear_system;
 
  int _deriv_order;
 
  const int _extra_deriv_order;
 
  const bool _enable_AD;
 
  bool _AD_1st_deriv;
 
  bool _AD_2nd_deriv;
 
  Tensor _dr_ds, _dr_dsn, _dr_df, _dr_dfn, _dr_dp;
 
#ifndef NDEBUG
  bool _evaluated_once;
#endif
};
} // namespace neml2