work dump, add onnx stuff

README.md

@@ -1,7 +1,7 @@
 
-# Lead Optimization via Fragment Prediction
+# Lead Optimization
 
-# Overview
+# Structure
 
 - `config`: configuration information (eg. TRAIN/TEST partitions)
 - `data`: training/inference data (see [`data/README.md`](data/README.md))

launch.py

@@ -0,0 +1,78 @@
+
+import argparse
+import sys
+import os
+import subprocess
+import tempfile
+
+
+RUN_DIR = '/zfs1/jdurrant/durrantlab/hag63/leadopt_pytorch'
+
+
+def main():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument('-t','--time',default='10:00:00')
+    parser.add_argument('-p','--partition',default='gtx1080')
+    parser.add_argument('-m','--mem',default='16g')
+    parser.add_argument('path')
+    parser.add_argument('script')
+
+    args = parser.parse_args()
+
+    run_path = os.path.join(RUN_DIR, args.path)
+    if os.path.exists(run_path):
+        print('[!] Run exists at %s' % run_path)
+        overwrite = input('- Overwrite? [Y/n]: ')
+        if overwrite.lower() == 'n':
+            print('Exiting...')
+            exit(0)
+    else:
+        print('[*] Creating run directory %s' % run_path)
+        os.mkdir(run_path)
+
+    script = '''#!/bin/bash
+#SBATCH --job-name={name}
+#SBATCH --output={run_path}/slurm_out.txt
+#SBATCH --error={run_path}/slurm_err.txt
+#SBATCH --time={time}
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cluster=gpu
+#SBATCH --partition={partition}
+#SBATCH [email protected]
+#SBATCH --mail-type=END,FAIL
+#SBATCH --gres=gpu:1
+#SBATCH --mem={mem}
+
+cd /ihome/jdurrant/hag63/cbio/leadopt_pytorch/
+./setup.sh
+
+export PYTHONPATH=/ihome/jdurrant/hag63/cbio/leadopt_pytorch/
+# export WANDB_DIR=/ihome/jdurrant/hag63/wandb_abs
+export WANDB_DISABLE_CODE=true
+
+cd {run_path}
+python /ihome/jdurrant/hag63/cbio/leadopt_pytorch/{script}
+    '''.format(
+        name='leadopt_%s' % args.path,
+        run_path=run_path,
+        time=args.time,
+        partition=args.partition,
+        mem=args.mem,
+        script=args.script
+    )
+
+    print('[*] Running script...')
+
+    with tempfile.NamedTemporaryFile('w') as f:
+        f.write(script)
+        f.flush()
+
+        r = subprocess.run('sbatch %s' % f.name, shell=True, stdin=subprocess.PIPE, stdout=subprocess.PIPE)
+        print(r)
+
+
+
+if __name__=='__main__':
+    main()

launch.sh

@@ -0,0 +1,47 @@
+#!/bin/bash
+
+DEFAULT_RUN_PATH=/zfs1/jdurrant/durrantlab/hag63/leadopt_pytorch/
+
+if [[ $# -ne 4 ]]; then
+    echo "Usage: $0 <run_name> <gpu_partition> <**args>"
+    exit 0
+fi
+
+ABS_SCRIPT=$(pwd)/train.py
+
+# navigate to runs directory
+RUNS_DIR="${RUNS_DIR:-$DEFAULT_RUN_PATH}"
+cd $RUNS_DIR
+
+if [[ -d $1 ]]; then
+    echo "Warning: run directory $1 already exists!"
+    exit -1
+fi
+
+echo "Creating run directory ($1)..."
+mkdir $1
+
+echo "Running script..."
+sbatch <<EOT
+#!/bin/bash
+#SBATCH --job-name=$1
+#SBATCH --output=$RUNS_DIR/$1/slurm_out.txt
+#SBATCH --error=$RUNS_DIR/$1/slurm_err.txt
+#SBATCH --time=10:00:00
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cluster=gpu
+#SBATCH --partition=$2
+#SBATCH [email protected]
+#SBATCH --mail-type=END,FAIL
+#SBATCH --gres=gpu:1
+#SBATCH --mem=80g
+
+cd /ihome/jdurrant/hag63/cbio/leadopt_pytorch/
+source ./setup.sh
+
+PYTHON_PATH=$PYTHON_PATH:/ihome/jdurrant/hag63/cbio/leadopt_pytorch/
+
+cd $RUNS_DIR/$1
+python train.py $4
+EOT

leadopt.py

@@ -13,52 +13,9 @@
 
 from leadopt.models.voxel import VoxelFingerprintNet2
 from leadopt.infer import infer_all
+from leadopt.pretrained import MODELS
 
 
-class SavedModel(object):
-
-    model_class = None
-    model_args = None
-
-    @classmethod
-    def load(cls, path):
-        m = cls.model_class(**cls.model_args).cuda()
-        m.load_state_dict(torch.load(path))
-        m.eval()
-        return m
-
-    @classmethod
-    def get_fingerprints(cls, path):
-        f = h5py.File(os.path.join(path, cls.fingerprint_data), 'r')
-
-        data = f['fingerprints'][()]
-        smiles = f['smiles'][()]
-
-        f.close()
-
-        return data, smiles
-
-class V2_RDK_M150(SavedModel):
-    model_class = VoxelFingerprintNet2
-    model_args = {
-        'in_channels': 18,
-        'output_size': 2048,
-        'batchnorm': True,
-        'sigmoid': True
-    }
-
-    grid_width=24
-    grid_res=1
-    receptor_types=[6,7,8,9,15,16,17,35,53]
-    parent_types=[6,7,8,9,15,16,17,35,53]
-
-    fingerprint_data = 'fingerprint_rdk_2048.h5'
-
-
-MODELS = {
-    'rdk_m150': V2_RDK_M150
-}
-
 def main():
     parser = argparse.ArgumentParser()