py: vectorized likelihood matrix evaluation

src/trackers.py

@@ -196,14 +196,23 @@ class GNN():
         # Entry for new targets
         np.fill_diagonal(association_matrix[:, Nc+ny:], np.log(self.logic_params['Bnt']))
 
-        for ti, track in enumerate(tracks): # Iterate over confirmed tracks
-            validation_matrix[:, ti] = self.gater.gate(track['x'][-1], track['P'][-1], meas)
-            # Entry for validated tracks
-            val_meas = meas[:, validation_matrix[:, ti]] # Get the validated measurements for this track
-            yhat = track['filt'].sensor_model['h'](track['x'][-1]) # Calculate the predicted measurement for this track
-            H = track['filt'].sensor_model['dhdx'](track['x'][-1])
-            py = stats.multivariate_normal.pdf(val_meas.squeeze().T, mean=yhat.flatten(), cov=H@track['P'][-1]@H.T+track['filt'].sensor_model['R'])
-            association_matrix[validation_matrix[:, ti], ti] = np.log(track['filt'].sensor_model['PD']*py/(1-track['filt'].sensor_model['PD'])) # PG assumed = 1
+        if tracks:
+            # All of the tracks are assumed to use the same sensor model!
+            x = np.vstack([track['x'][-1] for track in tracks]).T
+            yhat_t = tracks[0]['filt'].sensor_model['h'](x) # Returns a (ny x nx) matrix
+            H_t = tracks[0]['filt'].sensor_model['dhdx'](x) # Returns a (ny x nC x nx x nC) tensor
+            for ti, track in enumerate(tracks): # Iterate over confirmed tracks
+                validation_matrix[:, ti] = self.gater.gate(track['x'][-1], track['P'][-1], meas)
+                if validation_matrix[:, ti].any(): # If any measurements are validated
+                    val_meas = meas[:, validation_matrix[:, ti]] # Get the validated measurements for this track
+                    if Nc == 1: # Because of how numpy handles its matrices
+                        yhat = yhat_t
+                        H = H_t.squeeze()
+                    else:
+                        yhat = yhat_t[:, ti]
+                        H = H_t[:, ti, :, ti]
+                    py = stats.multivariate_normal.pdf(val_meas.squeeze().T, mean=yhat.flatten(), cov=H@track['P'][-1]@H.T+track['filt'].sensor_model['R'])
+                    association_matrix[validation_matrix[:, ti], ti] = np.log(track['filt'].sensor_model['PD']*py/(1-track['filt'].sensor_model['PD'])) # PG assumed = 1
         return association_matrix, validation_matrix
 
     def _update_track(self, meas, track):
@@ -387,16 +396,26 @@ class JPDA():
         validation_matrix = np.zeros((Nc, ny+1), dtype=bool)
         validation_matrix[:, 0] = 1
         likelihood_matrix = np.zeros((Nc, ny+1))
+        likelihood_matrix[:, 0] = 1-tracks[0]['filt'].sensor_model['PD'] # PG assumed 1
 
-        for ti, track in enumerate(tracks): # Iterate over confirmed tracks
-            validation_matrix[ti, 1:] = self.gater.gate(track['x'][-1], track['P'][-1], meas)
-            # Entry for validated tracks
-            val_meas = meas[:, validation_matrix[ti, 1:]] # Get the validated measurements for this track
-            yhat = track['filt'].sensor_model['h'](track['x'][-1]) # Calculate the predicted measurement for this track
-            H = track['filt'].sensor_model['dhdx'](track['x'][-1])
-            py = stats.multivariate_normal.pdf(val_meas.squeeze().T, mean=yhat.flatten(), cov=H@track['P'][-1]@H.T+track['filt'].sensor_model['R'])
-            likelihood_matrix[ti, np.where(validation_matrix[ti, 1:])[0]+1] = track['filt'].sensor_model['PD']*py
-            likelihood_matrix[ti, 0] = 1-track['filt'].sensor_model['PD'] # PG assumed 1
+        if tracks:
+            # All of the tracks are assumed to use the same sensor model!
+            x = np.vstack([track['x'][-1] for track in tracks]).T
+            yhat_t = tracks[0]['filt'].sensor_model['h'](x) # Returns a (ny x nx) matrix
+            H_t = tracks[0]['filt'].sensor_model['dhdx'](x) # Returns a (ny x nC x nx x nC) tensor
+            for ti, track in enumerate(tracks): # Iterate over confirmed tracks
+                validation_matrix[ti, 1:] = self.gater.gate(track['x'][-1], track['P'][-1], meas)
+                # Entry for validated tracks
+                if validation_matrix[ti, 1:].any(): # If any measurements are validated
+                    val_meas = meas[:, validation_matrix[ti, 1:]] # Get the validated measurements for this track
+                    if Nc == 1: # Because of how numpy handles its matrices
+                        yhat = yhat_t
+                        H = H_t.squeeze()
+                    else:
+                        yhat = yhat_t[:, ti]
+                        H = H_t[:, ti, :, ti]
+                    py = stats.multivariate_normal.pdf(val_meas.squeeze().T, mean=yhat.flatten(), cov=H@track['P'][-1]@H.T+track['filt'].sensor_model['R'])
+                    likelihood_matrix[ti, np.where(validation_matrix[ti, 1:])[0]+1] = track['filt'].sensor_model['PD']*py
         return likelihood_matrix, validation_matrix
 
     def _update_track(self, meas, track, association_probability):
@@ -624,8 +643,12 @@ class MHT():
                 # Entry for validated tracks
                 if validation_matrix[:, ti].any(): # If any measurements are validated
                     val_meas = meas[:, validation_matrix[:, ti]] # Get the validated measurements for this track
-                    yhat = yhat_t[:, ti]
-                    H = H_t[:, ti, :, ti]
+                    if Nc == 1: # Because of how numpy handles its matrices
+                        yhat = yhat_t
+                        H = H_t.squeeze()
+                    else:
+                        yhat = yhat_t[:, ti]
+                        H = H_t[:, ti, :, ti]
                     py = stats.multivariate_normal.pdf(val_meas.squeeze().T, mean=yhat.flatten(), cov=H@track['P'][-1]@H.T+track['filt'].sensor_model['R'])
                     association_matrix[validation_matrix[:, ti], ti] = np.log(track['filt'].sensor_model['PD']*py/(1-track['filt'].sensor_model['PD'])) # PG assumed = 1
                     likelihood_matrix[validation_matrix[:, ti], ti] = track['filt'].sensor_model['PD']*py