A Probabilistic Machine Learning Approach to Emulate Chemistry in Astrophysical Simulations

Authors: Lennart Buhlmann & Felix Rauprich

We are developing a machine learning framework to predict the chemical evolution in astrophysical simulations. Our model is based on a neural network that predicts probability distributions rather than single values, using a Gaussian Mixture Model approach (Bishop, 1995).

To train the model, we first generate a sample of training points from astrophysical simulations. The physical data in the sample cover multiple orders of magnitude. In order to reduce the data range, we transform the physical sample to the latent space. For this, we treat the chemistry data in logarithmic space and normalize the data with respect to the mean value and standard deviation. We further aim for a homogeneous distribution for each output variable.

At the beginning of the training process, the network assumes a random distribution of Gaussians that represent our probability function. During training, our network iteratively learns and refines these probability distributions based on continuous loss calculations.

To evaluate the model's performance, we compare the predicted probability distributions with the corresponding true values in the latent space. An example is shown in the figure: the x-axis represents the true values, while the y-axis shows the predicted values. The color scale indicates the summed probability density on a logarithmic scale. The closer the distribution is centered around the one-to-one line, the more accurate the prediction.

Based on the given input values, the trained network predicts probability distributions for the output values. However, simulations cannot handle probability distributions, as they expect discrete physical values. Therefore, we draw individual samples from the predicted probability distributions using stochastic sampling and transform the results back to the physical space.

At present, we are working on improving the training process by incorporating physical knowledge directly into the training procedure through the optimization of the loss function, e.g., using truncated Gaussian models to apply physical limits to the model.