3

I am asking a kind of generalisation of this question:

Best way to extract elements from nested lists.

It is also somehow related to these questions:

recursive function for extract elements from deep nested lists/tuples

Scheme - find most deeply values nested lists

Scheme - find most deeply nested lists

Essentially, I have some arbitrary nested list structure, where at the bottom there are various numpy arrays that are all the same shape. I want to iterate or slice all these bottom-level arrays whilst preserving the nested list structure in which they live. It is this part about preserving the nested structure in the output which doesn't seem to be answered in these other questions.

So, for example:

A = np.ones((3,3,3))
nest = [[A,A,A],[[A,A,[A,A]],[A,A]],A]

and we want, schematically,

nest[0,...] == [[A[0,...],A[0,...],A[0,...]],[[A[0,...],A[0,...],[A[0,...],A[0,...]]],[A[0,...],A[0,...]]],A[0,...]]

or

nest[1:3,5,:] == [[A[1:3,5,:],A[1:3,5,:],A[1:3,5,:]],[[A[1:3,5,:],A[1:3,5,:],[A[1:3,5,:],A[1:3,5,:]]],[A[1:3,5,:],A[1:3,5,:]]],A[1:3,5,:]]

I'm sure some clever recursive function or something can do this, my brain is just not coming up with it right now...

I guess it would also be best if this returns views onto the bottom level arrays rather than copies parts of them.

edit: Perhaps something like this can work: https://stackoverflow.com/a/43357135/1447953. That method would require that numpy slicing operations be converted into functions somehow, which I guess can be done on a case-by-case basis, so perhaps this is the way to go.

Ben Farmer
  • 2,387
  • 1
  • 25
  • 44
  • I see 2 tasks - a function, probably recursive, to work your way down the nest, and a function to act of arrays. Develop these in small steps in an interactive session. – hpaulj Mar 21 '18 at 17:16

3 Answers3

2

maybe a generator like:

def get_nested_elements(elements):
    if not elements or isinstance(elements[0], np.number):
        yield elements
    else:
        for node in elements:
            for e in get_nested_elements(node):
                yield e

will return the ndarray if the first element is of type number.

Totoro
  • 867
  • 9
  • 10
1

A quick first attempt:

def recurse(f, alist):
    def foo(f, item):
        if isinstance(item,list):
            return recurse(f, item)
        else:
            return f(item)
    return [foo(f, item) for item in alist]

test case:

In [1]: A = np.arange(10)
In [2]: alist = [[A,A],[[A],[A,A]],A]

In [6]: recurse(lambda A: A[3], alist)
Out[6]: [[3, 3], [[3], [3, 3]], 3]

In [7]: recurse(lambda A: A[:3], alist)
Out[7]: 
[[array([0, 1, 2]), array([0, 1, 2])],
 [[array([0, 1, 2])], [array([0, 1, 2]), array([0, 1, 2])]],
 array([0, 1, 2])]

not limited to indexing:

In [10]: recurse(lambda A: A.sum(), alist)
Out[10]: [[45, 45], [[45], [45, 45]], 45]
hpaulj
  • 221,503
  • 14
  • 230
  • 353
0

Appreciate the help! I came up with the following voodoo to do what I wanted, partly based on https://stackoverflow.com/a/43357135/1447953. I also generalised it to allow binary (and higher) operations with mirrored nested structures. It's basically a generalisation of 'map' to apply over nested structures. Could use better checking of the structure as it goes, for example it doesn't make sure that the lengths match, though I guess 'map' will throw an error if this is the case.

def apply_f(f,*iters):
    """Apply some function to matching 'bottom level' objects 
       in mirrored nested structure of lists,
       return the result in the same nested list structure.
       'iters' should be lists which have identical
       nested structure, and whose bottom level elements can
       be used as arguments to 'f'.
    """
    # We have to descend all list structures in lock-step!
    if all(isinstance(item,list) for item in iters):
        return list(map(lambda *items: apply_f(f,*items), *iters))
    elif any(isinstance(item,list) for item in iters):
        raise ValueError("Inconsistency in nested list structure of arguments detected! Nested structures must be identical in order to apply functions over them")
    else:
        return f(*iters)

It can be used like this:

A = np.ones((2,2))

a = [A,[A,A],A,[[A,A],A]]

def my_sum(a,b,c):
    return a + b + c

b = apply_f(my_sum,a,a,a)

print(a)
print(b)

# Check structure
print(apply_f(lambda A: A.shape, a))
print(apply_f(lambda A: A.shape, b))

Output:

[array([[1., 1.],
       [1., 1.]]), [array([[1., 1.],
       [1., 1.]]), array([[1., 1.],
       [1., 1.]])], array([[1., 1.],
       [1., 1.]]), [[array([[1., 1.],
       [1., 1.]]), array([[1., 1.],
       [1., 1.]])], array([[1., 1.],
       [1., 1.]])]]
[array([[3., 3.],
       [3., 3.]]), [array([[3., 3.],
       [3., 3.]]), array([[3., 3.],
       [3., 3.]])], array([[3., 3.],
       [3., 3.]]), [[array([[3., 3.],
       [3., 3.]]), array([[3., 3.],
       [3., 3.]])], array([[3., 3.],
       [3., 3.]])]]
[(2, 2), [(2, 2), (2, 2)], (2, 2), [[(2, 2), (2, 2)], (2, 2)]]
[(2, 2), [(2, 2), (2, 2)], (2, 2), [[(2, 2), (2, 2)], (2, 2)]]

Of course then the applied function can do whatever one likes to the bottom level arrays, whether it be extract elements, slice them, add them together, whatever.

Ben Farmer
  • 2,387
  • 1
  • 25
  • 44