Coverage for biobb_haddock/haddock/common.py: 0%

1"""Common functions for package biobb_haddock.haddock""" 

2 

3import shutil 

4import logging 

5import os 

6import jsonpickle 

7from pathlib import Path 

8from typing import Any, Optional 

9import biobb_common.tools.file_utils as fu 

10from biobb_common.generic.biobb_object import BiobbObject 

11from biobb_common.tools.file_utils import launchlogger 

12from .haddock3_config import load, save 

13 

14haddock_2_wf = { 

15 'ambig_fname': 'ambig_restraints_table_path', 

16 'unambig_fname': 'unambig_restraints_table_path', 

17 'hbond_fname': 'hb_restraints_table_path', 

18} 

19 

20 

21class HaddockStepBase(BiobbObject): 

22 """Base class for HADDOCK3 step modules with shared launch functionality.""" 

23 

24 def copy_step_output( 

25 self, filter_funct: callable, 

26 output_zip_path: str, sele_top: bool = False) -> None: 

27 """Copy the output files from the run directory to the output zip path. 

28 

29 Args: 

30 obj: The object containing the output paths. 

31 run_dir (str): The directory where the output files are located. 

32 filter_funct (callable): A function that accepts a Path and returns True for the files to be copied. 

33 output_zip_path (str): The path where the output zip file will be created.""" 

34 # Find the directories with the haddock step name 

35 haddock_output_list = [ 

36 str(path) 

37 for path in Path(self.run_dir).iterdir() 

38 if path.is_dir() and str(path).endswith(self.haddock_step_name) 

39 ] 

40 # Make the one with the highest step number the first one 

41 haddock_output_list.sort(reverse=True) 

42 # Select files with filter_funct 

43 output_file_list = [ 

44 str(path) 

45 for path in Path(haddock_output_list[0]).iterdir() 

46 if path.is_file() and filter_funct(path) 

47 ] 

48 if sele_top: 

49 with open(haddock_output_list[0]+'/io.json') as json_file: 

50 content = jsonpickle.decode(json_file.read()) 

51 output = content["output"] 

52 for file in output: 

53 rel_path = str(file.rel_path).split('/') 

54 output_file_list.extend(list(Path(self.run_dir+'/'+rel_path[-2]).glob(rel_path[-1]+'*'))) 

55 if len(output_file_list) == 0: 

56 fu.log("No output files found matching the criteria.", self.out_log, self.global_log) 

57 else: 

58 fu.zip_list(output_zip_path, output_file_list, self.out_log) 

59 

60 @launchlogger 

61 def launch(self) -> int: 

62 """Execute the HADDOCK step with common workflow.""" 

63 # Setup Biobb 

64 if self.check_restart(): 

65 return 0 

66 self.stage_files() 

67 

68 if self.stage_io_dict["in"]["input_haddock_wf_data"][-4:] == ".zip": 

69 # Unzip workflow data to workflow_data_out 

70 new_input = fu.create_unique_dir(self.stage_io_dict["unique_dir"], '_input_unzipped') 

71 fu.unzip_list(self.stage_io_dict["in"]["input_haddock_wf_data"], new_input) # , self.out_log) 

72 self.stage_io_dict["in"]["input_haddock_wf_data"] = new_input 

73 

74 self.run_dir = self.stage_io_dict["out"]["output_haddock_wf_data"] 

75 if self.stage_io_dict["in"]["input_haddock_wf_data"] != self.run_dir: 

76 # Different I/O folder 

77 shutil.copytree(self.stage_io_dict["in"]["input_haddock_wf_data"], 

78 self.run_dir, dirs_exist_ok=True) 

79 else: 

80 # Same I/O folder 

81 os.rename(self.stage_io_dict["in"]["input_haddock_wf_data"], self.run_dir) 

82 # Check if there are more than 9 numbered folders and rename them to add leading zeros 

83 input_wf = self.io_dict["in"]["input_haddock_wf_data"] 

84 numbered_dirs = [] 

85 for item in os.listdir(input_wf): 

86 item_path = os.path.join(input_wf, item) 

87 if os.path.isdir(item_path) and item[0].isdigit(): 

88 numbered_dirs.append(item) 

89 

90 if len(numbered_dirs) in [10, 100, 1000] and not self.disable_sandbox: 

91 for dirname in numbered_dirs: 

92 os.rename(os.path.join(input_wf, dirname), 

93 os.path.join(input_wf, f"0{dirname}")) 

94 

95 workflow_dict = {"haddock_step_name": self.haddock_step_name} 

96 workflow_dict.update(self.global_cfg) 

97 

98 if hasattr(self, '_handle_config_arguments'): 

99 self._handle_config_arguments() 

100 

101 # Create workflow configuration 

102 self.output_cfg_path = create_cfg( 

103 output_cfg_path=self.create_tmp_file('_haddock.cfg'), 

104 workflow_dict=workflow_dict, 

105 input_cfg_path=self.stage_io_dict["in"].get("haddock_config_path"), 

106 properties_cfg_dict=self.cfg, 

107 out_log=self.out_log, 

108 global_log=self.global_log, 

109 ) 

110 

111 if self.container_path: 

112 fu.log("Container execution enabled", self.out_log) 

113 move_to_container_path(self, self.run_dir) 

114 

115 self.cmd = [self.binary_path, self.output_cfg_path] 

116 

117 if self.haddock_step_name != "haddock3_run": 

118 self.cmd.extend(["--extend-run", os.path.abspath(self.run_dir)]) 

119 

120 # Run Biobb block 

121 with fu.change_dir(self.run_dir): 

122 self.run_biobb() 

123 

124 # Copy files to host 

125 if hasattr(self, '_handle_step_output'): 

126 self._handle_step_output() 

127 if self.io_dict["out"]["output_haddock_wf_data"][-4:] == ".zip": 

128 zip_wf_output(self) 

129 else: 

130 self.copy_to_host() 

131 

132 # Remove temporal files 

133 self.remove_tmp_files() 

134 

135 return self.return_code 

136 

137 

138def create_cfg( 

139 output_cfg_path: str, 

140 workflow_dict: dict[str, Any], 

141 input_cfg_path: Optional[str] = None, 

142 properties_cfg_dict: Optional[dict[str, str]] = None, 

143 out_log: Optional[logging.Logger] = None, 

144 global_log: Optional[logging.Logger] = None, 

145) -> str: 

146 """Creates an CFG file using the following hierarchy: properties_cfg_dict > input_cfg_path > preset_dict""" 

147 cfg_dict: dict[str, Any] = {} 

148 

149 # Handle input configuration if it exists 

150 if input_cfg_path: 

151 input_cfg = load(input_cfg_path) 

152 print(f"Input CFG: {input_cfg}") 

153 cfg_dict = input_cfg.copy() # Start with entire loaded config as base 

154 

155 # Apply single step configuration if specified 

156 haddock_step_name: str = workflow_dict["haddock_step_name"] 

157 if not haddock_step_name.startswith("haddock3_"): 

158 # Get preset properties for this step if any 

159 step_preset = cfg_preset(haddock_step_name) 

160 

161 # Create or update the step configuration 

162 if not cfg_dict: 

163 # No input config, create new structure with single step 

164 target_key = haddock_step_name 

165 cfg_dict = {target_key: step_preset or {}} 

166 else: 

167 # Update the specific step in the existing config 

168 target_key = f"{haddock_step_name}.1" 

169 if target_key not in cfg_dict: 

170 cfg_dict[target_key] = {} 

171 # Merge preset values while preserving existing values 

172 if step_preset: 

173 for k, v in step_preset.items(): 

174 if k not in cfg_dict[target_key]: # Only add if not already defined 

175 cfg_dict[target_key][k] = v 

176 

177 # Apply custom properties to the step 

178 if properties_cfg_dict: 

179 for k, v in properties_cfg_dict.items(): 

180 fu.log(f"CFG from properties: {k} = {v}", out_log, global_log) 

181 cfg_dict[target_key][k] = v 

182 # Multiple steps: haddock3_run and haddock3_extend 

183 else: 

