Surrogate Modeling for
Computationally Expensive
Simulations of Supernovae inHigh-Resolution Galaxy
Simulations
Keiya Hirashima1
With Kana Moriwaki1, Michiko S. Fujii1, Yutaka Hirai2,3, Takayuki R.
Saitoh4, Junichiro Makino4 , and Shirley Ho5
1University of Tokyo, Japan, 2Tohoku University, Japan, 3University of Notre Dame, USA
4Kobe University, Japan, 5CCA, Flatiron Institute, USA

2023/10/19

CPS Seminar

Credit:
ESO/PESSTO/S. Smartt
Butusova Elena/Shutterstock.com
Art tools design/Shutterstock.com
anttoniart/Shutterstock.com
Gorodenkoff/Shutterstock.com

Self-Introduction
Keiya Hirashima•PhD student (the University of Tokyo)•Research interest• Galaxy formation and evolution• Surrogate modeling for SN feedback with machine learning•Some other projects…
• Surrogate modeling with GNNs/transformers• Foundation model to describe the morphology of galaxies

1Outline• Introduction: Overheads of galaxy simulations• #1 A Hamiltonian splitting with deep learning• #2 Surrogate modeling for the SN feedback• Summary2

Galaxy Simulations Using SPH*
10 ~ 102 pc~102 yr

ASURA-FDPS (N-body/SPH)(based on Saitoh+08,09, Iwasawa+16 )• Gravity + Hydrodynamics• (DISPH; saitoh+13)• Radiative Cooling/Heating (Ferland+17)• Star formation (Hirai in prep.)• Feedback<102 kpc• SNe Ia/IIb, AGB, Neutron star merger~108 yr

The formation of the galaxy [1].
*SPH: Smoothed Particle Hydrodynamics[1] https://www.youtube.com/watch?v=DDsUXRfs6ZQ[2] Applebaum et al. (2021)[3] Grand et al. (2021)• Chemical evolution (CELib; Saitoh17)• FUV background• About to accomplish star-bystar simulations… but…
• Previous work [2, 3]： 10! 𝑀⊙• Our goal (ASURA-FDPS)： 10𝑀⊙
3

Overheads in Galaxy Formation Simulations
The parallelization eﬃciency saturates at ~103 CPU cores.
e.g.,• GADGET-4(Springel+21)• DC Justice League (Applebaum+21)• Fire-2(Hopkins+18)•Due to small timescale regions (e.g. SNe),the communication overhead occurs.

•Even the latest supercomputers cannot
solve it (e.g., Fugaku has ~106 CPU cores).

BAD

GOOD

The Number of CPU cores

Strong Scaling of GADGET-4(Based on Figure 63 in Springel et al. 2021)

4Communication Overheads
The integration with short timesteps needs a huge numberof calculations and inter-node communications.
Calculation cost of 1 global timestep (∆𝑡~10! yr)

𝒑𝟏

𝒑𝟐

𝒑𝟐

𝒑𝟑

𝒑𝟑
𝒑𝟒

𝒑𝟏

SN (∆𝑡~10! yr)

Idle time𝒑𝟒
Credit: NAOJ

5

How can I avoid the bottleneck?

6

Outline• Introduction: Overhead of galaxy simulations• #1 A Hamiltonian splitting with deep learning• #2 Surrogate modeling for the SN feedback• Summary7Hamiltonian Splitting Method<102 kpc

Use multiple integration methods
depending on the properties of the regions.
(Fujii et al. 2007; Saitoh et al. 2010; Pelupessy et al. 2012)• Long timescale regions require
the global (long) timestep.
• Short timescale regions (SNe)require short timesteps with thegravitational perturbations.
Perturbation

10 ~ 102 pc

credit:[Upper: NASA/JPL-Caltech/ESO/R. Hurt,Lower-Right: NAOJ, Lower-Left: ESO/M. Kornmesser]
8

Avoid Communication Overheads
𝒑𝟐

MPI̲COMM̲WORLD

𝒑𝟏

𝒑𝟏

𝒑𝟑

𝒑𝟐
𝒑𝟒

𝒑𝟑
𝒑𝟒

Idle time
MPI̲COMM̲WORLD
𝒑𝟏

𝒑𝟐
𝒑𝟑

𝒑𝟒

MPI̲COMM̲GLOBAL(>104 CPU cores)

MPI̲COMM̲SN(<103 CPU cores)

𝒑𝟏

𝒑𝟐

𝒑𝟒

𝒑𝟑

Idle time9

Avoid Communication Overheads
𝒑𝟐

𝒑𝟏
𝒑𝟑
𝒑𝟒

• Need to know which particles will have small
timesteps in advance• Predict the boundaries of the shells using deeplearning
MPI̲COMM̲WORLD

𝒑𝟏

𝒑𝟐
𝒑𝟑

𝒑𝟒

MPI̲COMM̲GLOBAL(>104 CPU cores)

MPI̲COMM̲SN(<103 CPU cores)

𝒑𝟏

𝒑𝟐

𝒑𝟒

𝒑𝟑

Idle time10

Training Data (3D cartesian grids)
SN simulations in inhomogeneous turbulent clouds

Temperature

10 [K]

Mean ambient density

40 ~ 60 [cm-3]

Input energy

1051 [erg]Total mass

106 [M⨀]

Mass of a gas particle

1 [M⨀]

Softening parameter

0.5 [pc](the equivalent resolution for our galaxy simulations)

60 pc

60 pc

Interpolation using
SPH summation

60 pc

323 voxels (density distribution)11

Deep Learning Model• Memory-In-Memory Network (Wang et al. 2018)• Generates 2D video in the future• Convolutional Neural Networks (CNNs)• Convolution is the operation like inner products.
• Increased the dimensionality: 2D à 3D

12

Letʼs learn the SN explosions with kernels!
Example. Edge detection

⨂0

10

1-4

10

10

=

Learnable parameters(Convolutional Neural Networks)
13

Forecast the Evolution of the SN Shells in Density• Taking the initial distribution, the ML model can forecast thespatiotemporal change in density due to a SN explosion.
• Duration: 0.2 Myr𝐸& '/) */)𝑅(𝑡) ∝𝑡
𝜌&60 pc

Cross section view
Left: Simulation data (ground truth)
Right: Prediction by the 3D-MIM

Hirashima et al. 2023, accepted
arXiv:2302.00026
14

Detect particles with small ∆𝑡• With image processing, select (d) pale-gray region using(a) Initial Condition and (c) prediction• The continuous region where the density becomes <90%
is selected as the target region.
• Calculate the selected particles locally.
15

To the further speed-up
Sub-grid models

Δt ~105 yr

SN Δt = ~102 yr

t=t0 yr
The rest of the region in the galaxy
Δt ~105 yr

t=t0 + 105 yr
16

We needmorephysics!!

17

Outline• Introduction: Overhead of galaxy simulations• #1 A Hamiltonian splitting with deep learning• #2 Surrogate modeling for the SN feedback• Summary18

Training Data (3D cartesian grids)•Mass resolution: 1M⨀•

Summation with the SPH kernel

19Simulations for SN feedback

Large number of calcs. and comms.

SN Δt = ~102 yrt=t0 yr

t=t0 + 105 yr

The rest of the region in the galaxy
Δt ~105 yr
20

Surrogate modeling for SN feedback
Unet
U-Net (CNN-based)

Physicalquantities

Bypass with ML and MCMCInterpolation
with SPHsummation

MCMC(Gibbs Sampling)

SN Δt = ~102 yr

t=t0 yrt=t0 + 105 yr

The rest of the region in the galaxy
Δt ~105 yr
21

Prediction
Inhomogeneous shell emerges by the dense filament blocking.

10 pc

22

Prediction #2
Dense region：Velocities are estimated at a lower values.

23Gibbs Sampling (2D)

Representation by Gibbs Sampling
The particle distribution is sampled by a Markov chain Monte Carlo method.
MCMC(Gibbs Sampling)

3D cartesian grids of density~ 3D probability density function25

ML + MCMC
t=t0 yr
Physicalquantities

U-Net

Gibbs Sampling
SPH kernel

t=t0 + 105 yr

26

Validation• Ground Truth: 1Msun sims（High-resolution）• Baseline: 10 Msun sims（“Low”-resolution）• High resolution for galaxy simulations, cannot resolve
supernova feedback (Hu+16, Steinwandel+20, Hirashima+23)• Compered total thermal energy and outer momentum of hot
particles (<103 K) on test data (100 results of independentsimulation)“Low”-resolution

High-resolution

Our method27

Fidelity Evaluation in Thermal Energy• Compered to the low-res. sims., our method can duplicate thethermal energy
Ours
Low-res.

High-res.

High-res.
28

Fidelity Evaluation in Outer Momentum• Almost similar to the low-res. sims.
• Need to get the bias removed
Low-res.

Ours

High-res.

High-res.
29

Experiment | Rerun the simulation
t=t0 yr
Physicalquantities

U-Net

SPH kernel?
~106 yr30

Preliminary Restart Run from 0.1 Myr
DNS

ML+MCMC31

Outline• Introduction: Overhead of galaxy simulations• A Hamiltonian splitting with deep learning• Surrogate modeling for the SN feedback• Summary32

To conclude
Surrogate modeling

Hamiltonian Splitting

Direct Numerical
Simulations

Hirashima et al. 2023, acceptedfor MNRAS
arXiv:2302.00026• Stochastic• Fast• Deterministic• Expensive33

Accelerates 3D CNN Inference
SoftNeuro® [1] optimized 3D Unet for Fugaku!

Elapsed time
CPU usage

TF (Desktop, OneDNN)(16 cores / 32 threads)

TF (Fugaku)(1 node / 48 cores)

Softneuro (Fugaku)(1 node / 48 cores)

710 [ms]

2820 [ms]147 [ms]

36%

3.5%

68%

x19 FASTER!
1. https://arxiv.org/abs/2110.06037
2. https://www.morphoinc.com/
34

Summary and Future WorkØ An approach for surrogating SN simulations with ML and MCMC
Ø Need more accurate predictions
Ø try generative models?
Ø Need higher-res simulations
Ø Preparing for higher resolution (0.1 Msun) simulations
Ø Try particle-based models (e.g. Graph Neural Networks)Thank You!!!!

35