module Data.Function.Affine (
Affine(..)
, evalAffine
, fromOffsetSlope
, lensEmbedding, correspondingDirections
) where
import Data.Semigroup
import Data.MemoTrie
import Data.VectorSpace
import Data.AffineSpace
import Data.Tagged
import Data.Manifold.Types.Primitive
import Data.Manifold.PseudoAffine
import Data.Manifold.Atlas
import Data.Embedding
import qualified Prelude
import qualified Control.Applicative as Hask
import Data.Functor (($>))
import Control.Category.Constrained.Prelude hiding ((^))
import Control.Category.Constrained.Reified
import Control.Arrow.Constrained
import Control.Monad.Constrained
import Data.Foldable.Constrained
import Math.LinearMap.Category
import Control.Lens
data Affine s d c where
Affine :: (ChartIndex d :->: (c, LinearMap s (Needle d) (Needle c)))
-> Affine s d c
instance Category (Affine s) where
type Object (Affine s) x = ( Manifold x, Interior x ~ x
, Atlas x, LinearSpace (Needle x)
, Scalar (Needle x) ~ s, HasTrie (ChartIndex x) )
id = Affine . trie $ chartReferencePoint >>> id &&& const id
Affine f . Affine g = Affine . trie
$ \ixa -> case untrie g ixa of
(b, ða'b) -> case untrie f $ lookupAtlas b of
(c, ðb'c) -> (c, ðb'c . ða'b)
instance ∀ s . Num' s => Cartesian (Affine s) where
type UnitObject (Affine s) = ZeroDim s
swap = Affine . trie $ chartReferencePoint >>> swap &&& const swap
attachUnit = Affine . trie $ chartReferencePoint >>> \a -> ((a,Origin), attachUnit)
detachUnit = Affine . trie $ chartReferencePoint
>>> \(a,Origin::ZeroDim s) -> (a, detachUnit)
regroup = Affine . trie $ chartReferencePoint >>> regroup &&& const regroup
regroup' = Affine . trie $ chartReferencePoint >>> regroup' &&& const regroup'
instance ∀ s . Num' s => Morphism (Affine s) where
Affine f *** Affine g = Affine . trie
$ \(ixα,ixβ) -> case (untrie f ixα, untrie g ixβ) of
((fα, ðα'f), (gβ,ðβ'g)) -> ((fα,gβ), ðα'f***ðβ'g)
instance ∀ s . Num' s => PreArrow (Affine s) where
Affine f &&& Affine g = Affine . trie
$ \ix -> case (untrie f ix, untrie g ix) of
((fα, ðα'f), (gβ,ðβ'g)) -> ((fα,gβ), ðα'f&&&ðβ'g)
terminal = Affine . trie $ \_ -> (Origin, zeroV)
fst = afst
where afst :: ∀ x y . ( Atlas x, Atlas y
, LinearSpace (Needle x), LinearSpace (Needle y)
, Scalar (Needle x) ~ s, Scalar (Needle y) ~ s
, HasTrie (ChartIndex x), HasTrie (ChartIndex y) )
=> Affine s (x,y) x
afst = Affine . trie $ chartReferencePoint >>> \(x,_::y) -> (x, fst)
snd = asnd
where asnd :: ∀ x y . ( Atlas x, Atlas y
, LinearSpace (Needle x), LinearSpace (Needle y)
, Scalar (Needle x) ~ s, Scalar (Needle y) ~ s
, HasTrie (ChartIndex x), HasTrie (ChartIndex y) )
=> Affine s (x,y) y
asnd = Affine . trie $ chartReferencePoint >>> \(_::x,y) -> (y, snd)
instance ∀ s . Num' s => WellPointed (Affine s) where
const x = Affine . trie $ const (x, zeroV)
unit = Tagged Origin
instance EnhancedCat (->) (Affine s) where
arr f = fst . evalAffine f
instance EnhancedCat (Affine s) (LinearMap s) where
arr = alarr (linearManifoldWitness, linearManifoldWitness)
where alarr :: ∀ x y . ( LinearSpace x, Atlas x, HasTrie (ChartIndex x)
, LinearSpace y
, Scalar x ~ s, Scalar y ~ s )
=> (LinearManifoldWitness x, LinearManifoldWitness y)
-> LinearMap s x y -> Affine s x y
alarr (LinearManifoldWitness _, LinearManifoldWitness _) f
= Affine . trie $ chartReferencePoint
>>> \x₀ -> let y₀ = f $ x₀
in (negateV y₀, f)
instance ( Atlas x, HasTrie (ChartIndex x), LinearSpace (Needle x), Scalar (Needle x) ~ s
, Manifold y, Scalar (Needle y) ~ s )
=> Semimanifold (Affine s x y) where
type Needle (Affine s x y) = Affine s x (Needle y)
toInterior = pure
fromInterior = id
(.+~^) = case ( semimanifoldWitness :: SemimanifoldWitness y
, boundarylessWitness :: BoundarylessWitness y ) of
(SemimanifoldWitness _, BoundarylessWitness) -> \(Affine f) (Affine g)
-> Affine . trie $ \ix -> case (untrie f ix, untrie g ix) of
((fx₀,f'), (gx₀,g')) -> (fx₀.+~^gx₀, f'^+^g')
translateP = Tagged (.+~^)
semimanifoldWitness = case semimanifoldWitness :: SemimanifoldWitness y of
SemimanifoldWitness _ -> SemimanifoldWitness BoundarylessWitness
instance ( Atlas x, HasTrie (ChartIndex x), LinearSpace (Needle x), Scalar (Needle x) ~ s
, Manifold y, Scalar (Needle y) ~ s )
=> PseudoAffine (Affine s x y) where
(.-~!) = case ( semimanifoldWitness :: SemimanifoldWitness y
, boundarylessWitness :: BoundarylessWitness y ) of
(SemimanifoldWitness _, BoundarylessWitness) -> \(Affine f) (Affine g)
-> Affine . trie $ \ix -> case (untrie f ix, untrie g ix) of
((fx₀,f'), (gx₀,g')) -> (fx₀.-~!gx₀, f'^-^g')
pseudoAffineWitness = case semimanifoldWitness :: SemimanifoldWitness y of
SemimanifoldWitness _ -> PseudoAffineWitness (SemimanifoldWitness BoundarylessWitness)
instance ( Atlas x, HasTrie (ChartIndex x), LinearSpace (Needle x), Scalar (Needle x) ~ s
, Manifold y, Scalar (Needle y) ~ s )
=> AffineSpace (Affine s x y) where
type Diff (Affine s x y) = Affine s x (Needle y)
(.+^) = (.+~^); (.-.) = (.-~!)
instance ( Atlas x, HasTrie (ChartIndex x), LinearSpace (Needle x), Scalar (Needle x) ~ s
, LinearSpace y, Scalar y ~ s, Num' s )
=> AdditiveGroup (Affine s x y) where
zeroV = case linearManifoldWitness :: LinearManifoldWitness y of
LinearManifoldWitness _ -> Affine . trie $ const (zeroV, zeroV)
(^+^) = case ( linearManifoldWitness :: LinearManifoldWitness y
, dualSpaceWitness :: DualSpaceWitness y ) of
(LinearManifoldWitness BoundarylessWitness, DualSpaceWitness) -> (.+~^)
negateV = case linearManifoldWitness :: LinearManifoldWitness y of
LinearManifoldWitness _ -> \(Affine f) -> Affine . trie $
untrie f >>> negateV***negateV
instance ( Atlas x, HasTrie (ChartIndex x), LinearSpace (Needle x), Scalar (Needle x) ~ s
, LinearSpace y, Scalar y ~ s, Num' s )
=> VectorSpace (Affine s x y) where
type Scalar (Affine s x y) = s
(*^) = case linearManifoldWitness :: LinearManifoldWitness y of
LinearManifoldWitness _ -> \μ (Affine f) -> Affine . trie $
untrie f >>> (μ*^)***(μ*^)
evalAffine :: ∀ s x y . ( Manifold x, Atlas x, HasTrie (ChartIndex x)
, Manifold y
, s ~ Scalar (Needle x), s ~ Scalar (Needle y) )
=> Affine s x y -> x -> (y, LinearMap s (Needle x) (Needle y))
evalAffine = ea (boundarylessWitness, boundarylessWitness)
where ea :: (BoundarylessWitness x, BoundarylessWitness y)
-> Affine s x y -> x -> (y, LinearMap s (Needle x) (Needle y))
ea (BoundarylessWitness, BoundarylessWitness)
(Affine f) x = (fx₀.+~^(ðx'f $ v), ðx'f)
where Just v = x .-~. chartReferencePoint chIx
chIx = lookupAtlas x
(fx₀, ðx'f) = untrie f chIx
fromOffsetSlope :: ∀ s x y . ( LinearSpace x, Atlas x, HasTrie (ChartIndex x)
, Manifold y
, s ~ Scalar x, s ~ Scalar (Needle y) )
=> y -> LinearMap s x (Needle y) -> Affine s x y
fromOffsetSlope = case ( linearManifoldWitness :: LinearManifoldWitness x
, boundarylessWitness :: BoundarylessWitness y ) of
(LinearManifoldWitness _, BoundarylessWitness)
-> \y0 ðx'y -> Affine . trie $ chartReferencePoint
>>> \x₀ -> let δy = ðx'y $ x₀
in (y0.+~^δy, ðx'y)
instance EnhancedCat (Embedding (Affine s)) (Embedding (LinearMap s)) where
arr (Embedding e p) = Embedding (arr e) (arr p)
lensEmbedding :: ∀ k s x c .
( Num' s
, LinearSpace x, LinearSpace c, Object k x, Object k c
, Scalar x ~ s, Scalar c ~ s
, EnhancedCat k (LinearMap s) )
=> Lens' x c -> Embedding k c x
lensEmbedding l = Embedding (arr $ (arr $ LinearFunction (\c -> zeroV & l .~ c)
:: LinearMap s c x) )
(arr $ (arr $ LinearFunction (^.l)
:: LinearMap s x c) )
correspondingDirections :: ∀ s x c t
. ( WithField s AffineManifold c
, WithField s AffineManifold x
, SemiInner (Needle c), SemiInner (Needle x)
, RealFrac' s
, Traversable t )
=> (Interior c, Interior x)
-> t (Needle c, Needle x) -> Maybe (Embedding (Affine s) c x)
correspondingDirections (c₀, x₀) dirMap
= freeEmbeddings $> Embedding (Affine . trie $ c2x boundarylessWitness)
(Affine . trie $ x2c boundarylessWitness)
where freeEmbeddings = fzip ( embedFreeSubspace $ fst<$>dirMap
, embedFreeSubspace $ snd<$>dirMap )
c2t :: Lens' (Needle c) (t s)
c2t = case freeEmbeddings of Just (Lens ct, _) -> ct
x2t :: Lens' (Needle x) (t s)
x2t = case freeEmbeddings of Just (_, Lens xt) -> xt
c2x :: BoundarylessWitness c -> ChartIndex c
-> (x, LinearMap s (Needle c) (Needle x))
c2x BoundarylessWitness ιc
= ( x₀ .+~^ (zeroV & x2t .~ δc^.c2t)
, arr . LinearFunction $ \dc -> zeroV & x2t .~ dc^.c2t )
where Just δc = chartReferencePoint ιc .-~. c₀
x2c :: BoundarylessWitness x -> ChartIndex x
-> (c, LinearMap s (Needle x) (Needle c))
x2c BoundarylessWitness ιx
= ( c₀ .+~^ (zeroV & c2t .~ δx^.x2t)
, arr . LinearFunction $ \dx -> zeroV & c2t .~ dx^.x2t )
where Just δx = chartReferencePoint ιx .-~. x₀