Skip to content
Snippets Groups Projects
plot_mask.py 11.7 KiB
Newer Older
import sys

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.image as mpimg

#import cartopy.crs as ccrs

from matplotlib import colors
from netCDF4 import Dataset
from numpy import less, greater

import argparse

_map_names = ['bit0', 'cloudy', 'cloud_mask_value', 'day', 'sunglint', 'snow_ice',
         'land', 'thin_cirrus_solar', 'snow_covered_from_anc', 'thin_cirrus_IR',
         'cloudy_plus_1_pixel_adj', 'ocean_BT_threshold', 'high_cloud_flag_138',
         'high_cloud_flag_39_12', 'high_cloud_flag_11_12',
         'cloud_flag_39_11_BTD_test', 'cloud_flag_VisNIR_test',
         'cloud_flag_VisNIR_ratio', 'shallow_water_clear_sky_restoral',
         'cloud_flag_16_21_test', 'high_cloud_flag_86_11_BTD_test',
         'ocean_clear_sky_restoral_test', 'land_sunglint_polar_clearsky_restoral',
         'cloud_flag_sfc_temp_test', 'cloud_flag_11um_var_test',
         'low_cloud_flag_39_11BTD_test', 'simple_IR_test',
         'high_cloud_CO2_test', 'high_cloud_067um_test',
         'cloud_vis_reflectance_test', 'cloud_near_IR_reflectance_test',
         'vza_mask_00', 'vza_mask_10', 'vza_mask_30', 'vza_mask_50', 'sct_mask_70',
         'sct_mask_90', 'sct_mask_110', 'sct_mask_130', 'sct_mask_150',
         'sct_mask_170', 'sct_mask_180']


_map_names = ['cloud_mask_value']

def cmap_discretize(cmap, N, N2=1024):
    """Return a discrete colormap from the continuous colormap cmap.

        cmap: colormap instance, eg. cm.jet.
        N: number of colors.

    Example
        x = resize(arange(100), (5,100))
        djet = cmap_discretize(cm.jet, 5)
        imshow(x, cmap=djet)
    """

    if type(cmap) == str:
        cmap = cm.get_cmap(cmap)

    colors_i = np.concatenate((np.linspace(0, 1., N), (0., 0., 0., 0.,)))
    colors_rgba = cmap(colors_i)
    indices = np.linspace(0, 1., N+1)
    cdict = {}
    for ki, key in enumerate(('red', 'green', 'blue')):
        cdict[key] = [(indices[i], colors_rgba[i-1, ki], colors_rgba[i, ki])
                      for i in range(N+1)]

    return colors.LinearSegmentedColormap(cmap.name + '_%d' % N, cdict, N2)


def main(fname, test_name='',
         c1='mediumblue', c2='skyblue', c3='yellowgreen', c4='olivedrab'):

    #data = read_data(fname, 18, 21, 102, 107)
    data = read_data(fname, -90, 90, -180, 180)
    masks = create_masks(data)

    for test_name in _map_names:
        print(test_name)
        x = masks[test_name]

        zz = np.zeros((x.shape[0], x.shape[1], 4), dtype=float)

        ind = np.where(x == 0)
        zz[ind[0], ind[1], :] = np.array(colors.to_rgba(c1))

        ind = np.where(x == 1)
        zz[ind[0], ind[1], :] = np.array(colors.to_rgba(c2))

        ind = np.where(x == 2)
        zz[ind[0], ind[1], :] = np.array(colors.to_rgba(c3))

        ind = np.where(x == 3)
        zz[ind[0], ind[1], :] = np.array(colors.to_rgba(c4))

        mpimg.imsave(f'{test_name}.png', np.uint8(zz*255))


def read_data(filename, lat1, lat2, lon1, lon2):
    f = Dataset(filename)
    lat = f['geolocation_data/latitude'][:]
    lon = f['geolocation_data/longitude'][:]
    ind = np.nonzero(np.bitwise_and(
                     np.bitwise_and(lat>lat1, lat<lat2),
                     np.bitwise_and(lon>lon1, lon<lon2) ))
    x1, x2, y1, y2 = np.min(ind[0]), np.max(ind[0]), np.min(ind[1]), np.max(ind[1])
    qa = f['geophysical_data/Quality_Assurance'][x1:x2, y1:y2, :]
    clearsky_conf = f['geophysical_data/Clear_Sky_Confidence'][x1:x2, y1:y2]
    cm = f['geophysical_data/Cloud_Mask'][:, x1:x2, y1:y2]
    vza = f['geolocation_data/sensor_zenith'][x1:x2, y1:y2]
    sza = f['geolocation_data/solar_zenith'][x1:x2, y1:y2]
    senaz = f['geolocation_data/sensor_azimuth'][x1:x2, y1:y2]
    solaz = f['geolocation_data/solar_azimuth'][x1:x2, y1:y2]

    f.close()

    data = {'QA': qa, 'clear_sky_confidence': clearsky_conf,
            'cloud_mask': cm, 'latitude': lat, 'longitude': lon,
            'sensor_zenith': vza, 'solar_zenith': sza,
            'sensor_azimuth': senaz, 'solar_azimuth': solaz}

    return data


def create_masks(data):
    # 2
    dtr = 3.14159 / 180.0
    rtd = 180.0 / 3.14159
    qa = data['QA']
    cm = data['cloud_mask']
    vza = data['sensor_zenith']
    sza = data['solar_zenith']
    senaz = data['sensor_azimuth']
    solaz = data['solar_azimuth']

    array_mask = cm[0, :, :].mask

    bit0 = read_bits(cm[0, :, :], 0)
    # clear_sky_confidence = 0

    # 3
    cloudy = read_bits(cm[0, :, :], [1, 2])
#   cloudy = np.array((np.array(cm[0, :, :], np.byte) &
#                      (pow(2, 1) + pow(2, 2))) / pow(2, 1))
    cm_flag = np.array((np.array(cm[0, :, :], np.byte) &
                       (pow(2, 1) + pow(2, 2)))/pow(2, 1))

    # 4
    day = read_bits(cm[0, :, :], 3)
    sunglint = ~read_bits(cm[0, :, :], 4)
    snow_ice = ~read_bits(cm[0, :, :], 5)
#    land = read_bits(cm[0, :, :], [6, 7])
    only_land = np.array((np.array(cm[0, :, :], np.byte) &
                         (pow(2, 6) + pow(2, 7)))/pow(2, 6))
    only_land = np.array(only_land, dtype='bool') + 0

    # only CM additional info
    simpleIRtest = read_bits(cm[1, :, :], 5);
    hicld_CO2 = read_bits(cm[1, :, :], 6);
    hicld_67um = read_bits(cm[1, :, :], 7);
    cld_visrefl = read_bits(cm[2, :, :], 4);
    cld_nearIR = read_bits(cm[2, :, :], 6);
    # 5
    thin_cirrus_solar = read_bits(qa[:, :, 1], 1) & ~read_bits(cm[1, :, :], 1)
    snow_cover = read_bits(qa[:, :, 1], 2) & ~read_bits(cm[1, :, :], 2)
    thin_cirrus_IR = read_bits(qa[:, :, 1], 3) & ~read_bits(cm[1, :, :], 3)
    cloudy_plus1 = read_bits(qa[:, :, 1], 4) & ~read_bits(cm[1, :, :], 4)
    oceanBTthresh = read_bits(qa[:, :, 1], 5) & ~read_bits(cm[1, :, :], 5)
    hicloud_1_38 = read_bits(qa[:, :, 2], 0) & ~read_bits(cm[2, :, :], 0)
    hicloud_39_12 = read_bits(qa[:, :, 2], 1) & ~read_bits(cm[2, :, :], 1)
    hicloud_11_12 = read_bits(qa[:, :, 2], 2) & ~read_bits(cm[2, :, :], 2)
    cloud_39_11BTD = read_bits(qa[:, :, 2], 3) & ~read_bits(cm[2, :, :], 3)
    cloud_VisNIR = read_bits(qa[:, :, 2], 4) & ~read_bits(cm[2, :, :], 4)
    cloud_NIRVis_ratio = read_bits(qa[:, :, 2], 5) & ~read_bits(cm[2, :, :], 5)
    shallow_clear = read_bits(qa[:, :, 2], 6) & read_bits(cm[2, :, :], 6)
    cloud_16_21 = read_bits(qa[:, :, 2], 7) & ~read_bits(cm[2, :, :], 7)
    hicloud_86_11BTD = read_bits(qa[:, :, 3], 0) & ~read_bits(cm[3, :, :], 0)
    ocean_clear = read_bits(qa[:, :, 3], 1) & read_bits(cm[3, :, :], 1)
    land_sunglint_polar = read_bits(qa[:, :, 3], 2) & read_bits(cm[3, :, :], 2)
    cloud_sfctemp = read_bits(qa[:, :, 3], 3) & ~read_bits(cm[3, :, :], 3)
    cloud_11umvar = read_bits(qa[:, :, 3], 6) & ~read_bits(cm[3, :, :], 6)
    lowcloud_39_11BTD = read_bits(qa[:, :, 3], 7) & ~read_bits(cm[3, :, :], 7)
    vzamask00 = (vza[:, :] <= 10).astype(int)
    vzamask10 = ((vza[:, :] > 10) & (vza[:, :] <= 30)).astype(int)
    vzamask30 = ((vza[:, :] > 30) & (vza[:, :] <= 50)).astype(int)
    vzamask50 = (vza[:, :] > 50).astype(int)

    raz = np.abs(180.0 - np.abs(senaz[:, :] - solaz[:, :]))
    cos_sctang = -1.0 * (np.cos(sza[:, :]*dtr) *np.cos(vza[:, :]*dtr) -
                         np.sin(sza[:, :]*dtr) *np.sin(vza[:, :]*dtr) *np.cos(raz[:, :]*dtr) )
    sct = ( np.arccos(cos_sctang[:, :]) * rtd);

    sctmask70 = (sct[:, :] <= 70).astype(int)
    sctmask90 = ((sct[:, :] > 70) & (sct[:, :] <= 90)).astype(int)
    sctmask110 = ((sct[:, :] > 90) & (sct[:, :] <= 110)).astype(int)
    sctmask130 = ((sct[:, :] > 110) & (sct[:, :] <= 130)).astype(int)
    sctmask150 = ((sct[:, :] > 130) & (sct[:, :] <= 150)).astype(int)
    sctmask170 = ((sct[:, :] > 150) & (sct[:, :] <= 170)).astype(int)
    sctmask180 = (sct[:, :] > 170).astype(int)

    masks = {'bit0': bit0,
             'cloudy': cloudy,
             'cloud_mask_value': cm_flag,
             'day': day,
             'sunglint': sunglint,
             'snow_ice': snow_ice,
             'land': only_land,
             'thin_cirrus_solar': thin_cirrus_solar,
             'snow_covered_from_anc': snow_cover,
             'thin_cirrus_IR': thin_cirrus_IR,
             'cloudy_plus_1_pixel_adj': cloudy_plus1,
             'ocean_BT_threshold': oceanBTthresh,
             'high_cloud_flag_138': hicloud_1_38,
             'high_cloud_flag_39_12': hicloud_39_12,
             'high_cloud_flag_11_12': hicloud_11_12,
             'cloud_flag_39_11_BTD_test': cloud_39_11BTD,
             'cloud_flag_VisNIR_test': cloud_VisNIR,
             'cloud_flag_VisNIR_ratio': cloud_NIRVis_ratio,
             'shallow_water_clear_sky_restoral': shallow_clear,
             'cloud_flag_16_21_test': cloud_16_21,
             'high_cloud_flag_86_11_BTD_test': hicloud_86_11BTD,
             'ocean_clear_sky_restoral_test': ocean_clear,
             'land_sunglint_polar_clearsky_restoral': land_sunglint_polar,
             'cloud_flag_sfc_temp_test': cloud_sfctemp,
             'cloud_flag_11um_var_test': cloud_11umvar,
             'low_cloud_flag_39_11BTD_test': lowcloud_39_11BTD,
             'simple_IR_test': simpleIRtest,
             'high_cloud_CO2_test': hicld_CO2,
             'high_cloud_067um_test': hicld_67um,
             'cloud_vis_reflectance_test': cld_visrefl,
             'cloud_near_IR_reflectance_test': cld_nearIR,
             'vza_mask_00': vzamask00,
             'vza_mask_10': vzamask10,
             'vza_mask_30': vzamask30,
             'vza_mask_50': vzamask50,
             'sct_mask_70' : sctmask70,
             'sct_mask_90' : sctmask90,
             'sct_mask_110' : sctmask110,
             'sct_mask_130' : sctmask130,
             'sct_mask_150' : sctmask150,
             'sct_mask_170' : sctmask170,
             'sct_mask_180' : sctmask180}

    for k in masks.keys():
        masks[k] = np.ma.array(masks[k], mask=array_mask)

    return masks


def read_bits(byte, bits):
    # orig_shape = byte.shape
    # if byte.ndim > 1:
    #     byte = byte.reshape((byte.shape[0]*byte.shape[1],))

    if type(bits) is int:
        flag = np.array((np.array(byte, np.byte) & pow(2, bits))/pow(2, bits),
                        dtype='bool')
    else:
        flag = (np.array(byte, np.byte) & pow(2, bits[0]) +
                         pow(2, bits[1]))/pow(2, bits[0])

        flag = np.array(np.trunc(1.0-flag/2.0), dtype='bool')
    # flag = flag.reshape((orig_shape))
    return flag


parser = argparse.ArgumentParser()

parser.add_argument('input')
parser.add_argument('test_name')
parser.add_argument('-c1', default='mediumblue')
parser.add_argument('-c2', default='skyblue')
parser.add_argument('-c3', default='yellowgreen')
parser.add_argument('-c4', default='olivedrab')

args = parser.parse_args()


if __name__ == "__main__":

    map_names = ['bit0', 'cloudy', 'cloud_mask_value', 'day', 'sunglint', 'snow_ice',
             'land', 'thin_cirrus_solar', 'snow_covered_from_anc', 'thin_cirrus_IR',
             'cloudy_plus_1_pixel_adj', 'ocean_BT_threshold', 'high_cloud_flag_138',
             'high_cloud_flag_39_12', 'high_cloud_flag_11_12',
             'cloud_flag_39_11_BTD_test', 'cloud_flag_VisNIR_test',
             'cloud_flag_VisNIR_ratio', 'shallow_water_clear_sky_restoral',
             'cloud_flag_16_21_test', 'high_cloud_flag_86_11_BTD_test',
             'ocean_clear_sky_restoral_test', 'land_sunglint_polar_clearsky_restoral',
             'cloud_flag_sfc_temp_test', 'cloud_flag_11um_var_test',
             'low_cloud_flag_39_11BTD_test', 'simple_IR_test',
             'high_cloud_CO2_test', 'high_cloud_067um_test',
             'cloud_vis_reflectance_test', 'cloud_near_IR_reflectance_test',
             'vza_mask_00, vza_mask_10, vza_mask_30, vza_mask_50, sct_mask_70',
             'sct_mask_90', 'sct_mask_110', 'sct_mask_130', 'vza_mask_150',
             'sct_mask170', 'sct_mask_180']

#    for n in map_names:
#        print(n)
    main(args.input,
#             n,