def get_sde()

in jcm/sde_lib.py [0:0]

• 30 lines of code
• 5 McCabe index (conditional complexity)

def get_sde(config):
    if config.training.sde.lower() == "vpsde":
        sde = VPSDE(
            beta_min=config.model.beta_min,
            beta_max=config.model.beta_max,
            N=config.model.num_scales,
        )
    elif config.training.sde.lower() == "subvpsde":
        sde = subVPSDE(
            beta_min=config.model.beta_min,
            beta_max=config.model.beta_max,
            N=config.model.num_scales,
        )
    elif config.training.sde.lower() == "vesde":
        sde = VESDE(
            sigma_min=config.model.sigma_min,
            sigma_max=config.model.sigma_max,
            N=config.model.num_scales,
        )
    elif config.training.sde.lower() == "kvesde":
        sde = KVESDE(
            t_min=config.model.t_min,
            t_max=config.model.t_max,
            N=config.model.num_scales,
            rho=config.model.rho,
            data_std=config.model.data_std,
        )
    else:
        raise NotImplementedError(f"SDE {config.training.sde} unknown.")

    return sde