Arithmetic.hpp

#pragma once
 
#include <stratax/core/Concepts.hpp>
#include <stratax/core/Exceptions.hpp>
#include <stratax/core/validation/Validation.hpp>
 
#include <functional>

template<Array A>
void require_same_arithmetic_shape(const A& lhs, const A& rhs)
{
    stratax::core::validation::require_same_shape(
        lhs,
        rhs,
        "Arithmetic operands must have the same shape.");
}

template<Array A, typename Op>
A binary_op(const A& lhs, const A& rhs, Op op, bool check_zero_divisor = false)
{
    require_same_arithmetic_shape(lhs, rhs);
 
    A result(lhs.shape());
 
    auto it1 = lhs.begin();
    auto it2 = rhs.begin();
    auto it3 = result.begin();
 
    for (; it1 != lhs.end(); ++it1, ++it2, ++it3)
    {
        if (check_zero_divisor && *it2 == typename A::value_type{})
        {
            throw Exceptions::ZeroDivisionError("Array division divisor element cannot be zero.");
        }
 
        *it3 = op(*it1, *it2);
    }
 
    return result;
}

template<Array A, Numeric Scalar, typename Op>
A binary_scalar_op(const A& lhs, const Scalar& rhs, Op op, bool check_zero_divisor = false)
{
    A result(lhs.shape());
 
    auto out = result.begin();
 
    if (check_zero_divisor && rhs == Scalar{})
    {
        throw Exceptions::ZeroDivisionError("Array division scalar divisor cannot be zero.");
    }
 
    for (auto it = lhs.begin(); it != lhs.end(); ++it, ++out)
    {
        *out = op(*it, rhs);
    }
 
    return result;
}

template<Numeric Scalar, Array A, typename Op>
A binary_scalar_op(const Scalar& lhs, const A& rhs, Op op, bool check_zero_divisor = false)
{
    A result(rhs.shape());
 
    auto out = result.begin();
 
    for (auto it = rhs.begin(); it != rhs.end(); ++it, ++out)
    {
        if (check_zero_divisor && *it == typename A::value_type{})
        {
            throw Exceptions::ZeroDivisionError("Scalar division divisor element cannot be zero.");
        }
 
        *out = op(lhs, *it);
    }
 
    return result;
}

template<Array A>
A operator+(const A& lhs, const A& rhs)
{
    return binary_op(lhs, rhs, std::plus<>{});
}

template<Array A>
A operator-(const A& lhs, const A& rhs)
{
    return binary_op(lhs, rhs, std::minus<>{});
}

template<Array A>
A operator*(const A& lhs, const A& rhs)
{
    return binary_op(lhs, rhs, std::multiplies<>{});
}

template<Array A>
A operator/(const A& lhs, const A& rhs)
{
    return binary_op(lhs, rhs, std::divides<>{}, true);
}

template<Array A, Numeric Scalar>
A operator+(const A& lhs, const Scalar& rhs)
{
    return binary_scalar_op(lhs, rhs, std::plus<>{});
}

template<Array A, Numeric Scalar>
A operator-(const A& lhs, const Scalar& rhs)
{
    return binary_scalar_op(lhs, rhs, std::minus<>{});
}

template<Array A, Numeric Scalar>
A operator*(const A& lhs, const Scalar& rhs)
{
    return binary_scalar_op(lhs, rhs, std::multiplies<>{});
}

template<Array A, Numeric Scalar>
A operator/(const A& lhs, const Scalar& rhs)
{
    return binary_scalar_op(lhs, rhs, std::divides<>{}, true);
}

template<Numeric Scalar, Array A>
A operator+(const Scalar& lhs, const A& rhs)
{
    return rhs + lhs;
}

template<Numeric Scalar, Array A>
A operator-(const Scalar& lhs, const A& rhs)
{
    return binary_scalar_op(lhs, rhs, std::minus<>{});
}

template<Numeric Scalar, Array A>
A operator*(const Scalar& lhs, const A& rhs)
{
    return rhs * lhs;
}

template<Numeric Scalar, Array A>
A operator/(const Scalar& lhs, const A& rhs)
{
    return binary_scalar_op(lhs, rhs, std::divides<>{}, true);
}

template<Array A>
A& operator+=(A& lhs, const A& rhs)
{
    lhs = lhs + rhs;
    return lhs;
}

template<Array A>
A& operator-=(A& lhs, const A& rhs)
{
    lhs = lhs - rhs;
    return lhs;
}

template<Array A>
A& operator*=(A& lhs, const A& rhs)
{
    lhs = lhs * rhs;
    return lhs;
}

template<Array A>
A& operator/=(A& lhs, const A& rhs)
{
    lhs = lhs / rhs;
    return lhs;
}

template<Array A, Numeric Scalar>
A& operator+=(A& lhs, const Scalar& rhs)
{
    lhs = lhs + rhs;
    return lhs;
}

template<Array A, Numeric Scalar>
A& operator-=(A& lhs, const Scalar& rhs)
{
    lhs = lhs - rhs;
    return lhs;
}

template<Array A, Numeric Scalar>
A& operator*=(A& lhs, const Scalar& rhs)
{
    lhs = lhs * rhs;
    return lhs;
}

template<Array A, Numeric Scalar>
A& operator/=(A& lhs, const Scalar& rhs)
{
    lhs = lhs / rhs;
    return lhs;
}

template<Array A>
A operator-(const A& arr)
{
    return arr * typename A::value_type{-1};
}

template<Array A>
A operator+(const A& arr)
{
    return arr;
}