0 | module Data.Container.Additive.Morphism.Instances
  1 |
  2 | import Data.Vect
  3 | import Data.List.Quantifiers
  4 | import Data.Vect.Quantifiers
  5 |
  6 | import Data.Container.Base
  7 | import Data.Container.Base.Morphism.Instances as Base
  8 | import Data.ComMonoid
  9 | import Data.Materialise
 10 | import Data.Num
 11 | import Data.Container.Additive.Object.Definition
 12 | import Data.Container.Additive.Object.Instances
 13 | import Data.Container.Additive.Product.Instances
 14 | import Data.Container.Additive.Extension.Definition
 15 | import Data.Container.Additive.Morphism.Definition
 16 | import Data.Container.Additive.Product.Definition
 17 | import Data.Container.Additive.Properties.Definition
 18 |
 19 | import Data.Container.Additive.Quantifiers
 20 |
 21 |
 22 | import Misc
 23 |
 24 | %hide Base.Object.Instances.Const
 25 | %hide Data.Vect.Quantifiers.All.index
 26 | %hide Base.Morphism.Instances.State.State
 27 | %hide Base.Morphism.Instances.Costate.Costate
 28 |
 29 | public export
 30 | pushIntoContinuation : {p : AddCont} ->
 31 |   (f : d >*< p =%+> l) ->
 32 |   (p =%+> Nap d.Shp >-+@ l)
 33 | pushIntoContinuation f = !%+ \param => (() <| \dShp => f.fwd (dShp, param) **
 34 |     fromGenerators (\(dShp ** grad) => snd (f.bwd (dShp, param) grad)))
 35 |
 36 | ||| Categorical product of additive containers
 37 | ||| On underlying containers computed as the hancock tensor product
 38 | namespace CategoricalProduct
 39 |   public export
 40 |   terminal : {c : AddCont} -> c =%+> UnitCont
 41 |   terminal = !%+ \x => (() ** \() => c.Zero x)
 42 |
 43 |   public export
 44 |   leftUnit : UnitCont >*< c =%+> c
 45 |   leftUnit = !% leftUnit
 46 |   
 47 |   public export
 48 |   rightUnit : c >*< UnitCont =%+> c
 49 |   rightUnit = !% rightUnit
 50 |
 51 |   public export
 52 |   leftUnitInv : c =%+> UnitCont >*< c
 53 |   leftUnitInv = !% leftUnitInv
 54 |   
 55 |   public export
 56 |   rightUnitInv : c =%+> c >*< UnitCont
 57 |   rightUnitInv = !% rightUnitInv
 58 |
 59 |   public export
 60 |   assocL : (a >*< b) >*< c =%+> a >*< (b >*< c)
 61 |   assocL = !% assocL {a=UC a, b=UC b, c=UC c}
 62 |
 63 |   public export
 64 |   assocR : a >*< (b >*< c) =%+> (a >*< b) >*< c
 65 |   assocR = !% assocR {a=UC a, b=UC b, c=UC c}
 66 |
 67 |   public export
 68 |   swap : a >*< b =%+> b >*< a
 69 |   swap = !% swap {a=UC a, b=UC b}
 70 |
 71 |   public export
 72 |   swapMiddle : (c1 >*< c2) >*< (c3 >*< c4) =%+> (c1 >*< c3) >*< (c2 >*< c4)
 73 |   swapMiddle = !% swapMiddle {c1=UC c1, c2=UC c2, c3=UC c3, c4=UC c4}
 74 |
 75 |   ||| These do not exist for ordinary containers!
 76 |   ||| Here we need `c` not to be erased since we're using its monoid structure
 77 |   public export
 78 |   copy : {c : AddCont} -> c =%+> c >*< c
 79 |   copy = !%+ \x => ((x, x) ** uncurry (c.Plus x))
 80 |   
 81 |   public export
 82 |   pairMaps : {c : AddCont} ->
 83 |     c =%+> d ->
 84 |     c =%+> e ->
 85 |     c =%+> d >*< e
 86 |   pairMaps f g = copy %+>> (f >*< g)
 87 |
 88 |
 89 |   ||| Materialises both the forward pass and the backward pass
 90 |   public export
 91 |   materialiseCont : Materialise c.Shp =>
 92 |     (ip : InterfaceOnPositions c Materialise) =>
 93 |     c =%+> c
 94 |   materialiseCont = !%+ \x => (materialise x ** \x' =>
 95 |     materialise @{GetInterface ip x} (replace {p=c.PosSet} materialiseIsId x'))
 96 |
 97 |   ||| Generally used for pairing up parameters
 98 |   public export
 99 |   constPair : ComMonoid s => ComMonoid t =>
100 |     Const (s, t) =%+> Const s >*< Const t
101 |   constPair = !%+ \x => (x ** id)
102 |
103 |   ||| The n-ary `constPair`
104 |   public export
105 |   constFinite : {n : Nat} -> {ps : Fin n -> Type} ->
106 |     (ms : (i : Fin n) -> ComMonoid (ps i)) ->
107 |     Const (Product (\i => Const (ps i) @{ms i})).Shp
108 |       @{finiteShpMon {f = \i => Const (ps i) @{ms i}} ms}
109 |       =%+> Product (\i => Const (ps i) @{ms i})
110 |   constFinite {n = 0} _ = !%+ \x => (x ** id)
111 |   constFinite {n = S k} ms
112 |     = constPair @{ms FZ} @{finiteShpMon {f = \i => Const (ps (FS i)) @{ms (FS i)}} (\i => ms (FS i))}
113 |       %+>> (id {c = Const (ps FZ) @{ms FZ}}
114 |             >*< constFinite {ps = \i => ps (FS i)} (\i => ms (FS i)))
115 |
116 |   ||| Projection out of an iterated product; backward is the coprojection
117 |   public export
118 |   projFinite : {n : Nat} -> {f : Fin n -> AddCont} ->
119 |     (i : Fin n) -> Product f =%+> f i
120 |   projFinite i = !%+ \s => (indexShp {f} i s ** injectPos {f} i s)
121 |
122 |   public export
123 |   projLeft : {d : AddCont} -> c >*< d =%+> c
124 |   projLeft = !%+ \(x, y) => (x ** \x' => (x', d.Zero y))
125 |   
126 |   public export
127 |   projRight : {c : AddCont} -> c >*< d =%+> d
128 |   projRight = !%+ \(x, y) => (y ** \y' => (c.Zero x, y'))
129 |
130 | namespace Morphism
131 |   ||| The action of the `Maybe` functor on lenses
132 |   public export
133 |   Maybe : c =%+> d -> Maybe c =%+> Maybe d
134 |   Maybe f = id >+< f
135 |
136 | ||| Universal maps of the dependent pair and the section of a family
137 | namespace Dependent
138 |   public export
139 |   fanOut : {c : AddCont} -> {0 ix : Type} -> {0 f : ix -> AddCont} ->
140 |     ((i : ix) -> c =%+> f i) -> c =%+> Section f
141 |   fanOut gs = !%+ \x =>
142 |     (\i => (gs i).fwd x **
143 |      fromGenerators {y = c.Pos x} (\(i ** g) => (gs i).bwd x g))
144 |
145 |   ||| The action of `Section` on a family of lenses
146 |   public export
147 |   mapSection : {0 ix : Type} -> {0 a, b : ix -> AddCont} ->
148 |     ((i : ix) -> a i =%+> b i) -> Section a =%+> Section b
149 |   mapSection fs = !%+ \s =>
150 |     (\i => (fs i).fwd (s i) **
151 |      map (\ig => (fst ig ** (fs (fst ig)).bwd (s (fst ig)) (snd ig))))
152 |
153 |   public export
154 |   sectionStrength : {0 g : ix -> AddCont} -> {p : AddCont} ->
155 |     Section g >*< p =%+> Section (\i => g i >*< p)
156 |   sectionStrength = !%+ \sq =>
157 |     (\i => (fst sq i, snd sq) ** \bag =>
158 |       ( map (\ig => (fst ig ** fst (snd ig))) bag
159 |       , sum @{UMon p (snd sq)} (map (\ig => snd (snd ig)) bag)))
160 |
161 |   public export
162 |   generators : {n : Nat} -> {0 f : Fin n -> AddCont} ->
163 |     (s : (i : Fin n) -> (f i).Shp) ->
164 |     (Product f).PosSet (tabulateShp {f} s) ->
165 |     List (i : Fin n ** (f i).PosSet (s i))
166 |   generators {n = 0} s () = []
167 |   generators {n = S k} s (g, gs) = (FZ ** g) ::
168 |     map (\ig => (FS (fst ig) ** snd ig)) (generators {f = f . FS} (\i => s (FS i)) gs)
169 |
170 |   ||| Evaluate a section over `Fin n` at every index
171 |   public export
172 |   evalAll : {n : Nat} -> {0 f : Fin n -> AddCont} -> Section f =%+> Product f
173 |   evalAll = !%+ \s => (tabulateShp {f} s ** \g => MkBag (generators {f} s g))
174 |
175 |   ||| The universal map out of a dependent pair
176 |   public export
177 |   copair : {0 ix : Type} -> {0 c : ix -> AddCont} -> {0 d : AddCont} ->
178 |     ((i : ix) -> c i =%+> d) -> AddContDPair c =%+> d
179 |   copair fs = !%+ \(i ** x) => (%!+) (fs i) x
180 |
181 |   ||| Evaluate a section at the index an indexed value carries and pair it with
182 |   ||| that value
183 |   public export
184 |   evalSection : {0 ix : Type} -> {0 a, p : ix -> AddCont} ->
185 |     Section a >*< AddContDPair p =%+> AddContDPair (\i => a i >*< p i)
186 |   evalSection = !%+ \(s, (i ** y)) =>
187 |     ((i ** (s i, y)) ** \(x', y') => (MkBag [(i ** x')], y'))
188 |
189 |   ||| The graph of a section, implemented through effects
190 |   public export
191 |   graph : {0 x : Type} -> {0 f : x -> AddCont} ->
192 |     Section f =%+> Nap x >-+@ AddContDPair f
193 |   graph = !%+ \s => (() <| (\i => (i ** s i)) ** id)
194 |
195 |   ||| A lens from two dependent pairs into one. When they agree
196 |   ||| on the index they're paired, when they don't we return failure
197 |   public export
198 |   matchIndex : DecEq a => {b, l : a -> AddCont} ->
199 |     AddContDPair b >*< AddContDPair l =%+> Maybe (AddContDPair (\i => b i >*< l i))
200 |   matchIndex = !%+ \((i ** y), (j ** lab)) => case decEq i j of
201 |     Yes Refl => (Right (i ** (y, lab)) ** id)
202 |     No _ => (Left () ** \() => ((b i).Zero y, (l j).Zero lab))
203 |
204 | ||| Structure maps of the left action `>-+@` of `(Cont, >@, Scalar)` on AddCont
205 | ||| They generally use the following components:
206 | ||| * `pureBw` : a position becomes the singleton bag containing it
207 | ||| * `sumBw` : a bag of positions is added up using its monoid structure
208 | ||| * `joinBwComp` : nested bags of positions are flattened
209 | namespace CompositionProductAction
210 |   ||| Backwards pass is ComMon-homomorphism on the nose
211 |   public export
212 |   unitor : {c : AddCont} -> c =%+> Scalar >-+@ c
213 |   unitor = !% sumBw c %>> (Bag <!> leftUnitInv {c=UC c})
214 |
215 |   ||| Backwards map is a ComMon-homomorphism only through the quotient
216 |   public export
217 |   unitorInv : Scalar >-+@ c =%+> c
218 |   unitorInv = !% (Bag <!> leftUnit {c=UC c}) %>> pureBw
219 |
220 |   public export
221 |   multiplicator : m >-+@ n >-+@ c =%+> (m >@ n) >-+@ c
222 |   multiplicator = !% Bag <!> ((id >@ pureBw {c = n >@ UC c}) %>> assocR {c=UC c})
223 |
224 |   public export
225 |   multiplicatorInv : ((m >@ n) >-+@ c) =%+> (m >-+@ (n >-+@ c))
226 |   multiplicatorInv = !% (Bag <!> assocL {c=UC c}) %>> joinBwComp {d = n >@ UC c}
227 |
228 | ||| `!*` and `- >-+@ Scalar` are isomorphic: they're both right adjoint to 
229 | ||| `UC`.  They are two presentations of the same free commutative monoid on
230 | ||| positions: `!*` stores a bag of positions directly, while `- >-+@ Scalar` 
231 | ||| stores generators tagged with `Nat` multiplicities.
232 | namespace CompositionActionBang
233 |   ||| Read each position as the generator it is, with multiplicity one
234 |   public export
235 |   actionToFree : {0 e : Cont} -> e >-+@ Scalar =%+> !* e
236 |   actionToFree = !%+ \ex => (shapeExt ex ** map (\mp => (mp ** 1)))
237 |
238 |   ||| Expand each generator into as many copies as its multiplicity says.
239 |   public export
240 |   freeToAction : {0 e : Cont} -> !* e =%+> e >-+@ Scalar
241 |   freeToAction = !%+ \x => (x <| \_ => () ** \gens =>
242 |     sum @{bagIsMonoid} ((\(mp ** n) => scale @{bagIsMonoid} n (pure mp)) <$> gens))
243 |
244 | ||| Structure maps of the left-skew monoidal product `>+@` on AddCont
245 | ||| These definitions follow Theorem 3.1 in https://arxiv.org/abs/2506.06847
246 | namespace SkewCompositionProduct
247 |   ||| Hom-set isomorphism of the adjunction, which is the general purpose
248 |   ||| `addContTranspose` read through the isomorphism above
249 |   public export
250 |   adjR : {c : AddCont} -> UC c =%> m -> c =%+> m >-+@ Scalar
251 |   adjR f = addContTranspose f %+>> freeToAction
252 |
253 |   ||| Inverse of the hom-set isomorphism of the adjunction
254 |   public export
255 |   adjL : c =%+> m >-+@ Scalar -> UC c =%> m
256 |   adjL g = addContTransposeInv (g %+>> actionToFree)
257 |
258 |   public export
259 |   epsilon : UC Scalar =%> Scalar
260 |   epsilon = adjL (unitor {c = Scalar})
261 |
262 |   public export
263 |   leftUnit : {c : AddCont} -> Scalar >+@ c =%+> c
264 |   leftUnit = (epsilon >-+@ id) %+>> unitorInv
265 |
266 |   public export
267 |   rightUnit : {c : AddCont} -> c =%+> c >+@ Scalar
268 |   rightUnit = adjR id
269 |
270 |   public export
271 |   associator : {b : AddCont} -> (a >+@ b) >+@ c =%+> a >+@ (b >+@ c)
272 |   associator = (adjL ((id >-+@ SkewCompositionProduct.rightUnit {c = b})
273 |     %+>> multiplicator {c = Scalar}) >-+@ id) %+>> multiplicatorInv 
274 |
275 |   {-
276 |   Beyond skew structure, only `leftUnitInv` exists, rightUnitInv and assocR do not.
277 |
278 |   leftUnitInv is also not inverse to leftUnit. We only have
279 |   `leftUnitInv %+>> leftUnit = id`, but not the other way around.
280 |
281 |   The right associator is not definable because the forward  part involves the
282 |   the function `g : List (aPos ** bPos) -> c.Shp` which would have to collapse
283 |   a whole list of of positions into a single shape
284 |   -}
285 |   public export
286 |   leftUnitInv : {c : AddCont} -> c =%+> Scalar >+@ c
287 |   leftUnitInv = unitor %+>> (toState () >-+@ id)
288 |
289 |
290 | namespace Coproduct
291 |   public export
292 |   elim : c >+< c =%+> c
293 |   elim = !% elim
294 |
295 | ||| Lax interchange between the categorical product `>*<` on AddCont and the
296 | ||| action `>-+@` of `(Cont, ><)` on it. Not an isomorphism.
297 | public export
298 | duoidal : (m >-+@ d) >*< (n >-+@ g) =%+> (m >< n) >-+@ (d >*< g)
299 | duoidal = !%+ \(exM, exN) =>
300 |   ((shapeExt exM, shapeExt exN) <| \(mp, np) => (index exM mp, index exN np) **
301 |     \bag => ((\((mp, np) ** (dp, gp)) => (mp ** dp)) <$> bag,
302 |              (\((mp, np) ** (dp, gp)) => (np ** gp)) <$> bag))
303 |
304 | ||| Specific distributive law we need
305 | public export
306 | distribute : {c : AddCont} ->
307 |   (f : c.Shp -> e =%> s) ->
308 |   c >*< (e >-+@ g) =%+> s >-+@ g
309 | distribute f = uncurry (!%+ \cs => (f cs >-+@ id {c=g} ** \_ => c.Zero cs))
310 |
311 | public export
312 | coprodDistrOverTensor : {q, p : AddCont} ->
313 |   (a >+< b) >*< (p >*< q) =%+> (a >*< p) >+< (b >*< q)
314 | coprodDistrOverTensor = !%+ \case
315 |   (Left a, (p, _)) => (Left (a, p) ** \(a', p') => (a', (p', q.Zero _)))
316 |   (Right b, (_, q)) => (Right (b, q) ** \(b', q') => (b', (p.Zero _, q')))
317 |
318 | {-
319 | ||| Not an isomorphism, arising from duoidal structure between >@ and ><
320 | public export
321 | rebracketcomptensor: {y : AddCont} -> (e >+@ y) >< y =%+> e >+@ (y >< y)
322 | rebracketcomptensor = (id {c=e >+@ y} >< leftUnitInv {c=y})
323 |                       %+>> duoidal {c=e} {d=y} {e=Scalar} {f=y}
324 |                       %+>> (rightUnit {c=e} >+@ id {c=(y><y)})
325 |
326 |
327 | public export
328 | extractEffect : {d : AddCont} ->
329 |   d >< (e >+@ f) =%+> e >+@ (d >< f)
330 | extractEffect = (leftUnitInv >< (id {c=e >+@ f}))
331 |             %+>> duoidal {c=Scalar}
332 |             %+>> (leftUnit >+@ (id {c=d><f}))
333 |
334 | -}
335 |
336 | ||| References for State
337 | ||| Bruno's PhD thesis: https://arxiv.org/abs/2403.13001
338 | ||| Towards Foundations of Cat. Cybernetics: https://arxiv.org/abs/2105.06332
339 | namespace State
340 |   ||| State here differers for the one in `Cont`, because `Scalar` is different
341 |   |||
342 |   |||       ┌─────────────┐
343 |   |||       │             ├──► (x : c.Shp)
344 |   |||       │    State    │
345 |   |||       │             ├◄── c.Pos x
346 |   |||       └─────────────┘
347 |   public export
348 |   State : AddCont -> Type
349 |   State c = Scalar =%+> c
350 |
351 |   public export
352 |   toState : (x : c.Shp) -> State c
353 |   toState x = ?somethingInterestingHmm
354 |   
355 |   -- public export
356 |   -- fromState : State c -> c.Shp
357 |   -- fromState f = f.fwd ()
358 |
359 | ||| References for Costate
360 | ||| Bruno's PhD thesis: https://arxiv.org/abs/2403.13001
361 | ||| Towards Foundations of Cat. Cybernetics: https://arxiv.org/abs/2105.06332
362 | namespace Costate
363 |   ||| Costate here differs from the one in `Cont`, because `Scalar` is different
364 |   |||                  ┌─────────────┐
365 |   |||  (x : c.Shp)  ──►┤             │
366 |   |||                  │   Costate   │
367 |   |||     c.Pos x   ◄──┤             │
368 |   |||                  └─────────────┘
369 |   public export
370 |   Costate : AddCont -> Type
371 |   Costate c = c =%+> Scalar
372 |   
373 |   public export
374 |   toCostate : {c : AddCont} ->
375 |     ((x : c.Shp) -> c.PosSet x) -> Costate c
376 |   toCostate s = !%+ \x => (() ** \n => scale @{UMon c x} n (s x))
377 |
378 |   -- public export
379 |   -- fromCostate : Costate c -> (x : c.Shp) -> c.PosSet x
380 |   -- fromCostate f x = f.bwd x ()
381 |
382 |   public export
383 |   constantOne : {c : AddCont} ->
384 |     InterfaceOnPositions c Num => Costate c
385 |   constantOne @{ip} = toCostate (\x => let numPos = GetInterface ip x in 1)
386 |
387 |   public export
388 |   Delete : {c : AddCont} -> Costate c 
389 |   Delete = toCostate c.Zero
390 |   
391 |
392 | public export
393 | sum : Num a =>
394 |   (Const a >*< Const a) =%+> Const a
395 | sum = !%+ \(x1, x2) => (x1 + x2 ** \x' => (x', x'))
396 |
397 | public export
398 | bwSumBag : {l : Type} ->
399 |   (xs : List l) ->
400 |   (d' : l) ->
401 |   All (const l) xs
402 | bwSumBag [] d' = []
403 | bwSumBag (x :: xs) d' = d' :: bwSumBag xs d'
404 |
405 | public export
406 | sumBag : {l : Type} -> ComMonoid l =>
407 |   BagAll (Const l) =%+> Const l
408 | sumBag = !%+ \(MkBag xs) => (sum (MkBag xs) ** \d' => bwSumBag xs d')
409 |
410 | public export
411 | negate : Num a => Neg a =>
412 |   Const a =%+> Const a
413 | negate = !%+ \x => (-x ** \x' => -x')
414 |
415 | public export
416 | zero : {c : AddCont} -> Num a =>
417 |   c =%+> Const a
418 | zero = !%+ \_ => (0 ** \_ => c.Zero _)
419 |
420 | public export
421 | mul : Num a =>
422 |   (Const a >*< Const a) =%+> Const a
423 | mul = !%+ \(x1, x2) => (x1 * x2 ** \x' => (x' * x2, x' * x1))
424 |
425 | ||| Mean squared error
426 | public export
427 | SquaredDifference : {a : Type} -> Num a => Neg a =>
428 |   (Const a >*< Const a) =%+> (Const a)
429 | SquaredDifference = ((id {c=Const a}) >*< negate) %+>> sum %+>> copy %+>> mul
430 |
431 | namespace Sample
432 |   ||| Select a shape from All to produce an Any at the given index
433 |   ||| Same as `index i (allAnies shapes)` but reduces better
434 |   public export
435 |   selectShape : {cs : Vect k AddCont} ->
436 |     (shapes : All (.Shp) cs) -> (i : Fin k) -> Any (.Shp) cs
437 |   selectShape (s :: ss) FZ = Here s
438 |   selectShape (s :: ss) (FS j) = There (selectShape ss j)
439 |
440 |   ||| Extract the position from an AnyPos at a given index
441 |   public export
442 |   extractPos : {n : Nat} -> {xs : Vect n AddCont} -> {shapes : All (.Shp) xs} ->
443 |     (i : Fin n) ->
444 |     AnyShpPos (selectShape shapes i) ->
445 |     (index i xs).PosSet (index i shapes)
446 |   extractPos {shapes = (_ :: _)} FZ (Here x') = x'
447 |   extractPos {shapes = (_ :: _)} (FS j) (There rest) = extractPos j rest
448 |
449 |
450 |
451 | -- parameters (f : Type -> Type)
452 | --   ||| These are all of the morphisms in the cokleisli category of (f <!> -)  
453 | --   public export
454 | --   MonLens : Cont -> Cont -> Type
455 | --   MonLens c d = (f <!> c) =%> d
456 | -- 
457 | --   public export
458 | --   counit : Monad f => f <!> c =%> c
459 | --   counit = !% \x => (x ** pure)
460 | -- 
461 | --   public export
462 | --   cojoin : Monad f => (f <!> c) =%> (f <!> (f <!> c))
463 | --   cojoin = !% \x => (x ** join)
464 |
465 |   
466 | -- public export
467 | -- record FCoAlgCont (f : Type -> Type) where
468 | --   constructor MkFCoAlgCont
469 | --   carrier : Cont
470 | --   coalg : (a : carrier.Shp) -> f (carrier.Pos a) -> carrier.Pos a
471 |
472 | -- public export
473 | -- coAlgMorphism : (c, d : FCoAlgCont f) -> Type
474 | -- coAlgMorphism c d = c.carrier =%> d.carrier
475 | -- 
476 | -- convert : FCoAlgCont List -> AddCont
477 | -- convert (MkFCoAlgCont carrier coalg) = MkAddCont
478 | --   carrier
479 | --   {mon=(MkI $ \s => MkComMonoid
480 | --     (\l, r => coalg s [l, r])
481 | --     (coalg s []))}