globalemu Tutorial
This tutorial will show you the basics of training and evaluating an instance
of globalemu
. If you are just interested in evaluating the released models
then take a look at the second part towards the bottom of the page.
If you are intending to work with neutral fraction histories then the frame
work for training and evaluating models is identical you just need to pass
the kwarg xHI=True
to the pre-processing function, process(),
and model building function, nn(), discussed below.
The tutorial can also be found as a Jupyter notebook here.
Training an instance of globalemu
This tutorial will show you how to train a globalemu
model on simulations
of the Global 21-cm signal.
The first thing we need to do is download some 21-cm signal models to train our network on. For this we will use the 21cmGEM models and the following code.
import requests
import os
import numpy as np
data_dir = 'downloaded_data/'
if not os.path.exists(data_dir):
os.mkdir(data_dir)
files = ['Par_test_21cmGEM.txt', 'Par_train_21cmGEM.txt', 'T21_test_21cmGEM.txt', 'T21_train_21cmGEM.txt']
saves = ['test_data.txt', 'train_data.txt', 'test_labels.txt', 'train_labels.txt']
for i in range(len(files)):
url = 'https://zenodo.org/record/4541500/files/' + files[i]
with open(data_dir + saves[i], 'wb') as f:
f.write(requests.get(url).content)
In order for globalemu
to work the training data needs to be saved in the
data_dir
and in the files ‘train_data.txt’ and ‘train_labels.txt’ which
are the inputs and outputs of the network respectively.
Once the files have been downloaded we can go ahead and perform the
preprocessing necessary for globalemu
to effectively train a model. We
do this with the process()
function found in globalemu.preprocess
.
from globalemu.preprocess import process
base_dir = 'results/'
z = np.linspace(5, 50, 451)
num = 1000
process(num, z, base_dir=base_dir, data_location=data_dir)
Since this tutorial is only ment to demonstrate how to train a model with
the globalemu
code we are only going to pre-process 1000 models and train
with 1000 models out of a possible ~24000. We do this by setting num=1000
above but if we wanted to train on all the models we would set num='full'
.
Importantly the pre-processing function takes the data in data_dir
and
saves a .csv
file in the base_dir
containing the preprocessed inputs
for the neural network. It also saves some files used for normalisation in
the base_dir
so that when evaluating the network the inputs and outputs
can be properly dealt with.
By default the network subtracts and astrophysics free baseline from the models and resamples the signals at a higher rate in regions of high variation across the training data. Both of these pre-processing techniques are detailed in the globalemu MNRAS preprint. Users can prevent this happening by passing the kwargs AFB=False and resampling=False to process() if required.
Once pre-processing has been performed we can train our network with the
nn()
function in globalemu.network
.
from globalemu.network import nn
nn(batch_size=451, epochs=10, base_dir=base_dir, layer_sizes=[8])
nn()
has a bunch of keyword arguments that can be passed if required.
All are documented and all have default values. However you will likely
need to change things like base_dir
which tells the code where the
pre-processed data is and also layer_sizes
which determines the network
architecture. epochs
is the number of training calls and often the default
will be insufficient for training the network.
The code saves the model and loss history every ten epochs in case your computer crashes or the program is interrupted for some unforeseen reason. If this happens or you reach the max number of epochs and need to continue training you can do the following and the code will resume from the last save.
nn(batch_size=451, epochs=10, base_dir=base_dir, layer_sizes=[8], resume=True)
You have now successfully trained an instance of globalemu.
Evaluating an instance of globalemu
We can go ahead and evaluate the model using the testing data that we downloaded earlier.
test_data = np.loadtxt(data_dir + 'test_data.txt')
test_labels = np.loadtxt(data_dir + 'test_labels.txt')
With the data loaded we will look at how the model performs when predicting
the first signal in the data set. We do this with the evaluate()
class
in globalemu.eval
which takes in a set of parameters and returns a signal.
The class must first, however, be initialised with a set of kwargs.
We supply a base_dir
which contains the pre-processed data,
normalisation factors and trained model. You can also pass a redshift range with the
z
kwarg however if this isn’t supplied than the function will return the
signal at the original redshifts that were used for training.
from globalemu.eval import evaluate
input_params = test_data[0, :]
true_signal = test_labels[0, :]
predictor = evaluate(base_dir=base_dir)
signal, z = predictor(input_params)
import matplotlib.pyplot as plt
plt.plot(z, true_signal, label='True Signal')
plt.plot(z, signal, label='Emulation')
plt.legend()
plt.ylabel(r'$\delta T$ [mK]')
plt.xlabel(r'$z$')
The emulation is pretty poor for several reasons; we didn’t run the training for long enough (only 20 epochs), the network size is small and we used very little of the available training data.
We can have a look at the same signal emulated with the released model on github. This was trained with a much more appropriately sized network, the full training data and a few hundred epochs. The results are therefore more similar to the true signal.
predictor = evaluate(base_dir='../T_release/')
signal, z = predictor(input_params)
plt.plot(z, true_signal, label='True Signal')
plt.plot(z, signal, label='Emulation')
plt.legend()
plt.ylabel(r'$\delta T$ [mK]')
plt.xlabel(r'$z$')
In addition to evaluating one model at a time a user can also evaluate a set of parameters using the emulator.
input_params = test_data[:5, :]
true_signal = test_labels[:5, :]
signal, z = predictor(input_params)
for i in range(len(true_signal)):
if i==0:
plt.plot(z, true_signal[i, :], c='k', ls='--', label='True Signal')
plt.plot(z, signal[i, :], c='r', label='Emulation')
else:
plt.plot(z, true_signal[i, :], c='k')
plt.plot(z, signal[i, :], c='r')
plt.legend()
plt.ylabel(r'$\delta T$ [mK]')
plt.xlabel(r'$z$')
Further Evaluation
The function globalemu.plotter.signal_plot()
can also be used to assess the
quality of emulation. This function is designed to plot the mean, 95th
percentile and worse emulations, based on a given loss function,
of a set of signals given their corresponding parameters.
from globalemu.eval import evaluate
from globalemu.plotter import signal_plot
predictor = evaluate(base_dir='../T_release/')
parameters = np.loadtxt('downloaded_data/test_data.txt')
labels = np.loadtxt('downloaded_data/test_labels.txt')
plotter = signal_plot(parameters, labels, 'rmse', predictor, '../T_release/',
loss_label='RMSE = {:.4f} [mK]')
This particular example uses the 'rmse'
loss function that is built into
the emulator but an alternative function can be provided by the user (see
documentation for details). The graph that is produced gets saved in the
provided base_dir
, in this case 'T_release/'
and looks like the
below figure.
Downloading Trained Models
The released trained models can be directly downloaded from the github pages or a built in helper function can be used to download the models. The function can be called like so
from globalemu.downloads import download
download().model() # Redshift-Temperature Network
download(xHI=True).model() # Redshift-Neutral Fraction Network
which will download the released models into the present working directory
and the files T_release/
and xHI_release
.