1 files changed, 130 insertions, 0 deletions
diff --git a/murl/particles/__init__.py b/murl/particles/__init__.py
new file mode 100644
index 0000000..e460c10
--- /dev/null
+++ b/murl/particles/__init__.py
@@ -0,0 +1,130 @@
+from abc import ABC, abstractmethod
+
+import numpy as np
+import rospy
+
+from geometry_msgs.msg import Point32
+from sensor_msgs.msg import ChannelFloat32, PointCloud
+
+class ParticleCloud(ABC):
+    PARTICLE_COUNT = 10000
+    RESAMPLE_PROB = 2**-10
+    RAY_SAMPLES = 3
+    ALPHAS = [0.01, 0.00, 0.01, 0.00]
+
+    def __init__(self, world, robot, rng):
+        self._world = world
+        self._particles = np.array([[0.0] * 4] * self.PARTICLE_COUNT)
+        self._publisher = rospy.Publisher('/robot{}/particles'.format(robot), PointCloud, queue_size=1)
+        self._frame = 'robot{}/odom'.format(robot)
+        self._robot = robot
+        self._rng = rng
+
+
+    def resample_all(self):
+        self.resample(np.array([True] * self.PARTICLE_COUNT))
+
+
+    def resample(self, mask):
+        tmp = self._particles[mask,:]
+        self._world.sample(tmp)
+        self._particles[mask,:] = tmp
+
+
+    def move(self, delta):
+        rot1 = np.arctan2(delta[1], delta[0])
+
+        # Assume large rot1 means reversing
+        if rot1 < -np.pi:
+            rot1 += np.pi
+        elif rot1 > np.pi:
+            rot1 -= np.pi
+
+        trans = np.sqrt(delta[0]**2 + delta[1]**2)
+        rot2 = delta[2] - rot1
+
+        rot1_var = self.ALPHAS[0] * rot1**2 + self.ALPHAS[1] * trans**2
+        trans_var = self.ALPHAS[2] * trans**2 + self.ALPHAS[3] * (rot1**2 + rot2**2)
+        rot2_var = self.ALPHAS[0] * rot2**2 + self.ALPHAS[1] * trans**2
+
+        rot1 += self._rng.normal(scale = np.sqrt(rot1_var), size=self.PARTICLE_COUNT)
+        trans += self._rng.normal(scale = np.sqrt(trans_var), size=self.PARTICLE_COUNT)
+        rot2 += self._rng.normal(scale = np.sqrt(rot2_var), size=self.PARTICLE_COUNT)
+
+        self._particles[:,0] += trans * np.cos(rot1 + self.yaws)
+        self._particles[:,1] += trans * np.sin(rot1 + self.yaws)
+        self._particles[:,2] += rot1 + rot2
+
+        mask = self._rng.random(self.PARTICLE_COUNT) < self.RESAMPLE_PROB
+        self.resample(mask)
+
+
+    @abstractmethod
+    def calculate_weight(self, angle_dists, robot_guesses, robot_scans):
+        raise NotImplementedError()
+
+
+    def average(self):
+        biggest = max(self._particles[:,3])
+        mask = self._particles[:,3] > biggest - 5
+        masked = self._particles[mask]
+
+        if len(masked) < 20 or not np.isfinite(biggest):
+            print('Massive collapse for robot', self._robot)
+            self._particles[:,3] = 0.
+            self.resample_all()
+            return
+
+        weights = np.exp(masked[:,3] - max(masked[:,3]))
+        self._particles = self._rng.choice(masked, size=self.PARTICLE_COUNT, p=weights/np.sum(weights))
+        self._particles[:,3] = 0.0
+
+
+    def guess_position(self):
+        weights = np.exp(self._particles[:,3])
+        weights /= np.sum(weights)
+
+        mean = np.average(self._particles[:,:2], weights=weights, axis=0)
+        cov = np.cov(self._particles[:,:2], rowvar=False, ddof=0, aweights=weights)
+
+        # x and y being too correlated is suspicious
+        if abs(cov[0,0]*cov[1,1] - cov[0,1] * cov[1,0]) < 1e-15:
+            print('Singular covariance matrix for robot', self._robot)
+            # pretend to uncorrelate the data slightly
+            cov[0,0] += 0.1
+            cov[1,1] += 0.1
+
+        return mean, cov
+
+
+    def publish(self, frame_id):
+        msg = PointCloud()
+        msg.header.seq = frame_id
+        msg.header.stamp = rospy.Time.now()
+        msg.header.frame_id = self._frame
+        intensity = ChannelFloat32()
+        intensity.name = 'intensity'
+        msg.channels.append(intensity)
+        msg.points = [Point32(x = particle[0], y = particle[1], z = 0.05) for particle in self._particles]
+        intensity.values = list(np.exp(self._particles[:,3]))
+        self._publisher.publish(msg)
+
+
+    @property
+    def particles(self):
+        return self._particles
+
+
+    @property
+    def yaws(self):
+        return self._particles[:,2]
+
+
+    @property
+    def log_weights(self):
+        return self._particles[:,3]
+
+
+    @property
+    def weights(self):
+        return np.exp(self.log_weights)