1#    Copyright (C) 2005 Jeremy S. Sanders
2#    Email: Jeremy Sanders <jeremy@jeremysanders.net>
3#
4#    This program is free software; you can redistribute it and/or modify
5#    it under the terms of the GNU General Public License as published by
6#    the Free Software Foundation; either version 2 of the License, or
7#    (at your option) any later version.
8#
9#    This program is distributed in the hope that it will be useful,
10#    but WITHOUT ANY WARRANTY; without even the implied warranty of
11#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
12#    GNU General Public License for more details.
13#
14#    You should have received a copy of the GNU General Public License along
15#    with this program; if not, write to the Free Software Foundation, Inc.,
16#    51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
17###############################################################################
18
19"""Contour plotting from 2d datasets.
20
21Contour plotting requires that the veusz_helpers package is installed,
22as a C routine (taken from matplotlib) is used to trace the contours.
23"""
24
25from __future__ import division, print_function
26import sys
27import math
28
29from ..compat import czip, crange
30from .. import qtall as qt
31import numpy as N
32
33from .. import setting
34from .. import document
35from .. import utils
36
37from . import plotters
38
39try:
40    from ..helpers._nc_cntr import Cntr
41    from ..helpers.qtloops import LineLabeller
42except ImportError:
43    Cntr = None
44    LineLabeller = object   # allow class definition below
45
46def _(text, disambiguation=None, context='Contour'):
47    """Translate text."""
48    return qt.QCoreApplication.translate(context, text, disambiguation)
49
50def finitePoly(poly):
51    """Remove non-finite coordinates from numpy arrays of coordinates."""
52    out = []
53    for line in poly:
54        finite = N.isfinite(line)
55        validrows = N.logical_and(finite[:,0], finite[:,1])
56        out.append( line[validrows] )
57    return out
58
59class ContourLineLabeller(LineLabeller):
60    def __init__(self, clip, rot, painter, font, doc):
61        LineLabeller.__init__(self, clip, rot)
62        self.clippath = qt.QPainterPath()
63        self.clippath.addRect(clip)
64        self.labels = []
65        self.painter = painter
66        self.font = font
67        self.document = doc
68
69    def drawAt(self, idx, rect):
70        """Called to draw the label with the index given."""
71        text = self.labels[idx]
72        if not text:
73            return
74
75        angle = rect.angle*180/math.pi
76        if angle < -90 or angle > 90:
77            angle += 180
78
79        rend = utils.Renderer(
80            self.painter, self.font,
81            rect.cx, rect.cy, text,
82            alignhorz=0, alignvert=0,
83            angle=angle,
84            doc=self.document)
85
86        rend.render()
87        if rect.xw > 0:
88            p = qt.QPainterPath()
89            p.addPolygon(rect.makePolygon())
90            self.clippath -= p
91
92class ContourFills(setting.Settings):
93    """Settings for contour fills."""
94    def __init__(self, name, **args):
95        setting.Settings.__init__(self, name, **args)
96        self.add( setting.FillSet(
97            'fills', [],
98            descr = _('Fill styles to plot between contours'),
99            usertext=_('Fill styles'),
100            formatting=True) )
101        self.add( setting.Bool(
102            'hide', False,
103            descr = _('Hide fills'),
104            usertext = _('Hide'),
105            formatting = True) )
106
107class ContourLines(setting.Settings):
108    """Settings for contour lines."""
109    def __init__(self, name, **args):
110        setting.Settings.__init__(self, name, **args)
111        self.add( setting.LineSet(
112            'lines',
113            [('solid', '1pt', 'black', False)],
114            descr = _('Line styles to plot the contours '
115                      'using'), usertext=_('Line styles'),
116            formatting=True) )
117        self.add( setting.Bool(
118            'hide', False,
119            descr = _('Hide lines'),
120            usertext = _('Hide'),
121            formatting = True) )
122
123class SubContourLines(setting.Settings):
124    """Sub-dividing contour line settings."""
125    def __init__(self, name, **args):
126        setting.Settings.__init__(self, name, **args)
127        self.add( setting.LineSet(
128            'lines',
129            [('dot1', '1pt', 'black', False)],
130            descr = _('Line styles used for sub-contours'),
131            usertext=_('Line styles'),
132            formatting=True) )
133        self.add( setting.Int(
134            'numLevels', 5,
135            minval=2,
136            descr=_('Number of sub-levels to plot between '
137                    'each contour'),
138            usertext='Levels') )
139        self.add( setting.Bool(
140            'hide', True,
141            descr=_('Hide lines'),
142            usertext=_('Hide'),
143            formatting=True) )
144
145class ContourLabel(setting.Text):
146    """For tick labels on axes."""
147
148    def __init__(self, name, **args):
149        setting.Text.__init__(self, name, **args)
150        self.add( setting.Str(
151            'format', '%.3Vg',
152            descr = _('Format of the tick labels'),
153            usertext=_('Format')) )
154        self.add( setting.Float(
155            'scale', 1.,
156            descr=_('A scale factor to apply to the values '
157                    'of the tick labels'),
158            usertext=_('Scale')) )
159        self.add( setting.Bool(
160            'rotate',
161            True,
162            descr=_('Rotate labels to follow lines'),
163            usertext=_('Rotate')) )
164
165        self.get('hide').newDefault(True)
166
167class Contour(plotters.GenericPlotter):
168    """A class which plots contours on a graph with a specified
169    coordinate system."""
170
171    typename='contour'
172    allowusercreation=True
173    description=_('Plot a 2d dataset as contours')
174
175    def __init__(self, parent, name=None):
176        """Initialise plotter with axes."""
177
178        plotters.GenericPlotter.__init__(self, parent, name=name)
179
180        if Cntr is None:
181            print(('WARNING: Veusz cannot import contour module\n'
182                   'Please run python setup.py build\n'
183                   'Contour support is disabled'), file=sys.stderr)
184
185        # keep track of settings so we recalculate when necessary
186        self.contsettings = None
187
188        # cached traced contours
189        self._cachedcontours = None
190        self._cachedpolygons = None
191        self._cachedsubcontours = None
192
193    @classmethod
194    def addSettings(klass, s):
195        """Construct list of settings."""
196        plotters.GenericPlotter.addSettings(s)
197
198        s.add( setting.DatasetExtended(
199            'data', '',
200            dimensions = 2,
201            descr = _('Dataset to plot'),
202            usertext=_('Dataset')),
203               0 )
204        s.add( setting.FloatOrAuto(
205            'min', 'Auto',
206            descr = _('Minimum value of contour scale'),
207            usertext=_('Min. value')),
208               1 )
209        s.add( setting.FloatOrAuto(
210            'max', 'Auto',
211            descr = _('Maximum value of contour scale'),
212            usertext=_('Max. value')),
213               2 )
214        s.add( setting.Int(
215            'numLevels', 5,
216            minval = 1,
217            descr = _('Number of contour levels to plot'),
218            usertext=_('Number levels')),
219               3 )
220        s.add( setting.Choice(
221            'scaling',
222            ['linear', 'sqrt', 'log', 'squared', 'manual'],
223            'linear',
224            descr = _('Scaling between contour levels'),
225            usertext=_('Scaling')),
226               4 )
227        s.add( setting.FloatList(
228            'manualLevels',
229            [],
230            descr = _('Levels to use for manual scaling'),
231            usertext=_('Manual levels')),
232               5 )
233
234        s.add( setting.Bool(
235            'keyLevels',
236            False,
237            descr=_('Show levels in key'),
238            usertext=_('Levels in key')),
239               6 )
240
241        s.add( setting.FloatList(
242            'levelsOut',
243            [],
244            descr = _('Levels used in the plot'),
245            usertext=_('Output levels')),
246               7, readonly=True )
247
248        s.add( ContourLabel(
249            'ContourLabels',
250            descr = _('Contour label settings'),
251            usertext = _('Contour labels')),
252               pixmap = 'settings_axisticklabels' )
253
254        s.add( ContourLines(
255            'Lines',
256            descr=_('Contour lines'),
257            usertext=_('Contour lines')),
258               pixmap = 'settings_contourline' )
259
260        s.add( ContourFills(
261            'Fills',
262            descr=_('Fill within contours'),
263            usertext=_('Contour fills')),
264               pixmap = 'settings_contourfill' )
265
266        s.add( SubContourLines(
267            'SubLines',
268            descr=_('Sub-contour lines'),
269            usertext=_('Sub-contour lines')),
270               pixmap = 'settings_subcontourline' )
271
272        s.add( setting.SettingBackwardCompat('lines', 'Lines/lines', None) )
273        s.add( setting.SettingBackwardCompat('fills', 'Fills/fills', None) )
274
275        s.remove('key')
276
277    @property
278    def userdescription(self):
279        """User friendly description."""
280        s = self.settings
281        out = []
282        if s.data:
283            out.append( s.data )
284        if s.scaling == 'manual':
285            out.append('manual levels (%s)' %  (
286                    ', '.join([str(i) for i in s.manualLevels])))
287        else:
288            out.append('%(numLevels)i %(scaling)s levels (%(min)s to %(max)s)' % s)
289        return ', '.join(out)
290
291    def calculateLevels(self):
292        """Calculate contour levels from data and settings.
293
294        Returns levels as 1d numpy
295        """
296
297        # get dataset
298        s = self.settings
299        d = self.document
300
301        minval, maxval = 0., 1.
302        # scan data
303        data = s.get('data').getData(d)
304        if data is None or data.dimensions != 2 or data.data.size == 0:
305            return
306
307        minval, maxval = N.nanmin(data.data), N.nanmax(data.data)
308        if not N.isfinite(minval):
309            minval = 0.
310        if not N.isfinite(maxval):
311            maxval = 1.
312
313        # override if not auto
314        if s.min != 'Auto':
315            minval = s.min
316        if s.max != 'Auto':
317            maxval = s.max
318
319        numlevels = s.numLevels
320        scaling = s.scaling
321
322        if numlevels == 1 and scaling != 'manual':
323            # calculations below assume numlevels > 1
324            levels = N.array([minval,])
325        else:
326            # trap out silly cases
327            if minval == maxval:
328                minval = 0.
329                maxval = 1.
330
331            # calculate levels for each scaling
332            if scaling == 'linear':
333                delta = (maxval - minval) / (numlevels-1)
334                levels = minval + N.arange(numlevels)*delta
335            elif scaling == 'sqrt':
336                delta = N.sqrt(maxval - minval) / (numlevels-1)
337                levels = minval + (N.arange(numlevels)*delta)**2
338            elif scaling == 'log':
339                if minval == 0.:
340                    minval = 1.
341                if minval == maxval:
342                    maxval = minval + 1
343                delta = N.log(maxval/minval) / (numlevels-1)
344                levels = N.exp(N.arange(numlevels)*delta)*minval
345            elif scaling == 'squared':
346                delta = (maxval - minval)**2 / (numlevels-1)
347                levels = minval + N.sqrt(N.arange(numlevels)*delta)
348            else:
349                # manual
350                levels = N.array(s.manualLevels)
351
352        # for the user later
353        # we do this to convert array to list of floats
354        s.levelsOut = [float(i) for i in levels]
355
356        return minval, maxval, levels
357
358    def calculateSubLevels(self, minval, maxval, levels):
359        """Calculate sublevels between contours."""
360        s = self.settings
361        num = s.SubLines.numLevels
362        if s.SubLines.hide or len(s.SubLines.lines) == 0 or len(levels) <= 1:
363            return N.array([])
364
365        # indices where contour levels should be placed
366        numcont = (len(levels)-1) * num
367        indices = N.arange(numcont)
368        indices = indices[indices % num != 0]
369
370        scaling = s.scaling
371        if scaling == 'linear':
372            delta = (maxval-minval) / numcont
373            slev = indices*delta + minval
374        elif scaling == 'log':
375            delta = N.log( maxval/minval ) / numcont
376            slev = N.exp(indices*delta) * minval
377        elif scaling == 'sqrt':
378            delta = N.sqrt( maxval-minval ) / numcont
379            slev = minval + (indices*delta)**2
380        elif scaling == 'squared':
381            delta = (maxval-minval)**2 / numcont
382            slev = minval + N.sqrt(indices*delta)
383        elif scaling == 'manual':
384            drange = N.arange(1, num)
385            out = [[]]
386            for conmin, conmax in czip(levels[:-1], levels[1:]):
387                delta = (conmax-conmin) / num
388                out.append( conmin+drange*delta )
389            slev = N.hstack(out)
390
391        return slev
392
393    def affectsAxisRange(self):
394        """Range information provided by widget."""
395        s = self.settings
396        return ( (s.xAxis, 'sx'), (s.yAxis, 'sy') )
397
398    def getRange(self, axis, depname, axrange):
399        """Automatically determine the ranges of variable on the axes."""
400
401        # this is copied from Image, probably should combine
402        s = self.settings
403        d = self.document
404
405        # return if no data or if the dataset isn't two dimensional
406        data = s.get('data').getData(d)
407        if data is None or data.dimensions != 2 or data.data.size == 0:
408            return
409
410        xr, yr = data.getDataRanges()
411        if depname == 'sx':
412            axrange[0] = min( axrange[0], xr[0] )
413            axrange[1] = max( axrange[1], xr[1] )
414        elif depname == 'sy':
415            axrange[0] = min( axrange[0], yr[0] )
416            axrange[1] = max( axrange[1], yr[1] )
417
418    def getNumberKeys(self):
419        """How many keys to show."""
420        self.checkContoursUpToDate()
421        if self.settings.keyLevels:
422            return len( self.settings.levelsOut )
423        else:
424            return 0
425
426    def getKeyText(self, number):
427        """Get key entry."""
428        s = self.settings
429        if s.keyLevels:
430            cl = s.get('ContourLabels')
431            return utils.formatNumber(
432                s.levelsOut[number] * cl.scale,
433                cl.format,
434                locale=self.document.locale )
435        else:
436            return ''
437
438    def drawKeySymbol(self, number, painter, x, y, width, height):
439        """Draw key for contour level."""
440        painter.setPen(
441            self.settings.Lines.get('lines').makePen(painter, number))
442        painter.drawLine(x, y+height/2, x+width, y+height/2)
443
444    def checkContoursUpToDate(self):
445        """Update contours if necessary.
446        Returns True if okay to plot contours, False if error
447        """
448
449        s = self.settings
450        d = self.document
451
452        # return if no data or if the dataset isn't two dimensional
453        data = s.get('data').getData(d)
454        if data is None or data.dimensions != 2 or data.data.size == 0:
455            self.contsettings = None
456            s.levelsOut = []
457            return False
458
459        hashval = hash(bytes(data.data))
460        contsettings = (
461            s.min, s.max, s.numLevels, s.scaling,
462            s.SubLines.numLevels,
463            len(s.Fills.fills) == 0 or s.Fills.hide,
464            len(s.SubLines.lines) == 0 or s.SubLines.hide,
465            tuple(s.manualLevels),
466            hashval
467        )
468
469        if contsettings != self.contsettings:
470            self.updateContours()
471            self.contsettings = contsettings
472
473        return True
474
475    def dataDraw(self, painter, axes, posn, cliprect):
476        """Draw the contours."""
477
478        # update contours if necessary
479        if not self.checkContoursUpToDate():
480            return
481
482        self.plotContourFills(painter, posn, axes, cliprect)
483        self.plotContours(painter, posn, axes, cliprect)
484        self.plotSubContours(painter, posn, axes, cliprect)
485
486    def updateContours(self):
487        """Update calculated contours."""
488
489        s = self.settings
490        d = self.document
491
492        minval, maxval, levels = self.calculateLevels()
493        sublevels = self.calculateSubLevels(minval, maxval, levels)
494
495        # find coordinates of image coordinate bounds
496        data = s.get('data').getData(d)
497        if data is None or data.dimensions != 2 or data.data.size == 0:
498            return
499
500        rangex, rangey = data.getDataRanges()
501        yw, xw = data.data.shape
502        xc, yc = data.getPixelCentres()
503        xpts = N.reshape( N.tile(xc, yw), (yw, xw) )
504        ypts = N.tile(yc[:, N.newaxis], xw)
505
506        # only keep finite data points
507        mask = N.logical_not(N.isfinite(data.data))
508
509        # iterate over the levels and trace the contours
510        self._cachedcontours = None
511        self._cachedpolygons = None
512        self._cachedsubcontours = None
513
514        if Cntr is not None:
515            c = Cntr(xpts, ypts, data.data, mask)
516
517            # trace the contour levels
518            if len(s.Lines.lines) != 0:
519                self._cachedcontours = []
520                for level in levels:
521                    linelist = c.trace(level)
522                    self._cachedcontours.append( finitePoly(linelist) )
523
524            # trace the polygons between the contours
525            if len(s.Fills.fills) != 0 and len(levels) > 1 and not s.Fills.hide:
526                self._cachedpolygons = []
527                for level1, level2 in czip(levels[:-1], levels[1:]):
528                    linelist = c.trace(level1, level2)
529                    self._cachedpolygons.append( finitePoly(linelist) )
530
531            # trace sub-levels
532            if len(sublevels) > 0:
533                self._cachedsubcontours = []
534                for level in sublevels:
535                    linelist = c.trace(level)
536                    self._cachedsubcontours.append( finitePoly(linelist) )
537
538    def _plotContours(self, painter, posn, axes, linestyles,
539                      contours, showlabels, hidelines, clip):
540        """Plot a set of contours.
541        """
542
543        s = self.settings
544
545        # no lines cached as no line styles
546        if contours is None:
547            return
548
549        cl = s.get('ContourLabels')
550        font = cl.makeQFont(painter)
551        labelpen = cl.makeQPen(painter)
552        descent = qt.QFontMetricsF(font).descent()
553
554        # linelabeller does clipping and labelling of contours
555        linelabeller = ContourLineLabeller(
556            clip, cl.rotate, painter, font, self.document)
557        levels = []
558
559        # iterate over each level, and list of lines
560        for num, linelist in enumerate(contours):
561
562            if showlabels and num<len(s.levelsOut):
563                number = s.levelsOut[num]
564                text = utils.formatNumber(
565                    number * cl.scale, cl.format,
566                    locale=self.document.locale)
567                rend = utils.Renderer(
568                    painter, font, 0, 0, text, alignhorz=0,
569                    alignvert=0, angle=0, doc=self.document)
570                textdims = qt.QSizeF(*rend.getDimensions())
571                textdims += qt.QSizeF(descent*2, descent*2)
572            else:
573                textdims = qt.QSizeF(0, 0)
574
575            # iterate over each complete line of the contour
576            for curve in linelist:
577                # convert coordinates from graph to plotter
578                xplt = axes[0].dataToPlotterCoords(posn, curve[:,0])
579                yplt = axes[1].dataToPlotterCoords(posn, curve[:,1])
580
581                pts = qt.QPolygonF()
582                utils.addNumpyToPolygonF(pts, xplt, yplt)
583                linelabeller.addLine(pts, textdims)
584
585                if showlabels:
586                    linelabeller.labels.append(text)
587                else:
588                    linelabeller.labels.append(None)
589                levels.append(num)
590
591        painter.save()
592        painter.setPen(labelpen)
593        linelabeller.process()
594        painter.setClipPath(linelabeller.clippath)
595
596        for i in crange(linelabeller.getNumPolySets()):
597            polyset = linelabeller.getPolySet(i)
598            painter.setPen(linestyles.makePen(painter, levels[i]))
599            for poly in polyset:
600                painter.drawPolyline(poly)
601
602        painter.restore()
603
604    def plotContours(self, painter, posn, axes, clip):
605        """Plot the traced contours on the painter."""
606        s = self.settings
607        self._plotContours(painter, posn, axes, s.Lines.get('lines'),
608                           self._cachedcontours,
609                           not s.ContourLabels.hide, s.Lines.hide, clip)
610
611    def plotSubContours(self, painter, posn, axes, clip):
612        """Plot sub contours on painter."""
613        s = self.settings
614        self._plotContours(painter, posn, axes, s.SubLines.get('lines'),
615                           self._cachedsubcontours,
616                           False, s.SubLines.hide, clip)
617
618    def plotContourFills(self, painter, posn, axes, clip):
619        """Plot the traced contours on the painter."""
620
621        s = self.settings
622
623        # don't draw if there are no cached polygons
624        if self._cachedpolygons is None or s.Fills.hide:
625            return
626
627        # iterate over each level, and list of lines
628        for num, polylist in enumerate(self._cachedpolygons):
629
630            # iterate over each complete line of the contour
631            path = qt.QPainterPath()
632            for poly in polylist:
633                # convert coordinates from graph to plotter
634                xplt = axes[0].dataToPlotterCoords(posn, poly[:,0])
635                yplt = axes[1].dataToPlotterCoords(posn, poly[:,1])
636
637                pts = qt.QPolygonF()
638                utils.addNumpyToPolygonF(pts, xplt, yplt)
639
640                clippedpoly = qt.QPolygonF()
641                utils.polygonClip(pts, clip, clippedpoly)
642                path.addPolygon(clippedpoly)
643
644            # fill polygons
645            brush = s.Fills.get('fills').returnBrushExtended(num)
646            utils.brushExtFillPath(painter, brush, path)
647
648# allow the factory to instantiate a contour
649document.thefactory.register( Contour )
650