Support `flag_masks`

[dcherian]

From http://cfconventions.org/Data/cf-conventions/cf-conventions-1.8/cf-conventions.html#flags

int basin(lat, lon);
       standard_name: region;
       flag_values: 1, 2, 3;
       flag_meanings:"atlantic_arctic_ocean indo_pacific_ocean global_ocean";
data:
   basin: 1, 1, 1, 1, 2, ..... ;

It'd be nice to provide .cf.__eq__ and allow

ds.where(basin.cf == "atlantic_arctic_ocean")  # becomes ds.where(basin == 1)

or ds.where(basin.cf.flag == "atlantic_arctic_ocean"))?

We could do something similar for .isin too

I don't understand flag_masks but we could probably do something similar

  byte sensor_status_qc(time, depth, lat, lon) ;
    sensor_status_qc:long_name = "Sensor Status" ;
    sensor_status_qc:_FillValue = 0b ;
    sensor_status_qc:valid_range = 1b, 63b ;
    sensor_status_qc:flag_masks = 1b, 2b, 4b, 8b, 16b, 32b ;
    sensor_status_qc:flag_meanings = "low_battery processor_fault
                                      memory_fault disk_fault
                                      software_fault
                                      maintenance_required" ;

[jbusecke]

Am I right, that this does not work yet? I am about to work over at jbusecke/xMIP#119 for a bit, and just wanted to check.

[dcherian]

No but you can implement it!

How are you planning to use it? in .where?

[jbusecke]

Yes. This is a prototype of the code I am writing right now:

from cmip6_preprocessing.regionmask import merged_mask, _default_merge_dict

def basin_mask(ds):
    basins = regionmask.defined_regions.natural_earth.ocean_basins_50
    basins = merged_mask(basins, ds)
    
    # add cf compatible metadata
    basins.name = 'basin'
    flag_meanings = [m.replace(' ','_').lower() for m in list(_default_merge_dict().keys())]
    flag_values = np.arange(len(flag_meanings)).astype(str)
    basins.attrs.update({'flag_values':','.join(flag_values), 'flag_meanings':' '.join(flag_meanings)})
    return basins

def mask_basin(ds, basin, drop=False):
    mask = basin_mask(ds)
    
    # this should be done with cf_xarray once https://github.com/xarray-contrib/cf-xarray/issues/199 is implemented.
    # but for now some hardcoded hack...
    flag_meanings = mask.attrs['flag_meanings'].split(' ')
    flag_values = mask.attrs['flag_values'].split(',')
    flag_zipped = zip(flag_meanings, flag_values)
    print(flag_zipped)
    if basin not in flag_meanings:
        raise ValueError(f'Did not find basin [{basin}] in given values [{flag_meanings}]')
    else:
        flag_matched = [f for f in flag_zipped if f[0] == basin][0]
    
    print(flag_matched)
    ds_masked = ds.where(mask==int(flag_matched[1]), drop=drop)
    ds_masked.attrs['basin_masked'] = flag_matched[0]
    
    return ds_masked

This works very nicely in my work context (with any CMIP6 model):

ds = data_dict['CMIP.MIROC.MIROC6.historical.Omon.gn']
# basin_test = basin_mask(ds)
ds_masked = mask_basin(ds, 'southern_ocean')
plt.figure()
ds.thetao.isel(time=0, lev=0, member_id=0).plot()
plt.figure()
ds_masked.thetao.isel(time=0, lev=0, member_id=0).plot()

I would like to replace these lines

    flag_meanings = mask.attrs['flag_meanings'].split(' ')
    flag_values = mask.attrs['flag_values'].split(',')
    flag_zipped = zip(flag_meanings, flag_values)
    print(flag_zipped)
    if basin not in flag_meanings:
        raise ValueError(f'Did not find basin [{basin}] in given values [{flag_meanings}]')
    else:
        flag_matched = [f for f in flag_zipped if f[0] == basin][0]
    
    print(flag_matched)
    ds_masked = ds.where(mask==int(flag_matched[1]), drop=drop)

with e.g.

ds_masked = ds.where(mask.cf == basin)

Thought?

[dcherian]

Looks nice. I think I like the more explicit version with a CfFlags class:

1. mask.cf.flag == basin and
2. mask.cf.flag.isin(["atlantic_ocean", "pacific_ocean"])

We should also add a flags section to the repr.

PR?

[jbusecke]

I dont think I can put one in today (real WTF energy over here haha). But would love to do this at a later point!

[dcherian]

Leaving this open for flag masks. Contributions welcome!

[mathause]

Oh nice, didn't know about this - maybe I could use this for 2D masks in regionmask somehow. (It's always a bit a bummer that there is no way to identify regions by name for 2D masks.)

[dcherian]

Can you provide an example of what you want?

The documentation is decent: https://cf-xarray.readthedocs.io/en/latest/flags.html

[jbusecke]

@mathause if we could work functionality like in my example above into regionmask that would be amazing. Together with this I might be able to completely remove the regionmask module from cmip6_preprocessing (I would really like that!).
Let me know if you would like some review, comments or anything else that would help.