import typing
from collections import abc
from functools import partial
import numpy as np
from stk._internal.building_block import BuildingBlock
from stk._internal.construction_result.construction_result import (
ConstructionResult,
)
from stk._internal.construction_state.construction_state import (
ConstructionState,
)
from stk._internal.optimizers.optimizer import Optimizer
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.parallel import (
_Parallel,
)
from stk._internal.topology_graphs.topology_graph.serial import (
_Serial,
)
from stk._internal.topology_graphs.topology_graph.utilities import (
_TopologyGraphImplementation,
)
from stk._internal.topology_graphs.vertex import Vertex
from stk._internal.utilities.utilities import flatten
[docs]
class TopologyGraph:
"""
An abstract base class for topology graphs.
It is responsible for the construction of molecules. To create
a new topology graph, you want to subclass and implement this
abstract base class.
Notes:
*Adding New Topology Graphs*
You might notice that some of the methods of this abstract base
class are implemented. This is purely for convenience when
implementing subclasses. The implemented public methods are
simply default implementations, which can safely be ignored or
overridden, when implementing subclasses. Any private methods
are implementation details of these default implementations.
Many classes, such as :class:`.Vertex`, :class:`.Edge`,
:class:`.EdgeGroup` and :class:`.ConstructionState`, exist as
implementation details of this default :class:`.TopologyGraph`
implementation. You could ignore all of them, and define a new
:meth:`.construct` method from scratch. In fact, your topology
graph does not have to be represented as a graph at all.
However, using the default implementation of
:class:`.TopologyGraph` makes it significantly easier to
implement a construction process. When using the default
implementation of :class:`.TopologyGraph`, you mostly just
need to implement a :class:`.Vertex` subclass, which is much
easier than figuring out the whole construction process from
scratch. In addition, you get benefits like parallel
construction for free, as it is included in the default
implementation.
Typically, adding a new topology graph will involve
implementing any pure virtual methods of
:class:`.TopologyGraph`, in a new subclass, as well as
implementing any pure virtual methods of :class:`.Vertex`,
again in a new subclass. Combined, this is just a handful of
simple methods to implement. Sometimes, rarely, you might also
want to subclass :class:`.ConstructionState`, when you want to
add additional hooks during construction, by extending
the methods of this class. If you do this, make sure to
override :meth:`._get_construction_state` to return your
subclass of :class:`.ConstructionState`, rather than the base
class, as is done by default. You can subclass and extend the
methods of any class as you wish, but it would be unusual if
this doesn't cover all your requirements.
*The Default Implementation*
The default implementation of :class:`.TopologyGraph`
represents the constructed molecule through a graph. The
vertices indicate where building blocks are placed and the
edges indicate which building blocks have bonds formed between
them by the construction process.
:class:`.Vertex` instances are responsible for placing the
building block molecules. By initializing the vertices with
different parameters, you can alter how they position the
building block molecules, and therefore allow the user to
easily specify a different structural isomer.
Once a building block is placed on a vertex, the functional
groups on the building block must be mapped to the different
edges connected to the vertex. The number of functional groups
in the building block must match the number of edges connected
to the vertex.
Once the functional groups are mapped to edges, the edges are
used to perform reactions on the building blocks. Edges are
grouped in an :class:`.EdgeGroup`, and all functional groups
present in the edge group are reacted together. Normally,
unless you are doing something very exotic, an
:class:`.EdgeGroup` will hold just one :class:`.Edge`, and the
two functional groups on that edge will be reacted together
through a single :class:`.Reaction`. This reaction will
normally add the bonds which are required to form the
joined-up constructed molecule, but note that it does not have
to add any bonds at all. In addition, a :class:`.Reaction` can
add and remove atoms from the constructed molecule. Which
reaction is selected to join the functional groups depends on
the :class:`.ReactionFactory` given to the
:class:`.TopologyGraph` during initialization.
Once this is done, you have a :class:`.ConstructedMolecule`.
Examples:
*Subclass Implementation*
The source code of subclasses, listed in
:mod:`~.topology_graph.topology_graph.topology_graph`, can
serve as good examples.
*Changing the Building Blocks of a Topology Graph*
To change the building blocks used by a topology graph you
can use :meth:`.with_building_blocks` to get a clone of the
topology graph holding the new building blocks
.. testcode:: changing-the-building-blocks-of-a-topology-graph
import stk
bb1 = stk.BuildingBlock('BrCCBr', [stk.BromoFactory()])
bb2 = stk.BuildingBlock('BrCCCBr', [stk.BromoFactory()])
linear = stk.polymer.Linear(
building_blocks=(bb1, bb2),
repeating_unit='A',
num_repeating_units=15,
)
bb3 = stk.BuildingBlock('BrCNCBr', [stk.BromoFactory()])
# All bb1 instances are replaced by bb3, but bb2 remains
# in place.
clone = linear.with_building_blocks({
bb1: bb3,
})
"""
_implementation: _TopologyGraphImplementation
def __init__(
self,
building_block_vertices: dict[
BuildingBlock,
abc.Sequence[Vertex],
],
edges: tuple[Edge, ...],
reaction_factory: ReactionFactory,
construction_stages: tuple[
abc.Callable[[Vertex], bool],
...,
],
num_processes: int,
optimizer: Optimizer,
edge_groups: tuple[EdgeGroup, ...] | None = None,
scale_multiplier: float = 1.0,
) -> None:
"""
Parameters:
building_block_vertices:
Maps each :class:`.BuildingBlock` to be placed, to the
:class:`.Vertex` instances, on which it should be
placed.
edges:
The edges which make up the topology graph.
reaction_factory:
Used to pick which :class:`.Reaction` is used on each
:class:`.EdgeGroup` of the topology graph.
construction_stages:
A collection of :class:`~collections.abc.Callable`,
each of which takes a :class:`.Vertex` and returns
``True`` or ``False``. If the first
:class:`~collections.abc.Callable` is applied
to a vertex in the topology graph, and the result is
``True``, that vertex is a part of the first
construction stage. The second :class:`callable` is
then applied to all vertices not in the first stage
and those which return ``True`` belong to the second
stage and so on.
Vertices which belong to the same construction stage
all place building blocks together in parallel, before
placement is done by any vertices which are part of a
later stage. This breaks down parallel construction
into serial stages if synchronization between stages
is needed.
If the topology graph is performing construction
serially, then all vertices which belong to an earlier
stage will place their building block before those at
a later stage.
num_processes:
The number of parallel processes to create during
:meth:`construct`.
optimizer:
Used to optimize the structure of the constructed
molecule.
edge_groups:
The edge groups of the topology graph, if ``None``,
every :class:`.Edge` is in its own edge group.
scale_multiplier:
Scales the positions of the vertices.
"""
self._scale_multiplier = scale_multiplier
self._scale = scale = self._get_scale(
building_block_vertices, self._scale_multiplier
)
def apply_scale(item):
return item.with_scale(scale)
self._building_block_vertices = {
building_block: tuple(map(apply_scale, vertices))
for building_block, vertices in building_block_vertices.items()
}
self._edges = tuple(map(apply_scale, edges))
self._reaction_factory = reaction_factory
if num_processes == 1:
self._implementation = _Serial(
stages=tuple(self._get_stages(construction_stages)),
)
else:
self._implementation = _Parallel(
stages=tuple(self._get_stages(construction_stages)),
num_processes=num_processes,
)
if edge_groups is None:
edge_groups = tuple(EdgeGroup((edge,)) for edge in self._edges)
self._edge_groups = edge_groups
self._optimizer = optimizer
def _with_building_blocks(
self,
building_block_map: dict[BuildingBlock, BuildingBlock],
) -> typing.Self:
"""
Modify the topology graph.
"""
# The original scaling first needs to be removed, so that when
# the scale is recalculated with the new building blocks, it
# has the same starting geometry.
def undo_scale(vertex):
return vertex.with_scale(1 / self._scale)
building_block_vertices: dict[BuildingBlock, abc.Sequence[Vertex]]
building_block_vertices = {
building_block_map.get(bb, bb): tuple(map(undo_scale, vertices))
for bb, vertices in self._building_block_vertices.items()
}
scale = self._get_scale(
building_block_vertices, self._scale_multiplier
)
def scale_vertex(vertex):
return vertex.with_scale(scale)
self._building_block_vertices = {
building_block: tuple(map(scale_vertex, vertices))
for building_block, vertices in building_block_vertices.items()
}
def scale_edge(edge):
# Remove the old scale and apply the new one.
return edge.with_scale(scale / self._scale)
self._edges = edges = tuple(map(scale_edge, self._edges))
def get_new_edge(edge_id):
return edges[edge_id]
self._edge_groups = tuple(
EdgeGroup(map(get_new_edge, edge_group.get_edge_ids()))
for edge_group in self._edge_groups
)
self._scale = scale
return self
[docs]
def with_building_blocks(
self,
building_block_map: dict[BuildingBlock, BuildingBlock],
) -> typing.Self:
"""
Return a clone holding different building blocks.
Parameters:
building_block_map:
Maps a building block in the current topology
graph to the building block which should replace
it in the clone. If a building block should be not
replaced in the clone, it can be omitted from the map.
Returns:
The clone.
"""
return self.clone()._with_building_blocks(building_block_map)
[docs]
def clone(self) -> typing.Self:
"""
Return a clone.
Returns:
The clone.
"""
return self._clone()
def _clone(self) -> typing.Self:
clone = self.__class__.__new__(self.__class__)
clone._scale = self._scale
clone._building_block_vertices = dict(self._building_block_vertices)
clone._edges = self._edges
clone._reaction_factory = self._reaction_factory
clone._implementation = self._implementation
clone._optimizer = self._optimizer
clone._edge_groups = self._edge_groups
clone._scale_multiplier = self._scale_multiplier
return clone
[docs]
def get_building_blocks(self) -> abc.Iterator[BuildingBlock]:
"""
Yield the building blocks.
Building blocks are yielded in an order based on their
position in the topology graph. For two equivalent
topology graphs, but with different building blocks,
equivalently positioned building blocks will be yielded at the
same time.
Yields:
A building block of the topology graph.
"""
vertex_building_blocks = {}
num_vertices = 0
for (
building_block,
vertices,
) in self._building_block_vertices.items():
for vertex in vertices:
num_vertices += 1
vertex_building_blocks[vertex.get_id()] = building_block
yielded = set()
for vertex_id in range(num_vertices):
building_block = vertex_building_blocks[vertex_id]
if building_block not in yielded:
yielded.add(building_block)
yield building_block
[docs]
def get_num_building_block(
self,
building_block: BuildingBlock,
) -> int:
"""
Get the number of times `building_block` is present.
Parameters:
building_block:
The building block whose frequency in the topology
graph is desired.
Returns:
The number of times `building_block` is present in the
topology graph.
"""
return len(self._building_block_vertices.get(building_block, []))
def _get_lattice_constants(self) -> abc.Iterator[np.ndarray]:
"""
Yield the lattice constants of the topology graph.
The a, b and c lattice constants are yielded, in that order.
By default, this is an empty generator.
Yields:
A lattice constant.
"""
return iter(())
[docs]
def construct(self) -> ConstructionResult:
"""
Construct a :class:`.ConstructedMolecule`.
Returns:
The data describing the :class:`.ConstructedMolecule`.
"""
state = self._get_construction_state()
state = self._place_building_blocks(state)
state = self._run_reactions(state)
state = self._optimizer.optimize(state)
return self._get_construction_result(state)
def _get_construction_result(
self,
state: ConstructionState,
) -> ConstructionResult:
"""
Get the result of the construction.
Parameters:
state:
The state of the molecule being constructed.
Returns:
The data describing the :class:`.ConstructedMolecule`.
"""
return ConstructionResult(state)
def _get_construction_state(self) -> ConstructionState:
return ConstructionState(
building_block_vertices=self._building_block_vertices,
edges=self._edges,
lattice_constants=tuple(
np.array(constant, dtype=np.float64) * self._scale
for constant in self._get_lattice_constants()
),
)
@staticmethod
def _get_scale(
building_block_vertices: dict[BuildingBlock, abc.Sequence[Vertex]],
scale_multiplier: float,
) -> float:
"""
Get the scale, which should be applied to topology graph.
The scale should be applied to the position of every vertex
and edge of topology graph. This allows to graph to adjust
based on the size of the building blocks.
Parameters:
building_block_vertices:
Maps every :class:`.BuildingBlock` of the topology
graph, to the :class:`.Vertex` instances it is meant
to be placed on.
Returns:
The scale.
"""
raise NotImplementedError()
def _place_building_blocks(
self,
state: ConstructionState,
) -> ConstructionState:
"""
Place the building blocks onto the vertices.
Parameters:
state:
Holds data necessary to construct the molecule.
Returns:
The new construction state, updated to account for the
placed building blocks.
"""
return self._implementation._place_building_blocks(state)
def _run_reactions(
self,
state: ConstructionState,
) -> ConstructionState:
"""
Perform the reactions on the building blocks.
Parameters:
state:
The current state of the construction process.
Returns:
The new construction state, updated to account for the
reactions between building blocks.
"""
get_reaction = partial(
self._reaction_factory.get_reaction,
state,
)
reactions = tuple(map(get_reaction, self._edge_groups))
results = map(
lambda reaction: reaction.get_result(),
reactions,
)
return state.with_reaction_results(reactions, results)
def _get_stages(
self,
construction_stages,
) -> abc.Iterator[tuple[int, ...]]:
"""
Yield the parallelizable stages of construction.
Yields:
Vertices ids, which can be placed in parallel.
"""
stages: tuple[list[int], ...] = tuple(
[] for i in range(len(construction_stages) + 1)
)
vertices = flatten(self._building_block_vertices.values())
for vertex in vertices:
placed = False
for i, stage in enumerate(construction_stages):
if stage(vertex):
stages[i].append(vertex.get_id())
placed = True
break
if not placed:
stages[-1].append(vertex.get_id())
yield from (tuple(stage) for stage in stages if stage)
def __str__(self) -> str:
return repr(self)
def __repr__(self) -> str:
raise NotImplementedError()