arithmetic.hpp

Go to the documentation of this file.

// Copyright (c) 2019-2024 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: Thomas Hahn, Olivier Parcollet, Nils Wentzell
 
#pragma once
 
#include "./concepts.hpp"
#include "./declarations.hpp"
#include "./linalg/matmul.hpp"
#include "./linalg/det_and_inverse.hpp"
#include "./macros.hpp"
#include "./stdutil/complex.hpp"
#include "./traits.hpp"
 
#include <functional>
#include <type_traits>
#include <utility>
 
#ifdef NDA_ENFORCE_BOUNDCHECK
#include "./exceptions.hpp"
#endif // NDA_ENFORCE_BOUNDCHECK
 
namespace nda {
 
  template <char OP, Array A>
  struct expr_unary {
    static_assert(OP == '-', "Error in nda::expr_unary: Only negation is supported");
 
    A a;
 
    template <typename... Args>
    auto operator()(Args &&...args) const {
      return -a(std::forward<Args>(args)...);
    }
 
    [[nodiscard]] constexpr auto shape() const { return a.shape(); }
 
    [[nodiscard]] constexpr long size() const { return a.size(); }
  };
 
  template <char OP, ArrayOrScalar L, ArrayOrScalar R>
  struct expr {
    L l;
 
    R r;
 
    using L_t = std::decay_t<L>;
 
    using R_t = std::decay_t<R>;
 
    // FIXME : we should use is_scalar_for_v but the trait needs work to accommodate scalar L or R
    static constexpr bool l_is_scalar = nda::is_scalar_v<L>;
 
    static constexpr bool r_is_scalar = nda::is_scalar_v<R>;
 
    static constexpr char algebra = (l_is_scalar ? get_algebra<R> : get_algebra<L>);
 
    static constexpr layout_info_t compute_layout_info() {
      if (l_is_scalar) return (algebra == 'A' ? get_layout_info<R> : layout_info_t{}); // 1 as an array has all flags, it is just 1
      if (r_is_scalar) return (algebra == 'A' ? get_layout_info<L> : layout_info_t{}); // 1 as a matrix does not, as it is diagonal only.
      return get_layout_info<R> & get_layout_info<L>;                                  // default case. Take the logical and of all flags
    }
 
    [[nodiscard]] constexpr decltype(auto) shape() const {
      if constexpr (l_is_scalar) {
        return r.shape();
      } else if constexpr (r_is_scalar) {
        return l.shape();
      } else {
        EXPECTS(l.shape() == r.shape());
        return l.shape();
      }
    }
 
    [[nodiscard]] constexpr long size() const {
      if constexpr (l_is_scalar) {
        return r.size();
      } else if constexpr (r_is_scalar) {
        return l.size();
      } else {
        EXPECTS(l.size() == r.size());
        return l.size();
      }
    }
 
    template <typename... Args>
    auto operator()(Args const &...args) const {
      // addition
      if constexpr (OP == '+') {
        if constexpr (l_is_scalar) {
          // lhs is a scalar
          if constexpr (algebra == 'M')
            // rhs is a matrix
            return (std::equal_to{}(args...) ? l + r(args...) : r(args...));
          else
            // rhs is an array
            return l + r(args...);
        } else if constexpr (r_is_scalar) {
          // rhs is a scalar
          if constexpr (algebra == 'M')
            // lhs is a matrix
            return (std::equal_to{}(args...) ? l(args...) + r : l(args...));
          else
            // lhs is an array
            return l(args...) + r;
        } else
          // both are arrays or matrices
          return l(args...) + r(args...);
      }
 
      // subtraction
      if constexpr (OP == '-') {
        if constexpr (l_is_scalar) {
          // lhs is a scalar
          if constexpr (algebra == 'M')
            // rhs is a matrix
            return (std::equal_to{}(args...) ? l - r(args...) : -r(args...));
          else
            // rhs is an array
            return l - r(args...);
        } else if constexpr (r_is_scalar) {
          // rhs is a scalar
          if constexpr (algebra == 'M')
            // lhs is a matrix
            return (std::equal_to{}(args...) ? l(args...) - r : l(args...));
          else
            // lhs is an array
            return l(args...) - r;
        } else
          // both are arrays or matrices
          return l(args...) - r(args...);
      }
 
      // multiplication
      if constexpr (OP == '*') {
        if constexpr (l_is_scalar)
          // lhs is a scalar
          return l * r(args...);
        else if constexpr (r_is_scalar)
          // rhs is a scalar
          return l(args...) * r;
        else {
          // both are arrays (matrix product is not supported here)
          static_assert(algebra != 'M', "Error in nda::expr: Matrix algebra not supported");
          return l(args...) * r(args...);
        }
      }
 
      // division
      if constexpr (OP == '/') {
        if constexpr (l_is_scalar) {
          // lhs is a scalar
          static_assert(algebra != 'M', "Error in nda::expr: Matrix algebra not supported");
          return l / r(args...);
        } else if constexpr (r_is_scalar)
          // rhs is a scalar
          return l(args...) / r;
        else {
          // both are arrays (matrix division is not supported here)
          static_assert(algebra != 'M', "Error in nda::expr: Matrix algebra not supported");
          return l(args...) / r(args...);
        }
      }
    }
 
    template <typename Arg>
    auto operator[](Arg &&arg) const {
      static_assert(get_rank<expr> == 1, "Error in nda::expr: Subscript operator only available for expressions of rank 1");
      return operator()(std::forward<Arg>(arg));
    }
  };
 
  template <Array A>
  expr_unary<'-', A> operator-(A &&a) {
    return {std::forward<A>(a)};
  }
 
