Coverage for model_workflow/utils/mda_spells.py: 100%

1import json 

2import numpy as np 

3import pickle 

4 

5from model_workflow.utils.type_hints import * 

6 

7#from MDAnalysis.topology.PDBParser import PDBParser # for class reference 

8from MDAnalysis.core.universe import Universe 

9from MDAnalysis.core.topology import Topology 

10from MDAnalysis.core.topologyattrs import ( 

11 Atomnames, 

12 Bonds, 

13 ChainIDs, 

14 Charges, 

15 Elements, 

16 Resids, 

17 Resnames, 

18 Resnums, 

19 Segids 

20) 

21 

22def to_MDAnalysis_topology(standard_topology_file : 'File') -> 'Topology': 

23 """ 

24 Creates a MDAnalysis topology from a json topology file. 

25 

26 The json file should contain the following keys: 

27 - atom_names 

28 - atom_elements 

29 - atom_charges 

30 - atom_bonds (list of lists of atom indices) 

31 - residue_names 

32 - residue_numbers 

33 - chain_names 

34 - atom_residue_indices 

35 - residue_chain_indices 

36 

37 :param topology: path to the json file 

38 :returns: a MDAnalysis topology object 

39 """ 

40 topology = json.load(open(standard_topology_file.path)) 

41 

42 # transform bond to non redundant tuples 

43 bonds = [] 

44 for bond_from, bond_tos in enumerate(topology['atom_bonds']): 

45 for bond_to in bond_tos: 

46 bond = tuple(sorted([bond_from, bond_to])) 

47 bonds.append(bond) 

48 # sorted(set(bonds)) if you want to check for duplicates 

49 topology['bonds'] = set(bonds) 

50 atom_2_chain = np.array(topology['residue_chain_indices'])[topology['atom_residue_indices']] 

51 topology['chainIDs'] = np.array(topology['chain_names'])[atom_2_chain] 

52 

53 

54 attrs = [Atomnames (topology['atom_names']), 

55 Elements (topology['atom_elements']), 

56 Charges (topology['atom_charges']), 

57 Bonds (topology['bonds']), 

58 Resnames (topology['residue_names']), 

59 Resnums (topology['residue_numbers']), 

60 ChainIDs (topology['chainIDs']), 

61 Segids (topology['chain_names']), 

62 Resids (np.arange(len(topology['residue_names']))), 

63 ] 

64 

65 mda_top = Topology(n_atoms=len(topology['atom_names']), 

66 n_res=len(topology['residue_names']), 

67 n_seg=len(topology['chain_names']), 

68 attrs=attrs, 

69 atom_resindex=topology['atom_residue_indices'], 

70 residue_segindex=topology['residue_chain_indices'] 

71 ) 

72 return mda_top 

73 

74def get_mda_universe (standard_topology_file : 'File', structure_file : 'File') -> 'Universe': 

75 """Create a MDAnalysis universe using data in the workflow.""" 

76 mda_topology = to_MDAnalysis_topology(standard_topology_file) 

77 # Create a MDAnalysis topology from the standard topology file 

78 return Universe(mda_topology, structure_file.path) 

79 

80# Get a cksum from a MDA universe for equality comparission 

81def get_mda_universe_cksum (universe) -> str: 

82 pickled = pickle.dumps(universe.atoms) 

83 return sum(pickled)