106 lines
2.7 KiB
Haskell
106 lines
2.7 KiB
Haskell
|
#!/usr/bin/env nix-script
|
||
|
|
#!>haskell
|
||
|
|
#! haskell | attoparsec hmatrix
|
||
|
|
|
||
|
|
{- Balance chemical reactions using linear algebra. -}
|
||
|
|
|
||
|
|
{-# LANGUAGE OverloadedStrings #-}
|
||
|
|
|
||
|
|
import Control.Applicative (many, (<|>))
|
||
|
|
import Data.Function (on)
|
||
|
|
import Data.Maybe (fromMaybe)
|
||
|
|
import Data.List (groupBy, sort, nub)
|
||
|
|
import Data.Text (pack)
|
||
|
|
import Data.Attoparsec.Text
|
||
|
|
import Numeric.LinearAlgebra (Field, ident, inv, rank)
|
||
|
|
import Numeric.LinearAlgebra.Data
|
||
|
|
|
||
|
|
type Atom = String
|
||
|
|
type Element = (Atom, Z)
|
||
|
|
type Compound = [Element]
|
||
|
|
type Reaction = ([Compound], [Compound])
|
||
|
|
|
||
|
|
|
||
|
|
-- Parsers --
|
||
|
|
|
||
|
|
element :: Parser Element
|
||
|
|
element = do
|
||
|
|
atom <- ((++) <$> upper <*> lower) <|> upper
|
||
|
|
quantity <- option 1 decimal
|
||
|
|
return (atom, quantity)
|
||
|
|
where upper = charRange "A-Z"
|
||
|
|
lower = charRange "a-z"
|
||
|
|
|
||
|
|
compound :: Parser Compound
|
||
|
|
compound = reduce <$> many element
|
||
|
|
|
||
|
|
hand :: Parser [Compound]
|
||
|
|
hand = compound `sepBy` spaced "+"
|
||
|
|
|
||
|
|
reaction :: Parser Reaction
|
||
|
|
reaction = (,) <$> (hand <* spaced "->") <*> hand
|
||
|
|
|
||
|
|
|
||
|
|
-- parser helpers --
|
||
|
|
|
||
|
|
-- | Parse while ignoring enclosing whitespace
|
||
|
|
spaced :: Parser a -> Parser ()
|
||
|
|
spaced p = skipSpace >> p >> skipSpace
|
||
|
|
|
||
|
|
-- | Create a parser from a characters range
|
||
|
|
charRange :: String -> Parser String
|
||
|
|
charRange = fmap pure . satisfy . inClass
|
||
|
|
|
||
|
|
|
||
|
|
-- stochiometrycal matrix --
|
||
|
|
|
||
|
|
-- | Reduce a 'Compound' to his empirical formula
|
||
|
|
reduce :: Compound -> Compound
|
||
|
|
reduce = map count . groupBy ((==) `on` fst) . sort
|
||
|
|
where count ((s,n):xs) = (,) s (n + sum (map snd xs))
|
||
|
|
|
||
|
|
-- | Find a vector basis for the reaction
|
||
|
|
basisFor :: Reaction -> [Atom]
|
||
|
|
basisFor = nub . map fst . concat . fst
|
||
|
|
|
||
|
|
-- | Calculate the coordinats vector
|
||
|
|
coords :: [Atom] -> Compound -> Vector Z
|
||
|
|
coords basis v = fromList (map (flip count v) basis)
|
||
|
|
where count e = fromMaybe 0 . lookup e
|
||
|
|
|
||
|
|
-- | Calculate stochoimetrical matrix
|
||
|
|
stoch :: Reaction -> Matrix Z
|
||
|
|
stoch r@(x, y) = fromColumns (a ++ b)
|
||
|
|
where col = coords (basisFor r)
|
||
|
|
a = map col x
|
||
|
|
b = map (negate . col) y
|
||
|
|
|
||
|
|
|
||
|
|
-- linear algebra --
|
||
|
|
|
||
|
|
-- | Find a base of Ker A
|
||
|
|
kernel :: Matrix Z -> Vector Z
|
||
|
|
kernel a = cmap (round . min) t
|
||
|
|
where a' = extend (fromZ a) :: Matrix R
|
||
|
|
t = toColumns (inv a') !! (cols a'-1)
|
||
|
|
min = (/ minimum (toList t))
|
||
|
|
|
||
|
|
-- | Extend a Matrix to a square
|
||
|
|
extend :: Field k => Matrix k -> Matrix k
|
||
|
|
extend a
|
||
|
|
| m == n = a
|
||
|
|
| otherwise = a === fromRows padd
|
||
|
|
where (m,n) = size a
|
||
|
|
id = ident n
|
||
|
|
padd = map (toColumns id !!) [rank a-1..m-1]
|
||
|
|
|
||
|
|
|
||
|
|
main :: IO ()
|
||
|
|
main = do
|
||
|
|
input <- pack <$> prompt
|
||
|
|
case parseOnly reaction input of
|
||
|
|
Right react -> print (kernel $ stoch react)
|
||
|
|
Left err -> putStrLn ("Failed to parse: " ++ err)
|
||
|
|
main
|
||
|
|
where prompt = putStr "eq> " >> getLine
|