Simple Vector Calculus in Haskell

README.md

Vector Calculus in Haskell

Why is this a thing?

This module started out as an idle side project, mostly to learn Haskell but also because I'm too lazy to calculate grad/divg/curl etc myself

Soo why am I posting this now?

Suprisingly enough (to me at least), I was unable to find any module similar to this one anywhere.

VectorCalculus.hs

{-# LANGUAGE MultiParamTypeClasses, FlexibleInstances, UndecidableInstances #-}

-- 3D Vector calculus
module VectorCalculus where

import Control.Monad

-- You can use any a, but you won't get much out of this module
-- if a isn't Field a or at least Ring a
data Vector3 a = Vector3 a a a

-- I can't figure out how to tell Haskell's type system
-- how to implement vector operations
-- sorry
-- If you use any vector and want the operations bound, 

data VectorExpr a = MinusV (VectorExpr a)
                    | (VectorExpr a) :*: (VectorExpr a) -- Cross product
                    | (VectorExpr a) :+: (VectorExpr a) -- Vector addition
                    | (ScalarExpr a) :$*: (VectorExpr a) -- Left scalar multiplication
                    | (VectorExpr a) :*$: (ScalarExpr a) -- Right scalar multiplication (differs from :$*_: for non-commutative rings
                    | Coords (ScalarExpr a) (ScalarExpr a) (ScalarExpr a)
                    deriving (Show, Eq)

data ScalarExpr a = Const a
                    | Var String
                    | MinusS (ScalarExpr a)
                    | Log (ScalarExpr a)
                    | (VectorExpr a) :.: (VectorExpr a)
                    | (ScalarExpr a) :$^$: (ScalarExpr a)
                    | (ScalarExpr a) :$*$: (ScalarExpr a)
                    | (ScalarExpr a) :$+$: (ScalarExpr a)
                    | CoordX (VectorExpr a) | CoordY (VectorExpr a) | CoordZ (VectorExpr a)
                    deriving (Show, Eq)

infixl 8 :$^$:
infixl 7 :*: , :$*: , :*$: , :.: , :$*$:
infixl 6 :+: , :$+$:

-- These still do some infinite recursion i.e. no good
simplifyOnceV :: (Real a) => VectorExpr a -> VectorExpr a
simplifyOnceV (Coords x y z) = Coords (simplifyOnceS x) (simplifyOnceS y) (simplifyOnceS z)
simplifyOnceV (u :+: v) = Coords (CoordX u :$+$: CoordX v) (CoordY u :$+$: CoordY v) (CoordZ u :$+$: CoordZ v)
simplifyOnceV (a :$*: v) = Coords (a :$*$: CoordX v) (a :$*$: CoordY v) (a :$*$: CoordZ v)
simplifyOnceV (v :*$: a) = Coords (CoordX v :$*$: a) (CoordY v :$*$: a) (CoordZ v :$*$: a)
simplifyOnceV (u :*: v) = Coords (CoordY u :$*$: CoordZ v :$+$: MinusS (CoordZ u :$*$: CoordY v))
                                         (CoordZ u :$*$: CoordX v :$+$: MinusS (CoordX u :$*$: CoordZ v))
                                         (CoordX u :$*$: CoordY v :$+$: MinusS (CoordY u :$*$: CoordX v))
simplifyOnceV (MinusV (Coords x y z)) = simplifyOnceV $ Coords (MinusS x) (MinusS y) (MinusS z)
simplifyOnceV (MinusV v) = MinusV $ simplifyOnceV v

simplifyOnceS :: (Real a) => ScalarExpr a -> ScalarExpr a
simplifyOnceS (Const s) = Const s
simplifyOnceS (Var s) = Var s
simplifyOnceS (Const 0 :$+$: q) = q
simplifyOnceS (p :$+$: Const 0) = p
simplifyOnceS (Const 1 :$*$: q) = q
simplifyOnceS (p :$*$: Const 1) = p
simplifyOnceS (Const 0 :$*$: q) = Const 0
simplifyOnceS (p :$*$: Const 0) = Const 0
simplifyOnceS (Const x :$+$: Const y) = Const (x + y)
simplifyOnceS (p :$+$: q) = simplifyOnceS  p :$+$: simplifyOnceS q
simplifyOnceS (Const x :$*$: Const y) = Const (x * y)
simplifyOnceS (p :$*$: (q :$+$: r)) = simplifyOnceS (p :$*$: q :$+$: p :$*$: r)
simplifyOnceS (b :$^$: p :$*$: b' :$^$: q) = if b == b'
                                             then b :$^$: (p :$+$: q)
                                             else (simplifyOnceS b :$^$: simplifyOnceS p :$*$: simplifyOnceS b' :$^$: simplifyOnceS q)
simplifyOnceS ((q :$+$: r) :$*$: p) = simplifyOnceS (q :$*$: p :$+$: r :$*$: p)
simplifyOnceS (p :$*$: q) = (simplifyOnceS p :$*$: simplifyOnceS q)
simplifyOnceS ((p :$*$: q) :$^$: e) = p :$^$: e :$*$: q :$^$: e
simplifyOnceS (Const 1 :$^$: _) = Const 1
simplifyOnceS (Const 0 :$^$: _) = Const 0
simplifyOnceS (_ :$^$: Const 0) = Const 1
simplifyOnceS (b :$^$: Const 1) = b
simplifyOnceS (b :$^$: e) = simplifyOnceS b :$^$: simplifyOnceS e
simplifyOnceS (Log (p :$*$: q)) = Log p :$+$: Log q
simplifyOnceS (Log (b :$^$: e)) = e :$*$: Log b
simplifyOnceS (Log p) = Log $ simplifyOnceS p
simplifyOnceS (p :.: q) = (CoordX p :$*$: CoordX q :$+$:
                           CoordY p :$*$: CoordY q :$+$:
                           CoordZ p :$*$: CoordZ q)
simplifyOnceS (CoordX (Coords x _ _)) = x
simplifyOnceS (CoordY (Coords _ y _)) = y
simplifyOnceS (CoordZ (Coords _ _ z)) = z
simplifyOnceS (CoordX v) = CoordX $ simplifyOnceV v
simplifyOnceS (CoordY v) = CoordY $ simplifyOnceV v
simplifyOnceS (CoordZ v) = CoordZ $ simplifyOnceV v
simplifyOnceS (MinusS (Const c)) = Const (-c)
simplifyOnceS (MinusS (p :$+$: q)) = MinusS p :$+$: MinusS q
simplifyOnceS (MinusS p) = MinusS $ simplifyOnceS p

-- Recursion must be added externally by computing the infinite functional power of simplifyV / simplifyV
-- Also for some reason Haskell's type checker thinks Real a is required, but only if Real a is made an instance of
-- the operator typeclasses
simplifyV :: (Real a) => VectorExpr a -> VectorExpr a
simplifyV expr = let simplifyStepR :: Real a => VectorExpr a -> VectorExpr a -> VectorExpr a
                     simplifyStepR curr prev = if curr == prev
                                                then curr
                                                else simplifyStepR (simplifyOnceV curr) curr
                 in simplifyStepR (simplifyOnceV expr) expr
simplifyS :: (Real a) => ScalarExpr a -> ScalarExpr a
simplifyS expr = let simplifyStepR :: Real a => ScalarExpr a -> ScalarExpr a -> ScalarExpr a
                     simplifyStepR curr prev = if curr == prev
                                                then curr
                                                else simplifyStepR (simplifyOnceS curr) curr
                 in simplifyStepR (simplifyOnceS expr) expr

-- THIS IS HOW YOU DERIVE STUFF
-- Straightforward translation of differentation rules
deriveV :: (Real a) => String -> VectorExpr a -> VectorExpr a
deriveV s (MinusV v) = MinusV $ deriveV s v
deriveV s (u :*: v) = (deriveV s u :*: v) :+: (u :*: deriveV s v)
deriveV s (u :+: v) = deriveV s u :+: deriveV s v
deriveV s (p :$*: v) = (deriveS s p :$*: v) :+: (p :$*: deriveV s v)
deriveV s (u :*$: p) = (deriveV s u :*$: p) :+: (u :*$: deriveS s p)
deriveV s (Coords x y z) = Coords (deriveS s x) (deriveS s y) (deriveS s z)

deriveS :: (Real a) => String -> ScalarExpr a -> ScalarExpr a
deriveS s (Const p) = Const 0
deriveS s (Var n) = if n == s then Const 1 else Const 0
deriveS s (MinusS p) = MinusS $ deriveS s p
deriveS s (u :.: v) = deriveV s u :.: deriveV s v
deriveS s (p :$+$: q) = deriveS s p :$+$: deriveS s q
deriveS s (p :$*$: q) = deriveS s p :$*$: q :$+$: p :$*$: deriveS s q
deriveS s (p :$^$: q) = deriveS s (q :$*$: Log p) :$*$: p :$^$: q
deriveS s (Log p) = (deriveS s p) :$*$: p :$^$: Const (-1)
deriveS s (CoordX v) = CoordX $ deriveV s v
deriveS s (CoordY v) = CoordY $ deriveV s v
deriveS s (CoordZ v) = CoordZ $ deriveV s v

-- The divergence of a vector field
divg :: (Real a) => VectorExpr a -> ScalarExpr a
divg f = simplifyS $ (deriveS "x" $ CoordX f) :$+$: (deriveS "y" $ CoordY f) :$+$: (deriveS "z" $ CoordZ f)

-- The gradient of a scalar field
grad :: (Real a) => ScalarExpr a -> VectorExpr a
grad f = simplifyV $ Coords (deriveS "x" f) (deriveS "y" f) (deriveS "z" f)

-- The curl of a vector field
curl :: (Real a) => VectorExpr a -> VectorExpr a
curl f = simplifyV $ Coords (deriveS "y" (CoordZ f) :$+$: MinusS (deriveS "z" (CoordY f)))
                            (deriveS "z" (CoordX f) :$+$: MinusS (deriveS "x" (CoordZ f)))
                            (deriveS "x" (CoordY f) :$+$: MinusS (deriveS "y" (CoordX f)))

uniq :: (Eq a) => [a] -> [a]
uniq [] = []
uniq (x:xs) = if x `elem` xs then xs else x : uniq xs

-- Check which variables a scalar/vector expression depends on
varsS :: (Real a) => ScalarExpr a -> [String]
varsS = let vars' (Const _) = []
            vars' (Var s) = [s]
            vars' (MinusS p) = vars' p
            vars' (u :.: v) = varsV u ++ varsV v
            vars' (p :$*$: q) = vars' p ++ vars' q
            vars' (p :$+$: q) = vars' p ++ vars' q
            -- Do not match against Coord* constructor, it canot be returned from simplifyS anyway
        in uniq . vars' . simplifyS

varsV :: (Real a) => VectorExpr a -> [String]
varsV v = let (Coords x y z) = simplifyV v
          in uniq $ varsS x ++ varsS y ++ varsS z

(|=) :: (Real a, Eq a) => VectorExpr a -> (ScalarExpr a, ScalarExpr a) -> VectorExpr a
(u :*: v) |= t = (u |= t) :*: (v |= t)
(u :+: v) |= t = (u |= t) :+: (v |= t)
(p :$*: v) |= t = (p ||= t) :$*: (v |= t)
(u :*$: q) |= t = (u |= t) :*$: (q ||= t)
(Coords x y z) |= t = Coords (x ||= t) (y ||= t) (z ||= t)

(||=) :: (Real a, Eq a) => ScalarExpr a -> (ScalarExpr a, ScalarExpr a) -> ScalarExpr a
Const p ||= t = Const p
ex ||= (old, nw)
                | simplifyS ex == simplifyS old = nw  
Var n ||= t = Var n
MinusS p ||= t = MinusS $ p ||= t
(u :.: v) ||= t = (u |= t) :.: (v |= t)
(p :$*$: q) ||= t = (p ||= t) :$*$: (q ||= t)
(p :$+$: q) ||= t = (p ||= t) :$+$: (q ||= t)
CoordX u ||= t = CoordX $ u |= t
CoordY u ||= t = CoordY $ u |= t
CoordZ u ||= t = CoordZ $ u |= t

(\\=:) :: ScalarExpr a -> ScalarExpr a -> (ScalarExpr a, ScalarExpr a)
(\\=:) = (,)

infixl 5 |=,||=
infix 9 \\=:

-- Getting Nothing here typically means a variable wasn't substituted.
evalV :: (Floating a, Real a) => VectorExpr a -> Maybe (Vector3 a)
evalV (Coords x y z) = liftM3 Vector3 (evalS x) (evalS y) (evalS z) -- Yay monads!
evalV v = evalV $ simplifyV v

evalS :: (Floating a, Real a) => ScalarExpr a -> Maybe a
evalS (Const c) = Just c
evalS (Var _) = Nothing
evalS (MinusS p) = liftM negate $ evalS p
evalS (Log p) = liftM log $ evalS p
evalS (u :.: v) = do (Vector3 x y z) <- evalV u
                     (Vector3 x' y' z') <- evalV u
                     Just $ x*x' + y*y' + z*z'
evalS (p :$+$: q) = liftM2 (+) (evalS p) (evalS q)
evalS (p :$*$: q) = liftM2 (*) (evalS p) (evalS q)
evalS (p :$^$: q) = liftM2 (**) (evalS p) (evalS q)
evalS (CoordX u) = evalV u >>= \(Vector3 x _ _) -> Just x
evalS (CoordY u) = evalV u >>= \(Vector3 _ y _) -> Just y
evalS (CoordZ u) = evalV u >>= \(Vector3 _ _ z) -> Just z

class Presentable a where
    present :: a -> String

instance (Show a) => Presentable (Vector3 a) where
    present (Vector3 x y z) = "(" ++ show x ++ " | " ++ show y ++ " | " ++ show z ++ ")"

instance (Show a) => Presentable (VectorExpr a) where
    present (MinusV u) = "-(" ++ present u ++ ")"
    present (u :*: v) = "(" ++ present u ++ ") x (" ++ present v ++ ")"
    present (u :+: v) = present u ++ " + " ++ present v
    present (p :$*: v) = "(" ++ present p ++ ") (" ++ present v ++ ")"
    present (u :*$: q) = "(" ++ present u ++ ") (" ++ present q ++ ")"
    present (Coords x y z) = "(" ++ present x ++ " | " ++ present y ++ " | " ++ present z ++ ")"

instance (Show a) => Presentable (ScalarExpr a) where
    present (Const c) = show c
    present (Var s) = "'" ++ s ++ "'"
    present (MinusS p) = "-(" ++ present p ++ ")"
    present (Log p) = "ln (" ++ present p ++ ")"
    present (u :.: v) = "(" ++ present u ++ ") . (" ++ present v ++ ")"
    present (p :$+$: q) = present p ++ " + " ++ present q
    present (p :$*$: q) = "(" ++ present p ++ ") * (" ++ present q ++ ")"
    present (p :$^$: q) = "(" ++ present p ++ ") ^ (" ++ present q ++ ")"
    present (CoordX u) = "X-of (" ++ present u ++ ")"
    present (CoordY u) = "Y-of (" ++ present u ++ ")"
    present (CoordZ u) = "Z-of (" ++ present u ++ ")"