patch: MomentumArray works for single entries #83

jacanchaplais · Feb 6, 2023 · 97458da · 97458da
1 parent 70d9f33
commit 97458da
Showing 1 changed file with 2 additions and 2 deletions.
diff --git 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