nrdk.tss

Time series metric statistics.

The high level API is broken up into four steps:

1. index: Index evaluation files using a regex pattern.

   Info

   If you are only interested in a subset of evaluation traces, you can filter them at this stage.

2. experiments_from_index: Load data for each indexed result file, or a subset of experiments.

3. stats_from_experiments: Compute statistics for each experiment. See NDStats and effective_sample_size for more details about how and what statistics are computed.

4. dataframe_from_stats: Aggregate the statistics into a readable dataframe, ready to be plotted or exported.

Tip

We also provide dataframe_from_index, which combines the last three steps into a single function for convenience.

nrdk.tss.Control dataclass

Specifications for a control variable.

Warning

Each control variable should have a unique name.

Attributes:

Name	Type	Description
name	str	the name of this control variable.
baselines	Mapping[str, str]	a mapping between experiment names and their respective baselines.

Source code in src/nrdk/tss/control.py

@dataclass
class Control:
    r"""Specifications for a control variable.

    !!! warning

        Each control variable should have a unique name.

    Attributes:
        name: the name of this control variable.
        baselines: a mapping between experiment names and their respective
            baselines.
    """

    name: str
    baselines: Mapping[str, str]

    @classmethod
    def from_index_rule(
        cls, name: str, experiments: Iterable[str | None],
        rule: Callable[[str], str]
    ) -> "Control":
        """Create a control by applying a rule to each experiment name.

        !!! example

            Pairing each `<run>/<split>` experiment against the baseline run
            at the same split:
            ```python
            Control.from_index_rule(
                "split", index, lambda e: f"base/{e.rsplit('/', 1)[1]}")
            ```

        Args:
            name: the name of this control variable.
            experiments: experiment names to build the mapping over; an
                [`index`][^^^.api.] can be passed directly, and any `None`
                key is skipped.
            rule: given an experiment name, returns the name of the
                experiment which it should be compared against.

        Returns:
            A control which pairs each experiment with `rule(experiment)`.
        """
        return cls(name, {e: rule(e) for e in experiments if e is not None})

    def resolve(self, experiments: Sequence[str]) -> list[str]:
        """Experiments which this control applies to.

        Args:
            experiments: all experiment names which were loaded.

        Returns:
            The experiments which this control pairs against a baseline, in
                the order given.

        Raises:
            ValueError: if this control covers none of `experiments`, or
                names a baseline which is not among them.
        """
        covered = [k for k in experiments if k in self.baselines]
        if len(covered) == 0:
            raise ValueError(
                f"Control '{self.name}' does not cover any of the available "
                f"experiments: {list(experiments)}")

        missing = sorted(
            {self.baselines[k] for k in covered} - set(experiments))
        if len(missing) > 0:
            raise ValueError(
                f"Control '{self.name}' refers to baselines which are not "
                f"present in the index: {missing}.")

        return covered

from_index_rule classmethod

from_index_rule(
    name: str, experiments: Iterable[str | None], rule: Callable[[str], str]
) -> Control

Create a control by applying a rule to each experiment name.

Example

Pairing each <run>/<split> experiment against the baseline run at the same split:

Control.from_index_rule(
    "split", index, lambda e: f"base/{e.rsplit('/', 1)[1]}")

Parameters:

Name	Type	Description	Default
name	str	the name of this control variable.	required
experiments	Iterable[str | None]	experiment names to build the mapping over; an index can be passed directly, and any None key is skipped.	required
rule	Callable[[str], str]	given an experiment name, returns the name of the experiment which it should be compared against.	required

Returns:

Type	Description
Control	A control which pairs each experiment with rule(experiment).

Source code in src/nrdk/tss/control.py

@classmethod
def from_index_rule(
    cls, name: str, experiments: Iterable[str | None],
    rule: Callable[[str], str]
) -> "Control":
    """Create a control by applying a rule to each experiment name.

    !!! example

        Pairing each `<run>/<split>` experiment against the baseline run
        at the same split:
        ```python
        Control.from_index_rule(
            "split", index, lambda e: f"base/{e.rsplit('/', 1)[1]}")
        ```

    Args:
        name: the name of this control variable.
        experiments: experiment names to build the mapping over; an
            [`index`][^^^.api.] can be passed directly, and any `None`
            key is skipped.
        rule: given an experiment name, returns the name of the
            experiment which it should be compared against.

    Returns:
        A control which pairs each experiment with `rule(experiment)`.
    """
    return cls(name, {e: rule(e) for e in experiments if e is not None})

