LabeledTensor.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/misc/types.h"
#include "neml2/tensors/LabeledAxis.h"
#include "neml2/tensors/Tensor.h"
 
namespace neml2
{
template <class Derived, Size D>
class LabeledTensor
{
public:
  LabeledTensor() = default;
 
  LabeledTensor(const torch::Tensor & tensor, const std::array<const LabeledAxis *, D> & axes);
 
  LabeledTensor(const Tensor & tensor, const std::array<const LabeledAxis *, D> & axes);
 
  LabeledTensor(const Derived & other);
 
  LabeledTensor<Derived, D> & operator=(const Derived & other);
 
  operator Tensor() const;
  operator torch::Tensor() const;
 
  [[nodiscard]] static Derived
  empty(TensorShapeRef batch_shape,
        const std::array<const LabeledAxis *, D> & axes,
        const torch::TensorOptions & options = default_tensor_options());
 
  [[nodiscard]] static Derived
  zeros(TensorShapeRef batch_shape,
        const std::array<const LabeledAxis *, D> & axes,
        const torch::TensorOptions & options = default_tensor_options());
 
  const Tensor & tensor() const { return _tensor; }
  Tensor & tensor() { return _tensor; }
 
  // These methods mirror TensorBase
  Derived clone(torch::MemoryFormat memory_format = torch::MemoryFormat::Contiguous) const;
  Derived detach() const;
  void detach_();
  Derived to(const torch::TensorOptions & options) const;
  void copy_(const torch::Tensor & other);
  void zero_();
  bool requires_grad() const;
  void requires_grad_(bool req = true);
  Derived operator-() const;
 
  // These methods mirror TensorBase
  torch::TensorOptions options() const;
  torch::Dtype scalar_type() const;
  torch::Device device() const;
  Size dim() const;
  TensorShapeRef sizes() const;
  Size size(Size dim) const;
  bool batched() const;
  Size batch_dim() const;
  constexpr Size base_dim() const { return D; }
  TensorShapeRef batch_sizes() const;
  Size batch_size(Size d) const;
  TensorShapeRef base_sizes() const;
  Size base_size(Size d) const;
  Size base_storage() const;
 
  // These methods mirror TensorBase
  Derived batch_index(indexing::TensorIndicesRef indices) const;
  Tensor base_index(indexing::TensorLabelsRef labels) const;
  void batch_index_put_(indexing::TensorIndicesRef indices, const torch::Tensor & other);
  void base_index_put_(indexing::TensorLabelsRef labels, const Tensor & other);
 
  template <typename T, typename = std::enable_if_t<std::is_base_of_v<TensorBase<T>, T>>>
  T reinterpret(indexing::TensorLabelsRef indices) const;
 
  const std::array<const LabeledAxis *, D> & axes() const { return _axes; }
  const LabeledAxis & axis(Size i = 0) const { return *_axes[i]; }
 
protected:
  Tensor _tensor;
 
  std::array<const LabeledAxis *, D> _axes;
 
private:
  TensorShape storage_shape(indexing::TensorLabelsRef) const;
 
  indexing::TensorIndices labels_to_indices(indexing::TensorLabelsRef) const;
};
 
template <class Derived, Size D>
template <typename T, typename>
T
LabeledTensor<Derived, D>::reinterpret(indexing::TensorLabelsRef indices) const
{
  return base_index(indices).base_reshape(T::const_base_sizes);
}
} // namespace neml2