borg_public/libLSS/tools/array_tools.hpp
2023-05-29 10:41:03 +02:00

316 lines
11 KiB
C++

/*+
ARES/HADES/BORG Package -- ./libLSS/tools/array_tools.hpp
Copyright (C) 2014-2020 Guilhem Lavaux <guilhem.lavaux@iap.fr>
Copyright (C) 2009-2020 Jens Jasche <jens.jasche@fysik.su.se>
Additional contributions from:
Guilhem Lavaux <guilhem.lavaux@iap.fr> (2023)
+*/
#ifndef __LIBLSS_ARRAY_TOOLS_HPP
#define __LIBLSS_ARRAY_TOOLS_HPP
#include "libLSS/tools/errors.hpp"
#include "libLSS/tools/align_helper.hpp"
#include <CosmoTool/omptl/omptl>
#include <CosmoTool/omptl/omptl_algorithm>
#include <boost/lambda/lambda.hpp>
#include "libLSS/tools/fused_array.hpp"
namespace LibLSS {
namespace array {
typedef boost::multi_array_types::extent_range erange;
typedef boost::multi_array_types::index_range irange;
template<typename InArray>
struct EigenMap {
typedef typename InArray::element T;
typedef Eigen::Array<T, Eigen::Dynamic, 1> E_Array;
typedef Eigen::Map<E_Array, Eigen::Unaligned> MapArray;
static MapArray map(InArray& a) {
return MapArray(a.data(), a.num_elements());
}
};
template<typename InArray>
typename EigenMap<InArray>::MapArray eigen_map(InArray& a) {
return EigenMap<InArray>::map(a);
}
template<typename Array>
void fill(Array& a, typename Array::element val) {
using boost::lambda::constant;
typedef typename Array::element VArrayType;
LibLSS::copy_array(a, b_fused<VArrayType,Array::dimensionality>(constant(val)));
}
template<typename Array>
void density_rescale(Array& a, typename Array::element val) {
using boost::lambda::_1;
LibLSS::copy_array(a, b_fused<typename Array::element>(a, _1 / val - 1));//omptl::for_each(a.data(), a.data()+a.num_elements(), _1 / val - 1);
}
template<typename InArray, typename OutArray>
void copyArray3d(OutArray& out, const InArray& in, bool in_padded = false)
{
if (!in_padded && (out.shape()[0] < in.shape()[0] ||
out.shape()[1] < in.shape()[1] ||
out.shape()[2] < in.shape()[2])) {
error_helper<ErrorBadState>("Invalid copy shape in copyArray3d");
}
LibLSS::copy_array(out, in);
}
template<typename InOutArray>
void scaleArray3d(InOutArray& out, typename InOutArray::element scale)
{
using boost::lambda::_1;
LibLSS::copy_array(out, b_fused<typename InOutArray::element>(out, _1*scale));
}
template<typename T, size_t n>
T product(std::array<T,n> const& d) {
T a = T(1);
for (size_t i = 0; i < n; i++) a *= d[i];
return a;
}
template<typename Iterator, typename T = typename std::remove_reference<decltype(*std::declval<Iterator>())>::type>
T product(Iterator b, Iterator e) {
T a = T(1);
while (b != e) {
a *= *b;
++b;
}
return a;
}
template<typename InOutPlane>
void scalePlane(InOutPlane plane, typename InOutPlane::element scale)
{
const typename InOutPlane::index *base = plane.index_bases();
const typename InOutPlane::size_type *exts = plane.shape();
for (long i = base[0]; i < base[0] + exts[0]; i++) {
for (long j = base[1]; j < base[1] + exts[1]; j++) {
plane[i][j] *= scale;
}
}
}
template<typename InOutLine>
void scaleLine(InOutLine line, typename InOutLine::element scale)
{
const typename InOutLine::index *base = line.index_bases();
const typename InOutLine::size_type *exts = line.shape();
for (long i = base[0]; i < base[0] + exts[0]; i++) {
line[i] *= scale;
}
}
template<typename InArray, typename OutArray>
void scaleAndCopyArray3d_rv(OutArray out, const InArray& in, typename OutArray::element scale)
{
using boost::format;
long N0 = out.shape()[0],
N1 = out.shape()[1],
N2 = out.shape()[2];
Console& cons = Console::instance();
N0 = std::min(N0,long(in.shape()[0]));
N1 = std::min(N1,long(in.shape()[1]));
N2 = std::min(N2,long(in.shape()[2]));
long s0 = out.index_bases()[0],
s1 = out.index_bases()[1],
s2 = out.index_bases()[2];
long i_s0 = in.index_bases()[0],
i_s1 = in.index_bases()[1],
i_s2 = in.index_bases()[2];
cons.print<LOG_DEBUG>(format("Copying (%d-%d, %d-%d, %d-%d) -> (%d-%d, %d-%d, %d-%d)")
% i_s0 % (i_s0+N0) % i_s1 % (i_s1+N1) % i_s2 % (i_s2+N2) %
s0 % (s0+N0) % s1 % (s1+N1) % s2 % (s2+N2));
#pragma omp parallel for
for (long n0 = 0; n0 < N0; n0++) {
typename OutArray::reference out0 = out[s0+n0];
typename InArray::const_reference in0 = in[i_s0+n0];
cons.print<LOG_DEBUG>(format("Line %d") % n0);
for (long n1 = 0; n1 < N1; n1++) {
typename OutArray::reference::reference out1 = out0[s1+n1];
typename InArray::const_reference::const_reference in1 = in0[i_s1+n1];
for (long n2 = 0; n2 < N2; n2++) {
out1[s2+n2] = in1[i_s2+n2]*scale;
}
}
}
cons.print<LOG_DEBUG>("Done copy");
}
template<typename InArray, typename OutArray>
void scaleAndCopyArray3d(OutArray& out, const InArray& in, typename OutArray::element scale, bool in_padded = false)
{
using boost::format;
size_t N0 = out.shape()[0],
N1 = out.shape()[1],
N2 = out.shape()[2];
if (!in_padded && (N0 < in.shape()[0] ||
N1 < in.shape()[1] ||
N2 < in.shape()[2])) {
error_helper<ErrorBadState>("Invalid copy shape in scaleAndcopyArray3d");
}
N0 = std::min(N0,in.shape()[0]);
N1 = std::min(N1,in.shape()[1]);
N2 = std::min(N2,in.shape()[2]);
ssize_t s0 = out.index_bases()[0],
s1 = out.index_bases()[1],
s2 = out.index_bases()[2];
Console::instance().print<LOG_DEBUG>(format("Copying (%d-%d, %d-%d, %d-%d)") % s0 % (s0+N0) % s1 % (s1+N1) % s2 % (s2+N2));
#pragma omp parallel for
for (size_t n0 = s0; n0 < s0+N0; n0++) {
auto out_0 = out[n0];
auto in_0 = in[n0];
for (size_t n1 = s1; n1 < s1+N1; n1++) {
auto out_1 = out_0[n1];
auto in_1 = in_0[n1];
for (size_t n2 = s2; n2 < s2+N2; n2++) {
out_1[n2] = in_1[n2]*scale;
}
}
}
Console::instance().print<LOG_DEBUG>("Done copy");
}
template<typename T>
auto generate_slice(T x[6]) {
typedef boost::multi_array_types::index_range i_range;
return boost::indices[i_range(x[0], x[1])][i_range(x[2],x[3])][i_range(x[4],x[5])];
}
template<typename A, typename RangeList>
auto slice_array(A&& a, const RangeList& rlist)
-> decltype(a[rlist])
{
auto v = a[rlist];
typedef typename std::remove_reference<A>::type A_t;
boost::array<size_t, A_t::dimensionality> ids;
size_t i = 0 ;
for (auto v: rlist.ranges_) {
ids[i] = v.get_start(a.index_bases()[i]);
i++;
}
v.reindex(ids);
return v;
}
namespace details {
typedef boost::multi_array_types::extent_range range;
template<size_t Nd>
struct make_extent {
template<typename E, typename IB, typename S>
static inline auto make(E e, IB ib, S s)
{
return make_extent<Nd-1>::make(e[range(*ib, *ib+*s)], ib+1, s+1);
}
};
template<>
struct make_extent<0> {
template<typename E, typename IB, typename S>
static inline auto make(E e, IB ib, S s) { return e; }
};
}
namespace star_index_detail {
template<typename I>
auto _make_star_indices(I indices, std::integer_sequence<size_t>) {
return indices;
}
template<typename I, size_t N0, size_t... N1>
auto _make_star_indices(I indices, std::integer_sequence<size_t,N0,N1...>) {
typedef boost::multi_array_types::index_range i_range;
return _make_star_indices(indices, std::integer_sequence<size_t,N1...>())[i_range()];
}
template<size_t N, typename I>
auto make_star_indices(I indices) {
return _make_star_indices(indices, std::make_integer_sequence<size_t,N>());
}
}
using star_index_detail::make_star_indices;
template<size_t N, typename IB, typename S, typename E = boost::multi_array_types::extent_gen>
auto make_extent(IB bases, S shape, E e = E()) {
return
details::make_extent<N>::make(
e, bases, shape
);
}
template<typename... I>
auto extent(I... i) {
std::array<size_t, sizeof...(I)> s{i...};
std::array<size_t, sizeof...(I)> b;
std::fill(b.begin(), b.end(), 0);
return make_extent<sizeof...(I)>(b.data(), s.data());
}
/**
* @brief build a new extent object from the dimensions of an existing array-like type.
*
* @return an extent
*/
template<typename A, typename E = boost::multi_array_types::extent_gen>
auto make_extent_from(A&& a, E e = E())
{
return make_extent<std::remove_reference<A>::type::dimensionality>(a.index_bases(), a.shape(), e);
}
template<typename IndexArray, typename SwapFunc>
void reorder(const IndexArray& part_idx, SwapFunc func)
{
size_t numPart = part_idx.shape()[0], i = 0;
boost::multi_array<typename IndexArray::element,1> sorter(boost::extents[numPart]);
LibLSS::copy_array(sorter, part_idx);
// Now that partices have been gathered back, we have to reorder them
// to "unsort".
while (i < numPart) {
typename IndexArray::reference swapper = sorter[i];
if (swapper == i) {
i++;
continue;
}
func(swapper, i);
std::swap(sorter[i], sorter[swapper]);
}
}
};
};
#endif