API reference

PolyStep

Gradient-free polytope direct search with optimal-transport weighting, for piecewise-constant and combinatorial objectives. See minimize, PolyStepES and PolyStepConfig.

minimize(f, dim; steps=200, callback=nothing, kwargs...) -> PolyStepES

Run steps ask/tell rounds of the batched objective f(X::(dim, N))::(N,). callback(es) -> true stops early. The final iterate (clamped and repaired) is then scored once as a (dim, num_particles) matrix, costing num_particles extra evaluations, so best_f is never worse than f at the point the run ended on. f and repair must therefore accept any column count, not just popsize. Result in es.best_x, es.best_f and mean(es) (from Statistics).

PolyStepES(dim; num_particles=1, epsilon=0.5, step_radius=0.5, radius_decay=1.0,
           solver=SoftmaxSolver(), scale_cost=:mean, x0=nothing,
           lb=nothing, ub=nothing, repair=nothing,
           rng=Xoshiro(0), T=float(eltype(x0)))

Ask/tell wrapper around the PolyStep update. Population size is num_particles * 2 * dim (orthoplex). x0 is a (dim,) vector (replicated to all particles) or a (dim, num_particles) matrix. lb/ub (scalar or (dim,)) box-clamp candidates. repair(X_cols) post-processes candidates after clamping, for example by rounding for integrality. It edits X_cols in place or returns the repaired matrix. It must keep the bounds, because its output is not clamped again. T defaults to the element type of x0 (Float64 without one). best_x/best_f track the best evaluated point: every tell! candidate, plus the final iterate that minimize and PolyStepOptimizer score once. They are always feasible when bounds and a bounds-preserving repair are active.

step_radius is absolute: candidates are X +/- step_radius * R. (In PolyStepConfig an unscheduled radius is multiplied by epsilon, so pass r / epsilon or a schedule there for the same step.) With the default radius_decay = 1 the radius is fixed, so the mean settles in an O(step_radius) neighborhood of a minimizer. With radius_decay < 1, every tell! multiplies step_radius by radius_decay, also when the update is skipped. :mean recenters and rescales the costs every round, so near a minimizer the iterate keeps moving by a fixed fraction of step_radius (set by epsilon and dim, not by the distance to the minimizer).

ask!(es) -> (dim, popsize) Matrix

Candidate points as columns (column v + (p-1)*2dim = vertex v of particle p). Errors when called twice without an intervening tell!. The returned matrix is an internal buffer: do not mutate it (tell! reads it back), and copy it to keep it, since the next ask! overwrites it in place.

tell!(es, fitness::AbstractVector)

Update particles from the fitness of the last ask! (lower is better). The barycentric projection runs over the cached candidates (identical to the evaluated points, including clamping/repair), normalized by realized column mass. The particles stay when no fitness is finite or the update is not finite. Costs are weighted in the wider of T and eltype(fitness). Iterative solvers warm-start from the duals of the previous tell!. Tracks best_x/best_f/evals and applies radius_decay.

popsize(es)

Candidates per ask!: num_particles * 2 * dim.

PolyStepOptimizer(; num_particles=1, epsilon=0.5, step_radius=0.5, radius_decay=1.0,
                  solver=SoftmaxSolver(), scale_cost=:mean)

Optimization.jl algorithm (load Optimization or OptimizationBase). It runs a PolyStepES with these settings. maxiters counts ask/tell rounds, maxiters * num_particles * 2 * length(u0) evaluations plus num_particles for the final iterate; maxtime is checked after each round; with neither, 1000 rounds run. abstol/reltol are ignored. lb/ub clamp the candidates. MaxSense works. Constraints and repair are not supported. Extra solve keywords: batched = f_batch evaluates the flat (length(u0), N) candidate matrix in one call (otherwise the scalar objective runs once per candidate, with a warning), and rng (default Xoshiro(0), so repeated solves are identical). Return codes: Success, Terminated (callback), MaxTime, Failure (no finite objective).

PolyStepConfig(; dim, kwargs...)

Configuration for init_state, step! and solve!. The objective is batched: f(X) maps (dim, N) candidate columns to N costs (lower is better); wrap a scalar function with columnwise.

