嵌套object类型ndarray中一维数组的遍历与函数应用方案问询

Problem Context #

I work with nested NumPy array structures for time-series data:

• type1: 1D float64 array (raw time-series signal)
• type2: Object-dtype array containing multiple type1 elements
• type3: Object-dtype array containing multiple type2 elements

Higher-order types are created by splitting signals into smaller segments. My goal is to apply a function (like a filter) to every type1 vector, with these requirements:

• Handle type1, type2, or type3 inputs seamlessly
• Support paired iteration over two arrays of matching dimensions
• Allow passing auxiliary data and extracting multi-output results

I wanted a function with this signature:

out = apply([main_array, aux_array], function_to_apply, **kwargs)

where main_array gets modified through the function application. I tried np.nditer without success, so I needed an alternative approach.

Implemented Recursive Solution #

With guidance from @JeromeRichards and @hpaulj, I built this recursive solution that meets all requirements:

def zip_none(data, aux):
    '''Zip routine with None broadcasting for auxiliary data'''
    return zip(data, *aux) if aux is not None else zip(data)

def apply_1d(data, func, aux=None, **kwargs):
    '''Apply a function to every 1D vector in nested data structures.
    
    Parameters:
        data: List or NumPy array (can be nested to type2/type3 levels)
        aux: List of auxiliary arrays, each matching the dimensionality of `data`
        func: Function to apply to each 1D `type1` vector
        kwargs: Additional keyword arguments to pass to `func`
    '''
    if isinstance(data, list):
        result = []
        multi_output = False
        for x, *y in zip_none(data, aux):
            current_aux = y if y else None
            tmp = apply_1d(x, func, aux=current_aux, **kwargs)
            result.append(tmp)
            
            # Track if the function returns multiple outputs
            if isinstance(tmp, tuple):
                multi_output = True
        
        # If multi-output, restructure results into separate tuples per output
        if multi_output:
            result = tuple(list(out) for out in zip(*result))
    else:
        # Base case: apply function to the 1D vector
        return func(data, *aux, **kwargs) if aux is not None else func(data, **kwargs)
    
    return result

Key Features #

• Recursive Traversal: Handles any depth of nested arrays (type1/type2/type3) by recursively unpacking until it reaches 1D vectors.
• Auxiliary Data Support: Uses zip_none to pair main and auxiliary data elements at each nested level, ensuring correct alignment.
• Multi-Output Handling: Automatically restructures results if the applied function returns multiple values (e.g., filtered signal + metadata).

Question sourced from Stack Exchange, asked by Piotr Herbut