From d16fc4c13c3afb786ca281f21bee4ba802db0a85 Mon Sep 17 00:00:00 2001 From: manishvenu Date: Thu, 26 Sep 2024 16:25:36 -0600 Subject: [PATCH] Black formatting --- regional_mom6/regional_mom6.py | 50 ++++++++++++++++++++++++---------- 1 file changed, 36 insertions(+), 14 deletions(-) diff --git a/regional_mom6/regional_mom6.py b/regional_mom6/regional_mom6.py index 618de52a..96e8987c 100644 --- a/regional_mom6/regional_mom6.py +++ b/regional_mom6/regional_mom6.py @@ -1204,7 +1204,7 @@ def setup_tides_rectangle_boundaries( ### Find Rough Horizontal Subset (with 0.5 Buffer)### - if tidal_constituents != 'read_from_expt_init': + if tidal_constituents != "read_from_expt_init": self.tidal_constituents = tidal_constituents tpxo_h = ( xr.open_dataset(os.path.join(path_to_td, f"h_{tidal_filename}")) @@ -2284,8 +2284,12 @@ def coords(self): rcoord = rcoord.rename_dims({"nxp": f"nx_{self.segment_name}"}) rcoord.attrs["perpendicular"] = "ny" rcoord.attrs["parallel"] = "nx" - rcoord.attrs["axis_to_expand"] = 2 ## Need to keep track of which axis the 'main' coordinate corresponds to for rectangular_brushcut on when re-adding the 'secondary' axis - rcoord.attrs["locations_name"] = f"nx_{self.segment_name}" # Legacy name of nx_... was locations. This provides a clear transform in regrid_tides + rcoord.attrs["axis_to_expand"] = ( + 2 ## Need to keep track of which axis the 'main' coordinate corresponds to for rectangular_brushcut on when re-adding the 'secondary' axis + ) + rcoord.attrs["locations_name"] = ( + f"nx_{self.segment_name}" # Legacy name of nx_... was locations. This provides a clear transform in regrid_tides + ) elif self.orientation == "north": rcoord = xr.Dataset( { @@ -2325,7 +2329,7 @@ def coords(self): rcoord.attrs["parallel"] = "ny" rcoord.attrs["axis_to_expand"] = 3 rcoord.attrs["locations_name"] = f"ny_{self.segment_name}" - + # Make lat and lon coordinates rcoord = rcoord.assign_coords(lat=rcoord["lat"], lon=rcoord["lon"]) @@ -2524,7 +2528,8 @@ def rectangular_brushcut(self): ## Re-add the secondary dimension (even though it represents one value..) segment_out[v] = segment_out[v].expand_dims( - f"{self.coords.attrs["perpendicular"]}_{self.segment_name}", axis=self.coords.attrs["axis_to_expand"] + f"{self.coords.attrs["perpendicular"]}_{self.segment_name}", + axis=self.coords.attrs["axis_to_expand"], ) ## Add the layer thicknesses @@ -2569,7 +2574,8 @@ def rectangular_brushcut(self): segment_out[f"eta_{self.segment_name}"] = segment_out[ f"eta_{self.segment_name}" ].expand_dims( - f"{self.coords.attrs["perpendicular"]}_{self.segment_name}", axis=self.coords.attrs["axis_to_expand"] - 1 + f"{self.coords.attrs["perpendicular"]}_{self.segment_name}", + axis=self.coords.attrs["axis_to_expand"] - 1, ) # Overwrite the actual lat/lon values in the dimensions, replace with incrementing integers @@ -2600,11 +2606,15 @@ def rectangular_brushcut(self): # Store actual lat/lon values here as variables rather than coordinates segment_out[f"lon_{self.segment_name}"] = ( [f"ny_{self.segment_name}", f"nx_{self.segment_name}"], - self.coords.lon.expand_dims(dim="blank",axis=self.coords.attrs["axis_to_expand"] - 2).data, + self.coords.lon.expand_dims( + dim="blank", axis=self.coords.attrs["axis_to_expand"] - 2 + ).data, ) segment_out[f"lat_{self.segment_name}"] = ( [f"ny_{self.segment_name}", f"nx_{self.segment_name}"], - self.coords.lon.expand_dims(dim="blank",axis=self.coords.attrs["axis_to_expand"] - 2).data, + self.coords.lon.expand_dims( + dim="blank", axis=self.coords.attrs["axis_to_expand"] - 2 + ).data, ) # Add units to the lat / lon to keep the `categorize_axis_from_units` checker happy @@ -2679,8 +2689,12 @@ def regrid_tides( # Fill missing data. # Need to do this first because complex would get converted to real - redest = redest.ffill(dim=self.coords.attrs["locations_name"], limit=None)["hRe"] - imdest = imdest.ffill(dim=self.coords.attrs["locations_name"], limit=None)["hIm"] + redest = redest.ffill(dim=self.coords.attrs["locations_name"], limit=None)[ + "hRe" + ] + imdest = imdest.ffill(dim=self.coords.attrs["locations_name"], limit=None)[ + "hIm" + ] # Convert complex cplex = redest + 1j * imdest @@ -2696,7 +2710,9 @@ def regrid_tides( # Add time coordinate and transpose so that time is first, # so that it can be the unlimited dimension ds_ap, _ = xr.broadcast(ds_ap, times) - ds_ap = ds_ap.transpose("time", "constituent", self.coords.attrs["locations_name"]) + ds_ap = ds_ap.transpose( + "time", "constituent", self.coords.attrs["locations_name"] + ) self.encode_tidal_files_and_output(ds_ap, "tz") @@ -2763,7 +2779,9 @@ def regrid_tides( # Need to transpose so that time is first, # so that it can be the unlimited dimension - ds_ap = ds_ap.transpose("time", "constituent", self.coords.attrs["locations_name"]) + ds_ap = ds_ap.transpose( + "time", "constituent", self.coords.attrs["locations_name"] + ) # Some things may have become missing during the transformation ds_ap = ds_ap.ffill(dim=self.coords.attrs["locations_name"], limit=None) @@ -2822,9 +2840,13 @@ def encode_tidal_files_and_output(self, ds, filename): if "z" in ds.coords: ds = ds.rename({"z": f"nz_{self.segment_name}"}) if self.orientation in ["south", "north"]: - ds = ds.rename({self.coords.attrs["locations_name"]: f"nx_{self.segment_name}"}) + ds = ds.rename( + {self.coords.attrs["locations_name"]: f"nx_{self.segment_name}"} + ) elif self.orientation in ["west", "east"]: - ds = ds.rename({self.coords.attrs["locations_name"]: f"ny_{self.segment_name}"}) + ds = ds.rename( + {self.coords.attrs["locations_name"]: f"ny_{self.segment_name}"} + ) ## Perform Encoding ## for v in ds: