Add xarray/dask support to DayNightCompositor

satpy/composites/__init__.py

@@ -770,60 +770,109 @@ class DayNightCompositor(GenericCompositor):
     """A compositor that takes one composite on the night side, another on day
     side, and then blends them together."""
 
-    def __call__(self, projectables, lim_low=85., lim_high=95., *args,
-                 **kwargs):
-        if len(projectables) != 3:
-            raise ValueError("Expected 3 datasets, got %d" %
-                             (len(projectables), ))
-        try:
-            day_data = projectables[0].copy()
-            night_data = projectables[1].copy()
-            coszen = np.cos(np.deg2rad(projectables[2]))
-
-            coszen -= min(np.cos(np.deg2rad(lim_high)),
-                          np.cos(np.deg2rad(lim_low)))
-            coszen /= np.abs(np.cos(np.deg2rad(lim_low)) -
-                             np.cos(np.deg2rad(lim_high)))
-            coszen = np.clip(coszen, 0, 1)
-
-            full_data = []
-
-            # Apply enhancements
-            day_data = enhance2dataset(day_data)
-            night_data = enhance2dataset(night_data)
-
-            # Match dimensions to the data with more channels
-            # There are only 1-channel and 3-channel composites
-            if day_data.shape[0] > night_data.shape[0]:
-                night_data = np.ma.repeat(night_data, 3, 0)
-            elif day_data.shape[0] < night_data.shape[0]:
-                day_data = np.ma.repeat(day_data, 3, 0)
-
-            for i in range(day_data.shape[0]):
-                day = day_data[i, :, :]
-                night = night_data[i, :, :]
-
-                data = (1 - coszen) * np.ma.masked_invalid(night).filled(0) + \
-                    coszen * np.ma.masked_invalid(day).filled(0)
-                data = np.ma.array(data, mask=np.logical_and(night.mask,
-                                                             day.mask),
-                                   copy=False)
-                data = Dataset(np.ma.masked_invalid(data),
-                               copy=True,
-                               **projectables[0].info)
-                full_data.append(data)
-
-            res = super(DayNightCompositor, self).__call__((full_data[0],
-                                                            full_data[1],
-                                                            full_data[2]),
-                                                           *args, **kwargs)
+    def __init__(self, lim_low=85., lim_high=95., **kwargs):
+        """Collect custom configuration values.
 
-        except ValueError:
-            raise IncompatibleAreas
+        Args:
+            lim_low (float): lower limit of Sun zenith angle for the
+                             blending of the given channels
+            lim_high (float): upper limit of Sun zenith angle for the
+                             blending of the given channels
+        """
+        self.lim_low = lim_low
+        self.lim_high = lim_high
+        super(DayNightCompositor, self).__init__(**kwargs)
+
+    def __call__(self, projectables, **kwargs):
+
+        day_data = projectables[0]
+        night_data = projectables[1]
+
+        lim_low = np.cos(np.deg2rad(self.lim_low))
+        lim_high = np.cos(np.deg2rad(self.lim_high))
+        try:
+            coszen = xu.cos(xu.deg2rad(projectables[2]))
+        except IndexError:
+            from pyorbital.astronomy import cos_zen
+            LOG.debug("Computing sun zenith angles.")
+            # Get chunking that matches the data
+            try:
+                chunks = day_data.sel(bands=day_data['bands'][0]).chunks
+            except KeyError:
+                chunks = day_data.chunks
+            lons, lats = day_data.attrs["area"].get_lonlats_dask(chunks)
+            coszen = xr.DataArray(cos_zen(day_data.attrs["start_time"],
+                                          lons, lats),
+                                  dims=['y', 'x'],
+                                  coords=[day_data['y'], day_data['x']])
+        # Calculate blending weights
+        coszen -= np.min((lim_high, lim_low))
+        coszen /= np.abs(lim_low - lim_high)
+        coszen = coszen.clip(0, 1)
+
+        # Apply enhancements to get images
+        day_data = enhance2dataset(day_data)
+        night_data = enhance2dataset(night_data)
+
+        # Adjust bands so that they match
+        # L/RGB -> RGB/RGB
+        # LA/RGB -> RGBA/RGBA
+        # RGB/RGBA -> RGBA/RGBA
+        day_data = add_bands(day_data, night_data['bands'])
+        night_data = add_bands(night_data, day_data['bands'])
+
+        # Get merged metadata
+        attrs = combine_metadata(day_data, night_data)
+
+        # Blend the two images together
+        data = (1 - coszen) * night_data + coszen * day_data
+        data.attrs = attrs
+
+        # Split to separate bands so the mode is correct
+        data = [data.sel(bands=b) for b in data['bands']]
+
+        res = super(DayNightCompositor, self).__call__(data, **kwargs)
 
         return res
 
 
+def enhance2dataset(dset):
+    """Apply enhancements to dataset *dset* and return the resulting data
+    array of the image."""
+    attrs = dset.attrs
+    img = get_enhanced_image(dset)
+    # Clip image data to interval [0.0, 1.0] and replace nan values
+    data = img.data.clip(0.0, 1.0).fillna(0.0)
+    data.attrs = attrs
+
+    return data
+
+
+def add_bands(data, bands):
+    """Add bands so that they match *bands*"""
+    # Add R, G and B bands, remove L band
+    if 'L' in data['bands'].data and 'R' in bands.data:
+        lum = data.sel(bands='L')
+        new_data = xr.concat((lum, lum, lum), dim='bands')
+        new_data['bands'] = ['R', 'G', 'B']
+        data = new_data
+    # Add alpha band
+    if 'A' not in data['bands'].data and 'A' in bands.data:
+        new_data = [data.sel(bands=band) for band in data['bands'].data]
+        # Create alpha band based on a copy of the first "real" band
+        alpha = new_data[0].copy()
+        alpha.data = da.ones((data.sizes['y'],
+                              data.sizes['x']),
+                             chunks=new_data[0].chunks)
+        # Rename band to indicate it's alpha
+        alpha['bands'] = 'A'
+        new_data.append(alpha)
+        new_data = xr.concat(new_data, dim='bands')
+        data = new_data
+
+    return data
+
+
 class Airmass(GenericCompositor):
 
     def __call__(self, projectables, *args, **kwargs):
@@ -978,24 +1027,6 @@ def __call__(self, projectables, **kwargs):
         return res
 
 
-def enhance2dataset(dset):
-    """Apply enhancements to dataset *dset* and convert the image data
-    back to Dataset object."""
-    img = get_enhanced_image(dset)
-
-    data = np.rollaxis(np.dstack(img.channels), axis=2)
-    mask = dset.mask
-    if mask.ndim < data.ndim:
-        mask = np.expand_dims(mask, 0)
-        mask = np.repeat(mask, 3, 0)
-    elif mask.ndim > data.ndim:
-        mask = mask[0, :, :]
-    data = Dataset(np.ma.masked_array(data, mask=mask),
-                   copy=False,
-                   **dset.info)
-    return data
-
-
 class RatioSharpenedRGB(GenericCompositor):
 
     def __init__(self, *args, **kwargs):
@@ -1079,6 +1110,7 @@ def __call__(self, datasets, optional_datasets=None, **info):
 
 
 class SelfSharpenedRGB(RatioSharpenedRGB):
+
     """Sharpen RGB with ratio of a band with a strided-version of itself.
 
     Example:

satpy/resample.py

@@ -483,6 +483,7 @@ def compute(self, data, fill_value=None, **kwargs):
 
 
 class NativeResampler(BaseResampler):
+
     """Expand or reduce input datasets to be the same shape.
 
     If `expand=True` (default) input datasets are replicated in both
@@ -492,6 +493,7 @@ class NativeResampler(BaseResampler):
     of the target area.
 
     """
+
     def resample(self, data, cache_dir=False, mask_area=False, **kwargs):
         # use 'mask_area' with a default of False. It wouldn't do anything.
         return super(NativeResampler, self).resample(data,