diff --git a/fre/app/regrid_xy/regrid_xy.py b/fre/app/regrid_xy/regrid_xy.py index 3e261c2e1..b08b2dd2c 100644 --- a/fre/app/regrid_xy/regrid_xy.py +++ b/fre/app/regrid_xy/regrid_xy.py @@ -1,6 +1,7 @@ import logging import os from pathlib import Path +import pprint import subprocess import tarfile import xarray as xr @@ -10,6 +11,12 @@ fre_logger = logging.getLogger(__name__) +# it appears that fregrid believes the grid_spec file must have a datetime attached to it, like +# YYYYMMDD.grid_spec.tile +# this is behavior in noaa-gfdl::fre-nctools==2022.02.01 that can't be changed, and so must be worked around. +# hence, "ATTACH_LEGACY_DT" +ATTACH_LEGACY_DT = True + # list of variables/fields that will not be regridded non_regriddable_variables = [ "geolon_c", @@ -47,7 +54,6 @@ def get_grid_spec(datadict: dict) -> str: - """ Gets the grid_spec.nc file from the tar file specified in yaml["postprocess"]["settings"]["pp_grid_spec"] @@ -66,30 +72,61 @@ def get_grid_spec(datadict: dict) -> str: input mosaic filename """ - #grid spec filename - grid_spec = "grid_spec.nc" + grid_spec = str(Path("grid_spec.nc").resolve()) + fre_logger.info('grid spec filename is: %s', grid_spec) - #get tar file containing the grid_spec file pp_grid_spec_tar = datadict["yaml"]["postprocess"]["settings"]["pp_grid_spec"] - fre_logger.debug(f"Going to untar this grid spec tarfile: {pp_grid_spec_tar}") + fre_logger.info('grid spec tar archive file name is: %s', pp_grid_spec_tar) + - #untar grid_spec tar file into the current work directory + fre_logger.debug('checking if %s is a tar file...', pp_grid_spec_tar) if tarfile.is_tarfile(pp_grid_spec_tar): + + fre_logger.debug('it is a tar file! attempting top open %s', pp_grid_spec_tar) with tarfile.open(pp_grid_spec_tar, "r") as tar: + + fre_logger.debug('opened! about to extract all from opened tar file object into %s', os.getcwd()) tar.extractall() + fre_logger.debug('everything extracted!') + fre_logger.debug('contents extracted are ... %s', str(os.listdir(os.getcwd())) ) #error if grid_spec file is not found after extracting from tar file if not Path(grid_spec).exists(): raise IOError(f"Cannot find {grid_spec} in tar file {pp_grid_spec_tar}") fre_logger.debug(f"Current directory: {Path.cwd()}") + fre_logger.debug('grid_spec = %s exists!', grid_spec) - fre_logger.debug(f"Found grid_spec file: {grid_spec}") return grid_spec + grid_spec_dt_symlink = None + if ATTACH_LEGACY_DT: + fre_logger.warning('creating symlink to account for legacy fre-nctools 2022.02.01 behavior') + grid_spec_dt_symlink = Path(grid_spec).parent / f'{datadict["input_date"]}.grid_spec.nc' -def get_input_mosaic(datadict: dict) -> str: + fre_logger.warning('grid_spec_dt_symlink = %s', grid_spec_dt_symlink) + + # what in the ever loving demon magic is going on here??? + try: + grid_spec_dt_symlink.symlink_to(Path(grid_spec)) + except FileExistsError: + pass + + # i am sorry python gods, i have failed out + if grid_spec_dt_symlink.exists(): + fre_logger.warning('great? how did this happen?') + + # continued inexplicable success. + if grid_spec_dt_symlink.is_symlink(): + fre_logger.warning('symlink created: %s', grid_spec_dt_symlink) + else: + raise Exception('problem with accounting for legacy fregrid/fre-nctools behavior, symbolic link creation ' + 'as-intended-side-effect failed. consult life choices.') + return grid_spec, grid_spec_dt_symlink + + +def get_input_mosaic(datadict: dict) -> str: """ Gets the input mosaic filename from the grid_spec file. @@ -104,8 +141,10 @@ def get_input_mosaic(datadict: dict) -> str: .. note:: The input mosaic filename is a required input argument for fregrid. The input mosaic contains the input grid information. """ - grid_spec = datadict["grid_spec"] + fre_logger.info('grid_spec is: %s', grid_spec) + if not Path(grid_spec).exists(): + raise FileNotFoundError(f'grid_spec = {grid_spec} does not exist') #gridspec variable name holding the mosaic filename information match datadict["inputRealm"]: @@ -115,7 +154,13 @@ def get_input_mosaic(datadict: dict) -> str: #get mosaic filename with xr.open_dataset(grid_spec) as dataset: - mosaic_file = str(dataset[mosaic_key].data.astype(str)) + mosaic_file = str( + Path( + str( + dataset[mosaic_key].data.astype(str) + ) + ).resolve() + ) #check if the mosaic file exists in the current directory if not Path(mosaic_file).exists(): @@ -124,8 +169,9 @@ def get_input_mosaic(datadict: dict) -> str: return mosaic_file -def get_input_file(datadict: dict, source: str) -> str: - +def get_input_file(datadict: dict, + source: str, + input_dir: str) -> str: """ Formats the input file name where the input file contains the variable data that will be regridded. @@ -133,8 +179,9 @@ def get_input_file(datadict: dict, source: str) -> str: :type datadict: dict :param source: history file type :type source: str + :param input_dir: input file directory - :return: formatted input file name + :return: formatted full path to input file name :rtype: str .. note:: The input filename is a required argument for fregrid and refer to the history files containing @@ -150,13 +197,20 @@ def get_input_file(datadict: dict, source: str) -> str: will return "20250805.atmos_daily_cmip" (3) Fregrid will append the tile numbers ("tile1.nc") for reading in the data """ - + fre_logger.debug('attempting to read input_date key from datadict') input_date = datadict["input_date"] - return source if input_date is None else f"{input_date}.{source}" + if input_date is None: + fre_logger.debug('input_date is None, resolve source = %s and return that absolute path', source) + input_file = f"{input_dir}/{source}" + else: + fre_logger.debug('input_date = %s, returning absolute path based on that', input_date) + input_file = f"{input_dir}/{input_date}.{source}" + fre_logger.debug('returning %s', input_file) + return input_file -def get_remap_file(datadict: dict) -> str: +def get_remap_file(datadict: dict) -> str: """ Determines the remap filename based on the input mosaic filename, output grid size, and conservative order. For example, this function will return the name @@ -198,7 +252,6 @@ def get_remap_file(datadict: dict) -> str: def get_scalar_fields(datadict: dict) -> tuple[str, bool]: - """ Returns the scalar_fields argument for fregrid. Scalar_fields is a string of comma separated list of variables @@ -215,7 +268,7 @@ def get_scalar_fields(datadict: dict) -> tuple[str, bool]: will be regridded. """ - input_dir = Path(datadict["input_dir"]) + #input_dir = Path(datadict["input_dir"]) mosaic_file = datadict["input_mosaic"] input_file = datadict["input_file"] @@ -225,7 +278,8 @@ def get_scalar_fields(datadict: dict) -> tuple[str, bool]: # xarray gives an error if variables in non_regriddable_variables do not exist in the dataset # The errors="ignore" overrides the error - with xr.open_dataset(input_dir/input_file) as dataset: + #with xr.open_dataset(input_dir/input_file) as dataset: + with xr.open_dataset(input_file) as dataset: regrid_dataset = dataset.drop_vars(non_regriddable_variables, errors="ignore") if len(regrid_dataset) == 0: @@ -236,7 +290,6 @@ def get_scalar_fields(datadict: dict) -> tuple[str, bool]: def write_summary(datadict): - """ Logs a summary of the component that will be regridded in a human-readable format This function will log only if the logging level is set to INFO or lower @@ -265,9 +318,7 @@ def regrid_xy(yamlfile: str, work_dir: str, remap_dir: str, source: str, - input_date: str = None, -): - + input_date: str = None): """ Calls fregrid to regrid data in the specified source data file. @@ -287,86 +338,124 @@ def regrid_xy(yamlfile: str, :type source: str :param input_date: Datestring where the first 8 characters correspond to YYYYMMDD Input_date[:8] represents the date prefix in the history files, - e.g., input_date=20250730T0000Z where the history filename is + e.g., input_date=20250730T0000Z where the history filename is 20250730.atmos_month_aer.tile1.nc :type input_date: str .. note:: All directories should be in absolute paths """ + fre_logger.info(f"yamlfile = {yamlfile}") + fre_logger.info(f"input_dir = {input_dir}") + fre_logger.info(f"output_dir = {output_dir}") + fre_logger.info(f"work_dir = {work_dir}") + fre_logger.info(f"remap_dir = {remap_dir}") + fre_logger.info(f"source = {source}") + fre_logger.info(f"input_date = {input_date}") + - #check if input_dir exists + fre_logger.debug('checking if input_dir = %s exists', input_dir) if not Path(input_dir).exists(): raise RuntimeError(f"Input directory {input_dir} containing the input data files does not exist") - #check if output_dir exists + fre_logger.debug('checking if output_dir = %s exists', output_dir) if not Path(output_dir).exists(): raise RuntimeError(f"Output directory {output_dir} where regridded data" \ "will be outputted does not exist") - #check if work_dir exists + fre_logger.debug('checking if work_dir = %s exists', work_dir) if not Path(work_dir).exists(): raise RuntimeError(f"Specified working directory {work_dir} does not exist") - #work in working directory + fre_logger.debug('cd\'ing to work_dir = %s', work_dir) with helpers.change_directory(work_dir): - #initialize datadict + fre_logger.debug('initializing datadict') datadict = {} - # load yamlfile to yamldict + fre_logger.info( 'opening yamlfile = %s to read into yamldict', yamlfile ) with open(yamlfile, "r") as openedfile: yamldict = yaml.safe_load(openedfile) + fre_logger.debug( 'the yamldict is: \n%s', pprint.pformat(yamldict) ) - # save arguments to datadict + fre_logger.debug( 'saving yamldict fields to datadict' ) datadict["yaml"] = yamldict - datadict["grid_spec"] = get_grid_spec(datadict) datadict["input_dir"] = input_dir datadict["output_dir"] = output_dir datadict["work_dir"] = work_dir datadict["remap_dir"] = remap_dir datadict["input_date"] = input_date[:8] + fre_logger.info( 'datadict is %s', pprint.pformat(datadict) ) + + datadict["grid_spec"], grid_spec_symlink = get_grid_spec(datadict) + if Path(datadict['grid_spec']).exists(): + fre_logger.info('grid_spec exists here: %s', datadict['grid_spec']) + + + fre_logger.debug('making component list from yamldict') components = [] for component in yamldict["postprocess"]["components"]: for this_source in component["sources"]: if this_source["history_file"] == source: components.append(component) + fre_logger.info('components list is: %s', pprint.pformat(components) ) - # submit fregrid job for each component + fre_logger.debug('assembling fregrid call arguments for each component') for component in components: + fre_logger.debug('component = %s', component) - # skip component if postprocess_on = False + fre_logger.debug('checking postprocess_on field in component dict') if not component["postprocess_on"]: fre_logger.warning(f"postprocess_on=False for {source} in component {component['type']}." \ - "Skipping {source}") + f"Skipping {source}") continue + fre_logger.debug( 'saving component-specific info to datadict' ) datadict["inputRealm"] = component["inputRealm"] + + fre_logger.debug('calling get_input_mosaic...') datadict["input_mosaic"] = get_input_mosaic(datadict) + fre_logger.debug('result was %s', datadict["input_mosaic"]) + datadict["output_nlat"], datadict["output_nlon"] = component["xyInterp"].split(",") datadict["interp_method"] = component["interpMethod"] + + fre_logger.debug('calling get_remap_file') datadict["remap_file"] = get_remap_file(datadict) - datadict["input_file"] = get_input_file(datadict, source) + fre_logger.debug('result is %s', datadict["remap_file"]) + + fre_logger.debug('calling get_input_file') + datadict["input_file"] = get_input_file(datadict, source, input_dir) + fre_logger.debug('result is %s', datadict["input_file"]) + + fre_logger.debug('calling get_scalar_fields') datadict["scalar_field"], regrid = get_scalar_fields(datadict) + fre_logger.debug('result is %s', regrid) + + fre_logger.info( 'datadict is now %s', pprint.pformat(datadict) ) # skip if there are no variables to regrid if regrid: write_summary(datadict) else: + fre_logger.warning('no variables to regrid, skipping component') continue # create the output dir output_subdir = Path(output_dir) / f"{datadict['output_nlat']}_{datadict['output_nlon']}.{datadict['interp_method']}" + fre_logger.debug('creating output_subdir... %s', output_subdir) output_subdir.mkdir(parents=True, exist_ok=True) #construct fregrid command + fre_logger.debug('constructing fregrid command...') fregrid_command = [ "fregrid", "--debug", "--standard_dimension", "--input_dir", input_dir, + "--associated_file_dir", input_dir, "--input_mosaic", datadict["input_mosaic"], - "--input_file", datadict["input_file"], + "--input_file", datadict["input_file"].split('/')[-1], # no dir with the input file "--interp_method", datadict["interp_method"], "--remap_file", datadict["remap_file"], "--nlon", datadict["output_nlon"], @@ -375,13 +464,19 @@ def regrid_xy(yamlfile: str, "--output_dir", output_subdir, "--associated_file_dir", input_dir ] - fre_logger.debug(f"fregrid command: {fregrid_command}") + + fre_logger.info('the fregrid command is: \n %s', + ' '.join(fregrid_command).replace(' --', ' \\\n --') ) #execute fregrid command + fre_logger.debug('using subprocess.run to execute fregrid...') fregrid_job = subprocess.run(fregrid_command, capture_output=True, text=True) #print job useful information if fregrid_job.returncode == 0: fre_logger.info(fregrid_job.stdout.split("\n")[-3:]) else: + fre_logger.error('fregrid return code is nonzero!') + fre_logger.error('return code is %s', fregrid_job.returncode) + fre_logger.error('before error raise: stdout was %s', fregrid_job.stdout) raise RuntimeError(fregrid_job.stderr) diff --git a/fre/app/regrid_xy/tests/generate_files.py b/fre/app/regrid_xy/tests/generate_files.py index a39de7b12..e39fdcd38 100644 --- a/fre/app/regrid_xy/tests/generate_files.py +++ b/fre/app/regrid_xy/tests/generate_files.py @@ -1,7 +1,12 @@ -import numpy as np +''' +setup routines for test_regrid_xy +''' + from pathlib import Path import shutil import tarfile + +import numpy as np import yaml import xarray as xr @@ -19,25 +24,27 @@ tar_list: list = None def cleanup(): + ''' + remove prev. generate files as needed + ''' + if Path(yamlfile).exists(): + Path(yamlfile).unlink() - if Path(yamlfile).exists(): - Path(yamlfile).unlink() + if Path("grid_spec.nc").exists(): + Path("grid_spec.nc").unlink() - if Path("grid_spec.nc").exists(): - Path("grid_spec.nc").unlink() + if Path(grid_spec_tar).exists(): + Path(grid_spec_tar).unlink() - if Path(grid_spec_tar).exists(): - Path(grid_spec_tar).unlink() + if Path(input_mosaic).exists(): + Path(input_mosaic).unlink() - if Path(input_mosaic).exists(): - Path(input_mosaic).unlink() + if Path(input_dir).exists(): + shutil.rmtree(input_dir) - if Path(input_dir).exists(): - shutil.rmtree(input_dir) - - for i in range(1, ntiles+1): - gridfile = Path(f"{input_grid}.tile{i}.nc") - if gridfile.exists(): gridfile.unlink() + for i in range(1, ntiles+1): + gridfile = Path(f"{input_grid}.tile{i}.nc") + if gridfile.exists(): gridfile.unlink() def set_test(components_in: dict, @@ -50,119 +57,148 @@ def set_test(components_in: dict, input_grid_in: str = None, input_dir_in: str = None): - global components - global nxyp, nxy, ntiles, grid_spec_tar, input_grid - global date, input_mosaic - global input_dir, yamlfile - global tar_list - - components = components_in - if nxy_in is not None: - nxy = nxy_in - nxyp = nxy_in+1 - input_grid = f"C{nxy}" - if ntiles_in is not None: ntiles = ntiles_in - if date_in is not None: date = date_in - if yamlfile_in is not None: yamlfile = yamlfile_in - if grid_spec_tar_in is not None: grid_spec_tar = grid_spec_tar_in - if input_grid_in is not None: input_grid = input_grid_in - if input_mosaic_in is not None: input_mosaic = input_mosaic_in - if input_dir_in is not None: input_dir = input_dir_in - - tar_list = [] + global components + global nxyp, nxy, ntiles, grid_spec_tar, input_grid + global date, input_mosaic + global input_dir, yamlfile + global tar_list + + components = components_in + if nxy_in is not None: + nxy = nxy_in + nxyp = nxy_in+1 + input_grid = f"C{nxy}" + if ntiles_in is not None: + ntiles = ntiles_in + if date_in is not None: + date = date_in + if yamlfile_in is not None: + yamlfile = yamlfile_in + if grid_spec_tar_in is not None: + grid_spec_tar = grid_spec_tar_in + if input_grid_in is not None: + input_grid = input_grid_in + if input_mosaic_in is not None: + input_mosaic = input_mosaic_in + if input_dir_in is not None: + input_dir = input_dir_in + + tar_list = [] def make_yaml(): + ''' + makes a yaml + ''' - ppyaml = {} - ppyaml["name"] = yamlfile + ppyaml= {} + ppyaml["name"] = yamlfile - directories = ppyaml["directories"] = {} - directories["history_dir"] = "./" - directories["pp_dir"] = "./" + directories = ppyaml["directories"] = {} + directories["history_dir"] = "./" + directories["pp_dir"] = "./" - postprocess = ppyaml["postprocess"] = {} - postprocess["settings"] = {"pp_grid_spec": grid_spec_tar} - postprocess["components"] = components + postprocess = ppyaml["postprocess"] = {} + postprocess["settings"] = {"pp_grid_spec": grid_spec_tar} + postprocess["components"] = components - with open(yamlfile, "w") as openedfile: - yaml.dump(ppyaml, openedfile, sort_keys=False) + with open(yamlfile, "w") as openedfile: + yaml.dump(ppyaml, openedfile, sort_keys=False) def make_grid_spec(): - xr.Dataset(data_vars={"atm_mosaic_file": f"{input_mosaic}".encode(), - "lnd_mosaic_file": f"{input_mosaic}".encode(), - "ocn_mosaic_file": "ocean_mosaic.nc".encode()} - ).to_netcdf("grid_spec.nc") + ''' + makes a grid spec file via xarray dataset functionality + ''' + xr.Dataset( + data_vars = { "atm_mosaic_file" : f"{input_mosaic}".encode(), + "lnd_mosaic_file" : f"{input_mosaic}".encode(), + "ocn_mosaic_file" : "ocean_mosaic.nc".encode() + } ).to_netcdf("grid_spec.nc") - tar_list.append("grid_spec.nc") + tar_list.append("grid_spec.nc") def make_mosaic(): + ''' + makes a mosaic file + ''' - if ntiles > 1: - gridfiles = [f"{input_grid}.tile{i}.nc".encode() for i in range(1,ntiles+1)] - gridtiles = [f"tile{i}".encode() for i in range(1,ntiles+1)] - else: - gridfiles = f"{input_grid}.nc".encode() - gridtiles = f"tile1".encode() + if ntiles > 1: + gridfiles = [f"{input_grid}.tile{i}.nc".encode() for i in range(1,ntiles+1)] + gridtiles = [f"tile{i}".encode() for i in range(1,ntiles+1)] + else: + gridfiles = f"{input_grid}.nc".encode() + gridtiles = f"tile1".encode() - data = dict(gridfiles = xr.DataArray(gridfiles, dims=["ntiles"]).astype("|S255"), - gridtiles = xr.DataArray(gridtiles, dims=["ntiles"]).astype("|S255") - ) + data = dict(gridfiles = xr.DataArray(gridfiles, dims=["ntiles"]).astype("|S255"), + gridtiles = xr.DataArray(gridtiles, dims=["ntiles"]).astype("|S255") + ) - xr.Dataset(data_vars=data).to_netcdf(f"{input_mosaic}") + xr.Dataset(data_vars=data).to_netcdf(f"{input_mosaic}") - tar_list.append(f"{input_mosaic}") + tar_list.append(f"{input_mosaic}") def make_grid(): + ''' + makes a grid + ''' + + xy = np.arange(0, nxyp, 1, dtype=np.float64) + area = np.ones((nxy, nxy), dtype=np.float64) - xy = np.arange(0, nxyp, 1, dtype=np.float64) - area = np.ones((nxy, nxy), dtype=np.float64) + x, y = np.meshgrid(xy, xy) - x, y = np.meshgrid(xy, xy) + data = dict(x = xr.DataArray(x, dims=["nyp", "nxp"]), + y = xr.DataArray(y, dims=["nyp", "nxp"]), + area = xr.DataArray(area, dims=["ny", "nx"]) + ) - data = dict(x = xr.DataArray(x, dims=["nyp", "nxp"]), - y = xr.DataArray(y, dims=["nyp", "nxp"]), - area = xr.DataArray(area, dims=["ny", "nx"]) - ) + for i in range(1, ntiles+1): + data["tile"] = xr.DataArray(f"tile{i}".encode()).astype("|S255") + xr.Dataset(data).to_netcdf(f"{input_grid}.tile{i}.nc") - for i in range(1, ntiles+1): - data["tile"] = xr.DataArray(f"tile{i}".encode()).astype("|S255") - xr.Dataset(data).to_netcdf(f"{input_grid}.tile{i}.nc") + tar_list.append(f"{input_grid}.tile{i}.nc") - tar_list.append(f"{input_grid}.tile{i}.nc") def make_data(): + ''' + makes data + ''' - data = {} - data["mister"] = xr.DataArray(np.full((nxy,nxy), 1.0, dtype=np.float64), dims=["ny", "nx"]) - data["darcy"] = xr.DataArray(np.full((nxy,nxy), 2.0, dtype=np.float64), dims=["ny", "nx"]) - data["wins"] = xr.DataArray(np.full((nxy,nxy), 3.0, dtype=np.float64), dims=["ny", "nx"]) - data["wet_c"] = xr.DataArray(np.full((nxy,nxy), 5.0, dtype=np.float64), dims=["ny", "nx"]) + data = {} + data["mister"] = xr.DataArray(np.full((nxy,nxy), 1.0, dtype=np.float64), dims=["ny", "nx"]) + data["darcy"] = xr.DataArray(np.full((nxy,nxy), 2.0, dtype=np.float64), dims=["ny", "nx"]) + data["wins"] = xr.DataArray(np.full((nxy,nxy), 3.0, dtype=np.float64), dims=["ny", "nx"]) + data["wet_c"] = xr.DataArray(np.full((nxy,nxy), 5.0, dtype=np.float64), dims=["ny", "nx"]) - coords = {"nx": np.arange(1,nxyp, dtype=np.float64), - "ny": np.arange(1,nxyp, dtype=np.float64)} + coords = {"nx": np.arange(1,nxyp, dtype=np.float64), + "ny": np.arange(1,nxyp, dtype=np.float64)} - dataset = xr.Dataset(data_vars=data, coords=coords) + dataset = xr.Dataset(data_vars=data, coords=coords) - for component in components: - for source in component["sources"]: - history_file = source["history_file"] - for i in range(1, ntiles+1): - dataset.to_netcdf(f"{input_dir}/{date}.{history_file}.tile{i}.nc") + for component in components: + for source in component["sources"]: + history_file = source["history_file"] + for i in range(1, ntiles+1): + #print(f"input_dir is {input_dir}") # input dir is /home/Ian.Laflotte/Working/fre-cli/test_inputs + dataset.to_netcdf(f"{input_dir}/{date}.{history_file}.tile{i}.nc") def make_all(): - make_yaml() - make_grid_spec() - make_mosaic() - make_grid() - make_data() - - with tarfile.open(grid_spec_tar, "w") as tar: - for ifile in tar_list: tar.add(ifile) - - for ifile in tar_list: - Path(ifile).unlink() + ''' + makes all + ''' + make_yaml() + make_grid_spec() + make_mosaic() + make_grid() + make_data() + + with tarfile.open(grid_spec_tar, "w") as tar: + for ifile in tar_list: + tar.add(ifile) + + for ifile in tar_list: + Path(ifile).unlink() diff --git a/fre/app/regrid_xy/tests/test_regrid_xy.py b/fre/app/regrid_xy/tests/test_regrid_xy.py index af5831d3e..9ba19c21e 100644 --- a/fre/app/regrid_xy/tests/test_regrid_xy.py +++ b/fre/app/regrid_xy/tests/test_regrid_xy.py @@ -1,13 +1,13 @@ -import numpy as np import os from pathlib import Path import shutil + +import numpy as np import xarray as xr import fre.app.regrid_xy.regrid_xy as regrid_xy import fre.app.regrid_xy.tests.generate_files as generate_files - nxy = 20 date = "20250729" @@ -47,140 +47,142 @@ def setup_test(): - - input_dir.mkdir(exist_ok=True) - output_dir.mkdir(exist_ok=True) - remap_dir.mkdir(exist_ok=True) - work_dir.mkdir(exist_ok=True) - - #generate test files - generate_files.set_test(components_in=components, - date_in=date, - grid_spec_tar_in=str(grid_spec_tar), - yamlfile_in=str(yamlfile), - input_dir_in=str(input_dir)) - generate_files.make_all() + ''' + create directories and generate test files + ''' + input_dir.mkdir(exist_ok=True) + output_dir.mkdir(exist_ok=True) + remap_dir.mkdir(exist_ok=True) + work_dir.mkdir(exist_ok=True) + + #generate test files + generate_files.set_test(components_in=components, + date_in=date, + grid_spec_tar_in=str(grid_spec_tar), + yamlfile_in=str(yamlfile), + input_dir_in=str(input_dir)) + generate_files.make_all() def cleanup_test(): + ''' + remove test directories + ''' + if output_dir.exists(): shutil.rmtree(output_dir) + if remap_dir.exists(): shutil.rmtree(remap_dir) + if work_dir.exists(): shutil.rmtree(work_dir) + generate_files.cleanup() - #remove test directories - if output_dir.exists(): shutil.rmtree(output_dir) - if remap_dir.exists(): shutil.rmtree(remap_dir) - if work_dir.exists(): shutil.rmtree(work_dir) - generate_files.cleanup() - - -def test_regrid_xy(): - """ - Tests the main function regrid_xy and ensures - data is regridded correctly - """ - - setup_test() - - #modify generate_files to change sources - for source_dict in pp_input_files + emma_input_files + here_input_files: - source = source_dict["history_file"] - regrid_xy.regrid_xy(yamlfile=str(yamlfile), - input_dir=str(input_dir), - output_dir=str(output_dir), - work_dir=str(work_dir), - remap_dir=str(remap_dir), - source=source, - input_date=date+"TTTT") - - #check answers - for source_dict in pp_input_files + emma_input_files: - # Files are now output to a subdirectory based on grid size and interpolation method - output_subdir = output_dir/f"{nxy}_{nxy}.conserve_order2" - outfile = output_subdir/f"{date}.{source_dict['history_file']}.nc" - - test = xr.load_dataset(outfile) - - assert "wet_c" not in test - assert "mister" in test - assert "darcy" in test - assert "wins" in test - - assert np.all(test["mister"].values==np.float64(1.0)) - assert np.all(test["darcy"].values==np.float64(2.0)) - assert np.all(test["wins"].values==np.float64(3.0)) - - #check answers, these shouldn't have been regridded - for source_dict in here_input_files: - ifile = source_dict["history_file"] - assert not (output_dir/f"{date}.{ifile}.nc").exists() - - #check remap_file exists and is not empty - remap_file = remap_dir/f"C{nxy}_mosaicX{nxy}by{nxy}_conserve_order2.nc" - assert remap_file.exists() - - cleanup_test() - - def test_get_input_mosaic(): + """ + Tests get_input_mosaic correctly copies the mosaic file to the input directory + """ + setup_test() - """ - Tests get_input_mosaic correctly copies the mosaic file to the input directory - """ + grid_spec = Path("grid_spec.nc").resolve() + mosaic_file = Path("ocean_mosaic.nc").resolve() - grid_spec = Path("grid_spec.nc") - mosaic_file = Path("ocean_mosaic.nc") + generate_files.make_grid_spec() + mosaic_file.touch() - generate_files.make_grid_spec() - mosaic_file.touch() + datadict=dict(grid_spec=grid_spec, inputRealm="ocean") - datadict=dict(grid_spec=grid_spec, inputRealm="ocean") + assert regrid_xy.get_input_mosaic(datadict) == str(mosaic_file) - assert regrid_xy.get_input_mosaic(datadict) == str(mosaic_file) - - mosaic_file.unlink() #clean up - grid_spec.unlink() #clean up + mosaic_file.unlink() #clean up + grid_spec.unlink() #clean up + cleanup_test() def test_get_input_file(): + """ + Tests get_input_file + """ - """ - Tests get_input_file - """ - - input_date = "20250807" - source = "pemberley" - datadict = {"input_date": input_date} - assert regrid_xy.get_input_file(datadict, source) == input_date + "." + source + input_dir = os.getcwd() + input_date = "20250807" + source = "pemberley" + datadict = {"input_date": input_date} + expected_answer_w_date = f"{input_dir}/{input_date}.{source}" + assert regrid_xy.get_input_file(datadict, source, input_dir) == expected_answer_w_date - datadict["input_date"] = None - assert regrid_xy.get_input_file(datadict, source) == source + datadict["input_date"] = None + expected_answer_no_date = f"{input_dir}/{source}" + assert regrid_xy.get_input_file(datadict, source, input_dir) == expected_answer_no_date def test_get_remap_file(): + """ + Tests get_remap_file + """ + + remap_dir = Path("remap_dir") + input_mosaic = "C20_mosaic" + nlon = 40 + nlat = 10 + interp_method = "conserve_order1" + + datadict = {"remap_dir": remap_dir.name, + "input_mosaic": input_mosaic+".nc", + "output_nlon": nlon, + "output_nlat": nlat, + "interp_method": interp_method} + + #check remap file from current directory is copied to input directory + remap_file = Path(f"remap_dir/{input_mosaic}X{nlon}by{nlat}_{interp_method}.nc") + assert regrid_xy.get_remap_file(datadict) == str(remap_dir/remap_file) + + remap_dir.mkdir(exist_ok=True) + remap_file.touch() + assert regrid_xy.get_remap_file(datadict) == str(remap_dir/remap_file) + + Path(remap_file).unlink() + shutil.rmtree(remap_dir) - """ - Tests get_remap_file - """ - - remap_dir = Path("remap_dir") - input_mosaic = "C20_mosaic" - nlon = 40 - nlat = 10 - interp_method = "conserve_order1" - - datadict = {"remap_dir": remap_dir.name, - "input_mosaic": input_mosaic+".nc", - "output_nlon": nlon, - "output_nlat": nlat, - "interp_method": interp_method} - #check remap file from current directory is copied to input directory - remap_file = Path(f"remap_dir/{input_mosaic}X{nlon}by{nlat}_{interp_method}.nc") - - regrid_xy.get_remap_file(datadict) == str(remap_dir/remap_file) - - remap_dir.mkdir(exist_ok=True) - remap_file.touch() - regrid_xy.get_remap_file(datadict) == str(remap_dir/remap_file) - - Path(remap_file).unlink() - shutil.rmtree(remap_dir) +def test_regrid_xy(): + """ + Tests the main function regrid_xy and ensures + data is regridded correctly + """ + + setup_test() + # assert False + + #modify generate_files to change sources + for source_dict in pp_input_files + emma_input_files + here_input_files: + source = source_dict["history_file"] + regrid_xy.regrid_xy(yamlfile=str(yamlfile), + input_dir=str(input_dir), + output_dir=str(output_dir), + work_dir=str(work_dir), + remap_dir=str(remap_dir), + source=source, + input_date=date+"TTTT") + + #check answers + for source_dict in pp_input_files + emma_input_files: + outfile = output_dir/f"{date}.{source_dict['history_file']}.nc" + + test = xr.load_dataset(outfile) + + assert "wet_c" not in test + assert "mister" in test + assert "darcy" in test + assert "wins" in test + + assert np.all(test["mister"].values==np.float64(1.0)) + assert np.all(test["darcy"].values==np.float64(2.0)) + assert np.all(test["wins"].values==np.float64(3.0)) + + #check answers, these shouldn't have been regridded + for source_dict in here_input_files: + ifile = source_dict["history_file"] + assert not (output_dir/f"{date}.{ifile}.nc").exists() + + #check remap_file exists and is not empty + remap_file = remap_dir/f"C{nxy}_mosaicX{nxy}by{nxy}_conserve_order2.nc" + assert remap_file.exists() + + cleanup_test()