Utilities & Typing

These are internal Utilities and variable types for type hinting. The end user is unlikely to need to worry about these, but these may of interest to developers

litmus._utils

These are internal utilities and convenience functions for use in LITMUS.

utils.py Handy internal utilities for brevity and convenience. Nothing in here is accesible in the public _init file

suppress_stdout()

Source code in litmus/_utils.py

@contextmanager
def suppress_stdout():
    ipython = get_ipython()

    if ipython and hasattr(ipython, 'kernel'):  # Likely in a Jupyter notebook
        with io.capture_output() as captured:
            yield
    else:
        # Standard Python environment
        original_stdout_fd = os.dup(sys.stdout.fileno())
        with open(os.devnull, 'w') as devnull:
            os.dup2(devnull.fileno(), sys.stdout.fileno())
            try:
                yield
            finally:
                os.dup2(original_stdout_fd, sys.stdout.fileno())
                os.close(original_stdout_fd)

isiter(x: any) -> bool

Checks to see if an object is itterable

Source code in litmus/_utils.py

def isiter(x: any) -> bool:
    """
    Checks to see if an object is itterable
    """
    if type(x) == dict:
        return len(x[list(x.keys())[0]]) > 1
    try:
        iter(x)
    except:
        return (False)
    else:
        return (True)

isiter_dict(DICT: dict) -> bool

like isiter but for a dictionary. Checks only the first element in DICT.keys

Source code in litmus/_utils.py

def isiter_dict(DICT: dict) -> bool:
    """
    like isiter but for a dictionary. Checks only the first element in DICT.keys
    """

    key = list(DICT.keys())[0]
    if isiter(DICT[key]):
        return True
    else:
        return False

dict_dim(DICT: dict) -> (int, int)

Checks the first element of a dictionary and returns its length

Source code in litmus/_utils.py

def dict_dim(DICT: dict) -> (int, int):
    """
    Checks the first element of a dictionary and returns its length
    """

    if isiter_dict(DICT):
        firstkey = list(DICT.keys())[0]
        return (len(list(DICT.keys())), len(DICT[firstkey]))
    else:
        return (len(list(DICT.keys())), 1)

dict_pack(DICT: dict, keys=None, recursive=True, H=None, d0={}) -> np.array

Packs a dictionary into an array format

Parameters:

Name	Type	Description	Default
DICT	dict	the dict to unpack	required
keys		the order in which to index the keyed elements. If none, will use DICT.keys(). Can be partial	None
recursive		whether to recurse into arrays	True
H		Matrix to scale parameters by	None
d0		Value to offset by before packing	{}

Returns:

Type	Description
array	(nkeys x len_array) np.arrayobject X = H (d-d0)

Source code in litmus/_utils.py

def dict_pack(DICT: dict, keys=None, recursive=True, H=None, d0={}) -> np.array:
    """
    Packs a dictionary into an array format
    :param DICT: the dict to unpack
    :param keys: the order in which to index the keyed elements. If none, will use DICT.keys(). Can be partial
    :param recursive: whether to recurse into arrays
    :param H: Matrix to scale parameters by
    :param d0: Value to offset by before packing
    :return: (nkeys x len_array) np.arrayobject

    X = H (d-d0)
    """

    nokeys = True if keys is None else 0
    keys = keys if keys is not None else DICT.keys()

    if d0 is {}: d0 = {key:0 for key in keys}

    for key in keys:
        if key in DICT.keys() and key not in d0.keys(): d0 |= {key: 0.0}

    if recursive and type(list(DICT.values())[0]) == dict:
        out = np.array(
            [dict_pack(DICT[key] - d0[key], keys=keys if not nokeys else None, recursive=recursive) for key in keys])
    else:
        if isiter(DICT[list(keys)[0]]):
            out = np.array([[DICT[key][i] - d0[key] for i in range(dict_dim(DICT)[1])] for key in keys])
        else:
            out = np.array([DICT[key] - d0[key] for key in keys])

    return (out)

dict_unpack(X: np.array, keys: [str], recursive=True, Hinv=None, x0=None) -> np.array

Unpacks an array into a dict

Parameters:

Name	Type	Description	Default
X	array	Array to unpack	required
keys	[str]	keys to unpack with	required

Returns:

Type	Description
array	Hinv(X) + x0

Source code in litmus/_utils.py

