Coverage for src/beamme/four_c/run_four_c.py: 91%

1# The MIT License (MIT) 

2# 

3# Copyright (c) 2018-2026 BeamMe Authors 

4# 

5# Permission is hereby granted, free of charge, to any person obtaining a copy 

6# of this software and associated documentation files (the "Software"), to deal 

7# in the Software without restriction, including without limitation the rights 

8# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 

9# copies of the Software, and to permit persons to whom the Software is 

10# furnished to do so, subject to the following conditions: 

11# 

12# The above copyright notice and this permission notice shall be included in 

13# all copies or substantial portions of the Software. 

14# 

15# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 

16# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 

17# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 

18# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 

19# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 

20# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 

21# THE SOFTWARE. 

22"""Provide a function that allows to run 4C.""" 

23 

24import os as _os 

25import shutil as _shutil 

26import subprocess as _subprocess # nosec B404 

27from pathlib import Path as _Path 

28 

29from beamme.utils.environment import get_env_variable as _get_env_variable 

30 

31 

32def run_four_c( 

33 input_file, 

34 output_dir, 

35 *, 

36 four_c_exe=None, 

37 mpi_command=None, 

38 n_proc=None, 

39 output_name="xxx", 

40 restart_step=None, 

41 restart_from=None, 

42): 

43 """Run a 4C simulation and return the exit code of the run. 

44 

45 This function looks into the environment variables for some parameters: 

46 "BEAMME_FOUR_C_EXE" 

47 "BEAMME_MPI_COMMAND" 

48 "BEAMME_MPI_NUM_PROC" 

49 If the corresponding keyword arguments are set, they overwrite the environment 

50 variable. 

51 

52 Args 

53 ---- 

54 input_file: str 

55 Path to the input file on the filesystem 

56 output_dir: str 

57 Directory where the simulation should be performed (will be created if 

58 it does not exist) 

59 four_c_exe: str 

60 Optionally explicitly specify path to the 4C executable 

61 mpi_command: str 

62 Command to launch MPI, defaults to "mpirun" 

63 n_proc: int 

64 Number of process used with MPI, defaults to 1 

65 output_name: str 

66 Base name of the output files 

67 restart_step: int 

68 Time step to restart from 

69 restart_from: str 

70 Path to initial simulation (relative to output_dir) 

71 

72 Return 

73 ---- 

74 return_code: int 

75 Return code of 4C run 

76 """ 

77 # Fist get all needed parameters 

78 if four_c_exe is None: 

79 four_c_exe = _get_env_variable("BEAMME_FOUR_C_EXE") 

80 if mpi_command is None: 

81 mpi_command = _get_env_variable("BEAMME_MPI_COMMAND", default="mpirun") 

82 if n_proc is None: 

83 n_proc = _get_env_variable("BEAMME_MPI_NUM_PROC", default="1") 

84 

85 # Setup paths and actual command to run 

86 _os.makedirs(output_dir, exist_ok=True) 

87 log_file_path = _os.path.join(output_dir, output_name + ".log") 

88 error_file_path = _os.path.join(output_dir, output_name + ".err") 

89 command = mpi_command.split(" ") + [ 

90 "-np", 

91 str(n_proc), 

92 four_c_exe, 

93 input_file, 

94 output_name, 

95 ] 

96 if restart_step is None and restart_from is None: 

97 pass 

98 elif restart_step is not None and restart_from is not None: 

99 command.extend([f"restart={restart_step}", f"restartfrom={restart_from}"]) 

100 else: 

101 raise ValueError( 

102 "Provide either both or no argument of [restart_step, restart_from]" 

103 ) 

104 

105 # Actually run the command 

106 with ( 

107 open(log_file_path, "w") as log_file, 

108 open(error_file_path, "w") as error_file, 

109 ): 

110 process = _subprocess.Popen( 

111 command, # nosec B603 

112 cwd=output_dir, 

113 stdout=log_file, 

114 stderr=error_file, 

115 ) 

116 return_code = process.wait() 

117 return return_code 

118 

119 

120def clean_simulation_directory(sim_dir, *, ask_before_clean=False): 

121 """Clear the simulation directory. 

122 

123 If it does not exist, it is created. Optionally the user can be 

124 asked before a deletion of files. 

125 

126 Args 

127 ---- 

128 sim_dir: 

129 Path to a directory 

130 ask_before_clean: bool 

131 Flag which indicates whether the user must confirm removal of files and directories 

132 """ 

133 # Check if simulation directory exists. 

134 if _os.path.exists(sim_dir): 

135 if ask_before_clean: 

136 print(f'Path "{sim_dir}" already exists') 

137 while True: 

138 if ask_before_clean: 

139 answer = input("DELETE all contents? (y/n): ") 

140 else: 

141 answer = "y" 

142 if answer.lower() == "y": 

143 for filename in _os.listdir(sim_dir): 

144 file_path = _os.path.join(sim_dir, filename) 

145 try: 

146 if _os.path.isfile(file_path) or _os.path.islink(file_path): 

147 _os.unlink(file_path) 

148 elif _os.path.isdir(file_path): 

149 _shutil.rmtree(file_path) 

150 except Exception as e: 

151 raise ValueError(f"Failed to delete {file_path}. Reason: {e}") 

152 return 

153 elif answer.lower() == "n": 

154 raise ValueError("Directory is not deleted!") 

155 else: 

156 _Path(sim_dir).mkdir(parents=True, exist_ok=True)