refactor to new recommended jaxtyping/beartype syntax

docs/source/tutorials.rst

@@ -45,7 +45,7 @@ Evaluating a pre-trained Model
     from graphein.protein.tensor.data import ProteinBatch
     from proteinworkshop.models.utils import get_aggregation
     from jaxtyping import jaxtyped
-    from beartype import beartype
+    from beartype import beartype as typechecker
 
 
     class IdentityModel(nn.Module):
@@ -58,8 +58,7 @@ Evaluating a pre-trained Model
             """This property describes the required attributes of the input batch."""
             return {"x", "batch"}
 
-        @jaxtyped
-        @beartype
+        @jaxtyped(typechecker=typechecker)
         def forward(self, batch: Union[Batch, ProteinBatch]) -> Dict[str, torch.Tensor]:
             """
             This method does the forward pass of the model.

notebooks/adding_new_task_tutorial.ipynb

@@ -67,7 +67,7 @@
     "    def required_attributes(self) -> Set[str]:\n",
     "        return {\"residue_type\"}\n",
     "\n",
-    "    @beartype\n",
+    "    @jaxtyped(typechecker=typechecker)\n",
     "    def __call__(self, x: Union[Data, Protein]) -> Union[Data, Protein]:\n",
     "        x.residue_type_uncorrupted = copy.deepcopy(x.residue_type)\n",
     "        # Get indices of residues to corrupt\n",

proteinworkshop/datasets/atom3d_datamodule.py

@@ -11,7 +11,7 @@
 import atom3d.datasets.datasets as da
 import lightning as L
 import torch
-from beartype import beartype
+from beartype import beartype as typechecker
 from graphein.protein.tensor.dataloader import ProteinDataLoader
 from loguru import logger as log
 from torch.utils.data import Dataset
@@ -34,7 +34,7 @@ def set_worker_sharing_strategy(worker_id: int):
     torch.multiprocessing.set_sharing_strategy(SHARING_STRATEGY)
 
 
-@beartype
+@typechecker
 def get_data_path(
     dataset: str,
     lba_split: int = 30,
@@ -57,7 +57,7 @@ def get_data_path(
     return data_paths[dataset]
 
 
-@beartype
+@typechecker
 def get_test_data_path(
     dataset: str,
     lba_split: int = 30,
@@ -85,7 +85,7 @@ def get_test_data_path(
     return data_paths[dataset]
 
 
-@beartype
+@typechecker
 def get_task_split(
     task: str,
     lba_split: int = 30,
@@ -317,7 +317,7 @@ def setup(self, stage: Optional[str] = None):
                 self.data_test,
             ) = self.get_datasets()
 
-    @beartype
+    @typechecker
     def get_dataloader(
         self,
         dataset: Union[da.LMDBDataset, PPIDataset, RESDataset],

proteinworkshop/datasets/components/atom3d_dataset.py

@@ -28,7 +28,7 @@
 NUM_ATOM_TYPES = len(_atom_types_dict)
 
 
-@beartype
+@typechecker
 def _element_mapping(x: str) -> int:
     return _atom_types_dict.get(x, 8)
 

proteinworkshop/datasets/components/ppi_dataset.py

@@ -7,7 +7,7 @@
 import scipy
 import torch
 from atom3d.datasets import LMDBDataset
-from beartype import beartype
+from beartype import beartype as typechecker
 from torch.utils.data import IterableDataset
 from torch_geometric.data import Data
 
@@ -23,14 +23,14 @@
 ]
 
 
-@beartype
+@typechecker
 def get_res(df: pd.DataFrame) -> pd.DataFrame:
     """Get all residues."""
     # Adapted from: https://github.com/drorlab/atom3d/blob/master/examples/ppi/dataset/neighbors.py
     return df[PPI_DF_INDEX_COLUMNS].drop_duplicates()
 
 
-@beartype
+@typechecker
 def _get_idx_to_res_mapping(
     df: pd.DataFrame,
 ) -> Tuple[pd.DataFrame, pd.Series]:
@@ -43,7 +43,7 @@ def _get_idx_to_res_mapping(
     return idx_to_res, res_to_idx
 
 
-@beartype
+@typechecker
 def get_subunits(
     ensemble: pd.DataFrame,
 ) -> Tuple[
@@ -79,7 +79,7 @@ def get_subunits(
     return names, (bdf0, bdf1, udf0, udf1)
 
 
-@beartype
+@typechecker
 def get_negatives(
     neighbors, df0: pd.DataFrame, df1: pd.DataFrame
 ) -> pd.DataFrame:
@@ -170,7 +170,7 @@ def __iter__(self):
             )
         return gen
 
-    @beartype
+    @typechecker
     def _df_to_graph(
         self, struct_df: pd.DataFrame, chain_res: Iterable, label: float
     ) -> Optional[Data]:
@@ -208,7 +208,7 @@ def _df_to_graph(
 
         return data
 
-    @beartype
+    @typechecker
     def _dataset_generator(
         self, indices: List[int], shuffle: bool = True
     ) -> Generator[Tuple[Data, Data], None, None]:
@@ -247,7 +247,7 @@ def _dataset_generator(
                             continue
                         yield graph1, graph2
 
-    @beartype
+    @typechecker
     def _create_labels(
         self,
         positives: pd.DataFrame,

proteinworkshop/features/edge_features.py

@@ -3,7 +3,7 @@
 
 import numpy as np
 import torch
-from beartype import beartype
+from beartype import beartype as typechecker
 from graphein.protein.tensor.types import CoordTensor, EdgeTensor
 from jaxtyping import jaxtyped
 from omegaconf import ListConfig
@@ -20,8 +20,7 @@
 """List of edge features that can be computed."""
 
 
-@jaxtyped
-@beartype
+@jaxtyped(typechecker=typechecker)
 def compute_scalar_edge_features(
     x: Union[Data, Batch], features: Union[List[str], ListConfig]
 ) -> torch.Tensor:
@@ -55,8 +54,7 @@ def compute_scalar_edge_features(
     return torch.cat(feats, dim=1)
 
 
-@jaxtyped
-@beartype
+@jaxtyped(typechecker=typechecker)
 def compute_vector_edge_features(
     x: Union[Data, Batch], features: Union[List[str], ListConfig]
 ) -> Union[Data, Batch]:
@@ -71,8 +69,7 @@ def compute_vector_edge_features(
     return x
 
 
-@jaxtyped
-@beartype
+@jaxtyped(typechecker=typechecker)
 def compute_edge_distance(
     pos: CoordTensor, edge_index: EdgeTensor
 ) -> torch.Tensor:
@@ -91,8 +88,7 @@ def compute_edge_distance(
     )
 
 
-@jaxtyped
-@beartype
+@jaxtyped(typechecker=typechecker)
 def pos_emb(edge_index: EdgeTensor, num_pos_emb: int = 16):
     # From https://github.com/jingraham/neurips19-graph-protein-design
     d = edge_index[0] - edge_index[1]

proteinworkshop/features/edges.py

@@ -4,13 +4,13 @@
 
 import graphein.protein.tensor.edges as gp
 import torch
-from beartype import beartype
+from beartype import beartype as typechecker
 from graphein.protein.tensor.data import Protein, ProteinBatch
 from omegaconf import ListConfig
 from torch_geometric.data import Batch, Data
 
 
-@beartype
+@typechecker
 def compute_edges(
     x: Union[Data, Batch, Protein, ProteinBatch],
     edge_types: Union[ListConfig, List[str]],
@@ -81,7 +81,7 @@ def compute_edges(
     return edges, indxs
 
 
-@beartype
+@typechecker
 def sequence_edges(
     b: Union[Data, Batch, Protein, ProteinBatch],
     chains: Optional[torch.Tensor] = None,

proteinworkshop/features/factory.py

@@ -2,7 +2,7 @@
 
 import torch
 import torch.nn as nn
-from beartype import beartype
+from beartype import beartype as typechecker
 from graphein.protein.tensor.data import ProteinBatch, get_random_batch
 from jaxtyping import jaxtyped
 from loguru import logger
@@ -75,8 +75,7 @@ def __init__(
         if "sequence_positional_encoding" in self.scalar_node_features:
             self.positional_encoding = PositionalEncoding(16)
 
-    @jaxtyped
-    @beartype
+    @jaxtyped(typechecker=typechecker)
     def forward(
         self, batch: Union[Batch, ProteinBatch]
     ) -> Union[Batch, ProteinBatch]:

proteinworkshop/features/node_features.py

@@ -3,7 +3,7 @@
 
 import torch
 import torch.nn.functional as F
-from beartype import beartype
+from beartype import beartype as typechecker
 from graphein.protein.tensor.angles import (
     alpha,
     dihedrals,
@@ -25,8 +25,7 @@
 from .utils import _normalize
 
 
-@jaxtyped
-@beartype
+@jaxtyped(typechecker=typechecker)
 def compute_scalar_node_features(
     x: Union[Batch, Data, Protein, ProteinBatch],
     node_features: Union[ListConfig, List[ScalarNodeFeature]],
@@ -86,8 +85,7 @@ def compute_scalar_node_features(
     return torch.cat(feats, dim=1) if feats else x.x
 
 
-@jaxtyped
-@beartype
+@jaxtyped(typechecker=typechecker)
 def compute_vector_node_features(
     x: Union[Batch, Data, Protein, ProteinBatch],
     vector_features: Union[ListConfig, List[str]],
@@ -114,8 +112,7 @@ def compute_vector_node_features(
     return x
 
 
-@jaxtyped
-@beartype
+@jaxtyped(typechecker=typechecker)
 def compute_surface_feat(
     coords: Union[CoordTensor, AtomTensor], k: int, sigma: List[float]
 ):
@@ -150,8 +147,7 @@ def compute_surface_feat(
     return torch.cat(feat, dim=1)
 
 
-@jaxtyped
-@beartype
+@jaxtyped(typechecker=typechecker)
 def orientations(
     X: Union[CoordTensor, AtomTensor], ca_idx: int = 1
 ) -> OrientationTensor:

proteinworkshop/features/representation.py

@@ -2,7 +2,7 @@
 from typing import Literal, Tuple
 
 import torch
-from beartype import beartype
+from beartype import beartype as typechecker
 from graphein.protein.tensor.types import AtomTensor, CoordTensor
 from jaxtyping import jaxtyped
 from torch_geometric.data import Batch, Data
@@ -11,8 +11,7 @@
 from proteinworkshop.configs.config import ExperimentConfigurationError
 
 
-@jaxtyped
-@beartype
+@jaxtyped(typechecker=typechecker)
 def get_full_atom_coords(
     atom_tensor: AtomTensor, fill_value: float = 1e-5
 ) -> Tuple[CoordTensor, torch.Tensor, torch.Tensor]:
@@ -40,8 +39,7 @@ def get_full_atom_coords(
     return coords, residue_index, atom_type
 
 
-@jaxtyped
-@beartype
+@jaxtyped(typechecker=typechecker)
 def transform_representation(
     x: Batch, representation_type: Literal["CA", "BB", "FA", "BB_SC", "CA_SC"]
 ) -> Batch:
@@ -87,7 +85,7 @@ def transform_representation(
         )
 
 
-@beartype
+@typechecker
 def _ca_to_fa_repr(x: Data) -> Data:
     """Converts CA representation to full atom representation."""
     coords, residue_index, atom_type = get_full_atom_coords(x.coords)
@@ -101,7 +99,7 @@ def _ca_to_fa_repr(x: Data) -> Data:
     return x
 
 
-@beartype
+@typechecker
 def _ca_to_bb_repr(x: Data) -> Data:
     """Converts CA representation to backbone representation."""
     x.pos = x.coords[:, :4, :].reshape(-1, 3)
@@ -117,7 +115,7 @@ def _ca_to_bb_repr(x: Data) -> Data:
     return x
 
 
-@beartype
+@typechecker
 def ca_to_bb_repr(batch: Batch) -> Batch:  # sourcery skip: assign-if-exp
     """
     Converts a batch of CA representations to backbone representations. I.e.
