import itertools as it
import typing
from collections import Counter, abc
from functools import partial
from operator import getitem
import numpy as np
from stk._internal.building_block import BuildingBlock
from stk._internal.construction_result.construction_result import (
ConstructionResult,
)
from stk._internal.construction_result.periodic import (
PeriodicConstructionResult,
)
from stk._internal.optimizers.null import (
NullOptimizer,
)
from stk._internal.optimizers.optimizer import Optimizer
from stk._internal.reaction_factories.generic_reaction_factory import (
GenericReactionFactory,
)
from stk._internal.reaction_factories.reaction_factory import ReactionFactory
from stk._internal.topology_graphs.edge import (
Edge,
)
from stk._internal.topology_graphs.edge_group import (
EdgeGroup,
)
from stk._internal.topology_graphs.topology_graph.topology_graph import (
TopologyGraph,
)
from stk._internal.topology_graphs.vertex import Vertex
from .edge import CofEdge
from .vertices import UnaligningVertex, _CofVertex
class UnoccupiedVertexError(Exception):
"""
When a COF vertex is not occupied by a building block.
"""
pass
class OverlyOccupiedVertexError(Exception):
"""
When a COF vertex is occupied by more than one building block.
"""
pass
[docs]
class Cof(TopologyGraph):
"""
Represents a COF topology graph.
Notes:
COF topologies are added by creating a subclass, which defines
the :attr:`_vertex_prototypes` and :attr:`_edge_prototypes` class
attributes.
.. _cof-topology-graph-examples:
Examples:
*Subclass Implementation*
The source code of the subclasses, listed in :mod:`~.cof.cof`,
can serve as good examples.
*Basic Construction*
:class:`.Cof` instances can be made by providing the building
block molecules and lattice size only (using :class:`.Honeycomb`
as an example)
.. testcode:: basic-construction
import stk
bb1 = stk.BuildingBlock('BrCCBr', [stk.BromoFactory()])
bb2 = stk.BuildingBlock('BrCC(CBr)CBr', [stk.BromoFactory()])
cof = stk.ConstructedMolecule(
topology_graph=stk.cof.Honeycomb((bb1, bb2), (3, 3, 1)),
)
.. moldoc::
import moldoc.molecule as molecule
import stk
bb1 = stk.BuildingBlock('BrCCBr', [stk.BromoFactory()])
bb2 = stk.BuildingBlock('BrCC(CBr)CBr', [stk.BromoFactory()])
cof = stk.ConstructedMolecule(
topology_graph=stk.cof.Honeycomb((bb1, bb2), (3, 3, 1)),
)
moldoc_display_molecule = molecule.Molecule(
atoms=(
molecule.Atom(
atomic_number=atom.get_atomic_number(),
position=position,
) for atom, position in zip(
cof.get_atoms(),
cof.get_position_matrix(),
)
),
bonds=(
molecule.Bond(
atom1_id=bond.get_atom1().get_id(),
atom2_id=bond.get_atom2().get_id(),
order=(
1
if bond.get_order() == 9
else bond.get_order()
),
) for bond in cof.get_bonds()
),
)
*Suggested Optimization*
For :class:`.Cof` topologies, it is recommend to use the
:class:`.Collapser` optimizer if using the nonperiodic form.
For periodic systems, the :class:`.PeriodicCollapser` is
recommended.
.. testcode:: suggested-optimization
import stk
bb1 = stk.BuildingBlock('BrCCBr', [stk.BromoFactory()])
bb2 = stk.BuildingBlock('BrCC(CBr)CBr', [stk.BromoFactory()])
# Nonperiodic.
cof = stk.ConstructedMolecule(
topology_graph=stk.cof.Honeycomb(
building_blocks=(bb1, bb2),
lattice_size=(3, 3, 1),
# Setting scale_steps to False tends to lead to a
# better structure.
optimizer=stk.Collapser(scale_steps=False),
),
)
# Periodic.
cof = stk.ConstructedMolecule(
topology_graph=stk.cof.PeriodicHoneycomb(
building_blocks=(bb1, bb2),
lattice_size=(3, 3, 1),
optimizer=stk.PeriodicCollapser(),
),
)
.. moldoc::
import moldoc.molecule as molecule
import stk
bb1 = stk.BuildingBlock('BrCCBr', [stk.BromoFactory()])
bb2 = stk.BuildingBlock('BrCC(CBr)CBr', [stk.BromoFactory()])
cof = stk.ConstructedMolecule(
topology_graph=stk.cof.Honeycomb(
building_blocks=(bb1, bb2),
lattice_size=(3, 3, 1),
# Setting scale_steps to False tends to lead to a
# better structure.
optimizer=stk.Collapser(scale_steps=False),
),
)
moldoc_display_molecule = molecule.Molecule(
atoms=(
molecule.Atom(
atomic_number=atom.get_atomic_number(),
position=position,
) for atom, position in zip(
cof.get_atoms(),
cof.get_position_matrix(),
)
),
bonds=(
molecule.Bond(
atom1_id=bond.get_atom1().get_id(),
atom2_id=bond.get_atom2().get_id(),
order=(
1
if bond.get_order() == 9
else bond.get_order()
),
) for bond in cof.get_bonds()
),
)
*Accessing the Periodic Information*
When building periodic :class:`.Cof` instances, the periodic
information, such as the unit cell, can be accessed if you use the
:class:`.PeriodicConstructionResult` returned by calling
:meth:`.Cof.construct`
.. testcode:: accessing-the-periodic-information
import stk
bb1 = stk.BuildingBlock('BrCCBr', [stk.BromoFactory()])
bb2 = stk.BuildingBlock('BrCC(CBr)CBr', [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,
)
periodic_info = construction_result.get_periodic_info()
cell_matrix = periodic_info.get_cell_matrix()
# Can access all unit-cell parameters.
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()
# Write to .pdb file.
writer = stk.PdbWriter()
writer.write(
molecule=cof,
path='cof.pdb',
periodic_info=periodic_info,
)
.. testcleanup:: accessing-the-periodic-information
import os
os.remove('cof.pdb')
*Structural Isomer Construction*
Different structural isomers of COFs can be made by using the
`vertex_alignments` optional parameter
.. testcode:: structural-isomer-construction
import stk
bb1 = stk.BuildingBlock('BrCCBr', [stk.BromoFactory()])
bb2 = stk.BuildingBlock('BrCC(CBr)CBr', [stk.BromoFactory()])
cof2 = stk.ConstructedMolecule(
topology_graph=stk.cof.Honeycomb(
building_blocks=(bb1, bb2),
lattice_size=(3, 3, 1),
vertex_alignments={0: 2, 1: 1, 2: 1},
),
)
The parameter maps the id of a vertex to a number
between 0 (inclusive) and the number of edges the vertex is
connected to (exclusive). So a vertex connected to three edges
can be mapped to ``0``, ``1`` or ``2``.
By changing which edge each vertex is aligned with, a different
structural isomer of the COF can be formed.
*Multi-Building Block COF Construction*
You can also build COFs with multiple building blocks, but,
if you have multiple building blocks with the same number of
functional groups, you have to assign each building block to the
vertex you want to place it on
.. testcode:: multi-building-block-cof-construction
import stk
bb1 = stk.BuildingBlock('BrCCBr', [stk.BromoFactory()])
bb2 = stk.BuildingBlock('BrCNCBr', [stk.BromoFactory()])
bb3 = stk.BuildingBlock('BrCC(CBr)CBr', [stk.BromoFactory()])
bb4 = stk.BuildingBlock('BrCC(NCBr)CBr', [stk.BromoFactory()])
cof = stk.ConstructedMolecule(
topology_graph=stk.cof.Honeycomb(
# building_blocks is now a dict, which maps building
# blocks to the id of the vertices it should be placed
# on. You can use ranges to specify the ids.
building_blocks={
bb1: (2, 4, 7, 8, 9, 12, 13, 14, 17, 18, 19),
bb2: 3,
bb3: (0, 10, 11, 15, 16),
bb4: (1, 5, 6),
},
lattice_size=(2, 2, 1),
),
)
You can combine this with the `vertex_alignments` parameter
.. testcode:: multi-building-block-cof-construction
cof2 = stk.ConstructedMolecule(
topology_graph=stk.cof.Honeycomb(
building_blocks={
bb1: (2, 4, 7, 8, 9, 12, 13, 14, 17, 18, 19),
bb2: 3,
bb3: (0, 10, 11, 15, 16),
bb4: (1, 5, 6),
},
lattice_size=(2, 2, 1),
vertex_alignments={0: 2, 1: 1, 2: 1},
),
)
"""
_allowed_degrees: typing.ClassVar[dict[int, int]]
_vertex_prototypes: typing.ClassVar[tuple[_CofVertex, ...]]
_edge_prototypes: typing.ClassVar[tuple[Edge, ...]]
_lattice_constants: typing.ClassVar[
tuple[np.ndarray, np.ndarray, np.ndarray]
]
def __init_subclass__(cls, **kwargs):
vertex_degrees = Counter(
vertex_id
for edge in cls._edge_prototypes
for vertex_id in edge.get_vertex_ids()
)
cls._allowed_degrees = set(vertex_degrees.values())
def __init__(
self,
building_blocks: (
abc.Iterable[BuildingBlock] | dict[BuildingBlock, tuple[int, ...]]
),
lattice_size: tuple[int, int, int],
periodic: bool = False,
vertex_alignments: dict[int, int] | None = None,
reaction_factory: ReactionFactory = GenericReactionFactory(),
num_processes: int = 1,
optimizer: Optimizer = NullOptimizer(),
scale_multiplier: float = 1.0,
) -> None:
"""
Initialize a :class:`.Cof` instance.
Parameters:
building_blocks:
Can be a :class:`tuple` of :class:`.BuildingBlock`
instances, which should be placed on the topology graph.
Can also be a :class:`dict` which maps the
:class:`.BuildingBlock` instances to the ids of the
vertices it should be placed on. A :class:`dict` is
required when there are multiple building blocks with the
same number of functional groups, because in this case
the desired placement is ambiguous.
lattice_size:
The size of the lattice in the x, y and z directions.
periodic:
Toggle the construction of a periodic molecule. If
``True``, periodic bonds will be made across the edges of
the lattice.
vertex_alignments:
A mapping from the id of a :class:`.Vertex`
to an :class:`.Edge` connected to it.
The :class:`.Edge` is used to align the first
:class:`.FunctionalGroup` of a :class:`.BuildingBlock`
placed on that vertex. Only vertices which need to have
their default edge changed need to be present in the
:class:`dict`. If ``None`` then the default edge is used
for each vertex. Changing which :class:`.Edge` is used will
mean that the topology graph represents different
structural isomers. The edge is referred to by a number
between ``0`` (inclusive) and the number of edges the
vertex is connected to (exclusive).
reaction_factory:
The reaction factory to use for creating bonds between
building blocks.
num_processes:
The number of parallel processes to create during
:meth:`construct`.
optimizer:
Used to optimize the structure of the constructed
molecule.
scale_multiplier:
Scales the positions of the vertices.
Raises:
:class:`AssertionError`
If the any building block does not have a
valid number of functional groups.
:class:`ValueError`
If the there are multiple building blocks with the
same number of functional_groups in `building_blocks`,
and they are not explicitly assigned to vertices. The
desired placement of building blocks is ambiguous in
this case.
:class:`~.cof.UnoccupiedVertexError`
If a vertex of the COF topology graph does not have a
building block placed on it.
:class:`~.cof.OverlyOccupiedVertexError`
If a vertex of the COF topology graph has more than one
building block placed on it.
"""
self._vertex_alignments = (
dict(vertex_alignments) if vertex_alignments is not None else {}
)
self._lattice_size = lattice_size
self._periodic = periodic
lattice = self._get_lattice(self._vertex_alignments)
edges = self._get_edges(lattice)
vertices = self._get_vertices(lattice)
if isinstance(building_blocks, dict):
for building_block in building_blocks:
assert (
building_block.get_num_functional_groups()
in self._allowed_degrees
), (
"The number of functional groups in "
f"{building_block} needs to be one of "
f"{tuple(self._allowed_degrees)}, but is "
"currently "
f"{building_block.get_num_functional_groups()}."
)
get_vertex = partial(getitem, vertices)
building_block_vertices = {
building_block: tuple(
map(
get_vertex,
# Account for the fact that a building block can
# be mapped to a single int.
(ids,) if isinstance(ids, int) else ids,
)
)
for building_block, ids in building_blocks.items()
}
else:
building_block_vertices = self._get_building_block_vertices( # type: ignore[assignment]
building_blocks=building_blocks,
vertices=vertices,
edges=edges,
)
building_block_vertices = self._with_unaligning_vertices(
building_block_vertices=building_block_vertices,
)
self._check_building_block_vertices(
num_vertices=(
np.prod(lattice_size) * len(self._vertex_prototypes)
),
building_block_vertices=building_block_vertices,
)
super().__init__(
building_block_vertices=typing.cast(
dict[BuildingBlock, abc.Sequence[Vertex]],
building_block_vertices,
),
edges=edges,
reaction_factory=reaction_factory,
construction_stages=(),
num_processes=num_processes,
optimizer=optimizer,
edge_groups=self._get_edge_groups(edges),
scale_multiplier=scale_multiplier,
)
@staticmethod
def _with_unaligning_vertices(building_block_vertices):
clone = dict(building_block_vertices)
for building_block, vertices in clone.items():
# Building blocks with 1 placer, cannot be aligned and
# must therefore use an UnaligningVertex.
if building_block.get_num_placers() == 1:
clone[building_block] = tuple(
UnaligningVertex(
id=v.get_id(),
position=v.get_position(),
aligner_edge=v.get_aligner_edge(),
cell=v.get_cell(),
)
for v in vertices
)
return clone
@classmethod
def _check_building_block_vertices(
cls,
num_vertices,
building_block_vertices,
):
unassigned_ids = set(range(num_vertices))
assigned_ids = set()
vertices = (
vertex
for vertices_ in building_block_vertices.values()
for vertex in vertices_
)
for vertex in vertices:
if vertex.get_id() in assigned_ids:
raise OverlyOccupiedVertexError(
f"Vertex {vertex.get_id()} has multiple building "
"blocks placed on it."
)
assigned_ids.add(vertex.get_id())
unassigned_ids.remove(vertex.get_id())
if unassigned_ids:
raise UnoccupiedVertexError(
"The following vertices are unoccupied " f"{unassigned_ids}."
)
[docs]
def clone(self):
clone = super().clone()
clone._vertex_alignments = dict(self._vertex_alignments)
clone._lattice_size = self._lattice_size
clone._periodic = self._periodic
return clone
def _get_edge_groups(
self,
edges: abc.Iterable[Edge],
) -> tuple[EdgeGroup, ...] | None:
"""
Get the edge groups for the COF.
Parameters:
edges:
The edges which are to be placed into edge groups.
Returns:
The edge groups. These are returned when the lattice is
not periodic, which means that some edges should not be
part of an edge group, in order to prevent bond formation
across the edges of the lattice.
None:
If an edge group is to be created for every single edge.
"""
if self._periodic:
return None
return tuple(
EdgeGroup((edge,)) for edge in edges if not edge.is_periodic()
)
def _get_vertices(
self,
lattice: list[list[list[dict[int, _CofVertex]]]],
) -> list[_CofVertex]:
"""
Get the vertices in the `lattice`.
Parameters:
lattice:
A nested list which can be in the form
``lattice[x][y][z][vertex_id]`` which returns a vertex
in the (x, y, z) cell of the lattice.
Returns:
All the vertices extracted from `lattice`.
"""
xdim, ydim, zdim = self._lattice_size
vertices = []
for x, y, z in it.product(
range(xdim),
range(ydim),
range(zdim),
):
for vertex in lattice[x][y][z].values():
vertices.append(vertex)
return sorted(vertices, key=lambda vertex: vertex.get_id())
def _get_lattice(
self,
vertex_alignments: dict[int, int],
) -> list[list[list[dict[int, _CofVertex]]]]:
"""
Get the vertices of the topology graph instance.
Parameters:
vertex_alignments:
A mapping from the id of a :class:`.Vertex`
to an :class:`.Edge` connected to it.
The :class:`.Edge` is used to align the first
:class:`.FunctionalGroup` of a :class:`.BuildingBlock`
placed on that vertex. Only vertices which need to have
their default edge changed need to be present in the
:class:`dict`. If ``None`` then the default edge is used
for each vertex. Changing which :class:`.Edge` is used will
mean that the topology graph represents different
structural isomers. The edge is referred to by a number
between ``0`` (inclusive) and the number of edges the
vertex is connected to (exclusive).
Returns:
A nested :class:`list` which can be indexed as
``vertices[x][y][z]``, which will return a :class:`dict`
for the unit cell at (x, y, z). The :class:`dict` maps
the vertices in :attr:`_vertex_prototypes` to its clone for
that unit cell.
"""
xdim, ydim, zdim = (range(dim) for dim in self._lattice_size)
# vertex_clones is indexed as vertex_clones[x][y][z]
lattice: list[list[list[dict[int, _CofVertex]]]] = [
[[{} for _ in zdim] for _ in ydim] for _ in xdim
]
# Make a clone of each vertex for each unit cell.
cells = it.product(xdim, ydim, zdim)
vertices = it.product(cells, self._vertex_prototypes)
for id_, (cell, vertex) in enumerate(vertices):
x, y, z = cell
shift = sum(
axis * dim for axis, dim in zip(cell, self._lattice_constants)
)
lattice[x][y][z][vertex.get_id()] = vertex.__class__(
id=id_,
position=vertex.get_position() + shift,
aligner_edge=vertex_alignments.get(id_, 0),
cell=cell,
)
return lattice
def _get_edges(self, lattice: list) -> tuple[Edge, ...]:
"""
Create the edges of the topology graph instance.
Parameters:
lattice:
A nested :class:`list` which can be indexed as
``vertices[x][y][z]``, which will return a :class:`dict`
for the unit cell at (x, y, z). The :class:`dict` maps
the vertices in :attr:`_vertex_prototypes` to the clones
for that unit cell.
Returns:
The edges of the topology graph instance.
"""
edge_clones = []
# Make a clone for each edge for each unit cell.
xdim, ydim, zdim = (range(dim) for dim in self._lattice_size)
cells = it.product(xdim, ydim, zdim)
edges = it.product(cells, self._edge_prototypes)
for id_, (cell, edge) in enumerate(edges):
x, y, z = cell
# The cell in which the second vertex of the edge is found.
periodic_cell = np.array(cell) + edge.get_periodicity()
# Wrap around periodic cells, ie those that are less than 0
# or greater than the lattice size along any dimension.
dims = zip(periodic_cell, self._lattice_size)
x2, y2, z2 = np.array(
[(dim + max_dim) % max_dim for dim, max_dim in dims]
)
# The edge is not periodic if periodic_cell did not
# have to wrap around.
dims = zip(periodic_cell, self._lattice_size)
edge_is_not_periodic = all(
dim >= 0 and dim < max_dim for dim, max_dim in dims
)
edge_clones.append(
CofEdge(
parent_id=edge.get_id(),
id=id_,
vertex1=lattice[x][y][z][edge.get_vertex1_id()],
vertex2=lattice[x2][y2][z2][edge.get_vertex2_id()],
periodicity=(
(0, 0, 0)
if edge_is_not_periodic
else edge.get_periodicity()
),
)
)
return tuple(edge_clones)
@classmethod
def _get_building_block_vertices(
cls,
building_blocks: abc.Iterable[BuildingBlock],
vertices: abc.Iterable[_CofVertex],
edges: abc.Iterable[Edge],
) -> dict[BuildingBlock, abc.Sequence[_CofVertex]]:
"""
Map building blocks to the vertices of the graph.
Parameters:
building_blocks:
The building blocks which need to be mapped to `vertices`.
vertices:
The vertices which need to have a building block map to
them.
edges:
The edges of the graph.
Returns:
Maps each building block in `building_blocks` to a
:class:`list` of the :class:`.Vertex` instances it should
be placed on.
Raises:
:class:`AssertionError`
If the any building block does not have a
valid number of functional groups.
:class:`ValueError`
If there are multiple building blocks with the same number
of functional groups.
"""
building_blocks_by_degree = {}
for building_block in building_blocks:
num_fgs = building_block.get_num_functional_groups()
assert num_fgs in cls._allowed_degrees, (
"The number of functional groups in "
f"{building_block} needs to be one of "
f"{tuple(cls._allowed_degrees)}, but is "
"currently "
f"{building_block.get_num_functional_groups()}."
)
if num_fgs in building_blocks_by_degree:
raise ValueError(
"If there are multiple building blocks with the "
"same number of functional groups, "
"building_block_vertices must be set explicitly."
)
building_blocks_by_degree[num_fgs] = building_block
vertex_degrees = Counter(
vertex_id for edge in edges for vertex_id in edge.get_vertex_ids()
)
building_block_vertices = {} # type: ignore[var-annotated]
for vertex in vertices:
vertex_degree = vertex_degrees[vertex.get_id()]
building_block = building_blocks_by_degree[vertex_degree]
building_block_vertices[building_block] = (
building_block_vertices.get(building_block, [])
)
building_block_vertices[building_block].append(vertex)
return building_block_vertices
def _get_lattice_constants(self) -> abc.Iterator[np.ndarray]:
yield from self._lattice_constants
[docs]
def construct(self) -> ConstructionResult:
"""
Construct a :class:`.ConstructedMolecule`.
Returns:
The data describing the :class:`.ConstructedMolecule`.
"""
return super().construct()
def _get_construction_result(self, state):
return PeriodicConstructionResult(state, self._lattice_size)
@staticmethod
def _get_scale(
building_block_vertices: dict[BuildingBlock, abc.Sequence[Vertex]],
scale_multiplier: float,
) -> float:
return (
scale_multiplier
* 5
* max(bb.get_maximum_diameter() for bb in building_block_vertices)
)
def __repr__(self):
x, y, z = self._lattice_size
periodic = ", periodic=True" if self._periodic else ""
vertex_alignments = (
f", vertex_alignments={self._vertex_alignments}"
if self._vertex_alignments
else ""
)
return (
f"cof.{self.__class__.__name__}("
f"({x}, {y}, {z})"
f"{vertex_alignments}"
f"{periodic})"
)