Skip to content

Commit

Permalink
Use Path objects instead of string manipulation (#75)
Browse files Browse the repository at this point in the history 
Importantly, reduces number of subprocess/os calls, in favour of using
Python-native operations.
  • Loading branch information
@angus-g
angus-g authored Oct 4, 2023
1 parent 4c62a58 commit 6b5817e
Showing 1 changed file with 49 additions and 67 deletions.
116 changes: 49 additions & 67 deletions regional_mom6/regional_mom6.py
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import numpy as np
from itertools import cycle
import os
from pathlib import Path
import dask.array as da
import dask.bag as db
import numpy as np
Expand Down Expand Up @@ -475,15 +475,12 @@ def __init__(
toolpath,
gridtype="even_spacing",
):
try:
os.mkdir(mom_run_dir)
except:
pass
self.mom_run_dir = Path(mom_run_dir)
self.mom_input_dir = Path(mom_input_dir)

self.mom_run_dir.mkdir(exist_ok=True)
self.mom_input_dir.mkdir(exist_ok=True)

try:
os.mkdir(mom_input_dir)
except:
pass
self.xextent = xextent
self.yextent = yextent
self.daterange = [
Expand All @@ -494,29 +491,21 @@ def __init__(
self.vlayers = vlayers
self.dz_ratio = dz_ratio
self.depth = depth
self.mom_run_dir = mom_run_dir
self.mom_input_dir = mom_input_dir
self.toolpath = toolpath
self.hgrid = self._make_hgrid(gridtype)
self.vgrid = self._make_vgrid()
self.gridtype = gridtype
# if "temp" not in os.listdir(inputdir):
# os.mkdir(inputdir + "temp")

if "weights" not in os.listdir(self.mom_input_dir):
os.mkdir(mom_input_dir + "weights")
if "forcing" not in os.listdir(self.mom_input_dir):
os.mkdir(self.mom_input_dir + "forcing")

# create a simlink from input directory to run directory and vv
subprocess.run(
f"ln -s {self.mom_input_dir} {self.mom_run_dir}/inputdir", shell=True
)
subprocess.run(
f"ln -s {self.mom_run_dir} {self.mom_input_dir}/rundir", shell=True
)
# create additional directories and links
(self.mom_input_dir / "weights").mkdir(exist_ok=True)
(self.mom_input_dir / "forcing").mkdir(exist_ok=True)

return
run_inputdir = self.mom_run_dir / "inputdir"
if not run_inputdir.exists():
run_inputdir.symlink_to(self.mom_input_dir.resolve())
input_rundir = self.mom_input_dir / "rundir"
if not input_rundir.exists():
input_rundir.symlink_to(self.mom_run_dir.resolve())

def _make_hgrid(self, gridtype):
"""Sets up hgrid based on users specification of
Expand Down Expand Up @@ -548,7 +537,7 @@ def _make_hgrid(self, gridtype):

y = np.linspace(self.yextent[0], self.yextent[1], ny)
hgrid = rectangular_hgrid(x, y)
hgrid.to_netcdf(self.mom_input_dir + "/hgrid.nc")
hgrid.to_netcdf(self.mom_input_dir / "hgrid.nc")

return hgrid

Expand All @@ -566,7 +555,7 @@ def _make_vgrid(self):
} ## THIS MIGHT BE WRONG REVISIT
)
vcoord["zi"].attrs = {"units": "meters"}
vcoord.to_netcdf(self.mom_input_dir + "/vcoord.nc")
vcoord.to_netcdf(self.mom_input_dir / "vcoord.nc")

return vcoord

Expand All @@ -577,7 +566,7 @@ def ocean_forcing(
boundaries (if specified) and for initial condition
Args:
path (str): Path to directory containing the forcing
path (Union[str, Path]): Path to directory containing the forcing
files. Files should be named
``north_segment_unprocessed`` for each boundary (for
the cardinal directions) and ``ic_unprocessed`` for the
Expand All @@ -593,10 +582,12 @@ def ocean_forcing(
"""

path = Path(path)

## Do initial condition

## pull out the initial velocity on MOM5's Bgrid
ic_raw = xr.open_dataset(path + "/ic_unprocessed")
ic_raw = xr.open_dataset(path / "ic_unprocessed")

if varnames["time"] in ic_raw.dims:
ic_raw = ic_raw.isel({varnames["time"]: 0})
Expand Down Expand Up @@ -806,7 +797,7 @@ def ocean_forcing(
eta_out = eta_out.isel(time=0).drop("time")

vel_out.fillna(0).to_netcdf(
self.mom_input_dir + "forcing/init_vel.nc",
self.mom_input_dir / "forcing/init_vel.nc",
mode="w",
encoding={
"u": {"_FillValue": netCDF4.default_fillvals["f4"]},
Expand All @@ -815,7 +806,7 @@ def ocean_forcing(
)

tracers_out.to_netcdf(
self.mom_input_dir + "forcing/init_tracers.nc",
self.mom_input_dir / "forcing/init_tracers.nc",
mode="w",
encoding={
"xh": {"_FillValue": None},
Expand All @@ -826,7 +817,7 @@ def ocean_forcing(
},
)
eta_out.to_netcdf(
self.mom_input_dir + "forcing/init_eta.nc",
self.mom_input_dir / "forcing/init_eta.nc",
mode="w",
encoding={
"xh": {"_FillValue": None},
Expand All @@ -840,18 +831,20 @@ def ocean_forcing(
self.ic_tracers = tracers_out
self.ic_vels = vel_out

if boundaries == None:
if boundaries is None:
return

print("BRUSHCUT BOUNDARIES")

## Generate a rectangular OBC domain. This is the default configuration. For fancier domains, need to use the segment class manually
## Generate a rectangular OBC domain. This is the default
## configuration. For fancier domains, need to use the segment
## class manually
for i, o in enumerate(boundaries):
print(f"Processing {o}...", end="")
seg = segment(
self.hgrid,
f"{path}/{o.lower()}_unprocessed", # location of raw boundary
f"{self.mom_input_dir}", # Save path
path / f"{o.lower()}_unprocessed", # location of raw boundary
self.mom_input_dir,
varnames,
"segment_{:03d}".format(i + 1),
o.lower(), # orienataion
Expand Down Expand Up @@ -956,7 +949,7 @@ def bathymetry(
bathyout.elevation.attrs["long_name"] = "Elevation relative to sea level"
bathyout.elevation.attrs["coordinates"] = "lon lat"
bathyout.to_netcdf(
f"{self.mom_input_dir}bathy_original.nc", mode="w", engine="netcdf4"
self.mom_input_dir / "bathy_original.nc", mode="w", engine="netcdf4"
)

tgrid = xr.Dataset(
Expand Down Expand Up @@ -1003,7 +996,7 @@ def bathymetry(
tgrid.lon.attrs["_FillValue"] = 1e20
tgrid.lat.attrs["units"] = "degrees_north"
tgrid.to_netcdf(
f"{self.mom_input_dir}topog_raw.nc", mode="w", engine="netcdf4"
self.mom_input_dir / "topog_raw.nc", mode="w", engine="netcdf4"
)
tgrid.close()

Expand All @@ -1020,12 +1013,12 @@ def bathymetry(

topog = regridder(bathyout)
topog.to_netcdf(
f"{self.mom_input_dir}topog_raw.nc", mode="w", engine="netcdf4"
self.mom_input_dir / "topog_raw.nc", mode="w", engine="netcdf4"
)

## reopen topography to modify
print("Reading in regridded bathymetry to fix up metadata...", end="")
topog = xr.open_dataset(self.mom_input_dir + "topog_raw.nc", engine="netcdf4")
topog = xr.open_dataset(self.mom_input_dir / "topog_raw.nc", engine="netcdf4")

## Ensure correct encoding
topog = xr.Dataset(
Expand Down Expand Up @@ -1181,32 +1174,26 @@ def bathymetry(
topog["depth"] = topog["depth"].where(topog["depth"] != 0, np.nan)

topog.expand_dims({"ntiles": 1}).to_netcdf(
self.mom_input_dir + "topog_deseas.nc",
self.mom_input_dir / "topog_deseas.nc",
mode="w",
encoding={"depth": {"_FillValue": None}},
)

subprocess.run(
"mv topog_deseas.nc topog.nc", shell=True, cwd=self.mom_input_dir
)
(self.mom_input_dir / "topog_deseas.nc").rename(self.mom_input_dir / "topog.nc")
print("done.")
self.topog = topog
return

def FRE_tools(self, layout):
"""
Just a wrapper for FRE Tools check_mask, make_solo_mosaic and make_quick_mosaic. User provides processor layout tuple of processing units.
"""

if "topog.nc" not in os.listdir(self.mom_input_dir):
if not (self.mom_input_dir / "topog.nc").exists():
print("No topography file! Need to run make_bathymetry first")
return
try:
os.remove(
"mask_table*"
) ## Removes old mask table so as not to clog up inputdir
except:
pass

for p in self.mom_input_dir.glob("mask_table*"):
p.unlink()

print(
"MAKE SOLO MOSAIC",
Expand Down Expand Up @@ -1240,7 +1227,6 @@ def FRE_tools(self, layout):
),
)
self.layout = layout
return


class segment:
Expand All @@ -1258,8 +1244,8 @@ class segment:
Args:
hgrid (xarray.Dataset): The horizontal grid used for domain
infile (str): Path to the raw, unprocessed boundary segment
outfolder (str): Path to folder where the model inputs will be stored
infile (Union[str, Path]): Path to the raw, unprocessed boundary segment
outfolder (Union[str, Path]): Path to folder where the model inputs will be stored
varnames (Dict[str, str]): Mapping between the
variable/dimension names and standard naming convension of
this pipeline, e.g. ``{"xq":"longitude, "yh":"latitude",
Expand Down Expand Up @@ -1322,9 +1308,7 @@ def __init__(

def brushcut(self, ryf=False):
### Implement brushcutter scheme on single segment ###
# print(self.infile + f"/{self.orientation}_segment_unprocessed")
rawseg = xr.open_dataset(self.infile, decode_times=False, engine="netcdf4")
# rawseg = xr.open_dataset(self.infile,decode_times=False,chunks={self.time:30,self.z:25})

## Depending on the orientation of the segment, cut out the right bit of the hgrid
## and define which coordinate is along or into the segment
Expand Down Expand Up @@ -1374,7 +1358,7 @@ def brushcut(self, ryf=False):
locstream_out=True,
reuse_weights=False,
filename=self.outfolder
+ f"weights/bilinear_velocity_weights_{self.orientation}.nc",
/ f"weights/bilinear_velocity_weights_{self.orientation}.nc",
)

segment_out = xr.merge(
Expand All @@ -1397,7 +1381,7 @@ def brushcut(self, ryf=False):
locstream_out=True,
reuse_weights=False,
filename=self.outfolder
+ f"weights/bilinear_velocity_weights_{self.orientation}.nc",
/ f"weights/bilinear_velocity_weights_{self.orientation}.nc",
)

regridder_tracer = xe.Regridder(
Expand All @@ -1407,7 +1391,7 @@ def brushcut(self, ryf=False):
locstream_out=True,
reuse_weights=False,
filename=self.outfolder
+ f"weights/bilinear_tracer_weights_{self.orientation}.nc",
/ f"weights/bilinear_tracer_weights_{self.orientation}.nc",
)

segment_out = xr.merge(
Expand Down Expand Up @@ -1438,7 +1422,7 @@ def brushcut(self, ryf=False):
locstream_out=True,
reuse_weights=False,
filename=self.outfolder
+ f"weights/bilinear_uvelocity_weights_{self.orientation}.nc",
/ f"weights/bilinear_uvelocity_weights_{self.orientation}.nc",
)

regridder_vvelocity = xe.Regridder(
Expand All @@ -1448,7 +1432,7 @@ def brushcut(self, ryf=False):
locstream_out=True,
reuse_weights=False,
filename=self.outfolder
+ f"weights/bilinear_vvelocity_weights_{self.orientation}.nc",
/ f"weights/bilinear_vvelocity_weights_{self.orientation}.nc",
)

regridder_tracer = xe.Regridder(
Expand All @@ -1458,7 +1442,7 @@ def brushcut(self, ryf=False):
locstream_out=True,
reuse_weights=False,
filename=self.outfolder
+ f"weights/bilinear_tracer_weights_{self.orientation}.nc",
/ f"weights/bilinear_tracer_weights_{self.orientation}.nc",
)

segment_out = xr.merge(
Expand Down Expand Up @@ -1626,14 +1610,12 @@ def brushcut(self, ryf=False):
hgrid_seg.y.data,
)

# del segment_out["depth"]

with ProgressBar():
segment_out["time"] = segment_out["time"].assign_attrs(
{"modulo": " "}
) ## Add modulo attribute for MOM6 to treat as repeat forcing
segment_out.load().to_netcdf(
self.outfolder + f"forcing/forcing_obc_{self.seg_name}.nc",
self.outfolder / f"forcing/forcing_obc_{self.seg_name}.nc",
encoding=encoding_dict,
unlimited_dims="time",
)
Expand Down

0 comments on commit 6b5817e

Please sign in to comment.