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

1from model_workflow.utils.pyt_spells import get_reduced_pytraj_trajectory 

2from model_workflow.utils.auxiliar import save_json 

3from model_workflow.utils.constants import OUTPUT_DENSITY_FILENAME 

4from model_workflow.utils.type_hints import * 

5 

6import pytraj as pt 

7 

8def density ( 

9 structure_file : 'File', 

10 trajectory_file : 'File', 

11 output_directory : str, 

12 membrane_map: dict, 

13 structure : 'Structure', 

14 snapshots : int, 

15 density_types = ['number', 'mass', 'charge', 'electron'], 

16 frames_limit = 1000): 

17 """Membrane density analysis.""" 

18 

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

20 print('-> Skipping density analysis') 

21 return 

22 

23 # Set the main output filepath 

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

25 

26 # Load 

27 tj, frame_step, frames_count = get_reduced_pytraj_trajectory( 

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

29 

30 # Set every selections to be analyzed separately 

31 components = [] 

32 for chain in structure.chains: 

33 components.append({ 

34 'name': chain.name, 

35 'selection': chain.get_selection(), 

36 'number': {}, 

37 'mass': {}, 

38 'charge': {}, # charge will be all 0 because we cannot add charges to pytraj topology 

39 'electron': {} 

40 }) 

41 # Parse selections to pytraj masks 

42 pytraj_masks = [ component['selection'].to_pytraj() for component in components ] 

43 # Add polar atoms selection 

44 polar_atoms = [] 

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

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

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

48 components.append({ 

49 'name': 'polar', 

50 'selection': polar_atoms, 

51 'number': {},'mass': {},'charge': {},'electron': {} 

52 }) 

53 pytraj_masks.append('@'+', '.join(map(str,polar_atoms))) 

54 

55 # Run pytraj 

56 for density_type in density_types: 

57 out = pt.density(tj, pytraj_masks, density_type) 

58 # Iterate pytraj results 

59 results = iter(out.values()) 

60 for component in components: 

61 # Mine pytraj results 

62 component[density_type]['dens'] = list(next(results)) 

63 component[density_type]['stdv'] = list(next(results)) 

64 

65 # Parse the selection to atom indices 

66 # Selections could be removed to make the file smaller 

67 for component in components: 

68 if component['name'] == 'polar': continue 

69 component['selection'] = component['selection'].atom_indices 

70 # Export results 

71 data = {'data': { 'comps': components, 'z': list(out['z']) } } 

72 save_json(data, output_analysis_filepath)