Coverage for model_workflow/tools/get_charges.py: 65%

1import pytraj as pt 

2 

3from json import load 

4 

5from model_workflow.utils.auxiliar import MISSING_TOPOLOGY, MISSING_CHARGES, load_json 

6from model_workflow.utils.constants import STANDARD_TOPOLOGY_FILENAME, RAW_CHARGES_FILENAME 

7from model_workflow.utils.gmx_spells import get_tpr_charges as get_tpr_charges_gromacs 

8from model_workflow.utils.type_hints import * 

9 

10from MDAnalysis.topology.TPRParser import TPRParser 

11from MDAnalysis.topology.TOPParser import TOPParser 

12 

13def get_charges (topology_file : Union['File', Exception], 

14 resorted_charges_file : Optional['File'] = None) -> List[float]: 

15 """ 

16 Extract charges from a source file. 

17 

18 Returns: 

19 List[float]: A list of atomic charges if extraction is successful,  

20 otherwise None if the file does not exist. 

21 

22 """ 

23 # If we have a resorted file then use it 

24 # Note that this is very excepcional 

25 if resorted_charges_file and resorted_charges_file.exists: 

26 print(' Using resorted atom charges') 

27 return load_json(resorted_charges_file.path) 

28 # If there is no topology at all 

29 if topology_file == MISSING_TOPOLOGY or not topology_file.exists: 

30 print(' No charges source file available') 

31 return MISSING_CHARGES 

32 print(f' Charges in the "{topology_file.path}" file will be used') 

33 charges = None 

34 # If we have the standard topology then get charges from it 

35 if topology_file.filename == STANDARD_TOPOLOGY_FILENAME: 

36 with open(topology_file.path, 'r') as file: 

37 standard_topology = load(file) 

38 charges = standard_topology['atom_charges'] 

39 # In some ocasions, charges may come inside a raw charges file 

40 elif topology_file.filename == RAW_CHARGES_FILENAME: 

41 charges = get_raw_charges(topology_file.path) 

42 # In some ocasions, charges may come inside a topology which can be parsed through pytraj 

43 elif topology_file.is_pytraj_supported(): 

44 charges = get_topology_charges(topology_file.path) 

45 # DANI: De momento ya no generaré más charges.txt ahora que las cargas estan en la topologia json 

46 #generate_raw_energies_file(charges) 

47 elif topology_file.format == 'tpr': 

48 charges = get_tpr_charges(topology_file.path) 

49 else: 

50 raise ValueError(f'Charges source file ({topology_file.path}) is in a non supported format') 

51 return charges 

52 

53# Given a topology which includes charges 

54# Extract those charges and save them in a list to be returned 

55# Use different tools, since ones may fail where others success 

56# Note that this not a matter of the format only, but also of the format version 

57# There are different versions of .top files, for instance 

58def get_topology_charges (topology_filename : str) -> list: 

59 try: 

60 topology_charges = get_topology_charges_pytraj(topology_filename) 

61 except Exception as err: 

62 print(err) 

63 print('The canonical charges mining (pytraj) failed. Retrying with alternative mining (mdanalysis)') 

64 topology_charges = get_topology_charges_mdanalysis(topology_filename) 

65 return topology_charges 

66 

67# Get topology charges using pytraj 

68# Supported formats (tested): prmtop, top, psf (standard psf, not from DESRES) 

69# Supported formats (not tested): mol2, cif, sdf 

70# Non supported formats: mae, tpr, pdb (has no charges) 

71def get_topology_charges_pytraj (topology_filename : str) -> list: 

72 topology = pt.load_topology(filename=topology_filename) 

73 # WARNING: We must convert this numpy ndarray to a normal list 

74 # Otherwise the search by index is extremly ineficient 

75 topology_charges = list(topology.charge) 

76 return topology_charges 

77 

78# Get topology charges using mdanalysis 

79def get_topology_charges_mdanalysis (topology_filename : str) -> list: 

80 parser = TOPParser(topology_filename) 

81 topology = parser.parse() 

82 charges = list(topology.charges.values) 

83 return charges 

84 

85# Write the raw charges file from a list of charges 

86def generate_raw_energies_file (charges : list, filename : str = RAW_CHARGES_FILENAME): 

87 with open(filename, 'w') as file: 

88 for charge in charges: 

89 file.write("{:.6f}".format(charge) + '\n') 

90 

91# Given a raw file with listed charges 

92# Extract those charges and save them in a list to be returned 

93def get_raw_charges (topology_filename : str) -> list: 

94 charges = [] 

95 with open(topology_filename, 'r') as file: 

96 lines = file.readlines() 

97 for line in lines: 

98 charges.append(float(line)) 

99 return charges 

100 

101# Given a tpr file, extract charges in a list 

102# Try 2 different methods and 1 of them should work 

103def get_tpr_charges (topology_filename : str) -> list: 

104 try: 

105 charges = get_tpr_charges_mdanalysis(topology_filename) 

106 except: 

107 print(' MDAnalysis failed to extract charges. Using manual extraction...') 

108 charges = get_tpr_charges_gromacs(topology_filename) 

109 return charges 

110 

111# This works for the old tpr format (tested in 112) 

112def get_tpr_charges_mdanalysis (topology_filename : str) -> list: 

113 parser = TPRParser(topology_filename) 

114 topology = parser.parse() 

115 charges = list(topology.charges.values) 

116 return charges