6 files changed, 453 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)
diff --git a/murl/particles/baseline.py b/murl/particles/baseline.py
new file mode 100644
index 0000000..2d8e045
--- /dev/null
+++ b/murl/particles/baseline.py
@@ -0,0 +1,21 @@
+from murl.particles import ParticleCloud
+from murl.sensors import CONE_WIDTH
+
+import numpy as np
+from scipy.stats import norm
+
+class Cloud(ParticleCloud):
+    def calculate_weight(self, angle_dists, _robot_guesses, _robot_scans):
+        self._particles[:,3] = 0.0
+        mask = self._world.is_valid(self._particles)
+        self._particles[~mask,3] = float('-inf')
+        count = np.ma.array(self._particles[:,3], mask=~mask).count()
+
+        for angle, dist in angle_dists:
+            angles = angle + self._rng.uniform(-CONE_WIDTH/2, CONE_WIDTH/2,
+                                               (self.RAY_SAMPLES, count))
+            true_dists = self._world.get_dist(self._particles[mask], angles)
+            if np.isfinite(dist):
+                self._particles[mask,3] += norm.logpdf(dist, true_dists, 0.8)
+            else:
+                self._particles[mask,3] += norm.logsf(3, true_dists, 0.8)
diff --git a/murl/particles/centralised.py b/murl/particles/centralised.py
new file mode 100644
index 0000000..33be0b3
--- /dev/null
+++ b/murl/particles/centralised.py
@@ -0,0 +1,30 @@
+from murl.particles import ParticleCloud
+from murl.sensors import CONE_WIDTH
+
+import numpy as np
+from scipy.stats import multivariate_normal, norm
+
+class Cloud(ParticleCloud):
+    BETA = 1.E+0
+
+    def calculate_weight(self, angle_dists, robot_guesses, robot_scans):
+        self._particles[:,3] = 0.0
+        mask = self._world.is_valid(self._particles)
+        self._particles[~mask,3] = float('-inf')
+        count = np.ma.array(self._particles[:,3], mask=~mask).count()
+
+        for angle, dist in angle_dists:
+            angles = angle + self._rng.uniform(-CONE_WIDTH/2, CONE_WIDTH/2,
+                                               (self.RAY_SAMPLES, count))
+            true_dists = self._world.get_dist(self._particles[mask], angles)
+
+            if np.isfinite(dist):
+                self._particles[mask,3] += norm.logpdf(dist, true_dists, 0.8)
+            else:
+                self._particles[mask,3] += norm.logsf(3, true_dists, 0.8)
+
+        for guess, scan in zip(robot_guesses, robot_scans):
+            angles = scan[1] + self._particles[mask,2]
+            poses = scan[0] * np.stack([np.cos(angles), np.sin(angles)]).transpose() + self._particles[mask,:2]
+            update = multivariate_normal.logpdf(poses, guess[0], guess[1])
+            self._particles[mask,3] += self.BETA * update
diff --git a/murl/sensors.py b/murl/sensors.py
new file mode 100644
index 0000000..700db35
--- /dev/null
+++ b/murl/sensors.py
@@ -0,0 +1,180 @@
+import numpy as np
+import rospy
+
+from gazebo_msgs.msg import ModelStates
+from nav_msgs.msg import Odometry
+from sensor_msgs.msg import LaserScan
+from tf.transformations import euler_from_quaternion
+
+SCAN_ANGLES = [0., np.pi / 4., -np.pi / 4., np.pi / 2., -np.pi / 2.]
+CONE_WIDTH = np.pi / 180. * 3.1 # 3.1 degrees cone of view (3 rays).
+
+class SimpleLaser:
+    def __init__(self, robot):
+        rospy.Subscriber('/robot{}/scan'.format(robot), LaserScan, self.callback)
+        self._angles = SCAN_ANGLES
+        self._width = CONE_WIDTH
+        self._measurements = [float('inf')] * len(self._angles)
+        self._indices = None
+
+    def callback(self, msg):
+        # Helper for angles.
+        def _within(x, a, b):
+            pi2 = np.pi * 2.
+            x %= pi2
+            a %= pi2
+            b %= pi2
+            if a < b:
+                return a <= x and x <= b
+            return a <= x or x <= b;
+
+        # Compute indices the first time.
+        if self._indices is None:
+            self._indices = [[] for _ in range(len(self._angles))]
+            for i, d in enumerate(msg.ranges):
+                angle = msg.angle_min + i * msg.angle_increment
+                for j, center_angle in enumerate(self._angles):
+                    if _within(angle, center_angle - self._width / 2., center_angle + self._width / 2.):
+                        self._indices[j].append(i)
+
+        ranges = np.array(msg.ranges)
+        for i, idx in enumerate(self._indices):
+            # We do not take the minimum range of the cone but the 10-th percentile for robustness.
+            self._measurements[i] = np.percentile(ranges[idx], 10)
+
+    @property
+    def ready(self):
+        return not np.isnan(self._measurements[0])
+
+    @property
+    def measurements(self):
+        return self._measurements
+
+
+class Motion:
+    def __init__(self, robot):
+        self._previous_time = None
+        self._delta_pose = np.array([0., 0., 0.], dtype=np.float32)
+        rospy.Subscriber('/robot{}/odom'.format(robot), Odometry, self.callback)
+
+    def callback(self, msg):
+        u = msg.twist.twist.linear.x
+        w = msg.twist.twist.angular.z
+        if self._previous_time is None:
+            self._previous_time = msg.header.stamp
+        current_time = msg.header.stamp
+        dt = (current_time - self._previous_time).to_sec()
+        self._delta_pose[0] += u * dt
+        self._delta_pose[1] += 0.
+        self._delta_pose[2] += w * dt
+        self._previous_time = current_time
+
+    @property
+    def ready(self):
+        return True
+
+    @property
+    def delta_pose(self):
+        ret = self._delta_pose.copy()
+        self._delta_pose[:] = 0
+        return ret
+
+
+class SimpleLaser:
+    def __init__(self, robot):
+        rospy.Subscriber('/robot{}/scan'.format(robot), LaserScan, self.callback)
+        self._angles = SCAN_ANGLES
+        self._width = CONE_WIDTH
+        self._measurements = [float('inf')] * len(self._angles)
+        self._indices = None
+
+    def callback(self, msg):
+        # Helper for angles.
+        def _within(x, a, b):
+            pi2 = np.pi * 2.
+            x %= pi2
+            a %= pi2
+            b %= pi2
+            if a < b:
+                return a <= x and x <= b
+            return a <= x or x <= b;
+
+        # Compute indices the first time.
+        if self._indices is None:
+            self._indices = [[] for _ in range(len(self._angles))]
+            for i, d in enumerate(msg.ranges):
+                angle = msg.angle_min + i * msg.angle_increment
+                for j, center_angle in enumerate(self._angles):
+                    if _within(angle, center_angle - self._width / 2., center_angle + self._width / 2.):
+                        self._indices[j].append(i)
+
+        ranges = np.array(msg.ranges)
+        for i, idx in enumerate(self._indices):
+            # We do not take the minimum range of the cone but the 10-th percentile for robustness.
+            self._measurements[i] = np.percentile(ranges[idx], 10)
+
+    @property
+    def ready(self):
+        return not np.isnan(self._measurements[0])
+
+    @property
+    def measurements(self):
+        return self._measurements
+
+
+class Motion(object):
+    def __init__(self, robot):
+        self._previous_time = None
+        self._delta_pose = np.array([0., 0., 0.], dtype=np.float32)
+        rospy.Subscriber('/robot{}/odom'.format(robot), Odometry, self.callback)
+
+    def callback(self, msg):
+        u = msg.twist.twist.linear.x
+        w = msg.twist.twist.angular.z
+        if self._previous_time is None:
+            self._previous_time = msg.header.stamp
+        current_time = msg.header.stamp
+        dt = (current_time - self._previous_time).to_sec()
+        self._delta_pose[0] += u * dt
+        self._delta_pose[1] += 0.
+        self._delta_pose[2] += w * dt
+        self._previous_time = current_time
+
+    @property
+    def ready(self):
+        return True
+
+    @property
+    def delta_pose(self):
+        ret = self._delta_pose.copy()
+        self._delta_pose[:] = 0
+        return ret
+
+
+class GroundtruthPose:
+    def __init__(self, robot):
+        rospy.Subscriber('/gazebo/model_states', ModelStates, self.callback)
+        self._pose = np.array([np.nan, np.nan, np.nan], dtype=np.float32)
+        self._name = 'Robot{}'.format(robot)
+
+    def callback(self, msg):
+        idx = [i for i, n in enumerate(msg.name) if n == self._name]
+        if not idx:
+            raise ValueError('Specified name "{}" does not exist.'.format(self._name))
+        idx = idx[0]
+        self._pose[0] = msg.pose[idx].position.x
+        self._pose[1] = msg.pose[idx].position.y
+        _, _, yaw = euler_from_quaternion([
+                msg.pose[idx].orientation.x,
+                msg.pose[idx].orientation.y,
+                msg.pose[idx].orientation.z,
+                msg.pose[idx].orientation.w])
+        self._pose[2] = yaw
+
+    @property
+    def ready(self):
+        return not np.isnan(self._pose[0])
+
+    @property
+    def pose(self):
+        return self._pose
diff --git a/murl/world/__init__.py b/murl/world/__init__.py
new file mode 100644
index 0000000..254a0f6
--- /dev/null
+++ b/murl/world/__init__.py
@@ -0,0 +1,15 @@
+from abc import ABC, abstractmethod
+
+class World(ABC):
+    ROBOT_RADIUS = 0.105 / 2.
+    @abstractmethod
+    def is_valid(self, pose):
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_dist(self, sources, angles):
+        raise NotImplementedError
+
+    @abstractmethod
+    def sample(self, particles):
+        raise NotImplementedError
diff --git a/murl/world/simple.py b/murl/world/simple.py
new file mode 100644
index 0000000..195c4db
--- /dev/null
+++ b/murl/world/simple.py
@@ -0,0 +1,77 @@
+import numpy as np
+
+from murl.world import World
+
+WALL_OFFSET = 2.
+CYLINDER_POSITION = np.array([.3, .2], dtype=np.float32)
+CYLINDER_OFFSET = .3
+
+SEGMENTS = list(map(lambda ps: (np.array(ps[0]), np.array(ps[1]), ps[2]),
+                    [[[-WALL_OFFSET, -WALL_OFFSET], [0, 1], 2*WALL_OFFSET],
+                     [[-WALL_OFFSET, WALL_OFFSET], [1, 0], 2*WALL_OFFSET],
+                     [[WALL_OFFSET, WALL_OFFSET], [0, -1], 2*WALL_OFFSET],
+                     [[WALL_OFFSET, -WALL_OFFSET], [-1, 0], 2*WALL_OFFSET]]))
+
+class SimpleWorld(World):
+    WALL_RADIUS = WALL_OFFSET - World.ROBOT_RADIUS
+    CYLINDER_RADIUS = CYLINDER_OFFSET + World.ROBOT_RADIUS
+
+    def __init__(self, rng):
+        self._rng = rng
+
+    def is_valid(self, poses):
+        return ((-self.WALL_RADIUS < poses[...,0])
+              & (poses[...,0] < self.WALL_RADIUS)
+              & (-self.WALL_RADIUS < poses[...,1])
+              & (poses[...,1] < self.WALL_RADIUS)
+              & (np.sum((poses[...,:2] - CYLINDER_POSITION) ** 2, axis=-1) > self.CYLINDER_RADIUS ** 2))
+
+    def get_dist(self, sources, angles):
+        def dot(x, y):
+            return x[...,0] * y[...,0] + x[...,1] * y[...,1]
+        def cross(x, y):
+            return x[...,0] * y[...,1] - x[...,1] * y[...,0]
+
+        def intersect_segment(sources, angle, start, direction, length):
+            Y = start - sources[:,:2]
+            r = np.stack((np.cos(angle + sources[:,2]),
+                          np.sin(angle + sources[:,2]))).T
+            t = cross(Y, direction) / cross(r, direction)
+            t1 = cross(Y, r) / cross(r, direction)
+
+            return np.where((t >= 0) & (t1 >= 0) & (t1 <= length), t, float('inf'))
+
+        def intersect_cylinder(sources, angle, center, radius):
+            Y = center - sources[:,:2]
+            d = np.stack((np.cos(angle + sources[:,2]),
+                          np.sin(angle + sources[:,2]))).T
+            b = dot(d, Y)
+            c = dot(Y, Y) - radius**2
+            disc = b**2 - dot(Y, Y) + radius**2
+
+            mask = (disc > 0) & (b > 0) & (b ** 2 > disc)
+
+            b[~mask] = float('inf')
+            b[mask] -= np.sqrt(disc[mask])
+
+            return b
+
+        intersections = np.stack([np.stack(
+                (intersect_segment(sources, angle, *SEGMENTS[0]),
+                 intersect_segment(sources, angle, *SEGMENTS[1]),
+                 intersect_segment(sources, angle, *SEGMENTS[2]),
+                 intersect_segment(sources, angle, *SEGMENTS[3]),
+                 intersect_cylinder(sources, angle, CYLINDER_POSITION, CYLINDER_OFFSET))
+            ) for angle in angles])
+
+        return np.amin(intersections, axis=(0,1))
+
+    def sample(self, particles):
+        particles[:,2] = self._rng.uniform(-np.pi, np.pi, size=len(particles))
+
+        invalid = np.ones(shape=len(particles), dtype=np.bool)
+        while np.any(invalid):
+            count = np.ma.array(particles[:,0], mask=~invalid).count()
+            particles[invalid,0] = self._rng.uniform(-self.WALL_RADIUS, self.WALL_RADIUS, size=count)
+            particles[invalid,1] = self._rng.uniform(-self.WALL_RADIUS, self.WALL_RADIUS, size=count)
+            invalid = ~self.is_valid(particles)