Curriculum Vitae

Wassim Kabalan — Research Associate, The Alan Turing Institute

Wassim Kabalan

Research Associate — AI for Physics, The Alan Turing Institute

I am a Research Associate at The Alan Turing Institute, where I work on AI for physics. During my PhD in computational cosmology I built differentiable simulation and inference software in JAX, from the statistical method down to the distributed GPU code, and several of these libraries are now used by other groups. Before the PhD I spent five years writing production C++ at Dassault Systèmes inside a multi-million-line codebase.

GitHub - ASKabalan LinkedIn - Wassim Kabalan Email - wkabalan@turing.ac.uk

Education & Professional Experience

Oct 2026 –
present 		Research Associate, The Alan Turing Institute, London, UK
AI for Physics team (Fundamental Research), with Jason McEwen. Physics-informed emulators and rigorous Bayesian uncertainty quantification.
Dec 2023 –
Oct 2026 		PhD in Computational Cosmology, Université Paris Cité — APC (CNRS/IN2P3), Paris, France
Thesis: Differentiable Field-Level Cosmology: CMB, Weak Lensing, and GPU-Accelerated Inference.
Advisors: Eric Aubourg, Alexandre Boucaud, Josquin Errard, François Lanusse.
  • Core of the thesis: developing JAX-based, GPU-accelerated, fully differentiable, high-dimensional statistical models to extract physical parameters from cosmological data.
  • Weak-lensing field-level inference (LSST). Building differentiable cosmological simulations as forward models, then running MCMC chains over a ~10⁹-dimensional Bayesian posterior on up to 512 H100 GPUs to recover cosmological parameters directly from large-scale-structure data.
  • CMB component separation (Simons Observatory). Built a robust GPU-accelerated optimization framework, and led the design of a spatial clustering strategy for upcoming CMB satellite missions; first author of the corresponding paper (arXiv:2604.08463).
Apr 2023 –
May 2023 		Advanced AI for Data Analysis, École Polytechnique (Executive Education), Palaiseau
Oct 2019 –
Oct 2023 		3DEXPERIENCE & CATIA Software Infrastructure Engineer, Dassault Systèmes, Vélizy-Villacoublay, France
  • Developed features and maintained core infrastructure inside a multi-million-line C++ codebase shared by hundreds of engineers across teams.
  • Built and extended core conversion services and the CATIA cache infrastructure in modern C++ (17/20) on Linux, including a multi-session distributed cache.
  • Built CI/CD pipelines and a test matrix for the conversion service infrastructure.
Jan 2019 –
Oct 2019 		Data Acquisition Engineer, SERMA (for Renault), Guyancourt, France
  • Built tooling for high-volume vehicle sensor logs, including a custom C decoder for ECU binary data.
  • Developed parallel Python post-processing workflows on an 88-core server.
  • Automated end-to-end conversion pipeline for vehicle network data.
Sep 2016 –
Sep 2018 		M2, Electronics, Electrical & Automation, Université Gustave Eiffel
Sep 2013 –
Nov 2016 		Licence, Engineering Science, Université Paris-Est Créteil (UPEC)

Collaborations

  • LSST Dark Energy Science Collaboration (DESC) — Member; weak-lensing forward models and scalable, differentiable simulation tools (JaxPM, jax-fli).
  • Simons Observatory — Member; component separation tooling (FURAX) for CMB analysis.

Open-Source Software

  • jax-fli — Author. JAX-based field-level inference package; fully distributed, high-resolution N-body and lensing simulations.
  • jaxDecomp — Author / Maintainer. JAX library for distributed 3D domain decomposition and FFTs: NVIDIA cuDecomp bindings (multi-GPU, NCCL) plus a pure-JAX backend for portability across CUDA / TPU / CPU.
  • JaxPM — Main contributor / Maintainer. Differentiable particle-mesh simulations in JAX with multi-accelerator support; used for weak-lensing forward modeling.
  • jax-healpy — Main contributor / Maintainer. JAX-native HEALPix utilities; AD-friendly, GPU-ready.
  • FURAX — Main contributor / Maintainer. JAX framework for inverse problems; used for Simons Observatory CMB component separation.

Talks & Tutorials

  • Generative AI with JAX [Tutorial] • AISSAI School 2025 • Oct 2025 slides
  • JAXPM: Scalable and Differentiable Particle-Mesh Simulations • Bayesian Deep Learning Workshop • May 2025 slides
  • Bayesian Inference for Cosmology with JAX [Tutorial] • Bayesian Deep Learning Workshop • May 2025 slides
  • Massively Parallel Computing in Cosmology with JAX [Tutorial] • CoPhy 2024 • Nov 2024 slides
  • Differentiable and Distributed Particle-Mesh N-body Simulations • LSST France 2024 • Jun 2024 slides

Publications

Refereed

Preprints

In Preparation

  • Kabalan, W., Lanusse, F., Boucaud, A., and Aubourg, E. 2026. Differentiable Forward Modelling for Field-Level Inference in Stage IV Weak Lensing Surveys.

Supervision

Jan 2025 –
Jun 2025 		Binh Nguyen, Master 1 Student
Project: Inference of Weak-Lensing Parameters from Blended Galaxies Using Generative Neural Networks. Final report.
Outcome: Talk at Rencontres du Vietnam.

Skills

Programming Python (7y), JAX (3y), C++ (5y), CUDA (3y), PyTorch (1y)

HPC & GPU Computing NCCL, MPI, Slurm, Nsight; multi-node GPU (up to 512 H100s), distributed FFTs, custom CUDA primitives via XLA FFI/CustomCall

Statistics & Machine Learning High-dimensional Bayesian inference (MCMC, HMC, NUTS) with NumPyro / BlackJAX; simulation-based and field-level inference; generative deep learning; uncertainty quantification

DevOps & Software Engineering GitHub/GitLab CI, packaging (PyPI), containers, Linux, CMake, TDD

Languages

French Native
English Professional proficiency (C1)
Arabic Native
German Basic (A2)