def cyclical_cosine_schedule_with_const_burnin(
    init_step_size: float, burnin_steps: int, cycle_length: int
) -> StepSchedule:
    """Create a cyclical cosine schedule with constant burn-in.

    Parameters
    ----------
    init_step_size: float
        The initial step size.
    burnin_steps: int
        The length of burn-in, in steps.
    cycle_length: int
        The length of the cosine cycle, in steps.

    Returns
    -------
    schedule_fn: StepSchedule
    """
    if not init_step_size >= 0:
        raise ValueError("`init_step_size` should be >= 0.")
    if not burnin_steps >= 0:
        raise ValueError("`burnin_steps` should be >= 0.")
    if not cycle_length >= 0:
        raise ValueError("`cycle_length` should be >= 0.")

    def schedule(step: Array):
        t = jnp.maximum(step - burnin_steps - 1, 0.0)
        t = (t % cycle_length) / cycle_length
        return 0.5 * init_step_size * (1 + jnp.cos(t * np.pi))

    return schedule
Create a cyclical cosine schedule with constant burn-in.

    Parameters
    ----------
    init_step_size: float
        The initial step size.
    burnin_steps: int
        The length of burn-in, in steps.
    cycle_length: int
        The length of the cosine cycle, in steps.

    Returns
    -------
    schedule_fn: StepSchedule
    
fortuna/prob_model/posterior/sgmcmc/sgmcmc_step_schedule.py