Skip to content

Commit

Permalink
black
Browse files Browse the repository at this point in the history
  • Loading branch information
@anton-seaice
anton-seaice committed Apr 11, 2024
1 parent f963def commit 43e3031
Show file tree
Hide file tree
Showing 20 changed files with 527 additions and 545 deletions.
160 changes: 74 additions & 86 deletions esmgrids/base_grid.py
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@

import numpy as np
import netCDF4 as nc

Expand All @@ -9,46 +8,46 @@ class BaseGrid(object):

def __init__(self, **kwargs):

x_t = kwargs.get('x_t', None)
y_t = kwargs.get('y_t', None)
self.x_u = kwargs.get('x_u', None)
self.y_u = kwargs.get('y_u', None)
self.x_v = kwargs.get('x_v', None)
self.y_v = kwargs.get('y_v', None)
x_t = kwargs.get("x_t", None)
y_t = kwargs.get("y_t", None)
self.x_u = kwargs.get("x_u", None)
self.y_u = kwargs.get("y_u", None)
self.x_v = kwargs.get("x_v", None)
self.y_v = kwargs.get("y_v", None)

dx_t = kwargs.get('dx_t', None)
dy_t = kwargs.get('dy_t', None)
dx_t = kwargs.get("dx_t", None)
dy_t = kwargs.get("dy_t", None)

self.dx_tn = kwargs.get('dx_tn', None)
self.dy_te = kwargs.get('dy_te', None)
self.dx_tn = kwargs.get("dx_tn", None)
self.dy_te = kwargs.get("dy_te", None)

self.dx_u = kwargs.get('dx_u', None)
self.dy_u = kwargs.get('dy_u', None)
self.dx_v = kwargs.get('dx_v', None)
self.dy_v = kwargs.get('dy_v', None)
self.dx_u = kwargs.get("dx_u", None)
self.dy_u = kwargs.get("dy_u", None)
self.dx_v = kwargs.get("dx_v", None)
self.dy_v = kwargs.get("dy_v", None)

self.area_t = kwargs.get('area_t', None)
self.area_u = kwargs.get('area_u', None)
self.area_v = kwargs.get('area_v', None)
self.area_t = kwargs.get("area_t", None)
self.area_u = kwargs.get("area_u", None)
self.area_v = kwargs.get("area_v", None)

self.clat_t = kwargs.get('clat_t', None)
self.clon_t = kwargs.get('clon_t', None)
self.clat_u = kwargs.get('clat_u', None)
self.clon_u = kwargs.get('clon_u', None)
self.clat_v = kwargs.get('clat_v', None)
self.clon_v = kwargs.get('clon_v', None)
self.clat_t = kwargs.get("clat_t", None)
self.clon_t = kwargs.get("clon_t", None)
self.clat_u = kwargs.get("clat_u", None)
self.clon_u = kwargs.get("clon_u", None)
self.clat_v = kwargs.get("clat_v", None)
self.clon_v = kwargs.get("clon_v", None)

self.angle_t = kwargs.get('angle_t', None)
self.angle_u = kwargs.get('angle_u', None)
self.angle_v = kwargs.get('angle_v', None)
self.angle_t = kwargs.get("angle_t", None)
self.angle_u = kwargs.get("angle_u", None)
self.angle_v = kwargs.get("angle_v", None)

self.mask_t = kwargs.get('mask_t', None)
self.mask_u = kwargs.get('mask_u', None)
self.mask_v = kwargs.get('mask_v', None)
self.mask_t = kwargs.get("mask_t", None)
self.mask_u = kwargs.get("mask_u", None)
self.mask_v = kwargs.get("mask_v", None)

self.z = kwargs.get('levels', [0])
self.description = kwargs.get('description', '')
self.calc_areas = kwargs.get('calc_areas', True)
self.z = kwargs.get("levels", [0])
self.description = kwargs.get("description", "")
self.calc_areas = kwargs.get("calc_areas", True)

if len(x_t.shape) == 1:
# Tile x_t
Expand All @@ -66,8 +65,8 @@ def __init__(self, **kwargs):
self.dy_t = self.dy_t.transpose()
elif dx_t is None:
# Repeat the first row/column using the values closest.
self.dx_t = np.ndarray((self.x_t.shape[0], self.x_t.shape[1]-1))
self.dy_t = np.ndarray((self.x_t.shape[0]-1, self.x_t.shape[1]))
self.dx_t = np.ndarray((self.x_t.shape[0], self.x_t.shape[1] - 1))
self.dy_t = np.ndarray((self.x_t.shape[0] - 1, self.x_t.shape[1]))
self.dx_t[:, :] = self.x_t[:, 1:] - self.x_t[:, :-1]
self.dy_t[:, :] = self.y_t[1:, :] - self.y_t[:-1, :]
else:
Expand All @@ -80,9 +79,7 @@ def __init__(self, **kwargs):

if self.mask_t is None:
# Default is all unmasked, up to user to mask.
self.mask_t = np.zeros((self.num_levels,
self.num_lat_points, self.num_lon_points),
dtype='int')
self.mask_t = np.zeros((self.num_levels, self.num_lat_points, self.num_lon_points), dtype="int")
if self.mask_u is None:
# FIXME
self.mask_u = self.mask_t
Expand All @@ -91,8 +88,7 @@ def __init__(self, **kwargs):
self.mask_v = self.mask_t

if self.clon_t is None or self.clat_t is None:
self.clat_t, self.clon_t, _, _, _, _ = \
make_corners(self.x_t, self.y_t, self.dx_t, self.dy_t)
self.clat_t, self.clon_t, _, _, _, _ = make_corners(self.x_t, self.y_t, self.dx_t, self.dy_t)

# The base class may implement this if it has a hole at the north
# pole to be fixed.
Expand All @@ -108,11 +104,9 @@ def __init__(self, **kwargs):
def fix_pole_holes(self):
pass


def fix_pole_overruns(self):
pass


@classmethod
def fromgrid(cls, grid):
"""
Expand All @@ -121,13 +115,12 @@ def fromgrid(cls, grid):

return cls(**grid.__dict__)

def write_scrip(self, filename, mask=None,
write_test_scrip=True, history=''):
def write_scrip(self, filename, mask=None, write_test_scrip=True, history=""):
"""
Write out grid in SCRIP format.
"""

f = nc.Dataset(filename, 'w')
f = nc.Dataset(filename, "w")

x = self.x_t
y = self.y_t
Expand All @@ -136,24 +129,24 @@ def write_scrip(self, filename, mask=None,
clon = self.clon_t
num_points = self.num_lat_points * self.num_lon_points

f.createDimension('grid_size', num_points)
f.createDimension('grid_corners', 4)
f.createDimension('grid_rank', 2)
f.createDimension("grid_size", num_points)
f.createDimension("grid_corners", 4)
f.createDimension("grid_rank", 2)

grid_dims = f.createVariable('grid_dims', 'i4', ('grid_rank'))
grid_dims = f.createVariable("grid_dims", "i4", ("grid_rank"))
# SCRIP likes lon, lat
grid_dims[:] = [self.num_lon_points, self.num_lat_points]

center_lat = f.createVariable('grid_center_lat', 'f8', ('grid_size'))
center_lat.units = 'degrees'
center_lat = f.createVariable("grid_center_lat", "f8", ("grid_size"))
center_lat.units = "degrees"
center_lat[:] = y[:].flatten()

center_lon = f.createVariable('grid_center_lon', 'f8', ('grid_size'))
center_lon.units = 'degrees'
center_lon = f.createVariable("grid_center_lon", "f8", ("grid_size"))
center_lon.units = "degrees"
center_lon[:] = x[:].flatten()

imask = f.createVariable('grid_imask', 'i4', ('grid_size'))
imask.units = 'unitless'
imask = f.createVariable("grid_imask", "i4", ("grid_size"))
imask.units = "unitless"

# Invert the mask. SCRIP uses zero for points that do not
# participate.
Expand All @@ -167,17 +160,15 @@ def write_scrip(self, filename, mask=None,
else:
imask[:] = 1 - mask[0, :, :].flatten()

corner_lat = f.createVariable('grid_corner_lat', 'f8',
('grid_size', 'grid_corners'))
corner_lat.units = 'degrees'
corner_lat = f.createVariable("grid_corner_lat", "f8", ("grid_size", "grid_corners"))
corner_lat.units = "degrees"
corner_lat[:, 0] = clat[0, :, :].flatten()
corner_lat[:, 1] = clat[1, :, :].flatten()
corner_lat[:, 2] = clat[2, :, :].flatten()
corner_lat[:, 3] = clat[3, :, :].flatten()

corner_lon = f.createVariable('grid_corner_lon', 'f8',
('grid_size', 'grid_corners'))
corner_lon.units = 'degrees'
corner_lon = f.createVariable("grid_corner_lon", "f8", ("grid_size", "grid_corners"))
corner_lon.units = "degrees"
corner_lon[:, 0] = clon[0, :, :].flatten()
corner_lon[:, 1] = clon[1, :, :].flatten()
corner_lon[:, 2] = clon[2, :, :].flatten()
Expand All @@ -188,53 +179,51 @@ def write_scrip(self, filename, mask=None,
f.close()

if write_test_scrip:
self.write_test_scrip(filename + '_test')
self.write_test_scrip(filename + "_test")

def write_test_scrip(self, filename):
"""
Write out SCRIP grid contents in a format which is easier to read/test.
"""

f = nc.Dataset(filename, 'w')
f = nc.Dataset(filename, "w")

x = self.x_t
y = self.y_t
clat = self.clat_t
clon = self.clon_t

f.createDimension('lats', self.num_lat_points)
f.createDimension('lons', self.num_lon_points)
f.createDimension('grid_corners', 4)
f.createDimension('grid_rank', 2)
f.createDimension("lats", self.num_lat_points)
f.createDimension("lons", self.num_lon_points)
f.createDimension("grid_corners", 4)
f.createDimension("grid_rank", 2)

center_lat = f.createVariable('center_lat', 'f8', ('lats', 'lons'))
center_lat.units = 'degrees'
center_lat = f.createVariable("center_lat", "f8", ("lats", "lons"))
center_lat.units = "degrees"
center_lat[:] = y[:]

center_lon = f.createVariable('center_lon', 'f8', ('lats', 'lons'))
center_lon.units = 'degrees'
center_lon = f.createVariable("center_lon", "f8", ("lats", "lons"))
center_lon.units = "degrees"
center_lon[:] = x[:]

imask = f.createVariable('mask', 'i4', ('lats', 'lons'))
imask.units = 'unitless'
imask = f.createVariable("mask", "i4", ("lats", "lons"))
imask.units = "unitless"
# Invert the mask. SCRIP uses zero for points that do not
# participate.
if len(self.mask_t.shape) == 2:
imask[:] = 1 - self.mask_t[:]
else:
imask[:] = 1 - self.mask_t[0, :, :]

corner_lat = f.createVariable('corner_lat', 'f8',
('lats', 'lons', 'grid_corners'))
corner_lat.units = 'degrees'
corner_lat = f.createVariable("corner_lat", "f8", ("lats", "lons", "grid_corners"))
corner_lat.units = "degrees"
corner_lat[:, :, 0] = clat[0, :, :].flatten()
corner_lat[:, :, 1] = clat[1, :, :].flatten()
corner_lat[:, :, 2] = clat[2, :, :].flatten()
corner_lat[:, :, 3] = clat[3, :, :].flatten()

corner_lon = f.createVariable('corner_lon', 'f8',
('lats', 'lons', 'grid_corners'))
corner_lon.units = 'degrees'
corner_lon = f.createVariable("corner_lon", "f8", ("lats", "lons", "grid_corners"))
corner_lon.units = "degrees"
corner_lon[:, :, 0] = clon[0, :, :].flatten()
corner_lon[:, :, 1] = clon[1, :, :].flatten()
corner_lon[:, :, 2] = clon[2, :, :].flatten()
Expand All @@ -243,7 +232,6 @@ def write_test_scrip(self, filename):
f.close()



def make_corners(x, y, dx, dy):

dx_half = dx / 2.0
Expand All @@ -266,7 +254,7 @@ def make_corners(x, y, dx, dy):
clon[2, :, :-1] = x[:, :-1] + dx_half[:, :]
clon[2, :, -1] = x[:, -1] + dx_half[:, -1]

assert(not np.isnan(np.sum(clon)))
assert not np.isnan(np.sum(clon))

clat = np.empty((4, nrow, ncol))
clat[:] = np.NAN
Expand All @@ -281,14 +269,14 @@ def make_corners(x, y, dx, dy):
clat[3, :-1, :] = y[:-1, :] + dy_half[:, :]
clat[3, -1, :] = y[-1, :] + dy_half[-1, :]

assert(not np.isnan(np.sum(clat)))
assert not np.isnan(np.sum(clat))

# The bottom latitude band should always be Southern extent.
assert(np.all(clat[0, 0, :] == np.min(y) - dy_half[0, :]))
assert(np.all(clat[1, 0, :] == np.min(y) - dy_half[0, :]))
assert np.all(clat[0, 0, :] == np.min(y) - dy_half[0, :])
assert np.all(clat[1, 0, :] == np.min(y) - dy_half[0, :])

# The top latitude band should always be Northern extent.
assert(np.all(clat[2, -1, :] == np.max(y) + dy_half[-1, :]))
assert(np.all(clat[3, -1, :] == np.max(y) + dy_half[-1, :]))
assert np.all(clat[2, -1, :] == np.max(y) + dy_half[-1, :])
assert np.all(clat[3, -1, :] == np.max(y) + dy_half[-1, :])

return clat, clon, None, None, None, None
Loading

0 comments on commit 43e3031

Please sign in to comment.