Azores_TidalForcing/Code/Calculations/kde_func.py at main · Parcels-code/Azores_TidalForcing · 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
# conda activate venv_kde_xarray

## Written on 18/01/2021 by Laura Gomez Navarro

# Modules:

from glob import glob
import xarray as xr
import numpy as np
from scipy.stats import gaussian_kde

def rem_nans(ds):
    """
    This renders lon and lat variables without nans for the last timestep.
    """
    bad_indices = np.isnan(ds['lon'][:,-1]) | np.isnan(ds['lat'][:,-1])
    good_indices = ~bad_indices
    lon_end_nonans = ds['lon'][:,-1][good_indices]
    lat_end_nonans = ds['lat'][:,-1][good_indices]

    return lon_end_nonans, lat_end_nonans

def kde_vals(lon_end_nonans, lat_end_nonans):
    """
    TO DO
    """
    x = lon_end_nonans.copy()
    y = lat_end_nonans.copy()

    xy = np.vstack([x, y])
    z = gaussian_kde(xy)(xy)

    idx = z.argsort()
    x, y, z = x[idx], y[idx], z[idx]

    return x, y, z

def kde_parcels(nfile):
    """
    Inputs
    nfile : filedir + filename

    Outputs
    KDE for end fields of lon and lat
    """

    ds = xr.open_dataset(nfile)

    # Remove nans:
    lon_end_nonans, lat_end_nonans = rem_nans(ds)

    kde_x, kde_y, kde_z = kde_vals(lon_end_nonans, lat_end_nonans)

    savedir = '/data/oceanparcels/output_data/data_LauraGN/outputs_parcels/kde_calcs/'
    savename = savedir + 'KDE_' + nfile.split('/')[-1].split('.nc')[0] + '.npz'
    np.savez(savename, kde_x=kde_x, kde_y=kde_y, kde_z=kde_z)

    return kde_x, kde_y, kde_z

def kde_parcels_bws(nfile, bw_set):
    """
    bw_set = np.arange(0.01, 0.26, 0.01)
    """
    ds = xr.open_dataset(nfile)

    # Remove nans:
    lon_end_nonans, lat_end_nonans = rem_nans(ds)

    x = lon_end_nonans.copy()
    y = lat_end_nonans.copy()

    xy = np.vstack([x, y])

    KDE_dict = {}
    for ii in range(0, len(bw_set)):
        KDE_dict["z%02d" %ii] =  gaussian_kde(xy, bw_method=bw_set[ii])(xy)

    # Saving:
    savedir = '/data/oceanparcels/output_data/data_LauraGN/outputs_parcels/kde_calcs/'
    savename = savedir + 'KDE_' + nfile.split('/')[-1].split('.nc')[0] + '_bws.npz'
    np.savez(savename, bw_set=bw_set, KDE_dict=KDE_dict)