helicon.optimize

Optimization: analytical screening, constraints, Gaussian process surrogate, Bayesian optimization, Pareto fronts.

Analytical Pre-Screening

helicon.optimize.analytical.screen_geometry(coils, *, z_min, z_max, n_pts=200, gamma=5.0 / 3.0, backend='auto')

Run the full Tier 1 analytical screening for a coil geometry.

Computes mirror ratio from Biot-Savart, then derives all analytical performance metrics in < 100 ms.

Parameters:

Name	Type	Description	Default
coils	list of Coil	Coil definitions.	required
z_min	float	Axial domain [m].	required
z_max	float	Axial domain [m].	required
n_pts	int	On-axis evaluation resolution.	200
gamma	float	Electron polytropic index.	5.0 / 3.0
backend	str	Biot-Savart backend.	'auto'

Returns:

Type	Description
NozzleScreeningResult

Source code in src/helicon/optimize/analytical.py

def screen_geometry(
    coils: list,
    *,
    z_min: float,
    z_max: float,
    n_pts: int = 200,
    gamma: float = 5.0 / 3.0,
    backend: str = "auto",
) -> NozzleScreeningResult:
    """Run the full Tier 1 analytical screening for a coil geometry.

    Computes mirror ratio from Biot-Savart, then derives all analytical
    performance metrics in < 100 ms.

    Parameters
    ----------
    coils : list of Coil
        Coil definitions.
    z_min, z_max : float
        Axial domain [m].
    n_pts : int
        On-axis evaluation resolution.
    gamma : float
        Electron polytropic index.
    backend : str
        Biot-Savart backend.

    Returns
    -------
    NozzleScreeningResult
    """
    R_B = mirror_ratio(coils, z_min=z_min, z_max=z_max, n_pts=n_pts, backend=backend)
    C_T = thrust_coefficient_paraxial(R_B, gamma=gamma)
    eta_T = thrust_efficiency(R_B, gamma=gamma)
    theta = divergence_half_angle(R_B)

    return NozzleScreeningResult(
        mirror_ratio=R_B,
        thrust_coefficient=C_T,
        divergence_half_angle_deg=theta,
        thrust_efficiency=eta_T,
    )

helicon.optimize.analytical.NozzleScreeningResult(mirror_ratio, thrust_coefficient, divergence_half_angle_deg, thrust_efficiency) dataclass

Fast analytical screening metrics for a coil configuration.

Attributes:

Name	Type	Description
mirror_ratio	float	R_B = B_throat / B_exit. Should be >> 1 for good confinement.
thrust_coefficient	float	C_T = F / (ṁ c_s). Dimensionless nozzle performance metric.
divergence_half_angle_deg	float	Plume half-angle estimate [°]. Lower is better.
thrust_efficiency	float	η_T = 1 − 1/√R_B. Fraction of thermal energy converted to thrust.

Constraints

helicon.optimize.constraints.CoilConstraints(max_total_mass_kg=None, max_total_power_W=None, max_B_conductor_T=None, current_density_Am2=10000000.0, conductor_resistivity_Ohm_m=1.72e-08, conductor_density_kg_m3=8960.0) dataclass

Engineering constraint specification for coil optimization.

Parameters:

Name	Type	Description	Default
max_total_mass_kg	float or None	Maximum total mass of all coils combined [kg]. None = unconstrained.	None
max_total_power_W	float or None	Maximum total resistive power dissipation [W]. None = unconstrained. For superconducting coils set this to None (zero resistive loss).	None
max_B_conductor_T	float or None	Maximum peak field at the conductor surface [T]. REBCO tapes have a practical limit of ~15-20 T; copper coils are limited by structural / Lorentz force considerations. None = unconstrained.	None
current_density_Am2	float	Maximum current density for conductor sizing [A/m²]. Typical values: 1-10 MA/m² for copper, 100-500 MA/m² for REBCO.	10000000.0
conductor_resistivity_Ohm_m	float	Electrical resistivity of conductor [Ω·m]. Copper at 20 °C: 1.72e-8 Ω·m. Set to 0.0 for superconducting (zero resistive loss).	1.72e-08
conductor_density_kg_m3	float	Mass density of conductor [kg/m³]. Copper: 8960 kg/m³.	8960.0

helicon.optimize.constraints.evaluate_constraints(coil_params, constraints, *, penalty_factor=1000.0)

Evaluate engineering constraints for a set of coil parameters.

Parameters:

Name	Type	Description	Default
coil_params	(array, shape(N_coils, 3))	Coil parameters with columns [z, r, I].	required
constraints	CoilConstraints	Active constraint specification.	required
penalty_factor	float	Scaling factor for quadratic penalty (used by optimizer).	1000.0

Returns:

Type	Description
CoilConstraintResult

Source code in src/helicon/optimize/constraints.py

def evaluate_constraints(
    coil_params: np.ndarray,
    constraints: CoilConstraints,
    *,
    penalty_factor: float = 1.0e3,
) -> CoilConstraintResult:
    """Evaluate engineering constraints for a set of coil parameters.

    Parameters
    ----------
    coil_params : array, shape (N_coils, 3)
        Coil parameters with columns ``[z, r, I]``.
    constraints : CoilConstraints
        Active constraint specification.
    penalty_factor : float
        Scaling factor for quadratic penalty (used by optimizer).

    Returns
    -------
    CoilConstraintResult
    """
    params = np.atleast_2d(np.asarray(coil_params, dtype=float))
    params[:, 0]
    r_arr = params[:, 1]
    I_arr = params[:, 2]

    J = constraints.current_density_Am2
    rho_elec = constraints.conductor_resistivity_Ohm_m
    rho_mass = constraints.conductor_density_kg_m3

    # Per-coil quantities
    abs_I = np.abs(I_arr)
    r_safe = np.maximum(r_arr, 1e-9)
    L_wire = 2.0 * math.pi * r_safe  # conductor path length per coil [m]
    A_cond = abs_I / max(J, 1e-30)  # conductor cross-section [m²]

    masses = rho_mass * A_cond * L_wire  # coil mass [kg]
    total_mass = float(np.sum(masses))

    R_coil = rho_elec * L_wire / np.maximum(A_cond, 1e-30)
    powers = abs_I**2 * R_coil  # coil power [W]
    total_power = float(np.sum(powers))

    B_peak = _MU0 * abs_I / (2.0 * r_safe)  # on-axis field at coil [T]

    violations: dict[str, float] = {}
    penalty = 0.0

    if constraints.max_total_mass_kg is not None:
        v = total_mass - constraints.max_total_mass_kg
        violations["total_mass_kg"] = float(v)
        if v > 0:
            penalty += penalty_factor * v**2

    if constraints.max_total_power_W is not None:
        v = total_power - constraints.max_total_power_W
        violations["total_power_W"] = float(v)
        if v > 0:
            penalty += penalty_factor * v**2

    if constraints.max_B_conductor_T is not None:
        B_max = float(np.max(B_peak))
        v = B_max - constraints.max_B_conductor_T
        violations["max_B_conductor_T"] = float(v)
        if v > 0:
            penalty += penalty_factor * v**2

    satisfied = all(v <= 0 for v in violations.values())

    return CoilConstraintResult(
        coil_masses_kg=masses.tolist(),
        coil_powers_W=powers.tolist(),
        coil_B_peak_T=B_peak.tolist(),
        total_mass_kg=total_mass,
        total_power_W=total_power,
        violations=violations,
        penalty=penalty,
        satisfied=satisfied,
    )

helicon.optimize.constraints.make_constrained_objective(objective_fn, constraints, *, penalty_factor=1000.0)

Wrap an MLX objective with a differentiable constraint penalty.

The returned function has the same signature as objective_fn but adds a quadratic penalty for each violated constraint:

f_constrained(x) = objective_fn(x) + Σ penalty_factor * max(g_i(x), 0)²

All operations use MLX so the result is differentiable via mlx.core.grad.

Parameters:

Name	Type	Description	Default
objective_fn	callable	MLX-differentiable objective; takes coil_params (shape (N, 3)) and returns a scalar mx.array.	required
constraints	CoilConstraints	Active constraints.	required
penalty_factor	float	Quadratic penalty weight.	1000.0

Returns:

Type	Description
callable	Penalized objective suitable for mlx.optimizers.Adam.

Examples:

::

bounds = constraints.CoilConstraints(max_total_mass_kg=50.0, max_B_conductor_T=15.0)
penalized = make_constrained_objective(throat_ratio_objective, bounds)
result = optimize_coils_mlx(init_params, penalized)

Source code in src/helicon/optimize/constraints.py

def make_constrained_objective(
    objective_fn,
    constraints: CoilConstraints,
    *,
    penalty_factor: float = 1.0e3,
):
    """Wrap an MLX objective with a differentiable constraint penalty.

    The returned function has the same signature as ``objective_fn`` but
    adds a quadratic penalty for each violated constraint:

        f_constrained(x) = objective_fn(x) + Σ penalty_factor * max(g_i(x), 0)²

    All operations use MLX so the result is differentiable via ``mlx.core.grad``.

    Parameters
    ----------
    objective_fn : callable
        MLX-differentiable objective; takes ``coil_params`` (shape (N, 3))
        and returns a scalar ``mx.array``.
    constraints : CoilConstraints
        Active constraints.
    penalty_factor : float
        Quadratic penalty weight.

    Returns
    -------
    callable
        Penalized objective suitable for ``mlx.optimizers.Adam``.

    Examples
    --------
    ::

        bounds = constraints.CoilConstraints(max_total_mass_kg=50.0, max_B_conductor_T=15.0)
        penalized = make_constrained_objective(throat_ratio_objective, bounds)
        result = optimize_coils_mlx(init_params, penalized)
    """
    try:
        import mlx.core as mx
    except ImportError as exc:
        msg = "MLX is required for make_constrained_objective. Install with: pip install mlx"
        raise ImportError(msg) from exc

    J = constraints.current_density_Am2
    rho_elec = constraints.conductor_resistivity_Ohm_m
    rho_mass = constraints.conductor_density_kg_m3
    mu0 = float(_MU0)

    def constrained(coil_params: mx.array) -> mx.array:
        obj = objective_fn(coil_params)

        r = coil_params[:, 1]
        I = coil_params[:, 2]
        r_safe = mx.maximum(r, 1e-9)
        abs_I = mx.abs(I)

        L_wire = 2.0 * math.pi * r_safe
        A_cond = abs_I / max(J, 1e-30)
        masses = rho_mass * A_cond * L_wire
        total_mass = mx.sum(masses)

        R_coil = rho_elec * L_wire / mx.maximum(A_cond, 1e-30)
        powers = abs_I**2 * R_coil
        total_power = mx.sum(powers)

        B_peak = mu0 * abs_I / (2.0 * r_safe)

        pen = mx.array(0.0)

        if constraints.max_total_mass_kg is not None:
            v = mx.maximum(total_mass - constraints.max_total_mass_kg, 0.0)
            pen = pen + penalty_factor * v**2

        if constraints.max_total_power_W is not None:
            v = mx.maximum(total_power - constraints.max_total_power_W, 0.0)
            pen = pen + penalty_factor * v**2

        if constraints.max_B_conductor_T is not None:
            v = mx.maximum(mx.max(B_peak) - constraints.max_B_conductor_T, 0.0)
            pen = pen + penalty_factor * v**2

        return obj + pen

    return constrained

GP Surrogate & Bayesian Optimization

helicon.optimize.surrogate.GPSurrogate(normalize_y=True, n_restarts=5)

Gaussian Process surrogate model backed by scikit-learn.

Uses a Matérn 5/2 kernel with automatic hyperparameter optimization. Provides Expected Improvement acquisition function for Bayesian optimization of noisy, expensive-to-evaluate objectives.

Parameters:

Name	Type	Description	Default
normalize_y	bool	Normalize target values before fitting (recommended).	True
n_restarts	int	Number of hyperparameter optimization restarts.	5

Source code in src/helicon/optimize/surrogate.py

def __init__(self, normalize_y: bool = True, n_restarts: int = 5):
    if not HAS_SKLEARN:
        raise ImportError(
            "scikit-learn is required for GPSurrogate. "
            "Install with: pip install scikit-learn"
        )
    # Wider bounds avoid hitting the lower limit when inputs span many
    # orders of magnitude (e.g. coil current 1e3–1e5 A).  After
    # StandardScaler normalization the effective length scale is ~O(1).
    kernel = Matern(nu=2.5, length_scale_bounds=(1e-2, 1e2))
    self._gp = GaussianProcessRegressor(
        kernel=kernel,
        n_restarts_optimizer=n_restarts,
        normalize_y=normalize_y,
        alpha=1e-6,
    )
    self._scaler: StandardScaler = StandardScaler()
    self._fitted = False

Functions

expected_improvement(X, y_best, xi=0.01)

Expected improvement acquisition function (for maximization).

Parameters:

Name	Type	Description	Default
X	(array, shape(n_points, n_params))		required
y_best	float	Best observed objective value so far.	required
xi	float	Exploration-exploitation trade-off (larger → more exploration).	0.01

Returns:

Name	Type	Description
EI	(array, shape(n_points))	Non-negative expected improvement at each candidate point.

Source code in src/helicon/optimize/surrogate.py

def expected_improvement(
    self, X: np.ndarray, y_best: float, xi: float = 0.01
) -> np.ndarray:
    """Expected improvement acquisition function (for maximization).

    Parameters
    ----------
    X : array, shape (n_points, n_params)
    y_best : float
        Best observed objective value so far.
    xi : float
        Exploration-exploitation trade-off (larger → more exploration).

    Returns
    -------
    EI : array, shape (n_points,)
        Non-negative expected improvement at each candidate point.
    """
    from scipy.stats import norm

    result = self.predict(X)
    mu, sigma = result.mean, result.std
    sigma = np.maximum(sigma, 1e-9)
    improvement = mu - y_best - xi
    Z = improvement / sigma
    ei = improvement * norm.cdf(Z) + sigma * norm.pdf(Z)
    return np.maximum(ei, 0.0)

fit(X, y)

Fit the GP to observations.

Inputs are normalized to zero mean and unit variance so the GP length-scale hyperparameter operates in a consistent numerical range regardless of the physical units of the parameters.

Parameters:

Name	Type	Description	Default
X	(array, shape(n_samples, n_params))		required
y	(array, shape(n_samples))		required

Source code in src/helicon/optimize/surrogate.py

def fit(self, X: np.ndarray, y: np.ndarray) -> None:
    """Fit the GP to observations.

    Inputs are normalized to zero mean and unit variance so the GP
    length-scale hyperparameter operates in a consistent numerical range
    regardless of the physical units of the parameters.

    Parameters
    ----------
    X : array, shape (n_samples, n_params)
    y : array, shape (n_samples,)
    """
    X_np = np.asarray(X, dtype=float)
    self._scaler.fit(X_np)
    self._gp.fit(self._scaler.transform(X_np), np.asarray(y, dtype=float))
    self._fitted = True

predict(X)

Predict mean and std at new points.

Parameters:

Name	Type	Description	Default
X	(array, shape(n_points, n_params))		required

Returns:

Type	Description
SurrogateResult with ``mean`` and ``std`` arrays.

Source code in src/helicon/optimize/surrogate.py

def predict(self, X: np.ndarray) -> SurrogateResult:
    """Predict mean and std at new points.

    Parameters
    ----------
    X : array, shape (n_points, n_params)

    Returns
    -------
    SurrogateResult with ``mean`` and ``std`` arrays.
    """
    if not self._fitted:
        raise RuntimeError("GPSurrogate must be fit() before predict()")
    X_scaled = self._scaler.transform(np.asarray(X, dtype=float))
    mean, std = self._gp.predict(X_scaled, return_std=True)
    return SurrogateResult(mean=mean, std=std)

helicon.optimize.surrogate.BayesianOptimizer(bounds, n_init=5, seed=0)

Sequential model-based optimizer using a GP surrogate.

Uses Expected Improvement as the acquisition function and maximizes it over a dense random candidate grid (no extra solver dependencies).

Parameters:

Name	Type	Description	Default
bounds	(array - like, shape(n_params, 2))	[[low_0, high_0], [low_1, high_1], ...] for each parameter.	required
n_init	int	Number of random evaluations before switching to GP-guided search.	5
seed	int	Random seed for initial sampling and acquisition maximization.	0

Source code in src/helicon/optimize/surrogate.py

def __init__(
    self,
    bounds: np.ndarray | list,
    n_init: int = 5,
    seed: int = 0,
):
    self.bounds = np.asarray(bounds, dtype=float)
    self.n_init = n_init
    self._rng = np.random.default_rng(seed)
    self._X: list[np.ndarray] = []
    self._y: list[float] = []
    self._surrogate: GPSurrogate | None = GPSurrogate() if HAS_SKLEARN else None

Attributes

n_evaluated property

Number of observations recorded so far.

Functions

ask(n=1)

Suggest the next point(s) to evaluate.

Returns random points until n_init observations are available, then maximizes Expected Improvement over a random candidate grid.

Returns:

Name	Type	Description
X	(array, shape(n, n_params))

Source code in src/helicon/optimize/surrogate.py

def ask(self, n: int = 1) -> np.ndarray:
    """Suggest the next point(s) to evaluate.

    Returns random points until ``n_init`` observations are available,
    then maximizes Expected Improvement over a random candidate grid.

    Returns
    -------
    X : array, shape (n, n_params)
    """
    n_params = self.bounds.shape[0]
    low = self.bounds[:, 0]
    high = self.bounds[:, 1]

    if self.n_evaluated < self.n_init or self._surrogate is None:
        return low + self._rng.uniform(size=(n, n_params)) * (high - low)

    X_train = np.array(self._X)
    y_train = np.array(self._y)
    self._surrogate.fit(X_train, y_train)

    n_candidates = max(1000, 200 * n_params)
    X_cand = low + self._rng.uniform(size=(n_candidates, n_params)) * (high - low)
    ei = self._surrogate.expected_improvement(X_cand, y_best=float(np.max(y_train)))
    top_idx = np.argsort(ei)[-n:][::-1]
    return X_cand[top_idx]

best()

Return the best observed point and its objective value.

Returns:

Name	Type	Description
x_best	(array, shape(n_params))
y_best	float

Source code in src/helicon/optimize/surrogate.py

def best(self) -> tuple[np.ndarray, float]:
    """Return the best observed point and its objective value.

    Returns
    -------
    x_best : array, shape (n_params,)
    y_best : float
    """
    if not self._y:
        raise RuntimeError("No observations yet. Call tell() first.")
    best_idx = int(np.argmax(self._y))
    return self._X[best_idx].copy(), self._y[best_idx]

tell(X, y)

Record new observations.

Parameters:

Name	Type	Description	Default
X	(array, shape(n, n_params) or (n_params,))		required
y	(array, shape(n) or scalar)		required

Source code in src/helicon/optimize/surrogate.py

def tell(self, X: np.ndarray, y: np.ndarray | float) -> None:
    """Record new observations.

    Parameters
    ----------
    X : array, shape (n, n_params) or (n_params,)
    y : array, shape (n,) or scalar
    """
    X = np.atleast_2d(X)
    y = np.atleast_1d(y)
    for xi, yi in zip(X, y):
        self._X.append(xi.copy())
        self._y.append(float(yi))

Gradient-Based Optimization (MLX)

helicon.optimize.gradient.GradientOptimizerConfig(n_steps=200, learning_rate=0.001, beta1=0.9, beta2=0.999, eps_adam=1e-08, tol=1e-06, history_every=10, n_phi=64) dataclass

Hyper-parameters for :class:GradientOptimizer.

Attributes:

Name	Type	Description
n_steps	int	Maximum number of gradient steps. Default: 200.
learning_rate	float	Adam base learning rate. Default: 1e-3.
beta1	float	Adam first-moment decay. Default: 0.9.
beta2	float	Adam second-moment decay. Default: 0.999.
eps_adam	float	Adam numerical stability term. Default: 1e-8.
tol	float	Convergence tolerance on gradient L2 norm. Optimization stops early when ||grad||_2 < tol. Default: 1e-6.
history_every	int	Record coil_params snapshot every this many steps. Default: 10.
n_phi	int	Azimuthal quadrature points for the MLX Biot-Savart backend. Default: 64.

helicon.optimize.gradient.GradientResult(coil_params_history, objective_history, final_coil_params, n_steps_run, converged) dataclass

Result from a :class:GradientOptimizer run.

Attributes:

Name	Type	Description
coil_params_history	list of np.ndarray	Parameter snapshots recorded every config.history_every steps. Each array has the same shape as the initial parameter array.
objective_history	list of float	Scalar objective value at each gradient step.
final_coil_params	ndarray	Optimized parameters at the last completed step.
n_steps_run	int	Actual number of steps executed (may be less than n_steps if converged early).
converged	bool	True if gradient norm dropped below config.tol.

helicon.optimize.gradient.GradientOptimizer(grid, objective_fn, config=None)

Gradient descent optimizer backed by mlx.core.grad.

Parameters:

