mirror of
https://github.com/slsdetectorgroup/aare.git
synced 2026-02-20 02:28:40 +01:00
273 lines
8.5 KiB
C++
273 lines
8.5 KiB
C++
// SPDX-License-Identifier: MPL-2.0
|
|
#pragma once
|
|
#include "aare/ArrayExpr.hpp"
|
|
#include "aare/defs.hpp"
|
|
|
|
#include <algorithm>
|
|
#include <array>
|
|
#include <cassert>
|
|
#include <cstdint>
|
|
#include <functional>
|
|
#include <iomanip>
|
|
#include <iostream>
|
|
#include <numeric>
|
|
#include <stdexcept>
|
|
#include <vector>
|
|
namespace aare {
|
|
|
|
template <ssize_t Ndim> using Shape = std::array<ssize_t, Ndim>;
|
|
|
|
// TODO! fix mismatch between signed and unsigned
|
|
template <ssize_t Ndim>
|
|
Shape<Ndim> make_shape(const std::vector<size_t> &shape) {
|
|
if (shape.size() != Ndim)
|
|
throw std::runtime_error("Shape size mismatch");
|
|
Shape<Ndim> arr;
|
|
std::copy_n(shape.begin(), Ndim, arr.begin());
|
|
return arr;
|
|
}
|
|
|
|
|
|
/**
|
|
* @brief Helper function to drop the first dimension of a shape.
|
|
* This is useful when you want to create a 2D view from a 3D array.
|
|
* @param shape The shape to drop the first dimension from.
|
|
* @return A new shape with the first dimension dropped.
|
|
*/
|
|
template<size_t Ndim>
|
|
Shape<Ndim-1> drop_first_dim(const Shape<Ndim> &shape) {
|
|
static_assert(Ndim > 1, "Cannot drop first dimension from a 1D shape");
|
|
Shape<Ndim - 1> new_shape;
|
|
std::copy(shape.begin() + 1, shape.end(), new_shape.begin());
|
|
return new_shape;
|
|
}
|
|
|
|
/**
|
|
* @brief Helper function when constructing NDArray/NDView. Calculates the number
|
|
* of elements in the resulting array from a shape.
|
|
* @param shape The shape to calculate the number of elements for.
|
|
* @return The number of elements in and NDArray/NDView of that shape.
|
|
*/
|
|
template <size_t Ndim>
|
|
size_t num_elements(const Shape<Ndim> &shape) {
|
|
return std::accumulate(shape.begin(), shape.end(), 1,
|
|
std::multiplies<size_t>());
|
|
}
|
|
|
|
template <ssize_t Dim = 0, typename Strides>
|
|
ssize_t element_offset(const Strides & /*unused*/) {
|
|
return 0;
|
|
}
|
|
|
|
template <ssize_t Dim = 0, typename Strides, typename... Ix>
|
|
ssize_t element_offset(const Strides &strides, ssize_t i, Ix... index) {
|
|
return i * strides[Dim] + element_offset<Dim + 1>(strides, index...);
|
|
}
|
|
|
|
template <ssize_t Ndim>
|
|
std::array<ssize_t, Ndim> c_strides(const std::array<ssize_t, Ndim> &shape) {
|
|
std::array<ssize_t, Ndim> strides{};
|
|
std::fill(strides.begin(), strides.end(), 1);
|
|
for (ssize_t i = Ndim - 1; i > 0; --i) {
|
|
strides[i - 1] = strides[i] * shape[i];
|
|
}
|
|
return strides;
|
|
}
|
|
|
|
template <ssize_t Ndim>
|
|
std::array<ssize_t, Ndim> make_array(const std::vector<ssize_t> &vec) {
|
|
assert(vec.size() == Ndim);
|
|
std::array<ssize_t, Ndim> arr{};
|
|
std::copy_n(vec.begin(), Ndim, arr.begin());
|
|
return arr;
|
|
}
|
|
|
|
template <typename T, ssize_t Ndim = 2>
|
|
class NDView : public ArrayExpr<NDView<T, Ndim>, Ndim> {
|
|
public:
|
|
NDView() = default;
|
|
~NDView() = default;
|
|
NDView(const NDView &) = default;
|
|
NDView(NDView &&) = default;
|
|
|
|
NDView(T *buffer, std::array<ssize_t, Ndim> shape)
|
|
: buffer_(buffer), strides_(c_strides<Ndim>(shape)), shape_(shape),
|
|
size_(std::accumulate(std::begin(shape), std::end(shape), 1,
|
|
std::multiplies<>())) {}
|
|
|
|
template <typename... Ix>
|
|
std::enable_if_t<sizeof...(Ix) == Ndim, T &> operator()(Ix... index) {
|
|
return buffer_[element_offset(strides_, index...)];
|
|
}
|
|
|
|
template <typename... Ix>
|
|
std::enable_if_t<sizeof...(Ix) == 1 && (Ndim > 1), NDView<T, Ndim - 1>> operator()(Ix... index) {
|
|
// return a view of the next dimension
|
|
std::array<ssize_t, Ndim - 1> new_shape{};
|
|
std::copy_n(shape_.begin() + 1, Ndim - 1, new_shape.begin());
|
|
return NDView<T, Ndim - 1>(&buffer_[element_offset(strides_, index...)],
|
|
new_shape);
|
|
|
|
}
|
|
|
|
template <typename... Ix>
|
|
std::enable_if_t<sizeof...(Ix) == Ndim, const T &> operator()(Ix... index) const {
|
|
return buffer_[element_offset(strides_, index...)];
|
|
}
|
|
|
|
|
|
ssize_t size() const { return static_cast<ssize_t>(size_); }
|
|
size_t total_bytes() const { return size_ * sizeof(T); }
|
|
std::array<ssize_t, Ndim> strides() const noexcept { return strides_; }
|
|
|
|
T *begin() { return buffer_; }
|
|
T *end() { return buffer_ + size_; }
|
|
T const *begin() const { return buffer_; }
|
|
T const *end() const { return buffer_ + size_; }
|
|
|
|
|
|
|
|
|
|
|
|
/**
|
|
* @brief Access element at index i.
|
|
*/
|
|
T &operator[](ssize_t i) { return buffer_[i]; }
|
|
|
|
/**
|
|
* @brief Access element at index i.
|
|
*/
|
|
const T &operator[](ssize_t i) const { return buffer_[i]; }
|
|
|
|
bool operator==(const NDView &other) const {
|
|
if (size_ != other.size_)
|
|
return false;
|
|
if (shape_ != other.shape_)
|
|
return false;
|
|
for (size_t i = 0; i != size_; ++i) {
|
|
if (buffer_[i] != other.buffer_[i])
|
|
return false;
|
|
}
|
|
return true;
|
|
}
|
|
|
|
NDView &operator+=(const T val) { return elemenwise(val, std::plus<T>()); }
|
|
NDView &operator-=(const T val) { return elemenwise(val, std::minus<T>()); }
|
|
NDView &operator*=(const T val) {
|
|
return elemenwise(val, std::multiplies<T>());
|
|
}
|
|
NDView &operator/=(const T val) {
|
|
return elemenwise(val, std::divides<T>());
|
|
}
|
|
|
|
NDView &operator/=(const NDView &other) {
|
|
return elemenwise(other, std::divides<T>());
|
|
}
|
|
|
|
template <size_t Size> NDView &operator=(const std::array<T, Size> &arr) {
|
|
if (size() != static_cast<ssize_t>(arr.size()))
|
|
throw std::runtime_error(LOCATION +
|
|
"Array and NDView size mismatch");
|
|
std::copy(arr.begin(), arr.end(), begin());
|
|
return *this;
|
|
}
|
|
|
|
NDView &operator=(const T val) {
|
|
for (auto it = begin(); it != end(); ++it)
|
|
*it = val;
|
|
return *this;
|
|
}
|
|
|
|
NDView &operator=(const NDView &other) {
|
|
if (this == &other)
|
|
return *this;
|
|
shape_ = other.shape_;
|
|
strides_ = other.strides_;
|
|
size_ = other.size_;
|
|
buffer_ = other.buffer_;
|
|
return *this;
|
|
}
|
|
|
|
NDView &operator=(NDView &&other) noexcept {
|
|
if (this == &other)
|
|
return *this;
|
|
shape_ = std::move(other.shape_);
|
|
strides_ = std::move(other.strides_);
|
|
size_ = other.size_;
|
|
buffer_ = other.buffer_;
|
|
other.buffer_ = nullptr;
|
|
return *this;
|
|
}
|
|
|
|
auto &shape() const { return shape_; }
|
|
auto shape(ssize_t i) const { return shape_[i]; }
|
|
|
|
T *data() { return buffer_; }
|
|
const T *data() const { return buffer_; }
|
|
void print_all() const;
|
|
|
|
/**
|
|
* @brief Create a subview of a range of the first dimension.
|
|
* This is useful for splitting a batches of frames in parallel processing.
|
|
* @param first The first index of the subview (inclusive).
|
|
* @param last The last index of the subview (exclusive).
|
|
* @return A new NDView that is a subview of the current view.
|
|
* @throws std::runtime_error if the range is invalid.
|
|
*/
|
|
NDView sub_view(ssize_t first, ssize_t last) const {
|
|
if (first < 0 || last > shape_[0] || first >= last)
|
|
throw std::runtime_error(LOCATION + "Invalid sub_view range");
|
|
auto new_shape = shape_;
|
|
new_shape[0] = last - first;
|
|
return NDView(buffer_ + first * strides_[0], new_shape);
|
|
}
|
|
|
|
private:
|
|
T *buffer_{nullptr};
|
|
std::array<ssize_t, Ndim> strides_{};
|
|
std::array<ssize_t, Ndim> shape_{};
|
|
uint64_t size_{};
|
|
|
|
template <class BinaryOperation>
|
|
NDView &elemenwise(T val, BinaryOperation op) {
|
|
for (uint64_t i = 0; i != size_; ++i) {
|
|
buffer_[i] = op(buffer_[i], val);
|
|
}
|
|
return *this;
|
|
}
|
|
template <class BinaryOperation>
|
|
NDView &elemenwise(const NDView &other, BinaryOperation op) {
|
|
for (uint64_t i = 0; i != size_; ++i) {
|
|
buffer_[i] = op(buffer_[i], other.buffer_[i]);
|
|
}
|
|
return *this;
|
|
}
|
|
};
|
|
|
|
template <typename T, ssize_t Ndim> void NDView<T, Ndim>::print_all() const {
|
|
for (auto row = 0; row < shape_[0]; ++row) {
|
|
for (auto col = 0; col < shape_[1]; ++col) {
|
|
std::cout << std::setw(3);
|
|
std::cout << (*this)(row, col) << " ";
|
|
}
|
|
std::cout << "\n";
|
|
}
|
|
}
|
|
|
|
template <typename T, ssize_t Ndim>
|
|
std::ostream &operator<<(std::ostream &os, const NDView<T, Ndim> &arr) {
|
|
for (auto row = 0; row < arr.shape(0); ++row) {
|
|
for (auto col = 0; col < arr.shape(1); ++col) {
|
|
os << std::setw(3);
|
|
os << arr(row, col) << " ";
|
|
}
|
|
os << "\n";
|
|
}
|
|
return os;
|
|
}
|
|
|
|
template <typename T> NDView<T, 1> make_view(std::vector<T> &vec) {
|
|
return NDView<T, 1>(vec.data(), {static_cast<ssize_t>(vec.size())});
|
|
}
|
|
|
|
} // namespace aare
|