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.Extension.Definition
 14 | import Data.Container.Additive.Morphism.Definition
 15 | import Data.Container.Additive.Product.Definition
 16 | import Data.Container.Additive.Properties.Definition
 17 |
 18 | import Data.Container.Additive.Quantifiers
 19 |
 20 |
 21 | import Misc
 22 |
 23 | %hide Base.Object.Instances.Const
 24 | %hide Data.Vect.Quantifiers.All.index
 25 | %hide Base.Morphism.Instances.State.State
 26 | %hide Base.Morphism.Instances.Costate.Costate
 27 |
 28 | public export
 29 | pushIntoContinuation : {p : AddCont} ->
 30 |   (f : d >*< p =%+> l) ->
 31 |   (p =%+> (pushDown d.Shp) >-+@ l)
 32 | pushIntoContinuation f = !%+ \param => (() <| \dShp => f.fwd (dShp, param) **
 33 |     fromGenerators (\(dShp ** grad=> snd (f.bwd (dShp, param) grad)))
 34 |
 35 | ||| Categorical product of additive containers
 36 | ||| On underlying containers computed as the hancock tensor product
 37 | namespace CategoricalProduct
 38 |   ||| The unique map to the terminal object; its backward is the zero
 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 |   ||| Materialises both the forward pass and the backward pass
 89 |   public export
 90 |   materialiseCont : Materialise c.Shp =>
 91 |     (ip : InterfaceOnPositions c Materialise) =>
 92 |     c =%+> c
 93 |   materialiseCont {ip=MkI i} = !%+ \x => (materialise x ** \x' =>
 94 |     materialise @{i x} (replace {p=c.PosSet} materialiseIsId x'))
 95 |
 96 |   ||| Generally used for pairing up parameters
 97 |   public export
 98 |   constPair : ComMonoid s => ComMonoid t =>
 99 |     Const (s, t) =%+> Const s >*< Const t
100 |   constPair = !%+ \x => (x ** id)
101 |
102 |   public export
103 |   projLeft : {d : AddCont} -> c >*< d =%+> c
104 |   projLeft = !%+ \(x, y) => (x ** \x' => (x', d.Zero y))
105 |   
106 |   public export
107 |   projRight : {c : AddCont} -> c >*< d =%+> d
108 |   projRight = !%+ \(x, y) => (y ** \y' => (c.Zero x, y'))
109 |
110 | ||| Structure maps of the left action `>-+@` of `(Cont, >@, Scalar)` on AddCont
111 | ||| They generally use the following components:
112 | ||| * `pureBw` : a position becomes the singleton bag containing it
113 | ||| * `sumBw` : a bag of positions is added up using its monoid structure
114 | ||| * `joinBwComp` : nested bags of positions are flattened
115 | namespace CompositionProductAction
116 |   ||| Backwards pass is ComMon-homomorphism on the nose
117 |   public export
118 |   unitor : {c : AddCont} -> c =%+> Scalar >-+@ c
119 |   unitor = !% (sumBw @{mon c} %>> (Bag <!> leftUnitInv {c=UC c}))
120 |
121 |   ||| Backwards map is a ComMon-homomorphism only through the quotient
122 |   public export
123 |   unitorInv : Scalar >-+@ c =%+> c
124 |   unitorInv = !% ((Bag <!> leftUnit {c=UC c}) %>> pureBw)
125 |
126 |   public export
127 |   multiplicator : (m >-+@ (n >-+@ c)) =%+> ((m >@ n) >-+@ c)
128 |   multiplicator = !% (Bag <!> ((id {c=m} >@ pureBw {c = n >@ UC c}) %>> assocR {a=m, b=n, c=UC c}))
129 |
130 |   public export
131 |   multiplicatorInv : ((m >@ n) >-+@ c) =%+> (m >-+@ (n >-+@ c))
132 |   multiplicatorInv = !% ((Bag <!> assocL {a=m, b=n, c=UC c}) %>> joinBwComp {c=m, d = n >@ UC c})
133 |
134 | ||| `!*` and `- >-+@ Scalar` are isomorphic: they're both right adjoint to 
135 | ||| `UC`.  They are two presentations of the same free commutative monoid on
136 | ||| positions: `!*` stores a bag of positions directly, while `- >-+@ Scalar` 
137 | ||| stores generators tagged with `Nat` multiplicities.
138 | namespace CompositionActionBang
139 |   ||| Read each position as the generator it is, with multiplicity one
140 |   public export
141 |   actionToFree : {0 e : Cont} -> e >-+@ Scalar =%+> !* e
142 |   actionToFree = !%+ \ex => (shapeExt ex ** map (\mp => (mp ** 1)))
143 |
144 |   ||| Expand each generator into as many copies as its multiplicity says.
145 |   public export
146 |   freeToAction : {0 e : Cont} -> !* e =%+> e >-+@ Scalar
147 |   freeToAction = !%+ \x => (x <| \_ => () ** \gens =>
148 |     sum @{bagIsMonoid} ((\(mp ** n=> scale @{bagIsMonoid} n (pure mp)) <$> gens))
149 |
150 | ||| Structure maps of the left-skew monoidal product `>+@` on AddCont
151 | ||| These definitions follow Theorem 3.1 in https://arxiv.org/abs/2506.06847
152 | namespace SkewCompositionProduct
153 |   ||| Hom-set isomorphism of the adjunction, which is the general purpose
154 |   ||| `addContTranspose` read through the isomorphism above
155 |   public export
156 |   adjR : {c : AddCont} -> UC c =%> m -> c =%+> m >-+@ Scalar
157 |   adjR f = addContTranspose f %+>> freeToAction
158 |
159 |   ||| Inverse of the hom-set isomorphism of the adjunction
160 |   public export
161 |   adjL : c =%+> m >-+@ Scalar -> UC c =%> m
162 |   adjL g = addContTransposeInv (g %+>> actionToFree)
163 |
164 |   public export
165 |   epsilon : UC Scalar =%> Scalar
166 |   epsilon = adjL (unitor {c = Scalar})
167 |
168 |   public export
169 |   leftUnit : {c : AddCont} -> Scalar >+@ c =%+> c
170 |   leftUnit = (epsilon >-+@ id) %+>> unitorInv
171 |
172 |   public export
173 |   rightUnit : {c : AddCont} -> c =%+> c >+@ Scalar
174 |   rightUnit = adjR id
175 |
176 |   public export
177 |   associator : {b : AddCont} -> (a >+@ b) >+@ c =%+> a >+@ (b >+@ c)
178 |   associator = (adjL ((id >-+@ SkewCompositionProduct.rightUnit {c = b})
179 |     %+>> multiplicator {c = Scalar}) >-+@ id) %+>> multiplicatorInv 
180 |
181 |   {-
182 |   Beyond skew structure, only `leftUnitInv` exists, rightUnitInv and assocR do not.
183 |
184 |   leftUnitInv is also not inverse to leftUnit. We only have
185 |   `leftUnitInv %+>> leftUnit = id`, but not the other way around.
186 |
187 |   The right associator is not definable because the forward  part involves the
188 |   the function `g : List (aPos ** bPos) -> c.Shp` which would have to collapse
189 |   a whole list of of positions into a single shape
190 |   -}
191 |   public export
192 |   leftUnitInv : {c : AddCont} -> c =%+> Scalar >+@ c
193 |   leftUnitInv = unitor %+>> (toState () >-+@ id)
194 |
195 |
196 | namespace Coproduct
197 |   public export
198 |   elim : c >+< c =%+> c
199 |   elim = !% elim
200 |
201 | ||| Lax interchange between the categorical product `>*<` on AddCont and the
202 | ||| action `>-+@` of `(Cont, ><)` on it. Not an isomorphism.
203 | public export
204 | duoidal : (m >-+@ d) >*< (n >-+@ g) =%+> (m >< n) >-+@ (d >*< g)
205 | duoidal = !%+ \(exM, exN) =>
206 |   ((shapeExt exM, shapeExt exN) <| \(mp, np) => (index exM mp, index exN np) **
207 |     \bag => ((\((mp, np) ** (dp, gp)) => (mp ** dp)) <$> bag,
208 |              (\((mp, np) ** (dp, gp)) => (np ** gp)) <$> bag))
209 |
210 | ||| Specific distributive law we need
211 | public export
212 | distribute : {c : AddCont} ->
213 |   (f : c.Shp -> e =%> s) ->
214 |   c >*< (e >-+@ g) =%+> s >-+@ g
215 | distribute f = uncurry (!%+ \cs => (f cs >-+@ id {c=g} ** \_ => c.Zero cs))
216 |
217 | public export
218 | coprodDistrOverTensor : {q, p : AddCont} ->
219 |   (a >+< b) >*< (p >*< q) =%+> (a >*< p) >+< (b >*< q)
220 | coprodDistrOverTensor = !%+ \case
221 |   (Left a, (p, _)) => (Left (a, p) ** \(a', p') => (a', (p', q.Zero _)))
222 |   (Right b, (_, q)) => (Right (b, q) ** \(b', q') => (b', (p.Zero _, q')))
223 |
224 | {-
225 | ||| Not an isomorphism, arising from duoidal structure between >@ and ><
226 | public export
227 | rebracketcomptensor: {y : AddCont} -> (e >+@ y) >< y =%+> e >+@ (y >< y)
228 | rebracketcomptensor = (id {c=e >+@ y} >< leftUnitInv {c=y})
229 |                       %+>> duoidal {c=e} {d=y} {e=Scalar} {f=y}
230 |                       %+>> (rightUnit {c=e} >+@ id {c=(y><y)})
231 |
232 |
233 | public export
234 | extractEffect : {d : AddCont} ->
235 |   d >< (e >+@ f) =%+> e >+@ (d >< f)
236 | extractEffect = (leftUnitInv >< (id {c=e >+@ f}))
237 |             %+>> duoidal {c=Scalar}
238 |             %+>> (leftUnit >+@ (id {c=d><f}))
239 |
240 | -}
241 |
242 | ||| References for State
243 | ||| Bruno's PhD thesis: https://arxiv.org/abs/2403.13001
244 | ||| Towards Foundations of Cat. Cybernetics: https://arxiv.org/abs/2105.06332
245 | namespace State
246 |   ||| State here differers for the one in `Cont`, because `Scalar` is different
247 |   |||
248 |   |||       ┌─────────────┐
249 |   |||       │             ├──► (x : c.Shp)
250 |   |||       │    State    │
251 |   |||       │             ├◄── c.Pos x
252 |   |||       └─────────────┘
253 |   public export
254 |   State : AddCont -> Type
255 |   State c = Scalar =%+> c
256 |
257 |   public export
258 |   toState : (x : c.Shp) -> State c
259 |   toState x = ?somethingInterestingHmm
260 |   
261 |   -- public export
262 |   -- fromState : State c -> c.Shp
263 |   -- fromState f = f.fwd ()
264 |
265 | ||| References for Costate
266 | ||| Bruno's PhD thesis: https://arxiv.org/abs/2403.13001
267 | ||| Towards Foundations of Cat. Cybernetics: https://arxiv.org/abs/2105.06332
268 | namespace Costate
269 |   ||| Costate here differs from the one in `Cont`, because `Scalar` is different
270 |   |||                  ┌─────────────┐
271 |   |||  (x : c.Shp)  ──►┤             │
272 |   |||                  │   Costate   │
273 |   |||     c.Pos x   ◄──┤             │
274 |   |||                  └─────────────┘
275 |   public export
276 |   Costate : AddCont -> Type
277 |   Costate c = c =%+> Scalar
278 |   
279 |   public export
280 |   toCostate : {c : AddCont} ->
281 |     ((x : c.Shp) -> c.PosSet x) -> Costate c
282 |   toCostate s = !%+ \x => (() ** \n => scale @{UMon c x} n (s x))
283 |
284 |   -- public export
285 |   -- fromCostate : Costate c -> (x : c.Shp) -> c.PosSet x
286 |   -- fromCostate f x = f.bwd x ()
287 |
288 |   public export
289 |   constantOne : {c : AddCont} ->
290 |     InterfaceOnPositions c Num => Costate c
291 |   constantOne @{MkI p} = toCostate (\x => let numPos = p x in 1)
292 |
293 |   public export
294 |   Delete : {c : AddCont} -> Costate c 
295 |   Delete = toCostate c.Zero
296 |   
297 |
298 | public export
299 | sum : Num a =>
300 |   (Const a >*< Const a) =%+> Const a
301 | sum = !%+ \(x1, x2) => (x1 + x2 ** \x' => (x', x'))
302 |
303 | public export
304 | bwSumBag : {l : Type} ->
305 |   (xs : List l) ->
306 |   (d' : l) ->
307 |   All (const l) xs
308 | bwSumBag [] d' = []
309 | bwSumBag (x :: xs) d' = d' :: bwSumBag xs d'
310 |
311 | public export
312 | sumBag : {l : Type} -> ComMonoid l =>
313 |   BagAll (Const l) =%+> Const l
314 | sumBag = !%+ \(MkBag xs) => (sum (MkBag xs) ** \d' => bwSumBag xs d')
315 |
316 | public export
317 | negate : Num a => Neg a =>
318 |   Const a =%+> Const a
319 | negate = !%+ \x => (-x ** \x' => -x')
320 |
321 | public export
322 | zero : {c : AddCont} -> Num a =>
323 |   c =%+> Const a
324 | zero = !%+ \_ => (0 ** \_ => c.Zero _)
325 |
326 | public export
327 | mul : Num a =>
328 |   (Const a >*< Const a) =%+> Const a
329 | mul = !%+ \(x1, x2) => (x1 * x2 ** \x' => (x' * x2, x' * x1))
330 |
331 | ||| Mean squared error
332 | public export
333 | SquaredDifference : {a : Type} -> Num a => Neg a =>
334 |   (Const a >*< Const a) =%+> (Const a)
335 | SquaredDifference = ((id {c=Const a}) >*< negate) %+>> sum %+>> copy %+>> mul
336 |
337 | namespace Sample
338 |   ||| Select a shape from All to produce an Any at the given index
339 |   ||| Same as `index i (allAnies shapes)` but reduces better
340 |   public export
341 |   selectShape : {cs : Vect k AddCont} ->
342 |     (shapes : All (.Shp) cs) -> (i : Fin k) -> Any (.Shp) cs
343 |   selectShape (s :: ss) FZ = Here s
344 |   selectShape (s :: ss) (FS j) = There (selectShape ss j)
345 |
346 |   ||| Extract the position from an AnyPos at a given index
347 |   public export
348 |   extractPos : {n : Nat} -> {xs : Vect n AddCont} -> {shapes : All (.Shp) xs} ->
349 |     (i : Fin n) ->
350 |     AnyShpPos (selectShape shapes i) ->
351 |     (index i xs).PosSet (index i shapes)
352 |   extractPos {shapes = (_ :: _)} FZ (Here x') = x'
353 |   extractPos {shapes = (_ :: _)} (FS j) (There rest) = extractPos j rest
354 |
355 |
356 |
357 | -- parameters (f : Type -> Type)
358 | --   ||| These are all of the morphisms in the cokleisli category of (f <!> -)  
359 | --   public export
360 | --   MonLens : Cont -> Cont -> Type
361 | --   MonLens c d = (f <!> c) =%> d
362 | -- 
363 | --   public export
364 | --   counit : Monad f => f <!> c =%> c
365 | --   counit = !% \x => (x ** pure)
366 | -- 
367 | --   public export
368 | --   cojoin : Monad f => (f <!> c) =%> (f <!> (f <!> c))
369 | --   cojoin = !% \x => (x ** join)
370 |
371 |   
372 | -- public export
373 | -- record FCoAlgCont (f : Type -> Type) where
374 | --   constructor MkFCoAlgCont
375 | --   carrier : Cont
376 | --   coalg : (a : carrier.Shp) -> f (carrier.Pos a) -> carrier.Pos a
377 |
378 | -- public export
379 | -- coAlgMorphism : (c, d : FCoAlgCont f) -> Type
380 | -- coAlgMorphism c d = c.carrier =%> d.carrier
381 | -- 
382 | -- convert : FCoAlgCont List -> AddCont
383 | -- convert (MkFCoAlgCont carrier coalg) = MkAddCont
384 | --   carrier
385 | --   {mon=(MkI $ \s => MkComMonoid
386 | --     (\l, r => coalg s [l, r])
387 | --     (coalg s []))}