Name	Type	Description	Default
grid	Grid	Axisymmetric computation grid — passed through to the objective function during :meth:run if needed. Not used internally by the optimizer itself; stored for user convenience.	required
objective_fn	callable	(coil_params: mx.array) -> mx.array — must return a scalar MLX array representing the loss to minimise. The function must be differentiable with respect to its argument via mlx.core.grad.	required
config	GradientOptimizerConfig	Optimizer hyper-parameters. Defaults to :class:GradientOptimizerConfig with all default values.	None

Notes

The optimizer performs gradient descent (minimisation). If you want to maximise a metric, negate it inside objective_fn.

Source code in src/helicon/optimize/gradient.py

def __init__(
    self,
    grid,
    objective_fn: Callable,
    config: GradientOptimizerConfig | None = None,
) -> None:
    self.grid = grid
    self.objective_fn = objective_fn
    self.config = config if config is not None else GradientOptimizerConfig()

Functions

run(initial_params)

Execute the optimization loop.

Parameters:

Name	Type	Description	Default
initial_params	ndarray	Starting coil parameters. Typically shape (N_coils, 3) for [z, r, I] parameterisation, but any shape is accepted as long as objective_fn handles it.	required

Returns:

Type	Description
GradientResult

Source code in src/helicon/optimize/gradient.py

def run(self, initial_params: np.ndarray) -> GradientResult:
    """Execute the optimization loop.

    Parameters
    ----------
    initial_params : np.ndarray
        Starting coil parameters.  Typically shape ``(N_coils, 3)`` for
        ``[z, r, I]`` parameterisation, but any shape is accepted as
        long as ``objective_fn`` handles it.

    Returns
    -------
    GradientResult
    """
    _require_mlx()

    cfg = self.config
    params = mx.array(initial_params.astype(np.float32))

    # Adam state
    m = mx.zeros_like(params)
    v = mx.zeros_like(params)

    objective_history: list[float] = []
    coil_params_history: list[np.ndarray] = []
    converged = False

    grad_fn = mx.grad(self.objective_fn)

    for step in range(cfg.n_steps):
        obj_val = self.objective_fn(params)
        g = grad_fn(params)
        mx.eval(obj_val, g)

        obj_scalar = float(np.array(obj_val))
        objective_history.append(obj_scalar)

        # Record parameter snapshot every history_every steps (and step 0)
        if step % cfg.history_every == 0:
            coil_params_history.append(np.array(params))

        # Adam update (descent: subtract update)
        t = step + 1
        m = cfg.beta1 * m + (1.0 - cfg.beta1) * g
        v = cfg.beta2 * v + (1.0 - cfg.beta2) * g * g
        m_hat = m / (1.0 - cfg.beta1**t)
        v_hat = v / (1.0 - cfg.beta2**t)
        params = params - cfg.learning_rate * m_hat / (mx.sqrt(v_hat) + cfg.eps_adam)
        mx.eval(params)

        # Convergence check: gradient L2 norm
        g_np = np.array(g)
        grad_norm = float(np.sqrt(np.sum(g_np**2)))
        if grad_norm < cfg.tol:
            converged = True
            # Record final snapshot on convergence
            coil_params_history.append(np.array(params))
            break

    final_params = np.array(params)

    return GradientResult(
        coil_params_history=coil_params_history,
        objective_history=objective_history,
        final_coil_params=final_params,
        n_steps_run=len(objective_history),
        converged=converged,
    )

helicon.optimize.gradient.optimize_mirror_ratio(base_coil_params, grid, *, target_mirror_ratio=5.0, n_steps=200, learning_rate=0.001, backend='auto')

Optimize coil currents to maximize mirror ratio via gradient descent.

Builds an objective_fn that calls :func:~helicon.fields.biot_savart.compute_bfield_mlx_differentiable and returns the negative mirror ratio (so minimising = maximising R_B = B_max / B_exit).

Parameters:

Name	Type	Description	Default
base_coil_params	(ndarray, shape(N_coils, 3))	Initial coil parameters [z_coil, radius, current] in SI units.	required
grid	Grid	Axisymmetric computation grid used to evaluate the field.	required
target_mirror_ratio	float	Currently unused — reserved for future penalised objectives. The optimization maximises R_B unconditionally.	5.0
n_steps	int	Maximum gradient steps. Default: 200.	200
learning_rate	float	Adam learning rate. Default: 1e-3.	0.001
backend	str	Backend selector — "auto" or "mlx". "numpy" is not supported because the objective must be differentiable; if "numpy" or any unsupported backend is requested, an :class:ImportError is raised. "auto" selects MLX when available, otherwise raises :class:ImportError.	'auto'

Returns:

Type	Description
GradientResult

Raises:

Type	Description
ImportError	When MLX is not available or an incompatible backend is requested.

Source code in src/helicon/optimize/gradient.py

def optimize_mirror_ratio(
    base_coil_params: np.ndarray,
    grid,
    *,
    target_mirror_ratio: float = 5.0,
    n_steps: int = 200,
    learning_rate: float = 1e-3,
    backend: str = "auto",
) -> GradientResult:
    """Optimize coil currents to maximize mirror ratio via gradient descent.

    Builds an ``objective_fn`` that calls
    :func:`~helicon.fields.biot_savart.compute_bfield_mlx_differentiable`
    and returns the *negative* mirror ratio (so minimising = maximising
    R_B = B_max / B_exit).

    Parameters
    ----------
    base_coil_params : np.ndarray, shape (N_coils, 3)
        Initial coil parameters ``[z_coil, radius, current]`` in SI units.
    grid : Grid
        Axisymmetric computation grid used to evaluate the field.
    target_mirror_ratio : float
        Currently unused — reserved for future penalised objectives.
        The optimization maximises R_B unconditionally.
    n_steps : int
        Maximum gradient steps.  Default: 200.
    learning_rate : float
        Adam learning rate.  Default: 1e-3.
    backend : str
        Backend selector — ``"auto"`` or ``"mlx"``.  ``"numpy"`` is not
        supported because the objective must be differentiable; if
        ``"numpy"`` or any unsupported backend is requested, an
        :class:`ImportError` is raised.  ``"auto"`` selects MLX when
        available, otherwise raises :class:`ImportError`.

    Returns
    -------
    GradientResult

    Raises
    ------
    ImportError
        When MLX is not available or an incompatible backend is requested.
    """
    from helicon.fields.biot_savart import (
        HAS_MLX as _HAS_MLX,
    )
    from helicon.fields.biot_savart import (
        compute_bfield_mlx_differentiable,
    )

    # Resolve backend
    if backend == "auto":
        if not _HAS_MLX:
            raise ImportError(
                "MLX is required for gradient-based mirror ratio optimization. "
                "Install with: pip install 'helicon[mlx]'"
            )
    elif backend != "mlx":
        raise ImportError(
            f"Backend {backend!r} does not support automatic differentiation. "
            "Use backend='mlx' or backend='auto' for gradient-based optimization."
        )

    _require_mlx()

    # Build grid arrays once (outside the hot loop)
    np.linspace(0.0, grid.r_max, grid.nr).astype(np.float32)
    z_np = np.linspace(grid.z_min, grid.z_max, grid.nz).astype(np.float32)
    # Use on-axis points only (r=0) for efficiency
    z_axis_np = z_np.copy()
    r_axis_np = np.zeros_like(z_axis_np)

    grid_r_mlx = mx.array(r_axis_np)
    grid_z_mlx = mx.array(z_axis_np)

    def objective_fn(coil_params: mx.array) -> mx.array:
        """Return negative mirror ratio (to minimise = maximise R_B)."""
        _, Bz = compute_bfield_mlx_differentiable(
            coil_params,
            grid_r_mlx,
            grid_z_mlx,
            n_phi=64,
        )
        Bz_abs = mx.abs(Bz)
        B_max = mx.max(Bz_abs)
        # B_exit: last on-axis point; guard against near-zero
        B_exit = mx.maximum(Bz_abs[-1], mx.array(1e-20))
        mirror_ratio_val = B_max / B_exit
        return -mirror_ratio_val  # negate: minimise = maximise R_B

    config = GradientOptimizerConfig(
        n_steps=n_steps,
        learning_rate=learning_rate,
    )
    optimizer = GradientOptimizer(grid, objective_fn, config=config)
    return optimizer.run(base_coil_params)

Pareto Front

helicon.optimize.pareto.ParetoResult(front_mask, front_indices, costs) dataclass

Pareto front from a multi-objective evaluation.

Attributes:

Name	Type	Description
front_mask	np.ndarray of bool, shape (n_points,)	True for Pareto-optimal (non-dominated) points.
front_indices	np.ndarray of int	Indices of Pareto-optimal points in the original array.
costs	(ndarray, shape(n_points, n_objectives))	Original cost matrix (minimization convention).

Attributes

front_costs property

Cost values for Pareto-optimal points only.

Functions

plot(*, labels=None, ax=None, figsize=(6, 5), dominated_color='lightgray', front_color='steelblue')

Plot the 2-objective Pareto front (minimization convention).

Parameters:

Name	Type	Description	Default
labels	tuple of str	Axis labels (x_label, y_label).	None
ax	matplotlib Axes	Axes to draw on. Creates new figure if None.	None
figsize	tuple	Figure size when creating a new figure.	(6, 5)
dominated_color	str	Color for dominated (non-Pareto) points.	'lightgray'
front_color	str	Color for Pareto-optimal points.	'steelblue'

Returns:

Type	Description
(fig, ax)

Source code in src/helicon/optimize/pareto.py

def plot(
    self,
    *,
    labels: tuple[str, str] | None = None,
    ax=None,
    figsize: tuple = (6, 5),
    dominated_color: str = "lightgray",
    front_color: str = "steelblue",
):
    """Plot the 2-objective Pareto front (minimization convention).

    Parameters
    ----------
    labels : tuple of str, optional
        Axis labels ``(x_label, y_label)``.
    ax : matplotlib Axes, optional
        Axes to draw on. Creates new figure if None.
    figsize : tuple
        Figure size when creating a new figure.
    dominated_color : str
        Color for dominated (non-Pareto) points.
    front_color : str
        Color for Pareto-optimal points.

    Returns
    -------
    fig, ax
    """
    try:
        import matplotlib.pyplot as plt
    except ImportError as exc:
        raise ImportError("matplotlib is required for ParetoResult.plot()") from exc

    if self.costs.shape[1] != 2:
        raise ValueError("plot() is only supported for 2-objective problems")

    if ax is None:
        fig, ax = plt.subplots(figsize=figsize)
    else:
        fig = ax.get_figure()

    dominated_mask = self.front_mask == False  # noqa: E712
    if np.any(dominated_mask):
        ax.scatter(
            self.costs[dominated_mask, 0],
            self.costs[dominated_mask, 1],
            c=dominated_color,
            s=30,
            label="Dominated",
            zorder=2,
        )

    front = self.front_costs
    # Sort by first objective for line connection
    order = np.argsort(front[:, 0])
    front_sorted = front[order]
    ax.plot(
        front_sorted[:, 0],
        front_sorted[:, 1],
        "o-",
        color=front_color,
        ms=7,
        lw=1.5,
        label="Pareto front",
        zorder=3,
    )

    if labels:
        ax.set_xlabel(labels[0])
        ax.set_ylabel(labels[1])
    ax.legend()
    ax.set_title("Pareto Front")
    fig.tight_layout()

    return fig, ax

