1# -*- coding: utf-8 -*-
2# Copyright (C) 2012, Keith Smith
3#
4# Visvis is distributed under the terms of the (new) BSD License.
5# The full license can be found in 'license.txt'.
6#
7# Thanks to Keith Smith for implementing the polar plot functionality.
8
9import numpy as np
10import visvis as vv
11
12from visvis.utils.pypoints import Pointset, is_Point, is_Pointset
13from visvis import PolarLine
14
15
16def makeArray(data):
17    if isinstance(data, np.ndarray):
18        return data
19    else:
20        # create numpy array
21        try:
22            l = len(data)
23            a = np.empty((l, 1))
24            for i in range(len(data)):
25                a[i] = data[i]
26            return a
27        except TypeError:
28            raise Exception("Cannot plot %s" % data.__class__.__name__)
29
30
31def _SetLimitsAfterDraw(event):
32    """ To be able to set the limits after the first draw. """
33    # Set limits
34    fig = event.owner
35    for axis in fig.FindObjects(vv.axises.PolarAxis2D):
36        limits = axis.GetLimits()
37        axis.SetLimits(rangeTheta=limits[0], rangeR=limits[1])
38    # Unsubscribe and redraw
39    fig.eventAfterDraw.Unbind(_SetLimitsAfterDraw)
40    fig.Draw()
41
42
43def polarplot(data1, data2=None, inRadians=False,
44            lw=1, lc='b', ls="-", mw=7, mc='b', ms='', mew=1, mec='k',
45            alpha=1, axesAdjust=True, axes=None, **kwargs):
46    """ polarplot(*args, inRadians=False,
47            lw=1, lc='b', ls="-", mw=7, mc='b', ms='', mew=1, mec='k',
48            alpha=1, axesAdjust=True, axes=None):
49
50    Plot 2D polar data, using a polar axis to draw a polar grid.
51
52    Usage
53    -----
54      * plot(Y, ...) plots a 1D polar signal.
55      * plot(X, Y, ...) also supplies angular coordinates
56      * plot(P, ...) plots using a Point or Pointset instance
57
58    Keyword arguments
59    -----------------
60    (The longer names for the line properties can also be used)
61    lw : scalar
62        lineWidth. The width of the line. If zero, no line is drawn.
63    mw : scalar
64        markerWidth. The width of the marker. If zero, no marker is drawn.
65    mew : scalar
66        markerEdgeWidth. The width of the edge of the marker.
67    lc : 3-element tuple or char
68        lineColor. The color of the line. A tuple should represent the RGB
69        values between 0 and 1. If a char is given it must be
70        one of 'rgbmcywk', for reg, green, blue, magenta, cyan, yellow,
71        white, black, respectively.
72    mc : 3-element tuple or char
73        markerColor. The color of the marker. See lineColor.
74    mec : 3-element tuple or char
75        markerEdgeColor. The color of the edge of the marker.
76    ls : string
77        lineStyle. The style of the line. (See below)
78    ms : string
79        markerStyle. The style of the marker. (See below)
80    axesAdjust : bool
81        If axesAdjust==True, this function will call axes.SetLimits(), and set
82        the camera type to 2D.
83    axes : Axes instance
84        Display the image in this axes, or the current axes if not given.
85
86    Line styles
87    -----------
88      * Solid line: '-'
89      * Dotted line: ':'
90      * Dashed line: '--'
91      * Dash-dot line: '-.' or '.-'
92      * A line that is drawn between each pair of points: '+'
93      * No line: '' or None.
94
95    Marker styles
96    -------------
97      * Plus: '+'
98      * Cross: 'x'
99      * Square: 's'
100      * Diamond: 'd'
101      * Triangle (pointing up, down, left, right): '^', 'v', '<', '>'
102      * Pentagram star: 'p' or '*'
103      * Hexgram: 'h'
104      * Point/cirle: 'o' or '.'
105      * No marker: '' or None
106
107    Polar axis
108    ----------
109    This polar axis has a few specialized methods for adjusting the polar
110    plot. Access these via vv.gca().axis.
111      * SetLimits(thetaRange, radialRange)
112      * thetaRange, radialRange = GetLimits()
113      * angularRefPos: Get and Set methods for the relative screen
114        angle of the 0 degree polar reference.  Default is 0 degs
115        which corresponds to the positive x-axis (y =0)
116      * isCW: Get and Set methods for the sense of rotation CCW or
117        CW. This method takes/returns a bool (True if the default CW).
118
119    Interaction
120    -----------
121      * Drag mouse up/down to translate radial axis.
122      * Drag mouse left/right to rotate angular ref position.
123      * Drag mouse + shift key up/down to rescale radial axis (min R fixed).
124
125    """
126
127    # create a dict from the properties and combine with kwargs
128    tmp = {'lineWidth': lw, 'lineColor': lc, 'lineStyle': ls,
129                'markerWidth': mw, 'markerColor': mc, 'markerStyle': ms,
130                'markerEdgeWidth': mew, 'markerEdgeColor': mec}
131    for i in tmp:
132        if not i in kwargs:
133            kwargs[i] = tmp[i]
134
135    ##  create the data
136    if is_Pointset(data1):
137        pp = data1
138    elif is_Point(data1):
139        pp = Pointset(data1.ndim)
140        pp.append(data1)
141    else:
142
143        if data1 is None:
144            raise ValueError("The first argument cannot be None!")
145        data1 = makeArray(data1)
146
147        if data2 is None:
148            # R data is given, thetadata must be
149            # a range starting from 0 degrees
150            data2 = data1
151            data1 = np.arange(0, data2.shape[0])
152        else:
153            data2 = makeArray(data2)
154
155        # check dimensions
156        L = data1.size
157        if L != data2.size:
158            raise ValueError("Array dimensions do not match! %i vs %i " %
159                    (data1.size, data2.size))
160
161        # build points
162        data1 = data1.reshape((data1.size, 1))
163        data2 = data2.reshape((data2.size, 1))
164
165    if not inRadians:
166        data1 = np.pi * data1 / 180.0
167
168    ## create the line
169    if axes is None:
170        axes = vv.gca()
171    axes.axisType = 'polar'
172    fig = axes.GetFigure()
173
174    l = PolarLine(axes, data1, data2)
175    l.lw = kwargs['lineWidth']
176    l.lc = kwargs['lineColor']
177    l.ls = kwargs['lineStyle']
178    l.mw = kwargs['markerWidth']
179    l.mc = kwargs['markerColor']
180    l.ms = kwargs['markerStyle']
181    l.mew = kwargs['markerEdgeWidth']
182    l.mec = kwargs['markerEdgeColor']
183    l.alpha = alpha
184
185    ## almost done...
186
187    # Init axis
188#     axes.axis.SetLimits()
189
190    if axesAdjust:
191        if axes.daspectAuto is None:
192            axes.daspectAuto = True
193        axes.cameraType = '2d'
194        axes.SetLimits()
195
196    # Subsribe after-draw event handler
197    # (unsubscribe first in case we do multiple plots)
198    fig.eventAfterDraw.Unbind(_SetLimitsAfterDraw)
199    fig.eventAfterDraw.Bind(_SetLimitsAfterDraw)
200
201    # Return
202    axes.Draw()
203    return l
204
205
206if __name__ == '__main__':
207    # Make data
208    angs = 0.1 + np.linspace(-90, 90, 181)  # 0.1+ get rid of singularity
209    angsRads = np.pi * angs / 180.0
210    mag = 10 * np.log10(np.abs(np.sin(10 * angsRads) / angsRads)) + angsRads
211    mag = mag - np.max(mag)
212    # Show data
213    vv.polarplot( angs, mag, lc='b')
214    vv.polarplot(angs+20, mag, lc='r', lw=2)
215    a = vv.gca() # Triggers an update required for polar plots
216