{-# OPTIONS_GHC -Wall #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE Safe #-}

{- | 
Module      :  Physics.Learn.Mechanics
Copyright   :  (c) Scott N. Walck 2014-2019
License     :  BSD3 (see LICENSE)
Maintainer  :  Scott N. Walck <walck@lvc.edu>
Stability   :  experimental

Newton's second law and all that
-}

module Physics.Learn.Mechanics
    ( TheTime
    , TimeStep
    , Velocity
    -- * Simple one-particle state
    , SimpleState
    , SimpleAccelerationFunction
    , simpleStateDeriv
    , simpleRungeKuttaStep
    -- * One-particle state
    , St(..)
    , DSt(..)
    , OneParticleSystemState
    , OneParticleAccelerationFunction
    , oneParticleStateDeriv
    , oneParticleRungeKuttaStep
    , oneParticleRungeKuttaSolution
    -- * Two-particle state
    , TwoParticleSystemState
    , TwoParticleAccelerationFunction
    , twoParticleStateDeriv
    , twoParticleRungeKuttaStep
    -- * Many-particle state
    , ManyParticleSystemState
    , ManyParticleAccelerationFunction
    , manyParticleStateDeriv
    , manyParticleRungeKuttaStep
    )
    where

import Data.VectorSpace
    ( AdditiveGroup(..)
    , VectorSpace(..)
    )
import Physics.Learn.StateSpace
    ( StateSpace(..)
    , Diff
    , DifferentialEquation
    )
import Physics.Learn.RungeKutta
    ( rungeKutta4
    , integrateSystem
    )
import Physics.Learn.Position
    ( Position
    )
import Physics.Learn.CarrotVec
    ( Vec
    )

-- | Time (in s).
type TheTime = Double

-- | A time step (in s).
type TimeStep = Double

-- | Velocity of a particle (in m/s).
type Velocity = Vec

-------------------------------
-- Simple one-particle state --
-------------------------------

-- | A simple one-particle state,
--   to get started quickly with mechanics of one particle.
type SimpleState = (TheTime,Position,Velocity)

-- | An acceleration function gives the particle's acceleration as
--   a function of the particle's state.
--   The specification of this function is what makes one single-particle
--   mechanics problem different from another.
--   In order to write this function, add all of the forces
--   that act on the particle, and divide this net force by the particle's mass.
--   (Newton's second law).
type SimpleAccelerationFunction = SimpleState -> Vec

-- | Time derivative of state for a single particle
--   with a constant mass.
simpleStateDeriv :: SimpleAccelerationFunction  -- ^ acceleration function for the particle
                 -> DifferentialEquation SimpleState  -- ^ differential equation
simpleStateDeriv :: SimpleAccelerationFunction -> DifferentialEquation SimpleState
simpleStateDeriv SimpleAccelerationFunction
a (Double
t, Position
r, Vec
v) = (Double
1, Vec
v, SimpleAccelerationFunction
a(Double
t, Position
r, Vec
v))

-- | Single Runge-Kutta step
simpleRungeKuttaStep :: SimpleAccelerationFunction  -- ^ acceleration function for the particle
                     -> TimeStep  -- ^ time step
                     -> SimpleState  -- ^ initial state
                     -> SimpleState  -- ^ state after one time step
simpleRungeKuttaStep :: SimpleAccelerationFunction -> Double -> SimpleState -> SimpleState
simpleRungeKuttaStep = forall p. StateSpace p => (p -> Diff p) -> Time p -> p -> p
rungeKutta4 forall b c a. (b -> c) -> (a -> b) -> a -> c
. SimpleAccelerationFunction -> DifferentialEquation SimpleState
simpleStateDeriv

------------------------
-- One-particle state --
------------------------

-- | The state of a single particle is given by
--   the position of the particle and the velocity of the particle.
data St = St { St -> Position
position :: Position
             , St -> Vec
velocity :: Velocity
             }
          deriving (Int -> St -> ShowS
[St] -> ShowS
St -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [St] -> ShowS
$cshowList :: [St] -> ShowS
show :: St -> String
$cshow :: St -> String
showsPrec :: Int -> St -> ShowS
$cshowsPrec :: Int -> St -> ShowS
Show)

-- | The associated vector space for the
--   state of a single particle.
data DSt = DSt Vec Vec
           deriving (Int -> DSt -> ShowS
[DSt] -> ShowS
DSt -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [DSt] -> ShowS
$cshowList :: [DSt] -> ShowS
show :: DSt -> String
$cshow :: DSt -> String
showsPrec :: Int -> DSt -> ShowS
$cshowsPrec :: Int -> DSt -> ShowS
Show)

instance AdditiveGroup DSt where
    zeroV :: DSt
zeroV = Vec -> Vec -> DSt
DSt forall v. AdditiveGroup v => v
zeroV forall v. AdditiveGroup v => v
zeroV
    negateV :: DSt -> DSt
negateV (DSt Vec
dr Vec
dv) = Vec -> Vec -> DSt
DSt (forall v. AdditiveGroup v => v -> v
negateV Vec
dr) (forall v. AdditiveGroup v => v -> v
negateV Vec
dv)
    DSt Vec
dr1 Vec
dv1 ^+^ :: DSt -> DSt -> DSt
^+^ DSt Vec
dr2 Vec
dv2 = Vec -> Vec -> DSt
DSt (Vec
dr1 forall v. AdditiveGroup v => v -> v -> v
^+^ Vec
dr2) (Vec
dv1 forall v. AdditiveGroup v => v -> v -> v
^+^ Vec
dv2)

instance VectorSpace DSt where
    type Scalar DSt = Double
    Scalar DSt
c *^ :: Scalar DSt -> DSt -> DSt
*^ DSt Vec
dr Vec
dv = Vec -> Vec -> DSt
DSt (Scalar DSt
cforall v. VectorSpace v => Scalar v -> v -> v
*^Vec
dr) (Scalar DSt
cforall v. VectorSpace v => Scalar v -> v -> v
*^Vec
dv)

instance StateSpace St where
    type Diff St = DSt
    St Position
r1 Vec
v1 .-. :: St -> St -> Diff St
.-. St Position
r2 Vec
v2  = Vec -> Vec -> DSt
DSt (Position
r1 forall p. StateSpace p => p -> p -> Diff p
.-. Position
r2) (Vec
v1 forall p. StateSpace p => p -> p -> Diff p
.-. Vec
v2)
    St Position
r1 Vec
v1 .+^ :: St -> Diff St -> St
.+^ DSt Vec
dr Vec
dv = Position -> Vec -> St
St (Position
r1 forall p. StateSpace p => p -> Diff p -> p
.+^ Vec
dr) (Vec
v1 forall p. StateSpace p => p -> Diff p -> p
.+^ Vec
dv)

-- | The state of a system of one particle is given by the current time,
--   the position of the particle, and the velocity of the particle.
--   Including time in the state like this allows us to
--   have time-dependent forces.
type OneParticleSystemState = (TheTime,St)

-- | An acceleration function gives the particle's acceleration as
--   a function of the particle's state.
type OneParticleAccelerationFunction = OneParticleSystemState -> Vec

-- | Time derivative of state for a single particle
--   with a constant mass.
oneParticleStateDeriv :: OneParticleAccelerationFunction  -- ^ acceleration function for the particle
                      -> DifferentialEquation OneParticleSystemState  -- ^ differential equation
oneParticleStateDeriv :: OneParticleAccelerationFunction
-> DifferentialEquation OneParticleSystemState
oneParticleStateDeriv OneParticleAccelerationFunction
a st :: OneParticleSystemState
st@(Double
_t, St Position
_r Vec
v) = (Double
1, Vec -> Vec -> DSt
DSt Vec
v (OneParticleAccelerationFunction
a OneParticleSystemState
st))

-- | Single Runge-Kutta step
oneParticleRungeKuttaStep :: OneParticleAccelerationFunction  -- ^ acceleration function for the particle
                          -> TimeStep  -- ^ time step
                          -> OneParticleSystemState  -- ^ initial state
                          -> OneParticleSystemState  -- ^ state after one time step
oneParticleRungeKuttaStep :: OneParticleAccelerationFunction
-> Double -> OneParticleSystemState -> OneParticleSystemState
oneParticleRungeKuttaStep = forall p. StateSpace p => (p -> Diff p) -> Time p -> p -> p
rungeKutta4 forall b c a. (b -> c) -> (a -> b) -> a -> c
. OneParticleAccelerationFunction
-> DifferentialEquation OneParticleSystemState
oneParticleStateDeriv

-- | List of system states
oneParticleRungeKuttaSolution :: OneParticleAccelerationFunction  -- ^ acceleration function for the particle
                              -> TimeStep  -- ^ time step
                              -> OneParticleSystemState  -- ^ initial state
                              -> [OneParticleSystemState]  -- ^ state after one time step
oneParticleRungeKuttaSolution :: OneParticleAccelerationFunction
-> Double -> OneParticleSystemState -> [OneParticleSystemState]
oneParticleRungeKuttaSolution = forall p. StateSpace p => (p -> Diff p) -> Time p -> p -> [p]
integrateSystem forall b c a. (b -> c) -> (a -> b) -> a -> c
. OneParticleAccelerationFunction
-> DifferentialEquation OneParticleSystemState
oneParticleStateDeriv

------------------------
-- Two-particle state --
------------------------

-- | The state of a system of two particles is given by the current time,
--   the position and velocity of particle 1,
--   and the position and velocity of particle 2.
type TwoParticleSystemState = (TheTime,St,St)

-- | An acceleration function gives a pair of accelerations
--   (one for particle 1, one for particle 2) as
--   a function of the system's state.
type TwoParticleAccelerationFunction = TwoParticleSystemState -> (Vec,Vec)

-- | Time derivative of state for two particles
--   with constant mass.
twoParticleStateDeriv :: TwoParticleAccelerationFunction  -- ^ acceleration function for two particles
                      -> DifferentialEquation TwoParticleSystemState  -- ^ differential equation
twoParticleStateDeriv :: TwoParticleAccelerationFunction
-> DifferentialEquation TwoParticleSystemState
twoParticleStateDeriv TwoParticleAccelerationFunction
af2 st2 :: TwoParticleSystemState
st2@(Double
_t, St Position
_r1 Vec
v1, St Position
_r2 Vec
v2) = (Double
1, Vec -> Vec -> DSt
DSt Vec
v1 Vec
a1, Vec -> Vec -> DSt
DSt Vec
v2 Vec
a2)
    where
      (Vec
a1,Vec
a2) = TwoParticleAccelerationFunction
af2 TwoParticleSystemState
st2

-- | Single Runge-Kutta step for two-particle system
twoParticleRungeKuttaStep :: TwoParticleAccelerationFunction  -- ^ acceleration function
                          -> TimeStep  -- ^ time step
                          -> TwoParticleSystemState  -- ^ initial state
                          -> TwoParticleSystemState  -- ^ state after one time step
twoParticleRungeKuttaStep :: TwoParticleAccelerationFunction
-> Double -> TwoParticleSystemState -> TwoParticleSystemState
twoParticleRungeKuttaStep = forall p. StateSpace p => (p -> Diff p) -> Time p -> p -> p
rungeKutta4 forall b c a. (b -> c) -> (a -> b) -> a -> c
. TwoParticleAccelerationFunction
-> DifferentialEquation TwoParticleSystemState
twoParticleStateDeriv

-------------------------
-- Many-particle state --
-------------------------

-- | The state of a system of many particles is given by the current time
--   and a list of one-particle states.
type ManyParticleSystemState = (TheTime,[St])

-- | An acceleration function gives a list of accelerations
--   (one for each particle) as
--   a function of the system's state.
type ManyParticleAccelerationFunction = ManyParticleSystemState -> [Vec]

-- | Time derivative of state for many particles
--   with constant mass.
manyParticleStateDeriv :: ManyParticleAccelerationFunction  -- ^ acceleration function for many particles
                       -> DifferentialEquation ManyParticleSystemState  -- ^ differential equation
manyParticleStateDeriv :: ManyParticleAccelerationFunction
-> DifferentialEquation ManyParticleSystemState
manyParticleStateDeriv ManyParticleAccelerationFunction
af st :: ManyParticleSystemState
st@(Double
_t, [St]
sts) = (Double
1, [Vec -> Vec -> DSt
DSt Vec
v Vec
a | (Vec
v,Vec
a) <- forall a b. [a] -> [b] -> [(a, b)]
zip [Vec]
vs [Vec]
as])
    where
      vs :: [Vec]
vs = forall a b. (a -> b) -> [a] -> [b]
map St -> Vec
velocity [St]
sts
      as :: [Vec]
as = ManyParticleAccelerationFunction
af ManyParticleSystemState
st

-- | Single Runge-Kutta step for many-particle system
manyParticleRungeKuttaStep :: ManyParticleAccelerationFunction  -- ^ acceleration function
                           -> TimeStep  -- ^ time step
                           -> ManyParticleSystemState  -- ^ initial state
                           -> ManyParticleSystemState  -- ^ state after one time step
manyParticleRungeKuttaStep :: ManyParticleAccelerationFunction
-> Double -> ManyParticleSystemState -> ManyParticleSystemState
manyParticleRungeKuttaStep = forall p. StateSpace p => (p -> Diff p) -> Time p -> p -> p
rungeKutta4 forall b c a. (b -> c) -> (a -> b) -> a -> c
. ManyParticleAccelerationFunction
-> DifferentialEquation ManyParticleSystemState
manyParticleStateDeriv



-- Can we automatically incorporate Newton's third law?