-
Notifications
You must be signed in to change notification settings - Fork 0
Problems
All optimization problems in dfbench represent gravitational-wave detector design tasks. The goal is to find optical component parameters (mirror reflectivities, laser power, cavity lengths, etc.) that minimize the detector's strain sensitivity across a frequency range of 20–5000 Hz.
Import:
from dfbench.problems import VoyagerProblem, VoyagerTuningProblem, ConstrainedVoyagerProblem, UIFOProblemContinuousProblem (ABC — core/problem.py)
└── OpticalSetupProblem (ABC — problems/base_problem.py)
├── VoyagerProblem
├── VoyagerTuningProblem
├── ConstrainedVoyagerProblem
└── UIFOProblem
Defines the minimal interface every problem must implement:
| Attribute / Method | Type | Description |
|---|---|---|
name |
str |
Human-readable identifier. |
objective_function |
Callable |
Loss in bounded parameter space. |
bounds |
Array[2, n_params] |
[lower_bounds, upper_bounds] for each parameter. |
optimization_pairs |
list[tuple[str, str]] |
(component_name, property_name) tuples mapping each parameter index to a Differometor component. |
n_params |
int |
Number of parameters = len(optimization_pairs). |
to_spec() → dict |
dict |
Reconstructive spec — a small, JSON-serialisable dict sufficient to rebuild an equivalent problem instance (see Reconstruction & Problem Spec below). |
Rationale — bounded problem contract: Problems expose the bounded loss only; Objective owns any mapping required by algorithms that search in unbounded coordinates.
Rationale — reconstructive spec: A checkpoint is only useful for resume or provenance if the originating problem can be rebuilt. to_spec() encodes the problem's constructor arguments so a saved run is fully self-describing (see Storage & Checkpointing).
Extends ContinuousProblem with optics-specific functionality shared by all Differometor problems:
-
Frequency grid: A log-spaced array of frequencies from 20 Hz to 5 kHz (configurable via
n_frequencies). -
Target sensitivity: Stored in
_target_sensitivities, computed from the reference detector design at initialization. -
calculate_sensitivity(params): Computes the sensitivity curve for a given parameter vector — used for plotting, not optimization. -
bounds_overrides: All concrete problems accept optional property-level bound overrides (narrowing only). -
signal_floor: All concrete problems floor detector signal magnitudes before sensitivity normalization. Defaults to1e-20. -
print_bounds(): Prints the effective per-parameter bounds currently used by the problem.
Note:
OpticalSetupProblemno longer has anoutput_to_filesmethod. Human-readable JSON/PNG output is now a derived view produced byRunDataExporterfrom aRunStatesnapshot (see Storage & Checkpointing). Keeping I/O on the problem was a responsibility violation — it mixed file layout, plotting, and timestamping into the mathematical problem definition.
| Property | Value |
|---|---|
| Setup | LIGO Voyager with balanced homodyne detection |
| Parameters | ~25 (reflectivities, tunings, squeezing, power, masses, lengths, phases) |
| Noise model | Single quantum noise source |
| Speed | ~12 ms/eval on A100 GPU |
| Difficulty | Moderate — loss < 0 achievable without physical constraints |
problem = VoyagerProblem(n_frequencies=100)problem = VoyagerProblem(
n_frequencies=100,
bounds_overrides={"tuning": (0, 45)},
signal_floor=1e-20,
)
problem.print_bounds()- The Voyager reference setup is simulated to get a target sensitivity curve.
- The target loss is
$\sum \log_{10}(\text{sensitivity}_{\text{target}})$ . - For a candidate parameter set, the loss is:
$$\text{loss} = \sum \log_{10}(\text{sensitivity}_{\text{candidate}}) - \text{target loss}$$ - Loss < 0 means the candidate has better sensitivity than the reference design.
Rationale for log-scale loss: Strain sensitivities span many orders of magnitude (
Each parameter corresponds to a (component, property) pair from the Voyager setup:
| Property | Bounds | Physical meaning |
|---|---|---|
reflectivity |
[0, 1] | Fraction of light reflected by a mirror |
tuning |
[0, 90] | Phase tuning of a mirror in degrees |
db |
[0.01, 20] | Squeezing level in decibels |
angle |
[-180, 180] | Squeezing angle in degrees |
power |
[0.01, 200] | Laser power in watts |
mass |
[0.01, 200] | Mirror suspension mass in kg |
length |
[1, 4000] | Cavity arm length in meters |
phase |
[-180, 180] | Phase offset in degrees |
VoyagerProblem does not enforce physical constraints (e.g. maximum mirror power absorption). A solution with loss < 0 may be physically unrealizable. For constrained optimization, use ConstrainedVoyagerProblem.
| Property | Value |
|---|---|
| Setup | LIGO Voyager with balanced homodyne detection |
| Parameters | 6 (tuning only: prm, itmy, etmy, itmx, etmx, srm) |
| Noise model | Single quantum noise source |
| Speed | ~12 ms/eval on A100 GPU |
| Difficulty | Moderate — lower-dimensional than VoyagerProblem, useful for quick prototyping |
problem = VoyagerTuningProblem(n_frequencies=100)problem = VoyagerTuningProblem(
n_frequencies=100,
bounds_overrides={"tuning": (0, 45)},
signal_floor=1e-20,
)
problem.print_bounds()- The Voyager reference setup is simulated to get a target sensitivity curve.
- For a candidate parameter set, the loss is: $$\text{loss} = \mathrm{mean}\left(\log_{10}\left(\frac{\text{sensitivity}{\text{candidate}}}{\text{sensitivity}{\text{target}}}\right)\right) $$
- Loss < 0 means the candidate has better average sensitivity than the reference design.
All optimized parameters are mirror tuning angles in degrees:
| Property | Bounds | Physical meaning |
|---|---|---|
tuning |
[-180, 180] | Phase tuning of selected Voyager optics (prm, itmy, etmy, itmx, etmx, srm) |
VoyagerTuningProblem does not enforce physical constraints (e.g. maximum mirror power absorption). For constrained optimization, use ConstrainedVoyagerProblem.
| Property | Value |
|---|---|
| Setup | LIGO Voyager with balanced homodyne detection |
| Parameters | ~25 (same as VoyagerProblem) |
| Noise model | Three sources: quantum, amplitude, frequency noise |
| Constraints | Power thresholds on mirrors, beamsplitters, and detectors |
| Speed | ~25 ms/eval on A100 GPU |
| Difficulty | Hard — loss < 0 is very difficult to achieve |
problem = ConstrainedVoyagerProblem(n_frequencies=100, signal_floor=1e-20)-
Realistic noise model: Uses three separate modulation modes (quantum noise, amplitude noise, frequency noise) and combines their contributions into a single sensitivity curve. This produces a more accurate picture of real detector performance.
-
Power constraints (ReLU-based): The loss includes a penalty term for violating optical power limits:
-
HARD_SIDE_POWER_THRESHOLD— maximum power on mirror/beamsplitter side ports -
SOFT_SIDE_POWER_THRESHOLD— softer limit with gradual penalty -
DETECTOR_POWER_THRESHOLD— maximum power on detector ports
For each component the configured
power_penalty_fn(value, threshold)is called and the results are summed. Three presets are provided:Preset Formula Import squashed_relu_penalty(default)$\frac{\max(v/t-1,,0)}{1+\max(v/t-1,,0)}$ from dfbench.problems import squashed_relu_penaltyrelu_penalty$\max(v/t-1,,0)$ from dfbench.problems import relu_penaltyzero_penalty$0$ from dfbench.problems import zero_penaltyYou can also pass any custom function with signature
fn(value, threshold) -> penalty:import jax.numpy as jnp def my_quadratic_penalty(value, threshold): relu = jnp.maximum(value / threshold - 1, 0) return relu ** 2 problem = ConstrainedVoyagerProblem(power_penalty_fn=my_quadratic_penalty)
The penalty function can also be swapped after the problem has been constructed (e.g. after wrapping it in an
Objective), viaObjective.set_penalty_fn(fn). This re-traces the problem's JIT-compiledobjective_functionand re-binds the Objective's cached evaluation callables, so the new penalty takes effect on subsequent evaluations. It must be called beforeObjective.start_logging():from dfbench import Objective from dfbench.problems import zero_penalty problem = ConstrainedVoyagerProblem() obj = Objective(problem) obj.set_penalty_fn(zero_penalty) # disable the penalty term obj.warmup_value() obj.start_logging()
Only problems with a power-constraint path accept
set_penalty_fn. The opt-in marker is the class attribute_supports_power_penalty;ConstrainedVoyagerProblemandUIFOProblemset it toTrue, whileVoyagerProblemandVoyagerTuningProblemleave itFalseand raiseRuntimeErrorif the call is attempted, even though they inherit the method from the optical base class. -
Rationale (penalty squashing): A raw penalty can become orders of magnitude larger than the sensitivity loss, making gradient-based optimizers ignore sensitivity entirely. The default squashed_relu_penalty bounds the penalty contribution while preserving its gradient direction.
It could very well be that other penalty functions work better for certain algorithms (or even Adam). Feel free to play around!
The constrained problems also expose a JIT-compiled objective_function_aux(params) alongside objective_function. It returns (loss, aux) where aux is a pytree dict carrying the loss decomposition and physical diagnostics for that eval:
| Key | Shape | Description |
|---|---|---|
sensitivity_loss |
scalar | The unpenalised sensitivity loss. |
penalty |
scalar | The summed penalty contribution. |
is_feasible |
scalar bool |
True iff every per-group power is at or below its threshold. This is a physical check, independent of the active power_penalty_fn preset, so it stays meaningful even when the penalty is disabled with zero_penalty. |
violations |
(n_constraints, n_freq) |
Per-constraint penalty values. |
power_values |
dict with hard, soft, detector leaves |
Raw per-group power arrays. |
Because aux is a JAX pytree, objective_function_aux vmapps cleanly: a batched call adds a leading batch dim to every leaf, including the power_values sub-arrays.
The Objective wraps this with value_aux, value_and_grad_aux, vmap_value_aux, and vmap_value_and_grad_aux. These thread aux through logging and the save-token system: each aux field has its own token (sensitivity_loss, penalty, is_feasible, power_values, violations) plus a batched_* variant and aux / batched_aux convenience aliases. When a batched_* token is off and the non-batched token is on, batched aux entries are reduced to the representative point picked by the loss minimum, so the recorded is_feasible and violations reflect that best point. Objective.best_is_feasible reports the feasibility of the best-loss point from that recorded history. Full reference in the Objective API Reference.
| Property | Value |
|---|---|
| Setup | Quasi-Universal Interferometer (UIFO) |
| Parameters | 50–250+ depending on grid size |
| Noise model | Three sources (same as constrained Voyager) |
| Constraints | Power thresholds (same as constrained Voyager) |
| Speed | ~500 ms/eval on A100 GPU |
| Difficulty | Hard but achievable — the UIFO is overparameterized |
# From a topology seed (random topology, deterministic from seed)
problem = UIFOProblem(size=3, n_frequencies=100, topology_seed=42)
# From a compact topology string
problem = UIFOProblem(size=3, topology="AECGCCHEG-SLLSSHLLLLS")
# From explicit dicts
problem = UIFOProblem(
size=3,
centers={"11": ("beamsplitter", "left"), ...},
boundaries={"01": "squeezer", ...},
)| Parameter | Default | Description |
|---|---|---|
size |
3 |
Grid dimensions (3 = 3×3). Larger grids have more components and parameters. |
n_frequencies |
100 |
Frequency points for sensitivity calculation. |
topology_seed |
42 |
Seed for random topology generation. Set to None (with no other topology args) to generate a truly random topology. The seed is always printed to the console. Mutually exclusive with topology and centers/boundaries. |
topology |
None |
Compact topology string (see below). Mutually exclusive with topology_seed. |
centers |
None |
Interior cell dict. Must be paired with boundaries. Mutually exclusive with topology_seed and topology. |
boundaries |
None |
Boundary cell dict. Must be paired with centers. Mutually exclusive with topology_seed and topology. |
power_penalty_fn |
squashed_relu_penalty |
Per-element penalty function fn(value, threshold). See presets above. |
signal_floor |
1e-20 |
Lower floor for detector signal magnitudes before sensitivity normalization. |
Backwards compatibility:
RandomUIFOProblemis an alias forUIFOProblem.
There are three mutually exclusive ways to specify a UIFO topology:
-
topology_seed— The simplest option. A random topology is generated deterministically from the seed (default:42). Passtopology_seed=Nonewith no other topology arguments to generate a truly random topology — the seed is printed so you can reproduce it. -
topologystring — A compact encoding ideal for configs, papers, and sharing. Uses single-character codes:-
Interior cells:
A–D= beamsplitter (left/right/top/bottom),E–H= directional_beamsplitter (left/right/top/bottom) -
Boundary cells:
L= laser,S= squeezer,D= detector,H= balanced_homodyne - Format:
"<interior_chars>-<boundary_chars>"in row-major order.
-
Interior cells:
-
centers+boundariesdicts — Explicit component placement, matching Differometor’s native format.
Conversion helpers are available:
from dfbench.problems.uifo import topology_to_string, topology_from_string
topology_str = topology_to_string(centers, boundaries, size=3)
centers, boundaries = topology_from_string(topology_str, size=3)A Quasi-Universal Interferometer Field Optimization (UIFO) is a grid-based interferometer where beamsplitters, mirrors, lasers, and squeezers are placed on a grid and connected by spaces. The topology (which components are placed where and how they connect) is generated randomly from topology_seed (printed on initialization for reproducibility). Once the topology is fixed, only the continuous parameters (reflectivities, tunings, lengths, etc.) are optimized.
Rationale — coupled grid-cell spaces: Horizontal and vertical spaces at the same grid positions are constrained to have equal lengths via constrain_inter_grid_cell_spaces(). This preserves the physical grid structure and prevents the optimizer from "folding" the interferometer into a degenerate geometry.
The reference sensitivity target is always the Voyager detector. Since the UIFO is overparameterized (many more degrees of freedom than Voyager), it can in principle achieve better sensitivity but the large parameter space makes optimization harder.
Every problem implements to_spec() → dict, which returns a small, JSON-serialisable dict capturing everything needed to rebuild an equivalent instance in a separate process. This is the reconstructive contract that makes checkpoints self-describing.
Starting with dfbench 0.1.1, the raw to_spec() dict is wrapped in a typed ProblemSpec container (dfbench.core.problem.ProblemSpec) that carries an explicit schema version and a separated params field. Checkpoints embed the container (ProblemSpec.to_dict() → {"type", "version", "params"}) in RunMetadata.extra["problem_spec"], so consumers get a stable, schema-validated identity instead of an untyped dict. Legacy flat specs ({"type", <kwargs>}) written by older versions are still accepted on load via ProblemSpec.from_dict.
- Each problem subclass implements
to_spec(), returning a dict with a"type"key (the registry name) plus its constructor arguments. - The
@register_problemdecorator registers the class in a module-level registry under its__name__(or a customspec_type). -
ContinuousProblem.to_problem_spec()wraps theto_spec()dict into a typedProblemSpeccontainer. Subclasses rarely need to override this; the default implementation is sufficient as long asto_spec()is correct. -
build_problem_from_spec(spec)accepts either aProblemSpecor a raw dict (typed container or legacy flat form; both are normalized viaProblemSpec.from_dict) and reconstructs the instance. -
Objective._build_metadata()callsproblem.to_problem_spec()and stores the resulting dict inRunMetadata.extra["problem_spec"], so every checkpoint records its originating problem.
from dfbench.core.problem import ProblemSpec, build_problem_from_spec
# A spec captured from a live problem
ps = problem.to_problem_spec()
# ProblemSpec(type="VoyagerProblem", params={"n_frequencies": 100, ...}, version=1)
# JSON-safe dict for embedding in checkpoint metadata
spec_dict = ps.to_dict()
# {"type": "VoyagerProblem", "version": 1, "params": {"n_frequencies": 100, ...}}
# Rebuild an equivalent problem later, in any process
problem2 = build_problem_from_spec(ps)
# or, equivalently, from the dict form:
problem2 = build_problem_from_spec(spec_dict)| Field | Type | Description |
|---|---|---|
type |
str |
Registry key matching a @register_problem-decorated class |
params |
dict[str, Any] |
Constructor keyword arguments forwarded to the problem class on reconstruction |
version |
int |
Container schema version (defaults to PROBLEM_SPEC_VERSION = 1); governs the type/version/params layout, not the per-problem constructor args |
ProblemSpec.from_dict accepts both the typed container and the legacy flat form, so checkpoints written before the typed container existed still load. ProblemSpec.__post_init__ validates that type is a non-empty string, params is a dict, and version is an int. A malformed or tampered spec becomes a deterministic ValueError at the trust boundary instead of a silent corruption downstream.
The params sub-dict is whatever each problem's to_spec() returns minus the "type" key:
| Problem |
params fields |
Reconstruction path |
|---|---|---|
VoyagerProblem |
n_frequencies, power_penalty_fn, bounds_overrides
|
Direct constructor call |
VoyagerTuningProblem |
n_frequencies, power_penalty_fn, bounds_overrides
|
Direct constructor call |
ConstrainedVoyagerProblem |
n_frequencies, power_penalty_fn, bounds_overrides
|
Direct constructor call |
UIFOProblem |
size, n_frequencies, topology (string), power_penalty_fn, bounds_overrides
|
Rebuilt from explicit topology string (deterministic, RNG-independent) |
Callables like power_penalty_fn are encoded by name via a registry of presets (squashed_relu_penalty, relu_penalty, zero_penalty). This keeps the spec JSON-safe. Custom penalty functions that are not registered presets will raise on to_spec() — register them or use the built-in presets.
Reconstruction is a two-step process that crosses the storage/problem layer boundary:
from dfbench.core.storage import CheckpointManager
from dfbench.core.problem import ProblemSpec, build_problem_from_spec
state = manager.load(path)
spec_dict = CheckpointManager.extract_problem_spec(state) # -> dict | None
if spec_dict is not None:
ps = ProblemSpec.from_dict(spec_dict) # typed container (accepts legacy flat too)
problem = build_problem_from_spec(ps) # or pass spec_dict directlyCheckpointManager.extract_problem_spec returns None if the run did not record a problem spec (e.g. the problem did not implement to_spec). The relevant problem module must be imported so its class is registered.
If you add a new ContinuousProblem subclass:
- Decorate it with
@register_problem(imported fromdfbench.core.problemordfbench.problems.base_problem). - Implement
to_spec()returning a dict with"type"(the class name) plus every constructor argument needed forbuild_problem_from_specto produce an equivalent instance. - Encode any callables by name against a registry (see the penalty-function pattern in
base_problem.py).
You do not need to override to_problem_spec(); the default implementation wraps to_spec() into the typed container automatically.
See Storage & Checkpointing for how the spec is embedded in checkpoints.
Artificial Scientist Lab | Website |University of Tübingen
Department of Computer Science
| Read our Documentation | Contact: laurin.sefa@student.uni-tuebingen.de, mario.krenn@uni-tuebingen.de, soham.basu@uni-tuebingen.de
Getting Started
Core API
Benchmarking
Contributing
Reference