Schedule Builder API¶
spectracles.optimise.schedule_builder.build_schedule(model, loss_fn, phases, params, default_optimizer='adam', managed=False, get_filter_spec_fn=get_opt_filter_spec, key=None)
¶
Build an optimization schedule using a parameter-centric specification.
This function provides a declarative way to specify which parameters should be free or fixed in each phase, with automatic pattern expansion for parameter paths.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
model
|
ShareModule
|
The ShareModule to optimize. |
required |
loss_fn
|
Callable[..., float]
|
Loss function to minimize. |
required |
phases
|
list[Union[PhaseSpec, PhaseSpecFull]]
|
List of phase specifications. Each is either: - (n_steps, learning_rate): Uses default optimizer - (n_steps, learning_rate, optimizer): Uses specified optimizer |
required |
params
|
dict[str, ParamSpec]
|
Dictionary mapping parameter patterns to ParamSpecs. Patterns support glob-style matching: - "" matches any single component - "*" matches any number of components |
required |
default_optimizer
|
Literal['sgd', 'adam']
|
Default optimizer type ("sgd" or "adam"). |
'adam'
|
managed
|
bool
|
If True, return a ManagedOptimiserSchedule with state tracking. |
False
|
get_filter_spec_fn
|
Callable[[ShareModule], Callable]
|
Function to get filter spec for optimization. |
get_opt_filter_spec
|
key
|
Optional[PRNGKeyArray]
|
PRNG key for random initializations. Required if any init_normal or init_uniform specs are used. |
None
|
Returns:
| Type | Description |
|---|---|
Union[OptimiserSchedule, ManagedOptimiserSchedule]
|
An OptimiserSchedule or ManagedOptimiserSchedule ready to run. |
Example
schedule = build_schedule( ... model, loss_fn, ... phases=[ ... (100, 0.1), # Phase 0: 100 steps, lr=0.1 ... (50, 0.01), # Phase 1: 50 steps, lr=0.01 ... (50, 0.001), # Phase 2: 50 steps, lr=0.001 ... ], ... params={ ... "gp.coefficients": free_in(0, 1, 2), ... "gp.kernel.*": free_after(1), ... "line.amplitude": free_in(1, 2) | init_normal(0), ... }, ... ) schedule.run_all(data)
Free/Fixed Helpers¶
spectracles.optimise.schedule_builder.free_in(*phases)
¶
Parameter is free (optimized) in the specified phases.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
*phases
|
int
|
Phase indices where the parameter should be free. |
()
|
Returns:
| Type | Description |
|---|---|
ParamSpec
|
ParamSpec with the specified free phases. |
Example
spec = free_in(0, 1, 2) # Free in phases 0, 1, and 2
spectracles.optimise.schedule_builder.free_after(phase)
¶
Parameter is free from the specified phase onwards.
Note: This creates a ParamSpec with a marker for "free from phase N onwards". The actual phases will be resolved when build_schedule is called.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
phase
|
int
|
Phase index from which the parameter should be free. |
required |
Returns:
| Type | Description |
|---|---|
ParamSpec
|
ParamSpec with a marker for late resolution. |
Example
spec = free_after(1) # Free in phases 1, 2, 3, ...
spectracles.optimise.schedule_builder.free_until(phase)
¶
Parameter is free up to and including the specified phase.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
phase
|
int
|
Last phase where the parameter should be free. |
required |
Returns:
| Type | Description |
|---|---|
ParamSpec
|
ParamSpec with the specified free phases. |
Example
spec = free_until(2) # Free in phases 0, 1, and 2
spectracles.optimise.schedule_builder.fixed_in(*phases)
¶
Parameter is fixed (not optimized) in the specified phases.
This is syntactic sugar - you specify which phases to FIX, and the parameter is free in all other phases.
Note: Requires knowing total number of phases at resolution time.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
*phases
|
int
|
Phase indices where the parameter should be fixed. |
()
|
Returns:
| Type | Description |
|---|---|
ParamSpec
|
ParamSpec with a marker for "fixed in these phases". |
Example
spec = fixed_in(0) # Fixed only in phase 0, free elsewhere
Initialization Helpers¶
spectracles.optimise.schedule_builder.init_normal(phase, mean=0.0, std=1.0)
¶
Initialize parameter to random normal values at the start of a phase.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
phase
|
int
|
Phase index at which to initialize. |
required |
mean
|
float
|
Mean of the normal distribution. |
0.0
|
std
|
float
|
Standard deviation of the normal distribution. |
1.0
|
Returns:
| Type | Description |
|---|---|
ParamSpec
|
ParamSpec with the initialization specification. |
Example
spec = free_in(0, 1) | init_normal(0, std=0.1)
spectracles.optimise.schedule_builder.init_value(phase, value)
¶
Initialize parameter to a specific value at the start of a phase.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
phase
|
int
|
Phase index at which to initialize. |
required |
value
|
float
|
The value to set. |
required |
Returns:
| Type | Description |
|---|---|
ParamSpec
|
ParamSpec with the initialization specification. |
Example
spec = free_after(1) | init_value(1, 1.5)
spectracles.optimise.schedule_builder.init_uniform(phase, low=0.0, high=1.0)
¶
Initialize parameter to random uniform values at the start of a phase.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
phase
|
int
|
Phase index at which to initialize. |
required |
low
|
float
|
Lower bound of the uniform distribution. |
0.0
|
high
|
float
|
Upper bound of the uniform distribution. |
1.0
|
Returns:
| Type | Description |
|---|---|
ParamSpec
|
ParamSpec with the initialization specification. |
Example
spec = free_in(0) | init_uniform(0, low=-1, high=1)