Rajeev Jain
Research Software Engineer | ML Infrastructure | Scientific Computing
Rajeev Jain
I build research software for climate analysis, AI workflows, and large simulation codes at Argonne National Laboratory, with a joint appointment at the University of Chicago.
My work sits where scientific computing meets usable engineering: climate and mesh analysis, machine-learning infrastructure, high-performance I/O, and tools that researchers can actually extend. I care most about software that survives beyond a single project cycle.
AI weather modeling on Aurora
Porting and stabilizing large weather-model training workflows on Intel GPUs, with attention to portability, runtime behavior, and scientific throughput.
UXarray
Mesh-aware analysis for next-generation climate grids, built so scientific users can inspect, subset, and reason about unstructured datasets without bespoke code for each mesh.
Scientific MCP workflows
Natural-language dataset discovery, plotting, workflows, and remote execution for scientific analysis across local machines and HPC systems.
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.
- Extending the project with an MCP server and AI-agent workflow for natural-language dataset exploration across local and HPC execution.
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
Capabilities across scientific software and systems
A compact view of the engineering areas I work in most often, from ML and analysis pipelines to HPC runtime and open-source delivery.
Data, models, and analysis workflows
- PyTorch, TensorFlow, NumPy, Pandas, Xarray, and Scikit-learn for model development and scientific analysis.
- Parsl and Swift/T for large experiment campaigns and repeatable workflow execution.
Performance, portability, and I/O
- MPI, OpenMP, HDF5, NetCDF, MOAB, and storage-aware design for leadership-class machines.
- Intel GPUs, CUDA portability, asynchronous checkpoint workflows, and runtime debugging on shared systems.
Software other researchers can extend
- Release engineering, CI pipelines, testing, and packaging for long-lived scientific projects.
- Roadmap ownership, collaboration across labs and universities, and mentoring around sustainable engineering practice.
Problem-driven engineering
- Climate modeling, unstructured mesh analysis, cancer pharmacogenomics, computational physics, and simulation software.
- AI infrastructure and reproducible workflows built around domain constraints rather than generic demos.
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
Links
Profiles and research links.
I keep the public site fairly minimal. LinkedIn is the best route for professional outreach.