Bead-Spring

AutoPoly.models.bead_spring

Bead-Spring Polymer Model Module

This module provides the BeadSpringPolymer class for generating coarse-grained bead-spring polymer models for molecular dynamics simulations in LAMMPS.

Features: - Multiple bead types (block copolymers) - Per-triplet angle stiffness - FENE bond support - Linear and ring topologies - Monte Carlo pre-equilibration for initial configurations

Classes:

Name	Description
BeadType	LJ parameters for a bead type.
AngleType	Angle parameters for a bead triplet.
SAWConfig	Configuration for Self-Avoiding Random Walk generation.
MCConfig	Monte Carlo equilibration configuration.
BeadSpringPolymer	Bead-spring polymer model generator for LAMMPS simulations.

Functions:

Name	Description
calculate_box_size	Calculate cubic box size from number of beads and target density.
compute_lj_energy	Compute total LJ energy for non-bonded pairs.
compute_lj_energy_vectorized	Vectorized LJ energy using numpy broadcasting.
compute_local_lj_energy	Compute LJ energy contribution from one bead to all others - O(N).
compute_local_bond_energy	Compute bond energy for bonds involving a specific bead - O(degree).
compute_bond_energy	Compute harmonic bond energy.
compute_bond_energy_vectorized	Vectorized bond energy calculation.
compute_total_energy	Compute total system energy for Metropolis criterion.
build_exclusion_structures	Build exclusion data structures for efficient energy calculations.
metropolis_accept	Metropolis acceptance criterion.
mc_single_bead_displacement	Displace a single bead by random vector.
mc_crankshaft_move	Rotate beads between i and j around the i-j axis.
mc_pivot_move	Rotate arm (pivot_idx+1 to chain_end) around pivot bead.
mc_reptation_move	Remove bead from one end (tail) and attach at other end (head).
place_chains_in_box	Place multiple chains randomly in periodic box.
mc_chain_translation	Translate entire chain by random displacement with PBC.
mc_chain_rotation	Rotate entire chain around its center of mass.
saw_grow_chain	Grow a single polymer chain using Self-Avoiding Random Walk with backtracking.
saw_generate_multi_chain	Generate multiple polymer chains using SAW.
mc_equilibrate	Pre-equilibrate polymer configuration using MC moves.

BeadType dataclass

LJ parameters for a bead type.

Source code in AutoPoly/models/bead_spring.py

@dataclass
class BeadType:
    """LJ parameters for a bead type."""
    name: str
    mass: float = 1.0
    epsilon: float = 1.0
    sigma: float = 1.0

AngleType dataclass

Angle parameters for a bead triplet.

Source code in AutoPoly/models/bead_spring.py

@dataclass
class AngleType:
    """Angle parameters for a bead triplet."""
    triplet: Tuple[str, str, str]
    k: float = 10.0
    theta0: float = 180.0

SAWConfig dataclass

Configuration for Self-Avoiding Random Walk generation.

Source code in AutoPoly/models/bead_spring.py

@dataclass
class SAWConfig:
    """Configuration for Self-Avoiding Random Walk generation."""
    collision_sigma: float = 1.0          # Bead diameter for collision
    collision_tolerance: float = 0.1      # Overlap tolerance
    n_trials: int = 50                    # Trial positions per bead
    max_backtrack_depth: int = 10         # Max beads to remove when stuck
    max_total_backtracks: int = 1000      # Total backtrack budget per chain
    bond_angle_min: float = 60.0          # Min bond angle (degrees)
    bond_angle_max: float = 180.0         # Max bond angle
    ring_closure_trials: int = 100        # Extra trials for ring closure
    ring_closure_tolerance: float = 0.2   # Distance tolerance for ring closure

MCConfig dataclass

Monte Carlo equilibration configuration.

Source code in AutoPoly/models/bead_spring.py

@dataclass
class MCConfig:
    """Monte Carlo equilibration configuration."""
    # Box sizing
    density: float = 0.85           # Bead density (beads per sigma^3)

    # MC move parameters
    max_displacement: float = 0.5   # Max single bead displacement
    max_angle: float = 0.3          # Max rotation angle (radians)
    temperature: float = 1.0        # Reduced temperature for Metropolis

    # Energy parameters
    lj_epsilon: float = 1.0         # LJ energy
    lj_sigma: float = 1.0           # LJ length
    lj_cutoff: float = 2.5          # LJ cutoff in sigma units
    bond_k: float = 100.0           # Bond spring constant
    bond_tolerance: float = 0.3     # Acceptable bond stretch

    # Equilibration settings
    n_steps: int = 10000            # Number of MC steps
    move_weights: Optional[Dict[str, float]] = None  # Move type probabilities

BeadSpringPolymer

Bead-spring polymer model generator for LAMMPS simulations.

Supports multiple bead types for block copolymers, per-triplet angle stiffness, and both harmonic and FENE bond styles.

Methods:

Name	Description
equilibrate	Pre-equilibrate the polymer configuration using Monte Carlo moves.
saw_generate	Generate configuration using Self-Avoiding Random Walk.
generate_data_file	Generate LAMMPS data file for bead-spring polymer.
get_system_info	Get comprehensive information about the bead-spring polymer system.
kremer_grest	Create a Kremer-Grest bead-spring polymer with standard parameters.

Attributes:

Name	Type	Description
n_beads	int	Number of beads per chain (derived from sequence).

Source code in AutoPoly/models/bead_spring.py

