Nicolas Audinet de Pieuchon

Re-sampling with correlation

2025-10-26T00:00:00+00:00

Thank you to my colleague and friend Ahmet Balcioglu for his feedback on this post! I also used Anthropic’s Sonnet 4.5 and OpenAI’s GPT-5 to brainstorm ideas for how to solve the problem.

Recently, as part of my research effort to replicate and understand the SHIFT method from the Sparse Feature Circuits paper by Marks et al, I found a way to create custom datasets with specific Pearson correlation coefficient between two binary labels. The aim of the blog post is to show you how I did it in a friendly and accessible way, and to give some intuitions about Pearson correlation coefficients for binary variables!

The Problem Setup

The problem has three inputs:

A dataset $\DD$ and with two binary labels $X \in \{0,1\}$ and $Y \in \{0,1\}$
A desired correlation coefficient $-1 \leq \rho \leq 1$
A desired number of samples $n$ for the output dataset

and a single output:

A dataset $\dd$ of size $n$ made up of samples from $\DD$ and where the $X$ and $Y$ labels have a Pearson correlation of $\rho$

Additionally, since I want to use $\dd$ to train machine learning models, I want it to have the following properties:

Property 1: Samples in $\dd$ should be independent and identically distributed to ensure that any statistics or machine learning models trained on the dataset are valid.

Property 2: There should be a roughly equal number of samples in each class for both $X$ and $Y$ to avoid arbitrary class imbalances influencing the performance of my machine learning model. In other words, I want that in $\dd$:

\[\begin{align} P(X=0) \approx P(X=1) & \approx 1/2 \\ P(Y=0) \approx P(Y=1) & \approx 1/2 \\ \end{align}\]

The Method

Since we’re dealing with a population with two binary labels, we can think of it in terms of four distinct sub-populations where each sub-population is defined by a particular combination of $X$ and $Y$. We can visualise the sub-populations for the input dataset $\DD$ as a contingency table:

where:

Each cell represents the sub-population for a particular label pair
$A$, $B$, $C$, and $D$ are the number of samples in each sub-population
$N = A + B + C + D$ is the total number of samples in the population
$p_A$, $p_B$, $p_C$, and $p_D$ are the estimated probabilities for belonging to a particular sub-population

The population defined by $\DD$ will also come with a set correlation between the $X$ and $Y$ labels which we do not control. We can use the following formula to approximate it based on the estimated probabilities of the sub-populations:

\[\rho_\DD = \frac{p_A p_D - p_B p_C}{\sqrt{(p_A + p_B)(p_A + p_C)(p_B + p_D)(p_C + p_D)}}\]

This is also known as the Phi coefficient and corresponds exactly to the Pearson correlation coefficient between two binary variables. Note that we could also directly used the sizes of the sub-populations rather than the estimated probabilities and reached the same result, since the factor of $1/N$ in the estimated probabilities cancels out.

If we didn’t care about the correlation between $X$ and $Y$ in the output dataset $\dd$, another thing we could do is to resample $\DD$ to create a new population from the same population distribution. Resampling is a common and widely-used technique to create confidence intervals (bootstrapping) or to validate machine learning models (cross-validation).

Since the sub-populations defined by $X$ and $Y$ are partitions, meaning that every member of the population belongs to one and only one sub-population, we can use stratified sampling to resample the population. To draw a new sample, we first choose one of the sub-populations based on the estimated probabilities, and then draw a sample from within the chosen sub-population uniformly at random. We can then repeat this process $n$ times with replacement to create the resampled dataset. Sampling with replacement here is important to guarantee Property 1, as otherwise the probability of drawing a sample would depend on previously drawn samples. This could result in drawing the same sample more than once in the resampled population, but a few repeated samples are typically not a big issue when training machine learning models so this is OK for our purposes.

More importantly, however, resampling in this way will lead to populations with the same correlation coefficient as $\DD$, which is not what we want. Instead, our goal is to create a new dataset where $X$ and $Y$ have any pre-determined correlation! To achieve this, we have to engineer a new population distribution where the labels have the correlation we desire and which we can sample from using the sub-populations from our original population. In our case this amounts to changing $p_A$, $p_B$, $p_C$ and $p_D$ to some new $p_a$, $p_b$, $p_c$ and $p_d$ based on the desired correlation, and then using the same resampling method as before with the new probabilities. But how should we modify the probabilities?

Lets start simple: what happens if we set all the sub-population probabilities to be the same? Since probabilities have to sum to 1 we have to set them to:

\[p_a = p_b = p_c = p_d = \frac{1}{4}\]

and we get the following correlation in the resampled population:

\[\rho_\dd = \frac{ \left( \frac{1}{4} \cdot \frac{1}{4} \right)^2 - \left( \frac{1}{4} \cdot \frac{1}{4} \right)^2 }{ \sqrt{ \left( \frac{1}{4} + \frac{1}{4} \right)^4} } = \frac{0}{\left( \frac{1}{2} \right) ^2} = 0\]

In this case, since we’ve perfectly balanced the probabilities between each sup-population, the two terms in the numerator cancel out to give us a correlation of 0, which implies that the $X$ and $Y$ labels are linearly independent in the resampled population. If we want to introduce arbitrary correlation we therefore need to skew the probability mass into one of the two terms.

A natural next question is then: what happens when we concentrate all the probability mass into one of the two terms? For instance, we could allocate all the probability mass into the $p_a p_d$ term. Then, we could allocate the probabilities like so:

\[p_a = p_d = \frac{1}{2} \quad \text{and} \quad p_b = p_c = 0\]

and we get the following correlation:

\[\rho_\dd = \frac{ \left( \frac{1}{2} \right)^2 - 0 }{ \sqrt{ \left( \frac{1}{2} + 0 \right)^4 } } = \frac{ \left( \frac{1}{2} \right)^2 }{ \left( \frac{1}{2} \right)^2 } = 1\]

Conversely, we could allocate all the probability mass to the $p_b p_c$ term. This means setting the probabilities to:

\[p_a = p_d = 0 \quad \text{and} \quad p_b = p_c = \frac{1}{2}\]

and results in the following correlation:

\[\rho_\dd = \frac{ 0 - \left( \frac{1}{2} \right)^2 }{ \sqrt{ \left( \frac{1}{2} + 0 \right)^4 } } = \frac{ - \left( \frac{1}{2} \right)^2 }{ \left( \frac{1}{2} \right)^2 } = -1\]

Intuitively, this makes sense. $p_a$ and $p_d$ are the probabilities for the sub-populations where $X = Y$, so if samples can only come from those two sub-populations then $X$ and $Y$ in $\dd$ will be exactly the same and exhibit a correlation of 1. Similarly, $p_b$ and $p_c$ are the probabilities for the sub-populations where $X \neq Y$, so if samples can only be selected from those sub-populations then $X$ and $Y$ in $\dd$ will have a perfect negative correlation of -1.

