BUG: fix load_uniform_grid with cell_widths and multiple fields

[chrishavlin]

Turns out #4732 did not test out what happens when the data dict supplied to load_uniform_grid contains multiple fields... and it turns out that it errors. But it's a simple fix to indentation level.

[chrishavlin]

For reference, here's the test failure when running the updated test on main:

yt/frontends/stream/tests/test_stream_stretched.py:116: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
yt/loaders.py:366: in load_uniform_grid
    ) = decompose_array(
yt/utilities/decompose.py:24: in decompose_array
    return split_array(bbox[:, 0], bbox[:, 1], shape, psize, cell_widths=cell_widths)
yt/utilities/decompose.py:134: in split_array
    offset_re.append(offset_le[idim] + np.sum(cws[idim]))
../../../.pyenv/versions/3.10.11/envs/yt_dev/lib/python3.10/site-packages/numpy/_core/fromnumeric.py:2485: in sum
    return _wrapreduction(
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

obj = [array([0.10411111, 0.02827083, 0.12350717, 0.05914868, 0.1226682 ,
       0.11587512, 0.07593358, 0.03014527, 0.04251...0433]), array([0.09959554, 0.04023969, 0.00440936, 0.05259157, 0.02975688,
       0.10836904, 0.01412674, 0.06676059])]
ufunc = <ufunc 'add'>, method = 'sum', axis = None, dtype = None, out = None, kwargs = {'initial': <no value>, 'keepdims': <no value>, 'where': <no value>}
passkwargs = {}

    def _wrapreduction(obj, ufunc, method, axis, dtype, out, **kwargs):
        passkwargs = {k: v for k, v in kwargs.items()
                      if v is not np._NoValue}
    
        if type(obj) is not mu.ndarray:
            try:
                reduction = getattr(obj, method)
            except AttributeError:
                pass
            else:
                # This branch is needed for reductions like any which don't
                # support a dtype.
                if dtype is not None:
                    return reduction(axis=axis, dtype=dtype, out=out, **passkwargs)
                else:
                    return reduction(axis=axis, out=out, **passkwargs)
    
>       return ufunc.reduce(obj, axis, dtype, out, **passkwargs)
E       ValueError: setting an array element with a sequence. The requested array has an inhomogeneous shape after 1 dimensions. The detected shape was (3,) + inhomogeneous part.

../../../.pyenv/versions/3.10.11/envs/yt_dev/lib/python3.10/site-packages/numpy/_core/fromnumeric.py:86: ValueError
====================================================================== short test summary info ======================================================================
FAILED yt/frontends/stream/tests/test_stream_stretched.py::test_cell_width_with_nproc - ValueError: setting an array element with a sequence. The requested array has an inhomogeneous shape after 1 dimensions. The detected shape was (3,) + inhomogen...

[neutrinoceros]

I don't think I understand what the problem is or why this fixes it.

[chrishavlin]

oops, sorry, should have included more details. the fix was obvious to me, but only after like 45 mins of debugging lol.

The following loop when nprocs > 1: (1) takes the continuous data arrays and splits them into separate grids and (2) records the grid edges, grid dimensions, and cell widths (if applicable) for each grid:

    if nprocs > 1:
        temp = {}
        new_data = {}  # type: ignore [var-annotated]
        for key in data.keys():
            psize = get_psize(np.array(data[key].shape), nprocs)
            (
                grid_left_edges,
                grid_right_edges,
                shapes,
                slices,
                grid_cell_widths,
            ) = decompose_array(
                data[key].shape,
                psize,
                bbox,
                cell_widths=cell_widths,
            )
            grid_dimensions = np.array(list(shapes), dtype="int32")
            temp[key] = [data[key][slice] for slice in slices]
            # setting cell_widths = grid_cell_widths would overwrite for next loop = bad
        cell_widths = grid_cell_widths

Only one copy of the grid related variables (grid_dimensions, grid_left_edges, grid_right_edges, cell_widths) is needed for the subsequent call to load_amr_grids below. When there are multiple fields to split, those grid related variables just get re-created every loop and only the last one is kept. The final grid_cell_widths needs to be renamed to cell_widths to match the behavior down in the rest of the function, but that rename has to happen outside the loop over fields or the call to decompose_array would use a cell_widths that was already split by grid.

[neutrinoceros]

Oh, I see it now. I somehow missed that cell_widths would be fed back to decompose_array. Thanks a lot !

[chrishavlin]

In addition to the extra details in the comment, here's a small reproducer modified from the updated tests:

import yt 
import numpy as np 

def data_cell_widths_N16():
    np.random.seed(0x4D3D3D3)
    N = 16
    data = {
        "density": np.random.random((N, N, N)),
        "temperature": np.random.random((N, N, N)),  # adding a second field fails on main
    }

    cell_widths = []
    for _ in range(3):
        cw = np.random.random(N)
        cw /= cw.sum()
        cell_widths.append(cw)
    return (data, cell_widths)

data, cell_widths = data_cell_widths_N16()
cell_widths = [cw.astype(np.float32) for cw in cell_widths]
ds = yt.load_uniform_grid(
    data,
    data["density"].shape,
    bbox=np.array([[0.0, 1.0], [0.0, 1.0], [0.0, 1.0]]),
    cell_widths=cell_widths,
)