{-# OPTIONS -Wall #-}
{-# LANGUAGE MultiParamTypeClasses #-}

{- | 
Module      :  LPFPCore.MultipleObjects
Copyright   :  (c) Scott N. Walck 2023
License     :  BSD3 (see LICENSE)
Maintainer  :  Scott N. Walck <walck@lvc.edu>
Stability   :  stable

Code from chapter 19 of the book Learn Physics with Functional Programming
-}

module LPFPCore.MultipleObjects where

import LPFPCore.SimpleVec
    ( Vec, R, (^+^), (^-^), (*^), (^*), (^/), zeroV, magnitude )
import LPFPCore.Mechanics1D
    ( RealVectorSpace(..), Diff(..), NumericalMethod, Mass, TimeStep, euler )
import LPFPCore.Mechanics3D
    ( OneBodyForce, ParticleState(..), DParticleState(..), HasTime(..)
    , defaultParticleState, newtonSecondPS )

type TwoBodyForce
    =  ParticleState  -- force is produced BY particle with this state
    -> ParticleState  -- force acts ON particle with this state
    -> ForceVector

type ForceVector = Vec

oneFromTwo :: ParticleState  -- state of particle PRODUCING the force
           -> TwoBodyForce
           -> OneBodyForce
oneFromTwo :: ParticleState -> TwoBodyForce -> OneBodyForce
oneFromTwo ParticleState
stBy TwoBodyForce
f = TwoBodyForce
f ParticleState
stBy

gravityMagnitude :: Mass -> Mass -> R -> R
gravityMagnitude :: R -> R -> R -> R
gravityMagnitude R
m1 R
m2 R
r = let gg :: R
gg = R
6.67408e-11  -- N m^2 / kg^2
                           in R
gg forall a. Num a => a -> a -> a
* R
m1 forall a. Num a => a -> a -> a
* R
m2 forall a. Fractional a => a -> a -> a
/ R
rforall a. Floating a => a -> a -> a
**R
2

universalGravity :: TwoBodyForce
universalGravity :: TwoBodyForce
universalGravity ParticleState
st1 ParticleState
st2
    = let gg :: R
gg = R
6.67408e-11  -- N m^2 / kg^2
          m1 :: R
m1 = ParticleState -> R
mass ParticleState
st1
          m2 :: R
m2 = ParticleState -> R
mass ParticleState
st2
          r1 :: Vec
r1 = OneBodyForce
posVec ParticleState
st1
          r2 :: Vec
r2 = OneBodyForce
posVec ParticleState
st2
          r21 :: Vec
r21 = Vec
r2 Vec -> Vec -> Vec
^-^ Vec
r1
      in (-R
gg) R -> Vec -> Vec
*^ R
m1 R -> Vec -> Vec
*^ R
m2 R -> Vec -> Vec
*^ Vec
r21 Vec -> R -> Vec
^/ Vec -> R
magnitude Vec
r21 forall a. Floating a => a -> a -> a
** R
3

constantRepulsiveForceWrong :: ForceVector -> TwoBodyForce
constantRepulsiveForceWrong :: Vec -> TwoBodyForce
constantRepulsiveForceWrong Vec
force = \ParticleState
_ ParticleState
_ -> Vec
force

constantRepulsiveForce :: R -> TwoBodyForce
constantRepulsiveForce :: R -> TwoBodyForce
constantRepulsiveForce R
force ParticleState
st1 ParticleState
st2
    = let r1 :: Vec
r1 = OneBodyForce
posVec ParticleState
st1
          r2 :: Vec
r2 = OneBodyForce
posVec ParticleState
st2
          r21 :: Vec
r21 = Vec
r2 Vec -> Vec -> Vec
^-^ Vec
r1
      in R
force R -> Vec -> Vec
*^ Vec
r21 Vec -> R -> Vec
^/ Vec -> R
magnitude Vec
r21

linearSpring :: R  -- spring constant
             -> R  -- equilibrium length
             -> TwoBodyForce
linearSpring :: R -> R -> TwoBodyForce
linearSpring R
k R
re ParticleState
st1 ParticleState
st2
    = let r1 :: Vec
r1 = OneBodyForce
posVec ParticleState
st1
          r2 :: Vec
r2 = OneBodyForce
posVec ParticleState
st2
          r21 :: Vec
r21 = Vec
r2 Vec -> Vec -> Vec
^-^ Vec
r1
          r21mag :: R
r21mag = Vec -> R
magnitude Vec
r21
      in (-R
k) R -> Vec -> Vec
*^ (R
r21mag forall a. Num a => a -> a -> a
- R
re) R -> Vec -> Vec
*^ Vec
r21 Vec -> R -> Vec
^/ R
r21mag

-- | Force provided by a spring that is fixed at one end.
fixedLinearSpring :: R -> R -> Vec -> OneBodyForce
fixedLinearSpring :: R -> R -> Vec -> OneBodyForce
fixedLinearSpring R
k R
re Vec
r1
    = ParticleState -> TwoBodyForce -> OneBodyForce
oneFromTwo (ParticleState
defaultParticleState { posVec :: Vec
posVec = Vec
r1 }) (R -> R -> TwoBodyForce
linearSpring R
k R
re)

centralForce :: (R -> R) -> TwoBodyForce
centralForce :: (R -> R) -> TwoBodyForce
centralForce R -> R
f ParticleState
st1 ParticleState
st2
    = let r1 :: Vec
r1 = OneBodyForce
posVec ParticleState
st1
          r2 :: Vec
r2 = OneBodyForce
posVec ParticleState
st2
          r21 :: Vec
r21 = Vec
r2 Vec -> Vec -> Vec
^-^ Vec
r1
          r21mag :: R
r21mag = Vec -> R
magnitude Vec
r21
      in R -> R
f R
r21mag R -> Vec -> Vec
*^ Vec
r21 Vec -> R -> Vec
^/ R
r21mag

linearSpringCentral :: R  -- spring constant
                    -> R  -- equilibrium length
                    -> TwoBodyForce
linearSpringCentral :: R -> R -> TwoBodyForce
linearSpringCentral R
k R
re = (R -> R) -> TwoBodyForce
centralForce (\R
r -> -R
k forall a. Num a => a -> a -> a
* (R
r forall a. Num a => a -> a -> a
- R
re))

billiardForce :: R  -- spring constant
              -> R  -- threshold center separation
              -> TwoBodyForce
billiardForce :: R -> R -> TwoBodyForce
billiardForce R
k R
re
    = (R -> R) -> TwoBodyForce
centralForce forall a b. (a -> b) -> a -> b
$ \R
r -> if R
r forall a. Ord a => a -> a -> Bool
>= R
re
                           then R
0
                           else (-R
k forall a. Num a => a -> a -> a
* (R
r forall a. Num a => a -> a -> a
- R
re))

data Force = ExternalForce Int OneBodyForce
           | InternalForce Int Int TwoBodyForce

data MultiParticleState
    = MPS { MultiParticleState -> [ParticleState]
particleStates :: [ParticleState] } deriving Int -> MultiParticleState -> ShowS
[MultiParticleState] -> ShowS
MultiParticleState -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [MultiParticleState] -> ShowS
$cshowList :: [MultiParticleState] -> ShowS
show :: MultiParticleState -> String
$cshow :: MultiParticleState -> String
showsPrec :: Int -> MultiParticleState -> ShowS
$cshowsPrec :: Int -> MultiParticleState -> ShowS
Show

instance HasTime MultiParticleState where
    timeOf :: MultiParticleState -> R
timeOf (MPS [ParticleState]
sts) = ParticleState -> R
time ([ParticleState]
sts forall a. [a] -> Int -> a
!! Int
0)

data DMultiParticleState = DMPS [DParticleState] deriving Int -> DMultiParticleState -> ShowS
[DMultiParticleState] -> ShowS
DMultiParticleState -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [DMultiParticleState] -> ShowS
$cshowList :: [DMultiParticleState] -> ShowS
show :: DMultiParticleState -> String
$cshow :: DMultiParticleState -> String
showsPrec :: Int -> DMultiParticleState -> ShowS
$cshowsPrec :: Int -> DMultiParticleState -> ShowS
Show

newtonSecondMPS :: [Force]
                -> MultiParticleState -> DMultiParticleState  -- a diff eqn

newtonSecondMPS :: [Force] -> MultiParticleState -> DMultiParticleState
newtonSecondMPS [Force]
fs mpst :: MultiParticleState
mpst@(MPS [ParticleState]
sts)
    = let deriv :: (Int, ParticleState) -> DParticleState
deriv (Int
n,ParticleState
st) = [OneBodyForce] -> ParticleState -> DParticleState
newtonSecondPS (Int -> MultiParticleState -> [Force] -> [OneBodyForce]
forcesOn Int
n MultiParticleState
mpst [Force]
fs) ParticleState
st
      in [DParticleState] -> DMultiParticleState
DMPS forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (Int, ParticleState) -> DParticleState
deriv (forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0..] [ParticleState]
sts)

forcesOn :: Int -> MultiParticleState -> [Force] -> [OneBodyForce]
forcesOn :: Int -> MultiParticleState -> [Force] -> [OneBodyForce]
forcesOn Int
n MultiParticleState
mpst = forall a b. (a -> b) -> [a] -> [b]
map (Int -> MultiParticleState -> Force -> OneBodyForce
forceOn Int
n MultiParticleState
mpst)

forceOn :: Int -> MultiParticleState -> Force -> OneBodyForce
forceOn :: Int -> MultiParticleState -> Force -> OneBodyForce
forceOn Int
n MultiParticleState
_         (ExternalForce Int
n0 OneBodyForce
fOneBody)
    | Int
n forall a. Eq a => a -> a -> Bool
== Int
n0    = OneBodyForce
fOneBody
    | Bool
otherwise  = forall a b. a -> b -> a
const Vec
zeroV
forceOn Int
n (MPS [ParticleState]
sts) (InternalForce Int
n0 Int
n1 TwoBodyForce
fTwoBody)
    | Int
n forall a. Eq a => a -> a -> Bool
== Int
n0    = ParticleState -> TwoBodyForce -> OneBodyForce
oneFromTwo ([ParticleState]
sts forall a. [a] -> Int -> a
!! Int
n1) TwoBodyForce
fTwoBody  -- n1 acts on n0
    | Int
n forall a. Eq a => a -> a -> Bool
== Int
n1    = ParticleState -> TwoBodyForce -> OneBodyForce
oneFromTwo ([ParticleState]
sts forall a. [a] -> Int -> a
!! Int
n0) TwoBodyForce
fTwoBody  -- n0 acts on n1
    | Bool
otherwise  = forall a b. a -> b -> a
const Vec
zeroV

instance RealVectorSpace DMultiParticleState where
    DMPS [DParticleState]
dsts1 +++ :: DMultiParticleState -> DMultiParticleState -> DMultiParticleState
+++ DMPS [DParticleState]
dsts2 = [DParticleState] -> DMultiParticleState
DMPS forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith forall ds. RealVectorSpace ds => ds -> ds -> ds
(+++) [DParticleState]
dsts1 [DParticleState]
dsts2
    scale :: R -> DMultiParticleState -> DMultiParticleState
scale R
w (DMPS [DParticleState]
dsts) = [DParticleState] -> DMultiParticleState
DMPS forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (forall ds. RealVectorSpace ds => R -> ds -> ds
scale R
w) [DParticleState]
dsts

instance Diff MultiParticleState DMultiParticleState where
    shift :: R
-> DMultiParticleState -> MultiParticleState -> MultiParticleState
shift R
dt (DMPS [DParticleState]
dsts) (MPS [ParticleState]
sts) = [ParticleState] -> MultiParticleState
MPS forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (forall s ds. Diff s ds => R -> ds -> s -> s
shift R
dt) [DParticleState]
dsts [ParticleState]
sts

eulerCromerMPS :: TimeStep        -- dt for stepping
               -> NumericalMethod MultiParticleState DMultiParticleState
eulerCromerMPS :: R -> NumericalMethod MultiParticleState DMultiParticleState
eulerCromerMPS R
dt MultiParticleState -> DMultiParticleState
deriv MultiParticleState
mpst0
    = let mpst1 :: MultiParticleState
mpst1 = forall s ds. Diff s ds => R -> (s -> ds) -> s -> s
euler R
dt MultiParticleState -> DMultiParticleState
deriv MultiParticleState
mpst0
          sts0 :: [ParticleState]
sts0 = MultiParticleState -> [ParticleState]
particleStates MultiParticleState
mpst0
          sts1 :: [ParticleState]
sts1 = MultiParticleState -> [ParticleState]
particleStates MultiParticleState
mpst1
          -- now update positions
          in [ParticleState] -> MultiParticleState
MPS forall a b. (a -> b) -> a -> b
$ [ ParticleState
st1 { posVec :: Vec
posVec = OneBodyForce
posVec ParticleState
st0 Vec -> Vec -> Vec
^+^ OneBodyForce
velocity ParticleState
st1 Vec -> R -> Vec
^* R
dt }
                         | (ParticleState
st0,ParticleState
st1) <- forall a b. [a] -> [b] -> [(a, b)]
zip [ParticleState]
sts0 [ParticleState]
sts1 ]

updateMPS :: NumericalMethod MultiParticleState DMultiParticleState
          -> [Force]
          -> MultiParticleState -> MultiParticleState
updateMPS :: NumericalMethod MultiParticleState DMultiParticleState
-> [Force] -> MultiParticleState -> MultiParticleState
updateMPS NumericalMethod MultiParticleState DMultiParticleState
method = NumericalMethod MultiParticleState DMultiParticleState
method forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Force] -> MultiParticleState -> DMultiParticleState
newtonSecondMPS

