Coverage for model_workflow/analyses/lipid_interactions.py: 100%

1from model_workflow.tools.get_reduced_trajectory import calculate_frame_step 

2from model_workflow.utils.mda_spells import to_MDAnalysis_topology 

3from model_workflow.utils.auxiliar import save_json 

4from model_workflow.utils.constants import OUTPUT_LIPID_INTERACTIONS_FILENAME 

5from model_workflow.utils.type_hints import * 

6from collections import Counter 

7import numpy as np 

8import MDAnalysis 

9 

10def lipid_interactions ( 

11 trajectory_file : 'File', 

12 standard_topology_file : 'File', 

13 output_directory : str, 

14 membrane_map: dict, 

15 lipid_map: dict, 

16 snapshots : int, 

17 frames_limit: int = 100): 

18 """ 

19 Lipid-protein interactions analysis. 

20 """ 

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

22 print('-> Skipping lipid-protein interactions analysis') 

23 return 

24 

25 # Set the main output filepath 

26 output_analysis_filepath = f'{output_directory}/{OUTPUT_LIPID_INTERACTIONS_FILENAME}' 

27 

28 mda_top = to_MDAnalysis_topology(standard_topology_file) 

29 u = MDAnalysis.Universe(mda_top, trajectory_file.path) 

30 frame_step, frame_count = calculate_frame_step(snapshots, frames_limit) 

31 lipids = set([ lipid['resname'] for lipid in lipid_map]) 

32 lipids_str = " ".join(lipids) 

33 resids = np.unique(u.select_atoms('protein').resindices) 

34 ocupancy = {resid: Counter() for resid in resids} 

35 

36 # Only iterate through the frames you need 

37 for ts in u.trajectory[0:snapshots:frame_step]: 

38 # Select non-protein atoms near the protein once 

39 non_protein_near = u.select_atoms(f'(around 6 protein) and (resname {lipids_str}) and not protein') 

40 # Make residue selections only once 

41 for resid in resids: 

42 residue_atoms = u.select_atoms(f'resid {resid}') 

43 # Find non-protein atoms near this specific residue 

44 nearby = non_protein_near.select_atoms(f'around 6 global group residuegroup', 

45 residuegroup=residue_atoms) 

46 # Count residue names 

47 ocupancy[resid].update(Counter(nearby.residues.resnames)) 

48 

49 ocupancy_arrs = {lipid: np.zeros(len(ocupancy)) for lipid in lipids} 

50 for resid, counter in ocupancy.items(): 

51 for lipid, count in counter.items(): 

52 ocupancy_arrs[lipid][resid] = count 

53 

54 # Normalize the occupancy arrays by dividing by the number of frames 

55 for lipid in lipids: 

56 ocupancy_arrs[lipid] /= frame_count 

57 ocupancy_arrs[lipid] = ocupancy_arrs[lipid].tolist() 

58 # Save the data 

59 data = { 'data': ocupancy_arrs} 

60 save_json(data, output_analysis_filepath)