class BeadSpringPolymer:
    """
    Bead-spring polymer model generator for LAMMPS simulations.

    Supports multiple bead types for block copolymers, per-triplet angle
    stiffness, and both harmonic and FENE bond styles.
    """

    VALID_TOPOLOGIES = ["linear", "ring"]
    VALID_BOND_STYLES = ["harmonic", "fene"]
    VALID_PAIR_STYLES = ["lj", "wca"]
    VALID_GENERATION_METHODS = ["geometric", "saw", "mc"]

    def __init__(
        self,
        name: str,
        system: object,
        n_chains: int,
        bead_types: List[BeadType],
        sequence: Union[List[str], List[Tuple[str, int]], str],
        topology: str = "linear",
        # Bond parameters
        bond_length: float = 1.0,
        bond_style: str = "harmonic",
        k_bond: float = 30.0,
        fene_r0: float = 1.5,
        # Pair style parameters
        pair_style: str = "lj",  # "lj" (full LJ 2.5) or "wca" (WCA 2^(1/6))
        # Angle parameters
        use_angles: bool = False,
        default_k_angle: float = 10.0,
        default_theta0: float = 180.0,
        angle_types: Optional[List[AngleType]] = None,
        # Box sizing
        density: Optional[float] = None,
        box_size: Optional[float] = None,
        # Generation method
        generation_method: str = "saw",  # "geometric", "saw", or "mc"
        saw_config: Optional[SAWConfig] = None,
        # MC equilibration (legacy, use generation_method="mc" instead)
        equilibrate: bool = False,
        mc_config: Optional[MCConfig] = None,
    ) -> None:
        """
        Initialize bead-spring polymer generator.

        Args:
            name: Name for output files.
            system: System object containing path information.
            n_chains: Number of polymer chains.
            bead_types: List of BeadType objects defining bead parameters.
            sequence: Chain structure as block pattern, explicit list, or string.
            topology: "linear" or "ring".
            bond_length: Equilibrium bond length.
            bond_style: "harmonic" or "fene".
            k_bond: Bond force constant.
            fene_r0: FENE maximum extension (only used if bond_style="fene").
            pair_style: Pair interaction style - "lj" for full LJ (cutoff 2.5) or
                "wca" for WCA purely repulsive (cutoff 2^(1/6)*sigma ≈ 1.12246).
                For Kremer-Grest melts at P=0, use "wca" (recommended).
            use_angles: Whether to include angle potentials.
            default_k_angle: Default angle force constant for unspecified triplets.
            default_theta0: Default equilibrium angle for unspecified triplets.
            angle_types: List of AngleType objects for specific triplet parameters.
            density: Target bead density (beads/sigma^3). If set, box size is calculated.
            box_size: Explicit box size. Overrides density if both are set.
            generation_method: Method for generating initial configurations:
                - "geometric": Simple geometric placement (fast, may have overlaps)
                - "saw": Self-Avoiding Random Walk (fast, overlap-free)
                - "mc": Monte Carlo equilibration (slow, equilibrated)
            saw_config: SAW configuration. Uses defaults if None.
            equilibrate: Whether to run MC equilibration (legacy, use generation_method="mc").
            mc_config: Monte Carlo configuration. Uses defaults if None.

        Raises:
            ValueError: If parameters are invalid.

        Note on FENE Bonds:
            The FENE potential has NO explicit r0 (equilibrium bond length) parameter.
            The equilibrium distance of ~0.97 for Kremer-Grest polymers EMERGES from
            the balance of FENE (attractive, wants r→0) + WCA (repulsive, prevents r<~1.0).
            LAMMPS FENE syntax: bond_coeff * K R0 epsilon sigma
        """
        if topology not in self.VALID_TOPOLOGIES:
            raise ValueError(f"Topology must be one of: {self.VALID_TOPOLOGIES}")
        if bond_style not in self.VALID_BOND_STYLES:
            raise ValueError(f"Bond style must be one of: {self.VALID_BOND_STYLES}")
        if generation_method not in self.VALID_GENERATION_METHODS:
            raise ValueError(f"Generation method must be one of: {self.VALID_GENERATION_METHODS}")
        if pair_style not in self.VALID_PAIR_STYLES:
            raise ValueError(f"Pair style must be one of: {self.VALID_PAIR_STYLES}")
        if not bead_types:
            raise ValueError("At least one bead type is required")

        self.name = name
        self.system = system
        self.path = f"{self.system.get_folder_path()}/{self.name}" if system else f"./{name}"
        self.n_chains = n_chains
        self.topology = topology

        # Store bead types
        self.bead_types = bead_types
        self._bead_type_map: Dict[str, BeadType] = {bt.name: bt for bt in bead_types}
        self._bead_type_id: Dict[str, int] = {bt.name: i + 1 for i, bt in enumerate(bead_types)}

        # Parse and validate sequence
        self._sequence = self._parse_sequence(sequence)
        self._validate()

        # Bond parameters
        self.bond_length = bond_length
        self.bond_style = bond_style
        self.k_bond = k_bond
        self.fene_r0 = fene_r0

        # Pair style parameter
        self._pair_style = pair_style

        # Angle parameters
        self.use_angles = use_angles
        self.default_k_angle = default_k_angle
        self.default_theta0 = default_theta0
        self._angle_types = angle_types or []

        # Box sizing parameters
        self.density = density
        self._explicit_box_size = box_size

        # Generation method parameters
        self._generation_method = generation_method
        self._saw_config = saw_config

        # MC equilibration parameters (legacy support)
        self._equilibrate = equilibrate
        self._mc_config = mc_config

        # If equilibrate=True is set, use mc generation method
        if equilibrate and generation_method == "geometric":
            self._generation_method = "mc"

        # Internal state for positions (populated during generate_data_file)
        self._positions: Optional[List[np.ndarray]] = None
        self._chain_indices: Optional[List[Tuple[int, int]]] = None
        self._bonds: Optional[List[Tuple[int, int]]] = None

        # Build internal maps
        self._pair_coeffs = self._build_pair_coeffs()
        if self.use_angles:
            self._angle_type_map = self._build_angle_type_map()

        # Create output directory
        Path(self.path).mkdir(parents=True, exist_ok=True)

        logger.info(
            f"Initialized bead-spring polymer: {n_chains} chains, "
            f"{len(self._sequence)} beads each, {topology} topology, "
            f"{len(bead_types)} bead type(s)"
        )

    @property
    def n_beads(self) -> int:
        """Number of beads per chain (derived from sequence)."""
        return len(self._sequence)

    def _parse_sequence(self, seq: Union[List[str], List[Tuple[str, int]], str]) -> List[str]:
        """
        Convert sequence input to explicit bead type list.

        Args:
            seq: Sequence as block pattern, explicit list, or string.

        Returns:
            List of bead type names.
        """
        if isinstance(seq, str):
            # String format: "AABB" -> ["A", "A", "B", "B"]
            return list(seq)

        result = []
        for item in seq:
            if isinstance(item, tuple):
                # Block format: ("A", 20) -> ["A"] * 20
                bead_name, count = item
                result.extend([bead_name] * count)
            else:
                # Explicit list format
                result.append(item)
        return result

    def _validate(self) -> None:
        """Validate bead types and sequence consistency."""
        if not self._sequence:
            raise ValueError("Sequence cannot be empty")

        for bead_name in self._sequence:
            if bead_name not in self._bead_type_map:
                raise ValueError(
                    f"Unknown bead type '{bead_name}' in sequence. "
                    f"Available types: {list(self._bead_type_map.keys())}"
                )

    def _get_canonical_triplet(self, t1: str, t2: str, t3: str) -> Tuple[str, str, str]:
        """
        Return canonical triplet (smaller endpoint first alphabetically).

        Args:
            t1, t2, t3: Bead type names.

        Returns:
            Canonical triplet tuple.
        """
        if t1 <= t3:
            return (t1, t2, t3)
        return (t3, t2, t1)

    def _build_pair_coeffs(self) -> Dict[Tuple[int, int], Tuple[float, float]]:
        """
        Calculate Lorentz-Berthelot mixing for all bead type pairs.

        Returns:
            Dict mapping (type_i, type_j) to (epsilon_ij, sigma_ij).
        """
        coeffs = {}
        n_types = len(self.bead_types)

        for i in range(n_types):
            for j in range(i, n_types):
                bt_i = self.bead_types[i]
                bt_j = self.bead_types[j]

                sigma_ij = (bt_i.sigma + bt_j.sigma) / 2
                epsilon_ij = math.sqrt(bt_i.epsilon * bt_j.epsilon)

                # Store with 1-based indices
                coeffs[(i + 1, j + 1)] = (epsilon_ij, sigma_ij)

        return coeffs

    def _build_angle_type_map(self) -> Dict[Tuple[str, str, str], int]:
        """
        Map canonical triplets to LAMMPS type IDs.

        Returns:
            Dict mapping canonical triplet to angle type ID.
        """
        # Collect all unique canonical triplets from the sequence
        unique_triplets = set()

        for i in range(len(self._sequence) - 2):
            triplet = self._get_canonical_triplet(
                self._sequence[i],
                self._sequence[i + 1],
                self._sequence[i + 2]
            )
            unique_triplets.add(triplet)

        # For ring topology, add wrap-around triplets
        if self.topology == "ring" and len(self._sequence) >= 3:
            # Last two beads + first bead
            triplet = self._get_canonical_triplet(
                self._sequence[-2],
                self._sequence[-1],
                self._sequence[0]
            )
            unique_triplets.add(triplet)
            # Last bead + first two beads
            triplet = self._get_canonical_triplet(
                self._sequence[-1],
                self._sequence[0],
                self._sequence[1]
            )
            unique_triplets.add(triplet)

        # Sort triplets for consistent ordering
        sorted_triplets = sorted(unique_triplets)

        return {triplet: i + 1 for i, triplet in enumerate(sorted_triplets)}

    def _get_angle_params(self, triplet: Tuple[str, str, str]) -> Tuple[float, float]:
        """
        Return (k, theta0) for a triplet, using default if not specified.

        Args:
            triplet: Canonical triplet tuple.

        Returns:
            Tuple of (k, theta0).
        """
        for angle_type in self._angle_types:
            canonical = self._get_canonical_triplet(*angle_type.triplet)
            if canonical == triplet:
                return (angle_type.k, angle_type.theta0)
        return (self.default_k_angle, self.default_theta0)

    def _calculate_box_size(self) -> float:
        """
        Calculate box size based on density or explicit setting.

        Returns:
            Box side length.
        """
        total_beads = self.n_chains * self.n_beads

        if self._explicit_box_size is not None:
            return self._explicit_box_size
        elif self.density is not None:
            return calculate_box_size(total_beads, self.density)
        else:
            # Default behavior (original logic)
            if self.topology == "ring":
                radius = self.bond_length * self.n_beads / (2 * np.pi)
                spacing = radius * 3
                n_per_dim = int(np.ceil(np.cbrt(self.n_chains)))
                return max(n_per_dim * spacing * 2, 50.0)
            else:
                return max(self.n_beads * self.bond_length * 2, 50.0)

    def _compute_total_energy(self) -> float:
        """
        Compute total system energy (LJ + bonds).

        Returns:
            Total energy.
        """
        if self._positions is None or self._bonds is None:
            raise ValueError("Positions and bonds must be initialized first")

        # Get average LJ parameters from bead types
        avg_sigma = np.mean([bt.sigma for bt in self.bead_types])
        avg_epsilon = np.mean([bt.epsilon for bt in self.bead_types])

        box_size = self._calculate_box_size()

        return compute_total_energy(
            self._positions,
            self._bonds,
            lj_sigma=avg_sigma,
            lj_epsilon=avg_epsilon,
            lj_cutoff=2.5,
            bond_k=self.k_bond,
            bond_r0=self.bond_length,
            box_size=box_size,
        )

    def _generate_initial_positions(self) -> None:
        """Generate initial chain positions and bonds."""
        n_beads = self.n_beads
        box_size = self._calculate_box_size()

        # Generate positions for each chain
        chain_positions = []

        for chain in range(self.n_chains):
            chain_pos = []
            if self.topology == "ring":
                radius = self.bond_length * n_beads / (2 * np.pi)
                for bead in range(n_beads):
                    angle = 2 * np.pi * bead / n_beads
                    pos = np.array([
                        radius * np.cos(angle),
                        radius * np.sin(angle),
                        0.0
                    ])
                    chain_pos.append(pos)
            else:
                # Linear chain
                for bead in range(n_beads):
                    pos = np.array([
                        bead * self.bond_length,
                        0.0,
                        0.0
                    ])
                    chain_pos.append(pos)

            chain_positions.append(chain_pos)

        # Place chains in box
        if self.n_chains > 1 and (self.density is not None or self._explicit_box_size is not None):
            # Use random placement with separation check
            min_sep = max(2.0 * self.bead_types[0].sigma, self.bond_length * 2)
            self._positions, self._chain_indices = place_chains_in_box(
                chain_positions, box_size, min_separation=min_sep
            )
        else:
            # Simple grid/linear placement
            self._positions = []
            self._chain_indices = []

            for chain_idx, chain_pos in enumerate(chain_positions):
                start_idx = len(self._positions)

                if self.topology == "ring":
                    n_per_dim = int(np.ceil(np.cbrt(self.n_chains)))
                    radius = self.bond_length * n_beads / (2 * np.pi)
                    spacing = radius * 3

                    ix = chain_idx % n_per_dim
                    iy = (chain_idx // n_per_dim) % n_per_dim
                    iz = chain_idx // (n_per_dim * n_per_dim)

                    center = np.array([
                        (ix - n_per_dim / 2 + 0.5) * spacing,
                        (iy - n_per_dim / 2 + 0.5) * spacing,
                        (iz - n_per_dim / 2 + 0.5) * spacing
                    ])

                    R = _random_rotation_matrix(np.pi)

                    for pos in chain_pos:
                        rotated = R @ pos
                        self._positions.append(rotated + center)
                else:
                    # Simple offset for linear chains
                    offset = np.array([0.0, chain_idx * self.bond_length * 2, 0.0])
                    for pos in chain_pos:
                        self._positions.append(pos + offset)

                end_idx = len(self._positions)
                self._chain_indices.append((start_idx, end_idx))

        # Generate bonds (0-indexed for internal use)
        self._bonds = []
        for start, end in self._chain_indices:
            chain_len = end - start
            for i in range(chain_len - 1):
                self._bonds.append((start + i, start + i + 1))
            if self.topology == "ring":
                self._bonds.append((end - 1, start))

    def equilibrate(self, mc_config: Optional[MCConfig] = None) -> None:
        """
        Pre-equilibrate the polymer configuration using Monte Carlo moves.

        Updates internal positions in-place.

        Args:
            mc_config: Monte Carlo configuration. Uses instance config or defaults.
        """
        config = mc_config or self._mc_config or MCConfig()

        # Initialize positions if not done
        if self._positions is None:
            self._generate_initial_positions()

        box_size = self._calculate_box_size()

        # Get average LJ parameters
        avg_sigma = np.mean([bt.sigma for bt in self.bead_types])
        avg_epsilon = np.mean([bt.epsilon for bt in self.bead_types])

        logger.info(
            f"Starting MC equilibration: {config.n_steps} steps, "
            f"T={config.temperature}, box_size={box_size:.3f}"
        )

        self._positions, acceptance_stats = mc_equilibrate(
            positions=self._positions,
            bonds=self._bonds,
            chain_indices=self._chain_indices,
            n_steps=config.n_steps,
            temperature=config.temperature,
            move_weights=config.move_weights,
            box_size=box_size,
            lj_sigma=config.lj_sigma if config.lj_sigma else avg_sigma,
            lj_epsilon=config.lj_epsilon if config.lj_epsilon else avg_epsilon,
            lj_cutoff=config.lj_cutoff,
            bond_k=config.bond_k,
            bond_r0=self.bond_length,
            max_displacement=config.max_displacement,
            max_angle=config.max_angle,
            verbose=True,
        )

        logger.info(f"MC equilibration complete. Acceptance rates: {acceptance_stats}")

    def saw_generate(self, saw_config: Optional[SAWConfig] = None) -> bool:
        """
        Generate configuration using Self-Avoiding Random Walk.

        This is a fast alternative to MC equilibration that generates
        overlap-free configurations directly.

        Args:
            saw_config: SAW configuration. Uses instance config or defaults.

        Returns:
            True if successful, False if SAW failed.
        """
        config = saw_config or self._saw_config or SAWConfig()

        # Set collision sigma based on bead type if not specified
        if config.collision_sigma == 1.0 and self.bead_types:
            config.collision_sigma = max(bt.sigma for bt in self.bead_types)

        box_size = self._calculate_box_size()

        logger.info(
            f"Starting SAW generation: {self.n_chains} chains, "
            f"{self.n_beads} beads each, box_size={box_size:.3f}"
        )

        positions, chain_indices, stats = saw_generate_multi_chain(
            n_chains=self.n_chains,
            n_beads_per_chain=self.n_beads,
            bond_length=self.bond_length,
            box_size=box_size,
            config=config,
            topology=self.topology,
        )

        if not stats["success"]:
            logger.warning(
                f"SAW generation failed: {stats['failure_reason']}. "
                f"Chains completed: {stats['chains_completed']}, "
                f"Backtracks: {stats['total_backtracks']}"
            )
            return False

        self._positions = positions
        self._chain_indices = chain_indices

        # Generate bonds (0-indexed)
        self._bonds = []
        for start, end in self._chain_indices:
            chain_len = end - start
            for i in range(chain_len - 1):
                self._bonds.append((start + i, start + i + 1))
            if self.topology == "ring":
                self._bonds.append((end - 1, start))

        logger.info(
            f"SAW generation complete. Backtracks: {stats['total_backtracks']}"
        )
        return True

    def generate_data_file(self) -> None:
        """Generate LAMMPS data file for bead-spring polymer."""
        n_beads = self.n_beads
        n_bonds_per_chain = n_beads - 1 if self.topology == "linear" else n_beads
        n_angles_per_chain = 0
        if self.use_angles:
            n_angles_per_chain = n_beads - 2 if self.topology == "linear" else n_beads

        total_atoms = self.n_chains * n_beads
        total_bonds = self.n_chains * n_bonds_per_chain
        total_angles = self.n_chains * n_angles_per_chain if self.use_angles else 0

        n_atom_types = len(self.bead_types)
        n_angle_types = len(self._angle_type_map) if self.use_angles else 0

        # Calculate box size using the new method
        box_size = self._calculate_box_size()

        # Generate initial positions based on generation method
        if self._positions is None:
            if self._generation_method == "saw":
                success = self.saw_generate(self._saw_config)
                if not success:
                    logger.warning("SAW failed, falling back to geometric placement")
                    self._generate_initial_positions()
            else:
                self._generate_initial_positions()

        # Run MC equilibration if requested (for "mc" method or legacy equilibrate flag)
        if self._generation_method == "mc" or self._equilibrate:
            self.equilibrate(self._mc_config)

        # Use stored positions for writing
        use_stored_positions = self._positions is not None

        with open(f"{self.path}/polymer.data", 'w') as f:
            # Header
            f.write("LAMMPS Bead-Spring Polymer Data File\n\n")
            f.write(f"{total_atoms} atoms\n")
            f.write(f"{total_bonds} bonds\n")
            if self.use_angles:
                f.write(f"{total_angles} angles\n")
            f.write("\n")
            f.write(f"{n_atom_types} atom types\n")
            f.write("1 bond types\n")
            if self.use_angles:
                f.write(f"{n_angle_types} angle types\n")
            f.write("\n")

            # Box dimensions
            f.write(f"{-box_size/2:.1f} {box_size/2:.1f} xlo xhi\n")
            f.write(f"{-box_size/2:.1f} {box_size/2:.1f} ylo yhi\n")
            f.write(f"{-box_size/2:.1f} {box_size/2:.1f} zlo zhi\n\n")

            # Masses
            f.write("Masses\n\n")
            for bt in self.bead_types:
                type_id = self._bead_type_id[bt.name]
                f.write(f"{type_id} {bt.mass:.3f}  # {bt.name}\n")
            f.write("\n")

            # Atoms section: atom-ID molecule-ID atom-type x y z
            f.write("Atoms  # molecular\n\n")
            atom_id = 1

            if use_stored_positions:
                # Use equilibrated/placed positions
                for chain_idx, (start, end) in enumerate(self._chain_indices):
                    for local_bead, global_idx in enumerate(range(start, end)):
                        pos = self._positions[global_idx]
                        bead_name = self._sequence[local_bead]
                        type_id = self._bead_type_id[bead_name]
                        f.write(f"{atom_id} {chain_idx + 1} {type_id} {pos[0]:.3f} {pos[1]:.3f} {pos[2]:.3f}\n")
                        atom_id += 1
            else:
                # Original position generation (fallback for compatibility)
                if self.topology == "ring":
                    radius = self.bond_length * n_beads / (2 * np.pi)
                    n_per_dim = int(np.ceil(np.cbrt(self.n_chains)))
                    spacing = radius * 3

                for chain in range(self.n_chains):
                    if self.topology == "ring":
                        ix = chain % n_per_dim
                        iy = (chain // n_per_dim) % n_per_dim
                        iz = chain // (n_per_dim * n_per_dim)

                        center_x = (ix - n_per_dim / 2 + 0.5) * spacing
                        center_y = (iy - n_per_dim / 2 + 0.5) * spacing
                        center_z = (iz - n_per_dim / 2 + 0.5) * spacing

                        theta = np.random.uniform(0, np.pi)
                        phi = np.random.uniform(0, 2 * np.pi)

                        Rx = np.array([
                            [1, 0, 0],
                            [0, np.cos(theta), -np.sin(theta)],
                            [0, np.sin(theta), np.cos(theta)]
                        ])
                        Rz = np.array([
                            [np.cos(phi), -np.sin(phi), 0],
                            [np.sin(phi), np.cos(phi), 0],
                            [0, 0, 1]
                        ])
                        R = Rz @ Rx

                        for bead in range(n_beads):
                            angle = 2 * np.pi * bead / n_beads
                            pos = np.array([
                                radius * np.cos(angle),
                                radius * np.sin(angle),
                                0.0
                            ])
                            pos = R @ pos + np.array([center_x, center_y, center_z])

                            bead_name = self._sequence[bead]
                            type_id = self._bead_type_id[bead_name]
                            f.write(f"{atom_id} {chain + 1} {type_id} {pos[0]:.3f} {pos[1]:.3f} {pos[2]:.3f}\n")
                            atom_id += 1
                    else:
                        for bead in range(n_beads):
                            x = bead * self.bond_length
                            y = chain * self.bond_length * 2
                            z = 0.0

                            bead_name = self._sequence[bead]
                            type_id = self._bead_type_id[bead_name]
                            f.write(f"{atom_id} {chain + 1} {type_id} {x:.3f} {y:.3f} {z:.3f}\n")
                            atom_id += 1

            # Bonds
            f.write("\nBonds\n\n")
            bond_id = 1
            for chain in range(self.n_chains):
                start_id = chain * n_beads + 1
                for bead in range(n_beads - 1):
                    f.write(f"{bond_id} 1 {start_id + bead} {start_id + bead + 1}\n")
                    bond_id += 1
                if self.topology == "ring":
                    f.write(f"{bond_id} 1 {start_id + n_beads - 1} {start_id}\n")
                    bond_id += 1

            # Angles
            if self.use_angles:
                f.write("\nAngles\n\n")
                angle_id = 1
                for chain in range(self.n_chains):
                    start_id = chain * n_beads + 1

                    # Internal angles
                    for bead in range(n_beads - 2):
                        triplet = self._get_canonical_triplet(
                            self._sequence[bead],
                            self._sequence[bead + 1],
                            self._sequence[bead + 2]
                        )
                        angle_type_id = self._angle_type_map[triplet]
                        f.write(f"{angle_id} {angle_type_id} {start_id + bead} {start_id + bead + 1} {start_id + bead + 2}\n")
                        angle_id += 1

                    # Wrap-around angles for ring
                    if self.topology == "ring" and n_beads >= 3:
                        # n-2, n-1, 0
                        triplet = self._get_canonical_triplet(
                            self._sequence[-2],
                            self._sequence[-1],
                            self._sequence[0]
                        )
                        angle_type_id = self._angle_type_map[triplet]
                        f.write(f"{angle_id} {angle_type_id} {start_id + n_beads - 2} {start_id + n_beads - 1} {start_id}\n")
                        angle_id += 1

                        # n-1, 0, 1
                        triplet = self._get_canonical_triplet(
                            self._sequence[-1],
                            self._sequence[0],
                            self._sequence[1]
                        )
                        angle_type_id = self._angle_type_map[triplet]
                        f.write(f"{angle_id} {angle_type_id} {start_id + n_beads - 1} {start_id} {start_id + 1}\n")
                        angle_id += 1

        # Generate LAMMPS input script
        self._generate_input_script()
        logger.info(f"Generated bead-spring polymer files in {self.path}")

    def _generate_input_script(self) -> None:
        """Generate LAMMPS input script for the bead-spring polymer.

        Note: For Kremer-Grest polymers using FENE bonds, the WCA cutoff
        (2^(1/6) * sigma ≈ 1.12246) should be used instead of full LJ cutoff.
        This is controlled by setting pair_style='wca' in the class constructor.
        The equilibrium bond length of ~0.97 emerges from the balance of
        FENE + WCA potentials - there is no explicit r0 parameter in FENE.
        """
        # Determine cutoff from max sigma
        max_sigma = max(bt.sigma for bt in self.bead_types)

        # Use WCA cutoff if pair_style is 'wca', otherwise use full LJ
        # For Kremer-Grest melts at P=0, WCA (purely repulsive) is correct
        if hasattr(self, '_pair_style') and self._pair_style == 'wca':
            # WCA cutoff: 2^(1/6) * sigma
            cutoff = (2 ** (1/6)) * max_sigma
            pair_modify = "pair_modify     shift yes\n"
        else:
            cutoff = 2.5 * max_sigma
            pair_modify = ""

        with open(f"{self.path}/in.polymer", 'w') as f:
            f.write("# LAMMPS input script for bead-spring polymer\n")
            f.write("# Auto-generated by AutoPoly BeadSpringPolymer\n\n")

            # Simulation settings
            f.write("units           lj\n")
            f.write("atom_style      molecular\n")
            f.write("boundary        p p p\n\n")

            f.write("read_data       polymer.data\n\n")

            # Pair style and coefficients
            f.write(f"pair_style      lj/cut {cutoff:.5f}\n")
            for (i, j), (eps, sig) in sorted(self._pair_coeffs.items()):
                f.write(f"pair_coeff      {i} {j} {eps:.4f} {sig:.4f}\n")
            if pair_modify:
                f.write(pair_modify)
            f.write("\n")

            # Bond style and coefficients
            if self.bond_style == "fene":
                # For FENE, use the first bead type's parameters for the attractive LJ part
                eps = self.bead_types[0].epsilon
                sig = self.bead_types[0].sigma
                f.write("bond_style      fene\n")
                # FENE bond: K, R0, epsilon, sigma
                # Note: No r0 parameter - equilibrium distance emerges from FENE+WCA balance
                f.write(f"bond_coeff      1 {self.k_bond:.1f} {self.fene_r0:.4f} {eps:.4f} {sig:.4f}\n")
                f.write("special_bonds   fene\n\n")
            else:
                f.write("bond_style      harmonic\n")
                f.write(f"bond_coeff      1 {self.k_bond:.1f} {self.bond_length:.4f}\n\n")

            # Angle style and coefficients
            if self.use_angles:
                f.write("angle_style     harmonic\n")
                for triplet, type_id in sorted(self._angle_type_map.items(), key=lambda x: x[1]):
                    k, theta0 = self._get_angle_params(triplet)
                    triplet_str = "-".join(triplet)
                    f.write(f"angle_coeff     {type_id} {k:.4f} {theta0:.1f}  # {triplet_str}\n")
                f.write("\n")

            # Neighbor settings - CRITICAL for FENE bonds
            # Use larger skin (2.0) for FENE to prevent lost atoms
            f.write("neighbor        2.0 bin\n")
            f.write("neigh_modify    every 2 delay 4 check yes\n\n")

            # Output settings
            f.write("thermo_style    custom step temp pe ke etotal press vol density\n")
            f.write("thermo          1000\n\n")

            # Trajectory output
            f.write("dump            1 all custom 1000 dump.lammpstrj id type mol x y z\n")
            f.write("dump_modify     1 sort id\n\n")

            # Minimization
            f.write("# Energy minimization\n")
            f.write("minimize        1.0e-4 1.0e-6 1000 10000\n")
            f.write("write_restart   min.restart\n")
            f.write("write_data      min.data\n")
            f.write("reset_timestep  0\n\n")

            # MD settings
            # Timestep: 0.001 is standard for Kremer-Grest with FENE bonds
            f.write("# Production MD\n")
            f.write("timestep        0.001\n")
            # Thermostat: Tdamp = 0.1 (100x timestep) for proper control
            f.write("fix             1 all nvt temp 1.0 1.0 0.1 tchain 3\n")
            f.write("run             100000\n")
            f.write("write_restart   prod.restart\n")
            f.write("write_data      prod.data\n")

    def get_system_info(self) -> dict:
        """
        Get comprehensive information about the bead-spring polymer system.

        Returns:
            Dictionary containing system properties.
        """
        n_beads = self.n_beads
        n_bonds_per_chain = n_beads - 1 if self.topology == "linear" else n_beads
        n_angles_per_chain = 0
        if self.use_angles:
            n_angles_per_chain = n_beads - 2 if self.topology == "linear" else n_beads

        total_atoms = self.n_chains * n_beads
        total_bonds = self.n_chains * n_bonds_per_chain
        total_angles = self.n_chains * n_angles_per_chain if self.use_angles else 0

        return {
            'name': self.name,
            'n_chains': self.n_chains,
            'n_beads_per_chain': n_beads,
            'topology': self.topology,
            'total_atoms': total_atoms,
            'total_bonds': total_bonds,
            'total_angles': total_angles,
            'bond_style': self.bond_style,
            'bond_length': self.bond_length,
            'k_bond': self.k_bond,
            'use_angles': self.use_angles,
            'bead_types': [bt.name for bt in self.bead_types],
            'sequence': self._sequence,
            'output_path': self.path,
        }

    @classmethod
    def kremer_grest(
        cls,
        name: str,
        system: object,
        n_chains: int,
        n_beads: int,
        topology: str = "linear",
        density: float = 0.74,
        generation_method: str = "saw",
    ) -> "BeadSpringPolymer":
        """
        Create a Kremer-Grest bead-spring polymer with standard parameters.

        The Kremer-Grest model is a standard coarse-grained polymer model used
        for universal polymer behavior studies. It uses:
        - FENE bonds (K=30, R0=1.5)
        - WCA purely repulsive interactions (cutoff = 2^(1/6)*sigma ≈ 1.12246)
        - Standard LJ parameters (epsilon=1.0, sigma=1.0, mass=1.0)

        The equilibrium bond length of ~0.97 EMERGES from the balance of
        FENE (attractive) + WCA (repulsive) - there is no explicit r0 parameter.

        Args:
            name: Name for output files.
            system: System object containing path information.
            n_chains: Number of polymer chains.
            n_beads: Number of beads per chain.
            topology: "linear" or "ring" topology.
            density: Target bead density (beads/sigma^3). Default 0.74 for
                initial placement (will compress to ~0.85-0.90 during NPT).
            generation_method: Method for generating initial configurations
                ("geometric", "saw", or "mc").

        Returns:
            BeadSpringPolymer instance configured for Kremer-Grest model.

        Example:
            >>> from AutoPoly import System
            >>> from AutoPoly import BeadSpringPolymer
            >>> system = System(out="kg_simulation")
            >>> kg_polymer = BeadSpringPolymer.kremer_grest(
            ...     name="kg_melt",
            ...     system=system,
            ...     n_chains=100,
            ...     n_beads=10,
            ...     topology="linear"
            ... )
            >>> kg_polymer.generate_data_file()

        References:
            Kremer, K., & Grest, G. S. (1990). Dynamics of entangled linear
            polymer melts: A molecular-dynamics simulation. J. Chem. Phys.
            92(8), 5057-5086.
        """
        # Standard Kremer-Grest bead type
        kg_bead_type = BeadType(name="KG", mass=1.0, epsilon=1.0, sigma=1.0)

        # Uniform sequence of KG beads
        sequence = [("KG", n_beads)]

        return cls(
            name=name,
            system=system,
            n_chains=n_chains,
            bead_types=[kg_bead_type],
            sequence=sequence,
            topology=topology,
            bond_length=0.97,  # Starting guess, equilibrium emerges from FENE+WCA
            bond_style="fene",
            k_bond=30.0,
            fene_r0=1.5,
            pair_style="wca",  # WCA purely repulsive for KG melts
            density=density,
            generation_method=generation_method,
        )

n_beads: int property

Number of beads per chain (derived from sequence).

equilibrate(mc_config: Optional[MCConfig] = None) -> None

Pre-equilibrate the polymer configuration using Monte Carlo moves.

Updates internal positions in-place.

Parameters:

Name	Type	Description	Default
mc_config	Optional[MCConfig]	Monte Carlo configuration. Uses instance config or defaults.	None

Source code in AutoPoly/models/bead_spring.py

def equilibrate(self, mc_config: Optional[MCConfig] = None) -> None:
    """
    Pre-equilibrate the polymer configuration using Monte Carlo moves.

    Updates internal positions in-place.

    Args:
        mc_config: Monte Carlo configuration. Uses instance config or defaults.
    """
    config = mc_config or self._mc_config or MCConfig()

    # Initialize positions if not done
    if self._positions is None:
        self._generate_initial_positions()

    box_size = self._calculate_box_size()

    # Get average LJ parameters
    avg_sigma = np.mean([bt.sigma for bt in self.bead_types])
    avg_epsilon = np.mean([bt.epsilon for bt in self.bead_types])

    logger.info(
        f"Starting MC equilibration: {config.n_steps} steps, "
        f"T={config.temperature}, box_size={box_size:.3f}"
    )

    self._positions, acceptance_stats = mc_equilibrate(
        positions=self._positions,
        bonds=self._bonds,
        chain_indices=self._chain_indices,
        n_steps=config.n_steps,
        temperature=config.temperature,
        move_weights=config.move_weights,
        box_size=box_size,
        lj_sigma=config.lj_sigma if config.lj_sigma else avg_sigma,
        lj_epsilon=config.lj_epsilon if config.lj_epsilon else avg_epsilon,
        lj_cutoff=config.lj_cutoff,
        bond_k=config.bond_k,
        bond_r0=self.bond_length,
        max_displacement=config.max_displacement,
        max_angle=config.max_angle,
        verbose=True,
    )

    logger.info(f"MC equilibration complete. Acceptance rates: {acceptance_stats}")

saw_generate(saw_config: Optional[SAWConfig] = None) -> bool

Generate configuration using Self-Avoiding Random Walk.

This is a fast alternative to MC equilibration that generates overlap-free configurations directly.

Parameters:

Name	Type	Description	Default
saw_config	Optional[SAWConfig]	SAW configuration. Uses instance config or defaults.	None

Returns:

Type	Description
bool	True if successful, False if SAW failed.

Source code in AutoPoly/models/bead_spring.py

def saw_generate(self, saw_config: Optional[SAWConfig] = None) -> bool:
    """
    Generate configuration using Self-Avoiding Random Walk.

    This is a fast alternative to MC equilibration that generates
    overlap-free configurations directly.

    Args:
        saw_config: SAW configuration. Uses instance config or defaults.

    Returns:
        True if successful, False if SAW failed.
    """
    config = saw_config or self._saw_config or SAWConfig()

    # Set collision sigma based on bead type if not specified
    if config.collision_sigma == 1.0 and self.bead_types:
        config.collision_sigma = max(bt.sigma for bt in self.bead_types)

    box_size = self._calculate_box_size()

    logger.info(
        f"Starting SAW generation: {self.n_chains} chains, "
        f"{self.n_beads} beads each, box_size={box_size:.3f}"
    )

    positions, chain_indices, stats = saw_generate_multi_chain(
        n_chains=self.n_chains,
        n_beads_per_chain=self.n_beads,
        bond_length=self.bond_length,
        box_size=box_size,
        config=config,
        topology=self.topology,
    )

    if not stats["success"]:
        logger.warning(
            f"SAW generation failed: {stats['failure_reason']}. "
            f"Chains completed: {stats['chains_completed']}, "
            f"Backtracks: {stats['total_backtracks']}"
        )
        return False

    self._positions = positions
    self._chain_indices = chain_indices

    # Generate bonds (0-indexed)
    self._bonds = []
    for start, end in self._chain_indices:
        chain_len = end - start
        for i in range(chain_len - 1):
            self._bonds.append((start + i, start + i + 1))
        if self.topology == "ring":
            self._bonds.append((end - 1, start))

    logger.info(
        f"SAW generation complete. Backtracks: {stats['total_backtracks']}"
    )
    return True

generate_data_file() -> None

Generate LAMMPS data file for bead-spring polymer.

Source code in AutoPoly/models/bead_spring.py

def generate_data_file(self) -> None:
    """Generate LAMMPS data file for bead-spring polymer."""
    n_beads = self.n_beads
    n_bonds_per_chain = n_beads - 1 if self.topology == "linear" else n_beads
    n_angles_per_chain = 0
    if self.use_angles:
        n_angles_per_chain = n_beads - 2 if self.topology == "linear" else n_beads

    total_atoms = self.n_chains * n_beads
    total_bonds = self.n_chains * n_bonds_per_chain
    total_angles = self.n_chains * n_angles_per_chain if self.use_angles else 0

    n_atom_types = len(self.bead_types)
    n_angle_types = len(self._angle_type_map) if self.use_angles else 0

    # Calculate box size using the new method
    box_size = self._calculate_box_size()

    # Generate initial positions based on generation method
    if self._positions is None:
        if self._generation_method == "saw":
            success = self.saw_generate(self._saw_config)
            if not success:
                logger.warning("SAW failed, falling back to geometric placement")
                self._generate_initial_positions()
        else:
            self._generate_initial_positions()

    # Run MC equilibration if requested (for "mc" method or legacy equilibrate flag)
    if self._generation_method == "mc" or self._equilibrate:
        self.equilibrate(self._mc_config)

    # Use stored positions for writing
    use_stored_positions = self._positions is not None

    with open(f"{self.path}/polymer.data", 'w') as f:
        # Header
        f.write("LAMMPS Bead-Spring Polymer Data File\n\n")
        f.write(f"{total_atoms} atoms\n")
        f.write(f"{total_bonds} bonds\n")
        if self.use_angles:
            f.write(f"{total_angles} angles\n")
        f.write("\n")
        f.write(f"{n_atom_types} atom types\n")
        f.write("1 bond types\n")
        if self.use_angles:
            f.write(f"{n_angle_types} angle types\n")
        f.write("\n")

        # Box dimensions
        f.write(f"{-box_size/2:.1f} {box_size/2:.1f} xlo xhi\n")
        f.write(f"{-box_size/2:.1f} {box_size/2:.1f} ylo yhi\n")
        f.write(f"{-box_size/2:.1f} {box_size/2:.1f} zlo zhi\n\n")

        # Masses
        f.write("Masses\n\n")
        for bt in self.bead_types:
            type_id = self._bead_type_id[bt.name]
            f.write(f"{type_id} {bt.mass:.3f}  # {bt.name}\n")
        f.write("\n")

        # Atoms section: atom-ID molecule-ID atom-type x y z
        f.write("Atoms  # molecular\n\n")
        atom_id = 1

        if use_stored_positions:
            # Use equilibrated/placed positions
            for chain_idx, (start, end) in enumerate(self._chain_indices):
                for local_bead, global_idx in enumerate(range(start, end)):
                    pos = self._positions[global_idx]
                    bead_name = self._sequence[local_bead]
                    type_id = self._bead_type_id[bead_name]
                    f.write(f"{atom_id} {chain_idx + 1} {type_id} {pos[0]:.3f} {pos[1]:.3f} {pos[2]:.3f}\n")
                    atom_id += 1
        else:
            # Original position generation (fallback for compatibility)
            if self.topology == "ring":
                radius = self.bond_length * n_beads / (2 * np.pi)
                n_per_dim = int(np.ceil(np.cbrt(self.n_chains)))
                spacing = radius * 3

            for chain in range(self.n_chains):
                if self.topology == "ring":
                    ix = chain % n_per_dim
                    iy = (chain // n_per_dim) % n_per_dim
                    iz = chain // (n_per_dim * n_per_dim)

                    center_x = (ix - n_per_dim / 2 + 0.5) * spacing
                    center_y = (iy - n_per_dim / 2 + 0.5) * spacing
                    center_z = (iz - n_per_dim / 2 + 0.5) * spacing

                    theta = np.random.uniform(0, np.pi)
                    phi = np.random.uniform(0, 2 * np.pi)

                    Rx = np.array([
                        [1, 0, 0],
                        [0, np.cos(theta), -np.sin(theta)],
                        [0, np.sin(theta), np.cos(theta)]
                    ])
                    Rz = np.array([
                        [np.cos(phi), -np.sin(phi), 0],
                        [np.sin(phi), np.cos(phi), 0],
                        [0, 0, 1]
                    ])
                    R = Rz @ Rx

                    for bead in range(n_beads):
                        angle = 2 * np.pi * bead / n_beads
                        pos = np.array([
                            radius * np.cos(angle),
                            radius * np.sin(angle),
                            0.0
                        ])
                        pos = R @ pos + np.array([center_x, center_y, center_z])

                        bead_name = self._sequence[bead]
                        type_id = self._bead_type_id[bead_name]
                        f.write(f"{atom_id} {chain + 1} {type_id} {pos[0]:.3f} {pos[1]:.3f} {pos[2]:.3f}\n")
                        atom_id += 1
                else:
                    for bead in range(n_beads):
                        x = bead * self.bond_length
                        y = chain * self.bond_length * 2
                        z = 0.0

                        bead_name = self._sequence[bead]
                        type_id = self._bead_type_id[bead_name]
                        f.write(f"{atom_id} {chain + 1} {type_id} {x:.3f} {y:.3f} {z:.3f}\n")
                        atom_id += 1

        # Bonds
        f.write("\nBonds\n\n")
        bond_id = 1
        for chain in range(self.n_chains):
            start_id = chain * n_beads + 1
            for bead in range(n_beads - 1):
                f.write(f"{bond_id} 1 {start_id + bead} {start_id + bead + 1}\n")
                bond_id += 1
            if self.topology == "ring":
                f.write(f"{bond_id} 1 {start_id + n_beads - 1} {start_id}\n")
                bond_id += 1

        # Angles
        if self.use_angles:
            f.write("\nAngles\n\n")
            angle_id = 1
            for chain in range(self.n_chains):
                start_id = chain * n_beads + 1

                # Internal angles
                for bead in range(n_beads - 2):
                    triplet = self._get_canonical_triplet(
                        self._sequence[bead],
                        self._sequence[bead + 1],
                        self._sequence[bead + 2]
                    )
                    angle_type_id = self._angle_type_map[triplet]
                    f.write(f"{angle_id} {angle_type_id} {start_id + bead} {start_id + bead + 1} {start_id + bead + 2}\n")
                    angle_id += 1

                # Wrap-around angles for ring
                if self.topology == "ring" and n_beads >= 3:
                    # n-2, n-1, 0
                    triplet = self._get_canonical_triplet(
                        self._sequence[-2],
                        self._sequence[-1],
                        self._sequence[0]
                    )
                    angle_type_id = self._angle_type_map[triplet]
                    f.write(f"{angle_id} {angle_type_id} {start_id + n_beads - 2} {start_id + n_beads - 1} {start_id}\n")
                    angle_id += 1

                    # n-1, 0, 1
                    triplet = self._get_canonical_triplet(
                        self._sequence[-1],
                        self._sequence[0],
                        self._sequence[1]
                    )
                    angle_type_id = self._angle_type_map[triplet]
                    f.write(f"{angle_id} {angle_type_id} {start_id + n_beads - 1} {start_id} {start_id + 1}\n")
                    angle_id += 1

    # Generate LAMMPS input script
    self._generate_input_script()
    logger.info(f"Generated bead-spring polymer files in {self.path}")

get_system_info() -> dict

Get comprehensive information about the bead-spring polymer system.

Returns:

Type	Description
dict	Dictionary containing system properties.

Source code in AutoPoly/models/bead_spring.py

def get_system_info(self) -> dict:
    """
    Get comprehensive information about the bead-spring polymer system.

    Returns:
        Dictionary containing system properties.
    """
    n_beads = self.n_beads
    n_bonds_per_chain = n_beads - 1 if self.topology == "linear" else n_beads
    n_angles_per_chain = 0
    if self.use_angles:
        n_angles_per_chain = n_beads - 2 if self.topology == "linear" else n_beads

    total_atoms = self.n_chains * n_beads
    total_bonds = self.n_chains * n_bonds_per_chain
    total_angles = self.n_chains * n_angles_per_chain if self.use_angles else 0

    return {
        'name': self.name,
        'n_chains': self.n_chains,
        'n_beads_per_chain': n_beads,
        'topology': self.topology,
        'total_atoms': total_atoms,
        'total_bonds': total_bonds,
        'total_angles': total_angles,
        'bond_style': self.bond_style,
        'bond_length': self.bond_length,
        'k_bond': self.k_bond,
        'use_angles': self.use_angles,
        'bead_types': [bt.name for bt in self.bead_types],
        'sequence': self._sequence,
        'output_path': self.path,
    }

kremer_grest(name: str, system: object, n_chains: int, n_beads: int, topology: str = 'linear', density: float = 0.74, generation_method: str = 'saw') -> BeadSpringPolymer classmethod

Create a Kremer-Grest bead-spring polymer with standard parameters.

The Kremer-Grest model is a standard coarse-grained polymer model used for universal polymer behavior studies. It uses: - FENE bonds (K=30, R0=1.5) - WCA purely repulsive interactions (cutoff = 2^(1/6)*sigma ≈ 1.12246) - Standard LJ parameters (epsilon=1.0, sigma=1.0, mass=1.0)

The equilibrium bond length of ~0.97 EMERGES from the balance of FENE (attractive) + WCA (repulsive) - there is no explicit r0 parameter.

Parameters:

Name	Type	Description	Default
name	str	Name for output files.	required
system	object	System object containing path information.	required
n_chains	int	Number of polymer chains.	required
n_beads	int	Number of beads per chain.	required
topology	str	"linear" or "ring" topology.	'linear'
density	float	Target bead density (beads/sigma^3). Default 0.74 for initial placement (will compress to ~0.85-0.90 during NPT).	0.74
generation_method	str	Method for generating initial configurations ("geometric", "saw", or "mc").	'saw'

Returns:

Type	Description
BeadSpringPolymer	BeadSpringPolymer instance configured for Kremer-Grest model.

Example

from AutoPoly import System from AutoPoly import BeadSpringPolymer system = System(out="kg_simulation") kg_polymer = BeadSpringPolymer.kremer_grest( ... name="kg_melt", ... system=system, ... n_chains=100, ... n_beads=10, ... topology="linear" ... ) kg_polymer.generate_data_file()

References

Kremer, K., & Grest, G. S. (1990). Dynamics of entangled linear polymer melts: A molecular-dynamics simulation. J. Chem. Phys. 92(8), 5057-5086.

Source code in AutoPoly/models/bead_spring.py

@classmethod
def kremer_grest(
    cls,
    name: str,
    system: object,
    n_chains: int,
    n_beads: int,
    topology: str = "linear",
    density: float = 0.74,
    generation_method: str = "saw",
) -> "BeadSpringPolymer":
    """
    Create a Kremer-Grest bead-spring polymer with standard parameters.

    The Kremer-Grest model is a standard coarse-grained polymer model used
    for universal polymer behavior studies. It uses:
    - FENE bonds (K=30, R0=1.5)
    - WCA purely repulsive interactions (cutoff = 2^(1/6)*sigma ≈ 1.12246)
    - Standard LJ parameters (epsilon=1.0, sigma=1.0, mass=1.0)

    The equilibrium bond length of ~0.97 EMERGES from the balance of
    FENE (attractive) + WCA (repulsive) - there is no explicit r0 parameter.

    Args:
        name: Name for output files.
        system: System object containing path information.
        n_chains: Number of polymer chains.
        n_beads: Number of beads per chain.
        topology: "linear" or "ring" topology.
        density: Target bead density (beads/sigma^3). Default 0.74 for
            initial placement (will compress to ~0.85-0.90 during NPT).
        generation_method: Method for generating initial configurations
            ("geometric", "saw", or "mc").

    Returns:
        BeadSpringPolymer instance configured for Kremer-Grest model.

    Example:
        >>> from AutoPoly import System
        >>> from AutoPoly import BeadSpringPolymer
        >>> system = System(out="kg_simulation")
        >>> kg_polymer = BeadSpringPolymer.kremer_grest(
        ...     name="kg_melt",
        ...     system=system,
        ...     n_chains=100,
        ...     n_beads=10,
        ...     topology="linear"
        ... )
        >>> kg_polymer.generate_data_file()

    References:
        Kremer, K., & Grest, G. S. (1990). Dynamics of entangled linear
        polymer melts: A molecular-dynamics simulation. J. Chem. Phys.
        92(8), 5057-5086.
    """
    # Standard Kremer-Grest bead type
    kg_bead_type = BeadType(name="KG", mass=1.0, epsilon=1.0, sigma=1.0)

    # Uniform sequence of KG beads
    sequence = [("KG", n_beads)]

    return cls(
        name=name,
        system=system,
        n_chains=n_chains,
        bead_types=[kg_bead_type],
        sequence=sequence,
        topology=topology,
        bond_length=0.97,  # Starting guess, equilibrium emerges from FENE+WCA
        bond_style="fene",
        k_bond=30.0,
        fene_r0=1.5,
        pair_style="wca",  # WCA purely repulsive for KG melts
        density=density,
        generation_method=generation_method,
    )

calculate_box_size(n_beads: int, density: float = DEFAULT_BEAD_DENSITY) -> float

Calculate cubic box size from number of beads and target density.

Parameters:

Name	Type	Description	Default
n_beads	int	Total number of beads in system.	required
density	float	Target bead density (beads per sigma^3).	DEFAULT_BEAD_DENSITY

Returns:

Type	Description
float	Side length of cubic box.

Formula: V = n_beads / density, box_size = V^(1/3)

Source code in AutoPoly/models/bead_spring.py

def calculate_box_size(n_beads: int, density: float = DEFAULT_BEAD_DENSITY) -> float:
    """
    Calculate cubic box size from number of beads and target density.

    Parameters:
        n_beads: Total number of beads in system.
        density: Target bead density (beads per sigma^3).

    Returns:
        Side length of cubic box.

    Formula: V = n_beads / density, box_size = V^(1/3)
    """
    volume = n_beads / density
    box_size = volume ** (1/3)
    return box_size

compute_lj_energy(positions: Union[List[np.ndarray], np.ndarray], sigma: float = 1.0, epsilon: float = 1.0, cutoff: float = 2.5, exclude_bonded: Optional[List[Tuple[int, int]]] = None, box_size: Optional[float] = None) -> float

Compute total LJ energy for non-bonded pairs.

Parameters:

Name	Type	Description	Default
positions	Union[List[ndarray], ndarray]	List of position arrays or (N, 3) array for each bead.	required
sigma	float	LJ sigma parameter.	1.0
epsilon	float	LJ epsilon parameter.	1.0
cutoff	float	Cutoff distance in units of sigma.	2.5
exclude_bonded	Optional[List[Tuple[int, int]]]	List of bonded pairs (i, j) to exclude from LJ calculation.	None
box_size	Optional[float]	Box size for periodic boundary conditions (None = no PBC).	None

Returns:

Type	Description
float	Total LJ energy.

Source code in AutoPoly/models/bead_spring.py

def compute_lj_energy(
    positions: Union[List[np.ndarray], np.ndarray],
    sigma: float = 1.0,
    epsilon: float = 1.0,
    cutoff: float = 2.5,
    exclude_bonded: Optional[List[Tuple[int, int]]] = None,
    box_size: Optional[float] = None,
) -> float:
    """
    Compute total LJ energy for non-bonded pairs.

    Parameters:
        positions: List of position arrays or (N, 3) array for each bead.
        sigma: LJ sigma parameter.
        epsilon: LJ epsilon parameter.
        cutoff: Cutoff distance in units of sigma.
        exclude_bonded: List of bonded pairs (i, j) to exclude from LJ calculation.
        box_size: Box size for periodic boundary conditions (None = no PBC).

    Returns:
        Total LJ energy.
    """
    n_beads = len(positions)
    cutoff_dist = cutoff * sigma
    cutoff_dist_sq = cutoff_dist ** 2
    energy = 0.0

    # Build set of excluded pairs for O(1) lookup
    excluded = set()
    if exclude_bonded:
        for i, j in exclude_bonded:
            excluded.add((min(i, j), max(i, j)))

    for i in range(n_beads):
        for j in range(i + 1, n_beads):
            if (i, j) in excluded:
                continue

            r_vec = positions[j] - positions[i]

            # Apply minimum image convention for PBC
            if box_size is not None:
                r_vec = r_vec - box_size * np.round(r_vec / box_size)

            r_sq = np.dot(r_vec, r_vec)

            if r_sq < cutoff_dist_sq and r_sq > 1e-10:
                r2_inv = (sigma * sigma) / r_sq
                r6_inv = r2_inv ** 3
                r12_inv = r6_inv ** 2
                energy += 4.0 * epsilon * (r12_inv - r6_inv)

    return energy

compute_lj_energy_vectorized(positions: np.ndarray, excluded_mask: np.ndarray, sigma: float = 1.0, epsilon: float = 1.0, cutoff: float = 2.5, box_size: Optional[float] = None) -> float

Vectorized LJ energy using numpy broadcasting.

Parameters:

Name	Type	Description	Default
positions	ndarray	(N, 3) array of bead positions.	required
excluded_mask	ndarray	(N, N) boolean array where True means pair is excluded.	required
sigma	float	LJ sigma parameter.	1.0
epsilon	float	LJ epsilon parameter.	1.0
cutoff	float	Cutoff distance in units of sigma.	2.5
box_size	Optional[float]	Box size for periodic boundary conditions.	None

Returns:

Type	Description
float	Total LJ energy.

Source code in AutoPoly/models/bead_spring.py

def compute_lj_energy_vectorized(
    positions: np.ndarray,
    excluded_mask: np.ndarray,
    sigma: float = 1.0,
    epsilon: float = 1.0,
    cutoff: float = 2.5,
    box_size: Optional[float] = None,
) -> float:
    """
    Vectorized LJ energy using numpy broadcasting.

    Parameters:
        positions: (N, 3) array of bead positions.
        excluded_mask: (N, N) boolean array where True means pair is excluded.
        sigma: LJ sigma parameter.
        epsilon: LJ epsilon parameter.
        cutoff: Cutoff distance in units of sigma.
        box_size: Box size for periodic boundary conditions.

    Returns:
        Total LJ energy.
    """
    cutoff_dist_sq = (cutoff * sigma) ** 2

    # Pairwise differences: (N, N, 3)
    diff = positions[:, None, :] - positions[None, :, :]

    # Apply PBC
    if box_size is not None:
        diff = diff - box_size * np.round(diff / box_size)

    # Squared distances: (N, N)
    r_sq = np.sum(diff ** 2, axis=2)

    # Mask for valid pairs (not excluded, within cutoff, not self)
    # Only consider upper triangle to avoid double counting
    upper_tri = np.triu(np.ones_like(r_sq, dtype=bool), k=1)
    mask = upper_tri & (r_sq < cutoff_dist_sq) & (r_sq > 1e-10) & ~excluded_mask

    # LJ calculation only for valid pairs
    r_sq_valid = np.where(mask, r_sq, 1.0)  # Avoid division by zero
    r2_inv = (sigma ** 2) / r_sq_valid
    r6_inv = r2_inv ** 3
    r12_inv = r6_inv ** 2

    energy = 4.0 * epsilon * np.sum(np.where(mask, r12_inv - r6_inv, 0.0))
    return energy

compute_local_lj_energy(positions: np.ndarray, bead_idx: int, excluded_set: set, sigma: float = 1.0, epsilon: float = 1.0, cutoff: float = 2.5, box_size: Optional[float] = None) -> float

Compute LJ energy contribution from one bead to all others - O(N).

This is the key optimization: instead of recomputing O(N²) energy, we only compute the O(N) interactions involving the moved bead.

Parameters:

Name	Type	Description	Default
positions	ndarray	(N, 3) array of bead positions.	required
bead_idx	int	Index of the bead to compute interactions for.	required
excluded_set	set	Set of bead indices that are bonded to bead_idx.	required
sigma	float	LJ sigma parameter.	1.0
epsilon	float	LJ epsilon parameter.	1.0
cutoff	float	Cutoff distance in units of sigma.	2.5
box_size	Optional[float]	Box size for periodic boundary conditions.	None

Returns:

Type	Description
float	LJ energy contribution from bead_idx to all other beads.

Source code in AutoPoly/models/bead_spring.py

def compute_local_lj_energy(
    positions: np.ndarray,
    bead_idx: int,
    excluded_set: set,
    sigma: float = 1.0,
    epsilon: float = 1.0,
    cutoff: float = 2.5,
    box_size: Optional[float] = None,
) -> float:
    """
    Compute LJ energy contribution from one bead to all others - O(N).

    This is the key optimization: instead of recomputing O(N²) energy,
    we only compute the O(N) interactions involving the moved bead.

    Parameters:
        positions: (N, 3) array of bead positions.
        bead_idx: Index of the bead to compute interactions for.
        excluded_set: Set of bead indices that are bonded to bead_idx.
        sigma: LJ sigma parameter.
        epsilon: LJ epsilon parameter.
        cutoff: Cutoff distance in units of sigma.
        box_size: Box size for periodic boundary conditions.

    Returns:
        LJ energy contribution from bead_idx to all other beads.
    """
    cutoff_dist_sq = (cutoff * sigma) ** 2
    pos_i = positions[bead_idx]
    n_beads = len(positions)

    # Vectorized distance to all other beads
    diff = positions - pos_i  # (N, 3)

    if box_size is not None:
        diff = diff - box_size * np.round(diff / box_size)

    r_sq = np.sum(diff ** 2, axis=1)  # (N,)

    # Build mask: valid pairs (not excluded, within cutoff, not self)
    mask = (r_sq < cutoff_dist_sq) & (r_sq > 1e-10)
    mask[bead_idx] = False  # Exclude self

    # Exclude bonded neighbors
    for j in excluded_set:
        mask[j] = False

    # LJ only for valid pairs
    r_sq_valid = np.where(mask, r_sq, 1.0)  # Avoid division by zero
    r2_inv = (sigma ** 2) / r_sq_valid
    r6_inv = r2_inv ** 3
    r12_inv = r6_inv ** 2

    energy = 4.0 * epsilon * np.sum(np.where(mask, r12_inv - r6_inv, 0.0))
    return energy

compute_local_bond_energy(positions: np.ndarray, bead_idx: int, bead_bonds: List[int], k_bond: float = 100.0, r0: float = 1.0, box_size: Optional[float] = None) -> float

Compute bond energy for bonds involving a specific bead - O(degree).

Parameters:

Name	Type	Description	Default
positions	ndarray	(N, 3) array of bead positions.	required
bead_idx	int	Index of the bead.	required
bead_bonds	List[int]	List of bead indices that are bonded to bead_idx.	required
k_bond	float	Bond spring constant.	100.0
r0	float	Equilibrium bond length.	1.0
box_size	Optional[float]	Box size for periodic boundary conditions.	None

Returns:

Type	Description
float	Bond energy contribution from bead_idx.

Source code in AutoPoly/models/bead_spring.py

def compute_local_bond_energy(
    positions: np.ndarray,
    bead_idx: int,
    bead_bonds: List[int],
    k_bond: float = 100.0,
    r0: float = 1.0,
    box_size: Optional[float] = None,
) -> float:
    """
    Compute bond energy for bonds involving a specific bead - O(degree).

    Parameters:
        positions: (N, 3) array of bead positions.
        bead_idx: Index of the bead.
        bead_bonds: List of bead indices that are bonded to bead_idx.
        k_bond: Bond spring constant.
        r0: Equilibrium bond length.
        box_size: Box size for periodic boundary conditions.

    Returns:
        Bond energy contribution from bead_idx.
    """
    if not bead_bonds:
        return 0.0

    pos_i = positions[bead_idx]
    energy = 0.0

    for j in bead_bonds:
        r_vec = positions[j] - pos_i

        if box_size is not None:
            r_vec = r_vec - box_size * np.round(r_vec / box_size)

        r = np.linalg.norm(r_vec)
        dr = r - r0
        energy += 0.5 * k_bond * dr * dr

    return energy

compute_bond_energy(positions: Union[List[np.ndarray], np.ndarray], bonds: List[Tuple[int, int]], k_bond: float = 100.0, r0: float = 1.0, box_size: Optional[float] = None) -> float

Compute harmonic bond energy.

Parameters:

Name	Type	Description	Default
positions	Union[List[ndarray], ndarray]	List of position arrays or (N, 3) array for each bead.	required
bonds	List[Tuple[int, int]]	List of (atom1_idx, atom2_idx) tuples (0-indexed).	required
k_bond	float	Bond spring constant.	100.0
r0	float	Equilibrium bond length.	1.0
box_size	Optional[float]	Box size for periodic boundary conditions.	None

Returns:

Type	Description
float	Total bond energy.

Source code in AutoPoly/models/bead_spring.py

def compute_bond_energy(
    positions: Union[List[np.ndarray], np.ndarray],
    bonds: List[Tuple[int, int]],
    k_bond: float = 100.0,
    r0: float = 1.0,
    box_size: Optional[float] = None,
) -> float:
    """
    Compute harmonic bond energy.

    Parameters:
        positions: List of position arrays or (N, 3) array for each bead.
        bonds: List of (atom1_idx, atom2_idx) tuples (0-indexed).
        k_bond: Bond spring constant.
        r0: Equilibrium bond length.
        box_size: Box size for periodic boundary conditions.

    Returns:
        Total bond energy.
    """
    energy = 0.0

    for i, j in bonds:
        r_vec = positions[j] - positions[i]

        # Apply minimum image convention for PBC
        if box_size is not None:
            r_vec = r_vec - box_size * np.round(r_vec / box_size)

        r = np.linalg.norm(r_vec)
        dr = r - r0
        energy += 0.5 * k_bond * dr * dr

    return energy

compute_bond_energy_vectorized(positions: np.ndarray, bonds: np.ndarray, k_bond: float = 100.0, r0: float = 1.0, box_size: Optional[float] = None) -> float

Vectorized bond energy calculation.

Parameters:

Name	Type	Description	Default
positions	ndarray	(N, 3) array of bead positions.	required
bonds	ndarray	(M, 2) array of bond pairs.	required
k_bond	float	Bond spring constant.	100.0
r0	float	Equilibrium bond length.	1.0
box_size	Optional[float]	Box size for periodic boundary conditions.	None

Returns:

Type	Description
float	Total bond energy.

Source code in AutoPoly/models/bead_spring.py

def compute_bond_energy_vectorized(
    positions: np.ndarray,
    bonds: np.ndarray,
    k_bond: float = 100.0,
    r0: float = 1.0,
    box_size: Optional[float] = None,
) -> float:
    """
    Vectorized bond energy calculation.

    Parameters:
        positions: (N, 3) array of bead positions.
        bonds: (M, 2) array of bond pairs.
        k_bond: Bond spring constant.
        r0: Equilibrium bond length.
        box_size: Box size for periodic boundary conditions.

    Returns:
        Total bond energy.
    """
    if len(bonds) == 0:
        return 0.0

    pos_i = positions[bonds[:, 0]]
    pos_j = positions[bonds[:, 1]]

    diff = pos_j - pos_i

    if box_size is not None:
        diff = diff - box_size * np.round(diff / box_size)

    r = np.linalg.norm(diff, axis=1)
    dr = r - r0
    return 0.5 * k_bond * np.sum(dr ** 2)

compute_total_energy(positions: Union[List[np.ndarray], np.ndarray], bonds: List[Tuple[int, int]], lj_sigma: float = 1.0, lj_epsilon: float = 1.0, lj_cutoff: float = 2.5, bond_k: float = 100.0, bond_r0: float = 1.0, box_size: Optional[float] = None) -> float

Compute total system energy for Metropolis criterion.

Parameters:

Name	Type	Description	Default
positions	Union[List[ndarray], ndarray]	List of position arrays or (N, 3) array for each bead.	required
bonds	List[Tuple[int, int]]	List of (atom1_idx, atom2_idx) tuples.	required
lj_sigma	float	LJ sigma parameter.	1.0
lj_epsilon	float	LJ epsilon parameter.	1.0
lj_cutoff	float	LJ cutoff in sigma units.	2.5
bond_k	float	Bond spring constant.	100.0
bond_r0	float	Equilibrium bond length.	1.0
box_size	Optional[float]	Box size for PBC.	None

Returns:

Type	Description
float	Total energy (LJ + bond).

Source code in AutoPoly/models/bead_spring.py

def compute_total_energy(
    positions: Union[List[np.ndarray], np.ndarray],
    bonds: List[Tuple[int, int]],
    lj_sigma: float = 1.0,
    lj_epsilon: float = 1.0,
    lj_cutoff: float = 2.5,
    bond_k: float = 100.0,
    bond_r0: float = 1.0,
    box_size: Optional[float] = None,
) -> float:
    """
    Compute total system energy for Metropolis criterion.

    Parameters:
        positions: List of position arrays or (N, 3) array for each bead.
        bonds: List of (atom1_idx, atom2_idx) tuples.
        lj_sigma: LJ sigma parameter.
        lj_epsilon: LJ epsilon parameter.
        lj_cutoff: LJ cutoff in sigma units.
        bond_k: Bond spring constant.
        bond_r0: Equilibrium bond length.
        box_size: Box size for PBC.

    Returns:
        Total energy (LJ + bond).
    """
    lj_energy = compute_lj_energy(
        positions, lj_sigma, lj_epsilon, lj_cutoff,
        exclude_bonded=bonds, box_size=box_size
    )
    bond_energy = compute_bond_energy(
        positions, bonds, bond_k, bond_r0, box_size
    )
    return lj_energy + bond_energy

build_exclusion_structures(n_beads: int, bonds: List[Tuple[int, int]]) -> Tuple[np.ndarray, List[set], List[List[int]]]

Build exclusion data structures for efficient energy calculations.

Parameters:

Name	Type	Description	Default
n_beads	int	Total number of beads.	required
bonds	List[Tuple[int, int]]	List of (atom1_idx, atom2_idx) tuples.	required

Returns:

Type	Description
ndarray	Tuple of:
List[set]	excluded_mask: (N, N) boolean array for vectorized LJ
List[List[int]]	excluded_neighbors: List of sets for local LJ energy
Tuple[ndarray, List[set], List[List[int]]]	bond_neighbors: List of lists for local bond energy

Source code in AutoPoly/models/bead_spring.py

def build_exclusion_structures(
    n_beads: int,
    bonds: List[Tuple[int, int]],
) -> Tuple[np.ndarray, List[set], List[List[int]]]:
    """
    Build exclusion data structures for efficient energy calculations.

    Parameters:
        n_beads: Total number of beads.
        bonds: List of (atom1_idx, atom2_idx) tuples.

    Returns:
        Tuple of:
        - excluded_mask: (N, N) boolean array for vectorized LJ
        - excluded_neighbors: List of sets for local LJ energy
        - bond_neighbors: List of lists for local bond energy
    """
    # Build excluded pair mask (N, N) for vectorized calculation
    excluded_mask = np.zeros((n_beads, n_beads), dtype=bool)
    for i, j in bonds:
        excluded_mask[i, j] = True
        excluded_mask[j, i] = True

    # Build per-bead exclusion sets for local LJ energy
    excluded_neighbors = [set() for _ in range(n_beads)]
    for i, j in bonds:
        excluded_neighbors[i].add(j)
        excluded_neighbors[j].add(i)

    # Build per-bead bond lists for local bond energy
    bond_neighbors = [[] for _ in range(n_beads)]
    for i, j in bonds:
        bond_neighbors[i].append(j)
        bond_neighbors[j].append(i)

    return excluded_mask, excluded_neighbors, bond_neighbors

metropolis_accept(delta_E: float, temperature: float = 1.0) -> bool

Metropolis acceptance criterion.

Accept if delta_E <= 0, else accept with probability exp(-delta_E/T).

Parameters:

Name	Type	Description	Default
delta_E	float	Energy change.	required
temperature	float	Reduced temperature.	1.0

Returns:

Type	Description
bool	True if move should be accepted.

Source code in AutoPoly/models/bead_spring.py

def metropolis_accept(delta_E: float, temperature: float = 1.0) -> bool:
    """
    Metropolis acceptance criterion.

    Accept if delta_E <= 0, else accept with probability exp(-delta_E/T).

    Parameters:
        delta_E: Energy change.
        temperature: Reduced temperature.

    Returns:
        True if move should be accepted.
    """
    if delta_E <= 0:
        return True
    return np.random.random() < np.exp(-delta_E / temperature)

mc_single_bead_displacement(positions: Union[List[np.ndarray], np.ndarray], bead_idx: int, max_disp: float = 0.5, box_size: Optional[float] = None) -> Tuple[Union[List[np.ndarray], np.ndarray], bool]

Displace a single bead by random vector.

Parameters:

Name	Type	Description	Default
positions	Union[List[ndarray], ndarray]	List of position arrays or (N, 3) array.	required
bead_idx	int	Index of bead to displace.	required
max_disp	float	Maximum displacement in each direction.	0.5
box_size	Optional[float]	Box size for PBC.	None

Returns:

Type	Description
Tuple[Union[List[ndarray], ndarray], bool]	Tuple of (new_positions, is_valid).

Source code in AutoPoly/models/bead_spring.py

def mc_single_bead_displacement(
    positions: Union[List[np.ndarray], np.ndarray],
    bead_idx: int,
    max_disp: float = 0.5,
    box_size: Optional[float] = None,
) -> Tuple[Union[List[np.ndarray], np.ndarray], bool]:
    """
    Displace a single bead by random vector.

    Parameters:
        positions: List of position arrays or (N, 3) array.
        bead_idx: Index of bead to displace.
        max_disp: Maximum displacement in each direction.
        box_size: Box size for PBC.

    Returns:
        Tuple of (new_positions, is_valid).
    """
    if isinstance(positions, np.ndarray):
        new_positions = positions.copy()
        delta = np.random.uniform(-max_disp, max_disp, 3)
        new_positions[bead_idx] = new_positions[bead_idx] + delta
        if box_size is not None:
            new_positions[bead_idx] = _apply_pbc(new_positions[bead_idx], box_size)
        return new_positions, True
    else:
        new_positions = [p.copy() for p in positions]
        delta = np.random.uniform(-max_disp, max_disp, 3)
        new_positions[bead_idx] = new_positions[bead_idx] + delta
        if box_size is not None:
            new_positions[bead_idx] = _apply_pbc(new_positions[bead_idx], box_size)
        return new_positions, True

mc_crankshaft_move(positions: Union[List[np.ndarray], np.ndarray], chain_start: int, chain_end: int, max_angle: float = 0.3, box_size: Optional[float] = None) -> Tuple[Union[List[np.ndarray], np.ndarray], bool]

Rotate beads between i and j around the i-j axis.

Keeps beads i and j fixed; rotates all beads strictly between them.

Parameters:

Name	Type	Description	Default
positions	Union[List[ndarray], ndarray]	List of position arrays or (N, 3) array.	required
chain_start	int	Start index of the chain.	required
chain_end	int	End index of the chain (exclusive).	required
max_angle	float	Maximum rotation angle in radians.	0.3
box_size	Optional[float]	Box size for PBC.	None

Returns:

Type	Description
Tuple[Union[List[ndarray], ndarray], bool]	Tuple of (new_positions, is_valid).

Source code in AutoPoly/models/bead_spring.py

def mc_crankshaft_move(
    positions: Union[List[np.ndarray], np.ndarray],
    chain_start: int,
    chain_end: int,
    max_angle: float = 0.3,
    box_size: Optional[float] = None,
) -> Tuple[Union[List[np.ndarray], np.ndarray], bool]:
    """
    Rotate beads between i and j around the i-j axis.

    Keeps beads i and j fixed; rotates all beads strictly between them.

    Parameters:
        positions: List of position arrays or (N, 3) array.
        chain_start: Start index of the chain.
        chain_end: End index of the chain (exclusive).
        max_angle: Maximum rotation angle in radians.
        box_size: Box size for PBC.

    Returns:
        Tuple of (new_positions, is_valid).
    """
    chain_length = chain_end - chain_start

    # Need at least 4 beads for crankshaft (2 endpoints + at least 2 in between)
    if chain_length < 4:
        return positions, False

    # Select two beads i, j with at least one bead between them
    i_local = np.random.randint(0, chain_length - 3)
    j_local = np.random.randint(i_local + 3, chain_length)

    i = chain_start + i_local
    j = chain_start + j_local

    is_numpy = isinstance(positions, np.ndarray)
    if is_numpy:
        new_positions = positions.copy()
        pos_i = positions[i]
        pos_j = positions[j]
    else:
        new_positions = [p.copy() for p in positions]
        pos_i = positions[i]
        pos_j = positions[j]

    # Rotation axis
    axis = pos_j - pos_i
    axis_len = np.linalg.norm(axis)
    if axis_len < 1e-10:
        return positions, False

    # Random rotation angle
    angle = np.random.uniform(-max_angle, max_angle)
    R = _rotation_matrix_around_axis(axis, angle)

    # Rotate beads between i and j (exclusive of endpoints)
    for k in range(i + 1, j):
        rel_pos = new_positions[k] - pos_i
        new_rel_pos = R @ rel_pos
        new_positions[k] = pos_i + new_rel_pos

        if box_size is not None:
            new_positions[k] = _apply_pbc(new_positions[k], box_size)

    return new_positions, True

mc_pivot_move(positions: Union[List[np.ndarray], np.ndarray], chain_start: int, chain_end: int, max_angle: float = 0.3, box_size: Optional[float] = None) -> Tuple[Union[List[np.ndarray], np.ndarray], bool]

Rotate arm (pivot_idx+1 to chain_end) around pivot bead.

Parameters:

Name	Type	Description	Default
positions	Union[List[ndarray], ndarray]	List of position arrays or (N, 3) array.	required
chain_start	int	Start index of the chain.	required
chain_end	int	End index of the chain (exclusive).	required
max_angle	float	Maximum rotation angle in radians.	0.3
box_size	Optional[float]	Box size for PBC.	None

Returns:

Type	Description
Tuple[Union[List[ndarray], ndarray], bool]	Tuple of (new_positions, is_valid).

Source code in AutoPoly/models/bead_spring.py

def mc_pivot_move(
    positions: Union[List[np.ndarray], np.ndarray],
    chain_start: int,
    chain_end: int,
    max_angle: float = 0.3,
    box_size: Optional[float] = None,
) -> Tuple[Union[List[np.ndarray], np.ndarray], bool]:
    """
    Rotate arm (pivot_idx+1 to chain_end) around pivot bead.

    Parameters:
        positions: List of position arrays or (N, 3) array.
        chain_start: Start index of the chain.
        chain_end: End index of the chain (exclusive).
        max_angle: Maximum rotation angle in radians.
        box_size: Box size for PBC.

    Returns:
        Tuple of (new_positions, is_valid).
    """
    chain_length = chain_end - chain_start

    # Need at least 2 beads
    if chain_length < 2:
        return positions, False

    # Select pivot (not the last bead)
    pivot_local = np.random.randint(0, chain_length - 1)
    pivot_idx = chain_start + pivot_local

    is_numpy = isinstance(positions, np.ndarray)
    if is_numpy:
        new_positions = positions.copy()
        pivot_pos = positions[pivot_idx].copy()
    else:
        new_positions = [p.copy() for p in positions]
        pivot_pos = positions[pivot_idx]

    # Random rotation
    R = _random_rotation_matrix(max_angle)

    # Rotate beads after pivot
    for k in range(pivot_idx + 1, chain_end):
        rel_pos = new_positions[k] - pivot_pos
        new_rel_pos = R @ rel_pos
        new_positions[k] = pivot_pos + new_rel_pos

        if box_size is not None:
            new_positions[k] = _apply_pbc(new_positions[k], box_size)

    return new_positions, True

mc_reptation_move(positions: Union[List[np.ndarray], np.ndarray], chain_start: int, chain_end: int, bond_length: float = 1.0, box_size: Optional[float] = None) -> Tuple[Union[List[np.ndarray], np.ndarray], bool]

Remove bead from one end (tail) and attach at other end (head).

Slithering snake move that maintains chain connectivity.

Parameters:

Name	Type	Description	Default
positions	Union[List[ndarray], ndarray]	List of position arrays or (N, 3) array.	required
chain_start	int	Start index of the chain.	required
chain_end	int	End index of the chain (exclusive).	required
bond_length	float	Bond length for new position.	1.0
box_size	Optional[float]	Box size for PBC.	None

Returns:

Type	Description
Tuple[Union[List[ndarray], ndarray], bool]	Tuple of (new_positions, is_valid).

Source code in AutoPoly/models/bead_spring.py

def mc_reptation_move(
    positions: Union[List[np.ndarray], np.ndarray],
    chain_start: int,
    chain_end: int,
    bond_length: float = 1.0,
    box_size: Optional[float] = None,
) -> Tuple[Union[List[np.ndarray], np.ndarray], bool]:
    """
    Remove bead from one end (tail) and attach at other end (head).

    Slithering snake move that maintains chain connectivity.

    Parameters:
        positions: List of position arrays or (N, 3) array.
        chain_start: Start index of the chain.
        chain_end: End index of the chain (exclusive).
        bond_length: Bond length for new position.
        box_size: Box size for PBC.

    Returns:
        Tuple of (new_positions, is_valid).
    """
    chain_length = chain_end - chain_start

    # Need at least 2 beads
    if chain_length < 2:
        return positions, False

    is_numpy = isinstance(positions, np.ndarray)
    if is_numpy:
        new_positions = positions.copy()
    else:
        new_positions = [p.copy() for p in positions]

    # Randomly choose direction
    forward = np.random.random() < 0.5

    if forward:
        # Move bead from tail (chain_start) to head (chain_end-1)
        # Shift all beads toward chain_start
        head_pos = positions[chain_end - 1].copy() if is_numpy else positions[chain_end - 1]

        for k in range(chain_start, chain_end - 1):
            new_positions[k] = positions[k + 1].copy()

        # Generate new position at head
        random_dir = np.random.randn(3)
        random_dir = random_dir / np.linalg.norm(random_dir)
        new_positions[chain_end - 1] = head_pos + bond_length * random_dir

    else:
        # Move bead from head (chain_end-1) to tail (chain_start)
        # Shift all beads toward chain_end
        tail_pos = positions[chain_start].copy() if is_numpy else positions[chain_start]

        for k in range(chain_end - 1, chain_start, -1):
            new_positions[k] = positions[k - 1].copy()

        # Generate new position at tail
        random_dir = np.random.randn(3)
        random_dir = random_dir / np.linalg.norm(random_dir)
        new_positions[chain_start] = tail_pos + bond_length * random_dir

    # Apply PBC
    if box_size is not None:
        for k in range(chain_start, chain_end):
            new_positions[k] = _apply_pbc(new_positions[k], box_size)

    return new_positions, True

place_chains_in_box(chain_positions: List[List[np.ndarray]], box_size: float, min_separation: float = 2.0, max_attempts: int = 1000) -> Tuple[List[np.ndarray], List[Tuple[int, int]]]

Place multiple chains randomly in periodic box.

Parameters:

Name	Type	Description	Default
chain_positions	List[List[ndarray]]	List of chain positions (each chain is list of np.array).	required
box_size	float	Cubic box side length.	required
min_separation	float	Minimum distance between chain COMs.	2.0
max_attempts	int	Maximum placement attempts per chain.	1000

Returns:

Type	Description
List[ndarray]	Tuple of (all_positions, chain_indices) where chain_indices is
List[Tuple[int, int]]	list of (start, end) indices for each chain.

Source code in AutoPoly/models/bead_spring.py

def place_chains_in_box(
    chain_positions: List[List[np.ndarray]],
    box_size: float,
    min_separation: float = 2.0,
    max_attempts: int = 1000,
) -> Tuple[List[np.ndarray], List[Tuple[int, int]]]:
    """
    Place multiple chains randomly in periodic box.

    Parameters:
        chain_positions: List of chain positions (each chain is list of np.array).
        box_size: Cubic box side length.
        min_separation: Minimum distance between chain COMs.
        max_attempts: Maximum placement attempts per chain.

    Returns:
        Tuple of (all_positions, chain_indices) where chain_indices is
        list of (start, end) indices for each chain.
    """
    all_positions = []
    chain_indices = []
    placed_coms = []

    for chain_idx, chain in enumerate(chain_positions):
        # Compute original COM
        orig_com = np.mean(chain, axis=0)

        # Center chain at origin
        centered_chain = [p - orig_com for p in chain]

        placed = False
        for attempt in range(max_attempts):
            # Random position in box
            new_com = np.random.uniform(-box_size/2, box_size/2, 3)

            # Check separation from existing chains
            too_close = False
            for existing_com in placed_coms:
                dist_vec = new_com - existing_com
                dist_vec = dist_vec - box_size * np.round(dist_vec / box_size)
                if np.linalg.norm(dist_vec) < min_separation:
                    too_close = True
                    break

            if too_close:
                continue

            # Random rotation
            R = _random_rotation_matrix(np.pi)

            # Transform chain
            start_idx = len(all_positions)
            for p in centered_chain:
                rotated = R @ p
                new_pos = rotated + new_com
                new_pos = _apply_pbc(new_pos, box_size)
                all_positions.append(new_pos)

            end_idx = len(all_positions)
            chain_indices.append((start_idx, end_idx))
            placed_coms.append(new_com)
            placed = True
            break

        if not placed:
            logger.warning(
                f"Could not place chain {chain_idx} with min_separation={min_separation}. "
                "Placing without separation constraint."
            )
            new_com = np.random.uniform(-box_size/2, box_size/2, 3)
            R = _random_rotation_matrix(np.pi)

            start_idx = len(all_positions)
            for p in centered_chain:
                rotated = R @ p
                new_pos = rotated + new_com
                new_pos = _apply_pbc(new_pos, box_size)
                all_positions.append(new_pos)

            end_idx = len(all_positions)
            chain_indices.append((start_idx, end_idx))
            placed_coms.append(new_com)

    return all_positions, chain_indices

mc_chain_translation(positions: Union[List[np.ndarray], np.ndarray], chain_indices: List[Tuple[int, int]], chain_idx: int, max_disp: float, box_size: float) -> Tuple[Union[List[np.ndarray], np.ndarray], bool]

Translate entire chain by random displacement with PBC.

Parameters:

Name	Type	Description	Default
positions	Union[List[ndarray], ndarray]	List of position arrays or (N, 3) array.	required
chain_indices	List[Tuple[int, int]]	List of (start, end) indices for each chain.	required
chain_idx	int	Index of chain to translate.	required
max_disp	float	Maximum displacement.	required
box_size	float	Box size for PBC.	required

Returns:

Type	Description
Tuple[Union[List[ndarray], ndarray], bool]	Tuple of (new_positions, is_valid).

Source code in AutoPoly/models/bead_spring.py

def mc_chain_translation(
    positions: Union[List[np.ndarray], np.ndarray],
    chain_indices: List[Tuple[int, int]],
    chain_idx: int,
    max_disp: float,
    box_size: float,
) -> Tuple[Union[List[np.ndarray], np.ndarray], bool]:
    """
    Translate entire chain by random displacement with PBC.

    Parameters:
        positions: List of position arrays or (N, 3) array.
        chain_indices: List of (start, end) indices for each chain.
        chain_idx: Index of chain to translate.
        max_disp: Maximum displacement.
        box_size: Box size for PBC.

    Returns:
        Tuple of (new_positions, is_valid).
    """
    start, end = chain_indices[chain_idx]
    is_numpy = isinstance(positions, np.ndarray)

    if is_numpy:
        new_positions = positions.copy()
    else:
        new_positions = [p.copy() for p in positions]

    delta = np.random.uniform(-max_disp, max_disp, 3)

    for k in range(start, end):
        new_positions[k] = new_positions[k] + delta
        new_positions[k] = _apply_pbc(new_positions[k], box_size)

    return new_positions, True

mc_chain_rotation(positions: Union[List[np.ndarray], np.ndarray], chain_indices: List[Tuple[int, int]], chain_idx: int, max_angle: float, box_size: Optional[float] = None) -> Tuple[Union[List[np.ndarray], np.ndarray], bool]

Rotate entire chain around its center of mass.

Parameters:

Name	Type	Description	Default
positions	Union[List[ndarray], ndarray]	List of position arrays or (N, 3) array.	required
chain_indices	List[Tuple[int, int]]	List of (start, end) indices for each chain.	required
chain_idx	int	Index of chain to rotate.	required
max_angle	float	Maximum rotation angle.	required
box_size	Optional[float]	Box size for PBC.	None

Returns:

Type	Description
Tuple[Union[List[ndarray], ndarray], bool]	Tuple of (new_positions, is_valid).

Source code in AutoPoly/models/bead_spring.py

def mc_chain_rotation(
    positions: Union[List[np.ndarray], np.ndarray],
    chain_indices: List[Tuple[int, int]],
    chain_idx: int,
    max_angle: float,
    box_size: Optional[float] = None,
) -> Tuple[Union[List[np.ndarray], np.ndarray], bool]:
    """
    Rotate entire chain around its center of mass.

    Parameters:
        positions: List of position arrays or (N, 3) array.
        chain_indices: List of (start, end) indices for each chain.
        chain_idx: Index of chain to rotate.
        max_angle: Maximum rotation angle.
        box_size: Box size for PBC.

    Returns:
        Tuple of (new_positions, is_valid).
    """
    start, end = chain_indices[chain_idx]
    is_numpy = isinstance(positions, np.ndarray)

    if is_numpy:
        new_positions = positions.copy()
        chain_pos = positions[start:end]
    else:
        new_positions = [p.copy() for p in positions]
        chain_pos = [positions[k] for k in range(start, end)]

    com = np.mean(chain_pos, axis=0)

    # Random rotation
    R = _random_rotation_matrix(max_angle)

    # Rotate around COM
    for k in range(start, end):
        rel_pos = new_positions[k] - com
        new_rel_pos = R @ rel_pos
        new_positions[k] = com + new_rel_pos

        if box_size is not None:
            new_positions[k] = _apply_pbc(new_positions[k], box_size)

    return new_positions, True

saw_grow_chain(n_beads: int, bond_length: float, start_position: np.ndarray, collision_detector: CollisionDetector, config: SAWConfig, topology: str = 'linear', chain_id_offset: int = 0) -> Tuple[Optional[List[np.ndarray]], int]

Grow a single polymer chain using Self-Avoiding Random Walk with backtracking.

Parameters:

Name	Type	Description	Default
n_beads	int	Number of beads in the chain.	required
bond_length	float	Distance between consecutive beads.	required
start_position	ndarray	Position of the first bead.	required
collision_detector	CollisionDetector	CollisionDetector for checking overlaps.	required
config	SAWConfig	SAWConfig with algorithm parameters.	required
topology	str	"linear" or "ring".	'linear'
chain_id_offset	int	Offset for bead IDs in collision detector.	0

Returns:

Type	Description
Tuple[Optional[List[ndarray]], int]	Tuple of (positions, backtracks_used) where positions is None if failed.

Source code in AutoPoly/models/bead_spring.py

def saw_grow_chain(
    n_beads: int,
    bond_length: float,
    start_position: np.ndarray,
    collision_detector: CollisionDetector,
    config: SAWConfig,
    topology: str = "linear",
    chain_id_offset: int = 0,
) -> Tuple[Optional[List[np.ndarray]], int]:
    """
    Grow a single polymer chain using Self-Avoiding Random Walk with backtracking.

    Parameters:
        n_beads: Number of beads in the chain.
        bond_length: Distance between consecutive beads.
        start_position: Position of the first bead.
        collision_detector: CollisionDetector for checking overlaps.
        config: SAWConfig with algorithm parameters.
        topology: "linear" or "ring".
        chain_id_offset: Offset for bead IDs in collision detector.

    Returns:
        Tuple of (positions, backtracks_used) where positions is None if failed.
    """
    if n_beads < 1:
        return [], 0

    positions = [start_position.copy()]
    backtrack_count = 0
    backtrack_depth = 1  # Current backtrack depth (exponential increase)

    # Add first bead to collision detector
    collision_detector.add_monomer(
        chain_id_offset,
        start_position,
        config.collision_sigma / 2
    )

    bead_idx = 1
    while bead_idx < n_beads:
        # Get previous direction for angle constraint
        if bead_idx >= 2:
            prev_direction = positions[bead_idx - 1] - positions[bead_idx - 2]
        else:
            prev_direction = None

        # For ring closure: last bead needs to connect back to first
        is_closing_ring = (topology == "ring" and bead_idx == n_beads - 1)

        # Generate trial positions
        n_trials = config.ring_closure_trials if is_closing_ring else config.n_trials
        trial_positions = _generate_trial_positions(
            positions[bead_idx - 1],
            bond_length,
            n_trials,
            prev_direction,
            config.bond_angle_min,
            config.bond_angle_max,
        )

        # Shuffle trials for randomness
        if len(trial_positions) > 0:
            np.random.shuffle(trial_positions)

        # Find valid position
        valid_position = None
        exclude_set = {chain_id_offset + bead_idx - 1}  # Exclude bonded neighbor

        for trial_pos in trial_positions:
            # Check collision with existing beads
            has_collision = collision_detector.check_collision(
                trial_pos,
                config.collision_sigma / 2,
                exclude_ids=exclude_set,
                tolerance=config.collision_tolerance,
            )

            if has_collision:
                continue

            # For ring closure: check distance to first bead
            if is_closing_ring:
                dist_to_first = np.linalg.norm(trial_pos - positions[0])
                if abs(dist_to_first - bond_length) > config.ring_closure_tolerance:
                    continue

            valid_position = trial_pos
            break

        if valid_position is not None:
            # Accept position
            positions.append(valid_position.copy())
            collision_detector.add_monomer(
                chain_id_offset + bead_idx,
                valid_position,
                config.collision_sigma / 2
            )
            bead_idx += 1
            backtrack_depth = 1  # Reset backtrack depth on success
        else:
            # Backtrack
            if backtrack_count >= config.max_total_backtracks:
                # Failed - remove all beads we added from detector
                for i in range(len(positions)):
                    collision_detector.remove_monomer(chain_id_offset + i)
                return None, backtrack_count

            # Determine how many beads to remove
            n_remove = min(backtrack_depth, bead_idx - 1, config.max_backtrack_depth)
            if n_remove == 0:
                # Can't backtrack further - failed
                for i in range(len(positions)):
                    collision_detector.remove_monomer(chain_id_offset + i)
                return None, backtrack_count

            # Remove beads from end
            for _ in range(n_remove):
                bead_idx -= 1
                collision_detector.remove_monomer(chain_id_offset + bead_idx)
                positions.pop()

            backtrack_count += 1
            backtrack_depth = min(backtrack_depth * 2, config.max_backtrack_depth)

    return positions, backtrack_count

saw_generate_multi_chain(n_chains: int, n_beads_per_chain: int, bond_length: float, box_size: float, config: SAWConfig, topology: str = 'linear', max_start_attempts: int = 100) -> Tuple[Optional[List[np.ndarray]], Optional[List[Tuple[int, int]]], Dict[str, any]]

Generate multiple polymer chains using SAW.

Parameters:

Name	Type	Description	Default
n_chains	int	Number of chains to generate.	required
n_beads_per_chain	int	Beads per chain.	required
bond_length	float	Bond length between consecutive beads.	required
box_size	float	Cubic box side length.	required
config	SAWConfig	SAWConfig with algorithm parameters.	required
topology	str	"linear" or "ring".	'linear'
max_start_attempts	int	Max attempts to find valid starting position per chain.	100

Returns:

Type	Description
Tuple[Optional[List[ndarray]], Optional[List[Tuple[int, int]]], Dict[str, any]]	Tuple of (all_positions, chain_indices, stats) or (None, None, stats) if failed.

Source code in AutoPoly/models/bead_spring.py

def saw_generate_multi_chain(
    n_chains: int,
    n_beads_per_chain: int,
    bond_length: float,
    box_size: float,
    config: SAWConfig,
    topology: str = "linear",
    max_start_attempts: int = 100,
) -> Tuple[Optional[List[np.ndarray]], Optional[List[Tuple[int, int]]], Dict[str, any]]:
    """
    Generate multiple polymer chains using SAW.

    Parameters:
        n_chains: Number of chains to generate.
        n_beads_per_chain: Beads per chain.
        bond_length: Bond length between consecutive beads.
        box_size: Cubic box side length.
        config: SAWConfig with algorithm parameters.
        topology: "linear" or "ring".
        max_start_attempts: Max attempts to find valid starting position per chain.

    Returns:
        Tuple of (all_positions, chain_indices, stats) or (None, None, stats) if failed.
    """
    half_box = box_size / 2
    box_bounds = ((-half_box, half_box), (-half_box, half_box), (-half_box, half_box))

    # Cell size should be at least collision diameter
    cell_size = max(config.collision_sigma * 2, bond_length * 2)
    collision_detector = CollisionDetector(box_bounds, cell_size)

    all_positions = []
    chain_indices = []
    total_backtracks = 0
    chains_completed = 0

    for chain_idx in range(n_chains):
        chain_offset = len(all_positions)

        # Find valid starting position
        start_found = False
        for attempt in range(max_start_attempts):
            # Random position within box (with margin)
            margin = config.collision_sigma * 2
            start_pos = np.random.uniform(
                -half_box + margin,
                half_box - margin,
                3
            )

            # Check if position is collision-free
            if not collision_detector.check_collision(
                start_pos,
                config.collision_sigma / 2,
                tolerance=config.collision_tolerance
            ):
                start_found = True
                break

        if not start_found:
            logger.warning(
                f"SAW: Could not find valid start position for chain {chain_idx}"
            )
            return None, None, {
                "success": False,
                "chains_completed": chains_completed,
                "total_backtracks": total_backtracks,
                "failure_reason": "start_position",
            }

        # Grow chain
        chain_positions, backtracks = saw_grow_chain(
            n_beads_per_chain,
            bond_length,
            start_pos,
            collision_detector,
            config,
            topology,
            chain_offset,
        )

        total_backtracks += backtracks

        if chain_positions is None:
            logger.warning(
                f"SAW: Failed to grow chain {chain_idx} after {backtracks} backtracks"
            )
            return None, None, {
                "success": False,
                "chains_completed": chains_completed,
                "total_backtracks": total_backtracks,
                "failure_reason": "chain_growth",
            }

        # Store chain
        all_positions.extend(chain_positions)
        chain_indices.append((chain_offset, chain_offset + len(chain_positions)))
        chains_completed += 1

    return all_positions, chain_indices, {
        "success": True,
        "chains_completed": chains_completed,
        "total_backtracks": total_backtracks,
    }

mc_equilibrate(positions: Union[List[np.ndarray], np.ndarray], bonds: List[Tuple[int, int]], chain_indices: Optional[List[Tuple[int, int]]] = None, n_steps: int = 10000, temperature: float = 1.0, move_weights: Optional[Dict[str, float]] = None, box_size: Optional[float] = None, lj_sigma: float = 1.0, lj_epsilon: float = 1.0, lj_cutoff: float = 2.5, bond_k: float = 100.0, bond_r0: float = 1.0, max_displacement: float = 0.5, max_angle: float = 0.3, verbose: bool = False) -> Tuple[List[np.ndarray], Dict[str, float]]

Pre-equilibrate polymer configuration using MC moves.

This implementation uses local energy updates for single-bead moves, reducing complexity from O(N²) to O(N) per step for the most common move type. For multi-bead moves, full energy recalculation is used.

Parameters:

Name	Type	Description	Default
positions	Union[List[ndarray], ndarray]	List of np.array or (N, 3) array, atom positions.	required
bonds	List[Tuple[int, int]]	List of (atom1_idx, atom2_idx) tuples (0-indexed).	required
chain_indices	Optional[List[Tuple[int, int]]]	List of (start, end) tuples for each chain. If None, treats entire system as one chain.	None
n_steps	int	Number of MC steps.	10000
temperature	float	Reduced temperature.	1.0
move_weights	Optional[Dict[str, float]]	Dict of move type probabilities.	None
box_size	Optional[float]	Apply PBC if specified.	None
lj_sigma	float	LJ sigma parameter.	1.0
lj_epsilon	float	LJ epsilon parameter.	1.0
lj_cutoff	float	LJ cutoff in sigma units.	2.5
bond_k	float	Bond spring constant.	100.0
bond_r0	float	Equilibrium bond length.	1.0
max_displacement	float	Max displacement for single bead moves.	0.5
max_angle	float	Max rotation angle for pivot/crankshaft.	0.3
verbose	bool	Print progress.	False

Returns:

Type	Description
Tuple[List[ndarray], Dict[str, float]]	Tuple of (equilibrated_positions, acceptance_stats).

Source code in AutoPoly/models/bead_spring.py

def mc_equilibrate(
    positions: Union[List[np.ndarray], np.ndarray],
    bonds: List[Tuple[int, int]],
    chain_indices: Optional[List[Tuple[int, int]]] = None,
    n_steps: int = 10000,
    temperature: float = 1.0,
    move_weights: Optional[Dict[str, float]] = None,
    box_size: Optional[float] = None,
    lj_sigma: float = 1.0,
    lj_epsilon: float = 1.0,
    lj_cutoff: float = 2.5,
    bond_k: float = 100.0,
    bond_r0: float = 1.0,
    max_displacement: float = 0.5,
    max_angle: float = 0.3,
    verbose: bool = False,
) -> Tuple[List[np.ndarray], Dict[str, float]]:
    """
    Pre-equilibrate polymer configuration using MC moves.

    This implementation uses local energy updates for single-bead moves,
    reducing complexity from O(N²) to O(N) per step for the most common
    move type. For multi-bead moves, full energy recalculation is used.

    Parameters:
        positions: List of np.array or (N, 3) array, atom positions.
        bonds: List of (atom1_idx, atom2_idx) tuples (0-indexed).
        chain_indices: List of (start, end) tuples for each chain.
            If None, treats entire system as one chain.
        n_steps: Number of MC steps.
        temperature: Reduced temperature.
        move_weights: Dict of move type probabilities.
        box_size: Apply PBC if specified.
        lj_sigma: LJ sigma parameter.
        lj_epsilon: LJ epsilon parameter.
        lj_cutoff: LJ cutoff in sigma units.
        bond_k: Bond spring constant.
        bond_r0: Equilibrium bond length.
        max_displacement: Max displacement for single bead moves.
        max_angle: Max rotation angle for pivot/crankshaft.
        verbose: Print progress.

    Returns:
        Tuple of (equilibrated_positions, acceptance_stats).
    """
    n_beads = len(positions)

    # Default chain indices
    if chain_indices is None:
        chain_indices = [(0, n_beads)]

    is_multi_chain = len(chain_indices) > 1

    # Default move weights - favor displacement moves for efficiency
    if move_weights is None:
        if is_multi_chain:
            move_weights = {
                "displacement": 0.4,  # Increased - most efficient move
                "crankshaft": 0.1,
                "pivot": 0.15,
                "reptation": 0.1,
                "chain_translation": 0.15,
                "chain_rotation": 0.1,
            }
        else:
            move_weights = {
                "displacement": 0.5,  # Increased - most efficient move
                "crankshaft": 0.15,
                "pivot": 0.2,
                "reptation": 0.15,
            }

    # Normalize weights
    total_weight = sum(move_weights.values())
    move_probs = {k: v / total_weight for k, v in move_weights.items()}

    # Build move list and cumulative probabilities
    moves = list(move_probs.keys())
    cum_probs = np.cumsum([move_probs[m] for m in moves])

    # Statistics
    move_attempts = {m: 0 for m in moves}
    move_accepts = {m: 0 for m in moves}

    # Convert to numpy array for efficient operations
    if isinstance(positions, list):
        current_positions = np.array([p.copy() for p in positions])
    else:
        current_positions = positions.copy()

    # Build exclusion structures once (O(N + M) where M = num bonds)
    excluded_mask, excluded_neighbors, bond_neighbors = build_exclusion_structures(
        n_beads, bonds
    )

    # Convert bonds to numpy array for vectorized bond energy
    bonds_array = np.array(bonds) if bonds else np.zeros((0, 2), dtype=int)

    # Compute initial total energy
    current_energy = compute_lj_energy_vectorized(
        current_positions, excluded_mask, lj_sigma, lj_epsilon, lj_cutoff, box_size
    ) + compute_bond_energy_vectorized(
        current_positions, bonds_array, bond_k, bond_r0, box_size
    )

    # Progress reporting interval
    report_interval = max(1, n_steps // 10)

    # MC loop
    for step in range(n_steps):
        # Select move type using searchsorted for efficiency
        r = np.random.random()
        move_idx = np.searchsorted(cum_probs, r)
        move_type = moves[move_idx]

        move_attempts[move_type] += 1

        # Select a chain for moves that operate on chains
        chain_idx = np.random.randint(len(chain_indices))
        start, end = chain_indices[chain_idx]

        # Perform move and compute energy change
        accepted = False

        if move_type == "displacement":
            # Single bead displacement - use LOCAL energy update (O(N) instead of O(N²))
            bead_idx = np.random.randint(start, end)

            # Store old position
            old_pos = current_positions[bead_idx].copy()

            # Compute old local energy (LJ + bonds involving this bead)
            old_lj = compute_local_lj_energy(
                current_positions, bead_idx, excluded_neighbors[bead_idx],
                lj_sigma, lj_epsilon, lj_cutoff, box_size
            )
            old_bond = compute_local_bond_energy(
                current_positions, bead_idx, bond_neighbors[bead_idx],
                bond_k, bond_r0, box_size
            )

            # Apply displacement in-place
            delta = np.random.uniform(-max_displacement, max_displacement, 3)
            current_positions[bead_idx] = old_pos + delta
            if box_size is not None:
                current_positions[bead_idx] = _apply_pbc(current_positions[bead_idx], box_size)

            # Compute new local energy
            new_lj = compute_local_lj_energy(
                current_positions, bead_idx, excluded_neighbors[bead_idx],
                lj_sigma, lj_epsilon, lj_cutoff, box_size
            )
            new_bond = compute_local_bond_energy(
                current_positions, bead_idx, bond_neighbors[bead_idx],
                bond_k, bond_r0, box_size
            )

            delta_E = (new_lj + new_bond) - (old_lj + old_bond)

            if metropolis_accept(delta_E, temperature):
                current_energy += delta_E
                accepted = True
            else:
                # Reject: restore old position
                current_positions[bead_idx] = old_pos

        else:
            # Multi-bead moves - use full energy recalculation
            new_positions = current_positions
            is_valid = False

            if move_type == "crankshaft":
                new_positions, is_valid = mc_crankshaft_move(
                    current_positions, start, end, max_angle, box_size
                )

            elif move_type == "pivot":
                new_positions, is_valid = mc_pivot_move(
                    current_positions, start, end, max_angle, box_size
                )

            elif move_type == "reptation":
                new_positions, is_valid = mc_reptation_move(
                    current_positions, start, end, bond_r0, box_size
                )

            elif move_type == "chain_translation":
                if box_size is not None:
                    new_positions, is_valid = mc_chain_translation(
                        current_positions, chain_indices, chain_idx,
                        max_displacement, box_size
                    )

            elif move_type == "chain_rotation":
                new_positions, is_valid = mc_chain_rotation(
                    current_positions, chain_indices, chain_idx,
                    max_angle, box_size
                )

            if is_valid:
                # Compute new total energy using vectorized functions
                new_energy = compute_lj_energy_vectorized(
                    new_positions, excluded_mask, lj_sigma, lj_epsilon, lj_cutoff, box_size
                ) + compute_bond_energy_vectorized(
                    new_positions, bonds_array, bond_k, bond_r0, box_size
                )

                delta_E = new_energy - current_energy
                if metropolis_accept(delta_E, temperature):
                    current_positions = new_positions
                    current_energy = new_energy
                    accepted = True

        if accepted:
            move_accepts[move_type] += 1

        # Progress report
        if verbose and (step + 1) % report_interval == 0:
            logger.info(
                f"MC step {step + 1}/{n_steps}, Energy: {current_energy:.4f}"
            )

    # Compute acceptance rates
    acceptance_stats = {}
    for m in moves:
        if move_attempts[m] > 0:
            acceptance_stats[m] = move_accepts[m] / move_attempts[m]
        else:
            acceptance_stats[m] = 0.0

    if verbose:
        logger.info(f"Final energy: {current_energy:.4f}")
        logger.info(f"Acceptance rates: {acceptance_stats}")

    # Convert back to list for backward compatibility
    result_positions = [current_positions[i].copy() for i in range(n_beads)]
    return result_positions, acceptance_stats