Data Dataset Weights

dataset_weights

FUNCTION	DESCRIPTION
stabilize_weights	Stabilize highly-variable sample weights using the specified method.
weights_log_space_clip	Clip weights in log-space.
weights_quantile_clip	Hard-clip raw weights at a given quantile.
weights_softmax_normalize	Apply softmax in log-weight space and rescale.
weights_ess_filter	Zero-out weights that fall below an effective sample size threshold.

Functions

stabilize_weights(weights: torch.Tensor, method: Literal['log_space_clip', 'quantile_clip', 'softmax_normalize', 'ess_filter'] = 'log_space_clip', normalize: bool = True, **kwargs) -> torch.Tensor

Stabilize highly-variable sample weights using the specified method.

PARAMETER	DESCRIPTION
weights	Sample weights (1D tensor). TYPE: Tensor
method	Stabilization algorithm. Can be: "log_space_clip": clip in log-space with a median + n_sigma * std threshold. Suitable for exponential weights; see weights_log_space_clip. "quantile_clip": hard-clip raw weights at a given quantile; see weights_quantile_clip. "softmax_normalize": apply softmax in log-weight space and rescale; see weights_softmax_normalize. "ess_filter": set weights of samples below an effective sample size threshold to zero; see weights_ess_filter. TYPE: str, optional, by default "log_space_clip" DEFAULT: 'log_space_clip'
normalize	If True, divide the result by its mean so the output mean equals one. Has no effect for "softmax_normalize" (already normalized) or "ess_filter" (binary mask). TYPE: bool, optional, by default True DEFAULT: True
**kwargs	Additional keyword arguments forwarded to the chosen method function. See the individual function docstrings for available parameters. DEFAULT: {}

RETURNS	DESCRIPTION
Tensor	1D tensor of stabilized weights.

Examples:

Default log-space clipping:

>>> import torch
>>> from spectre.data import stabilize_weights
>>> weights = torch.exp(torch.randn(500))
>>> stable = stabilize_weights(weights)
>>> stable.mean()
tensor(1.)

Quantile clipping at the 95^th percentile:

>>> stable = stabilize_weights(
...     weights,
...     method="quantile_clip",
...     quantile=0.95,
... )

Notes

• All methods retain every sample in the dataset; weights are modified in place of discarding samples.
• "log_space_clip" is recommended for exponential weights from metadynamics because clipping the log-weight is equivalent to capping the bias magnitude.
• "normalize" rescales to mean equal to 1 after stabilization to prevent scale drift.

Source code in spectre/data/dataset_weights.py

def stabilize_weights(
    weights: torch.Tensor,
    method: Literal[
        "log_space_clip",
        "quantile_clip",
        "softmax_normalize",
        "ess_filter",
    ] = "log_space_clip",
    normalize: bool = True,
    **kwargs,
) -> torch.Tensor:
    """
    Stabilize highly-variable sample weights using the specified method.

    Parameters
    ----------
    weights : torch.Tensor
        Sample weights (1D tensor).

    method : str, optional, by default "log_space_clip"
        Stabilization algorithm.

        Can be:

        - `"log_space_clip"`: clip in log-space with a `median + n_sigma * std`
          threshold. Suitable for exponential weights; see
          `weights_log_space_clip`.
        - `"quantile_clip"`: hard-clip raw weights at a given quantile; see
          `weights_quantile_clip`.
        - `"softmax_normalize"`: apply softmax in log-weight space and rescale;
          see `weights_softmax_normalize`.
        - `"ess_filter"`: set weights of samples below an effective sample size
          threshold to zero; see `weights_ess_filter`.

    normalize : bool, optional, by default True
        If True, divide the result by its mean so the output mean equals one.
        Has no effect for `"softmax_normalize"` (already normalized) or
        `"ess_filter"` (binary mask).

    **kwargs
        Additional keyword arguments forwarded to the chosen method function.
        See the individual function docstrings for available parameters.

    Returns
    -------
    torch.Tensor
        1D tensor of stabilized weights.

    Examples
    --------
    Default log-space clipping:

    >>> import torch
    >>> from spectre.data import stabilize_weights
    >>> weights = torch.exp(torch.randn(500))
    >>> stable = stabilize_weights(weights)
    >>> stable.mean()
    tensor(1.)

    Quantile clipping at the 95th percentile:

    >>> stable = stabilize_weights(
    ...     weights,
    ...     method="quantile_clip",
    ...     quantile=0.95,
    ... )

    Notes
    -----
    - All methods retain every sample in the dataset; weights are modified in
      place of discarding samples.
    - `"log_space_clip"` is recommended for exponential weights from
      metadynamics because clipping the log-weight is equivalent to capping the
      bias magnitude.
    - `"normalize"` rescales to mean equal to 1 after stabilization to prevent
      scale drift.
    """
    avail_methods: dict[str, Callable] = {
        "log_space_clip": weights_log_space_clip,
        "quantile_clip": weights_quantile_clip,
        "softmax_normalize": weights_softmax_normalize,
        "ess_filter": weights_ess_filter,
    }

    if method not in avail_methods:
        raise ValueError(
            f"Unknown method '{method}'. "
            f"Choose from: {', '.join(repr(m) for m in avail_methods)}."
        )

    stabilized = avail_methods[method](weights, **kwargs)

    if normalize and method not in ("softmax_normalize", "ess_filter"):
        mean = stabilized.mean()
        if mean > 0:
            stabilized = stabilized / mean

    return stabilized

