Coverage for wsimod\__main__.py: 0%

44 statements  

« prev     ^ index     » next       coverage.py v7.6.1, created at 2024-10-24 11:16 +0100

1"""The entry point for the myproject program.""" 

2 

3from argparse import ArgumentParser 

4from pathlib import Path 

5from typing import Any, cast 

6 

7import pandas as pd 

8 

9from wsimod.orchestration.model import Model 

10from wsimod.validation import ( 

11 _validate_input_dir, 

12 _validate_output_dir, 

13 assign_data_to_settings, 

14 evaluate_input_file, 

15 load_data_files, 

16 validate_io_args, 

17) 

18 

19 

20def create_parser() -> ArgumentParser: 

21 """Create the CLI argument parser.""" 

22 parser = ArgumentParser(prog="WSIMOD") 

23 parser.add_argument( 

24 "settings", 

25 type=Path, 

26 help="Path to the WSIMOD input file, in YAML format.", 

27 ) 

28 parser.add_argument( 

29 "--inputs", 

30 "-i", 

31 type=Path, 

32 help="Base directory for all input files. If present, overwrites value in the" 

33 " settings file.", 

34 ) 

35 parser.add_argument( 

36 "--outputs", 

37 "-o", 

38 type=Path, 

39 help="Base directory for all output files. If present, overwrites value in the" 

40 " settings file.", 

41 ) 

42 

43 return parser 

44 

45 

46def run_model(settings: dict[str, Any], outputs: Path) -> None: 

47 """Runs the mode with the chosen settings and saves the outputs as csv. 

48 

49 Args: 

50 settings (dict[str, Any]): Settings dictionary with loaded data. 

51 outputs(Path): Directory where to save the outputs. 

52 """ 

53 model = Model() 

54 

55 model.dates = cast(pd.Series, settings["dates"]).drop_duplicates() 

56 model.add_nodes(settings["nodes"]) 

57 model.add_arcs(settings["arcs"]) 

58 

59 flows, tanks, _, surfaces = model.run() 

60 

61 pd.DataFrame(flows).to_csv(outputs / "flows.csv") 

62 pd.DataFrame(tanks).to_csv(outputs / "tanks.csv") 

63 pd.DataFrame(surfaces).to_csv(outputs / "surfaces.csv") 

64 

65 

66def run_saved_model(settings: str, inputs: Path, outputs: Path) -> None: 

67 """Runs a previously saved model. 

68 

69 Args: 

70 settings (str): The name of the settings file to load. 

71 inputs (Path): The location of that file, as well as any other input file. 

72 outputs (Path): Directory where to save the outputs. 

73 """ 

74 model = Model() 

75 model.load(inputs, config_name=settings) 

76 flows, tanks, _, surfaces = model.run() 

77 

78 pd.DataFrame(flows).to_csv(outputs / "flows.csv") 

79 pd.DataFrame(tanks).to_csv(outputs / "tanks.csv") 

80 pd.DataFrame(surfaces).to_csv(outputs / "surfaces.csv") 

81 

82 

83def run() -> None: 

84 """Main entry point of the application.""" 

85 args = vars(create_parser().parse_args()) 

86 settings_type = evaluate_input_file(args["settings"]) 

87 

88 if settings_type == "custom": 

89 settings = validate_io_args(**args) 

90 inputs = settings.pop("inputs") 

91 outputs = settings.pop("outputs") 

92 loaded_data = load_data_files(settings.pop("data", {}), inputs) 

93 loaded_settings = assign_data_to_settings(settings, loaded_data) 

94 run_model(loaded_settings, outputs) 

95 else: 

96 settings_file: Path = args["settings"] 

97 inputs = _validate_input_dir(args["inputs"], settings_file.parent) 

98 outputs = _validate_output_dir(args["outputs"], settings_file.parent) 

99 run_saved_model(settings_file.name, inputs, outputs) 

100 

101 

102if __name__ == "__main__": 

103 run()