torchcast
torchcast

Covariance

class torchcast.covariance.Covariance(rank: int, method: str = 'log_cholesky', empty_idx: List[int] = (), predict_variance: Module | None = None, expected_kwargs: Sequence[str] | None = None, id: str | None = None, init_diag_multi: float = 0.1)

The Covariance can be used when you’d like more control over the covariance specification of a state- space model. For example, if you’re training on diverse time-serieses that vary in scale/behavior, you could use an torch.nn.Embedding to predict the variance of each series, with the group-ids as predictors:

kf = KalmanFilter(
    measures=measures,
    processes=processes,
    measure_covariance=Covariance.from_measures(
        measures,
        predict_variance=torch.nn.Sequential(
            torch.nn.Embedding(len(group_ids), len(measures), padding_idx=0),
            torch.nn.Softplus()
        ),
        expected_kwargs=['group_ids']
    ),
    process_covariance=Covariance.from_processes(
        processes,
        predict_variance=torch.nn.Sequential(
            torch.nn.Embedding(len(group_ids), Covariance.from_processes(processes).param_rank, padding_idx=0),
            torch.nn.Softplus()
        ),
        expected_kwargs=['group_ids']
    )
)
classmethod from_measures(measures: Sequence[str], predict_variance: bool | Module = None, **kwargs) → Covariance
Parameters:

  • measures – The measures used in your StateSpaceModel.

  • predict_variance – Will the variance be predicted upon calling forward()? This is implemented as a multiplier on the base variance given from the ‘method’. You can either pass True in which case it is expected you will pass multipliers as ‘measure_var_multi’ when forward() is called; or you can pass a torch.nn.Module that will predict the multipliers, in which case you’ll pass input(s) to this module at forward. Either way please note these should output strictly positive values with shape (num_groups, num_times, len(measures)).

  • kwargs – Other arguments passed to Covariance.__init__().

Returns:

A Covariance object that can be used in your StateSpaceModel.

classmethod from_processes(processes: Sequence[Process], cov_type: str = 'process', predict_variance: bool | Module = None, **kwargs) → Covariance
Parameters:

  • processes – The processes used in your StateSpaceModel.

  • cov_type – The type of covariance, either ‘process’ or ‘initial’ (default: ‘process’).

  • predict_variance – Will the variance be predicted upon calling forward()? This is implemented as a multiplier on the base variance given from the ‘method’. You can either pass True in which case it is expected you will pass multipliers as ‘process_var_multi’ when forward() is called; or you can pass a torch.nn.Module that will predict the multipliers, in which case you’ll pass input(s) to this module at forward. Either way please note these should output strictly positive values with shape (num_groups, num_times, self.param_rank).

  • kwargs – Other arguments passed to Covariance.__init__().

Returns:

A Covariance object that can be used in your StateSpaceModel.

Processes
Utils