diff --git a/examples/example_inference_inputs/query_multimer_cyclic.json b/examples/example_inference_inputs/query_multimer_cyclic.json new file mode 100644 index 000000000..1c99c20c6 --- /dev/null +++ b/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 + } + ] + } + } +} \ No newline at end of file diff --git a/openfold3/core/data/pipelines/featurization/structure.py b/openfold3/core/data/pipelines/featurization/structure.py index 96f4ef73c..4dd9f3b55 100644 --- a/openfold3/core/data/pipelines/featurization/structure.py +++ b/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: diff --git a/openfold3/core/data/primitives/structure/query.py b/openfold3/core/data/primitives/structure/query.py index 30ec1106d..fc4195abf 100644 --- a/openfold3/core/data/primitives/structure/query.py +++ b/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: diff --git a/openfold3/core/utils/relpos.py b/openfold3/core/utils/relpos.py index ab0208be5..7dab07292 100644 --- a/openfold3/core/utils/relpos.py +++ b/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) diff --git a/openfold3/projects/of3_all_atom/config/inference_query_format.py b/openfold3/projects/of3_all_atom/config/inference_query_format.py index 5276bd310..ecccc44b0 100644 --- a/openfold3/projects/of3_all_atom/config/inference_query_format.py +++ b/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): diff --git a/openfold3/tests/core/utils/test_relpos.py b/openfold3/tests/core/utils/test_relpos.py new file mode 100644 index 000000000..c32bb009a --- /dev/null +++ b/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) diff --git a/openfold3/tests/data_utils.py b/openfold3/tests/data_utils.py index 974fdb1d0..f8be948b3 100644 --- a/openfold3/tests/data_utils.py +++ b/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)) diff --git a/openfold3/tests/test_data/snapshots/test_structure_from_query/test_structure_from_query_non_canonical_peptide_.npz b/openfold3/tests/test_data/snapshots/test_structure_from_query/test_structure_from_query_non_canonical_peptide_.npz index 59d93bc06..5a6fcc03b 100644 Binary files a/openfold3/tests/test_data/snapshots/test_structure_from_query/test_structure_from_query_non_canonical_peptide_.npz and b/openfold3/tests/test_data/snapshots/test_structure_from_query/test_structure_from_query_non_canonical_peptide_.npz differ diff --git a/openfold3/tests/test_data/snapshots/test_structure_from_query/test_structure_from_query_standard_peptide_.npz b/openfold3/tests/test_data/snapshots/test_structure_from_query/test_structure_from_query_standard_peptide_.npz index f50fd3f50..b25c3d64a 100644 Binary files a/openfold3/tests/test_data/snapshots/test_structure_from_query/test_structure_from_query_standard_peptide_.npz and b/openfold3/tests/test_data/snapshots/test_structure_from_query/test_structure_from_query_standard_peptide_.npz differ diff --git a/openfold3/tests/test_data/snapshots/test_triangular_attention/nvidia/_snapshot_env.json b/openfold3/tests/test_data/snapshots/test_triangular_attention/nvidia/_snapshot_env.json index 429cc9a55..b382370d6 100644 --- a/openfold3/tests/test_data/snapshots/test_triangular_attention/nvidia/_snapshot_env.json +++ b/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" } diff --git a/openfold3/tests/test_data/snapshots/test_triangular_attention/nvidia/test_shape_cuda_False_.npz b/openfold3/tests/test_data/snapshots/test_triangular_attention/nvidia/test_shape_cuda_False_.npz index b71e3084c..372bf7241 100644 Binary files a/openfold3/tests/test_data/snapshots/test_triangular_attention/nvidia/test_shape_cuda_False_.npz and b/openfold3/tests/test_data/snapshots/test_triangular_attention/nvidia/test_shape_cuda_False_.npz differ diff --git a/openfold3/tests/test_data/snapshots/test_triangular_attention/nvidia/test_shape_cuda_True_.npz b/openfold3/tests/test_data/snapshots/test_triangular_attention/nvidia/test_shape_cuda_True_.npz index 9b1e196e7..5af6adb70 100644 Binary files a/openfold3/tests/test_data/snapshots/test_triangular_attention/nvidia/test_shape_cuda_True_.npz and b/openfold3/tests/test_data/snapshots/test_triangular_attention/nvidia/test_shape_cuda_True_.npz differ diff --git a/openfold3/tests/test_data/snapshots/test_triangular_multiplicative_update/nvidia/_snapshot_env.json b/openfold3/tests/test_data/snapshots/test_triangular_multiplicative_update/nvidia/_snapshot_env.json index 429cc9a55..b382370d6 100644 --- a/openfold3/tests/test_data/snapshots/test_triangular_multiplicative_update/nvidia/_snapshot_env.json +++ b/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" } diff --git a/openfold3/tests/test_data/snapshots/test_triangular_multiplicative_update/nvidia/test_shape_cuda_.npz b/openfold3/tests/test_data/snapshots/test_triangular_multiplicative_update/nvidia/test_shape_cuda_.npz index f63faf8b7..ca6959fe7 100644 Binary files a/openfold3/tests/test_data/snapshots/test_triangular_multiplicative_update/nvidia/test_shape_cuda_.npz and b/openfold3/tests/test_data/snapshots/test_triangular_multiplicative_update/nvidia/test_shape_cuda_.npz differ diff --git a/openfold3/tests/test_data/structure_from_query/structure-w-ref-mols_non-std-peptide.pkl b/openfold3/tests/test_data/structure_from_query/structure-w-ref-mols_non-std-peptide.pkl new file mode 100644 index 000000000..c385110e0 Binary files /dev/null and b/openfold3/tests/test_data/structure_from_query/structure-w-ref-mols_non-std-peptide.pkl differ diff --git a/openfold3/tests/test_data/structure_from_query/structure-w-ref-mols_std-peptide.pkl b/openfold3/tests/test_data/structure_from_query/structure-w-ref-mols_std-peptide.pkl new file mode 100644 index 000000000..7d61fc7f4 Binary files /dev/null and b/openfold3/tests/test_data/structure_from_query/structure-w-ref-mols_std-peptide.pkl differ diff --git a/openfold3/tests/test_diffusion_conditioning.py b/openfold3/tests/test_diffusion_conditioning.py index 3d912e008..b3e4a5027 100644 --- a/openfold3/tests/test_diffusion_conditioning.py +++ b/openfold3/tests/test_diffusion_conditioning.py @@ -52,6 +52,7 @@ def test_without_n_sample_channel(self): "sym_id": torch.zeros((batch_size, n_token)), "asym_id": torch.zeros((batch_size, n_token)), "entity_id": torch.zeros((batch_size, n_token)), + "cyclic_mask": torch.zeros((batch_size, 1, n_token)), } si, zij = dc( @@ -101,6 +102,7 @@ def test_with_different_schedule(self): "sym_id": torch.zeros((batch_size, 1, n_token)), "asym_id": torch.zeros((batch_size, 1, n_token)), "entity_id": torch.zeros((batch_size, 1, n_token)), + "cyclic_mask": torch.zeros((batch_size, 1, n_token)), } si, zij = dc( @@ -149,6 +151,7 @@ def test_with_same_schedule(self): "sym_id": torch.zeros((batch_size, 1, n_token)), "asym_id": torch.zeros((batch_size, 1, n_token)), "entity_id": torch.zeros((batch_size, 1, n_token)), + "cyclic_mask": torch.zeros((batch_size, 1, n_token)), } si, zij = dc(