Suggested change to std/var preprocessing to improve precision

[jemmajeffree]

Hi,
I've noticed that in a few rare situations, groupby and flox can return quite noisy standard deviations. In situations where the mean of an array is much larger than the standard deviation (such as deep ocean salinity, raised here), flox returns noisier values on dask arrays than on loaded numpy arrays. In extreme situations, the standard deviation of a dask array can contain NaNs from square-rooting negative variances.

I'm guessing it's the same idea as #386, in which case @dcherian has thought about this for much longer than I have. I've done a little bit of looking through the code, and could easily have missed something about how this works with neighbouring functions, but my thoughts on the potential problem and how it might be addressed are below.

Minimal complete verifiable example:

import numpy as np
import xarray as xr

l =12000
np.random.seed(1)
test_data = xr.DataArray(np.random.uniform(0,1,l)/100+1000000,dims=('time',)  # huge mean with relatively small variability
                        ).assign_coords({'month':xr.DataArray(np.arange(l)%12,dims=('time',))})

# with numpy arrays returns reasonable and consistent values
test_data.groupby('month').std('time')
# array([0.00283648, 0.00281895, 0.00287791, 0.00287652, 0.00287337,
#        0.00287037, 0.00289802, 0.00289441, 0.00285839, 0.00296478,
#        0.00284787, 0.00292089])

# using lazy computation/dask
dask_test_data = test_data.chunk({'time':100})
dask_test_data.groupby('month').std('time').load()
# array([0.01118034, 0.01118034, 0.01118034, 0.01581139, 0.        ,
#        0.01581139, 0.01118034, 0.01118034, 0.01118034,        nan,
#               nan, 0.        ])

A functional workaround is to subtract the mean before calculating standard deviation:

(dask_test_data.groupby('month')-dask_test_data.groupby('month').mean('time')).groupby('month').std('time').load()

My understanding is that the distinction comes from aggregate_npg.py improving precision by subtracting the first non-nan element of the array, a preprocessing step skipped by aggregations.py. This solution is probably not quite as stable as subtracting the mean, but the first element should be really close to the mean if the standard deviation is small, and it might be faster.

I’d suggest that to improve precision and match the numpy engine behaviour in aggregations_npg.py, the flox engine implementation for dask arrays of nanstd,nanvar,std,var could have a preprocessor that looks something like this:

def var_std_preprocess(array, axis): # Not sure of naming conventions, sorry
    """Subtracts first value of array from whole array, 
    to improve numerical precision of nanstd, nanvar, std, var

    Adapted from from argreduce_preprocess and _var_std_wrapper in aggregate_npg.py
    """
    import dask.array  # Copied from argreduce_preprocess, but maybe these shouldn’t be within the function? 
    import numpy as np # For either this function or argreduce_preprocess?
    
    # NEXT LINE IS PSEUDOCODE; I’m not entirely sure how to apply it lazily
    # If it doesn’t cost anything speed wise, then probably better to use mean. Happy to run some time tests on either
    first_elements = nanfirst(array,axis) 

    def subtract_first(array_, first_elements_):
        return array_-first_elements_

    return dask.array.map_blocks(
        subtract_first,
        array,
        first,
        dtype=array.dtype,
        meta=array._meta,
        name="groupby-var_std-preprocess",
    )

and is included in the Aggregations definition like so:

nanstd = Aggregation(
    "nanstd",
    preprocess=var_std_preprocess, #UPDATED LINE
    chunk=("nansum_of_squares", "nansum", "nanlen"),
    combine=("sum", "sum", "sum"),
    finalize=_std_finalize,
    fill_value=0,
    final_fill_value=np.nan,
    dtypes=(None, None, np.intp),
    final_dtype=np.floating,
)

It seems to work if first_elements is naively array[0] in the one-dimensional, no-nans case, but I’m not sure how to generalise it and apply nanfirst without the usual layers/wrappers. (aggregate_npg.py uses first = _get_aggregate(engine).aggregate(group_idx, array, func="nanfirst", axis=axis), but I don't think this syntax translates to the flox/dask implementation) . If you can give me a few tips or examples to work from, then I’m happy to try implement this behaviour.

Happy also to discuss alternatives, or to provide a pull request if that's easier to work with.

This is also my first time reading through flox code in detail (it's really nicely written and documented, by the way, was lovely to read), and one of the first times I’ve interacted with public github repos, so I’d appreciate any feedback or corrections on what's useful to provide when describing issues.

[dcherian]

First, thanks for the really well thought out post and deep-dive in to the code. That's not easy to do. I really appreciate it. (I also hope you're having fun, which it seems like you are :) )

