Set up code as a package and switch to relative imports

code/compute_rosetta_standardization.py

@@ -7,7 +7,7 @@
 import logging
 
 import pandas as pd
-import split_dataset as sd
+from . import split_dataset as sd
 
 logger = logging.getLogger("METL." + __name__)
 logger.setLevel(logging.DEBUG)

code/convert_ckpt.py

@@ -13,11 +13,9 @@
 import torch.nn as nn
 import torchinfo
 
-import constants
-import models
-import tasks
-import utils
-import encode as enc
+from . import models
+from . import utils
+from . import encode as enc
 
 
 def convert_checkpoint(ckpt_dict):

code/datamodules.py

@@ -15,13 +15,13 @@
 from torch.utils.data import DataLoader
 import pytorch_lightning as pl
 
-import datasets
-import pdb_sampler
-import utils
-import constants
-import split_dataset as sd
-import encode as enc
-from datasets import RosettaDatasetSQL
+from . import datasets
+from . import pdb_sampler
+from . import utils
+from . import constants
+from . import split_dataset as sd
+from . import encode as enc
+from .datasets import RosettaDatasetSQL
 
 
 class DMSDataModule(pl.LightningDataModule):

code/datasets.py

@@ -10,9 +10,9 @@
 import torch.utils.data
 from torch import Tensor
 
-import constants
-import split_dataset as sd
-import encode as enc
+from . import constants
+from . import split_dataset as sd
+from . import encode as enc
 
 
 def load_standardization_params(split_dir, train_only=True):

code/encode.py

@@ -7,9 +7,9 @@
 import numpy as np
 import pandas as pd
 
-import rosetta_data_utils as rd
-import constants
-import utils
+from . import rosetta_data_utils as rd
+from . import constants
+from . import utils
 
 
 def is_seq_level_encoding(encoding: str):

code/models.py

@@ -12,8 +12,8 @@
 import torch.nn.functional as F
 from torch import Tensor
 
-import relative_attention as ra
-import tasks
+from . import relative_attention as ra
+from . import tasks
 
 
 def reset_parameters_helper(m: nn.Module):

code/parse_raw_dms_data.py

@@ -6,7 +6,7 @@
 import numpy as np
 import pandas as pd
 
-import utils
+from . import utils
 
 
 def sort_and_save_to_csv(df, out_fn, precision=7, sort_muts=True, sort_variants=True, na_rep=""):

code/parse_rosetta_data.py

@@ -15,10 +15,9 @@
 import sqlite3
 from tqdm import tqdm
 
-import constants
-
-import utils
-import rosetta_data_utils as rd
+from . import constants
+from . import utils
+from . import rosetta_data_utils as rd
 
 logger = logging.getLogger("METL." + __name__)
 logger.setLevel(logging.DEBUG)

code/relative_attention.py

@@ -16,10 +16,10 @@
 from torch import Tensor
 from torch.nn import Linear, Dropout, LayerNorm
 import time
-
-import structure
 import networkx as nx
-import models
+
+from . import structure
+from . import models
 
 
 class RelativePosition3D(nn.Module):

code/rosetta_data_utils.py

@@ -3,7 +3,7 @@
 import pandas as pd
 import sqlalchemy as sqla
 
-import utils
+from . import utils
 
 
 def convert_dms_to_rosettafy_indexing(ds_name, variants, reverse=False):

code/split_dataset.py

@@ -13,7 +13,7 @@
 import pandas as pd
 from sklearn.model_selection import train_test_split
 
-import utils
+from . import utils
 
 
 logger = logging.getLogger("METL." + __name__)

code/structure.py

@@ -7,7 +7,7 @@
 import networkx as nx
 from biopandas.pdb import PandasPdb
 
-import utils
+from . import utils
 
 
 class GraphType(Enum):

code/tasks.py

@@ -10,9 +10,9 @@
 import pytorch_lightning as pl
 import torchmetrics
 
-import training_utils
-from training_utils import CosineWarmupScheduler, ConstantWarmupScheduler
-import models
+from . import training_utils
+from .training_utils import CosineWarmupScheduler, ConstantWarmupScheduler
+from . import models
 
 
 class RosettaTask(pl.LightningModule):

code/tests.py

@@ -1,10 +1,9 @@
 """ testing code """
 import metl
 import torch
-import utils
-
 from argparse import ArgumentParser
 
+from . import utils
 
 def load_checkpoint_run_inference(checkpoint_path, variants, dataset):
     """ loads a finetuned 3D model from a checkpoint and scores variants with the model """

code/train_source_model.py

@@ -16,11 +16,11 @@
 
 import wandb
 
-import utils
-from training_utils import BestMetricLogger, save_metrics_ptl, CondorStopping, create_log_dir, get_next_version
-from datamodules import RosettaDataModule
-import models
-import tasks
+from . import utils
+from .training_utils import BestMetricLogger, save_metrics_ptl, CondorStopping, create_log_dir, get_next_version
+from .datamodules import RosettaDataModule
+from . import models
+from . import tasks
 
 
 class ModelCheckpoint(pytorch_lightning.callbacks.ModelCheckpoint):

code/train_target_model.py

@@ -15,14 +15,14 @@
     LearningRateMonitor, Checkpoint, StochasticWeightAveraging, ModelSummary, RichProgressBar, RichModelSummary
 import numpy as np
 
-import models
-import training_utils
-import utils
-from datamodules import DMSDataModule
-import finetuning_callbacks
-from finetuning_callbacks import AnyFinetuning
-from tasks import DMSTask
-import analysis_utils as an
+from . import models
+from . import training_utils
+from . import utils
+from .datamodules import DMSDataModule
+from . import finetuning_callbacks
+from .finetuning_callbacks import AnyFinetuning
+from .tasks import DMSTask
+from . import analysis_utils as an
 
 logging.basicConfig(level=logging.INFO)
 

code/training_utils.py

@@ -23,9 +23,9 @@
 from torch import optim, Tensor
 from torch.optim.lr_scheduler import LambdaLR
 
-import utils
-import datamodules
-from metrics import compute_metrics
+from . import utils
+from . import datamodules
+from .metrics import compute_metrics
 
 
 def save_scatterplots(dm, predictions_d, log_dir, suffix=""):

code/utils.py

@@ -11,7 +11,7 @@
 from Bio import PDB
 from Bio.PDB.PDBParser import PDBParser
 
-import constants
+from . import constants
 
 
 def mkdir(d):
@@ -81,13 +81,14 @@ def load_dataset_metadata(metadata_fn: str = "data/dms_data/datasets.yml"):
 def load_dataset(ds_name: Optional[str] = None,
                  ds_fn: Optional[str] = None,
                  sort_mutations: bool = False,
-                 load_epistasis: bool = False):
+                 load_epistasis: bool = False,
+                 metadata_fn: str = "data/dms_data/datasets.yml"):
     """ load a dataset as pandas dataframe """
     if ds_name is None and ds_fn is None:
         raise ValueError("must provide either ds_name or ds_fn to load a dataset")
 
     if ds_fn is None:
-        datasets = load_dataset_metadata()
+        datasets = load_dataset_metadata(metadata_fn)
         ds_fn = datasets[ds_name]["ds_fn"]
 
     if not isfile(ds_fn):