Dynamic, heterogeneous dictionaries with tuple keys and non-constant tuple slicing

Numba Support: How do I do ...?
5

I feel I have made some progress by using the TupleKeyDict StructRef and the drop_elements function defined by @stuartarchibald in the related post. My code now looks like

import numpy as np
from scipy import spatial

from numba.typed import Dict, List
from numba import types


def f(X, simplices):
    """
    Parameters
    ===========
    X : N x D array
        Array of N Euclidean vectors in D dimensions

    simplices : n x (D + 1) array
        np.int32 (or np.int64??) array of indices, such as returned by
        scipy.spatial.Delaunay
    """
    D = X.shape[1]  # Top dimension
    filtration = {D: TupleKeyDict(Dict.empty(types.intp, types.float64)),
                  D - 1: TupleKeyDict(Dict.empty(types.intp, types.float64))}

    # Special iteration for highest dimensional simplices, but it's almost the
    # same as one of the iterations of the for loop in dim, below (for dim = D)
    for sigma in simplices:
        filtration[D][sigma] = np.sum(np.abs(X[np.asarray(sigma)]))
        for x in drop_elements(sigma):
            tau = x[1]
            if np.random.random() > 0.5:
                filtration[D - 1][tau] = filtration[D][sigma]
            else:
                filtration[D - 1][tau] = np.nan

    for dim in range(D - 1, 0, -1):
        filtration[dim - 1] = TupleKeyDict(Dict.empty(types.intp, types.float64))
        for sigma in filtration[dim]:
            if np.isnan(filtration[dim][sigma]):
                filtration[dim][sigma] = np.sum(np.abs(X[np.asarray(sigma)]))
            for x in drop_elements(sigma):
                tau = x[1]
                if np.random.random() > 0.5:
                    filtration[dim - 1][tau] = filtration[dim][sigma]
                else:
                    filtration[dim - 1][tau] = np.nan

    return filtration


X = np.random.random((100, 2))

simplices = List([tuple(simplex) for simplex in np.sort(spatial.Delaunay(X).simplices, axis=1)])

f(X, simplices)

Currently, this does not compile because looping over the keys of a TupleKeyDict and the contains function are not supported. Not only do I need to loop over the keys, I need to loop over them as tuples and not hashed integers.

Another question to @stuartarchibald: is the idea used in the overloaded setitem method of TupleKeyDict safe to collisions? It is not obvious to me that one could not pass two different tuples which hash to the same integer and replace a key-value pair instead of adding a new one.