1#    Copyright (C) 2014 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"""3D function plotting widget."""
20
21from __future__ import division, print_function
22import numpy as N
23
24from .. import qtall as qt
25from .. import setting
26from .. import document
27from .. import utils
28from ..helpers import threed
29
30from . import plotters3d
31
32def _(text, disambiguation=None, context='Function3D'):
33    """Translate text."""
34    return qt.QCoreApplication.translate(context, text, disambiguation)
35
36class FunctionSurface(setting.Surface3DWColorMap):
37    def __init__(self, *args, **argsv):
38        setting.Surface3DWColorMap.__init__(self, *args, **argsv)
39        self.get('color').newDefault(setting.Reference('../color'))
40
41class FunctionLine(setting.Line3DWColorMap):
42    def __init__(self, *args, **argsv):
43        setting.Line3DWColorMap.__init__(self, *args, **argsv)
44        self.get('color').newDefault(setting.Reference('../color'))
45        self.get('reflectivity').newDefault(20)
46
47class Function3D(plotters3d.GenericPlotter3D):
48    """Plotting functions in 3D."""
49
50    typename='function3d'
51    description=_('3D function')
52    allowusercreation=True
53
54    # list of the supported modes
55    _modes = [
56        'x=fn(y,z)', 'y=fn(x,z)', 'z=fn(x,y)',
57        'x,y,z=fns(t)',
58        'x,y=fns(z)', 'y,z=fns(x)', 'x,z=fns(y)'
59    ]
60
61    # which axes are affected by which modes
62    _affects = {
63        'z=fn(x,y)': (('zAxis', 'both'),),
64        'x=fn(y,z)': (('xAxis', 'both'),),
65        'y=fn(x,z)': (('yAxis', 'both'),),
66        'x,y,z=fns(t)': (('xAxis', 'sx'), ('yAxis', 'sy'), ('zAxis', 'sz')),
67        'x,y=fns(z)': (('xAxis', 'both'), ('yAxis', 'both')),
68        'y,z=fns(x)': (('yAxis', 'both'), ('zAxis', 'both')),
69        'x,z=fns(y)': (('xAxis', 'both'), ('zAxis', 'both')),
70    }
71
72    # which modes require which axes as inputs
73    _requires = {
74        'z=fn(x,y)': (('both', 'xAxis'), ('both', 'yAxis')),
75        'x=fn(y,z)': (('both', 'yAxis'), ('both', 'zAxis')),
76        'y=fn(x,z)': (('both', 'xAxis'), ('both', 'zAxis')),
77        'x,y,z=fns(t)': (),
78        'x,y=fns(z)': (('both', 'zAxis'),),
79        'y,z=fns(x)': (('both', 'xAxis'),),
80        'x,z=fns(y)': (('both', 'yAxis'),),
81    }
82
83    # which modes require which variables
84    _varmap = {
85        'x,y=fns(z)': ('x', 'y', 'z'),
86        'y,z=fns(x)': ('y', 'z', 'x'),
87        'x,z=fns(y)': ('x', 'z', 'y'),
88    }
89
90    @staticmethod
91    def _fnsetnshowhide(v):
92        """Return which function settings to show or hide depending on
93        mode."""
94        return {
95            'z=fn(x,y)': (('fnz',), ('fnx', 'fny')),
96            'x=fn(y,z)': (('fnx',), ('fny', 'fnz')),
97            'y=fn(x,z)': (('fny',), ('fnx', 'fnz')),
98            'x,y,z=fns(t)': (('fnx', 'fny', 'fnz'), ()),
99            'x,y=fns(z)': (('fnx', 'fny'), ('fnz',)),
100            'y,z=fns(x)': (('fny', 'fnz'), ('fnx',)),
101            'x,z=fns(y)': (('fnx', 'fnz'), ('fny',)),
102            }[v]
103
104    @classmethod
105    def addSettings(klass, s):
106        plotters3d.GenericPlotter3D.addSettings(s)
107
108        s.add(setting.Int(
109            'linesteps',
110            50,
111            minval = 3,
112            descr = _('Number of steps to evaluate the function over for lines'),
113            usertext=_('Line steps'),
114            formatting=True ))
115        s.add(setting.Int(
116            'surfacesteps',
117            20,
118            minval = 3,
119            descr = _('Number of steps to evaluate the function over for surfaces'
120                      ' in each direction'),
121            usertext=_('Surface steps'),
122            formatting=True ))
123        s.add(setting.ChoiceSwitch(
124            'mode', klass._modes,
125            'x,y,z=fns(t)',
126            descr=_('Type of function to plot'),
127            usertext=_('Mode'),
128            showfn=klass._fnsetnshowhide), 0)
129
130        s.add(setting.Str(
131            'fnx', '',
132            descr=_('Function for x coordinate'),
133            usertext=_('X function') ), 1)
134        s.add(setting.Str(
135            'fny', '',
136            descr=_('Function for y coordinate'),
137            usertext=_('Y function') ), 2)
138        s.add(setting.Str(
139            'fnz', '',
140            descr=_('Function for z coordinate'),
141            usertext=_('Z function') ), 3)
142        s.add(setting.Str(
143            'fncolor', '',
144            descr=_('Function to give color (0-1)'),
145            usertext=_('Color function') ), 4)
146
147        s.add( setting.Color(
148            'color',
149            'auto',
150            descr = _('Master color'),
151            usertext = _('Color'),
152            formatting=True), 0 )
153        s.add(FunctionLine(
154            'Line',
155            descr = _('Line settings'),
156            usertext = _('Plot line')),
157               pixmap = 'settings_plotline' )
158        s.add(setting.LineGrid3D(
159            'GridLine',
160            descr = _('Grid line settings'),
161            usertext = _('Grid line')),
162               pixmap = 'settings_gridline' )
163        s.add(FunctionSurface(
164            'Surface',
165            descr = _('Surface fill settings'),
166            usertext=_('Surface')),
167              pixmap='settings_bgfill' )
168
169    def affectsAxisRange(self):
170        """Which axes this widget affects."""
171        s = self.settings
172        affects = self._affects[s.mode]
173        return [(getattr(s, v[0]), v[1]) for v in affects]
174
175    def requiresAxisRange(self):
176        """Which axes this widget depends on."""
177        s = self.settings
178        requires = self._requires[s.mode]
179        return [(v[0], getattr(s, v[1])) for v in requires]
180
181    def getLineVals(self):
182        """Get vals for line plot by evaluating function."""
183        s = self.settings
184        mode = s.mode
185
186        if mode == 'x,y,z=fns(t)':
187            if not s.fnx or not s.fny or not s.fnz:
188                return None
189
190            xcomp = self.document.evaluate.compileCheckedExpression(s.fnx)
191            ycomp = self.document.evaluate.compileCheckedExpression(s.fny)
192            zcomp = self.document.evaluate.compileCheckedExpression(s.fnz)
193            if xcomp is None or ycomp is None or zcomp is None:
194                return None
195
196            # evaluate each expression
197            env = self.document.evaluate.context.copy()
198            env['t'] = N.linspace(0, 1, s.linesteps)
199            zeros = N.zeros(s.linesteps, dtype=N.float64)
200            try:
201                valsx = eval(xcomp, env) + zeros
202                valsy = eval(ycomp, env) + zeros
203                valsz = eval(zcomp, env) + zeros
204            except:
205                # something wrong in the evaluation
206                return None
207
208            fncolor = s.fncolor.strip()
209            if fncolor:
210                fncolor = self.document.evaluate.compileCheckedExpression(
211                    fncolor)
212                try:
213                    valscolor = eval(fncolor, env) + zeros
214                except:
215                    return None
216            else:
217                valscolor = None
218
219            retn = (valsx, valsy, valsz, valscolor)
220
221        else:
222            # lookup variables to go with function
223            var = self._varmap[mode]
224            fns = [getattr(s, 'fn'+var[0]), getattr(s, 'fn'+var[1])]
225            if not fns[0] or not fns[1]:
226                return None
227
228            # get points to evaluate functions over
229            axis = self.fetchAxis(var[2])
230            if not axis:
231                return
232            arange = axis.getPlottedRange()
233            if axis.settings.log:
234                evalpts = N.logspace(
235                    N.log10(arange[0]), N.log10(arange[1]), s.linesteps)
236            else:
237                evalpts = N.linspace(arange[0], arange[1], s.linesteps)
238
239            # evaluate expressions
240            env = self.document.evaluate.context.copy()
241            env[var[2]] = evalpts
242            zeros = N.zeros(s.linesteps, dtype=N.float64)
243            try:
244                vals1 = eval(fns[0], env) + zeros
245                vals2 = eval(fns[1], env) + zeros
246            except:
247                # something wrong in the evaluation
248                return None
249
250            fncolor = s.fncolor.strip()
251            if fncolor:
252                fncolor = self.document.evaluate.compileCheckedExpression(
253                    fncolor)
254                try:
255                    valscolor = eval(fncolor, env) + zeros
256                except:
257                    return None
258            else:
259                valscolor = None
260
261            # assign correct output points
262            retn = [None]*4
263            idxs = ('x', 'y', 'z')
264            retn[idxs.index(var[0])] = vals1
265            retn[idxs.index(var[1])] = vals2
266            retn[idxs.index(var[2])] = evalpts
267            retn[3] = valscolor
268
269        return retn
270
271    def getGridVals(self):
272        """Get values for 2D grid.
273
274        Return steps1, steps2, height, axidx, depvariable
275        axidx are the indices into the axes for height, step1, step2
276        """
277
278        s = self.settings
279        mode = s.mode
280
281        var, ovar1, ovar2, axidx = {
282            'x=fn(y,z)': ('x', 'y', 'z', (0, 1, 2)),
283            'y=fn(x,z)': ('y', 'z', 'x', (1, 2, 0)),
284            'z=fn(x,y)': ('z', 'x', 'y', (2, 0, 1)),
285        }[mode]
286
287        axes = self.fetchAxes()
288        if axes is None:
289            return None
290
291        # range of other axes
292        ax1, ax2 = axes[axidx[1]], axes[axidx[2]]
293        pr1 = ax1.getPlottedRange()
294        pr2 = ax2.getPlottedRange()
295        steps = s.surfacesteps
296        logax1, logax2 = ax1.settings.log, ax2.settings.log
297
298        # convert log ranges to linear temporarily
299        if logax1:
300            pr1 = N.log(pr1)
301        if logax2:
302            pr2 = N.log(pr2)
303
304        # set variables in environment
305        grid1, grid2 = N.indices((steps, steps))
306        del1 = (pr1[1]-pr1[0])/(steps-1.)
307        steps1 = N.arange(steps)*del1 + pr1[0]
308        grid1 = grid1*del1 + pr1[0]
309        del2 = (pr2[1]-pr2[0])/(steps-1.)
310        steps2 = N.arange(steps)*del2 + pr2[0]
311        grid2 = grid2*del2 + pr2[0]
312
313        fncolor = s.fncolor.strip()
314        if fncolor:
315            colgrid1 = 0.5*(grid1[1:,1:]+grid1[:-1,:-1])
316            colgrid2 = 0.5*(grid2[1:,1:]+grid2[:-1,:-1])
317            if logax1:
318                colgrid1 = N.exp(colgrid1)
319            if logax2:
320                colgrid2 = N.exp(colgrid2)
321
322        # convert back to log
323        if logax1:
324            grid1 = N.exp(grid1)
325        if logax2:
326            grid2 = N.exp(grid2)
327
328        env = self.document.evaluate.context.copy()
329        env[ovar1] = grid1
330        env[ovar2] = grid2
331
332        fn = getattr(s, 'fn%s' % var)  # get function from user
333        if not fn:
334            return
335        comp = self.document.evaluate.compileCheckedExpression(fn)
336        if comp is None:
337            return None
338
339        try:
340            height = eval(comp, env) + N.zeros(grid1.shape, dtype=N.float64)
341        except Exception:
342            # something wrong in the evaluation
343            return None
344
345        if fncolor:
346            compcolor = self.document.evaluate.compileCheckedExpression(
347                fncolor)
348            if not compcolor:
349                return
350            env[ovar1] = colgrid1
351            env[ovar2] = colgrid2
352
353            try:
354                colors = eval(compcolor, env) + N.zeros(
355                    colgrid1.shape, dtype=N.float64)
356            except Exception:
357                # something wrong in the evaluation
358                return None
359            colors = N.clip(colors, 0, 1)
360        else:
361            colors = None
362
363        return height, steps1, steps2, axidx, var, colors
364
365    def getRange(self, axis, depname, axrange):
366        """Get range of axis."""
367        mode = self.settings.mode
368        if mode == 'x,y,z=fns(t)':
369            # get range of each variable
370            retn = self.getLineVals()
371            if not retn:
372                return
373            valsx, valsy, valsz, valscolor = retn
374            coord = {'sx': valsx, 'sy': valsy, 'sz': valsz}[depname]
375
376        elif mode in ('x,y=fns(z)', 'y,z=fns(x)', 'x,z=fns(y)'):
377            # is this axis one of the ones we affect?
378            var = self._varmap[mode]
379            if self.fetchAxis(var[0]) is axis:
380                v = var[0]
381            elif self.fetchAxis(var[1]) is axis:
382                v = var[1]
383            else:
384                return
385
386            retn = self.getLineVals()
387            if not retn:
388                return
389            coord = retn[('x', 'y', 'z').index(v)]
390
391        elif mode in ('z=fn(x,y)', 'x=fn(y,z)', 'y=fn(x,z)'):
392            retn = self.getGridVals()
393            if not retn:
394                return
395            height, steps1, steps2, axidx, var, color = retn
396            if axis is not self.fetchAxis(var):
397                return
398            coord = height
399
400        finite = coord[N.isfinite(coord)]
401        if len(finite) > 0:
402            axrange[0] = min(axrange[0], finite.min())
403            axrange[1] = max(axrange[1], finite.max())
404
405    def updatePropColorMap(self, prop, setn, colorvals):
406        """Update line/surface properties given color map values.
407
408        prop is updated to use the data values colorvars (0-1) to apply
409        a color map from the setting setn given."""
410
411        cmap = self.document.evaluate.getColormap(
412            setn.colorMap, setn.colorMapInvert)
413        color2d = colorvals.reshape((1, colorvals.size))
414        colorimg = utils.applyColorMap(
415            cmap, 'linear', color2d, 0., 1., setn.transparency)
416        prop.setRGBs(colorimg)
417
418    def dataDrawSurface(self, painter, axes, container):
419        """Draw a surface plot."""
420        retn = self.getGridVals()
421        if not retn:
422            return
423        height, steps1, steps2, axidx, depvar, colors = retn
424        lheight = axes[axidx[0]].dataToLogicalCoords(height)
425        lsteps1 = axes[axidx[1]].dataToLogicalCoords(steps1)
426        lsteps2 = axes[axidx[2]].dataToLogicalCoords(steps2)
427
428        # draw grid over each axis
429        surfprop = lineprop = None
430        s = self.settings
431        if not s.Surface.hide:
432            surfprop = s.Surface.makeSurfaceProp(painter)
433            if colors is not None:
434                self.updatePropColorMap(surfprop, s.Surface, colors)
435
436        if not s.GridLine.hide:
437            lineprop = s.GridLine.makeLineProp(painter)
438
439        dirn = {'x': threed.Mesh.X_DIRN,
440                'y': threed.Mesh.Y_DIRN,
441                'z': threed.Mesh.Z_DIRN}[depvar]
442
443        mesh = threed.Mesh(
444            threed.ValVector(lsteps1), threed.ValVector(lsteps2),
445            threed.ValVector(N.ravel(lheight)),
446            dirn, lineprop, surfprop,
447            s.GridLine.hidehorz, s.GridLine.hidevert)
448        container.addObject(mesh)
449
450    def dataDrawLine(self, painter, axes, clipcontainer):
451        """Draw a line function."""
452
453        s = self.settings
454        if s.Line.hide:
455            return
456
457        retn = self.getLineVals()
458        if not retn:
459            return
460
461        valsx, valsy, valsz, valscolor = retn
462        lineprop = s.Line.makeLineProp(painter)
463        if valscolor is not None:
464            self.updatePropColorMap(lineprop, s.Line, valscolor)
465
466        lx = axes[0].dataToLogicalCoords(valsx)
467        ly = axes[1].dataToLogicalCoords(valsy)
468        lz = axes[2].dataToLogicalCoords(valsz)
469
470        line = threed.PolyLine(lineprop)
471        line.addPoints(
472            threed.ValVector(lx), threed.ValVector(ly),
473            threed.ValVector(lz))
474
475        clipcontainer.addObject(line)
476
477    def dataDrawToObject(self, painter, axes):
478        """Do actual drawing of function."""
479
480        s = self.settings
481        mode = s.mode
482
483        axes = self.fetchAxes()
484        if axes is None:
485            return
486
487        s = self.settings
488
489        clipcontainer = self.makeClipContainer(axes)
490        if mode in ('x,y,z=fns(t)', 'x,y=fns(z)', 'y,z=fns(x)', 'x,z=fns(y)'):
491            self.dataDrawLine(painter, axes, clipcontainer)
492        elif mode in ('z=fn(x,y)', 'x=fn(y,z)', 'y=fn(x,z)'):
493            self.dataDrawSurface(painter, axes, clipcontainer)
494
495        clipcontainer.assignWidgetId(id(self))
496        return clipcontainer
497
498document.thefactory.register(Function3D)
499