Coverage for cvpack/base_rmsd_content.py: 99%

1""" 

2.. class:: BaseRMSDContent 

3 :platform: Linux, MacOS, Windows 

4 :synopsis: Secondary-structure RMSD content of a sequence of residues 

5 

6.. classauthor:: Charlles Abreu <craabreu@gmail.com> 

7 

8""" 

9 

10from __future__ import annotations 

11 

12import re 

13from importlib import resources 

14 

15import numpy as np 

16import openmm 

17from openmm import app as mmapp 

18 

19from .collective_variable import CollectiveVariable 

20from .rmsd import RMSD 

21from .units import ScalarQuantity, value_in_md_units 

22 

23# Limits on the number of collective variables per OpenMM CustomCVForce 

24MAX_NUM_BLOCKS = 1024 

25CHUNK_SIZE = 32 

26 

27 

28class BaseRMSDContent(CollectiveVariable, openmm.CustomCVForce): 

29 """ 

30 Abstract class for secondary-structure RMSD content of a sequence of `n` residues. 

31 """ 

32 

33 def __init__( 

34 self, 

35 residue_blocks: list[int], 

36 references: list[list[openmm.Vec3]], 

37 residues: list[mmapp.topology.Residue], 

38 numAtoms: int, 

39 thresholdRMSD: ScalarQuantity, 

40 stepFunction: str = "(1+x^4)/(1+x^4+x^8)", 

41 normalize: bool = False, 

42 ): 

43 # As in PLUMED, the RMSD of a block is the smallest among all references. 

44 # Each reference takes one collective variable per block. 

45 num_references = len(references) 

46 chunk_size = CHUNK_SIZE // num_references 

47 max_num_blocks = MAX_NUM_BLOCKS // num_references 

48 num_residue_blocks = self._num_residue_blocks = len(residue_blocks) 

49 if not 1 <= num_residue_blocks <= max_num_blocks: 

50 raise ValueError( 

51 f"{num_residue_blocks} residue blocks were defined, which is not " 

52 f"between 1 and {max_num_blocks}" 

53 ) 

54 residue_atoms = list(map(self._getAtomList, residues)) 

55 block_atoms = [ 

56 sum([residue_atoms[index] for index in block], []) 

57 for block in residue_blocks 

58 ] 

59 

60 def get_expression(start): 

61 summands = [] 

62 definitions = [] 

63 for i in range(start, min(start + chunk_size, num_residue_blocks)): 

64 summands.append(re.sub(r"\bx\b", f"x{i}", stepFunction)) 

65 rmsd = rmsd_names[i][0] 

66 for name in rmsd_names[i][1:]: 

67 rmsd = f"min({rmsd},{name})" 

68 definitions.append(f"x{i}={rmsd}/{value_in_md_units(thresholdRMSD)}") 

69 summation = "+".join(summands) 

70 if normalize: 

71 summation = f"({summation})/{num_residue_blocks}" 

72 return ";".join([summation] + definitions) 

73 

74 rmsd_names = [ 

75 [f"rmsd{i}"] 

76 if num_references == 1 

77 else [f"rmsd{i}_{k}" for k in range(num_references)] 

78 for i in range(num_residue_blocks) 

79 ] 

80 if num_residue_blocks <= chunk_size: 

81 expression = get_expression(0) 

82 force = self 

83 else: 

84 expression = "+".join( 

85 f"chunk{i}" 

86 for i in range((num_residue_blocks + chunk_size - 1) // chunk_size) 

87 ) 

88 super().__init__(expression) 

89 for index in range(num_residue_blocks): 

90 if num_residue_blocks > chunk_size and index % chunk_size == 0: 

91 force = openmm.CustomCVForce(get_expression(index)) 

92 self.addCollectiveVariable(f"chunk{index // chunk_size}", force) 

93 for name, reference in zip(rmsd_names[index], references): 

94 force.addCollectiveVariable( 

95 name, 

96 RMSD( 

97 dict(zip(block_atoms[index], reference)), 

98 block_atoms[index], 

99 numAtoms, 

100 ), 

101 ) 

102 

103 @classmethod 

104 def _loadPositions(cls, filename: str) -> list[openmm.Vec3]: 

105 positions = 0.1 * np.loadtxt( 

106 str(resources.files("cvpack").joinpath("data").joinpath(filename)), 

107 delimiter=",", 

108 ) 

109 return [openmm.Vec3(*position) for position in positions] 

110 

111 @staticmethod 

112 def _getAtomList(residue: mmapp.topology.Residue) -> list[int]: 

113 residue_atoms = {atom.name: atom.index for atom in residue.atoms()} 

114 if residue.name == "GLY": 

115 # The hydrogen on the CB side is HA3 in PDB/Amber naming (HA2, HA3) and 

116 # HA1 in CHARMM naming (HA1, HA2), as in OpenMM's pdbNames.xml 

117 hydrogen = next((h for h in ("HA3", "HA1") if h in residue_atoms), None) 

118 if hydrogen is None: 

119 raise ValueError( 

120 f"Atom HA3 or HA1 not found in residue {residue.name}{residue.id}" 

121 ) 

122 residue_atoms["CB"] = residue_atoms[hydrogen] 

123 if "O" not in residue_atoms: 

124 # A C-terminal residue has OT1 in CHARMM naming and OC1 in GROMACS naming 

125 oxygen = next((o for o in ("OT1", "OC1") if o in residue_atoms), None) 

126 if oxygen is not None: 

127 residue_atoms["O"] = residue_atoms[oxygen] 

128 atom_list = [] 

129 for atom in ("N", "CA", "CB", "C", "O"): 

130 try: 

131 atom_list.append(residue_atoms[atom]) 

132 except KeyError as error: 

133 raise ValueError( 

134 f"Atom {atom} not found in residue {residue.name}{residue.id}" 

135 ) from error 

136 return atom_list 

137 

138 def getNumResidueBlocks(self) -> int: 

139 """ 

140 Get the number of residue blocks. 

141 

142 Returns 

143 ------- 

144 The number of residue blocks. 

145 """ 

146 return self._num_residue_blocks