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"
 
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 to(const torch::TensorOptions & options);
 
  void diagnose(std::vector<Diagnosis> &) const override;
 
  virtual bool is_nonlinear_system() const { return _nonlinear_system; }
 
  const std::vector<Model *> & registered_models() const { return _registered_models; }
  Model * registered_model(const std::string & name) const;
 
  std::set<VariableName> consumed_items() const override;
  std::set<VariableName> provided_items() const override;
 
  void clear_input() override;
  void clear_output() override;
  void zero_input() override;
  void zero_output() override;
 
  void request_AD(VariableBase & y, const VariableBase & u);
 
  void request_AD(VariableBase & y, const VariableBase & u1, const VariableBase & u2);
 
  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 ValueMap value(const ValueMap & in);
 
  virtual std::tuple<ValueMap, DerivMap> value_and_dvalue(const ValueMap & in);
 
  virtual DerivMap dvalue(const ValueMap & in);
 
  virtual std::tuple<ValueMap, DerivMap, SecDerivMap>
  value_and_dvalue_and_d2value(const ValueMap & in);
 
  virtual SecDerivMap d2value(const ValueMap & in);
 
  virtual std::tuple<DerivMap, SecDerivMap> dvalue_and_d2value(const ValueMap & in);
 
  friend class ParameterStore;
 
  friend class ComposedModel;
 
protected:
  void setup() override;
  virtual void link_input_variables();
  virtual void link_input_variables(Model * submodel);
  virtual void link_output_variables();
  virtual void link_output_variables(Model * submodel);
 
  virtual void request_AD() {}
 
  void diagnose_nl_sys(std::vector<Diagnosis> & diagnoses) const;
 
  virtual void set_value(bool out, bool dout_din, bool d2out_din2) = 0;
 
  template <typename T, typename = typename std::enable_if_t<std::is_base_of_v<Model, T>>>
  T & register_model(const std::string & name, bool nonlinear = false, bool merge_input = true)
  {
    neml_assert(name != this->name(),
                "Model named '",
                this->name(),
                "' is trying to register itself as a sub-model. This is not allowed.");
 
    OptionSet extra_opts;
    extra_opts.set<NEML2Object *>("_host") = host();
    extra_opts.set<bool>("_nonlinear_system") = nonlinear;
 
    auto model = Factory::get_object_ptr<Model>("Models", name, extra_opts);
 
    if (merge_input)
      for (auto && [name, var] : model->input_variables())
        clone_input_variable(var);
 
    _registered_models.push_back(model.get());
    return *(std::dynamic_pointer_cast<T>(model));
  }
 
  void set_guess(const Sol<false> &) override;
 
  void assemble(Res<false> *, Jac<false> *) override;
 
  std::vector<Model *> _registered_models;
 
private:
  bool AD_need_value(bool dout, bool d2out) const;
 
  void enable_AD();
 
  void extract_AD_derivatives(bool dout, bool d2out);
 
  bool _nonlinear_system;
 
  std::map<VariableBase *, std::set<const VariableBase *>> _ad_derivs;
  std::map<VariableBase *, std::map<const VariableBase *, std::set<const VariableBase *>>>
      _ad_secderivs;
  std::set<VariableBase *> _ad_args;
};
 
std::ostream & operator<<(std::ostream & os, const Model & model);
} // namespace neml2