array.hpp

Go to the documentation of this file.

// Copyright (c) 2018 Commissariat à l'énergie atomique et aux énergies alternatives (CEA)
// Copyright (c) 2018 Centre national de la recherche scientifique (CNRS)
// Copyright (c) 2018-2023 Simons Foundation
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0.txt
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// Authors: Olivier Parcollet, Nils Wentzell
 
#ifndef STDUTILS_ARRAY_H
#define STDUTILS_ARRAY_H
 
#include <array>
#include <ostream>
#include <sstream>
#include <string>
#include <type_traits>
#include <vector>
#include <utility>
 
namespace std { // has to be in the right namespace for ADL
 
  template <typename T, size_t R>
  std::ostream &operator<<(std::ostream &out, std::array<T, R> const &a) {
    return out << to_string(a);
  }
 
  template <typename T, size_t R>
  std::string to_string(std::array<T, R> const &a) {
    std::stringstream fs;
    fs << "(";
    for (int i = 0; i < R; ++i) fs << (i == 0 ? "" : " ") << a[i];
    fs << ")";
    return fs.str();
  }
 
  template <typename T, size_t R>
  constexpr std::array<T, R> operator+(std::array<T, R> const &lhs, std::array<T, R> const &rhs) {
    std::array<T, R> res;
    for (int i = 0; i < R; ++i) res[i] = lhs[i] + rhs[i];
    return res;
  }
 
  template <typename T, size_t R>
  constexpr std::array<T, R> operator-(std::array<T, R> const &lhs, std::array<T, R> const &rhs) {
    std::array<T, R> res;
    for (int i = 0; i < R; ++i) res[i] = lhs[i] - rhs[i];
    return res;
  }
 
  template <typename T, size_t R>
  constexpr std::array<T, R> operator*(std::array<T, R> const &lhs, std::array<T, R> const &rhs) {
    std::array<T, R> res;
    for (int i = 0; i < R; ++i) res[i] = lhs[i] * rhs[i];
    return res;
  }
 
  template <typename T, size_t R>
  constexpr std::array<T, R> operator-(std::array<T, R> const &a) {
    std::array<T, R> res;
    for (int i = 0; i < R; ++i) res[i] = -a[i];
    return res;
  }
 
  template <typename T, size_t R>
  constexpr std::array<T, R> operator*(T s, std::array<T, R> const &a) {
    std::array<T, R> res;
    for (int i = 0; i < R; ++i) res[i] = s * a[i];
    return res;
  }
 
} // namespace std
 
namespace nda::stdutil {
 
  template <size_t R, typename T>
  constexpr std::array<T, R> make_initialized_array(T v) {
    return [&v]<size_t... Is>(std::index_sequence<Is...>) {
      return std::array<T, R>{(Is, v)...}; // NOLINT (always v, just a trick to have the pack)
    }(std::make_index_sequence<R>{});
  }
 
  template <typename T, typename U, size_t R>
  constexpr std::array<T, R> make_std_array(std::array<U, R> const &a) {
    static_assert(std::is_constructible_v<T, U>, "Error in nda::stdutil::make_std_array: T must be constructible from U");
    std::array<T, R> result = make_initialized_array<R>(T{});
    for (int u = 0; u < R; ++u) result[u] = a[u];
    return result;
  }
 
  template <typename T, size_t R>
  constexpr std::vector<T> to_vector(std::array<T, R> const &a) {
    std::vector<T> V(R);
    for (int i = 0; i < R; ++i) V[i] = a[i];
    return V;
  }
 
  template <typename T, auto R, typename U>
  constexpr std::array<T, R + 1> append(std::array<T, R> const &a, U const &x) {
    std::array<T, R + 1> res;
    for (int i = 0; i < R; ++i) res[i] = a[i];
    res[R] = x;
    return res;
  }
 
  template <typename T, typename U, size_t R>
  constexpr std::array<T, R + 1> front_append(std::array<T, R> const &a, U const &x) {
    std::array<T, R + 1> res;
    res[0] = x;
    for (int i = 0; i < R; ++i) res[i + 1] = a[i];
    return res;
  }
 
  template <int N, typename T, size_t R>
  constexpr std::array<T, R - N> mpop(std::array<T, R> const &a) {
    std::array<T, R - N> res;
    for (int i = 0; i < R - N; ++i) res[i] = a[i];
    return res;
  }
 
  template <typename T, size_t R>
  constexpr std::array<T, R - 1> pop(std::array<T, R> const &a) {
    return mpop<1>(a);
  }
 
  template <int N, typename T, size_t R>
  constexpr std::array<T, R - N> front_mpop(std::array<T, R> const &a) {
    std::array<T, R - N> res;
    for (int i = N; i < R; ++i) res[i - N] = a[i];
    return res;
  }
 
  template <typename T, size_t R>
  constexpr std::array<T, R - 1> front_pop(std::array<T, R> const &a) {
    return front_mpop<1>(a);
  }
 
  template <typename T, size_t R1, size_t R2>
  constexpr std::array<T, R1 + R2> join(std::array<T, R1> const &a1, std::array<T, R2> const &a2) {
    std::array<T, R1 + R2> res;
    for (int i = 0; i < R1; ++i) res[i] = a1[i];
    for (int i = 0; i < R2; ++i) res[R1 + i] = a2[i];
    return res;
  }
 
  template <typename T, size_t R>
  constexpr auto sum(std::array<T, R> const &a) {
    if constexpr (R == 0)
      return T{};
    else {
      auto res = a[0];
      for (int i = 1; i < R; ++i) res += a[i];
      return res;
    }
  }
 
  template <typename T, size_t R>
  constexpr auto product(std::array<T, R> const &a) {
    static_assert(R > 0, "Error in nda::stdutil::product: Only defined for R > 0");
    auto res = a[0];
    for (int i = 1; i < R; ++i) res *= a[i];
    return res;
  }
 
  template <typename T, typename U, size_t R>
  constexpr auto dot_product(std::array<T, R> const &a1, std::array<U, R> const &a2) {
    if constexpr (R == 0)
      return T{};
    else {
      auto res = a1[0] * a2[0];
      for (int i = 1; i < R; ++i) res += a1[i] * a2[i];
      return res;
    }
  }
 
} // namespace nda::stdutil
 
#endif // STDUTILS_ARRAY_H