utils.py

# ##### 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 #####

print("version 3 imported")

import bpy
import time
import copy

from mathutils import *
from math import pi, sin, degrees, radians, atan2, copysign, cos, acos
from math import floor, ceil
from random import random, uniform, seed, choice, getstate, setstate, randint
from bpy.props import *
from collections import deque, OrderedDict

tau = 2 * pi

# 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))

# 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):
        self.spline = spline
        self.p = spline.bezier_points[-1]
        self.curv = curvature
        self.curvV = curvatureV
        self.vertAtt = attractUp
        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

    # 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')
        else:
            return ((self.spline.bezier_points[-1].co - self.spline.bezier_points[-2].co).normalized()).to_track_quat('Z', 'Y')
    # Determine the declination
    def dec(self):
        tempVec = zAxis.copy()
        tempVec.rotate(self.quat())
        return zAxis.angle(tempVec)
    # Update the end of the spline and increment the segment count
    def updateEnd(self):
        self.p = self.spline.bezier_points[-1]
        self.seg += 1

# 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):
        self.co = coords
        self.quat = quat
        self.radiusPar = radiusPar
        self.offset = offset
        self.stemOffset = sOfst
        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):
    if shape == 0:
        return 0.05 + 0.95*ratio #0.2 + 0.8*ratio
    elif shape == 1:
        return 0.2 + 0.8*sin(pi*ratio)
    elif shape == 2:
        return 0.2 + 0.8*sin(0.5*pi*ratio)
    elif shape == 3:
        return 1.0
    elif shape == 4:
        return 0.5 + 0.5*ratio
    elif shape == 5:
        if ratio <= 0.7:
            return 0.05 + 0.95 * ratio/0.7
        else:
            return 0.05 + 0.95 * (1.0 - ratio)/0.3
    elif shape == 6:
        return 1.0 - 0.8*ratio
    elif shape == 7:
        if ratio <= 0.7:
            return 0.5 + 0.5*ratio/0.7
        else:
            return 0.5 + 0.5*(1.0 - ratio)/0.3
    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:
        if (ratio < (1 - pruneWidthPeak)) and (ratio > 0.0):
            return ((ratio/(1 - pruneWidthPeak))**prunePowerHigh)
        elif (ratio >= (1 - pruneWidthPeak)) and (ratio < 1.0):
            return (((1 - ratio)/pruneWidthPeak)**prunePowerLow)
        else:
            return 0.0

    elif shape == 10:
        return 0.5 + 0.5 * (1 - ratio)

# 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)
    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

# Determine the declination from a given quaternion
def declination(quat):
    tempVec = zAxis.copy()
    tempVec.rotate(quat)
    tempVec.normalize()
    return degrees(acos(tempVec.z))

# Determines the angle of upward rotation of a segment due to attractUp
def curveUp(attractUp, quat, curveRes):
    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

# 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

# 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()

# Determine the range of t values along a splines length where child stems are formed
def findChildPoints(stemList, numChild):
    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))]