Merge pull request #18 from prescient-design/deps

dependencies

Author: kleinhenz

.github/workflows/push.yml

@@ -23,19 +23,17 @@ jobs:
           - "macos-latest"
           - "ubuntu-latest"
           # - "windows-latest"
-        python:
+        python-version:
           - "3.10"
     runs-on: ${{ matrix.platform }}
     steps:
       - uses: "actions/checkout@v4"
-      - uses: "actions/setup-python@v5"
+      - name: Install uv and set the python version
+        uses: astral-sh/setup-uv@v5
         with:
-          python-version: ${{ matrix.python }}
-      - run: "python -m pip install -r requirements.in"
-      - run: "python -m pip install -r requirements-dev.in"
-      - run: "python -m pip install -r requirements-mgm.in"
-      - run: "python -m pip install --editable ."
-      - run: "python -m pytest"
+          python-version: ${{ matrix.python-version }}
+      - run: uv sync --all-extras --dev
+      - run: uv run pytest
       - env:
           CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}
         uses: "codecov/codecov-action@v3"

requirements.in

@@ -16,11 +16,8 @@ datasets
 tokenizers
 transformers>=4.24.0
 scikit-learn
-matplotlib
-seaborn
 python-dotenv
 torcheval
-fastparquet
 datasketch
 peft
 icecream
@@ -29,5 +26,5 @@ pooch
 edlib
 onnx
 onnxscript
-beignet[all]
+beignet[datasets, mdtraj]
 fair-esm

src/lobster/cmdline/_intervene.py

@@ -260,13 +260,13 @@ def intervene(cfg: DictConfig) -> bool:
                     all_data = np.concatenate((all_data, data), axis=1) if all_data is not None else np.array(data)
 
                 if all_data.shape[1] > 1000:
-                    print(f"{concept} {i+1} of {len(predict_dataloader)} Saving")
+                    print(f"{concept} {i + 1} of {len(predict_dataloader)} Saving")
                     df = pd.DataFrame(all_data.transpose(), columns=columns)
                     df.to_csv(f"{data_folder_dir}/{index}.csv", index=False)
                     index += 1
                     all_data = None
 
             if all_data is not None:
-                print(f"{concept} {i+1} of {len(predict_dataloader)} Saving")
+                print(f"{concept} {i + 1} of {len(predict_dataloader)} Saving")
                 df = pd.DataFrame(all_data.transpose(), columns=columns)
                 df.to_csv(f"{data_folder_dir}/{index}.csv", index=False)

src/lobster/cmdline/_intervene_multiproperty.py

@@ -386,13 +386,13 @@ def intervene_multiproperty(cfg: DictConfig) -> bool:
                 all_data = np.concatenate((all_data, data), axis=1) if all_data is not None else np.array(data)
 
             if all_data.shape[1] > 1000:
-                print(f"{i+1} of {len(predict_dataloader)} Saving")
+                print(f"{i + 1} of {len(predict_dataloader)} Saving")
                 df = pd.DataFrame(all_data.transpose(), columns=columns)
-                df.to_csv(f"{data_folder_dir}/{index__+1}_indexed.csv", index=False)
+                df.to_csv(f"{data_folder_dir}/{index__ + 1}_indexed.csv", index=False)
                 index__ += 1
                 all_data = None
 
         if all_data is not None:
-            print(f"{i+1} of {len(predict_dataloader)} Saving")
+            print(f"{i + 1} of {len(predict_dataloader)} Saving")
             df = pd.DataFrame(all_data.transpose(), columns=columns)
-            df.to_csv(f"{data_folder_dir}/{index__+1}_indexed.csv", index=False)
+            df.to_csv(f"{data_folder_dir}/{index__ + 1}_indexed.csv", index=False)

src/lobster/data/__init__.py

@@ -7,12 +7,12 @@
     DataFrameDatasetInMemory,
     DataFrameLightningDataModule,
 )
-from ._farthest_first_traversal import FarthestFirstTraversal, ranked_fft
-from ._fasta_datamodule import FastaLightningDataModule  # nopycln: import
 from ._dyab_data import (  # nopycln: import
     DyAbDataFrameDatasetInMemory,
     DyAbDataFrameLightningDataModule,
 )
+from ._farthest_first_traversal import FarthestFirstTraversal, ranked_fft
+from ._fasta_datamodule import FastaLightningDataModule  # nopycln: import
 from ._minhasher import LobsterMinHasher
 from ._mmseqs import MMSeqsRunner
 from ._structure_datamodule import PDBDataModule

src/lobster/transforms/_convert_seqs.py

@@ -1,7 +1,18 @@
+from importlib.util import find_spec
 from typing import Callable, Dict, Optional
 
-import selfies as sf
-from rdkit import Chem
+_SELFIES_AVAILABLE = False
+_RDKIT_AVAILABLE = False
+
+if find_spec("selfies"):
+    import selfies as sf
+
+    _SELFIES_AVAILABLE = True
+
+if find_spec("rdkit"):
+    from rdkit import Chem
+
+    _RDKIT_AVAILABLE = True
 
 
 def convert_nt_to_aa(
@@ -51,6 +62,8 @@ def convert_aa_to_nt(
 
 
 def convert_aa_to_smiles(aa_seq: str, allowed_aa: set) -> Optional[str]:
+    assert _RDKIT_AVAILABLE, "rdkit not available. This dependency is part of the mgm extra"
+
     if not aa_seq.isupper():
         aa_seq = aa_seq.upper()
 
@@ -69,6 +82,8 @@ def convert_aa_to_smiles(aa_seq: str, allowed_aa: set) -> Optional[str]:
 def convert_smiles_to_aa(
     smiles_seq: str,
 ) -> Optional[str]:
+    assert _RDKIT_AVAILABLE, "rdkit not available. This dependency is part of the mgm extra"
+
     try:
         mol = Chem.MolFromSmiles(smiles_seq)
         aa_seq = Chem.MolToSequence(mol)
@@ -81,6 +96,7 @@ def convert_smiles_to_aa(
 def convert_smiles_to_selfies(
     smiles_seq: str,
 ) -> Optional[str]:
+    assert _SELFIES_AVAILABLE, "selfies not available. This dependency is part of the mgm extra"
     try:
         sf_seq = sf.encoder(smiles_seq)
         return sf_seq
@@ -92,6 +108,7 @@ def convert_selfies_to_smiles(
     selfies_seq: str,
 ) -> Optional[str]:
     # TODO: add conversion of unknown selfies tokens to Ala selfies
+    assert _SELFIES_AVAILABLE, "selfies not available. This dependency is part of the mgm extra"
     try:
         smiles_seq = sf.decoder(selfies_seq)
         return smiles_seq

tests/lobster/transforms/test__convert_seqs.py

@@ -115,7 +115,7 @@ def test_convert_aa_to_smiles(self):
         # print(smi_seq)
         assert len(smi_seq) == 100, f"Failed for AA seq {aa_seq}, smiles seq does not have the expected length"
         assert smi_seq == (
-            "CC(C)C[C@H](NC(=O)[C@H](CCC(N)=O)NC(=O)[C@@H](NC(=O)[C@@H](N)CCC(=O)O)" "C(C)C)C(=O)N[C@H](C(=O)O)C(C)C"
+            "CC(C)C[C@H](NC(=O)[C@H](CCC(N)=O)NC(=O)[C@@H](NC(=O)[C@@H](N)CCC(=O)O)C(C)C)C(=O)N[C@H](C(=O)O)C(C)C"
         )
 
     def test_convert_smiles_to_aa(self):