masked-verif/masked_verif.py at main · havardhhaugen/masked-verif · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
import xarray as xr
import numpy as np
import argparse
from sklearn.neighbors import BallTree

_EARTH_RADIUS_KM = 6371.0

def find_valid_locations(
        ds_template: xr.Dataset,
        ds_mask: xr.Dataset,
        mask_variable: str,
        mask_value: int = 1,
        cutoff: float = 5, # in kilometers
) -> list:
    """Find valid locations in ds_template based on a mask in ds_mask.

    A location in ds_template.location is considered valid if ds_mask[mask_variable] == mask_value for the point
    closest to the coordinates of the location.

    Parameters
    ----------
    ds_template : xr.Dataset
        Dateset with "location" (unique numbers) and corresponding coordinates ("lon","lat")
    ds_mask : xr.Dataset
        Dataset with mask. Coordinates given by ("longitude","latitude")
    mask_variable : str
        Variable that defines the mask in ds_mask (ds_mask[mask_variable])
    mask_value : int
        Mask value to define valid locations
    cutoff : float
        Cutoff distance  in km when checking if a location is valid (location must be < cutoff away from the closest point in mask)

    Returns
    -------
    list
        list of valid location numbers
    """

    lon_template = ds_template['lon'].values.flatten()
    lat_template = ds_template['lat'].values.flatten()
    locations = ds_template['location'].values.flatten()

    lon_mask = ds_mask['longitude'].values.flatten()
    lat_mask = ds_mask['latitude'].values.flatten()
    mask = ds_mask[mask_variable].values.flatten()

    mask_coords = np.deg2rad(np.column_stack([lat_mask, lon_mask]))
    tree = BallTree(mask_coords, metric='haversine')

    template_coords = np.deg2rad(np.column_stack([lat_template, lon_template]))

    dist_rad, ind = tree.query(template_coords, k=1)
    dist_km = dist_rad[:, 0] * _EARTH_RADIUS_KM
    nearest_idx = ind[:, 0]

    within_cutoff = dist_km <= cutoff
    correct_mask = mask[nearest_idx] == mask_value
    valid_locations = locations[within_cutoff & correct_mask].tolist()

    return valid_locations

def _main__():
    # Parse arguments from bash
    parser = argparse.ArgumentParser()
    parser.add_argument('--template-file', required=True, type=str, help='Path to the template dataset (NetCDF file)')
    parser.add_argument('--mask-file', required=True, type=str, help='Path to the mask dataset (NetCDF file)')
    parser.add_argument('--mask-variable', required=True, type=str, help='Variable name in the mask dataset to check for value 1')
    parser.add_argument('--mask-value', required=False, type=int, default=1, help='Value in the mask variable to consider as valid (default: 1)')
    parser.add_argument('--cutoff', required=False, type=float, default=5.0, help='Cutoff distance in kilometers (default: 5 km)')
    args = parser.parse_args()

    ds_template = xr.open_dataset(args.template_file)
    ds_mask = xr.open_dataset(args.mask_file)
    mask_variable = args.mask_variable
    mask_value = args.mask_value
    cutoff = args.cutoff

    # Find valid locations
    locations = find_valid_locations(ds_template, ds_mask, mask_variable, mask_value, cutoff)
    # Print valid locations (easiest way to send to bash-script)
    print(*locations)

if __name__ == "__main__":
    __main__()