Merge pull request #187 from dstl/statevector_indexing

Change StateVector indexing

stonesoup/feeder/geo.py

@@ -109,7 +109,7 @@ def detections_gen(self):
             utm_detections = set()
             for detection in detections:
                 easting, northing, zone_num, northern = utm.from_latlon(
-                    *detection.state_vector[self.mapping[::-1], 0],
+                    *detection.state_vector[self.mapping[::-1], :],
                     self.zone_number)
                 if self.zone_number is None:
                     self.zone_number = zone_num

stonesoup/feeder/tests/test_geo.py

@@ -35,7 +35,7 @@ def test_lla_reference_converter(detector, converter_class, reverse_func):
         detection = detections.pop()
 
         assert pytest.approx((50, i, 5000 + i*10), abs=1e-2, rel=1e-3) == \
-            reverse_func(*detection.state_vector[:, 0], 50, 0, 5000)
+            reverse_func(*detection.state_vector, 50, 0, 5000)
 
 
 def test_utm_converter(detector):
@@ -54,4 +54,4 @@ def test_utm_converter(detector):
         p_east = detection.state_vector[0]
 
         assert pytest.approx((50, long), rel=1e-2, abs=1e-4) == utm.to_latlon(
-            *detection.state_vector[0:2, 0], zone_number=30, northern=True)
+            *detection.state_vector[0:2], zone_number=30, northern=True)

stonesoup/functions.py

@@ -130,6 +130,8 @@ def gauss2sigma(state, alpha=1.0, beta=2.0, kappa=None):
 
     # Calculate sigma point locations
     sigma_points = np.tile(state.state_vector, (1, 2 * ndim_state + 1))
+    # as sigma_points is a 2d it should no longer be a StateVector
+    sigma_points = Matrix(sigma_points)
     # Can't use in place addition/subtraction as casting issues may arise when mixing float/int
     sigma_points[:, 1:(ndim_state + 1)] = \
         sigma_points[:, 1:(ndim_state + 1)] + sqrt_sigma*np.sqrt(c)

stonesoup/initiator/tests/test_simple.py

