bioemu2 0.1.0rc0__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 (183) hide show
  1. bioemu2-0.1.0rc0/LICENSE +21 -0
  2. bioemu2-0.1.0rc0/NOTICES.md +135 -0
  3. bioemu2-0.1.0rc0/PKG-INFO +355 -0
  4. bioemu2-0.1.0rc0/README.md +290 -0
  5. bioemu2-0.1.0rc0/licenses/LICENSE.BOLTZ-MIT +21 -0
  6. bioemu2-0.1.0rc0/licenses/LICENSE.GEOSTD +43 -0
  7. bioemu2-0.1.0rc0/licenses/LICENSE.GPL-3.0 +232 -0
  8. bioemu2-0.1.0rc0/licenses/LICENSE.LGPL-3.0 +165 -0
  9. bioemu2-0.1.0rc0/licenses/LICENSE.NVIDIA-FOLD-CP-MIT +20 -0
  10. bioemu2-0.1.0rc0/licenses/LICENSE.PYTORCH3D-BSD +30 -0
  11. bioemu2-0.1.0rc0/licenses/LICENSE.TORCH-SCATTER-MIT +19 -0
  12. bioemu2-0.1.0rc0/pyproject.toml +152 -0
  13. bioemu2-0.1.0rc0/setup.cfg +4 -0
  14. bioemu2-0.1.0rc0/src/bioemu2/__init__.py +87 -0
  15. bioemu2-0.1.0rc0/src/bioemu2/assets/examples/systems/ake.yaml +8 -0
  16. bioemu2-0.1.0rc0/src/bioemu2/assets/examples/systems/chignolin.yaml +8 -0
  17. bioemu2-0.1.0rc0/src/bioemu2/assets/examples/systems/meteorin_like_7tlw.yaml +11 -0
  18. bioemu2-0.1.0rc0/src/bioemu2/assets/examples/systems/myoglobin_101m.yaml +19 -0
  19. bioemu2-0.1.0rc0/src/bioemu2/assets/examples/systems/shp2_3tkz.yaml +15 -0
  20. bioemu2-0.1.0rc0/src/bioemu2/assets/polymer_templates.pt +0 -0
  21. bioemu2-0.1.0rc0/src/bioemu2/assets/stereo_chemistry.yaml +6567 -0
  22. bioemu2-0.1.0rc0/src/bioemu2/biomoa/__init__.py +0 -0
  23. bioemu2-0.1.0rc0/src/bioemu2/biomoa/models/__init__.py +0 -0
  24. bioemu2-0.1.0rc0/src/bioemu2/biomoa/models/base.py +12 -0
  25. bioemu2-0.1.0rc0/src/bioemu2/biomoa/models/edm_schedule.py +86 -0
  26. bioemu2-0.1.0rc0/src/bioemu2/biomoa/models/score_model.py +373 -0
  27. bioemu2-0.1.0rc0/src/bioemu2/biomoa/modules/__init__.py +0 -0
  28. bioemu2-0.1.0rc0/src/bioemu2/biomoa/modules/conditioning.py +125 -0
  29. bioemu2-0.1.0rc0/src/bioemu2/biomoa/modules/diffusion.py +866 -0
  30. bioemu2-0.1.0rc0/src/bioemu2/biomoa/modules/embedders.py +206 -0
  31. bioemu2-0.1.0rc0/src/bioemu2/biomoa/modules/encoders.py +457 -0
  32. bioemu2-0.1.0rc0/src/bioemu2/biomoa/modules/pairmixer.py +757 -0
  33. bioemu2-0.1.0rc0/src/bioemu2/biomoa/modules/relative_position_encoder.py +101 -0
  34. bioemu2-0.1.0rc0/src/bioemu2/biomoa/modules/time.py +78 -0
  35. bioemu2-0.1.0rc0/src/bioemu2/biomoa/modules/transformer.py +410 -0
  36. bioemu2-0.1.0rc0/src/bioemu2/biomoa/ops/__init__.py +0 -0
  37. bioemu2-0.1.0rc0/src/bioemu2/biomoa/ops/atom_token_ops.py +212 -0
  38. bioemu2-0.1.0rc0/src/bioemu2/biomoa/ops/attention_pair_bias.py +346 -0
  39. bioemu2-0.1.0rc0/src/bioemu2/biomoa/ops/dropout.py +70 -0
  40. bioemu2-0.1.0rc0/src/bioemu2/biomoa/ops/outer_product_mean.py +158 -0
  41. bioemu2-0.1.0rc0/src/bioemu2/biomoa/ops/pair_weighted_averaging.py +153 -0
  42. bioemu2-0.1.0rc0/src/bioemu2/biomoa/ops/transition.py +91 -0
  43. bioemu2-0.1.0rc0/src/bioemu2/biomoa/ops/triangle_multiplication.py +395 -0
  44. bioemu2-0.1.0rc0/src/bioemu2/biomoa/ops/triton/__init__.py +0 -0
  45. bioemu2-0.1.0rc0/src/bioemu2/biomoa/ops/triton/opm.py +271 -0
  46. bioemu2-0.1.0rc0/src/bioemu2/biomoa/utils/__init__.py +5 -0
  47. bioemu2-0.1.0rc0/src/bioemu2/biomoa/utils/feature_converter.py +122 -0
  48. bioemu2-0.1.0rc0/src/bioemu2/biomoa/utils/precision.py +85 -0
  49. bioemu2-0.1.0rc0/src/bioemu2/biomoa/utils/strategy.py +50 -0
  50. bioemu2-0.1.0rc0/src/bioemu2/biomoa/utils/tiling.py +254 -0
  51. bioemu2-0.1.0rc0/src/bioemu2/biomoa/utils/warnings.py +27 -0
  52. bioemu2-0.1.0rc0/src/bioemu2/checkpoints.py +114 -0
  53. bioemu2-0.1.0rc0/src/bioemu2/cli.py +734 -0
  54. bioemu2-0.1.0rc0/src/bioemu2/conditioning/__init__.py +0 -0
  55. bioemu2-0.1.0rc0/src/bioemu2/conditioning/contact.py +82 -0
  56. bioemu2-0.1.0rc0/src/bioemu2/conditioning/method.py +44 -0
  57. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/__init__.py +1 -0
  58. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/api.py +590 -0
  59. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/diagnostics.py +45 -0
  60. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/diffusion/__init__.py +3 -0
  61. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/distributed/__init__.py +0 -0
  62. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/distributed/comm.py +595 -0
  63. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/distributed/manager.py +891 -0
  64. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/distributed/utils.py +791 -0
  65. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/features.py +158 -0
  66. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/layers/__init__.py +17 -0
  67. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/layers/relative_position_encoder.py +88 -0
  68. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/layers/transformer_layers.py +224 -0
  69. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/models/__init__.py +50 -0
  70. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/models/score_model.py +310 -0
  71. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/modules/__init__.py +37 -0
  72. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/modules/conditioning.py +55 -0
  73. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/modules/diffusion.py +361 -0
  74. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/modules/embedders.py +31 -0
  75. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/modules/encoders.py +168 -0
  76. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/modules/pairmixer.py +424 -0
  77. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/ops/__init__.py +65 -0
  78. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/ops/atom_token_ops.py +101 -0
  79. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/ops/attention.py +164 -0
  80. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/ops/attention_impl.py +1325 -0
  81. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/ops/cat_and_chunk.py +31 -0
  82. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/ops/distributed_primitives/__init__.py +0 -0
  83. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/ops/distributed_primitives/dtensor_metadata_tools.py +138 -0
  84. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/ops/distributed_primitives/elementwise_op.py +438 -0
  85. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/ops/distributed_primitives/flatten_and_unflatten.py +797 -0
  86. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/ops/distributed_primitives/gather.py +532 -0
  87. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/ops/distributed_primitives/layernorm.py +402 -0
  88. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/ops/distributed_primitives/linear.py +433 -0
  89. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/ops/distributed_primitives/modules_utils.py +49 -0
  90. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/ops/distributed_primitives/outer_gather.py +1370 -0
  91. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/ops/distributed_primitives/redistribute_transpose.py +347 -0
  92. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/ops/distributed_primitives/sigmoid_gate.py +250 -0
  93. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/ops/distributed_primitives/utils.py +2212 -0
  94. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/ops/outer_product_mean.py +265 -0
  95. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/ops/pair_averaging.py +394 -0
  96. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/ops/pair_weighted_averaging.py +79 -0
  97. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/ops/readout_superposition.py +59 -0
  98. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/ops/scatter.py +603 -0
  99. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/ops/tiled_params_replicated.py +81 -0
  100. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/ops/transition.py +67 -0
  101. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/ops/triangle_multiplication.py +191 -0
  102. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/ops/triangular_mult_impl.py +696 -0
  103. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/ops/triton/__init__.py +0 -0
  104. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/ops/triton/flash_attn.py +1126 -0
  105. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/ops/triton/flash_attn_impl.py +323 -0
  106. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/ops/triton/opm.py +111 -0
  107. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/scripts/cp_env.sh +88 -0
  108. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/utils/__init__.py +81 -0
  109. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/utils/dtensor.py +164 -0
  110. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/utils/mesh.py +106 -0
  111. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/utils/padding.py +146 -0
  112. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/utils/rng.py +39 -0
  113. bioemu2-0.1.0rc0/src/bioemu2/context_parallel/utils/session.py +218 -0
  114. bioemu2-0.1.0rc0/src/bioemu2/cp.py +6 -0
  115. bioemu2-0.1.0rc0/src/bioemu2/data_interface/__init__.py +36 -0
  116. bioemu2-0.1.0rc0/src/bioemu2/data_interface/boltz.py +270 -0
  117. bioemu2-0.1.0rc0/src/bioemu2/data_interface/bond_reference.py +83 -0
  118. bioemu2-0.1.0rc0/src/bioemu2/data_interface/cif.py +362 -0
  119. bioemu2-0.1.0rc0/src/bioemu2/data_interface/const.py +578 -0
  120. bioemu2-0.1.0rc0/src/bioemu2/data_interface/ligand_validators.py +56 -0
  121. bioemu2-0.1.0rc0/src/bioemu2/data_interface/structure_hierarchy.py +614 -0
  122. bioemu2-0.1.0rc0/src/bioemu2/data_interface/system.py +516 -0
  123. bioemu2-0.1.0rc0/src/bioemu2/denoising/__init__.py +3 -0
  124. bioemu2-0.1.0rc0/src/bioemu2/denoising/augmentation.py +70 -0
  125. bioemu2-0.1.0rc0/src/bioemu2/denoising/cp_sync.py +84 -0
  126. bioemu2-0.1.0rc0/src/bioemu2/denoising/edm.py +185 -0
  127. bioemu2-0.1.0rc0/src/bioemu2/denoising/rigid_align.py +208 -0
  128. bioemu2-0.1.0rc0/src/bioemu2/denoising/samplers.py +181 -0
  129. bioemu2-0.1.0rc0/src/bioemu2/denoising/steering/__init__.py +5 -0
  130. bioemu2-0.1.0rc0/src/bioemu2/denoising/steering/conditioning.py +43 -0
  131. bioemu2-0.1.0rc0/src/bioemu2/denoising/steering/edm_fkc.py +464 -0
  132. bioemu2-0.1.0rc0/src/bioemu2/denoising/steering/edm_steered.py +381 -0
  133. bioemu2-0.1.0rc0/src/bioemu2/denoising/steering/feats.py +566 -0
  134. bioemu2-0.1.0rc0/src/bioemu2/denoising/steering/geometry_topology.py +591 -0
  135. bioemu2-0.1.0rc0/src/bioemu2/denoising/steering/potentials.py +1084 -0
  136. bioemu2-0.1.0rc0/src/bioemu2/denoising/steering/rdkit_constraints.py +243 -0
  137. bioemu2-0.1.0rc0/src/bioemu2/embedding_cache.py +154 -0
  138. bioemu2-0.1.0rc0/src/bioemu2/inference.py +236 -0
  139. bioemu2-0.1.0rc0/src/bioemu2/model.py +448 -0
  140. bioemu2-0.1.0rc0/src/bioemu2/msa/__init__.py +3 -0
  141. bioemu2-0.1.0rc0/src/bioemu2/msa/alignment.py +295 -0
  142. bioemu2-0.1.0rc0/src/bioemu2/msa/arrays.py +224 -0
  143. bioemu2-0.1.0rc0/src/bioemu2/msa/colabfold.py +528 -0
  144. bioemu2-0.1.0rc0/src/bioemu2/msa/features.py +49 -0
  145. bioemu2-0.1.0rc0/src/bioemu2/patches/__init__.py +0 -0
  146. bioemu2-0.1.0rc0/src/bioemu2/patches/mdtraj_patcher.py +371 -0
  147. bioemu2-0.1.0rc0/src/bioemu2/patches/openmm_patches.py +543 -0
  148. bioemu2-0.1.0rc0/src/bioemu2/paths.py +6 -0
  149. bioemu2-0.1.0rc0/src/bioemu2/physicality.py +1026 -0
  150. bioemu2-0.1.0rc0/src/bioemu2/precision.py +47 -0
  151. bioemu2-0.1.0rc0/src/bioemu2/propagation.py +475 -0
  152. bioemu2-0.1.0rc0/src/bioemu2/rng.py +11 -0
  153. bioemu2-0.1.0rc0/src/bioemu2/sampling.py +593 -0
  154. bioemu2-0.1.0rc0/src/bioemu2/shortcuts.py +7 -0
  155. bioemu2-0.1.0rc0/src/bioemu2/structure/__init__.py +31 -0
  156. bioemu2-0.1.0rc0/src/bioemu2/structure/all_atom_structure.py +1193 -0
  157. bioemu2-0.1.0rc0/src/bioemu2/structure/batch.py +237 -0
  158. bioemu2-0.1.0rc0/src/bioemu2/structure/builder.py +126 -0
  159. bioemu2-0.1.0rc0/src/bioemu2/structure/cif.py +230 -0
  160. bioemu2-0.1.0rc0/src/bioemu2/structure/inference_inputs.py +955 -0
  161. bioemu2-0.1.0rc0/src/bioemu2/structure/polymer_builder.py +135 -0
  162. bioemu2-0.1.0rc0/src/bioemu2/structure/preparation.py +80 -0
  163. bioemu2-0.1.0rc0/src/bioemu2/structure/ref_struct.py +280 -0
  164. bioemu2-0.1.0rc0/src/bioemu2/structure/scatter.py +113 -0
  165. bioemu2-0.1.0rc0/src/bioemu2/structure/templates.py +468 -0
  166. bioemu2-0.1.0rc0/src/bioemu2/trajectory_io.py +118 -0
  167. bioemu2-0.1.0rc0/src/bioemu2/visualization.py +153 -0
  168. bioemu2-0.1.0rc0/src/bioemu2.egg-info/PKG-INFO +355 -0
  169. bioemu2-0.1.0rc0/src/bioemu2.egg-info/SOURCES.txt +181 -0
  170. bioemu2-0.1.0rc0/src/bioemu2.egg-info/dependency_links.txt +1 -0
  171. bioemu2-0.1.0rc0/src/bioemu2.egg-info/entry_points.txt +2 -0
  172. bioemu2-0.1.0rc0/src/bioemu2.egg-info/requires.txt +53 -0
  173. bioemu2-0.1.0rc0/src/bioemu2.egg-info/top_level.txt +1 -0
  174. bioemu2-0.1.0rc0/tests/test_convert_boltz.py +153 -0
  175. bioemu2-0.1.0rc0/tests/test_import.py +15 -0
  176. bioemu2-0.1.0rc0/tests/test_physicality.py +786 -0
  177. bioemu2-0.1.0rc0/tests/test_propagation.py +699 -0
  178. bioemu2-0.1.0rc0/tests/test_release_contract.py +233 -0
  179. bioemu2-0.1.0rc0/tests/test_rng.py +19 -0
  180. bioemu2-0.1.0rc0/tests/test_sample_protein_protein_constraints.py +102 -0
  181. bioemu2-0.1.0rc0/tests/test_sampling.py +2040 -0
  182. bioemu2-0.1.0rc0/tests/test_trajectory_io.py +170 -0
  183. bioemu2-0.1.0rc0/tests/test_visualization.py +90 -0
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) Microsoft Corporation.
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,135 @@
1
+ # NOTICES
2
+
3
+ Except where otherwise noted, code authored for BioEmu2 by Microsoft is
4
+ licensed under the MIT License. This distribution also contains third-party
5
+ code under the licenses and notices described below and in the relevant source
6
+ files. See `LICENSE` and `licenses/LICENSE.BOLTZ-MIT`, `licenses/LICENSE.GEOSTD`,
7
+ `licenses/LICENSE.NVIDIA-FOLD-CP-MIT`, `licenses/LICENSE.PYTORCH3D-BSD`,
8
+ `licenses/LICENSE.TORCH-SCATTER-MIT`, `licenses/LICENSE.LGPL-3.0`, and
9
+ `licenses/LICENSE.GPL-3.0`.
10
+
11
+ ## Stereochemistry reference data
12
+
13
+ `src/bioemu2/assets/stereo_chemistry.yaml` contains bond-length and
14
+ bond-angle reference data exported from
15
+ [OpenStructure 2.11.1](https://git.scicore.unibas.ch/schwede/openstructure/-/tree/2.11.1)
16
+ (commit `aeebc45c0fc2e3255cb45178013dad5a00999a39`). The `monlib` section
17
+ comes from OpenStructure's `modules/mol/alg/pymod/stereo_data.json`, derived
18
+ from CCP4 MON_LIB, which is licensed under the GNU Lesser General Public
19
+ License, version 3.0. The license is reproduced in `licenses/LICENSE.LGPL-3.0`,
20
+ together with the GNU GPL version 3 in `licenses/LICENSE.GPL-3.0`.
21
+
22
+ The `monlib_link` section comes from OpenStructure's
23
+ `modules/mol/alg/pymod/stereochemistry.py` (`GetDefaultStereoLinkData`).
24
+ Its nucleotide link parameters are adapted from geostd's
25
+ `rna_dna/chain_link_rna2p.cif`. A
26
+ [historical revision](https://github.com/phenix-project/geostd/blob/aef2b63e65ffab32157da4d6422723982a3e7d6d/rna_dna/chain_link_rna2p.cif)
27
+ contains the corresponding values; OpenStructure does not specify the exact
28
+ geostd revision it used. The geostd license is reproduced in `licenses/LICENSE.GEOSTD`.
29
+ Its peptide and disulfide link parameters come from CCP4's
30
+ `standard_geometry.cif`. OpenStructure identifies both link-data sources in
31
+ the source file above. No legacy Engh-Huber section is included in this asset.
32
+
33
+ The following geostd notice is reproduced from its `COPYRIGHT.txt`:
34
+
35
+ ```text
36
+ geostd Copyright (c) 2009, The Regents of the University of
37
+ California, through Lawrence Berkeley National Laboratory (subject to
38
+ receipt of any required approvals from the U.S. Dept. of Energy). All
39
+ rights reserved.
40
+
41
+ If you have questions about your rights to use or distribute this
42
+ software, please contact Berkeley Lab's Technology Transfer Department
43
+ at TTD@lbl.gov referring to "geostd (LBNL Ref CR-2725)."
44
+
45
+ NOTICE. This software was developed under funding from the
46
+ U.S. Department of Energy. As such, the U.S. Government has been
47
+ granted for itself and others acting on its behalf a paid-up,
48
+ nonexclusive, irrevocable, worldwide license in the Software to
49
+ reproduce, prepare derivative works, and perform publicly and display
50
+ publicly. Beginning five (5) years after the date permission to
51
+ assert copyright is obtained from the U.S. Department of Energy, and
52
+ subject to any subsequent five (5) year renewals, the U.S. Government
53
+ is granted for itself and others acting on its behalf a paid-up,
54
+ nonexclusive, irrevocable, worldwide license in the Software to
55
+ reproduce, prepare derivative works, distribute copies to the public,
56
+ perform publicly and display publicly, and to permit others to do so.
57
+ ```
58
+
59
+ ## Boltz-derived code
60
+
61
+ BioMoA layers and operations, denoising and steering code, MSA feature
62
+ preparation, data-interface code, and all-atom feature preparation under
63
+ `src/bioemu2` contain portions copied or adapted from
64
+ [Boltz v1.0.0](https://github.com/jwohlwend/boltz/tree/v1.0.0). Those portions
65
+ are licensed under the MIT License reproduced in `licenses/LICENSE.BOLTZ-MIT`. This
66
+ notice continues to apply where modules and symbols use BioEmu2-local names.
67
+
68
+ Copyright (c) 2024 Jeremy Wohlwend, Gabriele Corso, Saro Passaro.
69
+
70
+ Contact-conditioning code and feature conventions in
71
+ `src/bioemu2/biomoa/modules/conditioning.py` and
72
+ `src/bioemu2/conditioning/contact.py` are adapted from
73
+ [Boltz at commit `b1ebfc46ecf57f5414e0d1a6f9027bbb122c53bc`](https://github.com/jwohlwend/boltz/tree/b1ebfc46ecf57f5414e0d1a6f9027bbb122c53bc).
74
+ Those portions are also licensed under the MIT License reproduced in
75
+ `licenses/LICENSE.BOLTZ-MIT`.
76
+
77
+ ## PyTorch3D rotation helpers
78
+
79
+ The quaternion and rotation helpers in
80
+ `src/bioemu2/denoising/augmentation.py` contain adapted portions of
81
+ [PyTorch3D](https://github.com/facebookresearch/pytorch3d) under the BSD
82
+ License reproduced in `licenses/LICENSE.PYTORCH3D-BSD`.
83
+
84
+ Copyright (c) Meta Platforms, Inc. and affiliates. All rights reserved.
85
+
86
+ ## torch_scatter-derived code
87
+
88
+ `src/bioemu2/structure/scatter.py` contains portions copied or adapted from
89
+ [`torch_scatter`](https://github.com/rusty1s/pytorch_scatter) at commit
90
+ `8ec9364b0bdcd99149952a25749ad211c2d0567b`, via the Microsoft-authored
91
+ `mdmlff.src.mlff.utils.scatter` wrapper. The upstream MIT License is reproduced
92
+ in `licenses/LICENSE.TORCH-SCATTER-MIT`.
93
+
94
+ Copyright (c) 2020 Matthias Fey <matthias.fey@tu-dortmund.de>.
95
+
96
+ ## NVIDIA Fold-CP-derived code
97
+
98
+ Context-parallel distributed primitives and operations under
99
+ `src/bioemu2/context_parallel` contain portions copied, adapted, or ported from
100
+ [NVIDIA Fold-CP](https://github.com/NVIDIA-BioNeMo/boltz-cp) at commit
101
+ `f76f37a77c854b56b6250b426af8c2d63b501f7f`. Source files identify their
102
+ upstream paths and describe material modifications. The upstream MIT License is
103
+ reproduced in `licenses/LICENSE.NVIDIA-FOLD-CP-MIT`.
104
+
105
+ Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
106
+
107
+ Fold-CP also identifies portions derived from
108
+ [Boltz](https://github.com/jwohlwend/boltz) at commit
109
+ `cb04aeccdd480fd4db707f0bbafde538397fa2ac`; the Boltz notice and license are
110
+ provided above and in `licenses/LICENSE.BOLTZ-MIT`.
111
+
112
+ ## MDTraj patches
113
+
114
+ `src/bioemu2/patches/mdtraj_patcher.py` contains modified portions of
115
+ [MDTraj](https://github.com/mdtraj/mdtraj), specifically the
116
+ `Topology.from_openmm` and `PDBxTrajectoryFile` implementations. The surrounding
117
+ patching and integration code was authored for BioEmu2.
118
+
119
+ The file is licensed under the GNU Lesser General Public License, version 3.0
120
+ or any later version. The copied MDTraj portions were originally offered under
121
+ the GNU Lesser General Public License, version 2.1 or any later version; this
122
+ distribution exercises the option to use version 3.0 or any later version.
123
+
124
+ The upstream MDTraj files also contain portions originating from OpenMM under
125
+ the MIT License. Their complete copyright, permission, and warranty notices are
126
+ retained in `src/bioemu2/patches/mdtraj_patcher.py`.
127
+
128
+ ## OpenMM patches
129
+
130
+ `src/bioemu2/patches/openmm_patches.py` contains modified portions of
131
+ [OpenMM](https://github.com/openmm/openmm) and is licensed under OpenMM's MIT
132
+ license. The complete OpenMM copyright and permission notice is reproduced in
133
+ that file.
134
+
135
+ Portions copyright (c) 2015-2023 Stanford University and the Authors.
@@ -0,0 +1,355 @@
1
+ Metadata-Version: 2.4
2
+ Name: bioemu2
3
+ Version: 0.1.0rc0
4
+ Summary: Local protein structure and conformational sampling with BioEmu2.
5
+ License-Expression: MIT
6
+ Classifier: Operating System :: POSIX :: Linux
7
+ Classifier: Programming Language :: Python :: 3.12
8
+ Classifier: Programming Language :: Python :: 3.13
9
+ Classifier: Programming Language :: Python :: 3.14
10
+ Requires-Python: >=3.12
11
+ Description-Content-Type: text/markdown
12
+ License-File: LICENSE
13
+ License-File: licenses/LICENSE.BOLTZ-MIT
14
+ License-File: licenses/LICENSE.GEOSTD
15
+ License-File: licenses/LICENSE.GPL-3.0
16
+ License-File: licenses/LICENSE.LGPL-3.0
17
+ License-File: licenses/LICENSE.NVIDIA-FOLD-CP-MIT
18
+ License-File: licenses/LICENSE.PYTORCH3D-BSD
19
+ License-File: licenses/LICENSE.TORCH-SCATTER-MIT
20
+ License-File: NOTICES.md
21
+ Requires-Dist: biopython>=1.88
22
+ Requires-Dist: biotite>=1.5
23
+ Requires-Dist: click
24
+ Requires-Dist: einops
25
+ Requires-Dist: huggingface_hub
26
+ Requires-Dist: hydra-core
27
+ Requires-Dist: mdtraj
28
+ Requires-Dist: numba
29
+ Requires-Dist: numpy<2.5
30
+ Requires-Dist: omegaconf
31
+ Requires-Dist: openmm
32
+ Requires-Dist: pdbecif
33
+ Requires-Dist: pydantic
34
+ Requires-Dist: pyyaml
35
+ Requires-Dist: rdkit
36
+ Requires-Dist: requests
37
+ Requires-Dist: scipy
38
+ Requires-Dist: tenacity
39
+ Requires-Dist: torch<2.14,>=2.11
40
+ Requires-Dist: tqdm
41
+ Requires-Dist: triton; sys_platform == "linux"
42
+ Provides-Extra: cuda
43
+ Requires-Dist: cuequivariance==0.11.1; sys_platform == "linux" and extra == "cuda"
44
+ Requires-Dist: cuequivariance-torch==0.11.1; sys_platform == "linux" and extra == "cuda"
45
+ Requires-Dist: cuequivariance-ops-torch-cu13==0.11.1; sys_platform == "linux" and extra == "cuda"
46
+ Requires-Dist: cuequivariance-ops-cu13==0.11.1; sys_platform == "linux" and extra == "cuda"
47
+ Provides-Extra: cuda12
48
+ Requires-Dist: torch==2.13.0+cu126; sys_platform == "linux" and extra == "cuda12"
49
+ Requires-Dist: cuequivariance==0.11.1; sys_platform == "linux" and extra == "cuda12"
50
+ Requires-Dist: cuequivariance-torch==0.11.1; sys_platform == "linux" and extra == "cuda12"
51
+ Requires-Dist: cuequivariance-ops-torch-cu12==0.11.1; sys_platform == "linux" and extra == "cuda12"
52
+ Requires-Dist: cuequivariance-ops-cu12==0.11.1; sys_platform == "linux" and extra == "cuda12"
53
+ Provides-Extra: dev
54
+ Requires-Dist: pre-commit; extra == "dev"
55
+ Requires-Dist: py3dmol>=2.5.5; extra == "dev"
56
+ Requires-Dist: pytest; extra == "dev"
57
+ Requires-Dist: pytest-cov; extra == "dev"
58
+ Provides-Extra: viz
59
+ Requires-Dist: ipykernel; extra == "viz"
60
+ Requires-Dist: ipywidgets>=8; extra == "viz"
61
+ Requires-Dist: matplotlib; extra == "viz"
62
+ Requires-Dist: playwright; extra == "viz"
63
+ Requires-Dist: py3Dmol; extra == "viz"
64
+ Dynamic: license-file
65
+
66
+ <div align="center">
67
+
68
+ <img src="assets/emu.png" alt="BioEmu2 logo" width="120">
69
+
70
+ # BioEmu2
71
+
72
+ ### A foundation model for biomolecular structure and dynamics
73
+
74
+ **Predict folded structures, sample equilibrium ensembles, and propagate structures forward in time.**
75
+
76
+ [Installation](#installation) · [Quick start](#quick-start) · [Tutorials](#tutorials) · [Documentation](docs/)
77
+
78
+ <br>
79
+
80
+ <img src="https://github.com/microsoft/bioemu2/releases/download/readme-media/mcl1_puma_clean_grid_v2.gif" alt="BioEmu2 simulating the PUMA BH3 peptide binding MCL-1" width="420">
81
+
82
+ <sub>BioEmu2 simulating the PUMA BH3 peptide binding MCL-1.</sub>
83
+
84
+ </div>
85
+
86
+ ---
87
+
88
+ BioEmu2, short for "biomolecular emulator", is developed by Microsoft Research
89
+ AI for Science.
90
+
91
+ Key features include:
92
+
93
+ - molecular systems containing proteins, nucleic acids, ligands, and other
94
+ chemical components;
95
+ - equilibrium sampling and structure prediction from the command line or
96
+ Python;
97
+ - time-lagged trajectory generation with implicit transfer operator models;
98
+ - contact conditioning and inference-time steering toward desired properties;
99
+ - context-parallel sampling of structures that do not fit on one GPU.
100
+
101
+ ## System requirements
102
+
103
+ BioEmu2 requires Linux and an NVIDIA GPU; macOS and Windows are not supported.
104
+ The `cuda` extra installs CUDA 13 cuEquivariance, which requires compute
105
+ capability 7.5 or newer, such as NVIDIA T4, A100, RTX A6000, L40S, H100, and B200
106
+ ([CUDA 13 release notes](https://docs.nvidia.com/cuda/archive/13.0.0/cuda-toolkit-release-notes/index.html#deprecated-architectures)).
107
+ GPUs with compute capability 7.0 (Volta, such as NVIDIA V100) need the CUDA 12
108
+ [installation](#installation) instead.
109
+ We have tested V100 (compute capability 7.0), T4 (7.5, with the CUDA 12
110
+ installation), A100 and RTX A6000 (8.0 and 8.6), and B200 (10.0); other eligible
111
+ GPUs have not been verified end to end. Check
112
+ your GPU's compute capability with
113
+ `nvidia-smi --query-gpu=name,compute_cap --format=csv`.
114
+
115
+ Context-parallel sampling requires at least four GPUs.
116
+
117
+ ## Installation
118
+
119
+ Use Python 3.12, 3.13, or 3.14 on a supported Linux machine. We recommend
120
+ installing BioEmu2 in an isolated
121
+ [virtual environment](https://docs.python.org/3/library/venv.html) or
122
+ [conda environment](https://docs.conda.io/projects/conda/en/latest/user-guide/tasks/manage-environments.html).
123
+ To install in an isolated virtual environment, create and activate it from the
124
+ repository root:
125
+
126
+ ```bash
127
+ python3.12 -m venv .venv
128
+ source .venv/bin/activate
129
+ ```
130
+
131
+ Replace `python3.12` with the name or path of your Python 3.12, 3.13, or 3.14
132
+ executable if needed. Use `python -m pip` so packages are installed into the
133
+ active Python environment rather than a different interpreter's `pip`.
134
+
135
+ Run the installation command for your GPU from the repository root.
136
+ For GPUs with compute capability 7.5 or newer, such as NVIDIA T4, A100,
137
+ RTX A6000, L40S, H100, and B200, use the `cuda` extra:
138
+
139
+ ```bash
140
+ python -m pip install ".[cuda]"
141
+ ```
142
+
143
+ Compute capability 7.0 (Volta, such as NVIDIA V100) is the oldest supported;
144
+ on these GPUs, use the `cuda12` extra with the PyTorch CUDA 12.6 package index:
145
+
146
+ ```bash
147
+ python -m pip install ".[cuda12]" \
148
+ --extra-index-url https://download.pytorch.org/whl/cu126
149
+ ```
150
+
151
+ GPUs with compute capability below 7.0 are not supported.
152
+
153
+ We tested BioEmu2 and evaluated its benchmark metrics with Python 3.12. Python
154
+ 3.13 and 3.14 pass the test suite, but we have not evaluated benchmark metrics
155
+ with them.
156
+
157
+ For development, see [DEVELOPER.md](DEVELOPER.md).
158
+
159
+ For Jupyter or VS Code notebooks, follow the [notebook
160
+ quickstart](#notebook-quickstart) to install visualization dependencies and
161
+ register a kernel.
162
+
163
+ ## Quick start
164
+
165
+ ### CLI
166
+
167
+ Sample eight equilibrium conformations of adenylate kinase (PDB 1AKE):
168
+
169
+ ```bash
170
+ bioemu2 sample-equilibrium \
171
+ --pdb-id 1AKE \
172
+ --out outputs/adk
173
+ ```
174
+
175
+ Or predict a folded structure from a PDB entry:
176
+
177
+ ```bash
178
+ bioemu2 predict-structure \
179
+ --pdb-id 1UBQ \
180
+ --num-samples 1 \
181
+ --out outputs/ubiquitin
182
+ ```
183
+
184
+ Propagate each sampled adenylate kinase conformation by 1 ns:
185
+
186
+ ```bash
187
+ bioemu2 propagate \
188
+ --x0-cif outputs/adk/topology.cif \
189
+ --x0-xtc outputs/adk/samples.xtc \
190
+ --lag-time-ns 1 \
191
+ --out outputs/adk-propagated
192
+ ```
193
+
194
+ Run `bioemu2 --help` to list all commands.
195
+
196
+ ### Python API using EmuSession
197
+
198
+ Construct an `EmuSession` to sample and propagate structures:
199
+
200
+ ```python
201
+ from pathlib import Path
202
+
203
+ from bioemu2 import EmuSession, System, load_model, save_frames
204
+
205
+ model = load_model()
206
+
207
+ # Predict a folded structure for ubiquitin.
208
+ ubiquitin = System.from_pdb_id("1UBQ")
209
+ prediction = EmuSession(ubiquitin, model=model, mode="structure").sample(num_samples=1)
210
+ save_frames(prediction, Path("outputs/ubiquitin"))
211
+
212
+ # Sample equilibrium conformations of adenylate kinase.
213
+ adk = System.from_pdb_id("1AKE")
214
+ emu = EmuSession(adk, model=model, mode="equilibrium")
215
+ samples = emu.sample(num_samples=8)
216
+ save_frames(samples, Path("outputs/adk"))
217
+
218
+ # Propagate each adenylate kinase conformation by 1 ns.
219
+ endpoints = emu.propagate(samples, lag_time_ns=1.0)
220
+ save_frames(endpoints, Path("outputs/adk-propagated"))
221
+ ```
222
+
223
+ `save_frames()` writes `samples.xtc` and `topology.cif` to each output
224
+ directory.
225
+
226
+ See [define molecular systems](docs/system.md), [predict and sample
227
+ structures](docs/sampling.md), and [propagate structures](docs/propagation.md)
228
+ for more options. The functional Python APIs provide
229
+ [high-level sampling](docs/functional-api.md#sample-structures) and
230
+ [propagation](docs/functional-api.md#propagate-structures) functions similar to
231
+ the CLI commands, plus lower-level workflows for explicit control of
232
+ embeddings and samplers.
233
+
234
+ ## More workflows
235
+
236
+ ### Steer sampling
237
+
238
+ Choose a steering sampler to favour structures with desired properties:
239
+
240
+ ```bash
241
+ bioemu2 sample-equilibrium \
242
+ --system src/bioemu2/assets/examples/systems/myoglobin_101m.yaml \
243
+ --out outputs/myoglobin \
244
+ --sampler edm_gd
245
+ ```
246
+
247
+ See [steer sampling](docs/steering.md) and [define steering
248
+ potentials](docs/potentials.md) to select an algorithm and configure its
249
+ objective.
250
+
251
+ ### Sample large structures across multiple GPUs
252
+
253
+ Context parallelism distributes the model's pair representations across a
254
+ square mesh of GPUs, reducing the memory required on each device:
255
+
256
+ ```bash
257
+ CUDA_VISIBLE_DEVICES=0,1,2,3 \
258
+ torchrun --standalone --nproc-per-node=4 \
259
+ --module bioemu2.cp sample-equilibrium \
260
+ --pdb-id 1AKE \
261
+ --checkpoint bioemu2-preview \
262
+ --num-samples 3 \
263
+ --out outputs/adk-cp
264
+ ```
265
+
266
+ See [sample large structures with context
267
+ parallelism](docs/context-parallel-sampling.md) for mesh requirements, the
268
+ Python API, and current limitations.
269
+
270
+ ### Inspect and subsample MSAs
271
+
272
+ Retrieve an alignment from ColabFold or load your own A3M files, inspect the
273
+ raw sequences, and select or randomly subsample rows with Python indexing.
274
+ Pass the resulting `MSA` directly to `EmuSession` to control the alignment
275
+ used for structure prediction, equilibrium sampling, or propagation without
276
+ manually constructing embeddings.
277
+
278
+ See [use multiple sequence alignments](docs/msa.md) for local-file loading,
279
+ [MSA subsampling](docs/msa.md#select-rows-and-subsample-an-msa), and
280
+ [custom-MSA inference](docs/msa.md#use-a-custom-msa-for-inference).
281
+
282
+ ## Tutorials
283
+
284
+ - [Define molecular systems](docs/system.md) from PDB IDs, structure files, or sequences.
285
+ - [Convert Boltz input YAML](docs/boltz.md) into systems you can sample from.
286
+ - [Customise multiple-sequence alignments](docs/msa.md) used to condition sampling.
287
+ - [Predict and sample structures](docs/sampling.md) for a system.
288
+ - [Improve sampling speed and scale to bigger systems](docs/speed-and-capacity.md) by choosing batch size and precision.
289
+ - [Load a model checkpoint](docs/checkpoints.md), choose where checkpoints are cached, or use a local checkpoint.
290
+ - [Propagate structures through time](docs/propagation.md).
291
+ - [Reuse trunk embeddings](docs/reuse-embeddings.md) across sampling and propagation calls, in memory or on disk.
292
+ - [Use the functional Python API](docs/functional-api.md) to sample and propagate without an `EmuSession`.
293
+ - [Steer sampling](docs/steering.md) towards structures with desired properties.
294
+ - [Define steering potentials](docs/potentials.md) that encode the properties you want to favour.
295
+ - [Sample large structures](docs/context-parallel-sampling.md) across multiple GPUs with context parallelism.
296
+
297
+ ## Notebook quickstart
298
+
299
+ The [`examples`](examples) directory contains complete workflows as regular
300
+ Python files split into executable cells with `# %%` markers. This Python
301
+ Interactive format provides notebook-style cells while keeping the examples as
302
+ plain `.py` files. They require a Linux machine with a supported NVIDIA GPU and
303
+ network access.
304
+
305
+ Activate a Python environment as described in [installation](#installation).
306
+ From the repository root, install BioEmu2 with the visualization dependencies
307
+ and register a kernel. The `viz` extra includes `ipykernel` and `py3Dmol`:
308
+
309
+ ```bash
310
+ python -m pip install ".[cuda,viz]"
311
+ python -m ipykernel install --user --name bioemu2 --display-name "BioEmu2"
312
+ ```
313
+
314
+ On compute capability 7.0 GPUs, such as NVIDIA V100, replace the first command
315
+ with:
316
+
317
+ ```bash
318
+ python -m pip install ".[cuda12,viz]" \
319
+ --extra-index-url https://download.pytorch.org/whl/cu126
320
+ ```
321
+
322
+ Select the **BioEmu2** kernel in Jupyter or VS Code. If registration succeeds
323
+ but the kernel does not appear in VS Code, restart VS Code and try again.
324
+ To check that it uses the environment you installed into, run
325
+ `import sys; print(sys.executable)` in a cell. The examples load system YAML
326
+ from the installed `bioemu2` package,
327
+ regardless of the working directory. Relative `outputs/` paths are written
328
+ under the current working directory.
329
+
330
+ Open a single-GPU example in VS Code and select **Run Cell**, or run **Jupyter:
331
+ Run Current File in Interactive Window** from the Command Palette. Launch the
332
+ context-parallel example with `torchrun` as described in the
333
+ [context-parallel sampling guide](docs/context-parallel-sampling.md#python-api).
334
+
335
+ To run without VS Code, convert an example with
336
+ [Jupytext](https://jupytext.readthedocs.io/) and open it in JupyterLab:
337
+
338
+ ```bash
339
+ python -m pip install jupytext jupyterlab
340
+ jupytext --to notebook examples/sample_structures.py
341
+ jupyter lab examples/sample_structures.ipynb
342
+ ```
343
+
344
+ Select the **BioEmu2** kernel, then run individual cells or choose **Run All
345
+ Cells** from the **Run** menu.
346
+
347
+ In VS Code, you can instead run **Jupyter: Export Current Python File as Jupyter
348
+ Notebook** from the Command Palette. Generated `.ipynb` files are ignored by
349
+ Git.
350
+
351
+ ## License
352
+
353
+ BioEmu2 is released under the [MIT License](LICENSE). It includes code and
354
+ reference data adapted from other open-source projects; the [third-party
355
+ notices](NOTICES.md) identify those components and their licenses.