TensorVectorConversion.hpp

Go to the documentation of this file. Or view the newest version in sPHENIX GitHub for file TensorVectorConversion.hpp

// This file is part of the Acts project.
//
// Copyright (C) 2023 CERN for the benefit of the Acts project
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

#pragma once

#include "Acts/Utilities/Concepts.hpp"

#include <cstdint>
#include <vector>

#include <torch/torch.h>

namespace Acts::detail {

template <typename T>
struct TorchTypeMap {};

template <>
struct TorchTypeMap<int64_t> {
  constexpr static torch::Dtype type = torch::kInt64;
};

template <>
struct TorchTypeMap<int32_t> {
  constexpr static torch::Dtype type = torch::kInt32;
};

template <>
struct TorchTypeMap<int16_t> {
  constexpr static torch::Dtype type = torch::kInt16;
};

template <>
struct TorchTypeMap<int8_t> {
  constexpr static torch::Dtype type = torch::kInt8;
};

template <>
struct TorchTypeMap<float> {
  constexpr static torch::Dtype type = torch::kFloat32;
};

template <>
struct TorchTypeMap<double> {
  constexpr static torch::Dtype type = torch::kFloat64;
};

template <typename T>
at::Tensor vectorToTensor2D(std::vector<T> &vec, std::size_t cols) {
  assert(vec.size() % cols == 0);

  auto opts =
      at::TensorOptions().dtype(TorchTypeMap<T>::type).device(torch::kCPU);

  return torch::from_blob(
      vec.data(),
      {static_cast<long>(vec.size() / cols), static_cast<long>(cols)}, opts);
}

template <typename T>
std::vector<T> tensor2DToVector(const at::Tensor &tensor) {
  assert(tensor.sizes().size() == 2);

  // clone to make sure we own the data
  // bring to CPU
  // convert to requested type
  // ensure the tensor is contiguous (e.g. not the case if indexed with step)

  at::Tensor transformedTensor =
      tensor.to(torch::kCPU).to(TorchTypeMap<T>::type).contiguous();

  std::vector<T> edgeIndex(
      transformedTensor.template data_ptr<T>(),
      transformedTensor.template data_ptr<T>() + transformedTensor.numel());

  return edgeIndex;
}

}  // namespace Acts::detail