def dict_unpack(X: np.array, keys: [str], recursive=True, Hinv=None, x0=None) -> np.array:
    """
    Unpacks an array into a dict
    :param X: Array to unpack
    :param keys: keys to unpack with
    :return:

    Hinv(X) + x0
    """
    if Hinv is not None: assert Hinv.shape[0] == len(keys), "Size of H must be equal to number of keys in dict_unpack"

    if recursive and isiter(X[0]):
        out = {key: dict_unpack(X[i], keys, recursive) for i, key in enumerate(list(keys))}
    else:
        X = X.copy()
        if Hinv is not None:
            X = np.dot(Hinv, X)
        if x0 is not None:
            X += x0
        out = {key: X[i] for i, key in enumerate(list(keys))}

    return (out)

dict_sortby(A: dict, B: dict, match_only=True) -> dict

Sorts dict A to match keys of dict B.

Parameters:

Name	Type	Description	Default
A	dict	Dict to be sorted	required
B	dict	Dict whose keys are will provide the ordering	required
match_only		If true, returns only for keys common to both A and B. Else, append un-sorted entries to end	True

Returns:

Type	Description
dict	{key: A[key] for key in B if key in A}

Source code in litmus/_utils.py

def dict_sortby(A: dict, B: dict, match_only=True) -> dict:
    """
    Sorts dict A to match keys of dict B.

    :param A: Dict to be sorted
    :param B: Dict whose keys are will provide the ordering
    :param match_only: If true, returns only for keys common to both A and B. Else, append un-sorted entries to end
    :return: {key: A[key] for key in B if key in A}
    """
    out = {key: A[key] for key in B if key in A}
    if not match_only:
        out |= {key: A[key] for key in A if key not in B}
    return (out)

dict_extend(A: dict, B: dict = None) -> dict

Extends all single-length entries of a dict to match the length of a non-singular element

Parameters:

Name	Type	Description	Default
A	dict	Dictionary whose elements are to be extended	required
B	dict	(optional) the array to extend by, equivalent to dict_extend(A|B)	None

Returns:

Type	Description
dict	Dict A with any singleton elements extended to the longest entry in A or B

Source code in litmus/_utils.py

def dict_extend(A: dict, B: dict = None) -> dict:
    """
    Extends all single-length entries of a dict to match the length of a non-singular element
    :param A: Dictionary whose elements are to be extended
    :param B: (optional) the array to extend by, equivalent to dict_extend(A|B)
    :return: Dict A with any singleton elements extended to the longest entry in A or B
    """

    out = A.copy()
    if B is not None: out |= B

    to_extend = [key for key in out if not isiter(out[key])]
    to_leave = [key for key in out if isiter(out[key])]

    if len(to_extend) == 0: return out
    if len(to_leave) == 0: return out

    N = len(out[to_leave[0]])
    for key in to_leave[1:]:
        assert len(out[key]) == N, "Tried to dict_extend() a dictionary with inhomogeneous lengths"

    for key in to_extend:
        out[key] = np.array([A[key]] * N)

    return (out)

dict_combine(X: [dict]) -> {str: [float]}

Combines an array, list etc. of dictionaries into a dictionary of arrays

Parameters:

Name	Type	Description	Default
X	[dict]	1D Iterable of dicts	required

Returns:

Type	Description
{str: [float]}	Dict of 1D iterables

Source code in litmus/_utils.py

def dict_combine(X: [dict]) -> {str: [float]}:
    """
    Combines an array, list etc. of dictionaries into a dictionary of arrays

    :param X: 1D Iterable of dicts
    :return: Dict of 1D iterables
    """

    N = len(X)
    keys = X[0].keys()

    out = {key: np.zeros(N) for key in keys}
    for n in range(N):
        for key in keys:
            out[key][n] = X[n][key]
    return (out)

dict_divide(X: dict) -> [dict]

Splits dict of arrays into array of dicts. Opposite of dict_combine

Parameters:

Name	Type	Description	Default
X	dict	Dict of 1D iterables	required

Returns:

Type	Description
[dict]	1D Iterable of dicts

Source code in litmus/_utils.py

def dict_divide(X: dict) -> [dict]:
    """
    Splits dict of arrays into array of dicts. Opposite of dict_combine

    :param X: Dict of 1D iterables
    :return: 1D Iterable of dicts
    """

    keys = list(X.keys())
    N = len(X[keys[0]])

    out = [{key: X[key][i] for key in X} for i in range(N)]

    return (out)

dict_split(X: dict, keys: [str]) -> (dict, dict)

Splits a dict in two based on keys

Parameters:

Name	Type	Description	Default
X	dict	Dict to be split into A,B	required
keys	[str]	Keys to be present in A, but not in B	required

Returns:

Type	Description
(dict, dict)	tuple of dicts (A,B)

Source code in litmus/_utils.py

def dict_split(X: dict, keys: [str]) -> (dict, dict):
    """
    Splits a dict in two based on keys

    :param X: Dict to be split into A,B
    :param keys: Keys to be present in A, but not in B
    :return: tuple of dicts (A,B)
    """
    assert type(X) is dict, "input to dict_split() must be of type dict"
    assert isiter(keys) and type(keys[0])==str, "in dict_split() keys must be list of strings"
    A = {key: X[key] for key in keys}
    B = {key: X[key] for key in X.keys() if key not in keys}
    return (A, B)

pack_function(func, packed_keys: [str], fixed_values: dict = {}, invert: bool = False, jit: bool = False, H: np.array = None, d0: dict = {}) -> _types.FunctionType

Re-arranges a function that takes dict arguments to tak array-like arguments instead, so as to be autograd friendly Takes a function f(D:dict, arg, kwargs) and returns f(X, D2, args, **kwargs), D2 is all elements of D not listed in 'packed_keys' or fixed_values.

Parameters:

Name	Type	Description	Default
func		Function to be unpacked	required
packed_keys	[str]	Keys in 'D' to be packed in an array	required
fixed_values	dict	Elements of 'D' to be fixed	{}
invert	bool	If true, will 'flip' the function upside down	False
jit	bool	If true, will 'jit' the function	False
H	array	(optional) scaling matrix to reparameterize H with	None
d0	dict	(optional) If given, will center the reparameterized function at x0	{}

Source code in litmus/_utils.py

def pack_function(func, packed_keys: ['str'], fixed_values: dict = {}, invert: bool = False, jit: bool = False,
                  H: np.array = None, d0: dict = {}) -> _types.FunctionType:
    """
    Re-arranges a function that takes dict arguments to tak array-like arguments instead, so as to be autograd friendly
    Takes a function f(D:dict, *arg, **kwargs) and returns f(X, D2, *args, **kwargs), D2 is all elements of D not
    listed in 'packed_keys' or fixed_values.

    :param func: Function to be unpacked
    :param packed_keys: Keys in 'D' to be packed in an array
    :param fixed_values: Elements of 'D' to be fixed
    :param invert:  If true, will 'flip' the function upside down
    :param jit: If true, will 'jit' the function
    :param H: (optional) scaling matrix to reparameterize H with
    :param d0: (optional) If given, will center the reparameterized  function at x0
    """

    if H is not None:
        assert H.shape[0] == len(packed_keys), "Scaling matrix H must be same length as packed_keys"
    else:
        H = jnp.eye(len(packed_keys))
    d0 = {key: 0.0 for key in packed_keys} | d0
    x0 = dict_pack(d0, packed_keys)

    # --------

    sign = -1 if invert else 1

    # --------
    def new_func(X, unpacked_params={}, *args, **kwargs):
        X = jnp.dot(H, X - x0)
        packed_dict = {key: x for key, x in zip(packed_keys, X)}
        packed_dict |= unpacked_params
        packed_dict |= fixed_values

        out = func(packed_dict, *args, **kwargs)
        return (sign * out)

    # --------
    if jit: new_func = jax.jit(new_func)

    return (new_func)

randint()

Quick utility to generate a random integer

Source code in litmus/_utils.py

def randint():
    """
    Quick utility to generate a random integer
    """
    return (np.random.randint(0, sys.maxsize // 1024))

litmus._types

Data types for type hinting. By convention in LITMUS, N indicates an over sample sites or observations while M indicates an index over number of parameters.

A single ref package to create and pull in types for type hinting

DType = TypeVar('DType', bound=generic) module-attribute

Any Dtype

ArrayN = Annotated[NDArray[DType], Literal['N']] module-attribute

1D array corresponding to N sample sites, e.g. an array of log densities

ArrayNxN = Annotated[NDArray[DType], Literal['N', 'N']] module-attribute

2D array corresponding to NxN sample sites, e.g. a GP covariance matrix

ArrayM = Annotated[NDArray[DType], Literal['M']] module-attribute

1D array corresponding to M parameters, e.g. a grad

ArrayMxM = Annotated[NDArray[DType], Literal['M', 'M']] module-attribute

2D array corresponding to MxM parameters, e.g. a hessian

ArrayNxMxM = Annotated[NDArray[DType], Literal['M', 'N', 'N']] module-attribute

3D array corresponding to N sheets of MxM arrays for N data points and M parameters, e.g. a plate of hessians