With small cost differences next to the temperature, a step is gradient descent with rate sr * (pr/2) / (dim * eps_ot * s) (s is the scale_cost divisor), so in high dimension lower ent_epsilon rather than epsilon.

Bounds clamp every evaluated point and box-project the iterate. repair runs after clamping on evaluated points only: the probes, the trust-region centers and the final iterate that solve! scores. It must keep them in the box. The iterate itself is not repaired. best_x/best_f hold the best evaluated point.

PolyStepState

Mutable optimizer state from init_state: particles X (d, P), the incumbent best_x/best_f, evals, iteration, and the reused step buffers.

init_state(ps::PolyStepConfig, X0::AbstractMatrix) -> PolyStepState

X0 is (dim, P): columns are particles; its element type sets the working precision (Float64 recommended for OR use; Float32 is the perf/GPU mode).

step!(f, ps::PolyStepConfig, st::PolyStepState; rng=Random.default_rng()) -> Float64

Run one PolyStep iteration. Return the mean vertex cost after sanitizing, or Inf when every vertex cost is non-finite. A vertex cost is the mean over its num_probe probes. f is the batched objective (d, N) -> (N,) and must return one cost per column.

solve!(f, ps::PolyStepConfig, st::PolyStepState; rng, callback=nothing) -> PolyStepState

Run up to ps.max_iterations steps. From step max(min_iterations, 3) on, stop when the mean squared step changes by less than threshold (relative) and is at most threshold times its peak. Stop after a step where every vertex cost is non-finite. callback(st) -> true stops early. The final iterate is then scored once (P evaluations, clamped and repaired) and can become st.best_x/st.best_f. Per-step diagnostics: st.costs, st.ot_converged, st.disp_sqnorms, st.trust_multipliers.

columnwise(f; parallel=:serial) -> batched objective

Adapt a scalar objective f(x::AbstractVector)::Real to the batched (d, N) -> (N,) protocol. parallel = :threads evaluates columns with Threads.@threads; off by default because OR simulators are frequently stateful and not thread-safe.

SoftmaxSolver()

One-sided softmax weighting: W = softmax(-C/eps, dims=1), plan = W .* a' (column sums equal a). No duals, single pass.

SinkhornSolver(; max_iterations=2000, threshold=3e-5, check_every=10,
               omega=1.0, anderson_depth=0, adaptive_omega=false,
               data_dependent_init=false)

