forge.optimise

Contains the principal FORGE optimisation procedure.

Copyright 2025-2026 Chris Marsden

This file is part of FORGE.

FORGE is free software: you can redistribute it and/or modify it under the terms of the GNU Lesser General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version.

FORGE is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public License for more details.

You should have received a copy of the GNU Lesser General Public License along with FORGE. If not, see <http://www.gnu.org/licenses/>.

class forge.optimise.Optimiser(eq, tokamak_initial, tokamak_opt=None, divertor_data=None, max_evals=3000, current_step_size_factor=0.05, estimate_initial_currents=False, initial_temperature=10.0, threshold_acceptance_rate_decay=2.0, initial_threshold_acceptance_rate=0.75, n_window=50, constraints=None, field_line_trace_step_size=0.15, field_line_trace_max_steps=1000, field_line_trace_psi_tollerance=0.01, buffer_intersection_penalty_factor=1.05, initial_total_connection_length_cost=1.0, initial_total_strike_point_distance_cost=1.0, initial_coil_currents_cost=1.0, initial_xpoint_regions_cost=1.0, use_xpoint_regions=False, xpoint_regions=None, max_magnetic_disconnection_factor=0.5, use_buffers=False, buffers=None, initial_alpha=50.0, alpha_update_factor=1.05, max_cooling_factor=0.99, min_cooling_factor=0.9, detailed_logging=False, cost_termination_fraction=0.01, min_temperature=1.0, on_iteration=None, stop_event=None)

Bases: object

Object representing the simulated annealing optimiser.

This Optimiser will carry out the simulated annealing-based optimisation of of the magnetic geometry of the divertor(s).

Parameters:

• eq (forge.equilibrium Equilibrium object) – The equilibrium that will be optimised.

• divertor_data (dict) – A dictionary of data pertaining to the divertor region(s) to be optimised.

• tokamak_initial (forge.machine Machine object) – Tokamak associated with the initial starting equilibrium. This machine should contain the PF coil set used to make the the initial equilibrium.

• tokamak_opt (forge.machine Machine object) – Tokamak that the user wishes to perform the optimisation on. This machine does not need to have the same PF coil set as tokamak_initial. If this is None, then tokamak_initial - the tokamak used to make the initial equilibrium - will be used. If tokamak_opt is not None then an estimation of an initial set of coil currents will be made.

• max_evals (int) – The maximum number of cost function evaluations to carry out (not including the initial evaluation).

• current_step_size_factor (float) – The step size taken when moving in current space is given by this factor of a typical coil current size taken from the set of initial coil currents.

• initial_temperature (float) – The initial temperature of the system. This pseudo-temperature is what is lowered during the annealing process - analogous of a hot metal being cooled.

• threshold_acceptance_rate_decay (float) – The annealing acceptance rate, over a specified window, must fall below a threshold value before cooling is allowed. This threshold acceptance rate is temperature dependent, and evolves as R = R0 * exp(-λ*(1-T/T0)) where R0 is the initial threshold value, T is the current temperature, T0 is the initial temperature and λ is this parameter, the threshold_acceptance_rate_decay constant. If λ is larger, the threshold drops more aggresively. If the optimiser should be more explorative, λ should be lower.

• initial_threshold_acceptance_rate (float) – The initial threshold value for the acceptance rate, R0.

• n_window (int) – The size of the window over which the rolling acceptance rate is calculated.

• constraints (dict) – A dictionary of constraints to be used. See the “generate_default_constraints” method.

• field_line_trace_step_size (float) – The poloidal step size (in metres) used by the RK4 field-line tracer.

• field_line_trace_max_steps (int) – The maximum number of steps the field-line tracer will take before abandoning a trace (e.g. if it never reaches the wall).

• field_line_trace_psi_tollerance (float) – Target maximum value of the relative error on the poloidal magnetic flux of the traced field line(s).

• buffer_intersection_penalty_factor (float) – A multiplier on the connection length cost term to be applied if the divertor leg intersects a buffer on the tokamak’s wall.

• initial_total_connection_length_cost (float) – The initial total combined connection length cost across all divertor regions.

• initial_total_strike_point_distance_cost (float) – The initial total combined strike point distance cost across all divertor regions.