From this, we can conclude that to get an arbitrary correlation in $\dd$ we somehow need to balance the probability mass between the two groups of probabilities. To do this, we can imagine starting from equal probabilities and then symmetrically adding some amount to $p_a$ and $p_d$ and taking the same amount away from $p_b$ and $p_c$. We can do it like this:

\[p_a = p_d = \frac{1 + \delta}{4} \quad \text{and} \quad p_b = p_c = \frac{1 - \delta}{4}\]

where $\delta \in \RR$ and $-1 \leq \delta \leq 1$, since probabilities have to be between 0 and 1. This is nice because it guarantees that the probabilities add up to 1 and that Property 2 holds:

\[P(X=0) = P(X=1) = P(Y=0) = P(Y=1) = \frac{(1 + \delta) + (1 - \delta)}{4} = \frac{1}{2}\]

Substituting our values into the Phi coefficient formula we can find the following relationship:

\[\rho_\dd = \frac{ \left( \frac{1+\delta}{4} \right)^2 - \left( \frac{1-\delta}{4} \right)^2}{\sqrt{ \left( \left( \frac{1 + \delta}{4} \right) + \left( \frac{1 - \delta}{4} \right) \right)^4}} = \frac{ \frac{1}{4^2} \left( 1 + 2\delta + \delta^2 - 1 +2\delta - \delta^2 \right)}{\frac{1}{4^2} \left( 1 + \delta + 1 - \delta \right)^2} = \frac{4\delta}{4} = \delta\]

Neat! In order to design a new population distribution with a particular correlation $\rho$, all we need to do is to set:

\[p_a = p_d = \frac{1 + \rho}{4} \quad \text{and} \quad p_b = p_c = \frac{1 - \rho}{4}\]

and the use stratified sampling on the original population with the new probabilities. Note however that for this method to work we need there to be samples in each sub-population in the original population to sample from.

An Implementation

Here’s an example for how to implement the algorithm in Python:

import numpy as np
import pandas as pd

def resample_correlated(
    data: pd.DataFrame,
    correlation: float,
    num_samples: int,
) -> pd.DataFrame:

    msg = f"Correlation {correlation} should be between 0 and 1"
    assert -1.0 <= correlation <= 1, msg

    A = data[(data["profession"] == 0) & (data["gender"] == 0)]
    B = data[(data["profession"] == 0) & (data["gender"] == 1)]
    C = data[(data["profession"] == 1) & (data["gender"] == 0)]
    D = data[(data["profession"] == 1) & (data["gender"] == 1)]

    probabilities = [
        (1 + correlation) / 4,
        (1 - correlation) / 4,
        (1 - correlation) / 4,
        (1 + correlation) / 4,
    ]

    counts = np.random.multinomial(n=num_samples, pvals=probabilities)

    a = A.sample(n=counts[0], replace=True)
    b = B.sample(n=counts[1], replace=True)
    c = C.sample(n=counts[2], replace=True)
    d = D.sample(n=counts[3], replace=True)

    return pd.concat([a, b, c, d], ignore_index=True)

In this case, we are dealing with a dataset with two binary labels: “profession” (nurse or professor) and “gender” (male or female). We use the number of examples in the sub-populations (counts) rather than the probabilities, and rather than performing one draw at a time we batch the sampling for each sub-population.

Can You Tell? - How To Play

2025-02-18T00:00:00+00:00

I recently finished the first version of Can You Tell?, a chat-based game with Large Language Models (LLMs). I’m very excited for you to try it out! You can find the game at canyoutell.eu. I have disabled account creation for testing. If you would like an account, contact me directly and I will make one for you!

Note: Don’t write anything important on here! This game is still in its testing phase and is unstable.

How To Play

Can You Tell? is a game where users take turns to send each other messages:

Every once in a while, rather than sending a normal message, you will be asked to write a prompt for a Large Language Model (LLM) instead:

The LLM will then write the message for you based on your prompt and previous messages:

At any time, you or your opponent can make a guess as to whether a message came from you or an LLM. Guesses are made by clicking on one of the two icons under messages from your opponent:

If a guess is correct, your opponent will lose a star. However, if a guess is incorrect then you will lose a star instead. The goal of the game is to make your opponent run out stars.

Note: There are still a few bugs I’m working out. If the game stops working as intended for whatever reason, do a hard refresh and log back in. This should fix most issues.

Starting a game

Once you login, you will be met with the main game page. The column titled Chats contains the current chats with other users.

To start a new game, press the + New Game button on the bottom left:

This will bring you to the Find User page:

To find a user to chat with, type their name (or part of it) in the input field and press the magnifying class button. This will bring up a list of names that match:

Clicking on a user will will bring you to the New Chat page. To initiate a new chat, write your initial message in the text input at the bottom of the page and press Enter (or click the send button).

You can only play a single game at a time with another user.

Additional details

The LLM used for the game is llama-3.3-70b-versatile provided by Groq.

Every LLM prompt will also have the following additional system prompt added to it:

You are playing a game where you have to imitate what a human would write in a chat. You will be given a chat history and a prompt. Write short, human-like responses based on the prompt

Water Sort in Haskell

2024-10-14T00:00:00+00:00

Water Sort is a puzzle game where you have to sort coloured water into bottles. If you haven’t tried it yet, there are many clones of it online or on your phone. I started playing the game a while ago and, after getting a little addicted to it, I decided to implement it in Haskell for fun (and to immunise myself against it).

In this post I’ll give an overview of how I implemented the game with a simple terminal-based UI and a list-based solver which takes advantage of laziness. My goal is to show people that are perhaps not very familiar with Haskell that implementing a simple game like this can be fun and rewarding, and to share a couple cool tricks I learned along the way.

All the code for the game is on GitHub, and I will be linking to various files there throughout the post.

Water Sort

The game starts with a bunch of bottles with colour layers stacked on top of one other along with some empty bottles. The version of the game I played has bottles with four layers of colours and two empty bottles:

The goal of the game is to sort the colours into single bottles:

You progress in the game by pouring the top colour layer of a bottle into another bottle:

However, pouring has some simple restrictions:

You can only pour a colour into an empty bottle or a bottle with the same colour on top
You cannot overfill bottles

Implementing the game

To implement the game I used a simple variation of the model-view-update pattern, also known as the Elm architecture. In brief, the model is the data structure containing the state of the game, the view is a function which takes the model and renders it to the screen, and the update is a function which takes user inputs and adjusts the model accordingly. The game itself consists of a game loop that repeats either until the bottles are sorted (the player wins) or there are no more available actions (the player loses).

The Model

My first step was to model the game state with data types. The whole game state is captured by the Bottles type, which is a Map of bottles to their identifiers:

type Bottles = M.Map BottleId Bottle

A Bottle is modelled as a list of Colors, where the head of the list corresponds to the topmost colour in the bottle. A BottleId is just a type alias for an Int:

type BottleId = Int
type Bottle = [Color]

Since the game will only support levels of a fixed size, our model only needs to represent about a dozen colours, so I decided to model Color with a sum type:

data Color
  = Yellow
  | White
  | Red
  | LightBlue
  | LightGreen
  | Pink
  | Brown
  | DarkGreen
  | Orange
  | DarkBlue
  | DarkRed

The full code for the model can be found in the Model.hs module.

The View

For simplicity, I chose to use the terminal as the UI for my implementation.

To print colours to the terminal, I used two space characters (which on my terminal roughly form a square) and coloured them using xterm-256 colour codes . The final outcome looks like this:

This worked well, but requires one to have a terminal that supports 256 colours and a mono-space font so that two spaces roughly form a square. In code, the function to print a colour looks like this:

showColor :: Color -> String
showColor color = "\x1b[48;5;" <> show (colorCode color) <> "m  \x1b[0m"
  where
    colorCode :: Color -> Int
    colorCode Yellow = 220
    colorCode LightBlue = 44
    colorCode DarkBlue = 21
    colorCode Brown = 130
    colorCode LightGreen = 47
    colorCode DarkGreen = 22
    colorCode Pink = 201
    colorCode White = 255
    colorCode Red = 196
    colorCode Orange = 208
    colorCode DarkRed = 88

colorCode here is a mapping from Color to xterm-256 colour codes. The weird-looking strings on either side of the colour code are the escape codes to set the background colour.

After figuring out how to print a colour, the next challenge was to print the bottles side by side. This is made a little tricky by the fact that the terminal can only print one line at time from top to bottom. My solution was to first transform our Bottles map into a list of Lines and then print them out in order. I modelled a Line as a list of Squares, where each square is either empty, a separator or a colour:

data Square = Empty | Separator | Full Color
type Line = [Square]

Then, I used an unfoldr to generate the list of Lines:

bottlesToGrid :: Bottles -> [Line]
bottlesToGrid = reverse . unfoldr makeLine . map reverse . M.elems
  where
    getSquare :: Bottle -> Square
    getSquare = maybe Empty Full . headMaybe

    makeLine :: [Bottle] -> Maybe (Line, [Bottle])
    makeLine xs
      | all null xs = Nothing
      | otherwise = Just (map getSquare xs, map tailSafe xs)

This function works by constructing the rows from the bottom of the bottles to the top, and then reversing the order of the rows at the end. The unfoldr uses a makeLine function which takes a list of Bottles and constructs a single line from it. The line is constructed by mapping over the bottles with the getSquare function, which returns an Empty square if the bottle is empty and a Full square with the colour if it’s not. If all the bottles are empty, makeLine returns Nothing to signal the end of the unfold. I also used two helper functions:

headMaybe :: [a] -> Maybe a is a safe version of head that returns Nothing when given an empty list
tailSafe :: [a] -> [a] is a version of tail that returns an empty list when given an empty list.

The final step was to actually make the functions to convert the lines to strings which can then be printed to the terminal:

showSquare :: Square -> String
showSquare Empty = "  "
showSquare Separator = "|"
showSquare (Full color) = showColor color

showLine :: Line -> String
showLine line =
  let sepLine = [Separator] <> intersperse Separator line <> [Separator]
   in concatMap showSquare sepLine

Check out View.hs file for the full code to show the game.

The Update

Players in this game can only make a single type of action: a pour from one bottle to another. Pours are modelled with a product type of two bottle ids:

data Pour = Pour
  { from :: BottleId
  , to :: BottleId
  }

Each turn, the user is asked to input a pour as two numbers separated by an arrow (for example "2 -> 3"). The game then takes the line as input, splits it on the "->" string and parses the two numbers on either side:

parsePour :: String -> Either GameError Pour
parsePour line = case splitOn "->" line of
  [from, to] -> do
    fromBottle <- maybeThrow (InvalidBottleId from) (readMaybe from)
    toBottle <- maybeThrow (InvalidBottleId to) (readMaybe to)
    pure (Pour fromBottle toBottle)
  _ -> Left (InvalidInput line)

Parsing errors are handled explicitly by returning an Either GameError Pour. GameError here is a sum type that captures all the logical errors that can occur in the game:

data GameError
  = InvalidPuzzleType String
  | InvalidBottleId String
  | InvalidInput String
  ...

In larger projects it could be a good idea to split up error types into several smaller types, but for a small project like this I decided that having a single error type was good enough.

I also chose to use Either as a simple mechanism to throw and handle errors in pure functions. “Throwing an error” is implemented by returning a Left, which then short-circuits the rest of the function thanks to the Either monad instance. Errors that crop up are handled in the top-level game loop. I also made two functions for convenience:

maybeThrow :: GameError -> Maybe a -> Either GameError a
maybeThrow err Nothing = Left err
maybeThrow _ (Just a) = Right a

whenThrow :: Bool -> GameError -> Either GameError ()
whenThrow True err = Left err
whenThrow False _ = Right ()

Once a Pour is received and parsed from the user it is passed to the update function which contains the main model update logic:

update :: Bottles -> Pour -> Either GameError Bottles
update bottles pour@(Pour from to) = do
  (fromBottle, toBottle) <- getPourBottles bottles pour
  (fromHead, fromTail) <- validate pour fromBottle toBottle
  let insertFrom = M.insert from fromTail
  let insertTo = M.insert to (fromHead <> toBottle)
  pure . insertFrom . insertTo $ bottles

This function essentially get the two bottles from the game state, check that the pour is valid, splits the colours in the from bottle into the part that will be poured into the to bottle and the part that will stay, and then updates the game state with the new bottles. getPourBottles and validate help with doing that and contain most of the logic that ensures that the pour follows what is allowed by the rules of the game.

The full code for the update is in the Update.hs module.

Putting it all together

Time to combine the components into the game loop! The body of the loop is contained in the step function:

step :: Bottles -> IO Bottles
step bottles = do
  putStrLn (showGame bottles)
  line <- getLine
  case (parsePour line >>= update bottles) of
    Right newBottles -> pure newBottles
    Left gameError -> do
      print gameError
      pure bottles

In other words, for each iteration:

Show the current state of the game
Ask the user for input
Parse the input
Update the game state based on the input

If an error comes up in step 3 or 4, the game shows the error to the user and reverts back to the start of the turn.

The last part is the actual game loop itself, which I implemented as a recursive function which repeatedly calls itself until the game is over:

