ERA5/regrid_xesmf.py at main · so-wise/ERA5 · 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
85
86
87
88
89
90
91
92
93
94
95
96
97

# Import modules
import matplotlib.pyplot as plt
import cartopy.crs as ccrs
import MITgcmutils as mit
from tqdm import tqdm
import xarray as xr
import numpy as np
import xesmf as xe
import os

# Data locations
fin  = 'netcdf_tmp/'
fout = 'bin_tmp/'

# List of all NetCDF files in the input directory
files = [f for f in os.listdir(fin) if f.endswith('.nc')]

#
# Input grid (ERA5)
#

# Load ERA5 dataset (just u10m for now, to calculate regridder)
ds = xr.open_dataset(fin + 'ERA5_sowise_u10m_1992.nc')

# Extract longitude and longitude, define input grid
XC_era = ds['longitude'].values
YC_era = ds['latitude'].values

# Define input grid
grid_in = {"lon": XC_era,
           "lat": YC_era}

# Load data and convert to Dask array (each time snapshot is a block)
var = list(ds.data_vars.values())[0].values

#
# Output grid (MITgcm SO-WISE configuration)
#

# Get the MITgcm model grid
XC_model = mit.mds.rdmds('../MITgcm/so-wise-gyre/grid/XC')[0,:] - 360
YC_model = mit.mds.rdmds('../MITgcm/so-wise-gyre/grid/YC')[:,0]
XG_model = mit.mds.rdmds('../MITgcm/so-wise-gyre/grid/XG')[0,:] - 360
YG_model = mit.mds.rdmds('../MITgcm/so-wise-gyre/grid/YG')[:,0]

# Define output grid
grid_out = {"lon": XC_model,
            "lat": YC_model,
            "lon_b": XG_model,
            "lat_b": YG_model}

#
# Regridder
#

# Calculate the regridding function (or load from existing weights)
regridder = xe.Regridder(grid_in, grid_out, "bilinear")
regridder

# Test the regridder out on an example slice
test_out = regridder(var[0,:,:])

# Loop over all NetCDF files in the input directory
for filename in files:

    # Display name of file that is being converted
    np.disp('Now working on:')
    np.disp(filename)

    # Initialise an array for the regridded output
    var_out = np.empty((var.shape[0], test_out.shape[0], test_out.shape[1]))

    # Open the NetCDF file
    ds = xr.open_dataset(os.path.join(fin, filename))

    # Get the plain numpy array
    var = list(ds.data_vars.values())[0].values

    # Loop over the time dimension, regrid each time slice
    for n in tqdm(range(var.shape[0])):

        # Apply regridder to each time slice
        var_out[n] = regridder(var[n])

    # Flip the latitude dimension to match MITgcm convention
    var_out = np.flip(var_out, 1)

    # Change from float64 to float32
    var_out = np.float32(var_out)

    # Transpose to (time, lat, lon) [I have tried all permutations, none of them work]
#   var_out = np.transpose(var_out,(1,2,0))
    var_out = np.transpose(var_out,(0,2,1))

    # Write results to 32-bit binary file
    var_out.astype('>f4').tofile(os.path.join(fout, filename.replace('.nc', '')))