boss_bench

boss_bench.trajectory

boss_bench.trajectory#

Trajectory benchmark: predict guiding-center trajectories from their first sample.

boss_bench.trajectory.embed_state(s, theta, phi, vpar)[source]#

Return the embedded form of a guiding-center state, in which models take and return states.

\[z = \left(\sqrt{s}\cos\theta,\ \sqrt{s}\sin\theta,\ \cos\varphi,\ \sin\varphi,\ v_\parallel / v_\alpha\right),\]

with \(v_\alpha\) the speed of a 3.5 MeV alpha particle.

Parameters:

  • s (np.ndarray) – normalized toroidal flux, any shape.

  • theta (np.ndarray) – Boozer poloidal angle [rad], the shape of s.

  • phi (np.ndarray) – Boozer toroidal angle [rad], the shape of s.

  • vpar (np.ndarray) – parallel velocity [m/s], the shape of s.

Returns:

the embedded states, shape (..., 5), where ... is the shape of s.

Return type:

np.ndarray

boss_bench.trajectory.load_trajectory(f, config, particle)[source]#

Read one trajectory of one configuration of a trajectory file.

The trajectory is returned as stored: its time t [s], normalized toroidal flux s, Boozer angles theta and zeta [rad], and parallel velocity vpar [m/s], each of shape (n_samples,). The samples are 1 µs apart, starting at \(t = 1\) µs; the trajectory of a lost particle ends with one more sample, at the time it was lost.

Parameters:

  • f (h5py.File) – an open trajectory file.

  • config (int or str) – the row index in the file or the uuid of the configuration.

  • particle (int) – the index of the trajectory among the configuration’s trajectories.

Returns:

t, s, theta, zeta, vpar, the samples of the trajectory.

Return type:

dict

Raises:

  • KeyError – if no configuration in the file has this uuid.

  • IndexError – if the configuration has no trajectory with index particle.

boss_bench.trajectory.to_xyz(record, s, theta, zeta)[source]#

Compute the Cartesian positions of points given in Boozer coordinates, such as the samples of a trajectory.

\(R\), \(Z\) and \(\nu\) are interpolated from the configuration’s Boozer grid, spectrally in the angles and with a cubic spline in \(\sqrt{s}\). With the cylindrical toroidal angle \(\phi = \zeta - \nu\),

\[(x, y, z) = (R \cos\phi,\ R \sin\phi,\ Z).\]
Parameters:

  • record (dict) – the configuration’s fields from data.load_record, including R_boozer, Z_boozer, nu_boozer, s1d and nfp.

  • s (np.ndarray) – normalized toroidal flux, any shape.

  • theta (np.ndarray) – Boozer poloidal angle [rad], the shape of s.

  • zeta (np.ndarray) – Boozer toroidal angle [rad], the shape of s.

Returns:

the Cartesian positions \((x, y, z)\) in meters, shape (..., 3), where ... is the shape of s.

Return type:

np.ndarray

Raises:

ValueError – if an input contains NaN or inf.

boss_bench.trajectory.read_trajectories(f, config)[source]#

Read all trajectories of one configuration of a trajectory file.

The trajectories are put on a common step axis and embedded with embed_state: z[i, k] is the state of trajectory i at step k, that is, \(k\) µs after its first sample. A lost particle stops early, so its rows after its last regular sample are NaN; its final sample, at the time it was lost, falls between steps and is not in z. n_valid is the number of filled rows of each trajectory, and t_loss is the time each particle was lost [s], or NaN if it stays confined. Trajectories with n_valid == 0 were lost before their first sample.

Parameters:

  • f (h5py.File) – an open trajectory file.

  • config (int or str) – the row index in the file or the uuid of the configuration.

Returns:

z, n_valid, t_loss, of shapes (n_trajectories, 1000, 5), (n_trajectories,) and (n_trajectories,).

Return type:

dict

Raises:

KeyError – if no configuration in the file has this uuid.

boss_bench.trajectory.trajectory_metrics(paths, predictions)[source]#

Score trajectory predictions with the metrics of the paper.

For every trajectory, a model predicts the state 10, 100 and 999 µs after its first sample, and the time the particle is lost. Trajectories lost before their first sample are not scored.

Over all predicted configurations, the function reports the nRMSE of the states at each lead time nrmse_10_microsec, nrmse_100_microsec and nrmse_999_microsec (nrmse); the mean absolute error of the loss fraction loss_fraction_mae (loss_fraction_mae); the Wasserstein-1 distance of the confinement times confinement_time_w1_microsec (confinement_time_w1); and the number of configurations scored n_configurations.

