import logging
import numpy as np
from stk._internal.building_block import BuildingBlock
from stk._internal.topology_graphs.vertex import Vertex
from stk._internal.utilities.utilities import normalize_vector
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)}
[docs]
class HelixVertex(LinearVertex):
"""Represents a vertex in a helical chain."""
[docs]
def place_building_block(
self,
building_block: BuildingBlock,
edges: tuple[Edge, ...],
) -> np.ndarray:
building_block = building_block.with_centroid(
position=self._position,
atom_ids=building_block.get_placer_ids(),
)
if len(edges) == 2:
fg_centroid = building_block.get_centroid(
atom_ids=next(
building_block.get_functional_groups()
).get_placer_ids(),
)
edge_position = edges[1 if self._flip else 0].get_position()
edge_centroid = sum(edge.get_position() for edge in edges) / len(
edges
)
building_block = building_block.with_rotation_between_vectors(
start=fg_centroid - self._position,
target=edge_position - edge_centroid,
origin=self._position,
)
core_centroid = building_block.get_centroid(
atom_ids=building_block.get_core_atom_ids(),
)
building_block = building_block.with_rotation_to_minimize_angle(
start=core_centroid - self._position,
target=self._position,
axis=normalize_vector(
edges[0].get_position() - edges[1].get_position()
),
origin=self._position,
)
elif len(edges) == 1:
if building_block.get_num_functional_groups() == 2:
(fg,) = building_block.get_functional_groups(
1 if self._flip else 0
)
else:
(fg,) = building_block.get_functional_groups()
fg_centroid = building_block.get_centroid(fg.get_placer_ids())
core_centroid = building_block.get_centroid(
atom_ids=building_block.get_core_atom_ids(),
)
edge_centroid = sum(edge.get_position() for edge in edges) / len(
edges
)
building_block = building_block.with_rotation_between_vectors(
start=fg_centroid - core_centroid,
target=edge_centroid - self._position,
origin=self._position,
)
return building_block.get_position_matrix()
[docs]
def map_functional_groups_to_edges(
self,
building_block: BuildingBlock,
edges: tuple[Edge, ...],
) -> dict[int, int]:
match len(edges):
case 2:
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,
}
case 1:
return {
fg_id: edge.get_id() for fg_id, edge in enumerate(edges)
}
case _:
msg = "Number of edges should be 1 or 2."
raise RuntimeError(msg)