resolve

resolve(experiments: Sequence[str]) -> list[str]

Experiments which this control applies to.

Parameters:

Name	Type	Description	Default
experiments	Sequence[str]	all experiment names which were loaded.	required

Returns:

Type	Description
list[str]	The experiments which this control pairs against a baseline, in the order given.

Raises:

Type	Description
ValueError	if this control covers none of experiments, or names a baseline which is not among them.

Source code in src/nrdk/tss/control.py

def resolve(self, experiments: Sequence[str]) -> list[str]:
    """Experiments which this control applies to.

    Args:
        experiments: all experiment names which were loaded.

    Returns:
        The experiments which this control pairs against a baseline, in
            the order given.

    Raises:
        ValueError: if this control covers none of `experiments`, or
            names a baseline which is not among them.
    """
    covered = [k for k in experiments if k in self.baselines]
    if len(covered) == 0:
        raise ValueError(
            f"Control '{self.name}' does not cover any of the available "
            f"experiments: {list(experiments)}")

    missing = sorted(
        {self.baselines[k] for k in covered} - set(experiments))
    if len(missing) > 0:
        raise ValueError(
            f"Control '{self.name}' refers to baselines which are not "
            f"present in the index: {missing}.")

    return covered

nrdk.tss.dataframe_from_index

dataframe_from_index(
    index: dict[str | None, dict[str | None, str]],
    key: str,
    timestamps: str | None = None,
    experiments: Sequence[str | None] | None = None,
    cut: float | None = None,
    baseline: str | None = None,
    workers: int = -1,
    t_max: int | None = None,
    controls: Sequence[Control] = (),
) -> DataFrame

Load and calculate statistics from indexed experiment results.

See (1) dataframe_from_stats, (2) stats_from_experiments, and (3) experiments_from_index.

Controlling for additional variables

By default, every experiment is compared against a single global baseline. A Control adds a second comparison which holds one of those axes fixed, pairing each experiment against the baseline sharing its value on that axis.