loop :: Bottles -> IO Bottles
loop bottles = do
  newBottles <- step bottles
  if gameOver newBottles
    then pure newBottles
    else loop newBottles

The code for the game loop is in Main.hs.

Creating puzzles

Now that I had a working game, I needed a way to generate new puzzles so that users (me, myself and I) can play the game to their heart’s content. To help learn the game, I also decided to support three puzzle sizes:

data PuzzleSize = Small | Medium | Large

Given a puzzle size, the next thing to do is to pick which (and how many) colours to use in the puzzle. I arbitrarily decided on 4, 7, and 10 colours for the small, medium, and large puzzles respectively, and used toEnum from the Enum class to make the list of colours of the appropriate size:

colorPalette :: PuzzleSize -> [Color]
colorPalette Small = map toEnum [0 .. 3]
colorPalette Medium = map toEnum [0 .. 6]
colorPalette Large = map toEnum [0 .. 9]

The puzzles start with the same number of full bottles as the number of colours. Mirroring the version of the game I’ve been playing, I decided to use bottles of height 4 and to add 2 extra empty bottles at the start.

I ended up with the following function to make a new random assortment of bottles:

randomBottles :: MonadRandom m => PuzzleSize -> m Bottles
randomBottles size = do
  let initColors = concat (replicate 4 (colorPalette size))
  randomColors <- shuffle initColors
  let bottles = chunksOf 4 randomColors <> [[], []]
  pure (M.fromList $ zip [0 ..] bottles)

In words, randomBottles:

Makes a list with 4 copies of each colour
Shuffles it
Splits it into a list of bottles of height 4
Adds two empty bottles
Converts the list into a Map

But, you might ask: how do you know that the puzzle is actually solvable? Interestingly, after some Googling I discovered that some mathematicians actually studied this problem and found some bounds for the minimum number of empty bottles needed to ensure the puzzle is always solvable. Plugging in our parameters, two empty bottles seems to be the exact lower bound for all three puzzle sizes. Cool! But that still doesn’t guarantee that the puzzle has a solution.

Instead, I went for a cruder approach and made a function that makes random puzzles and tries to solve them until it finds one that has a solution:

createPuzzle :: MonadRandom m => PuzzleSize -> m Bottles
createPuzzle size = do
  bottles <- randomBottles (sizeToInt size)
  if isJust (solve bottles)
  then pure bottles
  else createPuzzle size

In practice most random puzzles seem to be solvable, and in the worst case where the puzzle is not solvable or the solver is being really slow the user can just restart the game.

Solving the game

To solve the game one needs to find a sequence of pours that will sort an initial set of bottles (where possible). I implemented the solver as follows:

solve :: Bottles -> Maybe [Pour]
solve = fmap (reverse . history) . headMaybe . solutions . SolverState []

The first thing this function does is to transform the input Bottles into a SolverState, which is a record of the current game state and the history of pours that led to it:

data SolverState = SolverState
  { history :: [Pour]
  , current :: Bottles
  }

The function then passes the initial solver state to the solutions function, which is where most of the magic happens:

solutions :: SolverState -> [SolverState]
solutions state
  | gameOver (current state) = pure state
  | otherwise = do
      (pour, newBottles) <- possiblePours (current state)
      solutions (SolverState (pour : history state) newBottles)

The solutions function returns a list of all possible solution states coupled with their pour history. Given a solver state, it first checks if the game is over using the same function as before. If the game is indeed over, it returns the current solver state and stops the recursion. Otherwise, it finds the list of all valid pours from the current state and recursively calls solutions on each of them [1]. In this way, it traverses the entire tree of possible sequences of pours, gathering all the solutions as it goes.

The list of valid pours is created by the possiblePours function, which first creates a list of all possible pours and then filters out the invalid ones using mapMaybe and tryPour:

possiblePours :: Bottles -> [(Pour, Bottles)]
possiblePours bottles =
  let bottleIds = M.keys bottles
      pours = liftA2 Pour bottleIds bottleIds
   in mapMaybe (tryPour bottles) pours

As its name implies, the tryPour function “tries” to perform a pour in the current game state using the update function from before, and returns Nothing if the pour is not successful:

tryPour :: Bottles -> Pour -> Maybe (Pour, Bottles)
tryPour bottles pour =
  case update bottles pour of
    Left _ -> Nothing
    Right newBottles -> Just (pour, newBottles)

Infinite cycles, for example where a colour is poured back and forth between two bottles forever, are prevented by an extra check in the update function that ensures these types of pours are considered invalid. In turn, this ensures that the solutions function is always guaranteed to terminate.

However, traversing the entire tree of possible pour sequences can quickly become computationally infeasible. This is where the second bit of magic comes in: thanks to Haskell’s laziness, taking the head of the list of possible solutions with headMaybe means that we actually only compute the first solution and leave the rest as thunks, effectively morphing the solver into a depth-first search [2]! Neat right?

The final step taken by solve is to extract the pour history from our solution state and to reverse it to get the correct order of pours. The full code for the solver is contained in the Solver.hs module.

Next steps

Although the game is playable at this point there are still lots of things one could improve. Here are some ideas:

Undo - Add a way to undo a move. At the moment players have no way to backtrack if they make a mistake, which can quickly get annoying
Save - Add a way to save a puzzle and come back to it later. This could even be used to share puzzles with others
Web - Make it playable on the web! It should be straightforward to port the game to PureScript and make a simple UI for it
Tests - Add some tests to ensure that our implementation actually respects the rules of the game

And that’s it! Hope you enjoyed coming on this this little journey with me and maybe even learned a new trick or two along the way.

[1] I used the monad instance for lists here, but I could have just as easily used a list comprehension instead:

solutions :: SolverState -> [SolverState]
solutions state
  | gameOver (current state) = pure state
  | otherwise =
      concat
        [ solutions (SolverState (pour : history state) newBottles)
        | (pour, newBottles) <- possiblePours (current state)
        ]

[2] I used lists and laziness in this case because it was the simplest way I could think of to implement and explain the solver. However, using lists in this way could still consume tons of memory as it keeps track of previous unsuccessful branches in memory (I think). Instead, one could use Maybe and its MonadPlus instance to implement a constant-space search.

Reshape in Hmatrix

2024-03-03T00:00:00+00:00

The goal of this post is to implement a type-safe reshape function using the Hmatrix Static API.

I used GHC 9.4.8 throughout this post. I also had to enable a whole menagerie of language extensions:

If you’re not familiar with the Hmatrix Static API, I’ve also made a more introductory post about implementing the zeros function with the Hmatrix Static API.

The goal

Reshape is a common Numpy function that gives a new shape to a multi-dimensional array without changing the data. An example from the docs:

>>> a = np.arange(6).reshape((3, 2))
>>> a
array([[0, 1],
       [2, 3],
       [4, 5]])

At the moment, the Hmatrix Static API only supports vectors (1D) and matrices (2D). Since vectors are 1D they cannot be reshaped while maintaining the same elements. Functions to embed vectors into matrices and vice versa already exist in Hmatrix, although they are limited to single-row matrices:

row :: R n -> L 1 n: transform a vector into a single-row matrix
unrow :: L 1 n -> R n: transform a single-row matrix into a vector

What doesn’t exist yet in Hmatrix is a reshape function which transforms a matrix into another matrix of a different shape. More specifically, there isn’t a function that takes a matrix L n m and transforms it into another matrix L p q where n * m = p * q. So let’s make it!

Our strategy for implementing reshape will be to first implement two other functions:

flatten :: L n m -> R (n * m): a more general version of unrow which takes a matrix with any number of rows and transforms it into vector
unflatten :: R (n * m) -> L n m: a more general version of row which takes a vector and wraps it into a matrix

reshape will then be a defined as a composition of flatten and unflatten, which will end up looking something like this:

reshape :: (n * m ~ p * q) => L n m -> L p q
reshape = unflatten . flatten

For the purposes of this post, we will work in a row-wise fashion rather than column-wise. For example, row-wise flattening of a 3x3 matrix will look like this:

And unflattening the array into a 3x3 matrix row-wise looks like this:

The `flatten` function

