175 lines
7.2 KiB
Python
175 lines
7.2 KiB
Python
"""
|
|
Run this code by using the following wrapper script:
|
|
/tools/modify_input_files/mesh_mask_modifier
|
|
|
|
The wrapper script includes a full description and instructions.
|
|
"""
|
|
|
|
import logging
|
|
|
|
from math import isclose
|
|
import numpy as np
|
|
import xarray as xr
|
|
|
|
from ctsm.utils import abort
|
|
from ctsm.config_utils import lon_range_0_to_360
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class ModifyMeshMask:
|
|
"""
|
|
Description
|
|
-----------
|
|
"""
|
|
|
|
# The mesh_mask_modifier tool reads landmask, while the modify_fsurdat tool
|
|
# reads mod_lnd_props from the landmask file. Sample landmask
|
|
# file: fill_indianocean_slevis.nc located here as of 2022/8/1:
|
|
# /glade/work/slevis/git/mksurfdata_toolchain/tools/modify_input_files ...
|
|
# ... /islas_examples/modify_fsurdat/fill_indian_ocean/
|
|
# Read mod_lnd_props here only for consistency checks
|
|
def __init__(self, my_data, landmask_file, lat_dimname, lon_dimname, lat_varname, lon_varname):
|
|
|
|
self.file = my_data
|
|
|
|
# landmask from user-specified .nc file in the .cfg file
|
|
self._landmask_file = xr.open_dataset(landmask_file)
|
|
|
|
assert lat_varname in self._landmask_file.variables
|
|
assert lon_varname in self._landmask_file.variables
|
|
assert lat_dimname in self._landmask_file.dims
|
|
assert lon_dimname in self._landmask_file.dims
|
|
|
|
# latvar, lonvar same as right-hand-sides, which may be 1D or 2D
|
|
self.latvar = self._landmask_file[lat_varname][:, ...]
|
|
self.lonvar = self._landmask_file[lon_varname][..., :]
|
|
self.lsmlat = self._landmask_file.dims[lat_dimname]
|
|
self.lsmlon = self._landmask_file.dims[lon_dimname]
|
|
|
|
lonvar_first = self.lonvar[..., 0].data.max()
|
|
lonvar_last = self.lonvar[..., -1].data.max()
|
|
|
|
# If self.file["centerCoords"][0, 0] < 0, this suggests a
|
|
# -180 to 180 grid, so make self.lonvar -180 to 180 if not already.
|
|
# If self.file["centerCoords"][0, 0] >= 0, this suggests a
|
|
# 0 to 360 grid, so make self.lonvar 0 to 360 if not already.
|
|
if (self.file["centerCoords"][0, 0] < 0 <= lonvar_first) or (
|
|
self.file["centerCoords"][0, 0] >= 0 > lonvar_first
|
|
):
|
|
logger.info(
|
|
"first lon_mesh = %s, last lon_mesh = %s",
|
|
self.file["centerCoords"][0, 0].data,
|
|
self.file["centerCoords"][-1, 0].data,
|
|
)
|
|
logger.info("first lonvar = %s, last lonvar = %s.", lonvar_first, lonvar_last)
|
|
explanation = (
|
|
"For consistency in their order, changing lonvar to "
|
|
+ "lon_mesh's convention and later (in set_mesh_mask) "
|
|
+ "changing any negative longitude values to their "
|
|
+ "corresponding positive values."
|
|
)
|
|
logger.info(explanation)
|
|
self._landmask_file = self._landmask_file.roll(lsmlon=self.lsmlon // 2)
|
|
self.lonvar = self._landmask_file[lon_varname][..., :] # update lonvar
|
|
|
|
@classmethod
|
|
def init_from_file(
|
|
cls, file_in, landmask_file, lat_dimname, lon_dimname, lat_varname, lon_varname
|
|
):
|
|
"""Initialize a ModifyMeshMask object from file_in"""
|
|
logger.info("Opening file to be modified: %s", file_in)
|
|
my_file = xr.open_dataset(file_in)
|
|
return cls(my_file, landmask_file, lat_dimname, lon_dimname, lat_varname, lon_varname)
|
|
|
|
def set_mesh_mask(self, var):
|
|
"""
|
|
Sets 1d mask variable var = ocnmask = not(landmask).
|
|
Assumes the 1d vector is in the same south-to-north west-to-east
|
|
order as the 2d array.
|
|
"""
|
|
|
|
# Initialize
|
|
ncount = 0
|
|
mod_lnd_props = self._landmask_file.mod_lnd_props
|
|
landmask = self._landmask_file.landmask
|
|
|
|
for row in range(self.lsmlat): # rows from landmask file
|
|
logger.info("row = %d", row + 1)
|
|
for col in range(self.lsmlon): # cols from landmask file
|
|
errmsg = (
|
|
"landmask not 0 or 1 at row, col, value = "
|
|
+ f"{row} {col} {landmask[row, col]}"
|
|
)
|
|
assert landmask[row, col] == 0 or landmask[row, col] == 1, errmsg
|
|
errmsg = (
|
|
"mod_lnd_props not 0 or 1 at row, col, value = "
|
|
+ f"{row} {col} {mod_lnd_props[row, col]}"
|
|
)
|
|
assert mod_lnd_props[row, col] == 0 or mod_lnd_props[row, col] == 1, errmsg
|
|
if int(mod_lnd_props[row, col]) == 1:
|
|
errmsg = (
|
|
"landmask should = mod_lnd_props where the "
|
|
+ f"latter equals 1, but here landmask = 0 at row, col = {row} {col}"
|
|
)
|
|
assert int(landmask[row, col]) == 1, errmsg
|
|
# Reshape landmask into the
|
|
# elementCount dimension of the mesh file.
|
|
# In the process overwrite self.file[var].
|
|
ocnmask = np.logical_not(int(landmask[row, col]))
|
|
self.file[var][ncount] = ocnmask
|
|
|
|
# All else in this function supports error checking
|
|
|
|
# lon and lat from the landmask file
|
|
if len(self.latvar.sizes) == 2:
|
|
latvar_scalar = float(self.latvar[row, col])
|
|
lonvar_scalar = float(self.lonvar[row, col])
|
|
elif len(self.latvar.sizes) == 1:
|
|
latvar_scalar = float(self.latvar[row])
|
|
lonvar_scalar = float(self.lonvar[col])
|
|
else:
|
|
errmsg = (
|
|
"ERROR: Expecting latvar.sizes == 1 or 2, not "
|
|
+ f"{len(self.latvar.sizes)}"
|
|
)
|
|
abort(errmsg)
|
|
# ensure lon range of 0-360 rather than -180 to 180
|
|
lonvar_scalar = lon_range_0_to_360(lonvar_scalar)
|
|
# lon and lat from the mesh file
|
|
lat_mesh = float(self.file["centerCoords"][ncount, 1])
|
|
lon_mesh = float(self.file["centerCoords"][ncount, 0])
|
|
# ensure lon range of 0-360 rather than -180 to 180
|
|
lon_mesh = lon_range_0_to_360(lon_mesh)
|
|
|
|
errmsg = (
|
|
"Must be equal: "
|
|
" latvar_scalar = "
|
|
+ str(latvar_scalar)
|
|
+ " lat_mesh = "
|
|
+ str(lat_mesh)
|
|
+ " (at ncount = "
|
|
+ str(ncount)
|
|
+ ")"
|
|
)
|
|
assert isclose(latvar_scalar, lat_mesh, abs_tol=1e-5), errmsg
|
|
errmsg = (
|
|
"Must be equal: "
|
|
" lonvar_scalar = "
|
|
+ str(lonvar_scalar)
|
|
+ " lon_mesh = "
|
|
+ str(lon_mesh)
|
|
+ " (at ncount = "
|
|
+ str(ncount)
|
|
+ ")"
|
|
)
|
|
assert isclose(lonvar_scalar, lon_mesh, abs_tol=1e-5), errmsg
|
|
|
|
# increment counter
|
|
ncount = ncount + 1
|
|
|
|
# Error check
|
|
element_count = int(max((self.file["elementCount"])) + 1)
|
|
errmsg = "element_count =" + str(element_count) + "ncount =" + str(ncount) + "must be equal"
|
|
assert ncount == element_count, errmsg
|