{-# LANGUAGE RebindableSyntax #-}

-- | Metric-aware transpose for differentiable arrows.
--
-- A metric @g@ turns the bare dagger (transpose) of a 'Diff into the true
-- adjoint @g_a^-1 . J^T . g_b@.  The same combinator also supports
-- optimizer-style preconditioning, where @g@ is a diagonal metric
-- @diag(sqrt v + eps)@ and only the lowering half is used.
module Circuit.Diff.Metric
  ( -- * Adjoint with respect to domain/codomain metrics
    adjointWith,

    -- * Forward-only metric application
    raiseWith,
    lowerWith,

    -- * Common metrics
    diagonalMetric,
    euclideanMetric,
  )
where

import Circuit.Diff (Diff (..), Diff', runDiff)
import NumHask.Prelude

-- | Adjoint of @J : a -> b@ with respect to domain metric @g_a@ and
-- codomain metric @g_b@.
--
-- The metrics are themselves 'Diffs of type @Diff' (point, vector) vector@:
-- the forward pass lowers or raises a vector at the given point.  The
-- backward pass of the metric carries @∂g@ in its point-slot; 'adjointWith'
-- uses only the forward passes.
adjointWith ::
  -- | @g_a^-1@ — raise a covector on the domain to a vector
  Diff' (a, a) a ->
  -- | @g_b@ — lower a vector on the codomain to a covector
  Diff' (b, b) b ->
  -- | @J : a -> b@
  Diff' a b ->
  -- | @J@ with its pullback conjugated to the @g@-adjoint.  The type is
  -- still @Diff' a b@ because the lens stores the forward map @a -> b@ and
  -- the backward map @b -> a@; the adjoint lives in the pullback.
  Diff' a b
adjointWith :: forall a b.
Diff' (a, a) a -> Diff' (b, b) b -> Diff' a b -> Diff' a b
adjointWith Diff' (a, a) a
raiseA Diff' (b, b) b
lowerB (Diff a -> (b, b -> a)
f) = (a -> (b, b -> a)) -> Diff () a b
forall k (p :: k) a b. (a -> (b, b -> a)) -> Diff p a b
Diff ((a -> (b, b -> a)) -> Diff () a b)
-> (a -> (b, b -> a)) -> Diff () a b
forall a b. (a -> b) -> a -> b
$ \a
x ->
  let (b
y, b -> a
jt) = a -> (b, b -> a)
f a
x
   in ( b
y,
        \b
db ->
          let (b
lowered, b -> (b, b)
_) = Diff' (b, b) b -> (b, b) -> (b, b -> (b, b))
forall k (p :: k) a b. Diff p a b -> a -> (b, b -> a)
runDiff Diff' (b, b) b
lowerB (b
y, b
db)
              covectorA :: a
covectorA = b -> a
jt b
lowered
              (a
raised, a -> (a, a)
_) = Diff' (a, a) a -> (a, a) -> (a, a -> (a, a))
forall k (p :: k) a b. Diff p a b -> a -> (b, b -> a)
runDiff Diff' (a, a) a
raiseA (a
x, a
covectorA)
           in a
raised
      )

-- | Raise a covector using a metric at the given point.
raiseWith :: Diff' (s, a) a -> s -> a -> a
raiseWith :: forall s a. Diff' (s, a) a -> s -> a -> a
raiseWith Diff' (s, a) a
g s
s a
covec = (a, a -> (s, a)) -> a
forall a b. (a, b) -> a
fst (Diff' (s, a) a -> (s, a) -> (a, a -> (s, a))
forall k (p :: k) a b. Diff p a b -> a -> (b, b -> a)
runDiff Diff' (s, a) a
g (s
s, a
covec))

-- | Lower a vector using a metric at the given point.
lowerWith :: Diff' (s, a) a -> s -> a -> a
lowerWith :: forall s a. Diff' (s, a) a -> s -> a -> a
lowerWith Diff' (s, a) a
g s
s a
vec = (a, a -> (s, a)) -> a
forall a b. (a, b) -> a
fst (Diff' (s, a) a -> (s, a) -> (a, a -> (s, a))
forall k (p :: k) a b. Diff p a b -> a -> (b, b -> a)
runDiff Diff' (s, a) a
g (s
s, a
vec))

-- | Euclidean metric @g = δ@: raise and lower are both the identity.
euclideanMetric :: (Additive a) => Diff' (a, a) a
euclideanMetric :: forall a. Additive a => Diff' (a, a) a
euclideanMetric = ((a, a) -> (a, a -> (a, a))) -> Diff () (a, a) a
forall k (p :: k) a b. (a -> (b, b -> a)) -> Diff p a b
Diff (((a, a) -> (a, a -> (a, a))) -> Diff () (a, a) a)
-> ((a, a) -> (a, a -> (a, a))) -> Diff () (a, a) a
forall a b. (a -> b) -> a -> b
$ \(a
_, a
v) -> (a
v, (a, a) -> a -> (a, a)
forall a b. a -> b -> a
const (a
forall a. Additive a => a
zero, a
forall a. Additive a => a
zero))

-- | Diagonal metric @g(s) = diag(f s)@.
--
-- Lowering maps @x -> f s * x@; raising maps @c -> recip (f s) * c@.
-- The pullback's state-slot is zero because the typical consumer (an
-- optimizer) does not back-propagate through the metric coefficients.
-- If you need Christoffel-style @∂g@, write a custom metric 'Diff.
diagonalMetric ::
  (Additive s) =>
  -- | @f@ such that @g(s) = diag(f s)@
  (s -> a -> a) ->
  Diff' (s, a) a
diagonalMetric :: forall s a. Additive s => (s -> a -> a) -> Diff' (s, a) a
diagonalMetric s -> a -> a
f = ((s, a) -> (a, a -> (s, a))) -> Diff () (s, a) a
forall k (p :: k) a b. (a -> (b, b -> a)) -> Diff p a b
Diff (((s, a) -> (a, a -> (s, a))) -> Diff () (s, a) a)
-> ((s, a) -> (a, a -> (s, a))) -> Diff () (s, a) a
forall a b. (a -> b) -> a -> b
$ \(s
s, a
x) ->
  let gx :: a
gx = s -> a -> a
f s
s a
x
   in ( a
gx,
        \a
dc -> (s
forall a. Additive a => a
zero, s -> a -> a
f s
s a
dc)
      )