Each nRMSE is followed by its breakdown by channel group, nrmse_<k>_microsec_<group>; the nRMSE is the mean of the four groups. The groups are channels of the embedded state (see embed_state): pol_x is \(\sqrt{s}\cos\theta\), pol_y is \(\sqrt{s}\sin\theta\), phi is the pair \((\cos\varphi, \sin\varphi)\), and lambda is \(v_\parallel / v_\alpha\).

As in the paper, a predicted particle is lost the first time its predicted \(s \ge 0.98\), and its predicted state is held at the loss state afterwards.

Parameters:

  • paths (list of str) – the trajectory files holding the predicted configurations, usually data.list_files(data_dir, "trajectory", "test").

  • predictions (dict) – maps the uuid of each configuration to a tuple (z, t_loss), with one row per trajectory in the order of read_trajectories: the predicted embedded states at the three lead times, shape (n_trajectories, 3, 5), and the predicted loss times [s], NaN for particles that stay confined, shape (n_trajectories,). Loss times count from the start of the simulation, as in the data.

Returns:

nrmse_10_microsec, nrmse_100_microsec, nrmse_999_microsec, loss_fraction_mae, confinement_time_w1_microsec, n_configurations, the metrics above and the number of configurations scored, with the nRMSE breakdown after each nRMSE.

Return type:

dict

Raises:

  • KeyError – if a predicted configuration is not in the files.

  • ValueError – if a prediction has the wrong shape, or contains inf or NaN states.

boss_bench.trajectory.nrmse(z_pred, z_true, by_group=False)[source]#

Compute the normalized root-mean-square error of predicted states, averaged over configurations.

For the four channel groups \(g\) of the embedded state, pol_x (\(\sqrt{s}\cos\theta\)), pol_y (\(\sqrt{s}\sin\theta\)), phi (\(\cos\varphi, \sin\varphi\)) and lambda (\(v_\parallel / v_\alpha\)), with frozen training-set RMS \(\sigma_g\),

\[\mathrm{nRMSE}(\tau) = \frac{1}{4} \sum_g \frac{1}{\sigma_g} \frac{1}{D} \sum_d \left( \frac{1}{n_{d,\tau}} \sum_i \left\|\hat z_{i,g}(\tau) - z_{i,g}(\tau)\right\|^2 \right)^{1/2},\]

where \(i\) runs over the \(n_{d,\tau}\) particles of configuration \(d\) that are still confined at \(\tau\), and \(d\) over the \(D\) configurations with any.

Parameters:

  • z_pred (list of np.ndarray) – per configuration, the predicted embedded states at the lead times, shape (n_d, n_tau, 5).

  • z_true (list of np.ndarray) – per configuration, the true states, of the same shapes; NaN where the particle is no longer confined.

  • by_group (bool) – if True, return the nRMSE of each channel group instead of their mean.

Returns:

the nRMSE at each lead time, shape (n_tau,). With by_group, a dict that maps each group, pol_x, pol_y, phi, lambda, to its term \(\frac{1}{\sigma_g D} \sum_d (\dots)^{1/2}\), shape (n_tau,); their mean is the nRMSE.

Return type:

np.ndarray or dict

Raises:

ValueError – if a predicted state contains NaN or inf.

boss_bench.trajectory.loss_fraction_mae(t_loss_pred, t_loss_true)[source]#

Compute the mean absolute error of the fraction of particles lost within 1 ms, averaged over configurations.

With \(f_d\) the fraction of the particles of configuration \(d\) lost by \(t = 1\) ms,

\[\text{MAE} = \frac{1}{D} \sum_d \left|\hat f_d - f_d\right|.\]
Parameters:

  • t_loss_pred (list of np.ndarray) – per configuration, the predicted loss times [s], NaN if confined, shape (n_d,).

  • t_loss_true (list of np.ndarray) – per configuration, the true loss times [s], of the same shapes.

Returns:

the MAE.

Return type:

float

Raises:

ValueError – if a loss time is infinite.

boss_bench.trajectory.confinement_time_w1(t_loss_pred, t_loss_true)[source]#

Compute the Wasserstein-1 distance between predicted and true confinement times, averaged over configurations.

With \(T\) the confinement time of a particle, its loss time \(t_\text{loss}\) if it is lost within 1 ms and 1 ms otherwise, and \(T_{d,(i)}\) the \(i\)-th smallest of the \(n_d\) confinement times of configuration \(d\),

\[W_1 = \frac{1}{D} \sum_d \frac{1}{n_d} \sum_i \left|\hat T_{d,(i)} - T_{d,(i)}\right|.\]
Parameters:

  • t_loss_pred (list of np.ndarray) – per configuration, the predicted loss times [s], NaN if confined, shape (n_d,).

  • t_loss_true (list of np.ndarray) – per configuration, the true loss times [s], of the same shapes.

Returns:

\(W_1\) [µs].

Return type:

float

Raises:

ValueError – if a loss time is infinite.