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 if membrane_map is None or membrane_map['n_mems'] == 0: 

19 print('-> Skipping thickness analysis') 

20 return 

21 

22 # Set the main output filepath 

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

24 

25 tj, frame_step, frames_count = get_reduced_pytraj_trajectory( 

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

27 head_sel = [] 

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

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

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

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

32 # Run the analysis on the whole membrane 

33 prop = { 

34 'lipid_sel': head_sel_mda, 

35 'steps': frame_step, 

36 'height_sel': head_sel_mda, 

37 'ignore_no_box': True, 

38 'disable_logs': True, 

39 'disable_sandbox': True, 

40 } 

41 print(' Running BioBB LiPyphilic ZPositions') 

42 with redirect_stdout(None): 

43 lpp_zpositions(input_top_path=structure_file.path, 

44 input_traj_path=trajectory_file.path, 

45 output_positions_path='.zpositions.csv', 

46 properties=prop) 

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

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

49 

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

51 midplane_z = [] 

52 for i in range(0, frames_count): 

53 frame = tj[i] 

54 selected_atoms = frame[head_sel] 

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

56 midplane_z.append(float(mean_z)) 

57 # Save the data 

58 data = {'data': { 

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

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

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

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

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

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

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

66 'midplane_z': midplane_z, 

67 'step': frame_step 

68 } 

69 } 

70 save_json(data, output_analysis_filepath) 

71 

72 

73def plot_thickness(output_analysis_filepath): 

74 """Plot membrane thickness analysis results.""" 

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()