Rajeev Jain

Rajeev Jain

Principal Specialist, Research Software Engineering · AI Systems and Verification · HPC · Scientific Computing

The gap between prototype and production is where I work — ML training pipelines that scale on new accelerator hardware, I/O that doesn’t bottleneck at exascale, Python platforms research teams can actually maintain across institutions and years. Lately most of that is AI systems: giving agents typed, scoped, provenance-tracked tools instead of a chat box, and checking what comes back against a measurement rather than a plausible story. Sixteen years of numerical code makes me a hard sell on a benchmark. Principal Specialist at Argonne National Laboratory, with a joint appointment at the University of Chicago.

Rajeev Jain

Work

  • UXarrayLead developer · open-source climate analysis
    Python library for unstructured climate grid analysis — the standard tool for DOE labs, NCAR, and universities working with MPAS, ICON, SAM, and next-generation meshes. Conservative zonal averaging via Gauss-Legendre quadrature; grid I/O for ESMF, MPAS, SCRIP, and HEALPix; MCP server for AI-agent dataset exploration across local and HPC execution. Also the correctness and performance work underneath it — compensated summation in the spherical-geometry kernels after naive accumulation drifted 0.68 m on a real Earth mesh, and a 1.67× speedup in a Numba hot path once hidden heap allocation turned out to be the actual cost rather than the math. Docs · GitHub · MCP article
  • Local LLM inference, measuredBenchmark harness · open data and protocol
    Four models on one machine, taken apart to find what actually sets throughput. Bytes read per token parsed from GGUF tensor tables rather than taken from file size, which overstates by 64% on a model with gathered embeddings; the bandwidth denominator measured at 340 GB/s rather than read off the 400 GB/s datasheet figure, which is a pin rate and unreachable. All four models land between 36% and 41% of memory bandwidth regardless of architecture — the 8× spread in tokens per second is explained entirely by bytes per token. Scripts, raw data, and protocol published so every number regenerates. Book · Repo
  • Pangu-Weather on AuroraIntel XPU port · Argonne Leadership Computing Facility
    PyTorch climate emulator on the Spherical Fourier Neural Operator, ported to Aurora’s Intel XPU stack for DOE exascale Earth system modeling. First stable portable DDP baseline: PMIX/PALS environment mapping, XPU/CUDA device branching, device-aware mixed precision with gradient scaling on CUDA and bf16 on Intel XPU. Measured baseline is one node, 12 XPU ranks, ~12 s steady-state epochs; full-dataset scale-up is still ahead. Article
  • CANDLE / IMPROVECore contributor · R&D 100 Award 2023
    HPO and benchmarking infrastructure for cancer drug response models — 15+ researchers across Argonne, LLNL, and ORNL. 10,000+ training experiments across Summit, Theta, and Cori using Swift/T. Published in Briefings in Bioinformatics, 2025.
  • FLASH-XI/O and compression lead · R&D 100 Award 2022
    Checkpoint and restart redesign for a million-line multiphysics engine. Async HDF5 with Argobots plus SZ3/ZFP compression: 40–70% checkpoint overhead reduction and 50%+ storage savings on Summit. Cross-checkpoint restart between AMReX and Paramesh — removing a hard constraint that forced full restarts when switching solvers. SC24 paper
  • MeshKitPI and software lead · DOE NEAMS · 2009–2016
    Open-source C++ toolkit for automated nuclear reactor core mesh generation. Parallel CoreGen: 712 processors, 101 million hexahedral elements, 14 GB MONJU reactor mesh in under 7 minutes — a job the serial path couldn’t run at all. Blog post · Source

Writing

Selected papers

Full list on Google Scholar · 22+ publications

Recent talks

Recognition

  • R&D 100, 2023 CANDLE — cancer AI infrastructure across Argonne, LLNL, and ORNL
  • R&D 100, 2022 FLASH-X — multiphysics simulation engine
  • IMR 2010 Best Paper — reactor core mesh generation with lattice hierarchy encoding
  • ATPESC 2015 Scholar — Argonne training program on extreme-scale computing

Service

  • 2026 Program committee — AGENT4SC, 1st Workshop on Agentic AI for Large-scale Science, IEEE eScience 2026
  • Ongoing SBIR/STTR proposal reviewer, U.S. Department of Energy · reviewer, Journal of Open Research Software

Funding

  • Active DOE SEATS — Software Ecosystem for Advancing Climate Tools and Services
  • Active NSF Raijin — collaborative research in climate model analysis
  • 2017–2023 DOE ECP CANDLE — core contributor
  • 2009–2016 DOE NEAMS — principal investigator, MeshKit

Roles

  • 2009–present Argonne National Laboratory — Principal Specialist in Research Software Engineering
  • 2023–present University of Chicago — Staff At-Large, cancer pharmacogenomics and Earth system science
  • 2007–2009 Arizona State University — Research and teaching assistant, structural and computational mechanics

Education

  • 2020 M.S. Computer Science — University of Chicago
  • 2009 M.S. Structural Engineering — Arizona State University
  • 2006 B.Tech. Mechanical Engineering — IIT ISM Dhanbad