• initial_coil_currents_cost (float) – The initial cost due to the coil currents.

• self.initial_xpoint_regions_cost (float) – The initial cost due to regions in which additional X-points are ncouraged to be formed.

• use_xpoint_regions (bool) – Flag for using regions in which secondary divertor X-point(s) will be encouraged to form.

• xpoint_regions (dict) – A dictionary mapping divertor region names to X-point region definitions. Required when use_xpoint_regions is True. Each key is a divertor region name (e.g. "lower_outer") and each value is a dictionary containing the R and Z coordinates of the region boundary: {"lower_outer": {"R": [...], "Z": [...]}, ...}.

• max_magnetic_disconnection_factor (float) – Maximum relative difference in poloidal magnetic flux between a secondary divertor X-point and the separatrix, i.e. |ψ_xpt − ψ_lcfs| / |ψ_lcfs|. This sets the scale at which the magnetic disconnection cost saturates when forming an X-point target (XPT) geometry.

• use_buffers (bool) – Flag for adding buffer regions around parts of the wall.

• buffers (dict) –

  Buffer region definitions keyed by divertor region name. Required when use_buffers is True. Each value is a list of dictionaries containing the R and Z coordinates of the wall segment endpoints and the buffer distance:

  {"lower_outer": [{"R": [R1, R2], "Z": [Z1, Z2], "distance": d}, ...], ...}
  

• initial_alpha (float) – The initial value of the alpha parameter used in the Metropolis-Hastings acceptance criterion. Alpha controls the sensitivity of the acceptance probability to the energy change. FORGE will automatically tune alpha to meet the first threshold acceptance rate, after which it is frozen.

• alpha_update_factor (float) – The factor by which the parameter alpha is updated in order to induce a first cooling event.

• max_cooling_factor (float) – The maximum allowed temperature cooling multiplication factor, β_max.

• min_cooling_factor (float) – The minimum allowed temperature cooling multiplication factor, β_min.

• detailed_logging (bool) – Flag for carrying out a detailed logging of data.

• cost_termination_fraction (float) – If the incumbent cost drops below this fraction of the initial cost, the annealing will terminate early. For example, 0.01 means the optimisation stops once the cost reaches 1% of its initial value.

• min_temperature (float) – The minimum temperature below which the annealing will stop.

optimise()

Performs the simulated annealing-based optimisation of the divertor(s)’ magnetic geometry.

evaluate_neighbour(neighbour, previous_accepted_cost)

Calculates the cost of a neighbouring state and accepts/rejects it accordingly.

Evaluates whether or not a new neighbouring state is going to be kept or not. In doing so, the cost function for this candidate state is evaluated.

Parameters:

• neighbour (dict) – The new candidate neighbour state to be evaluated.

• previous_accepted_cost (float) – The cost of the previous state that was deemed acceptable [dict]

Returns:

state_data (dict) – Dictionary of data of the neighbouring state. Includes updated data related to the cost of the state.

generate_new_neighbour_in_current_space(currents, step_size)

Generates a new neighbouring state.

Generates a new position in coil current space for evaluation. Samples a random unit vector from the null space already generated and multiplies it by a given step size. This new perturbation vector is then added to the current vector to make a new set of coil currents. A new flux map is then generated for this configuration.

Parameters:

• currents (1D numpy array.) – The exisiting coil currents (A).

• step_size (float) – A factor to scale the step size taken in coil current space.

Returns:

new state dictionary – A dictionary of data representing the new state.

generate_eq_from_currents(new_currents)

Generates equilibrium data from a set of coil currents.

Parameters:

new_currents (1D numpy array) – The set of coil currents.

Returns:

dictionary of equilibrium data. – Contains the 2D grid of poloidal magnetic flux, as well as the machine and plasma poloidal magnetic fluxes, and the coil currents.

calculate_cost(currents, flux_map)

Evaluates the cost of a state.

Calculates the cost of an equilibrium state. This includes all cost components of all divertor regions being optimised, as well as cost terms not tied to specific divertor regions, such as the coil currents cost term. In calculating the cost, field lines are traced in the divertor regions as appropriate. If X-point regions are used, then any secondary divertor X-points in these regions are located.

Parameters:

• currents (1D numpy array) – The set of coil currents.

• flux_map (2D numpy array) – 2D array of the state’s poloidal magnetic flux.

Returns:

Dictionary of cost data and divertor data – A dictionary containing detailed cost data as well as divertor data (such as the strike point, connection length, traced field line etc.) for each divertor region.

trace_field_line(starting_position, step_size, max_steps=1000, direction=1.0, grid_data=None, divertor_region=None)

Trace a magnetic field line until it hits a wall segment.

Trace a field line from the specified starting location. Stop tracing either when the field line intersects the machine’s wall, or a maximum number of steps along the field line has been taken.

Uses RK4 integration with a fast bilinear interpolator on pre-computed derivative grids and a gradient-projection correction after every step to keep the trace pinned to the correct flux surface.

Performance notes

The original bottleneck was psi_func.ev() called as a scalar inside a Python loop (~30-100 us overhead per call). By pre-computing dpsi/dR and dpsi/dZ on the regular grid once (in calculate_cost) and interpolating them bilinearly here (~1 us per call), the per-step cost drops by 30-50x. Combined with RK4’s larger feasible step size this gives an overall speed-up of roughly 10-20x on the tracer.

The gradient-projection step ensures that the accumulated psi drift after each RK4 step is corrected to machine precision, so the trace stays on the target flux surface regardless of step size.

param starting_position:

[R,Z] of the starting position for the trace.

type starting_position:

list

param step_size:

The size of each poloidal step along the field line.

type step_size:

float

param max_steps:

Maximum number of steps to take along the field line.

type max_steps:

int

param direction:

Direction to trace the field line. +1/-1 corresponds to the poloidal helicity of the trace.

type direction:

float

param grid_data:

Pre-computed grids and metadata from calculate_cost. Keys: psi_grid, dpsi_dR_grid, dpsi_dZ_grid, R_min, Z_min, inv_dR, inv_dZ, nR_m1, nZ_m1.

type grid_data:

dict

param divertor_region:

Name of the divertor region being traced (e.g. "lower_outer"). When buffers are keyed per region, only the buffers for this region are checked for intersections.

type divertor_region:

str, optional

check_field_line_intersection(field_line_R_start, field_line_Z_start, field_line_R_end, field_line_Z_end, divertor_region=None)

Locates intersections between a traced field line segment and wall/buffer structures.

Checks for an intersection between the provided field line segment and the tokamak’s wall. If buffer regions are present, intersections with these are also examined. If multiple intersections are present, the point closest to the field line’s starting point is returned. Note that this is for only a single segment of a field line, i.e. between two traced points.

Uses the prepared (spatially-indexed) wall geometry created during __init__ for a fast boolean pre-check before computing the full intersection geometry. This avoids the expensive intersection computation on steps that are entirely inside the domain (the vast majority).

Parameters:

• field_line_R_start (float) – R coordinate of the starting location of the field line segment.

• field_line_Z_start (float) – Z coordinate of the starting location of the field line segment.

• field_line_R_end (float) – R coordinate of the end location of the field line segment.

• field_line_Z_end (float) – Z coordinate of the end location of the field line segment.

• divertor_region (str, optional) – Name of the divertor region being traced. When provided and buffers are keyed per region, only the buffers for this region are checked.

Returns:

intersection_data (dict) – Dictionary of intersection data containing a flag for whether or not there was an intersection, the strike point (R,Z) coordinates (if there is an intersection) and an additional flag for if an intersection occurs with a buffer.

compute_coilset_greens_on_grid()

Computes the coils’ Green’s functions on the 2D equilibrium grid.

Pre-calculates, for each coil in the coilset, the Green’s function matrices for:

• Poloidal magnetic flux, psi — via coil.control_psi.

• The R-derivative of psi, dpsi/dR — derived from the coil’s analytical control_Bz response using the identity dpsi/dR = 2 pi R Bz.

• The Z-derivative of psi, dpsi/dZ — derived from the coil’s analytical control_Br response using the identity dpsi/dZ = -2 pi R Br.

These derivative matrices are used inside calculate_cost to construct the full dpsi/dR and dpsi/dZ grids via a single np.einsum call (a currents-vector times the precomputed matrices), which is much faster than fitting a RectBivariateSpline and numerically differentiating it every iteration. The analytical Green’s function derivatives (via greens.Greens_dpsi_dR / greens.Greens_dpsi_dZ, which underpin control_Bz / control_Br) are also more accurate than spline-based numerical differentiation.

