From 7eeba58db783c87a0907450192b7eaea3e4b0690 Mon Sep 17 00:00:00 2001
From: Eachan Johnson <eachan.johnson@crick.ac.uk>
Date: Sun, 26 Apr 2026 22:04:58 +0100
Subject: [PATCH 1/2] Fix typo

---
 aspect/data.py | 17 ++++++++---------
 pyproject.toml |  2 +-
 2 files changed, 9 insertions(+), 10 deletions(-)

diff --git a/aspect/data.py b/aspect/data.py
index b42cbc1..a3758eb 100644
--- a/aspect/data.py
+++ b/aspect/data.py
@@ -401,15 +401,14 @@ def _featurize(
     def _unsqueeze(
         x: Mapping[str, ArrayLike]
     ) -> 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, 
diff --git a/pyproject.toml b/pyproject.toml
index e08f10b..707d39f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "aspect-data"
-version = "0.0.3"
+version = "0.0.3.post1"
 authors = [
   { name="Eachan Johnson", email="eachan.johnson@crick.ac.uk" },
 ]

From 3c612f8ca299264568dc7b72effbec854e4c4acc Mon Sep 17 00:00:00 2001
From: Eachan Johnson <eachan.johnson@crick.ac.uk>
Date: Sun, 26 Apr 2026 22:08:12 +0100
Subject: [PATCH 2/2] Fix unsqueeze behaviour

---
 aspect/data.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/aspect/data.py b/aspect/data.py
index a3758eb..5379740 100644
--- a/aspect/data.py
+++ b/aspect/data.py
@@ -399,7 +399,8 @@ def _featurize(
 
     @staticmethod
     def _unsqueeze(
-        x: Mapping[str, ArrayLike]
+        x: Mapping[str, ArrayLike],
+        columns: Optional[Iterable[str]] = None
     ) -> Dict[str, np.ndarray]:
         columns = columns or x.keys()
         for key in columns:
@@ -476,6 +477,7 @@ def __call__(
             .select_columns(all_output_columns)
             .map(
                 self._unsqueeze,
+                fn_kwargs={"columns": output_columns},
                 batched=True,
                 batch_size=batch_size,
                 desc="Unsqueezing",