Coverage for mddb_workflow/analyses/thickness.py: 62%

1from biobb_mem.lipyphilic_biobb.lpp_zpositions import lpp_zpositions, frame_df 

2from mddb_workflow.utils.pyt_spells import get_reduced_pytraj_trajectory 

3from mddb_workflow.utils.auxiliar import save_json, load_json 

4from mddb_workflow.utils.constants import OUTPUT_THICKNESS_FILENAME 

5from mddb_workflow.utils.type_hints import * 

6from contextlib import redirect_stdout 

7import os 

8 

9 

10def thickness ( 

11 structure_file : 'File', 

12 trajectory_file : 'File', 

13 output_directory : str, 

14 membrane_map: dict, 

15 snapshots : int, 

16 frames_limit: int = 100): 

17 """Membrane thickness analysis.""" 

18 

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

20 print('-> Skipping thickness analysis') 

21 return 

22 

23 # Set the main output filepath 

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

25 

26 tj, frame_step, frames_count = get_reduced_pytraj_trajectory( 

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

28 head_sel = [] 

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

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

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

32 head_sel_mda = 'index ' + " ".join(map(str,(head_sel))) 

33 # Run the analysis on the whole membrane 

34 prop = { 

35 'lipid_sel': head_sel_mda, 

36 'steps': frame_step, 

37 'height_sel': head_sel_mda, 

38 'ignore_no_box': True, 

39 'disable_logs': True, 

40 'disable_sandbox': True, 

41 } 

42 print(' Running BioBB LiPyphilic ZPositions') 

43 with redirect_stdout(None): 

44 lpp_zpositions(input_top_path=structure_file.path, 

45 input_traj_path=trajectory_file.path, 

46 output_positions_path='.zpositions.csv', 

47 properties=prop) 

48 df = frame_df('.zpositions.csv') # Per frame data 

49 os.remove('.zpositions.csv') 

50 

51 # Calculate the mean z position of midplane wrt the box axis using pytraj 

52 midplane_z = [] 

53 for i in range(0, frames_count): 

54 frame = tj[i] 

55 selected_atoms = frame[head_sel] 

56 mean_z = selected_atoms.mean(axis=0)[2] 

57 midplane_z.append(float(mean_z)) 

58 # Save the data 

59 data = { 'data':{ 

60 'frame': df.index.tolist(), 

61 'mean_positive': df['mean_positive'].tolist(), 

62 'mean_negative': df['mean_negative'].tolist(), 

63 'std_positive': df['std_positive'].tolist(), 

64 'std_negative': df['std_negative'].tolist(), 

65 'thickness': df['thickness'].tolist(), 

66 'std_thickness': df['std_thickness'].tolist(), 

67 'midplane_z': midplane_z, 

68 'step': frame_step 

69 } 

70 } 

71 save_json(data, output_analysis_filepath) 

72 

73 

74def plot_thickness(output_analysis_filepath): 

75 import matplotlib.pyplot as plt 

76 

77 data = load_json(output_analysis_filepath)["data"] 

78 

79 # Extract data for plotting 

80 frames = data['frame'] 

81 thickness = data['thickness'] 

82 std_thickness = data['std_thickness'] 

83 midplane_z = data['midplane_z'] 

84 mean_positive = data['mean_positive'] 

85 mean_negative = data['mean_negative'] 

86 std_positive = data['std_positive'] 

87 std_negative = data['std_negative'] 

88 

89 # Create the plot 

90 plt.figure(figsize=(12, 8)) 

91 

92 # Plot thickness with error bars 

93 plt.errorbar(frames, thickness, yerr=std_thickness, label='Thickness', fmt='-o', capsize=3) 

94 

95 # Plot midplane z position 

96 plt.plot(frames, midplane_z, label='Midplane Z', linestyle='--', color='orange') 

97 

98 # Plot mean positive and mean negative with error bars 

99 eb1= plt.errorbar(frames, mean_positive, yerr=std_positive, label='Mean Positive', fmt='-o', color='green', capsize=3) 

100 eb1[-1][0].set_linestyle('--') 

101 eb2 = plt.errorbar(frames, mean_negative, yerr=std_negative, label='Mean Negative', fmt='-o', color='red', capsize=3) 

102 eb2[-1][0].set_linestyle('--') 

103 

104 # Add labels, legend, and title 

105 plt.xlabel('Frame'); plt.ylabel('Value') 

106 plt.title('Membrane Thickness, Midplane Z Position, and Mean Positive/Negative with Errors') 

107 plt.legend(); plt.grid(); plt.show()