weights_log_space_clip(weights: torch.Tensor, n_sigma: float = 3.0) -> torch.Tensor

Clip weights in log-space.

Computes the logarithm of weights log(w) and clips any value above median(log_w) + n_sigma * std(log_w) before exponentiating back. The median is used instead of the mean as the log-space center so that the threshold itself is not skewed by the extreme values being suppressed.

PARAMETER	DESCRIPTION
weights	Sample weights (1D tensor). TYPE: Tensor
n_sigma	Number of standard deviations above the log-space median used as the clipping threshold. Larger values allow more extreme weights through; smaller values enforce stronger regularization. Typical range: 2–5. TYPE: float, optional, by default 3.0 DEFAULT: 3.0

RETURNS	DESCRIPTION
Tensor	1D tensor of clipped weights (not normalized).

Examples:

>>> import torch
>>> from spectre.data import weights_log_space_clip
>>> weights = torch.exp(torch.randn(500))
>>> clipped = weights_log_space_clip(weights, n_sigma=3.0)
>>> clipped.max() <= weights.max()
tensor(True)

Source code in spectre/data/dataset_weights.py

def weights_log_space_clip(
    weights: torch.Tensor,
    n_sigma: float = 3.0,
) -> torch.Tensor:
    """
    Clip weights in log-space.

    Computes the logarithm of weights `log(w)` and clips any value above
    `median(log_w) + n_sigma * std(log_w)` before exponentiating back. The
    median is used instead of the mean as the log-space center so that the
    threshold itself is not skewed by the extreme values being suppressed.

    Parameters
    ----------
    weights : torch.Tensor
        Sample weights (1D tensor).

    n_sigma : float, optional, by default 3.0
        Number of standard deviations above the log-space median used as the
        clipping threshold. Larger values allow more extreme weights through;
        smaller values enforce stronger regularization. Typical range: 2–5.

    Returns
    -------
    torch.Tensor
        1D tensor of clipped weights (not normalized).

    Examples
    --------
    >>> import torch
    >>> from spectre.data import weights_log_space_clip
    >>> weights = torch.exp(torch.randn(500))
    >>> clipped = weights_log_space_clip(weights, n_sigma=3.0)
    >>> clipped.max() <= weights.max()
    tensor(True)
    """
    check_1d(weights)
    check_sample_weights(weights)

    check_in_interval(n_sigma, "(0,inf)")

    log_w = torch.log(weights)
    threshold = torch.median(log_w) + n_sigma * log_w.std()

    return torch.exp(torch.clamp(log_w, max=threshold))

weights_quantile_clip(weights: torch.Tensor, quantile: float = 0.99) -> torch.Tensor

Hard-clip raw weights at a given quantile.

All weights above the quantile-th percentile are set to that percentile value (Winsorization).

PARAMETER	DESCRIPTION
weights	1D tensor of sample weights. TYPE: Tensor
quantile	Fraction of samples to keep unclipped. Must be in (0, 1). For example, 0.99 clips the top 1 % of weights. TYPE: float, optional, by default 0.99 DEFAULT: 0.99

RETURNS	DESCRIPTION
Tensor	1D tensor of clipped weights (not normalized).

Examples:

>>> import torch
>>> from spectre.data import weights_quantile_clip
>>> weights = torch.exp(torch.randn(500))
>>> clipped = weights_quantile_clip(weights, quantile=0.99)
>>> clipped.max() <= weights.max()
tensor(True)

Source code in spectre/data/dataset_weights.py

def weights_quantile_clip(
    weights: torch.Tensor,
    quantile: float = 0.99,
) -> torch.Tensor:
    """
    Hard-clip raw weights at a given quantile.

    All weights above the `quantile`-th percentile are set to that percentile
    value (Winsorization).

    Parameters
    ----------
    weights : torch.Tensor
        1D tensor of sample weights.

    quantile : float, optional, by default 0.99
        Fraction of samples to keep unclipped. Must be in (0, 1).
        For example, 0.99 clips the top 1 % of weights.

    Returns
    -------
    torch.Tensor
        1D tensor of clipped weights (not normalized).

    Examples
    --------
    >>> import torch
    >>> from spectre.data import weights_quantile_clip
    >>> weights = torch.exp(torch.randn(500))
    >>> clipped = weights_quantile_clip(weights, quantile=0.99)
    >>> clipped.max() <= weights.max()
    tensor(True)
    """
    check_1d(weights)
    check_sample_weights(weights)

    check_in_interval(quantile, "(0,1)")

    cap = torch.quantile(weights, quantile)

    return torch.clamp(weights, max=cap)

weights_softmax_normalize(weights: torch.Tensor) -> torch.Tensor

Apply softmax in log-weight space and rescale.

Computes softmax(log(weights)) then multiplies by the number of samples (equivalent to a mean of 1). When the range of log-weights is very large the distribution is effectively flattened; choose weights_log_space_clip if preserving the relative ranking of high-weight samples is important.

PARAMETER	DESCRIPTION
weights	1D tensor of sample weights. TYPE: Tensor

RETURNS	DESCRIPTION
Tensor	1D tensor of normalized weights that sum to len(weights).

Examples:

>>> import torch
>>> from spectre.data import weights_softmax_normalize
>>> weights = torch.exp(torch.randn(500))
>>> normed = weights_softmax_normalize(weights)
>>> normed.sum().round()
tensor(500.)

Source code in spectre/data/dataset_weights.py

def weights_softmax_normalize(
    weights: torch.Tensor,
) -> torch.Tensor:
    """
    Apply softmax in log-weight space and rescale.

    Computes `softmax(log(weights))` then multiplies by the number of samples
    (equivalent to a mean of 1). When the range of log-weights is very large
    the distribution is effectively flattened; choose `weights_log_space_clip`
    if preserving the relative ranking of high-weight samples is important.

    Parameters
    ----------
    weights : torch.Tensor
        1D tensor of sample weights.

    Returns
    -------
    torch.Tensor
        1D tensor of normalized weights that sum to `len(weights)`.

    Examples
    --------
    >>> import torch
    >>> from spectre.data import weights_softmax_normalize
    >>> weights = torch.exp(torch.randn(500))
    >>> normed = weights_softmax_normalize(weights)
    >>> normed.sum().round()
    tensor(500.)
    """
    check_1d(weights)
    check_sample_weights(weights)

    log_w = torch.log(weights)

    return torch.softmax(log_w, dim=0) * len(weights)

weights_ess_filter(weights: torch.Tensor, threshold: float = 0.3) -> torch.Tensor

Zero-out weights that fall below an effective sample size threshold.