@@ -114,7 +114,7 @@ def test_linear_measurement():
 
 def test_nonlinear_measurement():
     measurement_model = CartesianToBearingRange(
-        2, (0, 1), np.diag([np.radians(2), 30]))
+        2, [0, 1], np.diag([np.radians(2), 30]))
     measurement_initiator = SimpleMeasurementInitiator(
         GaussianState(np.array([[0], [0]]), np.diag([100, 10])),
         measurement_model

stonesoup/models/measurement/nonlinear.py

@@ -244,23 +244,20 @@ def function(self, state, noise=False, **kwargs):
         xyz_rot = self._rotation_matrix @ xyz
 
         # Convert to Spherical
-        rho, phi, theta = cart2sphere(*xyz_rot[:, 0])
+        rho, phi, theta = cart2sphere(*xyz_rot)
 
         return StateVector([[Elevation(theta)], [Bearing(phi)], [rho]]) + noise
 
     def inverse_function(self, detection, **kwargs):
 
-        theta, phi, rho = detection.state_vector[:, 0]
-        x, y, z = sphere2cart(rho, phi, theta)
+        theta, phi, rho = detection.state_vector
+        xyz = StateVector(sphere2cart(rho, phi, theta))
 
-        xyz = [[x], [y], [z]]
         inv_rotation_matrix = inv(self._rotation_matrix)
-        xyz_rot = inv_rotation_matrix @ xyz
-        xyz = [xyz_rot[0][0], xyz_rot[1][0], xyz_rot[2][0]]
-        x, y, z = xyz + self.translation_offset[:, 0]
+        xyz = inv_rotation_matrix @ xyz
 
         res = np.zeros((self.ndim_state, 1)).view(StateVector)
-        res[self.mapping, 0] = x, y, z
+        res[self.mapping, :] = xyz + self.translation_offset
 
         return res
 
@@ -343,23 +340,22 @@ def ndim_meas(self):
         return 2
 
     def inverse_function(self, detection, **kwargs):
-        if not ((self.rotation_offset[0][0] == 0)
-                and (self.rotation_offset[1][0] == 0)):
+        if not ((self.rotation_offset[0] == 0)
+                and (self.rotation_offset[1] == 0)):
             raise RuntimeError(
                 "Measurement model assumes 2D space. \
                 Rotation in 3D space is unsupported at this time.")
 
-        phi, rho = detection.state_vector[:, 0]
-        x, y = pol2cart(rho, phi)
+        phi, rho = detection.state_vector[:]
+        xy = StateVector(pol2cart(rho, phi))
 
-        xyz = [[x], [y], [0]]
+        xyz = np.concatenate((xy, StateVector([0])), axis=0)
         inv_rotation_matrix = inv(self._rotation_matrix)
-        xyz_rot = inv_rotation_matrix @ xyz
-        xy = [xyz_rot[0][0], xyz_rot[1][0]]
-        x, y = xy + self.translation_offset[:, 0]
+        xyz = inv_rotation_matrix @ xyz
+        xy = xyz[0:2]
 
         res = np.zeros((self.ndim_state, 1)).view(StateVector)
-        res[self.mapping, 0] = x, y
+        res[self.mapping, :] = xy + self.translation_offset
 
         return res
 
@@ -507,13 +503,13 @@ def function(self, state, noise=False, **kwargs):
                 noise = 0
 
         # Account for origin offset
-        xyz = state.state_vector[self.mapping] - self.translation_offset
+        xyz = state.state_vector[self.mapping, :] - self.translation_offset
 
         # Rotate coordinates
         xyz_rot = self._rotation_matrix @ xyz
 
         # Convert to Angles
-        phi, theta = cart2angles(*xyz_rot[:, 0])
+        phi, theta = cart2angles(*xyz_rot)
 
         return StateVector([[Elevation(theta)], [Bearing(phi)]]) + noise
 

stonesoup/sensor/radar/radar.py

@@ -455,13 +455,13 @@ def gen_probability(self, sky_state):
         [r, pos_az, pos_el] = cart2sphere(*relative_vector)
 
         # target position relative to beam position
-        relative_az = pos_az[0] - beam_az
-        relative_el = pos_el[0] - beam_el
+        relative_az = pos_az - beam_az
+        relative_el = pos_el - beam_el
         # calculate power directed towards target
         self.beam_shape.beam_width = spoiled_width  # beam spoiling to width
         directed_power = self.beam_shape.beam_power(relative_az, relative_el)
         # calculate signal to noise ratio
-        snr = self._snr_constant * rcs * spoiled_gain ** 2 * directed_power / (r[0] ** 4)
+        snr = self._snr_constant * rcs * spoiled_gain ** 2 * directed_power / (r ** 4)
         # calculate probability of detection using the North's approximation
         det_prob = 0.5 * erfc(
             (-np.log(self.probability_false_alarm)) ** 0.5 - (

stonesoup/sensor/tests/test_passive.py

@@ -48,7 +48,7 @@ def test_passive_sensor():
     xyz_rot = rot_mat @ xyz
 
     # Convert to Angles
-    phi, theta = cart2angles(*xyz_rot[:, 0])
+    phi, theta = cart2angles(*xyz_rot)
 
     # Assert correction of generated measurement
     assert (measurement.timestamp == target_state.timestamp)

stonesoup/simulator/tests/test_platform.py

@@ -85,7 +85,7 @@ def test_platform_ground_truth_detection_simulator(sensor_model1,
         for detection in detections:
             for i in range(0, len(detection.state_vector)):
                 # Detection at location of ground truth.
-                assert int(detection.state_vector[i][0]) == int(n/2)
+                assert int(detection.state_vector[i]) == int(n/2)
 
 
 def test_detection_simulator(sensor_model1,
@@ -108,4 +108,4 @@ def test_detection_simulator(sensor_model1,
         assert len(detections) == 2  # Detection count at each step.
         for detection in detections:
             # Detection at location of ground truth or at platform.
-            assert int(detection.state_vector[0][0]) in (int(n/3), 2*int(n/3))
+            assert int(detection.state_vector[0]) in (int(n/3), 2*int(n/3))

stonesoup/types/array.py

@@ -50,6 +50,17 @@ class StateVector(Matrix):
     ``StateVector([1., 2., 3.])``, ``StateVector ([[1., 2., 3.,]])``, and
     ``StateVector([[1.], [2.], [3.]])`` will all return the same 3x1 StateVector.
 
+    It also overrides the behaviour of indexing such that my_state_vector[1] returns the second
+    element (as `int`, `float` etc), rather than a StateVector of size (1, 1) as would be the case
+    without this override. Behaviour of indexing with lists, slices or other indexing is
+    unaffected (as you would expect those to return StateVectors). This override avoids the need
+    for client to specifically index with zero as the second element (`my_state_vector[1, 0]`) to
+    get a native numeric type. Iterating through the StateVector returns a sequence of numbers,
+    rather than a sequence of 1x1 StateVectors. This makes the class behave as would be expected
+    and avoids 'gotchas'.
+
+    Note that code using the pattern `my_state_vector[1, 0]` will continue to work.
+
     .. note ::
         It is not recommended to use a StateVector for indexing another vector. Doing so will lead
         to unexpected effects. Use a :class:`tuple`, :class:`list` or :class:`np.ndarray` for this.
@@ -70,6 +81,23 @@ def __new__(cls, *args, **kwargs):
                     array.shape))
         return array.view(cls)
 
+    def __getitem__(self, item):
+        # If item has two elements, it is a tuple and should be left alone.
+        # If item is a slice object, or an ndarray, we would expect a StateVector returned,
+        #   so leave it alone.
+        # If item is an int, we would expected a number returned, so we should append 0  to the
+        #   item and extract the first (and only) column
+        # Note that an ndarray of ints is an instance of int
+        #   i.e. isinstance(np.array([1]), int) == True
+        if isinstance(item, int):
+            item = (item, 0)
+        return super().__getitem__(item)
+
+    def __setitem__(self, key, value):
+        if isinstance(key, int):
+            key = (key, 0)
+        return super().__setitem__(key, value)
+
 
 class CovarianceMatrix(Matrix):
     """Covariance matrix wrapper for :class:`numpy.ndarray`.

stonesoup/types/tests/test_array.py

@@ -18,6 +18,49 @@ def test_statevector():
     assert np.array_equal(state_vector, state_vector_array)
     assert np.array_equal(StateVector([1, 2, 3, 4]), state_vector_array)
     assert np.array_equal(StateVector([[1, 2, 3, 4]]), state_vector_array)
+    assert np.array_equal(StateVector(state_vector_array), state_vector)
+
+
+def test_standard_statevector_indexing():
+    state_vector_array = np.array([[1], [2], [3], [4]])
+    state_vector = StateVector(state_vector_array)
+
+    # test standard indexing
+    assert state_vector[2, 0] == 3
+    assert not isinstance(state_vector[2, 0], StateVector)
+
+    # test Slicing
+    assert state_vector[1:2, 0] == 2
+    assert isinstance(state_vector[1:2, 0], StateVector)
+    assert np.array_equal(state_vector[:], state_vector)
+    assert isinstance(state_vector[:, 0], StateVector)
+    assert np.array_equal(state_vector[0:], state_vector)
+    assert isinstance(state_vector[0:, 0], StateVector)
+
+    # test list indices
+    assert np.array_equal(state_vector[[1, 3]], StateVector([2, 4]))
+    assert isinstance(state_vector[[1, 3], 0], StateVector)
+
+    # test int indexing
+    assert state_vector[2] == 3
+    assert not isinstance(state_vector[2], StateVector)
+
+
+def test_setting():
+    state_vector_array = np.array([[1], [2], [3], [4]])
+    state_vector = StateVector(state_vector_array.copy())
+
+    state_vector[2, 0] = 4
+    assert np.array_equal(state_vector, StateVector([1, 2, 4, 4]))
+
+    state_vector[2] = 5
+    assert np.array_equal(state_vector, StateVector([1, 2, 5, 4]))
+
+    state_vector[:] = state_vector_array[:]
+    assert np.array_equal(state_vector, StateVector([1, 2, 3, 4]))
+
+    state_vector[1:3] = StateVector([5, 6])
+    assert np.array_equal(state_vector, StateVector([1, 5, 6, 4]))
 
 
 def test_covariancematrix():