Entropic OT on the (V, P) cost: min_P <C,P> + eps*sum(P .* (log.(P) .- 1)) with vertex marginal b (row sums) and particle marginal a (column sums), solved by log-domain alternating dual updates with optional overrelaxation omega in [0.5, 1.95], warm-started duals, Anderson acceleration (accepted only when it raises the dual), the residual-ratio adaptive omega (estimated while omega == 1), and a divergence back-off for omega > 1.5. It stops when the relative marginal error max|expm1(log(P'1) - log(a))| drops below threshold, or unconverged once the error stalls at the dtype's precision. threshold <= 0 runs a fixed number of iterations and reports converged when the result is finite. Zero-mass entries of a/b get zero mass and -Inf duals. ent_cost is the dual value in the sanitized, scaled cost frame.

An instance keeps a workspace between calls: do not call solve on one instance from two tasks at once (init_state and PolyStepES copy the solver). The solve is a CPU path.

KLSoftmaxSolver(; lam=Inf, max_iterations=2000, threshold=3e-5)

One-sided KL relaxation of min_P <C,P> + eps*H(P) + lam*KL(P'1 || b) with particle marginal a: a log-domain fixed point with alpha = lam/(lam+eps), exact f-update and damped g = alpha * g_target. This is the unbalanced Sinkhorn iteration, so it converges to the stated minimizer for every lam: lam=0 is SoftmaxSolver, lam=Inf is Sinkhorn, and intermediate lam soft-relaxes the vertex marginal toward b. It stops when the dual step max(|df|, |dg|/alpha)/eps is at most threshold, or unconverged once that step stalls at the dtype's precision. For lam > 0, zero-mass entries of a/b get zero mass and -Inf duals. last_marginal_violation holds KL(P'1 || b) of the last solve.

MinCostGreedySolver()

All of each particle's mass on its lowest-cost vertex. Non-finite costs rank last. A column with no finite spread (constant, or all non-finite) spreads uniformly, so the particle holds.

TopKMeanSolver(; k=3)

Uniform mass over each particle's k_eff = min(k, V) lowest-cost vertices. Non-finite costs rank last, and selected invalid vertices pass their mass to the valid ones. A column with no finite spread spreads uniformly, so the particle holds.

OTResult

plan::(V,P) transport (column p sums to ~ a[p] for one-sided solvers), ent_cost a per-solver cost proxy on the sanitized, scaled cost (Sinkhorn reports the dual value, the other solvers the transport cost <C,P>; it is not the regularized primal <C,P> + eps*KL), warm-startable duals f::(P,) / g::(V,) (nothing for non-iterative solvers), converged, iters, per-check marginal errors (Sinkhorn diagnostics, else nothing).

PolyStep.solve(solver, C, eps; a=nothing, b=nothing, f0=nothing, g0=nothing,
               scale_cost=nothing) -> OTResult

Weight the (V, P) cost matrix C (vertices by particles) at temperature eps. a (length P) and b (length V) are the particle and vertex marginals, uniform by default. Only SinkhornSolver and KLSoftmaxSolver use b and f0/g0 (dual warm start); the other solvers warn on a non-uniform b and ignore it. scale_cost (nothing, a positive divisor, :mean or :max) rescales C first. Non-finite costs rank last.

LinearEpsilon(; target=1e-3, init=1.0, decay=0.01)

eps(t) = max(init - decay*t, target).

CosineEpsilon(; target=1e-3, init=1.0, decay=0.01, total_steps=0, restart_mult=1.0)

eps(t) = target + 0.5*(init-target)*(1 + cos(pi*t/T)) with T = total_steps > 0 ? total_steps : max(1, ceil((init-target)/decay)) (so the target lands where LinearEpsilon's does), plus optional SGDR warm restarts (restart_mult > 1): cycle k = 0, 1, ... lasts T * restart_mult^k steps.

ProgressiveEpsilon(; init=1.0, target=0.01, max_epsilon=5.0, increase_factor=1.2,
                   decrease_factor=0.95, fast_threshold=0.1, slow_threshold=0.5,
                   ema_alpha=0.7, total_steps=0)

Feedback-driven epsilon (ProgOT-inspired); update! reacts to solver convergence. With the default total_steps=0, epsilon_at ignores the iteration and returns the EMA-smoothed value. With total_steps>0, epsilon_at(t) returns a decreasing cosine baseline from init to target over total_steps, multiplied by the bounded feedback factor smoothed/init, so epsilon sharpens over the run while the feedback still raises it when the solver struggles and lowers it when it converges fast. Incompatible with non-iterative solvers (they always report converged in 1 iteration, which would only ever shrink epsilon).

epsilon_at(schedule, t)

Value of a constant or a schedule at step t (0-based); t = nothing gives the initial value.

update!(s::ProgressiveEpsilon; n_iters, max_iterations, converged)

Struggling solve (non-converged or n_iters/max > slow_threshold) -> increase; fast solve (ratio < fast_threshold) -> decrease; then EMA-smooth and clamp.

orthoplex_vertices([T,] dim; radius=1) -> (dim, 2dim)

Cross-polytope template. Vertex order (FD extractors index by it): columns 1..d are +e_i, columns d+1..2d are -e_i.

simplex_vertices([T,] dim; radius=1) -> (dim, dim+1)

Regular simplex centered at the origin, unit-norm vertices (as in the Python get_simplex_vertices).

cube_vertices([T,] dim; radius=1) -> (dim, 2^dim)

Hypercube template, /sqrt(dim) normalized. Vertex column v encodes the 0-based index i = v-1; coordinate a reads bit a-1: bit 0 -> +1, bit 1 -> -1.

polytope_vertices([T,] polytope::Symbol, dim; radius=1)

Dispatch on :orthoplex | :simplex | :cube.

num_vertices(polytope, dim)

Vertex count: 2dim (:orthoplex), dim + 1 (:simplex), 2^dim (:cube).

probe_scales(T, K) -> Vector{T}

Fractional probe radii k/(K+1), k = 1..K. K=1 -> [0.5].

haar_rotations!(R, Z, rng) -> R

Fill R::(d,d,P) with Haar-uniform SO(d) rotations. Z::(d,d,P) is Gaussian scratch for d = 3..8 and may alias R (each slice is read fully before it is written). All Gaussians come serially from one rng and the per-slice work is deterministic given them, so results are bit-identical for any thread count.

d=2: analytic rotation from theta ~ U[0,2 * pi), one uniform per particle. d=3..8: StaticArrays QR of a Gaussian matrix. d>8: Stewart (1980): reflector j of the Householder QR of a Gaussian matrix is drawn directly from d-j+1 fresh Gaussians (its law under QR), then Q is accumulated from the reflectors. d>2: Mezzadri sign fix, then a det=+1 flip of column 1.

biased_rotation!(R, bias, Vt = nothing) -> R

Turn each rotation so that a polytope vertex points along bias[:, p]. Without a template (the orthoplex), one Householder reflection maps R[:, 1] onto the bias and a column sign flip restores det=+1; for Haar R the other columns stay uniform on the bias complement. With a template Vt (unit-norm vertices as columns), the vertex v nearest the bias, argmax_v (R'b)'v, is sent onto b by the minimal rotation R + ((1+2c)b - u)/(1+c) v' - (u+b)/(1+c) (R'b)' with b scaled to unit norm, u = Rv and c = u'b. A slice keeps its rotation when the squared norm of its bias is below 0.25 or not finite, or when its rotation is not finite.

ParamLayout(params)

Flat layout of a model's parameter arrays, given as name => shape pairs, shapes, or arrays.

ParamEntry(key, shape, offset, numel)

One parameter array in a ParamLayout.

LayerSpec

Coordinate block a HybridSubspace assigns to one parameter array.

HybridSubspace(layout; rank, seed = 0, max_subspace_dim = nothing, T = Float64)

Per-layer fixed orthonormal basis: min(d_out*r + r*d_in, numel) coordinates, with r = min(rank, d_in, d_out), per array with 2 or more dims; 1-D and full-width arrays are searched directly. For N-D arrays d_out is the last axis and d_in the product of the others, so a Flux/Lux conv weight (kw, kh, cin, cout) has d_out = cout. The count is symmetric in d_out and d_in, so a 2-D (out, in) weight gets the same count either way. max_subspace_dim caps the total by shrinking every array with 2 or more dims, full-width ones included (nothing means no cap); a cap below the 1-D size plus one per weight array warns and stops there.

Each projected layer stores a dense numel x ncoords basis of eltype T, so it takes numel * ncoords * sizeof(T) bytes. There is no sparse basis for large layers.

subspace_dim(s) -> Int

Total number of subspace coordinates.

compression_ratio(s) -> Float64

subspace_dim / total_params.

expand(s, z) -> Vector

Full-parameter displacement of subspace coordinates z.

expand!(out, s, z) -> out

Write the full-parameter displacement of subspace coordinates z into out.

project(s, x) -> Vector

Subspace coordinates of the full-parameter displacement x.

project!(z, s, x) -> z

Write the subspace coordinates of the full-parameter displacement x into z.

reconstruct_batch(s, base, Z) -> Matrix

Columns base .+ expand(s, Z[:, j]) for every column of Z.

subspace_objective(f, s, base) -> g

Batched objective in subspace coordinates: g(Z) = f(reconstruct_batch(s, base, Z)).

The coordinates form one subspace_dim-dimensional polytope; the Python optimizer instead splits them into 8-D particles, so the step shrinks as d grows: about as 1/d with a fixed cost divisor, about as 1/sqrt(d) at large epsilon with scale_cost = :mean. Lower the temperature if the steps stall: ent_epsilon in PolyStepConfig (its epsilon also scales both radii), or epsilon in PolyStepES/minimize.

splitmix64(seed, extras...) -> UInt64

Stable seed derivation (SplitMix64). Unlike Base.hash, output is documented stable across Julia versions and processes, safe for reproducible RNG streams, e.g. Xoshiro(splitmix64(master_seed, iteration)).

cuda_objective(f_gpu; T=Float32)

GPU-objective adapter, implemented in the CUDA extension (load CUDA.jl).