Hugging Face's logo

Synthyra
/
Boltz2

Feature Extraction
Transformers
PyTorch
Safetensors
boltz2_automodel
protein-language-model
fastplms
custom_code
Model card Files
xet
 Community
lhallee commited on
Commit
961083d
·
verified ·
1 Parent(s): 827d9ec

Upload folder using huggingface_hub

Browse files
Files changed (32) hide show
  1. __init__.py +11 -11
  2. cif_writer.py +126 -126
  3. config.json +139 -139
  4. entrypoint_setup.py +22 -0
  5. minimal_featurizer.py +487 -487
  6. minimal_structures.py +20 -20
  7. modeling_boltz2.py +1 -0
  8. pytorch_model.bin +3 -0
  9. vb_const.py +1184 -1184
  10. vb_layers_attention.py +136 -136
  11. vb_layers_attentionv2.py +111 -111
  12. vb_layers_confidence_utils.py +231 -231
  13. vb_layers_dropout.py +34 -34
  14. vb_layers_initialize.py +86 -86
  15. vb_layers_outer_product_mean.py +98 -98
  16. vb_layers_pair_averaging.py +135 -135
  17. vb_layers_pairformer.py +335 -335
  18. vb_layers_transition.py +78 -78
  19. vb_layers_triangular_mult.py +215 -215
  20. vb_loss_diffusionv2.py +138 -138
  21. vb_modules_confidencev2.py +498 -498
  22. vb_modules_diffusion_conditioning.py +116 -116
  23. vb_modules_diffusionv2.py +693 -693
  24. vb_modules_encodersv2.py +565 -565
  25. vb_modules_transformersv2.py +262 -262
  26. vb_modules_trunkv2.py +828 -828
  27. vb_modules_utils.py +303 -303
  28. vb_potentials_potentials.py +787 -787
  29. vb_potentials_schedules.py +37 -37
  30. vb_tri_attn_attention.py +189 -189
  31. vb_tri_attn_primitives.py +411 -411
  32. vb_tri_attn_utils.py +380 -380
__init__.py CHANGED
@@ -1,11 +1,11 @@
1
- from boltz_automodel.modeling_boltz2 import (
2
- Boltz2Config,
3
- Boltz2Model,
4
- Boltz2StructureOutput,
5
- )
6
-
7
- __all__ = [
8
- "Boltz2Config",
9
- "Boltz2Model",
10
- "Boltz2StructureOutput",
11
- ]
 
1
+ from boltz_fastplms.modeling_boltz2 import (
2
+ Boltz2Config,
3
+ Boltz2Model,
4
+ Boltz2StructureOutput,
5
+ )
6
+
7
+ __all__ = [
8
+ "Boltz2Config",
9
+ "Boltz2Model",
10
+ "Boltz2StructureOutput",
11
+ ]
cif_writer.py CHANGED
@@ -1,126 +1,126 @@
1
- from pathlib import Path
2
- from typing import Optional
3
-
4
- import numpy as np
5
- import torch
6
-
7
- from .minimal_structures import ProteinStructureTemplate
8
-
9
-
10
- def _confidence_per_atom(
11
- plddt: Optional[torch.Tensor],
12
- atom_to_residue: list[int],
13
- num_atoms: int,
14
- sample_index: int,
15
- ) -> np.ndarray:
16
- if plddt is None:
17
- return np.ones((num_atoms,), dtype=np.float32) * 100.0
18
-
19
- values = plddt.detach().cpu()
20
- if values.ndim == 1:
21
- values = values.unsqueeze(0)
22
- assert values.ndim == 2, "Expected pLDDT with shape [samples, tokens/atoms]."
23
- assert sample_index < values.shape[0], "sample_index out of range for pLDDT."
24
-
25
- selected = values[sample_index]
26
- if selected.shape[0] == num_atoms:
27
- return (selected.numpy() * 100.0).astype(np.float32)
28
-
29
- num_residues = max(atom_to_residue) + 1
30
- if selected.shape[0] == num_residues:
31
- expanded = np.zeros((num_atoms,), dtype=np.float32)
32
- selected_np = selected.numpy()
33
- for atom_idx, residue_idx in enumerate(atom_to_residue):
34
- expanded[atom_idx] = selected_np[residue_idx] * 100.0
35
- return expanded
36
-
37
- return np.ones((num_atoms,), dtype=np.float32) * 100.0
38
-
39
-
40
- def write_cif(
41
- structure_template: ProteinStructureTemplate,
42
- atom_coords: torch.Tensor,
43
- atom_mask: torch.Tensor,
44
- output_path: str,
45
- plddt: Optional[torch.Tensor] = None,
46
- sample_index: int = 0,
47
- ) -> str:
48
- coords = atom_coords.detach().cpu()
49
- if coords.ndim == 2:
50
- coords = coords.unsqueeze(0)
51
- assert coords.ndim == 3, "Expected coordinates with shape [samples, atoms, 3]."
52
- assert sample_index < coords.shape[0], "sample_index out of range."
53
- selected_coords_tensor = coords[sample_index]
54
- all_non_finite = torch.logical_not(torch.isfinite(selected_coords_tensor))
55
- assert not torch.any(all_non_finite), (
56
- "CIF export received non-finite coordinates. "
57
- f"Non-finite count: {int(all_non_finite.sum().item())}"
58
- )
59
- selected_coords = selected_coords_tensor.numpy()
60
-
61
- mask = atom_mask.detach().cpu()
62
- if mask.ndim == 2:
63
- mask = mask[0]
64
- assert mask.ndim == 1, "Expected atom mask with shape [atoms]."
65
- assert mask.shape[0] == selected_coords.shape[0], "Atom mask/coord size mismatch."
66
- assert torch.any(mask > 0), "Atom mask has no valid atoms for CIF export."
67
- valid_non_finite = torch.logical_not(torch.isfinite(selected_coords_tensor[mask > 0]))
68
- assert not torch.any(valid_non_finite), (
69
- "CIF export has non-finite coordinates in unmasked atoms. "
70
- f"Non-finite count: {int(valid_non_finite.sum().item())}"
71
- )
72
-
73
- b_iso = _confidence_per_atom(
74
- plddt=plddt,
75
- atom_to_residue=structure_template.atom_residue_index,
76
- num_atoms=structure_template.num_atoms,
77
- sample_index=sample_index,
78
- )
79
- assert b_iso.shape[0] == structure_template.num_atoms
80
-
81
- lines = [
82
- "data_boltz2_prediction",
83
- "#",
84
- "loop_",
85
- "_atom_site.group_PDB",
86
- "_atom_site.id",
87
- "_atom_site.type_symbol",
88
- "_atom_site.label_atom_id",
89
- "_atom_site.label_comp_id",
90
- "_atom_site.label_asym_id",
91
- "_atom_site.label_seq_id",
92
- "_atom_site.Cartn_x",
93
- "_atom_site.Cartn_y",
94
- "_atom_site.Cartn_z",
95
- "_atom_site.occupancy",
96
- "_atom_site.B_iso_or_equiv",
97
- "_atom_site.pdbx_PDB_model_num",
98
- ]
99
-
100
- atom_id = 1
101
- for idx in range(structure_template.num_atoms):
102
- if mask[idx] <= 0:
103
- continue
104
-
105
- residue_idx = structure_template.atom_residue_index[idx]
106
- residue_name = structure_template.residue_names[residue_idx]
107
- atom_name = structure_template.atom_names[idx]
108
- element = structure_template.atom_elements[idx]
109
- chain_id = structure_template.atom_chain_id[idx]
110
- x_val, y_val, z_val = selected_coords[idx].tolist()
111
- b_factor = float(b_iso[idx])
112
-
113
- line = (
114
- f"ATOM {atom_id} {element} {atom_name} {residue_name} {chain_id} "
115
- f"{residue_idx + 1} {x_val:.3f} {y_val:.3f} {z_val:.3f} 1.00 {b_factor:.2f} 1"
116
- )
117
- lines.append(line)
118
- atom_id += 1
119
-
120
- lines.append("#")
121
- text = "\n".join(lines) + "\n"
122
-
123
- out_path = Path(output_path)
124
- out_path.parent.mkdir(parents=True, exist_ok=True)
125
- out_path.write_text(text, encoding="utf-8")
126
- return str(out_path)
 
1
+ from pathlib import Path
2
+ from typing import Optional
3
+
4
+ import numpy as np
5
+ import torch
6
+
7
+ from .minimal_structures import ProteinStructureTemplate
8
+
9
+
10
+ def _confidence_per_atom(
11
+ plddt: Optional[torch.Tensor],
12
+ atom_to_residue: list[int],
13
+ num_atoms: int,
14
+ sample_index: int,
15
+ ) -> np.ndarray:
16
+ if plddt is None:
17
+ return np.ones((num_atoms,), dtype=np.float32) * 100.0
18
+
19
+ values = plddt.detach().cpu()
20
+ if values.ndim == 1:
21
+ values = values.unsqueeze(0)
22
+ assert values.ndim == 2, "Expected pLDDT with shape [samples, tokens/atoms]."
23
+ assert sample_index < values.shape[0], "sample_index out of range for pLDDT."
24
+
25
+ selected = values[sample_index]
26
+ if selected.shape[0] == num_atoms:
27
+ return (selected.numpy() * 100.0).astype(np.float32)
28
+
29
+ num_residues = max(atom_to_residue) + 1
30
+ if selected.shape[0] == num_residues:
31
+ expanded = np.zeros((num_atoms,), dtype=np.float32)
32
+ selected_np = selected.numpy()
33
+ for atom_idx, residue_idx in enumerate(atom_to_residue):
34
+ expanded[atom_idx] = selected_np[residue_idx] * 100.0
35
+ return expanded
36
+
37
+ return np.ones((num_atoms,), dtype=np.float32) * 100.0
38
+
39
+
40
+ def write_cif(
41
+ structure_template: ProteinStructureTemplate,
42
+ atom_coords: torch.Tensor,
43
+ atom_mask: torch.Tensor,
44
+ output_path: str,
45
+ plddt: Optional[torch.Tensor] = None,
46
+ sample_index: int = 0,
47
+ ) -> str:
48
+ coords = atom_coords.detach().cpu()
49
+ if coords.ndim == 2:
50
+ coords = coords.unsqueeze(0)
51
+ assert coords.ndim == 3, "Expected coordinates with shape [samples, atoms, 3]."
52
+ assert sample_index < coords.shape[0], "sample_index out of range."
53
+ selected_coords_tensor = coords[sample_index]
54
+ all_non_finite = torch.logical_not(torch.isfinite(selected_coords_tensor))
55
+ assert not torch.any(all_non_finite), (
56
+ "CIF export received non-finite coordinates. "
57
+ f"Non-finite count: {int(all_non_finite.sum().item())}"
58
+ )
59
+ selected_coords = selected_coords_tensor.numpy()
60
+
61
+ mask = atom_mask.detach().cpu()
62
+ if mask.ndim == 2:
63
+ mask = mask[0]
64
+ assert mask.ndim == 1, "Expected atom mask with shape [atoms]."
65
+ assert mask.shape[0] == selected_coords.shape[0], "Atom mask/coord size mismatch."
66
+ assert torch.any(mask > 0), "Atom mask has no valid atoms for CIF export."
67
+ valid_non_finite = torch.logical_not(torch.isfinite(selected_coords_tensor[mask > 0]))
68
+ assert not torch.any(valid_non_finite), (
69
+ "CIF export has non-finite coordinates in unmasked atoms. "
70
+ f"Non-finite count: {int(valid_non_finite.sum().item())}"
71
+ )
72
+
73
+ b_iso = _confidence_per_atom(
74
+ plddt=plddt,
75
+ atom_to_residue=structure_template.atom_residue_index,
76
+ num_atoms=structure_template.num_atoms,
77
+ sample_index=sample_index,
78
+ )
79
+ assert b_iso.shape[0] == structure_template.num_atoms
80
+
81
+ lines = [
82
+ "data_boltz2_prediction",
83
+ "#",
84
+ "loop_",
85
+ "_atom_site.group_PDB",
86
+ "_atom_site.id",
87
+ "_atom_site.type_symbol",
88
+ "_atom_site.label_atom_id",
89
+ "_atom_site.label_comp_id",
90
+ "_atom_site.label_asym_id",
91
+ "_atom_site.label_seq_id",
92
+ "_atom_site.Cartn_x",
93
+ "_atom_site.Cartn_y",
94
+ "_atom_site.Cartn_z",
95
+ "_atom_site.occupancy",
96
+ "_atom_site.B_iso_or_equiv",
97
+ "_atom_site.pdbx_PDB_model_num",
98
+ ]
99
+
100
+ atom_id = 1
101
+ for idx in range(structure_template.num_atoms):
102
+ if mask[idx] <= 0:
103
+ continue
104
+
105
+ residue_idx = structure_template.atom_residue_index[idx]
106
+ residue_name = structure_template.residue_names[residue_idx]
107
+ atom_name = structure_template.atom_names[idx]
108
+ element = structure_template.atom_elements[idx]
109
+ chain_id = structure_template.atom_chain_id[idx]
110
+ x_val, y_val, z_val = selected_coords[idx].tolist()
111
+ b_factor = float(b_iso[idx])
112
+
113
+ line = (
114
+ f"ATOM {atom_id} {element} {atom_name} {residue_name} {chain_id} "
115
+ f"{residue_idx + 1} {x_val:.3f} {y_val:.3f} {z_val:.3f} 1.00 {b_factor:.2f} 1"
116
+ )
117
+ lines.append(line)
118
+ atom_id += 1
119
+
120
+ lines.append("#")
121
+ text = "\n".join(lines) + "\n"
122
+
123
+ out_path = Path(output_path)
124
+ out_path.parent.mkdir(parents=True, exist_ok=True)
125
+ out_path.write_text(text, encoding="utf-8")
126
+ return str(out_path)
config.json CHANGED
@@ -1,139 +1,139 @@
1
- {
2
- "architectures": [
3
- "Boltz2Model"
4
- ],
5
- "auto_map": {
6
- "AutoConfig": "modeling_boltz2.Boltz2Config",
7
- "AutoModel": "modeling_boltz2.Boltz2Model"
8
- },
9
- "core_kwargs": {
10
- "alpha_pae": 1,
11
- "atom_feature_dim": 388,
12
- "atom_s": 128,
13
- "atom_z": 16,
14
- "atoms_per_window_keys": 128,
15
- "atoms_per_window_queries": 32,
16
- "bond_type_feature": true,
17
- "conditioning_cutoff_max": 20.0,
18
- "conditioning_cutoff_min": 4.0,
19
- "confidence_model_args": {
20
- "add_s_input_to_s": true,
21
- "add_s_to_z_prod": true,
22
- "add_z_input_to_z": true,
23
- "confidence_args": {
24
- "num_pae_bins": 64,
25
- "num_pde_bins": 64,
26
- "num_plddt_bins": 50,
27
- "relative_confidence": "none",
28
- "use_separate_heads": true
29
- },
30
- "max_dist": 22,
31
- "no_trunk_feats": false,
32
- "num_dist_bins": 64,
33
- "pairformer_args": {
34
- "dropout": 0.25,
35
- "num_blocks": 8,
36
- "num_heads": 16,
37
- "v2": true
38
- },
39
- "use_gaussian": false,
40
- "use_miniformer": false,
41
- "use_s_diffusion": false
42
- },
43
- "confidence_prediction": true,
44
- "cyclic_pos_enc": true,
45
- "diffusion_process_args": {
46
- "P_mean": -1.2,
47
- "P_std": 1.5,
48
- "alignment_reverse_diff": true,
49
- "coordinate_augmentation": true,
50
- "gamma_0": 0.8,
51
- "gamma_min": 1.0,
52
- "mse_rotational_alignment": true,
53
- "noise_scale": 1.003,
54
- "rho": 7,
55
- "sigma_data": 16.0,
56
- "sigma_max": 160.0,
57
- "sigma_min": 0.0001,
58
- "step_scale": 1.5,
59
- "step_scale_random": [
60
- 1.0,
61
- 1.1,
62
- 1.2,
63
- 1.3,
64
- 1.4,
65
- 1.5
66
- ],
67
- "synchronize_sigmas": true
68
- },
69
- "embedder_args": {
70
- "add_cyclic_flag": true,
71
- "add_method_conditioning": true,
72
- "add_modified_flag": true,
73
- "add_mol_type_feat": true,
74
- "atom_encoder_depth": 3,
75
- "atom_encoder_heads": 4
76
- },
77
- "fix_sym_check": true,
78
- "msa_args": {
79
- "activation_checkpointing": true,
80
- "miniformer_blocks": false,
81
- "msa_blocks": 4,
82
- "msa_dropout": 0.15,
83
- "msa_s": 64,
84
- "pairwise_head_width": 32,
85
- "pairwise_num_heads": 4,
86
- "use_paired_feature": true,
87
- "use_trifast": true,
88
- "z_dropout": 0.25
89
- },
90
- "num_bins": 64,
91
- "pairformer_args": {
92
- "activation_checkpointing": true,
93
- "dropout": 0.25,
94
- "num_blocks": 64,
95
- "num_heads": 16,
96
- "post_layer_norm": false,
97
- "use_trifast": true,
98
- "v2": true
99
- },
100
- "run_trunk_and_structure": true,
101
- "score_model_args": {
102
- "activation_checkpointing": true,
103
- "atom_decoder_depth": 3,
104
- "atom_decoder_heads": 4,
105
- "atom_encoder_depth": 3,
106
- "atom_encoder_heads": 4,
107
- "conditioning_transition_layers": 2,
108
- "dim_fourier": 256,
109
- "sigma_data": 16,
110
- "token_transformer_depth": 24,
111
- "token_transformer_heads": 16,
112
- "transformer_post_ln": false
113
- },
114
- "skip_run_structure": false,
115
- "steering_args": {
116
- "contact_guidance_update": false,
117
- "fk_lambda": 4.0,
118
- "fk_resampling_interval": 3,
119
- "fk_steering": false,
120
- "num_gd_steps": 16,
121
- "num_particles": 3,
122
- "physical_guidance_update": false
123
- },
124
- "token_level_confidence": true,
125
- "token_s": 384,
126
- "token_z": 128,
127
- "use_atom_backbone_feat": false,
128
- "use_kernels": false,
129
- "use_no_atom_char": false,
130
- "use_residue_feats_atoms": false
131
- },
132
- "default_diffusion_samples": 5,
133
- "default_recycling_steps": 3,
134
- "default_sampling_steps": 200,
135
- "dtype": "float32",
136
- "model_type": "boltz2_automodel",
137
- "num_bins": 64,
138
- "transformers_version": "5.1.0"
139
- }
 
1
+ {
2
+ "architectures": [
3
+ "Boltz2Model"
4
+ ],
5
+ "auto_map": {
6
+ "AutoConfig": "modeling_boltz2.Boltz2Config",
7
+ "AutoModel": "modeling_boltz2.Boltz2Model"
8
+ },
9
+ "core_kwargs": {
10
+ "alpha_pae": 1,
11
+ "atom_feature_dim": 388,
12
+ "atom_s": 128,
13
+ "atom_z": 16,
14
+ "atoms_per_window_keys": 128,
15
+ "atoms_per_window_queries": 32,
16
+ "bond_type_feature": true,
17
+ "conditioning_cutoff_max": 20.0,
18
+ "conditioning_cutoff_min": 4.0,
19
+ "confidence_model_args": {
20
+ "add_s_input_to_s": true,
21
+ "add_s_to_z_prod": true,
22
+ "add_z_input_to_z": true,
23
+ "confidence_args": {
24
+ "num_pae_bins": 64,
25
+ "num_pde_bins": 64,
26
+ "num_plddt_bins": 50,
27
+ "relative_confidence": "none",
28
+ "use_separate_heads": true
29
+ },
30
+ "max_dist": 22,
31
+ "no_trunk_feats": false,
32
+ "num_dist_bins": 64,
33
+ "pairformer_args": {
34
+ "dropout": 0.25,
35
+ "num_blocks": 8,
36
+ "num_heads": 16,
37
+ "v2": true
38
+ },
39
+ "use_gaussian": false,
40
+ "use_miniformer": false,
41
+ "use_s_diffusion": false
42
+ },
43
+ "confidence_prediction": true,
44
+ "cyclic_pos_enc": true,
45
+ "diffusion_process_args": {
46
+ "P_mean": -1.2,
47
+ "P_std": 1.5,
48
+ "alignment_reverse_diff": true,
49
+ "coordinate_augmentation": true,
50
+ "gamma_0": 0.8,
51
+ "gamma_min": 1.0,
52
+ "mse_rotational_alignment": true,
53
+ "noise_scale": 1.003,
54
+ "rho": 7,
55
+ "sigma_data": 16.0,
56
+ "sigma_max": 160.0,
57
+ "sigma_min": 0.0001,
58
+ "step_scale": 1.5,
59
+ "step_scale_random": [
60
+ 1.0,
61
+ 1.1,
62
+ 1.2,
63
+ 1.3,
64
+ 1.4,
65
+ 1.5
66
+ ],
67
+ "synchronize_sigmas": true
68
+ },
69
+ "embedder_args": {
70
+ "add_cyclic_flag": true,
71
+ "add_method_conditioning": true,
72
+ "add_modified_flag": true,
73
+ "add_mol_type_feat": true,
74
+ "atom_encoder_depth": 3,
75
+ "atom_encoder_heads": 4
76
+ },
77
+ "fix_sym_check": true,
78
+ "msa_args": {
79
+ "activation_checkpointing": true,
80
+ "miniformer_blocks": false,
81
+ "msa_blocks": 4,
82
+ "msa_dropout": 0.15,
83
+ "msa_s": 64,
84
+ "pairwise_head_width": 32,
85
+ "pairwise_num_heads": 4,
86
+ "use_paired_feature": true,
87
+ "use_trifast": true,
88
+ "z_dropout": 0.25
89
+ },
90
+ "num_bins": 64,
91
+ "pairformer_args": {
92
+ "activation_checkpointing": true,
93
+ "dropout": 0.25,
94
+ "num_blocks": 64,
95
+ "num_heads": 16,
96
+ "post_layer_norm": false,
97
+ "use_trifast": true,
98
+ "v2": true
99
+ },
100
+ "run_trunk_and_structure": true,
101
+ "score_model_args": {
102
+ "activation_checkpointing": true,
103
+ "atom_decoder_depth": 3,
104
+ "atom_decoder_heads": 4,
105
+ "atom_encoder_depth": 3,
106
+ "atom_encoder_heads": 4,
107
+ "conditioning_transition_layers": 2,
108
+ "dim_fourier": 256,
109
+ "sigma_data": 16,
110
+ "token_transformer_depth": 24,
111
+ "token_transformer_heads": 16,
112
+ "transformer_post_ln": false
113
+ },
114
+ "skip_run_structure": false,
115
+ "steering_args": {
116
+ "contact_guidance_update": false,
117
+ "fk_lambda": 4.0,
118
+ "fk_resampling_interval": 3,
119
+ "fk_steering": false,
120
+ "num_gd_steps": 16,
121
+ "num_particles": 3,
122
+ "physical_guidance_update": false
123
+ },
124
+ "token_level_confidence": true,
125
+ "token_s": 384,
126
+ "token_z": 128,
127
+ "use_atom_backbone_feat": false,
128
+ "use_kernels": false,
129
+ "use_no_atom_char": false,
130
+ "use_residue_feats_atoms": false
131
+ },
132
+ "default_diffusion_samples": 5,
133
+ "default_recycling_steps": 3,
134
+ "default_sampling_steps": 200,
135
+ "model_type": "boltz2_automodel",
136
+ "num_bins": 64,
137
+ "torch_dtype": "float32",
138
+ "transformers_version": "4.48.1"
139
+ }
entrypoint_setup.py ADDED
@@ -0,0 +1,22 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import torch._inductor.config as inductor_config
3
+ import torch._dynamo as dynamo
4
+
5
+ # Enable TensorFloat32 tensor cores for float32 matmul (Ampere+ GPUs)
6
+ # Provides significant speedup with minimal precision loss
7
+ torch.set_float32_matmul_precision('high')
8
+
9
+ # Enable TF32 for matrix multiplications and cuDNN operations
10
+ torch.backends.cuda.matmul.allow_tf32 = True
11
+ torch.backends.cudnn.allow_tf32 = True
12
+
13
+ # Enable cuDNN autotuner - finds fastest algorithms for your hardware
14
+ # Best when input sizes are consistent; may slow down first iterations
15
+ torch.backends.cudnn.benchmark = True
16
+
17
+ # Deterministic operations off for speed (set True if reproducibility needed)
18
+ torch.backends.cudnn.deterministic = False
19
+ inductor_config.max_autotune_gemm_backends = "ATEN,CUTLASS,FBGEMM"
20
+
21
+ dynamo.config.capture_scalar_outputs = True
22
+ torch._dynamo.config.recompile_limit = 64
minimal_featurizer.py CHANGED
@@ -1,487 +1,487 @@
1
- import math
2
- from typing import Dict, List, Tuple
3
-
4
- import numpy as np
5
- import torch
6
- from torch.nn.functional import one_hot
7
-
8
- from .minimal_structures import ProteinStructureTemplate
9
- from . import vb_const as const
10
-
11
-
12
- _ELEMENT_TO_Z = {
13
- "H": 1,
14
- "C": 6,
15
- "N": 7,
16
- "O": 8,
17
- "P": 15,
18
- "S": 16,
19
- }
20
-
21
-
22
- def _normalize_sequence(sequence: str) -> str:
23
- seq = sequence.strip().upper()
24
- assert len(seq) > 0, "Amino acid sequence must be non-empty."
25
- for aa in seq:
26
- assert aa in const.prot_letter_to_token, f"Unsupported residue code '{aa}'."
27
- return seq
28
-
29
-
30
- def _atom_name_to_element(atom_name: str) -> str:
31
- name = atom_name.strip().upper()
32
- if len(name) == 0:
33
- return "C"
34
- if name[0].isdigit():
35
- name = name[1:]
36
- if len(name) >= 2 and name[0:2] in ("CL", "BR", "FE", "MG", "ZN", "NA", "CA"):
37
- return name[0]
38
- return name[0]
39
-
40
-
41
- def _atom_name_to_codes(atom_name: str) -> torch.Tensor:
42
- clipped = atom_name.strip()[:4]
43
- vals = [ord(ch) - 32 for ch in clipped]
44
- while len(vals) < 4:
45
- vals.append(0)
46
- out = torch.tensor(vals, dtype=torch.long)
47
- assert torch.all(out >= 0) and torch.all(out < 64), (
48
- f"Invalid atom-name encoding for '{atom_name}'."
49
- )
50
- return out
51
-
52
-
53
- def _default_atom_offset(atom_name: str, atom_idx: int) -> np.ndarray:
54
- if atom_name == "N":
55
- return np.array([-1.20, 0.80, 0.00], dtype=np.float32)
56
- if atom_name == "CA":
57
- return np.array([0.00, 0.00, 0.00], dtype=np.float32)
58
- if atom_name == "C":
59
- return np.array([1.35, 0.65, 0.00], dtype=np.float32)
60
- if atom_name == "O":
61
- return np.array([2.25, 0.10, 0.00], dtype=np.float32)
62
- if atom_name == "CB":
63
- return np.array([0.15, -1.20, 1.00], dtype=np.float32)
64
- angle = (atom_idx + 1) * 0.7
65
- radius = 1.4 + 0.03 * atom_idx
66
- return np.array(
67
- [
68
- radius * math.cos(angle),
69
- radius * math.sin(angle),
70
- 0.1 * ((atom_idx % 5) - 2),
71
- ],
72
- dtype=np.float32,
73
- )
74
-
75
-
76
- def _build_template(
77
- sequence: str,
78
- ) -> Tuple[
79
- ProteinStructureTemplate,
80
- List[str],
81
- List[int],
82
- List[int],
83
- List[int],
84
- List[np.ndarray],
85
- List[int],
86
- ]:
87
- residue_names: List[str] = []
88
- residue_token_ids: List[int] = []
89
- atom_names: List[str] = []
90
- atom_elements: List[str] = []
91
- atom_residue_index: List[int] = []
92
- atom_chain_id: List[str] = []
93
- atom_positions: List[np.ndarray] = []
94
- residue_center_atom_idx: List[int] = []
95
- residue_disto_atom_idx: List[int] = []
96
- residue_frame_atom_idx: List[int] = []
97
-
98
- global_atom_idx = 0
99
- for res_idx, aa in enumerate(sequence):
100
- token_name = const.prot_letter_to_token[aa]
101
- residue_names.append(token_name)
102
- residue_token_ids.append(const.token_ids[token_name])
103
-
104
- residue_atoms = const.ref_atoms[token_name]
105
- assert len(residue_atoms) > 0, f"No reference atoms for residue {token_name}."
106
- center_atom_name = const.res_to_center_atom[token_name]
107
- disto_atom_name = const.res_to_disto_atom[token_name]
108
-
109
- center_idx = -1
110
- disto_idx = -1
111
- n_idx = -1
112
- ca_idx = -1
113
- c_idx = -1
114
-
115
- for local_idx, atom_name in enumerate(residue_atoms):
116
- atom_names.append(atom_name)
117
- element = _atom_name_to_element(atom_name)
118
- atom_elements.append(element)
119
- atom_residue_index.append(res_idx)
120
- atom_chain_id.append("A")
121
-
122
- atom_pos = _default_atom_offset(atom_name, local_idx)
123
- atom_positions.append(atom_pos)
124
-
125
- if atom_name == center_atom_name:
126
- center_idx = global_atom_idx
127
- if atom_name == disto_atom_name:
128
- disto_idx = global_atom_idx
129
- if atom_name == "N":
130
- n_idx = global_atom_idx
131
- if atom_name == "CA":
132
- ca_idx = global_atom_idx
133
- if atom_name == "C":
134
- c_idx = global_atom_idx
135
- global_atom_idx += 1
136
-
137
- if center_idx == -1:
138
- center_idx = global_atom_idx - len(residue_atoms)
139
- if disto_idx == -1:
140
- disto_idx = center_idx
141
- if n_idx == -1:
142
- n_idx = center_idx
143
- if ca_idx == -1:
144
- ca_idx = center_idx
145
- if c_idx == -1:
146
- c_idx = center_idx
147
-
148
- residue_center_atom_idx.append(center_idx)
149
- residue_disto_atom_idx.append(disto_idx)
150
- residue_frame_atom_idx.extend([n_idx, ca_idx, c_idx])
151
-
152
- template = ProteinStructureTemplate(
153
- sequence=sequence,
154
- residue_names=residue_names,
155
- atom_names=atom_names,
156
- atom_elements=atom_elements,
157
- atom_residue_index=atom_residue_index,
158
- atom_chain_id=atom_chain_id,
159
- )
160
-
161
- return (
162
- template,
163
- residue_names,
164
- residue_token_ids,
165
- residue_center_atom_idx,
166
- residue_disto_atom_idx,
167
- atom_positions,
168
- residue_frame_atom_idx,
169
- )
170
-
171
-
172
- def _center_atoms_per_residue(
173
- atom_positions: torch.Tensor,
174
- atom_residue_index: List[int],
175
- num_residues: int,
176
- ) -> torch.Tensor:
177
- centered = atom_positions.clone()
178
- residue_index_tensor = torch.tensor(atom_residue_index, dtype=torch.long)
179
- for residue_idx in range(num_residues):
180
- residue_mask = residue_index_tensor == residue_idx
181
- assert torch.any(residue_mask), f"Residue index {residue_idx} has no atoms."
182
- residue_coords = centered[residue_mask]
183
- residue_center = residue_coords.mean(dim=0, keepdim=True)
184
- centered[residue_mask] = residue_coords - residue_center
185
- return centered
186
-
187
-
188
- def build_boltz2_features(
189
- amino_acid_sequence: str,
190
- num_bins: int = 64,
191
- atoms_per_window_queries: int = 32,
192
- ) -> Tuple[Dict[str, torch.Tensor], ProteinStructureTemplate]:
193
- sequence = _normalize_sequence(amino_acid_sequence)
194
- (
195
- template,
196
- residue_names,
197
- residue_token_ids,
198
- residue_center_atom_idx,
199
- residue_disto_atom_idx,
200
- atom_positions_np,
201
- residue_frame_atom_idx_flat,
202
- ) = _build_template(sequence)
203
-
204
- num_tokens = len(residue_names)
205
- num_atoms = len(atom_positions_np)
206
- assert num_tokens > 0 and num_atoms > 0
207
-
208
- atom_positions = torch.tensor(np.asarray(atom_positions_np), dtype=torch.float32)
209
- atom_positions = _center_atoms_per_residue(
210
- atom_positions=atom_positions,
211
- atom_residue_index=template.atom_residue_index,
212
- num_residues=num_tokens,
213
- )
214
-
215
- token_index = torch.arange(num_tokens, dtype=torch.long).unsqueeze(0)
216
- residue_index = torch.arange(num_tokens, dtype=torch.long).unsqueeze(0)
217
- asym_id = torch.zeros((1, num_tokens), dtype=torch.long)
218
- entity_id = torch.zeros((1, num_tokens), dtype=torch.long)
219
- sym_id = torch.zeros((1, num_tokens), dtype=torch.long)
220
- mol_type = torch.full(
221
- (1, num_tokens),
222
- fill_value=const.chain_type_ids["PROTEIN"],
223
- dtype=torch.long,
224
- )
225
-
226
- res_type_ids = torch.tensor(residue_token_ids, dtype=torch.long)
227
- res_type = one_hot(res_type_ids, num_classes=const.num_tokens).float().unsqueeze(0)
228
-
229
- token_bonds = torch.zeros((num_tokens, num_tokens), dtype=torch.float32)
230
- type_bonds = torch.zeros((num_tokens, num_tokens), dtype=torch.long)
231
- assert "SINGLE" in const.bond_type_ids
232
- assert "COVALENT" in const.bond_type_ids
233
- single_bond_id = const.bond_type_ids["SINGLE"] + 1
234
- covalent_bond_id = const.bond_type_ids["COVALENT"] + 1
235
- for idx in range(num_tokens):
236
- token_bonds[idx, idx] = 1.0
237
- type_bonds[idx, idx] = single_bond_id
238
- for idx in range(num_tokens - 1):
239
- token_bonds[idx, idx + 1] = 1.0
240
- token_bonds[idx + 1, idx] = 1.0
241
- type_bonds[idx, idx + 1] = covalent_bond_id
242
- type_bonds[idx + 1, idx] = covalent_bond_id
243
- token_bonds = token_bonds.unsqueeze(0).unsqueeze(-1)
244
- type_bonds = type_bonds.unsqueeze(0)
245
-
246
- token_pad_mask = torch.ones((1, num_tokens), dtype=torch.float32)
247
- token_resolved_mask = torch.ones((1, num_tokens), dtype=torch.float32)
248
- token_disto_mask = torch.ones((1, num_tokens), dtype=torch.float32)
249
-
250
- num_contact_classes = len(const.contact_conditioning_info)
251
- unspecified_id = const.contact_conditioning_info["UNSPECIFIED"]
252
- contact_ids = torch.full(
253
- (num_tokens, num_tokens),
254
- fill_value=unspecified_id,
255
- dtype=torch.long,
256
- )
257
- contact_conditioning = one_hot(
258
- contact_ids,
259
- num_classes=num_contact_classes,
260
- ).float().unsqueeze(0)
261
- contact_threshold = torch.zeros((1, num_tokens, num_tokens), dtype=torch.float32)
262
-
263
- assert "x-ray diffraction" in const.method_types_ids
264
- method_feature = torch.full(
265
- (1, num_tokens),
266
- fill_value=const.method_types_ids["x-ray diffraction"],
267
- dtype=torch.long,
268
- )
269
- modified = torch.zeros((1, num_tokens), dtype=torch.long)
270
- cyclic_period = torch.zeros((1, num_tokens), dtype=torch.float32)
271
- affinity_token_mask = torch.zeros((1, num_tokens), dtype=torch.float32)
272
-
273
- ref_pos = atom_positions.unsqueeze(0)
274
- atom_pad_mask = torch.ones((1, num_atoms), dtype=torch.float32)
275
- atom_resolved_mask = torch.ones((1, num_atoms), dtype=torch.float32)
276
-
277
- atom_name_codes = torch.stack(
278
- [_atom_name_to_codes(atom_name) for atom_name in template.atom_names],
279
- dim=0,
280
- )
281
- ref_atom_name_chars = one_hot(atom_name_codes, num_classes=64).float().unsqueeze(0)
282
-
283
- atomic_numbers = []
284
- for element in template.atom_elements:
285
- if element in _ELEMENT_TO_Z:
286
- z_value = _ELEMENT_TO_Z[element]
287
- else:
288
- z_value = _ELEMENT_TO_Z["C"]
289
- assert z_value < const.num_elements
290
- atomic_numbers.append(z_value)
291
- ref_element = one_hot(
292
- torch.tensor(atomic_numbers, dtype=torch.long),
293
- num_classes=const.num_elements,
294
- ).float().unsqueeze(0)
295
-
296
- ref_charge = torch.zeros((1, num_atoms), dtype=torch.float32)
297
- ref_chirality = torch.zeros((1, num_atoms), dtype=torch.long)
298
- ref_space_uid = torch.tensor(template.atom_residue_index, dtype=torch.long).unsqueeze(0)
299
-
300
- atom_to_token = one_hot(
301
- torch.tensor(template.atom_residue_index, dtype=torch.long),
302
- num_classes=num_tokens,
303
- ).float().unsqueeze(0)
304
- token_to_rep_atom = one_hot(
305
- torch.tensor(residue_disto_atom_idx, dtype=torch.long),
306
- num_classes=num_atoms,
307
- ).float().unsqueeze(0)
308
- token_to_center_atom = one_hot(
309
- torch.tensor(residue_center_atom_idx, dtype=torch.long),
310
- num_classes=num_atoms,
311
- ).float().unsqueeze(0)
312
- r_set_to_rep_atom = token_to_center_atom.clone()
313
-
314
- num_backbone_classes = (
315
- 1
316
- + len(const.protein_backbone_atom_index)
317
- + len(const.nucleic_backbone_atom_index)
318
- )
319
- backbone_ids = []
320
- for atom_name in template.atom_names:
321
- if atom_name in const.protein_backbone_atom_index:
322
- backbone_ids.append(const.protein_backbone_atom_index[atom_name] + 1)
323
- else:
324
- backbone_ids.append(0)
325
- atom_backbone_feat = one_hot(
326
- torch.tensor(backbone_ids, dtype=torch.long),
327
- num_classes=num_backbone_classes,
328
- ).float().unsqueeze(0)
329
-
330
- coords = ref_pos.unsqueeze(1).contiguous()
331
- disto_coords = torch.stack(
332
- [atom_positions[idx] for idx in residue_disto_atom_idx],
333
- dim=0,
334
- )
335
- disto_coords_ensemble = disto_coords.unsqueeze(0).unsqueeze(0).contiguous()
336
-
337
- bfactor = torch.zeros((1, num_atoms), dtype=torch.float32)
338
- atom_plddt = torch.ones((1, num_atoms), dtype=torch.float32)
339
-
340
- assert atoms_per_window_queries > 0
341
- pad_atoms = (
342
- ((num_atoms - 1) // atoms_per_window_queries + 1) * atoms_per_window_queries
343
- - num_atoms
344
- )
345
- if pad_atoms > 0:
346
- ref_pos = torch.nn.functional.pad(ref_pos, (0, 0, 0, pad_atoms), value=0.0)
347
- atom_pad_mask = torch.nn.functional.pad(atom_pad_mask, (0, pad_atoms), value=0.0)
348
- atom_resolved_mask = torch.nn.functional.pad(
349
- atom_resolved_mask,
350
- (0, pad_atoms),
351
- value=0.0,
352
- )
353
- ref_atom_name_chars = torch.nn.functional.pad(
354
- ref_atom_name_chars,
355
- (0, 0, 0, 0, 0, pad_atoms),
356
- value=0.0,
357
- )
358
- ref_element = torch.nn.functional.pad(ref_element, (0, 0, 0, pad_atoms), value=0.0)
359
- ref_charge = torch.nn.functional.pad(ref_charge, (0, pad_atoms), value=0.0)
360
- ref_chirality = torch.nn.functional.pad(ref_chirality, (0, pad_atoms), value=0)
361
- atom_backbone_feat = torch.nn.functional.pad(
362
- atom_backbone_feat,
363
- (0, 0, 0, pad_atoms),
364
- value=0.0,
365
- )
366
- ref_space_uid = torch.nn.functional.pad(ref_space_uid, (0, pad_atoms), value=0)
367
- coords = torch.nn.functional.pad(coords, (0, 0, 0, pad_atoms), value=0.0)
368
- atom_to_token = torch.nn.functional.pad(atom_to_token, (0, 0, 0, pad_atoms), value=0.0)
369
- token_to_rep_atom = torch.nn.functional.pad(
370
- token_to_rep_atom,
371
- (0, pad_atoms),
372
- value=0.0,
373
- )
374
- token_to_center_atom = torch.nn.functional.pad(
375
- token_to_center_atom,
376
- (0, pad_atoms),
377
- value=0.0,
378
- )
379
- r_set_to_rep_atom = torch.nn.functional.pad(
380
- r_set_to_rep_atom,
381
- (0, pad_atoms),
382
- value=0.0,
383
- )
384
- bfactor = torch.nn.functional.pad(bfactor, (0, pad_atoms), value=0.0)
385
- atom_plddt = torch.nn.functional.pad(atom_plddt, (0, pad_atoms), value=0.0)
386
-
387
- frames_idx = torch.tensor(
388
- residue_frame_atom_idx_flat,
389
- dtype=torch.long,
390
- ).reshape(num_tokens, 3)
391
- frames_idx = frames_idx.unsqueeze(0).unsqueeze(1)
392
- frame_resolved_mask = torch.ones((1, 1, num_tokens), dtype=torch.float32)
393
-
394
- msa = torch.tensor(residue_token_ids, dtype=torch.long).unsqueeze(0).unsqueeze(0)
395
- msa_paired = torch.ones((1, 1, num_tokens), dtype=torch.float32)
396
- deletion_value = torch.zeros((1, 1, num_tokens), dtype=torch.float32)
397
- has_deletion = torch.zeros((1, 1, num_tokens), dtype=torch.float32)
398
- msa_mask = torch.ones((1, 1, num_tokens), dtype=torch.float32)
399
- deletion_mean = torch.zeros((1, num_tokens), dtype=torch.float32)
400
- profile = one_hot(
401
- torch.tensor(residue_token_ids, dtype=torch.long),
402
- num_classes=const.num_tokens,
403
- ).float().unsqueeze(0)
404
-
405
- template_restype = one_hot(
406
- torch.zeros((1, 1, num_tokens), dtype=torch.long),
407
- num_classes=const.num_tokens,
408
- ).float()
409
- template_frame_rot = torch.zeros((1, 1, num_tokens, 3, 3), dtype=torch.float32)
410
- template_frame_t = torch.zeros((1, 1, num_tokens, 3), dtype=torch.float32)
411
- template_cb = torch.zeros((1, 1, num_tokens, 3), dtype=torch.float32)
412
- template_ca = torch.zeros((1, 1, num_tokens, 3), dtype=torch.float32)
413
- template_mask_cb = torch.zeros((1, 1, num_tokens), dtype=torch.float32)
414
- template_mask_frame = torch.zeros((1, 1, num_tokens), dtype=torch.float32)
415
- template_mask = torch.zeros((1, 1, num_tokens), dtype=torch.float32)
416
- query_to_template = torch.zeros((1, 1, num_tokens), dtype=torch.long)
417
- visibility_ids = torch.zeros((1, 1, num_tokens), dtype=torch.float32)
418
-
419
- disto_target = torch.zeros(
420
- (1, num_tokens, num_tokens, 1, num_bins),
421
- dtype=torch.float32,
422
- )
423
- disto_center = torch.stack(
424
- [atom_positions[idx] for idx in residue_disto_atom_idx],
425
- dim=0,
426
- ).unsqueeze(0)
427
-
428
- features: Dict[str, torch.Tensor] = {
429
- "token_index": token_index,
430
- "residue_index": residue_index,
431
- "asym_id": asym_id,
432
- "entity_id": entity_id,
433
- "sym_id": sym_id,
434
- "mol_type": mol_type,
435
- "res_type": res_type,
436
- "disto_center": disto_center,
437
- "token_bonds": token_bonds,
438
- "type_bonds": type_bonds,
439
- "token_pad_mask": token_pad_mask,
440
- "token_resolved_mask": token_resolved_mask,
441
- "token_disto_mask": token_disto_mask,
442
- "contact_conditioning": contact_conditioning,
443
- "contact_threshold": contact_threshold,
444
- "method_feature": method_feature,
445
- "modified": modified,
446
- "cyclic_period": cyclic_period,
447
- "affinity_token_mask": affinity_token_mask,
448
- "ref_pos": ref_pos,
449
- "atom_resolved_mask": atom_resolved_mask,
450
- "ref_atom_name_chars": ref_atom_name_chars,
451
- "ref_element": ref_element,
452
- "ref_charge": ref_charge,
453
- "ref_chirality": ref_chirality,
454
- "atom_backbone_feat": atom_backbone_feat,
455
- "ref_space_uid": ref_space_uid,
456
- "coords": coords,
457
- "atom_pad_mask": atom_pad_mask,
458
- "atom_to_token": atom_to_token,
459
- "token_to_rep_atom": token_to_rep_atom,
460
- "r_set_to_rep_atom": r_set_to_rep_atom,
461
- "token_to_center_atom": token_to_center_atom,
462
- "disto_target": disto_target,
463
- "disto_coords_ensemble": disto_coords_ensemble,
464
- "bfactor": bfactor,
465
- "plddt": atom_plddt,
466
- "frames_idx": frames_idx,
467
- "frame_resolved_mask": frame_resolved_mask,
468
- "msa": msa,
469
- "msa_paired": msa_paired,
470
- "deletion_value": deletion_value,
471
- "has_deletion": has_deletion,
472
- "deletion_mean": deletion_mean,
473
- "profile": profile,
474
- "msa_mask": msa_mask,
475
- "template_restype": template_restype,
476
- "template_frame_rot": template_frame_rot,
477
- "template_frame_t": template_frame_t,
478
- "template_cb": template_cb,
479
- "template_ca": template_ca,
480
- "template_mask_cb": template_mask_cb,
481
- "template_mask_frame": template_mask_frame,
482
- "template_mask": template_mask,
483
- "query_to_template": query_to_template,
484
- "visibility_ids": visibility_ids,
485
- }
486
-
487
- return features, template
 
1
+ import math
2
+ from typing import Dict, List, Tuple
3
+
4
+ import numpy as np
5
+ import torch
6
+ from torch.nn.functional import one_hot
7
+
8
+ from .minimal_structures import ProteinStructureTemplate
9
+ from . import vb_const as const
10
+
11
+
12
+ _ELEMENT_TO_Z = {
13
+ "H": 1,
14
+ "C": 6,
15
+ "N": 7,
16
+ "O": 8,
17
+ "P": 15,
18
+ "S": 16,
19
+ }
20
+
21
+
22
+ def _normalize_sequence(sequence: str) -> str:
23
+ seq = sequence.strip().upper()
24
+ assert len(seq) > 0, "Amino acid sequence must be non-empty."
25
+ for aa in seq:
26
+ assert aa in const.prot_letter_to_token, f"Unsupported residue code '{aa}'."
27
+ return seq
28
+
29
+
30
+ def _atom_name_to_element(atom_name: str) -> str:
31
+ name = atom_name.strip().upper()
32
+ if len(name) == 0:
33
+ return "C"
34
+ if name[0].isdigit():
35
+ name = name[1:]
36
+ if len(name) >= 2 and name[0:2] in ("CL", "BR", "FE", "MG", "ZN", "NA", "CA"):
37
+ return name[0]
38
+ return name[0]
39
+
40
+
41
+ def _atom_name_to_codes(atom_name: str) -> torch.Tensor:
42
+ clipped = atom_name.strip()[:4]
43
+ vals = [ord(ch) - 32 for ch in clipped]
44
+ while len(vals) < 4:
45
+ vals.append(0)
46
+ out = torch.tensor(vals, dtype=torch.long)
47
+ assert torch.all(out >= 0) and torch.all(out < 64), (
48
+ f"Invalid atom-name encoding for '{atom_name}'."
49
+ )
50
+ return out
51
+
52
+
53
+ def _default_atom_offset(atom_name: str, atom_idx: int) -> np.ndarray:
54
+ if atom_name == "N":
55
+ return np.array([-1.20, 0.80, 0.00], dtype=np.float32)
56
+ if atom_name == "CA":
57
+ return np.array([0.00, 0.00, 0.00], dtype=np.float32)
58
+ if atom_name == "C":
59
+ return np.array([1.35, 0.65, 0.00], dtype=np.float32)
60
+ if atom_name == "O":
61
+ return np.array([2.25, 0.10, 0.00], dtype=np.float32)
62
+ if atom_name == "CB":
63
+ return np.array([0.15, -1.20, 1.00], dtype=np.float32)
64
+ angle = (atom_idx + 1) * 0.7
65
+ radius = 1.4 + 0.03 * atom_idx
66
+ return np.array(
67
+ [
68
+ radius * math.cos(angle),
69
+ radius * math.sin(angle),
70
+ 0.1 * ((atom_idx % 5) - 2),
71
+ ],
72
+ dtype=np.float32,
73
+ )
74
+
75
+
76
+ def _build_template(
77
+ sequence: str,
78
+ ) -> Tuple[
79
+ ProteinStructureTemplate,
80
+ List[str],
81
+ List[int],
82
+ List[int],
83
+ List[int],
84
+ List[np.ndarray],
85
+ List[int],
86
+ ]:
87
+ residue_names: List[str] = []
88
+ residue_token_ids: List[int] = []
89
+ atom_names: List[str] = []
90
+ atom_elements: List[str] = []
91
+ atom_residue_index: List[int] = []
92
+ atom_chain_id: List[str] = []
93
+ atom_positions: List[np.ndarray] = []
94
+ residue_center_atom_idx: List[int] = []
95
+ residue_disto_atom_idx: List[int] = []
96
+ residue_frame_atom_idx: List[int] = []
97
+
98
+ global_atom_idx = 0
99
+ for res_idx, aa in enumerate(sequence):
100
+ token_name = const.prot_letter_to_token[aa]
101
+ residue_names.append(token_name)
102
+ residue_token_ids.append(const.token_ids[token_name])
103
+
104
+ residue_atoms = const.ref_atoms[token_name]
105
+ assert len(residue_atoms) > 0, f"No reference atoms for residue {token_name}."
106
+ center_atom_name = const.res_to_center_atom[token_name]
107
+ disto_atom_name = const.res_to_disto_atom[token_name]
108
+
109
+ center_idx = -1
110
+ disto_idx = -1
111
+ n_idx = -1
112
+ ca_idx = -1
113
+ c_idx = -1
114
+
115
+ for local_idx, atom_name in enumerate(residue_atoms):
116
+ atom_names.append(atom_name)
117
+ element = _atom_name_to_element(atom_name)
118
+ atom_elements.append(element)
119
+ atom_residue_index.append(res_idx)
120
+ atom_chain_id.append("A")
121
+
122
+ atom_pos = _default_atom_offset(atom_name, local_idx)
123
+ atom_positions.append(atom_pos)
124
+
125
+ if atom_name == center_atom_name:
126
+ center_idx = global_atom_idx
127
+ if atom_name == disto_atom_name:
128
+ disto_idx = global_atom_idx
129
+ if atom_name == "N":
130
+ n_idx = global_atom_idx
131
+ if atom_name == "CA":
132
+ ca_idx = global_atom_idx
133
+ if atom_name == "C":
134
+ c_idx = global_atom_idx
135
+ global_atom_idx += 1
136
+
137
+ if center_idx == -1:
138
+ center_idx = global_atom_idx - len(residue_atoms)
139
+ if disto_idx == -1:
140
+ disto_idx = center_idx
141
+ if n_idx == -1:
142
+ n_idx = center_idx
143
+ if ca_idx == -1:
144
+ ca_idx = center_idx
145
+ if c_idx == -1:
146
+ c_idx = center_idx
147
+
148
+ residue_center_atom_idx.append(center_idx)
149
+ residue_disto_atom_idx.append(disto_idx)
150
+ residue_frame_atom_idx.extend([n_idx, ca_idx, c_idx])
151
+
152
+ template = ProteinStructureTemplate(
153
+ sequence=sequence,
154
+ residue_names=residue_names,
155
+ atom_names=atom_names,
156
+ atom_elements=atom_elements,
157
+ atom_residue_index=atom_residue_index,
158
+ atom_chain_id=atom_chain_id,
159
+ )
160
+
161
+ return (
162
+ template,
163
+ residue_names,
164
+ residue_token_ids,
165
+ residue_center_atom_idx,
166
+ residue_disto_atom_idx,
167
+ atom_positions,
168
+ residue_frame_atom_idx,
169
+ )
170
+
171
+
172
+ def _center_atoms_per_residue(
173
+ atom_positions: torch.Tensor,
174
+ atom_residue_index: List[int],
175
+ num_residues: int,
176
+ ) -> torch.Tensor:
177
+ centered = atom_positions.clone()
178
+ residue_index_tensor = torch.tensor(atom_residue_index, dtype=torch.long)
179
+ for residue_idx in range(num_residues):
180
+ residue_mask = residue_index_tensor == residue_idx
181
+ assert torch.any(residue_mask), f"Residue index {residue_idx} has no atoms."
182
+ residue_coords = centered[residue_mask]
183
+ residue_center = residue_coords.mean(dim=0, keepdim=True)
184
+ centered[residue_mask] = residue_coords - residue_center
185
+ return centered
186
+
187
+
188
+ def build_boltz2_features(
189
+ amino_acid_sequence: str,
190
+ num_bins: int = 64,
191
+ atoms_per_window_queries: int = 32,
192
+ ) -> Tuple[Dict[str, torch.Tensor], ProteinStructureTemplate]:
193
+ sequence = _normalize_sequence(amino_acid_sequence)
194
+ (
195
+ template,
196
+ residue_names,
197
+ residue_token_ids,
198
+ residue_center_atom_idx,
199
+ residue_disto_atom_idx,
200
+ atom_positions_np,
201
+ residue_frame_atom_idx_flat,
202
+ ) = _build_template(sequence)
203
+
204
+ num_tokens = len(residue_names)
205
+ num_atoms = len(atom_positions_np)
206
+ assert num_tokens > 0 and num_atoms > 0
207
+
208
+ atom_positions = torch.tensor(np.asarray(atom_positions_np), dtype=torch.float32)
209
+ atom_positions = _center_atoms_per_residue(
210
+ atom_positions=atom_positions,
211
+ atom_residue_index=template.atom_residue_index,
212
+ num_residues=num_tokens,
213
+ )
214
+
215
+ token_index = torch.arange(num_tokens, dtype=torch.long).unsqueeze(0)
216
+ residue_index = torch.arange(num_tokens, dtype=torch.long).unsqueeze(0)
217
+ asym_id = torch.zeros((1, num_tokens), dtype=torch.long)
218
+ entity_id = torch.zeros((1, num_tokens), dtype=torch.long)
219
+ sym_id = torch.zeros((1, num_tokens), dtype=torch.long)
220
+ mol_type = torch.full(
221
+ (1, num_tokens),
222
+ fill_value=const.chain_type_ids["PROTEIN"],
223
+ dtype=torch.long,
224
+ )
225
+
226
+ res_type_ids = torch.tensor(residue_token_ids, dtype=torch.long)
227
+ res_type = one_hot(res_type_ids, num_classes=const.num_tokens).float().unsqueeze(0)
228
+
229
+ token_bonds = torch.zeros((num_tokens, num_tokens), dtype=torch.float32)
230
+ type_bonds = torch.zeros((num_tokens, num_tokens), dtype=torch.long)
231
+ assert "SINGLE" in const.bond_type_ids
232
+ assert "COVALENT" in const.bond_type_ids
233
+ single_bond_id = const.bond_type_ids["SINGLE"] + 1
234
+ covalent_bond_id = const.bond_type_ids["COVALENT"] + 1
235
+ for idx in range(num_tokens):
236
+ token_bonds[idx, idx] = 1.0
237
+ type_bonds[idx, idx] = single_bond_id
238
+ for idx in range(num_tokens - 1):
239
+ token_bonds[idx, idx + 1] = 1.0
240
+ token_bonds[idx + 1, idx] = 1.0
241
+ type_bonds[idx, idx + 1] = covalent_bond_id
242
+ type_bonds[idx + 1, idx] = covalent_bond_id
243
+ token_bonds = token_bonds.unsqueeze(0).unsqueeze(-1)
244
+ type_bonds = type_bonds.unsqueeze(0)
245
+
246
+ token_pad_mask = torch.ones((1, num_tokens), dtype=torch.float32)
247
+ token_resolved_mask = torch.ones((1, num_tokens), dtype=torch.float32)
248
+ token_disto_mask = torch.ones((1, num_tokens), dtype=torch.float32)
249
+
250
+ num_contact_classes = len(const.contact_conditioning_info)
251
+ unspecified_id = const.contact_conditioning_info["UNSPECIFIED"]
252
+ contact_ids = torch.full(
253
+ (num_tokens, num_tokens),
254
+ fill_value=unspecified_id,
255
+ dtype=torch.long,
256
+ )
257
+ contact_conditioning = one_hot(
258
+ contact_ids,
259
+ num_classes=num_contact_classes,
260
+ ).float().unsqueeze(0)
261
+ contact_threshold = torch.zeros((1, num_tokens, num_tokens), dtype=torch.float32)
262
+
263
+ assert "x-ray diffraction" in const.method_types_ids
264
+ method_feature = torch.full(
265
+ (1, num_tokens),
266
+ fill_value=const.method_types_ids["x-ray diffraction"],
267
+ dtype=torch.long,
268
+ )
269
+ modified = torch.zeros((1, num_tokens), dtype=torch.long)
270
+ cyclic_period = torch.zeros((1, num_tokens), dtype=torch.float32)
271
+ affinity_token_mask = torch.zeros((1, num_tokens), dtype=torch.float32)
272
+
273
+ ref_pos = atom_positions.unsqueeze(0)
274
+ atom_pad_mask = torch.ones((1, num_atoms), dtype=torch.float32)
275
+ atom_resolved_mask = torch.ones((1, num_atoms), dtype=torch.float32)
276
+
277
+ atom_name_codes = torch.stack(
278
+ [_atom_name_to_codes(atom_name) for atom_name in template.atom_names],
279
+ dim=0,
280
+ )
281
+ ref_atom_name_chars = one_hot(atom_name_codes, num_classes=64).float().unsqueeze(0)
282
+
283
+ atomic_numbers = []
284
+ for element in template.atom_elements:
285
+ if element in _ELEMENT_TO_Z:
286
+ z_value = _ELEMENT_TO_Z[element]
287
+ else:
288
+ z_value = _ELEMENT_TO_Z["C"]
289
+ assert z_value < const.num_elements
290
+ atomic_numbers.append(z_value)
291
+ ref_element = one_hot(
292
+ torch.tensor(atomic_numbers, dtype=torch.long),
293
+ num_classes=const.num_elements,
294
+ ).float().unsqueeze(0)
295
+
296
+ ref_charge = torch.zeros((1, num_atoms), dtype=torch.float32)
297
+ ref_chirality = torch.zeros((1, num_atoms), dtype=torch.long)
298
+ ref_space_uid = torch.tensor(template.atom_residue_index, dtype=torch.long).unsqueeze(0)
299
+
300
+ atom_to_token = one_hot(
301
+ torch.tensor(template.atom_residue_index, dtype=torch.long),
302
+ num_classes=num_tokens,
303
+ ).float().unsqueeze(0)
304
+ token_to_rep_atom = one_hot(
305
+ torch.tensor(residue_disto_atom_idx, dtype=torch.long),
306
+ num_classes=num_atoms,
307
+ ).float().unsqueeze(0)
308
+ token_to_center_atom = one_hot(
309
+ torch.tensor(residue_center_atom_idx, dtype=torch.long),
310
+ num_classes=num_atoms,
311
+ ).float().unsqueeze(0)
312
+ r_set_to_rep_atom = token_to_center_atom.clone()
313
+
314
+ num_backbone_classes = (
315
+ 1
316
+ + len(const.protein_backbone_atom_index)
317
+ + len(const.nucleic_backbone_atom_index)
318
+ )
319
+ backbone_ids = []
320
+ for atom_name in template.atom_names:
321
+ if atom_name in const.protein_backbone_atom_index:
322
+ backbone_ids.append(const.protein_backbone_atom_index[atom_name] + 1)
323
+ else:
324
+ backbone_ids.append(0)
325
+ atom_backbone_feat = one_hot(
326
+ torch.tensor(backbone_ids, dtype=torch.long),
327
+ num_classes=num_backbone_classes,
328
+ ).float().unsqueeze(0)
329
+
330
+ coords = ref_pos.unsqueeze(1).contiguous()
331
+ disto_coords = torch.stack(
332
+ [atom_positions[idx] for idx in residue_disto_atom_idx],
333
+ dim=0,
334
+ )
335
+ disto_coords_ensemble = disto_coords.unsqueeze(0).unsqueeze(0).contiguous()
336
+
337
+ bfactor = torch.zeros((1, num_atoms), dtype=torch.float32)
338
+ atom_plddt = torch.ones((1, num_atoms), dtype=torch.float32)
339
+
340
+ assert atoms_per_window_queries > 0
341
+ pad_atoms = (
342
+ ((num_atoms - 1) // atoms_per_window_queries + 1) * atoms_per_window_queries
343
+ - num_atoms
344
+ )
345
+ if pad_atoms > 0:
346
+ ref_pos = torch.nn.functional.pad(ref_pos, (0, 0, 0, pad_atoms), value=0.0)
347
+ atom_pad_mask = torch.nn.functional.pad(atom_pad_mask, (0, pad_atoms), value=0.0)
348
+ atom_resolved_mask = torch.nn.functional.pad(
349
+ atom_resolved_mask,
350
+ (0, pad_atoms),
351
+ value=0.0,
352
+ )
353
+ ref_atom_name_chars = torch.nn.functional.pad(
354
+ ref_atom_name_chars,
355
+ (0, 0, 0, 0, 0, pad_atoms),
356
+ value=0.0,
357
+ )
358
+ ref_element = torch.nn.functional.pad(ref_element, (0, 0, 0, pad_atoms), value=0.0)
359
+ ref_charge = torch.nn.functional.pad(ref_charge, (0, pad_atoms), value=0.0)
360
+ ref_chirality = torch.nn.functional.pad(ref_chirality, (0, pad_atoms), value=0)
361
+ atom_backbone_feat = torch.nn.functional.pad(
362
+ atom_backbone_feat,
363
+ (0, 0, 0, pad_atoms),
364
+ value=0.0,
365
+ )
366
+ ref_space_uid = torch.nn.functional.pad(ref_space_uid, (0, pad_atoms), value=0)
367
+ coords = torch.nn.functional.pad(coords, (0, 0, 0, pad_atoms), value=0.0)
368
+ atom_to_token = torch.nn.functional.pad(atom_to_token, (0, 0, 0, pad_atoms), value=0.0)
369
+ token_to_rep_atom = torch.nn.functional.pad(
370
+ token_to_rep_atom,
371
+ (0, pad_atoms),
372
+ value=0.0,
373
+ )
374
+ token_to_center_atom = torch.nn.functional.pad(
375
+ token_to_center_atom,
376
+ (0, pad_atoms),
377
+ value=0.0,
378
+ )
379
+ r_set_to_rep_atom = torch.nn.functional.pad(
380
+ r_set_to_rep_atom,
381
+ (0, pad_atoms),
382
+ value=0.0,
383
+ )
384
+ bfactor = torch.nn.functional.pad(bfactor, (0, pad_atoms), value=0.0)
385
+ atom_plddt = torch.nn.functional.pad(atom_plddt, (0, pad_atoms), value=0.0)
386
+
387
+ frames_idx = torch.tensor(
388
+ residue_frame_atom_idx_flat,
389
+ dtype=torch.long,
390
+ ).reshape(num_tokens, 3)
391
+ frames_idx = frames_idx.unsqueeze(0).unsqueeze(1)
392
+ frame_resolved_mask = torch.ones((1, 1, num_tokens), dtype=torch.float32)
393
+
394
+ msa = torch.tensor(residue_token_ids, dtype=torch.long).unsqueeze(0).unsqueeze(0)
395
+ msa_paired = torch.ones((1, 1, num_tokens), dtype=torch.float32)
396
+ deletion_value = torch.zeros((1, 1, num_tokens), dtype=torch.float32)
397
+ has_deletion = torch.zeros((1, 1, num_tokens), dtype=torch.float32)
398
+ msa_mask = torch.ones((1, 1, num_tokens), dtype=torch.float32)
399
+ deletion_mean = torch.zeros((1, num_tokens), dtype=torch.float32)
400
+ profile = one_hot(
401
+ torch.tensor(residue_token_ids, dtype=torch.long),
402
+ num_classes=const.num_tokens,
403
+ ).float().unsqueeze(0)
404
+
405
+ template_restype = one_hot(
406
+ torch.zeros((1, 1, num_tokens), dtype=torch.long),
407
+ num_classes=const.num_tokens,
408
+ ).float()
409
+ template_frame_rot = torch.zeros((1, 1, num_tokens, 3, 3), dtype=torch.float32)
410
+ template_frame_t = torch.zeros((1, 1, num_tokens, 3), dtype=torch.float32)
411
+ template_cb = torch.zeros((1, 1, num_tokens, 3), dtype=torch.float32)
412
+ template_ca = torch.zeros((1, 1, num_tokens, 3), dtype=torch.float32)
413
+ template_mask_cb = torch.zeros((1, 1, num_tokens), dtype=torch.float32)
414
+ template_mask_frame = torch.zeros((1, 1, num_tokens), dtype=torch.float32)
415
+ template_mask = torch.zeros((1, 1, num_tokens), dtype=torch.float32)
416
+ query_to_template = torch.zeros((1, 1, num_tokens), dtype=torch.long)
417
+ visibility_ids = torch.zeros((1, 1, num_tokens), dtype=torch.float32)
418
+
419
+ disto_target = torch.zeros(
420
+ (1, num_tokens, num_tokens, 1, num_bins),
421
+ dtype=torch.float32,
422
+ )
423
+ disto_center = torch.stack(
424
+ [atom_positions[idx] for idx in residue_disto_atom_idx],
425
+ dim=0,
426
+ ).unsqueeze(0)
427
+
428
+ features: Dict[str, torch.Tensor] = {
429
+ "token_index": token_index,
430
+ "residue_index": residue_index,
431
+ "asym_id": asym_id,
432
+ "entity_id": entity_id,
433
+ "sym_id": sym_id,
434
+ "mol_type": mol_type,
435
+ "res_type": res_type,
436
+ "disto_center": disto_center,
437
+ "token_bonds": token_bonds,
438
+ "type_bonds": type_bonds,
439
+ "token_pad_mask": token_pad_mask,
440
+ "token_resolved_mask": token_resolved_mask,
441
+ "token_disto_mask": token_disto_mask,
442
+ "contact_conditioning": contact_conditioning,
443
+ "contact_threshold": contact_threshold,
444
+ "method_feature": method_feature,
445
+ "modified": modified,
446
+ "cyclic_period": cyclic_period,
447
+ "affinity_token_mask": affinity_token_mask,
448
+ "ref_pos": ref_pos,
449
+ "atom_resolved_mask": atom_resolved_mask,
450
+ "ref_atom_name_chars": ref_atom_name_chars,
451
+ "ref_element": ref_element,
452
+ "ref_charge": ref_charge,
453
+ "ref_chirality": ref_chirality,
454
+ "atom_backbone_feat": atom_backbone_feat,
455
+ "ref_space_uid": ref_space_uid,
456
+ "coords": coords,
457
+ "atom_pad_mask": atom_pad_mask,
458
+ "atom_to_token": atom_to_token,
459
+ "token_to_rep_atom": token_to_rep_atom,
460
+ "r_set_to_rep_atom": r_set_to_rep_atom,
461
+ "token_to_center_atom": token_to_center_atom,
462
+ "disto_target": disto_target,
463
+ "disto_coords_ensemble": disto_coords_ensemble,
464
+ "bfactor": bfactor,
465
+ "plddt": atom_plddt,
466
+ "frames_idx": frames_idx,
467
+ "frame_resolved_mask": frame_resolved_mask,
468
+ "msa": msa,
469
+ "msa_paired": msa_paired,
470
+ "deletion_value": deletion_value,
471
+ "has_deletion": has_deletion,
472
+ "deletion_mean": deletion_mean,
473
+ "profile": profile,
474
+ "msa_mask": msa_mask,
475
+ "template_restype": template_restype,
476
+ "template_frame_rot": template_frame_rot,
477
+ "template_frame_t": template_frame_t,
478
+ "template_cb": template_cb,
479
+ "template_ca": template_ca,
480
+ "template_mask_cb": template_mask_cb,
481
+ "template_mask_frame": template_mask_frame,
482
+ "template_mask": template_mask,
483
+ "query_to_template": query_to_template,
484
+ "visibility_ids": visibility_ids,
485
+ }
486
+
487
+ return features, template
minimal_structures.py CHANGED
@@ -1,20 +1,20 @@
1
- from dataclasses import dataclass
2
- from typing import List
3
-
4
-
5
- @dataclass
6
- class ProteinStructureTemplate:
7
- sequence: str
8
- residue_names: List[str]
9
- atom_names: List[str]
10
- atom_elements: List[str]
11
- atom_residue_index: List[int]
12
- atom_chain_id: List[str]
13
-
14
- @property
15
- def num_atoms(self) -> int:
16
- return len(self.atom_names)
17
-
18
- @property
19
- def num_residues(self) -> int:
20
- return len(self.residue_names)
 
1
+ from dataclasses import dataclass
2
+ from typing import List
3
+
4
+
5
+ @dataclass
6
+ class ProteinStructureTemplate:
7
+ sequence: str
8
+ residue_names: List[str]
9
+ atom_names: List[str]
10
+ atom_elements: List[str]
11
+ atom_residue_index: List[int]
12
+ atom_chain_id: List[str]
13
+
14
+ @property
15
+ def num_atoms(self) -> int:
16
+ return len(self.atom_names)
17
+
18
+ @property
19
+ def num_residues(self) -> int:
20
+ return len(self.residue_names)
modeling_boltz2.py CHANGED
@@ -1,3 +1,4 @@
 
1
  import copy
2
  import inspect
3
  from collections.abc import Mapping, Sequence
 
1
+ import entrypoint_setup
2
  import copy
3
  import inspect
4
  from collections.abc import Mapping, Sequence
pytorch_model.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:5f090917cc00280b57246910eae790940cdeea01f68b0f2b793d81a844021505
3
+ size 2084301186
vb_const.py CHANGED
@@ -1,1184 +1,1184 @@
1
- ####################################################################################################
2
- # CHAINS
3
- ####################################################################################################
4
-
5
- chain_types = [
6
- "PROTEIN",
7
- "DNA",
8
- "RNA",
9
- "NONPOLYMER",
10
- ]
11
- chain_type_ids = {chain: i for i, chain in enumerate(chain_types)}
12
-
13
- out_types = [
14
- "dna_protein",
15
- "rna_protein",
16
- "ligand_protein",
17
- "dna_ligand",
18
- "rna_ligand",
19
- "intra_ligand",
20
- "intra_dna",
21
- "intra_rna",
22
- "intra_protein",
23
- "protein_protein",
24
- "modified",
25
- ]
26
-
27
- out_types_weights_af3 = {
28
- "dna_protein": 10.0,
29
- "rna_protein": 10.0,
30
- "ligand_protein": 10.0,
31
- "dna_ligand": 5.0,
32
- "rna_ligand": 5.0,
33
- "intra_ligand": 20.0,
34
- "intra_dna": 4.0,
35
- "intra_rna": 16.0,
36
- "intra_protein": 20.0,
37
- "protein_protein": 20.0,
38
- "modified": 0.0,
39
- }
40
-
41
- out_types_weights = {
42
- "dna_protein": 5.0,
43
- "rna_protein": 5.0,
44
- "ligand_protein": 20.0,
45
- "dna_ligand": 2.0,
46
- "rna_ligand": 2.0,
47
- "intra_ligand": 20.0,
48
- "intra_dna": 2.0,
49
- "intra_rna": 8.0,
50
- "intra_protein": 20.0,
51
- "protein_protein": 20.0,
52
- "modified": 0.0,
53
- }
54
-
55
-
56
- out_single_types = ["protein", "ligand", "dna", "rna"]
57
-
58
- clash_types = [
59
- "dna_protein",
60
- "rna_protein",
61
- "ligand_protein",
62
- "protein_protein",
63
- "dna_ligand",
64
- "rna_ligand",
65
- "ligand_ligand",
66
- "rna_dna",
67
- "dna_dna",
68
- "rna_rna",
69
- ]
70
-
71
- chain_types_to_clash_type = {
72
- frozenset(("PROTEIN", "DNA")): "dna_protein",
73
- frozenset(("PROTEIN", "RNA")): "rna_protein",
74
- frozenset(("PROTEIN", "NONPOLYMER")): "ligand_protein",
75
- frozenset(("PROTEIN",)): "protein_protein",
76
- frozenset(("NONPOLYMER", "DNA")): "dna_ligand",
77
- frozenset(("NONPOLYMER", "RNA")): "rna_ligand",
78
- frozenset(("NONPOLYMER",)): "ligand_ligand",
79
- frozenset(("DNA", "RNA")): "rna_dna",
80
- frozenset(("DNA",)): "dna_dna",
81
- frozenset(("RNA",)): "rna_rna",
82
- }
83
-
84
- chain_type_to_out_single_type = {
85
- "PROTEIN": "protein",
86
- "DNA": "dna",
87
- "RNA": "rna",
88
- "NONPOLYMER": "ligand",
89
- }
90
- ####################################################################################################
91
- # RESIDUES & TOKENS
92
- ####################################################################################################
93
-
94
-
95
- canonical_tokens = [
96
- "ALA",
97
- "ARG",
98
- "ASN",
99
- "ASP",
100
- "CYS",
101
- "GLN",
102
- "GLU",
103
- "GLY",
104
- "HIS",
105
- "ILE",
106
- "LEU",
107
- "LYS",
108
- "MET",
109
- "PHE",
110
- "PRO",
111
- "SER",
112
- "THR",
113
- "TRP",
114
- "TYR",
115
- "VAL",
116
- "UNK", # unknown protein token
117
- ]
118
-
119
- tokens = [
120
- "<pad>",
121
- "-",
122
- *canonical_tokens,
123
- "A",
124
- "G",
125
- "C",
126
- "U",
127
- "N", # unknown rna token
128
- "DA",
129
- "DG",
130
- "DC",
131
- "DT",
132
- "DN", # unknown dna token
133
- ]
134
-
135
- token_ids = {token: i for i, token in enumerate(tokens)}
136
- num_tokens = len(tokens)
137
- unk_token = {"PROTEIN": "UNK", "DNA": "DN", "RNA": "N"}
138
- unk_token_ids = {m: token_ids[t] for m, t in unk_token.items()}
139
-
140
- prot_letter_to_token = {
141
- "A": "ALA",
142
- "R": "ARG",
143
- "N": "ASN",
144
- "D": "ASP",
145
- "C": "CYS",
146
- "E": "GLU",
147
- "Q": "GLN",
148
- "G": "GLY",
149
- "H": "HIS",
150
- "I": "ILE",
151
- "L": "LEU",
152
- "K": "LYS",
153
- "M": "MET",
154
- "F": "PHE",
155
- "P": "PRO",
156
- "S": "SER",
157
- "T": "THR",
158
- "W": "TRP",
159
- "Y": "TYR",
160
- "V": "VAL",
161
- "X": "UNK",
162
- "J": "UNK",
163
- "B": "UNK",
164
- "Z": "UNK",
165
- "O": "UNK",
166
- "U": "UNK",
167
- "-": "-",
168
- }
169
-
170
- prot_token_to_letter = {v: k for k, v in prot_letter_to_token.items()}
171
- prot_token_to_letter["UNK"] = "X"
172
-
173
- rna_letter_to_token = {
174
- "A": "A",
175
- "G": "G",
176
- "C": "C",
177
- "U": "U",
178
- "N": "N",
179
- }
180
- rna_token_to_letter = {v: k for k, v in rna_letter_to_token.items()}
181
-
182
- dna_letter_to_token = {
183
- "A": "DA",
184
- "G": "DG",
185
- "C": "DC",
186
- "T": "DT",
187
- "N": "DN",
188
- }
189
- dna_token_to_letter = {v: k for k, v in dna_letter_to_token.items()}
190
-
191
- ####################################################################################################
192
- # ATOMS
193
- ####################################################################################################
194
-
195
- num_elements = 128
196
-
197
- chirality_types = [
198
- "CHI_UNSPECIFIED",
199
- "CHI_TETRAHEDRAL_CW",
200
- "CHI_TETRAHEDRAL_CCW",
201
- "CHI_SQUAREPLANAR",
202
- "CHI_OCTAHEDRAL",
203
- "CHI_TRIGONALBIPYRAMIDAL",
204
- "CHI_OTHER",
205
- ]
206
- chirality_type_ids = {chirality: i for i, chirality in enumerate(chirality_types)}
207
- unk_chirality_type = "CHI_OTHER"
208
-
209
- hybridization_map = [
210
- "S",
211
- "SP",
212
- "SP2",
213
- "SP2D",
214
- "SP3",
215
- "SP3D",
216
- "SP3D2",
217
- "OTHER",
218
- "UNSPECIFIED",
219
- ]
220
- hybridization_type_ids = {hybrid: i for i, hybrid in enumerate(hybridization_map)}
221
- unk_hybridization_type = "UNSPECIFIED"
222
-
223
- # fmt: off
224
- ref_atoms = {
225
- "PAD": [],
226
- "UNK": ["N", "CA", "C", "O", "CB"],
227
- "-": [],
228
- "ALA": ["N", "CA", "C", "O", "CB"],
229
- "ARG": ["N", "CA", "C", "O", "CB", "CG", "CD", "NE", "CZ", "NH1", "NH2"],
230
- "ASN": ["N", "CA", "C", "O", "CB", "CG", "OD1", "ND2"],
231
- "ASP": ["N", "CA", "C", "O", "CB", "CG", "OD1", "OD2"],
232
- "CYS": ["N", "CA", "C", "O", "CB", "SG"],
233
- "GLN": ["N", "CA", "C", "O", "CB", "CG", "CD", "OE1", "NE2"],
234
- "GLU": ["N", "CA", "C", "O", "CB", "CG", "CD", "OE1", "OE2"],
235
- "GLY": ["N", "CA", "C", "O"],
236
- "HIS": ["N", "CA", "C", "O", "CB", "CG", "ND1", "CD2", "CE1", "NE2"],
237
- "ILE": ["N", "CA", "C", "O", "CB", "CG1", "CG2", "CD1"],
238
- "LEU": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2"],
239
- "LYS": ["N", "CA", "C", "O", "CB", "CG", "CD", "CE", "NZ"],
240
- "MET": ["N", "CA", "C", "O", "CB", "CG", "SD", "CE"],
241
- "PHE": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ"],
242
- "PRO": ["N", "CA", "C", "O", "CB", "CG", "CD"],
243
- "SER": ["N", "CA", "C", "O", "CB", "OG"],
244
- "THR": ["N", "CA", "C", "O", "CB", "OG1", "CG2"],
245
- "TRP": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2", "NE1", "CE2", "CE3", "CZ2", "CZ3", "CH2"], # noqa: E501
246
- "TYR": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ", "OH"],
247
- "VAL": ["N", "CA", "C", "O", "CB", "CG1", "CG2"],
248
- "A": ["P", "OP1", "OP2", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "O2'", "C1'", "N9", "C8", "N7", "C5", "C6", "N6", "N1", "C2", "N3", "C4"], # noqa: E501
249
- "G": ["P", "OP1", "OP2", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "O2'", "C1'", "N9", "C8", "N7", "C5", "C6", "O6", "N1", "C2", "N2", "N3", "C4"], # noqa: E501
250
- "C": ["P", "OP1", "OP2", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "O2'", "C1'", "N1", "C2", "O2", "N3", "C4", "N4", "C5", "C6"], # noqa: E501
251
- "U": ["P", "OP1", "OP2", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "O2'", "C1'", "N1", "C2", "O2", "N3", "C4", "O4", "C5", "C6"], # noqa: E501
252
- "N": ["P", "OP1", "OP2", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "O2'", "C1'"], # noqa: E501
253
- "DA": ["P", "OP1", "OP2", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "C1'", "N9", "C8", "N7", "C5", "C6", "N6", "N1", "C2", "N3", "C4"], # noqa: E501
254
- "DG": ["P", "OP1", "OP2", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "C1'", "N9", "C8", "N7", "C5", "C6", "O6", "N1", "C2", "N2", "N3", "C4"], # noqa: E501
255
- "DC": ["P", "OP1", "OP2", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "C1'", "N1", "C2", "O2", "N3", "C4", "N4", "C5", "C6"], # noqa: E501
256
- "DT": ["P", "OP1", "OP2", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "C1'", "N1", "C2", "O2", "N3", "C4", "O4", "C5", "C7", "C6"], # noqa: E501
257
- "DN": ["P", "OP1", "OP2", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "C1'"]
258
- }
259
-
260
- protein_backbone_atom_names = ["N", "CA", "C", "O"]
261
- nucleic_backbone_atom_names = ["P", "OP1", "OP2", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "O2'", "C1'"]
262
-
263
- protein_backbone_atom_index = {name: i for i, name in enumerate(protein_backbone_atom_names)}
264
- nucleic_backbone_atom_index = {name: i for i, name in enumerate(nucleic_backbone_atom_names)}
265
-
266
- ref_symmetries = {
267
- "PAD": [],
268
- "ALA": [],
269
- "ARG": [],
270
- "ASN": [],
271
- "ASP": [[(6, 7), (7, 6)]],
272
- "CYS": [],
273
- "GLN": [],
274
- "GLU": [[(7, 8), (8, 7)]],
275
- "GLY": [],
276
- "HIS": [],
277
- "ILE": [],
278
- "LEU": [],
279
- "LYS": [],
280
- "MET": [],
281
- "PHE": [[(6, 7), (7, 6), (8, 9), (9, 8)]],
282
- "PRO": [],
283
- "SER": [],
284
- "THR": [],
285
- "TRP": [],
286
- "TYR": [[(6, 7), (7, 6), (8, 9), (9, 8)]],
287
- "VAL": [],
288
- "A": [[(1, 2), (2, 1)]],
289
- "G": [[(1, 2), (2, 1)]],
290
- "C": [[(1, 2), (2, 1)]],
291
- "U": [[(1, 2), (2, 1)]],
292
- #"N": [[(1, 2), (2, 1)]],
293
- "DA": [[(1, 2), (2, 1)]],
294
- "DG": [[(1, 2), (2, 1)]],
295
- "DC": [[(1, 2), (2, 1)]],
296
- "DT": [[(1, 2), (2, 1)]],
297
- #"DN": [[(1, 2), (2, 1)]]
298
- }
299
-
300
-
301
- res_to_center_atom = {
302
- "UNK": "CA",
303
- "ALA": "CA",
304
- "ARG": "CA",
305
- "ASN": "CA",
306
- "ASP": "CA",
307
- "CYS": "CA",
308
- "GLN": "CA",
309
- "GLU": "CA",
310
- "GLY": "CA",
311
- "HIS": "CA",
312
- "ILE": "CA",
313
- "LEU": "CA",
314
- "LYS": "CA",
315
- "MET": "CA",
316
- "PHE": "CA",
317
- "PRO": "CA",
318
- "SER": "CA",
319
- "THR": "CA",
320
- "TRP": "CA",
321
- "TYR": "CA",
322
- "VAL": "CA",
323
- "A": "C1'",
324
- "G": "C1'",
325
- "C": "C1'",
326
- "U": "C1'",
327
- "N": "C1'",
328
- "DA": "C1'",
329
- "DG": "C1'",
330
- "DC": "C1'",
331
- "DT": "C1'",
332
- "DN": "C1'"
333
- }
334
-
335
- res_to_disto_atom = {
336
- "UNK": "CB",
337
- "ALA": "CB",
338
- "ARG": "CB",
339
- "ASN": "CB",
340
- "ASP": "CB",
341
- "CYS": "CB",
342
- "GLN": "CB",
343
- "GLU": "CB",
344
- "GLY": "CA",
345
- "HIS": "CB",
346
- "ILE": "CB",
347
- "LEU": "CB",
348
- "LYS": "CB",
349
- "MET": "CB",
350
- "PHE": "CB",
351
- "PRO": "CB",
352
- "SER": "CB",
353
- "THR": "CB",
354
- "TRP": "CB",
355
- "TYR": "CB",
356
- "VAL": "CB",
357
- "A": "C4",
358
- "G": "C4",
359
- "C": "C2",
360
- "U": "C2",
361
- "N": "C1'",
362
- "DA": "C4",
363
- "DG": "C4",
364
- "DC": "C2",
365
- "DT": "C2",
366
- "DN": "C1'"
367
- }
368
-
369
- res_to_center_atom_id = {
370
- res: ref_atoms[res].index(atom)
371
- for res, atom in res_to_center_atom.items()
372
- }
373
-
374
- res_to_disto_atom_id = {
375
- res: ref_atoms[res].index(atom)
376
- for res, atom in res_to_disto_atom.items()
377
- }
378
-
379
- # fmt: on
380
-
381
- ####################################################################################################
382
- # BONDS
383
- ####################################################################################################
384
-
385
- atom_interface_cutoff = 5.0
386
- interface_cutoff = 15.0
387
-
388
- bond_types = [
389
- "OTHER",
390
- "SINGLE",
391
- "DOUBLE",
392
- "TRIPLE",
393
- "AROMATIC",
394
- "COVALENT",
395
- ]
396
- bond_type_ids = {bond: i for i, bond in enumerate(bond_types)}
397
- unk_bond_type = "OTHER"
398
-
399
-
400
- ####################################################################################################
401
- # Contacts
402
- ####################################################################################################
403
-
404
-
405
- pocket_contact_info = {
406
- "UNSPECIFIED": 0,
407
- "UNSELECTED": 1,
408
- "POCKET": 2,
409
- "BINDER": 3,
410
- }
411
-
412
- contact_conditioning_info = {
413
- "UNSPECIFIED": 0,
414
- "UNSELECTED": 1,
415
- "POCKET>BINDER": 2,
416
- "BINDER>POCKET": 3,
417
- "CONTACT": 4,
418
- }
419
-
420
-
421
- ####################################################################################################
422
- # MSA
423
- ####################################################################################################
424
-
425
- max_msa_seqs = 16384
426
- max_paired_seqs = 8192
427
-
428
-
429
- ####################################################################################################
430
- # CHUNKING
431
- ####################################################################################################
432
-
433
- chunk_size_threshold = 384
434
-
435
- ####################################################################################################
436
- # Method conditioning
437
- ####################################################################################################
438
-
439
- # Methods
440
- method_types_ids = {
441
- "MD": 0,
442
- "X-RAY DIFFRACTION": 1,
443
- "ELECTRON MICROSCOPY": 2,
444
- "SOLUTION NMR": 3,
445
- "SOLID-STATE NMR": 4,
446
- "NEUTRON DIFFRACTION": 4,
447
- "ELECTRON CRYSTALLOGRAPHY": 4,
448
- "FIBER DIFFRACTION": 4,
449
- "POWDER DIFFRACTION": 4,
450
- "INFRARED SPECTROSCOPY": 4,
451
- "FLUORESCENCE TRANSFER": 4,
452
- "EPR": 4,
453
- "THEORETICAL MODEL": 4,
454
- "SOLUTION SCATTERING": 4,
455
- "OTHER": 4,
456
- "AFDB": 5,
457
- "BOLTZ-1": 6,
458
- "FUTURE1": 7, # Placeholder for future supervision sources
459
- "FUTURE2": 8,
460
- "FUTURE3": 9,
461
- "FUTURE4": 10,
462
- "FUTURE5": 11,
463
- }
464
- method_types_ids = {k.lower(): v for k, v in method_types_ids.items()}
465
- num_method_types = len(set(method_types_ids.values()))
466
-
467
- # Temperature
468
- temperature_bins = [(265, 280), (280, 295), (295, 310)]
469
- temperature_bins_ids = {temp: i for i, temp in enumerate(temperature_bins)}
470
- temperature_bins_ids["other"] = len(temperature_bins)
471
- num_temp_bins = len(temperature_bins_ids)
472
-
473
-
474
- # pH
475
- ph_bins = [(0, 6), (6, 8), (8, 14)]
476
- ph_bins_ids = {ph: i for i, ph in enumerate(ph_bins)}
477
- ph_bins_ids["other"] = len(ph_bins)
478
- num_ph_bins = len(ph_bins_ids)
479
-
480
- ####################################################################################################
481
- # VDW_RADII
482
- ####################################################################################################
483
-
484
- # fmt: off
485
- vdw_radii = [
486
- 1.2, 1.4, 2.2, 1.9, 1.8, 1.7, 1.6, 1.55, 1.5, 1.54,
487
- 2.4, 2.2, 2.1, 2.1, 1.95, 1.8, 1.8, 1.88, 2.8, 2.4,
488
- 2.3, 2.15, 2.05, 2.05, 2.05, 2.05, 2.0, 2.0, 2.0, 2.1,
489
- 2.1, 2.1, 2.05, 1.9, 1.9, 2.02, 2.9, 2.55, 2.4, 2.3,
490
- 2.15, 2.1, 2.05, 2.05, 2.0, 2.05, 2.1, 2.2, 2.2, 2.25,
491
- 2.2, 2.1, 2.1, 2.16, 3.0, 2.7, 2.5, 2.48, 2.47, 2.45,
492
- 2.43, 2.42, 2.4, 2.38, 2.37, 2.35, 2.33, 2.32, 2.3, 2.28,
493
- 2.27, 2.25, 2.2, 2.1, 2.05, 2.0, 2.0, 2.05, 2.1, 2.05,
494
- 2.2, 2.3, 2.3, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.4,
495
- 2.0, 2.3, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0,
496
- 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0,
497
- 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0
498
- ]
499
- # fmt: on
500
-
501
- ####################################################################################################
502
- # Excluded ligands
503
- ####################################################################################################
504
-
505
- ligand_exclusion = {
506
- "144",
507
- "15P",
508
- "1PE",
509
- "2F2",
510
- "2JC",
511
- "3HR",
512
- "3SY",
513
- "7N5",
514
- "7PE",
515
- "9JE",
516
- "AAE",
517
- "ABA",
518
- "ACE",
519
- "ACN",
520
- "ACT",
521
- "ACY",
522
- "AZI",
523
- "BAM",
524
- "BCN",
525
- "BCT",
526
- "BDN",
527
- "BEN",
528
- "BME",
529
- "BO3",
530
- "BTB",
531
- "BTC",
532
- "BU1",
533
- "C8E",
534
- "CAD",
535
- "CAQ",
536
- "CBM",
537
- "CCN",
538
- "CIT",
539
- "CL",
540
- "CLR",
541
- "CM",
542
- "CMO",
543
- "CO3",
544
- "CPT",
545
- "CXS",
546
- "D10",
547
- "DEP",
548
- "DIO",
549
- "DMS",
550
- "DN",
551
- "DOD",
552
- "DOX",
553
- "EDO",
554
- "EEE",
555
- "EGL",
556
- "EOH",
557
- "EOX",
558
- "EPE",
559
- "ETF",
560
- "FCY",
561
- "FJO",
562
- "FLC",
563
- "FMT",
564
- "FW5",
565
- "GOL",
566
- "GSH",
567
- "GTT",
568
- "GYF",
569
- "HED",
570
- "IHP",
571
- "IHS",
572
- "IMD",
573
- "IOD",
574
- "IPA",
575
- "IPH",
576
- "LDA",
577
- "MB3",
578
- "MEG",
579
- "MES",
580
- "MLA",
581
- "MLI",
582
- "MOH",
583
- "MPD",
584
- "MRD",
585
- "MSE",
586
- "MYR",
587
- "N",
588
- "NA",
589
- "NH2",
590
- "NH4",
591
- "NHE",
592
- "NO3",
593
- "O4B",
594
- "OHE",
595
- "OLA",
596
- "OLC",
597
- "OMB",
598
- "OME",
599
- "OXA",
600
- "P6G",
601
- "PE3",
602
- "PE4",
603
- "PEG",
604
- "PEO",
605
- "PEP",
606
- "PG0",
607
- "PG4",
608
- "PGE",
609
- "PGR",
610
- "PLM",
611
- "PO4",
612
- "POL",
613
- "POP",
614
- "PVO",
615
- "SAR",
616
- "SCN",
617
- "SEO",
618
- "SEP",
619
- "SIN",
620
- "SO4",
621
- "SPD",
622
- "SPM",
623
- "SR",
624
- "STE",
625
- "STO",
626
- "STU",
627
- "TAR",
628
- "TBU",
629
- "TME",
630
- "TPO",
631
- "TRS",
632
- "UNK",
633
- "UNL",
634
- "UNX",
635
- "UPL",
636
- "URE",
637
- }
638
-
639
-
640
- ####################################################################################################
641
- # TEMPLATES
642
- ####################################################################################################
643
-
644
- min_coverage_residues = 10
645
- min_coverage_fraction = 0.1
646
-
647
-
648
- ####################################################################################################
649
- # Ambiguous atoms
650
- ####################################################################################################
651
-
652
- ambiguous_atoms = {
653
- "CA": {
654
- "*": "C",
655
- "OEX": "CA",
656
- "OEC": "CA",
657
- "543": "CA",
658
- "OC6": "CA",
659
- "OC1": "CA",
660
- "OC7": "CA",
661
- "OEY": "CA",
662
- "OC4": "CA",
663
- "OC3": "CA",
664
- "ICA": "CA",
665
- "CA": "CA",
666
- "OC2": "CA",
667
- "OC5": "CA",
668
- },
669
- "CD": {"*": "C", "CD": "CD", "CD3": "CD", "CD5": "CD", "CD1": "CD"},
670
- "BR": "BR",
671
- "CL": {
672
- "*": "CL",
673
- "C8P": "C",
674
- "L3T": "C",
675
- "TLC": "C",
676
- "TZ0": "C",
677
- "471": "C",
678
- "NLK": "C",
679
- "PGM": "C",
680
- "PNE": "C",
681
- "RCY": "C",
682
- "11F": "C",
683
- "PII": "C",
684
- "C1Q": "C",
685
- "4MD": "C",
686
- "R5A": "C",
687
- "KW2": "C",
688
- "I7M": "C",
689
- "R48": "C",
690
- "FC3": "C",
691
- "55V": "C",
692
- "KPF": "C",
693
- "SPZ": "C",
694
- "0TT": "C",
695
- "R9A": "C",
696
- "5NA": "C",
697
- "C55": "C",
698
- "NIX": "C",
699
- "5PM": "C",
700
- "PP8": "C",
701
- "544": "C",
702
- "812": "C",
703
- "NPM": "C",
704
- "KU8": "C",
705
- "A1AMM": "C",
706
- "4S0": "C",
707
- "AQC": "C",
708
- "2JK": "C",
709
- "WJR": "C",
710
- "A1AAW": "C",
711
- "85E": "C",
712
- "MB0": "C",
713
- "ZAB": "C",
714
- "85K": "C",
715
- "GBP": "C",
716
- "A1H80": "C",
717
- "A1AFR": "C",
718
- "L9M": "C",
719
- "MYK": "C",
720
- "MB9": "C",
721
- "38R": "C",
722
- "EKB": "C",
723
- "NKF": "C",
724
- "UMQ": "C",
725
- "T4K": "C",
726
- "3PT": "C",
727
- "A1A7S": "C",
728
- "1Q9": "C",
729
- "11R": "C",
730
- "D2V": "C",
731
- "SM8": "C",
732
- "IFC": "C",
733
- "DB5": "C",
734
- "L2T": "C",
735
- "GNB": "C",
736
- "PP7": "C",
737
- "072": "C",
738
- "P88": "C",
739
- "DRL": "C",
740
- "C9W": "C",
741
- "NTP": "C",
742
- "4HJ": "C",
743
- "7NA": "C",
744
- "LPC": "C",
745
- "T8W": "C",
746
- "63R": "C",
747
- "570": "C",
748
- "R4A": "C",
749
- "3BG": "C",
750
- "4RB": "C",
751
- "GSO": "C",
752
- "BQ6": "C",
753
- "R4P": "C",
754
- "5CP": "C",
755
- "TTR": "C",
756
- "6UZ": "C",
757
- "SPJ": "C",
758
- "0SA": "C",
759
- "ZL1": "C",
760
- "BYG": "C",
761
- "F0E": "C",
762
- "PC0": "C",
763
- "B2Q": "C",
764
- "KV6": "C",
765
- "NTO": "C",
766
- "CLG": "C",
767
- "R7U": "C",
768
- "SMQ": "C",
769
- "GM2": "C",
770
- "Z7P": "C",
771
- "NXF": "C",
772
- "C6Q": "C",
773
- "A1G": "C",
774
- "433": "C",
775
- "L9N": "C",
776
- "7OX": "C",
777
- "A1H84": "C",
778
- "97L": "C",
779
- "HDV": "C",
780
- "LUO": "C",
781
- "R6A": "C",
782
- "1PC": "C",
783
- "4PT": "C",
784
- "SBZ": "C",
785
- "EAB": "C",
786
- "FL4": "C",
787
- "OPS": "C",
788
- "C2X": "C",
789
- "SLL": "C",
790
- "BFC": "C",
791
- "GIP": "C",
792
- "7CP": "C",
793
- "CLH": "C",
794
- "34E": "C",
795
- "5NE": "C",
796
- "PBF": "C",
797
- "ABD": "C",
798
- "ABC": "C",
799
- "LPF": "C",
800
- "TIZ": "C",
801
- "4HH": "C",
802
- "AFC": "C",
803
- "WQH": "C",
804
- "9JL": "C",
805
- "CS3": "C",
806
- "NL0": "C",
807
- "KPY": "C",
808
- "DNA": "C",
809
- "B3C": "C",
810
- "TKL": "C",
811
- "KVS": "C",
812
- "HO6": "C",
813
- "NLH": "C",
814
- "1PB": "C",
815
- "CYF": "C",
816
- "G4M": "C",
817
- "R5B": "C",
818
- "N4S": "C",
819
- "N11": "C",
820
- "C8F": "C",
821
- "PIJ": "C",
822
- "WIN": "C",
823
- "NT1": "C",
824
- "WJW": "C",
825
- "HF7": "C",
826
- "TY1": "C",
827
- "VM1": "C",
828
- },
829
- "OS": {"*": "O", "DWC": "OS", "OHX": "OS", "OS": "OS", "8WV": "OS", "OS4": "OS"},
830
- "PB": {"*": "P", "ZN9": "PB", "ZN7": "PB", "PBM": "PB", "PB": "PB", "CSB": "PB"},
831
- "CE": {"*": "C", "CE": "CE"},
832
- "FE": {"*": "FE", "TFR": "F", "PF5": "F", "IFC": "F", "F5C": "F"},
833
- "NA": {"*": "N", "CGO": "NA", "R2K": "NA", "LVQ": "NA", "NA": "NA"},
834
- "ND": {"*": "N", "ND": "ND"},
835
- "CF": {"*": "C", "CF": "CF"},
836
- "RU": "RU",
837
- "BRAF": "BR",
838
- "EU": "EU",
839
- "CLAA": "CL",
840
- "CLBQ": "CL",
841
- "CM": {"*": "C", "ZCM": "CM"},
842
- "SN": {"*": "SN", "TAP": "S", "SND": "S", "TAD": "S", "XPT": "S"},
843
- "AG": "AG",
844
- "CLN": "CL",
845
- "CLM": "CL",
846
- "CLA": {"*": "CL", "PII": "C", "TDL": "C", "D0J": "C", "GM2": "C", "PIJ": "C"},
847
- "CLB": {
848
- "*": "CL",
849
- "TD5": "C",
850
- "PII": "C",
851
- "TDL": "C",
852
- "GM2": "C",
853
- "TD7": "C",
854
- "TD6": "C",
855
- "PIJ": "C",
856
- },
857
- "CR": {
858
- "*": "C",
859
- "BW9": "CR",
860
- "CQ4": "CR",
861
- "AC9": "CR",
862
- "TIL": "CR",
863
- "J7U": "CR",
864
- "CR": "CR",
865
- },
866
- "CLAY": "CL",
867
- "CLBC": "CL",
868
- "PD": {
869
- "*": "P",
870
- "F6Q": "PD",
871
- "SVP": "PD",
872
- "SXC": "PD",
873
- "U5U": "PD",
874
- "PD": "PD",
875
- "PLL": "PD",
876
- },
877
- "CO": {
878
- "*": "C",
879
- "J1S": "CO",
880
- "OCN": "CO",
881
- "OL3": "CO",
882
- "OL4": "CO",
883
- "B12": "CO",
884
- "XCO": "CO",
885
- "UFU": "CO",
886
- "CON": "CO",
887
- "OL5": "CO",
888
- "B13": "CO",
889
- "7KI": "CO",
890
- "PL1": "CO",
891
- "OCO": "CO",
892
- "J1R": "CO",
893
- "COH": "CO",
894
- "SIR": "CO",
895
- "6KI": "CO",
896
- "NCO": "CO",
897
- "9CO": "CO",
898
- "PC3": "CO",
899
- "BWU": "CO",
900
- "B1Z": "CO",
901
- "J83": "CO",
902
- "CO": "CO",
903
- "COY": "CO",
904
- "CNC": "CO",
905
- "3CO": "CO",
906
- "OCL": "CO",
907
- "R5Q": "CO",
908
- "X5Z": "CO",
909
- "CBY": "CO",
910
- "OLS": "CO",
911
- "F0X": "CO",
912
- "I2A": "CO",
913
- "OCM": "CO",
914
- },
915
- "CU": {
916
- "*": "C",
917
- "8ZR": "CU",
918
- "K7E": "CU",
919
- "CU3": "CU",
920
- "SI9": "CU",
921
- "35N": "CU",
922
- "C2O": "CU",
923
- "SI7": "CU",
924
- "B15": "CU",
925
- "SI0": "CU",
926
- "CUP": "CU",
927
- "SQ1": "CU",
928
- "CUK": "CU",
929
- "CUL": "CU",
930
- "SI8": "CU",
931
- "IC4": "CU",
932
- "CUM": "CU",
933
- "MM2": "CU",
934
- "B30": "CU",
935
- "S32": "CU",
936
- "V79": "CU",
937
- "IMF": "CU",
938
- "CUN": "CU",
939
- "MM1": "CU",
940
- "MP1": "CU",
941
- "IME": "CU",
942
- "B17": "CU",
943
- "C2C": "CU",
944
- "1CU": "CU",
945
- "CU6": "CU",
946
- "C1O": "CU",
947
- "CU1": "CU",
948
- "B22": "CU",
949
- "CUS": "CU",
950
- "RUQ": "CU",
951
- "CUF": "CU",
952
- "CUA": "CU",
953
- "CU": "CU",
954
- "CUO": "CU",
955
- "0TE": "CU",
956
- "SI4": "CU",
957
- },
958
- "CS": {"*": "C", "CS": "CS"},
959
- "CLQ": "CL",
960
- "CLR": "CL",
961
- "CLU": "CL",
962
- "TE": "TE",
963
- "NI": {
964
- "*": "N",
965
- "USN": "NI",
966
- "NFO": "NI",
967
- "NI2": "NI",
968
- "NFS": "NI",
969
- "NFR": "NI",
970
- "82N": "NI",
971
- "R5N": "NI",
972
- "NFU": "NI",
973
- "A1ICD": "NI",
974
- "NI3": "NI",
975
- "M43": "NI",
976
- "MM5": "NI",
977
- "BF8": "NI",
978
- "TCN": "NI",
979
- "NIK": "NI",
980
- "CUV": "NI",
981
- "MM6": "NI",
982
- "J52": "NI",
983
- "NI": "NI",
984
- "SNF": "NI",
985
- "XCC": "NI",
986
- "F0L": "NI",
987
- "UWE": "NI",
988
- "NFC": "NI",
989
- "3NI": "NI",
990
- "HNI": "NI",
991
- "F43": "NI",
992
- "RQM": "NI",
993
- "NFE": "NI",
994
- "NFB": "NI",
995
- "B51": "NI",
996
- "NI1": "NI",
997
- "WCC": "NI",
998
- "NUF": "NI",
999
- },
1000
- "SB": {"*": "S", "UJI": "SB", "SB": "SB", "118": "SB", "SBO": "SB", "3CG": "SB"},
1001
- "MO": "MO",
1002
- "SEG": "SE",
1003
- "CLL": "CL",
1004
- "CLAH": "CL",
1005
- "CLC": {
1006
- "*": "CL",
1007
- "TD5": "C",
1008
- "PII": "C",
1009
- "TDL": "C",
1010
- "GM2": "C",
1011
- "TD7": "C",
1012
- "TD6": "C",
1013
- "PIJ": "C",
1014
- },
1015
- "CLD": {"*": "CL", "PII": "C", "GM2": "C", "PIJ": "C"},
1016
- "CLAD": "CL",
1017
- "CLAE": "CL",
1018
- "LA": "LA",
1019
- "RH": "RH",
1020
- "BRAC": "BR",
1021
- "BRAD": "BR",
1022
- "CLBN": "CL",
1023
- "CLAC": "CL",
1024
- "BRAB": "BR",
1025
- "BRAE": "BR",
1026
- "MG": "MG",
1027
- "IR": "IR",
1028
- "SE": {
1029
- "*": "SE",
1030
- "HII": "S",
1031
- "NT2": "S",
1032
- "R2P": "S",
1033
- "S2P": "S",
1034
- "0IU": "S",
1035
- "QMB": "S",
1036
- "81S": "S",
1037
- "0QB": "S",
1038
- "UB4": "S",
1039
- "OHS": "S",
1040
- "Q78": "S",
1041
- "0Y2": "S",
1042
- "B3M": "S",
1043
- "NT1": "S",
1044
- "81R": "S",
1045
- },
1046
- "BRAG": "BR",
1047
- "CLF": {"*": "CL", "PII": "C", "GM2": "C", "PIJ": "C"},
1048
- "CLE": {"*": "CL", "PII": "C", "GM2": "C", "PIJ": "C"},
1049
- "BRAX": "BR",
1050
- "CLK": "CL",
1051
- "ZN": "ZN",
1052
- "AS": "AS",
1053
- "AU": "AU",
1054
- "PT": "PT",
1055
- "CLAS": "CL",
1056
- "MN": "MN",
1057
- "CLBE": "CL",
1058
- "CLBF": "CL",
1059
- "CLAF": "CL",
1060
- "NA'": {"*": "N", "CGO": "NA"},
1061
- "BRAH": "BR",
1062
- "BRAI": "BR",
1063
- "BRA": "BR",
1064
- "BRB": "BR",
1065
- "BRAV": "BR",
1066
- "HG": {
1067
- "*": "HG",
1068
- "BBA": "H",
1069
- "MID": "H",
1070
- "APM": "H",
1071
- "4QQ": "H",
1072
- "0ZG": "H",
1073
- "APH": "H",
1074
- },
1075
- "AR": "AR",
1076
- "D": "H",
1077
- "CLAN": "CL",
1078
- "SI": "SI",
1079
- "CLS": "CL",
1080
- "ZR": "ZR",
1081
- "CLAR": {"*": "CL", "ZM4": "C"},
1082
- "HO": "HO",
1083
- "CLI": {"*": "CL", "GM2": "C"},
1084
- "CLH": {"*": "CL", "GM2": "C"},
1085
- "CLAP": "CL",
1086
- "CLBL": "CL",
1087
- "CLBM": "CL",
1088
- "PR": {"*": "PR", "UF0": "P", "252": "P"},
1089
- "IN": "IN",
1090
- "CLJ": "CL",
1091
- "BRU": "BR",
1092
- "SC": {"*": "S", "SFL": "SC"},
1093
- "CLG": {"*": "CL", "GM2": "C"},
1094
- "BRAT": "BR",
1095
- "BRAR": "BR",
1096
- "CLAG": "CL",
1097
- "CLAB": "CL",
1098
- "CLV": "CL",
1099
- "TI": "TI",
1100
- "CLAX": "CL",
1101
- "CLAJ": "CL",
1102
- "CL'": {"*": "CL", "BNR": "C", "25A": "C", "BDA": "C"},
1103
- "CLAW": "CL",
1104
- "BRF": "BR",
1105
- "BRE": "BR",
1106
- "RE": "RE",
1107
- "GD": "GD",
1108
- "SM": {"*": "S", "SM": "SM"},
1109
- "CLBH": "CL",
1110
- "CLBI": "CL",
1111
- "CLAI": "CL",
1112
- "CLY": "CL",
1113
- "CLZ": "CL",
1114
- "AC": "AC",
1115
- "BR'": "BR",
1116
- "CLT": "CL",
1117
- "CLO": "CL",
1118
- "CLP": "CL",
1119
- "LU": "LU",
1120
- "BA": {"*": "B", "BA": "BA"},
1121
- "CLAU": "CL",
1122
- "RB": "RB",
1123
- "LI": "LI",
1124
- "MOM": "MO",
1125
- "BRAQ": "BR",
1126
- "SR": {"*": "S", "SR": "SR", "OER": "SR"},
1127
- "CLAT": "CL",
1128
- "BRAL": "BR",
1129
- "SEB": "SE",
1130
- "CLW": "CL",
1131
- "CLX": "CL",
1132
- "BE": "BE",
1133
- "BRG": "BR",
1134
- "SEA": "SE",
1135
- "BRAW": "BR",
1136
- "BRBB": "BR",
1137
- "ER": "ER",
1138
- "TH": "TH",
1139
- "BRR": "BR",
1140
- "CLBV": "CL",
1141
- "AL": "AL",
1142
- "CLAV": "CL",
1143
- "BRH": "BR",
1144
- "CLAQ": "CL",
1145
- "GA": "GA",
1146
- "X": "*",
1147
- "TL": "TL",
1148
- "CLBB": "CL",
1149
- "TB": "TB",
1150
- "CLAK": "CL",
1151
- "XE": {"*": "*", "XE": "XE"},
1152
- "SEL": "SE",
1153
- "PU": {"*": "P", "4PU": "PU"},
1154
- "CLAZ": "CL",
1155
- "SE'": "SE",
1156
- "CLBA": "CL",
1157
- "SEN": "SE",
1158
- "SNN": "SN",
1159
- "MOB": "MO",
1160
- "YB": "YB",
1161
- "BRC": "BR",
1162
- "BRD": "BR",
1163
- "CLAM": "CL",
1164
- "DA": "H",
1165
- "DB": "H",
1166
- "DC": "H",
1167
- "DXT": "H",
1168
- "DXU": "H",
1169
- "DXX": "H",
1170
- "DXY": "H",
1171
- "DXZ": "H",
1172
- "DY": "DY",
1173
- "TA": "TA",
1174
- "XD": "*",
1175
- "SED": "SE",
1176
- "CLAL": "CL",
1177
- "BRAJ": "BR",
1178
- "AM": "AM",
1179
- "CLAO": "CL",
1180
- "BI": "BI",
1181
- "KR": "KR",
1182
- "BRBJ": "BR",
1183
- "UNK": "*",
1184
- }
 
1
+ ####################################################################################################
2
+ # CHAINS
3
+ ####################################################################################################
4
+
5
+ chain_types = [
6
+ "PROTEIN",
7
+ "DNA",
8
+ "RNA",
9
+ "NONPOLYMER",
10
+ ]
11
+ chain_type_ids = {chain: i for i, chain in enumerate(chain_types)}
12
+
13
+ out_types = [
14
+ "dna_protein",
15
+ "rna_protein",
16
+ "ligand_protein",
17
+ "dna_ligand",
18
+ "rna_ligand",
19
+ "intra_ligand",
20
+ "intra_dna",
21
+ "intra_rna",
22
+ "intra_protein",
23
+ "protein_protein",
24
+ "modified",
25
+ ]
26
+
27
+ out_types_weights_af3 = {
28
+ "dna_protein": 10.0,
29
+ "rna_protein": 10.0,
30
+ "ligand_protein": 10.0,
31
+ "dna_ligand": 5.0,
32
+ "rna_ligand": 5.0,
33
+ "intra_ligand": 20.0,
34
+ "intra_dna": 4.0,
35
+ "intra_rna": 16.0,
36
+ "intra_protein": 20.0,
37
+ "protein_protein": 20.0,
38
+ "modified": 0.0,
39
+ }
40
+
41
+ out_types_weights = {
42
+ "dna_protein": 5.0,
43
+ "rna_protein": 5.0,
44
+ "ligand_protein": 20.0,
45
+ "dna_ligand": 2.0,
46
+ "rna_ligand": 2.0,
47
+ "intra_ligand": 20.0,
48
+ "intra_dna": 2.0,
49
+ "intra_rna": 8.0,
50
+ "intra_protein": 20.0,
51
+ "protein_protein": 20.0,
52
+ "modified": 0.0,
53
+ }
54
+
55
+
56
+ out_single_types = ["protein", "ligand", "dna", "rna"]
57
+
58
+ clash_types = [
59
+ "dna_protein",
60
+ "rna_protein",
61
+ "ligand_protein",
62
+ "protein_protein",
63
+ "dna_ligand",
64
+ "rna_ligand",
65
+ "ligand_ligand",
66
+ "rna_dna",
67
+ "dna_dna",
68
+ "rna_rna",
69
+ ]
70
+
71
+ chain_types_to_clash_type = {
72
+ frozenset(("PROTEIN", "DNA")): "dna_protein",
73
+ frozenset(("PROTEIN", "RNA")): "rna_protein",
74
+ frozenset(("PROTEIN", "NONPOLYMER")): "ligand_protein",
75
+ frozenset(("PROTEIN",)): "protein_protein",
76
+ frozenset(("NONPOLYMER", "DNA")): "dna_ligand",
77
+ frozenset(("NONPOLYMER", "RNA")): "rna_ligand",
78
+ frozenset(("NONPOLYMER",)): "ligand_ligand",
79
+ frozenset(("DNA", "RNA")): "rna_dna",
80
+ frozenset(("DNA",)): "dna_dna",
81
+ frozenset(("RNA",)): "rna_rna",
82
+ }
83
+
84
+ chain_type_to_out_single_type = {
85
+ "PROTEIN": "protein",
86
+ "DNA": "dna",
87
+ "RNA": "rna",
88
+ "NONPOLYMER": "ligand",
89
+ }
90
+ ####################################################################################################
91
+ # RESIDUES & TOKENS
92
+ ####################################################################################################
93
+
94
+
95
+ canonical_tokens = [
96
+ "ALA",
97
+ "ARG",
98
+ "ASN",
99
+ "ASP",
100
+ "CYS",
101
+ "GLN",
102
+ "GLU",
103
+ "GLY",
104
+ "HIS",
105
+ "ILE",
106
+ "LEU",
107
+ "LYS",
108
+ "MET",
109
+ "PHE",
110
+ "PRO",
111
+ "SER",
112
+ "THR",
113
+ "TRP",
114
+ "TYR",
115
+ "VAL",
116
+ "UNK", # unknown protein token
117
+ ]
118
+
119
+ tokens = [
120
+ "<pad>",
121
+ "-",
122
+ *canonical_tokens,
123
+ "A",
124
+ "G",
125
+ "C",
126
+ "U",
127
+ "N", # unknown rna token
128
+ "DA",
129
+ "DG",
130
+ "DC",
131
+ "DT",
132
+ "DN", # unknown dna token
133
+ ]
134
+
135
+ token_ids = {token: i for i, token in enumerate(tokens)}
136
+ num_tokens = len(tokens)
137
+ unk_token = {"PROTEIN": "UNK", "DNA": "DN", "RNA": "N"}
138
+ unk_token_ids = {m: token_ids[t] for m, t in unk_token.items()}
139
+
140
+ prot_letter_to_token = {
141
+ "A": "ALA",
142
+ "R": "ARG",
143
+ "N": "ASN",
144
+ "D": "ASP",
145
+ "C": "CYS",
146
+ "E": "GLU",
147
+ "Q": "GLN",
148
+ "G": "GLY",
149
+ "H": "HIS",
150
+ "I": "ILE",
151
+ "L": "LEU",
152
+ "K": "LYS",
153
+ "M": "MET",
154
+ "F": "PHE",
155
+ "P": "PRO",
156
+ "S": "SER",
157
+ "T": "THR",
158
+ "W": "TRP",
159
+ "Y": "TYR",
160
+ "V": "VAL",
161
+ "X": "UNK",
162
+ "J": "UNK",
163
+ "B": "UNK",
164
+ "Z": "UNK",
165
+ "O": "UNK",
166
+ "U": "UNK",
167
+ "-": "-",
168
+ }
169
+
170
+ prot_token_to_letter = {v: k for k, v in prot_letter_to_token.items()}
171
+ prot_token_to_letter["UNK"] = "X"
172
+
173
+ rna_letter_to_token = {
174
+ "A": "A",
175
+ "G": "G",
176
+ "C": "C",
177
+ "U": "U",
178
+ "N": "N",
179
+ }
180
+ rna_token_to_letter = {v: k for k, v in rna_letter_to_token.items()}
181
+
182
+ dna_letter_to_token = {
183
+ "A": "DA",
184
+ "G": "DG",
185
+ "C": "DC",
186
+ "T": "DT",
187
+ "N": "DN",
188
+ }
189
+ dna_token_to_letter = {v: k for k, v in dna_letter_to_token.items()}
190
+
191
+ ####################################################################################################
192
+ # ATOMS
193
+ ####################################################################################################
194
+
195
+ num_elements = 128
196
+
197
+ chirality_types = [
198
+ "CHI_UNSPECIFIED",
199
+ "CHI_TETRAHEDRAL_CW",
200
+ "CHI_TETRAHEDRAL_CCW",
201
+ "CHI_SQUAREPLANAR",
202
+ "CHI_OCTAHEDRAL",
203
+ "CHI_TRIGONALBIPYRAMIDAL",
204
+ "CHI_OTHER",
205
+ ]
206
+ chirality_type_ids = {chirality: i for i, chirality in enumerate(chirality_types)}
207
+ unk_chirality_type = "CHI_OTHER"
208
+
209
+ hybridization_map = [
210
+ "S",
211
+ "SP",
212
+ "SP2",
213
+ "SP2D",
214
+ "SP3",
215
+ "SP3D",
216
+ "SP3D2",
217
+ "OTHER",
218
+ "UNSPECIFIED",
219
+ ]
220
+ hybridization_type_ids = {hybrid: i for i, hybrid in enumerate(hybridization_map)}
221
+ unk_hybridization_type = "UNSPECIFIED"
222
+
223
+ # fmt: off
224
+ ref_atoms = {
225
+ "PAD": [],
226
+ "UNK": ["N", "CA", "C", "O", "CB"],
227
+ "-": [],
228
+ "ALA": ["N", "CA", "C", "O", "CB"],
229
+ "ARG": ["N", "CA", "C", "O", "CB", "CG", "CD", "NE", "CZ", "NH1", "NH2"],
230
+ "ASN": ["N", "CA", "C", "O", "CB", "CG", "OD1", "ND2"],
231
+ "ASP": ["N", "CA", "C", "O", "CB", "CG", "OD1", "OD2"],
232
+ "CYS": ["N", "CA", "C", "O", "CB", "SG"],
233
+ "GLN": ["N", "CA", "C", "O", "CB", "CG", "CD", "OE1", "NE2"],
234
+ "GLU": ["N", "CA", "C", "O", "CB", "CG", "CD", "OE1", "OE2"],
235
+ "GLY": ["N", "CA", "C", "O"],
236
+ "HIS": ["N", "CA", "C", "O", "CB", "CG", "ND1", "CD2", "CE1", "NE2"],
237
+ "ILE": ["N", "CA", "C", "O", "CB", "CG1", "CG2", "CD1"],
238
+ "LEU": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2"],
239
+ "LYS": ["N", "CA", "C", "O", "CB", "CG", "CD", "CE", "NZ"],
240
+ "MET": ["N", "CA", "C", "O", "CB", "CG", "SD", "CE"],
241
+ "PHE": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ"],
242
+ "PRO": ["N", "CA", "C", "O", "CB", "CG", "CD"],
243
+ "SER": ["N", "CA", "C", "O", "CB", "OG"],
244
+ "THR": ["N", "CA", "C", "O", "CB", "OG1", "CG2"],
245
+ "TRP": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2", "NE1", "CE2", "CE3", "CZ2", "CZ3", "CH2"], # noqa: E501
246
+ "TYR": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ", "OH"],
247
+ "VAL": ["N", "CA", "C", "O", "CB", "CG1", "CG2"],
248
+ "A": ["P", "OP1", "OP2", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "O2'", "C1'", "N9", "C8", "N7", "C5", "C6", "N6", "N1", "C2", "N3", "C4"], # noqa: E501
249
+ "G": ["P", "OP1", "OP2", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "O2'", "C1'", "N9", "C8", "N7", "C5", "C6", "O6", "N1", "C2", "N2", "N3", "C4"], # noqa: E501
250
+ "C": ["P", "OP1", "OP2", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "O2'", "C1'", "N1", "C2", "O2", "N3", "C4", "N4", "C5", "C6"], # noqa: E501
251
+ "U": ["P", "OP1", "OP2", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "O2'", "C1'", "N1", "C2", "O2", "N3", "C4", "O4", "C5", "C6"], # noqa: E501
252
+ "N": ["P", "OP1", "OP2", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "O2'", "C1'"], # noqa: E501
253
+ "DA": ["P", "OP1", "OP2", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "C1'", "N9", "C8", "N7", "C5", "C6", "N6", "N1", "C2", "N3", "C4"], # noqa: E501
254
+ "DG": ["P", "OP1", "OP2", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "C1'", "N9", "C8", "N7", "C5", "C6", "O6", "N1", "C2", "N2", "N3", "C4"], # noqa: E501
255
+ "DC": ["P", "OP1", "OP2", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "C1'", "N1", "C2", "O2", "N3", "C4", "N4", "C5", "C6"], # noqa: E501
256
+ "DT": ["P", "OP1", "OP2", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "C1'", "N1", "C2", "O2", "N3", "C4", "O4", "C5", "C7", "C6"], # noqa: E501
257
+ "DN": ["P", "OP1", "OP2", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "C1'"]
258
+ }
259
+
260
+ protein_backbone_atom_names = ["N", "CA", "C", "O"]
261
+ nucleic_backbone_atom_names = ["P", "OP1", "OP2", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "O2'", "C1'"]
262
+
263
+ protein_backbone_atom_index = {name: i for i, name in enumerate(protein_backbone_atom_names)}
264
+ nucleic_backbone_atom_index = {name: i for i, name in enumerate(nucleic_backbone_atom_names)}
265
+
266
+ ref_symmetries = {
267
+ "PAD": [],
268
+ "ALA": [],
269
+ "ARG": [],
270
+ "ASN": [],
271
+ "ASP": [[(6, 7), (7, 6)]],
272
+ "CYS": [],
273
+ "GLN": [],
274
+ "GLU": [[(7, 8), (8, 7)]],
275
+ "GLY": [],
276
+ "HIS": [],
277
+ "ILE": [],
278
+ "LEU": [],
279
+ "LYS": [],
280
+ "MET": [],
281
+ "PHE": [[(6, 7), (7, 6), (8, 9), (9, 8)]],
282
+ "PRO": [],
283
+ "SER": [],
284
+ "THR": [],
285
+ "TRP": [],
286
+ "TYR": [[(6, 7), (7, 6), (8, 9), (9, 8)]],
287
+ "VAL": [],
288
+ "A": [[(1, 2), (2, 1)]],
289
+ "G": [[(1, 2), (2, 1)]],
290
+ "C": [[(1, 2), (2, 1)]],
291
+ "U": [[(1, 2), (2, 1)]],
292
+ #"N": [[(1, 2), (2, 1)]],
293
+ "DA": [[(1, 2), (2, 1)]],
294
+ "DG": [[(1, 2), (2, 1)]],
295
+ "DC": [[(1, 2), (2, 1)]],
296
+ "DT": [[(1, 2), (2, 1)]],
297
+ #"DN": [[(1, 2), (2, 1)]]
298
+ }
299
+
300
+
301
+ res_to_center_atom = {
302
+ "UNK": "CA",
303
+ "ALA": "CA",
304
+ "ARG": "CA",
305
+ "ASN": "CA",
306
+ "ASP": "CA",
307
+ "CYS": "CA",
308
+ "GLN": "CA",
309
+ "GLU": "CA",
310
+ "GLY": "CA",
311
+ "HIS": "CA",
312
+ "ILE": "CA",
313
+ "LEU": "CA",
314
+ "LYS": "CA",
315
+ "MET": "CA",
316
+ "PHE": "CA",
317
+ "PRO": "CA",
318
+ "SER": "CA",
319
+ "THR": "CA",
320
+ "TRP": "CA",
321
+ "TYR": "CA",
322
+ "VAL": "CA",
323
+ "A": "C1'",
324
+ "G": "C1'",
325
+ "C": "C1'",
326
+ "U": "C1'",
327
+ "N": "C1'",
328
+ "DA": "C1'",
329
+ "DG": "C1'",
330
+ "DC": "C1'",
331
+ "DT": "C1'",
332
+ "DN": "C1'"
333
+ }
334
+
335
+ res_to_disto_atom = {
336
+ "UNK": "CB",
337
+ "ALA": "CB",
338
+ "ARG": "CB",
339
+ "ASN": "CB",
340
+ "ASP": "CB",
341
+ "CYS": "CB",
342
+ "GLN": "CB",
343
+ "GLU": "CB",
344
+ "GLY": "CA",
345
+ "HIS": "CB",
346
+ "ILE": "CB",
347
+ "LEU": "CB",
348
+ "LYS": "CB",
349
+ "MET": "CB",
350
+ "PHE": "CB",
351
+ "PRO": "CB",
352
+ "SER": "CB",
353
+ "THR": "CB",
354
+ "TRP": "CB",
355
+ "TYR": "CB",
356
+ "VAL": "CB",
357
+ "A": "C4",
358
+ "G": "C4",
359
+ "C": "C2",
360
+ "U": "C2",
361
+ "N": "C1'",
362
+ "DA": "C4",
363
+ "DG": "C4",
364
+ "DC": "C2",
365
+ "DT": "C2",
366
+ "DN": "C1'"
367
+ }
368
+
369
+ res_to_center_atom_id = {
370
+ res: ref_atoms[res].index(atom)
371
+ for res, atom in res_to_center_atom.items()
372
+ }
373
+
374
+ res_to_disto_atom_id = {
375
+ res: ref_atoms[res].index(atom)
376
+ for res, atom in res_to_disto_atom.items()
377
+ }
378
+
379
+ # fmt: on
380
+
381
+ ####################################################################################################
382
+ # BONDS
383
+ ####################################################################################################
384
+
385
+ atom_interface_cutoff = 5.0
386
+ interface_cutoff = 15.0
387
+
388
+ bond_types = [
389
+ "OTHER",
390
+ "SINGLE",
391
+ "DOUBLE",
392
+ "TRIPLE",
393
+ "AROMATIC",
394
+ "COVALENT",
395
+ ]
396
+ bond_type_ids = {bond: i for i, bond in enumerate(bond_types)}
397
+ unk_bond_type = "OTHER"
398
+
399
+
400
+ ####################################################################################################
401
+ # Contacts
402
+ ####################################################################################################
403
+
404
+
405
+ pocket_contact_info = {
406
+ "UNSPECIFIED": 0,
407
+ "UNSELECTED": 1,
408
+ "POCKET": 2,
409
+ "BINDER": 3,
410
+ }
411
+
412
+ contact_conditioning_info = {
413
+ "UNSPECIFIED": 0,
414
+ "UNSELECTED": 1,
415
+ "POCKET>BINDER": 2,
416
+ "BINDER>POCKET": 3,
417
+ "CONTACT": 4,
418
+ }
419
+
420
+
421
+ ####################################################################################################
422
+ # MSA
423
+ ####################################################################################################
424
+
425
+ max_msa_seqs = 16384
426
+ max_paired_seqs = 8192
427
+
428
+
429
+ ####################################################################################################
430
+ # CHUNKING
431
+ ####################################################################################################
432
+
433
+ chunk_size_threshold = 384
434
+
435
+ ####################################################################################################
436
+ # Method conditioning
437
+ ####################################################################################################
438
+
439
+ # Methods
440
+ method_types_ids = {
441
+ "MD": 0,
442
+ "X-RAY DIFFRACTION": 1,
443
+ "ELECTRON MICROSCOPY": 2,
444
+ "SOLUTION NMR": 3,
445
+ "SOLID-STATE NMR": 4,
446
+ "NEUTRON DIFFRACTION": 4,
447
+ "ELECTRON CRYSTALLOGRAPHY": 4,
448
+ "FIBER DIFFRACTION": 4,
449
+ "POWDER DIFFRACTION": 4,
450
+ "INFRARED SPECTROSCOPY": 4,
451
+ "FLUORESCENCE TRANSFER": 4,
452
+ "EPR": 4,
453
+ "THEORETICAL MODEL": 4,
454
+ "SOLUTION SCATTERING": 4,
455
+ "OTHER": 4,
456
+ "AFDB": 5,
457
+ "BOLTZ-1": 6,
458
+ "FUTURE1": 7, # Placeholder for future supervision sources
459
+ "FUTURE2": 8,
460
+ "FUTURE3": 9,
461
+ "FUTURE4": 10,
462
+ "FUTURE5": 11,
463
+ }
464
+ method_types_ids = {k.lower(): v for k, v in method_types_ids.items()}
465
+ num_method_types = len(set(method_types_ids.values()))
466
+
467
+ # Temperature
468
+ temperature_bins = [(265, 280), (280, 295), (295, 310)]
469
+ temperature_bins_ids = {temp: i for i, temp in enumerate(temperature_bins)}
470
+ temperature_bins_ids["other"] = len(temperature_bins)
471
+ num_temp_bins = len(temperature_bins_ids)
472
+
473
+
474
+ # pH
475
+ ph_bins = [(0, 6), (6, 8), (8, 14)]
476
+ ph_bins_ids = {ph: i for i, ph in enumerate(ph_bins)}
477
+ ph_bins_ids["other"] = len(ph_bins)
478
+ num_ph_bins = len(ph_bins_ids)
479
+
480
+ ####################################################################################################
481
+ # VDW_RADII
482
+ ####################################################################################################
483
+
484
+ # fmt: off
485
+ vdw_radii = [
486
+ 1.2, 1.4, 2.2, 1.9, 1.8, 1.7, 1.6, 1.55, 1.5, 1.54,
487
+ 2.4, 2.2, 2.1, 2.1, 1.95, 1.8, 1.8, 1.88, 2.8, 2.4,
488
+ 2.3, 2.15, 2.05, 2.05, 2.05, 2.05, 2.0, 2.0, 2.0, 2.1,
489
+ 2.1, 2.1, 2.05, 1.9, 1.9, 2.02, 2.9, 2.55, 2.4, 2.3,
490
+ 2.15, 2.1, 2.05, 2.05, 2.0, 2.05, 2.1, 2.2, 2.2, 2.25,
491
+ 2.2, 2.1, 2.1, 2.16, 3.0, 2.7, 2.5, 2.48, 2.47, 2.45,
492
+ 2.43, 2.42, 2.4, 2.38, 2.37, 2.35, 2.33, 2.32, 2.3, 2.28,
493
+ 2.27, 2.25, 2.2, 2.1, 2.05, 2.0, 2.0, 2.05, 2.1, 2.05,
494
+ 2.2, 2.3, 2.3, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.4,
495
+ 2.0, 2.3, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0,
496
+ 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0,
497
+ 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0
498
+ ]
499
+ # fmt: on
500
+
501
+ ####################################################################################################
502
+ # Excluded ligands
503
+ ####################################################################################################
504
+
505
+ ligand_exclusion = {
506
+ "144",
507
+ "15P",
508
+ "1PE",
509
+ "2F2",
510
+ "2JC",
511
+ "3HR",
512
+ "3SY",
513
+ "7N5",
514
+ "7PE",
515
+ "9JE",
516
+ "AAE",
517
+ "ABA",
518
+ "ACE",
519
+ "ACN",
520
+ "ACT",
521
+ "ACY",
522
+ "AZI",
523
+ "BAM",
524
+ "BCN",
525
+ "BCT",
526
+ "BDN",
527
+ "BEN",
528
+ "BME",
529
+ "BO3",
530
+ "BTB",
531
+ "BTC",
532
+ "BU1",
533
+ "C8E",
534
+ "CAD",
535
+ "CAQ",
536
+ "CBM",
537
+ "CCN",
538
+ "CIT",
539
+ "CL",
540
+ "CLR",
541
+ "CM",
542
+ "CMO",
543
+ "CO3",
544
+ "CPT",
545
+ "CXS",
546
+ "D10",
547
+ "DEP",
548
+ "DIO",
549
+ "DMS",
550
+ "DN",
551
+ "DOD",
552
+ "DOX",
553
+ "EDO",
554
+ "EEE",
555
+ "EGL",
556
+ "EOH",
557
+ "EOX",
558
+ "EPE",
559
+ "ETF",
560
+ "FCY",
561
+ "FJO",
562
+ "FLC",
563
+ "FMT",
564
+ "FW5",
565
+ "GOL",
566
+ "GSH",
567
+ "GTT",
568
+ "GYF",
569
+ "HED",
570
+ "IHP",
571
+ "IHS",
572
+ "IMD",
573
+ "IOD",
574
+ "IPA",
575
+ "IPH",
576
+ "LDA",
577
+ "MB3",
578
+ "MEG",
579
+ "MES",
580
+ "MLA",
581
+ "MLI",
582
+ "MOH",
583
+ "MPD",
584
+ "MRD",
585
+ "MSE",
586
+ "MYR",
587
+ "N",
588
+ "NA",
589
+ "NH2",
590
+ "NH4",
591
+ "NHE",
592
+ "NO3",
593
+ "O4B",
594
+ "OHE",
595
+ "OLA",
596
+ "OLC",
597
+ "OMB",
598
+ "OME",
599
+ "OXA",
600
+ "P6G",
601
+ "PE3",
602
+ "PE4",
603
+ "PEG",
604
+ "PEO",
605
+ "PEP",
606
+ "PG0",
607
+ "PG4",
608
+ "PGE",
609
+ "PGR",
610
+ "PLM",
611
+ "PO4",
612
+ "POL",
613
+ "POP",
614
+ "PVO",
615
+ "SAR",
616
+ "SCN",
617
+ "SEO",
618
+ "SEP",
619
+ "SIN",
620
+ "SO4",
621
+ "SPD",
622
+ "SPM",
623
+ "SR",
624
+ "STE",
625
+ "STO",
626
+ "STU",
627
+ "TAR",
628
+ "TBU",
629
+ "TME",
630
+ "TPO",
631
+ "TRS",
632
+ "UNK",
633
+ "UNL",
634
+ "UNX",
635
+ "UPL",
636
+ "URE",
637
+ }
638
+
639
+
640
+ ####################################################################################################
641
+ # TEMPLATES
642
+ ####################################################################################################
643
+
644
+ min_coverage_residues = 10
645
+ min_coverage_fraction = 0.1
646
+
647
+
648
+ ####################################################################################################
649
+ # Ambiguous atoms
650
+ ####################################################################################################
651
+
652
+ ambiguous_atoms = {
653
+ "CA": {
654
+ "*": "C",
655
+ "OEX": "CA",
656
+ "OEC": "CA",
657
+ "543": "CA",
658
+ "OC6": "CA",
659
+ "OC1": "CA",
660
+ "OC7": "CA",
661
+ "OEY": "CA",
662
+ "OC4": "CA",
663
+ "OC3": "CA",
664
+ "ICA": "CA",
665
+ "CA": "CA",
666
+ "OC2": "CA",
667
+ "OC5": "CA",
668
+ },
669
+ "CD": {"*": "C", "CD": "CD", "CD3": "CD", "CD5": "CD", "CD1": "CD"},
670
+ "BR": "BR",
671
+ "CL": {
672
+ "*": "CL",
673
+ "C8P": "C",
674
+ "L3T": "C",
675
+ "TLC": "C",
676
+ "TZ0": "C",
677
+ "471": "C",
678
+ "NLK": "C",
679
+ "PGM": "C",
680
+ "PNE": "C",
681
+ "RCY": "C",
682
+ "11F": "C",
683
+ "PII": "C",
684
+ "C1Q": "C",
685
+ "4MD": "C",
686
+ "R5A": "C",
687
+ "KW2": "C",
688
+ "I7M": "C",
689
+ "R48": "C",
690
+ "FC3": "C",
691
+ "55V": "C",
692
+ "KPF": "C",
693
+ "SPZ": "C",
694
+ "0TT": "C",
695
+ "R9A": "C",
696
+ "5NA": "C",
697
+ "C55": "C",
698
+ "NIX": "C",
699
+ "5PM": "C",
700
+ "PP8": "C",
701
+ "544": "C",
702
+ "812": "C",
703
+ "NPM": "C",
704
+ "KU8": "C",
705
+ "A1AMM": "C",
706
+ "4S0": "C",
707
+ "AQC": "C",
708
+ "2JK": "C",
709
+ "WJR": "C",
710
+ "A1AAW": "C",
711
+ "85E": "C",
712
+ "MB0": "C",
713
+ "ZAB": "C",
714
+ "85K": "C",
715
+ "GBP": "C",
716
+ "A1H80": "C",
717
+ "A1AFR": "C",
718
+ "L9M": "C",
719
+ "MYK": "C",
720
+ "MB9": "C",
721
+ "38R": "C",
722
+ "EKB": "C",
723
+ "NKF": "C",
724
+ "UMQ": "C",
725
+ "T4K": "C",
726
+ "3PT": "C",
727
+ "A1A7S": "C",
728
+ "1Q9": "C",
729
+ "11R": "C",
730
+ "D2V": "C",
731
+ "SM8": "C",
732
+ "IFC": "C",
733
+ "DB5": "C",
734
+ "L2T": "C",
735
+ "GNB": "C",
736
+ "PP7": "C",
737
+ "072": "C",
738
+ "P88": "C",
739
+ "DRL": "C",
740
+ "C9W": "C",
741
+ "NTP": "C",
742
+ "4HJ": "C",
743
+ "7NA": "C",
744
+ "LPC": "C",
745
+ "T8W": "C",
746
+ "63R": "C",
747
+ "570": "C",
748
+ "R4A": "C",
749
+ "3BG": "C",
750
+ "4RB": "C",
751
+ "GSO": "C",
752
+ "BQ6": "C",
753
+ "R4P": "C",
754
+ "5CP": "C",
755
+ "TTR": "C",
756
+ "6UZ": "C",
757
+ "SPJ": "C",
758
+ "0SA": "C",
759
+ "ZL1": "C",
760
+ "BYG": "C",
761
+ "F0E": "C",
762
+ "PC0": "C",
763
+ "B2Q": "C",
764
+ "KV6": "C",
765
+ "NTO": "C",
766
+ "CLG": "C",
767
+ "R7U": "C",
768
+ "SMQ": "C",
769
+ "GM2": "C",
770
+ "Z7P": "C",
771
+ "NXF": "C",
772
+ "C6Q": "C",
773
+ "A1G": "C",
774
+ "433": "C",
775
+ "L9N": "C",
776
+ "7OX": "C",
777
+ "A1H84": "C",
778
+ "97L": "C",
779
+ "HDV": "C",
780
+ "LUO": "C",
781
+ "R6A": "C",
782
+ "1PC": "C",
783
+ "4PT": "C",
784
+ "SBZ": "C",
785
+ "EAB": "C",
786
+ "FL4": "C",
787
+ "OPS": "C",
788
+ "C2X": "C",
789
+ "SLL": "C",
790
+ "BFC": "C",
791
+ "GIP": "C",
792
+ "7CP": "C",
793
+ "CLH": "C",
794
+ "34E": "C",
795
+ "5NE": "C",
796
+ "PBF": "C",
797
+ "ABD": "C",
798
+ "ABC": "C",
799
+ "LPF": "C",
800
+ "TIZ": "C",
801
+ "4HH": "C",
802
+ "AFC": "C",
803
+ "WQH": "C",
804
+ "9JL": "C",
805
+ "CS3": "C",
806
+ "NL0": "C",
807
+ "KPY": "C",
808
+ "DNA": "C",
809
+ "B3C": "C",
810
+ "TKL": "C",
811
+ "KVS": "C",
812
+ "HO6": "C",
813
+ "NLH": "C",
814
+ "1PB": "C",
815
+ "CYF": "C",
816
+ "G4M": "C",
817
+ "R5B": "C",
818
+ "N4S": "C",
819
+ "N11": "C",
820
+ "C8F": "C",
821
+ "PIJ": "C",
822
+ "WIN": "C",
823
+ "NT1": "C",
824
+ "WJW": "C",
825
+ "HF7": "C",
826
+ "TY1": "C",
827
+ "VM1": "C",
828
+ },
829
+ "OS": {"*": "O", "DWC": "OS", "OHX": "OS", "OS": "OS", "8WV": "OS", "OS4": "OS"},
830
+ "PB": {"*": "P", "ZN9": "PB", "ZN7": "PB", "PBM": "PB", "PB": "PB", "CSB": "PB"},
831
+ "CE": {"*": "C", "CE": "CE"},
832
+ "FE": {"*": "FE", "TFR": "F", "PF5": "F", "IFC": "F", "F5C": "F"},
833
+ "NA": {"*": "N", "CGO": "NA", "R2K": "NA", "LVQ": "NA", "NA": "NA"},
834
+ "ND": {"*": "N", "ND": "ND"},
835
+ "CF": {"*": "C", "CF": "CF"},
836
+ "RU": "RU",
837
+ "BRAF": "BR",
838
+ "EU": "EU",
839
+ "CLAA": "CL",
840
+ "CLBQ": "CL",
841
+ "CM": {"*": "C", "ZCM": "CM"},
842
+ "SN": {"*": "SN", "TAP": "S", "SND": "S", "TAD": "S", "XPT": "S"},
843
+ "AG": "AG",
844
+ "CLN": "CL",
845
+ "CLM": "CL",
846
+ "CLA": {"*": "CL", "PII": "C", "TDL": "C", "D0J": "C", "GM2": "C", "PIJ": "C"},
847
+ "CLB": {
848
+ "*": "CL",
849
+ "TD5": "C",
850
+ "PII": "C",
851
+ "TDL": "C",
852
+ "GM2": "C",
853
+ "TD7": "C",
854
+ "TD6": "C",
855
+ "PIJ": "C",
856
+ },
857
+ "CR": {
858
+ "*": "C",
859
+ "BW9": "CR",
860
+ "CQ4": "CR",
861
+ "AC9": "CR",
862
+ "TIL": "CR",
863
+ "J7U": "CR",
864
+ "CR": "CR",
865
+ },
866
+ "CLAY": "CL",
867
+ "CLBC": "CL",
868
+ "PD": {
869
+ "*": "P",
870
+ "F6Q": "PD",
871
+ "SVP": "PD",
872
+ "SXC": "PD",
873
+ "U5U": "PD",
874
+ "PD": "PD",
875
+ "PLL": "PD",
876
+ },
877
+ "CO": {
878
+ "*": "C",
879
+ "J1S": "CO",
880
+ "OCN": "CO",
881
+ "OL3": "CO",
882
+ "OL4": "CO",
883
+ "B12": "CO",
884
+ "XCO": "CO",
885
+ "UFU": "CO",
886
+ "CON": "CO",
887
+ "OL5": "CO",
888
+ "B13": "CO",
889
+ "7KI": "CO",
890
+ "PL1": "CO",
891
+ "OCO": "CO",
892
+ "J1R": "CO",
893
+ "COH": "CO",
894
+ "SIR": "CO",
895
+ "6KI": "CO",
896
+ "NCO": "CO",
897
+ "9CO": "CO",
898
+ "PC3": "CO",
899
+ "BWU": "CO",
900
+ "B1Z": "CO",
901
+ "J83": "CO",
902
+ "CO": "CO",
903
+ "COY": "CO",
904
+ "CNC": "CO",
905
+ "3CO": "CO",
906
+ "OCL": "CO",
907
+ "R5Q": "CO",
908
+ "X5Z": "CO",
909
+ "CBY": "CO",
910
+ "OLS": "CO",
911
+ "F0X": "CO",
912
+ "I2A": "CO",
913
+ "OCM": "CO",
914
+ },
915
+ "CU": {
916
+ "*": "C",
917
+ "8ZR": "CU",
918
+ "K7E": "CU",
919
+ "CU3": "CU",
920
+ "SI9": "CU",
921
+ "35N": "CU",
922
+ "C2O": "CU",
923
+ "SI7": "CU",
924
+ "B15": "CU",
925
+ "SI0": "CU",
926
+ "CUP": "CU",
927
+ "SQ1": "CU",
928
+ "CUK": "CU",
929
+ "CUL": "CU",
930
+ "SI8": "CU",
931
+ "IC4": "CU",
932
+ "CUM": "CU",
933
+ "MM2": "CU",
934
+ "B30": "CU",
935
+ "S32": "CU",
936
+ "V79": "CU",
937
+ "IMF": "CU",
938
+ "CUN": "CU",
939
+ "MM1": "CU",
940
+ "MP1": "CU",
941
+ "IME": "CU",
942
+ "B17": "CU",
943
+ "C2C": "CU",
944
+ "1CU": "CU",
945
+ "CU6": "CU",
946
+ "C1O": "CU",
947
+ "CU1": "CU",
948
+ "B22": "CU",
949
+ "CUS": "CU",
950
+ "RUQ": "CU",
951
+ "CUF": "CU",
952
+ "CUA": "CU",
953
+ "CU": "CU",
954
+ "CUO": "CU",
955
+ "0TE": "CU",
956
+ "SI4": "CU",
957
+ },
958
+ "CS": {"*": "C", "CS": "CS"},
959
+ "CLQ": "CL",
960
+ "CLR": "CL",
961
+ "CLU": "CL",
962
+ "TE": "TE",
963
+ "NI": {
964
+ "*": "N",
965
+ "USN": "NI",
966
+ "NFO": "NI",
967
+ "NI2": "NI",
968
+ "NFS": "NI",
969
+ "NFR": "NI",
970
+ "82N": "NI",
971
+ "R5N": "NI",
972
+ "NFU": "NI",
973
+ "A1ICD": "NI",
974
+ "NI3": "NI",
975
+ "M43": "NI",
976
+ "MM5": "NI",
977
+ "BF8": "NI",
978
+ "TCN": "NI",
979
+ "NIK": "NI",
980
+ "CUV": "NI",
981
+ "MM6": "NI",
982
+ "J52": "NI",
983
+ "NI": "NI",
984
+ "SNF": "NI",
985
+ "XCC": "NI",
986
+ "F0L": "NI",
987
+ "UWE": "NI",
988
+ "NFC": "NI",
989
+ "3NI": "NI",
990
+ "HNI": "NI",
991
+ "F43": "NI",
992
+ "RQM": "NI",
993
+ "NFE": "NI",
994
+ "NFB": "NI",
995
+ "B51": "NI",
996
+ "NI1": "NI",
997
+ "WCC": "NI",
998
+ "NUF": "NI",
999
+ },
1000
+ "SB": {"*": "S", "UJI": "SB", "SB": "SB", "118": "SB", "SBO": "SB", "3CG": "SB"},
1001
+ "MO": "MO",
1002
+ "SEG": "SE",
1003
+ "CLL": "CL",
1004
+ "CLAH": "CL",
1005
+ "CLC": {
1006
+ "*": "CL",
1007
+ "TD5": "C",
1008
+ "PII": "C",
1009
+ "TDL": "C",
1010
+ "GM2": "C",
1011
+ "TD7": "C",
1012
+ "TD6": "C",
1013
+ "PIJ": "C",
1014
+ },
1015
+ "CLD": {"*": "CL", "PII": "C", "GM2": "C", "PIJ": "C"},
1016
+ "CLAD": "CL",
1017
+ "CLAE": "CL",
1018
+ "LA": "LA",
1019
+ "RH": "RH",
1020
+ "BRAC": "BR",
1021
+ "BRAD": "BR",
1022
+ "CLBN": "CL",
1023
+ "CLAC": "CL",
1024
+ "BRAB": "BR",
1025
+ "BRAE": "BR",
1026
+ "MG": "MG",
1027
+ "IR": "IR",
1028
+ "SE": {
1029
+ "*": "SE",
1030
+ "HII": "S",
1031
+ "NT2": "S",
1032
+ "R2P": "S",
1033
+ "S2P": "S",
1034
+ "0IU": "S",
1035
+ "QMB": "S",
1036
+ "81S": "S",
1037
+ "0QB": "S",
1038
+ "UB4": "S",
1039
+ "OHS": "S",
1040
+ "Q78": "S",
1041
+ "0Y2": "S",
1042
+ "B3M": "S",
1043
+ "NT1": "S",
1044
+ "81R": "S",
1045
+ },
1046
+ "BRAG": "BR",
1047
+ "CLF": {"*": "CL", "PII": "C", "GM2": "C", "PIJ": "C"},
1048
+ "CLE": {"*": "CL", "PII": "C", "GM2": "C", "PIJ": "C"},
1049
+ "BRAX": "BR",
1050
+ "CLK": "CL",
1051
+ "ZN": "ZN",
1052
+ "AS": "AS",
1053
+ "AU": "AU",
1054
+ "PT": "PT",
1055
+ "CLAS": "CL",
1056
+ "MN": "MN",
1057
+ "CLBE": "CL",
1058
+ "CLBF": "CL",
1059
+ "CLAF": "CL",
1060
+ "NA'": {"*": "N", "CGO": "NA"},
1061
+ "BRAH": "BR",
1062
+ "BRAI": "BR",
1063
+ "BRA": "BR",
1064
+ "BRB": "BR",
1065
+ "BRAV": "BR",
1066
+ "HG": {
1067
+ "*": "HG",
1068
+ "BBA": "H",
1069
+ "MID": "H",
1070
+ "APM": "H",
1071
+ "4QQ": "H",
1072
+ "0ZG": "H",
1073
+ "APH": "H",
1074
+ },
1075
+ "AR": "AR",
1076
+ "D": "H",
1077
+ "CLAN": "CL",
1078
+ "SI": "SI",
1079
+ "CLS": "CL",
1080
+ "ZR": "ZR",
1081
+ "CLAR": {"*": "CL", "ZM4": "C"},
1082
+ "HO": "HO",
1083
+ "CLI": {"*": "CL", "GM2": "C"},
1084
+ "CLH": {"*": "CL", "GM2": "C"},
1085
+ "CLAP": "CL",
1086
+ "CLBL": "CL",
1087
+ "CLBM": "CL",
1088
+ "PR": {"*": "PR", "UF0": "P", "252": "P"},
1089
+ "IN": "IN",
1090
+ "CLJ": "CL",
1091
+ "BRU": "BR",
1092
+ "SC": {"*": "S", "SFL": "SC"},
1093
+ "CLG": {"*": "CL", "GM2": "C"},
1094
+ "BRAT": "BR",
1095
+ "BRAR": "BR",
1096
+ "CLAG": "CL",
1097
+ "CLAB": "CL",
1098
+ "CLV": "CL",
1099
+ "TI": "TI",
1100
+ "CLAX": "CL",
1101
+ "CLAJ": "CL",
1102
+ "CL'": {"*": "CL", "BNR": "C", "25A": "C", "BDA": "C"},
1103
+ "CLAW": "CL",
1104
+ "BRF": "BR",
1105
+ "BRE": "BR",
1106
+ "RE": "RE",
1107
+ "GD": "GD",
1108
+ "SM": {"*": "S", "SM": "SM"},
1109
+ "CLBH": "CL",
1110
+ "CLBI": "CL",
1111
+ "CLAI": "CL",
1112
+ "CLY": "CL",
1113
+ "CLZ": "CL",
1114
+ "AC": "AC",
1115
+ "BR'": "BR",
1116
+ "CLT": "CL",
1117
+ "CLO": "CL",
1118
+ "CLP": "CL",
1119
+ "LU": "LU",
1120
+ "BA": {"*": "B", "BA": "BA"},
1121
+ "CLAU": "CL",
1122
+ "RB": "RB",
1123
+ "LI": "LI",
1124
+ "MOM": "MO",
1125
+ "BRAQ": "BR",
1126
+ "SR": {"*": "S", "SR": "SR", "OER": "SR"},
1127
+ "CLAT": "CL",
1128
+ "BRAL": "BR",
1129
+ "SEB": "SE",
1130
+ "CLW": "CL",
1131
+ "CLX": "CL",
1132
+ "BE": "BE",
1133
+ "BRG": "BR",
1134
+ "SEA": "SE",
1135
+ "BRAW": "BR",
1136
+ "BRBB": "BR",
1137
+ "ER": "ER",
1138
+ "TH": "TH",
1139
+ "BRR": "BR",
1140
+ "CLBV": "CL",
1141
+ "AL": "AL",
1142
+ "CLAV": "CL",
1143
+ "BRH": "BR",
1144
+ "CLAQ": "CL",
1145
+ "GA": "GA",
1146
+ "X": "*",
1147
+ "TL": "TL",
1148
+ "CLBB": "CL",
1149
+ "TB": "TB",
1150
+ "CLAK": "CL",
1151
+ "XE": {"*": "*", "XE": "XE"},
1152
+ "SEL": "SE",
1153
+ "PU": {"*": "P", "4PU": "PU"},
1154
+ "CLAZ": "CL",
1155
+ "SE'": "SE",
1156
+ "CLBA": "CL",
1157
+ "SEN": "SE",
1158
+ "SNN": "SN",
1159
+ "MOB": "MO",
1160
+ "YB": "YB",
1161
+ "BRC": "BR",
1162
+ "BRD": "BR",
1163
+ "CLAM": "CL",
1164
+ "DA": "H",
1165
+ "DB": "H",
1166
+ "DC": "H",
1167
+ "DXT": "H",
1168
+ "DXU": "H",
1169
+ "DXX": "H",
1170
+ "DXY": "H",
1171
+ "DXZ": "H",
1172
+ "DY": "DY",
1173
+ "TA": "TA",
1174
+ "XD": "*",
1175
+ "SED": "SE",
1176
+ "CLAL": "CL",
1177
+ "BRAJ": "BR",
1178
+ "AM": "AM",
1179
+ "CLAO": "CL",
1180
+ "BI": "BI",
1181
+ "KR": "KR",
1182
+ "BRBJ": "BR",
1183
+ "UNK": "*",
1184
+ }
vb_layers_attention.py CHANGED
@@ -1,136 +1,136 @@
1
- from typing import Optional
2
-
3
- import torch
4
- from einops.layers.torch import Rearrange
5
- from torch import Tensor, nn
6
-
7
- from . import vb_layers_initialize as init
8
-
9
-
10
- class AttentionPairBias(nn.Module):
11
- """Attention pair bias layer."""
12
-
13
- def __init__(
14
- self,
15
- c_s: int,
16
- c_z: int,
17
- num_heads: int,
18
- inf: float = 1e6,
19
- initial_norm: bool = True,
20
- ) -> None:
21
- """Initialize the attention pair bias layer.
22
-
23
- Parameters
24
- ----------
25
- c_s : int
26
- The input sequence dimension.
27
- c_z : int
28
- The input pairwise dimension.
29
- num_heads : int
30
- The number of heads.
31
- inf : float, optional
32
- The inf value, by default 1e6
33
- initial_norm: bool, optional
34
- Whether to apply layer norm to the input, by default True
35
-
36
- """
37
- super().__init__()
38
-
39
- assert c_s % num_heads == 0
40
-
41
- self.c_s = c_s
42
- self.num_heads = num_heads
43
- self.head_dim = c_s // num_heads
44
- self.inf = inf
45
-
46
- self.initial_norm = initial_norm
47
- if self.initial_norm:
48
- self.norm_s = nn.LayerNorm(c_s)
49
-
50
- self.proj_q = nn.Linear(c_s, c_s)
51
- self.proj_k = nn.Linear(c_s, c_s, bias=False)
52
- self.proj_v = nn.Linear(c_s, c_s, bias=False)
53
- self.proj_g = nn.Linear(c_s, c_s, bias=False)
54
-
55
- self.proj_z = nn.Sequential(
56
- nn.LayerNorm(c_z),
57
- nn.Linear(c_z, num_heads, bias=False),
58
- Rearrange("b ... h -> b h ..."),
59
- )
60
-
61
- self.proj_o = nn.Linear(c_s, c_s, bias=False)
62
- init.final_init_(self.proj_o.weight)
63
-
64
- def forward(
65
- self,
66
- s: Tensor,
67
- z: Tensor,
68
- mask: Tensor,
69
- k_in: Optional[Tensor] = None,
70
- multiplicity: int = 1,
71
- to_keys=None,
72
- model_cache=None,
73
- ) -> Tensor:
74
- """Forward pass.
75
-
76
- Parameters
77
- ----------
78
- s : torch.Tensor
79
- The input sequence tensor (B, S, D)
80
- z : torch.Tensor
81
- The input pairwise tensor (B, N, N, D)
82
- mask : torch.Tensor
83
- The pairwise mask tensor (B, N)
84
- multiplicity : int, optional
85
- The diffusion batch size, by default 1
86
-
87
- Returns
88
- -------
89
- torch.Tensor
90
- The output sequence tensor.
91
-
92
- """
93
- B = s.shape[0]
94
-
95
- # Layer norms
96
- if self.initial_norm:
97
- s = self.norm_s(s)
98
-
99
- if to_keys is not None:
100
- k_in = to_keys(s)
101
- mask = to_keys(mask.unsqueeze(-1)).squeeze(-1)
102
- else:
103
- if k_in is None:
104
- k_in = s
105
-
106
- # Compute projections
107
- q = self.proj_q(s).view(B, -1, self.num_heads, self.head_dim)
108
- k = self.proj_k(k_in).view(B, -1, self.num_heads, self.head_dim)
109
- v = self.proj_v(k_in).view(B, -1, self.num_heads, self.head_dim)
110
-
111
- # Caching z projection during diffusion roll-out
112
- if model_cache is None or "z" not in model_cache:
113
- z = self.proj_z(z)
114
-
115
- if model_cache is not None:
116
- model_cache["z"] = z
117
- else:
118
- z = model_cache["z"]
119
- z = z.repeat_interleave(multiplicity, 0)
120
-
121
- g = self.proj_g(s).sigmoid()
122
-
123
- with torch.autocast("cuda", enabled=False):
124
- # Compute attention weights
125
- attn = torch.einsum("bihd,bjhd->bhij", q.float(), k.float())
126
- attn = attn / (self.head_dim**0.5) + z.float()
127
- # The pairwise mask tensor (B, N) is broadcasted to (B, 1, 1, N) and (B, H, N, N)
128
- attn = attn + (1 - mask[:, None, None].float()) * -self.inf
129
- attn = attn.softmax(dim=-1)
130
-
131
- # Compute output
132
- o = torch.einsum("bhij,bjhd->bihd", attn, v.float()).to(v.dtype)
133
- o = o.reshape(B, -1, self.c_s)
134
- o = self.proj_o(g * o)
135
-
136
- return o
 
1
+ from typing import Optional
2
+
3
+ import torch
4
+ from einops.layers.torch import Rearrange
5
+ from torch import Tensor, nn
6
+
7
+ from . import vb_layers_initialize as init
8
+
9
+
10
+ class AttentionPairBias(nn.Module):
11
+ """Attention pair bias layer."""
12
+
13
+ def __init__(
14
+ self,
15
+ c_s: int,
16
+ c_z: int,
17
+ num_heads: int,
18
+ inf: float = 1e6,
19
+ initial_norm: bool = True,
20
+ ) -> None:
21
+ """Initialize the attention pair bias layer.
22
+
23
+ Parameters
24
+ ----------
25
+ c_s : int
26
+ The input sequence dimension.
27
+ c_z : int
28
+ The input pairwise dimension.
29
+ num_heads : int
30
+ The number of heads.
31
+ inf : float, optional
32
+ The inf value, by default 1e6
33
+ initial_norm: bool, optional
34
+ Whether to apply layer norm to the input, by default True
35
+
36
+ """
37
+ super().__init__()
38
+
39
+ assert c_s % num_heads == 0
40
+
41
+ self.c_s = c_s
42
+ self.num_heads = num_heads
43
+ self.head_dim = c_s // num_heads
44
+ self.inf = inf
45
+
46
+ self.initial_norm = initial_norm
47
+ if self.initial_norm:
48
+ self.norm_s = nn.LayerNorm(c_s)
49
+
50
+ self.proj_q = nn.Linear(c_s, c_s)
51
+ self.proj_k = nn.Linear(c_s, c_s, bias=False)
52
+ self.proj_v = nn.Linear(c_s, c_s, bias=False)
53
+ self.proj_g = nn.Linear(c_s, c_s, bias=False)
54
+
55
+ self.proj_z = nn.Sequential(
56
+ nn.LayerNorm(c_z),
57
+ nn.Linear(c_z, num_heads, bias=False),
58
+ Rearrange("b ... h -> b h ..."),
59
+ )
60
+
61
+ self.proj_o = nn.Linear(c_s, c_s, bias=False)
62
+ init.final_init_(self.proj_o.weight)
63
+
64
+ def forward(
65
+ self,
66
+ s: Tensor,
67
+ z: Tensor,
68
+ mask: Tensor,
69
+ k_in: Optional[Tensor] = None,
70
+ multiplicity: int = 1,
71
+ to_keys=None,
72
+ model_cache=None,
73
+ ) -> Tensor:
74
+ """Forward pass.
75
+
76
+ Parameters
77
+ ----------
78
+ s : torch.Tensor
79
+ The input sequence tensor (B, S, D)
80
+ z : torch.Tensor
81
+ The input pairwise tensor (B, N, N, D)
82
+ mask : torch.Tensor
83
+ The pairwise mask tensor (B, N)
84
+ multiplicity : int, optional
85
+ The diffusion batch size, by default 1
86
+
87
+ Returns
88
+ -------
89
+ torch.Tensor
90
+ The output sequence tensor.
91
+
92
+ """
93
+ B = s.shape[0]
94
+
95
+ # Layer norms
96
+ if self.initial_norm:
97
+ s = self.norm_s(s)
98
+
99
+ if to_keys is not None:
100
+ k_in = to_keys(s)
101
+ mask = to_keys(mask.unsqueeze(-1)).squeeze(-1)
102
+ else:
103
+ if k_in is None:
104
+ k_in = s
105
+
106
+ # Compute projections
107
+ q = self.proj_q(s).view(B, -1, self.num_heads, self.head_dim)
108
+ k = self.proj_k(k_in).view(B, -1, self.num_heads, self.head_dim)
109
+ v = self.proj_v(k_in).view(B, -1, self.num_heads, self.head_dim)
110
+
111
+ # Caching z projection during diffusion roll-out
112
+ if model_cache is None or "z" not in model_cache:
113
+ z = self.proj_z(z)
114
+
115
+ if model_cache is not None:
116
+ model_cache["z"] = z
117
+ else:
118
+ z = model_cache["z"]
119
+ z = z.repeat_interleave(multiplicity, 0)
120
+
121
+ g = self.proj_g(s).sigmoid()
122
+
123
+ with torch.autocast("cuda", enabled=False):
124
+ # Compute attention weights
125
+ attn = torch.einsum("bihd,bjhd->bhij", q.float(), k.float())
126
+ attn = attn / (self.head_dim**0.5) + z.float()
127
+ # The pairwise mask tensor (B, N) is broadcasted to (B, 1, 1, N) and (B, H, N, N)
128
+ attn = attn + (1 - mask[:, None, None].float()) * -self.inf
129
+ attn = attn.softmax(dim=-1)
130
+
131
+ # Compute output
132
+ o = torch.einsum("bhij,bjhd->bihd", attn, v.float()).to(v.dtype)
133
+ o = o.reshape(B, -1, self.c_s)
134
+ o = self.proj_o(g * o)
135
+
136
+ return o
vb_layers_attentionv2.py CHANGED
@@ -1,111 +1,111 @@
1
- from typing import Optional
2
-
3
- import torch
4
- from einops.layers.torch import Rearrange
5
- from torch import Tensor, nn
6
-
7
- from . import vb_layers_initialize as init
8
-
9
-
10
- class AttentionPairBias(nn.Module):
11
- """Attention pair bias layer."""
12
-
13
- def __init__(
14
- self,
15
- c_s: int,
16
- c_z: Optional[int] = None,
17
- num_heads: Optional[int] = None,
18
- inf: float = 1e6,
19
- compute_pair_bias: bool = True,
20
- ) -> None:
21
- """Initialize the attention pair bias layer.
22
-
23
- Parameters
24
- ----------
25
- c_s : int
26
- The input sequence dimension.
27
- c_z : int
28
- The input pairwise dimension.
29
- num_heads : int
30
- The number of heads.
31
- inf : float, optional
32
- The inf value, by default 1e6
33
-
34
- """
35
- super().__init__()
36
-
37
- assert c_s % num_heads == 0
38
-
39
- self.c_s = c_s
40
- self.num_heads = num_heads
41
- self.head_dim = c_s // num_heads
42
- self.inf = inf
43
-
44
- self.proj_q = nn.Linear(c_s, c_s)
45
- self.proj_k = nn.Linear(c_s, c_s, bias=False)
46
- self.proj_v = nn.Linear(c_s, c_s, bias=False)
47
- self.proj_g = nn.Linear(c_s, c_s, bias=False)
48
-
49
- self.compute_pair_bias = compute_pair_bias
50
- if compute_pair_bias:
51
- self.proj_z = nn.Sequential(
52
- nn.LayerNorm(c_z),
53
- nn.Linear(c_z, num_heads, bias=False),
54
- Rearrange("b ... h -> b h ..."),
55
- )
56
- else:
57
- self.proj_z = Rearrange("b ... h -> b h ...")
58
-
59
- self.proj_o = nn.Linear(c_s, c_s, bias=False)
60
- init.final_init_(self.proj_o.weight)
61
-
62
- def forward(
63
- self,
64
- s: Tensor,
65
- z: Tensor,
66
- mask: Tensor,
67
- k_in: Tensor,
68
- multiplicity: int = 1,
69
- ) -> Tensor:
70
- """Forward pass.
71
-
72
- Parameters
73
- ----------
74
- s : torch.Tensor
75
- The input sequence tensor (B, S, D)
76
- z : torch.Tensor
77
- The input pairwise tensor or bias (B, N, N, D)
78
- mask : torch.Tensor
79
- The pairwise mask tensor (B, N, N)
80
-
81
- Returns
82
- -------
83
- torch.Tensor
84
- The output sequence tensor.
85
-
86
- """
87
- B = s.shape[0]
88
-
89
- # Compute projections
90
- q = self.proj_q(s).view(B, -1, self.num_heads, self.head_dim)
91
- k = self.proj_k(k_in).view(B, -1, self.num_heads, self.head_dim)
92
- v = self.proj_v(k_in).view(B, -1, self.num_heads, self.head_dim)
93
-
94
- bias = self.proj_z(z)
95
- bias = bias.repeat_interleave(multiplicity, 0)
96
-
97
- g = self.proj_g(s).sigmoid()
98
-
99
- with torch.autocast("cuda", enabled=False):
100
- # Compute attention weights
101
- attn = torch.einsum("bihd,bjhd->bhij", q.float(), k.float())
102
- attn = attn / (self.head_dim**0.5) + bias.float()
103
- attn = attn + (1 - mask[:, None, None].float()) * -self.inf
104
- attn = attn.softmax(dim=-1)
105
-
106
- # Compute output
107
- o = torch.einsum("bhij,bjhd->bihd", attn, v.float()).to(v.dtype)
108
- o = o.reshape(B, -1, self.c_s)
109
- o = self.proj_o(g * o)
110
-
111
- return o
 
1
+ from typing import Optional
2
+
3
+ import torch
4
+ from einops.layers.torch import Rearrange
5
+ from torch import Tensor, nn
6
+
7
+ from . import vb_layers_initialize as init
8
+
9
+
10
+ class AttentionPairBias(nn.Module):
11
+ """Attention pair bias layer."""
12
+
13
+ def __init__(
14
+ self,
15
+ c_s: int,
16
+ c_z: Optional[int] = None,
17
+ num_heads: Optional[int] = None,
18
+ inf: float = 1e6,
19
+ compute_pair_bias: bool = True,
20
+ ) -> None:
21
+ """Initialize the attention pair bias layer.
22
+
23
+ Parameters
24
+ ----------
25
+ c_s : int
26
+ The input sequence dimension.
27
+ c_z : int
28
+ The input pairwise dimension.
29
+ num_heads : int
30
+ The number of heads.
31
+ inf : float, optional
32
+ The inf value, by default 1e6
33
+
34
+ """
35
+ super().__init__()
36
+
37
+ assert c_s % num_heads == 0
38
+
39
+ self.c_s = c_s
40
+ self.num_heads = num_heads
41
+ self.head_dim = c_s // num_heads
42
+ self.inf = inf
43
+
44
+ self.proj_q = nn.Linear(c_s, c_s)
45
+ self.proj_k = nn.Linear(c_s, c_s, bias=False)
46
+ self.proj_v = nn.Linear(c_s, c_s, bias=False)
47
+ self.proj_g = nn.Linear(c_s, c_s, bias=False)
48
+
49
+ self.compute_pair_bias = compute_pair_bias
50
+ if compute_pair_bias:
51
+ self.proj_z = nn.Sequential(
52
+ nn.LayerNorm(c_z),
53
+ nn.Linear(c_z, num_heads, bias=False),
54
+ Rearrange("b ... h -> b h ..."),
55
+ )
56
+ else:
57
+ self.proj_z = Rearrange("b ... h -> b h ...")
58
+
59
+ self.proj_o = nn.Linear(c_s, c_s, bias=False)
60
+ init.final_init_(self.proj_o.weight)
61
+
62
+ def forward(
63
+ self,
64
+ s: Tensor,
65
+ z: Tensor,
66
+ mask: Tensor,
67
+ k_in: Tensor,
68
+ multiplicity: int = 1,
69
+ ) -> Tensor:
70
+ """Forward pass.
71
+
72
+ Parameters
73
+ ----------
74
+ s : torch.Tensor
75
+ The input sequence tensor (B, S, D)
76
+ z : torch.Tensor
77
+ The input pairwise tensor or bias (B, N, N, D)
78
+ mask : torch.Tensor
79
+ The pairwise mask tensor (B, N, N)
80
+
81
+ Returns
82
+ -------
83
+ torch.Tensor
84
+ The output sequence tensor.
85
+
86
+ """
87
+ B = s.shape[0]
88
+
89
+ # Compute projections
90
+ q = self.proj_q(s).view(B, -1, self.num_heads, self.head_dim)
91
+ k = self.proj_k(k_in).view(B, -1, self.num_heads, self.head_dim)
92
+ v = self.proj_v(k_in).view(B, -1, self.num_heads, self.head_dim)
93
+
94
+ bias = self.proj_z(z)
95
+ bias = bias.repeat_interleave(multiplicity, 0)
96
+
97
+ g = self.proj_g(s).sigmoid()
98
+
99
+ with torch.autocast("cuda", enabled=False):
100
+ # Compute attention weights
101
+ attn = torch.einsum("bihd,bjhd->bhij", q.float(), k.float())
102
+ attn = attn / (self.head_dim**0.5) + bias.float()
103
+ attn = attn + (1 - mask[:, None, None].float()) * -self.inf
104
+ attn = attn.softmax(dim=-1)
105
+
106
+ # Compute output
107
+ o = torch.einsum("bhij,bjhd->bihd", attn, v.float()).to(v.dtype)
108
+ o = o.reshape(B, -1, self.c_s)
109
+ o = self.proj_o(g * o)
110
+
111
+ return o
vb_layers_confidence_utils.py CHANGED
@@ -1,231 +1,231 @@
1
- import torch
2
- from torch import nn
3
-
4
- from . import vb_const as const
5
-
6
-
7
- def compute_collinear_mask(v1, v2):
8
- norm1 = torch.norm(v1, dim=1, keepdim=True)
9
- norm2 = torch.norm(v2, dim=1, keepdim=True)
10
- v1 = v1 / (norm1 + 1e-6)
11
- v2 = v2 / (norm2 + 1e-6)
12
- mask_angle = torch.abs(torch.sum(v1 * v2, dim=1)) < 0.9063
13
- mask_overlap1 = norm1.reshape(-1) > 1e-2
14
- mask_overlap2 = norm2.reshape(-1) > 1e-2
15
- return mask_angle & mask_overlap1 & mask_overlap2
16
-
17
-
18
- def compute_frame_pred(
19
- pred_atom_coords,
20
- frames_idx_true,
21
- feats,
22
- multiplicity,
23
- resolved_mask=None,
24
- inference=False,
25
- ):
26
- with torch.amp.autocast("cuda", enabled=False):
27
- asym_id_token = feats["asym_id"]
28
- asym_id_atom = torch.bmm(
29
- feats["atom_to_token"].float(), asym_id_token.unsqueeze(-1).float()
30
- ).squeeze(-1)
31
-
32
- B, N, _ = pred_atom_coords.shape
33
- pred_atom_coords = pred_atom_coords.reshape(B // multiplicity, multiplicity, -1, 3)
34
- frames_idx_pred = (
35
- frames_idx_true.clone()
36
- .repeat_interleave(multiplicity, 0)
37
- .reshape(B // multiplicity, multiplicity, -1, 3)
38
- )
39
-
40
- # Iterate through the batch and modify the frames for nonpolymers
41
- for i, pred_atom_coord in enumerate(pred_atom_coords):
42
- token_idx = 0
43
- atom_idx = 0
44
- for id in torch.unique(asym_id_token[i]):
45
- mask_chain_token = (asym_id_token[i] == id) * feats["token_pad_mask"][i]
46
- mask_chain_atom = (asym_id_atom[i] == id) * feats["atom_pad_mask"][i]
47
- num_tokens = int(mask_chain_token.sum().item())
48
- num_atoms = int(mask_chain_atom.sum().item())
49
- if (
50
- feats["mol_type"][i, token_idx] != const.chain_type_ids["NONPOLYMER"]
51
- or num_atoms < 3
52
- ):
53
- token_idx += num_tokens
54
- atom_idx += num_atoms
55
- continue
56
- dist_mat = (
57
- (
58
- pred_atom_coord[:, mask_chain_atom.bool()][:, None, :, :]
59
- - pred_atom_coord[:, mask_chain_atom.bool()][:, :, None, :]
60
- )
61
- ** 2
62
- ).sum(-1) ** 0.5
63
- if inference:
64
- resolved_pair = 1 - (
65
- feats["atom_pad_mask"][i][mask_chain_atom.bool()][None, :]
66
- * feats["atom_pad_mask"][i][mask_chain_atom.bool()][:, None]
67
- ).to(torch.float32)
68
- resolved_pair[resolved_pair == 1] = torch.inf
69
- indices = torch.sort(dist_mat + resolved_pair, axis=2).indices
70
- else:
71
- if resolved_mask is None:
72
- resolved_mask = feats["atom_resolved_mask"]
73
- resolved_pair = 1 - (
74
- resolved_mask[i][mask_chain_atom.bool()][None, :]
75
- * resolved_mask[i][mask_chain_atom.bool()][:, None]
76
- ).to(torch.float32)
77
- resolved_pair[resolved_pair == 1] = torch.inf
78
- indices = torch.sort(dist_mat + resolved_pair, axis=2).indices
79
- frames = (
80
- torch.cat(
81
- [
82
- indices[:, :, 1:2],
83
- indices[:, :, 0:1],
84
- indices[:, :, 2:3],
85
- ],
86
- dim=2,
87
- )
88
- + atom_idx
89
- )
90
- try:
91
- frames_idx_pred[i, :, token_idx : token_idx + num_atoms, :] = frames
92
- except Exception as e:
93
- print(f"Failed to process {feats['pdb_id']} due to {e}")
94
- token_idx += num_tokens
95
- atom_idx += num_atoms
96
-
97
- frames_expanded = pred_atom_coords[
98
- torch.arange(0, B // multiplicity, 1)[:, None, None, None].to(
99
- frames_idx_pred.device
100
- ),
101
- torch.arange(0, multiplicity, 1)[None, :, None, None].to(
102
- frames_idx_pred.device
103
- ),
104
- frames_idx_pred,
105
- ].reshape(-1, 3, 3)
106
-
107
- # Compute masks for collinearity / overlap
108
- mask_collinear_pred = compute_collinear_mask(
109
- frames_expanded[:, 1] - frames_expanded[:, 0],
110
- frames_expanded[:, 1] - frames_expanded[:, 2],
111
- ).reshape(B // multiplicity, multiplicity, -1)
112
- return frames_idx_pred, mask_collinear_pred * feats["token_pad_mask"][:, None, :]
113
-
114
-
115
- def compute_aggregated_metric(logits, end=1.0):
116
- # Compute aggregated metric from logits
117
- num_bins = logits.shape[-1]
118
- bin_width = end / num_bins
119
- bounds = torch.arange(
120
- start=0.5 * bin_width, end=end, step=bin_width, device=logits.device
121
- )
122
- probs = nn.functional.softmax(logits, dim=-1)
123
- plddt = torch.sum(
124
- probs * bounds.view(*((1,) * len(probs.shape[:-1])), *bounds.shape),
125
- dim=-1,
126
- )
127
- return plddt
128
-
129
-
130
- def tm_function(d, Nres):
131
- d0 = 1.24 * (torch.clip(Nres, min=19) - 15) ** (1 / 3) - 1.8
132
- return 1 / (1 + (d / d0) ** 2)
133
-
134
-
135
- def compute_ptms(logits, x_preds, feats, multiplicity):
136
- # It needs to take as input the mask of the frames as they are not used to compute the PTM
137
- _, mask_collinear_pred = compute_frame_pred(
138
- x_preds, feats["frames_idx"], feats, multiplicity, inference=True
139
- )
140
- # mask overlapping, collinear tokens and ions (invalid frames)
141
- mask_pad = feats["token_pad_mask"].repeat_interleave(multiplicity, 0)
142
- maski = mask_collinear_pred.reshape(-1, mask_collinear_pred.shape[-1])
143
- pair_mask_ptm = maski[:, :, None] * mask_pad[:, None, :] * mask_pad[:, :, None]
144
- asym_id = feats["asym_id"].repeat_interleave(multiplicity, 0)
145
- pair_mask_iptm = (
146
- maski[:, :, None]
147
- * (asym_id[:, None, :] != asym_id[:, :, None])
148
- * mask_pad[:, None, :]
149
- * mask_pad[:, :, None]
150
- )
151
- num_bins = logits.shape[-1]
152
- bin_width = 32.0 / num_bins
153
- end = 32.0
154
- pae_value = torch.arange(
155
- start=0.5 * bin_width, end=end, step=bin_width, device=logits.device
156
- ).unsqueeze(0)
157
- N_res = mask_pad.sum(dim=-1, keepdim=True)
158
- tm_value = tm_function(pae_value, N_res).unsqueeze(1).unsqueeze(2)
159
- probs = nn.functional.softmax(logits, dim=-1)
160
- tm_expected_value = torch.sum(
161
- probs * tm_value,
162
- dim=-1,
163
- ) # shape (B, N, N)
164
- ptm = torch.max(
165
- torch.sum(tm_expected_value * pair_mask_ptm, dim=-1)
166
- / (torch.sum(pair_mask_ptm, dim=-1) + 1e-5),
167
- dim=1,
168
- ).values
169
- iptm = torch.max(
170
- torch.sum(tm_expected_value * pair_mask_iptm, dim=-1)
171
- / (torch.sum(pair_mask_iptm, dim=-1) + 1e-5),
172
- dim=1,
173
- ).values
174
-
175
- # compute ligand and protein iPTM
176
- token_type = feats["mol_type"]
177
- token_type = token_type.repeat_interleave(multiplicity, 0)
178
- is_ligand_token = (token_type == const.chain_type_ids["NONPOLYMER"]).float()
179
- is_protein_token = (token_type == const.chain_type_ids["PROTEIN"]).float()
180
-
181
- ligand_iptm_mask = (
182
- maski[:, :, None]
183
- * (asym_id[:, None, :] != asym_id[:, :, None])
184
- * mask_pad[:, None, :]
185
- * mask_pad[:, :, None]
186
- * (
187
- (is_ligand_token[:, :, None] * is_protein_token[:, None, :])
188
- + (is_protein_token[:, :, None] * is_ligand_token[:, None, :])
189
- )
190
- )
191
- protein_ipmt_mask = (
192
- maski[:, :, None]
193
- * (asym_id[:, None, :] != asym_id[:, :, None])
194
- * mask_pad[:, None, :]
195
- * mask_pad[:, :, None]
196
- * (is_protein_token[:, :, None] * is_protein_token[:, None, :])
197
- )
198
-
199
- ligand_iptm = torch.max(
200
- torch.sum(tm_expected_value * ligand_iptm_mask, dim=-1)
201
- / (torch.sum(ligand_iptm_mask, dim=-1) + 1e-5),
202
- dim=1,
203
- ).values
204
- protein_iptm = torch.max(
205
- torch.sum(tm_expected_value * protein_ipmt_mask, dim=-1)
206
- / (torch.sum(protein_ipmt_mask, dim=-1) + 1e-5),
207
- dim=1,
208
- ).values
209
-
210
- # Compute pair chain ipTM
211
- chain_pair_iptm = {}
212
- asym_ids_list = torch.unique(asym_id).tolist()
213
- for idx1 in asym_ids_list:
214
- chain_iptm = {}
215
- for idx2 in asym_ids_list:
216
- mask_pair_chain = (
217
- maski[:, :, None]
218
- * (asym_id[:, None, :] == idx1)
219
- * (asym_id[:, :, None] == idx2)
220
- * mask_pad[:, None, :]
221
- * mask_pad[:, :, None]
222
- )
223
-
224
- chain_iptm[idx2] = torch.max(
225
- torch.sum(tm_expected_value * mask_pair_chain, dim=-1)
226
- / (torch.sum(mask_pair_chain, dim=-1) + 1e-5),
227
- dim=1,
228
- ).values
229
- chain_pair_iptm[idx1] = chain_iptm
230
-
231
- return ptm, iptm, ligand_iptm, protein_iptm, chain_pair_iptm
 
1
+ import torch
2
+ from torch import nn
3
+
4
+ from . import vb_const as const
5
+
6
+
7
+ def compute_collinear_mask(v1, v2):
8
+ norm1 = torch.norm(v1, dim=1, keepdim=True)
9
+ norm2 = torch.norm(v2, dim=1, keepdim=True)
10
+ v1 = v1 / (norm1 + 1e-6)
11
+ v2 = v2 / (norm2 + 1e-6)
12
+ mask_angle = torch.abs(torch.sum(v1 * v2, dim=1)) < 0.9063
13
+ mask_overlap1 = norm1.reshape(-1) > 1e-2
14
+ mask_overlap2 = norm2.reshape(-1) > 1e-2
15
+ return mask_angle & mask_overlap1 & mask_overlap2
16
+
17
+
18
+ def compute_frame_pred(
19
+ pred_atom_coords,
20
+ frames_idx_true,
21
+ feats,
22
+ multiplicity,
23
+ resolved_mask=None,
24
+ inference=False,
25
+ ):
26
+ with torch.amp.autocast("cuda", enabled=False):
27
+ asym_id_token = feats["asym_id"]
28
+ asym_id_atom = torch.bmm(
29
+ feats["atom_to_token"].float(), asym_id_token.unsqueeze(-1).float()
30
+ ).squeeze(-1)
31
+
32
+ B, N, _ = pred_atom_coords.shape
33
+ pred_atom_coords = pred_atom_coords.reshape(B // multiplicity, multiplicity, -1, 3)
34
+ frames_idx_pred = (
35
+ frames_idx_true.clone()
36
+ .repeat_interleave(multiplicity, 0)
37
+ .reshape(B // multiplicity, multiplicity, -1, 3)
38
+ )
39
+
40
+ # Iterate through the batch and modify the frames for nonpolymers
41
+ for i, pred_atom_coord in enumerate(pred_atom_coords):
42
+ token_idx = 0
43
+ atom_idx = 0
44
+ for id in torch.unique(asym_id_token[i]):
45
+ mask_chain_token = (asym_id_token[i] == id) * feats["token_pad_mask"][i]
46
+ mask_chain_atom = (asym_id_atom[i] == id) * feats["atom_pad_mask"][i]
47
+ num_tokens = int(mask_chain_token.sum().item())
48
+ num_atoms = int(mask_chain_atom.sum().item())
49
+ if (
50
+ feats["mol_type"][i, token_idx] != const.chain_type_ids["NONPOLYMER"]
51
+ or num_atoms < 3
52
+ ):
53
+ token_idx += num_tokens
54
+ atom_idx += num_atoms
55
+ continue
56
+ dist_mat = (
57
+ (
58
+ pred_atom_coord[:, mask_chain_atom.bool()][:, None, :, :]
59
+ - pred_atom_coord[:, mask_chain_atom.bool()][:, :, None, :]
60
+ )
61
+ ** 2
62
+ ).sum(-1) ** 0.5
63
+ if inference:
64
+ resolved_pair = 1 - (
65
+ feats["atom_pad_mask"][i][mask_chain_atom.bool()][None, :]
66
+ * feats["atom_pad_mask"][i][mask_chain_atom.bool()][:, None]
67
+ ).to(torch.float32)
68
+ resolved_pair[resolved_pair == 1] = torch.inf
69
+ indices = torch.sort(dist_mat + resolved_pair, axis=2).indices
70
+ else:
71
+ if resolved_mask is None:
72
+ resolved_mask = feats["atom_resolved_mask"]
73
+ resolved_pair = 1 - (
74
+ resolved_mask[i][mask_chain_atom.bool()][None, :]
75
+ * resolved_mask[i][mask_chain_atom.bool()][:, None]
76
+ ).to(torch.float32)
77
+ resolved_pair[resolved_pair == 1] = torch.inf
78
+ indices = torch.sort(dist_mat + resolved_pair, axis=2).indices
79
+ frames = (
80
+ torch.cat(
81
+ [
82
+ indices[:, :, 1:2],
83
+ indices[:, :, 0:1],
84
+ indices[:, :, 2:3],
85
+ ],
86
+ dim=2,
87
+ )
88
+ + atom_idx
89
+ )
90
+ try:
91
+ frames_idx_pred[i, :, token_idx : token_idx + num_atoms, :] = frames
92
+ except Exception as e:
93
+ print(f"Failed to process {feats['pdb_id']} due to {e}")
94
+ token_idx += num_tokens
95
+ atom_idx += num_atoms
96
+
97
+ frames_expanded = pred_atom_coords[
98
+ torch.arange(0, B // multiplicity, 1)[:, None, None, None].to(
99
+ frames_idx_pred.device
100
+ ),
101
+ torch.arange(0, multiplicity, 1)[None, :, None, None].to(
102
+ frames_idx_pred.device
103
+ ),
104
+ frames_idx_pred,
105
+ ].reshape(-1, 3, 3)
106
+
107
+ # Compute masks for collinearity / overlap
108
+ mask_collinear_pred = compute_collinear_mask(
109
+ frames_expanded[:, 1] - frames_expanded[:, 0],
110
+ frames_expanded[:, 1] - frames_expanded[:, 2],
111
+ ).reshape(B // multiplicity, multiplicity, -1)
112
+ return frames_idx_pred, mask_collinear_pred * feats["token_pad_mask"][:, None, :]
113
+
114
+
115
+ def compute_aggregated_metric(logits, end=1.0):
116
+ # Compute aggregated metric from logits
117
+ num_bins = logits.shape[-1]
118
+ bin_width = end / num_bins
119
+ bounds = torch.arange(
120
+ start=0.5 * bin_width, end=end, step=bin_width, device=logits.device
121
+ )
122
+ probs = nn.functional.softmax(logits, dim=-1)
123
+ plddt = torch.sum(
124
+ probs * bounds.view(*((1,) * len(probs.shape[:-1])), *bounds.shape),
125
+ dim=-1,
126
+ )
127
+ return plddt
128
+
129
+
130
+ def tm_function(d, Nres):
131
+ d0 = 1.24 * (torch.clip(Nres, min=19) - 15) ** (1 / 3) - 1.8
132
+ return 1 / (1 + (d / d0) ** 2)
133
+
134
+
135
+ def compute_ptms(logits, x_preds, feats, multiplicity):
136
+ # It needs to take as input the mask of the frames as they are not used to compute the PTM
137
+ _, mask_collinear_pred = compute_frame_pred(
138
+ x_preds, feats["frames_idx"], feats, multiplicity, inference=True
139
+ )
140
+ # mask overlapping, collinear tokens and ions (invalid frames)
141
+ mask_pad = feats["token_pad_mask"].repeat_interleave(multiplicity, 0)
142
+ maski = mask_collinear_pred.reshape(-1, mask_collinear_pred.shape[-1])
143
+ pair_mask_ptm = maski[:, :, None] * mask_pad[:, None, :] * mask_pad[:, :, None]
144
+ asym_id = feats["asym_id"].repeat_interleave(multiplicity, 0)
145
+ pair_mask_iptm = (
146
+ maski[:, :, None]
147
+ * (asym_id[:, None, :] != asym_id[:, :, None])
148
+ * mask_pad[:, None, :]
149
+ * mask_pad[:, :, None]
150
+ )
151
+ num_bins = logits.shape[-1]
152
+ bin_width = 32.0 / num_bins
153
+ end = 32.0
154
+ pae_value = torch.arange(
155
+ start=0.5 * bin_width, end=end, step=bin_width, device=logits.device
156
+ ).unsqueeze(0)
157
+ N_res = mask_pad.sum(dim=-1, keepdim=True)
158
+ tm_value = tm_function(pae_value, N_res).unsqueeze(1).unsqueeze(2)
159
+ probs = nn.functional.softmax(logits, dim=-1)
160
+ tm_expected_value = torch.sum(
161
+ probs * tm_value,
162
+ dim=-1,
163
+ ) # shape (B, N, N)
164
+ ptm = torch.max(
165
+ torch.sum(tm_expected_value * pair_mask_ptm, dim=-1)
166
+ / (torch.sum(pair_mask_ptm, dim=-1) + 1e-5),
167
+ dim=1,
168
+ ).values
169
+ iptm = torch.max(
170
+ torch.sum(tm_expected_value * pair_mask_iptm, dim=-1)
171
+ / (torch.sum(pair_mask_iptm, dim=-1) + 1e-5),
172
+ dim=1,
173
+ ).values
174
+
175
+ # compute ligand and protein iPTM
176
+ token_type = feats["mol_type"]
177
+ token_type = token_type.repeat_interleave(multiplicity, 0)
178
+ is_ligand_token = (token_type == const.chain_type_ids["NONPOLYMER"]).float()
179
+ is_protein_token = (token_type == const.chain_type_ids["PROTEIN"]).float()
180
+
181
+ ligand_iptm_mask = (
182
+ maski[:, :, None]
183
+ * (asym_id[:, None, :] != asym_id[:, :, None])
184
+ * mask_pad[:, None, :]
185
+ * mask_pad[:, :, None]
186
+ * (
187
+ (is_ligand_token[:, :, None] * is_protein_token[:, None, :])
188
+ + (is_protein_token[:, :, None] * is_ligand_token[:, None, :])
189
+ )
190
+ )
191
+ protein_ipmt_mask = (
192
+ maski[:, :, None]
193
+ * (asym_id[:, None, :] != asym_id[:, :, None])
194
+ * mask_pad[:, None, :]
195
+ * mask_pad[:, :, None]
196
+ * (is_protein_token[:, :, None] * is_protein_token[:, None, :])
197
+ )
198
+
199
+ ligand_iptm = torch.max(
200
+ torch.sum(tm_expected_value * ligand_iptm_mask, dim=-1)
201
+ / (torch.sum(ligand_iptm_mask, dim=-1) + 1e-5),
202
+ dim=1,
203
+ ).values
204
+ protein_iptm = torch.max(
205
+ torch.sum(tm_expected_value * protein_ipmt_mask, dim=-1)
206
+ / (torch.sum(protein_ipmt_mask, dim=-1) + 1e-5),
207
+ dim=1,
208
+ ).values
209
+
210
+ # Compute pair chain ipTM
211
+ chain_pair_iptm = {}
212
+ asym_ids_list = torch.unique(asym_id).tolist()
213
+ for idx1 in asym_ids_list:
214
+ chain_iptm = {}
215
+ for idx2 in asym_ids_list:
216
+ mask_pair_chain = (
217
+ maski[:, :, None]
218
+ * (asym_id[:, None, :] == idx1)
219
+ * (asym_id[:, :, None] == idx2)
220
+ * mask_pad[:, None, :]
221
+ * mask_pad[:, :, None]
222
+ )
223
+
224
+ chain_iptm[idx2] = torch.max(
225
+ torch.sum(tm_expected_value * mask_pair_chain, dim=-1)
226
+ / (torch.sum(mask_pair_chain, dim=-1) + 1e-5),
227
+ dim=1,
228
+ ).values
229
+ chain_pair_iptm[idx1] = chain_iptm
230
+
231
+ return ptm, iptm, ligand_iptm, protein_iptm, chain_pair_iptm
vb_layers_dropout.py CHANGED
@@ -1,34 +1,34 @@
1
- import torch
2
- from torch import Tensor
3
-
4
-
5
- def get_dropout_mask(
6
- dropout: float,
7
- z: Tensor,
8
- training: bool,
9
- columnwise: bool = False,
10
- ) -> Tensor:
11
- """Get the dropout mask.
12
-
13
- Parameters
14
- ----------
15
- dropout : float
16
- The dropout rate
17
- z : torch.Tensor
18
- The tensor to apply dropout to
19
- training : bool
20
- Whether the model is in training mode
21
- columnwise : bool, optional
22
- Whether to apply dropout columnwise
23
-
24
- Returns
25
- -------
26
- torch.Tensor
27
- The dropout mask
28
-
29
- """
30
- dropout = dropout * training
31
- v = z[:, 0:1, :, 0:1] if columnwise else z[:, :, 0:1, 0:1]
32
- d = torch.rand(v.shape, dtype=torch.float32, device=v.device) >= dropout
33
- d = d * 1.0 / (1.0 - dropout)
34
- return d
 
1
+ import torch
2
+ from torch import Tensor
3
+
4
+
5
+ def get_dropout_mask(
6
+ dropout: float,
7
+ z: Tensor,
8
+ training: bool,
9
+ columnwise: bool = False,
10
+ ) -> Tensor:
11
+ """Get the dropout mask.
12
+
13
+ Parameters
14
+ ----------
15
+ dropout : float
16
+ The dropout rate
17
+ z : torch.Tensor
18
+ The tensor to apply dropout to
19
+ training : bool
20
+ Whether the model is in training mode
21
+ columnwise : bool, optional
22
+ Whether to apply dropout columnwise
23
+
24
+ Returns
25
+ -------
26
+ torch.Tensor
27
+ The dropout mask
28
+
29
+ """
30
+ dropout = dropout * training
31
+ v = z[:, 0:1, :, 0:1] if columnwise else z[:, :, 0:1, 0:1]
32
+ d = torch.rand(v.shape, dtype=torch.float32, device=v.device) >= dropout
33
+ d = d * 1.0 / (1.0 - dropout)
34
+ return d
vb_layers_initialize.py CHANGED
@@ -1,86 +1,86 @@
1
- """Utility functions for initializing weights and biases."""
2
-
3
- # Copyright 2021 AlQuraishi Laboratory
4
- # Copyright 2021 DeepMind Technologies Limited
5
- #
6
- # Licensed under the Apache License, Version 2.0 (the "License");
7
- # you may not use this file except in compliance with the License.
8
- # You may obtain a copy of the License at
9
- #
10
- # http://www.apache.org/licenses/LICENSE-2.0
11
- #
12
- # Unless required by applicable law or agreed to in writing, software
13
- # distributed under the License is distributed on an "AS IS" BASIS,
14
- # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15
- # See the License for the specific language governing permissions and
16
- # limitations under the License.
17
-
18
- import math
19
-
20
- import torch
21
-
22
-
23
- def _calculate_fan(linear_weight_shape, fan="fan_in"):
24
- fan_out, fan_in = linear_weight_shape
25
-
26
- if fan == "fan_in":
27
- f = fan_in
28
- elif fan == "fan_out":
29
- f = fan_out
30
- elif fan == "fan_avg":
31
- f = (fan_in + fan_out) / 2
32
- else:
33
- raise ValueError("Invalid fan option")
34
-
35
- return f
36
-
37
-
38
- def trunc_normal_init_(weights, scale=1.0, fan="fan_in"):
39
- shape = weights.shape
40
- f = _calculate_fan(shape, fan)
41
- scale = scale / max(1, f)
42
- std = math.sqrt(scale)
43
- with torch.no_grad():
44
- torch.nn.init.trunc_normal_(weights, mean=0.0, std=std, a=-2 * std, b=2 * std)
45
-
46
-
47
- def lecun_normal_init_(weights):
48
- trunc_normal_init_(weights, scale=1.0)
49
-
50
-
51
- def he_normal_init_(weights):
52
- trunc_normal_init_(weights, scale=2.0)
53
-
54
-
55
- def glorot_uniform_init_(weights):
56
- torch.nn.init.xavier_uniform_(weights, gain=1)
57
-
58
-
59
- def final_init_(weights):
60
- with torch.no_grad():
61
- weights.fill_(0.0)
62
-
63
-
64
- def gating_init_(weights):
65
- with torch.no_grad():
66
- weights.fill_(0.0)
67
-
68
-
69
- def bias_init_zero_(bias):
70
- with torch.no_grad():
71
- bias.fill_(0.0)
72
-
73
-
74
- def bias_init_one_(bias):
75
- with torch.no_grad():
76
- bias.fill_(1.0)
77
-
78
-
79
- def normal_init_(weights):
80
- torch.nn.init.kaiming_normal_(weights, nonlinearity="linear")
81
-
82
-
83
- def ipa_point_weights_init_(weights):
84
- with torch.no_grad():
85
- softplus_inverse_1 = 0.541324854612918
86
- weights.fill_(softplus_inverse_1)
 
1
+ """Utility functions for initializing weights and biases."""
2
+
3
+ # Copyright 2021 AlQuraishi Laboratory
4
+ # Copyright 2021 DeepMind Technologies Limited
5
+ #
6
+ # Licensed under the Apache License, Version 2.0 (the "License");
7
+ # you may not use this file except in compliance with the License.
8
+ # You may obtain a copy of the License at
9
+ #
10
+ # http://www.apache.org/licenses/LICENSE-2.0
11
+ #
12
+ # Unless required by applicable law or agreed to in writing, software
13
+ # distributed under the License is distributed on an "AS IS" BASIS,
14
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15
+ # See the License for the specific language governing permissions and
16
+ # limitations under the License.
17
+
18
+ import math
19
+
20
+ import torch
21
+
22
+
23
+ def _calculate_fan(linear_weight_shape, fan="fan_in"):
24
+ fan_out, fan_in = linear_weight_shape
25
+
26
+ if fan == "fan_in":
27
+ f = fan_in
28
+ elif fan == "fan_out":
29
+ f = fan_out
30
+ elif fan == "fan_avg":
31
+ f = (fan_in + fan_out) / 2
32
+ else:
33
+ raise ValueError("Invalid fan option")
34
+
35
+ return f
36
+
37
+
38
+ def trunc_normal_init_(weights, scale=1.0, fan="fan_in"):
39
+ shape = weights.shape
40
+ f = _calculate_fan(shape, fan)
41
+ scale = scale / max(1, f)
42
+ std = math.sqrt(scale)
43
+ with torch.no_grad():
44
+ torch.nn.init.trunc_normal_(weights, mean=0.0, std=std, a=-2 * std, b=2 * std)
45
+
46
+
47
+ def lecun_normal_init_(weights):
48
+ trunc_normal_init_(weights, scale=1.0)
49
+
50
+
51
+ def he_normal_init_(weights):
52
+ trunc_normal_init_(weights, scale=2.0)
53
+
54
+
55
+ def glorot_uniform_init_(weights):
56
+ torch.nn.init.xavier_uniform_(weights, gain=1)
57
+
58
+
59
+ def final_init_(weights):
60
+ with torch.no_grad():
61
+ weights.fill_(0.0)
62
+
63
+
64
+ def gating_init_(weights):
65
+ with torch.no_grad():
66
+ weights.fill_(0.0)
67
+
68
+
69
+ def bias_init_zero_(bias):
70
+ with torch.no_grad():
71
+ bias.fill_(0.0)
72
+
73
+
74
+ def bias_init_one_(bias):
75
+ with torch.no_grad():
76
+ bias.fill_(1.0)
77
+
78
+
79
+ def normal_init_(weights):
80
+ torch.nn.init.kaiming_normal_(weights, nonlinearity="linear")
81
+
82
+
83
+ def ipa_point_weights_init_(weights):
84
+ with torch.no_grad():
85
+ softplus_inverse_1 = 0.541324854612918
86
+ weights.fill_(softplus_inverse_1)
vb_layers_outer_product_mean.py CHANGED
@@ -1,98 +1,98 @@
1
- import torch
2
- from torch import Tensor, nn
3
-
4
- from . import vb_layers_initialize as init
5
-
6
-
7
- class OuterProductMean(nn.Module):
8
- """Outer product mean layer."""
9
-
10
- def __init__(self, c_in: int, c_hidden: int, c_out: int) -> None:
11
- """Initialize the outer product mean layer.
12
-
13
- Parameters
14
- ----------
15
- c_in : int
16
- The input dimension.
17
- c_hidden : int
18
- The hidden dimension.
19
- c_out : int
20
- The output dimension.
21
-
22
- """
23
- super().__init__()
24
- self.c_hidden = c_hidden
25
- self.norm = nn.LayerNorm(c_in)
26
- self.proj_a = nn.Linear(c_in, c_hidden, bias=False)
27
- self.proj_b = nn.Linear(c_in, c_hidden, bias=False)
28
- self.proj_o = nn.Linear(c_hidden * c_hidden, c_out)
29
- init.final_init_(self.proj_o.weight)
30
- init.final_init_(self.proj_o.bias)
31
-
32
- def forward(self, m: Tensor, mask: Tensor, chunk_size: int = None) -> Tensor:
33
- """Forward pass.
34
-
35
- Parameters
36
- ----------
37
- m : torch.Tensor
38
- The sequence tensor (B, S, N, c_in).
39
- mask : torch.Tensor
40
- The mask tensor (B, S, N).
41
-
42
- Returns
43
- -------
44
- torch.Tensor
45
- The output tensor (B, N, N, c_out).
46
-
47
- """
48
- # Expand mask
49
- mask = mask.unsqueeze(-1).to(m)
50
-
51
- # Compute projections
52
- m = self.norm(m)
53
- a = self.proj_a(m) * mask
54
- b = self.proj_b(m) * mask
55
-
56
- # Compute outer product mean
57
- if chunk_size is not None and not self.training:
58
- # Compute pairwise mask
59
- for i in range(0, mask.shape[1], 64):
60
- if i == 0:
61
- num_mask = (
62
- mask[:, i : i + 64, None, :] * mask[:, i : i + 64, :, None]
63
- ).sum(1)
64
- else:
65
- num_mask += (
66
- mask[:, i : i + 64, None, :] * mask[:, i : i + 64, :, None]
67
- ).sum(1)
68
- num_mask = num_mask.clamp(min=1)
69
-
70
- # Compute squentially in chunks
71
- for i in range(0, self.c_hidden, chunk_size):
72
- a_chunk = a[:, :, :, i : i + chunk_size]
73
- sliced_weight_proj_o = self.proj_o.weight[
74
- :, i * self.c_hidden : (i + chunk_size) * self.c_hidden
75
- ]
76
-
77
- z = torch.einsum("bsic,bsjd->bijcd", a_chunk, b)
78
- z = z.reshape(*z.shape[:3], -1)
79
- z = z / num_mask
80
-
81
- # Project to output
82
- if i == 0:
83
- z_out = z.to(m) @ sliced_weight_proj_o.T
84
- else:
85
- z_out = z_out + z.to(m) @ sliced_weight_proj_o.T
86
-
87
- z_out = z_out + self.proj_o.bias # add bias
88
- return z_out
89
- else:
90
- mask = mask[:, :, None, :] * mask[:, :, :, None]
91
- num_mask = mask.sum(1).clamp(min=1)
92
- z = torch.einsum("bsic,bsjd->bijcd", a.float(), b.float())
93
- z = z.reshape(*z.shape[:3], -1)
94
- z = z / num_mask
95
-
96
- # Project to output
97
- z = self.proj_o(z.to(m))
98
- return z
 
1
+ import torch
2
+ from torch import Tensor, nn
3
+
4
+ from . import vb_layers_initialize as init
5
+
6
+
7
+ class OuterProductMean(nn.Module):
8
+ """Outer product mean layer."""
9
+
10
+ def __init__(self, c_in: int, c_hidden: int, c_out: int) -> None:
11
+ """Initialize the outer product mean layer.
12
+
13
+ Parameters
14
+ ----------
15
+ c_in : int
16
+ The input dimension.
17
+ c_hidden : int
18
+ The hidden dimension.
19
+ c_out : int
20
+ The output dimension.
21
+
22
+ """
23
+ super().__init__()
24
+ self.c_hidden = c_hidden
25
+ self.norm = nn.LayerNorm(c_in)
26
+ self.proj_a = nn.Linear(c_in, c_hidden, bias=False)
27
+ self.proj_b = nn.Linear(c_in, c_hidden, bias=False)
28
+ self.proj_o = nn.Linear(c_hidden * c_hidden, c_out)
29
+ init.final_init_(self.proj_o.weight)
30
+ init.final_init_(self.proj_o.bias)
31
+
32
+ def forward(self, m: Tensor, mask: Tensor, chunk_size: int = None) -> Tensor:
33
+ """Forward pass.
34
+
35
+ Parameters
36
+ ----------
37
+ m : torch.Tensor
38
+ The sequence tensor (B, S, N, c_in).
39
+ mask : torch.Tensor
40
+ The mask tensor (B, S, N).
41
+
42
+ Returns
43
+ -------
44
+ torch.Tensor
45
+ The output tensor (B, N, N, c_out).
46
+
47
+ """
48
+ # Expand mask
49
+ mask = mask.unsqueeze(-1).to(m)
50
+
51
+ # Compute projections
52
+ m = self.norm(m)
53
+ a = self.proj_a(m) * mask
54
+ b = self.proj_b(m) * mask
55
+
56
+ # Compute outer product mean
57
+ if chunk_size is not None and not self.training:
58
+ # Compute pairwise mask
59
+ for i in range(0, mask.shape[1], 64):
60
+ if i == 0:
61
+ num_mask = (
62
+ mask[:, i : i + 64, None, :] * mask[:, i : i + 64, :, None]
63
+ ).sum(1)
64
+ else:
65
+ num_mask += (
66
+ mask[:, i : i + 64, None, :] * mask[:, i : i + 64, :, None]
67
+ ).sum(1)
68
+ num_mask = num_mask.clamp(min=1)
69
+
70
+ # Compute squentially in chunks
71
+ for i in range(0, self.c_hidden, chunk_size):
72
+ a_chunk = a[:, :, :, i : i + chunk_size]
73
+ sliced_weight_proj_o = self.proj_o.weight[
74
+ :, i * self.c_hidden : (i + chunk_size) * self.c_hidden
75
+ ]
76
+
77
+ z = torch.einsum("bsic,bsjd->bijcd", a_chunk, b)
78
+ z = z.reshape(*z.shape[:3], -1)
79
+ z = z / num_mask
80
+
81
+ # Project to output
82
+ if i == 0:
83
+ z_out = z.to(m) @ sliced_weight_proj_o.T
84
+ else:
85
+ z_out = z_out + z.to(m) @ sliced_weight_proj_o.T
86
+
87
+ z_out = z_out + self.proj_o.bias # add bias
88
+ return z_out
89
+ else:
90
+ mask = mask[:, :, None, :] * mask[:, :, :, None]
91
+ num_mask = mask.sum(1).clamp(min=1)
92
+ z = torch.einsum("bsic,bsjd->bijcd", a.float(), b.float())
93
+ z = z.reshape(*z.shape[:3], -1)
94
+ z = z / num_mask
95
+
96
+ # Project to output
97
+ z = self.proj_o(z.to(m))
98
+ return z
vb_layers_pair_averaging.py CHANGED
@@ -1,135 +1,135 @@
1
- import torch
2
- from torch import Tensor, nn
3
-
4
- from . import vb_layers_initialize as init
5
-
6
-
7
- class PairWeightedAveraging(nn.Module):
8
- """Pair weighted averaging layer."""
9
-
10
- def __init__(
11
- self,
12
- c_m: int,
13
- c_z: int,
14
- c_h: int,
15
- num_heads: int,
16
- inf: float = 1e6,
17
- ) -> None:
18
- """Initialize the pair weighted averaging layer.
19
-
20
- Parameters
21
- ----------
22
- c_m: int
23
- The dimension of the input sequence.
24
- c_z: int
25
- The dimension of the input pairwise tensor.
26
- c_h: int
27
- The dimension of the hidden.
28
- num_heads: int
29
- The number of heads.
30
- inf: float
31
- The value to use for masking, default 1e6.
32
-
33
- """
34
- super().__init__()
35
- self.c_m = c_m
36
- self.c_z = c_z
37
- self.c_h = c_h
38
- self.num_heads = num_heads
39
- self.inf = inf
40
-
41
- self.norm_m = nn.LayerNorm(c_m)
42
- self.norm_z = nn.LayerNorm(c_z)
43
-
44
- self.proj_m = nn.Linear(c_m, c_h * num_heads, bias=False)
45
- self.proj_g = nn.Linear(c_m, c_h * num_heads, bias=False)
46
- self.proj_z = nn.Linear(c_z, num_heads, bias=False)
47
- self.proj_o = nn.Linear(c_h * num_heads, c_m, bias=False)
48
- init.final_init_(self.proj_o.weight)
49
-
50
- def forward(
51
- self, m: Tensor, z: Tensor, mask: Tensor, chunk_heads: False = bool
52
- ) -> Tensor:
53
- """Forward pass.
54
-
55
- Parameters
56
- ----------
57
- m : torch.Tensor
58
- The input sequence tensor (B, S, N, D)
59
- z : torch.Tensor
60
- The input pairwise tensor (B, N, N, D)
61
- mask : torch.Tensor
62
- The pairwise mask tensor (B, N, N)
63
-
64
- Returns
65
- -------
66
- torch.Tensor
67
- The output sequence tensor (B, S, N, D)
68
-
69
- """
70
- # Compute layer norms
71
- m = self.norm_m(m)
72
- z = self.norm_z(z)
73
-
74
- if chunk_heads and not self.training:
75
- # Compute heads sequentially
76
- o_chunks = []
77
- for head_idx in range(self.num_heads):
78
- sliced_weight_proj_m = self.proj_m.weight[
79
- head_idx * self.c_h : (head_idx + 1) * self.c_h, :
80
- ]
81
- sliced_weight_proj_g = self.proj_g.weight[
82
- head_idx * self.c_h : (head_idx + 1) * self.c_h, :
83
- ]
84
- sliced_weight_proj_z = self.proj_z.weight[head_idx : (head_idx + 1), :]
85
- sliced_weight_proj_o = self.proj_o.weight[
86
- :, head_idx * self.c_h : (head_idx + 1) * self.c_h
87
- ]
88
-
89
- # Project input tensors
90
- v: Tensor = m @ sliced_weight_proj_m.T
91
- v = v.reshape(*v.shape[:3], 1, self.c_h)
92
- v = v.permute(0, 3, 1, 2, 4)
93
-
94
- # Compute weights
95
- b: Tensor = z @ sliced_weight_proj_z.T
96
- b = b.permute(0, 3, 1, 2)
97
- b = b + (1 - mask[:, None]) * -self.inf
98
- w = torch.softmax(b, dim=-1)
99
-
100
- # Compute gating
101
- g: Tensor = m @ sliced_weight_proj_g.T
102
- g = g.sigmoid()
103
-
104
- # Compute output
105
- o = torch.einsum("bhij,bhsjd->bhsid", w, v)
106
- o = o.permute(0, 2, 3, 1, 4)
107
- o = o.reshape(*o.shape[:3], 1 * self.c_h)
108
- o_chunks = g * o
109
- if head_idx == 0:
110
- o_out = o_chunks @ sliced_weight_proj_o.T
111
- else:
112
- o_out += o_chunks @ sliced_weight_proj_o.T
113
- return o_out
114
- else:
115
- # Project input tensors
116
- v: Tensor = self.proj_m(m)
117
- v = v.reshape(*v.shape[:3], self.num_heads, self.c_h)
118
- v = v.permute(0, 3, 1, 2, 4)
119
-
120
- # Compute weights
121
- b: Tensor = self.proj_z(z)
122
- b = b.permute(0, 3, 1, 2)
123
- b = b + (1 - mask[:, None]) * -self.inf
124
- w = torch.softmax(b, dim=-1)
125
-
126
- # Compute gating
127
- g: Tensor = self.proj_g(m)
128
- g = g.sigmoid()
129
-
130
- # Compute output
131
- o = torch.einsum("bhij,bhsjd->bhsid", w, v)
132
- o = o.permute(0, 2, 3, 1, 4)
133
- o = o.reshape(*o.shape[:3], self.num_heads * self.c_h)
134
- o = self.proj_o(g * o)
135
- return o
 
1
+ import torch
2
+ from torch import Tensor, nn
3
+
4
+ from . import vb_layers_initialize as init
5
+
6
+
7
+ class PairWeightedAveraging(nn.Module):
8
+ """Pair weighted averaging layer."""
9
+
10
+ def __init__(
11
+ self,
12
+ c_m: int,
13
+ c_z: int,
14
+ c_h: int,
15
+ num_heads: int,
16
+ inf: float = 1e6,
17
+ ) -> None:
18
+ """Initialize the pair weighted averaging layer.
19
+
20
+ Parameters
21
+ ----------
22
+ c_m: int
23
+ The dimension of the input sequence.
24
+ c_z: int
25
+ The dimension of the input pairwise tensor.
26
+ c_h: int
27
+ The dimension of the hidden.
28
+ num_heads: int
29
+ The number of heads.
30
+ inf: float
31
+ The value to use for masking, default 1e6.
32
+
33
+ """
34
+ super().__init__()
35
+ self.c_m = c_m
36
+ self.c_z = c_z
37
+ self.c_h = c_h
38
+ self.num_heads = num_heads
39
+ self.inf = inf
40
+
41
+ self.norm_m = nn.LayerNorm(c_m)
42
+ self.norm_z = nn.LayerNorm(c_z)
43
+
44
+ self.proj_m = nn.Linear(c_m, c_h * num_heads, bias=False)
45
+ self.proj_g = nn.Linear(c_m, c_h * num_heads, bias=False)
46
+ self.proj_z = nn.Linear(c_z, num_heads, bias=False)
47
+ self.proj_o = nn.Linear(c_h * num_heads, c_m, bias=False)
48
+ init.final_init_(self.proj_o.weight)
49
+
50
+ def forward(
51
+ self, m: Tensor, z: Tensor, mask: Tensor, chunk_heads: False = bool
52
+ ) -> Tensor:
53
+ """Forward pass.
54
+
55
+ Parameters
56
+ ----------
57
+ m : torch.Tensor
58
+ The input sequence tensor (B, S, N, D)
59
+ z : torch.Tensor
60
+ The input pairwise tensor (B, N, N, D)
61
+ mask : torch.Tensor
62
+ The pairwise mask tensor (B, N, N)
63
+
64
+ Returns
65
+ -------
66
+ torch.Tensor
67
+ The output sequence tensor (B, S, N, D)
68
+
69
+ """
70
+ # Compute layer norms
71
+ m = self.norm_m(m)
72
+ z = self.norm_z(z)
73
+
74
+ if chunk_heads and not self.training:
75
+ # Compute heads sequentially
76
+ o_chunks = []
77
+ for head_idx in range(self.num_heads):
78
+ sliced_weight_proj_m = self.proj_m.weight[
79
+ head_idx * self.c_h : (head_idx + 1) * self.c_h, :
80
+ ]
81
+ sliced_weight_proj_g = self.proj_g.weight[
82
+ head_idx * self.c_h : (head_idx + 1) * self.c_h, :
83
+ ]
84
+ sliced_weight_proj_z = self.proj_z.weight[head_idx : (head_idx + 1), :]
85
+ sliced_weight_proj_o = self.proj_o.weight[
86
+ :, head_idx * self.c_h : (head_idx + 1) * self.c_h
87
+ ]
88
+
89
+ # Project input tensors
90
+ v: Tensor = m @ sliced_weight_proj_m.T
91
+ v = v.reshape(*v.shape[:3], 1, self.c_h)
92
+ v = v.permute(0, 3, 1, 2, 4)
93
+
94
+ # Compute weights
95
+ b: Tensor = z @ sliced_weight_proj_z.T
96
+ b = b.permute(0, 3, 1, 2)
97
+ b = b + (1 - mask[:, None]) * -self.inf
98
+ w = torch.softmax(b, dim=-1)
99
+
100
+ # Compute gating
101
+ g: Tensor = m @ sliced_weight_proj_g.T
102
+ g = g.sigmoid()
103
+
104
+ # Compute output
105
+ o = torch.einsum("bhij,bhsjd->bhsid", w, v)
106
+ o = o.permute(0, 2, 3, 1, 4)
107
+ o = o.reshape(*o.shape[:3], 1 * self.c_h)
108
+ o_chunks = g * o
109
+ if head_idx == 0:
110
+ o_out = o_chunks @ sliced_weight_proj_o.T
111
+ else:
112
+ o_out += o_chunks @ sliced_weight_proj_o.T
113
+ return o_out
114
+ else:
115
+ # Project input tensors
116
+ v: Tensor = self.proj_m(m)
117
+ v = v.reshape(*v.shape[:3], self.num_heads, self.c_h)
118
+ v = v.permute(0, 3, 1, 2, 4)
119
+
120
+ # Compute weights
121
+ b: Tensor = self.proj_z(z)
122
+ b = b.permute(0, 3, 1, 2)
123
+ b = b + (1 - mask[:, None]) * -self.inf
124
+ w = torch.softmax(b, dim=-1)
125
+
126
+ # Compute gating
127
+ g: Tensor = self.proj_g(m)
128
+ g = g.sigmoid()
129
+
130
+ # Compute output
131
+ o = torch.einsum("bhij,bhsjd->bhsid", w, v)
132
+ o = o.permute(0, 2, 3, 1, 4)
133
+ o = o.reshape(*o.shape[:3], self.num_heads * self.c_h)
134
+ o = self.proj_o(g * o)
135
+ return o
vb_layers_pairformer.py CHANGED
@@ -1,335 +1,335 @@
1
- from typing import Optional
2
-
3
- import torch
4
- from torch import Tensor, nn
5
-
6
- from . import vb_const as const
7
- from .vb_layers_attention import AttentionPairBias
8
- from .vb_layers_attentionv2 import AttentionPairBias as AttentionPairBiasV2
9
- from .vb_layers_dropout import get_dropout_mask
10
- from .vb_layers_transition import Transition
11
- from .vb_tri_attn_attention import (
12
- TriangleAttentionEndingNode,
13
- TriangleAttentionStartingNode,
14
- )
15
- from .vb_layers_triangular_mult import (
16
- TriangleMultiplicationIncoming,
17
- TriangleMultiplicationOutgoing,
18
- )
19
-
20
-
21
- class PairformerLayer(nn.Module):
22
- """Pairformer module."""
23
-
24
- def __init__(
25
- self,
26
- token_s: int,
27
- token_z: int,
28
- num_heads: int = 16,
29
- dropout: float = 0.25,
30
- pairwise_head_width: int = 32,
31
- pairwise_num_heads: int = 4,
32
- post_layer_norm: bool = False,
33
- v2: bool = False,
34
- ) -> None:
35
- super().__init__()
36
- self.token_z = token_z
37
- self.dropout = dropout
38
- self.num_heads = num_heads
39
- self.post_layer_norm = post_layer_norm
40
-
41
- self.pre_norm_s = nn.LayerNorm(token_s)
42
- if v2:
43
- self.attention = AttentionPairBiasV2(token_s, token_z, num_heads)
44
- else:
45
- self.attention = AttentionPairBias(token_s, token_z, num_heads)
46
-
47
- self.tri_mul_out = TriangleMultiplicationOutgoing(token_z)
48
- self.tri_mul_in = TriangleMultiplicationIncoming(token_z)
49
-
50
- self.tri_att_start = TriangleAttentionStartingNode(
51
- token_z, pairwise_head_width, pairwise_num_heads, inf=1e9
52
- )
53
- self.tri_att_end = TriangleAttentionEndingNode(
54
- token_z, pairwise_head_width, pairwise_num_heads, inf=1e9
55
- )
56
-
57
- self.transition_s = Transition(token_s, token_s * 4)
58
- self.transition_z = Transition(token_z, token_z * 4)
59
-
60
- self.s_post_norm = (
61
- nn.LayerNorm(token_s) if self.post_layer_norm else nn.Identity()
62
- )
63
-
64
- def forward(
65
- self,
66
- s: Tensor,
67
- z: Tensor,
68
- mask: Tensor,
69
- pair_mask: Tensor,
70
- chunk_size_tri_attn: Optional[int] = None,
71
- use_kernels: bool = False,
72
- use_cuequiv_mul: bool = False,
73
- use_cuequiv_attn: bool = False,
74
- ) -> tuple[Tensor, Tensor]:
75
- # Compute pairwise stack
76
- dropout = get_dropout_mask(self.dropout, z, self.training)
77
- z = z + dropout * self.tri_mul_out(
78
- z, mask=pair_mask, use_kernels=use_cuequiv_mul or use_kernels
79
- )
80
-
81
- dropout = get_dropout_mask(self.dropout, z, self.training)
82
- z = z + dropout * self.tri_mul_in(
83
- z, mask=pair_mask, use_kernels=use_cuequiv_mul or use_kernels
84
- )
85
-
86
- dropout = get_dropout_mask(self.dropout, z, self.training)
87
- z = z + dropout * self.tri_att_start(
88
- z,
89
- mask=pair_mask,
90
- chunk_size=chunk_size_tri_attn,
91
- use_kernels=use_cuequiv_attn or use_kernels,
92
- )
93
-
94
- dropout = get_dropout_mask(self.dropout, z, self.training, columnwise=True)
95
- z = z + dropout * self.tri_att_end(
96
- z,
97
- mask=pair_mask,
98
- chunk_size=chunk_size_tri_attn,
99
- use_kernels=use_cuequiv_attn or use_kernels,
100
- )
101
-
102
- z = z + self.transition_z(z)
103
-
104
- # Compute sequence stack
105
- with torch.autocast("cuda", enabled=False):
106
- s_normed = self.pre_norm_s(s.float())
107
- s = s.float() + self.attention(
108
- s=s_normed, z=z.float(), mask=mask.float(), k_in=s_normed
109
- )
110
- s = s + self.transition_s(s)
111
- s = self.s_post_norm(s)
112
-
113
- return s, z
114
-
115
-
116
- class PairformerModule(nn.Module):
117
- """Pairformer module."""
118
-
119
- def __init__(
120
- self,
121
- token_s: int,
122
- token_z: int,
123
- num_blocks: int,
124
- num_heads: int = 16,
125
- dropout: float = 0.25,
126
- pairwise_head_width: int = 32,
127
- pairwise_num_heads: int = 4,
128
- post_layer_norm: bool = False,
129
- activation_checkpointing: bool = False,
130
- v2: bool = False,
131
- **kwargs,
132
- ) -> None:
133
- super().__init__()
134
- self.token_z = token_z
135
- self.num_blocks = num_blocks
136
- self.dropout = dropout
137
- self.num_heads = num_heads
138
- self.post_layer_norm = post_layer_norm
139
- self.activation_checkpointing = activation_checkpointing
140
-
141
- self.layers = nn.ModuleList()
142
- for _ in range(num_blocks):
143
- self.layers.append(
144
- PairformerLayer(
145
- token_s,
146
- token_z,
147
- num_heads,
148
- dropout,
149
- pairwise_head_width,
150
- pairwise_num_heads,
151
- post_layer_norm,
152
- v2,
153
- ),
154
- )
155
-
156
- def forward(
157
- self,
158
- s: Tensor,
159
- z: Tensor,
160
- mask: Tensor,
161
- pair_mask: Tensor,
162
- use_kernels: bool = False,
163
- ) -> tuple[Tensor, Tensor]:
164
- """Perform the forward pass.
165
-
166
- Parameters
167
- ----------
168
- s : Tensor
169
- The sequence stack.
170
- z : Tensor
171
- The pairwise stack.
172
- mask : Tensor
173
- The mask.
174
- pair_mask : Tensor
175
- The pairwise mask.
176
- use_kernels : bool
177
- Whether to use kernels.
178
-
179
- """
180
- if not self.training:
181
- if z.shape[1] > const.chunk_size_threshold:
182
- chunk_size_tri_attn = 128
183
- else:
184
- chunk_size_tri_attn = 512
185
- else:
186
- chunk_size_tri_attn = None
187
-
188
- for layer in self.layers:
189
- if self.activation_checkpointing and self.training:
190
- s, z = torch.utils.checkpoint.checkpoint(
191
- layer,
192
- s,
193
- z,
194
- mask,
195
- pair_mask,
196
- chunk_size_tri_attn,
197
- use_kernels,
198
- )
199
- else:
200
- s, z = layer(s, z, mask, pair_mask, chunk_size_tri_attn, use_kernels)
201
- return s, z
202
-
203
-
204
- class PairformerNoSeqLayer(nn.Module):
205
- """Pairformer module without sequence track."""
206
-
207
- def __init__(
208
- self,
209
- token_z: int,
210
- dropout: float = 0.25,
211
- pairwise_head_width: int = 32,
212
- pairwise_num_heads: int = 4,
213
- post_layer_norm: bool = False,
214
- ) -> None:
215
- super().__init__()
216
- self.token_z = token_z
217
- self.dropout = dropout
218
- self.post_layer_norm = post_layer_norm
219
-
220
- self.tri_mul_out = TriangleMultiplicationOutgoing(token_z)
221
- self.tri_mul_in = TriangleMultiplicationIncoming(token_z)
222
-
223
- self.tri_att_start = TriangleAttentionStartingNode(
224
- token_z, pairwise_head_width, pairwise_num_heads, inf=1e9
225
- )
226
- self.tri_att_end = TriangleAttentionEndingNode(
227
- token_z, pairwise_head_width, pairwise_num_heads, inf=1e9
228
- )
229
-
230
- self.transition_z = Transition(token_z, token_z * 4)
231
-
232
- def forward(
233
- self,
234
- z: Tensor,
235
- pair_mask: Tensor,
236
- chunk_size_tri_attn: Optional[int] = None,
237
- use_kernels: bool = False,
238
- use_cuequiv_mul: bool = False,
239
- use_cuequiv_attn: bool = False,
240
- ) -> Tensor:
241
- # Compute pairwise stack
242
- dropout = get_dropout_mask(self.dropout, z, self.training)
243
- z = z + dropout * self.tri_mul_out(
244
- z, mask=pair_mask, use_kernels=use_cuequiv_mul or use_kernels
245
- )
246
-
247
- dropout = get_dropout_mask(self.dropout, z, self.training)
248
- z = z + dropout * self.tri_mul_in(
249
- z, mask=pair_mask, use_kernels=use_cuequiv_mul or use_kernels
250
- )
251
-
252
- dropout = get_dropout_mask(self.dropout, z, self.training)
253
- z = z + dropout * self.tri_att_start(
254
- z,
255
- mask=pair_mask,
256
- chunk_size=chunk_size_tri_attn,
257
- use_kernels=use_cuequiv_attn or use_kernels,
258
- )
259
-
260
- dropout = get_dropout_mask(self.dropout, z, self.training, columnwise=True)
261
- z = z + dropout * self.tri_att_end(
262
- z,
263
- mask=pair_mask,
264
- chunk_size=chunk_size_tri_attn,
265
- use_kernels=use_cuequiv_attn or use_kernels,
266
- )
267
-
268
- z = z + self.transition_z(z)
269
- return z
270
-
271
-
272
- class PairformerNoSeqModule(nn.Module):
273
- """Pairformer module without sequence track."""
274
-
275
- def __init__(
276
- self,
277
- token_z: int,
278
- num_blocks: int,
279
- dropout: float = 0.25,
280
- pairwise_head_width: int = 32,
281
- pairwise_num_heads: int = 4,
282
- post_layer_norm: bool = False,
283
- activation_checkpointing: bool = False,
284
- **kwargs,
285
- ) -> None:
286
- super().__init__()
287
- self.token_z = token_z
288
- self.num_blocks = num_blocks
289
- self.dropout = dropout
290
- self.post_layer_norm = post_layer_norm
291
- self.activation_checkpointing = activation_checkpointing
292
-
293
- self.layers = nn.ModuleList()
294
- for i in range(num_blocks):
295
- self.layers.append(
296
- PairformerNoSeqLayer(
297
- token_z,
298
- dropout,
299
- pairwise_head_width,
300
- pairwise_num_heads,
301
- post_layer_norm,
302
- ),
303
- )
304
-
305
- def forward(
306
- self,
307
- z: Tensor,
308
- pair_mask: Tensor,
309
- use_kernels: bool = False,
310
- ) -> Tensor:
311
- if not self.training:
312
- if z.shape[1] > const.chunk_size_threshold:
313
- chunk_size_tri_attn = 128
314
- else:
315
- chunk_size_tri_attn = 512
316
- else:
317
- chunk_size_tri_attn = None
318
-
319
- for layer in self.layers:
320
- if self.activation_checkpointing and self.training:
321
- z = torch.utils.checkpoint.checkpoint(
322
- layer,
323
- z,
324
- pair_mask,
325
- chunk_size_tri_attn,
326
- use_kernels,
327
- )
328
- else:
329
- z = layer(
330
- z,
331
- pair_mask,
332
- chunk_size_tri_attn,
333
- use_kernels,
334
- )
335
- return z
 
1
+ from typing import Optional
2
+
3
+ import torch
4
+ from torch import Tensor, nn
5
+
6
+ from . import vb_const as const
7
+ from .vb_layers_attention import AttentionPairBias
8
+ from .vb_layers_attentionv2 import AttentionPairBias as AttentionPairBiasV2
9
+ from .vb_layers_dropout import get_dropout_mask
10
+ from .vb_layers_transition import Transition
11
+ from .vb_tri_attn_attention import (
12
+ TriangleAttentionEndingNode,
13
+ TriangleAttentionStartingNode,
14
+ )
15
+ from .vb_layers_triangular_mult import (
16
+ TriangleMultiplicationIncoming,
17
+ TriangleMultiplicationOutgoing,
18
+ )
19
+
20
+
21
+ class PairformerLayer(nn.Module):
22
+ """Pairformer module."""
23
+
24
+ def __init__(
25
+ self,
26
+ token_s: int,
27
+ token_z: int,
28
+ num_heads: int = 16,
29
+ dropout: float = 0.25,
30
+ pairwise_head_width: int = 32,
31
+ pairwise_num_heads: int = 4,
32
+ post_layer_norm: bool = False,
33
+ v2: bool = False,
34
+ ) -> None:
35
+ super().__init__()
36
+ self.token_z = token_z
37
+ self.dropout = dropout
38
+ self.num_heads = num_heads
39
+ self.post_layer_norm = post_layer_norm
40
+
41
+ self.pre_norm_s = nn.LayerNorm(token_s)
42
+ if v2:
43
+ self.attention = AttentionPairBiasV2(token_s, token_z, num_heads)
44
+ else:
45
+ self.attention = AttentionPairBias(token_s, token_z, num_heads)
46
+
47
+ self.tri_mul_out = TriangleMultiplicationOutgoing(token_z)
48
+ self.tri_mul_in = TriangleMultiplicationIncoming(token_z)
49
+
50
+ self.tri_att_start = TriangleAttentionStartingNode(
51
+ token_z, pairwise_head_width, pairwise_num_heads, inf=1e9
52
+ )
53
+ self.tri_att_end = TriangleAttentionEndingNode(
54
+ token_z, pairwise_head_width, pairwise_num_heads, inf=1e9
55
+ )
56
+
57
+ self.transition_s = Transition(token_s, token_s * 4)
58
+ self.transition_z = Transition(token_z, token_z * 4)
59
+
60
+ self.s_post_norm = (
61
+ nn.LayerNorm(token_s) if self.post_layer_norm else nn.Identity()
62
+ )
63
+
64
+ def forward(
65
+ self,
66
+ s: Tensor,
67
+ z: Tensor,
68
+ mask: Tensor,
69
+ pair_mask: Tensor,
70
+ chunk_size_tri_attn: Optional[int] = None,
71
+ use_kernels: bool = False,
72
+ use_cuequiv_mul: bool = False,
73
+ use_cuequiv_attn: bool = False,
74
+ ) -> tuple[Tensor, Tensor]:
75
+ # Compute pairwise stack
76
+ dropout = get_dropout_mask(self.dropout, z, self.training)
77
+ z = z + dropout * self.tri_mul_out(
78
+ z, mask=pair_mask, use_kernels=use_cuequiv_mul or use_kernels
79
+ )
80
+
81
+ dropout = get_dropout_mask(self.dropout, z, self.training)
82
+ z = z + dropout * self.tri_mul_in(
83
+ z, mask=pair_mask, use_kernels=use_cuequiv_mul or use_kernels
84
+ )
85
+
86
+ dropout = get_dropout_mask(self.dropout, z, self.training)
87
+ z = z + dropout * self.tri_att_start(
88
+ z,
89
+ mask=pair_mask,
90
+ chunk_size=chunk_size_tri_attn,
91
+ use_kernels=use_cuequiv_attn or use_kernels,
92
+ )
93
+
94
+ dropout = get_dropout_mask(self.dropout, z, self.training, columnwise=True)
95
+ z = z + dropout * self.tri_att_end(
96
+ z,
97
+ mask=pair_mask,
98
+ chunk_size=chunk_size_tri_attn,
99
+ use_kernels=use_cuequiv_attn or use_kernels,
100
+ )
101
+
102
+ z = z + self.transition_z(z)
103
+
104
+ # Compute sequence stack
105
+ with torch.autocast("cuda", enabled=False):
106
+ s_normed = self.pre_norm_s(s.float())
107
+ s = s.float() + self.attention(
108
+ s=s_normed, z=z.float(), mask=mask.float(), k_in=s_normed
109
+ )
110
+ s = s + self.transition_s(s)
111
+ s = self.s_post_norm(s)
112
+
113
+ return s, z
114
+
115
+
116
+ class PairformerModule(nn.Module):
117
+ """Pairformer module."""
118
+
119
+ def __init__(
120
+ self,
121
+ token_s: int,
122
+ token_z: int,
123
+ num_blocks: int,
124
+ num_heads: int = 16,
125
+ dropout: float = 0.25,
126
+ pairwise_head_width: int = 32,
127
+ pairwise_num_heads: int = 4,
128
+ post_layer_norm: bool = False,
129
+ activation_checkpointing: bool = False,
130
+ v2: bool = False,
131
+ **kwargs,
132
+ ) -> None:
133
+ super().__init__()
134
+ self.token_z = token_z
135
+ self.num_blocks = num_blocks
136
+ self.dropout = dropout
137
+ self.num_heads = num_heads
138
+ self.post_layer_norm = post_layer_norm
139
+ self.activation_checkpointing = activation_checkpointing
140
+
141
+ self.layers = nn.ModuleList()
142
+ for _ in range(num_blocks):
143
+ self.layers.append(
144
+ PairformerLayer(
145
+ token_s,
146
+ token_z,
147
+ num_heads,
148
+ dropout,
149
+ pairwise_head_width,
150
+ pairwise_num_heads,
151
+ post_layer_norm,
152
+ v2,
153
+ ),
154
+ )
155
+
156
+ def forward(
157
+ self,
158
+ s: Tensor,
159
+ z: Tensor,
160
+ mask: Tensor,
161
+ pair_mask: Tensor,
162
+ use_kernels: bool = False,
163
+ ) -> tuple[Tensor, Tensor]:
164
+ """Perform the forward pass.
165
+
166
+ Parameters
167
+ ----------
168
+ s : Tensor
169
+ The sequence stack.
170
+ z : Tensor
171
+ The pairwise stack.
172
+ mask : Tensor
173
+ The mask.
174
+ pair_mask : Tensor
175
+ The pairwise mask.
176
+ use_kernels : bool
177
+ Whether to use kernels.
178
+
179
+ """
180
+ if not self.training:
181
+ if z.shape[1] > const.chunk_size_threshold:
182
+ chunk_size_tri_attn = 128
183
+ else:
184
+ chunk_size_tri_attn = 512
185
+ else:
186
+ chunk_size_tri_attn = None
187
+
188
+ for layer in self.layers:
189
+ if self.activation_checkpointing and self.training:
190
+ s, z = torch.utils.checkpoint.checkpoint(
191
+ layer,
192
+ s,
193
+ z,
194
+ mask,
195
+ pair_mask,
196
+ chunk_size_tri_attn,
197
+ use_kernels,
198
+ )
199
+ else:
200
+ s, z = layer(s, z, mask, pair_mask, chunk_size_tri_attn, use_kernels)
201
+ return s, z
202
+
203
+
204
+ class PairformerNoSeqLayer(nn.Module):
205
+ """Pairformer module without sequence track."""
206
+
207
+ def __init__(
208
+ self,
209
+ token_z: int,
210
+ dropout: float = 0.25,
211
+ pairwise_head_width: int = 32,
212
+ pairwise_num_heads: int = 4,
213
+ post_layer_norm: bool = False,
214
+ ) -> None:
215
+ super().__init__()
216
+ self.token_z = token_z
217
+ self.dropout = dropout
218
+ self.post_layer_norm = post_layer_norm
219
+
220
+ self.tri_mul_out = TriangleMultiplicationOutgoing(token_z)
221
+ self.tri_mul_in = TriangleMultiplicationIncoming(token_z)
222
+
223
+ self.tri_att_start = TriangleAttentionStartingNode(
224
+ token_z, pairwise_head_width, pairwise_num_heads, inf=1e9
225
+ )
226
+ self.tri_att_end = TriangleAttentionEndingNode(
227
+ token_z, pairwise_head_width, pairwise_num_heads, inf=1e9
228
+ )
229
+
230
+ self.transition_z = Transition(token_z, token_z * 4)
231
+
232
+ def forward(
233
+ self,
234
+ z: Tensor,
235
+ pair_mask: Tensor,
236
+ chunk_size_tri_attn: Optional[int] = None,
237
+ use_kernels: bool = False,
238
+ use_cuequiv_mul: bool = False,
239
+ use_cuequiv_attn: bool = False,
240
+ ) -> Tensor:
241
+ # Compute pairwise stack
242
+ dropout = get_dropout_mask(self.dropout, z, self.training)
243
+ z = z + dropout * self.tri_mul_out(
244
+ z, mask=pair_mask, use_kernels=use_cuequiv_mul or use_kernels
245
+ )
246
+
247
+ dropout = get_dropout_mask(self.dropout, z, self.training)
248
+ z = z + dropout * self.tri_mul_in(
249
+ z, mask=pair_mask, use_kernels=use_cuequiv_mul or use_kernels
250
+ )
251
+
252
+ dropout = get_dropout_mask(self.dropout, z, self.training)
253
+ z = z + dropout * self.tri_att_start(
254
+ z,
255
+ mask=pair_mask,
256
+ chunk_size=chunk_size_tri_attn,
257
+ use_kernels=use_cuequiv_attn or use_kernels,
258
+ )
259
+
260
+ dropout = get_dropout_mask(self.dropout, z, self.training, columnwise=True)
261
+ z = z + dropout * self.tri_att_end(
262
+ z,
263
+ mask=pair_mask,
264
+ chunk_size=chunk_size_tri_attn,
265
+ use_kernels=use_cuequiv_attn or use_kernels,
266
+ )
267
+
268
+ z = z + self.transition_z(z)
269
+ return z
270
+
271
+
272
+ class PairformerNoSeqModule(nn.Module):
273
+ """Pairformer module without sequence track."""
274
+
275
+ def __init__(
276
+ self,
277
+ token_z: int,
278
+ num_blocks: int,
279
+ dropout: float = 0.25,
280
+ pairwise_head_width: int = 32,
281
+ pairwise_num_heads: int = 4,
282
+ post_layer_norm: bool = False,
283
+ activation_checkpointing: bool = False,
284
+ **kwargs,
285
+ ) -> None:
286
+ super().__init__()
287
+ self.token_z = token_z
288
+ self.num_blocks = num_blocks
289
+ self.dropout = dropout
290
+ self.post_layer_norm = post_layer_norm
291
+ self.activation_checkpointing = activation_checkpointing
292
+
293
+ self.layers = nn.ModuleList()
294
+ for i in range(num_blocks):
295
+ self.layers.append(
296
+ PairformerNoSeqLayer(
297
+ token_z,
298
+ dropout,
299
+ pairwise_head_width,
300
+ pairwise_num_heads,
301
+ post_layer_norm,
302
+ ),
303
+ )
304
+
305
+ def forward(
306
+ self,
307
+ z: Tensor,
308
+ pair_mask: Tensor,
309
+ use_kernels: bool = False,
310
+ ) -> Tensor:
311
+ if not self.training:
312
+ if z.shape[1] > const.chunk_size_threshold:
313
+ chunk_size_tri_attn = 128
314
+ else:
315
+ chunk_size_tri_attn = 512
316
+ else:
317
+ chunk_size_tri_attn = None
318
+
319
+ for layer in self.layers:
320
+ if self.activation_checkpointing and self.training:
321
+ z = torch.utils.checkpoint.checkpoint(
322
+ layer,
323
+ z,
324
+ pair_mask,
325
+ chunk_size_tri_attn,
326
+ use_kernels,
327
+ )
328
+ else:
329
+ z = layer(
330
+ z,
331
+ pair_mask,
332
+ chunk_size_tri_attn,
333
+ use_kernels,
334
+ )
335
+ return z
vb_layers_transition.py CHANGED
@@ -1,78 +1,78 @@
1
- from typing import Optional
2
-
3
- from torch import Tensor, nn
4
-
5
- from . import vb_layers_initialize as init
6
-
7
-
8
- class Transition(nn.Module):
9
- """Perform a two-layer MLP."""
10
-
11
- def __init__(
12
- self,
13
- dim: int = 128,
14
- hidden: int = 512,
15
- out_dim: Optional[int] = None,
16
- ) -> None:
17
- """Initialize the TransitionUpdate module.
18
-
19
- Parameters
20
- ----------
21
- dim: int
22
- The dimension of the input, default 128
23
- hidden: int
24
- The dimension of the hidden, default 512
25
- out_dim: Optional[int]
26
- The dimension of the output, default None
27
-
28
- """
29
- super().__init__()
30
- if out_dim is None:
31
- out_dim = dim
32
-
33
- self.norm = nn.LayerNorm(dim, eps=1e-5)
34
- self.fc1 = nn.Linear(dim, hidden, bias=False)
35
- self.fc2 = nn.Linear(dim, hidden, bias=False)
36
- self.fc3 = nn.Linear(hidden, out_dim, bias=False)
37
- self.silu = nn.SiLU()
38
- self.hidden = hidden
39
-
40
- init.bias_init_one_(self.norm.weight)
41
- init.bias_init_zero_(self.norm.bias)
42
-
43
- init.lecun_normal_init_(self.fc1.weight)
44
- init.lecun_normal_init_(self.fc2.weight)
45
- init.final_init_(self.fc3.weight)
46
-
47
- def forward(self, x: Tensor, chunk_size: int = None) -> Tensor:
48
- """Perform a forward pass.
49
-
50
- Parameters
51
- ----------
52
- x: torch.Tensor
53
- The input data of shape (..., D)
54
-
55
- Returns
56
- -------
57
- x: torch.Tensor
58
- The output data of shape (..., D)
59
-
60
- """
61
- x = self.norm(x)
62
-
63
- if chunk_size is None or self.training:
64
- x = self.silu(self.fc1(x)) * self.fc2(x)
65
- x = self.fc3(x)
66
- return x
67
- else:
68
- # Compute in chunks
69
- for i in range(0, self.hidden, chunk_size):
70
- fc1_slice = self.fc1.weight[i : i + chunk_size, :]
71
- fc2_slice = self.fc2.weight[i : i + chunk_size, :]
72
- fc3_slice = self.fc3.weight[:, i : i + chunk_size]
73
- x_chunk = self.silu((x @ fc1_slice.T)) * (x @ fc2_slice.T)
74
- if i == 0:
75
- x_out = x_chunk @ fc3_slice.T
76
- else:
77
- x_out = x_out + x_chunk @ fc3_slice.T
78
- return x_out
 
1
+ from typing import Optional
2
+
3
+ from torch import Tensor, nn
4
+
5
+ from . import vb_layers_initialize as init
6
+
7
+
8
+ class Transition(nn.Module):
9
+ """Perform a two-layer MLP."""
10
+
11
+ def __init__(
12
+ self,
13
+ dim: int = 128,
14
+ hidden: int = 512,
15
+ out_dim: Optional[int] = None,
16
+ ) -> None:
17
+ """Initialize the TransitionUpdate module.
18
+
19
+ Parameters
20
+ ----------
21
+ dim: int
22
+ The dimension of the input, default 128
23
+ hidden: int
24
+ The dimension of the hidden, default 512
25
+ out_dim: Optional[int]
26
+ The dimension of the output, default None
27
+
28
+ """
29
+ super().__init__()
30
+ if out_dim is None:
31
+ out_dim = dim
32
+
33
+ self.norm = nn.LayerNorm(dim, eps=1e-5)
34
+ self.fc1 = nn.Linear(dim, hidden, bias=False)
35
+ self.fc2 = nn.Linear(dim, hidden, bias=False)
36
+ self.fc3 = nn.Linear(hidden, out_dim, bias=False)
37
+ self.silu = nn.SiLU()
38
+ self.hidden = hidden
39
+
40
+ init.bias_init_one_(self.norm.weight)
41
+ init.bias_init_zero_(self.norm.bias)
42
+
43
+ init.lecun_normal_init_(self.fc1.weight)
44
+ init.lecun_normal_init_(self.fc2.weight)
45
+ init.final_init_(self.fc3.weight)
46
+
47
+ def forward(self, x: Tensor, chunk_size: int = None) -> Tensor:
48
+ """Perform a forward pass.
49
+
50
+ Parameters
51
+ ----------
52
+ x: torch.Tensor
53
+ The input data of shape (..., D)
54
+
55
+ Returns
56
+ -------
57
+ x: torch.Tensor
58
+ The output data of shape (..., D)
59
+
60
+ """
61
+ x = self.norm(x)
62
+
63
+ if chunk_size is None or self.training:
64
+ x = self.silu(self.fc1(x)) * self.fc2(x)
65
+ x = self.fc3(x)
66
+ return x
67
+ else:
68
+ # Compute in chunks
69
+ for i in range(0, self.hidden, chunk_size):
70
+ fc1_slice = self.fc1.weight[i : i + chunk_size, :]
71
+ fc2_slice = self.fc2.weight[i : i + chunk_size, :]
72
+ fc3_slice = self.fc3.weight[:, i : i + chunk_size]
73
+ x_chunk = self.silu((x @ fc1_slice.T)) * (x @ fc2_slice.T)
74
+ if i == 0:
75
+ x_out = x_chunk @ fc3_slice.T
76
+ else:
77
+ x_out = x_out + x_chunk @ fc3_slice.T
78
+ return x_out
vb_layers_triangular_mult.py CHANGED
@@ -1,215 +1,215 @@
1
- import importlib
2
-
3
- import torch
4
- from torch import Tensor, nn
5
-
6
- from . import vb_layers_initialize as init
7
-
8
-
9
- @torch.compiler.disable
10
- def kernel_triangular_mult(
11
- x,
12
- direction,
13
- mask,
14
- norm_in_weight,
15
- norm_in_bias,
16
- p_in_weight,
17
- g_in_weight,
18
- norm_out_weight,
19
- norm_out_bias,
20
- p_out_weight,
21
- g_out_weight,
22
- eps,
23
- ):
24
- triangle_module = importlib.import_module("cuequivariance_torch.primitives.triangle")
25
- triangle_multiplicative_update = triangle_module.triangle_multiplicative_update
26
- return triangle_multiplicative_update(
27
- x,
28
- direction=direction,
29
- mask=mask,
30
- norm_in_weight=norm_in_weight,
31
- norm_in_bias=norm_in_bias,
32
- p_in_weight=p_in_weight,
33
- g_in_weight=g_in_weight,
34
- norm_out_weight=norm_out_weight,
35
- norm_out_bias=norm_out_bias,
36
- p_out_weight=p_out_weight,
37
- g_out_weight=g_out_weight,
38
- eps=eps,
39
- )
40
-
41
-
42
- class TriangleMultiplicationOutgoing(nn.Module):
43
- """TriangleMultiplicationOutgoing."""
44
-
45
- def __init__(self, dim: int = 128) -> None:
46
- """Initialize the TriangularUpdate module.
47
-
48
- Parameters
49
- ----------
50
- dim: int
51
- The dimension of the input, default 128
52
-
53
- """
54
- super().__init__()
55
-
56
- self.norm_in = nn.LayerNorm(dim, eps=1e-5)
57
- self.p_in = nn.Linear(dim, 2 * dim, bias=False)
58
- self.g_in = nn.Linear(dim, 2 * dim, bias=False)
59
-
60
- self.norm_out = nn.LayerNorm(dim)
61
- self.p_out = nn.Linear(dim, dim, bias=False)
62
- self.g_out = nn.Linear(dim, dim, bias=False)
63
-
64
- init.bias_init_one_(self.norm_in.weight)
65
- init.bias_init_zero_(self.norm_in.bias)
66
-
67
- init.lecun_normal_init_(self.p_in.weight)
68
- init.gating_init_(self.g_in.weight)
69
-
70
- init.bias_init_one_(self.norm_out.weight)
71
- init.bias_init_zero_(self.norm_out.bias)
72
-
73
- init.final_init_(self.p_out.weight)
74
- init.gating_init_(self.g_out.weight)
75
-
76
- def forward(self, x: Tensor, mask: Tensor, use_kernels: bool = False) -> Tensor:
77
- """Perform a forward pass.
78
-
79
- Parameters
80
- ----------
81
- x: torch.Tensor
82
- The input data of shape (B, N, N, D)
83
- mask: torch.Tensor
84
- The input mask of shape (B, N, N)
85
- use_kernels: bool
86
- Whether to use the kernel
87
-
88
- Returns
89
- -------
90
- x: torch.Tensor
91
- The output data of shape (B, N, N, D)
92
-
93
- """
94
- if use_kernels:
95
- return kernel_triangular_mult(
96
- x,
97
- direction="outgoing",
98
- mask=mask,
99
- norm_in_weight=self.norm_in.weight,
100
- norm_in_bias=self.norm_in.bias,
101
- p_in_weight=self.p_in.weight,
102
- g_in_weight=self.g_in.weight,
103
- norm_out_weight=self.norm_out.weight,
104
- norm_out_bias=self.norm_out.bias,
105
- p_out_weight=self.p_out.weight,
106
- g_out_weight=self.g_out.weight,
107
- eps=1e-5,
108
- )
109
-
110
- # Input gating: D -> D
111
- x = self.norm_in(x)
112
- x_in = x
113
- x = self.p_in(x) * self.g_in(x).sigmoid()
114
-
115
- # Apply mask
116
- x = x * mask.unsqueeze(-1)
117
-
118
- # Split input and cast to float
119
- a, b = torch.chunk(x.float(), 2, dim=-1)
120
-
121
- # Triangular projection
122
- x = torch.einsum("bikd,bjkd->bijd", a, b)
123
-
124
- # Output gating
125
- x = self.p_out(self.norm_out(x)) * self.g_out(x_in).sigmoid()
126
-
127
- return x
128
-
129
-
130
- class TriangleMultiplicationIncoming(nn.Module):
131
- """TriangleMultiplicationIncoming."""
132
-
133
- def __init__(self, dim: int = 128) -> None:
134
- """Initialize the TriangularUpdate module.
135
-
136
- Parameters
137
- ----------
138
- dim: int
139
- The dimension of the input, default 128
140
-
141
- """
142
- super().__init__()
143
-
144
- self.norm_in = nn.LayerNorm(dim, eps=1e-5)
145
- self.p_in = nn.Linear(dim, 2 * dim, bias=False)
146
- self.g_in = nn.Linear(dim, 2 * dim, bias=False)
147
-
148
- self.norm_out = nn.LayerNorm(dim)
149
- self.p_out = nn.Linear(dim, dim, bias=False)
150
- self.g_out = nn.Linear(dim, dim, bias=False)
151
-
152
- init.bias_init_one_(self.norm_in.weight)
153
- init.bias_init_zero_(self.norm_in.bias)
154
-
155
- init.lecun_normal_init_(self.p_in.weight)
156
- init.gating_init_(self.g_in.weight)
157
-
158
- init.bias_init_one_(self.norm_out.weight)
159
- init.bias_init_zero_(self.norm_out.bias)
160
-
161
- init.final_init_(self.p_out.weight)
162
- init.gating_init_(self.g_out.weight)
163
-
164
- def forward(self, x: Tensor, mask: Tensor, use_kernels: bool = False) -> Tensor:
165
- """Perform a forward pass.
166
-
167
- Parameters
168
- ----------
169
- x: torch.Tensor
170
- The input data of shape (B, N, N, D)
171
- mask: torch.Tensor
172
- The input mask of shape (B, N, N)
173
- use_kernels: bool
174
- Whether to use the kernel
175
-
176
- Returns
177
- -------
178
- x: torch.Tensor
179
- The output data of shape (B, N, N, D)
180
-
181
- """
182
- if use_kernels:
183
- return kernel_triangular_mult(
184
- x,
185
- direction="incoming",
186
- mask=mask,
187
- norm_in_weight=self.norm_in.weight,
188
- norm_in_bias=self.norm_in.bias,
189
- p_in_weight=self.p_in.weight,
190
- g_in_weight=self.g_in.weight,
191
- norm_out_weight=self.norm_out.weight,
192
- norm_out_bias=self.norm_out.bias,
193
- p_out_weight=self.p_out.weight,
194
- g_out_weight=self.g_out.weight,
195
- eps=1e-5,
196
- )
197
-
198
- # Input gating: D -> D
199
- x = self.norm_in(x)
200
- x_in = x
201
- x = self.p_in(x) * self.g_in(x).sigmoid()
202
-
203
- # Apply mask
204
- x = x * mask.unsqueeze(-1)
205
-
206
- # Split input and cast to float
207
- a, b = torch.chunk(x.float(), 2, dim=-1)
208
-
209
- # Triangular projection
210
- x = torch.einsum("bkid,bkjd->bijd", a, b)
211
-
212
- # Output gating
213
- x = self.p_out(self.norm_out(x)) * self.g_out(x_in).sigmoid()
214
-
215
- return x
 
1
+ import importlib
2
+
3
+ import torch
4
+ from torch import Tensor, nn
5
+
6
+ from . import vb_layers_initialize as init
7
+
8
+
9
+ @torch.compiler.disable
10
+ def kernel_triangular_mult(
11
+ x,
12
+ direction,
13
+ mask,
14
+ norm_in_weight,
15
+ norm_in_bias,
16
+ p_in_weight,
17
+ g_in_weight,
18
+ norm_out_weight,
19
+ norm_out_bias,
20
+ p_out_weight,
21
+ g_out_weight,
22
+ eps,
23
+ ):
24
+ triangle_module = importlib.import_module("cuequivariance_torch.primitives.triangle")
25
+ triangle_multiplicative_update = triangle_module.triangle_multiplicative_update
26
+ return triangle_multiplicative_update(
27
+ x,
28
+ direction=direction,
29
+ mask=mask,
30
+ norm_in_weight=norm_in_weight,
31
+ norm_in_bias=norm_in_bias,
32
+ p_in_weight=p_in_weight,
33
+ g_in_weight=g_in_weight,
34
+ norm_out_weight=norm_out_weight,
35
+ norm_out_bias=norm_out_bias,
36
+ p_out_weight=p_out_weight,
37
+ g_out_weight=g_out_weight,
38
+ eps=eps,
39
+ )
40
+
41
+
42
+ class TriangleMultiplicationOutgoing(nn.Module):
43
+ """TriangleMultiplicationOutgoing."""
44
+
45
+ def __init__(self, dim: int = 128) -> None:
46
+ """Initialize the TriangularUpdate module.
47
+
48
+ Parameters
49
+ ----------
50
+ dim: int
51
+ The dimension of the input, default 128
52
+
53
+ """
54
+ super().__init__()
55
+
56
+ self.norm_in = nn.LayerNorm(dim, eps=1e-5)
57
+ self.p_in = nn.Linear(dim, 2 * dim, bias=False)
58
+ self.g_in = nn.Linear(dim, 2 * dim, bias=False)
59
+
60
+ self.norm_out = nn.LayerNorm(dim)
61
+ self.p_out = nn.Linear(dim, dim, bias=False)
62
+ self.g_out = nn.Linear(dim, dim, bias=False)
63
+
64
+ init.bias_init_one_(self.norm_in.weight)
65
+ init.bias_init_zero_(self.norm_in.bias)
66
+
67
+ init.lecun_normal_init_(self.p_in.weight)
68
+ init.gating_init_(self.g_in.weight)
69
+
70
+ init.bias_init_one_(self.norm_out.weight)
71
+ init.bias_init_zero_(self.norm_out.bias)
72
+
73
+ init.final_init_(self.p_out.weight)
74
+ init.gating_init_(self.g_out.weight)
75
+
76
+ def forward(self, x: Tensor, mask: Tensor, use_kernels: bool = False) -> Tensor:
77
+ """Perform a forward pass.
78
+
79
+ Parameters
80
+ ----------
81
+ x: torch.Tensor
82
+ The input data of shape (B, N, N, D)
83
+ mask: torch.Tensor
84
+ The input mask of shape (B, N, N)
85
+ use_kernels: bool
86
+ Whether to use the kernel
87
+
88
+ Returns
89
+ -------
90
+ x: torch.Tensor
91
+ The output data of shape (B, N, N, D)
92
+
93
+ """
94
+ if use_kernels:
95
+ return kernel_triangular_mult(
96
+ x,
97
+ direction="outgoing",
98
+ mask=mask,
99
+ norm_in_weight=self.norm_in.weight,
100
+ norm_in_bias=self.norm_in.bias,
101
+ p_in_weight=self.p_in.weight,
102
+ g_in_weight=self.g_in.weight,
103
+ norm_out_weight=self.norm_out.weight,
104
+ norm_out_bias=self.norm_out.bias,
105
+ p_out_weight=self.p_out.weight,
106
+ g_out_weight=self.g_out.weight,
107
+ eps=1e-5,
108
+ )
109
+
110
+ # Input gating: D -> D
111
+ x = self.norm_in(x)
112
+ x_in = x
113
+ x = self.p_in(x) * self.g_in(x).sigmoid()
114
+
115
+ # Apply mask
116
+ x = x * mask.unsqueeze(-1)
117
+
118
+ # Split input and cast to float
119
+ a, b = torch.chunk(x.float(), 2, dim=-1)
120
+
121
+ # Triangular projection
122
+ x = torch.einsum("bikd,bjkd->bijd", a, b)
123
+
124
+ # Output gating
125
+ x = self.p_out(self.norm_out(x)) * self.g_out(x_in).sigmoid()
126
+
127
+ return x
128
+
129
+
130
+ class TriangleMultiplicationIncoming(nn.Module):
131
+ """TriangleMultiplicationIncoming."""
132
+
133
+ def __init__(self, dim: int = 128) -> None:
134
+ """Initialize the TriangularUpdate module.
135
+
136
+ Parameters
137
+ ----------
138
+ dim: int
139
+ The dimension of the input, default 128
140
+
141
+ """
142
+ super().__init__()
143
+
144
+ self.norm_in = nn.LayerNorm(dim, eps=1e-5)
145
+ self.p_in = nn.Linear(dim, 2 * dim, bias=False)
146
+ self.g_in = nn.Linear(dim, 2 * dim, bias=False)
147
+
148
+ self.norm_out = nn.LayerNorm(dim)
149
+ self.p_out = nn.Linear(dim, dim, bias=False)
150
+ self.g_out = nn.Linear(dim, dim, bias=False)
151
+
152
+ init.bias_init_one_(self.norm_in.weight)
153
+ init.bias_init_zero_(self.norm_in.bias)
154
+
155
+ init.lecun_normal_init_(self.p_in.weight)
156
+ init.gating_init_(self.g_in.weight)
157
+
158
+ init.bias_init_one_(self.norm_out.weight)
159
+ init.bias_init_zero_(self.norm_out.bias)
160
+
161
+ init.final_init_(self.p_out.weight)
162
+ init.gating_init_(self.g_out.weight)
163
+
164
+ def forward(self, x: Tensor, mask: Tensor, use_kernels: bool = False) -> Tensor:
165
+ """Perform a forward pass.
166
+
167
+ Parameters
168
+ ----------
169
+ x: torch.Tensor
170
+ The input data of shape (B, N, N, D)
171
+ mask: torch.Tensor
172
+ The input mask of shape (B, N, N)
173
+ use_kernels: bool
174
+ Whether to use the kernel
175
+
176
+ Returns
177
+ -------
178
+ x: torch.Tensor
179
+ The output data of shape (B, N, N, D)
180
+
181
+ """
182
+ if use_kernels:
183
+ return kernel_triangular_mult(
184
+ x,
185
+ direction="incoming",
186
+ mask=mask,
187
+ norm_in_weight=self.norm_in.weight,
188
+ norm_in_bias=self.norm_in.bias,
189
+ p_in_weight=self.p_in.weight,
190
+ g_in_weight=self.g_in.weight,
191
+ norm_out_weight=self.norm_out.weight,
192
+ norm_out_bias=self.norm_out.bias,
193
+ p_out_weight=self.p_out.weight,
194
+ g_out_weight=self.g_out.weight,
195
+ eps=1e-5,
196
+ )
197
+
198
+ # Input gating: D -> D
199
+ x = self.norm_in(x)
200
+ x_in = x
201
+ x = self.p_in(x) * self.g_in(x).sigmoid()
202
+
203
+ # Apply mask
204
+ x = x * mask.unsqueeze(-1)
205
+
206
+ # Split input and cast to float
207
+ a, b = torch.chunk(x.float(), 2, dim=-1)
208
+
209
+ # Triangular projection
210
+ x = torch.einsum("bkid,bkjd->bijd", a, b)
211
+
212
+ # Output gating
213
+ x = self.p_out(self.norm_out(x)) * self.g_out(x_in).sigmoid()
214
+
215
+ return x
vb_loss_diffusionv2.py CHANGED
@@ -1,138 +1,138 @@
1
- # started from code from https://github.com/lucidrains/alphafold3-pytorch, MIT License, Copyright (c) 2024 Phil Wang
2
-
3
- import torch
4
- import torch.nn.functional as F
5
- from einops import einsum, rearrange
6
-
7
-
8
- def weighted_rigid_align(
9
- true_coords, # Float['b n 3'], # true coordinates
10
- pred_coords, # Float['b n 3'], # predicted coordinates
11
- weights, # Float['b n'], # weights for each atom
12
- mask, # Bool['b n'] | None = None # mask for variable lengths
13
- ): # -> Float['b n 3']:
14
- """Algorithm 28 : note there is a problem with the pseudocode in the paper where predicted and
15
- GT are swapped in algorithm 28, but correct in equation (2)."""
16
-
17
- out_shape = torch.broadcast_shapes(true_coords.shape, pred_coords.shape)
18
- *batch_size, num_points, dim = out_shape
19
- weights = (mask * weights).unsqueeze(-1)
20
-
21
- # Compute weighted centroids
22
- true_centroid = (true_coords * weights).sum(dim=-2, keepdim=True) / weights.sum(
23
- dim=-2, keepdim=True
24
- )
25
- pred_centroid = (pred_coords * weights).sum(dim=-2, keepdim=True) / weights.sum(
26
- dim=-2, keepdim=True
27
- )
28
-
29
- # Center the coordinates
30
- true_coords_centered = true_coords - true_centroid
31
- pred_coords_centered = pred_coords - pred_centroid
32
-
33
- if torch.any(mask.sum(dim=-1) < (dim + 1)):
34
- print(
35
- "Warning: The size of one of the point clouds is <= dim+1. "
36
- + "`WeightedRigidAlign` cannot return a unique rotation."
37
- )
38
-
39
- # Compute the weighted covariance matrix
40
- cov_matrix = einsum(
41
- weights * pred_coords_centered,
42
- true_coords_centered,
43
- "... n i, ... n j -> ... i j",
44
- )
45
-
46
- # Compute the SVD of the covariance matrix, required float32 for svd and determinant
47
- original_dtype = cov_matrix.dtype
48
- cov_matrix_32 = cov_matrix.to(dtype=torch.float32)
49
-
50
- U, S, V = torch.linalg.svd(
51
- cov_matrix_32, driver="gesvd" if cov_matrix_32.is_cuda else None
52
- )
53
- V = V.mH
54
-
55
- # Catch ambiguous rotation by checking the magnitude of singular values
56
- if (S.abs() <= 1e-15).any() and not (num_points < (dim + 1)):
57
- print(
58
- "Warning: Excessively low rank of "
59
- + "cross-correlation between aligned point clouds. "
60
- + "`WeightedRigidAlign` cannot return a unique rotation."
61
- )
62
-
63
- # Compute the rotation matrix
64
- rot_matrix = torch.einsum("... i j, ... k j -> ... i k", U, V).to(
65
- dtype=torch.float32
66
- )
67
-
68
- # Ensure proper rotation matrix with determinant 1
69
- F = torch.eye(dim, dtype=cov_matrix_32.dtype, device=cov_matrix.device)[
70
- None
71
- ].repeat(*batch_size, 1, 1)
72
- F[..., -1, -1] = torch.det(rot_matrix)
73
- rot_matrix = einsum(U, F, V, "... i j, ... j k, ... l k -> ... i l")
74
- rot_matrix = rot_matrix.to(dtype=original_dtype)
75
-
76
- # Apply the rotation and translation
77
- aligned_coords = (
78
- einsum(true_coords_centered, rot_matrix, "... n i, ... j i -> ... n j")
79
- + pred_centroid
80
- )
81
- aligned_coords.detach_()
82
-
83
- return aligned_coords
84
-
85
-
86
- def smooth_lddt_loss(
87
- pred_coords, # Float['b n 3'],
88
- true_coords, # Float['b n 3'],
89
- is_nucleotide, # Bool['b n'],
90
- coords_mask, # Bool['b n'] | None = None,
91
- nucleic_acid_cutoff: float = 30.0,
92
- other_cutoff: float = 15.0,
93
- multiplicity: int = 1,
94
- ): # -> Float['']:
95
- """Algorithm 27
96
- pred_coords: predicted coordinates
97
- true_coords: true coordinates
98
- Note: for efficiency pred_coords is the only one with the multiplicity expanded
99
- TODO: add weighing which overweight the smooth lddt contribution close to t=0 (not present in the paper)
100
- """
101
- lddt = []
102
- for i in range(true_coords.shape[0]):
103
- true_dists = torch.cdist(true_coords[i], true_coords[i])
104
-
105
- is_nucleotide_i = is_nucleotide[i // multiplicity]
106
- coords_mask_i = coords_mask[i // multiplicity]
107
-
108
- is_nucleotide_pair = is_nucleotide_i.unsqueeze(-1).expand(
109
- -1, is_nucleotide_i.shape[-1]
110
- )
111
-
112
- mask = is_nucleotide_pair * (true_dists < nucleic_acid_cutoff).float()
113
- mask += (1 - is_nucleotide_pair) * (true_dists < other_cutoff).float()
114
- mask *= 1 - torch.eye(pred_coords.shape[1], device=pred_coords.device)
115
- mask *= coords_mask_i.unsqueeze(-1)
116
- mask *= coords_mask_i.unsqueeze(-2)
117
-
118
- valid_pairs = mask.nonzero()
119
- true_dists_i = true_dists[valid_pairs[:, 0], valid_pairs[:, 1]]
120
-
121
- pred_coords_i1 = pred_coords[i, valid_pairs[:, 0]]
122
- pred_coords_i2 = pred_coords[i, valid_pairs[:, 1]]
123
- pred_dists_i = F.pairwise_distance(pred_coords_i1, pred_coords_i2)
124
-
125
- dist_diff_i = torch.abs(true_dists_i - pred_dists_i)
126
-
127
- eps_i = (
128
- F.sigmoid(0.5 - dist_diff_i)
129
- + F.sigmoid(1.0 - dist_diff_i)
130
- + F.sigmoid(2.0 - dist_diff_i)
131
- + F.sigmoid(4.0 - dist_diff_i)
132
- ) / 4.0
133
-
134
- lddt_i = eps_i.sum() / (valid_pairs.shape[0] + 1e-5)
135
- lddt.append(lddt_i)
136
-
137
- # average over batch & multiplicity
138
- return 1.0 - torch.stack(lddt, dim=0).mean(dim=0)
 
1
+ # started from code from https://github.com/lucidrains/alphafold3-pytorch, MIT License, Copyright (c) 2024 Phil Wang
2
+
3
+ import torch
4
+ import torch.nn.functional as F
5
+ from einops import einsum, rearrange
6
+
7
+
8
+ def weighted_rigid_align(
9
+ true_coords, # Float['b n 3'], # true coordinates
10
+ pred_coords, # Float['b n 3'], # predicted coordinates
11
+ weights, # Float['b n'], # weights for each atom
12
+ mask, # Bool['b n'] | None = None # mask for variable lengths
13
+ ): # -> Float['b n 3']:
14
+ """Algorithm 28 : note there is a problem with the pseudocode in the paper where predicted and
15
+ GT are swapped in algorithm 28, but correct in equation (2)."""
16
+
17
+ out_shape = torch.broadcast_shapes(true_coords.shape, pred_coords.shape)
18
+ *batch_size, num_points, dim = out_shape
19
+ weights = (mask * weights).unsqueeze(-1)
20
+
21
+ # Compute weighted centroids
22
+ true_centroid = (true_coords * weights).sum(dim=-2, keepdim=True) / weights.sum(
23
+ dim=-2, keepdim=True
24
+ )
25
+ pred_centroid = (pred_coords * weights).sum(dim=-2, keepdim=True) / weights.sum(
26
+ dim=-2, keepdim=True
27
+ )
28
+
29
+ # Center the coordinates
30
+ true_coords_centered = true_coords - true_centroid
31
+ pred_coords_centered = pred_coords - pred_centroid
32
+
33
+ if torch.any(mask.sum(dim=-1) < (dim + 1)):
34
+ print(
35
+ "Warning: The size of one of the point clouds is <= dim+1. "
36
+ + "`WeightedRigidAlign` cannot return a unique rotation."
37
+ )
38
+
39
+ # Compute the weighted covariance matrix
40
+ cov_matrix = einsum(
41
+ weights * pred_coords_centered,
42
+ true_coords_centered,
43
+ "... n i, ... n j -> ... i j",
44
+ )
45
+
46
+ # Compute the SVD of the covariance matrix, required float32 for svd and determinant
47
+ original_dtype = cov_matrix.dtype
48
+ cov_matrix_32 = cov_matrix.to(dtype=torch.float32)
49
+
50
+ U, S, V = torch.linalg.svd(
51
+ cov_matrix_32, driver="gesvd" if cov_matrix_32.is_cuda else None
52
+ )
53
+ V = V.mH
54
+
55
+ # Catch ambiguous rotation by checking the magnitude of singular values
56
+ if (S.abs() <= 1e-15).any() and not (num_points < (dim + 1)):
57
+ print(
58
+ "Warning: Excessively low rank of "
59
+ + "cross-correlation between aligned point clouds. "
60
+ + "`WeightedRigidAlign` cannot return a unique rotation."
61
+ )
62
+
63
+ # Compute the rotation matrix
64
+ rot_matrix = torch.einsum("... i j, ... k j -> ... i k", U, V).to(
65
+ dtype=torch.float32
66
+ )
67
+
68
+ # Ensure proper rotation matrix with determinant 1
69
+ F = torch.eye(dim, dtype=cov_matrix_32.dtype, device=cov_matrix.device)[
70
+ None
71
+ ].repeat(*batch_size, 1, 1)
72
+ F[..., -1, -1] = torch.det(rot_matrix)
73
+ rot_matrix = einsum(U, F, V, "... i j, ... j k, ... l k -> ... i l")
74
+ rot_matrix = rot_matrix.to(dtype=original_dtype)
75
+
76
+ # Apply the rotation and translation
77
+ aligned_coords = (
78
+ einsum(true_coords_centered, rot_matrix, "... n i, ... j i -> ... n j")
79
+ + pred_centroid
80
+ )
81
+ aligned_coords.detach_()
82
+
83
+ return aligned_coords
84
+
85
+
86
+ def smooth_lddt_loss(
87
+ pred_coords, # Float['b n 3'],
88
+ true_coords, # Float['b n 3'],
89
+ is_nucleotide, # Bool['b n'],
90
+ coords_mask, # Bool['b n'] | None = None,
91
+ nucleic_acid_cutoff: float = 30.0,
92
+ other_cutoff: float = 15.0,
93
+ multiplicity: int = 1,
94
+ ): # -> Float['']:
95
+ """Algorithm 27
96
+ pred_coords: predicted coordinates
97
+ true_coords: true coordinates
98
+ Note: for efficiency pred_coords is the only one with the multiplicity expanded
99
+ TODO: add weighing which overweight the smooth lddt contribution close to t=0 (not present in the paper)
100
+ """
101
+ lddt = []
102
+ for i in range(true_coords.shape[0]):
103
+ true_dists = torch.cdist(true_coords[i], true_coords[i])
104
+
105
+ is_nucleotide_i = is_nucleotide[i // multiplicity]
106
+ coords_mask_i = coords_mask[i // multiplicity]
107
+
108
+ is_nucleotide_pair = is_nucleotide_i.unsqueeze(-1).expand(
109
+ -1, is_nucleotide_i.shape[-1]
110
+ )
111
+
112
+ mask = is_nucleotide_pair * (true_dists < nucleic_acid_cutoff).float()
113
+ mask += (1 - is_nucleotide_pair) * (true_dists < other_cutoff).float()
114
+ mask *= 1 - torch.eye(pred_coords.shape[1], device=pred_coords.device)
115
+ mask *= coords_mask_i.unsqueeze(-1)
116
+ mask *= coords_mask_i.unsqueeze(-2)
117
+
118
+ valid_pairs = mask.nonzero()
119
+ true_dists_i = true_dists[valid_pairs[:, 0], valid_pairs[:, 1]]
120
+
121
+ pred_coords_i1 = pred_coords[i, valid_pairs[:, 0]]
122
+ pred_coords_i2 = pred_coords[i, valid_pairs[:, 1]]
123
+ pred_dists_i = F.pairwise_distance(pred_coords_i1, pred_coords_i2)
124
+
125
+ dist_diff_i = torch.abs(true_dists_i - pred_dists_i)
126
+
127
+ eps_i = (
128
+ F.sigmoid(0.5 - dist_diff_i)
129
+ + F.sigmoid(1.0 - dist_diff_i)
130
+ + F.sigmoid(2.0 - dist_diff_i)
131
+ + F.sigmoid(4.0 - dist_diff_i)
132
+ ) / 4.0
133
+
134
+ lddt_i = eps_i.sum() / (valid_pairs.shape[0] + 1e-5)
135
+ lddt.append(lddt_i)
136
+
137
+ # average over batch & multiplicity
138
+ return 1.0 - torch.stack(lddt, dim=0).mean(dim=0)
vb_modules_confidencev2.py CHANGED
@@ -1,498 +1,498 @@
1
- import torch
2
- from torch import nn
3
- from torch.nn.functional import pad
4
-
5
- from . import vb_const as const
6
- from . import vb_layers_initialize as init
7
- from .vb_layers_confidence_utils import (
8
- compute_aggregated_metric,
9
- compute_ptms,
10
- )
11
- from .vb_layers_pairformer import PairformerModule
12
- from .vb_modules_encodersv2 import RelativePositionEncoder
13
- from .vb_modules_trunkv2 import (
14
- ContactConditioning,
15
- )
16
- from .vb_modules_utils import LinearNoBias
17
-
18
-
19
- class ConfidenceModule(nn.Module):
20
- """Algorithm 31"""
21
-
22
- def __init__(
23
- self,
24
- token_s,
25
- token_z,
26
- pairformer_args: dict,
27
- num_dist_bins=64,
28
- token_level_confidence=True,
29
- max_dist=22,
30
- add_s_to_z_prod=False,
31
- add_s_input_to_s=False,
32
- add_z_input_to_z=False,
33
- maximum_bond_distance=0,
34
- bond_type_feature=False,
35
- confidence_args: dict = None,
36
- compile_pairformer=False,
37
- fix_sym_check=False,
38
- cyclic_pos_enc=False,
39
- return_latent_feats=False,
40
- conditioning_cutoff_min=None,
41
- conditioning_cutoff_max=None,
42
- **kwargs,
43
- ):
44
- super().__init__()
45
- self.max_num_atoms_per_token = 23
46
- if "no_update_s" in pairformer_args:
47
- self.no_update_s = pairformer_args["no_update_s"]
48
- else:
49
- self.no_update_s = False
50
- boundaries = torch.linspace(2, max_dist, num_dist_bins - 1)
51
- self.register_buffer("boundaries", boundaries)
52
- self.dist_bin_pairwise_embed = nn.Embedding(num_dist_bins, token_z)
53
- init.gating_init_(self.dist_bin_pairwise_embed.weight)
54
- self.token_level_confidence = token_level_confidence
55
-
56
- self.s_to_z = LinearNoBias(token_s, token_z)
57
- self.s_to_z_transpose = LinearNoBias(token_s, token_z)
58
- init.gating_init_(self.s_to_z.weight)
59
- init.gating_init_(self.s_to_z_transpose.weight)
60
-
61
- self.add_s_to_z_prod = add_s_to_z_prod
62
- if add_s_to_z_prod:
63
- self.s_to_z_prod_in1 = LinearNoBias(token_s, token_z)
64
- self.s_to_z_prod_in2 = LinearNoBias(token_s, token_z)
65
- self.s_to_z_prod_out = LinearNoBias(token_z, token_z)
66
- init.gating_init_(self.s_to_z_prod_out.weight)
67
-
68
- self.s_inputs_norm = nn.LayerNorm(token_s)
69
- if not self.no_update_s:
70
- self.s_norm = nn.LayerNorm(token_s)
71
- self.z_norm = nn.LayerNorm(token_z)
72
-
73
- self.add_s_input_to_s = add_s_input_to_s
74
- if add_s_input_to_s:
75
- self.s_input_to_s = LinearNoBias(token_s, token_s)
76
- init.gating_init_(self.s_input_to_s.weight)
77
-
78
- self.add_z_input_to_z = add_z_input_to_z
79
- if add_z_input_to_z:
80
- self.rel_pos = RelativePositionEncoder(
81
- token_z, fix_sym_check=fix_sym_check, cyclic_pos_enc=cyclic_pos_enc
82
- )
83
- self.token_bonds = nn.Linear(
84
- 1 if maximum_bond_distance == 0 else maximum_bond_distance + 2,
85
- token_z,
86
- bias=False,
87
- )
88
- self.bond_type_feature = bond_type_feature
89
- if bond_type_feature:
90
- self.token_bonds_type = nn.Embedding(len(const.bond_types) + 1, token_z)
91
-
92
- self.contact_conditioning = ContactConditioning(
93
- token_z=token_z,
94
- cutoff_min=conditioning_cutoff_min,
95
- cutoff_max=conditioning_cutoff_max,
96
- )
97
- pairformer_args["v2"] = True
98
- self.pairformer_stack = PairformerModule(
99
- token_s,
100
- token_z,
101
- **pairformer_args,
102
- )
103
- self.return_latent_feats = return_latent_feats
104
-
105
- self.confidence_heads = ConfidenceHeads(
106
- token_s,
107
- token_z,
108
- token_level_confidence=token_level_confidence,
109
- **confidence_args,
110
- )
111
-
112
- def forward(
113
- self,
114
- s_inputs, # Float['b n ts']
115
- s, # Float['b n ts']
116
- z, # Float['b n n tz']
117
- x_pred, # Float['bm m 3']
118
- feats,
119
- pred_distogram_logits,
120
- multiplicity=1,
121
- run_sequentially=False,
122
- use_kernels: bool = False,
123
- ):
124
- if run_sequentially and multiplicity > 1:
125
- assert z.shape[0] == 1, "Not supported with batch size > 1"
126
- out_dicts = []
127
- for sample_idx in range(multiplicity):
128
- out_dicts.append( # noqa: PERF401
129
- self.forward(
130
- s_inputs,
131
- s,
132
- z,
133
- x_pred[sample_idx : sample_idx + 1],
134
- feats,
135
- pred_distogram_logits,
136
- multiplicity=1,
137
- run_sequentially=False,
138
- use_kernels=use_kernels,
139
- )
140
- )
141
-
142
- out_dict = {}
143
- for key in out_dicts[0]:
144
- if key != "pair_chains_iptm":
145
- out_dict[key] = torch.cat([out[key] for out in out_dicts], dim=0)
146
- else:
147
- pair_chains_iptm = {}
148
- for chain_idx1 in out_dicts[0][key]:
149
- chains_iptm = {}
150
- for chain_idx2 in out_dicts[0][key][chain_idx1]:
151
- chains_iptm[chain_idx2] = torch.cat(
152
- [out[key][chain_idx1][chain_idx2] for out in out_dicts],
153
- dim=0,
154
- )
155
- pair_chains_iptm[chain_idx1] = chains_iptm
156
- out_dict[key] = pair_chains_iptm
157
- return out_dict
158
-
159
- s_inputs = self.s_inputs_norm(s_inputs)
160
- if not self.no_update_s:
161
- s = self.s_norm(s)
162
-
163
- if self.add_s_input_to_s:
164
- s = s + self.s_input_to_s(s_inputs)
165
-
166
- z = self.z_norm(z)
167
-
168
- if self.add_z_input_to_z:
169
- relative_position_encoding = self.rel_pos(feats)
170
- z = z + relative_position_encoding
171
- z = z + self.token_bonds(feats["token_bonds"].float())
172
- if self.bond_type_feature:
173
- z = z + self.token_bonds_type(feats["type_bonds"].long())
174
- z = z + self.contact_conditioning(feats)
175
-
176
- s = s.repeat_interleave(multiplicity, 0)
177
-
178
- z = (
179
- z
180
- + self.s_to_z(s_inputs)[:, :, None, :]
181
- + self.s_to_z_transpose(s_inputs)[:, None, :, :]
182
- )
183
- if self.add_s_to_z_prod:
184
- z = z + self.s_to_z_prod_out(
185
- self.s_to_z_prod_in1(s_inputs)[:, :, None, :]
186
- * self.s_to_z_prod_in2(s_inputs)[:, None, :, :]
187
- )
188
-
189
- z = z.repeat_interleave(multiplicity, 0)
190
- s_inputs = s_inputs.repeat_interleave(multiplicity, 0)
191
-
192
- token_to_rep_atom = feats["token_to_rep_atom"]
193
- token_to_rep_atom = token_to_rep_atom.repeat_interleave(multiplicity, 0)
194
- if len(x_pred.shape) == 4:
195
- B, mult, N, _ = x_pred.shape
196
- x_pred = x_pred.reshape(B * mult, N, -1)
197
- else:
198
- BM, N, _ = x_pred.shape
199
- x_pred_repr = torch.bmm(token_to_rep_atom.float(), x_pred)
200
- d = torch.cdist(x_pred_repr, x_pred_repr)
201
- distogram = (d.unsqueeze(-1) > self.boundaries).sum(dim=-1).long()
202
- distogram = self.dist_bin_pairwise_embed(distogram)
203
- z = z + distogram
204
-
205
- mask = feats["token_pad_mask"].repeat_interleave(multiplicity, 0)
206
- pair_mask = mask[:, :, None] * mask[:, None, :]
207
-
208
- s_t, z_t = self.pairformer_stack(
209
- s, z, mask=mask, pair_mask=pair_mask, use_kernels=use_kernels
210
- )
211
-
212
- # AF3 has residual connections, we remove them
213
- s = s_t
214
- z = z_t
215
-
216
- out_dict = {}
217
-
218
- if self.return_latent_feats:
219
- out_dict["s_conf"] = s
220
- out_dict["z_conf"] = z
221
-
222
- # confidence heads
223
- out_dict.update(
224
- self.confidence_heads(
225
- s=s,
226
- z=z,
227
- x_pred=x_pred,
228
- d=d,
229
- feats=feats,
230
- multiplicity=multiplicity,
231
- pred_distogram_logits=pred_distogram_logits,
232
- )
233
- )
234
- return out_dict
235
-
236
-
237
- class ConfidenceHeads(nn.Module):
238
- def __init__(
239
- self,
240
- token_s,
241
- token_z,
242
- num_plddt_bins=50,
243
- num_pde_bins=64,
244
- num_pae_bins=64,
245
- token_level_confidence=True,
246
- use_separate_heads: bool = False,
247
- **kwargs,
248
- ):
249
- super().__init__()
250
- self.max_num_atoms_per_token = 23
251
- self.token_level_confidence = token_level_confidence
252
- self.use_separate_heads = use_separate_heads
253
-
254
- if self.use_separate_heads:
255
- self.to_pae_intra_logits = LinearNoBias(token_z, num_pae_bins)
256
- self.to_pae_inter_logits = LinearNoBias(token_z, num_pae_bins)
257
- else:
258
- self.to_pae_logits = LinearNoBias(token_z, num_pae_bins)
259
-
260
- if self.use_separate_heads:
261
- self.to_pde_intra_logits = LinearNoBias(token_z, num_pde_bins)
262
- self.to_pde_inter_logits = LinearNoBias(token_z, num_pde_bins)
263
- else:
264
- self.to_pde_logits = LinearNoBias(token_z, num_pde_bins)
265
-
266
- if self.token_level_confidence:
267
- self.to_plddt_logits = LinearNoBias(token_s, num_plddt_bins)
268
- self.to_resolved_logits = LinearNoBias(token_s, 2)
269
- else:
270
- self.to_plddt_logits = LinearNoBias(
271
- token_s, num_plddt_bins * self.max_num_atoms_per_token
272
- )
273
- self.to_resolved_logits = LinearNoBias(
274
- token_s, 2 * self.max_num_atoms_per_token
275
- )
276
-
277
- def forward(
278
- self,
279
- s, # Float['b n ts']
280
- z, # Float['b n n tz']
281
- x_pred, # Float['bm m 3']
282
- d,
283
- feats,
284
- pred_distogram_logits,
285
- multiplicity=1,
286
- ):
287
- if self.use_separate_heads:
288
- asym_id_token = feats["asym_id"]
289
- is_same_chain = asym_id_token.unsqueeze(-1) == asym_id_token.unsqueeze(-2)
290
- is_different_chain = ~is_same_chain
291
-
292
- if self.use_separate_heads:
293
- pae_intra_logits = self.to_pae_intra_logits(z)
294
- pae_intra_logits = pae_intra_logits * is_same_chain.float().unsqueeze(-1)
295
-
296
- pae_inter_logits = self.to_pae_inter_logits(z)
297
- pae_inter_logits = pae_inter_logits * is_different_chain.float().unsqueeze(
298
- -1
299
- )
300
-
301
- pae_logits = pae_inter_logits + pae_intra_logits
302
- else:
303
- pae_logits = self.to_pae_logits(z)
304
-
305
- if self.use_separate_heads:
306
- pde_intra_logits = self.to_pde_intra_logits(z + z.transpose(1, 2))
307
- pde_intra_logits = pde_intra_logits * is_same_chain.float().unsqueeze(-1)
308
-
309
- pde_inter_logits = self.to_pde_inter_logits(z + z.transpose(1, 2))
310
- pde_inter_logits = pde_inter_logits * is_different_chain.float().unsqueeze(
311
- -1
312
- )
313
-
314
- pde_logits = pde_inter_logits + pde_intra_logits
315
- else:
316
- pde_logits = self.to_pde_logits(z + z.transpose(1, 2))
317
- resolved_logits = self.to_resolved_logits(s)
318
- plddt_logits = self.to_plddt_logits(s)
319
-
320
- ligand_weight = 20
321
- non_interface_weight = 1
322
- interface_weight = 10
323
-
324
- token_type = feats["mol_type"]
325
- token_type = token_type.repeat_interleave(multiplicity, 0)
326
- is_ligand_token = (token_type == const.chain_type_ids["NONPOLYMER"]).float()
327
-
328
- if self.token_level_confidence:
329
- plddt = compute_aggregated_metric(plddt_logits)
330
- token_pad_mask = feats["token_pad_mask"].repeat_interleave(multiplicity, 0)
331
- complex_plddt = (plddt * token_pad_mask).sum(dim=-1) / token_pad_mask.sum(
332
- dim=-1
333
- )
334
-
335
- is_contact = (d < 8).float()
336
- is_different_chain = (
337
- feats["asym_id"].unsqueeze(-1) != feats["asym_id"].unsqueeze(-2)
338
- ).float()
339
- is_different_chain = is_different_chain.repeat_interleave(multiplicity, 0)
340
- token_interface_mask = torch.max(
341
- is_contact * is_different_chain * (1 - is_ligand_token).unsqueeze(-1),
342
- dim=-1,
343
- ).values
344
- token_non_interface_mask = (1 - token_interface_mask) * (
345
- 1 - is_ligand_token
346
- )
347
- iplddt_weight = (
348
- is_ligand_token * ligand_weight
349
- + token_interface_mask * interface_weight
350
- + token_non_interface_mask * non_interface_weight
351
- )
352
- complex_iplddt = (plddt * token_pad_mask * iplddt_weight).sum(
353
- dim=-1
354
- ) / torch.sum(token_pad_mask * iplddt_weight, dim=-1)
355
-
356
- else:
357
- # token to atom conversion for resolved logits
358
- B, N, _ = resolved_logits.shape
359
- resolved_logits = resolved_logits.reshape(
360
- B, N, self.max_num_atoms_per_token, 2
361
- )
362
-
363
- arange_max_num_atoms = (
364
- torch.arange(self.max_num_atoms_per_token)
365
- .reshape(1, 1, -1)
366
- .to(resolved_logits.device)
367
- )
368
- max_num_atoms_mask = (
369
- feats["atom_to_token"].sum(1).unsqueeze(-1) > arange_max_num_atoms
370
- )
371
- resolved_logits = resolved_logits[:, max_num_atoms_mask.squeeze(0)]
372
- resolved_logits = pad(
373
- resolved_logits,
374
- (
375
- 0,
376
- 0,
377
- 0,
378
- int(
379
- feats["atom_pad_mask"].shape[1]
380
- - feats["atom_pad_mask"].sum().item()
381
- ),
382
- ),
383
- value=0,
384
- )
385
- plddt_logits = plddt_logits.reshape(B, N, self.max_num_atoms_per_token, -1)
386
- plddt_logits = plddt_logits[:, max_num_atoms_mask.squeeze(0)]
387
- plddt_logits = pad(
388
- plddt_logits,
389
- (
390
- 0,
391
- 0,
392
- 0,
393
- int(
394
- feats["atom_pad_mask"].shape[1]
395
- - feats["atom_pad_mask"].sum().item()
396
- ),
397
- ),
398
- value=0,
399
- )
400
- atom_pad_mask = feats["atom_pad_mask"].repeat_interleave(multiplicity, 0)
401
- plddt = compute_aggregated_metric(plddt_logits)
402
-
403
- complex_plddt = (plddt * atom_pad_mask).sum(dim=-1) / atom_pad_mask.sum(
404
- dim=-1
405
- )
406
- token_type = feats["mol_type"].float()
407
- atom_to_token = feats["atom_to_token"].float()
408
- chain_id_token = feats["asym_id"].float()
409
- atom_type = torch.bmm(atom_to_token, token_type.unsqueeze(-1)).squeeze(-1)
410
- is_ligand_atom = (atom_type == const.chain_type_ids["NONPOLYMER"]).float()
411
- d_atom = torch.cdist(x_pred, x_pred)
412
- is_contact = (d_atom < 8).float()
413
- chain_id_atom = torch.bmm(
414
- atom_to_token, chain_id_token.unsqueeze(-1)
415
- ).squeeze(-1)
416
- is_different_chain = (
417
- chain_id_atom.unsqueeze(-1) != chain_id_atom.unsqueeze(-2)
418
- ).float()
419
-
420
- atom_interface_mask = torch.max(
421
- is_contact * is_different_chain * (1 - is_ligand_atom).unsqueeze(-1),
422
- dim=-1,
423
- ).values
424
- atom_non_interface_mask = (1 - atom_interface_mask) * (1 - is_ligand_atom)
425
- iplddt_weight = (
426
- is_ligand_atom * ligand_weight
427
- + atom_interface_mask * interface_weight
428
- + atom_non_interface_mask * non_interface_weight
429
- )
430
-
431
- complex_iplddt = (plddt * feats["atom_pad_mask"] * iplddt_weight).sum(
432
- dim=-1
433
- ) / torch.sum(feats["atom_pad_mask"] * iplddt_weight, dim=-1)
434
-
435
- # Compute the gPDE and giPDE
436
- pde = compute_aggregated_metric(pde_logits, end=32)
437
- pred_distogram_prob = nn.functional.softmax(
438
- pred_distogram_logits, dim=-1
439
- ).repeat_interleave(multiplicity, 0)
440
- contacts = torch.zeros((1, 1, 1, 64), dtype=pred_distogram_prob.dtype).to(
441
- pred_distogram_prob.device
442
- )
443
- contacts[:, :, :, :20] = 1.0
444
- prob_contact = (pred_distogram_prob * contacts).sum(-1)
445
- token_pad_mask = feats["token_pad_mask"].repeat_interleave(multiplicity, 0)
446
- token_pad_pair_mask = (
447
- token_pad_mask.unsqueeze(-1)
448
- * token_pad_mask.unsqueeze(-2)
449
- * (
450
- 1
451
- - torch.eye(
452
- token_pad_mask.shape[1], device=token_pad_mask.device
453
- ).unsqueeze(0)
454
- )
455
- )
456
- token_pair_mask = token_pad_pair_mask * prob_contact
457
- complex_pde = (pde * token_pair_mask).sum(dim=(1, 2)) / token_pair_mask.sum(
458
- dim=(1, 2)
459
- )
460
- asym_id = feats["asym_id"].repeat_interleave(multiplicity, 0)
461
- token_interface_pair_mask = token_pair_mask * (
462
- asym_id.unsqueeze(-1) != asym_id.unsqueeze(-2)
463
- )
464
- complex_ipde = (pde * token_interface_pair_mask).sum(dim=(1, 2)) / (
465
- token_interface_pair_mask.sum(dim=(1, 2)) + 1e-5
466
- )
467
- out_dict = dict(
468
- pde_logits=pde_logits,
469
- plddt_logits=plddt_logits,
470
- resolved_logits=resolved_logits,
471
- pde=pde,
472
- plddt=plddt,
473
- complex_plddt=complex_plddt,
474
- complex_iplddt=complex_iplddt,
475
- complex_pde=complex_pde,
476
- complex_ipde=complex_ipde,
477
- )
478
- out_dict["pae_logits"] = pae_logits
479
- out_dict["pae"] = compute_aggregated_metric(pae_logits, end=32)
480
-
481
- try:
482
- ptm, iptm, ligand_iptm, protein_iptm, pair_chains_iptm = compute_ptms(
483
- pae_logits, x_pred, feats, multiplicity
484
- )
485
- out_dict["ptm"] = ptm
486
- out_dict["iptm"] = iptm
487
- out_dict["ligand_iptm"] = ligand_iptm
488
- out_dict["protein_iptm"] = protein_iptm
489
- out_dict["pair_chains_iptm"] = pair_chains_iptm
490
- except Exception as e:
491
- print(f"Error in compute_ptms: {e}")
492
- out_dict["ptm"] = torch.zeros_like(complex_plddt)
493
- out_dict["iptm"] = torch.zeros_like(complex_plddt)
494
- out_dict["ligand_iptm"] = torch.zeros_like(complex_plddt)
495
- out_dict["protein_iptm"] = torch.zeros_like(complex_plddt)
496
- out_dict["pair_chains_iptm"] = torch.zeros_like(complex_plddt)
497
-
498
- return out_dict
 
1
+ import torch
2
+ from torch import nn
3
+ from torch.nn.functional import pad
4
+
5
+ from . import vb_const as const
6
+ from . import vb_layers_initialize as init
7
+ from .vb_layers_confidence_utils import (
8
+ compute_aggregated_metric,
9
+ compute_ptms,
10
+ )
11
+ from .vb_layers_pairformer import PairformerModule
12
+ from .vb_modules_encodersv2 import RelativePositionEncoder
13
+ from .vb_modules_trunkv2 import (
14
+ ContactConditioning,
15
+ )
16
+ from .vb_modules_utils import LinearNoBias
17
+
18
+
19
+ class ConfidenceModule(nn.Module):
20
+ """Algorithm 31"""
21
+
22
+ def __init__(
23
+ self,
24
+ token_s,
25
+ token_z,
26
+ pairformer_args: dict,
27
+ num_dist_bins=64,
28
+ token_level_confidence=True,
29
+ max_dist=22,
30
+ add_s_to_z_prod=False,
31
+ add_s_input_to_s=False,
32
+ add_z_input_to_z=False,
33
+ maximum_bond_distance=0,
34
+ bond_type_feature=False,
35
+ confidence_args: dict = None,
36
+ compile_pairformer=False,
37
+ fix_sym_check=False,
38
+ cyclic_pos_enc=False,
39
+ return_latent_feats=False,
40
+ conditioning_cutoff_min=None,
41
+ conditioning_cutoff_max=None,
42
+ **kwargs,
43
+ ):
44
+ super().__init__()
45
+ self.max_num_atoms_per_token = 23
46
+ if "no_update_s" in pairformer_args:
47
+ self.no_update_s = pairformer_args["no_update_s"]
48
+ else:
49
+ self.no_update_s = False
50
+ boundaries = torch.linspace(2, max_dist, num_dist_bins - 1)
51
+ self.register_buffer("boundaries", boundaries)
52
+ self.dist_bin_pairwise_embed = nn.Embedding(num_dist_bins, token_z)
53
+ init.gating_init_(self.dist_bin_pairwise_embed.weight)
54
+ self.token_level_confidence = token_level_confidence
55
+
56
+ self.s_to_z = LinearNoBias(token_s, token_z)
57
+ self.s_to_z_transpose = LinearNoBias(token_s, token_z)
58
+ init.gating_init_(self.s_to_z.weight)
59
+ init.gating_init_(self.s_to_z_transpose.weight)
60
+
61
+ self.add_s_to_z_prod = add_s_to_z_prod
62
+ if add_s_to_z_prod:
63
+ self.s_to_z_prod_in1 = LinearNoBias(token_s, token_z)
64
+ self.s_to_z_prod_in2 = LinearNoBias(token_s, token_z)
65
+ self.s_to_z_prod_out = LinearNoBias(token_z, token_z)
66
+ init.gating_init_(self.s_to_z_prod_out.weight)
67
+
68
+ self.s_inputs_norm = nn.LayerNorm(token_s)
69
+ if not self.no_update_s:
70
+ self.s_norm = nn.LayerNorm(token_s)
71
+ self.z_norm = nn.LayerNorm(token_z)
72
+
73
+ self.add_s_input_to_s = add_s_input_to_s
74
+ if add_s_input_to_s:
75
+ self.s_input_to_s = LinearNoBias(token_s, token_s)
76
+ init.gating_init_(self.s_input_to_s.weight)
77
+
78
+ self.add_z_input_to_z = add_z_input_to_z
79
+ if add_z_input_to_z:
80
+ self.rel_pos = RelativePositionEncoder(
81
+ token_z, fix_sym_check=fix_sym_check, cyclic_pos_enc=cyclic_pos_enc
82
+ )
83
+ self.token_bonds = nn.Linear(
84
+ 1 if maximum_bond_distance == 0 else maximum_bond_distance + 2,
85
+ token_z,
86
+ bias=False,
87
+ )
88
+ self.bond_type_feature = bond_type_feature
89
+ if bond_type_feature:
90
+ self.token_bonds_type = nn.Embedding(len(const.bond_types) + 1, token_z)
91
+
92
+ self.contact_conditioning = ContactConditioning(
93
+ token_z=token_z,
94
+ cutoff_min=conditioning_cutoff_min,
95
+ cutoff_max=conditioning_cutoff_max,
96
+ )
97
+ pairformer_args["v2"] = True
98
+ self.pairformer_stack = PairformerModule(
99
+ token_s,
100
+ token_z,
101
+ **pairformer_args,
102
+ )
103
+ self.return_latent_feats = return_latent_feats
104
+
105
+ self.confidence_heads = ConfidenceHeads(
106
+ token_s,
107
+ token_z,
108
+ token_level_confidence=token_level_confidence,
109
+ **confidence_args,
110
+ )
111
+
112
+ def forward(
113
+ self,
114
+ s_inputs, # Float['b n ts']
115
+ s, # Float['b n ts']
116
+ z, # Float['b n n tz']
117
+ x_pred, # Float['bm m 3']
118
+ feats,
119
+ pred_distogram_logits,
120
+ multiplicity=1,
121
+ run_sequentially=False,
122
+ use_kernels: bool = False,
123
+ ):
124
+ if run_sequentially and multiplicity > 1:
125
+ assert z.shape[0] == 1, "Not supported with batch size > 1"
126
+ out_dicts = []
127
+ for sample_idx in range(multiplicity):
128
+ out_dicts.append( # noqa: PERF401
129
+ self.forward(
130
+ s_inputs,
131
+ s,
132
+ z,
133
+ x_pred[sample_idx : sample_idx + 1],
134
+ feats,
135
+ pred_distogram_logits,
136
+ multiplicity=1,
137
+ run_sequentially=False,
138
+ use_kernels=use_kernels,
139
+ )
140
+ )
141
+
142
+ out_dict = {}
143
+ for key in out_dicts[0]:
144
+ if key != "pair_chains_iptm":
145
+ out_dict[key] = torch.cat([out[key] for out in out_dicts], dim=0)
146
+ else:
147
+ pair_chains_iptm = {}
148
+ for chain_idx1 in out_dicts[0][key]:
149
+ chains_iptm = {}
150
+ for chain_idx2 in out_dicts[0][key][chain_idx1]:
151
+ chains_iptm[chain_idx2] = torch.cat(
152
+ [out[key][chain_idx1][chain_idx2] for out in out_dicts],
153
+ dim=0,
154
+ )
155
+ pair_chains_iptm[chain_idx1] = chains_iptm
156
+ out_dict[key] = pair_chains_iptm
157
+ return out_dict
158
+
159
+ s_inputs = self.s_inputs_norm(s_inputs)
160
+ if not self.no_update_s:
161
+ s = self.s_norm(s)
162
+
163
+ if self.add_s_input_to_s:
164
+ s = s + self.s_input_to_s(s_inputs)
165
+
166
+ z = self.z_norm(z)
167
+
168
+ if self.add_z_input_to_z:
169
+ relative_position_encoding = self.rel_pos(feats)
170
+ z = z + relative_position_encoding
171
+ z = z + self.token_bonds(feats["token_bonds"].float())
172
+ if self.bond_type_feature:
173
+ z = z + self.token_bonds_type(feats["type_bonds"].long())
174
+ z = z + self.contact_conditioning(feats)
175
+
176
+ s = s.repeat_interleave(multiplicity, 0)
177
+
178
+ z = (
179
+ z
180
+ + self.s_to_z(s_inputs)[:, :, None, :]
181
+ + self.s_to_z_transpose(s_inputs)[:, None, :, :]
182
+ )
183
+ if self.add_s_to_z_prod:
184
+ z = z + self.s_to_z_prod_out(
185
+ self.s_to_z_prod_in1(s_inputs)[:, :, None, :]
186
+ * self.s_to_z_prod_in2(s_inputs)[:, None, :, :]
187
+ )
188
+
189
+ z = z.repeat_interleave(multiplicity, 0)
190
+ s_inputs = s_inputs.repeat_interleave(multiplicity, 0)
191
+
192
+ token_to_rep_atom = feats["token_to_rep_atom"]
193
+ token_to_rep_atom = token_to_rep_atom.repeat_interleave(multiplicity, 0)
194
+ if len(x_pred.shape) == 4:
195
+ B, mult, N, _ = x_pred.shape
196
+ x_pred = x_pred.reshape(B * mult, N, -1)
197
+ else:
198
+ BM, N, _ = x_pred.shape
199
+ x_pred_repr = torch.bmm(token_to_rep_atom.float(), x_pred)
200
+ d = torch.cdist(x_pred_repr, x_pred_repr)
201
+ distogram = (d.unsqueeze(-1) > self.boundaries).sum(dim=-1).long()
202
+ distogram = self.dist_bin_pairwise_embed(distogram)
203
+ z = z + distogram
204
+
205
+ mask = feats["token_pad_mask"].repeat_interleave(multiplicity, 0)
206
+ pair_mask = mask[:, :, None] * mask[:, None, :]
207
+
208
+ s_t, z_t = self.pairformer_stack(
209
+ s, z, mask=mask, pair_mask=pair_mask, use_kernels=use_kernels
210
+ )
211
+
212
+ # AF3 has residual connections, we remove them
213
+ s = s_t
214
+ z = z_t
215
+
216
+ out_dict = {}
217
+
218
+ if self.return_latent_feats:
219
+ out_dict["s_conf"] = s
220
+ out_dict["z_conf"] = z
221
+
222
+ # confidence heads
223
+ out_dict.update(
224
+ self.confidence_heads(
225
+ s=s,
226
+ z=z,
227
+ x_pred=x_pred,
228
+ d=d,
229
+ feats=feats,
230
+ multiplicity=multiplicity,
231
+ pred_distogram_logits=pred_distogram_logits,
232
+ )
233
+ )
234
+ return out_dict
235
+
236
+
237
+ class ConfidenceHeads(nn.Module):
238
+ def __init__(
239
+ self,
240
+ token_s,
241
+ token_z,
242
+ num_plddt_bins=50,
243
+ num_pde_bins=64,
244
+ num_pae_bins=64,
245
+ token_level_confidence=True,
246
+ use_separate_heads: bool = False,
247
+ **kwargs,
248
+ ):
249
+ super().__init__()
250
+ self.max_num_atoms_per_token = 23
251
+ self.token_level_confidence = token_level_confidence
252
+ self.use_separate_heads = use_separate_heads
253
+
254
+ if self.use_separate_heads:
255
+ self.to_pae_intra_logits = LinearNoBias(token_z, num_pae_bins)
256
+ self.to_pae_inter_logits = LinearNoBias(token_z, num_pae_bins)
257
+ else:
258
+ self.to_pae_logits = LinearNoBias(token_z, num_pae_bins)
259
+
260
+ if self.use_separate_heads:
261
+ self.to_pde_intra_logits = LinearNoBias(token_z, num_pde_bins)
262
+ self.to_pde_inter_logits = LinearNoBias(token_z, num_pde_bins)
263
+ else:
264
+ self.to_pde_logits = LinearNoBias(token_z, num_pde_bins)
265
+
266
+ if self.token_level_confidence:
267
+ self.to_plddt_logits = LinearNoBias(token_s, num_plddt_bins)
268
+ self.to_resolved_logits = LinearNoBias(token_s, 2)
269
+ else:
270
+ self.to_plddt_logits = LinearNoBias(
271
+ token_s, num_plddt_bins * self.max_num_atoms_per_token
272
+ )
273
+ self.to_resolved_logits = LinearNoBias(
274
+ token_s, 2 * self.max_num_atoms_per_token
275
+ )
276
+
277
+ def forward(
278
+ self,
279
+ s, # Float['b n ts']
280
+ z, # Float['b n n tz']
281
+ x_pred, # Float['bm m 3']
282
+ d,
283
+ feats,
284
+ pred_distogram_logits,
285
+ multiplicity=1,
286
+ ):
287
+ if self.use_separate_heads:
288
+ asym_id_token = feats["asym_id"]
289
+ is_same_chain = asym_id_token.unsqueeze(-1) == asym_id_token.unsqueeze(-2)
290
+ is_different_chain = ~is_same_chain
291
+
292
+ if self.use_separate_heads:
293
+ pae_intra_logits = self.to_pae_intra_logits(z)
294
+ pae_intra_logits = pae_intra_logits * is_same_chain.float().unsqueeze(-1)
295
+
296
+ pae_inter_logits = self.to_pae_inter_logits(z)
297
+ pae_inter_logits = pae_inter_logits * is_different_chain.float().unsqueeze(
298
+ -1
299
+ )
300
+
301
+ pae_logits = pae_inter_logits + pae_intra_logits
302
+ else:
303
+ pae_logits = self.to_pae_logits(z)
304
+
305
+ if self.use_separate_heads:
306
+ pde_intra_logits = self.to_pde_intra_logits(z + z.transpose(1, 2))
307
+ pde_intra_logits = pde_intra_logits * is_same_chain.float().unsqueeze(-1)
308
+
309
+ pde_inter_logits = self.to_pde_inter_logits(z + z.transpose(1, 2))
310
+ pde_inter_logits = pde_inter_logits * is_different_chain.float().unsqueeze(
311
+ -1
312
+ )
313
+
314
+ pde_logits = pde_inter_logits + pde_intra_logits
315
+ else:
316
+ pde_logits = self.to_pde_logits(z + z.transpose(1, 2))
317
+ resolved_logits = self.to_resolved_logits(s)
318
+ plddt_logits = self.to_plddt_logits(s)
319
+
320
+ ligand_weight = 20
321
+ non_interface_weight = 1
322
+ interface_weight = 10
323
+
324
+ token_type = feats["mol_type"]
325
+ token_type = token_type.repeat_interleave(multiplicity, 0)
326
+ is_ligand_token = (token_type == const.chain_type_ids["NONPOLYMER"]).float()
327
+
328
+ if self.token_level_confidence:
329
+ plddt = compute_aggregated_metric(plddt_logits)
330
+ token_pad_mask = feats["token_pad_mask"].repeat_interleave(multiplicity, 0)
331
+ complex_plddt = (plddt * token_pad_mask).sum(dim=-1) / token_pad_mask.sum(
332
+ dim=-1
333
+ )
334
+
335
+ is_contact = (d < 8).float()
336
+ is_different_chain = (
337
+ feats["asym_id"].unsqueeze(-1) != feats["asym_id"].unsqueeze(-2)
338
+ ).float()
339
+ is_different_chain = is_different_chain.repeat_interleave(multiplicity, 0)
340
+ token_interface_mask = torch.max(
341
+ is_contact * is_different_chain * (1 - is_ligand_token).unsqueeze(-1),
342
+ dim=-1,
343
+ ).values
344
+ token_non_interface_mask = (1 - token_interface_mask) * (
345
+ 1 - is_ligand_token
346
+ )
347
+ iplddt_weight = (
348
+ is_ligand_token * ligand_weight
349
+ + token_interface_mask * interface_weight
350
+ + token_non_interface_mask * non_interface_weight
351
+ )
352
+ complex_iplddt = (plddt * token_pad_mask * iplddt_weight).sum(
353
+ dim=-1
354
+ ) / torch.sum(token_pad_mask * iplddt_weight, dim=-1)
355
+
356
+ else:
357
+ # token to atom conversion for resolved logits
358
+ B, N, _ = resolved_logits.shape
359
+ resolved_logits = resolved_logits.reshape(
360
+ B, N, self.max_num_atoms_per_token, 2
361
+ )
362
+
363
+ arange_max_num_atoms = (
364
+ torch.arange(self.max_num_atoms_per_token)
365
+ .reshape(1, 1, -1)
366
+ .to(resolved_logits.device)
367
+ )
368
+ max_num_atoms_mask = (
369
+ feats["atom_to_token"].sum(1).unsqueeze(-1) > arange_max_num_atoms
370
+ )
371
+ resolved_logits = resolved_logits[:, max_num_atoms_mask.squeeze(0)]
372
+ resolved_logits = pad(
373
+ resolved_logits,
374
+ (
375
+ 0,
376
+ 0,
377
+ 0,
378
+ int(
379
+ feats["atom_pad_mask"].shape[1]
380
+ - feats["atom_pad_mask"].sum().item()
381
+ ),
382
+ ),
383
+ value=0,
384
+ )
385
+ plddt_logits = plddt_logits.reshape(B, N, self.max_num_atoms_per_token, -1)
386
+ plddt_logits = plddt_logits[:, max_num_atoms_mask.squeeze(0)]
387
+ plddt_logits = pad(
388
+ plddt_logits,
389
+ (
390
+ 0,
391
+ 0,
392
+ 0,
393
+ int(
394
+ feats["atom_pad_mask"].shape[1]
395
+ - feats["atom_pad_mask"].sum().item()
396
+ ),
397
+ ),
398
+ value=0,
399
+ )
400
+ atom_pad_mask = feats["atom_pad_mask"].repeat_interleave(multiplicity, 0)
401
+ plddt = compute_aggregated_metric(plddt_logits)
402
+
403
+ complex_plddt = (plddt * atom_pad_mask).sum(dim=-1) / atom_pad_mask.sum(
404
+ dim=-1
405
+ )
406
+ token_type = feats["mol_type"].float()
407
+ atom_to_token = feats["atom_to_token"].float()
408
+ chain_id_token = feats["asym_id"].float()
409
+ atom_type = torch.bmm(atom_to_token, token_type.unsqueeze(-1)).squeeze(-1)
410
+ is_ligand_atom = (atom_type == const.chain_type_ids["NONPOLYMER"]).float()
411
+ d_atom = torch.cdist(x_pred, x_pred)
412
+ is_contact = (d_atom < 8).float()
413
+ chain_id_atom = torch.bmm(
414
+ atom_to_token, chain_id_token.unsqueeze(-1)
415
+ ).squeeze(-1)
416
+ is_different_chain = (
417
+ chain_id_atom.unsqueeze(-1) != chain_id_atom.unsqueeze(-2)
418
+ ).float()
419
+
420
+ atom_interface_mask = torch.max(
421
+ is_contact * is_different_chain * (1 - is_ligand_atom).unsqueeze(-1),
422
+ dim=-1,
423
+ ).values
424
+ atom_non_interface_mask = (1 - atom_interface_mask) * (1 - is_ligand_atom)
425
+ iplddt_weight = (
426
+ is_ligand_atom * ligand_weight
427
+ + atom_interface_mask * interface_weight
428
+ + atom_non_interface_mask * non_interface_weight
429
+ )
430
+
431
+ complex_iplddt = (plddt * feats["atom_pad_mask"] * iplddt_weight).sum(
432
+ dim=-1
433
+ ) / torch.sum(feats["atom_pad_mask"] * iplddt_weight, dim=-1)
434
+
435
+ # Compute the gPDE and giPDE
436
+ pde = compute_aggregated_metric(pde_logits, end=32)
437
+ pred_distogram_prob = nn.functional.softmax(
438
+ pred_distogram_logits, dim=-1
439
+ ).repeat_interleave(multiplicity, 0)
440
+ contacts = torch.zeros((1, 1, 1, 64), dtype=pred_distogram_prob.dtype).to(
441
+ pred_distogram_prob.device
442
+ )
443
+ contacts[:, :, :, :20] = 1.0
444
+ prob_contact = (pred_distogram_prob * contacts).sum(-1)
445
+ token_pad_mask = feats["token_pad_mask"].repeat_interleave(multiplicity, 0)
446
+ token_pad_pair_mask = (
447
+ token_pad_mask.unsqueeze(-1)
448
+ * token_pad_mask.unsqueeze(-2)
449
+ * (
450
+ 1
451
+ - torch.eye(
452
+ token_pad_mask.shape[1], device=token_pad_mask.device
453
+ ).unsqueeze(0)
454
+ )
455
+ )
456
+ token_pair_mask = token_pad_pair_mask * prob_contact
457
+ complex_pde = (pde * token_pair_mask).sum(dim=(1, 2)) / token_pair_mask.sum(
458
+ dim=(1, 2)
459
+ )
460
+ asym_id = feats["asym_id"].repeat_interleave(multiplicity, 0)
461
+ token_interface_pair_mask = token_pair_mask * (
462
+ asym_id.unsqueeze(-1) != asym_id.unsqueeze(-2)
463
+ )
464
+ complex_ipde = (pde * token_interface_pair_mask).sum(dim=(1, 2)) / (
465
+ token_interface_pair_mask.sum(dim=(1, 2)) + 1e-5
466
+ )
467
+ out_dict = dict(
468
+ pde_logits=pde_logits,
469
+ plddt_logits=plddt_logits,
470
+ resolved_logits=resolved_logits,
471
+ pde=pde,
472
+ plddt=plddt,
473
+ complex_plddt=complex_plddt,
474
+ complex_iplddt=complex_iplddt,
475
+ complex_pde=complex_pde,
476
+ complex_ipde=complex_ipde,
477
+ )
478
+ out_dict["pae_logits"] = pae_logits
479
+ out_dict["pae"] = compute_aggregated_metric(pae_logits, end=32)
480
+
481
+ try:
482
+ ptm, iptm, ligand_iptm, protein_iptm, pair_chains_iptm = compute_ptms(
483
+ pae_logits, x_pred, feats, multiplicity
484
+ )
485
+ out_dict["ptm"] = ptm
486
+ out_dict["iptm"] = iptm
487
+ out_dict["ligand_iptm"] = ligand_iptm
488
+ out_dict["protein_iptm"] = protein_iptm
489
+ out_dict["pair_chains_iptm"] = pair_chains_iptm
490
+ except Exception as e:
491
+ print(f"Error in compute_ptms: {e}")
492
+ out_dict["ptm"] = torch.zeros_like(complex_plddt)
493
+ out_dict["iptm"] = torch.zeros_like(complex_plddt)
494
+ out_dict["ligand_iptm"] = torch.zeros_like(complex_plddt)
495
+ out_dict["protein_iptm"] = torch.zeros_like(complex_plddt)
496
+ out_dict["pair_chains_iptm"] = torch.zeros_like(complex_plddt)
497
+
498
+ return out_dict
vb_modules_diffusion_conditioning.py CHANGED
@@ -1,116 +1,116 @@
1
- from __future__ import annotations
2
-
3
- import torch
4
- from torch import nn
5
- from torch.nn import Module
6
-
7
- from .vb_modules_encodersv2 import (
8
- AtomEncoder,
9
- PairwiseConditioning,
10
- )
11
-
12
-
13
- class DiffusionConditioning(Module):
14
- def __init__(
15
- self,
16
- token_s: int,
17
- token_z: int,
18
- atom_s: int,
19
- atom_z: int,
20
- atoms_per_window_queries: int = 32,
21
- atoms_per_window_keys: int = 128,
22
- atom_encoder_depth: int = 3,
23
- atom_encoder_heads: int = 4,
24
- token_transformer_depth: int = 24,
25
- token_transformer_heads: int = 8,
26
- atom_decoder_depth: int = 3,
27
- atom_decoder_heads: int = 4,
28
- atom_feature_dim: int = 128,
29
- conditioning_transition_layers: int = 2,
30
- use_no_atom_char: bool = False,
31
- use_atom_backbone_feat: bool = False,
32
- use_residue_feats_atoms: bool = False,
33
- ) -> None:
34
- super().__init__()
35
-
36
- self.pairwise_conditioner = PairwiseConditioning(
37
- token_z=token_z,
38
- dim_token_rel_pos_feats=token_z,
39
- num_transitions=conditioning_transition_layers,
40
- )
41
-
42
- self.atom_encoder = AtomEncoder(
43
- atom_s=atom_s,
44
- atom_z=atom_z,
45
- token_s=token_s,
46
- token_z=token_z,
47
- atoms_per_window_queries=atoms_per_window_queries,
48
- atoms_per_window_keys=atoms_per_window_keys,
49
- atom_feature_dim=atom_feature_dim,
50
- structure_prediction=True,
51
- use_no_atom_char=use_no_atom_char,
52
- use_atom_backbone_feat=use_atom_backbone_feat,
53
- use_residue_feats_atoms=use_residue_feats_atoms,
54
- )
55
-
56
- self.atom_enc_proj_z = nn.ModuleList()
57
- for _ in range(atom_encoder_depth):
58
- self.atom_enc_proj_z.append(
59
- nn.Sequential(
60
- nn.LayerNorm(atom_z),
61
- nn.Linear(atom_z, atom_encoder_heads, bias=False),
62
- )
63
- )
64
-
65
- self.atom_dec_proj_z = nn.ModuleList()
66
- for _ in range(atom_decoder_depth):
67
- self.atom_dec_proj_z.append(
68
- nn.Sequential(
69
- nn.LayerNorm(atom_z),
70
- nn.Linear(atom_z, atom_decoder_heads, bias=False),
71
- )
72
- )
73
-
74
- self.token_trans_proj_z = nn.ModuleList()
75
- for _ in range(token_transformer_depth):
76
- self.token_trans_proj_z.append(
77
- nn.Sequential(
78
- nn.LayerNorm(token_z),
79
- nn.Linear(token_z, token_transformer_heads, bias=False),
80
- )
81
- )
82
-
83
- def forward(
84
- self,
85
- s_trunk, # Float['b n ts']
86
- z_trunk, # Float['b n n tz']
87
- relative_position_encoding, # Float['b n n tz']
88
- feats,
89
- ):
90
- z = self.pairwise_conditioner(
91
- z_trunk,
92
- relative_position_encoding,
93
- )
94
-
95
- q, c, p, to_keys = self.atom_encoder(
96
- feats=feats,
97
- s_trunk=s_trunk, # Float['b n ts'],
98
- z=z, # Float['b n n tz'],
99
- )
100
-
101
- atom_enc_bias = []
102
- for layer in self.atom_enc_proj_z:
103
- atom_enc_bias.append(layer(p))
104
- atom_enc_bias = torch.cat(atom_enc_bias, dim=-1)
105
-
106
- atom_dec_bias = []
107
- for layer in self.atom_dec_proj_z:
108
- atom_dec_bias.append(layer(p))
109
- atom_dec_bias = torch.cat(atom_dec_bias, dim=-1)
110
-
111
- token_trans_bias = []
112
- for layer in self.token_trans_proj_z:
113
- token_trans_bias.append(layer(z))
114
- token_trans_bias = torch.cat(token_trans_bias, dim=-1)
115
-
116
- return q, c, to_keys, atom_enc_bias, atom_dec_bias, token_trans_bias
 
1
+ from __future__ import annotations
2
+
3
+ import torch
4
+ from torch import nn
5
+ from torch.nn import Module
6
+
7
+ from .vb_modules_encodersv2 import (
8
+ AtomEncoder,
9
+ PairwiseConditioning,
10
+ )
11
+
12
+
13
+ class DiffusionConditioning(Module):
14
+ def __init__(
15
+ self,
16
+ token_s: int,
17
+ token_z: int,
18
+ atom_s: int,
19
+ atom_z: int,
20
+ atoms_per_window_queries: int = 32,
21
+ atoms_per_window_keys: int = 128,
22
+ atom_encoder_depth: int = 3,
23
+ atom_encoder_heads: int = 4,
24
+ token_transformer_depth: int = 24,
25
+ token_transformer_heads: int = 8,
26
+ atom_decoder_depth: int = 3,
27
+ atom_decoder_heads: int = 4,
28
+ atom_feature_dim: int = 128,
29
+ conditioning_transition_layers: int = 2,
30
+ use_no_atom_char: bool = False,
31
+ use_atom_backbone_feat: bool = False,
32
+ use_residue_feats_atoms: bool = False,
33
+ ) -> None:
34
+ super().__init__()
35
+
36
+ self.pairwise_conditioner = PairwiseConditioning(
37
+ token_z=token_z,
38
+ dim_token_rel_pos_feats=token_z,
39
+ num_transitions=conditioning_transition_layers,
40
+ )
41
+
42
+ self.atom_encoder = AtomEncoder(
43
+ atom_s=atom_s,
44
+ atom_z=atom_z,
45
+ token_s=token_s,
46
+ token_z=token_z,
47
+ atoms_per_window_queries=atoms_per_window_queries,
48
+ atoms_per_window_keys=atoms_per_window_keys,
49
+ atom_feature_dim=atom_feature_dim,
50
+ structure_prediction=True,
51
+ use_no_atom_char=use_no_atom_char,
52
+ use_atom_backbone_feat=use_atom_backbone_feat,
53
+ use_residue_feats_atoms=use_residue_feats_atoms,
54
+ )
55
+
56
+ self.atom_enc_proj_z = nn.ModuleList()
57
+ for _ in range(atom_encoder_depth):
58
+ self.atom_enc_proj_z.append(
59
+ nn.Sequential(
60
+ nn.LayerNorm(atom_z),
61
+ nn.Linear(atom_z, atom_encoder_heads, bias=False),
62
+ )
63
+ )
64
+
65
+ self.atom_dec_proj_z = nn.ModuleList()
66
+ for _ in range(atom_decoder_depth):
67
+ self.atom_dec_proj_z.append(
68
+ nn.Sequential(
69
+ nn.LayerNorm(atom_z),
70
+ nn.Linear(atom_z, atom_decoder_heads, bias=False),
71
+ )
72
+ )
73
+
74
+ self.token_trans_proj_z = nn.ModuleList()
75
+ for _ in range(token_transformer_depth):
76
+ self.token_trans_proj_z.append(
77
+ nn.Sequential(
78
+ nn.LayerNorm(token_z),
79
+ nn.Linear(token_z, token_transformer_heads, bias=False),
80
+ )
81
+ )
82
+
83
+ def forward(
84
+ self,
85
+ s_trunk, # Float['b n ts']
86
+ z_trunk, # Float['b n n tz']
87
+ relative_position_encoding, # Float['b n n tz']
88
+ feats,
89
+ ):
90
+ z = self.pairwise_conditioner(
91
+ z_trunk,
92
+ relative_position_encoding,
93
+ )
94
+
95
+ q, c, p, to_keys = self.atom_encoder(
96
+ feats=feats,
97
+ s_trunk=s_trunk, # Float['b n ts'],
98
+ z=z, # Float['b n n tz'],
99
+ )
100
+
101
+ atom_enc_bias = []
102
+ for layer in self.atom_enc_proj_z:
103
+ atom_enc_bias.append(layer(p))
104
+ atom_enc_bias = torch.cat(atom_enc_bias, dim=-1)
105
+
106
+ atom_dec_bias = []
107
+ for layer in self.atom_dec_proj_z:
108
+ atom_dec_bias.append(layer(p))
109
+ atom_dec_bias = torch.cat(atom_dec_bias, dim=-1)
110
+
111
+ token_trans_bias = []
112
+ for layer in self.token_trans_proj_z:
113
+ token_trans_bias.append(layer(z))
114
+ token_trans_bias = torch.cat(token_trans_bias, dim=-1)
115
+
116
+ return q, c, to_keys, atom_enc_bias, atom_dec_bias, token_trans_bias
vb_modules_diffusionv2.py CHANGED
@@ -1,693 +1,693 @@
1
- # started from code from https://github.com/lucidrains/alphafold3-pytorch, MIT License, Copyright (c) 2024 Phil Wang
2
-
3
- from __future__ import annotations
4
-
5
- from math import sqrt
6
-
7
- import numpy as np
8
- import torch
9
- import torch.nn.functional as F # noqa: N812
10
- from einops import rearrange
11
- from torch import nn
12
- from torch.nn import Module
13
-
14
- from . import vb_const as const
15
- from . import vb_layers_initialize as init
16
- from .vb_loss_diffusionv2 import (
17
- smooth_lddt_loss,
18
- weighted_rigid_align,
19
- )
20
- from .vb_modules_encodersv2 import (
21
- AtomAttentionDecoder,
22
- AtomAttentionEncoder,
23
- SingleConditioning,
24
- )
25
- from .vb_modules_transformersv2 import (
26
- DiffusionTransformer,
27
- )
28
- from .vb_modules_utils import (
29
- LinearNoBias,
30
- center_random_augmentation,
31
- compute_random_augmentation,
32
- default,
33
- log,
34
- )
35
- from .vb_potentials_potentials import get_potentials
36
-
37
-
38
- class DiffusionModule(Module):
39
- """Diffusion module"""
40
-
41
- def __init__(
42
- self,
43
- token_s: int,
44
- atom_s: int,
45
- atoms_per_window_queries: int = 32,
46
- atoms_per_window_keys: int = 128,
47
- sigma_data: int = 16,
48
- dim_fourier: int = 256,
49
- atom_encoder_depth: int = 3,
50
- atom_encoder_heads: int = 4,
51
- token_transformer_depth: int = 24,
52
- token_transformer_heads: int = 8,
53
- atom_decoder_depth: int = 3,
54
- atom_decoder_heads: int = 4,
55
- conditioning_transition_layers: int = 2,
56
- activation_checkpointing: bool = False,
57
- transformer_post_ln: bool = False,
58
- ) -> None:
59
- super().__init__()
60
-
61
- self.atoms_per_window_queries = atoms_per_window_queries
62
- self.atoms_per_window_keys = atoms_per_window_keys
63
- self.sigma_data = sigma_data
64
- self.activation_checkpointing = activation_checkpointing
65
-
66
- # conditioning
67
- self.single_conditioner = SingleConditioning(
68
- sigma_data=sigma_data,
69
- token_s=token_s,
70
- dim_fourier=dim_fourier,
71
- num_transitions=conditioning_transition_layers,
72
- )
73
-
74
- self.atom_attention_encoder = AtomAttentionEncoder(
75
- atom_s=atom_s,
76
- token_s=token_s,
77
- atoms_per_window_queries=atoms_per_window_queries,
78
- atoms_per_window_keys=atoms_per_window_keys,
79
- atom_encoder_depth=atom_encoder_depth,
80
- atom_encoder_heads=atom_encoder_heads,
81
- structure_prediction=True,
82
- activation_checkpointing=activation_checkpointing,
83
- transformer_post_layer_norm=transformer_post_ln,
84
- )
85
-
86
- self.s_to_a_linear = nn.Sequential(
87
- nn.LayerNorm(2 * token_s), LinearNoBias(2 * token_s, 2 * token_s)
88
- )
89
- init.final_init_(self.s_to_a_linear[1].weight)
90
-
91
- self.token_transformer = DiffusionTransformer(
92
- dim=2 * token_s,
93
- dim_single_cond=2 * token_s,
94
- depth=token_transformer_depth,
95
- heads=token_transformer_heads,
96
- activation_checkpointing=activation_checkpointing,
97
- # post_layer_norm=transformer_post_ln,
98
- )
99
-
100
- self.a_norm = nn.LayerNorm(
101
- 2 * token_s
102
- ) # if not transformer_post_ln else nn.Identity()
103
-
104
- self.atom_attention_decoder = AtomAttentionDecoder(
105
- atom_s=atom_s,
106
- token_s=token_s,
107
- attn_window_queries=atoms_per_window_queries,
108
- attn_window_keys=atoms_per_window_keys,
109
- atom_decoder_depth=atom_decoder_depth,
110
- atom_decoder_heads=atom_decoder_heads,
111
- activation_checkpointing=activation_checkpointing,
112
- # transformer_post_layer_norm=transformer_post_ln,
113
- )
114
-
115
- def forward(
116
- self,
117
- s_inputs, # Float['b n ts']
118
- s_trunk, # Float['b n ts']
119
- r_noisy, # Float['bm m 3']
120
- times, # Float['bm 1 1']
121
- feats,
122
- diffusion_conditioning,
123
- multiplicity=1,
124
- ):
125
- if self.activation_checkpointing and self.training:
126
- s, normed_fourier = torch.utils.checkpoint.checkpoint(
127
- self.single_conditioner,
128
- times,
129
- s_trunk.repeat_interleave(multiplicity, 0),
130
- s_inputs.repeat_interleave(multiplicity, 0),
131
- )
132
- else:
133
- s, normed_fourier = self.single_conditioner(
134
- times,
135
- s_trunk.repeat_interleave(multiplicity, 0),
136
- s_inputs.repeat_interleave(multiplicity, 0),
137
- )
138
-
139
- # Sequence-local Atom Attention and aggregation to coarse-grained tokens
140
- a, q_skip, c_skip, to_keys = self.atom_attention_encoder(
141
- feats=feats,
142
- q=diffusion_conditioning["q"].float(),
143
- c=diffusion_conditioning["c"].float(),
144
- atom_enc_bias=diffusion_conditioning["atom_enc_bias"].float(),
145
- to_keys=diffusion_conditioning["to_keys"],
146
- r=r_noisy, # Float['b m 3'],
147
- multiplicity=multiplicity,
148
- )
149
-
150
- # Full self-attention on token level
151
- a = a + self.s_to_a_linear(s)
152
-
153
- mask = feats["token_pad_mask"].repeat_interleave(multiplicity, 0)
154
- a = self.token_transformer(
155
- a,
156
- mask=mask.float(),
157
- s=s,
158
- bias=diffusion_conditioning[
159
- "token_trans_bias"
160
- ].float(), # note z is not expanded with multiplicity until after bias is computed
161
- multiplicity=multiplicity,
162
- )
163
- a = self.a_norm(a)
164
-
165
- # Broadcast token activations to atoms and run Sequence-local Atom Attention
166
- r_update = self.atom_attention_decoder(
167
- a=a,
168
- q=q_skip,
169
- c=c_skip,
170
- atom_dec_bias=diffusion_conditioning["atom_dec_bias"].float(),
171
- feats=feats,
172
- multiplicity=multiplicity,
173
- to_keys=to_keys,
174
- )
175
-
176
- return r_update
177
-
178
-
179
- class AtomDiffusion(Module):
180
- def __init__(
181
- self,
182
- score_model_args,
183
- num_sampling_steps: int = 5, # number of sampling steps
184
- sigma_min: float = 0.0004, # min noise level
185
- sigma_max: float = 160.0, # max noise level
186
- sigma_data: float = 16.0, # standard deviation of data distribution
187
- rho: float = 7, # controls the sampling schedule
188
- P_mean: float = -1.2, # mean of log-normal distribution from which noise is drawn for training
189
- P_std: float = 1.5, # standard deviation of log-normal distribution from which noise is drawn for training
190
- gamma_0: float = 0.8,
191
- gamma_min: float = 1.0,
192
- noise_scale: float = 1.003,
193
- step_scale: float = 1.5,
194
- step_scale_random: list = None,
195
- coordinate_augmentation: bool = True,
196
- coordinate_augmentation_inference=None,
197
- compile_score: bool = False,
198
- alignment_reverse_diff: bool = False,
199
- synchronize_sigmas: bool = False,
200
- ):
201
- super().__init__()
202
- self.score_model = DiffusionModule(
203
- **score_model_args,
204
- )
205
- if compile_score:
206
- self.score_model = torch.compile(
207
- self.score_model, dynamic=False, fullgraph=False
208
- )
209
-
210
- # parameters
211
- self.sigma_min = sigma_min
212
- self.sigma_max = sigma_max
213
- self.sigma_data = sigma_data
214
- self.rho = rho
215
- self.P_mean = P_mean
216
- self.P_std = P_std
217
- self.num_sampling_steps = num_sampling_steps
218
- self.gamma_0 = gamma_0
219
- self.gamma_min = gamma_min
220
- self.noise_scale = noise_scale
221
- self.step_scale = step_scale
222
- self.step_scale_random = step_scale_random
223
- self.coordinate_augmentation = coordinate_augmentation
224
- self.coordinate_augmentation_inference = (
225
- coordinate_augmentation_inference
226
- if coordinate_augmentation_inference is not None
227
- else coordinate_augmentation
228
- )
229
- self.alignment_reverse_diff = alignment_reverse_diff
230
- self.synchronize_sigmas = synchronize_sigmas
231
-
232
- self.token_s = score_model_args["token_s"]
233
- self.register_buffer("zero", torch.tensor(0.0), persistent=False)
234
-
235
- @property
236
- def device(self):
237
- return next(self.score_model.parameters()).device
238
-
239
- def c_skip(self, sigma):
240
- return (self.sigma_data**2) / (sigma**2 + self.sigma_data**2)
241
-
242
- def c_out(self, sigma):
243
- return sigma * self.sigma_data / torch.sqrt(self.sigma_data**2 + sigma**2)
244
-
245
- def c_in(self, sigma):
246
- return 1 / torch.sqrt(sigma**2 + self.sigma_data**2)
247
-
248
- def c_noise(self, sigma):
249
- return log(sigma / self.sigma_data) * 0.25
250
-
251
- def preconditioned_network_forward(
252
- self,
253
- noised_atom_coords, #: Float['b m 3'],
254
- sigma, #: Float['b'] | Float[' '] | float,
255
- network_condition_kwargs: dict,
256
- ):
257
- batch, device = noised_atom_coords.shape[0], noised_atom_coords.device
258
-
259
- if isinstance(sigma, float):
260
- sigma = torch.full((batch,), sigma, device=device)
261
-
262
- padded_sigma = rearrange(sigma, "b -> b 1 1")
263
-
264
- r_update = self.score_model(
265
- r_noisy=self.c_in(padded_sigma) * noised_atom_coords,
266
- times=self.c_noise(sigma),
267
- **network_condition_kwargs,
268
- )
269
-
270
- denoised_coords = (
271
- self.c_skip(padded_sigma) * noised_atom_coords
272
- + self.c_out(padded_sigma) * r_update
273
- )
274
- return denoised_coords
275
-
276
- def sample_schedule(self, num_sampling_steps=None):
277
- num_sampling_steps = default(num_sampling_steps, self.num_sampling_steps)
278
- inv_rho = 1 / self.rho
279
-
280
- steps = torch.arange(
281
- num_sampling_steps, device=self.device, dtype=torch.float32
282
- )
283
- sigmas = (
284
- self.sigma_max**inv_rho
285
- + steps
286
- / (num_sampling_steps - 1)
287
- * (self.sigma_min**inv_rho - self.sigma_max**inv_rho)
288
- ) ** self.rho
289
-
290
- sigmas = sigmas * self.sigma_data
291
-
292
- sigmas = F.pad(sigmas, (0, 1), value=0.0) # last step is sigma value of 0.
293
- return sigmas
294
-
295
- def sample(
296
- self,
297
- atom_mask,
298
- num_sampling_steps=None,
299
- multiplicity=1,
300
- max_parallel_samples=None,
301
- steering_args=None,
302
- **network_condition_kwargs,
303
- ):
304
- if steering_args is not None and (
305
- steering_args["fk_steering"]
306
- or steering_args["physical_guidance_update"]
307
- or steering_args["contact_guidance_update"]
308
- ):
309
- potentials = get_potentials(steering_args, boltz2=True)
310
-
311
- if steering_args["fk_steering"]:
312
- multiplicity = multiplicity * steering_args["num_particles"]
313
- energy_traj = torch.empty((multiplicity, 0), device=self.device)
314
- resample_weights = torch.ones(multiplicity, device=self.device).reshape(
315
- -1, steering_args["num_particles"]
316
- )
317
- if (
318
- steering_args["physical_guidance_update"]
319
- or steering_args["contact_guidance_update"]
320
- ):
321
- scaled_guidance_update = torch.zeros(
322
- (multiplicity, *atom_mask.shape[1:], 3),
323
- dtype=torch.float32,
324
- device=self.device,
325
- )
326
- if max_parallel_samples is None:
327
- max_parallel_samples = multiplicity
328
-
329
- num_sampling_steps = default(num_sampling_steps, self.num_sampling_steps)
330
- atom_mask = atom_mask.repeat_interleave(multiplicity, 0)
331
-
332
- shape = (*atom_mask.shape, 3)
333
-
334
- # get the schedule, which is returned as (sigma, gamma) tuple, and pair up with the next sigma and gamma
335
- sigmas = self.sample_schedule(num_sampling_steps)
336
- gammas = torch.where(sigmas > self.gamma_min, self.gamma_0, 0.0)
337
- sigmas_and_gammas = list(zip(sigmas[:-1], sigmas[1:], gammas[1:]))
338
- if self.training and self.step_scale_random is not None:
339
- step_scale = np.random.choice(self.step_scale_random)
340
- else:
341
- step_scale = self.step_scale
342
-
343
- # atom position is noise at the beginning
344
- init_sigma = sigmas[0]
345
- atom_coords = init_sigma * torch.randn(shape, device=self.device)
346
- token_repr = None
347
- atom_coords_denoised = None
348
-
349
- # gradually denoise
350
- for step_idx, (sigma_tm, sigma_t, gamma) in enumerate(sigmas_and_gammas):
351
- random_R, random_tr = compute_random_augmentation(
352
- multiplicity, device=atom_coords.device, dtype=atom_coords.dtype
353
- )
354
- atom_coords = atom_coords - atom_coords.mean(dim=-2, keepdims=True)
355
- atom_coords = (
356
- torch.einsum("bmd,bds->bms", atom_coords, random_R) + random_tr
357
- )
358
- if atom_coords_denoised is not None:
359
- atom_coords_denoised -= atom_coords_denoised.mean(dim=-2, keepdims=True)
360
- atom_coords_denoised = (
361
- torch.einsum("bmd,bds->bms", atom_coords_denoised, random_R)
362
- + random_tr
363
- )
364
- if (
365
- steering_args["physical_guidance_update"]
366
- or steering_args["contact_guidance_update"]
367
- ) and scaled_guidance_update is not None:
368
- scaled_guidance_update = torch.einsum(
369
- "bmd,bds->bms", scaled_guidance_update, random_R
370
- )
371
-
372
- sigma_tm, sigma_t, gamma = sigma_tm.item(), sigma_t.item(), gamma.item()
373
-
374
- t_hat = sigma_tm * (1 + gamma)
375
- steering_t = 1.0 - (step_idx / num_sampling_steps)
376
- noise_var = self.noise_scale**2 * (t_hat**2 - sigma_tm**2)
377
- eps = sqrt(noise_var) * torch.randn(shape, device=self.device)
378
- atom_coords_noisy = atom_coords + eps
379
-
380
- with torch.no_grad():
381
- atom_coords_denoised = torch.zeros_like(atom_coords_noisy)
382
- sample_ids = torch.arange(multiplicity).to(atom_coords_noisy.device)
383
- sample_ids_chunks = sample_ids.chunk(
384
- multiplicity % max_parallel_samples + 1
385
- )
386
-
387
- for sample_ids_chunk in sample_ids_chunks:
388
- atom_coords_denoised_chunk = self.preconditioned_network_forward(
389
- atom_coords_noisy[sample_ids_chunk],
390
- t_hat,
391
- network_condition_kwargs=dict(
392
- multiplicity=sample_ids_chunk.numel(),
393
- **network_condition_kwargs,
394
- ),
395
- )
396
- atom_coords_denoised[sample_ids_chunk] = atom_coords_denoised_chunk
397
-
398
- if steering_args["fk_steering"] and (
399
- (
400
- step_idx % steering_args["fk_resampling_interval"] == 0
401
- and noise_var > 0
402
- )
403
- or step_idx == num_sampling_steps - 1
404
- ):
405
- # Compute energy of x_0 prediction
406
- energy = torch.zeros(multiplicity, device=self.device)
407
- for potential in potentials:
408
- parameters = potential.compute_parameters(steering_t)
409
- if parameters["resampling_weight"] > 0:
410
- component_energy = potential.compute(
411
- atom_coords_denoised,
412
- network_condition_kwargs["feats"],
413
- parameters,
414
- )
415
- energy += parameters["resampling_weight"] * component_energy
416
- energy_traj = torch.cat((energy_traj, energy.unsqueeze(1)), dim=1)
417
-
418
- # Compute log G values
419
- if step_idx == 0:
420
- log_G = -1 * energy
421
- else:
422
- log_G = energy_traj[:, -2] - energy_traj[:, -1]
423
-
424
- # Compute ll difference between guided and unguided transition distribution
425
- if (
426
- steering_args["physical_guidance_update"]
427
- or steering_args["contact_guidance_update"]
428
- ) and noise_var > 0:
429
- ll_difference = (
430
- eps**2 - (eps + scaled_guidance_update) ** 2
431
- ).sum(dim=(-1, -2)) / (2 * noise_var)
432
- else:
433
- ll_difference = torch.zeros_like(energy)
434
-
435
- # Compute resampling weights
436
- resample_weights = F.softmax(
437
- (ll_difference + steering_args["fk_lambda"] * log_G).reshape(
438
- -1, steering_args["num_particles"]
439
- ),
440
- dim=1,
441
- )
442
-
443
- # Compute guidance update to x_0 prediction
444
- if (
445
- steering_args["physical_guidance_update"]
446
- or steering_args["contact_guidance_update"]
447
- ) and step_idx < num_sampling_steps - 1:
448
- guidance_update = torch.zeros_like(atom_coords_denoised)
449
- for guidance_step in range(steering_args["num_gd_steps"]):
450
- energy_gradient = torch.zeros_like(atom_coords_denoised)
451
- for potential in potentials:
452
- parameters = potential.compute_parameters(steering_t)
453
- if (
454
- parameters["guidance_weight"] > 0
455
- and (guidance_step) % parameters["guidance_interval"]
456
- == 0
457
- ):
458
- energy_gradient += parameters[
459
- "guidance_weight"
460
- ] * potential.compute_gradient(
461
- atom_coords_denoised + guidance_update,
462
- network_condition_kwargs["feats"],
463
- parameters,
464
- )
465
- guidance_update -= energy_gradient
466
- atom_coords_denoised += guidance_update
467
- scaled_guidance_update = (
468
- guidance_update
469
- * -1
470
- * self.step_scale
471
- * (sigma_t - t_hat)
472
- / t_hat
473
- )
474
-
475
- if steering_args["fk_steering"] and (
476
- (
477
- step_idx % steering_args["fk_resampling_interval"] == 0
478
- and noise_var > 0
479
- )
480
- or step_idx == num_sampling_steps - 1
481
- ):
482
- resample_indices = (
483
- torch.multinomial(
484
- resample_weights,
485
- resample_weights.shape[1]
486
- if step_idx < num_sampling_steps - 1
487
- else 1,
488
- replacement=True,
489
- )
490
- + resample_weights.shape[1]
491
- * torch.arange(
492
- resample_weights.shape[0], device=resample_weights.device
493
- ).unsqueeze(-1)
494
- ).flatten()
495
-
496
- atom_coords = atom_coords[resample_indices]
497
- atom_coords_noisy = atom_coords_noisy[resample_indices]
498
- atom_mask = atom_mask[resample_indices]
499
- if atom_coords_denoised is not None:
500
- atom_coords_denoised = atom_coords_denoised[resample_indices]
501
- energy_traj = energy_traj[resample_indices]
502
- if (
503
- steering_args["physical_guidance_update"]
504
- or steering_args["contact_guidance_update"]
505
- ):
506
- scaled_guidance_update = scaled_guidance_update[
507
- resample_indices
508
- ]
509
- if token_repr is not None:
510
- token_repr = token_repr[resample_indices]
511
-
512
- if self.alignment_reverse_diff:
513
- with torch.autocast("cuda", enabled=False):
514
- atom_coords_noisy = weighted_rigid_align(
515
- atom_coords_noisy.float(),
516
- atom_coords_denoised.float(),
517
- atom_mask.float(),
518
- atom_mask.float(),
519
- )
520
-
521
- atom_coords_noisy = atom_coords_noisy.to(atom_coords_denoised)
522
-
523
- denoised_over_sigma = (atom_coords_noisy - atom_coords_denoised) / t_hat
524
- atom_coords_next = (
525
- atom_coords_noisy + step_scale * (sigma_t - t_hat) * denoised_over_sigma
526
- )
527
-
528
- atom_coords = atom_coords_next
529
-
530
- return dict(sample_atom_coords=atom_coords, diff_token_repr=token_repr)
531
-
532
- def loss_weight(self, sigma):
533
- return (sigma**2 + self.sigma_data**2) / ((sigma * self.sigma_data) ** 2)
534
-
535
- def noise_distribution(self, batch_size):
536
- return (
537
- self.sigma_data
538
- * (
539
- self.P_mean
540
- + self.P_std * torch.randn((batch_size,), device=self.device)
541
- ).exp()
542
- )
543
-
544
- def forward(
545
- self,
546
- s_inputs,
547
- s_trunk,
548
- feats,
549
- diffusion_conditioning,
550
- multiplicity=1,
551
- ):
552
- # training diffusion step
553
- batch_size = feats["coords"].shape[0] // multiplicity
554
-
555
- if self.synchronize_sigmas:
556
- sigmas = self.noise_distribution(batch_size).repeat_interleave(
557
- multiplicity, 0
558
- )
559
- else:
560
- sigmas = self.noise_distribution(batch_size * multiplicity)
561
- padded_sigmas = rearrange(sigmas, "b -> b 1 1")
562
-
563
- atom_coords = feats["coords"]
564
-
565
- atom_mask = feats["atom_pad_mask"]
566
- atom_mask = atom_mask.repeat_interleave(multiplicity, 0)
567
-
568
- atom_coords = center_random_augmentation(
569
- atom_coords, atom_mask, augmentation=self.coordinate_augmentation
570
- )
571
-
572
- noise = torch.randn_like(atom_coords)
573
- noised_atom_coords = atom_coords + padded_sigmas * noise
574
-
575
- denoised_atom_coords = self.preconditioned_network_forward(
576
- noised_atom_coords,
577
- sigmas,
578
- network_condition_kwargs={
579
- "s_inputs": s_inputs,
580
- "s_trunk": s_trunk,
581
- "feats": feats,
582
- "multiplicity": multiplicity,
583
- "diffusion_conditioning": diffusion_conditioning,
584
- },
585
- )
586
-
587
- return {
588
- "denoised_atom_coords": denoised_atom_coords,
589
- "sigmas": sigmas,
590
- "aligned_true_atom_coords": atom_coords,
591
- }
592
-
593
- def compute_loss(
594
- self,
595
- feats,
596
- out_dict,
597
- add_smooth_lddt_loss=True,
598
- nucleotide_loss_weight=5.0,
599
- ligand_loss_weight=10.0,
600
- multiplicity=1,
601
- filter_by_plddt=0.0,
602
- ):
603
- with torch.autocast("cuda", enabled=False):
604
- denoised_atom_coords = out_dict["denoised_atom_coords"].float()
605
- sigmas = out_dict["sigmas"].float()
606
-
607
- resolved_atom_mask_uni = feats["atom_resolved_mask"].float()
608
-
609
- if filter_by_plddt > 0:
610
- plddt_mask = feats["plddt"] > filter_by_plddt
611
- resolved_atom_mask_uni = resolved_atom_mask_uni * plddt_mask.float()
612
-
613
- resolved_atom_mask = resolved_atom_mask_uni.repeat_interleave(
614
- multiplicity, 0
615
- )
616
-
617
- align_weights = denoised_atom_coords.new_ones(denoised_atom_coords.shape[:2])
618
- atom_type = (
619
- torch.bmm(
620
- feats["atom_to_token"].float(),
621
- feats["mol_type"].unsqueeze(-1).float(),
622
- )
623
- .squeeze(-1)
624
- .long()
625
- )
626
- atom_type_mult = atom_type.repeat_interleave(multiplicity, 0)
627
-
628
- align_weights = (
629
- align_weights
630
- * (
631
- 1
632
- + nucleotide_loss_weight
633
- * (
634
- torch.eq(atom_type_mult, const.chain_type_ids["DNA"]).float()
635
- + torch.eq(atom_type_mult, const.chain_type_ids["RNA"]).float()
636
- )
637
- + ligand_loss_weight
638
- * torch.eq(
639
- atom_type_mult, const.chain_type_ids["NONPOLYMER"]
640
- ).float()
641
- ).float()
642
- )
643
-
644
- atom_coords = out_dict["aligned_true_atom_coords"].float()
645
- atom_coords_aligned_ground_truth = weighted_rigid_align(
646
- atom_coords.detach(),
647
- denoised_atom_coords.detach(),
648
- align_weights.detach(),
649
- mask=feats["atom_resolved_mask"]
650
- .float()
651
- .repeat_interleave(multiplicity, 0)
652
- .detach(),
653
- )
654
-
655
- # Cast back
656
- atom_coords_aligned_ground_truth = atom_coords_aligned_ground_truth.to(
657
- denoised_atom_coords
658
- )
659
-
660
- # weighted MSE loss of denoised atom positions
661
- mse_loss = (
662
- (denoised_atom_coords - atom_coords_aligned_ground_truth) ** 2
663
- ).sum(dim=-1)
664
- mse_loss = torch.sum(
665
- mse_loss * align_weights * resolved_atom_mask, dim=-1
666
- ) / (torch.sum(3 * align_weights * resolved_atom_mask, dim=-1) + 1e-5)
667
-
668
- # weight by sigma factor
669
- loss_weights = self.loss_weight(sigmas)
670
- mse_loss = (mse_loss * loss_weights).mean()
671
-
672
- total_loss = mse_loss
673
-
674
- # proposed auxiliary smooth lddt loss
675
- lddt_loss = self.zero
676
- if add_smooth_lddt_loss:
677
- lddt_loss = smooth_lddt_loss(
678
- denoised_atom_coords,
679
- feats["coords"],
680
- torch.eq(atom_type, const.chain_type_ids["DNA"]).float()
681
- + torch.eq(atom_type, const.chain_type_ids["RNA"]).float(),
682
- coords_mask=resolved_atom_mask_uni,
683
- multiplicity=multiplicity,
684
- )
685
-
686
- total_loss = total_loss + lddt_loss
687
-
688
- loss_breakdown = {
689
- "mse_loss": mse_loss,
690
- "smooth_lddt_loss": lddt_loss,
691
- }
692
-
693
- return {"loss": total_loss, "loss_breakdown": loss_breakdown}
 
1
+ # started from code from https://github.com/lucidrains/alphafold3-pytorch, MIT License, Copyright (c) 2024 Phil Wang
2
+
3
+ from __future__ import annotations
4
+
5
+ from math import sqrt
6
+
7
+ import numpy as np
8
+ import torch
9
+ import torch.nn.functional as F # noqa: N812
10
+ from einops import rearrange
11
+ from torch import nn
12
+ from torch.nn import Module
13
+
14
+ from . import vb_const as const
15
+ from . import vb_layers_initialize as init
16
+ from .vb_loss_diffusionv2 import (
17
+ smooth_lddt_loss,
18
+ weighted_rigid_align,
19
+ )
20
+ from .vb_modules_encodersv2 import (
21
+ AtomAttentionDecoder,
22
+ AtomAttentionEncoder,
23
+ SingleConditioning,
24
+ )
25
+ from .vb_modules_transformersv2 import (
26
+ DiffusionTransformer,
27
+ )
28
+ from .vb_modules_utils import (
29
+ LinearNoBias,
30
+ center_random_augmentation,
31
+ compute_random_augmentation,
32
+ default,
33
+ log,
34
+ )
35
+ from .vb_potentials_potentials import get_potentials
36
+
37
+
38
+ class DiffusionModule(Module):
39
+ """Diffusion module"""
40
+
41
+ def __init__(
42
+ self,
43
+ token_s: int,
44
+ atom_s: int,
45
+ atoms_per_window_queries: int = 32,
46
+ atoms_per_window_keys: int = 128,
47
+ sigma_data: int = 16,
48
+ dim_fourier: int = 256,
49
+ atom_encoder_depth: int = 3,
50
+ atom_encoder_heads: int = 4,
51
+ token_transformer_depth: int = 24,
52
+ token_transformer_heads: int = 8,
53
+ atom_decoder_depth: int = 3,
54
+ atom_decoder_heads: int = 4,
55
+ conditioning_transition_layers: int = 2,
56
+ activation_checkpointing: bool = False,
57
+ transformer_post_ln: bool = False,
58
+ ) -> None:
59
+ super().__init__()
60
+
61
+ self.atoms_per_window_queries = atoms_per_window_queries
62
+ self.atoms_per_window_keys = atoms_per_window_keys
63
+ self.sigma_data = sigma_data
64
+ self.activation_checkpointing = activation_checkpointing
65
+
66
+ # conditioning
67
+ self.single_conditioner = SingleConditioning(
68
+ sigma_data=sigma_data,
69
+ token_s=token_s,
70
+ dim_fourier=dim_fourier,
71
+ num_transitions=conditioning_transition_layers,
72
+ )
73
+
74
+ self.atom_attention_encoder = AtomAttentionEncoder(
75
+ atom_s=atom_s,
76
+ token_s=token_s,
77
+ atoms_per_window_queries=atoms_per_window_queries,
78
+ atoms_per_window_keys=atoms_per_window_keys,
79
+ atom_encoder_depth=atom_encoder_depth,
80
+ atom_encoder_heads=atom_encoder_heads,
81
+ structure_prediction=True,
82
+ activation_checkpointing=activation_checkpointing,
83
+ transformer_post_layer_norm=transformer_post_ln,
84
+ )
85
+
86
+ self.s_to_a_linear = nn.Sequential(
87
+ nn.LayerNorm(2 * token_s), LinearNoBias(2 * token_s, 2 * token_s)
88
+ )
89
+ init.final_init_(self.s_to_a_linear[1].weight)
90
+
91
+ self.token_transformer = DiffusionTransformer(
92
+ dim=2 * token_s,
93
+ dim_single_cond=2 * token_s,
94
+ depth=token_transformer_depth,
95
+ heads=token_transformer_heads,
96
+ activation_checkpointing=activation_checkpointing,
97
+ # post_layer_norm=transformer_post_ln,
98
+ )
99
+
100
+ self.a_norm = nn.LayerNorm(
101
+ 2 * token_s
102
+ ) # if not transformer_post_ln else nn.Identity()
103
+
104
+ self.atom_attention_decoder = AtomAttentionDecoder(
105
+ atom_s=atom_s,
106
+ token_s=token_s,
107
+ attn_window_queries=atoms_per_window_queries,
108
+ attn_window_keys=atoms_per_window_keys,
109
+ atom_decoder_depth=atom_decoder_depth,
110
+ atom_decoder_heads=atom_decoder_heads,
111
+ activation_checkpointing=activation_checkpointing,
112
+ # transformer_post_layer_norm=transformer_post_ln,
113
+ )
114
+
115
+ def forward(
116
+ self,
117
+ s_inputs, # Float['b n ts']
118
+ s_trunk, # Float['b n ts']
119
+ r_noisy, # Float['bm m 3']
120
+ times, # Float['bm 1 1']
121
+ feats,
122
+ diffusion_conditioning,
123
+ multiplicity=1,
124
+ ):
125
+ if self.activation_checkpointing and self.training:
126
+ s, normed_fourier = torch.utils.checkpoint.checkpoint(
127
+ self.single_conditioner,
128
+ times,
129
+ s_trunk.repeat_interleave(multiplicity, 0),
130
+ s_inputs.repeat_interleave(multiplicity, 0),
131
+ )
132
+ else:
133
+ s, normed_fourier = self.single_conditioner(
134
+ times,
135
+ s_trunk.repeat_interleave(multiplicity, 0),
136
+ s_inputs.repeat_interleave(multiplicity, 0),
137
+ )
138
+
139
+ # Sequence-local Atom Attention and aggregation to coarse-grained tokens
140
+ a, q_skip, c_skip, to_keys = self.atom_attention_encoder(
141
+ feats=feats,
142
+ q=diffusion_conditioning["q"].float(),
143
+ c=diffusion_conditioning["c"].float(),
144
+ atom_enc_bias=diffusion_conditioning["atom_enc_bias"].float(),
145
+ to_keys=diffusion_conditioning["to_keys"],
146
+ r=r_noisy, # Float['b m 3'],
147
+ multiplicity=multiplicity,
148
+ )
149
+
150
+ # Full self-attention on token level
151
+ a = a + self.s_to_a_linear(s)
152
+
153
+ mask = feats["token_pad_mask"].repeat_interleave(multiplicity, 0)
154
+ a = self.token_transformer(
155
+ a,
156
+ mask=mask.float(),
157
+ s=s,
158
+ bias=diffusion_conditioning[
159
+ "token_trans_bias"
160
+ ].float(), # note z is not expanded with multiplicity until after bias is computed
161
+ multiplicity=multiplicity,
162
+ )
163
+ a = self.a_norm(a)
164
+
165
+ # Broadcast token activations to atoms and run Sequence-local Atom Attention
166
+ r_update = self.atom_attention_decoder(
167
+ a=a,
168
+ q=q_skip,
169
+ c=c_skip,
170
+ atom_dec_bias=diffusion_conditioning["atom_dec_bias"].float(),
171
+ feats=feats,
172
+ multiplicity=multiplicity,
173
+ to_keys=to_keys,
174
+ )
175
+
176
+ return r_update
177
+
178
+
179
+ class AtomDiffusion(Module):
180
+ def __init__(
181
+ self,
182
+ score_model_args,
183
+ num_sampling_steps: int = 5, # number of sampling steps
184
+ sigma_min: float = 0.0004, # min noise level
185
+ sigma_max: float = 160.0, # max noise level
186
+ sigma_data: float = 16.0, # standard deviation of data distribution
187
+ rho: float = 7, # controls the sampling schedule
188
+ P_mean: float = -1.2, # mean of log-normal distribution from which noise is drawn for training
189
+ P_std: float = 1.5, # standard deviation of log-normal distribution from which noise is drawn for training
190
+ gamma_0: float = 0.8,
191
+ gamma_min: float = 1.0,
192
+ noise_scale: float = 1.003,
193
+ step_scale: float = 1.5,
194
+ step_scale_random: list = None,
195
+ coordinate_augmentation: bool = True,
196
+ coordinate_augmentation_inference=None,
197
+ compile_score: bool = False,
198
+ alignment_reverse_diff: bool = False,
199
+ synchronize_sigmas: bool = False,
200
+ ):
201
+ super().__init__()
202
+ self.score_model = DiffusionModule(
203
+ **score_model_args,
204
+ )
205
+ if compile_score:
206
+ self.score_model = torch.compile(
207
+ self.score_model, dynamic=False, fullgraph=False
208
+ )
209
+
210
+ # parameters
211
+ self.sigma_min = sigma_min
212
+ self.sigma_max = sigma_max
213
+ self.sigma_data = sigma_data
214
+ self.rho = rho
215
+ self.P_mean = P_mean
216
+ self.P_std = P_std
217
+ self.num_sampling_steps = num_sampling_steps
218
+ self.gamma_0 = gamma_0
219
+ self.gamma_min = gamma_min
220
+ self.noise_scale = noise_scale
221
+ self.step_scale = step_scale
222
+ self.step_scale_random = step_scale_random
223
+ self.coordinate_augmentation = coordinate_augmentation
224
+ self.coordinate_augmentation_inference = (
225
+ coordinate_augmentation_inference
226
+ if coordinate_augmentation_inference is not None
227
+ else coordinate_augmentation
228
+ )
229
+ self.alignment_reverse_diff = alignment_reverse_diff
230
+ self.synchronize_sigmas = synchronize_sigmas
231
+
232
+ self.token_s = score_model_args["token_s"]
233
+ self.register_buffer("zero", torch.tensor(0.0), persistent=False)
234
+
235
+ @property
236
+ def device(self):
237
+ return next(self.score_model.parameters()).device
238
+
239
+ def c_skip(self, sigma):
240
+ return (self.sigma_data**2) / (sigma**2 + self.sigma_data**2)
241
+
242
+ def c_out(self, sigma):
243
+ return sigma * self.sigma_data / torch.sqrt(self.sigma_data**2 + sigma**2)
244
+
245
+ def c_in(self, sigma):
246
+ return 1 / torch.sqrt(sigma**2 + self.sigma_data**2)
247
+
248
+ def c_noise(self, sigma):
249
+ return log(sigma / self.sigma_data) * 0.25
250
+
251
+ def preconditioned_network_forward(
252
+ self,
253
+ noised_atom_coords, #: Float['b m 3'],
254
+ sigma, #: Float['b'] | Float[' '] | float,
255
+ network_condition_kwargs: dict,
256
+ ):
257
+ batch, device = noised_atom_coords.shape[0], noised_atom_coords.device
258
+
259
+ if isinstance(sigma, float):
260
+ sigma = torch.full((batch,), sigma, device=device)
261
+
262
+ padded_sigma = rearrange(sigma, "b -> b 1 1")
263
+
264
+ r_update = self.score_model(
265
+ r_noisy=self.c_in(padded_sigma) * noised_atom_coords,
266
+ times=self.c_noise(sigma),
267
+ **network_condition_kwargs,
268
+ )
269
+
270
+ denoised_coords = (
271
+ self.c_skip(padded_sigma) * noised_atom_coords
272
+ + self.c_out(padded_sigma) * r_update
273
+ )
274
+ return denoised_coords
275
+
276
+ def sample_schedule(self, num_sampling_steps=None):
277
+ num_sampling_steps = default(num_sampling_steps, self.num_sampling_steps)
278
+ inv_rho = 1 / self.rho
279
+
280
+ steps = torch.arange(
281
+ num_sampling_steps, device=self.device, dtype=torch.float32
282
+ )
283
+ sigmas = (
284
+ self.sigma_max**inv_rho
285
+ + steps
286
+ / (num_sampling_steps - 1)
287
+ * (self.sigma_min**inv_rho - self.sigma_max**inv_rho)
288
+ ) ** self.rho
289
+
290
+ sigmas = sigmas * self.sigma_data
291
+
292
+ sigmas = F.pad(sigmas, (0, 1), value=0.0) # last step is sigma value of 0.
293
+ return sigmas
294
+
295
+ def sample(
296
+ self,
297
+ atom_mask,
298
+ num_sampling_steps=None,
299
+ multiplicity=1,
300
+ max_parallel_samples=None,
301
+ steering_args=None,
302
+ **network_condition_kwargs,
303
+ ):
304
+ if steering_args is not None and (
305
+ steering_args["fk_steering"]
306
+ or steering_args["physical_guidance_update"]
307
+ or steering_args["contact_guidance_update"]
308
+ ):
309
+ potentials = get_potentials(steering_args, boltz2=True)
310
+
311
+ if steering_args["fk_steering"]:
312
+ multiplicity = multiplicity * steering_args["num_particles"]
313
+ energy_traj = torch.empty((multiplicity, 0), device=self.device)
314
+ resample_weights = torch.ones(multiplicity, device=self.device).reshape(
315
+ -1, steering_args["num_particles"]
316
+ )
317
+ if (
318
+ steering_args["physical_guidance_update"]
319
+ or steering_args["contact_guidance_update"]
320
+ ):
321
+ scaled_guidance_update = torch.zeros(
322
+ (multiplicity, *atom_mask.shape[1:], 3),
323
+ dtype=torch.float32,
324
+ device=self.device,
325
+ )
326
+ if max_parallel_samples is None:
327
+ max_parallel_samples = multiplicity
328
+
329
+ num_sampling_steps = default(num_sampling_steps, self.num_sampling_steps)
330
+ atom_mask = atom_mask.repeat_interleave(multiplicity, 0)
331
+
332
+ shape = (*atom_mask.shape, 3)
333
+
334
+ # get the schedule, which is returned as (sigma, gamma) tuple, and pair up with the next sigma and gamma
335
+ sigmas = self.sample_schedule(num_sampling_steps)
336
+ gammas = torch.where(sigmas > self.gamma_min, self.gamma_0, 0.0)
337
+ sigmas_and_gammas = list(zip(sigmas[:-1], sigmas[1:], gammas[1:]))
338
+ if self.training and self.step_scale_random is not None:
339
+ step_scale = np.random.choice(self.step_scale_random)
340
+ else:
341
+ step_scale = self.step_scale
342
+
343
+ # atom position is noise at the beginning
344
+ init_sigma = sigmas[0]
345
+ atom_coords = init_sigma * torch.randn(shape, device=self.device)
346
+ token_repr = None
347
+ atom_coords_denoised = None
348
+
349
+ # gradually denoise
350
+ for step_idx, (sigma_tm, sigma_t, gamma) in enumerate(sigmas_and_gammas):
351
+ random_R, random_tr = compute_random_augmentation(
352
+ multiplicity, device=atom_coords.device, dtype=atom_coords.dtype
353
+ )
354
+ atom_coords = atom_coords - atom_coords.mean(dim=-2, keepdims=True)
355
+ atom_coords = (
356
+ torch.einsum("bmd,bds->bms", atom_coords, random_R) + random_tr
357
+ )
358
+ if atom_coords_denoised is not None:
359
+ atom_coords_denoised -= atom_coords_denoised.mean(dim=-2, keepdims=True)
360
+ atom_coords_denoised = (
361
+ torch.einsum("bmd,bds->bms", atom_coords_denoised, random_R)
362
+ + random_tr
363
+ )
364
+ if (
365
+ steering_args["physical_guidance_update"]
366
+ or steering_args["contact_guidance_update"]
367
+ ) and scaled_guidance_update is not None:
368
+ scaled_guidance_update = torch.einsum(
369
+ "bmd,bds->bms", scaled_guidance_update, random_R
370
+ )
371
+
372
+ sigma_tm, sigma_t, gamma = sigma_tm.item(), sigma_t.item(), gamma.item()
373
+
374
+ t_hat = sigma_tm * (1 + gamma)
375
+ steering_t = 1.0 - (step_idx / num_sampling_steps)
376
+ noise_var = self.noise_scale**2 * (t_hat**2 - sigma_tm**2)
377
+ eps = sqrt(noise_var) * torch.randn(shape, device=self.device)
378
+ atom_coords_noisy = atom_coords + eps
379
+
380
+ with torch.no_grad():
381
+ atom_coords_denoised = torch.zeros_like(atom_coords_noisy)
382
+ sample_ids = torch.arange(multiplicity).to(atom_coords_noisy.device)
383
+ sample_ids_chunks = sample_ids.chunk(
384
+ multiplicity % max_parallel_samples + 1
385
+ )
386
+
387
+ for sample_ids_chunk in sample_ids_chunks:
388
+ atom_coords_denoised_chunk = self.preconditioned_network_forward(
389
+ atom_coords_noisy[sample_ids_chunk],
390
+ t_hat,
391
+ network_condition_kwargs=dict(
392
+ multiplicity=sample_ids_chunk.numel(),
393
+ **network_condition_kwargs,
394
+ ),
395
+ )
396
+ atom_coords_denoised[sample_ids_chunk] = atom_coords_denoised_chunk
397
+
398
+ if steering_args["fk_steering"] and (
399
+ (
400
+ step_idx % steering_args["fk_resampling_interval"] == 0
401
+ and noise_var > 0
402
+ )
403
+ or step_idx == num_sampling_steps - 1
404
+ ):
405
+ # Compute energy of x_0 prediction
406
+ energy = torch.zeros(multiplicity, device=self.device)
407
+ for potential in potentials:
408
+ parameters = potential.compute_parameters(steering_t)
409
+ if parameters["resampling_weight"] > 0:
410
+ component_energy = potential.compute(
411
+ atom_coords_denoised,
412
+ network_condition_kwargs["feats"],
413
+ parameters,
414
+ )
415
+ energy += parameters["resampling_weight"] * component_energy
416
+ energy_traj = torch.cat((energy_traj, energy.unsqueeze(1)), dim=1)
417
+
418
+ # Compute log G values
419
+ if step_idx == 0:
420
+ log_G = -1 * energy
421
+ else:
422
+ log_G = energy_traj[:, -2] - energy_traj[:, -1]
423
+
424
+ # Compute ll difference between guided and unguided transition distribution
425
+ if (
426
+ steering_args["physical_guidance_update"]
427
+ or steering_args["contact_guidance_update"]
428
+ ) and noise_var > 0:
429
+ ll_difference = (
430
+ eps**2 - (eps + scaled_guidance_update) ** 2
431
+ ).sum(dim=(-1, -2)) / (2 * noise_var)
432
+ else:
433
+ ll_difference = torch.zeros_like(energy)
434
+
435
+ # Compute resampling weights
436
+ resample_weights = F.softmax(
437
+ (ll_difference + steering_args["fk_lambda"] * log_G).reshape(
438
+ -1, steering_args["num_particles"]
439
+ ),
440
+ dim=1,
441
+ )
442
+
443
+ # Compute guidance update to x_0 prediction
444
+ if (
445
+ steering_args["physical_guidance_update"]
446
+ or steering_args["contact_guidance_update"]
447
+ ) and step_idx < num_sampling_steps - 1:
448
+ guidance_update = torch.zeros_like(atom_coords_denoised)
449
+ for guidance_step in range(steering_args["num_gd_steps"]):
450
+ energy_gradient = torch.zeros_like(atom_coords_denoised)
451
+ for potential in potentials:
452
+ parameters = potential.compute_parameters(steering_t)
453
+ if (
454
+ parameters["guidance_weight"] > 0
455
+ and (guidance_step) % parameters["guidance_interval"]
456
+ == 0
457
+ ):
458
+ energy_gradient += parameters[
459
+ "guidance_weight"
460
+ ] * potential.compute_gradient(
461
+ atom_coords_denoised + guidance_update,
462
+ network_condition_kwargs["feats"],
463
+ parameters,
464
+ )
465
+ guidance_update -= energy_gradient
466
+ atom_coords_denoised += guidance_update
467
+ scaled_guidance_update = (
468
+ guidance_update
469
+ * -1
470
+ * self.step_scale
471
+ * (sigma_t - t_hat)
472
+ / t_hat
473
+ )
474
+
475
+ if steering_args["fk_steering"] and (
476
+ (
477
+ step_idx % steering_args["fk_resampling_interval"] == 0
478
+ and noise_var > 0
479
+ )
480
+ or step_idx == num_sampling_steps - 1
481
+ ):
482
+ resample_indices = (
483
+ torch.multinomial(
484
+ resample_weights,
485
+ resample_weights.shape[1]
486
+ if step_idx < num_sampling_steps - 1
487
+ else 1,
488
+ replacement=True,
489
+ )
490
+ + resample_weights.shape[1]
491
+ * torch.arange(
492
+ resample_weights.shape[0], device=resample_weights.device
493
+ ).unsqueeze(-1)
494
+ ).flatten()
495
+
496
+ atom_coords = atom_coords[resample_indices]
497
+ atom_coords_noisy = atom_coords_noisy[resample_indices]
498
+ atom_mask = atom_mask[resample_indices]
499
+ if atom_coords_denoised is not None:
500
+ atom_coords_denoised = atom_coords_denoised[resample_indices]
501
+ energy_traj = energy_traj[resample_indices]
502
+ if (
503
+ steering_args["physical_guidance_update"]
504
+ or steering_args["contact_guidance_update"]
505
+ ):
506
+ scaled_guidance_update = scaled_guidance_update[
507
+ resample_indices
508
+ ]
509
+ if token_repr is not None:
510
+ token_repr = token_repr[resample_indices]
511
+
512
+ if self.alignment_reverse_diff:
513
+ with torch.autocast("cuda", enabled=False):
514
+ atom_coords_noisy = weighted_rigid_align(
515
+ atom_coords_noisy.float(),
516
+ atom_coords_denoised.float(),
517
+ atom_mask.float(),
518
+ atom_mask.float(),
519
+ )
520
+
521
+ atom_coords_noisy = atom_coords_noisy.to(atom_coords_denoised)
522
+
523
+ denoised_over_sigma = (atom_coords_noisy - atom_coords_denoised) / t_hat
524
+ atom_coords_next = (
525
+ atom_coords_noisy + step_scale * (sigma_t - t_hat) * denoised_over_sigma
526
+ )
527
+
528
+ atom_coords = atom_coords_next
529
+
530
+ return dict(sample_atom_coords=atom_coords, diff_token_repr=token_repr)
531
+
532
+ def loss_weight(self, sigma):
533
+ return (sigma**2 + self.sigma_data**2) / ((sigma * self.sigma_data) ** 2)
534
+
535
+ def noise_distribution(self, batch_size):
536
+ return (
537
+ self.sigma_data
538
+ * (
539
+ self.P_mean
540
+ + self.P_std * torch.randn((batch_size,), device=self.device)
541
+ ).exp()
542
+ )
543
+
544
+ def forward(
545
+ self,
546
+ s_inputs,
547
+ s_trunk,
548
+ feats,
549
+ diffusion_conditioning,
550
+ multiplicity=1,
551
+ ):
552
+ # training diffusion step
553
+ batch_size = feats["coords"].shape[0] // multiplicity
554
+
555
+ if self.synchronize_sigmas:
556
+ sigmas = self.noise_distribution(batch_size).repeat_interleave(
557
+ multiplicity, 0
558
+ )
559
+ else:
560
+ sigmas = self.noise_distribution(batch_size * multiplicity)
561
+ padded_sigmas = rearrange(sigmas, "b -> b 1 1")
562
+
563
+ atom_coords = feats["coords"]
564
+
565
+ atom_mask = feats["atom_pad_mask"]
566
+ atom_mask = atom_mask.repeat_interleave(multiplicity, 0)
567
+
568
+ atom_coords = center_random_augmentation(
569
+ atom_coords, atom_mask, augmentation=self.coordinate_augmentation
570
+ )
571
+
572
+ noise = torch.randn_like(atom_coords)
573
+ noised_atom_coords = atom_coords + padded_sigmas * noise
574
+
575
+ denoised_atom_coords = self.preconditioned_network_forward(
576
+ noised_atom_coords,
577
+ sigmas,
578
+ network_condition_kwargs={
579
+ "s_inputs": s_inputs,
580
+ "s_trunk": s_trunk,
581
+ "feats": feats,
582
+ "multiplicity": multiplicity,
583
+ "diffusion_conditioning": diffusion_conditioning,
584
+ },
585
+ )
586
+
587
+ return {
588
+ "denoised_atom_coords": denoised_atom_coords,
589
+ "sigmas": sigmas,
590
+ "aligned_true_atom_coords": atom_coords,
591
+ }
592
+
593
+ def compute_loss(
594
+ self,
595
+ feats,
596
+ out_dict,
597
+ add_smooth_lddt_loss=True,
598
+ nucleotide_loss_weight=5.0,
599
+ ligand_loss_weight=10.0,
600
+ multiplicity=1,
601
+ filter_by_plddt=0.0,
602
+ ):
603
+ with torch.autocast("cuda", enabled=False):
604
+ denoised_atom_coords = out_dict["denoised_atom_coords"].float()
605
+ sigmas = out_dict["sigmas"].float()
606
+
607
+ resolved_atom_mask_uni = feats["atom_resolved_mask"].float()
608
+
609
+ if filter_by_plddt > 0:
610
+ plddt_mask = feats["plddt"] > filter_by_plddt
611
+ resolved_atom_mask_uni = resolved_atom_mask_uni * plddt_mask.float()
612
+
613
+ resolved_atom_mask = resolved_atom_mask_uni.repeat_interleave(
614
+ multiplicity, 0
615
+ )
616
+
617
+ align_weights = denoised_atom_coords.new_ones(denoised_atom_coords.shape[:2])
618
+ atom_type = (
619
+ torch.bmm(
620
+ feats["atom_to_token"].float(),
621
+ feats["mol_type"].unsqueeze(-1).float(),
622
+ )
623
+ .squeeze(-1)
624
+ .long()
625
+ )
626
+ atom_type_mult = atom_type.repeat_interleave(multiplicity, 0)
627
+
628
+ align_weights = (
629
+ align_weights
630
+ * (
631
+ 1
632
+ + nucleotide_loss_weight
633
+ * (
634
+ torch.eq(atom_type_mult, const.chain_type_ids["DNA"]).float()
635
+ + torch.eq(atom_type_mult, const.chain_type_ids["RNA"]).float()
636
+ )
637
+ + ligand_loss_weight
638
+ * torch.eq(
639
+ atom_type_mult, const.chain_type_ids["NONPOLYMER"]
640
+ ).float()
641
+ ).float()
642
+ )
643
+
644
+ atom_coords = out_dict["aligned_true_atom_coords"].float()
645
+ atom_coords_aligned_ground_truth = weighted_rigid_align(
646
+ atom_coords.detach(),
647
+ denoised_atom_coords.detach(),
648
+ align_weights.detach(),
649
+ mask=feats["atom_resolved_mask"]
650
+ .float()
651
+ .repeat_interleave(multiplicity, 0)
652
+ .detach(),
653
+ )
654
+
655
+ # Cast back
656
+ atom_coords_aligned_ground_truth = atom_coords_aligned_ground_truth.to(
657
+ denoised_atom_coords
658
+ )
659
+
660
+ # weighted MSE loss of denoised atom positions
661
+ mse_loss = (
662
+ (denoised_atom_coords - atom_coords_aligned_ground_truth) ** 2
663
+ ).sum(dim=-1)
664
+ mse_loss = torch.sum(
665
+ mse_loss * align_weights * resolved_atom_mask, dim=-1
666
+ ) / (torch.sum(3 * align_weights * resolved_atom_mask, dim=-1) + 1e-5)
667
+
668
+ # weight by sigma factor
669
+ loss_weights = self.loss_weight(sigmas)
670
+ mse_loss = (mse_loss * loss_weights).mean()
671
+
672
+ total_loss = mse_loss
673
+
674
+ # proposed auxiliary smooth lddt loss
675
+ lddt_loss = self.zero
676
+ if add_smooth_lddt_loss:
677
+ lddt_loss = smooth_lddt_loss(
678
+ denoised_atom_coords,
679
+ feats["coords"],
680
+ torch.eq(atom_type, const.chain_type_ids["DNA"]).float()
681
+ + torch.eq(atom_type, const.chain_type_ids["RNA"]).float(),
682
+ coords_mask=resolved_atom_mask_uni,
683
+ multiplicity=multiplicity,
684
+ )
685
+
686
+ total_loss = total_loss + lddt_loss
687
+
688
+ loss_breakdown = {
689
+ "mse_loss": mse_loss,
690
+ "smooth_lddt_loss": lddt_loss,
691
+ }
692
+
693
+ return {"loss": total_loss, "loss_breakdown": loss_breakdown}
vb_modules_encodersv2.py CHANGED
@@ -1,565 +1,565 @@
1
- # started from code from https://github.com/lucidrains/alphafold3-pytorch, MIT License, Copyright (c) 2024 Phil Wang
2
- from functools import partial
3
- from math import pi
4
-
5
- import torch
6
- from einops import rearrange
7
- from torch import nn
8
- from torch.nn import Linear, Module, ModuleList
9
- from torch.nn.functional import one_hot
10
-
11
- from . import vb_layers_initialize as init
12
- from .vb_layers_transition import Transition
13
- from .vb_modules_transformersv2 import AtomTransformer
14
- from .vb_modules_utils import LinearNoBias
15
-
16
-
17
- class FourierEmbedding(Module):
18
- """Algorithm 22."""
19
-
20
- def __init__(self, dim):
21
- super().__init__()
22
- self.proj = nn.Linear(1, dim)
23
- torch.nn.init.normal_(self.proj.weight, mean=0, std=1)
24
- torch.nn.init.normal_(self.proj.bias, mean=0, std=1)
25
- self.proj.requires_grad_(False)
26
-
27
- def forward(
28
- self,
29
- times, # Float[' b'],
30
- ): # -> Float['b d']:
31
- times = rearrange(times, "b -> b 1")
32
- rand_proj = self.proj(times)
33
- return torch.cos(2 * pi * rand_proj)
34
-
35
-
36
- class RelativePositionEncoder(Module):
37
- """Algorithm 3."""
38
-
39
- def __init__(
40
- self, token_z, r_max=32, s_max=2, fix_sym_check=False, cyclic_pos_enc=False
41
- ):
42
- super().__init__()
43
- self.r_max = r_max
44
- self.s_max = s_max
45
- self.linear_layer = LinearNoBias(4 * (r_max + 1) + 2 * (s_max + 1) + 1, token_z)
46
- self.fix_sym_check = fix_sym_check
47
- self.cyclic_pos_enc = cyclic_pos_enc
48
-
49
- def forward(self, feats):
50
- b_same_chain = torch.eq(
51
- feats["asym_id"][:, :, None], feats["asym_id"][:, None, :]
52
- )
53
- b_same_residue = torch.eq(
54
- feats["residue_index"][:, :, None], feats["residue_index"][:, None, :]
55
- )
56
- b_same_entity = torch.eq(
57
- feats["entity_id"][:, :, None], feats["entity_id"][:, None, :]
58
- )
59
-
60
- d_residue = (
61
- feats["residue_index"][:, :, None] - feats["residue_index"][:, None, :]
62
- )
63
-
64
- if self.cyclic_pos_enc and torch.any(feats["cyclic_period"] > 0):
65
- period = torch.where(
66
- feats["cyclic_period"] > 0,
67
- feats["cyclic_period"],
68
- torch.zeros_like(feats["cyclic_period"]) + 10000,
69
- )
70
- d_residue = (d_residue - period * torch.round(d_residue / period)).long()
71
-
72
- d_residue = torch.clip(
73
- d_residue + self.r_max,
74
- 0,
75
- 2 * self.r_max,
76
- )
77
- d_residue = torch.where(
78
- b_same_chain, d_residue, torch.zeros_like(d_residue) + 2 * self.r_max + 1
79
- )
80
- a_rel_pos = one_hot(d_residue, 2 * self.r_max + 2)
81
-
82
- d_token = torch.clip(
83
- feats["token_index"][:, :, None]
84
- - feats["token_index"][:, None, :]
85
- + self.r_max,
86
- 0,
87
- 2 * self.r_max,
88
- )
89
- d_token = torch.where(
90
- b_same_chain & b_same_residue,
91
- d_token,
92
- torch.zeros_like(d_token) + 2 * self.r_max + 1,
93
- )
94
- a_rel_token = one_hot(d_token, 2 * self.r_max + 2)
95
-
96
- d_chain = torch.clip(
97
- feats["sym_id"][:, :, None] - feats["sym_id"][:, None, :] + self.s_max,
98
- 0,
99
- 2 * self.s_max,
100
- )
101
- d_chain = torch.where(
102
- (~b_same_entity) if self.fix_sym_check else b_same_chain,
103
- torch.zeros_like(d_chain) + 2 * self.s_max + 1,
104
- d_chain,
105
- )
106
- # Note: added | (~b_same_entity) based on observation of ProteinX manuscript
107
- a_rel_chain = one_hot(d_chain, 2 * self.s_max + 2)
108
-
109
- p = self.linear_layer(
110
- torch.cat(
111
- [
112
- a_rel_pos.float(),
113
- a_rel_token.float(),
114
- b_same_entity.unsqueeze(-1).float(),
115
- a_rel_chain.float(),
116
- ],
117
- dim=-1,
118
- )
119
- )
120
- return p
121
-
122
-
123
- class SingleConditioning(Module):
124
- """Algorithm 21."""
125
-
126
- def __init__(
127
- self,
128
- sigma_data: float,
129
- token_s: int = 384,
130
- dim_fourier: int = 256,
131
- num_transitions: int = 2,
132
- transition_expansion_factor: int = 2,
133
- eps: float = 1e-20,
134
- disable_times: bool = False,
135
- ) -> None:
136
- super().__init__()
137
- self.eps = eps
138
- self.sigma_data = sigma_data
139
- self.disable_times = disable_times
140
-
141
- self.norm_single = nn.LayerNorm(2 * token_s)
142
- self.single_embed = nn.Linear(2 * token_s, 2 * token_s)
143
- if not self.disable_times:
144
- self.fourier_embed = FourierEmbedding(dim_fourier)
145
- self.norm_fourier = nn.LayerNorm(dim_fourier)
146
- self.fourier_to_single = LinearNoBias(dim_fourier, 2 * token_s)
147
-
148
- transitions = ModuleList([])
149
- for _ in range(num_transitions):
150
- transition = Transition(
151
- dim=2 * token_s, hidden=transition_expansion_factor * 2 * token_s
152
- )
153
- transitions.append(transition)
154
-
155
- self.transitions = transitions
156
-
157
- def forward(
158
- self,
159
- times, # Float[' b'],
160
- s_trunk, # Float['b n ts'],
161
- s_inputs, # Float['b n ts'],
162
- ): # -> Float['b n 2ts']:
163
- s = torch.cat((s_trunk, s_inputs), dim=-1)
164
- s = self.single_embed(self.norm_single(s))
165
- if not self.disable_times:
166
- fourier_embed = self.fourier_embed(
167
- times
168
- ) # note: sigma rescaling done in diffusion module
169
- normed_fourier = self.norm_fourier(fourier_embed)
170
- fourier_to_single = self.fourier_to_single(normed_fourier)
171
-
172
- s = rearrange(fourier_to_single, "b d -> b 1 d") + s
173
-
174
- for transition in self.transitions:
175
- s = transition(s) + s
176
-
177
- return s, normed_fourier if not self.disable_times else None
178
-
179
-
180
- class PairwiseConditioning(Module):
181
- """Algorithm 21."""
182
-
183
- def __init__(
184
- self,
185
- token_z,
186
- dim_token_rel_pos_feats,
187
- num_transitions=2,
188
- transition_expansion_factor=2,
189
- ):
190
- super().__init__()
191
-
192
- self.dim_pairwise_init_proj = nn.Sequential(
193
- nn.LayerNorm(token_z + dim_token_rel_pos_feats),
194
- LinearNoBias(token_z + dim_token_rel_pos_feats, token_z),
195
- )
196
-
197
- transitions = ModuleList([])
198
- for _ in range(num_transitions):
199
- transition = Transition(
200
- dim=token_z, hidden=transition_expansion_factor * token_z
201
- )
202
- transitions.append(transition)
203
-
204
- self.transitions = transitions
205
-
206
- def forward(
207
- self,
208
- z_trunk, # Float['b n n tz'],
209
- token_rel_pos_feats, # Float['b n n 3'],
210
- ): # -> Float['b n n tz']:
211
- z = torch.cat((z_trunk, token_rel_pos_feats), dim=-1)
212
- z = self.dim_pairwise_init_proj(z)
213
-
214
- for transition in self.transitions:
215
- z = transition(z) + z
216
-
217
- return z
218
-
219
-
220
- def get_indexing_matrix(K, W, H, device):
221
- assert W % 2 == 0
222
- assert H % (W // 2) == 0
223
-
224
- h = H // (W // 2)
225
- assert h % 2 == 0
226
-
227
- arange = torch.arange(2 * K, device=device)
228
- index = ((arange.unsqueeze(0) - arange.unsqueeze(1)) + h // 2).clamp(
229
- min=0, max=h + 1
230
- )
231
- index = index.view(K, 2, 2 * K)[:, 0, :]
232
- onehot = one_hot(index, num_classes=h + 2)[..., 1:-1].transpose(1, 0)
233
- return onehot.reshape(2 * K, h * K).float()
234
-
235
-
236
- def single_to_keys(single, indexing_matrix, W, H):
237
- B, N, D = single.shape
238
- K = N // W
239
- single = single.view(B, 2 * K, W // 2, D)
240
- return torch.einsum("b j i d, j k -> b k i d", single, indexing_matrix).reshape(
241
- B, K, H, D
242
- ) # j = 2K, i = W//2, k = h * K
243
-
244
-
245
- class AtomEncoder(Module):
246
- def __init__(
247
- self,
248
- atom_s,
249
- atom_z,
250
- token_s,
251
- token_z,
252
- atoms_per_window_queries,
253
- atoms_per_window_keys,
254
- atom_feature_dim,
255
- structure_prediction=True,
256
- use_no_atom_char=False,
257
- use_atom_backbone_feat=False,
258
- use_residue_feats_atoms=False,
259
- ):
260
- super().__init__()
261
-
262
- self.embed_atom_features = Linear(atom_feature_dim, atom_s)
263
- self.embed_atompair_ref_pos = LinearNoBias(3, atom_z)
264
- self.embed_atompair_ref_dist = LinearNoBias(1, atom_z)
265
- self.embed_atompair_mask = LinearNoBias(1, atom_z)
266
- self.atoms_per_window_queries = atoms_per_window_queries
267
- self.atoms_per_window_keys = atoms_per_window_keys
268
- self.use_no_atom_char = use_no_atom_char
269
- self.use_atom_backbone_feat = use_atom_backbone_feat
270
- self.use_residue_feats_atoms = use_residue_feats_atoms
271
-
272
- self.structure_prediction = structure_prediction
273
- if structure_prediction:
274
- self.s_to_c_trans = nn.Sequential(
275
- nn.LayerNorm(token_s), LinearNoBias(token_s, atom_s)
276
- )
277
- init.final_init_(self.s_to_c_trans[1].weight)
278
-
279
- self.z_to_p_trans = nn.Sequential(
280
- nn.LayerNorm(token_z), LinearNoBias(token_z, atom_z)
281
- )
282
- init.final_init_(self.z_to_p_trans[1].weight)
283
-
284
- self.c_to_p_trans_k = nn.Sequential(
285
- nn.ReLU(),
286
- LinearNoBias(atom_s, atom_z),
287
- )
288
- init.final_init_(self.c_to_p_trans_k[1].weight)
289
-
290
- self.c_to_p_trans_q = nn.Sequential(
291
- nn.ReLU(),
292
- LinearNoBias(atom_s, atom_z),
293
- )
294
- init.final_init_(self.c_to_p_trans_q[1].weight)
295
-
296
- self.p_mlp = nn.Sequential(
297
- nn.ReLU(),
298
- LinearNoBias(atom_z, atom_z),
299
- nn.ReLU(),
300
- LinearNoBias(atom_z, atom_z),
301
- nn.ReLU(),
302
- LinearNoBias(atom_z, atom_z),
303
- )
304
- init.final_init_(self.p_mlp[5].weight)
305
-
306
- def forward(
307
- self,
308
- feats,
309
- s_trunk=None, # Float['bm n ts'],
310
- z=None, # Float['bm n n tz'],
311
- ):
312
- with torch.autocast("cuda", enabled=False):
313
- B, N, _ = feats["ref_pos"].shape
314
- atom_mask = feats["atom_pad_mask"].bool() # Bool['b m'],
315
-
316
- atom_ref_pos = feats["ref_pos"] # Float['b m 3'],
317
- atom_uid = feats["ref_space_uid"] # Long['b m'],
318
-
319
- atom_feats = [
320
- atom_ref_pos,
321
- feats["ref_charge"].unsqueeze(-1),
322
- feats["ref_element"],
323
- ]
324
- if not self.use_no_atom_char:
325
- atom_feats.append(feats["ref_atom_name_chars"].reshape(B, N, 4 * 64))
326
- if self.use_atom_backbone_feat:
327
- atom_feats.append(feats["atom_backbone_feat"])
328
- if self.use_residue_feats_atoms:
329
- res_feats = torch.cat(
330
- [
331
- feats["res_type"],
332
- feats["modified"].unsqueeze(-1),
333
- one_hot(feats["mol_type"], num_classes=4).float(),
334
- ],
335
- dim=-1,
336
- )
337
- atom_to_token = feats["atom_to_token"].float()
338
- atom_res_feats = torch.bmm(atom_to_token, res_feats)
339
- atom_feats.append(atom_res_feats)
340
-
341
- atom_feats = torch.cat(atom_feats, dim=-1)
342
-
343
- c = self.embed_atom_features(atom_feats)
344
-
345
- # note we are already creating the windows to make it more efficient
346
- W, H = self.atoms_per_window_queries, self.atoms_per_window_keys
347
- B, N = c.shape[:2]
348
- K = N // W
349
- keys_indexing_matrix = get_indexing_matrix(K, W, H, c.device)
350
- to_keys = partial(
351
- single_to_keys, indexing_matrix=keys_indexing_matrix, W=W, H=H
352
- )
353
-
354
- atom_ref_pos_queries = atom_ref_pos.view(B, K, W, 1, 3)
355
- atom_ref_pos_keys = to_keys(atom_ref_pos).view(B, K, 1, H, 3)
356
-
357
- d = atom_ref_pos_keys - atom_ref_pos_queries # Float['b k w h 3']
358
- d_norm = torch.sum(d * d, dim=-1, keepdim=True) # Float['b k w h 1']
359
- d_norm = 1 / (
360
- 1 + d_norm
361
- ) # AF3 feeds in the reciprocal of the distance norm
362
-
363
- atom_mask_queries = atom_mask.view(B, K, W, 1)
364
- atom_mask_keys = (
365
- to_keys(atom_mask.unsqueeze(-1).float()).view(B, K, 1, H).bool()
366
- )
367
- atom_uid_queries = atom_uid.view(B, K, W, 1)
368
- atom_uid_keys = (
369
- to_keys(atom_uid.unsqueeze(-1).float()).view(B, K, 1, H).long()
370
- )
371
- v = (
372
- (
373
- atom_mask_queries
374
- & atom_mask_keys
375
- & (atom_uid_queries == atom_uid_keys)
376
- )
377
- .float()
378
- .unsqueeze(-1)
379
- ) # Bool['b k w h 1']
380
-
381
- p = self.embed_atompair_ref_pos(d) * v
382
- p = p + self.embed_atompair_ref_dist(d_norm) * v
383
- p = p + self.embed_atompair_mask(v) * v
384
-
385
- q = c
386
-
387
- if self.structure_prediction:
388
- # run only in structure model not in initial encoding
389
- atom_to_token = feats["atom_to_token"].float() # Long['b m n'],
390
-
391
- s_to_c = self.s_to_c_trans(s_trunk.float())
392
- s_to_c = torch.bmm(atom_to_token, s_to_c)
393
- c = c + s_to_c.to(c)
394
-
395
- atom_to_token_queries = atom_to_token.view(
396
- B, K, W, atom_to_token.shape[-1]
397
- )
398
- atom_to_token_keys = to_keys(atom_to_token)
399
- z_to_p = self.z_to_p_trans(z.float())
400
- z_to_p = torch.einsum(
401
- "bijd,bwki,bwlj->bwkld",
402
- z_to_p,
403
- atom_to_token_queries,
404
- atom_to_token_keys,
405
- )
406
- p = p + z_to_p.to(p)
407
-
408
- p = p + self.c_to_p_trans_q(c.view(B, K, W, 1, c.shape[-1]))
409
- p = p + self.c_to_p_trans_k(to_keys(c).view(B, K, 1, H, c.shape[-1]))
410
- p = p + self.p_mlp(p)
411
- return q, c, p, to_keys
412
-
413
-
414
- class AtomAttentionEncoder(Module):
415
- def __init__(
416
- self,
417
- atom_s,
418
- token_s,
419
- atoms_per_window_queries,
420
- atoms_per_window_keys,
421
- atom_encoder_depth=3,
422
- atom_encoder_heads=4,
423
- structure_prediction=True,
424
- activation_checkpointing=False,
425
- transformer_post_layer_norm=False,
426
- ):
427
- super().__init__()
428
-
429
- self.structure_prediction = structure_prediction
430
- if structure_prediction:
431
- self.r_to_q_trans = LinearNoBias(3, atom_s)
432
- init.final_init_(self.r_to_q_trans.weight)
433
-
434
- self.atom_encoder = AtomTransformer(
435
- dim=atom_s,
436
- dim_single_cond=atom_s,
437
- attn_window_queries=atoms_per_window_queries,
438
- attn_window_keys=atoms_per_window_keys,
439
- depth=atom_encoder_depth,
440
- heads=atom_encoder_heads,
441
- activation_checkpointing=activation_checkpointing,
442
- post_layer_norm=transformer_post_layer_norm,
443
- )
444
-
445
- self.atom_to_token_trans = nn.Sequential(
446
- LinearNoBias(atom_s, 2 * token_s if structure_prediction else token_s),
447
- nn.ReLU(),
448
- )
449
-
450
- def forward(
451
- self,
452
- feats,
453
- q,
454
- c,
455
- atom_enc_bias,
456
- to_keys,
457
- r=None, # Float['bm m 3'],
458
- multiplicity=1,
459
- ):
460
- B, N, _ = feats["ref_pos"].shape
461
- atom_mask = feats["atom_pad_mask"].bool() # Bool['b m'],
462
-
463
- if self.structure_prediction:
464
- # only here the multiplicity kicks in because we use the different positions r
465
- q = q.repeat_interleave(multiplicity, 0)
466
- r_to_q = self.r_to_q_trans(r)
467
- q = q + r_to_q
468
-
469
- c = c.repeat_interleave(multiplicity, 0)
470
- atom_mask = atom_mask.repeat_interleave(multiplicity, 0)
471
-
472
- q = self.atom_encoder(
473
- q=q,
474
- mask=atom_mask,
475
- c=c,
476
- bias=atom_enc_bias,
477
- multiplicity=multiplicity,
478
- to_keys=to_keys,
479
- )
480
-
481
- with torch.autocast("cuda", enabled=False):
482
- q_to_a = self.atom_to_token_trans(q).float()
483
- atom_to_token = feats["atom_to_token"].float()
484
- atom_to_token = atom_to_token.repeat_interleave(multiplicity, 0)
485
- atom_to_token_mean = atom_to_token / (
486
- atom_to_token.sum(dim=1, keepdim=True) + 1e-6
487
- )
488
- a = torch.bmm(atom_to_token_mean.transpose(1, 2), q_to_a)
489
-
490
- a = a.to(q)
491
-
492
- return a, q, c, to_keys
493
-
494
-
495
- class AtomAttentionDecoder(Module):
496
- """Algorithm 6."""
497
-
498
- def __init__(
499
- self,
500
- atom_s,
501
- token_s,
502
- attn_window_queries,
503
- attn_window_keys,
504
- atom_decoder_depth=3,
505
- atom_decoder_heads=4,
506
- activation_checkpointing=False,
507
- transformer_post_layer_norm=False,
508
- ):
509
- super().__init__()
510
-
511
- self.a_to_q_trans = LinearNoBias(2 * token_s, atom_s)
512
- init.final_init_(self.a_to_q_trans.weight)
513
-
514
- self.atom_decoder = AtomTransformer(
515
- dim=atom_s,
516
- dim_single_cond=atom_s,
517
- attn_window_queries=attn_window_queries,
518
- attn_window_keys=attn_window_keys,
519
- depth=atom_decoder_depth,
520
- heads=atom_decoder_heads,
521
- activation_checkpointing=activation_checkpointing,
522
- post_layer_norm=transformer_post_layer_norm,
523
- )
524
-
525
- if transformer_post_layer_norm:
526
- self.atom_feat_to_atom_pos_update = LinearNoBias(atom_s, 3)
527
- init.final_init_(self.atom_feat_to_atom_pos_update.weight)
528
- else:
529
- self.atom_feat_to_atom_pos_update = nn.Sequential(
530
- nn.LayerNorm(atom_s), LinearNoBias(atom_s, 3)
531
- )
532
- init.final_init_(self.atom_feat_to_atom_pos_update[1].weight)
533
-
534
- def forward(
535
- self,
536
- a, # Float['bm n 2ts'],
537
- q, # Float['bm m as'],
538
- c, # Float['bm m as'],
539
- atom_dec_bias, # Float['bm m m az'],
540
- feats,
541
- to_keys,
542
- multiplicity=1,
543
- ):
544
- with torch.autocast("cuda", enabled=False):
545
- atom_to_token = feats["atom_to_token"].float()
546
- atom_to_token = atom_to_token.repeat_interleave(multiplicity, 0)
547
-
548
- a_to_q = self.a_to_q_trans(a.float())
549
- a_to_q = torch.bmm(atom_to_token, a_to_q)
550
-
551
- q = q + a_to_q.to(q)
552
- atom_mask = feats["atom_pad_mask"] # Bool['b m'],
553
- atom_mask = atom_mask.repeat_interleave(multiplicity, 0)
554
-
555
- q = self.atom_decoder(
556
- q=q,
557
- mask=atom_mask,
558
- c=c,
559
- bias=atom_dec_bias,
560
- multiplicity=multiplicity,
561
- to_keys=to_keys,
562
- )
563
-
564
- r_update = self.atom_feat_to_atom_pos_update(q)
565
- return r_update
 
1
+ # started from code from https://github.com/lucidrains/alphafold3-pytorch, MIT License, Copyright (c) 2024 Phil Wang
2
+ from functools import partial
3
+ from math import pi
4
+
5
+ import torch
6
+ from einops import rearrange
7
+ from torch import nn
8
+ from torch.nn import Linear, Module, ModuleList
9
+ from torch.nn.functional import one_hot
10
+
11
+ from . import vb_layers_initialize as init
12
+ from .vb_layers_transition import Transition
13
+ from .vb_modules_transformersv2 import AtomTransformer
14
+ from .vb_modules_utils import LinearNoBias
15
+
16
+
17
+ class FourierEmbedding(Module):
18
+ """Algorithm 22."""
19
+
20
+ def __init__(self, dim):
21
+ super().__init__()
22
+ self.proj = nn.Linear(1, dim)
23
+ torch.nn.init.normal_(self.proj.weight, mean=0, std=1)
24
+ torch.nn.init.normal_(self.proj.bias, mean=0, std=1)
25
+ self.proj.requires_grad_(False)
26
+
27
+ def forward(
28
+ self,
29
+ times, # Float[' b'],
30
+ ): # -> Float['b d']:
31
+ times = rearrange(times, "b -> b 1")
32
+ rand_proj = self.proj(times)
33
+ return torch.cos(2 * pi * rand_proj)
34
+
35
+
36
+ class RelativePositionEncoder(Module):
37
+ """Algorithm 3."""
38
+
39
+ def __init__(
40
+ self, token_z, r_max=32, s_max=2, fix_sym_check=False, cyclic_pos_enc=False
41
+ ):
42
+ super().__init__()
43
+ self.r_max = r_max
44
+ self.s_max = s_max
45
+ self.linear_layer = LinearNoBias(4 * (r_max + 1) + 2 * (s_max + 1) + 1, token_z)
46
+ self.fix_sym_check = fix_sym_check
47
+ self.cyclic_pos_enc = cyclic_pos_enc
48
+
49
+ def forward(self, feats):
50
+ b_same_chain = torch.eq(
51
+ feats["asym_id"][:, :, None], feats["asym_id"][:, None, :]
52
+ )
53
+ b_same_residue = torch.eq(
54
+ feats["residue_index"][:, :, None], feats["residue_index"][:, None, :]
55
+ )
56
+ b_same_entity = torch.eq(
57
+ feats["entity_id"][:, :, None], feats["entity_id"][:, None, :]
58
+ )
59
+
60
+ d_residue = (
61
+ feats["residue_index"][:, :, None] - feats["residue_index"][:, None, :]
62
+ )
63
+
64
+ if self.cyclic_pos_enc and torch.any(feats["cyclic_period"] > 0):
65
+ period = torch.where(
66
+ feats["cyclic_period"] > 0,
67
+ feats["cyclic_period"],
68
+ torch.zeros_like(feats["cyclic_period"]) + 10000,
69
+ )
70
+ d_residue = (d_residue - period * torch.round(d_residue / period)).long()
71
+
72
+ d_residue = torch.clip(
73
+ d_residue + self.r_max,
74
+ 0,
75
+ 2 * self.r_max,
76
+ )
77
+ d_residue = torch.where(
78
+ b_same_chain, d_residue, torch.zeros_like(d_residue) + 2 * self.r_max + 1
79
+ )
80
+ a_rel_pos = one_hot(d_residue, 2 * self.r_max + 2)
81
+
82
+ d_token = torch.clip(
83
+ feats["token_index"][:, :, None]
84
+ - feats["token_index"][:, None, :]
85
+ + self.r_max,
86
+ 0,
87
+ 2 * self.r_max,
88
+ )
89
+ d_token = torch.where(
90
+ b_same_chain & b_same_residue,
91
+ d_token,
92
+ torch.zeros_like(d_token) + 2 * self.r_max + 1,
93
+ )
94
+ a_rel_token = one_hot(d_token, 2 * self.r_max + 2)
95
+
96
+ d_chain = torch.clip(
97
+ feats["sym_id"][:, :, None] - feats["sym_id"][:, None, :] + self.s_max,
98
+ 0,
99
+ 2 * self.s_max,
100
+ )
101
+ d_chain = torch.where(
102
+ (~b_same_entity) if self.fix_sym_check else b_same_chain,
103
+ torch.zeros_like(d_chain) + 2 * self.s_max + 1,
104
+ d_chain,
105
+ )
106
+ # Note: added | (~b_same_entity) based on observation of ProteinX manuscript
107
+ a_rel_chain = one_hot(d_chain, 2 * self.s_max + 2)
108
+
109
+ p = self.linear_layer(
110
+ torch.cat(
111
+ [
112
+ a_rel_pos.float(),
113
+ a_rel_token.float(),
114
+ b_same_entity.unsqueeze(-1).float(),
115
+ a_rel_chain.float(),
116
+ ],
117
+ dim=-1,
118
+ )
119
+ )
120
+ return p
121
+
122
+
123
+ class SingleConditioning(Module):
124
+ """Algorithm 21."""
125
+
126
+ def __init__(
127
+ self,
128
+ sigma_data: float,
129
+ token_s: int = 384,
130
+ dim_fourier: int = 256,
131
+ num_transitions: int = 2,
132
+ transition_expansion_factor: int = 2,
133
+ eps: float = 1e-20,
134
+ disable_times: bool = False,
135
+ ) -> None:
136
+ super().__init__()
137
+ self.eps = eps
138
+ self.sigma_data = sigma_data
139
+ self.disable_times = disable_times
140
+
141
+ self.norm_single = nn.LayerNorm(2 * token_s)
142
+ self.single_embed = nn.Linear(2 * token_s, 2 * token_s)
143
+ if not self.disable_times:
144
+ self.fourier_embed = FourierEmbedding(dim_fourier)
145
+ self.norm_fourier = nn.LayerNorm(dim_fourier)
146
+ self.fourier_to_single = LinearNoBias(dim_fourier, 2 * token_s)
147
+
148
+ transitions = ModuleList([])
149
+ for _ in range(num_transitions):
150
+ transition = Transition(
151
+ dim=2 * token_s, hidden=transition_expansion_factor * 2 * token_s
152
+ )
153
+ transitions.append(transition)
154
+
155
+ self.transitions = transitions
156
+
157
+ def forward(
158
+ self,
159
+ times, # Float[' b'],
160
+ s_trunk, # Float['b n ts'],
161
+ s_inputs, # Float['b n ts'],
162
+ ): # -> Float['b n 2ts']:
163
+ s = torch.cat((s_trunk, s_inputs), dim=-1)
164
+ s = self.single_embed(self.norm_single(s))
165
+ if not self.disable_times:
166
+ fourier_embed = self.fourier_embed(
167
+ times
168
+ ) # note: sigma rescaling done in diffusion module
169
+ normed_fourier = self.norm_fourier(fourier_embed)
170
+ fourier_to_single = self.fourier_to_single(normed_fourier)
171
+
172
+ s = rearrange(fourier_to_single, "b d -> b 1 d") + s
173
+
174
+ for transition in self.transitions:
175
+ s = transition(s) + s
176
+
177
+ return s, normed_fourier if not self.disable_times else None
178
+
179
+
180
+ class PairwiseConditioning(Module):
181
+ """Algorithm 21."""
182
+
183
+ def __init__(
184
+ self,
185
+ token_z,
186
+ dim_token_rel_pos_feats,
187
+ num_transitions=2,
188
+ transition_expansion_factor=2,
189
+ ):
190
+ super().__init__()
191
+
192
+ self.dim_pairwise_init_proj = nn.Sequential(
193
+ nn.LayerNorm(token_z + dim_token_rel_pos_feats),
194
+ LinearNoBias(token_z + dim_token_rel_pos_feats, token_z),
195
+ )
196
+
197
+ transitions = ModuleList([])
198
+ for _ in range(num_transitions):
199
+ transition = Transition(
200
+ dim=token_z, hidden=transition_expansion_factor * token_z
201
+ )
202
+ transitions.append(transition)
203
+
204
+ self.transitions = transitions
205
+
206
+ def forward(
207
+ self,
208
+ z_trunk, # Float['b n n tz'],
209
+ token_rel_pos_feats, # Float['b n n 3'],
210
+ ): # -> Float['b n n tz']:
211
+ z = torch.cat((z_trunk, token_rel_pos_feats), dim=-1)
212
+ z = self.dim_pairwise_init_proj(z)
213
+
214
+ for transition in self.transitions:
215
+ z = transition(z) + z
216
+
217
+ return z
218
+
219
+
220
+ def get_indexing_matrix(K, W, H, device):
221
+ assert W % 2 == 0
222
+ assert H % (W // 2) == 0
223
+
224
+ h = H // (W // 2)
225
+ assert h % 2 == 0
226
+
227
+ arange = torch.arange(2 * K, device=device)
228
+ index = ((arange.unsqueeze(0) - arange.unsqueeze(1)) + h // 2).clamp(
229
+ min=0, max=h + 1
230
+ )
231
+ index = index.view(K, 2, 2 * K)[:, 0, :]
232
+ onehot = one_hot(index, num_classes=h + 2)[..., 1:-1].transpose(1, 0)
233
+ return onehot.reshape(2 * K, h * K).float()
234
+
235
+
236
+ def single_to_keys(single, indexing_matrix, W, H):
237
+ B, N, D = single.shape
238
+ K = N // W
239
+ single = single.view(B, 2 * K, W // 2, D)
240
+ return torch.einsum("b j i d, j k -> b k i d", single, indexing_matrix).reshape(
241
+ B, K, H, D
242
+ ) # j = 2K, i = W//2, k = h * K
243
+
244
+
245
+ class AtomEncoder(Module):
246
+ def __init__(
247
+ self,
248
+ atom_s,
249
+ atom_z,
250
+ token_s,
251
+ token_z,
252
+ atoms_per_window_queries,
253
+ atoms_per_window_keys,
254
+ atom_feature_dim,
255
+ structure_prediction=True,
256
+ use_no_atom_char=False,
257
+ use_atom_backbone_feat=False,
258
+ use_residue_feats_atoms=False,
259
+ ):
260
+ super().__init__()
261
+
262
+ self.embed_atom_features = Linear(atom_feature_dim, atom_s)
263
+ self.embed_atompair_ref_pos = LinearNoBias(3, atom_z)
264
+ self.embed_atompair_ref_dist = LinearNoBias(1, atom_z)
265
+ self.embed_atompair_mask = LinearNoBias(1, atom_z)
266
+ self.atoms_per_window_queries = atoms_per_window_queries
267
+ self.atoms_per_window_keys = atoms_per_window_keys
268
+ self.use_no_atom_char = use_no_atom_char
269
+ self.use_atom_backbone_feat = use_atom_backbone_feat
270
+ self.use_residue_feats_atoms = use_residue_feats_atoms
271
+
272
+ self.structure_prediction = structure_prediction
273
+ if structure_prediction:
274
+ self.s_to_c_trans = nn.Sequential(
275
+ nn.LayerNorm(token_s), LinearNoBias(token_s, atom_s)
276
+ )
277
+ init.final_init_(self.s_to_c_trans[1].weight)
278
+
279
+ self.z_to_p_trans = nn.Sequential(
280
+ nn.LayerNorm(token_z), LinearNoBias(token_z, atom_z)
281
+ )
282
+ init.final_init_(self.z_to_p_trans[1].weight)
283
+
284
+ self.c_to_p_trans_k = nn.Sequential(
285
+ nn.ReLU(),
286
+ LinearNoBias(atom_s, atom_z),
287
+ )
288
+ init.final_init_(self.c_to_p_trans_k[1].weight)
289
+
290
+ self.c_to_p_trans_q = nn.Sequential(
291
+ nn.ReLU(),
292
+ LinearNoBias(atom_s, atom_z),
293
+ )
294
+ init.final_init_(self.c_to_p_trans_q[1].weight)
295
+
296
+ self.p_mlp = nn.Sequential(
297
+ nn.ReLU(),
298
+ LinearNoBias(atom_z, atom_z),
299
+ nn.ReLU(),
300
+ LinearNoBias(atom_z, atom_z),
301
+ nn.ReLU(),
302
+ LinearNoBias(atom_z, atom_z),
303
+ )
304
+ init.final_init_(self.p_mlp[5].weight)
305
+
306
+ def forward(
307
+ self,
308
+ feats,
309
+ s_trunk=None, # Float['bm n ts'],
310
+ z=None, # Float['bm n n tz'],
311
+ ):
312
+ with torch.autocast("cuda", enabled=False):
313
+ B, N, _ = feats["ref_pos"].shape
314
+ atom_mask = feats["atom_pad_mask"].bool() # Bool['b m'],
315
+
316
+ atom_ref_pos = feats["ref_pos"] # Float['b m 3'],
317
+ atom_uid = feats["ref_space_uid"] # Long['b m'],
318
+
319
+ atom_feats = [
320
+ atom_ref_pos,
321
+ feats["ref_charge"].unsqueeze(-1),
322
+ feats["ref_element"],
323
+ ]
324
+ if not self.use_no_atom_char:
325
+ atom_feats.append(feats["ref_atom_name_chars"].reshape(B, N, 4 * 64))
326
+ if self.use_atom_backbone_feat:
327
+ atom_feats.append(feats["atom_backbone_feat"])
328
+ if self.use_residue_feats_atoms:
329
+ res_feats = torch.cat(
330
+ [
331
+ feats["res_type"],
332
+ feats["modified"].unsqueeze(-1),
333
+ one_hot(feats["mol_type"], num_classes=4).float(),
334
+ ],
335
+ dim=-1,
336
+ )
337
+ atom_to_token = feats["atom_to_token"].float()
338
+ atom_res_feats = torch.bmm(atom_to_token, res_feats)
339
+ atom_feats.append(atom_res_feats)
340
+
341
+ atom_feats = torch.cat(atom_feats, dim=-1)
342
+
343
+ c = self.embed_atom_features(atom_feats)
344
+
345
+ # note we are already creating the windows to make it more efficient
346
+ W, H = self.atoms_per_window_queries, self.atoms_per_window_keys
347
+ B, N = c.shape[:2]
348
+ K = N // W
349
+ keys_indexing_matrix = get_indexing_matrix(K, W, H, c.device)
350
+ to_keys = partial(
351
+ single_to_keys, indexing_matrix=keys_indexing_matrix, W=W, H=H
352
+ )
353
+
354
+ atom_ref_pos_queries = atom_ref_pos.view(B, K, W, 1, 3)
355
+ atom_ref_pos_keys = to_keys(atom_ref_pos).view(B, K, 1, H, 3)
356
+
357
+ d = atom_ref_pos_keys - atom_ref_pos_queries # Float['b k w h 3']
358
+ d_norm = torch.sum(d * d, dim=-1, keepdim=True) # Float['b k w h 1']
359
+ d_norm = 1 / (
360
+ 1 + d_norm
361
+ ) # AF3 feeds in the reciprocal of the distance norm
362
+
363
+ atom_mask_queries = atom_mask.view(B, K, W, 1)
364
+ atom_mask_keys = (
365
+ to_keys(atom_mask.unsqueeze(-1).float()).view(B, K, 1, H).bool()
366
+ )
367
+ atom_uid_queries = atom_uid.view(B, K, W, 1)
368
+ atom_uid_keys = (
369
+ to_keys(atom_uid.unsqueeze(-1).float()).view(B, K, 1, H).long()
370
+ )
371
+ v = (
372
+ (
373
+ atom_mask_queries
374
+ & atom_mask_keys
375
+ & (atom_uid_queries == atom_uid_keys)
376
+ )
377
+ .float()
378
+ .unsqueeze(-1)
379
+ ) # Bool['b k w h 1']
380
+
381
+ p = self.embed_atompair_ref_pos(d) * v
382
+ p = p + self.embed_atompair_ref_dist(d_norm) * v
383
+ p = p + self.embed_atompair_mask(v) * v
384
+
385
+ q = c
386
+
387
+ if self.structure_prediction:
388
+ # run only in structure model not in initial encoding
389
+ atom_to_token = feats["atom_to_token"].float() # Long['b m n'],
390
+
391
+ s_to_c = self.s_to_c_trans(s_trunk.float())
392
+ s_to_c = torch.bmm(atom_to_token, s_to_c)
393
+ c = c + s_to_c.to(c)
394
+
395
+ atom_to_token_queries = atom_to_token.view(
396
+ B, K, W, atom_to_token.shape[-1]
397
+ )
398
+ atom_to_token_keys = to_keys(atom_to_token)
399
+ z_to_p = self.z_to_p_trans(z.float())
400
+ z_to_p = torch.einsum(
401
+ "bijd,bwki,bwlj->bwkld",
402
+ z_to_p,
403
+ atom_to_token_queries,
404
+ atom_to_token_keys,
405
+ )
406
+ p = p + z_to_p.to(p)
407
+
408
+ p = p + self.c_to_p_trans_q(c.view(B, K, W, 1, c.shape[-1]))
409
+ p = p + self.c_to_p_trans_k(to_keys(c).view(B, K, 1, H, c.shape[-1]))
410
+ p = p + self.p_mlp(p)
411
+ return q, c, p, to_keys
412
+
413
+
414
+ class AtomAttentionEncoder(Module):
415
+ def __init__(
416
+ self,
417
+ atom_s,
418
+ token_s,
419
+ atoms_per_window_queries,
420
+ atoms_per_window_keys,
421
+ atom_encoder_depth=3,
422
+ atom_encoder_heads=4,
423
+ structure_prediction=True,
424
+ activation_checkpointing=False,
425
+ transformer_post_layer_norm=False,
426
+ ):
427
+ super().__init__()
428
+
429
+ self.structure_prediction = structure_prediction
430
+ if structure_prediction:
431
+ self.r_to_q_trans = LinearNoBias(3, atom_s)
432
+ init.final_init_(self.r_to_q_trans.weight)
433
+
434
+ self.atom_encoder = AtomTransformer(
435
+ dim=atom_s,
436
+ dim_single_cond=atom_s,
437
+ attn_window_queries=atoms_per_window_queries,
438
+ attn_window_keys=atoms_per_window_keys,
439
+ depth=atom_encoder_depth,
440
+ heads=atom_encoder_heads,
441
+ activation_checkpointing=activation_checkpointing,
442
+ post_layer_norm=transformer_post_layer_norm,
443
+ )
444
+
445
+ self.atom_to_token_trans = nn.Sequential(
446
+ LinearNoBias(atom_s, 2 * token_s if structure_prediction else token_s),
447
+ nn.ReLU(),
448
+ )
449
+
450
+ def forward(
451
+ self,
452
+ feats,
453
+ q,
454
+ c,
455
+ atom_enc_bias,
456
+ to_keys,
457
+ r=None, # Float['bm m 3'],
458
+ multiplicity=1,
459
+ ):
460
+ B, N, _ = feats["ref_pos"].shape
461
+ atom_mask = feats["atom_pad_mask"].bool() # Bool['b m'],
462
+
463
+ if self.structure_prediction:
464
+ # only here the multiplicity kicks in because we use the different positions r
465
+ q = q.repeat_interleave(multiplicity, 0)
466
+ r_to_q = self.r_to_q_trans(r)
467
+ q = q + r_to_q
468
+
469
+ c = c.repeat_interleave(multiplicity, 0)
470
+ atom_mask = atom_mask.repeat_interleave(multiplicity, 0)
471
+
472
+ q = self.atom_encoder(
473
+ q=q,
474
+ mask=atom_mask,
475
+ c=c,
476
+ bias=atom_enc_bias,
477
+ multiplicity=multiplicity,
478
+ to_keys=to_keys,
479
+ )
480
+
481
+ with torch.autocast("cuda", enabled=False):
482
+ q_to_a = self.atom_to_token_trans(q).float()
483
+ atom_to_token = feats["atom_to_token"].float()
484
+ atom_to_token = atom_to_token.repeat_interleave(multiplicity, 0)
485
+ atom_to_token_mean = atom_to_token / (
486
+ atom_to_token.sum(dim=1, keepdim=True) + 1e-6
487
+ )
488
+ a = torch.bmm(atom_to_token_mean.transpose(1, 2), q_to_a)
489
+
490
+ a = a.to(q)
491
+
492
+ return a, q, c, to_keys
493
+
494
+
495
+ class AtomAttentionDecoder(Module):
496
+ """Algorithm 6."""
497
+
498
+ def __init__(
499
+ self,
500
+ atom_s,
501
+ token_s,
502
+ attn_window_queries,
503
+ attn_window_keys,
504
+ atom_decoder_depth=3,
505
+ atom_decoder_heads=4,
506
+ activation_checkpointing=False,
507
+ transformer_post_layer_norm=False,
508
+ ):
509
+ super().__init__()
510
+
511
+ self.a_to_q_trans = LinearNoBias(2 * token_s, atom_s)
512
+ init.final_init_(self.a_to_q_trans.weight)
513
+
514
+ self.atom_decoder = AtomTransformer(
515
+ dim=atom_s,
516
+ dim_single_cond=atom_s,
517
+ attn_window_queries=attn_window_queries,
518
+ attn_window_keys=attn_window_keys,
519
+ depth=atom_decoder_depth,
520
+ heads=atom_decoder_heads,
521
+ activation_checkpointing=activation_checkpointing,
522
+ post_layer_norm=transformer_post_layer_norm,
523
+ )
524
+
525
+ if transformer_post_layer_norm:
526
+ self.atom_feat_to_atom_pos_update = LinearNoBias(atom_s, 3)
527
+ init.final_init_(self.atom_feat_to_atom_pos_update.weight)
528
+ else:
529
+ self.atom_feat_to_atom_pos_update = nn.Sequential(
530
+ nn.LayerNorm(atom_s), LinearNoBias(atom_s, 3)
531
+ )
532
+ init.final_init_(self.atom_feat_to_atom_pos_update[1].weight)
533
+
534
+ def forward(
535
+ self,
536
+ a, # Float['bm n 2ts'],
537
+ q, # Float['bm m as'],
538
+ c, # Float['bm m as'],
539
+ atom_dec_bias, # Float['bm m m az'],
540
+ feats,
541
+ to_keys,
542
+ multiplicity=1,
543
+ ):
544
+ with torch.autocast("cuda", enabled=False):
545
+ atom_to_token = feats["atom_to_token"].float()
546
+ atom_to_token = atom_to_token.repeat_interleave(multiplicity, 0)
547
+
548
+ a_to_q = self.a_to_q_trans(a.float())
549
+ a_to_q = torch.bmm(atom_to_token, a_to_q)
550
+
551
+ q = q + a_to_q.to(q)
552
+ atom_mask = feats["atom_pad_mask"] # Bool['b m'],
553
+ atom_mask = atom_mask.repeat_interleave(multiplicity, 0)
554
+
555
+ q = self.atom_decoder(
556
+ q=q,
557
+ mask=atom_mask,
558
+ c=c,
559
+ bias=atom_dec_bias,
560
+ multiplicity=multiplicity,
561
+ to_keys=to_keys,
562
+ )
563
+
564
+ r_update = self.atom_feat_to_atom_pos_update(q)
565
+ return r_update
vb_modules_transformersv2.py CHANGED
@@ -1,262 +1,262 @@
1
- # started from code from https://github.com/lucidrains/alphafold3-pytorch, MIT License, Copyright (c) 2024 Phil Wang
2
-
3
- import torch
4
- from torch import nn, sigmoid
5
- from torch.nn import (
6
- LayerNorm,
7
- Linear,
8
- Module,
9
- ModuleList,
10
- Sequential,
11
- )
12
-
13
- from .vb_layers_attentionv2 import AttentionPairBias
14
- from .vb_modules_utils import LinearNoBias, SwiGLU, default
15
-
16
-
17
- class AdaLN(Module):
18
- """Algorithm 26"""
19
-
20
- def __init__(self, dim, dim_single_cond):
21
- super().__init__()
22
- self.a_norm = LayerNorm(dim, elementwise_affine=False, bias=False)
23
- self.s_norm = LayerNorm(dim_single_cond, bias=False)
24
- self.s_scale = Linear(dim_single_cond, dim)
25
- self.s_bias = LinearNoBias(dim_single_cond, dim)
26
-
27
- def forward(self, a, s):
28
- a = self.a_norm(a)
29
- s = self.s_norm(s)
30
- a = sigmoid(self.s_scale(s)) * a + self.s_bias(s)
31
- return a
32
-
33
-
34
- class ConditionedTransitionBlock(Module):
35
- """Algorithm 25"""
36
-
37
- def __init__(self, dim_single, dim_single_cond, expansion_factor=2):
38
- super().__init__()
39
-
40
- self.adaln = AdaLN(dim_single, dim_single_cond)
41
-
42
- dim_inner = int(dim_single * expansion_factor)
43
- self.swish_gate = Sequential(
44
- LinearNoBias(dim_single, dim_inner * 2),
45
- SwiGLU(),
46
- )
47
- self.a_to_b = LinearNoBias(dim_single, dim_inner)
48
- self.b_to_a = LinearNoBias(dim_inner, dim_single)
49
-
50
- output_projection_linear = Linear(dim_single_cond, dim_single)
51
- nn.init.zeros_(output_projection_linear.weight)
52
- nn.init.constant_(output_projection_linear.bias, -2.0)
53
-
54
- self.output_projection = nn.Sequential(output_projection_linear, nn.Sigmoid())
55
-
56
- def forward(
57
- self,
58
- a, # Float['... d']
59
- s,
60
- ): # -> Float['... d']:
61
- a = self.adaln(a, s)
62
- b = self.swish_gate(a) * self.a_to_b(a)
63
- a = self.output_projection(s) * self.b_to_a(b)
64
-
65
- return a
66
-
67
-
68
- class DiffusionTransformer(Module):
69
- """Algorithm 23"""
70
-
71
- def __init__(
72
- self,
73
- depth,
74
- heads,
75
- dim=384,
76
- dim_single_cond=None,
77
- pair_bias_attn=True,
78
- activation_checkpointing=False,
79
- post_layer_norm=False,
80
- ):
81
- super().__init__()
82
- self.activation_checkpointing = activation_checkpointing
83
- dim_single_cond = default(dim_single_cond, dim)
84
- self.pair_bias_attn = pair_bias_attn
85
-
86
- self.layers = ModuleList()
87
- for _ in range(depth):
88
- self.layers.append(
89
- DiffusionTransformerLayer(
90
- heads,
91
- dim,
92
- dim_single_cond,
93
- post_layer_norm,
94
- )
95
- )
96
-
97
- def forward(
98
- self,
99
- a, # Float['bm n d'],
100
- s, # Float['bm n ds'],
101
- bias=None, # Float['b n n dp']
102
- mask=None, # Bool['b n'] | None = None
103
- to_keys=None,
104
- multiplicity=1,
105
- ):
106
- if self.pair_bias_attn:
107
- B, N, M, D = bias.shape
108
- L = len(self.layers)
109
- bias = bias.view(B, N, M, L, D // L)
110
-
111
- for i, layer in enumerate(self.layers):
112
- if self.pair_bias_attn:
113
- bias_l = bias[:, :, :, i]
114
- else:
115
- bias_l = None
116
-
117
- if self.activation_checkpointing and self.training:
118
- a = torch.utils.checkpoint.checkpoint(
119
- layer,
120
- a,
121
- s,
122
- bias_l,
123
- mask,
124
- to_keys,
125
- multiplicity,
126
- )
127
-
128
- else:
129
- a = layer(
130
- a, # Float['bm n d'],
131
- s, # Float['bm n ds'],
132
- bias_l, # Float['b n n dp']
133
- mask, # Bool['b n'] | None = None
134
- to_keys,
135
- multiplicity,
136
- )
137
- return a
138
-
139
-
140
- class DiffusionTransformerLayer(Module):
141
- """Algorithm 23"""
142
-
143
- def __init__(
144
- self,
145
- heads,
146
- dim=384,
147
- dim_single_cond=None,
148
- post_layer_norm=False,
149
- ):
150
- super().__init__()
151
-
152
- dim_single_cond = default(dim_single_cond, dim)
153
-
154
- self.adaln = AdaLN(dim, dim_single_cond)
155
- self.pair_bias_attn = AttentionPairBias(
156
- c_s=dim, num_heads=heads, compute_pair_bias=False
157
- )
158
-
159
- self.output_projection_linear = Linear(dim_single_cond, dim)
160
- nn.init.zeros_(self.output_projection_linear.weight)
161
- nn.init.constant_(self.output_projection_linear.bias, -2.0)
162
-
163
- self.output_projection = nn.Sequential(
164
- self.output_projection_linear, nn.Sigmoid()
165
- )
166
- self.transition = ConditionedTransitionBlock(
167
- dim_single=dim, dim_single_cond=dim_single_cond
168
- )
169
-
170
- if post_layer_norm:
171
- self.post_lnorm = nn.LayerNorm(dim)
172
- else:
173
- self.post_lnorm = nn.Identity()
174
-
175
- def forward(
176
- self,
177
- a, # Float['bm n d'],
178
- s, # Float['bm n ds'],
179
- bias=None, # Float['b n n dp']
180
- mask=None, # Bool['b n'] | None = None
181
- to_keys=None,
182
- multiplicity=1,
183
- ):
184
- b = self.adaln(a, s)
185
-
186
- k_in = b
187
- if to_keys is not None:
188
- k_in = to_keys(b)
189
- mask = to_keys(mask.unsqueeze(-1)).squeeze(-1)
190
-
191
- if self.pair_bias_attn:
192
- b = self.pair_bias_attn(
193
- s=b,
194
- z=bias,
195
- mask=mask,
196
- multiplicity=multiplicity,
197
- k_in=k_in,
198
- )
199
- else:
200
- b = self.no_pair_bias_attn(s=b, mask=mask, k_in=k_in)
201
-
202
- b = self.output_projection(s) * b
203
-
204
- a = a + b
205
- a = a + self.transition(a, s)
206
-
207
- a = self.post_lnorm(a)
208
- return a
209
-
210
-
211
- class AtomTransformer(Module):
212
- """Algorithm 7"""
213
-
214
- def __init__(
215
- self,
216
- attn_window_queries,
217
- attn_window_keys,
218
- **diffusion_transformer_kwargs,
219
- ):
220
- super().__init__()
221
- self.attn_window_queries = attn_window_queries
222
- self.attn_window_keys = attn_window_keys
223
- self.diffusion_transformer = DiffusionTransformer(
224
- **diffusion_transformer_kwargs
225
- )
226
-
227
- def forward(
228
- self,
229
- q, # Float['b m d'],
230
- c, # Float['b m ds'],
231
- bias, # Float['b m m dp']
232
- to_keys,
233
- mask, # Bool['b m'] | None = None
234
- multiplicity=1,
235
- ):
236
- W = self.attn_window_queries
237
- H = self.attn_window_keys
238
-
239
- B, N, D = q.shape
240
- NW = N // W
241
-
242
- # reshape tokens
243
- q = q.view((B * NW, W, -1))
244
- c = c.view((B * NW, W, -1))
245
- mask = mask.view(B * NW, W)
246
- bias = bias.repeat_interleave(multiplicity, 0)
247
- bias = bias.view((bias.shape[0] * NW, W, H, -1))
248
-
249
- to_keys_new = lambda x: to_keys(x.view(B, NW * W, -1)).view(B * NW, H, -1)
250
-
251
- # main transformer
252
- q = self.diffusion_transformer(
253
- a=q,
254
- s=c,
255
- bias=bias,
256
- mask=mask.float(),
257
- multiplicity=1, # bias term already expanded with multiplicity
258
- to_keys=to_keys_new,
259
- )
260
-
261
- q = q.view((B, NW * W, D))
262
- return q
 
1
+ # started from code from https://github.com/lucidrains/alphafold3-pytorch, MIT License, Copyright (c) 2024 Phil Wang
2
+
3
+ import torch
4
+ from torch import nn, sigmoid
5
+ from torch.nn import (
6
+ LayerNorm,
7
+ Linear,
8
+ Module,
9
+ ModuleList,
10
+ Sequential,
11
+ )
12
+
13
+ from .vb_layers_attentionv2 import AttentionPairBias
14
+ from .vb_modules_utils import LinearNoBias, SwiGLU, default
15
+
16
+
17
+ class AdaLN(Module):
18
+ """Algorithm 26"""
19
+
20
+ def __init__(self, dim, dim_single_cond):
21
+ super().__init__()
22
+ self.a_norm = LayerNorm(dim, elementwise_affine=False, bias=False)
23
+ self.s_norm = LayerNorm(dim_single_cond, bias=False)
24
+ self.s_scale = Linear(dim_single_cond, dim)
25
+ self.s_bias = LinearNoBias(dim_single_cond, dim)
26
+
27
+ def forward(self, a, s):
28
+ a = self.a_norm(a)
29
+ s = self.s_norm(s)
30
+ a = sigmoid(self.s_scale(s)) * a + self.s_bias(s)
31
+ return a
32
+
33
+
34
+ class ConditionedTransitionBlock(Module):
35
+ """Algorithm 25"""
36
+
37
+ def __init__(self, dim_single, dim_single_cond, expansion_factor=2):
38
+ super().__init__()
39
+
40
+ self.adaln = AdaLN(dim_single, dim_single_cond)
41
+
42
+ dim_inner = int(dim_single * expansion_factor)
43
+ self.swish_gate = Sequential(
44
+ LinearNoBias(dim_single, dim_inner * 2),
45
+ SwiGLU(),
46
+ )
47
+ self.a_to_b = LinearNoBias(dim_single, dim_inner)
48
+ self.b_to_a = LinearNoBias(dim_inner, dim_single)
49
+
50
+ output_projection_linear = Linear(dim_single_cond, dim_single)
51
+ nn.init.zeros_(output_projection_linear.weight)
52
+ nn.init.constant_(output_projection_linear.bias, -2.0)
53
+
54
+ self.output_projection = nn.Sequential(output_projection_linear, nn.Sigmoid())
55
+
56
+ def forward(
57
+ self,
58
+ a, # Float['... d']
59
+ s,
60
+ ): # -> Float['... d']:
61
+ a = self.adaln(a, s)
62
+ b = self.swish_gate(a) * self.a_to_b(a)
63
+ a = self.output_projection(s) * self.b_to_a(b)
64
+
65
+ return a
66
+
67
+
68
+ class DiffusionTransformer(Module):
69
+ """Algorithm 23"""
70
+
71
+ def __init__(
72
+ self,
73
+ depth,
74
+ heads,
75
+ dim=384,
76
+ dim_single_cond=None,
77
+ pair_bias_attn=True,
78
+ activation_checkpointing=False,
79
+ post_layer_norm=False,
80
+ ):
81
+ super().__init__()
82
+ self.activation_checkpointing = activation_checkpointing
83
+ dim_single_cond = default(dim_single_cond, dim)
84
+ self.pair_bias_attn = pair_bias_attn
85
+
86
+ self.layers = ModuleList()
87
+ for _ in range(depth):
88
+ self.layers.append(
89
+ DiffusionTransformerLayer(
90
+ heads,
91
+ dim,
92
+ dim_single_cond,
93
+ post_layer_norm,
94
+ )
95
+ )
96
+
97
+ def forward(
98
+ self,
99
+ a, # Float['bm n d'],
100
+ s, # Float['bm n ds'],
101
+ bias=None, # Float['b n n dp']
102
+ mask=None, # Bool['b n'] | None = None
103
+ to_keys=None,
104
+ multiplicity=1,
105
+ ):
106
+ if self.pair_bias_attn:
107
+ B, N, M, D = bias.shape
108
+ L = len(self.layers)
109
+ bias = bias.view(B, N, M, L, D // L)
110
+
111
+ for i, layer in enumerate(self.layers):
112
+ if self.pair_bias_attn:
113
+ bias_l = bias[:, :, :, i]
114
+ else:
115
+ bias_l = None
116
+
117
+ if self.activation_checkpointing and self.training:
118
+ a = torch.utils.checkpoint.checkpoint(
119
+ layer,
120
+ a,
121
+ s,
122
+ bias_l,
123
+ mask,
124
+ to_keys,
125
+ multiplicity,
126
+ )
127
+
128
+ else:
129
+ a = layer(
130
+ a, # Float['bm n d'],
131
+ s, # Float['bm n ds'],
132
+ bias_l, # Float['b n n dp']
133
+ mask, # Bool['b n'] | None = None
134
+ to_keys,
135
+ multiplicity,
136
+ )
137
+ return a
138
+
139
+
140
+ class DiffusionTransformerLayer(Module):
141
+ """Algorithm 23"""
142
+
143
+ def __init__(
144
+ self,
145
+ heads,
146
+ dim=384,
147
+ dim_single_cond=None,
148
+ post_layer_norm=False,
149
+ ):
150
+ super().__init__()
151
+
152
+ dim_single_cond = default(dim_single_cond, dim)
153
+
154
+ self.adaln = AdaLN(dim, dim_single_cond)
155
+ self.pair_bias_attn = AttentionPairBias(
156
+ c_s=dim, num_heads=heads, compute_pair_bias=False
157
+ )
158
+
159
+ self.output_projection_linear = Linear(dim_single_cond, dim)
160
+ nn.init.zeros_(self.output_projection_linear.weight)
161
+ nn.init.constant_(self.output_projection_linear.bias, -2.0)
162
+
163
+ self.output_projection = nn.Sequential(
164
+ self.output_projection_linear, nn.Sigmoid()
165
+ )
166
+ self.transition = ConditionedTransitionBlock(
167
+ dim_single=dim, dim_single_cond=dim_single_cond
168
+ )
169
+
170
+ if post_layer_norm:
171
+ self.post_lnorm = nn.LayerNorm(dim)
172
+ else:
173
+ self.post_lnorm = nn.Identity()
174
+
175
+ def forward(
176
+ self,
177
+ a, # Float['bm n d'],
178
+ s, # Float['bm n ds'],
179
+ bias=None, # Float['b n n dp']
180
+ mask=None, # Bool['b n'] | None = None
181
+ to_keys=None,
182
+ multiplicity=1,
183
+ ):
184
+ b = self.adaln(a, s)
185
+
186
+ k_in = b
187
+ if to_keys is not None:
188
+ k_in = to_keys(b)
189
+ mask = to_keys(mask.unsqueeze(-1)).squeeze(-1)
190
+
191
+ if self.pair_bias_attn:
192
+ b = self.pair_bias_attn(
193
+ s=b,
194
+ z=bias,
195
+ mask=mask,
196
+ multiplicity=multiplicity,
197
+ k_in=k_in,
198
+ )
199
+ else:
200
+ b = self.no_pair_bias_attn(s=b, mask=mask, k_in=k_in)
201
+
202
+ b = self.output_projection(s) * b
203
+
204
+ a = a + b
205
+ a = a + self.transition(a, s)
206
+
207
+ a = self.post_lnorm(a)
208
+ return a
209
+
210
+
211
+ class AtomTransformer(Module):
212
+ """Algorithm 7"""
213
+
214
+ def __init__(
215
+ self,
216
+ attn_window_queries,
217
+ attn_window_keys,
218
+ **diffusion_transformer_kwargs,
219
+ ):
220
+ super().__init__()
221
+ self.attn_window_queries = attn_window_queries
222
+ self.attn_window_keys = attn_window_keys
223
+ self.diffusion_transformer = DiffusionTransformer(
224
+ **diffusion_transformer_kwargs
225
+ )
226
+
227
+ def forward(
228
+ self,
229
+ q, # Float['b m d'],
230
+ c, # Float['b m ds'],
231
+ bias, # Float['b m m dp']
232
+ to_keys,
233
+ mask, # Bool['b m'] | None = None
234
+ multiplicity=1,
235
+ ):
236
+ W = self.attn_window_queries
237
+ H = self.attn_window_keys
238
+
239
+ B, N, D = q.shape
240
+ NW = N // W
241
+
242
+ # reshape tokens
243
+ q = q.view((B * NW, W, -1))
244
+ c = c.view((B * NW, W, -1))
245
+ mask = mask.view(B * NW, W)
246
+ bias = bias.repeat_interleave(multiplicity, 0)
247
+ bias = bias.view((bias.shape[0] * NW, W, H, -1))
248
+
249
+ to_keys_new = lambda x: to_keys(x.view(B, NW * W, -1)).view(B * NW, H, -1)
250
+
251
+ # main transformer
252
+ q = self.diffusion_transformer(
253
+ a=q,
254
+ s=c,
255
+ bias=bias,
256
+ mask=mask.float(),
257
+ multiplicity=1, # bias term already expanded with multiplicity
258
+ to_keys=to_keys_new,
259
+ )
260
+
261
+ q = q.view((B, NW * W, D))
262
+ return q
vb_modules_trunkv2.py CHANGED
@@ -1,828 +1,828 @@
1
- import torch
2
- from torch import Tensor, nn
3
- from torch.nn.functional import one_hot
4
-
5
- from . import vb_const as const
6
- from .vb_layers_outer_product_mean import OuterProductMean
7
- from .vb_layers_pair_averaging import PairWeightedAveraging
8
- from .vb_layers_pairformer import (
9
- PairformerNoSeqLayer,
10
- PairformerNoSeqModule,
11
- get_dropout_mask,
12
- )
13
- from .vb_layers_transition import Transition
14
- from .vb_modules_encodersv2 import (
15
- AtomAttentionEncoder,
16
- AtomEncoder,
17
- FourierEmbedding,
18
- )
19
-
20
-
21
- class ContactConditioning(nn.Module):
22
- def __init__(self, token_z: int, cutoff_min: float, cutoff_max: float):
23
- super().__init__()
24
-
25
- self.fourier_embedding = FourierEmbedding(token_z)
26
- self.encoder = nn.Linear(
27
- token_z + len(const.contact_conditioning_info) - 1, token_z
28
- )
29
- self.encoding_unspecified = nn.Parameter(torch.zeros(token_z))
30
- self.encoding_unselected = nn.Parameter(torch.zeros(token_z))
31
- self.cutoff_min = cutoff_min
32
- self.cutoff_max = cutoff_max
33
-
34
- def forward(self, feats):
35
- assert const.contact_conditioning_info["UNSPECIFIED"] == 0
36
- assert const.contact_conditioning_info["UNSELECTED"] == 1
37
- contact_conditioning = feats["contact_conditioning"][:, :, :, 2:]
38
- contact_threshold = feats["contact_threshold"]
39
- contact_threshold_normalized = (contact_threshold - self.cutoff_min) / (
40
- self.cutoff_max - self.cutoff_min
41
- )
42
- contact_threshold_fourier = self.fourier_embedding(
43
- contact_threshold_normalized.flatten()
44
- ).reshape(contact_threshold_normalized.shape + (-1,))
45
-
46
- contact_conditioning = torch.cat(
47
- [
48
- contact_conditioning,
49
- contact_threshold_normalized.unsqueeze(-1),
50
- contact_threshold_fourier,
51
- ],
52
- dim=-1,
53
- )
54
- contact_conditioning = self.encoder(contact_conditioning)
55
-
56
- contact_conditioning = (
57
- contact_conditioning
58
- * (
59
- 1
60
- - feats["contact_conditioning"][:, :, :, 0:2].sum(dim=-1, keepdim=True)
61
- )
62
- + self.encoding_unspecified * feats["contact_conditioning"][:, :, :, 0:1]
63
- + self.encoding_unselected * feats["contact_conditioning"][:, :, :, 1:2]
64
- )
65
- return contact_conditioning
66
-
67
-
68
- class InputEmbedder(nn.Module):
69
- def __init__(
70
- self,
71
- atom_s: int,
72
- atom_z: int,
73
- token_s: int,
74
- token_z: int,
75
- atoms_per_window_queries: int,
76
- atoms_per_window_keys: int,
77
- atom_feature_dim: int,
78
- atom_encoder_depth: int,
79
- atom_encoder_heads: int,
80
- activation_checkpointing: bool = False,
81
- add_method_conditioning: bool = False,
82
- add_modified_flag: bool = False,
83
- add_cyclic_flag: bool = False,
84
- add_mol_type_feat: bool = False,
85
- use_no_atom_char: bool = False,
86
- use_atom_backbone_feat: bool = False,
87
- use_residue_feats_atoms: bool = False,
88
- ) -> None:
89
- """Initialize the input embedder.
90
-
91
- Parameters
92
- ----------
93
- atom_s : int
94
- The atom embedding size.
95
- atom_z : int
96
- The atom pairwise embedding size.
97
- token_s : int
98
- The token embedding size.
99
-
100
- """
101
- super().__init__()
102
- self.token_s = token_s
103
- self.add_method_conditioning = add_method_conditioning
104
- self.add_modified_flag = add_modified_flag
105
- self.add_cyclic_flag = add_cyclic_flag
106
- self.add_mol_type_feat = add_mol_type_feat
107
-
108
- self.atom_encoder = AtomEncoder(
109
- atom_s=atom_s,
110
- atom_z=atom_z,
111
- token_s=token_s,
112
- token_z=token_z,
113
- atoms_per_window_queries=atoms_per_window_queries,
114
- atoms_per_window_keys=atoms_per_window_keys,
115
- atom_feature_dim=atom_feature_dim,
116
- structure_prediction=False,
117
- use_no_atom_char=use_no_atom_char,
118
- use_atom_backbone_feat=use_atom_backbone_feat,
119
- use_residue_feats_atoms=use_residue_feats_atoms,
120
- )
121
-
122
- self.atom_enc_proj_z = nn.Sequential(
123
- nn.LayerNorm(atom_z),
124
- nn.Linear(atom_z, atom_encoder_depth * atom_encoder_heads, bias=False),
125
- )
126
-
127
- self.atom_attention_encoder = AtomAttentionEncoder(
128
- atom_s=atom_s,
129
- token_s=token_s,
130
- atoms_per_window_queries=atoms_per_window_queries,
131
- atoms_per_window_keys=atoms_per_window_keys,
132
- atom_encoder_depth=atom_encoder_depth,
133
- atom_encoder_heads=atom_encoder_heads,
134
- structure_prediction=False,
135
- activation_checkpointing=activation_checkpointing,
136
- )
137
-
138
- self.res_type_encoding = nn.Linear(const.num_tokens, token_s, bias=False)
139
- self.msa_profile_encoding = nn.Linear(const.num_tokens + 1, token_s, bias=False)
140
-
141
- if add_method_conditioning:
142
- self.method_conditioning_init = nn.Embedding(
143
- const.num_method_types, token_s
144
- )
145
- self.method_conditioning_init.weight.data.fill_(0)
146
- if add_modified_flag:
147
- self.modified_conditioning_init = nn.Embedding(2, token_s)
148
- self.modified_conditioning_init.weight.data.fill_(0)
149
- if add_cyclic_flag:
150
- self.cyclic_conditioning_init = nn.Linear(1, token_s, bias=False)
151
- self.cyclic_conditioning_init.weight.data.fill_(0)
152
- if add_mol_type_feat:
153
- self.mol_type_conditioning_init = nn.Embedding(
154
- len(const.chain_type_ids), token_s
155
- )
156
- self.mol_type_conditioning_init.weight.data.fill_(0)
157
-
158
- def forward(self, feats: dict[str, Tensor], affinity: bool = False) -> Tensor:
159
- """Perform the forward pass.
160
-
161
- Parameters
162
- ----------
163
- feats : dict[str, Tensor]
164
- Input features
165
-
166
- Returns
167
- -------
168
- Tensor
169
- The embedded tokens.
170
-
171
- """
172
- # Load relevant features
173
- res_type = feats["res_type"].float()
174
- if affinity:
175
- profile = feats["profile_affinity"]
176
- deletion_mean = feats["deletion_mean_affinity"].unsqueeze(-1)
177
- else:
178
- profile = feats["profile"]
179
- deletion_mean = feats["deletion_mean"].unsqueeze(-1)
180
-
181
- # Compute input embedding
182
- q, c, p, to_keys = self.atom_encoder(feats)
183
- atom_enc_bias = self.atom_enc_proj_z(p)
184
- a, _, _, _ = self.atom_attention_encoder(
185
- feats=feats,
186
- q=q,
187
- c=c,
188
- atom_enc_bias=atom_enc_bias,
189
- to_keys=to_keys,
190
- )
191
-
192
- s = (
193
- a
194
- + self.res_type_encoding(res_type)
195
- + self.msa_profile_encoding(torch.cat([profile, deletion_mean], dim=-1))
196
- )
197
-
198
- if self.add_method_conditioning:
199
- s = s + self.method_conditioning_init(feats["method_feature"])
200
- if self.add_modified_flag:
201
- s = s + self.modified_conditioning_init(feats["modified"])
202
- if self.add_cyclic_flag:
203
- cyclic = feats["cyclic_period"].clamp(max=1.0).unsqueeze(-1)
204
- s = s + self.cyclic_conditioning_init(cyclic)
205
- if self.add_mol_type_feat:
206
- s = s + self.mol_type_conditioning_init(feats["mol_type"])
207
-
208
- return s
209
-
210
-
211
- class TemplateModule(nn.Module):
212
- """Template module."""
213
-
214
- def __init__(
215
- self,
216
- token_z: int,
217
- template_dim: int,
218
- template_blocks: int,
219
- dropout: float = 0.25,
220
- pairwise_head_width: int = 32,
221
- pairwise_num_heads: int = 4,
222
- post_layer_norm: bool = False,
223
- activation_checkpointing: bool = False,
224
- min_dist: float = 3.25,
225
- max_dist: float = 50.75,
226
- num_bins: int = 38,
227
- **kwargs,
228
- ) -> None:
229
- """Initialize the template module.
230
-
231
- Parameters
232
- ----------
233
- token_z : int
234
- The token pairwise embedding size.
235
-
236
- """
237
- super().__init__()
238
- self.min_dist = min_dist
239
- self.max_dist = max_dist
240
- self.num_bins = num_bins
241
- self.relu = nn.ReLU()
242
- self.z_norm = nn.LayerNorm(token_z)
243
- self.v_norm = nn.LayerNorm(template_dim)
244
- self.z_proj = nn.Linear(token_z, template_dim, bias=False)
245
- self.a_proj = nn.Linear(
246
- const.num_tokens * 2 + num_bins + 5,
247
- template_dim,
248
- bias=False,
249
- )
250
- self.u_proj = nn.Linear(template_dim, token_z, bias=False)
251
- self.pairformer = PairformerNoSeqModule(
252
- template_dim,
253
- num_blocks=template_blocks,
254
- dropout=dropout,
255
- pairwise_head_width=pairwise_head_width,
256
- pairwise_num_heads=pairwise_num_heads,
257
- post_layer_norm=post_layer_norm,
258
- activation_checkpointing=activation_checkpointing,
259
- )
260
-
261
- def forward(
262
- self,
263
- z: Tensor,
264
- feats: dict[str, Tensor],
265
- pair_mask: Tensor,
266
- use_kernels: bool = False,
267
- ) -> Tensor:
268
- """Perform the forward pass.
269
-
270
- Parameters
271
- ----------
272
- z : Tensor
273
- The pairwise embeddings
274
- feats : dict[str, Tensor]
275
- Input features
276
- pair_mask : Tensor
277
- The pair mask
278
-
279
- Returns
280
- -------
281
- Tensor
282
- The updated pairwise embeddings.
283
-
284
- """
285
- # Load relevant features
286
- asym_id = feats["asym_id"]
287
- res_type = feats["template_restype"]
288
- frame_rot = feats["template_frame_rot"]
289
- frame_t = feats["template_frame_t"]
290
- frame_mask = feats["template_mask_frame"]
291
- cb_coords = feats["template_cb"]
292
- ca_coords = feats["template_ca"]
293
- cb_mask = feats["template_mask_cb"]
294
- template_mask = feats["template_mask"].any(dim=2).float()
295
- num_templates = template_mask.sum(dim=1)
296
- num_templates = num_templates.clamp(min=1)
297
-
298
- # Compute pairwise masks
299
- b_cb_mask = cb_mask[:, :, :, None] * cb_mask[:, :, None, :]
300
- b_frame_mask = frame_mask[:, :, :, None] * frame_mask[:, :, None, :]
301
-
302
- b_cb_mask = b_cb_mask[..., None]
303
- b_frame_mask = b_frame_mask[..., None]
304
-
305
- # Compute asym mask, template features only attend within the same chain
306
- B, T = res_type.shape[:2] # noqa: N806
307
- asym_mask = (asym_id[:, :, None] == asym_id[:, None, :]).float()
308
- asym_mask = asym_mask[:, None].expand(-1, T, -1, -1)
309
-
310
- # Compute template features
311
- with torch.autocast(device_type="cuda", enabled=False):
312
- # Compute distogram
313
- cb_dists = torch.cdist(cb_coords, cb_coords)
314
- boundaries = torch.linspace(self.min_dist, self.max_dist, self.num_bins - 1)
315
- boundaries = boundaries.to(cb_dists.device)
316
- distogram = (cb_dists[..., None] > boundaries).sum(dim=-1).long()
317
- distogram = one_hot(distogram, num_classes=self.num_bins)
318
-
319
- # Compute unit vector in each frame
320
- frame_rot = frame_rot.unsqueeze(2).transpose(-1, -2)
321
- frame_t = frame_t.unsqueeze(2).unsqueeze(-1)
322
- ca_coords = ca_coords.unsqueeze(3).unsqueeze(-1)
323
- vector = torch.matmul(frame_rot, (ca_coords - frame_t))
324
- norm = torch.norm(vector, dim=-1, keepdim=True)
325
- unit_vector = torch.where(norm > 0, vector / norm, torch.zeros_like(vector))
326
- unit_vector = unit_vector.squeeze(-1)
327
-
328
- # Concatenate input features
329
- a_tij = [distogram, b_cb_mask, unit_vector, b_frame_mask]
330
- a_tij = torch.cat(a_tij, dim=-1)
331
- a_tij = a_tij * asym_mask.unsqueeze(-1)
332
-
333
- res_type_i = res_type[:, :, :, None]
334
- res_type_j = res_type[:, :, None, :]
335
- res_type_i = res_type_i.expand(-1, -1, -1, res_type.size(2), -1)
336
- res_type_j = res_type_j.expand(-1, -1, res_type.size(2), -1, -1)
337
- a_tij = torch.cat([a_tij, res_type_i, res_type_j], dim=-1)
338
- a_tij = self.a_proj(a_tij)
339
-
340
- # Expand mask
341
- pair_mask = pair_mask[:, None].expand(-1, T, -1, -1)
342
- pair_mask = pair_mask.reshape(B * T, *pair_mask.shape[2:])
343
-
344
- # Compute input projections
345
- v = self.z_proj(self.z_norm(z[:, None])) + a_tij
346
- v = v.view(B * T, *v.shape[2:])
347
- v = v + self.pairformer(v, pair_mask, use_kernels=use_kernels)
348
- v = self.v_norm(v)
349
- v = v.view(B, T, *v.shape[1:])
350
-
351
- # Aggregate templates
352
- template_mask = template_mask[:, :, None, None, None]
353
- num_templates = num_templates[:, None, None, None]
354
- u = (v * template_mask).sum(dim=1) / num_templates.to(v)
355
-
356
- # Compute output projection
357
- u = self.u_proj(self.relu(u))
358
- return u
359
-
360
-
361
- class TemplateV2Module(nn.Module):
362
- """Template module."""
363
-
364
- def __init__(
365
- self,
366
- token_z: int,
367
- template_dim: int,
368
- template_blocks: int,
369
- dropout: float = 0.25,
370
- pairwise_head_width: int = 32,
371
- pairwise_num_heads: int = 4,
372
- post_layer_norm: bool = False,
373
- activation_checkpointing: bool = False,
374
- min_dist: float = 3.25,
375
- max_dist: float = 50.75,
376
- num_bins: int = 38,
377
- **kwargs,
378
- ) -> None:
379
- """Initialize the template module.
380
-
381
- Parameters
382
- ----------
383
- token_z : int
384
- The token pairwise embedding size.
385
-
386
- """
387
- super().__init__()
388
- self.min_dist = min_dist
389
- self.max_dist = max_dist
390
- self.num_bins = num_bins
391
- self.relu = nn.ReLU()
392
- self.z_norm = nn.LayerNorm(token_z)
393
- self.v_norm = nn.LayerNorm(template_dim)
394
- self.z_proj = nn.Linear(token_z, template_dim, bias=False)
395
- self.a_proj = nn.Linear(
396
- const.num_tokens * 2 + num_bins + 5,
397
- template_dim,
398
- bias=False,
399
- )
400
- self.u_proj = nn.Linear(template_dim, token_z, bias=False)
401
- self.pairformer = PairformerNoSeqModule(
402
- template_dim,
403
- num_blocks=template_blocks,
404
- dropout=dropout,
405
- pairwise_head_width=pairwise_head_width,
406
- pairwise_num_heads=pairwise_num_heads,
407
- post_layer_norm=post_layer_norm,
408
- activation_checkpointing=activation_checkpointing,
409
- )
410
-
411
- def forward(
412
- self,
413
- z: Tensor,
414
- feats: dict[str, Tensor],
415
- pair_mask: Tensor,
416
- use_kernels: bool = False,
417
- ) -> Tensor:
418
- """Perform the forward pass.
419
-
420
- Parameters
421
- ----------
422
- z : Tensor
423
- The pairwise embeddings
424
- feats : dict[str, Tensor]
425
- Input features
426
- pair_mask : Tensor
427
- The pair mask
428
-
429
- Returns
430
- -------
431
- Tensor
432
- The updated pairwise embeddings.
433
-
434
- """
435
- # Load relevant features
436
- res_type = feats["template_restype"]
437
- frame_rot = feats["template_frame_rot"]
438
- frame_t = feats["template_frame_t"]
439
- frame_mask = feats["template_mask_frame"]
440
- cb_coords = feats["template_cb"]
441
- ca_coords = feats["template_ca"]
442
- cb_mask = feats["template_mask_cb"]
443
- visibility_ids = feats["visibility_ids"]
444
- template_mask = feats["template_mask"].any(dim=2).float()
445
- num_templates = template_mask.sum(dim=1)
446
- num_templates = num_templates.clamp(min=1)
447
-
448
- # Compute pairwise masks
449
- b_cb_mask = cb_mask[:, :, :, None] * cb_mask[:, :, None, :]
450
- b_frame_mask = frame_mask[:, :, :, None] * frame_mask[:, :, None, :]
451
-
452
- b_cb_mask = b_cb_mask[..., None]
453
- b_frame_mask = b_frame_mask[..., None]
454
-
455
- # Compute asym mask, template features only attend within the same chain
456
- B, T = res_type.shape[:2] # noqa: N806
457
- tmlp_pair_mask = (
458
- visibility_ids[:, :, :, None] == visibility_ids[:, :, None, :]
459
- ).float()
460
-
461
- # Compute template features
462
- with torch.autocast(device_type="cuda", enabled=False):
463
- # Compute distogram
464
- cb_dists = torch.cdist(cb_coords, cb_coords)
465
- boundaries = torch.linspace(self.min_dist, self.max_dist, self.num_bins - 1)
466
- boundaries = boundaries.to(cb_dists.device)
467
- distogram = (cb_dists[..., None] > boundaries).sum(dim=-1).long()
468
- distogram = one_hot(distogram, num_classes=self.num_bins)
469
-
470
- # Compute unit vector in each frame
471
- frame_rot = frame_rot.unsqueeze(2).transpose(-1, -2)
472
- frame_t = frame_t.unsqueeze(2).unsqueeze(-1)
473
- ca_coords = ca_coords.unsqueeze(3).unsqueeze(-1)
474
- vector = torch.matmul(frame_rot, (ca_coords - frame_t))
475
- norm = torch.norm(vector, dim=-1, keepdim=True)
476
- unit_vector = torch.where(norm > 0, vector / norm, torch.zeros_like(vector))
477
- unit_vector = unit_vector.squeeze(-1)
478
-
479
- # Concatenate input features
480
- a_tij = [distogram, b_cb_mask, unit_vector, b_frame_mask]
481
- a_tij = torch.cat(a_tij, dim=-1)
482
- a_tij = a_tij * tmlp_pair_mask.unsqueeze(-1)
483
-
484
- res_type_i = res_type[:, :, :, None]
485
- res_type_j = res_type[:, :, None, :]
486
- res_type_i = res_type_i.expand(-1, -1, -1, res_type.size(2), -1)
487
- res_type_j = res_type_j.expand(-1, -1, res_type.size(2), -1, -1)
488
- a_tij = torch.cat([a_tij, res_type_i, res_type_j], dim=-1)
489
- a_tij = self.a_proj(a_tij)
490
-
491
- # Expand mask
492
- pair_mask = pair_mask[:, None].expand(-1, T, -1, -1)
493
- pair_mask = pair_mask.reshape(B * T, *pair_mask.shape[2:])
494
-
495
- # Compute input projections
496
- v = self.z_proj(self.z_norm(z[:, None])) + a_tij
497
- v = v.view(B * T, *v.shape[2:])
498
- v = v + self.pairformer(v, pair_mask, use_kernels=use_kernels)
499
- v = self.v_norm(v)
500
- v = v.view(B, T, *v.shape[1:])
501
-
502
- # Aggregate templates
503
- template_mask = template_mask[:, :, None, None, None]
504
- num_templates = num_templates[:, None, None, None]
505
- u = (v * template_mask).sum(dim=1) / num_templates.to(v)
506
-
507
- # Compute output projection
508
- u = self.u_proj(self.relu(u))
509
- return u
510
-
511
-
512
- class MSAModule(nn.Module):
513
- """MSA module."""
514
-
515
- def __init__(
516
- self,
517
- msa_s: int,
518
- token_z: int,
519
- token_s: int,
520
- msa_blocks: int,
521
- msa_dropout: float,
522
- z_dropout: float,
523
- pairwise_head_width: int = 32,
524
- pairwise_num_heads: int = 4,
525
- activation_checkpointing: bool = False,
526
- use_paired_feature: bool = True,
527
- subsample_msa: bool = False,
528
- num_subsampled_msa: int = 1024,
529
- **kwargs,
530
- ) -> None:
531
- """Initialize the MSA module.
532
-
533
- Parameters
534
- ----------
535
- token_z : int
536
- The token pairwise embedding size.
537
-
538
- """
539
- super().__init__()
540
- self.msa_blocks = msa_blocks
541
- self.msa_dropout = msa_dropout
542
- self.z_dropout = z_dropout
543
- self.use_paired_feature = use_paired_feature
544
- self.activation_checkpointing = activation_checkpointing
545
- self.subsample_msa = subsample_msa
546
- self.num_subsampled_msa = num_subsampled_msa
547
-
548
- self.s_proj = nn.Linear(token_s, msa_s, bias=False)
549
- self.msa_proj = nn.Linear(
550
- const.num_tokens + 2 + int(use_paired_feature),
551
- msa_s,
552
- bias=False,
553
- )
554
- self.layers = nn.ModuleList()
555
- for i in range(msa_blocks):
556
- self.layers.append(
557
- MSALayer(
558
- msa_s,
559
- token_z,
560
- msa_dropout,
561
- z_dropout,
562
- pairwise_head_width,
563
- pairwise_num_heads,
564
- )
565
- )
566
-
567
- def forward(
568
- self,
569
- z: Tensor,
570
- emb: Tensor,
571
- feats: dict[str, Tensor],
572
- use_kernels: bool = False,
573
- ) -> Tensor:
574
- """Perform the forward pass.
575
-
576
- Parameters
577
- ----------
578
- z : Tensor
579
- The pairwise embeddings
580
- emb : Tensor
581
- The input embeddings
582
- feats : dict[str, Tensor]
583
- Input features
584
- use_kernels: bool
585
- Whether to use kernels for triangular updates
586
-
587
- Returns
588
- -------
589
- Tensor
590
- The output pairwise embeddings.
591
-
592
- """
593
- # Set chunk sizes
594
- if not self.training:
595
- if z.shape[1] > const.chunk_size_threshold:
596
- chunk_heads_pwa = True
597
- chunk_size_transition_z = 64
598
- chunk_size_transition_msa = 32
599
- chunk_size_outer_product = 4
600
- chunk_size_tri_attn = 128
601
- else:
602
- chunk_heads_pwa = False
603
- chunk_size_transition_z = None
604
- chunk_size_transition_msa = None
605
- chunk_size_outer_product = None
606
- chunk_size_tri_attn = 512
607
- else:
608
- chunk_heads_pwa = False
609
- chunk_size_transition_z = None
610
- chunk_size_transition_msa = None
611
- chunk_size_outer_product = None
612
- chunk_size_tri_attn = None
613
-
614
- # Load relevant features
615
- msa = feats["msa"]
616
- msa = torch.nn.functional.one_hot(msa, num_classes=const.num_tokens)
617
- has_deletion = feats["has_deletion"].unsqueeze(-1)
618
- deletion_value = feats["deletion_value"].unsqueeze(-1)
619
- is_paired = feats["msa_paired"].unsqueeze(-1)
620
- msa_mask = feats["msa_mask"]
621
- token_mask = feats["token_pad_mask"].float()
622
- token_mask = token_mask[:, :, None] * token_mask[:, None, :]
623
-
624
- # Compute MSA embeddings
625
- if self.use_paired_feature:
626
- m = torch.cat([msa, has_deletion, deletion_value, is_paired], dim=-1)
627
- else:
628
- m = torch.cat([msa, has_deletion, deletion_value], dim=-1)
629
-
630
- # Subsample the MSA
631
- if self.subsample_msa:
632
- msa_indices = torch.randperm(msa.shape[1])[: self.num_subsampled_msa]
633
- m = m[:, msa_indices]
634
- msa_mask = msa_mask[:, msa_indices]
635
-
636
- # Compute input projections
637
- m = self.msa_proj(m)
638
- m = m + self.s_proj(emb).unsqueeze(1)
639
-
640
- # Perform MSA blocks
641
- for i in range(self.msa_blocks):
642
- if self.activation_checkpointing and self.training:
643
- z, m = torch.utils.checkpoint.checkpoint(
644
- self.layers[i],
645
- z,
646
- m,
647
- token_mask,
648
- msa_mask,
649
- chunk_heads_pwa,
650
- chunk_size_transition_z,
651
- chunk_size_transition_msa,
652
- chunk_size_outer_product,
653
- chunk_size_tri_attn,
654
- use_kernels,
655
- )
656
- else:
657
- z, m = self.layers[i](
658
- z,
659
- m,
660
- token_mask,
661
- msa_mask,
662
- chunk_heads_pwa,
663
- chunk_size_transition_z,
664
- chunk_size_transition_msa,
665
- chunk_size_outer_product,
666
- chunk_size_tri_attn,
667
- use_kernels,
668
- )
669
- return z
670
-
671
-
672
- class MSALayer(nn.Module):
673
- """MSA module."""
674
-
675
- def __init__(
676
- self,
677
- msa_s: int,
678
- token_z: int,
679
- msa_dropout: float,
680
- z_dropout: float,
681
- pairwise_head_width: int = 32,
682
- pairwise_num_heads: int = 4,
683
- ) -> None:
684
- """Initialize the MSA module.
685
-
686
- Parameters
687
- ----------
688
- token_z : int
689
- The token pairwise embedding size.
690
-
691
- """
692
- super().__init__()
693
- self.msa_dropout = msa_dropout
694
- self.msa_transition = Transition(dim=msa_s, hidden=msa_s * 4)
695
- self.pair_weighted_averaging = PairWeightedAveraging(
696
- c_m=msa_s,
697
- c_z=token_z,
698
- c_h=32,
699
- num_heads=8,
700
- )
701
-
702
- self.pairformer_layer = PairformerNoSeqLayer(
703
- token_z=token_z,
704
- dropout=z_dropout,
705
- pairwise_head_width=pairwise_head_width,
706
- pairwise_num_heads=pairwise_num_heads,
707
- )
708
- self.outer_product_mean = OuterProductMean(
709
- c_in=msa_s,
710
- c_hidden=32,
711
- c_out=token_z,
712
- )
713
-
714
- def forward(
715
- self,
716
- z: Tensor,
717
- m: Tensor,
718
- token_mask: Tensor,
719
- msa_mask: Tensor,
720
- chunk_heads_pwa: bool = False,
721
- chunk_size_transition_z: int = None,
722
- chunk_size_transition_msa: int = None,
723
- chunk_size_outer_product: int = None,
724
- chunk_size_tri_attn: int = None,
725
- use_kernels: bool = False,
726
- ) -> tuple[Tensor, Tensor]:
727
- """Perform the forward pass.
728
-
729
- Parameters
730
- ----------
731
- z : Tensor
732
- The pairwise embeddings
733
- emb : Tensor
734
- The input embeddings
735
- feats : dict[str, Tensor]
736
- Input features
737
-
738
- Returns
739
- -------
740
- Tensor
741
- The output pairwise embeddings.
742
-
743
- """
744
- # Communication to MSA stack
745
- msa_dropout = get_dropout_mask(self.msa_dropout, m, self.training)
746
- m = m + msa_dropout * self.pair_weighted_averaging(
747
- m, z, token_mask, chunk_heads_pwa
748
- )
749
- m = m + self.msa_transition(m, chunk_size_transition_msa)
750
-
751
- z = z + self.outer_product_mean(m, msa_mask, chunk_size_outer_product)
752
-
753
- # Compute pairwise stack
754
- z = self.pairformer_layer(
755
- z, token_mask, chunk_size_tri_attn, use_kernels=use_kernels
756
- )
757
-
758
- return z, m
759
-
760
-
761
- class BFactorModule(nn.Module):
762
- """BFactor Module."""
763
-
764
- def __init__(self, token_s: int, num_bins: int) -> None:
765
- """Initialize the bfactor module.
766
-
767
- Parameters
768
- ----------
769
- token_s : int
770
- The token embedding size.
771
-
772
- """
773
- super().__init__()
774
- self.bfactor = nn.Linear(token_s, num_bins)
775
- self.num_bins = num_bins
776
-
777
- def forward(self, s: Tensor) -> Tensor:
778
- """Perform the forward pass.
779
-
780
- Parameters
781
- ----------
782
- s : Tensor
783
- The sequence embeddings
784
-
785
- Returns
786
- -------
787
- Tensor
788
- The predicted bfactor histogram.
789
-
790
- """
791
- return self.bfactor(s)
792
-
793
-
794
- class DistogramModule(nn.Module):
795
- """Distogram Module."""
796
-
797
- def __init__(self, token_z: int, num_bins: int, num_distograms: int = 1) -> None:
798
- """Initialize the distogram module.
799
-
800
- Parameters
801
- ----------
802
- token_z : int
803
- The token pairwise embedding size.
804
-
805
- """
806
- super().__init__()
807
- self.distogram = nn.Linear(token_z, num_distograms * num_bins)
808
- self.num_distograms = num_distograms
809
- self.num_bins = num_bins
810
-
811
- def forward(self, z: Tensor) -> Tensor:
812
- """Perform the forward pass.
813
-
814
- Parameters
815
- ----------
816
- z : Tensor
817
- The pairwise embeddings
818
-
819
- Returns
820
- -------
821
- Tensor
822
- The predicted distogram.
823
-
824
- """
825
- z = z + z.transpose(1, 2)
826
- return self.distogram(z).reshape(
827
- z.shape[0], z.shape[1], z.shape[2], self.num_distograms, self.num_bins
828
- )
 
1
+ import torch
2
+ from torch import Tensor, nn
3
+ from torch.nn.functional import one_hot
4
+
5
+ from . import vb_const as const
6
+ from .vb_layers_outer_product_mean import OuterProductMean
7
+ from .vb_layers_pair_averaging import PairWeightedAveraging
8
+ from .vb_layers_pairformer import (
9
+ PairformerNoSeqLayer,
10
+ PairformerNoSeqModule,
11
+ get_dropout_mask,
12
+ )
13
+ from .vb_layers_transition import Transition
14
+ from .vb_modules_encodersv2 import (
15
+ AtomAttentionEncoder,
16
+ AtomEncoder,
17
+ FourierEmbedding,
18
+ )
19
+
20
+
21
+ class ContactConditioning(nn.Module):
22
+ def __init__(self, token_z: int, cutoff_min: float, cutoff_max: float):
23
+ super().__init__()
24
+
25
+ self.fourier_embedding = FourierEmbedding(token_z)
26
+ self.encoder = nn.Linear(
27
+ token_z + len(const.contact_conditioning_info) - 1, token_z
28
+ )
29
+ self.encoding_unspecified = nn.Parameter(torch.zeros(token_z))
30
+ self.encoding_unselected = nn.Parameter(torch.zeros(token_z))
31
+ self.cutoff_min = cutoff_min
32
+ self.cutoff_max = cutoff_max
33
+
34
+ def forward(self, feats):
35
+ assert const.contact_conditioning_info["UNSPECIFIED"] == 0
36
+ assert const.contact_conditioning_info["UNSELECTED"] == 1
37
+ contact_conditioning = feats["contact_conditioning"][:, :, :, 2:]
38
+ contact_threshold = feats["contact_threshold"]
39
+ contact_threshold_normalized = (contact_threshold - self.cutoff_min) / (
40
+ self.cutoff_max - self.cutoff_min
41
+ )
42
+ contact_threshold_fourier = self.fourier_embedding(
43
+ contact_threshold_normalized.flatten()
44
+ ).reshape(contact_threshold_normalized.shape + (-1,))
45
+
46
+ contact_conditioning = torch.cat(
47
+ [
48
+ contact_conditioning,
49
+ contact_threshold_normalized.unsqueeze(-1),
50
+ contact_threshold_fourier,
51
+ ],
52
+ dim=-1,
53
+ )
54
+ contact_conditioning = self.encoder(contact_conditioning)
55
+
56
+ contact_conditioning = (
57
+ contact_conditioning
58
+ * (
59
+ 1
60
+ - feats["contact_conditioning"][:, :, :, 0:2].sum(dim=-1, keepdim=True)
61
+ )
62
+ + self.encoding_unspecified * feats["contact_conditioning"][:, :, :, 0:1]
63
+ + self.encoding_unselected * feats["contact_conditioning"][:, :, :, 1:2]
64
+ )
65
+ return contact_conditioning
66
+
67
+
68
+ class InputEmbedder(nn.Module):
69
+ def __init__(
70
+ self,
71
+ atom_s: int,
72
+ atom_z: int,
73
+ token_s: int,
74
+ token_z: int,
75
+ atoms_per_window_queries: int,
76
+ atoms_per_window_keys: int,
77
+ atom_feature_dim: int,
78
+ atom_encoder_depth: int,
79
+ atom_encoder_heads: int,
80
+ activation_checkpointing: bool = False,
81
+ add_method_conditioning: bool = False,
82
+ add_modified_flag: bool = False,
83
+ add_cyclic_flag: bool = False,
84
+ add_mol_type_feat: bool = False,
85
+ use_no_atom_char: bool = False,
86
+ use_atom_backbone_feat: bool = False,
87
+ use_residue_feats_atoms: bool = False,
88
+ ) -> None:
89
+ """Initialize the input embedder.
90
+
91
+ Parameters
92
+ ----------
93
+ atom_s : int
94
+ The atom embedding size.
95
+ atom_z : int
96
+ The atom pairwise embedding size.
97
+ token_s : int
98
+ The token embedding size.
99
+
100
+ """
101
+ super().__init__()
102
+ self.token_s = token_s
103
+ self.add_method_conditioning = add_method_conditioning
104
+ self.add_modified_flag = add_modified_flag
105
+ self.add_cyclic_flag = add_cyclic_flag
106
+ self.add_mol_type_feat = add_mol_type_feat
107
+
108
+ self.atom_encoder = AtomEncoder(
109
+ atom_s=atom_s,
110
+ atom_z=atom_z,
111
+ token_s=token_s,
112
+ token_z=token_z,
113
+ atoms_per_window_queries=atoms_per_window_queries,
114
+ atoms_per_window_keys=atoms_per_window_keys,
115
+ atom_feature_dim=atom_feature_dim,
116
+ structure_prediction=False,
117
+ use_no_atom_char=use_no_atom_char,
118
+ use_atom_backbone_feat=use_atom_backbone_feat,
119
+ use_residue_feats_atoms=use_residue_feats_atoms,
120
+ )
121
+
122
+ self.atom_enc_proj_z = nn.Sequential(
123
+ nn.LayerNorm(atom_z),
124
+ nn.Linear(atom_z, atom_encoder_depth * atom_encoder_heads, bias=False),
125
+ )
126
+
127
+ self.atom_attention_encoder = AtomAttentionEncoder(
128
+ atom_s=atom_s,
129
+ token_s=token_s,
130
+ atoms_per_window_queries=atoms_per_window_queries,
131
+ atoms_per_window_keys=atoms_per_window_keys,
132
+ atom_encoder_depth=atom_encoder_depth,
133
+ atom_encoder_heads=atom_encoder_heads,
134
+ structure_prediction=False,
135
+ activation_checkpointing=activation_checkpointing,
136
+ )
137
+
138
+ self.res_type_encoding = nn.Linear(const.num_tokens, token_s, bias=False)
139
+ self.msa_profile_encoding = nn.Linear(const.num_tokens + 1, token_s, bias=False)
140
+
141
+ if add_method_conditioning:
142
+ self.method_conditioning_init = nn.Embedding(
143
+ const.num_method_types, token_s
144
+ )
145
+ self.method_conditioning_init.weight.data.fill_(0)
146
+ if add_modified_flag:
147
+ self.modified_conditioning_init = nn.Embedding(2, token_s)
148
+ self.modified_conditioning_init.weight.data.fill_(0)
149
+ if add_cyclic_flag:
150
+ self.cyclic_conditioning_init = nn.Linear(1, token_s, bias=False)
151
+ self.cyclic_conditioning_init.weight.data.fill_(0)
152
+ if add_mol_type_feat:
153
+ self.mol_type_conditioning_init = nn.Embedding(
154
+ len(const.chain_type_ids), token_s
155
+ )
156
+ self.mol_type_conditioning_init.weight.data.fill_(0)
157
+
158
+ def forward(self, feats: dict[str, Tensor], affinity: bool = False) -> Tensor:
159
+ """Perform the forward pass.
160
+
161
+ Parameters
162
+ ----------
163
+ feats : dict[str, Tensor]
164
+ Input features
165
+
166
+ Returns
167
+ -------
168
+ Tensor
169
+ The embedded tokens.
170
+
171
+ """
172
+ # Load relevant features
173
+ res_type = feats["res_type"].float()
174
+ if affinity:
175
+ profile = feats["profile_affinity"]
176
+ deletion_mean = feats["deletion_mean_affinity"].unsqueeze(-1)
177
+ else:
178
+ profile = feats["profile"]
179
+ deletion_mean = feats["deletion_mean"].unsqueeze(-1)
180
+
181
+ # Compute input embedding
182
+ q, c, p, to_keys = self.atom_encoder(feats)
183
+ atom_enc_bias = self.atom_enc_proj_z(p)
184
+ a, _, _, _ = self.atom_attention_encoder(
185
+ feats=feats,
186
+ q=q,
187
+ c=c,
188
+ atom_enc_bias=atom_enc_bias,
189
+ to_keys=to_keys,
190
+ )
191
+
192
+ s = (
193
+ a
194
+ + self.res_type_encoding(res_type)
195
+ + self.msa_profile_encoding(torch.cat([profile, deletion_mean], dim=-1))
196
+ )
197
+
198
+ if self.add_method_conditioning:
199
+ s = s + self.method_conditioning_init(feats["method_feature"])
200
+ if self.add_modified_flag:
201
+ s = s + self.modified_conditioning_init(feats["modified"])
202
+ if self.add_cyclic_flag:
203
+ cyclic = feats["cyclic_period"].clamp(max=1.0).unsqueeze(-1)
204
+ s = s + self.cyclic_conditioning_init(cyclic)
205
+ if self.add_mol_type_feat:
206
+ s = s + self.mol_type_conditioning_init(feats["mol_type"])
207
+
208
+ return s
209
+
210
+
211
+ class TemplateModule(nn.Module):
212
+ """Template module."""
213
+
214
+ def __init__(
215
+ self,
216
+ token_z: int,
217
+ template_dim: int,
218
+ template_blocks: int,
219
+ dropout: float = 0.25,
220
+ pairwise_head_width: int = 32,
221
+ pairwise_num_heads: int = 4,
222
+ post_layer_norm: bool = False,
223
+ activation_checkpointing: bool = False,
224
+ min_dist: float = 3.25,
225
+ max_dist: float = 50.75,
226
+ num_bins: int = 38,
227
+ **kwargs,
228
+ ) -> None:
229
+ """Initialize the template module.
230
+
231
+ Parameters
232
+ ----------
233
+ token_z : int
234
+ The token pairwise embedding size.
235
+
236
+ """
237
+ super().__init__()
238
+ self.min_dist = min_dist
239
+ self.max_dist = max_dist
240
+ self.num_bins = num_bins
241
+ self.relu = nn.ReLU()
242
+ self.z_norm = nn.LayerNorm(token_z)
243
+ self.v_norm = nn.LayerNorm(template_dim)
244
+ self.z_proj = nn.Linear(token_z, template_dim, bias=False)
245
+ self.a_proj = nn.Linear(
246
+ const.num_tokens * 2 + num_bins + 5,
247
+ template_dim,
248
+ bias=False,
249
+ )
250
+ self.u_proj = nn.Linear(template_dim, token_z, bias=False)
251
+ self.pairformer = PairformerNoSeqModule(
252
+ template_dim,
253
+ num_blocks=template_blocks,
254
+ dropout=dropout,
255
+ pairwise_head_width=pairwise_head_width,
256
+ pairwise_num_heads=pairwise_num_heads,
257
+ post_layer_norm=post_layer_norm,
258
+ activation_checkpointing=activation_checkpointing,
259
+ )
260
+
261
+ def forward(
262
+ self,
263
+ z: Tensor,
264
+ feats: dict[str, Tensor],
265
+ pair_mask: Tensor,
266
+ use_kernels: bool = False,
267
+ ) -> Tensor:
268
+ """Perform the forward pass.
269
+
270
+ Parameters
271
+ ----------
272
+ z : Tensor
273
+ The pairwise embeddings
274
+ feats : dict[str, Tensor]
275
+ Input features
276
+ pair_mask : Tensor
277
+ The pair mask
278
+
279
+ Returns
280
+ -------
281
+ Tensor
282
+ The updated pairwise embeddings.
283
+
284
+ """
285
+ # Load relevant features
286
+ asym_id = feats["asym_id"]
287
+ res_type = feats["template_restype"]
288
+ frame_rot = feats["template_frame_rot"]
289
+ frame_t = feats["template_frame_t"]
290
+ frame_mask = feats["template_mask_frame"]
291
+ cb_coords = feats["template_cb"]
292
+ ca_coords = feats["template_ca"]
293
+ cb_mask = feats["template_mask_cb"]
294
+ template_mask = feats["template_mask"].any(dim=2).float()
295
+ num_templates = template_mask.sum(dim=1)
296
+ num_templates = num_templates.clamp(min=1)
297
+
298
+ # Compute pairwise masks
299
+ b_cb_mask = cb_mask[:, :, :, None] * cb_mask[:, :, None, :]
300
+ b_frame_mask = frame_mask[:, :, :, None] * frame_mask[:, :, None, :]
301
+
302
+ b_cb_mask = b_cb_mask[..., None]
303
+ b_frame_mask = b_frame_mask[..., None]
304
+
305
+ # Compute asym mask, template features only attend within the same chain
306
+ B, T = res_type.shape[:2] # noqa: N806
307
+ asym_mask = (asym_id[:, :, None] == asym_id[:, None, :]).float()
308
+ asym_mask = asym_mask[:, None].expand(-1, T, -1, -1)
309
+
310
+ # Compute template features
311
+ with torch.autocast(device_type="cuda", enabled=False):
312
+ # Compute distogram
313
+ cb_dists = torch.cdist(cb_coords, cb_coords)
314
+ boundaries = torch.linspace(self.min_dist, self.max_dist, self.num_bins - 1)
315
+ boundaries = boundaries.to(cb_dists.device)
316
+ distogram = (cb_dists[..., None] > boundaries).sum(dim=-1).long()
317
+ distogram = one_hot(distogram, num_classes=self.num_bins)
318
+
319
+ # Compute unit vector in each frame
320
+ frame_rot = frame_rot.unsqueeze(2).transpose(-1, -2)
321
+ frame_t = frame_t.unsqueeze(2).unsqueeze(-1)
322
+ ca_coords = ca_coords.unsqueeze(3).unsqueeze(-1)
323
+ vector = torch.matmul(frame_rot, (ca_coords - frame_t))
324
+ norm = torch.norm(vector, dim=-1, keepdim=True)
325
+ unit_vector = torch.where(norm > 0, vector / norm, torch.zeros_like(vector))
326
+ unit_vector = unit_vector.squeeze(-1)
327
+
328
+ # Concatenate input features
329
+ a_tij = [distogram, b_cb_mask, unit_vector, b_frame_mask]
330
+ a_tij = torch.cat(a_tij, dim=-1)
331
+ a_tij = a_tij * asym_mask.unsqueeze(-1)
332
+
333
+ res_type_i = res_type[:, :, :, None]
334
+ res_type_j = res_type[:, :, None, :]
335
+ res_type_i = res_type_i.expand(-1, -1, -1, res_type.size(2), -1)
336
+ res_type_j = res_type_j.expand(-1, -1, res_type.size(2), -1, -1)
337
+ a_tij = torch.cat([a_tij, res_type_i, res_type_j], dim=-1)
338
+ a_tij = self.a_proj(a_tij)
339
+
340
+ # Expand mask
341
+ pair_mask = pair_mask[:, None].expand(-1, T, -1, -1)
342
+ pair_mask = pair_mask.reshape(B * T, *pair_mask.shape[2:])
343
+
344
+ # Compute input projections
345
+ v = self.z_proj(self.z_norm(z[:, None])) + a_tij
346
+ v = v.view(B * T, *v.shape[2:])
347
+ v = v + self.pairformer(v, pair_mask, use_kernels=use_kernels)
348
+ v = self.v_norm(v)
349
+ v = v.view(B, T, *v.shape[1:])
350
+
351
+ # Aggregate templates
352
+ template_mask = template_mask[:, :, None, None, None]
353
+ num_templates = num_templates[:, None, None, None]
354
+ u = (v * template_mask).sum(dim=1) / num_templates.to(v)
355
+
356
+ # Compute output projection
357
+ u = self.u_proj(self.relu(u))
358
+ return u
359
+
360
+
361
+ class TemplateV2Module(nn.Module):
362
+ """Template module."""
363
+
364
+ def __init__(
365
+ self,
366
+ token_z: int,
367
+ template_dim: int,
368
+ template_blocks: int,
369
+ dropout: float = 0.25,
370
+ pairwise_head_width: int = 32,
371
+ pairwise_num_heads: int = 4,
372
+ post_layer_norm: bool = False,
373
+ activation_checkpointing: bool = False,
374
+ min_dist: float = 3.25,
375
+ max_dist: float = 50.75,
376
+ num_bins: int = 38,
377
+ **kwargs,
378
+ ) -> None:
379
+ """Initialize the template module.
380
+
381
+ Parameters
382
+ ----------
383
+ token_z : int
384
+ The token pairwise embedding size.
385
+
386
+ """
387
+ super().__init__()
388
+ self.min_dist = min_dist
389
+ self.max_dist = max_dist
390
+ self.num_bins = num_bins
391
+ self.relu = nn.ReLU()
392
+ self.z_norm = nn.LayerNorm(token_z)
393
+ self.v_norm = nn.LayerNorm(template_dim)
394
+ self.z_proj = nn.Linear(token_z, template_dim, bias=False)
395
+ self.a_proj = nn.Linear(
396
+ const.num_tokens * 2 + num_bins + 5,
397
+ template_dim,
398
+ bias=False,
399
+ )
400
+ self.u_proj = nn.Linear(template_dim, token_z, bias=False)
401
+ self.pairformer = PairformerNoSeqModule(
402
+ template_dim,
403
+ num_blocks=template_blocks,
404
+ dropout=dropout,
405
+ pairwise_head_width=pairwise_head_width,
406
+ pairwise_num_heads=pairwise_num_heads,
407
+ post_layer_norm=post_layer_norm,
408
+ activation_checkpointing=activation_checkpointing,
409
+ )
410
+
411
+ def forward(
412
+ self,
413
+ z: Tensor,
414
+ feats: dict[str, Tensor],
415
+ pair_mask: Tensor,
416
+ use_kernels: bool = False,
417
+ ) -> Tensor:
418
+ """Perform the forward pass.
419
+
420
+ Parameters
421
+ ----------
422
+ z : Tensor
423
+ The pairwise embeddings
424
+ feats : dict[str, Tensor]
425
+ Input features
426
+ pair_mask : Tensor
427
+ The pair mask
428
+
429
+ Returns
430
+ -------
431
+ Tensor
432
+ The updated pairwise embeddings.
433
+
434
+ """
435
+ # Load relevant features
436
+ res_type = feats["template_restype"]
437
+ frame_rot = feats["template_frame_rot"]
438
+ frame_t = feats["template_frame_t"]
439
+ frame_mask = feats["template_mask_frame"]
440
+ cb_coords = feats["template_cb"]
441
+ ca_coords = feats["template_ca"]
442
+ cb_mask = feats["template_mask_cb"]
443
+ visibility_ids = feats["visibility_ids"]
444
+ template_mask = feats["template_mask"].any(dim=2).float()
445
+ num_templates = template_mask.sum(dim=1)
446
+ num_templates = num_templates.clamp(min=1)
447
+
448
+ # Compute pairwise masks
449
+ b_cb_mask = cb_mask[:, :, :, None] * cb_mask[:, :, None, :]
450
+ b_frame_mask = frame_mask[:, :, :, None] * frame_mask[:, :, None, :]
451
+
452
+ b_cb_mask = b_cb_mask[..., None]
453
+ b_frame_mask = b_frame_mask[..., None]
454
+
455
+ # Compute asym mask, template features only attend within the same chain
456
+ B, T = res_type.shape[:2] # noqa: N806
457
+ tmlp_pair_mask = (
458
+ visibility_ids[:, :, :, None] == visibility_ids[:, :, None, :]
459
+ ).float()
460
+
461
+ # Compute template features
462
+ with torch.autocast(device_type="cuda", enabled=False):
463
+ # Compute distogram
464
+ cb_dists = torch.cdist(cb_coords, cb_coords)
465
+ boundaries = torch.linspace(self.min_dist, self.max_dist, self.num_bins - 1)
466
+ boundaries = boundaries.to(cb_dists.device)
467
+ distogram = (cb_dists[..., None] > boundaries).sum(dim=-1).long()
468
+ distogram = one_hot(distogram, num_classes=self.num_bins)
469
+
470
+ # Compute unit vector in each frame
471
+ frame_rot = frame_rot.unsqueeze(2).transpose(-1, -2)
472
+ frame_t = frame_t.unsqueeze(2).unsqueeze(-1)
473
+ ca_coords = ca_coords.unsqueeze(3).unsqueeze(-1)
474
+ vector = torch.matmul(frame_rot, (ca_coords - frame_t))
475
+ norm = torch.norm(vector, dim=-1, keepdim=True)
476
+ unit_vector = torch.where(norm > 0, vector / norm, torch.zeros_like(vector))
477
+ unit_vector = unit_vector.squeeze(-1)
478
+
479
+ # Concatenate input features
480
+ a_tij = [distogram, b_cb_mask, unit_vector, b_frame_mask]
481
+ a_tij = torch.cat(a_tij, dim=-1)
482
+ a_tij = a_tij * tmlp_pair_mask.unsqueeze(-1)
483
+
484
+ res_type_i = res_type[:, :, :, None]
485
+ res_type_j = res_type[:, :, None, :]
486
+ res_type_i = res_type_i.expand(-1, -1, -1, res_type.size(2), -1)
487
+ res_type_j = res_type_j.expand(-1, -1, res_type.size(2), -1, -1)
488
+ a_tij = torch.cat([a_tij, res_type_i, res_type_j], dim=-1)
489
+ a_tij = self.a_proj(a_tij)
490
+
491
+ # Expand mask
492
+ pair_mask = pair_mask[:, None].expand(-1, T, -1, -1)
493
+ pair_mask = pair_mask.reshape(B * T, *pair_mask.shape[2:])
494
+
495
+ # Compute input projections
496
+ v = self.z_proj(self.z_norm(z[:, None])) + a_tij
497
+ v = v.view(B * T, *v.shape[2:])
498
+ v = v + self.pairformer(v, pair_mask, use_kernels=use_kernels)
499
+ v = self.v_norm(v)
500
+ v = v.view(B, T, *v.shape[1:])
501
+
502
+ # Aggregate templates
503
+ template_mask = template_mask[:, :, None, None, None]
504
+ num_templates = num_templates[:, None, None, None]
505
+ u = (v * template_mask).sum(dim=1) / num_templates.to(v)
506
+
507
+ # Compute output projection
508
+ u = self.u_proj(self.relu(u))
509
+ return u
510
+
511
+
512
+ class MSAModule(nn.Module):
513
+ """MSA module."""
514
+
515
+ def __init__(
516
+ self,
517
+ msa_s: int,
518
+ token_z: int,
519
+ token_s: int,
520
+ msa_blocks: int,
521
+ msa_dropout: float,
522
+ z_dropout: float,
523
+ pairwise_head_width: int = 32,
524
+ pairwise_num_heads: int = 4,
525
+ activation_checkpointing: bool = False,
526
+ use_paired_feature: bool = True,
527
+ subsample_msa: bool = False,
528
+ num_subsampled_msa: int = 1024,
529
+ **kwargs,
530
+ ) -> None:
531
+ """Initialize the MSA module.
532
+
533
+ Parameters
534
+ ----------
535
+ token_z : int
536
+ The token pairwise embedding size.
537
+
538
+ """
539
+ super().__init__()
540
+ self.msa_blocks = msa_blocks
541
+ self.msa_dropout = msa_dropout
542
+ self.z_dropout = z_dropout
543
+ self.use_paired_feature = use_paired_feature
544
+ self.activation_checkpointing = activation_checkpointing
545
+ self.subsample_msa = subsample_msa
546
+ self.num_subsampled_msa = num_subsampled_msa
547
+
548
+ self.s_proj = nn.Linear(token_s, msa_s, bias=False)
549
+ self.msa_proj = nn.Linear(
550
+ const.num_tokens + 2 + int(use_paired_feature),
551
+ msa_s,
552
+ bias=False,
553
+ )
554
+ self.layers = nn.ModuleList()
555
+ for i in range(msa_blocks):
556
+ self.layers.append(
557
+ MSALayer(
558
+ msa_s,
559
+ token_z,
560
+ msa_dropout,
561
+ z_dropout,
562
+ pairwise_head_width,
563
+ pairwise_num_heads,
564
+ )
565
+ )
566
+
567
+ def forward(
568
+ self,
569
+ z: Tensor,
570
+ emb: Tensor,
571
+ feats: dict[str, Tensor],
572
+ use_kernels: bool = False,
573
+ ) -> Tensor:
574
+ """Perform the forward pass.
575
+
576
+ Parameters
577
+ ----------
578
+ z : Tensor
579
+ The pairwise embeddings
580
+ emb : Tensor
581
+ The input embeddings
582
+ feats : dict[str, Tensor]
583
+ Input features
584
+ use_kernels: bool
585
+ Whether to use kernels for triangular updates
586
+
587
+ Returns
588
+ -------
589
+ Tensor
590
+ The output pairwise embeddings.
591
+
592
+ """
593
+ # Set chunk sizes
594
+ if not self.training:
595
+ if z.shape[1] > const.chunk_size_threshold:
596
+ chunk_heads_pwa = True
597
+ chunk_size_transition_z = 64
598
+ chunk_size_transition_msa = 32
599
+ chunk_size_outer_product = 4
600
+ chunk_size_tri_attn = 128
601
+ else:
602
+ chunk_heads_pwa = False
603
+ chunk_size_transition_z = None
604
+ chunk_size_transition_msa = None
605
+ chunk_size_outer_product = None
606
+ chunk_size_tri_attn = 512
607
+ else:
608
+ chunk_heads_pwa = False
609
+ chunk_size_transition_z = None
610
+ chunk_size_transition_msa = None
611
+ chunk_size_outer_product = None
612
+ chunk_size_tri_attn = None
613
+
614
+ # Load relevant features
615
+ msa = feats["msa"]
616
+ msa = torch.nn.functional.one_hot(msa, num_classes=const.num_tokens)
617
+ has_deletion = feats["has_deletion"].unsqueeze(-1)
618
+ deletion_value = feats["deletion_value"].unsqueeze(-1)
619
+ is_paired = feats["msa_paired"].unsqueeze(-1)
620
+ msa_mask = feats["msa_mask"]
621
+ token_mask = feats["token_pad_mask"].float()
622
+ token_mask = token_mask[:, :, None] * token_mask[:, None, :]
623
+
624
+ # Compute MSA embeddings
625
+ if self.use_paired_feature:
626
+ m = torch.cat([msa, has_deletion, deletion_value, is_paired], dim=-1)
627
+ else:
628
+ m = torch.cat([msa, has_deletion, deletion_value], dim=-1)
629
+
630
+ # Subsample the MSA
631
+ if self.subsample_msa:
632
+ msa_indices = torch.randperm(msa.shape[1])[: self.num_subsampled_msa]
633
+ m = m[:, msa_indices]
634
+ msa_mask = msa_mask[:, msa_indices]
635
+
636
+ # Compute input projections
637
+ m = self.msa_proj(m)
638
+ m = m + self.s_proj(emb).unsqueeze(1)
639
+
640
+ # Perform MSA blocks
641
+ for i in range(self.msa_blocks):
642
+ if self.activation_checkpointing and self.training:
643
+ z, m = torch.utils.checkpoint.checkpoint(
644
+ self.layers[i],
645
+ z,
646
+ m,
647
+ token_mask,
648
+ msa_mask,
649
+ chunk_heads_pwa,
650
+ chunk_size_transition_z,
651
+ chunk_size_transition_msa,
652
+ chunk_size_outer_product,
653
+ chunk_size_tri_attn,
654
+ use_kernels,
655
+ )
656
+ else:
657
+ z, m = self.layers[i](
658
+ z,
659
+ m,
660
+ token_mask,
661
+ msa_mask,
662
+ chunk_heads_pwa,
663
+ chunk_size_transition_z,
664
+ chunk_size_transition_msa,
665
+ chunk_size_outer_product,
666
+ chunk_size_tri_attn,
667
+ use_kernels,
668
+ )
669
+ return z
670
+
671
+
672
+ class MSALayer(nn.Module):
673
+ """MSA module."""
674
+
675
+ def __init__(
676
+ self,
677
+ msa_s: int,
678
+ token_z: int,
679
+ msa_dropout: float,
680
+ z_dropout: float,
681
+ pairwise_head_width: int = 32,
682
+ pairwise_num_heads: int = 4,
683
+ ) -> None:
684
+ """Initialize the MSA module.
685
+
686
+ Parameters
687
+ ----------
688
+ token_z : int
689
+ The token pairwise embedding size.
690
+
691
+ """
692
+ super().__init__()
693
+ self.msa_dropout = msa_dropout
694
+ self.msa_transition = Transition(dim=msa_s, hidden=msa_s * 4)
695
+ self.pair_weighted_averaging = PairWeightedAveraging(
696
+ c_m=msa_s,
697
+ c_z=token_z,
698
+ c_h=32,
699
+ num_heads=8,
700
+ )
701
+
702
+ self.pairformer_layer = PairformerNoSeqLayer(
703
+ token_z=token_z,
704
+ dropout=z_dropout,
705
+ pairwise_head_width=pairwise_head_width,
706
+ pairwise_num_heads=pairwise_num_heads,
707
+ )
708
+ self.outer_product_mean = OuterProductMean(
709
+ c_in=msa_s,
710
+ c_hidden=32,
711
+ c_out=token_z,
712
+ )
713
+
714
+ def forward(
715
+ self,
716
+ z: Tensor,
717
+ m: Tensor,
718
+ token_mask: Tensor,
719
+ msa_mask: Tensor,
720
+ chunk_heads_pwa: bool = False,
721
+ chunk_size_transition_z: int = None,
722
+ chunk_size_transition_msa: int = None,
723
+ chunk_size_outer_product: int = None,
724
+ chunk_size_tri_attn: int = None,
725
+ use_kernels: bool = False,
726
+ ) -> tuple[Tensor, Tensor]:
727
+ """Perform the forward pass.
728
+
729
+ Parameters
730
+ ----------
731
+ z : Tensor
732
+ The pairwise embeddings
733
+ emb : Tensor
734
+ The input embeddings
735
+ feats : dict[str, Tensor]
736
+ Input features
737
+
738
+ Returns
739
+ -------
740
+ Tensor
741
+ The output pairwise embeddings.
742
+
743
+ """
744
+ # Communication to MSA stack
745
+ msa_dropout = get_dropout_mask(self.msa_dropout, m, self.training)
746
+ m = m + msa_dropout * self.pair_weighted_averaging(
747
+ m, z, token_mask, chunk_heads_pwa
748
+ )
749
+ m = m + self.msa_transition(m, chunk_size_transition_msa)
750
+
751
+ z = z + self.outer_product_mean(m, msa_mask, chunk_size_outer_product)
752
+
753
+ # Compute pairwise stack
754
+ z = self.pairformer_layer(
755
+ z, token_mask, chunk_size_tri_attn, use_kernels=use_kernels
756
+ )
757
+
758
+ return z, m
759
+
760
+
761
+ class BFactorModule(nn.Module):
762
+ """BFactor Module."""
763
+
764
+ def __init__(self, token_s: int, num_bins: int) -> None:
765
+ """Initialize the bfactor module.
766
+
767
+ Parameters
768
+ ----------
769
+ token_s : int
770
+ The token embedding size.
771
+
772
+ """
773
+ super().__init__()
774
+ self.bfactor = nn.Linear(token_s, num_bins)
775
+ self.num_bins = num_bins
776
+
777
+ def forward(self, s: Tensor) -> Tensor:
778
+ """Perform the forward pass.
779
+
780
+ Parameters
781
+ ----------
782
+ s : Tensor
783
+ The sequence embeddings
784
+
785
+ Returns
786
+ -------
787
+ Tensor
788
+ The predicted bfactor histogram.
789
+
790
+ """
791
+ return self.bfactor(s)
792
+
793
+
794
+ class DistogramModule(nn.Module):
795
+ """Distogram Module."""
796
+
797
+ def __init__(self, token_z: int, num_bins: int, num_distograms: int = 1) -> None:
798
+ """Initialize the distogram module.
799
+
800
+ Parameters
801
+ ----------
802
+ token_z : int
803
+ The token pairwise embedding size.
804
+
805
+ """
806
+ super().__init__()
807
+ self.distogram = nn.Linear(token_z, num_distograms * num_bins)
808
+ self.num_distograms = num_distograms
809
+ self.num_bins = num_bins
810
+
811
+ def forward(self, z: Tensor) -> Tensor:
812
+ """Perform the forward pass.
813
+
814
+ Parameters
815
+ ----------
816
+ z : Tensor
817
+ The pairwise embeddings
818
+
819
+ Returns
820
+ -------
821
+ Tensor
822
+ The predicted distogram.
823
+
824
+ """
825
+ z = z + z.transpose(1, 2)
826
+ return self.distogram(z).reshape(
827
+ z.shape[0], z.shape[1], z.shape[2], self.num_distograms, self.num_bins
828
+ )
vb_modules_utils.py CHANGED
@@ -1,303 +1,303 @@
1
- # started from code from https://github.com/lucidrains/alphafold3-pytorch, MIT License, Copyright (c) 2024 Phil Wang
2
-
3
- from functools import partial
4
- from typing import Optional
5
-
6
- import torch
7
- import torch.nn.functional as F
8
- from torch.nn import (
9
- Linear,
10
- Module,
11
- )
12
- from torch.types import Device
13
-
14
- LinearNoBias = partial(Linear, bias=False)
15
-
16
-
17
- def exists(v):
18
- return v is not None
19
-
20
-
21
- def default(v, d):
22
- return v if exists(v) else d
23
-
24
-
25
- def log(t, eps=1e-20):
26
- return torch.log(t.clamp(min=eps))
27
-
28
-
29
- class SwiGLU(Module):
30
- def forward(
31
- self,
32
- x, #: Float['... d']
33
- ): # -> Float[' ... (d//2)']:
34
- x, gates = x.chunk(2, dim=-1)
35
- return F.silu(gates) * x
36
-
37
-
38
- def center(atom_coords, atom_mask):
39
- atom_mean = torch.sum(
40
- atom_coords * atom_mask[:, :, None], dim=1, keepdim=True
41
- ) / torch.sum(atom_mask[:, :, None], dim=1, keepdim=True)
42
- atom_coords = atom_coords - atom_mean
43
- return atom_coords
44
-
45
-
46
- def compute_random_augmentation(
47
- multiplicity, s_trans=1.0, device=None, dtype=torch.float32
48
- ):
49
- R = random_rotations(multiplicity, dtype=dtype, device=device)
50
- random_trans = (
51
- torch.randn((multiplicity, 1, 3), dtype=dtype, device=device) * s_trans
52
- )
53
- return R, random_trans
54
-
55
-
56
- def randomly_rotate(coords, return_second_coords=False, second_coords=None):
57
- R = random_rotations(len(coords), coords.dtype, coords.device)
58
-
59
- if return_second_coords:
60
- return torch.einsum("bmd,bds->bms", coords, R), torch.einsum(
61
- "bmd,bds->bms", second_coords, R
62
- ) if second_coords is not None else None
63
-
64
- return torch.einsum("bmd,bds->bms", coords, R)
65
-
66
-
67
- def center_random_augmentation(
68
- atom_coords,
69
- atom_mask,
70
- s_trans=1.0,
71
- augmentation=True,
72
- centering=True,
73
- return_second_coords=False,
74
- second_coords=None,
75
- ):
76
- """Algorithm 19"""
77
- if centering:
78
- atom_mean = torch.sum(
79
- atom_coords * atom_mask[:, :, None], dim=1, keepdim=True
80
- ) / torch.sum(atom_mask[:, :, None], dim=1, keepdim=True)
81
- atom_coords = atom_coords - atom_mean
82
-
83
- if second_coords is not None:
84
- # apply same transformation also to this input
85
- second_coords = second_coords - atom_mean
86
-
87
- if augmentation:
88
- atom_coords, second_coords = randomly_rotate(
89
- atom_coords, return_second_coords=True, second_coords=second_coords
90
- )
91
- random_trans = torch.randn_like(atom_coords[:, 0:1, :]) * s_trans
92
- atom_coords = atom_coords + random_trans
93
-
94
- if second_coords is not None:
95
- second_coords = second_coords + random_trans
96
-
97
- if return_second_coords:
98
- return atom_coords, second_coords
99
-
100
- return atom_coords
101
-
102
-
103
- class ExponentialMovingAverage:
104
- """from https://github.com/yang-song/score_sde_pytorch/blob/main/models/ema.py, Apache-2.0 license
105
- Maintains (exponential) moving average of a set of parameters."""
106
-
107
- def __init__(self, parameters, decay, use_num_updates=True):
108
- """
109
- Args:
110
- parameters: Iterable of `torch.nn.Parameter`; usually the result of
111
- `model.parameters()`.
112
- decay: The exponential decay.
113
- use_num_updates: Whether to use number of updates when computing
114
- averages.
115
- """
116
- if decay < 0.0 or decay > 1.0:
117
- raise ValueError("Decay must be between 0 and 1")
118
- self.decay = decay
119
- self.num_updates = 0 if use_num_updates else None
120
- self.shadow_params = [p.clone().detach() for p in parameters if p.requires_grad]
121
- self.collected_params = []
122
-
123
- def update(self, parameters):
124
- """
125
- Update currently maintained parameters.
126
- Call this every time the parameters are updated, such as the result of
127
- the `optimizer.step()` call.
128
- Args:
129
- parameters: Iterable of `torch.nn.Parameter`; usually the same set of
130
- parameters used to initialize this object.
131
- """
132
- decay = self.decay
133
- if self.num_updates is not None:
134
- self.num_updates += 1
135
- decay = min(decay, (1 + self.num_updates) / (10 + self.num_updates))
136
- one_minus_decay = 1.0 - decay
137
- with torch.no_grad():
138
- parameters = [p for p in parameters if p.requires_grad]
139
- for s_param, param in zip(self.shadow_params, parameters):
140
- s_param.sub_(one_minus_decay * (s_param - param))
141
-
142
- def compatible(self, parameters):
143
- if len(self.shadow_params) != len(parameters):
144
- print(
145
- f"Model has {len(self.shadow_params)} parameter tensors, the incoming ema {len(parameters)}"
146
- )
147
- return False
148
-
149
- for s_param, param in zip(self.shadow_params, parameters):
150
- if param.data.shape != s_param.data.shape:
151
- print(
152
- f"Model has parameter tensor of shape {s_param.data.shape} , the incoming ema {param.data.shape}"
153
- )
154
- return False
155
- return True
156
-
157
- def copy_to(self, parameters):
158
- """
159
- Copy current parameters into given collection of parameters.
160
- Args:
161
- parameters: Iterable of `torch.nn.Parameter`; the parameters to be
162
- updated with the stored moving averages.
163
- """
164
- parameters = [p for p in parameters if p.requires_grad]
165
- for s_param, param in zip(self.shadow_params, parameters):
166
- if param.requires_grad:
167
- param.data.copy_(s_param.data)
168
-
169
- def store(self, parameters):
170
- """
171
- Save the current parameters for restoring later.
172
- Args:
173
- parameters: Iterable of `torch.nn.Parameter`; the parameters to be
174
- temporarily stored.
175
- """
176
- self.collected_params = [param.clone() for param in parameters]
177
-
178
- def restore(self, parameters):
179
- """
180
- Restore the parameters stored with the `store` method.
181
- Useful to validate the model with EMA parameters without affecting the
182
- original optimization process. Store the parameters before the
183
- `copy_to` method. After validation (or model saving), use this to
184
- restore the former parameters.
185
- Args:
186
- parameters: Iterable of `torch.nn.Parameter`; the parameters to be
187
- updated with the stored parameters.
188
- """
189
- for c_param, param in zip(self.collected_params, parameters):
190
- param.data.copy_(c_param.data)
191
-
192
- def state_dict(self):
193
- return dict(
194
- decay=self.decay,
195
- num_updates=self.num_updates,
196
- shadow_params=self.shadow_params,
197
- )
198
-
199
- def load_state_dict(self, state_dict, device):
200
- self.decay = state_dict["decay"]
201
- self.num_updates = state_dict["num_updates"]
202
- self.shadow_params = [
203
- tensor.to(device) for tensor in state_dict["shadow_params"]
204
- ]
205
-
206
- def to(self, device):
207
- self.shadow_params = [tensor.to(device) for tensor in self.shadow_params]
208
-
209
-
210
- # the following is copied from Torch3D, BSD License, Copyright (c) Meta Platforms, Inc. and affiliates.
211
-
212
-
213
- def _copysign(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
214
- """
215
- Return a tensor where each element has the absolute value taken from the,
216
- corresponding element of a, with sign taken from the corresponding
217
- element of b. This is like the standard copysign floating-point operation,
218
- but is not careful about negative 0 and NaN.
219
-
220
- Args:
221
- a: source tensor.
222
- b: tensor whose signs will be used, of the same shape as a.
223
-
224
- Returns:
225
- Tensor of the same shape as a with the signs of b.
226
- """
227
- signs_differ = (a < 0) != (b < 0)
228
- return torch.where(signs_differ, -a, a)
229
-
230
-
231
- def quaternion_to_matrix(quaternions: torch.Tensor) -> torch.Tensor:
232
- """
233
- Convert rotations given as quaternions to rotation matrices.
234
-
235
- Args:
236
- quaternions: quaternions with real part first,
237
- as tensor of shape (..., 4).
238
-
239
- Returns:
240
- Rotation matrices as tensor of shape (..., 3, 3).
241
- """
242
- r, i, j, k = torch.unbind(quaternions, -1)
243
- # pyre-fixme[58]: `/` is not supported for operand types `float` and `Tensor`.
244
- two_s = 2.0 / (quaternions * quaternions).sum(-1)
245
-
246
- o = torch.stack(
247
- (
248
- 1 - two_s * (j * j + k * k),
249
- two_s * (i * j - k * r),
250
- two_s * (i * k + j * r),
251
- two_s * (i * j + k * r),
252
- 1 - two_s * (i * i + k * k),
253
- two_s * (j * k - i * r),
254
- two_s * (i * k - j * r),
255
- two_s * (j * k + i * r),
256
- 1 - two_s * (i * i + j * j),
257
- ),
258
- -1,
259
- )
260
- return o.reshape(quaternions.shape[:-1] + (3, 3))
261
-
262
-
263
- def random_quaternions(
264
- n: int, dtype: Optional[torch.dtype] = None, device: Optional[Device] = None
265
- ) -> torch.Tensor:
266
- """
267
- Generate random quaternions representing rotations,
268
- i.e. versors with nonnegative real part.
269
-
270
- Args:
271
- n: Number of quaternions in a batch to return.
272
- dtype: Type to return.
273
- device: Desired device of returned tensor. Default:
274
- uses the current device for the default tensor type.
275
-
276
- Returns:
277
- Quaternions as tensor of shape (N, 4).
278
- """
279
- if isinstance(device, str):
280
- device = torch.device(device)
281
- o = torch.randn((n, 4), dtype=dtype, device=device)
282
- s = (o * o).sum(1)
283
- o = o / _copysign(torch.sqrt(s), o[:, 0])[:, None]
284
- return o
285
-
286
-
287
- def random_rotations(
288
- n: int, dtype: Optional[torch.dtype] = None, device: Optional[Device] = None
289
- ) -> torch.Tensor:
290
- """
291
- Generate random rotations as 3x3 rotation matrices.
292
-
293
- Args:
294
- n: Number of rotation matrices in a batch to return.
295
- dtype: Type to return.
296
- device: Device of returned tensor. Default: if None,
297
- uses the current device for the default tensor type.
298
-
299
- Returns:
300
- Rotation matrices as tensor of shape (n, 3, 3).
301
- """
302
- quaternions = random_quaternions(n, dtype=dtype, device=device)
303
- return quaternion_to_matrix(quaternions)
 
1
+ # started from code from https://github.com/lucidrains/alphafold3-pytorch, MIT License, Copyright (c) 2024 Phil Wang
2
+
3
+ from functools import partial
4
+ from typing import Optional
5
+
6
+ import torch
7
+ import torch.nn.functional as F
8
+ from torch.nn import (
9
+ Linear,
10
+ Module,
11
+ )
12
+ from torch.types import Device
13
+
14
+ LinearNoBias = partial(Linear, bias=False)
15
+
16
+
17
+ def exists(v):
18
+ return v is not None
19
+
20
+
21
+ def default(v, d):
22
+ return v if exists(v) else d
23
+
24
+
25
+ def log(t, eps=1e-20):
26
+ return torch.log(t.clamp(min=eps))
27
+
28
+
29
+ class SwiGLU(Module):
30
+ def forward(
31
+ self,
32
+ x, #: Float['... d']
33
+ ): # -> Float[' ... (d//2)']:
34
+ x, gates = x.chunk(2, dim=-1)
35
+ return F.silu(gates) * x
36
+
37
+
38
+ def center(atom_coords, atom_mask):
39
+ atom_mean = torch.sum(
40
+ atom_coords * atom_mask[:, :, None], dim=1, keepdim=True
41
+ ) / torch.sum(atom_mask[:, :, None], dim=1, keepdim=True)
42
+ atom_coords = atom_coords - atom_mean
43
+ return atom_coords
44
+
45
+
46
+ def compute_random_augmentation(
47
+ multiplicity, s_trans=1.0, device=None, dtype=torch.float32
48
+ ):
49
+ R = random_rotations(multiplicity, dtype=dtype, device=device)
50
+ random_trans = (
51
+ torch.randn((multiplicity, 1, 3), dtype=dtype, device=device) * s_trans
52
+ )
53
+ return R, random_trans
54
+
55
+
56
+ def randomly_rotate(coords, return_second_coords=False, second_coords=None):
57
+ R = random_rotations(len(coords), coords.dtype, coords.device)
58
+
59
+ if return_second_coords:
60
+ return torch.einsum("bmd,bds->bms", coords, R), torch.einsum(
61
+ "bmd,bds->bms", second_coords, R
62
+ ) if second_coords is not None else None
63
+
64
+ return torch.einsum("bmd,bds->bms", coords, R)
65
+
66
+
67
+ def center_random_augmentation(
68
+ atom_coords,
69
+ atom_mask,
70
+ s_trans=1.0,
71
+ augmentation=True,
72
+ centering=True,
73
+ return_second_coords=False,
74
+ second_coords=None,
75
+ ):
76
+ """Algorithm 19"""
77
+ if centering:
78
+ atom_mean = torch.sum(
79
+ atom_coords * atom_mask[:, :, None], dim=1, keepdim=True
80
+ ) / torch.sum(atom_mask[:, :, None], dim=1, keepdim=True)
81
+ atom_coords = atom_coords - atom_mean
82
+
83
+ if second_coords is not None:
84
+ # apply same transformation also to this input
85
+ second_coords = second_coords - atom_mean
86
+
87
+ if augmentation:
88
+ atom_coords, second_coords = randomly_rotate(
89
+ atom_coords, return_second_coords=True, second_coords=second_coords
90
+ )
91
+ random_trans = torch.randn_like(atom_coords[:, 0:1, :]) * s_trans
92
+ atom_coords = atom_coords + random_trans
93
+
94
+ if second_coords is not None:
95
+ second_coords = second_coords + random_trans
96
+
97
+ if return_second_coords:
98
+ return atom_coords, second_coords
99
+
100
+ return atom_coords
101
+
102
+
103
+ class ExponentialMovingAverage:
104
+ """from https://github.com/yang-song/score_sde_pytorch/blob/main/models/ema.py, Apache-2.0 license
105
+ Maintains (exponential) moving average of a set of parameters."""
106
+
107
+ def __init__(self, parameters, decay, use_num_updates=True):
108
+ """
109
+ Args:
110
+ parameters: Iterable of `torch.nn.Parameter`; usually the result of
111
+ `model.parameters()`.
112
+ decay: The exponential decay.
113
+ use_num_updates: Whether to use number of updates when computing
114
+ averages.
115
+ """
116
+ if decay < 0.0 or decay > 1.0:
117
+ raise ValueError("Decay must be between 0 and 1")
118
+ self.decay = decay
119
+ self.num_updates = 0 if use_num_updates else None
120
+ self.shadow_params = [p.clone().detach() for p in parameters if p.requires_grad]
121
+ self.collected_params = []
122
+
123
+ def update(self, parameters):
124
+ """
125
+ Update currently maintained parameters.
126
+ Call this every time the parameters are updated, such as the result of
127
+ the `optimizer.step()` call.
128
+ Args:
129
+ parameters: Iterable of `torch.nn.Parameter`; usually the same set of
130
+ parameters used to initialize this object.
131
+ """
132
+ decay = self.decay
133
+ if self.num_updates is not None:
134
+ self.num_updates += 1
135
+ decay = min(decay, (1 + self.num_updates) / (10 + self.num_updates))
136
+ one_minus_decay = 1.0 - decay
137
+ with torch.no_grad():
138
+ parameters = [p for p in parameters if p.requires_grad]
139
+ for s_param, param in zip(self.shadow_params, parameters):
140
+ s_param.sub_(one_minus_decay * (s_param - param))
141
+
142
+ def compatible(self, parameters):
143
+ if len(self.shadow_params) != len(parameters):
144
+ print(
145
+ f"Model has {len(self.shadow_params)} parameter tensors, the incoming ema {len(parameters)}"
146
+ )
147
+ return False
148
+
149
+ for s_param, param in zip(self.shadow_params, parameters):
150
+ if param.data.shape != s_param.data.shape:
151
+ print(
152
+ f"Model has parameter tensor of shape {s_param.data.shape} , the incoming ema {param.data.shape}"
153
+ )
154
+ return False
155
+ return True
156
+
157
+ def copy_to(self, parameters):
158
+ """
159
+ Copy current parameters into given collection of parameters.
160
+ Args:
161
+ parameters: Iterable of `torch.nn.Parameter`; the parameters to be
162
+ updated with the stored moving averages.
163
+ """
164
+ parameters = [p for p in parameters if p.requires_grad]
165
+ for s_param, param in zip(self.shadow_params, parameters):
166
+ if param.requires_grad:
167
+ param.data.copy_(s_param.data)
168
+
169
+ def store(self, parameters):
170
+ """
171
+ Save the current parameters for restoring later.
172
+ Args:
173
+ parameters: Iterable of `torch.nn.Parameter`; the parameters to be
174
+ temporarily stored.
175
+ """
176
+ self.collected_params = [param.clone() for param in parameters]
177
+
178
+ def restore(self, parameters):
179
+ """
180
+ Restore the parameters stored with the `store` method.
181
+ Useful to validate the model with EMA parameters without affecting the
182
+ original optimization process. Store the parameters before the
183
+ `copy_to` method. After validation (or model saving), use this to
184
+ restore the former parameters.
185
+ Args:
186
+ parameters: Iterable of `torch.nn.Parameter`; the parameters to be
187
+ updated with the stored parameters.
188
+ """
189
+ for c_param, param in zip(self.collected_params, parameters):
190
+ param.data.copy_(c_param.data)
191
+
192
+ def state_dict(self):
193
+ return dict(
194
+ decay=self.decay,
195
+ num_updates=self.num_updates,
196
+ shadow_params=self.shadow_params,
197
+ )
198
+
199
+ def load_state_dict(self, state_dict, device):
200
+ self.decay = state_dict["decay"]
201
+ self.num_updates = state_dict["num_updates"]
202
+ self.shadow_params = [
203
+ tensor.to(device) for tensor in state_dict["shadow_params"]
204
+ ]
205
+
206
+ def to(self, device):
207
+ self.shadow_params = [tensor.to(device) for tensor in self.shadow_params]
208
+
209
+
210
+ # the following is copied from Torch3D, BSD License, Copyright (c) Meta Platforms, Inc. and affiliates.
211
+
212
+
213
+ def _copysign(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
214
+ """
215
+ Return a tensor where each element has the absolute value taken from the,
216
+ corresponding element of a, with sign taken from the corresponding
217
+ element of b. This is like the standard copysign floating-point operation,
218
+ but is not careful about negative 0 and NaN.
219
+
220
+ Args:
221
+ a: source tensor.
222
+ b: tensor whose signs will be used, of the same shape as a.
223
+
224
+ Returns:
225
+ Tensor of the same shape as a with the signs of b.
226
+ """
227
+ signs_differ = (a < 0) != (b < 0)
228
+ return torch.where(signs_differ, -a, a)
229
+
230
+
231
+ def quaternion_to_matrix(quaternions: torch.Tensor) -> torch.Tensor:
232
+ """
233
+ Convert rotations given as quaternions to rotation matrices.
234
+
235
+ Args:
236
+ quaternions: quaternions with real part first,
237
+ as tensor of shape (..., 4).
238
+
239
+ Returns:
240
+ Rotation matrices as tensor of shape (..., 3, 3).
241
+ """
242
+ r, i, j, k = torch.unbind(quaternions, -1)
243
+ # pyre-fixme[58]: `/` is not supported for operand types `float` and `Tensor`.
244
+ two_s = 2.0 / (quaternions * quaternions).sum(-1)
245
+
246
+ o = torch.stack(
247
+ (
248
+ 1 - two_s * (j * j + k * k),
249
+ two_s * (i * j - k * r),
250
+ two_s * (i * k + j * r),
251
+ two_s * (i * j + k * r),
252
+ 1 - two_s * (i * i + k * k),
253
+ two_s * (j * k - i * r),
254
+ two_s * (i * k - j * r),
255
+ two_s * (j * k + i * r),
256
+ 1 - two_s * (i * i + j * j),
257
+ ),
258
+ -1,
259
+ )
260
+ return o.reshape(quaternions.shape[:-1] + (3, 3))
261
+
262
+
263
+ def random_quaternions(
264
+ n: int, dtype: Optional[torch.dtype] = None, device: Optional[Device] = None
265
+ ) -> torch.Tensor:
266
+ """
267
+ Generate random quaternions representing rotations,
268
+ i.e. versors with nonnegative real part.
269
+
270
+ Args:
271
+ n: Number of quaternions in a batch to return.
272
+ dtype: Type to return.
273
+ device: Desired device of returned tensor. Default:
274
+ uses the current device for the default tensor type.
275
+
276
+ Returns:
277
+ Quaternions as tensor of shape (N, 4).
278
+ """
279
+ if isinstance(device, str):
280
+ device = torch.device(device)
281
+ o = torch.randn((n, 4), dtype=dtype, device=device)
282
+ s = (o * o).sum(1)
283
+ o = o / _copysign(torch.sqrt(s), o[:, 0])[:, None]
284
+ return o
285
+
286
+
287
+ def random_rotations(
288
+ n: int, dtype: Optional[torch.dtype] = None, device: Optional[Device] = None
289
+ ) -> torch.Tensor:
290
+ """
291
+ Generate random rotations as 3x3 rotation matrices.
292
+
293
+ Args:
294
+ n: Number of rotation matrices in a batch to return.
295
+ dtype: Type to return.
296
+ device: Device of returned tensor. Default: if None,
297
+ uses the current device for the default tensor type.
298
+
299
+ Returns:
300
+ Rotation matrices as tensor of shape (n, 3, 3).
301
+ """
302
+ quaternions = random_quaternions(n, dtype=dtype, device=device)
303
+ return quaternion_to_matrix(quaternions)
vb_potentials_potentials.py CHANGED
@@ -1,787 +1,787 @@
1
- from abc import ABC, abstractmethod
2
- from typing import Optional, Dict, Any, Set, List, Union
3
-
4
- import torch
5
- import numpy as np
6
- from . import vb_const as const
7
- from .vb_potentials_schedules import (
8
- ParameterSchedule,
9
- ExponentialInterpolation,
10
- PiecewiseStepFunction,
11
- )
12
- from .vb_loss_diffusionv2 import weighted_rigid_align
13
-
14
-
15
- class Potential(ABC):
16
- def __init__(
17
- self,
18
- parameters: Optional[
19
- Dict[str, Union[ParameterSchedule, float, int, bool]]
20
- ] = None,
21
- ):
22
- self.parameters = parameters
23
-
24
- def compute(self, coords, feats, parameters):
25
- index, args, com_args, ref_args, operator_args = self.compute_args(
26
- feats, parameters
27
- )
28
-
29
- if index.shape[1] == 0:
30
- return torch.zeros(coords.shape[:-2], device=coords.device)
31
-
32
- if com_args is not None:
33
- com_index, atom_pad_mask = com_args
34
- unpad_com_index = com_index[atom_pad_mask]
35
- unpad_coords = coords[..., atom_pad_mask, :]
36
- coords = torch.zeros(
37
- (*unpad_coords.shape[:-2], unpad_com_index.max() + 1, 3),
38
- device=coords.device,
39
- ).scatter_reduce(
40
- -2,
41
- unpad_com_index.unsqueeze(-1).expand_as(unpad_coords),
42
- unpad_coords,
43
- "mean",
44
- )
45
- else:
46
- com_index, atom_pad_mask = None, None
47
-
48
- if ref_args is not None:
49
- ref_coords, ref_mask, ref_atom_index, ref_token_index = ref_args
50
- coords = coords[..., ref_atom_index, :]
51
- else:
52
- ref_coords, ref_mask, ref_atom_index, ref_token_index = (
53
- None,
54
- None,
55
- None,
56
- None,
57
- )
58
-
59
- if operator_args is not None:
60
- negation_mask, union_index = operator_args
61
- else:
62
- negation_mask, union_index = None, None
63
-
64
- value = self.compute_variable(
65
- coords,
66
- index,
67
- ref_coords=ref_coords,
68
- ref_mask=ref_mask,
69
- compute_gradient=False,
70
- )
71
- energy = self.compute_function(
72
- value, *args, negation_mask=negation_mask, compute_derivative=False
73
- )
74
-
75
- if union_index is not None:
76
- neg_exp_energy = torch.exp(-1 * parameters["union_lambda"] * energy)
77
- Z = torch.zeros(
78
- (*energy.shape[:-1], union_index.max() + 1), device=union_index.device
79
- ).scatter_reduce(
80
- -1,
81
- union_index.expand_as(neg_exp_energy),
82
- neg_exp_energy,
83
- "sum",
84
- )
85
- softmax_energy = neg_exp_energy / Z[..., union_index]
86
- softmax_energy[Z[..., union_index] == 0] = 0
87
- return (energy * softmax_energy).sum(dim=-1)
88
-
89
- return energy.sum(dim=tuple(range(1, energy.dim())))
90
-
91
- def compute_gradient(self, coords, feats, parameters):
92
- index, args, com_args, ref_args, operator_args = self.compute_args(
93
- feats, parameters
94
- )
95
- if index.shape[1] == 0:
96
- return torch.zeros_like(coords)
97
-
98
- if com_args is not None:
99
- com_index, atom_pad_mask = com_args
100
- unpad_coords = coords[..., atom_pad_mask, :]
101
- unpad_com_index = com_index[atom_pad_mask]
102
- coords = torch.zeros(
103
- (*unpad_coords.shape[:-2], unpad_com_index.max() + 1, 3),
104
- device=coords.device,
105
- ).scatter_reduce(
106
- -2,
107
- unpad_com_index.unsqueeze(-1).expand_as(unpad_coords),
108
- unpad_coords,
109
- "mean",
110
- )
111
- com_counts = torch.bincount(com_index[atom_pad_mask])
112
- else:
113
- com_index, atom_pad_mask = None, None
114
-
115
- if ref_args is not None:
116
- ref_coords, ref_mask, ref_atom_index, ref_token_index = ref_args
117
- coords = coords[..., ref_atom_index, :]
118
- else:
119
- ref_coords, ref_mask, ref_atom_index, ref_token_index = (
120
- None,
121
- None,
122
- None,
123
- None,
124
- )
125
-
126
- if operator_args is not None:
127
- negation_mask, union_index = operator_args
128
- else:
129
- negation_mask, union_index = None, None
130
-
131
- value, grad_value = self.compute_variable(
132
- coords,
133
- index,
134
- ref_coords=ref_coords,
135
- ref_mask=ref_mask,
136
- compute_gradient=True,
137
- )
138
- energy, dEnergy = self.compute_function(
139
- value,
140
- *args, negation_mask=negation_mask, compute_derivative=True
141
- )
142
- if union_index is not None:
143
- neg_exp_energy = torch.exp(-1 * parameters["union_lambda"] * energy)
144
- Z = torch.zeros(
145
- (*energy.shape[:-1], union_index.max() + 1), device=union_index.device
146
- ).scatter_reduce(
147
- -1,
148
- union_index.expand_as(energy),
149
- neg_exp_energy,
150
- "sum",
151
- )
152
- softmax_energy = neg_exp_energy / Z[..., union_index]
153
- softmax_energy[Z[..., union_index] == 0] = 0
154
- f = torch.zeros(
155
- (*energy.shape[:-1], union_index.max() + 1), device=union_index.device
156
- ).scatter_reduce(
157
- -1,
158
- union_index.expand_as(energy),
159
- energy * softmax_energy,
160
- "sum",
161
- )
162
- dSoftmax = (
163
- dEnergy
164
- * softmax_energy
165
- * (1 + parameters["union_lambda"] * (energy - f[..., union_index]))
166
- )
167
- prod = dSoftmax.tile(grad_value.shape[-3]).unsqueeze(
168
- -1
169
- ) * grad_value.flatten(start_dim=-3, end_dim=-2)
170
- if prod.dim() > 3:
171
- prod = prod.sum(dim=list(range(1, prod.dim() - 2)))
172
- grad_atom = torch.zeros_like(coords).scatter_reduce(
173
- -2,
174
- index.flatten(start_dim=0, end_dim=1)
175
- .unsqueeze(-1)
176
- .expand((*coords.shape[:-2], -1, 3)),
177
- prod,
178
- "sum",
179
- )
180
- else:
181
- prod = dEnergy.tile(grad_value.shape[-3]).unsqueeze(
182
- -1
183
- ) * grad_value.flatten(start_dim=-3, end_dim=-2)
184
- if prod.dim() > 3:
185
- prod = prod.sum(dim=list(range(1, prod.dim() - 2)))
186
- grad_atom = torch.zeros_like(coords).scatter_reduce(
187
- -2,
188
- index.flatten(start_dim=0, end_dim=1)
189
- .unsqueeze(-1)
190
- .expand((*coords.shape[:-2], -1, 3)), # 9 x 516 x 3
191
- prod,
192
- "sum",
193
- )
194
-
195
- if com_index is not None:
196
- grad_atom = grad_atom[..., com_index, :]
197
- elif ref_token_index is not None:
198
- grad_atom = grad_atom[..., ref_token_index, :]
199
-
200
- return grad_atom
201
-
202
- def compute_parameters(self, t):
203
- if self.parameters is None:
204
- return None
205
- parameters = {
206
- name: parameter
207
- if not isinstance(parameter, ParameterSchedule)
208
- else parameter.compute(t)
209
- for name, parameter in self.parameters.items()
210
- }
211
- return parameters
212
-
213
- @abstractmethod
214
- def compute_function(
215
- self, value, *args, negation_mask=None, compute_derivative=False
216
- ):
217
- raise NotImplementedError
218
-
219
- @abstractmethod
220
- def compute_variable(self, coords, index, compute_gradient=False):
221
- raise NotImplementedError
222
-
223
- @abstractmethod
224
- def compute_args(self, t, feats, **parameters):
225
- raise NotImplementedError
226
-
227
- def get_reference_coords(self, feats, parameters):
228
- return None, None
229
-
230
-
231
- class FlatBottomPotential(Potential):
232
- def compute_function(
233
- self,
234
- value,
235
- k,
236
- lower_bounds,
237
- upper_bounds,
238
- negation_mask=None,
239
- compute_derivative=False,
240
- ):
241
- if lower_bounds is None:
242
- lower_bounds = torch.full_like(value, float("-inf"))
243
- if upper_bounds is None:
244
- upper_bounds = torch.full_like(value, float("inf"))
245
- lower_bounds = lower_bounds.expand_as(value).clone()
246
- upper_bounds = upper_bounds.expand_as(value).clone()
247
-
248
- if negation_mask is not None:
249
- unbounded_below_mask = torch.isneginf(lower_bounds)
250
- unbounded_above_mask = torch.isposinf(upper_bounds)
251
- unbounded_mask = unbounded_below_mask + unbounded_above_mask
252
- assert torch.all(unbounded_mask + negation_mask)
253
- lower_bounds[~unbounded_above_mask * ~negation_mask] = upper_bounds[
254
- ~unbounded_above_mask * ~negation_mask
255
- ]
256
- upper_bounds[~unbounded_above_mask * ~negation_mask] = float("inf")
257
- upper_bounds[~unbounded_below_mask * ~negation_mask] = lower_bounds[
258
- ~unbounded_below_mask * ~negation_mask
259
- ]
260
- lower_bounds[~unbounded_below_mask * ~negation_mask] = float("-inf")
261
-
262
- neg_overflow_mask = value < lower_bounds
263
- pos_overflow_mask = value > upper_bounds
264
-
265
- energy = torch.zeros_like(value)
266
- energy[neg_overflow_mask] = (k * (lower_bounds - value))[neg_overflow_mask]
267
- energy[pos_overflow_mask] = (k * (value - upper_bounds))[pos_overflow_mask]
268
- if not compute_derivative:
269
- return energy
270
-
271
- dEnergy = torch.zeros_like(value)
272
- dEnergy[neg_overflow_mask] = (
273
- -1 * k.expand_as(neg_overflow_mask)[neg_overflow_mask]
274
- )
275
- dEnergy[pos_overflow_mask] = (
276
- 1 * k.expand_as(pos_overflow_mask)[pos_overflow_mask]
277
- )
278
-
279
- return energy, dEnergy
280
-
281
-
282
- class ReferencePotential(Potential):
283
- def compute_variable(
284
- self, coords, index, ref_coords, ref_mask, compute_gradient=False
285
- ):
286
- aligned_ref_coords = weighted_rigid_align(
287
- ref_coords.float(),
288
- coords[:, index].float(),
289
- ref_mask,
290
- ref_mask,
291
- )
292
-
293
- r = coords[:, index] - aligned_ref_coords
294
- r_norm = torch.linalg.norm(r, dim=-1)
295
-
296
- if not compute_gradient:
297
- return r_norm
298
-
299
- r_hat = r / r_norm.unsqueeze(-1)
300
- grad = (r_hat * ref_mask.unsqueeze(-1)).unsqueeze(1)
301
- return r_norm, grad
302
-
303
-
304
- class DistancePotential(Potential):
305
- def compute_variable(
306
- self, coords, index, ref_coords=None, ref_mask=None, compute_gradient=False
307
- ):
308
- r_ij = coords.index_select(-2, index[0]) - coords.index_select(-2, index[1])
309
- r_ij_norm = torch.linalg.norm(r_ij, dim=-1)
310
- r_hat_ij = r_ij / r_ij_norm.unsqueeze(-1)
311
-
312
- if not compute_gradient:
313
- return r_ij_norm
314
-
315
- grad_i = r_hat_ij
316
- grad_j = -1 * r_hat_ij
317
- grad = torch.stack((grad_i, grad_j), dim=1)
318
- return r_ij_norm, grad
319
-
320
-
321
- class DihedralPotential(Potential):
322
- def compute_variable(
323
- self, coords, index, ref_coords=None, ref_mask=None, compute_gradient=False
324
- ):
325
- r_ij = coords.index_select(-2, index[0]) - coords.index_select(-2, index[1])
326
- r_kj = coords.index_select(-2, index[2]) - coords.index_select(-2, index[1])
327
- r_kl = coords.index_select(-2, index[2]) - coords.index_select(-2, index[3])
328
-
329
- n_ijk = torch.cross(r_ij, r_kj, dim=-1)
330
- n_jkl = torch.cross(r_kj, r_kl, dim=-1)
331
-
332
- r_kj_norm = torch.linalg.norm(r_kj, dim=-1)
333
- n_ijk_norm = torch.linalg.norm(n_ijk, dim=-1)
334
- n_jkl_norm = torch.linalg.norm(n_jkl, dim=-1)
335
-
336
- sign_phi = torch.sign(
337
- r_kj.unsqueeze(-2) @ torch.cross(n_ijk, n_jkl, dim=-1).unsqueeze(-1)
338
- ).squeeze(-1, -2)
339
- phi = sign_phi * torch.arccos(
340
- torch.clamp(
341
- (n_ijk.unsqueeze(-2) @ n_jkl.unsqueeze(-1)).squeeze(-1, -2)
342
- / (n_ijk_norm * n_jkl_norm),
343
- -1 + 1e-8,
344
- 1 - 1e-8,
345
- )
346
- )
347
-
348
- if not compute_gradient:
349
- return phi
350
-
351
- a = (
352
- (r_ij.unsqueeze(-2) @ r_kj.unsqueeze(-1)).squeeze(-1, -2) / (r_kj_norm**2)
353
- ).unsqueeze(-1)
354
- b = (
355
- (r_kl.unsqueeze(-2) @ r_kj.unsqueeze(-1)).squeeze(-1, -2) / (r_kj_norm**2)
356
- ).unsqueeze(-1)
357
-
358
- grad_i = n_ijk * (r_kj_norm / n_ijk_norm**2).unsqueeze(-1)
359
- grad_l = -1 * n_jkl * (r_kj_norm / n_jkl_norm**2).unsqueeze(-1)
360
- grad_j = (a - 1) * grad_i - b * grad_l
361
- grad_k = (b - 1) * grad_l - a * grad_i
362
- grad = torch.stack((grad_i, grad_j, grad_k, grad_l), dim=1)
363
- return phi, grad
364
-
365
-
366
- class AbsDihedralPotential(DihedralPotential):
367
- def compute_variable(
368
- self, coords, index, ref_coords=None, ref_mask=None, compute_gradient=False
369
- ):
370
- if not compute_gradient:
371
- phi = super().compute_variable(
372
- coords, index, compute_gradient=compute_gradient
373
- )
374
- phi = torch.abs(phi)
375
- return phi
376
-
377
- phi, grad = super().compute_variable(
378
- coords, index, compute_gradient=compute_gradient
379
- )
380
- grad[(phi < 0)[..., None, :, None].expand_as(grad)] *= -1
381
- phi = torch.abs(phi)
382
-
383
- return phi, grad
384
-
385
-
386
- class PoseBustersPotential(FlatBottomPotential, DistancePotential):
387
- def compute_args(self, feats, parameters):
388
- pair_index = feats["rdkit_bounds_index"][0]
389
- lower_bounds = feats["rdkit_lower_bounds"][0].clone()
390
- upper_bounds = feats["rdkit_upper_bounds"][0].clone()
391
- bond_mask = feats["rdkit_bounds_bond_mask"][0]
392
- angle_mask = feats["rdkit_bounds_angle_mask"][0]
393
-
394
- lower_bounds[bond_mask * ~angle_mask] *= 1.0 - parameters["bond_buffer"]
395
- upper_bounds[bond_mask * ~angle_mask] *= 1.0 + parameters["bond_buffer"]
396
- lower_bounds[~bond_mask * angle_mask] *= 1.0 - parameters["angle_buffer"]
397
- upper_bounds[~bond_mask * angle_mask] *= 1.0 + parameters["angle_buffer"]
398
- lower_bounds[bond_mask * angle_mask] *= 1.0 - min(
399
- parameters["bond_buffer"], parameters["angle_buffer"]
400
- )
401
- upper_bounds[bond_mask * angle_mask] *= 1.0 + min(
402
- parameters["bond_buffer"], parameters["angle_buffer"]
403
- )
404
- lower_bounds[~bond_mask * ~angle_mask] *= 1.0 - parameters["clash_buffer"]
405
- upper_bounds[~bond_mask * ~angle_mask] = float("inf")
406
-
407
- vdw_radii = torch.zeros(
408
- const.num_elements, dtype=torch.float32, device=pair_index.device
409
- )
410
- vdw_radii[1:119] = torch.tensor(
411
- const.vdw_radii, dtype=torch.float32, device=pair_index.device
412
- )
413
- atom_vdw_radii = (
414
- feats["ref_element"].float() @ vdw_radii.unsqueeze(-1)
415
- ).squeeze(-1)[0]
416
- bond_cutoffs = 0.35 + atom_vdw_radii[pair_index].mean(dim=0)
417
- lower_bounds[~bond_mask] = torch.max(lower_bounds[~bond_mask], bond_cutoffs[~bond_mask])
418
- upper_bounds[bond_mask] = torch.min(upper_bounds[bond_mask], bond_cutoffs[bond_mask])
419
-
420
- k = torch.ones_like(lower_bounds)
421
-
422
- return pair_index, (k, lower_bounds, upper_bounds), None, None, None
423
-
424
-
425
- class ConnectionsPotential(FlatBottomPotential, DistancePotential):
426
- def compute_args(self, feats, parameters):
427
- pair_index = feats["connected_atom_index"][0]
428
- lower_bounds = None
429
- upper_bounds = torch.full(
430
- (pair_index.shape[1],), parameters["buffer"], device=pair_index.device
431
- )
432
- k = torch.ones_like(upper_bounds)
433
-
434
- return pair_index, (k, lower_bounds, upper_bounds), None, None, None
435
-
436
-
437
- class VDWOverlapPotential(FlatBottomPotential, DistancePotential):
438
- def compute_args(self, feats, parameters):
439
- atom_chain_id = (
440
- torch.bmm(
441
- feats["atom_to_token"].float(), feats["asym_id"].unsqueeze(-1).float()
442
- )
443
- .squeeze(-1)
444
- .long()
445
- )[0]
446
- atom_pad_mask = feats["atom_pad_mask"][0].bool()
447
- chain_sizes = torch.bincount(atom_chain_id[atom_pad_mask])
448
- single_ion_mask = (chain_sizes > 1)[atom_chain_id]
449
-
450
- vdw_radii = torch.zeros(
451
- const.num_elements, dtype=torch.float32, device=atom_chain_id.device
452
- )
453
- vdw_radii[1:119] = torch.tensor(
454
- const.vdw_radii, dtype=torch.float32, device=atom_chain_id.device
455
- )
456
- atom_vdw_radii = (
457
- feats["ref_element"].float() @ vdw_radii.unsqueeze(-1)
458
- ).squeeze(-1)[0]
459
-
460
- pair_index = torch.triu_indices(
461
- atom_chain_id.shape[0],
462
- atom_chain_id.shape[0],
463
- 1,
464
- device=atom_chain_id.device,
465
- )
466
-
467
- pair_pad_mask = atom_pad_mask[pair_index].all(dim=0)
468
- pair_ion_mask = single_ion_mask[pair_index[0]] * single_ion_mask[pair_index[1]]
469
-
470
- num_chains = atom_chain_id.max() + 1
471
- connected_chain_index = feats["connected_chain_index"][0]
472
- connected_chain_matrix = torch.eye(
473
- num_chains, device=atom_chain_id.device, dtype=torch.bool
474
- )
475
- connected_chain_matrix[connected_chain_index[0], connected_chain_index[1]] = (
476
- True
477
- )
478
- connected_chain_matrix[connected_chain_index[1], connected_chain_index[0]] = (
479
- True
480
- )
481
- connected_chain_mask = connected_chain_matrix[
482
- atom_chain_id[pair_index[0]], atom_chain_id[pair_index[1]]
483
- ]
484
-
485
- pair_index = pair_index[
486
- :, pair_pad_mask * pair_ion_mask * ~connected_chain_mask
487
- ]
488
-
489
- lower_bounds = atom_vdw_radii[pair_index].sum(dim=0) * (
490
- 1.0 - parameters["buffer"]
491
- )
492
- upper_bounds = None
493
- k = torch.ones_like(lower_bounds)
494
-
495
- return pair_index, (k, lower_bounds, upper_bounds), None, None, None
496
-
497
-
498
- class SymmetricChainCOMPotential(FlatBottomPotential, DistancePotential):
499
- def compute_args(self, feats, parameters):
500
- atom_chain_id = (
501
- torch.bmm(
502
- feats["atom_to_token"].float(), feats["asym_id"].unsqueeze(-1).float()
503
- )
504
- .squeeze(-1)
505
- .long()
506
- )[0]
507
- atom_pad_mask = feats["atom_pad_mask"][0].bool()
508
- chain_sizes = torch.bincount(atom_chain_id[atom_pad_mask])
509
- single_ion_mask = chain_sizes > 1
510
-
511
- pair_index = feats["symmetric_chain_index"][0]
512
- pair_ion_mask = single_ion_mask[pair_index[0]] * single_ion_mask[pair_index[1]]
513
- pair_index = pair_index[:, pair_ion_mask]
514
- lower_bounds = torch.full(
515
- (pair_index.shape[1],),
516
- parameters["buffer"],
517
- dtype=torch.float32,
518
- device=pair_index.device,
519
- )
520
- upper_bounds = None
521
- k = torch.ones_like(lower_bounds)
522
-
523
- return (
524
- pair_index,
525
- (k, lower_bounds, upper_bounds),
526
- (atom_chain_id, atom_pad_mask),
527
- None,
528
- None,
529
- )
530
-
531
-
532
- class StereoBondPotential(FlatBottomPotential, AbsDihedralPotential):
533
- def compute_args(self, feats, parameters):
534
- stereo_bond_index = feats["stereo_bond_index"][0]
535
- stereo_bond_orientations = feats["stereo_bond_orientations"][0].bool()
536
-
537
- lower_bounds = torch.zeros(
538
- stereo_bond_orientations.shape, device=stereo_bond_orientations.device
539
- )
540
- upper_bounds = torch.zeros(
541
- stereo_bond_orientations.shape, device=stereo_bond_orientations.device
542
- )
543
- lower_bounds[stereo_bond_orientations] = torch.pi - parameters["buffer"]
544
- upper_bounds[stereo_bond_orientations] = float("inf")
545
- lower_bounds[~stereo_bond_orientations] = float("-inf")
546
- upper_bounds[~stereo_bond_orientations] = parameters["buffer"]
547
-
548
- k = torch.ones_like(lower_bounds)
549
-
550
- return stereo_bond_index, (k, lower_bounds, upper_bounds), None, None, None
551
-
552
-
553
- class ChiralAtomPotential(FlatBottomPotential, DihedralPotential):
554
- def compute_args(self, feats, parameters):
555
- chiral_atom_index = feats["chiral_atom_index"][0]
556
- chiral_atom_orientations = feats["chiral_atom_orientations"][0].bool()
557
-
558
- lower_bounds = torch.zeros(
559
- chiral_atom_orientations.shape, device=chiral_atom_orientations.device
560
- )
561
- upper_bounds = torch.zeros(
562
- chiral_atom_orientations.shape, device=chiral_atom_orientations.device
563
- )
564
- lower_bounds[chiral_atom_orientations] = parameters["buffer"]
565
- upper_bounds[chiral_atom_orientations] = float("inf")
566
- upper_bounds[~chiral_atom_orientations] = -1 * parameters["buffer"]
567
- lower_bounds[~chiral_atom_orientations] = float("-inf")
568
-
569
- k = torch.ones_like(lower_bounds)
570
- return chiral_atom_index, (k, lower_bounds, upper_bounds), None, None, None
571
-
572
-
573
- class PlanarBondPotential(FlatBottomPotential, AbsDihedralPotential):
574
- def compute_args(self, feats, parameters):
575
- double_bond_index = feats["planar_bond_index"][0].T
576
- double_bond_improper_index = torch.tensor(
577
- [
578
- [1, 2, 3, 0],
579
- [4, 5, 0, 3],
580
- ],
581
- device=double_bond_index.device,
582
- ).T
583
- improper_index = (
584
- double_bond_index[:, double_bond_improper_index]
585
- .swapaxes(0, 1)
586
- .flatten(start_dim=1)
587
- )
588
- lower_bounds = None
589
- upper_bounds = torch.full(
590
- (improper_index.shape[1],),
591
- parameters["buffer"],
592
- device=improper_index.device,
593
- )
594
- k = torch.ones_like(upper_bounds)
595
-
596
- return improper_index, (k, lower_bounds, upper_bounds), None, None, None
597
-
598
-
599
- class TemplateReferencePotential(FlatBottomPotential, ReferencePotential):
600
- def compute_args(self, feats, parameters):
601
- if "template_mask_cb" not in feats or "template_force" not in feats:
602
- return torch.empty([1, 0]), None, None, None, None
603
-
604
- template_mask = feats["template_mask_cb"][feats["template_force"]]
605
- if template_mask.shape[0] == 0:
606
- return torch.empty([1, 0]), None, None, None, None
607
-
608
- ref_coords = feats["template_cb"][feats["template_force"]].clone()
609
- ref_mask = feats["template_mask_cb"][feats["template_force"]].clone()
610
- ref_atom_index = (
611
- torch.bmm(
612
- feats["token_to_rep_atom"].float(),
613
- torch.arange(
614
- feats["atom_pad_mask"].shape[1],
615
- device=feats["atom_pad_mask"].device,
616
- dtype=torch.float32,
617
- )[None, :, None],
618
- )
619
- .squeeze(-1)
620
- .long()
621
- )[0]
622
- ref_token_index = (
623
- torch.bmm(
624
- feats["atom_to_token"].float(),
625
- feats["token_index"].unsqueeze(-1).float(),
626
- )
627
- .squeeze(-1)
628
- .long()
629
- )[0]
630
-
631
- index = torch.arange(
632
- template_mask.shape[-1], dtype=torch.long, device=template_mask.device
633
- )[None]
634
- upper_bounds = torch.full(
635
- template_mask.shape, float("inf"), device=index.device, dtype=torch.float32
636
- )
637
- ref_idxs = torch.argwhere(template_mask).T
638
- upper_bounds[ref_idxs.unbind()] = feats["template_force_threshold"][
639
- feats["template_force"]
640
- ][ref_idxs[0]]
641
-
642
- lower_bounds = None
643
- k = torch.ones_like(upper_bounds)
644
- return (
645
- index,
646
- (k, lower_bounds, upper_bounds),
647
- None,
648
- (ref_coords, ref_mask, ref_atom_index, ref_token_index),
649
- None,
650
- )
651
-
652
-
653
- class ContactPotentital(FlatBottomPotential, DistancePotential):
654
- def compute_args(self, feats, parameters):
655
- index = feats["contact_pair_index"][0]
656
- union_index = feats["contact_union_index"][0]
657
- negation_mask = feats["contact_negation_mask"][0]
658
- lower_bounds = None
659
- upper_bounds = feats["contact_thresholds"][0].clone()
660
- k = torch.ones_like(upper_bounds)
661
- return (
662
- index,
663
- (k, lower_bounds, upper_bounds),
664
- None,
665
- None,
666
- (negation_mask, union_index),
667
- )
668
-
669
-
670
- def get_potentials(steering_args, boltz2=False):
671
- potentials = []
672
- if steering_args["fk_steering"] or steering_args["physical_guidance_update"]:
673
- potentials.extend(
674
- [
675
- SymmetricChainCOMPotential(
676
- parameters={
677
- "guidance_interval": 4,
678
- "guidance_weight": 0.5
679
- if steering_args["physical_guidance_update"]
680
- else 0.0,
681
- "resampling_weight": 0.5,
682
- "buffer": ExponentialInterpolation(
683
- start=1.0, end=5.0, alpha=-2.0
684
- ),
685
- }
686
- ),
687
- VDWOverlapPotential(
688
- parameters={
689
- "guidance_interval": 5,
690
- "guidance_weight": (
691
- PiecewiseStepFunction(thresholds=[0.4], values=[0.125, 0.0])
692
- if steering_args["physical_guidance_update"]
693
- else 0.0
694
- ),
695
- "resampling_weight": PiecewiseStepFunction(
696
- thresholds=[0.6], values=[0.01, 0.0]
697
- ),
698
- "buffer": 0.225,
699
- }
700
- ),
701
- ConnectionsPotential(
702
- parameters={
703
- "guidance_interval": 1,
704
- "guidance_weight": 0.15
705
- if steering_args["physical_guidance_update"]
706
- else 0.0,
707
- "resampling_weight": 1.0,
708
- "buffer": 2.0,
709
- }
710
- ),
711
- PoseBustersPotential(
712
- parameters={
713
- "guidance_interval": 1,
714
- "guidance_weight": 0.01
715
- if steering_args["physical_guidance_update"]
716
- else 0.0,
717
- "resampling_weight": 0.1,
718
- "bond_buffer": 0.125,
719
- "angle_buffer": 0.125,
720
- "clash_buffer": 0.10,
721
- }
722
- ),
723
- ChiralAtomPotential(
724
- parameters={
725
- "guidance_interval": 1,
726
- "guidance_weight": 0.1
727
- if steering_args["physical_guidance_update"]
728
- else 0.0,
729
- "resampling_weight": 1.0,
730
- "buffer": 0.52360,
731
- }
732
- ),
733
- StereoBondPotential(
734
- parameters={
735
- "guidance_interval": 1,
736
- "guidance_weight": 0.05
737
- if steering_args["physical_guidance_update"]
738
- else 0.0,
739
- "resampling_weight": 1.0,
740
- "buffer": 0.52360,
741
- }
742
- ),
743
- PlanarBondPotential(
744
- parameters={
745
- "guidance_interval": 1,
746
- "guidance_weight": 0.05
747
- if steering_args["physical_guidance_update"]
748
- else 0.0,
749
- "resampling_weight": 1.0,
750
- "buffer": 0.26180,
751
- }
752
- ),
753
- ]
754
- )
755
- if boltz2 and (
756
- steering_args["fk_steering"] or steering_args["contact_guidance_update"]
757
- ):
758
- potentials.extend(
759
- [
760
- ContactPotentital(
761
- parameters={
762
- "guidance_interval": 4,
763
- "guidance_weight": (
764
- PiecewiseStepFunction(
765
- thresholds=[0.25, 0.75], values=[0.0, 0.5, 1.0]
766
- )
767
- if steering_args["contact_guidance_update"]
768
- else 0.0
769
- ),
770
- "resampling_weight": 1.0,
771
- "union_lambda": ExponentialInterpolation(
772
- start=8.0, end=0.0, alpha=-2.0
773
- ),
774
- }
775
- ),
776
- TemplateReferencePotential(
777
- parameters={
778
- "guidance_interval": 2,
779
- "guidance_weight": 0.1
780
- if steering_args["contact_guidance_update"]
781
- else 0.0,
782
- "resampling_weight": 1.0,
783
- }
784
- ),
785
- ]
786
- )
787
- return potentials
 
1
+ from abc import ABC, abstractmethod
2
+ from typing import Optional, Dict, Any, Set, List, Union
3
+
4
+ import torch
5
+ import numpy as np
6
+ from . import vb_const as const
7
+ from .vb_potentials_schedules import (
8
+ ParameterSchedule,
9
+ ExponentialInterpolation,
10
+ PiecewiseStepFunction,
11
+ )
12
+ from .vb_loss_diffusionv2 import weighted_rigid_align
13
+
14
+
15
+ class Potential(ABC):
16
+ def __init__(
17
+ self,
18
+ parameters: Optional[
19
+ Dict[str, Union[ParameterSchedule, float, int, bool]]
20
+ ] = None,
21
+ ):
22
+ self.parameters = parameters
23
+
24
+ def compute(self, coords, feats, parameters):
25
+ index, args, com_args, ref_args, operator_args = self.compute_args(
26
+ feats, parameters
27
+ )
28
+
29
+ if index.shape[1] == 0:
30
+ return torch.zeros(coords.shape[:-2], device=coords.device)
31
+
32
+ if com_args is not None:
33
+ com_index, atom_pad_mask = com_args
34
+ unpad_com_index = com_index[atom_pad_mask]
35
+ unpad_coords = coords[..., atom_pad_mask, :]
36
+ coords = torch.zeros(
37
+ (*unpad_coords.shape[:-2], unpad_com_index.max() + 1, 3),
38
+ device=coords.device,
39
+ ).scatter_reduce(
40
+ -2,
41
+ unpad_com_index.unsqueeze(-1).expand_as(unpad_coords),
42
+ unpad_coords,
43
+ "mean",
44
+ )
45
+ else:
46
+ com_index, atom_pad_mask = None, None
47
+
48
+ if ref_args is not None:
49
+ ref_coords, ref_mask, ref_atom_index, ref_token_index = ref_args
50
+ coords = coords[..., ref_atom_index, :]
51
+ else:
52
+ ref_coords, ref_mask, ref_atom_index, ref_token_index = (
53
+ None,
54
+ None,
55
+ None,
56
+ None,
57
+ )
58
+
59
+ if operator_args is not None:
60
+ negation_mask, union_index = operator_args
61
+ else:
62
+ negation_mask, union_index = None, None
63
+
64
+ value = self.compute_variable(
65
+ coords,
66
+ index,
67
+ ref_coords=ref_coords,
68
+ ref_mask=ref_mask,
69
+ compute_gradient=False,
70
+ )
71
+ energy = self.compute_function(
72
+ value, *args, negation_mask=negation_mask, compute_derivative=False
73
+ )
74
+
75
+ if union_index is not None:
76
+ neg_exp_energy = torch.exp(-1 * parameters["union_lambda"] * energy)
77
+ Z = torch.zeros(
78
+ (*energy.shape[:-1], union_index.max() + 1), device=union_index.device
79
+ ).scatter_reduce(
80
+ -1,
81
+ union_index.expand_as(neg_exp_energy),
82
+ neg_exp_energy,
83
+ "sum",
84
+ )
85
+ softmax_energy = neg_exp_energy / Z[..., union_index]
86
+ softmax_energy[Z[..., union_index] == 0] = 0
87
+ return (energy * softmax_energy).sum(dim=-1)
88
+
89
+ return energy.sum(dim=tuple(range(1, energy.dim())))
90
+
91
+ def compute_gradient(self, coords, feats, parameters):
92
+ index, args, com_args, ref_args, operator_args = self.compute_args(
93
+ feats, parameters
94
+ )
95
+ if index.shape[1] == 0:
96
+ return torch.zeros_like(coords)
97
+
98
+ if com_args is not None:
99
+ com_index, atom_pad_mask = com_args
100
+ unpad_coords = coords[..., atom_pad_mask, :]
101
+ unpad_com_index = com_index[atom_pad_mask]
102
+ coords = torch.zeros(
103
+ (*unpad_coords.shape[:-2], unpad_com_index.max() + 1, 3),
104
+ device=coords.device,
105
+ ).scatter_reduce(
106
+ -2,
107
+ unpad_com_index.unsqueeze(-1).expand_as(unpad_coords),
108
+ unpad_coords,
109
+ "mean",
110
+ )
111
+ com_counts = torch.bincount(com_index[atom_pad_mask])
112
+ else:
113
+ com_index, atom_pad_mask = None, None
114
+
115
+ if ref_args is not None:
116
+ ref_coords, ref_mask, ref_atom_index, ref_token_index = ref_args
117
+ coords = coords[..., ref_atom_index, :]
118
+ else:
119
+ ref_coords, ref_mask, ref_atom_index, ref_token_index = (
120
+ None,
121
+ None,
122
+ None,
123
+ None,
124
+ )
125
+
126
+ if operator_args is not None:
127
+ negation_mask, union_index = operator_args
128
+ else:
129
+ negation_mask, union_index = None, None
130
+
131
+ value, grad_value = self.compute_variable(
132
+ coords,
133
+ index,
134
+ ref_coords=ref_coords,
135
+ ref_mask=ref_mask,
136
+ compute_gradient=True,
137
+ )
138
+ energy, dEnergy = self.compute_function(
139
+ value,
140
+ *args, negation_mask=negation_mask, compute_derivative=True
141
+ )
142
+ if union_index is not None:
143
+ neg_exp_energy = torch.exp(-1 * parameters["union_lambda"] * energy)
144
+ Z = torch.zeros(
145
+ (*energy.shape[:-1], union_index.max() + 1), device=union_index.device
146
+ ).scatter_reduce(
147
+ -1,
148
+ union_index.expand_as(energy),
149
+ neg_exp_energy,
150
+ "sum",
151
+ )
152
+ softmax_energy = neg_exp_energy / Z[..., union_index]
153
+ softmax_energy[Z[..., union_index] == 0] = 0
154
+ f = torch.zeros(
155
+ (*energy.shape[:-1], union_index.max() + 1), device=union_index.device
156
+ ).scatter_reduce(
157
+ -1,
158
+ union_index.expand_as(energy),
159
+ energy * softmax_energy,
160
+ "sum",
161
+ )
162
+ dSoftmax = (
163
+ dEnergy
164
+ * softmax_energy
165
+ * (1 + parameters["union_lambda"] * (energy - f[..., union_index]))
166
+ )
167
+ prod = dSoftmax.tile(grad_value.shape[-3]).unsqueeze(
168
+ -1
169
+ ) * grad_value.flatten(start_dim=-3, end_dim=-2)
170
+ if prod.dim() > 3:
171
+ prod = prod.sum(dim=list(range(1, prod.dim() - 2)))
172
+ grad_atom = torch.zeros_like(coords).scatter_reduce(
173
+ -2,
174
+ index.flatten(start_dim=0, end_dim=1)
175
+ .unsqueeze(-1)
176
+ .expand((*coords.shape[:-2], -1, 3)),
177
+ prod,
178
+ "sum",
179
+ )
180
+ else:
181
+ prod = dEnergy.tile(grad_value.shape[-3]).unsqueeze(
182
+ -1
183
+ ) * grad_value.flatten(start_dim=-3, end_dim=-2)
184
+ if prod.dim() > 3:
185
+ prod = prod.sum(dim=list(range(1, prod.dim() - 2)))
186
+ grad_atom = torch.zeros_like(coords).scatter_reduce(
187
+ -2,
188
+ index.flatten(start_dim=0, end_dim=1)
189
+ .unsqueeze(-1)
190
+ .expand((*coords.shape[:-2], -1, 3)), # 9 x 516 x 3
191
+ prod,
192
+ "sum",
193
+ )
194
+
195
+ if com_index is not None:
196
+ grad_atom = grad_atom[..., com_index, :]
197
+ elif ref_token_index is not None:
198
+ grad_atom = grad_atom[..., ref_token_index, :]
199
+
200
+ return grad_atom
201
+
202
+ def compute_parameters(self, t):
203
+ if self.parameters is None:
204
+ return None
205
+ parameters = {
206
+ name: parameter
207
+ if not isinstance(parameter, ParameterSchedule)
208
+ else parameter.compute(t)
209
+ for name, parameter in self.parameters.items()
210
+ }
211
+ return parameters
212
+
213
+ @abstractmethod
214
+ def compute_function(
215
+ self, value, *args, negation_mask=None, compute_derivative=False
216
+ ):
217
+ raise NotImplementedError
218
+
219
+ @abstractmethod
220
+ def compute_variable(self, coords, index, compute_gradient=False):
221
+ raise NotImplementedError
222
+
223
+ @abstractmethod
224
+ def compute_args(self, t, feats, **parameters):
225
+ raise NotImplementedError
226
+
227
+ def get_reference_coords(self, feats, parameters):
228
+ return None, None
229
+
230
+
231
+ class FlatBottomPotential(Potential):
232
+ def compute_function(
233
+ self,
234
+ value,
235
+ k,
236
+ lower_bounds,
237
+ upper_bounds,
238
+ negation_mask=None,
239
+ compute_derivative=False,
240
+ ):
241
+ if lower_bounds is None:
242
+ lower_bounds = torch.full_like(value, float("-inf"))
243
+ if upper_bounds is None:
244
+ upper_bounds = torch.full_like(value, float("inf"))
245
+ lower_bounds = lower_bounds.expand_as(value).clone()
246
+ upper_bounds = upper_bounds.expand_as(value).clone()
247
+
248
+ if negation_mask is not None:
249
+ unbounded_below_mask = torch.isneginf(lower_bounds)
250
+ unbounded_above_mask = torch.isposinf(upper_bounds)
251
+ unbounded_mask = unbounded_below_mask + unbounded_above_mask
252
+ assert torch.all(unbounded_mask + negation_mask)
253
+ lower_bounds[~unbounded_above_mask * ~negation_mask] = upper_bounds[
254
+ ~unbounded_above_mask * ~negation_mask
255
+ ]
256
+ upper_bounds[~unbounded_above_mask * ~negation_mask] = float("inf")
257
+ upper_bounds[~unbounded_below_mask * ~negation_mask] = lower_bounds[
258
+ ~unbounded_below_mask * ~negation_mask
259
+ ]
260
+ lower_bounds[~unbounded_below_mask * ~negation_mask] = float("-inf")
261
+
262
+ neg_overflow_mask = value < lower_bounds
263
+ pos_overflow_mask = value > upper_bounds
264
+
265
+ energy = torch.zeros_like(value)
266
+ energy[neg_overflow_mask] = (k * (lower_bounds - value))[neg_overflow_mask]
267
+ energy[pos_overflow_mask] = (k * (value - upper_bounds))[pos_overflow_mask]
268
+ if not compute_derivative:
269
+ return energy
270
+
271
+ dEnergy = torch.zeros_like(value)
272
+ dEnergy[neg_overflow_mask] = (
273
+ -1 * k.expand_as(neg_overflow_mask)[neg_overflow_mask]
274
+ )
275
+ dEnergy[pos_overflow_mask] = (
276
+ 1 * k.expand_as(pos_overflow_mask)[pos_overflow_mask]
277
+ )
278
+
279
+ return energy, dEnergy
280
+
281
+
282
+ class ReferencePotential(Potential):
283
+ def compute_variable(
284
+ self, coords, index, ref_coords, ref_mask, compute_gradient=False
285
+ ):
286
+ aligned_ref_coords = weighted_rigid_align(
287
+ ref_coords.float(),
288
+ coords[:, index].float(),
289
+ ref_mask,
290
+ ref_mask,
291
+ )
292
+
293
+ r = coords[:, index] - aligned_ref_coords
294
+ r_norm = torch.linalg.norm(r, dim=-1)
295
+
296
+ if not compute_gradient:
297
+ return r_norm
298
+
299
+ r_hat = r / r_norm.unsqueeze(-1)
300
+ grad = (r_hat * ref_mask.unsqueeze(-1)).unsqueeze(1)
301
+ return r_norm, grad
302
+
303
+
304
+ class DistancePotential(Potential):
305
+ def compute_variable(
306
+ self, coords, index, ref_coords=None, ref_mask=None, compute_gradient=False
307
+ ):
308
+ r_ij = coords.index_select(-2, index[0]) - coords.index_select(-2, index[1])
309
+ r_ij_norm = torch.linalg.norm(r_ij, dim=-1)
310
+ r_hat_ij = r_ij / r_ij_norm.unsqueeze(-1)
311
+
312
+ if not compute_gradient:
313
+ return r_ij_norm
314
+
315
+ grad_i = r_hat_ij
316
+ grad_j = -1 * r_hat_ij
317
+ grad = torch.stack((grad_i, grad_j), dim=1)
318
+ return r_ij_norm, grad
319
+
320
+
321
+ class DihedralPotential(Potential):
322
+ def compute_variable(
323
+ self, coords, index, ref_coords=None, ref_mask=None, compute_gradient=False
324
+ ):
325
+ r_ij = coords.index_select(-2, index[0]) - coords.index_select(-2, index[1])
326
+ r_kj = coords.index_select(-2, index[2]) - coords.index_select(-2, index[1])
327
+ r_kl = coords.index_select(-2, index[2]) - coords.index_select(-2, index[3])
328
+
329
+ n_ijk = torch.cross(r_ij, r_kj, dim=-1)
330
+ n_jkl = torch.cross(r_kj, r_kl, dim=-1)
331
+
332
+ r_kj_norm = torch.linalg.norm(r_kj, dim=-1)
333
+ n_ijk_norm = torch.linalg.norm(n_ijk, dim=-1)
334
+ n_jkl_norm = torch.linalg.norm(n_jkl, dim=-1)
335
+
336
+ sign_phi = torch.sign(
337
+ r_kj.unsqueeze(-2) @ torch.cross(n_ijk, n_jkl, dim=-1).unsqueeze(-1)
338
+ ).squeeze(-1, -2)
339
+ phi = sign_phi * torch.arccos(
340
+ torch.clamp(
341
+ (n_ijk.unsqueeze(-2) @ n_jkl.unsqueeze(-1)).squeeze(-1, -2)
342
+ / (n_ijk_norm * n_jkl_norm),
343
+ -1 + 1e-8,
344
+ 1 - 1e-8,
345
+ )
346
+ )
347
+
348
+ if not compute_gradient:
349
+ return phi
350
+
351
+ a = (
352
+ (r_ij.unsqueeze(-2) @ r_kj.unsqueeze(-1)).squeeze(-1, -2) / (r_kj_norm**2)
353
+ ).unsqueeze(-1)
354
+ b = (
355
+ (r_kl.unsqueeze(-2) @ r_kj.unsqueeze(-1)).squeeze(-1, -2) / (r_kj_norm**2)
356
+ ).unsqueeze(-1)
357
+
358
+ grad_i = n_ijk * (r_kj_norm / n_ijk_norm**2).unsqueeze(-1)
359
+ grad_l = -1 * n_jkl * (r_kj_norm / n_jkl_norm**2).unsqueeze(-1)
360
+ grad_j = (a - 1) * grad_i - b * grad_l
361
+ grad_k = (b - 1) * grad_l - a * grad_i
362
+ grad = torch.stack((grad_i, grad_j, grad_k, grad_l), dim=1)
363
+ return phi, grad
364
+
365
+
366
+ class AbsDihedralPotential(DihedralPotential):
367
+ def compute_variable(
368
+ self, coords, index, ref_coords=None, ref_mask=None, compute_gradient=False
369
+ ):
370
+ if not compute_gradient:
371
+ phi = super().compute_variable(
372
+ coords, index, compute_gradient=compute_gradient
373
+ )
374
+ phi = torch.abs(phi)
375
+ return phi
376
+
377
+ phi, grad = super().compute_variable(
378
+ coords, index, compute_gradient=compute_gradient
379
+ )
380
+ grad[(phi < 0)[..., None, :, None].expand_as(grad)] *= -1
381
+ phi = torch.abs(phi)
382
+
383
+ return phi, grad
384
+
385
+
386
+ class PoseBustersPotential(FlatBottomPotential, DistancePotential):
387
+ def compute_args(self, feats, parameters):
388
+ pair_index = feats["rdkit_bounds_index"][0]
389
+ lower_bounds = feats["rdkit_lower_bounds"][0].clone()
390
+ upper_bounds = feats["rdkit_upper_bounds"][0].clone()
391
+ bond_mask = feats["rdkit_bounds_bond_mask"][0]
392
+ angle_mask = feats["rdkit_bounds_angle_mask"][0]
393
+
394
+ lower_bounds[bond_mask * ~angle_mask] *= 1.0 - parameters["bond_buffer"]
395
+ upper_bounds[bond_mask * ~angle_mask] *= 1.0 + parameters["bond_buffer"]
396
+ lower_bounds[~bond_mask * angle_mask] *= 1.0 - parameters["angle_buffer"]
397
+ upper_bounds[~bond_mask * angle_mask] *= 1.0 + parameters["angle_buffer"]
398
+ lower_bounds[bond_mask * angle_mask] *= 1.0 - min(
399
+ parameters["bond_buffer"], parameters["angle_buffer"]
400
+ )
401
+ upper_bounds[bond_mask * angle_mask] *= 1.0 + min(
402
+ parameters["bond_buffer"], parameters["angle_buffer"]
403
+ )
404
+ lower_bounds[~bond_mask * ~angle_mask] *= 1.0 - parameters["clash_buffer"]
405
+ upper_bounds[~bond_mask * ~angle_mask] = float("inf")
406
+
407
+ vdw_radii = torch.zeros(
408
+ const.num_elements, dtype=torch.float32, device=pair_index.device
409
+ )
410
+ vdw_radii[1:119] = torch.tensor(
411
+ const.vdw_radii, dtype=torch.float32, device=pair_index.device
412
+ )
413
+ atom_vdw_radii = (
414
+ feats["ref_element"].float() @ vdw_radii.unsqueeze(-1)
415
+ ).squeeze(-1)[0]
416
+ bond_cutoffs = 0.35 + atom_vdw_radii[pair_index].mean(dim=0)
417
+ lower_bounds[~bond_mask] = torch.max(lower_bounds[~bond_mask], bond_cutoffs[~bond_mask])
418
+ upper_bounds[bond_mask] = torch.min(upper_bounds[bond_mask], bond_cutoffs[bond_mask])
419
+
420
+ k = torch.ones_like(lower_bounds)
421
+
422
+ return pair_index, (k, lower_bounds, upper_bounds), None, None, None
423
+
424
+
425
+ class ConnectionsPotential(FlatBottomPotential, DistancePotential):
426
+ def compute_args(self, feats, parameters):
427
+ pair_index = feats["connected_atom_index"][0]
428
+ lower_bounds = None
429
+ upper_bounds = torch.full(
430
+ (pair_index.shape[1],), parameters["buffer"], device=pair_index.device
431
+ )
432
+ k = torch.ones_like(upper_bounds)
433
+
434
+ return pair_index, (k, lower_bounds, upper_bounds), None, None, None
435
+
436
+
437
+ class VDWOverlapPotential(FlatBottomPotential, DistancePotential):
438
+ def compute_args(self, feats, parameters):
439
+ atom_chain_id = (
440
+ torch.bmm(
441
+ feats["atom_to_token"].float(), feats["asym_id"].unsqueeze(-1).float()
442
+ )
443
+ .squeeze(-1)
444
+ .long()
445
+ )[0]
446
+ atom_pad_mask = feats["atom_pad_mask"][0].bool()
447
+ chain_sizes = torch.bincount(atom_chain_id[atom_pad_mask])
448
+ single_ion_mask = (chain_sizes > 1)[atom_chain_id]
449
+
450
+ vdw_radii = torch.zeros(
451
+ const.num_elements, dtype=torch.float32, device=atom_chain_id.device
452
+ )
453
+ vdw_radii[1:119] = torch.tensor(
454
+ const.vdw_radii, dtype=torch.float32, device=atom_chain_id.device
455
+ )
456
+ atom_vdw_radii = (
457
+ feats["ref_element"].float() @ vdw_radii.unsqueeze(-1)
458
+ ).squeeze(-1)[0]
459
+
460
+ pair_index = torch.triu_indices(
461
+ atom_chain_id.shape[0],
462
+ atom_chain_id.shape[0],
463
+ 1,
464
+ device=atom_chain_id.device,
465
+ )
466
+
467
+ pair_pad_mask = atom_pad_mask[pair_index].all(dim=0)
468
+ pair_ion_mask = single_ion_mask[pair_index[0]] * single_ion_mask[pair_index[1]]
469
+
470
+ num_chains = atom_chain_id.max() + 1
471
+ connected_chain_index = feats["connected_chain_index"][0]
472
+ connected_chain_matrix = torch.eye(
473
+ num_chains, device=atom_chain_id.device, dtype=torch.bool
474
+ )
475
+ connected_chain_matrix[connected_chain_index[0], connected_chain_index[1]] = (
476
+ True
477
+ )
478
+ connected_chain_matrix[connected_chain_index[1], connected_chain_index[0]] = (
479
+ True
480
+ )
481
+ connected_chain_mask = connected_chain_matrix[
482
+ atom_chain_id[pair_index[0]], atom_chain_id[pair_index[1]]
483
+ ]
484
+
485
+ pair_index = pair_index[
486
+ :, pair_pad_mask * pair_ion_mask * ~connected_chain_mask
487
+ ]
488
+
489
+ lower_bounds = atom_vdw_radii[pair_index].sum(dim=0) * (
490
+ 1.0 - parameters["buffer"]
491
+ )
492
+ upper_bounds = None
493
+ k = torch.ones_like(lower_bounds)
494
+
495
+ return pair_index, (k, lower_bounds, upper_bounds), None, None, None
496
+
497
+
498
+ class SymmetricChainCOMPotential(FlatBottomPotential, DistancePotential):
499
+ def compute_args(self, feats, parameters):
500
+ atom_chain_id = (
501
+ torch.bmm(
502
+ feats["atom_to_token"].float(), feats["asym_id"].unsqueeze(-1).float()
503
+ )
504
+ .squeeze(-1)
505
+ .long()
506
+ )[0]
507
+ atom_pad_mask = feats["atom_pad_mask"][0].bool()
508
+ chain_sizes = torch.bincount(atom_chain_id[atom_pad_mask])
509
+ single_ion_mask = chain_sizes > 1
510
+
511
+ pair_index = feats["symmetric_chain_index"][0]
512
+ pair_ion_mask = single_ion_mask[pair_index[0]] * single_ion_mask[pair_index[1]]
513
+ pair_index = pair_index[:, pair_ion_mask]
514
+ lower_bounds = torch.full(
515
+ (pair_index.shape[1],),
516
+ parameters["buffer"],
517
+ dtype=torch.float32,
518
+ device=pair_index.device,
519
+ )
520
+ upper_bounds = None
521
+ k = torch.ones_like(lower_bounds)
522
+
523
+ return (
524
+ pair_index,
525
+ (k, lower_bounds, upper_bounds),
526
+ (atom_chain_id, atom_pad_mask),
527
+ None,
528
+ None,
529
+ )
530
+
531
+
532
+ class StereoBondPotential(FlatBottomPotential, AbsDihedralPotential):
533
+ def compute_args(self, feats, parameters):
534
+ stereo_bond_index = feats["stereo_bond_index"][0]
535
+ stereo_bond_orientations = feats["stereo_bond_orientations"][0].bool()
536
+
537
+ lower_bounds = torch.zeros(
538
+ stereo_bond_orientations.shape, device=stereo_bond_orientations.device
539
+ )
540
+ upper_bounds = torch.zeros(
541
+ stereo_bond_orientations.shape, device=stereo_bond_orientations.device
542
+ )
543
+ lower_bounds[stereo_bond_orientations] = torch.pi - parameters["buffer"]
544
+ upper_bounds[stereo_bond_orientations] = float("inf")
545
+ lower_bounds[~stereo_bond_orientations] = float("-inf")
546
+ upper_bounds[~stereo_bond_orientations] = parameters["buffer"]
547
+
548
+ k = torch.ones_like(lower_bounds)
549
+
550
+ return stereo_bond_index, (k, lower_bounds, upper_bounds), None, None, None
551
+
552
+
553
+ class ChiralAtomPotential(FlatBottomPotential, DihedralPotential):
554
+ def compute_args(self, feats, parameters):
555
+ chiral_atom_index = feats["chiral_atom_index"][0]
556
+ chiral_atom_orientations = feats["chiral_atom_orientations"][0].bool()
557
+
558
+ lower_bounds = torch.zeros(
559
+ chiral_atom_orientations.shape, device=chiral_atom_orientations.device
560
+ )
561
+ upper_bounds = torch.zeros(
562
+ chiral_atom_orientations.shape, device=chiral_atom_orientations.device
563
+ )
564
+ lower_bounds[chiral_atom_orientations] = parameters["buffer"]
565
+ upper_bounds[chiral_atom_orientations] = float("inf")
566
+ upper_bounds[~chiral_atom_orientations] = -1 * parameters["buffer"]
567
+ lower_bounds[~chiral_atom_orientations] = float("-inf")
568
+
569
+ k = torch.ones_like(lower_bounds)
570
+ return chiral_atom_index, (k, lower_bounds, upper_bounds), None, None, None
571
+
572
+
573
+ class PlanarBondPotential(FlatBottomPotential, AbsDihedralPotential):
574
+ def compute_args(self, feats, parameters):
575
+ double_bond_index = feats["planar_bond_index"][0].T
576
+ double_bond_improper_index = torch.tensor(
577
+ [
578
+ [1, 2, 3, 0],
579
+ [4, 5, 0, 3],
580
+ ],
581
+ device=double_bond_index.device,
582
+ ).T
583
+ improper_index = (
584
+ double_bond_index[:, double_bond_improper_index]
585
+ .swapaxes(0, 1)
586
+ .flatten(start_dim=1)
587
+ )
588
+ lower_bounds = None
589
+ upper_bounds = torch.full(
590
+ (improper_index.shape[1],),
591
+ parameters["buffer"],
592
+ device=improper_index.device,
593
+ )
594
+ k = torch.ones_like(upper_bounds)
595
+
596
+ return improper_index, (k, lower_bounds, upper_bounds), None, None, None
597
+
598
+
599
+ class TemplateReferencePotential(FlatBottomPotential, ReferencePotential):
600
+ def compute_args(self, feats, parameters):
601
+ if "template_mask_cb" not in feats or "template_force" not in feats:
602
+ return torch.empty([1, 0]), None, None, None, None
603
+
604
+ template_mask = feats["template_mask_cb"][feats["template_force"]]
605
+ if template_mask.shape[0] == 0:
606
+ return torch.empty([1, 0]), None, None, None, None
607
+
608
+ ref_coords = feats["template_cb"][feats["template_force"]].clone()
609
+ ref_mask = feats["template_mask_cb"][feats["template_force"]].clone()
610
+ ref_atom_index = (
611
+ torch.bmm(
612
+ feats["token_to_rep_atom"].float(),
613
+ torch.arange(
614
+ feats["atom_pad_mask"].shape[1],
615
+ device=feats["atom_pad_mask"].device,
616
+ dtype=torch.float32,
617
+ )[None, :, None],
618
+ )
619
+ .squeeze(-1)
620
+ .long()
621
+ )[0]
622
+ ref_token_index = (
623
+ torch.bmm(
624
+ feats["atom_to_token"].float(),
625
+ feats["token_index"].unsqueeze(-1).float(),
626
+ )
627
+ .squeeze(-1)
628
+ .long()
629
+ )[0]
630
+
631
+ index = torch.arange(
632
+ template_mask.shape[-1], dtype=torch.long, device=template_mask.device
633
+ )[None]
634
+ upper_bounds = torch.full(
635
+ template_mask.shape, float("inf"), device=index.device, dtype=torch.float32
636
+ )
637
+ ref_idxs = torch.argwhere(template_mask).T
638
+ upper_bounds[ref_idxs.unbind()] = feats["template_force_threshold"][
639
+ feats["template_force"]
640
+ ][ref_idxs[0]]
641
+
642
+ lower_bounds = None
643
+ k = torch.ones_like(upper_bounds)
644
+ return (
645
+ index,
646
+ (k, lower_bounds, upper_bounds),
647
+ None,
648
+ (ref_coords, ref_mask, ref_atom_index, ref_token_index),
649
+ None,
650
+ )
651
+
652
+
653
+ class ContactPotentital(FlatBottomPotential, DistancePotential):
654
+ def compute_args(self, feats, parameters):
655
+ index = feats["contact_pair_index"][0]
656
+ union_index = feats["contact_union_index"][0]
657
+ negation_mask = feats["contact_negation_mask"][0]
658
+ lower_bounds = None
659
+ upper_bounds = feats["contact_thresholds"][0].clone()
660
+ k = torch.ones_like(upper_bounds)
661
+ return (
662
+ index,
663
+ (k, lower_bounds, upper_bounds),
664
+ None,
665
+ None,
666
+ (negation_mask, union_index),
667
+ )
668
+
669
+
670
+ def get_potentials(steering_args, boltz2=False):
671
+ potentials = []
672
+ if steering_args["fk_steering"] or steering_args["physical_guidance_update"]:
673
+ potentials.extend(
674
+ [
675
+ SymmetricChainCOMPotential(
676
+ parameters={
677
+ "guidance_interval": 4,
678
+ "guidance_weight": 0.5
679
+ if steering_args["physical_guidance_update"]
680
+ else 0.0,
681
+ "resampling_weight": 0.5,
682
+ "buffer": ExponentialInterpolation(
683
+ start=1.0, end=5.0, alpha=-2.0
684
+ ),
685
+ }
686
+ ),
687
+ VDWOverlapPotential(
688
+ parameters={
689
+ "guidance_interval": 5,
690
+ "guidance_weight": (
691
+ PiecewiseStepFunction(thresholds=[0.4], values=[0.125, 0.0])
692
+ if steering_args["physical_guidance_update"]
693
+ else 0.0
694
+ ),
695
+ "resampling_weight": PiecewiseStepFunction(
696
+ thresholds=[0.6], values=[0.01, 0.0]
697
+ ),
698
+ "buffer": 0.225,
699
+ }
700
+ ),
701
+ ConnectionsPotential(
702
+ parameters={
703
+ "guidance_interval": 1,
704
+ "guidance_weight": 0.15
705
+ if steering_args["physical_guidance_update"]
706
+ else 0.0,
707
+ "resampling_weight": 1.0,
708
+ "buffer": 2.0,
709
+ }
710
+ ),
711
+ PoseBustersPotential(
712
+ parameters={
713
+ "guidance_interval": 1,
714
+ "guidance_weight": 0.01
715
+ if steering_args["physical_guidance_update"]
716
+ else 0.0,
717
+ "resampling_weight": 0.1,
718
+ "bond_buffer": 0.125,
719
+ "angle_buffer": 0.125,
720
+ "clash_buffer": 0.10,
721
+ }
722
+ ),
723
+ ChiralAtomPotential(
724
+ parameters={
725
+ "guidance_interval": 1,
726
+ "guidance_weight": 0.1
727
+ if steering_args["physical_guidance_update"]
728
+ else 0.0,
729
+ "resampling_weight": 1.0,
730
+ "buffer": 0.52360,
731
+ }
732
+ ),
733
+ StereoBondPotential(
734
+ parameters={
735
+ "guidance_interval": 1,
736
+ "guidance_weight": 0.05
737
+ if steering_args["physical_guidance_update"]
738
+ else 0.0,
739
+ "resampling_weight": 1.0,
740
+ "buffer": 0.52360,
741
+ }
742
+ ),
743
+ PlanarBondPotential(
744
+ parameters={
745
+ "guidance_interval": 1,
746
+ "guidance_weight": 0.05
747
+ if steering_args["physical_guidance_update"]
748
+ else 0.0,
749
+ "resampling_weight": 1.0,
750
+ "buffer": 0.26180,
751
+ }
752
+ ),
753
+ ]
754
+ )
755
+ if boltz2 and (
756
+ steering_args["fk_steering"] or steering_args["contact_guidance_update"]
757
+ ):
758
+ potentials.extend(
759
+ [
760
+ ContactPotentital(
761
+ parameters={
762
+ "guidance_interval": 4,
763
+ "guidance_weight": (
764
+ PiecewiseStepFunction(
765
+ thresholds=[0.25, 0.75], values=[0.0, 0.5, 1.0]
766
+ )
767
+ if steering_args["contact_guidance_update"]
768
+ else 0.0
769
+ ),
770
+ "resampling_weight": 1.0,
771
+ "union_lambda": ExponentialInterpolation(
772
+ start=8.0, end=0.0, alpha=-2.0
773
+ ),
774
+ }
775
+ ),
776
+ TemplateReferencePotential(
777
+ parameters={
778
+ "guidance_interval": 2,
779
+ "guidance_weight": 0.1
780
+ if steering_args["contact_guidance_update"]
781
+ else 0.0,
782
+ "resampling_weight": 1.0,
783
+ }
784
+ ),
785
+ ]
786
+ )
787
+ return potentials
vb_potentials_schedules.py CHANGED
@@ -1,37 +1,37 @@
1
- import math
2
- from abc import ABC
3
-
4
-
5
- class ParameterSchedule(ABC):
6
- def compute(self, t):
7
- raise NotImplementedError
8
-
9
-
10
- class ExponentialInterpolation(ParameterSchedule):
11
- def __init__(self, start, end, alpha):
12
- self.start = start
13
- self.end = end
14
- self.alpha = alpha
15
-
16
- def compute(self, t):
17
- if self.alpha != 0:
18
- return self.start + (self.end - self.start) * (
19
- math.exp(self.alpha * t) - 1
20
- ) / (math.exp(self.alpha) - 1)
21
- else:
22
- return self.start + (self.end - self.start) * t
23
-
24
-
25
- class PiecewiseStepFunction(ParameterSchedule):
26
- def __init__(self, thresholds, values):
27
- self.thresholds = thresholds
28
- self.values = values
29
-
30
- def compute(self, t):
31
- assert len(self.thresholds) > 0
32
- assert len(self.values) == len(self.thresholds) + 1
33
-
34
- idx = 0
35
- while idx < len(self.thresholds) and t > self.thresholds[idx]:
36
- idx += 1
37
- return self.values[idx]
 
1
+ import math
2
+ from abc import ABC
3
+
4
+
5
+ class ParameterSchedule(ABC):
6
+ def compute(self, t):
7
+ raise NotImplementedError
8
+
9
+
10
+ class ExponentialInterpolation(ParameterSchedule):
11
+ def __init__(self, start, end, alpha):
12
+ self.start = start
13
+ self.end = end
14
+ self.alpha = alpha
15
+
16
+ def compute(self, t):
17
+ if self.alpha != 0:
18
+ return self.start + (self.end - self.start) * (
19
+ math.exp(self.alpha * t) - 1
20
+ ) / (math.exp(self.alpha) - 1)
21
+ else:
22
+ return self.start + (self.end - self.start) * t
23
+
24
+
25
+ class PiecewiseStepFunction(ParameterSchedule):
26
+ def __init__(self, thresholds, values):
27
+ self.thresholds = thresholds
28
+ self.values = values
29
+
30
+ def compute(self, t):
31
+ assert len(self.thresholds) > 0
32
+ assert len(self.values) == len(self.thresholds) + 1
33
+
34
+ idx = 0
35
+ while idx < len(self.thresholds) and t > self.thresholds[idx]:
36
+ idx += 1
37
+ return self.values[idx]
vb_tri_attn_attention.py CHANGED
@@ -1,189 +1,189 @@
1
- # Copyright 2021 AlQuraishi Laboratory
2
- # Copyright 2021 DeepMind Technologies Limited
3
- #
4
- # Licensed under the Apache License, Version 2.0 (the "License");
5
- # you may not use this file except in compliance with the License.
6
- # You may obtain a copy of the License at
7
- #
8
- # http://www.apache.org/licenses/LICENSE-2.0
9
- #
10
- # Unless required by applicable law or agreed to in writing, software
11
- # distributed under the License is distributed on an "AS IS" BASIS,
12
- # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
- # See the License for the specific language governing permissions and
14
- # limitations under the License.
15
-
16
- from functools import partial, partialmethod
17
- from typing import Optional
18
-
19
- import torch
20
- import torch.nn as nn
21
-
22
- from .vb_tri_attn_primitives import (
23
- Attention,
24
- LayerNorm,
25
- Linear,
26
- )
27
- from .vb_tri_attn_utils import (
28
- chunk_layer,
29
- permute_final_dims,
30
- )
31
-
32
-
33
- class TriangleAttention(nn.Module):
34
- """Implement Algorithm 12."""
35
-
36
- def __init__(
37
- self,
38
- c_in: int,
39
- c_hidden: int,
40
- no_heads: int,
41
- starting: bool = True,
42
- inf: float = 1e9,
43
- ) -> None:
44
- super().__init__()
45
-
46
- self.c_in = c_in
47
- self.c_hidden = c_hidden
48
- self.no_heads = no_heads
49
- self.starting = starting
50
- self.inf = inf
51
-
52
- self.layer_norm = LayerNorm(self.c_in)
53
-
54
- self.linear = Linear(c_in, self.no_heads, bias=False, init="normal")
55
-
56
- self.mha = Attention(
57
- self.c_in, self.c_in, self.c_in, self.c_hidden, self.no_heads
58
- )
59
-
60
- @torch.jit.ignore
61
- def _chunk(
62
- self,
63
- x: torch.Tensor,
64
- tri_bias: torch.Tensor,
65
- mask_bias: torch.Tensor,
66
- mask: torch.Tensor,
67
- chunk_size: int,
68
- use_kernels: bool = False,
69
- ) -> torch.Tensor:
70
- """Compute triangle attention.
71
-
72
- Parameters
73
- ----------
74
- x : torch.Tensor
75
- Input tensor of shape [*, I, J, C_in]
76
- biases : list[torch.Tensor]
77
- List of bias tensors of shape [*, H, I, J]
78
- chunk_size : int
79
- Size of chunks for memory efficient computation
80
- use_kernels : bool, default=False
81
- Whether to use optimized CUDA kernels
82
-
83
- Returns
84
- -------
85
- torch.Tensor
86
- Output tensor of shape [*, I, J, C_in]
87
-
88
- """
89
- mha_inputs = {
90
- "q_x": x,
91
- "kv_x": x,
92
- "tri_bias": tri_bias,
93
- "mask_bias": mask_bias,
94
- "mask": mask,
95
- }
96
-
97
- return chunk_layer(
98
- partial(
99
- self.mha,
100
- use_kernels=use_kernels,
101
- ),
102
- mha_inputs,
103
- chunk_size=chunk_size,
104
- no_batch_dims=len(x.shape[:-2]),
105
- _out=None,
106
- )
107
-
108
- def forward(
109
- self,
110
- x: torch.Tensor,
111
- mask: Optional[torch.Tensor] = None,
112
- chunk_size: Optional[int] = None,
113
- use_kernels: bool = False,
114
- ) -> torch.Tensor:
115
- """Compute triangle attention.
116
-
117
- Parameters
118
- ----------
119
- x : torch.Tensor
120
- Input tensor of shape [*, I, J, C_in]
121
- mask : torch.Tensor, optional
122
- Attention mask of shape [*, I, J]
123
- chunk_size : int, optional
124
- Size of chunks for memory efficient computation
125
- use_kernels : bool, default=False
126
- Whether to use optimized CUDA kernels
127
-
128
- Returns
129
- -------
130
- torch.Tensor
131
- Output tensor of shape [*, I, J, C_in]
132
-
133
- """
134
- if mask is None:
135
- # [*, I, J]
136
- mask = x.new_ones(
137
- x.shape[:-1],
138
- )
139
-
140
- if not self.starting:
141
- x = x.transpose(-2, -3)
142
- mask = mask.transpose(-1, -2)
143
-
144
- # [*, I, J, C_in]
145
- x = self.layer_norm(x)
146
-
147
- # [*, I, 1, 1, J]
148
- mask = mask[..., :, None, None, :]
149
- mask_bias = self.inf * (mask - 1)
150
-
151
- # [*, H, I, J]
152
- triangle_bias = permute_final_dims(self.linear(x), (2, 0, 1))
153
-
154
- # [*, 1, H, I, J]
155
- triangle_bias = triangle_bias.unsqueeze(-4)
156
-
157
- if chunk_size is not None and not use_kernels:
158
- x = self._chunk(
159
- x,
160
- triangle_bias,
161
- mask_bias,
162
- mask,
163
- chunk_size,
164
- use_kernels=use_kernels,
165
- )
166
- else:
167
- x = self.mha(
168
- x,
169
- x,
170
- triangle_bias,
171
- mask_bias,
172
- mask,
173
- use_kernels=use_kernels,
174
- )
175
-
176
- if not self.starting:
177
- x = x.transpose(-2, -3)
178
-
179
- return x
180
-
181
-
182
- # Implements Algorithm 13
183
- TriangleAttentionStartingNode = TriangleAttention
184
-
185
-
186
- class TriangleAttentionEndingNode(TriangleAttention):
187
- """Implement Algorithm 14."""
188
-
189
- __init__ = partialmethod(TriangleAttention.__init__, starting=False)
 
1
+ # Copyright 2021 AlQuraishi Laboratory
2
+ # Copyright 2021 DeepMind Technologies Limited
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+
16
+ from functools import partial, partialmethod
17
+ from typing import Optional
18
+
19
+ import torch
20
+ import torch.nn as nn
21
+
22
+ from .vb_tri_attn_primitives import (
23
+ Attention,
24
+ LayerNorm,
25
+ Linear,
26
+ )
27
+ from .vb_tri_attn_utils import (
28
+ chunk_layer,
29
+ permute_final_dims,
30
+ )
31
+
32
+
33
+ class TriangleAttention(nn.Module):
34
+ """Implement Algorithm 12."""
35
+
36
+ def __init__(
37
+ self,
38
+ c_in: int,
39
+ c_hidden: int,
40
+ no_heads: int,
41
+ starting: bool = True,
42
+ inf: float = 1e9,
43
+ ) -> None:
44
+ super().__init__()
45
+
46
+ self.c_in = c_in
47
+ self.c_hidden = c_hidden
48
+ self.no_heads = no_heads
49
+ self.starting = starting
50
+ self.inf = inf
51
+
52
+ self.layer_norm = LayerNorm(self.c_in)
53
+
54
+ self.linear = Linear(c_in, self.no_heads, bias=False, init="normal")
55
+
56
+ self.mha = Attention(
57
+ self.c_in, self.c_in, self.c_in, self.c_hidden, self.no_heads
58
+ )
59
+
60
+ @torch.jit.ignore
61
+ def _chunk(
62
+ self,
63
+ x: torch.Tensor,
64
+ tri_bias: torch.Tensor,
65
+ mask_bias: torch.Tensor,
66
+ mask: torch.Tensor,
67
+ chunk_size: int,
68
+ use_kernels: bool = False,
69
+ ) -> torch.Tensor:
70
+ """Compute triangle attention.
71
+
72
+ Parameters
73
+ ----------
74
+ x : torch.Tensor
75
+ Input tensor of shape [*, I, J, C_in]
76
+ biases : list[torch.Tensor]
77
+ List of bias tensors of shape [*, H, I, J]
78
+ chunk_size : int
79
+ Size of chunks for memory efficient computation
80
+ use_kernels : bool, default=False
81
+ Whether to use optimized CUDA kernels
82
+
83
+ Returns
84
+ -------
85
+ torch.Tensor
86
+ Output tensor of shape [*, I, J, C_in]
87
+
88
+ """
89
+ mha_inputs = {
90
+ "q_x": x,
91
+ "kv_x": x,
92
+ "tri_bias": tri_bias,
93
+ "mask_bias": mask_bias,
94
+ "mask": mask,
95
+ }
96
+
97
+ return chunk_layer(
98
+ partial(
99
+ self.mha,
100
+ use_kernels=use_kernels,
101
+ ),
102
+ mha_inputs,
103
+ chunk_size=chunk_size,
104
+ no_batch_dims=len(x.shape[:-2]),
105
+ _out=None,
106
+ )
107
+
108
+ def forward(
109
+ self,
110
+ x: torch.Tensor,
111
+ mask: Optional[torch.Tensor] = None,
112
+ chunk_size: Optional[int] = None,
113
+ use_kernels: bool = False,
114
+ ) -> torch.Tensor:
115
+ """Compute triangle attention.
116
+
117
+ Parameters
118
+ ----------
119
+ x : torch.Tensor
120
+ Input tensor of shape [*, I, J, C_in]
121
+ mask : torch.Tensor, optional
122
+ Attention mask of shape [*, I, J]
123
+ chunk_size : int, optional
124
+ Size of chunks for memory efficient computation
125
+ use_kernels : bool, default=False
126
+ Whether to use optimized CUDA kernels
127
+
128
+ Returns
129
+ -------
130
+ torch.Tensor
131
+ Output tensor of shape [*, I, J, C_in]
132
+
133
+ """
134
+ if mask is None:
135
+ # [*, I, J]
136
+ mask = x.new_ones(
137
+ x.shape[:-1],
138
+ )
139
+
140
+ if not self.starting:
141
+ x = x.transpose(-2, -3)
142
+ mask = mask.transpose(-1, -2)
143
+
144
+ # [*, I, J, C_in]
145
+ x = self.layer_norm(x)
146
+
147
+ # [*, I, 1, 1, J]
148
+ mask = mask[..., :, None, None, :]
149
+ mask_bias = self.inf * (mask - 1)
150
+
151
+ # [*, H, I, J]
152
+ triangle_bias = permute_final_dims(self.linear(x), (2, 0, 1))
153
+
154
+ # [*, 1, H, I, J]
155
+ triangle_bias = triangle_bias.unsqueeze(-4)
156
+
157
+ if chunk_size is not None and not use_kernels:
158
+ x = self._chunk(
159
+ x,
160
+ triangle_bias,
161
+ mask_bias,
162
+ mask,
163
+ chunk_size,
164
+ use_kernels=use_kernels,
165
+ )
166
+ else:
167
+ x = self.mha(
168
+ x,
169
+ x,
170
+ triangle_bias,
171
+ mask_bias,
172
+ mask,
173
+ use_kernels=use_kernels,
174
+ )
175
+
176
+ if not self.starting:
177
+ x = x.transpose(-2, -3)
178
+
179
+ return x
180
+
181
+
182
+ # Implements Algorithm 13
183
+ TriangleAttentionStartingNode = TriangleAttention
184
+
185
+
186
+ class TriangleAttentionEndingNode(TriangleAttention):
187
+ """Implement Algorithm 14."""
188
+
189
+ __init__ = partialmethod(TriangleAttention.__init__, starting=False)
vb_tri_attn_primitives.py CHANGED
@@ -1,411 +1,411 @@
1
- # Copyright 2021 AlQuraishi Laboratory
2
- # Copyright 2021 DeepMind Technologies Limited
3
- #
4
- # Licensed under the Apache License, Version 2.0 (the "License");
5
- # you may not use this file except in compliance with the License.
6
- # You may obtain a copy of the License at
7
- #
8
- # http://www.apache.org/licenses/LICENSE-2.0
9
- #
10
- # Unless required by applicable law or agreed to in writing, software
11
- # distributed under the License is distributed on an "AS IS" BASIS,
12
- # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
- # See the License for the specific language governing permissions and
14
- # limitations under the License.
15
-
16
- import math
17
- import importlib
18
- from typing import Callable, List, Optional, Tuple
19
-
20
- import torch
21
- from einops import rearrange
22
- from torch import nn
23
-
24
- from . import vb_layers_initialize as initialize
25
- from .vb_tri_attn_utils import (
26
- flatten_final_dims,
27
- permute_final_dims,
28
- )
29
-
30
-
31
- class Linear(nn.Linear):
32
- """
33
- A Linear layer with built-in nonstandard initializations. Called just
34
- like torch.nn.Linear.
35
-
36
- Implements the initializers in 1.11.4, plus some additional ones found
37
- in the code.
38
- """
39
-
40
- def __init__(
41
- self,
42
- in_dim: int,
43
- out_dim: int,
44
- bias: bool = True,
45
- init: str = "default",
46
- init_fn: Optional[Callable[[torch.Tensor, torch.Tensor], None]] = None,
47
- precision=None,
48
- ):
49
- """Initialize the linear layer.
50
-
51
- Parameters
52
- ----------
53
- in_dim : int
54
- The final dimension of inputs to the layer
55
- out_dim : int
56
- The final dimension of layer outputs
57
- bias : bool, default=True
58
- Whether to learn an additive bias
59
- init : str, default='default'
60
- The initializer to use. Choose from:
61
-
62
- - "default": LeCun fan-in truncated normal initialization
63
- - "relu": He initialization w/ truncated normal distribution
64
- - "glorot": Fan-average Glorot uniform initialization
65
- - "gating": Weights=0, Bias=1
66
- - "normal": Normal initialization with std=1/sqrt(fan_in)
67
- - "final": Weights=0, Bias=0
68
-
69
- Overridden by init_fn if the latter is not None.
70
- init_fn : callable, optional
71
- A custom initializer taking weight and bias as inputs.
72
- Overrides init if not None.
73
-
74
- """
75
- super().__init__(in_dim, out_dim, bias=bias)
76
-
77
- if bias:
78
- with torch.no_grad():
79
- self.bias.fill_(0)
80
-
81
- with torch.no_grad():
82
- if init_fn is not None:
83
- init_fn(self.weight, self.bias)
84
- else:
85
- if init == "default":
86
- initialize.lecun_normal_init_(self.weight)
87
- elif init == "relu":
88
- initialize.he_normal_init_(self.weight)
89
- elif init == "glorot":
90
- initialize.glorot_uniform_init_(self.weight)
91
- elif init == "gating":
92
- initialize.gating_init_(self.weight)
93
- if bias:
94
- self.bias.fill_(1.0)
95
- elif init == "normal":
96
- initialize.normal_init_(self.weight)
97
- elif init == "final":
98
- initialize.final_init_(self.weight)
99
- else:
100
- raise ValueError("Invalid init string.")
101
-
102
- self.precision = precision
103
-
104
- def forward(self, input: torch.Tensor) -> torch.Tensor:
105
- d = input.dtype
106
- if self.precision is not None:
107
- with torch.autocast("cuda", enabled=False):
108
- bias = (
109
- self.bias.to(dtype=self.precision)
110
- if self.bias is not None
111
- else None
112
- )
113
- return nn.functional.linear(
114
- input.to(dtype=self.precision),
115
- self.weight.to(dtype=self.precision),
116
- bias,
117
- ).to(dtype=d)
118
-
119
- if d is torch.bfloat16:
120
- with torch.autocast("cuda", enabled=False):
121
- bias = self.bias.to(dtype=d) if self.bias is not None else None
122
- return nn.functional.linear(input, self.weight.to(dtype=d), bias)
123
-
124
- return nn.functional.linear(input, self.weight, self.bias)
125
-
126
-
127
- class LayerNorm(nn.Module):
128
- def __init__(self, c_in, eps=1e-5):
129
- super(LayerNorm, self).__init__()
130
-
131
- self.c_in = (c_in,)
132
- self.eps = eps
133
-
134
- self.weight = nn.Parameter(torch.ones(c_in))
135
- self.bias = nn.Parameter(torch.zeros(c_in))
136
-
137
- def forward(self, x):
138
- d = x.dtype
139
- if d is torch.bfloat16:
140
- with torch.autocast("cuda", enabled=False):
141
- out = nn.functional.layer_norm(
142
- x,
143
- self.c_in,
144
- self.weight.to(dtype=d),
145
- self.bias.to(dtype=d),
146
- self.eps,
147
- )
148
- else:
149
- out = nn.functional.layer_norm(
150
- x,
151
- self.c_in,
152
- self.weight,
153
- self.bias,
154
- self.eps,
155
- )
156
-
157
- return out
158
-
159
-
160
- @torch.jit.ignore
161
- def softmax_no_cast(t: torch.Tensor, dim: int = -1) -> torch.Tensor:
162
- """
163
- Softmax, but without automatic casting to fp32 when the input is of
164
- type bfloat16
165
- """
166
- d = t.dtype
167
- if d is torch.bfloat16:
168
- with torch.autocast("cuda", enabled=False):
169
- s = torch.nn.functional.softmax(t, dim=dim)
170
- else:
171
- s = torch.nn.functional.softmax(t, dim=dim)
172
-
173
- return s
174
-
175
-
176
- # @torch.jit.script
177
- def _attention(
178
- query: torch.Tensor,
179
- key: torch.Tensor,
180
- value: torch.Tensor,
181
- biases: List[torch.Tensor],
182
- ) -> torch.Tensor:
183
- # [*, H, C_hidden, K]
184
- key = permute_final_dims(key, (1, 0))
185
-
186
- # [*, H, Q, K]
187
- a = torch.matmul(query, key)
188
-
189
- for b in biases:
190
- a += b
191
-
192
- a = softmax_no_cast(a, -1)
193
-
194
- # [*, H, Q, C_hidden]
195
- a = torch.matmul(a, value)
196
-
197
- return a
198
-
199
-
200
- @torch.compiler.disable
201
- def kernel_triangular_attn(q, k, v, tri_bias, mask, scale):
202
- triangle_module = importlib.import_module("cuequivariance_torch.primitives.triangle")
203
- triangle_attention = triangle_module.triangle_attention
204
- return triangle_attention(q, k, v, tri_bias, mask=mask, scale=scale)
205
-
206
-
207
- class Attention(nn.Module):
208
- """
209
- Standard multi-head attention using AlphaFold's default layer
210
- initialization. Allows multiple bias vectors.
211
- """
212
-
213
- def __init__(
214
- self,
215
- c_q: int,
216
- c_k: int,
217
- c_v: int,
218
- c_hidden: int,
219
- no_heads: int,
220
- gating: bool = True,
221
- ):
222
- """Initialize the attention layer.
223
-
224
- Parameters
225
- ----------
226
- c_q : int
227
- Input dimension of query data
228
- c_k : int
229
- Input dimension of key data
230
- c_v : int
231
- Input dimension of value data
232
- c_hidden : int
233
- Per-head hidden dimension
234
- no_heads : int
235
- Number of attention heads
236
- gating : bool, default=True
237
- Whether the output should be gated using query data
238
-
239
- """
240
- super().__init__()
241
-
242
- self.c_q = c_q
243
- self.c_k = c_k
244
- self.c_v = c_v
245
- self.c_hidden = c_hidden
246
- self.no_heads = no_heads
247
- self.gating = gating
248
-
249
- # DISCREPANCY: c_hidden is not the per-head channel dimension, as
250
- # stated in the supplement, but the overall channel dimension.
251
-
252
- self.linear_q = Linear(
253
- self.c_q, self.c_hidden * self.no_heads, bias=False, init="glorot"
254
- )
255
- self.linear_k = Linear(
256
- self.c_k, self.c_hidden * self.no_heads, bias=False, init="glorot"
257
- )
258
- self.linear_v = Linear(
259
- self.c_v, self.c_hidden * self.no_heads, bias=False, init="glorot"
260
- )
261
- self.linear_o = Linear(
262
- self.c_hidden * self.no_heads, self.c_q, bias=False, init="final"
263
- )
264
-
265
- self.linear_g = None
266
- if self.gating:
267
- self.linear_g = Linear(
268
- self.c_q, self.c_hidden * self.no_heads, bias=False, init="gating"
269
- )
270
-
271
- self.sigmoid = nn.Sigmoid()
272
-
273
- def _prep_qkv(
274
- self, q_x: torch.Tensor, kv_x: torch.Tensor, apply_scale: bool = True
275
- ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
276
- # [*, Q/K/V, H * C_hidden]
277
- q = self.linear_q(q_x)
278
- k = self.linear_k(kv_x)
279
- v = self.linear_v(kv_x)
280
-
281
- # [*, Q/K, H, C_hidden]
282
- q = q.view(q.shape[:-1] + (self.no_heads, -1))
283
- k = k.view(k.shape[:-1] + (self.no_heads, -1))
284
- v = v.view(v.shape[:-1] + (self.no_heads, -1))
285
-
286
- # [*, H, Q/K, C_hidden]
287
- q = q.transpose(-2, -3)
288
- k = k.transpose(-2, -3)
289
- v = v.transpose(-2, -3)
290
-
291
- if apply_scale:
292
- q /= math.sqrt(self.c_hidden)
293
-
294
- return q, k, v
295
-
296
- def _wrap_up(self, o: torch.Tensor, q_x: torch.Tensor) -> torch.Tensor:
297
- if self.linear_g is not None:
298
- g = self.sigmoid(self.linear_g(q_x))
299
-
300
- # [*, Q, H, C_hidden]
301
- g = g.view(g.shape[:-1] + (self.no_heads, -1))
302
- o = o * g
303
-
304
- # [*, Q, H * C_hidden]
305
- o = flatten_final_dims(o, 2)
306
-
307
- # [*, Q, C_q]
308
- o = self.linear_o(o)
309
-
310
- return o
311
-
312
- def forward(
313
- self,
314
- q_x: torch.Tensor,
315
- kv_x: torch.Tensor,
316
- tri_bias: torch.Tensor,
317
- mask_bias: torch.Tensor,
318
- mask: torch.Tensor,
319
- use_kernels: bool = False,
320
- ) -> torch.Tensor:
321
- """Compute attention.
322
-
323
- Parameters
324
- ----------
325
- q_x : torch.Tensor
326
- [*, Q, C_q] query data
327
- kv_x : torch.Tensor
328
- [*, K, C_k] key data
329
- tri_bias : torch.Tensor
330
- [*, H, Q, K] triangular bias
331
- mask_bias : torch.Tensor
332
- [*, H, Q, K] mask bias
333
- mask : torch.Tensor
334
- [*, Q, K] mask
335
- use_kernels : bool, default=False
336
- Whether to use optimized CUDA kernels
337
-
338
- Returns
339
- -------
340
- [*, Q, C_q] attention update
341
-
342
- """
343
- # Attention kernel applies scaling internally
344
- q, k, v = self._prep_qkv(
345
- q_x,
346
- kv_x,
347
- apply_scale=not use_kernels,
348
- )
349
-
350
- if use_kernels:
351
- scale = 1.0 / math.sqrt(self.c_hidden)
352
- o = kernel_triangular_attn(
353
- q,
354
- k,
355
- v,
356
- tri_bias=tri_bias,
357
- mask=mask.bool(),
358
- scale=scale,
359
- )
360
- o = o.transpose(-2, -3)
361
- else:
362
- biases = [mask_bias, tri_bias]
363
- o = _attention(q, k, v, biases)
364
- o = o.transpose(-2, -3)
365
-
366
- o = self._wrap_up(o, q_x)
367
-
368
- return o
369
-
370
-
371
- def _trifast_attn(q, k, v, biases):
372
- orig_n_dims = len(q.shape)
373
-
374
- if len(biases) != 2:
375
- raise ValueError(f"Trifast expects two bias terms, found {len(biases)}")
376
-
377
- mask, b = biases
378
-
379
- if len(b.shape) == 5:
380
- # Sometimes there is an extra batch dim -- why?
381
- b = b.squeeze(1)
382
-
383
- if orig_n_dims == 4:
384
- # add fake batch dim
385
- q = q.unsqueeze(0)
386
- k = k.unsqueeze(0)
387
- v = v.unsqueeze(0)
388
- # b = b.unsqueeze(0) not sure why this and only this has a batch dim?
389
- mask = mask.unsqueeze(0)
390
-
391
- if len(q.shape) != 5:
392
- raise ValueError(f"Trifast expects q/k/v to be 5D, found {len(q.shape)}")
393
-
394
- # Reorder q/k/v
395
- q = rearrange(q, "b i h j d -> b h i j d")
396
- k = rearrange(k, "b i h j d -> b h i j d")
397
- v = rearrange(v, "b i h j d -> b h i j d")
398
-
399
- # Make mask the right shape.
400
- mask = rearrange(mask, "b i () () j -> b i j").bool()
401
-
402
- # Delay import to here to avoid initializing cuda too early
403
- from trifast import triangle_attention
404
-
405
- o = triangle_attention(q, k, v, b, mask)
406
- o = rearrange(o, "b h i j d -> b i j h d")
407
-
408
- # Remove the batch dim if we added it.
409
- if orig_n_dims == 4:
410
- o = o.squeeze(0)
411
- return o
 
1
+ # Copyright 2021 AlQuraishi Laboratory
2
+ # Copyright 2021 DeepMind Technologies Limited
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+
16
+ import math
17
+ import importlib
18
+ from typing import Callable, List, Optional, Tuple
19
+
20
+ import torch
21
+ from einops import rearrange
22
+ from torch import nn
23
+
24
+ from . import vb_layers_initialize as initialize
25
+ from .vb_tri_attn_utils import (
26
+ flatten_final_dims,
27
+ permute_final_dims,
28
+ )
29
+
30
+
31
+ class Linear(nn.Linear):
32
+ """
33
+ A Linear layer with built-in nonstandard initializations. Called just
34
+ like torch.nn.Linear.
35
+
36
+ Implements the initializers in 1.11.4, plus some additional ones found
37
+ in the code.
38
+ """
39
+
40
+ def __init__(
41
+ self,
42
+ in_dim: int,
43
+ out_dim: int,
44
+ bias: bool = True,
45
+ init: str = "default",
46
+ init_fn: Optional[Callable[[torch.Tensor, torch.Tensor], None]] = None,
47
+ precision=None,
48
+ ):
49
+ """Initialize the linear layer.
50
+
51
+ Parameters
52
+ ----------
53
+ in_dim : int
54
+ The final dimension of inputs to the layer
55
+ out_dim : int
56
+ The final dimension of layer outputs
57
+ bias : bool, default=True
58
+ Whether to learn an additive bias
59
+ init : str, default='default'
60
+ The initializer to use. Choose from:
61
+
62
+ - "default": LeCun fan-in truncated normal initialization
63
+ - "relu": He initialization w/ truncated normal distribution
64
+ - "glorot": Fan-average Glorot uniform initialization
65
+ - "gating": Weights=0, Bias=1
66
+ - "normal": Normal initialization with std=1/sqrt(fan_in)
67
+ - "final": Weights=0, Bias=0
68
+
69
+ Overridden by init_fn if the latter is not None.
70
+ init_fn : callable, optional
71
+ A custom initializer taking weight and bias as inputs.
72
+ Overrides init if not None.
73
+
74
+ """
75
+ super().__init__(in_dim, out_dim, bias=bias)
76
+
77
+ if bias:
78
+ with torch.no_grad():
79
+ self.bias.fill_(0)
80
+
81
+ with torch.no_grad():
82
+ if init_fn is not None:
83
+ init_fn(self.weight, self.bias)
84
+ else:
85
+ if init == "default":
86
+ initialize.lecun_normal_init_(self.weight)
87
+ elif init == "relu":
88
+ initialize.he_normal_init_(self.weight)
89
+ elif init == "glorot":
90
+ initialize.glorot_uniform_init_(self.weight)
91
+ elif init == "gating":
92
+ initialize.gating_init_(self.weight)
93
+ if bias:
94
+ self.bias.fill_(1.0)
95
+ elif init == "normal":
96
+ initialize.normal_init_(self.weight)
97
+ elif init == "final":
98
+ initialize.final_init_(self.weight)
99
+ else:
100
+ raise ValueError("Invalid init string.")
101
+
102
+ self.precision = precision
103
+
104
+ def forward(self, input: torch.Tensor) -> torch.Tensor:
105
+ d = input.dtype
106
+ if self.precision is not None:
107
+ with torch.autocast("cuda", enabled=False):
108
+ bias = (
109
+ self.bias.to(dtype=self.precision)
110
+ if self.bias is not None
111
+ else None
112
+ )
113
+ return nn.functional.linear(
114
+ input.to(dtype=self.precision),
115
+ self.weight.to(dtype=self.precision),
116
+ bias,
117
+ ).to(dtype=d)
118
+
119
+ if d is torch.bfloat16:
120
+ with torch.autocast("cuda", enabled=False):
121
+ bias = self.bias.to(dtype=d) if self.bias is not None else None
122
+ return nn.functional.linear(input, self.weight.to(dtype=d), bias)
123
+
124
+ return nn.functional.linear(input, self.weight, self.bias)
125
+
126
+
127
+ class LayerNorm(nn.Module):
128
+ def __init__(self, c_in, eps=1e-5):
129
+ super(LayerNorm, self).__init__()
130
+
131
+ self.c_in = (c_in,)
132
+ self.eps = eps
133
+
134
+ self.weight = nn.Parameter(torch.ones(c_in))
135
+ self.bias = nn.Parameter(torch.zeros(c_in))
136
+
137
+ def forward(self, x):
138
+ d = x.dtype
139
+ if d is torch.bfloat16:
140
+ with torch.autocast("cuda", enabled=False):
141
+ out = nn.functional.layer_norm(
142
+ x,
143
+ self.c_in,
144
+ self.weight.to(dtype=d),
145
+ self.bias.to(dtype=d),
146
+ self.eps,
147
+ )
148
+ else:
149
+ out = nn.functional.layer_norm(
150
+ x,
151
+ self.c_in,
152
+ self.weight,
153
+ self.bias,
154
+ self.eps,
155
+ )
156
+
157
+ return out
158
+
159
+
160
+ @torch.jit.ignore
161
+ def softmax_no_cast(t: torch.Tensor, dim: int = -1) -> torch.Tensor:
162
+ """
163
+ Softmax, but without automatic casting to fp32 when the input is of
164
+ type bfloat16
165
+ """
166
+ d = t.dtype
167
+ if d is torch.bfloat16:
168
+ with torch.autocast("cuda", enabled=False):
169
+ s = torch.nn.functional.softmax(t, dim=dim)
170
+ else:
171
+ s = torch.nn.functional.softmax(t, dim=dim)
172
+
173
+ return s
174
+
175
+
176
+ # @torch.jit.script
177
+ def _attention(
178
+ query: torch.Tensor,
179
+ key: torch.Tensor,
180
+ value: torch.Tensor,
181
+ biases: List[torch.Tensor],
182
+ ) -> torch.Tensor:
183
+ # [*, H, C_hidden, K]
184
+ key = permute_final_dims(key, (1, 0))
185
+
186
+ # [*, H, Q, K]
187
+ a = torch.matmul(query, key)
188
+
189
+ for b in biases:
190
+ a += b
191
+
192
+ a = softmax_no_cast(a, -1)
193
+
194
+ # [*, H, Q, C_hidden]
195
+ a = torch.matmul(a, value)
196
+
197
+ return a
198
+
199
+
200
+ @torch.compiler.disable
201
+ def kernel_triangular_attn(q, k, v, tri_bias, mask, scale):
202
+ triangle_module = importlib.import_module("cuequivariance_torch.primitives.triangle")
203
+ triangle_attention = triangle_module.triangle_attention
204
+ return triangle_attention(q, k, v, tri_bias, mask=mask, scale=scale)
205
+
206
+
207
+ class Attention(nn.Module):
208
+ """
209
+ Standard multi-head attention using AlphaFold's default layer
210
+ initialization. Allows multiple bias vectors.
211
+ """
212
+
213
+ def __init__(
214
+ self,
215
+ c_q: int,
216
+ c_k: int,
217
+ c_v: int,
218
+ c_hidden: int,
219
+ no_heads: int,
220
+ gating: bool = True,
221
+ ):
222
+ """Initialize the attention layer.
223
+
224
+ Parameters
225
+ ----------
226
+ c_q : int
227
+ Input dimension of query data
228
+ c_k : int
229
+ Input dimension of key data
230
+ c_v : int
231
+ Input dimension of value data
232
+ c_hidden : int
233
+ Per-head hidden dimension
234
+ no_heads : int
235
+ Number of attention heads
236
+ gating : bool, default=True
237
+ Whether the output should be gated using query data
238
+
239
+ """
240
+ super().__init__()
241
+
242
+ self.c_q = c_q
243
+ self.c_k = c_k
244
+ self.c_v = c_v
245
+ self.c_hidden = c_hidden
246
+ self.no_heads = no_heads
247
+ self.gating = gating
248
+
249
+ # DISCREPANCY: c_hidden is not the per-head channel dimension, as
250
+ # stated in the supplement, but the overall channel dimension.
251
+
252
+ self.linear_q = Linear(
253
+ self.c_q, self.c_hidden * self.no_heads, bias=False, init="glorot"
254
+ )
255
+ self.linear_k = Linear(
256
+ self.c_k, self.c_hidden * self.no_heads, bias=False, init="glorot"
257
+ )
258
+ self.linear_v = Linear(
259
+ self.c_v, self.c_hidden * self.no_heads, bias=False, init="glorot"
260
+ )
261
+ self.linear_o = Linear(
262
+ self.c_hidden * self.no_heads, self.c_q, bias=False, init="final"
263
+ )
264
+
265
+ self.linear_g = None
266
+ if self.gating:
267
+ self.linear_g = Linear(
268
+ self.c_q, self.c_hidden * self.no_heads, bias=False, init="gating"
269
+ )
270
+
271
+ self.sigmoid = nn.Sigmoid()
272
+
273
+ def _prep_qkv(
274
+ self, q_x: torch.Tensor, kv_x: torch.Tensor, apply_scale: bool = True
275
+ ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
276
+ # [*, Q/K/V, H * C_hidden]
277
+ q = self.linear_q(q_x)
278
+ k = self.linear_k(kv_x)
279
+ v = self.linear_v(kv_x)
280
+
281
+ # [*, Q/K, H, C_hidden]
282
+ q = q.view(q.shape[:-1] + (self.no_heads, -1))
283
+ k = k.view(k.shape[:-1] + (self.no_heads, -1))
284
+ v = v.view(v.shape[:-1] + (self.no_heads, -1))
285
+
286
+ # [*, H, Q/K, C_hidden]
287
+ q = q.transpose(-2, -3)
288
+ k = k.transpose(-2, -3)
289
+ v = v.transpose(-2, -3)
290
+
291
+ if apply_scale:
292
+ q /= math.sqrt(self.c_hidden)
293
+
294
+ return q, k, v
295
+
296
+ def _wrap_up(self, o: torch.Tensor, q_x: torch.Tensor) -> torch.Tensor:
297
+ if self.linear_g is not None:
298
+ g = self.sigmoid(self.linear_g(q_x))
299
+
300
+ # [*, Q, H, C_hidden]
301
+ g = g.view(g.shape[:-1] + (self.no_heads, -1))
302
+ o = o * g
303
+
304
+ # [*, Q, H * C_hidden]
305
+ o = flatten_final_dims(o, 2)
306
+
307
+ # [*, Q, C_q]
308
+ o = self.linear_o(o)
309
+
310
+ return o
311
+
312
+ def forward(
313
+ self,
314
+ q_x: torch.Tensor,
315
+ kv_x: torch.Tensor,
316
+ tri_bias: torch.Tensor,
317
+ mask_bias: torch.Tensor,
318
+ mask: torch.Tensor,
319
+ use_kernels: bool = False,
320
+ ) -> torch.Tensor:
321
+ """Compute attention.
322
+
323
+ Parameters
324
+ ----------
325
+ q_x : torch.Tensor
326
+ [*, Q, C_q] query data
327
+ kv_x : torch.Tensor
328
+ [*, K, C_k] key data
329
+ tri_bias : torch.Tensor
330
+ [*, H, Q, K] triangular bias
331
+ mask_bias : torch.Tensor
332
+ [*, H, Q, K] mask bias
333
+ mask : torch.Tensor
334
+ [*, Q, K] mask
335
+ use_kernels : bool, default=False
336
+ Whether to use optimized CUDA kernels
337
+
338
+ Returns
339
+ -------
340
+ [*, Q, C_q] attention update
341
+
342
+ """
343
+ # Attention kernel applies scaling internally
344
+ q, k, v = self._prep_qkv(
345
+ q_x,
346
+ kv_x,
347
+ apply_scale=not use_kernels,
348
+ )
349
+
350
+ if use_kernels:
351
+ scale = 1.0 / math.sqrt(self.c_hidden)
352
+ o = kernel_triangular_attn(
353
+ q,
354
+ k,
355
+ v,
356
+ tri_bias=tri_bias,
357
+ mask=mask.bool(),
358
+ scale=scale,
359
+ )
360
+ o = o.transpose(-2, -3)
361
+ else:
362
+ biases = [mask_bias, tri_bias]
363
+ o = _attention(q, k, v, biases)
364
+ o = o.transpose(-2, -3)
365
+
366
+ o = self._wrap_up(o, q_x)
367
+
368
+ return o
369
+
370
+
371
+ def _trifast_attn(q, k, v, biases):
372
+ orig_n_dims = len(q.shape)
373
+
374
+ if len(biases) != 2:
375
+ raise ValueError(f"Trifast expects two bias terms, found {len(biases)}")
376
+
377
+ mask, b = biases
378
+
379
+ if len(b.shape) == 5:
380
+ # Sometimes there is an extra batch dim -- why?
381
+ b = b.squeeze(1)
382
+
383
+ if orig_n_dims == 4:
384
+ # add fake batch dim
385
+ q = q.unsqueeze(0)
386
+ k = k.unsqueeze(0)
387
+ v = v.unsqueeze(0)
388
+ # b = b.unsqueeze(0) not sure why this and only this has a batch dim?
389
+ mask = mask.unsqueeze(0)
390
+
391
+ if len(q.shape) != 5:
392
+ raise ValueError(f"Trifast expects q/k/v to be 5D, found {len(q.shape)}")
393
+
394
+ # Reorder q/k/v
395
+ q = rearrange(q, "b i h j d -> b h i j d")
396
+ k = rearrange(k, "b i h j d -> b h i j d")
397
+ v = rearrange(v, "b i h j d -> b h i j d")
398
+
399
+ # Make mask the right shape.
400
+ mask = rearrange(mask, "b i () () j -> b i j").bool()
401
+
402
+ # Delay import to here to avoid initializing cuda too early
403
+ from trifast import triangle_attention
404
+
405
+ o = triangle_attention(q, k, v, b, mask)
406
+ o = rearrange(o, "b h i j d -> b i j h d")
407
+
408
+ # Remove the batch dim if we added it.
409
+ if orig_n_dims == 4:
410
+ o = o.squeeze(0)
411
+ return o
vb_tri_attn_utils.py CHANGED
@@ -1,380 +1,380 @@
1
- # Copyright 2021 AlQuraishi Laboratory
2
- #
3
- # Licensed under the Apache License, Version 2.0 (the "License");
4
- # you may not use this file except in compliance with the License.
5
- # You may obtain a copy of the License at
6
- #
7
- # http://www.apache.org/licenses/LICENSE-2.0
8
- #
9
- # Unless required by applicable law or agreed to in writing, software
10
- # distributed under the License is distributed on an "AS IS" BASIS,
11
- # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
- # See the License for the specific language governing permissions and
13
- # limitations under the License.
14
-
15
- from functools import partial
16
- from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple
17
-
18
- import torch
19
-
20
-
21
- def add(m1, m2, inplace):
22
- # The first operation in a checkpoint can't be in-place, but it's
23
- # nice to have in-place addition during inference. Thus...
24
- if not inplace:
25
- m1 = m1 + m2
26
- else:
27
- m1 += m2
28
-
29
- return m1
30
-
31
-
32
- def permute_final_dims(tensor: torch.Tensor, inds: List[int]):
33
- zero_index = -1 * len(inds)
34
- first_inds = list(range(len(tensor.shape[:zero_index])))
35
- return tensor.permute(first_inds + [zero_index + i for i in inds])
36
-
37
-
38
- def is_fp16_enabled():
39
- # Autocast world
40
- fp16_enabled = torch.get_autocast_gpu_dtype() == torch.float16
41
- fp16_enabled = fp16_enabled and torch.is_autocast_enabled()
42
-
43
- return fp16_enabled
44
-
45
-
46
- # With tree_map, a poor man's JAX tree_map
47
- def dict_map(fn, dic, leaf_type):
48
- new_dict = {}
49
- for k, v in dic.items():
50
- if type(v) is dict:
51
- new_dict[k] = dict_map(fn, v, leaf_type)
52
- else:
53
- new_dict[k] = tree_map(fn, v, leaf_type)
54
-
55
- return new_dict
56
-
57
-
58
- def tree_map(fn, tree, leaf_type):
59
- if isinstance(tree, dict):
60
- return dict_map(fn, tree, leaf_type)
61
- elif isinstance(tree, list):
62
- return [tree_map(fn, x, leaf_type) for x in tree]
63
- elif isinstance(tree, tuple):
64
- return tuple([tree_map(fn, x, leaf_type) for x in tree])
65
- elif isinstance(tree, leaf_type):
66
- return fn(tree)
67
- else:
68
- raise ValueError(f"Tree of type {type(tree)} not supported")
69
-
70
-
71
- tensor_tree_map = partial(tree_map, leaf_type=torch.Tensor)
72
-
73
-
74
- def flatten_final_dims(t: torch.Tensor, no_dims: int):
75
- return t.reshape(t.shape[:-no_dims] + (-1,))
76
-
77
-
78
- def _fetch_dims(tree):
79
- shapes = []
80
- tree_type = type(tree)
81
- if tree_type is dict:
82
- for v in tree.values():
83
- shapes.extend(_fetch_dims(v))
84
- elif tree_type is list or tree_type is tuple:
85
- for t in tree:
86
- shapes.extend(_fetch_dims(t))
87
- elif tree_type is torch.Tensor:
88
- shapes.append(tree.shape)
89
- else:
90
- raise ValueError("Not supported")
91
-
92
- return shapes
93
-
94
-
95
- @torch.jit.ignore
96
- def _flat_idx_to_idx(
97
- flat_idx: int,
98
- dims: Tuple[int],
99
- ) -> Tuple[int]:
100
- idx = []
101
- for d in reversed(dims):
102
- idx.append(flat_idx % d)
103
- flat_idx = flat_idx // d
104
-
105
- return tuple(reversed(idx))
106
-
107
-
108
- @torch.jit.ignore
109
- def _get_minimal_slice_set(
110
- start: Sequence[int],
111
- end: Sequence[int],
112
- dims: int,
113
- start_edges: Optional[Sequence[bool]] = None,
114
- end_edges: Optional[Sequence[bool]] = None,
115
- ) -> Sequence[Tuple[int]]:
116
- """
117
- Produces an ordered sequence of tensor slices that, when used in
118
- sequence on a tensor with shape dims, yields tensors that contain every
119
- leaf in the contiguous range [start, end]. Care is taken to yield a
120
- short sequence of slices, and perhaps even the shortest possible (I'm
121
- pretty sure it's the latter).
122
-
123
- end is INCLUSIVE.
124
- """
125
-
126
- # start_edges and end_edges both indicate whether, starting from any given
127
- # dimension, the start/end index is at the top/bottom edge of the
128
- # corresponding tensor, modeled as a tree
129
- def reduce_edge_list(l):
130
- tally = 1
131
- for i in range(len(l)):
132
- reversed_idx = -1 * (i + 1)
133
- l[reversed_idx] *= tally
134
- tally = l[reversed_idx]
135
-
136
- if start_edges is None:
137
- start_edges = [s == 0 for s in start]
138
- reduce_edge_list(start_edges)
139
- if end_edges is None:
140
- end_edges = [e == (d - 1) for e, d in zip(end, dims)]
141
- reduce_edge_list(end_edges)
142
-
143
- # Base cases. Either start/end are empty and we're done, or the final,
144
- # one-dimensional tensor can be simply sliced
145
- if len(start) == 0:
146
- return [tuple()]
147
- elif len(start) == 1:
148
- return [(slice(start[0], end[0] + 1),)]
149
-
150
- slices = []
151
- path = []
152
-
153
- # Dimensions common to start and end can be selected directly
154
- for s, e in zip(start, end):
155
- if s == e:
156
- path.append(slice(s, s + 1))
157
- else:
158
- break
159
-
160
- path = tuple(path)
161
- divergence_idx = len(path)
162
-
163
- # start == end, and we're done
164
- if divergence_idx == len(dims):
165
- return [tuple(path)]
166
-
167
- def upper():
168
- sdi = start[divergence_idx]
169
- return [
170
- path + (slice(sdi, sdi + 1),) + s
171
- for s in _get_minimal_slice_set(
172
- start[divergence_idx + 1 :],
173
- [d - 1 for d in dims[divergence_idx + 1 :]],
174
- dims[divergence_idx + 1 :],
175
- start_edges=start_edges[divergence_idx + 1 :],
176
- end_edges=[1 for _ in end_edges[divergence_idx + 1 :]],
177
- )
178
- ]
179
-
180
- def lower():
181
- edi = end[divergence_idx]
182
- return [
183
- path + (slice(edi, edi + 1),) + s
184
- for s in _get_minimal_slice_set(
185
- [0 for _ in start[divergence_idx + 1 :]],
186
- end[divergence_idx + 1 :],
187
- dims[divergence_idx + 1 :],
188
- start_edges=[1 for _ in start_edges[divergence_idx + 1 :]],
189
- end_edges=end_edges[divergence_idx + 1 :],
190
- )
191
- ]
192
-
193
- # If both start and end are at the edges of the subtree rooted at
194
- # divergence_idx, we can just select the whole subtree at once
195
- if start_edges[divergence_idx] and end_edges[divergence_idx]:
196
- slices.append(path + (slice(start[divergence_idx], end[divergence_idx] + 1),))
197
- # If just start is at the edge, we can grab almost all of the subtree,
198
- # treating only the ragged bottom edge as an edge case
199
- elif start_edges[divergence_idx]:
200
- slices.append(path + (slice(start[divergence_idx], end[divergence_idx]),))
201
- slices.extend(lower())
202
- # Analogous to the previous case, but the top is ragged this time
203
- elif end_edges[divergence_idx]:
204
- slices.extend(upper())
205
- slices.append(
206
- path + (slice(start[divergence_idx] + 1, end[divergence_idx] + 1),)
207
- )
208
- # If both sides of the range are ragged, we need to handle both sides
209
- # separately. If there's contiguous meat in between them, we can index it
210
- # in one big chunk
211
- else:
212
- slices.extend(upper())
213
- middle_ground = end[divergence_idx] - start[divergence_idx]
214
- if middle_ground > 1:
215
- slices.append(
216
- path + (slice(start[divergence_idx] + 1, end[divergence_idx]),)
217
- )
218
- slices.extend(lower())
219
-
220
- return [tuple(s) for s in slices]
221
-
222
-
223
- @torch.jit.ignore
224
- def _chunk_slice(
225
- t: torch.Tensor,
226
- flat_start: int,
227
- flat_end: int,
228
- no_batch_dims: int,
229
- ) -> torch.Tensor:
230
- """
231
- Equivalent to
232
-
233
- t.reshape((-1,) + t.shape[no_batch_dims:])[flat_start:flat_end]
234
-
235
- but without the need for the initial reshape call, which can be
236
- memory-intensive in certain situations. The only reshape operations
237
- in this function are performed on sub-tensors that scale with
238
- (flat_end - flat_start), the chunk size.
239
- """
240
-
241
- batch_dims = t.shape[:no_batch_dims]
242
- start_idx = list(_flat_idx_to_idx(flat_start, batch_dims))
243
- # _get_minimal_slice_set is inclusive
244
- end_idx = list(_flat_idx_to_idx(flat_end - 1, batch_dims))
245
-
246
- # Get an ordered list of slices to perform
247
- slices = _get_minimal_slice_set(
248
- start_idx,
249
- end_idx,
250
- batch_dims,
251
- )
252
-
253
- sliced_tensors = [t[s] for s in slices]
254
-
255
- return torch.cat([s.view((-1,) + t.shape[no_batch_dims:]) for s in sliced_tensors])
256
-
257
-
258
- def chunk_layer(
259
- layer: Callable,
260
- inputs: Dict[str, Any],
261
- chunk_size: int,
262
- no_batch_dims: int,
263
- low_mem: bool = False,
264
- _out: Any = None,
265
- _add_into_out: bool = False,
266
- ) -> Any:
267
- """
268
- Implements the "chunking" procedure described in section 1.11.8.
269
-
270
- Layer outputs and inputs are assumed to be simple "pytrees,"
271
- consisting only of (arbitrarily nested) lists, tuples, and dicts with
272
- torch.Tensor leaves.
273
-
274
- Args:
275
- layer:
276
- The layer to be applied chunk-wise
277
- inputs:
278
- A (non-nested) dictionary of keyworded inputs. All leaves must
279
- be tensors and must share the same batch dimensions.
280
- chunk_size:
281
- The number of sub-batches per chunk. If multiple batch
282
- dimensions are specified, a "sub-batch" is defined as a single
283
- indexing of all batch dimensions simultaneously (s.t. the
284
- number of sub-batches is the product of the batch dimensions).
285
- no_batch_dims:
286
- How many of the initial dimensions of each input tensor can
287
- be considered batch dimensions.
288
- low_mem:
289
- Avoids flattening potentially large input tensors. Unnecessary
290
- in most cases, and is ever so slightly slower than the default
291
- setting.
292
- Returns:
293
- The reassembled output of the layer on the inputs.
294
- """
295
- if not (len(inputs) > 0):
296
- raise ValueError("Must provide at least one input")
297
-
298
- initial_dims = [shape[:no_batch_dims] for shape in _fetch_dims(inputs)]
299
- orig_batch_dims = tuple([max(s) for s in zip(*initial_dims)])
300
-
301
- def _prep_inputs(t):
302
- if not low_mem:
303
- if not sum(t.shape[:no_batch_dims]) == no_batch_dims:
304
- t = t.expand(orig_batch_dims + t.shape[no_batch_dims:])
305
- t = t.reshape(-1, *t.shape[no_batch_dims:])
306
- else:
307
- t = t.expand(orig_batch_dims + t.shape[no_batch_dims:])
308
- return t
309
-
310
- prepped_inputs = tensor_tree_map(_prep_inputs, inputs)
311
- prepped_outputs = None
312
- if _out is not None:
313
- reshape_fn = lambda t: t.view([-1] + list(t.shape[no_batch_dims:]))
314
- prepped_outputs = tensor_tree_map(reshape_fn, _out)
315
-
316
- flat_batch_dim = 1
317
- for d in orig_batch_dims:
318
- flat_batch_dim *= d
319
-
320
- no_chunks = flat_batch_dim // chunk_size + (flat_batch_dim % chunk_size != 0)
321
-
322
- i = 0
323
- out = prepped_outputs
324
- for _ in range(no_chunks):
325
- # Chunk the input
326
- if not low_mem:
327
- select_chunk = lambda t: t[i : i + chunk_size] if t.shape[0] != 1 else t
328
- else:
329
- select_chunk = partial(
330
- _chunk_slice,
331
- flat_start=i,
332
- flat_end=min(flat_batch_dim, i + chunk_size),
333
- no_batch_dims=len(orig_batch_dims),
334
- )
335
-
336
- chunks = tensor_tree_map(select_chunk, prepped_inputs)
337
-
338
- # Run the layer on the chunk
339
- output_chunk = layer(**chunks)
340
-
341
- # Allocate space for the output
342
- if out is None:
343
- allocate = lambda t: t.new_zeros((flat_batch_dim,) + t.shape[1:])
344
- out = tensor_tree_map(allocate, output_chunk)
345
-
346
- # Put the chunk in its pre-allocated space
347
- out_type = type(output_chunk)
348
- if out_type is dict:
349
-
350
- def assign(d1, d2):
351
- for k, v in d1.items():
352
- if type(v) is dict:
353
- assign(v, d2[k])
354
- else:
355
- if _add_into_out:
356
- v[i : i + chunk_size] += d2[k]
357
- else:
358
- v[i : i + chunk_size] = d2[k]
359
-
360
- assign(out, output_chunk)
361
- elif out_type is tuple:
362
- for x1, x2 in zip(out, output_chunk):
363
- if _add_into_out:
364
- x1[i : i + chunk_size] += x2
365
- else:
366
- x1[i : i + chunk_size] = x2
367
- elif out_type is torch.Tensor:
368
- if _add_into_out:
369
- out[i : i + chunk_size] += output_chunk
370
- else:
371
- out[i : i + chunk_size] = output_chunk
372
- else:
373
- raise ValueError("Not supported")
374
-
375
- i += chunk_size
376
-
377
- reshape = lambda t: t.view(orig_batch_dims + t.shape[1:])
378
- out = tensor_tree_map(reshape, out)
379
-
380
- return out
 
1
+ # Copyright 2021 AlQuraishi Laboratory
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ from functools import partial
16
+ from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple
17
+
18
+ import torch
19
+
20
+
21
+ def add(m1, m2, inplace):
22
+ # The first operation in a checkpoint can't be in-place, but it's
23
+ # nice to have in-place addition during inference. Thus...
24
+ if not inplace:
25
+ m1 = m1 + m2
26
+ else:
27
+ m1 += m2
28
+
29
+ return m1
30
+
31
+
32
+ def permute_final_dims(tensor: torch.Tensor, inds: List[int]):
33
+ zero_index = -1 * len(inds)
34
+ first_inds = list(range(len(tensor.shape[:zero_index])))
35
+ return tensor.permute(first_inds + [zero_index + i for i in inds])
36
+
37
+
38
+ def is_fp16_enabled():
39
+ # Autocast world
40
+ fp16_enabled = torch.get_autocast_gpu_dtype() == torch.float16
41
+ fp16_enabled = fp16_enabled and torch.is_autocast_enabled()
42
+
43
+ return fp16_enabled
44
+
45
+
46
+ # With tree_map, a poor man's JAX tree_map
47
+ def dict_map(fn, dic, leaf_type):
48
+ new_dict = {}
49
+ for k, v in dic.items():
50
+ if type(v) is dict:
51
+ new_dict[k] = dict_map(fn, v, leaf_type)
52
+ else:
53
+ new_dict[k] = tree_map(fn, v, leaf_type)
54
+
55
+ return new_dict
56
+
57
+
58
+ def tree_map(fn, tree, leaf_type):
59
+ if isinstance(tree, dict):
60
+ return dict_map(fn, tree, leaf_type)
61
+ elif isinstance(tree, list):
62
+ return [tree_map(fn, x, leaf_type) for x in tree]
63
+ elif isinstance(tree, tuple):
64
+ return tuple([tree_map(fn, x, leaf_type) for x in tree])
65
+ elif isinstance(tree, leaf_type):
66
+ return fn(tree)
67
+ else:
68
+ raise ValueError(f"Tree of type {type(tree)} not supported")
69
+
70
+
71
+ tensor_tree_map = partial(tree_map, leaf_type=torch.Tensor)
72
+
73
+
74
+ def flatten_final_dims(t: torch.Tensor, no_dims: int):
75
+ return t.reshape(t.shape[:-no_dims] + (-1,))
76
+
77
+
78
+ def _fetch_dims(tree):
79
+ shapes = []
80
+ tree_type = type(tree)
81
+ if tree_type is dict:
82
+ for v in tree.values():
83
+ shapes.extend(_fetch_dims(v))
84
+ elif tree_type is list or tree_type is tuple:
85
+ for t in tree:
86
+ shapes.extend(_fetch_dims(t))
87
+ elif tree_type is torch.Tensor:
88
+ shapes.append(tree.shape)
89
+ else:
90
+ raise ValueError("Not supported")
91
+
92
+ return shapes
93
+
94
+
95
+ @torch.jit.ignore
96
+ def _flat_idx_to_idx(
97
+ flat_idx: int,
98
+ dims: Tuple[int],
99
+ ) -> Tuple[int]:
100
+ idx = []
101
+ for d in reversed(dims):
102
+ idx.append(flat_idx % d)
103
+ flat_idx = flat_idx // d
104
+
105
+ return tuple(reversed(idx))
106
+
107
+
108
+ @torch.jit.ignore
109
+ def _get_minimal_slice_set(
110
+ start: Sequence[int],
111
+ end: Sequence[int],
112
+ dims: int,
113
+ start_edges: Optional[Sequence[bool]] = None,
114
+ end_edges: Optional[Sequence[bool]] = None,
115
+ ) -> Sequence[Tuple[int]]:
116
+ """
117
+ Produces an ordered sequence of tensor slices that, when used in
118
+ sequence on a tensor with shape dims, yields tensors that contain every
119
+ leaf in the contiguous range [start, end]. Care is taken to yield a
120
+ short sequence of slices, and perhaps even the shortest possible (I'm
121
+ pretty sure it's the latter).
122
+
123
+ end is INCLUSIVE.
124
+ """
125
+
126
+ # start_edges and end_edges both indicate whether, starting from any given
127
+ # dimension, the start/end index is at the top/bottom edge of the
128
+ # corresponding tensor, modeled as a tree
129
+ def reduce_edge_list(l):
130
+ tally = 1
131
+ for i in range(len(l)):
132
+ reversed_idx = -1 * (i + 1)
133
+ l[reversed_idx] *= tally
134
+ tally = l[reversed_idx]
135
+
136
+ if start_edges is None:
137
+ start_edges = [s == 0 for s in start]
138
+ reduce_edge_list(start_edges)
139
+ if end_edges is None:
140
+ end_edges = [e == (d - 1) for e, d in zip(end, dims)]
141
+ reduce_edge_list(end_edges)
142
+
143
+ # Base cases. Either start/end are empty and we're done, or the final,
144
+ # one-dimensional tensor can be simply sliced
145
+ if len(start) == 0:
146
+ return [tuple()]
147
+ elif len(start) == 1:
148
+ return [(slice(start[0], end[0] + 1),)]
149
+
150
+ slices = []
151
+ path = []
152
+
153
+ # Dimensions common to start and end can be selected directly
154
+ for s, e in zip(start, end):
155
+ if s == e:
156
+ path.append(slice(s, s + 1))
157
+ else:
158
+ break
159
+
160
+ path = tuple(path)
161
+ divergence_idx = len(path)
162
+
163
+ # start == end, and we're done
164
+ if divergence_idx == len(dims):
165
+ return [tuple(path)]
166
+
167
+ def upper():
168
+ sdi = start[divergence_idx]
169
+ return [
170
+ path + (slice(sdi, sdi + 1),) + s
171
+ for s in _get_minimal_slice_set(
172
+ start[divergence_idx + 1 :],
173
+ [d - 1 for d in dims[divergence_idx + 1 :]],
174
+ dims[divergence_idx + 1 :],
175
+ start_edges=start_edges[divergence_idx + 1 :],
176
+ end_edges=[1 for _ in end_edges[divergence_idx + 1 :]],
177
+ )
178
+ ]
179
+
180
+ def lower():
181
+ edi = end[divergence_idx]
182
+ return [
183
+ path + (slice(edi, edi + 1),) + s
184
+ for s in _get_minimal_slice_set(
185
+ [0 for _ in start[divergence_idx + 1 :]],
186
+ end[divergence_idx + 1 :],
187
+ dims[divergence_idx + 1 :],
188
+ start_edges=[1 for _ in start_edges[divergence_idx + 1 :]],
189
+ end_edges=end_edges[divergence_idx + 1 :],
190
+ )
191
+ ]
192
+
193
+ # If both start and end are at the edges of the subtree rooted at
194
+ # divergence_idx, we can just select the whole subtree at once
195
+ if start_edges[divergence_idx] and end_edges[divergence_idx]:
196
+ slices.append(path + (slice(start[divergence_idx], end[divergence_idx] + 1),))
197
+ # If just start is at the edge, we can grab almost all of the subtree,
198
+ # treating only the ragged bottom edge as an edge case
199
+ elif start_edges[divergence_idx]:
200
+ slices.append(path + (slice(start[divergence_idx], end[divergence_idx]),))
201
+ slices.extend(lower())
202
+ # Analogous to the previous case, but the top is ragged this time
203
+ elif end_edges[divergence_idx]:
204
+ slices.extend(upper())
205
+ slices.append(
206
+ path + (slice(start[divergence_idx] + 1, end[divergence_idx] + 1),)
207
+ )
208
+ # If both sides of the range are ragged, we need to handle both sides
209
+ # separately. If there's contiguous meat in between them, we can index it
210
+ # in one big chunk
211
+ else:
212
+ slices.extend(upper())
213
+ middle_ground = end[divergence_idx] - start[divergence_idx]
214
+ if middle_ground > 1:
215
+ slices.append(
216
+ path + (slice(start[divergence_idx] + 1, end[divergence_idx]),)
217
+ )
218
+ slices.extend(lower())
219
+
220
+ return [tuple(s) for s in slices]
221
+
222
+
223
+ @torch.jit.ignore
224
+ def _chunk_slice(
225
+ t: torch.Tensor,
226
+ flat_start: int,
227
+ flat_end: int,
228
+ no_batch_dims: int,
229
+ ) -> torch.Tensor:
230
+ """
231
+ Equivalent to
232
+
233
+ t.reshape((-1,) + t.shape[no_batch_dims:])[flat_start:flat_end]
234
+
235
+ but without the need for the initial reshape call, which can be
236
+ memory-intensive in certain situations. The only reshape operations
237
+ in this function are performed on sub-tensors that scale with
238
+ (flat_end - flat_start), the chunk size.
239
+ """
240
+
241
+ batch_dims = t.shape[:no_batch_dims]
242
+ start_idx = list(_flat_idx_to_idx(flat_start, batch_dims))
243
+ # _get_minimal_slice_set is inclusive
244
+ end_idx = list(_flat_idx_to_idx(flat_end - 1, batch_dims))
245
+
246
+ # Get an ordered list of slices to perform
247
+ slices = _get_minimal_slice_set(
248
+ start_idx,
249
+ end_idx,
250
+ batch_dims,
251
+ )
252
+
253
+ sliced_tensors = [t[s] for s in slices]
254
+
255
+ return torch.cat([s.view((-1,) + t.shape[no_batch_dims:]) for s in sliced_tensors])
256
+
257
+
258
+ def chunk_layer(
259
+ layer: Callable,
260
+ inputs: Dict[str, Any],
261
+ chunk_size: int,
262
+ no_batch_dims: int,
263
+ low_mem: bool = False,
264
+ _out: Any = None,
265
+ _add_into_out: bool = False,
266
+ ) -> Any:
267
+ """
268
+ Implements the "chunking" procedure described in section 1.11.8.
269
+
270
+ Layer outputs and inputs are assumed to be simple "pytrees,"
271
+ consisting only of (arbitrarily nested) lists, tuples, and dicts with
272
+ torch.Tensor leaves.
273
+
274
+ Args:
275
+ layer:
276
+ The layer to be applied chunk-wise
277
+ inputs:
278
+ A (non-nested) dictionary of keyworded inputs. All leaves must
279
+ be tensors and must share the same batch dimensions.
280
+ chunk_size:
281
+ The number of sub-batches per chunk. If multiple batch
282
+ dimensions are specified, a "sub-batch" is defined as a single
283
+ indexing of all batch dimensions simultaneously (s.t. the
284
+ number of sub-batches is the product of the batch dimensions).
285
+ no_batch_dims:
286
+ How many of the initial dimensions of each input tensor can
287
+ be considered batch dimensions.
288
+ low_mem:
289
+ Avoids flattening potentially large input tensors. Unnecessary
290
+ in most cases, and is ever so slightly slower than the default
291
+ setting.
292
+ Returns:
293
+ The reassembled output of the layer on the inputs.
294
+ """
295
+ if not (len(inputs) > 0):
296
+ raise ValueError("Must provide at least one input")
297
+
298
+ initial_dims = [shape[:no_batch_dims] for shape in _fetch_dims(inputs)]
299
+ orig_batch_dims = tuple([max(s) for s in zip(*initial_dims)])
300
+
301
+ def _prep_inputs(t):
302
+ if not low_mem:
303
+ if not sum(t.shape[:no_batch_dims]) == no_batch_dims:
304
+ t = t.expand(orig_batch_dims + t.shape[no_batch_dims:])
305
+ t = t.reshape(-1, *t.shape[no_batch_dims:])
306
+ else:
307
+ t = t.expand(orig_batch_dims + t.shape[no_batch_dims:])
308
+ return t
309
+
310
+ prepped_inputs = tensor_tree_map(_prep_inputs, inputs)
311
+ prepped_outputs = None
312
+ if _out is not None:
313
+ reshape_fn = lambda t: t.view([-1] + list(t.shape[no_batch_dims:]))
314
+ prepped_outputs = tensor_tree_map(reshape_fn, _out)
315
+
316
+ flat_batch_dim = 1
317
+ for d in orig_batch_dims:
318
+ flat_batch_dim *= d
319
+
320
+ no_chunks = flat_batch_dim // chunk_size + (flat_batch_dim % chunk_size != 0)
321
+
322
+ i = 0
323
+ out = prepped_outputs
324
+ for _ in range(no_chunks):
325
+ # Chunk the input
326
+ if not low_mem:
327
+ select_chunk = lambda t: t[i : i + chunk_size] if t.shape[0] != 1 else t
328
+ else:
329
+ select_chunk = partial(
330
+ _chunk_slice,
331
+ flat_start=i,
332
+ flat_end=min(flat_batch_dim, i + chunk_size),
333
+ no_batch_dims=len(orig_batch_dims),
334
+ )
335
+
336
+ chunks = tensor_tree_map(select_chunk, prepped_inputs)
337
+
338
+ # Run the layer on the chunk
339
+ output_chunk = layer(**chunks)
340
+
341
+ # Allocate space for the output
342
+ if out is None:
343
+ allocate = lambda t: t.new_zeros((flat_batch_dim,) + t.shape[1:])
344
+ out = tensor_tree_map(allocate, output_chunk)
345
+
346
+ # Put the chunk in its pre-allocated space
347
+ out_type = type(output_chunk)
348
+ if out_type is dict:
349
+
350
+ def assign(d1, d2):
351
+ for k, v in d1.items():
352
+ if type(v) is dict:
353
+ assign(v, d2[k])
354
+ else:
355
+ if _add_into_out:
356
+ v[i : i + chunk_size] += d2[k]
357
+ else:
358
+ v[i : i + chunk_size] = d2[k]
359
+
360
+ assign(out, output_chunk)
361
+ elif out_type is tuple:
362
+ for x1, x2 in zip(out, output_chunk):
363
+ if _add_into_out:
364
+ x1[i : i + chunk_size] += x2
365
+ else:
366
+ x1[i : i + chunk_size] = x2
367
+ elif out_type is torch.Tensor:
368
+ if _add_into_out:
369
+ out[i : i + chunk_size] += output_chunk
370
+ else:
371
+ out[i : i + chunk_size] = output_chunk
372
+ else:
373
+ raise ValueError("Not supported")
374
+
375
+ i += chunk_size
376
+
377
+ reshape = lambda t: t.view(orig_batch_dims + t.shape[1:])
378
+ out = tensor_tree_map(reshape, out)
379
+
380
+ return out