Coverage for mddb_workflow / analyses / channels.py: 96%

1from biobb_mem.mdanalysis_biobb.mda_hole import mda_hole 

2from mddb_workflow.tools.get_reduced_trajectory import calculate_frame_step 

3from mddb_workflow.utils.auxiliar import save_json 

4from mddb_workflow.utils.constants import OUTPUT_CHANNELS_FILENAME 

5from mddb_workflow.utils.type_hints import * 

6import re 

7import numpy as np 

8 

9 

10def channels( 

11 structure_file: 'File', 

12 trajectory_file: 'File', 

13 output_directory: str, 

14 membrane_map: dict, 

15 snapshots: int, 

16 frames_limit: int 

17): 

18 """Analyze channels in a membrane protein using MDAnalysis' mda_hole.""" 

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

20 print('-> Skipping channels analysis') 

21 return 

22 

23 frame_step, frame_count = calculate_frame_step(snapshots, frames_limit) 

24 

25 prop = { 

26 'select': 'protein', 

27 'steps': frame_step, 

28 'sample': 0.2, 

29 'dot_density': 13, 

30 'disable_logs': True, 

31 'disable_sandbox': True, 

32 } 

33 

34 mda_hole(input_top_path=structure_file.path, 

35 input_traj_path=trajectory_file.path, 

36 output_hole_path=output_directory + '/hole.vmd', 

37 output_csv_path=output_directory + '/mda.hole_profile.csv', 

38 properties=prop) 

39 

40 with open(output_directory + '/hole.vmd', 'r') as f: 

41 lines = f.readlines() 

42 

43 # Find lines with triangle coordinates 

44 trinorms = [] 

45 for i, line in enumerate(lines): 

46 if i > 3 and 'set triangle' in line: 

47 vmd_set = re.sub(r'set triangles\(\d+\)', '', line) # Remove set triangles(i) 

48 vmd_set = re.sub(r'\{(\s*-?\d[^\s]*)(\s*-?\d[^\s]*)(\s*-?\d[^}]*)\}', r'[\1,\2,\3]', vmd_set) # Convert { x y z } to [x,y,z] 

49 vmd_set = vmd_set.replace('{', '[').replace('}', ']') # Convert { to [ and } to ] 

50 vmd_set = re.sub(r'\]\s*\[', '], [', vmd_set) # Add commas between brackets 

51 vmd_set = eval(vmd_set.strip()) # Evaluate string as list 

52 # different hole colors 

53 trinorms.append(vmd_set) 

54 # Create a list of positions, colors, and normals 

55 assert frame_count == len(trinorms), f'Frame count {frame_count} does not match trinorms length {len(trinorms)}' 

56 data = {} 

57 for frame in range(frame_count): 

58 poss = [] 

59 cols = [0, 0, 0] 

60 z_range = 0 

61 for i in range(3): # RGB 

62 # Get triangle coordinates for this frame and color 

63 # Red are low radiues, green medium, blue high 

64 trinorms_cl = np.array(trinorms[frame][i]) # (N, 6, 3) # 6: 3 positions + 3 normals, 3: vertex 

65 if len(trinorms_cl) == 0: 

66 continue 

67 pos = trinorms_cl[:, :3, :] 

68 z_range = max(pos[..., 2].max() - pos[..., 2].min(), z_range) 

69 poss.append(pos.flatten()) 

70 cols[i] = pos.size 

71 poss = np.concatenate(poss) 

72 if z_range < 40: 

73 # failed to find channel on this frame 

74 continue 

75 # compute cumulative offsets for colors 

76 data[frame] = { 

77 'position': poss.tolist(), 

78 'color_offset': np.cumsum(cols).tolist(), 

79 } 

80 if len(data) == 0: 

81 print('-> No channels found') 

82 return 

83 data_to_save = {'data': data, 'n_frames': frame_count} 

84 save_json(data_to_save, output_directory + '/' + OUTPUT_CHANNELS_FILENAME)