Source code for stk._internal.topology_graphs.polymer.vertices

import logging

import numpy as np

from stk._internal.building_block import BuildingBlock
from stk._internal.topology_graphs.vertex import Vertex

from ..edge import Edge

logger = logging.getLogger(__name__)


[docs] class LinearVertex(Vertex): """ Represents a vertex in the middle of a linear polymer chain. """ def __init__( self, id: int, position: tuple[float, float, float] | np.ndarray, flip: bool, ) -> None: """ Initialize a :class:`.LinearVertex` instance. Parameters: id: The id of the vertex. position: The position of the vertex. flip: If ``True`` any building block placed by the vertex will have its orientation along the chain flipped. """ super().__init__(id, position) self._flip = flip
[docs] def get_flip(self) -> bool: """ Return whether the vertex flips building blocks it places. Returns: ``True`` if the vertex flips building blocks it places. """ return self._flip
[docs] def clone(self): clone = super().clone() clone._flip = self._flip return clone
[docs] def place_building_block( self, building_block: BuildingBlock, edges: tuple[Edge, ...], ) -> np.ndarray: assert building_block.get_num_functional_groups() == 2, ( f"{building_block} needs to have exactly 2 functional " "groups but has " f"{building_block.get_num_functional_groups()}." ) building_block = building_block.with_centroid( position=self._position, atom_ids=building_block.get_placer_ids(), ) fg1, fg2 = building_block.get_functional_groups() fg1_position = building_block.get_centroid( atom_ids=fg1.get_placer_ids(), ) fg2_position = building_block.get_centroid( atom_ids=fg2.get_placer_ids(), ) return building_block.with_rotation_between_vectors( start=fg2_position - fg1_position, target=np.array([-1 if self._flip else 1, 0, 0]), origin=self._position, ).get_position_matrix()
[docs] def map_functional_groups_to_edges( self, building_block: BuildingBlock, edges: tuple[Edge, ...], ) -> dict[int, int]: fg1_id, fg2_id = self._sort_functional_groups(building_block) edge1_id, edge2_id = self._sort_edges(edges) return { fg1_id: edge1_id, fg2_id: edge2_id, }
@staticmethod def _sort_functional_groups( building_block: BuildingBlock, ) -> tuple[int, int]: fg1, fg2 = building_block.get_functional_groups() x1, y1, z1 = building_block.get_centroid( atom_ids=fg1.get_placer_ids(), ) x2, y2, z2 = building_block.get_centroid( atom_ids=fg2.get_placer_ids(), ) return (0, 1) if x1 < x2 else (1, 0) @staticmethod def _sort_edges(edges: tuple[Edge, ...]) -> tuple[int, ...]: edge1, edge2 = edges x1, y1, z1 = edge1.get_position() x2, y2, z2 = edge2.get_position() if x1 < x2: return edge1.get_id(), edge2.get_id() else: return edge2.get_id(), edge1.get_id() def __str__(self): return ( f"Vertex(id={self._id}, " f"position={self._position.tolist()}, " f"flip={self._flip})" )
[docs] class TerminalVertex(LinearVertex): """ Represents a vertex at the end of a polymer chain. Do not instantiate this class directly, use :class:`.HeadVertex` or :class:`.TailVertex` instead. """ _cap_direction = 0
[docs] def place_building_block( self, building_block: BuildingBlock, edges: tuple[Edge, ...], ) -> np.ndarray: if ( building_block.get_num_functional_groups() != 1 and building_block.get_num_placers() > 1 ): return super().place_building_block(building_block, edges) building_block = building_block.with_centroid( position=self._position, atom_ids=building_block.get_placer_ids(), ) fg, *_ = building_block.get_functional_groups() fg_centroid = building_block.get_centroid( atom_ids=fg.get_placer_ids(), ) core_centroid = building_block.get_centroid( atom_ids=building_block.get_core_atom_ids(), ) return building_block.with_rotation_between_vectors( start=fg_centroid - core_centroid, # _cap_direction is defined by a subclass. target=np.array([self._cap_direction, 0, 0]), origin=self._position, ).get_position_matrix()
[docs] def map_functional_groups_to_edges( self, building_block: BuildingBlock, edges: tuple[Edge, ...], ) -> dict[int, int]: if building_block.get_num_functional_groups() == 2: functional_groups = self._sort_functional_groups( building_block=building_block, ) index = 1 if self._cap_direction == 1 else 0 return {functional_groups[index]: edges[0].get_id()} elif building_block.get_num_functional_groups() == 1: return {0: edges[0].get_id()} else: raise ValueError( "The building block of a polymer " "must have 1 or 2 functional groups." )
[docs] class HeadVertex(TerminalVertex): """ Represents a vertex at the head of a polymer chain. """ # The direction to use if the building block placed on the # vertex only has 1 FunctionalGroup. _cap_direction = 1
[docs] class TailVertex(TerminalVertex): """ Represents a vertex at the tail of a polymer chain. """ # The direction to use if the building block placed on the # vertex only has 1 FunctionalGroup. _cap_direction = -1
[docs] class UnaligningVertex(LinearVertex): """ Just places a building block, does not align. """
[docs] def place_building_block( self, building_block: BuildingBlock, edges: tuple[Edge, ...], ) -> np.ndarray: return building_block.with_centroid( position=self._position, atom_ids=building_block.get_placer_ids(), ).get_position_matrix()
[docs] def map_functional_groups_to_edges( self, building_block: BuildingBlock, edges: tuple[Edge, ...], ) -> dict[int, int]: return {fg_id: edge.get_id() for fg_id, edge in enumerate(edges)}