FastSLAM_2.py

import numpy as np 
from FastSLAM.Particle_2 import Particle, Landmark

class FastSLAM:

    def __init__(self, initial_pose, n_particles, R, Q, sensor_range, p0=0.01): 
        self.initial_pose = initial_pose
        self.M = n_particles
        self.R = R
        self.Q = Q
        self.sensor_range = sensor_range
        self.p0 = p0
        self.particles = [Particle(self.initial_pose) for _ in range(n_particles)]
        self.next_landmark_id = 0
        self.association_threshold = 5.99
        
        self.R_inv = np.linalg.inv(R)
        self.Q_inv = np.linalg.inv(Q)
        self.det_2pi_Q = np.linalg.det(2 * np.pi * Q)
        self.sqrt_det_2pi_Q = np.sqrt(self.det_2pi_Q)

    def update(self, z, u, dt):
        if len(z) > 0:
            distances = z[:, 0]
            closest_indices = np.argsort(distances)
            z = z[closest_indices]
        
        weights = np.zeros(self.M)

        for p in range(self.M):
            particle = self.particles[p]
            observed_landmarks = []

            if len(z) == 0:
                x_hat = particle.predict_pose(u, dt)
                noise = np.random.multivariate_normal(np.zeros(3), self.R)
                particle.pose = x_hat + noise
                particle.pose[2] = particle.wrap(particle.pose[2])
                particle.weight = 1.0

            else:
                x_hat = particle.predict_pose(u, dt)
                
                mu_proposal = x_hat.copy() 
                Sigma_proposal = self.R.copy()
                
                particle.weight = 1.0
                
                observations_with_landmarks = []
                observations_new_landmarks = []
                
                for z_i in z:
                    best_pi = self.p0
                    best_landmark = None
                    best_proposal_update = None
                    
                    for landmark in particle.landmarks:
                        mu_updated, Sigma_updated, Q_j, Q_j_inv = particle.incremental_proposal_update(
                            z_i, landmark, x_hat, mu_proposal, Sigma_proposal, self.Q, self.R_inv
                        )
                        
                        x_temp = np.random.multivariate_normal(mu_updated, Sigma_updated)
                        x_temp[2] = particle.wrap(x_temp[2])

                        z_hat = particle.predict_measurement(landmark, x_temp)
                        
                        nu = z_i - z_hat
                        nu[1] = particle.wrap(nu[1])

                        det = np.linalg.det(2 * np.pi * Q_j)
                        
                        if det <= 0:
                            continue
                        
                        exponent = -0.5 * (nu @ Q_j_inv @ nu)
                        pi_j = np.exp(exponent) / np.sqrt(det)

                        if pi_j > best_pi:
                            best_pi = pi_j
                            best_landmark = landmark
                            best_proposal_update = (mu_updated, Sigma_updated)
                    
                    if best_landmark is not None:
                        mu_proposal, Sigma_proposal = best_proposal_update
                        observations_with_landmarks.append((z_i, best_landmark))
                    else:
                        observations_new_landmarks.append(z_i)
                
                particle.pose = np.random.multivariate_normal(mu_proposal, Sigma_proposal)
                particle.pose[2] = particle.wrap(particle.pose[2])
                
                for z_i, landmark in observations_with_landmarks:
                    particle.update_landmark(landmark.landmark_id, z_i, self.Q, self.Q_inv)
                    observed_landmarks.append(landmark.landmark_id)
                    
                    w_i = particle.importance_weight(z_i, landmark, self.Q, self.R)
                    particle.weight *= w_i

                for z_new in observations_new_landmarks:
                    particle.initialize_landmark(z_new, self.Q, self.next_landmark_id)
                    observed_landmarks.append(self.next_landmark_id)
                    self.next_landmark_id += 1
                    particle.weight *= self.p0

            particle.update_landmark_counters(observed_landmarks, self.sensor_range)
            weights[p] = particle.weight

        weights_sum = weights.sum()
        if weights_sum > 0:
            weights /= weights_sum
        else:
            weights[:] = 1.0 / self.M

        indices = np.random.choice(self.M, size=self.M, p=weights, replace=True)

        new_particles = [self.copy_particle(self.particles[i]) for i in indices]
        for particle in new_particles:
            particle.weight = 1.0

        self.particles = new_particles

    def copy_particle(self, particle):
        new_particle = Particle(particle.pose.copy())
        new_particle.weight = particle.weight
        new_particle.n_landmarks = particle.n_landmarks

        for landmark in particle.landmarks:
            new_landmark = Landmark(
                mu=landmark.mu.copy(),
                covariance=landmark.covariance.copy(), 
                counter=landmark.counter,
                landmark_id=landmark.landmark_id
            )
            new_particle.landmarks.append(new_landmark)

        return new_particle
        
    def get_best_particle(self):
        best_idx = np.argmax([p.weight for p in self.particles])
        return self.particles[best_idx]
    
    def get_map_estimate(self):

        poses = np.array([p.pose for p in self.particles])

        x_est = np.mean(poses[:, 0])
        y_est = np.mean(poses[:, 1])
        
        sin_theta = np.mean(np.sin(poses[:, 2]))
        cos_theta = np.mean(np.cos(poses[:, 2]))
        theta_est = np.arctan2(sin_theta, cos_theta)
        
        pose = np.array([x_est, y_est, theta_est])

        landmark_dict = {}
        
        for particle in self.particles:
            for lm in particle.landmarks:
                if lm.landmark_id not in landmark_dict:
                    landmark_dict[lm.landmark_id] = []
                landmark_dict[lm.landmark_id].append((lm.mu, lm.covariance, lm.counter))
        
        landmarks_list = []
        for lm_id, observations in landmark_dict.items():
            n_obs = len(observations)
            
            if n_obs > 0:
                mu_avg = sum(obs[0] for obs in observations) / n_obs
                cov_avg = sum(obs[1] for obs in observations) / n_obs
                counter_avg = int(sum(obs[2] for obs in observations) / n_obs)
                
                landmarks_list.append({
                    'id': lm_id,
                    'mu': mu_avg,
                    'covariance': cov_avg,
                    'counter': counter_avg
                })
    
        landmarks_list.sort(key=lambda x: x['id'])
        landmarks = np.array([lm['mu'] for lm in landmarks_list]) if landmarks_list else np.array([])
    
        return pose, landmarks