Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Added functionality to Split #54

Merged
merged 4 commits into from
Dec 14, 2024
Merged

Added functionality to Split #54

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
ab7cfa4
fixed a couple PEP8 errors in MarsVars
falconstryker Dec 10, 2024
2f0cbf2
fixed a couple PEP8 errors in MarsVars
falconstryker Dec 10, 2024
c3a461a
Adding functionality to split that allows selection of just one
falconstryker Dec 13, 2024
bc6addf
updated help function to describe new split functionality
falconstryker Dec 14, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
111 changes: 70 additions & 41 deletions bin/MarsFiles.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -120,20 +120,25 @@

parser.add_argument("-split", "--split", nargs="+",
help=(
f"Extract values between min and max solar longitudes 0-360 [°]\n"
f"This assumes all values in the file are from only one Mars Year.\n"
f"Extract a range of values along a dimension. Defaults to Ls, unless "
f"otherwise specified using --dim. If the file contains multiple Mars "
f"Years of data, this function splits the file according to the Ls "
f"values from the first Mars Year.\n"
f"{Yellow}Use [-dim, --dim] to specify the dimension (see below).\n"
f"{Green}Usage:\n"
f"> MarsFiles.py 00668.atmos_average.nc --split 0 90 \n"
f"> MarsFiles.py 00668.atmos_average.nc --split 0 90"
f"> MarsFiles.py 00668.atmos_average.nc --split 270"
f"{Nclr}\n\n"
)
)

parser.add_argument("-dim", "--dim", type=str, default = 'areo',
help=(
f"Flag to specify dimension to split on. Acceptable values are \n"
f"Flag to specify the dimension to split. Acceptable values are \n"
f"areo, lat, lon, lev. For use with --split.\n"
f"{Green}Usage:\n"
f"> MarsFiles.py 00668.atmos_average.nc --split 0 90 --dim areo"
f"> MarsFiles.py 00668.atmos_average.nc --split -70 --dim lat"
f"{Nclr}\n\n"
)
)
Expand Down Expand Up @@ -456,8 +461,11 @@ def split_files(file_list, split_dim):
split_dim = 'time'
bounds = np.asarray(parser.parse_args().split).astype(float)

if len(np.atleast_1d(bounds)) != 2:
print(f"{Red}Requires two values: [lower_bound] [upper_bound]{Nclr}")
if len(np.atleast_1d(bounds)) > 2 or len(np.atleast_1d(bounds)) < 1:
print(f"{Red}Accepts only ONE or TWO values:"
f"[bound] to reduce one dimension to a single value"
f"[lower_bound] [upper_bound] to reduce one dimension to "
f"a range{Nclr}")
exit()

# Add path unless full path is provided
Expand All @@ -470,50 +478,71 @@ def split_files(file_list, split_dim):
fNcdf = Dataset(input_file_name, 'r', format = 'NETCDF4_CLASSIC')
var_list = filter_vars(fNcdf, parser.parse_args().include)

dim_var = fNcdf.variables[split_dim][:]
# reducing_dim = fNcdf.variables[split_dim][:]

# Get file type (diurn, average, daily, etc.)
f_type, _ = FV3_file_type(fNcdf)

# Remove all single dimensions from areo (scalar_axis)
if f_type == 'diurn':
if split_dim == 'time':
# size = areo (133, 24, 1)
bounds_limit = np.squeeze(fNcdf.variables['areo'][:, 0, :]) % 360
# size areo = (time, tod, scalar_axis)
reducing_dim = np.squeeze(fNcdf.variables['areo'][:, 0, :]) % 360
else:
bounds_limit = np.squeeze(fNcdf.variables[split_dim][:, 0])
reducing_dim = np.squeeze(fNcdf.variables[split_dim][:, 0])
else:
if split_dim == 'time':
# size = areo (133, 1)
bounds_limit = np.squeeze(fNcdf.variables['areo'][:]) % 360
# size areo = (time, scalar_axis)
reducing_dim = np.squeeze(fNcdf.variables['areo'][:]) % 360
else:
bounds_limit = np.squeeze(fNcdf.variables[split_dim][:])

