spd_learn.modules.TimeFrequencyGraph#

class spd_learn.modules.TimeFrequencyGraph(coordinates: Tensor, *, time_radius: float, frequency_radius: float, components: Tensor | None = None, bandwidth: float = 1.0, neighborhood: Literal['box', 'forward'] = 'box')[source]#

Bases: Module

Construct a time-frequency graph with training-fitted edge weights.

See [Ju and Guan, 2024].

Parameters:

  • coordinates (torch.Tensor) – Node coordinates of shape (nodes, 2): time in seconds and frequency in Hz.

  • time_radius (float) – Temporal neighborhood radius in seconds.

  • frequency_radius (float) – Frequency neighborhood radius in Hz.

  • components (torch.Tensor or None) – Optional component IDs of shape (nodes,); edges stay within a component.

  • bandwidth (float, default=1.0) – Positive divisor of squared AIRM distances in Gaussian edge weights.

  • neighborhood ({'box', 'forward'}, default='box') – Symmetric midpoint box or the released forward-grid neighborhood.

Notes

fit estimates arithmetic training lattices and Gaussian AIRM edge weights. Forward returns the stored row-stochastic propagation without refitting. Coordinates must match covariance node order.

coordinates: Tensor#
fit(input: Tensor) → TimeFrequencyGraph[source]#

Fit edge weights from training covariance nodes.

Parameters:

input (torch.Tensor) – Finite SPD matrices of shape (training_trials, nodes, channels, channels).

Returns:

self – The graph with fitted lattice and propagation buffers.

Return type:

TimeFrequencyGraph

fitted: Tensor#
forward() → Tensor[source]#

Return the stored graph propagation.

Returns:

Row-stochastic propagation of shape (nodes, nodes).

Return type:

torch.Tensor

lattice: Tensor#
load_state_dict(state_dict, strict=True, assign=False)[source]#

Load a checkpoint while preserving state after a failed load.

Parameters:

  • state_dict (dict) – Parameters and buffers returned by state_dict.

  • strict (bool, default=True) – Require the checkpoint keys to match the module state.

  • assign (bool, default=False) – Use checkpoint tensor properties when assigning state.

Returns:

Missing and unexpected checkpoint keys.

Return type:

NamedTuple

Notes

Reconstruct the same graph and segmentation configuration before loading. Failed loads restore tensor values and ownership, including assign=True.

propagation: Tensor#
topology: Tensor#