GitHub

OptimisersExt

Optimization rules and proximal operators for constrained inverse design.

Overview

The OptimisersExt module extends Optimisers.jl with:

  • Custom optimization rules: FISTA, NoDescent
  • Proximal operators: Constraints and regularization (TV, sparsity, clamping)
  • Per-parameter rules: Different optimizers for different components
  • Proximal-gradient methods: Combine optimization with constraints

Quick Example

using FluxOptics

# Define source
xv, yv = spatial_vectors(128, 128, 1.0, 1.0)
u = ScalarField(Gaussian(20.0)(xv, yv), (1.0, 1.0), 1.064)
source = ScalarSource(u)

# Define optical components
phasemask = Phase(u, (x, y) -> 0.0; trainable=true)
mask = FourierMask(u, (fx, fy) -> 1.0; trainable=true)

# Define optical system
system = source |> phasemask |> mask

# Per-component optimization rules
rules = make_rules(
    phasemask => ProxRule(Descent(0.01), ClampProx(-π, π)),  # Constrained phase
    mask => Momentum(0.1, 0.9)                               # Momentum for mask
)

# Setup
opt_state = setup(rules, system)
(source = (u0 = (),), components = (optical_components = ((ϕ = Leaf(ProxRule{Descent{Float64}, PointwiseProx{var"#5#6"{Float64, Irrational{:π}}}}(Descent(0.01), PointwiseProx{var"#5#6"{Float64, Irrational{:π}}}(#5)), (nothing, ())),), (optical_components = ((p_f = (), s = (), direct = ()), (m = Leaf(Momentum(eta=0.1, rho=0.9), ComplexF64[0.0+0.0im 0.0+0.0im … 0.0+0.0im 0.0+0.0im; 0.0+0.0im 0.0+0.0im … 0.0+0.0im 0.0+0.0im; … ; 0.0+0.0im 0.0+0.0im … 0.0+0.0im 0.0+0.0im; 0.0+0.0im 0.0+0.0im … 0.0+0.0im 0.0+0.0im]),), (p_f = (), s = (), direct = ())),)),))

Key Types

  • ProxRule: Combine optimizer with proximal operator
  • Fista: Fast iterative shrinkage-thresholding
  • NoDescent: No-op optimizer for fixed parameters

Proximal Operators

Key Functions

  • make_rules: Create per-parameter optimization rules
  • setup: Initialize optimization state with custom rules
  • update!: Apply optimization step (from Optimisers.jl)

See Also

Index