Correction to Automatic Differentiation Code

I've uncovered a *large* inefficiency in the OCaml code I posted for Automatic Differentiation. The problem is that I don't take any care to avoid duplication of tags in a nested D(tag, x, dx) expression. That's not a problem in a computation with only one tag, but if you're trying to take the jacobian of some multi-dimensional function, you have derivatives wrt each of the arguments floating through the computation. If you duplicate tags every time you construct a new D(...) expression, then the nested length of the expression grows *really* fast. The key to avoiding duplication is to think of a given D(tag, x, dx) expression as a tree which has the heap property: all tags in the expressions x and dx are larger than tag. Then you just have to re-write the binary operations (these are the only ones which can combine D(...) expressions with different tags) to preserve the heap property. Updated code is posted below (sorry for another long post, but this change makes a big difference)! I've added the jacobian and lower_jacobian methods because I needed them to test for symplecticity (which can be expressed as the condition J^T.S.J = S, where S is the standard symplectic form [[0, I], [-I, 0]], where I is the identity matrix, and 0 is the zero matrix, and J is the jacobian of a map which is symplecitc). As before, this code is released under the GPL.

(** Code to implement differentation on arbitrary types. 

   diff.ml: Library to compute exact derivatives of arbitrary
   mathematical functions.

   Copyright (C) 2006 Will M. Farr 
   
   This program is free software; you can redistribute it and/or modify
   it under the terms of the GNU General Public License as published by
   the Free Software Foundation; either version 2 of the License, or
   (at your option) any later version.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU General Public License for more details.

   You should have received a copy of the GNU General Public License along
   with this program; if not, write to the Free Software Foundation, Inc.,
   51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.

*)

(** Input signature for [Make].*)
module type In = sig
  type t

(** The multiplicative identity *)
  val one : t

(** The additive identity *)
  val zero : t

(** Comparisons *)
  val compare : t -> t -> int

  val (+) : t -> t -> t
  val (-) : t -> t -> t
  val ( * ) : t -> t -> t
  val (/) : t -> t -> t
(** Arithmetic operations *)


  val sin : t -> t
  val cos : t -> t
  val tan : t -> t 
(** Trigonometric functions *)

  val asin : t -> t
  val acos : t -> t
  val atan : t -> t
(** Inverse trigonometric functions *)

  val sqrt : t -> t
  val log : t -> t
  val exp : t -> t
  val ( ** ) : t -> t -> t
(** Powers, exponents and logs. *)

  val print_t : out_channel -> t -> unit
end

(** Output signature for [Make]. *)
module type Out = sig 
(** Input from [In] *)
  type t 

(** Opaque tags for variables being differentiated. *)
  type tag

(** Output diff type.  A [diff] is either a constant [t] value [C(x)]
   or a tuple of [D(tag, x, dx)] where [x] and [dx] can also be
   [diff].  [D(tag, x, dx)] represents a small increment [dx] in the
   value labeled by [tag] about the value [x]. *)
  type diff = 
    | C of t 
    | D of tag * diff * diff

  val zero : diff
  val one : diff
(** Additive and multiplicative identities in [diff]. *)

(** [lift f f'] takes [f : t -> t] to [diff -> diff] using [f'] to
   compute the derivative.  For example, if we had already defined
   [(/) : diff -> diff -> diff], and the input signature didn't
   provide [log], we could define it using [let log = lift log ((/)
   one)]. *)
  val lift : (t -> t) -> (diff -> diff) -> (diff -> diff)

(** [lower f] takes a function defined on [diff] to the equivalent one
   defined on [t].  It is an error if [f (C x)] does not evaluate to
   [C y] for some [y].  The function is called [lower] because it is
   the inverse of the function [lift] which lifts [f : t -> t] to
   [diff -> diff].  *)
  val lower : (diff -> diff) -> (t -> t)

(** [lower_multi f] lowers [f] from [diff array -> diff] to [t array ->
   t].  It is an error if [f \[|(C x0); ...; (C xn)|\]] does not
   evaluate to [\[|(C y0); ...; (C yn)|\]].*)
  val lower_multi : (diff array -> diff) -> (t array -> t)

(** [lower_multi_multi f] lowers [f] from [diff array -> diff array]
   to [t array -> t array].  *)
  val lower_multi_multi : (diff array -> diff array) -> 
    (t array -> t array)

(** [lower_jacobian j] lowers the jacobian [j] from [diff array ->
   diff array array] to [t array -> t array array]. *)
  val lower_jacobian : (diff array -> diff array array) -> 
    (t array -> t array array)

(** [d f] returns the function which computes the derivative of
   [f]. *)
  val d : (diff -> diff) -> (diff -> diff)

(** [partial i f] returns the function which computes the derivative
    of [f] with respect to its [i]th argument (counting from 0). *)
  val partial : int -> (diff array -> diff) -> (diff array -> diff)

(** [jacobian f] returns the jacobian matrix with elements [m.(i).(j)
   = d f.(i)/d x.(j)]. *)
  val jacobian : (diff array -> diff array) -> 
    (diff array -> diff array array)

(** [compare d1 d2] compares only the values stored in the
   derivative---that is either the [C x] or [D(_, x, _)].  [(<)],
   [(>)], ... are defined in terms of compare, which is defined in
   terms of [compare] from [In]. *)
  val compare : diff -> diff -> int

  val (<) : diff -> diff -> bool
  val (<=) : diff -> diff -> bool
  val (>) : diff -> diff -> bool
  val (>=) : diff -> diff -> bool
  val (=) : diff -> diff -> bool
  val (<>) : diff -> diff -> bool
(** Comparison functions *)

  val (+) : diff -> diff -> diff
  val (-) : diff -> diff -> diff
  val ( * ) : diff -> diff -> diff
  val (/) : diff -> diff -> diff
(** Algebra. *)

  val cos : diff -> diff
  val sin : diff -> diff
  val tan : diff -> diff
(** Trig *) 

  val acos : diff -> diff
  val asin : diff -> diff
  val atan : diff -> diff
(** Inverse trig *)

  val sqrt : diff -> diff
  val log : diff -> diff
  val exp : diff -> diff
  val ( ** ) : diff -> diff -> diff
(** Powers, exponents and logs *)

  val print_diff : out_channel -> diff -> unit
end

module Make(I : In) : Out with type t = I.t = struct
  type t = I.t

  type tag = int

(* Terms are either constant or of the form (x + dx), with x
   represented by tag.  We are careful to maintain the invariant that
   all the tags in x and dx are larger than the tag of (x + dx).
   (That is, you can think of a D(tag, x, dx) as a tree : 

          tag 
         /           x     dx

   which satisfies the heap property tag < tagx and tag < tagdx.
*)
  type diff = 
    | C of t
    | D of tag * diff * diff

  let rec print_diff out = function 
    | C(x) -> Printf.fprintf out "C(";
        I.print_t out x;
        Printf.fprintf out ")"
    | D(tag, x, dx) -> 
        Printf.fprintf out "D(%d," tag;
        print_diff out x;
        Printf.fprintf out ", ";
        print_diff out dx;
        Printf.fprintf out ")"

(* Additive and multiplicative identities in derivatives. *)
  let zero = C(I.zero)
  let one = C(I.one)

(* Unique tags *)
  let new_tag = 
    let count = ref 0 in 
    fun () -> 
      count := !count + 1;
      !count

(* Have to define the arithmetic operators first because they are used
   in [lift] and friends.  To maintain the heap property of the tags,
   we select the smallest of tagx and tagy when we're operating on two
   D(_,_,_) objects.  We know that we can directly construct
   D(smallest, _, _), where _ and _ can be any of the sub derivatives.
   But, we can only use the sub-derivatives from larger in any direct
   construction D(larger, _, _); there is no guarantee that the tags
   in the sub-derivatives from D(smaller, _, _) are in any relation to
   larger.  This is the reason for the somewhat obfuscated code
   below. *)
  let rec (+) d1 d2 = 
    match d1, d2 with 
    | C(x), C(y) -> C(I.(+) x y)
    | C(x), D(tag, y, dy) -> 
        D(tag, d1 + y, dy)
    | D(tag, x, dx), C(y) -> 
        D(tag, x + d2, dx)
    | D(tagx, x, dx), D(tagy, y, dy) when tagx = tagy ->
        D(tagx, x + y, dx + dy)
    | D(tagx, x, dx), D(tagy, y, dy) when tagx < tagy ->
        D(tagx, x + d2, dx)
    | D(tagx, x, dx), D(tagy, y, dy) -> 
        D(tagy, d1 + y, dy)

  let rec (-) d1 d2 = 
    match d1, d2 with 
    | C(x), C(y) -> C(I.(-) x y)
    | C(x), D(tag, y, dy) -> 
        D(tag, d1 - y, zero - dy)
    | D(tag, x, dx), C(y) -> 
        D(tag, x - d2, dx)
    | D(tagx, x, dx), D(tagy, y, dy) when tagx = tagy -> 
        D(tagx, x - y, dx - dy)
    | D(tagx, x, dx), D(tagy, y, dy) when tagx < tagy -> 
        D(tagx, x - d2, dx)
    | D(tagx, x, dx), D(tagy, y, dy) -> 
        D(tagy, d1 - y, zero - dy)

  let rec ( * ) d1 d2 = 
    match d1, d2 with 
    | C(x), C(y) -> C(I.( * ) x y)
    | C(x), D(tag, y, dy) -> D(tag, d1 * y, d1 * dy)
    | D(tag, x, dx), C(y) -> D(tag, x * d2, dx * d2)
    | D(tagx, x, dx), D(tagy, y, dy) when tagx = tagy -> 
        D(tagx, x * y, x*dy + dx*y)
    | D(tagx, x, dx), D(tagy, y, dy) when tagx < tagy -> 
        D(tagx, x*d2, dx*d2)
    | D(tagx, x, dx), D(tagy, y, dy) -> 
        D(tagy, d1*y, d1*dy)

  let rec (/) d1 d2 = 
    match d1, d2 with 
    | C(x), C(y) -> C(I.(/) x y)
    | D(tag, x, dx), C(y) -> D(tag, x/d2, dx/d2)
    | C(x), D(tag, y, dy) -> D(tag, d1/y, zero - d1*dy/(y*y))
    | D(tagx, x, dx), D(tagy, y, dy) when tagx = tagy -> 
        D(tagx, x/y, dx/y - x*dy/(y*y))
    | D(tagx, x, dx), D(tagy, y, dy) when tagx < tagy -> 
        let y2 = y * y in 
        let mdyoy2 = zero - dy/y2 and 
            ooy = one/y in 
        D(tagx, x * (D(tagy, ooy, mdyoy2)),
          dx*(D(tagy, ooy, mdyoy2)))
    | D(tagx, x, dx), D(tagy, y, dy) -> 
        D(tagy, d1/y, zero - d1*dy/(y*y))

  let lift f f' = 
    let rec lf = function 
      | C(x) -> C(f x)
      | D(tag, x, dx) -> D(tag, lf x, (f' x)*dx) in 
    lf

(* Now that we have the algebra of derivatives worked out, we can
   define lift, lift2, lower and lower2. *)
  let lower f = 
    fun x -> 
      match f (C x) with 
      | C(y) -> y
      | _ -> raise (Failure "lower expects numerical output")

  let rec extract_deriv tag = function 
    | C(_) -> zero
    | D(tagx, x, dx) when tagx = tag -> 
        dx
    | D(tagx, x, dx) -> 
        D(tagx, extract_deriv tag x, extract_deriv tag dx) 
  and drop_deriv tag = function 
    | D(tagx, x, dx) when tagx = tag -> 
        x
    | D(tagx, x, dx) -> 
        D(tagx, drop_deriv tag x, drop_deriv tag dx)
    | x -> x (* Only matches C(_) *)
          
  let d f = 
    fun x -> 
      let tag = new_tag () in 
      let res = f (D(tag, x, one)) in 
      extract_deriv tag res

  let rec compare d1 d2 = 
    match d1, d2 with 
    | C(x), C(y) -> I.compare x y
    | C(x), D(_, y, _) -> compare d1 y
    | D(_, x, _), C(y) -> compare x d2
    | D(_, x, _), D(_, y, _) -> compare x y

  let rec cos = function 
    | C(x) -> C(I.cos x)
    | D(tag, x, dx) -> D(tag, cos x, zero - dx * (sin x)) and 
      sin = function 
        | C(x) -> C(I.sin x)
        | D(tag, x, dx) -> D(tag, sin x, dx * (cos x))

  let rec tan = function 
    | C(x) -> C(I.tan x)
    | D(tag, x, dx) -> 
        let tx = tan x in 
        D(tag, tx, (one + tx*tx)*dx)

  let rec sqrt = function 
    | C(x) -> C(I.sqrt x)
    | D(tag, x, dx) -> 
        let sx = sqrt x in 
        D(tag, sx, dx / (sx + sx))

  let rec acos = function 
    | C(x) -> C(I.acos x)
    | D(tag, x, dx) -> 
        D(tag, acos x, zero - dx / (sqrt (one - x * x)))

  let rec asin = function 
    | C(x) -> C(I.asin x)
    | D(tag, x, dx) -> 
        D(tag, asin x, dx / (sqrt (one - x * x)))

  let rec atan = function 
    | C(x) -> C(I.atan x)
    | D(tag, x, dx) -> 
        D(tag, atan x, dx / (one + x*x))

  let rec log = function 
    | C(x) -> C(I.log x)
    | D(tag, x, dx) -> D(tag, log x, dx/x)

  let rec exp = function 
    | C(x) -> C(I.exp x)
    | D(tag, x, dx) -> 
        let ex = exp x in 
        D(tag, ex, dx*ex)

  let rec ( ** ) d1 d2 = 
    match d1, d2 with 
    | C(x), C(y) -> C(I.( ** ) x y)
    | C(x), D(tag, y, dy) -> 
        D(tag, d1**y, d1**y * (log d1) * dy)
    | D(tag, x, dx), C(y) -> 
        D(tag, x**d2, x**(d2 - one)*d2*dx)
    | D(tagx, x, dx), D(tagy, y, dy) when tagx = tagy -> 
        D(tagx, x**y, x**(y - one) * (dx*y + x*(log x)*dy))
    | D(tagx, x, dx), D(tagy, y, dy) when tagx < tagy -> 
        D(tagx, x**y*(one + (log x)*(D(tagy, zero, dy))),
          dx*x**(y - one)*(y + (one + y*(log x))*(D(tagy, zero, dy))))
    | D(tagx, x, dx), D(tagy, y, dy) -> 
        D(tagy, x**y + y*x**(y-one)*(D(tagx, zero, dx)),
          dy*(x**y*(log x) + x**(y-one)*(one+y*(log x))*(D(tagx, zero, dx))))
                      

  let replace arr i x = 
    Array.mapi 
      (fun j y -> 
        if j <> i then y else x)
      arr

  let partial i f = 
    fun args -> 
      let x = args.(i) in 
      let one_d_f x = 
        f (replace args i x) in 
      (d one_d_f) x

  let jacobian f args = 
    let n = Array.length args in 
    let tags = Array.init n (fun _ -> new_tag ()) in 
    let dargs = Array.mapi (fun i arg -> D(tags.(i), arg, one)) args in 
    let result = f dargs in 
    Array.init n 
      (fun i -> 
        let fi = result.(i) in 
        Array.init n 
          (fun j -> 
            let tj = tags.(j) in 
            Array.fold_left 
              (fun res tag -> 
                if tag = tj then 
                  extract_deriv tag res
                else
                  drop_deriv tag res)
              fi
              tags))
   
  let c_ify arr = 
    Array.map (fun x -> (C x)) arr

  let de_c_ify arr = 
    Array.map 
      (function 
        | C(x) -> x
        | d -> 
            raise (Failure "cannot lower [|D(...); ...; D(...)|]"))
      arr

  let lower_multi f = 
    fun args -> 
      match (f (c_ify args)) with 
      | C(x) -> x
      | _ -> raise (Failure "cannot lower D(...)")

  let lower_multi_multi f = 
    fun args -> 
      de_c_ify (f (c_ify args))

  let lower_jacobian j args = 
     Array.map 
      (fun carr -> 
        de_c_ify carr)
      (j (c_ify args))

(* Define these all at once (and last) so as not to spoil the
   comparison operator namespace. *)
let (<) d1 d2 = 
  compare d1 d2 < 0 and 
    (<=) d1 d2 = 
  not (compare d1 d2 > 0) and 
    (>) d1 d2 = 
  compare d1 d2 > 0 and
    (>=) d1 d2 = 
  not (compare d1 d2 < 0) and 
    (=) d1 d2 = 
  compare d1 d2 = 0 and 
    (<>) d1 d2 = 
  not (compare d1 d2 = 0)
end

module DFloat = Make(struct 
  type t = float 

  include Pervasives

  let (+) = (+.)
  let (-) = (-.)
  let (/) = (/.)
  let ( * ) = ( *. )

  let one = 1.0
  let zero = 0.0

  let print_t out = Printf.fprintf out "%g"
end)