{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- | Column-style Jacobian construction via reverse-mode AD.
--
-- Build ∂f/∂x for a list-valued function by seeding one input component at a
-- time through 'Circuit.Diff.Circuit.Diff' and reading back the pullbacks.
-- This is the reusable pattern underlying the Daisyworld / DubinsChase oracle
-- style: write the physics once, polymorphic in a carrier with 'Lit', then
-- instantiate the Jacobian helper at 'Double'.
module Circuit.Diff.Jacobian
  ( jacobianFrom,
  )
where

import Circuit.Diff.Carrier (Lit (..))
import Circuit.Diff.Circuit (Diff (..), runDiff)
import Prelude

-- | Jacobian of @f : R^n -> R^m@ at a point, built column-by-column.
--
-- The supplied function must already be instantiated at the 'Diff' carrier;
-- callers typically wrap a polymorphic model:
--
-- @
-- jacobianFrom s0 (\s -> daisyRHS (lit l) s)
-- @
--
-- Returns an @m × n@ list of lists: outer index is output component, inner
-- index is input component.
jacobianFrom ::
  forall tag.
  [Double] ->
  ([Diff tag Double Double] -> [Diff tag Double Double]) ->
  [[Double]]
jacobianFrom :: forall {k} (tag :: k).
[Double]
-> ([Diff tag Double Double] -> [Diff tag Double Double])
-> [[Double]]
jacobianFrom [Double]
s0 [Diff tag Double Double] -> [Diff tag Double Double]
f =
  let n :: Int
n = [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
s0
      m :: Int
m = [Diff tag Double Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Diff tag Double Double] -> [Diff tag Double Double]
f ((Double -> Diff tag Double Double)
-> [Double] -> [Diff tag Double Double]
forall a b. (a -> b) -> [a] -> [b]
map Double -> Diff tag Double Double
forall a. Lit a => Double -> a
lit [Double]
s0))
      col :: Int -> [Double]
col Int
k =
        let seed :: Diff tag Double Double
seed = (Double -> (Double, Double -> Double)) -> Diff tag Double Double
forall k (p :: k) a b. (a -> (b, b -> a)) -> Diff p a b
Diff (,Double -> Double
forall a. a -> a
id) :: Diff tag Double Double
            s :: [Diff tag Double Double]
s = [if Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
k then Diff tag Double Double
seed else Double -> Diff tag Double Double
forall a. Lit a => Double -> a
lit ([Double]
s0 [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
i) | Int
i <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
            outs :: [Diff tag Double Double]
outs = [Diff tag Double Double] -> [Diff tag Double Double]
f [Diff tag Double Double]
s
         in [let (Double
_, Double -> Double
pb) = Diff tag Double Double -> Double -> (Double, Double -> Double)
forall k (p :: k) a b. Diff p a b -> a -> (b, b -> a)
runDiff ([Diff tag Double Double]
outs [Diff tag Double Double] -> Int -> Diff tag Double Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
i) ([Double]
s0 [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
k) in Double -> Double
pb Double
1 | Int
i <- [Int
0 .. Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
      cols :: [[Double]]
cols = (Int -> [Double]) -> [Int] -> [[Double]]
forall a b. (a -> b) -> [a] -> [b]
map Int -> [Double]
col [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
   in [[([[Double]]
cols [[Double]] -> Int -> [Double]
forall a. HasCallStack => [a] -> Int -> a
!! Int
j) [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
i | Int
j <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]] | Int
i <- [Int
0 .. Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]