"""
Graph State
===========
"""
from collections import defaultdict
import numpy as np
[docs]
class GraphState:
"""
The topology graph of a molecule under construction.
"""
__slots__ = [
"_vertex_building_blocks",
"_vertices",
"_edges",
"_lattice_constants",
"_vertex_edges",
"_num_building_blocks",
]
def __init__(
self,
building_block_vertices,
edges,
lattice_constants,
):
"""
Initialize a :class:`._GraphState` instance.
Parameters
----------
building_block_vertices : :class:`dict`
Maps each :class:`.BuildingBlock` to be placed, to a
:class:`tuple` of :class:`.Vertex` instances, on which
it should be placed.
edges : :class:`tuple` of :class:`.Edge`
The edges which make up the topology graph.
lattice_constants : :class:`tuple` of :class:`numpy.ndarray`
A :class:`numpy.ndarray` for each lattice constant.
Can be an empty :class:`tuple` if the topology graph is
not periodic.
"""
self._vertex_building_blocks = {
vertex.get_id(): building_block
for building_block, vertices in building_block_vertices.items()
for vertex in vertices
}
self._num_building_blocks = {
building_block: len(vertices)
for building_block, vertices in building_block_vertices.items()
}
self._vertices = {
vertex.get_id(): vertex
for vertices in building_block_vertices.values()
for vertex in vertices
}
self._edges = edges
self._lattice_constants = tuple(
map(
np.array,
lattice_constants,
)
)
self._vertex_edges = self._get_vertex_edges()
def _get_vertex_edges(self):
"""
Get the edges connected to each vertex.
Returns
-------
:class:`dict`
Maps the id of every vertex to a :class:`tuple` of
:class:`.Edge` instances connected to it.
"""
vertex_edges = defaultdict(list)
for edge in self._edges:
if edge.is_periodic():
for vertex_id in edge.get_vertex_ids():
periodic_edge = self._get_periodic_edge(
edge=edge,
reference=vertex_id,
)
vertex_edges[vertex_id].append(periodic_edge)
else:
for vertex_id in edge.get_vertex_ids():
vertex_edges[vertex_id].append(edge)
return vertex_edges
def _get_periodic_edge(self, edge, reference):
"""
Get an :class:`.Edge`, with its position correctly set.
For a periodic edge, its correct position is not at the
midpoint of the two vertices it connects. Instead, its
correction position is different for each vertex. The get the
correct position from the perspective of *vertex1*, *vertex2*
must first be shifted to its periodic position, and only then
can the midpoint of the vertices be used to get the edge
position. An analogous calculation must be done to get the
position of the edge from the perspective of *vertex2*.
Parameters
----------
edge : :class:`.Edge`
The edge whose periodic position must be set.
reference : :class:`int`
The id of the vertex, relative to which the edge position
is being calculated.
Returns
-------
:class:`.Edge`
A clone of `edge`, shifted to have the correct periodic
position relative to `reference`.
"""
vertex1 = self._vertices[reference]
id1, id2 = edge.get_vertex_ids()
vertex2 = self._vertices[id1 if reference == id2 else id2]
direction = 1 if reference == id1 else -1
periodicity = np.array(edge.get_periodicity())
end_cell = vertex1.get_cell() + direction * periodicity
cell_shift = end_cell - vertex2.get_cell()
shift = sum(
axis_shift * constant
for axis_shift, constant in zip(
cell_shift,
self._lattice_constants,
)
)
position = (
vertex2.get_position() + shift + vertex1.get_position()
) / 2
return edge.with_position(position)
[docs]
def clone(self):
"""
Get a clone.
Returns
-------
:class:`._GraphState`
The clone. Has the same type as the original instance.
"""
clone = self.__class__.__new__(self.__class__)
clone._vertex_building_blocks = dict(self._vertex_building_blocks)
clone._vertices = dict(self._vertices)
clone._vertex_edges = dict(self._vertex_edges)
clone._edges = self._edges
clone._lattice_constants = tuple(
map(
np.array,
self._lattice_constants,
)
)
clone._num_building_blocks = dict(self._num_building_blocks)
return clone
[docs]
def get_building_block(self, vertex_id):
"""
Get the building block to be placed on a given vertex.
Parameters
----------
vertex_id : :class:`int`
The id of the vertex, on which the building block is to
be placed.
Returns
-------
:class:`.BuildingBlock`
The building block.
"""
return self._vertex_building_blocks[vertex_id]
[docs]
def get_vertices(self, vertex_ids=None):
"""
Yield the topology graph vertices.
Parameters
----------
vertex_ids : :class:`iterable` of :class:`int`, optional
The ids of vertices to yield. If ``None``, all vertices
will be yielded. Can be a single :class:`int` if a
single vertex is to be yielded.
Yields
------
:class:`.Vertex`
A vertex.
"""
if vertex_ids is None:
vertex_ids = range(len(self._vertices))
elif isinstance(vertex_ids, int):
vertex_ids = (vertex_ids,)
for vertex_id in vertex_ids:
yield self._vertices[vertex_id]
[docs]
def get_num_vertices(self):
"""
Get the number of vertices in the topology graph.
Returns
-------
:class:`int`
The number of vertices in the topology graph.
"""
return len(self._vertices)
[docs]
def get_edge(self, edge_id):
"""
Get an edge.
Parameters
----------
edge_id : :class:`int`
The id of an edge.
Returns
-------
:class:`.Edge`
An edge.
"""
return self._edges[edge_id]
[docs]
def get_num_edges(self):
"""
Get the number of edges in the topology graph.
Returns
-------
:class:`int`
The number of edges.
"""
return len(self._edges)
[docs]
def get_lattice_constants(self):
"""
Get the lattice constants of the state.
Returns
-------
:class:`tuple` of :class:`numpy.ndarray`
The lattice constants.
"""
return tuple(map(np.array, self._lattice_constants))
[docs]
def get_edges(self, vertex_id):
"""
Get the edges connect to a vertex.
Parameters
----------
vertex_id : class:`int`
The id of a vertex.
Returns
-------
:class:`tuple` of :class:`.Edge`
The connected edges.
"""
return self._vertex_edges[vertex_id]
def _with_vertices(self, vertices):
"""
Modify the instance.
"""
self._vertices = {vertex.get_id(): vertex for vertex in vertices}
return self
[docs]
def with_vertices(self, vertices):
"""
Returns a clone holding `vertices`.
Parameters
----------
vertices : :class:`iterable` of :class:`.Vertex`
The vertices the clone should hold.
Returns
-------
:class:`._GraphState`
The clone. Has the same type as the original instance.
"""
return self.clone()._with_vertices(vertices)
def _with_lattice_constants(self, lattice_constants):
"""
Modify the instance.
"""
self._lattice_constants = tuple(
map(
np.array,
lattice_constants,
)
)
return self
[docs]
def with_lattice_constants(self, lattice_constants):
"""
Return a clone holding the `lattice_constants`.
Parameters
----------
lattice_constants : :class:`tuple` of :class:`numpy.ndarray`
The lattice constants of the clone. Requires 3 arrays of
size``(3, )``.
Returns
-------
:class:`._GraphState`
The clone holding the new lattice constants. Has the same
type as the original instance.
"""
return self.clone()._with_lattice_constants(lattice_constants)
[docs]
def get_num_building_block(self, building_block):
"""
Get the number of times `building_block` is present.
Parameters
----------
building_block : :class:`.BuildingBlock`
The building block whose frequency in the topology graph
is desired.
Returns
-------
:class:`int`
The number of times `building_block` is present in the
topology graph.
"""
return self._num_building_blocks[building_block]
[docs]
def get_building_blocks(self):
"""
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
------
:class:`.BuildingBlock`
A building block of the topology graph.
"""
yielded = set()
for vertex_id in range(max(self._vertex_building_blocks) + 1):
building_block = self._vertex_building_blocks[vertex_id]
if building_block not in yielded:
yielded.add(building_block)
yield building_block