My first thought is that nanfirst is one full pass through the data which we should try hard to avoid (if we can). To that end, the idea in #386 is to use Welford's method (though the second post in that thread links to a paper that does something different). I would first understand what intermediates we need to track during the compute and then modify

flox/flox/aggregations.py

Lines 349 to 358 in ca57681

	var = Aggregation(
	"var",
	chunk=("sum_of_squares", "sum", "nanlen"),
	combine=("sum", "sum", "sum"),
	finalize=_var_finalize,
	fill_value=0,
	final_fill_value=np.nan,
	dtypes=(None, None, np.intp),
	final_dtype=np.floating,
	)

to match. Is this something you are up for exploring?

For reference, Cubed implements the Welford algo, I think: https://github.com/cubed-dev/cubed/blob/745f564964933d6178666c62e71d92fbc0b3fd2b/cubed/array_api/statistical_functions.py#L209. This would be a useful reference

[jemmajeffree]

Thanks @dcherian, I'm definitely having fun.

I see your point about avoiding a first full pass through the data. I guess I was hoping that there was a way to pick any datapoint that's not a NaN in O(1) time, but I can see how that would be difficult to implement

I'm not sure how well Welford's method (at least, that described by the blog post you linked) would adapt to parallelization, because calculations for each additional point rely on all those that have come before. However, the variant described in the Schubert and Gertz (2018) paper you linked would work for combining variance calculated on individual chunks. Cubed seems to be doing something along these lines, merging variance from subsets of data with more than one datapoint.

How much is touching data within a chunk twice a problem? I'm wondering if applying the simpler two-pass solution (subtract mean, then sum squares) within each chunk, and then combining these as described in Schubert and Gertz (2018) increases precision with minimal loss of speed for the typical xarray use case. I'll run some tests for speed, but I think this is what the "proposed minibatch" in Schubert and Gertz is describing.

I'm keen to explore :) It's looking to me like implementing anything along these lines would require defining something new that can be applied blockwise to dask chunks — would you agree or have I missed something here?

[dcherian]

Thanks @dcherian, I'm definitely having fun.

🥳 !

because calculations for each additional point rely on all those that have come before.

In general these parallel algorithms are quite similar, apply something blockwise ("chunk"), combine intermediates ("combine"), and potentially add a final transformation ("aggregate"). The phrase to look for is "incremental update".

For example, the snippet at the end of https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Welford's_online_algorithm updates the estimate for the new value, and scrolling down "Chan's algorithm" has formulae for combining two intermediate estimates.

It's looking to me like implementing anything along these lines would require defining something new that can be applied blockwise to dask chunks

We already have this in the chunk parameter to Aggregation. For simplicity in many cases, it is simply a string, but it can be a function. Or at least it should accept any function though I do see that there are no examples of that sort.

The principle behind welford's algorithm seems to be to track the mean and variance instead of sum & sum-of-squares which we do currently. I haven't read Schubert and Gertz but I'd encourage you to experiment. If you open a PR i'm happy to help with comments or tiny code changes to unblock you.

[jemmajeffree]

Just noting that I’m travelling for the next month, so while I’m still thinking about how to implement this functionality I’m unlikely to make progress in the immediate future

[jemmajeffree]

Coming back to this, and having read the Schubert and Gertz paper in detail, I like the algorithm and think it would work here. Below is an explanation of how I understand the algorithm and would consider implementing it, with some preliminary code. This isb little more than pseudocode — it does run, but I'm also passing variables around as lists, and other poor habits, just to check the basic idea before implementing properly.

My understanding of how such an implementation would work is as follows:

$S_A$ - weighted sum of squared deviations in chunk A (sum_deviations)
$N_A$ - sum of weights in chunk A (sum_w; would just be number of elements if unweighted)

so, variance can be computed as $S_{all}$/ $N_{all}$

in order to combine the chunks, we also need to keep track of one other intermediate number:
$V_A$ - weighted sum of values in chunk A (sum_x)

Equation 22 from Schubert and Gertz can then be simplified (for covariance, rather than variance) to allow combining of results from chunks:

$$ S_{AB} = S_A+S_B+\frac{(N_AV_B-N_BV_A)^2}{N_AN_B(N_A+N_B)}$$

In the code below, I'm referring to the fraction from the line above as the "adjustment term" (adj_terms)

The below diagram summarised the flow of information through my suggested implementation, where the compute_intermediates function acts on each chunk (A-D) and the combine intermediates function combines the outputs from that first function, in a cascading add-each-chunk-on-the-end manner

def compute_intermediates(x,w=None): #this is the computation
    ''' 
    function that operates on a chunk
    take x, which is a chunk
    return:
        Sum of deviations (ie, VW_P from paper)
        Sum of X*weights
        Sum of weights
        '''
    
    # I've left the weights in for the rest of the code, but I suspect they're not
    # going to stay so I haven't tested them, just using ones (ie, equal weighting)
    if w is None:
        w = np.ones(x.shape)
         
    # Calculate intermediates
    sum_x = np.sum(x*w)
    sum_w = np.sum(w)
    
    mean_x = sum_X/sum_w
    sum_deviations = np.sum(w*(x-mean_x)**2)
    
    # All three of these should be floats (ie, scalars)
    return sum_deviations, sum_X, sum_w


def combine_intermediates(intermediates_list): #combine intermediates
    ''' 
    function that combines results from chunks to get end result
    take output from get_block_pieces,
    in sets of sum_deviations, sum_X, sum_w, 
    and combine these to return the same for the set'''
    
    # Crude implementation of unwrapping the way intermediates are passed around. 
    # Basically a nested list, while there's definitely better options I'm sure 
    # flox has defaults anyway so there's not much point me optimising
    sum_deviations = []
    sum_X = []
    sum_w = []

    for item in intermediates_list:
        sum_deviations.append(item[0])
        sum_X.append(item[1])
        sum_w.append(item[2])

    # Calculate parts needed for cascading combination
    cumsum_X = np.cumsum(sum_X)
    cumsum_w = np.cumsum(sum_w)
    
    # cumsum_w or _X are the terms from the left (ie, lots of previous chunks) and sum_w or _X terms
    # the vertical/from right arrows, (ie, single chunk being combined)
    adj_terms = (cumsum_w[:-1]*sum_X[1:]-sum_w[1:]*cumsum_X[:-1])**2/(cumsum_w[:-1]*sum_w[1:]*(cumsum_w[:-1]+sum_w[1:]))

    return np.sum(sum_deviations)+np.sum(adj_terms)

I've tested that this implementation works, and doesn't suffer the same catastrophic cancellation as the $E(X^2)-E(X)^2$ method. I haven't bothered timing it, because I'm still passing intermediates around as lists which I suspect would take more time than everything else combined

# Build a giant dataset that we expect to maybe cancel catestrophically
N = 1000
np.random.seed(1)
test_data2 = np.random.normal(0,1,N)
test_data2 = test_data2+100000000

# numpy standard deviation
np.var(test_data2)
# 0.962369110543228

# Two-pass method (accurate, slow)
# (I think this is what numpy does under the hood)
np.mean((test_data2-np.mean(test_data2))**2)
# 0.9623691105432282

# Naive implementation, less memory touching
np.mean(test_data2**2)-np.mean(test_data2)**2
# 2.0 (low precision; wrong)

# This implementation, in chunks
intermediates = []
for i in range(0,N,N//10):
    intermediates.append(compute_intermediates(test_data2[i:i+N//10]))
combine_intermediates(intermediates)/N
# 0.9623691103776513

I've got a couple questions:

• Do you think the above implementation would work, or otherwise what changes would you suggest?
• In particular, I suspect it'd be slightly more stable if the intermediates were combined pairwise, as a tree, but not enough for it to be worthwhile implementing the more complicated combination
• Does the implementation need to work for weighted variance, or covariance, or just for unweighted variance (and standard deviation)?
• Can I use numpy functions in the combine function (or is it too far down the stack)?

[dcherian]

👏 👏 Very nice.

What you call "cascading combine" is called a "fold". flox instead uses a tree reduction.

graph TD
    A1[a]
    A2[b]
    A3[c]
    A4[d]

    B1[a + b]
    B2[c + d]
    C[a+b + c+d]


    A1 --> B1
    A2 --> B1

    A3 --> B2
    A4 --> B2

    B1 --> C
    B2 --> C

Loading

In particular, I suspect it'd be slightly more stable if the intermediates were combined pairwise, as a tree, but not enough for it to be worthwhile implementing the more complicated combination

flox already does this for you :) More below.

Does the implementation need to work for weighted variance, or covariance, or just for unweighted variance (and standard deviation)?

flox doesn't support weights yet

Can I use numpy functions in the combine function (or is it too far down the stack)?

Yes, but you'll have to alter your thinking a bit.

Adapting to flox

What you've done is quite good, but needs to adapt to how flox is designed (which sadly was done keeping much simpler things in mind). The complexity here is that to update your estimate your combine function needs access to multiple intermediates.

flox doesn't let you do that (yet).

Current implementation

If you look at the current implementation for var

flox/flox/aggregations.py

Lines 349 to 358 in ca57681

	var = Aggregation(
	"var",
	chunk=("sum_of_squares", "sum", "nanlen"),
	combine=("sum", "sum", "sum"),
	finalize=_var_finalize,
	fill_value=0,
	final_fill_value=np.nan,
	dtypes=(None, None, np.intp),
	final_dtype=np.floating,
	)

it calculates three intermediates, and accumulates them separately; and then _var_finalize combines them to create the result.

The sequential accumulation of intermediates is in this loop:

flox/flox/core.py

Lines 1451 to 1452 in 619a390

	for idx, combine in enumerate(agg.simple_combine):
	array = _conc2(x_chunk, key1="intermediates", key2=idx, axis=axis_)

Approaches

So to adapt this algorithm, we have two options:
(1) There is a single chunk function that creates numpy record arrays with 3 intermediates, the combine and finalize functions know how to process those structured arrays. The memory layout is quite inefficient though :( .

(2) We could create some kind of custom array thing whose implementation would look very similar to https://github.com/pydata/xarray/blob/59e98f1efddcd866a4ba43634fe61c4893504cb4/xarray/tests/arrays.py#L129:

class MultiArray:
    arrays: tuple[np.ndarray, ...]

   def __array_function__(self, ...):
        ...

   # These are the only two "array" method we need to work AFAICT
   def concatenate(self, ...): 
        # concatenate all the arrays in `arrays`     
        ...

   def expand_dims(self, ...):
       # expand_dims for all the arrays in `arrays`     
       ...

expand_dims and concatenate are the only two operations flox needs to perform to set up the combination of intermediates.

Proposal

So my proposal

1. Create a _var_chunk(array: np.ndarray) -> MultiArray function that takes numpy array as input and returns the intermediates for each block
2. Create a _var_combine(intermediates: MultiArray, axis: int) -> MultiArray that runs the combination along axis and returns a new MultiArray
3. Create a _var_finalize(array: MultiArray) -> np.ndarray that does the final aggregation.

Are you up for this?

[jemmajeffree]

Thanks @dcherian for the very detailed explanations and suggestions

That makes sense, and I'm keen to give it a go.

The best example to work off, in terms of code structure, would probably be the existing var aggregation? Is there anything in particular you'd suggest keeping in mind when writing to handle multiple dimensions?

Out of curiosity, why does flox combine as a tree? I guess it'd be more versatile for huge datasets, but wouldn't the overhead for calling so many functions which just do one operation (i.e. sum two numbers) be substantial?

[dcherian]

I'd suggest taking your nice example and modifying it so:

stride = 4
state = intermediates_list[0]
for items in itertools.batched(intermediates_list[1:], n=stride):
    # implement this by concatenating `state` and `*items` along a new "dummy axis" = -2.
    # and then reducing that to one MultiArray
    state = combine_intermediates(state, *items)

This will be similar to what happens.

---

The reduction tree exposes more parallelism, the fold is inherently sequential. So if you have associative & commutative reductions, the tree is better (IIUC).

Though note that you can adapt your version to run well in parallel sincecumsum can be executed in parallel (section 39.2.4) [still more expensive than a simple reduction tree though].

[dcherian]

In looking for a better answer to your fold vs reduce question I found https://dl.acm.org/doi/10.1145/1596550.1596551 (see video below)

A short summary from here:

Since the (left or right) fold commits us to sequential evaluation, Guy Steele called it `slightly harmful'. He urged using reduce, as far as possible, because it is flexible and decouples the algorithm from the execution strategy. The strategy (sequential, parallel, distributed) and data partitioning can be chosen later, depending on circumstances and available resources.

[dcherian]

Hi @jemmajeffree is there anything I can do to help here?

[jemmajeffree]

Hi Deepak, I'm slowly moving forwards and getting a feel for the flox code. Thank you for your support and detailed answers to my questions.

The current sticking point is working out how to turn my pseudocode into a function that can handle dimensions and groups, and the inputs that individual functions should expect (which I can sort-of figure out by throwing different data at .groupby().var() and chasing it through with a debugger, but that's not super efficient)

I was starting by trying to implement something that sums the sum-of-squared-deviations from different chunks, and ignoring the adjustment terms for the moment. I tried (in a separate git branch for messing around) adding the following into aggregations.py, but I think I'm missing some sort of generic wrapper that handles the grouping.

def sum_squared_deviations(group_idx, array,**kwargs):
    out = (array-np.mean(array))**2 #probably shouldn't be using np mean here, but some kind of partial?
    return out

nanvar = Aggregation(
    "nanvar",
    chunk=(sum_squared_deviations,),
    combine=("sum",),
    #finalize=_var_finalize,
    fill_value=0,
    final_fill_value=np.nan,
    dtypes=(None, ),
    final_dtype=np.floating,
)

It also looks like sum_squared_deviations would actually need to be in aggregate_flox.py (and the other aggregate_*.py files), but I can't work out where the lookup table is that turns strings like "sum_of_squares" into the sum_of_squares function in aggregate_flox.py.

While I'm here, when is the -2 axis created, and why is it -2?
And am I correct in my interpretation that flox combines four intermediates per tree/combine step?

[dcherian]

#probably shouldn't be using np mean here, but some kind of partial?

You'll need the grouped mean in aggregate_flox so something like

means = mean(group_idx, array, ...)
deviations = array - means[..., group_idx]

It also looks like sum_squared_deviations would actually need to be in aggregate_flox.py (and the other aggregate_*.py files),

yes. and you'll need to propagate 3 intermediates back instead of just 1 which is what all the current functions do.

but I can't work out where the lookup table is that turns strings like "sum_of_squares" into the sum_of_squares function in aggregate_flox.py.

It's pretty yucky, see here:

flox/flox/aggregations.py

Lines 83 to 109 in 619a390

	elif engine == "flox":
	try:
	method = getattr(aggregate_flox, func)
	except AttributeError:
	# logger.debug(f"Couldn't find {func} for engine='flox'. Falling back to numpy")
	method = get_npg_aggregation(func, engine="numpy")
	elif engine == "numbagg":
	from . import aggregate_numbagg
	try:
	if "var" in func or "std" in func:
	ddof = kwargs.get("ddof", 0)
	if aggregate_numbagg.NUMBAGG_SUPPORTS_DDOF or (ddof != 0):
	method = getattr(aggregate_numbagg, func)
	else:
	logger.debug(f"numbagg too old for ddof={ddof}. Falling back to numpy")
	method = get_npg_aggregation(func, engine="numpy")
	else:
	method = getattr(aggregate_numbagg, func)
	except AttributeError:
	# logger.debug(f"Couldn't find {func} for engine='numbagg'. Falling back to numpy")
	method = get_npg_aggregation(func, engine="numpy")
	elif engine in ["numpy", "numba"]:
	method = get_npg_aggregation(func, engine=engine)

While I'm here, when is the -2 axis created, and why is it -2?

flox/flox/core.py

Lines 1874 to 1876 in 619a390

	if do_simple_combine:
	# Add a dummy dimension that then gets reduced over
	blockwise_method = tlz.compose(_expand_dims, blockwise_method)

-1 is the "group" axis i.e. shape[-1] == len(groups) so, I stuck the dummy axis right next to it.

And am I correct in my interpretation that flox combines four intermediates per tree/combine step?

yes this is the default for dask'ssplit_every

[jemmajeffree]

Okay, so I created a prototype MultiArray Class:

class MultiArray:
    arrays: tuple[np.ndarray, ...]

    def __init__(self, arrays):
        self.arrays = arrays # something else probably needed here to be more careful about types (not sure what)
        
    def astype(self,dt,**kwargs):
        for array in self.arrays: # Do we care about trying to avoid for loops here? three separate lines would be faster, but harder to read
            array = array.astype(dt,**kwargs)# Seems kind of awkward, like something that wouldn't work in C and while it works in python it's icky
        return self
    
    def reshape(self, *args, **kwargs):
        for array in self.arrays:
            array = array.reshape(*args,**kwargs)
        return self
    
    def expand_dims(self, *args, **kwargs):
        ## expand_dims for all the arrays in `arrays`     
        for array in self.arrays:
            array = array.expand_dims(*args,**kwargs)
        return self       

    #def concatenate(self, ...): #Hadn't been needed yet; I'll add it in when it is
        ## concatenate all the arrays in `arrays`  

(astype and reshape needed by here and here respectively)

And a chunk function in aggregate_flox.py that returns a MultiArray:

def var_chunk(group_idx, array, *, axis=-1, size=None, fill_value=None, dtype=None):
    # Calculate length and sum - important for the adjustment terms to sum squared deviations
    array_lens = nanlen(
        group_idx,
        array,
        axis=axis,
        size=size,
        fill_value=fill_value,
        dtype=dtype,
    )
    
    array_sums = sum(
        group_idx,
        array,
        axis=axis,
        size=size,
        fill_value=fill_value,
        dtype=dtype,
    )
    
    # Calculate sum squared deviations - the main part of variance sum
    array_means = array_sums/array_lens # Does this risk being run eagerly because it's not wrapped in anything?
    
    sum_squared_deviations = sum(
        group_idx,
        (array-array_means[..., group_idx])**2,
        axis=axis,
        size=size,
        fill_value=fill_value,
        dtype=dtype,
    ) 

    return MultiArray((sum_squared_deviations,array_sums,array_lens))

I'm now hitting an error on this line, I assume because np.expand_dims doesn't like the MultiArray

def _expand_dims(results: IntermediateDict) -> IntermediateDict:
    results["intermediates"] = tuple(
        np.expand_dims(array, axis=DUMMY_AXIS) for array in results["intermediates"] # <— PROBLEM LINE
    )
    return results

full error

File "/Users/u6955431/Documents/*PhD/code_environment/flox_std/debugging flox.py", line 21, in <module>
  print(std.load())
File "/Users/u6955431/mambaforge/envs/flox-tests/lib/python3.13/site-packages/xarray/core/dataarray.py", line 1175, in load
  ds = self._to_temp_dataset().load(**kwargs)
File "/Users/u6955431/mambaforge/envs/flox-tests/lib/python3.13/site-packages/xarray/core/dataset.py", line 541, in load
  evaluated_data: tuple[np.ndarray[Any, Any], ...] = chunkmanager.compute(
File "/Users/u6955431/mambaforge/envs/flox-tests/lib/python3.13/site-packages/xarray/namedarray/daskmanager.py", line 85, in compute
  return compute(*data, **kwargs) # type: ignore[no-untyped-call, no-any-return]
File "/Users/u6955431/mambaforge/envs/flox-tests/lib/python3.13/site-packages/dask/base.py", line 681, in compute
  results = schedule(expr, keys, **kwargs)
File "/Users/u6955431/mambaforge/envs/flox-tests/lib/python3.13/site-packages/toolz/functoolz.py", line 489, in __call__
  ret = f(ret)
File "/Users/u6955431/Documents/*PhD/code_environment/flox_std/flox/flox/core.py", line 1393, in _expand_dims
  results["intermediates"] = tuple(
File "/Users/u6955431/Documents/*PhD/code_environment/flox_std/flox/flox/core.py", line 1394, in <genexpr>
  np.expand_dims(array, axis=DUMMY_AXIS) for array in results["intermediates"]
File "/Users/u6955431/mambaforge/envs/flox-tests/lib/python3.13/site-packages/numpy/lib/_shape_base_impl.py", line 593, in expand_dims
  axis = normalize_axis_tuple(axis, out_ndim)
File "/Users/u6955431/mambaforge/envs/flox-tests/lib/python3.13/site-packages/numpy/_core/numeric.py", line 1435, in normalize_axis_tuple
  axis = tuple([normalize_axis_index(ax, ndim, argname) for ax in axis])

numpy.exceptions.AxisError: axis -2 is out of bounds for array of dimension 1

I considered replacing the line with array.expand_dims(axis=DUMMY_AXIS) for array in results["intermediates"], but it turns out np.ndarray.expand_dims isn't a function.

I also tried looking a step down into np.expand_dims, and it initially seemed that adding a ndim attribute to the MultiArray might help, but numpy is wrapping the whole thing up as an array of a single object before it tries to expand dims.

Any suggestions?

[dcherian]

Yes, great progress! 💪

Now we need to teach numpy to defer to MultipleArrays for expand_dims. Let's use NEP-18 or __array_function__ protocol. See here for an example. Notice you'll need to do the same for concatenate.

[jemmajeffree]

Thanks for the pointer. I think I've gotten MultiArray to handle the numpy functions (code below). I've a couple more questions about types, and generally about trying to write good code

• As it currently stands, MultiArray.__init__ doesn't do anything to coerce non-tuples into tuples. On it's own, this is potentially problematic, but it's definitely problematic when combined with the fact that I've used list comprehension to implement several numpy functions. What would you recommend changing here?
• For that matter, is there a significant time cost associated with using list comprehension? I'm used to avoiding for loops in favour of vectorisation, but I can's see a smart way to avoid them here
• for shape, ndims, dtype, I've assumed that these properties are the same across all arrays in the MultiArray. Is this a reasonable assumption? Would you check it, and if so where?
• I've put some naive checking into concatenate, to make sure all the MultiArrays at least have the same number of arrays. Can you take a look, and let me know if it should be replaced by some other construct
• should MultiArray be defined in its own python script so it can be used by the different aggregate_*.py scripts?

MULTIARRAY_HANDLED_FUNCTIONS = {}

class MultiArray:
    arrays: tuple[np.ndarray, ...]

    def __init__(self, arrays):
        self.arrays = arrays # something else needed here to be more careful about types (not sure what)
        # Do we want to co-erce arrays into a tuple and make sure it's immutable? Do we want it to be immutable?
        
    def astype(self,dt,**kwargs):
        return MultiArray([array.astype(dt,**kwargs) for array in self.arrays])  # Is this the tidiest construct to do this kind of thing?
    
    def reshape(self,shape,**kwargs):
        return MultiArray([array.reshape(shape,**kwargs) for array in self.arrays])
    
    def __array_function__(self, func, types, args, kwargs):
        if func not in MULTIARRAY_HANDLED_FUNCTIONS:
            return NotImplemented
        return MULTIARRAY_HANDLED_FUNCTIONS[func](*args, **kwargs)    

        
   # Shape is needed, seems likely that the other two might be
   # Making some strong assumptions here that all the arrays are the same shape, and I don't really like this
    @property
    def dtype(self) -> np.dtype:
        return self.arrays[0].dtype

    @property
    def shape(self) -> tuple[int, ...]:
        return self.arrays[0].shape

    @property
    def ndim(self) -> int:
        return self.arrays[0].ndim

def implements(numpy_function):
    """Register an __array_function__ implementation for MultiArray objects."""
    def decorator(func):
        MULTIARRAY_HANDLED_FUNCTIONS[numpy_function] = func
        return func
    return decorator

@implements(np.expand_dims)
def expand_dims_MultiArray(multiarray,axis):
    return MultiArray([np.expand_dims(a,axis) for a in multiarray.arrays]) #This is gonna spit out a list and I'm not sure if I'm okay with that

@implements(np.concatenate)
def concatenate_MultiArray(multiarrays,axis):
    
    n_arrays = len(multiarrays[0].arrays)
    for ma in multiarrays[1:]:
        if not (len(ma.arrays) == n_arrays): # I don't know what trying to concatenate MultiArrays with different numbers of arrays would even mean
            raise NotImplementedError
      
    # There's the potential for problematic different shapes coming in here.  
    # Probably warrants some defensive programming, but I'm not sure what to check for while still being generic
    
    return MultiArray([np.concatenate(arrays,axis) for arrays in zip([ma.arrays for ma in multiarrays])]) # This is nasty, sorry, probably could be written nicer. Very open to suggestions

None of these questions are blockers — I should have the information I need to implement the combine/chunk function — and I know it's SciPy next week, so no rush to reply.