Skip to content

Commit

Permalink
refactor(path_handling): move to pathlib.Path
Browse files Browse the repository at this point in the history
Transform all path variables to the more modern `pathlib.Path` object.
This helps with cross-platform compatibility and makes the code more
readable.

Some type annotations were also passed.
Meant to be applied in conjunction with a move to Pydantic for configuration validation.

NOTE: If cherry-picked, needs to be revised carefully as the commit is
meant to be applied with a move to Pydantic.

Fixes cryotools#81
  • Loading branch information
benatouba committed Nov 26, 2024
1 parent a198368 commit 67c0550
Show file tree
Hide file tree
Showing 11 changed files with 152 additions and 109 deletions.
57 changes: 34 additions & 23 deletions COSIPY.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,11 +22,14 @@
Correspondence: [email protected]
"""
import cProfile

from __future__ import annotations

import logging
import os
from datetime import datetime
from itertools import product
from pathlib import Path
from typing import Union

import numpy as np
import pandas as pd
Expand Down Expand Up @@ -116,8 +119,8 @@ def main():
#encoding[var] = dict(zlib=True, complevel=compression_level, dtype=dtype, scale_factor=scale_factor, add_offset=add_offset, _FillValue=FillValue)
encoding[var] = dict(zlib=True, complevel=Config.compression_level)
output_netcdf = set_output_netcdf_path()
output_path = create_data_directory(path='output')
IO.get_result().to_netcdf(os.path.join(output_path,output_netcdf), encoding=encoding, mode='w')
output_path = create_data_directory(name="output")
IO.get_result().to_netcdf(output_path / output_netcdf, encoding=encoding, mode="w")

encoding = dict()
for var in IO.get_restart().data_vars:
Expand All @@ -129,8 +132,8 @@ def main():
#encoding[var] = dict(zlib=True, complevel=compression_level, dtype=dtype, scale_factor=scale_factor, add_offset=add_offset, _FillValue=FillValue)
encoding[var] = dict(zlib=True, complevel=Config.compression_level)

restart_path = create_data_directory(path='restart')
IO.get_restart().to_netcdf(os.path.join(restart_path,f'restart_{timestamp}.nc'), encoding=encoding)
restart_path = create_data_directory(name='restart')
IO.get_restart().to_netcdf(restart_path / f'restart_{timestamp}.nc', encoding=encoding)

#-----------------------------------------------
# Stop time measurement
Expand Down Expand Up @@ -326,19 +329,29 @@ def run_cosipy(cluster, IO, DATA, RESULT, RESTART, futures):

if Config.stake_evaluation:
# Save the statistics and the mass balance simulations at the stakes to files
output_path = create_data_directory(path='output')
df_stat.to_csv(os.path.join(output_path,'stake_statistics.csv'),sep='\t', float_format='%.2f')
df_val.to_csv(os.path.join(output_path,'stake_simulations.csv'),sep='\t', float_format='%.2f')


def create_data_directory(path: str) -> str:
output_path = create_data_directory(name="output")
df_stat.to_csv(
output_path / "stake_statistics.csv",
sep="\t",
float_format="%.2f",
)
df_val.to_csv(
output_path / "stake_simulations.csv",
sep="\t",
float_format="%.2f",
)


def create_data_directory(name: Union[Path, str]) -> Path:
"""Create a directory in the configured data folder.
Returns:
Path to the created directory.
"""
dir_path = os.path.join(Config.data_path, path)
os.makedirs(dir_path, exist_ok=True)
if isinstance(name, Path):
name = name.name
dir_path = Path(Config.data_path) / str(name)
dir_path.mkdir(parents=True, exist_ok=True)

return dir_path

Expand All @@ -355,25 +368,23 @@ def get_timestamp_label(timestamp: str) -> str:
return (timestamp[0:10]).replace("-", "")


def set_output_netcdf_path() -> str:
def set_output_netcdf_path() -> Path:
"""Set the file path for the output netCDF file.
Returns:
The path to the output netCDF file.
"""
time_start = get_timestamp_label(timestamp=Config.time_start)
time_end = get_timestamp_label(timestamp=Config.time_end)
output_path = f"{Config.output_prefix}_{time_start}-{time_end}.nc"

return output_path
return Path(f"{Config.output_prefix}_{time_start}-{time_end}.nc")


def start_logging():
"""Start the python logging"""

if os.path.exists('./cosipy.yaml'):
with open('./cosipy.yaml', 'rt') as f:
config = yaml.load(f.read(),Loader=yaml.SafeLoader)
"""Start the python logging."""
log_config_path = Path("./cosipy.yaml")
if log_config_path.exists():
with log_config_path.open() as f:
config = yaml.load(f.read(), Loader=yaml.SafeLoader)
logging.config.dictConfig(config)
else:
logging.basicConfig(level=logging.INFO)
Expand Down
26 changes: 17 additions & 9 deletions convert_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
import configparser
import inspect
import sys
from pathlib import Path

import config
import constants
Expand Down Expand Up @@ -284,12 +285,12 @@ def get_utilities_params() -> dict:
return params


def write_toml(parameters: dict, filename: str):
def write_toml(parameters: dict, filename: Path | str):
"""Write parameters to .toml file."""

with open(f"{filename}.toml", "w") as f:
if isinstance(filename, str):
filename = Path(filename)
with filename.with_suffix(".toml").open("w") as f:
toml.dump(parameters, f)


print(f"Generated {filename}.toml")

Expand All @@ -308,15 +309,22 @@ def main():

print_warning()

script_path = inspect.getfile(inspect.currentframe())
toml_suffix = script_path.split("/")[-2] # avoid overwrite
frame = inspect.currentframe()
if frame is None:
msg = "Could not find the current frame. This is likely due to a bug in the code."
raise RuntimeError(msg)
try:
script_path = Path(inspect.getfile(frame)).resolve()
finally:
del frame
_ = script_path.parent.name # HACK: avoid overwrite (Why is this here?)

config_params = get_config_params()
write_toml(parameters=config_params, filename=f"config")
write_toml(parameters=config_params, filename="config")
constants_params = get_constants_params()
write_toml(parameters=constants_params, filename=f"constants")
write_toml(parameters=constants_params, filename="constants")
utilities_params = get_utilities_params()
write_toml(parameters=utilities_params, filename=f"utilities_config")
write_toml(parameters=utilities_params, filename="utilities_config")


if __name__ == "__main__":
Expand Down
20 changes: 14 additions & 6 deletions cosipy/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,12 +5,20 @@
import argparse
import sys
from importlib.metadata import entry_points
from pathlib import Path

if sys.version_info >= (3, 11):
import tomllib
else:
import tomli as tomllib # backwards compatibility

# FIXME: Will this work for all occasions or do we need to use frame?
cwd = Path.cwd()
default_path = cwd / "config.toml"
default_slurm_path = cwd / "slurm_config.toml"
default_constants_path = cwd / "constants.toml"
default_utilities_path = cwd / "utilities_config.toml"


def set_parser() -> argparse.ArgumentParser:
"""Set argument parser for COSIPY."""
Expand All @@ -23,9 +31,9 @@ def set_parser() -> argparse.ArgumentParser:
parser.add_argument(
"-c",
"--config",
default="./config.toml",
default=default_path,
dest="config_path",
type=str,
type=Path,
metavar="<path>",
required=False,
help="relative path to configuration file",
Expand All @@ -34,9 +42,9 @@ def set_parser() -> argparse.ArgumentParser:
parser.add_argument(
"-x",
"--constants",
default="./constants.toml",
default=default_constants_path,
dest="constants_path",
type=str,
type=Path,
metavar="<path>",
required=False,
help="relative path to constants file",
Expand All @@ -45,9 +53,9 @@ def set_parser() -> argparse.ArgumentParser:
parser.add_argument(
"-s",
"--slurm",
default="./slurm_config.toml",
default=default_slurm_path,
dest="slurm_path",
type=str,
type=Path,
metavar="<path>",
required=False,
help="relative path to Slurm configuration file",
Expand Down
2 changes: 2 additions & 0 deletions cosipy/constants.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,6 @@
import sys
from pathlib import Path
from typing import Literal

from cosipy.config import Config, TomlLoader, get_user_arguments

Expand Down
25 changes: 14 additions & 11 deletions cosipy/cpkernel/io.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,11 @@
"""
Read the input data (model forcing) and write the output to netCDF file.
"""
"""Read the input data (model forcing) and write the output to netCDF file."""

from __future__ import annotations

import os
import warnings
from datetime import datetime
from pathlib import Path

import numpy as np
import xarray as xr
Expand All @@ -14,8 +15,7 @@


class IOClass:

def __init__(self, DATA=None):
def __init__(self, data: xr.Dataset | None = None):
"""Initialise the IO Class.
Attributes:
Expand Down Expand Up @@ -134,11 +134,15 @@ def create_data_file(self) -> xr.Dataset:
start_timestamp = self.get_datetime(time_start)
end_timestamp = self.get_datetime(time_end)
timestamp = start_timestamp.strftime("%Y-%m-%dT%H-%M")
restart_path = os.path.join(
Config.data_path, "restart", f"restart_{timestamp}.nc"
)
restart_path = Path(Config.data_path) / "restart" / f"restart_{timestamp}.nc"
if not restart_path.is_file():
msg = f"No restart file available at {restart_path}"
raise FileNotFoundError(msg)
if start_timestamp == end_timestamp:
msg = f"Start date {time_start} equals end date {time_end}"
raise IndexError(msg)
try:
if not os.path.isfile(restart_path):
if not restart_path.is_file():
raise FileNotFoundError
elif start_timestamp == end_timestamp:
raise IndexError
Expand Down Expand Up @@ -269,9 +273,8 @@ def init_data_dataset(self):
:U2: Wind speed (magnitude) [|m s^-1|].
:HGT: Elevation [m].
"""

input_path = Path(Config.data_path) / "input" / Config.input_netcdf
try:
input_path = os.path.join(Config.data_path, "input", Config.input_netcdf)
self.DATA = xr.open_dataset(input_path)
except FileNotFoundError:
raise SystemExit(f"Input file not found at: {input_path}")
Expand Down
3 changes: 2 additions & 1 deletion cosipy/postprocessing/field_plots/plot_cosipy_fields.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
"""

