Coverage for mddb_workflow/analyses/density.py: 58%

1from mddb_workflow.utils.pyt_spells import get_reduced_pytraj_trajectory 

2from mddb_workflow.utils.auxiliar import save_json, load_json 

3from mddb_workflow.utils.constants import OUTPUT_DENSITY_FILENAME 

4from mddb_workflow.utils.type_hints import * 

5 

6import pytraj as pt 

7 

8def density ( 

9 structure_file : 'File', 

10 trajectory_file : 'File', 

11 output_directory : str, 

12 membrane_map: dict, 

13 structure : 'Structure', 

14 snapshots : int, 

15 density_types = ['number', 'mass', 'charge', 'electron'], 

16 frames_limit = 1000): 

17 """ Membrane density analysis. """ 

18 

19 if membrane_map is None or membrane_map['n_mems'] == 0: 

20 print('-> Skipping density analysis') 

21 return 

22 

23 # Set the main output filepath 

24 output_analysis_filepath = f'{output_directory}/{OUTPUT_DENSITY_FILENAME}' 

25 

26 # Load 

27 tj, frame_step, frames_count = get_reduced_pytraj_trajectory( 

28 structure_file.path, trajectory_file.path, snapshots, frames_limit) 

29 

30 # Set every selections to be analyzed separately 

31 components = [] 

32 for chain in structure.chains: 

33 components.append({ 

34 'name': chain.name, 

35 'selection': chain.get_selection(), 

36 'number': {}, 

37 'mass': {}, 

38 'charge': {}, # charge will be all 0 because we cannot add charges to pytraj topology 

39 'electron': {} 

40 }) 

41 # Parse selections to pytraj masks 

42 pytraj_masks = [ component['selection'].to_pytraj() for component in components ] 

43 # Add polar atoms selection 

44 polar_atoms = [] 

45 for n in range(membrane_map['n_mems']): 

46 polar_atoms.extend(membrane_map['mems'][str(n)]['polar_atoms']['top']) 

47 polar_atoms.extend(membrane_map['mems'][str(n)]['polar_atoms']['bot']) 

48 components.append({ 

49 'name': 'polar', 

50 'selection': polar_atoms, 

51 'number': {},'mass': {},'charge': {},'electron': {} 

52 }) 

53 pytraj_masks.append('@'+', '.join(map(str,polar_atoms))) 

54 

55 # Run pytraj 

56 for density_type in density_types: 

57 out = pt.density(tj, pytraj_masks, density_type) 

58 # Iterate pytraj results 

59 results = iter(out.values()) 

60 for component in components: 

61 # Mine pytraj results 

62 component[density_type]['dens'] = list(next(results)) 

63 component[density_type]['stdv'] = list(next(results)) 

64 

65 # Parse the selection to atom indices 

66 # Selections could be removed to make the file smaller 

67 for component in components: 

68 if component['name'] == 'polar': continue 

69 component['selection'] = component['selection'].atom_indices 

70 # Export results 

71 data = {'data': { 'comps': components, 'z': list(out['z']) } } 

72 save_json(data, output_analysis_filepath) 

73 

74def plot_density(output_analysis_filepath): 

75 """Plot density analysis grouped by density type.""" 

76 import matplotlib.pyplot as plt 

77 # Load the density analysis results 

78 data = load_json(output_analysis_filepath) 

79 

80 components = data['data']['comps'] 

81 z = data['data']['z'] 

82 

83 # Group plots by density type in a 2x2 grid 

84 density_types = ['number', 'mass', 'charge', 'electron'] 

85 fig, axes = plt.subplots(2, 2, figsize=(12, 10)) 

86 axes = axes.flatten() 

87 

88 for i, density_type in enumerate(density_types): 

89 ax = axes[i] 

90 for component in components: 

91 name = component['name'] 

92 if density_type in component: 

93 dens = component[density_type].get('dens', []) 

94 stdv = component[density_type].get('stdv', []) 

95 if dens: 

96 ax.plot(z, dens, label=f"{name}") 

97 ax.fill_between(z, [d - s for d, s in zip(dens, stdv)], 

98 [d + s for d, s in zip(dens, stdv)], alpha=0.2) 

99 

100 ax.set_title(f"{density_type.capitalize()} Density Analysis") 

101 ax.set_xlabel("Z-axis") 

102 ax.set_ylabel(f"{density_type.capitalize()} Density") 

103 ax.legend() 

104 ax.grid() 

105 

106 plt.tight_layout() 

107 plt.show()