Lenses and probability

:e BlockArguments
:e LambdaCase
:e OverloadedStrings

import Data.Aeson.Lens
import Control.Lens hiding (nth, (#))
import Data.Aeson
import Data.Maybe (fromMaybe)
import Control.Monad.Bayes.Class
import Control.Monad.Bayes.Sampler.Strict
import Control.Monad.Bayes.Weighted
import Data.Monoid 
import qualified Data.ByteString.Lazy as B
import Text.Pretty.Simple
import qualified Data.Text as T
import Control.Monad

Lenses are getters and setters, i.e. ways of extracting a value from inside a data structure or updating that value.

Because of their type in the Haskell lens library, namely forall f. (a -> f b) -> (c -> f d), we can specialize f to MonadDistribution m => m, and so use any lens to probabilistically update a data structure's elements in complex ways.

As an example of how easy and extensible this, we provide some examples using lenses to manipulate a JSON file. This relies on the aeson and lens-aeson packages.

json <- fromMaybe undefined . decode <$> B.readFile "../file.json" :: IO Value

pPrintCustom = pPrintOpt CheckColorTty defaultOutputOptionsNoColor {outputOptionsCompact = True, outputOptionsIndentAmount = 2} 
pPrintCustom json

Object
  ( fromList
    [
      ( "address"
      , Object
        ( fromList
          [ ( "id", Number 5.4 ), ( "streetAddress", String "21 2nd Street" ) ]
        )
      )
    ,
      ( "age", Number 27.0 )
    ,
      ( "height", Number 1.5 )
    ,
      ( "isAlive", Bool True )
    ,
      ( "name", String "John" )
    ]
  )

For example, if I have a function that randomly flips a boolean or a function that puts noise into a string, like

randomlyFlip :: MonadDistribution m => Bool -> m Bool
randomlyFlip True = bernoulli 0.01
randomlyFlip False = bernoulli 0.9

noisifyString :: MonadDistribution m => T.Text -> m T.Text
noisifyString t = fmap T.pack $ forM (T.unpack t) $ \letter -> do
    x <- bernoulli 0.2
    if x then uniformD "abcdefghijklmnopqrstuvwxyz" else return letter

then I can use it to update the "isAlive" value as follows, using the lens (_Object . traverse . _Bool) which points to the json object's list of element's boolean element:

pPrintCustom =<< sampleIOfixed ((_Object . traverse . _Bool) randomlyFlip json)

Object
  ( fromList
    [
      ( "address"
      , Object
        ( fromList
          [ ( "id", Number 5.4 ), ( "streetAddress", String "21 2nd Street" ) ]
        )
      )
    ,
      ( "age", Number 27.0 )
    ,
      ( "height", Number 1.5 )
    ,
      ( "isAlive", Bool False )
    ,
      ( "name", String "John" )
    ]
  )

Using the Plated instance of json values, I can do more powerful operations, like updating all doubles, bools and strings at arbitrarily nested positions in the json

jsonDist :: Distribution Value
jsonDist = 
    ((transformM . _Double) (\case x -> normal x 0.001) >=>
    (transformM . _Bool) randomlyFlip >=>
    (transformM . _String) noisifyString
    )
    json

pPrintCustom =<< sampleIOfixed (unweighted jsonDist)

Object
  ( fromList
    [
      ( "address"
      , Object
        ( fromList
          [
            ( "id", Number 5.398324403347256605911752558313310146331787109375 )
          ,
            ( "streetAddress", String "21 fnd Stredt" )
          ]
        )
      )
    ,
      ( "age", Number 27.001264842576262026341282762587070465087890625 )
    ,
      ( "height"
      , Number 1.5021398495523310412380624256911687552928924560546875
      )
    ,
      ( "isAlive", Bool False )
    ,
      ( "name", String "lohn" )
    ]
  )