Software

These are some personal open-source software projects I’ve worked on. Some are from my time at university so quite old, but recently I have also been working on building a bunch of numerical simulation and constrained optimization tooling for JAX (more physics, less statistics). You can find all software projects on my GitHub.

Numerical linear algebra

splineax

Sparse linear operators and direct solvers for Lineax. The JAX ecosystem has been lacking in support for sparse matrices for a while, and this is a big limitation for many physics-motivated systems which are naturally very sparse (e.g. a grid where there are many cells but each only touches its neighbors). I was motivated to create this library after finding asdex, which can deduce structural sparsity in JAX functions and calculate sparse Jacobians efficiently if the sparsity is known. This made it possible to create sparse linear operators backed by JAX BCOO and BCSR formats, and interface with sparse direct solvers.

I have added solver support for SuiteSparse KLU and oneMKL Pardiso on CPU and NVIDIA cuDSS and GPU, and an API extension to Lineax that allows for passing through solver state across multiple solve calls, in order to allow efficient reuse of the factorization objects these direct sparse solvers create. It also contains some usability features, such as a “solve profile” you can use to check whether you are able to reuse factorizations or not, and a “stateful solve transform” that allows you to take functions written against the plain Lineax API and transform to reuse factorizations across calls. This makes it possible to take e.g. an existing optimization algorithm which uses dense linear solves and “sparsify” it without changing its code.

jax-nansparse

JAX is a great package, but its origin in machine learning applications and focus on GPUs means it is not a very good fit for array computations with a lot of sparsity. This is a shame, since especially when there is structural sparsity (sparsity that is determined by the structure of the computation, not data-dependent) there is really no reason to have large dense arrays with lots of zeroes around. While there are some libraries that can make use of the sparsity when you know it (like sparsejac), there are not many tools for deducing structural sparsity from a JAX computation. This package was an attempt to do so, by making use of the “dominant” property of NaNs in floating point: if you put a NaN into most computations, you’ll get a NaN out. This means that in many cases, entering a NaN as a tangent value allows for deducing whether a derivative is possibly nonzero. The library is however in early stages and there are many edge cases it doesn’t handle, and I haven’t found the time to develop it further.

pardiso-mkl-jax

A JAX interface to the Pardiso direct sparse solver from Intel oneMKL, just glue code to enable splineax.

entangle-jax

A small library that helps with ordering operations in JAX/XLA programs. XLA is free to reorder operations without any data dependencies on each other, so entangle creates an artificial data dependency between two values. This is useful when dealing with libraries that have internal state and define operations that are side-effecting but not visible to XLA, such as some of the sparse solver libraries used in splineax.

Mathematical optimization

cyxla

A package for running XLA programs from Cython extension modules through the PJRT C API. This can be beneficial if you have a native library that has to execute your JAX/XLA functions repeatedly, as it removes the overhead of going back through the Python interpreter, and doesn’t require taking the GIL.

xipopt

Python bindings to the Ipopt nonlinear optimization solver, built around JAX. The motivation for this was to have a simple and performant way to optimize systems defined via JAX with Ipopt. cyipopt is an existing package, but it requires some work to integrate efficiently, and requires building from source to install which can be cumbersome. xipopt rewrites objective and constraint functions into a fused primal/tangent form, compiles/executes them with cyxla to avoid Python overhead during optimization, and caches repeated evaluations of the same point by Ipopt. It also builds fully bundled wheels for Linux/Windows/macOS so it is easy to install.

cyhighs

Cython bindings to the HiGHS MILP/QP optimization solver. The motivation for this library was to be able to use the simple scipy.optimize.linprog interface which only requires the problem in matrix form, but allowing to decouple the version of HiGHS used from the version of SciPy (which bundles HiGHS). Wraps the HiGHS C API with Cython. Mostly glue code.

jaxormick

Computes McCormick relaxations of JAX functions, which allow a function to be bounded over a certain part of the domain by e.g. an optimization solver. This uses function rewriting similar to how automatic differentiation works (transforming a function into a function that calculates its derivative). Primarily created to enable jax-cip. A similar project is AutoBound from Google, but this does not support all JAX primitives and is not actively maintained.

