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