Repository navigation
Expand file tree
/
Copy pathutils.py
More file actions
453 lines (401 loc) · 18.5 KB
/
Copy pathutils.py
File metadata and controls
453 lines (401 loc) · 18.5 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
import argparse
import collections.abc
import json
import yaml
import sys
from typing import Any, Dict, List, Optional
import os
import logging
import subprocess
from collections import OrderedDict
from functools import reduce
logger = logging.getLogger(__name__)
def init_logger(log_file=None, level=logging.INFO):
root_logger = logging.getLogger()
for h in root_logger.handlers:
root_logger.removeHandler(h)
basic_fmt = "[%(asctime)s] %(levelname)-8s | %(filename)-12s:%(lineno)-4d | %(message)s"
logging.basicConfig(
filename=None, filemode="a", format=basic_fmt, level=level,
)
if log_file:
root_logger = logging.getLogger()
fhlr = logging.FileHandler(log_file)
formatter = logging.Formatter(basic_fmt)
fhlr.setFormatter(formatter)
root_logger.addHandler(fhlr)
def recursive_update(d, u):
"""recursivly update a dictionary
https://stackoverflow.com/questions/3232943/update-value-of-a-nested-dictionary-of-varying-depth
Args:
d (dict): dictionary being updated
u (dict): dictionary storing updates
Returns:
dict: updated dictionary
"""
for k, v in u.items():
if isinstance(v, collections.abc.Mapping):
d[k] = recursive_update(d.get(k, {}), v)
else:
d[k] = v
return d
class ConfigManager:
"""Singleton General Config Manager for Tinker GUI/CLI applications.
`CONFIG_DEFINITION` serves as the single source of truth for all configuration parameters.
It defines each parameter's name, default value, type, and help description.
"""
CONFIG_DEFINITION = [
{'name': 'tinker_path', 'default': '/path/to/tinker/bin', 'type': str, 'help': 'Path to the Tinker executables.'},
{'name': 'amoeba_prm', 'default': '/path/to/amoeba.prm', 'type': str, 'help': 'Path to the AMOEBA parameter file.'},
{'name': 'output_prefix', 'default': 'my_system', 'type': str, 'help': 'Prefix for output files.'},
{'name': 'solutes.protein', 'default': ['/path/to/my_protein.pdb', '/path/to/my_protein_2.pdb'], 'type': list, 'help': 'List of protein PDB files.'},
{'name': 'solutes.nucleic_acid', 'default': ['/path/to/my_dna.pdb'], 'type': list, 'help': 'List of nucleic acid PDB files.'},
# {'name': 'solutes.ligand', 'default': ['/path/to/my_ligand.sdf'], 'type': list, 'help': 'List of ligand SDF files.'},
{'name': 'solvent.name', 'default': 'water', 'type': str, 'help': 'Solvent name.'},
# {'name': 'solvent.density', 'default': 1.000, 'type': float, 'help': 'Solvent density in g/cm^3.'},
{'name': 'ions.neutralizers', 'default': ['Na+', 'Cl-'], 'type': list, 'help': 'List of neutralizing ions.'},
# {'name': 'ions.salts', 'default': {'K': 0.15, 'Cl': 0.15}, 'type': dict, 'help': 'Concentration of salts in mol/L.'},
{'name': 'ions.salts.names', 'default': ['K+', 'Cl-'], 'type': list, 'help': 'List of salt names.'},
{'name': 'ions.salts.concentrations', 'default': [0.15, 0.15], 'type': list, 'help': 'List of respective concentrations in mol/L.'},
{'name': 'box.type', 'default': 'rectangular', 'type': str, 'help': 'Box type. Options: cubic, rectangular.', 'choices': ['cubic', 'rectangular']},
{'name': 'box.buffer', 'default': 15.0, 'type': float, 'help': 'Box buffer size in Å.'},
{'name': 'box.size', 'default': [], 'type': list, 'help': 'Cuboid box size in Å for each dimension (this one will take precedence over box.buffer and box.type). (Example: 40.0 50.0 60.0)'},
]
_instance = None
_initialized = False
def __new__(cls):
if cls._instance is None:
cls._instance = super(ConfigManager, cls).__new__(cls)
return cls._instance
def __init__(self):
# Only initialize once due to singleton pattern
if ConfigManager._initialized:
return
self.config = argparse.Namespace()
self.parser = argparse.ArgumentParser(description="Tinker-GUI CLI Config Manager")
self._add_arguments()
ConfigManager._initialized = True
def _add_arguments(self):
"""Add command-line arguments based on CONFIG_DEFINITION."""
self.parser.add_argument(
'-c', '--config',
type=str,
help='Path to the YAML config file.'
)
self.parser.add_argument(
'--generate-yaml-template',
action='store_true',
help='Generate a sample YAML config and exit.'
)
for param in ConfigManager.CONFIG_DEFINITION:
# Allow hierarchical names using dot-paths; use dashes on CLI and map to safe dest
path_name = param['name']
option = f"--{path_name.replace('.', '-')}"
dest = path_name.replace('.', '__')
if param['type'] == list:
self.parser.add_argument(
option,
dest=dest,
type=type(param['default'][0]) if param['default'] else str,
nargs='+',
default=None, # Use None to detect if user provided it
help=param['help'] + f" (Default: {' '.join(map(str, param['default']))})"
)
elif param['type'] == dict:
self.parser.add_argument(
option,
dest=dest,
type=json.loads,
default=None, # Use None to detect if user provided it
help=param['help'] + f" (Default: '{json.dumps(param['default'])}')"
)
else:
self.parser.add_argument(
option,
dest=dest,
type=param['type'],
default=None, # Use None to detect if user provided it
help=param['help'] + f" (Default: {param['default']})"
)
def _check_arguments_validity(self):
"""Check validity of certain arguments after parsing."""
has_solute = False
for param in ConfigManager.CONFIG_DEFINITION:
# check tinker_path, amoeba_prm must be different from default and exist
if param['name'] in ['tinker_path', 'amoeba_prm']:
value = reduce(getattr, param['name'].split('.'), self.config)
if value == param['default']:
logger.error(f"Configuration parameter '{param['name']}' must be set to a valid path, not the default value.")
sys.exit(1)
if not os.path.exists(value):
logger.error(f"Path specified for '{param['name']}' does not exist: {value}")
sys.exit(1)
# check at least one of protein/nucleic_acid/ligand must be provided, different from default, and exist
if param['name'].startswith('solutes.'):
try:
value = reduce(getattr, param['name'].split('.'), self.config)
if value and value != param['default']:
for fpath in value:
if not os.path.exists(fpath):
logger.error(f"Solute file specified for '{param['name']}' does not exist: {fpath}")
sys.exit(1)
has_solute = True
except AttributeError:
continue
# assign valid choices if applicable and check validity
if 'choices' in param:
value = reduce(getattr, param['name'].split('.'), self.config)
if value not in param['choices']:
logger.error(f"Configuration parameter '{param['name']}' has invalid value '{value}'. Valid choices are: {param['choices']}")
sys.exit(1)
if not has_solute:
logger.error("At least one valid solute (protein, nucleic acid, or ligand) must be provided.")
sys.exit(1)
def _generate_yaml_template(self):
"""Generate a sample YAML configuration file based on CONFIG_DEFINITION."""
# Use OrderedDict to maintain the order of parameters
template_dict: Dict[str, Any] = OrderedDict()
for param in ConfigManager.CONFIG_DEFINITION:
# Support nested keys with dot-paths
self._set_by_path(template_dict, param['name'].split('.'), param['default'])
yaml_string = yaml.dump(dict(template_dict), sort_keys=False)
if os.path.exists("sample_config.yaml"):
logger.error("sample_config.yaml already exists. Please remove it before generating a new template.")
return
else:
with open("sample_config.yaml", "w") as f:
f.write(yaml_string)
# ---------- Helpers for hierarchical configuration ----------
def _set_by_path(self, d: Dict[str, Any], path: List[str], value: Any) -> None:
cur = d
for key in path[:-1]:
if key not in cur or not isinstance(cur[key], (dict, OrderedDict)):
cur[key] = {}
cur = cur[key]
cur[path[-1]] = value
def _remove_by_path(self, d: Dict[str, Any], path: List[str]) -> None:
cur = d
for key in path[:-1]:
if key not in cur or not isinstance(cur[key], (dict, OrderedDict)):
return
cur = cur[key]
if path[-1] in cur:
del cur[path[-1]]
def _deep_merge(self, base: Dict[str, Any], inc: Dict[str, Any]) -> Dict[str, Any]:
"""Recursively merge inc into base and return a new dict."""
if not isinstance(base, (dict, OrderedDict)):
base = {}
if not isinstance(inc, (dict, OrderedDict)):
return dict(base)
result = dict()
all_keys = list(base.keys()) + list(inc.keys())
for k in sorted(set(all_keys), key=lambda x: all_keys.index(x)):
bv = base.get(k)
iv = inc.get(k)
if isinstance(bv, (dict, OrderedDict)) and isinstance(iv, (dict, OrderedDict)):
result[k] = self._deep_merge(bv, iv)
elif iv is not None:
result[k] = iv
else:
result[k] = bv
return result
def _to_namespace(self, data: Dict[str, Any]) -> argparse.Namespace:
"""Convert nested dicts into nested Namespaces for attribute access."""
if not isinstance(data, (dict, OrderedDict)):
return argparse.Namespace()
result = OrderedDict()
for k, v in data.items():
if isinstance(v, (dict, OrderedDict)):
result[k] = self._to_namespace(v)
else:
result[k] = v
return argparse.Namespace(**result)
def _to_plain_dict(self, ns_or_dict: Any) -> Dict[str, Any]:
"""Convert Namespace (possibly nested) to plain dict recursively."""
if isinstance(ns_or_dict, argparse.Namespace):
ns_or_dict = vars(ns_or_dict)
if not isinstance(ns_or_dict, (dict, OrderedDict)):
return {}
out = dict()
for k, v in ns_or_dict.items():
if isinstance(v, argparse.Namespace) or isinstance(v, (dict, OrderedDict)):
out[k] = self._to_plain_dict(v)
else:
out[k] = v
return out
# ---------- Public Methods ----------
def parse_args(self):
"""Parse command-line arguments, load YAML config, and merge them with defaults.
Priority: Command-line > YAML file > Defaults"""
# Build defaults as possibly nested dict based on dotted names
defaults = OrderedDict()
for param in ConfigManager.CONFIG_DEFINITION:
self._set_by_path(defaults, param['name'].split('.'), param['default'])
if "choices" in param:
self._set_by_path(defaults, ["choices"] + param['name'].split('.'), param['choices'])
cmd_args = self.parser.parse_args()
# Check for special flags first
if cmd_args.generate_yaml_template:
self._generate_yaml_template()
sys.exit(0)
yaml_config = OrderedDict()
if cmd_args.config:
with open(cmd_args.config, 'r') as f:
yaml_config = yaml.safe_load(f) or {}
# Start with defaults and deep-merge YAML config values
final_config = self._deep_merge(defaults, yaml_config)
# Finally, override with command-line arguments if provided
for param in ConfigManager.CONFIG_DEFINITION:
path = param['name']
dest = path.replace('.', '__')
cmd_value = getattr(cmd_args, dest, None)
if cmd_value is not None:
self._set_by_path(final_config, path.split('.'), cmd_value)
# Remove solutes entries that were not provided via CLI or YAML thus exactly the same as examplary defaults
for param in ConfigManager.CONFIG_DEFINITION:
if param['name'].startswith('solutes.'):
path = param['name']
default_value = param['default']
final_value = final_config
for key in path.split('.'):
final_value = final_value.get(key, None)
if final_value is None:
break
if final_value == default_value:
self._remove_by_path(final_config, path.split('.'))
recursive_update(self.config.__dict__, self._to_namespace(final_config).__dict__)
self._check_arguments_validity()
logger.info(f"Final Configuration: \n{yaml.dump(final_config, sort_keys=False)}")
def get_config(self) -> argparse.Namespace:
"""Return the current configuration as an argparse.Namespace object."""
return self.config
def save_yaml(self, path: str):
data = self._to_plain_dict(self.config)
with open(path, 'w') as f:
yaml.safe_dump(data, f, sort_keys=False)
logger.info(f"Configuration saved to {path}")
class TinkerKeyFile:
"""Class for Tinker key files."""
def __init__(self, key_file_path: str):
self.key_file_path = key_file_path
self.latest_path = key_file_path
if not os.path.isfile(self.key_file_path):
self.keys = []
logger.info(f"Key file {self.key_file_path} does not exist. Starting with empty keys.")
else:
self.keys = self._load_key_file()
logger.info(f"Loaded key file from {self.key_file_path}.")
def _load_key_file(self):
"""Load and parse the Tinker key file."""
keys = []
with open(self.key_file_path, 'r') as f:
for line in f:
line = line.strip()
if line and not line.startswith('#'):
keys.append(line.split())
return keys
def has_key(self, key: str) -> bool:
"""Check if a parameter key exists."""
for k in self.keys:
if k[0] == key:
return True
return False
def get_key(self, key: str):
"""Get a parameter value by key."""
for k in self.keys:
if k[0] == key:
return k
return []
def set_key(self, key: str, value: str):
"""Set a parameter value by key."""
for k in self.keys:
if k[0] == key:
logger.info(f"Updating key in key file: `{key} {value}`")
k[1:] = value.strip().split()
break
else:
logger.info(f"Adding new key to key file: `{key} {value}`")
self.keys.append([key] + value.strip().split())
def save_key_file(self, output_path: Optional[str] = None):
"""Save the current keys to a key file."""
# if output_path is None, create a new file with untaken numbered suffix
if output_path is None:
if not os.path.exists(self.key_file_path):
path = self.key_file_path
else:
base_path = self.key_file_path[:-4] # remove .key
i = 1
while os.path.exists(f"{base_path}_{i}.key"):
i += 1
path = f"{base_path}_{i}.key"
else:
path = output_path
with open(path, 'w') as f:
for key in self.keys:
f.write(' '.join(key) + "\n")
logger.info(f"Key file saved to {path}")
self.latest_path = path
class AtomTypeFinder:
"""Class for finding atom types from Tinker parameter files."""
def __init__(self, prm_file_path: str):
self.prm_file_path = prm_file_path
self.atom_defs = self._load_atom_definitions()
self.multipole_defs = self._load_multipole_definitions()
def _load_atom_definitions(self):
"""Load and parse atom definitions from the Tinker parameter file."""
atom_defs = []
with open(self.prm_file_path, 'r') as f:
for line in f:
line = line.strip()
if line.startswith('atom'):
atom_defs.append(line)
return atom_defs
def _load_multipole_definitions(self):
"""Load and parse multipole definitions from the Tinker parameter file."""
multipole_defs = []
with open(self.prm_file_path, 'r') as f:
for line in f:
line = line.strip()
if line.startswith('multipole'):
multipole_defs.append(line)
return multipole_defs
def find_atom_type(self, description):
"""Find the atom type based on description."""
for atom_def in self.atom_defs:
parts = atom_def.split()
if len(parts) >= 4:
atom_type = parts[1]
# atomic_number = parts[5]
if description in atom_def:
return atom_type
return 0
def find_atom_charge(self, atom_type):
"""Find the atom charge based on atom type."""
for multipole_def in self.multipole_defs:
parts = multipole_def.split()
m_type = parts[1]
charge = parts[-1]
if m_type == atom_type:
return float(charge)
return None
def make_readable_ids(ids, fill_gaps_under=1):
ids = sorted(ids)
readable_ids = []
start = ids[0]
last = ids[0]
for i in ids[1:]:
if i - last <= fill_gaps_under:
last = i
continue
if last == start:
readable_ids.append(last)
else:
readable_ids.append([start, last])
start = i
last = i
# append the last one
if last == start:
readable_ids.append(last)
else:
readable_ids.append([start, last])
return readable_ids