import argparse
from pathlib import Path
import re

import cartopy.crs as ccrs
Expand Down Expand Up @@ -409,7 +410,7 @@ def parse_arguments() -> argparse.Namespace:
dest="file",
required=True,
default=None,
type=str,
type=Path,
metavar="<path>",
help="Path to .nc file",
)
Expand Down
8 changes: 6 additions & 2 deletions cosipy/postprocessing/profile_plots/plot_profiles.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,7 @@
import argparse
import os
from pathlib import Path

import matplotlib
# matplotlib.use("TkAgg")
import matplotlib.pyplot as plt
Expand Down Expand Up @@ -242,14 +246,14 @@ def naive_fast(latvar,lonvar,lat0,lon0):
if __name__ == "__main__":

parser = argparse.ArgumentParser(description='Quick plot of the results file.')
parser.add_argument('-f', '--file', dest='file', help='Path to the result file')
parser.add_argument('-f', '--file', dest='file', type=Path, help='Path to the result file')
parser.add_argument('-d', '--date', dest='pdate', help='Date of the profile plot')
parser.add_argument('-v', '--var', dest='var', default='RHO', help='Which variable to plot (e.g. T, RHO, etc.)')
parser.add_argument('-n', '--lat', dest='lat', default=None, help='Latitude value in case of 2D simulation', type=float)
parser.add_argument('-m', '--lon', dest='lon', default=None, help='Longitude value in case of 2D simulation', type=float)
parser.add_argument('-s', '--start', dest='start', default=None, help='Start date for the time plot')
parser.add_argument('-e', '--end', dest='end', default=None, help='End date for the time plot')
parser.add_argument('--stake-file', dest='stake_file', default=None, help='Path to the stake data file')
parser.add_argument('--stake-file', dest='stake_file', default=None, type=Path, help='Path to the stake data file')
parser.add_argument('--pit', dest='pit_name', default=None, help='Name of the pit in the stake data file')
parser.add_argument('--depth', dest='d', nargs='+', default=None, help='An array with depth values for which the corresponding values are to be displayed', type=float)

Expand Down
Loading

0 comments on commit 67c0550

Please sign in to comment.