add_curve_sapling/utils.py

   1 # ##### BEGIN GPL LICENSE BLOCK #####
   2 #
   3 #  This program is free software; you can redistribute it and/or
   4 #  modify it under the terms of the GNU General Public License
   5 #  as published by the Free Software Foundation; either version 2
   6 #  of the License, or (at your option) any later version.
   7 #
   8 #  This program is distributed in the hope that it will be useful,
   9 #  but WITHOUT ANY WARRANTY; without even the implied warranty of
  10 #  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  11 #  GNU General Public License for more details.
  12 #
  13 #  You should have received a copy of the GNU General Public License
  14 #  along with this program; if not, write to the Free Software Foundation,
  15 #  Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
  16 #
  17 # ##### END GPL LICENSE BLOCK #####
  18
  19
  20 import bpy
  21 import time
  22 import copy
  23
  24 from mathutils import *
  25 from math import pi,sin,degrees,radians,atan2,copysign,cos,acos
  26 from random import random,uniform,seed,choice,getstate,setstate
  27 from bpy.props import *
  28 from collections import deque
  29
  30 # Initialise the split error and axis vectors
  31 splitError = 0.0
  32 zAxis = Vector((0,0,1))
  33 yAxis = Vector((0,1,0))
  34 xAxis = Vector((1,0,0))
  35
  36 # This class will contain a part of the tree which needs to be extended and the required tree parameters
  37 class stemSpline:
  38     def __init__(self,spline,curvature,curvatureV,segments,maxSegs,segLength,childStems,stemRadStart,stemRadEnd,splineNum):
  39         self.spline = spline
  40         self.p = spline.bezier_points[-1]
  41         self.curv = curvature
  42         self.curvV = curvatureV
  43         self.seg = segments
  44         self.segMax = maxSegs
  45         self.segL = segLength
  46         self.children = childStems
  47         self.radS = stemRadStart
  48         self.radE = stemRadEnd
  49         self.splN = splineNum
  50     # This method determines the quaternion of the end of the spline
  51     def quat(self):
  52         if len(self.spline.bezier_points) == 1:
  53             return ((self.spline.bezier_points[-1].handle_right - self.spline.bezier_points[-1].co).normalized()).to_track_quat('Z','Y')
  54         else:
  55             return ((self.spline.bezier_points[-1].co - self.spline.bezier_points[-2].co).normalized()).to_track_quat('Z','Y')
  56     # Determine the declination
  57     def dec(self):
  58         tempVec = zAxis.copy()
  59         tempVec.rotate(self.quat())
  60         return zAxis.angle(tempVec)
  61     # Update the end of the spline and increment the segment count
  62     def updateEnd(self):
  63         self.p = self.spline.bezier_points[-1]
  64         self.seg += 1
  65     # Determine the spread angle for a split
  66     def spreadAng(self):
  67         return radians(choice([-1,1])*(20 + 0.75*(30 + abs(degrees(self.dec()) - 90))*random()**2))
  68     # Determine the splitting angle for a split
  69     def splitAngle(self,splitAng,splitAngV):
  70         return max(0,splitAng+uniform(-splitAngV,splitAngV)-self.dec())
  71     # This is used to change the the curvature per segment of the spline
  72     def curvAdd(self,curvD):
  73         self.curv += curvD
  74
  75 # This class contains the data for a point where a new branch will sprout
  76 class childPoint:
  77     def __init__(self,coords,quat,radiusPar,offset,lengthPar,parBone):
  78         self.co = coords
  79         self.quat = quat
  80         self.radiusPar = radiusPar
  81         self.offset = offset
  82         self.lengthPar = lengthPar
  83         self.parBone = parBone
  84
  85
  86 # This function calculates the shape ratio as defined in the paper
  87 def shapeRatio(shape,ratio,pruneWidthPeak=0.0,prunePowerHigh=0.0,prunePowerLow=0.0):
  88     if shape == 0:
  89         return 0.2 + 0.8*ratio
  90     elif shape == 1:
  91         return 0.2 + 0.8*sin(pi*ratio)
  92     elif shape == 2:
  93         return 0.2 + 0.8*sin(0.5*pi*ratio)
  94     elif shape == 3:
  95         return 1.0
  96     elif shape == 4:
  97         return 0.5 + 0.5*ratio
  98     elif shape == 5:
  99         if ratio <= 0.7:
 100             return ratio/0.7
 101         else:
 102             return (1.0 - ratio)/0.3
 103     elif shape == 6:
 104         return 1.0 - 0.8*ratio
 105     elif shape == 7:
 106         if ratio <= 0.7:
 107             return 0.5 + 0.5*ratio/0.7
 108         else:
 109             return 0.5 + 0.5*(1.0 - ratio)/0.3
 110     elif shape == 8:
 111         if (ratio < (1 - pruneWidthPeak)) and (ratio > 0.0):
 112             return ((ratio/(1 - pruneWidthPeak))**prunePowerHigh)
 113         elif (ratio >= (1 - pruneWidthPeak)) and (ratio < 1.0):
 114             return (((1 - ratio)/pruneWidthPeak)**prunePowerLow)
 115         else:
 116             return 0.0
 117
 118 # This function determines the actual number of splits at a given point using the global error
 119 def splits(n):
 120     global splitError
 121     nEff = round(n + splitError,0)
 122     splitError -= (nEff - n)
 123     return int(nEff)
 124
 125 # Determine the declination from a given quaternion
 126 def declination(quat):
 127     tempVec = zAxis.copy()
 128     tempVec.rotate(quat)
 129     tempVec.normalize()
 130     return degrees(acos(tempVec.z))
 131
 132 # Determine the length of a child stem
 133 def lengthChild(lMax,offset,lPar,shape=False,lBase=None):
 134     if shape:
 135         return lPar*lMax*shapeRatio(shape,(lPar - offset)/(lPar - lBase))
 136     else:
 137         return lMax*(lPar - 0.6*offset)
 138
 139 # Find the actual downAngle taking into account the special case
 140 def downAngle(downAng,downAngV,lPar=None,offset=None,lBase=None):
 141     if downAngV < 0:
 142         return downAng + (uniform(-downAngV,downAngV)*(1 - 2*shapeRatio(0,(lPar - offset)/(lPar - lBase))))
 143     else:
 144         return downAng + uniform(-downAngV,downAngV)
 145
 146 # Returns the rotation matrix equivalent to i rotations by 2*pi/(n+1)
 147 def splitRotMat(n,i):
 148     return Matrix.Rotation(2*i*pi/(n+1),3,'Z')
 149
 150 # Returns the split angle
 151 def angleSplit(splitAng,splitAngV,quat):
 152     return max(0,splitAng+uniform(-splitAngV,splitAngV)-declination(quat))
 153
 154 # Returns number of stems a stem will sprout
 155 def stems(stemsMax,lPar,offset,lChild=False,lChildMax=None):
 156     if lChild:
 157         return stemsMax*(0.2 + 0.8*(lChild/lPar)/lChildMax)
 158     else:
 159         return stemsMax*(1.0 - 0.5*offset/lPar)
 160
 161 # Returns the spreading angle
 162 def spreadAng(dec):
 163     return radians(choice([-1,1])*(20 + 0.75*(30 + abs(dec - 90))*random()**2))
 164
 165 # Determines the angle of upward rotation of a segment due to attractUp
 166 def curveUp(attractUp,quat,curveRes):
 167     tempVec = yAxis.copy()
 168     tempVec.rotate(quat)
 169     tempVec.normalize()
 170     return attractUp*radians(declination(quat))*abs(tempVec.z)/curveRes
 171
 172 # Evaluate a bezier curve for the parameter 0<=t<=1 along its length
 173 def evalBez(p1,h1,h2,p2,t):
 174     return ((1-t)**3)*p1 + (3*t*(1-t)**2)*h1 + (3*(t**2)*(1-t))*h2 + (t**3)*p2
 175
 176 # Evaluate the unit tangent on a bezier curve for t
 177 def evalBezTan(p1,h1,h2,p2,t):
 178     return ((-3*(1-t)**2)*p1 + (-6*t*(1-t) + 3*(1-t)**2)*h1 + (-3*(t**2) + 6*t*(1-t))*h2 + (3*t**2)*p2).normalized()
 179
 180 # Determine the range of t values along a splines length where child stems are formed
 181 def findChildPoints(stemList,numChild):
 182     numPoints = sum([len(n.spline.bezier_points) for n in stemList])
 183     numSplines = len(stemList)
 184     numSegs = numPoints - numSplines
 185     numPerSeg = numChild/numSegs
 186     numMain = round(numPerSeg*stemList[0].segMax,0)
 187     return [(a+1)/(numMain) for a in range(int(numMain))]
 188
 189 # Find the coordinates, quaternion and radius for each t on the stem
 190 def interpStem(stem,tVals,lPar,parRad):
 191     tempList = deque()
 192     addpoint = tempList.append
 193     checkVal = (stem.segMax - len(stem.spline.bezier_points) + 1)/stem.segMax
 194     points = stem.spline.bezier_points
 195     numPoints = len(stem.spline.bezier_points)
 196     # Loop through all the parametric values to be determined
 197     for t in tVals:
 198         if (t >= checkVal) and (t < 1.0):
 199             scaledT = (t-checkVal)/(tVals[-1]-checkVal)
 200             length = (numPoints-1)*t#scaledT
 201             index = int(length)
 202             if scaledT == 1.0:
 203                 coord = points[-1].co
 204                 quat = (points[-1].handle_right - points[-1].co).to_track_quat('Z','Y')
 205                 radius = parRad#points[-2].radius
 206             else:
 207                 tTemp = length - index
 208                 coord = evalBez(points[index].co,points[index].handle_right,points[index+1].handle_left,points[index+1].co,tTemp)
 209                 quat = (evalBezTan(points[index].co,points[index].handle_right,points[index+1].handle_left,points[index+1].co,tTemp)).to_track_quat('Z','Y')
 210                 radius = (1-tTemp)*points[index].radius + tTemp*points[index+1].radius # Not sure if this is the parent radius at the child point or parent start radius
 211             addpoint(childPoint(coord,quat,(parRad, radius),t*lPar,lPar,'bone'+(str(stem.splN).rjust(3,'0'))+'.'+(str(index).rjust(3,'0'))))
 212     return tempList
 213
 214 # Convert a list of degrees to radians
 215 def toRad(list):
 216     return [radians(a) for a in list]
 217
 218 # This is the function which extends (or grows) a given stem.
 219 def growSpline(stem,numSplit,splitAng,splitAngV,splineList,attractUp,hType,splineToBone):
 220     # First find the current direction of the stem
 221     dir = stem.quat()
 222     # If the stem splits, we need to add new splines etc
 223     if numSplit > 0:
 224         # Get the curve data
 225         cuData = stem.spline.id_data.name
 226         cu = bpy.data.curves[cuData]
 227         # Now for each split add the new spline and adjust the growth direction
 228         for i in range(numSplit):
 229             newSpline = cu.splines.new('BEZIER')
 230             newPoint = newSpline.bezier_points[-1]
 231             (newPoint.co,newPoint.handle_left_type,newPoint.handle_right_type) = (stem.p.co,'VECTOR','VECTOR')
 232             newPoint.radius = stem.radS*(1 - stem.seg/stem.segMax) + stem.radE*(stem.seg/stem.segMax)
 233
 234             # Here we make the new "sprouting" stems diverge from the current direction
 235             angle = stem.splitAngle(splitAng,splitAngV)
 236             divRotMat = Matrix.Rotation(angle + stem.curv + uniform(-stem.curvV,stem.curvV),3,'X')#CurveUP should go after curve is applied
 237             dirVec = zAxis.copy()
 238             dirVec.rotate(divRotMat)
 239             dirVec.rotate(splitRotMat(numSplit,i+1))
 240             dirVec.rotate(dir)
 241
 242             # if attractUp != 0.0: # Shouldn't have a special case as this will mess with random number generation
 243             divRotMat = Matrix.Rotation(angle + stem.curv + uniform(-stem.curvV,stem.curvV),3,'X')
 244             dirVec = zAxis.copy()
 245             dirVec.rotate(divRotMat)
 246             dirVec.rotate(splitRotMat(numSplit,i+1))
 247             dirVec.rotate(dir)
 248
 249             #Different version of the commented code above. We could use the inbuilt vector rotations but given this is a special case, it can be quicker to initialise the vector to the correct value.
 250 #            angle = stem.splitAngle(splitAng,splitAngV)
 251 #            curveUpAng = curveUp(attractUp,dir,stem.segMax)
 252 #            angleX = angle + stem.curv + uniform(-stem.curvV,stem.curvV) - curveUpAng
 253 #            angleZ = 2*i*pi/(numSplit+1)
 254 #            dirVec = Vector((sin(angleX)*sin(angleZ), -sin(angleX)*cos(angleZ), cos(angleX)))
 255 #            dirVec.rotate(dir)
 256
 257             # Spread the stem out in a random fashion
 258             spreadMat = Matrix.Rotation(spreadAng(degrees(dirVec.z)),3,'Z')
 259             dirVec.rotate(spreadMat)
 260             # Introduce upward curvature
 261             upRotAxis = xAxis.copy()
 262             upRotAxis.rotate(dirVec.to_track_quat('Z','Y'))
 263             curveUpAng = curveUp(attractUp,dirVec.to_track_quat('Z','Y'),stem.segMax)
 264             upRotMat = Matrix.Rotation(-curveUpAng,3,upRotAxis)
 265             dirVec.rotate(upRotMat)
 266             # Make the growth vec the length of a stem segment
 267             dirVec.normalize()
 268             dirVec *= stem.segL
 269
 270             # Get the end point position
 271             end_co = stem.p.co.copy()
 272
 273             # Add the new point and adjust its coords, handles and radius
 274             newSpline.bezier_points.add()
 275             newPoint = newSpline.bezier_points[-1]
 276             (newPoint.co,newPoint.handle_left_type,newPoint.handle_right_type) = (end_co + dirVec,hType,hType)
 277             newPoint.radius = stem.radS*(1 - (stem.seg + 1)/stem.segMax) + stem.radE*((stem.seg + 1)/stem.segMax)
 278             # If this isn't the last point on a stem, then we need to add it to the list of stems to continue growing
 279             if stem.seg != stem.segMax:
 280                 splineList.append(stemSpline(newSpline,stem.curv-angle/(stem.segMax-stem.seg),stem.curvV,stem.seg+1,stem.segMax,stem.segL,stem.children,stem.radS,stem.radE,len(cu.splines)-1))
 281                 splineToBone.append('bone'+(str(stem.splN)).rjust(3,'0')+'.'+(str(len(stem.spline.bezier_points)-2)).rjust(3,'0'))
 282         # The original spline also needs to keep growing so adjust its direction too
 283         angle = stem.splitAngle(splitAng,splitAngV)
 284         divRotMat = Matrix.Rotation(angle + stem.curv + uniform(-stem.curvV,stem.curvV),3,'X')
 285         dirVec = zAxis.copy()
 286         dirVec.rotate(divRotMat)
 287         dirVec.rotate(dir)
 288         spreadMat = Matrix.Rotation(spreadAng(degrees(dirVec.z)),3,'Z')
 289         dirVec.rotate(spreadMat)
 290     else:
 291         # If there are no splits then generate the growth direction without accounting for spreading of stems
 292         dirVec = zAxis.copy()
 293         #curveUpAng = curveUp(attractUp,dir,stem.segMax)
 294         divRotMat = Matrix.Rotation(stem.curv + uniform(-stem.curvV,stem.curvV),3,'X')
 295         dirVec.rotate(divRotMat)
 296         #dirVec = Vector((0,-sin(stem.curv - curveUpAng),cos(stem.curv - curveUpAng)))
 297         dirVec.rotate(dir)
 298     upRotAxis = xAxis.copy()
 299     upRotAxis.rotate(dirVec.to_track_quat('Z','Y'))
 300     curveUpAng = curveUp(attractUp,dirVec.to_track_quat('Z','Y'),stem.segMax)
 301     upRotMat = Matrix.Rotation(-curveUpAng,3,upRotAxis)
 302     dirVec.rotate(upRotMat)
 303     dirVec.normalize()
 304     dirVec *= stem.segL
 305
 306     # Get the end point position
 307     end_co = stem.p.co.copy()
 308
 309     stem.spline.bezier_points.add()
 310     newPoint = stem.spline.bezier_points[-1]
 311     (newPoint.co,newPoint.handle_left_type,newPoint.handle_right_type) = (end_co + dirVec,hType,hType)
 312     newPoint.radius = stem.radS*(1 - (stem.seg + 1)/stem.segMax) + stem.radE*((stem.seg + 1)/stem.segMax)
 313     # There are some cases where a point cannot have handles as VECTOR straight away, set these now.
 314     if numSplit != 0:
 315         tempPoint = stem.spline.bezier_points[-2]
 316         (tempPoint.handle_left_type,tempPoint.handle_right_type) = ('VECTOR','VECTOR')
 317     if len(stem.spline.bezier_points) == 2:
 318         tempPoint = stem.spline.bezier_points[0]
 319         (tempPoint.handle_left_type,tempPoint.handle_right_type) = ('VECTOR','VECTOR')
 320     # Update the last point in the spline to be the newly added one
 321     stem.updateEnd()
 322     #return splineList
 323
 324 def genLeafMesh(leafScale,leafScaleX,loc,quat,index,downAngle,downAngleV,rotate,rotateV,oldRot,bend,leaves, leafShape):
 325     if leafShape == 'hex':
 326         verts = [Vector((0,0,0)),Vector((0.5,0,1/3)),Vector((0.5,0,2/3)),Vector((0,0,1)),Vector((-0.5,0,2/3)),Vector((-0.5,0,1/3))]
 327         edges = [[0,1],[1,2],[2,3],[3,4],[4,5],[5,0],[0,3]]
 328         faces = [[0,1,2,3],[0,3,4,5]]
 329     elif leafShape == 'rect':
 330         verts = [Vector((1,0,0)),Vector((1,0,1)),Vector((-1,0,1)),Vector((-1,0,0))]
 331         edges = [[0,1],[1,2],[2,3],[3,0]]
 332         faces = [[0,1,2,3],]
 333     #faces = [[0,1,5],[1,2,4,5],[2,3,4]]
 334
 335     vertsList = []
 336     facesList = []
 337
 338     # If the special -ve flag is used we need a different rotation of the leaf geometry
 339     if leaves < 0:
 340         rotMat = Matrix.Rotation(oldRot,3,'Y')
 341         oldRot += rotate/(abs(leaves)-1)
 342     else:
 343         oldRot += rotate+uniform(-rotateV,rotateV)
 344         downRotMat = Matrix.Rotation(downAngle+uniform(-downAngleV,downAngleV),3,'X')
 345         rotMat = Matrix.Rotation(oldRot,3,'Z')
 346
 347     normal = yAxis.copy()
 348     #dirVec = zAxis.copy()
 349     orientationVec = zAxis.copy()
 350
 351     # If the bending of the leaves is used we need to rotated them differently
 352     if (bend != 0.0) and (leaves >= 0):
 353 #        normal.rotate(downRotMat)
 354 #        orientationVec.rotate(downRotMat)
 355 #
 356 #        normal.rotate(rotMat)
 357 #        orientationVec.rotate(rotMat)
 358
 359         normal.rotate(quat)
 360         orientationVec.rotate(quat)
 361
 362         thetaPos = atan2(loc.y,loc.x)
 363         thetaBend = thetaPos - atan2(normal.y,normal.x)
 364         rotateZ = Matrix.Rotation(bend*thetaBend,3,'Z')
 365         normal.rotate(rotateZ)
 366         orientationVec.rotate(rotateZ)
 367
 368         phiBend = atan2((normal.xy).length,normal.z)
 369         orientation = atan2(orientationVec.y,orientationVec.x)
 370         rotateZOrien = Matrix.Rotation(orientation,3,'X')
 371
 372         rotateX = Matrix.Rotation(bend*phiBend,3,'Z')
 373
 374         rotateZOrien2 = Matrix.Rotation(-orientation,3,'X')
 375
 376     # For each of the verts we now rotate and scale them, then append them to the list to be added to the mesh
 377     for v in verts:
 378
 379         v.z *= leafScale
 380         v.x *= leafScaleX*leafScale
 381
 382         if leaves > 0:
 383             v.rotate(downRotMat)
 384
 385         v.rotate(rotMat)
 386         v.rotate(quat)
 387
 388         if (bend != 0.0) and (leaves > 0):
 389             # Correct the rotation
 390             v.rotate(rotateZ)
 391             v.rotate(rotateZOrien)
 392             v.rotate(rotateX)
 393             v.rotate(rotateZOrien2)
 394
 395         #v.rotate(quat)
 396     for v in verts:
 397         v += loc
 398         vertsList.append([v.x,v.y,v.z])
 399
 400     for f in faces:
 401         facesList.append([f[0] + index,f[1] + index,f[2] + index,f[3] + index])
 402     return vertsList,facesList,oldRot
 403
 404 def addTree(props):
 405     global splitError
 406     #startTime = time.time()
 407     # Set the seed for repeatable results
 408     seed(props.seed)#
 409
 410     # Set all other variables
 411     levels = props.levels#
 412     length = props.length#
 413     lengthV = props.lengthV#
 414     branches = props.branches#
 415     curveRes = props.curveRes#
 416     curve = toRad(props.curve)#
 417     curveV = toRad(props.curveV)#
 418     curveBack = toRad(props.curveBack)#
 419     baseSplits = props.baseSplits#
 420     segSplits = props.segSplits#
 421     splitAngle = toRad(props.splitAngle)#
 422     splitAngleV = toRad(props.splitAngleV)#
 423     scale = props.scale#
 424     scaleV = props.scaleV#
 425     attractUp = props.attractUp#
 426     shape = int(props.shape)#
 427     baseSize = props.baseSize
 428     ratio = props.ratio
 429     taper = props.taper#
 430     ratioPower = props.ratioPower#
 431     downAngle = toRad(props.downAngle)#
 432     downAngleV = toRad(props.downAngleV)#
 433     rotate = toRad(props.rotate)#
 434     rotateV = toRad(props.rotateV)#
 435     scale0 = props.scale0#
 436     scaleV0 = props.scaleV0#
 437     prune = props.prune#
 438     pruneWidth = props.pruneWidth#
 439     pruneWidthPeak = props.pruneWidthPeak#
 440     prunePowerLow = props.prunePowerLow#
 441     prunePowerHigh = props.prunePowerHigh#
 442     pruneRatio = props.pruneRatio#
 443     leafScale = props.leafScale#
 444     leafScaleX = props.leafScaleX#
 445     leafShape = props.leafShape
 446     bend = props.bend#
 447     leafDist = int(props.leafDist)#
 448     bevelRes = props.bevelRes#
 449     resU = props.resU#
 450     useArm = props.useArm
 451
 452     frameRate = props.frameRate
 453     windSpeed = props.windSpeed
 454     windGust = props.windGust
 455     armAnim = props.armAnim
 456
 457     leafObj = None
 458
 459     # Some effects can be turned ON and OFF, the necessary variables are changed here
 460     if not props.bevel:
 461         bevelDepth = 0.0
 462     else:
 463         bevelDepth = 1.0
 464
 465     if not props.showLeaves:
 466         leaves = 0
 467     else:
 468         leaves = props.leaves
 469
 470     if props.handleType == '0':
 471         handles = 'AUTO'
 472     else:
 473         handles = 'VECTOR'
 474
 475     for ob in bpy.data.objects:
 476         ob.select = False
 477
 478     childP = []
 479     stemList = []
 480
 481     # Initialise the tree object and curve and adjust the settings
 482     cu = bpy.data.curves.new('tree','CURVE')
 483     treeOb = bpy.data.objects.new('tree',cu)
 484     bpy.context.scene.objects.link(treeOb)
 485
 486     treeOb.location=bpy.context.scene.cursor_location
 487
 488     cu.dimensions = '3D'
 489     cu.fill_mode = 'FULL'
 490     cu.bevel_depth = bevelDepth
 491     cu.bevel_resolution = bevelRes
 492
 493     # Fix the scale of the tree now
 494     scaleVal = scale + uniform(-scaleV,scaleV)
 495
 496     # If pruning is turned on we need to draw the pruning envelope
 497     if prune:
 498         enHandle = 'VECTOR'
 499         enNum = 128
 500         enCu = bpy.data.curves.new('envelope','CURVE')
 501         enOb = bpy.data.objects.new('envelope',enCu)
 502         enOb.parent = treeOb
 503         bpy.context.scene.objects.link(enOb)
 504         newSpline = enCu.splines.new('BEZIER')
 505         newPoint = newSpline.bezier_points[-1]
 506         newPoint.co = Vector((0,0,scaleVal))
 507         (newPoint.handle_right_type,newPoint.handle_left_type) = (enHandle,enHandle)
 508         # Set the coordinates by varying the z value, envelope will be aligned to the x-axis
 509         for c in range(enNum):
 510             newSpline.bezier_points.add()
 511             newPoint = newSpline.bezier_points[-1]
 512             ratioVal = (c+1)/(enNum)
 513             zVal = scaleVal - scaleVal*(1-baseSize)*ratioVal
 514             newPoint.co = Vector((scaleVal*pruneWidth*shapeRatio(8,ratioVal,pruneWidthPeak,prunePowerHigh,prunePowerLow),0,zVal))
 515             (newPoint.handle_right_type,newPoint.handle_left_type) = (enHandle,enHandle)
 516         newSpline = enCu.splines.new('BEZIER')
 517         newPoint = newSpline.bezier_points[-1]
 518         newPoint.co = Vector((0,0,scaleVal))
 519         (newPoint.handle_right_type,newPoint.handle_left_type) = (enHandle,enHandle)
 520         # Create a second envelope but this time on the y-axis
 521         for c in range(enNum):
 522             newSpline.bezier_points.add()
 523             newPoint = newSpline.bezier_points[-1]
 524             ratioVal = (c+1)/(enNum)
 525             zVal = scaleVal - scaleVal*(1-baseSize)*ratioVal
 526             newPoint.co = Vector((0,scaleVal*pruneWidth*shapeRatio(8,ratioVal,pruneWidthPeak,prunePowerHigh,prunePowerLow),zVal))
 527             (newPoint.handle_right_type,newPoint.handle_left_type) = (enHandle,enHandle)
 528
 529     leafVerts = []
 530     leafFaces = []
 531     levelCount = []
 532
 533     splineToBone = deque([''])
 534     addsplinetobone = splineToBone.append
 535
 536     # Each of the levels needed by the user we grow all the splines
 537     for n in range(levels):
 538         storeN = n
 539         stemList = deque()
 540         addstem = stemList.append
 541         # If n is used as an index to access parameters for the tree it must be at most 3 or it will reference outside the array index
 542         n = min(3,n)
 543         vertAtt = attractUp
 544         splitError = 0.0
 545         # If this is the first level of growth (the trunk) then we need some special work to begin the tree
 546         if n == 0:
 547             vertAtt = 0.0
 548             newSpline = cu.splines.new('BEZIER')
 549             cu.resolution_u = resU
 550             newPoint = newSpline.bezier_points[-1]
 551             newPoint.co = Vector((0,0,0))
 552             newPoint.handle_right = Vector((0,0,1))
 553             newPoint.handle_left = Vector((0,0,-1))
 554             #(newPoint.handle_right_type,newPoint.handle_left_type) = ('VECTOR','VECTOR')
 555             branchL = (scaleVal)*(length[0] + uniform(-lengthV[0],lengthV[0]))
 556             childStems = branches[1]
 557             startRad = branchL*ratio*(scale0 + uniform(-scaleV0,scaleV0))
 558             endRad = startRad*(1 - taper[0])
 559             newPoint.radius = startRad
 560             addstem(stemSpline(newSpline,curve[0]/curveRes[0],curveV[0]/curveRes[0],0,curveRes[0],branchL/curveRes[0],childStems,startRad,endRad,0))
 561         # If this isn't the trunk then we may have multiple stem to intialise
 562         else:
 563             # Store the old rotation to allow new stems to be rotated away from the previous one.
 564             oldRotate = 0
 565             # For each of the points defined in the list of stem starting points we need to grow a stem.
 566             for p in childP:
 567                 # Add a spline and set the coordinate of the first point.
 568                 newSpline = cu.splines.new('BEZIER')
 569                 cu.resolution_u = resU
 570                 newPoint = newSpline.bezier_points[-1]
 571                 newPoint.co = p.co
 572                 tempPos = zAxis.copy()
 573                 # If the -ve flag for downAngle is used we need a special formula to find it
 574                 if downAngleV[n] < 0.0:
 575                     downV = downAngleV[n]*(1 - 2*shapeRatio(0,(p.lengthPar - p.offset)/(p.lengthPar - baseSize*scaleVal)))
 576                     random()
 577                 # Otherwise just find a random value
 578                 else:
 579                     downV = uniform(-downAngleV[n],downAngleV[n])
 580                 downRotMat = Matrix.Rotation(downAngle[n]+downV,3,'X')
 581                 tempPos.rotate(downRotMat)
 582                 # If the -ve flag for rotate is used we need to find which side of the stem the last child point was and then grow in the opposite direction.
 583                 if rotate[n] < 0.0:
 584                     oldRotate = -copysign(rotate[n] + uniform(-rotateV[n],rotateV[n]),oldRotate)
 585                 # Otherwise just generate a random number in the specified range
 586                 else:
 587                     oldRotate += rotate[n]+uniform(-rotateV[n],rotateV[n])
 588                 # Rotate the direction of growth and set the new point coordinates
 589                 rotMat = Matrix.Rotation(oldRotate,3,'Z')
 590                 tempPos.rotate(rotMat)
 591                 tempPos.rotate(p.quat)
 592                 newPoint.handle_right = p.co + tempPos
 593                 # If this is the first level of branching then upward attraction has no effect and a special formula is used to find branch length and the number of child stems
 594                 if n == 1:
 595                     vertAtt = 0.0
 596                     lMax = length[1] + uniform(-lengthV[1],lengthV[1])
 597                     branchL = p.lengthPar*lMax*shapeRatio(shape,(p.lengthPar - p.offset)/(p.lengthPar - baseSize*scaleVal))
 598                     childStems = branches[2]*(0.2 + 0.8*(branchL/p.lengthPar)/lMax)
 599                 elif storeN <= levels - 2:
 600                     branchL = (length[n] + uniform(-lengthV[n],lengthV[n]))*(p.lengthPar - 0.6*p.offset)
 601                     childStems = branches[min(3,n+1)]*(1.0 - 0.5*p.offset/p.lengthPar)
 602                 # If this is the last level before leaves then we need to generate the child points differently
 603                 else:
 604                     branchL = (length[n] + uniform(-lengthV[n],lengthV[n]))*(p.lengthPar - 0.6*p.offset)
 605                     if leaves < 0:
 606                         childStems = False
 607                     else:
 608                         childStems = leaves*shapeRatio(leafDist,p.offset/p.lengthPar)
 609                 # Determine the starting and ending radii of the stem using the tapering of the stem
 610                 startRad = min(p.radiusPar[0]*((branchL/p.lengthPar)**ratioPower), p.radiusPar[1])
 611                 endRad = startRad*(1 - taper[n])
 612                 newPoint.radius = startRad
 613                 # If curveBack is used then the curviness of the stem is different for the first half
 614                 if curveBack[n] == 0:
 615                     curveVal = curve[n]/curveRes[n]
 616                 else:
 617                     curveVal = 2*curve[n]/curveRes[n]
 618                 # Add the new stem to list of stems to grow and define which bone it will be parented to
 619                 addstem(stemSpline(newSpline,curveVal,curveV[n]/curveRes[n],0,curveRes[n],branchL/curveRes[n],childStems,startRad,endRad,len(cu.splines)-1))
 620                 addsplinetobone(p.parBone)
 621
 622         childP = []
 623         # Now grow each of the stems in the list of those to be extended
 624         for st in stemList:
 625             # When using pruning, we need to ensure that the random effects will be the same for each iteration to make sure the problem is linear.
 626             randState = getstate()
 627             startPrune = True
 628             lengthTest = 0.0
 629             # Store all the original values for the stem to make sure we have access after it has been modified by pruning
 630             originalLength = st.segL
 631             originalCurv = st.curv
 632             originalCurvV = st.curvV
 633             originalSeg = st.seg
 634             originalHandleR = st.p.handle_right.copy()
 635             originalHandleL = st.p.handle_left.copy()
 636             originalCo = st.p.co.copy()
 637             currentMax = 1.0
 638             currentMin = 0.0
 639             currentScale = 1.0
 640             oldMax = 1.0
 641             deleteSpline = False
 642             orginalSplineToBone = copy.copy(splineToBone)
 643             forceSprout = False
 644             # Now do the iterative pruning, this uses a binary search and halts once the difference between upper and lower bounds of the search are less than 0.005
 645             while startPrune and ((currentMax - currentMin) > 0.005):
 646                 setstate(randState)
 647
 648                 # If the search will halt after this iteration, then set the adjustment of stem length to take into account the pruning ratio
 649                 if (currentMax - currentMin) < 0.01:
 650                     currentScale = (currentScale - 1)*pruneRatio + 1
 651                     startPrune = False
 652                     forceSprout = True
 653                 # Change the segment length of the stem by applying some scaling
 654                 st.segL = originalLength*currentScale
 655                 # To prevent millions of splines being created we delete any old ones and replace them with only their first points to begin the spline again
 656                 if deleteSpline:
 657                     for x in splineList:
 658                         cu.splines.remove(x.spline)
 659                     newSpline = cu.splines.new('BEZIER')
 660                     newPoint = newSpline.bezier_points[-1]
 661                     newPoint.co = originalCo
 662                     newPoint.handle_right = originalHandleR
 663                     newPoint.handle_left = originalHandleL
 664                     (newPoint.handle_left_type,newPoint.handle_right_type) = ('VECTOR','VECTOR')
 665                     st.spline = newSpline
 666                     st.curv = originalCurv
 667                     st.curvV = originalCurvV
 668                     st.seg = originalSeg
 669                     st.p = newPoint
 670                     newPoint.radius = st.radS
 671                     splineToBone = orginalSplineToBone
 672
 673                 # Initialise the spline list for those contained in the current level of branching
 674                 splineList = [st]
 675                 # For each of the segments of the stem which must be grown we have to add to each spline in splineList
 676                 for k in range(curveRes[n]):
 677                     # Make a copy of the current list to avoid continually adding to the list we're iterating over
 678                     tempList = splineList[:]
 679                     #print('Leng: ',len(tempList))
 680                     # For each of the splines in this list set the number of splits and then grow it
 681                     for spl in tempList:
 682                         if k == 0:
 683                             numSplit = 0
 684                         elif (k == 1) and (n == 0):
 685                             numSplit = baseSplits
 686                         else:
 687                             numSplit = splits(segSplits[n])
 688                         if (k == int(curveRes[n]/2)) and (curveBack[n] != 0):
 689                             spl.curvAdd(-2*curve[n]/curveRes[n] + 2*curveBack[n]/curveRes[n])
 690                         growSpline(spl,numSplit,splitAngle[n],splitAngleV[n],splineList,vertAtt,handles,splineToBone)# Add proper refs for radius and attractUp
 691
 692                 # If pruning is enabled then we must to the check to see if the end of the spline is within the evelope
 693                 if prune:
 694                     # Check each endpoint to see if it is inside
 695                     for s in splineList:
 696                         coordMag = (s.spline.bezier_points[-1].co.xy).length
 697                         ratio = (scaleVal - s.spline.bezier_points[-1].co.z)/(scaleVal*(1 - baseSize))
 698                         # Don't think this if part is needed
 699                         if (n == 0) and (s.spline.bezier_points[-1].co.z < baseSize*scaleVal):
 700                             pass#insideBool = True
 701                         else:
 702                             insideBool = ((coordMag/scaleVal) < pruneWidth*shapeRatio(8,ratio,pruneWidthPeak,prunePowerHigh,prunePowerLow))
 703                         # If the point is not inside then we adjust the scale and current search bounds
 704                         if not insideBool:
 705                             oldMax = currentMax
 706                             currentMax = currentScale
 707                             currentScale = 0.5*(currentMax + currentMin)
 708                             break
 709                     # If the scale is the original size and the point is inside then we need to make sure it won't be pruned or extended to the edge of the envelope
 710                     if insideBool and (currentScale != 1):
 711                         currentMin = currentScale
 712                         currentMax = oldMax
 713                         currentScale = 0.5*(currentMax + currentMin)
 714                     if insideBool and ((currentMax - currentMin) == 1):
 715                         currentMin = 1
 716                 # If the search will halt on the next iteration then we need to make sure we sprout child points to grow the next splines or leaves
 717                 if (((currentMax - currentMin) < 0.005) or not prune) or forceSprout:
 718                     tVals = findChildPoints(splineList,st.children)
 719                     # If leaves is -ve then we need to make sure the only point which sprouts is the end of the spline
 720                     #if not st.children:
 721                     if not st.children:
 722                         tVals = [0.9]
 723                     # If this is the trunk then we need to remove some of the points because of baseSize
 724                     if n == 0:
 725                         trimNum = int(baseSize*(len(tVals)+1))
 726                         tVals = tVals[trimNum:]
 727
 728                     # For all the splines, we interpolate them and add the new points to the list of child points
 729                     for s in splineList:
 730                         #print(str(n)+'level: ',s.segMax*s.segL)
 731                         childP.extend(interpStem(s,tVals,s.segMax*s.segL,s.radS))
 732
 733                 # Force the splines to be deleted
 734                 deleteSpline = True
 735                 # If pruning isn't enabled then make sure it doesn't loop
 736                 if not prune:
 737                     startPrune = False
 738
 739         levelCount.append(len(cu.splines))
 740         # If we need to add leaves, we do it here
 741         if (storeN == levels-1) and leaves:
 742             oldRot = 0.0
 743             n = min(3,n+1)
 744             # For each of the child points we add leaves
 745             for cp in childP:
 746                 # If the special flag is set then we need to add several leaves at the same location
 747                 if leaves < 0:
 748                     oldRot = -rotate[n]/2
 749                     for g in range(abs(leaves)):
 750                         (vertTemp,faceTemp,oldRot) = genLeafMesh(leafScale,leafScaleX,cp.co,cp.quat,len(leafVerts),downAngle[n],downAngleV[n],rotate[n],rotateV[n],oldRot,bend,leaves, leafShape)
 751                         leafVerts.extend(vertTemp)
 752                         leafFaces.extend(faceTemp)
 753                 # Otherwise just add the leaves like splines.
 754                 else:
 755                     (vertTemp,faceTemp,oldRot) = genLeafMesh(leafScale,leafScaleX,cp.co,cp.quat,len(leafVerts),downAngle[n],downAngleV[n],rotate[n],rotateV[n],oldRot,bend,leaves, leafShape)
 756                     leafVerts.extend(vertTemp)
 757                     leafFaces.extend(faceTemp)
 758             # Create the leaf mesh and object, add geometry using from_pydata, edges are currently added by validating the mesh which isn't great
 759             leafMesh = bpy.data.meshes.new('leaves')
 760             leafObj = bpy.data.objects.new('leaves',leafMesh)
 761             bpy.context.scene.objects.link(leafObj)
 762             leafObj.parent = treeOb
 763             leafMesh.from_pydata(leafVerts,(),leafFaces)
 764
 765             if leafShape == 'rect':
 766                 leafMesh.uv_textures.new("leafUV")
 767                 uvlayer = leafMesh.uv_layers.active.data
 768
 769                 for i in range(0, len(leafFaces)):
 770                     uvlayer[i*4 + 0].uv = Vector((1, 0))
 771                     uvlayer[i*4 + 1].uv = Vector((1, 1))
 772                     uvlayer[i*4 + 2].uv = Vector((1 - leafScaleX, 1))
 773                     uvlayer[i*4 + 3].uv = Vector((1 - leafScaleX, 0))
 774
 775             leafMesh.validate()
 776
 777 # This can be used if we need particle leaves
 778 #            if (storeN == levels-1) and leaves:
 779 #                normalList = []
 780 #                oldRot = 0.0
 781 #                n = min(3,n+1)
 782 #                oldRot = 0.0
 783 #                # For each of the child points we add leaves
 784 #                for cp in childP:
 785 #                    # Here we make the new "sprouting" stems diverge from the current direction
 786 #                    dirVec = zAxis.copy()
 787 #                    oldRot += rotate[n]+uniform(-rotateV[n],rotateV[n])
 788 #                    downRotMat = Matrix.Rotation(downAngle[n]+uniform(-downAngleV[n],downAngleV[n]),3,'X')
 789 #                    rotMat = Matrix.Rotation(oldRot,3,'Z')
 790 #                    dirVec.rotate(downRotMat)
 791 #                    dirVec.rotate(rotMat)
 792 #                    dirVec.rotate(cp.quat)
 793 #                    normalList.extend([dirVec.x,dirVec.y,dirVec.z])
 794 #                    leafVerts.append(cp.co)
 795 #                # Create the leaf mesh and object, add geometry using from_pydata, edges are currently added by validating the mesh which isn't great
 796 #                edgeList = [(a,a+1) for a in range(len(childP)-1)]
 797 #                leafMesh = bpy.data.meshes.new('leaves')
 798 #                leafObj = bpy.data.objects.new('leaves',leafMesh)
 799 #                bpy.context.scene.objects.link(leafObj)
 800 #                leafObj.parent = treeOb
 801 #                leafMesh.from_pydata(leafVerts,edgeList,())
 802 #                leafMesh.vertices.foreach_set('normal',normalList)
 803
 804     # If we need and armature we add it
 805     if useArm:
 806         # Create the armature and objects
 807         arm = bpy.data.armatures.new('tree')
 808         armOb = bpy.data.objects.new('treeArm',arm)
 809         bpy.context.scene.objects.link(armOb)
 810
 811         # Create a new action to store all animation
 812         newAction = bpy.data.actions.new(name='windAction')
 813         armOb.animation_data_create()
 814         armOb.animation_data.action = newAction
 815
 816         arm.draw_type = 'STICK'
 817         arm.use_deform_delay = True
 818
 819         # Add the armature modifier to the curve
 820         armMod = treeOb.modifiers.new('windSway','ARMATURE')
 821         #armMod.use_apply_on_spline = True
 822         armMod.object = armOb
 823
 824         # If there are leaves then they need a modifier
 825         if leaves:
 826             armMod = leafObj.modifiers.new('windSway','ARMATURE')
 827             armMod.object = armOb
 828
 829         # Make sure all objects are deselected (may not be required?)
 830         for ob in bpy.data.objects:
 831             ob.select = False
 832
 833         # Set the armature as active and go to edit mode to add bones
 834         bpy.context.scene.objects.active = armOb
 835         bpy.ops.object.mode_set(mode='EDIT')
 836
 837         masterBones = []
 838
 839         offsetVal = 0
 840
 841         # For all the splines in the curve we need to add bones at each bezier point
 842         for i,parBone in enumerate(splineToBone):
 843             s = cu.splines[i]
 844             b = None
 845             # Get some data about the spline like length and number of points
 846             numPoints = len(s.bezier_points)-1
 847             splineL = numPoints*((s.bezier_points[0].co-s.bezier_points[1].co).length)
 848             # Set the random phase difference of the animation
 849             bxOffset = uniform(0,2*pi)
 850             byOffset = uniform(0,2*pi)
 851             # Set the phase multiplier for the spline
 852             bMult = (s.bezier_points[0].radius/splineL)*(1/15)*(1/frameRate)
 853             # For all the points in the curve (less the last) add a bone and name it by the spline it will affect
 854             for n in range(numPoints):
 855                 oldBone = b
 856                 boneName = 'bone'+(str(i)).rjust(3,'0')+'.'+(str(n)).rjust(3,'0')
 857                 b = arm.edit_bones.new(boneName)
 858                 b.head = s.bezier_points[n].co
 859                 b.tail = s.bezier_points[n+1].co
 860
 861                 b.head_radius = s.bezier_points[n].radius
 862                 b.tail_radius = s.bezier_points[n+1].radius
 863                 b.envelope_distance = 0.001#0.001
 864
 865                 # If there are leaves then we need a new vertex group so they will attach to the bone
 866                 if (len(levelCount) > 1) and (i >= levelCount[-2]) and leafObj:
 867                     leafObj.vertex_groups.new(boneName)
 868                 elif (len(levelCount) == 1) and leafObj:
 869                     leafObj.vertex_groups.new(boneName)
 870                 # If this is first point of the spline then it must be parented to the level above it
 871                 if n == 0:
 872                     if parBone:
 873                         b.parent = arm.edit_bones[parBone]
 874 #                            if len(parBone) > 11:
 875 #                                b.use_connect = True
 876                 # Otherwise, we need to attach it to the previous bone in the spline
 877                 else:
 878                     b.parent = oldBone
 879                     b.use_connect = True
 880                 # If there isn't a previous bone then it shouldn't be attached
 881                 if not oldBone:
 882                     b.use_connect = False
 883                 #tempList.append(b)
 884
 885                 # Add the animation to the armature if required
 886                 if armAnim:
 887                     # Define all the required parameters of the wind sway by the dimension of the spline
 888                     a0 = 4*splineL*(1-n/(numPoints+1))/s.bezier_points[n].radius
 889                     a1 = (windSpeed/50)*a0
 890                     a2 = (windGust/50)*a0 + a1/2
 891
 892                     # Add new fcurves for each sway  as well as the modifiers
 893                     swayX = armOb.animation_data.action.fcurves.new('pose.bones["' + boneName + '"].rotation_euler',0)
 894                     swayY = armOb.animation_data.action.fcurves.new('pose.bones["' + boneName + '"].rotation_euler',2)
 895
 896                     swayXMod1 = swayX.modifiers.new(type='FNGENERATOR')
 897                     swayXMod2 = swayX.modifiers.new(type='FNGENERATOR')
 898
 899                     swayYMod1 = swayY.modifiers.new(type='FNGENERATOR')
 900                     swayYMod2 = swayY.modifiers.new(type='FNGENERATOR')
 901
 902                     # Set the parameters for each modifier
 903                     swayXMod1.amplitude = radians(a1)/numPoints
 904                     swayXMod1.phase_offset = bxOffset
 905                     swayXMod1.phase_multiplier = degrees(bMult)
 906
 907                     swayXMod2.amplitude = radians(a2)/numPoints
 908                     swayXMod2.phase_offset = 0.7*bxOffset
 909                     swayXMod2.phase_multiplier = 0.7*degrees(bMult) # This shouldn't have to be in degrees but it looks much better in animation
 910                     swayXMod2.use_additive = True
 911
 912                     swayYMod1.amplitude = radians(a1)/numPoints
 913                     swayYMod1.phase_offset = byOffset
 914                     swayYMod1.phase_multiplier = degrees(bMult) # This shouldn't have to be in degrees but it looks much better in animation
 915
 916                     swayYMod2.amplitude = radians(a2)/numPoints
 917                     swayYMod2.phase_offset = 0.7*byOffset
 918                     swayYMod2.phase_multiplier = 0.7*degrees(bMult) # This shouldn't have to be in degrees but it looks much better in animation
 919                     swayYMod2.use_additive = True
 920
 921         # If there are leaves we need to assign vertices to their vertex groups
 922         if leaves:
 923             offsetVal = 0
 924             leafVertSize = 6
 925             if leafShape == 'rect':
 926                 leafVertSize = 4
 927             for i,cp in enumerate(childP):
 928                 for v in leafMesh.vertices[leafVertSize*i:(leafVertSize*i+leafVertSize)]:
 929                     leafObj.vertex_groups[cp.parBone].add([v.index],1.0,'ADD')
 930
 931         # Now we need the rotation mode to be 'XYZ' to ensure correct rotation
 932         bpy.ops.object.mode_set(mode='OBJECT')
 933         for p in armOb.pose.bones:
 934             p.rotation_mode = 'XYZ'
 935         treeOb.parent = armOb
 936     #print(time.time()-startTime)