gillespie.hs

import qualified Data.Map as Map
import qualified Data.Set as Set
import Data.Maybe
import Data.List
import Random (randomRIO)
import qualified Data.Enumerator as E
import qualified Data.Enumerator.List as EL

type Mol = String
type Mixture = Map.Map Mol Integer
type Rxn = (Mixture, (Mixture -> Mixture), Double) -- (Lhs, Action, Rate)
type State = (Double, Integer, Mixture) -- (Time, NumSteps, Mixture)

binomialCoef n 0 = 1
binomialCoef 0 k = 0
binomialCoef n k = (binomialCoef (n-1) (k-1)) * n `div` k

getActivities :: [Rxn] -> Mixture -> [Double]
getActivities rxns mixture =
  let h mol 1 = Map.findWithDefault 0 mol mixture
      h mol sc = binomialCoef (Map.findWithDefault 0 mol mixture) sc -- sc is stoichiometric coefficient
      activity (lhs, _, rate) = Map.foldlWithKey
        (\acc mol sc -> acc * (fromIntegral $ h mol sc)) rate lhs
  in map activity rxns

choice :: (Fractional b, Ord b) => (a -> b) -> [a] -> b -> a
choice f xs n =
  let chc (h:[]) acc = h
      chc (h:t) acc =
        let nacc = acc+(f h) in
        if n < nacc
          then h
          else chc t nacc
  in chc xs 0.0

getIndex xs n =
  let cum = scanl1 (+) xs
  in findIndex (>n) cum

-- Monads
selectRxn :: [Rxn] -> [Double] -> Double -> IO Rxn
selectRxn rxns activities totalActivity =
  do rnd <- randomRIO (0.0, totalActivity)
     return $ rxns !! fromJust (getIndex activities rnd)

step :: [Rxn] -> State -> IO State
step rxns (t, ns, mixture) =
  let activities = getActivities rxns mixture
      totalActivity = sum activities
  in if totalActivity == 0
    then return (t, ns+1, mixture)
    else do r@(_, action, _) <- selectRxn rxns activities totalActivity
            rnd <- randomRIO (0.0, 1.0) 
            let dt = log (1.0/rnd) / totalActivity
            return (t+dt, ns+1, action mixture)

simulate :: [Rxn] -> Mixture -> Integer -> IO [State]
simulate rxns mixture numSteps = E.run_
  (EL.iterateM (step rxns) (0.0, 0, mixture) E.$$ EL.take numSteps)
-- End of Monads

addToEntry :: (Ord a, Num b) => b -> a -> Map.Map a b -> Map.Map a b
addToEntry x k m = Map.alter (\v -> Just $ fromMaybe 0 v + x) k m

frequencies :: (Ord a, Num b) => [a] -> Map.Map a b
frequencies xs = foldl' (flip $ addToEntry 1) Map.empty xs

mixture :: String -> Mixture
mixture s = frequencies $ words s

action :: Mixture -> Mixture -> (Mixture -> Mixture)
action lhs rhs =
  let molecules = Set.union (Map.keysSet lhs) (Map.keysSet rhs)
      act mol = let r = Map.findWithDefault 0 mol rhs
                    diff = r - (Map.findWithDefault 0 mol lhs)
                in if diff == 0
                  then Nothing
                  else Just (addToEntry diff mol)
      actions = map act (Set.toList molecules)
  in foldl1' (.) (catMaybes actions)

rxns :: [String] -> [Rxn]
rxns ss =
  let rxn s = let ws = words s
                  lhs = takeWhile (/= "->") ws
                  rr_ = tail $ dropWhile (/= "->") ws
                  rhs = takeWhile (/= "@") rr_
                  rate = (read $ head $ tail $ dropWhile (/= "@") rr_) :: Double
                  lhsCounts = frequencies lhs
                  rhsCounts = frequencies rhs
              in (lhsCounts, action lhsCounts rhsCounts, rate)
  in map rxn ss

main :: IO ()
main = do simulate (rxns ["X0 E1 -> X0-E1 @ 1",
                          "X0-E1 -> X0 E1 @ 1",
                          "X0-E1 -> X1-E1 @ 1",
                          "X1-E1 -> X0-E1 @ 1",
                          "X1-E1 -> X1 E1 @ 1",
                          "X1 E1 -> X1-E1 @ 1"])
                   (Map.fromList [("X0", 10000), ("E1", 20)])
                   1000000;
          return ();