Each control adds a rel_{name}/*, pct_{name}/*, and p0.05_{name} column group, computed independently of the others, which report the effect of every axis except the one it holds fixed. Experiments which a control does not cover are NaN in its columns, and every baseline it names must itself be loaded.

Parameters:

Name	Type	Description	Default
index	dict[str | None, dict[str | None, str]]	2-level dictionary with experiment names, sequence/trace names, and paths to the result files; see index.	required
key	str	name of the metric to load from the result files.	required
timestamps	str | None	name of the timestamps to load from the result files.	None
experiments	Sequence[str | None] | None	list of experiment names to load from the index; loads all experiments if not specified.	None
cut	float | None	cut each time series when there is a gap in the timestamps larger than this value if provided; see cut_trace.	None
baseline	str | None	baseline experiment for relative statistics.	None
workers	int	number of worker threads to use when loading. If <0, load all in parallel; if =0, load all in the main thread.	-1
t_max	int | None	maximum time delay to consider when computing effective sample size; if None, do not use any additional constraints.	None
controls	Sequence[Control]	additional control variables to compute paired statistics for; see Control. Requires a baseline.	()

Returns:

Type	Description
DataFrame	Dataframe with statistics for each experiment.

Source code in src/nrdk/tss/api.py

def dataframe_from_index(
    index: dict[str | None, dict[str | None, str]],
    key: str, timestamps: str | None = None,
    experiments: Sequence[str | None] | None = None,
    cut: float | None = None, baseline: str | None = None, workers: int = -1,
    t_max: int | None = None, controls: Sequence[Control] = ()
) -> pd.DataFrame:
    """Load and calculate statistics from indexed experiment results.

    See (1) [`dataframe_from_stats`][^.], (2) [`stats_from_experiments`][^.],
    and (3) [`experiments_from_index`][^.].

    !!! info "Controlling for additional variables"

        By default, every experiment is compared against a single global
        `baseline`. A [`Control`][^.] adds a second comparison which
        holds one of those axes fixed, pairing each experiment against the
        baseline sharing its value on that axis.

        Each control adds a `rel_{name}/*`, `pct_{name}/*`, and
        `p0.05_{name}` column group, computed independently of the others,
        which report the effect of every axis *except* the one it holds
        fixed. Experiments which a control does not cover are `NaN` in its
        columns, and every baseline it names must itself be loaded.

    Args:
        index: 2-level dictionary with experiment names, sequence/trace names,
            and paths to the result files; see [`index`][^^.api.].
        key: name of the metric to load from the result files.
        timestamps: name of the timestamps to load from the result files.
        experiments: list of experiment names to load from the index; loads all
            experiments if not specified.
        cut: cut each time series when there is a gap in the timestamps larger
            than this value if provided; see [`cut_trace`][^^.utils.].
        baseline: baseline experiment for relative statistics.
        workers: number of worker threads to use when loading. If `<0`, load
            all in parallel; if `=0`, load all in the main thread.
        t_max: maximum time delay to consider when computing effective sample
            size; if `None`, do not use any additional constraints.
        controls: additional control variables to compute paired statistics
            for; see [`Control`][^.]. Requires a `baseline`.

    Returns:
        Dataframe with statistics for each experiment.
    """
    if len(controls) > 0 and baseline is None:
        raise ValueError(
            "Provided `controls`, but no global `baseline`; a baseline is "
            "required to compute the `pct_{name}/*` columns.")

    duplicates = sorted({
        k for k, v in Counter(c.name for c in controls).items() if v > 1})
    if len(duplicates) > 0:
        raise ValueError(
            f"Controls must have unique names, since each control adds its "
            f"own column group; got duplicates: {duplicates}")

    y, t, _ = experiments_from_index(
        index, key, timestamps=timestamps, experiments=experiments,
        cut=cut, workers=workers)
    names, stats_abs, stats_rel = stats_from_experiments(
        y, t, baseline=baseline, workers=workers, t_max=t_max)
    df = dataframe_from_stats(names, stats_abs, stats_rel, baseline=baseline)

    for control in controls:
        df = _append_control(
            df, control, y, t, names, baseline,  # type: ignore
            workers=workers, t_max=t_max)

    return df

nrdk.tss.dataframe_from_stats

dataframe_from_stats(
    names: list[str],
    abs: NDStats,
    rel: NDStats | None = None,
    baseline: str | None = None,
) -> DataFrame

Create a dataframe from (possibly un-aggregated) experiment statistics.

Returns a dataframe where each row is a different experiment.

• abs/(mean|std|stderr|n|ess): absolute statistics for the provided metric for each experiment.
• rel/(mean|std|stderr|n|ess): relative statistics for the provided metric for each experiment, relative to the baseline. If no baseline is provided, these columns are not included.
• pct/(mean|stderr): percent difference and standard error relative to the baseline, computed as 100 * <rel/mean>/<abs/mean> and 100 * <rel/stderr>/<abs/mean>, where <abs/mean> is the baseline's absolute mean.
• p0.05: whether the difference from the baseline is significant at the 5% level (two-sided), Bonferroni-corrected by the number of experiments compared against the baseline; pd.NA where no comparison was made.

Parameters:

Name	Type	Description	Default
names	list[str]	names of the experiments corresponding to the leading axis in the input statistics.	required
abs	NDStats	absolute statistics for the provided metric for each experiment.	required
rel	NDStats | None	optional relative statistics.	None
baseline	str | None	name of the experiment used as the baseline.	None

Returns:

Type	Description
DataFrame	Dataframe with statistics for each experiment.

Source code in src/nrdk/tss/api.py

def dataframe_from_stats(
    names: list[str], abs: NDStats, rel: NDStats | None = None,
    baseline: str | None = None
) -> pd.DataFrame:
    """Create a dataframe from (possibly un-aggregated) experiment statistics.

    Returns a dataframe where each row is a different experiment.

    - `abs/(mean|std|stderr|n|ess)`: absolute statistics for the
        provided metric for each experiment.
    - `rel/(mean|std|stderr|n|ess)`: relative statistics for the
        provided metric for each experiment, relative to the `baseline`. If no
        `baseline` is provided, these columns are not included.
    - `pct/(mean|stderr)`: percent difference and standard error relative to
        the `baseline`, computed as `100 * <rel/mean>/<abs/mean>` and
        `100 * <rel/stderr>/<abs/mean>`, where `<abs/mean>` is the
        *baseline's* absolute mean.
    - `p0.05`: whether the difference from the `baseline` is significant at
        the 5% level (two-sided), Bonferroni-corrected by the number of
        experiments compared against the baseline; `pd.NA` where no
        comparison was made.

    Args:
        names: names of the experiments corresponding to the leading axis in
            the input statistics.
        abs: absolute statistics for the provided metric for each experiment.
        rel: optional relative statistics.
        baseline: name of the experiment used as the baseline.

    Returns:
        Dataframe with statistics for each experiment.
    """
    df = abs.reshape(
        len(names), -1).sum(axis=-1).as_df(names, prefix="abs/")

    if rel is not None and baseline is None:
        raise ValueError(
            "Provided relative statistics `rel`, but the `baseline` used is "
            "not specified.")

    if rel is not None:
        df_rel = rel.reshape(
            len(names), -1).sum(axis=-1).as_df(names, prefix="rel/")
        df = df.merge(df_rel, on='name')
        _baseline = float(df.at[baseline, 'abs/mean'])  # type: ignore
        df['pct/mean'] = df['rel/mean'] / _baseline * 100
        df['pct/stderr'] = df['rel/stderr'] / _baseline * 100

        # Two-sided, Bonferroni-corrected by the number of experiments
        # compared against the baseline
        z = norm.ppf(1 - 0.05 / 2 / max(len(names) - 1, 1))
        df['p0.05'] = (
            (df['rel/mean'].abs() / df['rel/stderr']) > z
        ).where(df['rel/mean'].notna()).astype("boolean")

    return df

nrdk.tss.experiments_from_index

experiments_from_index(
    index: dict[str | None, dict[str | None, str]],
    key: str,
    timestamps: str | None = None,
    experiments: Sequence[str | None] | str | None = None,
    cut: float | None = None,
    workers: int = -1,
) -> tuple[
    Mapping[str, NestedValues[Num[ndarray, _N]]],
    Mapping[str, NestedValues[Float64[ndarray, _N]]] | None,
    list[str | None],
]

Load experiment results from indexed result files.

Each results file is expected to be a .npz file containing metric and metadata arrays; the keys for these arrays should be specified by key and timestamps, respectively.

• These arrays should all have the same leading axis length.
• The metric array should have only a single axis.

Warning

Only sequences which are present in all experiments will be loaded. Check the returned common list to make sure it matchse what you expect!

Tip

A timestamps key can optionally be provided.

• If not provided, the metrics are assumed to be at identical timestamps.
• If multiple timestamps are present, the last one is used.

Parameters:

Name	Type	Description	Default
index	dict[str | None, dict[str | None, str]]	2-level dictionary with experiment names, sequence/trace names, and paths to the result files; see index.	required
key	str	name of the metric to load from the result files.	required
timestamps	str | None	name of the timestamps to load from the result files.	None
experiments	Sequence[str | None] | str | None	list of experiment names to load from the index (or a regex filter); loads all experiments if not specified.	None
cut	float | None	cut each time series when there is a gap in the timestamps larger than this value if provided; see cut_trace.	None
workers	int	number of worker threads to use when loading. If <0, load all in parallel; if =0, load all in the main thread.	-1

Returns:

Type	Description
Mapping[str, NestedValues[Num[ndarray, _N]]]	A dictionary of metric values (as a list of metric values by sequence).
Mapping[str, NestedValues[Float64[ndarray, _N]]] | None	A dictionary of timestamps (or None if not specified).
list[str | None]	A list of the common sequence/trace names which correspond to the loaded metrics.

Source code in src/nrdk/tss/api.py

def experiments_from_index(
    index: dict[str | None, dict[str | None, str]],
    key: str, timestamps: str | None = None,
    experiments: Sequence[str | None] | str | None = None,
    cut: float | None = None, workers: int = -1
) -> tuple[
    Mapping[str, NestedValues[Num[np.ndarray, "_N"]]],
    Mapping[str, NestedValues[Float64[np.ndarray, "_N"]]] | None,
    list[str | None]
]:
    """Load experiment results from indexed result files.

    Each results file is expected to be a `.npz` file containing metric and
    metadata arrays; the keys for these arrays should be specified by `key` and
    `timestamps`, respectively.

    - These arrays should all have the same leading axis length.
    - The metric array should have only a single axis.

    !!! warning

        Only sequences which are present in all experiments will be loaded.
        Check the returned `common` list to make sure it matchse what you
        expect!

    !!! tip

        A `timestamps` key can optionally be provided.

        - If not provided, the metrics are assumed to be at identical
            timestamps.
        - If multiple timestamps are present, the last one is used.

    Args:
        index: 2-level dictionary with experiment names, sequence/trace names,
            and paths to the result files; see [`index`][^.].
        key: name of the metric to load from the result files.
        timestamps: name of the timestamps to load from the result files.
        experiments: list of experiment names to load from the index (or a
            regex filter); loads all experiments if not specified.
        cut: cut each time series when there is a gap in the timestamps larger
            than this value if provided; see [`cut_trace`][^^.utils.].
        workers: number of worker threads to use when loading. If `<0`, load
            all in parallel; if `=0`, load all in the main thread.

    Returns:
        A dictionary of metric values (as a list of metric values by sequence).
        A dictionary of timestamps (or `None` if not specified).
        A list of the common sequence/trace names which correspond to the
            loaded metrics.
    """
    if len(index) == 0:
        raise ValueError("Could not fetch experiments: the index is empty.")

    if experiments is None:
        experiments = list(index.keys())
    elif isinstance(experiments, str):
        re_filter = re.compile(experiments)
        experiments = [x for x in index.keys() if re_filter.match(str(x))]
        if len(experiments) == 0:
            raise ValueError(
                f"No experiments found matching the filter: {experiments}")

    common = list(set.intersection(
        *[set(index[k].keys()) for k in experiments]))
    if workers < 0:
        workers = len(common) * len(experiments)

    def _load(path: str):
        data = np.load(path)

        if key not in data:
            raise KeyError(
                f"Key '{key}' not found in file: {path}. Available keys: "
                f"{list(data.keys())}")

        if timestamps is not None:
            t = data[timestamps]
            t = t.reshape(t.shape[0], -1)[:, -1]
            if cut is not None:
                ytyt = cut_trace(t, (data[key], t), gap=cut)
                return list(zip(*ytyt))
            else:
                return [data[key]], [t]
        else:
            return [data[key]]

    iterload = [(x, s) for x in experiments for s in common]
    if workers == 0:
        loaded = [_load(index[x][s]) for x, s in iterload]
    else:
        with pool.ThreadPool(workers) as p:
            loaded = list(p.map(_load, [index[x][s] for x, s in iterload]))

    if timestamps is not None:
        yy, tt = {}, {}
        for (x, s), (y, t) in zip(iterload, loaded):
            yy.setdefault(x, []).extend(y)
            tt.setdefault(x, []).extend(t)
        return yy, tt, common
    else:
        yy = {}
        for (x, s), y in zip(iterload, loaded):
            yy.setdefault(x, []).extend(y)
        return yy, None, common

nrdk.tss.index

index(
    path: str, pattern: str | Pattern, follow_symlinks: bool = False
) -> dict[str | None, dict[str | None, str]]

Recursively find all evaluations matching the given pattern.

The pattern can have two groups: experiment, and trace, which respectively indicate the name of the experiment and trace. If either group is omitted, it is set as None.

Example

# match `<experiment>/eval/<trace>.npz`
index(path, r'^(?P<experiment>(.*))/eval/(?P<trace>(.*))\.npz$')
# match `<experiment>.npz`
index(path, r'^(?P<experiment>(.*))\.npz$')

Parameters:

Name	Type	Description	Default
path	str	directory to start searching from.	required
pattern	str | Pattern	regex pattern to match the evaluation directories.	required
follow_symlinks	bool	whether to follow symbolic links.	False

Returns:

Type	Description
dict[str | None, dict[str | None, str]]	A two-level dictionary, where the first level keys are the experiment names, the second level keys are the trace names, and the values are paths to the matching files.

Source code in src/nrdk/tss/api.py

def index(
    path: str, pattern: str | re.Pattern, follow_symlinks: bool = False
) -> dict[str | None, dict[str | None, str]]:
    r"""Recursively find all evaluations matching the given pattern.

    The pattern can have two groups: `experiment`, and `trace`, which
    respectively indicate the name of the experiment and trace. If either group
    is omitted, it is set as `None`.

    !!! example

        ```python
        # match `<experiment>/eval/<trace>.npz`
        index(path, r'^(?P<experiment>(.*))/eval/(?P<trace>(.*))\.npz$')
        # match `<experiment>.npz`
        index(path, r'^(?P<experiment>(.*))\.npz$')
        ```

    Args:
        path: directory to start searching from.
        pattern: regex pattern to match the evaluation directories.
        follow_symlinks: whether to follow symbolic links.

    Returns:
        A two-level dictionary, where the first level keys are the experiment
            names, the second level keys are the trace names, and the values
            are paths to the matching files.
    """
    if isinstance(pattern, str):
        pattern = re.compile(pattern)
    manifest = {}

    def _find(path, base):
        matches = pattern.match(os.path.relpath(path, base))
        if matches is not None:
            groups = matches.groupdict()
            manifest.setdefault(
                groups.get('experiment', None), {}
            )[groups.get('trace', None)] = path
        elif os.path.isdir(path):
            if follow_symlinks or not os.path.islink(path):
                for p in os.listdir(path):
                    _find(os.path.join(path, p), base)

    _find(path, path)
    return manifest

nrdk.tss.stats_from_experiments

stats_from_experiments(
    y: Mapping[str, NestedValues[Num[ndarray, _N]]],
    t: Mapping[str, NestedValues[Float64[ndarray, _N]]] | None = None,
    baseline: str | None = None,
    workers: int = -1,
    t_max: int | None = None,
) -> tuple[list[str], NDStats, NDStats | None]

Calculate statistics from experiment results.

Parameters:

Name	Type	Description	Default
y	Mapping[str, NestedValues[Num[ndarray, _N]]]	mapping of experiment names and metric values.	required
t	Mapping[str, NestedValues[Float64[ndarray, _N]]] | None	mapping of experiment names and timestamps. If not provided, the metrics are assumed to be at identical timestamps.	None
baseline	str | None	baseline experiment for relative statistics.	None
workers	int	number of worker threads to use for computation.	-1
t_max	int | None	maximum time delay to consider when computing effective sample size; if None, do not use any additional constraints.	None

Returns:

Type	Description
list[str]	Names of each experiment corresponding to leading axis in the output statistics.
NDStats	Absolute statistics for the provided metric.
NDStats | None	Relative statistics (difference relative to the specified baseline), if provided.

Source code in src/nrdk/tss/api.py

def stats_from_experiments(
    y: Mapping[str, NestedValues[Num[np.ndarray, "_N"]]],
    t: Mapping[str, NestedValues[Float64[np.ndarray, "_N"]]] | None = None,
    baseline: str | None = None, workers: int = -1, t_max: int | None = None
) -> tuple[list[str], NDStats, NDStats | None]:
    """Calculate statistics from experiment results.

    Args:
        y: mapping of experiment names and metric values.
        t: mapping of experiment names and timestamps. If not provided, the
            metrics are assumed to be at identical timestamps.
        baseline: baseline experiment for relative statistics.
        workers: number of worker threads to use for computation.
        t_max: maximum time delay to consider when computing effective sample
            size; if `None`, do not use any additional constraints.

    Returns:
        Names of each experiment corresponding to leading axis in the output
            statistics.
        Absolute statistics for the provided metric.
        Relative statistics (difference relative to the specified baseline), if
            provided.
    """
    n_sorted = sorted(y.keys())
    if t is not None and set(n_sorted) != set(t.keys()):
        raise ValueError(
            f"Keys of `y` and `t` must match if `t` is provided: "
            f"y:{list(y.keys())}, t:{list(t.keys())}")

    y_sorted = [y[k] for k in n_sorted]
    stats_abs = NDStats.from_values(y_sorted, workers=workers, t_max=t_max)
    if baseline is not None:
        stats_rel = _relative_stats(
            y, t, n_sorted, {k: baseline for k in n_sorted},
            workers=workers, t_max=t_max)
    else:
        stats_rel = None
    return n_sorted, stats_abs, stats_rel