Forward and Reverse Mode Automatic Differenatiation in C#

using System;

namespace HigherLogics
{
    // Copyright 2020-05-24 Sandro Magi

    //---------------- Forward Mode Automatic Differentiation ---------------------

    public readonly struct Fwd
    {
        public readonly double Magnitude;
        public readonly double Derivative;

        public Fwd(double mag, double deriv)
        {
            this.Magnitude = mag;
            this.Derivative = deriv;
        }

        public Fwd Pow(int k) =>
            new Fwd(Math.Pow(Magnitude, k), k * Math.Pow(Magnitude, k - 1) * Derivative);

        public static Fwd operator +(Fwd lhs, Fwd rhs) =>
            new Fwd(lhs.Magnitude + rhs.Magnitude, lhs.Derivative + rhs.Derivative);

        public static Fwd operator *(Fwd lhs, Fwd rhs) =>
            new Fwd(lhs.Magnitude + rhs.Magnitude, lhs.Derivative * rhs.Magnitude + rhs.Derivative * lhs.Magnitude);

        public static Func<double, Fwd> Differentiate(Func<Fwd, Fwd> f) =>
            x => f(new Fwd(x, 1));

        public static Func<double, double, Fwd> DifferentiateX0(Func<Fwd, Fwd, Fwd> f) =>
            (x0, x1) => f(new Fwd(x0, 1), new Fwd(x1, 0));

        public static Func<double, double, Fwd> DifferentiateX1(Func<Fwd, Fwd, Fwd> f) =>
            (x0, x1) => f(new Fwd(x0, 0), new Fwd(x1, 1));
    }



    //----------------- Reverse Mode Automatic Differentiation -------------------

    public readonly struct Rev
    {
        public readonly double Magnitude;
        readonly Action<double> Derivative;

        public Rev(double y, Action<double> dy)
        {
            this.Magnitude = y;
            this.Derivative = dy;
        }

        public Rev Pow(int e)
        {
            var x = Magnitude;
            var k = Derivative;
            return new Rev(Math.Pow(Magnitude, e), dx => k(e * Math.Pow(x, e - 1) * dx));
        }

        public static Rev operator +(Rev lhs, Rev rhs) =>
            new Rev(lhs.Magnitude + rhs.Magnitude, dx =>
            {
                lhs.Derivative(dx);
                rhs.Derivative(dx);
            });

        public static Rev operator *(Rev lhs, Rev rhs) =>
            new Rev(lhs.Magnitude * rhs.Magnitude,
                   dx =>
                   {
                       lhs.Derivative(dx * rhs.Magnitude);
                       rhs.Derivative(dx * lhs.Magnitude);
                   });

        public static Func<double, (double, double)> Differentiate(Func<Rev, Rev> f) =>
            x =>
            {
                double dx = 1;
                var y = f(new Rev(x, dy => dx = dy));
                y.Derivative(1);
                return (y.Magnitude, dx);
            };

        public static Func<double, double, (double, double, double)> Differentiate(Func<Rev, Rev, Rev> f) =>
            (x0, x1) =>
            {
                double dx0 = 1, dx1 = 1;
                var y = f(new Rev(x0, dy => dx0 = dy), new Rev(x1, dy => dx1 = dy));
                y.Derivative(1);
                return (y.Magnitude, dx0, dx1);
            };

        public static Func<double, double, double, (double, double, double, double)> Differentiate(Func<Rev, Rev, Rev, Rev> f) =>
            (x0, x1, x2) =>
            {
                double dx0 = -1, dx1 = -1, dx2 = -1;
                var y = f(new Rev(x0, dy => dx0 = dy),
                          new Rev(x1, dy => dx1 = dy),
                          new Rev(x2, dy => dx2 = dy));
                y.Derivative(1);
                return (y.Magnitude, dx0, dx1, dx2);
            };
    }
}