tree log

tree_log.py

def tree_log(tree, root, data, *, iter=None):
    if tree is None:
        return
    
    elif isinstance(tree, list):
        for (i, val) in enumerate(tree):
            tree_log(val, f"{root}/{i}", data, iter=iter)
    
    # handle namedtuples
    elif isinstance(tree, list) and hasattr(tree, "_fields"):
        tree_log(iter, f"{root}/iter", data)
        for key in tree._fields:
            tree_log(getattr(tree, key), f"root/{key}", data)

    elif isinstance(tree, tuple):
        tree_log(iter, f"{root}/iter", data)
        for (i, val) in enumerate(tree):
            tree_log(val, f"{root}/{i}", data)

    elif isinstance(tree, dict):
        for key, value in tree.items():
            tree_log(value, f"{root}/{key}", data, iter=iter)  # noqa: F722

    elif hasattr(tree, "to_compound"):
        tree_log(tree.to_dict(), f"{root}/iter", data)
        tree_log(tree.to_compound()[1], root, data)  # noqa: F722

    elif hasattr(tree, "to_dict"):
        tree_log(tree.to_dict(), f"{root}/iter", data)
        tree_log(iter, root, data)  # noqa: F722

    else:
        if iter is not None:
            tree_log(iter, f"{root}/iter", data)
            root = f"{root}/value"
    
        value = np.asarray(tree)
        if root in data:
            f_value = data[root]
            f_value.resize(f_value.shape[0]+1, axis=0)
            f_value[-1] = value
        else:
            maxshape = (None, ) + value.shape
            data.create_dataset(root, data=[value], maxshape = maxshape)