lower_bound = np.argmin(np.abs(bounds[0] - bounds_limit))
upper_bound = np.argmin(np.abs(bounds[1] - bounds_limit))

if lower_bound == upper_bound:
print(f"{Red}Warning, requested {split_dim} min, max ({bounds[0]}, "
f"{bounds[1]}) are out of file range ({split_dim} = "
f"{bounds_limit[0]:.1f}-{bounds_limit[-1]:.1f})")
exit()

dim_out = dim_var[lower_bound:upper_bound]
print(f"{Yellow}dim_var = {dim_var}")
print(f"{Yellow}dim_out = {dim_out}")
reducing_dim = np.squeeze(fNcdf.variables[split_dim][:])

print(f"\n{Yellow}All values in dimension:\n{reducing_dim}\n")
if len(np.atleast_1d(bounds)) < 2:
indices = [(np.abs(reducing_dim - bounds[0])).argmin()]
dim_out = reducing_dim[indices]
print(f"Requested value = {bounds[0]}\n"
f"Nearest value = {dim_out[0]}\n")
else:
indices = np.where((reducing_dim >= bounds[0]) & (reducing_dim <= bounds[1]))[0]
dim_out = reducing_dim[indices]
print(f"Requested range = {bounds[0]} - {bounds[1]}\n"
f"Corresponding values = {dim_out}\n")
if len(indices) == 0:
print(f"{Red}Warning, no values were found in the range {split_dim} "
f"{bounds[0]}, {bounds[1]}) ({split_dim} values range from "
f"{reducing_dim[0]:.1f} to {reducing_dim[-1]:.1f})")
exit()

if split_dim == 'time':
time_dim = (np.squeeze(fNcdf.variables['time'][:]))[indices]
print(f"time_dim = {time_dim}")

fpath, fname = extract_path_basename(input_file_name)
if split_dim == 'time':
output_file_name = (f"{fpath}/{int(dim_out[0]):05d}{fname[5:-3]}_"
f"Ls{int(bounds[0]):03d}_{int(bounds[1]):03d}.nc")
if len(np.atleast_1d(bounds)) < 2:
output_file_name = (f"{fpath}/{int(time_dim):05d}{fname[5:-3]}_"
f"nearest_Ls{int(bounds[0]):03d}.nc")
else:
output_file_name = (f"{fpath}/{int(time_dim[0]):05d}{fname[5:-3]}_"
f"Ls{int(bounds[0]):03d}_{int(bounds[1]):03d}.nc")
elif split_dim == 'lat':
new_bounds = [str(abs(int(b)))+"S" if b < 0 else str(int(b))+"N" for b in bounds]
print(f"{Yellow}bounds = {bounds[0]} {bounds[1]}")
print(f"{Yellow}new_bounds = {new_bounds[0]} {new_bounds[1]}")
output_file_name = (f"{fpath}/{original_date}{fname[5:-3]}_{split_dim}"
f"{new_bounds[0]}_{new_bounds[1]}.nc")
if len(np.atleast_1d(bounds)) < 2:
output_file_name = (f"{fpath}/{original_date}{fname[5:-3]}_{split_dim}"
f"nearest_{new_bounds[0]}.nc")
else:
print(f"{Yellow}bounds = {bounds[0]} {bounds[1]}")
print(f"{Yellow}new_bounds = {new_bounds[0]} {new_bounds[1]}")
output_file_name = (f"{fpath}/{original_date}{fname[5:-3]}_{split_dim}"
f"{new_bounds[0]}_{new_bounds[1]}.nc")
else:
output_file_name = (f"{fpath}/{original_date}{fname[5:-3]}_{split_dim}"
f"{int(bounds[0]):03d}_{int(bounds[1]):03d}.nc")
if len(np.atleast_1d(bounds)) < 2:
output_file_name = (f"{fpath}/{original_date}{fname[5:-3]}_{split_dim}"
f"nearest_{int(bounds[0]):03d}.nc")
else:
output_file_name = (f"{fpath}/{original_date}{fname[5:-3]}_{split_dim}"
f"{int(bounds[0]):03d}_{int(bounds[1]):03d}.nc")

