Useful functions for converting between different types (dicts, lists, tuples, etc.)

"""Useful functions for converting between different types (dicts, lists, tuples, etc.)
"""

import operator
from copy import deepcopy
from typing import Union
from uuid import uuid4

import numpy as np
import pandas as pd
import ray
from ray.remote_function import RemoteFunction as RayRemoteFun

from vflow.subkey import Subkey
from vflow.vfunc import VfuncPromise

PREV_KEY = "__prev__"


def init_step(idx, cols):
    """Name for the 'init' column at `idx`: "init-" plus the first non-init
    column after it (the Vset the data went into), or None if there is none.

    Parameters
    ----------
    idx: int
        Index of 'init' column in cols.
    cols: list[str]
        List of column names.
    """
    for i in range(idx, len(cols)):
        if cols[i] != "init":
            return "init-" + cols[i]
    return None


def base_dict(d: dict):
    """Remove PREV_KEY from dict d if present"""
    return {k: v for k, v in d.items() if k != PREV_KEY}


def dict_data(d: dict):
    """Returns a list containing all data in dict d"""
    return list(base_dict(d).values())


def dict_keys(d: dict):
    """Returns a list containing all keys in dict d"""
    return list(base_dict(d).keys())


def to_tuple(lists: list):
    """Convert from lists to unpacked tuple

    Allows us to write `X, y = to_tuple([[x1, y1], [x2, y2], [x3, y3]])`

    Parameters
    ----------
    lists: list
        list of objects to convert to unpacked tuple

    Examples
    --------
    >>> to_tuple([[x1, y1], [x2, y2], [x3, y3]])
    ([x1, x2, x3], [y1, y2, y3])
    >>> to_tuple([[x1, y1]])
    ([x1], [y1])
    >>> to_tuple([m1, m2, m3])
    [m1, m2, m3]
    """
    n_mods = len(lists)
    if n_mods <= 1:
        return lists
    if not isinstance(lists[0], list):
        return lists
    n_tup = len(lists[0])
    tup = [[] for _ in range(n_tup)]
    for i in range(n_mods):
        for j in range(n_tup):
            tup[j].append(lists[i][j])
    return tuple(tup)


