aspect-data 0.0.3__tar.gz → 0.0.4__tar.gz

{aspect_data-0.0.3 → aspect_data-0.0.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aspect-data
-Version: 0.0.3
+Version: 0.0.4
 Summary: General purpose, chemistry, and taxonomic featurization of tabular data with caching.
 Author-email: Eachan Johnson <eachan.johnson@crick.ac.uk>
 License: MIT License
@@ -48,14 +48,17 @@ Requires-Dist: platformdirs>=4.4.0
 Requires-Dist: tqdm
 Provides-Extra: chem
 Requires-Dist: schemist>=0.0.6; extra == "chem"
+Provides-Extra: torch
+Requires-Dist: torch>=2.4; extra == "torch"
 Provides-Extra: chemprop
+Requires-Dist: aspect-data[chem,torch]; extra == "chemprop"
 Requires-Dist: chemprop>=2.1.0; extra == "chemprop"
 Provides-Extra: deep
 Requires-Dist: transformers; extra == "deep"
 Provides-Extra: bio
-Requires-Dist: vectome>=0.0.4; extra == "bio"
+Requires-Dist: vectome>=0.0.8; extra == "bio"
 Provides-Extra: dev
-Requires-Dist: aspect-data[bio,chem]; extra == "dev"
+Requires-Dist: aspect-data[bio,chemprop]; extra == "dev"
 Requires-Dist: pytest>=7.4; extra == "dev"
 Requires-Dist: flake8; extra == "dev"
 Dynamic: license-file
@@ -66,6 +69,8 @@ Dynamic: license-file
 ![PyPI - Python Version](https://img.shields.io/pypi/pyversions/aspect-data)
 ![PyPI](https://img.shields.io/pypi/v/aspect-data)
 
+<img src="docs/source/_static/logo.svg" width="200">
+
 **aspect** is a lightweight featurisation layer for tabular machine-learning data.
 
 It lets you define independent feature pipelines for different input columns, then apply them to dictionaries, Pandas data frames, local files, Hugging Face datasets, or saved dataset checkpoints.
@@ -83,6 +88,15 @@ input table
 output dataset with named feature columns
 ```
 
+Aspect datasets store Arrow-compatible feature representations.
+
+Before model execution, an optional collator converts stored examples into
+runtime representations. A runtime batch is always a mapping from feature
+column names to batched values. Values may be tensors or structured mappings.
+
+Downstream model libraries should treat runtime batch values as opaque and
+should not perform dataset preprocessing or deserialization.
+
 ---
 
 ## Contents

{aspect_data-0.0.3 → aspect_data-0.0.4}/README.md

@@ -4,6 +4,8 @@
 ![PyPI - Python Version](https://img.shields.io/pypi/pyversions/aspect-data)
 ![PyPI](https://img.shields.io/pypi/v/aspect-data)
 
+<img src="docs/source/_static/logo.svg" width="200">
+
 **aspect** is a lightweight featurisation layer for tabular machine-learning data.
 
 It lets you define independent feature pipelines for different input columns, then apply them to dictionaries, Pandas data frames, local files, Hugging Face datasets, or saved dataset checkpoints.
@@ -21,6 +23,15 @@ input table
 output dataset with named feature columns
 ```
 
+Aspect datasets store Arrow-compatible feature representations.
+
+Before model execution, an optional collator converts stored examples into
+runtime representations. A runtime batch is always a mapping from feature
+column names to batched values. Values may be tensors or structured mappings.
+
+Downstream model libraries should treat runtime batch values as opaque and
+should not perform dataset preprocessing or deserialization.
+
 ---
 
 ## Contents

aspect_data-0.0.4/aspect/collate/__init__.py

@@ -0,0 +1,7 @@
+from .base import ColumnCollator
+from .chemprop import chemprop_collate
+
+__all__ = [
+    "ColumnCollator",
+    "chemprop_collate",
+]

aspect_data-0.0.4/aspect/collate/base.py

@@ -0,0 +1,42 @@
+"""Specialised column collation."""
+
+from collections.abc import Callable, Mapping, Iterable
+from typing import Any, TypeAlias
+
+from ..typing import Batch
+
+
+CollateFn: TypeAlias = Callable[[Iterable[Any]], Any]
+
+
+class ColumnCollator:
+    """Override collation for selected dataset columns."""
+
+    def __init__(
+        self,
+        collators: Mapping[str, CollateFn] | None = None,
+    ):
+        self.collators = collators or {}
+
+    def __call__(
+        self,
+        rows: Batch,
+    ) -> dict[str, ...]:
+        from torch.utils.data import default_collate
+
+        out = {}
+        if not rows:
+            return out
+
+        for column in rows[0]:
+            values = [
+                row[column]
+                for row in rows
+            ]
+
+            if column in self.collators:
+                out[column] = self.collators[column](values)
+            else:
+                out[column] = default_collate(values)
+
+        return out

aspect_data-0.0.4/aspect/collate/chemprop.py

@@ -0,0 +1,115 @@
+"""Chemprop collation helpers."""
+
+from collections.abc import Mapping, Iterable
+
+from chemprop.data import MolGraph
+from chemprop.data.collate import BatchMolGraph
+import numpy as np
+import torch
+from torch import Tensor
+
+from ..typing import Datum
+
+
+def _as_tensor(x) -> Tensor:
+    """Convert an value to a tensor."""
+    if x is None:
+        return None
+    if torch.is_tensor(x):
+        return x
+    return torch.as_tensor(x)
+
+
+def _as_numpy(x) -> np.ndarray:
+    """Convert an value to a tensor."""
+    if x is None:
+        return None
+    if isinstance(x, np.ndarray):
+        return x
+    return np.asarray(x)
+
+
+def _molgraph_from_mapping(
+    x: Mapping[str, ...]
+) -> MolGraph:
+    """Reconstruct a Chemprop MolGraph from a cached mapping."""
+    return MolGraph(
+        V=_as_numpy(x["V"]),
+        E=_as_numpy(x["E"]),
+        edge_index=_as_numpy(x["edge_index"]).astype(np.int32),
+        rev_edge_index=_as_numpy(x["rev_edge_index"]).astype(np.int32),
+    )
+
+
+def _collate_optional(
+    values: Iterable[...],
+    *,
+    stack: bool,
+):
+    """Collate an optional Chemprop feature."""
+    if all(value is None for value in values):
+        return None
+
+    if any(value is None for value in values):
+        raise ValueError(
+            "Cannot collate Chemprop feature containing a mixture "
+            "of None and non-None values."
+        )
+
+    tensors = [
+        _as_tensor(value)
+        for value in values
+    ]
+
+    if stack:
+        return torch.stack(tensors, dim=0)
+
+    return torch.cat(tensors, dim=0)
+
+
+def chemprop_collate(
+    values: Iterable[Datum],
+) -> dict[str, ...]:
+    """Collate cached Chemprop features into a runtime model input.
+
+    Each input value is expected to have the form::
+
+        {
+            "bmg": {
+                "V": ...,
+                "E": ...,
+                "edge_index": ...,
+                "rev_edge_index": ...,
+            },
+            "V_d": ...,
+            "X_d": ...,
+        }
+
+    Returns
+    =======
+    dict
+        Runtime representation suitable for ``ChempropEncoder(**batch)``.
+    
+    """
+    if len(values) == 0:
+        raise ValueError("Cannot collate an empty Chemprop batch.")
+
+    molgraphs = [
+        _molgraph_from_mapping(value["bmg"])
+        for value in values
+    ]
+    bmg = BatchMolGraph(molgraphs)
+
+    V_d = _collate_optional(
+        [value.get("V_d") for value in values],
+        stack=False,
+    )
+    X_d = _collate_optional(
+        [value.get("X_d") for value in values],
+        stack=True,
+    )
+    return {
+        "bmg": bmg,
+        "V_d": V_d,
+        "X_d": X_d,
+    }

{aspect_data-0.0.3 → aspect_data-0.0.4}/aspect/data.py

@@ -399,17 +399,17 @@ class DataPipeline:
 
     @staticmethod
     def _unsqueeze(
-        x: Mapping[str, ArrayLike]
+        x: Mapping[str, ArrayLike],
+        columns: Optional[Iterable[str]] = None
     ) -> Dict[str, np.ndarray]:
-        def _unsqueeze(x, columns=None):
-            columns = columns or x.keys()
-            for key in columns:
-                vals = x[key]
-                if not isinstance(vals, dict):
-                    vals = np.asarray(x[key])
-                    if vals.ndim == 1 and np.issubdtype(vals.dtype, np.number):
-                        x[key] = vals[:, None]
-            return x
+        columns = columns or x.keys()
+        for key in columns:
+            vals = x[key]
+            if not isinstance(vals, dict):
+                vals = np.asarray(x[key])
+                if vals.ndim == 1 and np.issubdtype(vals.dtype, np.number):
+                    x[key] = vals[:, None]
+        return x
 
     def __call__(
         self, 
@@ -477,6 +477,7 @@ class DataPipeline:
             .select_columns(all_output_columns)
             .map(
                 self._unsqueeze,
+                fn_kwargs={"columns": output_columns},
                 batched=True,
                 batch_size=batch_size,
                 desc="Unsqueezing",

{aspect_data-0.0.3 → aspect_data-0.0.4}/aspect/typing.py

@@ -1,6 +1,7 @@
 """Custom types."""
 
-from typing import TYPE_CHECKING, Any, Iterable, Mapping, Union
+from typing import TYPE_CHECKING, Any, TypeAlias, Union
+from collections.abc import Iterable, Mapping
 
 from datasets import Dataset, IterableDataset
 from pandas import DataFrame
@@ -11,14 +12,14 @@ if TYPE_CHECKING:
 else:
     ColumnTransform = Any
 
-DataLike = Union[
+DataLike: TypeAlias = Union[
     str, 
     DataFrame, 
     Mapping[str, ArrayLike], 
     Dataset, 
     IterableDataset
 ]
-FeatureLike = Union[
+FeatureLike: TypeAlias = Union[
     str, 
     Mapping[str, ColumnTransform], 
     Mapping[str, Mapping[str, Any]], 
@@ -29,3 +30,6 @@ FeatureLike = Union[
     ]]
 ]
 StrOrIterableOfStr = Union[str, Iterable[str]]
+
+Datum: TypeAlias = Mapping[str, Any]
+Batch: TypeAlias = Iterable[Mapping[str, Any]]

{aspect_data-0.0.3 → aspect_data-0.0.4}/aspect_data.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aspect-data
-Version: 0.0.3
+Version: 0.0.4
 Summary: General purpose, chemistry, and taxonomic featurization of tabular data with caching.
 Author-email: Eachan Johnson <eachan.johnson@crick.ac.uk>
 License: MIT License
@@ -48,14 +48,17 @@ Requires-Dist: platformdirs>=4.4.0
 Requires-Dist: tqdm
 Provides-Extra: chem
 Requires-Dist: schemist>=0.0.6; extra == "chem"
+Provides-Extra: torch
+Requires-Dist: torch>=2.4; extra == "torch"
 Provides-Extra: chemprop
+Requires-Dist: aspect-data[chem,torch]; extra == "chemprop"
 Requires-Dist: chemprop>=2.1.0; extra == "chemprop"
 Provides-Extra: deep
 Requires-Dist: transformers; extra == "deep"
 Provides-Extra: bio
-Requires-Dist: vectome>=0.0.4; extra == "bio"
+Requires-Dist: vectome>=0.0.8; extra == "bio"
 Provides-Extra: dev
-Requires-Dist: aspect-data[bio,chem]; extra == "dev"
+Requires-Dist: aspect-data[bio,chemprop]; extra == "dev"
 Requires-Dist: pytest>=7.4; extra == "dev"
 Requires-Dist: flake8; extra == "dev"
 Dynamic: license-file
@@ -66,6 +69,8 @@ Dynamic: license-file
 ![PyPI - Python Version](https://img.shields.io/pypi/pyversions/aspect-data)
 ![PyPI](https://img.shields.io/pypi/v/aspect-data)
 
+<img src="docs/source/_static/logo.svg" width="200">
+
 **aspect** is a lightweight featurisation layer for tabular machine-learning data.
 
 It lets you define independent feature pipelines for different input columns, then apply them to dictionaries, Pandas data frames, local files, Hugging Face datasets, or saved dataset checkpoints.
@@ -83,6 +88,15 @@ input table
 output dataset with named feature columns
 ```
 
+Aspect datasets store Arrow-compatible feature representations.
+
+Before model execution, an optional collator converts stored examples into
+runtime representations. A runtime batch is always a mapping from feature
+column names to batched values. Values may be tensors or structured mappings.
+
+Downstream model libraries should treat runtime batch values as opaque and
+should not perform dataset preprocessing or deserialization.
+
 ---
 
 ## Contents

{aspect_data-0.0.3 → aspect_data-0.0.4}/aspect_data.egg-info/SOURCES.txt

@@ -11,6 +11,9 @@ aspect/cli_module/__init__.py
 aspect/cli_module/cli.py
 aspect/cli_module/featurize.py
 aspect/cli_module/io.py
+aspect/collate/__init__.py
+aspect/collate/base.py
+aspect/collate/chemprop.py
 aspect/transform/__init__.py
 aspect/transform/base.py
 aspect/transform/deep_functions.py
@@ -21,4 +24,6 @@ aspect_data.egg-info/SOURCES.txt
 aspect_data.egg-info/dependency_links.txt
 aspect_data.egg-info/entry_points.txt
 aspect_data.egg-info/requires.txt
-aspect_data.egg-info/top_level.txt
+aspect_data.egg-info/top_level.txt
+test/test_collate.py
+test/test_collate_chemprop.py

{aspect_data-0.0.3 → aspect_data-0.0.4}/aspect_data.egg-info/requires.txt

@@ -7,18 +7,22 @@ platformdirs>=4.4.0
 tqdm
 
 [bio]
-vectome>=0.0.4
+vectome>=0.0.8
 
 [chem]
 schemist>=0.0.6
 
 [chemprop]
+aspect-data[chem,torch]
 chemprop>=2.1.0
 
 [deep]
 transformers
 
 [dev]
-aspect-data[bio,chem]
+aspect-data[bio,chemprop]
 pytest>=7.4
 flake8
+
+[torch]
+torch>=2.4

{aspect_data-0.0.3 → aspect_data-0.0.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "aspect-data"
-version = "0.0.3"
+version = "0.0.4"
 authors = [
   { name="Eachan Johnson", email="eachan.johnson@crick.ac.uk" },
 ]
@@ -44,10 +44,11 @@ dependencies = [
 
 [project.optional-dependencies]
 chem = ["schemist>=0.0.6"]
-chemprop = ["chemprop>=2.1.0"]
+torch = ["torch>=2.4"]
+chemprop = ["aspect-data[torch,chem]", "chemprop>=2.1.0"]
 deep = ["transformers"]
-bio = ["vectome>=0.0.4"]
-dev = ["aspect-data[chem,bio]", "pytest>=7.4", "flake8"]
+bio = ["vectome>=0.0.8"]
+dev = ["aspect-data[chemprop,bio]", "pytest>=7.4", "flake8"]
 
 [project.urls]
 "Homepage" = "https://github.com/scbirlab/aspect"

aspect_data-0.0.4/test/test_collate.py

@@ -0,0 +1,78 @@
+
+from aspect.collate import ColumnCollator
+
+def test_column_collator_default():
+    import torch
+
+    collator = ColumnCollator()
+
+    batch = collator([
+        {
+            "x": torch.tensor([1., 2.]),
+            "y": torch.tensor(1.),
+        },
+        {
+            "x": torch.tensor([3., 4.]),
+            "y": torch.tensor(2.),
+        },
+    ])
+
+    assert batch["x"].shape == (2, 2)
+    assert batch["y"].shape == (2,)
+
+
+def test_column_collator_override():
+    collator = ColumnCollator(
+        collators={
+            "special": lambda values: {
+                "values": values,
+            }
+        }
+    )
+
+    batch = collator([
+        {
+            "x": 1,
+            "special": "a",
+        },
+        {
+            "x": 2,
+            "special": "b",
+        },
+    ])
+
+    assert batch["special"] == {
+        "values": ["a", "b"],
+    }
+
+
+def test_column_collator_mixed():
+    import torch
+
+    collator = ColumnCollator(
+        collators={
+            "graph": lambda values: {
+                "graph_batch": values,
+            }
+        }
+    )
+
+    batch = collator([
+        {
+            "fp": torch.tensor([1., 2.]),
+            "graph": {"id": 1},
+            "target": torch.tensor(0.),
+        },
+        {
+            "fp": torch.tensor([3., 4.]),
+            "graph": {"id": 2},
+            "target": torch.tensor(1.),
+        },
+    ])
+
+    assert batch["fp"].shape == (2, 2)
+    assert batch["target"].shape == (2,)
+    assert batch["graph"]["graph_batch"] == [
+        {"id": 1},
+        {"id": 2},
+    ]

aspect_data-0.0.4/test/test_collate_chemprop.py

@@ -0,0 +1,104 @@
+"""Tests for Chemprop collation."""
+
+from chemprop.data.collate import BatchMolGraph
+from chemprop.featurizers import SimpleMoleculeMolGraphFeaturizer
+import pytest
+from rdkit import Chem
+import torch
+
+from aspect.collate.chemprop import chemprop_collate
+
+
+def _cached_molgraph(smiles: str) -> dict:
+    """Create an HF-like serialized MolGraph."""
+    featurizer = SimpleMoleculeMolGraphFeaturizer()
+    molgraph = featurizer(
+        Chem.MolFromSmiles(smiles)
+    )
+
+    return {
+        "V": molgraph.V.tolist(),
+        "E": molgraph.E.tolist(),
+        "edge_index": molgraph.edge_index.tolist(),
+        "rev_edge_index": molgraph.rev_edge_index.tolist(),
+    }
+
+
+def test_chemprop_collate_graphs():
+    smiles = ["CCO", "CCN", "CCC"]
+    values = [
+        {
+            "bmg": _cached_molgraph(s),
+            "V_d": None,
+            "X_d": None,
+        }
+        for s in smiles
+    ]
+    batch = chemprop_collate(values)
+
+    assert set(batch) == {
+        "bmg",
+        "V_d",
+        "X_d",
+    }
+    assert isinstance(
+        batch["bmg"],
+        BatchMolGraph,
+    )
+
+    assert batch["V_d"] is None
+    assert batch["X_d"] is None
+
+    assert batch["bmg"].V.ndim == 2
+    assert batch["bmg"].edge_index.ndim == 2
+
+    assert torch.unique(
+        batch["bmg"].batch
+    ).numel() == 3
+
+
+def test_chemprop_collate_X_d():
+    values = [
+        {
+            "bmg": _cached_molgraph("CCO"),
+            "V_d": None,
+            "X_d": [1.0, 2.0],
+        },
+        {
+            "bmg": _cached_molgraph("CCN"),
+            "V_d": None,
+            "X_d": [3.0, 4.0],
+        },
+    ]
+
+    batch = chemprop_collate(values)
+
+    assert batch["X_d"].shape == (2, 2)
+    assert torch.equal(
+        batch["X_d"],
+        torch.tensor([
+            [1., 2.],
+            [3., 4.],
+        ]),
+    )
+
+
+def test_chemprop_collate_rejects_mixed_optional_values():
+    values = [
+        {
+            "bmg": _cached_molgraph("CCO"),
+            "V_d": None,
+            "X_d": [1.0, 2.0],
+        },
+        {
+            "bmg": _cached_molgraph("CCN"),
+            "V_d": None,
+            "X_d": None,
+        },
+    ]
+
+    with pytest.raises(
+        ValueError,
+        match="mixture of None and non-None",
+    ):
+        chemprop_collate(values)