jax-cip

An interface between JAX and the SCIP solver for global optimization. Objective and constraint functions can be general nonlinear functions defined via JAX, and then derivatives can be automatically deduced, as well as McCormick relaxations via jaxormick. The relaxations (convex/concave under/over-estimators) allow SCIP to bound the function over certain regions and perform global optimization, which also means we can constrain variables to binary/integer values. This is a bit of an experimental project, not yet applied to any real use case.

aulax

Constrained nonlinear optimization in JAX using the augmented Lagrangian method. It allows adding general nonlinear equality and inequality constraints on top of the unconstrained minimizers in Optimistix. This is a bit of an experiment, to try a constrained optimization method that is fully JAX-based (so no compatibility issues, runs on any accelerator) and can in principle be first-order/gradient-only (depending on choice of minimizer). Unfortunately it seems to be in many cases outperformed by interior-point methods such as Ipopt, but I have not done much testing with large problems on GPUs yet, so it could have a leg up there.

Machine learning

surreax

A library for fitting surrogate models to JAX functions. The motivation for this library was to have an easy-to-use surrogate model training toolkit that integrates well with models written in JAX, and that has a focus on physics simulation models and amortized optimization. Concretely this means there is dedicated support for fitting surrogates of implicitly defined models (e.g. training an autoencoder over the state space by looking at the residual function directly) and model forms that lend themselves to (global) optimization (e.g. a ReLU network that can be expressed as an MILP for which efficient optimization packages exist).

Kuramoto.jl

An implementation of the Kuramoto model for use with the DynamicalSystems library in Julia. I used it to generate test data for DimensionGrouping.jl.

DimensionGrouping.jl

A dimensionality reduction algorithm I wrote for a project in university. The goal was to reduce grouped trajectory data, for example: given the position of a number of animals moving in packs, find the trajectories of the packs.

This lead to some interesting requirements, such as the need for the trajectory to remain feasible (a simple mean of a point cloud might generate impossible trajectories). Also, the determination of the groups to which animals belong is a kind of clustering problem, but spread out over time: animals will stick with the same group for a good amount of time. My idea was to write a clustering algorithm as an unsupervised learning algorithm, and then smooth the cost (distance) of the animal joining each group over time, to create a configurable tradeoff between having “optimal” groups and any moment and stable groups over time. And it worked!

The code was written in Julia and the project report with theoretical details is available here. To evaluate it I also wrote a small simulation of the Kuramoto model in Julia.

Other libraries and tools

meson-jll

A Rust tool that generates Meson wrap files from Julia JLL packages. The Julia ecosystem has precompiled binaries for many scientific C, C++ and Fortran libraries, and I thought it would be nice to be able to use these binaries in non-Julia projects. meson-jll parses the Julia manifests and writes corresponding Meson “wrap” files that allow a user to simply declare the library via Meson dependency and use it in any Meson project. Some of my libraries use this method to load and bundle native binaries, e.g. xipopt, jax-cip and cyhighs.

polder

This is a project I’m quite excited about. In short, its a “labeled array” package similar to Xarray, but compatible with both modern DataFrame libraries (Polars, DuckDB) for the labels and any array API library (like NumPy, JAX, CuPy, Dask, Sparse) for the array. It’s quite simple, but I think an abstraction like this (if done right) can be very useful in many applications.

rust-jaxpr-interpreter

This project in large part exists for me to learn Rust in a familiar context. The idea comes from the fact that JAX is a big package: it provides useful computational transformations such as autodiff as well as an interface to the XLA compiler for efficient execution. But, if you only want to use the Jaxpr expression framework and transformations built on it (like autodiff), you then need to accept the overhead, limitiations and intransparency of the XLA compiler as well, if you want anything that resembles efficient execution. For that reason I thought it might be useful to have an alternative way of excecuting JAX programs, that is a little bit more efficient pure-Python execution, but more flexible and with less overhead than XLA computation. There could also be opportunities for novel interpreters with different backing data structures, such as native sparse arrays.

python-delayed-import

