def pdfs_aligned_nonzero()

in src/smclarify/util/__init__.py [0:0]

• 19 lines of code
• 8 McCabe index (conditional complexity)

def pdfs_aligned_nonzero(*args) -> List[np.ndarray]:
    """
    Convert a list of discrete pdfs / freq counts to aligned numpy arrays of the same size for common non-zero elements
    :return: pair of numpy arrays of the same size with the aligned pdfs
    """
    num_pdfs = len(args)
    pdfs = []
    for x in args:
        pdfs.append(pdf(x))

    def keys(_xs):
        return seq(_xs).map(lambda x: x[0])

    # Extract union of keys
    all_keys = seq(pdfs).flat_map(keys).distinct().sorted()

    # Index all pdfs by value
    dict_pdfs = seq(pdfs).map(dict).list()

    # result aligned lists
    aligned_lists: List[List] = [[] for x in range(num_pdfs)]

    # fill keys present in all pdfs
    for i, key in enumerate(all_keys):
        for j, d in enumerate(dict_pdfs):
            if d.get(key, 0) == 0:
                break
        else:
            # All keys exist and are != 0
            for j, d in enumerate(dict_pdfs):
                aligned_lists[j].append(d[key])
    np_arrays = seq(aligned_lists).map(np.array).list()
    return np_arrays