From 97458da63e62f36e396c1eb8edc4ba5f068eeb70 Mon Sep 17 00:00:00 2001
From: Jacan Chaplais <jacan.chaplais@gmail.com>
Date: Mon, 6 Feb 2023 18:08:22 +0000
Subject: [PATCH] patch: MomentumArray works for single entries #83

---
 graphicle/data.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/graphicle/data.py b/graphicle/data.py
index 1a7ba63..8bfdd98 100644
--- a/graphicle/data.py
+++ b/graphicle/data.py
@@ -1054,14 +1054,14 @@ def copy(self) -> "MomentumArray":
 
     @property
     def _xy_pol(self) -> base.ComplexVector:
-        return self._data[:, :2].view(dtype="<c16").reshape(-1)
+        return self._data[..., :2].view(dtype="<c16").squeeze(-1)
 
     @fn.cached_property
     def _zt_pol(self) -> base.ComplexVector:
         return (
             np.stack((self.data["z"], np.abs(self._xy_pol)), axis=-1)
             .view(dtype="<c16")
-            .reshape(-1)
+            .squeeze(-1)
         )
 
     @fn.cached_property