1"""
2Colormap
3--------
4
5Utility module for dealing with colormaps.
6
7"""
8
9import json
10import math
11import os
12
13from jinja2 import Template
14
15from branca.element import ENV, Figure, JavascriptLink, MacroElement
16from branca.utilities import legend_scaler
17
18rootpath = os.path.abspath(os.path.dirname(__file__))
19
20with open(os.path.join(rootpath, '_cnames.json')) as f:
21    _cnames = json.loads(f.read())
22
23with open(os.path.join(rootpath, '_schemes.json')) as f:
24    _schemes = json.loads(f.read())
25
26
27def _is_hex(x):
28    return x.startswith('#') and len(x) == 7
29
30
31def _parse_hex(color_code):
32    return (int(color_code[1:3], 16),
33            int(color_code[3:5], 16),
34            int(color_code[5:7], 16))
35
36
37def _parse_color(x):
38    if isinstance(x, (tuple, list)):
39        color_tuple = tuple(x)[:4]
40    elif isinstance(x, (str, bytes)) and _is_hex(x):
41        color_tuple = _parse_hex(x)
42    elif isinstance(x, (str, bytes)):
43        cname = _cnames.get(x.lower(), None)
44        if cname is None:
45            raise ValueError('Unknown color {!r}.'.format(cname))
46        color_tuple = _parse_hex(cname)
47    else:
48        raise ValueError('Unrecognized color code {!r}'.format(x))
49    if max(color_tuple) > 1.:
50        color_tuple = tuple(u/255. for u in color_tuple)
51    return tuple(map(float, (color_tuple+(1.,))[:4]))
52
53
54def _base(x):
55    if x > 0:
56        base = pow(10, math.floor(math.log10(x)))
57        return round(x/base)*base
58    else:
59        return 0
60
61
62class ColorMap(MacroElement):
63    """A generic class for creating colormaps.
64
65    Parameters
66    ----------
67    vmin: float
68        The left bound of the color scale.
69    vmax: float
70        The right bound of the color scale.
71    caption: str
72        A caption to draw with the colormap.
73    """
74    _template = ENV.get_template('color_scale.js')
75
76    def __init__(self, vmin=0., vmax=1., caption=''):
77        super(ColorMap, self).__init__()
78        self._name = 'ColorMap'
79
80        self.vmin = vmin
81        self.vmax = vmax
82        self.caption = caption
83        self.index = [vmin, vmax]
84
85    def render(self, **kwargs):
86        """Renders the HTML representation of the element."""
87        self.color_domain = [self.vmin + (self.vmax-self.vmin) * k/499. for
88                             k in range(500)]
89        self.color_range = [self.__call__(x) for x in self.color_domain]
90        self.tick_labels = legend_scaler(self.index)
91
92        super(ColorMap, self).render(**kwargs)
93
94        figure = self.get_root()
95        assert isinstance(figure, Figure), ('You cannot render this Element '
96                                            'if it is not in a Figure.')
97
98        figure.header.add_child(JavascriptLink("https://cdnjs.cloudflare.com/ajax/libs/d3/3.5.5/d3.min.js"), name='d3')  # noqa
99
100    def rgba_floats_tuple(self, x):
101        """
102        This class has to be implemented for each class inheriting from
103        Colormap. This has to be a function of the form float ->
104        (float, float, float, float) describing for each input float x,
105        the output color in RGBA format;
106        Each output value being between 0 and 1.
107        """
108        raise NotImplementedError
109
110    def rgba_bytes_tuple(self, x):
111        """Provides the color corresponding to value `x` in the
112        form of a tuple (R,G,B,A) with int values between 0 and 255.
113        """
114        return tuple(int(u*255.9999) for u in self.rgba_floats_tuple(x))
115
116    def rgb_bytes_tuple(self, x):
117        """Provides the color corresponding to value `x` in the
118        form of a tuple (R,G,B) with int values between 0 and 255.
119        """
120        return self.rgba_bytes_tuple(x)[:3]
121
122    def rgb_hex_str(self, x):
123        """Provides the color corresponding to value `x` in the
124        form of a string of hexadecimal values "#RRGGBB".
125        """
126        return '#%02x%02x%02x' % self.rgb_bytes_tuple(x)
127
128    def rgba_hex_str(self, x):
129        """Provides the color corresponding to value `x` in the
130        form of a string of hexadecimal values "#RRGGBBAA".
131        """
132        return '#%02x%02x%02x%02x' % self.rgba_bytes_tuple(x)
133
134    def __call__(self, x):
135        """Provides the color corresponding to value `x` in the
136        form of a string of hexadecimal values "#RRGGBBAA".
137        """
138        return self.rgba_hex_str(x)
139
140    def _repr_html_(self):
141        return (
142            '<svg height="50" width="500">' +
143            ''.join(
144                [('<line x1="{i}" y1="0" x2="{i}" '
145                  'y2="20" style="stroke:{color};stroke-width:3;" />').format(
146                      i=i*1,
147                      color=self.rgba_hex_str(
148                          self.vmin +
149                          (self.vmax-self.vmin)*i/499.)
150                  )
151                 for i in range(500)]) +
152            '<text x="0" y="35">{}</text>'.format(self.vmin) +
153            '<text x="500" y="35" style="text-anchor:end;">{}</text>'.format(
154                self.vmax) +
155            '</svg>')
156
157
158class LinearColormap(ColorMap):
159    """Creates a ColorMap based on linear interpolation of a set of colors
160    over a given index.
161
162    Parameters
163    ----------
164
165    colors : list-like object with at least two colors.
166        The set of colors to be used for interpolation.
167        Colors can be provided in the form:
168        * tuples of RGBA ints between 0 and 255 (e.g: `(255, 255, 0)` or
169        `(255, 255, 0, 255)`)
170        * tuples of RGBA floats between 0. and 1. (e.g: `(1.,1.,0.)` or
171        `(1., 1., 0., 1.)`)
172        * HTML-like string (e.g: `"#ffff00`)
173        * a color name or shortcut (e.g: `"y"` or `"yellow"`)
174    index : list of floats, default None
175        The values corresponding to each color.
176        It has to be sorted, and have the same length as `colors`.
177        If None, a regular grid between `vmin` and `vmax` is created.
178    vmin : float, default 0.
179        The minimal value for the colormap.
180        Values lower than `vmin` will be bound directly to `colors[0]`.
181    vmax : float, default 1.
182        The maximal value for the colormap.
183        Values higher than `vmax` will be bound directly to `colors[-1]`."""
184
185    def __init__(self, colors, index=None, vmin=0., vmax=1., caption=''):
186        super(LinearColormap, self).__init__(vmin=vmin, vmax=vmax,
187                                             caption=caption)
188
189        n = len(colors)
190        if n < 2:
191            raise ValueError('You must provide at least 2 colors.')
192        if index is None:
193            self.index = [vmin + (vmax-vmin)*i*1./(n-1) for i in range(n)]
194        else:
195            self.index = list(index)
196        self.colors = [_parse_color(x) for x in colors]
197
198    def rgba_floats_tuple(self, x):
199        """Provides the color corresponding to value `x` in the
200        form of a tuple (R,G,B,A) with float values between 0. and 1.
201        """
202        if x <= self.index[0]:
203            return self.colors[0]
204        if x >= self.index[-1]:
205            return self.colors[-1]
206
207        i = len([u for u in self.index if u < x])  # 0 < i < n.
208        if self.index[i-1] < self.index[i]:
209            p = (x - self.index[i-1])*1./(self.index[i]-self.index[i-1])
210        elif self.index[i-1] == self.index[i]:
211            p = 1.
212        else:
213            raise ValueError('Thresholds are not sorted.')
214
215        return tuple((1.-p) * self.colors[i-1][j] + p*self.colors[i][j] for j
216                     in range(4))
217
218    def to_step(self, n=None, index=None, data=None, method=None,
219                quantiles=None, round_method=None):
220        """Splits the LinearColormap into a StepColormap.
221
222        Parameters
223        ----------
224        n : int, default None
225            The number of expected colors in the ouput StepColormap.
226            This will be ignored if `index` is provided.
227        index : list of floats, default None
228            The values corresponding to each color bounds.
229            It has to be sorted.
230            If None, a regular grid between `vmin` and `vmax` is created.
231        data : list of floats, default None
232            A sample of data to adapt the color map to.
233        method : str, default 'linear'
234            The method used to create data-based colormap.
235            It can be 'linear' for linear scale, 'log' for logarithmic,
236            or 'quant' for data's quantile-based scale.
237        quantiles : list of floats, default None
238            Alternatively, you can provide explicitely the quantiles you
239            want to use in the scale.
240        round_method : str, default None
241            The method used to round thresholds.
242            * If 'int', all values will be rounded to the nearest integer.
243            * If 'log10', all values will be rounded to the nearest
244            order-of-magnitude integer. For example, 2100 is rounded to
245            2000, 2790 to 3000.
246
247        Returns
248        -------
249        A StepColormap with `n=len(index)-1` colors.
250
251        Examples:
252        >> lc.to_step(n=12)
253        >> lc.to_step(index=[0, 2, 4, 6, 8, 10])
254        >> lc.to_step(data=some_list, n=12)
255        >> lc.to_step(data=some_list, n=12, method='linear')
256        >> lc.to_step(data=some_list, n=12, method='log')
257        >> lc.to_step(data=some_list, n=12, method='quantiles')
258        >> lc.to_step(data=some_list, quantiles=[0, 0.3, 0.7, 1])
259        >> lc.to_step(data=some_list, quantiles=[0, 0.3, 0.7, 1],
260        ...           round_method='log10')
261
262        """
263        msg = 'You must specify either `index` or `n`'
264        if index is None:
265            if data is None:
266                if n is None:
267                    raise ValueError(msg)
268                else:
269                    index = [self.vmin + (self.vmax-self.vmin)*i*1./n for
270                             i in range(1+n)]
271                    scaled_cm = self
272            else:
273                max_ = max(data)
274                min_ = min(data)
275                scaled_cm = self.scale(vmin=min_, vmax=max_)
276                method = ('quantiles' if quantiles is not None
277                          else method if method is not None
278                          else 'linear'
279                          )
280                if method.lower().startswith('lin'):
281                    if n is None:
282                        raise ValueError(msg)
283                    index = [min_ + i*(max_-min_)*1./n for i in range(1+n)]
284                elif method.lower().startswith('log'):
285                    if n is None:
286                        raise ValueError(msg)
287                    if min_ <= 0:
288                        msg = ('Log-scale works only with strictly '
289                               'positive values.')
290                        raise ValueError(msg)
291                    index = [math.exp(
292                        math.log(min_) + i * (math.log(max_) -
293                                              math.log(min_)) * 1./n
294                    ) for i in range(1+n)]
295                elif method.lower().startswith('quant'):
296                    if quantiles is None:
297                        if n is None:
298                            msg = ('You must specify either `index`, `n` or'
299                                   '`quantiles`.')
300                            raise ValueError(msg)
301                        else:
302                            quantiles = [i*1./n for i in range(1+n)]
303                    p = len(data)-1
304                    s = sorted(data)
305                    index = [s[int(q*p)] * (1.-(q*p) % 1) +
306                             s[min(int(q*p) + 1, p)] * ((q*p) % 1) for
307                             q in quantiles]
308                else:
309                    raise ValueError('Unknown method {}'.format(method))
310        else:
311            scaled_cm = self.scale(vmin=min(index), vmax=max(index))
312
313        n = len(index)-1
314
315        if round_method == 'int':
316            index = [round(x) for x in index]
317
318        if round_method == 'log10':
319            index = [_base(x) for x in index]
320
321        colors = [scaled_cm.rgba_floats_tuple(index[i] * (1.-i/(n-1.)) +
322                                              index[i+1] * i/(n-1.)) for
323                  i in range(n)]
324
325        return StepColormap(colors, index=index, vmin=index[0], vmax=index[-1])
326
327    def scale(self, vmin=0., vmax=1.):
328        """Transforms the colorscale so that the minimal and maximal values
329        fit the given parameters.
330        """
331        return LinearColormap(
332            self.colors,
333            index=[vmin + (vmax-vmin)*(x-self.vmin)*1./(self.vmax-self.vmin) for x in self.index],  # noqa
334            vmin=vmin,
335            vmax=vmax,
336            caption=self.caption,
337            )
338
339
340class StepColormap(ColorMap):
341    """Creates a ColorMap based on linear interpolation of a set of colors
342    over a given index.
343
344    Parameters
345    ----------
346    colors : list-like object
347        The set of colors to be used for interpolation.
348        Colors can be provided in the form:
349        * tuples of int between 0 and 255 (e.g: `(255,255,0)` or
350        `(255, 255, 0, 255)`)
351        * tuples of floats between 0. and 1. (e.g: `(1.,1.,0.)` or
352        `(1., 1., 0., 1.)`)
353        * HTML-like string (e.g: `"#ffff00`)
354        * a color name or shortcut (e.g: `"y"` or `"yellow"`)
355    index : list of floats, default None
356        The values corresponding to each color.
357        It has to be sorted, and have the same length as `colors`.
358        If None, a regular grid between `vmin` and `vmax` is created.
359    vmin : float, default 0.
360        The minimal value for the colormap.
361        Values lower than `vmin` will be bound directly to `colors[0]`.
362    vmax : float, default 1.
363        The maximal value for the colormap.
364        Values higher than `vmax` will be bound directly to `colors[-1]`.
365
366    """
367    def __init__(self, colors, index=None, vmin=0., vmax=1., caption=''):
368        super(StepColormap, self).__init__(vmin=vmin, vmax=vmax,
369                                           caption=caption)
370
371        n = len(colors)
372        if n < 1:
373            raise ValueError('You must provide at least 1 colors.')
374        if index is None:
375            self.index = [vmin + (vmax-vmin)*i*1./n for i in range(n+1)]
376        else:
377            self.index = list(index)
378        self.colors = [_parse_color(x) for x in colors]
379
380    def rgba_floats_tuple(self, x):
381        """
382        Provides the color corresponding to value `x` in the
383        form of a tuple (R,G,B,A) with float values between 0. and 1.
384
385        """
386        if x <= self.index[0]:
387            return self.colors[0]
388        if x >= self.index[-1]:
389            return self.colors[-1]
390
391        i = len([u for u in self.index if u < x])  # 0 < i < n.
392        return tuple(self.colors[i-1])
393
394    def to_linear(self, index=None):
395        """
396        Transforms the StepColormap into a LinearColormap.
397
398        Parameters
399        ----------
400        index : list of floats, default None
401                The values corresponding to each color in the output colormap.
402                It has to be sorted.
403                If None, a regular grid between `vmin` and `vmax` is created.
404
405        """
406        if index is None:
407            n = len(self.index)-1
408            index = [self.index[i]*(1.-i/(n-1.))+self.index[i+1]*i/(n-1.) for
409                     i in range(n)]
410
411        colors = [self.rgba_floats_tuple(x) for x in index]
412        return LinearColormap(colors, index=index,
413                              vmin=self.vmin, vmax=self.vmax)
414
415    def scale(self, vmin=0., vmax=1.):
416        """Transforms the colorscale so that the minimal and maximal values
417        fit the given parameters.
418        """
419        return StepColormap(
420            self.colors,
421            index=[vmin + (vmax-vmin)*(x-self.vmin)*1./(self.vmax-self.vmin) for x in self.index],  # noqa
422            vmin=vmin,
423            vmax=vmax,
424            caption=self.caption,
425            )
426
427
428class _LinearColormaps(object):
429    """A class for hosting the list of built-in linear colormaps."""
430    def __init__(self):
431        self._schemes = _schemes.copy()
432        self._colormaps = {key: LinearColormap(val) for
433                           key, val in _schemes.items()}
434        for key, val in _schemes.items():
435            setattr(self, key, LinearColormap(val))
436
437    def _repr_html_(self):
438        return Template("""
439        <table>
440        {% for key,val in this._colormaps.items() %}
441        <tr><td>{{key}}</td><td>{{val._repr_html_()}}</td></tr>
442        {% endfor %}</table>
443        """).render(this=self)
444
445
446linear = _LinearColormaps()
447
448
449class _StepColormaps(object):
450    """A class for hosting the list of built-in step colormaps."""
451    def __init__(self):
452        self._schemes = _schemes.copy()
453        self._colormaps = {key: StepColormap(val) for
454                           key, val in _schemes.items()}
455        for key, val in _schemes.items():
456            setattr(self, key, StepColormap(val))
457
458    def _repr_html_(self):
459        return Template("""
460        <table>
461        {% for key,val in this._colormaps.items() %}
462        <tr><td>{{key}}</td><td>{{val._repr_html_()}}</td></tr>
463        {% endfor %}</table>
464        """).render(this=self)
465
466
467step = _StepColormaps()
468