Coverage for model_workflow/analyses/dihedral_energies.py: 14%

1import mdtraj as mdt 

2import numpy as np 

3 

4from model_workflow.tools.get_reduced_trajectory import get_reduced_trajectory 

5from model_workflow.utils.type_hints import * 

6from model_workflow.utils.auxiliar import save_json, mean 

7from model_workflow.utils.constants import OUTPUT_DIHEDRAL_ENERGIES_FILENAME 

8 

9 

10def compute_dihedral_energies ( 

11 structure_file : 'File', 

12 trajectory_file : 'File', 

13 output_directory : str, 

14 dihedrals : List[dict], 

15 snapshots : int, 

16 frames_limit : int, 

17): 

18 """Calculate torsions and then dihedral energies for every dihedral along the trajectory.""" 

19 # Set the main output filepath 

20 output_analysis_filepath = f'{output_directory}/{OUTPUT_DIHEDRAL_ENERGIES_FILENAME}' 

21 

22 # Set a dict with the results of energies calculations 

23 # Atom indices are used as keys 

24 dihedral_energies = {} 

25 for dihedral_data in dihedrals: 

26 atom_indices = dihedral_data['atom_indices'] 

27 dihedral_energies[atom_indices] = { 

28 'atom_indices': list(atom_indices), 

29 'torsion': [], 

30 'ee': [], 

31 'vdw': [], 

32 } 

33 

34 # If trajectory frames number is bigger than the limit we create a reduced trajectory 

35 reduced_trajectory_filepath, step, frames = get_reduced_trajectory( 

36 structure_file, trajectory_file, snapshots, frames_limit) 

37 

38 # Load the reduced trajectory 

39 traj = mdt.load(reduced_trajectory_filepath, top=structure_file.path) 

40 

41 # MDtraj is prepared to analyze all dihedrals at once 

42 # To do so, prepare a list with every group of atoms conforming a dihedral 

43 dihedral_atom_indices = [ data['atom_indices'] for data in dihedrals ] 

44 

45 # Compute dihedral torsions using MDtraj 

46 trajectory_torsions = mdt.compute_dihedrals(traj, dihedral_atom_indices) 

47 

48 # Results are returned by frame 

49 # DANI: Creo, porque he hecho las pruebas con una trayectoria de 1 frame xD 

50 # Iterate frames and calculate dihedral energies for each frame 

51 for frame_torsions in trajectory_torsions: 

52 

53 # Iterate every dihedral with its torsion 

54 for dihedral_data, torsion in zip(dihedrals, frame_torsions): 

55 atom_indices = dihedral_data['atom_indices'] 

56 # Dihedral torsion energy is the sum of the torsion energy of its terms 

57 torsion_energy = 0 

58 # Iterate dihedral terms 

59 for term in dihedral_data['terms']: 

60 # Get dihedral term parameters 

61 k = term['force'] 

62 if k == 0: continue 

63 n = term['period'] 

64 δ = term['phase'] 

65 φ = torsion 

66 # Calculate the dihedral torsion energy according to the formula 

67 term_energy = (k/2) * (1 + np.cos(n*φ - δ)) 

68 torsion_energy += term_energy 

69 # Add torison energies to the dihedral energies object 

70 dihedral_energies[atom_indices]['torsion'].append(torsion_energy) 

71 

72 # Now calculate non covalent energies 

73 # These energies are always computed between atoms 1-4 in the dihedral 

74 # i.e. atoms in both ends 

75 dihedral_1_4_indices = [ [ data['atom_indices'][0], data['atom_indices'][3] ] for data in dihedrals ] 

76 

77 # Compute atom distance along the trajectory using MDtraj 

78 trajectory_distances = mdt.compute_distances(traj, dihedral_1_4_indices) 

79 

80 # Results are returned by frame 

81 # Iterate frames and calculate dihedral energies for each frame 

82 for frame_distances in trajectory_distances: 

83 

84 # Iterate every dihedral with its 1-4 distance 

85 for dihedral_data, distance in zip(dihedrals, frame_distances): 

86 # Get scaling factors for both electrostatic and Van Der Waals parameters 

87 # Note that these may be missing for some dihedrals since all its terms are to be ignored 

88 # SANTI: Cuando esto pasa hay que caluclar la energía igualmente, pero sin escalarla 

89 scee = dihedral_data.get('ee_scaling', 1) 

90 scnb = dihedral_data.get('vdw_scaling', 1) 

91 # get Leonard-Johns constants to calculate Van Der Waals energies 

92 acoef = dihedral_data['lj_acoef'] 

93 bcoef = dihedral_data['lj_bcoef'] 

94 # Note that charges are not actual atom charges, but they are already scaled 

95 q1 = dihedral_data['atom1_charge'] 

96 q4 = dihedral_data['atom4_charge'] 

97 # MDtraj outputs in nm and we want distance in Angstroms 

98 r14 = distance * 10 

99 # Calculate the dihedral electrostatic energy according to the formula 

100 ee_energy = (1/scee) * (q1 * q4) / r14 

101 # Calculate the dihedral Var Der Waals energy according to the formula 

102 vdw_energy = (1/scnb) * (( acoef / (r14 ** 12) ) - ( bcoef / (r14 ** 6) )) 

103 # Add non covalent energies to the dihedral energies object 

104 atom_indices = dihedral_data['atom_indices'] 

105 dihedral_energies[atom_indices]['ee'].append(ee_energy) 

106 dihedral_energies[atom_indices]['vdw'].append(vdw_energy) 

107 

108 # Reformat output data by calculating average values instead of per-frame values 

109 output_data = [] 

110 for energies in dihedral_energies.values(): 

111 output_data.append({ 

112 'indices': energies['atom_indices'], 

113 'torsion': mean(energies['torsion']), 

114 'ee': mean(energies['ee']) if len(energies['ee']) > 0 else None, 

115 'vdw': mean(energies['vdw']) if len(energies['vdw']) > 0 else None, 

116 }) 

117 

118 # Write results to disk 

119 save_json(output_data, output_analysis_filepath)