
   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.Num
  10 | import Data.Container.Additive.Object.Definition
  11 | import Data.Container.Additive.Object.Instances
  12 | import Data.Container.Additive.Extension.Definition
  13 | import Data.Container.Additive.Morphism.Definition
  14 | import Data.Container.Additive.Product.Definitions
  15 | import Data.Container.Additive.Properties.Definitions
  16 | 
  17 | import Data.Container.Additive.Quantifiers
  18 | 
  19 | import Control.Monad.Distribution
  20 | import Control.Monad.Sample.Definition
  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 | 
  30 | ||| If we model the idea of a container (S !> P) as a box
  31 | |||  ┌──────┐
  32 | |||  │ s:S  │
  33 | |||  ├──────┤
  34 | |||  │  Ps  │
  35 | |||  └──────┘
  36 | ||| then `pushDown` is interpreted as pushing down the container,
  37 | ||| pruning anything that goes out of the box, and using `Unit` for
  38 | ||| anything new that appears:
  39 | |||  ┌──────┐
  40 | |||  │ Unit │
  41 | |||  ├──────┤
  42 | |||  │ s:S  │
  43 | |||  └──────┘
  44 | |||     Ps
  45 | ||| For additive containers we need to take the free commutative monoid
  46 | public export
  47 | pushDown : AddCont -> AddCont
  48 | pushDown c = !* pushDown (UC c)
  49 | 
  50 | public export
  51 | pushIntoContinuationBag : {p : AddCont} -> {0 d, l : AddCont} ->
  52 |   d >< p =%+> l ->
  53 |   p =%+> (pushDown d) >+@ (Bag l)
  54 | pushIntoContinuationBag f = !%+ \param => (() <|
  55 |   map (\dShp => f.fwd (dShp, param)) **
  56 |     \ll => sum @{UMon p param} $ ll >>=
  57 |       \(ds ** grads) => extractPGradsBag param ds grads)
  58 |   where
  59 |     extractPGrads : (param : p.Shp) ->
  60 |       (ds : List d.Shp) ->
  61 |       All l.Pos ((\dShp => f.fwd (dShp, param)) <$> ds) ->
  62 |       List (p.Pos param)
  63 |     extractPGrads param [] [] = []
  64 |     extractPGrads param (dShp :: ds) (grad :: grads) =
  65 |       snd (f.bwd (dShp, param) grad) :: extractPGrads param ds grads
  66 | 
  67 |     extractPGradsBag : (param : p.Shp) ->
  68 |       (ds : Bag d.Shp) ->
  69 |       All l.Pos ((\dShp => f.fwd (dShp, param)) <$> ds) ->
  70 |       Bag (p.Pos param)
  71 |     extractPGradsBag param (MkBag dsl) grads
  72 |       = MkBag $ extractPGrads param dsl grads
  73 | 
  74 | 
  75 | public export
  76 | pushIntoContinuation : {p : AddCont} -> (flat : IsFlat l) => Num l.Shp =>
  77 |   (f : d >< p =%+> l) ->
  78 |   (p =%+> (pushDown d) >+@ l)
  79 | pushIntoContinuation {flat = MkIsFlat lp} f = !%+ \param => (() <|
  80 |   \ds => sum @{numIsMonoid} ((\dShp => f.fwd (dShp, param)) <$> ds) **
  81 |     \bb => sum @{UMon p param} (bb >>=
  82 |       \(ds ** grad) => ds <&> (\dShp => snd (f.bwd (dShp, param) grad))))
  83 | 
  84 | ||| This is also the categorical product since our containers are additive
  85 | namespace HancockTensorProduct
  86 |   public export
  87 |   leftUnit : Scalar >< c =%+> c
  88 |   leftUnit = !% leftUnit
  89 |   
  90 |   public export
  91 |   rightUnit : c >< Scalar =%+> c
  92 |   rightUnit = !% rightUnit
  93 | 
  94 |   public export
  95 |   leftUnitInv : c =%+> Scalar >< c
  96 |   leftUnitInv = !% leftUnitInv
  97 |   
  98 |   public export
  99 |   rightUnitInv : c =%+> c >< Scalar
 100 |   rightUnitInv = !% rightUnitInv
 101 | 
 102 |   public export
 103 |   assocL : (a >< b) >< c =%+> a >< (b >< c)
 104 |   assocL = !% assocL
 105 | 
 106 |   public export
 107 |   assocR : a >< (b >< c) =%+> (a >< b) >< c
 108 |   assocR = !% assocR
 109 | 
 110 |   public export
 111 |   swap : a >< b =%+> b >< a
 112 |   swap = !% swap
 113 | 
 114 |   public export
 115 |   swapMiddle : (c1 >< c2) >< (c3 >< c4) =%+> (c1 >< c3) >< (c2 >< c4)
 116 |   swapMiddle = !% swapMiddle
 117 | 
 118 |   ||| These do not exist for ordinary containers!
 119 |   ||| Here we need `c` not to be erased since we're using its monoid structure
 120 |   public export
 121 |   copy : {c : AddCont} -> c =%+> c >< c
 122 |   copy = !%+ \x => ((x, x) ** uncurry (c.Plus x))
 123 |   
 124 |   public export
 125 |   pairMaps : {c : AddCont} ->
 126 |     c =%+> d ->
 127 |     c =%+> e ->
 128 |     c =%+> d >< e
 129 |   pairMaps f g = copy %+>> (f >< g)
 130 |   
 131 |   public export
 132 |   projLeft : {d : AddCont} -> c >< d =%+> c
 133 |   projLeft = !%+ \(x, y) => (x ** \x' => (x', d.Zero y))
 134 |   
 135 |   public export
 136 |   projRight : {c : AddCont} -> c >< d =%+> d
 137 |   projRight = !%+ \(x, y) => (y ** \y' => (c.Zero x, y'))
 138 | 
 139 | namespace CompositionProduct
 140 |   public export
 141 |   leftUnit : Scalar >+@ c =%+> c
 142 |   leftUnit = !% pureBw %>> leftUnit
 143 | 
 144 |   public export
 145 |   rightUnit : c >+@ Scalar =%+> c
 146 |   rightUnit = !% pureBw %>> rightUnit
 147 | 
 148 |   ||| Left unit inverse: c =%+> Scalar >+@ c
 149 |   public export
 150 |   leftUnitInv : {c : AddCont} -> c =%+> Scalar >+@ c
 151 |   leftUnitInv = !% sumBw @{mon c} %>> (Bag <!> leftUnitInv)
 152 | 
 153 |   ||| Right unit inverse: c =%+> c >@ I
 154 |   public export
 155 |   rightUnitInv : {c : AddCont} -> c =%+> c >+@ Scalar
 156 |   rightUnitInv = !% sumBw @{mon c} %>> (Bag <!> rightUnitInv)
 157 | 
 158 |   public export
 159 |   assocL : (a >+@ b) >+@ c =%+> a >+@ (b >+@ c)
 160 |   assocL = !%+ \((aShp <| f) <| g) =>
 161 |     (aShp <| \aPos => f aPos <| \bPos => g (MkBag [(aPos ** bPos)]) **
 162 |       \ll => join $ ll <&> \(aPos ** lbc) =>
 163 |         lbc <&> \(bPos ** cPos) => (MkBag [(aPos ** bPos)] ** cPos))
 164 | 
 165 |   ||| Associator, "un-flatten" direction. NOT definable as a total lens in
 166 |   ||| general: the forward would have to produce the target's outer index
 167 |   ||| `g : List (aPos ** bPos) -> c.Shp`, i.e. collapse a whole list of
 168 |   ||| (a,b)-positions into a single c-shape. All we have is one c-shape per
 169 |   ||| element (`index (f aPos) bPos`), and c-shapes carry no monoid/default,
 170 |   ||| so the empty-list case has no answer. This is the precise sense in which
 171 |   ||| the free composition product is only laxly (one-directionally) associative.
 172 |   public export
 173 |   assocR : a >+@ (b >+@ c) =%+> (a >+@ b) >+@ c
 174 |   assocR = !%+ \(aShp <| f) => (((aShp <| shapeExt . f) <|
 175 |     \ll => let ff : (aPos : a.Pos aShp ** b.Pos (Ext.shapeExt $ f aPos)) -> c.Shp 
 176 |                ff (aPos ** bPos) = index (f aPos) bPos
 177 |            in ?llb) ** ?fififi)
 178 | 
 179 | 
 180 | namespace Coproduct
 181 |   public export
 182 |   elim : c >+< c =%+> c
 183 |   elim = !% elim
 184 | 
 185 | public export
 186 | duoidal : (c >+@ d) >< (e >+@ f) =%+> (c >< e) >+@ (d >< f)
 187 | duoidal = !%+ \((sc <| idxC), (se <| idxE)) =>
 188 |   ((sc, se) <| \(cp, ep) => (idxC cp, idxE ep) **
 189 |     \ll => ((\((cp, ep) ** (dp, fp)) => (cp ** dp)) <$> ll,
 190 |             (\((cp, ep) ** (dp, fp)) => (ep ** fp)) <$> ll))
 191 | 
 192 | 
 193 | public export
 194 | coprodDistrOverTensor : {q, p : AddCont} ->
 195 |   (a >+< b) >< (p >< q) =%+> (a >< p) >+< (b >< q)
 196 | coprodDistrOverTensor = !%+ \case
 197 |   (Left a, (p, _)) => (Left (a, p) ** \(a', p') => (a', (p', q.Zero _)))
 198 |   (Right b, (_, q)) => (Right (b, q) ** \(b', q') => (b', (p.Zero _, q')))
 199 | 
 200 | ||| Not an isomorphism, arising from duoidal structure between >@ and ><
 201 | public export
 202 | rebracketcomptensor: {y : AddCont} -> (e >+@ y) >< y =%+> e >+@ (y >< y)
 203 | rebracketcomptensor = (id {c=e >+@ y} >< leftUnitInv {c=y})
 204 |                       %+>> duoidal {c=e} {d=y} {e=Scalar} {f=y}
 205 |                       %+>> (rightUnit {c=e} >+@ id {c=(y><y)})
 206 | 
 207 | 
 208 | public export
 209 | distribute : {c : AddCont} ->
 210 |   c >< e =%+> s ->
 211 |   c >< (e >+@ g) =%+> s >+@ g
 212 | distribute f = (rightUnitInv >< id {c=e >+@ g})
 213 |              %+>> duoidal {d = Scalar}
 214 |              %+>> (f >+@ leftUnit)
 215 | 
 216 | public export
 217 | extractEffect : {d : AddCont} ->
 218 |   d >< (e >+@ f) =%+> e >+@ (d >< f)
 219 | extractEffect = (leftUnitInv >< (id {c=e >+@ f}))
 220 |             %+>> duoidal {c=Scalar}
 221 |             %+>> (leftUnit >+@ (id {c=d><f}))
 222 | 
 223 | 
 224 | namespace State
 225 |   ||| "State" as defined in https://arxiv.org/abs/2403.13001 and open games 
 226 |   |||
 227 |   |||       ┌─────────────┐
 228 |   |||       │             ├──► (x : c.Shp)
 229 |   |||       │    State    │
 230 |   |||       │             ├◄── c.Pos x
 231 |   |||       └─────────────┘
 232 |   public export
 233 |   State : AddCont -> Type
 234 |   State c = Scalar =%+> c
 235 | 
 236 |   public export
 237 |   toState : (x : c.Shp) -> State c
 238 |   toState x = !% toState x
 239 |   
 240 |   public export
 241 |   fromState : State c -> c.Shp
 242 |   fromState f = f.fwd ()
 243 | 
 244 | namespace Costate
 245 |   ||| "Costate" as defined in https://arxiv.org/abs/2403.13001 and open games 
 246 |   |||
 247 |   |||                  ┌─────────────┐
 248 |   |||  (x : c.Shp)  ──►┤             │
 249 |   |||                  │   Costate   │
 250 |   |||     c.Pos x   ◄──┤             │
 251 |   |||                  └─────────────┘
 252 |   public export
 253 |   Costate : AddCont -> Type
 254 |   Costate c = c =%+> Scalar
 255 |   
 256 |   public export
 257 |   toCostate : ((x : c.Shp) -> c.Pos x) -> Costate c
 258 |   toCostate s = !% toCostate s
 259 |   
 260 |   public export
 261 |   fromCostate : Costate c -> (x : c.Shp) -> c.Pos x
 262 |   fromCostate f x = f.bwd x ()
 263 | 
 264 |   public export
 265 |   constantOne : InterfaceOnPositions c Num => Costate c
 266 |   constantOne @{MkI p} = toCostate (\x => let numPos = p x in 1)
 267 | 
 268 |   public export
 269 |   Delete : {c : AddCont} -> Costate c 
 270 |   Delete = toCostate c.Zero
 271 |   
 272 | 
 273 |   
 274 | public export
 275 | sum : Num a =>
 276 |   (Const a >< Const a) =%+> Const a
 277 | sum = !%+ \(x1, x2) => (x1 + x2 ** \x' => (x', x'))
 278 | 
 279 | public export
 280 | bwSumList : {l : Type} -> ComMonoid l =>
 281 |   (xs : List l) ->
 282 |   (d' : l) ->
 283 |   All (const l) xs
 284 | bwSumList [] d' = []
 285 | bwSumList (x :: xs) d' = x :: bwSumList xs x
 286 | 
 287 | public export
 288 | bwSumBag : {l : Type} -> ComMonoid l =>
 289 |   (xs : Bag l) ->
 290 |   (d' : l) ->
 291 |   All (const l) xs
 292 | bwSumBag (MkBag xs) d' = bwSumList xs d'
 293 | 
 294 | 
 295 | public export
 296 | sumList : {l : Type} -> ComMonoid l =>
 297 |   Bag (Const l) =%+> Const l
 298 | sumList = !%+ \xs => (sum xs ** \d' => bwSumBag xs d')
 299 | 
 300 | public export
 301 | negate : Num a => Neg a =>
 302 |   Const a =%+> Const a
 303 | negate = !%+ \x => (-x ** \x' => -x')
 304 | 
 305 | public export
 306 | zero : {c : AddCont} -> Num a =>
 307 |   c =%+> Const a
 308 | zero = !%+ \_ => (0 ** \_ => c.Zero _)
 309 | 
 310 | public export
 311 | mul : Num a =>
 312 |   (Const a >< Const a) =%+> Const a
 313 | mul = !%+ \(x1, x2) => (x1 * x2 ** \x' => (x' * x2, x' * x1))
 314 | 
 315 | ||| Mean squared error
 316 | public export
 317 | SquaredDifference : {a : Type} -> Num a => Neg a =>
 318 |   (Const a >< Const a) =%+> (Const a)
 319 | SquaredDifference = ((id {c=Const a}) >< negate) %+>> sum %+>> copy %+>> mul
 320 | 
 321 | namespace Sample
 322 |   ||| Select a shape from All to produce an Any at the given index
 323 |   ||| Same as `index i (allAnies shapes)` but reduces better
 324 |   public export
 325 |   selectShape : {cs : Vect k AddCont} ->
 326 |     (shapes : All (.Shp) cs) -> (i : Fin k) -> Any (.Shp) cs
 327 |   selectShape (s :: ss) FZ = Here s
 328 |   selectShape (s :: ss) (FS j) = There (selectShape ss j)
 329 | 
 330 |   ||| Extract the position from an AnyPos at a given index
 331 |   public export
 332 |   extractPos : {n : Nat} -> {xs : Vect n AddCont} -> {shapes : All (.Shp) xs} ->
 333 |     (i : Fin n) ->
 334 |     AnyShpPos (selectShape shapes i) ->
 335 |     (index i xs).Pos (index i shapes)
 336 |   extractPos {shapes = (_ :: _)} FZ (Here x') = x'
 337 |   extractPos {shapes = (_ :: _)} (FS j) (There rest) = extractPos j rest
 338 | 
 339 | -- parameters (f : Type -> Type)
 340 | --   ||| These are all of the morphisms in the cokleisli category of (f <!> -)  
 341 | --   public export
 342 | --   MonLens : Cont -> Cont -> Type
 343 | --   MonLens c d = (f <!> c) =%> d
 344 | -- 
 345 | --   public export
 346 | --   counit : Monad f => f <!> c =%> c
 347 | --   counit = !% \x => (x ** pure)
 348 | -- 
 349 | --   public export
 350 | --   cojoin : Monad f => (f <!> c) =%> (f <!> (f <!> c))
 351 | --   cojoin = !% \x => (x ** join)
 352 | 
 353 |   
 354 | -- public export
 355 | -- record FCoAlgCont (f : Type -> Type) where
 356 | --   constructor MkFCoAlgCont
 357 | --   carrier : Cont
 358 | --   coalg : (a : carrier.Shp) -> f (carrier.Pos a) -> carrier.Pos a
 359 | 
 360 | -- public export
 361 | -- coAlgMorphism : (c, d : FCoAlgCont f) -> Type
 362 | -- coAlgMorphism c d = c.carrier =%> d.carrier
 363 | -- 
 364 | -- convert : FCoAlgCont List -> AddCont
 365 | -- convert (MkFCoAlgCont carrier coalg) = MkAddCont
 366 | --   carrier
 367 | --   {mon=(MkI $ \s => MkComMonoid
 368 | --     (\l, r => coalg s [l, r])
 369 | --     (coalg s []))}