diff --git a/regional_mom6/regional_mom6.py b/regional_mom6/regional_mom6.py index 7c93c37c..5a632816 100644 --- a/regional_mom6/regional_mom6.py +++ b/regional_mom6/regional_mom6.py @@ -22,10 +22,8 @@ import copy from . import regridding as rgd from . import rotation as rot -from .utils import quadrilateral_areas, ap2ep, ep2ap, is_rectilinear_hgrid, setup_logger -import logging +from .utils import quadrilateral_areas, ap2ep, ep2ap, is_rectilinear_hgrid -rm6_logger = setup_logger(__name__, set_handler=True) warnings.filterwarnings("ignore") @@ -100,14 +98,14 @@ def create_experiment_from_config( Returns: experiment: An experiment object with the fields from the config loaded in. """ - rm6_logger.info("Reading from config file....") + print("Reading from config file....") with open(config_file_path, "r") as f: config_dict = json.load(f) - rm6_logger.info("Creating Empty Experiment Object....") + print("Creating Empty Experiment Object....") expt = experiment.create_empty() - rm6_logger.info("Setting Default Variables.....") + print("Setting Default Variables.....") expt.expt_name = config_dict["expt_name"] try: expt.longitude_extent = tuple(config_dict["longitude_extent"]) @@ -148,13 +146,13 @@ def create_experiment_from_config( expt.boundaries = config_dict["boundaries"] if create_hgrid_and_vgrid: - rm6_logger.info("Creating hgrid and vgrid....") + print("Creating hgrid and vgrid....") expt.hgrid = expt._make_hgrid() expt.vgrid = expt._make_vgrid() else: - rm6_logger.info("Skipping hgrid and vgrid creation....") + print("Skipping hgrid and vgrid creation....") - rm6_logger.info("Done!") + print("Done!") return expt @@ -716,7 +714,7 @@ def __init__( float(self.hgrid.y.max()), ) except: - rm6_logger.error( + print( "Error while reading in existing horizontal grid!\n\n" + f"Make sure `hgrid.nc`exists in {self.mom_input_dir} directory." ) @@ -740,7 +738,7 @@ def __init__( vgrid_from_file = xr.open_dataset(vgrid_path) except: - rm6_logger.error( + print( "Error while reading in existing vertical coordinates!\n\n" + f"Make sure `vcoord.nc`exists in {self.mom_input_dir} directory." ) @@ -782,7 +780,7 @@ def __getattr__(self, name): decode_times=False, ) else: - rm6_logger.error( + print( f"bathymetry.nc file not found! Make sure you've successfully run the setup_bathmetry method, or copied your own bathymetry.nc file into {self.mom_input_dir}." ) return None @@ -794,7 +792,7 @@ def __getattr__(self, name): decode_times=False, ) else: - rm6_logger.error( + print( f"init_vel.nc file not found! Make sure you've successfully run the setup_initial_condition method, or copied your own init_vel.nc file into {self.mom_input_dir}." ) return @@ -807,7 +805,7 @@ def __getattr__(self, name): decode_times=False, ) else: - rm6_logger.error( + print( f"init_tracers.nc file not found! Make sure you've successfully run the setup_initial_condition method, or copied your own init_tracers.nc file into {self.mom_input_dir}." ) return @@ -820,7 +818,7 @@ def __getattr__(self, name): decode_cf=False, ) except: - rm6_logger.error( + print( f"{name} files not found! Make sure you've successfully run the setup_ocean_state_boundaries method, or copied your own segment files file into {self.mom_input_dir}." ) return None @@ -953,7 +951,7 @@ def _make_vgrid(self, thicknesses=None): ## Check whether the minimum depth is less than the first three layers if self.minimum_depth < zi[2]: - rm6_logger.warning( + print( f"Warning: Minimum depth of {self.minimum_depth}m is less than the depth of the third interface ({zi[2]}m)!\n" + "This means that some areas may only have one or two layers between the surface and sea floor. \n" + "For increased stability, consider increasing the minimum depth, or adjusting the vertical coordinate to add more layers near the surface." @@ -1092,7 +1090,7 @@ def write_config_file(self, path=None, export=True, quiet=False): Dict: A dictionary containing the configuration information. """ if not quiet: - rm6_logger.info("Writing Config File.....") + print("Writing Config File.....") try: date_range = [ self.date_range[0].strftime("%Y-%m-%d %H:%M:%S"), @@ -1129,7 +1127,7 @@ def write_config_file(self, path=None, export=True, quiet=False): indent=4, ) if not quiet: - rm6_logger.info("Done.") + print("Done.") return config_dict def setup_initial_condition( @@ -1338,11 +1336,11 @@ def setup_initial_condition( # vgrid = rgd.get_hgrid_arakawa_c_points(self.hgrid, "v").rename({"vlon": "lon", "vlat": "lat"}).set_coords(["lat", "lon"]) ## Construct the cell centre grid for tracers (xh, yh). - rm6_logger.info("Setting up Initial Conditions") + print("Setting up Initial Conditions") ## Regrid all fields horizontally. - rm6_logger.info("Regridding Velocities... ") + print("Regridding Velocities... ", end="") regridded_u = regridder_u(ic_raw_u) regridded_v = regridder_v(ic_raw_v) if rotational_method == rot.RotationMethod.GIVEN_ANGLE: @@ -1385,7 +1383,7 @@ def setup_initial_condition( ] ) - rm6_logger.info("Done.\nRegridding Tracers... ") + print("Done.\nRegridding Tracers... ", end="") tracers_out = ( xr.merge( @@ -1409,7 +1407,7 @@ def setup_initial_condition( } ) - rm6_logger.info("Done.\nRegridding Free surface... ") + print("Done.\nRegridding Free surface... ", end="") eta_out = ( regridder_t(ic_raw_eta) @@ -1417,7 +1415,7 @@ def setup_initial_condition( .rename("eta_t") .transpose("ny", "nx") ) ## eta_t is the name set in MOM_input by default - rm6_logger.info("Done.") + print("Done.") ## Return attributes to arrays @@ -1449,7 +1447,7 @@ def setup_initial_condition( tracers_out = tracers_out.interp({"zl": self.vgrid.zl.values}) vel_out = vel_out.interp({"zl": self.vgrid.zl.values}) - rm6_logger.info("Saving outputs... ") + print("Saving outputs... ", end="") vel_out.fillna(0).to_netcdf( self.mom_input_dir / "init_vel.nc", @@ -1485,7 +1483,7 @@ def setup_initial_condition( self.ic_tracers = tracers_out self.ic_vels = vel_out - rm6_logger.info("done setting up initial condition.") + print("done setting up initial condition.") return @@ -1618,7 +1616,7 @@ def setup_ocean_state_boundaries( ) if len(self.boundaries) < 4: - rm6_logger.warning( + print( "NOTE: the 'setup_run_directories' method does understand the less than four boundaries but be careful. Please check the MOM_input/override file carefully to reflect the number of boundaries you have, and their orientations. You should be able to find the relevant section in the MOM_input/override file by searching for 'segment_'. Ensure that the segment names match those in your inputdir/forcing folder" ) @@ -1678,7 +1676,7 @@ def setup_single_boundary( rotational_method (Optional[str]): Method to use for rotating the boundary velocities. Default is 'GIVEN_ANGLE'. """ - rm6_logger.info("Processing {} boundary...".format(orientation)) + print("Processing {} boundary...".format(orientation), end="") if not path_to_bc.exists(): raise FileNotFoundError( f"Boundary file not found at {path_to_bc}. Please ensure that the files are named in the format `east_unprocessed.nc`." @@ -1702,7 +1700,7 @@ def setup_single_boundary( rotational_method=rotational_method ) - rm6_logger.info("Done.") + print("Done.") return def setup_boundary_tides( @@ -1789,7 +1787,7 @@ def setup_boundary_tides( ) # Initialize or find boundary segment for b in self.boundaries: - rm6_logger.info("Processing {} boundary...".format(b)) + print("Processing {} boundary...".format(b), end="") # If the GLORYS ocean_state has already created segments, we don't create them again. seg = segment( @@ -1810,7 +1808,7 @@ def setup_boundary_tides( seg.regrid_tides( tpxo_v, tpxo_u, tpxo_h, times, rotational_method=rotational_method ) - rm6_logger.info("Done") + print("Done") def setup_bathymetry( self, @@ -1994,7 +1992,7 @@ def setup_bathymetry( ) self.tidy_bathymetry(fill_channels, positive_down, bathymetry=bathymetry) - rm6_logger.info("setup bathymetry has finished successfully.") + print("setup bathymetry has finished successfully.") return bathymetry @@ -2027,8 +2025,8 @@ def tidy_bathymetry( """ ## reopen bathymetry to modify - rm6_logger.info( - "Tidy bathymetry: Reading in regridded bathymetry to fix up metadata..." + print( + "Tidy bathymetry: Reading in regridded bathymetry to fix up metadata...", end="" ) if read_bathy_from_file := bathymetry is None: bathymetry = xr.open_dataset( @@ -2054,7 +2052,7 @@ def tidy_bathymetry( land_mask = np.abs(ocean_mask - 1) ## REMOVE INLAND LAKES - rm6_logger.info("done. Filling in inland lakes and channels... ") + print("done. Filling in inland lakes and channels... ", end="") changed = True ## keeps track of whether solution has converged or not @@ -2202,7 +2200,7 @@ def tidy_bathymetry( encoding={"depth": {"_FillValue": None}}, ) - rm6_logger.info("done.") + print("done.") return def run_FRE_tools(self, layout=None): @@ -2210,11 +2208,11 @@ def run_FRE_tools(self, layout=None): User provides processor ``layout`` tuple of processing units. """ - rm6_logger.info( + print( "Running GFDL's FRE Tools. The following information is all printed by the FRE tools themselves" ) if not (self.mom_input_dir / "bathymetry.nc").exists(): - rm6_logger.error( + print( "No bathymetry file! Need to run setup_bathymetry method first" ) return @@ -2222,7 +2220,7 @@ def run_FRE_tools(self, layout=None): for p in self.mom_input_dir.glob("mask_table*"): p.unlink() - rm6_logger.info( + print( "OUTPUT FROM MAKE SOLO MOSAIC:", subprocess.run( str(self.toolpath_dir / "make_solo_mosaic/make_solo_mosaic") @@ -2233,7 +2231,7 @@ def run_FRE_tools(self, layout=None): sep="\n\n", ) - rm6_logger.info( + print( "OUTPUT FROM QUICK MOSAIC:", subprocess.run( str(self.toolpath_dir / "make_quick_mosaic/make_quick_mosaic") @@ -2253,7 +2251,7 @@ def configure_cpu_layout(self, layout): ``layout`` tuple of processing units. """ - rm6_logger.info( + print( "OUTPUT FROM CHECK MASK:\n\n", subprocess.run( str(self.toolpath_dir / "check_mask/check_mask") @@ -2295,11 +2293,11 @@ def setup_run_directory( ) if not premade_rundir_path.exists(): - rm6_logger.info( + print( "Could not find premade run directories at ", premade_rundir_path ) - rm6_logger.info( - "Perhaps the package was imported directly rather than installed with conda. Checking if this is the case... " + print( + "Perhaps the package was imported directly rather than installed with conda. Checking if this is the case... ", end="" ) premade_rundir_path = Path( @@ -2313,7 +2311,7 @@ def setup_run_directory( + "There may be an issue with package installation. Check that the `premade_run_directory` folder is present in one of these two locations" ) else: - rm6_logger.info("Found run files. Continuing...") + print("Found run files. Continuing...") # Define the locations of the directories we'll copy files across from. Base contains most of the files, and overwrite replaces files in the base directory. base_run_dir = Path(premade_rundir_path / "common_files") @@ -2385,7 +2383,7 @@ def setup_run_directory( mask_table = None for p in self.mom_input_dir.glob("mask_table.*"): if mask_table != None: - rm6_logger.warning( + print( f"WARNING: Multiple mask tables found. Defaulting to {mask_table}. If this is not what you want, remove it from the run directory and try again." ) break @@ -2399,7 +2397,7 @@ def setup_run_directory( y, ) # This is a local variable keeping track of the layout as read from the mask table. Not to be confused with self.layout which is unchanged and may differ. - rm6_logger.info( + print( f"Mask table {p.name} read. Using this to infer the cpu layout {layout}, total masked out cells {masked}, and total number of CPUs {ncpus}." ) # Case where there's no mask table. Either because user hasn't run FRE tools, or because the domain is mostly water. @@ -2409,11 +2407,11 @@ def setup_run_directory( # in case the user accidentally loads in the wrong mask table. layout = self.layout if layout == None: - rm6_logger.warning( + print( "WARNING: No mask table found, and the cpu layout has not been set. \nAt least one of these is requiret to set up the experiment if you're running MOM6 standalone with the FMS coupler. \nIf you're running within CESM, ignore this message." ) else: - rm6_logger.warning( + print( f"No mask table found, but the cpu layout has been set to {self.layout} This suggests the domain is mostly water, so there are " + "no `non compute` cells that are entirely land. If this doesn't seem right, " + "ensure you've already run the `FRE_tools` method which sets up the cpu mask table. Keep an eye on any errors that might print while" @@ -2456,7 +2454,7 @@ def setup_run_directory( # OBC Adjustments # Delete MOM_input OBC stuff that is indexed because we want them only in MOM_override. - rm6_logger.info( + print( "Deleting indexed OBC keys from MOM_input_dict in case we have a different number of segments" ) keys_to_delete = [key for key in MOM_input_dict if "_SEGMENT_00" in key] @@ -2556,7 +2554,7 @@ def setup_run_directory( if not using_payu and os.path.exists(f"{self.mom_run_dir}/config.yaml"): os.remove(f"{self.mom_run_dir}/config.yaml") elif ncpus == None: - rm6_logger.warning( + print( "WARNING: Layout has not been set! Cannot create payu configuration file. Run the FRE_tools first." ) else: @@ -2634,7 +2632,7 @@ def change_MOM_parameter( if param_name in MOM_override_dict.keys(): original_val = MOM_override_dict[param_name]["value"] - rm6_logger.info( + print( "This parameter {} is being replaced from {} to {} in MOM_override".format( param_name, original_val, param_value ) @@ -2646,12 +2644,12 @@ def change_MOM_parameter( else: if param_name in MOM_override_dict.keys(): original_val = MOM_override_dict[param_name]["value"] - rm6_logger.info( + print( "Deleting parameter {} from MOM_override".format(param_name) ) del MOM_override_dict[param_name] else: - rm6_logger.info( + print( "Key to be deleted {} was not in MOM_override to begin with.".format( param_name ) @@ -2720,7 +2718,7 @@ def write_MOM_file(self, MOM_file_dict): # As in there wasn't an override before but we want one if MOM_file_dict[var]["override"]: lines[jj] = "#override " + lines[jj] - rm6_logger.info("Added override to variable " + var + "!") + print("Added override to variable " + var + "!") if var in MOM_file_dict.keys() and ( str(MOM_file_dict[var]["value"]) ) != str(original_MOM_file_dict[var]["value"]): @@ -2741,7 +2739,7 @@ def write_MOM_file(self, MOM_file_dict): + "\n" ) - rm6_logger.info( + print( "Changed " + str(var) + " from " @@ -2763,7 +2761,7 @@ def write_MOM_file(self, MOM_file_dict): lines.append( f"{key} = {MOM_file_dict[key]['value']} !{MOM_file_dict[key]['comment']}\n" ) - rm6_logger.info( + print( "Added", key, "to", @@ -2785,7 +2783,7 @@ def write_MOM_file(self, MOM_file_dict): for line in lines if not all(word in line for word in search_words) ] - rm6_logger.info( + print( "Removed", key, "in",