Diffusion-based Galaxy Simulations for the Roman High Latitude Survey

Summary

This repository contain the dataset used to train and evaluate the diffusion model for Roman-like galaxy simulations, galaxies generated by the model, the trained model weights, and a tagged release of the software.

• train_2p0cut.hdf5
  • 10,876 galaxies for model training (described in Section 3).
• test_2p0cut.hdf5
  • 2,734 galaxies for model evaluation (described in Section 3).
• model.zip
  • An exported Hugging Face pipeline, including configuration JSONs and a model weight .safetensor
  • Created with train.py
• generated_2p0cut.hdf5
  • 10,000 generated galaxies of shape [10000, 3, 56, 56] (described in Section 4).
  • Created with predict_fast.py
• roman-galaxy-ddpm-1.0.0.zip
  • Source code as released at https://github.com/diana-sco/roman-galaxy-ddpm/releases/tag/v1.0.0

Data Description

train_2p0cut.hdf5 and test_2p0cut.hdf5 contain the following datasets:

% h5ls train_2p0cut.hdf5
ids                      Dataset {10876}
stamps                   Dataset {10876, 3, 56, 56}
% h5ls test_2p0cut.hdf5 
ids                      Dataset {2734}
stamps                   Dataset {2734, 3, 56, 56}

where ids contain a unique ID for each galaxy for internal tracking, and stamps contains the 3-channel galaxy data.

generated_2p0cut.hdf5 contain only the stamps:

% h5ls generated_2p0cut.hdf5 
stamps                   Dataset {10000, 3, 56, 56}

model.zip contains the following:

model
├── model_index.json
├── scheduler
│   └── scheduler_config.json
└── unet
    ├── config.json
    └── diffusion_pytorch_model.safetensors

predict_fast.py expects a path to the root directory, not just the safetensor model weights.

Please refer to the origin GitHub repository for the software release zip.

Acknowledgements

The research was carried out at the Jet Propulsion Laboratory, California Institute of Technology, under a contract with the National Aeronautics and Space Administration (80NM0018D0004), © 2026. All rights reserved. In particular, this work was funded through the Jet Propulsion Laboratory's Spontaneous Concept Research and Technology Development program, which supported this research. The High Performance Computing resources used in this work were provided by funding from the JPL Enterprise Technology, Strategy, and Cybersecurity Directorate. The authors also acknowledge the Texas Advanced Computing Center (TACC) at The University of Texas at Austin for providing computational resources that have contributed to the research results reported within this paper. Portions of this work were completed at Duke University. D.S. thanks Michael Troxel and Arun Kannawadi for insightful comments that improved the manuscript. This work is based on observations made with the NASA/ESA/CSA James Webb Space Telescope. The data were obtained from the Mikulski Archive for Space Telescopes at the Space Telescope Science Institute, which is operated by the Association of Universities for Research in Astronomy, Inc., under NASA contract NAS 5-03127 for JWST. These observations are associated with program #3215 and #1963.