@@ -180,7 +178,7 @@ def ca_to_bb_repr(batch: Batch) -> Batch:  # sourcery skip: assign-if-exp
     return batch
 
 
-@beartype
+@typechecker
 def ca_to_bb_sc_repr(batch: Batch) -> Batch:
     """Converts a batch of CA representations to backbone + sidechain representations."""
     # Get centroids
@@ -190,7 +188,7 @@ def ca_to_bb_sc_repr(batch: Batch) -> Batch:
     return ca_to_fa_repr(batch)
 
 
-@beartype
+@typechecker
 def ca_to_ca_sc_repr(batch: Batch) -> Batch:
     """Converts a batch of CA representations to C + sidechain representations."""
     # Get centroids
@@ -200,7 +198,7 @@ def ca_to_ca_sc_repr(batch: Batch) -> Batch:
     return batch
 
 
-@beartype
+@typechecker
 def coarsen_sidechain(x: Data, aggr: str = "mean") -> CoordTensor:
     """Returns tensor of sidechain centroids: L x 3"""
     # sourcery skip: remove-unnecessary-else, swap-if-else-branches
@@ -216,7 +214,7 @@ def coarsen_sidechain(x: Data, aggr: str = "mean") -> CoordTensor:
     return sc_points
 
 
-@beartype
+@typechecker
 def ca_to_fa_repr(batch: Batch) -> Batch:  # sourcery skip: assign-if-exp
     """Converts a batch of CA representations to full atom representations."""
     if "sidechain_torsion" in batch.keys: