algorithms.hpp

Go to the documentation of this file.

// Copyright (c) 2019--present, The Simons Foundation
// This file is part of TRIQS/nda and is licensed under the Apache License, Version 2.0.
// SPDX-License-Identifier: Apache-2.0
// See LICENSE in the root of this distribution for details.

 
#pragma once
 
#include "./basic_functions.hpp"
#include "./concepts.hpp"
#include "./layout/for_each.hpp"
#include "./layout/range.hpp"
#include "./macros.hpp"
#include "./map.hpp"
#include "./traits.hpp"
 
#include <algorithm>
#include <array>
#include <cmath>
#include <concepts>
#include <cstdlib>
#include <functional>
#include <type_traits>
#include <utility>
#include <vector>
 
namespace nda {

 
  // FIXME : CHECK ORDER of the LOOP !
  template <Array A, typename F, typename R>
  auto fold(F f, A const &a, R r) {
    // cast the initial value to the return type of f to avoid narrowing
    using res_t = std::decay_t<decltype(make_regular(f(r, get_value_t<A>{})))>;
    auto res    = res_t{r};
    nda::for_each(a.shape(), [&a, &res, &f](auto &&...args) { res = f(res, a(args...)); });
    return res;
  }

  template <Array A, typename F>
  auto fold(F f, A const &a) {
    return fold(std::move(f), a, get_value_t<A>{});
  }

  template <Array A>
  bool any(A const &a) {
    static_assert(std::is_same_v<get_value_t<A>, bool>, "Error in nda::any: Value type of the array must be bool");
    return fold([](bool r, auto const &x) -> bool { return r or bool(x); }, a, false);
  }

  template <Array A>
  bool all(A const &a) {
    static_assert(std::is_same_v<get_value_t<A>, bool>, "Error in nda::all: Value type of the array must be bool");
    return fold([](bool r, auto const &x) -> bool { return r and bool(x); }, a, true);
  }

  template <Array A>
  auto max_element(A const &a) {
    return fold(
       [](auto const &x, auto const &y) {
         using std::max;
         return max(x, y);
       },
       a, get_first_element(a));
  }

  template <Array A>
  auto min_element(A const &a) {
    return fold(
       [](auto const &x, auto const &y) {
         using std::min;
         return min(x, y);
       },
       a, get_first_element(a));
  }

  template <ArrayOfRank<2> A>
  double frobenius_norm(A const &a) {
    return std::sqrt(fold(
       [](double r, auto const &x) -> double {
         auto ab = std::abs(x);
         return r + ab * ab;
       },
       a, double(0)));
  }

  template <Array A, typename Value = get_value_t<A>>
  auto sum(A const &a)
    requires(nda::Scalar<Value> or nda::Array<Value>)
  {
    if constexpr (nda::Scalar<Value>) {
      return fold(std::plus<>{}, a);
    } else { // Array<Value>
      return fold(std::plus<>{}, a, Value::zeros(get_first_element(a).shape()));
    }
  }

  template <Array A, std::integral I, size_t N>
    requires(get_rank<A> >= N and nda::Scalar<get_value_t<A>>)
  auto sum(A const &a, std::array<I, N> axes) {
    // if no axes are specified, return a copy of the input
    if constexpr (N == 0) {
      return make_regular(a);
    } else {
      // sort axes and check validity
      std::ranges::sort(axes);
      EXPECTS(std::ranges::adjacent_find(axes) == axes.end());
      EXPECTS(axes.front() >= 0 and axes.back() < get_rank<A>);
 
      constexpr int res_rank = get_rank<A> - static_cast<int>(N);
 
      if constexpr (res_rank == 0) {
        return sum(a);
      } else {
        // get the result shape and the axes that we keep (are not summed over)
        std::array<long, res_rank> keep_axes, res_shape;
        for (int i = 0; auto ax : nda::range(get_rank<A>)) {
          if (!std::ranges::binary_search(axes, ax)) {
            keep_axes[i]   = ax;
            res_shape[i++] = a.shape()[ax];
          }
        }
 
        // create the result array initialized to zero
        auto res = array<get_value_t<A>, res_rank>::zeros(res_shape);
 
        // loop over all indices of the input array and sum over the specified axes
        nda::for_each(a.shape(), [&](auto... idxs) {
          auto idx_arr = std::array{idxs...};
          std::apply([&](auto... keep) { res(idx_arr[keep]...) += a(idxs...); }, keep_axes);
        });
 
        return res;
      }
    }
  }

  template <Array A>
    requires(get_rank<A> >= 1 and nda::Scalar<get_value_t<A>>)
  auto sum(A const &a, int axis) {
    return sum(a, std::array{axis});
  }

  template <Array A, typename Value = get_value_t<A>>
  auto product(A const &a)
    requires(nda::Scalar<Value> or nda::Array<Value>)
  {
    if constexpr (nda::Scalar<Value>) {
      return fold(std::multiplies<>{}, a, get_value_t<A>{1});
    } else { // Array<Value>
      return fold(std::multiplies<>{}, a, Value::ones(get_first_element(a).shape()));
    }
  }

  template <Array A, Array B>
    requires(nda::get_rank<A> == nda::get_rank<B>)
  [[nodiscard]] constexpr auto hadamard(A &&a, B &&b) {
    return nda::map([](auto const &x, auto const &y) { return x * y; })(std::forward<A>(a), std::forward<B>(b));
  }

  template <typename T, typename U, size_t R>
  [[nodiscard]] constexpr auto hadamard(std::array<T, R> const &a, std::array<U, R> const &b) {
    return a * b;
  }

  template <typename T, typename U>
  [[nodiscard]] constexpr auto hadamard(std::vector<T> const &a, std::vector<U> const &b) {
    using TU = decltype(std::declval<T>() * std::declval<U>());
    EXPECTS(a.size() == b.size());
 
    std::vector<TU> c(a.size());
    for (auto i : range(c.size())) c[i] = a[i] * b[i];
    return c;
  }

  constexpr auto hadamard(nda::Scalar auto a, nda::Scalar auto b) { return a * b; }

 
} // namespace nda