statesMPS :: NumericalMethod MultiParticleState DMultiParticleState
          -> [Force]
          -> MultiParticleState -> [MultiParticleState]
statesMPS :: NumericalMethod MultiParticleState DMultiParticleState
-> [Force] -> MultiParticleState -> [MultiParticleState]
statesMPS NumericalMethod MultiParticleState DMultiParticleState
method = forall a. (a -> a) -> a -> [a]
iterate forall b c a. (b -> c) -> (a -> b) -> a -> c
. NumericalMethod MultiParticleState DMultiParticleState
method forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Force] -> MultiParticleState -> DMultiParticleState
newtonSecondMPS

speed :: ParticleState -> R
speed :: ParticleState -> R
speed ParticleState
st = forall a. HasCallStack => a
undefined ParticleState
st

universalGravity' :: TwoBodyForce
universalGravity' :: TwoBodyForce
universalGravity' (ParticleState R
m1 R
_ R
_ Vec
r1 Vec
_) (ParticleState R
m2 R
_ R
_ Vec
r2 Vec
_)
    = forall a. HasCallStack => a
undefined R
m1 Vec
r1 R
m2 Vec
r2

universalGravityCentral :: TwoBodyForce
universalGravityCentral :: TwoBodyForce
universalGravityCentral = forall a. HasCallStack => a
undefined

