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/models/DependencyDefinition.h"
#include "neml2/base/DiagnosticsInterface.h"
 
#include "neml2/models/Data.h"
#include "neml2/models/ParameterStore.h"
#include "neml2/models/VariableStore.h"
#include "neml2/models/Variable.h"
#include "neml2/solvers/NonlinearSystem.h"
#include "neml2/models/NonlinearParameter.h"
 
// These headers are not directly used by Model, but are included here so that derived classes do
// not have to include them separately. This is a convenience for the user, and is a reasonable
// choice since these headers are light and bring in little dependency.
#include "neml2/base/LabeledAxis.h"
 
namespace neml2
{
class Model;

using ModelCallback =
    std::function<void(const Model &,
                       const std::map<VariableName, std::unique_ptr<VariableBase>> &,
                       const std::map<VariableName, std::unique_ptr<VariableBase>> &)>;

std::shared_ptr<Model> load_model(const std::filesystem::path & path, const std::string & mname);

class Model : public std::enable_shared_from_this<Model>,
              public Data,
              public ParameterStore,
              public VariableStore,
              public NonlinearSystem,
              public DependencyDefinition<VariableName>,
              public DiagnosticsInterface
{
public:
  struct TraceSchema
  {
    std::vector<Size> batch_dims;
    at::DispatchKey dispatch_key;
    bool operator==(const TraceSchema & other) const;
    bool operator<(const TraceSchema & other) const;
  };
 
  static OptionSet expected_options();

  Model(const OptionSet & options);
 
  void setup() override;
 
  void diagnose() const override;

  virtual bool defines_values() const { return _defines_value; }

  virtual bool defines_derivatives() const { return _defines_dvalue; }

  virtual bool defines_second_derivatives() const { return _defines_d2value; }

  virtual bool is_nonlinear_system() const { return _nonlinear_system; }

  virtual bool is_jit_enabled() const { return _jit; }

  virtual void to(const TensorOptions & options);

  const std::vector<std::shared_ptr<Model>> & registered_models() const
  {
    return _registered_models;
  }

  std::shared_ptr<Model> registered_model(const std::string & name) const;

  void register_nonlinear_parameter(const std::string & pname, const NonlinearParameter & param);

  bool has_nl_param(bool recursive = false) const;

  const VariableBase * nl_param(const std::string &) const;

  virtual std::map<std::string, NonlinearParameter>
  named_nonlinear_parameters(bool recursive = false) const;

  std::set<VariableName> consumed_items() const override;
  std::set<VariableName> provided_items() const override;

  void request_AD(VariableBase & y, const VariableBase & u);

  void request_AD(VariableBase & y, const VariableBase & u1, const VariableBase & u2);

  void register_callback(const ModelCallback & callback);

  void register_callback_recursive(const ModelCallback & callback);

  void forward(bool out, bool dout, bool d2out);

  void forward_maybe_jit(bool out, bool dout, bool d2out);

  std::string variable_name_lookup(const ATensor & var);

  virtual ValueMap value(const ValueMap & in);
  virtual ValueMap value(ValueMap && in);

  virtual std::tuple<ValueMap, DerivMap> value_and_dvalue(const ValueMap & in);
  virtual std::tuple<ValueMap, DerivMap> value_and_dvalue(ValueMap && in);

  virtual DerivMap dvalue(const ValueMap & in);
  virtual DerivMap dvalue(ValueMap && in);

  virtual std::tuple<ValueMap, DerivMap, SecDerivMap>
  value_and_dvalue_and_d2value(const ValueMap & in);
  virtual std::tuple<ValueMap, DerivMap, SecDerivMap> value_and_dvalue_and_d2value(ValueMap && in);

  virtual SecDerivMap d2value(const ValueMap & in);
  virtual SecDerivMap d2value(ValueMap && in);

  virtual std::tuple<DerivMap, SecDerivMap> dvalue_and_d2value(const ValueMap & in);
  virtual std::tuple<DerivMap, SecDerivMap> dvalue_and_d2value(ValueMap && in);

  friend class ParameterStore;

  friend class ComposedModel;
 
protected:
  void diagnostic_assert_state(const VariableBase & v) const;
  void diagnostic_assert_old_state(const VariableBase & v) const;
  void diagnostic_assert_force(const VariableBase & v) const;
  void diagnostic_assert_old_force(const VariableBase & v) const;
  void diagnostic_assert_residual(const VariableBase & v) const;
  void diagnostic_check_input_variable(const VariableBase & v) const;
  void diagnostic_check_output_variable(const VariableBase & v) const;

  void diagnose_nl_sys() const;
 
  virtual void link_input_variables();
  virtual void link_input_variables(Model * submodel);
  virtual void link_output_variables();
  virtual void link_output_variables(Model * submodel);
 
  void clear_input() override;
  void clear_output() override;
  void zero_input() override;
  void zero_output() override;

  void check_precision() const;

  virtual void request_AD() {}

  virtual void set_value(bool out, bool dout_din, bool d2out_din2) = 0;

  template <typename T = Model, 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)
  {
    auto model_name =
        input_options().contains(name) ? input_options().get<std::string>(name) : name;
    if (model_name == this->name())
      throw SetupException("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;
 
    if (!host()->factory())
      throw SetupException("Internal error: Host object '" + host()->name() +
                           "' does not have a factory set.");
    auto model = host()->factory()->get_object<T>("Models", model_name, extra_opts);
    if (std::find(_registered_models.begin(), _registered_models.end(), model) !=
        _registered_models.end())
      throw SetupException("Model named '" + model_name + "' has already been registered.");
 
    if (merge_input)
      for (auto && [name, var] : model->input_variables())
        clone_input_variable(*var);
 
    _registered_models.push_back(model);
    return *model;
  }
 
  void assign_input_stack(jit::Stack & stack);
 
  jit::Stack collect_input_stack() const;
 
  void set_guess(const Sol<false> &) override;
 
  void assemble(Res<false> *, Jac<false> *) override;

  std::vector<std::shared_ptr<Model>> _registered_models;
 
private:
  void call_callbacks() const;
 
  template <typename T>
  void forward_helper(T && in, bool out, bool dout, bool d2out)
  {
    check_precision();
    zero_input();
    assign_input(std::forward<T>(in));
    zero_output();
    forward_maybe_jit(out, dout, d2out);
  }

  bool AD_need_value(bool dout, bool d2out) const;

  void enable_AD();

  void extract_AD_derivatives(bool dout, bool d2out);

  std::size_t forward_operator_index(bool out, bool dout, bool d2out) const;

  TraceSchema compute_trace_schema() const;

  bool _defines_value;
  bool _defines_dvalue;
  bool _defines_d2value;

  bool _nonlinear_system;

  std::map<std::string, NonlinearParameter> _nl_params;

  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;

  const bool _jit;

  const bool _production;

  std::array<std::map<TraceSchema, std::unique_ptr<jit::GraphFunction>>, 8> _traced_functions;

  std::array<std::map<TraceSchema, std::unique_ptr<jit::GraphFunction>>, 8>
      _traced_functions_nl_sys;

  std::vector<ModelCallback> _callbacks;
};
 
std::ostream & operator<<(std::ostream & os, const Model & model);
} // namespace neml2