pdft.training.cosine_with_warmup

pdft.training.cosine_with_warmup(step, total_steps, *, warmup_frac=0.05, lr_peak=0.01, lr_final=0.001)

Linear warmup followed by cosine decay.

Mirror of ParametricDFT.jl/src/training.jl::_cosine_with_warmup. step is 0-indexed conceptually but Julia uses 1-indexed; we match Julia’s behavior: the warmup ramp ends exactly at step == warmup_steps where warmup_steps = max(1, round(warmup_frac * total_steps)).

Parameters:

• step (int)

• total_steps (int)

• warmup_frac (float)

• lr_peak (float)

• lr_final (float)

Return type:

float