SYMBOL INDEX (249 symbols across 20 files)

FILE: bin/guidance_experiment.py
  function load_model (line 32) | def load_model(config):
  function evaluate_guidance (line 54) | def evaluate_guidance(
  function main (line 126) | def main(config: dict, log_dir: str):

FILE: bin/refinement_experiment.py
  function load_model (line 41) | def load_model(config):
  function get_best_diffusion_step (line 63) | def get_best_diffusion_step(model: TSDiff, data_loader, device):
  function train_and_forecast_base_model (line 81) | def train_and_forecast_base_model(dataset, base_model_name, config):
  function forecast_guidance (line 124) | def forecast_guidance(
  function main (line 158) | def main(config: dict, log_dir: str):

FILE: bin/train_cond_model.py
  function create_model (line 33) | def create_model(config):
  function evaluate_conditional (line 50) | def evaluate_conditional(
  function main (line 101) | def main(config, log_dir):

FILE: bin/train_model.py
  function create_model (line 36) | def create_model(config):
  function evaluate_guidance (line 52) | def evaluate_guidance(
  function main (line 123) | def main(config, log_dir):

FILE: bin/tstr_experiment.py
  function load_model (line 44) | def load_model(config):
  function sample_synthetic (line 66) | def sample_synthetic(
  function sample_real (line 83) | def sample_real(
  function evaluate_tstr (line 108) | def evaluate_tstr(
  function train_and_evaluate (line 143) | def train_and_evaluate(
  function main (line 214) | def main(config: dict, log_dir: str, samples_path: str):

FILE: src/uncond_ts_diff/arch/backbones.py
  class SinusoidalPositionEmbeddings (line 11) | class SinusoidalPositionEmbeddings(nn.Module):
    method __init__ (line 12) | def __init__(self, dim):
    method forward (line 16) | def forward(self, time):
  class S4Layer (line 28) | class S4Layer(nn.Module):
    method __init__ (line 29) | def __init__(
    method forward (line 48) | def forward(self, x):
    method default_state (line 63) | def default_state(self, *args, **kwargs):
    method step (line 66) | def step(self, x, state, **kwargs):
  class S4Block (line 77) | class S4Block(nn.Module):
    method __init__ (line 78) | def __init__(self, d_model, dropout=0.0, expand=2, num_features=0):
    method forward (line 93) | def forward(self, x, t, features=None):
  function Conv1dKaiming (line 105) | def Conv1dKaiming(in_channels, out_channels, kernel_size):
  class BackboneModel (line 111) | class BackboneModel(nn.Module):
    method __init__ (line 112) | def __init__(
    method forward (line 155) | def forward(self, input, t, features=None):

FILE: src/uncond_ts_diff/arch/s4.py
  function get_logger (line 18) | def get_logger(name=__name__, level=logging.INFO) -> logging.Logger:
  function _broadcast_dims (line 60) | def _broadcast_dims(*tensors):
  function cauchy_conj (line 68) | def cauchy_conj(v, z, w):
  function log_vandermonde (line 92) | def log_vandermonde(v, x, L):
  function log_vandermonde_transpose (line 114) | def log_vandermonde_transpose(u, v, x, L):
  function cauchy_naive (line 154) | def cauchy_naive(v, z, w):
  function log_vandermonde (line 170) | def log_vandermonde(v, x, L):
  function log_vandermonde_transpose (line 184) | def log_vandermonde_transpose(u, v, x, L):
  function _conj (line 197) | def _conj(x):
  function _resolve_conj (line 205) | def _resolve_conj(x):
  function _resolve_conj (line 210) | def _resolve_conj(x):
  function Activation (line 217) | def Activation(activation=None, dim=-1):
  function LinearActivation (line 238) | def LinearActivation(
  class DropoutNd (line 261) | class DropoutNd(nn.Module):
    method __init__ (line 262) | def __init__(self, p: float = 0.5, tie=True, transposed=True):
    method forward (line 277) | def forward(self, X):
  function power (line 296) | def power(L, A, v=None):
  function transition (line 345) | def transition(measure, N):
  function rank_correction (line 412) | def rank_correction(measure, N, rank=1, dtype=torch.float):
  function nplr (line 449) | def nplr(measure, N, rank=1, dtype=torch.float, diagonalize_precision=Tr...
  function dplr (line 518) | def dplr(
  function ssm (line 595) | def ssm(measure, N, R, H, **ssm_args):
  function combination (line 628) | def combination(measures, N, R, S, **ssm_args):
  class OptimModule (line 650) | class OptimModule(nn.Module):
    method register (line 653) | def register(self, name, tensor, lr=None):
  class SSKernelNPLR (line 667) | class SSKernelNPLR(OptimModule):
    method _setup_C (line 671) | def _setup_C(self, L):
    method _omega (line 707) | def _omega(self, L, dtype, device, cache=True):
    method __init__ (line 731) | def __init__(
    method _w_init (line 811) | def _w_init(self, w_real):
    method _w (line 828) | def _w(self):
    method forward (line 845) | def forward(self, state=None, rate=1.0, L=None):
    method _setup_linear (line 981) | def _setup_linear(self):
    method _step_state_linear (line 1022) | def _step_state_linear(self, u=None, state=None):
    method _setup_state (line 1078) | def _setup_state(self):
    method _step_state (line 1101) | def _step_state(self, u, state):
    method _setup_step (line 1108) | def _setup_step(self, mode="dense"):
    method default_state (line 1157) | def default_state(self, *batch_shape):
    method step (line 1199) | def step(self, u, state):
  class SSKernelDiag (line 1210) | class SSKernelDiag(OptimModule):
    method __init__ (line 1213) | def __init__(
    method _A_init (line 1254) | def _A_init(self, A_real):
    method _A (line 1271) | def _A(self):
    method forward (line 1289) | def forward(self, L, state=None, rate=1.0, u=None):
    method _setup_step (line 1368) | def _setup_step(self):
    method default_state (line 1390) | def default_state(self, *batch_shape):
    method step (line 1397) | def step(self, u, state):
    method forward_state (line 1404) | def forward_state(self, u, state):
  class SSKernel (line 1413) | class SSKernel(nn.Module):
    method __init__ (line 1423) | def __init__(
    method forward (line 1548) | def forward(self, state=None, L=None, rate=1.0):
    method forward_state (line 1552) | def forward_state(self, u, state):
    method _setup_step (line 1582) | def _setup_step(self, **kwargs):
    method step (line 1590) | def step(self, u, state, **kwargs):
    method default_state (line 1594) | def default_state(self, *args, **kwargs):
  class S4 (line 1598) | class S4(nn.Module):
    method __init__ (line 1599) | def __init__(
    method forward (line 1721) | def forward(self, u, state=None, rate=1.0, lengths=None, **kwargs):
    method setup_step (line 1807) | def setup_step(self, **kwargs):
    method step (line 1810) | def step(self, u, state):
    method default_state (line 1829) | def default_state(self, *batch_shape, device=None):
    method d_output (line 1835) | def d_output(self):

FILE: src/uncond_ts_diff/dataset.py
  function get_gts_dataset (line 15) | def get_gts_dataset(dataset_name):

FILE: src/uncond_ts_diff/metrics/linear_pred_score.py
  function linear_pred_score (line 20) | def linear_pred_score(

FILE: src/uncond_ts_diff/model/callback.py
  class GradNormCallback (line 21) | class GradNormCallback(Callback):
    method __init__ (line 22) | def __init__(self) -> None:
    method on_before_optimizer_step (line 25) | def on_before_optimizer_step(
    method grad_norm (line 36) | def grad_norm(self, parameters):
  class PredictiveScoreCallback (line 48) | class PredictiveScoreCallback(Callback):
    method __init__ (line 49) | def __init__(
    method _generate_real_samples (line 72) | def _generate_real_samples(
    method _generate_synth_samples (line 104) | def _generate_synth_samples(
    method on_train_epoch_end (line 117) | def on_train_epoch_end(self, trainer, pl_module):
  class EvaluateCallback (line 169) | class EvaluateCallback(Callback):
    method __init__ (line 170) | def __init__(
    method on_train_epoch_end (line 208) | def on_train_epoch_end(self, trainer, pl_module):

FILE: src/uncond_ts_diff/model/diffusion/_base.py
  class TSDiffBase (line 28) | class TSDiffBase(pl.LightningModule):
    method __init__ (line 29) | def __init__(
    method _extract_features (line 97) | def _extract_features(self, data):
    method configure_optimizers (line 100) | def configure_optimizers(self):
    method log (line 107) | def log(self, name, value, **kwargs):
    method get_logs (line 116) | def get_logs(self):
    method q_sample (line 121) | def q_sample(self, x_start, t, noise=None):
    method p_losses (line 137) | def p_losses(
    method p_sample (line 167) | def p_sample(self, x, t, t_index, features=None):
    method p_sample_ddim (line 187) | def p_sample_ddim(self, x, t, features=None, noise=None):
    method p_sample_genddim (line 209) | def p_sample_genddim(
    method sample (line 266) | def sample(self, noise, features=None):
    method fast_denoise (line 283) | def fast_denoise(self, xt, t, features=None, noise=None):
    method forward (line 294) | def forward(self, x, mask):
    method training_step (line 297) | def training_step(self, data, idx):
    method training_epoch_end (line 314) | def training_epoch_end(self, outputs):
    method validation_step (line 320) | def validation_step(self, data, idx):
    method validation_epoch_end (line 335) | def validation_epoch_end(self, outputs):

FILE: src/uncond_ts_diff/model/diffusion/tsdiff.py
  class TSDiff (line 13) | class TSDiff(TSDiffBase):
    method __init__ (line 14) | def __init__(
    method _extract_features (line 74) | def _extract_features(self, data):
    method sample_n (line 133) | def sample_n(
    method on_train_batch_end (line 156) | def on_train_batch_end(self, outputs, batch, batch_idx):
  function update_ema (line 161) | def update_ema(target_state_dict, source_state_dict, rate=0.99):

FILE: src/uncond_ts_diff/model/diffusion/tsdiff_cond.py
  class TSDiffCond (line 15) | class TSDiffCond(TSDiffBase):
    method __init__ (line 16) | def __init__(
    method _extract_features (line 73) | def _extract_features(self, data):
    method step (line 136) | def step(self, x, t, features, loss_mask):
    method training_step (line 158) | def training_step(self, data, idx):
    method validation_step (line 177) | def validation_step(self, data, idx):
    method forecast (line 199) | def forecast(self, observation, observation_mask, features=None):
    method forward (line 220) | def forward(
    method get_predictor (line 270) | def get_predictor(self, input_transform, batch_size=40, device=None):

FILE: src/uncond_ts_diff/model/linear/_estimator.py
  function stack (line 40) | def stack(data):
  function batchify (line 48) | def batchify(data: List[dict]):
  class LinearModel (line 54) | class LinearModel:
    method __init__ (line 55) | def __init__(self, weight, bias, scaler, num_parallel_samples=100) -> ...
    method _linear (line 62) | def _linear(self, x, A, b):
    method __call__ (line 65) | def __call__(self, x, mask):
  function _ (line 73) | def _(prediction_net, args) -> np.ndarray:
  class LinearPredictor (line 77) | class LinearPredictor(Predictor):
    method __init__ (line 78) | def __init__(
    method predict (line 95) | def predict(self, dataset: Dataset, num_samples: Optional[int] = None):
  class LinearEstimator (line 112) | class LinearEstimator(Estimator):
    method __init__ (line 149) | def __init__(
    method create_transformation (line 174) | def create_transformation(self) -> Transformation:
    method _create_instance_splitter (line 188) | def _create_instance_splitter(self, mode: str):
    method _create_training_samples (line 213) | def _create_training_samples(self, training_data) -> np.ndarray:
    method create_predictor (line 242) | def create_predictor(self, transformation, model):
    method train (line 252) | def train(

FILE: src/uncond_ts_diff/model/linear/_scaler.py
  class MeanScaler (line 8) | class MeanScaler:
    method __init__ (line 11) | def __init__(
    method __call__ (line 24) | def __call__(
  class NOPScaler (line 63) | class NOPScaler:
    method __init__ (line 68) | def __init__(self, axis: int, keepdims: bool = False):
    method __call__ (line 73) | def __call__(

FILE: src/uncond_ts_diff/predictor.py
  class PyTorchPredictorWGrads (line 12) | class PyTorchPredictorWGrads(PyTorchPredictor):
    method predict (line 13) | def predict(

FILE: src/uncond_ts_diff/sampler/_base.py
  function grad_fn (line 10) | def grad_fn(fn, x):
  function langevin_dynamics (line 16) | def langevin_dynamics(
  function leapfrog (line 67) | def leapfrog(
  function hmc (line 106) | def hmc(
  function linear_midpoint_em_step (line 150) | def linear_midpoint_em_step(
  function udld (line 161) | def udld(

FILE: src/uncond_ts_diff/sampler/observation_guidance.py
  class Guidance (line 23) | class Guidance(torch.nn.Module):
    method __init__ (line 26) | def __init__(
    method quantile_loss (line 47) | def quantile_loss(self, y_prediction, y_target):
    method energy_func (line 64) | def energy_func(self, y, t, observation, observation_mask, features):
    method score_func (line 79) | def score_func(self, y, t, observation, observation_mask, features):
    method scale_func (line 87) | def scale_func(self, y, t, base_scale):
    method guide (line 90) | def guide(self, observation, observation_mask, features, scale):
    method forward (line 93) | def forward(
    method get_predictor (line 171) | def get_predictor(self, input_transform, batch_size=40, device=None):
  class DDPMGuidance (line 182) | class DDPMGuidance(Guidance):
    method __init__ (line 183) | def __init__(
    method scale_func (line 203) | def scale_func(self, y, t, base_scale):
    method _reverse_diffusion (line 207) | def _reverse_diffusion(
    method guide (line 228) | def guide(self, observation, observation_mask, features, base_scale):
  class DDIMGuidance (line 234) | class DDIMGuidance(Guidance):
    method __init__ (line 237) | def __init__(
    method scale_func (line 264) | def scale_func(self, y, t, base_scale):
    method _get_timesteps (line 270) | def _get_timesteps(self):
    method _reverse_ddim (line 286) | def _reverse_ddim(
    method guide (line 322) | def guide(self, observation, observation_mask, features, base_scale):

FILE: src/uncond_ts_diff/sampler/refiner.py
  class Refiner (line 27) | class Refiner(torch.nn.Module):
    method __init__ (line 28) | def __init__(
    method quantile_loss (line 49) | def quantile_loss(self, y_prediction, y_target):
    method prior (line 66) | def prior(self, y_prediction, obs, obs_mask):
    method refine (line 79) | def refine(self, observation, observation_mask):
    method forward (line 82) | def forward(
    method get_predictor (line 190) | def get_predictor(self, input_transform, batch_size=40, device=None):
  class MostLikelyRefiner (line 201) | class MostLikelyRefiner(Refiner):
    method __init__ (line 202) | def __init__(
    method _most_likely (line 228) | def _most_likely(self, observation, observation_mask):
    method refine (line 254) | def refine(self, observation, observation_mask):
  class MCMCRefiner (line 258) | class MCMCRefiner(Refiner):
    method __init__ (line 261) | def __init__(
    method _mcmc (line 290) | def _mcmc(self, observation, observation_mask):
    method refine (line 355) | def refine(self, observation, observation_mask):

FILE: src/uncond_ts_diff/utils.py
  function filter_metrics (line 47) | def filter_metrics(metrics, select={"ND", "NRMSE", "mean_wQuantileLoss"}):
  function extract (line 51) | def extract(a, t, x_shape):
  function cosine_beta_schedule (line 57) | def cosine_beta_schedule(timesteps, s=0.008):
  function linear_beta_schedule (line 71) | def linear_beta_schedule(timesteps):
  function plot_train_stats (line 77) | def plot_train_stats(df: pd.DataFrame, y_keys=None, skip_first_epoch=True):
  function get_lags_for_freq (line 98) | def get_lags_for_freq(freq_str: str):
  function create_transforms (line 116) | def create_transforms(
  function create_splitter (line 189) | def create_splitter(past_length: int, future_length: int, mode: str = "t...
  function get_next_file_num (line 214) | def get_next_file_num(
  function str2bool (line 255) | def str2bool(v):
  function add_config_to_argparser (line 266) | def add_config_to_argparser(config: Dict, parser: ArgumentParser):
  class AddMeanAndStdFeature (line 280) | class AddMeanAndStdFeature(MapTransformation):
    method __init__ (line 282) | def __init__(
    method map_transform (line 292) | def map_transform(self, data: DataEntry, is_train: bool) -> DataEntry:
  class ScaleAndAddMeanFeature (line 300) | class ScaleAndAddMeanFeature(MapTransformation):
    method __init__ (line 301) | def __init__(
    method map_transform (line 322) | def map_transform(self, data, is_train: bool):
  class ScaleAndAddMinMaxFeature (line 336) | class ScaleAndAddMinMaxFeature(MapTransformation):
    method __init__ (line 337) | def __init__(
    method map_transform (line 358) | def map_transform(self, data, is_train: bool):
  function descale (line 371) | def descale(data, scale, scaling_type):
  function predict_and_descale (line 381) | def predict_and_descale(predictor, dataset, num_samples, scaling_type):
  function to_dataframe_and_descale (line 420) | def to_dataframe_and_descale(input_label, scaling_type) -> pd.DataFrame:
  function make_evaluation_predictions_with_scaling (line 447) | def make_evaluation_predictions_with_scaling(
  class PairDataset (line 489) | class PairDataset(Dataset):
    method __init__ (line 490) | def __init__(self, x, y) -> None:
    method __getitem__ (line 496) | def __getitem__(self, index):
    method __len__ (line 499) | def __len__(self):
  class GluonTSNumpyDataset (line 503) | class GluonTSNumpyDataset:
    method __init__ (line 514) | def __init__(
    method __iter__ (line 520) | def __iter__(self):
    method __len__ (line 525) | def __len__(self):
  class MaskInput (line 529) | class MaskInput(MapTransformation):
    method __init__ (line 531) | def __init__(
    method map_transform (line 548) | def map_transform(self, data: DataEntry, is_train: bool) -> DataEntry:
  class ConcatDataset (line 573) | class ConcatDataset:
    method __init__ (line 574) | def __init__(self, test_pairs, axis=-1) -> None:
    method _concat (line 578) | def _concat(self, test_pairs):
    method __iter__ (line 587) | def __iter__(self):