protenix 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (130) hide show
  1. protenix-0.1.0/LICENSE +413 -0
  2. protenix-0.1.0/PKG-INFO +11 -0
  3. protenix-0.1.0/README.md +255 -0
  4. protenix-0.1.0/configs/__init__.py +0 -0
  5. protenix-0.1.0/configs/configs_base.py +353 -0
  6. protenix-0.1.0/configs/configs_data.py +169 -0
  7. protenix-0.1.0/configs/configs_inference.py +27 -0
  8. protenix-0.1.0/protenix/__init__.py +0 -0
  9. protenix-0.1.0/protenix/config/__init__.py +15 -0
  10. protenix-0.1.0/protenix/config/config.py +286 -0
  11. protenix-0.1.0/protenix/config/extend_types.py +56 -0
  12. protenix-0.1.0/protenix/data/__init__.py +0 -0
  13. protenix-0.1.0/protenix/data/ccd.py +480 -0
  14. protenix-0.1.0/protenix/data/constants.py +965 -0
  15. protenix-0.1.0/protenix/data/data_pipeline.py +311 -0
  16. protenix-0.1.0/protenix/data/dataloader.py +373 -0
  17. protenix-0.1.0/protenix/data/dataset.py +1101 -0
  18. protenix-0.1.0/protenix/data/featurizer.py +803 -0
  19. protenix-0.1.0/protenix/data/infer_data_pipeline.py +221 -0
  20. protenix-0.1.0/protenix/data/json_parser.py +609 -0
  21. protenix-0.1.0/protenix/data/json_to_feature.py +324 -0
  22. protenix-0.1.0/protenix/data/msa_featurizer.py +1163 -0
  23. protenix-0.1.0/protenix/data/msa_utils.py +1417 -0
  24. protenix-0.1.0/protenix/data/parser.py +763 -0
  25. protenix-0.1.0/protenix/data/substructure_perms.py +258 -0
  26. protenix-0.1.0/protenix/data/tokenizer.py +197 -0
  27. protenix-0.1.0/protenix/data/utils.py +481 -0
  28. protenix-0.1.0/protenix/metrics/__init__.py +0 -0
  29. protenix-0.1.0/protenix/metrics/clash.py +273 -0
  30. protenix-0.1.0/protenix/metrics/lddt_metrics.py +278 -0
  31. protenix-0.1.0/protenix/metrics/rmsd.py +261 -0
  32. protenix-0.1.0/protenix/model/__init__.py +0 -0
  33. protenix-0.1.0/protenix/model/generator.py +333 -0
  34. protenix-0.1.0/protenix/model/layer_norm/__init__.py +16 -0
  35. protenix-0.1.0/protenix/model/layer_norm/kernel/compat.h +24 -0
  36. protenix-0.1.0/protenix/model/layer_norm/kernel/layer_norm_cuda.cpp +138 -0
  37. protenix-0.1.0/protenix/model/layer_norm/kernel/layer_norm_cuda_kernel.cu +409 -0
  38. protenix-0.1.0/protenix/model/layer_norm/kernel/type_shim.h +246 -0
  39. protenix-0.1.0/protenix/model/layer_norm/layer_norm.py +134 -0
  40. protenix-0.1.0/protenix/model/layer_norm/torch_ext_compile.py +56 -0
  41. protenix-0.1.0/protenix/model/loss.py +1807 -0
  42. protenix-0.1.0/protenix/model/modules/__init__.py +0 -0
  43. protenix-0.1.0/protenix/model/modules/confidence.py +295 -0
  44. protenix-0.1.0/protenix/model/modules/diffusion.py +542 -0
  45. protenix-0.1.0/protenix/model/modules/embedders.py +261 -0
  46. protenix-0.1.0/protenix/model/modules/frames.py +109 -0
  47. protenix-0.1.0/protenix/model/modules/head.py +54 -0
  48. protenix-0.1.0/protenix/model/modules/pairformer.py +828 -0
  49. protenix-0.1.0/protenix/model/modules/primitives.py +888 -0
  50. protenix-0.1.0/protenix/model/modules/transformer.py +947 -0
  51. protenix-0.1.0/protenix/model/protenix.py +726 -0
  52. protenix-0.1.0/protenix/model/sample_confidence.py +772 -0
  53. protenix-0.1.0/protenix/model/utils.py +472 -0
  54. protenix-0.1.0/protenix/openfold_local/__init__.py +3 -0
  55. protenix-0.1.0/protenix/openfold_local/data/__init__.py +0 -0
  56. protenix-0.1.0/protenix/openfold_local/data/data_transforms.py +221 -0
  57. protenix-0.1.0/protenix/openfold_local/data/errors.py +24 -0
  58. protenix-0.1.0/protenix/openfold_local/data/mmcif_parsing.py +494 -0
  59. protenix-0.1.0/protenix/openfold_local/data/msa_identifiers.py +81 -0
  60. protenix-0.1.0/protenix/openfold_local/data/msa_pairing.py +502 -0
  61. protenix-0.1.0/protenix/openfold_local/data/parsers.py +642 -0
  62. protenix-0.1.0/protenix/openfold_local/data/templates.py +592 -0
  63. protenix-0.1.0/protenix/openfold_local/data/tools/__init__.py +0 -0
  64. protenix-0.1.0/protenix/openfold_local/data/tools/jackhmmer.py +245 -0
  65. protenix-0.1.0/protenix/openfold_local/data/tools/utils.py +46 -0
  66. protenix-0.1.0/protenix/openfold_local/model/__init__.py +0 -0
  67. protenix-0.1.0/protenix/openfold_local/model/dropout.py +78 -0
  68. protenix-0.1.0/protenix/openfold_local/model/outer_product_mean.py +156 -0
  69. protenix-0.1.0/protenix/openfold_local/model/primitives.py +859 -0
  70. protenix-0.1.0/protenix/openfold_local/model/triangular_attention.py +163 -0
  71. protenix-0.1.0/protenix/openfold_local/model/triangular_multiplicative_update.py +626 -0
  72. protenix-0.1.0/protenix/openfold_local/np/__init__.py +0 -0
  73. protenix-0.1.0/protenix/openfold_local/np/residue_constants.py +1430 -0
  74. protenix-0.1.0/protenix/openfold_local/utils/__init__.py +0 -0
  75. protenix-0.1.0/protenix/openfold_local/utils/all_atom_multimer.py +468 -0
  76. protenix-0.1.0/protenix/openfold_local/utils/checkpointing.py +96 -0
  77. protenix-0.1.0/protenix/openfold_local/utils/chunk_utils.py +411 -0
  78. protenix-0.1.0/protenix/openfold_local/utils/feats.py +268 -0
  79. protenix-0.1.0/protenix/openfold_local/utils/geometry/__init__.py +28 -0
  80. protenix-0.1.0/protenix/openfold_local/utils/geometry/quat_rigid.py +42 -0
  81. protenix-0.1.0/protenix/openfold_local/utils/geometry/rigid_matrix_vector.py +185 -0
  82. protenix-0.1.0/protenix/openfold_local/utils/geometry/rotation_matrix.py +189 -0
  83. protenix-0.1.0/protenix/openfold_local/utils/geometry/test_utils.py +106 -0
  84. protenix-0.1.0/protenix/openfold_local/utils/geometry/utils.py +22 -0
  85. protenix-0.1.0/protenix/openfold_local/utils/geometry/vector.py +257 -0
  86. protenix-0.1.0/protenix/openfold_local/utils/kernel/__init__.py +0 -0
  87. protenix-0.1.0/protenix/openfold_local/utils/kernel/attention_core.py +101 -0
  88. protenix-0.1.0/protenix/openfold_local/utils/precision_utils.py +23 -0
  89. protenix-0.1.0/protenix/openfold_local/utils/rigid_utils.py +1381 -0
  90. protenix-0.1.0/protenix/openfold_local/utils/tensor_utils.py +115 -0
  91. protenix-0.1.0/protenix/utils/__init__.py +0 -0
  92. protenix-0.1.0/protenix/utils/cropping.py +699 -0
  93. protenix-0.1.0/protenix/utils/distributed.py +73 -0
  94. protenix-0.1.0/protenix/utils/file_io.py +104 -0
  95. protenix-0.1.0/protenix/utils/geometry.py +72 -0
  96. protenix-0.1.0/protenix/utils/logger.py +85 -0
  97. protenix-0.1.0/protenix/utils/lr_scheduler.py +147 -0
  98. protenix-0.1.0/protenix/utils/metrics.py +75 -0
  99. protenix-0.1.0/protenix/utils/permutation/__init__.py +0 -0
  100. protenix-0.1.0/protenix/utils/permutation/atom_permutation.py +824 -0
  101. protenix-0.1.0/protenix/utils/permutation/chain_permutation/__init__.py +127 -0
  102. protenix-0.1.0/protenix/utils/permutation/chain_permutation/heuristic.py +947 -0
  103. protenix-0.1.0/protenix/utils/permutation/chain_permutation/pocket_based_permutation.py +265 -0
  104. protenix-0.1.0/protenix/utils/permutation/chain_permutation/utils.py +64 -0
  105. protenix-0.1.0/protenix/utils/permutation/permutation.py +482 -0
  106. protenix-0.1.0/protenix/utils/permutation/utils.py +140 -0
  107. protenix-0.1.0/protenix/utils/scatter_utils.py +201 -0
  108. protenix-0.1.0/protenix/utils/seed.py +35 -0
  109. protenix-0.1.0/protenix/utils/torch_utils.py +155 -0
  110. protenix-0.1.0/protenix/utils/training.py +119 -0
  111. protenix-0.1.0/protenix/web_service/__init__.py +0 -0
  112. protenix-0.1.0/protenix/web_service/colab_request_parser.py +422 -0
  113. protenix-0.1.0/protenix/web_service/colab_request_utils.py +260 -0
  114. protenix-0.1.0/protenix/web_service/dependency_url.py +20 -0
  115. protenix-0.1.0/protenix/web_service/prediction_visualization.py +264 -0
  116. protenix-0.1.0/protenix/web_service/viewer.py +616 -0
  117. protenix-0.1.0/protenix.egg-info/PKG-INFO +11 -0
  118. protenix-0.1.0/protenix.egg-info/SOURCES.txt +128 -0
  119. protenix-0.1.0/protenix.egg-info/dependency_links.txt +1 -0
  120. protenix-0.1.0/protenix.egg-info/entry_points.txt +2 -0
  121. protenix-0.1.0/protenix.egg-info/requires.txt +25 -0
  122. protenix-0.1.0/protenix.egg-info/top_level.txt +3 -0
  123. protenix-0.1.0/requirements.txt +25 -0
  124. protenix-0.1.0/runner/__init__.py +0 -0
  125. protenix-0.1.0/runner/dumper.py +175 -0
  126. protenix-0.1.0/runner/ema.py +64 -0
  127. protenix-0.1.0/runner/inference.py +267 -0
  128. protenix-0.1.0/runner/train.py +569 -0
  129. protenix-0.1.0/setup.cfg +4 -0
  130. protenix-0.1.0/setup.py +49 -0