184 if properties_cfg_dict: 

185 for key, value in properties_cfg_dict.items(): 

186 if isinstance(value, dict): 

187 # If the value is a dictionary, update the corresponding section in cfg_dict 

188 if key not in cfg_dict: 

189 cfg_dict[key] = {} 

190 for sub_key, sub_value in value.items(): 

191 fu.log(f"CFG: {key}.{sub_key} = {sub_value}", out_log, global_log) 

192 cfg_dict[key][sub_key] = sub_value 

193 else: 

194 # If the value is not a dictionary, treat it as a top-level property 

195 fu.log(f"CFG: {key} = {value}", out_log, global_log) 

196 cfg_dict[key] = value 

197 # Add workflow_dict properties to cfg_dict 

198 for key, value in cfg_dict.items(): 

199 if isinstance(value, dict): 

200 for sub_key, sub_value in value.items(): 

201 mapped_key = haddock_2_wf.get(sub_key) 

202 if mapped_key and mapped_key in workflow_dict: 

203 sub_value = workflow_dict[mapped_key] 

204 cfg_dict[key][sub_key] = sub_value 

205 if haddock_step_name == "haddock3_run" and "run_dir" not in cfg_dict: 

206 cfg_dict["run_dir"] = "run" 

207 

208 # Add molecules and run_dir if provided 

209 for key, value in workflow_dict.items(): 

210 if key == 'haddock_step_name' or key in haddock_2_wf.values(): 

211 continue 

212 fu.log(f"CFG: {key} = {value}", out_log, global_log) 

213 cfg_dict[key] = value 

214 

215 # Use haddock save 

216 save(cfg_dict, output_cfg_path) 

217 

218 return output_cfg_path 

219 

220 

221def cfg_preset(haddock_step_name: str) -> dict[str, Any]: 

222 cfg_dict: dict[str, Any] = {} 

223 # cfg_dict["debug"] = True 

224 

225 if haddock_step_name == "topoaa": 

226 cfg_dict["autohis"] = True 

227 cfg_dict["delenph"] = True 

228 cfg_dict["log_level"] = "quiet" 

229 cfg_dict["iniseed"] = 917 

230 cfg_dict["ligand_param_fname"] = "" 

231 cfg_dict["ligand_top_fname"] = "" 

232 cfg_dict["limit"] = True 

233 cfg_dict["tolerance"] = 0 

234 

235 elif haddock_step_name == "rigidbody": 

236 cfg_dict["sampling"] = 20 

237 cfg_dict["tolerance"] = 20 

238 

239 elif haddock_step_name == "seletop": 

240 cfg_dict["select"] = 5 

241 

242 elif haddock_step_name == "flexref": 

243 cfg_dict["tolerance"] = 20 

244 

245 elif haddock_step_name == "emref": 

246 cfg_dict["tolerance"] = 20 

247 

248 return cfg_dict 

249 

250 

251def move_to_container_path(obj, run_dir=None): 

252 """Move configuration and run directory to container path.""" 

253 shutil.copy2(obj.output_cfg_path, obj.stage_io_dict.get("unique_dir", "")) 

254 obj.output_cfg_path = str( 

255 Path(obj.container_volume_path).joinpath( 

256 Path(obj.output_cfg_path).name 

257 ) 

258 ) 

259 if run_dir: 

260 shutil.copytree( 

261 run_dir, 

262 str( 

263 Path(obj.stage_io_dict.get("unique_dir", "")).joinpath( 

264 Path(run_dir).name 

265 ) 

266 ), 

267 ) 

268 run_dir = str(Path(obj.stage_io_dict.get("unique_dir", "")).joinpath(Path(run_dir).name)) 

269 

270 

271def zip_wf_output(obj): 

272 """Zip all the files in the run directory and save it to the output path.""" 

273 dest_path = str(Path(obj.io_dict["out"]["output_haddock_wf_data"]).with_suffix('')) 

274 fu.log(f"Zipping {obj.run_dir} to {dest_path} ", obj.out_log, obj.global_log) 

275 shutil.make_archive(dest_path, "zip", obj.run_dir)