Computes the effective sample size (ESS) after cumulatively including each sample in time order, and sets the weight of all samples to zero until the cumulative ESS fraction ESS / n_included exceeds threshold. This identifies the earliest point at which the weight distribution is sufficiently converged and masks out everything before it.

ESS is defined as (sum(w))^2 / sum(w^2), which ranges from 1 (all weight on one sample) to N (uniform weights).

PARAMETER	DESCRIPTION
weights	1D tensor of sample weights ordered by time (earliest first). TYPE: Tensor
threshold	Minimum ESS fraction ESS / n_included required to include a sample. Must be in (0, 1). Lower values are more permissive (include more early samples); higher values are stricter. If the threshold is never reached, all weights are returned as-is with a warning. TYPE: float, optional, by default 0.3 DEFAULT: 0.3

RETURNS	DESCRIPTION
Tensor	1D tensor with the same values as weights for samples that passed the ESS filter, and zero for samples that did not.

Examples:

>>> import torch
>>> from spectre.data import weights_ess_filter
>>> weights = torch.exp(torch.randn(500))
>>> filtered = weights_ess_filter(weights, threshold=0.3)
>>> (filtered == 0).sum() >= 0  # some early samples may be zeroed
tensor(True)

Source code in spectre/data/dataset_weights.py

def weights_ess_filter(
    weights: torch.Tensor,
    threshold: float = 0.3,
) -> torch.Tensor:
    """
    Zero-out weights that fall below an effective sample size threshold.

    Computes the effective sample size (ESS) after cumulatively including each
    sample in time order, and sets the weight of all samples to zero until the
    cumulative ESS fraction `ESS / n_included` exceeds `threshold`. This
    identifies the earliest point at which the weight distribution is
    sufficiently converged and masks out everything before it.

    ESS is defined as `(sum(w))^2 / sum(w^2)`, which ranges from 1 (all weight
    on one sample) to N (uniform weights).

    Parameters
    ----------
    weights : torch.Tensor
        1D tensor of sample weights ordered by time (earliest first).

    threshold : float, optional, by default 0.3
        Minimum ESS fraction `ESS / n_included` required to include a sample.
        Must be in (0, 1).

        Lower values are more permissive (include more early samples); higher
        values are stricter. If the threshold is never reached, all weights are
        returned as-is with a warning.

    Returns
    -------
    torch.Tensor
        1D tensor with the same values as `weights` for samples that passed the
        ESS filter, and zero for samples that did not.

    Examples
    --------
    >>> import torch
    >>> from spectre.data import weights_ess_filter
    >>> weights = torch.exp(torch.randn(500))
    >>> filtered = weights_ess_filter(weights, threshold=0.3)
    >>> (filtered == 0).sum() >= 0  # some early samples may be zeroed
    tensor(True)
    """
    check_1d(weights)
    check_sample_weights(weights)

    check_in_interval(threshold, "(0,1)")

    n = len(weights)
    mask = torch.zeros(n, dtype=weights.dtype, device=weights.device)

    # Scan forward and find the first index i (1-based) at which the cumulative
    # ESS fraction recovers to >= threshold after having dropped below it.
    # ESS/i begins at 1.0 for a single sample, dips while extreme early weights
    # dominate, and recovers once enough equilibrated samples are included. We
    # want to zero out all samples before that recovery point.
    below_threshold_seen = False
    cutoff = 0  # default: keep all samples

    for i in range(1, n + 1):
        w_i = weights[:i]
        ess = w_i.sum().pow(2) / w_i.pow(2).sum()
        ess_fraction = (ess / i).item()

        if ess_fraction < threshold:
            below_threshold_seen = True

        if below_threshold_seen and ess_fraction >= threshold:
            cutoff = i - 1
            break
    else:
        if below_threshold_seen:
            # ESS dropped but never recovered — zero out everything
            cutoff = n
        # If ESS never dropped below threshold the weights are already good;
        # cutoff stays 0 and all samples are kept.

    mask[cutoff:] = 1.0

    return weights * mask