helicon.optimize.pareto.pareto_front(costs)

Compute the Pareto front from a set of objective vectors.

Parameters:

Name	Type	Description	Default
costs	(array, shape(n_points, n_objectives))	Objective values in minimization convention. To maximize an objective, pass its negation.	required

Returns:

Type	Description
ParetoResult	Contains front_mask (bool array) and front_indices.

Examples:

Maximize thrust (F) and detachment efficiency (η_d) simultaneously::

costs = np.column_stack([-thrust_values, -eta_d_values])
result = pareto_front(costs)
best_configs = scan_points[result.front_indices]

Source code in src/helicon/optimize/pareto.py

def pareto_front(costs: np.ndarray) -> ParetoResult:
    """Compute the Pareto front from a set of objective vectors.

    Parameters
    ----------
    costs : array, shape (n_points, n_objectives)
        Objective values in **minimization** convention.
        To maximize an objective, pass its negation.

    Returns
    -------
    ParetoResult
        Contains ``front_mask`` (bool array) and ``front_indices``.

    Examples
    --------
    Maximize thrust (F) and detachment efficiency (η_d) simultaneously::

        costs = np.column_stack([-thrust_values, -eta_d_values])
        result = pareto_front(costs)
        best_configs = scan_points[result.front_indices]
    """
    costs = np.atleast_2d(np.asarray(costs, dtype=float))
    mask = ~is_dominated(costs)
    indices = np.where(mask)[0]
    return ParetoResult(front_mask=mask, front_indices=indices, costs=costs)

helicon.optimize.pareto.is_dominated(costs)

Return a boolean mask: True if a point is dominated by any other.

A point i is dominated by point j if j is at least as good in all objectives and strictly better in at least one.

Parameters:

Name	Type	Description	Default
costs	(array, shape(n_points, n_objectives))	Objective values. Minimization convention.	required

Returns:

Name	Type	Description
dominated	bool array, shape (n_points,)

Source code in src/helicon/optimize/pareto.py

def is_dominated(costs: np.ndarray) -> np.ndarray:
    """Return a boolean mask: True if a point is dominated by any other.

    A point *i* is dominated by point *j* if *j* is at least as good in
    all objectives and strictly better in at least one.

    Parameters
    ----------
    costs : array, shape (n_points, n_objectives)
        Objective values. **Minimization** convention.

    Returns
    -------
    dominated : bool array, shape (n_points,)
    """
    costs = np.asarray(costs, dtype=float)
    n = costs.shape[0]
    dominated = np.zeros(n, dtype=bool)

    for i in range(n):
        if dominated[i]:
            continue
        for j in range(n):
            if i == j:
                continue
            # j dominates i: all j ≤ i AND at least one j < i
            if np.all(costs[j] <= costs[i]) and np.any(costs[j] < costs[i]):
                dominated[i] = True
                break

    return dominated

Parameter Scans

helicon.optimize.scan.ParameterRange(path, low, high, n) dataclass

One parameter axis in a scan.

Parameters:

Name	Type	Description	Default
path	str	Dot-notation path into the config dict, e.g. "coils.0.I" or "plasma.T_i_eV".	required
low	float	Inclusive range boundaries.	required
high	float	Inclusive range boundaries.	required
n	int	Number of points along this axis.	required

Functions

from_string(s) classmethod

Parse "path:low:high:n" format used by the CLI.

Source code in src/helicon/optimize/scan.py

@classmethod
def from_string(cls, s: str) -> ParameterRange:
    """Parse ``"path:low:high:n"`` format used by the CLI."""
    parts = s.split(":")
    if len(parts) != 4:
        raise ValueError(f"Expected 'path:low:high:n', got {s!r}")
    return cls(path=parts[0], low=float(parts[1]), high=float(parts[2]), n=int(parts[3]))

values()

Return n linearly spaced values from low to high.

Source code in src/helicon/optimize/scan.py

def values(self) -> np.ndarray:
    """Return ``n`` linearly spaced values from ``low`` to ``high``."""
    return np.linspace(self.low, self.high, self.n)

helicon.optimize.scan.ScanResult(points, metrics, param_names, base_config, n_screened=0, objectives=None) dataclass

Results from a completed parameter scan.

Functions

plot_pareto(*, x_key='thrust_N', y_key=None, ax=None)

Plot the Pareto front for this scan result.

Extracts metric values for x_key and y_key from :attr:metrics, builds a cost matrix (negated for maximization), computes the Pareto front and calls :meth:ParetoResult.plot.

Parameters:

Name	Type	Description	Default
x_key	str	Metric key to use for the x-axis (default "thrust_N").	'thrust_N'
y_key	str	Metric key to use for the y-axis. Defaults to the second objective in :attr:objectives if available, else "beam_efficiency".	None
ax	matplotlib Axes	Axes to draw on. A new figure is created if None.	None