generate_constraints()

Generates constraint matrices/vector for simulated annealing and Tikhonov regularisation (if used).

Generates various constraint matrices/vectors. In general, there will exist a matrix of Green’s functions related to the various constraints, as well as a corresponding vector of the constraints themselves. These constraints will be things such as the value of magnetic field or flux from the coils at a point.

There will exist up to two sets of these constraints. One set of constraints are used during the simulated annealing. The other set is used in the determination of the initial coil currents. Determining the initial coil currents is itself optional, but is particularly needed if the user decides to perform an optimisation using a set of PF coils that were different to those that were used to make the initial equilibrium. If the user is using the same PF set as was used to make the initial equilibrium, then this determination of the initial coil currents is not required. In such a case the user can still choose to estimate the initial coil currents. They may want to do this in such a way as to incroporate new divertor constraints, such as those describing the position of the divertor legs, as such a resultant equilibrium may provide a more useful initial starting point for their particular optimisation problem.

During the annealing a null-space approach is utilised, wherein there is a requirement that number of coils - number of constraints >= 1. It may not be possible for the user to include lots of additional divetor geometry constraints in the annealing, if the number of coils they have supplied is insufficient. However, if the user chooses to estimate the initial coil currents using the Tikhonov regulariastion approach, then during this process they are not subject to the strict requirement on the number of constraints, and can try to include more divertor constraints. In such a case, a larger set of cosntraints can be used, with a smaller set of constraints then later used during the annealing process. In this way, the user may be able to bias the initial equilibrium into having favourable features in the divertor, which whilst not constrained during the annealing, may be more likely to be produced naturally, having started from an initial equilibrium with such features.

The set of constraints used during the annealing are reffered to as the annealing constraints, with the set used during the Tikhonov regularisation reffered to as the Tikhonov constraints.

create_buffers()

Creates buffer regions around the wall of the machine.

The user can choose to create buffer regions around parts of the wall. These are useful if the user wants the divertor leg(s) to not come within a certain distance of parts of the wall. An example of this might be regions such as the divertor nose in tightly baffled geometries, which may not be equipped to handle heat fluxes as high as the main divertor targets.

Buffer definitions are keyed by divertor region name so that, at runtime, only the buffers belonging to the region currently being traced are checked for intersections. This improves efficiency and avoids false-positive buffer hits from unrelated regions.

Raises:

• ValueError – If use_buffers is True but no buffer data has been supplied.

• TypeError – If buffers is not a dict keyed by region name.

create_xpoint_regions()

Defines regions for additional divertor X-points.

The user can choose to define regions within which X-points are encouraged to be formed under optimisation. The regions are recorded as a dictionary mapping divertor region names to Shapely LineString objects.

Raises:

• ValueError – If use_xpoint_regions is True but no region data has been supplied.

• TypeError – If xpoint_regions is not a dictionary keyed by divertor region name.

init_xpoint_regions()

Initialises secondary divertor X-point regions.

Initialises regions in which the optimiser will encourage additional X-points to form. This is particularly useful if X-points target or snowflake divertor configurations are desired.

A list of regions in the form of Shapely LineStrings will have been pre-defined as an attribute of the tokamak being used in the optimisation. For each of these regions, those points on the 2D R,Z equilibrium grid that lie within the region(s) are identified and recorded. There will be N_points number of points that lie within the X-points region(s).

For each of these points a set of 2x2 Jacobian matrices for the poloidal field from each of the coils in the coilset is created, along with a corresponding 2x2 Jacobian matrix for the background poloidal field from the plasma.

The coil Jacobian matrices are themselves stored in a N_points x N_coils matrix, with the background field Jacobian matrices stored in an N_points vector.

These data will later be used to identify the location of additional secondary divertor X-points that may form within these regions.

We also record the background poloidal field components (Br,Bz) themselves, along with coil Br/Bz control responses, as a means of later quickly calculating the poloidal field strength inside the regions.

init_current_step_size(currents)

Sets the typical step size for perturbing the coil currents.

Parameters:

currents (1D numpy array.) – The coil currents (A).

estimate_currents()

Estimates the initial currents in the coils using Tikhonov regularisation.

Performs Tikhonov reguralisation to estimate initial coil currents, automatically tuning the regularisation parameter, alpha, to meet the equilibrium constraints within a specified threshold. This is optionally used, and is particularly useful when using a coil set that differs from the coil set used to produce the initial equilibrium. The simulated annealing will move us around in the subspace of coil currents that satisfy our constraints - this only works by perturbing an initial set of coil currents that themselves satisfy the constraints. Hence, if you have a new coil set, and don’t know what the initial currents are, you can use this to estimate such currents.

tikhonov_min_residual(G=None, b=None, threshold=0.01, alpha_start=1e-10, alpha_min=1e-35, decay_factor=0.8, max_iter=1000)

Performs Tikhonov regularisation on the coil curents.

Performs Tikhonov regularisation to solve G I = b for I, where I is a vector of coil currents, G is a matrix of Green’s functions related to a set of equilibrium constraints on field/flux, denoted by the vector b. The regularisation parameter, alpha, is automatically tuned (becoming smaller each iteration), such that the residuals are minimised until the maximum error value drops below some threshold value set by the user.

Parameters:

• G (2D numpy array) – 2D matrix of Green’s functions/coil responses.

• b (1D numpy array) – Vector of constraints.

• threshold (float) – Value that the maximum relative error must drop below.

• alpha_start (float) – The initial value of alpha.

• alpha_min (float) – The minimum allowed value of alpha.

• decay_factor (float) – Value by which alpha is multiplied during each tuning.

• max_iter (int) – Maximum number of interations/tunings of alpha to take.

Returns:

I (1D numpy array) – Vector of coil currents.

update_cooling_factor()

Updates the cooling factor, β.

Sets the cooling factor, β, for use in the cooling schedule (T_new = β * T_old). If many iterations at the current temperature were required to reach the threshold acceptance rate, then this indicates that the optimiser is still acting in an exploratory manner, hence β will remain high (the temperature will not decrease much). This aims to prevent premature cooling. Likewise, if not that many iterations were required, this indicates the optimiser is not being as explorative, hence β will remain low (the temperature will decrease more rapidly). β is bounded between β_min and β_max.

generate_default_constraints()

If the user does not specify any constraints, a dictionary of default constraints will be generated and returned.

init_divertor_data()

Initialises key pieces of data related to the divertor regions.

The starting position of the field line trace for each divertor region is determined, as is the direction of the field line trace. Divertor-related cost terms are also initialised.

request_stop()

Request the optimisation loop to stop gracefully.

The loop will terminate at the start of the next iteration. This is safe to call from any thread.

get_tracking_snapshot()

Return a snapshot dict of all tracking data for external consumers (e.g. GUI).

Returns:

dict – A dictionary containing copies of the current tracking lists and key state needed by forge.plotting.update_tracking_plots().

generate_optimised_eq_machine()

Creates new Equilibrium and Machine objects of the final state.

Creates a new forge.equilibrium.Equilibrium object for the optimised equilibrium. Also creates a new forge.machine.Machine object for the tokamak with the updated coil currents.

separatrix_coil_flux_change()

Quantify the change in coil flux along the separatrix after optimisation.

Evaluates the poloidal magnetic flux from the coils (psi_mach) along the initial separatrix for both the initial and optimised equilibria. If the coil flux is unchanged on the separatrix, the separatrix position is preserved and the background plasma equilibrium remains self-consistent.

Returns:

dict – A dictionary with the following keys:

• R_lcfs : ndarray — R coordinates of the separatrix evaluation points.

• Z_lcfs : ndarray — Z coordinates of the separatrix evaluation points.

• psi_mach_initial : ndarray — Coil flux along the separatrix (initial).

• psi_mach_optimised : ndarray — Coil flux along the separatrix (optimised).

• delta_psi_mach : ndarray — Absolute change (optimised - initial).

• delta_psi_mach_rel : ndarray — Relative change (optimised - initial) / initial.

• max_abs_change : float — Maximum absolute change along the separatrix.

• max_rel_change : float — Maximum relative change (dimensionless).

• mean_abs_change : float — Mean absolute change along the separatrix.

• mean_rel_change : float — Mean relative change (dimensionless).