Skip to content
Snippets Groups Projects
utils.py 77.7 KiB
Newer Older
  • Learn to ignore specific revisions
  • # -*- coding: utf-8 -*-
    
    Andrew Hale's avatar
    Andrew Hale committed
    # ##### BEGIN GPL LICENSE BLOCK #####
    #
    #  This program is free software; you can redistribute it and/or
    #  modify it under the terms of the GNU General Public License
    #  as published by the Free Software Foundation; either version 2
    #  of the License, or (at your option) any later version.
    #
    #  This program is distributed in the hope that it will be useful,
    #  but WITHOUT ANY WARRANTY; without even the implied warranty of
    #  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    #  GNU General Public License for more details.
    #
    #  You should have received a copy of the GNU General Public License
    #  along with this program; if not, write to the Free Software Foundation,
    #  Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
    #
    # ##### END GPL LICENSE BLOCK #####
    
    
    import bpy
    
    Andrew Hale's avatar
    Andrew Hale committed
    import time
    import copy
    
    
    from mathutils import (
            Euler,
            Matrix,
            Vector,
            )
    
    from math import pi, sin, degrees, radians, atan2, copysign, cos, acos
    
    from math import floor
    
    from random import random, uniform, seed, choice, getstate, setstate, randint
    from collections import deque, OrderedDict
    
    tau = 2 * pi
    
    Andrew Hale's avatar
    Andrew Hale committed
    
    # Initialise the split error and axis vectors
    splitError = 0.0
    
    zAxis = Vector((0, 0, 1))
    yAxis = Vector((0, 1, 0))
    xAxis = Vector((1, 0, 0))
    
    Andrew Hale's avatar
    Andrew Hale committed
    
    
    Andrew Hale's avatar
    Andrew Hale committed
    # This class will contain a part of the tree which needs to be extended and the required tree parameters
    class stemSpline:
    
        def __init__(self, spline, curvature, curvatureV, attractUp, segments, maxSegs,
                     segLength, childStems, stemRadStart, stemRadEnd, splineNum, ofst, pquat):
    
    Andrew Hale's avatar
    Andrew Hale committed
            self.spline = spline
            self.p = spline.bezier_points[-1]
            self.curv = curvature
            self.curvV = curvatureV
    
            self.vertAtt = attractUp
    
    Andrew Hale's avatar
    Andrew Hale committed
            self.seg = segments
            self.segMax = maxSegs
            self.segL = segLength
            self.children = childStems
            self.radS = stemRadStart
            self.radE = stemRadEnd
            self.splN = splineNum
    
            self.offsetLen = ofst
            self.patentQuat = pquat
            self.curvSignx = 1
            self.curvSigny = 1
    
    
    Andrew Hale's avatar
    Andrew Hale committed
        # This method determines the quaternion of the end of the spline
        def quat(self):
            if len(self.spline.bezier_points) == 1:
    
                return ((self.spline.bezier_points[-1].handle_right -
                         self.spline.bezier_points[-1].co).normalized()).to_track_quat('Z', 'Y')
    
    Andrew Hale's avatar
    Andrew Hale committed
            else:
    
                return ((self.spline.bezier_points[-1].co -
                         self.spline.bezier_points[-2].co).normalized()).to_track_quat('Z', 'Y')
    
    
    Andrew Hale's avatar
    Andrew Hale committed
        # Determine the declination
        def dec(self):
            tempVec = zAxis.copy()
            tempVec.rotate(self.quat())
            return zAxis.angle(tempVec)
    
    Andrew Hale's avatar
    Andrew Hale committed
        # Update the end of the spline and increment the segment count
        def updateEnd(self):
            self.p = self.spline.bezier_points[-1]
            self.seg += 1
    
    
    Andrew Hale's avatar
    Andrew Hale committed
    # This class contains the data for a point where a new branch will sprout
    class childPoint:
    
        def __init__(self, coords, quat, radiusPar, offset, sOfst, lengthPar, parBone):
    
    Andrew Hale's avatar
    Andrew Hale committed
            self.co = coords
            self.quat = quat
            self.radiusPar = radiusPar
            self.offset = offset
    
            self.stemOffset = sOfst
    
    Andrew Hale's avatar
    Andrew Hale committed
            self.lengthPar = lengthPar
            self.parBone = parBone
    
    
    # This function calculates the shape ratio as defined in the paper
    
    def shapeRatio(shape, ratio, pruneWidthPeak=0.0, prunePowerHigh=0.0, prunePowerLow=0.0, custom=None):
    
    Andrew Hale's avatar
    Andrew Hale committed
        if shape == 0:
    
            return 0.05 + 0.95 * ratio  # 0.2 + 0.8 * ratio
    
    Andrew Hale's avatar
    Andrew Hale committed
        elif shape == 1:
    
            return 0.2 + 0.8 * sin(pi * ratio)
    
    Andrew Hale's avatar
    Andrew Hale committed
        elif shape == 2:
    
            return 0.2 + 0.8 * sin(0.5 * pi * ratio)
    
    Andrew Hale's avatar
    Andrew Hale committed
        elif shape == 3:
            return 1.0
        elif shape == 4:
    
            return 0.5 + 0.5 * ratio
    
    Andrew Hale's avatar
    Andrew Hale committed
        elif shape == 5:
            if ratio <= 0.7:
    
                return 0.05 + 0.95 * ratio / 0.7
    
    Andrew Hale's avatar
    Andrew Hale committed
            else:
    
                return 0.05 + 0.95 * (1.0 - ratio) / 0.3
    
    Andrew Hale's avatar
    Andrew Hale committed
        elif shape == 6:
    
            return 1.0 - 0.8 * ratio
    
    Andrew Hale's avatar
    Andrew Hale committed
        elif shape == 7:
            if ratio <= 0.7:
    
                return 0.5 + 0.5 * ratio / 0.7
    
    Andrew Hale's avatar
    Andrew Hale committed
            else:
    
                return 0.5 + 0.5 * (1.0 - ratio) / 0.3
    
    Andrew Hale's avatar
    Andrew Hale committed
        elif shape == 8:
    
            r = 1 - ratio
            if r == 1:
                v = custom[3]
            elif r >= custom[2]:
                pos = (r - custom[2]) / (1 - custom[2])
    
                # if (custom[0] >= custom[1] <= custom[3]) or (custom[0] <= custom[1] >= custom[3]):
    
                pos = pos * pos
                v = (pos * (custom[3] - custom[1])) + custom[1]
            else:
                pos = r / custom[2]
    
                # if (custom[0] >= custom[1] <= custom[3]) or (custom[0] <= custom[1] >= custom[3]):
    
                pos = 1 - (1 - pos) * (1 - pos)
                v = (pos * (custom[1] - custom[0])) + custom[0]
    
            return v
    
        elif shape == 9:
    
    Andrew Hale's avatar
    Andrew Hale committed
            if (ratio < (1 - pruneWidthPeak)) and (ratio > 0.0):
    
                return ((ratio / (1 - pruneWidthPeak))**prunePowerHigh)
    
    Andrew Hale's avatar
    Andrew Hale committed
            elif (ratio >= (1 - pruneWidthPeak)) and (ratio < 1.0):
    
                return (((1 - ratio) / pruneWidthPeak)**prunePowerLow)
    
    Andrew Hale's avatar
    Andrew Hale committed
            else:
                return 0.0
    
    
        elif shape == 10:
            return 0.5 + 0.5 * (1 - ratio)
    
    
    Andrew Hale's avatar
    Andrew Hale committed
    # This function determines the actual number of splits at a given point using the global error
    def splits(n):
        global splitError
    
        nEff = round(n + splitError, 0)
    
    Andrew Hale's avatar
    Andrew Hale committed
        splitError -= (nEff - n)
        return int(nEff)
    
    
    def splits2(n):
        r = random()
        if r < n:
            return 1
        else:
            return 0
    
    
    def splits3(n):
        ni = int(n)
        nf = n - int(n)
        r = random()
        if r < nf:
            return ni + 1
        else:
            return ni + 0
    
    
    Andrew Hale's avatar
    Andrew Hale committed
    # Determine the declination from a given quaternion
    def declination(quat):
        tempVec = zAxis.copy()
        tempVec.rotate(quat)
        tempVec.normalize()
        return degrees(acos(tempVec.z))
    
    
    Andrew Hale's avatar
    Andrew Hale committed
    # Determines the angle of upward rotation of a segment due to attractUp
    
    def curveUp(attractUp, quat, curveRes):
    
    Andrew Hale's avatar
    Andrew Hale committed
        tempVec = yAxis.copy()
        tempVec.rotate(quat)
        tempVec.normalize()
    
    
        dec = radians(declination(quat))
    
        curveUpAng = attractUp * dec * abs(tempVec.z) / curveRes
    
        if (-dec + curveUpAng) < -pi:
            curveUpAng = -pi + dec
        if (dec - curveUpAng) < 0:
            curveUpAng = dec
        return curveUpAng
    
    Andrew Hale's avatar
    Andrew Hale committed
    # Evaluate a bezier curve for the parameter 0<=t<=1 along its length
    
    def evalBez(p1, h1, h2, p2, t):
    
        return ((1 - t)**3) * p1 + (3 * t * (1 - t)**2) * h1 + (3 * (t**2) * (1 - t)) * h2 + (t**3) * p2
    
    
    Andrew Hale's avatar
    Andrew Hale committed
    
    # Evaluate the unit tangent on a bezier curve for t
    
    def evalBezTan(p1, h1, h2, p2, t):
    
        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()
    
    
    Andrew Hale's avatar
    Andrew Hale committed
    
    # Determine the range of t values along a splines length where child stems are formed
    
    def findChildPoints(stemList, numChild):
    
    Andrew Hale's avatar
    Andrew Hale committed
        numPoints = sum([len(n.spline.bezier_points) for n in stemList])
        numSplines = len(stemList)
        numSegs = numPoints - numSplines
    
        numPerSeg = numChild / numSegs
        numMain = round(numPerSeg * stemList[0].segMax, 0)
        return [(a + 1) / (numMain) for a in range(int(numMain))]
    
    
    def findChildPoints2(stemList, numChild):
    
        return [(a + 1) / (numChild) for a in range(int(numChild))]
    
    
    Andrew Hale's avatar
    Andrew Hale committed
    # Find the coordinates, quaternion and radius for each t on the stem
    
    def interpStem1(stem, tVals, lPar, parRad):
        points = stem.spline.bezier_points
        numPoints = len(points)
        checkVal = (stem.segMax - (numPoints - 1)) / stem.segMax
        # Loop through all the parametric values to be determined
    
    Andrew Hale's avatar
    Andrew Hale committed
        tempList = deque()
    
        for t in tVals:
            if t == 1.0:
    
                index = numPoints - 2
    
                coord = points[-1].co
                quat = (points[-1].handle_right - points[-1].co).to_track_quat('Z', 'Y')
                radius = points[-1].radius
    
    
                tempList.append(
                        childPoint(coord, quat, (parRad, radius), t, lPar, 'bone' +
                                  (str(stem.splN).rjust(3, '0')) + '.' + (str(index).rjust(3, '0')))
                        )
    
    
            elif (t >= checkVal) and (t < 1.0):
                scaledT = (t - checkVal) / ((1 - checkVal) + .0001)
    
                length = (numPoints - 1) * scaledT
    
                index = int(length)
    
                tTemp = length - index
    
                coord = evalBez(
                            points[index].co, points[index].handle_right,
                            points[index + 1].handle_left, points[index + 1].co, tTemp
                            )
                quat = (
                    evalBezTan(
                        points[index].co, points[index].handle_right,
                        points[index + 1].handle_left, points[index + 1].co, tTemp)
                        ).to_track_quat('Z', 'Y')
                # Not sure if this is the parent radius at the child point or parent start radius
                radius = (1 - tTemp) * points[index].radius + tTemp * points[index + 1].radius
    
                tempList.append(
                        childPoint(
                                coord, quat, (parRad, radius), t, lPar, 'bone' +
                                (str(stem.splN).rjust(3, '0')) + '.' + (str(index).rjust(3, '0')))
                                )
    
    def interpStem(stem, tVals, lPar, parRad, maxOffset, baseSize):
    
    Andrew Hale's avatar
    Andrew Hale committed
        points = stem.spline.bezier_points
    
        numSegs = len(points) - 1
        stemLen = stem.segL * numSegs
    
        checkBottom = stem.offsetLen / maxOffset
        checkTop = checkBottom + (stemLen / maxOffset)
    
    
    Andrew Hale's avatar
    Andrew Hale committed
        # Loop through all the parametric values to be determined
    
        tempList = deque()
    
    Andrew Hale's avatar
    Andrew Hale committed
        for t in tVals:
    
            if (t >= checkBottom) and (t <= checkTop) and (t < 1.0):
                scaledT = (t - checkBottom) / (checkTop - checkBottom)
                ofst = ((t - baseSize) / (checkTop - baseSize)) * (1 - baseSize) + baseSize
    
                length = numSegs * scaledT
    
    Andrew Hale's avatar
    Andrew Hale committed
                index = int(length)
    
                tTemp = length - index
    
    
                coord = evalBez(
                            points[index].co, points[index].handle_right,
                            points[index + 1].handle_left, points[index + 1].co, tTemp
                            )
                quat = (
                    evalBezTan(
                        points[index].co, points[index].handle_right,
                        points[index + 1].handle_left, points[index + 1].co, tTemp
                        )
                    ).to_track_quat('Z', 'Y')
                # Not sure if this is the parent radius at the child point or parent start radius
                radius = (1 - tTemp) * points[index].radius + tTemp * points[index + 1].radius
    
                tempList.append(
                        childPoint(
                            coord, quat, (parRad, radius), t, ofst, lPar,
                            'bone' + (str(stem.splN).rjust(3, '0')) + '.' + (str(index).rjust(3, '0')))
                        )
    
        # add stem at tip
        index = numSegs - 1
    
        coord = points[-1].co
        quat = (points[-1].handle_right - points[-1].co).to_track_quat('Z', 'Y')
        radius = points[-1].radius
    
        tempList.append(
                    childPoint(
                            coord, quat, (parRad, radius), 1, 1, lPar,
                            'bone' + (str(stem.splN).rjust(3, '0')) + '.' + (str(index).rjust(3, '0'))
                            )
                        )
    
    Andrew Hale's avatar
    Andrew Hale committed
        return tempList
    
    
    # round down bone number
    def roundBone(bone, step):
        bone_i = bone[:-3]
        bone_n = int(bone[-3:])
        bone_n = int(bone_n / step) * step
        return bone_i + str(bone_n).rjust(3, '0')
    
    
    Andrew Hale's avatar
    Andrew Hale committed
    # Convert a list of degrees to radians
    def toRad(list):
        return [radians(a) for a in list]
    
    
    def anglemean(a1, a2, fac):
        x1 = sin(a1)
        y1 = cos(a1)
        x2 = sin(a2)
        y2 = cos(a2)
        x = x1 + (x2 - x1) * fac
        y = y1 + (y2 - y1) * fac
        return atan2(x, y)
    
    
    
    Andrew Hale's avatar
    Andrew Hale committed
    # This is the function which extends (or grows) a given stem.
    
    def growSpline(n, stem, numSplit, splitAng, splitAngV, splineList,
                   hType, splineToBone, closeTip, kp, splitHeight, outAtt, stemsegL,
    
                   lenVar, taperCrown, boneStep, rotate, rotateV):
    
    
        # curv at base
    
        sCurv = stem.curv
        if (n == 0) and (kp <= splitHeight):
            sCurv = 0.0
    
    
        # curveangle = sCurv + (uniform(-stem.curvV, stem.curvV) * kp)
        # curveVar = uniform(-stem.curvV, stem.curvV) * kp
    
        curveangle = sCurv + (uniform(0, stem.curvV) * kp * stem.curvSignx)
        curveVar = uniform(0, stem.curvV) * kp * stem.curvSigny
        stem.curvSignx *= -1
        stem.curvSigny *= -1
    
        curveVarMat = Matrix.Rotation(curveVar, 3, 'Y')
    
    
    Andrew Hale's avatar
    Andrew Hale committed
        # First find the current direction of the stem
        dir = stem.quat()
    
    
        if n == 0:
            adir = zAxis.copy()
            adir.rotate(dir)
    
            ry = atan2(adir[0], adir[2])
            adir.rotate(Euler((0, -ry, 0)))
            rx = atan2(adir[1], adir[2])
    
            dir = Euler((-rx, ry, 0), 'XYZ')
    
    
        # length taperCrown
    
        if n == 0:
            dec = declination(dir) / 180
            dec = dec ** 2
            tf = 1 - (dec * taperCrown * 30)
            tf = max(.1, tf)
        else:
            tf = 1.0
    
    
        # outward attraction
    
        if (n > 0) and (kp > 0) and (outAtt > 0):
            p = stem.p.co.copy()
            d = atan2(p[0], -p[1]) + tau
            edir = dir.to_euler('XYZ', Euler((0, 0, d), 'XYZ'))
            d = anglemean(edir[2], d, (kp * outAtt))
            dirv = Euler((edir[0], edir[1], d), 'XYZ')
            dir = dirv.to_quaternion()
    
        """
        # parent weight
        parWeight = kp * degrees(stem.curvV) * pi
        parWeight = parWeight * kp
        parWeight = kp
        if (n > 1) and (parWeight != 0):
            d1 = zAxis.copy()
            d2 = zAxis.copy()
            d1.rotate(dir)
            d2.rotate(stem.patentQuat)
    
            x = d1[0] + ((d2[0] - d1[0]) * parWeight)
            y = d1[1] + ((d2[1] - d1[1]) * parWeight)
            z = d1[2] + ((d2[2] - d1[2]) * parWeight)
    
            d3 = Vector((x, y, z))
            dir = d3.to_track_quat('Z', 'Y')
        """
    
    Andrew Hale's avatar
    Andrew Hale committed
        # If the stem splits, we need to add new splines etc
        if numSplit > 0:
            # Get the curve data
            cuData = stem.spline.id_data.name
            cu = bpy.data.curves[cuData]
    
            # calc split angles
    
            angle = choice([-1, 1]) * (splitAng + uniform(-splitAngV, splitAngV))
            if n > 0:
    
                # make branches flatter
    
                angle *= max(1 - declination(dir) / 90, 0) * .67 + .33
            spreadangle = choice([-1, 1]) * (splitAng + uniform(-splitAngV, splitAngV))
    
    
            # branchRotMat = Matrix.Rotation(radians(uniform(0, 360)), 3, 'Z')
    
            if not hasattr(stem, 'rLast'):
                stem.rLast = radians(uniform(0, 360))
    
            br = rotate[0] + uniform(-rotateV[0], rotateV[0])
            branchRot = stem.rLast + br
            branchRotMat = Matrix.Rotation(branchRot, 3, 'Z')
            stem.rLast = branchRot
    
    
    Andrew Hale's avatar
    Andrew Hale committed
            # Now for each split add the new spline and adjust the growth direction
            for i in range(numSplit):
    
                # find split scale
    
                lenV = uniform(1 - lenVar, 1 + lenVar)
                bScale = min(lenV * tf, 1)
    
    
    Andrew Hale's avatar
    Andrew Hale committed
                newSpline = cu.splines.new('BEZIER')
                newPoint = newSpline.bezier_points[-1]
    
                (newPoint.co, newPoint.handle_left_type, newPoint.handle_right_type) = (stem.p.co, 'VECTOR', 'VECTOR')
    
                newPoint.radius = (
                            stem.radS * (1 - stem.seg / stem.segMax) + stem.radE * (stem.seg / stem.segMax)
                            ) * bScale
    
    Andrew Hale's avatar
    Andrew Hale committed
                # Here we make the new "sprouting" stems diverge from the current direction
    
                divRotMat = Matrix.Rotation(angle + curveangle, 3, 'X')
    
    Andrew Hale's avatar
    Andrew Hale committed
                dirVec = zAxis.copy()
                dirVec.rotate(divRotMat)
    
                # horizontal curvature variation
    
                dirVec.rotate(curveVarMat)
    
    
                if n == 0:  # Special case for trunk splits
    
                    dirVec.rotate(branchRotMat)
    
    
                    ang = pi - ((tau) / (numSplit + 1)) * (i + 1)
    
                    dirVec.rotate(Matrix.Rotation(ang, 3, 'Z'))
    
    Andrew Hale's avatar
    Andrew Hale committed
    
                # Spread the stem out in a random fashion
    
                spreadMat = Matrix.Rotation(spreadangle, 3, 'Y')
    
                if n != 0:  # Special case for trunk splits
    
                    dirVec.rotate(spreadMat)
    
                dirVec.rotate(dir)
    
    
    Andrew Hale's avatar
    Andrew Hale committed
                # Introduce upward curvature
                upRotAxis = xAxis.copy()
    
                upRotAxis.rotate(dirVec.to_track_quat('Z', 'Y'))
                curveUpAng = curveUp(stem.vertAtt, dirVec.to_track_quat('Z', 'Y'), stem.segMax)
                upRotMat = Matrix.Rotation(-curveUpAng, 3, upRotAxis)
    
    Andrew Hale's avatar
    Andrew Hale committed
                dirVec.rotate(upRotMat)
    
    Andrew Hale's avatar
    Andrew Hale committed
                # Make the growth vec the length of a stem segment
                dirVec.normalize()
    
                # split length variation
    
                stemL = stemsegL * lenV
                dirVec *= stemL * tf
                ofst = stem.offsetLen + (stem.segL * (len(stem.spline.bezier_points) - 1))
    
    
                # dirVec *= stem.segL
    
    
                # Get the end point position
                end_co = stem.p.co.copy()
    
    
    Andrew Hale's avatar
    Andrew Hale committed
                # Add the new point and adjust its coords, handles and radius
    
                newSpline.bezier_points.add(1)
    
    Andrew Hale's avatar
    Andrew Hale committed
                newPoint = newSpline.bezier_points[-1]
    
                (newPoint.co, newPoint.handle_left_type, newPoint.handle_right_type) = (end_co + dirVec, hType, hType)
    
                newPoint.radius = (
                            stem.radS * (1 - (stem.seg + 1) / stem.segMax) +
                            stem.radE * ((stem.seg + 1) / stem.segMax)
                            ) * bScale
                if (stem.seg == stem.segMax - 1) and closeTip:
    
                    newPoint.radius = 0.0
    
                # If this isn't the last point on a stem, then we need to add it
                # to the list of stems to continue growing
                # print(stem.seg != stem.segMax, stem.seg, stem.segMax)
    
                # if stem.seg != stem.segMax: # if probs not necessary
    
                nstem = stemSpline(
                            newSpline, stem.curv, stem.curvV, stem.vertAtt, stem.seg + 1,
                            stem.segMax, stemL, stem.children,
                            stem.radS * bScale, stem.radE * bScale, len(cu.splines) - 1, ofst, stem.quat()
                            )
                nstem.splitlast = 1  # numSplit  # keep track of numSplit for next stem
    
                nstem.rLast = branchRot + pi
                splineList.append(nstem)
    
                bone = 'bone' + (str(stem.splN)).rjust(3, '0') + '.' + \
                        (str(len(stem.spline.bezier_points) - 2)).rjust(3, '0')
    
                bone = roundBone(bone, boneStep[n])
    
                splineToBone.append((bone, False, True, len(stem.spline.bezier_points) - 2))
    
    Andrew Hale's avatar
    Andrew Hale committed
            # The original spline also needs to keep growing so adjust its direction too
    
            divRotMat = Matrix.Rotation(-angle + curveangle, 3, 'X')
    
    Andrew Hale's avatar
    Andrew Hale committed
            dirVec = zAxis.copy()
            dirVec.rotate(divRotMat)
    
            # horizontal curvature variation
    
            dirVec.rotate(curveVarMat)
    
    
            if n == 0:  # Special case for trunk splits
    
                dirVec.rotate(branchRotMat)
    
    
            spreadMat = Matrix.Rotation(-spreadangle, 3, 'Y')
    
            if n != 0:  # Special case for trunk splits
    
                dirVec.rotate(spreadMat)
    
    
    Andrew Hale's avatar
    Andrew Hale committed
            dirVec.rotate(dir)
    
            stem.splitlast = 1  # numSplit #keep track of numSplit for next stem
    
    Andrew Hale's avatar
    Andrew Hale committed
        else:
            # If there are no splits then generate the growth direction without accounting for spreading of stems
            dirVec = zAxis.copy()
    
            divRotMat = Matrix.Rotation(curveangle, 3, 'X')
    
    Andrew Hale's avatar
    Andrew Hale committed
            dirVec.rotate(divRotMat)
    
            # horizontal curvature variation
    
            dirVec.rotate(curveVarMat)
    
    
    Andrew Hale's avatar
    Andrew Hale committed
            dirVec.rotate(dir)
    
            stem.splitlast = 0  # numSplit #keep track of numSplit for next stem
    
    
        # Introduce upward curvature
    
    Andrew Hale's avatar
    Andrew Hale committed
        upRotAxis = xAxis.copy()
    
        upRotAxis.rotate(dirVec.to_track_quat('Z', 'Y'))
        curveUpAng = curveUp(stem.vertAtt, dirVec.to_track_quat('Z', 'Y'), stem.segMax)
        upRotMat = Matrix.Rotation(-curveUpAng, 3, upRotAxis)
    
    Andrew Hale's avatar
    Andrew Hale committed
        dirVec.rotate(upRotMat)
    
    Andrew Hale's avatar
    Andrew Hale committed
        dirVec.normalize()
    
        dirVec *= stem.segL * tf
    
    
        # Get the end point position
        end_co = stem.p.co.copy()
    
    
        stem.spline.bezier_points.add(1)
    
    Andrew Hale's avatar
    Andrew Hale committed
        newPoint = stem.spline.bezier_points[-1]
    
        (newPoint.co, newPoint.handle_left_type, newPoint.handle_right_type) = (end_co + dirVec, hType, hType)
    
        newPoint.radius = stem.radS * (1 - (stem.seg + 1) / stem.segMax) + \
                          stem.radE * ((stem.seg + 1) / stem.segMax)
    
        if (stem.seg == stem.segMax - 1) and closeTip:
    
            newPoint.radius = 0.0
    
        # There are some cases where a point cannot have handles as VECTOR straight away, set these now
    
    Andrew Hale's avatar
    Andrew Hale committed
        if len(stem.spline.bezier_points) == 2:
            tempPoint = stem.spline.bezier_points[0]
    
            (tempPoint.handle_left_type, tempPoint.handle_right_type) = ('VECTOR', 'VECTOR')
    
    Andrew Hale's avatar
    Andrew Hale committed
        # Update the last point in the spline to be the newly added one
        stem.updateEnd()
    
        # return splineList
    
    
    def genLeafMesh(leafScale, leafScaleX, leafScaleT, leafScaleV, loc, quat,
                    offset, index, downAngle, downAngleV, rotate, rotateV, oldRot,
    
                    bend, leaves, leafShape, leafangle, horzLeaves):
    
            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))
                ]
    
            edges = [[0, 1], [1, 2], [2, 3], [3, 4], [4, 5], [5, 0], [0, 3]]
            faces = [[0, 1, 2, 3], [0, 3, 4, 5]]
    
        elif leafShape == 'rect':
    
            # verts = [Vector((1, 0, 0)), Vector((1, 0, 1)), Vector((-1, 0, 1)), Vector((-1, 0, 0))]
    
            verts = [Vector((.5, 0, 0)), Vector((.5, 0, 1)), Vector((-.5, 0, 1)), Vector((-.5, 0, 0))]
            edges = [[0, 1], [1, 2], [2, 3], [3, 0]]
            faces = [[0, 1, 2, 3]]
        elif leafShape == 'dFace':
            verts = [Vector((.5, .5, 0)), Vector((.5, -.5, 0)), Vector((-.5, -.5, 0)), Vector((-.5, .5, 0))]
            edges = [[0, 1], [1, 2], [2, 3], [3, 0]]
            faces = [[0, 3, 2, 1]]
        elif leafShape == 'dVert':
            verts = [Vector((0, 0, 1))]
            edges = []
            faces = []
    
    Andrew Hale's avatar
    Andrew Hale committed
    
        vertsList = []
        facesList = []
    
        normal = Vector((0, 0, 1))
    
    Andrew Hale's avatar
    Andrew Hale committed
    
        if leaves < 0:
    
            rotMat = Matrix.Rotation(oldRot, 3, 'Y')
    
    Andrew Hale's avatar
    Andrew Hale committed
        else:
    
            rotMat = Matrix.Rotation(oldRot, 3, 'Z')
    
        # 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
    
        if rotate < 0.0:
            oldRot = -copysign(rotate + uniform(-rotateV, rotateV), oldRot)
        else:
            # If the special -ve flag for leaves is used we need a different rotation of the leaf geometry
            if leaves == -1:
    
                # oldRot = 0
    
                rotMat = Matrix.Rotation(0, 3, 'Y')
            elif leaves < -1:
                oldRot += rotate / (-leaves - 1)
            else:
                oldRot += rotate + uniform(-rotateV, rotateV)
    
        """
        if leaves < 0:
            rotMat = Matrix.Rotation(oldRot, 3, 'Y')
        else:
            rotMat = Matrix.Rotation(oldRot, 3, 'Z')
        """
    
            # downRotMat = Matrix.Rotation(downAngle+uniform(-downAngleV, downAngleV), 3, 'X')
    
    
            if downAngleV > 0.0:
                downV = -downAngleV * offset
            else:
                downV = uniform(-downAngleV, downAngleV)
            downRotMat = Matrix.Rotation(downAngle + downV, 3, 'X')
    
    
        # leaf scale variation
    
        if (leaves < -1) and (rotate != 0):
            f = 1 - abs((oldRot - (rotate / (-leaves - 1))) / (rotate / 2))
        else:
            f = offset
    
        if leafScaleT < 0:
            leafScale = leafScale * (1 - (1 - f) * -leafScaleT)
        else:
            leafScale = leafScale * (1 - f * leafScaleT)
    
        leafScale = leafScale * uniform(1 - leafScaleV, 1 + leafScaleV)
    
        if leafShape == 'dFace':
            leafScale = leafScale * .1
    
        # If the bending of the leaves is used we need to rotate them differently
    
    Andrew Hale's avatar
    Andrew Hale committed
        if (bend != 0.0) and (leaves >= 0):
    
            normal = yAxis.copy()
            orientationVec = zAxis.copy()
    
    Andrew Hale's avatar
    Andrew Hale committed
    
            normal.rotate(quat)
            orientationVec.rotate(quat)
    
    
            thetaPos = atan2(loc.y, loc.x)
            thetaBend = thetaPos - atan2(normal.y, normal.x)
    
            rotateZ = Matrix.Rotation(bend * thetaBend, 3, 'Z')
    
    Andrew Hale's avatar
    Andrew Hale committed
            normal.rotate(rotateZ)
            orientationVec.rotate(rotateZ)
    
    
            phiBend = atan2((normal.xy).length, normal.z)
            orientation = atan2(orientationVec.y, orientationVec.x)
            rotateZOrien = Matrix.Rotation(orientation, 3, 'X')
    
            rotateX = Matrix.Rotation(bend * phiBend, 3, 'Z')
    
            rotateZOrien2 = Matrix.Rotation(-orientation, 3, 'X')
    
    Andrew Hale's avatar
    Andrew Hale committed
    
        # For each of the verts we now rotate and scale them, then append them to the list to be added to the mesh
        for v in verts:
            v.z *= leafScale
    
            v.x *= leafScaleX * leafScale
    
            v.rotate(Euler((0, 0, radians(180))))
    
    
            # leafangle
    
            v.rotate(Matrix.Rotation(radians(-leafangle), 3, 'X'))
    
            if rotate < 0:
                v.rotate(Euler((0, 0, radians(90))))
                if oldRot < 0:
                    v.rotate(Euler((0, 0, radians(180))))
    
            if (leaves > 0) and (rotate > 0) and horzLeaves:
                nRotMat = Matrix.Rotation(-oldRot + rotate, 3, 'Z')
                v.rotate(nRotMat)
    
    
    Andrew Hale's avatar
    Andrew Hale committed
            if leaves > 0:
                v.rotate(downRotMat)
    
            v.rotate(rotMat)
            v.rotate(quat)
    
            if (bend != 0.0) and (leaves > 0):
                # Correct the rotation
                v.rotate(rotateZ)
                v.rotate(rotateZOrien)
                v.rotate(rotateX)
                v.rotate(rotateZOrien2)
    
    
        if leafShape == 'dVert':
            normal = verts[0]
            normal.normalize()
            v = loc
            vertsList.append([v.x, v.y, v.z])
        else:
            for v in verts:
                v += loc
                vertsList.append([v.x, v.y, v.z])
            for f in faces:
                facesList.append([f[0] + index, f[1] + index, f[2] + index, f[3] + index])
    
        return vertsList, facesList, normal, oldRot
    
    def create_armature(armAnim, leafP, cu, frameRate, leafMesh, leafObj, leafVertSize, leaves,
                        levelCount, splineToBone, treeOb, wind, gust, gustF, af1, af2, af3,
                        leafAnim, loopFrames, previewArm, armLevels, makeMesh, boneStep):
    
        arm = bpy.data.armatures.new('tree')
        armOb = bpy.data.objects.new('treeArm', arm)
    
        bpy.context.scene.collection.objects.link(armOb)
    
        # Create a new action to store all animation
        newAction = bpy.data.actions.new(name='windAction')
        armOb.animation_data_create()
        armOb.animation_data.action = newAction
    
        arm.display_type = 'STICK'
    
        # Add the armature modifier to the curve
        armMod = treeOb.modifiers.new('windSway', 'ARMATURE')
    
        if previewArm:
            armMod.show_viewport = False
    
            arm.display_type = 'WIRE'
    
            treeOb.hide_viewport = True
    
        armMod.use_apply_on_spline = True
    
        armMod.use_bone_envelopes = True
    
        armMod.use_vertex_groups = False  # curves don't have vertex groups (yet)
    
        # If there are leaves then they need a modifier
        if leaves:
            armMod = leafObj.modifiers.new('windSway', 'ARMATURE')
            armMod.object = armOb
    
            armMod.use_bone_envelopes = False
            armMod.use_vertex_groups = True
    
        # Make sure all objects are deselected (may not be required?)
    
        for ob in bpy.context.view_layer.objects:
    
            ob.select_set(state=False)
    
        fps = bpy.context.scene.render.fps
        animSpeed = (24 / fps) * frameRate
    
    
        # Set the armature as active and go to edit mode to add bones
    
        bpy.context.view_layer.objects.active = armOb
    
        bpy.ops.object.mode_set(mode='EDIT')
        # For all the splines in the curve we need to add bones at each bezier point
        for i, parBone in enumerate(splineToBone):
    
            if (i < levelCount[armLevels]) or (armLevels == -1) or (not makeMesh):
                s = cu.splines[i]
                b = None
                # Get some data about the spline like length and number of points
                numPoints = len(s.bezier_points) - 1
    
    
                # find branching level
    
                level = 0
                for l, c in enumerate(levelCount):
                    if i < c:
                        level = l
                        break
                level = min(level, 3)
    
                step = boneStep[level]
    
                # Calculate things for animation
    
                    splineL = numPoints * ((s.bezier_points[0].co - s.bezier_points[1].co).length)
                    # Set the random phase difference of the animation
                    bxOffset = uniform(0, tau)
                    byOffset = uniform(0, tau)
                    # Set the phase multiplier for the spline
    
                    # bMult_r = (s.bezier_points[0].radius / max(splineL, 1e-6)) * (1 / 15) * (1 / frameRate)
                    # This shouldn't have to be in degrees but it looks much better in animation
                    # bMult = degrees(bMult_r)
    
                    bMult = (1 / max(splineL ** .5, 1e-6)) * (1 / 4)
    
                    # print((1 / bMult) * tau) #print wavelength in frames
    
    
                    windFreq1 = bMult * animSpeed
                    windFreq2 = 0.7 * bMult * animSpeed
                    if loopFrames != 0:
                        bMult_l = 1 / (loopFrames / tau)
                        fRatio = max(1, round(windFreq1 / bMult_l))
                        fgRatio = max(1, round(windFreq2 / bMult_l))
                        windFreq1 = fRatio * bMult_l
                        windFreq2 = fgRatio * bMult_l
    
                # For all the points in the curve (less the last) add a bone and name it by the spline it will affect
                nx = 0
                for n in range(0, numPoints, step):
                    oldBone = b
                    boneName = 'bone' + (str(i)).rjust(3, '0') + '.' + (str(n)).rjust(3, '0')
                    b = arm.edit_bones.new(boneName)
                    b.head = s.bezier_points[n].co
                    nx += step
                    nx = min(nx, numPoints)
                    b.tail = s.bezier_points[nx].co
    
                    b.head_radius = s.bezier_points[n].radius
                    b.tail_radius = s.bezier_points[n + 1].radius
                    b.envelope_distance = 0.001
    
                    """
                    # If there are leaves then we need a new vertex group so they will attach to the bone
                    if not leafAnim:
                        if (len(levelCount) > 1) and (i >= levelCount[-2]) and leafObj:
    
                            leafObj.vertex_groups.new(name=boneName)
    
                        elif (len(levelCount) == 1) and leafObj:
    
                            leafObj.vertex_groups.new(name=boneName)
    
                    # If this is first point of the spline then it must be parented to the level above it
                    if n == 0:
                        if parBone:
                            b.parent = arm.edit_bones[parBone]
                    # Otherwise, we need to attach it to the previous bone in the spline
                    else:
                        b.parent = oldBone
                        b.use_connect = True
                    # If there isn't a previous bone then it shouldn't be attached
                    if not oldBone:
                        b.use_connect = False
    
                    # Add the animation to the armature if required
                    if armAnim:
                        # Define all the required parameters of the wind sway by the dimension of the spline
    
                        # a0 = 4 * splineL * (1 - n / (numPoints + 1)) / max(s.bezier_points[n].radius, 1e-6)
    
                        a0 = 2 * (splineL / numPoints) * (1 - n / (numPoints + 1)) / max(s.bezier_points[n].radius, 1e-6)
                        a0 = a0 * min(step, numPoints)
    
                        # a0 = (splineL / numPoints) / max(s.bezier_points[n].radius, 1e-6)
    
                        a1 = (wind / 50) * a0
    
                        a2 = a1 * .65  # (windGust / 50) * a0 + a1 / 2
    
    
                        p = s.bezier_points[nx].co - s.bezier_points[n].co
                        p.normalize()
                        ag = (wind * gust / 50) * a0
                        a3 = -p[0] * ag
                        a4 = p[2] * ag
    
                        a1 = radians(a1)
                        a2 = radians(a2)
                        a3 = radians(a3)
                        a4 = radians(a4)
    
    
                        # wind bending
    
                        if loopFrames == 0:
    
                            swayFreq = gustF * (tau / fps) * frameRate  # animSpeed # .075 # 0.02
    
                        else:
                            swayFreq = 1 / (loopFrames / tau)
    
                        # Prevent tree base from rotating
                        if (boneName == "bone000.000") or (boneName == "bone000.001"):
                            a1 = 0
                            a2 = 0
                            a3 = 0
                            a4 = 0
    
                        # Add new fcurves for each sway as well as the modifiers
    
                        swayX = armOb.animation_data.action.fcurves.new(
    
                                                'pose.bones["' + boneName + '"].rotation_euler', index=0
    
                                                )
                        swayY = armOb.animation_data.action.fcurves.new(
    
                                                'pose.bones["' + boneName + '"].rotation_euler', index=2
    
                        swayXMod1 = swayX.modifiers.new(type='FNGENERATOR')
                        swayXMod2 = swayX.modifiers.new(type='FNGENERATOR')
    
                        swayYMod1 = swayY.modifiers.new(type='FNGENERATOR')
                        swayYMod2 = swayY.modifiers.new(type='FNGENERATOR')
    
                        # Set the parameters for each modifier
                        swayXMod1.amplitude = a1
                        swayXMod1.phase_offset = bxOffset
                        swayXMod1.phase_multiplier = windFreq1
    
                        swayXMod2.amplitude = a2
                        swayXMod2.phase_offset = 0.7 * bxOffset
                        swayXMod2.phase_multiplier = windFreq2
                        swayXMod2.use_additive = True
    
                        swayYMod1.amplitude = a1
                        swayYMod1.phase_offset = byOffset
                        swayYMod1.phase_multiplier = windFreq1
    
                        swayYMod2.amplitude = a2
                        swayYMod2.phase_offset = 0.7 * byOffset
                        swayYMod2.phase_multiplier = windFreq2
                        swayYMod2.use_additive = True
    
    
                        # wind bending
    
                        swayYMod3 = swayY.modifiers.new(type='FNGENERATOR')
                        swayYMod3.amplitude = a3
                        swayYMod3.phase_multiplier = swayFreq
                        swayYMod3.value_offset = .6 * a3
                        swayYMod3.use_additive = True
    
                        swayXMod3 = swayX.modifiers.new(type='FNGENERATOR')
                        swayXMod3.amplitude = a4
                        swayXMod3.phase_multiplier = swayFreq
                        swayXMod3.value_offset = .6 * a4
                        swayXMod3.use_additive = True
    
    
            bonelist = [b.name for b in arm.edit_bones]
            vertexGroups = OrderedDict()
            for i, cp in enumerate(leafP):
                # find leafs parent bone
                leafParent = roundBone(cp.parBone, boneStep[armLevels])
                idx = int(leafParent[4:-4])
                while leafParent not in bonelist:
    
                    # find parent bone of parent bone
    
                    leafParent = splineToBone[idx]
                    idx = int(leafParent[4:-4])
    
                if leafAnim:
                    bname = "leaf" + str(i)
                    b = arm.edit_bones.new(bname)
                    b.head = cp.co
                    b.tail = cp.co + Vector((0, 0, .02))
                    b.envelope_distance = 0.0
                    b.parent = arm.edit_bones[leafParent]
    
    
                    vertexGroups[bname] = [
                                        v.index for v in
                                        leafMesh.vertices[leafVertSize * i:(leafVertSize * i + leafVertSize)]
                                        ]
    
    
                    if armAnim:
                        # Define all the required parameters of the wind sway by the dimension of the spline
                        a1 = wind * .25
                        a1 *= af1
    
                        bMult = (1 / animSpeed) * 6
                        bMult *= 1 / max(af2, .001)
    
                        ofstRand = af3
                        bxOffset = uniform(-ofstRand, ofstRand)
                        byOffset = uniform(-ofstRand, ofstRand)
    
                        # Add new fcurves for each sway as well as the modifiers
    
                        swayX = armOb.animation_data.action.fcurves.new(
    
                                                    'pose.bones["' + bname + '"].rotation_euler', index=0
    
                                                    )
                        swayY = armOb.animation_data.action.fcurves.new(
    
                                                    'pose.bones["' + bname + '"].rotation_euler', index=2
    
                        # Add keyframe so noise works
    
                        swayX.keyframe_points.add(1)
                        swayY.keyframe_points.add(1)
    
                        swayX.keyframe_points[0].co = (0, 0)
                        swayY.keyframe_points[0].co = (0, 0)
    
                        # Add noise modifiers
                        swayXMod = swayX.modifiers.new(type='NOISE')
                        swayYMod = swayY.modifiers.new(type='NOISE')
    
                        if loopFrames != 0:
                            swayXMod.use_restricted_range = True
                            swayXMod.frame_end = loopFrames
                            swayXMod.blend_in = 4
                            swayXMod.blend_out = 4
                            swayYMod.use_restricted_range = True
                            swayYMod.frame_end = loopFrames
                            swayYMod.blend_in = 4
                            swayYMod.blend_out = 4
    
                        swayXMod.scale = bMult
                        swayXMod.strength = a1
                        swayXMod.offset = bxOffset
    
                        swayYMod.scale = bMult
                        swayYMod.strength = a1
                        swayYMod.offset = byOffset
    
                else:
                    if leafParent not in vertexGroups:
                        vertexGroups[leafParent] = []
    
                    vertexGroups[leafParent].extend(
                                            [v.index for v in
                                            leafMesh.vertices[leafVertSize * i:(leafVertSize * i + leafVertSize)]]
                                            )
    
    
            for group in vertexGroups: