23 October 2006

Better Jacobians in Computing Derivatives

And another fix to the code computing derivatives: before I was trying to compute the Jacobian of functions with only one call to the function:

```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))
```
That's efficient in the sense that it doesn't re-compute results within f. Unfortunately, it introduces n tags, where n is the dimensionality of the argument to f. Each tag labels a fundamental ``incremental value'' dx, dy, .... A particular D(...) expression is a tree which represents an expression a0 + a1*dx + a2*dx*dy + ... which potentially contains one term for each of the 2^n different combinations of the dx, dy, .... This means that, though we save on repeated invocations of f, the cost for computing the jacobian becomes exponential in the dimension of the argument because each operation (+, sin, etc) has to manipulate these trees to a depth of 2^n.

The following simple change wastes effort computing f dargs multiple times, but scales only linearly in complexity with the dimension of args. This is a huge win. I've also added a function insert_tag which inserts a new D(tag, x, one) term within the derivative-computing wrappers (but does so in way consistent with the heap property of the diff-tree). The modified functions are below:

```  let rec insert_tag tag = function
| (C(x) as d) -> D(tag, d, one)
| (D(tagd, x, dx) as d) when tag < tagd ->
D(tag, d, one)
| D(tagd, x, dx) when tagd = tag ->
raise (Failure "Cannot insert tag which is already present.")
| D(tagd, x, dx) ->
D(tagd, insert_tag tag x, dx)

let d f =
fun x ->
let tag = new_tag () in
let res = f (insert_tag tag x) in
extract_deriv tag res

let jacobian f args =
let n = Array.length args in
let jac = Array.make_matrix n n zero in
for j = 0 to Pervasives.(-) n 1 do
let tag = new_tag () in
let dargs = Array.mapi
(fun i arg -> if Pervasives.(=) i j then insert_tag tag arg else arg)
args in
let result = f dargs in
Array.iteri (fun i res ->
jac.(i).(j) <- extract_deriv tag res)
result
done;
jac
```