Coverage for mddb_workflow/utils/mda_spells.py: 62%

1import json 

2import numpy as np 

3import pickle 

4 

5from mddb_workflow.utils.type_hints import * 

6from mddb_workflow.utils.constants import GREY_HEADER, COLOR_END 

7from mddb_workflow.utils.auxiliar import MISSING_CHARGES, MISSING_BONDS 

8 

9from MDAnalysis.topology.TPRParser import TPRParser 

10#from MDAnalysis.topology.PDBParser import PDBParser # for class reference 

11from MDAnalysis.core.universe import Universe 

12from MDAnalysis.core.topology import Topology 

13from MDAnalysis.core.topologyattrs import ( 

14 Atomnames, 

15 Bonds, 

16 ChainIDs, 

17 Charges, 

18 Elements, 

19 Resids, 

20 Resnames, 

21 Resnums, 

22 Segids 

23) 

24 

25def to_MDAnalysis_topology(standard_topology_path : str) -> 'Topology': 

26 """ 

27 Creates a MDAnalysis topology from a json topology file. 

28 

29 The json file should contain the following keys: 

30 - atom_names 

31 - atom_elements 

32 - atom_charges 

33 - atom_bonds (list of lists of atom indices) 

34 - residue_names 

35 - residue_numbers 

36 - chain_names 

37 - atom_residue_indices 

38 - residue_chain_indices 

39 

40 :param topology: path to the json file 

41 :returns: a MDAnalysis topology object 

42 """ 

43 topology = json.load(open(standard_topology_path)) 

44 

45 # transform bond to non redundant tuples 

46 bonds = [] 

47 for bond_from, bond_tos in enumerate(topology['atom_bonds']): 

48 for bond_to in bond_tos: 

49 bond = tuple(sorted([bond_from, bond_to])) 

50 bonds.append(bond) 

51 # sorted(set(bonds)) if you want to check for duplicates 

52 topology['bonds'] = set(bonds) 

53 atom_2_chain = np.array(topology['residue_chain_indices'])[topology['atom_residue_indices']] 

54 topology['chainIDs'] = np.array(topology['chain_names'])[atom_2_chain] 

55 

56 

57 attrs = [Atomnames (topology['atom_names']), 

58 Elements (topology['atom_elements']), 

59 Charges (topology['atom_charges']), 

60 Bonds (topology['bonds']), 

61 Resnames (topology['residue_names']), 

62 Resnums (topology['residue_numbers']), 

63 ChainIDs (topology['chainIDs']), 

64 Segids (topology['chain_names']), 

65 Resids (np.arange(len(topology['residue_names']))), 

66 ] 

67 

68 mda_top = Topology(n_atoms=len(topology['atom_names']), 

69 n_res=len(topology['residue_names']), 

70 n_seg=len(topology['chain_names']), 

71 attrs=attrs, 

72 atom_resindex=topology['atom_residue_indices'], 

73 residue_segindex=topology['residue_chain_indices'] 

74 ) 

75 return mda_top 

76 

77def get_mda_universe_from_stopology ( 

78 standard_topology_path : str, coordinates_file : str) -> 'Universe': 

79 """Create a MDAnalysis universe using data in the workflow. 

80  

81 Args: 

82 standard_topology_path (str): Path to the standard topology file. 

83 coordinates_file (str): Path to the coordinates file (e.g., PDB, XTC).""" 

84 mda_topology = to_MDAnalysis_topology(standard_topology_path) 

85 # Create a MDAnalysis topology from the standard topology file 

86 return Universe(mda_topology, coordinates_file) 

87 

88def get_mda_universe (structure_file : 'File', # To load in MDAnalysis 

89 trajectory_file : 'File', # To load in MDAnalysis 

90 reference_bonds : list[list[int]], # To set the bonds 

91 charges : list[float]) -> 'Universe': # To set the charges 

92 """Create a MDAnalysis universe using data in the workflow.""" 

93 

94 # Make MDAnalysis warnings and logs grey 

95 print(GREY_HEADER, end='\r') 

96 universe = Universe(structure_file.path, trajectory_file.path) 

97 # Set the atom bonds 

98 bonds = [] 

99 for bond_from, bond_tos in enumerate(reference_bonds): 

100 if bond_tos == MISSING_BONDS: continue 

101 for bond_to in bond_tos: 

102 bond = tuple(sorted([bond_from, bond_to])) 

103 bonds.append(bond) 

104 universe.add_TopologyAttr('bonds', set(bonds)) 

105 

106 # Set the charges 

107 if charges and charges != MISSING_CHARGES and len(charges) > 0: 

108 universe.add_TopologyAttr('charges', np.array(charges, dtype=np.float32)) 

109 print(COLOR_END, end='\r') 

110 return universe 

111 

112# Get a cksum from a MDA universe for equality comparission 

113def get_mda_universe_cksum (universe) -> str: 

114 pickled = pickle.dumps(universe.atoms) 

115 return sum(pickled) 

116 

117# Get TPR bonds using MDAnalysis 

118# WARNING: Sometimes this function takes additional constrains as actual bonds 

119# DANI: si miras los topology.bonds.values estos enlaces falsos van al final 

120# DANI: Lo veo porque los índices están en orden ascendente y veulven a empezar 

121# DANI: He pedido ayuda aquí https://github.com/MDAnalysis/mdanalysis/pull/463 

122def get_tpr_bonds_mdanalysis (tpr_filepath : str) -> list[ tuple[int, int] ]: 

123 parser = TPRParser(tpr_filepath) 

124 topology = parser.parse() 

125 bonds = list(topology.bonds.values) 

126 return bonds