Haskell heat engine simulation. (Not quite complete.)

Main.hs

{-# LANGUAGE FlexibleContexts #-} 

import Data.Array
import Data.Array.IO
import Data.Foldable
import Data.List hiding (sum)
import System.Random
import Control.Monad.State
import Prelude hiding (sum)
--import Debug.Trace

-- Find number of solutions to e ≤ n²β ≤ f, n ≥ 0
count' :: Double -> Double -> Double -> Int
count' β e f =
    let nmin = if e >= 0 then ceiling (sqrt (e/β)) else 0
        nmax = if f >= 0 then floor (sqrt (f/β)) else -1
    in if nmax >= nmin then nmax-nmin+1 else 0

-- Find number of solutions to e ≤ m²α+n²β ≤ f, n ≥ 0
count :: Double -> Double -> Double -> Double -> Int
count α β e f =
    sum [count' β (e-α*fromIntegral m*fromIntegral m) (f-α*fromIntegral m*fromIntegral m) | m <- [0::Int ..floor (sqrt (f/α))]]

-- Pick the ith solution to e ≤ α*m²+β*n² ≤ f counting from first
-- solution of form (m, ·)
pick' :: Int -> Int -> Double -> Double -> Double -> Double -> (Int, Int)
pick' i m α β e f = do
    let e' = e-α*fromIntegral m*fromIntegral m
        f' = f-α*fromIntegral m*fromIntegral m
        c = count' β e' f'
    if i >= c
        then pick' (i-c) (m+1) α β e f
        else (m, i+if e' >= 0 then ceiling (sqrt (e'/β)) else 0)

-- Pick a solution to e ≤ α*m²+β*n² ≤ f uniformly
pick :: (MonadState StdGen m) => Double -> Double -> Double -> Double -> m (Int, Int)
pick α β e f = do
    let n = count α β e f
    i <- state $ randomR (0, n-1)
    return $ pick' i 0 α β e f

-- Collide particles with slack
shareEnergy :: (MonadState StdGen m) =>
                Double -> Double -> Double -> Double -> Int -> Int ->
                m (Int, Int, Double)