A small package I put together recently that allows you to “delay” Python imports without too much hassle. Theoretically this could help with improving startup times on large codebases, but in practice I haven’t found any real-world use cases yet where there is large performance gain. There are a couple similar packages to this around, all with their own caveats and tradeoffs, but I wanted to try and roll my own. Also a way to try out some modern Python packaging tools like uv. Update: a similar feature will be in Python 3.15, so now this package is even more useless!

Apps

btt

BigTrainTrip, an app for planning multi-day train trips. This is a more recreational project. I often make 1-3 week “road trips by train”, going to various places around Europe, and always spend quite some time piecing together the route, especially when long-distance (night) trains are involved with only a few connections and large price differences. Most train planning software I could find is very single-journey focused, and most multi-day trip planning software is not very well integrated with public transport data. For tickets I would often just buy an Interrail pass to make it easier, but end up not really saving that much money or even paying more, since ticket prices and reservation requirements differ so much between countries.

So, I thought I would make something to help me with this. btt is an offline planning app that contains train schedule data for pretty much all European operators via Transitous, uses a RAPTOR-based algorithm for route planning, can look up live ticket prices for a number of operators and contains a online-learning price model that is used during planning. There’s also a web viewer for checking your planned route on the go. I haven’t yet used it for a real trip myself but look forward to do so!

Early projects

lark-translator

A translation add-on for the Lark parser, inspired by the Grammatical Framework. A grammar describes the structure of a language, so you can use it both to parse text and to generate it. If you parse text with a grammar for one language and then reconstruct text from the same parse tree with a grammar for another, you get a kind of translation. This idea came from my thesis work on translating sign language via a grammatical approach, but never got used for anything.

pobx

A small library for reactive programming in Python. Basically a port of MobX, though it’s probably a few year out of date by now. I believe the concept is still interesting though and don’t really know of any real alternatives, especially in Python. One recent tool I can think of that gives a very nice application of reactivity in the Python ecosystem is marimo, which provides reactive notebooks.

ajay

A small actor-like library but using agent terminology. Not that different from many actor/message passing frameworks already out there, but during my masters we had one or two course on agent programming and I thought it was a shame there was no modern library out there for a language like Python that allowed for interacting with these concepts (e.g. standards like FIPA-ACL) in a sort of ergonomic way. It makes of use of a lot of interesting constructs: transparent sinlge/multi-process networking using ZeroMQ, message parsing/generating, and tries to explicitly enforce the concept that an agents behavior can be reproduced just by replaying it’s messages/“observations” (which is basically something that holds in any functional language, but I digress). However, in the end I had no real reason to develop this further.

preparejs

Prepare.js runs JavaScript on the server in a browser-like environment before a page is sent to the client. You mark script tags in a .js.html file with type="server". Prepare executes them with a DOM from jsdom, removes them, and returns the resulting HTML. It was a basic framework for what is now called isomorphic JavaScript. A nice idea, but not really relevant anymore.

MobileNose

Mobile apps for DataNose, a timetable and course system for the University of Amsterdam. The code is a cross-platform .NET library with separate front-ends per platform. It was used by a good number of students for a few years, until a usable web app was produced by the university itself.

rbkey-midi

A small piece of software from way back that reads USB signals coming from a Rock Band 3 keyboard controller (keytar) and translates them to a virtual MIDI device. These keyboards were available for pretty cheap at one point (like 20 euros) and they are portable and wireless, so I thought it would be nice if you could use them as a generic MIDI keyboard. It was a good learning experience, since the way they got it to work on game consoles is by packing MIDI data into floating point numbers sent over the controller API (stuff like, the 3rd bit set in the left stick X position = key A2 pressed). So noticing that and then unpacking the data was very insightful. And I know it’s been used by at least one band in a live performance!

OscilloscopeEmulator

This is something I wrote way back in secondary school as part of my final project: a simple oscilloscope simulator. It just takes the left and right audio channels and continuously draws a line where the left channel gives the X position and the right the Y position. You can make Lissajous figures and stuff like that. I was trying to make something like Oscillofun and wanted a way to test it without sitting by the oscilloscope all day. Probably by now there is much better software that does the same thing!