Model.h

// Copyright 2023, 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/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:
  static OptionSet expected_options();
 
  Model(const OptionSet & options);
 
  virtual std::vector<Diagnosis> preflight() const;
 
  virtual bool is_nonlinear_system() const { return _nonlinear_system; }
 
  virtual void reinit(TorchShapeRef batch_shape,
                      int deriv_order = 0,
                      const torch::Device & device = torch::kCPU,
                      const torch::Dtype & dtype = NEML2_DTYPE);
 
  virtual void reinit(const BatchTensor & tensor, int deriv_order);
 
  bool requires_grad() const { return _deriv_order >= 1; }
 
  bool requires_2nd_grad() const { return _deriv_order >= 2; }
 
  TorchSize batch_dim() const { return _batch_sizes.size(); }
 
  TorchShapeRef batch_sizes() const { return _batch_sizes; }
 
  const torch::TensorOptions & options() const { return _options; }
 
  const std::vector<Model *> & registered_models() const { return _registered_models; }
 
  Model * registered_model(const std::string & name) const;
 
  virtual const std::set<VariableName> consumed_items() const override;
 
  virtual const std::set<VariableName> provided_items() const override;
 
  virtual void check_AD_limitation() const;
 
  void input_requires_grad_(bool req = true);
 
  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);
 
  virtual void set_input(const LabeledVector & in);
 
  virtual LabeledVector get_output();
 
  virtual LabeledMatrix get_doutput_dinput();
 
  virtual LabeledTensor3D get_d2output_dinput2();
 
  virtual LabeledVector value(const LabeledVector & in);
 
  virtual std::tuple<LabeledVector, LabeledMatrix> value_and_dvalue(const LabeledVector & in);
 
  virtual std::tuple<LabeledVector, LabeledMatrix, LabeledTensor3D>
  value_and_dvalue_and_d2value(const LabeledVector & in);
 
  virtual void value();
  virtual void value_and_dvalue();
  virtual void value_and_dvalue_and_d2value();
 
  static enum Stage { SOLVING, UPDATING } stage;
 
  friend class ParameterStore;
 
  friend class ComposedModel;
 
protected:
  virtual void setup() override;
 
  using VariableStore::allocate_variables;
 
  virtual void allocate_variables(int deriv_order, bool options_changed);
 
  virtual void setup_input_views() override;
  virtual void setup_submodel_input_views();
 
  virtual void setup_output_views() override;
  virtual void setup_submodel_output_views();
 
  virtual void reinit_input_views() override;
 
  virtual void reinit_output_views(bool out, bool dout_din = true, bool d2out_din2 = true) override;
 
  virtual void detach_and_zero(bool out, bool dout_din = true, bool d2out_din2 = true) override;
 
  virtual void set_solution(const BatchTensor & x) override;
 
  virtual void set_value(bool out, bool dout_din, bool d2out_din2) = 0;
 
  using VariableStore::cache;
 
  virtual void cache(TorchShapeRef batch_shape) override;
 
  virtual void cache(const torch::TensorOptions & options);
 
  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;
 
    auto model = Factory::get_object_ptr<Model>("Models", name, extra_opts, /*force_create=*/true);
 
    if (merge_input)
      for (auto && [name, var] : model->input_views())
        declare_input_variable(var.base_storage(), name);
 
    _registered_models.push_back(model.get());
    return *(std::dynamic_pointer_cast<T>(model));
  }
 
  virtual void assemble(bool residual, bool Jacobian) override;
 
  std::vector<Model *> _registered_models;
 
  bool _AD_1st_deriv;
 
  bool _AD_2nd_deriv;
 
private:
  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);
 
  TorchShape _batch_sizes;
 
  torch::TensorOptions _options;
 
  int _deriv_order;
 
  const int _extra_deriv_order;
 
  bool _nonlinear_system;
};
} // namespace neml2