enzymetk 0.0.1__tar.gz → 0.0.6__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.
- {enzymetk-0.0.1 → enzymetk-0.0.6}/PKG-INFO +61 -24
- {enzymetk-0.0.1 → enzymetk-0.0.6}/README.md +57 -16
- {enzymetk-0.0.1 → enzymetk-0.0.6}/enzymetk/__init__.py +1 -24
- {enzymetk-0.0.1 → enzymetk-0.0.6}/enzymetk/annotateEC_CLEAN_step.py +2 -2
- {enzymetk-0.0.1 → enzymetk-0.0.6}/enzymetk/annotateEC_CREEP_step.py +1 -1
- enzymetk-0.0.6/enzymetk/dock_boltz_step.py +46 -0
- {enzymetk-0.0.1 → enzymetk-0.0.6}/enzymetk/dock_chai_step.py +5 -3
- {enzymetk-0.0.1 → enzymetk-0.0.6}/enzymetk/dock_vina_step.py +49 -16
- enzymetk-0.0.6/enzymetk/embedprotein_esm3_step.py +71 -0
- {enzymetk-0.0.1 → enzymetk-0.0.6}/enzymetk/embedprotein_esm_step.py +5 -4
- {enzymetk-0.0.1 → enzymetk-0.0.6}/enzymetk/inpaint_ligandMPNN_step.py +3 -2
- {enzymetk-0.0.1 → enzymetk-0.0.6}/enzymetk/predict_catalyticsite_run.py +1 -1
- {enzymetk-0.0.1 → enzymetk-0.0.6}/enzymetk/predict_catalyticsite_step.py +20 -12
- {enzymetk-0.0.1 → enzymetk-0.0.6}/enzymetk/sequence_search_blast.py +33 -4
- {enzymetk-0.0.1 → enzymetk-0.0.6}/enzymetk/similarity_foldseek_step.py +26 -6
- {enzymetk-0.0.1 → enzymetk-0.0.6}/enzymetk/similarity_reaction_step.py +12 -8
- {enzymetk-0.0.1 → enzymetk-0.0.6}/enzymetk/similarity_substrate_step.py +8 -2
- {enzymetk-0.0.1 → enzymetk-0.0.6}/enzymetk/step.py +5 -3
- {enzymetk-0.0.1 → enzymetk-0.0.6}/enzymetk.egg-info/PKG-INFO +61 -24
- {enzymetk-0.0.1 → enzymetk-0.0.6}/enzymetk.egg-info/SOURCES.txt +2 -0
- enzymetk-0.0.6/enzymetk.egg-info/requires.txt +6 -0
- {enzymetk-0.0.1 → enzymetk-0.0.6}/setup.py +3 -7
- enzymetk-0.0.1/enzymetk.egg-info/requires.txt +0 -11
- {enzymetk-0.0.1 → enzymetk-0.0.6}/LICENSE +0 -0
- {enzymetk-0.0.1 → enzymetk-0.0.6}/enzymetk/annotateEC_proteinfer_step.py +0 -0
- {enzymetk-0.0.1 → enzymetk-0.0.6}/enzymetk/embedchem_chemberta_step.py +0 -0
- {enzymetk-0.0.1 → enzymetk-0.0.6}/enzymetk/embedchem_rxnfp_run.py +0 -0
- {enzymetk-0.0.1 → enzymetk-0.0.6}/enzymetk/embedchem_rxnfp_step.py +0 -0
- {enzymetk-0.0.1 → enzymetk-0.0.6}/enzymetk/embedchem_selformer_run.py +0 -0
- {enzymetk-0.0.1 → enzymetk-0.0.6}/enzymetk/embedchem_selformer_step.py +0 -0
- {enzymetk-0.0.1 → enzymetk-0.0.6}/enzymetk/embedchem_unimol_step.py +0 -0
- {enzymetk-0.0.1 → enzymetk-0.0.6}/enzymetk/esm-extract.py +0 -0
- {enzymetk-0.0.1 → enzymetk-0.0.6}/enzymetk/filter_sequence_step.py +0 -0
- {enzymetk-0.0.1 → enzymetk-0.0.6}/enzymetk/filter_structure_step.py +0 -0
- {enzymetk-0.0.1 → enzymetk-0.0.6}/enzymetk/generate_msa_step.py +0 -0
- {enzymetk-0.0.1 → enzymetk-0.0.6}/enzymetk/generate_oligopool_step.py +0 -0
- {enzymetk-0.0.1 → enzymetk-0.0.6}/enzymetk/generate_tree_step.py +0 -0
- {enzymetk-0.0.1 → enzymetk-0.0.6}/enzymetk/main.py +0 -0
- {enzymetk-0.0.1 → enzymetk-0.0.6}/enzymetk/metagenomics_porechop_trim_reads_step.py +0 -0
- {enzymetk-0.0.1 → enzymetk-0.0.6}/enzymetk/metagenomics_prokka_annotate_genes.py +0 -0
- {enzymetk-0.0.1 → enzymetk-0.0.6}/enzymetk/pipeline.py +0 -0
- {enzymetk-0.0.1 → enzymetk-0.0.6}/enzymetk/predict_activity_step.py +0 -0
- {enzymetk-0.0.1 → enzymetk-0.0.6}/enzymetk/reducedim_pca_run.py +0 -0
- {enzymetk-0.0.1 → enzymetk-0.0.6}/enzymetk/reducedim_vae_run.py +0 -0
- {enzymetk-0.0.1 → enzymetk-0.0.6}/enzymetk/reducedim_vae_step.py +0 -0
- {enzymetk-0.0.1 → enzymetk-0.0.6}/enzymetk/save_step.py +0 -0
- {enzymetk-0.0.1 → enzymetk-0.0.6}/enzymetk/similarity_mmseqs_step.py +0 -0
- {enzymetk-0.0.1 → enzymetk-0.0.6}/enzymetk.egg-info/dependency_links.txt +0 -0
- {enzymetk-0.0.1 → enzymetk-0.0.6}/enzymetk.egg-info/entry_points.txt +0 -0
- {enzymetk-0.0.1 → enzymetk-0.0.6}/enzymetk.egg-info/top_level.txt +0 -0
- {enzymetk-0.0.1 → enzymetk-0.0.6}/setup.cfg +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
|
-
Metadata-Version: 2.
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
2
|
Name: enzymetk
|
|
3
|
-
Version: 0.0.
|
|
3
|
+
Version: 0.0.6
|
|
4
4
|
Home-page: https://github.com/arianemora/enzyme-tk/
|
|
5
5
|
Author: Ariane Mora
|
|
6
6
|
Author-email: ariane.n.mora@gmail.com
|
|
@@ -18,17 +18,12 @@ Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
|
|
|
18
18
|
Requires-Python: >=3.8
|
|
19
19
|
Description-Content-Type: text/markdown
|
|
20
20
|
License-File: LICENSE
|
|
21
|
-
Requires-Dist: fair-esm
|
|
22
21
|
Requires-Dist: scikit-learn
|
|
23
22
|
Requires-Dist: numpy
|
|
24
23
|
Requires-Dist: seaborn
|
|
25
24
|
Requires-Dist: sciutil
|
|
26
|
-
Requires-Dist: pandas
|
|
25
|
+
Requires-Dist: pandas
|
|
27
26
|
Requires-Dist: biopython
|
|
28
|
-
Requires-Dist: sentence_transformers
|
|
29
|
-
Requires-Dist: pubchempy
|
|
30
|
-
Requires-Dist: pyfaidx
|
|
31
|
-
Requires-Dist: spacy
|
|
32
27
|
Dynamic: author
|
|
33
28
|
Dynamic: author-email
|
|
34
29
|
Dynamic: classifier
|
|
@@ -37,6 +32,7 @@ Dynamic: description-content-type
|
|
|
37
32
|
Dynamic: home-page
|
|
38
33
|
Dynamic: keywords
|
|
39
34
|
Dynamic: license
|
|
35
|
+
Dynamic: license-file
|
|
40
36
|
Dynamic: project-url
|
|
41
37
|
Dynamic: requires-dist
|
|
42
38
|
Dynamic: requires-python
|
|
@@ -45,26 +41,36 @@ Dynamic: requires-python
|
|
|
45
41
|
|
|
46
42
|
Enzyme-tk is a collection of tools for enzyme engineering, setup as interoperable modules that act on dataframes. These modules are designed to be imported into pipelines for specific function. For this reason, `steps` as each module is called (e.g. finding similar proteins with `BLAST` would be considered a step) are designed to be as light as possible. An example of a pipeline is the [annotate-e](https://github.com/ArianeMora/annotate-e) ` pipeline, this acts to annotate a fasta with an ensemble of methods (each is designated as an Enzyme-tk step).
|
|
47
43
|
|
|
44
|
+
|
|
45
|
+
**If you have any issues installing, let me know - this has been tested only on Linux/Ubuntu. Please post an issue!**
|
|
46
|
+
|
|
48
47
|
## Installation
|
|
49
48
|
|
|
49
|
+
## Install base package to import modules
|
|
50
|
+
|
|
50
51
|
```bash
|
|
51
|
-
|
|
52
|
+
pip install enzymetk
|
|
52
53
|
```
|
|
53
54
|
|
|
54
|
-
|
|
55
|
+
### Install only the specific requirements you need (recomended)
|
|
55
56
|
|
|
57
|
+
For this clone the repo and then install the requirements for the specific modules you use
|
|
56
58
|
```bash
|
|
57
59
|
git clone git@github.com:ArianeMora/enzyme-tk.git
|
|
58
|
-
|
|
59
|
-
|
|
60
|
+
cd enzymetk/conda_envs/ # would recommend looking at thes
|
|
61
|
+
# e.g. to install all from within that folder you would do
|
|
62
|
+
source install_all.sh
|
|
60
63
|
```
|
|
61
64
|
|
|
62
65
|
## Usage
|
|
63
66
|
|
|
64
67
|
If you have any issues at all just email me using my caltech email: `amora at caltech . edu`
|
|
65
68
|
|
|
69
|
+
This is a work-in progress! e.g. some tools (e.g. proteInfer and CLEAN) require extra data to be downloaded in order to run (like model weights.) I'm working on integrating these atm, buzz me if you need this!
|
|
70
|
+
|
|
66
71
|
Here are some of the tools that have been implemented to be chained together as a pipeline:
|
|
67
72
|
|
|
73
|
+
[boltz2](https://github.com/jwohlwend/boltz)
|
|
68
74
|
[mmseqs2](https://github.com/soedinglab/mmseqs2)
|
|
69
75
|
[foldseek](https://github.com/steineggerlab/foldseek)
|
|
70
76
|
[diamond](https://github.com/bbuchfink/diamond)
|
|
@@ -83,6 +89,7 @@ Here are some of the tools that have been implemented to be chained together as
|
|
|
83
89
|
[fasttree](https://morgannprice.github.io/fasttree/)
|
|
84
90
|
[Porechop](https://github.com/rrwick/Porechop)
|
|
85
91
|
[prokka](https://github.com/tseemann/prokka)
|
|
92
|
+
|
|
86
93
|
## Things to note
|
|
87
94
|
|
|
88
95
|
All the tools use the conda env of `enzymetk` by default.
|
|
@@ -114,6 +121,8 @@ The steps are the main building blocks of the pipeline. They are responsible for
|
|
|
114
121
|
|
|
115
122
|
BLAST is a tool for searching a database of sequences for similar sequences. Here you can either pass a database that you have already created or pass the sequences as part of your dataframe and pass the label column (this needs to have two values: reference and query) reference refers to sequences that you want to search against and query refers to sequences that you want to search for.
|
|
116
123
|
|
|
124
|
+
Note you need to have installed the BLAST environment.
|
|
125
|
+
|
|
117
126
|
```python
|
|
118
127
|
id_col = 'Entry'
|
|
119
128
|
seq_col = 'Sequence'
|
|
@@ -142,6 +151,34 @@ df = pd.DataFrame(rows, columns=[id_col, seq_col])
|
|
|
142
151
|
print(df)
|
|
143
152
|
df << (ActiveSitePred(id_col, seq_col, squidly_dir, num_threads) >> Save('tmp/squidly_as_pred.pkl'))
|
|
144
153
|
|
|
154
|
+
```
|
|
155
|
+
### Boltz2
|
|
156
|
+
|
|
157
|
+
Boltz2 is a model for predicting structures. Note you need docko installed as I run via that.
|
|
158
|
+
|
|
159
|
+
Below is an example using boltz with 4 threads, and uses a cofactor (intermediate in this case). Just set to be None for a single substrate version.
|
|
160
|
+
```
|
|
161
|
+
import sys
|
|
162
|
+
from enzymetk.dock_boltz_step import Boltz
|
|
163
|
+
from enzymetk.save_step import Save
|
|
164
|
+
import pandas as pd
|
|
165
|
+
import os
|
|
166
|
+
os.environ['MKL_THREADING_LAYER'] = 'GNU'
|
|
167
|
+
|
|
168
|
+
output_dir = 'tmp/'
|
|
169
|
+
num_threads = 4
|
|
170
|
+
id_col = 'Entry'
|
|
171
|
+
seq_col = 'Sequence'
|
|
172
|
+
substrate_col = 'Substrate'
|
|
173
|
+
intermediate_col = 'Intermediate'
|
|
174
|
+
|
|
175
|
+
rows = [['P0DP23_boltz_8999', 'MALWMRLLPLLALLALWGPDPAAAMALWMRLLPLLALLALWGPDPAAAMALWMRLLPLLALLALWGPDPAAA', 'CCCCC(CC)COC(=O)C1=CC=CC=C1C(=O)OCC(CC)CCCC', 'CC1=C(C2=CC3=C(C(=C([N-]3)C=C4C(=C(C(=N4)C=C5C(=C(C(=N5)C=C1[N-]2)C)C=C)C)C=C)C)CCC(=O)[O-])CCC(=O)[O-].[Fe]'],
|
|
176
|
+
['P0DP24_boltz_p1', 'MALWMRLLPLLALLALWGPDPAAAMALWMRLLPLLALLALWGPDPAAAMALWMRLLPLLALLALWGPDPAAA', 'CCCCC(CC)COC(=O)C1=CC=CC=C1C(=O)OCC(CC)CCCC', 'CC1=C(C2=CC3=C(C(=C([N-]3)C=C4C(=C(C(=N4)C=C5C(=C(C(=N5)C=C1[N-]2)C)C=C)C)C=C)C)CCC(=O)[O-])CCC(=O)[O-].[Fe]'],
|
|
177
|
+
['P0DP23_boltz_p2', 'MALWMRLLPLLALLALWGPDPAAAMALWMRLLPLLALLALWGPDPAAAMALWMRLLPLLALLALWGPDPAAA', 'CCCCC(CC)COC(=O)C1=CC=CC=C1C(=O)OCC(CC)CCCC', 'CC1=C(C2=CC3=C(C(=C([N-]3)C=C4C(=C(C(=N4)C=C5C(=C(C(=N5)C=C1[N-]2)C)C=C)C)C=C)C)CCC(=O)[O-])CCC(=O)[O-].[Fe]'],
|
|
178
|
+
['P0DP24_boltz_p3', 'MALWMRLLPLLALLALWGPDPAAAMALWMRLLPLLALLALWGPDPAAAMALWMRLLPLLALLALWGPDPAAA', 'CCCCC(CC)COC(=O)C1=CC=CC=C1C(=O)OCC(CC)CCCC', 'CC1=C(C2=CC3=C(C(=C([N-]3)C=C4C(=C(C(=N4)C=C5C(=C(C(=N5)C=C1[N-]2)C)C=C)C)C=C)C)CCC(=O)[O-])CCC(=O)[O-].[Fe]'],
|
|
179
|
+
['P0DP24_boltz_p4', 'MALWMRLLPLLALLALWGPDPAAAMALWMRLLPLLALLALWGPDPAAAMALWMRLLPLLALLALWGPDPAAA', 'CCCCC(CC)COC(=O)C1=CC=CC=C1C(=O)OCC(CC)CCCC', 'CC1=C(C2=CC3=C(C(=C([N-]3)C=C4C(=C(C(=N4)C=C5C(=C(C(=N5)C=C1[N-]2)C)C=C)C)C=C)C)CCC(=O)[O-])CCC(=O)[O-].[Fe]']]
|
|
180
|
+
df = pd.DataFrame(rows, columns=[id_col, seq_col, substrate_col, intermediate_col])
|
|
181
|
+
df << (Boltz(id_col, seq_col, substrate_col, intermediate_col, f'{output_dir}', num_threads) >> Save(f'{output_dir}test.pkl'))
|
|
145
182
|
```
|
|
146
183
|
|
|
147
184
|
### Chai
|
|
@@ -169,8 +206,8 @@ df << (Chai(id_col, seq_col, substrate_col, f'{output_dir}', num_threads) >> Sav
|
|
|
169
206
|
ChemBERTa2 encodes reactions and SMILES strings into a vector space. Note this requires the base environment, i.e. `enzymetk` conda env.
|
|
170
207
|
|
|
171
208
|
```python
|
|
172
|
-
from
|
|
173
|
-
from
|
|
209
|
+
from enzymetk.embedchem_chemberta_step import ChemBERT
|
|
210
|
+
from enzymetk.save_step import Save
|
|
174
211
|
|
|
175
212
|
output_dir = 'tmp/'
|
|
176
213
|
num_threads = 1
|
|
@@ -180,7 +217,7 @@ substrate_col = 'Substrate'
|
|
|
180
217
|
rows = [['P0DP23', 'MALWMRLLPLLALLALWGPDPAAAMALWMRLLPLLALLALWGPDPAAAMALWMRLLPLLALLALWGPDPAAA', 'CCCCC(CC)COC(=O)C1=CC=CC=C1C(=O)OCC(CC)CCCC'],
|
|
181
218
|
['P0DP24', 'MALWMRLLPLLALLALWGPDPAAAMALWMRLLPLLALLALWGPDPAAAMALWMRLLPLLALLALWGPDPAAA', 'CCCCC(CC)COC(=O)C1=CC=CC=C1C(=O)OCC(CC)CCCC']]
|
|
182
219
|
df = pd.DataFrame(rows, columns=[id_col, seq_col, substrate_col])
|
|
183
|
-
df << (ChemBERT(id_col, substrate_col, num_threads) >> Save(f'{output_dir}chemberta.pkl'))
|
|
220
|
+
new_df = (df << (ChemBERT(id_col, substrate_col, num_threads) >> Save(f'{output_dir}chemberta.pkl')))
|
|
184
221
|
```
|
|
185
222
|
|
|
186
223
|
### CLEAN
|
|
@@ -206,11 +243,11 @@ df << (CLEAN(id_col, seq_col, clean_dir, num_threads=num_threads) >> Save(f'clea
|
|
|
206
243
|
```
|
|
207
244
|
### ClustalOmega
|
|
208
245
|
|
|
209
|
-
ClustalOmega is a tool for aligning a set of sequences. This gets installed to the system (expecting a linux machine) and added to the bash path.
|
|
246
|
+
ClustalOmega is a tool for aligning a set of sequences. This gets installed to the system (expecting a linux machine) and added to the bash path. You need to have installed it first (check out the `conda_envs` directory in enzymetk.)
|
|
210
247
|
|
|
211
248
|
```python
|
|
212
|
-
from
|
|
213
|
-
from
|
|
249
|
+
from enzymetk.generate_msa_step import ClustalOmega
|
|
250
|
+
from enzymetk.save_step import Save
|
|
214
251
|
import pandas as pd
|
|
215
252
|
|
|
216
253
|
id_col = 'Entry'
|
|
@@ -230,8 +267,8 @@ df << (ClustalOmega(id_col, seq_col) >> Save('tmp/clustalomega_test.pkl'))
|
|
|
230
267
|
CREEP is a tool for predicting the EC number of a reaction. At the moment it only supports reactions to EC however we are extending this to other modalities.
|
|
231
268
|
|
|
232
269
|
```python
|
|
233
|
-
from
|
|
234
|
-
from
|
|
270
|
+
from enzymetk.annotateEC_CREEP_step import CREEP
|
|
271
|
+
from enzymetk.save_step import Save
|
|
235
272
|
import pandas as pd
|
|
236
273
|
|
|
237
274
|
# CREEP expects you to have downloaded the data from the zotero page and put it in the data/CREEP folder
|
|
@@ -252,8 +289,8 @@ df << (CREEP(id_col, reaction_col, CREEP_cache_dir='/disk1/share/software/CREEP/
|
|
|
252
289
|
EmbedESM is a tool for embedding a set of sequences using ESM2.
|
|
253
290
|
|
|
254
291
|
```python
|
|
255
|
-
from
|
|
256
|
-
from
|
|
292
|
+
from enzymetk.embedprotein_esm_step import EmbedESM
|
|
293
|
+
from enzymetk.save_step import Save
|
|
257
294
|
import pandas as pd
|
|
258
295
|
|
|
259
296
|
id_col = 'Entry'
|
|
@@ -280,8 +317,8 @@ If you pass a database, you need to pass the path to the database.
|
|
|
280
317
|
The columns expect a path to a pdb file i.e. the output from the `Chai` step.
|
|
281
318
|
|
|
282
319
|
```python
|
|
283
|
-
from
|
|
284
|
-
from
|
|
320
|
+
from enzymetk.similarity_foldseek_step import FoldSeek
|
|
321
|
+
from enzymetk.save_step import Save
|
|
285
322
|
import pandas as pd
|
|
286
323
|
|
|
287
324
|
# id_col: str, seq_col: str, proteinfer_dir: str,
|
|
@@ -2,26 +2,36 @@
|
|
|
2
2
|
|
|
3
3
|
Enzyme-tk is a collection of tools for enzyme engineering, setup as interoperable modules that act on dataframes. These modules are designed to be imported into pipelines for specific function. For this reason, `steps` as each module is called (e.g. finding similar proteins with `BLAST` would be considered a step) are designed to be as light as possible. An example of a pipeline is the [annotate-e](https://github.com/ArianeMora/annotate-e) ` pipeline, this acts to annotate a fasta with an ensemble of methods (each is designated as an Enzyme-tk step).
|
|
4
4
|
|
|
5
|
+
|
|
6
|
+
**If you have any issues installing, let me know - this has been tested only on Linux/Ubuntu. Please post an issue!**
|
|
7
|
+
|
|
5
8
|
## Installation
|
|
6
9
|
|
|
10
|
+
## Install base package to import modules
|
|
11
|
+
|
|
7
12
|
```bash
|
|
8
|
-
|
|
13
|
+
pip install enzymetk
|
|
9
14
|
```
|
|
10
15
|
|
|
11
|
-
|
|
16
|
+
### Install only the specific requirements you need (recomended)
|
|
12
17
|
|
|
18
|
+
For this clone the repo and then install the requirements for the specific modules you use
|
|
13
19
|
```bash
|
|
14
20
|
git clone git@github.com:ArianeMora/enzyme-tk.git
|
|
15
|
-
|
|
16
|
-
|
|
21
|
+
cd enzymetk/conda_envs/ # would recommend looking at thes
|
|
22
|
+
# e.g. to install all from within that folder you would do
|
|
23
|
+
source install_all.sh
|
|
17
24
|
```
|
|
18
25
|
|
|
19
26
|
## Usage
|
|
20
27
|
|
|
21
28
|
If you have any issues at all just email me using my caltech email: `amora at caltech . edu`
|
|
22
29
|
|
|
30
|
+
This is a work-in progress! e.g. some tools (e.g. proteInfer and CLEAN) require extra data to be downloaded in order to run (like model weights.) I'm working on integrating these atm, buzz me if you need this!
|
|
31
|
+
|
|
23
32
|
Here are some of the tools that have been implemented to be chained together as a pipeline:
|
|
24
33
|
|
|
34
|
+
[boltz2](https://github.com/jwohlwend/boltz)
|
|
25
35
|
[mmseqs2](https://github.com/soedinglab/mmseqs2)
|
|
26
36
|
[foldseek](https://github.com/steineggerlab/foldseek)
|
|
27
37
|
[diamond](https://github.com/bbuchfink/diamond)
|
|
@@ -40,6 +50,7 @@ Here are some of the tools that have been implemented to be chained together as
|
|
|
40
50
|
[fasttree](https://morgannprice.github.io/fasttree/)
|
|
41
51
|
[Porechop](https://github.com/rrwick/Porechop)
|
|
42
52
|
[prokka](https://github.com/tseemann/prokka)
|
|
53
|
+
|
|
43
54
|
## Things to note
|
|
44
55
|
|
|
45
56
|
All the tools use the conda env of `enzymetk` by default.
|
|
@@ -71,6 +82,8 @@ The steps are the main building blocks of the pipeline. They are responsible for
|
|
|
71
82
|
|
|
72
83
|
BLAST is a tool for searching a database of sequences for similar sequences. Here you can either pass a database that you have already created or pass the sequences as part of your dataframe and pass the label column (this needs to have two values: reference and query) reference refers to sequences that you want to search against and query refers to sequences that you want to search for.
|
|
73
84
|
|
|
85
|
+
Note you need to have installed the BLAST environment.
|
|
86
|
+
|
|
74
87
|
```python
|
|
75
88
|
id_col = 'Entry'
|
|
76
89
|
seq_col = 'Sequence'
|
|
@@ -99,6 +112,34 @@ df = pd.DataFrame(rows, columns=[id_col, seq_col])
|
|
|
99
112
|
print(df)
|
|
100
113
|
df << (ActiveSitePred(id_col, seq_col, squidly_dir, num_threads) >> Save('tmp/squidly_as_pred.pkl'))
|
|
101
114
|
|
|
115
|
+
```
|
|
116
|
+
### Boltz2
|
|
117
|
+
|
|
118
|
+
Boltz2 is a model for predicting structures. Note you need docko installed as I run via that.
|
|
119
|
+
|
|
120
|
+
Below is an example using boltz with 4 threads, and uses a cofactor (intermediate in this case). Just set to be None for a single substrate version.
|
|
121
|
+
```
|
|
122
|
+
import sys
|
|
123
|
+
from enzymetk.dock_boltz_step import Boltz
|
|
124
|
+
from enzymetk.save_step import Save
|
|
125
|
+
import pandas as pd
|
|
126
|
+
import os
|
|
127
|
+
os.environ['MKL_THREADING_LAYER'] = 'GNU'
|
|
128
|
+
|
|
129
|
+
output_dir = 'tmp/'
|
|
130
|
+
num_threads = 4
|
|
131
|
+
id_col = 'Entry'
|
|
132
|
+
seq_col = 'Sequence'
|
|
133
|
+
substrate_col = 'Substrate'
|
|
134
|
+
intermediate_col = 'Intermediate'
|
|
135
|
+
|
|
136
|
+
rows = [['P0DP23_boltz_8999', 'MALWMRLLPLLALLALWGPDPAAAMALWMRLLPLLALLALWGPDPAAAMALWMRLLPLLALLALWGPDPAAA', 'CCCCC(CC)COC(=O)C1=CC=CC=C1C(=O)OCC(CC)CCCC', 'CC1=C(C2=CC3=C(C(=C([N-]3)C=C4C(=C(C(=N4)C=C5C(=C(C(=N5)C=C1[N-]2)C)C=C)C)C=C)C)CCC(=O)[O-])CCC(=O)[O-].[Fe]'],
|
|
137
|
+
['P0DP24_boltz_p1', 'MALWMRLLPLLALLALWGPDPAAAMALWMRLLPLLALLALWGPDPAAAMALWMRLLPLLALLALWGPDPAAA', 'CCCCC(CC)COC(=O)C1=CC=CC=C1C(=O)OCC(CC)CCCC', 'CC1=C(C2=CC3=C(C(=C([N-]3)C=C4C(=C(C(=N4)C=C5C(=C(C(=N5)C=C1[N-]2)C)C=C)C)C=C)C)CCC(=O)[O-])CCC(=O)[O-].[Fe]'],
|
|
138
|
+
['P0DP23_boltz_p2', 'MALWMRLLPLLALLALWGPDPAAAMALWMRLLPLLALLALWGPDPAAAMALWMRLLPLLALLALWGPDPAAA', 'CCCCC(CC)COC(=O)C1=CC=CC=C1C(=O)OCC(CC)CCCC', 'CC1=C(C2=CC3=C(C(=C([N-]3)C=C4C(=C(C(=N4)C=C5C(=C(C(=N5)C=C1[N-]2)C)C=C)C)C=C)C)CCC(=O)[O-])CCC(=O)[O-].[Fe]'],
|
|
139
|
+
['P0DP24_boltz_p3', 'MALWMRLLPLLALLALWGPDPAAAMALWMRLLPLLALLALWGPDPAAAMALWMRLLPLLALLALWGPDPAAA', 'CCCCC(CC)COC(=O)C1=CC=CC=C1C(=O)OCC(CC)CCCC', 'CC1=C(C2=CC3=C(C(=C([N-]3)C=C4C(=C(C(=N4)C=C5C(=C(C(=N5)C=C1[N-]2)C)C=C)C)C=C)C)CCC(=O)[O-])CCC(=O)[O-].[Fe]'],
|
|
140
|
+
['P0DP24_boltz_p4', 'MALWMRLLPLLALLALWGPDPAAAMALWMRLLPLLALLALWGPDPAAAMALWMRLLPLLALLALWGPDPAAA', 'CCCCC(CC)COC(=O)C1=CC=CC=C1C(=O)OCC(CC)CCCC', 'CC1=C(C2=CC3=C(C(=C([N-]3)C=C4C(=C(C(=N4)C=C5C(=C(C(=N5)C=C1[N-]2)C)C=C)C)C=C)C)CCC(=O)[O-])CCC(=O)[O-].[Fe]']]
|
|
141
|
+
df = pd.DataFrame(rows, columns=[id_col, seq_col, substrate_col, intermediate_col])
|
|
142
|
+
df << (Boltz(id_col, seq_col, substrate_col, intermediate_col, f'{output_dir}', num_threads) >> Save(f'{output_dir}test.pkl'))
|
|
102
143
|
```
|
|
103
144
|
|
|
104
145
|
### Chai
|
|
@@ -126,8 +167,8 @@ df << (Chai(id_col, seq_col, substrate_col, f'{output_dir}', num_threads) >> Sav
|
|
|
126
167
|
ChemBERTa2 encodes reactions and SMILES strings into a vector space. Note this requires the base environment, i.e. `enzymetk` conda env.
|
|
127
168
|
|
|
128
169
|
```python
|
|
129
|
-
from
|
|
130
|
-
from
|
|
170
|
+
from enzymetk.embedchem_chemberta_step import ChemBERT
|
|
171
|
+
from enzymetk.save_step import Save
|
|
131
172
|
|
|
132
173
|
output_dir = 'tmp/'
|
|
133
174
|
num_threads = 1
|
|
@@ -137,7 +178,7 @@ substrate_col = 'Substrate'
|
|
|
137
178
|
rows = [['P0DP23', 'MALWMRLLPLLALLALWGPDPAAAMALWMRLLPLLALLALWGPDPAAAMALWMRLLPLLALLALWGPDPAAA', 'CCCCC(CC)COC(=O)C1=CC=CC=C1C(=O)OCC(CC)CCCC'],
|
|
138
179
|
['P0DP24', 'MALWMRLLPLLALLALWGPDPAAAMALWMRLLPLLALLALWGPDPAAAMALWMRLLPLLALLALWGPDPAAA', 'CCCCC(CC)COC(=O)C1=CC=CC=C1C(=O)OCC(CC)CCCC']]
|
|
139
180
|
df = pd.DataFrame(rows, columns=[id_col, seq_col, substrate_col])
|
|
140
|
-
df << (ChemBERT(id_col, substrate_col, num_threads) >> Save(f'{output_dir}chemberta.pkl'))
|
|
181
|
+
new_df = (df << (ChemBERT(id_col, substrate_col, num_threads) >> Save(f'{output_dir}chemberta.pkl')))
|
|
141
182
|
```
|
|
142
183
|
|
|
143
184
|
### CLEAN
|
|
@@ -163,11 +204,11 @@ df << (CLEAN(id_col, seq_col, clean_dir, num_threads=num_threads) >> Save(f'clea
|
|
|
163
204
|
```
|
|
164
205
|
### ClustalOmega
|
|
165
206
|
|
|
166
|
-
ClustalOmega is a tool for aligning a set of sequences. This gets installed to the system (expecting a linux machine) and added to the bash path.
|
|
207
|
+
ClustalOmega is a tool for aligning a set of sequences. This gets installed to the system (expecting a linux machine) and added to the bash path. You need to have installed it first (check out the `conda_envs` directory in enzymetk.)
|
|
167
208
|
|
|
168
209
|
```python
|
|
169
|
-
from
|
|
170
|
-
from
|
|
210
|
+
from enzymetk.generate_msa_step import ClustalOmega
|
|
211
|
+
from enzymetk.save_step import Save
|
|
171
212
|
import pandas as pd
|
|
172
213
|
|
|
173
214
|
id_col = 'Entry'
|
|
@@ -187,8 +228,8 @@ df << (ClustalOmega(id_col, seq_col) >> Save('tmp/clustalomega_test.pkl'))
|
|
|
187
228
|
CREEP is a tool for predicting the EC number of a reaction. At the moment it only supports reactions to EC however we are extending this to other modalities.
|
|
188
229
|
|
|
189
230
|
```python
|
|
190
|
-
from
|
|
191
|
-
from
|
|
231
|
+
from enzymetk.annotateEC_CREEP_step import CREEP
|
|
232
|
+
from enzymetk.save_step import Save
|
|
192
233
|
import pandas as pd
|
|
193
234
|
|
|
194
235
|
# CREEP expects you to have downloaded the data from the zotero page and put it in the data/CREEP folder
|
|
@@ -209,8 +250,8 @@ df << (CREEP(id_col, reaction_col, CREEP_cache_dir='/disk1/share/software/CREEP/
|
|
|
209
250
|
EmbedESM is a tool for embedding a set of sequences using ESM2.
|
|
210
251
|
|
|
211
252
|
```python
|
|
212
|
-
from
|
|
213
|
-
from
|
|
253
|
+
from enzymetk.embedprotein_esm_step import EmbedESM
|
|
254
|
+
from enzymetk.save_step import Save
|
|
214
255
|
import pandas as pd
|
|
215
256
|
|
|
216
257
|
id_col = 'Entry'
|
|
@@ -237,8 +278,8 @@ If you pass a database, you need to pass the path to the database.
|
|
|
237
278
|
The columns expect a path to a pdb file i.e. the output from the `Chai` step.
|
|
238
279
|
|
|
239
280
|
```python
|
|
240
|
-
from
|
|
241
|
-
from
|
|
281
|
+
from enzymetk.similarity_foldseek_step import FoldSeek
|
|
282
|
+
from enzymetk.save_step import Save
|
|
242
283
|
import pandas as pd
|
|
243
284
|
|
|
244
285
|
# id_col: str, seq_col: str, proteinfer_dir: str,
|
|
@@ -22,34 +22,11 @@ Date: March 2025
|
|
|
22
22
|
__title__ = 'enzymetk'
|
|
23
23
|
__description__ = 'Toolkit for enzymes and what not'
|
|
24
24
|
__url__ = 'https://github.com/arianemora/enzyme-tk/'
|
|
25
|
-
__version__ = '0.0.
|
|
25
|
+
__version__ = '0.0.6'
|
|
26
26
|
__author__ = 'Ariane Mora'
|
|
27
27
|
__author_email__ = 'ariane.n.mora@gmail.com'
|
|
28
28
|
__license__ = 'GPL3'
|
|
29
29
|
|
|
30
|
-
# from enzymetk.step import *
|
|
31
|
-
# from enzymetk.generate_msa_step import ClustalOmega
|
|
32
|
-
# from enzymetk.annotateEC_CLEAN_step import CLEAN
|
|
33
|
-
# from enzymetk.annotateEC_proteinfer_step import ProteInfer
|
|
34
|
-
# from enzymetk.dock_chai_step import Chai
|
|
35
|
-
# from enzymetk.dock_vina_step import Vina
|
|
36
|
-
# from enzymetk.embedchem_chemberta_step import ChemBERT
|
|
37
|
-
# from enzymetk.embedchem_rxnfp_step import RxnFP
|
|
38
|
-
# from enzymetk.embedchem_selformer_step import SelFormer
|
|
39
|
-
# from enzymetk.embedchem_unimol_step import UniMol
|
|
40
|
-
# from enzymetk.embedprotein_esm_step import EmbedESM
|
|
41
|
-
# from enzymetk.generate_tree_step import FastTree
|
|
42
|
-
# from enzymetk.inpaint_ligandMPNN_step import LigandMPNN
|
|
43
|
-
# from enzymetk.metagenomics_porechop_trim_reads_step import PoreChop
|
|
44
|
-
# from enzymetk.metagenomics_prokka_annotate_genes import Prokka
|
|
45
|
-
# #from enzymetk.predict_activity_step import
|
|
46
|
-
# from enzymetk.predict_catalyticsite_step import ActiveSitePred
|
|
47
|
-
# from enzymetk.sequence_search_blast import BLAST
|
|
48
|
-
# from enzymetk.similarity_foldseek_step import FoldSeek
|
|
49
|
-
# from enzymetk.similarity_mmseqs_step import MMseqs
|
|
50
|
-
# from enzymetk.similarity_reaction_step import ReactionDist
|
|
51
|
-
# from enzymetk.similarity_substrate_step import SubstrateDist
|
|
52
|
-
|
|
53
30
|
|
|
54
31
|
|
|
55
32
|
|
|
@@ -116,7 +116,7 @@ class CLEAN(Step):
|
|
|
116
116
|
print(output_filenames)
|
|
117
117
|
for sub_df in output_filenames:
|
|
118
118
|
df = pd.concat([df, sub_df])
|
|
119
|
-
return df
|
|
119
|
+
return self.__filter_df(df)
|
|
120
120
|
else:
|
|
121
|
-
return self.__execute([df, tmp_dir])
|
|
121
|
+
return self.__filter_df(self.__execute([df, tmp_dir]))
|
|
122
122
|
return df
|
|
@@ -38,7 +38,7 @@ class CREEP(Step):
|
|
|
38
38
|
self.args_extract = args_extract
|
|
39
39
|
self.args_retrieval = args_retrieval
|
|
40
40
|
|
|
41
|
-
def __execute(self, df: pd.DataFrame, tmp_dir: str)
|
|
41
|
+
def __execute(self, df: pd.DataFrame, tmp_dir: str):
|
|
42
42
|
tmp_dir = '/disk1/ariane/vscode/degradeo/pipeline/tmp/'
|
|
43
43
|
input_filename = f'{tmp_dir}/creepasjkdkajshdkja.csv'
|
|
44
44
|
df.to_csv(input_filename, index=False)
|
|
@@ -0,0 +1,46 @@
|
|
|
1
|
+
from enzymetk.step import Step
|
|
2
|
+
import pandas as pd
|
|
3
|
+
from docko.boltz import run_boltz_affinity
|
|
4
|
+
import logging
|
|
5
|
+
import numpy as np
|
|
6
|
+
from multiprocessing.dummy import Pool as ThreadPool
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
logger = logging.getLogger(__name__)
|
|
10
|
+
logger.setLevel(logging.INFO)
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
class Boltz(Step):
|
|
14
|
+
|
|
15
|
+
def __init__(self, id_col: str, seq_col: str, substrate_col: str, intermediate_col: str, output_dir: str, num_threads: int):
|
|
16
|
+
self.id_col = id_col
|
|
17
|
+
self.seq_col = seq_col
|
|
18
|
+
self.substrate_col = substrate_col
|
|
19
|
+
self.intermediate_col = intermediate_col
|
|
20
|
+
self.output_dir = output_dir or None
|
|
21
|
+
self.num_threads = num_threads or 1
|
|
22
|
+
|
|
23
|
+
def __execute(self, df: pd.DataFrame) -> pd.DataFrame:
|
|
24
|
+
output_filenames = []
|
|
25
|
+
|
|
26
|
+
for run_id, seq, substrate, intermediate in df[[self.id_col, self.seq_col, self.substrate_col, self.intermediate_col]].values:
|
|
27
|
+
# Might have an issue if the things are not correctly installed in the same dicrectory
|
|
28
|
+
if not isinstance(substrate, str):
|
|
29
|
+
substrate = ''
|
|
30
|
+
print(run_id, seq, substrate)
|
|
31
|
+
run_boltz_affinity(run_id, seq, substrate, self.output_dir, intermediate)
|
|
32
|
+
output_filenames.append(f'{self.output_dir}/{run_id}/')
|
|
33
|
+
return output_filenames
|
|
34
|
+
|
|
35
|
+
def execute(self, df: pd.DataFrame) -> pd.DataFrame:
|
|
36
|
+
if self.output_dir:
|
|
37
|
+
if self.num_threads > 1:
|
|
38
|
+
pool = ThreadPool(self.num_threads)
|
|
39
|
+
df_list = np.array_split(df, self.num_threads)
|
|
40
|
+
results = pool.map(self.__execute, df_list)
|
|
41
|
+
else:
|
|
42
|
+
results = self.__execute(df)
|
|
43
|
+
df['output_dir'] = results
|
|
44
|
+
return df
|
|
45
|
+
else:
|
|
46
|
+
print('No output directory provided')
|
|
@@ -11,16 +11,17 @@ logger.setLevel(logging.INFO)
|
|
|
11
11
|
|
|
12
12
|
class Chai(Step):
|
|
13
13
|
|
|
14
|
-
def __init__(self, id_col: str, seq_col: str, substrate_col: str, output_dir: str, num_threads: int):
|
|
14
|
+
def __init__(self, id_col: str, seq_col: str, substrate_col: str, cofactor_col: str, output_dir: str, num_threads: int):
|
|
15
15
|
self.id_col = id_col
|
|
16
16
|
self.seq_col = seq_col
|
|
17
17
|
self.substrate_col = substrate_col
|
|
18
|
+
self.cofactor_col = cofactor_col
|
|
18
19
|
self.output_dir = output_dir or None
|
|
19
20
|
self.num_threads = num_threads or 1
|
|
20
21
|
|
|
21
22
|
def __execute(self, df: pd.DataFrame, tmp_dir: str) -> pd.DataFrame:
|
|
22
23
|
output_filenames = []
|
|
23
|
-
for run_id, seq, substrate in df[[self.id_col, self.seq_col, self.substrate_col]].values:
|
|
24
|
+
for run_id, seq, substrate, cofactor in df[[self.id_col, self.seq_col, self.substrate_col, self.cofactor_col]].values:
|
|
24
25
|
# Might have an issue if the things are not correctly installed in the same dicrectory
|
|
25
26
|
if not isinstance(substrate, str):
|
|
26
27
|
substrate = ''
|
|
@@ -28,7 +29,8 @@ class Chai(Step):
|
|
|
28
29
|
run_chai(run_id, # name
|
|
29
30
|
seq, # sequence
|
|
30
31
|
substrate, # ligand as smiles
|
|
31
|
-
tmp_dir
|
|
32
|
+
tmp_dir,
|
|
33
|
+
cofactor) # cofactor as smiles
|
|
32
34
|
output_filenames.append(f'{tmp_dir}/{run_id}/')
|
|
33
35
|
return output_filenames
|
|
34
36
|
|
|
@@ -4,6 +4,7 @@ from docko.docko import *
|
|
|
4
4
|
import logging
|
|
5
5
|
import numpy as np
|
|
6
6
|
import os
|
|
7
|
+
from pathlib import Path
|
|
7
8
|
from multiprocessing.dummy import Pool as ThreadPool
|
|
8
9
|
|
|
9
10
|
logger = logging.getLogger(__name__)
|
|
@@ -21,34 +22,66 @@ class Vina(Step):
|
|
|
21
22
|
self.substrate_col = substrate_col
|
|
22
23
|
self.substrate_name_col = substrate_name_col
|
|
23
24
|
self.active_site_col = active_site_col # Expects active site residues as a string separated by |
|
|
24
|
-
self.output_dir = output_dir or None
|
|
25
|
+
self.output_dir = Path( output_dir) or None
|
|
25
26
|
self.num_threads = num_threads or 1
|
|
26
27
|
|
|
27
28
|
def __execute(self, df: pd.DataFrame) -> pd.DataFrame:
|
|
28
29
|
output_filenames = []
|
|
29
30
|
# ToDo: update to create from sequence if the path doesn't exist.
|
|
30
31
|
for label, structure_path, seq, substrate_smiles, substrate_name, residues in df[[self.id_col, self.structure_col, self.sequence_col, self.substrate_col, self.substrate_name_col, self.active_site_col]].values:
|
|
31
|
-
|
|
32
|
+
|
|
32
33
|
try:
|
|
34
|
+
structure_path = str(structure_path)
|
|
33
35
|
residues = str(residues)
|
|
34
36
|
residues = [int(r) + 1 for r in residues.split('|')]
|
|
35
|
-
|
|
36
|
-
|
|
37
|
-
|
|
38
|
-
|
|
39
|
-
|
|
40
|
-
|
|
41
|
-
|
|
42
|
-
|
|
43
|
-
|
|
44
|
-
|
|
45
|
-
|
|
37
|
+
|
|
38
|
+
label_dir = self.output_dir / label
|
|
39
|
+
label_dir.mkdir(parents=True, exist_ok=True)
|
|
40
|
+
structure_path = Path(structure_path)
|
|
41
|
+
|
|
42
|
+
if not structure_path.exists():
|
|
43
|
+
# Try to download AF2 structure
|
|
44
|
+
get_alphafold_structure(label, label_dir / f"{label}_AF2.pdb")
|
|
45
|
+
structure_path = label_dir / f"{label}_AF2.pdb"
|
|
46
|
+
|
|
47
|
+
# Skip if still not found
|
|
48
|
+
if not structure_path.exists():
|
|
49
|
+
print(f"Skipping {label}: AF2 structure not found.")
|
|
50
|
+
output_filenames.append(None)
|
|
51
|
+
continue
|
|
52
|
+
|
|
53
|
+
# Proceed with docking
|
|
54
|
+
pdb_path = label_dir / f"{label}.pdb"
|
|
55
|
+
pdbqt_path = label_dir / f"{label}.pdbqt"
|
|
56
|
+
|
|
57
|
+
clean_one_pdb(str(structure_path), str(pdb_path))
|
|
58
|
+
pdb_to_pdbqt_protein(str(pdb_path), str(pdbqt_path))
|
|
59
|
+
|
|
60
|
+
score = dock(
|
|
61
|
+
sequence='',
|
|
62
|
+
protein_name=label,
|
|
63
|
+
smiles=substrate_smiles,
|
|
64
|
+
ligand_name=substrate_name,
|
|
65
|
+
residues=residues,
|
|
66
|
+
protein_dir=str(self.output_dir),
|
|
67
|
+
ligand_dir=str(self.output_dir),
|
|
68
|
+
output_dir=str(label_dir),
|
|
69
|
+
pH=7.4,
|
|
70
|
+
method='vina',
|
|
71
|
+
size_x=10.0,
|
|
72
|
+
size_y=10.0,
|
|
73
|
+
size_z=10.0
|
|
74
|
+
)
|
|
75
|
+
|
|
76
|
+
output_filenames.append(str(pdb_path))
|
|
77
|
+
|
|
46
78
|
except Exception as e:
|
|
47
79
|
print(f'Error docking {label}: {e}')
|
|
48
80
|
output_filenames.append(None)
|
|
81
|
+
|
|
49
82
|
return output_filenames
|
|
50
|
-
|
|
51
|
-
|
|
83
|
+
|
|
84
|
+
|
|
52
85
|
def execute(self, df: pd.DataFrame) -> pd.DataFrame:
|
|
53
86
|
if self.output_dir:
|
|
54
87
|
if self.num_threads > 1:
|
|
@@ -60,4 +93,4 @@ class Vina(Step):
|
|
|
60
93
|
df['output_dir'] = results
|
|
61
94
|
return df
|
|
62
95
|
else:
|
|
63
|
-
print('No output directory provided')
|
|
96
|
+
print('No output directory provided')
|
|
@@ -0,0 +1,71 @@
|
|
|
1
|
+
# ESM 3 script
|
|
2
|
+
from esm.sdk.api import ESMProtein
|
|
3
|
+
from tempfile import TemporaryDirectory
|
|
4
|
+
import torch
|
|
5
|
+
import os
|
|
6
|
+
import pandas as pd
|
|
7
|
+
from esm.models.esm3 import ESM3
|
|
8
|
+
from esm.sdk.api import ESMProtein, SamplingConfig
|
|
9
|
+
from huggingface_hub import login
|
|
10
|
+
from enzymetk.step import Step
|
|
11
|
+
import numpy as np
|
|
12
|
+
from tqdm import tqdm
|
|
13
|
+
|
|
14
|
+
# CUDA setup
|
|
15
|
+
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID" # see issue #152
|
|
16
|
+
cuda = True
|
|
17
|
+
DEVICE = torch.device("cuda" if cuda else "cpu")
|
|
18
|
+
device = DEVICE
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
class EmbedESM3(Step):
|
|
22
|
+
|
|
23
|
+
def __init__(self, id_col: str, seq_col: str, extraction_method='mean', num_threads=1,
|
|
24
|
+
tmp_dir: str = None, env_name: str = 'enzymetk', save_tensors=False): # type: ignore
|
|
25
|
+
login()
|
|
26
|
+
self.client = ESM3.from_pretrained("esm3-open").to("cuda")
|
|
27
|
+
self.seq_col = seq_col
|
|
28
|
+
self.id_col = id_col
|
|
29
|
+
self.num_threads = num_threads or 1
|
|
30
|
+
self.extraction_method = extraction_method
|
|
31
|
+
self.tmp_dir = tmp_dir
|
|
32
|
+
self.env_name = env_name
|
|
33
|
+
self.save_tensors = save_tensors
|
|
34
|
+
|
|
35
|
+
def __execute(self, df: pd.DataFrame, tmp_dir: str) -> pd.DataFrame:
|
|
36
|
+
client = self.client
|
|
37
|
+
means = []
|
|
38
|
+
for id, seq in tqdm(df[[self.id_col, self.seq_col]].values):
|
|
39
|
+
protein = ESMProtein(
|
|
40
|
+
sequence=(
|
|
41
|
+
seq
|
|
42
|
+
)
|
|
43
|
+
)
|
|
44
|
+
protein_tensor = client.encode(protein)
|
|
45
|
+
output = client.forward_and_sample(
|
|
46
|
+
protein_tensor, SamplingConfig(return_per_residue_embeddings=True)
|
|
47
|
+
)
|
|
48
|
+
if self.save_tensors:
|
|
49
|
+
torch.save(output.per_residue_embedding, os.path.join(tmp_dir, f'{id}.pt'))
|
|
50
|
+
means.append(np.array(output.per_residue_embedding.mean(dim=0).cpu()))
|
|
51
|
+
df['esm3_mean'] = means
|
|
52
|
+
return df
|
|
53
|
+
|
|
54
|
+
def execute(self, df: pd.DataFrame) -> pd.DataFrame:
|
|
55
|
+
if self.tmp_dir is None:
|
|
56
|
+
with TemporaryDirectory() as tmp_dir:
|
|
57
|
+
if self.num_threads > 1:
|
|
58
|
+
dfs = []
|
|
59
|
+
df_list = np.array_split(df, self.num_threads)
|
|
60
|
+
for df_chunk in tqdm(df_list):
|
|
61
|
+
dfs.append(self.__execute(df_chunk, tmp_dir))
|
|
62
|
+
df = pd.DataFrame()
|
|
63
|
+
for tmp_df in tqdm(dfs):
|
|
64
|
+
df = pd.concat([df, tmp_df])
|
|
65
|
+
return df
|
|
66
|
+
else:
|
|
67
|
+
df = self.__execute(df, tmp_dir)
|
|
68
|
+
return df
|
|
69
|
+
else:
|
|
70
|
+
df = self.__execute(df, self.tmp_dir)
|
|
71
|
+
return df
|