AI-based D-amino acid substitution for optimizing antimicrobial peptides to treat multidrug-resistant bacterial infection
This repository contains the code for the paper "AI-based D-amino acid substitution for optimizing antimicrobial peptides to treat multidrug-resistant bacterial infection"
Requirements
mamba_ssm==2.2.4
numpy==1.26.3
pandas==2.1.4
rdkit==2024.3.5
scikit_learn==1.4.1.post1
scipy==1.13.0
torch==2.2.0
torchmetrics==1.3.1
torchvision==0.17.0
You can install them with pip install -r requirements.txt
Additionally, mamba_ssm is optional since it is not used for our final method.
You can comment mamba_ssm==2.2.4 in requirements.txt and from mamba_ssm import Mamba in network.py out if you don't want to install it and avoid use --q-encoder mamba.
Training
There are two .py file for training: main.py and main_simple.py.
main.py: Can train model with Classification and Regression tasks. Prefered with regression task.
main_simple.py: Can ONLY train model with Classification task. Prefered with classification task. simple means a simple dataset that direct loads pre-processed data.
example:
python main-simple.py \
--q-encoder cnn \ # Encoder, can be cnn, lstm, gru, mamba, mha
--channels 16 \ # Encoder channels
--side-enc lstm \ # Side sequence Encoder, only lstm implemented, only use with cnn encoder
--fusion att \ # Fusion method, can be att, mlp or diff
--task cls \ # Task, can be cls or reg
--loss ce \ # Loss, can be ce or mse, some other losses can be found in code
--batch-size 32 \ # Batch size
--epochs 35 \ # Epochs
--gpu 0 \ # GPU index to use, -1 for cpu
# ===CNN only options=== \
--pcs \ # Enable protease cleavage site dyeing for input pictures
--resize 768 \ # Resize input pictures, can be 1 or 2 numbers like 768 or 768 512
# ===main_simple.py only options=== \
--llm-data # Use LLM augmented training data
Corresponding model weight checkpoints will be saved in the subdirectory of run-cls or run-reg, e.g. /run-cls/cnn-att-16-lstm-pcs-simple-llm-768-oneway-ce-32-0.001-35/
For more arguments, please refer to the code of main.py or main_simple.py
Inference
You can simple replace main.py with infer.py in your training command to do inference. Remember to add --simple if you used checkpoints trained from main_simple.py
For case study scanning, please use infer_case.py with an additional argument --case r2 or --case YOUR_PEPTIDE_SEQUENCE
Inference results will be saved in the weights directory in csv format, e.g. /run-cls/cnn-att-16-lstm-pcs-simple-llm-768-oneway-ce-32-0.001-35/preds_test.csv