"""
PDB Writer
==========
"""
from __future__ import annotations
import pathlib
import typing
from stk._internal.molecule import Molecule
from stk._internal.periodic_info import PeriodicInfo
from stk._internal.utilities.utilities import OneOrMany
[docs]
class PdbWriter:
"""
A writer class for ``.pdb`` files.
Examples:
*Writing to a File with a Unit Cell*
This writer can write to a file with the unit
cell included for periodic molecules. Note that this always
assumes P1 space group.
.. testcode:: writing-to-a-file-with-a-unit-cell
import stk
bb1 = stk.BuildingBlock('BrCCBr', [stk.BromoFactory()])
bb2 = stk.BuildingBlock(
smiles='BrCC(CBr)CBr',
functional_groups=[stk.BromoFactory()],
)
topology_graph = stk.cof.PeriodicHoneycomb(
building_blocks=(bb1, bb2),
lattice_size=(3, 3, 1),
)
construction_result = topology_graph.construct()
cof = (
stk.ConstructedMolecule.init_from_construction_result(
construction_result=construction_result,
)
)
writer = stk.PdbWriter()
writer.write(
molecule=cof,
path='cof.pdb',
periodic_info=construction_result.get_periodic_info(),
)
.. testcode:: writing-to-a-file-with-a-unit-cell
:hide:
import os
assert os.path.exists('cof.pdb')
.. testcleanup:: writing-to-a-file-with-a-unit-cell
os.remove('cof.pdb')
"""
def _write_content(
self,
molecule: Molecule,
atom_ids: typing.Optional[OneOrMany[int]],
periodic_info: typing.Optional[PeriodicInfo] = None,
) -> list[str]:
if atom_ids is None:
atom_ids = range(molecule.get_num_atoms())
elif isinstance(atom_ids, int):
atom_ids = (atom_ids,)
content = []
if periodic_info is not None:
# Input unit cell information.
a = periodic_info.get_a()
b = periodic_info.get_b()
c = periodic_info.get_c()
alpha = periodic_info.get_alpha()
beta = periodic_info.get_beta()
gamma = periodic_info.get_gamma()
content.append(
f"CRYST1 {a:>8.3f} {b:>8.3f} {c:>8.3f}"
f" {alpha:>6.2f} {beta:>6.2f} {gamma:>6.2f} "
f"P 1\n"
)
atom_counts: dict[str, int] = {}
hetatm = "HETATM"
alt_loc = ""
res_name = "UNL"
chain_id = ""
res_seq = "1"
i_code = ""
occupancy = "1.00"
temp_factor = "0.00"
coords = molecule.get_position_matrix()
# This set will be used by bonds.
atoms = set()
for atom_id in atom_ids:
atoms.add(atom_id)
(atom,) = molecule.get_atoms(atom_ids=atom_id)
serial = atom_id + 1
element = atom.__class__.__name__
charge = atom.get_charge()
atom_counts[element] = atom_counts.get(element, 0) + 1
name = f"{element}{atom_counts[element]}"
# Make sure the coords are no more than 8 columns wide
# each.
x, y, z = (i for i in coords[atom_id])
content.append(
f"{hetatm:<6}{serial:>5} {name:<4}"
f"{alt_loc:<1}{res_name:<3} {chain_id:<1}"
f"{res_seq:>4}{i_code:<1} "
f" {x:>7.3f} {y:>7.3f} {z:>7.3f}"
f"{occupancy:>6}{temp_factor:>6} "
f"{element:>2}{charge:>2}\n"
)
conect = "CONECT"
for bond in molecule.get_bonds():
a1 = bond.get_atom1().get_id()
a2 = bond.get_atom2().get_id()
if a1 in atoms and a2 in atoms:
content.append(
f"{conect:<6}{a1+1:>5}{a2+1:>5} \n"
)
content.append("END\n")
return content
[docs]
def to_string(
self,
molecule: Molecule,
atom_ids: typing.Optional[OneOrMany[int]] = None,
periodic_info: typing.Optional[PeriodicInfo] = None,
) -> str:
"""
Get a ``.pdb`` file format string of `molecule`.
Parameters:
molecule:
Molecule to write to ``.pdb`` format.
atom_ids:
The atom ids of atoms to write. Can be a single
:class:`int`, if a single atom is to be used, or
``None``, if all atoms are to be used.
periodic_info:
Information about the periodic cell.
Returns:
A string holding the content of a ``.pdf`` file.
"""
content = self._write_content(
molecule=molecule,
atom_ids=atom_ids,
periodic_info=periodic_info,
)
return "".join(content)
[docs]
def write(
self,
molecule: Molecule,
path: pathlib.Path | str,
atom_ids: typing.Optional[OneOrMany[int]] = None,
periodic_info: typing.Optional[PeriodicInfo] = None,
) -> None:
"""
Write `molecule` to ``.pdb`` file format.
Parameters:
molecule:
Molecule to write to ``.pdb`` format.
path:
The full path to the file being written.
atom_ids:
The atom ids of atoms to write. Can be a single
:class:`int`, if a single atom is to be used, or
``None``, if all atoms are to be used.
periodic_info:
Information about the periodic cell.
"""
content = self._write_content(
molecule=molecule,
atom_ids=atom_ids,
periodic_info=periodic_info,
)
with open(path, "w") as f:
f.write("".join(content))