protenix-0.1.0/LICENSE ADDED
@@ -0,0 +1,413 @@
1
+ Attribution-NonCommercial 4.0 International
2
+
3
+ =======================================================================
4
+
5
+ Creative Commons Corporation ("Creative Commons") is not a law firm and
6
+ does not provide legal services or legal advice. Distribution of
7
+ Creative Commons public licenses does not create a lawyer-client or
8
+ other relationship. Creative Commons makes its licenses and related
9
+ information available on an "as-is" basis. Creative Commons gives no
10
+ warranties regarding its licenses, any material licensed under their
11
+ terms and conditions, or any related information. Creative Commons
12
+ disclaims all liability for damages resulting from their use to the
13
+ fullest extent possible.
14
+
15
+ Using Creative Commons Public Licenses
16
+
17
+ Creative Commons public licenses provide a standard set of terms and
18
+ conditions that creators and other rights holders may use to share
19
+ original works of authorship and other material subject to copyright
20
+ and certain other rights specified in the public license below. The
21
+ following considerations are for informational purposes only, are not
22
+ exhaustive, and do not form part of our licenses.
23
+
24
+ Considerations for licensors: Our public licenses are
25
+ intended for use by those authorized to give the public
26
+ permission to use material in ways otherwise restricted by
27
+ copyright and certain other rights. Our licenses are
28
+ irrevocable. Licensors should read and understand the terms
29
+ and conditions of the license they choose before applying it.
30
+ Licensors should also secure all rights necessary before
31
+ applying our licenses so that the public can reuse the
32
+ material as expected. Licensors should clearly mark any
33
+ material not subject to the license. This includes other CC-
34
+ licensed material, or material used under an exception or
35
+ limitation to copyright. More considerations for licensors:
36
+ wiki.creativecommons.org/Considerations_for_licensors
37
+
38
+ Considerations for the public: By using one of our public
39
+ licenses, a licensor grants the public permission to use the
40
+ licensed material under specified terms and conditions. If
41
+ the licensor's permission is not necessary for any reason--for
42
+ example, because of any applicable exception or limitation to
43
+ copyright--then that use is not regulated by the license. Our
44
+ licenses grant only permissions under copyright and certain
45
+ other rights that a licensor has authority to grant. Use of
46
+ the licensed material may still be restricted for other
47
+ reasons, including because others have copyright or other
48
+ rights in the material. A licensor may make special requests,
49
+ such as asking that all changes be marked or described.
50
+ Although not required by our licenses, you are encouraged to
51
+ respect those requests where reasonable. More considerations
52
+ for the public:
53
+ wiki.creativecommons.org/Considerations_for_licensees
54
+
55
+ =======================================================================
56
+
57
+ Creative Commons Attribution-NonCommercial 4.0 International Public
58
+ License
59
+
60
+ By exercising the Licensed Rights (defined below), You accept and agree
61
+ to be bound by the terms and conditions of this Creative Commons
62
+ Attribution-NonCommercial 4.0 International Public License ("Public
63
+ License"). To the extent this Public License may be interpreted as a
64
+ contract, You are granted the Licensed Rights in consideration of Your
65
+ acceptance of these terms and conditions, and the Licensor grants You
66
+ such rights in consideration of benefits the Licensor receives from
67
+ making the Licensed Material available under these terms and
68
+ conditions.
69
+
70
+
71
+ Section 1 -- Definitions.
72
+
73
+ a. Adapted Material means material subject to Copyright and Similar
74
+ Rights that is derived from or based upon the Licensed Material
75
+ and in which the Licensed Material is translated, altered,
76
+ arranged, transformed, or otherwise modified in a manner requiring
77
+ permission under the Copyright and Similar Rights held by the
78
+ Licensor. For purposes of this Public License, where the Licensed
79
+ Material is a musical work, performance, or sound recording,
80
+ Adapted Material is always produced where the Licensed Material is
81
+ synched in timed relation with a moving image.
82
+
83
+ b. Adapter's License means the license You apply to Your Copyright
84
+ and Similar Rights in Your contributions to Adapted Material in
85
+ accordance with the terms and conditions of this Public License.
86
+
87
+ c. Copyright and Similar Rights means copyright and/or similar rights
88
+ closely related to copyright including, without limitation,
89
+ performance, broadcast, sound recording, and Sui Generis Database
90
+ Rights, without regard to how the rights are labeled or
91
+ categorized. For purposes of this Public License, the rights
92
+ specified in Section 2(b)(1)-(2) are not Copyright and Similar
93
+ Rights.
94
+ d. Effective Technological Measures means those measures that, in the
95
+ absence of proper authority, may not be circumvented under laws
96
+ fulfilling obligations under Article 11 of the WIPO Copyright
97
+ Treaty adopted on December 20, 1996, and/or similar international
98
+ agreements.
99
+
100
+ e. Exceptions and Limitations means fair use, fair dealing, and/or
101
+ any other exception or limitation to Copyright and Similar Rights
102
+ that applies to Your use of the Licensed Material.
103
+
104
+ f. Licensed Material means the artistic or literary work, database,
105
+ or other material to which the Licensor applied this Public
106
+ License.
107
+
108
+ g. Licensed Rights means the rights granted to You subject to the
109
+ terms and conditions of this Public License, which are limited to
110
+ all Copyright and Similar Rights that apply to Your use of the
111
+ Licensed Material and that the Licensor has authority to license.
112
+
113
+ h. Licensor means the individual(s) or entity(ies) granting rights
114
+ under this Public License.
115
+
116
+ i. NonCommercial means not primarily intended for or directed towards
117
+ commercial advantage or monetary compensation. For purposes of
118
+ this Public License, the exchange of the Licensed Material for
119
+ other material subject to Copyright and Similar Rights by digital
120
+ file-sharing or similar means is NonCommercial provided there is
121
+ no payment of monetary compensation in connection with the
122
+ exchange.
123
+
124
+ j. Share means to provide material to the public by any means or
125
+ process that requires permission under the Licensed Rights, such
126
+ as reproduction, public display, public performance, distribution,
127
+ dissemination, communication, or importation, and to make material
128
+ available to the public including in ways that members of the
129
+ public may access the material from a place and at a time
130
+ individually chosen by them.
131
+
132
+ k. Sui Generis Database Rights means rights other than copyright
133
+ resulting from Directive 96/9/EC of the European Parliament and of
134
+ the Council of 11 March 1996 on the legal protection of databases,
135
+ as amended and/or succeeded, as well as other essentially
136
+ equivalent rights anywhere in the world.
137
+
138
+ l. You means the individual or entity exercising the Licensed Rights
139
+ under this Public License. Your has a corresponding meaning.
140
+
141
+
142
+ Section 2 -- Scope.
143
+
144
+ a. License grant.
145
+
146
+ 1. Subject to the terms and conditions of this Public License,
147
+ the Licensor hereby grants You a worldwide, royalty-free,
148
+ non-sublicensable, non-exclusive, irrevocable license to
149
+ exercise the Licensed Rights in the Licensed Material to:
150
+
151
+ a. reproduce and Share the Licensed Material, in whole or
152
+ in part, for NonCommercial purposes only; and
153
+
154
+ b. produce, reproduce, and Share Adapted Material for
155
+ NonCommercial purposes only.
156
+
157
+ 2. Exceptions and Limitations. For the avoidance of doubt, where
158
+ Exceptions and Limitations apply to Your use, this Public
159
+ License does not apply, and You do not need to comply with
160
+ its terms and conditions.
161
+
162
+ 3. Term. The term of this Public License is specified in Section
163
+ 6(a).
164
+
165
+ 4. Media and formats; technical modifications allowed. The
166
+ Licensor authorizes You to exercise the Licensed Rights in
167
+ all media and formats whether now known or hereafter created,
168
+ and to make technical modifications necessary to do so. The
169
+ Licensor waives and/or agrees not to assert any right or
170
+ authority to forbid You from making technical modifications
171
+ necessary to exercise the Licensed Rights, including
172
+ technical modifications necessary to circumvent Effective
173
+ Technological Measures. For purposes of this Public License,
174
+ simply making modifications authorized by this Section 2(a)
175
+ (4) never produces Adapted Material.
176
+
177
+ 5. Downstream recipients.
178
+
179
+ a. Offer from the Licensor -- Licensed Material. Every
180
+ recipient of the Licensed Material automatically
181
+ receives an offer from the Licensor to exercise the
182
+ Licensed Rights under the terms and conditions of this
183
+ Public License.
184
+
185
+ b. No downstream restrictions. You may not offer or impose
186
+ any additional or different terms or conditions on, or
187
+ apply any Effective Technological Measures to, the
188
+ Licensed Material if doing so restricts exercise of the
189
+ Licensed Rights by any recipient of the Licensed
190
+ Material.
191
+
192
+ 6. No endorsement. Nothing in this Public License constitutes or
193
+ may be construed as permission to assert or imply that You
194
+ are, or that Your use of the Licensed Material is, connected
195
+ with, or sponsored, endorsed, or granted official status by,
196
+ the Licensor or others designated to receive attribution as
197
+ provided in Section 3(a)(1)(A)(i).
198
+
199
+ b. Other rights.
200
+
201
+ 1. Moral rights, such as the right of integrity, are not
202
+ licensed under this Public License, nor are publicity,
203
+ privacy, and/or other similar personality rights; however, to
204
+ the extent possible, the Licensor waives and/or agrees not to
205
+ assert any such rights held by the Licensor to the limited
206
+ extent necessary to allow You to exercise the Licensed
207
+ Rights, but not otherwise.
208
+
209
+ 2. Patent and trademark rights are not licensed under this
210
+ Public License.
211
+
212
+ 3. To the extent possible, the Licensor waives any right to
213
+ collect royalties from You for the exercise of the Licensed
214
+ Rights, whether directly or through a collecting society
215
+ under any voluntary or waivable statutory or compulsory
216
+ licensing scheme. In all other cases the Licensor expressly
217
+ reserves any right to collect such royalties, including when
218
+ the Licensed Material is used other than for NonCommercial
219
+ purposes.
220
+
221
+
222
+ Section 3 -- License Conditions.
223
+
224
+ Your exercise of the Licensed Rights is expressly made subject to the
225
+ following conditions.
226
+
227
+ a. Attribution.
228
+
229
+ 1. If You Share the Licensed Material (including in modified
230
+ form), You must:
231
+
232
+ a. retain the following if it is supplied by the Licensor
233
+ with the Licensed Material:
234
+
235
+ i. identification of the creator(s) of the Licensed
236
+ Material and any others designated to receive
237
+ attribution, in any reasonable manner requested by
238
+ the Licensor (including by pseudonym if
239
+ designated);
240
+
241
+ ii. a copyright notice;
242
+
243
+ iii. a notice that refers to this Public License;
244
+
245
+ iv. a notice that refers to the disclaimer of
246
+ warranties;
247
+
248
+ v. a URI or hyperlink to the Licensed Material to the
249
+ extent reasonably practicable;
250
+
251
+ b. indicate if You modified the Licensed Material and
252
+ retain an indication of any previous modifications; and
253
+
254
+ c. indicate the Licensed Material is licensed under this
255
+ Public License, and include the text of, or the URI or
256
+ hyperlink to, this Public License.
257
+
258
+ 2. You may satisfy the conditions in Section 3(a)(1) in any
259
+ reasonable manner based on the medium, means, and context in
260
+ which You Share the Licensed Material. For example, it may be
261
+ reasonable to satisfy the conditions by providing a URI or
262
+ hyperlink to a resource that includes the required
263
+ information.
264
+
265
+ 3. If requested by the Licensor, You must remove any of the
266
+ information required by Section 3(a)(1)(A) to the extent
267
+ reasonably practicable.
268
+
269
+ 4. If You Share Adapted Material You produce, the Adapter's
270
+ License You apply must not prevent recipients of the Adapted
271
+ Material from complying with this Public License.
272
+
273
+
274
+ Section 4 -- Sui Generis Database Rights.
275
+
276
+ Where the Licensed Rights include Sui Generis Database Rights that
277
+ apply to Your use of the Licensed Material:
278
+
279
+ a. for the avoidance of doubt, Section 2(a)(1) grants You the right
280
+ to extract, reuse, reproduce, and Share all or a substantial
281
+ portion of the contents of the database for NonCommercial purposes
282
+ only;
283
+
284
+ b. if You include all or a substantial portion of the database
285
+ contents in a database in which You have Sui Generis Database
286
+ Rights, then the database in which You have Sui Generis Database
287
+ Rights (but not its individual contents) is Adapted Material; and
288
+
289
+ c. You must comply with the conditions in Section 3(a) if You Share
290
+ all or a substantial portion of the contents of the database.
291
+
292
+ For the avoidance of doubt, this Section 4 supplements and does not
293
+ replace Your obligations under this Public License where the Licensed
294
+ Rights include other Copyright and Similar Rights.
295
+
296
+
297
+ Section 5 -- Disclaimer of Warranties and Limitation of Liability.
298
+
299
+ a. UNLESS OTHERWISE SEPARATELY UNDERTAKEN BY THE LICENSOR, TO THE
300
+ EXTENT POSSIBLE, THE LICENSOR OFFERS THE LICENSED MATERIAL AS-IS
301
+ AND AS-AVAILABLE, AND MAKES NO REPRESENTATIONS OR WARRANTIES OF
302
+ ANY KIND CONCERNING THE LICENSED MATERIAL, WHETHER EXPRESS,
303
+ IMPLIED, STATUTORY, OR OTHER. THIS INCLUDES, WITHOUT LIMITATION,
304
+ WARRANTIES OF TITLE, MERCHANTABILITY, FITNESS FOR A PARTICULAR
305
+ PURPOSE, NON-INFRINGEMENT, ABSENCE OF LATENT OR OTHER DEFECTS,
306
+ ACCURACY, OR THE PRESENCE OR ABSENCE OF ERRORS, WHETHER OR NOT
307
+ KNOWN OR DISCOVERABLE. WHERE DISCLAIMERS OF WARRANTIES ARE NOT
308
+ ALLOWED IN FULL OR IN PART, THIS DISCLAIMER MAY NOT APPLY TO YOU.
309
+
310
+ b. TO THE EXTENT POSSIBLE, IN NO EVENT WILL THE LICENSOR BE LIABLE
311
+ TO YOU ON ANY LEGAL THEORY (INCLUDING, WITHOUT LIMITATION,
312
+ NEGLIGENCE) OR OTHERWISE FOR ANY DIRECT, SPECIAL, INDIRECT,
313
+ INCIDENTAL, CONSEQUENTIAL, PUNITIVE, EXEMPLARY, OR OTHER LOSSES,
314
+ COSTS, EXPENSES, OR DAMAGES ARISING OUT OF THIS PUBLIC LICENSE OR
315
+ USE OF THE LICENSED MATERIAL, EVEN IF THE LICENSOR HAS BEEN
316
+ ADVISED OF THE POSSIBILITY OF SUCH LOSSES, COSTS, EXPENSES, OR
317
+ DAMAGES. WHERE A LIMITATION OF LIABILITY IS NOT ALLOWED IN FULL OR
318
+ IN PART, THIS LIMITATION MAY NOT APPLY TO YOU.
319
+
320
+ c. The disclaimer of warranties and limitation of liability provided
321
+ above shall be interpreted in a manner that, to the extent
322
+ possible, most closely approximates an absolute disclaimer and
323
+ waiver of all liability.
324
+
325
+
326
+ Section 6 -- Term and Termination.
327
+
328
+ a. This Public License applies for the term of the Copyright and
329
+ Similar Rights licensed here. However, if You fail to comply with
330
+ this Public License, then Your rights under this Public License
331
+ terminate automatically.
332
+
333
+ b. Where Your right to use the Licensed Material has terminated under
334
+ Section 6(a), it reinstates:
335
+
336
+ 1. automatically as of the date the violation is cured, provided
337
+ it is cured within 30 days of Your discovery of the
338
+ violation; or
339
+
340
+ 2. upon express reinstatement by the Licensor.
341
+
342
+ For the avoidance of doubt, this Section 6(b) does not affect any
343
+ right the Licensor may have to seek remedies for Your violations
344
+ of this Public License.
345
+
346
+ c. For the avoidance of doubt, the Licensor may also offer the
347
+ Licensed Material under separate terms or conditions or stop
348
+ distributing the Licensed Material at any time; however, doing so
349
+ will not terminate this Public License.
350
+
351
+ d. Sections 1, 5, 6, 7, and 8 survive termination of this Public
352
+ License.
353
+
354
+
355
+ Section 7 -- Other Terms and Conditions.
356
+
357
+ a. The Licensor shall not be bound by any additional or different
358
+ terms or conditions communicated by You unless expressly agreed.
359
+
360
+ b. Any arrangements, understandings, or agreements regarding the
361
+ Licensed Material not stated herein are separate from and
362
+ independent of the terms and conditions of this Public License.
363
+
364
+
365
+ Section 8 -- Interpretation.
366
+
367
+ a. For the avoidance of doubt, this Public License does not, and
368
+ shall not be interpreted to, reduce, limit, restrict, or impose
369
+ conditions on any use of the Licensed Material that could lawfully
370
+ be made without permission under this Public License.
371
+
372
+ b. To the extent possible, if any provision of this Public License is
373
+ deemed unenforceable, it shall be automatically reformed to the
374
+ minimum extent necessary to make it enforceable. If the provision
375
+ cannot be reformed, it shall be severed from this Public License
376
+ without affecting the enforceability of the remaining terms and
377
+ conditions.
378
+
379
+ c. No term or condition of this Public License will be waived and no
380
+ failure to comply consented to unless expressly agreed to by the
381
+ Licensor.
382
+
383
+ d. Nothing in this Public License constitutes or may be interpreted
384
+ as a limitation upon, or waiver of, any privileges and immunities
385
+ that apply to the Licensor or You, including from the legal
386
+ processes of any jurisdiction or authority.
387
+
388
+ =======================================================================
389
+
390
+ Creative Commons is not a party to its public
391
+ licenses. Notwithstanding, Creative Commons may elect to apply one of
392
+ its public licenses to material it publishes and in those instances
393
+ will be considered the “Licensor.” The text of the Creative Commons
394
+ public licenses is dedicated to the public domain under the CC0 Public
395
+ Domain Dedication. Except for the limited purpose of indicating that
396
+ material is shared under a Creative Commons public license or as
397
+ otherwise permitted by the Creative Commons policies published at
398
+ creativecommons.org/policies, Creative Commons does not authorize the
399
+ use of the trademark "Creative Commons" or any other trademark or logo
400
+ of Creative Commons without its prior written consent including,
401
+ without limitation, in connection with any unauthorized modifications
402
+ to any of its public licenses or any other arrangements,
403
+ understandings, or agreements concerning use of licensed material. For
404
+ the avoidance of doubt, this paragraph does not form part of the
405
+ public licenses.
406
+
407
+ Creative Commons may be contacted at creativecommons.org.
408
+
409
+ Implementation of the `LayerNorm` operators (in protenix/model/layer_norm/) referred to [OneFlow]
410
+ (https://github.com/Oneflow-Inc/oneflow) and [FastFold](https://github.com/hpcaitech/FastFold).
411
+ We used [OpenFold](https://github.com/aqlaboratory/openfold) for some
412
+ (in protenix/openfold_local/) implementations, except the `LayerNorm` part. the worker OneFlow,
413
+ FastFold and openfold are licensed under Apache License 2.0.
@@ -0,0 +1,11 @@
1
+ Metadata-Version: 2.1
2
+ Name: protenix
3
+ Version: 0.1.0
4
+ Summary: A trainable PyTorch reproduction of AlphaFold 3.
5
+ Home-page: https://github.com/bytedance/Protenix
6
+ Author: Bytedance Inc.
7
+ Author-email: ai4s-bio@bytedance.com
8
+ License: Attribution-NonCommercial 4.0 International
9
+ Platform: manylinux1
10
+ Requires-Python: >=3.10
11
+ License-File: LICENSE
@@ -0,0 +1,255 @@
1
+ # Protenix: Protein + X
2
+
3
+ A trainable PyTorch reproduction of [AlphaFold 3](https://www.nature.com/articles/s41586-024-07487-w).
4
+
5
+ For more information on the model's performance and capabilities, see our [technical report](Protenix_Technical_Report.pdf).
6
+
7
+ You can follow our [twitter](https://x.com/ai4s_protenix) or join the conversation in the [discord server](https://discord.gg/Utgk4Ykw).
8
+
9
+ ![Protenix predictions](assets/protenix_predictions.gif)
10
+
11
+ ## ⚡ Try it online
12
+ - [Web server link](http://101.126.11.40:8000/)
13
+
14
+ ## Installation and Preparations
15
+
16
+ ### Installing Protenix
17
+
18
+ Follow these steps to set up and run Protenix:
19
+
20
+ #### Run with Docker
21
+
22
+ 1. Install Docker (with GPU Support)
23
+ Ensure that Docker is installed and configured with GPU support. Follow these steps:
24
+ - Install [Docker](https://www.docker.com/) if not already installed.
25
+ * Install the [NVIDIA Container Toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html) to enable GPU support.
26
+ * Verify the setup with:
27
+ ```bash
28
+ docker run --rm --gpus all nvidia/cuda:12.1.0-base-ubuntu22.04 nvidia-smi
29
+ ```
30
+
31
+ 1. Pull the Docker image, which was built based on this [Dockerfile](Dockerfile)
32
+ ```bash
33
+ docker pull ai4s-cn-beijing.cr.volces.com/infra/protenix:v0.0.1
34
+ ```
35
+
36
+ 1. Clone this repository and `cd` into it
37
+ ```bash
38
+ git clone https://github.com/bytedance/protenix.git
39
+ cd ./protenix
40
+ pip install -e .
41
+ ```
42
+
43
+ 1. Run Docker with an interactive shell
44
+ ```bash
45
+ docker run --gpus all -it -v $(pwd):/workspace -v /dev/shm:/dev/shm ai4s-cn-beijing.cr.volces.com/infra/protenix:v0.0.1 /bin/bash
46
+ ```
47
+
48
+ After running above commands, you’ll be inside the container’s environment and can execute commands as you would on a normal Linux terminal.
49
+
50
+ #### Run directly
51
+
52
+ ```bash
53
+ # maybe you need to update libxrender1 and libxext6 firstly, run as following for Debian:
54
+ # apt-get update
55
+ # apt-get install libxrender1
56
+ # apt-get install libxext6
57
+ pip3 install protenix
58
+ ```
59
+
60
+ ### Setting up kernels
61
+
62
+ - **Custom CUDA layernorm kernels** modified from [FastFold](https://github.com/hpcaitech/FastFold) and [Oneflow](https://github.com/Oneflow-Inc/oneflow) accelerate about 30%-50% during different training stages. To use this feature, run the following command:
63
+ ```bash
64
+ export LAYERNORM_TYPE=fast_layernorm
65
+ ```
66
+ If the environment variable `LAYERNORM_TYPE` is set to `fast_layernorm`, the model will employ the layernorm we have developed; otherwise, the naive PyTorch layernorm will be adopted. The kernels will be compiled when `fast_layernorm` is called for the first time.
67
+ - **[DeepSpeed DS4Sci_EvoformerAttention kernel](https://www.deepspeed.ai/tutorials/ds4sci_evoformerattention/)** is a memory-efficient attention kernel developed as part of a collaboration between OpenFold and the DeepSpeed4Science initiative. To use this feature, simply pass:
68
+ ```bash
69
+ --use_deepspeed_evo_attention true
70
+ ```
71
+ into the command line. DS4Sci_EvoformerAttention is implemented based on [CUTLASS](https://github.com/NVIDIA/cutlass). You need to clone the CUTLASS repository and specify the path to it in the environment variable CUTLASS_PATH. The [Dockerfile](Dockerfile) has already include this setting:
72
+ ```bash
73
+ RUN git clone -b v3.5.1 https://github.com/NVIDIA/cutlass.git /opt/cutlass
74
+ ENV CUTLASS_PATH=/opt/cutlass
75
+ ```
76
+ The kernels will be compiled when DS4Sci_EvoformerAttention is called for the first time.
77
+
78
+
79
+ ### Preparing the datasets
80
+ To download the [wwPDB dataset](https://www.wwpdb.org/) and proprecessed training data, you need at least 1T disk space.
81
+
82
+ Use the following command to download the preprocessed wwpdb training databases:
83
+
84
+ ```bash
85
+ wget -P /af3-dev/release_data/ https://af3-dev.tos-cn-beijing.volces.com/release_data.tar.gz
86
+ tar -xzvf /af3-dev/release_data/release_data.tar.gz -C /af3-dev/release_data/
87
+ rm /af3-dev/release_data/release_data.tar.gz
88
+ ```
89
+
90
+
91
+ The data should be placed in the `/af3-dev/release_data/` directory. You can also download it to a different directory, but remember to modify the `DATA_ROOT_DIR` in [configs/configs_data.py](configs/configs_data.py) correspondingly. Data hierarchy after extraction is as follows:
92
+
93
+ ```bash
94
+ ├── components.v20240608.cif [408M] # ccd source file
95
+ ├── components.v20240608.cif.rdkit_mol.pkl [121M] # rdkit Mol object generated by ccd source file
96
+ ├── indices [33M] # chain or interface entries
97
+ ├── mmcif [283G] # raw mmcif data
98
+ ├── mmcif_bioassembly [36G] # preprocessed wwPDB structural data
99
+ ├── mmcif_msa [450G] # msa files
100
+ ├── posebusters_bioassembly [42M] # preprocessed posebusters structural data
101
+ ├── posebusters_mmcif [361M] # raw mmcif data
102
+ ├── recentPDB_bioassembly [1.5G] # preprocessed recentPDB structural data
103
+ └── seq_to_pdb_index.json [45M] # sequence to pdb id mapping file
104
+ ```
105
+ With the above data, you can run the training demo from scratch. `components.v20240608.cif` and `components.v20240608.cif.rdkit_mol.pkl` is also used in inference pipeline for generating [ccd](https://www.wwpdb.org/data/ccd) reference feature. If you only want to run inference, the full released data is not necessary, you can download these two files separately.
106
+ ```bash
107
+ wget -P /af3-dev/release_data/ https://af3-dev.tos-cn-beijing.volces.com/release_data/components.v20240608.cif
108
+ wget -P /af3-dev/release_data/ https://af3-dev.tos-cn-beijing.volces.com/release_data/components.v20240608.cif.rdkit_mol.pkl
109
+ ```
110
+ Data processing scripts are still being organized and prepared, and distillation data will be released in the future.
111
+
112
+ ## Running your first prediction
113
+
114
+ ### Model checkpoints
115
+
116
+ Use the following command to download pretrained checkpoint [1.4G]:
117
+
118
+ ```bash
119
+ wget -P /af3-dev/release_model/ https://af3-dev.tos-cn-beijing.volces.com/release_model/model_v1.pt
120
+
121
+ ```
122
+ the checkpoint should be placed in the `/af3-dev/release_model/` directory.
123
+
124
+ ### Notebook demo
125
+ You can use [notebooks/protenix_inference.ipynb](notebooks/protenix_inference.ipynb) to run the model inference.
126
+
127
+ ### Inference demo
128
+ You can run the script `inference_demo.sh` to do model inference:
129
+
130
+ ```bash
131
+ bash inference_demo.sh
132
+ ```
133
+
134
+ Arguments in this scripts are explained as follows:
135
+ * `load_checkpoint_path`: path to the model checkpoints.
136
+ * `input_json_path`: path to a JSON file that fully describes the input.
137
+ * `dump_dir`: path to a directory where the results of the inference will be saved.
138
+ * `dtype`: data type used in inference. Valid options include `"bf16"` and `"fp32"`.
139
+ * `use_deepspeed_evo_attention`: whether use the EvoformerAttention provided by DeepSpeed.
140
+ * `use_msa`: whether to use the MSA feature, the default is true. If you want to disable the MSA feature, add `--use_msa false` to the [inference_demo.sh](inference_demo.sh) script.
141
+
142
+ or you can run inference with:
143
+
144
+ ```bash
145
+ # run with examples floder
146
+ protenix_infer --input_json_path examples/example.json --dump_dir ./output
147
+ ```
148
+
149
+ **Detailed information on the format of the input JSON file and the output files can be found [here](runner/infer_json_format.md)**.
150
+
151
+ Predicted structures for the posebusters set are available at:
152
+ ```bash
153
+ https://af3-dev.tos-cn-beijing.volces.com/pb_samples_release.tar.gz
154
+ ```
155
+
156
+ ## Training and Finetuning
157
+ ### Training demo
158
+ After the installation and data preparations, you can run the following command to train the model from scratch:
159
+
160
+ ```bash
161
+ bash train_demo.sh
162
+ ```
163
+ Key arguments in this scripts are explained as follows:
164
+ * `dtype`: data type used in training. Valid options include `"bf16"` and `"fp32"`.
165
+ * `--dtype fp32`: the model will be trained in full FP32 precision.
166
+ * `--dtype bf16`: the model will be trained in BF16 Mixed precision, by default, the `SampleDiffusion`,`ConfidenceHead`, `Mini-rollout` and `Loss` part will still be training in FP32 precision. if you want to train and infer the model in full BF16 Mixed precision, pass the following arguments to the [train_demo.sh](train_demo.sh):
167
+ ```bash
168
+ --skip_amp.sample_diffusion_training false \
169
+ --skip_amp.confidence_head false \
170
+ --skip_amp.sample_diffusion false \
171
+ --skip_amp.loss false \
172
+ ```
173
+ * `use_deepspeed_evo_attention`: whether use the EvoformerAttention provided by DeepSpeed as mentioned above.
174
+ * `ema_decay`: the decay rate of the EMA, default is 0.999.
175
+ * `sample_diffusion.N_step`: during evalutaion, the number of steps for the diffusion process is reduced to 20 to improve efficiency.
176
+ * `data.train_sets/data.test_sets`: the datasets used for training and evaluation. If there are multiple datasets, separate them with commas.
177
+
178
+ * Some settings follow those in the [AlphaFold 3](https://www.nature.com/articles/s41586-024-07487-w) paper, The table below shows the training settings for different fine-tuning stages:
179
+
180
+ | Arguments | Initial training | Fine tuning 1 | Fine tuning 2 | Fine tuning 3 |
181
+ |-----------------------------------------|--------|---------|-------|-----|
182
+ | `train_crop_size` | 384 | 640 | 768 | 768 |
183
+ | `diffusion_batch_size` | 48 | 32 | 32 | 32 |
184
+ | `loss.weight.alpha_pae` | 0 | 0 | 0 | 1.0 |
185
+ | `loss.weight.alpha_bond` | 0 | 1.0 | 1.0 | 0 |
186
+ | `loss.weight.smooth_lddt` | 1.0 | 0 | 0 | 0 |
187
+ | `loss.weight.alpha_confidence` | 1e-4 | 1e-4 | 1e-4 | 1e-4|
188
+ | `loss.weight.alpha_diffusion` | 4.0 | 4.0 | 4.0 | 0 |
189
+ | `loss.weight.alpha_distogram` | 0.03 | 0.03 | 0.03 | 0 |
190
+ | `train_confidence_only` | False | False | False | True|
191
+ | full BF16-mixed speed(A100, s/step) | ~12 | ~30 | ~44 | ~13 |
192
+ | full BF16-mixed peak memory (G) | ~34 | ~35 | ~48 | ~24 |
193
+
194
+ We recommend carrying out the training on A100-80G or H20/H100 GPUs. If utilizing full BF16-Mixed precision training, the initial training stage can also be performed on A800-40G GPUs. GPUs with smaller memory, such as A30, you'll need to reduce the model size, such as decreasing `model.pairformer.nblocks` and `diffusion_batch_size`.
195
+ * In this version, we do not use the template and RNA MSA feature for training. As the default settings in [configs/configs_base.py](configs/configs_base.py) and [configs/configs_data.py](configs/configs_data.py):
196
+ ```bash
197
+ --model.template_embedder.n_blocks 0 \
198
+ --data.msa.enable_rna_msa false \
199
+ ```
200
+ This will be considered in our future work.
201
+
202
+ * The model also supports distributed training with PyTorch’s [`torchrun`](https://pytorch.org/docs/stable/elastic/run.html). For example, if you’re running distributed training on a single node with 4 GPUs, you can use:
203
+ ```bash
204
+ torchrun --nproc_per_node=4 runner/train.py
205
+ ```
206
+ You can also pass other arguments with `--<ARGS_KEY> <ARGS_VALUE>` as you want.
207
+
208
+
209
+ ### Finetune demo
210
+
211
+ If you want to fine-tune the model on a specific subset, such as an antibody dataset, you only need to provide a PDB list file and load the pretrained weights as [finetune_demo.sh](finetune_demo.sh) shows:
212
+
213
+ ```bash
214
+ checkpoint_path="/af3-dev/release_model/model_v1.pt"
215
+ ...
216
+
217
+ --load_checkpoint_path ${checkpoint_path} \
218
+ --load_checkpoint_ema_path ${checkpoint_path} \
219
+ --data.weightedPDB_before2109_wopb_nometalc_0925.base_info.pdb_list examples/subset.txt \
220
+ ```
221
+
222
+ , where the `subset.txt` is a file containing the PDB IDs like:
223
+ ```bash
224
+ 6hvq
225
+ 5mqc
226
+ 5zin
227
+ 3ew0
228
+ 5akv
229
+ ```
230
+
231
+ ## Acknowledgements
232
+
233
+ Implementation of the layernorm operators referred to [OneFlow](https://github.com/Oneflow-Inc/oneflow) and [FastFold](https://github.com/hpcaitech/FastFold). We used [OpenFold](https://github.com/aqlaboratory/openfold) for some [module](protenix/openfold_local/) implementations, except the [`LayerNorm`](protenix/model/layer_norm/).
234
+
235
+
236
+ ## Contribution
237
+
238
+ Please check [Contributing](CONTRIBUTING.md) for more details. If you encounter problems using Protenix, feel free to create an issue! We also welcome pull requests from the community.
239
+
240
+ ## Code of Conduct
241
+
242
+ Please check [Code of Conduct](CODE_OF_CONDUCT.md) for more details.
243
+
244
+ ## Security
245
+
246
+ If you discover a potential security issue in this project, or think you may
247
+ have discovered a security issue, we ask that you notify Bytedance Security via our [security center](https://security.bytedance.com/src) or [vulnerability reporting email](sec@bytedance.com).
248
+
249
+ Please do **not** create a public GitHub issue.
250
+
251
+ ## License
252
+
253
+ This project, including code and model parameters are made available under the terms of the Creative Commons Attribution-NonCommercial 4.0 International License. You can find details at: https://creativecommons.org/licenses/by-nc/4.0/
254
+
255
+ For commercial use, please reach out to us at ai4s-bio@bytedance.com for the commercial license. We welcome all types of collaborations.
File without changes