rc-foundry 0.1.1__py3-none-any.whl

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 (180) hide show
  1. foundry/__init__.py +57 -0
  2. foundry/callbacks/__init__.py +5 -0
  3. foundry/callbacks/callback.py +116 -0
  4. foundry/callbacks/health_logging.py +419 -0
  5. foundry/callbacks/metrics_logging.py +211 -0
  6. foundry/callbacks/timing_logging.py +67 -0
  7. foundry/callbacks/train_logging.py +278 -0
  8. foundry/common.py +108 -0
  9. foundry/constants.py +28 -0
  10. foundry/hydra/resolvers.py +77 -0
  11. foundry/inference_engines/base.py +235 -0
  12. foundry/inference_engines/checkpoint_registry.py +66 -0
  13. foundry/metrics/__init__.py +12 -0
  14. foundry/metrics/losses.py +30 -0
  15. foundry/metrics/metric.py +319 -0
  16. foundry/model/layers/blocks.py +47 -0
  17. foundry/testing/__init__.py +6 -0
  18. foundry/testing/fixtures.py +19 -0
  19. foundry/testing/pytest_hooks.py +15 -0
  20. foundry/trainers/fabric.py +923 -0
  21. foundry/training/EMA.py +67 -0
  22. foundry/training/checkpoint.py +61 -0
  23. foundry/training/schedulers.py +91 -0
  24. foundry/utils/alignment.py +86 -0
  25. foundry/utils/components.py +415 -0
  26. foundry/utils/datasets.py +405 -0
  27. foundry/utils/ddp.py +103 -0
  28. foundry/utils/instantiators.py +72 -0
  29. foundry/utils/logging.py +279 -0
  30. foundry/utils/rigid.py +1460 -0
  31. foundry/utils/rotation_augmentation.py +65 -0
  32. foundry/utils/squashfs.py +172 -0
  33. foundry/utils/torch.py +317 -0
  34. foundry/utils/weights.py +271 -0
  35. foundry/version.py +34 -0
  36. foundry_cli/__init__.py +3 -0
  37. foundry_cli/download_checkpoints.py +281 -0
  38. mpnn/__init__.py +1 -0
  39. mpnn/collate/feature_collator.py +265 -0
  40. mpnn/inference.py +53 -0
  41. mpnn/inference_engines/mpnn.py +549 -0
  42. mpnn/loss/nll_loss.py +122 -0
  43. mpnn/metrics/nll.py +369 -0
  44. mpnn/metrics/sequence_recovery.py +440 -0
  45. mpnn/model/layers/graph_embeddings.py +2372 -0
  46. mpnn/model/layers/message_passing.py +332 -0
  47. mpnn/model/layers/position_wise_feed_forward.py +44 -0
  48. mpnn/model/layers/positional_encoding.py +98 -0
  49. mpnn/model/mpnn.py +2632 -0
  50. mpnn/pipelines/mpnn.py +162 -0
  51. mpnn/samplers/samplers.py +167 -0
  52. mpnn/train.py +341 -0
  53. mpnn/trainers/mpnn.py +193 -0
  54. mpnn/transforms/feature_aggregation/mpnn.py +184 -0
  55. mpnn/transforms/feature_aggregation/polymer_ligand_interface.py +76 -0
  56. mpnn/transforms/feature_aggregation/token_encodings.py +132 -0
  57. mpnn/transforms/feature_aggregation/user_settings.py +347 -0
  58. mpnn/transforms/polymer_ligand_interface.py +164 -0
  59. mpnn/utils/inference.py +2397 -0
  60. mpnn/utils/probability.py +37 -0
  61. mpnn/utils/weights.py +309 -0
  62. rc_foundry-0.1.1.dist-info/METADATA +239 -0
  63. rc_foundry-0.1.1.dist-info/RECORD +180 -0
  64. rc_foundry-0.1.1.dist-info/WHEEL +4 -0
  65. rc_foundry-0.1.1.dist-info/entry_points.txt +5 -0
  66. rc_foundry-0.1.1.dist-info/licenses/LICENSE.md +28 -0
  67. rf3/__init__.py +3 -0
  68. rf3/_version.py +33 -0
  69. rf3/alignment.py +79 -0
  70. rf3/callbacks/dump_validation_structures.py +101 -0
  71. rf3/callbacks/metrics_logging.py +324 -0
  72. rf3/chemical.py +1529 -0
  73. rf3/cli.py +77 -0
  74. rf3/data/cyclic_transform.py +78 -0
  75. rf3/data/extra_xforms.py +36 -0
  76. rf3/data/ground_truth_template.py +463 -0
  77. rf3/data/paired_msa.py +206 -0
  78. rf3/data/pipeline_utils.py +128 -0
  79. rf3/data/pipelines.py +558 -0
  80. rf3/diffusion_samplers/inference_sampler.py +222 -0
  81. rf3/inference.py +65 -0
  82. rf3/inference_engines/__init__.py +5 -0
  83. rf3/inference_engines/rf3.py +735 -0
  84. rf3/kinematics.py +354 -0
  85. rf3/loss/af3_confidence_loss.py +515 -0
  86. rf3/loss/af3_losses.py +655 -0
  87. rf3/loss/loss.py +179 -0
  88. rf3/metrics/chiral.py +179 -0
  89. rf3/metrics/clashing_chains.py +68 -0
  90. rf3/metrics/distogram.py +421 -0
  91. rf3/metrics/lddt.py +523 -0
  92. rf3/metrics/metadata.py +43 -0
  93. rf3/metrics/metric_utils.py +192 -0
  94. rf3/metrics/predicted_error.py +134 -0
  95. rf3/metrics/rasa.py +108 -0
  96. rf3/metrics/selected_distances.py +91 -0
  97. rf3/model/RF3.py +527 -0
  98. rf3/model/RF3_blocks.py +92 -0
  99. rf3/model/RF3_structure.py +303 -0
  100. rf3/model/layers/af3_auxiliary_heads.py +255 -0
  101. rf3/model/layers/af3_diffusion_transformer.py +544 -0
  102. rf3/model/layers/attention.py +313 -0
  103. rf3/model/layers/layer_utils.py +127 -0
  104. rf3/model/layers/mlff.py +118 -0
  105. rf3/model/layers/outer_product.py +59 -0
  106. rf3/model/layers/pairformer_layers.py +783 -0
  107. rf3/model/layers/structure_bias.py +56 -0
  108. rf3/scoring.py +1787 -0
  109. rf3/symmetry/resolve.py +284 -0
  110. rf3/train.py +194 -0
  111. rf3/trainers/rf3.py +570 -0
  112. rf3/util_module.py +47 -0
  113. rf3/utils/frames.py +109 -0
  114. rf3/utils/inference.py +665 -0
  115. rf3/utils/io.py +198 -0
  116. rf3/utils/loss.py +72 -0
  117. rf3/utils/predict_and_score.py +165 -0
  118. rf3/utils/predicted_error.py +673 -0
  119. rf3/utils/recycling.py +42 -0
  120. rf3/validate.py +140 -0
  121. rfd3/.gitignore +7 -0
  122. rfd3/Makefile +76 -0
  123. rfd3/__init__.py +12 -0
  124. rfd3/callbacks.py +66 -0
  125. rfd3/cli.py +41 -0
  126. rfd3/constants.py +212 -0
  127. rfd3/engine.py +543 -0
  128. rfd3/inference/datasets.py +193 -0
  129. rfd3/inference/input_parsing.py +1123 -0
  130. rfd3/inference/legacy_input_parsing.py +717 -0
  131. rfd3/inference/parsing.py +165 -0
  132. rfd3/inference/symmetry/atom_array.py +298 -0
  133. rfd3/inference/symmetry/checks.py +241 -0
  134. rfd3/inference/symmetry/contigs.py +63 -0
  135. rfd3/inference/symmetry/frames.py +355 -0
  136. rfd3/inference/symmetry/symmetry_utils.py +398 -0
  137. rfd3/metrics/design_metrics.py +465 -0
  138. rfd3/metrics/hbonds_hbplus_metrics.py +308 -0
  139. rfd3/metrics/hbonds_metrics.py +389 -0
  140. rfd3/metrics/losses.py +325 -0
  141. rfd3/metrics/metrics_utils.py +118 -0
  142. rfd3/metrics/sidechain_metrics.py +349 -0
  143. rfd3/model/RFD3.py +105 -0
  144. rfd3/model/RFD3_diffusion_module.py +387 -0
  145. rfd3/model/cfg_utils.py +81 -0
  146. rfd3/model/inference_sampler.py +635 -0
  147. rfd3/model/layers/attention.py +577 -0
  148. rfd3/model/layers/block_utils.py +580 -0
  149. rfd3/model/layers/blocks.py +777 -0
  150. rfd3/model/layers/chunked_pairwise.py +377 -0
  151. rfd3/model/layers/encoders.py +417 -0
  152. rfd3/model/layers/layer_utils.py +197 -0
  153. rfd3/model/layers/pairformer_layers.py +128 -0
  154. rfd3/run_inference.py +45 -0
  155. rfd3/testing/debug.py +139 -0
  156. rfd3/testing/debug_utils.py +73 -0
  157. rfd3/testing/testing_utils.py +356 -0
  158. rfd3/train.py +194 -0
  159. rfd3/trainer/dump_validation_structures.py +154 -0
  160. rfd3/trainer/fabric_trainer.py +923 -0
  161. rfd3/trainer/recycling.py +42 -0
  162. rfd3/trainer/rfd3.py +485 -0
  163. rfd3/trainer/trainer_utils.py +502 -0
  164. rfd3/transforms/conditioning_base.py +508 -0
  165. rfd3/transforms/conditioning_utils.py +200 -0
  166. rfd3/transforms/design_transforms.py +807 -0
  167. rfd3/transforms/dna_crop.py +523 -0
  168. rfd3/transforms/hbonds.py +407 -0
  169. rfd3/transforms/hbonds_hbplus.py +246 -0
  170. rfd3/transforms/ncaa_transforms.py +153 -0
  171. rfd3/transforms/pipelines.py +632 -0
  172. rfd3/transforms/ppi_transforms.py +541 -0
  173. rfd3/transforms/rasa.py +116 -0
  174. rfd3/transforms/symmetry.py +76 -0
  175. rfd3/transforms/training_conditions.py +552 -0
  176. rfd3/transforms/util_transforms.py +498 -0
  177. rfd3/transforms/virtual_atoms.py +305 -0
  178. rfd3/utils/inference.py +648 -0
  179. rfd3/utils/io.py +245 -0
  180. rfd3/utils/vizualize.py +276 -0
