jtobin.io

Runtime Analysis (ppad-fixed)

2025-12-26T00:00:00+04:00

As noted in my last post, I recently released ppad-fixed, a high-performance fixed-width word library meant for use in cryptographic applications.

At present I use this library to drive elliptic curve arithmetic in ppad-secp256k1, as well as stuff that depends on that, e.g. ppad-bip32. It’s safe to say I have few users at the moment, but, as they say, “the way you do anything is the way you do everything,” so in this post I want to rigorously analyze some of its runtime and execution properties, to increase assurance both in general, and especially to myself, that it’s really fit for purpose.

My goal here is to illustrate constant-resource execution of this library on aarch64 when compiled with optimizations via GHC 9.8.1’s LLVM backend, first by presentation of the code in general, then some analysis of the STG IR, and then analysis of the produced assembly. I’ll also look at overall sanity checks, i.e. Criterion benchmarks and allocation measurement via ‘weigh’.

Data Representation

To start, ppad-fixed rolls its own four-limb word type (i.e. 256-bit, on a 64-bit system), consisting of a heap-allocated data constructor wrapping a strict, unboxed tuple:

data Wider = Wider !(# Limb, Limb, Limb, Limb #)

A ‘Limb’ here is just a newtype over an unboxed word, i.e. ‘Word#’, which is register-allocated. Such a ‘Wider’ allocates precisely 40 bytes on a 64-bit system: 8 bytes for each limb, and another 8 bytes for the ‘Wider’ data constructor itself. Library functions that return a ‘Wider’ value should allocate 40 bytes exactly, no matter what intermediate calculations they perform.

Operations on individual ‘Limbs’ just use available primitives from GHC.Exts, e.g. the following carrying addition function, and otherwise are written from scratch:

add_c#
  :: Limb             -- ^ augend
  -> Limb             -- ^ addend
  -> Limb             -- ^ carry
  -> (# Limb, Limb #) -- ^ (# sum, new carry #)
add_c# (Limb a) (Limb b) (Limb c) =
  let !(# c0, s0 #) = Exts.plusWord2# a b
      !(# c1,  s #) = Exts.plusWord2# s0 c
  in  (# Limb s, Limb (Exts.or# c0 c1) #)
{-# INLINE add_c# #-}

For ‘Wider’ and friends, the library provides two API’s; each function has an unboxed variant that works with unboxed tuples, avoiding all heap allocation whatsoever, and a boxed variant that allocates data constructors only at the end of the computation, simply to box up the result. Internal computations are always performed exclusively in terms of unboxed tuples.

For example, building on the above carrying addition function on single limbs, we have the following overflowing addition function on four-limb tuples:

add_o#
  :: (# Limb, Limb, Limb, Limb #)             -- ^ augend
  -> (# Limb, Limb, Limb, Limb #)             -- ^ addend
  -> (# (# Limb, Limb, Limb, Limb #), Limb #) -- ^ (# sum, carry bit #)
add_o# (# a0, a1, a2, a3 #) (# b0, b1, b2, b3 #) =
  let !(# s0, c0 #) = L.add_o# a0 b0
      !(# s1, c1 #) = L.add_c# a1 b1 c0
      !(# s2, c2 #) = L.add_c# a2 b2 c1
      !(# s3, c3 #) = L.add_c# a3 b3 c2
  in  (# (# s0, s1, s2, s3 #), c3 #)
{-# INLINE add_o# #-}

and then the boxed variant, which just uses the above function internally:

add_o
  :: Wider
  -> Wider
  -> (Wider, Word)
add_o (Wider a) (Wider b) =
  let !(# s, Limb c #) = add_o# a b
  in  (Wider s, W# c)

This is actually the same pattern that one gets when using GHC’s UNBOX pragma, and you can typically see this kind of separation between boxed and unboxed variants generated in compiled Core. For this work I wanted to write in terms of unboxed primitives specifically, just to be sure I knew exactly what was going on (I mentioned that one may want to take this approach in Fast Haskell Redux, as well).

The crux here in any case is that, unlike GHC’s native unbounded-size ‘Integer’ or ‘Natural’, ‘Wider’ values are always represented in the exact same manner, whether they have value 2 or 2 ^ 256 - 1, just like ‘Word’ or ‘Int’ are. The same is true for the ‘Montgomery’ types that are used to represent words in Montgomery form:

data Montgomery = Montgomery !(# Limb, Limb, Limb, Limb #)

(N.b., Montgomery form is a representation typically used to avoid expensive modular reductions when working with multiplication over a prime field. The primes used to define the elliptic curve secp256k1 apparently have some special structure that allow fast reductions anyway, eschewing the need for something like Montgomery form, but it works well with them regardless.)

ppad-fixed defines two of these types at present, one for each of the secp256k1 primes. The unboxed/boxed API split present for ‘Wider’ words also exists for the ‘Montgomery’ types.

In any case: the data representation ensures that there should be absolutely no input-dependent variability in allocation when producing values. Using the “unboxed API” internally means we should allocate solely when boxing results, and nowhere else.

Algorithm-Level Semantics

The algorithm-level invariant upheld in constant-time code is that it is “straight-line” code that avoids any branching or indexing on inputs. In ppad-fixed, following crypto-bigint, all functions are constant-time by default unless marked otherwise (and again, following crypto-bigint’s convention, I suffix all variable-time functions with ‘vartime’).

A good example of this kind of algorithmic invariant is seen in the implementation of modular inverse via Fermat inversion, in which there’s an unrolled loop of Montgomery squares and multiplies that proceeds according to a fixed schedule:

-- generated by etc/generate_inv.sh
inv#
  :: (# Limb, Limb, Limb, Limb #)
  -> (# Limb, Limb, Limb, Limb #)
inv# a =
  let -- montgomery 'one'
      !t0 = (# Limb 0x1000003D1##, Limb 0##, Limb 0##, Limb 0## #)
      !t1 = sqr# t0
      !t2 = mul# a t1
      !t3 = sqr# t2
      !t4 = mul# a t3
      --
      -- many intermediate steps omitted..
      --
      !t502 = mul# a t501
      !t503 = sqr# t502
      !t504 = sqr# t503
      !t505 = mul# a t504
      !r = t505
  in  r
{-# INLINE inv# #-}

As you can see from the comment, I generated this function via a bash script that loops through the binary expansion of the (constant, public, known) exponent used for Fermat inversion, printing ‘mul#’ and ‘sqr#’ bindings as appropriate. Since the schedule of multiplies and squares is fixed, then, by construction, the whole sequence runs in constant time, assuming that ‘sqr#’ and ‘mul#’ also do.

It’s also illustrative to look at a generic Montgomery exponentiation function:

exp#
  :: (# Limb, Limb, Limb, Limb #)
  -> (# Limb, Limb, Limb, Limb #)
  -> (# Limb, Limb, Limb, Limb #)
exp# b e =
  let !o = (# Limb 0x1000003D1##, Limb 0##, Limb 0##, Limb 0## #)
      loop !r !m !ex n = case n of
        0 -> r
        _ ->
          let !(# ne, bit #) = Wider.shr1_c# ex
              !candidate = mul# r m
              !nr = select# r candidate bit
              !nm = sqr# m
          in  loop nr nm ne (n - 1)
  in  loop o b e (256 :: Word)
{-# INLINE exp# #-}

This one has no fixed schedule, and so we don’t know at compile time how many ‘mul#’ and ‘sqr#’ calls we’ll need to perform (we introduce a ‘mul#’ call when, looping through the bits of the exponent, the least-significant bit is set). Here, then, instead of computing one or the other branch depending on the LSB, which would introduce argument-dependent variability to the function’s runtime, we compute both branches, and perform a constant-time selection using ‘select#’ to continue the computation with the appropriate value.

(N.b., I use a ‘Data.Choice’ module to implement constant-time primitives like ‘select#’, again mostly following the lead of crypto-bigint. It might be worth splitting this out into its own library at some point.)

So, you get the point: looking at this code at a high, algorithmic level, it’s written explicitly to run in constant time. There is no input-dependent branching or indexing to be found, except in functions that are clearly marked as running in variable time.

Evaluation-Level Semantics

So, the constant-time semantics check out at the algorithm level, but the code needs to pass through GHC, which is well-known to optimize aggressively. It’s thus worth checking out an intermediate representation to verify that GHC doesn’t plan on doing anything undesirable to the code during compilation.

The usual thing that one wants to examine when doing low-level optimisation is Core, which is basically stripped-down Haskell with some optimizations applied, but for execution semantics it’s more useful to analyse an STG dump (where STG, as I’m sure everyone knows, means “Spineless Tagless G-machine,” GHC’s IR that makes operational semantics explicit).

As far as we’re concerned, raw STG differs from Core primarily in that 1) it is fully normalized to A-normal form, 2) thunks are explicit (in that every let binding corresponds directly to a heap allocation), and 3) case expressions represent both sequencing, i.e. forcing of evaluation, and branching.

An STG dump can be procured by building with the ‘-ddump-stg-final’ flag. Take a look at the STG corresponding to the overflowing addition function on ‘Wider’ values shown above, for example. First, the unboxed variant:

add_o#
  :: (# Limb, Limb, Limb, Limb #)
     -> (# Limb, Limb, Limb, Limb #)
     -> (# (# Limb, Limb, Limb, Limb #), Limb #) =
    \r [us us us us us us us us]
        case plusWord2# [us us] of {
        (#,#) c s ->
        case plusWord2# [us us] of {
        (#,#) c0 s0 ->
        case plusWord2# [s0 c] of {
        (#,#) c1 s1 ->
        case plusWord2# [us us] of {
        (#,#) c2 s2 ->
        case or# [c0 c1] of sat {
        __DEFAULT ->
        case plusWord2# [s2 sat] of {
        (#,#) c3 s3 ->
        case plusWord2# [us us] of {
        (#,#) c4 s4 ->
        case or# [c2 c3] of sat {
        __DEFAULT ->
        case plusWord2# [s4 sat] of {
        (#,#) c5 s5 ->
        case or# [c4 c5] of sat {
        __DEFAULT -> (#,,,,#) [s s1 s3 s5 sat];
        [..]
        };

Notice it contains no let bindings, and thus doesn’t allocate (recall raw limbs are stored in registers), and also that each case expression has only a single branch, such that each represents straightforward evaluation sequencing. This is exactly what one wants to see.

For the boxed wrapper, in which the above unboxed function has already been inlined:

add_o :: Wider -> Wider -> (Wider, Word) =
    \r [ds ds1]
        case ds of {
        Wider us us us us ->
        case ds1 of {
        Wider us us us us ->
        case plusWord2# [us us] of {
        (#,#) c s ->
        case plusWord2# [us us] of {
        (#,#) c0 s0 ->
        case plusWord2# [s0 c] of {
        (#,#) c1 s1 ->
        case plusWord2# [us us] of {
        (#,#) c2 s2 ->
        case or# [c0 c1] of sat {
        __DEFAULT ->
        case plusWord2# [s2 sat] of {
        (#,#) c3 s3 ->
        case plusWord2# [us us] of {
        (#,#) c4 s4 ->
        case or# [c2 c3] of sat {
        __DEFAULT ->
        case plusWord2# [s4 sat] of {
        (#,#) c5 s5 ->
        case or# [c4 c5] of sat {
        __DEFAULT ->
        let { sat :: Word = W#! [sat]; } in
        let { sat :: Wider = Wider! [s s1 s3 s5]; } in  (,) [sat sat];
        [..]
        };

This function returns a boxed value, so it allocates, as can be seen by the let bindings at the end. Recall that this function returns a value of type (Wider, Word), so it needs to allocate precisely one ‘Wider’ data constructor, one ‘Word’ data constructor, and one tuple constructor. This is all constant with respect to the function’s inputs, and there is no additional allocation to be found anywhere.

It’s also worth taking a look at the STG for the loop in the exponentiation function shown previously:

$wloop
  :: (# Limb, Limb, Limb, Limb #)
     -> (# Limb, Limb, Limb, Limb #)
     -> (# Limb, Limb, Limb, Limb #)
     -> Word#
     -> (# Limb, Limb, Limb, Limb #) =
    \r [us us us us us us us us us us us us ww]
        case ww<TagProper> of wild2 {
          __DEFAULT ->
              case sqr# us us us us of nm {
              (#,,,#) _ _ _ _ ->
              case mul# us us us us us us us us of {
              (#,,,#) ipv4 ipv5 ipv6 ipv7 ->
              case uncheckedShiftL# [us 63#] of sat {
              __DEFAULT ->
              case uncheckedShiftRL# [sat 63#] of sat {
              __DEFAULT ->
              case not# [sat] of sat {
              __DEFAULT ->
              case plusWord# [sat 1##] of ipv8 {
              __DEFAULT ->
              case minusWord# [wild2 1##] of sat {
              __DEFAULT ->
              case xor# [us ipv7] of sat {
              __DEFAULT ->
              case and# [ipv8 sat] of sat {
              __DEFAULT ->
              case xor# [us sat] of sat {
              __DEFAULT ->
              [..]
              case or# [sat sat] of sat {
              __DEFAULT ->
              exp_$s$wloop1
                  sat sat sat sat ipv ipv1 ipv2 ipv3 sat sat sat sat sat;
              };
              [..]
              };
          0## -> (#,,,#) [us us us us];
        };

This contains no let bindings at all, but you can see two branches in the first case expression, This corresponds to whether or not to continue or stop the loop (it stops when the looping index hits 0), which is deterministic and constant with respect to inputs, so this branch is unproblematic. Note in the meantime that the constant-time select operation used in the loop produces no such branching, but instead is captured by the salad of ‘xor#’ and ‘and#’ expressions in the body of the function (many omitted above), as intended.

For an example of something that does not run in constant time, take a look at the STG for the aptly-named eq_vartime function, which compares for equality in variable time:

eq_vartime :: Montgomery -> Montgomery -> Bool =
    \r [ds ds1]
        case ds of {
        Montgomery us us us us ->
        case ds1 of {
        Montgomery us us us us ->
        case eqWord# [us us] of {
          __DEFAULT -> False [];
          1# ->
              case eqWord# [us us] of {
                __DEFAULT -> False [];
                1# ->
                    case eqWord# [us us] of {
                      __DEFAULT -> False [];
                      1# -> eq_vartime# us us;
                    };
              };
        };
        };
        };

Here you can see a number of branches, each of which depends on the inputs to the function. This is what one would not want to see in code that’s supposed to run in constant time.

Assembly Analysis

All well thus far, in that the data representation is fixed-size, the algorithms used have constant-time and constant-allocation semantics, and that GHC’s evaluation model demonstrably doesn’t want to introduce any unexpected variability to them. But there’s still the matter of the assembly produced – one wants to be sure that the machine instructions that the code ultimately compiles down to are free of any unexpected input-dependent variation.

By my observation, GHC’s LLVM backend produces dramatically faster assembly than does the native code generator, so I’ve adopted it for all of my ppad libraries. To inspect the assembly, we can dump the LLVM IR, and then run it through ‘llc’ with optimizations, examining the result.

Assuming we’re compiling with the LLVM backend, an LLVM IR dump can be gotten via the ‘-ddump-llvm’ flag, and, after removing some headers GHC inserts, we can run it through llc -O3 -mcpu=native -filetype=asm to produce assembly.

Let’s first take a look at the (aarch64) assembly for the overflowing addition function. I’ll paste the entire thing, for completeness:

; %bb.0:                                ; %naoh
  stp x25, x26, [sp, #40]
  stp x23, x24, [sp, #24]
  stp x20, x22, [sp, #8]
  ldp x8, x9, [x20]
  stp x9, x8, [sp, #176]
  ldr x10, [x20, #16]
  str x10, [sp, #168]
  adds  x8, x24, x8
  cset  w10, hs
  adds  x11, x23, x27
  cset  w12, hs
  stp x11, x12, [sp, #152]
  stp x8, x10, [sp, #136]
  adcs  x8, x8, xzr
  cset  w11, hs
  stp x8, x11, [sp, #120]
  adds  x8, x25, x9
  cset  w9, hs
  stp x8, x9, [sp, #104]
  orr x10, x10, x11
  adds  x24, x8, x10
  cset  w8, hs
  stp x24, x8, [sp, #88]
  ldr x10, [sp, #48]
  ldr x11, [sp, #168]
  adds  x10, x10, x11
  cset  w11, hs
  stp x10, x11, [sp, #72]
  orr x8, x9, x8
  adds  x25, x10, x8
  cset  w8, hs
  stp x25, x8, [sp, #56]
  orr x26, x11, x8
  stp x25, x26, [sp, #40]
  ldr x23, [sp, #120]
  stp x23, x24, [sp, #24]
  ldr x22, [sp, #152]
  ldr x8, [sp, #8]
  add x20, x8, #24
  stp x20, x22, [sp, #8]
  ldr x8, [x8, #24]
  blr x8
  ret

The ‘str’ and ‘ldr’ instructions mean store register and load register, and ‘stp’ and ‘ldp’ mean store pair of registers and load pair of registers, respectively, and they all indicate stuff being shuffled around between registers and the stack. These all run in constant time; note that Jon’s Arm Reference contains the following blurb for each:

This instruction is a data-independent-time instruction

The branch with link to register (blr) and return from subroutine (ret) instructions at the end of the computation are “whole program” control-flow instructions that branch unconditionally, and not on secret data. BLR’s target register contains a GHC continuation pointer (i.e., “what we’re doing next”) and RET indicates that this is a subroutine return.

So, that’s all administrative stuff. The actual overflowing addition algorithm is better-captured by the following core sequence of instructions:

  adds  x8, x24, x8
  cset  w10, hs
  adds  x11, x23, x27
  cset  w12, hs
  adcs  x8, x8, xzr
  cset  w11, hs
  adds  x8, x25, x9
  cset  w9, hs
  orr x10, x10, x11
  adds  x24, x8, x10
  cset  w8, hs
  adds  x10, x10, x11
  cset  w11, hs
  orr x8, x9, x8
  adds  x25, x10, x8
  cset  w8, hs
  orr x26, x11, x8
  add x20, x8, #24

The add (add), add, setting flags (adds), add with carry, setting flags (adcs), and bitwise or (orr) instructions are all “data-independent-time,” and “cset” is actually an alias for conditional select increment (csinc), which is a constant-time branchless select. This is exactly what we want to see: that every instruction in this program runs in time constant with respect to its inputs.

The assembly for the variable-time equality function makes for an excellent comparison. Here’s the meat of it, with the “administrative” instructions omitted:

; %bb.0:
  cmp x8, x9
  b.ne  LBB31_4
; %bb.1:
  cmp x8, x9
  b.ne  LBB31_4
; %bb.2:
  cmp x8, x9
  b.ne  LBB31_4
; %bb.3:
  add x20, x8, #40

The core variable-time instruction here is branch conditionally (b.cond), where the condition of interest is “not equal to” (thus ‘b.ne’). This returns the subroutine early if a limb inequality check (performed by ‘cmp’) succeeds. If the subroutine can be timed reliably, especially with chosen inputs, then the timing differential will expose which limb differs first.

Benchmarks and Allocation Measurement

With the code, IR, and assembly all appearing to be in order, the benchmark suites serve mostly as sanity checks. They’re immensely useful, though. I’m reminded a bit of that quote attributed to Knuth, which I’ll paraphrase as “beware of bugs – I’ve only proved the program correct, not tested it.”

ppad-fixed has a Criterion suite that benchmarks functions of interest on varying-sized inputs. We expect variation in wall-clock time to be readily attributable to noise; here’s an example excerpt of the output, for a Montgomery modular exponentiation function:

benchmarking exp/curve:  M(2) ^ 2
time                 5.217 μs   (5.206 μs .. 5.224 μs)
                     1.000 R²   (1.000 R² .. 1.000 R²)
mean                 5.208 μs   (5.198 μs .. 5.215 μs)
std dev              27.24 ns   (20.87 ns .. 37.18 ns)

benchmarking exp/curve:  M(2 ^ 255 - 19) ^ (2 ^ 255 - 19)
time                 5.214 μs   (5.207 μs .. 5.220 μs)
                     1.000 R²   (1.000 R² .. 1.000 R²)
mean                 5.211 μs   (5.201 μs .. 5.217 μs)
std dev              27.82 ns   (16.79 ns .. 51.12 ns)

With a three-nanosecond average differential, the function runs in the same time (up to noise) when called with either small and large inputs, as expected.

A ‘weigh’ suite similarly ensures that there are no surprises when it comes to allocation. Allocations match exactly what is expected across the board, e.g.:

exp

  Case                            Allocated  GCs
  curve:  M(2) ^ 2                       40    0
  curve:  M(2) ^ (2 ^ 255 - 19)          40    0
  scalar:  M(2) ^ 2                      40    0
  scalar:  M(2) ^ (2 ^ 255 - 19)         40    0

You’ll recall that 40 bytes is what a four-limb strict, unboxed tuple behind a single constructor (‘Montgomery’, in this case) should allocate on a 64-bit system.

For functions that return something more complicated, we can always account for the additional allocation precisely. For the variable-time modular square root function on one of the Montgomery domains, for example:

sqrt_vartime

  Case                                  Allocated  GCs
  curve:  sqrt_vartime M(2)                    56    0
  curve:  sqrt_vartime M(2 ^ 255 - 19)         56    0

Here the function returns a ‘Maybe Montgomery’ value, so the additional 16 bytes are for the ‘Just’ constructor and the pointer to its payload.

Conclusion

So: I’ve analysed the entirety of this library in the above manner, and by my appraisal, all functions in this library will run in constant time on aarch64 when compiled with optimizations with GHC 9.8.1, unless marked otherwise. I haven’t performed the same analysis for other compiler versions or flags, or for other target architectures, but the above gives me confidence that the results are also likely to hold true for those too. The Criterion suite should be particularly useful as a sanity check for those cases.

I’ve also performed the same analysis for ppad-secp256k1, and intend to do it for all my other libraries where constant-time execution is expected. Feel free to reproduce my analysis yourself, if curious – and let me know if you do!

Much Faster Signatures on secp256k1

2025-12-22T00:00:00+04:00

I’ve just released a new version of ppad-secp256k1 that supports dramatically faster Schnorr and ECDSA signature schemes, powered by faster elliptic curve arithmetic under the hood.

This follows on my previous posts in what are starting to constitute a sort of secp256k1 series. In the last post I had used the wNAF method for elliptic curve scalar multiplication to glean a large performance boost across the board. This time there have been two significant changes that have resulted in performance wins: first, the use of fixed-width, register-allocated words exclusively, in place of GHC’s heap-allocated ‘Integer’, and second, the use of Montgomery-form arithmetic to avoid performing expensive division operations when working over the secp256k1 fields.

Both of these features are provided by a new library I’ve released, ppad-fixed, which provides such stuff, as well as constant-time selection and arithmetic. ppad-fixed is essentially a Haskell translation of parts of Rust’s crypto-bigint library, except specialized to ‘wide’ (two-limb) and ‘wider’ (four-limb) words, as, unlike in Rust or Golang, we don’t have stack-allocated arrays in Haskell, and so everything is implemented in terms of fixed-size unboxed tuples. I’ve carefully verified that the assembly for many of the arithmetic operations is effectively optimal, at least for darwin-aarch64, when produced with GHC’s LLVM backend, so expect speediness – typically a few nanoseconds for addition and subtraction, a few more for multiplication, and a few microseconds for modular exponentiation, inversion, and square root.

I’d originally mentioned my desire for better fixed-width support in my first post on secp256k1:

Arbitrary-precision Integers are still slow, compared to fixed-width stuff that modern computers can positively chew through. I achieved a staggering speedup on some basic integer parsing by using a custom Word256 type (built from a bunch of Word64 values) under the hood, and converting to Integer only at the end.

They can also be annoying when one wants to achieve constant-time execution of cryptographically-sensitive functions, for obvious reasons. It would be nice to have good fixed-width support for stuff like Word128, Word256, Word512, and so on – I briefly considered implementing and using a custom Word256 type for everything, but this would be a ton of work, and I’m not sure I could beat GHC’s native bigint support for e.g. modular multiplication, exponentiation, and inversion anyway. We’ll stick with plain-old Integer for the time being – it’s still no slouch.

Well, at this point we beat Integer very handily, and provide solid guarantees of constant-time and constant-allocation execution to boot. We’re now getting into the realm of libsecp256k1-tier performance:

  benchmarking schnorr/sign_schnorr' (large)
  time                 43.34 μs   (43.21 μs .. 43.47 μs)
                       1.000 R²   (1.000 R² .. 1.000 R²)
  mean                 43.21 μs   (43.13 μs .. 43.30 μs)
  std dev              280.8 ns   (221.0 ns .. 357.9 ns)

  benchmarking schnorr/verify_schnorr'
  time                 94.57 μs   (94.27 μs .. 94.90 μs)
                       1.000 R²   (1.000 R² .. 1.000 R²)
  mean                 94.13 μs   (93.92 μs .. 94.38 μs)
  std dev              777.4 ns   (644.8 ns .. 925.8 ns)

  benchmarking ecdsa/sign_ecdsa' (large)
  time                 35.21 μs   (35.11 μs .. 35.32 μs)
                       1.000 R²   (1.000 R² .. 1.000 R²)
  mean                 35.14 μs   (35.09 μs .. 35.20 μs)
  std dev              184.4 ns   (135.9 ns .. 241.5 ns)

  benchmarking ecdsa/verify_ecdsa'
  time                 88.70 μs   (88.44 μs .. 88.93 μs)
                       1.000 R²   (1.000 R² .. 1.000 R²)
  mean                 88.93 μs   (88.75 μs .. 89.13 μs)
  std dev              635.0 ns   (517.3 ns .. 963.9 ns)

And yeah, I also implemented Montgomery arithmetic for two of the secp256k1 domains of interest, i.e. those with moduli defined by the secp256k1 field prime and scalar group order, respectively. This allowed me to cut out a bunch of expensive modP/modQ operations everywhere in the Schnorr and ECDSA implementations.

Haskell’s much-maligned ‘Num’ typeclass actually makes these Montgomery form numbers a joy to work with. To crib an illustrative GHCi session from the ppad-fixed README:

  > import qualified Numeric.Montgomery.Secp256k1.Curve as C
  >
  > let one = 1 :: C.Montgomery
  > one
  1
  > putStrLn (C.render one)
  (4294968273, 0, 0, 0)
  >
  > let two = one + one
  > putStrLn (C.render two)
  (8589936546, 0, 0, 0)
  >
  > let big = 2 ^ (128 :: Int) :: C.Montgomery
  > big
  340282366920938463463374607431768211456
  > putStrLn (C.render big)
  (0, 0, 4294968273, 0)
  >
  > let inv = C.inv big
  > inv
  85349562743316995932995116683053049354367560536510302240860302699983992117553
  > putStrLn (C.render inv)
  (0, 0, 1, 0)
  >
  > inv * big
  1

(N.b., I’ve been writing Haskell for something like fourteen years, but to this day I still become amazed by what you can do with Num. It’s great to be able to write in terms of plain integers, but get fixed-width Montgomery stuff automatically handled under the hood, similar to the application I described for automatic differentiation some years ago.)

There’s still room to squeeze out further performance gains here. Some stuff that comes to mind: using Barrett reduction for my modQ operation, in place of the cute Montgomery round-trip I presently do to calculate it; employing the GLV endomorphism that secp256k1 admits in wNAF scalar multiplication; probably going for Bernstein-Yang inversion instead of the fixed-schedule Fermat inversion that I presently use. I’ll get to all of them eventually, as the need or interest arises.

Check it out for all of your Haskell secp256k1 needs, and merely let me know if there are any additional features you could use!

Reproducible Build Artifacts

2025-11-28T00:00:00+04:00

Taproot Assets endeavours to provide completely reproducible and verifiable build artifacts for all releases. These include binaries for a number of target architectures/platforms, along with source code and vendored dependencies, all packaged up in gzipped tarballs, or in zip(1) files. Each release comes with a manifest of these artifacts along with accompanying SHA256 digests.

To attest the artifacts for a release, trusted contributors verify that they can produce the same manifest, and then they cryptographically sign it:

$ gpg --verify manifest-jtobin-v0.7.0.sig manifest-v0.7.0.txt
gpg: Signature made Thu 20 Nov 14:07:09 2025 +04
gpg:                using EDDSA key 7363F82F43B1324E1B0761310E4647D58F8A69E4
gpg: Good signature from "Jared Tobin (https://ppad.tech) <jared@ppad.tech>" [ultimate]
gpg:                 aka "Jared Tobin (jtobin) <jared@jtobin.io>" [ultimate]

If one is operating at an appropriate paranoia level, he can thus personally verify that, first, a quorum of trusted contributors has attested to a certain set of artifacts, and, second, that he can produce, from source and his own tool set, those very same artifacts, before running the software.

Interestingly, the binaries themselves are the least tricky part to verify, at least in principle. Modern Go compilers are able to produce completely reproducible binaries for any target, independent of the build host architecture. That is: given appropriate environment variables and linker flags (namely CGO_ENABLED="0" and -trimpath, probably -s and -w too), a modern Go compiler will produce, bit-for-bit, the same output, given consistent source, other flags, and so on.

You have to do a bit more legwork for the tarballs, etc. Putting together reproducible tarballs involves a salad of flags, e.g.:

  export TZ=UTC

  find "${dir}" -print0 | LC_ALL=C sort -r -z | tar \
    "--mtime=${BUILD_DATE}" --no-recursion --null --mode=u+rw,go+r-w,a+X \
    --owner=0 --group=0 --numeric-owner -c -T - | gzip -9n > "${dir}.tar.gz"

You have to be careful to lock down a bunch of stuff, like timestamps and user/group permissions. What’s more, you have to ensure you’re using the right variant of tar (BSD vs GNU), and ditto for at least gzip and zip.

The problem of producing a hermetic build environment is solved by Nix, of course (or Docker, more on which another time). But some interesting problems can actually pop up when using Nix to build the kinds of reproducible Go binaries I mentioned above.

Let’s start with the tarballs and zip files. Timestamps and permissions are normalized within the Nix store, so, easy: we merely need to put the following in the right place in a Nix flake:

  { nativeBuildInputs = with pkgs; [ coreutils gnutar gzip zip ]; }

Add a minimal set of reproducibility flags to the tar call, and then anyone building the output with the same Nix flake metadata should be able to produce consistent gzipped tarballs containing consistent compiled binaries.

Or so it would seem. When I first tried this, I was able to produce reproducible tarballs for both the project source and its vendored dependencies. The source is declared as follows:

  package = "taproot-assets";
  tag = "v0.7.0";

  # source info
  org = "lightninglabs";
  repo = package;
  tap_src = pkgs.fetchFromGitHub {
    owner = org;
    repo = repo;
    rev = tag;
    hash = "sha256-z4lNeVmy0AEDM8ivdfXfXqi/V1LIDHV2KtS5/19+Jlk=";
  };

and given a ‘tap’ derivation built by Nixpkgs’s buildGoModule, the following, in the buildPhase, can later get us the tarballs of interest:

  tar_gz $out/vendor.tar.gz ${tap.goModules} .
  tar_gz $out/taproot-assets-source-${tag}.tar.gz ${tap_src} .

But strangely, to me, I couldn’t generate reproducible binaries. That is: even though I pinned to a specific Go toolchain, I would produce binaries for different targets that differed depending on the architecture of the host. This was strange, since I know that the Go compiler shouldn’t have any issue doing this, and I was clearly using the right settings:

  # tapd builder
  build_tap = {
        goos ? null
      , goarch ? null
      , goarm ? null
      , tags ? tags_base
      , ldflags ? ldflags_base
      , gcflags ? [ ]
      }: pkgs.buildGoModule {
    pname = package;
    version = tag;
    src = tap_src;
    vendorHash = "sha256-p75eZoM6tSayrxcKTCoXR8Jlc3y2UyzfPCtRJmDy9ew=";
    subPackages = [ "cmd/tapd" "cmd/tapcli" ];

    inherit tags ldflags gcflags;

    preBuild = with pkgs; ''
      export CGO_ENABLED="0"
      export GOAMD64="v1"
      ${lib.optionalString (goos != null) "export GOOS=${goos}"}
      ${lib.optionalString (goarch != null) "export GOARCH=${goarch}"}
      ${lib.optionalString (goarm != null) "export GOARM=${goarm}"}
    '';

    doCheck = false;
  };

(‘-trimpath’, another requirement for reproducible builds, is added by default when using Nixpkgs’s buildGoModule, as is ‘-buildid=’, ensuring no rogue metadata creeps in.)

I dug into this by inspecting a couple of the produced binaries for a random architecture, one built on an amd64 Linux machine and the other on aarch64 Darwin. I produced a couple of dumps of ELF section headers via readelf -S $MY_BINARY and diffed them:

$ diff sections-linux.txt sections-darwin.txt
1c1
< There are 14 section headers, starting at offset 0xd4:
---
> There are 14 section headers, starting at offset 0x386ed30:
18c18
<   [13] .shstrtab         STRTAB          00000000 3870000 000094 00      0   0  1
---
>   [13] .shstrtab         STRTAB          00000000 386ecb0 00007d 00      0   0  1

The .shstrtab (“section name string table”) sections, which are pure ELF metadata, differed, so I produced hexdumps of those, this time via readelf -x .shstrtab $MY_BINARY, and diffed them:

$ diff shstrtab-linux.txt shstrtab-darwin.txt
3,12c3,10
<   0x00000000 002e7465 7874002e 6e6f7074 72646174 ..text..noptrdat
<   0x00000010 61002e64 61746100 2e627373 002e6e6f a..data..bss..no
<   0x00000020 70747262 7373002e 676f2e66 757a7a63 ptrbss..go.fuzzc
<   0x00000030 6e747273 002e676f 2e627569 6c64696e ntrs..go.buildin
<   0x00000040 666f002e 676f2e66 69707369 6e666f00 fo..go.fipsinfo.
<   0x00000050 2e656c66 64617461 002e726f 64617461 .elfdata..rodata
<   0x00000060 002e7479 70656c69 6e6b002e 69746162 ..typelink..itab
<   0x00000070 6c696e6b 002e676f 73796d74 6162002e link..gosymtab..
<   0x00000080 676f7063 6c6e7461 62002e73 68737472 gopclntab..shstr
<   0x00000090 74616200                            tab.
---
>   0x00000000 002e7465 7874002e 6e6f7074 72627373 ..text..noptrbss
>   0x00000010 002e6273 73002e67 6f2e6669 7073696e ..bss..go.fipsin
>   0x00000020 666f002e 676f2e62 75696c64 696e666f fo..go.buildinfo
>   0x00000030 002e7479 70656c69 6e6b002e 69746162 ..typelink..itab
>   0x00000040 6c696e6b 002e7368 73747274 6162002e link..shstrtab..
>   0x00000050 676f7063 6c6e7461 62002e67 6f73796d gopclntab..gosym
>   0x00000060 74616200 2e6e6f70 74726461 7461002e tab..noptrdata..
>   0x00000070 726f6461 7461002e 64617461 00       rodata..data.

Not much to go on there, but I eventually traced this to the fact that mkDerivation in Nixpkgs’s stdenv runs a “fixupPhase” by default, which performs its own stripping of ELF metadata, separate from what the Go compiler does when passed the ‘-s’ and ‘-w’ linker flags (for stripping symbol and debug information, plus DWARF debugging info, respectively). The stripping done by the fixupPhase seems to vary depending on the package set used, e.g. between aarch64-darwin and x86_64-linux.

I disabled the fixupPhase’s stripping by adding

  dontStrip = true;

to my builder function, and tried again. This time the resulting artifacts’ file sizes normalized, but they still produced different SHA256 digests. I produced another hexdump of each, this time using good old xxd(1), and diffed them again:

2028124,2028126c2028124,2028126
< 01ef25b0: 7265 2f32 7334 6870 7137 3368 6e34 396a  re/2s4hpq73hn49j
< 01ef25c0: 6438 3478 3937 366d 3361 6370 3372 6433  d84x976m3acp3rd3
< 01ef25d0: 6b31 782d 747a 6461 7461 2d32 3032 3562  k1x-tzdata-2025b
---
> 01ef25b0: 7265 2f78 686e 6231 6434 7767 6a33 3878  re/xhnb1d4wgj38x
> 01ef25c0: 6833 3833 7973 3162 7969 676e 3079 6861  h383ys1byign0yha
> 01ef25d0: 6371 682d 747a 6461 7461 2d32 3032 3562  cqh-tzdata-2025b
2028308,2028310c2028308,2028310
< 01ef3130: 2069 740a 2f6e 6978 2f73 746f 7265 2f34   it./nix/store/4
< 01ef3140: 6b77 3938 7138 3470 6276 326e 7771 3361  kw98q84pbv2nwq3a
< 01ef3150: 7931 3261 6176 6661 397a 6934 6439 7a2d  y12aavfa9zi4d9z-
---
> 01ef3130: 2069 740a 2f6e 6978 2f73 746f 7265 2f7a   it./nix/store/z
> 01ef3140: 6a62 3467 7763 7971 7670 3862 3236 716b  jb4gwcyqvp8b26qk
> 01ef3150: 7832 7163 3033 3776 6a6a 3631 7838 302d  x2qc037vjj61x80-
2028473,2028475c2028473,2028475
< 01ef3b80: 7374 6f72 652f 7279 6d64 7764 386b 3461  store/rymdwd8k4a
< 01ef3b90: 6268 3430 6233 6d77 3470 6636 3967 3739  bh40b3mw4pf69g79
< 01ef3ba0: 336d 6471 6269 2d69 616e 612d 6574 632d  3mdqbi-iana-etc-
---
> 01ef3b80: 7374 6f72 652f 7038 6369 3239 7637 6b37  store/p8ci29v7k7
> 01ef3b90: 6777 7773 786a 3138 3170 3035 6a63 6362  gwwsxj181p05jccb
> 01ef3ba0: 6e71 6236 3439 2d69 616e 612d 6574 632d  nqb649-iana-etc-
2028626,2028628c2028626,2028628
< 01ef4510: 6f72 652f 7279 6d64 7764 386b 3461 6268  ore/rymdwd8k4abh
< 01ef4520: 3430 6233 6d77 3470 6636 3967 3739 336d  40b3mw4pf69g793m
< 01ef4530: 6471 6269 2d69 616e 612d 6574 632d 3230  dqbi-iana-etc-20
---
> 01ef4510: 6f72 652f 7038 6369 3239 7637 6b37 6777  ore/p8ci29v7k7gw
> 01ef4520: 7773 786a 3138 3170 3035 6a63 6362 6e71  wsxj181p05jccbnq
> 01ef4530: 6236 3439 2d69 616e 612d 6574 632d 3230  b649-iana-etc-20

The difference here was immediately discernable: different binaries contained different hard-coded paths to the Nix store. This was at first immensely confusing to me, but it turns out that the default Go toolchains in Nixpkgs are not “official” Go toolchains, but patched ones. You can see this in the Go derivation for e.g. the nixos-25.05 package set, which contains the following ‘patches’ attribute:

  patches = [
    (replaceVars ./iana-etc-1.25.patch {
      iana = iana-etc;
    })
    # Patch the mimetype database location which is missing on NixOS.
    # but also allow static binaries built with NixOS to run outside nix
    (replaceVars ./mailcap-1.17.patch {
      inherit mailcap;
    })
    # prepend the nix path to the zoneinfo files but also leave the
    # original value for static binaries that run outside a nix server
    (replaceVars ./tzdata-1.19.patch {
      inherit tzdata;
    })
    ./remove-tools-1.11.patch
    ./go_no_vendor_checks-1.23.patch
  ];

The idea here is that the Go compiler bakes into the binaries it produces a number of fallback runtime paths for stuff like ‘tzdata’, e.g. ‘/usr/share/zoneinfo’. On NixOS this, and other paths, do not exist, so the ‘tzdata-1.19.patch’, for example, hard codes a path to the tzdata package in the Nix store that depends on the package set used. Since the package sets differ between platforms, these paths will also differ, and so will any binaries produced by the Go compilers they contain.

The solution to this is easy in Nix, provided you know you actually have to do it. Just build with an unpatched Go toolchain:

  # nixpkgs applies various patches to go toolchains; we need to remove
  # them to produce host-architecture-independent binaries.
  go-unpatched = pkgs.go.overrideAttrs (_: {
    patches = [ ];
  });

  buildGoModule = pkgs.buildGoModule.override {
    go = go-unpatched;
  };

  # tapd builder
  build_tap = {
        goos ? null
      , goarch ? null
      , ..
      , gcflags ? [ ]
      }: buildGoModule {
    ..
  };

This works as expected. The produced manifests look like this:

$ head -3 manifest-v0.7.0-darwin.txt | cut -f1 -d' '
ebebc38f71c7a902341bf4247d8ee2752d0d233deec2034953dab4b44ecb622a
737eb475c5435beedcdb49ae7adf0716c9980523e89682fea4622b59c5d132df
8570b3c659fc2df9cff365fc06bd95459d903b6a3e7c8fd2e5e4efc1a77f7adf

and diffing the manifests produced on both macOS and Linux hosts illustrates that the artifacts match, bit-for-bit, precisely:

$ diff manifest-v0.7.0-darwin.txt manifest-v0.7.0-linux.txt
<empty>

Here’s a gist to a flake for producing various tapd builds, as well as release artifacts, if you’re curious.

Revenge of the Builders

2025-01-17T00:00:00+04:00

I wanted to make a quick addendum to my last post in that I did some further experimenting with Data.ByteString.Builder on both the base16 encoding and decoding tasks, comparing results with the impure direct-allocation-and-write-based implementations I mentioned previously.

I had remarked that builders can be pretty efficient if you’re careful to pack your data aggressively, such that you don’t wind up needing to use too many builders in the first place. I decided to try minimising, as in objectively minimising, the number of builders required in both base16 encoding & decoding, to see what kind of performance I could squeeze out while sticking to the pure API. The builders didn’t disappoint.

How would one “objectively minimise” the number of builders required here? Simply by processing the biggest-sized chunk possible at a time, given we always want to write a Word64. If we can’t do that, we’ll write a Word32 and a Word16 and a Word8. If we can’t do that, we’ll write a Word32 and a Word16. And so on. We can figure all this out just by doing some checks on the length of the input bytestring when starting out: we trade some additional cheap arithmetic operations on machine integers up-front for fewer expensive builder allocations further down the line.

For encoding, this means doing the following checks:

encode :: BS.ByteString -> BS.ByteString
encode bs@(BI.PS _ _ l)
    | l < 64    = to_strict_small loop
    | otherwise = to_strict loop
  where
    loop
      | l `rem` 4 == 0 = go64 bs
      | (l - 3) `rem` 4 == 0 = case BS.splitAt (l - 3) bs of
          (chunk, etc) ->
               go64 chunk
            <> go32 (BU.unsafeTake 2 etc)
            <> go16 (BU.unsafeDrop 2 etc)
      | (l - 2) `rem` 4 == 0 = case BS.splitAt (l - 2) bs of
          (chunk, etc) -> go64 chunk <> go32 etc
      | (l - 1) `rem` 4 == 0 = case BS.splitAt (l - 1) bs of
          (chunk, etc) -> go64 chunk <> go16 etc

      | l `rem` 2 == 0 = go32 bs
      | (l - 1) `rem` 2 == 0 = case BS.splitAt (l - 1) bs of
          (chunk, etc) -> go32 chunk <> go16 etc

      | otherwise = go16 bs

where each ‘go’ function writes words with the indicated number of bits at a time, e.g.:

go64 b = case BS.splitAt 4 b of
  (chunk, etc)
    | BS.null chunk -> mempty
    | otherwise ->
        let !w16_0 = expand_w8 (BU.unsafeIndex chunk 0)
            !w16_1 = expand_w8 (BU.unsafeIndex chunk 1)
            !w16_2 = expand_w8 (BU.unsafeIndex chunk 2)
            !w16_3 = expand_w8 (BU.unsafeIndex chunk 3)

            !w64 = fi w16_0 `B.shiftL` 48
               .|. fi w16_1 `B.shiftL` 32
               .|. fi w16_2 `B.shiftL` 16
               .|. fi w16_3

        in  BSB.word64BE w64 <> go64 etc

and where expand_w8 is just a variant of the previous ‘hilo’ function that returns a Word16 directly, rather than a pair of Word8’s.

‘go32’ works on a chunk of size two, writing a single Word32, and ‘go16’ on a chunk of size one, writing a Word16. The point of having all these functions is that we can now always write the largest-sized word that we can, instead of writing Word8’s or Word16’s exclusively.

(Decoding works similarly, except we need more checks, and an additional ‘go8’ function to handle the one-byte case. I won’t paste the salad of conditionals required here.)

In any case, by employing this strategy we can come pretty close to the performance of the impure implementations. Here are some benchmark results for encoding a 1kb input:

benchmarking encode/ppad-base16
time                 5.929 μs   (5.847 μs .. 6.013 μs)
                     0.999 R²   (0.998 R² .. 0.999 R²)
mean                 5.975 μs   (5.913 μs .. 6.057 μs)
std dev              233.1 ns   (172.4 ns .. 310.0 ns)

benchmarking encode/base16-bytestring
time                 3.246 μs   (3.233 μs .. 3.262 μs)
                     1.000 R²   (1.000 R² .. 1.000 R²)
mean                 3.284 μs   (3.268 μs .. 3.303 μs)
std dev              61.05 ns   (47.83 ns .. 76.77 ns)

benchmarking encode/base16
time                 3.236 μs   (3.221 μs .. 3.253 μs)
                     1.000 R²   (1.000 R² .. 1.000 R²)
mean                 3.244 μs   (3.233 μs .. 3.256 μs)
std dev              37.31 ns   (31.87 ns .. 47.21 ns)

Case                        Allocated  GCs
ppad-base16 (encode)           53,704    0
base16-bytestring (encode)      2,272    0
base16 (encode)                 2,256    0

We’re allocating 25x more than the impure versions, but are only 2x slower or less. Here’s the decoding story:

benchmarking decode/ppad-base16
time                 4.942 μs   (4.884 μs .. 4.995 μs)
                     0.999 R²   (0.998 R² .. 0.999 R²)
mean                 4.908 μs   (4.854 μs .. 4.964 μs)
std dev              176.8 ns   (150.3 ns .. 214.3 ns)
variance introduced by outliers: 46% (moderately inflated)

benchmarking decode/base16-bytestring
time                 540.8 ns   (533.7 ns .. 548.2 ns)
                     0.999 R²   (0.999 R² .. 0.999 R²)
mean                 541.6 ns   (536.9 ns .. 547.5 ns)
std dev              17.64 ns   (13.87 ns .. 22.24 ns)
variance introduced by outliers: 47% (moderately inflated)

benchmarking decode/base16
time                 555.8 ns   (549.7 ns .. 560.9 ns)
                     0.999 R²   (0.999 R² .. 1.000 R²)
mean                 550.7 ns   (546.7 ns .. 555.5 ns)
std dev              15.46 ns   (13.11 ns .. 18.97 ns)
variance introduced by outliers: 39% (moderately inflated)

Case                        Allocated  GCs
ppad-base16 (decode)           21,960    0
base16-bytestring (decode)        128    0
base16 (decode)                 2,440    0

We’re allocating less (we’re writing less), but are closer to 10x slower. Not too bad, all things considered!

(N.b., it’s worth noting that the impure decoding functions also use what appears to be a more efficient lookup table to convert from hex characters back to Word8, so that may account for some of the differential there.)

Fast Haskell, Redux

2025-01-16T00:00:00+04:00

In this post I’m going to incrementally optimise a simple base16 (hexadecimal) encoding routine and illustrate what sort of performance boost each optimisation yields. Hopefully it can be used to glean a bit about what tends to make Haskell code fast – especially code that deals with bytestrings.

You can think of this as a kind of supplement to Chris Done’s Fast Haskell: Competing with C at parsing XML post from a few years ago. Here, like in Chris’s example, we’re going to focus a lot on bytestring handling, though we’ll deal with some different issues than he faced, and also eventually go a little lower-level on the bytestring side of things.

Base16 Encoding

Let’s get right into it. The basic idea here is: for each byte (Word8) in an input, extract its high and low bits, and then map each (effectively a Word4) to a character from the hex alphabet:

import qualified Data.Bits as B
import qualified Data.ByteString as BS

hex_charset :: BS.ByteString
hex_charset = "0123456789abcdef"

hilo :: Word8 -> (Word8, Word8)
hilo b =
  let hi = BS.index hex_charset (fromIntegral b `B.shiftR` 4)
      lo = BS.index hex_charset (fromIntegral b .&. 0b00001111)
  in  (hi, lo)

You then get a base16-encoded output by gluing the resulting characters together in the appropriate fashion:

encode :: BS.ByteString -> BS.ByteString
encode bs = go mempty 0
  where
    l = BS.length bs
    go !acc j
      | j == l = BS.reverse acc
      | otherwise =
          let (hi, lo) = hilo (BS.index bs j)
          in  go (BS.cons lo (BS.cons hi acc)) (succ j)

There are some things here that might stick out to someone accustomed to writing performant Haskell code, but it’s an otherwise reasonable-looking first take. How does it perform on a 1kb input?

benchmarking base16/basic
time                 114.2 μs   (113.1 μs .. 115.4 μs)
                     0.999 R²   (0.998 R² .. 1.000 R²)
mean                 116.1 μs   (115.1 μs .. 118.3 μs)
std dev              4.453 μs   (2.609 μs .. 8.263 μs)
variance introduced by outliers: 38% (moderately inflated)

                 Allocated  GCs
basic            2,326,160    0

Well, according to weigh it seems to allocate a ton. The primary issue is that every invocation of the ‘cons’ function creates a copy of its input bytestring; this is the main thing that would scream out at an experienced Haskeller if they were to glance at the above code. We’re not dealing with O(1) ‘cons’ in bytestring-land, as we are when we use lists.

(As a side note: although ‘BS.reverse’ might raise an eyebrow, it’s actually really fast. It’s a FFI call to a C reverse routine.)

Builders

A more efficient way to construct bytestrings is via Data.ByteString.Builder, which supports constant-time concatenation of sequences of bytes. Here’s a version of ‘encode’ that uses builders:

import qualified Data.ByteString.Builder as BSB

to_strict :: BSB.Builder -> BS.ByteString
to_strict = BS.toStrict . BSB.toLazyByteString
{-# INLINE to_strict #-}

encode :: BS.ByteString -> BS.ByteString
encode bs = to_strict (go 0)
  where
    l = BS.length bs
    go j
      | j == l = mempty
      | otherwise =
          let (hi, lo) = hilo (BS.index bs j)
          in  BSB.word8 hi <> BSB.word8 lo <> go (succ j)

There’s a new function to convert the builder back to a strict bytestring, and now we concatenate builder singletons in order. Simple enough. How does it compare in terms of performance?

benchmarking base16/builder
time                 42.54 μs   (42.01 μs .. 43.27 μs)
                     0.999 R²   (0.999 R² .. 0.999 R²)
mean                 42.88 μs   (42.57 μs .. 43.22 μs)
std dev              1.105 μs   (946.6 ns .. 1.387 μs)
variance introduced by outliers: 24% (moderately inflated)

                 Allocated  GCs
builder            397,768    0

Much better. It allocates about 6x less and is almost 3x faster.

Builders are definitely worth knowing about when dealing with bytestrings, as they’re easy to use, and allow one to write pure code that performs reasonably well. There’s also some fine-tining you can do in order to squeeze additional performance out of them in certain cases. For small inputs, you can use a custom strategy to more efficiently convert builders to lazy bytestrings en route to a strict one, e.g.:

import qualified Data.ByteString.Builder.Extra as BE

to_strict_small :: BSB.Builder -> BS.ByteString
to_strict_small = BS.toStrict
  . BE.toLazyByteStringWith (BE.safeStrategy 128 BE.smallChunkSize) mempty

Using less builders helps as well, and probably even moreso. Consider the following, in which the loop writes a single Word16 at a time instead of two Word8’s:

encode :: BS.ByteString -> BS.ByteString
encode bs@(BI.PS _ _ l)
    | l < 128 = to_strict_small (go 0)
    | otherwise = to_strict (go 0)
  where
    go j
      | j == l = mempty
      | otherwise =
          let (hi, lo) = hilo (BS.index bs j)
              w16 = fromIntegral hi `B.shiftL` 8
                .|. fromIntegral lo
          in  BSB.word16BE w16 <> go (succ j)

It allocates slightly less, and is a microsecond peppier, because there’s just less building going on:

benchmarking base16/better builder
time                 40.96 μs   (40.64 μs .. 41.33 μs)
                     0.999 R²   (0.999 R² .. 1.000 R²)
mean                 41.12 μs   (40.79 μs .. 41.48 μs)
std dev              1.163 μs   (969.0 ns .. 1.494 μs)
variance introduced by outliers: 28% (moderately inflated)

                 Allocated  GCs
better builder     389,592    0

(Note that I’ve also introduced the use of the ‘BI.PS’ pattern synonym here, but only to more easily grab the input bytestring’s length. It has nothing to do with performance.)

Unsafe Functions

Another easy gain can be won by replacing the calls to bytestring’s ‘index’ with its ‘unsafeIndex’ variant:

import qualified Data.ByteString.Unsafe as BU

hilo :: Word8 -> (Word8, Word8)
hilo b =
  let hi = BU.unsafeIndex hex_charset (fromIntegral b `B.shiftR` 4)
      lo = BU.unsafeIndex hex_charset (fromIntegral b .&. 0b00001111)
  in  (hi, lo)

encode :: BS.ByteString -> BS.ByteString
encode bs@(BI.PS _ _ l)
    | l < 128 = to_strict_small (go 0)
    | otherwise = to_strict (go 0)
  where
    go j
      | j == l = mempty
      | otherwise =
          let (hi, lo) = hilo (BU.unsafeIndex bs j)
              w16 = fromIntegral hi `B.shiftL` 8
                .|. fromIntegral lo
          in  BSB.word16BE w16 <> go (succ j)

It often makes sense to do this so long as you can prove that the call is actually safe (the compiler, of course, can’t), as ‘unsafeIndex’ reliably yields a decent performance boost. In this case, the unsafe indexing into the hex alphabet in ‘hilo’ is being done with what are effectively four-bit words, which will always be safe to use as indices in a 16-length bytestring (there are 2^4 = 16 distinct Word4’s). The unsafe index called in the body of the loop can similarly be verified to remain safely within the input bytestring’s bounds, since the loop terminates when its index argument hits the end.

The performance now:

benchmarking base16/unsafe index
time                 25.58 μs   (25.25 μs .. 25.89 μs)
                     0.998 R²   (0.997 R² .. 0.999 R²)
mean                 25.69 μs   (25.41 μs .. 26.03 μs)
std dev              1.051 μs   (852.3 ns .. 1.396 μs)
variance introduced by outliers: 47% (moderately inflated)

                 Allocated  GCs
unsafe index       233,944    0

Another substantial win. Much less allocation and a great reduction in wall-clock time.

Note however that not all unsafe functions will yield a boost as impressive as bytestring’s ‘unsafeIndex’. I’ve never found the ‘unsafeShiftL’ and ‘unsafeShiftR’ functions in Data.Bits to ever really seem to do much at all, for example, so we’ll keep the plain ‘B.shiftR’ call in ‘hilo’ above.

Unboxed Primitives

Next up is another optimisation that any seasoned Haskeller should know about: use unboxed types, unless there’s some unusual reason not to.

Unboxed values require no allocation. To quote the GHC user guide:

The representation of a Haskell Int, for example, is a two-word heap object. An unboxed type, however, is represented by the value itself, no pointers or heap allocation are involved.

Unboxed types correspond to the “raw machine” types you would use in C: Int# (long int), Double# (double), Addr# (void *), etc. The primitive operations (PrimOps) on these types are what you might expect; e.g., (+#) is addition on Int#s, and is the machine-addition that we all know and love—usually one instruction.

You can work with unboxed types and values explicitly by using the appropriate imports and the MagicHash pragma, which can actually be pretty nice, because you express what’s really going on, but more commonly, unboxed types are denoted via strictness annotations and the UNPACK pragma, like so:

data W8Pair = Pair
  {-# UNPACK #-} !Word8
  {-# UNPACK #-} !Word8

hilo :: Word8 -> W8Pair
hilo b =
  let !hi = BU.unsafeIndex hex_charset (fromIntegral b `B.shiftR` 4)
      !lo = BU.unsafeIndex hex_charset (fromIntegral b .&. 0b00001111)
  in  Pair hi lo

encode :: BS.ByteString -> BS.ByteString
encode bs@(BI.PS _ _ l)
    | l < 128 = to_strict_small (go 0)
    | otherwise = to_strict (go 0)
  where
    go j
      | j == l = mempty
      | otherwise =
          let !(Pair hi lo) = hilo (BU.unsafeIndex bs j)
          in  BSB.word8 hi <> BSB.word8 lo <> go (succ j)

Now ‘hilo’ as a whole – so long as one compiles with optimisation – simply doesn’t allocate at all, and the overall allocation and wall-clock time are trimmed by a decent chunk yet again:

benchmarking base16/strict, unpack
time                 20.90 μs   (20.56 μs .. 21.25 μs)
                     0.998 R²   (0.998 R² .. 0.999 R²)
mean                 20.97 μs   (20.76 μs .. 21.19 μs)
std dev              742.1 ns   (619.9 ns .. 938.1 ns)
variance introduced by outliers: 41% (moderately inflated)

                 Allocated  GCs
strict, unpack     176,600    0

At this point the allocation is absolutely dominated by the builders, so if we want to do better we’ll need to do something about them.

Direct Allocation and Writes

A strict ByteString is just a wrapper around some memory. It’s defined in recent versions via:

data ByteString = BS
  {-# UNPACK #-} !(ForeignPtr Word8) -- payload
  {-# UNPACK #-} !Int                -- length

So, a “foreign pointer” to some memory location and a length (where a foreign pointer means a pointer to something that isn’t managed by the RTS in the same way normal data is). To most efficiently create a bytestring, one can thus do it in the same way one would do so in C: allocate some memory and write to it directly.

This is the line beyond which one probably can’t consider his code to be “pure Haskell” anymore. An ‘unsafePerformIO’ call or similar isn’t technically going to be required (it could be masked via ‘Data.ByteString.Internal.unsafeCreate’, for example), and one still doesn’t need to type the {-# LANGUAGE FFI #-} pragma at the top of his module. But it’s safe to say that any assertions of purity would at this point be at least somewhat controversial.

But if one does want to wield this particular Ring of Power, it can be done like so (this is basically what Bryan O’Sullivan’s base16-bytestring or Emily Pillmore’s base16 package do, for example):

import Foreign.Ptr
import Foreign.Storable
import GHC.ForeignPtr
import GHC.Word
import System.IO.Unsafe

encode :: BS.ByteString -> BS.ByteString
encode (BI.PS bs _ l) = unsafeDupablePerformIO $ do
    buffer <- mallocPlainForeignPtrBytes (l * 2)

    withForeignPtr buffer $ \p_buf ->
      withForeignPtr bs $ \p_src ->
        go p_buf p_src (p_src `plusPtr` l)

    pure (BI.BS buffer (l * 2))
  where
    go !buf !src !end
      | src == end = pure ()
      | otherwise = do
          !b <- peek src
          let !(Pair hi lo) = hilo b
          poke buf hi
          poke (buf `plusPtr` 1) lo
          go (buf `plusPtr` 2) (src `plusPtr` 1) end

Here we allocate a buffer explicitly and loop through the input bytestring’s allocated memory, rather than its Haskell representation, in order to populate it, wrapping the result up in a new bytestring via the raw ‘BI.BS’ constructor. It’s worth noting that we could still just make use of ‘unsafeIndex’ on the input bytestring if we wanted, rather than making direct use of its foreign pointer, but since we’re already going big, why not go all the way?

This function allocates minimally, and indeed is at parity with the relevant “encode” functions found in both the base16 and base16-bytestring packages in terms of wall-clock time. It’s about 50x faster than the initial naïve version, and perhaps 5x faster than the most efficient version that used builders:

benchmarking base16/pointer ops
time                 2.929 μs   (2.903 μs .. 2.959 μs)
                     1.000 R²   (0.999 R² .. 1.000 R²)
mean                 2.950 μs   (2.936 μs .. 2.967 μs)
std dev              52.87 ns   (43.32 ns .. 64.89 ns)
variance introduced by outliers: 18% (moderately inflated)

                 Allocated  GCs
pointer ops            120    0

I’m guessing this is probably about as fast as one can make this function via “normal” means. Other techniques that I tried while preparing this article don’t seem to move the needle much on this problem, if at all. I’d be pretty impressed by anything that produced another order-of-magnitude performance boost, though – if anyone can achieve that, I’d love to hear about it!

Recap

So to conclude with a quick, advice-laden summary of the techniques used here:

Avoid using BS.cons, BS.append, and so on, as each requires making a copy of their input bytestring or bytestrings. Instead, prefer builders, and try to pack your data so that you can get away with as few of them as possible. If you can write an occasional Word64, that’s much better than writing many Word8’s.

It’s worth nothing, though, that occasional uses of cons, append, etc. on small inputs can be very benign, and even hard to beat by clever rewriting. You don’t always need to reach for builders or harder-core optimisations every time you want to glue a few bytestrings together. Consider this HMAC implementation from ppad-sha256:
```
hmac
  :: BS.ByteString
  -> BS.ByteString
  -> BS.ByteString
hmac mk@(BI.PS _ _ l) text =
    let step1 = k <> BS.replicate (64 - lk) 0x00
        step2 = BS.map (B.xor 0x36) step1
        step3 = step2 <> text
        step4 = hash step3
        step5 = BS.map (B.xor 0x5C) step1
        step6 = step5 <> step4
    in  hash step6
  where
    !(KeyAndLen k lk)
      | l > 64    = KeyAndLen (hash mk) 32
      | otherwise = KeyAndLen mk l
```
It looks like there’s a lot of unnecessary copying going on here, but in the grand scheme of things, it’s simply not that much, and the bytestring functions are highly optimised, after all. I wasn’t able to improve this function’s performance either by builders or even low-level allocations and writes, including manually minimising allocations, reusing memory areas, and so on. The story would be different if I were calling these functions many times in a loop, though, at which point builders or direct writes would definitely prove useful.
Don’t necessarily shy away from unsafe functions if said use can be proven to be perfectly safe. Aside from ‘unsafeIndex’, other things can be useful, particularly ‘unsafeTake’ and ‘unsafeDrop’. These can provide great performance boosts, though if one is not gunning for absolute, maximum performance, then sure, it might be better to avoid them.

It’s also worth nothing that there are some useful unsafe functions that don’t exist in the bytestring API, but that can be assembled manually. Take this home-cooked unsafe splitAt, for example, that does no bounds checking:
```
data SSPair = SSPair
  {-# UNPACK #-} !BS.ByteString
  {-# UNPACK #-} !BS.ByteString

unsafe_splitAt :: Int -> BS.ByteString -> SSPair
unsafe_splitAt n (BI.BS x l) =
  SSPair (BI.BS x n) (BI.BS (plusForeignPtr x n) (l - n))
```
Make data strict unless you for some reason need it to be lazy, and always unpack strict fields when it’s possible to do so. You can also use -funbox-strict-fields or -funbox-small-strict-fields to unbox all strict fields at the module level, though I generally prefer the more explicit local use of UNPACK pragmas.
Don’t be afraid to use low-level allocation and writes where it makes sense to do so. It may be unnecessary, even in highly-optimised library code (I chose to use builders in ppad-bech32 because we’re never dealing with large outputs there, for example), but it’s a great technique to have around for when performance really counts.

If you’d like to play with the code yourself, I’ve pushed it to this GitHub repo. Presuming you have flake support enabled, you can use ‘nix develop’ to enter a Nix shell, and then e.g. ‘cabal repl’ to open GHCi, or ‘cabal bench’ to run the benchmark suite. Enjoy!

Retvrning to X

2024-12-21T00:00:00+04:00

(UPDATE 2026/02/06: the nukes were redeployed; I rode the bomb down whooping, Major Kong style.)

I nuked all of my social media accounts years ago and never looked back. My general take is that all legacy web 2.0-style social media stuff needs to be ground under a boot, and then the boot filled with cement and sank to the bottom of a particularly deep ocean. Just in case.

That said, I recently recreated an account on X using my old handle, @jaredtobin (some other dude has nabbed @jtobin, unfortunately), and will be experimenting with it for traditional microblogging purposes for things that don’t warrant a full post on this, my trusty long-form blog.

Why not nostr, etc.? Well, I could. I like nostr, and maybe I’ll mirror my stuff there just for the lulz. But, pragmatically, X appears to be where the action is, and I want my stuff to be easy to access and interact with.

It’s an experiment. I’m aloof by nature, so there’s no guarantee I won’t at some point turn up my nose and redeploy the nukes. But we’ll see how it goes. Follow me there, if you’re into that kind of thing!

Faster Signatures on secp256k1

2024-11-25T00:00:00+04:00

In my last post I mentioned that both the Schnorr and ECDSA signature schemes on secp256k1 could be made faster via the so-called wNAF method for elliptic curve point multiplication (short for the cumbersome “w-ary non-adjacent form”). I implemented wNAF for ppad-secp256k1 afterwards, adding a bunch of functions that use it internally.

The only “downside” to the use of wNAF is that one needs to supply a context argument that consists of a bunch of precomputed multiples of the secp256k1 base point (this is at least one of the reasons why libsecp256k1, as well as any bindings to it, generally require you to pass a context argument everywhere). It’s not much of a downside, but it does complicate the API to some degree, so I added these functions on top of the existing ones – the original functions are preserved for when terse code, rather than raw speed, is desired.

The wNAF method really shines when it comes to ECDSA. Here are some benchmarks that illustrate the improvement – the wNAF-powered functions are differentiated by a trailing apostrophe:

benchmarking ecdsa/sign_ecdsa
time                 1.795 ms   (1.767 ms .. 1.822 ms)
                     0.998 R²   (0.997 R² .. 0.999 R²)
mean                 1.806 ms   (1.785 ms .. 1.849 ms)
std dev              93.43 μs   (58.86 μs .. 163.7 μs)

benchmarking ecdsa/sign_ecdsa'
time                 243.1 μs   (237.6 μs .. 249.6 μs)
                     0.996 R²   (0.993 R² .. 0.999 R²)
mean                 241.4 μs   (238.1 μs .. 245.3 μs)
std dev              12.06 μs   (9.492 μs .. 16.17 μs)

benchmarking ecdsa/verify_ecdsa
time                 2.473 ms   (2.409 ms .. 2.541 ms)
                     0.995 R²   (0.991 R² .. 0.997 R²)
mean                 2.432 ms   (2.396 ms .. 2.480 ms)
std dev              140.2 μs   (110.4 μs .. 187.0 μs)

benchmarking ecdsa/verify_ecdsa'
time                 1.460 ms   (1.418 ms .. 1.497 ms)
                     0.994 R²   (0.989 R² .. 0.997 R²)
mean                 1.419 ms   (1.398 ms .. 1.446 ms)
std dev              80.76 μs   (66.73 μs .. 104.9 μs)

So, a 7.5x improvement on ECDSA signature creation, and almost a 2x improvement on signature verification. Not bad.

For Schnorr signatures the improvements are less pronounced, but still substantial:

benchmarking schnorr/sign_schnorr
time                 5.368 ms   (5.327 ms .. 5.423 ms)
                     0.999 R²   (0.999 R² .. 1.000 R²)
mean                 5.436 ms   (5.410 ms .. 5.464 ms)
std dev              84.83 μs   (72.17 μs .. 101.4 μs)

benchmarking schnorr/sign_schnorr'
time                 2.557 ms   (2.534 ms .. 2.596 ms)
                     0.998 R²   (0.997 R² .. 0.999 R²)
mean                 2.579 ms   (2.556 ms .. 2.605 ms)
std dev              83.75 μs   (69.50 μs .. 100.5 μs)

benchmarking schnorr/verify_schnorr
time                 2.338 ms   (2.309 ms .. 2.382 ms)
                     0.998 R²   (0.995 R² .. 0.999 R²)
mean                 2.339 ms   (2.316 ms .. 2.366 ms)
std dev              82.73 μs   (65.83 μs .. 105.9 μs)

benchmarking schnorr/verify_schnorr'
time                 1.429 ms   (1.381 ms .. 1.482 ms)
                     0.993 R²   (0.987 R² .. 0.998 R²)
mean                 1.372 ms   (1.355 ms .. 1.396 ms)
std dev              67.72 μs   (50.44 μs .. 110.5 μs)

So here about a 2x improvement, plus or minus change, across the board. I hope to hammer these down a good bit further in the future if at all possible!

Signatures on secp256k1

2024-10-19T00:00:00+04:00

I’ve released a library supporting BIP340 Schnorr signatures and deterministic ECDSA on the elliptic curve secp256k1. Get it while it’s hot – for when you just aren’t feeling libsecp256k1!

This is another “minimal” library in the ppad suite of libraries I’m working on. Minimal in the sense that it is pure Haskell (no FFI – you can check out ppad-csecp256k1 if you want that) and depends only on ‘base’, ‘bytestring’, and my own HMAC-DRBG and SHA256 libraries. The feature set also intentionally remains rather lean for the time being (though if you could use other features in there, let me know!).

Performance is decent, though unsurprisingly it still pales in comparison to the low-level and battle-hardened libsecp256k1 (think 5ms vs 50μs to create a Schnorr signature, for example). There’s ample room for optimisation, though. Probably the lowest-hanging fruit is that scalar multiplication on secp256k1 can seemingly be made much more efficient via the so-called wNAF method that relies on precomputed points, such that we might be looking at more like 500μs to create a Schnorr signature, with a similar improvement for ECDSA. It would require slightly more annoying UX, probably warranting its own set of user-facing functions that would also accept a context argument, but does not appear difficult to implement.

A few things I observed or noted while writing this library:

The modular arithmetic functions for arbitrary-precision Integers contained in GHC.Num.Integer can be extremely fast, compared to hand-rolled alternatives. Things like integerPowMod# and integerRecipMod# absolutely fly, and will probably beat any hand-rolled variant.
Arbitrary-precision Integers are still slow, compared to fixed-width stuff that modern computers can positively chew through (this is not news to me, but still). I achieved a staggering speedup on some basic integer parsing by using a custom Word256 type (built from a bunch of Word64 values) under the hood, and converting to Integer only at the end.

They can also be annoying when one wants to achieve constant-time execution of cryptographically-sensitive functions, for obvious reasons. It would be nice to have good fixed-width support for stuff like Word128, Word256, Word512, and so on – I briefly considered implementing and using a custom Word256 type for everything, but this would be a ton of work, and I’m not sure I could beat GHC’s native bigint support for e.g. modular multiplication, exponentiation, and inversion anyway. We’ll stick with plain-old Integer for the time being – it’s still no slouch.
Algorithmically constant-time operations can still fail to be constant-time in practice due to factors largely outside the programmer’s control. A good example is found in bit operations; looping through a bytestring and performing some “equivalent-looking” work on every byte may still result in execution time discrepancies between zero and nonzero bytes, for example, and these can be very hard to eliminate. This stuff can depend on the compiler or runtime, the architecture/processor used, etc.
The necessary organization of this kind of “catch-all” library is kind of unsatisfying. Rather than picking a single curve, and then implementing every feature one can possibly think of for it, it would intuitively be better to implement a generic curve library or libraries (for Weierstrass, Edwards, Montgomery, etc.), and then implement e.g. ECDSA or EdDSA or ECDH or or whatever as separate libraries depending on those generic curves as appropriate. One could then use everything in more of a plug-and-play fashion – this might be a design I’ll explore further in the future.
ByteString seems to provide a good user-facing “interface” for libraries like this. It’s reliable, familiar to Haskellers, has a great API, and is very well-tuned. It’s possible one might want to standardize on something else for internals, though; unboxed vectors are an obvious choice, though I would actually be inclined to use the primitive library’s PrimArrays directly, favouring simplicity, and eschewing vector’s harder-core optimisations.

The idea here in any case would be that one would use ByteString only at the user-facing layer, and then work with PrimArrays (or whatever) everywhere internally. It’s perhaps worth exploring further – bytestring is very fast (strict bytestrings are, after all, merely pointers to cstrings), but so are PrimArrays, and mutation à la MutablePrimArray could be very helpful to have here and there.

(FWIW, though, I’ve benchmarked PrimArray/MutablePrimArray in ppad-sha256 and found them to yield equivalent-or-slower performance compared to ByteString in that setting.)

The library has been tested on the Project Wycheproof and official BIP340 test vectors, as well as noble-secp256k1’s test suite, and care has been taken around timing of functions that operate on secret data. Kick the tires on it, if you feel so inclined!

(I mentioned this in my last post as well, but I’m indebted to Paul Miller’s noble cryptography project for this work, both as inspiration and also as a reference.)

New HMAC-DRBG and SHA-2 Libraries

2024-10-07T00:00:00+04:00

Just FYI, I’ve dropped a few simple libraries supporting SHA-{256,512}, HMAC-SHA{256, 512}, and HMAC-DRBG. You can find the repos here:

Each is packaged there as a Nix flake, and each is also available on Hackage.

This is the first battery of a series of libraries I’m writing that were primarily inspired by noble-cryptography after the death (or at least deprecation) of cryptonite. The libraries are pure, readable, concise GHC Haskell, having minimal dependencies, and aim for clarity, security, performance, and user-friendliness.

I finally got around to going through most of the famous cryptopals challenges last year and have since felt like writing some “foundational” cryptography (and cryptography-adjacent) libraries that I myself would want to use. I’d like to understand them well, test and benchmark them myself, eke out performance and UX wins where I can get them, etc. etc.

An example is found in the case of ppad-hmac-drbg – I want to use this DRBG in the manner I’m accustomed to using generators from e.g. mwc-random, in which I can utilise any PrimMonad to handle the generator state:

ghci> :set -XOverloadedStrings
ghci> import qualified Crypto.DRBG.HMAC as DRBG
ghci> import qualified Crypto.Hash.SHA256 as SHA256
ghci>
ghci> let entropy = "very random"
ghci> let nonce = "very unused"
ghci> let personalization_string = "very personal"
ghci>
ghci> drbg <- DRBG.new SHA256.hmac entropy nonce personalization_string
ghci> bytes <- DRBG.gen mempty 32 drbg
ghci> more_bytes <- DRBG.gen mempty 16 drbg

I haven’t actually tried the other DRBG library I found on Hackage, but it has different UX, a lot of dependencies, and has since apparently been deprecated. The generator in ppad-hmac-drbg matches my preferred UX, passes the official DRBGVS vectors, depends only on ‘base’, ‘bytestring’, and ‘primitive’, and can be used with arbitrary appropriate HMAC functions (and is maintained!).

The ppad-sha256 and ppad-sha512 libraries depend only on ‘base’ and ‘bytestring’ and are faster than any other pure Haskell SHA-2 implementations I’m aware of, even if the performance differences are moral, rather than practical, victories:

ppad-sha256’s SHA-256, vs SHA’s, on a 32B input:

  benchmarking ppad-sha256/SHA256 (32B input)/hash
  time                 1.898 μs   (1.858 μs .. 1.941 μs)
                       0.997 R²   (0.996 R² .. 0.999 R²)
  mean                 1.874 μs   (1.856 μs .. 1.902 μs)
  std dev              75.90 ns   (60.30 ns .. 101.8 ns)
  variance introduced by outliers: 55% (severely inflated)

  benchmarking SHA/SHA256 (32B input)/sha256
  time                 2.929 μs   (2.871 μs .. 2.995 μs)
                       0.997 R²   (0.995 R² .. 0.998 R²)
  mean                 2.879 μs   (2.833 μs .. 2.938 μs)
  std dev              170.4 ns   (130.4 ns .. 258.9 ns)
  variance introduced by outliers: 71% (severely inflated)

And the same, but now a HMAC-SHA256 battle:

  benchmarking ppad-sha256/HMAC-SHA256 (32B input)/hmac
  time                 7.287 μs   (7.128 μs .. 7.424 μs)
                       0.996 R²   (0.995 R² .. 0.998 R²)
  mean                 7.272 μs   (7.115 μs .. 7.455 μs)
  std dev              565.2 ns   (490.9 ns .. 689.7 ns)
  variance introduced by outliers: 80% (severely inflated)

  benchmarking SHA/HMAC-SHA256 (32B input)/hmacSha256
  time                 11.42 μs   (11.09 μs .. 11.80 μs)
                       0.994 R²   (0.992 R² .. 0.997 R²)
  mean                 11.36 μs   (11.09 μs .. 11.61 μs)
  std dev              903.5 ns   (766.5 ns .. 1.057 μs)
  variance introduced by outliers: 79% (severely inflated)

The performance differential is larger on larger inputs; I think the difference between the two on a contrived 1GB input was 22 vs 32s, all on my mid-2020 MacBook Air. I haven’t bothered to implement e.g. SHA-224 and SHA-384, which are trivial adjustments of SHA-256 and SHA-512, but if anyone could use them for some reason, please just let me know.

Anyway: enjoy, and let me know if you get any use out of these. Expect more releases in this spirit!

Reservoir Sampling

2024-09-22T00:00:00+04:00

I have a little library called sampling floating around for general-purpose sampling from arbitrary foldable collections. It’s a bit of a funny project: I originally hacked it together quickly, just to get something done, so it’s not a very good library qua library – it has plenty of unnecessary dependencies, and it’s not at all tuned for performance. But it was always straightforward to use, and good enough for my needs, so I’ve never felt any particular urge to update it. Lately it caught my attention again, though, and I started thinking about possible ways to revamp it, as well as giving it some much-needed quality-of-life improvements, in order to make it more generally useful.

The library supports sampling with and without replacement in both the equal and unequal-probability cases, from collections such as lists, maps, vectors, etc. – again, anything with a Foldable instance. In particular: the equal-probability, sampling-without-replacement case depends on some code that Gabriella Gonzalez wrote for reservoir sampling, a common online method for sampling from a potentially-unbounded stream. I started looking into alternative algorithms for reservoir sampling, to see what else was out there, and discovered one that was extremely fast, extremely clever, and that I’d never heard of before. It doesn’t seem to be that well-known, so I simply want to illustrate it here, just so others are aware of it.

I learned about the algorithm from Erik Erlandson’s blog, but, as he points out, it goes back almost forty years to one J.S Vitter, who apparently also popularised the “basic” reservoir sampling algorithm that everyone uses today (though it was not invented by him).

The basic reservoir sampling algorithm is simple. We’re sampling some number of elements from a stream of unknown size: if we haven’t collected enough elements yet, we just dump everything we see into our “reservoir” (i.e., the collected sample); otherwise, we just generate a random number and do a basic comparison to determine whether or not to eject a value in our reservoir in favour of the new element we’ve just seen. This is evidently also known as Algorithm R, and can apparently be found in Knuth’s Art of Computer Programming. Here’s a basic implementation that treats the reservoir as a mutable vector:

algo_r len stream prng = do
    reservoir <- VM.new len
    loop reservoir 0 stream
  where
    loop !sample index = \case
      [] ->
        V.freeze sample

      (h:t)
        | index < len -> do
            VM.write sample index h
            loop sample (succ index) t

        | otherwise -> do
            j <- S.uniformRM (0, index - 1) prng
            when (j < len)
              (VM.write sample j h)
            loop sample (succ index) t

Here’s a quick, informal benchmark of it. Sampling 100 elements from a stream of 10M 64-bit integers, using Marsaglia’s MWC256 PRNG, yields the following runtime statistics on my mid-2020 MacBook Air:

  73,116,027,160 bytes allocated in the heap
      14,279,384 bytes copied during GC
          45,960 bytes maximum residency (2 sample(s))
          31,864 bytes maximum slop
               6 MiB total memory in use (0 MiB lost due to fragmentation)

                                     Tot time (elapsed)  Avg pause  Max pause
  Gen  0     17639 colls,     0 par    0.084s   0.123s     0.0000s    0.0004s
  Gen  1         2 colls,     0 par    0.000s   0.000s     0.0002s    0.0003s

  INIT    time    0.007s  (  0.006s elapsed)
  MUT     time   18.794s  ( 18.651s elapsed)
  GC      time    0.084s  (  0.123s elapsed)
  RP      time    0.000s  (  0.000s elapsed)
  PROF    time    0.000s  (  0.000s elapsed)
  EXIT    time    0.000s  (  0.000s elapsed)
  Total   time   18.885s  ( 18.781s elapsed)

  %GC     time       0.0%  (0.0% elapsed)

  Alloc rate    3,890,333,621 bytes per MUT second

  Productivity  99.5% of total user, 99.3% of total elapsed

It’s fairly slow, and pretty much all we’re doing is generating 10M random numbers.

In my experience, the most effective optimisations that can be made to a numerical algorithm like this tend to be “mechanical” in nature – avoiding allocation, cache misses, branch prediction failures, etc. I find it exceptionally pleasing when there’s some domain-specific intuition that admits substantial optimisation of an algorithm.

Vitter’s optimisation is along these lines. I find it as ingenious as e.g. the kernel trick in the support vector machine context. The crucial observation Vitter made is that one doesn’t need to consider whether or not to add every single element he encounters to the reservoir; the “gap” between entries also follows a well-defined probability distribution, and one can just instead sample that in order to determine the next element to add. Erik Erlandson points out that when the size of the reservoir is small relative to the size of the stream, as is typically the case, this distribution is well-approximated by the geometric distribution – that which describes the number of coin flips required before one observes a head (it has come up in a couple of my previous blog posts).

So: the algorithm is adjusted so that, while processing the stream, we sample how many elements to skip, and do so, then adding the next element encountered after that to the reservoir. Here’s a version of that, using Erlandson’s fast geometric approximation for sampling what Vitter calls the skip distance:

data Loop =
    Next
  | Skip !Int

algo_r_optim len stream prng = do
    reservoir <- VM.new len
    go Next reservoir 0 stream
  where
    go what !sample !index = \case
      [] ->
        V.freeze sample

      s@(h:t) -> case what of
        Next
          -- below the reservoir size, just write elements
          | index < len -> do
              VM.write sample index h
              go Next sample (succ index) t

          -- sample skip distance
          | otherwise -> do
              u <- S.uniformDouble01M prng
              let p = fi len / fi index
                  skip = floor (log u / log (1 - p))
              go (Skip skip) sample index s

        Skip skip
          -- stop skipping, use this element
          | skip == 0 -> do
              j <- S.uniformRM (0, len - 1) prng
              VM.write sample j h
              go Next sample (succ index) t

          -- skip (d - 1) more elements
          | otherwise ->
              go (Skip (pred d)) sample (succ index) t

And now the same simple benchmark:

   1,852,883,584 bytes allocated in the heap
         210,440 bytes copied during GC
          45,960 bytes maximum residency (2 sample(s))
          31,864 bytes maximum slop
               6 MiB total memory in use (0 MiB lost due to fragmentation)

                                     Tot time (elapsed)  Avg pause  Max pause
  Gen  0       445 colls,     0 par    0.002s   0.003s     0.0000s    0.0002s
  Gen  1         2 colls,     0 par    0.000s   0.000s     0.0002s    0.0003s

  INIT    time    0.007s  (  0.007s elapsed)
  MUT     time    0.286s  (  0.283s elapsed)
  GC      time    0.002s  (  0.003s elapsed)
  RP      time    0.000s  (  0.000s elapsed)
  PROF    time    0.000s  (  0.000s elapsed)
  EXIT    time    0.000s  (  0.000s elapsed)
  Total   time    0.296s  (  0.293s elapsed)

  %GC     time       0.0%  (0.0% elapsed)

  Alloc rate    6,477,345,638 bytes per MUT second

  Productivity  96.8% of total user, 96.4% of total elapsed

Both Vitter and Erlandson estimated orders of magnitude improvement in sampling time, given we need to spend much less time iterating our PRNG; here we see a 65x performance gain, with 40x less allocation. Very impressive, and again, the optimisation is entirely probabilistic, rather than “mechanical,” in nature (indeed, I haven’t tested any mechanical optimisations to these implementations at all).

It turns out there are extensions in this spirit to unequal-probability reservoir sampling as well, as is the method of “exponential jumps” described in a 2006 paper by Efraimidis and Spirakis. I’ll probably benchmark that algorithm too, and, if it fits the bill, and I ever really do get around to properly updating my ‘sampling’ library, I’ll refactor things to make use of these high-performance algorithms. Let me know if you could use this sort of thing!

More Recursive Stochastic Processes

2024-09-01T00:00:00+04:00

Some years ago I wrote about using recursion schemes to encode stochastic processes in an embedded probabilistic programming setting. The crux of it was that recursion schemes allow one to “factor out” the probabilistic phenomena from the recursive structure of the process; the probabilistic stuff typically sits in the so-called coalgebra of the recursion scheme, while the recursion scheme itself dictates the manner in which the process evolves.

While it’s arguable as to whether this stuff is useful from a strictly practical perspective (I would suggest “not”), I think the intuition one gleans from it can be somewhat worthwhile, and my curious brain finds itself wandering back to the topic from time to time.

I happened to take a look at this sort of framework again recently and discovered that I couldn’t easily seem to implement a Chinese Restaurant Process (CRP) – a stochastic process famous from the setting of nonparametric Bayesian models – via either of the “standard” patterns I used throughout my Recursive Stochastic Processes post. This indicated that the recursive structure of the CRP differs from the others I studied previously, at least in the manner I was attempting to encode it in my particular embedded language setting.

So let’s take a look to see what the initial problem was, how one can resolve it, and what insights we can take away from it all.

Framework

Here’s a simple and slightly more refined version of the embedded probabilistic programming framework I introduced in some of my older posts. I’ll elide incidental details, such as common imports and noisy code that might distract from the main points.

First, some unfamiliar imports and minimal core types:

import qualified Control.Monad.Trans.Free as TF
import qualified System.Random.MWC.Probability as MWC

data ModelF r =
    BernoulliF Double (Bool -> r)
  | UniformF (Double -> r)
  deriving Functor

type Model = Free ModelF

As a quick refresher, an expression of type ‘Model’ denotes a probability distribution over some carrier type, and terms that construct, manipulate, or interpret values of type ‘Model’ constitute those of a simple embedded probabilistic programming language.

Importantly, here we’re using a very minimal such language, consisting only of two primitives: the Bernoulli distribution (which is a probability distribution over a coin flip) and the uniform distribution over the interval [0, 1]. A trivial model that draws a probability of success from a uniform distribution, and then flips a coin conditional on that probability, would look like this, for example:

model :: Model Bool
model = Free (UniformF (\u ->
  Free (BernoulliF pure))

(We could create helper functions such that programs in this embedded language would be nicer to write, but that’s not the focus of this post.)

We can construct expressions that denote stochastic processes by using the normal salad of recursion schemes. Here’s how we can encode a geometric distribution, for example, in which one counts the number of coin flips required to observe a head:

geometric :: Double -> Model Int
geometric p = apo coalg 1 where
  coalg n = TF.Free (BernoulliF p (\accept ->
    if   accept
    then Left (pure n)
    else Right (n + 1)))

The coalgebra isolates the probabilistic phenomena (a Bernoulli draw, i.e. a coin flip), and the recursion scheme determines how the process evolves (halting if a Bernoulli proposal is accepted). The coalgebra is defined in terms of the so-called base or pattern functor of the free monad type, defined in ‘Control.Monad.Trans.Free’ (see Practical Recursion Schemes for a refresher on base functors if you’re rusty).

The result is an expression in our embedded language, of course, and is completely abstract. If we want to sample from the model it encodes, we can use an interpreter like the following:

prob :: Model a -> MWC.Prob IO a
prob = iterM $ \case
  BernoulliF p f -> MWC.bernoulli p >>= f
  UniformF f     -> MWC.uniform >>= f

where the ‘MWC’-prefixed functions are sampling functions from the mwc-probability library, and ‘iterM’ is the familiar monadic catamorphism-like recursion scheme over the free monad. This will produce a function that can be sampled with ‘MWC.sample’ when provided with a PRNG.

Here’s what some samples look like:

ghci> gen <- MWC.create
ghci> replicateM 10 (MWC.sample (prob (geometric 0.1)) gen)
[1,9,13,3,4,3,13,1,4,17]

Chinese Restaurant Process

The CRP is described technically as a stochastic process over “finite integer partitions,” and, more memorably, over configurations of a particular sort of indefinitely-large Chinese restaurant. One is to imagine customers entering the restaurant sequentially; the first sits at the first table available, and each additional customer is either seated at a new table with probability proportional to some dispersion parameter, or is seated at an occupied table with probability proportional to the number of other customers already sitting there.

If each customer is labelled by how many others were in the restaurant when they entered, the result is, for ‘n’ total customers, a random partition of the natural numbers up to ‘n’. A particular realisation of the process, following the arrival of 10 customers, might look like the following:

[[9,8,7,6,4,2,0], [1], [5,3]]

This is a restaurant configuration after 10 arrivals, where each element is a table populated by the labelled customers.

One of the rules of explaining the CRP is that you always have to include a visualization like this:

It corresponds to the realised sample, and we certainly wouldn’t want to break any pedagogical regulations.

Encoding, First Attempt

A natural way to encode the process is to seat each arriving customer at a new table with the appropriate probability, and then, if it turns out he is to be sat at an occupied table, to do that with the appropriate conditional probability (i.e., conditional on the fact that he’s not going to be seated at a new table).

So let’s imagine encoding a CRP with dispersion parameter ‘a’ and total number of customers ‘n’. Our first attempt might go something like this:

crp n a = ana coalg (1, <initial restaurant>) where
  coalg (customer, tables)
    | customer >= n = TF.Pure tables
    | otherwise =
        let p = <probability of seating customer at new table>
        in  TF.Free (BernoulliF p (\accept ->
              if   accept
              then (succ customer, <seat at new table>)
              else ???
            ))

We run into a problem when we hit the ‘else’ branch of the conditional. Here we want to express another random choice – viz., at which occupied table do we seat the arriving customer. We’d want to do something like ‘TF.Free (UniformF (\u -> …))’ in that ‘else’ branch and then use the produced uniform value to choose between the occupied tables based on their conditional probabilities. But that will prove to be impossible, given the type requirements.

There’s similarly no way to restructure things by producing the desired uniform value earlier, before the conditional expression. For appropriate type ‘t’, the coalgebra used for the anamorphism must have type:

coalg :: a -> Base t a

which in our case reduces to:

coalg :: a -> TF.FreeF ModelF a a

You’ll find that there’s simply no way to add another TF.Free constructor to the mix while satisfying the above type. So it seems that with an anamorphism (or apomorphism, which encounters the same problem) we’re limited to denoting a single probabilistic operation on any recursive call.

Encoding, Correctly

We thus need a recursion scheme that allows us to add multiple levels of the base functor at a time. The most appropriate scheme appears to me to be the futumorphism, which I also wrote about in Time Traveling Recursion Schemes.

(As Patrick Thomson pointed out in his sublime series on recursion schemes, both anamorphisms and apomorphisms are special cases of the futumorphism.)

The coalgebra used by a futumorphism has a different type than that used by an ana- or apomorphism, namely:

coalg :: a -> Base t (Free (Base t) a)

In our case, this is:

coalg :: a -> TF.FreeF ModelF a (Free (TF.FreeF ModelF a) a)

Note that here we’ll be able to use additional ‘TF.Free’ constructors inside other expressions involving the base functor. In Time Traveling Recursion Schemes I referred to this as “working with values that don’t exist yet, while pretending like they do” – but better would be to say that one is simply defining values to be produced during (co)recursion, using separate monadic code.

Here’s how we’d encode the CRP using a futumorphism, in pseudocode:

crp n a = futu coalg (1, <initial restaurant>) where
  coalg (customer, tables)
    | customer >= n = TF.Pure tables
    | otherwise =
        let p = <probability of seating customer at new table>
        in  TF.Free (BernoulliF p (\accept ->
              if   accept
              then pure (succ customer, <seat at new table>)
              else do
                res <- liftF (TF.Free (UniformF (\u ->
                         <seat amongst occupied tables using 'u'>
                       )))
                pure (succ customer, res)))

Note that recursive points are expressed under a ‘TF.Free’ constructor by returning a value (using ‘pure’) in the free monad itself. This effectively allows us to use more than one embedded language construct on each recursive call – as many as we want, as a matter of fact. The familiar ‘liftF’ function lifts each such expression into the free monad for us.

I mentioned in Time Traveling Recursion Schemes that this sort of thing can look a little bit nicer if you have the appropriate embedded language terms floating around. If we define the following:

uniform :: Free (TF.FreeF ModelF a) Double
uniform = liftF (TF.Free (UniformF id))

then we can use it to tidy things up a bit:

crp n a = futu coalg (1, <initial restaurant>) where
  coalg (customer, tables)
    | customer >= n = TF.Pure tables
    | otherwise =
        let p = <probability of seating customer at new table>
        in  TF.Free (BernoulliF p (\accept ->
              if   accept
              then pure (succ customer, <seat at new table>)
              else do
                u <- uniform
                let res = <seat amongst occupied tables using 'u'>
                pure (succ customer, res)))

In any case, the futumorphism gets us to a faithfully-encoded CRP. It allows us to express multiple probabilistic operations in every recursive call, which is what’s required here due to our use of conditional probabilities.

Alternative Encodings

Now. Recall that I mentioned the following:

A natural way to encode the process is to seat each arriving customer at a new table with the appropriate probability, and then, if it turns out he is to be sat at an occupied table, to do that with the appropriate conditional probability (i.e., conditional on the fact that he’s not going to be seated at a new table).

This is probably the most natural way to encode the process, and indeed, if we only have Bernoulli and uniform language terms (or beta, Gaussian, etc. in place of uniform – something that can at least be transformed to produce a uniform, in any case), this seems to be the best we can do. But it is not the only way to encode the CRP. If we have a language term corresponding to a categorical distribution, then we can instead choose between a new table and any of the occupied tables simultaneously using the unconditional probabilities for each.

Let’s adjust our ModelF functor, adding a language term corresponding to a categorical distribution. It will have as its parameter a list of probabilities, one for each possible outcome under consideration:

data ModelF r =
    BernoulliF Double (Bool -> r)
  | UniformF (Double -> r)
  | CategoricalF [Double] (Int -> r)
  deriving Functor

With this we’ll be able to encode the CRP using a mere anamorphism, as the following pseudocode describes:

crp n a = ana coalg (1, <initial restaurant>) where
  coalg (customer, tables)
    | customer >= n = TF.Pure tables
    | otherwise =
        let ps = <unconditional categorical probabilities>
        in  TF.Free (CategoricalF ps (\i ->
              if   i == 0
              then (succ customer, <seat at new table>)
              else (succ customer, <seat at occupied table 'i - 1'>)
            ))

Since here we only ever express a single probabilistic term during recursion, the “standard” pattern we’ve used previously applies here. But it’s worth noting that we could still employ the “unconditional, then conditional” approach using a categorical distribution – we’d just use the conditional probabilities to sample from the occupied tables directly, rather than doing it in more low-level fashion using the single uniform draw. Given an appropriate ‘categorical’ term, that would look more like:

crp n a = futu coalg (1, <initial restaurant>) where
  coalg (customer, tables)
    | customer >= n = TF.Pure tables
    | otherwise =
        let p = <probability of seating customer at new table>
        in  TF.Free (BernoulliF p (\accept ->
              if   accept
              then pure (succ customer, <seat at new table>)
              else do
                let ps = <conditional categorical probabilities>
                i <- categorical ps
                pure (succ customer, <seat at occupied table 'i'>)))

The recursive structure in this case is more complicated, requiring a futumorphism instead of an anamorphism, but typically the conditional probabilities are easier to compute.

Fin

So, some takeaways here, if one wants to indulge in this sort of framework:

If we want to express a stochastic process using multiple probabilistic operations in any given recursive call, we may need to employ a scheme that supports richer (co)recursion than a plain anamorphism or apomorphism. Here that’s captured nicely by a futumorphism, which naturally captures what we’re looking for.

The same might be true if our functor is sufficiently limited, as was the case here if we supported only the Bernoulli and uniform distributions. There we had no option but to express the recursion using condiitonal probabilistic operations, and so needed the richer recursive structure that a futumorphism provides.

On the other hand, it may not be the case that one needs the richer structure provided by a futumorphism, if instead one can express the coalgebra using only a single layer of the base functor. Adding a primitive categorical distribution to our embedded language eliminated the need to use conditional probabilities when describing the recursion, allowing us to drop back to a “basic” anamorphism.

Here is a gist containing a fleshed-out version of the code above, if you’d like to play with it. Enjoy!

Kelvin Versioning

2020-02-25T00:00:00+04:00

Long ago, in the distant past, Curtis introduced the idea of kelvin versioning in an informal blog post about Urbit. Imagining the idea of an ancient and long-frozen form of Martian computing, he described this versioning scheme as follows:

Some standards are extensible or versionable, but some are not. ASCII, for instance, is perma-frozen. So is IPv4 (its relationship to IPv6 is little more than nominal - if they were really the same protocol, they’d have the same ethertype). Moreover, many standards render themselves incompatible in practice through excessive enthusiasm for extensibility. They may not be perma-frozen, but they probably should be.

The true, Martian way to perma-freeze a system is what I call Kelvin versioning. In Kelvin versioning, releases count down by integer degrees Kelvin. At absolute zero, the system can no longer be changed. At 1K, one more modification is possible. And so on. For instance, Nock is at 9K. It might change, though it probably won’t. Nouns themselves are at 0K - it is impossible to imagine changing anything about those three sentences.

Understood in this way, kelvin versioning is very simple. One simply counts downwards, and at absolute zero (i.e. 0K) no other releases are legal. It is no more than a versioning scheme designed for abstract components that should eventually freeze.

Many years later, the Urbit blog described kelvin versioning once more in the post Towards a Frozen Operating System. This presented a significant refinement of the original scheme, introducing both recursive and so-called “telescoping” mechanics to it:

The right way for this trunk to approach absolute zero is to “telescope” its Kelvin versions. The rules of telescoping are simple:

If tool B sits on platform A, either both A and B must be at absolute zero, or B must be warmer than A.

Whenever the temperature of A (the platform) declines, the temperature of B (the tool) must also decline.

B must state the version of A it was developed against. A, when loading B, must state its own current version, and the warmest version of itself with which it’s backward-compatible.

Of course, if B itself is a platform on which some higher-level tool C depends, it must follow the same constraints recursively.

This is more or less a complete characterisation of kelvin versioning, but it’s still not quite precise enough. If one looks at other versioning schemes that try to communicate some specific semantic content (the most obvious example being semver), it’s obvious that they take great pains to be formal and precise about their mechanics.

Experience has demonstrated to me that such formality is necessary. Even the excerpt above has proven to be ambiguous or underspecified re: the details of various situations or corner cases that one might run into. These confusions can be resolved by a rigorous protocol specification, which, in this case isn’t very difficult to put together.

Kelvin versioning and its use in Urbit is the subject of the currently-evolving UP9, recent proposed updates to which have not yet been ratified. The following is my own personal take on and simple formal specification of kelvin versioning – I believe it resolves any major ambiguities that the original descriptions may have introduced.

Kelvin Versioning (Specification)

(The key words “MUST”, “MUST NOT”, “REQUIRED”, “SHALL”, “SHALL NOT”, “SHOULD”, “SHOULD NOT”, “RECOMMENDED”, “MAY”, and “OPTIONAL” in this document are to be interpreted as described in RFC 2119.)

For any component A following kelvin versioning,

A’s version SHALL be a nonnegative integer.
A, at any specific version, MUST NOT be modified after release.
At version 0, new versions of A MUST NOT be released.
New releases of A MUST be assigned a new version, and this version MUST be strictly less than the previous one.
If A supports another component B that also follows kelvin versioning, then:
- Either both A and B MUST be at version 0, or B’s version MUST be strictly greater than A’s version.
- If a new version of A is released and that version supports B, then a new version of B MUST be released.

These rules apply recursively for any kelvin-versioned component C that is supported by B, and so on.

Examples

Examples are particularly useful here, so let me go through a few.

Let’s take the following four components, sitting in three layers, as a running example. Here’s our initial state:

A       10K
  B     20K
    C   21K
    D   30K

So we have A at 10K supporting B at 20K. B in turn supports both C at 21K and D at 30K.

State 1

Imagine we have some patches lying around for D and want to release a new version of it. That’s easy to do; we push out a new version of D. In this case it will have version one less than 30, i.e. 29K:

A       10K
  B     20K
    C   21K
    D   29K  <-- cools from 30K to 29K

Easy peasy. This is the most trivial example.

The only possible point of confusion here is: well, what kind of change warrants a version decrement? And the answer is: any (released) change whatsoever. Anything with an associated kelvin version is immutable after being released at that version, analogous to how things are done in any other versioning scheme.

State 2

For a second example, imagine that we now have completed a major refactoring of A and want to release a new version of that.

Since A supports B, releasing a new version of A obligates us to release a new version of B as well. And since B supports both C and D, we are obligated, recursively, to release new versions of those to boot.

The total effect of a new A release is thus the following:

A       9K   <-- cools from 10K to 9K
  B     19K  <-- cools from 20K to 19K
    C   20K  <-- cools from 21K to 20K
    D   28K  <-- cools from 29K to 28K

This demonstrates the recursive mechanic of kelvin versioning.

An interesting effect of the above mechanic, as described in Toward a Frozen Operating System is that anything that depends on (say) A, B, and C only needs to express its dependency on some version of C. Depending on C at e.g. 20K implicitly specifies a dependency on its supporting component, B, at 19K, and then A at 9K as well (since any change to A or B must also result in a change to C).

State 3

Now imagine that someone has contributed a performance enhancement to C, and we’d like to release a new version of that.

The interesting thing here is that we’re prohibited from releasing a new version of C. Recall our current state:

A       9K
  B     19K
    C   20K  <-- one degree K warmer than B
    D   28K

Releasing a new version of C would require us to cool it by at least one kelvin, resulting in the warmest possible version of 19K. But since its supporting component, B, is already at 19K, this would constitute an illegal state under kelvin versioning. A supporting component must always be strictly cooler than anything it supports, or be at absolute zero conjointly with anything it supports.

This illustrates the so-called telescoping mechanic of kelvin versioning – one is to imagine one of those handheld telescopes made of segments that flatten into each other when collapsed.

State 4

But now, say that we’re finally going to release our new API for B. We release a new version of B, this one at 18K, which obligates us to in turn release new versions of C and D:

A       9K
  B     18K  <-- cools from 19K to 18K
    C   19K  <-- cools from 20K to 19K
    D   27K  <-- cools from 28K to 27K

In particular, the new version of B gives us the necessary space to release a new version of C, and, indeed, obligates us to release a new version of it. In releasing C at 19K, presumably we’d include the performance enhancement that we were prohibited from releasing in State 3.

State 5

A final example that’s simple, but useful to illustrate explicitly, involves introducing a new component, or replacing a component entirely.

For example: say that we’ve decided to deprecate C and D and replace them with a single new component, E, supported by B. This is as easy as it sounds:

A       9K
  B     18K
    E   40K  <-- initial release at 40K

We just swap in E at the desired initial kelvin version. The initial kelvin can be chosen arbitrarily; the only restriction is that it be warmer than the the component that supports it (or be at absolute zero conjointly with it).

It’s important to remember that, in this component-resolution of kelvin versioning, there is no notion of the “total temperature” of the stack. Some third party could write another component, F, supported by E, with initial version at 1000K, for example. It doesn’t introduce any extra burden or responsibility on the maintainers of components A through E.

Collective Kelvin Versioning

So – all that is well and good for what I’ll call the component-level mechanics of kelvin versioning. But it’s useful to touch on a related construct, that of collectively versioning a stack of kelvin-versioned components. This minor innovation on Curtis’s original idea was put together by myself and my colleague Philip Monk.

If you have a collection of kelvin-versioned things, e.g. the things in our initial state from the prior examples:

A       10K
  B     20K
    C   21K
    D   30K

then you may want to release all these things, together, as some abstract thing. Notably, this happens in the case of the Urbit kernel, where the stack consists of a functional VM, an unapologetically amathematical purely functional programming language, special-purpose kernel modules, etc. It’s useful to be able to describe the whole kernel with a single version number.

To do this in a consistent way, you can select one component in your stack to serve as a primary index of sorts, and then capture everything it supports via a patch-like, monotonically decreasing “fractional temperature” suffix.

This is best illustrated via example. If we choose B as our primary index in the initial state above, for example, we could version the stack collectively as 20.9K. B provides the 20K, and everything it supports is just lumped into the “patch version” 9.

If we then consider the example given in State 1, i.e.:

A       10K
  B     20K
    C   21K
    D   29K

in which D has cooled by a degree kelvin, then we can version this stack collectively as 20.8K. If we were to then release a new version of C at 20K, then we could release the stack collectively as 20.7K. And so on.

There is no strictly prescribed schedule as to how to decrease the fractional temperature, but the following schedule is recommended:

.9, .8, .7, .., .1, .01, .001, .0001, ..

Similarly, the fractional temperature should reset to .9 whenever the primary index cools. If we consider the State 2, for example, where a new release of A led to every other component in the stack cooling, we had this:

Note that B has cooled by a kelvin, so we would version this stack collectively as 19.9K. The primary index has decreased by a kelvin, and the fractional temperature has been reset to .9.

While I think examples illustrate this collective scheme most clearly, after my schpeel about the pitfalls of ambiguity it would be remiss of me not to include a more formal spec:

Collective Kelvin Versioning (Specification)

For a collection of kelvin-versioned components K:

K’s version SHALL be characterised by a primary index, chosen from a component in K, and and a real number in the interval [0, 1) (the “fractional temperature”), determined by all components that the primary index component supports.

The fractional temperature MAY be 0 only if the primary index’s version is 0.
K, at any particular version, MUST NOT be modified after release.
At primary index version 0 and fractional temperature 0, new versions of K MUST NOT be released.
New releases of K MUST be assigned a new version, and this version MUST be strictly less than the previous one.
When a new release of K includes new versions of any component supported by the primary index, but not a new version of the primary index proper, its fractional temperature MUST be less than the previous version.

Given constant primary index versions, fractional temperatures corresponding to new releases SHOULD decrease according to the following schedule:
```
.9, .8, .7, .., .1, .01, .001, .0001, ..
```
When a new release of K includes a new version of the primary index, the fractional temperature of SHOULD be reset to 9.
New versions of K MAY be indexed by components other than the primary index (i.e., K may be “reindexed” at any point). However, the new chosen component MUST either be colder than the primary index it replaces, or be at version 0 conjointly with the primary index it replaces.

Etc.

In my experience, the major concern in adopting a kelvin versioning scheme is that one will accidentally initialise everything with a set of temperatures (i.e. versions) that are too cold (i.e. too close to 0), and thus burn through too many version numbers too quickly on the path to freezing. To alleviate this, it helps to remember that one has an infinite number of release candidates available for every component at every temperature.

The convention around release candidates is just to prepend a suffix to the next release version along the lines of .rc1, .rc2, etc. One should feel comfortable using these liberally, iterating through release candidates as necessary before finally committing to a new version at a properly cooler temperature.

The applications that might want to adopt kelvin versioning are probably pretty limited, and may indeed even be restricted to the Urbit kernel itself (Urbit has been described by some as “that operating system with kernel that eventually reaches absolute zero under kelvin versioning”). Nonetheless: I believe this scheme to certainly be more than a mere marketing gimmick or what have you, and, at minimum, it makes for an interesting change of pace from semver.

Email for l33t h4x0rz

2020-02-11T00:00:00+04:00

(UPDATE 2024/09/08: while hosting your own mailserver is not covered in this post, I recommend you check out Simple NixOS Mailserver for a borderline trivial way to do it.)

A couple of people recently asked about my email setup, so I figured it might be best to simply document some of it here.

I run my own mail server for jtobin.io, plus another domain or two, and usually wind up interacting with gmail for work. I use offlineimap to fetch and sync mail with these remote servers, msmtp and msmtpq to send mail, mutt as my MUA, notmuch for search, and tarsnap for backups.

There are other details; vim for writing emails, urlview for dealing with URLs, w3m for viewing HTML, pass for password storage, etc. etc. But the mail setup proper is as above.

I’ll just spell out some of the major config below, and will focus on the configuration that works with gmail, since that’s probably of broader appeal. You can get all the software for it via the following in nixpkgs:

mutt offlineimap msmtp notmuch notmuch-mutt

offlineimap

offlineimap is used to sync local and remote email; I use it to manually grab emails occasionally throughout the day. You could of course set it up to run automatically as a cron job or what have you, but I like deliberately fetching my email only when I actually want to deal with it.

Here’s a tweaked version of one of my .offlineimaprc files:

[general]
accounts = work

[Account work]
localrepository = work-local
remoterepository = work-remote
postsynchook = notmuch new

[Repository work-local]
type = Maildir
localfolders = ~/mail/work
sep = /
restoreatime = no

[Repository work-remote]
type = Gmail
remoteuser = FIXME_user@domain.tld
remotepass = FIXME_google_app_password
realdelete = no
ssl = yes
sslcacertfile = /usr/local/etc/openssl/cert.pem
folderfilter = lambda folder: folder not in\
  ['[Gmail]/All Mail', '[Gmail]/Important', '[Gmail]/Starred']

You should be able to figure out the gist of this. Pay particular attention to the ‘remoteuser’, ‘remotepass’, and ‘folderfilter’ options. For ‘remotepass’ in particular you’ll want to generate an app-specific password from Google. The ‘folderfilter’ option lets you specify the gmail folders that you actually want to sync; folder in [..] and folder not in [..] are probably all you’ll want here.

If you don’t want to store your password in cleartext, and instead want to grab it from an encrypted store, you can use the ‘remotepasseval’ option. I don’t bother with this for Google accounts that have app-specific passwords, but do for others.

This involves a little bit of extra setup. First, you can make some Python functions available to the config file with ‘pythonfile’:

[general]
accounts = work
pythonfile = ~/.offlineimap.py

Here’s a version of that file that I keep, which grabs the desired from pass(1):

#! /usr/bin/env python2
from subprocess import check_output

def get_pass():
    return check_output("pass FIXME_PASSWORD", shell=True).strip("\n")

Then you can just call the get_pass function in ‘remotepasseval’ back in .offlineimaprc:

[Repository work-remote]
type = Gmail
remoteuser = FIXME_user@domain.tld
remotepasseval = get_pass()
realdelete = no
ssl = yes
sslcacertfile = /usr/local/etc/openssl/cert.pem
folderfilter = lambda folder: folder not in\
  ['[Gmail]/All Mail', '[Gmail]/Important', '[Gmail]/Starred']

When you’ve got this set up, you should just be able to run offlineimap to fetch your email. If you maintain multiple configuration files, it’s helpful to specify a specific one using -c, e.g. offlineimap -c .offlineimaprc-foo.

msmtp, msmtpq

msmtp is used to send emails. It’s a very simple SMTP client. Here’s a version of my .msmtprc:

defaults
auth           on
tls            on
tls_starttls   on
tls_trust_file /usr/local/etc/openssl/cert.pem
logfile        ~/.msmtp.log

account work
host           smtp.gmail.com
port           587
from           FIXME_user@domain.tld
user           FIXME_user@domain.tld
password       FIXME_google_app_password

account default: work

Again, very simple.

You can do a similar thing here if you don’t want to store passwords in cleartext. Just use ‘passwordeval’ and the desired shell command directly, e.g.:

account work
host           smtp.gmail.com
port           587
from           FIXME_user@domain.tld
user           FIXME_user@domain.tld
passwordeval   "pass FIXME_PASSWORD"

I occasionally like to work offline, so I use msmtpq to queue up emails to send later. Normally you don’t have to deal with any of this directly, but occasionally it’s nice to be able to check the queue. You can do that with msmtp-queue -d:

$ msmtp-queue -d

  no mail in queue

If there is something stuck in the queue, you can force it to send with msmtp-queue -r or -R. FWIW, this has happened to me while interacting with gmail under a VPN in the past.

mutt

Mutt is a fantastic MUA. Its tagline is “all mail clients suck, this one just sucks less,” but I really love mutt. It may come as a surprise that working with email can be a pleasure, especially if you’re accustomed to working with clunky webmail UIs, but mutt makes it so.

Here’s a pruned-down version of one of my .muttrc files:

set realname      = "MyReal Name"
set from          = "user@domain.tld"
set use_from      = yes
set envelope_from = yes

set mbox_type     = Maildir
set sendmail      = "~/.nix-profile/bin/msmtpq -a work"
set sendmail_wait = -1
set folder        = "~/mail/work"
set spoolfile     = "+INBOX"
set record        = "+[Gmail]/Sent Mail"
set postponed     = "+[Gmail]/Drafts"

set smtp_pass = "FIXME_google_app_password"
set imap_pass = "FIXME_google_app_password"

set signature = "~/.mutt/.signature-work"

set editor = "vim"

set sort     = threads
set sort_aux = reverse-last-date-received

set pgp_default_key          = "my_default_pgp@key"

set crypt_use_gpgme          = yes
set crypt_autosign           = yes
set crypt_replysign          = yes
set crypt_replyencrypt       = yes
set crypt_replysignencrypted = yes

bind index gg first-entry
bind index G last-entry
bind index B imap-fetch-mail

bind index - collapse-thread
bind index _ collapse-all

set alias_file = ~/.mutt/aliases
set sort_alias = alias
set reverse_alias = yes
source $alias_file

auto_view text/html
alternative_order text/plain text/enriched text/html

subscribe my_favourite@mailing.list

macro index <F8> \
"<enter-command>set my_old_pipe_decode=\$pipe_decode my_old_wait_key=\$wait_key nopipe_decode nowait_key<enter>\
<shell-escape>notmuch-mutt -r --prompt search<enter>\
<change-folder-readonly>`echo ${XDG_CACHE_HOME:-$HOME/.cache}/notmuch/mutt/results`<enter>\
<enter-command>set pipe_decode=\$my_old_pipe_decode wait_key=\$my_old_wait_key<enter>i" \
"notmuch: search mail"

macro index <F9> \
"<enter-command>set my_old_pipe_decode=\$pipe_decode my_old_wait_key=\$wait_key nopipe_decode nowait_key<enter>\
<pipe-message>notmuch-mutt -r thread<enter>\
<change-folder-readonly>`echo ${XDG_CACHE_HOME:-$HOME/.cache}/notmuch/mutt/results`<enter>\
<enter-command>set pipe_decode=\$my_old_pipe_decode wait_key=\$my_old_wait_key<enter>" \
"notmuch: reconstruct thread"

macro index l "<enter-command>unset wait_key<enter><shell-escape>read -p 'notmuch query: ' x; echo \$x >~/.cache/mutt_terms<enter><limit>~i \"\`notmuch search --output=messages \$(cat ~/.cache/mutt_terms) | head -n 600 | perl -le '@a=<>;chomp@a;s/\^id:// for@a;$,=\"|\";print@a'\`\"<enter>" "show only messages matching a notmuch pattern"

# patch rendering
# https://karelzak.blogspot.com/2010/02/highlighted-patches-inside-mutt.html
color body green default "^diff \-.*"
color body green default "^index [a-f0-9].*"
color body green default "^\-\-\- .*"
color body green default "^[\+]{3} .*"
color body cyan default "^[\+][^\+]+.*"
color body red  default "^\-[^\-]+.*"
color body brightblue default "^@@ .*"

# vim: ft=muttrc:

Some comments on all that:

set mbox_type     = Maildir
set sendmail      = "~/.nix-profile/bin/msmtpq -a work"
set sendmail_wait = -1
set folder        = "~/mail/work"
set spoolfile     = "+INBOX"
set record        = "+[Gmail]/Sent Mail"
set postponed     = "+[Gmail]/Drafts"

Note here that we’re specifying msmtpq as our sendmail program. The -a work command here refers to the account defined in your .msmtprc file, so if you change the name of it there, you have to do it here as well. Ditto for the folder.

(If you’re tweaking these config files for your own use, I’d recommend just substituting all instances of ‘work’ with your own preferred account name.)

The negative ‘sendmail_wait’ value handles queueing mails up appropriately when offline, IIRC.

set smtp_pass = "FIXME_google_app_password"
set imap_pass = "FIXME_google_app_password"

Here are the usual cleartext app passwords. If you want to store them encrypted, there’s a usual method for doing that: add the following to the top of your .muttrc:

source "gpg -d ~/.mutt/my-passwords.gpg |"

where .mutt/my-passwords.gpg should be the above smtp_pass and imap_pass assignments, encrypted with your desired private key.

Continuing with the file at hand:

set signature = "~/.mutt/.signature-work"

set editor = "vim"

These should be self-explanatory. The signature file should just contain the signature you want appended to your mails (it will be appended under a pair of dashes). And if you want to use some other editor to compose your emails, just specify it here.

set pgp_default_key          = "my_default_pgp@key"

set crypt_use_gpgme          = yes
set crypt_autosign           = yes
set crypt_replysign          = yes
set crypt_replyencrypt       = yes
set crypt_replysignencrypted = yes

Mutt is one of the few programs that has great built-in support for PGP. It can easily encrypt, decrypt, and sign messages, grab public keys, etc. Here you can see that I’ve set it to autosign messages, reply to encrypted messages with encrypted messages, and so on.

bind index gg first-entry
bind index G last-entry
bind index B imap-fetch-mail

bind index - collapse-thread
bind index _ collapse-all

These are a few key bindings that I find helpful. The first bunch are familiar to vim users and are useful for navigating around; the second two are really useful for compressing or expanding the view of your mailbox.

set alias_file = ~/.mutt/aliases
set sort_alias = alias
set reverse_alias = yes
source $alias_file

auto_view text/html
alternative_order text/plain text/enriched text/html

The alias file lets you define common shortcuts for single or multiple addresses. I get a lot of use out of multiple address aliases, e.g.:

alias chums socrates@ago.ra, plato@acade.my, aristotle@lyce.um

The MIME type stuff below the alias config is just a sane set of defaults for viewing common mail formats.

subscribe my_favourite@mailing.list

Mutt makes interacting with mailing lists very easy just by default, but you can also indicate addresses that you’re subscribed to, as above, to unlock a few extra features for them (‘list reply’ being a central one). To tell mutt that you’re subscribed to haskell-cafe, for example, you’d use:

subscribe haskell-cafe@haskell.org

The three longer macros that follow are for notmuch. I really only find myself using the last one, ‘l’, for search. FWIW, notmuch’s search functionality is fantastic; I’ve found it to be more useful than gmail’s web UI search, I think.

The patch rendering stuff at the end is just a collection of heuristics for rendering in-body .patch files well. This is useful if you subscribe to a patch-heavy mailing list, e.g. LKML or git@vger.kernel.org, or if you just want to be able to better-communicate about diffs in your day-to-day emails with your buddies.

Fin

There are obviously endless ways you can configure all this stuff, especially mutt, and common usage patterns that you’ll quickly find yourself falling into. But whatever you find those to be, the above should at least get you up and running pretty quickly with 80% of the desired feature set.

Basic Hoonery

2019-02-04T00:00:00+04:00

In my last post I first introduced hnock, a little interpreter for Nock, and then demonstrated it on a hand-rolled decrement function. In this post I’ll look at how one can handle the same (contrived, but illustrative) task in Hoon.

Hoon is the higher- or application-level programming language for working with Arvo, the operating system of Urbit. The best way I can describe it is something like “Haskell meets C meets J meets the environment is always explicit.”

As a typed, functional language, Hoon feels surprisingly low-level. One is never allocating or deallocating memory explicitly when programming in Hoon, but the experience somehow feels similar to working in C. The idea is that the language should be simple and straightforward and support a fairly limited level of abstraction. There are the usual low-level functional idioms (map, reduce, etc.), as well as a structural type system to keep the programmer honest, but at its core, Hoon is something of a functional Go (a language which, I happen to think, is not good).

It’s not a complex language, like Scala or Rust, nor a language that overtly supports sky-high abstraction, like Haskell or Idris. Hoon is supposed to exist at a sweet spot for getting work done. And I am at least willing to buy the argument that it is pretty good for getting work done in Urbit.

Recall our naïve decrement function in Haskell. It looked like this:

dec :: Integer -> Integer
dec m =
  let loop n
        | succ n == m = n
        | otherwise   = loop (succ n)
  in  loop 0

Let’s look at a number of ways to write this in Hoon, showing off some of the most important Hoon programming concepts in the process.

Cores

Here’s a Hoon version of decrement. Note that to the uninitiated, Hoon looks gnarly:

|=  m=@
=/  n=@  0
=/  loop
|%
  ++  recur
    ?:  =(+(n) m)
      n
    recur(n +(n))
  --
recur:loop

We can read it as follows:

Define a function that takes an argument, ‘m’, having type atom (recall that an atom is an unsigned integer).
Define a local variable called ‘n’, having type atom and value 0, and add it to the environment (or, if you recall our Nock terminology, to the subject).
Define a local variable called ‘loop’, with precise definition to follow, and add it to the environment.
‘loop’ is a core, i.e. more or less a named collection of functions. Define one such function (or arm), ‘recur’, that checks to see if the increment of ‘n’ is equal to ‘m’, returning ‘n’ if so, and calling itself, except with the value of ‘n’ in the environment changed to ‘n + 1’, if not.
Evaluate ‘recur’ as defined in ‘loop’.

(To test this, you can enter the Hoon line-by-line into the Arvo dojo. Just preface it with something like =core-dec to give it a name, and call it via e.g. (core-dec 20).)

Hoon may appear to be a write-only language, though I’ve found this to not necessarily be the case (just to note, at present I’ve read more Hoon code than I’ve written). Good Hoon has a terse and very vertical style. The principle that keeps it readable is that, roughly, each line should contain one important logical operation. These operations are denoted by runes, the =/ and ?: and similar ASCII digraphs sprinkled along the left hand columns of the above example. This makes it look similar to e.g. J – a language I have long loved, but never mastered – although in J the rough ‘one operator per line’ convention is not typically in play.

In addition to the standard digraph runes, there is also a healthy dose of ‘irregular’ syntax in most Hoon code for simple operations that one uses frequently. Examples used above include =(a b) for equality testing, +(n) for incrementing an atom, and foo(a b) for evaluating ‘foo’ with the value of ‘a’ in the environment changed to ‘b’. Each of these could be replaced with a more standard rune-based expression, though for such operations the extra verbosity is not usually warranted.

Cores like ‘loop’ seem, to me, to be the mainstay workhorse of Hoon programming. A core is more or less a structure, or object, or dictionary, or whatever, of functions. One defines them liberally, constructs a subject (i.e. environment) to suit, and then evaluates them, or some part of them, against the subject.

To be more precise, a core is a Nock expression; like every non-atomic value in Nock, it is a tree. Starting from the cell [l r], the left subtree, ‘l’, is a tree of Nock formulas (i.e. the functions, like ‘recur’, defined in the core). The right subtree, ‘r’ is all the data required to evaluate those Nock formulas. The traditional name for the left subtree, ‘l’, is the battery of the core; the traditional name for the right subtree is the payload.

One is always building up a local environment in Hoon and then evaluating some value against it. Aside from the arm ‘recur’, the core ‘loop’ also contains in its payload the values ‘m’ and ‘n’. The expression ‘recur:loop’ – irregular syntax for =< recur loop – means “use ‘loop’ as the environment and evaluate ‘recur’.” Et voilà, that’s how we get our decrement.

You’ll note that this should feel very similar to the way we defined decrement in Nock. Our hand-assembled Nock code, slightly cleaned up, looked like this:

[8
  [1 0]
  8
    [1
      6
        [5 [4 0 6] [0 7]]
        [0 6]
        2 [[0 2] [4 0 6] [0 7]] [0 2]
    ]
    2 [0 1] [0 2]
]

This formula, when evaluated against an atom subject, creates another subject from it, defining a ‘loop’ analogue that looks in specific addresses in the subject for itself, as well as the ‘m’ and ‘n’ variables, such that it produces the decrement of the original subject. Our Hoon code does much the same – every ‘top-level’ rune expression adds something to the subject, until we get to the final expression, ‘recur:loop’, which evaluates ‘recur’ against the subject, ‘loop’.

The advantage of Hoon, in comparison to Nock, is that we can work with names, instead of raw tree addresses, as well as with higher-level abstractions like cores. The difference between Hoon and Nock really is like the difference between C and assembly!

For what it’s worth, here is the compiled Nock corresponding to our above decrement function:

[8
  [1 0]
  8
    [8
      [1
        6
          [5 [4 0 6] 0 30]
          [0 6]
          9 2 10 [6 4 0 6] 0 1
      ]
      0 1
    ]
    7 [0 2] 9 2 0 1
]

It’s similar, though not identical, to our hand-rolled Nock. In particular, you can see that it is adding a constant conditional formula, including the familiar equality check, to the subject (note that the equality check, using Nock-5, refers to address 30 instead of 7 – presumably this is because I have more junk floating around in my dojo subject). Additionally, the formulas using Nock-9 and Nock-10 reduce to Nock-2 and Nock-0, just like our hand-rolled code does.

But our Hoon is doing more than the bespoke Nock version did, so we’re not getting quite the same code. Worth noting is the ‘extra’ use of Nock-8, which is presumably required because I’ve defined both ‘recur’, the looping function, and ‘loop’, the core to hold it, and the hand-rolled Nock obviously didn’t involve a core.

Doors

Here’s another way to write decrement, using another fundamental Hoon construct, the door:

|=  m=@
=/  loop
|_  n=@
  ++  recur
    ?:  =(+(n) m)
      n
    ~(recur ..recur +(n))
  --
~(recur loop 0)

A door is a core that takes an argument. Here we’ve used the |_ rune, instead than |%, to define ‘loop’, and note that it takes ‘n’ as an argument. So instead of ‘n’ being defined external to the core, as it was in the previous example, here we have to specify it explicitly when we call ‘recur’. Note that this is more similar to our Haskell example, in which ‘loop’ was defined as a function taking ‘n’ as an argument.

The two other novel things here are the ~(recur ..recur +(n)) and ~(recur loop 0) expressions, which actually turn out to be mostly the same thing. The syntax:

~(arm door argument)

is irregular, and means “evaluate ‘arm’ in ‘door’ using ‘argument’”. So in the last line, ~(recur loop 0) means “evaluate ‘recur’ in ‘loop’ with n set to 0.” In the definition of ‘recur’, on the other hand, we need to refer to the door that contains it, but are in the very process of defining that thing. The ‘..recur’ syntax means “the door that contains ‘recur’,” and is useful for exactly this task, given we can’t yet refer to ‘loop’. The syntax ~(recur ..recur +(n)) means “evaluate ‘recur’ in its parent door with n set to n + 1.”

Let’s check the compiled Nock of this version:

[8
  [8
    [1 0]
    [1
      6
        [5 [4 0 6] 0 30]
        [0 6]
        8
          [0 1]
          9 2 10 [6 7 [0 3] 4 0 6] 0 2
    ]
    0 1
  ]
  8
    [0 2]
    9 2 10 [6 7 [0 3] 1 0] 0 2
]

There’s even more going on here than in our core-implemented decrement, but doors are a generalisation of cores, so that’s to be expected.

Hoon has special support, though, for one-armed doors. This is precisely how functions (also called gates or traps, depending on the context) are implemented in Hoon. The following is probably the most idiomatic version of naïve decrement:

|=  m=@
=/  n  0
|-
?:  =(+(n) m)
  n
$(n +(n))

The |= rune that we’ve been using throughout these examples really defines a door, taking the specified argument, with a single arm called ‘$’. The |- rune here does the same, except it immediately calls the ‘$’ arm after defining it. The last line, $(n +(n)), is analogous to the recur(n +(n)) line in our first example: it evaluates the ‘$’ arm, except changing the value of ‘n’ to ‘n + 1’ in the environment.

(Note that there are two ‘$’ arms defined in the above code – one via the use of |=, and one via the use of |-. But there is no confusion as to which one we mean, since the latter has been the latest to be added to the subject. Additions to the subject are always prepended in Hoon – i.e. they are placed at address 2. As the topmost ‘$’ in the subject is the one that corresponds to |-, it is resolved first.)

The compiled Nock for this version looks like the following:

[8
  [1 0]
  8
    [1
      6
        [5 [4 0 6] 0 30]
        [0 6]
        9 2 10 [6 4 0 6] 0 1
    ]
    9 2 0 1
]

And it is possible (see the appendix) to show that, modulo some different addressing, this reduces exactly to our hand-rolled Nock code.

UPDATE: my colleague Ted Blackman, an actual Hoon programmer, recommended the following as a slightly more idiomatic version of naïve decrement:

=|  n=@
|=  m=@
^-  @
?:  =(+(n) m)
  n
$(n +(n))

Note that here we’re declaring ‘n’ outside of the gate itself by using another rune, =|, that gives the variable a default value based on its type (an atom’s default value is 0). There’s also an explicit type cast via ^- @, indicating that the gate produces an atom (like type signatures in Haskell, it is considered good practice to include these, even though they may not strictly be required).

Declaring ‘n’ outside the gate is interesting. It has an imperative feel, as if one were writing the code in Python, or were using a monad like ‘State’ or a ‘PrimMonad’ in Haskell. Like in the Haskell case, we aren’t actually doing any mutation here, of course – we’re creating new subjects to evaluate each iteration of our Nock formula against. And the resulting Nock is very succinct:

[6
  [5 [4 0 14] 0 6]
  [0 14]
  9 2 10 [14 4 0 14] 0 1
]

Basic Generators

If you tested the above examples, I instructed you to do so by typing them into Arvo’s dojo. I’ve come to believe that, in general, this is a poor way to teach Hoon. It shouldn’t be done for all but the most introductory examples (such as the ones I’ve provided here).

If you’ve learned Haskell, you are familiar with the REPL provided by GHCi, the Glasgow Haskell Compiler’s interpreter. Code running in GHCi is implicitly running in the IO monad, and I think this leads to confusion amongst newcomers who must then mentally separate “Haskell in GHC” from “Haskell in GHCi.”

I think there is a similar problem in Hoon. Expressions entered into the dojo implicitly grow or shrink or otherwise manipulate the dojo’s subject, which is not, in general, available to standalone Hoon programs. Such standalone Hoon programs are called generators. In general, they’re what you will use when working in Hoon and Arvo.

There are four kinds of generators: naked, %say, %ask, and %get. In this post we’ll just look at the first two; the last couple are out of scope, for now.

Naked Generators

The simplest kind of generator is the ‘naked’ generator, which just exists in a file somewhere in your Urbit’s “desk.” If you save the following as naive-decrement.hoon in an Urbit’s home/gen directory, for example:

|=  m=@
=/  n  0
|-
?:  =(+(n) m)
  n
$(n +(n))

Then you’ll be able to run it in a dojo via:

~zod:dojo> +naive-decrement 20
19

A naked generator can only be a simple function (technically, a gate) that produces a noun. It has no access to any external environment – it’s basically just a self-contained function in a file. It must have an argument, and it must have only one argument; to pass multiple values to a naked generator, one must use a cell.

Say Generators

Hoon is a purely functional language, but, unlike Haskell, it also has no IO monad to demarcate I/O effects. Hoon programs do not produce effects on their own at all – instead, they construct nouns that tell Arvo how to produce some effect or other.

A %say generator (where %say is a symbol) produces a noun, but it can also make use of provided environment data (e.g. date information, entropy, etc.). The idea is that the generator has a specific structure that Arvo knows how to handle, in order to supply it with the requisite information. Specifically, %say generators have the structure:

:-  %say
|=  [<environment data> <list of arguments> <list of optional arguments>]
:-  %noun
<code>

I’ll avoid discussing what a list is in Hoon at the moment, and we won’t actually use any environment data in any examples here. But if you dump the following in home/gen/naive-decrement.hoon, for example:

:-  %say
|=  [* [m=@ ~] ~]
:-  %noun
=/  n  0
|-
?:  =(+(n) m)
  n
$(n +(n))

you can call it from the dojo via the mechanism as before:

~zod:dojo> +naive-decrement 20
19

The generator itself actually returns a particularly-structured noun; a cell with the symbol %say as its head, and a gate returning a pair of the symbol %noun and a noun as its tail. The %noun symbol describes the data produced by the generator. But note that this is not displayed when evaluating the generator in the dojo – instead, we just get the noun itself, but this behaviour is dojo-dependent.

I think one should almost get in the habit of writing %say generators for most Hoon code, even if a simple naked generator or throwaway dojo command would do the trick. They are so important for getting things done in Hoon that it helps to learn about & start using them sooner than later.

Fin

I’ve introduced Hoon and given a brief tour of what I think are some of the most important tools for getting work done in the language. Cores, doors, and gates will get you plenty far, and early exposure to generators, in the form of the basic naked and %say variants, will help you avoid the habit of programming in the dojo, and get you writing more practically-structured Hoon code from the get-go.

I haven’t had time in this post to describe Hoon’s type system, which is another very important topic when it comes to getting work done in the language. I’ll probably write one more to create a small trilogy of sorts – stay tuned.

Appendix

Let’s demonstrate that the compiled Nock code from our door-implemented decrement reduces to the same as our hand-rolled Nock, save different address use. Recall that our compile Nock code was:

[8
  [1 0]
  8
    [1
      6
        [5 [4 0 6] 0 30]
        [0 6]
        9 2 10 [6 4 0 6] 0 1
    ]
    9 2 0 1
]

An easy reduction is from Nock-9 to Nock-2. Note that *[a 9 b c] is the same as *[*[a c] 2 [0 1] 0 b]. When ‘c’ is [0 1], we have that *[a c] = a, such that *[a 9 b [0 1]] is the same as *[a 2 [0 1] 0 b], i.e. that the formula [9 b c] is the same as the formula [2 [0 1] 0 b]. We can thus reduce the use of Nock-9 on the last line to:

[8
  [1 0]
  8
    [1
      6
        [5 [4 0 6] 0 30]
        [0 6]
        9 2 10 [6 4 0 6] 0 1
    ]
    2 [0 1] 0 2
]

The remaining formula involving Nock-9 evaluates [10 [6 4 0 6] 0 1] against the subject, and then evaluates [2 [0 1] [0 2]] against the result. Note that, for some subject ‘a’, we have:

*[a 10 [6 4 0 6] 0 1]
= #[6 *[a 4 0 6] *[a 0 1]]
= #[6 *[a 4 0 6] a]
= #[3 [*[a 4 0 6] /[7 a]] a]
= #[1 [/[2 a] [*[a 4 0 6] /[7 a]]] a]
= [/[2 a] [*[a 4 0 6] /[7 a]]]
= [*[a 0 2] [*[a 4 0 6] *[a 0 7]]]
= *[a [0 2] [4 0 6] [0 7]]

such that [10 [6 4 0 6] 0 1] = [[0 2] [4 0 6] [0 7]]. And for c = [[0 2] [4 0 6] [0 7]] and some subject ‘a’, we have:

*[a 9 2 c]
= *[*[a c] 2 [0 1] 0 2]

and for b = [2 [0 1] 0 2]:

*[*[a c] b]
= *[a 7 c b]
= *[a 7 [[0 2] [4 0 6] [0 7]] [2 [0 1] 0 2]]

such that:

[9 2 [0 2] [4 0 6] [0 7]] = [7 [[0 2] [4 0 6] [0 7]] [2 [0 1] 0 2]]

Now. Note that for any subject ‘a’ we have:

*[a 7 [[0 2] [4 0 6] [0 7]] [2 [0 1] 0 2]]
= *[a 7 [[0 2] [4 0 6] [0 7]] *[a 0 2]]

since *[a 2 [0 1] 0 2] = *[a *[a 0 2]]. Thus, we can reduce:

*[a 7 [[0 2] [4 0 6] [0 7]] *[a 0 2]]
= *[*[a [0 2] [4 0 6] [0 7]] *[a 0 2]]
= *[a 2 [[0 2] [4 0 6] [0 7]] [0 2]]

such that

[7 [[0 2] [4 0 6] [0 7]] [2 [0 1] 0 2]] = [2 [[0 2] [4 0 6] [0 7]] [0 2]]

and, so that, finally, we can reduce the compiled Nock to:

[8
  [1 0]
  8
    [1
      6
        [5 [4 0 6] 0 30]
        [0 6]
        2 [[0 2] [4 0 6] [0 7]] 0 2
    ]
    2 [0 1] 0 2
]

which, aside from the use of the dojo-assigned address 30 (and any reduction errors on this author’s part), is the same as our hand-rolled Nock.

A Nock Interpreter

2019-01-31T00:00:00+04:00

I wrote a little Nock interpreter called hnock some months ago and just yesterday updated it to support the latest version of Nock, 4K. Nock – the base layer VM of Urbit – is a very simple little “functional assembly language” of sorts. It is of particular interest in that it is capable of practical computation (indeed, it is Turing complete) but is not defined in terms of the lambda calculus. So, no variable naming and capture-avoidance and the like to deal with, which is kind of neat.

Nock (at 4K) itself has a small specification, repeated below. Famously, it fits on a t-shirt:

nock(a)             *a
[a b c]             [a [b c]]

?[a b]              0
?a                  1
+[a b]              +[a b]
+a                  1 + a
=[a a]              0
=[a b]              1

/[1 a]              a
/[2 a b]            a
/[3 a b]            b
/[(a + a) b]        /[2 /[a b]]
/[(a + a + 1) b]    /[3 /[a b]]
/a                  /a

#[1 a b]            a
#[(a + a) b c]      #[a [b /[(a + a + 1) c]] c]
#[(a + a + 1) b c]  #[a [/[(a + a) c] b] c]
#a                  #a

*[a [b c] d]        [*[a b c] *[a d]]

*[a 0 b]            /[b a]
*[a 1 b]            b
*[a 2 b c]          *[*[a b] *[a c]]
*[a 3 b]            ?*[a b]
*[a 4 b]            +*[a b]
*[a 5 b c]          =[*[a b] *[a c]]

*[a 6 b c d]        *[a *[[c d] 0 *[[2 3] 0 *[a 4 4 b]]]]
*[a 7 b c]          *[*[a b] c]
*[a 8 b c]          *[[*[a b] a] c]
*[a 9 b c]          *[*[a c] 2 [0 1] 0 b]
*[a 10 [b c] d]     #[b *[a c] *[a d]]

*[a 11 [b c] d]     *[[*[a c] *[a d]] 0 3]
*[a 11 b c]         *[a c]

*a                  *a

Perhaps you are a neophyte Haskeller and have never implemented an interpreter for a language before. Nock makes for an excellent target to practice on.

Expressions

A noun in Nock is either an atom, i.e. an unsigned integer, or a cell, i.e. an ordered pair of nouns:

data Noun =
    Atom Integer
  | Cell Noun Noun
  deriving Eq

Very simple. A Nock expression is then a noun, or some operator applied to a noun. Per the spec, the operators are denoted by ?, +, =, /, #, and *, and the fun, convenient, and authentically Urbit-y way to pronounce these are wut, lus, tis, net, hax, and tar respectively:

data Expr =
    Noun Noun
  | Wut Noun
  | Lus Noun
  | Tis Noun
  | Net Noun
  | Hax Noun
  | Tar Noun
  deriving Eq

So, equipped with the above definitions, you can write Nock expressions in Haskell. For example, the following expression:

*[57 [4 0 1]]

maps to:

Tar
  (Cell (Atom 57)
  (Cell (Atom 4)
  (Cell (Atom 0) (Atom 1))))

Note per the spec that cells associate to the right, thus something like [4 0 1] maps to [4 [0 1]], and thus Cell (Atom 4) (Cell (Atom 0) (Atom 1)).

Evaluation and Semantics

For evaluation one can more or less copy out the production rules from the Nock spec. We can use an Either type to denote a lack of defined semantics:

data Error = Error Noun
  deriving Show

type Possibly = Either Error

and then simply define the various evaluation functions appropriately. ‘wut’, for example, checks if something is a cell:

wut :: Noun -> Possibly Noun
wut noun = return $ case noun of
  Cell {} -> Atom 0
  Atom {} -> Atom 1

(note that in Nock and Urbit more generally, 0 denotes ‘true’).

‘tis’ checks if the elements of a cell are equal:

tis :: Noun -> Possibly Noun
tis noun = case noun of
  Atom {}  -> Left (Error noun)
  Cell m n -> return $
    if   m == n
    then Atom 0
    else Atom 1

And so on. There are other operators for lookups, substitution, et cetera. The most involved operator is ‘tar’, which constitutes the Nock interpreter proper. One can simply match on the cases appropriately:

tar :: Noun -> Possibly Noun
tar noun = case noun of
  Cell a (Cell (Cell b c) d) -> do
    tard0 <- tar (Cell a (Cell b c))
    tard1 <- tar (Cell a d)
    return (Cell tard0 tard1)

  Cell a (Cell (Atom 0) b) ->
    net (Cell b a)

  Cell _ (Cell (Atom 1) b) ->
    return b

  Cell a (Cell (Atom 2) (Cell b c)) -> do
    tard0 <- tar (Cell a b)
    tard1 <- tar (Cell a c)
    tar (Cell tard0 tard1)

  -- ... and so on

It is particularly useful to look at ‘tar’ to get an idea of what is going on in Nock. One is always evaluating a noun to another noun. Semantics are only defined when the noun being evaluated is a cell, and that cell is said to have the structure:

[subject formula]

The subject is essentially the environment, reified explicitly. The formula is the code to execute against the subject.

The resulting noun is called a ‘product’. So, evaluation proceeds as:

[subject formula] -> product

In any valid application of ‘tar’, i.e. for some *[x y], ‘x’ is the subject, and ‘y’ is the formula. The first case of ‘tar’, for example, is defined via:

*[a 0 b]            /[b a]

Recall that cells in Nock are right-associative, so [a 0 b] is really [a [0 b]]. In other words: a is the subject, [0 b] is the formula. The formula [0 b], read as ‘Nock-zero’, means “look up the value at address ‘b’ in the subject,” the subject here being ‘a’. So *[a 0 b] means “look up the value at address ‘b’ in ‘a’.”

Other formulas in the various ‘tar’ cases denote additional useful language features. A formula [1 b] is the constant function, [6 b c d] is a conditional, [7 b c] is function (viz. formula) composition, and so on. There is a Turing-complete amount of power packed onto that t-shirt!

In any case, once one has constructed all of these requisite evaluation functions, we can stitch them together in a single evaluator:

eval :: Expr -> Possibly Noun
eval expr = case expr of
  Noun noun -> return noun
  Wut  e    -> wut e
  Lus  e    -> lus e
  Tis  e    -> tis e
  Net  e    -> net e
  Hax  e    -> hax e
  Tar  e    -> tar e

Parsing

Writing a combinator-based parser is an easy job. You can use parsec, for example, and define your combinators as follows.

Given low-level ‘atom’ and ‘cell’ parsers, a noun is just an atom or a cell:

noun :: Monad m => P.ParsecT T.Text u m Noun
noun =
      P.try cell
  <|> atom

and an expression is either an operator, followed by an noun, or just a noun:

expr :: Monad m => P.ParsecT T.Text u m Expr
expr =
      P.try operator
  <|> fmap Noun noun

operator :: Monad m => P.ParsecT T.Text u m Expr
operator = do
  op <- P.oneOf "?+=/#*"
  case op of
    '?' -> fmap Wut noun
    '+' -> fmap Lus noun
    '=' -> fmap Tis noun
    '/' -> fmap Net noun
    '#' -> fmap Hax noun
    '*' -> fmap Tar noun
    _   -> fail "op: bad token"

Note that we don’t allow expressions like the following, taken from the Nock spec:

*[a *[[c d] 0 *[[2 3] 0 *[a 4 4 b]]]]

These are used to define the production rules, but are not valid Nock expressions per se.

The end parser is simply:

parse :: T.Text -> Either P.ParseError Expr
parse = P.runParser expr [] "input"

Example

The final interpreter is pretty simple. Here we distinguish between cases in order to report either parsing or evaluation errors in basic fashion, respectively:

hnock :: T.Text -> Noun
hnock input = case parse input of
  Left perr -> error (show perr)
  Right ex  -> case eval ex of
    Left err -> error (show err)
    Right e  -> e

It’s easy to test. Take the simple Nock-zero example given previously, where *[a 0 b] means “look up the value at address ‘b’ in ‘a’”. Nock’s addressing scheme is simple; address ‘2’ means the leftmost component of a cell, and ‘3’ means the rightmost component. For every additional cell you want to recurse into, you just multiply by two and add one as is necessary. Address ‘1’ is just the cell itself.

So, let’s test:

~ hnock "*[[[1 2] [3 4]] [0 1]"
[[1 2] [3 4]]

Again, the formula [0 1] means look up at address one, which is the cell itself.

~ hnock "*[[[1 2] [3 4]] [0 2]]"
[1 2]

The formula [0 2] means look up at address 2, which is the leftmost component of the cell.

~ hnock "*[[[1 2] [3 4]] [0 7]]"
4

To understand [0 7], we first look up the rightmost component [3 4] at address 3, then multiply by two and add one for the rightmost component of that.

The most famous Nock learner’s problem is to implement a function for decrementing an atom (if you could do this in ~2010 you could receive an Urbit galaxy for it, which is now probably worth a serious chunk of change). Nock’s only arithmetic operator is increment, so the trick is to decrement by actually counting up from zero.

(N.b. in practice this is not how one actually evaluates decrement in Nock. The compiler is ‘jetted’, such that alternate implementations, e.g. the CPU’s decrement instruction directly, can be used to evaluate semantically-identical expressions.)

But the idea is simple enough. To warm up, we can implement this kind of decrement in Haskell. Don’t handle the underflow condition, just to keep things simple:

dec :: Integer -> Integer
dec m =
  let loop n
        | succ n == m = n
        | otherwise   = loop (succ n)
  in  loop 0

Look at the language features we’re using to accomplish this. There’s:

Variable naming / declaration, for the ‘m’, ‘loop’, and ‘n’ arguments
Integer increment, via the built-in ‘succ’
Equality testing, via ‘==’
Conditional expressions, via the ‘|’ guards
A function call, in ‘loop 0’

Let’s try and duplicate this in Nock, step by step. We’ll want analogues for all of our declared variables, to start.

A Nock formula that simply returns the subject it’s evaluated against is easy: [0 1]. We used it above.

~ hnock "*[10 [0 1]]"
10

That’s the equivalent of our ‘m’ declaration in the Haskell code.

Similarly, we can just ignore the subject and return a ‘0’ via [1 0].

~ hnock "*[10 [1 0]]"
0

That’s the initial value, 0, that appears in ‘loop 0’.

We can use Nock-8 to pair these together. It evaluates a formula against the subject, and then uses the original subject, augmented with the product, to evaluate another one:

~ hnock "*[10 [8 [1 0] [0 1]]]"
[0 10]

The last thing we need define is the ‘loop’ function itself. We’re going to place it at the head of the subject, i.e. at address two. Our ‘n’ and ‘m’ variables will then be at addresses six and seven, respectively. When we define our ‘loop’ analogue, we just need to refer to those addresses, instead of our variable names.

(N.b. Nock’s addressing scheme seems more or less equivalent to a tree-flavoured variant of de Bruijn indices.)

The ‘loop’ function in the Haskell code tests to see the successor of ‘n’ is equal to ‘m’, and if so, returns ‘n’. Otherwise it loops on the successor of ‘n’.

Nock-4 computes the successor. Remember that ‘n’ is at address six, so our analogue to ‘succ n’ is [4 0 6]. We want to check if it’s equal to ‘m’, at address seven, gotten via [0 7]. The equality check is then handled by Nock-5:

[5 [4 0 6] [0 7]]

Nock-6 is the conditional operator. If the equality check is true, we just want to return ‘n’ via [0 6]. The whole formula will look like this:

[6 [5 [4 0 6] [0 7]] [0 6] <loop on successor of n>]

To loop, we’ll use Nock-2. We’ll compute a new subject – consisting of our ‘loop’ analogue at address two, an updated ‘n’ at address six, and our original ‘m’ at address seven – and then evaluate the same looping formula against it. In Nock, the formula is:

[2 [[0 2] [4 0 6] [0 7]] [0 2]]

And our ‘loop’ analogue in full is thus:

[6 [5 [4 0 6] [0 7]] [0 6] [2 [[0 2] [4 0 6] [0 7]] [0 2]]]

Finally, we now need to actually store this in the subject. Again, we’ll ignore the subject with Nock-1, and use Nock-8 to tuple the formula up with the ‘n’ and ‘m’ analogues:

(N.b. at this point we’re pretty much at line noise – spacing it out doesn’t make it that much more readable, but it does avoid word wrap.)

[8
  [1 0]
  [8
    [1
      [6
        [5 [4 0 6] [0 7]]
        [0 6]
        [2 [[0 2] [4 0 6] [0 7]] [0 2]]
      ]
    ]
    [2 [0 1] [0 2]]
  ]
]

Let’s check if it works. We’re embedded in Haskell, after all, so let’s get some support from our host, for convenience:

~ let dec = \subj -> "*[" <> subj <> " [8 [1 0] [8 [1 [6 [5 [4 0 6] [0 7]] [0 6] [2 [[0 2] [4 0 6] [0 7]] [0 2]]]] [2 [0 1] [0 2]]]]]"
~ hnock (dec "10")
9
~ hnock (dec "25")
24
~ hnock (dec "200")
199

Phew!

Fin

Nock makes for an interesting “functional assembly language” of sorts. Its spec is so simple that writing a naïve interpreter is more or less a process of copying out the production rules. You can do it without a lot of thought.

The time-honoured decrement-writing exercise, on the other hand, really makes one think about Nock. That the environment is always explicit, you can access it via a nice addressing scheme, that you can refer to functions via pointers (viz. addresses), etc. It’s pretty intuitive, when you dig into the mechanics.

In any case – writing a decrement (or similar function) and writing your own interpreter to test it is a fun exercise. From a practical perspective, it also gives one a lot of intuition about programming in Hoon – it’s probably not too far from the truth to compare the relationship between Nock and Hoon to that between, say, x86 and C. If one is programming in Hoon, it can’t hurt to to know your Nock!

Crushing ISAAC

2018-10-07T00:00:00+04:00

(UPDATE 2020/06/30: the good people at tweag.io have since published a Nix shell environment that appears to make testing arbitrary PRNGs much less of a pain. I recommend you check it out!)

I recently needed a good cryptographically-secure and seedable pseudorandom number generator for Javascript. This didn’t turn out to be as trivial a procedure as I figured it’d be: most Javascript CSPRNGs I found didn’t appear to be manually seedable, instead automatically seeding themselves behind-the-scenes using a trusted high-quality entropy source like /dev/urandom (as per the WebCryptoAPI spec).

But! I want to use my own entropy source. So I could either implement my own cryptographically-secure PRNG, which I would obviously need to test rigorously, or I could find an existing implementation that had already been vetted widely by use. I settled on the ISAAC PRNG/stream cipher, both because it seemed like a reasonable choice of PRNG (it is used in things like coreutils’ shred, and there are no known practical attacks against it), and also because there was a Javascript implementation floating around on the internet. But the interesting thing was that the implementation did not really seem to be vetted widely – it hadn’t been updated in five years or so, and both the Github repository and the npm package seem to show very little activity around it in general.

(Update: Stefano, the author of the fork I used, later emailed me to point out that the original version of this ISAAC code is located here. His fork, on the other hand, is node-friendly.)

I was going to be using this thing in an application that made some nontrivial demands in terms of security. So while the PRNG itself seemed to fit the bill, I wanted to be assured that the implementation satisfied at least some basic criteria for pseudorandomness. I assumed that it probably works – I just wanted to have some reasonable confidence in making that assumption.

There are a number of statistical suites for testing PRNGs out there: I am aware of at least the NIST statistical test suite, the DieHard suite, and TestU01, being most familiar with the latter (this is not saying much). The TestU01 suite is implemented in C; you provide it with a PRNG represented as a function pointer that returns 32-bit integers, and it will run the desired battery of tests (SmallCrush, Crush, or BigCrush – each of increasing size and scope) for you. These more or less consist of frequentist tests involving the chi-square distribution, with more specialised test statistics appearing on occasion. Results are provided in terms of p-values on these statistics.

I found this page that gave some instructions for using TestU01, but the ISAAC implementation I wanted to test is of course written in Javascript. So I knew there was some FFI hackery to be done in order to get the two codebases to play nicely together. I also discovered that Jan de Mooij did some work on testing JS’s basic ‘Math.random’ generator with TestU01 using Emscripten, an LLVM-to-JS compiler, so these two resources seemed a useful place to start.

After several hours of ~~bashing my way through emcc compilation and linker errors~~ careful and methodical programming, I managed to get everything to work. Since the documentation for these tools can be kind of sparse, I figured I’d outline the steps to run the tests here, and hopefully save somebody else a little bit of time and effort in the future. But that said, this is inevitably a somewhat tricky procedure – when trying to reproduce my steps, I ran into a few strange errors that required additional manual fiddling. The following should be about right, but may require some tweaking on your end.

Emscripten and TestU01

Install and activate Emscripten in the current terminal. To go from the official git repo:

$ git clone https://github.com/juj/emsdk.git
$ cd emsdk
$ ./emsdk install latest
$ ./emsdk activate latest
$ source ./emsdk_env.sh

Emscripten provides a couple of wrappers over existing tools; notably, there’s emconfigure and emmake for specialising make builds for compilation via emcc, the Emscripten compiler itself.

In some other directory, grab the TestU01 suite from the Université de Montréal website and extract it:

$ wget http://simul.iro.umontreal.ca/testu01/TestU01.zip
$ unzip -q TestU01.zip

This is some oldish, gnarly academic C code that uses a very wonky autoconf/automake-based build system. There is probably a better way to do it, but the easiest way to get this thing built without too much grief is to build it twice – once as per normal, specifying the appropriate base directory, and once again to specialise it for Emscripten’s use:

$ basedir=`pwd`
$ cd TestU01-1.2.3
$ ./configure --prefix="$basedir"
$ make -j 2
$ make -j 2 install

If all goes well you’ll see ‘bin’, ‘include’, ‘lib’, and ‘share’ directories pop up in ‘basedir’. Repeat the analogous steps for emscripten; note that you’ll get some “no input files” errors here when the configure script checks dynamic linker characteristics, but these are inconsequential:

$ emconfigure ./configure --prefix="$basedir"
$ emmake make -j 2
$ emmake make -j 2 install

Similarly, you’ll notice some warnings re: dynamic linking when building. We’ll handle these later. In the meantime, you can return to your ‘basedir’ to continue working:

$ cd $basedir

A Test Shim for ISAAC

Check out the raw ISAAC code. It’s structured in a sort-of-object-y way; the state of the PRNG is held in a bunch of opaque internal variables, and the whole thing is initialised by calling the ‘isaac’ function and binding the result as a variable. One then iterates the PRNG by calling either the ‘rand’ or ‘random’ property of that variable for a random integer or double, respectively.

We need to write the actual testing code in C. You can get away with the following, which I’ve adapted from M.E. O’Neill’s code – call it something like ‘test-isaac.c’:

#include <emscripten.h>
#include "TestU01.h"

extern void isaac_init(void);
extern unsigned int isaac_rand(void);

int main()
{
    isaac_init();

    unif01_Gen* gen = unif01_CreateExternGenBits("ISAAC", isaac_rand);

    bbattery_SmallCrush(gen);

    unif01_DeleteExternGenBits(gen);

    return 0;
}

Note the two external functions I’m declaring here: the first mimics calling the ‘isaac’ function in Javascript and binding it to a variable, ‘isaac’, and the second mimics a call to isaac.rand(). The testing code follows the same pattern: isaac_init() initialises the generator state, and isaac_rand produces a value from it. The surrounding code passes isaac_rand in as the generator to use for the SmallCrush battery of tests.

C to LLVM Bitcode

We can compile this to LLVM IR as it is, via Emscripten. But first, recall those dynamic linker warnings from the initial setup step. Emscripten deals with a lot of files, compile targets, etc. based on file extension. We thus need to rename all the .dylib files in the ‘lib’ directory, which are in fact all LLVM bitcode, to have the appropraite .bc extension. From the ‘lib’ directory itself, one can just do the following in bash:

$ for old in *.dylib; do mv $old `basename $old .dylib`.bc; done

When that’s done, we can compile the C code to LLVM with emcc:

$ emcc -O3 -o test-isaac.bc \
    -I/$basedir/include \
    -L/$basedir/lib \
    -ltestu01 -lprobdist -lmylib -lm -I/usr/local/include \
    test-isaac.c

Again, Emscripten decides what its compile target should be via its file extension – thus here, an output with the .bc extension means we’re compiling to LLVM IR.

LLVM to Javascript and WASM

Now, to provide the requisite isaac_init and isaac_rand symbols to the compiled LLVM bitcode, we need to pass the ISAAC library itself. This is another finicky procedure, but there is a method to the madness, and one can find a bit of documentation on it here.

Helpfully, Evan Wallace at Figma wrote an emscripten JS library generation helper that makes this task a little less painful. Install that via npm as follows:

$ npm install -g emscripten-library-generator

To wrap the ISAAC code up in the appropriate format, one needs to make a few small modifications to it. I won’t elaborate on this too much, but one needs to:

Change the String.prototype.toIntArray function declaration to a simple function toIntArray(string) declaration, and alter its use in the code appropriately,
Change the var isaac = ... function declaration/execution binding to a simple function isaac() declaration, and,

Declare the two functions used in our C shim:

function isaac_init() {
  var isaac_initialised = isaac();
}

function isaac_rand() {
  return isaac_initialised.rand();
}

Then we can package it up in an emscripten-friendly format as follows:

$ emscripten-library-generator isaac-mod.js > isaac-lib.js

You’ll need to make one last adjustment. In the isaac-lib.js file just generated for us, add the following emscripten ‘postset’ instruction above the ‘isaac’ property:

isaac__postset: 'var isaac_initialised = _isaac();',   // add me
isaac: function () {                                   // leave me alone

This makes sure that the isaac_initialised variable referred to in isaac_rand is accessible.

Whew. Ok, with all that done, we’ll compile our LLVM bytecode to a Javascript and wasm pair. You’ll need to bump up the total memory option in order to run the resulting code – I think I grabbed the amount I used from Jan de Mooij’s example:

$ emcc --js-library isaac-lib.js \
    lib/libtestu01.0.0.1.bc \
    -o test-isaac.js \
    -s TOTAL_MEMORY=536870912 \
    test-isaac.bc

Running SmallCrush

That’s about it. If all has gone well, you should have seen no compilation or linking errors when running emcc, and you should have a couple of ‘test-isaac.js’ and ‘test-isaac.wasm’ files kicking around in your ‘basedir’.

To (finally) run the Small Crush suite, execute ‘test-isaac.js’ with node:

$ node test-isaac.js
xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
                 Starting SmallCrush
                 Version: TestU01 1.2.3
xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx


***********************************************************
HOST = emscripten, Emscripten

ISAAC


smarsa_BirthdaySpacings test:
-----------------------------------------------
   N =  1,  n = 5000000,  r =  0,    d = 1073741824,    t = 2,    p = 1


      Number of cells = d^t = 1152921504606846976
      Lambda = Poisson mean =      27.1051


----------------------------------------------------
Total expected number = N*Lambda      :      27.11
Total observed number                 :      27
p-value of test                       :    0.53


-----------------------------------------------
CPU time used                    :  00:00:00.00

Generator state:

<more output omitted>

========= Summary results of SmallCrush =========

 Version:          TestU01 1.2.3
 Generator:        ISAAC
 Number of statistics:  15
 Total CPU time:   00:00:00.00

 All tests were passed

Et voilà, your SmallCrush battery output is dumped to stdout. You need only to tweak the C code shim and recompile if you want to run the more intensive Crush or BigCrush suites. Similarly, you can dump generator output to stdout with console.log() if you want to reassure yourself that the generator is running correctly.

Fin

So: the Javascript ISAAC PRNG indeed passes TestU01. Nice! It was satisfying enough to get this hacky sequence of steps to actually run, but it was even better to see the tests actually pass, as I’d both hoped and expected.

I did a few extra things to convince myself that ISAAC was really passing the tests as it seemed to be doing. I ran TestU01 on a cheap little xorshift generator (which failed several tests) and also did some ad-hoc analysis of values that I had ISAAC log to stdout. And, I even looked at the code, and compared it with a reference implementation by eye. At this point, I’m convinced it operates as advertised, and I feel very comfortable using it in my application.

Transforming to CPS

2018-08-04T00:00:00+04:00

I recently picked up Appel’s classic Compiling with Continuations and have been refreshing my continuation-fu more generally.

Continuation-passing style (CPS) itself is nothing uncommon to the functional programmer; it simply involves writing in a manner such that functions never return, instead passing control over to something else (a continuation) to finish the job. The simplest example is just the identity function, which in CPS looks like this:

id :: a -> (a -> b) -> b
id x k = k x

The first argument is the conventional identity function argument – the second is the continuation. I wrote a little about continuations in the context of the Giry monad, which is a somewhat unfamiliar setting, but one that follows the same principles as anything else.

In this post I just want to summarise a few useful CPS transforms and related techniques in one place.

Manual CPS Transformation

Consider a binary tree type. We’ll keep things simple here:

data Tree a =
    Leaf a
  | Branch a (Tree a) (Tree a)

Calculating the depth of a tree is done very easily:

depth :: Tree a -> Int
depth = loop where
  loop tree = case tree of
    Leaf _       -> 1
    Branch _ l r ->
      let dl = loop l
          dr = loop r
      in  succ (max dl dr)

Note however that this is not a tail-recursive function – that is, it does not end with a call to itself (instead it ends with a call to something like ‘succ . uncurry max’). This isn’t necessarily a big deal – the function is easy to read and write and everything, and certainly has fine performance characteristics in Haskell – but it is less easy to deal with for, say, an optimising compiler that may want to handle evaluation in this or that alternative way (primarily related to memory management).

One can construct a tail-recursive (depth-first) version of ‘depth’ via a manual CPS transformation. The looping function is simply augmented to take an additional continuation argument, like so:

depth :: Tree a -> Int
depth tree = loop tree id where
  loop cons k = case cons of
    Leaf _       -> k 1
    Branch _ l r ->
      loop l $ \dl ->
        loop r $ \dr ->
          k (succ (max dl dr))

Notice now that the ‘loop’ function terminates with a call to itself (or just passes control to a supplied continuation), and is thus tail-recursive.

Due to the presence of the continuation argument, ‘loop’ is a higher-order function. This is fine and dandy in Haskell, but there is a neat technique called defunctionalisation that allows us to avoid the jump to higher-order and makes sure things stay KILO (“keep it lower order”), which can be simpler to deal with more generally.

The idea is just to reify the continuations as abstract syntax, and then evaluate them as one would any embedded language. Note the continuation \dl -> .., for example – the free parameters ‘r’ and ‘k’ occuring in the function body correspond to a tree (the right subtree) and another continuation, respectively. And in \dr -> .. one has the free parameters ‘dl’ and ‘k’ – now the depth of the left subtree, and the other continuation again. We also have ‘id’ used on the initial call to ‘loop’. These can all be reified via the following data type:

data DCont a =
    DContL (Tree a) (DCont a)
  | DContR Int (DCont a)
  | DContId

Note that this is a very simple recursive type – it has a simple list-like pattern of recursion, in which each ‘level’ of a value is either a constructor, carrying both a field of some type and a recursive point, or is the ‘DContId’ constructor, which simply terminates the recursion. The reified continuations are, on a suitable level of abstraction, more or less the sequential operations to be performed in the computation. In other words: by reifying the continuations, we also reify the stack of the computation.

Now ‘depth’ can be rewritten such that its looping function is not higher-order; the cost is that another function is needed, one that lets us evaluate items (again, reified continuations) on the stack:

depth :: Tree a -> Int
depth tree = loop tree DContId where
  loop cons k = case cons of
    Leaf _       -> eval k 1
    Branch _ l r -> loop l (DContL r k)

  eval cons d = case cons of
    DContL r k  -> loop r (DContR d k)
    DContR dl k -> eval k (succ (max dl d))
    DContId     -> d

The resulting function is mutually tail-recursive in terms of both ‘loop’ and ‘eval’, neither of which are higher-order.

One can do a little better in this particular case and reify the stack using an actual Haskell list, which simplifies evaluation somewhat – it just requires that the list elements have a type along the lines of ‘(Tree a, Int)’ rather than something like ‘Either (Tree a) Int’, which is effectively what we get from ‘DCont a’. You can see an example of this in this StackOverflow answer by Chris Taylor.

Mechanical CPS Transformation

“Mechanical CPS transformation” might be translated as simply “compiling with continuations.” Matt Might has quite a few posts on this topic; in particular he has one very nice post on mechanical CPS conversion that summarises various transformations described in Appel, etc.

Matt describes three transformations that I think illustrate the general mechanical CPS business very well (he describes more, but they are more specialised). The first is a “naive” transformation, which is simple, but produces a lot of noisy “administrative redexes” that must be cleaned up in another pass. The second is a higher-order transformation, which makes use of the host language’s facilities for function definition and application – it produces simpler code, but some unnecessary noise still leaks through. The last is a “hybrid” transformation, which makes use of both the naive and higher-order transformations, depending on which is more appropriate.

Let’s take a look at these in Haskell. First let’s get some imports out of the way:

{-# LANGUAGE OverloadedStrings #-}

import Data.Monoid
import Data.Text (Text)
import qualified Data.Text as T
import Data.Unique
import qualified Text.PrettyPrint.Leijen.Text as PP

I’ll also make use of a simple, Racket-like ‘gensym’ function:

gensym :: IO Text
gensym = fmap render newUnique where
  render u =
    let hu = hashUnique u
    in  T.pack ("$v" <> show hu)

We’ll use a bare-bones lambda calculus as our input language. Many examples – Appel’s especially – use significantly more complex languages when illustrating CPS transforms, but I think this distracts from the meat of the topic. Lambda does just fine:

data Expr =
    Lam Text Expr
  | Var Text
  | App Expr Expr

I want to render expressions in my input and output languages in a Lisp-like manner. This is very easy to do using a good pretty-printing library; here I’m using the excellent wl-pprint-text, and will omit the ‘Pretty’ instances in the body of my post. But I’ll link to a gist including them at the bottom.

When performing a mechanical CPS transform, one targets both “atomic” expressions – i.e., variables and lambda abstractions – and “complex” expressions, i.e. function application. The target language is thus a combination of the ‘AExpr’ and ‘CExpr’ types:

data AExpr =
    AVar Text
  | ALam [Text] CExpr

data CExpr =
    CApp AExpr [AExpr]

All the mechanical CPS transformations use variants on two functions going by the cryptic names m and t. m is responsible for converting atomic expressions in the input languages (i.e., variables and lambda abstractions) into atomic expressions in the target language (an atomic CPS expression). t is the actual CPS transformation; it converts an expression in the input language into CPS, invoking a specified continuation (already in the target language) on the result.

Let’s look at the naive transform. Here are m and t, prefixed by ‘n’ to indicate that they are naive. First, m:

nm :: Expr -> IO AExpr
nm expr = case expr of
  Lam var cexpr0 -> do
    k      <- gensym
    cexpr1 <- nt cexpr0 (AVar k)
    return (ALam [var, k] cexpr1)

  Var var -> return (AVar var)

  App {} -> error "non-atomic expression"

(N.b. you almost never want to use ‘error’ in a production implementation of anything. It’s trivial to wrap e.g. ‘MaybeT’ around the appropriate functions to handle the bogus pattern match on ‘App’ totally, but I just want to keep the types super simple here.)

The only noteworthy thing that m does here is in the case of a lambda abstraction: a new abstract continuation is generated, and the body of the abstraction is converted to CPS via t, such that the freshly-generated continuation is called on the result. Remember, m is really just mapping atomic expressions in the input language to atomic expressions in the target language.

Here’s t for the naive transform. Remember, t is responsible for converting expressions to continuation-passing style:

nt :: Expr -> AExpr -> IO CExpr
nt expr cont = case expr of
  Lam {} -> do
    aexpr <- m expr
    return (CApp cont [aexpr])

  Var _ -> do
    aexpr <- m expr
    return (CApp cont [aexpr])

  App f e -> do
    fs <- gensym
    es <- gensym

    let aexpr0 = ALam [es] (CApp (AVar fs) [AVar es, cont])

    cexpr <- nt e aexpr0

    let aexpr1 = ALam [fs] cexpr

    nt f aexpr1

For both kinds of atomic expressions (lambda and variable), the expression is converted to the target language via m, and then the supplied continuation is applied to it. Very simple.

In the case of function application (a “complex”, or non-atomic expression), both the function to be applied, and the argument it is to be applied to, must be converted to CPS. This is done by generating two fresh continuations, transforming the argument, and then transforming the function. The control flow here is always handled by stitching continuations together; notice when transforming the function ‘f’ that the continuation to be applied has already handled its argument.

Next, the higher-order transform. Here are m and t:

hom :: Expr -> IO AExpr
hom expr = case expr of
  Lam var e -> do
    k  <- gensym
    ce <- hot e (\rv -> return (CApp (AVar k) [rv]))
    return (ALam [var, k] ce)

  Var n  -> return (AVar n)

  App {} -> error "non-atomic expression"

hot :: Expr -> (AExpr -> IO CExpr) -> IO CExpr
hot expr k = case expr of
  Lam {} -> do
    aexpr <- m expr
    k aexpr

  Var {} -> do
    aexpr <- m expr
    k aexpr

  App f e -> do
    rv      <- gensym
    xformed <- k (AVar rv)

    let cont = ALam [rv] xformed

        cexpr fs = hot e (\es -> return (CApp fs [es, cont]))

    hot f cexpr

Both of these have the same form as they do in the naive transform – the difference here is simply that the continuation to be applied to a transformed expression is expressed in the host language – i.e., here, Haskell. Thus the transform is “higher-order,” in exactly the same sense that higher-order abstract syntax is higher-order.

The final transformation I’ll illustrate here, the hybrid transform, applies the naive transformation to lambda and variable expressions, and applies the higher-order transformation to function applications. Here t is split up into tc and tk to handle these cases accordingly:

m :: Expr -> IO AExpr
m expr = case expr of
  Lam var cexpr -> do
    k       <- gensym
    xformed <- tc cexpr (AVar k)
    return (ALam [var, k] xformed)

  Var n  -> return (AVar n)

  App {} -> error "non-atomic expression"

tc :: Expr -> AExpr -> IO CExpr
tc expr c = case expr of
  Lam {} -> do
    aexpr <- m expr
    return (CApp c [aexpr])

  Var _ -> do
    aexpr <- m expr
    return (CApp c [aexpr])

  App f e -> do
    let cexpr fs = tk e (\es -> return (CApp fs [es, c]))
    tk f cexpr

tk :: Expr -> (AExpr -> IO CExpr) -> IO CExpr
tk expr k = case expr of
  Lam {} -> do
    aexpr <- m expr
    k aexpr

  Var {} -> do
    aexpr <- m expr
    k aexpr

  App f e -> do
    rv <- gensym
    xformed <- k (AVar rv)

    let cont     = ALam [rv] xformed
        cexpr fs = tk e (\es -> return (CApp fs [es, cont]))

    tk f cexpr

Matt illustrates these transformations on a simple expression: (g a). We can do the same:

test :: Expr
test = App (Var "g") (Var "a")

First, the naive transform. Note all the noisy administrative redexes that come along with it:

> cexpr <- nt test (AVar "halt")
> PP.pretty cexpr
((λ ($v1).
  ((λ ($v2).
    ($v1 $v2 halt)) a)) g)

The higher-order transform does better, containing only one such redex (an eta-expansion). Note that since the supplied continuation must be expressed in terms of a Haskell function, we need to write it in a more HOAS-y style:

> cexpr <- hot test (\ans -> return (CApp (AVar "halt") [ans]))
> PP.pretty cexpr
(g a (λ ($v3).
  (halt $v3)))

Finally the hybrid transform, which, here, is literally perfect. We don’t even need to deal with the minor annoyance of the HOAS-style continuation when calling it:

> cexpr <- tc test (AVar "halt")
> PP.pretty cexpr
(g a halt)

Matt goes on to describe a “partioned CPS transform” that can be used to recover a stack, in (seemingly) much the same manner that the defunctionalised manual CPS transform worked in the previous section. Very neat, but something I’ll have to look at in another post.

Fin

CPS is pretty gnarly. My experience in compiling with continuations is not substantial, but I dig learning it. Appel’s book, in particular, is meaty – expect more posts on the subject here eventually, probably.

‘Til next time! I’ve dumped the code from the latter section into a gist.

Embedded DSLs for Bayesian Modelling and Inference: a Retrospective

2018-07-02T00:00:00+04:00

Why does my blog often feature its typical motley mix of probability, functional programming, and computer science anyway?

From 2011 through 2017 I slogged through a Ph.D. in statistics, working on it full time in 2012, and part-time in every other year. It was an interesting experience. Although everything worked out for me in the end – I managed to do a lot of good and interesting work in industry while still picking up a Ph.D. on the side – it’s not something I’d necessarily recommend to others. The smart strategy is surely to choose one thing and give it one’s maximum effort; by splitting my time between work and academia, both obviously suffered to some degree.

That said, at the end of the day I was pretty happy with the results on both fronts. On the academic side of things, the main product was a dissertation, Embedded Domain-Specific Languages for Bayesian Modelling and Inference, supporting my thesis: that novel and useful DSLs for solving problems in Bayesian statistics can be embedded in statically-typed, purely functional programming languages.

It helps to remember that in this day and age, one can still typically graduate by, uh, “merely” submitting and defending a dissertation. Publishing in academic venues certainly helps focus one’s work, and is obviously necessary for a career in academia (or, increasingly, industrial research). But it’s optional when it comes to getting your degree, so if it doesn’t help you achieve your goals, you may want to reconsider it, as I did.

The problem with the dissertation-first approach, of course, is that nobody reads your work. To some extent I think I’ve mitigated that; most of the content in my dissertation is merely a fleshed-out version of various ideas I’ve written about on this blog. Here I’ll continue that tradition and write a brief, informal summary of my dissertation and Ph.D. more broadly – what I did, how I approached it, and what my thoughts are on everything after the fact.

The Idea

Following the advice of Olin Shivers (by way of Matt Might), I oriented my work around a concrete thesis, which wound up more or less being that embedding DSLs in a Haskell-like language can be a useful technique for solving statistical problems. This thesis wasn’t born into the world fully-formed, of course – it began as quite a vague (or misguided) thing, but matured naturally over time. Using the tools of programming languages and compilers to do statistics and machine learning is the motivation behind probabilistic programming in general; what I was interested in was exploring the problem in the setting of languages embedded in a purely functional host. Haskell was the obvious choice of host for all of my implementations.

It may sound obvious that putting together a thesis is a good strategy for a Ph.D. But here I’m talking about a thesis in the original (Greek) sense of a proposition, i.e. a falsifiable idea or claim (in contrast to a dissertation, from the Latin disserere, i.e. to examine or to discuss). Having a central idea to orient your work around can be immensely useful in terms of focus. When you read a dissertation with a clear thesis, it’s easy to know what the writer is generally on about – without one it can (increasingly) be tricky.

My thesis is pretty easy to defend in the abstract. A DSL really exposes the structure of one’s problem while also constraining it appropriately, and embedding one in a host language means that one doesn’t have to implement an entire compiler toolchain to support it. I reckoned that simply pointing the artillery of “language engineering” at the statistical domain would lead to some interesting insight on structure, and maybe even produce some useful tools. And it did!

The Contributions

Of course, one needs to do a little more defending than that to satisfy his or her examination committee. Doctoral research is supposed to be substantial and novel. In my experience, reviewers are concerned with your answers to the following questions:

What, specifically, are your claims?
Are they novel contributions to your field?
Have you backed them up sufficiently?

At the end of the day, I claimed the following advances from my work.

Novel probabilistic interpretations of the Giry monad’s algebraic structure. The Giry monad (Lawvere, 1962; Giry, 1981) is the “canonical” probability monad, in a meaningful sense, and I demonstrated that one can characterise the measure-theoretic notion of image measure by its functorial structure, as well as the notion of product measure by its monoidal structure. Having the former around makes it easy to transform the support of a probability distribution while leaving its density structure invariant, and the latter lets one encode probabilistic independence, enabling things like measure convolution and the like. What’s more, the analogous semantics carry over to other probability monads – for example the well-known sampling monad, or more abstract variants.
A novel characterisation of the Giry monad as a restricted continuation monad. Ramsey & Pfeffer (2002) discussed an “expectation monad,” and I had independently come up with my own “measure monad” based on continuations. But I showed both reduce to a restricted form of the continuation monad of Wadler (1994) – and that indeed, when the return type of Wadler’s continuation monad is restricted to the reals, it is the Giry monad.

To be precise it’s actually somewhat more general – it permits integration with respect to any measure, not only a probability measure – but that definition strictly subsumes the Giry monad. I also showed that product measure, via the applicative instance, yields measure convolution and associated operations.
A novel technique for embedding a statically-typed probabilistic programming language in a purely functional language. The general idea itself is well-known to those who have worked with DSLs in Haskell: one constructs a base functor and wraps it in the free monad. But the reason that technique is appropriate in the probabilistic programming domain is that probabilistic models are fundamentally monadic constructs – merely recall the existence of the Giry monad for proof!

To construct the requisite base functor, one maps some core set of concrete probability distributions denoted by the Giry monad to a collection of abstract probability distributions represented only by unique names. These constitute the branches of one’s base functor, which is then wrapped in the familiar ‘Free’ machinery that gives one access to the functorial, applicative, and monadic structure that I talked about above. This abstract representation of a probabilistic model allows one to implement other probability monads, such as the well-known sampling monad (Ramsey & Pfeffer, 2002; Park et al., 2008) or the Giry monad, by way of interpreters.

(N.b. Ścibior et al. (2015) did some very similar work to this, although the monad they used was arguably more operational in its flavour.)
A novel characterisation of execution traces as cofree comonads. The idea of an “execution trace” is that one runs a probabilistic program (typically generating a sample) and then records how it executed – what randomness was used, the execution path of the program, etc. To do Bayesian inference, one then runs a Markov chain over the space of possible execution traces, calculating statistics about the resulting distribution in trace space (Wingate et al., 2011).

Remarkably, a cofree comonad over the same abstract probabilistic base functor described above allows us to represent an execution trace at the embedded language level itself. In practical terms, that means one can denote a probabilistic model, and then run a Markov chain over the space of possible ways it could have executed, without leaving GHCi. You can alternatively examine and perturb the way the program executes, stepping through it piece by piece, as I believe was originally a feature in Venture (Mansinghka et al., 2014).

(N.b. this really blew my mind when I first started toying with it, converting programs into execution traces and then manipulating them as first-class values, defining other probabilistic programs over spaces of execution traces, etc. Meta.)
A novel technique for statically encoding conditional independence of terms in this kind of embedded probabilistic programming language. If you recall that I previously demonstrated the monoidal (i.e. applicative) structure of the Giry monad encodes the notion of product measure, it will not be too surprising to hear that I used the free applicative functor (Capriotti & Kaposi, 2014) (again, over the same kind of abstract probabilistic base functor) to reify applicative expressions such that they can be identified statically.
A novel shallowly-embedded language for building custom transition operators for use in Markov chain Monte Carlo. MCMC is the de-facto standard way to perform inference on Bayesian models (although it is not limited to Bayesian models in particular). By wrapping a simple state monad transformer around a probability monad, one can denote Markov transition operators, combine them, and transform them in a few ways that are useful for doing MCMC.

The framework here was inspired by the old parallel “strategies” idea of Trinder et al. (1998). The idea is that you want to “evaluate” a posterior via MCMC, and want to choose a strategy by which to do so – e.g. Metropolis (Metropolis, 1953), slice sampling (Neal, 2003), Hamiltonian (Neal, 2011), etc. Since Markov transition operators are closed under composition and convex combinations, it is easy to write a little shallowly-embedded combinator language for working with them – effectively building evaluation strategies in a manner familiar to those who’ve worked with Haskell’s parallel library.

(N.b. although this was the most trivial part of my research, theoretical or implementation-wise, it remains the most useful for day-to-day practical work.)

The Execution

One needs to stitch his or her contributions together in some kind of over-arching narrative that supports the underlying thesis. Mine went something like this:

The Giry monad is appropriate for denoting probabilistic semantics in languages with purely-functional hosts. Its functorial, applicative, and monadic structure denote probability distributions, independence, and marginalisation, respectively, and these are necessary and sufficient for encoding probabilistic models. An embedded language based on the Giry monad is type-safe and composable.

Probabilistic models in an embedded language, semantically denoted in terms of the Giry monad, can be made abstract and interpretation-independent by defining them in terms of a probabilistic base functor and a free monad instead. They can be forward-interpreted using standard free monad recursion schemes in order to compute probabilities (via a measure intepretation) or samples (via a sampling interpretation); the latter interpretation is useful for performing limited forms of Bayesian inference, in particular. These free-encoded models can also be transformed into cofree-encoded models, under which they represent execution traces that can be perturbed arbitrarily by standard comonadic machinery. This representation is amenable to more elaborate forms of Bayesian inference. To accurately denote conditional independence in the embedded language, the free applicative functor can also be used.

One can easily construct a shallowly-embedded language for building custom Markov transitions. Markov chains that use these compound transitions can outperform those that use only “primitive” transitions in certain settings. The shallowly embedded language guarantees that transitions can only be composed in well-defined, type-safe ways that preserve the properties desirable for MCMC. What’s more, one can implement “transition transformers” for implementing still more complex inference techniques, e.g. annealing or tempering, over existing transitions.

Thus: novel and useful domain-specific languages for solving problems in Bayesian statistics can be embedded in statically-typed, purely-functional programming languages.

I used the twenty-minute talk period of my defence to go through this narrative and point out my claims, after which I was grilled on them for an hour or two. The defence was probably the funnest part of my whole Ph.D.

The Product

In the end, I mainly produced a dissertation, a few blog posts, and some code. By my count, the following repos came out of the work:

deanie: An embedded probabilistic programming language.
http://github.com/jtobin/deanie
declarative: DIY Markov Chains.
http://github.com/jtobin/declarative
flat-mcmc: Painless general-purpose sampling.
http://github.com/jtobin/flat-mcmc
hasty-hamiltonian: Speedy traversal through parameter space.
http://github.com/jtobin/hasty-hamiltonian
hnuts: Automatic gradient-based sampling.
http://github.com/jtobin/hnuts
lazy-langevin: Gradient-based diffusion.
http://github.com/jtobin/lazy-langevin
mcmc-types: Common types for implementing MCMC algorithms.
https://github.com/jtobin/mcmc-types
measurable: A shallowly-embedded DSL for basic measure wrangling.
http://github.com/jtobin/measurable
mighty-metropolis: The Metropolis sampling algorithm.
http://github.com/jtobin/mighty-metropolis
mwc-probability: Sampling function-based probability distributions.
http://github.com/jtobin/mwc-probability
sampling: Tools for sampling from collections.
https://github.com/jtobin/sampling
speedy-slice: Speedy slice sampling.
http://github.com/jtobin/speedy-slice

If any of this stuff is or was useful to you, that’s great! I still use the declarative libraries, flat-mcmc, mwc-probability, and sampling pretty regularly. They’re fast and convenient for practical work.

Some of the other stuff, e.g. measurable, is useful for building intuition, but not so much in practice, and deanie, for example, is a work-in-progress that will probably not see much more progress (from me, at least). Continuing from where I left off might be a good idea for someone who wants to explore problems in this kind of setting in the future.

General Thoughts

When I first read about probabilistic (functional) programming in Dan Roy’s 2011 dissertation I was absolutely blown away by the idea. It seemed that, since there was such an obvious connection between the structure of Bayesian models and programming languages (via the underlying semantic graph structure, something that has been exploited to some degree as far back as BUGS), it was only a matter of time until someone was able to really create a tool that would revolutionize the practice of Bayesian statistics.

Now I’m much more skeptical. It’s true that probabilistic programming tends to expose some beautiful structure in statistical models, and that a probabilistic programming language that was easy to use and “just worked” for inference would be a very useful tool. But putting something expressive and usable together that also “just works” for that inference step is very, very difficult. Very difficult indeed.

Almost every probabilistic programming framework of the past ten years, from Church down to my own stuff, has more or less wound up as “thesisware,” or remains the exclusive publication-generating mechanism of a single research group. The exceptions are almost unexceptional in of themselves: JAGS and Stan are probably the most-used such frameworks, certainly in statistics (I will mention the very honourable PyMC here as well), but they innovate little, if at all, over the original BUGS in terms of expressiveness. Similarly it’s very questionable whether the fancy MCMC algo du jour is really any better than some combination of Metropolis-Hastings (even plain Metropolis), Gibbs (or its approximate variant, slice sampling), or nested sampling in anything outside of favourably-engineered examples (I will note that Hamiltonian Monte Carlo could probably be counted in there too, but it can still be quite a pain to use, its variants are probably overrated, and it is comparatively expensive).

Don’t get me wrong. I am a militant Bayesian. Bayesian statistics, i.e., as far as I’m concerned, probability theory, describes the world accurately. And there’s nothing wrong with thesisware, either. Research is research, and this is a very thorny problem area. I hope to see more abandoned, innovative software that moves the ball up the field, or kicks it into another stadium entirely. Not less. The more ingenious implementations and sampling schemes out there, the better.

But more broadly, I often find myself in the camp of Leo Breiman, who in 2001 characterised the two predominant cultures in statistics as those of data modelling and algorithmic modelling respectively, the latter now known as machine learning, of course. The crux of the data modelling argument, which is of course predominant in probabilistic programming research and Bayesian statistics more generally, is that a practitioner, by means of his or her ingenuity, is able to suss out the essence of a problem and distill it into a useful equation or program. Certainly there is something to this: science is a matter of creating hypotheses, testing them against the world, and iterating on that, and the “data modelling” procedure is absolutely scientific in principle. Moreover, with a hat tip to Max Dama, one often wants to impose a lot of structure on a problem, especially if the problem is in a domain where there is a tremendous amount of noise. There are many areas where this approach is just the thing one is looking for.

That said, it seems to me that a lot of the data modelling-flavoured side of probabilistic programming, Bayesian nonparametrics, etc., is to some degree geared more towards being, uh, “research paper friendly” than anything else. These are extremely seductive areas for curious folks who like to play at the intersection of math, statistics, and computer science (raises hand), and one can easily spend a lifetime chasing this or that exquisite theoretical construct into any number of rabbit holes. But at the end of the day, the data modelling culture, per Breiman:

.. has at its heart the belief that a statistician, by imagination and by looking at the data, can invent a reasonably good parametric class of models for a complex mechanism devised by nature.

Certainly the traditional statistics that Breiman wrote about in 2001 was very different from probabilistic programming and similar fields in 2018. But I think there is the same element of hubris in them, and to some extent, a similar dissociation from reality. I have cultivated some of the applied bent of a Breiman, or a Dama, or a Locklin, so perhaps this should not be too surprising.

I feel that the 2012-ish resurgence of neural networks jolted the machine learning community out of a large-scale descent into some rather dubious Bayesian nonparametrics research, which, much as I enjoy that subject area, seemed more geared towards generating fun machine learning summer school lectures and NIPS papers than actually getting much practical work done. I can’t help but feel that probabilistic programming may share a few of those same characteristics. When all is said and done, answering the question is this stuff useful? often feels like a stretch.

So: onward & upward and all that, but my enthusiasm has been tempered somewhat, is all.

Fini

Administrative headaches and the existential questions associated with grad school aside, I had a great time working in this area for a few years, if in my own aloof and eccentric way.

If you ever interacted with this area of my work, I hope you got some utility out of it: research ideas, use of my code, or just some blog post that you thought was interesting during a slow day at the office. If you’re working in the area, or are considering it, I wish you success, whether your goal is to build practical tools, or to publish sexy papers. :-)

Fubini and Applicatives

2018-06-27T00:00:00+04:00

Take an iterated integral, e.g. $\int_X \int_Y f(x, y) dy dx$. Fubini’s Theorem describes the conditions under which the order of integration can be swapped on this kind of thing while leaving its value invariant. If Fubini’s conditions are met, you can convert your integral into $\int_Y \int_X f(x, y) dx dy$ and be guaranteed to obtain the same result you would have gotten by going the other way.

What are these conditions? Just that you can glue your individual measures together as a product measure, and that $f$ is integrable with respect to it. I.e.,

\[\int_{X \times Y} | f(x, y) | d(x \times y) < \infty.\]

Say you have a Giry monad implementation kicking around and you want to see how Fubini’s Theorem works in terms of applicative functors, monads, continuations, and all that. It’s pretty easy. You could start with my old measurable library that sits on GitHub and attracts curious stars from time to time and cook up the following example:

import Control.Applicative ((<$>), (<*>))
import Measurable

dx :: Measure Int
dx = bernoulli 0.5

dy :: Measure Double
dy = beta 1 1

dprod :: Measure (Int, Double)
dprod = (,) <$> dx <*> dy

Note that ‘dprod’ is clearly a product measure (I’ve constructed it using the Applicative instance for the Giry monad, so it must be a product measure) and take a simple, obviously integrable function:

add :: (Int, Double) -> Double
add (m, x) = fromIntegral m + x

Since ‘dprod’ is a product measure, Fubini’s Theorem guarantees that the following are equivalent:

i0 :: Double
i0 = integrate add dprod

i1 :: Double
i1 = integrate (\x -> integrate (curry add x) dy) dx

i2 :: Double
i2 = integrate (\y -> integrate (\x -> curry add x y) dx) dy

And indeed they are – you can verify them yourself if you don’t believe me (or our boy Fubini).

For an example of a where interchanging the order of integration would be impossible, we can construct some other measure:

dpair :: Measure (Int, Double)
dpair = do
  x <- dx
  y <- fmap (* fromIntegral x) dy
  return (x, y)

It can be integrated as follows:

i3 :: Double
i3 = integrate (\x -> integrate (curry add x) (fmap (* fromIntegral x) dy)) dx

But notice how ‘dpair’ is constructed: it is strictly monadic, not applicative, so the order of the expressions matters. Since ‘dpair’ can’t be expressed as a product measure (i.e. by an applicative expression), Fubini says that swapping the order of integration is a no-no.

Note that if you were to just look at the types of ‘dprod’ and ‘dpair’ – both ‘Measure (Int, Double)’ – you wouldn’t be able to tell immediately that one represents a product measure while the other one does not. If being able to tell these things apart statically is important to you (say, you want to statically apply order-of-integration optimisations to integral expressions or what have you), you need look no further than the free applicative functor to help you out.

Fun fact: there is a well-known variant of Fubini’s Theorem, called Tonelli’s Theorem, that was developed by another Italian guy at around the same time. I’m not sure how early-20th century Italy became so strong in order-of-integration research, exactly.

Byzantine Generals and Nakamoto Consensus

2018-01-22T00:00:00+04:00

You can recognize truth by its beauty and simplicity.

– Richard Feynman (attributed)

In one of his early emails on the Cryptography mailing list, Satoshi claimed that the proof-of-work chain is a solution to the Byzantine Generals Problem (BGP). He describes this via an example where a bunch of generals – Byzantine ones, of course – collude to break a king’s wifi.

It’s interesting to look at this a little closer in the language of the originally-stated BGP itself. One doesn’t need to be too formal to glean useful intuition here.

What, more precisely, did Satoshi claim?

The Decentralized Timestamp Server

Satoshi’s problem is that of a decentralized timestamp server (DTS). Namely, he posits that any number of nodes, following some protocol, can together act as a timestamping server – producing some consistent ordering on what we’ll consider to be abstract ‘blocks’.

The decentralized timestamp server reduces to an instance of the Byzantine Generals Problem as follows. There are a bunch of nodes, who could each be honest or dishonest. All honest nodes want to agree on some ordering – a history – of blocks, and a small number of dishonest nodes should not easily be able to compromise that history – say, by convincing the honest nodes to adopt some alternate one of their choosing.

(N.b. it’s unimportant here to be concerned about the contents of blocks. Since the decentralized timestamp server problem is only concerned about block orderings, we don’t need to consider the case of invalid transactions within blocks or what have you, and can safely assume that any history must be internally consistent. We only need to assume that child blocks depend utterly on their parents, so that rewriting a history by altering some parent block also necessitates rewriting its children, and that honest nodes are constantly trying to append blocks.)

As demonstrated in the introduction to the original paper, the Byzantine Generals Problem can be reduced to the problem of how any given node communicates its information to others. In our context, it reduces to the following:

Byzantine Generals Problem (DTS)

A node must broadcast a history of blocks to its peers, such that:

(IC1) All honest peers agree on the history.

(IC2) If the node is honest, then all honest peers agree with the history it broadcasts.

To produce consensus, every node will communicate its history to others by using a solution to the Byzantine Generals Problem.

Longest Proof-of-Work Chain

Satoshi’s proposed solution to the BGP has since come to be known as ‘Nakamoto Consensus’. It is the following protocol:

Nakamoto Consensus

Always use the longest history.

Appending a block to any history requires a proof that a certain amount of work – proportional in expectation to the total ‘capability’ of the network – has been completed.

To examine how it works, consider an abstract network and communication medium. We can assume that messages are communicated instantly (it suffices that communication is dwarfed in time by actually producing a proof of work) and that the network is static and fixed, so that only active or ‘live’ nodes actually contribute to consensus.

The crux of Nakamoto consensus is that nodes must always use the longest available history – the one that provably has the largest amount of work invested in it – and appending to any history requires a nontrivial amount of work in of itself. Consider a set of nodes, each having some (not necessarily shared) history. Whenever any node broadcasts a one-block longer history, all honest nodes will immediately agree on it, and conditions (IC1) and (IC2) are thus automatically satisfied whether or not the broadcasting node is honest. Nakamoto Consensus trivially solves the BGP in this most important case; we can examine other cases by examining how they reduce to this one.

If two or more nodes broadcast longer histories at approximately the same time, then honest nodes may not agree on a single history for as long as it takes a longer history to be produced and broadcast. As soon as this occurs (which, in all probability, is only a matter of time), we reduce to the previous case in which all honest nodes agree with each other again, and the BGP is resolved.

The ‘bad’ outcome we’re primarily concerned about is that of dishonest nodes rewriting history in their favour, i.e. by replacing some history $\{\ldots, B_1, B_2, B_3, \ldots\}$ by another one $\{\ldots, B_1, B_2', B_3', \ldots\}$ that somehow benefits them. The idea here is that some dishonest node (or nodes) intends to use block $B_2$ as some sort of commitment, but later wants to renege. To do so, the node needs to rewrite not only $B_2$, but all other blocks that depend on $B_2$ (here $B_3$, etc.), ultimately producing a longer history than is currently agreed upon by honest peers.

Moreover, it needs to do this faster than honest nodes are able to produce longer histories on their own. Catching up to and exceeding the honest nodes becomes exponentially unlikely in the number of blocks to be rewritten, and so a measure of confidence can be ascribed to agreement on the state of any sub-history that has been ‘buried’ by a certain number of blocks (see the penultimate section of Satoshi’s paper for details).

Dishonest nodes that seek to replace some well-established, agreed-upon history with another will thus find it effectively impossible (i.e. the probability is negligible) unless they control a majority of the network’s capability – at which point they no longer constitute a small number of peers.

Summary

So in the language of the originally-stated BGP: Satoshi claimed that the decentralized timestamp server is an instance of the Byzantine Generals Problem, and that Nakamoto Consensus (as it came to be known) is a solution to the Byzantine Generals Problem. Because Nakamoto Consensus solves the BGP, honest nodes that always use the longest proof-of-work history in the decentralized timestamp network will eventually come to consensus on the ordering of blocks.

Recursive Stochastic Processes

2017-03-01T00:00:00+04:00

Last week Dan Peebles asked me on Twitter if I knew of any writing on the use of recursion schemes for expressing stochastic processes or other probability distributions. And I don’t! So I’ll write some of what I do know myself.

There are a number of popular statistical models or stochastic processes that have an overtly recursive structure, and when one has some recursive structure lying around, the elegant way to represent it is by way of a recursion scheme. In the case of stochastic processes, this typically boils down to using an anamorphism to drive things. Or, if you actually want to be able to observe the thing (note: you do), an apomorphism.

By representing a stochastic process in this way one can really isolate the probabilistic phenomena involved in it. One bundles up the essence of a process in a coalgebra, and then drives it via some appropriate recursion scheme.

Let’s take a look at three stochastic processes and examine their probabilistic and recursive structures.

Foundations

To start, I’m going to construct a simple embedded language in the spirit of the ones used in my simple probabilistic programming and comonadic inference posts. Check those posts out if this stuff looks too unfamiliar. Here’s a preamble that constitutes the skeleton of the code we’ll be working with.

{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TypeFamilies #-}

import Control.Monad
import Control.Monad.Free
import qualified Control.Monad.Trans.Free as TF
import Data.Functor.Foldable
import Data.Random (RVar, sample)
import qualified Data.Random.Distribution.Bernoulli as RF
import qualified Data.Random.Distribution.Beta as RF
import qualified Data.Random.Distribution.Normal as RF

-- probabilistic instruction set, program definitions

data ModelF a r =
    BernoulliF Double (Bool -> r)
  | GaussianF Double Double (Double -> r)
  | BetaF Double Double (Double -> r)
  | DiracF a
  deriving Functor

type Program a = Free (ModelF a)

type Model b = forall a. Program a b

type Terminating a = Program a a

-- core language terms

bernoulli :: Double -> Model Bool
bernoulli p = liftF (BernoulliF vp id) where
  vp
    | p < 0 = 0
    | p > 1 = 1
    | otherwise = p

gaussian :: Double -> Double -> Model Double
gaussian m s
  | s <= 0    = error "gaussian: variance out of bounds"
  | otherwise = liftF (GaussianF m s id)

beta :: Double -> Double -> Model Double
beta a b
  | a <= 0 || b <= 0 = error "beta: parameter out of bounds"
  | otherwise        = liftF (BetaF a b id)

dirac :: a -> Program a b
dirac x = liftF (DiracF x)

-- interpreter

rvar :: Program a a -> RVar a
rvar = iterM $ \case
  BernoulliF p f  -> RF.bernoulli p >>= f
  GaussianF m s f -> RF.normal m s >>= f
  BetaF a b f     -> RF.beta a b >>= f
  DiracF x        -> return x

-- utilities

free :: Functor f => Fix f -> Free f a
free = cata Free

affine :: Num a => a -> a -> a -> a
affine translation scale = (+ translation) . (* scale)

Just as a quick review, we’ve got:

A probabilistic instruction set defined by ‘ModelF’. Each constructor represents a foundational probability distribution that we can use in our embedded programs.
Three types corresponding to probabilistic programs. The ‘Program’ type simply wraps our instruction set up in a naïve free monad. The ‘Model’ type denotes probabilistic programs that may not necessarily terminate (in some weak sense), while the ‘Terminating’ type denotes probabilistic programs that terminate (ditto).
A bunch of embedded language terms. These are just probability distributions; here we’ll manage with the Bernouli, Gaussian, and beta distributions. We also have a ‘dirac’ term for constructing a Dirac distribution at a point.
A single interpeter ‘rvar’ that interprets a probabilistic program into a random variable (where the ‘RVar’ type is provided by random-fu). Typically I use mwc-probability for this but random-fu is quite nice. When a program has been interpreted into a random variable we can use ‘sample’ to sample from it.

So: we can write simple probabilistic programs in standard monadic fashion, like so:

betaBernoulli :: Double -> Double -> Model Bool
betaBernoulli a b = do
  p <- beta a b
  bernoulli p

and then interpret them as needed:

> replicateM 10 (sample (rvar (betaBernoulli 1 8)))
[False,False,False,False,False,False,False,True,True,False]

The Geometric Distribution

The geometric distribution is not a stochastic process per se, but it can be represented by one. If we repeatedly flip a coin and then count the number of flips until the first head, and then consider the probability distribution over that count, voilà. That’s the geometric distribution. You might see a head right away, or you might be infinitely unlucky and never see a head. So the distribution is supported over the entirety of the natural numbers.

For illustration, we can encode the coin flipping process in a straightforward recursive manner:

simpleGeometric :: Double -> Terminating Int
simpleGeometric p = loop 1 where
  loop n = do
    accept <- bernoulli p
    if   accept
    then dirac n
    else loop (n + 1)

We start flipping Bernoulli-distributed coins, and if we observe a head we stop and return the number of coins flipped thus far. Otherwise we keep flipping.

The underlying probabilistic phenomena here are the Bernoulli draw, which determines if we’ll terminate, and the dependent Dirac return, which will wrap a terminating value in a point mass. The recursive procedure itself has the pattern of:

If some condition is met, abort the recursion and return a value.
Otherwise, keep recursing.

This pattern describes an apomorphism, and the recursion-schemes type signature of ‘apo’ is:

apo :: Corecursive t => (a -> Base t (Either t a)) -> a -> t

It takes a coalgebra that returns an ‘Either’ value wrapped up in a base functor, and uses that coalgebra to drive the recursion. A ‘Left’-returned value halts the recursion, while a ‘Right’-returned value keeps it going.

Don’t be put off by the type of the coalgebra if you’re unfamiliar with apomorphisms - its bark is worse than its bite. Check out my older post on apomorphisms for a brief introduction to them.

With reference to the ‘apo’ type signature, The main thing to choose here is the recursive type that we’ll use to wrap up the ‘ModelF’ base functor. ‘Fix’ might be conceivably simpler to start, so I’ll begin with that. The coalgebra defining the model looks like this:

geoCoalg p n = BernoulliF p (\accept ->
  if   accept
  then Left (Fix (DiracF n))
  else Right (n + 1))

Then given the coalgebra, we can just wrap it up in ‘apo’ to represent the geometric distribution.

geometric :: Double -> Terminating Int
geometric p = free (apo (geoCoalg p) 1)

Since the geometric distribution (weakly) terminates, the program has return type ‘Terminating Int’.

Since we’ve encoded the coalgebra using ‘Fix’, we have to explicitly convert to ‘Free’ via the ‘free’ utility function I defined in the preamble. Recent versions of recursion-schemes have added a ‘Corecursive’ instance for ‘Free’, though, so the superior alternative is to just use that:

geometric :: Double -> Terminating Int
geometric p = apo coalg 1 where
  coalg n = TF.Free (BernoulliF p (\accept ->
    if   accept
    then Left (dirac n)
    else Right (n + 1)))

The point of all this is that we can isolate the core probabilistic phenomena of the recursive process by factoring it out into a coalgebra. The recursion itself takes the form of an apomorphism, which knows nothing about probability or flipping coins or what have you - it just knows how to recurse, or stop.

For illustration, here’s a histogram of samples drawn from the geometric via:

> replicateM 100 (sample (rvar (geometric 0.2)))

An Autoregressive Process

Autoregressive (AR) processes simply use a previous epoch’s output as the current epoch’s input; the number of previous epochs used as input on any given epoch is called the order of the process. An AR(1) process looks like this, for example:

\[y_t = \alpha + \beta y_{t - 1} + \epsilon_t\]

Here $\epsilon_t$ are independent and identically-distributed random variables that follow some error distribution. In other words, in this model the value $\alpha + \beta y_{t - 1}$ follows some probability distribution given the last epoch’s output $y_{t - 1}$ and some parameters $\alpha$ and $\beta$.

An autoregressive process doesn’t have any notion of termination built into it, so the purest way to represent one is via an anamorphism. We’ll focus on AR(1) processes in this example:

ar1 :: Double -> Double -> Double -> Double -> Model Double
ar1 a b s = ana coalg where
  coalg x = TF.Free (GaussianF (affine a b x) s (affine a b))

Each epoch is just a Gaussian-distributed affine transformation of the previous epochs’s output. But the problem with using an anamorphism here is that it will just shoot off to infinity, recursing endlessly. This doesn’t do us a ton of good if we want to actually observe the process, so if we want to do that we’ll need to bake in our own conditions for termination. Again we’ll rely on an apomorphism for this; we can just specify how many periods we want to observe the process for, and stop recursing as soon as we exceed that.

There are two ways to do this. We can either get a view of the process at $n$ periods in the future, or we can get a view of the process over $n$ periods in the future. I’ll write both, for illustration. The coalgebra for the first is simpler, and looks like:

arCoalg (n, x) = TF.Free (GaussianF (affine a b x) s (\y ->
    if   n <= 0
    then Left (dirac x)
    else Right (pred m, y)))

The coalgebra is saying:

Given $x$, let $z$ have a Gaussian distribution with mean $\alpha + \beta x$ and standard deviation $s$.
If we’re on the last epoch, return $x$ as a Dirac point mass.
Otherwise, continue recursing with $z$ as input to the next epoch.

Now, to observe the process over the next $n$ periods we can just collect the observations we’ve seen so far in a list. An implementation of the process, apomorphism and all, looks like this:

ar :: Int -> Double -> Double -> Double -> Double -> Terminating [Double]
ar n a b s origin = apo coalg (n, [origin]) where
  coalg (epochs, history@(x:_)) =
    TF.Free (GaussianF (affine a b x) s (\y ->
      if   epochs <= 0
      then Left (dirac (reverse history))
      else Right (pred epochs, y:history)))

(Note that I’m deliberately not handling the error condition here so as to focus on the essence of the coalgebra.)

We can generate some traces for it in the standard way. Here’s how we’d sample a 100-long trace from an AR(1) process originating at 0 with $\alpha = 0$, $\beta = 1$, and $s = 1$:

> sample (rvar (ar 100 0 1 1 0))

and here’s a visualization of 10 of those traces:

The Stick-Breaking Process

The stick breaking process is one of any number of whimsical stochastic processes used as prior distributions in nonparametric Bayesian models. The idea here is that we want to take a stick and endlessly break it into smaller and smaller pieces. Every time we break a stick, we recursively take the rest of the stick and break it again, ad infinitum.

Again, if we wanted to represent this endless process very faithfully, we’d use an anamorphism to drive it. But in practice we’re going to only want to break a stick some finite number of times, so we’ll follow the same pattern as the AR process and use an apomorphism to do that:

sbp :: Int -> Double -> Terminating [Double]
sbp n a = apo coalg (n, 1, []) where
  coalg (epochs, stick, sticks) = TF.Free (BetaF 1 a (\p ->
    if   epochs <= 0
    then Left (dirac (reverse (stick : sticks)))
    else Right (pred epochs, (1 - p) * stick, (p * stick):sticks)))

The coalgebra that defines the process says the following:

Let the location $p$ of the break on the next (normalized) stick be beta$(1, \alpha)$-distributed.
If we’re on the last epoch, return all the pieces of the stick that we broke as a Dirac point mass.
Otherwise, break the stick again and recurse.

Here’s a plot of five separate draws from a stick breaking process with $\alpha = 0.2$, each one observed for five breaks. Note that each draw encodes a categorical distribution over the set $\{1, \ldots, 6\}$; the stick breaking process is a ‘distribution over distributions’ in that sense:

The stick breaking process is useful for developing mixture models with an unknown number of components, for example. The $\alpha$ parameter can be tweaked to concentrate or disperse probability mass as needed.

Conclusion

This seems like enough for now. I’d be interested in exploring other models generated by recursive processes just to see how they can be encoded, exactly. Basically all of Bayesian nonparametrics is based on using recursive processses as prior distributions, so the Dirichlet process, Chinese Restaurant Process, Indian Buffet Process, etc. should work beautifully in this setting.

Fun fact: back in 2011 before ~~neural networks~~ deep learning had taken over machine learning, Bayesian nonparametrics was probably the hottest research area in town. I used to joke that I’d create a new prior called the Malaysian Takeaway Process for some esoteric nonparametric model and thus achieve machine learning fame, but never did get around to that.

Addendum

I got a question about how I produce these plots. And the answer is the only sane way when it comes to visualization in Haskell: dump the output to disk and plot it with something else. I use R for most of my interactive/exploratory data science-fiddling, as well as for visualization. Python with matplotlib is obviously a good choice too.

Here’s how I made the autoregressive process plot, for example. First, I just produced the actual samples in GHCi:

> samples <- replicateM 10 (sample (rvar (ar 100 0 1 1 0)))

Then I wrote them to disk:

> let render handle = hPutStrLn handle . filter (`notElem` "[]") . show
> withFile "trace.dat" WriteMode (\handle -> mapM_ (render handle) samples)

The following R script will then get you the plot:

require(ggplot2)
require(reshape2)

raw = read.csv('trace.dat', header = F)

d = data.frame(t(raw), x = seq_along(raw))
m = melt(d, id.vars = 'x')

ggplot(m, aes(x, value, colour = variable)) + geom_line()

I used ggplot2 for the other plots as well; check out the ggplot2 functions geom_histogram, geom_bar, and facet_grid in particular.

The Applicative Structure of the Giry Monad

2017-02-26T00:00:00+04:00

In my last two posts about the Giry monad I derived the thing from its categorical and measure-theoretic foundations. I kind of thought that those posts wouldn’t be of much interest to people but they turned out to be a hit. I clearly can’t tell what the internet likes.

Anyway, something I left out was the theoretical foundation of the Giry monad’s Applicative instance, which seemed a bit odd. I also pointed out that applicativeness in the context of probability implies independence between probability measures.

In this article I’m going to look at each of these issues. After playing around with the foundations, it looks like the applicative instance for the Giry monad can be put on a sound footing in terms of the standard measure-theoretic concept of product measure. Also, it turns out that the claim I made of applicativeness $\implies$ independence is somewhat ill-posed. But, using the shiny new intuition we’ll glean from a better understanding of the applicative structure, we can put that on a solid footing too.

So let’s look at both of these things and see what’s up.

Monoidal Functors

The foundational categorical concept behind applicative functors is the monoidal functor, which is a functor between monoidal categories that preserves monoidal structure.

Formally: for monoidal categories $(C, \otimes, I)$ and $(D, \oplus, J)$, a monoidal functor $F : C \to D$ is a functor and associated natural transformations $\phi : F(A) \oplus F(B) \to F(A \otimes B)$ and $i : J \to F(I)$ that satisfy some coherence conditions that I won’t mention here. Notably, if $\phi$ and $i$ are isomorphisms (i.e. are invertible) then $F$ is called a strong monoidal functor. Otherwise it’s called lax. Applicative functors in particular are lax monoidal functors.

This can be made much clearer for endofunctors on a monoidal category $(C, \otimes, I)$. Then you only have $F : C \to C$ and $\phi : F(A) \otimes F(B) \to F(A \otimes B)$ to worry about. If we sub in the Giry monad $\mathcal{P}$ from the last couple of posts, we’d want $\mathcal{P} : \textbf{Meas} \to \textbf{Meas}$ and $\phi : \mathcal{P}(M) \otimes \mathcal{P}(N) \to \mathcal{P}(M \otimes N)$.

Does the category of measurable spaces $\textbf{Meas}$ have a monoidal structure? Yup. Take measurable spaces $M = (X, \mathcal{X})$ and $N = (Y, \mathcal{Y})$. From the Giry monad derivation we already have that the monoidal identity $i : M \to \mathcal{P}(M)$ corresponds to a Dirac measure at a point, so that’s well and good. And we can define the tensor product $\otimes$ between $M$ and $N$ as follows: let $X \times Y$ be the standard Cartesian product on $X$ and $Y$ and let $\mathcal{X} \otimes \mathcal{Y}$ be the smallest $\sigma$-algebra generated by the Cartesian product $A \times B$ of measurable sets $A \in \mathcal{X}$ and $B \in \mathcal{Y}$. Then $(X \times Y, \mathcal{X} \otimes \mathcal{Y})$ is a measurable space, and so $(\textbf{Meas}, \otimes, i)$ is monoidal.

Recall that $\mathcal{P}(M)$ and $\mathcal{P}(N)$ - the space of measures over $M$ and $N$ respectively - are themselves objects in $\textbf{Meas}$. So, clearly $\mathcal{P}(M) \otimes \mathcal{P}(N)$ is a measurable space, and if $\mathcal{P}$ is monoidal then there must exist a natural transformation that can take us from there to $\mathcal{P}(M \otimes N)$. This is the space of measures over the product $M \otimes N$.

So the question is: does $\mathcal{P}$ have the required monoidal structure?

Yes. It must, since $\mathcal{P}$ is a monad, and any monad can generate the required natural transformation. Let $\mu$ be the monadic ‘join’ operator $\mathcal{P}^2 \to \mathcal{P}$ and $\eta$ be the monadic identity $I \to \mathcal{P}$. We have, evaluating right-to-left:

\[\phi_{\nu \times \rho} = \mu \mathcal{P} \left\{ \lambda m . \mu \mathcal{P}\left(\lambda n. \eta_{m \times n}\right)\mathcal{P}(\rho) \right\} \mathcal{P}(\nu).\]

Using $\gg\!\!=$ makes this much easier to read:

\[\phi_{\nu \times \rho} = \nu \gg\!\!= \lambda m. \rho \gg\!\!= \lambda n. \eta_{m \times n}\]

or in code, just:

phi :: Monad m => (m a, m b) -> m (a, b)
phi (m, n) = liftM2 (,) m n

So with that we have that $(\mathcal{P}, \phi, i)$ is a (lax) monoidal functor. And you can glean a monad-generated applicative operator from that immediately (this leads to the function called ‘ap’ in Control.Monad):

ap :: Monad m => m (a -> b) -> m a -> m b
ap f x = fmap (\(g, z) -> g z) (phi f x)

(Note: I won’t refer to $\mu$ as the join operator from this point out in order to free it up for denoting measures.)

Probabilistic Interpretations

Product Measure

The correct probabilistic interpretation here is that $\phi$ takes a pair of probability measures to the product measure over the appropriate product space. For probability measures $\mu$ and $\nu$ on measurable spaces $M$ and $N$ respectively, the product measure is the (unique) measure $\mu \times \nu$ on $M \otimes N$ such that:

\[(\mu \times \nu)(A \times B) = \mu(A) \nu(B)\]

for $A \times B$ a measurable set in $M \otimes N$.

Going through the monoidal functor route seems to put the notion of the Giry applicative instance on a more firm measure-theoretic foundation. Instead of considering the following from the Giry monad foundations article:

\[(\rho \, \langle \ast \rangle \, \nu)(f) = \int_{\mathcal{P}(M \to N)} \left\{\lambda T . \int_{M \to N} (f \circ T) d\nu \right\} d \rho\]

which is defined in terms of the dubious space of measures over measurable functions $M \to N$, we can better view things using the monoidal structure-preserving natural transformation $\phi$. For measures $\mu$ and $\nu$ on $(X, \mathcal{X})$ and $(Y, \mathcal{Y})$ respectively, we have:

\[\phi(\mu, \nu)(f) = \int_{X \times Y}f d(\mu \times \nu)\]

and then for $g : Z \to X \otimes Y$ we can use the functor structure of $\mathcal{P}$ to do:

\[(\text{fmap} \, g \, \phi(\mu, \nu))(f) = \int_{Z} (f \circ g) d((\mu \times \nu) \circ g^{-1})\]

which corresponds to a standard applicative expression g <$> mu <*> nu. I suspect there’s then some sort of Yoneda argument or something that makes currying and partial function application acceptable.

Independence

Now. What does this have to say about independence?

In particular, it’s too fast and loose to claim measures can be ‘independent’ at all. Independence is a property of measurable sets, measurable functions, and $\sigma$-algebras. Not of measures! But there is a really useful connection, so let’s illustrate that.

First, let’s define independence formally as follows. Take a probability space $(X, \mathcal{X}, \mathbb{P})$. Then any measurable sets $A$ and $B$ in $\mathcal{X}$ are independent if

\[\mathbb{P}(A \cap B) = \mathbb{P}(A)\mathbb{P}(B).\]

That’s the simplest notion.

Next, consider two sub-$\sigma$-algebras $\mathcal{A}$ and $\mathcal{B}$ of $\mathcal{X}$ (a sub-$\sigma$-algebra is just a a subset of a $\sigma$-algebra that itself happens to be a $\sigma$ algebra). Then $\mathcal{A}$ and $\mathcal{B}$ are independent if, for any $A$ in $\mathcal{A}$ and any $B$ in $\mathcal{B}$, we have that $A$ and $B$ are independent.

The final example is independence of measurable functions. Take measurable functions $f$ and $g$ both from $X$ to the real numbers equipped with some appropriate $\sigma$-algebra $\mathcal{B}$. Then each of these functions generates a sub-$\sigma$ algebra of $\mathcal{X}$ as follows:

\[\begin{align*} \mathcal{X}_{f} & = \{ f^{-1}(B) : B \in \mathcal{B} \} \\ \mathcal{X}_{g} & = \{ g^{-1}(B) : B \in \mathcal{B} \}. \end{align*}\]

Then $f$ and $g$ are independent if the generated $\sigma$-algebras $\mathcal{X}_{f}$ and $\mathcal{X}_{g}$ are independent.

Note that in every case independence is defined in terms of a single measure, $\mathbb{P}$. We can’t talk about different measures being independent. To paraphrase Terry Tao here:

The notion of independence between [measurable functions] does not make sense if the [measurable functions] are being modeled by separate probability spaces; they have to be coupled together into a single probability space before independence becomes a meaningful notion.

To be pedantic and explicitly specify the measure by which some things are independent, some authors state that measurable functions $f$ and $g$ are $\mathbb{P}$-independent, for example.

We can see a connection to independence when we look at convolution and associated operators. Recall that for measures $\mu$ and $\nu$ on the same measurable space $M = (X, \mathcal{X})$ that supports some notion of addition, their convolution looks like:

\[(\mu + \nu)(f) = \int_{X}\int_{X} f(x + y) d\mu(x) d\nu(y).\]

The probabilistic interpretation here (see Terry Tao on this too) is that $\mu + \nu$ is the measure corresponding to the sum of independent measurable functions $g$ and $h$ with corresponding measures $\mu$ and $\nu$ respectively.

That looks weird though, since we clearly defined independence between measurable functions using a single probability measure. How is it we can talk about independent measurable functions $g$ and $h$ having different corresponding measures?

We first need to couple everything together into a single probability space as per Terry’s quote. Complete $M$ with some abstract probability measure $\mathbb{P}$ to form the probability space $(X, \mathcal{X}, \mathbb{P})$. Now we have $g$ and $h$ measurable functions from $X$ to $\mathbb{R}$.

To say that $g$ and $h$ are independent is to say that their generated $\sigma$-algebras are $\mathbb{P}$-independent. And the measures that they correspond to are the pushforwards of $\mathbb{P}$ under $g$ and $h$ respectively. So, $\mu = \mathbb{P} \circ g^{-1}$ and $\nu = \mathbb{P} \circ h^{-1}$. The result is that the measurable functions correspond to different (pushforward) measures $\mu$ and $\nu$, but are independent with respect to the same underlying probability measure $\mathbb{P}$.

The monoidal structure of $\mathcal{P}$ then gets us to convolution. Given a product of measures $\mu$ and $\nu$ each on $(X, \mathcal{X})$ we can immediately retrieve their product measure $\mu \times \nu$ via $\phi$. And from there we can get to $\mu + \nu$ via the functor structure of $\mathcal{P}$ - we just find the pushforward of $\mu \times \nu$ with respect to a function $\rho$ that collapses a product via addition. So $\rho : X \times X \to \mathbb{R}$ is defined as:

\[\rho(a \times b) = a + b\]

and then the convolution $\mu + \nu$ is thus:

\[\mu + \nu = (\mu \times \nu) \circ \rho^{-1}.\]

Other operations can be defined similarly, e.g. for $\sigma(a \times b) = a - b$ we get:

\[\mu - \nu = (\mu \times \nu) \circ \sigma^{-1}.\]

The crux of all this is whenever we apply a measurable function to a product measure, we can always extract notions of independent measurable functions from the result. And the measures corresponding to those independent measurable functions will be the components of the product measure respectively.

This is super useful and lets one claim something stronger than what the monadic structure gives you. In an expression like g <$> mu <*> nu <*> rho, you are guaranteed that the corresponding random variables $g_\mu$, $g_\nu$, $g_\rho$ (for suitable projections) are independent. The same cannot be said if you use the monadic structure to do something like g mu nu rho where the product structure is not enforced - in that case you’re not guaranteed anything of the sort. This is why the applicative structure is useful for encoding independence in a way that the monadic structure is not.

Conclusion

So there you have it. Applicativeness can seemingly be put on a straightforward measure-theoretic grounding and has some useful implications for independence.

It’s worth noting that, in the case of the Giry monad, we don’t need to go through its monadic structure in order to recover an applicative instance. We can do so entirely by hacking together continuations without using a single monadic bind. This is actually how I defined the applicative instance in the Giry monad implementation article previously:

instance Applicative Measure where
  pure x = Measure (\f -> f x)
  Measure g <*> Measure h = Measure $ \f ->
    g (\k -> h (f . k))

Teasing out the exact structure of this and its relation to the codensity monad is again something I’ll leave to others.

Implementing the Giry Monad

2017-02-13T00:00:00+04:00

In my last post I went over the categorical and measure-theoretic foundations of the Giry monad, the ‘canonical’ probability monad that operates on the level of probability measures.

In this post I’ll pick up from where I left off and talk about a neat and faithful (if impractical) implementation of the Giry monad that one can put together in Haskell.

Measure, Integral, and Continuation

So. For a quick review, we’ve established the Giry monad as a triple $(\mathcal{P}, \mu, \eta)$, where $\mathcal{P}$ is an endofunctor on the category of measurable spaces $\textbf{Meas}$, $\mu$ is a marginalizing integration operation defined by:

\[\mu(\rho)(A) = \int_{\mathcal{P}(M)} \left\{\lambda \nu . \int_M \chi_A d \nu \right\} d \rho\]

and $\eta$ is a monoidal identity, defined by the Dirac measure at a point:

\[\eta(x)(A) = \chi_A(x).\]

How do we actually implement this beast? If we’re looking to be suitably general then it is unlikely that we’re going to be able to easily represent something like a $\sigma$-algebra over some space of measures on a computer, so that route is sort of a non-starter.

But it can be done. The key to implementing a general-purpose Giry monad is to notice that the fundamental operation involved in it is integration, and that we can avoid working with $\sigma$-algebras and measurable spaces directly if we focus on dealing with measurable functions instead of measurable sets.

Consider the integration map on measurable functions $\tau_f$ that we’ve been using this whole time. For some measurable function $f$, $\tau_f$ takes a measure on some measurable space $M = (X, \mathcal{X})$ and uses it to integrate $f$ over $X$. In other words:

\[\tau_f(\nu) = \int_X f d\nu.\]

A measure in $\mathcal{P}(M)$ has type $X \to \mathbb{R}$, so $\tau_f$ has corresponding type $(X \to \mathbb{R}) \to \mathbb{R}$.

This might look familiar to you; it’s very similar to the type signature for a continuation:

newtype Cont a r = Cont ((a -> r) -> r)

Indeed, if we restrict the carrier type of ‘Cont’ to the reals, we can be really faithful to the type:

newtype Integral a = Integral ((a -> Double) -> Double)

Now, let’s overload notation and call the integration map $\tau_f$ itself a measure. That is, $\tau_f$ is a mapping $\nu \mapsto \int_{X}fd\nu$, so we’ll just interpret the notation $\nu(f)$ to mean the same thing - $\int_{X}fd\nu$. This is convenient because we can dispense with $\tau$ and just pretend measures can be applied directly to measurable functions. There’s no way we can get confused here; measures operate on sets, not functions, so notation like $\nu(f)$ is not currently in use. We just set $\nu(f) = \tau_f(\nu)$ and that’s that. Let’s rename the ‘Integral’ type to match:

newtype Measure a = Measure ((a -> Double) -> Double)

We can extract a very nice shallowly-embedded language for integration here, the core of which is a single term:

integrate :: (a -> Double) -> Measure a -> Double
integrate f (Measure nu) = nu f

Note that this is the same way we’d express integration mathematically; we specify that we want to integrate a measurable function $f$ with respect to some measure $\nu$:

\[\int f d\nu = \texttt{integrate f nu}.\]

The only subtle difference here is that we don’t specify the space we’re integrating over in the integral expression - instead, we’ll bake that into the definition of the measures we create themselves. Details in a bit.

What’s interesting here is that the Giry monad is the continuation monad with the carrier type restricted to the reals. This isn’t surprising when you think about what’s going on here - we’re representing measures as integration procedures, that is, programs that take a measurable function as input and then compute its integral in some particular way. A measure, as we’ve implemented it here, is just a ‘program with a missing piece’. And this is exactly the essence of the continuation monad in Haskell.

Typeclass Instances

We can fill out the functor, applicative, and monad instances mechanically by reference to the a standard continuation monad implementation, and each instance gives us some familiar conceptual structure or operation on probability measures. Let’s take a look.

The functor instance lets us transform the support of a measurable space while keeping its density structure invariant. If we have:

\[\nu(f) = \int_X f d\nu\]

then mapping a measurable function over the measure corresponds to:

\[(\texttt{fmap} \, g \, \nu)(f) = \int_{X} (f \circ g) d\nu.\]

The functor structure allows us to precisely express a pushforward measure or distribution of $\nu$ under $g$. It lets us ‘adapt’ a measure to other measurable spaces, just like a good functor should.

In Haskell, the functor instance corresponds exactly to the math:

instance Functor Measure where
  fmap g nu = Measure $ \f ->
    integrate (f . g) nu

The monad instance is exactly the Giry monad structure that we developed previously, and it allows us to sequence probability measures together by marginalizing one into another. We’ll write it in terms of bind, of course, which went like:

\[(\rho \gg\!\!= g)(f) = \int_M \left\{\lambda m . \int_N f dg(m) \right\} d \rho.\]

The Haskell translation is verbatim:

instance Monad Measure where
  return x  = Measure (\f -> f x)
  rho >>= g = Measure $ \f ->
    integrate (\m -> integrate f (g m)) rho

Finally there’s the Applicative instance, which as I mentioned in the last post is sort of conceptually weird here. So in the spirit of that comment, I’m going to dodge any formal justification for now and just use the following instance which works in practice:

instance Applicative Measure where
  pure x = Measure (\f -> f x)
  Measure g <*> Measure h = Measure $ \f ->
    g (\k -> h (f . k))

Conceptual Example

It’s worth taking a look at an example of how things should conceivably work here. Consider the following probabilistic model:

\[\begin{align*} \pi & \sim \text{beta}(\alpha, \beta) \\ \mu \, | \, \pi & \sim \text{binomial}(n, \pi) \end{align*}\]

It’s a standard hierarchical presentation. A ‘compound’ measure can be obtained here by marginalizing over the beta measure $\pi$, and that’s called the beta-binomial measure. Let’s find it.

The beta distribution has support on the $[0, 1]$ subset of the reals, and the binomial distribution with argument $n$ has support on the $\{0, \ldots, n\}$ subset of the integers, so we know that things should proceed like so:

\[\begin{align*} \psi(f) & = (\pi \gg\!\!= \mu)(f) \\ & = \int_{\mathbb{R}} \left\{\lambda p . \int_{\mathbb{Z}} f d\mu(p) \right\} d \pi. \end{align*}\]

Eliding some theory of integration, I can tell you that $\pi$ is absolutely continuous with respect to Lebesgue measure and that $\mu(p)$ is absolutely continuous w/respect to counting measure for appropriate $p$. So, $\pi$ admits a density $d\pi/dx = g_\pi$ and $\mu(p)$ admits a density $d\mu(p)/d\# = g_{\mu(p)}$, defined as:

\[g_\pi(p \, | \, \alpha, \beta) = \frac{1}{B(\alpha, \beta)} p^{\alpha - 1} (1 - p)^{\beta - 1}\]

and

\[g_{\mu(p)}(x \, | \, n, p) = \binom{n}{x} p^x (1 - p)^{n - x}\]

respectively, for $B$ the beta function and $\binom{n}{x}$ a binomial coefficient. Again, we can reduce the integral as follows, transforming the outermost integral into a standard Riemann integral and the innermost integral into a simple sum of products:

\[\psi(f) = \int_{0}^{1} \lambda p. \left\{ \lambda \alpha. \lambda \beta. g_{\pi}(p \, | \alpha, \beta) \sum_{z \in \{0, \ldots, n\}} f(z) \left( \lambda n. g_{\mu(p)}(z \, | \, n, p) \right) \right\} dx.\]

where $dx$ denotes Lebesgue measure. I could expand this further or simplify things a little more (the beta and binomial are conjugates) but you get the point, which is that we have a way to evaluate the integral.

What is really required here then is to be able to encode into the definitions of measures like $\pi$ and $\mu(p)$ the method of integration to use when evaluating them. For measures absolutely continuous w/respect to Lebesgue measure, we can use the Riemann integral over the reals. For measures absolutely continuous w/respect to counting measure, we can use a sum of products. In both cases, we’ll also need to supply the density or mass function by which the integral should be evaluated.

Creating Measures

Recall that we are representing measures as integration procedures. So to create one is to define a program by which we’ll perform integration.

Let’s start with the conceptually simpler case of a probability measure that’s absolutely continuous with respect to counting measure. We need to provide a support (the region for which probability is greater than 0) and a probability mass function (so that we can weight every point appropriately). Then we just want to integrate a function by evaluating it at every point in the support, multiplying the result by that point’s probability mass, and summing everything up. In code, this translates to:

fromMassFunction :: (a -> Double) -> [a] -> Measure a
fromMassFunction f support = Measure $ \g ->
  foldl' (\acc x -> acc + f x * g x) 0 support

So if we want to construct a binomial measure, we can do that like so (where choose comes from Numeric.SpecFunctions):

binomial :: Int -> Double -> Measure Int
binomial n p = fromMassFunction (pmf n p) [0..n] where
  pmf n p x
    | x < 0 || n < x = 0
    | otherwise = choose n x * p ^^ x * (1 - p) ^^ (n - x)

The second example involves measures over the real line that are absolutely continuous with respect to Lebesgue measure. In this case we want to evaluate a Riemann integral over the entire real line, which is going to necessitate approximation on our part. There are a bunch of methods out there for approximating integrals, but a simple one for one-dimensional problems like this is quadrature, an implementation for which Ed Kmett has handily packaged up in his integration package:

fromDensityFunction :: (Double -> Double) -> Measure Double
fromDensityFunction d = Measure $ \f ->
    quadratureTanhSinh (\x -> f x * d x)
  where
    quadratureTanhSinh = result . last . everywhere trap

Here we’re using quadrature to approximate the integral, but otherwise it has a similar form as ‘fromMassFunction’. The difference here is that we’re integrating over the entire real line, and so don’t have to supply a support explicitly.

We can use this to create a beta measure (where logBeta again comes from Numeric.SpecFunctions):

beta :: Double -> Double -> Measure Double
beta a b = fromDensityFunction (density a b) where
  density a b p
    | p < 0 || p > 1 = 0
    | otherwise = 1 / exp (logBeta a b) * p ** (a - 1) * (1 - p) ** (b - 1)

Note that since we’re going to be integrating over the entire real line and the beta distribution has support only over $[0, 1]$, we need to implicitly define the support here by specifying which regions of the domain will lead to a density of 0.

In any case, now that we’ve constructed those things we can just use a monadic bind to create the beta-binomial measure we described before. It masks a lot of under-the-hood complexity.

betaBinomial :: Int -> Double -> Double -> Measure Int
betaBinomial n a b = beta a b >>= binomial n

There are a couple of other useful ways to create measures, but the most notable is to use a sample in order to create an empirical measure. This is equivalent to passing in a specific support for which the mass function assigns equal probability to every element; I’ll use Gabriel Gonzalez’s foldl package here as it’s pretty elegant:

fromSample :: Foldable f => f a -> Measure a
fromSample = Measure . flip weightedAverage

weightedAverage :: (Foldable f, Fractional r) => (a -> r) -> f a -> r
weightedAverage f = Foldl.fold (weightedAverageFold f) where
  weightedAverageFold :: Fractional r => (a -> r) -> Fold a r
  weightedAverageFold f = Foldl.premap f averageFold

  averageFold :: Fractional a => Fold a a
  averageFold = (/) <$> Foldl.sum <*> Foldl.genericLength

Using ‘fromSample’ you can create an empirical measure using just about anything you’d like:

data Foo = Foo | Bar | Baz

foos :: [Foo]
foos = [Foo, Foo, Bar, Foo, Baz, Foo, Bar, Foo, Foo, Foo, Bar]

nu :: Measure Foo
nu = fromSample foos

Though I won’t demonstrate it here, you can use this approach to also create measures from sampling functions or random variables that use a source of randomness - just draw a sample from the function and pipe the result into ‘fromSample’.

Querying Measures

To query a measure is to simply get some result out of it, and we do that by integrating some measurable function against it. The easiest thing to do is to just take a straightforward expectation by integrating the identity function; for example, here’s the expected value of a beta(10, 10) measure:

> integrate id (beta 10 10)
0.49999999999501316

The expected value of a beta($\alpha$, $\beta$) distribution is $\alpha / (\alpha + \beta)$, so we can verify analytically that the result should be 0.5. We observe a bit of numerical imprecision here because, if you’ll recall, we’re just approximating the integral via quadrature. For measures created via ‘fromMassFunction’ we don’t need to use quadrature, so we won’t observe the same kind of approximation error. Here’s the expected value of a binomial(10, 0.5) measure, for example:

> integrate fromIntegral (binomial 10 0.5)
5.0

Note here that we’re integrating the ‘fromIntegral’ function against the binomial measure. This is because the binomial measure is defined over the integers, rather than the reals, and we always need to evaluate to a real when we integrate. That’s part of the definition of a measure!

Let’s calculate the expectation of the beta-binomial distribution with $n = 10$, $\alpha = 1$, and $\beta = 8$:

> integrate fromIntegral (betaBinomial 10 1 8)
1.108635884924813

Neato. And since we can integrate like this, we can really compute any of the moments of a measure. The first raw moment is what we’ve been doing here, and is called the expectation:

expectation :: Measure Double -> Double
expectation = integrate id

The second (central) moment is the variance. Here I mean variance in the moment-based sense, rather than as the possibly better-known sample variance:

variance :: Measure Double -> Double
variance nu = integrate (^ 2) nu - expectation nu ^ 2

The variance of a binomial($n$, $p$) distribution is known to be $np(1-p)$, so for $n = 10$ and $p = 0.5$ we should get 2.5:

> variance (binomial 10 0.5)
<interactive>:87:11: error:
    • Couldn't match type ‘Int’ with ‘Double’
      Expected type: Measure Double
        Actual type: Measure Int
    • In the first argument of ‘variance’, namely ‘(binomial 10 0.5)’
      In the expression: variance (binomial 10 0.5)
      In an equation for ‘it’: it = variance (binomial 10 0.5)

Ahhh, but remember: the binomial measure is defined over the integers, so we can’t integrate it directly. No matter - the functorial structure allows us to adapt it to any other measurable space via a measurable function:

> variance (fmap fromIntegral (binomial 10 0.5))
2.5

Expectation and variance (and other moments) are pretty well-known, but you can do more exotic things as well. You can calculate the moment generating function for a measure, for example:

momentGeneratingFunction :: Measure Double -> Double -> Double
momentGeneratingFunction nu t = integrate (\x -> exp (t * x)) nu

and the cumulant generating function follows naturally:

cumulantGeneratingFunction :: Measure Double -> Double -> Double
cumulantGeneratingFunction nu = log . momentGeneratingFunction nu

A particularly useful construct is the cumulative distribution function for a measure, which calculates the probability of a region less than or equal to some number:

cdf :: Measure Double -> Double -> Double
cdf nu x = integrate (negativeInfinity `to` x) nu

negativeInfinity :: Double
negativeInfinity = negate (1 / 0)

to :: (Num a, Ord a) => a -> a -> a -> a
to a b x
  | x >= a && x <= b = 1
  | otherwise        = 0

The beta(2, 2) distribution is symmetric around its mean 0.5, so the probability of the region $[0, 0.5]$ should itself be 0.5. This checks out as expected, modulo approximation error due to quadrature:

> cdf (beta 2 2) 0.5
0.4951814897381374

Similarly for measurable spaces without any notion of order, there’s a simple CDF analogue that calculates the probability of a region that contains the given points:

containing :: (Num a, Eq b) => [b] -> b -> a
containing xs x
  | x `elem` xs = 1
  | otherwise   = 0

And probably the least interesting query of all is the simple ‘volume’, which calculates the total measure of a space. For any probability measure this must obviously be one, so it can at least be used as a quick sanity check:

volume :: Measure Double -> Double
volume = integrate (const 1)

Convolution and Friends

I mentioned in the last post that applicativeness corresponds to independence in some sense, and that independent measures over the same measurable space can be convolved together, à la:

\[(\nu + \zeta)(f) = \int_{M}\int_{M}f(x + y)d\nu(x)d\zeta(y)\]

for measures $\nu$ and $\zeta$ on $M$. In Haskell-land it’s well-known that any applicative instance gives you a free ‘Num’ instance, and the story is no different here:

instance Num a => Num (Measure a) where
  (+)         = liftA2 (+)
  (-)         = liftA2 (-)
  (*)         = liftA2 (*)
  abs         = fmap abs
  signum      = fmap signum
  fromInteger = pure . fromInteger

There are a few neat ways to demonstrate this kind of thing. Let’s use a Gaussian measure here as a running example:

gaussian :: Double -> Double -> Measure Double
gaussian m s = fromDensityFunction (density m s) where
  density m s x
    | s <= 0    = 0
    | otherwise =
        1 / (s * sqrt (2 * pi)) * exp (negate ((x - m) ^^ 2) / (2 * (s ^^ 2)))

First, consider a chi-squared measure with $k$ degrees of freedom. We could create this directly using a density function, but instead we can represent it by summing up independent squared Gaussian measures:

chisq :: Int -> Measure Double
chisq k = sum (replicate k normal) where
  normal = fmap (^ 2) (gaussian 0 1)

To sanity check the result, we can compute the mean and variance of a $\chi^2(2)$ measure, which should be $k$ and $2k$ respectively for $k = 2$:

> expectation (chisq 2)
2.0000000000000004
> variance (chisq 2)
4.0

As a second example, consider a product of independent Gaussian measures. This is a trickier distribution to deal with analytically (see here), but we can use some well-known identities for general independent measures in order to verify our results. For any independent measures $\mu$ and $\nu$, we have:

\[\mathbb{E}(\mu\nu) = \mathbb{E}\mu \mathbb{E}\nu\]

and

\[\text{var}(\mu\nu) = \text{var}(\mu)\text{var}(\nu) + \text{var}(\mu)(\mathbb{E}\nu)^2 + \text{var}(\nu)(\mathbb{E}\mu)^2\]

for the expectation and variance of their product. So for a product of independent Gaussians w/parameters (1, 2) and (2, 3) respectively, we expect to see 2 for its expectation and 61 for its variance:

> expectation (gaussian 1 2 * gaussian 2 3)
2.0000000000000001
> variance (gaussian 1 2 * gaussian 2 3)
61.00000000000003

Woop!

Wrapping Up

And there you have it, a continuation-based implementation of the Giry monad. You can find a bunch of code with similar functionality to this packaged up in my old measurable library on GitHub if you’d like to play around with the concepts.

That library has accumulated a few stars since I first pushed it up in 2013. I think a lot of people are curious about these weird measure things, and this framework at least gives you the ability to play around with a representation for measures directly. I found it particularly useful for really grokking, say, that integrating some function $f$ against a probability measure $\nu$ is identical to integrating the identity function against the probability measure $\texttt{fmap} \, f \, \nu$. And there are a few similar concepts there that I find really pop out when you play with measures directly, rather than when one just works with them on paper.

But let me now tell you why the Giry monad sucks in practice.

Take a look at this integral expression, which is brought about due to a monadic bind:

\[(\nu \gg\!\!= \mu)(f) = \int_{M} \left\{\lambda m . \int_{M} f d\mu(m) \right\} d \nu.\]

For simplicitly, let’s assume that $M$ is discrete and has cardinality $|M|$. This means that the integral reduces to

\[(\nu \gg\!\!= \mu)(f) = \underbrace{\sum_{m \in M} d\nu(m) \underbrace{ \sum_{n \in M} f(n) d\mu(m)(n) }_{O(|M|)}}_{O(|M|)}\]

for $d\mu(m)$ and $d\nu$ the appropriate Radon-Nikodym derivatives. You can see that the total number of operations involved in the integral is $O(|M|^2)$, and indeed, for $p$ monadic binds the computational complexity involved in evaluating all the integrals involved is exponential, on the order of $|M|^p$. It was no coincidence that I demonstrated a variance calculation for a $\chi^2(2)$ distribution instead of for a $\chi^2(10)$.

This isn’t really much of a surprise - the cottage industry of approximating integrals exists because integration is hard in practice, and integration is surely best avoided whenever one can get away with doing so. Vikash Mansinghka’s quote on this topic is fitting: “don’t calculate probabilities - sample good guesses.” I’ll also add: relegate the measures to measure theory, where they seem to belong.

The Giry monad is a lovely abstract construction for formalizing the monadic structure of probability, and as canonical probabilistic objects, measures and integrals are tremendously useful when working theoretically. But they’re a complete non-starter when it comes to getting anything nontrivial done in practice. For that, there are far more useful representations for probability distributions in Haskell - notably, the sampling function or random variable representation found in things like mwc-probability/mwc-random-monad and random-fu, or even better, the structural representation based on free or operational monads like I’ve written about before, or that you can find in something like monad-bayes.

The intuitions gleaned from playing with the Giry monad carry over precisely to other representations for the probability monad. In all cases, ‘return’ will correspond, semantically, to constructing a Dirac distribution at a point, while ‘bind’ will correspond to a marginalizing operator. The same is true for the underlying (applicative) functor structure: ‘fmap’ always corresponds to a density-preserving transformation of the support, while applicativeness corresponds to independence (yielding convolution, etc.). And you have to admit, the connection to continuations is pretty cool.

There is clearly some connection to the codensity monad as well, but I think I’ll let someone else figure out the specifics of that one. Something something right-Kan extension..

Foundations of the Giry Monad

2017-02-10T00:00:00+04:00

The Giry monad is the canonical probability monad that operates on the level of measures, which are the abstract constructs that canonically represent probability distributions. It’s sort of the baseline by which all other probability monads can be judged.

In this article I’m going to go through the categorical and measure-theoretic foundations of the Giry monad. In another article, I’ll describe how you can implement it in a very faithful sense in Haskell.

I was putting some notes together for another project and wound up writing up things up in a somewhat blog-friendly style, but this isn’t intended to be a tutorial per se. Really this isn’t the kind of content I’d usually post here, but since I’ve jotted everything up, I figured I may as well. If you like extremely dry mathematics and computer science, you’re in the right place.

I won’t define everything under the sun here - for properties or coherence conditions or other things that I’ve elided details on, check out something like Mac Lane or Aliprantis & Border. I’ll include some references at the end.

This is the game plan we’re working with:

Define monads and their supporting machinery in a categorical sense.
Define probability measures and some required background around that.
Construct the functor that maps a measurable space to the collection of all probability measures on that space.
Demonstrate that it’s a monad.

Let’s get started.

Categorical Foundations

A category $C$ is a collection of objects and morphisms between them. So if $W$, $X$, $Y$, and $Z$ are objects in $C$, then $f : W \to X$, $g : X \to Y$, and $h : Y \to Z$ are examples of morphisms. These morphisms can be composed in the obvious associative way, i.e.

\[f \circ (g \circ h) = (f \circ g) \circ h\]

and there exist identity morphisms (or automorphisms) that simply map objects to themselves.

A functor is a mapping between categories (equivalently, it’s a morphism in the category of so-called ‘small’ categories). The functor $F : C \to D$ takes every object in $C$ to some object in $D$, and every morphism in $C$ to some morphism in $D$, such that the structure of morphism composition is preserved. An endofunctor is a functor from a category to itself, and a bifunctor is a functor from a pair of categories to another category, i.e. $F : A \times B \to C$.

A natural transformation is a mapping between functors. So for two functors $F, G : C \to D$, a natural transformation $\epsilon : F \to G$ associates to every object $c$ in $C$ a morphism $\epsilon_c : F(c) \to G(c)$ in $D$.

A monoidal category $C$ is a category with some additional monoidal structure, namely an identity object $I$ and a bifunctor $\otimes : C \times C \to C$ called the tensor product, plus several natural isomorphisms that provide the associativity of the tensor product and its right and left identity with the identity object $I$.

A monoid $(M, \mu, \eta)$ in a monoidal category $C$ is an object $M$ in $C$ together with two morphisms (obeying the standard associativity and identity properties) that make use of the category’s monoidal structure: the associative binary operator $\mu : M \otimes M \to M$, and the identity $\eta : I \to M$.

A monad is (infamously) a ‘monoid in the category of endofunctors’. So take the category of endofunctors $\mathcal{F}$ whose objects are endofunctors and whose morphisms are natural transformations between them. This is a monoidal category; there exists an identity endofunctor $1_\mathcal{F}(F) = F$ for all $F$ in $\mathcal{F}$, plus a tensor product $\otimes : \mathcal{F} \times \mathcal{F} \to \mathcal{F}$ defined by functor composition such that the required associativity and identity properties hold. $\mathcal{F}$ is thus a monoidal category, and any specific monoid $(F, \mu, \eta)$ we construct on it is a specific monad.

Probabilistic Foundations

A measurable space $(X, \mathcal{X})$ is a set $X$ equipped with a topology-like structure called a $\sigma$-algebra $\mathcal{X}$ that essentially contains every well-behaved subset of $X$ in some sense. A measure $\nu : \mathcal{X} \to \mathbb{R}$ is a particular kind of set function from the $\sigma$-algebra to the nonnegative real line. A measure just assigns a generalized notion of area or volume to well-behaved subsets of $X$. In particular, if the total possible area or volume of the underlying set is 1 then we’re dealing with a probability measure. A measurable space completed with a measure, e.g. $(X, \mathcal{X}, \nu)$ is called a measure space, and a measurable space completed with a probability measure is called a probability space.

There is a lot of overloaded lingo around the word ‘measurable’. A ‘measurable set’ is an element of a $\sigma$-algebra in a measurable space. A measurable mapping is a mapping between measurable spaces. Given a ‘source’ measurable space $(X, \mathcal{X})$ and ‘target’ measurable space $(Y, \mathcal{Y})$, a measurable mapping $(X, \mathcal{X}) \to (Y, \mathcal{Y})$ is a map $T : X \to Y$ with the property that, for any measurable set in the target, the inverse image is measurable in the source. Or, formally, for any $B$ in $\mathcal{Y}$, you have that $T^{-1}(B)$ is in $\mathcal{X}$.

The Space of Probability Measures on a Measurable Space

If you consider the collection of all measurable spaces and measurable mappings between them, you get a category. Define $\textbf{Meas}$ to be the category of measurable spaces. So, objects are measurable spaces and morphisms are the measurable mappings between them.

For any specific measurable space $M$ in $\textbf{Meas}$, we can consider the space of all possible probability measures that could be placed on it and denote that $\mathcal{P}(M)$. To be clear, $\mathcal{P}(M)$ is a space of measures - that is, a space in which the points themselves are probability measures.

What’s remarkable about $\mathcal{P}(M)$ is that it is itself a measurable space. Let me explain.

As a probability measure, any element of $\mathcal{P}(M)$ is a function from measurable subsets of $M$ to the interval $[0, 1]$ in $\mathbb{R}$. That is: if $M$ is the measurable space $(X, \mathcal{X})$, then a point $\nu$ in $\mathcal{P}(M)$ is a function $\mathcal{X} \to \mathbb{R}$. For any measurable $A$ in $M$, there just naturally exists a sort of ‘evaluation’ mapping I’ll call $\tau_A: \mathcal{P}(M) \to \mathbb{R}$ that takes a measure on $M$ and evaluates it on the set $A$. To be explicit: if $\nu$ is a measure in $\mathcal{P}(M)$, then $\tau_A$ simply evaluates $\nu(A)$. It ‘runs’ the measure in a sense; in Haskell, $\tau_A$ would be analogous to a function like \f -> f a for some a.

This evaluation map $\tau_A$ corresponds to an integral. If you have a measurable space $(X, \mathcal{X})$, then for any $A$ a subset in $\mathcal{X}$, $\tau_A(\nu) = \nu(A) = \int_{X}\chi_A d\nu$ for $\chi$ the characteristic or indicator function of $A$ (where $\chi(x)$ is $1$ if $x$ is in $A$, and is $0$ otherwise). And we can actually extend $\tau$ to operate over measurable mappings from $(X, \mathcal{X})$ to $(\mathbb{R}, \mathcal{B}(\mathbb{R}))$, where $\mathcal{B}(\mathbb{R})$ is a suitable $\sigma$-algebra on $\mathbb{R}$. Here we typically use what’s called the Borel $\sigma$-algebra, which takes a topology on the set and then generates a $\sigma$-algebra from the open sets in the topology (for $\mathbb{R}$ we can just use the ‘usual’ topology generated by the Euclidean metric). For $f : X \to \mathbb{R}$ a measurable function, we can define the evaluation mapping $\tau_f : \mathcal{P}(M) \to \mathbb{R}$ as $\tau_f(\nu) = \int_X f d\nu$.

We can abuse notation here a bit and just use $\tau$ to refer to ‘duck typed’ mappings that evaluate measures over measurable sets or measurable functions depending on context. If we treat $\tau_A(\nu)$ as a function $\tau(\nu)(A)$, then $\tau(\nu)$ has type $\mathcal{X} \to \mathbb{R}$. If we treat $\tau_f(\nu)$ as a function $\tau(\nu)(f)$, then $\tau(\nu)$ has type $(X \to \mathbb{R}) \to \mathbb{R}$. I’ll say $\tau_{\{A, f\}}$ to refer to the mappings that accept either measurable sets or functions.

In any case. For a measurable space $M$, there exists a topology on $\mathcal{P}(M)$ called the weak-* topology that makes all the evaluation mappings $\tau_{\{A, f\}}$ continuous for any measurable set $A$ or measurable function $f$. From there, we can generate the Borel $\sigma$-algebra $\mathcal{B}(\mathcal{P}(M))$ that makes the evaluation functions $\tau_{\{A, f\}}$ measurable. The result is that $(\mathcal{P}(M), \mathcal{B}(\mathcal{P}(M)))$ is itself a measurable space, and thus an object in $\textbf{Meas}$.

The space $\mathcal{P}(M)$ actually has all sorts of insane properties that one wouldn’t expect - there are implications on convexity, completeness, compactness and such that carry over from $M$. But I digress.

$\mathcal{P}$ is a Functor

So: for any $M$ an object in $\textbf{Meas}$, we have that $\mathcal{P}(M)$ is also an object in $\textbf{Meas}$. And if you look at $\mathcal{P}$ like a functor, you notice that it takes objects of $\textbf{Meas}$ to objects of $\textbf{Meas}$. Indeed, you can define an analogous procedure on morphisms in $\textbf{Meas}$ as follows. Take $N$ to be another object (read: measurable space) in $\textbf{Meas}$ and $T : M \to N$ to be a morphism (read: measurable mapping) between them. Now, for any measure $\nu$ in $\mathcal{P}(M)$ we can define $\mathcal{P}(T)(\nu) = \nu \circ T^{-1}$ (this is called the image, distribution, or pushforward of $\nu$ under $T$). For some $T$ and $\nu$, $\mathcal{P}(T)(\nu)$ thus takes measurable sets in $N$ to a value in the interval $[0, 1]$ - that is, it is a measure on $\mathcal{P}(N)$. So we have that:

\[\mathcal{P}(T) : \mathcal{P}(M) \to \mathcal{P}(N)\]

and so $\mathcal{P}$ is an endofunctor on $\textbf{Meas}$.

$\mathcal{P}$ is a Monad

See where we’re going here? If we can define natural transformations $\mu$ and $\eta$ such that $(\mathcal{P}, \mu, \eta)$ is a monoid in the category of endofunctors, we’ll have defined a monad. We thus need to come up with a suitable monoidal structure, et voilà.

First the identity. We want a natural transformation $\eta$ between the identity functor $1_{\mathcal{F}}$ and the functor $\mathcal{P}$ such that $\eta_M : 1_{\mathcal{F}}(M) \to \mathcal{P}(M)$ for any measurable space $M$ in $\textbf{Meas}$. Evaluating the identity functor simplifies things to $\eta_M : M \to \mathcal{P}(M)$.

We can define this concretely as follows. Grab a measurable space $M$ in $\textbf{Meas}$ and define $\eta(x)(A) = \chi_A(x)$ for any point $x \in M$ and any measurable set $A \subseteq M$. $\eta(x)$ is thus a probability measure on $M$ - we assign $1$ to measurable sets that contain $x$, and 0 to those that don’t. If we peel away another argument, we have that $\eta : M \to \mathcal{P}(M)$, as required.

So $\eta$ takes points in measurable spaces to probability measures on those spaces. In technical parlance, it takes a point $x$ to the Dirac measure at $x$ - the probability measure that places the entirety of its mass at $x$.

Now for the other part of the monoidal structure, $\mu$. I initially found this next part to be a bit of a trip, but let me see what I can do about that.

Recall that the category of endofunctors, $\mathcal{F}$, is monoidal, so there exists a tensor product $\otimes : \mathcal{F} \times \mathcal{F} \to \mathcal{F}$ that we can deal with, which here just corresponds to functor composition. We’re looking for a natural transformation:

\[\mu : \mathcal{P} \circ \mathcal{P} \to \mathcal{P}\]

which is often written as:

\[\mu : \mathcal{P}^2 \to \mathcal{P}.\]

Take $M = (X, \mathcal{X})$ a measurable space in $\textbf{Meas}$ and then consider the space of probability measures over it, $\mathcal{P}(M)$. Then take the space of probability measures over the space of probability measures on $M$, $\mathcal{P}(\mathcal{P}(M))$. Since $\mathcal{P}$ is an endofunctor, this is again a measurable space, and for any measurable subset $A$ of $M$ we again have a family of mappings $\tau_A$ that take a probability measure in $\mathcal{P}(\mathcal{P}(M))$ and evaluate it on $A$. We want $\mu$ to be the thing that turns a measure over measures $\rho$ into a plain old probability measure on $\mathcal{P}(M)$.

In the context of probability theory, this kind of semigroup action is a marginalizing operator. We’re taking the ‘uncertainty’ captured in $\mathcal{P}(\mathcal{P}(M))$ via the probability measure $\rho$ and smearing it into the probability measures in $\mathcal{P}(M)$.

Take $\rho$ in $\mathcal{P}(\mathcal{P}(M))$ and some $A$ a measurable subset of $M$. We can define $\mu$ as follows:

\[\mu(\rho)(A) = \int_{\mathcal{P}(M)} \tau_A d\rho.\]

Using some lambda calculus notation to see the argument for $\tau_A$, we can expand the integrals to get the following gnarly expression:

\[\mu(\rho)(A) = \int_{\mathcal{P}(M)} \left\{\lambda \nu . \int_M \chi_A d \nu \right\} d \rho.\]

Notice what’s happening here. For $M$ a measurable space, we’re integrating over $\mathcal{P}(M)$ the space of probability measures on $M$, with respect to the probability measure $\rho$, which itself is a point in the space of probability measures over probability measures on $M$, $\mathcal{P}(\mathcal{P}(M))$. Whew.

The spaces we’re integrating over here are unusual, but $\rho$ is still a probability measure, so when applied to a measurable set in $\mathcal{B}(\mathcal{P}(M))$ it results in a probability in $[0, 1]$. So, peeling back an argument, we have that $\mu(\rho)$ has type $\mathcal{X} \to \mathbb{R}$. In other words, it’s a probability measure on $M$, and thus is in $\mathcal{P}(M)$. And if we peel back another argument, we find that:

\[\mu_M : \mathcal{P}(\mathcal{P}(M)) \to \mathcal{P}(M)\]

so, as required, that

\[\mu : \mathcal{P}^{2} \to \mathcal{P}.\]

It’s also worth noting that we can overload the notation for $\mu$ in the same way we did for $\tau$, i.e. to supply measurable functions in addition to measurable sets:

\[\mu(\rho)(f) = \int_{\mathcal{P}(M)} \left\{\lambda \nu . \int_M f d \nu \right\} d \rho.\]

Combining the three components, we get $(\mathcal{P}, \mu, \eta)$, the canonical Giry monad.

In Haskell, when we’re dealing with monads we typically use the bind operator $\gg\!\!=$ instead of manually dealing with the functorial structure and $\mu$ (called ‘join’). Bind has the type:

\[\gg\!\!= : \mathcal{P}(M) \to (M \to \mathcal{P}(N)) \to \mathcal{P}(N)\]

and for illustration, we can define $\gg\!\!=$ for the Giry monad like so:

\[(\rho \gg\!\!= g)(f) = \int_{M} \left\{ \lambda m . \int_N f d g(m) \right\} d\rho.\]

Here $\rho$ is in $\mathcal{P}(M)$, $g$ is in $M \to \mathcal{P}(N)$, and $f$ is in $N \to \mathbb{R}$, so note that we potentially simplify the outermost integral enormously. It now operates over a general measurable space, rather than a space of measures in particular, and this will come in handy when we get to implementation details in the next post.

Wrapping Up

That’s about it for now. It’s worth noting as a kind of footnote here that the existence of the Giry monad also obviously implies the existence of a Giry applicative functor. But the official situation for applicative functors seems kind of weird in this context, and I’m not yet up to the task of dealing with it formally.

Intuitively, one should be able to define the binary applicative operator characteristic of its lax monoidal structure as follows:

\[(\rho \, \langle \ast \rangle \, \nu)(f) = \int_{\mathcal{P}(M \to N)} \left\{\lambda T . \int_{M \to N} (f \circ T) d\nu \right\} d \rho.\]

But this has some really weird measure-theoretic implications - namely, that it assumes the existence of a space of probability measures over the space of all measurable functions $M \to N$, which is not trivial to define and indeed may not even exist. It seems like some people are looking into this problem as I just happened to stumble on this paper on the arXiv while doing some googling. I notice that some people on e.g. nLab require categories with additional structure beyond $\textbf{Meas}$ for the development of the Giry monad as well, for example the category of Polish (separable, completely metrizable) spaces $\textbf{Pol}$, so maybe the extra structure there takes care of the quirks.

Anyway. Applicatives are neat here because applicative probability measures are independent probability measures. And the existence of applicativeness means you can do all the things with independent probability measures that you might be used to. Measure convolution and friends are good examples. Given a measurable space $M$ that supports some notion of addition and two probability measures $\nu$ and $\zeta$ in $\mathcal{P}(M)$, we can add measures together via:

\[(\nu + \zeta)(f) = \int_{M}\int_{M}f(x + y)d\nu(x)d\zeta(y)\]

where $x$ and $y$ are both points in $M$. Subtraction and multiplication translate trivially as well.

In another article I’ll detail how the Giry monad can be implemented in Haskell and point out some neat extensions. There are some cool connections to continuations and codensity monads, and seemingly de Finetti’s theorem and exchangeability. That kind of thing. It’d also be worth trying to justify independence of probability measures from a categorical perspective, which seems easier than resolving the nitty-gritty measurability qualms I mentioned above.

‘Til then! Thanks to Jason Forbes for sifting through this stuff and providing some great comments.

References:

Rotating Squares

2017-01-04T00:00:00+04:00

Here’s a short one.

I use Colin Percival’s Hacker News Daily to catch the top ten articles of the day on Hacker News. Today an article called Why Recursive Data Structures? popped up, which illustrates that recursive algorithms can become both intuitive and borderline trivial when a suitable data structure is used to implement them. This is exactly the motivation for using recursion schemes.

In the above article, Reginald rotates squares by representing them via a quadtree. If we have a square of bits, something like:

.x..
..x.
xxx.
....

then we want to be able to easily rotate it 90 degrees clockwise, for example. So let’s define a quadtree in Haskell:

{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE LambdaCase #-}

import Data.Functor.Foldable
import Data.List.Split

data QuadTreeF a r =
    NodeF r r r r
  | LeafF a
  | EmptyF
  deriving (Show, Functor)

type QuadTree a = Fix (QuadTreeF a)

The four fields of the ‘NodeF’ constructor correspond to the upper left, upper right, lower right, and lower left quadrants of the tree respectively.

Gimme some embedded language terms:

node :: QuadTree a -> QuadTree a -> QuadTree a -> QuadTree a -> QuadTree a
node ul ur lr ll = Fix (NodeF ul ur lr ll)

leaf :: a -> QuadTree a
leaf = Fix . LeafF

empty :: QuadTree a
empty = Fix EmptyF

That lets us define quadtrees easily. Here’s the tree that the previous diagram corresponds to:

tree :: QuadTree Bool
tree = node ul ur lr ll where
  ul = node (leaf False) (leaf True) (leaf False) (leaf False)
  ur = node (leaf False) (leaf False) (leaf False) (leaf True)
  lr = node (leaf True) (leaf False) (leaf False) (leaf False)
  ll = node (leaf True) (leaf True) (leaf False) (leaf False)

Rotating is then really easy - we rotate each quadrant recursively. Just reach for a catamorphism:

rotate :: QuadTree a -> QuadTree a
rotate = cata $ \case
  NodeF ul ur lr ll -> node ll ul ur lr
  LeafF a           -> leaf a
  EmptyF            -> empty

Notice that you just have to shift each field of ‘NodeF’ rightward, with wraparound. Then if you rotate and render the original tree you get:

.x..
.x.x
.xx.
....

Rotating things more times yields predictable results.

If you want to rotate another structure - say, a flat list - you can go through a quadtree as an intermediate representation using the same pattern I described in Sorting with Style. Build yourself a coalgebra and algebra pair:

builder :: [a] -> QuadTreeF a [a]
builder = \case
  []  -> EmptyF
  [x] -> LeafF x
  xs  -> NodeF a b c d where
    [a, b, c, d] = chunksOf (length xs `div` 4) xs

consumer :: QuadTreeF a [a] -> [a]
consumer = \case
  EmptyF            -> []
  LeafF a           -> [a]
  NodeF ul ur lr ll -> concat [ll, ul, ur, lr]

and then glue them together with a hylomorphism:

rotateList :: [a] -> [a]
rotateList = hylo consumer builder

Neato.

For a recent recursion scheme resource I’ve spotted on the Twitters, check out Pascal Hartig’s compendium in progress.

Promorphisms, Pre and Post

2016-11-26T00:00:00+04:00

To the.. uh, ‘layperson’, pre- and postpromorphisms are probably well into the WTF category of recursion schemes. This is a mistake - they’re simple and useful, and I’m going to try and convince you of this in short order.

Preliminaries:

{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE LambdaCase #-}

import Data.Functor.Foldable
import Prelude hiding (sum)

For simplicity, let’s take a couple of standard interpreters on lists. We’ll define ‘sumAlg’ as an interpreter for adding up list contents and ‘lenAlg’ for just counting the number of elements present:

sumAlg :: Num a => ListF a a -> a
sumAlg = \case
  Cons h t -> h + t
  Nil      -> 0

lenAlg :: ListF a Int -> Int
lenAlg = \case
  Cons h t -> 1 + t
  Nil      -> 0

Easy-peasy. We can use cata to make these things useful:

sum :: Num a => [a] -> a
sum = cata sumAlg

len :: [a] -> Int
len = cata lenAlg

Nothing new there; ‘sum [1..10]’ will give you 55 and ‘len [1..10]’ will give you 10.

An interesting twist is to consider only small elements in some sense; say, we only want to add or count elements that are less than or equal to 10, and ignore any others.

We could rewrite the previous interpreters, manually checking for the condition we’re interested in and handling it accordingly:

smallSumAlg :: (Ord a, Num a) => ListF a a -> a
smallSumAlg = \case
  Cons h t ->
    if   h <= 10
    then h + t
    else 0
  Nil      -> 0

smallLenAlg :: (Ord a, Num a) => ListF a Int -> Int
smallLenAlg = \case
  Cons h t ->
    if   h <= 10
    then 1 + t
    else 0
  Nil      -> 0

And you get ‘smallSum’ and ‘smallLen’ by using ‘cata’ on them respectively. They work like you’d expect - ‘smallLen [1, 5, 20]’ ignores the 20 and just returns 2, for example.

You can do better though. Enter the prepromorphism.

Instead of writing additional special-case interpreters for the ‘small’ case, consider the following natural transformation on the list base functor. It maps the list base functor to itself, without needing to inspect the carrier type:

small :: (Ord a, Num a) => ListF a b -> ListF a b
small Nil = Nil
small term@(Cons h t)
  | h <= 10   = term
  | otherwise = Nil

A prepromorphism is a ‘cata’-like recursion scheme that proceeds by first applying a natural transformation before interpreting via a supplied algebra. That’s.. surprisingly simple. Here are ‘smallSum’ and ‘smallLen’, defined without needing to clumsily create new special-case algebras:

smallSum :: (Ord a, Num a) => [a] -> a
smallSum = prepro small sumAlg

smallLen :: (Ord a, Num a) => [a] -> Int
smallLen = prepro small lenAlg

They work great:

> smallSum [1..100]
55
> smallLen [1..100]
10

In pseudo category-theoretic notation you visualize how a prepromorphism works via the following commutative diagram:

The only difference, when compared to a standard catamorphism, is the presence of the natural transformation applied via the looping arrow in the top left. The natural transformation ‘h’ has type ‘forall r. Base t r -> Base t r’, and ‘embed’ has type ‘Base t t -> t’, so their composition gets you exactly the type you need for an algebra, which is then the input to ‘cata’ there. Mapping the catamorphism over the type ‘Base t t’ brings it right back to ‘Base t t’.

A postpromorphism is dual to a prepromorphism. It’s ‘ana’-like; proceed with your corecursive production, applying natural transformations as you go.

Here’s a streaming coalgebra:

streamCoalg :: Enum a => a -> ListF a a
streamCoalg n = Cons n (succ n)

A normal anamorphism would just send this thing shooting off into infinity, but we can use the existing ‘small’ natural transformation to cap it at 10:

smallStream :: (Ord a, Num a, Enum a) => a -> [a]
smallStream = postpro small streamCoalg

You get what you might expect:

> smallStream 3
[3,4,5,6,7,8,9,10]

And similarly, you can visualize a postpromorphism like so:

In this case the natural transformation is applied after mapping the postpromorphism over the base functor (hence the ‘post’ namesake).

Comonadic Markov Chain Monte Carlo

2016-10-26T00:00:00+04:00

Some time ago I came across a way to in-principle perform inference on certain probabilistic programs using comonadic structures and operations.

I decided to dig it up and try to use it to extend the simple probabilistic programming language I talked about a few days ago with a stateful, experimental inference backend. In this post we’ll

Represent probabilistic programs as recursive types parameterized by a terminating instruction set.
Represent execution traces of probabilistic programs via a simple transformation of our program representation.
Implement the Metropolis-Hastings algorithm over this space of execution traces and thus do some inference.

Let’s get started!

Representing Programs That Terminate

I like thinking of embedded languages in terms of instruction sets. That is: I want to be able to construct my embedded language by first defining a collection of abstract instructions and then using some appropriate recursive structure to represent programs over that set.

In the case of probabilistic programs, our instructions are probability distributions. Last time we used the following simple instruction set to define our embedded language:

data ModelF r =
    BernoulliF Double (Bool -> r)
  | BetaF Double Double (Double -> r)
  deriving Functor

We then created an embedded language by just wrapping it up in the higher-kinded Free type to denote programs of type Model.

data Free f a =
    Pure a
  | Free (f (Free f a))

type Model = Free ModelF

Recall that Free represents programs that can terminate, either by some instruction in the underlying instruction set, or via the Pure constructor of the Free type itself. The language defined by Free ModelF is expressive enough to easily construct a ‘forward-sampling’ interpreter, as well as a simple rejection sampler for performing inference.

Notice that we don’t have a terminating instruction in ModelF itself - if we’re using it, then we need to rely on the Pure constructor of Free to terminate programs. Otherwise they’d just have to recurse forever. This can be a bit limiting if we want to transform a program of type Free ModelF to something else that doesn’t have a notion of termination baked-in (Fix, for example).

Let’s tweak the ModelF type to get the following:

data ModelF a r =
    BernoulliF Double (Bool -> r)
  | BetaF Double Double (Double -> r)
  | NormalF Double Double (Double -> r)
  | DiracF a
  deriving Functor

Aside from adding another foundational distribution - NormalF - we’ve also added a new constructor, DiracF, which carries a parameter with type a. We need to incorporate this carrier type in the overall type of ModelF as well, so ModelF itself also gets a new type parameter to carry around.

The DiracF instruction is a terminating instruction; it has no recursive point and just terminates with a value of type a when reached. It’s structurally equivalent to the Pure a branch of Free that we were relying on to terminate our programs previously - the only thing we’ve done is add it to our instruction set proper.

Why DiracF? A Dirac distribution places the entirety of its probability mass on a single point, and this is the exact probabilistic interpretation of the applicative pure or monadic return that one encounters with an appropriate probability type. Intuitively, if I sample a value $x$ from a uniform distribution, then that is indistinguishable from sampling $x$ from said uniform distribution and then sampling from a Dirac distribution with parameter $x$.

Make sense? If not, it might be helpful to note that there is no difference between any of the following (to which uniform and dirac are analogous):

> action :: m a
> action >>= return :: m a
> action >>= return >>= return >>= return :: m a

Wrapping ModelF a up in Free, we get the following general type for our programs:

type Program a = Free (ModelF a)

And we can construct a bunch of embedded language terms in the standard way:

beta :: Double -> Double -> Program a Double
beta a b = liftF (BetaF a b id)

bernoulli :: Double -> Program a Bool
bernoulli p = liftF (BernoulliF p id)

normal :: Double -> Double -> Program a Double
normal m s = liftF (NormalF m s id)

dirac :: a -> Program a b
dirac x = liftF (DiracF x)

Program is a general type, capturing both terminating and nonterminating programs via its type parameters. What do I mean by this? Note that in Program a b, the a type parameter can only be concretely instantiated via use of the terminating dirac term. On the other hand, the b type parameter is unaffected by the dirac term; it can only be instantiated by the other nonterminating terms: beta, bernoulli, normal, or compound expressions of these.

We can thus distinguish between terminating and nonterminating programs at the type level, like so:

type Terminating a = Program a Void

type Model b = forall a. Program a b

Void is the uninhabited type, brought into scope via Data.Void or simply defined via data Void = Void Void. Any program that ends via a dirac instruction must be Terminating, and any program that doesn’t end with a dirac instruction can not be Terminating. We’ll just continue to call a nonterminating program a Model, as before.

Good. So if it’s not clear: from a user’s perspective, nothing has changed. We still write probabilistic programs using simple monadic language terms. Here’s a Gaussian mixture model where the mixing parameter follows a beta distribution, for example:

mixture :: Double -> Double -> Model Double
mixture a b = do
  prob   <- beta a b
  accept <- bernoulli prob
  if   accept
  then normal (negate 2) 0.5
  else normal 2 0.5

Meanwhile the syntax tree generated looks something like the following. It’s more or less a traditional probabilistic graphical model description of our program:

It’s important to note that in this embedded framework, the only pieces of the syntax tree that we can observe are those related directly to our primitive instructions. For our purposes this is excellent - we can focus on programs entirely at the level of their probabilistic components, and ignore the deterministic parts that would otherwise be distractions.

To collect samples from mixture, we can first interpret it into a sampling function, and then simulate from it. The toSampler function from last time doesn’t change much:

toSampler :: Program a a -> Prob IO a
toSampler = iterM $ \case
  BernoulliF p f -> Prob.bernoulli p >>= f
  BetaF a b f    -> Prob.beta a b >>= f
  NormalF m s f  -> Prob.normal m s >>= f
  DiracF x       -> return x

Sampling from mixture 2 3 a thousand times yields the following

> simulate (toSampler (mixture 2 3))

Note that the rightmost component gets more traffic due to the hyperparameter combination of 2 and 3 that we provided to mixture.

Also, a note - since we have general recursion in Haskell, so-called ‘terminating’ programs here can actually.. uh, fail to terminate. They must only terminate as far as we can express the sentiment at the embedded language level. Consider the following, for example:

foo :: Terminating a
foo = (loop 1) >>= dirac where
  loop a = do
    p <- beta a 1
    loop p

foo here doesn’t actually terminate. But at least this kind of weird case can be picked up in the types:

> :t simulate (toSampler foo)
simulate (toSampler foo) :: IO Void

If you try to sample from a distribution over Void or forall a. a then I can’t be held responsible for what you get up to. But there are other cases, sadly, where we’re also out of luck:

trollGeometric :: Double -> Model Int
trollGeometric p = loop where
  loop = do
    accept <- return False
    if   accept
    then return 1
    else fmap succ loop

A geometric distribution that actually used its argument $p$, for $0 < p \leq 1$, could be guaranteed to terminate with probability 1. This one doesn’t, so trollGeometric undefined >>= dirac won’t.

At the end of the day we’re stuck with what our host language offers us. So, take the termination guarantees for our embedded language with a grain of salt.

Stateful Inference

In the previous post we used a simple rejection sampler to sample from a conditional distribution. ‘Vanilla’ Monte Carlo algorithms like rejection and importance sampling are stateless. This makes them nice in some ways - they tend to be simple to implement and are embarrassingly parallel, for example. But the curse of dimensionality prevents them from scaling well to larger problems. I won’t go into detail on that here - for a deep dive on the topic, you probably won’t find anything better than this phenomenal couple of talks on MCMC that Iain Murray gave at a MLSS session in Cambridge in 2009. I think they’re unparalleled to this day.

The point is that in higher dimensions we tend to get a lot out of state. Essentially, if one finds an interesting region of high-dimensional parameter space, then it’s better to remember where that is, rather than forgetting it exists as soon as one stumbles onto it. The manifold hypothesis conjectures that interesting regions of space tend to be near other interesting regions of space, so exploring neighbourhoods of interesting places tends to pay off. Stateful Monte Carlo methods - namely, the family of Markov chain Monte Carlo algorithms - handle exactly this, by using a Markov chain to wander over parameter space. I’ve written on MCMC in the past - you can check out some of those articles if you’re interested.

In the stateless rejection sampler we just performed conditional inference via the following algorithm:

Sample from a parameter model.
Sample from a data model, using the sample from the parameter model as input.
If the sample from the data model matches the provided observations, return the sample from the parameter model.

By repeating this many times we get a sample of arbitrary size from the appropriate conditional, inverse, or posterior distribution (whatever you want to call it).

In a stateful inference routine - here, the good old Metropolis-Hastings algorithm - we’re instead going to do the following repeatedly:

Sample from a parameter model, recording the way the program executed in order to return the sample that it did.
Compute the cost, in some sense, of generating the provided observations, using the sample from the parameter model as input.
Propose a new sample from the parameter model by perturbing the way the program executed and recording the new sample the program outputs.
Compute the cost of generating the provided observations using this new sample from the parameter model as input.
Compare the costs of generating the provided observations under the respective samples from the parameter models.
With probability depending on the ratio of the costs, flip a coin. If you see a head, then move to the new, proposed execution trace of the program. Otherwise, stay at the old execution trace.

This procedure generates a Markov chain over the space of possible execution traces of the program - essentially, plausible ways that the program could have executed in order to generate the supplied observations.

Implementations of Church use variations of this method to do inference, the most famous of which is a low-overhead transformational compilation procedure described in a great and influential 2011 paper by David Wingate et al.

Representing Running Programs

To perform inference on probabilistic programs according to the aforementioned Metropolis-Hastings algorithm, we need to represent executing programs somehow, in a form that enables us to examine and modify their internal state.

How can we do that? We’ll pluck another useful recursive structure from our repertoire and consider the humble Cofree:

data Cofree f a = a :< f (Cofree f a)

Recall that Cofree allows one to annotate programs with arbitrary information at each internal node. This is a great feature; if we can annotate each internal node with important information about its state - its current value, the current state of its generator, the ‘cost’ associated with it - then we can walk through the program and examine it as required. So, it can capture a ‘running’ program in exactly the way we need.

Let’s describe running programs as values having the following Execution type:

type Execution a = Cofree (ModelF a) Node

The Node type is what we’ll use to describe the internal state of each node on the program. I’ll define it like so:

data Node = Node {
    nodeCost    :: Double
  , nodeValue   :: Dynamic
  , nodeSeed    :: MWC.Seed
  , nodeHistory :: [Dynamic]
  } deriving Show

I’ll elaborate on this type below, but you can see that it captures a bunch of information about the state of each node.

One can mechanically transform any Free-encoded program into a Cofree-encoded program, so long as the original Free-encoded program can terminate of its own accord, i.e. on the level of its own instructions. Hence the need for our Terminating type and all that.

In our case, setting everything up just right takes a bit of code, mainly around handling pseudo-random number generators in a pure fashion. So I won’t talk about every little detail of it right here. The general idea is to write a function that takes instructions to the appropriate state captured by a Node value, like so:

initialize :: Typeable a => MWC.Seed -> ModelF a b -> Node
initialize seed = \case
  BernoulliF p _ -> runST $ do
    (nodeValue, nodeSeed) <- samplePurely (Prob.bernoulli p) seed
    let nodeCost    = logDensityBernoulli p (unsafeFromDyn nodeValue)
        nodeHistory = mempty
    return Node {..}

  BetaF a b _ -> runST $ do
    (nodeValue, nodeSeed) <- samplePurely (Prob.beta a b) seed
    let nodeCost    = logDensityBeta a b (unsafeFromDyn nodeValue)
        nodeHistory = mempty
    return Node {..}

  ...

You can see that for each node, I sample from it, calculate its cost, and then initialize its ‘history’ as an empty list.

Here it’s worth going into a brief aside.

There are two mildly annoying things we have to deal with in this situation. First, individual nodes in the program typically sample values at different types, and second, we can’t easily use effects when annotating. This means that we have to pack heterogeneously-typed things into a homogeneously-typed container, and also use pure random number generation facilities to sample them.

A quick-and-dirty answer for the first case is to just use dynamic typing when storing the values. It works and is easy, but of course is subject to the standard caveats. I use a function called unsafeFromDyn to convert dynamically-typed values back to a typed form, so you can gauge the safety of all this for yourself.

For the second case, I just use the ST monad, along with manual state snapshotting, to execute and iterate a random number generator. Pretty simple.

Also: in terms of efficiency, keeping a node’s history on-site at each execution falls into the ‘completely insane’ category, but let’s not worry much about efficiency right now. Prototypes gonna prototype and all that.

Anyway.

Given this initialize function, we can transform a terminating program into a running program by simple recursion. Again, we can only transform programs with type Terminating a because we need to rule out the case of ever visiting the Pure constructor of Free. We handle that by the absurd function provided by Data.Void:

execute :: Typeable a => Terminating a -> Execution a
execute = annotate defaultSeed where
  defaultSeed         = (42, 108512)
  annotate seeds term = case term of
    Pure r -> absurd r
    Free instruction ->
      let (nextSeeds, generator) = xorshift seeds
          seed  = MWC.toSeed (V.singleton generator)
          node  = initialize seed instruction
      in  node :< fmap (annotate nextSeeds) instruction

And there you have it - execute takes a terminating program as input and returns a running program - an execution trace - as output. The syntax tree we had previously gets turned into something like this:

Perturbing Running Programs

Given an execution trace, we’re able to step through it sequentially and investigate the program’s internal state. But to do inference we also need to modify it as well. What’s the answer here?

Just as Free has a monadic structure that allows us to write embedded programs using built-in monadic combinators and do-notation, Cofree has a comonadic structure that is amenable to use with the various comonadic combinators found in Control.Comonad. The most important for our purposes is the comonadic ‘extend’ operation that’s dual to monad’s ‘bind’:

extend :: Comonad w => (w a -> b) -> w a -> w b
extend f = fmap f . duplicate

To perturb a running program, we can thus write a function that perturbs any given annotated node, and then extend it over the entire execution trace.

The perturbNode function can be similar to the initialize function from earlier; it describes how to perturb every node based on the instruction found there:

perturbNode :: Execution a -> Node
perturbNode (node@Node {..} :< cons) = case cons of
  BernoulliF p _ -> runST $ do
    (nvalue, nseed) <- samplePurely (Prob.bernoulli p) nodeSeed
    let nscore   = logDensityBernoulli p (unsafeFromDyn nvalue)
    return $! Node nscore nvalue nseed nodeHistory

  BetaF a b _ -> runST $ do
    (nvalue, nseed) <- samplePurely (Prob.beta a b) nodeSeed
    let nscore   = logDensityBeta a b (unsafeFromDyn nvalue)
    return $! Node nscore nvalue nseed nodeHistory

  ...

Note that this is a very crude way to perturb nodes - we’re just sampling from whatever distribution we find at each one. A more refined procedure would sample from each node on a more local basis, sampling from its respective domain in a neighbourhood of its current location. For example, to perturb a BetaF node we might sample from a tiny Gaussian bubble around its current location, repeating the process if we happen to ‘fall off’ the support. I’ll leave matters like that for another post.

Perturbing an entire trace is then as easy as I claimed it to be:

perturb :: Execution a -> Execution a
perturb = extend perturbNode

For some comonadic intuition: when we ‘extend’ a function over an execution, the trace itself gets ‘duplicated’ in a comonadic context. Each node in the program becomes annotated with a view of the rest of the execution trace from that point forward. It can be difficult to visualize at first, but I reckon the following image is pretty faithful:

Each annotation then has perturbNode applied to it, which reduces the trace back to the standard annotated version we saw before.

Iterating the Markov Chain

So: to move around in parameter space, we’ll propose state changes by perturbing the current state, and then accept or reject proposals according to local economic conditions.

If you already have no idea what I’m talking about, then the phrase ‘local economic conditions’ probably didn’t help you much. But it’s a useful analogy to have in one’s head. Each state in parameter space has a cost associated with it - the cost of generating the observations that we’re conditioning on while doing inference. If certain parameter values yield a data model that is unlikely to generate the provided observations, then those observations will be expensive to generate when measured in terms of log-likelihood. Parameter values that yield data models more likely to generate the supplied observations will be comparatively cheaper.

If a proposed execution trace is significantly cheaper than the trace we’re currently at, then we usually want to move to it. We allow some randomness in our decision to keep everything nice and measure-preserving.

We can thus construct the conditional distribution over execution traces using the following invert function, using the same nomenclature as the rejection sampler we used previously. To focus on the main points, I’ll elide some of its body:

invert
  :: (Eq a, Typeable a, Typeable b)
  => Int -> [a] -> Model b -> (b -> a -> Double)
  -> Model (Execution b)
invert epochs obs prior ll = loop epochs (execute terminated) where
  terminated = prior >>= dirac
  loop n current
    | n == 0    = return current
    | otherwise = do
        let proposal = perturb current

            -- calculate costs and movement probability here

        accept <- bernoulli prob
        let next = if accept then proposal else stepGenerators current
        loop (pred n) (snapshot next)

There are a few things to comment on here.

First, notice how the return type of invert is Model (Execution b)? Using the semantics of our embedded language, it’s literally a standard model over execution traces. The above function returns a first-class value that is completely uninterpreted and abstract. Cool.

We’re also dealing with things a little differently from the rejection sampler that we built previously. Here, the data model is expressed by a cost function; that is, a function that takes a parameter value and observation as input, and returns the cost of generating the observation (conditional on the supplied parameter value) as output. This is the approach used in the excellent Practical Probabilistic Programming with Monads paper by Adam Scibior et al and also mentioned by Dan Roy in his recent talk at the Simons Institute. Ideally we’d just reify the cost function here from the description of a model directly (to keep the interface similar to the one used in the rejection sampler implementation), but I haven’t yet found a way to do this in a type-safe fashion.

Regardless of whether or not we accept a proposed move, we need to snapshot the current value of each node and add it to that node’s history. This can be done using another comonadic extend:

snapshotValue :: Cofree f Node -> Node
snapshotValue (Node {..} :< cons) = Node { nodeHistory = history, .. } where
  history = nodeValue : nodeHistory

snapshot :: Functor f => Cofree f Node -> Cofree f Node
snapshot = extend snapshotValue

The other point of note is minor, but an extremely easy detail to overlook. Since we’re handling random value generation at each node purely, using on-site PRNGs, we need to iterate the generators forward a step in the event that we don’t accept a proposal. Otherwise we’d propose a new execution based on the same generator states that we’d used previously! For now I’ll just iterate the generators by forcing a sample of a uniform variate at each node, and then throwing away the result. To do this we can use the now-standard comonadic pattern:

stepGenerator :: Cofree f Node -> Node
stepGenerator (Node {..} :< cons) = runST $ do
  (nval, nseed) <- samplePurely (Prob.beta 1 1) nodeSeed
  return Node {nodeSeed = nseed, ..}

stepGenerators :: Functor f => Cofree f Node -> Cofree f Node
stepGenerators = extend stepGenerator

Inspecting Execution Traces

Alright so let’s see how this all works. Let’s write a model, condition it on some observations, and do inference.

We’ll choose our simple Gaussian mixture model from earlier, where the mixing probability follows a beta distribution, and cluster assignment itself follows a Bernoulli distribution. We thus choose the ‘leftmost’ component of the mixture with the appropriate mixture probability.

We can break the mixture model up as follows:

prior :: Double -> Double -> Model Bool
prior a b = do
  p <- beta a b
  bernoulli p

likelihood :: Bool -> Model Double
likelihood left
  | left      = normal (negate 2) 0.5
  | otherwise = normal 2 0.5

Let’s take a look at some samples from the marginal distribution. This time I’ll flip things and assign hyperparameters of 3 and 2 for the prior:

> simulate (toSampler (prior 3 2 >>= likelihood))

It looks like we’re slightly more likely to sample from the left mixture component than the right one. Again, this makes sense - the mean of a beta(3, 2) distribution is 0.6.

Now, what about inference? I’ll define the conditional model as follows:

posterior :: Model (Execution Bool)
posterior = invert 1000 obs prior ll where
  obs = [ -1.7, -1.8, -2.01, -2.4
        , 1.9, 1.8
        ]

  ll left
    | left      = logDensityNormal (negate 2) 0.5
    | otherwise = logDensityNormal 2 0.5

Here we have four observations that presumably arise from the leftmost component, and only two that match up with the rightmost. Note also that I’ve replaced the likelihood model with its appropriate cost function due to reasons I mentioned in the last section. (It would be easy to reify this model as its cost function, but doing it for general models is trickier)

Anyway, let’s sample from the conditional distribution:

> simulate (toSampler posterior)

Sampling returns a running program, of course, and we can step through it to examine its structure. We can use the supplied values recorded at each node to ‘automatically’ step through execution, or we can supply our own values to investigate arbitrary branches.

The conditional distribution we’ve found over the mixing probability is as follows:

Looks like we’re in the right ballpark.

We can examine the traces of other elements of the program as well. Here’s the recorded distribution over component assignments, for example - note that the rightmost bar here corresponds to the leftmost component in the mixture:

You can see that whenever we wandered into the rightmost component, we’d swiftly wind up jumping back out of it:

Comments

This is a fun take on probabilistic programming. In particular I find a few aspects of the whole setup to be pretty attractive:

We use a primitive, limited instruction set to parameterize both programs - via Free - and running programs - via Cofree. These off-the-shelf recursive types are used to wrap things up and provide most of our required control flow automatically. It’s easy to transparently add structure to embedded programs built in this way; for example, we can statically encode independence by replacing our ModelF a type with something like:

data InstructionF a = Coproduct (ModelF a) (Ap (ModelF a))

This can be hidden from the user so that we’re left with the same simple monadic syntax we presently enjoy, but we also get to take independence into account when performing inference, or any other structural interpretation for that matter.

When it comes to inference, the program representation is completely separate from whatever inference backend we choose to augment it with. We can deal with traces as first-class values that can be directly stored, inspected, manipulated, and so on. And everything is done in a typed and purely-functional framework. I’ve used dynamic typing functionality from Data.Dynamic to store values in execution traces here, but we could similarly just define a concrete Value type with the appropriate constructors for integers, doubles, bools, etc., and use that to store everything.

At the same time, this is a pretty early concept - doing inference efficiently in this setting is another matter, and there are a couple of computational and statistical issues here that need to be ironed out to make further progress.

The current way I’ve organized Markov chain generation and iteration is just woefully inefficient. Storing the history of each node on-site is needlessly costly and I’m sure results in a ton of unnecessary allocation. On a semantic level, it also ‘complects’ state and identity: why, after all, should a single execution trace know anything about traces that preceded it? Clearly this should be accumulated in another data structure. There is a lot of other low-hanging fruit around strictness and PRNG management as well.

From a more statistical angle, the present implementation does a poor job when it comes to perturbing execution traces. Some changes - such as improving the proposal mechanism for a given instruction - are easy to implement, and representing distributions as instructions indeed makes it easy to tailor local proposal distributions in a context-independent way. But another problem is that, by using a ‘blunt’ comonadic extend, we perturb an execution by perturbing every node in it. In general it’s better to make small perturbations rather than large ones to ensure a reasonable acceptance ratio, but to do that we’d need to perturb single nodes (or at least subsets of nodes) at a time.

There may be some inroads here via comonad transformers like StoreT or lenses that would allow us to zoom in on a particular node and perturb it, rather than perturbing everything at once. But my comonad-fu is not yet quite at the required level to evaluate this, so I’ll come back to that idea some other time.

I’m interested in playing with this concept some more in the future, though I’m not yet sure how much I expect it to be a tenable way to do inference at scale. If you’re interested in playing with it, I’ve dumped the code from this post into this gist.

Thanks to Niffe Hermansson and Fredrik Olsen for reviewing a draft of this post and providing helpful comments.

A Simple Embedded Probabilistic Programming Language

2016-10-17T00:00:00+04:00

What does a dead-simple probabilistic programming language look like? The simplest thing I can imagine involves three components:

A representation for probabilistic models.
A way to simulate from those models (‘forward’ sampling).
A way to sample from a conditional model (‘backward’ sampling).

Rob Zinkov wrote an article on this type of thing around a year ago, and Dan Roy recently gave a talk on the topic as well. In the spirit of unabashed unoriginality, I’ll give a sort of composite example of the two. Most of the material here comes directly from Dan’s talk; definitely check it out if you’re curious about this whole probabilistic programming mumbojumbo.

Let’s whip together a highly-structured, typed, embedded probabilistic programming language - the core of which will encompass a tiny amount of code.

Some preliminaries - note that you’ll need my simple little mwc-probability library handy for when it comes time to do sampling:

{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE LambdaCase #-}

import Control.Monad
import Control.Monad.Free
import qualified System.Random.MWC.Probability as MWC

Representing Probabilistic Models

Step one is to represent the fundamental constructs found in probabilistic programs. These are abstract probability distributions; I like to call them models:

data ModelF r =
    BernoulliF Double (Bool -> r)
  | BetaF Double Double (Double -> r)
  deriving Functor

type Model = Free ModelF

Each foundational probability distribution we want to consider is represented as a constructor of the ModelF type. You can think of them as probabilistic instructions, in a sense. A Model itself is a program parameterized by this probabilistic instruction set.

In a more sophisticated implementation you’d probably want to add more primitives, but you can get pretty far with the beta and Bernoulli distributions alone. Here are some embedded language terms, only two of which correspond one-to-one with to the constructors themselves:

bernoulli :: Double -> Model Bool
bernoulli p = liftF (BernoulliF p id)

beta :: Double -> Double -> Model Double
beta a b = liftF (BetaF a b id)

uniform :: Model Double
uniform = beta 1 1

binomial :: Int -> Double -> Model Int
binomial n p = fmap count coins where
  count = length . filter id
  coins = replicateM n (bernoulli p)

betaBinomial :: Int -> Double -> Double -> Model Int
betaBinomial n a b = do
  p <- beta a b
  binomial n p

You can build a lot of other useful distributions by just starting from the beta and Bernoulli as well. And technically I guess the more foundational distributions to use here would be the Dirichlet and categorical, of which the beta and Bernoulli are special cases. But I digress. The point is that other distributions are easy to construct from a set of reliable primitives; you can check out the old lambda-naught paper by Park et al for more examples.

See how binomial and betaBinomial are defined? In the case of binomial we’re using the property that models have a functorial structure by just mapping a counting function over the result of a bunch of Bernoulli random variables. For betaBinomial we’re directly making use of our monadic structure, first describing a weight parameter via a beta distribution and then using it as an input to a binomial distribution.

Note in particular that we’ve expressed betaBinomial by binding a parameter model to a data model. This is a foundational pattern in Bayesian statistics; in the more usual lingo, the parameter model corresponds to the prior distribution, and the data model is the likelihood.

Forward-Mode Sampling

So we have our representation. Next up, we want to simulate from these models. Thus far they’re purely abstract, and don’t encode any information about probability or sampling or what have you. We have to ascribe that ourselves.

mwc-probability defines a monadic sampling-based probability distribution type called Prob, and we can use a basic recursion scheme on free monads to adapt our own model type to that:

toSampler :: Model a -> MWC.Prob IO a
toSampler = iterM $ \case
  BernoulliF p f -> MWC.bernoulli p >>= f
  BetaF a b f    -> MWC.beta a b >>= f

We can glue that around the relevant mwc-probability functionality to simulate from models directly:

simulate :: Model a -> IO a
simulate model = MWC.withSystemRandom . MWC.asGenIO $
  MWC.sample (toSampler model)

And this can be used with standard monadic combinators like replicateM to collect larger samples:

> replicateM 10 $ simulate (betaBinomial 10 1 4)
[5,7,1,4,4,1,1,0,4,2]

Reverse-Mode Sampling

Now. Here we want to condition our model on some observations and then recover the conditional distribution over its internal parameters.

This part - inference - is what makes probabilistic programming hard, and doing it really well remains an unsolved problem. One of the neat theoretical results in this space due to Ackerman, Freer, and Roy is that in the general case the problem is actually unsolvable, in that one can encode as a probabilistic program a conditional distribution that computes the halting problem. Similarly, in general it’s impossible to do this sort of thing efficiently even for computable conditional distributions. Consider the case of a program that returns the hash of a random n-long binary string, and then try to infer the distribution over strings given some hashes, for example. This is never going to be a tractable problem.

For now let’s use a simple rejection sampler to encode a conditional distribution. We’ll require some observations, a proposal distribution, and the model that we want to invert:

invert :: (Monad m, Eq b) => m a -> (a -> m b) -> [b] -> m a
invert proposal model observed = loop where
  loop = do
    parameters <- proposal
    generated  <- replicateM (length observed) (model parameters)
    if   generated == observed
    then return parameters
    else loop

Let’s use it to compute the posterior or inverse model of an (apparently) biased coin, given a few observations. We’ll just use a uniform distribution as our proposal:

posterior :: Model Double
posterior = invert [True, True, False, True] uniform bernoulli

Let’s grab some samples from the posterior distribution:

> replicateM 1000 (simulate posterior)

The central tendency of the posterior floats about 0.75, which is what we’d expect, given our observations. This has been inferred from only four points; let’s try adding a few more. But before we do that, note that the present way the rejection sampling algorithm works is:

Propose a parameter value according to the supplied proposal distribution.
Generate a sample from the model, of equal size to the supplied observations.
Compare the collected sample to the supplied observations. If they’re equal, then return the proposed parameter value. Otherwise start over.

Rejection sampling isn’t exactly efficient in nontrivial settings anyway, but it’s supremely inefficient for our present case. The random variables we’re interested in are exchangeable, so what we’re concerned about is the total number of True or False values observed - not any specific order they appear in.

We can add an ‘assistance’ function to the rejection sampler to help us out in this case:

invertWithAssistance
  :: (Monad m, Eq c) => ([a] -> c) -> m b -> (b -> m a) -> [a] -> m b
invertWithAssistance assister proposal model observed = loop where
  loop = do
    parameters <- proposal
    generated  <- replicateM (length observed) (model parameters)
    if   assister generated == assister observed
    then return parameters
    else loop

The assister summarizes both our observations and collected sample to ensure they’re efficiently comparable. In our situation, we can use a simple counting function to tally up the number of True values we observe:

count :: [Bool] -> Int
count = length . filter id

Now let’s create another posterior by conditioning on a few more observations:

posterior0 :: Model Double
posterior0 = invertWithAssitance count uniform bernoulli obs where
  obs =
    [True, True, True, False, True, True, False, True, True, True, True, False]

and collect another thousand samples from it. This would likely take an annoying amount of time without the use of our count function for assistance above:

> replicateM 1000 (simulate posterior0)

Note that with more information to condition on, we get a more informative posterior.

Conclusion

This is a really basic formulation - too basic to be useful in any meaningful way - but it illustrates some of the most important concepts in probabilistic programming. Representation, simulation, and inference.

I think it’s also particularly nice to do this in Haskell, rather than something like Python (which Dan used in his talk) - it provides us with a lot of extensible structure in a familiar framework for language hacking. It sort of demands you’re a fan of all these higher-kinded types and structured recursions and all that, but if you’re reading this blog then you’re probably in that camp anyway.

I’ll probably write a few more little articles like this over time. There are a ton of improvements that we can make to this basic setup - encoding independence, sampling via MCMC, etc. - and it might be fun to grow everything out piece by piece.

I’ve dropped the code from this post into this gist.

Randomness in Haskell

2016-10-01T00:00:00+04:00

Randomness is a constant nuisance point for Haskell beginners who may be coming from a language like Python or R. While in Python you can just get away with something like:

In [2]: numpy.random.rand(3)
Out[2]: array([ 0.61426175,  0.05309224,  0.38861597])

or in R:

> runif(3)
[1] 0.49473012 0.68436352 0.04135914

In Haskell, the situation is more complicated. It’s not too much worse when you get the hang of things, but it’s certainly one of those things that throws beginners for a loop - and for good reason.

In this article I want to provide a simple guide, with examples, for getting started and becoming comfortable with randomness in Haskell. Hopefully it helps!

I’m writing this from a hotel during my girlfriend’s birthday, so it’s being slapped together very rapidly with a kind of get-it-done attitude. If anything is unclear or you have any questions, feel free to shoot me a ping and I’ll try to improve it when I get a chance.

Randomness on Computers in General

Check out the R code I posted previously. If you just open R and type runif(3) on your machine, then odds are you’ll get a different triple of numbers than what I got above.

These numbers are being generated based on R’s global random number generator (RNG), which, absent any fiddling by the user, is initialized as needed based on the system time and ID of the R process. So: if you open up the R interpreter and call runif(3), then behind the scenes R will initialize the RNG based on the time and process ID, and then use a particular algorithm to generate random numbers based on that initialized value (called the ‘seed’).

These numbers aren’t truly random - they’re pseudo-random, which means they’re generated by a deterministic algorithm such that the resulting values appear random over time. The default algorithm used by R, for example, is the famous Mersenne Twister, which you can verify as follows:

> RNGkind()
[1] "Mersenne-Twister" "Inversion"

You can also set the seed yourself in R, using the set.seed function. Then if you type something like runif(3), R will use this initialized RNG rather than coming up with its own seed based on the time and process ID. Setting the seed allows you to reproduce operations involving pseudo-random numbers; just re-set the seed and perform the same operations again:

> set.seed(42)
> runif(3)
[1] 0.9148060 0.9370754 0.2861395
> set.seed(42)
> runif(3)
[1] 0.9148060 0.9370754 0.2861395

(It’s good practice to always initialize the RNG using some known seed before running an experiment, simulation, and so on.)

So the big thing to notice here, in any case, is that R uses a global RNG. It maintains the state of this RNG implicitly and behind the scenes. When you type runif(3), R consults this implicit RNG, gives you your pseudo-random numbers based on its value, and updates the global RNG without you needing to worry about any of this plumbing yourself. The same is generally true for randomness in most programming languages - Python, C, Ruby, and so on.

Explicit RNG Management

But let’s come back to Haskell. Haskell, unlike R or Python, is purely-functional. State, or effects in general, are never implicit in the same way that R updates its global RNG. We need to either explicitly pass around a RNG ourselves, or at least allow some explicit monad to do it for us.

Passing around a RNG manually is annoying, so in practice this means everyone uses a monad to handle RNG state. This means that one needs to be comfortable working with monadic code in order to practically use random numbers in Haskell, which presents a big hurdle for beginners who may have been able to ignore monads thus far on their Haskell journey.

Let’s see what I mean by all of this by going through a few examples. Make sure you have stack installed, and then grab a few libraries that we’ll make use of in the remainder of this post:

$ stack install random mwc-random primitive

The Really Annoying Method - Manual RNG Management

Let me demonstrate the simplest conceptual method for dealing with random numbers: manually grabbing and passing around a RNG without involving any monads whatsoever.

First, open up GHCi:

$ stack ghci

And let’s also get some quick preliminaries out of the way:

Prelude> :set prompt "> "
> import System.Random
> import Control.Monad
> let runif_pure = randomR (0 :: Double, 1)
> let runif n = replicateM n (randomRIO (0 :: Double, 1))
> let set_seed = setStdGen . mkStdGen

We’ll first use the basic System.Random module for illustration. To initialize a RNG, we can make one by providing the mkStdGen function with an integer seed:

> let rng = mkStdGen 42
> rng
43 1

We can use this thing to generate random numbers. A simple function to do that is randomR, which will generate pseudo-random values for some ordered type in a given range. We’ll use the runif_pure alias for it that we defined previously, just to make things look similar to the previous R example and also emphasize that this one is a pure function:

> runif_pure rng
(1.0663729393723398e-2,2060101257 2103410263)

You can see that we got back a pair of values, the first element of which is our random number 1.0663729393723398e-2. Cool. Let’s try to generate another:

> runif_pure rng
(1.0663729393723398e-2,2060101257 2103410263)

Hmm. We generated the same number again. This is because the value of rng hasn’t changed - it’s still the same value we made via mkStdGen 42. Since we’re using the same random number generator to generate a pseudo-random value, we get the same pseudo-random value.

If we want to make new random numbers, then we need to use a different generator. And the second element of the pair returned from our call to runif_pure is exactly that - an updated RNG that we can use to generate additional random numbers.

Let’s try that all again, using the generator we get back from the first function call as an input to the second:

> let (x, rng1) = runif_pure rng
> x
1.0663729393723398e-2
> let (y, rng2) = runif_pure rng1
> y
0.9827538369038856

Success!

I mean.. sort of. It works and all, and it does constitute a general-purpose solution. But manually binding updated RNG states to names and swapping those in for new values is still pretty annoying.

You could also generate an infinite list of random numbers using the randomRs function and just take from it as needed, but you still probably need to manage that list to make sure you don’t re-use any numbers. You kind of trade off managing the RNG for managing an infinite list of random numbers, which isn’t much better.

The Less-Annoying Method - Get A Monad To Do It

The good news is that we can offload the job of managing the RNG state to a monad. I won’t actually explain how that works in detail here - I think most people facing this problem are initially more concerned with getting something working, rather than deeply grokking monads off the bat - so I’ll just claim that we can get a monad to handle the RNG state for us, and that will hopefully (mostly) suffice for now.

Still rolling with the System.Random module for the time being, we’ll use the runif alias for the randomRIO function that we defined previously to generate some new random numbers:

> runif 3
[0.9873934690803106,0.3794382930121829,0.2285653405908732]
> runif 3
[0.7651878964537555,0.2623159001635825,0.7683468476766804]

Simpler! Notice we haven’t had to do anything with a generator manually - we just ask for random numbers and then get them, just like in R. And if we want to set the value of the RNG being used here, we can use the setStdGen function with an RNG that we’ve already created. Here let’s just use the set_seed alias we defined earlier, to mimic R’s set.seed function:

> set_seed 42
> runif 3
[1.0663729393723398e-2,0.9827538369038856,0.7042944187434987]
> set_seed 42
> runif 3
[1.0663729393723398e-2,0.9827538369038856,0.7042944187434987]

So things are similar to how they work in R here - we have a global RNG of sorts, and we can set its state as desired using the set_seed function. But since this is Haskell, the effects of creating and updating the generator state must still be explicit. And they are explicit - it’s just that they’re explicit in the type of runif:

> :t runif
runif :: Int -> IO [Double]

Note that runif returns a value that’s wrapped up in IO. This is how we indicate explicitly - at the type level - that something is being done with the generator in the background. IO is a monad, and it happens to be the thing that’s dealing with the generator for us here.

What this means for you, the practitioner, is that you can’t just mix values of some type a with values of type IO a willy-nilly. You may be writing a function f with type [Double] -> Double, where the input list of doubles is intended to be randomly-generated. But if you just go ahead and generate a list xs of random numbers, they’ll have type IO [Double], and you’ll stare in confusion at some type error from GHC when you try to apply f to xs.

Here’s what I mean. Take the example of just generating some random numbers and then summing them up. First, in R:

> xs = runif(3)
> sum(xs)
[1] 1.20353

And now in Haskell, using the same mechanism we tried earlier:

> let xs = runif 3
> :t xs
xs :: IO [Double]
> sum xs
<interactive>:16:1:
    No instance for (Num [Double]) arising from a use of ‘sum’
    In the expression: sum xs
    In an equation for ‘it’: it = sum xs

This means that to deal with the numbers we generate, we have to treat them a little differently than we would in R, or compared to the situation where we were managing the RNG explicitly in Haskell. Concretely: if we use a monad to manage the RNG for us, then the numbers we generate will be ‘tagged’ by the monad. So we need to do something or other to make those tagged numbers work with ‘untagged’ numbers, or functions designed to work with ‘untagged’ numbers.

This is where things get confusing for beginners. Here’s how we could add up some random numbers in GHCi:

> xs <- runif 3
> sum xs
1.512024272587933

We’ve used the <- symbol to bind the result of runif 3 to the name xs, rather than let xs = .... But this is sort of particular to running code in GHCi; if you try to do this in a generic Haskell function, you’ll possibly wind up with some more weird type errors. To do this in regular ol’ Haskell code, you need to both use <--style binding and also acknowledge the ‘tagged’ nature of randomly-generated values.

The crux is that, when you’re using a monad to generate random numbers in Haskell, you need to separate generating them from using them. Rather than try to explain what I mean here precisely, let’s rely on example, and implement a simple Metropolis sampler for illustration.

A Metropolis Sampler

The Metropolis algorithm will help you approximate expectations over certain probability spaces. Here’s how it works. Picture yourself strolling around some bumpy landscape; you want to walk around it in such a fashion that you visit regions of it with probability proportional to their altitude. To do that, you can repeatedly:

Pick a random point near your current location.
Compare your present altitude to the altitude of that point you picked. Calculate a probability based on their ratio.
Flip a coin where the chance of observing a head is equal to that probability. If you get a head, move to the location you picked. Otherwise, stay put.

Let’s implement it in Haskell, using a monadic random number generator to do so. This time we’re going to use mwc-random - a more industrial-strength randomness library that you can confidently use in production code.

mwc-random uses Marsaglia’s multiply-with-carry algorithm to generate pseudo-random numbers. It requires you to explicitly create and pass a RNG to functions that need to generate random numbers, but it uses a monad to update the RNG state itself. This winds up being pretty nice; let’s dive in to see.

Create a module called Metropolis.hs and get some imports out of the way:

module Metropolis where

import Control.Monad
import Control.Monad.Primitive
import System.Random.MWC as MWC
import System.Random.MWC.Distributions as MWC

Step One

The first thing we want to do is implement is point (1) from above:

Pick a random point near your current location.

We’ll just use a standard normal distribution of the appropriate dimension to do this - we just want to take a location, perturb it, and return the perturbed location.

propose :: [Double] -> Gen RealWorld -> IO [Double]
propose location rng = traverse (perturb rng) location where
  perturb gen x = MWC.normal x 1 gen

So at finer detail: we’re walking over the coordinates of the current location and generating a normally-distributed value centered at each coordinate. The MWC.normal function will do this for a given mean and standard deviation, and we can use the traverse function to walk over each coordinate.

Note that we pass a mwc-random RNG - the value with type Gen RealWorld - to the propose function. We need to supply this generator anywhere we want to generate random numbers, but we don’t need to manually worry about tracking and updating its state. The IO monad will do that for us. The resulting randomly-generated values will be tagged with IO, so we’ll need to deal with that appropriately.

Step Two

Now let’s implement point (2):

Compare your present altitude to the altitude of that point you picked. Calculate a probability based on their ratio.

So, we need a function that will compare the altitude of our current point to the altitude of a proposed point and compute a probability from that. The following will do: it takes a function that will compute a (log-scale) altitude for us, as well as the current and proposed locations, and returns a probability.

moveProbability :: ([Double] -> Double) -> [Double] -> [Double] -> Double
moveProbability altitude current proposed =
    whenNaN 0 (exp (min 0 (altitude proposed - altitude current)))
  where
    whenNaN val x
      | isNaN x   = val
      | otherwise = x

Step Three

Finally, the third step of the algorithm:

Flip a coin where the chance of observing a head is equal to that probability. If you get a head, move to the location you picked. Otherwise stay put.

So let’s get to it:

decide :: [Double] -> [Double] -> Double -> Gen RealWorld -> IO [Double]
decide current proposed prob rng = do
  accept <- MWC.bernoulli prob rng
  return $
    if   accept
    then proposed
    else current

Here we need to flip a coin, so we require a source of randomness again. The decide function thus takes another generator of type Gen RealWorld that we then supply to the MWC.bernoulli function, and the result - the final location - is once again wrapped in IO.

This function clearly demonstrates the typical way that you’ll deal with random numbers in Haskell code. decide is a monadic function, so it proceeds using do-notation. When you need to generate a random value - here we generate a random True or False value according to a Bernoulli distribution - you bind the result to a name using the <- symbol. Then afterwards, in the scope of the function, you can use the bound value as if it were pure. But the entire function must still return a ‘wrapped-up’ value that makes the effect of passing the generator explicit at the type level; right here, that means that the value will be wrapped up in IO.

Putting Everything Together

The final Metropolis transition is a combination of steps one through three. We can put them together like so:

metropolis :: ([Double] -> Double) -> [Double] -> Gen RealWorld -> IO [Double]
metropolis altitude current rng = do
  proposed <- propose current rng
  let prob = moveProbability altitude current proposed
  decide current proposed prob rng

Again, metropolis is monadic, so we start off with a do to make monadic programming easy on us. Whenever we need a random value, we bind the result of a random number-returning function using the <- notation.

The propose function returns a random location, so we bind its result to the name proposed using the <- symbol. The moveProbability function, on the other hand, is pure - so we bind that using a let prob = ... expression. The decide function returns a random value, so we can just plop it right on the end here. The entire result of the metropolis function is random, so it is wrapped up in IO.

The result of metropolis is just a single transition of the Metropolis algorithm, which involves doing this kind of thing over and over. If we do that, we observe a bunch of points that trace out a particular realization of a Markov chain, which we can generate as follows:

chain
  :: Int -> ([Double] -> Double) -> [Double] -> Gen RealWorld -> IO [[Double]]
chain epochs altitude origin rng = loop epochs [origin] where
  loop n history@(current:_)
    | n <= 0    = return history
    | otherwise = do
        next <- metropolis altitude current rng
        loop (n - 1) (next:history)

An Example

Now that we have our chain function, we can use it to trace out a collection of points visited on a realization of a Markov chain. Remember that we’re supposed to be wandering over some particular abstract landscape; here, let’s stroll over the one defined by the following function:

landscape :: [Double] -> Double
landscape [x0, x1] =
  -0.5 * (x0 ^ 2 * x1 ^ 2 + x0 ^ 2 + x1 ^ 2 - 8 * x0 - 8 * x1)

What we’ll now do is pick an origin to start from, wander over the landscape for some number of steps, and then print the resulting realization of the Markov chain to stdout. We’ll do all that through the following main function:

main :: IO ()
main = do
  rng <- MWC.createSystemRandom
  let origin = [-0.2, 0.3]
  trace <- chain 1000 landscape origin rng
  mapM_ print trace

Running that will dump a trace to stdout. If you clean it up and plot it, you’ll see that the visited points have traced out a rough approximation of the landscape:

Fini

Hopefully this gives a broad idea of how to go about using random numbers in Haskell. I’ve talked about:

Why randomness in Haskell isn’t as simple as randomness in (say) Python or R.
How to handle randomness in Haskell, either by manual generator management or by offloading that job to a monad.
How to get thing done when a monad manages the generator for you - separating random number generation from random number processing.
Doing all the above with an industrial-strength RNG, using a simple Metropolis algorithm as an example.

Hopefully the example gives you an idea of how to work with random numbers in practice.

I’ll be the first to admit that randomness in Haskell requires more work than randomness in a language like R, which to this day remains my go-to interactive data analysis language of choice. Using randomness effectively in Haskell requires a decent understanding of how to work with monadic code, even if one doesn’t quite understand monads entirely yet.

What I can say is that when one has developed some intuition for monads - acquiring a ‘feel’ for how to work with monadic functions and values - the difficulty and awkwardness drop off a bit, and working with randomness feels no different than working with any other effect.

Happy generating! I’ve dumped the code for the Metropolis example into a gist.

For a more production-quality Metropolis sampler, you can check out my mighty-metropolis library, which is a member of the declarative suite of MCMC algos.

On Measurability

2016-07-18T00:00:00+04:00

.. this one is pretty dry, I’ll admit. David Williams said it best:

.. Measure theory, that most arid of subjects when done for its own sake, becomes amazingly more alive when used in probability, not only because it is then applied, but also because it is immensely enriched.

Unfortunately for you, dear reader, we won’t be talking about probability.

Moving on. What does it mean for something to be measurable in the mathematical sense? Take some arbitrary collection $X$ and slap an appropriate algebraic structure $\mathcal{X}$ on it - usually an algebra or $\sigma$-algebra, etc. Then we can refer to a few different objects as ‘measurable’, going roughly as follows.

The elements of the structure $\mathcal{X}$ are called measurable sets. They’re called this because they can literally be assigned a notion of measure, whcih is a kind of generalized volume. If we’re just talking about some subset of $X$ out of the context of its structure then we can be pedantic and call it measurable with respect to $\mathcal{X}$, say. You could also call a set $\mathcal{X}$-measurable, to be similarly precise.

The product of the original collection and its associated structure $(X, \mathcal{X})$ is called a measurable space. It’s called that because it can be completed with a measuring function $\mu$ - itself called a measure - that assigns notions of measure to measurable sets.

Now take some other measurable space $(Y, \mathcal{Y})$ and consider a function $f$ from $X$ to $Y$. This is a measurable function if it satisfies the following technical requirement: that for any $\mathcal{Y}$-measurable set, its preimage under $f$ is an element of $\mathcal{X}$ (so: the preimage under $f$ is $\mathcal{X}$-measurable).

The concept of measurability for functions probably feels the least intuitive of the three - like one of those dry taxonomical classifications that we just have to keep on the books. The ‘make sure your function is measurable and everything will be ok’ heuristic will get you pretty far. But there is some good intuition available, if you want to look for it.

Here’s an example: define a set $X$ that consists of the elements $A$, $B$, and $C$. To talk about measurable functions, we first need to define our measurable sets. The de-facto default structure used for this is a $\sigma$-algebra, and we can always generate one from some underlying class of sets. Let’s do that from the following plain old partition that splits the original collection into a couple of disjoint ‘slices’:

\[H = \{\{A, B\}, \{C\}\}\]

The $\sigma$-algebra $\mathcal{X}$ generated from this partition will just be the partition itself, completed with the whole set $X$ and the empty set. To be clear, it’s the following:

\[\mathcal{X} = \left\{\{A, B, C\}, \{A, B\}, \{C\}, \emptyset\right\}\]

The resulting measurable space is $(X, \mathcal{X})$. So we could assign a notion of generalized volume to any element of $\mathcal{X}$, though I won’t actually worry about doing that here.

Now. Let’s think about some functions from $X$ to the real numbers, which we’ll assume to be endowed with a suitable $\sigma$-algebra of their own (one typically assumes the standard topology on $\mathbb{R}$ and the associated Borel $\sigma$-algebra).

How about this - a simple indicator function on the slice containing $C$:

\[f(x) = \begin{cases} 0, \, x \in \{A, B\} \\ 1, \, x \in \{C\} \end{cases}\]

Is it measurable? That’s easy to check. The preimage of $\{0\}$ is $\{A, B\}$, the preimage of $\{1\}$ is $\{C\}$, and the preimage of $\{0, 1\}$ is $X$ itself. Those are all in $\mathcal{X}$, and the preimage of the empty set is the empty set, so we’re good.

Ok. What about this one:

\[g(x) = \begin{cases} 0, \, x \in \{A\} \\ 1, \, x \in \{B\} \\ 2, \, x \in \{C\} \end{cases}\]

Check the preimage of $\{1, 2\}$ and you’ll find it’s $\{B, C\}$. But that’s not a member of $\mathcal{X}$, so $g$ is not measurable!

What happened here? Failing to satisfying technical requirements aside: what, intuitively, made $f$ measurable where $g$ wasn’t?

The answer is a problem of resolution. Look again at $\mathcal{X}$:

\[\left\{\{A, B, C\}, \{A, B\}, \{C\}, \emptyset\right\}\]

The structure $\mathcal{X}$ that we’ve endowed our collection $X$ with is too coarse to permit distinguishing between elements of the slice $\{A, B\}$. There is no measurable set $A$, nor a measurable set $B$ - just a measurable set $\{A, B\}$. And as a result, if we define a function that says something about either $A$ or $B$ without saying the same thing about the other, that function won’t be measurable. The function $f$ assigned the same value to both $A$ and $B$, so we didn’t have any problem there.

If we want to be able to distinguish between $A$ and $B$, we’ll need to equip $X$ with some structure that has a finer resolution. You can check that if you make a measurable space out of $X$ and its power set (the set of all subsets of $X$) then $g$ will be measurable there, for example.

So if we’re using partitions to define our measurable sets, we get a neat little property: for any measurable function, elements in the same slice of the partition must have the same value when passed through the function. So if you have a function $h : X \to H$ that takes an element to its respective slice in a partition, you know that $h(x_{0}) = h(x_{1})$ for any $x_{0}$, $x_{1}$ in $X$ implies that $f(x_{0}) = f(x_{1})$ for any measurable function $f$.

Addendum

Whipping together a measurable space using a $\sigma$-algebra generated by a partition of sets occurs naturally when we talk about correlated equilibrium, a solution concept in non-cooperative game theory. It’s common to say a function - in that context a correlated strategy - must be measurable ‘with respect to the partition’, which sort of elides the fact that we still need to generate a $\sigma$-algebra from it anyway.

Some oldschool authors (Halmos, at least) developed their measure theory using $\sigma$-rings, but this doesn’t seem very popular nowadays. Since a ring doesn’t require including the entire set $X$, you need to go through an awkward extra hoop when defining measurability on functions. But regardless, it’s interesting to think about what happens when one uses different structures to define measurable sets!

Making a Market

2016-04-20T00:00:00+04:00

Suppose you’re in the derivatives business. You are interested in making a market on some events; say, whether or not your friend Jay will win tomorrow night’s poker game, or that the winning pot will be at least USD 100. Let’s examine some rules about how you should do business if you want this venture to succeed.

What do I mean by ‘make a market’? I mean that you will be willing to buy and sell units of a particular security that will be redeemable from the seller at some particular value after tomorrow’s poker game has ended (you will be making a simple prediction market, in other words). You can make bid offers to buy securities at some price, and ask offers to sell securities at some price.

To keep things simple let’s say you’re doing this gratis; society rewards you extrinsically for facilitating the market - your friends will give you free pizza at the game, maybe - so you won’t levy any transaction fees for making trades. Also scarcity isn’t a huge issue, so you’re willing to buy or sell any finite number of securities.

Consider the possible outcomes of the game (one and only one of which must occur):

(A) Jay wins and the pot is at least USD 100.
(B) Jay wins and the pot is less than USD 100.
(C) Jay loses and the pot is at least USD 100.
(D) Jay loses and the pot is less than USD 100.

The securities you are making a market on pay USD 1 if an event occurs, and USD 0 otherwise. So: if I buy 5 securities on outcome $A$ from you, and outcome $A$ occurs, I’ll be able to go to you and redeem my securities for a total of USD 5. Alternatively, if I sell you 5 securities on outcome $A$, and outcome $A$ occurs, you’ll be able to come to me and redeem your securities for a total of USD 5.

Consider what that implies: as a market maker, you face the prospect of making hefty payments to customers who redeem valuable securities. For example, imagine the situation where you charge USD 0.50 for a security on outcome $A$, but outcome $A$ is almost certain to occur in some sense (Jay is a beast when it comes to poker and a lot of high rollers are playing); if your customers exclusively load up on 100 cheap securities on outcome $A$, and outcome $A$ occurs, then you stand to owe them a total payment of USD 100 against the USD 50 that they have paid for the securities. You thus have a heavy incentive to price your securities as accurately as possible - where ‘accurate’ means to minimize your expected loss.

It may always be the case, however, that it is difficult to price your securities accurately. For example, if some customer has more information than you (say, she privately knows that Jay is unusually bad at poker) then she potentially stands to profit from holding said information in lieu of your ignorance on the matter (and that of your prices). Such is life for a market maker. But there are particular prices you could offer - independent of any participant’s private information - that are plainly stupid or ruinous for you (a set of prices like this is sometimes called a Dutch book). Consider selling securities on outcome $A$ for the price of USD -1; then anyone who buys one of these securities not only stands to redeem USD 1 in the event outcome $A$ occurs, but also gains USD 1 simply from the act of buying the security in the first place.

Setting a negative price like this is irrational on your part; customers will realize an arbitrage opportunity on securities for outcome $A$ and will happily buy as many as they can get their hands on, to your ruin. In other words - and to nobody’s surprise - by setting a negative price, you can be made a sure loser in the market.

There are other prices you should avoid setting as well, if you want to avoid arbitrage opportunities like these. For starters:

For any outcome $E$, you must set the price of a security on $E$ to be at least USD 0.
For any certain outcome $E$, you must set the price of a security on $E$ to be exactly USD 1.

The first condition rules out negative prices, and the second ensures that your books balance when it comes time to settle payment for securities on a certain event.

What’s more, the price that you set on any given security doesn’t exist in isolation. Given the outcomes $A$, $B$, $C$, and $D$ listed previously, at least one must occur. So as per the second rule, the price of a synthetic derivative on the outcome “Jay wins or loses, and the pot is any value” must be set to USD 1. This places constraints on the prices that you can set for individual securities. It suffices that:

For any countable set of mutually exclusive outcomes $E_{1}, E_{2}, \ldots$, you must set the price of the security on outcome “$E_{1}$ or $E_{2}$ or..” to exactly the sum of the prices of the individual outcomes.

This eliminates the possibility that your customers will make you a certain loser by buying elaborate combinations of securities on different outcomes.

There are other rules that your prices must obey as well, but they fall out as corollaries of these three. If you broke any of them you’d also be breaking one of these.

It turns out that you cannot be made a sure loser if, and only if, your prices obey these three rules. That is:

If your prices follow these rules, then you will offer customers no arbitrage opportunities.
Any market absent of arbitrage opportunities must have prices that conform to these rules.

These prices are called coherent, and absence of coherence implies the existence of arbitrage opportunities for your customers.

But Why Male Models

The trick, of course, is that these prices correspond to probabilities, and the rules for avoiding arbitrage correspond to the standard Kolmogorov axioms of probability theory.

The consequence is that if your description of uncertain phenomena does not involve probability theory, or does not behave exactly like probability theory, then it is an incoherent representation of information you have about those phenomena.

As a result, probability theory should be your tool of choice when it comes to describing uncertain phenomena. Granted you may not have to worry about market making in return for pizza, but you’d like to be assured that there are no structural problems with your description.

Comments

This is a summary of the development of probability presented in Jay Kadane’s brilliant Principles of Uncertainty. The original argument was developed by de Finetti and Savage in the mid-20th century.

Kadane’s book makes for an exceptional read, and it’s free to boot. I recommend checking it out if it has flown under your radar.

An interesting characteristic of this development of probability is that there is no way to guarantee the nonexistence of arbitrage opportunities for a countably infinite number of purchased securities. That is: if you’re a market maker, you could be made a sure loser in the market when it came time for you to settle a countably infinite number of redemption claims. The quirk here is that you could also be made a sure winner as well; whether you win or lose with certainty depends on the order in which the claims are settled! (Fortunately this doesn’t tend to be an issue in practice.)

Thanks to Fredrik Olsen for reviewing a draft of this post.

References

flat-mcmc Update and v1.0.0 Release

2016-04-07T00:00:00+04:00

I’ve updated my old flat-mcmc library for ensemble sampling in Haskell and have pushed out a v1.0.0 release.

History

I wrote flat-mcmc in 2012, and it was the first serious-ish size project I attempted in Haskell. It’s an implementation of Goodman & Weare’s affine invariant ensemble sampler, a Monte Carlo algorithm that works by running a Markov chain over an ensemble of particles. It’s easy to get started with (there are no tuning parameters, etc.) and is sufficiently robust for a lot of purposes. The algorithm became somewhat famous in the astrostatistics community, where some of its members implemented it via the very nice and polished Python library, emcee.

The library has become my second-most starred repo on Github, with a whopping 10 stars as of this writing (the Haskell MCMC community is pretty niche, bro). Lately someone emailed me and asked if I wouldn’t mind pushing it to Stackage, so I figured it was due for an update and gave it a little modernizing along the way.

I’m currently on sabbatical and am traveling through Vietnam; I started the rewrite in Hanoi and finished it in Saigon, so it was a kind of nice side project to do while sipping coffees and the like during downtime.

What Is It

I wrote a little summary of the library in 2012, which you can still find tucked away on my personal site. Check that out if you’d like a description of the algorithm and why you might want to use it.

Since I wrote the initial version my astrostatistics-inclined friends David Huijser and Brendon Brewer wrote a paper about some limitations they discovered when using this algorithm in high-dimensional settings. So caveat emptor, buyer beware and all that.

In general this is an extremely easy-to-use algorithm that will probably get you decent samples from arbitrary targets without tedious tuning/fiddling.

What’s New

I’ve updated and standardized the API in line with my other MCMC projects huddled around the declarative library. That means that, like the others, there are two primary ways to use the library: via an mcmc function that will print a trace to stdout, or a flat transition operator that can be used to work with chains in memory.

Regrettably you can’t use the flat transition operator with others in the declarative ecosystem as it operates over ensembles, whereas the others are single-particle algorithms.

The README over at the Github repo contains a brief usage example. If there’s some feature you’d like to see or documentation/examples you could stand to have added then don’t hestitate to ping me and I’ll be happy to whip something up.

In the meantime I’ve pushed a new version to Hackage and added the library to Stackage, so it should show up in an LTS release soon enough.

Cheers!

Encoding Statistical Independence, Statically

2016-02-16T00:00:00+04:00

Applicative functors are useful for encoding context-free effects. This typically gets put to work around things like parsing or validation, but if you have a statistical bent then an applicative structure will be familiar to you as an encoder of independence.

In this article I’ll give a whirlwind tour of probability monads and algebraic freeness, and demonstrate that applicative functors can be used to represent independence between probability distributions in a way that can be verified statically.

I’ll use the following preamble for the code in the rest of this article. You’ll need the free and mwc-probability libraries if you’re following along at home:

{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE LambdaCase #-}

import Control.Applicative
import Control.Applicative.Free
import Control.Monad
import Control.Monad.Free
import Control.Monad.Primitive
import System.Random.MWC.Probability (Prob)
import qualified System.Random.MWC.Probability as MWC

Probability Distributions and Algebraic Freeness

Many functional programmers (though fewer statisticians) know that probability has a monadic structure. This can be expressed in multiple ways; the discrete probability distribution type found in the PFP framework, the sampling function representation used in the lambda-naught paper (and implemented here, for example), and even an obscure measure-based representation first described by Ramsey and Pfeffer, which doesn’t have a ton of practical use.

The monadic structure allows one to sequence distributions together. That is: if some distribution ‘foo’ has a parameter which itself has the probability distribution ‘bar’ attached to it, the compound distribution can be expressed by the monadic expression ‘bar »= foo’.

At a larger scale, monadic programs like this correspond exactly to what you’d typically see in a run-of-the-mill visualization of a probabilistic model:

In this classical kind of visualization the nodes represent probability distributions and the arrows describe the dependence structure. Translating it to a monadic program is mechanical: the nodes become monadic expressions and the arrows become binds. You’ll see a simple example in this article shortly.

The monadic structure of probability implies that it also has a functorial structure. Mapping a function over some probability distrubution will transform its support while leaving its probability density structure invariant in some sense. If the function ‘uniform’ defines a uniform probability distribution over the interval (0, 1), then the function ‘fmap (+ 1) uniform’ will define a probability distribution over the interval (1, 2).

I’ll come back to probability shortly, but the point is that probability distributions have a clear and well-defined algebraic structure in terms of things like functors and monads.

Recently free objects have become fashionable in functional programming. I won’t talk about it in detail here, but algebraic ‘freeness’ corresponds to a certain preservation of structure, and exploiting this kind of preserved structure is a useful technique for writing and interpreting programs.

Gabriel Gonzalez famously wrote about freeness in an oft-cited article about free monads, John De Goes wrote a compelling piece on the topic in the excellent A Modern Architecture for Functional Programming, and just today I noticed that Chris Stucchio had published an article on using Free Boolean Algebras for implementing a kind of constraint DSL. The last article included the following quote, which IMO sums up much of the raison d’être to exploit freeness in your day-to-day work:

.. if you find yourself re-implementing the same algebraic structure over and over, it might be worth asking yourself if a free version of that algebraic structure exists. If so, you might save yourself a lot of work by using that.

If a free version of some structure exists, then it embodies the ‘essence’ of that structure, and you can encode specific instances of it by just layering the required functionality over the free object itself.

A Type for Probabilistic Models

Back to probability. Since probability distributions are monads, we can use a free monad to encode them in a structure-preserving way. Here I’ll define a simple probability base functor for which each constructor is a particular ‘named’ probability distribution:

data ProbF r =
    BetaF Double Double (Double -> r)
  | BernoulliF Double (Bool -> r)
  deriving Functor

type Model = Free ProbF

Here we’ll only work with two simple named distributions - the beta and the Bernoulli - but the sky is the limit.

The ‘Model’ type wraps up this probability base functor in the free monad, ‘Free’. In this sense a ‘Model’ can be viewed as a program parameterized by the underlying probabilistic instruction set defined by ‘ProbF’ (a technique I described recently).

Expressions with the type ‘Model’ are terms in an embedded language. We can create some user-friendly ones for our beta-bernoulli language like so:

beta :: Double -> Double -> Model Double
beta a b = liftF (BetaF a b id)

bernoulli :: Double -> Model Bool
bernoulli p = liftF (BernoulliF p id)

Those primitive terms can then be used to construct expressions.

The beta and Bernoulli distributions enjoy an algebraic property called conjugacy that ensures (amongst other things) that the compound distribution formed by combining the two of them is analytically tractable. Here’s a parameterized coin constructed by doing just that:

coin :: Double -> Double -> Model Bool
coin a b = beta a b >>= bernoulli

By tweaking the parameters ‘a’ and ‘b’ we can bias the coin in particular ways, making it more or less likely to observe a head when it’s inspected.

A simple evaluator for the language goes like this:

eval :: PrimMonad m => Model a -> Prob m a
eval = iterM $ \case
  BetaF a b k    -> MWC.beta a b >>= k
  BernoulliF p k -> MWC.bernoulli p >>= k

‘iterM’ is a monadic, catamorphism-like recursion scheme that can be used to succinctly consume a ‘Model’. Here I’m using it to propagate uncertainty through the model by sampling from it ancestrally in a top-down manner. The ‘MWC.beta’ and ‘MWC.bernoulli’ functions are sampling functions from the mwc-probability library, and the resulting type ‘Prob m a’ is a simple probability monad type based on sampling functions.

To actually sample from the resulting ‘Prob’ type we can use the system’s PRNG for randomness. Here are some simple coin tosses with various biases as an example; you can mentally substitute ‘Head’ for ‘True’ if you’d like:

> gen <- MWC.createSystemRandom
> replicateM 10 $ MWC.sample (eval (coin 1 1)) gen
[False,True,False,False,False,False,False,True,False,False]
> replicateM 10 $ MWC.sample (eval (coin 1 8)) gen
[False,False,False,False,False,False,False,False,False,False]
> replicateM 10 $ MWC.sample (eval (coin 8 1)) gen
[True,True,True,False,True,True,True,True,True,True]

As a side note: encoding probability distributions in this way means that the other ‘forms’ of probability monad described previously happen to fall out naturally in the form of specific interpreters over the free monad itself. A measure-based probability monad could be achieved by using a different ‘eval’ function; the important monadic structure is already preserved ‘for free’:

measureEval :: Model a -> Measure a
measureEval = iterM $ \case
  BetaF a b k    -> Measurable.beta a b >>= k
  BernoulliF p k -> Measurable.bernoulli p >>= k

Independence and Applicativeness

So that’s all cool stuff. But in some cases the monadic structure is more than what we actually require. Consider flipping two coins and then returning them in a pair, for example:

flips :: Model (Bool, Bool)
flips = do
  c0 <- coin 1 8
  c1 <- coin 8 1
  return (c0, c1)

These coins are independent - they don’t affect each other whatsoever and enjoy the probabilistic/statistical property that formalizes that relationship. But the monadic program above doesn’t actually capture this independence in any sense; desugared, the program actually proceeds like this:

flips =
  coin 1 8 >>= \c0 ->
  coin 8 1 >>= \c1 ->
  return (c0, c1)

On the right side of any monadic bind we just have a black box - an opaque function that can’t be examined statically. Each monadic expression binds its result to the rest of the program, and we - programming ‘at the surface’ - can’t look at it without going ahead and evaluating it. In particular we can’t guarantee that the coins are truly independent - it’s just a mental invariant that can’t be transferred to an interpreter.

But this is the well-known motivation for applicative functors, so we can do a little better here by exploiting them. Applicatives are strictly less powerful than monads, so they let us write a probabilistic program that can guarantee the independence of expressions.

Let’s bring in the free applicative, ‘Ap’. I’ll define a type called ‘Sample’ by layering ‘Ap’ over our existing ‘Model’ type:

type Sample = Ap Model

So an expression with type ‘Sample’ is a free applicative over the ‘Model’ base functor. I chose the namesake because typically we talk about samples that are independent and identically-distributed draws from some probability distribution, though we could use ‘Ap’ to collect samples that are independently-but-not-identically distributed as well.

To use our existing embedded language terms with the free applicative, we can create the following helper function as an alias for ‘liftAp’ from ‘Control.Applicative.Free’:

independent :: f a -> Ap f a
independent = liftAp

With that in hand, we can write programs that statically encode independence. Here are the two coin flips from earlier (and if you’re applicative-savvy I’ll avoid using ‘liftA2’ here for clarity):

flips :: Sample (Bool, Bool)
flips = (,) <$> independent (coin 1 8) <*> independent (coin 8 1)

The applicative structure enforces exactly what we want: that no part of the effectful computation can depend on a previous part of the effectful computation. Or in probability-speak: that the distributions involved do not depend on each other in any way (they would be captured by the plate notation in the visualization shown previously).

To wrap up, we can reuse our previous evaluation function to interpret a ‘Sample’ into a value with the ‘Prob’ type:

evalIndependent :: PrimMonad m => Sample a -> Prob m a
evalIndependent = runAp eval

And from here it can just be evaluated as before:

> MWC.sample (evalIndependent flips) gen
(False,True)

Conclusion

That applicativeness embodies context-freeness seems to be well-known when it comes to parsing, but its relation to independence in probability theory seems less so.

Why might this be useful, you ask? Because preserving structure is mandatory for performing inference on probabilistic programs, and it’s safe to bet that the more structure you can capture, the easier that job will be.

In particular, algorithms for sampling from independent distributions tend to be simpler and more efficient than those useful for sampling from dependent distributions (where you might want something like Hamiltonian Monte Carlo or NUTS). Identifying independent components of a probabilistic program statically could thus conceptually simplify the task of sampling from some conditioned programs quite a bit - and that matters.

Enjoy! I’ve dumped the code from this article into a gist.

Time Traveling Recursion Schemes

2016-02-09T00:00:00+04:00

In Practical Recursion Schemes I talked about recursion schemes, describing them as elegant and useful patterns for expressing general computation. In that article I introduced a number of things relevant to working with the recursion-schemes package in Haskell.

In particular, I went over:

factoring the recursion out of recursive types using base functors and a fixed-point wrapper
the ‘Foldable’ and ‘Unfoldable’ typeclasses corresponding to recursive and corecursive data types, plus their ‘project’ and ‘embed’ functions respectively
the ‘Base’ type family that maps recursive types to their base functors
some of the most common and useful recursion schemes: cata, ana, para, and hylo.

In A Tour of Some Useful Recursive Types I went into further detail on ‘Fix’, ‘Free’, and ‘Cofree’ - three higher-kinded recursive types that are useful for encoding programs defined by some underlying instruction set.

I’ve also posted a couple of minor notes - I described the apo scheme in Sorting Slower With Style (as well as how to use recursion-schemes with regular Haskell lists) and chatted about monadic versions of the various schemes in Monadic Recursion Schemes.

Here I want to clue up this whole recursion series by briefly talking about two other recursion schemes - histo and futu - that work by looking at the past or future of the recursion respectively.

Here’s a little preamble for the examples to come:

{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE TypeFamilies #-}

import Control.Comonad.Cofree
import Control.Monad.Free
import Data.Functor.Foldable

Histomorphisms

Histomorphisms are terrifically simple - they just give you access to arbitrary previously-computed values of the recursion at any given point (its history, hence the namesake). They’re perfectly suited to dynamic programming problems, or anything where you might need to re-use intermediate computations later.

Histo needs a data structure to store the history of the recursion in. The the natural choice there is ‘Cofree’, which allows one to annotate recursive types with arbitrary metadata. Brian McKenna wrote a great article on making practical use of these kind of annotations awhile back.

But yeah, histomorphisms are very easy to use. Check out the following function that returns all the odd-indexed elements of a list:

oddIndices :: [a] -> [a]
oddIndices = histo $ \case
  Nil                           -> []
  Cons h (_ :< Nil)             -> [h]
  Cons h (_ :< Cons _ (t :< _)) -> h:t

The value to the left of a ‘:<’ constructor is an annotation provided by ‘Cofree’, and the value to right is the (similarly annotated) next step of the recursion. The annotations at any point are the previously computed values of the recursion corresponding to that point.

So in the case above, we’re just grabbing some elements from the input list and ignoring others. The algebra is saying:

if the input list is empty, return an empty list
if the input list has only one element, return that one-element list
if the input list has at least two elements, return the list built by cons-ing the first element together with the list computed two steps ago

The list computed two steps ago is available as the annotation on the constructor two steps down - I’ve matched it as ‘t’ in the last line of the above example. Like cata, histo works from the bottom-up.

A function that computes even indices is similar:

evenIndices :: [a] -> [a]
evenIndices = histo $ \case
  Nil                           -> []
  Cons _ (_ :< Nil)             -> []
  Cons _ (_ :< Cons h (t :< _)) -> h:t

Futumorphisms

Like histomorphisms, futumorphisms are also simple. They give you access to a particular computed part of the recursion at any given point.

However I’ll concede that the perceived simplicity probably comes with experience, and there is likely some conceptual weirdness to be found here. Just as histo gives you access to previously-computed values, futu gives you access to values that the recursion will compute in the future.

So yeah, that sounds crazy. But the reality is more mundane, if you’re familiar with the underlying concepts.

For histo, the recursion proceeds from the bottom up. At each point, the part of the recursive type you’re working with is annotated with the value of the recursion at that point (using ‘Cofree’), so you can always just reach back and grab it for use in the present.

With futu, the recursion proceeds from the top down. At each point, you construct an expression that can make use of a value to be supplied later. When the value does become available, you can use it to evaluate the expression.

A histomorphism makes use of ‘Cofree’ to do its annotation, so it should be no surprise that a futumorphism uses the dual structure - ‘Free’ - to create its expressions. The well-known ‘free monad’ is tremendously useful for working with small embedded languages. I also wrote about ‘Free’ in the same article mentioned previously, so I’ll link it again in case you want to refer to it.

As an example, we’ll use futu to implement the same two functions that we did for histo. First, the function that grabs all odd-indexed elements:

oddIndicesF :: [a] -> [a]
oddIndicesF = futu coalg where
  coalg list = case project list of
    Nil      -> Nil
    Cons x s -> Cons x $ do
      return $ case project s of
        Nil      -> s
        Cons _ t -> t

The coalgebra is saying the following:

if the input list is empty, return an empty list
if the input list has at least one element, construct an expression that will use a value to be supplied later.
if the supplied value turns out to be empty, return the one-element list.
if the supplied value turns out to have at least one more element, return the list constructed by skipping it, and using the value from another step in the future.

You can write that function more concisely, and indeed HLint will complain at you for writing it as I’ve done above. But I think this one makes it clear what parts are happening based on values to be supplied in the future. Namely, anything that occurs after ‘do’.

It’s kind of cool - you Enter The Monad™ and can suddenly work with values that don’t exist yet, while treating them as if they do.

Here’s futu-implemented ‘evenIndices’ for good measure:

evenIndicesF :: [a] -> [a]
evenIndicesF = futu coalg where
  coalg list = case project list of
    Nil      -> Nil
    Cons _ s -> case project s of
      Nil -> Nil
      Cons h t -> Cons h $ return t

Sort of a neat feature is that ‘Free’ part of the coalgebra can be written a little cleaner if you have ‘Free’-encoded embedded language terms floating around. Here are a couple of such terms, plus a ‘twiddle’ function that uses them to permute elements of an input list as they’re encountered:

nil :: Free (Prim [a]) b
nil = liftF Nil

cons :: a -> b -> Free (Prim [a]) b
cons h t = liftF (Cons h t)

twiddle :: [a] -> [a]
twiddle = futu coalg where
  coalg r = case project r of
    Nil      -> Nil
    Cons x l -> case project l of
      Nil      -> Cons x nil
      Cons h t -> Cons h $ cons x t

If you’ve been looking to twiddle elements of a recursive type then you’ve found a classy way to do it:

> take 20 $ twiddle [1..]
[2,1,4,3,6,5,8,7,10,9,12,11,14,13,16,15,18,17,20,19]

Enjoy! You can find the code from this article in this gist.

Monadic Recursion Schemes

2016-01-20T00:00:00+04:00

I have another few posts that I’d like to write before cluing up the whole recursion schemes kick I’ve been on. The first is a simple note about monadic versions of the schemes introduced thus far.

In practice you often want to deal with effectful versions of something like cata. Take a very simple embedded language, for example (“Hutton’s Razor”, with variables):

{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveFoldable #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE LambdaCase #-}

import           Control.Monad              ((<=<), liftM2)
import           Control.Monad.Trans.Class  (lift)
import           Control.Monad.Trans.Reader (ReaderT, ask, runReaderT)
import           Data.Functor.Foldable      hiding (Foldable, Unfoldable)
import qualified Data.Functor.Foldable      as RS (Foldable, Unfoldable)
import           Data.Map.Strict            (Map)
import qualified Data.Map.Strict            as Map

data ExprF r =
    VarF String
  | LitF Int
  | AddF r r
  deriving (Show, Functor, Foldable, Traversable)

type Expr = Fix ExprF

var :: String -> Expr
var = Fix . VarF

lit :: Int -> Expr
lit = Fix . LitF

add :: Expr -> Expr -> Expr
add a b = Fix (AddF a b)

(Note: Make sure you import ‘Data.Functor.Foldable.Foldable’ with a qualifier because GHC’s ‘DeriveFoldable’ pragma will become confused if there are multiple ‘Foldables’ in scope.)

Take proper error handling over an expression of type ‘Expr’ as an example; at present we’d have to write an ‘eval’ function as something like

eval :: Expr -> Int
eval = cata $ \case
  LitF j   -> j
  AddF i j -> i + j
  VarF _   -> error "free variable in expression"

This is a bit of a non-starter in a serious or production implementation, where errors are typically handled using a higher-kinded type like ‘Maybe’ or ‘Either’ instead of by just blowing up the program on the spot. If we hit an unbound variable during evaluation, we’d be better suited to return an error value that can be dealt with in a more appropriate place.

Look at the algebra used in ‘eval’; what would be useful is something like

monadicAlgebra = \case
  LitF j   -> return j
  AddF i j -> return (i + j)
  VarF v   -> Left (FreeVar v)

data Error =
    FreeVar String
  deriving Show

This won’t fly with cata as-is, and recursion-schemes doesn’t appear to include any support for monadic variants out of the box. But we can produce a monadic cata - as well as monadic versions of the other schemes I’ve talked about to date - without a lot of trouble.

To begin, I’ll stoop to a level I haven’t yet descended to and include a commutative diagram that defines a catamorphism:

To read it, start in the bottom left corner and work your way to the bottom right. You can see that we can go from a value of type ‘t’ to one of type ‘a’ by either applying ‘cata alg’ directly, or by composing a bunch of other functions together.

If we’re trying to define cata, we’ll obviously want to do it in terms of the compositions:

cata:: (RS.Foldable t) => (Base t a -> a) -> t ->  a
cata alg = alg . fmap (cata alg) . project

Note that in practice it’s typically more efficient to write recursive functions using a non-recursive wrapper, like so:

cata:: (RS.Foldable t) => (Base t a -> a) -> t ->  a
cata alg = c where c = alg . fmap c . project

This ensures that the function can be inlined. Indeed, this is the version that recursion-schemes uses internally.

To get to a monadic version we need to support a monadic algebra - that is, a function with type ‘Base t a -> m a’ for appropriate ‘t’. To translate the commutative diagram, we need to replace ‘fmap’ with ‘traverse’ (requiring a ‘Traversable’ instance) and the final composition with monadic (or Kleisli) composition:

The resulting function can be read straight off the diagram, modulo additional constraints on type variables. I’ll go ahead and write it directly in the inline-friendly way:

cataM
  :: (Monad m, Traversable (Base t), RS.Foldable t)
  => (Base t a -> m a) -> t ->  m a
cataM alg = c where
  c = alg <=< traverse c . project

Going back to the previous example, we can now define a proper ‘eval’ as follows:

eval :: Expr -> Either Error Int
eval = cataM $ \case
  LitF j   -> return j
  AddF i j -> return (i + j)
  VarF v   -> Left (FreeVar v)

This will of course work for any monad. A common pattern on an ‘eval’ function is to additionally slap on a ‘ReaderT’ layer to supply an environment, for example:

eval :: Expr -> ReaderT (Map String Int) (Either Error) Int
eval = cataM $ \case
  LitF j   -> return j
  AddF i j -> return (i + j)
  VarF v   -> do
    env <- ask
    case Map.lookup v env of
      Nothing -> lift (Left (FreeVar v))
      Just j  -> return j

And just an example of how that works:

> let open = add (var "x") (var "y")
> runReaderT (eval open) (Map.singleton "x" 1)
Left (FreeVar "y")
> runReaderT (eval open) (Map.fromList [("x", 1), ("y", 5)])
Right 6

You can follow the same formula to create the other monadic recursion schemes. Here’s monadic ana:

anaM
  :: (Monad m, Traversable (Base t), RS.Unfoldable t)
  => (a -> m (Base t a)) -> a -> m t
anaM coalg = a where
  a = (return . embed) <=< traverse a <=< coalg

and monadic para, apo, and hylo follow in much the same way:

paraM
  :: (Monad m, Traversable (Base t), RS.Foldable t)
  => (Base t (t, a) -> m a) -> t -> m a
paraM alg = p where
  p   = alg <=< traverse f . project
  f t = liftM2 (,) (return t) (p t)

apoM
  :: (Monad m, Traversable (Base t), RS.Unfoldable t)
  => (a -> m (Base t (Either t a))) -> a -> m t
apoM coalg = a where
  a = (return . embed) <=< traverse f <=< coalg
  f = either return a

hyloM
  :: (Monad m, Traversable t)
  => (t b -> m b) -> (a -> m (t a)) -> a -> m b
hyloM alg coalg = h
  where h = alg <=< traverse h <=< coalg

These are straightforward extensions from the basic schemes. A good exercise is to try putting together the commutative diagrams corresponding to each scheme yourself, and then use them to derive the monadic versions. That’s pretty fun to do for para and apo in particular.

If you’re using these monadic versions in your own project, you may want to drop them into a module like ‘Data.Functor.Foldable.Extended’ as recommended by my colleague Jasper Van der Jeugt. Additionally, there is an old issue floating around on the recursion-schemes repo that proposes adding them to the library itself. So maybe they’ll turn up in there eventually.

Sorting Slower with Style

2016-01-19T00:00:00+04:00

I previously wrote about implementing merge sort using recursion schemes. By using a hylomorphism we could express the algorithm concisely and true to its high-level description.

Insertion sort can be implemented in a similar way - this time by putting one recursion scheme inside of another.

Read on for details.

Apomorphisms

These guys don’t seem to get a lot of love in the recursion scheme tutorial du jour, probably because they might be the first scheme you encounter that looks truly weird on first glance. But apo is really not bad at all - I’d go so far as to call apomorphisms straightforward and practical.

So: if you remember your elementary recursion schemes, you can say that apo is to ana as para is to cata. A paramorphism gives you access to a value of the original input type at every point of the recursion; an apomorphism lets you terminate using a value of the original input type at any point of the recursion.

This is pretty useful - often when traversing some structure you just want to bail out and return some value on the spot, rather than continuing on recursing needlessly.

A good introduction is the toy ‘mapHead’ function that maps some other function over the head of a list and leaves the rest of it unchanged. Let’s first rig up a hand-rolled list type to illustrate it on:

{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE TypeFamilies #-}

import Data.Functor.Foldable

data ListF a r =
    ConsF a r
  | NilF
  deriving (Show, Functor)

type List a = Fix (ListF a)

fromList :: [a] -> List a
fromList = ana coalg . project where
  coalg Nil        = NilF
  coalg (Cons h t) = ConsF h t

(I’ll come back to the implementation of ‘fromList’ later, but for now you can see it’s implemented via an anamorphism.)

Example One

Here’s ‘mapHead’ for our hand-rolled list type, implemented via apo:

mapHead :: (a -> a) -> List a -> List a
mapHead f = apo coalg . project where
  coalg NilF        = NilF
  coalg (ConsF h t) = ConsF (f h) (Left t)

Before I talk you through it, here’s a trivial usage example:

> fromList [1..3]
Fix (ConsF 1 (Fix (ConsF 2 (Fix (ConsF 3 (Fix NilF))))))
> mapHead succ (fromList [1..3])
Fix (ConsF 2 (Fix (ConsF 2 (Fix (ConsF 3 (Fix NilF))))))

Now. Take a look at the coalgebra involved in writing ‘mapHead’. It has the type ‘a -> Base t (Either t a)’, which for our hand-rolled list case simplifies to ‘a -> ListF a (Either (List a) a)’.

Just as a reminder, you can check this in GHCi using the ‘:kind!’ command:

> :set -XRankNTypes
> :kind! forall a. a -> Base (List a) (Either (List a) a)
forall a. a -> Base (List a) (Either (List a) a) :: *
= a -> ListF a (Either (List a) a)

So, inside any base functor on the right hand side we’re going to be dealing with some ‘Either’ values. The ‘Left’ branch indicates that we’re going to terminate the recursion using whatever value we pass, whereas the ‘Right’ branch means we’ll continue recursing as per normal.

In the case of ‘mapHead’, the coalgebra is saying:

deconstruct the list; if it has no elements just return an empty list
if the list has at least one element, return the list constructed by prepending ‘f h’ to the existing tail.

Here the ‘Left’ branch is used to terminate the recursion and just return the existing tail on the spot.

Example Two

That was pretty easy, so let’s take it up a notch and implement list concatenation:

cat :: List a -> List a -> List a
cat l0 l1 = apo coalg (project l0) where
  coalg NilF = case project l1 of
    NilF      -> NilF
    ConsF h t -> ConsF h (Left t)

  coalg (ConsF x l) = case project l of
    NilF      -> ConsF x (Left l1)
    ConsF h t -> ConsF x (Right (ConsF h t))

This one is slightly more involved, but the principles are almost entirely the same. If both lists are empty we just return an empty list, and if the first list has at most one element we return the list constructed by jamming the second list onto it. The ‘Left’ branch again just terminates the recursion and stops everything there.

If both lists are nonempty? Then we actually do some work and recurse, which is what the ‘Right’ branch indicates.

So hopefully you can see there’s nothing too weird going on - the coalgebras are really simple once you get used to the Either constructors floating around in there.

Paramorphisms involve an algebra that gives you access to a value of the original input type in a pair - a product of two values. Since apomorphisms are their dual, it’s no surprise that you can give them a value of the original input type using ‘Either’ - a sum of two values.

Insertion Sort

So yeah, my favourite example of an apomorphism is for implementing the ‘inner loop’ of insertion sort, a famous worst-case $O(n^2)$ comparison-based sort. Granted that insertion sort itself is a bit of a toy algorithm, but the pattern used to implement its internals is pretty interesting and more broadly applicable.

This animation found on Wikipedia illustrates how insertion sort works:

We’ll actually be doing this thing in reverse - starting from the right-hand side and scanning left - but here’s the inner loop that we’ll be concerned with: if we’re looking at two elements that are out of sorted order, slide the offending element to where it belongs by pushing it to the right until it hits either a bigger element or the end of the list.

As an example, picture the following list:

[3, 1, 1, 2, 4, 3, 5, 1, 6, 2, 1]

The first two elements are out of sorted order, so we want to slide the 3 rightwards along the list until it bumps up against a larger element - here that’s the 4.

The following function describes how to do that in general for our hand-rolled list type:

coalg NilF        = NilF
coalg (ConsF x l) = case project l of
  NilF          -> ConsF x (Left l)
  ConsF h t
    | x <= h    -> ConsF x (Left l)
    | otherwise -> ConsF h (Right (ConsF x t))

It says:

deconstruct the list; if it has no elements just return an empty list
if the list has only one element, or has at least two elements that are in sorted order, terminate with the original list by passing the tail of the list in the ‘Left’ branch
if the list has at least two elements that are out of sorted order, swap them and recurse using the ‘Right’ branch

And with that in place, we can use an apomorphism to put it to work:

knockback :: Ord a => List a -> List a
knockback = apo coalg . project where
  coalg NilF        = NilF
  coalg (ConsF x l) = case project l of
    NilF          -> ConsF x (Left l)
    ConsF h t
      | x <= h    -> ConsF x (Left l)
      | otherwise -> ConsF h (Right (ConsF x t))

Check out how it works on our original list, slotting the leading 3 in front of the 4 as required. I’ll use a regular list here for readability:

> let test = [3, 1, 1, 2, 4, 3, 5, 1, 6, 2, 1]
> knockbackL test
[1, 1, 2, 3, 4, 3, 5, 1, 6, 2, 1]

Now to implement insertion sort we just want to do this repeatedly like in the animation above.

This isn’t something you’d likely notice at first glance, but check out the type of ‘knockback . embed’:

> :t knockback . embed
knockback . embed :: Ord a => ListF a (List a) -> List a

That’s an algebra in the ‘ListF a’ base functor, so we can drop it into cata:

insertionSort :: Ord a => List a -> List a
insertionSort = cata (knockback . embed)

And voila, we have our sort.

If it’s not clear how the thing works, you can visualize the whole process as working from the back of the list, knocking back unsorted elements and recursing towards the front like so:

[]
[1]
[2, 1] -> [1, 2]
[6, 1, 2] -> [1, 2, 6]
[1, 1, 2, 6]
[5, 1, 1, 2, 6] -> [1, 1, 2, 5, 6]
[3, 1, 1, 2, 5, 6] -> [1, 1, 2, 3, 5, 6]
[4, 1, 1, 2, 3, 5, 6] -> [1, 1, 2, 3, 4, 5, 6]
[2, 1, 1, 2, 3, 4, 5, 6] -> [1, 1, 2, 2, 3, 4, 5, 6]
[1, 1, 1, 2, 2, 3, 4, 5, 6]
[1, 1, 1, 1, 2, 2, 3, 4, 5, 6]
[3, 1, 1, 1, 1, 2, 2, 3, 4, 5, 6] -> [1, 1, 1, 1, 2, 2, 3, 3, 4, 5, 6]
[1, 1, 1, 1, 2, 2, 3, 3, 4, 5, 6]

Wrapping Up

And that’s it! If you’re unlucky you may be sorting asymptotically worse than if you had used mergesort. But at least you’re doing it with style.

The ‘mapHead’ and ‘cat’ examples come from the unreadable Vene and Uustalu paper that first described apomorphisms. The insertion sort example comes from Tim Williams’s excellent recursion schemes talk.

As always, I’ve dumped the code for this article into a gist.

Addendum: Using Regular Lists

You’ll note that the ‘fromList’ and ‘knockbackL’ functions above operate on regular Haskell lists. The short of it is that it’s easy to do this; recursion-schemes defines a data family called ‘Prim’ that basically endows lists with base functor constructors of their own. You just need to use ‘Nil’ in place of ‘[]’ and ‘Cons’ in place of ‘(:)’.

Here’s insertion sort implemented in the same way, but for regular lists:

knockbackL :: Ord a => [a] -> [a]
knockbackL = apo coalg . project where
  coalg Nil        = Nil
  coalg (Cons x l) = case project l of
    Nil           -> Cons x (Left l)
    Cons h t
      | x <= h    -> Cons x (Left l)
      | otherwise -> Cons h (Right (Cons x t))

insertionSortL :: Ord a => [a] -> [a]
insertionSortL = cata (knockbackL . embed)

Yo Dawg We Heard You Like Derivatives

2016-01-08T00:00:00+04:00

I noticed this article by Tom Ellis today that provides an excellent ‘demystified’ introduction to automatic differentiation. His exposition is exceptionally clear and simple.

Hopefully not in the spirit of re-mystifying things too much, I wanted to demonstrate that this kind of forward-mode automatic differentiation can be implemented using a catamorphism, which cleans up the various let statements found in Tom’s version (at the expense of slightly more pattern matching).

Let me first duplicate his setup using the standard recursion scheme machinery:

{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE LambdaCase #-}

import Data.Functor.Foldable

data ExprF r =
    VarF
  | ZeroF
  | OneF
  | NegateF r
  | SumF r r
  | ProductF r r
  | ExpF r
  deriving (Show, Functor)

type Expr = Fix ExprF

Since my expression type uses a fixed-point wrapper I’ll define my own embedded language terms to get around it:

var :: Expr
var = Fix VarF

zero :: Expr
zero = Fix ZeroF

one :: Expr
one = Fix OneF

neg :: Expr -> Expr
neg x = Fix (NegateF x)

add :: Expr -> Expr -> Expr
add a b = Fix (SumF a b)

prod :: Expr -> Expr -> Expr
prod a b = Fix (ProductF a b)

e :: Expr -> Expr
e x = Fix (ExpF x)

To implement a corresponding eval function we can use a catamorphism:

eval :: Double -> Expr -> Double
eval x = cata $ \case
  VarF         -> x
  ZeroF        -> 0
  OneF         -> 1
  NegateF a    -> negate a
  SumF a b     -> a + b
  ProductF a b -> a * b
  ExpF a       -> exp a

Very clear. We just match things mechanically.

Now, symbolic differentiation. If you refer to the original diff function you’ll notice that in cases like Product or Exp there are uses of both an original expression and also its derivative. This can be captured simply by a paramorphism:

diff :: Expr -> Expr
diff = para $ \case
  VarF                     -> one
  ZeroF                    -> zero
  OneF                     -> zero
  NegateF (_, x')          -> neg x'
  SumF (_, x') (_, y')     -> add x' y'
  ProductF (x, x') (y, y') -> add (prod x y') (prod x' y)
  ExpF (x, x')             -> prod (e x) x'

Here the primes indicate derivatives in the usual fashion, and the standard rules of differentiation are self-explanatory.

For automatic differentiation we just do sort of the same thing, except we’re also also going to lug around the evaluated function value itself at each point and evaluate to doubles instead of other expressions.

It’s worth noting here: why doubles? Because the expression type that we’ve defined has no notion of sharing, and thus the expressions will blow up à la diff (to see what I mean, try printing the analogue of diff bigExpression in GHCi). This could probably be mitigated by incorporating sharing into the embedded language in some way, but that’s a topic for another post.

Anyway, a catamorphism will do the trick:

ad :: Double -> Expr -> (Double, Double)
ad x = cata $ \case
  VarF                     -> (x, 1)
  ZeroF                    -> (0, 0)
  OneF                     -> (1, 0)
  NegateF (x, x')          -> (negate x, negate x')
  SumF (x, x') (y, y')     -> (x + y, x' + y')
  ProductF (x, x') (y, y') -> (x * y, x * y' + x' * y)
  ExpF (x, x')             -> (exp x, exp x * x')

Take a look at the pairs to the right of the pattern matches; the first element in each is just the corresponding term from eval, and the second is just the corresponding term from diff (made ‘Double’-friendly). The catamorphism gives us access to all the terms we need, and we can avoid a lot of work on the right-hand side by doing some more pattern matching on the left.

Some sanity checks to make sure that these functions match up with Tom’s:

*Main> map (snd . (`ad` testSmall)) [0.0009, 1.0, 1.0001]
[0.12254834896191881,1.0,1.0003000600100016]
*Main> map (snd . (`ad` testBig)) [0.00009, 1.0, 1.00001]
[3.2478565715996756e-6,1.0,1.0100754777229357]

UPDATE:

I had originally defined ad using a paramorphism but noticed that we can get by just fine with cata.

A Tour of Some Useful Recursive Types

2015-12-09T00:00:00+04:00

I’m presently at NIPS and so felt like writing about some appropriate machine learning topic, but along the way I wound up talking about parameterized recursive types, and here we are. Enjoy!

One starts to see common ‘shapes’ in algebraic data types after working with them for a while. Take the natural numbers and a standard linked list, for example:

data Natural =
    One
  | Succ Natural

data List a =
    Empty
  | Cons a (List a)

These are similar in some sense. There are some differences - a list has an additional type parameter, and each recursive point in the list is tagged with a value of that type - but the nature of the recursion in each is the same. There is a single recursive point wrapped up in a single constructor, plus a single base case.

Consider a recursive type that is parameterized by a functor with kind ‘* -> *’, such that the kind of the resulting type is something like ‘(* -> *) -> *’ or ‘(* -> *) -> * -> *’ or so on. It’s interesting to look at the ‘shapes’ of some useful types like this and see what kind of similarities and differences in recursive structure that we can find.

In this article we’ll look at three such recursive types: ‘Fix’, ‘Free’, and ‘Cofree’. I’ll demonstrate that each can be viewed as a kind of program parameterized by some underlying instruction set.

Fix

To start, let’s review the famous fixed-point type ‘Fix’. I’ve talked about it before, but will go into a bit more detail here.

{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE UndecideableInstances #-}

newtype Fix f = Fix (f (Fix f))

deriving instance (Show (f (Fix f))) => Show (Fix f)

Note: I’ll omit interpreter output for examples throughout this article, but feel free to try the code yourself in GHCi. I’ll post some gists at the bottom. The above code block also contains some pragmas that you can ignore; they’re just there to help GHC derive some instances for us.

Anyway. ‘Fix’ is in some sense a template recursive structure. It relies on some underlying functor ‘f’ to define the scope of recursion that you can expect a value with type ‘Fix f’ to support. There is the degenerate constant case, for example, which supports no recursion:

data DegenerateF r = DegenerateF
  deriving (Functor, Show)

type Degenerate = Fix DegenerateF

degenerate :: Degenerate
degenerate = Fix DegenerateF

Then you have the case like the one below, where only an infinitely recursive value exists:

newtype InfiniteF r = InfiniteF r
  deriving (Functor, Show)

type Infinite = Fix InfiniteF

infinite :: Infinite
infinite = Fix (InfiniteF infinite)

But in practice you’ll have something in between; a type with at least one recursive point or ‘running’ case and also at least one base or ‘terminating’ case. Take the natural numbers, for example:

data NatF r =
    OneF
  | SuccF r
  deriving (Functor, Show)

type Nat = Fix NatF

one :: Nat
one = Fix OneF

succ :: Nat -> Nat
succ = Fix . SuccF

Here ‘NatF’ provides both a ‘running’ case - ‘SuccF’ - and a ‘terminating’ case in - ‘OneF’. ‘Fix’ just lets ‘NatF’ do whatever it wants, having no say of its own about termination. In fact, we could have defined ‘Fix’ like this:

data Program f = Running (f (Program f))

Indeed, you can think of ‘Fix’ as defining a program that runs until ‘f’ decides to terminate. In turn, you can think of ‘f’ as an instruction set for the program. The whole shebang of ‘Fix f’ may only terminate if ‘f’ contains a terminating instruction.

Here’s a simple set of instructions, for example:

data Instruction r =
    Increment r
  | Decrement r
  | Terminate
  deriving (Functor, Show)

increment :: Program Instruction -> Program Instruction
increment = Running . Increment

decrement :: Program Instruction -> Program Instruction
decrement = Running . Decrement

terminate :: Program Instruction
terminate = Running Terminate

And we can write a sort of stack-based program like so:

program :: Program Instruction
program =
    increment
  . increment
  . decrement
  $ terminate

Richness of ‘Fix’

It’s worthwhile to review two functions that are useful for working with ‘Fix’, unimaginatively named ‘fix’ and ‘unfix’:

fix :: f (Fix f) -> Fix f
fix = Fix

unfix :: Fix f -> f (Fix f)
unfix (Fix f) = f

You can think of them like so: ‘fix’ embeds a value of type ‘f’ into a recursive structure by adding a new layer of recursion, while ‘unfix’ projects a value of type ‘f’ out of a recursive structure by peeling back a layer of recursion.

This is a pretty rich recursive structure - we have a guarantee that we can always embed into or project out of something with type ‘Fix f’, no matter what ‘f’ is.

Free

Next up is ‘Free’, which is really just ‘Fix’ with some added structure. It is defined as follows:

data Free f a =
    Free (f (Free f a))
  | Pure a
  deriving Functor

deriving instance (Show a, Show (f (Free f a))) => Show (Free f a)

The ‘Free’ constructor has an analogous definition to the ‘Fix’ constructor, meaning we can use ‘Free’ to implement the same things we did previously. Here are the natural numbers redux, for example:

type NatFree = Free NatF

oneFree :: NatFree a
oneFree = Free OneF

succFree :: NatFree a -> NatFree a
succFree = Free . SuccF

There’s also another branch here called ‘Pure’, though, that just bluntly wraps a value of type ‘a’, and has nothing to do with the parameter ‘f’. This has an interesting consequence: it means that ‘Free’ can have an opinion of its own about termination, regardless about what ‘f’ might decree:

type NotSoInfinite = Free InfiniteF

notSoInfinite :: NotSoInfinite ()
notSoInfinite = Free (InfiniteF (Free (InfiniteF (Pure ()))))

(Note that here I’ve returned the value of type unit when terminating under the ‘Pure’ branch, but you could pick whatever else you’d like.)

You’ll recall that ‘InfiniteF’ provides no terminating instruction, and left to its own devices will just recurse endlessly.

So: instead of being forced to choose a branch of the underlying functor to recurse on, ‘Free’ can just bail out on a whim and return some value wrapped up in ‘Pure’. We could have defined the whole type like this:

data Program f a =
    Running (f (Program f a))
  | Terminated a
  deriving Functor

Again, it’s ‘Fix’ with more structure. It’s a program that runs until ‘f’ decides to terminate, or that terminates and returns a value of type ‘a’

As a quick illustration, take our simple stack-based instruction set again. We can define the following embedded language terms:

increment :: Program Instruction a -> Program Instruction a
increment = Running . Increment

decrement :: Program Instruction a -> Program Instruction a
decrement = Running . Decrement

terminate :: Program Instruction a
terminate = Running Terminate

sigkill :: Program f Int
sigkill = Terminated 1

So note that ‘sigkill’ is independent of whatever instruction set we’re working with. We can thus write another simple program like before, except this time have ‘sigkill’ terminate it:

program :: Program Instruction Int
program =
    increment
  . increment
  . decrement
  $ sigkill

Richness of ‘Free’

Try to define the equivalent versions of ‘fix’ and ‘unfix’ for ‘Free’. The equivalent to ‘fix’ is easy:

free :: f (Free f a) -> Free f a
free = Free

You’ll hit a wall, though, if you want to implement the (total) analogue to ‘unfix’. One wants a function of type ‘Free f a -> f (Free f a)’, but the existence of the ‘Pure’ branch makes this impossible to implement totally. In general there is not going to be an ‘f’ to pluck out:

unfree :: Free f a -> f (Free f a)
unfree (Free f) = f
unfree (Pure a) = error "kaboom"

The recursion provided by ‘Free’ is thus a little less rich than that provided by ‘Fix’. With ‘Fix’ one can always project a value out of its recursive structure - but that’s not the case with ‘Free’.

It’s well-known that ‘Free’ is monadic, and indeed it’s usually called the ‘free monad’. The namesake ‘free’ comes from an algebraic definition; roughly, a free ‘foo’ is a ‘foo’ that satisfies the minimum possible constraints to make it a ‘foo’, and nothing else. Check out the slides from Dan Piponi’s excellent talk from Bayhac a few years back for a deeper dive on algebraic freeness.

Cofree

‘Cofree’ is also like ‘Fix’, but again with some extra structure. It can be defined as follows:

data Cofree f a = Cofree a (f (Cofree f a))
  deriving Functor

deriving instance (Show a, Show (f (Cofree f a))) => Show (Cofree f a)

Again, part of the definition - the second field of the ‘Cofree’ constructor - looks just like ‘Fix’. So predictably we can do a redux-redux of the natural numbers using ‘Cofree’:

type NatCofree = Cofree NatF

oneCofree :: NatCofree ()
oneCofree = Cofree () OneF

succFree :: NatCofree () -> NatCofree ()
succFree f = Cofree () (SuccF f)

(Note that here I’ve again used unit to fill in the first field - you could of course choose whatever you’d like.)

This looks a lot like ‘Free’, and in fact it’s the categorical dual of ‘Free’. Whereas ‘Free’ is a sum type with two branches, ‘Cofree’ is a product type with two fields. In the case of ‘Free’, we could have a program that either runs an instruction from a set ‘f’, or terminates with a value having type ‘a’. In the case of ‘Cofree’, we have a program that runs an instruction from a set ‘f’ and returns a value of type ‘a’.

A ‘Free’ value thus contains at most one recursive point wrapping the value with type ‘a’, while a ‘Cofree’ value contains potentially infinite recursive points - each one of which is tagged with a value of type ‘a’.

Rolling with the ‘Program’ analogy, we could have written this alternate definition for ‘Cofree’:

data Program f a = Program {
    annotation :: a
  , running    :: f (Program f a)
  } deriving Show

A ‘Cofree’ value is thus a program in which every instruction is annotated with a value of type ‘a’. This means that, unlike ‘Free’, it can’t have its own opinion on termination. Like ‘Fix’, it has to let ‘f’ decide how to do that.

We’ll use the stack-based instruction set example to wrap up. Here we can annotate instructions with progress about how many instructions remain to execute. First our new embedded language terms:

increment :: Program Instruction Int -> Program Instruction Int
increment p = Program (remaining p) (Increment p)

decrement :: Program Instruction Int -> Program Instruction Int
decrement p = Program (remaining p) (Decrement p)

terminate :: Program Instruction Int
terminate = Program 0 Terminate

Notice that two of these terms use a helper function ‘remaining’ that counts the number of instructions left in the program. It’s defined as follows:

remaining :: Program Instruction Int -> Int
remaining = loop where
  loop (Program a f) = case f of
    Increment p -> succ (loop p)
    Decrement p -> succ (loop p)
    Terminate   -> succ a

And we can write our toy program like so:

program :: Program Instruction Int
program =
    increment
  . increment
  . decrement
  $ terminate

Evaluate it in GHCi to see what the resulting value looks like.

Richness of ‘Cofree’

If you try and implement the ‘fix’ and ‘unfix’ analogues for ‘Cofree’ you’ll rapidly infer that we have the opposite situation to ‘Free’ here. Implementing the ‘unfix’ analogue is easy:

uncofree :: Cofree f a -> f (Cofree f a)
uncofree (Cofree _ f) = f

But implementing a total function corresponding to ‘fix’ is impossible - we can’t just come up with something of arbitrary type ‘a’ to tag an instruction ‘f’ with, so, like before, we can’t do any better than define something partially:

cofree :: f (Cofree f a) -> Cofree f a
cofree f = Cofree (error "kaboom") f

Just as how ‘Free’ forms a monad, ‘Cofree’ forms a comonad. It’s thus known as the ‘cofree comonad’, though I can’t claim to really have any idea what the algebraic notion of ‘cofreeness’ captures, exactly.

Wrapping Up

So: ‘Fix’, ‘Free’, and ‘Cofree’ all share a similar sort of recursive structure that make them useful for encoding programs, given some instruction set. And while their definitions are similar, ‘Fix’ supports the richest recursion of the three in some sense - it can both ‘embed’ things into and ‘project’ things out of its recursive structure, while ‘Free’ supports only embedding and ‘Cofree’ supports only projecting.

This has a practical implication: it means one can’t make use of certain recursion schemes for ‘Free’ and ‘Cofree’ in the same way that one can for ‘Fix’. There do exist analogues, but they’re sort of out-of-scope for this post.

I haven’t actually mentioned any truly practical uses of ‘Free’ and ‘Cofree’ here, but they’re wonderful things to keep in your toolkit if you’re doing any work with embedded languages, and I’ll likely write more about them in the future. In the meantime, Dave Laing wrote an excellent series of posts on ‘Free’ and ‘Cofree’ that are more than worth reading. They go into much more interesting detail than I’ve done here - in particular he details a nice pairing that exists between ‘Free’ and ‘Cofree’ (also discussed by Dan Piponi), plus a whack of examples.

You can also find industrial-strength infrastructure for both ‘Free’ and ‘Cofree’ in Edward Kmett’s excellent free library, and for ‘Fix’ in recursion-schemes.

I’ve dumped the code for this article into a few gists. Here’s one of everything excluding the running ‘Program’ examples, and here are the corresponding ‘Program’ examples for the Fix, Free, and Cofree cases respectively.

Thanks to Fredrik Olsen for review and great feedback.

Sorting with Style

2015-12-02T00:00:00+04:00

Merge sort is a famous comparison-based sorting algorithm that starts by first recursively dividing a collection of orderable elements into smaller subcollections, and then finishes by recursively sorting and merging the smaller subcollections together to reconstruct the (now sorted) original.

A clear implementation of mergesort should by definition be as faithful to that high-level description as possible. We can get pretty close to that using the whole recursion schemes business that I’ve talked about in the past. Near the end of that article I briefly mentioned the idea of implementing mergesort via a hylomorphism, and here I just want to elaborate on that a little.

Start with a collection of orderable elements. We can divide the collection into a bunch of smaller collections by using a binary tree:

{-# LANGUAGE DeriveFunctor #-}

import Data.Functor.Foldable (hylo)
import Data.List.Ordered (merge)

data Tree a r =
    Empty
  | Leaf a
  | Node r r
  deriving Functor

The idea is that each node in the tree holds two subtrees, each of which contains half of the remaining elements. We can build a tree like this from a collection - say, a basic Haskell list. The following unfolder function defines what part of a tree to build for any corresponding part of a list:

unfolder []  = Empty
unfolder [x] = Leaf x
unfolder xs  = Node l r where
  (l, r) = splitAt (length xs `div` 2) xs

On the other hand, we can also collapse an existing tree back into a list. The following folder function defines how to collapse any given part of a tree into the corresponding part of a list; again we just pattern match on whatever part of the tree we’re looking at, and construct the complementary list:

folder Empty      = []
folder (Leaf x)   = [x]
folder (Node l r) = merge l r

Now to sort a list we can just glue these instructions together using a hylomorphism:

mergesort :: Ord a => [a] -> [a]
mergesort = hylo folder unfolder

And it works just like you’d expect:

> mergesort [1,10,3,4,5]
[1,3,4,5,10]
> mergesort "aloha"
"aahlo"
> mergesort [True, False, False, True, False]
[False, False, False, True, True]

Pretty concise!

The code is eminently clean and faithful to the high-level algorithm description: first recursively divide a collection into smaller subcollections

via a binary tree and unfolder - and then recursively sort and merge the subcollections to reconstruct the (now sorted) original one - via folder.

A version of this post originally appeared on the Fugue blog.

Markov Chains à la Carte

2015-10-14T00:00:00+04:00

I’ve released a number of libraries for doing Markov Chain Monte Carlo (MCMC) in Haskell.

You can get at them via a ‘frontend’ library, declarative, but each can also be used fruitfully on its own. À la carte, if you will.

Some background: MCMC is a family of stateful algorithms for sampling from a large class of probability distributions. Typically one is interested in doing this to approximate difficult integrals; instead of choosing some suitable grid of points in parameter space over which to approximate an integral, just offload the problem to probability theory and use a Markov chain to find them for you.

For an excellent introduction to MCMC you won’t find better than Iain Murray’s lectures from MLSS ’09 in Cambridge, so check those out if you’re interested in more details.

I’ve put together a handful of popular MCMC algorithms as well as an easy way to glue them together in a couple of useful ways. At present these implementations are useful in cases where you can write your target function in closed form, and that’s pretty much all that’s required (aside from the standard algorithm-specific tuning parameters).

The API should be pretty easy to work with — write your target as a function of its parameters, specify a start location, and away you go. It’s also cool if your target accepts its parameters via most common traversable functors — lists, vectors, sequences, maps, etc.

That’s sort of the goal of this first release: if you can give me a target function, I’ll do my best to give you samples from it. Less is more and all that.

What‘s In The Box

There are a number of libraries involved. I have a few more in the queue and there are a number of additional features I plan to support for these ones in particular, but without further ado:

mwc-probability, a sampling-function based probability monad implemented as a thin wrapper over the excellent mwc-random library.
mcmc-types, housing a number of types used by the the whole family.
mighty-metropolis, an implementation of the famous Metropolis algorithm.
speedy-slice, a slice sampling implementation suitable for both continuous & discrete parameter spaces.
hasty-hamiltonian, an implementation of the gradient-based Hamiltonian Monte Carlo algorithm.
declarative, the one ring to rule them all.

Pull down declarative if you just want to have access to all of them. If you’re a Haskell neophyte you can find installation instructions at the Github repo.

Motivation

MCMC is fundamentally about observing Markov chains over probability spaces. In this context a chain is a stochastic process that wanders around a state space, eventually visiting regions of the space in proportion to their probability.

Markov chains are constructed by transition operators that obey the Markov property: that the probability of transitioning to the next location — conditional on the history of the chain — depends only on the current location. For MCMC we’re also interested in operators that satisfy the reversibility property — that the probability a transition from state A to state B occurs is the same as that a transition from state B to state A occurs. A chain is characterized by a transition operator T that drives it from state to state, and for MCMC we want the stationary or limiting distribution of the chain to be the distribution we’re sampling from.

One of the major cottage industries in Bayesian research is inventing new transition operators to drive the Markov chains used in MCMC. This has been fruitful, but it could likely be aided by a practical way to make existing transition operators work together.

This is easy to do in theory: there are a couple of ways to combine transition operators such that the resulting composite operator preserves all the properties we’re interested in for MCMC — the stationary distribution, reversibility, and Markov property. See Geyer, 2005 for details here, but the crux is that we can establish the following simple grammar for transition operators:

transition ::= primitive <transition>
             | concat transition transition
             | sample transition transition

A transition is either some primitive operator, a deterministic concatenation of operators (via ‘concat’), or a probabilistic concatenation of operators (via ‘sample’). A deterministic concatenation works by just transitioning through two operators one after the other; a probabilistic concatenation works by randomly choosing one transition operator or the other to use on any given transition. These kinds of concatenation preserve all the properties we’re interested in.

We can trivially generalize this further by adding a term that concatenates n transition operators together deterministically, or another for probabilistically concatenating a bunch of operators according to some desired probability distribution.

The idea here is that there are tradeoffs involved in different transition operators. Some may be more computationally expensive than others (perhaps requiring a gradient evaluation, or evaluation of some inner loop) but have better ability to make ‘good’ transitions in certain situations. Other operators are cheap, but can be inefficient (taking a long time to visit certain regions of the space).

By employing deterministic or probabilistic concatenation, one can concoct a Markov chain that uses a varied range of tuning parameters, for example. Or only occasionally employs a computationally expensive transition, otherwise preferring some cheaper, reliable operator.

Usage

The declarative library implements this simple language for transition operators, and the mighty-metropolis, speedy-slice, and hasty-hamiltonian libraries provide some primitive transitions that you can combine as needed.

The Metropolis and slice sampling transitions are cheap and require little information, whereas Hamiltonian Monte Carlo exploits information about the target’s gradient and also involves evaluation of an inner loop (the length of which is determined by a tuning parameter). Feel free to use one that suits your problem, or combine them together using the combinators supplied in declarative to build a custom solution.

As an example, the Rosenbrock density is a great test dummy as it’s simple, low-dimensional, and can be easily visualized, but it still exhibits a pathological anisotropic structure that makes it somewhat tricky to sample from.

Getting started via declarative is pretty simple:

import Numeric.MCMC

You’ll want to supply a target to sample over, and if you want to use an algorithm like Hamiltonian Monte Carlo you’ll also need to provide a gradient. If you can’t be bothered to calculate gradients by hand, you can always turn to your friend automatic differentiation:

import Numeric.AD

The Rosenbrock log-density and its gradient can then be written as follows:

target :: Num a => [a] -> a
target [x0, x1] = negate (100 * (x1 — x0 ^ 2) ^ 2 + (1 — x0) ^ 2)

gTarget :: Num a => [a] -> [a]
gTarget = grad target

All you need to do here is provide a function proportional to a log-probability density. The logarithmic scale is important; various internals expect to be passed (something proportional to) a log-probability density.

To package these guys up together we can wrap them in a Target. Note that we don’t always care about including a gradient, so that part is optional:

rosenbrock :: Target [Double]
rosenbrock = Target target (Just gTarget)

The Target type is parameterized over the shape of the parameter space. You could similarly have a Target (Seq Double), Target (Map String Double), and so on. Your target may be implemented using a boxed vector for efficiency, for example. Or using a Map or HashMap with string/text keys such that parameter names are preserved. They should work just fine.

Given a target, we can sample from it a bunch of times using a simple Metropolis transition via the mcmc function. Aside from the target and a PRNG, provide it with the desired number of transitions, a starting point, and the transition operator to use:

> -- haskell
> prng <- create
> mcmc 10000 [0, 0] (metropolis 1) rosenbrock prng

In return you’ll get the desired trace of the chain dumped to stdout:

8.136972300105949e-2,0.273896953404261
0.4657348148676972,0.17462596647788464
-0.48609414127836326,9.465052854751566e-2
-0.49781488399832785,0.42092910345708523
-0.3019713424699155,0.39135350029173566
0.12058426470979189,0.12485407390388925
..

The intent is for the chain to be processed elsewhere — if you’re me, that will usually be in R. Libraries like coda have a ton of functionality useful for working with Markov chain traces, and ggplot2 as a library for static statistical graphics can’t really be beat:

> # r
> d = read.csv(‘rosenbrock-trace.dat’, header = F)
> names(d) = c(‘x’, ‘y’)
> require(ggplot2)
> ggplot(d, aes(x, y)) + geom_point(colour = ‘darkblue’, alpha = 0.2)

You get the following trace over the Rosenbrock density, taken for 10k iterations. This is using a Metropolis transition with variance 1:

If you do want to work with chains in memory in Haskell you can do that by writing your own handling code around the supplied transition operators. I’ll probably make this a little easier in later versions.

The implementations are reasonably quick and don’t leak memory — the traces are streamed to stdout as the chains are traversed. Compiling the above with ‘-O2’ and running it for 100k iterations yields the following performance characteristics on my mid-2011 model MacBook Air:

$ ./test/Rosenbrock +RTS -s > /dev/null

3,837,201,632 bytes allocated in the heap
    8,453,696 bytes copied during GC
       89,600 bytes maximum residency (2 sample(s))
       23,288 bytes maximum slop
         1 MB total memory in use (0 MB lost due to fragmentation)

 INIT time 0.000s ( 0.000s elapsed)
  MUT time 3.539s ( 3.598s elapsed)
   GC time 0.049s ( 0.058s elapsed)
 EXIT time 0.000s ( 0.000s elapsed)
Total time 3.591s ( 3.656s elapsed)

%GC time 1.4% (1.6% elapsed)

Alloc rate 1,084,200,280 bytes per MUT second

Productivity 98.6% of total user, 96.8% of total elapsed

The beauty is that rather than running a chain solely on something like the simple Metropolis operator used above, you can sort of ‘hedge your sampling risk’ and use a composite operator that proposes transitions using a multitude of ways. Consider this guy, for example:

transition =
  concatT
    (sampleT (metropolis 0.5) (metropolis 1.0))
    (sampleT (slice 2.0) (slice 3.0))

Here concatT and sampleT correspond to the concat and sample terms in the BNF description in the previous section. This operator performs two transitions back-to-back; the first is randomly a Metropolis transition with standard deviation 0.5 or 1 respectively, and the second is a slice sampling transition using a step size of 2 or 3, randomly.

Running it for 5000 iterations (to keep the total computation approximately constant), we see a chain that has traversed the space a little better:

> mcmc 5000 [0, 0] transition rosenbrock prng

It’s worth noting that I didn’t put any work into coming up with this composite transition: this was just the first example I thought up, and a lot of the benefits here probably come primarily from including the eminently-reliable slice sampling transition. But from informal experimentation, it does seem that chains driven by composite transitions involving numerous operators and tuning parameter settings often seem to perform better on average than a given chain driven by a single (poorly-selected) transition.

I know exactly how meticulous proofs and benchmarks must be so I haven’t rigorously established any properties around this, but hey: it ‘seems to be the case’, and intuitively, including varied transition operators surely hedges your bets when compared to using a single one.

Try it out and see how your mileage varies, and be sure to let me know if you find some killer apps where composite transitions really seem to win.

Implementation Notes

If you’re just interested in using the libraries you can skip the following section, but I just want to point out how easy this is to implement.

The implementations are defined using a small set of types living in mcmc-types:

type Transition m a = StateT a (Prob m) ()

data Chain a b = Chain {
    chainTarget   :: Target a
  , chainScore    :: Double
  , chainPosition :: a
  , chainTunables :: Maybe b
  }

data Target a = Target {
    lTarget  :: a -> Double
  , glTarget :: Maybe (a -> a)
  }

Most important here is the Transition type, which is just a state transformer over a probability monad (itself defined in mwc-probability). The probability monad is the source of randomness used to define transition operators useful for MCMC, and values with type Transition are the transition operators in question.

The Chain type is the state of the Markov chain at any given iteration. All that’s really required here is the chainPosition field, which represents the location of the chain in parameter space. But adding some additional information here is convenient; chainScore caches the most recent score of the chain (which is typically used in internal calculations, and caching avoids recomputing things needlessly) and chainTunables is an optional record intended to be used for stateful tuning parameters (used by adaptive algorithms or in burn-in phases and the like). Additionally the target being sampled from itself — chainTarget — is included in the state.

Undisciplined use of chainTarget and chainTunables can have all sorts of nasty consequences — you can use them to change the stationary distribution you’re sampling from or invalidate the Markov property — but keeping them around is useful for implementing some desirable features. Tweaking chainTarget, for example, allows one to easily implement annealing, which can be very useful for sampling from annoying multi-modal densities.

Setting everything up like this makes it trivial to mix-and-match transition operators as required — the state and probability monad stack provides everything we need. Deterministic concatenation is implemented as follows, for example:

concatT = (>>)

and a generalized version of probabilistic concatenation just requires a coin flip:

bernoulliT p t0 t1 = do
  heads <- lift (MWC.bernoulli p)
  if heads then t0 else t1

A uniform probabilistic concatenation over two operators, implemented in sampleT, is then just bernoulliT 0.5.

The difficulty of implementing primitive operators just depends on the operator itself; the surrounding framework is extremely lightweight. Here’s the Metropolis transition, for example (with type signatures omitted to keep the noise down):

metropolis radial = do
  Chain {..} <- get
  proposal   <- lift (propose radial chainPosition)
  let proposalScore = lTarget chainTarget proposal
      acceptProb    = whenNaN 0
        (exp (min 0 (proposalScore - chainScore)))

  accept <- lift (MWC.bernoulli acceptProb)
  when accept
    (put (Chain chainTarget proposalScore proposal chainTunables))

propose radial = traverse perturb where
  perturb m = MWC.normal m radial

And the excellent pipes library is used to generate a Markov chain:

chain radial = loop where
  loop state prng = do
    next <- lift
      (MWC.sample (execStateT (metropolis radial) state) prng)
    yield next
    loop next prng

The mcmc functions are also implemented using pipes. Take the first n iterations of a chain and print them to stdout. That simple.

Future Work

In the near term I plan on updating some old MCMC implementations I have kicking around on Github (flat-mcmc, lazy-langevin, hnuts) and releasing them within this framework. Additionally I’ve got some code for building annealed operators that I want to release — it has been useful in some situations when sampling from things like the Himmelblau density, which has a few disparate clumps of probability that make it tricky to sample from with conventional algorithms.

This framework is also useful as an inference backend to languages for working with directed graphical models (think BUGS/Stan). The idea here is that you don’t need to specify your target function (typically a posterior density) explicitly: just describe your model and I’ll give you samples from the posterior distribution. A similar version has been put to use around the BayesHive project.

Longer term — I’ll have to see what’s up in terms of demand. There are performance improvements and straightforward extensions to things like parallel tempering, but I’m growing more interested in ‘online’ methods like particle MCMC and friends that are proving useful for inference in more general probabilistic programs (think those expressible by Church and its ilk).

Let me know if you get any use out of these things, or please file an issue if there’s some particular feature you’d like to see supported.

Thanks to Niffe Hermansson for review and helpful comments.

Practical Recursion Schemes

2015-09-06T00:00:00+04:00

Recursion schemes are elegant and useful patterns for expressing general computation. In particular, they allow you to ‘factor recursion out’ of whatever semantics you may be trying to express when interpreting programs, keeping your interpreters concise, your concerns separated, and your code more maintainable.

What’s more, formulating programs in terms of recursion schemes seems to help suss out particular similarities in structure between what might be seen as disparate problems in other domains. So aside from being a practical computational tool, they seem to be of some use when it comes to ‘hacking understanding’ in varied areas.

Unfortunately, they come with a pretty forbidding barrier to entry. While there are a few nice resources out there for learning about recursion schemes and how they work, most literature around them is quite academic and awash in some astoundingly technical jargon (more on this later). Fortunately, the accessible resources out there do a great job of explaining what recursion schemes are and how you might use them, so they go through some effort to build up the required machinery from scratch.

In this article I want to avoid building up the machinery meticulously and instead concentrate mostly on understanding and using Edward Kmett’s recursion-schemes library, which, while lacking in documentation, is very well put together and implements all the background plumbing one needs to get started.

In particular, to feel comfortable using recursion-schemes I found that there were a few key patterns worth understanding:

Factoring recursion out of your data types using pattern functors and a fixed-point wrapper.
Using the ‘Foldable’ & ‘Unfoldable’ classes, plus navigating the ‘Base’ type family.
How to use some of the more common recursion schemes out there for everyday tasks.

The Basics

If you’re following along in GHCi, I’m going to first bring in some imports and add a useful pragma. I’ll dump a gist at the bottom; note that this article targets GHC 7.10.2 and recursion-schemes-4.1.2, plus I’ll also require data-ordlist-0.4.7.0 for an example later. Here’s the requisite boilerplate:

{-# LANGUAGE DeriveFunctor #-}

import Data.Functor.Foldable
import Data.List.Ordered (merge)
import Prelude hiding (Foldable, succ)

So, let’s get started.

Recursion schemes are applicable to data types that have a suitable recursive structure. Lists, trees, and natural numbers are illustrative candidates.

Being so dead-simple, let’s take the natural numbers as an illustrative/toy example. We can define them recursively as follows:

data Natural =
    Zero
  | Succ Natural

This is a fine definition, but many such recursive structures can also be defined in a different way: we can first ‘factor out’ the recursion by defining some base structure, and then ‘add it back in’ by using a recursive wrapper type.

The price of this abstraction is a slightly more involved type definition, but it unlocks some nice benefits — namely, the ability to reason about recursion and base structures separate from each other. This turns out to be a very useful pattern for getting up and running with recursion schemes.

The trick is to create a different, parameterized type, in which the new parameter takes the place of all recursive points in the original type. We can create this kind of base structure for the natural numbers example as follows:

data NatF r =
    ZeroF
  | SuccF r
  deriving (Show, Functor)

This type must be a functor in this new parameter, so the type is often called a ‘pattern functor’ for some other type. I like to use the notation ‘F’ when defining constructors for pattern functors.

We can define pattern functors for lists and trees in the same way:

data ListF a r =
    NilF
  | ConsF a r
  deriving (Show, Functor)

data TreeF a r =
    EmptyF
  | LeafF a
  | NodeF r r
  deriving (Show, Functor)

Now, to add recursion to these pattern functors we’re going to use the famous fixed-point type, ‘Fix’, to wrap them in:

type Nat    = Fix NatF
type List a = Fix (ListF a)
type Tree a = Fix (TreeF a)

‘Fix’ is a standard fixed-point type imported from the recursion-schemes library. You can get a ton of mileage from it. It introduces the ‘Fix’ constructor everywhere, but that’s actually not much of an issue in practice. One thing I typically like to do is add some smart constructors to get around it:

zero :: Nat
zero = Fix ZeroF

succ :: Nat -> Nat
succ = Fix . SuccF

nil :: List a
nil = Fix NilF

cons :: a -> List a -> List a
cons x xs = Fix (ConsF x xs)

Then you can write expressions like ‘succ (succ (succ zero))’ without having to deal with the ‘Fix’ constructor explicitly. Note also that these expressions are Showable à la carte, for example in GHCi:

> succ (succ (succ zero))
Fix (SuccF (Fix (SuccF (Fix (SuccF (Fix ZeroF))))))

A Short Digression on ‘Fix’

The ‘Fix’ type is brought into scope from ‘Data.Functor.Foldable’, but it’s worth looking at it in some detail. It can be defined as follows, along with two helpful functions for working with it:

newtype Fix f = Fix (f (Fix f))

fix :: f (Fix f) -> Fix f
fix = Fix

unfix :: Fix f -> f (Fix f)
unfix (Fix f) = f

‘Fix’ has a simple recursive structure. For a given value, you can think of ‘fix’ as adding one level of recursion to it. ‘unfix’ in turn removes one level of recursion.

This generic recursive structure is what makes ‘Fix’ so useful: we can write some nominally recursive type we’re interested in without actually using recursion, but then package it up in ‘Fix’ to hijack the recursion it provides automatically.

Understanding Some Internal Plumbing

If we wrap a pattern functor in ‘Fix’ then the underlying machinery of recursion-schemes should ‘just work’. Here it’s worth explaining a little as to why that’s the case.

There are two fundamental type classes involved in recursion-schemes: ‘Foldable’ and ‘Unfoldable’. These serve to tease apart the recursive structure of something like ‘Fix’ even more: loosely, ‘Foldable’ corresponds to types that can be ‘unfixed’, and ‘Unfoldable’ corresponds to types that can be ‘fixed’. That is, we can add more layers of recursion to instances of ‘Unfoldable’, and we can peel off layers of recursion from instances of ‘Foldable’.

In particular ‘Foldable’ and ‘Unfoldable’ contain functions called ‘project’ and ‘embed’ respectively, corresponding to more general forms of ‘unfix’ and ‘fix’. Their types are as follows:

project :: Foldable t   => t -> Base t t
embed   :: Unfoldable t => Base t t -> t

I’ve found it useful while using recursion-schemes to have a decent understanding of how to interpret the type family ‘Base’. It appears frequently in type signatures of various recursion schemes and being able to reason about it can help a lot.

‘Base’ and Basic Type Families

Type families are type-level functions; they take types as input and return types as output. The ‘Base’ definition in recursion-schemes looks like this:

type family Base t :: * -> *

You can interpret this as a function that takes one type ‘t’ as input and returns some other type. An implementation of this function is called an instance of the family. The instance for ‘Fix’, for example, looks like:

type instance Base (Fix f) = f

In particular, a type family like ‘Base’ is a synonym for instances of the family. So using the above example: anywhere you see something like ‘Base (Fix f)’ you can mentally replace it with ‘f’.

Instances of the ‘Base’ type family have a structure like ‘Fix’, but using ‘Base’ enables all the internal machinery of recursion-schemes to work out-of-the-box for types other than ‘Fix’ alone. This has a typical Kmettian flavour: first solve the most general problem, and then recover useful, specific solutions to it automatically.

Predictably, ‘Fix f’ is an instance of ‘Base’, ‘Foldable’, and ‘Unfoldable’ for some functor ‘f’, so if you use it, you can freely use all of recursion-schemes’s innards without needing to manually write any instances for your own data types. But as mentioned above, it’s worth noting that you can exploit the various typeclass & type family machinery to get by without using ‘Fix’ at all: see e.g. Danny Gratzer’s recursion-schemes post for an example of this.

Some Useful Schemes

So, with some discussion of the internals out of the way, we can look at some of the more common and useful recursion schemes. I’ll concentrate on the following four, as they’re the ones I’ve found the most use for:

catamorphisms, implemented via ‘cata’, are generalized folds.
anamorphisms, implemented via ‘ana’, are generalized unfolds.
hylomorphisms, implemented via ‘hylo’, are anamorphisms followed by catamorphisms (corecursive production followed by recursive consumption).
paramorphisms, implemented via ‘para’, are generalized folds with access to the input argument corresponding to the most recent state of the computation.

Let me digress slightly on nomenclature.

Yes, the names of these things are celebrations of the ridiculous. There’s no getting around it; they look like self-parody to almost anyone not pre-acquainted with categorical concepts. They have been accused — probably correctly — of being off-putting.

That said, they communicate important technical details and are actually not so bad when you get used to them. It’s perfectly fine and even encouraged to arm-wave about folds or unfolds when speaking informally, but the moment someone distinguishes one particular style of fold from another via a prefix like e.g. para, I know exactly the relevant technical distinctions required to understand the discussion. The names might be silly, but they have their place.

Anyway.

There are many other more exotic schemes that I’m sure are quite useful (see Tim Williams’s recursion schemes talk, for example), but I haven’t made use of any outside of these four just yet. The recursion-schemes library contains a plethora of unfamiliar schemes just waiting to be grokked, but in the interim even cata and ana alone will get you plenty far.

Now let’s use the motley crew of schemes to do some useful computation on our example data types.

Catamorphisms

Take our natural numbers type, ‘Nat’. To start, we can use a catamorphism to represent a ‘Nat’ as an ‘Int’ by summing it up.

natsum :: Nat -> Int
natsum = cata alg where
  alg ZeroF     = 0
  alg (SuccF n) = n + 1

Here ‘alg’ refers to ‘algebra’, which is the function that we use to define our reducing semantics. Notice that the semantics are not defined recursively! The recursion present in ‘Nat’ has been decoupled and is handled for us by ‘cata’. And as a plus, we still don’t have to deal with the ‘Fix’ constructor anywhere.

As a brief aside: I like to write my recursion schemes in this way, but your mileage may vary. If you’d like to enable the ‘LambdaCase’ extension, then another option is to elide mentioning the algebra altogether using a very simple case statement:

{-# LANGUAGE LambdaCase #-}

natsum :: Nat -> Int
natsum = cata $ \case ->
  ZeroF   -> 0
  SuccF n -> n + 1

Some people find this more readable.

To understand how we used ‘cata’ to build this function, take a look at its type:

cata :: Foldable t => (Base t a -> a) -> t -> a

The ‘Base t a -> a’ term is the algebra; ‘t’ is our recursive datatype (i.e. ‘Nat’), and ‘a’ is whatever type we’re reducing a value to.

Historically I’ve found ‘Base’ here to be confusing, but here’s a neat trick to help reason through it.

Remember that ‘Base’ is a type family, so for some appropriate ‘t’ and ‘a’, ‘Base t a’ is going to be a synonym for some other type. To figure out what ‘Base t a’ corresponds to for some concrete ‘t’ and ‘a’, we can ask GHCi via this lesser-known command that evaluates type families:

> :kind! Base Nat Int
Base Nat Int :: *
= NatF Int

So in the ‘natsum’ example the algebra used with ‘cata’ must have type ‘NatF Int -> Int’. This is pretty obvious for ‘cata’, but I initially found that figuring out what type should be replaced for ‘Base’ exactly could be confusing for some of the more exotic schemes.

As another example, we can use a catamorphism to implement ‘filter’ for our list type:

filterL :: (a -> Bool) -> List a -> List a
filterL p = cata alg where
  alg NilF = nil
  alg (ConsF x xs)
    | p x       = cons x xs
    | otherwise = xs

It follows the same simple pattern: we define our semantics by interpreting recursion-less constructors through an algebra, then pump it through ‘cata’.

Anamorphisms

These running examples are toys, but even here it’s really annoying to have to type ‘succ (succ (succ (succ (succ (succ zero)))))’ to get a natural number corresponding to six for debugging or what have you.

We can use an anamorphism to build a ‘Nat’ value from an ‘Int’:

nat :: Int -> Nat
nat = ana coalg where
  coalg n
    | n <= 0    = ZeroF
    | otherwise = SuccF (n - 1)

Just as a small detail: to be descriptive, here I’ve used ‘coalg’ as the argument to ‘ana’, for ‘coalgebra’.

Now the expression ‘nat 6’ will do the same for us as the more verbose example above. As always, recursion is not part of the semantics; to have the integer ‘n’ we pass in correspond to the correct natural number, we use the successor value of ‘n — 1’.

Paramorphisms

As an example, try to express a factorial on a natural number in terms of ‘cata’. It’s (apparently) doable, but an implementation is not immediately clear.

A paramorphism will operate on an algebra that provides access to the input argument corresponding to the running state of the recursion. Check out the type of ‘para’ below:

para :: Foldable t => (Base t (t, a) -> a) -> t -> a

If we’re implementing a factorial on ‘Nat’ values then ‘t’ is going to correspond to ‘Nat’ and ‘a’ is going to correspond to (say) ‘Integer’. Here it might help to use the ‘:kind!’ trick to help reason through the requirements of the algebra. We can ask GHCi to help us out:

> :kind! Base Nat (Nat, Int)
Base Nat (Nat, Int) :: *
= NatF (Nat, Int)

Side note: after doing this trick a few times you’ll probably find it much easier to reason about type families sans-GHCi. In any case, we can now implement an algebra corresponding to the required type:

natfac :: Nat -> Int
natfac = para alg where
 alg ZeroF = 1
 alg (SuccF (n, f)) = natsum (succ n) * f

Here there are some details to point out.

The type of our algebra is ‘NatF (Nat, Int) -> Int’; the value with the ‘Nat’ type, ‘n’, holds the most recent input argument used to compute the state of the computation, ‘f’.

If you picture a factorial defined as

0!       = 1
(k + 1)! = (k + 1) * k!

Then ‘n’ corresponds to ‘k’ and ‘f’ corresponds to ‘k!’. To compute the factorial of the successor to ‘n’, we just convert ‘succ n’ to an integer (via ‘natsum’) and multiply it by ‘f’.

Paramorphisms tend to be pretty useful for a lot of mundane tasks. We can easily implement ‘pred’ on natural numbers via ‘para’:

natpred :: Nat -> Nat
natpred = para alg where
  alg ZeroF          = zero
  alg (SuccF (n, _)) = n

We can also implement ‘tail’ on lists. To check the type of the required algebra we can again get some help from GHCi; here I’ll evaluate a general type family, for illustration:

> :set -XRankNTypes
> :kind! forall a b. Base (List a) (List a, b)
forall a b. Base (List a) (List a, b) :: *
= forall a b. ListF a (Fix (ListF a), b)

Providing an algebra of the correct structure lets ‘tailL’ fall out as follows:

tailL :: List a -> List a
tailL = para alg where
  alg NilF             = nil
  alg (ConsF _ (l, _)) = l

You can check that ‘tailL’ indeed returns the tail of its argument.

Hylomorphisms

Hylomorphisms can express general computation — corecursive production followed by recursive consumption. Compared to the other type signatures in recursion-schemes, ‘hylo’ is quite simple:

hylo :: Functor f => (f b -> b) -> (a -> f a) -> a -> b

It doesn’t even require the full structure built up for e.g. ‘cata’ and ‘ana’; just very simple F-{co}algebras.

My favourite example hylomorphism is an absolutely beautiful implementation of mergesort. I think it helps illustrate how recursion schemes can help tease out incredibly simple structure in what could otherwise be a more involved problem.

Our input will be a Haskell list containing some orderable type. We’ll use it to build a balanced binary tree via an anamorphism and then tear it down with a catamorphism, merging lists together and sorting them as we go.

The resulting code looks like this:

mergeSort :: Ord a => [a] -> [a]
mergeSort = hylo alg coalg where
  alg EmptyF      = []
  alg (LeafF c)   = [c]
  alg (NodeF l r) = merge l r

  coalg []  = EmptyF
  coalg [x] = LeafF x
  coalg xs  = NodeF l r where
    (l, r) = splitAt (length xs `div` 2) xs

What’s more, the fusion achieved via this technique is really quite lovely.

Wrapping Up

Hopefully this article helps fuel any ‘programming via structured recursion’ trend that might be ever-so-slowly growing.

When programming in a language like Haskell, a very natural pattern is to write little embedded languages and mini-interpreters or compilers to accomplish tasks. Typically these tiny embedded languages have a recursive structure, and when you’re interpreting a recursive structure, you have use all these lovely off-the-shelf strategies for recursion available to you to keep your programs concise, modular, and efficient. The recursion-schemes library really has all the built-in machinery you need to start using these things for real.

If you’re interested about hearing about using recursion schemes ‘for real’ I recommend Tim Williams’s Exotic Tools For Exotic Trades talk (for a motivating example for the use of recursion schemes in production) or his talk on recursion schemes from the London Haskell User’s Group a few years ago.

So happy recursing! I’ve dropped the code from this article into a gist.

Thanks to Maciej Woś for review and helpful comments.

Automasymbolic Differentiation

2014-07-06T00:00:00+04:00

Automatic differentiation is one of those things that’s famous for not being as famous as it should be (uh..). It’s useful, it’s convenient, and yet fewer know about it than one would think.

This article (by one of the guys working on Venture) is the single best introduction you’ll probably find to AD, anywhere. It gives a wonderful introduction to the basics, the subtleties, and the gotchas of the subject. You should read it. I’m going to assume you have.

In particular, you should note this part:

[computing gradients] can be done — the method is called reverse mode — but it introduces both code complexity and runtime cost in the form of managing this storage (traditionally called the “tape”). In particular, the space requirements of raw reverse mode are proportional to the runtime of f.

In some applications this can be somewhat inconvenient - picture iterative gradient-based sampling algorithms like Hamiltonian Monte Carlo, its famous auto-tuning version NUTS, and Riemannian manifold variants. Tom noted in this reddit thread that symbolic differentiation - which doesn’t need to deal with tapes and infinitesimal-tracking - can often be orders of magnitude faster than AD for this kind of problem. When running these algos we calculate gradients many, many times, and the additional runtime cost of the reverse-mode AD dance can add up.

An interesting question is whether or not this can this be mitigated at all, and to what degree. In particular: can we use automatic differentiation to implement efficient symbolic differentiation? Here’s a possible attack plan:

use an existing automatic differentiation implementation to calculate the gradient for some target function at a point
capture the symbolic expression for the gradient and optimize it by eliminating common subexpressions or whatnot
reify that expression as a function that we can evaluate for any input
voila, (more) efficient gradient

Ed Kmett’s ‘ad’ library is the best automatic differentiation library I know of, in terms of its balance between power and accessibility. It can be used on arbitrary Haskell types that are instances of the Num typeclass and can carry out automatic differentiation via a number of modes, so it’s very easy to get started with. Rather than rolling my own AD implementation to try to do this sort of thing, I’d much rather use his.

Start with a really basic language. In practice we’d be interested in working with more expressive things, but this is sort of the minimal interesting language capable of illustrating the problem:

import Numeric.AD

data Expr a =
    Lit a
  | Var String
  | Add (Expr a) (Expr a)
  | Sub (Expr a) (Expr a)
  | Mul (Expr a) (Expr a)
  deriving (Eq, Show)

instance Num a => Num (Expr a) where
  fromInteger = Lit . fromInteger
  e0 + e1 = Add e0 e1
  e0 - e1 = Sub e0 e1
  e0 * e1 = Mul e0 e1

We can already use Kmett’s ad library on these expressions to generate symbolic expressions. We just have to write functions we’re interested in generically (using +, -, and *) and then call diff or grad or whatnot on them with a concretely-typed argument. Some examples:

> :t diff (\x -> 2 * x ^ 2)
diff (\x -> 2 * x ^ 2) :: Num a => a -> a

> diff (\x -> 2 * x ^ 2) (Lit 1)
Mul (Add (Mul (Lit 1) (Lit 1)) (Mul (Lit 1) (Lit 1))) (Lit 2)

> grad (\[x, y] -> 2 * x ^ 2 + 3 * y) [Lit 1, Lit 2]
[ Add (Lit 0) (Add (Add (Lit 0) (Mul (Lit 1) (Add (Lit 0) (Mul (Lit 2) (Add
  (Lit 0) (Mul (Lit 1) (Lit 1))))))) (Mul (Lit 1) (Add (Lit 0) (Mul (Lit 2) (Add
  (Lit 0) (Mul (Lit 1) (Lit 1)))))))
, Add (Lit 0) (Add (Lit 0) (Mul (Lit 3) (Add (Lit 0) (Mul (Lit 1) (Lit 1)))))
]

It’s really easy to extract a proper derivative/gradient ‘function’ by doing something like this:

> diff (\x -> 2 * x ^ 2) (Var "x")
Mul (Add (Mul (Var "x") (Lit 1)) (Mul (Lit 1) (Var "x"))) (Lit 2)

and then, given that expression, reifying a direct function for the derivative by substituting over the captured variable and evaluating everything. Here’s some initial machinery to handle that:

-- | Close an expression over some variable.
close :: Expr a -> String -> a -> Expr a
close (Add e0 e1) s x = Add (close e0 s x) (close e1 s x)
close (Sub e0 e1) s x = Sub (close e0 s x) (close e1 s x)
close (Mul e0 e1) s x = Mul (close e0 s x) (close e1 s x)

close (Var v) s x
  | v == s    = Lit x
  | otherwise = Var v

close e _ _ = e

-- | Evaluate a closed expression.
eval :: Num a => Expr a -> a
eval (Lit d) = d
eval (Var _) = error "expression not closed"
eval (Add e0 e1) = eval e0 + eval e1
eval (Sub e0 e1) = eval e0 - eval e1
eval (Mul e0 e1) = eval e0 * eval e1

So, using this on the example above yields

> let testExpr = diff (\x -> 2 * x ^ 2) (Var "x")
> eval . close testExpr "x" $ 1
2

and it looks like a basic toDerivative function for expressions could be implemented as follows:

-- | Do some roundabout AD.
toDerivative expr = eval . close diffExpr "x" where
  diffExpr = diff expr (Var "x")

But that’s a no go. ‘ad’ throws a type error, as presumably we’d be at risk of perturbation confusion:

Couldn't match expected type ‘AD
       s (Numeric.AD.Internal.Forward.Forward (Expr c))
    -> AD s (Numeric.AD.Internal.Forward.Forward (Expr c))’
   with actual type ‘t’
  because type variable ‘s’ would escape its scope
This (rigid, skolem) type variable is bound by
  a type expected by the context:
    AD s (Numeric.AD.Internal.Forward.Forward (Expr c))
    -> AD s (Numeric.AD.Internal.Forward.Forward (Expr c))
  at ParamBasic.hs:174:14-32
Relevant bindings include
  diffExpr :: Expr c (bound at ParamBasic.hs:174:3)
  expr :: t (bound at ParamBasic.hs:173:14)
  toDerivative :: t -> c -> c (bound at ParamBasic.hs:173:1)
In the first argument of ‘diff’, namely ‘expr’
In the expression: diff expr (Var "x")

Instead, we can use ad’s auto combinator to write an alternate eval function:

autoEval :: Mode a => String -> Expr (Scalar a) -> a -> a
autoEval x expr = (`go` expr) where
  go _ (Lit d) = auto d
  go v (Var s)
    | s == x    = v
    | otherwise = error "expression not closed"

  go v (Add e0 e1) = go v e0 + go v e1
  go v (Sub e0 e1) = go v e0 - go v e1
  go v (Mul e0 e1) = go v e0 * go v e1

and using that, implement a working toDerivative:

toDerivative :: Num a => String -> Expr (Expr a) -> Expr a
toDerivative v expr = diff (autoEval v expr)

which, though it has a weird-looking type, typechecks and does the trick:

> let d = toDerivative "x" (Mul (Lit 2) (Mul (Var "x") (Var "x"))) (Var "x")
Mul (Add (Mul (Var "x") (Lit 1)) (Mul (Lit 1) (Var "x"))) (Lit 2)

> eval . close "x" 1 $ d
4

So now we have access to a reified AST for a derivative (or gradient), which can be tweaked and optimized as needed. Cool.

The available optimizations depend heavily on the underlying language. For starters, there’s easy and universal stuff like this:

-- | Reduce superfluous expressions.
elimIdent :: (Num a, Eq a) => Expr a -> Expr a
elimIdent (Add (Lit 0) e) = elimIdent e
elimIdent (Add e (Lit 0)) = elimIdent e
elimIdent (Add e0 e1)     = Add (elimIdent e0) (elimIdent e1)

elimIdent (Sub (Lit 0) e) = elimIdent e
elimIdent (Sub e (Lit 0)) = elimIdent e
elimIdent (Sub e0 e1)     = Sub (elimIdent e0) (elimIdent e1)

elimIdent (Mul (Lit 1) e) = elimIdent e
elimIdent (Mul e (Lit 1)) = elimIdent e
elimIdent (Mul e0 e1)     = Mul (elimIdent e0) (elimIdent e1)

elimIdent e = e

Which lets us do some work up-front:

> let e = 2 * Var "x" ^ 2 + Var "x" ^ 4
Add (Mul (Lit 2) (Mul (Var "x") (Var "x"))) (Mul (Mul (Var "x") (Var "x"))
(Mul (Var "x") (Var "x")))

> let ge = toDerivative "x" e
Add (Mul (Add (Mul (Var "x") (Lit 1)) (Mul (Lit 1) (Var "x"))) (Lit 2))
(Add (Mul (Mul (Var "x") (Var "x")) (Add (Mul (Var "x") (Lit 1)) (Mul
(Lit 1) (Var "x")))) (Mul (Add (Mul (Var "x") (Lit 1)) (Mul (Lit 1)
(Var "x"))) (Mul (Var "x") (Var "x"))))

> let geOptim = elimIdent ge
Add (Mul (Add (Var "x") (Var "x")) (Lit 2)) (Add (Mul (Mul (Var "x")
(Var "x")) (Add (Var "x") (Var "x"))) (Mul (Add (Var "x") (Var "x")) (Mul
(Var "x") (Var "x"))))

> eval . close "x" 1 $ geOptim
8

But there are also some more involved optimizations that can be useful for some languages. The basic language I’ve been using above, for example, has no explicit support for sharing common subexpressions. You’ll recall from one of my previous posts that we have a variety of methods to do that in Haskell EDSLs, including some that allow sharing to be observed without modifying the underlying language. We can use data-reify, for example, to observe any implicit sharing in expressions:

> reifyGraph $ geOptim
let [(1,AddF 2 6),(6,AddF 7 10),(10,MulF 11 12),(12,MulF 4 4),(11,AddF 4 4),
(7,MulF 8 9),(9,AddF 4 4),(8,MulF 4 4),(2,MulF 3 5),(5,LitF 2),(3,AddF 4 4),
(4,VarF "x")] in 1

And even make use of a handy library found on Hackage for performing common subexpression elimination on graphs returned by reifyGraph:

> cse <$> reifyGraph geOptim
let [(5,LitF 2),(1,AddF 2 6),(3,AddF 4 4),(6,AddF 7 10),(2,MulF 3 5),
(10,MulF 3 8),(8,MulF 4 4),(7,MulF 8 3),(4,VarF "x")] in 1

With an appropriate graph evaluator we can cut down the size of the syntax we have to traverse substantially.

Happy automasymbolic differentiating!

Sharing in Haskell EDSLs

2014-05-30T00:00:00+04:00

Lately I’ve been trying to do some magic by way of nonstandard interpretations of abstract syntax. One of the things that I’ve managed to grok along the way has been the problem of sharing in deeply-embedded languages.

Here’s a simple illustration of the ‘vanilla’ sharing problem by way of plain Haskell; a function that computes 2^n:

naiveTree :: (Eq a, Num a, Num b) => a -> a
naiveTree 0 = 1
naiveTree n = naiveTree (n - 1) + naiveTree (n - 1)

This naive implementation is a poor way to roll as it is exponentially complex in n. Look at how evaluation proceeds for something like naiveTree 4:

> naiveTree 4
> naiveTree 3 + naiveTree 3
> naiveTree 2 + naiveTree 2 + naiveTree 2 + naiveTree 2
> naiveTree 1 + naiveTree 1 + naiveTree 1 + naiveTree 1
  + naiveTree 1 + naiveTree 1 + naiveTree 1 + naiveTree 1
> naiveTree 0 + naiveTree 0 + naiveTree 0 + naiveTree 0
  + naiveTree 0 + naiveTree 0 + naiveTree 0 + naiveTree 0
  + naiveTree 0 + naiveTree 0 + naiveTree 0 + naiveTree 0
  + naiveTree 0 + naiveTree 0 + naiveTree 0 + naiveTree 0
> 1 + 1 + 1 + 1 + 1 + 1 + 1 + 1 + 1 + 1 + 1 + 1 + 1 + 1 + 1 + 1
> 16

Each recursive call doubles the number of function evaluations we need to make. Don’t wait up for naiveTree 50 to return a value.

A better way to write this function would be:

tree :: (Eq a, Num a, Num b) => a -> a
tree 0 = 1
tree n =
  let shared = tree (n - 1)
  in  shared + shared

Here we store solutions to subproblems, and thus avoid having to recompute things over and over. Look at how tree 4 proceeds now:

> tree 4
> let shared0 =
      let shared1 =
          let shared2 =
              let shared3 = 1
              in  shared3 + shared3
          in  shared2 + shared2
      in  shared1 + shared1
  in  shared0 + shared0
> let shared0 =
      let shared1 =
          let shared2 = 2
          in  shared2 + shared2
      in  shared1 + shared1
  in  shared0 + shared0
> let shared0 =
      let shared1 = 4
      in  shared1 + shared1
  in  shared0 + shared0
> let shared0 = 8
  in  shared0 + shared0
> 16

You could say that Haskell’s let syntax enables sharing between computations; using it reduces the complexity of our tree implementation from $O(2^n)$ to $O(n)$. tree 50 now returns instantly:

> tree 50
1125899906842624

So let’s move everything over to an abstract syntax setting and see how the results translate there.

Let’s start with a minimalist language, known in some circles as Hutton’s Razor. While puny, it is sufficiently expressive to illustrate the subtleties of this whole sharing business:

data Expr =
    Lit Int
  | Add Expr Expr
  deriving (Eq, Ord, Show)

instance Num Expr where
  fromInteger = Lit . fromInteger
  (+)         = Add

eval :: Expr -> Int
eval (Lit d)     = d
eval (Add e0 e1) = eval e0 + eval e1

I’ve provided a Num instance so that we can conveniently write expressions in this language. We can use conventional notation and extract abstract syntax for free by specifying a particular type signature:

> 1 + 1 :: Expr
Add (Lit 1) (Lit 1)

And of course we can use eval to evaluate things:

> eval (1 + 1 :: Expr)
2

Due to the Num instance and the polymorphic definitions of naiveTree and tree, these functions will automatically work on our expression type. Check them out:

> naiveTree 2 :: Expr
Add (Add (Lit 1) (Lit 1)) (Add (Lit 1) (Lit 1))

> tree 2 :: Expr
Add (Add (Lit 1) (Lit 1)) (Add (Lit 1) (Lit 1))

Notice there’s a quirk here: each of these functions - having wildly different complexities - yields the same abstract syntax, implying that tree is no more efficient than naiveTree when it comes to dealing with this expression type. That means..

> eval (tree 50 :: Expr)
-- ain't happening

So there is a big problem here: Haskell’s let syntax doesn’t carry its sharing over to our embedded language. Equivalently, the embedded language is not expressive enough to represent sharing in its own abstract syntax.

There are a few ways to get around this.

Memoizing Evaluation

For some interpretations (like evaluation) we can use a memoization library. Here we can use Data.StableMemo to define a clean and simple evaluator:

import Data.StableMemo

memoEval :: Expr -> Int
memoEval = go where
  go = memo eval
  eval (Lit i)     = i
  eval (Add e0 e1) = go e0 + go e1

This will very conveniently handle any grimy details of caching intermediate computations. It passes the tree 50 test just fine:

> memoEval (tree 50 :: Expr)
1125899906842624

Some other interpretations are still inefficient; a similar memoPrint function will still dump out a massive syntax tree due to the limited expressiveness of the embedded language. The memoizer doesn’t really allow us to observe sharing, if we’re interested in doing that for some reason.

We can actually use GHC’s internal sharing analysis to recover any implicit sharing present in an embedded expression. This is the technique introduced by Andy Gill’s Type Safe Observable Sharing In Haskell and implemented in the data-reify library on Hackage. It’s as technically unsafe as it sounds, but in practice has the benefits of being both relatively benign and minimally intrusive on the existing language.

Here is the extra machinery required to observe implicit sharing in our Expr type:

{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE TypeFamilies #-}

import Control.Applicative
import Data.Reify hiding (Graph)
import qualified Data.Reify as Reify
import System.IO.Unsafe

data ExprF e =
    LitF Int
  | AddF e e
  deriving (Eq, Ord, Show, Functor)

instance MuRef Expr where
  type DeRef Expr        = ExprF
  mapDeRef f (Add e0 e1) = AddF <$> f e0 <*> f e1
  mapDeRef _ (Lit v)     = pure (LitF v)

We need to make Expr an instance of the MuRef class, which loosely provides a mapping between the Expr and ExprF types. ExprF itself is a so-called ‘pattern functor’, which is a parameterized type in which recursive points are indicated by the parameter. We need the TypeFamilies pragma for instantiating the MuRef class, and DeriveFunctor eliminates the need to manually instantiate a Functor instance for ExprF.

Writing MuRef instances is pretty easy. For more complicated types you can often use Data.Traversable.traverse in order to provide the required implementation for mapDeRef (example).

With this in place we can use reifyGraph from data-reify in order to observe the implicit sharing. Let’s try this on a bite-sized tree 2 and note that it is an IO action:

> reifyGraph (tree 2 :: Expr)
let [(1,AddF 2 2),(2,AddF 3 3),(3,LitF 1)] in 1

Here we get an abstract syntax graph - rather than a tree - and sharing has been made explicit.

We can write an interpreter for expressions by internally reifying them as graphs and then working on those. reifyGraph is an IO action, but since its effects are pretty tame I don’t feel too bad about wrapping calls to it in unsafePerformIO. Interpreting these graphs must be handled a little differently from interpreting a tree; a naive ‘tree-like’ evaluator will eliminate sharing undesirably:

naiveEval :: Expr -> Int
naiveEval expr = gEval reified where
  reified = unsafePerformIO $ reifyGraph expr
  gEval (Reify.Graph env r) = go r where
    go j = case lookup j env of
      Just (AddF a b) -> go a + go b
      Just (LitF d)   -> d
      Nothing         -> 0

This evaluator fails the tree 50 test:

> naiveEval (tree 50)
-- hang

We need to use a more appropriately graph-y method to traverse and interpret this (directed, acyclic) graph. Here’s an idea:

topologically sort the graph, yielding a linear ordering of vertices such that for every edge $u \to v$, $v$ is ordered before $u$.
iterate through the sorted vertices, interpreting them as desired and storing the interpretation
look up the previously-interpreted vertices as needed

We can use the Data.Graph module from the containers library to deal with the topological sorting and vertex lookups. The following graph-based evaluator gets the job done:

import Data.Graph
import Data.Maybe

graphEval :: Expr -> Int
graphEval expr = consume reified where
  reified = unsafePerformIO (toGraph <$> reifyGraph expr)
  toGraph (Reify.Graph env _) = graphFromEdges . map toNode $ env
  toNode (j, AddF a b) = (AddF a b, j, [a, b])
  toNode (j, LitF d)   = (LitF d, j, [])

consume :: Eq a => (Graph, Vertex -> (ExprF a, a, b), c) -> Int
consume (g, vmap, _) = go (reverse . topSort $ g) [] where
  go [] acc = snd $ head acc
  go (v:vs) acc =
    let nacc = evalNode (vmap v) acc : acc
    in  go vs nacc

evalNode :: Eq a => (ExprF a, b, c) -> [(a, Int)] -> (b, Int)
evalNode (LitF d, k, _)   _ = (k, d)
evalNode (AddF a b, k, _) l =
  let v = fromJust ((+) <$> lookup a l <*> lookup b l)
  in  (k, v)

In a serious implementation I’d want to use a more appropriate caching structure and avoid the partial functions like fromJust and head, but you get the point. In any case, this evaluator passes the tree 50 test without issue:

> graphEval (tree 50)
1125899906842624

Instead of working around the lack of sharing in our language, we can augment it by adding the necessary sharing constructs. In particular, we can add our own let-binding that piggybacks on Haskell’s let. Here’s an enhanced language (using the same Num instance as before):

data Expr =
    Lit Int
  | Add Expr Expr
  | Let Expr (Expr -> Expr)

The new Let constructor implements higher-order abstract syntax, or HOAS. There are some immediate consequences of this: we can’t derive instances of our language for typeclasses like Eq, Ord, and Show, and interpreting everything becomes a bit more painful. But, we don’t need to make any use of data-reify in order to share expressions, since the language now handles that 'a la carte. Here’s an efficient evaluator:

eval :: Expr -> Int
eval (Lit d)     = d
eval (Add e0 e1) = eval e0 + eval e1
eval (Let e0 e1) =
  let shared = Lit (eval e0)
  in  eval (e1 shared)

In particular, note that we need a sort of back-interpreter to re-embed shared expressions into our language while interpreting them. Here we use Lit to do that, but this is more painful if we want to implement, say, a pretty printer; in that case we need a parser such that print (parse x) == x (see here).

We also can’t use the existing tree function. Here’s the HOAS equivalent, which is no longer polymorphic in its return type:

tree :: (Num a, Eq a) => a -> Expr
tree 0 = 1
tree n = Let (tree (n - 1)) (\shared -> shared + shared)

Using that, we can see that sharing is preserved just fine:

> eval (tree 50)
1125899906842624

Another way to make sharing explicit is to use a parameterized HOAS, known as PHOAS. This requires the greatest augmentation of the original language (recycling the same Num instance):

data Expr a =
    Lit Int
  | Add (Expr a) (Expr a)
  | Let (Expr a) (a -> Expr a)
  | Var a

eval :: Expr Int -> Int
eval (Lit d)     = d
eval (Var v)     = v
eval (Add e0 e1) = eval e0 + eval e1
eval (Let e f)   = eval (f (eval e))

Here we parameterize the expression type and add both Let and Var constructors to the language. Sharing expressions explicitly now takes a slightly different form than in the HOAS version:

tree :: (Num a, Eq a) => a -> Expr b
tree 0 = 1
tree n = Let (tree (n - 1)) ((\shared -> shared + shared) . Var)

The Var term wraps the intermediate computation, which is then passed to the semantics-defining lambda. Sharing is again preserved in the language:

> eval $ tree 50
1125899906842624

Here, however, we don’t need the same kind of back-interpreter that we did when using HOAS. It’s easy to write a pretty-printer that observes sharing, for example (from here):

text e = go e 0 where
  go (Lit j)     _ = show j
  go (Add e0 e1) c = "(Add " ++ go e0 c ++ " " ++ go e1 c ++ ")"
  go (Var x) _     = x
  go (Let e0 e1) c = "(Let " ++ v ++ " " ++ go e0 (c + 1) ++
                     " in " ++ go (e1 v) (c + 1) ++ ")"
    where v = "v" ++ show c

Which yields the following string representation of our syntax:

> putStrLn . text $ tree 2
(Let v0 (Let v1 1 in (Add v1 v1)) in (Add v0 v0))

Cluing up

I’ve gone over several methods of handling sharing in embedded languages: an external memoizer, observable (implicit) sharing, and adding explicit sharing via adding a HOAS or PHOAS let-binding to the original language. Some may be more convenient than others, depending on what you’re trying to do.

I’ve dumped code for the minimal, HOAS, and PHOAS examples in some gists.

Basic EC2 Management with Ansible

2013-12-21T00:00:00+04:00

(UPDATE 2016/08/15: Here be monsters. This code is ancient, the style is not really idiomatic Ansible, and it’s likely that nothing works anymore.)

EC2 is cool. The ability to dynamically spin up a whack of free-to-cheap server instances anywhere in the world at any time, is.. well, pretty mean. Need to run a long computation job? Scale up a distributed system? Reduce latency to clients in a particular geographical region? YEAH WE CAN DO THAT.

The EC2 Management Console is a pretty great tool in of itself. Well laid-out and very responsive. But for a hacker’s hacker, a great tool to manage EC2 instances (amongst other things) is Ansible, which provides a way to automate tasks over an arbitrary number of servers, concurrently.

With EC2 and Ansible you can rapidly find yourself controlling an amorphous, globally-distributed network of servers that are eager to do your bidding.

Here’s a quick example that elides most of the nitty-gritty details. I’m going to spin up three micro instances in Asia Pacific. To do that, I’m going to use an Ansible playbook, which is essentially a YAML file that describes a sequence of commands to be performed. I’m going to delegate my local machine to handle that task, so I’m first going to store the following inventory in /etc/ansible/local:

[localhost]
127.0.0.1

The following playbook, spin-up-instances.yml is what actually launches these guys. Here’s its header:

---
- name: Spin up some EC2 instances
  hosts: 127.0.0.1
  connection: local
  tasks:
    - name: Create security group
      local_action:
        module: ec2_group
        name: my-security-group
        description: Access my-security-group
        region: ap-southeast-2
        rules:
          - proto: tcp
            from_port: 22
            to_port: 22
            cidr_ip: 0.0.0.0/0

    - name: Launch instances
      local_action:
        module: ec2
        region: ap-southeast-2
        keypair: jtobin-aws
        group: my-security-group
        instance_type: t1.micro
        image: ami-3d128f07
        count: 3
        wait: yes
      register: ec2

    - name: Add instances to host group
      local_action: add_host hostname= groupname=my-security-group
      with_items: ec2.instances

    - name: Tag instances
      local_action: ec2_tag resource= region=ap-southeast-2 state=present
      with_items: ec2.instances
      args:
        tags:
          Name: Abrek

    - name: Give everyone a minute
      pause: minutes=1

Roughly, the tasks I want performed are each declared with a name and follow the ‘tasks:’ line. They’re relatively self-explanatory. When Ansible runs this playbook, it will execute the tasks in the order they appear in the playbook.

First I create a security group in Asia Pacific (Sydney) for all the instances I want to launch, and then go ahead and actually launch the instances. You can see that I launch them as a local action on my machine. I’m using micro instances (the most lightweight instance type available) and pick an Ubuntu LTS Server machine image for each. I then do some bookkeeping and tag each instance with the name ‘Abrek’. The final task just pauses the playbook execution long enough for the instances to get up and running.

Fun fact: ‘Abrek’ was the name of one of the first two Soviet monkeys shot into space. Today is apparently the 30th anniversary of his safe return.

Now, I also want to install some software on each of these guys. I’ll separate all that into two groups: some essentials, and a specialized stack consisting of 0MQ and supporting libraries. To do that, I’ll create two separate files called ‘install-essentials.yml’ and ‘install-specialized.yml’.

I’ll keep the essentials bare for now: git, gcc/g++, and make. Here’s install-essentials.yml:

---
- name: Install git
  apt: pkg=git update_cache=yes

- name: Install gcc
  apt: pkg=gcc

- name: Install g++
  apt: pkg=g++

- name: Install make
  apt: pkg=make

I can grab all of those via apt. ‘update_cache’ is equivalent to ‘apt-get update’, which only needs to be done once.

Next, the specialized stuff in install-specialized.yml:

---
- name: Grab 0MQ
  command: >
    wget http://download.zeromq.org/zeromq-4.0.3.tar.gz
    creates=zeromq-4.0.3.tar.gz

- name: Unpack 0MQ
  command: >
    tar -xzf zeromq-4.0.3.tar.gz
    creates=zeromq-4.0.3

- name: Get libsodium
  command: >
    wget https://download.libsodium.org/libsodium/releases/libsodium-0.4.5.tar.gz
    creates=libsodium-0.4.5
    chdir=zeromq-4.0.3

- name: Install libsodium
  shell: >
    tar xzf libsodium-0.4.5.tar.gz;
    cd libsodium-0.4.5;
    ./configure && make && make check && make install
    chdir=zeromq-4.0.3

- name: Install 0MQ
  shell: >
    ./configure; make; make install
    chdir=zeromq-4.0.3

- name: Install libtool
  apt: pkg=libtool

- name: Install automake
  apt: pkg=automake

- name: Install automake
  apt: pkg=autoconf

- name: Install uuid-dev
  apt: pkg=uuid-dev

- name: Grab CZMQ
  command: >
    wget http://download.zeromq.org/czmq-2.0.3.tar.gz
    creates=czmq-2.0.3.tar.gz

- name: Unpack CZMQ
  command: >
    tar xzf czmq-2.0.3.tar.gz
    creates=czmq-2.0.3

- name: Install CZMQ
  shell: >
    ./configure && make;
    ldconfig
    chdir=czmq-2.0.3

Lots going on here. I use a variety of apt and shell commands to download and install everything I need.

Now to add those tasks back into the spin-up-instances.yml playbook so that the software gets installed right after the instances boot up. I can append the following to that file:

- name: Install essential and specialized software
  hosts: my-security-group
  user: ubuntu
  sudo: True
  tasks:
    - include: tasks/install-essentials.yml
    - include: tasks/install-specialized.yml

Let’s run the playbook and see those instances get launched. I need to use the ‘local’ inventory that I set up, so I pass that to ‘ansible-playbook’ explicitly.

Running it, we can see the security group being created, the instances popping up, and tags getting assigned:

Our essential software getting pulled down:

And the tail end of our 0MQ stack showing up before the play ends with a summary.

For a quick sanity check to ensure that everything really did go as planned, I can look for the CZMQ header on each instance. This time I’ll run a quick ad-hoc command, identifying the hosts via the ‘Abrek’ tag:

Voila, three servers ready to roll. Great stuff.

To fill the missing details, you might want to check out the excellent Ansible documentation, as well as the great tutorials at AnswersForAws.