Introducing GeometryIndexed as a base class for MapHistory

src/c3nav/mapdata/cache.py

@@ -3,7 +3,6 @@ import math
 import os
 import struct
 import threading
-import traceback
 from itertools import chain
 
 import numpy as np
@@ -12,6 +11,7 @@ from django.db.models.signals import m2m_changed, post_delete
 from PIL import Image
 from shapely import prepared
 from shapely.geometry import box
+from shapely.geometry.base import BaseGeometry
 from shapely.ops import unary_union
 
 from c3nav.mapdata.models import MapUpdate
@@ -20,47 +20,197 @@ from c3nav.mapdata.utils.models import get_submodels
 logger = logging.getLogger('c3nav')
 
 
-class MapHistory:
+class GeometryIndexed:
     # binary format (everything little-endian):
+    # 1 byte (uint8): variant id
     # 1 byte (uint8): resolution
     # 2 bytes (uint16): origin x
     # 2 bytes (uint16): origin y
     # 2 bytes (uint16): origin width
     # 2 bytes (uint16): origin height
-    # 2 bytes (uint16): number of updates
+    # (optional meta data, depending on subclass)
-    # n uptates times:
+    # x bytes data, line after line. (cell size depends on subclass)
-    #     4 bytes (uint32): update id
+    dtype = np.uint16
-    #     4 bytes (uint32): timestamp
+    variant_id = 0
-    # width*height*2 bytes:
-    #     data array (line after line) with uint16 cells
-    empty_array = np.empty((0, 0), dtype=np.uint16)
 
-    def __init__(self, resolution=settings.CACHE_RESOLUTION, x=0, y=0, updates=None, data=empty_array, filename=None):
+    def __init__(self, resolution=settings.CACHE_RESOLUTION, x=0, y=0, data=None, filename=None):
         self.resolution = resolution
         self.x = x
         self.y = y
-        self.updates = updates
+        self.data = data if data is not None else self._get_empty_array()
-        self.data = data
         self.filename = filename
