feat: add cyclic peptide support to relative position encoding

examples/example_inference_inputs/query_multimer_cyclic.json

@@ -0,0 +1,23 @@
+{
+    "queries": {
+        "mdm2": {
+            "chains": [
+                {
+                    "molecule_type": "protein",
+                    "chain_ids": [
+                        "A"
+                    ],
+                    "sequence": "EQETLVRPKPLLLKLLKSVGAQKDTYTMKEVLFYLGQYIMTKRLYDEKQQHIVYCSNDLLGDLFGVPSFSVKEHRKIYTMIYRNLVVVNQQE"
+                },
+                {
+                    "molecule_type": "protein",
+                    "chain_ids": [
+                        "B"
+                    ],
+                    "sequence": "EALKKESLLL",
+                    "cyclic": true
+                }
+            ]
+        }
+    }
+}

openfold3/core/data/pipelines/featurization/structure.py

@@ -191,6 +191,10 @@ def featurize_structure_of3(
             num_atoms_per_token=features["num_atoms_per_token"],
         )
 
+        features["cyclic_mask"] = torch.tensor(
+            atom_array.is_cyclic[token_starts], dtype=torch.bool
+        )
+
     # Ground-truth-specific features
     # TODO reorganize GT feature logic
     if is_gt:

openfold3/core/data/primitives/structure/query.py

@@ -630,6 +630,10 @@ def structure_with_ref_mols_from_query(query: Query) -> StructureWithReferenceMo
                 "entity_id",
                 np.repeat(entity_to_id[representation], len(segment_atom_array)),
             )
+            segment_atom_array.set_annotation(
+                "is_cyclic",
+                np.repeat(chain.cyclic, len(segment_atom_array)),
+            )
 
             # Append atom array to end
             if atom_array is None:

openfold3/core/utils/relpos.py

@@ -17,6 +17,39 @@
 from openfold3.core.utils.tensor_utils import binned_one_hot
 
 
+def cyclic_offset(residue_index: torch.Tensor) -> torch.Tensor:
+    """Calculate the cyclic offset for the given residue index.
+    Args:
+        residue_index:
+            [*, N_token] Token index
+
+    Returns:
+        cyclic_offset_array:
+            [N_token, N_token] token by token index distances
+
+    Example:
+        >>> import torch
+        >>> residue_index = torch.tensor([0,1,2,3,4,5,6])
+        >>> cyclic_offset_array = cyclic_offset(residue_index)
+        >>> cyclic_offset_array:
+            tensor([[ 0, -1, -2, -3,  2,  1],
+                    [ 1,  0, -1, -2, -3,  2],
+                    [ 2,  1,  0, -1, -2, -3],
+                    [-3,  2,  1,  0, -1, -2],
+                    [-2, -3,  2,  1,  0, -1],
+                    [-1, -2, -3,  2,  1,  0]], device='cuda:0', dtype=torch.int32)
+
+    """
+    peptide_length = residue_index.shape[0]
+    cyclic_offset_array = torch.zeros((peptide_length, peptide_length))
+    cyc_row = torch.arange(0, -peptide_length, -1)
+    pc = int(torch.round(torch.tensor(peptide_length / 2)))  # Get centre
+    cyc_row[pc + 1 :] = torch.arange(len(cyc_row[pc + 1 :]), 0, -1)
+    for i in range(len(cyclic_offset_array)):
+        cyclic_offset_array[i] = torch.roll(cyc_row, i)
+    return cyclic_offset_array.type(torch.int).to(residue_index.device)
+
+
 def relpos_complex(
     batch: dict, max_relative_idx: int, max_relative_chain: int
 ) -> torch.Tensor:
@@ -34,13 +67,19 @@ def relpos_complex(
     """
     res_idx = batch["residue_index"]
     asym_id = batch["asym_id"]
+    cyclic_mask = batch["cyclic_mask"]
     entity_id = batch["entity_id"]
     same_chain = asym_id[..., None] == asym_id[..., None, :]
+
     same_res = res_idx[..., None] == res_idx[..., None, :]
     same_entity = entity_id[..., None] == entity_id[..., None, :]
 
     def relpos(
-        pos: torch.Tensor, condition: torch.BoolTensor, rel_clip_idx: int
+        pos: torch.Tensor,
+        condition: torch.BoolTensor,
+        rel_clip_idx: int,
+        cyclic_mask: torch.Tensor,
+        asym_id: torch.Tensor,
     ) -> torch.Tensor:
         """
         Args:
@@ -50,11 +89,34 @@ def relpos(
                 [*, N_token, N_token] Condition for clipping
             rel_clip_idx:
                 Max idx for clipping (max_relative_idx or max_relative_chain)
+            cyclic_mask:
+                [*, N_token] Boolean tensor for cyclic residues
+            asym_id:
+                [*, N_token] Used by cyclic mask for multi-chain cyclic
         Returns:
             rel_pos:
                 [*, N_token, N_token, 2 * rel_clip_idx + 2] Relative position embedding
         """
         offset = pos[..., None] - pos[..., None, :]
+        if cyclic_mask is not None and cyclic_mask.any():
+            for chain_id in torch.unique(asym_id):
+                chain_cyclic_mask = cyclic_mask & (asym_id == chain_id)
+                pair_cyclic = (
+                    chain_cyclic_mask[..., None] & chain_cyclic_mask[..., None, :]
+                )
+
+                if not pair_cyclic.any():
+                    continue
+                cyc_mask_1d = cyclic_mask.view(-1, cyclic_mask.shape[-1])[0]
+                cyc_indices = torch.where(
+                    cyc_mask_1d & (asym_id.squeeze(0) == chain_id)
+                )[0]
+                cyc_pos = pos.view(-1)[cyc_indices]
+                cyc_off = cyclic_offset(cyc_pos).to(dtype=offset.dtype)
+                full_cyc_off = offset.new_zeros(offset.shape)
+                full_cyc_off[..., cyc_indices[:, None], cyc_indices[None, :]] = cyc_off
+                offset = torch.where(pair_cyclic, full_cyc_off, offset)
+
         clipped_offset = torch.clamp(offset + rel_clip_idx, min=0, max=2 * rel_clip_idx)
         final_offset = torch.where(
             condition,
@@ -71,16 +133,27 @@ def relpos(
 
         return rel_pos
 
-    rel_pos = relpos(pos=res_idx, condition=same_chain, rel_clip_idx=max_relative_idx)
+    rel_pos = relpos(
+        pos=res_idx,
+        condition=same_chain,
+        rel_clip_idx=max_relative_idx,
+        cyclic_mask=cyclic_mask,
+        asym_id=asym_id,
+    )
+
     rel_token = relpos(
         pos=batch["token_index"],
         condition=same_chain & same_res,
         rel_clip_idx=max_relative_idx,
+        cyclic_mask=cyclic_mask,
+        asym_id=asym_id,
     )
     rel_chain = relpos(
         pos=batch["sym_id"],
         condition=same_entity,
         rel_clip_idx=max_relative_chain,
+        cyclic_mask=cyclic_mask,
+        asym_id=asym_id,
     )
 
     same_entity = same_entity[..., None].to(dtype=rel_pos.dtype)

openfold3/projects/of3_all_atom/config/inference_query_format.py

@@ -74,6 +74,7 @@ class Chain(BaseModel):
         Annotated[list[str | None], BeforeValidator(_ensure_list)] | None
     ) = None
     sdf_file_path: FilePath | None = None
+    cyclic: bool = False
 
     @field_serializer("molecule_type", return_type=str)
     def serialize_enum_name(self, v: MoleculeType, _info):

openfold3/tests/core/utils/test_relpos.py

@@ -0,0 +1,173 @@
+# Copyright 2026 AlQuraishi Laboratory
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+import torch
+
+from openfold3.core.utils.relpos import cyclic_offset, relpos_complex
+
+
+def _make_batch(n_token, asym_ids, cyclic_mask, batch_size=1):
+    """Build a minimal feature dict for relpos_complex."""
+    residue_index = torch.arange(n_token).unsqueeze(0).repeat(batch_size, 1)
+    token_index = torch.arange(n_token).unsqueeze(0).repeat(batch_size, 1)
+    asym_id = (
+        torch.tensor(asym_ids, dtype=torch.int32).unsqueeze(0).repeat(batch_size, 1)
+    )
+    entity_id = asym_id.clone()
+    sym_id = torch.ones(batch_size, n_token, dtype=torch.int32)
+    cm = torch.tensor(cyclic_mask, dtype=torch.bool).unsqueeze(0).repeat(batch_size, 1)
+    return {
+        "residue_index": residue_index,
+        "token_index": token_index,
+        "asym_id": asym_id,
+        "entity_id": entity_id,
+        "sym_id": sym_id,
+        "cyclic_mask": cm,
+    }
+
+
+class TestCyclicOffset:
+    def test_diagonal_is_zero(self):
+        # A residue's distance to itself is always 0.
+        idx = torch.arange(8)
+        off = cyclic_offset(idx)
+        assert (off.diagonal() == 0).all()
+
+    def test_antisymmetry(self):
+        # Cyclic offset is antisymmetric: off[i,j] == -off[j,i].
+        # Equivalently, the magnitudes are symmetric.
+        idx = torch.arange(6)
+        off = cyclic_offset(idx)
+        assert torch.equal(off.abs(), off.T.abs())
+
+    def test_max_distance_at_midpoint(self):
+        # For an even-length chain the maximum offset magnitude is peptide_length // 2.
+        n = 10
+        idx = torch.arange(n)
+        off = cyclic_offset(idx)
+        assert int(off.abs().max()) == n // 2
+
+    def test_odd_length(self):
+        # For an odd-length chain all entries are <= (n-1)//2 away from 0.
+        n = 7
+        idx = torch.arange(n)
+        off = cyclic_offset(idx)
+        assert int(off.abs().max()) <= (n - 1) // 2 + 1
+
+    def test_output_shape(self):
+        n = 5
+        idx = torch.arange(n)
+        off = cyclic_offset(idx)
+        assert off.shape == (n, n)
+
+    def test_values_wrap_correctly(self):
+        # For n=6 the cyclic row starting at 0 should be:
+        #   [0, -1, -2, -3, 2, 1]
+        # i.e. going forward costs +dist, wrapping back costs negative dist past midpoint.
+        n = 6
+        idx = torch.arange(n)
+        off = cyclic_offset(idx)
+        row0 = off[0].tolist()
+        # Distance from 0 to 3 (midpoint) is -3; to 4 wraps back: +2; to 5: +1.
+        assert row0[0] == 0
+        assert row0[1] == -1
+        assert row0[2] == -2
+        assert row0[3] == -3
+        assert row0[4] == 2
+        assert row0[5] == 1
+
+
+class TestRelposComplex:
+    MAX_IDX = 32
+    MAX_CHAIN = 2
+
+    def _relpos(self, batch):
+        return relpos_complex(batch, self.MAX_IDX, self.MAX_CHAIN)
+
+    @pytest.mark.parametrize("is_cyclic", [True, False])
+    def test_relpos_shape(self, is_cyclic):
+        n = 10
+        batch = _make_batch(n, [1] * n, [is_cyclic] * n)
+        out = self._relpos(batch)
+        expected_last = (2 * self.MAX_IDX + 2) * 2 + 1 + (2 * self.MAX_CHAIN + 2)
+        assert out.shape == (1, n, n, expected_last)
+
+    def test_cyclic_changes_encoding_vs_linear(self):
+        # A cyclic chain should produce different rel-pos encodings than a linear one.
+        n = 10
+        linear_batch = _make_batch(n, [1] * n, [False] * n)
+        cyclic_batch = _make_batch(n, [1] * n, [True] * n)
+        linear_out = self._relpos(linear_batch)
+        cyclic_out = self._relpos(cyclic_batch)
+        assert not torch.equal(linear_out, cyclic_out)
+
+    def test_cyclic_self_pairs_get_center_bin(self):
+        # Self-pairs always have offset=0, which clamps to MAX_IDX → one-hot at bin MAX_IDX.
+        n = 8
+        batch = _make_batch(n, [1] * n, [True] * n)
+        out = self._relpos(batch)
+        rel_pos_slice = out[0, :, :, : 2 * self.MAX_IDX + 2]
+        diag = torch.stack([rel_pos_slice[i, i] for i in range(n)])
+        assert (diag[:, self.MAX_IDX] == 1.0).all()
+        assert (diag[:, self.MAX_IDX] == diag.sum(dim=-1)).all()
+
+    def test_non_cyclic_chain_unchanged_in_multimer(self):
+        # In a multimer where only chain 2 is cyclic, chain 1's encoding should
+        # be identical to an all-linear batch.
+        n_lin = 6
+        n_cyc = 6
+        n = n_lin + n_cyc
+        asym_ids = [1] * n_lin + [2] * n_cyc
+        cyclic_mask_mixed = [False] * n_lin + [True] * n_cyc
+        cyclic_mask_none = [False] * n
+
+        mixed_batch = _make_batch(n, asym_ids, cyclic_mask_mixed)
+        linear_batch = _make_batch(n, asym_ids, cyclic_mask_none)
+
+        mixed_out = self._relpos(mixed_batch)
+        linear_out = self._relpos(linear_batch)
+
+        # Chain 1 rows/cols (indices 0..n_lin-1) should be unchanged.
+        ch1_mixed = mixed_out[0, :n_lin, :n_lin, :]
+        ch1_linear = linear_out[0, :n_lin, :n_lin, :]
+        assert torch.equal(ch1_mixed, ch1_linear)
+
+    def test_cross_chain_pairs_unaffected_by_cyclic(self):
+        # Cross-chain token pairs should use the "different chain" sentinel regardless
+        # of whether either chain is cyclic (same_chain=False → clipped to sentinel).
+        n_a = 5
+        n_b = 5
+        n = n_a + n_b
+        asym_ids = [1] * n_a + [2] * n_b
+        cyclic_mask = [True] * n_a + [False] * n_b
+
+        batch = _make_batch(n, asym_ids, cyclic_mask)
+        out = self._relpos(batch)
+
+        # The cross-chain rel_pos block uses the sentinel bin (2*MAX_IDX+1).
+        # After one-hot encoding each row sums to 1; check the sentinel column.
+        rel_pos_slice = out[0, :n_a, n_a:, : 2 * self.MAX_IDX + 2]
+        sentinel_col = 2 * self.MAX_IDX + 1
+        assert (rel_pos_slice[..., sentinel_col] == 1).all()
+
+    def test_no_cyclic_mask_does_not_raise(self):
+        # cyclic_mask all-False should run without error and match linear baseline.
+        n = 8
+        batch_false = _make_batch(n, [1] * n, [False] * n)
+        batch_zeros = _make_batch(n, [1] * n, [False] * n)
+        batch_zeros["cyclic_mask"] = torch.zeros(1, n, dtype=torch.bool)
+        out_false = self._relpos(batch_false)
+        out_zeros = self._relpos(batch_zeros)
+        assert torch.equal(out_false, out_zeros)

openfold3/tests/data_utils.py

@@ -256,6 +256,7 @@ def random_of3_features(batch_size, n_token, n_msa, n_templ, is_eval=False):
         # Additional features
         "token_mask": token_mask.unsqueeze(0).repeat((batch_size, 1)),
         "atom_mask": atom_mask.unsqueeze(0).repeat((batch_size, 1)),
+        "cyclic_mask": torch.zeros(n_token).repeat((batch_size, 1)).type(torch.long),
         "start_atom_index": start_atom_index.unsqueeze(0).repeat((batch_size, 1)).int(),
         "num_atoms_per_token": num_atoms_per_token.unsqueeze(0)
         .repeat((batch_size, 1))

openfold3/tests/test_data/snapshots/test_triangular_attention/nvidia/_snapshot_env.json

@@ -1,7 +1,7 @@
 {
-  "torch_version": "2.10.0+cu130",
-  "python_version": "3.13.12",
-  "cuda_version": "13.0",
-  "cudnn_version": "91501",
-  "gpu_name": "NVIDIA GB10"
+  "torch_version": "2.10.0",
+  "python_version": "3.14.6",
+  "cuda_version": "12.9",
+  "cudnn_version": "91002",
+  "gpu_name": "NVIDIA RTX 6000 Ada Generation"
 }

openfold3/tests/test_data/snapshots/test_triangular_multiplicative_update/nvidia/_snapshot_env.json

@@ -1,7 +1,7 @@
 {
-  "torch_version": "2.10.0+cu130",
-  "python_version": "3.13.12",
-  "cuda_version": "13.0",
-  "cudnn_version": "91501",
-  "gpu_name": "NVIDIA GB10"
+  "torch_version": "2.10.0",
+  "python_version": "3.14.6",
+  "cuda_version": "12.9",
+  "cudnn_version": "91002",
+  "gpu_name": "NVIDIA RTX 6000 Ada Generation"
 }