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/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/TensorName.h"
#include "neml2/base/LabeledAxis.h"
#include "neml2/tensors/TensorValue.h"
#include "neml2/models/Variable.h"
 
namespace neml2
{
class Model;

Model & get_model(const std::string & mname);

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

Model & reload_model(const std::filesystem::path & path, const std::string & mname);

void check_precision();

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;
  };
  struct TraceSchema {…};
 
  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<Model *> & registered_models() const { return _registered_models; }
  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 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;

  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)
  {
    if (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;
 
    auto model = Factory::get_object_ptr<T>("Models", name, extra_opts);
    if (std::find(_registered_models.begin(), _registered_models.end(), model.get()) !=
        _registered_models.end())
      throw SetupException("Model named '" + name + "' has already been registered.");
 
    if (merge_input)
      for (auto && [name, var] : model->input_variables())
        clone_input_variable(*var);
 
    _registered_models.push_back(model.get());
    return *model;
  }
  T & register_model(const std::string & name, bool nonlinear = false, bool merge_input = true) {…}
 
  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<Model *> _registered_models;
 
private:
  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;
};
class Model : public std::enable_shared_from_this<Model>, {…};
 
std::ostream & operator<<(std::ostream & os, const Model & model);
} // namespace neml2