Coverage for eminus/io/pdb.py: 98.39%

1# SPDX-FileCopyrightText: 2021 The eminus developers 

2# SPDX-License-Identifier: Apache-2.0 

3"""PDB file handling.""" 

4 

5import numpy as np 

6from scipy.linalg import norm 

7 

8from ..logger import log 

9from ..units import bohr2ang 

10from ..utils import vector_angle 

11 

12 

13def write_pdb(obj, filename, fods=None, elec_symbols=("X", "He"), trajectory=False): 

14 """Generate PDB files from atoms objects. 

15 

16 See :func:`~eminus.io.pdb.create_pdb_str` for more information about the PDB file format. 

17 

18 Args: 

19 obj: Atoms or SCF object. 

20 filename: PDB output file path/name. 

21 

22 Keyword Args: 

23 fods: FOD coordinates to write. 

24 elec_symbols: Identifier for up and down FODs. 

25 trajectory: Allow appending to a file to create trajectories. 

26 """ 

27 atoms = obj._atoms 

28 

29 if not filename.endswith(".pdb"): 

30 filename += ".pdb" 

31 

32 if "He" in atoms.atom and atoms.unrestricted: 

33 log.warning( 

34 'You need to modify "elec_symbols" to write helium with FODs in the spin-' 

35 "polarized case." 

36 ) 

37 

38 atom = atoms.atom 

39 pos = atoms.pos 

40 if fods is not None: 

41 if len(fods[0]) != 0: 

42 atom = atom + [elec_symbols[0]] * len(fods[0]) 

43 pos = np.vstack((pos, fods[0])) 

44 if len(fods) > 1 and len(fods[1]) != 0: 

45 atom = atom + [elec_symbols[1]] * len(fods[1]) 

46 pos = np.vstack((pos, fods[1])) 

47 

48 # Append to a file when using the trajectory keyword 

49 if trajectory: 

50 mode = "a" 

51 else: 

52 mode = "w" 

53 

54 with open(filename, mode, encoding="utf-8") as fp: 

55 fp.write(create_pdb_str(atom, pos, a=atoms.a)) 

56 

57 

58def create_pdb_str(atom, pos, a=None): 

59 """Convert atom symbols and positions to the PDB format. 

60 

61 File format definitions: 

62 CRYST1: https://wwpdb.org/documentation/file-format-content/format33/sect8.html#CRYST1 

63 

64 ATOM: https://wwpdb.org/documentation/file-format-content/format33/sect9.html#ATOM 

65 

66 Args: 

67 atom: Atom symbols. 

68 pos: Atom positions. 

69 

70 Keyword Args: 

71 a: Cell size. 

72 

73 Returns: 

74 PDB file format. 

75 """ 

76 # Convert Bohr to Angstrom 

77 pos = bohr2ang(pos) 

78 if a is not None: 

79 a = bohr2ang(a) 

80 

81 # PDB files have specific numbers of characters for every data with changing justification 

82 # Write everything explicitly down to not lose track of line lengths 

83 pdb = "" 

84 # Create data for a cuboidal cell 

85 if a is not None: 

86 pdb += "CRYST1" # 1-6 "CRYST1" 

87 pdb += f"{norm(a[0]):>9,.3f}" # 7-15 a 

88 pdb += f"{norm(a[1]):>9,.3f}" # 16-24 b 

89 pdb += f"{norm(a[2]):>9,.3f}" # 25-33 c 

90 pdb += f"{vector_angle(a[1], a[2]):>7,.2f}" # 34-40 alpha 

91 pdb += f"{vector_angle(a[0], a[2]):>7,.2f}" # 41-47 beta 

92 pdb += f"{vector_angle(a[0], a[1]):>7,.2f}" # 48-54 gamma 

93 pdb += " " 

94 pdb += "P 1 " # 56-66 Space group 

95 # pdb += " 1" # 67-70 Z value 

96 pdb += "\n" 

97 

98 # Create molecule data 

99 pdb += "MODEL 1" 

100 for ia in range(len(atom)): 

101 pdb += "\nATOM " # 1-6 "ATOM" 

102 pdb += f"{ia + 1:>5}" # 7-11 Atom serial number 

103 pdb += " " 

104 pdb += f"{atom[ia]:>4}" # 13-16 Atom name 

105 pdb += " " # 17 Alternate location indicator 

106 pdb += "MOL" # 18-20 Residue name 

107 pdb += " " 

108 pdb += " " # 22 Chain identifier 

109 pdb += " 1" # 23-26 Residue sequence number 

110 pdb += " " # 27 Code for insertions of residues 

111 pdb += " " 

112 pdb += f"{pos[ia, 0]:>8,.3f}" # 31-38 X orthogonal coordinate 

113 pdb += f"{pos[ia, 1]:>8,.3f}" # 39-46 Y orthogonal coordinate 

114 pdb += f"{pos[ia, 2]:>8,.3f}" # 47-54 Z orthogonal coordinate 

115 pdb += f"{1:>6,.2f}" # 55-60 Occupancy 

116 pdb += f"{0:>6,.2f}" # 61-66 Temperature factor 

117 pdb += " " 

118 pdb += f"{atom[ia]:>2}" # 77-78 Element symbol 

119 # pdb += " " # 79-80 Charge 

120 return f"{pdb}\nENDMDL\n"