Enable parallel image writes to S3

[Jan-Willem]

Currently, parallel image writes are done as follows:

• Before the creation of the graph in imaging/cube_imaging_niter0.py (runs in serial):
  • img_xds = make_empty_sky_image(...) #Creates a dataset with only coordinates and no data variables.
  • write_image(img_xds, imagename=image_name, out_format="zarr") #Write image to disk.
  • zarr_meta = create_data_variable_meta_data_on_disk(image_name, data_varaibles_and_dims_sel, xds_dims, parallel_coords, compressor) # Creates the JSON files (.zattrs, .zmetadata) that describe the data variables (sky, psf, etc).
• In the mapping stage of the graph imaging/_utils/_make_image.py:
  • for data_variable, meta in input_params["zarr_meta"].items(): write_chunk(img_xds,meta,parallel_dims_chunk_id,input_params["compressor"],input_params["image_file"]) #Writes the binary chunks to disk. For example 0.1.0.0.

This should be changed to make use of Zarr functionality so that parallel writes S3 are supported.

[amcnicho]

As of 6fcad36a42ca8acf16f512c408e8e4e750c6e7fa there is a new file tests/domain/imaging/test_cube_imaging_niter0.py that does basic integrity checks of image writes to disk. The branch also contains the changes I found necessary in order to get the new tests to pass with the most recent version of the Antennae dataset (i.e., for cube_imaging_niter0 not to raise exceptions when trying to write an image following the mosaic notebook).

I'll use these tests as the basis for switching write_chunk to use zarr.open, so @Jan-Willem it would help if you ran them to confirm the output still looks like you'd expect after the recent schema changes.