Reduced the complexity of accum

master
Per A Brodtkorb 8 years ago
parent 110d3df086
commit ea2591d89d

@ -168,20 +168,25 @@ def accum(accmap, a, func=None, shape=None, fill_value=0, dtype=None):
return vals return vals
# Check for bad arguments and handle the defaults. def _initialize(accmap, a, func, shape, dtype):
if accmap.shape[:a.ndim] != a.shape: """Check for bad arguments and handle the defaults."""
raise ValueError(
"The initial dimensions of accmap must be the same as a.shape") _assert(accmap.shape[:a.ndim] == a.shape,
if func is None: "The initial dimensions of accmap must be the same as a.shape")
func = np.sum if accmap.shape == a.shape:
if dtype is None: accmap = np.expand_dims(accmap, -1)
dtype = a.dtype if func is None:
if accmap.shape == a.shape: func = np.sum
accmap = np.expand_dims(accmap, -1) if dtype is None:
adims = tuple(range(a.ndim)) dtype = a.dtype
if shape is None: adims = tuple(range(a.ndim))
shape = 1 + np.squeeze(np.apply_over_axes(np.max, accmap, axes=adims)) if shape is None:
shape = np.atleast_1d(shape) shape = 1 + np.squeeze(np.apply_over_axes(np.max, accmap,
axes=adims))
shape = np.atleast_1d(shape)
return accmap, shape, dtype, func
accmap, shape, dtype, func = _initialize(accmap, a, func, shape, dtype)
# Create an array of python lists of values. # Create an array of python lists of values.
vals = create_array_of_python_lists(accmap, a, shape) vals = create_array_of_python_lists(accmap, a, shape)

Loading…
Cancel
Save