shareEnergy α β slack maxSlack m n = do
    let e = α*fromIntegral m*fromIntegral m+β*fromIntegral n*fromIntegral n
    (m', n') <- pick α β (e-maxSlack+slack) (e+slack)
    let e' = α*fromIntegral m'*fromIntegral m'+β*fromIntegral n'*fromIntegral n'
    let slack' = slack-(e'-e)
    return (m', n', slack')

{-
equilibriate1 :: Int -> Double -> Double -> Double -> Array Int Int -> 
                 StateT StdGen IO (Double, Array Int Int)
equilibriate1 n slack maxSlack α states = do
    arr <- lift $ thaw states :: StateT StdGen IO (IOArray Int Int)

    slack' <- flip execStateT slack $ replicateM_ n $ do
        i <- lift $ state $ randomR (bounds states)
        j <- lift $ state $ randomR (bounds states)
        when (i /= j) $ do
            u <- lift $ lift $ readArray arr i
            v <- lift $ lift $ readArray arr j
            slack0 <- get
            (u', v', slack') <- lift $ shareEnergy α α slack0 maxSlack u v
            put slack'
            lift $ lift $ writeArray arr i u'
            lift $ lift $ writeArray arr j v'

    states' <- lift $ freeze arr
    return (slack', states')
-}

listPick :: (MonadState StdGen m) => [Int] -> m (Int, Int)
listPick lens = do
    let sums = scanl1 (+) lens :: [Int]
    let s = last sums
    i <- state $ randomR (0, s-1)
    let Just j = findIndex (> i) sums
    return (j, if j == 0 then i else i-sums!!(j-1))

equilibriate :: Int -> Double -> Double -> [Double] -> [Array Int Int] -> 
                 StateT StdGen IO (Double, [Array Int Int])
equilibriate n slack maxSlack αs states = do
    let lens = map (\(a, b) -> b-a+1) $ map bounds states
    arr <- lift $ mapM thaw states :: StateT StdGen IO [IOArray Int Int]

    slack' <- flip execStateT slack $ replicateM_ n $ do
        i@(arrayI, indexI) <- lift $ listPick lens
        j@(arrayJ, indexJ) <- lift $ listPick lens
        when (i /= j) $ do
            u <- lift $ lift $ readArray (arr!!arrayI) indexI
            v <- lift $ lift $ readArray (arr!!arrayJ) indexJ
            slack0 <- get
            (u', v', slack') <- lift $ shareEnergy (αs!!arrayI) (αs!!arrayJ) slack0 maxSlack u v
            put slack'
            lift $ lift $ writeArray (arr!!arrayI) indexI u'
            lift $ lift $ writeArray (arr!!arrayJ) indexJ v'

    states' <- lift $ mapM freeze arr
    return (slack', states')
    --return (slack', states')

totalEnergy :: Double -> Array Int Int -> Double
totalEnergy α arr = α*sum (fmap (\x -> fromIntegral (x*x)) arr)

data EngineState = E {
        slack1 :: Double,
        slack2 :: Double,
        reservoir1 :: Array Int Int,
        piston :: Array Int Int,
        reservoir2 :: Array Int Int,
        input :: Double,
        output :: Double
    } deriving Show

len :: Array Int a -> Int
len a = let (b, e) = bounds a in (e-b+1)

dump :: MonadIO m => Double -> EngineState -> m ()
dump α engineState = do
    liftIO $ putStrLn $ "reservoir1 energy/particle = " ++ show (totalEnergy 1.0 (reservoir1 engineState)/(fromIntegral (len (reservoir1 engineState))))
    liftIO $ putStrLn $ "piston energy/particle = " ++ show (totalEnergy α (piston engineState)/(fromIntegral (len (piston engineState))))
    liftIO $ putStrLn $ "reservoir2 energy/particle = " ++ show (totalEnergy 1.0 (reservoir2 engineState)/(fromIntegral (len (reservoir2 engineState))))
    liftIO $ putStrLn $ "total input = " ++ show (input engineState)
    liftIO $ putStrLn $ "total output = " ++ show (output engineState)
    liftIO $ putStrLn $ "efficiency = " ++ show (output engineState/input engineState)

equilibriate1 :: Int -> Double -> Double -> EngineState -> StateT StdGen IO EngineState
equilibriate1 n α maxSlack state = do
    (slack, [r1', r2']) <- equilibriate n (slack1 state)
                               maxSlack
                               [1.0, α]
                               [reservoir1 state, piston state]
    return $ state { slack1 = slack, reservoir1 = r1', piston = r2' }

equilibriate2 :: Int -> Double -> Double -> EngineState -> StateT StdGen IO EngineState
equilibriate2 n β maxSlack state = do
    (slack, [r1', r2']) <- equilibriate n (slack2 state)
                               maxSlack
                               [β, 1.0]
                               [piston state, reservoir2 state]
    return $ state { slack2 = slack, piston = r1', reservoir2 = r2' }

main :: IO ()
main = 
    flip evalStateT (mkStdGen 670) $ do
        r1 <- replicateM 1000 $ state $ randomR (0 :: Int, 99)
        let r1' = array (0, 999) $ zip [0..999] r1
        p <- replicateM 100 $ state $ randomR (0 :: Int, 49)
        let p' = array (0, 99) $ zip [0..99] p
        r2 <- replicateM 1000 $ state $ randomR (0 :: Int, 49)
        let r2' = array (0, 999) $ zip [0..999] r2
        let engineState = E 0.0 0.0 r1' p' r2' 0.0 0.0

        flip evalStateT engineState $ replicateM_ 10 $ do

            engineState <- get
            liftIO $ dump 1.0 engineState
            liftIO $ putStrLn "Equilibriating with reservoir1"
            engineState' <- lift $ equilibriate1 100000 1.0 100.0 engineState
            liftIO $ dump 1.0 engineState'
            liftIO $ putStrLn "Expanding"
            liftIO $ dump 0.5 engineState'
            liftIO $ putStrLn "Equilibriating with reservoir2"
            engineState'' <- lift $ equilibriate2 100000 0.5 100.0 engineState'
            liftIO $ dump 0.5 engineState''
            liftIO $ putStrLn "Compressing"
            liftIO $ dump 1.0 engineState''
            put engineState''