What changed, file by file.

def cyclize_relpos(relpos, cyclic_mask, L):
    """CYCLIC MODIFICATION — NEW function (JAX).

    Port of RFpeptide's Embeddings.PositionalEncoding2D.forward cyclic
    wrap-around (Rettie et al. 2025, rfpeptides branch).
    """
    import jax.numpy as jnp

    abs_sep = jnp.abs(relpos)
    sign = jnp.sign(relpos)
    wrapped = jnp.where(abs_sep > L / 2.0,
                        sign * (abs_sep - L),
                        relpos)
    both_cyclic = cyclic_mask[:, None] & cyclic_mask[None, :]
    if relpos.ndim == 3 and both_cyclic.ndim == 2:
        both_cyclic = both_cyclic[..., None]
    return jnp.where(both_cyclic, wrapped, relpos)

def add_cyclic_bond_loss(self, weight=1.0,
                         target_N_C=TARGET_N_C_DISTANCE,
                         target_CA_CA=TARGET_CA_CA_DISTANCE,
                         target_omega=TARGET_OMEGA_DEG,
                         nc_tol=NC_TOLERANCE,
                         caca_tol=CA_CA_TOLERANCE,
                         omega_tol=OMEGA_TOLERANCE_DEG):
    """CYCLIC MODIFICATION — NEW function (JAX).

    Replace the soft add_termini_distance_loss (threshold 7 A CA-CA, ELU
    penalty) with a hard amide-bond geometry loss. Three terms:
      1. N-C bond length   — quadratic when |dist(N_1, C_L) - 1.33| > 0.10
      2. CA-CA across bond — quadratic when |dist(CA_1, CA_L) - 3.80| > 0.50
      3. omega dihedral    — quadratic when |omega_deg - 180| > 30
    """
    def loss_fn(inputs, outputs):
        xyz = outputs["structure_module"]
        atoms = xyz["final_atom_positions"][-self._binder_len:]

        N_IDX = residue_constants.atom_order["N"]
        CA_IDX = residue_constants.atom_order["CA"]
        C_IDX = residue_constants.atom_order["C"]

        N1 = atoms[0, N_IDX];   CA1 = atoms[0, CA_IDX]
        CL = atoms[-1, C_IDX];  CAL = atoms[-1, CA_IDX]

        nc_dist = jnp.linalg.norm(N1 - CL)
        nc_dev = jax.nn.relu(jnp.abs(nc_dist - target_N_C) - nc_tol)
        nc_loss = nc_dev ** 2

        caca_dist = jnp.linalg.norm(CA1 - CAL)
        caca_dev = jax.nn.relu(jnp.abs(caca_dist - target_CA_CA) - caca_tol)
        caca_loss = caca_dev ** 2

        omega_rad = _cyclic_dihedral_rad(CAL, CL, N1, CA1)
        omega_deg = jnp.degrees(omega_rad)
        omega_wrapped = ((omega_deg + 180.0) % 360.0) - 180.0
        omega_dev = jax.nn.relu(jnp.abs(omega_wrapped - target_omega) - omega_tol)
        omega_loss = omega_dev ** 2

        total = nc_loss + caca_loss + 0.01 * omega_loss
        return {"cyclic_bond": total}

    self._callbacks["model"]["loss"].append(loss_fn)
    self.opt["weights"]["cyclic_bond"] = weight

def declare_cyclic_bond(pose, binder_chain: str = "B") -> bool:
    """Declare a covalent peptide bond between the binder chain's residue-1
    N atom and residue-L C atom on a PyRosetta pose."""
    try:
        from pyrosetta.rosetta.core.kinematics import DeclareBond
    except Exception:
        return False

    try:
        from pyrosetta.rosetta.core.pose import get_chain_id_from_chain
        cid = get_chain_id_from_chain(binder_chain, pose)
        chain_begin = pose.conformation().chain_begin(cid)
        chain_end = pose.conformation().chain_end(cid)

        declare = DeclareBond()
        declare.set(chain_end, chain_begin, "C", "N", False)
        declare.apply(pose)
        return True
    except Exception:
        return False

# CYCLIC MODIFICATION: pre-compute the {first_res, last_res} pair
# of the binder chain and skip those atom pairs in the clash loop.
cyclic_neighbor_pair = None
if cyclic:
    try:
        binder_chain_obj = structure[0][binder_chain]
        binder_residues = [r for r in binder_chain_obj
                           if is_aa(r, standard=True)]
        if len(binder_residues) >= 2:
            first_res = binder_residues[0].id[1]
            last_res = binder_residues[-1].id[1]
            cyclic_neighbor_pair = frozenset((first_res, last_res))
    except Exception:
        cyclic_neighbor_pair = None

# ... inside the pair loop:
if (cyclic_neighbor_pair is not None
        and chain_i == chain_j == binder_chain
        and frozenset((res_i, res_j)) == cyclic_neighbor_pair):
    continue   # skip — this is the cyclic amide bond, not a clash

def cyclic_bond_conect_records(pdb_file: str,
                               binder_chain: str = "B") -> str:
    """Read a PDB and return a CONECT record string that documents
    the head-to-tail cyclic amide bond between residue-1 N and
    residue-L C."""
    from Bio.PDB import PDBParser, is_aa

    parser = PDBParser(QUIET=True)
    structure = parser.get_structure("conect", pdb_file)
    chain = structure[0][binder_chain]
    residues = [r for r in chain if is_aa(r, standard=True)]
    if len(residues) < 2:
        return ""

    first = residues[0]; last = residues[-1]
    if "N" not in first or "C" not in last:
        return ""

    n_serial = first["N"].get_serial_number()
    c_serial = last["C"].get_serial_number()
    if n_serial is None or c_serial is None:
        return ""

    return f"CONECT{n_serial:>5d}{c_serial:>5d}{'':>70s}\n".rstrip() + "\n"