Assigned groups, but haven't used it to overhaul color plotting yet

hail/python/hail/plot2/ggplot.py

@@ -2,13 +2,14 @@
 import plotly.graph_objects as go
 
 from pprint import pprint
+from collections import defaultdict
 
 import hail as hl
 
 from .coord_cartesian import CoordCartesian
 from .geoms import Geom, FigureAttribute
 from .labels import Labels
-from .scale import Scale, scale_x_continuous, scale_x_genomic, scale_y_continuous, scale_x_discrete, scale_y_discrete
+from .scale import Scale, ScaleContinuous, ScaleDiscrete, scale_x_continuous, scale_x_genomic, scale_y_continuous, scale_x_discrete, scale_y_discrete
 from .aes import Aesthetic, aes
 
 
@@ -61,24 +62,28 @@ def is_genomic_type(dtype):
 
         for aesthetic_str, mapped_expr in aesthetic.items():
             dtype = mapped_expr.dtype
-            if aesthetic_str in self.scales:
-                pass
-            else:
+            if aesthetic_str not in self.scales:
+                is_continuous = is_continuous_type(dtype)
                 # We only know how to come up with a few default scales.
                 if aesthetic_str == "x":
-                    if is_continuous_type(dtype):
+                    if is_continuous:
                         self.scales["x"] = scale_x_continuous()
                     elif is_genomic_type(dtype):
                         self.scales["x"] = scale_x_genomic(reference_genome=dtype.reference_genome)
                     else:
                         self.scales["x"] = scale_x_discrete()
                 elif aesthetic_str == "y":
-                    if is_continuous_type(dtype):
+                    if is_continuous:
                         self.scales["y"] = scale_y_continuous()
                     elif is_genomic_type(dtype):
                         raise ValueError("Don't yet support y axis genomic")
                     else:
                         self.scales["y"] = scale_y_discrete()
+                else:
+                    if is_continuous:
+                        self.scales[aesthetic_str] = ScaleContinuous(aesthetic_str)
+                    else:
+                        self.scales[aesthetic_str] = ScaleDiscrete(aesthetic_str)
 
     def copy(self):
         return GGPlot(self.ht, self.aes, self.geoms[:], self.labels, self.coord_cartesian, self.scales,
@@ -114,6 +119,25 @@ def render(self):
 
         for geom, (label, agg_result) in zip(self.geoms, aggregated.items()):
             listified_agg_result = labels_to_stats[label].listify(agg_result)
+
+            # Ok, need to identify every possible value of every discrete scale.
+            if listified_agg_result:
+                relevant_aesthetics = [scale_name for scale_name in list(listified_agg_result[0]) if scale_name in self.scales and self.scales[scale_name].is_discrete()]
+                subsetted_to_relevant = tuple([one_struct.select(*relevant_aesthetics) for one_struct in listified_agg_result])
+                group_counter = 0
+
+                def increment_and_return_old():
+                    nonlocal group_counter
+                    group_counter += 1
+                    return group_counter - 1
+                numberer = defaultdict(increment_and_return_old)
+                numbering = [numberer[data_tuple] for data_tuple in subsetted_to_relevant]
+
+                numbered_result = []
+                for s, i in zip(listified_agg_result, numbering):
+                    numbered_result.append(s.annotate(group=i))
+                listified_agg_result = numbered_result
+
             geom.apply_to_fig(self, listified_agg_result, fig)
 
         # Important to update axes after labels, axes names take precedence.

hail/python/hail/plot2/scale.py

@@ -12,6 +12,10 @@ def __init__(self, aesthetic_name):
     def transform_data(self, field_expr):
         pass
 
+    @abc.abstractmethod
+    def is_discrete(self):
+        pass
+
 
 class PositionScale(Scale):
     def __init__(self, aesthetic_name, name, breaks, labels):
@@ -55,6 +59,9 @@ def apply_to_fig(self, parent, fig_so_far):
     def transform_data(self, field_expr):
         return field_expr.global_position()
 
+    def is_discrete(self):
+        return False
+
 
 class PositionScaleContinuous(PositionScale):
 
@@ -76,6 +83,9 @@ def apply_to_fig(self, parent, fig_so_far):
     def transform_data(self, field_expr):
         return field_expr
 
+    def is_discrete(self):
+        return False
+
 
 class PositionScaleDiscrete(PositionScale):
     def __init__(self, axis=None, name=None, breaks=None, labels=None):
@@ -87,6 +97,31 @@ def apply_to_fig(self, parent, fig_so_far):
     def transform_data(self, field_expr):
         return field_expr
 
+    def is_discrete(self):
+        return True
+
+
+class ScaleContinuous(Scale):
+    def __init__(self, aesthetic_name):
+        super().__init__(aesthetic_name)
+
+    def transform_data(self, field_expr):
+        return field_expr
+
+    def is_discrete(self):
+        return False
+
+
+class ScaleDiscrete(Scale):
+    def __init__(self, aesthetic_name):
+        super().__init__(aesthetic_name)
+
+    def transform_data(self, field_expr):
+        return field_expr
+
+    def is_discrete(self):
+        return True
+
 
 def scale_x_log10():
     return PositionScaleContinuous("x", transformation="log10")