Program Listing for File Tensor.hpp¶
↰ Return to documentation for file (include/jet/Tensor.hpp)
#pragma once
#include <complex>
#include <random>
#include <string>
#include <unordered_map>
#include <vector>
#include "TensorHelpers.hpp"
#include "Utilities.hpp"
#include "permute/PermuterIncludes.hpp"
namespace Jet {
template <class T = std::complex<float>> class Tensor {
static_assert(TensorHelpers::is_supported_data_type<T>,
"Tensor data type must be one of std::complex<float>, "
"std::complex<double>");
public:
using scalar_type_t = T;
Tensor() : data_(1) {}
Tensor(const std::vector<size_t> &shape)
: data_(Jet::Utilities::ShapeToSize(shape))
{
using namespace Utilities;
std::vector<std::string> indices(shape.size());
for (size_t i = 0; i < indices.size(); i++) {
indices[i] = "?" + GenerateStringIndex(i);
}
InitIndicesAndShape(indices, shape);
}
Tensor(const std::vector<std::string> &indices,
const std::vector<size_t> &shape)
: data_(Jet::Utilities::ShapeToSize(shape))
{
InitIndicesAndShape(indices, shape);
}
Tensor(const std::vector<std::string> &indices,
const std::vector<size_t> &shape, const std::vector<T> &data)
: Tensor(indices, shape)
{
Utilities::FastCopy(data, data_);
}
Tensor(const Tensor &other)
{
InitIndicesAndShape(other.GetIndices(), other.GetShape());
Utilities::FastCopy(other.data_, data_);
}
Tensor(Tensor &&other)
{
indices_ = std::move(other.indices_);
shape_ = std::move(other.shape_);
index_to_dimension_ = std::move(other.index_to_dimension_);
data_ = std::move(other.data_);
}
virtual ~Tensor() {}
void InitIndicesAndShape(const std::vector<std::string> &indices,
const std::vector<size_t> &shape) noexcept
{
indices_ = indices;
shape_ = shape;
index_to_dimension_.clear();
for (size_t i = 0; i < shape_.size(); ++i)
index_to_dimension_[indices[i]] = shape[i];
}
void SetShape(const std::vector<size_t> &shape) noexcept { shape_ = shape; }
const std::vector<size_t> &GetShape() const noexcept { return shape_; }
T &operator[](size_t pos) { return data_[pos]; }
const T &operator[](size_t pos) const { return data_[pos]; }
void RenameIndex(size_t pos, std::string new_label) noexcept
{
const auto old_label = indices_[pos];
const auto dimension = index_to_dimension_[old_label];
index_to_dimension_.erase(old_label);
indices_[pos] = new_label;
index_to_dimension_.emplace(new_label, dimension);
}
bool operator==(const Tensor<T> &other) const noexcept
{
return shape_ == other.GetShape() && indices_ == other.GetIndices() &&
index_to_dimension_ == other.GetIndexToDimension() &&
data_ == other.GetData();
}
bool operator!=(const Tensor<T> &other) const { return !(*this == other); }
const Tensor<T> &operator=(const Tensor<T> &other)
{
InitIndicesAndShape(other.GetIndices(), other.GetShape());
Utilities::FastCopy(other.GetData(), GetData());
return *this;
}
const Tensor<T> &operator=(Tensor<T> &&other)
{
indices_ = std::move(other.indices_);
shape_ = std::move(other.shape_);
index_to_dimension_ = std::move(other.index_to_dimension_);
data_ = std::move(other.data_);
return *this;
}
const std::unordered_map<std::string, size_t> &GetIndexToDimension() const
{
return index_to_dimension_;
}
void SetValue(const std::vector<size_t> &indices, const T &value)
{
data_[Jet::Utilities::RavelIndex(indices, shape_)] = value;
}
T GetValue(const std::vector<size_t> &indices) const
{
return data_[Jet::Utilities::RavelIndex(indices, shape_)];
}
void SetData(const std::vector<T> &data)
{
JET_ABORT_IF_NOT(data.size() == GetSize(),
"Size of data and tensor do not match.");
Utilities::FastCopy(data, data_);
}
const std::vector<T> &GetData() const noexcept { return data_; }
std::vector<T> &GetData() { return data_; }
const std::vector<std::string> &GetIndices() const noexcept
{
return indices_;
}
size_t GetSize() const { return data_.size(); }
const T &GetScalar() const { return data_[0]; }
bool IsScalar() const noexcept { return GetSize() == 1; }
void FillRandom(size_t seed)
{
std::mt19937 mt_engine(seed);
std::uniform_real_distribution<typename T::value_type> r_dist(-1, 1);
for (size_t i = 0; i < GetSize(); i++) {
data_[i] = {r_dist(mt_engine), r_dist(mt_engine)};
}
}
void FillRandom()
{
static std::mt19937 mt_engine(std::random_device{}());
static std::uniform_real_distribution<typename T::value_type> r_dist(-1,
1);
for (size_t i = 0; i < GetSize(); i++) {
data_[i] = {r_dist(mt_engine), r_dist(mt_engine)};
}
}
template <class U = T>
static Tensor<U> AddTensors(const Tensor<U> &A, const Tensor<U> &B)
{
static const Tensor<U> zero;
// The zero tensor is used in reductions where the shape of an
// accumulator is not known beforehand.
if (A == zero) {
return B;
}
else if (B == zero) {
return A;
}
const auto disjoint_indices = Jet::Utilities::VectorDisjunctiveUnion(
A.GetIndices(), B.GetIndices());
JET_ABORT_IF_NOT(
disjoint_indices.empty(),
"Tensor addition with disjoint indices is not supported.");
const auto &indices = A.GetIndices();
const auto &shape = A.GetShape();
// Align the underlying data vectors of `A` and `B`.
const auto &&Bt = Transpose(B, indices);
Tensor<U> C(indices, shape);
const auto size = C.GetSize();
auto c_ptr = C.GetData().data();
auto a_ptr = A.GetData().data();
auto bt_ptr = Bt.GetData().data();
#if defined _OPENMP
#pragma omp parallel for schedule(static, 1024) // MAX_RIGHT_DIM)
#endif
for (size_t i = 0; i < size; i++) {
c_ptr[i] = a_ptr[i] + bt_ptr[i];
}
return C;
}
Tensor<T> AddTensor(const Tensor<T> &other) const
{
return AddTensors<T>(*this, other);
}
template <class U = T>
static Tensor<U> SliceIndex(const Tensor<U> &tensor,
const std::string &index, size_t value)
{
std::vector<std::string> new_ordering = tensor.GetIndices();
auto it = find(new_ordering.begin(), new_ordering.end(), index);
size_t index_num = std::distance(new_ordering.begin(), it);
new_ordering.erase(new_ordering.begin() + index_num);
new_ordering.insert(new_ordering.begin(), index);
auto &&tensor_trans = Transpose(tensor, new_ordering);
std::vector<std::string> sliced_indices(
tensor_trans.GetIndices().begin() + 1,
tensor_trans.GetIndices().end());
std::vector<size_t> sliced_dimensions(
tensor_trans.GetShape().begin() + 1, tensor_trans.GetShape().end());
Tensor<U> tensor_sliced(sliced_indices, sliced_dimensions);
size_t projection_size = tensor_sliced.GetSize();
size_t projection_begin = projection_size * value;
auto data_ptr = tensor_trans.GetData().data();
auto tensor_sliced_ptr = tensor_sliced.GetData().data();
#if defined _OPENMP
int max_right_dim = 1024;
#pragma omp parallel for schedule(static, max_right_dim)
#endif
for (size_t p = 0; p < projection_size; ++p)
tensor_sliced_ptr[p] = data_ptr[projection_begin + p];
return tensor_sliced;
}
Tensor<T> SliceIndex(const std::string &index, size_t value) const
{
return SliceIndex<T>(*this, index, value);
}
template <class U = T>
static Tensor<U> Reshape(const Tensor<U> &old_tensor,
const std::vector<size_t> &new_shape)
{
using namespace Utilities;
JET_ABORT_IF_NOT(old_tensor.GetSize() ==
Jet::Utilities::ShapeToSize(new_shape),
"Size is inconsistent between tensors.");
Tensor<U> new_tensor(new_shape);
Utilities::FastCopy(old_tensor.GetData(), new_tensor.GetData());
return new_tensor;
}
Tensor<T> Reshape(const std::vector<size_t> &new_shape) const
{
return Reshape<T>(*this, new_shape);
}
template <class U = T, size_t BLOCKSIZE = 1024, size_t MINSIZE = 32>
static Tensor<U> Transpose(const Tensor<U> &A,
const std::vector<std::string> &new_indices)
{
using namespace Jet::Utilities;
if (A.GetIndices() == new_indices)
return A;
if (A.GetIndices().size() == 0)
JET_ABORT("Number of indices cannot be zero.");
std::vector<size_t> new_shape(A.GetShape().size());
for (size_t i = 0; i < new_indices.size(); i++)
new_shape[i] = A.GetIndexToDimension().at(new_indices[i]);
if (is_pow_2(A.GetSize())) {
auto permuter = Permuter<QFlexPermuter<BLOCKSIZE, MINSIZE>>();
try {
return Tensor<U>{new_indices, new_shape,
permuter.Transpose(A.GetData(), A.GetShape(),
A.GetIndices(),
new_indices)};
}
catch (Jet::Exception &e) {
std::cerr << "Error in fast transpose with given parameters:="
<< e.what() << std::endl;
std::cerr << "Using fall-back default transpose." << std::endl;
}
}
auto permuter = Permuter<DefaultPermuter<BLOCKSIZE>>();
return Tensor<U>{new_indices, new_shape,
permuter.Transpose(A.GetData(), A.GetShape(),
A.GetIndices(), new_indices)};
}
template <class U = T, size_t BLOCKSIZE = 1024, size_t MINSIZE = 32>
static Tensor<U> Transpose(const Tensor<U> &A,
const std::vector<size_t> &new_ordering)
{
const size_t num_indices = A.GetIndices().size();
JET_ABORT_IF_NOT(
num_indices == new_ordering.size(),
"Size of ordering must match number of tensor indices.");
std::vector<std::string> new_indices(num_indices);
const auto &old_indices = A.GetIndices();
for (size_t i = 0; i < num_indices; i++) {
new_indices[i] = old_indices[new_ordering[i]];
}
return Transpose<U, BLOCKSIZE, MINSIZE>(A, new_indices);
}
Tensor<T> Transpose(const std::vector<size_t> &new_ordering) const
{
return Transpose<T>(*this, new_ordering);
}
Tensor<T> Transpose(const std::vector<std::string> &new_indices) const
{
return Transpose<T>(*this, new_indices);
}
template <class U = T> static Tensor<U> Conj(const Tensor<U> &A)
{
Tensor<U> A_conj(A.GetIndices(), A.GetShape());
for (size_t i = 0; i < A.GetSize(); i++) {
A_conj[i] = std::conj(A[i]);
}
return A_conj;
}
Tensor<T> Conj() const { return Conj<T>(*this); }
template <class U = T>
static Tensor<U> ContractTensors(const Tensor<U> &A, const Tensor<U> &B)
{
using namespace Jet::Utilities;
using namespace Jet::TensorHelpers;
auto &&left_indices = VectorSubtraction(A.GetIndices(), B.GetIndices());
auto &&right_indices =
VectorSubtraction(B.GetIndices(), A.GetIndices());
auto &&common_indices =
VectorIntersection(A.GetIndices(), B.GetIndices());
size_t left_dim = 1, right_dim = 1, common_dim = 1;
for (size_t i = 0; i < left_indices.size(); ++i) {
left_dim *= A.GetIndexToDimension().at(left_indices[i]);
}
for (size_t i = 0; i < right_indices.size(); ++i) {
right_dim *= B.GetIndexToDimension().at(right_indices[i]);
}
for (size_t i = 0; i < common_indices.size(); ++i) {
size_t a_dim = A.GetIndexToDimension().at(common_indices[i]);
common_dim *= a_dim;
}
auto &&a_new_ordering = VectorUnion(left_indices, common_indices);
auto &&b_new_ordering = VectorUnion(common_indices, right_indices);
auto &&C_indices = VectorUnion(left_indices, right_indices);
std::vector<size_t> C_dimensions(C_indices.size());
for (size_t i = 0; i < left_indices.size(); ++i)
C_dimensions[i] = A.GetIndexToDimension().at(left_indices[i]);
for (size_t i = 0; i < right_indices.size(); ++i)
C_dimensions[i + left_indices.size()] =
B.GetIndexToDimension().at(right_indices[i]);
Tensor<U> C(C_indices, C_dimensions);
auto &&At = Transpose<U, 1024UL, 32UL>(A, a_new_ordering);
auto &&Bt = Transpose<U, 1024UL, 32UL>(B, b_new_ordering);
TensorHelpers::MultiplyTensorData<U>(
At.GetData(), Bt.GetData(), C.GetData(), left_indices,
right_indices, left_dim, right_dim, common_dim);
return C;
}
Tensor<T> ContractWithTensor(const Tensor<T> &other) const
{
return ContractTensors<T>(*this, other);
}
private:
std::vector<std::string> indices_;
std::vector<size_t> shape_;
std::unordered_map<std::string, size_t> index_to_dimension_;
std::vector<T> data_;
};
template <class T>
inline std::ostream &operator<<(std::ostream &out, const Tensor<T> &tensor)
{
using namespace Jet::Utilities;
out << "Size = " << tensor.GetSize() << std::endl;
out << "Indices = " << tensor.GetIndices() << std::endl;
out << "Data = " << tensor.GetData();
return out;
}
}; // namespace Jet
api/program_listing_file_include_jet_Tensor.hpp
Download Python script
Download Notebook
View on GitHub
Downloads