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.

Work
- UXarrayLead developer · open-source climate analysisPython 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 protocolFour 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 FacilityPyTorch 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 2023HPO 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 2022Checkpoint 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–2016Open-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
- 2026 I decided to write a basics book while setting up Meta Muse Glimmer 30B locally on my M1 Max Book-length, from what a model file is to a measured explanation of why a 30B model gives about six words per second on a 64 GB laptop — and why no setting you change will fix that.
- 2026 The Bug Hiding in Every Sum: Kahan Summation, Compensated Arithmetic, and Where They Actually Matter A $0.0085 drift in a $50 million ledger and a 0.68-meter error on a real Earth mesh, both from the same one-line bug: adding floats in the naive order.
- 2026 Loops: The Oldest Primitive, and How One Nearly Wrecked a C++-to-Python Port Agentic AI runs the world in a while-true, but the humble for-loop still hides the most common way a faithful C++-to-Python port turns slow.
- 2026 Where Your Numbers Live: Killing Hidden Heap Allocations in a Numba Hot Path A 1.67× speedup that came not from better math but from where the numbers lived — heap versus registers, and why NumPy's ndarray is always the former.
- 2026 MCP in 2026: From Anthropic Side Project to Industry Infrastructure MCP's move to the Linux Foundation, its SEP-driven governance, and the five spec changes reshaping what agents can do — timeline, governance, and active proposals.
- 2026 The AI ROI Problem Is Real. The Fix Isn't Cutting Licenses. Why most enterprise AI fails to deliver returns, and the architecture — shared by the UXarray MCP work — that the projects creating real value actually use.
- 2026 UXarray MCP Server: AI-Agent Dataset Exploration with Globus Compute How the UXarray MCP server lets AI agents explore, analyze, and visualize unstructured climate grids — locally and on HPC via Globus Compute.
- 2025 IMPROVE: Building Rigorous Benchmark Infrastructure for Cancer Drug Response Prediction The improvelib package, cross-study analysis framework, GitHub Actions CI/CD, and the UNO dual-branch neural network for drug response prediction.
- 2024 Pangu-Weather on Aurora: Porting a Weather Foundation Model to Intel GPUs Device abstraction, DDP setup, PMIX/PALS environment mapping, and mixed-precision on Intel XPU to get a stable training baseline on Aurora — plus the first measured one-node run.
- 2023 From RGG and MeshKit to the MOOSE Reactor Module Parallel CoreGen generated a 101M-element MONJU reactor mesh on 712 processors in under 7 minutes — a job the serial path couldn't run at all.
- 2022 CANDLE/Supervisor: Running Cancer AI Research at Scale on DOE Supercomputers mlrMBO, DEAP, Hyperopt, and Swift/T ran 10,000+ experiments on Summit; noise injection and counterfactuals revealed which genes drive tumor classification.
- 2020 Urban Microclimate at Scale: Array of Things, EnergyPlus, and CFD for Chicago Coupling Chicago's IoT sensor network, WRF mesoscale weather, EnergyPlus building simulation, and Nek5000 wall-resolved LES into a city-scale workflow.
Selected papers
- Jain, R., Jacob, R. Beyond Tool Execution: Evaluating Scientific MCP Interfaces with UXarray. AGENT4SC — 1st Workshop on Agentic AI for Large-scale Science, IEEE eScience 2026, Naples. Accepted; to be presented.
- Partin, A., ..., Jain, R., et al. Benchmarking community drug response prediction models. Briefings in Bioinformatics, 2025.
- Jain, R., Tang, H., Dhruv, A., Byna, S. Enabling Data Reduction for FLASH-X Simulations. DRBSD-10 Workshop, SC24, 2024.
- Jain, R., Wozniak, J.M., Partin, A., et al. Cross-HPO: Optimizing Neural Networks for Cancer Drug Response. CAFCW24, SC24, 2024.
- Wozniak, J.M., ..., Jain, R., et al. CANDLE/Supervisor: A workflow framework for machine learning applied to cancer research. BMC Bioinformatics, 2018.
- Tautges, T.J., Jain, R. Creating Geometry and Mesh Models for Nuclear Reactor Core Geometries. Engineering with Computers, 2011.
Full list on Google Scholar · 22+ publications
Recent talks
- eScience 2026 AGENT4SC: Beyond Tool Execution — Evaluating Scientific MCP Interfaces with UXarray — Naples, September 2026
- SciFM26 Presented UXarray MCP Server for agentic analysis of unstructured Earth-system meshes
- SC24 Tutorial: UXarray for Analysis of Unstructured Climate Data
- SC24 DRBSD-10: Enabling Data Reduction for FLASH-X
- SC24 CAFCW24: Cross-HPO for Cancer Drug Response
- AMS 2024 UXarray: Extending Xarray with Support for Unstructured Grids
- SciPy 2023 UXarray for Unstructured Climate Data
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