index modification as reviewer requested

dstl · May 6, 2020 · e175f0e · e175f0e
1 parent 9087bce
commit e175f0e
Showing 1 changed file with 9 additions and 8 deletions.
diff --git a/stonesoup/models/measurement/nonlinear.py b/stonesoup/models/measurement/nonlinear.py
@@ -680,7 +680,7 @@ def function(self, state, noise=False, **kwargs) -> StateVector:
                 noise = 0
 
         # Account for origin offset in position to enable range and angles to be determined
-        xy_pos = state.state_vector[self.mapping] - self.translation_offset
+        xy_pos = state.state_vector[self.mapping, :] - self.translation_offset
 
         # Rotate coordinates based upon the sensor_velocity
         xy_rot = self._rotation_matrix @ xy_pos
@@ -825,7 +825,7 @@ def function(self, state, noise=False, **kwargs) -> StateVector:
                 noise = 0
 
         # Account for origin offset in position to enable range and angles to be determined
-        xyz_pos = state.state_vector[self.mapping] - self.translation_offset
+        xyz_pos = state.state_vector[self.mapping, :] - self.translation_offset
 
         # Rotate coordinates based upon the sensor_velocity
         xyz_rot = self._rotation_matrix @ xyz_pos
@@ -857,14 +857,15 @@ def inverse_function(self, detection, **kwargs) -> StateVector:
         inv_rotation_matrix = inv(self._rotation_matrix)
 
         out_vector = StateVector([[0.], [0.], [0.], [0.], [0.], [0.]])
-        out_vector[self.mapping] = x, y, z
-        out_vector[self.velocity_mapping] = x_rate, y_rate, z_rate
+        out_vector[self.mapping, :] = x, y, z
+        out_vector[self.velocity_mapping, :] = x_rate, y_rate, z_rate
 
-        out_vector[self.mapping] = inv_rotation_matrix @ out_vector[self.mapping]
-        out_vector[self.velocity_mapping] = inv_rotation_matrix @ out_vector[self.velocity_mapping]
+        out_vector[self.mapping, :] = inv_rotation_matrix @ out_vector[self.mapping, :]
+        out_vector[self.velocity_mapping, :] = \
+            inv_rotation_matrix @ out_vector[self.velocity_mapping, :]
+
+        out_vector[self.mapping, :] = out_vector[self.mapping, :] + self.translation_offset
 
-        out_vector[self.mapping] = out_vector[self.mapping] + self.translation_offset
-        print(out_vector)
         return out_vector
 
     def rvs(self, num_samples=1, **kwargs) -> np.ndarray: