load_uniform_grid with nprocs>1 and cell_widths

[chrishavlin]

Bug report

Calling load_uniform_grid with nprocs>1 and cell_widths does not work.

Code for reproduction

As of now, the following illustrates appears to work:

import numpy as np
import yt

N = 16
data = {"density": np.random.random((N, N, N))}

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

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,
        nprocs=4,
    )

yt : [INFO     ] 2023-02-09 11:20:23,001 Parameters: current_time              = 0.0
yt : [INFO     ] 2023-02-09 11:20:23,002 Parameters: domain_dimensions         = [16 16 16]
yt : [INFO     ] 2023-02-09 11:20:23,002 Parameters: domain_left_edge          = [0. 0. 0.]
yt : [INFO     ] 2023-02-09 11:20:23,003 Parameters: domain_right_edge         = [1. 1. 1.]
yt : [INFO     ] 2023-02-09 11:20:23,003 Parameters: cosmological_simulation   = 0

but do anything with that dataset:

dens = ds.all_data()[("stream", "density")]

and you get

Traceback (most recent call last):
  File "/home/chris/miniconda3/envs/yt_dev_39/lib/python3.9/site-packages/IPython/core/interactiveshell.py", line 3398, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "<ipython-input-3-e6b5073a1feb>", line 1, in <cell line: 1>
    dens = ds.all_data()[("stream", "density")]
  File "/home/chris/src/yt_/yt_dev/yt/yt/data_objects/static_output.py", line 1194, in all_data
    self.index
  File "/home/chris/src/yt_/yt_dev/yt/yt/data_objects/static_output.py", line 604, in index
    self._instantiated_index = self._index_class(
  File "/home/chris/src/yt_/yt_dev/yt/yt/frontends/stream/data_structures.py", line 167, in __init__
    GridIndex.__init__(self, ds, dataset_type)
  File "/home/chris/src/yt_/yt_dev/yt/yt/geometry/geometry_handler.py", line 39, in __init__
    self._setup_geometry()
  File "/home/chris/src/yt_/yt_dev/yt/yt/geometry/grid_geometry_handler.py", line 41, in _setup_geometry
    self._parse_index()
  File "/home/chris/src/yt_/yt_dev/yt/yt/frontends/stream/data_structures.py", line 209, in _parse_index
    self.grids.append(self.grid(id, self))
  File "/home/chris/src/yt_/yt_dev/yt/yt/frontends/stream/data_structures.py", line 87, in __init__
    cell_widths = index.grid_cell_widths[id - self._id_offset]
IndexError: list index out of range

NOTE: in #4328 I added a check for when nprocs>1 being used with cell_widths in order to raise an error earlier during the initial call to load_uniform_grid. So when that gets merged,

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,
        nprocs=4,
    )

will result in the following (or some similar variation):

NotImplementedError                       Traceback (most recent call last)
<ipython-input-1-bbf1cc62ea5c> in <module>
     11     cell_widths.append(cw)
     12 
---> 13 ds = yt.load_uniform_grid(
     14         data,
     15         data["density"].shape,

~/src/yt_/yt_dev/yt/yt/loaders.py in load_uniform_grid(data, domain_dimensions, length_unit, bbox, nprocs, sim_time, mass_unit, time_unit, velocity_unit, magnetic_unit, periodicity, geometry, unit_system, default_species_fields, axis_order, cell_widths, parameters)
    354         # cell_widths left as an empty guard value if None
    355         if nprocs != 1:
--> 356             raise NotImplementedError("nprocs must equal 1 if supplying cell_widths.")
    357         cell_widths = _validate_cell_widths(cell_widths, domain_dimensions)
    358 

NotImplementedError: nprocs must equal 1 if supplying cell_widths.

Expected outcome

A stretched grid decomponsed by nprocs.

Conceptually, I think the issue is just that cell_widths needs to be decomposed similar to how the data fields are already decomposed. Splitting a single stretched grid into multiple grid objects should be OK as long as they're all level 0, so I think it comes down to refactoring load_uniform_grid to also decompose the cell_widths.