tSZ DeepSet

Learn electron pressure in clusters from a gravity-only simulation. The network architecture consists mostly of DeepSets and MLPs, enabling a modular, interpretable design that operates directly on the simulation representation as a set of dark matter (DM) particles.

Dependencies

• cosmological simulation with gravity-only and full-physics runs from same initial conditions. We used IllustrisTNG 300-1.
• halo finder. We used Rockstar.
• PyTorch
• Optuna
• Voxelize for create_boxes.py
• group_particles for collect_particles.cpp

Use guide

0. Code that needs to be compiled:

   • collect_particles.cpp: compile into both collect_particles_DM and collect_particles_TNG executables, depending on whether preprocessor macros DM or TNG are defined. This script finds the simulation particles in the vicinity of halos.
   • sort_particles.cpp into sort_particles executable.
   • prtfinder.cpp into libprtfinder.so.

1. WORKFLOW describes the steps that need to be taken for pre-processing of the simulation data. The required scripts are:

   • ascii_to_hdf5.py
   • collect_particles.cpp
   • sort_particles.py
   • create_halo_catalog.py
   • create_boxes.py
   • create_Nprt.py
   • create_residuals.py

2. Optional: run create_normalization.py and analyse_normalization.py, if used on a different simulation than IllustrisTNG 300-1. The output can be used to tweak the hardcoded numbers in normalization.py to bring inputs to zero mean, unit variance.

3. Internal data handling is defined in

   • data_modes.py
   • data_loader.py
   • data_set.py
   • data_batch.py
   • data_item.py
   • halo_catalog.py
   • halo.py

4. The cluster-scale scalars and vectors are defined in

   • global_fields.py
   • basis.py
   • fixed_len_vec.py

5. Runtime configuration is defined in cfg.py. The idea is that this module is populated from either command line settings or the environment variable TSZ_DEEP_SET_CFG by init_proc.py. For later reference, archive_cfg.py is used. Quite often we set default arguments to something from cfg.py, this is implemented in default_from_cfg.py.

6. There is some basic support for distributed training, although more work would be required to get this to work. Currently training is fast enough on a single GPU and the required RAM per training process is about 60G. Some possibly buggy definitions to set up the distributed environment are in mpi_env_types.py.

7. The network architecture is defined in

   • network.py combined architecture
   • network_batt12.py GNFW module
   • network_origin.py Origin module
   • network_local.py Local module
   • network_vae.py Stochastic module
   • network_decoder.py Aggregator module
   • network_deformer.py deviations from spherical symmetry
   • network_encoder.py cluster-scale vector DeepSet
   • network_scalarencoder.py cluster-scale scalar DeepSet
   • network_mlp.py single MLP
   • network_layer.py DeepSet primitive

   Initialization of the network in init_model.py.

8. Driver code is in

   • training.py Can be directly executed with command line arguments setting the cfg.py settings.
   • testing.py To test a trained model on either validation or testing set.
   • optuna_driver.py Wrapper around training.py with Optuna sampling.

9. The files generate_cl_*.py are used for Optuna hyperparameter searches.

10. Auxiliary files for training:

    • training_loss.py
    • training_optimizer.py
    • training_loss_record.py
    • plot_loss.py

11. Auxiliary files for testing:

    • testing_loss_record.py
    • cubify_prediction.py
    • plot_testing_loss.py
    • plot_prediction.py
    • plot_profiles.py

12. Files paper_plot_*.py were used to generate publication figures.

13. Various other files are there, most are buggy and not for use. In particular, anything with FOF in the name was for our initial attempt to work with friends-of-friends instead of Rockstar halos.