Training and using a new surrogate in fiesta#
This notebook demonstrates how to train a new fiesta surrogate model from scratch, using the kilonova model of Dietrich et al. (2020) (Bu2019lm) as an example: preprocessing the raw POSSIS training data into fiesta’s training format, training MLP and CVAE surrogates on it, and benchmarking the result.
This needs a working environment with fiesta installed.
Checking for an existing surrogate#
Before training a new surrogate from scratch, it’s worth checking whether fiesta already ships one for the model we’re after. Fiesta has both built-in surrogates (already installed) and downloadable ones (hosted on Hugging Face).
[1]:
from fiesta.surrogates import print_built_in_surrogates
print_built_in_surrogates()
17:08 fiesta INFO : Available built-in surrogates in fiesta are:
17:08 fiesta INFO : afgpy_gaussian_CVAE (GRB)
17:08 fiesta INFO : pbag_gaussian_CVAE (GRB)
17:08 fiesta INFO : pbag_tophat_CVAE (GRB)
17:08 fiesta INFO : Bu2026_MLP (KN)
And which surrogates are available for download:
[2]:
from fiesta.surrogates import print_downloadable_surrogates
print_downloadable_surrogates()
17:08 fiesta INFO : Downloadable surrogates for fiesta are:
17:08 fiesta INFO : ====================
17:08 fiesta INFO : -KN-
17:08 fiesta INFO : Brethauer2025_MLP
17:08 fiesta INFO : Bu2025_MLP
17:08 fiesta INFO : Brethauer2025_CVAE
17:08 fiesta INFO : Bu2026_MLP
17:08 fiesta INFO : Bu2025_CVAE
17:08 fiesta INFO : Bu2026_CVAE
17:08 fiesta INFO : Bu2025_lc
17:08 fiesta INFO : ====================
17:08 fiesta INFO : ====================
17:08 fiesta INFO : -GRB-
17:08 fiesta INFO : afgpy_tophat_CVAE
17:08 fiesta INFO : afgpy_gaussian_CVAE
17:08 fiesta INFO : blastwave_gaussian_CVAE
17:08 fiesta INFO : pbag_gaussian_CVAE
17:08 fiesta INFO : pbag_gaussian_MLP
17:08 fiesta INFO : afgpy_tophat_MLP
17:08 fiesta INFO : blastwave_rs_gaussian_CVAE
17:08 fiesta INFO : pbag_tophat_CVAE
17:08 fiesta INFO : afgpy_gaussian_MLP
17:08 fiesta INFO : ====================
Training a Bu2019lm surrogate from scratch#
Now we turn to the case where we want to train a new surrogate.
Specifically, we want to train a surrogate for the kilonova model from Dietrich et al. (2020) (https://arxiv.org/abs/2002.11355), commonly referred to as Bu2019lm in the literature.
Before training it ourselves, let’s check whether fiesta already ships this as a built-in or downloadable surrogate (see the print_built_in_surrogates and print_downloadable_surrogates calls above): the downloadable list only contains Bu2026_*, Brethauer2025_*. and Bu2019 models. But Bu2019 is easy to re-train and huts a good candidate for this demo.
The training data (the raw POSSIS spectral simulations, the .dat lightcurve files) is hosted as a zip file on Google Drive. Below we download that zip file into ./Bu2019_training_data using gdown. The original set of files is available here, the link we point to is just pointing to a single zip file containing all those files, since downloading each file individually might take a long time.
This requires gdown (a small utility for downloading files from Google Drive), which is listed as an optional dev dependency of fiestaEM rather than a core one. If it’s not already installed in your environment, install it first with pip install -e ".[dev]" (or uv sync --extra dev if you’re using uv), then re-run the cell below. (or just pip install gdown).
We then use it to download the training-data zip file and extract it into ./Bu2019_training_data, next to this notebook. Everything under ./Bu2019* is git-ignored, so running this notebook won’t leave any downloaded or generated files tracked by git.
[3]:
import os
import zipfile
import gdown
os.makedirs("./Bu2019_training_data", exist_ok=True)
# https://drive.google.com/file/d/1u8XVxE-lOc8sbOyL5Xfd_QE7La1mtzXD/view?usp=sharing
zip_path = "./Bu2019_training_data/2020_dietrich_Bu2019lm.zip"
gdown.download(id="1u8XVxE-lOc8sbOyL5Xfd_QE7La1mtzXD", output=zip_path, quiet=False)
with zipfile.ZipFile(zip_path) as zf:
zf.extractall("./Bu2019_training_data/Bu2019lm_training_data")
Downloading...
From (original): https://drive.google.com/uc?id=1u8XVxE-lOc8sbOyL5Xfd_QE7La1mtzXD
From (redirected): https://drive.google.com/uc?id=1u8XVxE-lOc8sbOyL5Xfd_QE7La1mtzXD&confirm=t&uuid=86243753-cdfb-4ac8-b212-95f44046695c
To: /home/koehn/Documents/Fahren_retreat_2026/retreat/hauke_presentation_fiesta/fiestaEM/docs/user_guide/training/Bu2019_training_data/2020_dietrich_Bu2019lm.zip
100%|██████████| 258M/258M [00:04<00:00, 61.6MB/s]
a) Creating a training data file#
In fiesta, training a machine-learning surrogate starts from a .h5 file that contains all the raw training data plus the metadata.
This .h5 file must have a certain layout. Specifically, it has to include the following data sets as metadata:
times: An array for the time domain of the data in days.nus: An array for the frequency domain of the data in Hz.parameter_names: A list of strings that contains the parameter names. This determines which parameter names need to be present in the param-dict that is the argument for the surrogate prediction.parameter_distributions: A string-converted dictionary that hasparameter_namesas keys and the values are tuplestuple[float, float, str]. The first two numbers are the minimum and maximum range of this parameter in the training data, i.e., the range in which the trained surrogate will be valid. The string should indicate which distribution the training parameter samples follow, though there are no negative side-effects should the distribution provided here be inaccurate.
Further, there should be three data sets train, val, and test. Each of these contains an array X and an array y.
The X array is the array with the parameter vectors. The y array is the array with the flux densities at 10 pc (but zero redshift), in units of log10(mJy).
For fiesta it does not really matter, how this file is created, as long as it adheres to the aforementioned format. But of course, there are utility functions to convert the outputs of common codes like POSSIS or SEDONA to the fiesta format.
The raw Bu2019 data is not in the format fiesta's own convert_POSSIS_outputs_to_h5 utility function expects. Instead, it’s in POSSIS’s older plain-text ascii format. For this reason, we have to define some helper functions:
[4]:
import re
import numpy as np
import tqdm
import jax.numpy as jnp
from fiesta.conversions import Flambda_to_Fnu
from fiesta.train.utils import train_test_split, write_training_data
def read_POSSIS_ascii_parameters(filename: str) -> tuple[float, float, float]:
"""
Reads parameters from filenames
``nph<N>_mejdyn<value>_mejwind<value>_phi<value>.txt``.
"""
match = re.search(r"mejdyn([\d.]+)_mejwind([\d.]+)_phi(\d+)", os.path.basename(filename))
if match is None:
raise ValueError(f"Could not parse mejdyn/mejwind/phi from filename {filename}.")
mej_dyn, mej_wind, phi = match.groups()
return float(mej_dyn), float(mej_wind), float(phi)
def read_POSSIS_ascii_file(filename: str, cos_theta_grid: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
"""
Reads a single POSSIS ascii output file, as used for the Dietrich et al. 2020 ``Bu2019lm``
kilonova grid. The ascii layout is: line 1 is the number of viewing angles ``nobs`` stored in the
file, line 2 the number of wavelength bins, line 3 the number of time bins with its (tmin, tmax)
range in days. The remaining lines contain, per viewing angle (stacked one after another), one row
per wavelength bin: the wavelength in Angstrom followed by the flux density in erg/s/cm^2/Angstrom
at a fiducial distance of 10 pc for every time bin.
Some files (the axisymmetric phi=0 and phi=90 deg limits) only contain a single viewing angle
because the light curve does not depend on the observer angle there; these are replicated across
``cos_theta_grid`` so every returned sample has the same parameter dimensionality.
Args:
filename (str): Path to the ascii file.
cos_theta_grid (np.ndarray): Grid of cos(viewing angle) values that multi-angle files are expected to be sampled on (uniform in [0, 1]).
Returns:
tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: ``X`` (shape ``(n_obs, 4)``, columns ``[mej_dyn, mej_wind, phi, cos_theta]``), ``y`` (shape ``(n_obs, n_nus, n_times)``, log10 of the spectral flux density in mJy at 10 pc), ``times`` (days), ``nus`` (Hz).
"""
mej_dyn, mej_wind, phi = read_POSSIS_ascii_parameters(filename)
with open(filename, "r") as f:
nobs = int(f.readline().split()[0])
nwave = int(f.readline().split()[0])
ntime, tmin, tmax = f.readline().split()
ntime, tmin, tmax = int(ntime), float(tmin), float(tmax)
data = np.loadtxt(f, max_rows=nobs * nwave)
dt = (tmax - tmin) / ntime
times = tmin + dt * (np.arange(ntime) + 0.5) # bin centers, see the grid's README
data = data.reshape(nobs, nwave, ntime + 1)
waves = data[0, :, 0] # Angstrom, shared by every viewing angle block
flux = np.maximum(data[:, :, 1:], 1e-15) # erg / s / cm^2 / Angstrom at 10 pc
y_per_obs, nus = [], None
for i in range(nobs):
mJys, nus = Flambda_to_Fnu(jnp.array(flux[i]), jnp.array(waves))
y_per_obs.append(np.log10(np.array(mJys)))
y_per_obs = np.stack(y_per_obs)
if nobs == 1:
y = np.repeat(y_per_obs, len(cos_theta_grid), axis=0)
elif nobs == len(cos_theta_grid):
y = y_per_obs
else:
raise ValueError(f"{filename} has {nobs} viewing angles, expected 1 or {len(cos_theta_grid)}.")
X = np.array([[mej_dyn, mej_wind, phi, ct] for ct in cos_theta_grid])
return X, y, times, nus
def convert_POSSIS_ascii_outputs_to_h5(
dirs: str | list[str],
outfile: str,
parameter_names: list[str],
log_arguments: list[int],
n_inclinations: int = 11,
clip: float = 6.5144,
) -> None:
"""
Merges a directory (or several directories) full of POSSIS ascii outputs, as used for the
Dietrich et al. 2020 ``Bu2019lm`` kilonova grid, into a single training data file in the fiesta
format. See ``read_POSSIS_ascii_file`` for the expected ascii layout and file naming.
Args:
dirs (str | list[str]): directory or list of directories with the ascii files to merge.
outfile (str): Name of the ``.h5`` training data file to create.
parameter_names (list[str]): Parameter names for ``[mej_dyn, mej_wind, phi, cos_theta]``, in that order.
log_arguments (list[int]): Indices of parameters that should be converted to log10 for training (typically mej_dyn and mej_wind).
n_inclinations (int): Number of viewing angle grid points (cos(theta), uniform in [0, 1]) expected in the multi-angle files. Defaults to 11.
clip (float): Lower floor value for the minimum log10(mJy) at 10 pc. Every flux density below that will be set to that value. Defaults to 6.5144 (appr. 0 abs. mag).
"""
if isinstance(dirs, str):
dirs = [dirs]
files = []
for dir in dirs:
files.extend([os.path.join(dir, f) for f in os.listdir(dir) if f.endswith(".txt")])
inclination_EM_grid = np.arccos(np.linspace(0, 1, n_inclinations))
X, y, times, nus = [], [], None, None
for file in tqdm.tqdm(files):
X_file, y_file, times_file, nus_file = read_POSSIS_ascii_file(file, inclination_EM_grid)
if times is None:
times, nus = times_file, nus_file
X.append(X_file)
y.append(y_file)
X = np.concatenate(X)
y = np.concatenate(y)
if X.shape[1] != len(parameter_names):
raise ValueError(f"parameter_names do not match parameters stored in the ascii files ({X.shape[1]} parameters found).")
X[:, log_arguments] = np.log10(X[:, log_arguments])
y = np.maximum(y, clip)
train_X, val_X, train_y, val_y = train_test_split(X, y, train_size=0.8)
val_X, test_X, val_y, test_y = train_test_split(val_X, val_y, train_size=0.5)
parameter_distributions = {p: (np.min(train_X[:, j]).item(), np.max(train_X[:, j]).item(), "uniform") for j, p in enumerate(parameter_names)}
write_training_data(outfile, train_X, train_y, val_X, val_y, test_X, test_y, times, nus, parameter_names, parameter_distributions)
These will convert the older possis files to a fiesta training data file. Specifically, they
parse
mej_dyn,mej_wind, andphifrom the file name,read the ascii block(s) and converts the
F_lambdaflux densities to spectral flux densities in mJy with fiesta’s ownFlambda_to_Fnuconversionreplicate the single-angle files across the full
cos(theta)grid so every training sample has the same 4 parameters[log10_mej_dyn, log10_mej_wind, phi, cos_theta],and finally write everything into the
train/val/testlayout fiesta expects.
Since this ascii format is specific to this Bu2019lm demo (fiesta’s own converters, convert_POSSIS_outputs_to_h5 and convert_SEDONA_outputs_to_h5 in fiesta.train.utils, already cover the more common formats), we define the small preprocessing helper functions below right here in the notebook, rather than adding them to the fiesta package itself.
In the end we just need to call one function to convert the possis outputs to the desired format:
[5]:
convert_POSSIS_ascii_outputs_to_h5(
dirs="./Bu2019_training_data/Bu2019lm_training_data/2020_dietrich_Bu2019lm",
outfile="./Bu2019_training_data/Bu2019lm_training_data.h5",
parameter_names=["log10_mej_dyn", "log10_mej_wind", "phi", "inclination_EM"],
log_arguments=[0, 1], # mej_dyn and mej_wind are given in linear mass in the file names, convert to log10 for training
)
100%|██████████| 196/196 [00:05<00:00, 32.94it/s]
b) Training surrogates#
Now that we have a training data file, we can continue with our training procedure.
The intended way to interact with the training data file is through the DataManager class:
[6]:
from fiesta.train import DataLoader
data = DataLoader(
file="./Bu2019_training_data/Bu2019lm_training_data.h5",
tmin=0.1,
tmax=15,
numin=5e13,
numax=1e15,
)
data.print_file_info()
17:08 fiesta INFO : File info for ./Bu2019_training_data/Bu2019lm_training_data.h5:
17:08 fiesta INFO : Time range in file: 0.10 19.90 days
17:08 fiesta INFO : Frequency range in file: 3.00e+13 3.00e+16 Hz
17:08 fiesta INFO : Parameter distributions: {'log10_mej_dyn': (-3.0, -1.6989700043360187, 'uniform'), 'log10_mej_wind': (-2.0, -0.8860566476931632, 'uniform'), 'phi': (0.0, 90.0, 'uniform'), 'inclination_EM': (0.0, 1.5707963267948966, 'uniform')}
17:08 fiesta INFO :
17:08 fiesta INFO : Training data: 1724
17:08 fiesta INFO : Validation data: 216
17:08 fiesta INFO : Test data: 216
17:08 fiesta INFO : Special data:
17:08 fiesta INFO :
Here, the tmin, tmax, numin, and numax arguments are there to cut the flux densities to the desired time and frequency range.
We can now use the trainer classes to have a (more or less automatized) training process.
We first initialize a simple feed-forward multi-layer perceptron (MLP) neural network.
[7]:
from fiesta.train import NeuralnetConfig, MLP
import jax
config = NeuralnetConfig(
name="MLP",
input_size=len(data.parameter_names),
output_size=30, # determines the number of PCA components kept (see below)
hidden_layer_sizes=[64, 64],
learning_rate = 1e-2,
batch_size=128,
nb_epochs=15_000,
)
nn = MLP(
config=config,
key=jax.random.key(5892),
)
Now, we will use the FluxSurrogateTrainer to get the training data to the neural network we just initialized.
Note that internally, the raw data will be reduced to a lower dimensional representation in a preprocessing step. In case of the MLP architecture, the 2D flux densities are decomposed through principal component analysis (PCA). The number of PCA components that will be kept is determined by the output_size of the MLP network, so choose wisely. Usually, somewhere between 20 - 100 components is sufficient to ensure good reconstruction of the flux densities from the reduced representation.
The FluxSurroagateTrainer will link the data and the neural network. Its fit() method then starts the training process, where the learning rate and number of epochs have been determined in the NN`s config.
The fitting procedure here runs a little over 1 minute.
[8]:
from fiesta.train import FluxSurrogateTrainer, NeuralnetConfig
trainer_mlp = FluxSurrogateTrainer(
surrogate_name="Bu2019_MLP",
outdir="./Bu2019_MLP",
data=data,
network=nn,
plots_dir="./Bu2019_MLP",
)
trainer_mlp.fit()
trainer_mlp.save() # save the trained surrogate to outdir
trainer_mlp.plot_example_lc(["besselli", "bessellux"]) # plot two example light curves
17:08 fiesta INFO : Using the following data from ./Bu2019_training_data/Bu2019lm_training_data.h5 for training the surrogate:
17:08 fiesta INFO : Time range loaded: 0.10 15.10 days
17:08 fiesta INFO : Number of points in the time array: 76
17:08 fiesta INFO : Frequency range loaded: 4.99e+13 1.03e+15 Hz
17:08 fiesta INFO : Number of points in the frequency array: 287
17:08 fiesta INFO : Parameter names: ['log10_mej_dyn', 'log10_mej_wind', 'phi', 'inclination_EM']
17:08 fiesta INFO :
17:08 fiesta INFO : Training data: 1724
17:08 fiesta INFO : Validation data: 216
17:08 fiesta INFO :
100%|██████████| 1724/1724 [00:03<00:00, 440.16it/s]
17:09 fiesta INFO : PCA decomposition accounts for 89.05 % of the total variance in the training data. This value is hopefully close to 1.
17:09 fiesta INFO : Preprocessing data . . . done
17:09 fiesta INFO : Train loss at step 1: 8547.1201171875
17:09 fiesta INFO : Valid loss at step 1: 8901.6787109375
17:09 fiesta INFO : Best valid loss so far: 8901.6787109375
17:09 fiesta INFO : Learning rate: 0.01
17:09 fiesta INFO : ---
17:09 fiesta INFO : Train loss at step 1501: 488.8849792480469
17:09 fiesta INFO : Valid loss at step 1501: 597.4097900390625
17:09 fiesta INFO : Best valid loss so far: 594.9603881835938
17:09 fiesta INFO : Learning rate: 0.01
17:09 fiesta INFO : ---
17:09 fiesta INFO : Train loss at step 3001: 304.52960205078125
17:09 fiesta INFO : Valid loss at step 3001: 351.1684875488281
17:09 fiesta INFO : Best valid loss so far: 340.25201416015625
17:09 fiesta INFO : Learning rate: 0.01
17:09 fiesta INFO : ---
17:09 fiesta INFO : Train loss at step 4501: 204.88148498535156
17:09 fiesta INFO : Valid loss at step 4501: 226.63316345214844
17:09 fiesta INFO : Best valid loss so far: 225.09967041015625
17:09 fiesta INFO : Learning rate: 0.01
17:09 fiesta INFO : ---
17:09 fiesta INFO : Train loss at step 6001: 162.4371337890625
17:09 fiesta INFO : Valid loss at step 6001: 177.8101348876953
17:09 fiesta INFO : Best valid loss so far: 173.78521728515625
17:09 fiesta INFO : Learning rate: 0.01
17:09 fiesta INFO : ---
17:09 fiesta INFO : Train loss at step 7501: 148.68121337890625
17:09 fiesta INFO : Valid loss at step 7501: 166.36874389648438
17:09 fiesta INFO : Best valid loss so far: 140.80392456054688
17:09 fiesta INFO : Learning rate: 0.01
17:09 fiesta INFO : ---
17:09 fiesta INFO : Train loss at step 9001: 118.12957000732422
17:09 fiesta INFO : Valid loss at step 9001: 125.47672271728516
17:09 fiesta INFO : Best valid loss so far: 124.30897521972656
17:09 fiesta INFO : Learning rate: 0.01
17:09 fiesta INFO : ---
17:09 fiesta INFO : Train loss at step 10501: 123.16339111328125
17:09 fiesta INFO : Valid loss at step 10501: 124.9992446899414
17:09 fiesta INFO : Best valid loss so far: 112.96234893798828
17:09 fiesta INFO : Learning rate: 0.01
17:09 fiesta INFO : ---
17:09 fiesta INFO : Train loss at step 12001: 137.02357482910156
17:09 fiesta INFO : Valid loss at step 12001: 146.00357055664062
17:09 fiesta INFO : Best valid loss so far: 106.81938934326172
17:09 fiesta INFO : Learning rate: 0.01
17:09 fiesta INFO : ---
17:09 fiesta INFO : Train loss at step 13501: 101.27723693847656
17:09 fiesta INFO : Valid loss at step 13501: 107.45780944824219
17:09 fiesta INFO : Best valid loss so far: 101.57323455810547
17:09 fiesta INFO : Learning rate: 0.01
17:09 fiesta INFO : ---
17:10 fiesta INFO : Training for 15000 took 61.37080001831055 seconds.
17:10 fiesta INFO : Best validation loss: 98.61186981201172
17:10 fiesta INFO : Loading surrogate Bu2019_MLP. This surrogate should only be used in the following parameter ranges:
17:10 fiesta INFO : log10_mej_dyn: (-3.0, -1.6989700043360187)
17:10 fiesta INFO : log10_mej_wind: (-2.0, -0.8860566476931632)
17:10 fiesta INFO : phi: (0.0, 90.0)
17:10 fiesta INFO : inclination_EM: (0.0, 1.5707963267948966)
17:10 fiesta INFO : Surrogate Bu2019_MLP is loading with source-frame time range [ 0.1 15.1] days.
17:10 fiesta INFO : Surrogate Bu2019_MLP is loading with the following filters: ['besselli', 'bessellux'].
17:10 fiesta INFO : Loaded for surrogate Bu2019_MLP from ./Bu2019_MLP.
The other neural network architecture implemented is the conditional variational autoencoder (CVAE). However, its training is rather slow, so we do not execute it in this notebook: feel free to uncomment and train it yourself.
The fitting procedure here runs roughly 2 and a half minutes.
[9]:
from fiesta.train import CVAE
image_size = np.array([34, 45])
config = NeuralnetConfig(
name="CVAE",
hidden_layer_sizes=[200, 50, 30],
conditional_dim=len(data.parameter_names),
learning_rate = 3e-3,
batch_size=256,
nb_epochs=2_500,
)
cvae = CVAE(
config=config,
image_size=image_size,
key=jax.random.key(245890),
)
trainer_cvae = FluxSurrogateTrainer(
surrogate_name="Bu2019_CVAE",
outdir="./Bu2019_CVAE",
data=data,
network=cvae,
)
trainer_cvae.fit()
trainer_cvae.save()
trainer_cvae.plot_example_lc(["besselli", "bessellux"])
17:10 fiesta INFO : Using the following data from ./Bu2019_training_data/Bu2019lm_training_data.h5 for training the surrogate:
17:10 fiesta INFO : Time range loaded: 0.10 15.10 days
17:10 fiesta INFO : Number of points in the time array: 76
17:10 fiesta INFO : Frequency range loaded: 4.99e+13 1.03e+15 Hz
17:10 fiesta INFO : Number of points in the frequency array: 287
17:10 fiesta INFO : Parameter names: ['log10_mej_dyn', 'log10_mej_wind', 'phi', 'inclination_EM']
17:10 fiesta INFO :
17:10 fiesta INFO : Training data: 1724
17:10 fiesta INFO : Validation data: 216
17:10 fiesta INFO :
100%|██████████| 1724/1724 [00:00<00:00, 1987.42it/s]
17:10 fiesta INFO : Preprocessing data . . . done
17:10 fiesta INFO : Train loss at step 1: 1745.2486572265625
17:10 fiesta INFO : Valid loss at step 1: 1819.2696533203125
17:10 fiesta INFO : Best valid loss so far: 1819.2696533203125
17:10 fiesta INFO : Learning rate: 0.003
17:10 fiesta INFO : ---
17:10 fiesta INFO : Train loss at step 251: 93.3531723022461
17:10 fiesta INFO : Valid loss at step 251: 96.58218383789062
17:10 fiesta INFO : Best valid loss so far: 96.22515106201172
17:10 fiesta INFO : Learning rate: 0.003
17:10 fiesta INFO : ---
17:10 fiesta INFO : Train loss at step 501: 76.37451171875
17:10 fiesta INFO : Valid loss at step 501: 79.16888427734375
17:10 fiesta INFO : Best valid loss so far: 78.80060577392578
17:10 fiesta INFO : Learning rate: 0.003
17:10 fiesta INFO : ---
17:10 fiesta INFO : Train loss at step 751: 64.38070678710938
17:10 fiesta INFO : Valid loss at step 751: 68.26271057128906
17:10 fiesta INFO : Best valid loss so far: 67.34375762939453
17:10 fiesta INFO : Learning rate: 0.003
17:10 fiesta INFO : ---
17:11 fiesta INFO : Train loss at step 1001: 56.16806411743164
17:11 fiesta INFO : Valid loss at step 1001: 59.840110778808594
17:11 fiesta INFO : Best valid loss so far: 59.25658416748047
17:11 fiesta INFO : Learning rate: 0.003
17:11 fiesta INFO : ---
17:11 fiesta INFO : Train loss at step 1251: 50.14777374267578
17:11 fiesta INFO : Valid loss at step 1251: 54.3632698059082
17:11 fiesta INFO : Best valid loss so far: 53.23629379272461
17:11 fiesta INFO : Learning rate: 0.003
17:11 fiesta INFO : ---
17:11 fiesta INFO : Train loss at step 1501: 44.99894332885742
17:11 fiesta INFO : Valid loss at step 1501: 49.07686996459961
17:11 fiesta INFO : Best valid loss so far: 49.07686996459961
17:11 fiesta INFO : Learning rate: 0.003
17:11 fiesta INFO : ---
17:11 fiesta INFO : Train loss at step 1751: 41.67242431640625
17:11 fiesta INFO : Valid loss at step 1751: 46.59922790527344
17:11 fiesta INFO : Best valid loss so far: 46.25434112548828
17:11 fiesta INFO : Learning rate: 0.003
17:11 fiesta INFO : ---
17:12 fiesta INFO : Train loss at step 2001: 39.023563385009766
17:12 fiesta INFO : Valid loss at step 2001: 44.00343322753906
17:12 fiesta INFO : Best valid loss so far: 43.91112518310547
17:12 fiesta INFO : Learning rate: 0.003
17:12 fiesta INFO : ---
17:12 fiesta INFO : Train loss at step 2251: 36.99894714355469
17:12 fiesta INFO : Valid loss at step 2251: 42.911766052246094
17:12 fiesta INFO : Best valid loss so far: 42.076622009277344
17:12 fiesta INFO : Learning rate: 0.003
17:12 fiesta INFO : ---
17:12 fiesta INFO : Training for 2500 took 151.5548334121704 seconds.
17:12 fiesta INFO : Best validation loss: 40.4495735168457
17:12 fiesta INFO : Loading surrogate Bu2019_CVAE. This surrogate should only be used in the following parameter ranges:
17:12 fiesta INFO : log10_mej_dyn: (-3.0, -1.6989700043360187)
17:12 fiesta INFO : log10_mej_wind: (-2.0, -0.8860566476931632)
17:12 fiesta INFO : phi: (0.0, 90.0)
17:12 fiesta INFO : inclination_EM: (0.0, 1.5707963267948966)
17:12 fiesta INFO : Surrogate Bu2019_CVAE is loading with source-frame time range [ 0.1 15.1] days.
17:12 fiesta INFO : Surrogate Bu2019_CVAE is loading with the following filters: ['besselli', 'bessellux'].
17:12 fiesta INFO : Loaded for surrogate Bu2019_CVAE from ./Bu2019_CVAE.
c) Benchmarking a surrogate#
Now, to check whether the training was actually successful, we need to do some benchmarking of our freshly made surrogates. For this purpose, fiesta has a benchmarker functionality that will use the test data from the training data file.
We first load our trained surrogates:
[10]:
from fiesta.models import FluxSurrogate
FILTERS = ["besselli", "bessellv", "bessellux"]
bu2019_mlp = FluxSurrogate("Bu2019_MLP", filters=FILTERS, directory="./Bu2019_MLP")
bu2019_cvae = FluxSurrogate("Bu2019_CVAE", filters=FILTERS, directory="./Bu2019_CVAE")
17:12 fiesta INFO : Loading surrogate Bu2019_MLP. This surrogate should only be used in the following parameter ranges:
17:12 fiesta INFO : log10_mej_dyn: (-3.0, -1.6989700043360187)
17:12 fiesta INFO : log10_mej_wind: (-2.0, -0.8860566476931632)
17:12 fiesta INFO : phi: (0.0, 90.0)
17:12 fiesta INFO : inclination_EM: (0.0, 1.5707963267948966)
17:12 fiesta INFO : Surrogate Bu2019_MLP is loading with source-frame time range [ 0.1 15.1] days.
17:12 fiesta INFO : Surrogate Bu2019_MLP is loading with the following filters: ['besselli', 'bessellv', 'bessellux'].
17:12 fiesta INFO : Loaded for surrogate Bu2019_MLP from ./Bu2019_MLP.
17:12 fiesta INFO : Loading surrogate Bu2019_CVAE. This surrogate should only be used in the following parameter ranges:
17:12 fiesta INFO : log10_mej_dyn: (-3.0, -1.6989700043360187)
17:12 fiesta INFO : log10_mej_wind: (-2.0, -0.8860566476931632)
17:12 fiesta INFO : phi: (0.0, 90.0)
17:12 fiesta INFO : inclination_EM: (0.0, 1.5707963267948966)
17:12 fiesta INFO : Surrogate Bu2019_CVAE is loading with source-frame time range [ 0.1 15.1] days.
17:12 fiesta INFO : Surrogate Bu2019_CVAE is loading with the following filters: ['besselli', 'bessellv', 'bessellux'].
17:12 fiesta INFO : Loaded for surrogate Bu2019_CVAE from ./Bu2019_CVAE.
[16]:
from fiesta.train import Benchmarker
benchmarker_mlp = Benchmarker(
model=bu2019_mlp,
data=data,
outdir="./Bu2019_MLP",
output_format="png"
)
benchmarker_cvae = Benchmarker(
model=bu2019_cvae,
data=data,
outdir="./Bu2019_CVAE",
output_format="png"
)
17:21 fiesta INFO : Initialized benchmarker for model Bu2019_MLP.
17:21 fiesta INFO : Loaded filters are: ['besselli', 'bessellv', 'bessellux'].
17:21 fiesta INFO : Initialized benchmarker for model Bu2019_CVAE.
17:21 fiesta INFO : Loaded filters are: ['besselli', 'bessellv', 'bessellux'].
Creating the Benchmarker already computes the errors on the test split. Each of its plotting methods saves a PNG into outdir (here ./Bu2019lm_MLP); we call them one at a time below instead of the aggregate benchmark() call so each plot can be inspected separately.
There are two metrics, the benchmarker usually considers: The absolute highest deviation of the predicted light curve from the test data sample across all times and the average mean squared deviation from the test sample integrated across time. Both are provided in the unit mag.
Here, we first plot the worst test light curves according to either metric:
[17]:
# For each filter, plots the single worst-fit test lightcurve (truth vs. surrogate prediction).
benchmarker_mlp.plot_worst_lightcurves()
from IPython.display import Image, display
display(Image(filename="./Bu2019_MLP/worst_lightcurves_highest_lc_error.png"))
display(Image(filename="./Bu2019_MLP/worst_lightcurves_mean_square_lc_error.png"))
Now, we look how the error is distributed across the parameter space. Again, we consider both metrics, the highest absolute deviation and the mean squared error.
We focus on the besselli filter, but the benchmark plot is also created for the two other filters with which we have loaded the model.
Note: This plot does not look so great here, because the test data is discretely distributed. When the training data is spread more randomly across the parameter space, this can be a very useful diagnostic plot.
[18]:
# For each filter, a corner-style scatter of the test points colored by mismatch, plus a histogram of the error metric.
benchmarker_mlp.plot_lightcurves_mismatch()
display(Image(filename=f"./Bu2019_MLP/benchmark_besselli_highest_lc_error.png"))
display(Image(filename=f"./Bu2019_MLP/benchmark_besselli_mean_square_lc_error.png"))
Finally, arguably as most important check, we look how the error distribution evolves across time in each filter. This does not use a specific metric, as the deviation in magnitude is only calculated at specific points in time. These are then plotted here.
[19]:
# For each filter, violin plots of the magnitude error at a set of log-spaced times.
benchmarker_mlp.plot_error_over_time()
display(Image(filename="./Bu2019_MLP/benchmark_error_over_time.png"))
Equivalently, on can just call the .benchmark() method to generate everything all at once:
[20]:
benchmarker_cvae.benchmark()