Global IR Cloud Composite from Geostationary Satellites

cloud_composite.py

import datetime
from dateutil.rrule import rrule, MINUTELY
from os.path import join as pjoin
import numpy as np
import xarray as xr

import pyresample.geometry
import satpy
import satpy.dataset
import satpy.composites
import satpy.writers


LON_LIMITS = {'Meteosat-11': [-37.5, 20.75],
              'Meteosat-08 IODC': [20.75, 91.1],
              'HIMAWARI-08': [91.1, -177.15],
              'GOES-17': [-177.15, -105.],
              'GOES-16': [-105., -37.5],
}


def harmonize(scenes, common_id):
    """Harmonize scenes so that they are recognized by MultiScene."""
    harmonized_scenes = []
    for scene in scenes:
        new_scene = satpy.Scene()
        ds_copy = scene[scene.keys()[0]].copy()
        ds_copy.attrs['wavelength'] = common_id.wavelength
        new_scene[common_id] = ds_copy
        harmonized_scenes.append(new_scene)
    return harmonized_scenes


def add_lonlats(dataset):
    """Replace x/y by lat/lon coordinates in the given dataset."""
    lon, lat = dataset.area.get_lonlats()

    dataset = dataset.rename({'x': 'lon', 'y': 'lat'})
    try:
        dataset = dataset.drop(['x_image', 'y_image'])
    except ValueError:
        pass

    dataset['lat'] = lat[:, 0]
    dataset['lat'].attrs['standard_name'] = 'Latitude'
    dataset['lat'].attrs['units'] = 'degrees_north'

    dataset['lon'] = lon[0, :]
    dataset['lon'].attrs['standard_name'] = 'Longitude'
    dataset['lon'].attrs['units'] = 'degrees_east'
    
    return dataset


def remove_lonlats(dataset):
    """Replace lat/lon by x/y coordinates in the given dataset."""
    dataset = dataset.rename({'lon': 'x', 'lat': 'y'})
    dataset['x'] = np.arange(dataset.shape[1])
    dataset['x'].attrs = {}
    dataset['y'] = np.arange(dataset.shape[0])
    dataset['x'].attrs = {}
    return dataset


def stack(datasets):
    # Add lon/lat coordinates to facilitate data selection. Ideally you would
    # do that using a MultiIndex, but vectorized selection is not yet supported
    # for MultiIndex, see https://github.com/pydata/xarray/issues/1603
    datasets = [add_lonlats(ds) for ds in datasets]

    # Initialize target array with nans
    res = xr.DataArray(data=np.full(datasets[0].shape, np.nan),
                       dims=datasets[0].dims,
                       coords=datasets[0].coords)
    
    # Fill in data from the different platforms
    for dataset in datasets:
        # Cut longitude range
        left, right = LON_LIMITS[dataset.attrs['platform_name']]
        mask = xr.DataArray(data=np.zeros(dataset.shape),
                            dims=dataset.dims,
                            coords=dataset.coords)
        if right > left:
            mask.loc[dict(lon=slice(left, right))] = 1
        else:
            # Crosses the date line
            mask.loc[dict(lon=slice(-180, right))] = 1
            mask.loc[dict(lon=slice(left, 180))] = 1
        cut = dataset.where(mask)

        # Fill in data (keeping existing data)
        res = res.where(res.notnull(), cut)

    # Cut high latitudes
    mask = xr.DataArray(data=np.zeros(res.shape),
                        dims=res.dims,
                        coords=res.coords)
    mask.loc[dict(lat=slice(73, -73))] = 1
    res = res.where(mask)

    # Reset coordinates (compositors require x/y coordinates)
    res = remove_lonlats(res)

    # Update attributes
    tstart = min([ds.start_time for ds in datasets])
    tend = max([ds.end_time for ds in datasets])
    res.name = 'IR_110'
    res.attrs = {'area': datasets[0].area,
                 'standard_name': 'toa_brightness_temperature',
                 'long_name': '11um Brightness Temperature',
                 'units': 'K',
                 'start_time': tstart.isoformat(),
                 'end_time': tend.isoformat(),
                 'platforms': [ds.attrs['platform_name'] for ds in datasets]}

    return res


def create_cloud_composite(dataset, filename):
    # Setup blue marble background image (matches the global 0.1 deg grid in this case)
    bluemarble_comp = satpy.composites.StaticImageCompositor(
        'bluemarble', filename='/path/to/bluemarble_0.1_deg_framed.jpeg')
    bluemarble_comp.area = global_grid
    bluemarble = bluemarble_comp()

    # Hack: x/y coordinates from StaticImageCompositor are off by 0.5; align them
    bluemarble.coords['x'].data = dataset.coords['x'].data
    bluemarble.coords['y'].data = dataset.coords['y'].data

    # Create IR cloud composite: Cold -> opaque white, warm -> transparent
    cloud_comp = satpy.composites.CloudCompositor('clouds',
                                                  transition_min=258.15,
                                                  transition_max=298.15,
                                                  transition_gamma=3.0)
    clouds = cloud_comp([dataset])
    im_clouds = satpy.writers.to_image(clouds)
    im_clouds.invert([True, False])  # hack to facilitate usage in BackgroundCompositor
    clouds = im_clouds.data

    # Put clouds on top of blue marble
    bg_comp = satpy.composites.BackgroundCompositor('my_composite')
    composite = bg_comp([clouds, bluemarble])

    # Convert composite to image, apply enhancements and save to file
    img = satpy.writers.to_image(composite)
    img.stretch('linear')
    img.save(filename)


if __name__ == '__main__':
    output_dir = '.'
    
    # Load IR scenes. 
    # Note: Each dataset must provide the "platform_name" attribute
    scenes = [scene1, scene2, ...]

    # Wavelengths are slightly different and therefore not grouped by MultiScene. 
    # Harmonize scenes using a common dataset ID.
    common_id = satpy.dataset.DatasetID(name='IR_110', wavelength=(10, 11, 12))
    harmonized = harmonize(scenes, common_id)

    # Create multi-scene
    multi_scene = satpy.MultiScene(harmonized)

    # Resample all scenes to the same global grid
    global_grid = pyresample.geometry.AreaDefinition(...)
    resampled = multi_scene.resample(global_grid,
                                     resampler='nearest',
                                     radius_of_influence=60000.0,
                                     reduce_data=False,
                                     cache_dir=pjoin(output_dir, 'cache'))

    # Blend observations
    blended = resampled.blend(blend_function=stack)

    # Create cloud composite
    create_cloud_composite(blended[common_id], 'mycomposite.png')