Rajeev Jain
Scientific software for climate, AI, and exascale computing
Rajeev Jain
Research software engineer at Argonne National Laboratory, with a joint appointment at the University of Chicago.
I work on research software for climate analysis, AI workflows, and large simulation codes. Over 16 years I have built software across national labs, universities, and multi-institutional collaborations, with an emphasis on reliable tools that other researchers can use and extend.
GitHub LinkedIn Google Scholar ORCID
Selected work
Selected projects
Representative work across climate science, cancer AI, and simulation software.
UXarray
Python library for unstructured climate grid analysis used by DOE labs, NCAR, and universities working with MPAS, ICON, SAM, and other next-generation meshes.
- Built conservative analysis operators, including zonal averaging via Gauss-Legendre quadrature.
- Shipped support for ESMF, MPAS, SCRIP, and HEALPix grid formats, with repeatable releases and CI.
- Currently extending the project with an MCP server and AI-agent workflow for natural-language dataset exploration.
Pangu-Weather on Aurora
PyTorch-based reimplementation of Pangu-Weather using the Spectral Fourier Neural Operator for deployment on more than 60,000 Intel GPUs.
- Ported the workflow to Intel GPUs and ran it at large scale on Aurora.
- Contributed to DOE exascale work in Earth system modeling and forecasting.
CANDLE / IMPROVE
Hyperparameter optimization and benchmarking infrastructure for cancer drug response models at supercomputer scale.
- Ran more than 10,000 training experiments across Summit, Theta, and Cori.
- Built GitHub Actions workflows for cross-study validation in a 15+ researcher collaboration.
- Published in Briefings in Bioinformatics in 2025.
FLASH-X
Optimization of checkpoint and restart workflows for a million-line multiphysics simulation engine used in astrophysics, combustion, and fluid dynamics.
- Implemented asynchronous HDF5 I/O with Argobots plus SZ3 and ZFP compression.
- Reduced checkpoint overhead by 40-70% on Summit and delivered 50%+ storage savings.
- Enabled cross-checkpoint restart between AMReX and Paramesh solvers.
MeshKit
Open-source C++ toolkit for automated nuclear reactor core mesh generation and lattice hierarchy modeling.
- Led the design of parallel meshing and multi-format I/O for reactor simulation teams at Argonne.
- Won Best Paper Award at the International Meshing Roundtable in 2010.
Technical expertise
Breadth across research software and systems
Tools and systems I work with most often across research software, data, and HPC.
Python to Fortran
Python, C++, Fortran, R, Bash, and SQL for analysis pipelines, simulation code, build systems, and automation.
Framework and workflow depth
PyTorch, TensorFlow, NumPy, Pandas, Xarray, Scikit-learn, Parsl, and Swift/T for model development and large experiment campaigns.
Performance and portability
MPI, OpenMP, HDF5, NetCDF, MOAB, Docker, Singularity, GitHub Actions, and storage-aware I/O design for leadership-class machines.
Science-driven software
Climate modeling, cancer pharmacogenomics, computational physics, mesh generation, AI infrastructure, and reproducible workflows.
Software that lasts
Release engineering, CI pipelines, open-source governance, multi-institution coordination, mentoring, and roadmap ownership.
Publications and talks
Selected papers, workshops, and presentations
More than 22 publications across HPC, machine learning, and computational science.
Benchmarking community drug response prediction models
Partin, A., ..., Jain, R., et al.
Enabling Data Reduction for FLASH-X Simulations
Jain, R., Tang, H., Dhruv, A., Byna, S.
Cross-HPO: Optimizing Neural Networks for Cancer Drug Response
Jain, R., Wozniak, J.M., Partin, A., et al.
CANDLE/Supervisor: A workflow framework for machine learning applied to cancer research
Wozniak, J.M., ..., Jain, R., et al.
Creating Geometry and Mesh Models for Nuclear Reactor Core Geometries
Tautges, T.J., Jain, R.
Recognition
Recognition and service
Awards, funding, and community work related to the software and collaborations above.
U.S. Permanent Resident
Authorized to work in the United States without sponsorship.
CANDLE
Project recognized by R&D World in 2023.
FLASH-X
Project recognized by R&D World in 2022.
ATPESC Scholar
Selected in 2015 for Argonne's training program on extreme-scale computing.
Best Paper Award
International Meshing Roundtable, 2010, for automated reactor core mesh generation research.
- DOE SEATS Active Software Ecosystem for Advancing Climate Tools and Services.
- NSF Raijin Active Collaborative research in climate model analysis.
- DOE ECP CANDLE Core contributor from 2017 to 2023.
- DOE NEAMS Principal investigator for MeshKit from 2009 to 2016.
- SBIR/STTR Proposal Reviewer U.S. Department of Energy.
- Panelist 5th Infraday Midwest Event on public infrastructure and AI.
- Reviewer Journal of Open Research Software and NumGrid.
- Committee Member NumGrid 2020 Program Committee.
Mentored students and doctoral researchers on scientific Python, HPC techniques, and open-source development practices.
Background
Roles, education, and collaboration style
Research roles across labs and universities, centered on long-lived software and collaborative delivery.
2009-present
Argonne National Laboratory
Principal Specialist in Research Software Engineering, working across UXarray, FLASH-X, CANDLE/IMPROVE, MeshKit, and urban simulation software efforts.
2023-present
University of Chicago
Staff At-Large with joint research activity spanning cancer pharmacogenomics and Earth system science.
2007-2009
Arizona State University
Research and teaching assistant in structural and computational mechanics, focused on blast mitigation and FEM-based design optimization.
- M.S. Computer Science University of Chicago, 2020
- M.S. Structural Engineering Arizona State University, 2009
- B.Tech. Mechanical Engineering IIT ISM Dhanbad, 2006
Contact
Available for conversations about research software and scientific computing.
I am interested in work across scientific computing, AI for health, climate modeling, reproducible workflows, and long-lived open-source systems.