The State Monad in Python terms

Consider a function like sin. It is a “pure” function, which is to say that it is referentially transparent – every time you call sin(pi) it will return the same value. You could in fact in C say #define sin(pi) = 0.0 in complete confidence that your program would continue to function as before. And it doesn’t need to touch any of the system’s global state – that is to say, it leaves no evidence of having run, apart from having consumed some CPU (OK so this will update various counters in the kernel – there’s no such thing as a truly pure function).

Now let’s consider a case in which we would want calculating a sine to maintain some state. Say we are prototyping a system in Python that we will later implement in Forth on an embedded platform on which sin is expensive and we want to count how many times it is called (I’m just making this up!). So we would like to give sin an angle to calculate the sine of, and a variable representing how many times it has been called already.

It will return both the result we are interested in and an updated counter. In Python:

from math import pi, sin as _sin 

def sin(theta, state): 
    return (_sin(theta), state + 1)     
    
(x, count) = sin(pi, 0)  
(x, count) = sin(pi/2, count) 
(x, count) = sin(pi/4, count) 
          
print "sin was called %d times" % count

We can say this because Python has mutable variables, but it if didn’t, like Haskell, we would have to say something like:

from math import pi, sin as _sin 
   
def sin(theta, state):  
    return (_sin(theta), state + 1) 
                                                                                                                  
(x, count1) = sin(pi, 0)    
(x2, count2) = sin(pi/2, count1)  
(x3, count3) = sin(pi/4, count2) 
                                                                                                                  
print "sin was called %d times" % count3  

Obviously that is completely nonsensical – to do that we’d have to know already how many states we planned to have! And we’ve changed the way sin works. Far better to “hide” the counter:

from math import pi, sin as _sin 
                   
sin_counter = 0
     
def sin(theta):
    global sin_counter 
    sin_counter = sin_counter + 1 
    return _sin(theta) 
      
x1 = sin(pi)  
x2 = sin(pi/2) 
x3 = sin(pi/4) 

print "sin was called %d times" % sin_counter 

But now we have introduced some global state – what if we had chosen the name x_counter and someone else also used it? The conjunction of these ideas – passing a state between invocations of a function, avoiding global state, and abstracting away the management of the state – results in what Haskell calls the State Monad, simply a function that takes a state and returns a result + a new state. So the equivalent code in Haskell:

module Main where 
        
import Control.Monad.State
import Prelude hiding (sin) 
import qualified GHC.Float as F  

sin:: Float -> State Int Float 
sin x = do c <- get 
           put (c + 1)  
           return $ F.sin x 
                                                                                                                  
sins:: ([Float], Int)           
sins = runState (do  
        x <- mapM sin [pi, pi/2, pi/4]
        return x) 0 -- initial state
                                                                                                                  
main = do 
    let (x, c) = sins   
               
    putStrLn $ "sin was called " ++ (show c) ++ " times"  

That’s quite verbose and isn’t entirely seamless (e.g. replacing = with ← in calls to sin). But we have managed to achieve (or at least fake) a mutable variable that persists between invocations of a function. The important keywords are:

  • get : retrieve the value that is in the monad – in this case the State which is an Int (it could be any type). In our (a, s) pair, it is equivalent to snd
  • put : put the value back into the monad
  • <- (line 14) : equivalent to fst on our (a, s) pair
  • return : this is not like return in Python – it is to ← as put is to get. IMHO this keyword was poorly chosen, it’s already encumbered with too much meaning (or state!) from other languages.
*Main> main
sin was called 3 times
*Main> sins
([-8.742278e-8,1.0,0.70710677],3)
*Main>

Having written this I now understand why in the examples random number generation is usually chosen to illustrate the State monad. The RNG in Haskell is interesting; because it executes in pure code, it obviously has no access to sources of entropy – reading /dev/urandom would constitute I/O. And as it has no state of its own it cannot “move on” from its first generated value. If you call it twice you get the same random number twice! But the return type of randomR is actually:

Prelude Control.Monad.State> :m +System.Random
Prelude Control.Monad.State System.Random> :t randomR
randomR :: (Random a, RandomGen g) => (a, a) -> g -> (a, g)

It wants a range (e.g. 0-10) and it wants an inital random number generator, and it will return a random number and a new RNG. That fits perfectly with what the State monad expects:

Prelude Control.Monad.State System.Random> :t State
State :: (s -> (a, s)) -> State s a

Which means that:

randInt::   Int -> State StdGen Int  
randInt x = do g <- get 
               (x', g') <- return $ randomR (0, x) g 
               put g'  
               return x'

Can neatly be replaced with:

randInt x = State $ randomR (0, x)

Thus impressing everyone. I suppose the morals of this story are twofold:

  • Monads are pretty straightforward if you come from a language like Python where returning a tuple is natural
  • Think about Monads up front or you will end up in a situation where your Haskell is more complicated than your Python!

About Gaius

Jus' a good ol' boy, never meanin' no harm
This entry was posted in Haskell, Python. Bookmark the permalink.

5 Responses to The State Monad in Python terms

  1. Pingback: Scope in database code | So I decided to take my work back underground

  2. Pingback: Scope in database code (2) | So I decided to take my work back underground

  3. Pingback: LDAP (2) | So I decided to take my work back underground

  4. Vladimir says:

    Hi, It is steal not clear how to use it in python. Could you please show an example – how to define such monad ?

    • Gaius says:

      The point is not to use it in Python, but to understand that the State monad is simply syntactic sugar for in Python:

      def stateful_func (oldstate, x):
          # do something that turn x to y and that changes oldstate to newstate
          return (newstate, y)
      

      It means a function that takes argument x and returns value y doesn’t need to explicitly pass state around. HTH.

Leave a Reply

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out /  Change )

Twitter picture

You are commenting using your Twitter account. Log Out /  Change )

Facebook photo

You are commenting using your Facebook account. Log Out /  Change )

Connecting to %s