lennardJones :: R  -- dissociation energy
             -> R  -- equilibrium length
             -> TwoBodyForce
lennardJones :: R -> R -> TwoBodyForce
lennardJones R
de R
re = (R -> R) -> TwoBodyForce
centralForce forall a b. (a -> b) -> a -> b
$ \R
r -> forall a. HasCallStack => a
undefined R
de R
re R
r

systemKE :: MultiParticleState -> R
systemKE :: MultiParticleState -> R
systemKE MultiParticleState
mpst = forall a. HasCallStack => a
undefined MultiParticleState
mpst

forcesOn' :: Int -> MultiParticleState -> [Force] -> [OneBodyForce]
forcesOn' :: Int -> MultiParticleState -> [Force] -> [OneBodyForce]
forcesOn' Int
n MultiParticleState
mpst [Force]
fs = Int -> [Force] -> [OneBodyForce]
externalForcesOn Int
n [Force]
fs forall a. [a] -> [a] -> [a]
++ Int -> MultiParticleState -> [Force] -> [OneBodyForce]
internalForcesOn Int
n MultiParticleState
mpst [Force]
fs

externalForcesOn :: Int -> [Force] -> [OneBodyForce]
externalForcesOn :: Int -> [Force] -> [OneBodyForce]
externalForcesOn Int
n [Force]
fs = forall a. HasCallStack => a
undefined Int
n [Force]
fs

internalForcesOn :: Int -> MultiParticleState -> [Force] -> [OneBodyForce]
internalForcesOn :: Int -> MultiParticleState -> [Force] -> [OneBodyForce]
internalForcesOn Int
n (MPS [ParticleState]
sts) [Force]
fs
    = [ParticleState -> TwoBodyForce -> OneBodyForce
oneFromTwo ([ParticleState]
sts forall a. [a] -> Int -> a
!! Int
n1) TwoBodyForce
f | InternalForce Int
n0 Int
n1 TwoBodyForce
f <- [Force]
fs, Int
n forall a. Eq a => a -> a -> Bool
== Int
n0] forall a. [a] -> [a] -> [a]
++
      [ParticleState -> TwoBodyForce -> OneBodyForce
oneFromTwo ([ParticleState]
sts forall a. [a] -> Int -> a
!! Int
n0) TwoBodyForce
f | InternalForce Int
n0 Int
n1 TwoBodyForce
f <- [Force]
fs, Int
n forall a. Eq a => a -> a -> Bool
== Int
n1]