`rayleigh.py` needs refactoring - mainly concerning the handling of Dask vs non-Dask input

[adybbroe]

Code Sample, a minimal, complete, and verifiable piece of code

The code conserning the atm correction in the visible spectrum ("Rayleigh" correction) needs refactoring. the get_reflectance method for instance, is rather complex and long. And the handling of Dask or pure Numpy inputs is difficult to maintain. See for instance #137

Also, the code below, that mimics the example on the RTD pages (which works - see second example here) but here with "real" data, fails with a rather odd error.

# Your code here
    scn = Scene(filenames=FILES[2], reader='viirs_sdr')
    scn.load(['M03', 'M04', 'M05', 'solar_zenith_angle', 'solar_azimuth_angle',
              'satellite_zenith_angle', 'satellite_azimuth_angle'])

    lcn = scn.resample(areaid)

    viirs = Rayleigh('Suomi-NPP', 'viirs')
    sunz = lcn['solar_zenith_angle']
    satz = lcn['satellite_zenith_angle']
    ssadiff = lcn['solar_azimuth_angle'] - lcn['satellite_azimuth_angle']

    nsunz = sunz.compute().data
    nsatz = satz.compute().data
    nssadiff = ssadiff.compute().data
    redband = lcn['M05'].compute().data

    refl_cor_m3 = viirs.get_reflectance(sunz, satz, ssadiff, 'M3', redband)

Problem description

The above code should work, returning an array of reflectances. The same call without the redband as input works fine.

In general it would be nice to have support for both Dask arrays, Numpy arrays and Xarrays (with Dask or not).

Expected Output

No error as below.

Actual Result, Traceback if applicable

KeyError                                  Traceback (most recent call last)
<ipython-input-4-546faf561721> in <module>
----> 1 refl_cor_m3 = viirs.get_reflectance(sunz, satz, ssadiff, 'M3', redband)

~/usr/src/forks/pyspectral/pyspectral/rayleigh.py in get_reflectance(self, sun_zenith, sat_zenith, azidiff, bandname, redband)
    254 
    255         # force dask arrays
--> 256         compute = False
    257         if HAVE_DASK and not isinstance(sun_zenith, Array):
    258             compute = True

~/miniconda3/envs/python38/lib/python3.8/site-packages/dask/base.py in compute(self, **kwargs)
    286         dask.base.compute
    287         """
--> 288         (result,) = compute(self, traverse=False, **kwargs)
    289         return result
    290 

~/miniconda3/envs/python38/lib/python3.8/site-packages/dask/base.py in compute(*args, **kwargs)
    562     )
    563 
--> 564     dsk = collections_to_dsk(collections, optimize_graph, **kwargs)
    565     keys, postcomputes = [], []
    566     for x in collections:

~/miniconda3/envs/python38/lib/python3.8/site-packages/dask/base.py in collections_to_dsk(collections, optimize_graph, optimizations, **kwargs)
    338         for opt, val in groups.items():
    339             dsk, keys = _extract_graph_and_keys(val)
--> 340             dsk = opt(dsk, keys, **kwargs)
    341 
    342             for opt_inner in optimizations:

~/miniconda3/envs/python38/lib/python3.8/site-packages/dask/array/optimization.py in optimize(dsk, keys, fuse_keys, fast_functions, inline_functions_fast_functions, rename_fused_keys, **kwargs)
     43         dsk = HighLevelGraph.from_collections(id(dsk), dsk, dependencies=())
     44 
---> 45     dsk = optimize_blockwise(dsk, keys=keys)
     46     dsk = fuse_roots(dsk, keys=keys)
     47     dsk = dsk.cull(set(keys))

~/miniconda3/envs/python38/lib/python3.8/site-packages/dask/blockwise.py in optimize_blockwise(graph, keys)
   1205     rewrite_blockwise
   1206     """
-> 1207     out = _optimize_blockwise(graph, keys=keys)
   1208     while out.dependencies != graph.dependencies:
   1209         graph = out

~/miniconda3/envs/python38/lib/python3.8/site-packages/dask/blockwise.py in _optimize_blockwise(full_graph, keys)
   1283 
   1284             # Merge these Blockwise layers into one
-> 1285             new_layer = rewrite_blockwise([layers[l] for l in blockwise_layers])
   1286             out[layer] = new_layer
   1287 

~/miniconda3/envs/python38/lib/python3.8/site-packages/dask/blockwise.py in rewrite_blockwise(inputs)
   1386             # Update new_axes
   1387             for k, v in inputs[dep].new_axes.items():
-> 1388                 new_axes[sub[k]] = v
   1389 
   1390             # Bump new inputs up in list

KeyError: 'j'

Versions of Python, package at hand and relevant dependencies

dask                      2021.11.1          pyhd8ed1ab_0    conda-forge
dask-core                 2021.11.1          pyhd8ed1ab_0    conda-forge
numpy                     1.20.3           py38h9894fe3_1    conda-forge
python                    3.8.12          hb7a2778_2_cpython    conda-forge

Thank you for reporting an issue !