{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE RebindableSyntax #-}
{-# LANGUAGE TypeFamilies #-}

-- | Coordinate charts and coordinate patches.
--
-- A 'Chart' is the differential-geometric notion of a homeomorphism from an
-- open subset of a manifold @m@ to an open subset of a model space of
-- coordinates @c@.  In this library the coordinate patch is represented by a
-- 'Space' (typically a 'Rect' for 2-D coordinates), so the existing
-- 'Space'/'Rect'/'Point' machinery becomes the chart infrastructure for free.
module NumHask.Space.Chart
  ( -- * Charts
    Chart (..),
    SomeSpace (..),
    chartMap,
    chartInverse,
    chartDomain,
    inverseChart,
    transition,

    -- * Common charts
    affineChart,
    polarChart,
  )
where

import NumHask.Prelude
import NumHask.Space.Point
import NumHask.Space.Rect
import NumHask.Space.Types

-- $setup
--
-- >>> :m -Prelude
-- >>> :set -XRebindableSyntax
-- >>> import NumHask.Prelude
-- >>> import NumHask.Space

-- | A coordinate chart on a manifold/model type @m@ with coordinates @c@.
--
-- The coordinate patch is any 'Space' whose elements have type @c@.  This is
-- usually a 'Rect' for 2-D coordinates @(Point a)@ or a 'Range' for 1-D
-- coordinates @a@.
data Chart m c where
  Chart ::
    (Space s, Element s ~ c) =>
    -- | coordinate-domain patch
    s ->
    -- | forward map: coordinates -> manifold point
    (c -> m) ->
    -- | inverse map: manifold point -> coordinates
    (m -> c) ->
    Chart m c

-- | An existential space whose elements have type @c@.
data SomeSpace c where
  SomeSpace :: (Space s, Element s ~ c) => s -> SomeSpace c

-- | Extract the forward map from a chart.
chartMap :: Chart m c -> c -> m
chartMap :: forall m c. Chart m c -> c -> m
chartMap (Chart s
_ c -> m
f m -> c
_) = c -> m
f

-- | Extract the inverse map from a chart.
chartInverse :: Chart m c -> m -> c
chartInverse :: forall m c. Chart m c -> m -> c
chartInverse (Chart s
_ c -> m
_ m -> c
g) = m -> c
g

-- | Extract the coordinate-domain patch from a chart.
chartDomain :: Chart m c -> SomeSpace c
chartDomain :: forall m c. Chart m c -> SomeSpace c
chartDomain (Chart s
d c -> m
_ m -> c
_) = s -> SomeSpace c
forall s c. (Space s, Element s ~ c) => s -> SomeSpace c
SomeSpace s
d

-- | Invert a chart, supplying the codomain patch for the inverse.
inverseChart :: (Space s, Element s ~ m) => s -> Chart m c -> Chart c m
inverseChart :: forall s m c.
(Space s, Element s ~ m) =>
s -> Chart m c -> Chart c m
inverseChart s
s (Chart s
_ c -> m
f m -> c
g) = s -> (m -> c) -> (c -> m) -> Chart c m
forall s c m.
(Space s, Element s ~ c) =>
s -> (c -> m) -> (m -> c) -> Chart m c
Chart s
s m -> c
g c -> m
f

-- | Transition map from the first chart to the second.
--
-- Both charts must cover (at least overlap on) the same manifold.  The result
-- is @φ₂ ∘ φ₁⁻¹@.
--
-- >>> let polar = polarChart (Rect 0.1 5 (-pi) pi)
-- >>> let cartesian = affineChart (Rect (-5) 5 (-5) 5) (Rect 0 1 0 1)
-- >>> transition cartesian polar (Point 1 0)
-- Point 0.6 0.5
transition :: Chart m c2 -> Chart m c1 -> c1 -> c2
transition :: forall m c2 c1. Chart m c2 -> Chart m c1 -> c1 -> c2
transition Chart m c2
c2 Chart m c1
c1 = Chart m c2 -> m -> c2
forall m c. Chart m c -> m -> c
chartInverse Chart m c2
c2 (m -> c2) -> (c1 -> m) -> c1 -> c2
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Chart m c1 -> c1 -> m
forall m c. Chart m c -> c -> m
chartMap Chart m c1
c1

-- | An affine chart between two rectangular patches.
--
-- >>> let c = affineChart (Rect 0 1 0 1) (Rect 0 2 0 2)
-- >>> chartMap c (Point 0.5 0.5)
-- Point 1.0 1.0
affineChart :: (Field a, Ord a) => Rect a -> Rect a -> Chart (Point a) (Point a)
affineChart :: forall a.
(Field a, Ord a) =>
Rect a -> Rect a -> Chart (Point a) (Point a)
affineChart Rect a
dom Rect a
cod = Rect a
-> (Point a -> Point a)
-> (Point a -> Point a)
-> Chart (Point a) (Point a)
forall s c m.
(Space s, Element s ~ c) =>
s -> (c -> m) -> (m -> c) -> Chart m c
Chart Rect a
dom (Rect a -> Rect a -> Element (Rect a) -> Element (Rect a)
forall s.
(Space s, Field (Element s)) =>
s -> s -> Element s -> Element s
project Rect a
dom Rect a
cod) (Rect a -> Rect a -> Element (Rect a) -> Element (Rect a)
forall s.
(Space s, Field (Element s)) =>
s -> s -> Element s -> Element s
project Rect a
cod Rect a
dom)

-- | Polar coordinates @(r, θ)@ to cartesian points @(x, y)@.
--
-- >>> let c = polarChart (Rect 0.1 5 (-pi) pi)
-- >>> chartMap c (Point 1 0)
-- Point 1.0 0.0
polarChart :: (TrigField a, ExpField a, Ord a) => Rect a -> Chart (Point a) (Point a)
polarChart :: forall a.
(TrigField a, ExpField a, Ord a) =>
Rect a -> Chart (Point a) (Point a)
polarChart Rect a
dom = Rect a
-> (Point a -> Point a)
-> (Point a -> Point a)
-> Chart (Point a) (Point a)
forall s c m.
(Space s, Element s ~ c) =>
s -> (c -> m) -> (m -> c) -> Chart m c
Chart Rect a
dom Point a -> Point a
forall {a}. TrigField a => Point a -> Point a
fwd Point a -> Point a
forall {a}. (ExpField a, TrigField a) => Point a -> Point a
back
  where
    fwd :: Point a -> Point a
fwd (Point a
r a
theta) = a -> a -> Point a
forall a. a -> a -> Point a
Point (a
r a -> a -> a
forall a. Multiplicative a => a -> a -> a
* a -> a
forall a. TrigField a => a -> a
cos a
theta) (a
r a -> a -> a
forall a. Multiplicative a => a -> a -> a
* a -> a
forall a. TrigField a => a -> a
sin a
theta)
    back :: Point a -> Point a
back (Point a
x a
y) = a -> a -> Point a
forall a. a -> a -> Point a
Point (a -> a
forall a. ExpField a => a -> a
sqrt (a
x a -> a -> a
forall a. Multiplicative a => a -> a -> a
* a
x a -> a -> a
forall a. Additive a => a -> a -> a
+ a
y a -> a -> a
forall a. Multiplicative a => a -> a -> a
* a
y)) (a -> a -> a
forall a. TrigField a => a -> a -> a
atan2 a
y a
x)