  template <Array L, Array R>
  Array auto operator+(L &&l, R &&r) {
    static_assert(get_rank<L> == get_rank<R>, "Error in lazy nda::operator+: Rank mismatch");
    return expr<'+', L, R>{std::forward<L>(l), std::forward<R>(r)};
  }
 
  template <Array A, Scalar S>
  Array auto operator+(A &&a, S &&s) { // NOLINT (S&& is mandatory for proper concept Array <: typename to work)
    return expr<'+', A, std::decay_t<S>>{std::forward<A>(a), s};
  }
 
  template <Scalar S, Array A>
  Array auto operator+(S &&s, A &&a) { // NOLINT (S&& is mandatory for proper concept Array <: typename to work)
    return expr<'+', std::decay_t<S>, A>{s, std::forward<A>(a)};
  }
 
  template <Array L, Array R>
  Array auto operator-(L &&l, R &&r) {
    static_assert(get_rank<L> == get_rank<R>, "Error in lazy nda::operator-: Rank mismatch");
    return expr<'-', L, R>{std::forward<L>(l), std::forward<R>(r)};
  }
 
  template <Array A, Scalar S>
  Array auto operator-(A &&a, S &&s) { // NOLINT (S&& is mandatory for proper concept Array <: typename to work)
    return expr<'-', A, std::decay_t<S>>{std::forward<A>(a), s};
  }
 
  template <Scalar S, Array A>
  Array auto operator-(S &&s, A &&a) { // NOLINT (S&& is mandatory for proper concept Array <: typename to work)
    return expr<'-', std::decay_t<S>, A>{s, std::forward<A>(a)};
  }
 
  template <Array L, Array R>
  auto operator*(L &&l, R &&r) {
    // allowed algebras: A * A or M * M or M * V
    static constexpr char l_algebra = get_algebra<L>;
    static constexpr char r_algebra = get_algebra<R>;
    static_assert(l_algebra != 'V', "Error in nda::operator*: Can not multiply vector by an array or a matrix");
 
    // two arrays: A * A
    if constexpr (l_algebra == 'A') {
      static_assert(r_algebra == 'A', "Error in nda::operator*: Both types need to be arrays");
      static_assert(get_rank<L> == get_rank<R>, "Error in nda::operator*: Rank mismatch");
#ifdef NDA_ENFORCE_BOUNDCHECK
      if (l.shape() != r.shape()) NDA_RUNTIME_ERROR << "Error in nda::operator*: Dimension mismatch: " << l.shape() << " != " << r.shape();
#endif
      return expr<'*', L, R>{std::forward<L>(l), std::forward<R>(r)};
    }
 
    // two matrices: M * M
    if constexpr (l_algebra == 'M') {
      static_assert(r_algebra != 'A', "Error in nda::operator*: Can not multiply a matrix by an array");
      if constexpr (r_algebra == 'M')
        // matrix * matrix
        return matmul(std::forward<L>(l), std::forward<R>(r));
      else
        // matrix * vector
        return matvecmul(std::forward<L>(l), std::forward<R>(r));
    }
  }
 
  template <Array A, Scalar S>
  Array auto operator*(A &&a, S &&s) { // NOLINT (S&& is mandatory for proper concept Array <: typename to work)
    return expr<'*', A, std::decay_t<S>>{std::forward<A>(a), s};
  }
 
  template <Scalar S, Array A>
  Array auto operator*(S &&s, A &&a) { // NOLINT (S&& is mandatory for proper concept Array <: typename to work)
    return expr<'*', std::decay_t<S>, A>{s, std::forward<A>(a)};
  }
 
  template <Array L, Array R>
  Array auto operator/(L &&l, R &&r) {
    // allowed algebras: A / A or M / M
    static constexpr char l_algebra = get_algebra<L>;
    static constexpr char r_algebra = get_algebra<R>;
    static_assert(l_algebra != 'V', "Error in nda::operator/: Can not divide a vector by an array or a matrix");
 
    // two arrays: A / A
    if constexpr (l_algebra == 'A') {
      static_assert(r_algebra == 'A', "Error in nda::operator/: Both types need to be arrays");
      static_assert(get_rank<L> == get_rank<R>, "Error in nda::operator/: Rank mismatch");
#ifdef NDA_ENFORCE_BOUNDCHECK
      if (l.shape() != r.shape()) NDA_RUNTIME_ERROR << "Error in nda::operator/: Dimension mismatch: " << l.shape() << " != " << r.shape();
#endif
      return expr<'/', L, R>{std::forward<L>(l), std::forward<R>(r)};
    }
 
    // two matrices: M / M
    if constexpr (l_algebra == 'M') {
      static_assert(r_algebra == 'M', "Error in nda::operator*: Can not divide a matrix by an array/vector");
      return std::forward<L>(l) * inverse(matrix<get_value_t<R>>{std::forward<R>(r)});
    }
  }
 
  template <Array A, Scalar S>
  Array auto operator/(A &&a, S &&s) { // NOLINT (S&& is mandatory for proper concept Array <: typename to work)
    return expr<'/', A, std::decay_t<S>>{std::forward<A>(a), s};
  }
 
  template <Scalar S, Array A>
  Array auto operator/(S &&s, A &&a) { // NOLINT (S&& is mandatory for proper concept Array <: typename to work)
    static constexpr char algebra = get_algebra<A>;
    if constexpr (algebra == 'M')
      return s * inverse(matrix<get_value_t<A>>{std::forward<A>(a)});
    else
      return expr<'/', std::decay_t<S>, A>{s, std::forward<A>(a)};
  }
 
} // namespace nda