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 

9def channels ( 

10 structure_file : 'File', 

11 trajectory_file : 'File', 

12 output_directory : str, 

13 membrane_map: dict, 

14 snapshots: int, 

15 frames_limit: int): 

16 

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

18 print('-> Skipping channels analysis') 

19 return 

20 

21 frame_step, frame_count = calculate_frame_step(snapshots, frames_limit) 

22 

23 prop = { 

24 'select': 'protein', 

25 'steps': frame_step, 

26 'sample': 0.2, 

27 'dot_density': 13, 

28 'disable_logs': True, 

29 'disable_sandbox': True, 

30 } 

31 

32 mda_hole(input_top_path=structure_file.path, 

33 input_traj_path=trajectory_file.path, 

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

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

36 properties=prop) 

37 

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

39 lines = f.readlines() 

40 

41 # Find lines with triangle coordinates 

42 trinorms = [] 

43 for i, line in enumerate(lines): 

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

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

46 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] 

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

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

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

50 # different hole colors 

51 trinorms.append(vmd_set) 

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

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

54 data = {} 

55 for frame in range(frame_count): 

56 poss = [] 

57 cols = [0,0,0] 

58 z_range = 0 

59 for i in range(3): # RGB 

60 # Get triangle coordinates for this frame and color 

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

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

63 if len(trinorms_cl) == 0: 

64 continue 

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

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

67 poss.append(pos.flatten()) 

68 cols[i] = pos.size 

69 poss = np.concatenate(poss) 

70 if z_range < 40: 

71 # failed to find channel on this frame 

72 continue 

73 # compute cumulative offsets for colors 

74 data[frame] = { 

75 'position': poss.tolist(), 

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

77 } 

78 if len(data) == 0: 

79 print('-> No channels found') 

80 return 

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

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