Add vertex positions argument to cage topologies (#513)

This commit will add a `vertex_positions` argument to cage topology 
graphs. This allows the users to overwrite the scaling and positioning
of any vertex (using a similar `dict[int, position]` format as the
`vertex_alignments` arguments).

It will not work for `OnePlusOne` graph currently because of strict edge
positioning.

Indeed. some choices in the code changes stem from me trying to avoid
editing code outside of the cage namespace.

A few type hints were added too on any file I edited.

I have added tests for a few cage graphs, but fundamentally it is hard
to test the success or failure of this except for consistency, which I
believe is working.

src/stk/_internal/ea/plotters/progress.py

@@ -333,7 +333,7 @@ def write(self, path: pathlib.Path | str, dpi: int = 500) -> typing.Self:
             ProgressPlotter: The plotter is returned.
         """
         sns.set(style="darkgrid")
-        fig = plt.figure(figsize=[8, 4.5])
+        fig = plt.figure(figsize=(8, 4.5))
         palette = sns.color_palette("deep")
 
         # It's possible that all values were filtered out, and trying

src/stk/_internal/topology_graphs/cage/cage.py

@@ -1,9 +1,9 @@
-from __future__ import annotations
-
 import typing
 from collections import Counter, abc, defaultdict
 from functools import partial
 
+import numpy as np
+
 from stk._internal.building_block import BuildingBlock
 from stk._internal.optimizers.null import NullOptimizer
 from stk._internal.optimizers.optimizer import Optimizer
@@ -319,7 +319,6 @@ class attributes.
                 ),
             )
 
-
         You can combine this with the `vertex_alignments` parameter
 
         .. testcode:: multi-building-block-cage-construction
@@ -395,6 +394,95 @@ class attributes.
                 ),
             )
 
+        *Construction with Custom Vertex Positions*
+
+        For :class:`.Cage` topologies, it is possible to redefine the
+        vertex positions by hand with the `vertex_positions` argument.
+
+        The parameter maps the id of a vertex to a numpy array for its
+        new position. The alignment should be modifed to match the new
+        vertex position.
+
+        It is possible to change some or all vertex positions.
+
+        Consider that the vertex positions that are provided by the user
+        are not scaled like the default ideal topology positions.
+        Additionally, existing placement rules for other vertices are
+        maintained; particularly, the effect of `vertex.init_at_center`.
+
+        .. testcode:: change-vertex-positions
+
+            import stk
+            import numpy as np
+
+            bb1 = stk.BuildingBlock(
+                smiles='NCCN',
+                functional_groups=stk.PrimaryAminoFactory(),
+            )
+            bb2 = stk.BuildingBlock(
+                smiles='O=CC(C=O)C=O',
+                functional_groups=stk.AldehydeFactory(),
+            )
+
+            cage = stk.ConstructedMolecule(
+                topology_graph=stk.cage.FourPlusSix(
+                    building_blocks=[bb1, bb2],
+                    # Build tetrahedron with tilt.
+                    vertex_positions={
+                        0: 5 * np.array([0, 1.5, 1.2]),
+                        1: 5 * np.array([-1, -0.6, -0.41]),
+                        2: 5 * np.array([1, -0.6, -0.41]),
+                        3: 5 * np.array([0, 1.2, -0.41]),
+                    },
+                ),
+            )
+
+        .. moldoc::
+
+            import moldoc.molecule as molecule
+            import stk
+            import numpy as np
+
+            bb1 = stk.BuildingBlock(
+                smiles='NCCN',
+                functional_groups=stk.PrimaryAminoFactory(),
+            )
+            bb2 = stk.BuildingBlock(
+                smiles='O=CC(C=O)C=O',
+                functional_groups=stk.AldehydeFactory(),
+            )
+
+            cage = stk.ConstructedMolecule(
+                topology_graph=stk.cage.FourPlusSix(
+                    building_blocks=[bb1, bb2],
+                    vertex_positions={
+                        0: 5 * np.array([0, 1.5, 1.2]),
+                        1: 5 * np.array([-1, -0.6, -0.41]),
+                        2: 5 * np.array([1, -0.6, -0.41]),
+                        3: 5 * np.array([0, 1.2, -0.41]),
+                    },
+                ),
+            )
+
+            moldoc_display_molecule = molecule.Molecule(
+                atoms=(
+                    molecule.Atom(
+                        atomic_number=atom.get_atomic_number(),
+                        position=position,
+                    ) for atom, position in zip(
+                        cage.get_atoms(),
+                        cage.get_position_matrix(),
+                    )
+                ),
+                bonds=(
+                    molecule.Bond(
+                        atom1_id=bond.get_atom1().get_id(),
+                        atom2_id=bond.get_atom2().get_id(),
+                        order=bond.get_order(),
+                    ) for bond in cage.get_bonds()
+                ),
+            )
+
         *Metal-Organic Cage Construction*
 
         A series of common metal-organic cage topologies are provided
@@ -878,18 +966,15 @@ def __init_subclass__(cls, **kwargs) -> None:
 
     def __init__(
         self,
-        building_blocks: typing.Union[
-            typing.Iterable[BuildingBlock],
-            dict[BuildingBlock, tuple[int, ...]],
-        ],
-        vertex_alignments: typing.Optional[dict[int, int]] = None,
+        building_blocks: abc.Iterable[BuildingBlock]
+        | dict[BuildingBlock, tuple[int, ...]],
+        vertex_alignments: dict[int, int] | None = None,
+        vertex_positions: dict[int, np.ndarray] | None = None,
         reaction_factory: ReactionFactory = GenericReactionFactory(),
         num_processes: int = 1,
         optimizer: Optimizer = NullOptimizer(),
     ) -> None:
         """
-        Initialize a :class:`.Cage`.
-
         Parameters:
 
             building_blocks:
@@ -918,6 +1003,18 @@ def __init__(
                 by a number between ``0`` (inclusive) and the number of
                 edges the vertex is connected to (exclusive).
 
+            vertex_positions:
+                A mapping from the id of a :class:`.Vertex` to a custom
+                :class:`.BuildingBlock` position. The default vertex
+                alignment algorithm is still applied. Only vertices
+                which need to have their default position changed need
+                to be present in the :class:`dict`. Note that any
+                vertices with modified positions will not be scaled like
+                the rest of the building block positions and will not
+                use neighbor placements in its positioning if requested
+                by the default topology. If ``None`` then the default
+                placement algorithm is used for each vertex.
+
             reaction_factory:
                 The reaction factory to use for creating bonds between
                 building blocks.
@@ -953,26 +1050,43 @@ def __init__(
 
         """
 
-        building_block_vertices = self._normalize_building_blocks(
-            building_blocks=building_blocks,
-        )
         self._vertex_alignments = (
             dict(vertex_alignments) if vertex_alignments is not None else {}
         )
+        self._vertex_positions = (
+            dict(vertex_positions) if vertex_positions is not None else {}
+        )
+        building_block_vertices = self._normalize_building_blocks(
+            building_blocks=building_blocks,
+        )
         building_block_vertices = self._with_unaligning_vertices(
             building_block_vertices=building_block_vertices,
         )
+        building_block_vertices = self._with_positioned_vertices(
+            building_block_vertices=building_block_vertices,
+            vertex_positions=self._vertex_positions,
+        )
         building_block_vertices = self._assign_aligners(
             building_block_vertices=building_block_vertices,
             vertex_alignments=self._vertex_alignments,
         )
         self._check_building_block_vertices(building_block_vertices)
+        edges = tuple(
+            self._normalize_edge_prototypes(
+                vertex_positions=self._vertex_positions,
+                vertices={
+                    vertex.get_id(): vertex
+                    for vertices_ in building_block_vertices.values()
+                    for vertex in vertices_
+                },
+            )
+        )
         super().__init__(
             building_block_vertices=typing.cast(
                 dict[BuildingBlock, abc.Sequence[Vertex]],
                 building_block_vertices,
             ),
-            edges=self._edge_prototypes,
+            edges=edges,
             reaction_factory=reaction_factory,
             construction_stages=tuple(
                 partial(self._has_degree, degree)
@@ -983,12 +1097,64 @@ def __init__(
             edge_groups=None,
         )
 
+    def _with_positioned_vertices(
+        self,
+        building_block_vertices: dict[
+            BuildingBlock, abc.Sequence[_CageVertex]
+        ],
+        vertex_positions: dict[int, np.ndarray],
+    ) -> dict[BuildingBlock, abc.Sequence[_CageVertex]]:
+        clone = dict(building_block_vertices)
+        for building_block, vertices in clone.items():
+            new_vertices = []
+            for vertex in vertices:
+                if vertex.get_id() in self._vertex_positions:
+                    scale = self._get_scale(
+                        building_block_vertices  # type: ignore
+                    )
+                    # Pre-reversing the scale
+                    # because altering the scale code is topology level,
+                    # which I am trying to avoid.
+                    new_vertex = vertex.with_position(
+                        vertex_positions[vertex.get_id()] / scale
+                    )
+
+                    new_vertex = new_vertex.with_use_neighbor_placement(False)
+                    new_vertices.append(new_vertex)
+                else:
+                    new_vertices.append(vertex)
+
+            clone[building_block] = tuple(new_vertices)
+
+        return clone
+
+    @classmethod
+    def _normalize_edge_prototypes(
+        cls,
+        vertex_positions: dict[int, np.ndarray],
+        vertices: dict[int, Vertex],
+    ) -> abc.Iterator[Edge]:
+        for edge in cls._edge_prototypes:
+            vertex1_id = edge.get_vertex1_id()
+            vertex2_id = edge.get_vertex2_id()
+            if any(i in vertex_positions for i in (vertex1_id, vertex2_id)):
+                new_edge = Edge(
+                    id=edge.get_id(),
+                    vertex1=vertices[vertex1_id],
+                    vertex2=vertices[vertex2_id],
+                    periodicity=edge.get_periodicity(),
+                )
+                yield new_edge
+            else:
+                yield edge
+
     @classmethod
     def _normalize_building_blocks(
         cls,
-        building_blocks: typing.Union[
-            typing.Iterable[BuildingBlock],
-            dict[BuildingBlock, tuple[int, ...]],
+        building_blocks: typing.Iterable[BuildingBlock]
+        | dict[
+            BuildingBlock,
+            tuple[int, ...],
         ],
     ) -> dict[BuildingBlock, abc.Sequence[_CageVertex]]:
         # Use tuple here because it prints nicely.
@@ -1086,15 +1252,16 @@ def _check_building_block_vertices(
                 "The following vertices are unoccupied " f"{unassigned_ids}."
             )
 
-    def clone(self) -> Cage:
+    def clone(self) -> typing.Self:
         clone = self._clone()
         clone._vertex_alignments = dict(self._vertex_alignments)
+        clone._vertex_positions = dict(self._vertex_positions)
         return clone
 
     @classmethod
     def _get_vertices(
         cls,
-        vertex_ids: typing.Union[int, typing.Iterable[int]],
+        vertex_ids: int | typing.Iterable[int],
     ) -> typing.Iterator[_CageVertex]:
         """
         Yield vertex prototypes.
@@ -1233,14 +1400,16 @@ def __repr__(self) -> str:
     def with_building_blocks(
         self,
         building_block_map: dict[BuildingBlock, BuildingBlock],
-    ) -> Cage:
+    ) -> typing.Self:
         return self.clone()._with_building_blocks(building_block_map)
 
     def get_vertex_alignments(self) -> dict[int, int]:
         """
         Get the vertex alignments.
 
         Returns:
+
             The vertex alignments.
+
         """
         return dict(self._vertex_alignments)

src/stk/_internal/topology_graphs/cage/eight_plus_sixteen.py

@@ -116,14 +116,14 @@ class EightPlusSixteen(Cage):
 
     _x = 2
     _non_linears = (
-        NonLinearVertex(0, [-0.5 * _x, 0.5 * _x, -0.35 * _x]),
-        NonLinearVertex(1, [-0.5 * _x, -0.5 * _x, -0.35 * _x]),
-        NonLinearVertex(2, [0.5 * _x, -0.5 * _x, -0.35 * _x]),
-        NonLinearVertex(3, [0.5 * _x, 0.5 * _x, -0.35 * _x]),
-        NonLinearVertex(4, [-_x * np.sqrt(2) / 2, 0, _x * 0.35]),
-        NonLinearVertex(5, [0, -_x * np.sqrt(2) / 2, _x * 0.35]),
-        NonLinearVertex(6, [_x * np.sqrt(2) / 2, 0, _x * 0.35]),
-        NonLinearVertex(7, [0, _x * np.sqrt(2) / 2, _x * 0.35]),
+        NonLinearVertex(0, np.array([-0.5 * _x, 0.5 * _x, -0.35 * _x])),
+        NonLinearVertex(1, np.array([-0.5 * _x, -0.5 * _x, -0.35 * _x])),
+        NonLinearVertex(2, np.array([0.5 * _x, -0.5 * _x, -0.35 * _x])),
+        NonLinearVertex(3, np.array([0.5 * _x, 0.5 * _x, -0.35 * _x])),
+        NonLinearVertex(4, np.array([-_x * np.sqrt(2) / 2, 0, _x * 0.35])),
+        NonLinearVertex(5, np.array([0, -_x * np.sqrt(2) / 2, _x * 0.35])),
+        NonLinearVertex(6, np.array([_x * np.sqrt(2) / 2, 0, _x * 0.35])),
+        NonLinearVertex(7, np.array([0, _x * np.sqrt(2) / 2, _x * 0.35])),
     )
 
     _vertex_prototypes = (

src/stk/_internal/topology_graphs/cage/eight_plus_twelve.py

@@ -4,6 +4,8 @@
 
 """
 
+import numpy as np
+
 from stk._internal.topology_graphs.edge import Edge
 
 from .cage import Cage
@@ -113,14 +115,14 @@ class EightPlusTwelve(Cage):
     """
 
     _non_linears = (
-        NonLinearVertex(0, [-1, 1, -1]),
-        NonLinearVertex(1, [-1, -1, -1]),
-        NonLinearVertex(2, [1, 1, -1]),
-        NonLinearVertex(3, [1, -1, -1]),
-        NonLinearVertex(4, [-1, 1, 1]),
-        NonLinearVertex(5, [-1, -1, 1]),
-        NonLinearVertex(6, [1, 1, 1]),
-        NonLinearVertex(7, [1, -1, 1]),
+        NonLinearVertex(0, np.array([-1, 1, -1])),
+        NonLinearVertex(1, np.array([-1, -1, -1])),
+        NonLinearVertex(2, np.array([1, 1, -1])),
+        NonLinearVertex(3, np.array([1, -1, -1])),
+        NonLinearVertex(4, np.array([-1, 1, 1])),
+        NonLinearVertex(5, np.array([-1, -1, 1])),
+        NonLinearVertex(6, np.array([1, 1, 1])),
+        NonLinearVertex(7, np.array([1, -1, 1])),
     )
 
     _vertex_prototypes = (