Curriculum Vitae

Wassim Kabalan — R&D Engineer & PhD Candidate (APC, CNRS/IN2P3)

Wassim Kabalan

PhD Candidate & Software Engineer — Computational Cosmology

I am a PhD candidate and software engineer specializing in computational physics, high-dimensional Bayesian inference, and differentiable programming. I develop JAX-based,differentiable and distributed statistical models to extract physical parameters from complex cosmological data.

GitHub - ASKabalan LinkedIn - Wassim Kabalan Email - wassim@apc.in2p3.fr

Education & Professional Experience

Dec 2023 –
Dec 2026
(expected) 		Graduate Researcher (PhD, Physics of the Universe), Université Paris Cité — APC (CNRS/IN2P3), Paris, France
Thesis: Automatically differentiable and distributed probabilistic programming for weak-lensing 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

Software

  • Kabalan, W., Lanusse, F., Boucaud, A., and Aubourg, E. 2025. jaxDecomp: JAX Library for 3D Domain Decomposition and Parallel FFTs. Submitted to JOSS.

In Preparation

  • Kabalan, W., Lanusse, F., Boucaud, A., and Aubourg, E. 2026. JAXPM: A JAX-Based Framework for Scalable and Differentiable Particle Mesh Simulations.

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)