-
-
Save PiDelport/abe82cb96a8ac5a86ac1 to your computer and use it in GitHub Desktop.
Variation and feedback on Michael Bernstein's Markov chain implementation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
module Chain where | |
import Control.Applicative | |
import Data.List (tails, unfoldr) | |
import System.Random (RandomGen, newStdGen, randomR) | |
import qualified Data.Map.Strict as M | |
import Data.Sequence as S (Seq, (><), (|>)) | |
import qualified Data.Sequence as S | |
-- A map of prefixes to possible successors. | |
type PrefixMap k = M.Map (Seq k) (Seq k) | |
-- Build a prefix map of width n for the given sequence. | |
build :: Ord k => Int -> [k] -> PrefixMap k | |
build n ks = M.fromListWith (><) prefixes | |
where prefixes = [ (S.fromList prefix, S.singleton k) | (prefix, k:_) <- splitAt n `map` tails ks ] | |
-- Generate a Markov chain from the given transition map and initial state. | |
-- The state should be a fixed-width prefix. | |
generate :: (Ord k, RandomGen g) => PrefixMap k -> Seq k -> g -> [k] | |
generate transitions initState rng = unfoldr step (initState, rng) | |
where | |
-- Choose a random successor k for prefix state. | |
step (state, g) | Just candidates <- M.lookup state transitions, | |
Just (k, g') <- pick candidates g, | |
state' <- S.drop 1 state |> k -- new prefix | |
= Just (k, (state', g')) -- output k | |
-- If state has no successor, stop. | |
| otherwise = Nothing | |
-- Like generate, but use the global random generator in IO. | |
generateIO :: Ord k => PrefixMap k -> Seq k -> IO [k] | |
generateIO t i = generate t i <$> newStdGen | |
-- Helper: Pick a random element from a list. | |
pick :: RandomGen g => Seq k -> g -> Maybe (k, g) | |
pick xs _ | S.null xs = Nothing | |
pick xs g = Just (xs `S.index` i, g') | |
where (i, g') = randomR (0, S.length xs - 1) g | |
main :: IO () | |
main = do | |
print' =<< unwords <$> demonstrate 1 (words "a man a plan a canal panama") (words "a") | |
print' =<< demonstrate 1 " abracadabra " "a" | |
print' =<< demonstrate 2 "--a--b--c--d--e++A++B++C++D++E--" "--" | |
where | |
print' = putStrLn . take 70 | |
-- Generate a Markov chain for the given prefix width, input, and seed. | |
demonstrate n ks s = (s ++) <$> generateIO (build n ks) (S.fromList s) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
This version generalizes the element type to any
Ord k
, simplifiesbuild
, and letsgenerate
produce an infinite lazy output stream.There are two variations:
Sample output: