mzamyatina / phd_ukca Goto Github PK

# Read ATom geo coords
DEFAULT_GROUP = "WAS"
species_deps_groups = {
    "ch4": {"deps": ["CH4_NOAA"], "group": "NOAA-Picarro"},
    "c2h6": {
        "deps": ["Ethane_WAS", "Ethene_WAS", "Ethyne_WAS"],
    },
    "c3h8": {
        "deps": ["Propane_WAS", "Propene_WAS"],
    },
    "no": {"deps": ["NO_CL"], "group": "NOyO3-NO"},
    "o3": {"deps": ["O3_CL"], "group": "NOyO3-O3"},
    "oh": {"deps": ["OH_pptv"], "group": "ATHOS-HOx"},
    "meono2": {
        "deps": ["MeONO2_WAS"],
    },
    "etono2": {
        "deps": ["EtONO2_WAS"],
    },
    "nprono2": {
        "deps": ["n-PrONO2_WAS"],
    },
    "iprono2": {
        "deps": ["i-PrONO2_WAS"],
    },
}

sample_lons, sample_lats = {}, {}  # for easy access of lon, lat

atom_data = {} # Yay, another dictionary!

for season in SEASONS:
    # Deal with lon, lat separately
    ds = xr.open_mfdataset(sorted(path_to_atom.glob(f'MER-WAS_DC8_20{season}*.nc')), group='MMS', decode_cf=True)
    sample_lons[season] = ds.G_LONG.values
    sample_lats[season] = ds.G_LAT.values
    # Create an empty dict for each season to store data for each of the groups
    atom_data[season] = {}

    for key in ['ch4', 'o3', 'oh']:
        # Loop over selected groups
        atom_data[season][key] = xr.open_mfdataset(sorted(path_to_atom.glob(f'MER-WAS_DC8_20{season}*.nc')), group=species_deps_groups[key].get(group, DEFAULT_GROUP), decode_cf=True)

        _lon_lat = np.logical_or(*[np.isnan(atom_data[season][key][i].values) for i in atom_data[season][key][species_deps[the_species]]])
        # Mask out coordinates of missing data
        sample_lons[season][_lon_lat] = np.nan
        sample_lats[season][_lon_lat] = np.nan