From 554f89fef52bb2090d262931f46c0a5cda47b305 Mon Sep 17 00:00:00 2001 From: CalCraven Date: Wed, 6 May 2026 17:53:46 -0500 Subject: [PATCH 01/11] update simulation.py for openmm better minimization, add path count functions such as angle and bond distributions, as well as freud rdf --- mbuild/path/build.py | 197 ++++++++++++++++++++++++++++++++++++++++--- mbuild/simulation.py | 189 ++++++++++++++++++++++++++++++++++++++--- 2 files changed, 362 insertions(+), 24 deletions(-) diff --git a/mbuild/path/build.py b/mbuild/path/build.py index 140c0df13..7151827e1 100644 --- a/mbuild/path/build.py +++ b/mbuild/path/build.py @@ -3,6 +3,7 @@ import logging import math import time +from itertools import combinations_with_replacement import networkx as nx import numpy as np @@ -23,6 +24,7 @@ get_second_point, ) from mbuild.path.termination import NumSites, Termination, Terminator +from mbuild.utils.io import import_ logger = logging.getLogger(__name__) @@ -318,11 +320,14 @@ def to_mol(self): return _to_mol(self) - def to_mol3000(self): + def to_mol3000(self, G=None): """Convert mBuild Path to SDF/MOL V3000 format.""" from mbuild.path.formats import to_mol3000 as _to_mol3000 - return _to_mol3000(self) + if G is None: + G = self.bond_graph + + return _to_mol3000(self, G) def visualize(self, radius, hide_periodic_bonds=False): """Visualize a path with py3Dmol. @@ -338,17 +343,34 @@ def visualize(self, radius, hide_periodic_bonds=False): return visualize_path(self, radius, hide_periodic_bonds) - def relax(self, bead_radius, bond_length=None, steps=1000, seed=1, nthreads=1): - """Perform a dpd simulation to relax the current path. + def relax( + self, radius, bonds=None, angles=None, box=None, steps=1000, seed=1, nthreads=1 + ): + """Perform an OpenMM simulation to relax the current path. Runs a short energy minimization simulation in OpenMM. Parameters ---------- - bead_radius : float - Bead size set in the simulation. - bond_length : float, optional - Bond length used for all bonds. + radius : float or dictionary + Bead sigma size set in the simulation in nm. + If a dict is passed key="bead_name" and value is float in nm for the given bead radius. + bonds : float or dictionary or string, optional + Bond lengths used for all bonds. + If a dictionary is passed, key=(bead_name, bead_name) and + value is {'r': float, 'k': float}. r is in units of nm and k is in units of kcal/nm**2. + If a float is passed, that is set as r in nanometers for a harmonic bond. + If 'constrain' is passed, then bonds are constrained. + If None is passed, then no harmonic bond forces are applied.. + angles : float or dictionary or string, optional + Harmonic angle forces used for all angles. + If a dictionary is passed, key=(bead_name, central_bead_name, bead_name) and + value is {'theta': float, 'k': float}. r is in units of nm and k is in units of degrees. + If a float is passed, that is set as theta in degress for a harmonic angle. + If 'constrain' is passed, then angles are constrained. + If None is passed, no angle force is used. + box : mb.Box, optional + Box to use for periodic boundaries. steps : int, optional, default 1,000 Number of simulation steps to run. seed : int, optional, default 1 @@ -358,9 +380,139 @@ def relax(self, bead_radius, bond_length=None, steps=1000, seed=1, nthreads=1): """ from mbuild.simulation import energy_minimize_path - energy_minimize_path(self, bead_radius, bond_length, steps, seed, nthreads) + energy_minimize_path(self, radius, bonds, angles, box, steps, seed, nthreads) return + def print_bond_lengths(self, box=None): + """Compute and return info about path bonds. + + Parameters + ---------- + box : list, optional + list of [Lx, Ly, Lz] to use for checking for periodic bonds. Assume centered at (0,0,0) + + Returns + ------- + bonds : dict + Dictionary mapping (i, center, k) tuples to their angles in nm. + bonds_typesDict : dict + Dictionary mapping sorted bead-type tuples to lists of bonds. + """ + positions = self.coordinates + bond_lengths = {} + beads = set(self.beads) + bond_types = list(combinations_with_replacement(beads, 2)) + bond_typesDict = { + tuple(sorted((str(b1), str(b2)))): [] for b1, b2 in bond_types + } + + for i, j in self.bond_graph.edges(): + delta = positions[j] - positions[i] + if box is not None: + box_arr = np.array(box) + delta -= np.round(delta / box_arr) * box_arr + bl = np.linalg.norm(delta) + bond_lengths[(i, j)] = bl + bond_name = tuple(sorted((str(self.beads[i]), str(self.beads[j])))) + bond_typesDict[bond_name].append(bl) + + # Summary stats + print(f"Min bond length: {min(bond_lengths.values()):.4f}") + print(f"Max bond length: {max(bond_lengths.values()):.4f}") + print(f"Mean bond length: {np.mean(list(bond_lengths.values())):.4f}") + for key in bond_typesDict: + if not len(bond_typesDict[key]): + continue + print( + f"{key}: Max={max(bond_typesDict[key]):.2f}, Min={min(bond_typesDict[key]):.2f}" + ) + + return bond_lengths, bond_typesDict + + def print_angle_lengths(self, box=None): + """Compute and return info about path angles. + + Parameters + ---------- + box : array-like of shape (3,), optional + Box dimensions [Lx, Ly, Lz]. If None, no periodic wrapping is applied. + + Returns + ------- + angles : dict + Dictionary mapping (i, center, k) tuples to their angles in degrees. + angle_typesDict : dict + Dictionary mapping sorted bead-type tuples to lists of angles. + """ + from itertools import combinations, combinations_with_replacement + + positions = self.coordinates + angles = {} + beads = set(self.beads) + + # Build angle type keys: (outer, center, outer) with outer pair sorted + angle_types = set() + for center_bead in beads: + for b1, b2 in combinations_with_replacement(beads, 2): + angle_types.add((str(b1), str(center_bead), str(b2))) + angle_typesDict = {key: [] for key in angle_types} + + for center_node in self.bond_graph.nodes(): + neighbors = list(self.bond_graph.neighbors(center_node)) + if len(neighbors) < 2: + continue + + for i_node, k_node in combinations(neighbors, 2): + # Vector from center to i + delta_i = positions[i_node] - positions[center_node] + if box is not None: + box_arr = np.array(box) + delta_i -= np.round(delta_i / box_arr) * box_arr + + # Vector from center to k + delta_k = positions[k_node] - positions[center_node] + if box is not None: + delta_k -= np.round(delta_k / box_arr) * box_arr + + # Compute angle via dot product + cos_angle = np.dot(delta_i, delta_k) / ( + np.linalg.norm(delta_i) * np.linalg.norm(delta_k) + ) + cos_angle = np.clip(cos_angle, -1.0, 1.0) + angle_deg = np.degrees(np.arccos(cos_angle)) + + angles[(i_node, center_node, k_node)] = angle_deg + + # Categorize by type (sort outer beads for consistency) + outer_beads = tuple( + sorted((str(self.beads[i_node]), str(self.beads[k_node]))) + ) + center_bead = str(self.beads[center_node]) + angle_key = (outer_beads[0], center_bead, outer_beads[1]) + angle_typesDict[angle_key].append(angle_deg) + + # Summary stats + print(f"Min angle: {min(angles.values()):.2f}°") + print(f"Max angle: {max(angles.values()):.2f}°") + print(f"Mean angle: {np.mean(list(angles.values())):.2f}°") + for key in angle_typesDict: + if not len(angle_typesDict[key]): + continue + print( + f"{key}: Max={max(angle_typesDict[key]):.2f}°, Min={min(angle_typesDict[key]):.2f}°, Mean={np.mean(angle_typesDict[key]):.2f}°" + ) + + return angles, angle_typesDict + + def freud_rdf(self, box, bins=50, r_max=1): + freud = import_("freud") + if len(box) == 3: + box = np.array([*box, 0, 0, 0]) # assume orthorhombic + + rdf = freud.density.RDF(bins=bins, r_max=r_max) + rdf.compute(system=(box, self.coordinates)) + return rdf + def lamellar( path=None, @@ -1497,7 +1649,32 @@ def get_reference_points(path, initial_point): # Calculate position for new crosslink node (centroid of selected beads) selected_coords = np.array(path.coordinates[selected_nodes]) - crosslink_position = np.mean(selected_coords, axis=0) + if volume_constraint is not None: + # Unwrap all selected coords relative to the first one + ref = selected_coords[0] + unwrapped = np.empty_like(selected_coords) + unwrapped[0] = ref + for k in range(1, len(selected_coords)): + delta = selected_coords[k] - ref + if pbc[0]: + delta[0] -= np.round(delta[0] / box_lengths[0]) * box_lengths[0] + if pbc[1]: + delta[1] -= np.round(delta[1] / box_lengths[1]) * box_lengths[1] + if pbc[2]: + delta[2] -= np.round(delta[2] / box_lengths[2]) * box_lengths[2] + unwrapped[k] = ref + delta + + crosslink_position = np.mean(unwrapped, axis=0) + + # Wrap back into box [-box/2, box/2] + for dim in range(3): + if pbc[dim]: + crosslink_position[dim] -= ( + np.round(crosslink_position[dim] / box_lengths[dim]) + * box_lengths[dim] + ) + else: + crosslink_position = np.mean(selected_coords, axis=0) # Add new node to path path.append_coordinates(crosslink_position, bead_name) diff --git a/mbuild/simulation.py b/mbuild/simulation.py index dae819e60..59af0eccf 100644 --- a/mbuild/simulation.py +++ b/mbuild/simulation.py @@ -3,6 +3,7 @@ import logging import os import tempfile +from itertools import combinations from warnings import warn import gmso @@ -1255,8 +1256,46 @@ def _energy_minimize_openbabel( def energy_minimize_path( - path, bead_size=0.3, bond_length=None, steps=1000, seed=1, nthreads=1 + path, + radius=0.3, + bonds=None, + angles=None, + box=None, + steps=1000, + seed=1, + nthreads=1, ): + """Perform an OpenMM simulation to relax the current path. + + Runs a short energy minimization simulation in OpenMM. + + Parameters + ---------- + radius : float + Bead sigma size set in the simulation. + bonds : float or dictionary or string, optional + Bond lengths used for all bonds. + If a dictionary is passed, key=(bead_name, bead_name) and + value is {'r': float, 'k': float}. r is in units of nm and k is in units of kcal/nm**2. + If a float is passed, that is set as r in nanometers for a harmonic bond. + If 'constrain' is passed, then bonds are constrained. + If None is passed, then no harmonic bond forces are applied.. + angles : float or dictionary or string, optional + Harmonic angle forces used for all angles. + If a dictionary is passed, key=(bead_name, central_bead_name, bead_name) and + value is {'theta': float, 'k': float}. r is in units of nm and k is in units of degrees. + If a float is passed, that is set as theta in degress for a harmonic angle. + If 'constrain' is passed, then angles are constrained. + If None is passed, no angle force is used. + box : mb.Box, optional + Box to use for periodic boundaries. + steps : int, optional, default 1,000 + Number of simulation steps to run. + seed : int, optional, default 1 + Random seed for integrator. + nthreads : int, optional, default 1 + Number of threads to use during OpenMM simulation. + """ # TODO: Set equilibrium or constrained angle positions = path.coordinates n_particles = len(positions) @@ -1274,7 +1313,7 @@ def energy_minimize_path( system = openmm.System() for i in range(n_particles): - system.addParticle(1.0 * u.amu) + system.addParticle(12.0 * u.amu) # default mass of carbon # Create a Langenvin Integrator in OpenMM integrator = LangevinIntegrator( @@ -1286,29 +1325,142 @@ def energy_minimize_path( # Set nonbonded force for each particle nonbonded_force = openmm.NonbondedForce() - nonbonded_force.setNonbondedMethod(openmm.NonbondedForce.NoCutoff) - sigma = bead_size * u.nanometer + nonbonded_force.setNonbondedMethod(openmm.NonbondedForce.CutoffPeriodic) + nonbonded_force.setCutoffDistance(0.3 * u.nanometer) epsilon = 0.1 * u.kilocalories_per_mole - for i in range(n_particles): - nonbonded_force.addParticle(0.0, sigma, epsilon) + mass = 12 * u.amu + if isinstance(radius, dict): + # Per-bead-type sigma + for i in range(n_particles): + bead_name = path.beads[i] + if bead_name in radius: + sigma = radius[bead_name] * u.nanometer + else: + sigma = 0.1 * u.nanometer # fallback default + nonbonded_force.addParticle(mass, sigma, epsilon) + else: + # Single sigma for all particles + sigma = radius * u.nanometer + for i in range(n_particles): + nonbonded_force.addParticle(mass, sigma, epsilon) + system.addForce(nonbonded_force) - if not bond_length: + # --- Bond forces --- + if bonds == "constrain": + # Constrain all bonds at current lengths for i, j in path.bond_graph.edges(): - bond_length = np.linalg.norm(positions[j] - positions[i]) - system.addConstraint(i, j, bond_length * u.nanometer) - else: - # Use HarmonicBondForce instead of constraints + delta = positions[j] - positions[i] + if box: + box_arr = np.array(box) + delta -= np.round(delta / box_arr) * box_arr + bl = np.linalg.norm(delta) + system.addConstraint(i, j, bl * u.nanometer) + + elif isinstance(bonds, dict): + # Per-type harmonic bonds harmonic_force = openmm.HarmonicBondForce() + if box: + harmonic_force.setUsesPeriodicBoundaryConditions(True) + + for i, j in path.bond_graph.edges(): + i_bead = path.beads[i] + j_bead = path.beads[j] + + bond_key = (i_bead, j_bead) + bond_key_rev = (j_bead, i_bead) + + params = None + if bond_key in bonds: + params = bonds[bond_key] + elif bond_key_rev in bonds: + params = bonds[bond_key_rev] + + if params is not None: + harmonic_force.addBond( + i, + j, + params["r"] * u.nanometer, + params["k"] * u.kilocalories_per_mole / u.nanometer**2, + ) + else: + # Fallback: constrain at current distance + delta = positions[j] - positions[i] + if box: + box_arr = np.array(box) + delta -= np.round(delta / box_arr) * box_arr + bl = np.linalg.norm(delta) + system.addConstraint(i, j, bl * u.nanometer) + + system.addForce(harmonic_force) + + elif isinstance(bonds, (int, float)): + # Single float: same equilibrium length for all bonds + harmonic_force = openmm.HarmonicBondForce() + if box: + harmonic_force.setUsesPeriodicBoundaryConditions(True) for i, j in path.bond_graph.edges(): harmonic_force.addBond( i, j, - bond_length * u.nanometer, + bonds * u.nanometer, 300 * u.kilocalories_per_mole / u.nanometer**2, ) system.addForce(harmonic_force) + # --- Angle forces --- + if angles == "constrain": + # Constrain angles by fixing the 1-3 distance + for center_node in path.bond_graph.nodes(): + neighbors = list(path.bond_graph.neighbors(center_node)) + if len(neighbors) < 2: + continue + for i_node, k_node in combinations(neighbors, 2): + delta = positions[k_node] - positions[i_node] + if box: + box_arr = np.array(box) + delta -= np.round(delta / box_arr) * box_arr + dist_13 = np.linalg.norm(delta) + system.addConstraint(i_node, k_node, dist_13 * u.nanometer) + + elif isinstance(angles, dict): + # Per-type harmonic angles + angle_force = openmm.HarmonicAngleForce() + if box: + angle_force.setUsesPeriodicBoundaryConditions(True) + + for center_node in path.bond_graph.nodes(): + neighbors = list(path.bond_graph.neighbors(center_node)) + if len(neighbors) < 2: + continue + + center_bead = path.beads[center_node] + + for i_node, k_node in combinations(neighbors, 2): + i_bead = path.beads[i_node] + k_bead = path.beads[k_node] + + angle_key = (i_bead, center_bead, k_bead) + angle_key_rev = (k_bead, center_bead, i_bead) + + params = None + if angle_key in angles: + params = angles[angle_key] + elif angle_key_rev in angles: + params = angles[angle_key_rev] + + if params is not None: + theta_rad = np.radians(params["theta"]) + angle_force.addAngle( + i_node, + center_node, + k_node, + theta_rad * u.radians, + params["k"] * u.kilocalories_per_mole / u.radian**2, + ) + + system.addForce(angle_force) + topology = Topology() chain = topology.addChain() for i in range(n_particles): @@ -1320,14 +1472,23 @@ def energy_minimize_path( topology.addBond(atomsList[i], atomsList[j]) simulation = Simulation(topology, system, integrator, platform) - simulation.context.setPositions(positions) + simulation.context.setPositions(positions * u.nanometer) + if box: + simulation.context.setPeriodicBoxVectors( + [box[0], 0, 0] * u.nanometer, + [0, box[1], 0] * u.nanometer, + [0, 0, box[2]] * u.nanometer, + ) # Run energy minimization through OpenMM - simulation.minimizeEnergy(maxIterations=steps) + simulation.minimizeEnergy( + maxIterations=steps, tolerance=1e-4 * u.kilojoules_per_mole / u.nanometer + ) # Get positions directly state = simulation.context.getState(getPositions=True) pos = np.array(state.getPositions(asNumpy=True)) + if np.any(np.isnan(pos)): logger.warning("Unable to energy minimize. Nan values detected.") else: From 2691d6d94811b8f129e344b2a940ef31a472faa2 Mon Sep 17 00:00:00 2001 From: CalCraven Date: Wed, 6 May 2026 18:07:22 -0500 Subject: [PATCH 02/11] replace charge with mass --- mbuild/simulation.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/mbuild/simulation.py b/mbuild/simulation.py index 59af0eccf..58402c6dd 100644 --- a/mbuild/simulation.py +++ b/mbuild/simulation.py @@ -1328,7 +1328,6 @@ def energy_minimize_path( nonbonded_force.setNonbondedMethod(openmm.NonbondedForce.CutoffPeriodic) nonbonded_force.setCutoffDistance(0.3 * u.nanometer) epsilon = 0.1 * u.kilocalories_per_mole - mass = 12 * u.amu if isinstance(radius, dict): # Per-bead-type sigma for i in range(n_particles): @@ -1337,12 +1336,12 @@ def energy_minimize_path( sigma = radius[bead_name] * u.nanometer else: sigma = 0.1 * u.nanometer # fallback default - nonbonded_force.addParticle(mass, sigma, epsilon) + nonbonded_force.addParticle(0.0, sigma, epsilon) else: # Single sigma for all particles sigma = radius * u.nanometer for i in range(n_particles): - nonbonded_force.addParticle(mass, sigma, epsilon) + nonbonded_force.addParticle(0.0, sigma, epsilon) system.addForce(nonbonded_force) From 640e926a4ed11ae522d19b980a8491a5a72f7cbd Mon Sep 17 00:00:00 2001 From: CalCraven Date: Fri, 15 May 2026 11:28:53 -0500 Subject: [PATCH 03/11] Implement replace_sites function, used by crosslinker methods to generate multi-site crosslinks. Adjust some of the args in relax_path, add utility functions that helps to optimally place new sites into current path --- mbuild/path/build.py | 111 +- mbuild/path/crosslink.py | 2312 +++++++++++++++++++++++++++++++ mbuild/simulation.py | 23 +- mbuild/tests/test_crosslinks.py | 478 +++++++ 4 files changed, 2909 insertions(+), 15 deletions(-) create mode 100644 mbuild/path/crosslink.py create mode 100644 mbuild/tests/test_crosslinks.py diff --git a/mbuild/path/build.py b/mbuild/path/build.py index 7151827e1..412b4fb69 100644 --- a/mbuild/path/build.py +++ b/mbuild/path/build.py @@ -9,7 +9,7 @@ import numpy as np from scipy.interpolate import interp1d -from mbuild import Compound +from mbuild import Box, Compound from mbuild.exceptions import PathConvergenceError from mbuild.path.constraints import CuboidConstraint, CylinderConstraint from mbuild.path.path_utils import ( @@ -119,6 +119,9 @@ def __add__(self, other): ) # TODO: Don't overwrite nodes in bg return Path(coordinates, bond_graph, beads) + def __len__(self): + return len(self.coordinates) + @classmethod def from_compound(cls, compound): coordinates = compound.xyz @@ -268,7 +271,7 @@ def form_linear_bond_graph(self, indices=None): def add_edge(self, u, v): """Add an edge to the Path's bond graph.""" - bond_vec = self.coordinates[v] - self.coordinates[u] + bond_vec = (self.coordinates[v] - self.coordinates[u]).astype(float) bond_length = np.linalg.norm(bond_vec) bond_vec /= bond_length self.bond_graph.add_edge( @@ -504,15 +507,113 @@ def print_angle_lengths(self, box=None): return angles, angle_typesDict - def freud_rdf(self, box, bins=50, r_max=1): + def freud_rdf(self, box=None, bins=50, r_max=1): freud = import_("freud") - if len(box) == 3: + if box is None: + box = np.array([np.inf, np.inf, np.inf, 0, 0, 0]) # infinite box + elif isinstance(box, Box): + box = np.array([*box.lengths, 0, 0, 0]) + elif len(box) == 3: box = np.array([*box, 0, 0, 0]) # assume orthorhombic rdf = freud.density.RDF(bins=bins, r_max=r_max) rdf.compute(system=(box, self.coordinates)) return rdf + def replace_sites( + self, + replacement, + sites=None, + bead_name=None, + bond_length=None, + tolerance=0.1, + volume_constraint=None, + n_rotation_samples=36, + overlap_radius=None, + seed=42, + ): + """Replace sites in this Path with a multi-site substructure, in place. + + Uses rigid-body alignment to guarantee bond lengths from connection + sites to neighbors. + + Parameters + ---------- + replacement : CrosslinkerGeometry + The replacement structure. + sites : list of int, optional + Node indices to replace. + bead_name : str, optional + Replace all sites matching this name. + bond_length : float, optional + Desired bond length from connection sites to neighbors. + If None, preserves existing distances. + tolerance : float, default 0.1 + Fractional tolerance on bond lengths. + volume_constraint : optional + For periodic boundary handling. + n_rotation_samples : int, default 36 + Angular samples for overlap minimization. + overlap_radius : float, optional + Radius for overlap detection. + seed : int, default 42 + + Returns + ------- + Path + self (modified in place). + """ + from mbuild.path.crosslink import replace_sites + + return replace_sites( + self, + replacement=replacement, + sites=sites, + bead_name=bead_name, + bond_length=bond_length, + tolerance=tolerance, + volume_constraint=volume_constraint, + n_rotation_samples=n_rotation_samples, + overlap_radius=overlap_radius, + seed=seed, + ) + + def crosslink( + self, + crosslinker=None, + bead_name="_R", + backbone_name="_A", + radius=0.1, + excluded_bond_depth=2, + n_connection_sites=2, + volume_constraint=None, + initial_point=None, + seed=42, + chunk_size=512, + run_on_gpu=False, + n_rotation_samples=36, + overlap_radius=None, + ): + """Add a crosslink to this path. See module-level crosslink() for details.""" + from mbuild.path.crosslink import crosslink as crosslink2 + + return crosslink2( + self, + crosslinker=crosslinker, + bead_name=bead_name, + backbone_name=backbone_name, + radius=radius, + excluded_bond_depth=excluded_bond_depth, + n_connection_sites=n_connection_sites, + volume_constraint=volume_constraint, + initial_point=initial_point, + seed=seed, + chunk_size=chunk_size, + run_on_gpu=run_on_gpu, + n_rotation_samples=n_rotation_samples, + overlap_radius=overlap_radius, + ) + def lamellar( path=None, @@ -1018,7 +1119,7 @@ def hard_sphere_random_walk( radius : float, default 0.1 nm Radius of sites used in checking for overlaps. rw_angles : tuple or dict or np.array or AnglesSampler, default None - Set the angle sampling method. A tuple of (min_val, max_val) sets the uniform distribution. + Set the angle sampling method in units of radians. A tuple of (min_val, max_val) sets the uniform distribution. The default value of None sets the uniform angle sampling from (np.pi/2, np.pi). Can also use a Gaussian distribution by passing a dict with keys {'loc':mean, 'scale':std}. Finally, a numpy array of 1D or 2D array of numpy values can be passed, which will be sampled diff --git a/mbuild/path/crosslink.py b/mbuild/path/crosslink.py new file mode 100644 index 000000000..c6fb51491 --- /dev/null +++ b/mbuild/path/crosslink.py @@ -0,0 +1,2312 @@ +from itertools import permutations + +import networkx as nx +import numpy as np + +from mbuild.exceptions import PathConvergenceError +from mbuild.path.build import Path +from mbuild.path.constraints import CuboidConstraint, CylinderConstraint +from mbuild.path.path_utils import calculate_sq_distances + + +class CrosslinkerGeometry(Path): + """A Path subclass representing a crosslinker with connection site metadata. + + Inherits all Path functionality (coordinates, bond_graph, beads) and adds + metadata about which sites form bonds to external beads. + + Parameters + ---------- + coordinates : array-like, shape (N, 3) + Positions of crosslinker sites. Will be re-centered at origin. + bond_graph : networkx.Graph, optional + Internal connectivity. + bead_name : str or np.ndarray, default "_R" + Name(s) for each site. + connection_sites : list of int + Node indices that form bonds to external beads. + Length determines how many external bonds this geometry requires. + Repeated indices allow one site to bond to multiple external beads. + + Examples + -------- + >>> cl = CrosslinkerGeometry.equilateral_triangle(bond_length=0.27) + >>> cl.connection_sites + [0, 1, 2] + >>> cl.bond_graph.edges() + EdgeView([(0, 1), (0, 2), (1, 2)]) + """ + + def __init__( + self, + coordinates=None, + bond_graph=None, + bead_name="_R", + connection_sites=None, + ): + super().__init__( + coordinates=coordinates, bond_graph=bond_graph, bead_name=bead_name + ) + + if connection_sites is None: + connection_sites = list(range(len(self.coordinates))) + + self.connection_sites = list(connection_sites) + self._validate_connection_sites() + + # Center positions at origin + if len(self.coordinates) > 0: + centroid = np.mean(self.coordinates, axis=0) + self.coordinates = (self.coordinates - centroid).astype(np.float32) + + def _validate_connection_sites(self): + n = len(self.coordinates) + for idx in self.connection_sites: + if not (0 <= idx < n): + raise ValueError( + f"Connection site {idx} is out of range (must be 0 to {n - 1})" + ) + + @property + def n_sites(self) -> int: + return len(self.coordinates) + + @property + def n_connections(self) -> int: + return len(self.connection_sites) + + @property + def unique_connection_sites(self): + return list(dict.fromkeys(self.connection_sites)) + + @property + def internal_bonds(self): + return list(self.bond_graph.edges()) + + def copy(self): + """Return a deep copy of this crosslinker geometry.""" + return CrosslinkerGeometry( + coordinates=self.coordinates.copy(), + bond_graph=self.bond_graph.copy(), + bead_name=self.beads.copy(), + connection_sites=list(self.connection_sites), + ) + + def recenter(self): + if len(self.coordinates) > 0: + centroid = np.mean(self.coordinates, axis=0) + self.coordinates = self.coordinates - centroid + + @classmethod + def from_path(cls, path, connection_sites): + """Create from an existing Path. + + Parameters + ---------- + path : Path + An existing path defining the structure. + connection_sites : list of int + Which nodes bond to external beads. + """ + return cls( + coordinates=path.coordinates.copy(), + bond_graph=path.bond_graph.copy(), + bead_name=path.beads.copy(), + connection_sites=connection_sites, + ) + + @classmethod + def single_site(cls, bead_name="_R", n_connections=2): + """Single-site (original crosslink behavior).""" + coords = np.array([[0.0, 0.0, 0.0]], dtype=np.float32) + return cls( + coordinates=coords, + bead_name=bead_name, + connection_sites=[0] * n_connections, + ) + + @classmethod + def linear(cls, n_sites=2, bond_length=0.27, bead_name="_R", connection_sites=None): + """Linear chain of sites.""" + positions = np.zeros((n_sites, 3), dtype=np.float32) + for i in range(n_sites): + positions[i, 0] = i * bond_length + + bead_names = ( + np.array([bead_name] * n_sites, dtype="U10") + if isinstance(bead_name, str) + else np.array(bead_name, dtype="U10") + ) + if connection_sites is None: + connection_sites = [0, n_sites - 1] + + G = nx.Graph() + for i in range(n_sites): + G.add_node(i) + for i in range(n_sites - 1): + G.add_edge(i, i + 1) + + return cls( + coordinates=positions, + bond_graph=G, + bead_name=bead_names, + connection_sites=connection_sites, + ) + + @classmethod + def equilateral_triangle( + cls, bond_length=0.27, bead_name="_R", connection_sites=None + ): + """Three sites in a flat equilateral triangle.""" + R = bond_length / np.sqrt(3) + angles = [0, 2 * np.pi / 3, 4 * np.pi / 3] + positions = np.array( + [[R * np.cos(a), R * np.sin(a), 0.0] for a in angles], dtype=np.float32 + ) + + bead_names = ( + np.array([bead_name] * 3, dtype="U10") + if isinstance(bead_name, str) + else np.array(bead_name, dtype="U10") + ) + if connection_sites is None: + connection_sites = [0, 1, 2] + + G = nx.Graph() + for i in range(3): + G.add_node(i) + G.add_edge(0, 1) + G.add_edge(1, 2) + G.add_edge(0, 2) + + return cls( + coordinates=positions, + bond_graph=G, + bead_name=bead_names, + connection_sites=connection_sites, + ) + + @classmethod + def square(cls, bond_length=0.27, bead_name="_R", connection_sites=None): + """Four sites in a flat square.""" + half = bond_length / 2.0 + positions = np.array( + [[half, half, 0], [-half, half, 0], [-half, -half, 0], [half, -half, 0]], + dtype=np.float32, + ) + bead_names = ( + np.array([bead_name] * 4, dtype="U10") + if isinstance(bead_name, str) + else np.array(bead_name, dtype="U10") + ) + if connection_sites is None: + connection_sites = [0, 1, 2, 3] + + G = nx.Graph() + for i in range(4): + G.add_node(i) + for i in range(4): + G.add_edge(i, (i + 1) % 4) + + return cls( + coordinates=positions, + bond_graph=G, + bead_name=bead_names, + connection_sites=connection_sites, + ) + + @classmethod + def tetrahedral(cls, bond_length=0.27, bead_name="_R", connection_sites=None): + """Four sites in a regular tetrahedron.""" + positions = np.array( + [[1, 1, 1], [1, -1, -1], [-1, 1, -1], [-1, -1, 1]], dtype=np.float32 + ) + current_dist = np.linalg.norm(positions[0] - positions[1]) + positions *= bond_length / current_dist + + bead_names = ( + np.array([bead_name] * 4, dtype="U10") + if isinstance(bead_name, str) + else np.array(bead_name, dtype="U10") + ) + if connection_sites is None: + connection_sites = [0, 1, 2, 3] + + G = nx.Graph() + for i in range(4): + G.add_node(i) + for i in range(4): + for j in range(i + 1, 4): + G.add_edge(i, j) + + return cls( + coordinates=positions, + bond_graph=G, + bead_name=bead_names, + connection_sites=connection_sites, + ) + + @classmethod + def trigonal_bipyramidal( + cls, bond_length=0.27, bead_name="_R", connection_sites=None + ): + """Five sites in trigonal bipyramidal arrangement (0-2 equatorial, 3-4 axial).""" + R_eq = bond_length / np.sqrt(3) + eq_angles = [0, 2 * np.pi / 3, 4 * np.pi / 3] + equatorial = [[R_eq * np.cos(a), R_eq * np.sin(a), 0.0] for a in eq_angles] + axial_dist = bond_length * np.sqrt(2.0 / 3.0) + axial = [[0.0, 0.0, axial_dist], [0.0, 0.0, -axial_dist]] + positions = np.array(equatorial + axial, dtype=np.float32) + + bead_names = ( + np.array([bead_name] * 5, dtype="U10") + if isinstance(bead_name, str) + else np.array(bead_name, dtype="U10") + ) + if connection_sites is None: + connection_sites = [0, 1, 2, 3, 4] + + G = nx.Graph() + for i in range(5): + G.add_node(i) + G.add_edge(0, 1) + G.add_edge(1, 2) + G.add_edge(0, 2) + for eq in range(3): + G.add_edge(eq, 3) + G.add_edge(eq, 4) + + return cls( + coordinates=positions, + bond_graph=G, + bead_name=bead_names, + connection_sites=connection_sites, + ) + + @classmethod + def pentagon(cls, bond_length=0.27, bead_name="_R", connection_sites=None): + """Five sites in a flat regular pentagon.""" + R = bond_length / (2 * np.sin(np.pi / 5)) + angles = [2 * np.pi * i / 5 for i in range(5)] + positions = np.array( + [[R * np.cos(a), R * np.sin(a), 0.0] for a in angles], dtype=np.float32 + ) + + bead_names = ( + np.array([bead_name] * 5, dtype="U10") + if isinstance(bead_name, str) + else np.array(bead_name, dtype="U10") + ) + if connection_sites is None: + connection_sites = [0, 1, 2, 3, 4] + + G = nx.Graph() + for i in range(5): + G.add_node(i) + for i in range(5): + G.add_edge(i, (i + 1) % 5) + + return cls( + coordinates=positions, + bond_graph=G, + bead_name=bead_names, + connection_sites=connection_sites, + ) + + @classmethod + def from_edges(cls, coordinates, edges, bead_name="_R", connection_sites=None): + """Create from explicit positions and edge list. + + Parameters + ---------- + coordinates : array-like, shape (N, 3) + edges : list of tuple(int, int) + bead_name : str or list of str + connection_sites : list of int, optional + """ + coordinates = np.asarray(coordinates, dtype=np.float32) + n = len(coordinates) + bead_names = ( + np.array([bead_name] * n, dtype="U10") + if isinstance(bead_name, str) + else np.array(bead_name, dtype="U10") + ) + if connection_sites is None: + connection_sites = list(range(n)) + + G = nx.Graph() + for i in range(n): + G.add_node(i) + for i, j in edges: + G.add_edge(i, j) + + return cls( + coordinates=coordinates, + bond_graph=G, + bead_name=bead_names, + connection_sites=connection_sites, + ) + + +# ============================================================================= +# Orientation / Rotation Utilities +# ============================================================================= + + +def _kabsch_rotation(P, Q): + """Optimal rotation aligning P onto Q (Kabsch algorithm).""" + H = P.T @ Q + U, S, Vt = np.linalg.svd(H) + d = np.linalg.det(Vt.T @ U.T) + sign_matrix = np.diag([1.0, 1.0, d]) + R = Vt.T @ sign_matrix @ U.T + return R.astype(np.float32) + + +def _rotation_between_vectors(a, b): + """Rotation matrix rotating unit vector a onto b (Rodrigues).""" + v = np.cross(a, b) + c = float(np.dot(a, b)) + s = float(np.linalg.norm(v)) + + if s < 1e-10: + if c > 0: + return np.eye(3, dtype=np.float32) + perp = np.array([1, 0, 0], dtype=np.float32) + if abs(np.dot(a, perp)) > 0.9: + perp = np.array([0, 1, 0], dtype=np.float32) + perp = perp - np.dot(perp, a) * a + perp /= np.linalg.norm(perp) + return (2 * np.outer(perp, perp) - np.eye(3)).astype(np.float32) + + vx = np.array( + [[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]], dtype=np.float32 + ) + R = np.eye(3, dtype=np.float32) + vx + vx @ vx * ((1 - c) / (s * s)) + return R + + +def _random_rotation_matrix(rng): + """Uniform random rotation via QR decomposition.""" + H = rng.standard_normal((3, 3)).astype(np.float32) + Q, R = np.linalg.qr(H) + Q *= np.sign(np.diag(R)) + if np.linalg.det(Q) < 0: + Q[:, 0] *= -1 + return Q.astype(np.float32) + + +def _rotation_about_axis(axis, angle): + """Rotation matrix for rotation by `angle` radians about `axis`.""" + axis = np.asarray(axis, dtype=np.float32) + axis = axis / np.linalg.norm(axis) + K = np.array( + [ + [0, -axis[2], axis[1]], + [axis[2], 0, -axis[0]], + [-axis[1], axis[0], 0], + ], + dtype=np.float32, + ) + R = np.eye(3, dtype=np.float32) + np.sin(angle) * K + (1 - np.cos(angle)) * (K @ K) + return R + + +def _orient_replacement( + replacement, + neighbor_vectors, + nearby_coords=None, + center=None, + rng=None, + n_rotation_samples=36, +): + """Orient a replacement geometry so connection sites point toward neighbors. + + Performs Kabsch alignment then searches residual rotational DOF to + minimize overlaps with nearby non-bonded sites. + + Parameters + ---------- + replacement : CrosslinkerGeometry + The replacement structure (centered at origin). + neighbor_vectors : np.ndarray, shape (n_connections, 3) + Vectors from placement center to each neighbor being bonded. + nearby_coords : np.ndarray, shape (M, 3), optional + Coordinates of nearby non-bonded sites (for overlap minimization). + Should be relative to the placement center. + center : np.ndarray, shape (3,), optional + The absolute position where replacement is centered (for overlap calc). + rng : numpy.random.Generator, optional + n_rotation_samples : int, default 36 + Number of rotational samples for overlap minimization. + + Returns + ------- + oriented_positions : np.ndarray, shape (n_sites, 3) + Rotated replacement positions (still centered at origin). + assignment : list of int + assignment[i] = j means connection_sites[i] bonds to neighbor j. + """ + if rng is None: + rng = np.random.default_rng(42) + + positions = replacement.coordinates.copy() + unique_conn = replacement.unique_connection_sites + conn_positions = positions[unique_conn] + + n_unique = len(unique_conn) + + # --- Degenerate case: all connection sites at origin --- + if np.allclose(conn_positions, 0, atol=1e-8): + R = _random_rotation_matrix(rng) + oriented = (R @ positions.T).T + return oriented, list(range(replacement.n_connections)) + + # --- Group target vectors by unique connection site (for repeated sites) --- + conn_site_to_targets = {} + for i, site_idx in enumerate(replacement.connection_sites): + conn_site_to_targets.setdefault(site_idx, []).append(i) + + # Use centroid of targets for each unique connection site + unique_targets = np.array( + [ + np.mean(neighbor_vectors[conn_site_to_targets[s]], axis=0) + for s in unique_conn + ], + dtype=np.float32, + ) + + # --- Normalize for angular alignment --- + P_norms = np.linalg.norm(conn_positions, axis=1, keepdims=True) + P_normed = conn_positions / np.maximum(P_norms, 1e-10) + + T_norms = np.linalg.norm(unique_targets, axis=1, keepdims=True) + T_normed = unique_targets / np.maximum(T_norms, 1e-10) + + # --- Single connection site case --- + if n_unique == 1: + R = _rotation_between_vectors(P_normed[0], T_normed[0]) + best_R = R + best_perm = [0] + else: + # --- Try all permutations (feasible for n <= 5) --- + best_perm = list(range(n_unique)) + best_R = np.eye(3, dtype=np.float32) + best_cost = np.inf + + for perm in permutations(range(n_unique)): + Q_perm = T_normed[list(perm)] + R = _kabsch_rotation(P_normed, Q_perm) + rotated = (R @ P_normed.T).T + cost = np.sum((rotated - Q_perm) ** 2) + if cost < best_cost: + best_cost = cost + best_perm = list(perm) + best_R = R + + oriented = (best_R @ positions.T).T + + # --- Overlap minimization: search residual rotation --- + if nearby_coords is not None and len(nearby_coords) > 0 and n_rotation_samples > 1: + # Determine a rotation axis (mean direction of connection sites after alignment) + rotated_conn = oriented[unique_conn] + mean_axis = np.mean(rotated_conn, axis=0) + mean_axis_norm = np.linalg.norm(mean_axis) + + if mean_axis_norm > 1e-8: + rotation_axis = mean_axis / mean_axis_norm + else: + # Use normal to the plane of connection sites or random axis + rotation_axis = np.array([0, 0, 1], dtype=np.float32) + + best_oriented = oriented.copy() + best_min_dist = _min_distance_to_neighbors(oriented, nearby_coords) + + angles = np.linspace(0, 2 * np.pi, n_rotation_samples, endpoint=False) + for angle in angles: + R_trial = _rotation_about_axis(rotation_axis, angle) + trial = (R_trial @ oriented.T).T + min_dist = _min_distance_to_neighbors(trial, nearby_coords) + if min_dist > best_min_dist: + best_min_dist = min_dist + best_oriented = trial + + oriented = best_oriented + + # --- Build full assignment --- + full_assignment = _build_full_assignment( + replacement, best_perm, unique_conn, conn_site_to_targets + ) + + return oriented, full_assignment + + +def _orient_replacement_partial( + replacement, + partial_conn_sites, + neighbor_vectors, + nearby_coords=None, + center=None, + rng=None, + n_rotation_samples=36, +): + """Orient a replacement when fewer neighbors exist than connection sites. + + Uses only the first ``len(neighbor_vectors)`` connection sites for alignment. + The remaining connection sites are left dangling (not bonded). + + Parameters + ---------- + replacement : CrosslinkerGeometry + The replacement structure (centered at origin). + partial_conn_sites : list of int + Subset of replacement.connection_sites to use for alignment. + Length matches len(neighbor_vectors). + neighbor_vectors : np.ndarray, shape (degree, 3) + Vectors from placement center to each neighbor being bonded. + nearby_coords : np.ndarray, shape (M, 3), optional + Coordinates of nearby non-bonded sites for overlap minimization. + center : np.ndarray, shape (3,), optional + Absolute position of placement center. + rng : numpy.random.Generator, optional + n_rotation_samples : int, default 36 + Number of rotational samples for overlap minimization. + + Returns + ------- + oriented_positions : np.ndarray, shape (n_sites, 3) + Rotated replacement positions (still centered at origin). + assignment : list of int + assignment[i] = j means partial_conn_sites[i] bonds to neighbor j. + Length equals len(partial_conn_sites) (i.e., degree). + """ + if rng is None: + rng = np.random.default_rng(42) + + positions = replacement.coordinates.copy() + degree = len(neighbor_vectors) + + if degree == 0: + R = _random_rotation_matrix(rng) + return (R @ positions.T).T, [] + + # Get positions of the partial connection sites + unique_partial = list(dict.fromkeys(partial_conn_sites)) + conn_positions = positions[unique_partial] + + # If all connection sites at origin, random rotation + if np.allclose(conn_positions, 0, atol=1e-8): + R = _random_rotation_matrix(rng) + oriented = (R @ positions.T).T + return oriented, list(range(degree)) + + # Group partial connections by unique site + conn_site_to_targets = {} + for i, site_idx in enumerate(partial_conn_sites): + conn_site_to_targets.setdefault(site_idx, []).append(i) + + # Use centroid of targets for each unique connection site + unique_targets = np.array( + [ + np.mean(neighbor_vectors[conn_site_to_targets[s]], axis=0) + for s in unique_partial + ], + dtype=np.float32, + ) + + # Normalize for angular alignment + P_norms = np.linalg.norm(conn_positions, axis=1, keepdims=True) + P_normed = conn_positions / np.maximum(P_norms, 1e-10) + + T_norms = np.linalg.norm(unique_targets, axis=1, keepdims=True) + T_normed = unique_targets / np.maximum(T_norms, 1e-10) + + n_unique = len(unique_partial) + + if n_unique == 1: + R = _rotation_between_vectors(P_normed[0], T_normed[0]) + best_R = R + best_perm = [0] + else: + # Try all permutations of unique sites -> unique targets + best_perm = list(range(n_unique)) + best_R = np.eye(3, dtype=np.float32) + best_cost = np.inf + + for perm in permutations(range(n_unique)): + Q_perm = T_normed[list(perm)] + R = _kabsch_rotation(P_normed, Q_perm) + rotated = (R @ P_normed.T).T + cost = np.sum((rotated - Q_perm) ** 2) + if cost < best_cost: + best_cost = cost + best_perm = list(perm) + best_R = R + + oriented = (best_R @ positions.T).T + + # --- Overlap minimization --- + if nearby_coords is not None and len(nearby_coords) > 0 and n_rotation_samples > 1: + rotated_conn = oriented[unique_partial] + mean_axis = np.mean(rotated_conn, axis=0) + mean_axis_norm = np.linalg.norm(mean_axis) + + if mean_axis_norm > 1e-8: + rotation_axis = mean_axis / mean_axis_norm + else: + rotation_axis = np.array([0, 0, 1], dtype=np.float32) + + best_oriented = oriented.copy() + best_min_dist = _min_distance_to_neighbors(oriented, nearby_coords) + + angles = np.linspace(0, 2 * np.pi, n_rotation_samples, endpoint=False) + for angle in angles: + R_trial = _rotation_about_axis(rotation_axis, angle) + trial = (R_trial @ oriented.T).T + min_dist = _min_distance_to_neighbors(trial, nearby_coords) + if min_dist > best_min_dist: + best_min_dist = min_dist + best_oriented = trial + + oriented = best_oriented + + # Build assignment: maps conn_idx (0..degree-1) -> neighbor_idx + full_assignment = _build_full_assignment_partial( + partial_conn_sites, best_perm, unique_partial, conn_site_to_targets + ) + + return oriented, full_assignment + + +def _build_full_assignment_partial( + partial_conn_sites, perm, unique_partial, conn_site_to_targets +): + """Build assignment for partial connection case. + + Parameters + ---------- + partial_conn_sites : list of int + The subset of connection sites being used. + perm : list of int + Permutation of unique sites to unique targets. + unique_partial : list of int + Unique site indices in partial_conn_sites. + conn_site_to_targets : dict + Maps site_idx -> list of connection indices using that site. + + Returns + ------- + full_assignment : list of int + assignment[i] = neighbor_index for partial_conn_sites[i]. + """ + degree = len(partial_conn_sites) + + if degree == len(unique_partial) and all( + len(v) == 1 for v in conn_site_to_targets.values() + ): + # Simple case: no repeated sites, 1:1 mapping + return perm + + # General case with possible repeated sites + full_assignment = [0] * degree + site_counter = {} + + for conn_idx, site_idx in enumerate(partial_conn_sites): + unique_pos = unique_partial.index(site_idx) + target_neighbor_group = perm[unique_pos] + + count = site_counter.get(site_idx, 0) + target_site = unique_partial[target_neighbor_group] + target_indices = conn_site_to_targets[target_site] + full_assignment[conn_idx] = target_indices[count % len(target_indices)] + site_counter[site_idx] = count + 1 + + return full_assignment + + +def _min_distance_to_neighbors(positions, nearby_coords): + """Compute minimum distance between any replacement site and nearby sites.""" + if len(nearby_coords) == 0: + return np.inf + # Pairwise distances + diffs = positions[:, None, :] - nearby_coords[None, :, :] + dists = np.linalg.norm(diffs, axis=-1) + return np.min(dists) + + +def _build_full_assignment(replacement, perm, unique_conn, conn_site_to_targets): + """Map each connection_sites entry to a neighbor index using the permutation. + + Parameters + ---------- + replacement : CrosslinkerGeometry + perm : list of int + Permutation mapping unique connection site positions to neighbor group positions. + unique_conn : list of int + Unique connection site indices. + conn_site_to_targets : dict + Maps site_idx -> list of connection indices using that site. + + Returns + ------- + full_assignment : list of int + assignment[i] = neighbor_index for connection_sites[i]. + """ + n_connections = replacement.n_connections + + if n_connections == len(unique_conn): + # No repeated sites: perm directly maps connection -> neighbor + return perm + + # With repeated sites: need to assign individual connections within groups + full_assignment = [0] * n_connections + + # perm[unique_pos] tells us which neighbor group this unique site maps to + # We need to iterate through connection_sites and assign appropriately + site_counter = {} # track how many times each site has been assigned + for conn_idx, site_idx in enumerate(replacement.connection_sites): + unique_pos = unique_conn.index(site_idx) + target_neighbor_group = perm[unique_pos] + + # Round-robin within the target group + count = site_counter.get(site_idx, 0) + # The target_neighbor_group-th unique site's target list + target_site = unique_conn[target_neighbor_group] + target_indices = conn_site_to_targets[target_site] + full_assignment[conn_idx] = target_indices[count % len(target_indices)] + site_counter[site_idx] = count + 1 + + return full_assignment + + +# ============================================================================= +# General replace_sites function +# ============================================================================= + + +def replace_sites( + path, + replacement, + sites=None, + bead_name=None, + bond_length=None, + tolerance=0.1, + min_separation=None, + volume_constraint=None, + n_rotation_samples=36, + overlap_radius=None, + seed=42, +): + """Replace sites in a Path with a multi-site substructure, in place. + + Uses rigid-body alignment to place the replacement such that: + - Connection sites are at specified bond_length from their neighbors + - Internal geometry is preserved exactly (rigid body) + - Placed sites are at least min_separation from non-bonded neighbors + + The optimization sweeps residual rotational degrees of freedom and, if + needed, offsets the placement perpendicular to the constraint axis to + find a clash-free configuration. + + Parameters + ---------- + path : Path + The path to modify in place. + replacement : CrosslinkerGeometry + The replacement structure. + sites : list of int, optional + Explicit node indices to replace. Mutually exclusive with ``bead_name``. + bead_name : str, optional + Replace all sites matching this bead name. Mutually exclusive with ``sites``. + bond_length : float, optional + Desired bond length from each connection site to its neighbor. + If None, uses the existing distance from the replaced site to each neighbor. + tolerance : float, default 0.1 + Fractional tolerance on bond lengths (±10%). + min_separation : float, optional + Minimum allowed distance between any placed crosslinker site and any + existing non-bonded site. If None, no overlap checking is performed. + volume_constraint : CuboidConstraint or CylinderConstraint, optional + For periodic boundary handling. + n_rotation_samples : int, default 36 + Angular samples for residual DOF optimization. + overlap_radius : float, optional + Radius around each replaced site to gather nearby sites for overlap check. + If None, defaults to 3x max extent of replacement geometry. + seed : int, default 42 + Random seed. + + Returns + ------- + Path + The same path object (modified in place), for method chaining. + """ + if sites is None and bead_name is None: + raise ValueError("Must specify either `sites` or `bead_name`.") + if sites is not None and bead_name is not None: + raise ValueError("Cannot specify both `sites` and `bead_name`.") + + rng = np.random.default_rng(seed) + + if bead_name is not None: + sites = [ + node for node in path.bond_graph.nodes() if path.beads[node] == bead_name + ] + sites = list(sites) + + if len(sites) == 0: + return path + + for site in sites: + degree = path.bond_graph.degree[site] + if degree > replacement.n_connections: + raise ValueError( + f"Site {site} has degree {degree} but replacement only has " + f"{replacement.n_connections} connection sites. " + f"Degree must not exceed len(replacement.connection_sites)." + ) + + pbc, box_lengths = _get_pbc_info(volume_constraint) + + if overlap_radius is None: + if replacement.n_sites > 0: + max_extent = np.max(np.linalg.norm(replacement.coordinates, axis=1)) + overlap_radius = max(max_extent * 3, 0.5) + else: + overlap_radius = 0.5 + + sites_set = set(sites) + + # --- Compute all replacements --- + replacement_data = {} + + for site in sites: + neighbors = list(path.bond_graph.neighbors(site)) + site_pos = path.coordinates[site] + degree = len(neighbors) + + # Unwrapped neighbor positions (absolute) + neighbor_coords = np.zeros((degree, 3), dtype=np.float32) + for i, nb in enumerate(neighbors): + delta = path.coordinates[nb] - site_pos + delta = _apply_mic(delta, pbc, box_lengths) + neighbor_coords[i] = site_pos + delta + + # Target bond lengths + if bond_length is not None: + target_bond_lengths = np.full(degree, bond_length, dtype=np.float32) + else: + target_bond_lengths = np.array( + [np.linalg.norm(neighbor_coords[i] - site_pos) for i in range(degree)], + dtype=np.float32, + ) + + # Nearby non-bonded coords (absolute positions for overlap checking) + nearby_coords = _get_nearby_nonbonded_absolute( + path, site, neighbors, site_pos, overlap_radius, pbc, box_lengths, sites_set + ) + + # Active connection sites (may be fewer than total if degree < n_connections) + active_conn_sites = replacement.connection_sites[:degree] + + # Place via rigid body alignment + overlap optimization + placed_positions, assignment = _place_replacement_rigid( + replacement, + active_conn_sites, + neighbor_coords, + target_bond_lengths, + nearby_coords, + rng, + n_rotation_samples, + tolerance, + min_separation, + ) + + # Wrap into PBC + placed_positions = _wrap_positions(placed_positions, pbc, box_lengths) + + replacement_data[site] = { + "positions": placed_positions, + "neighbors": neighbors, + "assignment": assignment, + } + + # --- Rebuild path in place --- + kept_nodes = [n for n in sorted(path.bond_graph.nodes()) if n not in sites_set] + + new_coords = [] + new_beads = [] + old_to_new = {} + current_idx = 0 + + for old_node in kept_nodes: + old_to_new[old_node] = current_idx + new_coords.append(path.coordinates[old_node]) + new_beads.append(path.beads[old_node]) + current_idx += 1 + + site_new_indices = {} + for site in sites: + data = replacement_data[site] + indices = [] + for i in range(replacement.n_sites): + indices.append(current_idx) + new_coords.append(data["positions"][i]) + new_beads.append(replacement.beads[i]) + current_idx += 1 + site_new_indices[site] = indices + + # --- Build edges --- + new_edges = [] + + # Kept-to-kept + for u, v, data in path.bond_graph.edges(data=True): + if u in sites_set or v in sites_set: + continue + new_edges.append((old_to_new[u], old_to_new[v], data)) + + # Internal replacement bonds + for site in sites: + indices = site_new_indices[site] + for i, j in replacement.bond_graph.edges(): + edge_data = { + "bond_type": (str(replacement.beads[i]), str(replacement.beads[j])), + } + new_edges.append((indices[i], indices[j], edge_data)) + + # Connection-to-neighbor bonds + for site in sites: + data = replacement_data[site] + indices = site_new_indices[site] + neighbors = data["neighbors"] + assignment = data["assignment"] + + for conn_idx, neighbor_idx in enumerate(assignment): + repl_site_local = replacement.connection_sites[conn_idx] + repl_new_node = indices[repl_site_local] + + external_neighbor = neighbors[neighbor_idx] + + if external_neighbor in old_to_new: + ext_new_node = old_to_new[external_neighbor] + ext_bead_name = path.beads[external_neighbor] + elif external_neighbor in site_new_indices: + other_data = replacement_data[external_neighbor] + other_neighbors = other_data["neighbors"] + try: + back_idx = other_neighbors.index(site) + except ValueError: + continue + other_assignment = other_data["assignment"] + found = False + for other_conn_idx, other_target in enumerate(other_assignment): + if other_target == back_idx: + other_repl_site_local = replacement.connection_sites[ + other_conn_idx + ] + ext_new_node = site_new_indices[external_neighbor][ + other_repl_site_local + ] + ext_bead_name = replacement.beads[other_repl_site_local] + found = True + break + if not found: + continue + else: + continue + + edge_data = { + "bond_type": ( + str(replacement.beads[repl_site_local]), + str(ext_bead_name), + ), + } + new_edges.append((repl_new_node, ext_new_node, edge_data)) + + # --- Overwrite path --- + path.coordinates = np.array(new_coords, dtype=np.float32) + path.beads = np.array(new_beads, dtype="U10") + + path.bond_graph = nx.Graph() + for i in range(len(path.coordinates)): + path.bond_graph.add_node(i, name=str(path.beads[i]), xyz=path.coordinates[i]) + + seen_edges = set() + for u, v, data in new_edges: + edge_key = (min(u, v), max(u, v)) + if edge_key not in seen_edges: + seen_edges.add(edge_key) + path.bond_graph.add_edge(u, v, **data) + + return path + + +def _get_nearby_nonbonded_absolute( + path, site, neighbors, site_pos, radius, pbc, box_lengths, excluded_sites +): + """Get absolute positions of nearby non-bonded sites for overlap checking. + + Unlike the relative version, returns absolute (unwrapped relative to site_pos) + positions suitable for direct distance comparison with placed crosslinker sites. + + Parameters + ---------- + path : Path + site : int + neighbors : list of int + site_pos : np.ndarray, shape (3,) + radius : float + pbc : array-like of bool + box_lengths : np.ndarray + excluded_sites : set of int + + Returns + ------- + np.ndarray, shape (M, 3) + """ + neighbors_set = set(neighbors) | {site} | excluded_sites + + nearby = [] + for node in path.bond_graph.nodes(): + if node in neighbors_set: + continue + delta = path.coordinates[node] - site_pos + delta = _apply_mic(delta, pbc, box_lengths) + dist = np.linalg.norm(delta) + if dist <= radius: + # Return absolute unwrapped position + nearby.append(site_pos + delta) + + if len(nearby) == 0: + return np.empty((0, 3), dtype=np.float32) + return np.array(nearby, dtype=np.float32) + + +def _place_replacement_rigid( + replacement, + active_conn_sites, + neighbor_coords, + target_bond_lengths, + nearby_coords, + rng, + n_rotation_samples, + tolerance, + min_separation=None, +): + """Place replacement via rigid-body alignment guaranteeing bond lengths and no overlaps. + + Strategy: + 1. Compute target positions for connection sites (at bond_length from neighbors). + 2. Solve Procrustes alignment (rotation + translation) to map connection sites to targets. + 3. For residual rotational DOF (1 DOF for 2 unique sites, full DOF for 1 unique site): + - Sweep angles and pick the one that: + a) Keeps all connection-to-neighbor distances within tolerance + b) Maximizes minimum distance to nearby non-bonded sites (avoids overlaps) + 4. If no rotation satisfies min_separation, offset the center along the rotation axis + (push the crosslinker "above" or "below" the backbone plane) and re-sweep. + 5. As a last resort, return the best-scoring rotation even if below min_separation. + + Parameters + ---------- + replacement : CrosslinkerGeometry + The replacement structure (centered at origin). + active_conn_sites : list of int + Subset of replacement.connection_sites to bond (length = degree). + neighbor_coords : np.ndarray, shape (degree, 3) + Unwrapped absolute positions of neighboring beads. + target_bond_lengths : np.ndarray, shape (degree,) + Desired bond length from each connection site to its neighbor. + nearby_coords : np.ndarray, shape (M, 3) + Absolute positions of nearby non-bonded sites for overlap avoidance. + rng : numpy.random.Generator + n_rotation_samples : int + tolerance : float + Fractional tolerance on bond lengths. + min_separation : float, optional + Minimum allowed distance between any placed site and nearby sites. + If None, no overlap checking is performed. + + Returns + ------- + placed_positions : np.ndarray, shape (n_sites, 3) + Final absolute positions for all replacement sites. + assignment : list of int + assignment[i] = neighbor_index for active_conn_sites[i]. + """ + degree = len(active_conn_sites) + + if degree == 0: + center = ( + np.mean(neighbor_coords, axis=0) + if len(neighbor_coords) > 0 + else np.zeros(3) + ) + R = _random_rotation_matrix(rng) + placed = (R @ replacement.coordinates.T).T + center + return placed.astype(np.float32), [] + + # --- Compute target positions for connection sites --- + neighbor_centroid = np.mean(neighbor_coords, axis=0) + + target_positions = np.zeros((degree, 3), dtype=np.float32) + for i in range(degree): + direction = neighbor_centroid - neighbor_coords[i] + d = np.linalg.norm(direction) + if d > 1e-10: + direction_hat = direction / d + else: + direction_hat = rng.standard_normal(3).astype(np.float32) + direction_hat /= np.linalg.norm(direction_hat) + target_positions[i] = ( + neighbor_coords[i] + direction_hat * target_bond_lengths[i] + ) + + # --- Get unique connection site info --- + unique_active = list(dict.fromkeys(active_conn_sites)) + unique_conn_positions = replacement.coordinates[unique_active] + n_unique = len(unique_active) + + conn_to_target_indices = {} + for i, site_idx in enumerate(active_conn_sites): + conn_to_target_indices.setdefault(site_idx, []).append(i) + + unique_targets = np.array( + [ + np.mean(target_positions[conn_to_target_indices[s]], axis=0) + for s in unique_active + ], + dtype=np.float32, + ) + + # --- Initial Procrustes alignment --- + if n_unique == 1: + base_placed, base_perm, rotation_type, rotation_info = _initial_align_single( + replacement, + unique_active, + unique_conn_positions, + unique_targets, + neighbor_centroid, + rng, + ) + elif n_unique >= 2: + base_placed, base_perm, rotation_type, rotation_info = _initial_align_multi( + replacement, + unique_active, + unique_conn_positions, + unique_targets, + rng, + ) + else: + R = _random_rotation_matrix(rng) + center = np.mean(neighbor_coords, axis=0) + placed = (R @ replacement.coordinates.T).T + center + return placed.astype(np.float32), [] + + # --- Optimize residual DOF for overlap avoidance + bond validity --- + placed = _optimize_residual_rotation( + base_placed, + rotation_type, + rotation_info, + replacement, + active_conn_sites, + neighbor_coords, + target_bond_lengths, + nearby_coords, + tolerance, + min_separation, + n_rotation_samples, + rng, + ) + + # --- Build assignment --- + assignment = _build_assignment_from_perm( + active_conn_sites, unique_active, conn_to_target_indices, base_perm, degree + ) + + return placed.astype(np.float32), assignment + + +def _initial_align_single( + replacement, + unique_active, + unique_conn_positions, + unique_targets, + neighbor_centroid, + rng, +): + """Initial alignment for 1 unique connection site. + + Returns base placement and rotation info for residual DOF. + + Returns + ------- + base_placed : np.ndarray, shape (n_sites, 3) + perm : list + rotation_type : str + 'axis' for rotation about an axis, 'full' for full sphere + rotation_info : dict + Contains 'pivot' and 'axis' for residual rotation. + """ + c = unique_conn_positions[0] + c_norm = np.linalg.norm(c) + + if c_norm < 1e-10: + R = _random_rotation_matrix(rng) + t = unique_targets[0] + base_placed = (R @ replacement.coordinates.T).T + t + return base_placed, [0], "full", {"pivot": unique_targets[0]} + + c_hat = c / c_norm + desired_dir = unique_targets[0] - neighbor_centroid + desired_norm = np.linalg.norm(desired_dir) + if desired_norm < 1e-10: + desired_dir = rng.standard_normal(3).astype(np.float32) + desired_norm = np.linalg.norm(desired_dir) + desired_hat = desired_dir / desired_norm + + R = _rotation_between_vectors(c_hat, desired_hat) + t = unique_targets[0] - R @ c + base_placed = (R @ replacement.coordinates.T).T + t + + # Residual DOF: rotation about axis from center to connection site + conn_placed = base_placed[unique_active[0]] + center = t # crosslinker center position + axis = conn_placed - center + axis_norm = np.linalg.norm(axis) + if axis_norm > 1e-10: + axis_hat = axis / axis_norm + else: + axis_hat = desired_hat + + return base_placed, [0], "axis", {"pivot": conn_placed, "axis": axis_hat} + + +def _initial_align_multi( + replacement, + unique_active, + unique_conn_positions, + unique_targets, + rng, +): + """Initial alignment for 2+ unique connection sites via Procrustes. + + Returns + ------- + base_placed : np.ndarray + perm : list + rotation_type : str + 'axis' for 2 sites, 'none' for 3+ + rotation_info : dict + """ + n_unique = len(unique_active) + + best_R = np.eye(3, dtype=np.float32) + best_t = np.zeros(3, dtype=np.float32) + best_cost = np.inf + best_perm = list(range(n_unique)) + + for perm in permutations(range(n_unique)): + perm_targets = unique_targets[list(perm)] + + src_centroid = np.mean(unique_conn_positions, axis=0) + tgt_centroid = np.mean(perm_targets, axis=0) + + src_centered = unique_conn_positions - src_centroid + tgt_centered = perm_targets - tgt_centroid + + R = _kabsch_rotation(src_centered, tgt_centered) + t = tgt_centroid - R @ src_centroid + + placed_conn = (R @ unique_conn_positions.T).T + t + cost = np.sum((placed_conn - perm_targets) ** 2) + + if cost < best_cost: + best_cost = cost + best_R = R + best_t = t + best_perm = list(perm) + + base_placed = (best_R @ replacement.coordinates.T).T + best_t + + if n_unique == 2: + # Residual DOF: rotation about axis connecting the two connection sites + conn0 = base_placed[unique_active[0]] + conn1 = base_placed[unique_active[1]] + axis = conn1 - conn0 + axis_norm = np.linalg.norm(axis) + if axis_norm > 1e-10: + axis_hat = axis / axis_norm + else: + axis_hat = np.array([0, 0, 1], dtype=np.float32) + midpoint = (conn0 + conn1) / 2.0 + return base_placed, best_perm, "axis", {"pivot": midpoint, "axis": axis_hat} + else: + # n_unique >= 3: fully determined, no residual DOF + return base_placed, best_perm, "none", {} + + +def _optimize_residual_rotation( + base_placed, + rotation_type, + rotation_info, + replacement, + active_conn_sites, + neighbor_coords, + target_bond_lengths, + nearby_coords, + tolerance, + min_separation, + n_rotation_samples, + rng, +): + """Optimize residual rotational DOF to avoid overlaps while maintaining bond lengths. + + Tries multiple rotations (and offsets if needed), scoring each by: + 1. Bond length validity (hard constraint) + 2. Minimum distance to nearby sites (soft maximize) + + If no rotation passes min_separation, tries pushing the crosslinker along + the perpendicular/normal direction at fractional offsets, then re-sweeps. + + Parameters + ---------- + base_placed : np.ndarray, shape (n_sites, 3) + Initial placement from Procrustes. + rotation_type : str + 'axis', 'full', or 'none' + rotation_info : dict + Contains 'pivot' and 'axis' as needed. + replacement : CrosslinkerGeometry + active_conn_sites : list of int + neighbor_coords : np.ndarray, shape (degree, 3) + target_bond_lengths : np.ndarray, shape (degree,) + nearby_coords : np.ndarray, shape (M, 3) + tolerance : float + min_separation : float or None + n_rotation_samples : int + rng : numpy.random.Generator + + Returns + ------- + best_placed : np.ndarray, shape (n_sites, 3) + """ + if rotation_type == "none" or n_rotation_samples <= 1: + # No residual DOF or no sampling requested + return base_placed + + has_nearby = nearby_coords is not None and len(nearby_coords) > 0 + if not has_nearby and min_separation is None: + return base_placed + + # Generate candidate rotations + if rotation_type == "axis": + pivot = rotation_info["pivot"] + axis_hat = rotation_info["axis"] + angles = np.linspace(0, 2 * np.pi, n_rotation_samples, endpoint=False) + + def generate_candidate(angle): + R_rot = _rotation_about_axis(axis_hat, angle) + return (R_rot @ (base_placed - pivot).T).T + pivot + + elif rotation_type == "full": + pivot = rotation_info["pivot"] + + # Sample random rotations about the pivot + def generate_candidate(idx): + R_rot = _random_rotation_matrix(rng) + return (R_rot @ (base_placed - pivot).T).T + pivot + + angles = list(range(n_rotation_samples)) + else: + return base_placed + + # --- Score each candidate --- + best_placed = base_placed.copy() + best_min_dist = -np.inf + best_valid = False + + bond_tol = tolerance # fractional + + for angle in angles: + candidate = generate_candidate(angle) + + # Check bond validity + bonds_ok = _check_bonds_valid( + candidate, active_conn_sites, neighbor_coords, target_bond_lengths, bond_tol + ) + if not bonds_ok: + continue + + # Score: minimum distance to nearby sites + if has_nearby: + min_dist = _min_distance_to_neighbors(candidate, nearby_coords) + else: + min_dist = np.inf + + # Track best valid placement + is_valid = min_separation is None or min_dist >= min_separation + + if is_valid and (not best_valid or min_dist > best_min_dist): + best_placed = candidate.copy() + best_min_dist = min_dist + best_valid = True + elif not best_valid and min_dist > best_min_dist: + # No valid found yet, track best-effort + best_placed = candidate.copy() + best_min_dist = min_dist + + # --- If no valid rotation found, try offset along perpendicular --- + if not best_valid and min_separation is not None and rotation_type == "axis": + best_placed, best_valid = _try_offset_placements( + base_placed, + rotation_info, + replacement, + active_conn_sites, + neighbor_coords, + target_bond_lengths, + nearby_coords, + bond_tol, + min_separation, + n_rotation_samples, + rng, + ) + + return best_placed + + +def _try_offset_placements( + base_placed, + rotation_info, + replacement, + active_conn_sites, + neighbor_coords, + target_bond_lengths, + nearby_coords, + bond_tol, + min_separation, + n_rotation_samples, + rng, +): + """Try offsetting the crosslinker perpendicular to the rotation axis. + + When the crosslinker is in the plane of the backbone and all rotations + produce overlaps, pushing it "above" or "below" the plane can resolve + the clash while maintaining bond lengths (connection sites rotate to + maintain the correct distance). + + Tries offsets in both directions at increasing magnitudes. + + Parameters + ---------- + ... (same as parent) + + Returns + ------- + best_placed : np.ndarray, shape (n_sites, 3) + found_valid : bool + """ + axis_hat = rotation_info["axis"] + pivot = rotation_info["pivot"] + + # Find a perpendicular direction to offset along + # Use the cross product of axis with a non-parallel vector + perp = np.array([1, 0, 0], dtype=np.float32) + if abs(np.dot(axis_hat, perp)) > 0.9: + perp = np.array([0, 1, 0], dtype=np.float32) + offset_dir = np.cross(axis_hat, perp) + offset_dir /= np.linalg.norm(offset_dir) + + # Try offsets at increasing distances + # Max offset: limited by how much bond length can stretch within tolerance + max_offset = np.mean(target_bond_lengths) * bond_tol * 3 + offsets = np.linspace(0.02 * max_offset, max_offset, 5) + + best_placed = base_placed.copy() + best_min_dist = -np.inf + found_valid = False + + for sign in [1.0, -1.0]: + for offset_mag in offsets: + offset = offset_dir * offset_mag * sign + shifted = base_placed + offset + + # Re-sweep rotation at this offset + angles = np.linspace(0, 2 * np.pi, n_rotation_samples, endpoint=False) + shifted_pivot = pivot + offset + + for angle in angles: + R_rot = _rotation_about_axis(axis_hat, angle) + candidate = (R_rot @ (shifted - shifted_pivot).T).T + shifted_pivot + + bonds_ok = _check_bonds_valid( + candidate, + active_conn_sites, + neighbor_coords, + target_bond_lengths, + bond_tol, + ) + if not bonds_ok: + continue + + if nearby_coords is not None and len(nearby_coords) > 0: + min_dist = _min_distance_to_neighbors(candidate, nearby_coords) + else: + min_dist = np.inf + + is_valid = min_dist >= min_separation + + if is_valid and (not found_valid or min_dist > best_min_dist): + best_placed = candidate.copy() + best_min_dist = min_dist + found_valid = True + elif not found_valid and min_dist > best_min_dist: + best_placed = candidate.copy() + best_min_dist = min_dist + + if found_valid: + break + if found_valid: + break + + return best_placed, found_valid + + +def _check_bonds_valid( + placed, active_conn_sites, neighbor_coords, target_bond_lengths, tolerance +): + """Check if all connection-to-neighbor bonds are within tolerance. + + Parameters + ---------- + placed : np.ndarray, shape (n_sites, 3) + active_conn_sites : list of int + neighbor_coords : np.ndarray, shape (degree, 3) + target_bond_lengths : np.ndarray, shape (degree,) + tolerance : float (fractional) + + Returns + ------- + valid : bool + """ + for i, site_idx in enumerate(active_conn_sites): + delta = placed[site_idx] - neighbor_coords[i] + actual = np.linalg.norm(delta) + target = target_bond_lengths[i] + if abs(actual - target) > target * tolerance: + return False + return True + + +def _min_distance_to_neighbors(positions, nearby_coords): + """Minimum distance from any position to any nearby coord.""" + if len(nearby_coords) == 0: + return np.inf + diffs = positions[:, None, :] - nearby_coords[None, :, :] + dists = np.linalg.norm(diffs, axis=-1) + return float(np.min(dists)) + + +def _build_assignment_from_perm( + active_conn_sites, unique_active, conn_to_target_indices, best_perm, degree +): + """Build full assignment list from permutation over unique sites. + + Parameters + ---------- + active_conn_sites : list of int + Connection sites being used (length = degree). + unique_active : list of int + Unique site indices. + conn_to_target_indices : dict + site_idx -> [target indices using that site]. + best_perm : list of int + Permutation mapping unique site position -> unique target position. + degree : int + Number of connections being made. + + Returns + ------- + assignment : list of int + assignment[i] = neighbor_index for active_conn_sites[i]. + """ + if degree == len(unique_active) and all( + len(v) == 1 for v in conn_to_target_indices.values() + ): + # Simple: 1:1, no repeated sites + # best_perm[i] maps unique site i to unique target i + # But we need to map conn_idx -> neighbor_idx + # active_conn_sites[conn_idx] -> unique_active.index(that) -> perm -> target + assignment = [] + for conn_idx, site_idx in enumerate(active_conn_sites): + unique_pos = unique_active.index(site_idx) + assignment.append(best_perm[unique_pos]) + return assignment + + # General case with possible repeated sites + assignment = [0] * degree + site_counter = {} + + for conn_idx, site_idx in enumerate(active_conn_sites): + unique_pos = unique_active.index(site_idx) + target_group_pos = best_perm[unique_pos] + target_site = unique_active[target_group_pos] + target_indices = conn_to_target_indices[target_site] + + count = site_counter.get(site_idx, 0) + assignment[conn_idx] = target_indices[count % len(target_indices)] + site_counter[site_idx] = count + 1 + + return assignment + + +def _verify_and_fix_bonds( + placed_positions, + active_conn_sites, + assignment, + neighbor_coords, + target_bond_lengths, + tolerance, +): + """Verify connection-to-neighbor distances and nudge if outside tolerance. + + Moves each connection site along the line to its neighbor to achieve + exactly target_bond_length if it deviates too much. Then shifts the + entire rigid body to maintain internal geometry as much as possible. + + Parameters + ---------- + placed_positions : np.ndarray, shape (n_sites, 3) + active_conn_sites : list of int + assignment : list of int + neighbor_coords : np.ndarray, shape (degree, 3) + target_bond_lengths : np.ndarray, shape (degree,) + tolerance : float + + Returns + ------- + placed_positions : np.ndarray, shape (n_sites, 3) + Corrected positions. + """ + placed = placed_positions.copy() + degree = len(active_conn_sites) + + needs_correction = False + corrections = [] + + for conn_idx, neighbor_idx in enumerate(assignment): + site_local = active_conn_sites[conn_idx] + conn_pos = placed[site_local] + nb_pos = neighbor_coords[neighbor_idx] + target_bl = target_bond_lengths[conn_idx] + + delta = conn_pos - nb_pos + actual_bl = np.linalg.norm(delta) + + if actual_bl < 1e-10: + # Degenerate: connection site is on top of neighbor + needs_correction = True + corrections.append((site_local, neighbor_idx, conn_idx)) + continue + + error = abs(actual_bl - target_bl) + if error > target_bl * tolerance: + needs_correction = True + corrections.append((site_local, neighbor_idx, conn_idx)) + + if not needs_correction: + return placed + + # Strategy: compute ideal connection positions and shift entire rigid body + # to best match them (least-squares center shift) + ideal_conn_positions = np.zeros((degree, 3), dtype=np.float32) + for conn_idx, neighbor_idx in enumerate(assignment): + site_local = active_conn_sites[conn_idx] + nb_pos = neighbor_coords[neighbor_idx] + target_bl = target_bond_lengths[conn_idx] + + conn_pos = placed[site_local] + delta = conn_pos - nb_pos + d = np.linalg.norm(delta) + if d > 1e-10: + direction = delta / d + else: + # Pick direction from crosslinker center to this site + center = np.mean(placed, axis=0) + direction = placed[site_local] - center + d2 = np.linalg.norm(direction) + if d2 > 1e-10: + direction = direction / d2 + else: + direction = np.array([1, 0, 0], dtype=np.float32) + + ideal_conn_positions[conn_idx] = nb_pos + direction * target_bl + + # Compute offset: difference between where connection sites are and should be + current_conn_positions = np.array( + [placed[active_conn_sites[i]] for i in range(degree)], dtype=np.float32 + ) + offsets = ideal_conn_positions - current_conn_positions + mean_offset = np.mean(offsets, axis=0) + + # Apply rigid shift to all sites + placed += mean_offset + + # Final per-site correction for any remaining error + for conn_idx, neighbor_idx in enumerate(assignment): + site_local = active_conn_sites[conn_idx] + nb_pos = neighbor_coords[neighbor_idx] + target_bl = target_bond_lengths[conn_idx] + + delta = placed[site_local] - nb_pos + actual_bl = np.linalg.norm(delta) + + if actual_bl < 1e-10: + center = np.mean(placed, axis=0) + direction = placed[site_local] - center + d2 = np.linalg.norm(direction) + direction = direction / d2 if d2 > 1e-10 else np.array([1, 0, 0]) + placed[site_local] = nb_pos + direction * target_bl + elif abs(actual_bl - target_bl) > target_bl * tolerance: + direction = delta / actual_bl + placed[site_local] = nb_pos + direction * target_bl + + return placed + + +def _get_pbc_info(volume_constraint): + """Extract PBC flags and box lengths from volume constraint.""" + if isinstance(volume_constraint, CuboidConstraint): + pbc = np.array(volume_constraint.pbc, dtype=bool) + box_lengths = volume_constraint.box_lengths.astype(np.float32) + elif isinstance(volume_constraint, CylinderConstraint): + pbc = np.array([False, False, volume_constraint.periodic_height], dtype=bool) + box_lengths = np.array( + [ + volume_constraint.radius * 2, + volume_constraint.radius * 2, + volume_constraint.height, + ], + dtype=np.float32, + ) + else: + pbc = np.array([False, False, False], dtype=bool) + box_lengths = np.array([np.inf, np.inf, np.inf], dtype=np.float32) + return pbc, box_lengths + + +def _apply_mic(delta, pbc, box_lengths): + """Apply minimum image convention to a displacement vector.""" + delta = np.asarray(delta, dtype=np.float32) + for dim in range(3): + if pbc[dim]: + delta[dim] -= np.round(delta[dim] / box_lengths[dim]) * box_lengths[dim] + return delta + + +def _wrap_positions(positions, pbc, box_lengths): + """Wrap positions into periodic box [-L/2, L/2].""" + positions = positions.copy() + for dim in range(3): + if pbc[dim]: + positions[:, dim] -= ( + np.round(positions[:, dim] / box_lengths[dim]) * box_lengths[dim] + ) + return positions + + +def _get_nearby_nonbonded( + path, site, neighbors, site_pos, radius, pbc, box_lengths, excluded_sites +): + """Get coordinates of nearby non-bonded sites for overlap computation. + + Returns positions relative to site_pos. + """ + neighbors_set = set(neighbors) | {site} | excluded_sites + + nearby = [] + for node in path.bond_graph.nodes(): + if node in neighbors_set: + continue + delta = path.coordinates[node] - site_pos + delta = _apply_mic(delta, pbc, box_lengths) + dist = np.linalg.norm(delta) + if dist <= radius: + nearby.append(delta) + + if len(nearby) == 0: + return np.empty((0, 3), dtype=np.float32) + return np.array(nearby, dtype=np.float32) + + +def crosslink( + path, + crosslinker=None, + bead_name="_R", + backbone_name="_A", + crosslink_bond_length=0.2, + tolerance=0.1, + excluded_bond_depth=2, + n_connection_sites=2, + volume_constraint=None, + initial_point=None, + seed=42, + chunk_size=512, + run_on_gpu=False, + n_rotation_samples=36, + overlap_radius=None, + min_separation=None, +): + """ + Create a crosslink bonded to backbone beads, modifying path in place. + + Selects backbone beads whose pairwise spacing matches the crosslinker + geometry, places a temporary placeholder, then calls ``replace_sites`` + which handles rigid-body alignment, bond length guarantees, and overlap + avoidance in a single pass. + + Parameters + ---------- + path : Path + The Path object to modify in place. + crosslinker : CrosslinkerGeometry, optional + If None, uses single-site from ``bead_name`` and ``n_connection_sites``. + bead_name : str, default "_R" + backbone_name : str, default "_A" + crosslink_bond_length : float, default 0.2 + Desired bond length from connection sites to backbone beads (nm). + tolerance : float, default 0.1 + Fractional tolerance on bond lengths (±10%). + excluded_bond_depth : int, default 2 + n_connection_sites : int, default 2 + volume_constraint : optional + initial_point : int or array-like, optional + seed : int, default 42 + chunk_size : int, default 512 + run_on_gpu : bool, default False + n_rotation_samples : int, default 36 + overlap_radius : float, optional + min_separation : float, optional + Minimum distance from placed sites to existing sites. + Defaults to ``crosslink_bond_length * 0.5``. + + Returns + ------- + Path + Same path, modified in place. + """ + if crosslinker is None: + crosslinker = CrosslinkerGeometry.single_site( + bead_name=bead_name, n_connections=n_connection_sites + ) + + if min_separation is None: + min_separation = crosslink_bond_length * 0.5 + + n_backbone_connections = crosslinker.n_connections + rng = np.random.default_rng(seed + len(path.coordinates)) + + # --- Compute search geometry --- + ideal_bb_positions, ideal_pairwise_distances, geom_search_radius = ( + _compute_ideal_backbone_distances(crosslinker, crosslink_bond_length) + ) + + if n_backbone_connections > 1: + max_pair_dist = np.max(ideal_pairwise_distances) + else: + max_pair_dist = 0.0 + + search_radius = max_pair_dist * (1 + tolerance) + crosslink_bond_length * tolerance + + # --- Find backbone targets --- + backbone_nodes = [ + node for node in path.bond_graph.nodes() if path.beads[node] == backbone_name + ] + backbone_subgraph = path.bond_graph.subgraph(backbone_nodes) + + if len(backbone_nodes) == 0: + raise ValueError(f"No backbone beads with name '{backbone_name}' found in path") + if len(backbone_nodes) < n_backbone_connections: + raise ValueError( + f"Not enough backbone beads ({len(backbone_nodes)}) for " + f"{n_backbone_connections} connection sites" + ) + + pbc, box_lengths = _get_pbc_info(volume_constraint) + + candidate_nodes = [ + node for node in backbone_nodes if path.bond_graph.degree[node] <= 2 + ] + candidate_coords = np.array(path.coordinates[candidate_nodes], dtype=np.float32) + + def get_reference_points(path, initial_point): + if initial_point is not None: + if isinstance(initial_point, (int, np.integer)): + if initial_point not in path.bond_graph.nodes: + raise ValueError(f"Node {initial_point} not found in bond_graph") + return np.array([initial_point]), np.array( + [path.coordinates[initial_point]] + ) + else: + initial_point32 = np.asarray(initial_point, dtype=np.float32) + sq_distances = calculate_sq_distances( + initial_point32, candidate_coords, pbc=pbc, box_lengths=box_lengths + ) + order = np.argsort(sq_distances) + return ( + np.array(candidate_nodes)[order], + path.coordinates[np.array(candidate_nodes)[order]], + ) + else: + shuffled = rng.choice( + candidate_nodes, size=len(candidate_nodes), replace=False + ) + return shuffled, path.coordinates[shuffled] + + ref_nodes, ref_coords = get_reference_points(path, initial_point) + + found_ref = False + selected_nodes = [] + + for ref_node, ref_coord in zip(ref_nodes, ref_coords): + selected_nodes = [ref_node] + + sq_distances = calculate_sq_distances( + ref_coord, candidate_coords, pbc=pbc, box_lengths=box_lengths + ) + distances = np.sqrt(sq_distances) + + within_radius_mask = distances <= search_radius + possible_pairs = np.where(within_radius_mask)[0] + + if len(possible_pairs) < n_backbone_connections - 1: + continue + + closest_paired_nodes = possible_pairs[np.argsort(distances[possible_pairs])] + excluded_nodes = set( + nx.single_source_shortest_path_length( + backbone_subgraph, ref_node, cutoff=excluded_bond_depth + ).keys() + ) + + for idx in closest_paired_nodes: + node = candidate_nodes[idx] + if node in excluded_nodes: + continue + + is_excluded = False + for selected in selected_nodes[1:]: + if selected in backbone_subgraph: + sel_excluded = set( + nx.single_source_shortest_path_length( + backbone_subgraph, selected, cutoff=excluded_bond_depth + ).keys() + ) + if node in sel_excluded: + is_excluded = True + break + if is_excluded: + continue + + selected_nodes.append(int(node)) + if len(selected_nodes) >= n_backbone_connections: + # Validate pairwise geometry + sel_coords = np.array( + [path.coordinates[n] for n in selected_nodes], dtype=np.float32 + ) + unwrapped_check = _unwrap_coords(sel_coords, pbc, box_lengths) + valid, _ = _validate_backbone_group( + unwrapped_check, + ideal_pairwise_distances, + crosslink_bond_length, + tolerance, + pbc, + box_lengths, + ) + if valid: + found_ref = True + break + else: + selected_nodes.pop() + + if found_ref: + break + + if not found_ref: + n_clinks = sum(1 for b in path.beads if b in set(crosslinker.beads)) + raise PathConvergenceError( + f"Could not find {n_backbone_connections} non-neighboring backbone beads " + f"matching crosslinker geometry with bond_length={crosslink_bond_length} " + f"(tolerance=±{tolerance * 100:.0f}%)." + f"\nCurrent crosslinks: {n_clinks}. " + "Ways to increase crosslinking:\n" + " - Increase crosslink_bond_length\n" + " - Increase tolerance\n" + " - Decrease excluded_bond_depth\n" + " - Pack at higher density\n" + " - Relax structure" + ) + + # --- Place placeholder at centroid --- + selected_coords = np.array(path.coordinates[selected_nodes], dtype=np.float32) + unwrapped = _unwrap_coords(selected_coords, pbc, box_lengths) + crosslink_center = np.mean(unwrapped, axis=0) + + for dim in range(3): + if pbc[dim]: + crosslink_center[dim] -= ( + np.round(crosslink_center[dim] / box_lengths[dim]) * box_lengths[dim] + ) + + _placeholder_name = "__placeholder__" + path.append_coordinates(crosslink_center, _placeholder_name) + placeholder_idx = len(path.coordinates) - 1 + + for backbone_node in selected_nodes: + path.bond_graph.add_edge( + int(placeholder_idx), + int(backbone_node), + bond_type=(_placeholder_name, backbone_name), + ) + + # --- Single-site shortcut --- + if crosslinker.n_sites == 1: + path.beads[placeholder_idx] = crosslinker.beads[0] + for backbone_node in selected_nodes: + path.bond_graph.edges[placeholder_idx, backbone_node]["bond_type"] = ( + str(crosslinker.beads[0]), + backbone_name, + ) + if "name" in path.bond_graph.nodes[placeholder_idx]: + path.bond_graph.nodes[placeholder_idx]["name"] = str(crosslinker.beads[0]) + + # Fix single-site bond length if needed + for backbone_node in selected_nodes: + delta = path.coordinates[placeholder_idx] - path.coordinates[backbone_node] + delta = _apply_mic(delta, pbc, box_lengths) + actual_bl = np.linalg.norm(delta) + if ( + abs(actual_bl - crosslink_bond_length) + > crosslink_bond_length * tolerance + ): + if actual_bl > 1e-10: + direction = delta / actual_bl + else: + direction = rng.standard_normal(3).astype(np.float32) + direction /= np.linalg.norm(direction) + path.coordinates[placeholder_idx] = path.coordinates[ + backbone_node + ] + _apply_mic(direction * crosslink_bond_length, pbc, box_lengths) + return path + + # --- Multi-site: replace_sites handles everything --- + replace_sites( + path, + replacement=crosslinker, + sites=[placeholder_idx], + bond_length=crosslink_bond_length, + tolerance=tolerance, + min_separation=min_separation, + volume_constraint=volume_constraint, + n_rotation_samples=n_rotation_samples, + overlap_radius=overlap_radius, + seed=seed, + ) + + return path + + +def _compute_ideal_backbone_distances(crosslinker, bond_length): + """Compute ideal pairwise distances between backbone targets given geometry and bond length. + + For a crosslinker with connection sites at positions c_i (relative to center), + the ideal backbone position for site i is: + b_i = c_i + bond_length * (c_i / |c_i|) + i.e., the backbone sits along the same radial direction, at distance bond_length + beyond the connection site. + + Parameters + ---------- + crosslinker : CrosslinkerGeometry + The crosslinker geometry. + bond_length : float + Desired bond length from connection site to backbone bead. + + Returns + ------- + ideal_bb_positions : np.ndarray, shape (n_connections, 3) + Ideal backbone positions relative to crosslinker center. + pairwise_distances : np.ndarray, shape (n_connections, n_connections) + Matrix of ideal pairwise distances between backbone targets. + search_radius : float + Maximum distance from center to any ideal backbone position. + """ + conn_sites = crosslinker.connection_sites + conn_positions = crosslinker.coordinates[conn_sites] + + # Compute ideal backbone positions + ideal_bb_positions = np.zeros_like(conn_positions) + for i in range(len(conn_sites)): + c_i = conn_positions[i] + c_norm = np.linalg.norm(c_i) + if c_norm > 1e-10: + # Backbone sits beyond connection site along radial direction + c_hat = c_i / c_norm + ideal_bb_positions[i] = c_i + bond_length * c_hat + else: + # Connection site at origin (single-site case): backbone at bond_length in any direction + ideal_bb_positions[i] = np.array([bond_length, 0, 0], dtype=np.float32) + + # Pairwise distances + n = len(conn_sites) + pairwise_distances = np.zeros((n, n), dtype=np.float32) + for i in range(n): + for j in range(i + 1, n): + d = np.linalg.norm(ideal_bb_positions[i] - ideal_bb_positions[j]) + pairwise_distances[i, j] = d + pairwise_distances[j, i] = d + + # Search radius: max distance from center to any backbone + search_radius = np.max(np.linalg.norm(ideal_bb_positions, axis=1)) + + return ideal_bb_positions, pairwise_distances, search_radius + + +def _unwrap_coords(coords, pbc, box_lengths): + """Unwrap coordinates relative to the first point.""" + unwrapped = np.empty_like(coords) + unwrapped[0] = coords[0] + ref = coords[0] + for k in range(1, len(coords)): + delta = coords[k] - ref + for dim in range(3): + if pbc[dim]: + delta[dim] -= np.round(delta[dim] / box_lengths[dim]) * box_lengths[dim] + unwrapped[k] = ref + delta + return unwrapped + + +def _compute_optimal_center( + crosslinker, backbone_coords, bond_length, pbc, box_lengths +): + """Compute the optimal crosslinker center so connection sites are at bond_length from backbones. + + The center is placed such that each connection site (after rotation) sits + at exactly bond_length from its target backbone bead. + + For the ideal placement: + backbone_i = center + c_i_hat * (|c_i| + bond_length) + + So: center = backbone_i - c_i_hat * (|c_i| + bond_length) + + We average over all connection sites for robustness. + + Parameters + ---------- + crosslinker : CrosslinkerGeometry + The crosslinker. + backbone_coords : np.ndarray, shape (n_connections, 3) + Unwrapped backbone coordinates. + bond_length : float + Desired bond length. + pbc, box_lengths : PBC info. + + Returns + ------- + center : np.ndarray, shape (3,) + """ + conn_sites = crosslinker.connection_sites + conn_positions = crosslinker.coordinates[conn_sites] + n = len(conn_sites) + + if n == 0: + return np.mean(backbone_coords, axis=0) + + # For single-site crosslinker (all connection sites at same position): + if np.allclose(conn_positions, conn_positions[0], atol=1e-8): + # Center is at bond_length distance from centroid toward... just use centroid + return np.mean(backbone_coords, axis=0) + + # Compute directions from crosslinker center to each connection site + conn_norms = np.linalg.norm(conn_positions, axis=1, keepdims=True) + conn_hat = conn_positions / np.maximum(conn_norms, 1e-10) + + # Ideal: backbone_i is at direction c_i_hat, distance |c_i| + bond_length from center + # So center = backbone_i - c_i_hat * (|c_i| + bond_length) + # But we haven't rotated the crosslinker yet, so we need to figure out the rotation. + + # Target vectors: from centroid of backbones to each backbone + bb_centroid = np.mean(backbone_coords, axis=0) + target_vectors = backbone_coords - bb_centroid + target_norms = np.linalg.norm(target_vectors, axis=1, keepdims=True) + target_hat = target_vectors / np.maximum(target_norms, 1e-10) + + # The ideal distance from center to each backbone: + # |c_i| + bond_length + ideal_radii = conn_norms.flatten() + bond_length + + # Center = backbone_i - ideal_radius_i * target_hat_i (if we orient along target directions) + # Average: + centers = backbone_coords - ideal_radii[:, None] * target_hat + center = np.mean(centers, axis=0) + + return center.astype(np.float32) + + +def _validate_backbone_group( + selected_coords, + ideal_pairwise_distances, + bond_length, + tolerance, + pbc, + box_lengths, +): + """Check if backbone beads match expected geometry within tolerance. + + Returns + ------- + valid : bool + best_perm : list of int or None + """ + n = len(selected_coords) + if n <= 1: + return True, [0] if n == 1 else [] + + abs_tol = bond_length * tolerance + + # Compute actual pairwise distances (PBC-aware) + actual_distances = np.zeros((n, n), dtype=np.float32) + for i in range(n): + for j in range(i + 1, n): + delta = selected_coords[j] - selected_coords[i] + for dim in range(3): + if pbc[dim]: + delta[dim] -= ( + np.round(delta[dim] / box_lengths[dim]) * box_lengths[dim] + ) + d = np.linalg.norm(delta) + actual_distances[i, j] = d + actual_distances[j, i] = d + + best_perm = None + best_cost = np.inf + + for perm in permutations(range(n)): + valid = True + cost = 0.0 + for i in range(n): + for j in range(i + 1, n): + ideal_d = ideal_pairwise_distances[i, j] + actual_d = actual_distances[perm[i], perm[j]] + error = abs(actual_d - ideal_d) + if error > abs_tol: + valid = False + break + cost += error + if not valid: + break + + if valid and cost < best_cost: + best_cost = cost + best_perm = list(perm) + + return best_perm is not None, best_perm diff --git a/mbuild/simulation.py b/mbuild/simulation.py index 58402c6dd..4fe957127 100644 --- a/mbuild/simulation.py +++ b/mbuild/simulation.py @@ -13,7 +13,7 @@ from ele.exceptions import ElementError from gmso.parameterization import apply -from mbuild import Compound +from mbuild import Box, Compound from mbuild.exceptions import MBuildError from mbuild.utils.io import import_ @@ -247,7 +247,7 @@ class ForcesHandler: The value passed to `dpd` is used as the the repulsion force constant. scale_bonds : float, default 1.0 scale_angles : float, default 1.0 - scale_lj : float, default 0.0 + scale_lj : float, default 1.0 scale_periodic : float, default 0.0 scale_opls : float, default 0.0 scale_improper : float, default 0.0 @@ -1287,7 +1287,7 @@ def energy_minimize_path( If a float is passed, that is set as theta in degress for a harmonic angle. If 'constrain' is passed, then angles are constrained. If None is passed, no angle force is used. - box : mb.Box, optional + box : mb.Box or list or numpy.ndarray, optional Box to use for periodic boundaries. steps : int, optional, default 1,000 Number of simulation steps to run. @@ -1300,6 +1300,9 @@ def energy_minimize_path( positions = path.coordinates n_particles = len(positions) + if isinstance(box, Box): + box = np.array([box.Lx, box.Ly, box.Lz]) + import openmm import openmm.unit as u from openmm import Platform @@ -1350,7 +1353,7 @@ def energy_minimize_path( # Constrain all bonds at current lengths for i, j in path.bond_graph.edges(): delta = positions[j] - positions[i] - if box: + if box is not None: box_arr = np.array(box) delta -= np.round(delta / box_arr) * box_arr bl = np.linalg.norm(delta) @@ -1359,7 +1362,7 @@ def energy_minimize_path( elif isinstance(bonds, dict): # Per-type harmonic bonds harmonic_force = openmm.HarmonicBondForce() - if box: + if box is not None: harmonic_force.setUsesPeriodicBoundaryConditions(True) for i, j in path.bond_graph.edges(): @@ -1385,7 +1388,7 @@ def energy_minimize_path( else: # Fallback: constrain at current distance delta = positions[j] - positions[i] - if box: + if box is not None: box_arr = np.array(box) delta -= np.round(delta / box_arr) * box_arr bl = np.linalg.norm(delta) @@ -1396,7 +1399,7 @@ def energy_minimize_path( elif isinstance(bonds, (int, float)): # Single float: same equilibrium length for all bonds harmonic_force = openmm.HarmonicBondForce() - if box: + if box is not None: harmonic_force.setUsesPeriodicBoundaryConditions(True) for i, j in path.bond_graph.edges(): harmonic_force.addBond( @@ -1416,7 +1419,7 @@ def energy_minimize_path( continue for i_node, k_node in combinations(neighbors, 2): delta = positions[k_node] - positions[i_node] - if box: + if box is not None: box_arr = np.array(box) delta -= np.round(delta / box_arr) * box_arr dist_13 = np.linalg.norm(delta) @@ -1425,7 +1428,7 @@ def energy_minimize_path( elif isinstance(angles, dict): # Per-type harmonic angles angle_force = openmm.HarmonicAngleForce() - if box: + if box is not None: angle_force.setUsesPeriodicBoundaryConditions(True) for center_node in path.bond_graph.nodes(): @@ -1472,7 +1475,7 @@ def energy_minimize_path( simulation = Simulation(topology, system, integrator, platform) simulation.context.setPositions(positions * u.nanometer) - if box: + if box is not None: simulation.context.setPeriodicBoxVectors( [box[0], 0, 0] * u.nanometer, [0, box[1], 0] * u.nanometer, diff --git a/mbuild/tests/test_crosslinks.py b/mbuild/tests/test_crosslinks.py new file mode 100644 index 000000000..d350402c3 --- /dev/null +++ b/mbuild/tests/test_crosslinks.py @@ -0,0 +1,478 @@ +import numpy as np +import pytest + +from mbuild.exceptions import PathConvergenceError +from mbuild.path.build import Path +from mbuild.path.crosslink import ( + CrosslinkerGeometry, + crosslink, + replace_sites, +) +from mbuild.tests.base_test import BaseTest + + +class TestCLinks(BaseTest): + @pytest.fixture + def linear_path(self): + coords = np.array([[0, 0, 0], [1, 0, 0], [3, 0, 0], [5, 0, 0]]) + return Path(coords) + + def test_clink_equilateral(self, linear_path): + cl = CrosslinkerGeometry.equilateral_triangle( + bond_length=0.27, + bead_name="_R", + connection_sites=[0, 1], # only sites 0 and 1 bond to backbone + ) + outPath = crosslink(linear_path, crosslinker=cl, crosslink_bond_length=0.5) + assert len(outPath) == 7 + assert len(outPath.bond_graph.edges()) == 5 + assert outPath.bond_graph.has_edge(0, 5) + assert outPath.bond_graph.has_edge(1, 4) # consistent edges form + + def test_replace_site(self): + coords = np.column_stack( + (np.arange(0, 5), np.zeros(5, dtype=int), np.zeros(5, dtype=int)) + ) + path = Path(coords) + path._connect_edges("linear") + assert len(path.bond_graph.edges()) == 4 + + triangle = CrosslinkerGeometry.equilateral_triangle( + bond_length=0.4, connection_sites=[0, 1] + ) + path = path.replace_sites(triangle, 1) + assert len(path) == 7 + + +class TestCrosslinkerGeometry(BaseTest): + """Tests for CrosslinkerGeometry construction and factory methods.""" + + def test_single_site_default(self): + cl = CrosslinkerGeometry.single_site() + assert cl.n_sites == 1 + assert cl.n_connections == 2 + assert cl.connection_sites == [0, 0] + assert len(cl.internal_bonds) == 0 + + def test_single_site_custom_connections(self): + cl = CrosslinkerGeometry.single_site(bead_name="_X", n_connections=3) + assert cl.n_sites == 1 + assert cl.n_connections == 3 + assert cl.beads[0] == "_X" + + def test_equilateral_triangle_geometry(self): + cl = CrosslinkerGeometry.equilateral_triangle(bond_length=0.27) + assert cl.n_sites == 3 + assert cl.n_connections == 3 + assert len(cl.internal_bonds) == 3 + # Verify equilateral: all edge lengths should be ~0.27 + for i, j in cl.bond_graph.edges(): + dist = np.linalg.norm(cl.coordinates[i] - cl.coordinates[j]) + assert abs(dist - 0.27) < 1e-4 + + def test_equilateral_triangle_partial_connections(self): + cl = CrosslinkerGeometry.equilateral_triangle( + bond_length=0.27, connection_sites=[0, 1] + ) + assert cl.n_connections == 2 + assert cl.connection_sites == [0, 1] + + def test_linear_geometry(self): + cl = CrosslinkerGeometry.linear(n_sites=3, bond_length=0.5) + assert cl.n_sites == 3 + assert cl.n_connections == 2 # default: first and last + assert len(cl.internal_bonds) == 2 + # Check bond lengths + for i, j in cl.bond_graph.edges(): + dist = np.linalg.norm(cl.coordinates[i] - cl.coordinates[j]) + assert abs(dist - 0.5) < 1e-4 + + def test_square_geometry(self): + cl = CrosslinkerGeometry.square(bond_length=0.27) + assert cl.n_sites == 4 + assert cl.n_connections == 4 + assert len(cl.internal_bonds) == 4 + for i, j in cl.bond_graph.edges(): + dist = np.linalg.norm(cl.coordinates[i] - cl.coordinates[j]) + assert abs(dist - 0.27) < 1e-4 + + def test_tetrahedral_geometry(self): + cl = CrosslinkerGeometry.tetrahedral(bond_length=0.27) + assert cl.n_sites == 4 + assert cl.n_connections == 4 + assert len(cl.internal_bonds) == 6 # complete graph K4 + for i, j in cl.bond_graph.edges(): + dist = np.linalg.norm(cl.coordinates[i] - cl.coordinates[j]) + assert abs(dist - 0.27) < 1e-4 + + def test_trigonal_bipyramidal_geometry(self): + cl = CrosslinkerGeometry.trigonal_bipyramidal(bond_length=0.27) + assert cl.n_sites == 5 + assert cl.n_connections == 5 + assert len(cl.internal_bonds) == 9 + + def test_pentagon_geometry(self): + cl = CrosslinkerGeometry.pentagon(bond_length=0.27) + assert cl.n_sites == 5 + assert cl.n_connections == 5 + assert len(cl.internal_bonds) == 5 + for i, j in cl.bond_graph.edges(): + dist = np.linalg.norm(cl.coordinates[i] - cl.coordinates[j]) + assert abs(dist - 0.27) < 1e-4 + + def test_from_edges(self): + coords = [[0, 0, 0], [1, 0, 0], [0.5, 0.866, 0]] + edges = [(0, 1), (1, 2)] + cl = CrosslinkerGeometry.from_edges( + coords, edges, bead_name="_R", connection_sites=[0, 2] + ) + assert cl.n_sites == 3 + assert cl.n_connections == 2 + assert len(cl.internal_bonds) == 2 + assert cl.connection_sites == [0, 2] + + def test_from_path(self): + coords = np.array([[0, 0, 0], [0.27, 0, 0]], dtype=np.float32) + p = Path(coords, bead_name="_R") + p._connect_edges("linear") + cl = CrosslinkerGeometry.from_path(p, connection_sites=[0, 1]) + assert cl.n_sites == 2 + assert cl.n_connections == 2 + assert len(cl.internal_bonds) == 1 + + def test_centered_at_origin(self): + cl = CrosslinkerGeometry.equilateral_triangle(bond_length=0.27) + centroid = np.mean(cl.coordinates, axis=0) + assert np.allclose(centroid, [0, 0, 0], atol=1e-6) + + def test_invalid_connection_site_raises(self): + coords = np.array([[0, 0, 0], [1, 0, 0]], dtype=np.float32) + with pytest.raises(ValueError, match="out of range"): + CrosslinkerGeometry( + coordinates=coords, bead_name="_R", connection_sites=[0, 5] + ) + + def test_unique_connection_sites(self): + cl = CrosslinkerGeometry.single_site(n_connections=3) + assert cl.unique_connection_sites == [0] + cl2 = CrosslinkerGeometry.equilateral_triangle(connection_sites=[0, 1, 2]) + assert cl2.unique_connection_sites == [0, 1, 2] + + def test_copy(self): + cl = CrosslinkerGeometry.equilateral_triangle(bond_length=0.27) + cl2 = cl.copy() + assert np.allclose(cl.coordinates, cl2.coordinates) + assert cl.connection_sites == cl2.connection_sites + # Mutate copy, original unchanged + cl2.coordinates[0] = [99, 99, 99] + assert not np.allclose(cl.coordinates, cl2.coordinates) + + +class TestCrosslink(BaseTest): + """Tests for the crosslink function.""" + + @pytest.fixture + def linear_path(self): + coords = np.array( + [[0, 0, 0], [1, 0, 0], [3, 0, 0], [5, 0, 0]], dtype=np.float32 + ) + p = Path(coords, bead_name="_A") + p._connect_edges("linear") + return p + + @pytest.fixture + def long_linear_path(self): + n = 20 + coords = np.column_stack( + (np.linspace(0, 5, n), np.zeros(n), np.zeros(n)) + ).astype(np.float32) + p = Path(coords, bead_name="_A") + p._connect_edges("linear") + return p + + def test_crosslink_single_site_default(self, linear_path): + out = crosslink(linear_path, crosslink_bond_length=3.0, n_connection_sites=2) + # Original 4 + 1 crosslink = 5 + assert len(out.coordinates) == 5 + # 3 backbone edges + 2 crosslink edges = 5 + assert len(out.bond_graph.edges()) == 5 + + def test_crosslink_equilateral_2conn(self, linear_path): + cl = CrosslinkerGeometry.equilateral_triangle( + bond_length=0.27, connection_sites=[0, 1] + ) + out = crosslink(linear_path, crosslinker=cl, crosslink_bond_length=3.0) + # 4 backbone + 3 triangle sites = 7 + assert len(out.coordinates) == 7 + # 3 backbone + 3 internal + 2 connections = 8 + assert len(out.bond_graph.edges()) == 8 + + def test_crosslink_equilateral_3conn(self, long_linear_path): + cl = CrosslinkerGeometry.equilateral_triangle( + bond_length=0.27, connection_sites=[0, 1, 2] + ) + out = crosslink( + long_linear_path, + crosslinker=cl, + crosslink_bond_length=2.0, + excluded_bond_depth=1, + ) + # 20 backbone + 3 triangle = 23 + assert len(out.coordinates) == 23 + # 19 backbone edges + 3 internal + 3 connections = 25 + assert len(out.bond_graph.edges()) == 25 + + def test_crosslink_linear_dimer(self, linear_path): + cl = CrosslinkerGeometry.linear( + n_sites=2, bond_length=0.27, connection_sites=[0, 1] + ) + out = crosslink(linear_path, crosslinker=cl, crosslink_bond_length=3.0) + # 4 + 2 = 6 + assert len(out.coordinates) == 6 + # 3 backbone + 1 internal + 2 connections = 6 + assert len(out.bond_graph.edges()) == 6 + + def test_crosslink_no_backbone_raises(self): + coords = np.array([[0, 0, 0], [1, 0, 0]], dtype=np.float32) + p = Path(coords, bead_name="_B") # not "_A" + p._connect_edges("linear") + with pytest.raises(ValueError, match="No backbone beads"): + crosslink(p, backbone_name="_A", crosslink_bond_length=1.0) + + def test_crosslink_radius_too_small_raises(self, linear_path): + with pytest.raises(PathConvergenceError): + crosslink(linear_path, crosslink_bond_length=0.001, n_connection_sites=2) + + def test_crosslink_preserves_backbone_edges(self, linear_path): + out = crosslink(linear_path, crosslink_bond_length=3.0, n_connection_sites=2) + # Backbone edges 0-1, 1-2, 2-3 should still exist + assert out.bond_graph.has_edge(0, 1) + assert out.bond_graph.has_edge(1, 2) + assert out.bond_graph.has_edge(2, 3) + + def test_crosslink_bead_names_correct(self, linear_path): + out = crosslink( + linear_path, bead_name="_R", crosslink_bond_length=3.0, n_connection_sites=2 + ) + # Last bead should be the crosslinker + assert out.beads[4] == "_R" + # Backbone names unchanged + for i in range(4): + assert out.beads[i] == "_A" + + def test_crosslink_index_consistency(self, linear_path): + out = crosslink(linear_path, crosslink_bond_length=3.0, n_connection_sites=2) + # Every node in bond_graph should be a valid coordinate index + for node in out.bond_graph.nodes(): + assert 0 <= node < len(out.coordinates) + + def test_crosslink_with_initial_point(self, long_linear_path): + # Should find crosslink near specified point + out = crosslink( + long_linear_path, + crosslink_bond_length=2.0, + n_connection_sites=2, + initial_point=5, + ) + assert len(out.coordinates) == 21 + + def test_clink_repeated_replaces(self, long_linear_path): + triangle = CrosslinkerGeometry.equilateral_triangle( + bond_length=0.27, connection_sites=[0, 1, 2], bead_name="_U" + ) + crosslink(long_linear_path, triangle, crosslink_bond_length=1) + crosslink(long_linear_path, triangle, crosslink_bond_length=1) + assert len(long_linear_path) == 26 + assert sum(long_linear_path.beads == "_U") == 6 + assert sum(long_linear_path.beads == "_A") == 20 + + replace_sites(long_linear_path, triangle, bead_name="_A") + assert len(long_linear_path) == 66 + assert sum(long_linear_path.beads == "_U") == 66 + + def test_error_on_excess_degree(self, long_linear_path): + triangle = CrosslinkerGeometry.equilateral_triangle( + bond_length=0.27, connection_sites=[0, 1], bead_name="_U" + ) + crosslink(long_linear_path, triangle, crosslink_bond_length=1) + with pytest.raises(ValueError, match="degree"): + replace_sites(long_linear_path, triangle, bead_name="_A") + + +class TestReplaceSites(BaseTest): + """Tests for the replace_sites function.""" + + @pytest.fixture + def chain_with_crosslinks(self): + """A linear backbone with two single-site crosslinks.""" + n = 10 + coords = np.column_stack( + (np.linspace(0, 5, n), np.zeros(n), np.zeros(n)) + ).astype(np.float32) + p = Path(coords, bead_name="_A") + p._connect_edges("linear") + # Add two crosslinks + p = crosslink( + p, bead_name="_R", crosslink_bond_length=3.0, n_connection_sites=2, seed=42 + ) + p = crosslink( + p, bead_name="_R", crosslink_bond_length=3.0, n_connection_sites=2, seed=99 + ) + return p + + def test_replace_single_site_with_triangle(self): + # Build a simple path with a degree-2 node in the middle + coords = np.array( + [[0, 0, 0], [1, 0, 0], [2, 0, 0], [3, 0, 0], [4, 0, 0]], + dtype=np.float32, + ) + p = Path(coords, bead_name="_A") + p._connect_edges("linear") + # Change middle node to _R + p.beads[2] = "_R" + + triangle = CrosslinkerGeometry.equilateral_triangle( + bond_length=0.27, connection_sites=[0, 1] + ) + out = replace_sites(p, triangle, bead_name="_R") + # 5 - 1 removed + 3 added = 7 + assert len(out.coordinates) == 7 + # Check all node indices are valid + for node in out.bond_graph.nodes(): + assert 0 <= node < len(out.coordinates) + + def test_replace_by_index(self): + coords = np.array( + [[0, 0, 0], [1, 0, 0], [2, 0, 0], [3, 0, 0]], dtype=np.float32 + ) + p = Path(coords, bead_name="_A") + p._connect_edges("linear") + + linear_repl = CrosslinkerGeometry.linear( + n_sites=2, bond_length=0.3, connection_sites=[0, 1] + ) + # Replace node 1 (degree 2) + out = replace_sites(p, linear_repl, sites=[1]) + # 4 - 1 + 2 = 5 + assert len(out.coordinates) == 5 + assert ( + len(out.bond_graph.edges()) == 4 + ) # 2 kept + 1 internal + 1 external? let's just check > 0 + assert len(out.bond_graph.edges()) >= 4 + + def test_replace_preserves_nonreplaced_edges(self): + coords = np.array( + [[0, 0, 0], [1, 0, 0], [2, 0, 0], [3, 0, 0], [4, 0, 0]], + dtype=np.float32, + ) + p = Path(coords, bead_name="_A") + p._connect_edges("linear") + p.beads[2] = "_R" + + triangle = CrosslinkerGeometry.equilateral_triangle( + bond_length=0.27, connection_sites=[0, 1] + ) + out = replace_sites(p, triangle, bead_name="_R") + + # Original edge 0-1 and 3-4 should still exist in some form + # (nodes may be renumbered but connectivity is preserved) + # Check total number of edges: 2 kept backbone + 3 internal + 2 external = 7 + # Actually: edges 0-1, 1-2, 2-3, 3-4. Remove node 2. Kept: 0-1, 3-4. + # Triangle: 3 internal + 2 connections to (what was 1 and 3) + # Total: 2 + 3 + 2 = 7 + assert len(out.bond_graph.edges()) == 7 + + def test_replace_multiple_sites(self): + coords = np.array( + [[0, 0, 0], [1, 0, 0], [2, 0, 0], [3, 0, 0], [4, 0, 0], [5, 0, 0]], + dtype=np.float32, + ) + p = Path(coords, bead_name="_A") + p._connect_edges("linear") + p.beads[1] = "_R" + p.beads[4] = "_R" + + linear_repl = CrosslinkerGeometry.linear( + n_sites=2, bond_length=0.3, connection_sites=[0, 1] + ) + out = replace_sites(p, linear_repl, bead_name="_R") + # 6 - 2 + 4 = 8 + assert len(out.coordinates) == 8 + + def test_replace_no_sites_noop(self): + coords = np.array([[0, 0, 0], [1, 0, 0], [2, 0, 0]], dtype=np.float32) + p = Path(coords, bead_name="_A") + p._connect_edges("linear") + + triangle = CrosslinkerGeometry.equilateral_triangle( + bond_length=0.27, connection_sites=[0, 1] + ) + # No "_R" beads exist + out = replace_sites(p, triangle, bead_name="_R") + assert len(out.coordinates) == 3 + assert len(out.bond_graph.edges()) == 2 + + def test_replace_neither_sites_nor_bead_raises(self): + coords = np.array([[0, 0, 0], [1, 0, 0]], dtype=np.float32) + p = Path(coords, bead_name="_A") + triangle = CrosslinkerGeometry.equilateral_triangle(bond_length=0.27) + with pytest.raises(ValueError, match="Must specify"): + replace_sites(p, triangle) + + def test_replace_both_sites_and_bead_raises(self): + coords = np.array([[0, 0, 0], [1, 0, 0]], dtype=np.float32) + p = Path(coords, bead_name="_A") + triangle = CrosslinkerGeometry.equilateral_triangle(bond_length=0.27) + with pytest.raises(ValueError, match="Cannot specify both"): + replace_sites(p, triangle, sites=[0], bead_name="_A") + + def test_replace_renumbers_bond_graph(self): + coords = np.array( + [[0, 0, 0], [1, 0, 0], [2, 0, 0], [3, 0, 0]], dtype=np.float32 + ) + p = Path(coords, bead_name="_A") + p._connect_edges("linear") + p.beads[2] = "_R" + + linear_repl = CrosslinkerGeometry.linear( + n_sites=2, bond_length=0.3, connection_sites=[0, 1] + ) + out = replace_sites(p, linear_repl, bead_name="_R") + + # All node indices should be contiguous 0..N-1 + nodes = sorted(out.bond_graph.nodes()) + assert nodes == list(range(len(out.coordinates))) + + def test_replace_via_path_method(self): + coords = np.array( + [[0, 0, 0], [1, 0, 0], [2, 0, 0], [3, 0, 0]], dtype=np.float32 + ) + p = Path(coords, bead_name="_A") + p._connect_edges("linear") + p.beads[1] = "_R" + + linear_repl = CrosslinkerGeometry.linear( + n_sites=2, bond_length=0.3, connection_sites=[0, 1] + ) + out = p.replace_sites(linear_repl, bead_name="_R") + assert len(out.coordinates) == 5 + + def test_crosslink_then_replace(self, chain_with_crosslinks): + """Replace single-site crosslinks with triangles.""" + p = chain_with_crosslinks + n_R = sum(1 for b in p.beads if b == "_R") + # Each _R replaced: remove 1, add 3 -> net +2 per replacement + expected_len = len(p.coordinates) - n_R + n_R * 3 + assert n_R >= 1 + + triangle = CrosslinkerGeometry.equilateral_triangle( + bond_length=0.27, connection_sites=[0, 1] + ) + out = replace_sites(p, triangle, bead_name="_R") + assert len(out.coordinates) == expected_len + + # No _R beads remain (they've been replaced with triangle beads) + # Triangle beads are also "_R" so they should still exist + # But original single-site _R nodes are gone + for node in out.bond_graph.nodes(): + assert 0 <= node < len(out.coordinates) From 9ca9005b6cba3315d49f3baca10d2a79fd2e5b0a Mon Sep 17 00:00:00 2001 From: CalCraven Date: Thu, 21 May 2026 23:28:37 -0500 Subject: [PATCH 04/11] Initial implementation for crosslinking and replace sites in path --- mbuild/path/crosslink.py | 2455 ++++++++++++------------------- mbuild/tests/test_crosslinks.py | 151 +- 2 files changed, 1114 insertions(+), 1492 deletions(-) diff --git a/mbuild/path/crosslink.py b/mbuild/path/crosslink.py index c6fb51491..9a4a7c94c 100644 --- a/mbuild/path/crosslink.py +++ b/mbuild/path/crosslink.py @@ -348,442 +348,6 @@ def from_edges(cls, coordinates, edges, bead_name="_R", connection_sites=None): ) -# ============================================================================= -# Orientation / Rotation Utilities -# ============================================================================= - - -def _kabsch_rotation(P, Q): - """Optimal rotation aligning P onto Q (Kabsch algorithm).""" - H = P.T @ Q - U, S, Vt = np.linalg.svd(H) - d = np.linalg.det(Vt.T @ U.T) - sign_matrix = np.diag([1.0, 1.0, d]) - R = Vt.T @ sign_matrix @ U.T - return R.astype(np.float32) - - -def _rotation_between_vectors(a, b): - """Rotation matrix rotating unit vector a onto b (Rodrigues).""" - v = np.cross(a, b) - c = float(np.dot(a, b)) - s = float(np.linalg.norm(v)) - - if s < 1e-10: - if c > 0: - return np.eye(3, dtype=np.float32) - perp = np.array([1, 0, 0], dtype=np.float32) - if abs(np.dot(a, perp)) > 0.9: - perp = np.array([0, 1, 0], dtype=np.float32) - perp = perp - np.dot(perp, a) * a - perp /= np.linalg.norm(perp) - return (2 * np.outer(perp, perp) - np.eye(3)).astype(np.float32) - - vx = np.array( - [[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]], dtype=np.float32 - ) - R = np.eye(3, dtype=np.float32) + vx + vx @ vx * ((1 - c) / (s * s)) - return R - - -def _random_rotation_matrix(rng): - """Uniform random rotation via QR decomposition.""" - H = rng.standard_normal((3, 3)).astype(np.float32) - Q, R = np.linalg.qr(H) - Q *= np.sign(np.diag(R)) - if np.linalg.det(Q) < 0: - Q[:, 0] *= -1 - return Q.astype(np.float32) - - -def _rotation_about_axis(axis, angle): - """Rotation matrix for rotation by `angle` radians about `axis`.""" - axis = np.asarray(axis, dtype=np.float32) - axis = axis / np.linalg.norm(axis) - K = np.array( - [ - [0, -axis[2], axis[1]], - [axis[2], 0, -axis[0]], - [-axis[1], axis[0], 0], - ], - dtype=np.float32, - ) - R = np.eye(3, dtype=np.float32) + np.sin(angle) * K + (1 - np.cos(angle)) * (K @ K) - return R - - -def _orient_replacement( - replacement, - neighbor_vectors, - nearby_coords=None, - center=None, - rng=None, - n_rotation_samples=36, -): - """Orient a replacement geometry so connection sites point toward neighbors. - - Performs Kabsch alignment then searches residual rotational DOF to - minimize overlaps with nearby non-bonded sites. - - Parameters - ---------- - replacement : CrosslinkerGeometry - The replacement structure (centered at origin). - neighbor_vectors : np.ndarray, shape (n_connections, 3) - Vectors from placement center to each neighbor being bonded. - nearby_coords : np.ndarray, shape (M, 3), optional - Coordinates of nearby non-bonded sites (for overlap minimization). - Should be relative to the placement center. - center : np.ndarray, shape (3,), optional - The absolute position where replacement is centered (for overlap calc). - rng : numpy.random.Generator, optional - n_rotation_samples : int, default 36 - Number of rotational samples for overlap minimization. - - Returns - ------- - oriented_positions : np.ndarray, shape (n_sites, 3) - Rotated replacement positions (still centered at origin). - assignment : list of int - assignment[i] = j means connection_sites[i] bonds to neighbor j. - """ - if rng is None: - rng = np.random.default_rng(42) - - positions = replacement.coordinates.copy() - unique_conn = replacement.unique_connection_sites - conn_positions = positions[unique_conn] - - n_unique = len(unique_conn) - - # --- Degenerate case: all connection sites at origin --- - if np.allclose(conn_positions, 0, atol=1e-8): - R = _random_rotation_matrix(rng) - oriented = (R @ positions.T).T - return oriented, list(range(replacement.n_connections)) - - # --- Group target vectors by unique connection site (for repeated sites) --- - conn_site_to_targets = {} - for i, site_idx in enumerate(replacement.connection_sites): - conn_site_to_targets.setdefault(site_idx, []).append(i) - - # Use centroid of targets for each unique connection site - unique_targets = np.array( - [ - np.mean(neighbor_vectors[conn_site_to_targets[s]], axis=0) - for s in unique_conn - ], - dtype=np.float32, - ) - - # --- Normalize for angular alignment --- - P_norms = np.linalg.norm(conn_positions, axis=1, keepdims=True) - P_normed = conn_positions / np.maximum(P_norms, 1e-10) - - T_norms = np.linalg.norm(unique_targets, axis=1, keepdims=True) - T_normed = unique_targets / np.maximum(T_norms, 1e-10) - - # --- Single connection site case --- - if n_unique == 1: - R = _rotation_between_vectors(P_normed[0], T_normed[0]) - best_R = R - best_perm = [0] - else: - # --- Try all permutations (feasible for n <= 5) --- - best_perm = list(range(n_unique)) - best_R = np.eye(3, dtype=np.float32) - best_cost = np.inf - - for perm in permutations(range(n_unique)): - Q_perm = T_normed[list(perm)] - R = _kabsch_rotation(P_normed, Q_perm) - rotated = (R @ P_normed.T).T - cost = np.sum((rotated - Q_perm) ** 2) - if cost < best_cost: - best_cost = cost - best_perm = list(perm) - best_R = R - - oriented = (best_R @ positions.T).T - - # --- Overlap minimization: search residual rotation --- - if nearby_coords is not None and len(nearby_coords) > 0 and n_rotation_samples > 1: - # Determine a rotation axis (mean direction of connection sites after alignment) - rotated_conn = oriented[unique_conn] - mean_axis = np.mean(rotated_conn, axis=0) - mean_axis_norm = np.linalg.norm(mean_axis) - - if mean_axis_norm > 1e-8: - rotation_axis = mean_axis / mean_axis_norm - else: - # Use normal to the plane of connection sites or random axis - rotation_axis = np.array([0, 0, 1], dtype=np.float32) - - best_oriented = oriented.copy() - best_min_dist = _min_distance_to_neighbors(oriented, nearby_coords) - - angles = np.linspace(0, 2 * np.pi, n_rotation_samples, endpoint=False) - for angle in angles: - R_trial = _rotation_about_axis(rotation_axis, angle) - trial = (R_trial @ oriented.T).T - min_dist = _min_distance_to_neighbors(trial, nearby_coords) - if min_dist > best_min_dist: - best_min_dist = min_dist - best_oriented = trial - - oriented = best_oriented - - # --- Build full assignment --- - full_assignment = _build_full_assignment( - replacement, best_perm, unique_conn, conn_site_to_targets - ) - - return oriented, full_assignment - - -def _orient_replacement_partial( - replacement, - partial_conn_sites, - neighbor_vectors, - nearby_coords=None, - center=None, - rng=None, - n_rotation_samples=36, -): - """Orient a replacement when fewer neighbors exist than connection sites. - - Uses only the first ``len(neighbor_vectors)`` connection sites for alignment. - The remaining connection sites are left dangling (not bonded). - - Parameters - ---------- - replacement : CrosslinkerGeometry - The replacement structure (centered at origin). - partial_conn_sites : list of int - Subset of replacement.connection_sites to use for alignment. - Length matches len(neighbor_vectors). - neighbor_vectors : np.ndarray, shape (degree, 3) - Vectors from placement center to each neighbor being bonded. - nearby_coords : np.ndarray, shape (M, 3), optional - Coordinates of nearby non-bonded sites for overlap minimization. - center : np.ndarray, shape (3,), optional - Absolute position of placement center. - rng : numpy.random.Generator, optional - n_rotation_samples : int, default 36 - Number of rotational samples for overlap minimization. - - Returns - ------- - oriented_positions : np.ndarray, shape (n_sites, 3) - Rotated replacement positions (still centered at origin). - assignment : list of int - assignment[i] = j means partial_conn_sites[i] bonds to neighbor j. - Length equals len(partial_conn_sites) (i.e., degree). - """ - if rng is None: - rng = np.random.default_rng(42) - - positions = replacement.coordinates.copy() - degree = len(neighbor_vectors) - - if degree == 0: - R = _random_rotation_matrix(rng) - return (R @ positions.T).T, [] - - # Get positions of the partial connection sites - unique_partial = list(dict.fromkeys(partial_conn_sites)) - conn_positions = positions[unique_partial] - - # If all connection sites at origin, random rotation - if np.allclose(conn_positions, 0, atol=1e-8): - R = _random_rotation_matrix(rng) - oriented = (R @ positions.T).T - return oriented, list(range(degree)) - - # Group partial connections by unique site - conn_site_to_targets = {} - for i, site_idx in enumerate(partial_conn_sites): - conn_site_to_targets.setdefault(site_idx, []).append(i) - - # Use centroid of targets for each unique connection site - unique_targets = np.array( - [ - np.mean(neighbor_vectors[conn_site_to_targets[s]], axis=0) - for s in unique_partial - ], - dtype=np.float32, - ) - - # Normalize for angular alignment - P_norms = np.linalg.norm(conn_positions, axis=1, keepdims=True) - P_normed = conn_positions / np.maximum(P_norms, 1e-10) - - T_norms = np.linalg.norm(unique_targets, axis=1, keepdims=True) - T_normed = unique_targets / np.maximum(T_norms, 1e-10) - - n_unique = len(unique_partial) - - if n_unique == 1: - R = _rotation_between_vectors(P_normed[0], T_normed[0]) - best_R = R - best_perm = [0] - else: - # Try all permutations of unique sites -> unique targets - best_perm = list(range(n_unique)) - best_R = np.eye(3, dtype=np.float32) - best_cost = np.inf - - for perm in permutations(range(n_unique)): - Q_perm = T_normed[list(perm)] - R = _kabsch_rotation(P_normed, Q_perm) - rotated = (R @ P_normed.T).T - cost = np.sum((rotated - Q_perm) ** 2) - if cost < best_cost: - best_cost = cost - best_perm = list(perm) - best_R = R - - oriented = (best_R @ positions.T).T - - # --- Overlap minimization --- - if nearby_coords is not None and len(nearby_coords) > 0 and n_rotation_samples > 1: - rotated_conn = oriented[unique_partial] - mean_axis = np.mean(rotated_conn, axis=0) - mean_axis_norm = np.linalg.norm(mean_axis) - - if mean_axis_norm > 1e-8: - rotation_axis = mean_axis / mean_axis_norm - else: - rotation_axis = np.array([0, 0, 1], dtype=np.float32) - - best_oriented = oriented.copy() - best_min_dist = _min_distance_to_neighbors(oriented, nearby_coords) - - angles = np.linspace(0, 2 * np.pi, n_rotation_samples, endpoint=False) - for angle in angles: - R_trial = _rotation_about_axis(rotation_axis, angle) - trial = (R_trial @ oriented.T).T - min_dist = _min_distance_to_neighbors(trial, nearby_coords) - if min_dist > best_min_dist: - best_min_dist = min_dist - best_oriented = trial - - oriented = best_oriented - - # Build assignment: maps conn_idx (0..degree-1) -> neighbor_idx - full_assignment = _build_full_assignment_partial( - partial_conn_sites, best_perm, unique_partial, conn_site_to_targets - ) - - return oriented, full_assignment - - -def _build_full_assignment_partial( - partial_conn_sites, perm, unique_partial, conn_site_to_targets -): - """Build assignment for partial connection case. - - Parameters - ---------- - partial_conn_sites : list of int - The subset of connection sites being used. - perm : list of int - Permutation of unique sites to unique targets. - unique_partial : list of int - Unique site indices in partial_conn_sites. - conn_site_to_targets : dict - Maps site_idx -> list of connection indices using that site. - - Returns - ------- - full_assignment : list of int - assignment[i] = neighbor_index for partial_conn_sites[i]. - """ - degree = len(partial_conn_sites) - - if degree == len(unique_partial) and all( - len(v) == 1 for v in conn_site_to_targets.values() - ): - # Simple case: no repeated sites, 1:1 mapping - return perm - - # General case with possible repeated sites - full_assignment = [0] * degree - site_counter = {} - - for conn_idx, site_idx in enumerate(partial_conn_sites): - unique_pos = unique_partial.index(site_idx) - target_neighbor_group = perm[unique_pos] - - count = site_counter.get(site_idx, 0) - target_site = unique_partial[target_neighbor_group] - target_indices = conn_site_to_targets[target_site] - full_assignment[conn_idx] = target_indices[count % len(target_indices)] - site_counter[site_idx] = count + 1 - - return full_assignment - - -def _min_distance_to_neighbors(positions, nearby_coords): - """Compute minimum distance between any replacement site and nearby sites.""" - if len(nearby_coords) == 0: - return np.inf - # Pairwise distances - diffs = positions[:, None, :] - nearby_coords[None, :, :] - dists = np.linalg.norm(diffs, axis=-1) - return np.min(dists) - - -def _build_full_assignment(replacement, perm, unique_conn, conn_site_to_targets): - """Map each connection_sites entry to a neighbor index using the permutation. - - Parameters - ---------- - replacement : CrosslinkerGeometry - perm : list of int - Permutation mapping unique connection site positions to neighbor group positions. - unique_conn : list of int - Unique connection site indices. - conn_site_to_targets : dict - Maps site_idx -> list of connection indices using that site. - - Returns - ------- - full_assignment : list of int - assignment[i] = neighbor_index for connection_sites[i]. - """ - n_connections = replacement.n_connections - - if n_connections == len(unique_conn): - # No repeated sites: perm directly maps connection -> neighbor - return perm - - # With repeated sites: need to assign individual connections within groups - full_assignment = [0] * n_connections - - # perm[unique_pos] tells us which neighbor group this unique site maps to - # We need to iterate through connection_sites and assign appropriately - site_counter = {} # track how many times each site has been assigned - for conn_idx, site_idx in enumerate(replacement.connection_sites): - unique_pos = unique_conn.index(site_idx) - target_neighbor_group = perm[unique_pos] - - # Round-robin within the target group - count = site_counter.get(site_idx, 0) - # The target_neighbor_group-th unique site's target list - target_site = unique_conn[target_neighbor_group] - target_indices = conn_site_to_targets[target_site] - full_assignment[conn_idx] = target_indices[count % len(target_indices)] - site_counter[site_idx] = count + 1 - - return full_assignment - - -# ============================================================================= -# General replace_sites function -# ============================================================================= - - def replace_sites( path, replacement, @@ -798,48 +362,7 @@ def replace_sites( seed=42, ): """Replace sites in a Path with a multi-site substructure, in place. - - Uses rigid-body alignment to place the replacement such that: - - Connection sites are at specified bond_length from their neighbors - - Internal geometry is preserved exactly (rigid body) - - Placed sites are at least min_separation from non-bonded neighbors - - The optimization sweeps residual rotational degrees of freedom and, if - needed, offsets the placement perpendicular to the constraint axis to - find a clash-free configuration. - - Parameters - ---------- - path : Path - The path to modify in place. - replacement : CrosslinkerGeometry - The replacement structure. - sites : list of int, optional - Explicit node indices to replace. Mutually exclusive with ``bead_name``. - bead_name : str, optional - Replace all sites matching this bead name. Mutually exclusive with ``sites``. - bond_length : float, optional - Desired bond length from each connection site to its neighbor. - If None, uses the existing distance from the replaced site to each neighbor. - tolerance : float, default 0.1 - Fractional tolerance on bond lengths (±10%). - min_separation : float, optional - Minimum allowed distance between any placed crosslinker site and any - existing non-bonded site. If None, no overlap checking is performed. - volume_constraint : CuboidConstraint or CylinderConstraint, optional - For periodic boundary handling. - n_rotation_samples : int, default 36 - Angular samples for residual DOF optimization. - overlap_radius : float, optional - Radius around each replaced site to gather nearby sites for overlap check. - If None, defaults to 3x max extent of replacement geometry. - seed : int, default 42 - Random seed. - - Returns - ------- - Path - The same path object (modified in place), for method chaining. + ... """ if sites is None and bead_name is None: raise ValueError("Must specify either `sites` or `bead_name`.") @@ -852,7 +375,10 @@ def replace_sites( sites = [ node for node in path.bond_graph.nodes() if path.beads[node] == bead_name ] - sites = list(sites) + if isinstance(sites, int): + sites = [sites] + else: + sites = list(sites) if len(sites) == 0: return path @@ -867,14 +393,6 @@ def replace_sites( ) pbc, box_lengths = _get_pbc_info(volume_constraint) - - if overlap_radius is None: - if replacement.n_sites > 0: - max_extent = np.max(np.linalg.norm(replacement.coordinates, axis=1)) - overlap_radius = max(max_extent * 3, 0.5) - else: - overlap_radius = 0.5 - sites_set = set(sites) # --- Compute all replacements --- @@ -885,7 +403,7 @@ def replace_sites( site_pos = path.coordinates[site] degree = len(neighbors) - # Unwrapped neighbor positions (absolute) + # Unwrapped neighbor positions neighbor_coords = np.zeros((degree, 3), dtype=np.float32) for i, nb in enumerate(neighbors): delta = path.coordinates[nb] - site_pos @@ -901,21 +419,21 @@ def replace_sites( dtype=np.float32, ) - # Nearby non-bonded coords (absolute positions for overlap checking) - nearby_coords = _get_nearby_nonbonded_absolute( - path, site, neighbors, site_pos, overlap_radius, pbc, box_lengths, sites_set - ) - - # Active connection sites (may be fewer than total if degree < n_connections) active_conn_sites = replacement.connection_sites[:degree] + # Indices that are bonded to this site (excluded from overlap check) + bonded_indices = set(neighbors) | {site} | sites_set + # Place via rigid body alignment + overlap optimization - placed_positions, assignment = _place_replacement_rigid( + placed_positions, assignment = _place_replacement_rigid_full( replacement, active_conn_sites, neighbor_coords, target_bond_lengths, - nearby_coords, + path.coordinates, # pass ALL coordinates + bonded_indices, + pbc, + box_lengths, rng, n_rotation_samples, tolerance, @@ -931,7 +449,7 @@ def replace_sites( "assignment": assignment, } - # --- Rebuild path in place --- + # --- Rebuild path in place (same as before) --- kept_nodes = [n for n in sorted(path.bond_graph.nodes()) if n not in sites_set] new_coords = [] @@ -956,25 +474,21 @@ def replace_sites( current_idx += 1 site_new_indices[site] = indices - # --- Build edges --- new_edges = [] - # Kept-to-kept for u, v, data in path.bond_graph.edges(data=True): if u in sites_set or v in sites_set: continue new_edges.append((old_to_new[u], old_to_new[v], data)) - # Internal replacement bonds for site in sites: indices = site_new_indices[site] for i, j in replacement.bond_graph.edges(): edge_data = { - "bond_type": (str(replacement.beads[i]), str(replacement.beads[j])), + "bond_type": (str(replacement.beads[i]), str(replacement.beads[j])) } new_edges.append((indices[i], indices[j], edge_data)) - # Connection-to-neighbor bonds for site in sites: data = replacement_data[site] indices = site_new_indices[site] @@ -984,7 +498,6 @@ def replace_sites( for conn_idx, neighbor_idx in enumerate(assignment): repl_site_local = replacement.connection_sites[conn_idx] repl_new_node = indices[repl_site_local] - external_neighbor = neighbors[neighbor_idx] if external_neighbor in old_to_new: @@ -1019,11 +532,10 @@ def replace_sites( "bond_type": ( str(replacement.beads[repl_site_local]), str(ext_bead_name), - ), + ) } new_edges.append((repl_new_node, ext_new_node, edge_data)) - # --- Overwrite path --- path.coordinates = np.array(new_coords, dtype=np.float32) path.beads = np.array(new_beads, dtype="U10") @@ -1041,534 +553,464 @@ def replace_sites( return path -def _get_nearby_nonbonded_absolute( - path, site, neighbors, site_pos, radius, pbc, box_lengths, excluded_sites +def crosslink( + path, + crosslinker=None, + bead_name="_R", + backbone_name="_A", + crosslink_bond_length=0.2, + max_backbone_degree=2, + tolerance=0.1, + excluded_bond_depth=2, + n_connection_sites=2, + volume_constraint=None, + initial_point=None, + seed=42, + n_rotation_samples=36, + overlap_radius=None, + min_separation=None, ): - """Get absolute positions of nearby non-bonded sites for overlap checking. + """ + Create a crosslink bonded to backbone beads, modifying path in place. - Unlike the relative version, returns absolute (unwrapped relative to site_pos) - positions suitable for direct distance comparison with placed crosslinker sites. + Selects backbone beads whose pairwise spacing matches the crosslinker + geometry, places a temporary placeholder, then calls ``replace_sites`` + which handles rigid-body alignment, bond length guarantees, and overlap + avoidance in a single pass. Parameters ---------- path : Path - site : int - neighbors : list of int - site_pos : np.ndarray, shape (3,) - radius : float - pbc : array-like of bool - box_lengths : np.ndarray - excluded_sites : set of int - - Returns - ------- - np.ndarray, shape (M, 3) - """ - neighbors_set = set(neighbors) | {site} | excluded_sites - - nearby = [] - for node in path.bond_graph.nodes(): - if node in neighbors_set: - continue - delta = path.coordinates[node] - site_pos - delta = _apply_mic(delta, pbc, box_lengths) - dist = np.linalg.norm(delta) - if dist <= radius: - # Return absolute unwrapped position - nearby.append(site_pos + delta) - - if len(nearby) == 0: - return np.empty((0, 3), dtype=np.float32) - return np.array(nearby, dtype=np.float32) - - -def _place_replacement_rigid( - replacement, - active_conn_sites, - neighbor_coords, - target_bond_lengths, - nearby_coords, - rng, - n_rotation_samples, - tolerance, - min_separation=None, -): - """Place replacement via rigid-body alignment guaranteeing bond lengths and no overlaps. - - Strategy: - 1. Compute target positions for connection sites (at bond_length from neighbors). - 2. Solve Procrustes alignment (rotation + translation) to map connection sites to targets. - 3. For residual rotational DOF (1 DOF for 2 unique sites, full DOF for 1 unique site): - - Sweep angles and pick the one that: - a) Keeps all connection-to-neighbor distances within tolerance - b) Maximizes minimum distance to nearby non-bonded sites (avoids overlaps) - 4. If no rotation satisfies min_separation, offset the center along the rotation axis - (push the crosslinker "above" or "below" the backbone plane) and re-sweep. - 5. As a last resort, return the best-scoring rotation even if below min_separation. - - Parameters - ---------- - replacement : CrosslinkerGeometry - The replacement structure (centered at origin). - active_conn_sites : list of int - Subset of replacement.connection_sites to bond (length = degree). - neighbor_coords : np.ndarray, shape (degree, 3) - Unwrapped absolute positions of neighboring beads. - target_bond_lengths : np.ndarray, shape (degree,) - Desired bond length from each connection site to its neighbor. - nearby_coords : np.ndarray, shape (M, 3) - Absolute positions of nearby non-bonded sites for overlap avoidance. - rng : numpy.random.Generator - n_rotation_samples : int - tolerance : float - Fractional tolerance on bond lengths. + The Path object to modify in place. + crosslinker : CrosslinkerGeometry, optional + If None, uses single-site from ``bead_name`` and ``n_connection_sites``. + bead_name : str, default "_R" + backbone_name : str, default "_A" + crosslink_bond_length : float, default 0.2 + Desired bond length from connection sites to backbone beads (nm). + tolerance : float, default 0.1 + Fractional tolerance on bond lengths (±10%). + excluded_bond_depth : int, default 2 + n_connection_sites : int, default 2 + volume_constraint : optional + initial_point : int or array-like, optional + seed : int, default 42 + chunk_size : int, default 512 + run_on_gpu : bool, default False + n_rotation_samples : int, default 36 + overlap_radius : float, optional min_separation : float, optional - Minimum allowed distance between any placed site and nearby sites. - If None, no overlap checking is performed. + Minimum distance from placed sites to existing sites. + Defaults to ``crosslink_bond_length * 0.5``. Returns ------- - placed_positions : np.ndarray, shape (n_sites, 3) - Final absolute positions for all replacement sites. - assignment : list of int - assignment[i] = neighbor_index for active_conn_sites[i]. + Path + Same path, modified in place. """ - degree = len(active_conn_sites) - - if degree == 0: - center = ( - np.mean(neighbor_coords, axis=0) - if len(neighbor_coords) > 0 - else np.zeros(3) - ) - R = _random_rotation_matrix(rng) - placed = (R @ replacement.coordinates.T).T + center - return placed.astype(np.float32), [] - - # --- Compute target positions for connection sites --- - neighbor_centroid = np.mean(neighbor_coords, axis=0) - - target_positions = np.zeros((degree, 3), dtype=np.float32) - for i in range(degree): - direction = neighbor_centroid - neighbor_coords[i] - d = np.linalg.norm(direction) - if d > 1e-10: - direction_hat = direction / d - else: - direction_hat = rng.standard_normal(3).astype(np.float32) - direction_hat /= np.linalg.norm(direction_hat) - target_positions[i] = ( - neighbor_coords[i] + direction_hat * target_bond_lengths[i] + if crosslinker is None: + crosslinker = CrosslinkerGeometry.single_site( + bead_name=bead_name, n_connections=n_connection_sites ) - # --- Get unique connection site info --- - unique_active = list(dict.fromkeys(active_conn_sites)) - unique_conn_positions = replacement.coordinates[unique_active] - n_unique = len(unique_active) + if min_separation is None: + min_separation = crosslink_bond_length * 0.2 - conn_to_target_indices = {} - for i, site_idx in enumerate(active_conn_sites): - conn_to_target_indices.setdefault(site_idx, []).append(i) + n_backbone_connections = crosslinker.n_connections + rng = np.random.default_rng(seed + len(path.coordinates)) - unique_targets = np.array( - [ - np.mean(target_positions[conn_to_target_indices[s]], axis=0) - for s in unique_active - ], - dtype=np.float32, + # --- Compute search geometry --- + ideal_bb_positions, ideal_pairwise_distances, geom_search_radius = ( + _compute_ideal_backbone_distances(crosslinker, crosslink_bond_length) ) - # --- Initial Procrustes alignment --- - if n_unique == 1: - base_placed, base_perm, rotation_type, rotation_info = _initial_align_single( - replacement, - unique_active, - unique_conn_positions, - unique_targets, - neighbor_centroid, - rng, - ) - elif n_unique >= 2: - base_placed, base_perm, rotation_type, rotation_info = _initial_align_multi( - replacement, - unique_active, - unique_conn_positions, - unique_targets, - rng, - ) + if n_backbone_connections > 1: + max_pair_dist = np.max(ideal_pairwise_distances) else: - R = _random_rotation_matrix(rng) - center = np.mean(neighbor_coords, axis=0) - placed = (R @ replacement.coordinates.T).T + center - return placed.astype(np.float32), [] + max_pair_dist = 0.0 - # --- Optimize residual DOF for overlap avoidance + bond validity --- - placed = _optimize_residual_rotation( - base_placed, - rotation_type, - rotation_info, - replacement, - active_conn_sites, - neighbor_coords, - target_bond_lengths, - nearby_coords, - tolerance, - min_separation, - n_rotation_samples, - rng, - ) + search_radius = max_pair_dist * (1 + tolerance) - # --- Build assignment --- - assignment = _build_assignment_from_perm( - active_conn_sites, unique_active, conn_to_target_indices, base_perm, degree - ) + # --- Find backbone targets --- + backbone_nodes = [ + node for node in path.bond_graph.nodes() if path.beads[node] == backbone_name + ] + backbone_subgraph = path.bond_graph.subgraph(backbone_nodes) - return placed.astype(np.float32), assignment + if len(backbone_nodes) == 0: + raise ValueError(f"No backbone beads with name '{backbone_name}' found in path") + if len(backbone_nodes) < n_backbone_connections: + raise ValueError( + f"Not enough backbone beads ({len(backbone_nodes)}) for " + f"{n_backbone_connections} connection sites" + ) + pbc, box_lengths = _get_pbc_info(volume_constraint) -def _initial_align_single( - replacement, - unique_active, - unique_conn_positions, - unique_targets, - neighbor_centroid, - rng, -): - """Initial alignment for 1 unique connection site. + candidate_nodes = [ + node + for node in backbone_nodes + if path.bond_graph.degree[node] <= max_backbone_degree + ] + candidate_coords = np.array(path.coordinates[candidate_nodes], dtype=np.float32) - Returns base placement and rotation info for residual DOF. + ref_nodes, ref_coords = _get_reference_points( + path, initial_point, candidate_nodes, pbc, box_lengths, rng + ) - Returns - ------- - base_placed : np.ndarray, shape (n_sites, 3) - perm : list - rotation_type : str - 'axis' for rotation about an axis, 'full' for full sphere - rotation_info : dict - Contains 'pivot' and 'axis' for residual rotation. - """ - c = unique_conn_positions[0] - c_norm = np.linalg.norm(c) + found_ref = False + selected_nodes = [] - if c_norm < 1e-10: - R = _random_rotation_matrix(rng) - t = unique_targets[0] - base_placed = (R @ replacement.coordinates.T).T + t - return base_placed, [0], "full", {"pivot": unique_targets[0]} + for ref_node, ref_coord in zip(ref_nodes, ref_coords): + selected_nodes = [ref_node] - c_hat = c / c_norm - desired_dir = unique_targets[0] - neighbor_centroid - desired_norm = np.linalg.norm(desired_dir) - if desired_norm < 1e-10: - desired_dir = rng.standard_normal(3).astype(np.float32) - desired_norm = np.linalg.norm(desired_dir) - desired_hat = desired_dir / desired_norm + sq_distances = calculate_sq_distances( + ref_coord, candidate_coords, pbc=pbc, box_lengths=box_lengths + ) + distances = np.sqrt(sq_distances) - R = _rotation_between_vectors(c_hat, desired_hat) - t = unique_targets[0] - R @ c - base_placed = (R @ replacement.coordinates.T).T + t + within_radius_mask = distances <= search_radius + possible_pairs = np.where(within_radius_mask)[0] - # Residual DOF: rotation about axis from center to connection site - conn_placed = base_placed[unique_active[0]] - center = t # crosslinker center position - axis = conn_placed - center - axis_norm = np.linalg.norm(axis) - if axis_norm > 1e-10: - axis_hat = axis / axis_norm - else: - axis_hat = desired_hat + if len(possible_pairs) < n_backbone_connections - 1: + continue - return base_placed, [0], "axis", {"pivot": conn_placed, "axis": axis_hat} + closest_paired_nodes = possible_pairs[np.argsort(distances[possible_pairs])] + excluded_nodes = set( + nx.single_source_shortest_path_length( + backbone_subgraph, ref_node, cutoff=excluded_bond_depth + ).keys() + ) + for idx in closest_paired_nodes: + node = candidate_nodes[idx] + if node in excluded_nodes: + continue -def _initial_align_multi( - replacement, - unique_active, - unique_conn_positions, - unique_targets, - rng, -): - """Initial alignment for 2+ unique connection sites via Procrustes. + is_excluded = False + for selected in selected_nodes[1:]: + if selected in backbone_subgraph: + sel_excluded = set( + nx.single_source_shortest_path_length( + backbone_subgraph, selected, cutoff=excluded_bond_depth + ).keys() + ) + if node in sel_excluded: + is_excluded = True + break + if is_excluded: + continue - Returns - ------- - base_placed : np.ndarray - perm : list - rotation_type : str - 'axis' for 2 sites, 'none' for 3+ - rotation_info : dict - """ - n_unique = len(unique_active) + selected_nodes.append(int(node)) + if len(selected_nodes) >= n_backbone_connections: + # Validate pairwise geometry + sel_coords = np.array( + [path.coordinates[n] for n in selected_nodes], dtype=np.float32 + ) + unwrapped_check = _unwrap_coords(sel_coords, pbc, box_lengths) + valid, _ = _validate_backbone_group( + unwrapped_check, + ideal_pairwise_distances, + crosslink_bond_length, + tolerance, + pbc, + box_lengths, + ) + valid_overlaps, clink_centroid = _validate_clinker( + path, + crosslinker, + crosslink_bond_length, + selected_nodes, + overlap_radius=min_separation, + ) + if valid and valid_overlaps: + found_ref = True + break + else: + selected_nodes.pop() - best_R = np.eye(3, dtype=np.float32) - best_t = np.zeros(3, dtype=np.float32) - best_cost = np.inf - best_perm = list(range(n_unique)) + if found_ref: + break - for perm in permutations(range(n_unique)): - perm_targets = unique_targets[list(perm)] + if not found_ref: + n_clinks = sum(1 for b in path.beads if b in set(crosslinker.beads)) + raise PathConvergenceError( + f"Could not find {n_backbone_connections} non-neighboring backbone beads " + f"matching crosslinker geometry with bond_length={crosslink_bond_length} " + f"(tolerance=±{tolerance * 100:.0f}%)." + f"\nCurrent crosslinks: {n_clinks}. " + "Ways to increase crosslinking:\n" + " - Increase crosslink_bond_length\n" + " - Increase tolerance\n" + " - Decrease excluded_bond_depth\n" + " - Pack at higher density\n" + " - Relax structure" + ) - src_centroid = np.mean(unique_conn_positions, axis=0) - tgt_centroid = np.mean(perm_targets, axis=0) + for dim in range(3): + if pbc[dim]: + clink_centroid[dim] -= ( + np.round(clink_centroid[dim] / box_lengths[dim]) * box_lengths[dim] + ) - src_centered = unique_conn_positions - src_centroid - tgt_centered = perm_targets - tgt_centroid + _placeholder_name = "__placeholder__" + path.append_coordinates(clink_centroid, _placeholder_name) + placeholder_idx = len(path.coordinates) - 1 - R = _kabsch_rotation(src_centered, tgt_centered) - t = tgt_centroid - R @ src_centroid + for backbone_node in selected_nodes: + path.bond_graph.add_edge( + int(placeholder_idx), + int(backbone_node), + bond_type=(_placeholder_name, backbone_name), + ) - placed_conn = (R @ unique_conn_positions.T).T + t - cost = np.sum((placed_conn - perm_targets) ** 2) + # --- Single-site shortcut --- + if crosslinker.n_sites == 1: + path.beads[placeholder_idx] = crosslinker.beads[0] + for backbone_node in selected_nodes: + path.bond_graph.edges[placeholder_idx, backbone_node]["bond_type"] = ( + str(crosslinker.beads[0]), + backbone_name, + ) + if "name" in path.bond_graph.nodes[placeholder_idx]: + path.bond_graph.nodes[placeholder_idx]["name"] = str(crosslinker.beads[0]) - if cost < best_cost: - best_cost = cost - best_R = R - best_t = t - best_perm = list(perm) + # Fix single-site bond length if needed + for backbone_node in selected_nodes: + delta = path.coordinates[placeholder_idx] - path.coordinates[backbone_node] + delta = _apply_mic(delta, pbc, box_lengths) + actual_bl = np.linalg.norm(delta) + if ( + abs(actual_bl - crosslink_bond_length) + > crosslink_bond_length * tolerance + ): + if actual_bl > 1e-10: + direction = delta / actual_bl + else: + direction = rng.standard_normal(3).astype(np.float32) + direction /= np.linalg.norm(direction) + path.coordinates[placeholder_idx] = path.coordinates[ + backbone_node + ] + _apply_mic(direction * crosslink_bond_length, pbc, box_lengths) + return path - base_placed = (best_R @ replacement.coordinates.T).T + best_t + # --- Multi-site: replace_sites handles everything --- + replace_sites( + path, + replacement=crosslinker, + sites=[placeholder_idx], + bond_length=crosslink_bond_length, + tolerance=tolerance, + min_separation=min_separation, + volume_constraint=volume_constraint, + n_rotation_samples=n_rotation_samples, + overlap_radius=overlap_radius, + seed=seed, + ) - if n_unique == 2: - # Residual DOF: rotation about axis connecting the two connection sites - conn0 = base_placed[unique_active[0]] - conn1 = base_placed[unique_active[1]] - axis = conn1 - conn0 - axis_norm = np.linalg.norm(axis) - if axis_norm > 1e-10: - axis_hat = axis / axis_norm + return path + + +def _get_reference_points(path, initial_point, candidate_nodes, pbc, box_lengths, rng): + candidate_coords = np.array(path.coordinates[candidate_nodes], dtype=np.float32) + if initial_point is not None: + if isinstance(initial_point, (int, np.integer)): + if initial_point not in path.bond_graph.nodes: + raise ValueError(f"Node {initial_point} not found in bond_graph") + return np.array([initial_point]), np.array( + [path.coordinates[initial_point]] + ) else: - axis_hat = np.array([0, 0, 1], dtype=np.float32) - midpoint = (conn0 + conn1) / 2.0 - return base_placed, best_perm, "axis", {"pivot": midpoint, "axis": axis_hat} + initial_point32 = np.asarray(initial_point, dtype=np.float32) + sq_distances = calculate_sq_distances( + initial_point32, candidate_coords, pbc=pbc, box_lengths=box_lengths + ) + order = np.argsort(sq_distances) + return ( + np.array(candidate_nodes)[order], + path.coordinates[np.array(candidate_nodes)[order]], + ) else: - # n_unique >= 3: fully determined, no residual DOF - return base_placed, best_perm, "none", {} + shuffled = rng.choice(candidate_nodes, size=len(candidate_nodes), replace=False) + return shuffled, path.coordinates[shuffled] -def _optimize_residual_rotation( - base_placed, - rotation_type, - rotation_info, - replacement, - active_conn_sites, - neighbor_coords, - target_bond_lengths, - nearby_coords, - tolerance, - min_separation, - n_rotation_samples, - rng, -): - """Optimize residual rotational DOF to avoid overlaps while maintaining bond lengths. +def _kabsch_rotation(P, Q): + """Optimal rotation aligning P onto Q (Kabsch algorithm).""" + H = P.T @ Q + U, S, Vt = np.linalg.svd(H) + d = np.linalg.det(Vt.T @ U.T) + sign_matrix = np.diag([1.0, 1.0, d]) + R = Vt.T @ sign_matrix @ U.T + return R.astype(np.float32) - Tries multiple rotations (and offsets if needed), scoring each by: - 1. Bond length validity (hard constraint) - 2. Minimum distance to nearby sites (soft maximize) - If no rotation passes min_separation, tries pushing the crosslinker along - the perpendicular/normal direction at fractional offsets, then re-sweeps. +def _rotation_between_vectors(a, b): + """Rotation matrix rotating unit vector a onto b (Rodrigues).""" + v = np.cross(a, b) + c = float(np.dot(a, b)) + s = float(np.linalg.norm(v)) - Parameters - ---------- - base_placed : np.ndarray, shape (n_sites, 3) - Initial placement from Procrustes. - rotation_type : str - 'axis', 'full', or 'none' - rotation_info : dict - Contains 'pivot' and 'axis' as needed. - replacement : CrosslinkerGeometry - active_conn_sites : list of int - neighbor_coords : np.ndarray, shape (degree, 3) - target_bond_lengths : np.ndarray, shape (degree,) - nearby_coords : np.ndarray, shape (M, 3) - tolerance : float - min_separation : float or None - n_rotation_samples : int - rng : numpy.random.Generator + if s < 1e-10: + if c > 0: + return np.eye(3, dtype=np.float32) + perp = np.array([1, 0, 0], dtype=np.float32) + if abs(np.dot(a, perp)) > 0.9: + perp = np.array([0, 1, 0], dtype=np.float32) + perp = perp - np.dot(perp, a) * a + perp /= np.linalg.norm(perp) + return (2 * np.outer(perp, perp) - np.eye(3)).astype(np.float32) - Returns - ------- - best_placed : np.ndarray, shape (n_sites, 3) - """ - if rotation_type == "none" or n_rotation_samples <= 1: - # No residual DOF or no sampling requested - return base_placed + vx = np.array( + [[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]], dtype=np.float32 + ) + R = np.eye(3, dtype=np.float32) + vx + vx @ vx * ((1 - c) / (s * s)) + return R - has_nearby = nearby_coords is not None and len(nearby_coords) > 0 - if not has_nearby and min_separation is None: - return base_placed - # Generate candidate rotations - if rotation_type == "axis": - pivot = rotation_info["pivot"] - axis_hat = rotation_info["axis"] - angles = np.linspace(0, 2 * np.pi, n_rotation_samples, endpoint=False) +def _random_rotation_matrix(rng): + """Uniform random rotation via QR decomposition.""" + H = rng.standard_normal((3, 3)).astype(np.float32) + Q, R = np.linalg.qr(H) + Q *= np.sign(np.diag(R)) + if np.linalg.det(Q) < 0: + Q[:, 0] *= -1 + return Q.astype(np.float32) - def generate_candidate(angle): - R_rot = _rotation_about_axis(axis_hat, angle) - return (R_rot @ (base_placed - pivot).T).T + pivot - elif rotation_type == "full": - pivot = rotation_info["pivot"] +def _rotation_about_axis(axis, angle): + """Rotation matrix for rotation by `angle` radians about `axis`.""" + axis = np.asarray(axis, dtype=np.float32) + axis = axis / np.linalg.norm(axis) + K = np.array( + [ + [0, -axis[2], axis[1]], + [axis[2], 0, -axis[0]], + [-axis[1], axis[0], 0], + ], + dtype=np.float32, + ) + R = np.eye(3, dtype=np.float32) + np.sin(angle) * K + (1 - np.cos(angle)) * (K @ K) + return R - # Sample random rotations about the pivot - def generate_candidate(idx): - R_rot = _random_rotation_matrix(rng) - return (R_rot @ (base_placed - pivot).T).T + pivot - angles = list(range(n_rotation_samples)) - else: - return base_placed +def _initial_align_single( + replacement, + unique_active, + unique_conn_positions, + unique_targets, + neighbor_centroid, + rng, +): + """Initial alignment for 1 unique connection site. - # --- Score each candidate --- - best_placed = base_placed.copy() - best_min_dist = -np.inf - best_valid = False + Returns base placement and rotation info for residual DOF. - bond_tol = tolerance # fractional + Returns + ------- + base_placed : np.ndarray, shape (n_sites, 3) + perm : list + rotation_type : str + 'axis' for rotation about an axis, 'full' for full sphere + rotation_info : dict + Contains 'pivot' and 'axis' for residual rotation. + """ + c = unique_conn_positions[0] + c_norm = np.linalg.norm(c) + + if c_norm < 1e-10: + R = _random_rotation_matrix(rng) + t = unique_targets[0] + base_placed = (R @ replacement.coordinates.T).T + t + return base_placed, [0], "full", {"pivot": unique_targets[0]} - for angle in angles: - candidate = generate_candidate(angle) + c_hat = c / c_norm + desired_dir = unique_targets[0] - neighbor_centroid + desired_norm = np.linalg.norm(desired_dir) + if desired_norm < 1e-10: + desired_dir = rng.standard_normal(3).astype(np.float32) + desired_norm = np.linalg.norm(desired_dir) + desired_hat = desired_dir / desired_norm - # Check bond validity - bonds_ok = _check_bonds_valid( - candidate, active_conn_sites, neighbor_coords, target_bond_lengths, bond_tol - ) - if not bonds_ok: - continue + R = _rotation_between_vectors(c_hat, desired_hat) + t = unique_targets[0] - R @ c + base_placed = (R @ replacement.coordinates.T).T + t - # Score: minimum distance to nearby sites - if has_nearby: - min_dist = _min_distance_to_neighbors(candidate, nearby_coords) - else: - min_dist = np.inf - - # Track best valid placement - is_valid = min_separation is None or min_dist >= min_separation - - if is_valid and (not best_valid or min_dist > best_min_dist): - best_placed = candidate.copy() - best_min_dist = min_dist - best_valid = True - elif not best_valid and min_dist > best_min_dist: - # No valid found yet, track best-effort - best_placed = candidate.copy() - best_min_dist = min_dist - - # --- If no valid rotation found, try offset along perpendicular --- - if not best_valid and min_separation is not None and rotation_type == "axis": - best_placed, best_valid = _try_offset_placements( - base_placed, - rotation_info, - replacement, - active_conn_sites, - neighbor_coords, - target_bond_lengths, - nearby_coords, - bond_tol, - min_separation, - n_rotation_samples, - rng, - ) + # Residual DOF: rotation about axis from center to connection site + conn_placed = base_placed[unique_active[0]] + center = t # crosslinker center position + axis = conn_placed - center + axis_norm = np.linalg.norm(axis) + if axis_norm > 1e-10: + axis_hat = axis / axis_norm + else: + axis_hat = desired_hat - return best_placed + return base_placed, [0], "axis", {"pivot": conn_placed, "axis": axis_hat} -def _try_offset_placements( - base_placed, - rotation_info, +def _initial_align_multi( replacement, - active_conn_sites, - neighbor_coords, - target_bond_lengths, - nearby_coords, - bond_tol, - min_separation, - n_rotation_samples, + unique_active, + unique_conn_positions, + unique_targets, rng, ): - """Try offsetting the crosslinker perpendicular to the rotation axis. - - When the crosslinker is in the plane of the backbone and all rotations - produce overlaps, pushing it "above" or "below" the plane can resolve - the clash while maintaining bond lengths (connection sites rotate to - maintain the correct distance). - - Tries offsets in both directions at increasing magnitudes. - - Parameters - ---------- - ... (same as parent) + """Initial alignment for 2+ unique connection sites via Procrustes. Returns ------- - best_placed : np.ndarray, shape (n_sites, 3) - found_valid : bool + base_placed : np.ndarray + perm : list + rotation_type : str + 'axis' for 2 sites, 'none' for 3+ + rotation_info : dict """ - axis_hat = rotation_info["axis"] - pivot = rotation_info["pivot"] + n_unique = len(unique_active) - # Find a perpendicular direction to offset along - # Use the cross product of axis with a non-parallel vector - perp = np.array([1, 0, 0], dtype=np.float32) - if abs(np.dot(axis_hat, perp)) > 0.9: - perp = np.array([0, 1, 0], dtype=np.float32) - offset_dir = np.cross(axis_hat, perp) - offset_dir /= np.linalg.norm(offset_dir) + best_R = np.eye(3, dtype=np.float32) + best_t = np.zeros(3, dtype=np.float32) + best_cost = np.inf + best_perm = list(range(n_unique)) - # Try offsets at increasing distances - # Max offset: limited by how much bond length can stretch within tolerance - max_offset = np.mean(target_bond_lengths) * bond_tol * 3 - offsets = np.linspace(0.02 * max_offset, max_offset, 5) + for perm in permutations(range(n_unique)): + perm_targets = unique_targets[list(perm)] - best_placed = base_placed.copy() - best_min_dist = -np.inf - found_valid = False + src_centroid = np.mean(unique_conn_positions, axis=0) + tgt_centroid = np.mean(perm_targets, axis=0) - for sign in [1.0, -1.0]: - for offset_mag in offsets: - offset = offset_dir * offset_mag * sign - shifted = base_placed + offset - - # Re-sweep rotation at this offset - angles = np.linspace(0, 2 * np.pi, n_rotation_samples, endpoint=False) - shifted_pivot = pivot + offset - - for angle in angles: - R_rot = _rotation_about_axis(axis_hat, angle) - candidate = (R_rot @ (shifted - shifted_pivot).T).T + shifted_pivot - - bonds_ok = _check_bonds_valid( - candidate, - active_conn_sites, - neighbor_coords, - target_bond_lengths, - bond_tol, - ) - if not bonds_ok: - continue + src_centered = unique_conn_positions - src_centroid + tgt_centered = perm_targets - tgt_centroid - if nearby_coords is not None and len(nearby_coords) > 0: - min_dist = _min_distance_to_neighbors(candidate, nearby_coords) - else: - min_dist = np.inf + R = _kabsch_rotation(src_centered, tgt_centered) + t = tgt_centroid - R @ src_centroid - is_valid = min_dist >= min_separation + placed_conn = (R @ unique_conn_positions.T).T + t + cost = np.sum((placed_conn - perm_targets) ** 2) - if is_valid and (not found_valid or min_dist > best_min_dist): - best_placed = candidate.copy() - best_min_dist = min_dist - found_valid = True - elif not found_valid and min_dist > best_min_dist: - best_placed = candidate.copy() - best_min_dist = min_dist + if cost < best_cost: + best_cost = cost + best_R = R + best_t = t + best_perm = list(perm) - if found_valid: - break - if found_valid: - break + base_placed = (best_R @ replacement.coordinates.T).T + best_t - return best_placed, found_valid + if n_unique == 2: + # Residual DOF: rotation about axis connecting the two connection sites + conn0 = base_placed[unique_active[0]] + conn1 = base_placed[unique_active[1]] + axis = conn1 - conn0 + axis_norm = np.linalg.norm(axis) + if axis_norm > 1e-10: + axis_hat = axis / axis_norm + else: + axis_hat = np.array([0, 0, 1], dtype=np.float32) + midpoint = (conn0 + conn1) / 2.0 + return base_placed, best_perm, "axis", {"pivot": midpoint, "axis": axis_hat} + else: + # n_unique >= 3: fully determined, no residual DOF + return base_placed, best_perm, "none", {} def _check_bonds_valid( @@ -1597,15 +1039,6 @@ def _check_bonds_valid( return True -def _min_distance_to_neighbors(positions, nearby_coords): - """Minimum distance from any position to any nearby coord.""" - if len(nearby_coords) == 0: - return np.inf - diffs = positions[:, None, :] - nearby_coords[None, :, :] - dists = np.linalg.norm(diffs, axis=-1) - return float(np.min(dists)) - - def _build_assignment_from_perm( active_conn_sites, unique_active, conn_to_target_indices, best_perm, degree ): @@ -1659,120 +1092,6 @@ def _build_assignment_from_perm( return assignment -def _verify_and_fix_bonds( - placed_positions, - active_conn_sites, - assignment, - neighbor_coords, - target_bond_lengths, - tolerance, -): - """Verify connection-to-neighbor distances and nudge if outside tolerance. - - Moves each connection site along the line to its neighbor to achieve - exactly target_bond_length if it deviates too much. Then shifts the - entire rigid body to maintain internal geometry as much as possible. - - Parameters - ---------- - placed_positions : np.ndarray, shape (n_sites, 3) - active_conn_sites : list of int - assignment : list of int - neighbor_coords : np.ndarray, shape (degree, 3) - target_bond_lengths : np.ndarray, shape (degree,) - tolerance : float - - Returns - ------- - placed_positions : np.ndarray, shape (n_sites, 3) - Corrected positions. - """ - placed = placed_positions.copy() - degree = len(active_conn_sites) - - needs_correction = False - corrections = [] - - for conn_idx, neighbor_idx in enumerate(assignment): - site_local = active_conn_sites[conn_idx] - conn_pos = placed[site_local] - nb_pos = neighbor_coords[neighbor_idx] - target_bl = target_bond_lengths[conn_idx] - - delta = conn_pos - nb_pos - actual_bl = np.linalg.norm(delta) - - if actual_bl < 1e-10: - # Degenerate: connection site is on top of neighbor - needs_correction = True - corrections.append((site_local, neighbor_idx, conn_idx)) - continue - - error = abs(actual_bl - target_bl) - if error > target_bl * tolerance: - needs_correction = True - corrections.append((site_local, neighbor_idx, conn_idx)) - - if not needs_correction: - return placed - - # Strategy: compute ideal connection positions and shift entire rigid body - # to best match them (least-squares center shift) - ideal_conn_positions = np.zeros((degree, 3), dtype=np.float32) - for conn_idx, neighbor_idx in enumerate(assignment): - site_local = active_conn_sites[conn_idx] - nb_pos = neighbor_coords[neighbor_idx] - target_bl = target_bond_lengths[conn_idx] - - conn_pos = placed[site_local] - delta = conn_pos - nb_pos - d = np.linalg.norm(delta) - if d > 1e-10: - direction = delta / d - else: - # Pick direction from crosslinker center to this site - center = np.mean(placed, axis=0) - direction = placed[site_local] - center - d2 = np.linalg.norm(direction) - if d2 > 1e-10: - direction = direction / d2 - else: - direction = np.array([1, 0, 0], dtype=np.float32) - - ideal_conn_positions[conn_idx] = nb_pos + direction * target_bl - - # Compute offset: difference between where connection sites are and should be - current_conn_positions = np.array( - [placed[active_conn_sites[i]] for i in range(degree)], dtype=np.float32 - ) - offsets = ideal_conn_positions - current_conn_positions - mean_offset = np.mean(offsets, axis=0) - - # Apply rigid shift to all sites - placed += mean_offset - - # Final per-site correction for any remaining error - for conn_idx, neighbor_idx in enumerate(assignment): - site_local = active_conn_sites[conn_idx] - nb_pos = neighbor_coords[neighbor_idx] - target_bl = target_bond_lengths[conn_idx] - - delta = placed[site_local] - nb_pos - actual_bl = np.linalg.norm(delta) - - if actual_bl < 1e-10: - center = np.mean(placed, axis=0) - direction = placed[site_local] - center - d2 = np.linalg.norm(direction) - direction = direction / d2 if d2 > 1e-10 else np.array([1, 0, 0]) - placed[site_local] = nb_pos + direction * target_bl - elif abs(actual_bl - target_bl) > target_bl * tolerance: - direction = delta / actual_bl - placed[site_local] = nb_pos + direction * target_bl - - return placed - - def _get_pbc_info(volume_constraint): """Extract PBC flags and box lengths from volume constraint.""" if isinstance(volume_constraint, CuboidConstraint): @@ -1814,499 +1133,653 @@ def _wrap_positions(positions, pbc, box_lengths): return positions -def _get_nearby_nonbonded( - path, site, neighbors, site_pos, radius, pbc, box_lengths, excluded_sites -): - """Get coordinates of nearby non-bonded sites for overlap computation. +def _compute_ideal_backbone_distances(crosslinker, bond_length): + """Compute ideal pairwise distances between backbone targets given geometry and bond length. + + For a crosslinker with connection sites at positions c_i (relative to center), + the ideal backbone position for site i is: + b_i = c_i + bond_length * (c_i / |c_i|) + i.e., the backbone sits along the same radial direction, at distance bond_length + beyond the connection site. + + Parameters + ---------- + crosslinker : CrosslinkerGeometry + The crosslinker geometry. + bond_length : float + Desired bond length from connection site to backbone bead. - Returns positions relative to site_pos. + Returns + ------- + ideal_bb_positions : np.ndarray, shape (n_connections, 3) + Ideal backbone positions relative to crosslinker center. + pairwise_distances : np.ndarray, shape (n_connections, n_connections) + Matrix of ideal pairwise distances between backbone targets. + search_radius : float + Maximum distance from center to any ideal backbone position. """ - neighbors_set = set(neighbors) | {site} | excluded_sites + conn_sites = crosslinker.connection_sites + conn_positions = crosslinker.coordinates[conn_sites] - nearby = [] - for node in path.bond_graph.nodes(): - if node in neighbors_set: - continue - delta = path.coordinates[node] - site_pos - delta = _apply_mic(delta, pbc, box_lengths) - dist = np.linalg.norm(delta) - if dist <= radius: - nearby.append(delta) + # Compute ideal backbone positions + ideal_bb_positions = np.zeros_like(conn_positions) + for i in range(len(conn_sites)): + c_i = conn_positions[i] + c_norm = np.linalg.norm(c_i) + if c_norm > 1e-10: + # Backbone sits beyond connection site along radial direction + c_hat = c_i / c_norm + ideal_bb_positions[i] = c_i + bond_length * c_hat + else: + # Connection site at origin (single-site case): backbone at bond_length in any direction + ideal_bb_positions[i] = np.array([bond_length, 0, 0], dtype=np.float32) - if len(nearby) == 0: - return np.empty((0, 3), dtype=np.float32) - return np.array(nearby, dtype=np.float32) + # Maximum valid pairwise distances (geometric feasibility bound) + n = len(conn_sites) + pairwise_distances = np.zeros((n, n), dtype=np.float32) + for i in range(n): + for j in range(i + 1, n): + d_internal = np.linalg.norm(conn_positions[i] - conn_positions[j]) + max_d = d_internal + 2 * bond_length + pairwise_distances[i, j] = max_d + pairwise_distances[j, i] = max_d + # Search radius: max distance from center to any backbone + search_radius = np.max(np.linalg.norm(ideal_bb_positions, axis=1)) -def crosslink( - path, - crosslinker=None, - bead_name="_R", - backbone_name="_A", - crosslink_bond_length=0.2, - tolerance=0.1, - excluded_bond_depth=2, - n_connection_sites=2, - volume_constraint=None, - initial_point=None, - seed=42, - chunk_size=512, - run_on_gpu=False, - n_rotation_samples=36, - overlap_radius=None, - min_separation=None, -): - """ - Create a crosslink bonded to backbone beads, modifying path in place. + return ideal_bb_positions, pairwise_distances, search_radius - Selects backbone beads whose pairwise spacing matches the crosslinker - geometry, places a temporary placeholder, then calls ``replace_sites`` - which handles rigid-body alignment, bond length guarantees, and overlap - avoidance in a single pass. - Parameters - ---------- - path : Path - The Path object to modify in place. - crosslinker : CrosslinkerGeometry, optional - If None, uses single-site from ``bead_name`` and ``n_connection_sites``. - bead_name : str, default "_R" - backbone_name : str, default "_A" - crosslink_bond_length : float, default 0.2 - Desired bond length from connection sites to backbone beads (nm). - tolerance : float, default 0.1 - Fractional tolerance on bond lengths (±10%). - excluded_bond_depth : int, default 2 - n_connection_sites : int, default 2 - volume_constraint : optional - initial_point : int or array-like, optional - seed : int, default 42 - chunk_size : int, default 512 - run_on_gpu : bool, default False - n_rotation_samples : int, default 36 - overlap_radius : float, optional - min_separation : float, optional - Minimum distance from placed sites to existing sites. - Defaults to ``crosslink_bond_length * 0.5``. +def _unwrap_coords(coords, pbc, box_lengths): + """Unwrap coordinates relative to the first point.""" + unwrapped = np.empty_like(coords) + unwrapped[0] = coords[0] + ref = coords[0] + for k in range(1, len(coords)): + delta = coords[k] - ref + for dim in range(3): + if pbc[dim]: + delta[dim] -= np.round(delta[dim] / box_lengths[dim]) * box_lengths[dim] + unwrapped[k] = ref + delta + return unwrapped + + +def _validate_backbone_group( + selected_coords, + ideal_pairwise_distances, + bond_length, + tolerance, + pbc, + box_lengths, +): + """Check if backbone beads match expected geometry within tolerance. Returns ------- - Path - Same path, modified in place. + valid : bool + best_perm : list of int or None """ - if crosslinker is None: - crosslinker = CrosslinkerGeometry.single_site( - bead_name=bead_name, n_connections=n_connection_sites - ) - - if min_separation is None: - min_separation = crosslink_bond_length * 0.5 - - n_backbone_connections = crosslinker.n_connections - rng = np.random.default_rng(seed + len(path.coordinates)) - - # --- Compute search geometry --- - ideal_bb_positions, ideal_pairwise_distances, geom_search_radius = ( - _compute_ideal_backbone_distances(crosslinker, crosslink_bond_length) - ) - - if n_backbone_connections > 1: - max_pair_dist = np.max(ideal_pairwise_distances) - else: - max_pair_dist = 0.0 + n = len(selected_coords) + if n <= 1: + return True, [0] if n == 1 else [] - search_radius = max_pair_dist * (1 + tolerance) + crosslink_bond_length * tolerance + abs_tol = bond_length * tolerance - # --- Find backbone targets --- - backbone_nodes = [ - node for node in path.bond_graph.nodes() if path.beads[node] == backbone_name - ] - backbone_subgraph = path.bond_graph.subgraph(backbone_nodes) + # Compute actual pairwise distances (PBC-aware) + actual_distances = np.zeros((n, n), dtype=np.float32) + for i in range(n): + for j in range(i + 1, n): + delta = selected_coords[j] - selected_coords[i] + for dim in range(3): + if pbc[dim]: + delta[dim] -= ( + np.round(delta[dim] / box_lengths[dim]) * box_lengths[dim] + ) + d = np.linalg.norm(delta) + actual_distances[i, j] = d + actual_distances[j, i] = d - if len(backbone_nodes) == 0: - raise ValueError(f"No backbone beads with name '{backbone_name}' found in path") - if len(backbone_nodes) < n_backbone_connections: - raise ValueError( - f"Not enough backbone beads ({len(backbone_nodes)}) for " - f"{n_backbone_connections} connection sites" - ) + best_perm = None + best_cost = np.inf - pbc, box_lengths = _get_pbc_info(volume_constraint) + for perm in permutations(range(n)): + valid = True + cost = 0.0 + for i in range(n): + for j in range(i + 1, n): + max_d = ideal_pairwise_distances[i, j] * (1 + tolerance) + actual_d = actual_distances[perm[i], perm[j]] + if actual_d > max_d: + valid = False + break + cost += actual_d + if not valid: + break - candidate_nodes = [ - node for node in backbone_nodes if path.bond_graph.degree[node] <= 2 - ] - candidate_coords = np.array(path.coordinates[candidate_nodes], dtype=np.float32) + if valid and cost < best_cost: + best_cost = cost + best_perm = list(perm) - def get_reference_points(path, initial_point): - if initial_point is not None: - if isinstance(initial_point, (int, np.integer)): - if initial_point not in path.bond_graph.nodes: - raise ValueError(f"Node {initial_point} not found in bond_graph") - return np.array([initial_point]), np.array( - [path.coordinates[initial_point]] - ) - else: - initial_point32 = np.asarray(initial_point, dtype=np.float32) - sq_distances = calculate_sq_distances( - initial_point32, candidate_coords, pbc=pbc, box_lengths=box_lengths - ) - order = np.argsort(sq_distances) - return ( - np.array(candidate_nodes)[order], - path.coordinates[np.array(candidate_nodes)[order]], - ) - else: - shuffled = rng.choice( - candidate_nodes, size=len(candidate_nodes), replace=False - ) - return shuffled, path.coordinates[shuffled] + return best_perm is not None, best_perm - ref_nodes, ref_coords = get_reference_points(path, initial_point) - found_ref = False - selected_nodes = [] +def _validate_clinker( + path, + crosslinker, + crosslink_bond_length, + selected_nodes, + overlap_radius=0.1, + n_samples=200, + pbc=None, + box_lengths=None, + seed=None, + tolerance=1e-5, +): + """Verify a point exists within constraints to place the centroid of the crosslinker. - for ref_node, ref_coord in zip(ref_nodes, ref_coords): - selected_nodes = [ref_node] + Treats the crosslinker as a cylindrical bead whose radius is the maximum + distance from its centroid to any constituent site. Samples candidate + center positions within the geometric feasibility region (intersection of + spheres around selected backbone nodes) and uses ``check_path`` to verify + no hard-sphere overlaps exist with the rest of the path. - sq_distances = calculate_sq_distances( - ref_coord, candidate_coords, pbc=pbc, box_lengths=box_lengths - ) - distances = np.sqrt(sq_distances) + Parameters + ---------- + path : Path + The current path with all existing coordinates. + crosslinker : CrosslinkerGeometry + The crosslinker geometry to be placed. + crosslink_bond_length : float + Desired bond length from crosslinker connection sites to backbone beads. + selected_nodes : list of int + Backbone node indices that the crosslinker would connect to. + overlap_radius : float, default 0.1 + Additional exclusion radius beyond the crosslinker's own bead radius. + The total check radius passed to ``check_path`` is + ``bead_radius + overlap_radius``. + n_samples : int, default 200 + Number of candidate center positions to test. + pbc : np.ndarray of bool or None + Periodic boundary conditions per axis. If None, assumes no PBC. + box_lengths : np.ndarray or None + Box dimensions for minimum image convention. + seed : int or None + Random seed for reproducibility. + tolerance : float, default 1e-5 + Tolerance passed to ``check_path`` for rounding in distance checks. - within_radius_mask = distances <= search_radius - possible_pairs = np.where(within_radius_mask)[0] + Returns + ------- + valid : bool + True if at least one overlap-free center position exists. + best_center : np.ndarray, shape (3,) or None + The first valid candidate center position found, or the best + (maximum clearance) candidate if none fully pass. + """ + from mbuild.path.path_utils import check_path - if len(possible_pairs) < n_backbone_connections - 1: - continue + if pbc is None: + pbc = np.array([False, False, False]) + if box_lengths is None: + box_lengths = np.array([0.0, 0.0, 0.0], dtype=np.float32) - closest_paired_nodes = possible_pairs[np.argsort(distances[possible_pairs])] - excluded_nodes = set( - nx.single_source_shortest_path_length( - backbone_subgraph, ref_node, cutoff=excluded_bond_depth - ).keys() - ) + rng = np.random.default_rng(seed) - for idx in closest_paired_nodes: - node = candidate_nodes[idx] - if node in excluded_nodes: - continue + # --- Compute crosslinker bead radius (cylindrical approximation) --- + crosslinker_centroid = np.mean(crosslinker.coordinates, axis=0) + site_offsets = crosslinker.coordinates - crosslinker_centroid + bead_radius = float(np.max(np.linalg.norm(site_offsets, axis=1))) - is_excluded = False - for selected in selected_nodes[1:]: - if selected in backbone_subgraph: - sel_excluded = set( - nx.single_source_shortest_path_length( - backbone_subgraph, selected, cutoff=excluded_bond_depth - ).keys() + # --- Max distance from centroid to any connection site --- + conn_sites = crosslinker.connection_sites + conn_positions = crosslinker.coordinates[conn_sites] + conn_offsets = conn_positions - crosslinker_centroid + max_conn_radius = float(np.max(np.linalg.norm(conn_offsets, axis=1))) + + # --- Feasibility: centroid must be within this distance of each backbone node --- + feasibility_radius = crosslink_bond_length + max_conn_radius + + # --- Total hard-sphere check radius --- + check_radius = bead_radius + overlap_radius + + # --- Build existing_points array (all path coords except selected nodes) --- + n_total = len(path.coordinates) + selected_set = set(selected_nodes) + mask = np.ones(n_total, dtype=bool) + for node in selected_set: + if 0 <= node < n_total: + mask[node] = False + existing_points = np.ascontiguousarray(path.coordinates[mask], dtype=np.float32) + # existing_points = path.coordinates + + # --- Unwrap selected node coordinates relative to first --- + selected_coords = np.array( + [path.coordinates[n] for n in selected_nodes], dtype=np.float32 + ) + if len(selected_coords) > 1: + ref = selected_coords[0].copy() + for k in range(1, len(selected_coords)): + delta = selected_coords[k] - ref + for dim in range(3): + if pbc[dim]: + delta[dim] -= ( + np.round(delta[dim] / box_lengths[dim]) * box_lengths[dim] ) - if node in sel_excluded: - is_excluded = True - break - if is_excluded: - continue - - selected_nodes.append(int(node)) - if len(selected_nodes) >= n_backbone_connections: - # Validate pairwise geometry - sel_coords = np.array( - [path.coordinates[n] for n in selected_nodes], dtype=np.float32 - ) - unwrapped_check = _unwrap_coords(sel_coords, pbc, box_lengths) - valid, _ = _validate_backbone_group( - unwrapped_check, - ideal_pairwise_distances, - crosslink_bond_length, - tolerance, - pbc, - box_lengths, - ) - if valid: - found_ref = True - break - else: - selected_nodes.pop() - - if found_ref: - break + selected_coords[k] = ref + delta - if not found_ref: - n_clinks = sum(1 for b in path.beads if b in set(crosslinker.beads)) - raise PathConvergenceError( - f"Could not find {n_backbone_connections} non-neighboring backbone beads " - f"matching crosslinker geometry with bond_length={crosslink_bond_length} " - f"(tolerance=±{tolerance * 100:.0f}%)." - f"\nCurrent crosslinks: {n_clinks}. " - "Ways to increase crosslinking:\n" - " - Increase crosslink_bond_length\n" - " - Increase tolerance\n" - " - Decrease excluded_bond_depth\n" - " - Pack at higher density\n" - " - Relax structure" - ) + centroid_of_selected = np.mean(selected_coords, axis=0) - # --- Place placeholder at centroid --- - selected_coords = np.array(path.coordinates[selected_nodes], dtype=np.float32) - unwrapped = _unwrap_coords(selected_coords, pbc, box_lengths) - crosslink_center = np.mean(unwrapped, axis=0) + # --- Compute sampling sphere radius --- + max_dist_to_centroid = float( + np.max(np.linalg.norm(selected_coords - centroid_of_selected, axis=1)) + ) + sample_radius = feasibility_radius - max_dist_to_centroid + if sample_radius <= 0: + sample_radius = feasibility_radius * 0.05 + + # --- Generate candidate center positions --- + candidates = np.zeros((n_samples, 3), dtype=np.float32) + candidates[0] = centroid_of_selected + for i in range(1, n_samples): + r = sample_radius * (rng.random() ** (1.0 / 3.0)) + direction = rng.standard_normal(3).astype(np.float32) + norm = np.linalg.norm(direction) + if norm > 1e-10: + direction /= norm + else: + direction = np.array([1.0, 0.0, 0.0], dtype=np.float32) + candidates[i] = centroid_of_selected + r * direction - for dim in range(3): - if pbc[dim]: - crosslink_center[dim] -= ( - np.round(crosslink_center[dim] / box_lengths[dim]) * box_lengths[dim] - ) + # --- Evaluate each candidate --- + best_center = None - _placeholder_name = "__placeholder__" - path.append_coordinates(crosslink_center, _placeholder_name) - placeholder_idx = len(path.coordinates) - 1 + for cand in candidates: + # Check feasibility: centroid within reach of all selected backbone nodes + feasible = True + for sc in selected_coords: + delta = cand - sc + for dim in range(3): + if pbc[dim]: + delta[dim] -= ( + np.round(delta[dim] / box_lengths[dim]) * box_lengths[dim] + ) + dist = np.linalg.norm(delta) + if dist > feasibility_radius: + feasible = False + break + if not feasible: + continue - for backbone_node in selected_nodes: - path.bond_graph.add_edge( - int(placeholder_idx), - int(backbone_node), - bond_type=(_placeholder_name, backbone_name), + # Use check_path for hard-sphere overlap detection + new_point = np.ascontiguousarray(cand, dtype=np.float32) + no_overlap = check_path( + existing_points=existing_points, + new_point=new_point, + radius=check_radius, + tolerance=tolerance, ) - # --- Single-site shortcut --- - if crosslinker.n_sites == 1: - path.beads[placeholder_idx] = crosslinker.beads[0] - for backbone_node in selected_nodes: - path.bond_graph.edges[placeholder_idx, backbone_node]["bond_type"] = ( - str(crosslinker.beads[0]), - backbone_name, - ) - if "name" in path.bond_graph.nodes[placeholder_idx]: - path.bond_graph.nodes[placeholder_idx]["name"] = str(crosslinker.beads[0]) - - # Fix single-site bond length if needed - for backbone_node in selected_nodes: - delta = path.coordinates[placeholder_idx] - path.coordinates[backbone_node] - delta = _apply_mic(delta, pbc, box_lengths) - actual_bl = np.linalg.norm(delta) - if ( - abs(actual_bl - crosslink_bond_length) - > crosslink_bond_length * tolerance - ): - if actual_bl > 1e-10: - direction = delta / actual_bl - else: - direction = rng.standard_normal(3).astype(np.float32) - direction /= np.linalg.norm(direction) - path.coordinates[placeholder_idx] = path.coordinates[ - backbone_node - ] + _apply_mic(direction * crosslink_bond_length, pbc, box_lengths) - return path - - # --- Multi-site: replace_sites handles everything --- - replace_sites( - path, - replacement=crosslinker, - sites=[placeholder_idx], - bond_length=crosslink_bond_length, - tolerance=tolerance, - min_separation=min_separation, - volume_constraint=volume_constraint, - n_rotation_samples=n_rotation_samples, - overlap_radius=overlap_radius, - seed=seed, - ) + if no_overlap: + return True, cand - return path + # Track first feasible candidate as fallback + if best_center is None: + best_center = cand.copy() + return False, best_center -def _compute_ideal_backbone_distances(crosslinker, bond_length): - """Compute ideal pairwise distances between backbone targets given geometry and bond length. - For a crosslinker with connection sites at positions c_i (relative to center), - the ideal backbone position for site i is: - b_i = c_i + bond_length * (c_i / |c_i|) - i.e., the backbone sits along the same radial direction, at distance bond_length - beyond the connection site. +def _place_replacement_rigid_full( + replacement, + active_conn_sites, + neighbor_coords, + target_bond_lengths, + all_path_coords, + bonded_indices, + pbc, + box_lengths, + rng, + n_rotation_samples, + tolerance, + min_separation=None, +): + """Place replacement with full-system overlap checking. + + Unlike the radius-limited version, this checks placed sites against ALL + existing coordinates in the path (minus bonded neighbors). Parameters ---------- - crosslinker : CrosslinkerGeometry - The crosslinker geometry. - bond_length : float - Desired bond length from connection site to backbone bead. + replacement : CrosslinkerGeometry + active_conn_sites : list of int + neighbor_coords : np.ndarray, shape (degree, 3) + target_bond_lengths : np.ndarray, shape (degree,) + all_path_coords : np.ndarray, shape (N, 3) + ALL coordinates in the current path. + bonded_indices : set of int + Indices to exclude from overlap checking. + pbc : array-like of bool + box_lengths : np.ndarray + rng : numpy.random.Generator + n_rotation_samples : int + tolerance : float + min_separation : float or None Returns ------- - ideal_bb_positions : np.ndarray, shape (n_connections, 3) - Ideal backbone positions relative to crosslinker center. - pairwise_distances : np.ndarray, shape (n_connections, n_connections) - Matrix of ideal pairwise distances between backbone targets. - search_radius : float - Maximum distance from center to any ideal backbone position. + placed_positions : np.ndarray, shape (n_sites, 3) + assignment : list of int """ - conn_sites = crosslinker.connection_sites - conn_positions = crosslinker.coordinates[conn_sites] - - # Compute ideal backbone positions - ideal_bb_positions = np.zeros_like(conn_positions) - for i in range(len(conn_sites)): - c_i = conn_positions[i] - c_norm = np.linalg.norm(c_i) - if c_norm > 1e-10: - # Backbone sits beyond connection site along radial direction - c_hat = c_i / c_norm - ideal_bb_positions[i] = c_i + bond_length * c_hat - else: - # Connection site at origin (single-site case): backbone at bond_length in any direction - ideal_bb_positions[i] = np.array([bond_length, 0, 0], dtype=np.float32) - - # Pairwise distances - n = len(conn_sites) - pairwise_distances = np.zeros((n, n), dtype=np.float32) - for i in range(n): - for j in range(i + 1, n): - d = np.linalg.norm(ideal_bb_positions[i] - ideal_bb_positions[j]) - pairwise_distances[i, j] = d - pairwise_distances[j, i] = d + degree = len(active_conn_sites) - # Search radius: max distance from center to any backbone - search_radius = np.max(np.linalg.norm(ideal_bb_positions, axis=1)) + if degree == 0: + center = ( + np.mean(neighbor_coords, axis=0) + if len(neighbor_coords) > 0 + else np.zeros(3) + ) + R = _random_rotation_matrix(rng) + placed = (R @ replacement.coordinates.T).T + center + return placed.astype(np.float32), [] - return ideal_bb_positions, pairwise_distances, search_radius + # --- Compute target positions --- + neighbor_centroid = np.mean(neighbor_coords, axis=0) + target_positions = np.zeros((degree, 3), dtype=np.float32) + for i in range(degree): + direction = neighbor_centroid - neighbor_coords[i] + d = np.linalg.norm(direction) + if d > 1e-10: + direction_hat = direction / d + else: + direction_hat = rng.standard_normal(3).astype(np.float32) + direction_hat /= np.linalg.norm(direction_hat) + target_positions[i] = ( + neighbor_coords[i] + direction_hat * target_bond_lengths[i] + ) + # --- Initial alignment --- + unique_active = list(dict.fromkeys(active_conn_sites)) + unique_conn_positions = replacement.coordinates[unique_active] + n_unique = len(unique_active) -def _unwrap_coords(coords, pbc, box_lengths): - """Unwrap coordinates relative to the first point.""" - unwrapped = np.empty_like(coords) - unwrapped[0] = coords[0] - ref = coords[0] - for k in range(1, len(coords)): - delta = coords[k] - ref - for dim in range(3): - if pbc[dim]: - delta[dim] -= np.round(delta[dim] / box_lengths[dim]) * box_lengths[dim] - unwrapped[k] = ref + delta - return unwrapped + conn_to_target_indices = {} + for i, site_idx in enumerate(active_conn_sites): + conn_to_target_indices.setdefault(site_idx, []).append(i) + unique_targets = np.array( + [ + np.mean(target_positions[conn_to_target_indices[s]], axis=0) + for s in unique_active + ], + dtype=np.float32, + ) -def _compute_optimal_center( - crosslinker, backbone_coords, bond_length, pbc, box_lengths -): - """Compute the optimal crosslinker center so connection sites are at bond_length from backbones. + if n_unique == 1: + base_placed, base_perm, rotation_type, rotation_info = _initial_align_single( + replacement, + unique_active, + unique_conn_positions, + unique_targets, + neighbor_centroid, + rng, + ) + elif n_unique >= 2: + base_placed, base_perm, rotation_type, rotation_info = _initial_align_multi( + replacement, + unique_active, + unique_conn_positions, + unique_targets, + rng, + ) + else: + R = _random_rotation_matrix(rng) + placed = (R @ replacement.coordinates.T).T + neighbor_centroid + assignment = _build_assignment_from_perm( + active_conn_sites, unique_active, conn_to_target_indices, [0], degree + ) + return placed.astype(np.float32), assignment + + # --- Precompute check coords (all minus bonded) --- + n_existing = len(all_path_coords) + check_mask = np.ones(n_existing, dtype=bool) + for idx in bonded_indices: + if 0 <= idx < n_existing: + check_mask[idx] = False + check_coords = all_path_coords[check_mask] + + # --- Sweep rotations with full overlap check --- + best_placed, best_min_dist, found_valid = _sweep_rotations_full( + base_placed, + rotation_type, + rotation_info, + active_conn_sites, + neighbor_coords, + target_bond_lengths, + check_coords, + pbc, + box_lengths, + tolerance, + min_separation, + n_rotation_samples, + rng, + ) - The center is placed such that each connection site (after rotation) sits - at exactly bond_length from its target backbone bead. + # --- If not valid, try offsets --- + if not found_valid and min_separation is not None and rotation_type == "axis": + best_placed, found_valid = _try_offsets_full( + base_placed, + rotation_info, + replacement, + active_conn_sites, + neighbor_coords, + target_bond_lengths, + check_coords, + pbc, + box_lengths, + tolerance, + min_separation, + n_rotation_samples, + rng, + ) - For the ideal placement: - backbone_i = center + c_i_hat * (|c_i| + bond_length) + # Build assignment + assignment = _build_assignment_from_perm( + active_conn_sites, unique_active, conn_to_target_indices, base_perm, degree + ) - So: center = backbone_i - c_i_hat * (|c_i| + bond_length) + return best_placed.astype(np.float32), assignment - We average over all connection sites for robustness. - Parameters - ---------- - crosslinker : CrosslinkerGeometry - The crosslinker. - backbone_coords : np.ndarray, shape (n_connections, 3) - Unwrapped backbone coordinates. - bond_length : float - Desired bond length. - pbc, box_lengths : PBC info. +def _sweep_rotations_full( + base_placed, + rotation_type, + rotation_info, + active_conn_sites, + neighbor_coords, + target_bond_lengths, + check_coords, + pbc, + box_lengths, + tolerance, + min_separation, + n_rotation_samples, + rng, +): + """Sweep rotations checking against full coordinate set. Returns ------- - center : np.ndarray, shape (3,) + best_placed : np.ndarray + best_min_dist : float + found_valid : bool """ - conn_sites = crosslinker.connection_sites - conn_positions = crosslinker.coordinates[conn_sites] - n = len(conn_sites) + best_placed = base_placed.copy() + best_min_dist = _min_dist_pbc(base_placed, check_coords, pbc, box_lengths) + found_valid = min_separation is None or best_min_dist >= min_separation + + # Check if base placement bonds are valid + if not _check_bonds_valid( + base_placed, active_conn_sites, neighbor_coords, target_bond_lengths, tolerance + ): + best_min_dist = -np.inf + found_valid = False - if n == 0: - return np.mean(backbone_coords, axis=0) + if rotation_type == "axis": + pivot = rotation_info["pivot"] + axis_hat = rotation_info["axis"] + angles = np.linspace(0, 2 * np.pi, n_rotation_samples, endpoint=False) - # For single-site crosslinker (all connection sites at same position): - if np.allclose(conn_positions, conn_positions[0], atol=1e-8): - # Center is at bond_length distance from centroid toward... just use centroid - return np.mean(backbone_coords, axis=0) + for angle in angles: + R_rot = _rotation_about_axis(axis_hat, angle) + candidate = (R_rot @ (base_placed - pivot).T).T + pivot + + if not _check_bonds_valid( + candidate, + active_conn_sites, + neighbor_coords, + target_bond_lengths, + tolerance, + ): + continue - # Compute directions from crosslinker center to each connection site - conn_norms = np.linalg.norm(conn_positions, axis=1, keepdims=True) - conn_hat = conn_positions / np.maximum(conn_norms, 1e-10) + min_dist = _min_dist_pbc(candidate, check_coords, pbc, box_lengths) - # Ideal: backbone_i is at direction c_i_hat, distance |c_i| + bond_length from center - # So center = backbone_i - c_i_hat * (|c_i| + bond_length) - # But we haven't rotated the crosslinker yet, so we need to figure out the rotation. + if min_separation is not None and min_dist >= min_separation: + if not found_valid or min_dist > best_min_dist: + best_placed = candidate.copy() + best_min_dist = min_dist + found_valid = True + elif not found_valid and min_dist > best_min_dist: + best_placed = candidate.copy() + best_min_dist = min_dist - # Target vectors: from centroid of backbones to each backbone - bb_centroid = np.mean(backbone_coords, axis=0) - target_vectors = backbone_coords - bb_centroid - target_norms = np.linalg.norm(target_vectors, axis=1, keepdims=True) - target_hat = target_vectors / np.maximum(target_norms, 1e-10) + elif rotation_type == "full": + pivot = rotation_info.get("pivot", np.mean(base_placed, axis=0)) + for _ in range(n_rotation_samples): + R = _random_rotation_matrix(rng) + candidate = (R @ (base_placed - pivot).T).T + pivot + + if not _check_bonds_valid( + candidate, + active_conn_sites, + neighbor_coords, + target_bond_lengths, + tolerance, + ): + continue - # The ideal distance from center to each backbone: - # |c_i| + bond_length - ideal_radii = conn_norms.flatten() + bond_length + min_dist = _min_dist_pbc(candidate, check_coords, pbc, box_lengths) - # Center = backbone_i - ideal_radius_i * target_hat_i (if we orient along target directions) - # Average: - centers = backbone_coords - ideal_radii[:, None] * target_hat - center = np.mean(centers, axis=0) + if min_separation is not None and min_dist >= min_separation: + if not found_valid or min_dist > best_min_dist: + best_placed = candidate.copy() + best_min_dist = min_dist + found_valid = True + elif not found_valid and min_dist > best_min_dist: + best_placed = candidate.copy() + best_min_dist = min_dist - return center.astype(np.float32) + return best_placed, best_min_dist, found_valid -def _validate_backbone_group( - selected_coords, - ideal_pairwise_distances, - bond_length, - tolerance, +def _try_offsets_full( + base_placed, + rotation_info, + replacement, + active_conn_sites, + neighbor_coords, + target_bond_lengths, + check_coords, pbc, box_lengths, + tolerance, + min_separation, + n_rotation_samples, + rng, ): - """Check if backbone beads match expected geometry within tolerance. + """Try offset placements with full overlap checking.""" + axis_hat = rotation_info["axis"] + pivot = rotation_info["pivot"] - Returns - ------- - valid : bool - best_perm : list of int or None - """ - n = len(selected_coords) - if n <= 1: - return True, [0] if n == 1 else [] + perp1 = _get_perpendicular(axis_hat) + perp2 = np.cross(axis_hat, perp1).astype(np.float32) + perp2 /= np.linalg.norm(perp2) - abs_tol = bond_length * tolerance + max_extent = float(np.max(np.linalg.norm(replacement.coordinates, axis=1))) + offset_magnitudes = np.linspace(0.05, max_extent * 1.5, 6) - # Compute actual pairwise distances (PBC-aware) - actual_distances = np.zeros((n, n), dtype=np.float32) - for i in range(n): - for j in range(i + 1, n): - delta = selected_coords[j] - selected_coords[i] - for dim in range(3): - if pbc[dim]: - delta[dim] -= ( - np.round(delta[dim] / box_lengths[dim]) * box_lengths[dim] - ) - d = np.linalg.norm(delta) - actual_distances[i, j] = d - actual_distances[j, i] = d + best_placed = base_placed.copy() + best_min_dist = _min_dist_pbc(base_placed, check_coords, pbc, box_lengths) + found_valid = False - best_perm = None - best_cost = np.inf + for offset_mag in offset_magnitudes: + for sign in [1.0, -1.0]: + for perp_dir in [perp1, perp2]: + offset = perp_dir * offset_mag * sign + + # Shift the whole placement + shifted = base_placed + offset + shifted_pivot = pivot + offset + + # Sweep rotations at this offset + angles = np.linspace(0, 2 * np.pi, n_rotation_samples, endpoint=False) + for angle in angles: + R_rot = _rotation_about_axis(axis_hat, angle) + candidate = (R_rot @ (shifted - shifted_pivot).T).T + shifted_pivot + + if not _check_bonds_valid( + candidate, + active_conn_sites, + neighbor_coords, + target_bond_lengths, + tolerance, + ): + continue + + min_dist = _min_dist_pbc(candidate, check_coords, pbc, box_lengths) + + if min_dist >= min_separation: + if not found_valid or min_dist > best_min_dist: + best_placed = candidate.copy() + best_min_dist = min_dist + found_valid = True + elif not found_valid and min_dist > best_min_dist: + best_placed = candidate.copy() + best_min_dist = min_dist + + if found_valid: + return best_placed, found_valid - for perm in permutations(range(n)): - valid = True - cost = 0.0 - for i in range(n): - for j in range(i + 1, n): - ideal_d = ideal_pairwise_distances[i, j] - actual_d = actual_distances[perm[i], perm[j]] - error = abs(actual_d - ideal_d) - if error > abs_tol: - valid = False - break - cost += error - if not valid: - break + return best_placed, found_valid - if valid and cost < best_cost: - best_cost = cost - best_perm = list(perm) - return best_perm is not None, best_perm +def _min_dist_pbc(positions, check_coords, pbc, box_lengths): + """Minimum distance from any position to any check coord, with PBC.""" + if len(check_coords) == 0: + return np.inf + + global_min_sq = np.inf + for site_pos in positions: + deltas = check_coords - site_pos + for dim in range(3): + if pbc[dim]: + deltas[:, dim] -= ( + np.round(deltas[:, dim] / box_lengths[dim]) * box_lengths[dim] + ) + sq_dists = np.sum(deltas * deltas, axis=1) + local_min = float(np.min(sq_dists)) + if local_min < global_min_sq: + global_min_sq = local_min + + return float(np.sqrt(global_min_sq)) + + +def _get_perpendicular(axis): + """Get a unit vector perpendicular to axis.""" + axis = np.asarray(axis, dtype=np.float32) + perp = np.array([1, 0, 0], dtype=np.float32) + if abs(np.dot(axis, perp)) > 0.9: + perp = np.array([0, 1, 0], dtype=np.float32) + perp = perp - np.dot(perp, axis) * axis + perp /= np.linalg.norm(perp) + return perp diff --git a/mbuild/tests/test_crosslinks.py b/mbuild/tests/test_crosslinks.py index d350402c3..f4288e7a8 100644 --- a/mbuild/tests/test_crosslinks.py +++ b/mbuild/tests/test_crosslinks.py @@ -1,3 +1,4 @@ +import networkx as nx import numpy as np import pytest @@ -294,10 +295,158 @@ def test_error_on_excess_degree(self, long_linear_path): triangle = CrosslinkerGeometry.equilateral_triangle( bond_length=0.27, connection_sites=[0, 1], bead_name="_U" ) - crosslink(long_linear_path, triangle, crosslink_bond_length=1) + crosslink(long_linear_path, triangle, crosslink_bond_length=5) with pytest.raises(ValueError, match="degree"): replace_sites(long_linear_path, triangle, bead_name="_A") + def test_non_centroid_candidate(self, long_linear_path): + triangle = CrosslinkerGeometry.equilateral_triangle( + bond_length=0.27, connection_sites=[0, 1], bead_name="_U" + ) + crosslink(long_linear_path, triangle, crosslink_bond_length=5) + assert long_linear_path + + def test_high_density_crosslinking_overlaps_linear(self): + # Make a gridded bond graph + X, Y, Z = np.meshgrid(np.arange(5), np.arange(5), np.arange(5), indexing="ij") + coords = np.column_stack((X.ravel(), Y.ravel(), Z.ravel())).astype(np.float32) + path = Path(coordinates=coords) + G = path.bond_graph + for i in range(len(coords)): + for j in range(i + 1, len(coords)): + distance = np.linalg.norm(coords[i] - coords[j]) + if distance < 1.01: + G.add_edge(i, j) + + bond_length = np.sqrt(3) / 2 # half the diagonal bond length + clink_coords = np.array([[0, 0, 0], [0, 0, bond_length / 2]], dtype=np.float32) + clinkG = nx.Graph() + clinkG.add_nodes_from([0, 1]) + clinkG.add_edge(0, 1) + clink = CrosslinkerGeometry( + clink_coords, clinkG, bead_name="_CR", connection_sites=[0, 0] + ) + min_sep = 0.99 * (bond_length / 2) + for i in range(4 * 4 * 4): + try: + crosslink( + path, + clink, + excluded_bond_depth=2, + min_separation=min_sep, + crosslink_bond_length=bond_length, + tolerance=0.01, + max_backbone_degree=8, + ) + except PathConvergenceError: + raise ValueError(f"Failed after {i} attempts to crosslink") + + # Find the non-connection crosslinker sites (site index 1 in the geometry) + # These are "_CR" beads that are NOT the connection site (site 0) + clink_nodes = [n for n in path.bond_graph.nodes() if path.beads[n] == "_CR"] + # The connection site (index 0) has degree 2+1=3 (two backbone + one internal) + # The internal site (index 1) has degree 1 (just the internal bond) + internal_nodes = [n for n in clink_nodes if path.bond_graph.degree[n] == 1] + + assert len(internal_nodes) > 0, "No internal crosslinker sites found" + + # For each internal node, compute min distance to all non-bonded nodes + for node in internal_nodes: + bonded = set(path.bond_graph.neighbors(node)) | {node} + node_pos = path.coordinates[node] + + other_positions = np.array( + [ + path.coordinates[i] + for i in path.bond_graph.nodes() + if i not in bonded + ], + dtype=np.float32, + ) + + distances = np.linalg.norm(other_positions - node_pos, axis=1) + min_dist = float(np.min(distances)) + + assert min_dist >= min_sep, ( + f"Internal crosslinker node {node} has overlap: " + f"min distance {min_dist:.4f} < min_separation {min_sep}" + ) + + def test_high_density_crosslinking_overlaps_eqtriangle(self): + # Make a gridded bond graph + X, Y, Z = np.meshgrid(np.arange(5), np.arange(5), np.arange(5), indexing="ij") + coords = np.column_stack((X.ravel(), Y.ravel(), Z.ravel())).astype(np.float32) + path = Path(coordinates=coords) + G = path.bond_graph + for i in range(len(coords)): + for j in range(i + 1, len(coords)): + distance = np.linalg.norm(coords[i] - coords[j]) + if distance < 1.01: + G.add_edge(i, j) + + bond_length = np.sqrt(3) / 3 # 1/3 the diagonal bond length + clink = CrosslinkerGeometry.equilateral_triangle( + bond_length=bond_length, connection_sites=[0, 1] + ) + min_sep = 0.1 + for i in range(4 * 4 * 4): + crosslink( + path, + clink, + excluded_bond_depth=2, + min_separation=min_sep, + crosslink_bond_length=bond_length, + tolerance=0.05, + max_backbone_degree=8, + ) + + # Find the non-connection crosslinker sites (site index 1 in the geometry) + # These are "_CR" beads that are NOT the connection site (site 0) + clink_nodes = [n for n in path.bond_graph.nodes() if path.beads[n] == "_R"] + # The connection site (index 0) has degree 2+1=3 (two backbone + one internal) + # The internal site (index 1) has degree 1 (just the internal bond) + internal_nodes = [n for n in clink_nodes if path.bond_graph.degree[n] == 2] + + assert len(internal_nodes) > 0, "No internal crosslinker sites found" + + # For each internal node, compute min distance to all non-bonded nodes + for node in internal_nodes: + bonded = set(path.bond_graph.neighbors(node)) | {node} + node_pos = path.coordinates[node] + + other_positions = np.array( + [ + path.coordinates[i] + for i in path.bond_graph.nodes() + if i not in bonded + ], + dtype=np.float32, + ) + + distances = np.linalg.norm(other_positions - node_pos, axis=1) + min_dist = float(np.min(distances)) + + assert min_dist >= min_sep, ( + f"Internal crosslinker node {node} has overlap: " + f"min distance {min_dist:.4f} < min_separation {min_sep}" + ) + + def test_failing_overlap_radius(self, long_linear_path): + clink_coords = np.array([[0, 0, 0], [0, 0, 0.5]], dtype=np.float32) + clinkG = nx.Graph() + clinkG.add_nodes_from([0, 1]) + clinkG.add_edge(0, 1) + linear_clink = CrosslinkerGeometry( + clink_coords, clinkG, connection_sites=[0, 1] + ) + with pytest.raises(PathConvergenceError): + crosslink( + long_linear_path, + linear_clink, + crosslink_bond_length=0.5, + min_separation=0.2, + ) + class TestReplaceSites(BaseTest): """Tests for the replace_sites function.""" From 6e53cfe528bfe447705acb5c156c95307d7423d0 Mon Sep 17 00:00:00 2001 From: chrisjonesBSU Date: Wed, 27 May 2026 15:29:39 +0100 Subject: [PATCH 05/11] Check for overlaps when passing in new point --- mbuild/path/points.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/mbuild/path/points.py b/mbuild/path/points.py index f2bf04b1b..33e2f897e 100644 --- a/mbuild/path/points.py +++ b/mbuild/path/points.py @@ -137,11 +137,22 @@ def get_initial_point(state, existing_points, beads, check_path, next_step): if state.include_compound: existing_points = np.concat((existing_points, state.include_compound.xyz)) - # An initial point was manuallyl given in hard_sphere_random_walk, use that. + # An initial point was manually given in hard_sphere_random_walk, use that. + # Check if this point causes any overlaps, if so, raise error. if isinstance(state.initial_point, np.ndarray) and state.initial_point.shape == ( 3, ): - return state.initial_point + if check_path( + existing_points=existing_points, + new_point=state.initial_point, + radius=state.radius, + tolerance=state.tolerance, + ): + return state.initial_point + raise PathConvergenceError( + f"The provided initial_point {state.initial_point} overlaps with " + "existing particles. Try a different starting point." + ) # Passing in an index to specify an initial point from already defined set of coordinates elif isinstance(state.initial_point, int): From 3a0d3a4d217a0442358596d6020e923f1b30a1d4 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 27 May 2026 14:29:58 +0000 Subject: [PATCH 06/11] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- mbuild/path/points.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mbuild/path/points.py b/mbuild/path/points.py index 33e2f897e..cbc8f1cad 100644 --- a/mbuild/path/points.py +++ b/mbuild/path/points.py @@ -138,7 +138,7 @@ def get_initial_point(state, existing_points, beads, check_path, next_step): existing_points = np.concat((existing_points, state.include_compound.xyz)) # An initial point was manually given in hard_sphere_random_walk, use that. - # Check if this point causes any overlaps, if so, raise error. + # Check if this point causes any overlaps, if so, raise error. if isinstance(state.initial_point, np.ndarray) and state.initial_point.shape == ( 3, ): From 8636f73b525afc26eebaacdd2770ac6734942930 Mon Sep 17 00:00:00 2001 From: CalCraven Date: Wed, 10 Jun 2026 18:17:40 -0500 Subject: [PATCH 07/11] Fixes to scaling other params in ffhandler --- mbuild/path/build.py | 4 ---- mbuild/simulation.py | 15 ++++++++------- 2 files changed, 8 insertions(+), 11 deletions(-) diff --git a/mbuild/path/build.py b/mbuild/path/build.py index 14bd80246..6c75e53dd 100644 --- a/mbuild/path/build.py +++ b/mbuild/path/build.py @@ -1496,10 +1496,6 @@ def __init__( ): self.bond_length = bond_length self.radius = radius - if bond_length < radius: - raise ValueError( - "Bond length should be greater than radius to prevent overlaps." - ) if angles_sampler is None: self.angles = AnglesSampler( "uniform", {"low": np.pi / 2, "high": np.pi}, seed diff --git a/mbuild/simulation.py b/mbuild/simulation.py index c02db1ae1..45be85628 100644 --- a/mbuild/simulation.py +++ b/mbuild/simulation.py @@ -329,10 +329,10 @@ def scale_sim(self, sim): "Iterate through HOOMD force objects and apply scaling factors." sim.active_forces = [] # reset active forces forcesDict = { - "lj": (hoomd.md.pair.LJ, ("epsilon")), - "charge": (hoomd.md.special_pair.Coulomb, ("alpha")), - "bond": (hoomd.md.bond.Harmonic, ("k")), - "angle": (hoomd.md.angle.Harmonic, ("k")), + "lj": (hoomd.md.pair.LJ, ("epsilon", "sigma")), + "charge": (hoomd.md.special_pair.Coulomb, ("alpha",)), + "bond": (hoomd.md.bond.Harmonic, ("k",)), + "angle": (hoomd.md.angle.Harmonic, ("k",)), "opls": ( hoomd.md.dihedral.OPLS, ( @@ -342,8 +342,8 @@ def scale_sim(self, sim): "k4", ), ), - "periodic": (hoomd.md.dihedral.Periodic, ("k")), - "improper": (hoomd.md.improper.Periodic, ("k")), + "periodic": (hoomd.md.dihedral.Periodic, ("k",)), + "improper": (hoomd.md.improper.Periodic, ("k",)), } for key, scalar in self.scale_forces.items(): if not scalar: # skip scalars of 0 @@ -358,10 +358,11 @@ def scale_sim(self, sim): force = sim.get_force(forcesDict[key][0]) orig_params = sim._orig_force_params.get(id(force), {}) for param in force.params: - for term in forcesDict[key][1:]: + for term in forcesDict[key][1]: if param in orig_params and term in orig_params[param]: force.params[param][term] = orig_params[param][term] * scalar else: + import pdb; pdb.set_trace() force.params[param][term] *= scalar sim.active_forces.append(force) self.forcesDict[key] = force # store for usage elsewhere From c4680029dae999ebb2699139887ab5794c36c412 Mon Sep 17 00:00:00 2001 From: CalCraven Date: Sun, 14 Jun 2026 16:11:45 -0500 Subject: [PATCH 08/11] Fix visualize system to handle different systems using R1..., Adjust crosslink function to take multiple connection points --- mbuild/path/build.py | 43 +- mbuild/path/constraints.py | 30 +- mbuild/path/crosslink.py | 2520 ++++++++++++------------------- mbuild/path/formats.py | 10 +- mbuild/path/points.py | 9 +- mbuild/simulation.py | 2 +- mbuild/tests/test_crosslinks.py | 19 + mbuild/tests/test_path.py | 26 + mbuild/utils/visualize.py | 11 +- 9 files changed, 1102 insertions(+), 1568 deletions(-) diff --git a/mbuild/path/build.py b/mbuild/path/build.py index 6c75e53dd..1465bed67 100644 --- a/mbuild/path/build.py +++ b/mbuild/path/build.py @@ -412,7 +412,13 @@ def print_bond_lengths(self, box=None): for i, j in self.bond_graph.edges(): delta = positions[j] - positions[i] - if box is not None: + if isinstance(box, CuboidConstraint): + box_arr = box.box_lengths + delta -= np.round(delta / box_arr) * box_arr + elif isinstance(box, Box): + box_arr = box.lengths + delta -= np.round(delta / box_arr) * box_arr + elif isinstance(box, list): box_arr = np.array(box) delta -= np.round(delta / box_arr) * box_arr bl = np.linalg.norm(delta) @@ -615,6 +621,37 @@ def crosslink( overlap_radius=overlap_radius, ) + def wrap_inside_box(self, constraint): + """ + Wrap coordinates inside a cubic box constraint with periodic boundary conditions. + + Args: + coordinates: ndarray of shape (N, 3) with 3D coordinates + box_length: length of the cubic box + center: center of the box as tuple (x, y, z) + pbc: tuple of booleans indicating PBC for each dimension (x, y, z) + + Returns: + ndarray of wrapped coordinates with same shape as input + """ + coordinates = np.array(self.coordinates, dtype=float) + wrapped = coordinates.copy() + center = np.array(constraint.center) + pbc = constraint.pbc + + # Calculate box bounds relative to center + half_lengths = constraint.box_lengths / 2.0 + box_min = center - half_lengths + box_max = center + half_lengths + + # Apply periodic boundary conditions to each dimension + for dim in range(3): + if pbc[dim]: + # Shift coordinates to box range and wrap using modulo + shifted = wrapped[:, dim] - box_min[dim] + wrapped[:, dim] = (shifted % constraint.box_lengths[dim]) + box_min[dim] + + self.coordinates = wrapped def lamellar( path=None, @@ -1360,7 +1397,7 @@ def hard_sphere_random_walk( # Create mask for particles inside volume constraint, allows for PBC if state.volume_constraint: is_inside_mask = volume_constraint.is_inside( - points=candidates, buffer=radius + points=candidates, buffer=0.0 ) candidates = candidates[is_inside_mask] # If there is a bias, sort candidates according to the bias @@ -1397,7 +1434,7 @@ def hard_sphere_random_walk( # Iterate through current state of candidates, break after first accept for xyz in candidates: if check_path_cpu( - existing_points=existing_points, + existing_points=existing_points[:-1], # skip previously bonded to coordinate new_point=xyz, radius=radius, tolerance=tolerance, diff --git a/mbuild/path/constraints.py b/mbuild/path/constraints.py index 5dcca495c..950446c46 100644 --- a/mbuild/path/constraints.py +++ b/mbuild/path/constraints.py @@ -98,7 +98,7 @@ def is_inside(self, points, buffer): pbc=self.pbc, ) - def sample_candidates(self, points, n_candidates, buffer, k=10): + def sample_candidates(self, points, n_candidates, buffer, rng=None, k=10): """Generate candidate points uniformly distributed inside the box, optionally ranked by lowest local density around existing points. @@ -124,8 +124,12 @@ def sample_candidates(self, points, n_candidates, buffer, k=10): the array is sorted so that points with the greatest distance to the nearest neighbor (lowest local density) appear first. """ + if rng is None: + rng = np.random.default_rng(seed=1) + elif isinstance(rng, int): + rng = np.random.default_rng(seed=rng) # Create random candidates inside the box to test and sample from - candidates = np.random.uniform( + candidates = rng.uniform( self.mins + buffer, self.maxs - buffer, size=(n_candidates, 3) ) if points is None or len(points) == 0: @@ -176,7 +180,7 @@ def is_inside(self, points, buffer): """ return is_inside_sphere(points=points, sphere_radius=self.radius, buffer=buffer) - def sample_candidates(self, points, n_candidates, buffer, k=10): + def sample_candidates(self, points, n_candidates, buffer, rng=None, k=10): """Generate candidate points uniformly distributed inside the sphere, optionally ranked by lowest local density around existing points. @@ -202,10 +206,14 @@ def sample_candidates(self, points, n_candidates, buffer, k=10): the array is sorted so that points with the greatest distance to the nearest neighbor (lowest local density) appear first. """ + if rng is None: + rng = np.random.default_rng(seed=1) + elif isinstance(rng, int): + rng = np.random.default_rng(seed=rng) effective_radius = self.radius - buffer - dirs = np.random.normal(size=(n_candidates, 3)) + dirs = rng.normal(size=(n_candidates, 3)) dirs /= np.linalg.norm(dirs, axis=1)[:, None] - u = np.random.random(size=n_candidates) + u = rng.random(size=n_candidates) radii = effective_radius * (u ** (1 / 3)) candidates = self.center + dirs * radii[:, None] # If just sampling from the volume, no KDTree needed @@ -282,7 +290,7 @@ def is_inside(self, points, buffer): periodic=self.periodic_height, ) - def sample_candidates(self, points, n_candidates, buffer, k=10): + def sample_candidates(self, points, n_candidates, buffer, rng=None, k=10): """Generate candidate points uniformly distributed inside the cylinder, optionally ranked by lowest local density around existing points. @@ -308,13 +316,17 @@ def sample_candidates(self, points, n_candidates, buffer, k=10): the array is sorted so that points with the greatest distance to the nearest neighbor (lowest local density) appear first. """ + if rng is None: + rng = np.random.default_rng(seed=1) + elif isinstance(rng, int): + rng = np.random.default_rng(seed=rng) eff_radius = max(self.radius - buffer, 0.0) eff_half_height = max(self.height * 0.5 - buffer, 0.0) - z = self.center[2] + np.random.uniform( + z = self.center[2] + rng.uniform( -eff_half_height, eff_half_height, size=n_candidates ) - r = eff_radius * np.sqrt(np.random.random(size=n_candidates)) - theta = 2 * np.pi * np.random.random(size=n_candidates) + r = eff_radius * np.sqrt(rng.random(size=n_candidates)) + theta = 2 * np.pi * rng.random(size=n_candidates) x = self.center[0] + r * np.cos(theta) y = self.center[1] + r * np.sin(theta) candidates = np.column_stack((x, y, z)) diff --git a/mbuild/path/crosslink.py b/mbuild/path/crosslink.py index 9a4a7c94c..8bb2d4780 100644 --- a/mbuild/path/crosslink.py +++ b/mbuild/path/crosslink.py @@ -1,169 +1,73 @@ -from itertools import permutations +"""Crosslinking module for polymer path structures. + +Handles placement of crosslinker geometries between backbone beads, +with proper overlap avoidance and bond length enforcement under PBC. +""" -import networkx as nx import numpy as np +import networkx as nx from mbuild.exceptions import PathConvergenceError from mbuild.path.build import Path from mbuild.path.constraints import CuboidConstraint, CylinderConstraint -from mbuild.path.path_utils import calculate_sq_distances -class CrosslinkerGeometry(Path): - """A Path subclass representing a crosslinker with connection site metadata. +# ============================================================================= +# CrosslinkerGeometry +# ============================================================================= - Inherits all Path functionality (coordinates, bond_graph, beads) and adds - metadata about which sites form bonds to external beads. + +class CrosslinkerGeometry(Path): + """Rigid crosslinker geometry with connection site metadata. Parameters ---------- coordinates : array-like, shape (N, 3) - Positions of crosslinker sites. Will be re-centered at origin. + Positions of crosslinker sites (re-centered at origin). bond_graph : networkx.Graph, optional Internal connectivity. bead_name : str or np.ndarray, default "_R" Name(s) for each site. connection_sites : list of int Node indices that form bonds to external beads. - Length determines how many external bonds this geometry requires. - Repeated indices allow one site to bond to multiple external beads. - - Examples - -------- - >>> cl = CrosslinkerGeometry.equilateral_triangle(bond_length=0.27) - >>> cl.connection_sites - [0, 1, 2] - >>> cl.bond_graph.edges() - EdgeView([(0, 1), (0, 2), (1, 2)]) + May contain duplicates (same node bonds to multiple groups). """ - def __init__( - self, - coordinates=None, - bond_graph=None, - bead_name="_R", - connection_sites=None, - ): - super().__init__( - coordinates=coordinates, bond_graph=bond_graph, bead_name=bead_name - ) - - if connection_sites is None: - connection_sites = list(range(len(self.coordinates))) - - self.connection_sites = list(connection_sites) - self._validate_connection_sites() - - # Center positions at origin - if len(self.coordinates) > 0: - centroid = np.mean(self.coordinates, axis=0) - self.coordinates = (self.coordinates - centroid).astype(np.float32) - - def _validate_connection_sites(self): - n = len(self.coordinates) - for idx in self.connection_sites: - if not (0 <= idx < n): - raise ValueError( - f"Connection site {idx} is out of range (must be 0 to {n - 1})" - ) - - @property - def n_sites(self) -> int: - return len(self.coordinates) - - @property - def n_connections(self) -> int: - return len(self.connection_sites) - - @property - def unique_connection_sites(self): - return list(dict.fromkeys(self.connection_sites)) - - @property - def internal_bonds(self): - return list(self.bond_graph.edges()) - - def copy(self): - """Return a deep copy of this crosslinker geometry.""" - return CrosslinkerGeometry( - coordinates=self.coordinates.copy(), - bond_graph=self.bond_graph.copy(), - bead_name=self.beads.copy(), - connection_sites=list(self.connection_sites), - ) - - def recenter(self): - if len(self.coordinates) > 0: - centroid = np.mean(self.coordinates, axis=0) - self.coordinates = self.coordinates - centroid + def __init__(self, coordinates, bond_graph=None, bead_name="_R", connection_sites=None): + coordinates = np.asarray(coordinates, dtype=np.float32) + centroid = coordinates.mean(axis=0) + self.coordinates = coordinates - centroid - @classmethod - def from_path(cls, path, connection_sites): - """Create from an existing Path. - - Parameters - ---------- - path : Path - An existing path defining the structure. - connection_sites : list of int - Which nodes bond to external beads. - """ - return cls( - coordinates=path.coordinates.copy(), - bond_graph=path.bond_graph.copy(), - bead_name=path.beads.copy(), - connection_sites=connection_sites, - ) + self.n_sites = len(self.coordinates) - @classmethod - def single_site(cls, bead_name="_R", n_connections=2): - """Single-site (original crosslink behavior).""" - coords = np.array([[0.0, 0.0, 0.0]], dtype=np.float32) - return cls( - coordinates=coords, - bead_name=bead_name, - connection_sites=[0] * n_connections, - ) + if bond_graph is None: + bond_graph = nx.Graph() + for i in range(self.n_sites): + bond_graph.add_node(i) + self.bond_graph = bond_graph - @classmethod - def linear(cls, n_sites=2, bond_length=0.27, bead_name="_R", connection_sites=None): - """Linear chain of sites.""" - positions = np.zeros((n_sites, 3), dtype=np.float32) - for i in range(n_sites): - positions[i, 0] = i * bond_length + if isinstance(bead_name, str): + self.beads = np.array([bead_name] * self.n_sites, dtype="U10") + else: + self.beads = np.array(bead_name, dtype="U10") - bead_names = ( - np.array([bead_name] * n_sites, dtype="U10") - if isinstance(bead_name, str) - else np.array(bead_name, dtype="U10") - ) if connection_sites is None: - connection_sites = [0, n_sites - 1] - - G = nx.Graph() - for i in range(n_sites): - G.add_node(i) - for i in range(n_sites - 1): - G.add_edge(i, i + 1) + connection_sites = list(range(self.n_sites)) + self.connection_sites = list(connection_sites) - return cls( - coordinates=positions, - bond_graph=G, - bead_name=bead_names, - connection_sites=connection_sites, - ) + @property + def n_connections(self): + return len(self.connection_sites) @classmethod - def equilateral_triangle( - cls, bond_length=0.27, bead_name="_R", connection_sites=None - ): - """Three sites in a flat equilateral triangle.""" + def equilateral_triangle(cls, bond_length=0.27, bead_name="_R", connection_sites=None): + """Three sites in a flat equilateral triangle with edge = bond_length.""" R = bond_length / np.sqrt(3) angles = [0, 2 * np.pi / 3, 4 * np.pi / 3] positions = np.array( - [[R * np.cos(a), R * np.sin(a), 0.0] for a in angles], dtype=np.float32 + [[R * np.cos(a), R * np.sin(a), 0.0] for a in angles], + dtype=np.float32, ) - bead_names = ( np.array([bead_name] * 3, dtype="U10") if isinstance(bead_name, str) @@ -171,14 +75,12 @@ def equilateral_triangle( ) if connection_sites is None: connection_sites = [0, 1, 2] - G = nx.Graph() for i in range(3): G.add_node(i) G.add_edge(0, 1) G.add_edge(1, 2) G.add_edge(0, 2) - return cls( coordinates=positions, bond_graph=G, @@ -187,1599 +89,1133 @@ def equilateral_triangle( ) @classmethod - def square(cls, bond_length=0.27, bead_name="_R", connection_sites=None): - """Four sites in a flat square.""" - half = bond_length / 2.0 + def linear(cls, bond_length=0.27, bead_name="_R", connection_sites=None): + """Two sites separated by bond_length.""" positions = np.array( - [[half, half, 0], [-half, half, 0], [-half, -half, 0], [half, -half, 0]], + [[-bond_length / 2, 0, 0], [bond_length / 2, 0, 0]], dtype=np.float32, ) - bead_names = ( - np.array([bead_name] * 4, dtype="U10") - if isinstance(bead_name, str) - else np.array(bead_name, dtype="U10") - ) if connection_sites is None: - connection_sites = [0, 1, 2, 3] - + connection_sites = [0, 1] G = nx.Graph() - for i in range(4): - G.add_node(i) - for i in range(4): - G.add_edge(i, (i + 1) % 4) - + G.add_node(0) + G.add_node(1) + G.add_edge(0, 1) return cls( - coordinates=positions, - bond_graph=G, - bead_name=bead_names, - connection_sites=connection_sites, - ) - - @classmethod - def tetrahedral(cls, bond_length=0.27, bead_name="_R", connection_sites=None): - """Four sites in a regular tetrahedron.""" - positions = np.array( - [[1, 1, 1], [1, -1, -1], [-1, 1, -1], [-1, -1, 1]], dtype=np.float32 - ) - current_dist = np.linalg.norm(positions[0] - positions[1]) - positions *= bond_length / current_dist - - bead_names = ( - np.array([bead_name] * 4, dtype="U10") - if isinstance(bead_name, str) - else np.array(bead_name, dtype="U10") + coordinates=positions, bond_graph=G, + bead_name=bead_name, connection_sites=connection_sites, ) - if connection_sites is None: - connection_sites = [0, 1, 2, 3] - G = nx.Graph() - for i in range(4): - G.add_node(i) - for i in range(4): - for j in range(i + 1, 4): - G.add_edge(i, j) - - return cls( - coordinates=positions, - bond_graph=G, - bead_name=bead_names, - connection_sites=connection_sites, - ) - @classmethod - def trigonal_bipyramidal( - cls, bond_length=0.27, bead_name="_R", connection_sites=None - ): - """Five sites in trigonal bipyramidal arrangement (0-2 equatorial, 3-4 axial).""" - R_eq = bond_length / np.sqrt(3) - eq_angles = [0, 2 * np.pi / 3, 4 * np.pi / 3] - equatorial = [[R_eq * np.cos(a), R_eq * np.sin(a), 0.0] for a in eq_angles] - axial_dist = bond_length * np.sqrt(2.0 / 3.0) - axial = [[0.0, 0.0, axial_dist], [0.0, 0.0, -axial_dist]] - positions = np.array(equatorial + axial, dtype=np.float32) +# ============================================================================= +# PBC Utilities +# ============================================================================= - bead_names = ( - np.array([bead_name] * 5, dtype="U10") - if isinstance(bead_name, str) - else np.array(bead_name, dtype="U10") - ) - if connection_sites is None: - connection_sites = [0, 1, 2, 3, 4] - G = nx.Graph() - for i in range(5): - G.add_node(i) - G.add_edge(0, 1) - G.add_edge(1, 2) - G.add_edge(0, 2) - for eq in range(3): - G.add_edge(eq, 3) - G.add_edge(eq, 4) +def _pbc_delta(a, b, box_lengths, pbc): + """Minimum-image displacement: a - b.""" + delta = np.asarray(a, dtype=np.float64) - np.asarray(b, dtype=np.float64) + for dim in range(3): + if pbc[dim] and box_lengths[dim] > 0: + delta[..., dim] -= ( + np.round(delta[..., dim] / box_lengths[dim]) * box_lengths[dim] + ) + return delta.astype(np.float32) - return cls( - coordinates=positions, - bond_graph=G, - bead_name=bead_names, - connection_sites=connection_sites, - ) - @classmethod - def pentagon(cls, bond_length=0.27, bead_name="_R", connection_sites=None): - """Five sites in a flat regular pentagon.""" - R = bond_length / (2 * np.sin(np.pi / 5)) - angles = [2 * np.pi * i / 5 for i in range(5)] - positions = np.array( - [[R * np.cos(a), R * np.sin(a), 0.0] for a in angles], dtype=np.float32 - ) +def _pbc_distance(a, b, box_lengths, pbc): + """Minimum-image scalar distance.""" + delta = _pbc_delta(a, b, box_lengths, pbc) + return np.linalg.norm(delta, axis=-1) - bead_names = ( - np.array([bead_name] * 5, dtype="U10") - if isinstance(bead_name, str) - else np.array(bead_name, dtype="U10") - ) - if connection_sites is None: - connection_sites = [0, 1, 2, 3, 4] - G = nx.Graph() - for i in range(5): - G.add_node(i) - for i in range(5): - G.add_edge(i, (i + 1) % 5) +def _wrap_positions(positions, box_lengths, pbc): + """Wrap positions into periodic box centered at origin.""" + positions = positions.copy() + for dim in range(3): + if pbc[dim] and box_lengths[dim] > 0: + half = box_lengths[dim] / 2 + positions[:, dim] = ((positions[:, dim] + half) % box_lengths[dim]) - half + return positions - return cls( - coordinates=positions, - bond_graph=G, - bead_name=bead_names, - connection_sites=connection_sites, - ) - @classmethod - def from_edges(cls, coordinates, edges, bead_name="_R", connection_sites=None): - """Create from explicit positions and edge list. - - Parameters - ---------- - coordinates : array-like, shape (N, 3) - edges : list of tuple(int, int) - bead_name : str or list of str - connection_sites : list of int, optional - """ - coordinates = np.asarray(coordinates, dtype=np.float32) - n = len(coordinates) - bead_names = ( - np.array([bead_name] * n, dtype="U10") - if isinstance(bead_name, str) - else np.array(bead_name, dtype="U10") - ) - if connection_sites is None: - connection_sites = list(range(n)) +# ============================================================================= +# Backbone Specification Parsing +# ============================================================================= - G = nx.Graph() - for i in range(n): - G.add_node(i) - for i, j in edges: - G.add_edge(i, j) - return cls( - coordinates=coordinates, - bond_graph=G, - bead_name=bead_names, - connection_sites=connection_sites, - ) +def _parse_backbone_spec(backbone_name, connection_sites): + """Parse backbone_name into per-connection-site specifications. + Parameters + ---------- + backbone_name : str or tuple + - String: every connection site bonds to one bead of that type. + - Tuple of length == len(connection_sites): per-site spec. + - String entry: one backbone bead. + - Tuple entry ("_B", "_B"): multiple neighboring backbone beads. -def replace_sites( - path, - replacement, - sites=None, - bead_name=None, - bond_length=None, - tolerance=0.1, - min_separation=None, - volume_constraint=None, - n_rotation_samples=36, - overlap_radius=None, - seed=42, -): - """Replace sites in a Path with a multi-site substructure, in place. - ... + Returns + ------- + specs : list of list[str] + specs[i] is a list of bead names that connection site i bonds to. """ - if sites is None and bead_name is None: - raise ValueError("Must specify either `sites` or `bead_name`.") - if sites is not None and bead_name is not None: - raise ValueError("Cannot specify both `sites` and `bead_name`.") - - rng = np.random.default_rng(seed) - - if bead_name is not None: - sites = [ - node for node in path.bond_graph.nodes() if path.beads[node] == bead_name - ] - if isinstance(sites, int): - sites = [sites] - else: - sites = list(sites) + n_conn = len(connection_sites) - if len(sites) == 0: - return path + if isinstance(backbone_name, str): + return [[backbone_name]] * n_conn - for site in sites: - degree = path.bond_graph.degree[site] - if degree > replacement.n_connections: - raise ValueError( - f"Site {site} has degree {degree} but replacement only has " - f"{replacement.n_connections} connection sites. " - f"Degree must not exceed len(replacement.connection_sites)." - ) + if not isinstance(backbone_name, (tuple, list)): + raise ValueError(f"backbone_name must be str or tuple, got {type(backbone_name)}") - pbc, box_lengths = _get_pbc_info(volume_constraint) - sites_set = set(sites) - - # --- Compute all replacements --- - replacement_data = {} - - for site in sites: - neighbors = list(path.bond_graph.neighbors(site)) - site_pos = path.coordinates[site] - degree = len(neighbors) - - # Unwrapped neighbor positions - neighbor_coords = np.zeros((degree, 3), dtype=np.float32) - for i, nb in enumerate(neighbors): - delta = path.coordinates[nb] - site_pos - delta = _apply_mic(delta, pbc, box_lengths) - neighbor_coords[i] = site_pos + delta - - # Target bond lengths - if bond_length is not None: - target_bond_lengths = np.full(degree, bond_length, dtype=np.float32) + if len(backbone_name) != n_conn: + raise ValueError( + f"backbone_name length ({len(backbone_name)}) must match " + f"number of connection sites ({n_conn})." + ) + + specs = [] + for entry in backbone_name: + if isinstance(entry, str): + specs.append([entry]) + elif isinstance(entry, (tuple, list)): + specs.append(list(entry)) else: - target_bond_lengths = np.array( - [np.linalg.norm(neighbor_coords[i] - site_pos) for i in range(degree)], - dtype=np.float32, - ) + raise ValueError(f"Entry must be str or tuple, got {type(entry)}") - active_conn_sites = replacement.connection_sites[:degree] - - # Indices that are bonded to this site (excluded from overlap check) - bonded_indices = set(neighbors) | {site} | sites_set - - # Place via rigid body alignment + overlap optimization - placed_positions, assignment = _place_replacement_rigid_full( - replacement, - active_conn_sites, - neighbor_coords, - target_bond_lengths, - path.coordinates, # pass ALL coordinates - bonded_indices, - pbc, - box_lengths, - rng, - n_rotation_samples, - tolerance, - min_separation, - ) + return specs - # Wrap into PBC - placed_positions = _wrap_positions(placed_positions, pbc, box_lengths) - - replacement_data[site] = { - "positions": placed_positions, - "neighbors": neighbors, - "assignment": assignment, - } - - # --- Rebuild path in place (same as before) --- - kept_nodes = [n for n in sorted(path.bond_graph.nodes()) if n not in sites_set] - - new_coords = [] - new_beads = [] - old_to_new = {} - current_idx = 0 - - for old_node in kept_nodes: - old_to_new[old_node] = current_idx - new_coords.append(path.coordinates[old_node]) - new_beads.append(path.beads[old_node]) - current_idx += 1 - - site_new_indices = {} - for site in sites: - data = replacement_data[site] - indices = [] - for i in range(replacement.n_sites): - indices.append(current_idx) - new_coords.append(data["positions"][i]) - new_beads.append(replacement.beads[i]) - current_idx += 1 - site_new_indices[site] = indices - - new_edges = [] - - for u, v, data in path.bond_graph.edges(data=True): - if u in sites_set or v in sites_set: - continue - new_edges.append((old_to_new[u], old_to_new[v], data)) - - for site in sites: - indices = site_new_indices[site] - for i, j in replacement.bond_graph.edges(): - edge_data = { - "bond_type": (str(replacement.beads[i]), str(replacement.beads[j])) - } - new_edges.append((indices[i], indices[j], edge_data)) - - for site in sites: - data = replacement_data[site] - indices = site_new_indices[site] - neighbors = data["neighbors"] - assignment = data["assignment"] - - for conn_idx, neighbor_idx in enumerate(assignment): - repl_site_local = replacement.connection_sites[conn_idx] - repl_new_node = indices[repl_site_local] - external_neighbor = neighbors[neighbor_idx] - - if external_neighbor in old_to_new: - ext_new_node = old_to_new[external_neighbor] - ext_bead_name = path.beads[external_neighbor] - elif external_neighbor in site_new_indices: - other_data = replacement_data[external_neighbor] - other_neighbors = other_data["neighbors"] - try: - back_idx = other_neighbors.index(site) - except ValueError: - continue - other_assignment = other_data["assignment"] - found = False - for other_conn_idx, other_target in enumerate(other_assignment): - if other_target == back_idx: - other_repl_site_local = replacement.connection_sites[ - other_conn_idx - ] - ext_new_node = site_new_indices[external_neighbor][ - other_repl_site_local - ] - ext_bead_name = replacement.beads[other_repl_site_local] - found = True - break - if not found: - continue - else: - continue - edge_data = { - "bond_type": ( - str(replacement.beads[repl_site_local]), - str(ext_bead_name), - ) - } - new_edges.append((repl_new_node, ext_new_node, edge_data)) +# ============================================================================= +# Geometry & Graph Helpers +# ============================================================================= - path.coordinates = np.array(new_coords, dtype=np.float32) - path.beads = np.array(new_beads, dtype="U10") - path.bond_graph = nx.Graph() - for i in range(len(path.coordinates)): - path.bond_graph.add_node(i, name=str(path.beads[i]), xyz=path.coordinates[i]) +def _get_excluded_indices(path, node_indices, excluded_bond_depth): + """Get indices within excluded_bond_depth bonds of any node in node_indices.""" + excluded = set(node_indices) + frontier = set(node_indices) + for _ in range(excluded_bond_depth): + next_frontier = set() + for node in frontier: + if node in path.bond_graph: + for neighbor in path.bond_graph.neighbors(node): + if neighbor not in excluded: + excluded.add(neighbor) + next_frontier.add(neighbor) + frontier = next_frontier + return excluded - seen_edges = set() - for u, v, data in new_edges: - edge_key = (min(u, v), max(u, v)) - if edge_key not in seen_edges: - seen_edges.add(edge_key) - path.bond_graph.add_edge(u, v, **data) - return path +def _get_perpendicular(v): + """Return a unit vector perpendicular to v.""" + v = np.asarray(v, dtype=np.float64) + if abs(v[0]) < 0.9: + perp = np.cross(v, [1, 0, 0]) + else: + perp = np.cross(v, [0, 1, 0]) + norm = np.linalg.norm(perp) + return (perp / max(norm, 1e-10)).astype(np.float32) -def crosslink( - path, - crosslinker=None, - bead_name="_R", - backbone_name="_A", - crosslink_bond_length=0.2, - max_backbone_degree=2, - tolerance=0.1, - excluded_bond_depth=2, - n_connection_sites=2, - volume_constraint=None, - initial_point=None, - seed=42, - n_rotation_samples=36, - overlap_radius=None, - min_separation=None, -): - """ - Create a crosslink bonded to backbone beads, modifying path in place. +# ============================================================================= +# Candidate Search +# ============================================================================= - Selects backbone beads whose pairwise spacing matches the crosslinker - geometry, places a temporary placeholder, then calls ``replace_sites`` - which handles rigid-body alignment, bond length guarantees, and overlap - avoidance in a single pass. - Parameters - ---------- - path : Path - The Path object to modify in place. - crosslinker : CrosslinkerGeometry, optional - If None, uses single-site from ``bead_name`` and ``n_connection_sites``. - bead_name : str, default "_R" - backbone_name : str, default "_A" - crosslink_bond_length : float, default 0.2 - Desired bond length from connection sites to backbone beads (nm). - tolerance : float, default 0.1 - Fractional tolerance on bond lengths (±10%). - excluded_bond_depth : int, default 2 - n_connection_sites : int, default 2 - volume_constraint : optional - initial_point : int or array-like, optional - seed : int, default 42 - chunk_size : int, default 512 - run_on_gpu : bool, default False - n_rotation_samples : int, default 36 - overlap_radius : float, optional - min_separation : float, optional - Minimum distance from placed sites to existing sites. - Defaults to ``crosslink_bond_length * 0.5``. +def _find_neighbor_groups(path, target_names, max_backbone_degree): + """Find bonded groups of backbone beads matching target_names spec. + + For ("_B", "_B"), finds pairs of directly bonded beads where both + have degree < max_backbone_degree and names match. Returns ------- - Path - Same path, modified in place. + groups : list of list[int] """ - if crosslinker is None: - crosslinker = CrosslinkerGeometry.single_site( - bead_name=bead_name, n_connections=n_connection_sites - ) + n_needed = len(target_names) + + all_eligible = set() + for i in range(len(path.beads)): + if path.bond_graph.degree(i) < max_backbone_degree: + if path.beads[i] in target_names: + all_eligible.add(i) + + if n_needed == 1: + return [[i] for i in all_eligible if path.beads[i] == target_names[0]] + + if n_needed == 2: + groups = [] + seen = set() + for i in all_eligible: + for j in path.bond_graph.neighbors(i): + if j not in all_eligible: + continue + key = (min(i, j), max(i, j)) + if key in seen: + continue + seen.add(key) + if (path.beads[i] == target_names[0] + and path.beads[j] == target_names[1]): + groups.append([i, j]) + elif (path.beads[i] == target_names[1] + and path.beads[j] == target_names[0]): + groups.append([j, i]) + return groups + + # General case for n_needed > 2 + groups = [] + found_keys = set() + for start in all_eligible: + _dfs_groups(path, start, all_eligible, target_names, n_needed, + groups, found_keys) + return groups + + +def _dfs_groups(path, start, eligible, target_names, size, groups, found_keys): + """DFS to find connected groups of given size from start node.""" + stack = [([start], {start})] + while stack: + current_group, visited = stack.pop() + if len(current_group) == size: + key = tuple(sorted(current_group)) + if key not in found_keys: + group_names = sorted(path.beads[n] for n in current_group) + if group_names == sorted(target_names): + found_keys.add(key) + groups.append(list(current_group)) + continue + last = current_group[-1] + for neighbor in path.bond_graph.neighbors(last): + if neighbor in eligible and neighbor not in visited: + stack.append((current_group + [neighbor], visited | {neighbor})) - if min_separation is None: - min_separation = crosslink_bond_length * 0.2 - n_backbone_connections = crosslinker.n_connections - rng = np.random.default_rng(seed + len(path.coordinates)) +def _all_beads_from_groups(candidate_group): + """Flatten a candidate group into a single list of all bead indices.""" + all_beads = [] + for group in candidate_group: + all_beads.extend(group) + return all_beads - # --- Compute search geometry --- - ideal_bb_positions, ideal_pairwise_distances, geom_search_radius = ( - _compute_ideal_backbone_distances(crosslinker, crosslink_bond_length) - ) - if n_backbone_connections > 1: - max_pair_dist = np.max(ideal_pairwise_distances) - else: - max_pair_dist = 0.0 +def _can_reach_all_beads( + bead_positions, crosslink_bond_length, tolerance, pbc, box_lengths, +): + """Check if there exists a point at crosslink_bond_length from all beads. - search_radius = max_pair_dist * (1 + tolerance) + Uses the circumcenter/equidistant-point approach. A necessary condition + is that all pairwise distances are <= 2 * crosslink_bond_length. - # --- Find backbone targets --- - backbone_nodes = [ - node for node in path.bond_graph.nodes() if path.beads[node] == backbone_name - ] - backbone_subgraph = path.bond_graph.subgraph(backbone_nodes) + For exact checking, computes the geometric median constraint. - if len(backbone_nodes) == 0: - raise ValueError(f"No backbone beads with name '{backbone_name}' found in path") - if len(backbone_nodes) < n_backbone_connections: - raise ValueError( - f"Not enough backbone beads ({len(backbone_nodes)}) for " - f"{n_backbone_connections} connection sites" - ) + Returns + ------- + feasible : bool + candidate_point : np.ndarray or None + An approximate valid point if feasible. + """ + n = len(bead_positions) + r = crosslink_bond_length + tol = tolerance - pbc, box_lengths = _get_pbc_info(volume_constraint) + # Necessary condition: all pairwise distances <= 2*r*(1+tol) + max_pair = 2 * r * (1 + tol) + for i in range(n): + for j in range(i + 1, n): + d = float(_pbc_distance(bead_positions[i], bead_positions[j], + box_lengths, pbc)) + if d > max_pair: + return False, None + + # Find the point equidistant from all beads via iterative projection + # Start from centroid + centroid = np.mean(bead_positions, axis=0) + point = centroid.copy().astype(np.float64) + + for _ in range(50): + shift = np.zeros(3, dtype=np.float64) + for i in range(n): + delta = _pbc_delta(bead_positions[i], point, box_lengths, pbc).astype(np.float64) + dist = np.linalg.norm(delta) + if dist > 1e-10: + # Project point onto sphere of radius r around bead i + target_on_sphere = point + delta * (1 - r / dist) + shift += (target_on_sphere - point) + shift /= n + point += shift * 0.8 # damped + + if np.linalg.norm(shift) < r * 0.001: + break - candidate_nodes = [ - node - for node in backbone_nodes - if path.bond_graph.degree[node] <= max_backbone_degree - ] - candidate_coords = np.array(path.coordinates[candidate_nodes], dtype=np.float32) + # Validate + max_error = 0.0 + for i in range(n): + d = float(_pbc_distance(point, bead_positions[i], box_lengths, pbc)) + max_error = max(max_error, abs(d - r) / r) - ref_nodes, ref_coords = _get_reference_points( - path, initial_point, candidate_nodes, pbc, box_lengths, rng - ) + if max_error <= tol: + return True, point.astype(np.float32) - found_ref = False - selected_nodes = [] + return False, None - for ref_node, ref_coord in zip(ref_nodes, ref_coords): - selected_nodes = [ref_node] - sq_distances = calculate_sq_distances( - ref_coord, candidate_coords, pbc=pbc, box_lengths=box_lengths +def _find_candidate_groups( + path, + crosslinker, + backbone_specs, + crosslink_bond_length, + tolerance, + excluded_bond_depth, + max_backbone_degree, + pbc, + box_lengths, + rng, +): + """Find groups of backbone nodes that can be connected by the crosslinker. + + Each candidate is a list of groups (one per connection site), where each + group is a list of backbone node indices. + + The key constraint: for each connection site, the crosslinker node at that + site must be placeable at crosslink_bond_length from EACH bead in its group. + """ + n_conn = len(crosslinker.connection_sites) + coords = path.coordinates + + # --- Gather eligible groups per connection site --- + eligible_per_site = [] + for site_idx in range(n_conn): + spec = backbone_specs[site_idx] + groups = _find_neighbor_groups(path, spec, max_backbone_degree) + eligible_per_site.append(groups) + + if any(len(e) == 0 for e in eligible_per_site): + return [] + + for e in eligible_per_site: + rng.shuffle(e) + + # --- Determine which connection sites share the same physical node --- + # Groups that share a connection site node have coupled constraints + conn_site_nodes = crosslinker.connection_sites # e.g., [0, 0] or [0, 1] + + # Find unique physical nodes and which connection indices map to them + node_to_conn_indices = {} + for conn_idx, node in enumerate(conn_site_nodes): + node_to_conn_indices.setdefault(node, []).append(conn_idx) + + # --- Search for valid combinations --- + if n_conn == 2: + return _pairwise_candidate_search( + path, crosslinker, eligible_per_site, node_to_conn_indices, + crosslink_bond_length, tolerance, excluded_bond_depth, + pbc, box_lengths, rng, ) - distances = np.sqrt(sq_distances) - within_radius_mask = distances <= search_radius - possible_pairs = np.where(within_radius_mask)[0] + return _general_candidate_search( + path, crosslinker, eligible_per_site, node_to_conn_indices, + crosslink_bond_length, tolerance, excluded_bond_depth, + pbc, box_lengths, rng, + ) - if len(possible_pairs) < n_backbone_connections - 1: - continue - closest_paired_nodes = possible_pairs[np.argsort(distances[possible_pairs])] - excluded_nodes = set( - nx.single_source_shortest_path_length( - backbone_subgraph, ref_node, cutoff=excluded_bond_depth - ).keys() - ) +def _pairwise_candidate_search( + path, crosslinker, eligible_per_site, node_to_conn_indices, + crosslink_bond_length, tolerance, excluded_bond_depth, + pbc, box_lengths, rng, +): + """Search for valid 2-connection-site candidates. + + Handles two cases: + 1. Both connection sites are DIFFERENT physical nodes → check that each + group is reachable from its own connection site, plus geometric + compatibility between sites. + 2. Both connection sites are the SAME physical node → all beads from + both groups must be reachable from a single point. + """ + coords = path.coordinates + candidates = [] - for idx in closest_paired_nodes: - node = candidate_nodes[idx] - if node in excluded_nodes: - continue + # Are both connection sites the same physical node? + conn_sites = crosslinker.connection_sites + same_node = (conn_sites[0] == conn_sites[1]) - is_excluded = False - for selected in selected_nodes[1:]: - if selected in backbone_subgraph: - sel_excluded = set( - nx.single_source_shortest_path_length( - backbone_subgraph, selected, cutoff=excluded_bond_depth - ).keys() - ) - if node in sel_excluded: - is_excluded = True - break - if is_excluded: + for group_a in eligible_per_site[0]: + excluded_a = _get_excluded_indices(path, group_a, excluded_bond_depth) + + for group_b in eligible_per_site[1]: + # No overlap between groups + if set(group_b) & set(group_a): + continue + # Exclusion zone + if set(group_b) & excluded_a: + continue + if set(group_a) & _get_excluded_indices(path, group_b, excluded_bond_depth): continue - selected_nodes.append(int(node)) - if len(selected_nodes) >= n_backbone_connections: - # Validate pairwise geometry - sel_coords = np.array( - [path.coordinates[n] for n in selected_nodes], dtype=np.float32 - ) - unwrapped_check = _unwrap_coords(sel_coords, pbc, box_lengths) - valid, _ = _validate_backbone_group( - unwrapped_check, - ideal_pairwise_distances, - crosslink_bond_length, - tolerance, - pbc, - box_lengths, + # --- Geometric feasibility check --- + if same_node: + # Both groups bond to the SAME crosslinker bead. + # All beads from both groups must be reachable from one point. + all_bead_indices = list(group_a) + list(group_b) + all_bead_positions = coords[all_bead_indices].copy() + feasible, _ = _can_reach_all_beads( + all_bead_positions, crosslink_bond_length, tolerance, + pbc, box_lengths, ) - valid_overlaps, clink_centroid = _validate_clinker( - path, - crosslinker, - crosslink_bond_length, - selected_nodes, - overlap_radius=min_separation, + else: + # Different physical nodes. Each group independently checked, + # plus pairwise distance between connection sites must match geometry. + feasible = _check_separate_sites_feasibility( + crosslinker, coords, group_a, group_b, + crosslink_bond_length, tolerance, pbc, box_lengths, ) - if valid and valid_overlaps: - found_ref = True - break - else: - selected_nodes.pop() - if found_ref: - break + if feasible: + candidates.append([group_a, group_b]) - if not found_ref: - n_clinks = sum(1 for b in path.beads if b in set(crosslinker.beads)) - raise PathConvergenceError( - f"Could not find {n_backbone_connections} non-neighboring backbone beads " - f"matching crosslinker geometry with bond_length={crosslink_bond_length} " - f"(tolerance=±{tolerance * 100:.0f}%)." - f"\nCurrent crosslinks: {n_clinks}. " - "Ways to increase crosslinking:\n" - " - Increase crosslink_bond_length\n" - " - Increase tolerance\n" - " - Decrease excluded_bond_depth\n" - " - Pack at higher density\n" - " - Relax structure" - ) + if len(candidates) >= 200: + return candidates - for dim in range(3): - if pbc[dim]: - clink_centroid[dim] -= ( - np.round(clink_centroid[dim] / box_lengths[dim]) * box_lengths[dim] - ) + return candidates - _placeholder_name = "__placeholder__" - path.append_coordinates(clink_centroid, _placeholder_name) - placeholder_idx = len(path.coordinates) - 1 - for backbone_node in selected_nodes: - path.bond_graph.add_edge( - int(placeholder_idx), - int(backbone_node), - bond_type=(_placeholder_name, backbone_name), - ) +def _check_separate_sites_feasibility( + crosslinker, coords, group_a, group_b, + crosslink_bond_length, tolerance, pbc, box_lengths, +): + """Check feasibility when connection sites are different physical nodes. - # --- Single-site shortcut --- - if crosslinker.n_sites == 1: - path.beads[placeholder_idx] = crosslinker.beads[0] - for backbone_node in selected_nodes: - path.bond_graph.edges[placeholder_idx, backbone_node]["bond_type"] = ( - str(crosslinker.beads[0]), - backbone_name, - ) - if "name" in path.bond_graph.nodes[placeholder_idx]: - path.bond_graph.nodes[placeholder_idx]["name"] = str(crosslinker.beads[0]) - - # Fix single-site bond length if needed - for backbone_node in selected_nodes: - delta = path.coordinates[placeholder_idx] - path.coordinates[backbone_node] - delta = _apply_mic(delta, pbc, box_lengths) - actual_bl = np.linalg.norm(delta) - if ( - abs(actual_bl - crosslink_bond_length) - > crosslink_bond_length * tolerance - ): - if actual_bl > 1e-10: - direction = delta / actual_bl - else: - direction = rng.standard_normal(3).astype(np.float32) - direction /= np.linalg.norm(direction) - path.coordinates[placeholder_idx] = path.coordinates[ - backbone_node - ] + _apply_mic(direction * crosslink_bond_length, pbc, box_lengths) - return path - - # --- Multi-site: replace_sites handles everything --- - replace_sites( - path, - replacement=crosslinker, - sites=[placeholder_idx], - bond_length=crosslink_bond_length, - tolerance=tolerance, - min_separation=min_separation, - volume_constraint=volume_constraint, - n_rotation_samples=n_rotation_samples, - overlap_radius=overlap_radius, - seed=seed, + Verifies: + 1. Each group's beads can be reached from a point at crosslink_bond_length. + 2. The distance between the two reach-points is compatible with the + rigid crosslinker geometry (distance between connection site nodes). + """ + conn_sites = crosslinker.connection_sites + conn_positions = crosslinker.coordinates[conn_sites] + internal_distance = float(np.linalg.norm(conn_positions[0] - conn_positions[1])) + + # Check group_a reachability + positions_a = coords[group_a].copy() + ok_a, point_a = _can_reach_all_beads( + positions_a, crosslink_bond_length, tolerance, pbc, box_lengths ) + if not ok_a: + return False - return path + # Check group_b reachability + positions_b = coords[group_b].copy() + ok_b, point_b = _can_reach_all_beads( + positions_b, crosslink_bond_length, tolerance, pbc, box_lengths + ) + if not ok_b: + return False + # Check that the reach-points are at the correct internal distance apart + dist_ab = float(_pbc_distance(point_a, point_b, box_lengths, pbc)) + if abs(dist_ab - internal_distance) > internal_distance * tolerance + crosslink_bond_length * tolerance: + return False -def _get_reference_points(path, initial_point, candidate_nodes, pbc, box_lengths, rng): - candidate_coords = np.array(path.coordinates[candidate_nodes], dtype=np.float32) - if initial_point is not None: - if isinstance(initial_point, (int, np.integer)): - if initial_point not in path.bond_graph.nodes: - raise ValueError(f"Node {initial_point} not found in bond_graph") - return np.array([initial_point]), np.array( - [path.coordinates[initial_point]] - ) - else: - initial_point32 = np.asarray(initial_point, dtype=np.float32) - sq_distances = calculate_sq_distances( - initial_point32, candidate_coords, pbc=pbc, box_lengths=box_lengths - ) - order = np.argsort(sq_distances) - return ( - np.array(candidate_nodes)[order], - path.coordinates[np.array(candidate_nodes)[order]], - ) - else: - shuffled = rng.choice(candidate_nodes, size=len(candidate_nodes), replace=False) - return shuffled, path.coordinates[shuffled] + return True -def _kabsch_rotation(P, Q): - """Optimal rotation aligning P onto Q (Kabsch algorithm).""" - H = P.T @ Q - U, S, Vt = np.linalg.svd(H) - d = np.linalg.det(Vt.T @ U.T) - sign_matrix = np.diag([1.0, 1.0, d]) - R = Vt.T @ sign_matrix @ U.T - return R.astype(np.float32) +def _general_candidate_search( + path, crosslinker, eligible_per_site, node_to_conn_indices, + crosslink_bond_length, tolerance, excluded_bond_depth, + pbc, box_lengths, rng, +): + """General search for n-connection-site crosslinkers.""" + coords = path.coordinates + candidates = [] + max_tries = 5000 + + n_conn = len(crosslinker.connection_sites) + + for _ in range(max_tries): + idx_0 = rng.integers(len(eligible_per_site[0])) + group_0 = eligible_per_site[0][idx_0] + excluded = _get_excluded_indices(path, group_0, excluded_bond_depth) + selected = [group_0] + all_nodes = set(group_0) + valid = True + for site_idx in range(1, n_conn): + found = False + for group in eligible_per_site[site_idx]: + if set(group) & (all_nodes | excluded): + continue -def _rotation_between_vectors(a, b): - """Rotation matrix rotating unit vector a onto b (Rodrigues).""" - v = np.cross(a, b) - c = float(np.dot(a, b)) - s = float(np.linalg.norm(v)) + # Check feasibility of adding this group + # Collect all beads assigned to same physical node + phys_node = crosslinker.connection_sites[site_idx] + sibling_indices = [ + ci for ci in range(site_idx) + if crosslinker.connection_sites[ci] == phys_node + ] + + if sibling_indices: + # Same physical node as a previous connection site + all_beads = list(group) + for ci in sibling_indices: + all_beads.extend(selected[ci]) + all_positions = coords[all_beads].copy() + feasible, _ = _can_reach_all_beads( + all_positions, crosslink_bond_length, tolerance, + pbc, box_lengths, + ) + else: + positions = coords[group].copy() + feasible, _ = _can_reach_all_beads( + positions, crosslink_bond_length, tolerance, + pbc, box_lengths, + ) - if s < 1e-10: - if c > 0: - return np.eye(3, dtype=np.float32) - perp = np.array([1, 0, 0], dtype=np.float32) - if abs(np.dot(a, perp)) > 0.9: - perp = np.array([0, 1, 0], dtype=np.float32) - perp = perp - np.dot(perp, a) * a - perp /= np.linalg.norm(perp) - return (2 * np.outer(perp, perp) - np.eye(3)).astype(np.float32) + if feasible: + selected.append(group) + all_nodes |= set(group) + excluded |= _get_excluded_indices(path, group, excluded_bond_depth) + found = True + break - vx = np.array( - [[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]], dtype=np.float32 - ) - R = np.eye(3, dtype=np.float32) + vx + vx @ vx * ((1 - c) / (s * s)) - return R + if not found: + valid = False + break + if valid: + candidates.append(selected) + if len(candidates) >= 50: + break -def _random_rotation_matrix(rng): - """Uniform random rotation via QR decomposition.""" - H = rng.standard_normal((3, 3)).astype(np.float32) - Q, R = np.linalg.qr(H) - Q *= np.sign(np.diag(R)) - if np.linalg.det(Q) < 0: - Q[:, 0] *= -1 - return Q.astype(np.float32) - - -def _rotation_about_axis(axis, angle): - """Rotation matrix for rotation by `angle` radians about `axis`.""" - axis = np.asarray(axis, dtype=np.float32) - axis = axis / np.linalg.norm(axis) - K = np.array( - [ - [0, -axis[2], axis[1]], - [axis[2], 0, -axis[0]], - [-axis[1], axis[0], 0], - ], - dtype=np.float32, - ) - R = np.eye(3, dtype=np.float32) + np.sin(angle) * K + (1 - np.cos(angle)) * (K @ K) - return R + return candidates -def _initial_align_single( - replacement, - unique_active, - unique_conn_positions, - unique_targets, - neighbor_centroid, - rng, +# ============================================================================= +# Crosslinker Placement +# ============================================================================= + + +def _compute_optimal_placement( + crosslinker, candidate_group, path_coords, + crosslink_bond_length, pbc, box_lengths, ): - """Initial alignment for 1 unique connection site. + """Compute optimal crosslinker placement satisfying bond length constraints. - Returns base placement and rotation info for residual DOF. + For each PHYSICAL node in the crosslinker, collects all backbone beads + that bond to it (across possibly multiple connection sites) and finds + the position that minimizes bond length error. Returns ------- - base_placed : np.ndarray, shape (n_sites, 3) - perm : list - rotation_type : str - 'axis' for rotation about an axis, 'full' for full sphere - rotation_info : dict - Contains 'pivot' and 'axis' for residual rotation. + placed_coords : np.ndarray, shape (n_sites, 3) + Crosslinker coordinates after optimal placement. + target_node_positions : dict of {int: np.ndarray} + Computed position for each physical crosslinker node that has + external bonds. Keys are node indices in the crosslinker. """ - c = unique_conn_positions[0] - c_norm = np.linalg.norm(c) - - if c_norm < 1e-10: - R = _random_rotation_matrix(rng) - t = unique_targets[0] - base_placed = (R @ replacement.coordinates.T).T + t - return base_placed, [0], "full", {"pivot": unique_targets[0]} - - c_hat = c / c_norm - desired_dir = unique_targets[0] - neighbor_centroid - desired_norm = np.linalg.norm(desired_dir) - if desired_norm < 1e-10: - desired_dir = rng.standard_normal(3).astype(np.float32) - desired_norm = np.linalg.norm(desired_dir) - desired_hat = desired_dir / desired_norm - - R = _rotation_between_vectors(c_hat, desired_hat) - t = unique_targets[0] - R @ c - base_placed = (R @ replacement.coordinates.T).T + t - - # Residual DOF: rotation about axis from center to connection site - conn_placed = base_placed[unique_active[0]] - center = t # crosslinker center position - axis = conn_placed - center - axis_norm = np.linalg.norm(axis) - if axis_norm > 1e-10: - axis_hat = axis / axis_norm + conn_sites = crosslinker.connection_sites + + # --- Group backbone beads by physical crosslinker node --- + # node_to_beads[phys_node] = list of all backbone bead indices bonding to it + node_to_beads = {} + for conn_idx, phys_node in enumerate(conn_sites): + node_to_beads.setdefault(phys_node, []) + node_to_beads[phys_node].extend(candidate_group[conn_idx]) + + # --- For each physical node, find optimal position --- + # Use first group's first bead as PBC reference + reference = path_coords[candidate_group[0][0]].copy() + + target_node_positions = {} + for phys_node, bead_indices in node_to_beads.items(): + # Unwrap bead positions relative to reference + bead_positions = [] + for idx in bead_indices: + delta = _pbc_delta(path_coords[idx], reference, box_lengths, pbc) + bead_positions.append(reference + delta) + bead_positions = np.array(bead_positions, dtype=np.float64) + + # Find point equidistant (at crosslink_bond_length) from all beads + target = _find_equidistant_point( + bead_positions, crosslink_bond_length, pbc, box_lengths + ) + target_node_positions[phys_node] = target + + # --- Place the crosslinker based on computed node positions --- + unique_nodes = list(target_node_positions.keys()) + + if len(unique_nodes) == 1: + # Single physical node: just translate the crosslinker + node = unique_nodes[0] + node_offset = crosslinker.coordinates[node] + centroid_shift = target_node_positions[node] - node_offset + placed_coords = crosslinker.coordinates.astype(np.float64) + centroid_shift + + elif len(unique_nodes) >= 2: + # Multiple physical nodes: use Kabsch alignment + # Source: crosslinker node positions (relative coords) + source_points = np.array( + [crosslinker.coordinates[n] for n in unique_nodes], dtype=np.float64 + ) + # Target: computed positions (relative to their centroid, for alignment) + target_points = np.array( + [target_node_positions[n] for n in unique_nodes], dtype=np.float64 + ) + + # Align via Kabsch: find R, t such that R @ source + t ≈ target + source_centroid = source_points.mean(axis=0) + target_centroid = target_points.mean(axis=0) + + P = source_points - source_centroid + Q = target_points - target_centroid + + if len(unique_nodes) >= 2 and np.linalg.norm(P) > 1e-10: + R = _kabsch_rotation(P, Q) + else: + R = np.eye(3, dtype=np.float64) + + # Apply to all crosslinker coordinates + all_coords_centered = crosslinker.coordinates.astype(np.float64) - source_centroid + placed_coords = (R @ all_coords_centered.T).T + target_centroid else: - axis_hat = desired_hat + # No connections (shouldn't happen, but be safe) + placed_coords = crosslinker.coordinates.astype(np.float64) - return base_placed, [0], "axis", {"pivot": conn_placed, "axis": axis_hat} + return placed_coords.astype(np.float32), target_node_positions -def _initial_align_multi( - replacement, - unique_active, - unique_conn_positions, - unique_targets, - rng, -): - """Initial alignment for 2+ unique connection sites via Procrustes. +def _find_equidistant_point(bead_positions, target_distance, pbc, box_lengths): + """Find a point at target_distance from all given bead positions. + + Uses iterative projection onto the intersection of spheres of radius + target_distance centered at each bead. + + Parameters + ---------- + bead_positions : np.ndarray, shape (N, 3) + Positions of backbone beads (already PBC-unwrapped). + target_distance : float + Desired distance from the result to each bead. Returns ------- - base_placed : np.ndarray - perm : list - rotation_type : str - 'axis' for 2 sites, 'none' for 3+ - rotation_info : dict + point : np.ndarray, shape (3,) """ - n_unique = len(unique_active) + n = len(bead_positions) + bead_positions = np.asarray(bead_positions, dtype=np.float64) - best_R = np.eye(3, dtype=np.float32) - best_t = np.zeros(3, dtype=np.float32) - best_cost = np.inf - best_perm = list(range(n_unique)) + if n == 1: + # Any point at target_distance works; pick one perpendicular to z + return bead_positions[0] + np.array([target_distance, 0, 0], dtype=np.float64) - for perm in permutations(range(n_unique)): - perm_targets = unique_targets[list(perm)] + # Start from centroid of beads + centroid = bead_positions.mean(axis=0) - src_centroid = np.mean(unique_conn_positions, axis=0) - tgt_centroid = np.mean(perm_targets, axis=0) + # If centroid is equidistant from all beads, project outward to correct radius + # Otherwise, iterate + point = centroid.copy() - src_centered = unique_conn_positions - src_centroid - tgt_centered = perm_targets - tgt_centroid + for iteration in range(100): + # Compute gradient: move toward the surface of each sphere + shift = np.zeros(3, dtype=np.float64) + for i in range(n): + delta = point - bead_positions[i] + dist = np.linalg.norm(delta) + if dist < 1e-12: + # Point is at bead position; nudge it + delta = np.array([1e-6, 1e-6, 1e-6], dtype=np.float64) + dist = np.linalg.norm(delta) + # Project point onto sphere i: move to radius target_distance + desired = bead_positions[i] + delta * (target_distance / dist) + shift += (desired - point) + + shift /= n + point += shift + + if np.linalg.norm(shift) < target_distance * 1e-8: + break - R = _kabsch_rotation(src_centered, tgt_centered) - t = tgt_centroid - R @ src_centroid + return point.astype(np.float32) - placed_conn = (R @ unique_conn_positions.T).T + t - cost = np.sum((placed_conn - perm_targets) ** 2) - if cost < best_cost: - best_cost = cost - best_R = R - best_t = t - best_perm = list(perm) +# ============================================================================= +# Rotation Matrix Utilities +# ============================================================================= - base_placed = (best_R @ replacement.coordinates.T).T + best_t - if n_unique == 2: - # Residual DOF: rotation about axis connecting the two connection sites - conn0 = base_placed[unique_active[0]] - conn1 = base_placed[unique_active[1]] - axis = conn1 - conn0 - axis_norm = np.linalg.norm(axis) - if axis_norm > 1e-10: - axis_hat = axis / axis_norm - else: - axis_hat = np.array([0, 0, 1], dtype=np.float32) - midpoint = (conn0 + conn1) / 2.0 - return base_placed, best_perm, "axis", {"pivot": midpoint, "axis": axis_hat} - else: - # n_unique >= 3: fully determined, no residual DOF - return base_placed, best_perm, "none", {} +def _kabsch_rotation(P, Q): + """Optimal rotation aligning source P to target Q (both centered). + Returns R such that R @ P ≈ Q in least-squares sense. + """ + P = np.asarray(P, dtype=np.float64) + Q = np.asarray(Q, dtype=np.float64) + H = P.T @ Q + U, S, Vt = np.linalg.svd(H) + d = np.linalg.det(Vt.T @ U.T) + sign_matrix = np.diag([1.0, 1.0, np.sign(d) if abs(d) > 1e-10 else 1.0]) + R = Vt.T @ sign_matrix @ U.T + return R -def _check_bonds_valid( - placed, active_conn_sites, neighbor_coords, target_bond_lengths, tolerance -): - """Check if all connection-to-neighbor bonds are within tolerance. - Parameters - ---------- - placed : np.ndarray, shape (n_sites, 3) - active_conn_sites : list of int - neighbor_coords : np.ndarray, shape (degree, 3) - target_bond_lengths : np.ndarray, shape (degree,) - tolerance : float (fractional) +def _rotation_between_vectors(src, tgt): + """Rotation matrix taking unit vector src to unit vector tgt.""" + src = np.asarray(src, dtype=np.float64) + tgt = np.asarray(tgt, dtype=np.float64) + v = np.cross(src, tgt) + c = float(np.dot(src, tgt)) + if c < -0.9999: + perp = _get_perpendicular(src) + return (2 * np.outer(perp, perp) - np.eye(3)).astype(np.float32) + s = np.linalg.norm(v) + if s < 1e-10: + return np.eye(3, dtype=np.float32) + vx = np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]]) + R = np.eye(3) + vx + (vx @ vx) * (1.0 / (1.0 + c)) + return R.astype(np.float32) - Returns - ------- - valid : bool - """ - for i, site_idx in enumerate(active_conn_sites): - delta = placed[site_idx] - neighbor_coords[i] - actual = np.linalg.norm(delta) - target = target_bond_lengths[i] - if abs(actual - target) > target * tolerance: - return False - return True +def _rotate_around_axis(points, axis, angle): + """Rodrigues rotation of points around axis through the origin.""" + axis = np.asarray(axis, dtype=np.float64) + norm = np.linalg.norm(axis) + if norm < 1e-10: + return points.copy() + axis = axis / norm + cos_a = np.cos(angle) + sin_a = np.sin(angle) + K = np.array([ + [0, -axis[2], axis[1]], + [axis[2], 0, -axis[0]], + [-axis[1], axis[0], 0], + ]) + R = np.eye(3) * cos_a + sin_a * K + (1 - cos_a) * np.outer(axis, axis) + return (R @ np.asarray(points, dtype=np.float64).T).T.astype(np.float32) -def _build_assignment_from_perm( - active_conn_sites, unique_active, conn_to_target_indices, best_perm, degree -): - """Build full assignment list from permutation over unique sites. - Parameters - ---------- - active_conn_sites : list of int - Connection sites being used (length = degree). - unique_active : list of int - Unique site indices. - conn_to_target_indices : dict - site_idx -> [target indices using that site]. - best_perm : list of int - Permutation mapping unique site position -> unique target position. - degree : int - Number of connections being made. +# ============================================================================= +# Overlap Checking +# ============================================================================= - Returns - ------- - assignment : list of int - assignment[i] = neighbor_index for active_conn_sites[i]. - """ - if degree == len(unique_active) and all( - len(v) == 1 for v in conn_to_target_indices.values() - ): - # Simple: 1:1, no repeated sites - # best_perm[i] maps unique site i to unique target i - # But we need to map conn_idx -> neighbor_idx - # active_conn_sites[conn_idx] -> unique_active.index(that) -> perm -> target - assignment = [] - for conn_idx, site_idx in enumerate(active_conn_sites): - unique_pos = unique_active.index(site_idx) - assignment.append(best_perm[unique_pos]) - return assignment - - # General case with possible repeated sites - assignment = [0] * degree - site_counter = {} - - for conn_idx, site_idx in enumerate(active_conn_sites): - unique_pos = unique_active.index(site_idx) - target_group_pos = best_perm[unique_pos] - target_site = unique_active[target_group_pos] - target_indices = conn_to_target_indices[target_site] - - count = site_counter.get(site_idx, 0) - assignment[conn_idx] = target_indices[count % len(target_indices)] - site_counter[site_idx] = count + 1 - - return assignment +def _has_overlaps(placed_coords, all_coords, excluded_indices, min_separation, + pbc, box_lengths): + """Check if any placed bead overlaps with non-excluded existing beads.""" + if len(all_coords) == 0: + return False -def _get_pbc_info(volume_constraint): - """Extract PBC flags and box lengths from volume constraint.""" - if isinstance(volume_constraint, CuboidConstraint): - pbc = np.array(volume_constraint.pbc, dtype=bool) - box_lengths = volume_constraint.box_lengths.astype(np.float32) - elif isinstance(volume_constraint, CylinderConstraint): - pbc = np.array([False, False, volume_constraint.periodic_height], dtype=bool) - box_lengths = np.array( - [ - volume_constraint.radius * 2, - volume_constraint.radius * 2, - volume_constraint.height, - ], - dtype=np.float32, - ) - else: - pbc = np.array([False, False, False], dtype=bool) - box_lengths = np.array([np.inf, np.inf, np.inf], dtype=np.float32) - return pbc, box_lengths + mask = np.ones(len(all_coords), dtype=bool) + for idx in excluded_indices: + if 0 <= idx < len(all_coords): + mask[idx] = False + active_coords = all_coords[mask] + if len(active_coords) == 0: + return False -def _apply_mic(delta, pbc, box_lengths): - """Apply minimum image convention to a displacement vector.""" - delta = np.asarray(delta, dtype=np.float32) - for dim in range(3): - if pbc[dim]: - delta[dim] -= np.round(delta[dim] / box_lengths[dim]) * box_lengths[dim] - return delta + for point in placed_coords: + deltas = _pbc_delta(point[np.newaxis, :], active_coords, box_lengths, pbc) + distances_sq = np.sum(deltas ** 2, axis=1) + if np.any(distances_sq < min_separation ** 2): + return True + return False -def _wrap_positions(positions, pbc, box_lengths): - """Wrap positions into periodic box [-L/2, L/2].""" - positions = positions.copy() - for dim in range(3): - if pbc[dim]: - positions[:, dim] -= ( - np.round(positions[:, dim] / box_lengths[dim]) * box_lengths[dim] - ) - return positions +# ============================================================================= +# Bond Length Validation +# ============================================================================= -def _compute_ideal_backbone_distances(crosslinker, bond_length): - """Compute ideal pairwise distances between backbone targets given geometry and bond length. - For a crosslinker with connection sites at positions c_i (relative to center), - the ideal backbone position for site i is: - b_i = c_i + bond_length * (c_i / |c_i|) - i.e., the backbone sits along the same radial direction, at distance bond_length - beyond the connection site. +def _compute_max_bond_error( + placed_coords, crosslinker, candidate_group, path_coords, + crosslink_bond_length, pbc, box_lengths, +): + """Compute the maximum fractional bond length error for a placement. - Parameters - ---------- - crosslinker : CrosslinkerGeometry - The crosslinker geometry. - bond_length : float - Desired bond length from connection site to backbone bead. + Checks each connection site against each backbone bead in its group. Returns ------- - ideal_bb_positions : np.ndarray, shape (n_connections, 3) - Ideal backbone positions relative to crosslinker center. - pairwise_distances : np.ndarray, shape (n_connections, n_connections) - Matrix of ideal pairwise distances between backbone targets. - search_radius : float - Maximum distance from center to any ideal backbone position. + max_error : float + Maximum |actual_dist - target| / target across all bonds. """ conn_sites = crosslinker.connection_sites - conn_positions = crosslinker.coordinates[conn_sites] + max_error = 0.0 - # Compute ideal backbone positions - ideal_bb_positions = np.zeros_like(conn_positions) - for i in range(len(conn_sites)): - c_i = conn_positions[i] - c_norm = np.linalg.norm(c_i) - if c_norm > 1e-10: - # Backbone sits beyond connection site along radial direction - c_hat = c_i / c_norm - ideal_bb_positions[i] = c_i + bond_length * c_hat - else: - # Connection site at origin (single-site case): backbone at bond_length in any direction - ideal_bb_positions[i] = np.array([bond_length, 0, 0], dtype=np.float32) + for conn_idx, phys_node in enumerate(conn_sites): + site_pos = placed_coords[phys_node] + for bb_idx in candidate_group[conn_idx]: + dist = float(_pbc_distance(site_pos, path_coords[bb_idx], + box_lengths, pbc)) + error = abs(dist - crosslink_bond_length) / crosslink_bond_length + max_error = max(max_error, error) - # Maximum valid pairwise distances (geometric feasibility bound) - n = len(conn_sites) - pairwise_distances = np.zeros((n, n), dtype=np.float32) - for i in range(n): - for j in range(i + 1, n): - d_internal = np.linalg.norm(conn_positions[i] - conn_positions[j]) - max_d = d_internal + 2 * bond_length - pairwise_distances[i, j] = max_d - pairwise_distances[j, i] = max_d - - # Search radius: max distance from center to any backbone - search_radius = np.max(np.linalg.norm(ideal_bb_positions, axis=1)) - - return ideal_bb_positions, pairwise_distances, search_radius - - -def _unwrap_coords(coords, pbc, box_lengths): - """Unwrap coordinates relative to the first point.""" - unwrapped = np.empty_like(coords) - unwrapped[0] = coords[0] - ref = coords[0] - for k in range(1, len(coords)): - delta = coords[k] - ref - for dim in range(3): - if pbc[dim]: - delta[dim] -= np.round(delta[dim] / box_lengths[dim]) * box_lengths[dim] - unwrapped[k] = ref + delta - return unwrapped - - -def _validate_backbone_group( - selected_coords, - ideal_pairwise_distances, - bond_length, - tolerance, - pbc, - box_lengths, + return max_error + + +# ============================================================================= +# Placement with Rotation Search +# ============================================================================= + + +def _try_placement_with_rotations( + crosslinker, candidate_group, path_coords, + crosslink_bond_length, excluded_indices, pbc, box_lengths, + tolerance, min_separation, n_rotation_samples, rng, ): - """Check if backbone beads match expected geometry within tolerance. + """Attempt to place crosslinker with rotation search for overlap avoidance. + + 1. Computes the optimal placement (translation + alignment). + 2. If single-site or the base placement works, returns it. + 3. Otherwise, tries rotations around the centroid to avoid overlaps. Returns ------- - valid : bool - best_perm : list of int or None + placed_coords : np.ndarray or None + Final placed coordinates, or None if no valid placement found. """ - n = len(selected_coords) - if n <= 1: - return True, [0] if n == 1 else [] + # Compute base placement + base_coords, target_node_positions = _compute_optimal_placement( + crosslinker, candidate_group, path_coords, + crosslink_bond_length, pbc, box_lengths, + ) - abs_tol = bond_length * tolerance + # Check bond lengths on base placement + base_error = _compute_max_bond_error( + base_coords, crosslinker, candidate_group, path_coords, + crosslink_bond_length, pbc, box_lengths, + ) - # Compute actual pairwise distances (PBC-aware) - actual_distances = np.zeros((n, n), dtype=np.float32) - for i in range(n): - for j in range(i + 1, n): - delta = selected_coords[j] - selected_coords[i] - for dim in range(3): - if pbc[dim]: - delta[dim] -= ( - np.round(delta[dim] / box_lengths[dim]) * box_lengths[dim] - ) - d = np.linalg.norm(delta) - actual_distances[i, j] = d - actual_distances[j, i] = d + if base_error > tolerance: + # Base placement doesn't satisfy bond lengths — skip + return None - best_perm = None - best_cost = np.inf + # Check overlaps on base placement + if not _has_overlaps(base_coords, path_coords, excluded_indices, + min_separation, pbc, box_lengths): + return base_coords - for perm in permutations(range(n)): - valid = True - cost = 0.0 - for i in range(n): - for j in range(i + 1, n): - max_d = ideal_pairwise_distances[i, j] * (1 + tolerance) - actual_d = actual_distances[perm[i], perm[j]] - if actual_d > max_d: - valid = False - break - cost += actual_d - if not valid: - break + # --- Need rotational search to avoid overlaps --- + # For single-site crosslinkers, rotation doesn't help + if crosslinker.n_sites == 1: + return None + + # Determine rotation axis + unique_nodes = list(target_node_positions.keys()) + centroid = base_coords.mean(axis=0) + + if len(unique_nodes) >= 2: + # Rotate around axis connecting the physical nodes + positions = np.array([target_node_positions[n] for n in unique_nodes]) + axis = positions[-1] - positions[0] + elif len(unique_nodes) == 1: + # Rotate around the axis from centroid to the single physical node + axis = target_node_positions[unique_nodes[0]] - centroid + else: + axis = np.array([0, 0, 1], dtype=np.float32) - if valid and cost < best_cost: - best_cost = cost - best_perm = list(perm) + axis_norm = np.linalg.norm(axis) + if axis_norm < 1e-10: + axis = np.array([0, 0, 1], dtype=np.float32) + else: + axis = axis / axis_norm - return best_perm is not None, best_perm + # Rotation pivot: the physical node(s) must stay fixed + # Rotate around the axis through the centroid of physical nodes + if len(unique_nodes) >= 1: + pivot = np.mean( + [target_node_positions[n] for n in unique_nodes], axis=0 + ) + else: + pivot = centroid + best_coords = None + best_error = float("inf") -def _validate_clinker( - path, - crosslinker, - crosslink_bond_length, - selected_nodes, - overlap_radius=0.1, - n_samples=200, - pbc=None, - box_lengths=None, - seed=None, - tolerance=1e-5, -): - """Verify a point exists within constraints to place the centroid of the crosslinker. + angles = np.linspace(0, 2 * np.pi, n_rotation_samples, endpoint=False) + rng.shuffle(angles) - Treats the crosslinker as a cylindrical bead whose radius is the maximum - distance from its centroid to any constituent site. Samples candidate - center positions within the geometric feasibility region (intersection of - spheres around selected backbone nodes) and uses ``check_path`` to verify - no hard-sphere overlaps exist with the rest of the path. + for angle in angles: + # Rotate around axis through pivot + relative = base_coords - pivot + rotated = _rotate_around_axis(relative, axis, angle) + pivot - Parameters - ---------- - path : Path - The current path with all existing coordinates. - crosslinker : CrosslinkerGeometry - The crosslinker geometry to be placed. - crosslink_bond_length : float - Desired bond length from crosslinker connection sites to backbone beads. - selected_nodes : list of int - Backbone node indices that the crosslinker would connect to. - overlap_radius : float, default 0.1 - Additional exclusion radius beyond the crosslinker's own bead radius. - The total check radius passed to ``check_path`` is - ``bead_radius + overlap_radius``. - n_samples : int, default 200 - Number of candidate center positions to test. - pbc : np.ndarray of bool or None - Periodic boundary conditions per axis. If None, assumes no PBC. - box_lengths : np.ndarray or None - Box dimensions for minimum image convention. - seed : int or None - Random seed for reproducibility. - tolerance : float, default 1e-5 - Tolerance passed to ``check_path`` for rounding in distance checks. + # Verify bond lengths still valid after rotation + error = _compute_max_bond_error( + rotated, crosslinker, candidate_group, path_coords, + crosslink_bond_length, pbc, box_lengths, + ) + if error > tolerance: + continue - Returns - ------- - valid : bool - True if at least one overlap-free center position exists. - best_center : np.ndarray, shape (3,) or None - The first valid candidate center position found, or the best - (maximum clearance) candidate if none fully pass. - """ - from mbuild.path.path_utils import check_path + # Check overlaps + if not _has_overlaps(rotated, path_coords, excluded_indices, + min_separation, pbc, box_lengths): + if error < best_error: + best_error = error + best_coords = rotated.copy() + break # First overlap-free is good enough - if pbc is None: - pbc = np.array([False, False, False]) - if box_lengths is None: - box_lengths = np.array([0.0, 0.0, 0.0], dtype=np.float32) + return best_coords - rng = np.random.default_rng(seed) - # --- Compute crosslinker bead radius (cylindrical approximation) --- - crosslinker_centroid = np.mean(crosslinker.coordinates, axis=0) - site_offsets = crosslinker.coordinates - crosslinker_centroid - bead_radius = float(np.max(np.linalg.norm(site_offsets, axis=1))) +# ============================================================================= +# Path Insertion +# ============================================================================= - # --- Max distance from centroid to any connection site --- - conn_sites = crosslinker.connection_sites - conn_positions = crosslinker.coordinates[conn_sites] - conn_offsets = conn_positions - crosslinker_centroid - max_conn_radius = float(np.max(np.linalg.norm(conn_offsets, axis=1))) - - # --- Feasibility: centroid must be within this distance of each backbone node --- - feasibility_radius = crosslink_bond_length + max_conn_radius - - # --- Total hard-sphere check radius --- - check_radius = bead_radius + overlap_radius - - # --- Build existing_points array (all path coords except selected nodes) --- - n_total = len(path.coordinates) - selected_set = set(selected_nodes) - mask = np.ones(n_total, dtype=bool) - for node in selected_set: - if 0 <= node < n_total: - mask[node] = False - existing_points = np.ascontiguousarray(path.coordinates[mask], dtype=np.float32) - # existing_points = path.coordinates - - # --- Unwrap selected node coordinates relative to first --- - selected_coords = np.array( - [path.coordinates[n] for n in selected_nodes], dtype=np.float32 - ) - if len(selected_coords) > 1: - ref = selected_coords[0].copy() - for k in range(1, len(selected_coords)): - delta = selected_coords[k] - ref - for dim in range(3): - if pbc[dim]: - delta[dim] -= ( - np.round(delta[dim] / box_lengths[dim]) * box_lengths[dim] - ) - selected_coords[k] = ref + delta - centroid_of_selected = np.mean(selected_coords, axis=0) +def _insert_crosslinker(path, crosslinker, placed_coords, candidate_group): + """Insert crosslinker into path: add nodes, internal edges, and external bonds. - # --- Compute sampling sphere radius --- - max_dist_to_centroid = float( - np.max(np.linalg.norm(selected_coords - centroid_of_selected, axis=1)) - ) - sample_radius = feasibility_radius - max_dist_to_centroid - if sample_radius <= 0: - sample_radius = feasibility_radius * 0.05 - - # --- Generate candidate center positions --- - candidates = np.zeros((n_samples, 3), dtype=np.float32) - candidates[0] = centroid_of_selected - for i in range(1, n_samples): - r = sample_radius * (rng.random() ** (1.0 / 3.0)) - direction = rng.standard_normal(3).astype(np.float32) - norm = np.linalg.norm(direction) - if norm > 1e-10: - direction /= norm - else: - direction = np.array([1.0, 0.0, 0.0], dtype=np.float32) - candidates[i] = centroid_of_selected + r * direction - - # --- Evaluate each candidate --- - best_center = None - - for cand in candidates: - # Check feasibility: centroid within reach of all selected backbone nodes - feasible = True - for sc in selected_coords: - delta = cand - sc - for dim in range(3): - if pbc[dim]: - delta[dim] -= ( - np.round(delta[dim] / box_lengths[dim]) * box_lengths[dim] - ) - dist = np.linalg.norm(delta) - if dist > feasibility_radius: - feasible = False - break - if not feasible: - continue + Parameters + ---------- + path : Path + crosslinker : CrosslinkerGeometry + placed_coords : np.ndarray, shape (n_sites, 3) + candidate_group : list of list[int] + candidate_group[conn_idx] = list of backbone node indices for that + connection site. + """ + n_existing = len(path.coordinates) + n_new = crosslinker.n_sites + conn_sites = crosslinker.connection_sites - # Use check_path for hard-sphere overlap detection - new_point = np.ascontiguousarray(cand, dtype=np.float32) - no_overlap = check_path( - existing_points=existing_points, - new_point=new_point, - radius=check_radius, - tolerance=tolerance, + # Append coordinates and bead names + path.coordinates = np.vstack([path.coordinates, placed_coords]) + path.beads = np.append(path.beads, crosslinker.beads) + + # Add nodes + new_indices = list(range(n_existing, n_existing + n_new)) + for i, new_idx in enumerate(new_indices): + path.bond_graph.add_node( + new_idx, + name=str(crosslinker.beads[i]), + xyz=placed_coords[i], ) - if no_overlap: - return True, cand + # Add internal crosslinker edges + for u, v in crosslinker.bond_graph.edges(): + bond_type = (str(crosslinker.beads[u]), str(crosslinker.beads[v])) + path.bond_graph.add_edge(new_indices[u], new_indices[v], bond_type=bond_type) + + # Add external bonds: each connection site bonds to its backbone group + for conn_idx, phys_node in enumerate(conn_sites): + new_conn_node = new_indices[phys_node] + for bb_node in candidate_group[conn_idx]: + # Avoid duplicate edges (same physical node may appear multiple times) + if not path.bond_graph.has_edge(new_conn_node, bb_node): + bond_type = ( + str(crosslinker.beads[phys_node]), + str(path.beads[bb_node]), + ) + path.bond_graph.add_edge(new_conn_node, bb_node, bond_type=bond_type) - # Track first feasible candidate as fallback - if best_center is None: - best_center = cand.copy() - return False, best_center +# ============================================================================= +# Main Crosslink Function +# ============================================================================= -def _place_replacement_rigid_full( - replacement, - active_conn_sites, - neighbor_coords, - target_bond_lengths, - all_path_coords, - bonded_indices, - pbc, - box_lengths, - rng, - n_rotation_samples, - tolerance, +def crosslink( + path, + crosslinker=None, + bead_name="_R", + backbone_name="_B", + crosslink_bond_length=0.2, + max_backbone_degree=4, + tolerance=0.1, + excluded_bond_depth=2, + n_connection_sites=2, + volume_constraint=None, + initial_point=None, + seed=42, + n_rotation_samples=36, + overlap_radius=None, min_separation=None, ): - """Place replacement with full-system overlap checking. - - Unlike the radius-limited version, this checks placed sites against ALL - existing coordinates in the path (minus bonded neighbors). + """Place a crosslinker bonded to backbone beads, modifying path in place. Parameters ---------- - replacement : CrosslinkerGeometry - active_conn_sites : list of int - neighbor_coords : np.ndarray, shape (degree, 3) - target_bond_lengths : np.ndarray, shape (degree,) - all_path_coords : np.ndarray, shape (N, 3) - ALL coordinates in the current path. - bonded_indices : set of int - Indices to exclude from overlap checking. - pbc : array-like of bool - box_lengths : np.ndarray - rng : numpy.random.Generator - n_rotation_samples : int - tolerance : float - min_separation : float or None + path : Path + The Path object to modify in place. + crosslinker : CrosslinkerGeometry, optional + If None, auto-generated from bead_name and n_connection_sites. + bead_name : str, default "_R" + Name for auto-generated crosslinker beads. + backbone_name : str or tuple + Specifies what each connection site bonds to: + - String: every connection site bonds to one bead of that type. + - Tuple of length == n_connection_sites: + - String entry: bonds to one bead. + - Tuple entry (e.g., ("_B", "_B")): bonds to multiple neighboring beads. + + Examples:: + + backbone_name = "_B" + # All connection sites bond to one "_B" each. + + backbone_name = (("_B", "_B"), "_B") + # Site 0 bonds to two neighboring "_B" beads. + # Site 1 bonds to one "_B" bead. + + backbone_name = (("_B", "_B"), ("_B", "_B")) + # Both sites bond to two neighboring "_B" beads each. + # If connection_sites=[0,0], ALL four beads must be at + # crosslink_bond_length from the same physical node. + + crosslink_bond_length : float, default 0.2 + Desired distance from crosslinker connection node to each backbone bead. + max_backbone_degree : int, default 4 + Maximum graph degree a backbone node may have to remain eligible. + tolerance : float, default 0.1 + Fractional tolerance on bond lengths (0.1 = ±10%). + excluded_bond_depth : int, default 2 + Beads within this many bonds of selected nodes are excluded from + being selected as additional connection points. + n_connection_sites : int, default 2 + Used only when crosslinker is None. + volume_constraint : optional + Provides PBC information. + seed : int, default 42 + n_rotation_samples : int, default 36 + min_separation : float, optional + Minimum distance to existing beads. Defaults to crosslink_bond_length * 0.5. Returns ------- - placed_positions : np.ndarray, shape (n_sites, 3) - assignment : list of int + path : Path (modified in place) """ - degree = len(active_conn_sites) + rng = np.random.default_rng(seed + len(path.coordinates)) - if degree == 0: - center = ( - np.mean(neighbor_coords, axis=0) - if len(neighbor_coords) > 0 - else np.zeros(3) - ) - R = _random_rotation_matrix(rng) - placed = (R @ replacement.coordinates.T).T + center - return placed.astype(np.float32), [] - - # --- Compute target positions --- - neighbor_centroid = np.mean(neighbor_coords, axis=0) - target_positions = np.zeros((degree, 3), dtype=np.float32) - for i in range(degree): - direction = neighbor_centroid - neighbor_coords[i] - d = np.linalg.norm(direction) - if d > 1e-10: - direction_hat = direction / d + # --- Default crosslinker --- + if crosslinker is None: + if n_connection_sites == 2: + crosslinker = CrosslinkerGeometry.linear( + bond_length=crosslink_bond_length, bead_name=bead_name + ) + elif n_connection_sites == 3: + crosslinker = CrosslinkerGeometry.equilateral_triangle( + bond_length=crosslink_bond_length, bead_name=bead_name + ) else: - direction_hat = rng.standard_normal(3).astype(np.float32) - direction_hat /= np.linalg.norm(direction_hat) - target_positions[i] = ( - neighbor_coords[i] + direction_hat * target_bond_lengths[i] - ) + raise ValueError( + f"No default crosslinker for n_connection_sites={n_connection_sites}" + ) + + # --- Parse backbone specification --- + backbone_specs = _parse_backbone_spec(backbone_name, crosslinker.connection_sites) - # --- Initial alignment --- - unique_active = list(dict.fromkeys(active_conn_sites)) - unique_conn_positions = replacement.coordinates[unique_active] - n_unique = len(unique_active) - - conn_to_target_indices = {} - for i, site_idx in enumerate(active_conn_sites): - conn_to_target_indices.setdefault(site_idx, []).append(i) - - unique_targets = np.array( - [ - np.mean(target_positions[conn_to_target_indices[s]], axis=0) - for s in unique_active - ], - dtype=np.float32, + # --- PBC setup --- + pbc, box_lengths = _get_pbc_info(volume_constraint) + + # --- Min separation --- + if min_separation is None: + if overlap_radius is not None: + min_separation = overlap_radius + else: + min_separation = crosslink_bond_length * 0.5 + + # --- Find candidate backbone groups --- + candidate_groups = _find_candidate_groups( + path, crosslinker, backbone_specs, + crosslink_bond_length, tolerance, excluded_bond_depth, + max_backbone_degree, pbc, box_lengths, rng, ) - if n_unique == 1: - base_placed, base_perm, rotation_type, rotation_info = _initial_align_single( - replacement, - unique_active, - unique_conn_positions, - unique_targets, - neighbor_centroid, - rng, - ) - elif n_unique >= 2: - base_placed, base_perm, rotation_type, rotation_info = _initial_align_multi( - replacement, - unique_active, - unique_conn_positions, - unique_targets, - rng, - ) - else: - R = _random_rotation_matrix(rng) - placed = (R @ replacement.coordinates.T).T + neighbor_centroid - assignment = _build_assignment_from_perm( - active_conn_sites, unique_active, conn_to_target_indices, [0], degree + if not candidate_groups: + n_clinks = sum(1 for b in path.beads if b in set(crosslinker.beads)) + raise PathConvergenceError( + f"Could not find backbone beads matching crosslinker geometry with " + f"bond_length={crosslink_bond_length} (tolerance=±{tolerance * 100:.0f}%).\n" + f"max_backbone_degree={max_backbone_degree}, " + f"excluded_bond_depth={excluded_bond_depth}.\n" + f"Current crosslinks: {n_clinks}.\n" + "Ways to increase crosslinking:\n" + " - Increase crosslink_bond_length or tolerance\n" + " - Increase max_backbone_degree\n" + " - Decrease excluded_bond_depth\n" + " - Pack at higher density\n" ) - return placed.astype(np.float32), assignment - - # --- Precompute check coords (all minus bonded) --- - n_existing = len(all_path_coords) - check_mask = np.ones(n_existing, dtype=bool) - for idx in bonded_indices: - if 0 <= idx < n_existing: - check_mask[idx] = False - check_coords = all_path_coords[check_mask] - - # --- Sweep rotations with full overlap check --- - best_placed, best_min_dist, found_valid = _sweep_rotations_full( - base_placed, - rotation_type, - rotation_info, - active_conn_sites, - neighbor_coords, - target_bond_lengths, - check_coords, - pbc, - box_lengths, - tolerance, - min_separation, - n_rotation_samples, - rng, - ) - # --- If not valid, try offsets --- - if not found_valid and min_separation is not None and rotation_type == "axis": - best_placed, found_valid = _try_offsets_full( - base_placed, - rotation_info, - replacement, - active_conn_sites, - neighbor_coords, - target_bond_lengths, - check_coords, - pbc, - box_lengths, - tolerance, - min_separation, - n_rotation_samples, - rng, + # --- Try each candidate group --- + rng.shuffle(candidate_groups) + placed = False + + for candidate_group in candidate_groups: + # Collect all backbone nodes involved + all_backbone_nodes = _all_beads_from_groups(candidate_group) + + # Excluded indices for overlap check: backbone nodes + their immediate neighbors + excluded_indices = set(all_backbone_nodes) + for node in all_backbone_nodes: + if node in path.bond_graph: + for neighbor in path.bond_graph.neighbors(node): + excluded_indices.add(neighbor) + + # Attempt placement + placed_coords = _try_placement_with_rotations( + crosslinker, candidate_group, path.coordinates, + crosslink_bond_length, excluded_indices, pbc, box_lengths, + tolerance, min_separation, n_rotation_samples, rng, ) - # Build assignment - assignment = _build_assignment_from_perm( - active_conn_sites, unique_active, conn_to_target_indices, base_perm, degree - ) + if placed_coords is None: + continue - return best_placed.astype(np.float32), assignment + # --- Success: insert into path --- + _insert_crosslinker(path, crosslinker, placed_coords, candidate_group) + if volume_constraint is not None: + path.wrap_inside_box(volume_constraint) -def _sweep_rotations_full( - base_placed, - rotation_type, - rotation_info, - active_conn_sites, - neighbor_coords, - target_bond_lengths, - check_coords, - pbc, - box_lengths, - tolerance, - min_separation, - n_rotation_samples, - rng, -): - """Sweep rotations checking against full coordinate set. + placed = True + break - Returns - ------- - best_placed : np.ndarray - best_min_dist : float - found_valid : bool - """ - best_placed = base_placed.copy() - best_min_dist = _min_dist_pbc(base_placed, check_coords, pbc, box_lengths) - found_valid = min_separation is None or best_min_dist >= min_separation - - # Check if base placement bonds are valid - if not _check_bonds_valid( - base_placed, active_conn_sites, neighbor_coords, target_bond_lengths, tolerance - ): - best_min_dist = -np.inf - found_valid = False - - if rotation_type == "axis": - pivot = rotation_info["pivot"] - axis_hat = rotation_info["axis"] - angles = np.linspace(0, 2 * np.pi, n_rotation_samples, endpoint=False) - - for angle in angles: - R_rot = _rotation_about_axis(axis_hat, angle) - candidate = (R_rot @ (base_placed - pivot).T).T + pivot - - if not _check_bonds_valid( - candidate, - active_conn_sites, - neighbor_coords, - target_bond_lengths, - tolerance, - ): - continue + if not placed: + n_clinks = sum(1 for b in path.beads if b in set(crosslinker.beads)) + raise PathConvergenceError( + f"Could not place crosslinker without overlaps. " + f"Tried {len(candidate_groups)} candidate groups.\n" + f"Current crosslinks: {n_clinks}.\n" + "Consider increasing tolerance, n_rotation_samples, or decreasing min_separation." + ) - min_dist = _min_dist_pbc(candidate, check_coords, pbc, box_lengths) - - if min_separation is not None and min_dist >= min_separation: - if not found_valid or min_dist > best_min_dist: - best_placed = candidate.copy() - best_min_dist = min_dist - found_valid = True - elif not found_valid and min_dist > best_min_dist: - best_placed = candidate.copy() - best_min_dist = min_dist - - elif rotation_type == "full": - pivot = rotation_info.get("pivot", np.mean(base_placed, axis=0)) - for _ in range(n_rotation_samples): - R = _random_rotation_matrix(rng) - candidate = (R @ (base_placed - pivot).T).T + pivot - - if not _check_bonds_valid( - candidate, - active_conn_sites, - neighbor_coords, - target_bond_lengths, - tolerance, - ): - continue + return path - min_dist = _min_dist_pbc(candidate, check_coords, pbc, box_lengths) - if min_separation is not None and min_dist >= min_separation: - if not found_valid or min_dist > best_min_dist: - best_placed = candidate.copy() - best_min_dist = min_dist - found_valid = True - elif not found_valid and min_dist > best_min_dist: - best_placed = candidate.copy() - best_min_dist = min_dist +# ============================================================================= +# PBC Info Extraction +# ============================================================================= - return best_placed, best_min_dist, found_valid +def _get_pbc_info(volume_constraint): + """Extract PBC flags and box lengths from a volume constraint. -def _try_offsets_full( - base_placed, - rotation_info, - replacement, - active_conn_sites, - neighbor_coords, - target_bond_lengths, - check_coords, - pbc, - box_lengths, - tolerance, - min_separation, - n_rotation_samples, - rng, -): - """Try offset placements with full overlap checking.""" - axis_hat = rotation_info["axis"] - pivot = rotation_info["pivot"] - - perp1 = _get_perpendicular(axis_hat) - perp2 = np.cross(axis_hat, perp1).astype(np.float32) - perp2 /= np.linalg.norm(perp2) - - max_extent = float(np.max(np.linalg.norm(replacement.coordinates, axis=1))) - offset_magnitudes = np.linspace(0.05, max_extent * 1.5, 6) - - best_placed = base_placed.copy() - best_min_dist = _min_dist_pbc(base_placed, check_coords, pbc, box_lengths) - found_valid = False - - for offset_mag in offset_magnitudes: - for sign in [1.0, -1.0]: - for perp_dir in [perp1, perp2]: - offset = perp_dir * offset_mag * sign - - # Shift the whole placement - shifted = base_placed + offset - shifted_pivot = pivot + offset - - # Sweep rotations at this offset - angles = np.linspace(0, 2 * np.pi, n_rotation_samples, endpoint=False) - for angle in angles: - R_rot = _rotation_about_axis(axis_hat, angle) - candidate = (R_rot @ (shifted - shifted_pivot).T).T + shifted_pivot - - if not _check_bonds_valid( - candidate, - active_conn_sites, - neighbor_coords, - target_bond_lengths, - tolerance, - ): - continue - - min_dist = _min_dist_pbc(candidate, check_coords, pbc, box_lengths) - - if min_dist >= min_separation: - if not found_valid or min_dist > best_min_dist: - best_placed = candidate.copy() - best_min_dist = min_dist - found_valid = True - elif not found_valid and min_dist > best_min_dist: - best_placed = candidate.copy() - best_min_dist = min_dist - - if found_valid: - return best_placed, found_valid - - return best_placed, found_valid - - -def _min_dist_pbc(positions, check_coords, pbc, box_lengths): - """Minimum distance from any position to any check coord, with PBC.""" - if len(check_coords) == 0: - return np.inf - - global_min_sq = np.inf - for site_pos in positions: - deltas = check_coords - site_pos - for dim in range(3): - if pbc[dim]: - deltas[:, dim] -= ( - np.round(deltas[:, dim] / box_lengths[dim]) * box_lengths[dim] - ) - sq_dists = np.sum(deltas * deltas, axis=1) - local_min = float(np.min(sq_dists)) - if local_min < global_min_sq: - global_min_sq = local_min - - return float(np.sqrt(global_min_sq)) - - -def _get_perpendicular(axis): - """Get a unit vector perpendicular to axis.""" - axis = np.asarray(axis, dtype=np.float32) - perp = np.array([1, 0, 0], dtype=np.float32) - if abs(np.dot(axis, perp)) > 0.9: - perp = np.array([0, 1, 0], dtype=np.float32) - perp = perp - np.dot(perp, axis) * axis - perp /= np.linalg.norm(perp) - return perp + Returns + ------- + pbc : np.ndarray, shape (3,), dtype bool + box_lengths : np.ndarray, shape (3,), dtype float32 + """ + if volume_constraint is None: + return np.array([False, False, False], dtype=bool), np.zeros(3, dtype=np.float32) + + if isinstance(volume_constraint, CuboidConstraint): + pbc = np.array(volume_constraint.pbc, dtype=bool) + box_lengths = np.array(volume_constraint.box_lengths, dtype=np.float32) + elif isinstance(volume_constraint, CylinderConstraint): + pbc = np.array([False, False, getattr(volume_constraint, 'pbc_z', False)], dtype=bool) + box_lengths = np.array([0.0, 0.0, getattr(volume_constraint, 'length', 0.0)], dtype=np.float32) + elif hasattr(volume_constraint, 'pbc'): + pbc = np.array(volume_constraint.pbc, dtype=bool) + box_lengths = np.array([ + getattr(volume_constraint, 'Lx', 0), + getattr(volume_constraint, 'Ly', 0), + getattr(volume_constraint, 'Lz', 0), + ], dtype=np.float32) + else: + pbc = np.array([False, False, False], dtype=bool) + box_lengths = np.zeros(3, dtype=np.float32) + + return pbc, box_lengths \ No newline at end of file diff --git a/mbuild/path/formats.py b/mbuild/path/formats.py index 8ecaa1880..3f5bdf027 100644 --- a/mbuild/path/formats.py +++ b/mbuild/path/formats.py @@ -119,10 +119,16 @@ def to_mol3000(path, G=None): # Atom block lines.append("M V30 BEGIN ATOM\n") + atypesDict = dict() for i, (coord, bead_name) in enumerate(zip(path.coordinates, path.beads), start=1): atom_type = bead_name.strip("_") + if atom_type in atypesDict: + r_type = atypesDict[atom_type] + else: + r_type = len(atypesDict) + 1 + atypesDict[atom_type] = r_type lines.append( - f"M V30 {i} {atom_type} {coord[0]:.4f} {coord[1]:.4f} {coord[2]:.4f} 0\n" + f"M V30 {i} R{r_type:.0f} {coord[0]:.4f} {coord[1]:.4f} {coord[2]:.4f} 0\n" ) lines.append("M V30 END ATOM\n") @@ -139,4 +145,4 @@ def to_mol3000(path, G=None): # End of record lines.append("M END\n") - return "".join(lines) + return "".join(lines), atypesDict diff --git a/mbuild/path/points.py b/mbuild/path/points.py index cbc8f1cad..83a1c1efb 100644 --- a/mbuild/path/points.py +++ b/mbuild/path/points.py @@ -63,7 +63,7 @@ def get_second_point(state, existing_points, beads, check_path, next_step): ) if state.volume_constraint: is_inside_mask = state.volume_constraint.is_inside( - points=xyzs, buffer=state.radius + points=xyzs, buffer=0.0 ) xyzs = xyzs[is_inside_mask] @@ -72,7 +72,7 @@ def get_second_point(state, existing_points, beads, check_path, next_step): for xyz in xyzs: if check_path( - existing_points=existing_points, + existing_points=existing_points[:-1], new_point=xyz, radius=state.radius, tolerance=state.tolerance, @@ -174,7 +174,7 @@ def get_initial_point(state, existing_points, beads, check_path, next_step): ) if state.volume_constraint: is_inside_mask = state.volume_constraint.is_inside( - points=xyzs, buffer=state.radius + points=xyzs, buffer=0.0 ) xyzs = xyzs[is_inside_mask] @@ -220,7 +220,8 @@ def get_initial_point(state, existing_points, beads, check_path, next_step): # TODO: Use find_low_density_point here instead? elif state.volume_constraint: xyzs = state.volume_constraint.sample_candidates( - points=existing_points, n_candidates=300, buffer=state.radius + 0.1 + points=existing_points, n_candidates=300, buffer=state.radius, + rng=state.rng ) for xyz in xyzs: if check_path( diff --git a/mbuild/simulation.py b/mbuild/simulation.py index 45be85628..8b774666a 100644 --- a/mbuild/simulation.py +++ b/mbuild/simulation.py @@ -125,7 +125,7 @@ def _to_hoomd_snap_forces(self): apply(top, forcefields=self.forcefield, ignore_params=["dihedral", "improper"]) # Get hoomd snapshot and force objects forces, _ = gmso.external.to_hoomd_forcefield(top, r_cut=self.r_cut) - snap, _ = gmso.external.to_gsd_snapshot(top=top) + snap, _ = gmso.external.to_gsd_snapshot(top=top, shift_coords=True) forces = list(set().union(*forces.values())) return snap, forces diff --git a/mbuild/tests/test_crosslinks.py b/mbuild/tests/test_crosslinks.py index f4288e7a8..717a42f57 100644 --- a/mbuild/tests/test_crosslinks.py +++ b/mbuild/tests/test_crosslinks.py @@ -447,6 +447,25 @@ def test_failing_overlap_radius(self, long_linear_path): min_separation=0.2, ) + def test_multiple_connection_sites(self): + coordinates = np.array([[0,0,0],[1,1,0], [1,0,0], [0,1,0]]) + path = Path(coordinates, bead_name="_B") + path.bond_graph.add_edges_from([[0,2], [1,3]]) + cl = CrosslinkerGeometry( + bead_name="_CROSS", + connection_sites=[0, 0], + coordinates=np.array([[0,0,0]]) + ) + crosslink( + path, crosslinker=cl, backbone_name=(("_B","_B"),("_B","_B")), + crosslink_bond_length=np.sqrt(2)/2, + tolerance=0.01, seed=1 + ) + assert len(path) == 5 # 4 _B and one _CR + assert len(path.bond_graph.edges) == 6 # 4 _B-_CROSS bonds + assert np.allclose(path.coordinates[-1], np.array([0.5,0.5,0])) # CROSS Bead + assert np.linalg.norm(path.coordinates[-1] - path.coordinates[-2]) == np.sqrt(2)/2 + class TestReplaceSites(BaseTest): """Tests for the replace_sites function.""" diff --git a/mbuild/tests/test_path.py b/mbuild/tests/test_path.py index 2d7f7f171..dfd1db17c 100644 --- a/mbuild/tests/test_path.py +++ b/mbuild/tests/test_path.py @@ -751,6 +751,32 @@ def test_rw_namer_beads_in_bond_graph(self): node_names = [d["name"] for _, d in path.bond_graph.nodes(data=True)] assert node_names == ["_A", "_B", "_A", "_B"] + def test_consistent_behavior(self): + num_sites = NumSites(5) + conditions = Termination((num_sites, )) + constraint = CuboidConstraint(3.2,3.2,3.2, center=(0,0,0), pbc=(True, True, True)) + + path1 = hard_sphere_random_walk( + bead_name="_B", + radius=0.2, + bond_length=0.27, + termination=conditions, + volume_constraint=constraint, + rw_angles=(np.pi/2, np.pi), + seed=1 + ) + + path2 = hard_sphere_random_walk( + bead_name="_B", + radius=0.2, + bond_length=0.27, + termination=conditions, + volume_constraint=constraint, + rw_angles=(np.pi/2, np.pi), + seed=1 + ) + assert path1 == path2 + class TestPathUtils(BaseTest): def test_target_sq_distances_no_pbc(self): diff --git a/mbuild/utils/visualize.py b/mbuild/utils/visualize.py index 027a61c3a..7ad32d3ec 100644 --- a/mbuild/utils/visualize.py +++ b/mbuild/utils/visualize.py @@ -30,10 +30,6 @@ def visualize_path(path, radius=0.1, hide_periodic_bonds=False): G.remove_edges_from(remove_edges) view = py3Dmol.view(width=600, height=600) - # view.addModel(mol2_string, "mol2", keepH=True) - - # Get unique bead names - unique_names = list(dict.fromkeys(node for node in path.beads)) # Color palette colors = [ @@ -61,12 +57,13 @@ def visualize_path(path, radius=0.1, hide_periodic_bonds=False): "#000000", ] - data = path.to_mol3000(G) + data, unique_beadnames = path.to_mol3000(G) view = py3Dmol.view(data=data) - for i, name in enumerate(unique_names): + for i, name in enumerate(unique_beadnames.values()): color = colors[i % len(colors)] + label = "R" + str(name) view.addStyle( - {"elem": name.strip("_")}, # Select all atoms with this name + {"elem": label}, # Select all atoms with this name { "sphere": {"color": color, "radius": radius, "scale": 0.5}, "stick": {"radius": radius / 4, "color": "grey"}, From c6ce494915d19f9da08bc3c9a664f5b293211868 Mon Sep 17 00:00:00 2001 From: CalCraven Date: Mon, 15 Jun 2026 14:12:41 -0500 Subject: [PATCH 09/11] Fixes to creating Path from list of bead_names instead of array --- mbuild/path/build.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/mbuild/path/build.py b/mbuild/path/build.py index 1465bed67..5d16e889f 100644 --- a/mbuild/path/build.py +++ b/mbuild/path/build.py @@ -56,7 +56,7 @@ def __init__(self, coordinates=None, bond_graph=None, bead_name="_A"): if ( coordinates is not None and bond_graph is not None - and isinstance(bead_name, np.ndarray) + and isinstance(bead_name, (np.ndarray, list)) ): assert len(coordinates) == len(bond_graph), ( len(coordinates), @@ -68,7 +68,7 @@ def __init__(self, coordinates=None, bond_graph=None, bead_name="_A"): ) self.bond_graph = bond_graph self.coordinates = coordinates - self.beads = bead_name.astype("U10") + self.beads = np.array(bead_name, dtype="U10") # Passing in an array of coordinates, bond graph, and single bead name elif coordinates is not None and bond_graph is not None: assert len(coordinates) == len(bond_graph) @@ -95,8 +95,8 @@ def __init__(self, coordinates=None, bond_graph=None, bead_name="_A"): if isinstance(bead_name, str): self.beads = np.array([bead_name for _ in coordinates], dtype="U10") # Passed array of bead names, cast to U10 dtype - elif isinstance(bead_name, np.ndarray): - self.beads = bead_name.astype("U10") + elif isinstance(bead_name, (np.ndarray, list)): + self.beads = np.array(bead_name, dtype="U10") for idx in range(len((self.coordinates))): self.bond_graph.add_node(idx) # Nothing is defined, create empty place holders for coords, bond graph and bead names @@ -115,7 +115,7 @@ def __eq__(self, other): def __add__(self, other): coordinates = np.concat((self.coordinates, other.coordinates)) beads = np.concat((self.beads, other.beads)) - bond_graph = nx.compose( + bond_graph = nx.disjoint_union( self.bond_graph, other.bond_graph ) # TODO: Don't overwrite nodes in bg return Path(coordinates, bond_graph, beads) From 2919aabd7fceddc9e05d310cdf270ae5838fba1c Mon Sep 17 00:00:00 2001 From: CalCraven Date: Tue, 16 Jun 2026 09:49:06 -0500 Subject: [PATCH 10/11] Tweaking API for crosslinks, remove tests temporarily for replace sites --- mbuild/path/build.py | 22 +- mbuild/path/crosslink.py | 920 ++++++++++++++++++++++++-------- mbuild/tests/test_crosslinks.py | 484 +++++++++-------- 3 files changed, 938 insertions(+), 488 deletions(-) diff --git a/mbuild/path/build.py b/mbuild/path/build.py index 5d16e889f..75df06090 100644 --- a/mbuild/path/build.py +++ b/mbuild/path/build.py @@ -128,7 +128,8 @@ def from_compound(cls, compound): coordinates = compound.xyz # Create the path with coordinates and bond graph - path = cls(coordinates=coordinates, bead_name=compound.name) + bead_names = ["_" + part.name for part in compound.particles()] + path = cls(coordinates=coordinates, bead_name=bead_names) path.bond_graph = nx.Graph() # Ensure all nodes have xyz and name attributes @@ -624,13 +625,13 @@ def crosslink( def wrap_inside_box(self, constraint): """ Wrap coordinates inside a cubic box constraint with periodic boundary conditions. - + Args: coordinates: ndarray of shape (N, 3) with 3D coordinates box_length: length of the cubic box center: center of the box as tuple (x, y, z) pbc: tuple of booleans indicating PBC for each dimension (x, y, z) - + Returns: ndarray of wrapped coordinates with same shape as input """ @@ -638,21 +639,22 @@ def wrap_inside_box(self, constraint): wrapped = coordinates.copy() center = np.array(constraint.center) pbc = constraint.pbc - + # Calculate box bounds relative to center half_lengths = constraint.box_lengths / 2.0 box_min = center - half_lengths box_max = center + half_lengths - + # Apply periodic boundary conditions to each dimension for dim in range(3): if pbc[dim]: # Shift coordinates to box range and wrap using modulo shifted = wrapped[:, dim] - box_min[dim] wrapped[:, dim] = (shifted % constraint.box_lengths[dim]) + box_min[dim] - + self.coordinates = wrapped + def lamellar( path=None, num_layers=1, @@ -1396,9 +1398,7 @@ def hard_sphere_random_walk( ) # Create mask for particles inside volume constraint, allows for PBC if state.volume_constraint: - is_inside_mask = volume_constraint.is_inside( - points=candidates, buffer=0.0 - ) + is_inside_mask = volume_constraint.is_inside(points=candidates, buffer=0.0) candidates = candidates[is_inside_mask] # If there is a bias, sort candidates according to the bias if state.bias: @@ -1434,7 +1434,9 @@ def hard_sphere_random_walk( # Iterate through current state of candidates, break after first accept for xyz in candidates: if check_path_cpu( - existing_points=existing_points[:-1], # skip previously bonded to coordinate + existing_points=existing_points[ + :-1 + ], # skip previously bonded to coordinate new_point=xyz, radius=radius, tolerance=tolerance, diff --git a/mbuild/path/crosslink.py b/mbuild/path/crosslink.py index 8bb2d4780..5b22c9ab2 100644 --- a/mbuild/path/crosslink.py +++ b/mbuild/path/crosslink.py @@ -4,14 +4,13 @@ with proper overlap avoidance and bond length enforcement under PBC. """ -import numpy as np import networkx as nx +import numpy as np from mbuild.exceptions import PathConvergenceError from mbuild.path.build import Path from mbuild.path.constraints import CuboidConstraint, CylinderConstraint - # ============================================================================= # CrosslinkerGeometry # ============================================================================= @@ -33,13 +32,13 @@ class CrosslinkerGeometry(Path): May contain duplicates (same node bonds to multiple groups). """ - def __init__(self, coordinates, bond_graph=None, bead_name="_R", connection_sites=None): + def __init__( + self, coordinates, bond_graph=None, bead_name="_R", connection_sites=None + ): coordinates = np.asarray(coordinates, dtype=np.float32) centroid = coordinates.mean(axis=0) self.coordinates = coordinates - centroid - self.n_sites = len(self.coordinates) - if bond_graph is None: bond_graph = nx.Graph() for i in range(self.n_sites): @@ -55,12 +54,68 @@ def __init__(self, coordinates, bond_graph=None, bead_name="_R", connection_site connection_sites = list(range(self.n_sites)) self.connection_sites = list(connection_sites) + @property + def n_sites(self): + return len(self.coordinates) + @property def n_connections(self): return len(self.connection_sites) + @property + def unique_connection_sites(self): + return list(dict.fromkeys(self.connection_sites)) + + @property + def internal_bonds(self): + return list(self.bond_graph.edges()) + + def copy(self): + """Return a deep copy of this crosslinker geometry.""" + return CrosslinkerGeometry( + coordinates=self.coordinates.copy(), + bond_graph=self.bond_graph.copy(), + bead_name=self.beads.copy(), + connection_sites=list(self.connection_sites), + ) + + def recenter(self): + if len(self.coordinates) > 0: + centroid = np.mean(self.coordinates, axis=0) + self.coordinates = self.coordinates - centroid + + @classmethod + def from_path(cls, path, connection_sites): + """Create from an existing Path. + + Parameters + ---------- + path : Path + An existing path defining the structure. + connection_sites : list of int + Which nodes bond to external beads. + """ + return cls( + coordinates=path.coordinates.copy(), + bond_graph=path.bond_graph.copy(), + bead_name=path.beads.copy(), + connection_sites=connection_sites, + ) + + @classmethod + def single_site(cls, bead_name="_R", n_connections=2): + """Single-site (original crosslink behavior).""" + coords = np.array([[0.0, 0.0, 0.0]], dtype=np.float32) + return cls( + coordinates=coords, + bead_name=bead_name, + connection_sites=[0] * n_connections, + ) + @classmethod - def equilateral_triangle(cls, bond_length=0.27, bead_name="_R", connection_sites=None): + def equilateral_triangle( + cls, bond_length=0.27, bead_name="_R", connection_sites=None + ): """Three sites in a flat equilateral triangle with edge = bond_length.""" R = bond_length / np.sqrt(3) angles = [0, 2 * np.pi / 3, 4 * np.pi / 3] @@ -89,21 +144,192 @@ def equilateral_triangle(cls, bond_length=0.27, bead_name="_R", connection_sites ) @classmethod - def linear(cls, bond_length=0.27, bead_name="_R", connection_sites=None): - """Two sites separated by bond_length.""" + def linear(cls, n_sites=2, bond_length=0.27, bead_name="_R", connection_sites=None): + """Linear chain of sites.""" + positions = np.zeros((n_sites, 3), dtype=np.float32) + for i in range(n_sites): + positions[i, 0] = i * bond_length + + bead_names = ( + np.array([bead_name] * n_sites, dtype="U10") + if isinstance(bead_name, str) + else np.array(bead_name, dtype="U10") + ) + if connection_sites is None: + connection_sites = [0, n_sites - 1] + + G = nx.Graph() + for i in range(n_sites): + G.add_node(i) + for i in range(n_sites - 1): + G.add_edge(i, i + 1) + + return cls( + coordinates=positions, + bond_graph=G, + bead_name=bead_names, + connection_sites=connection_sites, + ) + + @classmethod + def square(cls, bond_length=0.27, bead_name="_R", connection_sites=None): + """Four sites in a flat square.""" + half = bond_length / 2.0 positions = np.array( - [[-bond_length / 2, 0, 0], [bond_length / 2, 0, 0]], + [[half, half, 0], [-half, half, 0], [-half, -half, 0], [half, -half, 0]], dtype=np.float32, ) + bead_names = ( + np.array([bead_name] * 4, dtype="U10") + if isinstance(bead_name, str) + else np.array(bead_name, dtype="U10") + ) if connection_sites is None: - connection_sites = [0, 1] + connection_sites = [0, 1, 2, 3] + G = nx.Graph() - G.add_node(0) - G.add_node(1) + for i in range(4): + G.add_node(i) + for i in range(4): + G.add_edge(i, (i + 1) % 4) + + return cls( + coordinates=positions, + bond_graph=G, + bead_name=bead_names, + connection_sites=connection_sites, + ) + + @classmethod + def tetrahedral(cls, bond_length=0.27, bead_name="_R", connection_sites=None): + """Four sites in a regular tetrahedron.""" + positions = np.array( + [[1, 1, 1], [1, -1, -1], [-1, 1, -1], [-1, -1, 1]], dtype=np.float32 + ) + current_dist = np.linalg.norm(positions[0] - positions[1]) + positions *= bond_length / current_dist + + bead_names = ( + np.array([bead_name] * 4, dtype="U10") + if isinstance(bead_name, str) + else np.array(bead_name, dtype="U10") + ) + if connection_sites is None: + connection_sites = [0, 1, 2, 3] + + G = nx.Graph() + for i in range(4): + G.add_node(i) + for i in range(4): + for j in range(i + 1, 4): + G.add_edge(i, j) + + return cls( + coordinates=positions, + bond_graph=G, + bead_name=bead_names, + connection_sites=connection_sites, + ) + + @classmethod + def trigonal_bipyramidal( + cls, bond_length=0.27, bead_name="_R", connection_sites=None + ): + """Five sites in trigonal bipyramidal arrangement (0-2 equatorial, 3-4 axial).""" + R_eq = bond_length / np.sqrt(3) + eq_angles = [0, 2 * np.pi / 3, 4 * np.pi / 3] + equatorial = [[R_eq * np.cos(a), R_eq * np.sin(a), 0.0] for a in eq_angles] + axial_dist = bond_length * np.sqrt(2.0 / 3.0) + axial = [[0.0, 0.0, axial_dist], [0.0, 0.0, -axial_dist]] + positions = np.array(equatorial + axial, dtype=np.float32) + + bead_names = ( + np.array([bead_name] * 5, dtype="U10") + if isinstance(bead_name, str) + else np.array(bead_name, dtype="U10") + ) + if connection_sites is None: + connection_sites = [0, 1, 2, 3, 4] + + G = nx.Graph() + for i in range(5): + G.add_node(i) G.add_edge(0, 1) + G.add_edge(1, 2) + G.add_edge(0, 2) + for eq in range(3): + G.add_edge(eq, 3) + G.add_edge(eq, 4) + + return cls( + coordinates=positions, + bond_graph=G, + bead_name=bead_names, + connection_sites=connection_sites, + ) + + @classmethod + def pentagon(cls, bond_length=0.27, bead_name="_R", connection_sites=None): + """Five sites in a flat regular pentagon.""" + R = bond_length / (2 * np.sin(np.pi / 5)) + angles = [2 * np.pi * i / 5 for i in range(5)] + positions = np.array( + [[R * np.cos(a), R * np.sin(a), 0.0] for a in angles], dtype=np.float32 + ) + + bead_names = ( + np.array([bead_name] * 5, dtype="U10") + if isinstance(bead_name, str) + else np.array(bead_name, dtype="U10") + ) + if connection_sites is None: + connection_sites = [0, 1, 2, 3, 4] + + G = nx.Graph() + for i in range(5): + G.add_node(i) + for i in range(5): + G.add_edge(i, (i + 1) % 5) + + return cls( + coordinates=positions, + bond_graph=G, + bead_name=bead_names, + connection_sites=connection_sites, + ) + + @classmethod + def from_edges(cls, coordinates, edges, bead_name="_R", connection_sites=None): + """Create from explicit positions and edge list. + + Parameters + ---------- + coordinates : array-like, shape (N, 3) + edges : list of tuple(int, int) + bead_name : str or list of str + connection_sites : list of int, optional + """ + coordinates = np.asarray(coordinates, dtype=np.float32) + n = len(coordinates) + bead_names = ( + np.array([bead_name] * n, dtype="U10") + if isinstance(bead_name, str) + else np.array(bead_name, dtype="U10") + ) + if connection_sites is None: + connection_sites = list(range(n)) + + G = nx.Graph() + for i in range(n): + G.add_node(i) + for i, j in edges: + G.add_edge(i, j) + return cls( - coordinates=positions, bond_graph=G, - bead_name=bead_name, connection_sites=connection_sites, + coordinates=coordinates, + bond_graph=G, + bead_name=bead_names, + connection_sites=connection_sites, ) @@ -166,7 +392,9 @@ def _parse_backbone_spec(backbone_name, connection_sites): return [[backbone_name]] * n_conn if not isinstance(backbone_name, (tuple, list)): - raise ValueError(f"backbone_name must be str or tuple, got {type(backbone_name)}") + raise ValueError( + f"backbone_name must be str or tuple, got {type(backbone_name)}" + ) if len(backbone_name) != n_conn: raise ValueError( @@ -255,11 +483,15 @@ def _find_neighbor_groups(path, target_names, max_backbone_degree): if key in seen: continue seen.add(key) - if (path.beads[i] == target_names[0] - and path.beads[j] == target_names[1]): + if ( + path.beads[i] == target_names[0] + and path.beads[j] == target_names[1] + ): groups.append([i, j]) - elif (path.beads[i] == target_names[1] - and path.beads[j] == target_names[0]): + elif ( + path.beads[i] == target_names[1] + and path.beads[j] == target_names[0] + ): groups.append([j, i]) return groups @@ -267,8 +499,9 @@ def _find_neighbor_groups(path, target_names, max_backbone_degree): groups = [] found_keys = set() for start in all_eligible: - _dfs_groups(path, start, all_eligible, target_names, n_needed, - groups, found_keys) + _dfs_groups( + path, start, all_eligible, target_names, n_needed, groups, found_keys + ) return groups @@ -300,62 +533,110 @@ def _all_beads_from_groups(candidate_group): def _can_reach_all_beads( - bead_positions, crosslink_bond_length, tolerance, pbc, box_lengths, + bead_positions, crosslink_bond_length, tolerance, pbc, box_lengths ): - """Check if there exists a point at crosslink_bond_length from all beads. + """Check if a point exists within tolerance of crosslink_bond_length from all beads.""" + n = len(bead_positions) + bead_positions = np.asarray(bead_positions, dtype=np.float64) - Uses the circumcenter/equidistant-point approach. A necessary condition - is that all pairwise distances are <= 2 * crosslink_bond_length. + if n == 0: + return True, np.zeros(3, dtype=np.float32) - For exact checking, computes the geometric median constraint. + if n == 1: + offset = np.array([0.0, 0.0, crosslink_bond_length], dtype=np.float64) + return True, (bead_positions[0] + offset).astype(np.float32) - Returns - ------- - feasible : bool - candidate_point : np.ndarray or None - An approximate valid point if feasible. - """ - n = len(bead_positions) - r = crosslink_bond_length - tol = tolerance - - # Necessary condition: all pairwise distances <= 2*r*(1+tol) - max_pair = 2 * r * (1 + tol) - for i in range(n): - for j in range(i + 1, n): - d = float(_pbc_distance(bead_positions[i], bead_positions[j], - box_lengths, pbc)) - if d > max_pair: - return False, None - - # Find the point equidistant from all beads via iterative projection - # Start from centroid - centroid = np.mean(bead_positions, axis=0) - point = centroid.copy().astype(np.float64) - - for _ in range(50): - shift = np.zeros(3, dtype=np.float64) - for i in range(n): - delta = _pbc_delta(bead_positions[i], point, box_lengths, pbc).astype(np.float64) - dist = np.linalg.norm(delta) - if dist > 1e-10: - # Project point onto sphere of radius r around bead i - target_on_sphere = point + delta * (1 - r / dist) - shift += (target_on_sphere - point) - shift /= n - point += shift * 0.8 # damped + centroid = bead_positions.mean(axis=0) + spread = bead_positions - centroid + principal = ( + spread[0] if np.linalg.norm(spread[0]) > 1e-10 else np.array([1.0, 0.0, 0.0]) + ) + perp = _get_perpendicular(principal).astype(np.float64) + + # Multiple starting points — centroid FIRST (finds minimum-radius in-plane solution) + starts = [ + centroid.copy(), + centroid + perp * crosslink_bond_length, + centroid - perp * crosslink_bond_length, + centroid + perp * crosslink_bond_length * 0.1, + ] + + best_point = None + best_radius_error = float("inf") + + for start in starts: + point = start.copy() + + # Equidistant solver: converge to point with equal distances to all beads + for iteration in range(300): + distances = np.array( + [np.linalg.norm(point - bead_positions[i]) for i in range(n)] + ) + mean_dist = distances.mean() + if mean_dist < 1e-12: + point += perp * 0.01 + continue + shift = np.zeros(3, dtype=np.float64) + for i in range(n): + delta = point - bead_positions[i] + dist = distances[i] + if dist < 1e-12: + delta = perp * 0.001 + dist = np.linalg.norm(delta) + desired = bead_positions[i] + delta * (mean_dist / dist) + shift += desired - point + shift /= n + point += shift + if np.linalg.norm(shift) < 1e-10: + break - if np.linalg.norm(shift) < r * 0.001: - break + # Check quality + distances = np.array( + [np.linalg.norm(point - bead_positions[i]) for i in range(n)] + ) + natural_radius = distances.mean() + max_spread = np.max(np.abs(distances - natural_radius)) - # Validate - max_error = 0.0 - for i in range(n): - d = float(_pbc_distance(point, bead_positions[i], box_lengths, pbc)) - max_error = max(max_error, abs(d - r) / r) + if max_spread / max(natural_radius, 1e-10) > 0.05: + continue # didn't converge to equidistant point + + radius_error = ( + abs(natural_radius - crosslink_bond_length) / crosslink_bond_length + ) + if radius_error < best_radius_error: + best_radius_error = radius_error + best_point = point.copy() + + # Also try fixed-radius solver from centroid + for start in starts[:2]: + point = start.copy() + for iteration in range(300): + shift = np.zeros(3, dtype=np.float64) + for i in range(n): + delta = point - bead_positions[i] + dist = np.linalg.norm(delta) + if dist < 1e-12: + delta = perp * 0.001 + dist = np.linalg.norm(delta) + desired = bead_positions[i] + delta * (crosslink_bond_length / dist) + shift += desired - point + shift /= n + point += shift + if np.linalg.norm(shift) < crosslink_bond_length * 1e-8: + break + + distances = np.array( + [np.linalg.norm(point - bead_positions[i]) for i in range(n)] + ) + max_err = ( + np.max(np.abs(distances - crosslink_bond_length)) / crosslink_bond_length + ) + if max_err < best_radius_error: + best_radius_error = max_err + best_point = point.copy() - if max_error <= tol: - return True, point.astype(np.float32) + if best_point is not None and best_radius_error <= tolerance: + return True, best_point.astype(np.float32) return False, None @@ -408,22 +689,43 @@ def _find_candidate_groups( # --- Search for valid combinations --- if n_conn == 2: return _pairwise_candidate_search( - path, crosslinker, eligible_per_site, node_to_conn_indices, - crosslink_bond_length, tolerance, excluded_bond_depth, - pbc, box_lengths, rng, + path, + crosslinker, + eligible_per_site, + node_to_conn_indices, + crosslink_bond_length, + tolerance, + excluded_bond_depth, + pbc, + box_lengths, + rng, ) return _general_candidate_search( - path, crosslinker, eligible_per_site, node_to_conn_indices, - crosslink_bond_length, tolerance, excluded_bond_depth, - pbc, box_lengths, rng, + path, + crosslinker, + eligible_per_site, + node_to_conn_indices, + crosslink_bond_length, + tolerance, + excluded_bond_depth, + pbc, + box_lengths, + rng, ) def _pairwise_candidate_search( - path, crosslinker, eligible_per_site, node_to_conn_indices, - crosslink_bond_length, tolerance, excluded_bond_depth, - pbc, box_lengths, rng, + path, + crosslinker, + eligible_per_site, + node_to_conn_indices, + crosslink_bond_length, + tolerance, + excluded_bond_depth, + pbc, + box_lengths, + rng, ): """Search for valid 2-connection-site candidates. @@ -439,7 +741,7 @@ def _pairwise_candidate_search( # Are both connection sites the same physical node? conn_sites = crosslinker.connection_sites - same_node = (conn_sites[0] == conn_sites[1]) + same_node = conn_sites[0] == conn_sites[1] for group_a in eligible_per_site[0]: excluded_a = _get_excluded_indices(path, group_a, excluded_bond_depth) @@ -461,15 +763,24 @@ def _pairwise_candidate_search( all_bead_indices = list(group_a) + list(group_b) all_bead_positions = coords[all_bead_indices].copy() feasible, _ = _can_reach_all_beads( - all_bead_positions, crosslink_bond_length, tolerance, - pbc, box_lengths, + all_bead_positions, + crosslink_bond_length, + tolerance, + pbc, + box_lengths, ) else: # Different physical nodes. Each group independently checked, # plus pairwise distance between connection sites must match geometry. feasible = _check_separate_sites_feasibility( - crosslinker, coords, group_a, group_b, - crosslink_bond_length, tolerance, pbc, box_lengths, + crosslinker, + coords, + group_a, + group_b, + crosslink_bond_length, + tolerance, + pbc, + box_lengths, ) if feasible: @@ -482,8 +793,14 @@ def _pairwise_candidate_search( def _check_separate_sites_feasibility( - crosslinker, coords, group_a, group_b, - crosslink_bond_length, tolerance, pbc, box_lengths, + crosslinker, + coords, + group_a, + group_b, + crosslink_bond_length, + tolerance, + pbc, + box_lengths, ): """Check feasibility when connection sites are different physical nodes. @@ -512,24 +829,38 @@ def _check_separate_sites_feasibility( if not ok_b: return False - # Check that the reach-points are at the correct internal distance apart - dist_ab = float(_pbc_distance(point_a, point_b, box_lengths, pbc)) - if abs(dist_ab - internal_distance) > internal_distance * tolerance + crosslink_bond_length * tolerance: + # Check geometric feasibility via triangle inequality. + # A rigid rod of length d (internal distance) must have one end on + # sphere(centroid_a, r) and the other on sphere(centroid_b, r). + # This is feasible when: (D-d)/2 <= r, i.e., h² = r² - ((D-d)/2)² >= 0 + centroid_a = np.mean(coords[group_a], axis=0).astype(np.float64) + centroid_b = np.mean(coords[group_b], axis=0).astype(np.float64) + D = float(_pbc_distance(centroid_a, centroid_b, box_lengths, pbc)) + r = crosslink_bond_length * (1 + tolerance) + + half_diff = abs(D - internal_distance) / 2.0 + if half_diff > r: return False return True def _general_candidate_search( - path, crosslinker, eligible_per_site, node_to_conn_indices, - crosslink_bond_length, tolerance, excluded_bond_depth, - pbc, box_lengths, rng, + path, + crosslinker, + eligible_per_site, + node_to_conn_indices, + crosslink_bond_length, + tolerance, + excluded_bond_depth, + pbc, + box_lengths, + rng, ): """General search for n-connection-site crosslinkers.""" coords = path.coordinates candidates = [] - max_tries = 5000 - + max_tries = 1000 n_conn = len(crosslinker.connection_sites) for _ in range(max_tries): @@ -537,57 +868,77 @@ def _general_candidate_search( group_0 = eligible_per_site[0][idx_0] excluded = _get_excluded_indices(path, group_0, excluded_bond_depth) selected = [group_0] + + # Track accumulated beads per physical node + node_to_beads = {} # phys_node_id -> list of bead indices + phys_0 = crosslinker.connection_sites[0] + node_to_beads[phys_0] = list(group_0) + all_nodes = set(group_0) valid = True for site_idx in range(1, n_conn): + phys_node = crosslinker.connection_sites[site_idx] found = False - for group in eligible_per_site[site_idx]: + + iterGroups = list(rng.permutation(eligible_per_site[site_idx])) + for group in iterGroups: if set(group) & (all_nodes | excluded): continue - # Check feasibility of adding this group - # Collect all beads assigned to same physical node - phys_node = crosslinker.connection_sites[site_idx] - sibling_indices = [ - ci for ci in range(site_idx) - if crosslinker.connection_sites[ci] == phys_node - ] - - if sibling_indices: - # Same physical node as a previous connection site - all_beads = list(group) - for ci in sibling_indices: - all_beads.extend(selected[ci]) - all_positions = coords[all_beads].copy() - feasible, _ = _can_reach_all_beads( - all_positions, crosslink_bond_length, tolerance, - pbc, box_lengths, - ) - else: - positions = coords[group].copy() - feasible, _ = _can_reach_all_beads( - positions, crosslink_bond_length, tolerance, - pbc, box_lengths, - ) - - if feasible: - selected.append(group) - all_nodes |= set(group) - excluded |= _get_excluded_indices(path, group, excluded_bond_depth) - found = True - break + # --- Accumulate ALL beads for this physical node --- + accumulated = node_to_beads.get(phys_node, []) + list(group) + accumulated_positions = coords[accumulated].copy() + + # --- Check if ALL accumulated beads are reachable from ONE point --- + feasible, _ = _can_reach_all_beads( + accumulated_positions, + crosslink_bond_length, + tolerance, + pbc, + box_lengths, + ) + if not feasible: + continue + + # --- Passed: update state --- + selected.append(group) + all_nodes.update(group) + node_to_beads.setdefault(phys_node, []).extend(group) + excluded |= _get_excluded_indices(path, group, excluded_bond_depth) + found = True + break if not found: valid = False + # print(f"Bad Candidate: {idx_0=}\t{group_0=}\t{selected=}") break if valid: candidates.append(selected) + # print(f"Good Candidate: {idx_0=}\t{group_0=}\t{selected=}") if len(candidates) >= 50: break - return candidates + # --- Deduplicate candidates --- + # Group indices by physical node, sort within each group, then use as hashable key + unique_candidates = [] + seen = set() + for candidate in candidates: + # Build normalized key: for each physical node, collect and sort all bead indices + node_to_sorted_beads = {} + for site_idx, group in enumerate(candidate): + phys_node = crosslinker.connection_sites[site_idx] + node_to_sorted_beads.setdefault(phys_node, []).extend(int(x) for x in group) + # Sort beads within each physical node, then sort by physical node + key = tuple( + tuple(sorted(beads)) for _, beads in sorted(node_to_sorted_beads.items()) + ) + if key not in seen: + seen.add(key) + unique_candidates.append(candidate) + + return unique_candidates # ============================================================================= @@ -596,8 +947,13 @@ def _general_candidate_search( def _compute_optimal_placement( - crosslinker, candidate_group, path_coords, - crosslink_bond_length, pbc, box_lengths, + crosslinker, + candidate_group, + path_coords, + crosslink_bond_length, + pbc, + box_lengths, + tolerance, ): """Compute optimal crosslinker placement satisfying bond length constraints. @@ -636,9 +992,14 @@ def _compute_optimal_placement( bead_positions = np.array(bead_positions, dtype=np.float64) # Find point equidistant (at crosslink_bond_length) from all beads - target = _find_equidistant_point( - bead_positions, crosslink_bond_length, pbc, box_lengths + # If impossible, find natural equidistant point within tolerance + _, target = _can_reach_all_beads( + bead_positions, crosslink_bond_length, tolerance, pbc, box_lengths ) + if target is None: + target = _find_equidistant_point( + bead_positions, crosslink_bond_length, pbc, box_lengths + ) target_node_positions[phys_node] = target # --- Place the crosslinker based on computed node positions --- @@ -650,33 +1011,80 @@ def _compute_optimal_placement( node_offset = crosslinker.coordinates[node] centroid_shift = target_node_positions[node] - node_offset placed_coords = crosslinker.coordinates.astype(np.float64) + centroid_shift - elif len(unique_nodes) >= 2: - # Multiple physical nodes: use Kabsch alignment - # Source: crosslinker node positions (relative coords) + # Get unwrapped bead centroids for each physical node + node_bead_centroids = {} + for phys_node, bead_indices in node_to_beads.items(): + bead_positions = [] + for idx in bead_indices: + delta = _pbc_delta(path_coords[idx], reference, box_lengths, pbc) + bead_positions.append(reference + delta) + node_bead_centroids[phys_node] = np.mean(bead_positions, axis=0).astype( + np.float64 + ) + + if len(unique_nodes) == 2: + n0, n1 = unique_nodes[0], unique_nodes[1] + bead_a = node_bead_centroids[n0] + bead_b = node_bead_centroids[n1] + + # Internal distance between connection sites (rigid) + conn_0_local = crosslinker.coordinates[n0].astype(np.float64) + conn_1_local = crosslinker.coordinates[n1].astype(np.float64) + d = float(np.linalg.norm(conn_1_local - conn_0_local)) + + # Backbone distance + D = float(np.linalg.norm(bead_b - bead_a)) + r = float(crosslink_bond_length) + + # Geometric solution: midpoint of targets is at midpoint of beads, + # offset perpendicular by h = sqrt(r² - ((D-d)/2)²) + half_diff = (D - d) / 2.0 + h_sq = r * r - half_diff * half_diff + h = np.sqrt(max(h_sq, 0.0)) + + # Coordinate frame + midpoint_bb = (bead_a + bead_b) / 2.0 + bb_axis = (bead_b - bead_a) / max(D, 1e-10) + perp = _get_perpendicular(bb_axis).astype(np.float64) + + # Target positions for the two connection sites + target_0 = midpoint_bb - (d / 2.0) * bb_axis + h * perp + target_1 = midpoint_bb + (d / 2.0) * bb_axis + h * perp + + target_node_positions[n0] = target_0.astype(np.float32) + target_node_positions[n1] = target_1.astype(np.float32) + + else: + # 3+ nodes: use independently computed targets (best effort) + pass # target_node_positions already computed above + + # Kabsch alignment: map crosslinker node positions → target positions source_points = np.array( - [crosslinker.coordinates[n] for n in unique_nodes], dtype=np.float64 + [crosslinker.coordinates[n].astype(np.float64) for n in unique_nodes] ) - # Target: computed positions (relative to their centroid, for alignment) target_points = np.array( - [target_node_positions[n] for n in unique_nodes], dtype=np.float64 + [target_node_positions[n].astype(np.float64) for n in unique_nodes] ) - # Align via Kabsch: find R, t such that R @ source + t ≈ target - source_centroid = source_points.mean(axis=0) - target_centroid = target_points.mean(axis=0) + src_centroid = source_points.mean(axis=0) + tgt_centroid = target_points.mean(axis=0) - P = source_points - source_centroid - Q = target_points - target_centroid + src_centered = source_points - src_centroid + tgt_centered = target_points - tgt_centroid - if len(unique_nodes) >= 2 and np.linalg.norm(P) > 1e-10: - R = _kabsch_rotation(P, Q) - else: - R = np.eye(3, dtype=np.float64) + H = src_centered.T @ tgt_centered + U, S, Vt = np.linalg.svd(H) + det = np.linalg.det(Vt.T @ U.T) + sign_matrix = np.diag([1.0, 1.0, np.sign(det) if abs(det) > 1e-10 else 1.0]) + R = Vt.T @ sign_matrix @ U.T + + all_centered = crosslinker.coordinates.astype(np.float64) - src_centroid + placed_coords = (all_centered @ R.T) + tgt_centroid - # Apply to all crosslinker coordinates - all_coords_centered = crosslinker.coordinates.astype(np.float64) - source_centroid - placed_coords = (R @ all_coords_centered.T).T + target_centroid + # Update target_node_positions to actual placed positions + for n in unique_nodes: + target_node_positions[n] = placed_coords[n].astype(np.float32) else: # No connections (shouldn't happen, but be safe) placed_coords = crosslinker.coordinates.astype(np.float64) @@ -705,15 +1113,16 @@ def _find_equidistant_point(bead_positions, target_distance, pbc, box_lengths): bead_positions = np.asarray(bead_positions, dtype=np.float64) if n == 1: - # Any point at target_distance works; pick one perpendicular to z - return bead_positions[0] + np.array([target_distance, 0, 0], dtype=np.float64) + offset = np.array([0.0, 0.0, target_distance], dtype=np.float64) + return (bead_positions[0] + offset).astype(np.float32) - # Start from centroid of beads centroid = bead_positions.mean(axis=0) - - # If centroid is equidistant from all beads, project outward to correct radius - # Otherwise, iterate - point = centroid.copy() + spread = bead_positions - centroid + principal = ( + spread[0] if np.linalg.norm(spread[0]) > 1e-10 else np.array([1.0, 0.0, 0.0]) + ) + perp = _get_perpendicular(principal).astype(np.float64) + point = centroid + perp * target_distance for iteration in range(100): # Compute gradient: move toward the surface of each sphere @@ -727,7 +1136,7 @@ def _find_equidistant_point(bead_positions, target_distance, pbc, box_lengths): dist = np.linalg.norm(delta) # Project point onto sphere i: move to radius target_distance desired = bead_positions[i] + delta * (target_distance / dist) - shift += (desired - point) + shift += desired - point shift /= n point += shift @@ -784,11 +1193,13 @@ def _rotate_around_axis(points, axis, angle): axis = axis / norm cos_a = np.cos(angle) sin_a = np.sin(angle) - K = np.array([ - [0, -axis[2], axis[1]], - [axis[2], 0, -axis[0]], - [-axis[1], axis[0], 0], - ]) + K = np.array( + [ + [0, -axis[2], axis[1]], + [axis[2], 0, -axis[0]], + [-axis[1], axis[0], 0], + ] + ) R = np.eye(3) * cos_a + sin_a * K + (1 - cos_a) * np.outer(axis, axis) return (R @ np.asarray(points, dtype=np.float64).T).T.astype(np.float32) @@ -798,8 +1209,9 @@ def _rotate_around_axis(points, axis, angle): # ============================================================================= -def _has_overlaps(placed_coords, all_coords, excluded_indices, min_separation, - pbc, box_lengths): +def _has_overlaps( + placed_coords, all_coords, excluded_indices, minimum_separation, pbc, box_lengths +): """Check if any placed bead overlaps with non-excluded existing beads.""" if len(all_coords) == 0: return False @@ -815,8 +1227,8 @@ def _has_overlaps(placed_coords, all_coords, excluded_indices, min_separation, for point in placed_coords: deltas = _pbc_delta(point[np.newaxis, :], active_coords, box_lengths, pbc) - distances_sq = np.sum(deltas ** 2, axis=1) - if np.any(distances_sq < min_separation ** 2): + distances_sq = np.sum(deltas**2, axis=1) + if np.any(distances_sq < minimum_separation**2): return True return False @@ -828,8 +1240,13 @@ def _has_overlaps(placed_coords, all_coords, excluded_indices, min_separation, def _compute_max_bond_error( - placed_coords, crosslinker, candidate_group, path_coords, - crosslink_bond_length, pbc, box_lengths, + placed_coords, + crosslinker, + candidate_group, + path_coords, + crosslink_bond_length, + pbc, + box_lengths, ): """Compute the maximum fractional bond length error for a placement. @@ -846,8 +1263,7 @@ def _compute_max_bond_error( for conn_idx, phys_node in enumerate(conn_sites): site_pos = placed_coords[phys_node] for bb_idx in candidate_group[conn_idx]: - dist = float(_pbc_distance(site_pos, path_coords[bb_idx], - box_lengths, pbc)) + dist = float(_pbc_distance(site_pos, path_coords[bb_idx], box_lengths, pbc)) error = abs(dist - crosslink_bond_length) / crosslink_bond_length max_error = max(max_error, error) @@ -860,9 +1276,17 @@ def _compute_max_bond_error( def _try_placement_with_rotations( - crosslinker, candidate_group, path_coords, - crosslink_bond_length, excluded_indices, pbc, box_lengths, - tolerance, min_separation, n_rotation_samples, rng, + crosslinker, + candidate_group, + path_coords, + crosslink_bond_length, + excluded_indices, + pbc, + box_lengths, + tolerance, + minimum_separation, + n_rotation_samples, + rng, ): """Attempt to place crosslinker with rotation search for overlap avoidance. @@ -877,14 +1301,24 @@ def _try_placement_with_rotations( """ # Compute base placement base_coords, target_node_positions = _compute_optimal_placement( - crosslinker, candidate_group, path_coords, - crosslink_bond_length, pbc, box_lengths, + crosslinker, + candidate_group, + path_coords, + crosslink_bond_length, + pbc, + box_lengths, + tolerance, ) # Check bond lengths on base placement base_error = _compute_max_bond_error( - base_coords, crosslinker, candidate_group, path_coords, - crosslink_bond_length, pbc, box_lengths, + base_coords, + crosslinker, + candidate_group, + path_coords, + crosslink_bond_length, + pbc, + box_lengths, ) if base_error > tolerance: @@ -892,8 +1326,9 @@ def _try_placement_with_rotations( return None # Check overlaps on base placement - if not _has_overlaps(base_coords, path_coords, excluded_indices, - min_separation, pbc, box_lengths): + if not _has_overlaps( + base_coords, path_coords, excluded_indices, minimum_separation, pbc, box_lengths + ): return base_coords # --- Need rotational search to avoid overlaps --- @@ -924,9 +1359,7 @@ def _try_placement_with_rotations( # Rotation pivot: the physical node(s) must stay fixed # Rotate around the axis through the centroid of physical nodes if len(unique_nodes) >= 1: - pivot = np.mean( - [target_node_positions[n] for n in unique_nodes], axis=0 - ) + pivot = np.mean([target_node_positions[n] for n in unique_nodes], axis=0) else: pivot = centroid @@ -943,15 +1376,21 @@ def _try_placement_with_rotations( # Verify bond lengths still valid after rotation error = _compute_max_bond_error( - rotated, crosslinker, candidate_group, path_coords, - crosslink_bond_length, pbc, box_lengths, + rotated, + crosslinker, + candidate_group, + path_coords, + crosslink_bond_length, + pbc, + box_lengths, ) if error > tolerance: continue # Check overlaps - if not _has_overlaps(rotated, path_coords, excluded_indices, - min_separation, pbc, box_lengths): + if not _has_overlaps( + rotated, path_coords, excluded_indices, minimum_separation, pbc, box_lengths + ): if error < best_error: best_error = error best_coords = rotated.copy() @@ -1020,19 +1459,15 @@ def _insert_crosslinker(path, crosslinker, placed_coords, candidate_group): def crosslink( path, crosslinker=None, - bead_name="_R", - backbone_name="_B", + backbone_name="_A", crosslink_bond_length=0.2, max_backbone_degree=4, tolerance=0.1, excluded_bond_depth=2, - n_connection_sites=2, volume_constraint=None, - initial_point=None, seed=42, n_rotation_samples=36, - overlap_radius=None, - min_separation=None, + minimum_separation=None, ): """Place a crosslinker bonded to backbone beads, modifying path in place. @@ -1042,8 +1477,6 @@ def crosslink( The Path object to modify in place. crosslinker : CrosslinkerGeometry, optional If None, auto-generated from bead_name and n_connection_sites. - bead_name : str, default "_R" - Name for auto-generated crosslinker beads. backbone_name : str or tuple Specifies what each connection site bonds to: - String: every connection site bonds to one bead of that type. @@ -1074,35 +1507,28 @@ def crosslink( excluded_bond_depth : int, default 2 Beads within this many bonds of selected nodes are excluded from being selected as additional connection points. - n_connection_sites : int, default 2 - Used only when crosslinker is None. volume_constraint : optional Provides PBC information. seed : int, default 42 n_rotation_samples : int, default 36 - min_separation : float, optional - Minimum distance to existing beads. Defaults to crosslink_bond_length * 0.5. + minimum_separation : float, optional + Minimum distance to existing beads. + Defaults to crosslink_bond_length / 4. + This will only apply to non-bonded beads, + i.e. you can set a separation larger than + 1/2 the crosslink_bond_length and still + achieve valid placements. Returns ------- path : Path (modified in place) + Path object with a single added crosslink. """ rng = np.random.default_rng(seed + len(path.coordinates)) # --- Default crosslinker --- - if crosslinker is None: - if n_connection_sites == 2: - crosslinker = CrosslinkerGeometry.linear( - bond_length=crosslink_bond_length, bead_name=bead_name - ) - elif n_connection_sites == 3: - crosslinker = CrosslinkerGeometry.equilateral_triangle( - bond_length=crosslink_bond_length, bead_name=bead_name - ) - else: - raise ValueError( - f"No default crosslinker for n_connection_sites={n_connection_sites}" - ) + if crosslinker is None: # Single two site crosslinker + crosslinker = CrosslinkerGeometry.single_site() # --- Parse backbone specification --- backbone_specs = _parse_backbone_spec(backbone_name, crosslinker.connection_sites) @@ -1111,21 +1537,30 @@ def crosslink( pbc, box_lengths = _get_pbc_info(volume_constraint) # --- Min separation --- - if min_separation is None: - if overlap_radius is not None: - min_separation = overlap_radius - else: - min_separation = crosslink_bond_length * 0.5 + if minimum_separation is None: + minimum_separation = crosslink_bond_length * 0.25 # --- Find candidate backbone groups --- candidate_groups = _find_candidate_groups( - path, crosslinker, backbone_specs, - crosslink_bond_length, tolerance, excluded_bond_depth, - max_backbone_degree, pbc, box_lengths, rng, + path, + crosslinker, + backbone_specs, + crosslink_bond_length, + tolerance, + excluded_bond_depth, + max_backbone_degree, + pbc, + box_lengths, + rng, ) if not candidate_groups: - n_clinks = sum(1 for b in path.beads if b in set(crosslinker.beads)) + n_beads_in_cl = len( + crosslinker.coordinates + ) # divisor to approx previous number of crosslinks + n_clinks = ( + sum(1 for b in path.beads if b in set(crosslinker.beads)) // n_beads_in_cl + ) raise PathConvergenceError( f"Could not find backbone beads matching crosslinker geometry with " f"bond_length={crosslink_bond_length} (tolerance=±{tolerance * 100:.0f}%).\n" @@ -1156,9 +1591,17 @@ def crosslink( # Attempt placement placed_coords = _try_placement_with_rotations( - crosslinker, candidate_group, path.coordinates, - crosslink_bond_length, excluded_indices, pbc, box_lengths, - tolerance, min_separation, n_rotation_samples, rng, + crosslinker, + candidate_group, + path.coordinates, + crosslink_bond_length, + excluded_indices, + pbc, + box_lengths, + tolerance, + minimum_separation, + n_rotation_samples, + rng, ) if placed_coords is None: @@ -1179,7 +1622,7 @@ def crosslink( f"Could not place crosslinker without overlaps. " f"Tried {len(candidate_groups)} candidate groups.\n" f"Current crosslinks: {n_clinks}.\n" - "Consider increasing tolerance, n_rotation_samples, or decreasing min_separation." + "Consider increasing tolerance, n_rotation_samples, or decreasing minimum_separation." ) return path @@ -1199,23 +1642,32 @@ def _get_pbc_info(volume_constraint): box_lengths : np.ndarray, shape (3,), dtype float32 """ if volume_constraint is None: - return np.array([False, False, False], dtype=bool), np.zeros(3, dtype=np.float32) + return np.array([False, False, False], dtype=bool), np.zeros( + 3, dtype=np.float32 + ) if isinstance(volume_constraint, CuboidConstraint): pbc = np.array(volume_constraint.pbc, dtype=bool) box_lengths = np.array(volume_constraint.box_lengths, dtype=np.float32) elif isinstance(volume_constraint, CylinderConstraint): - pbc = np.array([False, False, getattr(volume_constraint, 'pbc_z', False)], dtype=bool) - box_lengths = np.array([0.0, 0.0, getattr(volume_constraint, 'length', 0.0)], dtype=np.float32) - elif hasattr(volume_constraint, 'pbc'): + pbc = np.array( + [False, False, getattr(volume_constraint, "pbc_z", False)], dtype=bool + ) + box_lengths = np.array( + [0.0, 0.0, getattr(volume_constraint, "length", 0.0)], dtype=np.float32 + ) + elif hasattr(volume_constraint, "pbc"): pbc = np.array(volume_constraint.pbc, dtype=bool) - box_lengths = np.array([ - getattr(volume_constraint, 'Lx', 0), - getattr(volume_constraint, 'Ly', 0), - getattr(volume_constraint, 'Lz', 0), - ], dtype=np.float32) + box_lengths = np.array( + [ + getattr(volume_constraint, "Lx", 0), + getattr(volume_constraint, "Ly", 0), + getattr(volume_constraint, "Lz", 0), + ], + dtype=np.float32, + ) else: pbc = np.array([False, False, False], dtype=bool) box_lengths = np.zeros(3, dtype=np.float32) - return pbc, box_lengths \ No newline at end of file + return pbc, box_lengths diff --git a/mbuild/tests/test_crosslinks.py b/mbuild/tests/test_crosslinks.py index 717a42f57..5aad1080c 100644 --- a/mbuild/tests/test_crosslinks.py +++ b/mbuild/tests/test_crosslinks.py @@ -7,7 +7,6 @@ from mbuild.path.crosslink import ( CrosslinkerGeometry, crosslink, - replace_sites, ) from mbuild.tests.base_test import BaseTest @@ -30,19 +29,19 @@ def test_clink_equilateral(self, linear_path): assert outPath.bond_graph.has_edge(0, 5) assert outPath.bond_graph.has_edge(1, 4) # consistent edges form - def test_replace_site(self): - coords = np.column_stack( - (np.arange(0, 5), np.zeros(5, dtype=int), np.zeros(5, dtype=int)) - ) - path = Path(coords) - path._connect_edges("linear") - assert len(path.bond_graph.edges()) == 4 + # def test_replace_site(self): + # coords = np.column_stack( + # (np.arange(0, 5), np.zeros(5, dtype=int), np.zeros(5, dtype=int)) + # ) + # path = Path(coords) + # path._connect_edges("linear") + # assert len(path.bond_graph.edges()) == 4 - triangle = CrosslinkerGeometry.equilateral_triangle( - bond_length=0.4, connection_sites=[0, 1] - ) - path = path.replace_sites(triangle, 1) - assert len(path) == 7 + # triangle = CrosslinkerGeometry.equilateral_triangle( + # bond_length=0.4, connection_sites=[0, 1] + # ) + # path = path.replace_sites(triangle, 1) + # assert len(path) == 7 class TestCrosslinkerGeometry(BaseTest): @@ -192,7 +191,7 @@ def long_linear_path(self): return p def test_crosslink_single_site_default(self, linear_path): - out = crosslink(linear_path, crosslink_bond_length=3.0, n_connection_sites=2) + out = crosslink(linear_path, crosslink_bond_length=3.0) # Original 4 + 1 crosslink = 5 assert len(out.coordinates) == 5 # 3 backbone edges + 2 crosslink edges = 5 @@ -242,10 +241,10 @@ def test_crosslink_no_backbone_raises(self): def test_crosslink_radius_too_small_raises(self, linear_path): with pytest.raises(PathConvergenceError): - crosslink(linear_path, crosslink_bond_length=0.001, n_connection_sites=2) + crosslink(linear_path, crosslink_bond_length=0.001) def test_crosslink_preserves_backbone_edges(self, linear_path): - out = crosslink(linear_path, crosslink_bond_length=3.0, n_connection_sites=2) + out = crosslink(linear_path, crosslink_bond_length=3.0) # Backbone edges 0-1, 1-2, 2-3 should still exist assert out.bond_graph.has_edge(0, 1) assert out.bond_graph.has_edge(1, 2) @@ -253,7 +252,8 @@ def test_crosslink_preserves_backbone_edges(self, linear_path): def test_crosslink_bead_names_correct(self, linear_path): out = crosslink( - linear_path, bead_name="_R", crosslink_bond_length=3.0, n_connection_sites=2 + linear_path, + crosslink_bond_length=3.0, ) # Last bead should be the crosslinker assert out.beads[4] == "_R" @@ -262,42 +262,32 @@ def test_crosslink_bead_names_correct(self, linear_path): assert out.beads[i] == "_A" def test_crosslink_index_consistency(self, linear_path): - out = crosslink(linear_path, crosslink_bond_length=3.0, n_connection_sites=2) + out = crosslink(linear_path, crosslink_bond_length=3.0) # Every node in bond_graph should be a valid coordinate index for node in out.bond_graph.nodes(): assert 0 <= node < len(out.coordinates) - def test_crosslink_with_initial_point(self, long_linear_path): - # Should find crosslink near specified point - out = crosslink( - long_linear_path, - crosslink_bond_length=2.0, - n_connection_sites=2, - initial_point=5, - ) - assert len(out.coordinates) == 21 - - def test_clink_repeated_replaces(self, long_linear_path): - triangle = CrosslinkerGeometry.equilateral_triangle( - bond_length=0.27, connection_sites=[0, 1, 2], bead_name="_U" - ) - crosslink(long_linear_path, triangle, crosslink_bond_length=1) - crosslink(long_linear_path, triangle, crosslink_bond_length=1) - assert len(long_linear_path) == 26 - assert sum(long_linear_path.beads == "_U") == 6 - assert sum(long_linear_path.beads == "_A") == 20 - - replace_sites(long_linear_path, triangle, bead_name="_A") - assert len(long_linear_path) == 66 - assert sum(long_linear_path.beads == "_U") == 66 - - def test_error_on_excess_degree(self, long_linear_path): - triangle = CrosslinkerGeometry.equilateral_triangle( - bond_length=0.27, connection_sites=[0, 1], bead_name="_U" - ) - crosslink(long_linear_path, triangle, crosslink_bond_length=5) - with pytest.raises(ValueError, match="degree"): - replace_sites(long_linear_path, triangle, bead_name="_A") + # def test_clink_repeated_replaces(self, long_linear_path): + # triangle = CrosslinkerGeometry.equilateral_triangle( + # bond_length=0.27, connection_sites=[0, 1, 2], bead_name="_U" + # ) + # crosslink(long_linear_path, triangle, crosslink_bond_length=1) + # crosslink(long_linear_path, triangle, crosslink_bond_length=1) + # assert len(long_linear_path) == 26 + # assert sum(long_linear_path.beads == "_U") == 6 + # assert sum(long_linear_path.beads == "_A") == 20 + + # replace_sites(long_linear_path, triangle, bead_name="_A") + # assert len(long_linear_path) == 66 + # assert sum(long_linear_path.beads == "_U") == 66 + + # def test_error_on_excess_degree(self, long_linear_path): + # triangle = CrosslinkerGeometry.equilateral_triangle( + # bond_length=0.27, connection_sites=[0, 1], bead_name="_U" + # ) + # crosslink(long_linear_path, triangle, crosslink_bond_length=5) + # with pytest.raises(ValueError, match="degree"): + # replace_sites(long_linear_path, triangle, bead_name="_A") def test_non_centroid_candidate(self, long_linear_path): triangle = CrosslinkerGeometry.equilateral_triangle( @@ -333,7 +323,7 @@ def test_high_density_crosslinking_overlaps_linear(self): path, clink, excluded_bond_depth=2, - min_separation=min_sep, + minimum_separation=min_sep, crosslink_bond_length=bond_length, tolerance=0.01, max_backbone_degree=8, @@ -369,7 +359,7 @@ def test_high_density_crosslinking_overlaps_linear(self): assert min_dist >= min_sep, ( f"Internal crosslinker node {node} has overlap: " - f"min distance {min_dist:.4f} < min_separation {min_sep}" + f"min distance {min_dist:.4f} < minimum_separation {min_sep}" ) def test_high_density_crosslinking_overlaps_eqtriangle(self): @@ -394,9 +384,9 @@ def test_high_density_crosslinking_overlaps_eqtriangle(self): path, clink, excluded_bond_depth=2, - min_separation=min_sep, + minimum_separation=min_sep, crosslink_bond_length=bond_length, - tolerance=0.05, + tolerance=1, max_backbone_degree=8, ) @@ -428,7 +418,7 @@ def test_high_density_crosslinking_overlaps_eqtriangle(self): assert min_dist >= min_sep, ( f"Internal crosslinker node {node} has overlap: " - f"min distance {min_dist:.4f} < min_separation {min_sep}" + f"min distance {min_dist:.4f} < minimum_separation {min_sep}" ) def test_failing_overlap_radius(self, long_linear_path): @@ -444,203 +434,209 @@ def test_failing_overlap_radius(self, long_linear_path): long_linear_path, linear_clink, crosslink_bond_length=0.5, - min_separation=0.2, + minimum_separation=0.2, ) def test_multiple_connection_sites(self): - coordinates = np.array([[0,0,0],[1,1,0], [1,0,0], [0,1,0]]) + coordinates = np.array([[0, 0, 0], [1, 1, 0], [1, 0, 0], [0, 1, 0]]) path = Path(coordinates, bead_name="_B") - path.bond_graph.add_edges_from([[0,2], [1,3]]) + path.bond_graph.add_edges_from([[0, 2], [1, 3]]) cl = CrosslinkerGeometry( bead_name="_CROSS", connection_sites=[0, 0], - coordinates=np.array([[0,0,0]]) + coordinates=np.array([[0, 0, 0]]), ) crosslink( - path, crosslinker=cl, backbone_name=(("_B","_B"),("_B","_B")), - crosslink_bond_length=np.sqrt(2)/2, - tolerance=0.01, seed=1 - ) - assert len(path) == 5 # 4 _B and one _CR - assert len(path.bond_graph.edges) == 6 # 4 _B-_CROSS bonds - assert np.allclose(path.coordinates[-1], np.array([0.5,0.5,0])) # CROSS Bead - assert np.linalg.norm(path.coordinates[-1] - path.coordinates[-2]) == np.sqrt(2)/2 - - -class TestReplaceSites(BaseTest): - """Tests for the replace_sites function.""" - - @pytest.fixture - def chain_with_crosslinks(self): - """A linear backbone with two single-site crosslinks.""" - n = 10 - coords = np.column_stack( - (np.linspace(0, 5, n), np.zeros(n), np.zeros(n)) - ).astype(np.float32) - p = Path(coords, bead_name="_A") - p._connect_edges("linear") - # Add two crosslinks - p = crosslink( - p, bead_name="_R", crosslink_bond_length=3.0, n_connection_sites=2, seed=42 - ) - p = crosslink( - p, bead_name="_R", crosslink_bond_length=3.0, n_connection_sites=2, seed=99 - ) - return p - - def test_replace_single_site_with_triangle(self): - # Build a simple path with a degree-2 node in the middle - coords = np.array( - [[0, 0, 0], [1, 0, 0], [2, 0, 0], [3, 0, 0], [4, 0, 0]], - dtype=np.float32, - ) - p = Path(coords, bead_name="_A") - p._connect_edges("linear") - # Change middle node to _R - p.beads[2] = "_R" - - triangle = CrosslinkerGeometry.equilateral_triangle( - bond_length=0.27, connection_sites=[0, 1] - ) - out = replace_sites(p, triangle, bead_name="_R") - # 5 - 1 removed + 3 added = 7 - assert len(out.coordinates) == 7 - # Check all node indices are valid - for node in out.bond_graph.nodes(): - assert 0 <= node < len(out.coordinates) - - def test_replace_by_index(self): - coords = np.array( - [[0, 0, 0], [1, 0, 0], [2, 0, 0], [3, 0, 0]], dtype=np.float32 - ) - p = Path(coords, bead_name="_A") - p._connect_edges("linear") - - linear_repl = CrosslinkerGeometry.linear( - n_sites=2, bond_length=0.3, connection_sites=[0, 1] - ) - # Replace node 1 (degree 2) - out = replace_sites(p, linear_repl, sites=[1]) - # 4 - 1 + 2 = 5 - assert len(out.coordinates) == 5 + path, + crosslinker=cl, + backbone_name=(("_B", "_B"), ("_B", "_B")), + crosslink_bond_length=np.sqrt(2) / 2, + tolerance=0.01, + seed=1, + ) + assert len(path) == 5 # 4 _B and one _CR + assert len(path.bond_graph.edges) == 6 # 4 _B-_CROSS bonds + assert np.allclose(path.coordinates[-1], np.array([0.5, 0.5, 0])) # CROSS Bead assert ( - len(out.bond_graph.edges()) == 4 - ) # 2 kept + 1 internal + 1 external? let's just check > 0 - assert len(out.bond_graph.edges()) >= 4 - - def test_replace_preserves_nonreplaced_edges(self): - coords = np.array( - [[0, 0, 0], [1, 0, 0], [2, 0, 0], [3, 0, 0], [4, 0, 0]], - dtype=np.float32, - ) - p = Path(coords, bead_name="_A") - p._connect_edges("linear") - p.beads[2] = "_R" - - triangle = CrosslinkerGeometry.equilateral_triangle( - bond_length=0.27, connection_sites=[0, 1] - ) - out = replace_sites(p, triangle, bead_name="_R") - - # Original edge 0-1 and 3-4 should still exist in some form - # (nodes may be renumbered but connectivity is preserved) - # Check total number of edges: 2 kept backbone + 3 internal + 2 external = 7 - # Actually: edges 0-1, 1-2, 2-3, 3-4. Remove node 2. Kept: 0-1, 3-4. - # Triangle: 3 internal + 2 connections to (what was 1 and 3) - # Total: 2 + 3 + 2 = 7 - assert len(out.bond_graph.edges()) == 7 - - def test_replace_multiple_sites(self): - coords = np.array( - [[0, 0, 0], [1, 0, 0], [2, 0, 0], [3, 0, 0], [4, 0, 0], [5, 0, 0]], - dtype=np.float32, - ) - p = Path(coords, bead_name="_A") - p._connect_edges("linear") - p.beads[1] = "_R" - p.beads[4] = "_R" - - linear_repl = CrosslinkerGeometry.linear( - n_sites=2, bond_length=0.3, connection_sites=[0, 1] - ) - out = replace_sites(p, linear_repl, bead_name="_R") - # 6 - 2 + 4 = 8 - assert len(out.coordinates) == 8 - - def test_replace_no_sites_noop(self): - coords = np.array([[0, 0, 0], [1, 0, 0], [2, 0, 0]], dtype=np.float32) - p = Path(coords, bead_name="_A") - p._connect_edges("linear") - - triangle = CrosslinkerGeometry.equilateral_triangle( - bond_length=0.27, connection_sites=[0, 1] - ) - # No "_R" beads exist - out = replace_sites(p, triangle, bead_name="_R") - assert len(out.coordinates) == 3 - assert len(out.bond_graph.edges()) == 2 - - def test_replace_neither_sites_nor_bead_raises(self): - coords = np.array([[0, 0, 0], [1, 0, 0]], dtype=np.float32) - p = Path(coords, bead_name="_A") - triangle = CrosslinkerGeometry.equilateral_triangle(bond_length=0.27) - with pytest.raises(ValueError, match="Must specify"): - replace_sites(p, triangle) - - def test_replace_both_sites_and_bead_raises(self): - coords = np.array([[0, 0, 0], [1, 0, 0]], dtype=np.float32) - p = Path(coords, bead_name="_A") - triangle = CrosslinkerGeometry.equilateral_triangle(bond_length=0.27) - with pytest.raises(ValueError, match="Cannot specify both"): - replace_sites(p, triangle, sites=[0], bead_name="_A") - - def test_replace_renumbers_bond_graph(self): - coords = np.array( - [[0, 0, 0], [1, 0, 0], [2, 0, 0], [3, 0, 0]], dtype=np.float32 - ) - p = Path(coords, bead_name="_A") - p._connect_edges("linear") - p.beads[2] = "_R" - - linear_repl = CrosslinkerGeometry.linear( - n_sites=2, bond_length=0.3, connection_sites=[0, 1] - ) - out = replace_sites(p, linear_repl, bead_name="_R") - - # All node indices should be contiguous 0..N-1 - nodes = sorted(out.bond_graph.nodes()) - assert nodes == list(range(len(out.coordinates))) - - def test_replace_via_path_method(self): - coords = np.array( - [[0, 0, 0], [1, 0, 0], [2, 0, 0], [3, 0, 0]], dtype=np.float32 - ) - p = Path(coords, bead_name="_A") - p._connect_edges("linear") - p.beads[1] = "_R" - - linear_repl = CrosslinkerGeometry.linear( - n_sites=2, bond_length=0.3, connection_sites=[0, 1] - ) - out = p.replace_sites(linear_repl, bead_name="_R") - assert len(out.coordinates) == 5 - - def test_crosslink_then_replace(self, chain_with_crosslinks): - """Replace single-site crosslinks with triangles.""" - p = chain_with_crosslinks - n_R = sum(1 for b in p.beads if b == "_R") - # Each _R replaced: remove 1, add 3 -> net +2 per replacement - expected_len = len(p.coordinates) - n_R + n_R * 3 - assert n_R >= 1 - - triangle = CrosslinkerGeometry.equilateral_triangle( - bond_length=0.27, connection_sites=[0, 1] - ) - out = replace_sites(p, triangle, bead_name="_R") - assert len(out.coordinates) == expected_len - - # No _R beads remain (they've been replaced with triangle beads) - # Triangle beads are also "_R" so they should still exist - # But original single-site _R nodes are gone - for node in out.bond_graph.nodes(): - assert 0 <= node < len(out.coordinates) + np.linalg.norm(path.coordinates[-1] - path.coordinates[-2]) + == np.sqrt(2) / 2 + ) + + +# class TestReplaceSites(BaseTest): +# """Tests for the replace_sites function.""" + +# @pytest.fixture +# def chain_with_crosslinks(self): +# """A linear backbone with two single-site crosslinks.""" +# n = 10 +# coords = np.column_stack( +# (np.linspace(0, 5, n), np.zeros(n), np.zeros(n)) +# ).astype(np.float32) +# p = Path(coords, bead_name="_A") +# p._connect_edges("linear") +# # Add two crosslinks +# p = crosslink( +# p, crosslink_bond_length=3.0,seed=42 +# ) +# p = crosslink( +# p, crosslink_bond_length=3.0, seed=99 +# ) +# return p + +# def test_replace_single_site_with_triangle(self): +# # Build a simple path with a degree-2 node in the middle +# coords = np.array( +# [[0, 0, 0], [1, 0, 0], [2, 0, 0], [3, 0, 0], [4, 0, 0]], +# dtype=np.float32, +# ) +# p = Path(coords, bead_name="_A") +# p._connect_edges("linear") +# # Change middle node to _R +# p.beads[2] = "_R" + +# triangle = CrosslinkerGeometry.equilateral_triangle( +# bond_length=0.27, connection_sites=[0, 1] +# ) +# out = replace_sites(p, triangle, bead_name="_R") +# # 5 - 1 removed + 3 added = 7 +# assert len(out.coordinates) == 7 +# # Check all node indices are valid +# for node in out.bond_graph.nodes(): +# assert 0 <= node < len(out.coordinates) + +# def test_replace_by_index(self): +# coords = np.array( +# [[0, 0, 0], [1, 0, 0], [2, 0, 0], [3, 0, 0]], dtype=np.float32 +# ) +# p = Path(coords, bead_name="_A") +# p._connect_edges("linear") + +# linear_repl = CrosslinkerGeometry.linear( +# n_sites=2, bond_length=0.3, connection_sites=[0, 1] +# ) +# # Replace node 1 (degree 2) +# out = replace_sites(p, linear_repl, sites=[1]) +# # 4 - 1 + 2 = 5 +# assert len(out.coordinates) == 5 +# assert ( +# len(out.bond_graph.edges()) == 4 +# ) # 2 kept + 1 internal + 1 external? let's just check > 0 +# assert len(out.bond_graph.edges()) >= 4 + +# def test_replace_preserves_nonreplaced_edges(self): +# coords = np.array( +# [[0, 0, 0], [1, 0, 0], [2, 0, 0], [3, 0, 0], [4, 0, 0]], +# dtype=np.float32, +# ) +# p = Path(coords, bead_name="_A") +# p._connect_edges("linear") +# p.beads[2] = "_R" + +# triangle = CrosslinkerGeometry.equilateral_triangle( +# bond_length=0.27, connection_sites=[0, 1] +# ) +# out = replace_sites(p, triangle, bead_name="_R") + +# # Original edge 0-1 and 3-4 should still exist in some form +# # (nodes may be renumbered but connectivity is preserved) +# # Check total number of edges: 2 kept backbone + 3 internal + 2 external = 7 +# # Actually: edges 0-1, 1-2, 2-3, 3-4. Remove node 2. Kept: 0-1, 3-4. +# # Triangle: 3 internal + 2 connections to (what was 1 and 3) +# # Total: 2 + 3 + 2 = 7 +# assert len(out.bond_graph.edges()) == 7 + +# def test_replace_multiple_sites(self): +# coords = np.array( +# [[0, 0, 0], [1, 0, 0], [2, 0, 0], [3, 0, 0], [4, 0, 0], [5, 0, 0]], +# dtype=np.float32, +# ) +# p = Path(coords, bead_name="_A") +# p._connect_edges("linear") +# p.beads[1] = "_R" +# p.beads[4] = "_R" + +# linear_repl = CrosslinkerGeometry.linear( +# n_sites=2, bond_length=0.3, connection_sites=[0, 1] +# ) +# out = replace_sites(p, linear_repl, bead_name="_R") +# # 6 - 2 + 4 = 8 +# assert len(out.coordinates) == 8 + +# def test_replace_no_sites_noop(self): +# coords = np.array([[0, 0, 0], [1, 0, 0], [2, 0, 0]], dtype=np.float32) +# p = Path(coords, bead_name="_A") +# p._connect_edges("linear") + +# triangle = CrosslinkerGeometry.equilateral_triangle( +# bond_length=0.27, connection_sites=[0, 1] +# ) +# # No "_R" beads exist +# out = replace_sites(p, triangle, bead_name="_R") +# assert len(out.coordinates) == 3 +# assert len(out.bond_graph.edges()) == 2 + +# def test_replace_neither_sites_nor_bead_raises(self): +# coords = np.array([[0, 0, 0], [1, 0, 0]], dtype=np.float32) +# p = Path(coords, bead_name="_A") +# triangle = CrosslinkerGeometry.equilateral_triangle(bond_length=0.27) +# with pytest.raises(ValueError, match="Must specify"): +# replace_sites(p, triangle) + +# def test_replace_both_sites_and_bead_raises(self): +# coords = np.array([[0, 0, 0], [1, 0, 0]], dtype=np.float32) +# p = Path(coords, bead_name="_A") +# triangle = CrosslinkerGeometry.equilateral_triangle(bond_length=0.27) +# with pytest.raises(ValueError, match="Cannot specify both"): +# replace_sites(p, triangle, sites=[0], bead_name="_A") + +# def test_replace_renumbers_bond_graph(self): +# coords = np.array( +# [[0, 0, 0], [1, 0, 0], [2, 0, 0], [3, 0, 0]], dtype=np.float32 +# ) +# p = Path(coords, bead_name="_A") +# p._connect_edges("linear") +# p.beads[2] = "_R" + +# linear_repl = CrosslinkerGeometry.linear( +# n_sites=2, bond_length=0.3, connection_sites=[0, 1] +# ) +# out = replace_sites(p, linear_repl, bead_name="_R") + +# # All node indices should be contiguous 0..N-1 +# nodes = sorted(out.bond_graph.nodes()) +# assert nodes == list(range(len(out.coordinates))) + +# def test_replace_via_path_method(self): +# coords = np.array( +# [[0, 0, 0], [1, 0, 0], [2, 0, 0], [3, 0, 0]], dtype=np.float32 +# ) +# p = Path(coords, bead_name="_A") +# p._connect_edges("linear") +# p.beads[1] = "_R" + +# linear_repl = CrosslinkerGeometry.linear( +# n_sites=2, bond_length=0.3, connection_sites=[0, 1] +# ) +# out = p.replace_sites(linear_repl, bead_name="_R") +# assert len(out.coordinates) == 5 + +# def test_crosslink_then_replace(self, chain_with_crosslinks): +# """Replace single-site crosslinks with triangles.""" +# p = chain_with_crosslinks +# n_R = sum(1 for b in p.beads if b == "_R") +# # Each _R replaced: remove 1, add 3 -> net +2 per replacement +# expected_len = len(p.coordinates) - n_R + n_R * 3 +# assert n_R >= 1 + +# triangle = CrosslinkerGeometry.equilateral_triangle( +# bond_length=0.27, connection_sites=[0, 1] +# ) +# out = replace_sites(p, triangle, bead_name="_R") +# assert len(out.coordinates) == expected_len + +# # No _R beads remain (they've been replaced with triangle beads) +# # Triangle beads are also "_R" so they should still exist +# # But original single-site _R nodes are gone +# for node in out.bond_graph.nodes(): +# assert 0 <= node < len(out.coordinates) From 9ec2182ddcd81a9a91948feb5f3cfdee308a82bb Mon Sep 17 00:00:00 2001 From: CalCraven Date: Tue, 16 Jun 2026 10:10:18 -0500 Subject: [PATCH 11/11] Add height and width to path visualization --- mbuild/path/build.py | 4 ++-- mbuild/utils/visualize.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/mbuild/path/build.py b/mbuild/path/build.py index 75df06090..03a6dff5e 100644 --- a/mbuild/path/build.py +++ b/mbuild/path/build.py @@ -334,7 +334,7 @@ def to_mol3000(self, G=None): return _to_mol3000(self, G) - def visualize(self, radius, hide_periodic_bonds=False): + def visualize(self, radius, hide_periodic_bonds=False, width=600, height=600): """Visualize a path with py3Dmol. Parameters @@ -346,7 +346,7 @@ def visualize(self, radius, hide_periodic_bonds=False): """ from mbuild.utils.visualize import visualize_path - return visualize_path(self, radius, hide_periodic_bonds) + return visualize_path(self, radius, hide_periodic_bonds, width, height) def relax( self, radius, bonds=None, angles=None, box=None, steps=1000, seed=1, nthreads=1 diff --git a/mbuild/utils/visualize.py b/mbuild/utils/visualize.py index 7ad32d3ec..e5cfd16ec 100644 --- a/mbuild/utils/visualize.py +++ b/mbuild/utils/visualize.py @@ -7,7 +7,7 @@ from mbuild.utils.io import import_ -def visualize_path(path, radius=0.1, hide_periodic_bonds=False): +def visualize_path(path, radius=0.1, hide_periodic_bonds=False, width=600, height=600): """Visualize in 3D space using py3Dmol of the Path as a Compound. Parameters @@ -29,7 +29,7 @@ def visualize_path(path, radius=0.1, hide_periodic_bonds=False): print(f"Hiding {len(remove_edges)} periodic edges") G.remove_edges_from(remove_edges) - view = py3Dmol.view(width=600, height=600) + view = py3Dmol.view(width=width, height=height) # Color palette colors = [