def to_list(tup: tuple):
    """Convert from tuple to packed list

    Allows us to call function with arguments in a loop

    Parameters
    ----------
    tup: tuple
        tuple of objects to convert to packed list

    Raises
    ------
    ValueError
        If passed uneven number of arguments without a list. Please wrap your args in a list.

    Examples
    --------
    >>> to_list(([x1, x2, x3], [y1, y2, y3]))
    [[x1, y1], [x2, y2], [x3, y3]]
    >>> to_list(([x1], [y1]))
    [[x1, y1]]
    >>> to_list(([x1, x2, x3], ))
    [[x1], [x2], [x3]]
    >>> to_list((x1, ))
    [[x1]]
    >>> to_list((x1, y1))
    [[x1, y1]]
    >>> to_list((x1, x2, x3, y1, y2, y3))
    [[x1, y1], [x2, y2], [x3, y3]]
    """
    n_tup = len(tup)
    if n_tup == 0:
        return []
    if not isinstance(tup[0], list):
        # the first element is data
        if n_tup == 1:
            return [list(tup)]
        if n_tup % 2 != 0:
            raise ValueError(
                "Don't know how to handle uneven number of args "
                "without a list. Please wrap your args in a list."
            )
        # assume first half of args is input and second half is outcome
        return [list(el) for el in zip(tup[: (n_tup // 2)], tup[(n_tup // 2) :])]
    if n_tup == 1:
        return [[x] for x in tup[0]]
    n_mods = len(tup[0])
    lists_packed = [[] for _ in range(n_mods)]
    for i in range(n_mods):
        for j in range(n_tup):
            lists_packed[i].append(tup[j][i])
    return lists_packed


def sep_dicts(d: dict, n_out: int = 1, keys=None):
    """Converts dictionary with value being saved as an iterable into multiple dictionaries

    Assumes every value has same length n_out

    Parameters
    ----------
    d: dict
        Dictionary with iterable values to be converted.
    n_out: int, default 1
        The number of dictionaries to separate d into.
    keys: list-like, default None
        Optional list of keys to use in output dicts.

    Returns
    -------
    sep_dicts_list: list
        List of separated dictionaries.

    Examples
    --------
    >>> sep_dicts({k1: (x1, y1), k2: (x2, y2), ...,  '__prev__': p})
    [{k1: x1, k2: x2, ..., '__prev__': p}, {k1: y1, k2: y2, ..., '__prev__': p}]
    """
    if keys and len(keys) != n_out:
        raise ValueError(f"keys should be empty or have length n_out={n_out}")
    if n_out <= 1:
        return d  # nothing to separate
    sep_dicts_id = str(uuid4())  # w/ high prob, uuid4 is unique
    sep_dicts_list = [{} for _ in range(n_out)]
    for key, value in d.items():
        if key != PREV_KEY:
            for i in range(n_out):
                # assumes the correct sub-key for item i is in the i-th position
                if not keys:
                    new_key = (key[i],) + key[n_out:]
                else:
                    new_sub = Subkey(
                        value=keys[i], origin=key[-1].origin + "-" + str(i)
                    )
                    new_key = (new_sub,) + key
                new_key[-1].sep_dicts_id = sep_dicts_id
                if isinstance(value, VfuncPromise):
                    # a promise of item i of the original promise's value
                    value_i = VfuncPromise(operator.getitem, value, i)
                else:
                    value_i = value[i]
                sep_dicts_list[i][new_key] = value_i

    return sep_dicts_list


def dict_to_df(d: dict, param_key=None):
    """Converts a dictionary with tuple keys
    into a pandas DataFrame, optionally separating
    parameters in `param_key` if not None

    Parameters
    ----------
    d: dict
        Output dictionary with tuple keys from a Vset.
    param_key: str (optional), default None
        Name of parameter to separate into multiple columns.

    Returns
    -------
    df: pandas.DataFrame
        A DataFrame with `d` tuple keys separated into columns.
    """
    d_copy = {tuple(sk.value for sk in k): d[k] for k in d if k != PREV_KEY}
    df = pd.Series(d_copy).reset_index()
    if d_copy:
        first_key = next(k for k in d if k != PREV_KEY)
        cols = [sk.origin for sk in first_key] + ["out"]
        # set each init col to init-{next_vfunc_set}
        cols = [
            c if c != "init" else init_step(idx, cols) for idx, c in enumerate(cols)
        ]
        df = df.set_axis(cols, axis=1)
        if param_key:
            param_keys = df[param_key].tolist()
            if param_key == "out" and hasattr(param_keys[0], "__iter__"):
                param_df = pd.DataFrame(param_keys)
                param_df.columns = [f"{param_key}-{col}" for col in param_df.columns]
                df = df.join(param_df)
            else:
                param_loc = df.columns.get_loc(param_key)
                # parse every row on its own: vfuncs built from different funcs
                # can have different parameter names
                rows = []
                for t in param_keys:
                    if isinstance(t, str) or not all(
                        isinstance(s, str) and "=" in s for s in t
                    ):
                        raise ValueError(
                            f"param_key={param_key!r} needs keys that are tuples "
                            "of 'name=value' strings, as build_vset(param_dict=...) "
                            f"makes; got {t!r}"
                        )
                    # split on the first '=' only, so values may contain '='
                    rows.append(dict(s.split("=", 1) for s in t))
                param_df = pd.DataFrame(rows, index=df.index)
                param_df.columns = [f"{c}-{param_key}" for c in param_df.columns]
                df = pd.concat(
                    [df.iloc[:, :param_loc], param_df, df.iloc[:, param_loc + 1 :]],
                    axis=1,
                )
    return df


def perturbation_stats(
    data: Union[pd.DataFrame, dict],
    *group_by: str,
    wrt: str = "out",
    func=None,
    prefix: str = None,
    split: bool = False,
):
    """Compute statistics for `wrt` in `data`, conditional on `group_by`

    Parameters
    ----------
    data: Union[pandas.DataFrame, dict]
        DataFrame, as from calling `dict_to_df` on an output dict from a Vset,
        or the output dict itself.
    *group_by: str
        Vset names in `data` to group on. If none provided, treats everything as one big
        group.
    wrt: str (optional)
        Column name in `data` or `dict_to_df(data)` on which to compute statistics.
        Defaults to `'out'`, the values of the original Vset output dict.
    func: function, str, list or dict (optional), default None
        A list of functions or function names to use for computing
        statistics, analogous to the parameter of the same name in
        pandas.core.groupby.DataFrameGroupBy.aggregate. If `None`, defaults to
        `['count', 'mean', 'std']`.
    prefix: str (optional), default None
        A string to prefix to new columns in output DataFrame. If `None`,
        uses the value of `wrt`.
    split: bool (optional), default False
        If `True` and `wrt` in `data` has `list` or `numpy.ndarray` entries, will
        attempt to split the entries into multiple columns for the output.

    Returns
    -------
    df: pandas.DataFrame
        A DataFrame with summary statistics on `wrt`.
    """
    if func is None:
        func = ["count", "mean", "std"]
    if isinstance(func, str):
        func = [func]
    # for scalar outputs, columns are prefixed only when a prefix is asked for,
    # so existing callers keep their column names
    prefix_scalar_cols = prefix is not None
    if prefix is None:
        prefix = wrt
    if isinstance(data, dict):
        df = dict_to_df(data)
    else:
        df = data
    group_by = list(group_by)
    if len(group_by) == 1:
        # a bare column name keeps scalar group keys (iterating a one-item
        # list groupby gives tuple keys)
        gb = df.groupby(group_by[0])[wrt]
    elif len(group_by) > 1:
        gb = df.groupby(group_by)[wrt]
    else:
        gb = df.groupby(lambda x: True)[wrt]
    array_entries = type(df[wrt].iloc[0]) in (list, np.ndarray)
    if array_entries and isinstance(func, list) and ("mean" in func or "std" in func):
        unsupported = [f for f in func if f not in ("count", "mean", "std")]
        if unsupported:
            raise ValueError(
                f"for list or array entries in {wrt!r}, func supports only "
                f"'count', 'mean' and 'std'; got {unsupported}"
            )
        # iterating the groups gives the same order as gb.agg; get_group with
        # a single group key changed meaning in pandas 3
        wrt_arrays = [np.stack(d.tolist()) for _, d in gb]
        df_out = pd.DataFrame(gb.agg("count"))
        df_out.columns = [f"{prefix}-count"]
        stats = {
            "mean": lambda arr: arr.mean(axis=0),
            "std": lambda arr: arr.std(axis=0, ddof=1),
        }
        for stat, compute in stats.items():
            if stat not in func:
                continue
            per_group = [compute(arr) for arr in wrt_arrays]
            # build on df_out's index: group keys may be tuples, which a fresh
            # index would turn into a MultiIndex
            if split:
                # one column per entry position: out0-mean, ... for 1-D entries,
                # out0_1-mean, ... for 2-D entries such as predict_proba output
                cols = [
                    f"{prefix}{'_'.join(map(str, ix))}-{stat}"
                    for ix in np.ndindex(per_group[0].shape)
                ]
                block = pd.DataFrame(
                    [v.reshape(-1) for v in per_group], columns=cols, index=df_out.index
                )
            else:
                cells = np.empty(len(per_group), dtype=object)
                for i, v in enumerate(per_group):
                    cells[i] = v
                block = pd.DataFrame({f"{prefix}-{stat}": cells}, index=df_out.index)
            df_out = df_out.join(block)
        if "count" not in func:
            df_out = df_out.drop(columns=f"{prefix}-count")
    else:
        df_out = gb.agg(func)
        if isinstance(df_out, pd.Series):  # a single callable gives a Series
            df_out = df_out.to_frame(getattr(func, "__name__", wrt))
        if prefix_scalar_cols:
            df_out.columns = [f"{prefix}-{c}" for c in df_out.columns]
    df_out = df_out.reindex(sorted(df_out.columns), axis=1)
    df_out.reset_index(inplace=True)
    if len(group_by) > 0:
        return df_out.sort_values(group_by[0])
    return df_out


def combine_keys(left_key, right_key):
    """Combines `left_key` and `right_key`, attempting to match on any `Subkey` where
    `vflow.subkey.Subkey.is_matching` is `True`.

    Returns an empty key on failed matches when
    `vflow.subkey.Subkey.mismatches` is `True`. Always filters on `right_key`
    and returns `combined_key` with `left_key` prefix.

    Parameters
    ----------
    left_key: tuple
        Left tuple key to combine.
    right_key: tuple
        Right tuple key to combine.

    Returns
    -------
    combined_key: tuple
        Combined tuple key filtered according to `vflow.subkey.Subkey.matches` rules,
        which is empty according to `vflow.subkey.Subkey.mismatches` rule.

    """
    # look for the shorter key's matching subkeys in the longer key
    if len(left_key) < len(right_key):
        match_key, compare_key = left_key, right_key
    else:
        match_key, compare_key = right_key, left_key
    matched = []
    for subkey in (s for s in match_key if s.is_matching()):
        for c_subkey in compare_key:
            if subkey.matches(c_subkey):
                matched.append(subkey)
                break
            if subkey.mismatches(c_subkey):
                # subkeys with same origin but different values are rejected
                return ()
    # a matched subkey appears once: drop it from the right key
    return left_key + tuple(s for s in right_key if s not in matched)


def combine_dicts(*args: dict, base_case=True):
    """Combines any number of dictionaries into a single dictionary. Dictionaries
    are combined left to right matching all keys according to `combine_keys`

    Parameters
    ----------
    *args: dict
        Dictionaries to recursively combine left to right.
    base_case: bool, default True
        Internal: True when the values of `args[0]` are single values; False
        when they are tuples already combined from earlier dicts.

    Returns
    -------
    combined_dict: dict
        Combined dictionary. Its values are tuples of arguments, so a vfunc
        can be called with `vfunc(*value)`.
    """
    n_args = len(args)
    combined_dict = {}
    if n_args == 0:
        return combined_dict
    if n_args == 1:
        return {k: v if k == PREV_KEY else (v,) for k, v in args[0].items()}
    if n_args == 2:
        for k0 in args[0]:
            for k1 in args[1]:
                if PREV_KEY in (k0, k1):
                    continue

                combined_key = combine_keys(k0, k1)

                if len(combined_key) > 0:
                    if base_case:
                        combined_dict[combined_key] = (args[0][k0], args[1][k1])
                    else:
                        combined_dict[combined_key] = args[0][k0] + (args[1][k1],)

        return combined_dict
    # combine the first two dicts and call recursively with remaining args
    return combine_dicts(combine_dicts(args[0], args[1]), *args[2:], base_case=False)


def apply_vfuncs(vfuncs: dict, data_dict: dict, lazy: bool = False):
    """Apply a dictionary of functions `vfuncs` to each item of `data_dict`,
    optionally returning a dictionary of `vflow.vfunc.VfuncPromise` objects if `lazy` is True

    Output keys are determined by applying `combine_keys` to each pair of items from
    `vfuncs` and `data_dict`. This function is used by all Vsets to apply functions.

    Parameters
    ----------
    vfuncs: dict
        Dictionary of functions to apply to `data_dict`.
    data_dict: dict
        Dictionary of parameters to call each function in `vfuncs`.
    lazy: bool (optional), default False
        If True, `vfuncs` are applied lazily, returning `vflow.vfunc.VfuncPromise`
        objects.

    Returns
    -------
    out_dict: dict
        Output dictionary of applying `vfuncs` to `data_dict`.
    """
    if not data_dict:
        # no inputs: call each vfunc with no arguments
        return {
            k: VfuncPromise(deepcopy(f)) if lazy else deepcopy(f)()
            for k, f in vfuncs.items()
        }
    out_dict = {}
    for vf_k in vfuncs:
        for data_k in data_dict:
            if PREV_KEY in (vf_k, data_k):
                continue

            combined_key = combine_keys(data_k, vf_k)
            if not combined_key:
                continue

            func = deepcopy(vfuncs[vf_k])
            if lazy:
                # return a promise
                out_dict[combined_key] = VfuncPromise(func, *data_dict[data_k])
            else:
                data_list = list(data_dict[data_k])
                for i, data in enumerate(data_list):
                    if isinstance(data, VfuncPromise):
                        data_list[i] = data()
                    if isinstance(func, RayRemoteFun) and not isinstance(
                        data_list[i], ray.ObjectRef
                    ):
                        # send data to Ray's remote object store
                        data_list[i] = ray.put(data_list[i])
                    elif isinstance(data_list[i], ray.ObjectRef):
                        # this is not a remote function so get the data
                        data_list[i] = ray.get(data_list[i])
                out_dict[combined_key] = func(*data_list)

    return out_dict

Functions

def apply_vfuncs(vfuncs: dict, data_dict: dict, lazy: bool = False)

Apply a dictionary of functions vfuncs to each item of data_dict, optionally returning a dictionary of VfuncPromise objects if lazy is True

Output keys are determined by applying combine_keys() to each pair of items from vfuncs and data_dict. This function is used by all Vsets to apply functions.

Parameters

vfuncs : dict

Dictionary of functions to apply to data_dict.

data_dict : dict

Dictionary of parameters to call each function in vfuncs.

lazy : bool (optional), default False

If True, vfuncs are applied lazily, returning VfuncPromise objects.

Returns

out_dict : dict

Output dictionary of applying vfuncs to data_dict.

Expand source code

def apply_vfuncs(vfuncs: dict, data_dict: dict, lazy: bool = False):
    """Apply a dictionary of functions `vfuncs` to each item of `data_dict`,
    optionally returning a dictionary of `vflow.vfunc.VfuncPromise` objects if `lazy` is True

    Output keys are determined by applying `combine_keys` to each pair of items from
    `vfuncs` and `data_dict`. This function is used by all Vsets to apply functions.

    Parameters
    ----------
    vfuncs: dict
        Dictionary of functions to apply to `data_dict`.
    data_dict: dict
        Dictionary of parameters to call each function in `vfuncs`.
    lazy: bool (optional), default False
        If True, `vfuncs` are applied lazily, returning `vflow.vfunc.VfuncPromise`
        objects.

    Returns
    -------
    out_dict: dict
        Output dictionary of applying `vfuncs` to `data_dict`.
    """
    if not data_dict:
        # no inputs: call each vfunc with no arguments
        return {
            k: VfuncPromise(deepcopy(f)) if lazy else deepcopy(f)()
            for k, f in vfuncs.items()
        }
    out_dict = {}
    for vf_k in vfuncs:
        for data_k in data_dict:
            if PREV_KEY in (vf_k, data_k):
                continue

            combined_key = combine_keys(data_k, vf_k)
            if not combined_key:
                continue

            func = deepcopy(vfuncs[vf_k])
            if lazy:
                # return a promise
                out_dict[combined_key] = VfuncPromise(func, *data_dict[data_k])
            else:
                data_list = list(data_dict[data_k])
                for i, data in enumerate(data_list):
                    if isinstance(data, VfuncPromise):
                        data_list[i] = data()
                    if isinstance(func, RayRemoteFun) and not isinstance(
                        data_list[i], ray.ObjectRef
                    ):
                        # send data to Ray's remote object store
                        data_list[i] = ray.put(data_list[i])
                    elif isinstance(data_list[i], ray.ObjectRef):
                        # this is not a remote function so get the data
                        data_list[i] = ray.get(data_list[i])
                out_dict[combined_key] = func(*data_list)

    return out_dict

def base_dict(d: dict)

Remove PREV_KEY from dict d if present

Expand source code

def base_dict(d: dict):
    """Remove PREV_KEY from dict d if present"""
    return {k: v for k, v in d.items() if k != PREV_KEY}

def combine_dicts(*args: dict, base_case=True)

Combines any number of dictionaries into a single dictionary. Dictionaries are combined left to right matching all keys according to combine_keys()

Parameters

*args : dict

Dictionaries to recursively combine left to right.

base_case : bool, default True

Internal: True when the values of args[0] are single values; False when they are tuples already combined from earlier dicts.

Returns

combined_dict : dict

Combined dictionary. Its values are tuples of arguments, so a vfunc can be called with vfunc(*value).

Expand source code

def combine_dicts(*args: dict, base_case=True):
    """Combines any number of dictionaries into a single dictionary. Dictionaries
    are combined left to right matching all keys according to `combine_keys`

    Parameters
    ----------
    *args: dict
        Dictionaries to recursively combine left to right.
    base_case: bool, default True
        Internal: True when the values of `args[0]` are single values; False
        when they are tuples already combined from earlier dicts.

    Returns
    -------
    combined_dict: dict
        Combined dictionary. Its values are tuples of arguments, so a vfunc
        can be called with `vfunc(*value)`.
    """
    n_args = len(args)
    combined_dict = {}
    if n_args == 0:
        return combined_dict
    if n_args == 1:
        return {k: v if k == PREV_KEY else (v,) for k, v in args[0].items()}
    if n_args == 2:
        for k0 in args[0]:
            for k1 in args[1]:
                if PREV_KEY in (k0, k1):
                    continue

                combined_key = combine_keys(k0, k1)

                if len(combined_key) > 0:
                    if base_case:
                        combined_dict[combined_key] = (args[0][k0], args[1][k1])
                    else:
                        combined_dict[combined_key] = args[0][k0] + (args[1][k1],)

        return combined_dict
    # combine the first two dicts and call recursively with remaining args
    return combine_dicts(combine_dicts(args[0], args[1]), *args[2:], base_case=False)

def combine_keys(left_key, right_key)

Combines left_key and right_key, attempting to match on any Subkey where Subkey.is_matching() is True.

Returns an empty key on failed matches when Subkey.mismatches() is True. Always filters on right_key and returns combined_key with left_key prefix.

Parameters

left_key : tuple

Left tuple key to combine.

right_key : tuple

Right tuple key to combine.

Returns

combined_key : tuple

Combined tuple key filtered according to Subkey.matches() rules, which is empty according to Subkey.mismatches() rule.

Expand source code

def combine_keys(left_key, right_key):
    """Combines `left_key` and `right_key`, attempting to match on any `Subkey` where
    `vflow.subkey.Subkey.is_matching` is `True`.

    Returns an empty key on failed matches when
    `vflow.subkey.Subkey.mismatches` is `True`. Always filters on `right_key`
    and returns `combined_key` with `left_key` prefix.

    Parameters
    ----------
    left_key: tuple
        Left tuple key to combine.
    right_key: tuple
        Right tuple key to combine.

    Returns
    -------
    combined_key: tuple
        Combined tuple key filtered according to `vflow.subkey.Subkey.matches` rules,
        which is empty according to `vflow.subkey.Subkey.mismatches` rule.

    """
    # look for the shorter key's matching subkeys in the longer key
    if len(left_key) < len(right_key):
        match_key, compare_key = left_key, right_key
    else:
        match_key, compare_key = right_key, left_key
    matched = []
    for subkey in (s for s in match_key if s.is_matching()):
        for c_subkey in compare_key:
            if subkey.matches(c_subkey):
                matched.append(subkey)
                break
            if subkey.mismatches(c_subkey):
                # subkeys with same origin but different values are rejected
                return ()
    # a matched subkey appears once: drop it from the right key
    return left_key + tuple(s for s in right_key if s not in matched)

def dict_data(d: dict)

Returns a list containing all data in dict d

Expand source code

def dict_data(d: dict):
    """Returns a list containing all data in dict d"""
    return list(base_dict(d).values())

def dict_keys(d: dict)

Returns a list containing all keys in dict d

Expand source code

def dict_keys(d: dict):
    """Returns a list containing all keys in dict d"""
    return list(base_dict(d).keys())

def dict_to_df(d: dict, param_key=None)

Converts a dictionary with tuple keys into a pandas DataFrame, optionally separating parameters in param_key if not None

Parameters

d : dict

Output dictionary with tuple keys from a Vset.

param_key : str (optional), default None

Name of parameter to separate into multiple columns.

Returns

df : pandas.DataFrame

A DataFrame with d tuple keys separated into columns.

Expand source code

def dict_to_df(d: dict, param_key=None):
    """Converts a dictionary with tuple keys
    into a pandas DataFrame, optionally separating
    parameters in `param_key` if not None

    Parameters
    ----------
    d: dict
        Output dictionary with tuple keys from a Vset.
    param_key: str (optional), default None
        Name of parameter to separate into multiple columns.

    Returns
    -------
    df: pandas.DataFrame
        A DataFrame with `d` tuple keys separated into columns.
    """
    d_copy = {tuple(sk.value for sk in k): d[k] for k in d if k != PREV_KEY}
    df = pd.Series(d_copy).reset_index()
    if d_copy:
        first_key = next(k for k in d if k != PREV_KEY)
        cols = [sk.origin for sk in first_key] + ["out"]
        # set each init col to init-{next_vfunc_set}
        cols = [
            c if c != "init" else init_step(idx, cols) for idx, c in enumerate(cols)
        ]
        df = df.set_axis(cols, axis=1)
        if param_key:
            param_keys = df[param_key].tolist()
            if param_key == "out" and hasattr(param_keys[0], "__iter__"):
                param_df = pd.DataFrame(param_keys)
                param_df.columns = [f"{param_key}-{col}" for col in param_df.columns]
                df = df.join(param_df)
            else:
                param_loc = df.columns.get_loc(param_key)
                # parse every row on its own: vfuncs built from different funcs
                # can have different parameter names
                rows = []
                for t in param_keys:
                    if isinstance(t, str) or not all(
                        isinstance(s, str) and "=" in s for s in t
                    ):
                        raise ValueError(
                            f"param_key={param_key!r} needs keys that are tuples "
                            "of 'name=value' strings, as build_vset(param_dict=...) "
                            f"makes; got {t!r}"
                        )
                    # split on the first '=' only, so values may contain '='
                    rows.append(dict(s.split("=", 1) for s in t))
                param_df = pd.DataFrame(rows, index=df.index)
                param_df.columns = [f"{c}-{param_key}" for c in param_df.columns]
                df = pd.concat(
                    [df.iloc[:, :param_loc], param_df, df.iloc[:, param_loc + 1 :]],
                    axis=1,
                )
    return df

def init_step(idx, cols)

Name for the 'init' column at idx: "init-" plus the first non-init column after it (the Vset the data went into), or None if there is none.

Parameters

idx : int

Index of 'init' column in cols.

cols : list[str]

List of column names.

Expand source code

def init_step(idx, cols):
    """Name for the 'init' column at `idx`: "init-" plus the first non-init
    column after it (the Vset the data went into), or None if there is none.

    Parameters
    ----------
    idx: int
        Index of 'init' column in cols.
    cols: list[str]
        List of column names.
    """
    for i in range(idx, len(cols)):
        if cols[i] != "init":
            return "init-" + cols[i]
    return None

def perturbation_stats(data: pandas.DataFrame | dict, *group_by: str, wrt: str = 'out', func=None, prefix: str = None, split: bool = False)

Compute statistics for wrt in data, conditional on group_by

Parameters

data : Union[pandas.DataFrame, dict]

DataFrame, as from calling dict_to_df() on an output dict from a Vset, or the output dict itself.

*group_by : str

Vset names in data to group on. If none provided, treats everything as one big group.

wrt : str (optional)

Column name in data or dict_to_df()(data) on which to compute statistics. Defaults to 'out', the values of the original Vset output dict.

func : function, str, list or dict (optional), default None

A list of functions or function names to use for computing statistics, analogous to the parameter of the same name in pandas.core.groupby.DataFrameGroupBy.aggregate. If None, defaults to ['count', 'mean', 'std'].

prefix : str (optional), default None

A string to prefix to new columns in output DataFrame. If None, uses the value of wrt.

split : bool (optional), default False

If True and wrt in data has list or numpy.ndarray entries, will attempt to split the entries into multiple columns for the output.

Returns

df : pandas.DataFrame

A DataFrame with summary statistics on wrt.

Expand source code

def perturbation_stats(
    data: Union[pd.DataFrame, dict],
    *group_by: str,
    wrt: str = "out",
    func=None,
    prefix: str = None,
    split: bool = False,
):
    """Compute statistics for `wrt` in `data`, conditional on `group_by`

    Parameters
    ----------
    data: Union[pandas.DataFrame, dict]
        DataFrame, as from calling `dict_to_df` on an output dict from a Vset,
        or the output dict itself.
    *group_by: str
        Vset names in `data` to group on. If none provided, treats everything as one big
        group.
    wrt: str (optional)
        Column name in `data` or `dict_to_df(data)` on which to compute statistics.
        Defaults to `'out'`, the values of the original Vset output dict.
    func: function, str, list or dict (optional), default None
        A list of functions or function names to use for computing
        statistics, analogous to the parameter of the same name in
        pandas.core.groupby.DataFrameGroupBy.aggregate. If `None`, defaults to
        `['count', 'mean', 'std']`.
    prefix: str (optional), default None
        A string to prefix to new columns in output DataFrame. If `None`,
        uses the value of `wrt`.
    split: bool (optional), default False
        If `True` and `wrt` in `data` has `list` or `numpy.ndarray` entries, will
        attempt to split the entries into multiple columns for the output.

    Returns
    -------
    df: pandas.DataFrame
        A DataFrame with summary statistics on `wrt`.
    """
    if func is None:
        func = ["count", "mean", "std"]
    if isinstance(func, str):
        func = [func]
    # for scalar outputs, columns are prefixed only when a prefix is asked for,
    # so existing callers keep their column names
    prefix_scalar_cols = prefix is not None
    if prefix is None:
        prefix = wrt
    if isinstance(data, dict):
        df = dict_to_df(data)
    else:
        df = data
    group_by = list(group_by)
    if len(group_by) == 1:
        # a bare column name keeps scalar group keys (iterating a one-item
        # list groupby gives tuple keys)
        gb = df.groupby(group_by[0])[wrt]
    elif len(group_by) > 1:
        gb = df.groupby(group_by)[wrt]
    else:
        gb = df.groupby(lambda x: True)[wrt]
    array_entries = type(df[wrt].iloc[0]) in (list, np.ndarray)
    if array_entries and isinstance(func, list) and ("mean" in func or "std" in func):
        unsupported = [f for f in func if f not in ("count", "mean", "std")]
        if unsupported:
            raise ValueError(
                f"for list or array entries in {wrt!r}, func supports only "
                f"'count', 'mean' and 'std'; got {unsupported}"
            )
        # iterating the groups gives the same order as gb.agg; get_group with
        # a single group key changed meaning in pandas 3
        wrt_arrays = [np.stack(d.tolist()) for _, d in gb]
        df_out = pd.DataFrame(gb.agg("count"))
        df_out.columns = [f"{prefix}-count"]
        stats = {
            "mean": lambda arr: arr.mean(axis=0),
            "std": lambda arr: arr.std(axis=0, ddof=1),
        }
        for stat, compute in stats.items():
            if stat not in func:
                continue
            per_group = [compute(arr) for arr in wrt_arrays]
            # build on df_out's index: group keys may be tuples, which a fresh
            # index would turn into a MultiIndex
            if split:
                # one column per entry position: out0-mean, ... for 1-D entries,
                # out0_1-mean, ... for 2-D entries such as predict_proba output
                cols = [
                    f"{prefix}{'_'.join(map(str, ix))}-{stat}"
                    for ix in np.ndindex(per_group[0].shape)
                ]
                block = pd.DataFrame(
                    [v.reshape(-1) for v in per_group], columns=cols, index=df_out.index
                )
            else:
                cells = np.empty(len(per_group), dtype=object)
                for i, v in enumerate(per_group):
                    cells[i] = v
                block = pd.DataFrame({f"{prefix}-{stat}": cells}, index=df_out.index)
            df_out = df_out.join(block)
        if "count" not in func:
            df_out = df_out.drop(columns=f"{prefix}-count")
    else:
        df_out = gb.agg(func)
        if isinstance(df_out, pd.Series):  # a single callable gives a Series
            df_out = df_out.to_frame(getattr(func, "__name__", wrt))
        if prefix_scalar_cols:
            df_out.columns = [f"{prefix}-{c}" for c in df_out.columns]
    df_out = df_out.reindex(sorted(df_out.columns), axis=1)
    df_out.reset_index(inplace=True)
    if len(group_by) > 0:
        return df_out.sort_values(group_by[0])
    return df_out

def sep_dicts(d: dict, n_out: int = 1, keys=None)

Converts dictionary with value being saved as an iterable into multiple dictionaries

Assumes every value has same length n_out

Parameters

d : dict

Dictionary with iterable values to be converted.

n_out : int, default 1

The number of dictionaries to separate d into.

keys : list-like, default None

Optional list of keys to use in output dicts.

Returns

sep_dicts_list : list

List of separated dictionaries.

Examples

>>> sep_dicts({k1: (x1, y1), k2: (x2, y2), ...,  '__prev__': p})
[{k1: x1, k2: x2, ..., '__prev__': p}, {k1: y1, k2: y2, ..., '__prev__': p}]

Expand source code

def sep_dicts(d: dict, n_out: int = 1, keys=None):
    """Converts dictionary with value being saved as an iterable into multiple dictionaries

    Assumes every value has same length n_out

    Parameters
    ----------
    d: dict
        Dictionary with iterable values to be converted.
    n_out: int, default 1
        The number of dictionaries to separate d into.
    keys: list-like, default None
        Optional list of keys to use in output dicts.

    Returns
    -------
    sep_dicts_list: list
        List of separated dictionaries.

    Examples
    --------
    >>> sep_dicts({k1: (x1, y1), k2: (x2, y2), ...,  '__prev__': p})
    [{k1: x1, k2: x2, ..., '__prev__': p}, {k1: y1, k2: y2, ..., '__prev__': p}]
    """
    if keys and len(keys) != n_out:
        raise ValueError(f"keys should be empty or have length n_out={n_out}")
    if n_out <= 1:
        return d  # nothing to separate
    sep_dicts_id = str(uuid4())  # w/ high prob, uuid4 is unique
    sep_dicts_list = [{} for _ in range(n_out)]
    for key, value in d.items():
        if key != PREV_KEY:
            for i in range(n_out):
                # assumes the correct sub-key for item i is in the i-th position
                if not keys:
                    new_key = (key[i],) + key[n_out:]
                else:
                    new_sub = Subkey(
                        value=keys[i], origin=key[-1].origin + "-" + str(i)
                    )
                    new_key = (new_sub,) + key
                new_key[-1].sep_dicts_id = sep_dicts_id
                if isinstance(value, VfuncPromise):
                    # a promise of item i of the original promise's value
                    value_i = VfuncPromise(operator.getitem, value, i)
                else:
                    value_i = value[i]
                sep_dicts_list[i][new_key] = value_i

    return sep_dicts_list

def to_list(tup: tuple)

Convert from tuple to packed list

Allows us to call function with arguments in a loop

Parameters

tup : tuple

tuple of objects to convert to packed list

Raises

ValueError

If passed uneven number of arguments without a list. Please wrap your args in a list.

Examples

>>> to_list(([x1, x2, x3], [y1, y2, y3]))
[[x1, y1], [x2, y2], [x3, y3]]
>>> to_list(([x1], [y1]))
[[x1, y1]]
>>> to_list(([x1, x2, x3], ))
[[x1], [x2], [x3]]
>>> to_list((x1, ))
[[x1]]
>>> to_list((x1, y1))
[[x1, y1]]
>>> to_list((x1, x2, x3, y1, y2, y3))
[[x1, y1], [x2, y2], [x3, y3]]

Expand source code

def to_list(tup: tuple):
    """Convert from tuple to packed list

    Allows us to call function with arguments in a loop

    Parameters
    ----------
    tup: tuple
        tuple of objects to convert to packed list

    Raises
    ------
    ValueError
        If passed uneven number of arguments without a list. Please wrap your args in a list.

    Examples
    --------
    >>> to_list(([x1, x2, x3], [y1, y2, y3]))
    [[x1, y1], [x2, y2], [x3, y3]]
    >>> to_list(([x1], [y1]))
    [[x1, y1]]
    >>> to_list(([x1, x2, x3], ))
    [[x1], [x2], [x3]]
    >>> to_list((x1, ))
    [[x1]]
    >>> to_list((x1, y1))
    [[x1, y1]]
    >>> to_list((x1, x2, x3, y1, y2, y3))
    [[x1, y1], [x2, y2], [x3, y3]]
    """
    n_tup = len(tup)
    if n_tup == 0:
        return []
    if not isinstance(tup[0], list):
        # the first element is data
        if n_tup == 1:
            return [list(tup)]
        if n_tup % 2 != 0:
            raise ValueError(
                "Don't know how to handle uneven number of args "
                "without a list. Please wrap your args in a list."
            )
        # assume first half of args is input and second half is outcome
        return [list(el) for el in zip(tup[: (n_tup // 2)], tup[(n_tup // 2) :])]
    if n_tup == 1:
        return [[x] for x in tup[0]]
    n_mods = len(tup[0])
    lists_packed = [[] for _ in range(n_mods)]
    for i in range(n_mods):
        for j in range(n_tup):
            lists_packed[i].append(tup[j][i])
    return lists_packed

def to_tuple(lists: list)

Convert from lists to unpacked tuple

Allows us to write X, y = to_tuple([[x1, y1], [x2, y2], [x3, y3]])

Parameters

lists : list

list of objects to convert to unpacked tuple

Examples

>>> to_tuple([[x1, y1], [x2, y2], [x3, y3]])
([x1, x2, x3], [y1, y2, y3])
>>> to_tuple([[x1, y1]])
([x1], [y1])
>>> to_tuple([m1, m2, m3])
[m1, m2, m3]

Expand source code

def to_tuple(lists: list):
    """Convert from lists to unpacked tuple

    Allows us to write `X, y = to_tuple([[x1, y1], [x2, y2], [x3, y3]])`

    Parameters
    ----------
    lists: list
        list of objects to convert to unpacked tuple

    Examples
    --------
    >>> to_tuple([[x1, y1], [x2, y2], [x3, y3]])
    ([x1, x2, x3], [y1, y2, y3])
    >>> to_tuple([[x1, y1]])
    ([x1], [y1])
    >>> to_tuple([m1, m2, m3])
    [m1, m2, m3]
    """
    n_mods = len(lists)
    if n_mods <= 1:
        return lists
    if not isinstance(lists[0], list):
        return lists
    n_tup = len(lists[0])
    tup = [[] for _ in range(n_tup)]
    for i in range(n_mods):
        for j in range(n_tup):
            tup[j].append(lists[i][j])
    return tuple(tup)