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

1# SPDX-FileCopyrightText: 2023 The eminus developers 

2# SPDX-License-Identifier: Apache-2.0 

3"""PDB file handling.""" 

4 

5from eminus import backend as xp 

6from eminus.logger import log 

7from eminus.units import bohr2ang 

8from eminus.utils import vector_angle 

9 

10 

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

12 """Generate PDB files from atoms objects. 

13 

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

15 

16 Args: 

17 obj: Atoms or SCF object. 

18 filename: PDB output file path/name. 

19 

20 Keyword Args: 

21 fods: FOD coordinates to write. 

22 elec_symbols: Identifier for up and down FODs. 

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

24 """ 

25 atoms = obj._atoms 

26 

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

28 filename += ".pdb" 

29 

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

31 log.warning( 

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

33 "polarized case." 

34 ) 

35 

36 atom = atoms.atom 

37 pos = atoms.pos 

38 if fods is not None: 

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

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

41 pos = xp.vstack((pos, fods[0])) 

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

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

44 pos = xp.vstack((pos, fods[1])) 

45 

46 # Append to a file when using the trajectory keyword 

47 if trajectory: 

48 mode = "a" 

49 else: 

50 mode = "w" 

51 

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

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

54 

55 

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

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

58 

59 File format definitions: 

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

61 

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

63 

64 Args: 

65 atom: Atom symbols. 

66 pos: Atom positions. 

67 

68 Keyword Args: 

69 a: Cell size. 

70 

71 Returns: 

72 PDB file format. 

73 """ 

74 # Convert Bohr to Angstrom 

75 pos = bohr2ang(pos) 

76 if a is not None: 

77 a = bohr2ang(a) 

78 

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

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

81 pdb = "" 

82 # Create data for a cuboidal cell 

83 if a is not None: 

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

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

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

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

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

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

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

91 pdb += " " 

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

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

94 pdb += "\n" 

95 

96 # Create molecule data 

97 pdb += "MODEL 1" 

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

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

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

101 pdb += " " 

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

103 pdb += " " # 17 Alternate location indicator 

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

105 pdb += " " 

106 pdb += " " # 22 Chain identifier 

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

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

109 pdb += " " 

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

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

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

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

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

115 pdb += " " 

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

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

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