Coverage for wsimod/__main__.py: 0%

44 statements  

« prev     ^ index     » next       coverage.py v7.3.2, created at 2024-01-11 09:32 +0000

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

2from argparse import ArgumentParser 

3from pathlib import Path 

4from typing import Any, cast 

5 

6import pandas as pd 

7 

8from wsimod.orchestration.model import Model 

9from wsimod.validation import ( 

10 _validate_input_dir, 

11 _validate_output_dir, 

12 assign_data_to_settings, 

13 evaluate_input_file, 

14 load_data_files, 

15 validate_io_args, 

16) 

17 

18 

19def create_parser() -> ArgumentParser: 

20 """Create the CLI argument parser.""" 

21 parser = ArgumentParser(prog="WSIMOD") 

22 parser.add_argument( 

23 "settings", 

24 type=Path, 

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

26 ) 

27 parser.add_argument( 

28 "--inputs", 

29 "-i", 

30 type=Path, 

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

32 " settings file.", 

33 ) 

34 parser.add_argument( 

35 "--outputs", 

36 "-o", 

37 type=Path, 

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

39 " settings file.", 

40 ) 

41 

42 return parser 

43 

44 

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

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

47 

48 Args: 

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

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

51 """ 

52 model = Model() 

53 

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

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

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

57 

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

59 

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

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

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

63 

64 

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

66 """Runs a previously saved model. 

67 

68 Args: 

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

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

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

72 """ 

73 model = Model() 

74 model.load(inputs, config_name=settings) 

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

76 

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

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

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

80 

81 

82def run() -> None: 

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

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

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

86 

87 if settings_type == "custom": 

88 settings = validate_io_args(**args) 

89 inputs = settings.pop("inputs") 

90 outputs = settings.pop("outputs") 

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

92 loaded_settings = assign_data_to_settings(settings, loaded_data) 

93 run_model(loaded_settings, outputs) 

94 else: 

95 settings_file: Path = args["settings"] 

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

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

98 run_saved_model(settings_file.name, inputs, outputs) 

99 

100 

101if __name__ == "__main__": 

102 run()