LabeledAxis.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 <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<Size, Size>> AxisLayout;
 
  // Custom comparator for sorting assembly indices
  struct AssemblySliceCmp
  {
    bool operator()(const indexing::TensorIndex & a, const indexing::TensorIndex & b) const
    {
      neml_assert(a.is_slice() && b.is_slice(), "Comparator must be used on slices");
      neml_assert(a.slice().step().expect_int() == 1 && b.slice().step().expect_int() == 1,
                  "Slices must have step == 1");
      return a.slice().start().expect_int() < b.slice().start().expect_int();
    }
  };
 
  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, Size sz);
 
  void clear();
 
  void setup_layout();
 
  bool has_state() const { return _has_state; }
  bool has_old_state() const { return _has_old_state; }
  bool has_forces() const { return _has_forces; }
  bool has_old_forces() const { return _has_old_forces; }
  bool has_residual() const { return _has_residual; }
  bool has_parameters() const { return _has_parameters; }
 
  size_t nvariable(bool recursive = true) const;
 
  size_t nsubaxis(bool recursive = false) const;
 
  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;
 
  Size storage_size(const LabeledAxisAccessor & name = {}) const;
 
  const AxisLayout & layout() const { return _layout; }
 
  indexing::TensorIndex indices(const LabeledAxisAccessor & accessor) const;
 
  std::vector<std::pair<indexing::TensorIndex, indexing::TensorIndex>>
  common_indices(const LabeledAxis & other, bool recursive = true) const;
 
  std::vector<LabeledAxisAccessor>
  sort_by_assembly_order(const std::set<LabeledAxisAccessor> &) const;
 
  const std::map<std::string, Size> & variables() const { return _variables; }
 
  std::set<LabeledAxisAccessor> variable_names(bool recursive = true) const;
 
  const std::map<std::string, std::shared_ptr<LabeledAxis>> & subaxes() const { return _subaxes; }
 
  std::set<LabeledAxisAccessor> subaxis_names(bool recursive = false) const;
 
  const LabeledAxis & subaxis(const LabeledAxisAccessor & name) const;
 
  LabeledAxis & subaxis(const LabeledAxisAccessor & name);
 
  bool equals(const LabeledAxis & other) const;
 
  friend std::ostream & operator<<(std::ostream & os, const LabeledAxis & axis);
 
private:
  void add(LabeledAxis & axis,
           Size sz,
           const LabeledAxisAccessor::const_iterator & cur,
           const LabeledAxisAccessor::const_iterator & end) const;
 
  Size storage_size(const LabeledAxisAccessor::const_iterator & cur,
                    const LabeledAxisAccessor::const_iterator & end) const;
 
  indexing::TensorIndex indices(Size offset,
                                const LabeledAxisAccessor::const_iterator & cur,
                                const LabeledAxisAccessor::const_iterator & end) const;
 
  void indices(const LabeledAxis & other,
               bool recursive,
               bool inclusive,
               std::vector<Size> & idx,
               Size offset) const;
 
  void common_indices(const LabeledAxis & other,
                      bool recursive,
                      std::vector<Size> & idxa,
                      std::vector<Size> & idxb,
                      Size offseta,
                      Size offsetb) const;
 
  std::map<std::string, Size> _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.
  Size _offset;
 
  bool _has_state;
  bool _has_old_state;
  bool _has_forces;
  bool _has_old_forces;
  bool _has_residual;
  bool _has_parameters;
};
 
std::ostream & operator<<(std::ostream & os, const LabeledAxis & axis);
 
bool operator==(const LabeledAxis & a, const LabeledAxis & b);
 
bool operator!=(const LabeledAxis & a, const LabeledAxis & b);
} // namespace neml2