Merge pull request #60 from mwcraig/add-marker-sets

Add ability to name markers

astrowidgets/core.py

@@ -8,7 +8,7 @@
 import numpy as np
 from astropy.coordinates import SkyCoord
 from astropy.io import fits
-from astropy.table import Table
+from astropy.table import Table, vstack
 
 # Jupyter widgets
 import ipywidgets as ipyw
@@ -24,6 +24,9 @@
 # Allowed locations for cursor display
 ALLOWED_CURSOR_LOCATIONS = ['top', 'bottom', None]
 
+# List of marker names that are for internal use only
+RESERVED_MARKER_SET_NAMES = ['all']
+
 
 class ImageWidget(ipyw.VBox):
     """
@@ -111,7 +114,12 @@ def __init__(self, logger=None, image_width=500, image_height=500,
 
         # Marker
         self.marker = {'type': 'circle', 'color': 'cyan', 'radius': 20}
-        self._marktag = 'marktag'
+        # Maintain marker tags as a set because we do not want
+        # duplicate names.
+        self._marktags = set()
+        # Let's have a default name for the tag too:
+        self._default_mark_tag_name = 'default-marker-name'
+        self._interactive_marker_set_name = 'interactive-markers'
 
         # coordinates display
         self._jup_coord = ipyw.HTML('Coordinates show up here')
@@ -195,21 +203,23 @@ def _mouse_click_cb(self, viewer, event, data_x, data_y):
         Callback to handle mouse clicks.
         """
         if self.is_marking:
+            marker_name = self._interactive_marker_set_name
             objs = []
             try:
-                c_mark = viewer.canvas.get_object_by_tag(self._marktag)
+                c_mark = viewer.canvas.get_object_by_tag(marker_name)
             except Exception:  # Nothing drawn yet
                 pass
             else:  # Add to existing marks
                 objs = c_mark.objects
-                viewer.canvas.delete_object_by_tag(self._marktag)
+                viewer.canvas.delete_object_by_tag(marker_name)
 
             # NOTE: By always using CompoundObject, marker handling logic
             # is simplified.
             obj = self.marker(x=data_x, y=data_y)
             objs.append(obj)
-            self._marktag = viewer.canvas.add(self.dc.CompoundObject(*objs))
-
+            viewer.canvas.add(self.dc.CompoundObject(*objs),
+                              tag=marker_name)
+            self._marktags.add(marker_name)
             with self.print_out:
                 print('Selected {} {}'.format(obj.x, obj.y))
 
@@ -421,7 +431,8 @@ def marker(self, val):
                 'Marker type "{}" not supported'.format(marker_type))
 
     def get_markers(self, x_colname='x', y_colname='y',
-                    skycoord_colname='coord'):
+                    skycoord_colname='coord',
+                    marker_name=None):
         """
         Return the locations of existing markers.
 
@@ -443,16 +454,48 @@ def get_markers(self, x_colname='x', y_colname='y',
             Table of markers, if any, or ``None``.
 
         """
+        if marker_name is None:
+            marker_name = self._default_mark_tag_name
+
+        if marker_name == 'all':
+            # If it wasn't for the fact that SKyCoord columns can't
+            # be stacked this would all fit nicely into a list
+            # comprehension. But they can't, so we delete the
+            # SkyCoord column if it is present, then add it
+            # back after we have stacked.
+            coordinates = []
+            tables = []
+            for name in self._marktags:
+                table = self.get_markers(x_colname=x_colname,
+                                         y_colname=y_colname,
+                                         skycoord_colname=skycoord_colname,
+                                         marker_name=name)
+                try:
+                    coordinates.extend(c for c in table[skycoord_colname])
+                except KeyError:
+                    pass
+                else:
+                    del table[skycoord_colname]
+                tables.append(table)
+
+            stacked = vstack(tables, join_type='exact')
+
+            if coordinates:
+                stacked[skycoord_colname] = SkyCoord(coordinates)
+
+            return stacked
+
         try:
-            c_mark = self._viewer.canvas.get_object_by_tag(self._marktag)
+            c_mark = self._viewer.canvas.get_object_by_tag(marker_name)
         except Exception as e:  # No markers
             self.logger.warning(str(e))
             return
 
         image = self._viewer.get_image()
         xy_col = []
 
-        if image.wcs.wcs is None:  # Do not include SkyCoord column
+        if (image is None) or (image.wcs.wcs is None):
+            # Do not include SkyCoord column
             include_skycoord = False
         else:
             include_skycoord = True
@@ -500,12 +543,13 @@ def get_markers(self, x_colname='x', y_colname='y',
                 [xy_col[:, 0], xy_col[:, 1], sky_col],
                 names=(x_colname, y_colname, skycoord_colname))
         else:
-            markers_table = Table(xy_col.T, names=(x_colname, y_colname))
+            markers_table = Table(xy_col, names=(x_colname, y_colname))
 
         return markers_table
 
     def add_markers(self, table, x_colname='x', y_colname='y',
-                    skycoord_colname='coord', use_skycoord=False):
+                    skycoord_colname='coord', use_skycoord=False,
+                    marker_name=None):
         """
         Creates markers in the image at given points.
 
@@ -531,13 +575,27 @@ def add_markers(self, table, x_colname='x', y_colname='y',
             If `True`, use ``skycoord_colname`` to mark.
             Otherwise, use ``x_colname`` and ``y_colname``.
 
+        marker_name : str, optional
+            Name to assign the markers in the table. Providing a name
+            allows markers to be removed by name at a later time.
         """
         # TODO: Resolve https://github.com/ejeschke/ginga/issues/672
 
         # For now we always convert marker locations to pixels; see
         # comment below.
         coord_type = 'data'
 
+        if marker_name is None:
+            marker_name = self._default_mark_tag_name
+
+        if marker_name in RESERVED_MARKER_SET_NAMES:
+            raise ValueError('The marker name {} is not allowed. Any name is '
+                             'allowed except these: '
+                             '{}'.format(marker_name,
+                                         ', '.join(RESERVED_MARKER_SET_NAMES)))
+
+        self._marktags.add(marker_name)
+
         # Extract coordinates from table.
         # They are always arrays, not scalar.
         if use_skycoord:
@@ -569,41 +627,64 @@ def add_markers(self, table, x_colname='x', y_colname='y',
         # Prepare canvas and retain existing marks
         objs = []
         try:
-            c_mark = self._viewer.canvas.get_object_by_tag(self._marktag)
+            c_mark = self._viewer.canvas.get_object_by_tag(marker_name)
         except Exception:
             pass
         else:
             objs = c_mark.objects
-            self._viewer.canvas.delete_object_by_tag(self._marktag)
+            self._viewer.canvas.delete_object_by_tag(marker_name)
 
         # TODO: Test to see if we can mix WCS and data on the same canvas
         objs += [self.marker(x=x, y=y, coord=coord_type)
                  for x, y in zip(coord_x, coord_y)]
-        self._marktag = self._viewer.canvas.add(self.dc.CompoundObject(*objs))
+        self._viewer.canvas.add(self.dc.CompoundObject(*objs),
+                                tag=marker_name)
 
-    # TODO: Future work?
-    def remove_markers(self, arr):
+    def remove_markers(self, marker_name=None):
         """
-        Remove some but not all of the markers.
+        Remove some but not all of the markers by name used when
+        adding the markers
 
         Parameters
         ----------
-        arr : ``SkyCoord`` or array-like
-            Sky coordinates or 2xN array.
 
+        marker_name : str, optional
+            Name used when the markers were added.
         """
+        # TODO:
+        #   arr : ``SkyCoord`` or array-like
+        #   Sky coordinates or 2xN array.
+        #
         # NOTE: How to match? Use np.isclose?
         #       What if there are 1-to-many matches?
-        return NotImplementedError
+
+        if marker_name is None:
+            marker_name = self._default_mark_tag_name
+
+        if marker_name not in self._marktags:
+            # This shouldn't have happened, raise an error
+            raise ValueError('Marker name {} not found in current markers.'
+                             ' Markers currently in use are '
+                             '{}'.format(marker_name,
+                                         sorted(self._marktags)))
+
+        try:
+            self._viewer.canvas.delete_object_by_tag(marker_name)
+        except KeyError:
+            raise KeyError('Unable to remove markers named {} from image. '
+                           ''.format(marker_name))
+        else:
+            self._marktags.remove(marker_name)
 
     def reset_markers(self):
         """
         Delete all markers.
         """
-        try:
-            self._viewer.canvas.delete_object_by_tag(self._marktag)
-        except Exception:
-            pass
+
+        # Grab the entire list of marker names before iterating
+        # otherwise what we are iterating over changes.
+        for marker_name in list(self._marktags):
+            self.remove_markers(marker_name)
 
     @property
     def stretch_options(self):

astrowidgets/tests/test_api.py

@@ -104,6 +104,15 @@ def test_reset_markers():
     # TODO add test of actually removing markers...
 
 
+def test_remove_markers():
+    image = ImageWidget()
+    # Add a tag name...
+    image._marktags.add(image._default_mark_tag_name)
+    with pytest.raises(ValueError) as e:
+        image.remove_markers('arf')
+    assert 'arf' in str(e)
+
+
 def test_stretch():
     image = ImageWidget()
     with pytest.raises(ValueError) as e: