LabeledAxis.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 <unordered_map>
#include <type_traits>
 
#include "neml2/misc/types.h"
 
#include "neml2/tensors/LabeledAxisAccessor.h"
#include "neml2/tensors/Scalar.h"
#include "neml2/tensors/SR2.h"
 
namespace neml2
{
class LabeledAxis
{
public:
  typedef std::unordered_map<std::string, std::pair<TorchSize, TorchSize>> AxisLayout;
 
  LabeledAxis();
 
  LabeledAxis(const LabeledAxis & other);
 
  template <typename T>
  LabeledAxis & add(const LabeledAxisAccessor & accessor)
  {
    // Add an empty subaxis
    if constexpr (std::is_same_v<LabeledAxis, T>)
    {
      if (!accessor.empty() && !has_subaxis(accessor.slice(0, 1)))
      {
        _subaxes.emplace(accessor.vec()[0], std::make_shared<LabeledAxis>());
        subaxis(accessor.vec()[0]).add<LabeledAxis>(accessor.slice(1));
      }
      return *this;
    }
    else
    {
      // The storage is *flat* -- will need to reshape when we return!
      // All NEML2 primitive data types will have the member const_base_sizes
      auto sz = utils::storage_size(T::const_base_sizes);
      add(accessor, sz);
      return *this;
    }
  }
 
  LabeledAxis & add(const LabeledAxisAccessor & accessor, TorchSize sz);
 
  LabeledAxis & rename(const std::string & original, const std::string & rename);
 
  LabeledAxis & remove(const std::string & name);
 
  LabeledAxis & clear();
 
  std::vector<LabeledAxisAccessor> merge(LabeledAxis & other);
 
  void setup_layout();
 
  size_t nitem() const { return nvariable() + nsubaxis(); }
 
  size_t nvariable() const { return _variables.size(); }
 
  size_t nsubaxis() const { return _subaxes.size(); }
 
  bool has_item(const LabeledAxisAccessor & name) const
  {
    return has_variable(name) || has_subaxis(name);
  }
 
  template <typename T>
  bool has_variable(const LabeledAxisAccessor & var) const
  {
    if (!has_variable(var))
      return false;
 
    // also check size
    return storage_size(var) == utils::storage_size(T::const_base_sizes);
  }
 
  bool has_variable(const LabeledAxisAccessor & var) const;
 
  bool has_subaxis(const LabeledAxisAccessor & s) const;
 
  TorchSize storage_size() const { return _offset; }
  TorchSize storage_size(const LabeledAxisAccessor & accessor) const;
 
  const AxisLayout & layout() const { return _layout; }
 
  TorchIndex indices(const LabeledAxisAccessor & accessor) const;
  TorchIndex indices(const LabeledAxis & other, bool recursive = true, bool inclusive = true) const;
 
  std::vector<std::pair<TorchIndex, TorchIndex>> common_indices(const LabeledAxis & other,
                                                                bool recursive = true) const;
 
  std::vector<std::string> item_names() const;
 
  const std::map<std::string, TorchSize> & variables() const { return _variables; }
 
  const std::map<std::string, std::shared_ptr<LabeledAxis>> & subaxes() const { return _subaxes; }
 
  std::set<LabeledAxisAccessor> variable_accessors(bool recursive = false,
                                                   const LabeledAxisAccessor & subaxis = {}) const;
 
  const LabeledAxis & subaxis(const std::string & name) const;
  LabeledAxis & subaxis(const std::string & name);
 
  bool equals(const LabeledAxis & other) const;
 
  friend std::ostream & operator<<(std::ostream & os, const LabeledAxis & info);
 
  void to_dot(std::ostream & os,
              int & id,
              std::string name = "",
              bool subgraph = false,
              bool node_handle = false) const;
 
private:
  void add(LabeledAxis & axis,
           TorchSize sz,
           const std::vector<std::string>::const_iterator & cur,
           const std::vector<std::string>::const_iterator & end) const;
 
  void merge(LabeledAxis & other,
             std::vector<std::string> subaxes,
             std::vector<LabeledAxisAccessor> & merged_vars);
 
  TorchSize storage_size(const std::vector<std::string>::const_iterator & cur,
                         const std::vector<std::string>::const_iterator & end) const;
 
  TorchIndex indices(TorchSize offset,
                     const std::vector<std::string>::const_iterator & cur,
                     const std::vector<std::string>::const_iterator & end) const;
 
  void indices(const LabeledAxis & other,
               bool recursive,
               bool inclusive,
               std::vector<TorchSize> & idx,
               TorchSize offset) const;
 
  void common_indices(const LabeledAxis & other,
                      bool recursive,
                      std::vector<TorchSize> & idxa,
                      std::vector<TorchSize> & idxb,
                      TorchSize offseta,
                      TorchSize offsetb) const;
 
  void variable_accessors(std::set<LabeledAxisAccessor> & accessors,
                          LabeledAxisAccessor cur,
                          bool recursive,
                          const LabeledAxisAccessor & subaxis) const;
 
  std::map<std::string, TorchSize> _variables;
 
  // Each sub-axis can contain its own variables and sub-axes
  std::map<std::string, std::shared_ptr<LabeledAxis>> _subaxes;
 
  // After all the `LabeledAxis`s are setup, we need to setup the layout once and only once. This is
  // important for performance considerations, as we need to use the layout to construct many,
  // many LabeledVector and LabeledMatrix at runtime, and so we don't want to waste time on setting
  // up the layout over and over again.
  AxisLayout _layout;
 
  // Similar considerations as `_layout`, i.e., the _offset will be zero during the setup stage,
  // and will have a fixed (hopefully correct) size after the layout have been setup.
  TorchSize _offset;
};
 
std::ostream & operator<<(std::ostream & os, const LabeledAxis & info);
 
bool operator==(const LabeledAxis & a, const LabeledAxis & b);
 
bool operator!=(const LabeledAxis & a, const LabeledAxis & b);
} // namespace neml2