Posts

A Primer on Just-In-Time (JIT) Compilation

Internal presentation
JIT
JAX
Numba
Julia
Agentic engineering
LLM

Numba, JAX and Julia are the three just-in-time compilers most likely to be found in production scientific computing. Run the same small problems through all three and what separates them turns out not to be speed: each language design comes bundled with a different set of abstraction layers — its own tower of intermediate representations, its own introspection tools, its own notion of what you are allowed to say — and the optimisation strategy is baked into that bundle rather than chosen independently of it. The useful question about any of them is never “is it fast?” but which language am I actually writing? — and the answer is never the one in the file extension. Every speed-up below is a rewrite that computes exactly the same thing: a compiler pass performed by hand. What the three differ in is how much of that narrowing they leave to you.

Each system is walked down its own tower, with a runnable notebook behind every figure: why fusing four array allocations into two is where the speed actually comes from; why naming an intermediate costs Numba a 127 MB round trip that JAX’s tracer never pays, and why hand-writing the loop then beats array notation; why Julia’s fusion is syntax you write rather than an inference the compiler makes, why attaching physical units can cost exactly nothing, how a routine ustrip on a symmetric sparse matrix silently densified a matrix that would have been 450 TB, and what it means that a language is its own metaprogramming language. Then what the layering itself buys — separation of concerns, a place to verify, a place to optimise, a place to express intent — what compiling late adds on top of it, and what it costs: no artefact to archive, no code to sign, verifiability traded for performance.

Then a digression — an LLM read as a JIT compiler: the ultimate late binding and the ultimate metaprogramming language, but performing enormous narrowing with no specification licensing it and no guard checking it. Bun’s 535,000-line Zig-to-Rust rewrite worked because it had both.

The slides, this write-up and the PDF are generated from a single source, so everything that was on screen is here, together with what was said around it.

August 5th, 2026

Journal Club

Internal presentation
LLM
Agentic engineering

Journal club discussion on agentic coding: the shift from autocomplete to autonomous agents (Claude Code, Codex CLI), vibe coding vs. agentic engineering, and the implications for software quality.

May 6th, 2026

Reproducibility in Computational Research

Internal presentation
Reproducibility

Reproducibility is a cornerstone of computational research, yet achieving it across diverse computing environments remains challenging. This examines the landscape of software environment management tools and their role in enabling reproducible research workflows. We begin by clarifying the terminology around reproducibility, distinguishing between re-runnability, repeatability, reproducibility, reusability, and replicability. We then explore the multifaceted nature of package managers—as software tools, ecosystems, indices, and distributions—and categorize them by scope, distribution method, platform support, and linking strategy.

The core focus is on practical solutions to three critical problems: reproducing research software environments across different systems and platforms, customizing build processes while maintaining reproducibility (particularly for high-performance computing contexts), and distributing these environments effectively. We provide in-depth analysis of several key technologies: Conda/Mamba for cross-platform binary package management, (TODO: Pixi for modern Python project workflows with lock files and task automation, and brief overviews of Nix, Spack, and Docker). Through concrete examples ranging from pure Python packages to complex scientific software stacks and HPC system environments, we demonstrate how these tools address real-world reproducibility challenges. The presentation emphasizes the importance of understanding the ecosystem around package managers, including channels like conda-forge, and discusses advanced topics such as build customization, platform-specific optimizations, (TODO: and the trade-offs between source and binary distributions).

November 26th, 2025

JIT compilers for scientific computing in Python: Numba vs. JAX

Public talk
JAX
JWST
PyCon UK

Accelerating scientific Python with JITs. We share our journey migrating a gravitational lensing likelihood calculation from Numba to JAX. Learn about performance gains, automatic differentiation benefits, and practical lessons for high-performance scientific computing in Python.

Python is widely used in scientific research, but pure Python can sometimes be too slow for computationally intensive tasks. Just-In-Time (JIT) compilers are essential tools for boosting performance, allowing Python code to run closer to native speeds. While libraries like Numba have long been popular for accelerating numerical Python functions, JAX offers a new paradigm, combining JIT compilation with powerful features like automatic differentiation (auto-diff) and execution across different hardware (CPUs, GPUs, TPUs).

This talk will take you on a journey through our experience optimizing a critical component of an astrophysics analysis pipeline: the calculation of the likelihood function for gravitational lensing models, used with data from the James Webb Space Telescope. We initially used Numba to accelerate this calculation, but the need for performance portability across hardwares, and the potential speed up from gradient computation for model fitting led us to explore JAX’s unique capabilities.

This session will walk through the practical steps, challenges, and insights gained from migrating this complex scientific code from its existing Numba implementation to a JAX-based one.

You will learn:

  • Why leveraging performance tools like JITs is crucial for cutting-edge scientific analysis in Python.
  • The practical considerations when migrating existing numerical code from Numba to JAX, including syntax changes and managing state.
  • How JAX’s auto-differentiation simplifies gradient calculations essential for scientific optimization and sampling tasks.
  • The significant performance improvements achieved in our specific gravitational lensing case study by using JAX’s compiled functions.
  • Broader lessons learned about structuring scientific Python projects to effectively use modern JIT compilers and harness capabilities like auto-diff.

We’ll conclude by comparing Numba and JAX in benchmark performance, developer ergonomics, and tradeoffs between the two, providing you with practical guidance for choosing the right tool for your scientific computing needs.

This case study offers a concrete example of how evolving Python libraries are enabling researchers to perform complex, high-performance computations directly within the Python ecosystem. Join us to see how tools like JAX are empowering scientific discoveries, one optimised function at a time.

This talk is suitable for intermediate Python programmers familiar with NumPy.

September 21st, 2025

Durham HPC Days 2025

Internal presentation
Durham HPC Days

Conference recap of Durham HPC Days 2025, covering themes of GPU acceleration, AI/HPC convergence, benchmarking, and the RSE community.

June 7th, 2025

Scientific Computing with JAX

Public talk
JAX
Durham HPC Days
JWST

Presented at Durham HPC Days 2025.

JAX is a Python library that combines just-in-time (JIT) compilation with automatic differentiation, powered by XLA (Accelerated Linear Algebra compiler), to target multiple hardware architectures including CPU, GPU (NVIDIA, AMD, Intel, Apple), and Google TPU. While primarily designed for machine learning research, JAX presents compelling advantages for scientific computing in the HPC landscape.

As HPC increasingly depends on heterogeneous accelerators, JAX’s “write once, deploy anywhere” approach enables scientific libraries to efficiently utilize diverse computing resources. Python programmers familiar with array programming and other JIT solutions like Numba can transition to JAX with minimal code changes, allowing straightforward adoption. Additionally, automatic differentiation provides gradient calculations for optimization problems with minimal developer effort — automatically generating derivatives that would be prohibitively complex to implement manually, while enabling optimization algorithms to converge more rapidly with gradient information.

This presentation examines the adaptation of a likelihood function calculation from Numba to JAX in the context of gravitational lensing analysis of James Webb Space Telescope observations. The case study specifically addresses the convergence of AI/ML techniques with traditional scientific computing, highlighting performance gains, implementation challenges, and practical considerations when migrating existing scientific code to JAX’s programming model. We’ll discuss how JAX enables scalable processing across different HPC hardware resources while maintaining scientific accuracy and computational efficiency.

June 4th, 2025

Quantum Mechanics references

References
QM

Annotated bibliography of quantum mechanics textbooks and lecture notes, with recommendations for different starting points and learning goals.

April 11th, 2025
No matching items