utils.py

# -*- coding: utf-8 -*-
# ##### 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

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

# 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