hexagon/jet.hpp

#ifndef JET_HPP
#define JET_HPP
#include <cmath>
#include <iostream>
#include <Eigen/Core>
using Eigen::Vector;
template<typename T, unsigned _n_vars = 1, typename diff_type = T>
struct jet {
    T value;
    constexpr static unsigned int n_vars = _n_vars;
    Eigen::Array<diff_type, n_vars, 1> deriv;
    jet() = default;
    jet(T v) : value(v), deriv(Eigen::Array<diff_type, n_vars, 1>::Zero()){}
    jet(T v, Eigen::Array<T, n_vars, 1> d) : value(v), deriv(d){}
    jet(T v, T d) requires(n_vars == 1) : value(v), deriv(d){}
    operator T()const noexcept{
        return value;
    }
        // Addition operator
    jet<T, n_vars> operator+(const jet<T, n_vars>& other) const {
        return {value + other.value, deriv + other.deriv};
    }

    // Subtraction operator
    jet<T, n_vars> operator-(const jet<T, n_vars>& other) const {
        return {value - other.value, deriv - other.deriv};
    }

    // Multiplication operator
    jet<T, n_vars> operator*(const jet<T, n_vars>& other) const {
        return {value * other.value, value * other.deriv + deriv * other.value};
    }

    // Division operator
    jet<T, n_vars> operator/(const jet<T, n_vars>& other) const {
        T denominator = other.value * other.value;
        return {value / other.value, (deriv * other.value - value * other.deriv) / denominator};
    }

    // Compound assignment operators

    // +=
    jet<T, n_vars>& operator+=(const jet<T, n_vars>& other) {
        value += other.value;
        deriv += other.deriv;
        return *this;
    }

    // -=
    jet<T, n_vars>& operator-=(const jet<T, n_vars>& other) {
        value -= other.value;
        deriv -= other.deriv;
        return *this;
    }

    // *=
    jet<T, n_vars>& operator*=(const jet<T, n_vars>& other) {
        T oldValue = value;
        value *= other.value;
        deriv = oldValue * other.deriv + deriv * other.value;
        return *this;
    }

    // /=
    jet<T, n_vars>& operator/=(const jet<T, n_vars>& other) {
        T denominator = other.value * other.value;
        T oldValue = value;
        value /= other.value;
        deriv = (deriv * other.value - oldValue * other.deriv) / denominator;
        return *this;
    }

    template<std::convertible_to<T> R>
    jet<T, n_vars> operator+(const R& scalar) const {
        return {value + scalar, deriv};
    }

    // Subtraction operator with T value
    template<std::convertible_to<T> R>
    jet<T, n_vars> operator-(const R& scalar) const {
        return {value - scalar, deriv};
    }

    // Multiplication operator with T value
    template<std::convertible_to<T> R>
    jet<T, n_vars> operator*(const R& scalar) const {
        return {value * scalar, deriv * scalar};
    }

    // Division operator with T value
    template<std::convertible_to<T> R>
    jet<T, n_vars> operator/(const R& scalar) const {
        return {value / scalar, deriv / scalar};
    }

    // Compound assignment operators with T values

    // +=
    template<std::convertible_to<T> R>
    jet<T, n_vars>& operator+=(const R& scalar) {
        value += scalar;
        return *this;
    }

    // -=
    template<std::convertible_to<T> R>
    jet<T, n_vars>& operator-=(const R& scalar) {
        value -= scalar;
        return *this;
    }

    // *=
    template<std::convertible_to<T> R>
    jet<T, n_vars>& operator*=(const R& scalar) {
        value *= scalar;
        deriv *= scalar;
        return *this;
    }

    // /=
    template<std::convertible_to<T> R>
    jet<T, n_vars>& operator/=(const R& scalar) {
        value /= scalar;
        deriv /= scalar;
        return *this;
    }
    template<typename stream_t>
    friend stream_t& operator<<(stream_t& str, const jet<T, n_vars, diff_type>& j){
        str << '(' << j.value << ", d/dx = " << j.deriv.transpose() << ')';
        return str;
    }
};
template<typename T, unsigned N>
jet<T, N> sin(const jet<T, N>& x) {
    using std::sin;
    using std::cos;
    return {sin(x.value), x.deriv * cos(x.value)};
}

template<typename T, unsigned N>
jet<T, N> cos(const jet<T, N>& x) {
    using std::sin;
    using std::cos;
    return {cos(x.value), -x.deriv * sin(x.value)};
}

template<typename T, unsigned N>
jet<T, N> exp(const jet<T, N>& x) {
    using std::exp;
    T expVal = exp(x.value);
    return {expVal, x.deriv * expVal};
}

template<typename T, unsigned N>
jet<T, N> log(const jet<T, N>& x) {
    using std::log;
    return {log(x.value), x.deriv / x.value};
}
template<typename T, unsigned N>
jet<T, N> sqrt(const jet<T, N>& x) {
    T sqrtVal = sqrt(x.value);
    return {sqrtVal, x.deriv / (T(2) * sqrtVal)};
}
#endif