jet

jet.hpp

@@ -0,0 +1,164 @@
+#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