print(f"{Cyan}new filename = {output_file_name}")
Log = Ncdf(output_file_name)
Expand Down Expand Up @@ -554,19 +583,19 @@ def split_files(file_list, split_dim):
# ivar is a dim of ivar but ivar is not ivar
print(f'{Cyan}Processing: {ivar}...{Nclr}')
if split_dim == 'time':
var_out = varNcf[lower_bound:upper_bound, ...]
var_out = varNcf[indices, ...]
elif split_dim == 'lat' and varNcf.ndim == 5:
var_out = varNcf[:, :, :, lower_bound:upper_bound, :]
var_out = varNcf[:, :, :, indices, :]
elif split_dim == 'lat' and varNcf.ndim == 4:
var_out = varNcf[:, :, lower_bound:upper_bound, :]
var_out = varNcf[:, :, indices, :]
elif split_dim == 'lat' and varNcf.ndim == 3:
var_out = varNcf[:, lower_bound:upper_bound, :]
var_out = varNcf[:, indices, :]
elif split_dim == 'lat' and varNcf.ndim == 2:
var_out = varNcf[lower_bound:upper_bound, ...]
var_out = varNcf[indices, ...]
elif split_dim == 'lon' and varNcf.ndim > 2:
var_out = varNcf[..., lower_bound:upper_bound]
var_out = varNcf[..., indices]
elif split_dim == 'lon' and varNcf.ndim == 2:
var_out = varNcf[lower_bound:upper_bound, ...]
var_out = varNcf[indices, ...]
longname_txt, units_txt = get_longname_units(fNcdf, ivar)
Log.log_variable(ivar, var_out, varNcf.dimensions,
longname_txt, units_txt)
Expand Down
40 changes: 34 additions & 6 deletions bin/MarsVars.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -332,13 +332,41 @@
# ===========================

def err_req_interpolated_file(ivar, ifile):
"""
Print message to the screen when user tries to add a variable to
an incompatible file type.

:param ivar: The requested variable to add
:type ivar: array [time, lev, lat, lon]
:param ifile: The file to add the variable into
:type ifile: str

:raises:

:return: print statement
:rtype: str
"""
return(
print(f"{Red}ERROR: variable {ivar} can only be added to a "
f"pressure-interpolated file.\nRun {Yellow}'MarsInterp.py "
f"{ifile} -t pstd' {Red}before trying again.{Nclr}")
)

def err_req_non_interpolated_file(ivar, ifile):
"""
Print message to the screen when user tries to add a variable to
an incompatible file type.

:param ivar: The requested variable to add
:type ivar: array [time, lev, lat, lon]
:param ifile: The file to add the variable into
:type ifile: str

:raises:

:return: print statement
:rtype: str
"""
return(
print(f"{Red}ERROR: variable {ivar} cannot be added to {ifile} "
f"as it is an interpolated file.\n Please add the "
Expand Down Expand Up @@ -1314,25 +1342,25 @@ def main():
if ivar == "div":
if interp_type not in ("pstd", "zstd", "zagl"):
OUT = spherical_div(ucomp, vcomp, lon, lat,
R = 3400*1000.,
spacing = "regular")
R=3400*1000.,
spacing="regular")
else:
err_req_non_interpolated_file(ivar, ifile)

if ivar == "curl":
if interp_type not in ("pstd", "zstd", "zagl"):
OUT = spherical_curl(ucomp, vcomp, lon, lat,
R = 3400*1000.,
spacing = "regular")
R=3400*1000.,
spacing="regular")
else:
err_req_non_interpolated_file(ivar, ifile)

if ivar == "fn":
if interp_type not in ("pstd", "zstd", "zagl"):
theta = fileNC.variables["theta"][:]
OUT = frontogenesis(ucomp, vcomp, theta, lon, lat,
R = 3400*1000.,
spacing = "regular")
R=3400*1000.,
spacing="regular")
else:
err_req_non_interpolated_file(ivar, ifile)

Expand Down
2 changes: 1 addition & 1 deletion setup.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,4 +36,4 @@
include_package_data=True,
zip_safe=False,
python_requires=">=3.7",
)
)