+
+    @classmethod
+    def _get_empty_array(cls):
+        return np.empty((0, 0), dtype=cls.dtype)
+
+    @classmethod
+    def open(cls, filename):
+        with open(filename, 'rb') as f:
+            instance = cls.read(f)
+        instance.filename = filename
+        return instance
+
+    @classmethod
+    def read(cls, f):
+        variant_id, resolution, x, y, width, height = struct.unpack('<BBHHHH', f.read(10))
+        if variant_id != cls.variant_id:
+            raise ValueError('variant id does not match')
+
+        kwargs = {
+            'resolution': resolution,
+            'x': x,
+            'y': y,
+        }
+        cls._read_metadata(f, kwargs)
+
+        # noinspection PyTypeChecker
+        kwargs['data'] = np.fromstring(f.read(width*height*cls.dtype().itemsize), cls.dtype).reshape((height, width))
+        return cls(**kwargs)
+
+    @classmethod
+    def _read_metadata(cls, f, kwargs):
+        pass
+
+    def save(self, filename=None):
+        if filename is None:
+            filename = self.filename
+        if filename is None:
+            raise ValueError('Missing filename.')
+
+        with open(filename, 'wb') as f:
+            self.write(f)
+
+    def write(self, f):
+        f.write(struct.pack('<BBHHHH', self.variant_id, self.resolution, self.x, self.y, *reversed(self.data.shape)))
+        self._write_metadata(f)
+        f.write(self.data.tobytes('C'))
+
+    def _write_metadata(cls, f):
+        pass
+
+    def _get_geometry_bounds(self, geometry):
+        minx, miny, maxx, maxy = geometry.bounds
+        return (
+            int(math.floor(minx / self.resolution)),
+            int(math.floor(miny / self.resolution)),
+            int(math.ceil(maxx / self.resolution)),
+            int(math.ceil(maxy / self.resolution)),
+        )
+
+    def fit_bounds(self, minx, miny, maxx, maxy):
+        height, width = self.data.shape
+
+        if self.data.size:
+            minx = min(self.x, minx)
+            miny = min(self.y, miny)
+            maxx = max(self.x + width, maxx - minx)
+            maxy = max(self.y + height, maxy - miny)
+
+        new_data = np.zeros((maxy - miny, maxx - minx), dtype=self.dtype)
+
+        if self.data.size:
+            dx = self.x - minx
+            dy = self.y - miny
+            new_data[dy:(dy + height), dx:(dx + width)] = self.data
+
+        self.data = new_data
+        self.x = minx
+        self.y = miny
+
+    def get_geometry_cells(self, geometry, bounds=None):
+        if bounds is None:
+            bounds = self._get_geometry_bounds(geometry)
+        minx, miny, maxx, maxy = bounds
+
+        height, width = self.data.shape
+        minx = max(minx, self.x)
+        miny = max(miny, self.y)
+        maxx = min(maxx, self.x + width)
+        maxy = min(maxy, self.y + height)
+
+        cells = np.zeros_like(self.data, dtype=np.bool)
+        prep = prepared.prep(geometry)
+        res = self.resolution
+        for iy, y in enumerate(range(miny * res, maxy * res, res), start=miny - self.y):
+            for ix, x in enumerate(range(minx * res, maxx * res, res), start=minx - self.x):
+                if prep.intersects(box(x, y, x + res, y + res)):
+                    cells[iy, ix] = True
+
+        return cells
+
+    @property
+    def bounds(self):
+        height, width = self.data.shape
+        return self.x, self.y, self.x+width, self.y+height
+
+    def __getitem__(self, key):
+        if isinstance(key, BaseGeometry):
+            bounds = self._get_geometry_bounds(key)
+            return self.data[self.get_geometry_cells(key, bounds)]
+
+        raise TypeError('GeometryIndexed index must be a shapely geometry, not %s' % type(key).__name__)
+
+    def __setitem__(self, key, value):
+        if isinstance(key, BaseGeometry):
+            bounds = self._get_geometry_bounds(key)
+            self.fit_bounds(*bounds)
+            self.data[self.get_geometry_cells(key, bounds)] = value
+            return
+
+        raise TypeError('GeometryIndexed index must be a shapely geometry, not %s' % type(key).__name__)
+
+    def to_image(self):
+        from c3nav.mapdata.models import Source
+        (minx, miny), (maxx, maxy) = Source.max_bounds()
+
+        height, width = self.data.shape
+        image_data = np.zeros((int(math.ceil((maxy-miny)/self.resolution)),
+                               int(math.ceil((maxx-minx)/self.resolution))), dtype=np.uint8)
+
+        minval = max(self.data.min(), 0)
+        maxval = max(self.data.max(), minval+0.01)
+        visible_data = (self.data.astype(float)-minval*255/(maxval-minval)).clip(0, 255).astype(np.uint8)
+        image_data[self.y:self.y+height, self.x:self.x+width] = visible_data
+
+        return Image.fromarray(np.flip(image_data, axis=0), 'L')
+
+
+class MapHistory(GeometryIndexed):
+    # metadata format:
+    # 2 bytes (uint16): number of updates
+    # n uptates times:
+    #     4 bytes (uint32): update id
+    #     8 bytes (uint64): timestamp
+    dtype = np.uint16
+    variant_id = 1
+
+    def __init__(self, updates, **kwargs):
+        super().__init__(**kwargs)
+        self.updates = updates
         self.unfinished = False
 
+    @classmethod
+    def _read_metadata(cls, f, kwargs):
+        num_updates = struct.unpack('<H', f.read(2))[0]
+        updates = struct.unpack('<'+'II'*num_updates, f.read(num_updates*8))
+        updates = list(zip(updates[0::2], updates[1::2]))
+        kwargs['updates'] = updates
+
+    def _write_metadata(self, f):
+        f.write(struct.pack('<H', len(self.updates)))
+        f.write(struct.pack('<'+'II'*len(self.updates), *chain(*self.updates)))
+
+    # todo: continue
     @classmethod
     def open(cls, filename, default_update=None):
         try:
-            with open(filename, 'rb') as f:
+            instance = super().open(filename)
-                resolution, x, y, width, height, num_updates = struct.unpack('<BHHHHH', f.read(11))
+        except FileNotFoundError:
-                updates = struct.unpack('<'+'II'*num_updates, f.read(num_updates*8))
-                updates = list(zip(updates[0::2], updates[1::2]))
-                # noinspection PyTypeChecker
-                data = np.fromstring(f.read(width*height*2), np.uint16).reshape((height, width))
-                return cls(resolution, x, y, list(updates), data, filename)
-        except (FileNotFoundError, struct.error) as e:
-            logger.info('Exception in MapHistory loading! %s' % traceback.format_exc())
             if default_update is None:
-                default_update = MapUpdate.last_update()
+                default_update = MapUpdate.last_processed_update()
-            new_empty = cls(updates=[default_update], filename=filename)
+            instance = cls(updates=[default_update], filename=filename)
-            new_empty.save(filename)
+        return instance
-            return new_empty
 
     @staticmethod
     def level_filename(level_id, mode):
@@ -90,126 +240,38 @@ class MapHistory:
             cls.cached[(level_id, mode)] = result
             return result
 
-    def save(self, filename=None):
+    def add_geometry(self, geometry, update):
-        if filename is None:
+        if self.updates[-1] != update:
-            filename = self.filename
+            self.updates.append(update)
-        with open(filename, 'wb') as f:
-            self.write(f)
 
-    def write(self, f):
+        self[geometry] = len(self.updates) - 1
-        f.write(struct.pack('<BHHHHH', self.resolution, self.x, self.y, *reversed(self.data.shape),
-                            len(self.updates)))
-        f.write(struct.pack('<'+'II'*len(self.updates), *chain(*self.updates)))
-        f.write(self.data.tobytes('C'))
-
-    def add_new(self, geometry, data=None):
-        logging.info('add_new called, res=%s, x=%s, y=%s, shape=%s, updates=%s' %
-                     (self.resolution, self.x, self.y, self.data.shape, self.updates))
-
-        prep = prepared.prep(geometry)
-        minx, miny, maxx, maxy = geometry.bounds
-        res = self.resolution
-        minx = int(math.floor(minx/res))
-        miny = int(math.floor(miny/res))
-        maxx = int(math.ceil(maxx/res))
-        maxy = int(math.ceil(maxy/res))
-
-        logging.info('minx=%d, miny=%d, maxx=%d, maxy=%d' % (minx, miny, maxx, maxy))
-
-        direct = data is None
-
-        if direct:
-            logging.info('direct!')
-            data = self.data
-            if self.resolution != settings.CACHE_RESOLUTION:
-                logging.info('cache_resolution does not match')
-                data = None
-                self.updates = self.updates[-1:]
-
-            if not data.size:
-                logging.info('data is empty, creating new map')
-                data = np.zeros(((maxy-miny), (maxx-minx)), dtype=np.uint16)
-                logging.info('data is empty, created new! shape=%s' % (data.shape, ))
-                self.x, self.y = minx, miny
-            else:
-                logging.info('resize?')
-                orig_height, orig_width = data.shape
-                if minx < self.x or miny < self.y or maxx > self.x+orig_width or maxy > self.y+orig_height:
-                    logging.info('resize!')
-                    new_x, new_y = min(minx, self.x), min(miny, self.y)
-                    new_width = max(maxx, self.x+orig_width)-new_x
-                    new_height = max(maxy, self.y+orig_height)-new_y
-                    new_data = np.zeros((new_height, new_width), dtype=np.uint16)
-                    dx, dy = self.x-new_x, self.y-new_y
-                    new_data[dy:(dy+orig_height), dx:(dx+orig_width)] = data
-                    data = new_data
-                    self.x, self.y = new_x, new_y
-                    logging.info('resized dx=%d, dy=%d, x=%d, y=%d, shape=%s' %
-                                 (dx, dy, self.x, self.y, data.shape))
-
-        else:
-            logging.info('not direct!')
-            height, width = data.shape
-            minx, miny = max(minx, self.x), max(miny, self.y)
-            maxx, maxy = min(maxx, self.x+width), min(maxy, self.y+height)
-
-        new_val = len(self.updates) if direct else 1
-        i = 0
-        for iy, y in enumerate(range(miny*res, maxy*res, res), start=miny-self.y):
-            for ix, x in enumerate(range(minx*res, maxx*res, res), start=minx-self.x):
-                if prep.intersects(box(x, y, x+res, y+res)):
-                    data[iy, ix] = new_val
-                    i += 1
-        logging.info('%d points changed' % i)
-
-        if direct:
-            logging.info('saved data')
-            self.data = data
-            self.unfinished = True
-        else:
-            return data
-
-    def finish(self, update):
-        self.unfinished = False
-        self.updates.append(update)
-        self.simplify()
 
     def simplify(self):