Returns:

Type	Description
(fig, ax)	Matplotlib figure and axes objects.

Raises:

Type	Description
ImportError	If matplotlib is not installed (skipped gracefully — returns (None, None)).

Source code in src/helicon/optimize/scan.py

def plot_pareto(
    self,
    *,
    x_key: str = "thrust_N",
    y_key: str | None = None,
    ax=None,
) -> tuple:
    """Plot the Pareto front for this scan result.

    Extracts metric values for *x_key* and *y_key* from
    :attr:`metrics`, builds a cost matrix (negated for maximization),
    computes the Pareto front and calls :meth:`ParetoResult.plot`.

    Parameters
    ----------
    x_key : str
        Metric key to use for the x-axis (default ``"thrust_N"``).
    y_key : str, optional
        Metric key to use for the y-axis.  Defaults to the second
        objective in :attr:`objectives` if available, else
        ``"beam_efficiency"``.
    ax : matplotlib Axes, optional
        Axes to draw on.  A new figure is created if *None*.

    Returns
    -------
    fig, ax
        Matplotlib figure and axes objects.

    Raises
    ------
    ImportError
        If matplotlib is not installed (skipped gracefully — returns
        ``(None, None)``).
    """
    try:
        import matplotlib  # noqa: F401
    except ImportError:
        return (None, None)

    from helicon.optimize.pareto import pareto_front

    # Determine y_key
    if y_key is None:
        if self.objectives and len(self.objectives) >= 2:
            y_key = self.objectives[1]
        else:
            y_key = "beam_efficiency"

    # Extract values; skip points where either metric is missing
    x_vals: list[float] = []
    y_vals: list[float] = []
    for m in self.metrics:
        xv = m.get(x_key)
        yv = m.get(y_key)
        if xv is not None and yv is not None:
            x_vals.append(float(xv))
            y_vals.append(float(yv))

    if not x_vals:
        return (None, None)

    # Negate for maximization convention (pareto_front uses minimization)
    costs = np.column_stack([-np.asarray(x_vals), -np.asarray(y_vals)])
    result = pareto_front(costs)

    return result.plot(labels=(x_key, y_key), ax=ax)

helicon.optimize.scan.generate_scan_points(base_config, ranges, *, method='grid', seed=0, prescreening=False, min_mirror_ratio=1.5)

Generate scan points from parameter ranges.

Parameters:

Name	Type	Description	Default
base_config	SimConfig	Base configuration to modify.	required
ranges	list of ParameterRange	Parameter axes to vary.	required
method	``"grid"`` | ``"lhc"``	"grid" — full Cartesian product. "lhc" — Latin Hypercube sampling with total n = product of all n values.	'grid'
seed	int	Random seed for LHC sampling.	0
prescreening	bool	If True, run Tier 1 analytical pre-screening after generating points. Points with mirror_ratio < min_mirror_ratio are marked screened_out=True but still included in the returned list.	False
min_mirror_ratio	float	Minimum acceptable mirror ratio when prescreening=True.	1.5

Returns:

Type	Description
list of ScanPoint

Source code in src/helicon/optimize/scan.py

def generate_scan_points(
    base_config: SimConfig,
    ranges: list[ParameterRange],
    *,
    method: str = "grid",
    seed: int = 0,
    prescreening: bool = False,
    min_mirror_ratio: float = 1.5,
) -> list[ScanPoint]:
    """Generate scan points from parameter ranges.

    Parameters
    ----------
    base_config : SimConfig
        Base configuration to modify.
    ranges : list of ParameterRange
        Parameter axes to vary.
    method : ``"grid"`` | ``"lhc"``
        ``"grid"`` — full Cartesian product.
        ``"lhc"`` — Latin Hypercube sampling with total n = product of all n values.
    seed : int
        Random seed for LHC sampling.
    prescreening : bool
        If True, run Tier 1 analytical pre-screening after generating points.
        Points with ``mirror_ratio < min_mirror_ratio`` are marked
        ``screened_out=True`` but still included in the returned list.
    min_mirror_ratio : float
        Minimum acceptable mirror ratio when ``prescreening=True``.

    Returns
    -------
    list of ScanPoint
    """
    if method == "grid":
        value_lists = [r.values() for r in ranges]
        combos = list(itertools.product(*value_lists))
        points = []
        for i, combo in enumerate(combos):
            params = {r.path: float(v) for r, v in zip(ranges, combo)}
            config = _apply_params(base_config, params)
            points.append(ScanPoint(index=i, params=params, config=config))

    elif method == "lhc":
        n_total = math.prod(r.n for r in ranges)
        rng = np.random.default_rng(seed)
        n_params = len(ranges)
        lhc = np.zeros((n_total, n_params))
        for j in range(n_params):
            perm = rng.permutation(n_total)
            lhc[:, j] = (perm + rng.uniform(size=n_total)) / n_total
        points = []
        for i in range(n_total):
            params: dict[str, float] = {}
            for j, r in enumerate(ranges):
                params[r.path] = float(r.low + lhc[i, j] * (r.high - r.low))
            config = _apply_params(base_config, params)
            points.append(ScanPoint(index=i, params=params, config=config))

    else:
        raise ValueError(f"Unknown scan method {method!r}. Use 'grid' or 'lhc'.")

    if prescreening:
        _apply_prescreening(points, base_config, min_mirror_ratio)

    return points