The flatten function can be thought of as a fold which recursively takes the top row of the matrix and concatenates it to the front of the flattened version of the rest of the matrix. The recursion stops when we encounter a matrix with a single row, in which case we can just use the unrow function from before. We can use natVal and Proxy to inspect the size of the matrix at the value level to decide whether to stop the recursion, and splitRows and (#) from the Static API to split the matrix and concatenate vectors, respectively. Bringing all of this together gives us our first implementation:

import Data.Proxy
import GHC.TypeLits
import Numeric.LinearAlgebra.Static

flatten :: forall n m. (KnownNat n, KnownNat m) => L n m -> R (n * m)
flatten mat =
  case natVal (Proxy @n) of
    1 -> unrow mat
    _ ->
      let (r, rest) = splitRows m :: (L 1 m, L (n - 1) m)
       in unrow r # flatten rest

Unfortunately, this doesn’t work. We get the following compiler error:

src/Reshape/Flatten.hs:19:16-18: error:
    • Couldn't match type ‘n’ with ‘1’
      Expected: L 1 m
        Actual: L n m
      ‘n’ is a rigid type variable bound by
        the type signature for:
          flatten :: forall (n :: Nat) (m :: Nat).
                     (KnownNat n, KnownNat m) =>
                     L n m -> R (n * m)
        at src/Reshape/Flatten.hs:16:1-71
    • In the first argument of ‘unrow’, namely ‘mat’
      In the expression: unrow mat
      In a case alternative: 1 -> unrow mat
    • Relevant bindings include
        mat :: L n m (bound at src/Reshape/Flatten.hs:17:9)
        flatten :: L n m -> R (n * m)
          (bound at src/Reshape/Flatten.hs:17:1)
   |
19 |     1 -> unrow mat
   |                ^^^
src/Reshape/Flatten.hs:21:23-31: error:
    • Cannot satisfy: 1 <= n
    • In the expression: splitRows mat :: (L 1 m, L (n - 1) m)
      In a pattern binding:
        (r, rest) = splitRows mat :: (L 1 m, L (n - 1) m)
      In the expression:
        let (r, rest) = splitRows mat :: (L 1 m, L (n - 1) m)
        in unrow r # flatten rest
   |
21 |       let (r, rest) = splitRows mat :: (L 1 m, L (n - 1) m)
   |                       ^^^^^^^^^

These errors are essentially telling us that, due to a limitation of the Haskell type system, we cannot implement the logic of choosing the recursive or non-recursive step at the value level. We will get similar errors if we try to implement unflatten in the same way. Instead, we need to encode the logic of the case statement at the type level, for which we will need to use type classes and some trickery.

The `Flattenable` type class

This time, we will implement flatten as a method in a type class called Flattenable:

class Flattenable n m where
  flatten :: L n m -> R (n * m)

The type class has two parameters, n and m, which represent the size of the matrix being flattened. We will also implement two instances for Flattenable:

-- the single-row case
instance Flattenable 1 m where
  flatten = unrow

-- the multi-row case
instance {-# OVERLAPS #-} Flattenable n m where
  flatten = undefined

These two instances are the key to encoding the case statement logic from before at the type level. When the compiler sees a call to flatten, it will try to match the type of the input matrix to an appropriate instance. If the input matrix has more than one row then the compiler will match the second instance (which matches matrices of any size). If the matrix has a single row then the compiler will match both instances. This is a problem, because if the compiler has more than one option to choose from it becomes unclear which one it should choose (and we want compilation to be predictable so the compiler can’t randomly choose one of them). One way to get around this is to use the {-# OVERLAPS #-} pragma. This tells the compiler that, in the case where both instances match, it can ignore the second instance in favour of the first one.

Implementing the instance for single-row matrices is straightforward: it just becomes a call to unrow. Implementing the instance for multi-row matrices is a bit trickier however. Our strategy will be the same as before: we will flatten the matrix by recursively taking the top row and concatenating it to the flattened version of the rest of the matrix. Here’s our second attempt:

instance {-# OVERLAPS #-} Flattenable n m where
  flatten mat =
    let (r, rest) = splitRows mat :: (L 1 m, L (n - 1) m)
     in unrow r # flatten rest

Again, this doesn’t work. Compiling this code gives us the following errors:

src/Reshape/Flatten.hs:49:21-29: error:
    • Cannot satisfy: 1 <= n
    • In the expression: splitRows mat :: (L 1 m, L (n - 1) m)
      In a pattern binding:
        (r, rest) = splitRows mat :: (L 1 m, L (n - 1) m)
      In the expression:
        let (r, rest) = splitRows mat :: (L 1 m, L (n - 1) m)
        in unrow r # flatten rest
   |
49 |     let (r, rest) = splitRows mat :: (L 1 m, L (n - 1) m)
   |                     ^^^^^^^^^
src/Reshape/Flatten.hs:50:9-30: error:
    • Couldn't match type: m + ((n - 1) * m)
                     with: n * m
      Expected: R (n * m)
        Actual: R (m + ((n - 1) * m))
    • In the expression: unrow r # flatten rest
      In the expression:
        let (r, rest) = splitRows mat :: (L 1 m, L (n - 1) m)
        in unrow r # flatten rest
      In an equation for ‘flatten’:
          flatten mat = let (r, rest) = ... in unrow r # flatten rest
    • Relevant bindings include
        r :: L 1 m (bound at src/Reshape/Flatten.hs:49:10)
        rest :: L (n - 1) m (bound at src/Reshape/Flatten.hs:49:13)
        mat :: L n m (bound at src/Reshape/Flatten.hs:48:11)
        flatten :: L n m -> R (n * m)
          (bound at src/Reshape/Flatten.hs:48:3)
   |
50 |      in unrow r # flatten rest
   |         ^^^^^^^^^^^^^^^^^^^^^^

The first error says that the compiler doesn’t known that the type variable n is greater than 0, which it needs to be since we call flatten on m2 :: L (n-1) m and we can’t have a matrix with a negative number of rows. To fix the error we need to add a 1 <= n constraint to prevent the compiler from matching the instance with matrices with 0 rows.

The second error tells us that the compiler doesn’t know that the equality m + ((n-1) * m) = n * m holds for numbers at the type level. By default the Haskell type system is not great at simple algebraic proofs! Thankfully we can convince the type checker that the equality holds with equality constraints), where t1 ~ t2 means that type t1 is the same as type t2.

instance
  {-# OVERLAPS #-}
  ( 1 <= n,
    m + ((n - 1) * m) ~ n * m
  ) =>
  Flattenable n m where

  flatten mat =
    let (r, rest) = splitRows mat :: (L 1 m, L (n - 1) m)
     in unrow r # flatten rest

Compiling again we get several new errors. Let’s tackle them one by one:

src/Reshape/Flatten.hs:52:21-29: error:
    • Could not deduce (KnownNat n) arising from a use of ‘splitRows’
      from the context: (1 <= n, (m + ((n - 1) * m)) ~ (n * m))
        bound by the instance declaration
        at src/Reshape/Flatten.hs:(36,3)-(49,17)
      Possible fix:
        add (KnownNat n) to the context of the instance declaration
    • In the expression: splitRows mat :: (L 1 m, L (n - 1) m)
      In a pattern binding:
        (r, rest) = splitRows mat :: (L 1 m, L (n - 1) m)
      In the expression:
        let (r, rest) = splitRows mat :: (L 1 m, L (n - 1) m)
        in unrow r # flatten rest
   |
52 |     let (r, rest) = splitRows mat :: (L 1 m, L (n - 1) m)
   |                     ^^^^^^^^^

This first error tell us that the compiler doesn’t know that n corresponds to a number (i.e. that it belongs to the KnownNat type class). We can fix this by adding a KnownNat n constraint to the instance.

src/Reshape/Flatten.hs:53:17: error:
    • Could not deduce (KnownNat m) arising from a use of ‘#’
      from the context: (1 <= n, (m + ((n - 1) * m)) ~ (n * m))
        bound by the instance declaration
        at src/Reshape/Flatten.hs:(36,3)-(49,17)
      Possible fix:
        add (KnownNat m) to the context of the instance declaration
    • In the expression: unrow r # flatten rest
      In the expression:
        let (r, rest) = splitRows mat :: (L 1 m, L (n - 1) m)
        in unrow r # flatten rest
      In an equation for ‘flatten’:
          flatten mat = let (r, rest) = ... in unrow r # flatten rest
   |
53 |      in unrow r # flatten rest
   |                 ^

Similarly, the compiler also doesn’t know that m corresponds to a number. Same fix: add a KnownNat m constraint to the instance.

src/Reshape/Flatten.hs:53:19-25: error:
    • Overlapping instances for Flattenable (n - 1) m
        arising from a use of ‘flatten’
      Matching instance:
        instance [overlap ok] (1 <= n, (m + ((n - 1) * m)) ~ (n * m)) =>
                              Flattenable n m
          -- Defined at src/Reshape/Flatten.hs:36:3
      Potentially matching instance:
        instance Flattenable 1 m -- Defined at src/Reshape/Flatten.hs:29:10
      (The choice depends on the instantiation of ‘n, m’
       and the result of evaluating ‘-’
       To pick the first instance above, use IncoherentInstances
       when compiling the other instance declarations)
    • In the second argument of ‘(#)’, namely ‘flatten rest’
      In the expression: unrow r # flatten rest
      In the expression:
        let (r, rest) = splitRows mat :: (L 1 m, L (n - 1) m)
        in unrow r # flatten rest
   |
53 |      in unrow r # flatten rest
   |                   ^^^^^^^

Here, the compiler is saying that it doesn’t know that m2 :: L (n-1) m will match an instance of Flattenable since it might have more than one row. The fix is to add a constraint for Flattenable (n-1) m. Adding this constraint also means that we can remove the {-# OVERLAPS #=} pragma since the instance won’t match with single-row matrices any more.

After adding the new constraints, our instance now looks like this:

instance
  ( 1 <= n,
    m + ((n - 1) * m) ~ n * m,
    KnownNat n,
    KnownNat m,
    Flattenable (n - 1) m
  ) =>
  Flattenable n m where

  flatten mat =
    let (r, rest) = splitRows mat :: (L 1 m, L (n - 1) m)
     in unrow r # flatten rest

Compiling our code now returns the following error:

src/Reshape/Flatten.hs:53:17: error:
    • Could not deduce (KnownNat ((n - 1) * m))
        arising from a use of ‘#’
      from the context: (1 <= n, (m + ((n - 1) * m)) ~ (n * m),
                         KnownNat n, KnownNat m, Flattenable (n - 1) m)
        bound by the instance declaration
        at src/Reshape/Flatten.hs:(36,3)-(49,17)
    • In the expression: unrow r # flatten rest
      In the expression:
        let (r, rest) = splitRows mat :: (L 1 m, L (n - 1) m)
        in unrow r # flatten rest
      In an equation for ‘flatten’:
          flatten mat = let (r, rest) = ... in unrow r # flatten rest
   |
53 |      in unrow r # flatten rest
   |                 ^

Now the compiler is telling us that it is not smart enough to recognise that if n and m are KnownNats then so is (n-1) * m. Let’s add that as yet another constraint:

instance
  ( 1 <= n,
    m + ((n - 1) * m) ~ n * m,
    KnownNat n,
    KnownNat m,
    Flattenable (n - 1) m,
    KnownNat ((n - 1) * m)
  ) =>
  Flattenable n m where

  flatten mat =
    let (r, rest) = splitRows mat :: (L 1 m, L (n - 1) m)
     in unrow r # flatten rest

And hooray, our code compiles! At last the mighty compiler has been appeased and we have successfully defined the row-wise flatten function for all matrices in a type-safe way.

The `unflatten` function

Now that we’ve implemented flatten we can extend the Flattenable type class to include unflatten as well:

class Flattenable n m where
  flatten :: L n m -> R (n * m)
  unflatten :: R (n * m) -> L n m

To recap, unflatten is a function that takes a vector and transforms it into a matrix. This operation can also be seen as a fold where we recursively take the first M elements from the input vector and stack them on top of the matrix constructed by unflattening the rest of the vector.

Like before, the instance for the single-row matrices is straightforward: all we need to do is invoke the row :: R n -> L 1 n function from the Hmatrix Static API:

instance Flattenable 1 m where
  flatten = unrow
  unflatten = row

Also like before, the instance for multi-row matrices is more complicated and involves recursively calling the unflatten function. The body of the function will look as follows, where split and (===) are functions to split vectors and stack matrices, respectively:

unflatten vec =
    let (r, rest) = split vec :: (R m, R ((n - 1) T.* m))
     in row r === unflatten rest

We will also need to provide several more constraints to the instance. These were implemented following the same process as for flatten: compile, get an error, add a constraint to fix the error, repeat.

The final code for the multi-row instance ended up looking like this:

instance
  ( 1 <= n,
    m + ((n - 1) * m) ~ n * m,
    KnownNat n,
    KnownNat m,
    Flattenable (n - 1) m,
    KnownNat ((n - 1) * m),
    -- New constraints:
    m <= n * m,
    (n * m) - m ~ m * (n - 1),
    m * (n - 1) ~ (n - 1) * m,
    1 + (n - 1) ~ n,
    KnownNat (n * m),
    KnownNat (n - 1)
  ) =>
  Flattenable n m where

  flatten mat =
    let (r, rest) = splitRows mat :: (L 1 m, L (n - 1) m)
     in unrow r # flatten rest

  unflatten vec =
    let (r, rest) = split vec :: (R m, R ((n - 1) T.* m))
     in row r === unflatten rest

A type-safe `reshape` function

Finally, we have all the pieces to implement our type-safe reshape function! Here it is in all its glory:

reshape ::
    ( Flattenable n m
    , Flattenable p q
    , n * m ~ p * q
    ) => L n m -> L p q
reshape = unflatten . flatten

Our previous work on the Flattenable type class has paid off. To reshape a matrix, all we need to do is flatten and unflatten our input into the correct shape, and the compiler take care of the rest.

Here’s the final code for everything we’ve covered in this post:

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}

import GHC.TypeLits
import Numeric.LinearAlgebra.Static

class Flattenable n m where
  flatten :: L n m -> R (n * m)
  unflatten :: R (n * m) -> L n m

instance Flattenable 1 m where
  flatten = unrow
  unflatten = row

instance
  ( 1 <= n,
    m + ((n - 1) * m) ~ n * m,
    KnownNat n,
    KnownNat m,
    Flattenable (n - 1) m,
    KnownNat ((n - 1) * m),
    m <= n * m,
    (n * m) - m ~ m * (n - 1),
    m * (n - 1) ~ (n - 1) * m,
    1 + (n - 1) ~ n,
    KnownNat (n * m),
    KnownNat (n - 1)
  ) =>
  Flattenable n m where

  flatten mat =
    let (r, rest) = splitRows mat :: (L 1 m, L (n - 1) m)
     in unrow r # flatten rest

  unflatten vec =
    let (r, rest) = split vec :: (R m, R ((n - 1) T.* m))
     in row r === unflatten rest

reshape ::
  ( Flattenable n m,
    Flattenable p q,
    n * m ~ p * q
  ) => L n m -> L p q
reshape = unflatten . flatten

Hmatrix - from zeros to hero

2024-02-11T00:00:00+00:00

The goal of this post is to give a brief introduction to hmatrix’s Static API and show how to implement a type-safe zeros function in two different ways.

I used GHC 9.4.8 throughout.

The `hmatrix` Static API

Hmatrix is a Haskell library that provides a nice functional interface for working with matrices. It is built on top of BLAS and LAPACK.

The library includes a Static API which allows one to define the size of the matrices at the type level. Here’s what creating a random 3x2 matrix looks like using the static API:

{-# LANGUAGE DataKinds #-} -- For using type-level numbers

import Numeric.LinearAlgebra.Static

main :: IO ()
main = do
    x <- rand :: IO (L 3 2)
    print x

which gives the following output:

(matrix
 [ 0.33926174432936207,  0.7822155900216734
 ,  0.9783547953601716, 0.40214259103040334
 , 0.15344813892312728,  0.2366858740508025 ] :: L 3 2)

The benefit of storing the size at the type level is that the compiler can use the information to prevent impossible operations at compile time. For example, trying to compile this program (where (<>) is matrix multiplication)

main :: IO ()
main = do
    x <- rand :: IO (L 3 2)
    y <- rand :: IO (L 5 10)
    print (x <> y)

results in the following compiler error:

exe/Main.hs:24:18: error:
    • Couldn't match type ‘5’ with ‘2’
      Expected: L 2 10
        Actual: L 5 10
    • In the second argument of ‘(LA.<>)’, namely ‘y’
      In the first argument of ‘print’, namely ‘(x LA.<> y)’
      In a stmt of a 'do' block: print (x LA.<> y)
   |
24 |   print (x LA.<> y)
   |

That is, the compiler is telling us that it is impossible to multiply a 3x2 matrix by a 5x10 matrix since 2 != 5. Admittedly, the ergonomics of the type error are not great, but the guarantees are nice!

On top of that, keeping the size information at the type level means we don’t need to check for different input sizes in the body of the function: the size of the inputs is guaranteed by the compiler.

However, the hmatrix static API suffers from lack of development and documentation. In the rest of this post we will close the gap a little by implementing the zeros function from Python’s numpy two ways.

The `zeros` function in two ways

numpy.zeros in Python is a common function for creating matrices with all elements set to 0. However, this function is missing from the Hmatrix Static API. The goal of this section is to implement it in two ways. The first method introduces some key type-level programming patters for working with the Static API. The second method shows a cleaner alternative and is included for completeness.

The type of the zeros function we want to build is:

import Numeric.LinearAlgebra.Static

zeros :: L n m
zeros = undefined

In other words, the zeros function returns a matrix with all elements set to 0 of size NxM.

Implementation using `matrix`

Our first strategy for implementing zeros will be to make a wrapper around the unsafe matrix function from the Static API:

matrix :: (KnownNat m, KnownNat n) => [Double] -> L n m

The matrix function takes a list of Doubles of size N*M and returns a matrix of size NxM. This function is unsafe, because giving matrix a list of the wrong size results in a runtime error. To solve this, we will make use of natVal from GHC.TypeLits, which allows us to bring information from the type level to the value level. For example:

ghci> :set -XDataKinds -XTypeApplications
ghci> import Data.Proxy
ghci> import GHC.TypeLits
ghci> x = Proxy @5
ghci> y = Proxy @6
ghci> natVal x + natVal y
11

Here, we first create two values x and y using Proxy. From the docs:

Proxy is a type that holds no data, but has a phantom parameter of arbitrary type (or even kind). Its use is to provide type information, even though there is no value available of that type (or it may be too costly to create one).

In the example, we use Proxy to carry numerical information at the type level: x carries the number 5 and y carries the number 6. We can then use natVal to bring the numbers carried by the proxies from the type level to the value level and add them together. We also used the TypeApplications language extension to use the @ notation for specifying the value of type variables.

In the zeros function, we can use natVal in order to automatically create a list of the correct size using the standard replicate function.

Let’s take a first stab at it:

{-# LANGUAGE TypeApplications #-}

import Numeric.LinearAlgebra.Static
import GHC.TypeLits (natVal)
import Data.Data (Proxy(..))

zeros :: L n m
zeros =
  let n_val = fromIntegral (natVal @n Proxy)
      m_val = fromIntegral (natVal @m Proxy)
   in matrix (replicate (n_val * m_val) 0)

Trying to run this code gives us the following compiler error:

src/Const/Matrix.hs:13:37: error: Not in scope: type variable ‘n’
   |
13 |   let n_val = fromIntegral (natVal @n Proxy)
   |                                     ^
src/Const/Matrix.hs:14:37: error: Not in scope: type variable ‘m’
   |
14 |       m_val = fromIntegral (natVal @m Proxy)
   |                                     ^

This error occurs because by default the compiler does not recognise that the type variable n in the let binding is the same as the type variable n in the function’s type definition (and the same with m). To fix this we can use the ScopedTypeVariables language extension which allows the body of the function to access type variables from the function’s type definition when they are explicitly introduced with the forall keyword:

{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE ScopedTypeVariables #-}

import Numeric.LinearAlgebra.Static
import GHC.TypeLits (natVal)
import Data.Data (Proxy(..))

zeros :: forall n m. L n m
zeros =
  let n_val = fromIntegral (natVal @n Proxy)
      m_val = fromIntegral (natVal @m Proxy)
   in matrix (replicate (n_val * m_val) 0)

Now we get another scary-looking set of compiler errors:

src/Const/Matrix.hs:13:29-34: error:
    • No instance for (KnownNat n) arising from a use of ‘natVal’
      Possible fix:
        add (KnownNat n) to the context of
          the type signature for:
            zeros :: forall (n :: GHC.TypeNats.Nat) (m :: GHC.TypeNats.Nat).
                     L n m
    • In the first argument of ‘fromIntegral’, namely
        ‘(natVal @n Proxy)’
      In the expression: fromIntegral (natVal @n Proxy)
      In an equation for ‘n_val’: n_val = fromIntegral (natVal @n Proxy)
   |
13 |   let n_val = fromIntegral (natVal @n Proxy)
   |                             ^^^^^^
src/Const/Matrix.hs:14:29-34: error:
    • No instance for (KnownNat m) arising from a use of ‘natVal’
      Possible fix:
        add (KnownNat m) to the context of
          the type signature for:
            zeros :: forall (n :: GHC.TypeNats.Nat) (m :: GHC.TypeNats.Nat).
                     L n m
    • In the first argument of ‘fromIntegral’, namely
        ‘(natVal @m Proxy)’
      In the expression: fromIntegral (natVal @m Proxy)
      In an equation for ‘m_val’: m_val = fromIntegral (natVal @m Proxy)
   |
14 |       m_val = fromIntegral (natVal @m Proxy)
   |                             ^^^^^^
src/Const/Matrix.hs:15:7-12: error:
    • No instance for (KnownNat n) arising from a use of ‘matrix’
      Possible fix:
        add (KnownNat n) to the context of
          the type signature for:
            zeros :: forall (n :: GHC.TypeNats.Nat) (m :: GHC.TypeNats.Nat).
                     L n m
    • In the expression: matrix (replicate (n_val * m_val) 0)
      In the expression:
        let
          n_val = fromIntegral (natVal @n Proxy)
          m_val = fromIntegral (natVal @m Proxy)
        in matrix (replicate (n_val * m_val) 0)
      In an equation for ‘zeros’:
          zeros
            = let
                n_val = fromIntegral (natVal @n Proxy)
                m_val = fromIntegral (natVal @m Proxy)
              in matrix (replicate (n_val * m_val) 0)
   |
15 |    in matrix (replicate (n_val * m_val) 0)
   |       ^^^^^^

Here, the compiler is saying that it doesn’t know whether the type variables n and m are associated with a particular integer. This can be fixed by giving them a KnownNat constraint:

{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE ScopedTypeVariables #-}

import Numeric.LinearAlgebra.Static
import GHC.TypeLits (natVal, KnownNat)
import Data.Data (Proxy(..))

zeros :: forall n m. (KnownNat n, KnownNat m) => L n m
zeros =
  let n_val = fromIntegral (natVal @n Proxy)
      m_val = fromIntegral (natVal @m Proxy)
   in matrix (replicate (n_val * m_val) 0)

And voilá! We have implemented a type-safe function to create constant matrices using the hmatrix static API. Here’s a complete program that creates and prints a matrix full of zeros:

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

import Data.Data (Proxy (..))
import GHC.TypeLits (KnownNat, natVal)
import Numeric.LinearAlgebra.Static as LA

zeros :: forall n m. (KnownNat m, KnownNat n) => L n m
zeros =
  let n_val = fromIntegral (natVal @n Proxy)
      m_val = fromIntegral (natVal @m Proxy)
   in matrix (replicate (n_val * m_val) 0)

main :: IO ()
main = print (zeros @3 @2)

which returns:

(matrix
 [ 0.0, 0.0
 , 0.0, 0.0
 , 0.0, 0.0 ] :: L 3 2)

Implementation using `build`

A more straightforward (but less fun) way to implement the zeros function in Hmatrix is using the build function from the Static API:

build :: forall m n. (KnownNat n, KnownNat m)
    => (Double -> Double -> Double) -> L m n

build is a higher-order function which takes a function as input and returns a matrix of the correct size. The input function takes two arguments (the row and column index of the element) and returns a value. The build function then creates the matrix of the correct size (using similar forall and natVal tricks as before) and creates the values by applying the input function to the indices of each cell. Therefore, to implement zeros all we need to do is define a function that takes two arguments of type Double and always returns 0 and pass it as an argument to build:

constZero :: Double -> Double -> Double
constZero _ _ = 0

zeros :: (KnownNat n, KnownNat m) => L n m
zeros = build constZero

Here’s the full program:

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeApplications #-}

import GHC.TypeLits (KnownNat)
import Numeric.LinearAlgebra.Static

constZero :: Double -> Double -> Double
constZero _ _ = 0

zeros :: (KnownNat n, KnownNat m) => L n m
zeros = build constZero

main :: IO ()
main = print (zeros @3 @2)