docstrings, restructure

proteinshake/adapter.py

@@ -0,0 +1,6 @@
+class Adapter:
+    """
+    Downloads raw pdb files and/or meta data from a source and formats it to the shake database schema.
+    """
+
+    pass

proteinshake/database.py

@@ -1,13 +1,16 @@
 from pathlib import Path
-from .collection import Collection
 
 
 class Database:
+    """
+    Spins up a redis database
+    """
+
     def __init__(self, storage: Path) -> None:
         pass
 
     def update(self) -> None:
         pass
 
-    def query(self, query: str) -> Collection:
+    def query(self, query: str):
         pass

proteinshake/framework.py

@@ -1,3 +1,10 @@
 class Framework:
+    """
+    Abstract class for a framework. Used as Mixin with a Transform.
+    """
+
     def create_loader(self, iterator):
-        pass
+        """
+        Creates a framework-specific dataloader from an iterator.
+        """
+        raise NotImplementedError

proteinshake/metric.py

@@ -1,4 +1,6 @@
 class Metric:
-    """For a collection of predictions and target values, return set of performance metrics.,"""
+    """
+    Computes a set of relevant metrics for a task.
+    """
 
     pass

proteinshake/metrics/__init__.py

@@ -1 +1 @@
-from .evaluator import *
+from .dummy_metric import *

proteinshake/metrics/dummy_metric.py

@@ -0,0 +1,7 @@
+from proteinshake.metric import Metric
+import numpy as np
+
+
+class DummyMetric(Metric):
+    def __call__(self, y_true, y_pred):
+        return {"Accuracy": np.random.random()}

proteinshake/representation.py

@@ -1,2 +1,6 @@
 class Representation:
+    """
+    Abstract class for a representation. Used as Mixin with a Transform.
+    """
+
     pass

proteinshake/split.py

@@ -1,8 +1,17 @@
+from typing import Dict, Iterator
+
+
 class Split:
     """
-    Abstract class for selecting train/val/test indices given a dataset.
+    Abstract class to create data splits from a dataset.
     """
 
+    def __call__(self, dataset: Iterator) -> Dict[str, Iterator]:
+        """
+        Takes an Xy iterator and returns a dictionary of Xy iterators, where each key denotes the split name (usually 'train', 'test', and 'val').
+        """
+        raise NotImplementedError
+
     @property
     def hash(self):
         return self.__class__.__name__

proteinshake/splits/__init__.py

@@ -1 +1 @@
-from .splitter import *
+from .dummy_split import *

proteinshake/splits/dummy_split.py

@@ -0,0 +1,13 @@
+from proteinshake.split import Split
+import itertools
+
+
+class DummySplit(Split):
+    def __call__(self, Xy):
+        train, testval = itertools.tee(Xy)
+        test, val = itertools.tee(testval)
+        return {
+            "train": filter(lambda Xy: Xy[0][0]["split"] == "train", train),
+            "test": filter(lambda Xy: Xy[0][0]["split"] == "test", test),
+            "val": filter(lambda Xy: Xy[0][0]["split"] == "val", val),
+        }

proteinshake/target.py

@@ -1,6 +1,13 @@
+from typing import Dict, Iterator
+
+
 class Target:
-    """Returns the attribute to predict for a single instance, given arbitrary inputs.
-    Different tasks will have target computations on different types and numbers of entitites.
+    """
+    Abstract class for reshaping a dataset into the correct data-target structure for a task.
     """
 
-    pass
+    def __call__(self, dataset: Iterator[dict]) -> Dict[str, Iterator]:
+        """
+        Takes a dataset iterator and returns an Xy iterator, whose elements are ((X1,X2,...), y) pairs of data tuples and targets.
+        """
+        raise NotImplementedError

proteinshake/targets/__init__.py

@@ -1 +1 @@
-from .target import Target
+from .attribute_target import AttributeTarget

proteinshake/targets/attribute_target.py

@@ -0,0 +1,10 @@
+from proteinshake.target import Target
+
+
+class AttributeTarget(Target):
+    def __init__(self, attribute) -> None:
+        super().__init__()
+        self.attribute = attribute
+
+    def __call__(self, dataset):
+        return (((p,), p[self.attribute]) for p in dataset)

proteinshake/task.py

@@ -6,7 +6,7 @@
 from proteinshake.split import Split
 from proteinshake.target import Target
 from proteinshake.metric import Metric
-from proteinshake.transform import Transform, Compose
+from proteinshake.transform import Transform, Compose, IdentityTransform
 from proteinshake.util import amino_acid_alphabet, sharded, save_shards, load, warn
 
 
@@ -15,7 +15,7 @@ class Task:
     split: Split = None
     target: Target = None
     metrics: Metric = None
-    augmentation: Transform = None
+    augmentation: Transform = IdentityTransform
 
     def __init__(
         self,
@@ -46,6 +46,7 @@ def __init__(
     @property
     def proteins(self):
         # return dataset iterator
+        # this is a dummy for now. It will load a dataset from file in the future.
         rng = np.random.default_rng(42)
         return (
             {
@@ -62,7 +63,7 @@ def proteins(self):
 
     def transform(self, *transforms) -> None:
         Xy = self.target(self.proteins)
-        partitions = self.split(Xy)  # returns dict of generators[(X,...),y]
+        partitions = self.split(Xy)
         self.transform = Compose(*[self.augmentation, *transforms])
         # cache from here
         self.transform.fit(partitions["train"])
@@ -73,7 +74,7 @@ def transform(self, *transforms) -> None:
             )
             save_shards(
                 data_transformed,
-                self.root / self.split.hash / self.transform.hash / "shards",
+                self.root / self.split.hash / name / self.transform.hash / "shards",
             )
             setattr(self, f"{name}_loader", partial(self.loader, split=name))
         return self
@@ -87,7 +88,7 @@ def loader(
         **kwargs,
     ):
         rng = np.random.default_rng(random_seed)
-        path = self.root / self.split.hash / self.transform.hash / "shards"
+        path = self.root / self.split.hash / split / self.transform.hash / "shards"
         shard_index = load(path / "index.npy")
         if self.shard_size % batch_size != 0 and batch_size % self.shard_size != 0:
             warn(

proteinshake/tasks/__init__.py

@@ -0,0 +1 @@
+from .dummy_task import DummyTask

proteinshake/tasks/dummy_task.py

@@ -0,0 +1,11 @@
+from proteinshake.task import Task
+from proteinshake.metrics import DummyMetric
+from proteinshake.targets import AttributeTarget
+from proteinshake.splits import DummySplit
+
+
+class DummyTask(Task):
+    dataset = "test"
+    split = DummySplit
+    target = AttributeTarget
+    metrics = DummyMetric

proteinshake/transform.py

@@ -3,6 +3,10 @@
 
 
 class BaseTransform:
+    """
+    Abstract class for transforms. A transform can be stochastic or deterministic, which decides whether the transformed result can be precomputed and saved to disk (deterministic), or if it needs to be computed when retrieving a data item (stochastic). Transforms generally take a batch of Xy tuples, some subclasses exist that facilitate reshaping (see below). Transforms can be fit beforehand (on the 'train' partition).
+    """
+
     stochastic = False
 
     def __call__(self, Xy):
@@ -55,7 +59,16 @@ def inverse_transform(self, y):
         return y
 
 
+class IdentityTransform(BaseTransform):
+    def __call__(self, Xy):
+        return Xy
+
+
 class Compose:
+    """
+    Composes multiple transforms into one object. Takes care of splitting the deterministic and stochastic part, as well as storing the framework create_dataloader method.
+    """
+
     def __init__(self, *transforms):
         self.transforms = transforms
         self.deterministic_transforms = []

tests/task.py

@@ -1,63 +1,26 @@
 import unittest
-import numpy as np
-import itertools
-from proteinshake.metric import Metric
-from proteinshake.target import Target
-from proteinshake.split import Split
-from proteinshake.task import Task
+from proteinshake.tasks import DummyTask
 from proteinshake.transform import *
 from proteinshake.transforms import *
 
 
 class TestTask(unittest.TestCase):
     def test_task(self):
-        # CONTRIBUTOR
-        class MyTarget(Target):
-            def __call__(self, dataset):
-                return (((p,), p["label"]) for p in dataset)
-
-        class MyMetric(Metric):
-            def __call__(self, y_true, y_pred):
-                return {"Accuracy": np.random.random()}
-
-        class MySplit(Split):
-            def __call__(self, Xy):
-                # this implementation looks a bit inefficient
-                train, testval = itertools.tee(Xy)
-                test, val = itertools.tee(testval)
-                return {
-                    "train": filter(lambda Xy: Xy[0][0]["split"] == "train", train),
-                    "test": filter(lambda Xy: Xy[0][0]["split"] == "test", test),
-                    "val": filter(lambda Xy: Xy[0][0]["split"] == "val", val),
-                }
-
-        class MyAugmentation(Transform):
-            def transform(self, X):
-                return X
-
-        class MyTask(Task):
-            dataset = "test"
-            split = MySplit
-            target = MyTarget
-            metrics = MyMetric
-            augmentation = MyAugmentation
-
-        # END USER
         class MyLabelTransform(LabelTransform):
             def transform(self, y):
                 return -y
 
             def inverse_transform(self, y):
                 return -y
 
-        task = MyTask(shard_size=8).transform(
+        task = DummyTask(target_kwargs={"attribute": "label"}, shard_size=8).transform(
             MyLabelTransform(),
             PointRepresentationTransform(),
             TorchFrameworkTransform(),
         )
 
         for epoch in range(5):
-            for X, y in task.train_loader(batch_size=64, random_seed=0):
+            for X, y in task.train_loader(batch_size=16, shuffle=True, random_seed=0):
                 print("X", X.shape)
                 print("y", y.shape)
                 break