-        logging.info('simplify!')
         # remove updates that have no longer any array cells
         new_updates = ((i, update, (self.data == i)) for i, update in enumerate(self.updates))
-        logging.info('before: %s' % (self.updates, ))
         self.updates, new_affected = zip(*((update, affected) for i, update, affected in new_updates
                                            if i == 0 or affected.any()))
-        logging.info('after: %s' % (self.updates, ))
         for i, affected in enumerate(new_affected):
             self.data[affected] = i
 
-    def composite(self, other, mask_geometry):
+    def write(self, *args, **kwargs):
-        if other.resolution != other.resolution:
+        self.simplify()
-            return
+        super().write(*args, **kwargs)
 
-        # check overlapping area
+    def composite(self, other, mask_geometry):
-        self_height, self_width = self.data.shape
+        if self.resolution != other.resolution:
-        other_height, other_width = other.data.shape
+            raise ValueError('Cannot composite with different resolutions.')
-        minx, miny = max(self.x, other.x), max(self.y, other.y)
+
-        maxx = min(self.x+self_width-1, other.x+other_width-1)
+        self.fit_bounds(*other.bounds)
-        maxy = min(self.y+self_height-1, other.y+other_height-1)
+        other.fit_bounds(*self.bounds)
-        if maxx < minx or maxy < miny:
-            return
 
         # merge update lists
         self_update_i = {update: i for i, update in enumerate(self.updates)}
         other_update_i = {update: i for i, update in enumerate(other.updates)}
         new_updates = sorted(set(self_update_i.keys()) | set(other_update_i.keys()))
 
-        # create slices
+        # reindex according to merged update list
-        self_slice = slice(miny-self.y, maxy-self.y+1), slice(minx-self.x, maxx-self.x+1)
+        other_data = other.data.copy()
-        other_slice = slice(miny-other.y, maxy-other.y+1), slice(minx-other.x, maxx-other.x+1)
-
-        # reindex according to new update list
-        other_data = np.zeros_like(self.data)
-        other_data[self_slice] = other.data[other_slice]
         for i, update in enumerate(new_updates):
             if update in self_update_i:
                 self.data[self.data == self_update_i[update]] = i
@@ -221,38 +283,20 @@ class MapHistory:
 
         # add with mask
         if mask_geometry is not None:
-            mask = self.add_new(mask_geometry.buffer(1), data=np.zeros_like(self.data, dtype=np.bool))
+            mask = self.get_geometry_cells(mask_geometry)
             self.data[mask] = maximum[mask]
         else:
             self.data = maximum
 
         # write new updates
         self.updates = new_updates
-
         self.simplify()
 
-    def to_image(self):
-        from c3nav.mapdata.models import Source
-        (minx, miny), (maxx, maxy) = Source.max_bounds()
-
-        height, width = self.data.shape
-        image_data = np.zeros((int(math.ceil((maxy-miny)/self.resolution)),
-                               int(math.ceil((maxx-minx)/self.resolution))), dtype=np.uint8)
-        visible_data = (self.data.astype(float)*255/(len(self.updates)-1)).clip(0, 255).astype(np.uint8)
-        image_data[self.y:self.y+height, self.x:self.x+width] = visible_data
-
-        return Image.fromarray(np.flip(image_data, axis=0), 'L')
-
     def last_update(self, minx, miny, maxx, maxy):
-        res = self.resolution
+        cells = self[box(minx, miny, maxx, maxy)]
-        height, width = self.data.shape
+        if cells.size:
-        minx = max(int(math.floor(minx/res)), self.x)-self.x
+            return self.updates[cells.max()]
-        miny = max(int(math.floor(miny/res)), self.y)-self.y
+        return self.updates[0]
-        maxx = min(int(math.ceil(maxx/res)), self.x+width)-self.x
-        maxy = min(int(math.ceil(maxy/res)), self.y+height)-self.y
-        if minx >= maxx or miny >= maxy:
-            return self.updates[0]
-        return self.updates[self.data[miny:maxy, minx:maxx].max()]
 
 
 class GeometryChangeTracker:
@@ -299,8 +343,7 @@ class GeometryChangeTracker:
             if geometries.is_empty:
                 continue
             history = MapHistory.open_level(level_id, mode='base', default_update=last_update)
-            history.add_new(geometries.buffer(1))
+            history.add_geometry(geometries.buffer(1), new_update)
-            history.finish(new_update)
             history.save()
         self.reset()