From ea5a9c34a601b8698e99cf67a769285e5c1da921 Mon Sep 17 00:00:00 2001 From: manishvenu Date: Fri, 27 Sep 2024 15:41:02 -0600 Subject: [PATCH] Add properties and change function names --- demos/reanalysis-forced.ipynb | 2 +- regional_mom6/regional_mom6.py | 95 +++++++++++++++++++++++++--- tests/test_grid_generation.py | 4 +- tests/test_tides_functions_config.py | 18 +++++- 4 files changed, 106 insertions(+), 13 deletions(-) diff --git a/demos/reanalysis-forced.ipynb b/demos/reanalysis-forced.ipynb index df57fe3d..70f515c2 100644 --- a/demos/reanalysis-forced.ipynb +++ b/demos/reanalysis-forced.ipynb @@ -324,7 +324,7 @@ "metadata": {}, "outputs": [], "source": [ - "expt.FRE_tools(layout=(10, 10)) ## Here the tuple defines the processor layout" + "expt.run_FRE_tools(layout=(10, 10)) ## Here the tuple defines the processor layout" ] }, { diff --git a/regional_mom6/regional_mom6.py b/regional_mom6/regional_mom6.py index 22543a56..772abc27 100644 --- a/regional_mom6/regional_mom6.py +++ b/regional_mom6/regional_mom6.py @@ -23,13 +23,15 @@ ) import pandas as pd import re +from pathlib import Path +import glob warnings.filterwarnings("ignore") __all__ = [ "longitude_slicer", "hyperbolictan_thickness_profile", - "rectangular_hgrid", + "calculate_rectangular_hgrid", "experiment", "segment", ] @@ -212,9 +214,6 @@ def find_MOM6_rectangular_orientation(input): raise ValueError("Invalid type of Input, can only be string or int.") -from pathlib import Path - - def get_glorys_data( longitude_extent, latitude_extent, @@ -375,7 +374,7 @@ def hyperbolictan_thickness_profile(nlayers, ratio, total_depth): return layer_thicknesses -def rectangular_hgrid(lons, lats): +def calculate_rectangular_hgrid(lons, lats): """ Construct a horizontal grid with all the metadata required by MOM6, based on arrays of longitudes (``lons``) and latitudes (``lats``) on the supergrid. @@ -668,7 +667,7 @@ def _make_hgrid(self): self.latitude_extent[0], self.latitude_extent[1], ny ) # latitudes in degrees - hgrid = rectangular_hgrid(lons, lats) + hgrid = calculate_rectangular_hgrid(lons, lats) hgrid.to_netcdf(self.mom_input_dir / "hgrid.nc") return hgrid @@ -699,6 +698,74 @@ def _make_vgrid(self): return vcoord + @property + def ocean_state_boundaries(self): + """ + Read the ocean state files from disk, and print 'em + """ + ocean_state_path = self.mom_input_dir / "forcing" + try: + # Use glob to find all tides files + patterns = ["forcing_*", "weights/bi*"] + all_files = [] + for pattern in patterns: + all_files.extend(glob.glob(os.path.join(ocean_state_path, pattern))) + + if len(all_files) == 0: + return "No ocean state files set up yet (or files misplaced from {}). Call `setup_ocean_state_boundaries` method to set up ocean state.".format( + ocean_state_path + ) + + # Open the files as xarray datasets + datasets = [xr.open_dataset(file) for file in all_files] + return datasets + except: + return "Error retrieving ocean state files" + + @property + def tides_boundaries(self): + """ + Read the tides from disk, and print 'em + """ + tides_path = self.mom_input_dir / "forcing" + try: + # Use glob to find all tides files + patterns = ["regrid*", "tu_*", "tz_*"] + all_files = [] + for pattern in patterns: + all_files.extend(glob.glob(os.path.join(tides_path, pattern))) + + if len(all_files) == 0: + return "No tides files set up yet (or files misplaced from {}). Call `setup_tides_boundaries` method to set up tides.".format( + tides_path + ) + + # Open the files as xarray datasets + datasets = [xr.open_dataset(file) for file in all_files] + return datasets + except: + return "Error retrieving tides files" + + @property + def era5(self): + """ + Read the era5's from disk, and print 'em + """ + era5_path = self.mom_input_dir / "forcing" + try: + # Use glob to find all *_ERA5.nc files + all_files = glob.glob(os.path.join(era5_path, "*_ERA5.nc")) + if len(all_files) == 0: + return "No era5 files set up yet (or files misplaced from {}). Call `setup_era5` method to set up era5.".format( + era5_path + ) + + # Open the files as xarray datasets + datasets = [xr.open_dataset(file) for file in all_files] + return datasets + except: + return "Error retrieving ERA5 files" + @property def initial_condition(self): """ @@ -709,13 +776,25 @@ def initial_condition(self): ic_tracers = xr.open_dataset(self.mom_input_dir / "forcing/init_tracers.nc") ic_vel = xr.open_dataset(self.mom_input_dir / "forcing/init_vel.nc") ic_eta = xr.open_dataset(self.mom_input_dir / "forcing/init_eta.nc") - return ic_tracers, ic_vel, ic_eta + return [ic_tracers, ic_vel, ic_eta] except: return "No initial condition set up yet (or files misplaced from {}). Call `setup_initial_condition` method to set up initial conditions.".format( self.mom_input_dir / "forcing" ) - return + @property + def bathymetry_property(self): + """ + Read the bathymetry from disk, and print 'em + """ + + try: + bathy = xr.open_dataset(self.mom_input_dir / "bathymetry.nc") + return [bathy] + except: + return "No bathymetry set up yet (or files misplaced from {}). Call `setup_bathymetry` method to set up bathymetry.".format( + self.mom_input_dir + ) def setup_initial_condition( self, diff --git a/tests/test_grid_generation.py b/tests/test_grid_generation.py index d5158420..538cf2ad 100644 --- a/tests/test_grid_generation.py +++ b/tests/test_grid_generation.py @@ -2,7 +2,7 @@ import pytest from regional_mom6 import hyperbolictan_thickness_profile -from regional_mom6 import rectangular_hgrid +from regional_mom6 import calculate_rectangular_hgrid from regional_mom6 import longitude_slicer from regional_mom6.utils import angle_between @@ -129,7 +129,7 @@ def test_quadrilateral_areas(lat, lon, true_area): ], ) def test_rectangular_hgrid(lat, lon): - assert isinstance(rectangular_hgrid(lat, lon), xr.Dataset) + assert isinstance(calculate_rectangular_hgrid(lat, lon), xr.Dataset) def test_longitude_slicer(): diff --git a/tests/test_tides_functions_config.py b/tests/test_tides_functions_config.py index d7ec55f3..9bb9b37c 100644 --- a/tests/test_tides_functions_config.py +++ b/tests/test_tides_functions_config.py @@ -99,5 +99,19 @@ def test_initial_condition(self): ocean_varnames, arakawa_grid="A", ) - d1, d2, d3 = self.expt.initial_condition - print(d1, d2, d3) + dss = self.expt.initial_condition + print(dss) + + @pytest.mark.skipif( + IN_GITHUB_ACTIONS, reason="Test doesn't work in Github Actions." + ) + def test_properties(self): + """ + Test the properties + """ + dss = self.expt.era5 + dss_2 = self.expt.tides_boundaries + dss_3 = self.expt.ocean_state_boundaries + dss_4 = self.expt.initial_condition + dss_5 = self.expt.bathymetry_property + print(dss_2, dss_3, dss_4, dss_5)