Usage#
We’ve designed the package so it is easy to write scripts that run MOBO for a particular JETTO problem. To do so, you need to:
Define your inputs
Define your objective functions
Write the main evaluation script
In this section, we run through the ecrh_q_optimisation example from theo-brown/jetto-mobo; this formed the basis of our SOFE2023 poster, looking at MOBO of the ECRH input to find good q-profiles.
1. Define inputs#
First, we define the ECRH parameterisation to use.
As ECRH is a plasma profile, we decorate it with @jetto_mobo.inputs.plasma_profile, which tells the optimiser to expect a function of the form f(xrho, parameters).
1from typing import Union
2
3import numpy as np
4from scipy.interpolate import CubicSpline
5from scipy.special import comb
6
7from jetto_mobo.inputs import plasma_profile
8
9
10@plasma_profile
11def marsden_piecewise_linear(xrho: np.ndarray, parameters: np.ndarray) -> np.ndarray:
12 """
13 12-parameter piecewise linear profile, originally designed for ECRH profiles by Stephen Marsden.
14
15 Parameters
16 ----------
17 xrho : np.ndarray
18 Normalised radial position.
19 parameters : np.ndarray
20 Profile parameters:
21 - parameters[0]: on-axis peak power
22 - parameters[1]: on-axis descent end power
23 - parameters[2]: on-axis descent end xrho
24 - parameters[3]: minimum power fraction
25 - parameters[4]: minimum power
26 - parameters[5]: minimum xrho
27 - parameters[6]: off-axis shaper 1 fraction
28 - parameters[7]: off-axis shaper 2 fraction
29 - parameters[8]: off-axis peak power
30 - parameters[9]: off-axis peak xrho
31 - parameters[10]: turn-off shaper fraction
32 - parameters[11]: turn-off xrho
33
34 Raises
35 ------
36 ValueError
37 If the number of parameters is not 12.
38 """
39 if len(parameters) != 12:
40 raise ValueError(f"Expected 12 parameters, got {len(parameters)}.")
41
42 lower_bounds = np.array([0, 0.05, 0.01, 0.1, 0, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1])
43 upper_bounds = np.array([1, 1, 0.09, 1, 1, 0.29, 0.9, 0.9, 1, 0.75, 0.9, 0.45])
44 is_outside_bounds = (parameters < lower_bounds) | (parameters > upper_bounds)
45 if np.any(is_outside_bounds):
46 raise ValueError(
47 f"Parameter(s) outside of bounds at indices {np.nonzero(is_outside_bounds)}"
48 )
49
50 on_axis_peak_power = parameters[0]
51 on_axis_descent_end_power = parameters[1] * on_axis_peak_power
52 minimum_power = parameters[4]
53 minimum_shaper_power = (
54 parameters[3] * minimum_power + (1 - parameters[3]) * on_axis_descent_end_power
55 )
56 off_axis_peak_power = parameters[8]
57 off_axis_shaper_2_power = (
58 parameters[7] * off_axis_peak_power + (1 - parameters[7]) * minimum_power
59 )
60 off_axis_shaper_1_power = (
61 parameters[6] * off_axis_shaper_2_power + (1 - parameters[6]) * minimum_power
62 )
63 turn_off_shaper_power = parameters[10] * off_axis_peak_power
64 turn_off_power = 0
65 on_axis_peak_xrho = 0
66 on_axis_descent_end_xrho = parameters[2]
67 minimum_xrho = parameters[5]
68 minimum_shaper_xrho = (minimum_xrho + on_axis_descent_end_xrho) / 2
69 off_axis_peak_xrho = minimum_xrho + parameters[9]
70 turn_off_xrho = off_axis_peak_xrho + parameters[11]
71 off_axis_shaper_1_xrho = 2 / 3 * minimum_xrho + 1 / 3 * off_axis_peak_xrho
72 off_axis_shaper_2_xrho = 1 / 3 * minimum_xrho + 2 / 3 * off_axis_peak_xrho
73 turn_off_shaper_xrho = 1 / 2 * off_axis_peak_xrho + 1 / 2 * turn_off_xrho
74
75 xdata = [
76 on_axis_peak_xrho,
77 on_axis_descent_end_xrho,
78 minimum_shaper_xrho,
79 minimum_xrho,
80 off_axis_shaper_1_xrho,
81 off_axis_shaper_2_xrho,
82 off_axis_peak_xrho,
83 turn_off_shaper_xrho,
84 turn_off_xrho,
85 ]
86 ydata = [
87 on_axis_peak_power,
88 on_axis_descent_end_power,
89 minimum_shaper_power,
90 minimum_power,
91 off_axis_shaper_1_power,
92 off_axis_shaper_2_power,
93 off_axis_peak_power,
94 turn_off_shaper_power,
95 turn_off_power,
96 ]
97
98 # Linearly interpolate between each of the points.
99 qece = np.zeros(len(xrho))
100 for i_point in range(len(xdata) - 1):
101 # Get just the pair of points.
102 points_xrho = [xdata[i_point], xdata[i_point + 1]]
103 points_power = [ydata[i_point], ydata[i_point + 1]]
104 # Fit with a straight line between them.
105 cs = CubicSpline(
106 points_xrho, points_power, bc_type=[(2, 0), (2, 0)], extrapolate=False
107 )
108 # Set the values in the profile.
109 qece = np.array([v if v > 0 else qece[i] for i, v in enumerate(cs(xrho))])
110
We also need to define the bounds on the parameter values.
114marsden_piecewise_linear_bounds = np.array(
115 [
116 [0, 0.05, 0.01, 0.1, 0, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1],
117 [1, 1, 0.09, 1, 1, 0.29, 0.9, 0.9, 1, 0.75, 0.9, 0.45],
118 ]
119)
120
121
122def bernstein(i: Union[int, np.ndarray], n: int, t: np.ndarray) -> np.ndarray:
123 """
124 Evaluate the Bernstein basis polynomials at t.
125
126 $$
127 b_{i,n}(t) = \sum_{i=0}^n \binom{n}{i} t^i (1-t)^{n-i}
128 $$
129 """
130 return comb(n, i) * np.power.outer(t, i) * np.power.outer((1 - t), (n - i))
131
132
133def bezier_parametric(t: np.ndarray, control_points: np.ndarray) -> np.ndarray:
134 """
135 Parameters
136 ----------
137 t : np.ndarray
138 Array of parametric coordinate values, length M.
139 control_points : np.ndarray
140 Array of shape (N, 2), giving N control points in the form [[x1, y1], ..., [xN, yN]].
141
142 Returns
143 -------
144 np.ndarray
145 Bezier curve defined by the control points, evaluated at t. Shape (M, 2), where [:, 0] is x and [:, 1] y coordinates.
146 """
147 # Check control points
148 if not len(control_points.shape) == 2 or not control_points.shape[1] == 2:
149 raise ValueError(
150 f"control_points must be an array of shape (N, 2) (got {control_points.shape})."
151 )
152
153 n = control_points.shape[0]
154
155 # Evaluate basis functions
156 i = np.arange(n)
157 basis_functions = bernstein(i, n - 1, t)
158
159 # Evaluate bezier curve
160 return basis_functions @ control_points
161
162
163def bezier_x(
164 x: np.ndarray, control_points: np.ndarray, parametric_resolution: int = int(1e3)
165) -> np.ndarray:
166 """
167 Parameters
168 ----------
169 x : np.ndarray
170 Array of x coordinates to evaluate the bezier curve at.
171 control_points : np.ndarray
172 Array of shape (N, 2), giving N control points in the form [[x1, y1], ..., [xN, yN]].
173 parametric_resolution : int, default int(1e3)
174 Number of points to evaluate the bezier curve at to obtain the implicit/parametric coordinates.
175
176 Returns
177 -------
178 np.ndarray
179 Array of y coordinates, evaluated at x.
180 """
181 t = np.linspace(0, 1, parametric_resolution)
182 b = bezier_parametric(t, control_points)
183 implicit_points_x = b[:, 0]
184 implicit_points_y = b[:, 1]
185
186 # Interpolate into x coordinates
187 t_x = np.interp(
188 x, implicit_points_x, t
189 ) # t_x is the implicit/parametric value corresponding to x
190 b = bezier_parametric(t_x, control_points)
191 x_ = b[:, 0]
192 y = b[:, 1]
193
194 return y
195
196
197on_axis_smoothness = 0.05
198off_axis_smoothness = 0.05
199
200
201@plasma_profile
202def constrained_bezier_profile(xrho: np.ndarray, parameters: np.ndarray) -> np.ndarray:
203 """
204 Bezier curve evaluated at x, with added hyperparameters to control the behaviour at the ends of the curve.
205
206 Parameters are of the form [y0, x1, y1, ..., xn].
207 Control points are [[0, y0], [on_axis_smoothness, y_0], [x1, y1], ..., [xn - off_axis_smoothness, 0], [xn, 0]].
208 """
209 y0 = parameters[0]
210 x_parameters = parameters[1:-1:2]
211 y_parameters = parameters[2:-1:2]
212 xn = parameters[-1]
213
214 control_points_x = np.concatenate(
215 [
216 [0, on_axis_smoothness],
217 # Sort the x coordinates of the control points so that they are monotonically increasing in x.
218 # np.sort(x_parameters),
219 # OR
220 # Product the x coordinates of the control points so that they are monotonically increasing in x.
221 # x coordinates will be [x0*x1*...*xn, x1*...*xn, ..., xn]
222 # np.cumprod(x_parameters[::-1])[::-1],
223 # OR
224 # Use the x coordinates of the control points as they are.
225 x_parameters * xn,
226 [xn - off_axis_smoothness, xn],
227 ]
228 )
229 control_points_y = np.concatenate([[y0, y0], y_parameters, [0, 0]])
230 if not control_points_x.shape == control_points_y.shape:
231 raise ValueError(
232 f"control_points_x and control_points_y must have the same shape (got {control_points_x.shape} and {control_points_y.shape})."
233 )
234 control_points = np.array([control_points_x, control_points_y]).T
235 return bezier_x(xrho, control_points, parametric_resolution=int(1e3))
236
237
238@plasma_profile
239def bezier_profile(xrho: np.ndarray, parameters: np.ndarray) -> np.ndarray:
240 """
241 Bezier curve evaluated at x.
242
243 Parameters are of the form [y0, x1, y1, ..., xn].
244 Control points are [[0, y0], [x1, y1], ..., [xn, 0]].
245 """
246 y0 = parameters[0]
247 x_parameters = parameters[1:-1:2]
248 y_parameters = parameters[2:-1:2]
249 xn = parameters[-1]
250
251 control_points_x = np.concatenate([[0], x_parameters * xn, [xn]])
252 control_points_y = np.concatenate([[y0], y_parameters, [0]])
253 if not control_points_x.shape == control_points_y.shape:
254 raise ValueError(
255 f"control_points_x and control_points_y must have the same shape (got {control_points_x.shape} and {control_points_y.shape})."
256 )
257 control_points = np.array([control_points_x, control_points_y]).T
258 return bezier_x(xrho, control_points, parametric_resolution=int(1e3))
1. Define objective functions#
Next, we define the objective functions. For our q-profile optimisation problem, we want to do multi-objective optimisation. Each of the objectives can be computed from the JETTO profiles and timetraces datasets:
1from typing import Union
2
3import numpy as np
4from jetto_tools.results import JettoResults
5from netCDF4 import Dataset
6
7from jetto_mobo.objectives import objective
8
9
10def soft_hat(
11 x: Union[float, np.ndarray],
12 x_lower: float = 0,
13 y_lower: float = 1e-3,
14 x_plateau_start: float = 0,
15 x_plateau_end: float = 0,
16 x_upper: float = 0,
17 y_upper: float = 1e-3,
18) -> np.ndarray:
19 """
20 Smooth tophat function.
21
22 Passes through (x_lower, y_lower), (x_plateau_start, 1), (x_plateau_end, 1), (x_upper, y_upper).
23 Squared exponential decay from 0 to x_plateau_start and from x_plateau_end to infinity, with rate of decay such that y=y_lower at x=x_lower and y=y_upper at x=x_upper.
24
25 Parameters
26 ----------
27 x : Union[float, np.ndarray]
28 Input value
29 x_lower : float, optional
30 x-value at which y=y_lower (default: 0)
31 y_lower : float, optional
32 y-value at x=x_lower (default: 1e-3)
33 x_plateau_start : float, optional
34 x-value at which the plateau starts (default: 0)
35 x_plateau_end : float, optional
36 x-value at which the plateau ends (default: 0)
37 x_upper : float, optional
38 x-value at which y=y_upper (default: 0)
39 y_upper : float, optional
40 y-value at x=x_upper (default: 1e-3)
41
42 Returns
43 -------
44 np.ndarray
45 Smooth objective value
46 """
47 k_lower = -np.log(y_lower) / np.power(x_lower - x_plateau_start, 2)
48 k_upper = -np.log(y_upper) / np.power(x_upper - x_plateau_end, 2)
49 return np.piecewise(
50 x,
51 [
52 x < x_plateau_start,
53 (x >= x_plateau_start) & (x <= x_plateau_end),
54 x > x_plateau_end,
55 ],
56 [
57 lambda x: np.exp(-k_lower * np.power(x - x_plateau_start, 2)),
58 1,
59 lambda x: np.exp(-k_upper * np.power(x - x_plateau_end, 2)),
60 ],
We can then use @jetto_mobo.objectives.objective to decorate a function that takes a JettoResults object and returns the vector of objective values:
62
63
64def softmax(x: np.ndarray) -> np.ndarray:
65 exp = np.exp(x)
66 return exp / np.sum(exp)
67
68
69def q0_close_to_qmin(profiles: Dataset, timetraces: Dataset) -> np.ndarray:
70 """1 if q0 = qmin, decaying to 0.5 at ||q0 - qmin|| = 2"""
71 distance = np.abs(timetraces["Q0"][-1].data - timetraces["QMIN"][-1].data)
72 return soft_hat(
73 distance,
74 x_lower=-1, # Not used, as 0 < x < 1
75 y_lower=1e-3, # Not used, as 0 < x < 1
76 x_plateau_start=0,
77 x_plateau_end=0,
78 x_upper=2,
79 y_upper=0.5,
80 )
81
82
83def qmin_close_to_centre(profiles: Dataset, timetraces: Dataset) -> np.ndarray:
84 """1 if argmin(q) = 0, decaying to 1e-3 at ||argmin(q)|| = 1"""
85 return soft_hat(
86 timetraces["ROQM"][-1].data,
87 x_lower=-1, # Not used, as 0 < x < 1
88 y_lower=1e-3, # Not used, as 0 < x < 1
89 x_plateau_start=0,
90 x_plateau_end=0,
91 x_upper=1,
92 y_upper=1e-3,
93 )
94
95
96def qmin_in_safe_region(profiles: Dataset, timetraces: Dataset) -> float:
97 """1 if qmin between 2.2 and 2.5, decaying to 0.5 at 2 and 3"""
If instead we wanted to do single-objective optimisation, we can use jetto_mobo.objectives.objective(weights=True) to decorate a scalar weighted version of the vector objective function:
100 x_lower=2.2,
101 y_lower=0.5,
102 x_plateau_start=2.2,
103 x_plateau_end=2.5,
104 x_upper=3,
105 y_upper=0.5,
106 )
107
108
109def q_increasing(profiles: Dataset, timetraces: Dataset) -> np.ndarray:
110 """1 if q is increasing at every radial point, decaying to 1e-3 if q is non-increasing at every radial point"""
111 is_increasing = np.gradient(profiles["Q"][-1].data) > 0
112 return soft_hat(
113 np.mean(is_increasing), # Fraction of curve where q is increasing
114 x_lower=0,
115 y_lower=1e-3,
116 x_plateau_start=1,
117 x_plateau_end=1,
118 x_upper=2, # Not used, as 0 < x < 1
119 y_upper=1e-3, # Not used, as 0 < x < 1
120 )
121
122
123def maximise_radius_at_which_q_is_value(
124 profiles: Dataset, timetraces: Dataset, value: float
125) -> np.ndarray:
126 """1 if q=value at r>0.8 decaying to 0.5 if q=value at r=0.5
127
128 Note that this excludes points where q=value but r < argmin(q)."""
129 xrho = profiles["XRHO"][-1].data
130 condition_1 = profiles["Q"][-1].data >= value
131 condition_2 = xrho >= timetraces["ROQM"][-1].data
132 i = np.where(condition_1 & condition_2)[0][0]
133 radius_of_q_is_value = xrho[i]
134 return soft_hat(
135 radius_of_q_is_value,
136 x_lower=0.5,
137 y_lower=0.5,
138 x_plateau_start=0.8,
139 x_plateau_end=1,
140 x_upper=2, # Not used, as 0 < x < 1
141 y_upper=2, # Not used, as 0 < x < 1
142 )
143
144
145@objective
146def q_vector_objective(results: JettoResults) -> np.ndarray:
147 """
148 Vector of 6 objective functions relating to the shape of the q-profile.
149
150 Returns
151 -------
152 np.ndarray
153 Vector of objective values:
154 - Reduced 'height' of reversed shear at axis (q0 close to qmin)
155 - Reduced 'width' of reversed shear at axis (qmin close to r=0)
156 - Monotonic q (q increasing at every radial point)
157 - qmin in safe region (2 < qmin < 3)
158 - Maximise radius at which q=3
159 - Maximise radius at which q=4
160 """
161 profiles = results.load_profiles()
162 timetraces = results.load_timetraces()
163
164 return q_vector_objective_from_cdf(profiles, timetraces)
165
166
167def q_vector_objective_from_cdf(profiles: Dataset, timetraces: Dataset) -> np.ndarray:
168 return np.array(
169 [
170 q0_close_to_qmin(profiles, timetraces),
171 qmin_close_to_centre(profiles, timetraces),
172 q_increasing(profiles, timetraces),
173 qmin_in_safe_region(profiles, timetraces),
174 maximise_radius_at_which_q_is_value(profiles, timetraces, 3),
175 maximise_radius_at_which_q_is_value(profiles, timetraces, 4),
176 ]
177 )
178
179
180def q_constraints(results: JettoResults) -> np.ndarray:
181 """Vector of constraints on the q profile.
182
183 Constraint functions are of the form ``g(x) <= 0`` (i.e. negative if the constraint is satisfied).
184
185 Returns
186 -------
187 np.ndarray
188 Vector of constraint values:
189 - qmin > 2
190 - qmin < 3
191 """
192 return q_constraints_from_cdf(results.load_profiles(), results.load_timetraces())
193
194
195def q_constraints_from_cdf(profiles: Dataset, timetraces: Dataset) -> np.ndarray:
196 return np.array(
197 [
198 2 - timetraces["QMIN"][-1].data,
199 timetraces["QMIN"][-1].data - 3,
200 ]
201 )
202
203
204@objective(weights=True)
205def q_scalar_objective(results: JettoResults, weights: np.ndarray) -> np.ndarray:
206 """
207 Weighted sum of q_vector_objective.
208 """
209 v = q_vector_objective(results)
210 return v @ weights
211
212
213def q_scalar_objective_from_cdf(
214 profiles: Dataset, timetraces: Dataset, weights: np.ndarray
215) -> np.ndarray:
216 v = q_vector_objective_from_cdf(profiles, timetraces)
217 return v @ weights
3. Write the main evaluation script#
As the evaluation of the objective functions depends on the particular problem at hand, we haven’t yet implemented a general framework to run the MOBO loop. (If you think it would be useful, do get in touch!)
Consequently, you’ll have to write the evaluation by hand, using our pre-built wrappers.
3.1 Evaluation helper function#
We define a helper function that takes a set of parameters, creates a JETTO config, sends the config off to be run, and returns the relevant values (input, output, objective) on completion.
3.2 Data storage#
We also define a helper function to save our results to a file. Use your team’s preferred data format!
3.3 Argument parsing#
Because we want it to be easy to use our final script for multiple different runs, we use argparse to parse the command line arguments:
1import argparse
2import asyncio
3import logging
4import sys
5from os import sched_getaffinity
6from pathlib import Path
7from typing import Tuple
8
9import h5py
10import jetto_tools
11import numpy as np
12import torch
13from ecrh_inputs import (
14 bezier_profile,
15 constrained_bezier_profile,
16 marsden_piecewise_linear,
17 marsden_piecewise_linear_bounds,
18)
19from q_objectives import q_vector_objective
20from utils import write_to_file
21
22from jetto_mobo import acquisition, simulation, surrogate, utils
23
24logger = logging.getLogger(__name__)
25logger.setLevel(logging.DEBUG)
26logger.addHandler(logging.StreamHandler(sys.stdout))
27
3.4 Initialisation#
Important
Parameter bounds must be a tensor! If you initialised them as a numpy array, cast them to a tensor before continuing. We do this with:
30parser.add_argument(
Before starting, we need to generate some initial candidates.
32)
33parser.add_argument("--jetto_timelimit", type=int, default=10400)
34parser.add_argument("--jetto_fail_value", type=float, default=0)
35parser.add_argument("--discard_failures", action="store_true")
36parser.add_argument("--sobol_only", action="store_true")
37parser.add_argument("--n_xrho_points", type=int, default=150)
38parser.add_argument("--batch_size", type=int, default=10)
39parser.add_argument("--initial_batch_size", type=int, default=30)
40parser.add_argument("--n_iterations", type=int, default=16)
41parser.add_argument("--reference_values", type=float, nargs="+", default=None)
42parser.add_argument(
43 "--device",
44 type=str,
45 default=torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu"),
46)
47parser.add_argument("--dtype", type=torch.dtype, default=torch.float64)
48parser.add_argument(
49 "--output_dir", type=Path, default=Path("data/piecewise_linear_mobo")
50)
51parser.add_argument(
52 "--parameterisation",
53 choices=["piecewise_linear", "bezier", "bezier2"],
54 default="piecewise_linear",
55)
56parser.add_argument("--n_parameters", type=int, default=12)
57parser.add_argument("--alpha", type=float, default=0.0)
58parser.add_argument("--seed", type=int, default=0)
59parser.add_argument("--resume", action="store_true")
60args = parser.parse_args()
61
62torch.manual_seed(args.seed)
63
64# Objectives
65n_objectives = 6
66
67# Input parameterisation
68if args.parameterisation == "piecewise_linear":
69 ecrh_function = marsden_piecewise_linear
70 parameter_bounds = torch.tensor(marsden_piecewise_linear_bounds)
71elif args.parameterisation == "bezier":
72 ecrh_function = constrained_bezier_profile
73 parameter_bounds = torch.tensor([[0] * args.n_parameters, [1] * args.n_parameters])
74elif args.parameterisation == "bezier2":
75 ecrh_function = bezier_profile
76 parameter_bounds = torch.tensor([[0] * args.n_parameters, [1] * args.n_parameters])
77
78
79# Reference values
80if args.reference_values is None:
81 reference_values = torch.tensor([0.0] * n_objectives)
3.5 Main loop#
Now we can bring it all together in the main loop.
83 reference_values = torch.tensor(args.reference_values)
84else:
85 raise ValueError(
86 f"Length of reference values ({len(args.reference_values)}) does not match number of objectives ({n_objectives})."
87 )
88
89# Failure behaviour
90if not args.discard_failures:
91 failure_objective_value = args.jetto_fail_value
92
93
94def evaluate(
95 ecrh_parameters_batch: np.ndarray,
96 batch_directory: Path,
97) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
98 """
99 Evaluate a batch of ECRH parameters.
100
101 Runs JETTO in parallel for each set of ECRH parameters, and returns the converged ECRH profiles and objective values.
102
103 Parameters
104 ----------
105 ecrh_parameters_batch: np.ndarray
106 batch_directory: Path
107
108 Returns
109 -------
110 Tuple[np.ndarray, np.ndarray, np.ndarray]
111 """
112 configs = {}
113
114 for i, ecrh_parameters in enumerate(ecrh_parameters_batch):
115 # Initialise config object
116 config_directory = Path(f"{batch_directory}/candidate_{i}")
117 config_object = simulation.create_config(
118 template=args.jetto_template, directory=config_directory
119 )
120
121 # Set the ECRH function
122 exfile = jetto_tools.binary.read_binary_file(config_object.exfile)
123 exfile["QECE"][0] = ecrh_function(
124 xrho=exfile["XRHO"][0], parameters=ecrh_parameters
125 )
126 jetto_tools.binary.write_binary_exfile(exfile, config_object.exfile)
127
128 # Store in dict
129 # Currently this is necessary as the JettoTools RunConfig does not store the directory path
130 configs[config_object] = config_directory
131
132 # Run asynchronously in parallel
133 batch_output = asyncio.run(
134 simulation.run_many(
135 jetto_image=args.jetto_image,
136 run_configs=configs,
137 timelimit=args.jetto_timelimit,
138 )
139 )
140
141 # Parse outputs
142 converged_ecrh = []
143 converged_q = []
144 objective_values = []
145 for results in batch_output:
146 if results is not None:
147 try:
148 profiles = results.load_profiles()
149 except:
150 logger.warning("Failed to load profiles. Maybe the file was corrupted?")
151 converged_ecrh.append(np.full(args.n_xrho_points, np.nan))
152 converged_q.append(np.full(args.n_xrho_points, np.nan))
153 objective_values.append(np.full(n_objectives, np.nan))
154 else:
155 converged_ecrh.append(profiles["QECE"][-1])
156 converged_q.append(profiles["Q"][-1])
157 objective_values.append(q_vector_objective(results))
158 else:
159 logger.warning("JETTO failed to converge.")
160 converged_ecrh.append(np.full(args.n_xrho_points, np.nan))
161 converged_q.append(np.full(args.n_xrho_points, np.nan))
162 objective_values.append(np.full(n_objectives, np.nan))
163
164 # Compress outputs
165 for _, config_directory in configs.items():
166 simulation.compress_jetto_dir(config_directory, delete=True)
167
168 return (
169 np.array(converged_ecrh),
170 np.array(converged_q),
171 np.array(objective_values),
172 )
173
174
175if args.resume:
176 # TODO: Update to permit batched initialisation data
177 logger.info("Resuming from file...")
178 with h5py.File(args.output_dir / "results.h5", "r") as f:
179 # Read initialisation data
180 ecrh_parameters = [torch.tensor(f["initialisation/ecrh_parameters"][:])]
181 objective_values = [torch.tensor(f["initialisation/objective_values"][:])]
182 # Read any additional optimisation steps
183 completed_optimisation_steps = len(f.keys()) - 1
184 if completed_optimisation_steps > 0:
185 for i in range(1, completed_optimisation_steps + 1):
186 logger.info(f"Loading data from optimisation step {i}...")
187 ecrh_parameters.append(
188 torch.tensor(f[f"optimisation_step_{i}/ecrh_parameters"][:])
189 )
190 objective_values.append(
191 torch.tensor(f[f"optimisation_step_{i}/objective_values"][:])
192 )
193
194 # Concatenate
195 ecrh_parameters = torch.cat(ecrh_parameters).to(
196 device=args.device, dtype=args.dtype
197 )
198 objective_values = torch.cat(objective_values).to(
199 device=args.device, dtype=args.dtype
200 )
201
202else:
203 # Generate initial data
204 logger.info("Generating initial data...")
205
206 # Get number of cores
207 n_cores = len(sched_getaffinity(0))
208 n_batches = args.initial_batch_size // n_cores
209
210 # Evaluate initial candidates in batches of size n_cores
211 for batch_index in range(n_batches):
212 # Sobol sampling for initial candidates
213 ecrh_parameters = acquisition.generate_initial_candidates(
214 bounds=parameter_bounds,
215 n=n_cores,
216 device=args.device,
217 dtype=args.dtype,
218 )
219
220 logger.info(f"Evaluating initial candidates (batch {batch_index})...")
221 # Save initial candidates to file
222 write_to_file(
223 output_file=args.output_dir / "results.h5",
224 root_label=f"initialisation/batch_{batch_index}",
225 ecrh_parameters_batch=ecrh_parameters.detach().cpu().numpy(),
226 preconverged_ecrh=np.array(
227 [
228 ecrh_function(
229 xrho=np.linspace(0, 1, args.n_xrho_points), parameters=p
230 )
231 for p in ecrh_parameters.detach().cpu().numpy()
232 ]
233 ),
234 )
235 (
236 converged_ecrh,
237 converged_q,
238 objective_values,
239 ) = evaluate(
240 ecrh_parameters_batch=ecrh_parameters.detach().cpu().numpy(),
241 batch_directory=args.output_dir / f"0_initialisation_batch_{batch_index}",
242 )
243 # Save evaluated results to file
244 write_to_file(
245 output_file=args.output_dir / "results.h5",
246 root_label=f"initialisation/batch_{batch_index}",
247 converged_ecrh=converged_ecrh,
248 converged_q=converged_q,
249 objective_values=objective_values,
250 )
251
252 # Rearrange the file structure
253 with h5py.File(args.output_dir / "results.h5", "a") as h5file:
254 n_evaluations = n_cores * n_batches
255 # Create new datasets
256 h5file.create_dataset(
257 "initialisation/converged_ecrh",
258 (n_evaluations, args.n_xrho_points),
259 )
260 h5file.create_dataset(
261 "initialisation/converged_q", (n_evaluations, args.n_xrho_points)
262 )
263 h5file.create_dataset(
264 "initialisation/ecrh_parameters",
265 (n_evaluations, args.n_parameters),
266 )
267 h5file.create_dataset(
268 "initialisation/objective_values",
269 (n_evaluations, n_objectives),
270 )
271 h5file.create_dataset(
272 "initialisation/preconverged_ecrh",
273 (n_evaluations, args.n_xrho_points),
274 )
275
276 # Copy data across
277 for batch_index in range(n_batches):
278 lower_index = batch_index * n_cores
279 upper_index = lower_index + n_cores
280 h5file["initialisation/converged_ecrh"][lower_index:upper_index] = h5file[
281 f"initialisation/batch_{batch_index}/converged_ecrh"
282 ]
283 h5file["initialisation/converged_q"][lower_index:upper_index] = h5file[
284 f"initialisation/batch_{batch_index}/converged_q"
285 ]
286 h5file["initialisation/ecrh_parameters"][lower_index:upper_index] = h5file[
287 f"initialisation/batch_{batch_index}/ecrh_parameters"
288 ]
289 h5file["initialisation/objective_values"][lower_index:upper_index] = h5file[
290 f"initialisation/batch_{batch_index}/objective_values"
291 ]
292 h5file["initialisation/preconverged_ecrh"][
293 lower_index:upper_index
294 ] = h5file[f"initialisation/batch_{batch_index}/preconverged_ecrh"]
295
296 del h5file[f"initialisation/batch_{batch_index}"]
297
298 # Load back into a tensor
299 ecrh_parameters = torch.tensor(
300 h5file["initialisation/ecrh_parameters"][:],
301 device=args.device,
302 dtype=args.dtype,
303 )
304 objective_values = torch.tensor(
305 h5file["initialisation/objective_values"][:],
306 device=args.device,
307 dtype=args.dtype,
308 )
309
310 # Compute and save the HVI
311 log_hypervolume = utils.compute_pareto_loghypervolume(
312 objective_values=objective_values,
313 reference_point=reference_values,
314 )
315 h5file["initialisation"].attrs["log_hypervolume"] = log_hypervolume
316
317 # Initialise optimisation step
318 completed_optimisation_steps = 0
319
320
321for optimisation_step in range(
322 completed_optimisation_steps + 1,
323 completed_optimisation_steps + 1 + args.n_iterations,
324):
325 logger.info(f"Optimisation step {optimisation_step}")
326
327 # Drop any NaNs
328 logger.info("Handling NaNs...")
329 if args.discard_failures:
330 mask = ~torch.isnan(objective_values).any(dim=1)
331 ecrh_parameters = ecrh_parameters[mask]
332 objective_values = objective_values[mask]
333 else:
334 # Replace NaNs with failure values
335 objective_values[torch.isnan(objective_values)] = failure_objective_value
336
337 # Generate trial candidates
338 if args.sobol_only or objective_values.nelement() == 0:
339 # Use quasirandom Sobol sampling to generate trial candidates
340 logger.info("Generating trial candidates using Sobol sampling...")
341 new_ecrh_parameters = acquisition.generate_initial_candidates(
342 bounds=parameter_bounds,
343 n=args.batch_size,
344 device=args.device,
345 dtype=args.dtype,
346 )
347 else:
348 logger.info("Fitting surrogate model...")
349 model = surrogate.fit_surrogate_model(
350 inputs=ecrh_parameters,
351 input_bounds=parameter_bounds,
352 objective_values=objective_values,
353 device=args.device,
354 dtype=args.dtype,
355 )
356
357 # Use qNEHVI to generate trial candidates
358 logger.info("Generating trial candidates using qNEHVI...")
359 new_ecrh_parameters = acquisition.generate_trial_candidates(
360 observed_inputs=ecrh_parameters,
361 bounds=parameter_bounds,
362 model=model,
363 acquisition_function=acquisition.qNoisyExpectedHypervolumeImprovement,
364 n_constraints=0,
365 device=args.device,
366 dtype=args.dtype,
367 batch_size=args.batch_size,
368 mode="sequential" if args.batch_size > 5 else "joint",
369 acqf_kwargs={
370 "ref_point": reference_values,
371 "alpha": args.alpha,
372 "prune_baseline": True,
373 },
374 )
375 write_to_file(
376 output_file=args.output_dir / "results.h5",
377 root_label=f"optimisation_step_{optimisation_step}",
378 ecrh_parameters_batch=new_ecrh_parameters.detach().cpu().numpy(),
379 preconverged_ecrh=np.array(
380 [
381 ecrh_function(xrho=np.linspace(0, 1, args.n_xrho_points), parameters=p)
382 for p in new_ecrh_parameters.detach().cpu().numpy()
383 ]
384 ),
385 )
386
387 # Evaluate candidates
388 logger.info("Evaluating trial candidates...")
389 (
390 converged_ecrh,
391 converged_q,
392 new_objective_values,
393 ) = evaluate(
394 ecrh_parameters_batch=new_ecrh_parameters.detach().cpu().numpy(),
395 batch_directory=args.output_dir / str(optimisation_step),
396 )
397 write_to_file(
398 output_file=args.output_dir / "results.h5",
399 root_label=f"optimisation_step_{optimisation_step}",
400 converged_ecrh=converged_ecrh,
401 converged_q=converged_q,
402 objective_values=new_objective_values,
403 )
404
405 # Concatenate new data
406 ecrh_parameters = torch.cat([ecrh_parameters, new_ecrh_parameters])
407 # Have to convert new_objective_values to tensor, because it is a np.ndarray output from reading JettoResults
408 objective_values = torch.cat(
409 [
410 objective_values,
411 torch.tensor(
412 new_objective_values,
413 device=args.device,
414 dtype=args.dtype,
415 ),
416 ]
417 )
418
419 # Compute and save the HVI
420 log_hypervolume = utils.compute_pareto_loghypervolume(
421 objective_values=objective_values,
422 reference_point=reference_values,
423 )
424 write_to_file(
425 output_file=args.output_dir / "results.h5",
426 root_label=f"optimisation_step_{optimisation_step}",
427 log_hypervolume=log_hypervolume,
428 )