Automatic Differentiation for PDEs

[Sophia Gold]

Hi All,

Following up on a post from last year about a library I wrote for forward mode automatic differentiation, I'd like to now share a version that works on functions of multiple variables: https://github.com/Sophia-Gold/Madhava-v2. 

It's greatly stripped down and at this point lacks polynomial division and composition, although I do plan on adding those functions later. I'm not aware of similar AD packages in Clojure, but they've become quite popular in Haskell (Edward Kmett's being impressively comprehensive) and I'm told OCaml and F# as well. However, this version takes a different approach from all the Haskell ones I've seen that are inspired by the work of Jerzy Karczmarczuk. They use functions as primitives whereas I use streams of data, generating several orders of partials at once and storing them in hash-maps. I believe this should make it faster for certain applications, perhaps nonlinear systems requiring higher order derivatives such as fluid dynamics simulations, although I'm embarrassingly lacking in education in these areas.

I'd love for any of you to play with it if you find it useful for your own applications, although I would maybe recommend waiting a bit to clone it until I do some performance tuning. On that note, I'd love any advice with regards to efficiency. Possible optimizations that already come to mind are swapping out Clojure vectors with Mikera's vectorz-clj library, although I wonder if any gains there would reach a bottleneck in converting from the list output of "for" (oddly enough, I don't get any speedup from using vector-of), as well as using Zachary Tellman's primitive-math.

Best,

Sophia

[Fu Yong Quah]

Have been playing with `diff`, and I have a few questions:

1. Why use an atom (I simply assumed that seeing that you use swap!) rather than a return value?
2. When using the function diff, given

(def m (atom []))

(diff [[3 7]] m 3)

@m

=> [[[3 7]] {1 [[21 6]]} {1 {1 [[126 5]]}} {1 {1 {1 [[630 4]]}}}]

Why not return [[3 7] [21 6] [126 5] [630 4]] instead?

3. Some simply documentation on usage would be helpful :)

[Sophia Gold]

Thanks for the comments!

1. Due to the nature of how partial derivatives of several orders "fan out," this means they're generated out of order. I don't think it makes sense to have to sort through them in the output to find the ones you need, especially considering how in practical usage they can easily number into the six figure range. And regardless, they'd still need to be tagged with keys in order to tell them apart. Finally, the speed difference in the approach you suggest is negligible and actually much slower when you want to do something with the data like calculate a linear transformation.
2. Are you suggesting I combine four partial derivatives into one polynomial expression? That doesn't make any sense to me. 
3. Writing quality documentation is a high priority of mine. Would you mind pointing out anything you find confusing or missing in the existing documentation so I can improve it?

Thanks,

Sophia

[Sophia Gold]

Edit: that should read *three* figure range. I can't imagine any application requiring hundreds of thousands of partials :)

[Fu Yong Quah]

1. I was thinking something along the lines of:

(defn new-partial-diff [p idx]
(vec
(for [expr p
:let [v (get expr idx)]
:when (not (zero? v))]
(-> expr
(update 0 * v)
(update idx dec)))))


(defn new-diff [p order]
(letfn [(diff-vars [p idx]
(map #(vector (conj idx %)
(new-partial-diff p %))
(range 1 (count (first p)))))
(diff-loop [n outer-acc polynomials]
(reduce (fn [acc [dervied-orders partial-derive]]
(if (< n order)
(let [next-ps (diff-vars partial-derive
(update dervied-orders 0 inc))]
(diff-loop (inc n)
(reduce (fn [x [y z]] (assoc-in x y z))
acc next-ps)
next-ps))
(assoc-in acc dervied-orders partial-derive)))
outer-acc polynomials))]
(diff-loop 0 {0 p} [[[0] p]])))


2, 3: My bad - I didn't realize that there's a usage guide under README.

On benchmarking, I believe there might be a slight error:

When using the original code:

Evaluation count : 490692 in 6 samples of 81782 calls.

Execution time mean : 1.235797 µs

Execution time std-deviation : 32.931879 ns

Execution time lower quantile : 1.198058 µs ( 2.5%)

Execution time upper quantile : 1.277707 µs (97.5%)

Overhead used : 11.260916 ns

When I changed

(map #(diff-loop (inc n) (diff-vars (first %) m (update (second %) 0 inc))) p)

to

(dorun (map #(diff-loop (inc n) (diff-vars (first %) m (update (second
%) 0 inc))) p))


Evaluation count : 1092 in 6 samples of 182 calls.
Execution time mean : 563.245309 µs
Execution time std-deviation : 18.833803 µs
Execution time lower quantile : 540.161599 µs ( 2.5%)
Execution time upper quantile : 582.791399 µs (97.5%)
Overhead used : 11.260916 ns

> --
> You received this message because you are subscribed to the Google Groups
> "Numerical Clojure" group.
> To unsubscribe from this group and stop receiving emails from it, send an
> email to numerical-cloj...@googlegroups.com.
> For more options, visit https://groups.google.com/d/optout.



--
Fu Yong Quah

[Fu Yong Quah]

I assumed that you intended it to be lazy, hence the reason you used
`map`. If that wasn't the case, you might want to consider

(doseq [x p]
(diff-loop (inc n) (diff-vars (first x) m (update (second x) 0 inc))))

It has a couple of advantages:
1. It returns a single nil
2. It doesn't implicitly create a function, which removes a (probably
insignificant) overhead.

--
Fu Yong Quah

[Sophia Gold]

Thanks for the advice! 

I assumed wrapping the function call in doall would make sure it was realized for benchmarking purposes and even checked with some folks on Slack about that, but it seems to not to be the case. The ~.5ms time does seem much more realistic for that kind of input.

I like your version that returns the output, but given it runs in roughly the same time I don't see the advantage of going that route. Much of my thinking with using hash-maps was to essentially automatically memoize all the partials of a given function for applications that involve multiple operations on them, for example calculating surface normals. 

And I actually don't need laziness here. One reason for using map initially was for parallelization with r/map, which turned out to not be worth it. I could use mapv, but I like your version with doseq better. The speed difference is negligible, but the single nil is certainly preferable :)