rf3/chemical.py ADDED
@@ -0,0 +1,1529 @@
1
+ # TODO: Deprecate in favor of AtomWorks
2
+ # (This code is a holdover from previous frameworks)
3
+
4
+ import os
5
+
6
+ import torch
7
+
8
+ script_dir = os.path.dirname(os.path.realpath(__file__)) + "/"
9
+
10
+ NHEAVY = 23
11
+ NTOTAL = 36
12
+ NNAPROTAAS = 32
13
+ NAATOKENS = 20 + 2 + 10 + 1 + 47 # 20 AAs, UNK, MASK, 8 NAs,HIS_D, 47 atoms
14
+ costgtNA = -0.2744
15
+
16
+ aa2long = [
17
+ ( # 0 ala
18
+ " N ",
19
+ " CA ",
20
+ " C ",
21
+ " O ",
22
+ " CB ",
23
+ None,
24
+ None,
25
+ None,
26
+ None,
27
+ None,
28
+ None,
29
+ None,
30
+ None,
31
+ None,
32
+ None,
33
+ None,
34
+ None,
35
+ None,
36
+ None,
37
+ None,
38
+ None,
39
+ None,
40
+ None,
41
+ " H ",
42
+ " HA ",
43
+ "1HB ",
44
+ "2HB ",
45
+ "3HB ",
46
+ None,
47
+ None,
48
+ None,
49
+ None,
50
+ None,
51
+ None,
52
+ None,
53
+ None,
54
+ ), # 0 ala
55
+ ( # 1 arg
56
+ " N ",
57
+ " CA ",
58
+ " C ",
59
+ " O ",
60
+ " CB ",
61
+ " CG ",
62
+ " CD ",
63
+ " NE ",
64
+ " CZ ",
65
+ " NH1",
66
+ " NH2",
67
+ None,
68
+ None,
69
+ None,
70
+ None,
71
+ None,
72
+ None,
73
+ None,
74
+ None,
75
+ None,
76
+ None,
77
+ None,
78
+ None,
79
+ " H ",
80
+ " HA ",
81
+ "1HB ",
82
+ "2HB ",
83
+ "1HG ",
84
+ "2HG ",
85
+ "1HD ",
86
+ "2HD ",
87
+ " HE ",
88
+ "1HH1",
89
+ "2HH1",
90
+ "1HH2",
91
+ "2HH2",
92
+ ), # 1 arg
93
+ ( # 2 asn
94
+ " N ",
95
+ " CA ",
96
+ " C ",
97
+ " O ",
98
+ " CB ",
99
+ " CG ",
100
+ " OD1",
101
+ " ND2",
102
+ None,
103
+ None,
104
+ None,
105
+ None,
106
+ None,
107
+ None,
108
+ None,
109
+ None,
110
+ None,
111
+ None,
112
+ None,
113
+ None,
114
+ None,
115
+ None,
116
+ None,
117
+ " H ",
118
+ " HA ",
119
+ "1HB ",
120
+ "2HB ",
121
+ "1HD2",
122
+ "2HD2",
123
+ None,
124
+ None,
125
+ None,
126
+ None,
127
+ None,
128
+ None,
129
+ None,
130
+ ), # 2 asn
131
+ (
132
+ " N ",
133
+ " CA ",
134
+ " C ",
135
+ " O ",
136
+ " CB ",
137
+ " CG ",
138
+ " OD1",
139
+ " OD2",
140
+ None,
141
+ None,
142
+ None,
143
+ None,
144
+ None,
145
+ None,
146
+ None,
147
+ None,
148
+ None,
149
+ None,
150
+ None,
151
+ None,
152
+ None,
153
+ None,
154
+ None,
155
+ " H ",
156
+ " HA ",
157
+ "1HB ",
158
+ "2HB ",
159
+ None,
160
+ None,
161
+ None,
162
+ None,
163
+ None,
164
+ None,
165
+ None,
166
+ None,
167
+ None,
168
+ ), # 3 asp
169
+ (
170
+ " N ",
171
+ " CA ",
172
+ " C ",
173
+ " O ",
174
+ " CB ",
175
+ " SG ",
176
+ None,
177
+ None,
178
+ None,
179
+ None,
180
+ None,
181
+ None,
182
+ None,
183
+ None,
184
+ None,
185
+ None,
186
+ None,
187
+ None,
188
+ None,
189
+ None,
190
+ None,
191
+ None,
192
+ None,
193
+ " H ",
194
+ " HA ",
195
+ "1HB ",
196
+ "2HB ",
197
+ " HG ",
198
+ None,
199
+ None,
200
+ None,
201
+ None,
202
+ None,
203
+ None,
204
+ None,
205
+ None,
206
+ ), # 4 cys
207
+ (
208
+ " N ",
209
+ " CA ",
210
+ " C ",
211
+ " O ",
212
+ " CB ",
213
+ " CG ",
214
+ " CD ",
215
+ " OE1",
216
+ " NE2",
217
+ None,
218
+ None,
219
+ None,
220
+ None,
221
+ None,
222
+ None,
223
+ None,
224
+ None,
225
+ None,
226
+ None,
227
+ None,
228
+ None,
229
+ None,
230
+ None,
231
+ " H ",
232
+ " HA ",
233
+ "1HB ",
234
+ "2HB ",
235
+ "1HG ",
236
+ "2HG ",
237
+ "1HE2",
238
+ "2HE2",
239
+ None,
240
+ None,
241
+ None,
242
+ None,
243
+ None,
244
+ ), # 5 gln
245
+ (
246
+ " N ",
247
+ " CA ",
248
+ " C ",
249
+ " O ",
250
+ " CB ",
251
+ " CG ",
252
+ " CD ",
253
+ " OE1",
254
+ " OE2",
255
+ None,
256
+ None,
257
+ None,
258
+ None,
259
+ None,
260
+ None,
261
+ None,
262
+ None,
263
+ None,
264
+ None,
265
+ None,
266
+ None,
267
+ None,
268
+ None,
269
+ " H ",
270
+ " HA ",
271
+ "1HB ",
272
+ "2HB ",
273
+ "1HG ",
274
+ "2HG ",
275
+ None,
276
+ None,
277
+ None,
278
+ None,
279
+ None,
280
+ None,
281
+ None,
282
+ ), # 6 glu
283
+ (
284
+ " N ",
285
+ " CA ",
286
+ " C ",
287
+ " O ",
288
+ None,
289
+ None,
290
+ None,
291
+ None,
292
+ None,
293
+ None,
294
+ None,
295
+ None,
296
+ None,
297
+ None,
298
+ None,
299
+ None,
300
+ None,
301
+ None,
302
+ None,
303
+ None,
304
+ None,
305
+ None,
306
+ None,
307
+ " H ",
308
+ "1HA ",
309
+ "2HA ",
310
+ None,
311
+ None,
312
+ None,
313
+ None,
314
+ None,
315
+ None,
316
+ None,
317
+ None,
318
+ None,
319
+ None,
320
+ ), # 7 gly
321
+ (
322
+ " N ",
323
+ " CA ",
324
+ " C ",
325
+ " O ",
326
+ " CB ",
327
+ " CG ",
328
+ " ND1",
329
+ " CD2",
330
+ " CE1",
331
+ " NE2",
332
+ None,
333
+ None,
334
+ None,
335
+ None,
336
+ None,
337
+ None,
338
+ None,
339
+ None,
340
+ None,
341
+ None,
342
+ None,
343
+ None,
344
+ None,
345
+ " H ",
346
+ " HA ",
347
+ "1HB ",
348
+ "2HB ",
349
+ "2HD ",
350
+ "1HE ",
351
+ "2HE ",
352
+ None,
353
+ None,
354
+ None,
355
+ None,
356
+ None,
357
+ None,
358
+ ), # 8 his
359
+ (
360
+ " N ",
361
+ " CA ",
362
+ " C ",
363
+ " O ",
364
+ " CB ",
365
+ " CG1",
366
+ " CG2",
367
+ " CD1",
368
+ None,
369
+ None,
370
+ None,
371
+ None,
372
+ None,
373
+ None,
374
+ None,
375
+ None,
376
+ None,
377
+ None,
378
+ None,
379
+ None,
380
+ None,
381
+ None,
382
+ None,
383
+ " H ",
384
+ " HA ",
385
+ " HB ",
386
+ "1HG2",
387
+ "2HG2",
388
+ "3HG2",
389
+ "1HG1",
390
+ "2HG1",
391
+ "1HD1",
392
+ "2HD1",
393
+ "3HD1",
394
+ None,
395
+ None,
396
+ ), # 9 ile
397
+ (
398
+ " N ",
399
+ " CA ",
400
+ " C ",
401
+ " O ",
402
+ " CB ",
403
+ " CG ",
404
+ " CD1",
405
+ " CD2",
406
+ None,
407
+ None,
408
+ None,
409
+ None,
410
+ None,
411
+ None,
412
+ None,
413
+ None,
414
+ None,
415
+ None,
416
+ None,
417
+ None,
418
+ None,
419
+ None,
420
+ None,
421
+ " H ",
422
+ " HA ",
423
+ "1HB ",
424
+ "2HB ",
425
+ " HG ",
426
+ "1HD1",
427
+ "2HD1",
428
+ "3HD1",
429
+ "1HD2",
430
+ "2HD2",
431
+ "3HD2",
432
+ None,
433
+ None,
434
+ ), # 10 leu
435
+ (
436
+ " N ",
437
+ " CA ",
438
+ " C ",
439
+ " O ",
440
+ " CB ",
441
+ " CG ",
442
+ " CD ",
443
+ " CE ",
444
+ " NZ ",
445
+ None,
446
+ None,
447
+ None,
448
+ None,
449
+ None,
450
+ None,
451
+ None,
452
+ None,
453
+ None,
454
+ None,
455
+ None,
456
+ None,
457
+ None,
458
+ None,
459
+ " H ",
460
+ " HA ",
461
+ "1HB ",
462
+ "2HB ",
463
+ "1HG ",
464
+ "2HG ",
465
+ "1HD ",
466
+ "2HD ",
467
+ "1HE ",
468
+ "2HE ",
469
+ "1HZ ",
470
+ "2HZ ",
471
+ "3HZ ",
472
+ ), # 11 lys
473
+ (
474
+ " N ",
475
+ " CA ",
476
+ " C ",
477
+ " O ",
478
+ " CB ",
479
+ " CG ",
480
+ " SD ",
481
+ " CE ",
482
+ None,
483
+ None,
484
+ None,
485
+ None,
486
+ None,
487
+ None,
488
+ None,
489
+ None,
490
+ None,
491
+ None,
492
+ None,
493
+ None,
494
+ None,
495
+ None,
496
+ None,
497
+ " H ",
498
+ " HA ",
499
+ "1HB ",
500
+ "2HB ",
501
+ "1HG ",
502
+ "2HG ",
503
+ "1HE ",
504
+ "2HE ",
505
+ "3HE ",
506
+ None,
507
+ None,
508
+ None,
509
+ None,
510
+ ), # 12 met
511
+ (
512
+ " N ",
513
+ " CA ",
514
+ " C ",
515
+ " O ",
516
+ " CB ",
517
+ " CG ",
518
+ " CD1",
519
+ " CD2",
520
+ " CE1",
521
+ " CE2",
522
+ " CZ ",
523
+ None,
524
+ None,
525
+ None,
526
+ None,
527
+ None,
528
+ None,
529
+ None,
530
+ None,
531
+ None,
532
+ None,
533
+ None,
534
+ None,
535
+ " H ",
536
+ " HA ",
537
+ "1HB ",
538
+ "2HB ",
539
+ "1HD ",
540
+ "2HD ",
541
+ "1HE ",
542
+ "2HE ",
543
+ " HZ ",
544
+ None,
545
+ None,
546
+ None,
547
+ None,
548
+ ), # 13 phe
549
+ (
550
+ " N ",
551
+ " CA ",
552
+ " C ",
553
+ " O ",
554
+ " CB ",
555
+ " CG ",
556
+ " CD ",
557
+ None,
558
+ None,
559
+ None,
560
+ None,
561
+ None,
562
+ None,
563
+ None,
564
+ None,
565
+ None,
566
+ None,
567
+ None,
568
+ None,
569
+ None,
570
+ None,
571
+ None,
572
+ None,
573
+ " HA ",
574
+ "1HB ",
575
+ "2HB ",
576
+ "1HG ",
577
+ "2HG ",
578
+ "1HD ",
579
+ "2HD ",
580
+ None,
581
+ None,
582
+ None,
583
+ None,
584
+ None,
585
+ None,
586
+ ), # 14 pro
587
+ (
588
+ " N ",
589
+ " CA ",
590
+ " C ",
591
+ " O ",
592
+ " CB ",
593
+ " OG ",
594
+ None,
595
+ None,
596
+ None,
597
+ None,
598
+ None,
599
+ None,
600
+ None,
601
+ None,
602
+ None,
603
+ None,
604
+ None,
605
+ None,
606
+ None,
607
+ None,
608
+ None,
609
+ None,
610
+ None,
611
+ " H ",
612
+ " HG ",
613
+ " HA ",
614
+ "1HB ",
615
+ "2HB ",
616
+ None,
617
+ None,
618
+ None,
619
+ None,
620
+ None,
621
+ None,
622
+ None,
623
+ None,
624
+ ), # 15 ser
625
+ (
626
+ " N ",
627
+ " CA ",
628
+ " C ",
629
+ " O ",
630
+ " CB ",
631
+ " OG1",
632
+ " CG2",
633
+ None,
634
+ None,
635
+ None,
636
+ None,
637
+ None,
638
+ None,
639
+ None,
640
+ None,
641
+ None,
642
+ None,
643
+ None,
644
+ None,
645
+ None,
646
+ None,
647
+ None,
648
+ None,
649
+ " H ",
650
+ " HG1",
651
+ " HA ",
652
+ " HB ",
653
+ "1HG2",
654
+ "2HG2",
655
+ "3HG2",
656
+ None,
657
+ None,
658
+ None,
659
+ None,
660
+ None,
661
+ None,
662
+ ), # 16 thr
663
+ (
664
+ " N ",
665
+ " CA ",
666
+ " C ",
667
+ " O ",
668
+ " CB ",
669
+ " CG ",
670
+ " CD1",
671
+ " CD2",
672
+ " CE2",
673
+ " CE3",
674
+ " NE1",
675
+ " CZ2",
676
+ " CZ3",
677
+ " CH2",
678
+ None,
679
+ None,
680
+ None,
681
+ None,
682
+ None,
683
+ None,
684
+ None,
685
+ None,
686
+ None,
687
+ " H ",
688
+ " HA ",
689
+ "1HB ",
690
+ "2HB ",
691
+ "1HD ",
692
+ "1HE ",
693
+ " HZ2",
694
+ " HH2",
695
+ " HZ3",
696
+ " HE3",
697
+ None,
698
+ None,
699
+ None,
700
+ ), # 17 trp
701
+ (
702
+ " N ",
703
+ " CA ",
704
+ " C ",
705
+ " O ",
706
+ " CB ",
707
+ " CG ",
708
+ " CD1",
709
+ " CD2",
710
+ " CE1",
711
+ " CE2",
712
+ " CZ ",
713
+ " OH ",
714
+ None,
715
+ None,
716
+ None,
717
+ None,
718
+ None,
719
+ None,
720
+ None,
721
+ None,
722
+ None,
723
+ None,
724
+ None,
725
+ " H ",
726
+ " HA ",
727
+ "1HB ",
728
+ "2HB ",
729
+ "1HD ",
730
+ "1HE ",
731
+ "2HE ",
732
+ "2HD ",
733
+ " HH ",
734
+ None,
735
+ None,
736
+ None,
737
+ None,
738
+ ), # 18 tyr
739
+ (
740
+ " N ",
741
+ " CA ",
742
+ " C ",
743
+ " O ",
744
+ " CB ",
745
+ " CG1",
746
+ " CG2",
747
+ None,
748
+ None,
749
+ None,
750
+ None,
751
+ None,
752
+ None,
753
+ None,
754
+ None,
755
+ None,
756
+ None,
757
+ None,
758
+ None,
759
+ None,
760
+ None,
761
+ None,
762
+ None,
763
+ " H ",
764
+ " HA ",
765
+ " HB ",
766
+ "1HG1",
767
+ "2HG1",
768
+ "3HG1",
769
+ "1HG2",
770
+ "2HG2",
771
+ "3HG2",
772
+ None,
773
+ None,
774
+ None,
775
+ None,
776
+ ), # 19 val
777
+ (
778
+ " N ",
779
+ " CA ",
780
+ " C ",
781
+ " O ",
782
+ " CB ",
783
+ None,
784
+ None,
785
+ None,
786
+ None,
787
+ None,
788
+ None,
789
+ None,
790
+ None,
791
+ None,
792
+ None,
793
+ None,
794
+ None,
795
+ None,
796
+ None,
797
+ None,
798
+ None,
799
+ None,
800
+ None,
801
+ " H ",
802
+ " HA ",
803
+ "1HB ",
804
+ "2HB ",
805
+ "3HB ",
806
+ None,
807
+ None,
808
+ None,
809
+ None,
810
+ None,
811
+ None,
812
+ None,
813
+ None,
814
+ ), # 20 unk
815
+ (
816
+ " N ",
817
+ " CA ",
818
+ " C ",
819
+ " O ",
820
+ " CB ",
821
+ None,
822
+ None,
823
+ None,
824
+ None,
825
+ None,
826
+ None,
827
+ None,
828
+ None,
829
+ None,
830
+ None,
831
+ None,
832
+ None,
833
+ None,
834
+ None,
835
+ None,
836
+ None,
837
+ None,
838
+ None,
839
+ " H ",
840
+ " HA ",
841
+ "1HB ",
842
+ "2HB ",
843
+ "3HB ",
844
+ None,
845
+ None,
846
+ None,
847
+ None,
848
+ None,
849
+ None,
850
+ None,
851
+ None,
852
+ ), # 21 mask
853
+ (
854
+ " O4'",
855
+ " C1'",
856
+ " C2'",
857
+ " OP1",
858
+ " P ",
859
+ " OP2",
860
+ " O5'",
861
+ " C5'",
862
+ " C4'",
863
+ " C3'",
864
+ " O3'",
865
+ " N9 ",
866
+ " C4 ",
867
+ " N3 ",
868
+ " C2 ",
869
+ " N1 ",
870
+ " C6 ",
871
+ " C5 ",
872
+ " N7 ",
873
+ " C8 ",
874
+ " N6 ",
875
+ None,
876
+ None,
877
+ "H5''",
878
+ " H5'",
879
+ " H4'",
880
+ " H3'",
881
+ "H2''",
882
+ " H2'",
883
+ " H1'",
884
+ " H2 ",
885
+ " H61",
886
+ " H62",
887
+ " H8 ",
888
+ None,
889
+ None,
890
+ ), # 22 DA
891
+ (
892
+ " O4'",
893
+ " C1'",
894
+ " C2'",
895
+ " OP1",
896
+ " P ",
897
+ " OP2",
898
+ " O5'",
899
+ " C5'",
900
+ " C4'",
901
+ " C3'",
902
+ " O3'",
903
+ " N1 ",
904
+ " C2 ",
905
+ " O2 ",
906
+ " N3 ",
907
+ " C4 ",
908
+ " N4 ",
909
+ " C5 ",
910
+ " C6 ",
911
+ None,
912
+ None,
913
+ None,
914
+ None,
915
+ "H5''",
916
+ " H5'",
917
+ " H4'",
918
+ " H3'",
919
+ "H2''",
920
+ " H2'",
921
+ " H1'",
922
+ " H42",
923
+ " H41",
924
+ " H5 ",
925
+ " H6 ",
926
+ None,
927
+ None,
928
+ ), # 23 DC
929
+ (
930
+ " O4'",
931
+ " C1'",
932
+ " C2'",
933
+ " OP1",
934
+ " P ",
935
+ " OP2",
936
+ " O5'",
937
+ " C5'",
938
+ " C4'",
939
+ " C3'",
940
+ " O3'",
941
+ " N9 ",
942
+ " C4 ",
943
+ " N3 ",
944
+ " C2 ",
945
+ " N1 ",
946
+ " C6 ",
947
+ " C5 ",
948
+ " N7 ",
949
+ " C8 ",
950
+ " N2 ",
951
+ " O6 ",
952
+ None,
953
+ "H5''",
954
+ " H5'",
955
+ " H4'",
956
+ " H3'",
957
+ "H2''",
958
+ " H2'",
959
+ " H1'",
960
+ " H1 ",
961
+ " H22",
962
+ " H21",
963
+ " H8 ",
964
+ None,
965
+ None,
966
+ ), # 24 DG
967
+ (
968
+ " O4'",
969
+ " C1'",
970
+ " C2'",
971
+ " OP1",
972
+ " P ",
973
+ " OP2",
974
+ " O5'",
975
+ " C5'",
976
+ " C4'",
977
+ " C3'",
978
+ " O3'",
979
+ " N1 ",
980
+ " C2 ",
981
+ " O2 ",
982
+ " N3 ",
983
+ " C4 ",
984
+ " O4 ",
985
+ " C5 ",
986
+ " C7 ",
987
+ " C6 ",
988
+ None,
989
+ None,
990
+ None,
991
+ "H5''",
992
+ " H5'",
993
+ " H4'",
994
+ " H3'",
995
+ "H2''",
996
+ " H2'",
997
+ " H1'",
998
+ " H3 ",
999
+ " H71",
1000
+ " H72",
1001
+ " H73",
1002
+ " H6 ",
1003
+ None,
1004
+ ), # 25 DT
1005
+ (
1006
+ " O4'",
1007
+ " C1'",
1008
+ " C2'",
1009
+ " OP1",
1010
+ " P ",
1011
+ " OP2",
1012
+ " O5'",
1013
+ " C5'",
1014
+ " C4'",
1015
+ " C3'",
1016
+ " O3'",
1017
+ None,
1018
+ None,
1019
+ None,
1020
+ None,
1021
+ None,
1022
+ None,
1023
+ None,
1024
+ None,
1025
+ None,
1026
+ None,
1027
+ None,
1028
+ None,
1029
+ "H5''",
1030
+ " H5'",
1031
+ " H4'",
1032
+ " H3'",
1033
+ "H2''",
1034
+ " H2'",
1035
+ " H1'",
1036
+ None,
1037
+ None,
1038
+ None,
1039
+ None,
1040
+ None,
1041
+ None,
1042
+ ), # 26 DX (unk DNA)
1043
+ (
1044
+ " O4'",
1045
+ " C1'",
1046
+ " C2'",
1047
+ " OP1",
1048
+ " P ",
1049
+ " OP2",
1050
+ " O5'",
1051
+ " C5'",
1052
+ " C4'",
1053
+ " C3'",
1054
+ " O3'",
1055
+ " O2'",
1056
+ " N1 ",
1057
+ " C2 ",
1058
+ " N3 ",
1059
+ " C4 ",
1060
+ " C5 ",
1061
+ " C6 ",
1062
+ " N6 ",
1063
+ " N7 ",
1064
+ " C8 ",
1065
+ " N9 ",
1066
+ None,
1067
+ " H5'",
1068
+ "H5''",
1069
+ " H4'",
1070
+ " H3'",
1071
+ " H2'",
1072
+ "HO2'",
1073
+ " H1'",
1074
+ " H2 ",
1075
+ " H61",
1076
+ " H62",
1077
+ " H8 ",
1078
+ None,
1079
+ None,
1080
+ ), # 27 A
1081
+ (
1082
+ " O4'",
1083
+ " C1'",
1084
+ " C2'",
1085
+ " OP1",
1086
+ " P ",
1087
+ " OP2",
1088
+ " O5'",
1089
+ " C5'",
1090
+ " C4'",
1091
+ " C3'",
1092
+ " O3'",
1093
+ " O2'",
1094
+ " N1 ",
1095
+ " C2 ",
1096
+ " O2 ",
1097
+ " N3 ",
1098
+ " C4 ",
1099
+ " N4 ",
1100
+ " C5 ",
1101
+ " C6 ",
1102
+ None,
1103
+ None,
1104
+ None,
1105
+ " H5'",
1106
+ "H5''",
1107
+ " H4'",
1108
+ " H3'",
1109
+ " H2'",
1110
+ "HO2'",
1111
+ " H1'",
1112
+ " H42",
1113
+ " H41",
1114
+ " H5 ",
1115
+ " H6 ",
1116
+ None,
1117
+ None,
1118
+ ), # 28 C
1119
+ (
1120
+ " O4'",
1121
+ " C1'",
1122
+ " C2'",
1123
+ " OP1",
1124
+ " P ",
1125
+ " OP2",
1126
+ " O5'",
1127
+ " C5'",
1128
+ " C4'",
1129
+ " C3'",
1130
+ " O3'",
1131
+ " O2'",
1132
+ " N1 ",
1133
+ " C2 ",
1134
+ " N2 ",
1135
+ " N3 ",
1136
+ " C4 ",
1137
+ " C5 ",
1138
+ " C6 ",
1139
+ " O6 ",
1140
+ " N7 ",
1141
+ " C8 ",
1142
+ " N9 ",
1143
+ " H5'",
1144
+ "H5''",
1145
+ " H4'",
1146
+ " H3'",
1147
+ " H2'",
1148
+ "HO2'",
1149
+ " H1'",
1150
+ " H1 ",
1151
+ " H22",
1152
+ " H21",
1153
+ " H8 ",
1154
+ None,
1155
+ None,
1156
+ ), # 29 G
1157
+ (
1158
+ " O4'",
1159
+ " C1'",
1160
+ " C2'",
1161
+ " OP1",
1162
+ " P ",
1163
+ " OP2",
1164
+ " O5'",
1165
+ " C5'",
1166
+ " C4'",
1167
+ " C3'",
1168
+ " O3'",
1169
+ " O2'",
1170
+ " N1 ",
1171
+ " C2 ",
1172
+ " O2 ",
1173
+ " N3 ",
1174
+ " C4 ",
1175
+ " O4 ",
1176
+ " C5 ",
1177
+ " C6 ",
1178
+ None,
1179
+ None,
1180
+ None,
1181
+ " H5'",
1182
+ "H5''",
1183
+ " H4'",
1184
+ " H3'",
1185
+ " H2'",
1186
+ "HO2'",
1187
+ " H1'",
1188
+ " H3 ",
1189
+ " H5 ",
1190
+ " H6 ",
1191
+ None,
1192
+ None,
1193
+ None,
1194
+ ), # 30 U
1195
+ (
1196
+ " O4'",
1197
+ " C1'",
1198
+ " C2'",
1199
+ " OP1",
1200
+ " P ",
1201
+ " OP2",
1202
+ " O5'",
1203
+ " C5'",
1204
+ " C4'",
1205
+ " C3'",
1206
+ " O3'",
1207
+ " O2'",
1208
+ None,
1209
+ None,
1210
+ None,
1211
+ None,
1212
+ None,
1213
+ None,
1214
+ None,
1215
+ None,
1216
+ None,
1217
+ None,
1218
+ None,
1219
+ " H5'",
1220
+ "H5''",
1221
+ " H4'",
1222
+ " H3'",
1223
+ " H2'",
1224
+ "HO2'",
1225
+ " H1'",
1226
+ None,
1227
+ None,
1228
+ None,
1229
+ None,
1230
+ None,
1231
+ None,
1232
+ ), # 31 RX (unk RNA)
1233
+ (
1234
+ " N ",
1235
+ " CA ",
1236
+ " C ",
1237
+ " O ",
1238
+ " CB ",
1239
+ " CG ",
1240
+ " NE2",
1241
+ " CD2",
1242
+ " CE1",
1243
+ " ND1",
1244
+ None,
1245
+ None,
1246
+ None,
1247
+ None,
1248
+ None,
1249
+ None,
1250
+ None,
1251
+ None,
1252
+ None,
1253
+ None,
1254
+ None,
1255
+ None,
1256
+ None,
1257
+ " H ",
1258
+ " HA ",
1259
+ "1HB ",
1260
+ "2HB ",
1261
+ "2HD ",
1262
+ "1HE ",
1263
+ "1HD ",
1264
+ None,
1265
+ None,
1266
+ None,
1267
+ None,
1268
+ None,
1269
+ None,
1270
+ ), # -1 his_d
1271
+ ]
1272
+
1273
+ # frames for generic FAPE
1274
+ ##NOTE: 1st entry must be "backbone frame"
1275
+ # fmt: on
1276
+ frames = [
1277
+ [[" N ", " CA ", " C "], [" CA ", " C ", " O "]], # ala
1278
+ [
1279
+ [" N ", " CA ", " C "],
1280
+ [" CA ", " C ", " O "],
1281
+ [" N ", " CA ", " CB "],
1282
+ [" CA ", " CB ", " CG "],
1283
+ [" CB ", " CG ", " CD "],
1284
+ [" CG ", " CD ", " NE "],
1285
+ ], # arg
1286
+ [
1287
+ [" N ", " CA ", " C "],
1288
+ [" CA ", " C ", " O "],
1289
+ [" N ", " CA ", " CB "],
1290
+ [" CA ", " CB ", " CG "],
1291
+ ], # asn
1292
+ [
1293
+ [" N ", " CA ", " C "],
1294
+ [" CA ", " C ", " O "],
1295
+ [" N ", " CA ", " CB "],
1296
+ [" CA ", " CB ", " CG "],
1297
+ ], # asp
1298
+ [
1299
+ [" N ", " CA ", " C "],
1300
+ [" CA ", " C ", " O "],
1301
+ [" N ", " CA ", " CB "],
1302
+ ], # cys
1303
+ [
1304
+ [" N ", " CA ", " C "],
1305
+ [" CA ", " C ", " O "],
1306
+ [" N ", " CA ", " CB "],
1307
+ [" CA ", " CB ", " CG "],
1308
+ [" CB ", " CG ", " CD "],
1309
+ ], # gln
1310
+ [
1311
+ [" N ", " CA ", " C "],
1312
+ [" CA ", " C ", " O "],
1313
+ [" N ", " CA ", " CB "],
1314
+ [" CA ", " CB ", " CG "],
1315
+ [" CB ", " CG ", " CD "],
1316
+ ], # glu
1317
+ [[" N ", " CA ", " C "], [" CA ", " C ", " O "]], # gly
1318
+ [
1319
+ [" N ", " CA ", " C "],
1320
+ [" CA ", " C ", " O "],
1321
+ [" N ", " CA ", " CB "],
1322
+ [" CA ", " CB ", " CG "],
1323
+ ], # his
1324
+ [
1325
+ [" N ", " CA ", " C "],
1326
+ [" CA ", " C ", " O "],
1327
+ [" N ", " CA ", " CB "],
1328
+ [" CA ", " CB ", " CG1"],
1329
+ ], # ile
1330
+ [
1331
+ [" N ", " CA ", " C "],
1332
+ [" CA ", " C ", " O "],
1333
+ [" N ", " CA ", " CB "],
1334
+ [" CA ", " CB ", " CG "],
1335
+ ], # leu
1336
+ [
1337
+ [" N ", " CA ", " C "],
1338
+ [" CA ", " C ", " O "],
1339
+ [" N ", " CA ", " CB "],
1340
+ [" CA ", " CB ", " CG "],
1341
+ [" CB ", " CG ", " CD "],
1342
+ [" CG ", " CD ", " CE "],
1343
+ ], # lys
1344
+ [
1345
+ [" N ", " CA ", " C "],
1346
+ [" CA ", " C ", " O "],
1347
+ [" N ", " CA ", " CB "],
1348
+ [" CA ", " CB ", " CG "],
1349
+ [" CB ", " CG ", " SD "],
1350
+ ], # met
1351
+ [
1352
+ [" N ", " CA ", " C "],
1353
+ [" CA ", " C ", " O "],
1354
+ [" N ", " CA ", " CB "],
1355
+ [" CA ", " CB ", " CG "],
1356
+ ], # phe
1357
+ [
1358
+ [" N ", " CA ", " C "],
1359
+ [" CA ", " C ", " O "],
1360
+ [" N ", " CA ", " CB "],
1361
+ [" CA ", " CB ", " CG "],
1362
+ [" CB ", " CG ", " CD "],
1363
+ ], # pro
1364
+ [
1365
+ [" N ", " CA ", " C "],
1366
+ [" CA ", " C ", " O "],
1367
+ [" N ", " CA ", " CB "],
1368
+ [" CA ", " CB ", " OG "],
1369
+ ], # ser
1370
+ [
1371
+ [" N ", " CA ", " C "],
1372
+ [" CA ", " C ", " O "],
1373
+ [" N ", " CA ", " CB "],
1374
+ [" CA ", " CB ", " OG1"],
1375
+ ], # thr
1376
+ [
1377
+ [" N ", " CA ", " C "],
1378
+ [" CA ", " C ", " O "],
1379
+ [" N ", " CA ", " CB "],
1380
+ [" CA ", " CB ", " CG "],
1381
+ ], # trp
1382
+ [
1383
+ [" N ", " CA ", " C "],
1384
+ [" CA ", " C ", " O "],
1385
+ [" N ", " CA ", " CB "],
1386
+ [" CA ", " CB ", " CG "],
1387
+ ], # tyr
1388
+ [
1389
+ [" N ", " CA ", " C "],
1390
+ [" CA ", " C ", " O "],
1391
+ [" N ", " CA ", " CB "],
1392
+ ], # val
1393
+ [[" N ", " CA ", " C "], [" CA ", " C ", " O "]], # unk
1394
+ [[" N ", " CA ", " C "], [" CA ", " C ", " O "]], # mask
1395
+ [
1396
+ [" O4'", " C1'", " C2'"],
1397
+ [" OP1", " P ", " O5'"],
1398
+ [" P ", " O5'", " C5'"],
1399
+ [" O5'", " C5'", " C4'"],
1400
+ [" C5'", " C4'", " C3'"],
1401
+ [" C5'", " C4'", " O4'"],
1402
+ [" C4'", " O4'", " C1'"],
1403
+ [" C2'", " C1'", " N9 "],
1404
+ [" C4'", " C3'", " O3'"],
1405
+ ], # DA
1406
+ [
1407
+ [" O4'", " C1'", " C2'"],
1408
+ [" OP1", " P ", " O5'"],
1409
+ [" P ", " O5'", " C5'"],
1410
+ [" O5'", " C5'", " C4'"],
1411
+ [" C5'", " C4'", " C3'"],
1412
+ [" C5'", " C4'", " O4'"],
1413
+ [" C4'", " O4'", " C1'"],
1414
+ [" C2'", " C1'", " N1 "],
1415
+ [" C4'", " C3'", " O3'"],
1416
+ ], # DC
1417
+ [
1418
+ [" O4'", " C1'", " C2'"],
1419
+ [" OP1", " P ", " O5'"],
1420
+ [" P ", " O5'", " C5'"],
1421
+ [" O5'", " C5'", " C4'"],
1422
+ [" C5'", " C4'", " C3'"],
1423
+ [" C5'", " C4'", " O4'"],
1424
+ [" C4'", " O4'", " C1'"],
1425
+ [" C2'", " C1'", " N9 "],
1426
+ [" C4'", " C3'", " O3'"],
1427
+ ], # DG
1428
+ [
1429
+ [" O4'", " C1'", " C2'"],
1430
+ [" OP1", " P ", " O5'"],
1431
+ [" P ", " O5'", " C5'"],
1432
+ [" O5'", " C5'", " C4'"],
1433
+ [" C5'", " C4'", " C3'"],
1434
+ [" C5'", " C4'", " O4'"],
1435
+ [" C4'", " O4'", " C1'"],
1436
+ [" C2'", " C1'", " N1 "],
1437
+ [" C4'", " C3'", " O3'"],
1438
+ ], # DT
1439
+ [
1440
+ [" O4'", " C1'", " C2'"],
1441
+ [" OP1", " P ", " O5'"],
1442
+ [" P ", " O5'", " C5'"],
1443
+ [" O5'", " C5'", " C4'"],
1444
+ [" C5'", " C4'", " C3'"],
1445
+ [" C5'", " C4'", " O4'"],
1446
+ [" C4'", " O4'", " C1'"],
1447
+ [" C4'", " C3'", " O3'"],
1448
+ ], # DX
1449
+ [
1450
+ [" O4'", " C1'", " C2'"],
1451
+ [" OP1", " P ", " O5'"],
1452
+ [" P ", " O5'", " C5'"],
1453
+ [" O5'", " C5'", " C4'"],
1454
+ [" C5'", " C4'", " C3'"],
1455
+ [" C5'", " C4'", " O4'"],
1456
+ [" C4'", " O4'", " C1'"],
1457
+ [" C2'", " C1'", " N9 "],
1458
+ [" C4'", " C3'", " O3'"],
1459
+ ], # A
1460
+ [
1461
+ [" O4'", " C1'", " C2'"],
1462
+ [" OP1", " P ", " O5'"],
1463
+ [" P ", " O5'", " C5'"],
1464
+ [" O5'", " C5'", " C4'"],
1465
+ [" C5'", " C4'", " C3'"],
1466
+ [" C5'", " C4'", " O4'"],
1467
+ [" C4'", " O4'", " C1'"],
1468
+ [" C2'", " C1'", " N1 "],
1469
+ [" C4'", " C3'", " O3'"],
1470
+ ], # C
1471
+ [
1472
+ [" O4'", " C1'", " C2'"],
1473
+ [" OP1", " P ", " O5'"],
1474
+ [" P ", " O5'", " C5'"],
1475
+ [" O5'", " C5'", " C4'"],
1476
+ [" C5'", " C4'", " C3'"],
1477
+ [" C5'", " C4'", " O4'"],
1478
+ [" C4'", " O4'", " C1'"],
1479
+ [" C2'", " C1'", " N9 "],
1480
+ [" C4'", " C3'", " O3'"],
1481
+ ], # G
1482
+ [
1483
+ [" O4'", " C1'", " C2'"],
1484
+ [" OP1", " P ", " O5'"],
1485
+ [" P ", " O5'", " C5'"],
1486
+ [" O5'", " C5'", " C4'"],
1487
+ [" C5'", " C4'", " C3'"],
1488
+ [" C5'", " C4'", " O4'"],
1489
+ [" C4'", " O4'", " C1'"],
1490
+ [" C2'", " C1'", " N1 "],
1491
+ [" C4'", " C3'", " O3'"],
1492
+ ], # U
1493
+ [
1494
+ [" O4'", " C1'", " C2'"],
1495
+ [" OP1", " P ", " O5'"],
1496
+ [" P ", " O5'", " C5'"],
1497
+ [" O5'", " C5'", " C4'"],
1498
+ [" C5'", " C4'", " C3'"],
1499
+ [" C5'", " C4'", " O4'"],
1500
+ [" C4'", " O4'", " C1'"],
1501
+ [" C4'", " C3'", " O3'"],
1502
+ ], # RX
1503
+ ]
1504
+ # fmt: on
1505
+ NFRAMES = max([len(f) for f in frames])
1506
+
1507
+ # general FAPE parameters
1508
+ frame_indices = torch.full((NAATOKENS, NFRAMES, 3, 2), 0, dtype=torch.long)
1509
+ for i in range(NNAPROTAAS):
1510
+ i_l = aa2long[i]
1511
+ for j, x in enumerate(frames[i]):
1512
+ if x is not None:
1513
+ # frames are stored as (residue offset, atom position)
1514
+ frame_indices[i, j, 0] = torch.tensor((0, i_l.index(x[0])))
1515
+ frame_indices[i, j, 1] = torch.tensor((0, i_l.index(x[1])))
1516
+ frame_indices[i, j, 2] = torch.tensor((0, i_l.index(x[2])))
1517
+
1518
+
1519
+ allatom_mask = torch.zeros((NAATOKENS, NTOTAL), dtype=torch.bool)
1520
+ long2alt = torch.zeros((NAATOKENS, NTOTAL), dtype=torch.long)
1521
+ for i in range(NNAPROTAAS):
1522
+ for j, a in enumerate(aa2long[i]):
1523
+ if a is not None:
1524
+ allatom_mask[i, j] = True
1525
+
1526
+ allatom_mask[NNAPROTAAS:, 1] = True
1527
+ allatom_mask[:, 1] = True
1528
+ heavyatom_mask = allatom_mask.clone()
1529
+ heavyatom_mask[:, NHEAVY:] = False