feat: change shape property of PhaseSpacePosition

src/galax/dynamics/_base.py

@@ -51,13 +51,20 @@ def _shape_tuple(self) -> tuple[tuple[int, ...], tuple[int, ...]]:
     @property
     def shape(self) -> tuple[int, ...]:
         """Shape of the position and velocity arrays."""
-        batch_shape, component_shapes = self._shape_tuple
-        return (*batch_shape, sum(component_shapes))
+        return self._shape_tuple[0]
 
     def __len__(self) -> int:
         """Return the number of particles."""
         return self.shape[0]
 
+    # ==========================================================================
+
+    @property
+    def full_shape(self) -> tuple[int, ...]:
+        """Shape of the position and velocity arrays."""
+        batch_shape, component_shapes = self._shape_tuple
+        return (*batch_shape, sum(component_shapes))
+
     # ==========================================================================
     # Convenience properties
 

tests/unit/potential/test_base.py

@@ -142,7 +142,7 @@ def test_integrate_orbit(self, pot: AbstractPotentialBase, xv: Vec6) -> None:
 
         orbit = pot.integrate_orbit(xv, ts)
         assert isinstance(orbit, gd.Orbit)
-        assert orbit.shape == (len(ts), 7)
+        assert orbit.shape == (len(ts),)
         assert array_equal(orbit.t, ts)
 
     def test_integrate_orbit_batch(self, pot: AbstractPotentialBase, xv: Vec6) -> None:
@@ -152,13 +152,13 @@ def test_integrate_orbit_batch(self, pot: AbstractPotentialBase, xv: Vec6) -> No
         # Simple batch
         orbits = pot.integrate_orbit(xv[None, :], ts)
         assert isinstance(orbits, gd.Orbit)
-        assert orbits.shape == (1, len(ts), 7)
+        assert orbits.shape == (1, len(ts))
         assert array_equal(orbits.t, ts[None, :])
 
         # More complicated batch
         xv2 = xp.stack([xv, xv], axis=0)
         orbits = pot.integrate_orbit(xv2, ts)
         assert isinstance(orbits, gd.Orbit)
-        assert orbits.shape == (2, len(ts), 7)
+        assert orbits.shape == (2, len(ts))
         assert array_equal(orbits.t[0], ts)
         assert array_equal(orbits.t[1], ts)