lab/treegen/grow.py

# This code is called when instances of this SOP cook.
#scaleFactor = node.parm("scaleFactor").eval()

import random
import hou



node = hou.pwd()
geo = node.geometry()

# map user input
GENERATIONS = hou.ch('generations')
TROPISM = hou.Vector3(hou.parmTuple("tropDir").eval())
TROPISMFACTOR = hou.ch('tropStrength')
GRAVITYFACTOR = hou.ch('gFactor')
SEED = hou.ch('seed')

#hou.ch("falloff")

# the "seed"
# this is the first point from which the tree will grow.
point = geo.createPoint()
pointers = [point]

# init the point attributes to store the values in, as these must exist.
# As we start, the first point's normal will point upward to grow in that direction.
geo.addAttrib(hou.attribType.Point, "N", hou.Vector3([0, 1, 0]))
geo.addAttrib(hou.attribType.Point, "Branch", 0)
geo.addAttrib(hou.attribType.Point, "Generation", 0)
geo.addAttrib(hou.attribType.Point, "BranchPoint", 0)
geo.addAttrib(hou.attribType.Point, 'NormalizedPosition', 0.0)
geo.addAttrib(hou.attribType.Point, 'Diameter', 1.0)

# Just in case, if we need groups, this is how to do it.
#defaultGrp=geo.createPointGroup('name')
#defaultGrp.add(point)


def duplicate_point(p):
    # hmm. maybe there is sth like this already, but i missed it...
    # this duplicates a point and all it's attributes.
    dupe=geo.createPoint()
    dupe.setPosition(p.position())
    allAttr=geo.pointAttribs()
    for i in allAttr:
        holder = None
        holder = p.attribValue(i)
        dupe.setAttribValue(i, holder)
    return dupe


def tropism(pt, vector, factor):
    '''Move point towards another point, as in growing towards a light source (sun)
       or a food source (roots)
    '''
    P = pt.position()
    vector = vector.normalized()
    pt.setPosition(P + (1*factor) * vector)


def move_along_normal(pt, factor):
    P = pt.position()
    #P = hou.Vector3(P)
    # be careful - this returns a value you cant simply add to P, must convert it to a hou.Vector3
    N = pt.attribValue("N")
    N = hou.Vector3(N)
    pt.setPosition(P + (1*factor) * N)


def gravity(pt, factor):
    '''Convenience function to move down a point, multiplied by diameter.'''
    P = pt.position()
    # The downforce D
    D = (1 - pt.attribValue('Diameter'))*factor
    G = hou.Vector3([0, D, 0])
    pt.setPosition(P - G)


def create_random_vector():
    rx = (random.random()-0.5)*2
    ry = (random.random()-0.5)*2
    rz = (random.random()-0.5)*2
    return hou.Vector3([rx, ry, rz])


def add_noise_to_normal(pt, factor):
    N = pt.attribValue("N")
    N = hou.Vector3(N)
    rx = (random.random()-0.5)*2
    ry = (random.random()-0.5)*2
    rz = (random.random()-0.5)*2
    rvec = hou.Vector3([rx*factor, ry*factor, rz*factor])
    newN = rvec + 1 * N
    pt.setAttribValue("N", newN)
    return pt


def rotate_normal(pt, factor, seed):
    N = pt.attribValue("N")
    N = hou.Vector3(N)
    
    x = (random.random()-0.5)*2
    y = (random.random()-0.5)*2
    z = (random.random()-0.5)*2
    print 'p random x {0} y {1} z {2}'.format(x, y, z)

    x = (hou.hmath.rand(seed+12*43)-0.5)*2
    y = (hou.hmath.rand(seed+2+2*3)-0.5)*2
    z = (hou.hmath.rand(seed+33*2)-0.5)*2
    print 'h random x {0} y {1} z {2}'.format(x, y, z)
 
    rvec = hou.Vector3([x*factor, y*factor, z*factor])
    newN = rvec + 1 * N
    pt.setAttribValue("N", newN)
    return pt


def get_points_by_attrib_value(attr, val):
    plist = []
    for p in geo.iterPoints():
        if p.attribValue(attr) == val:
            plist.append(p)
    return plist


def get_highest_numeric_attr_val(points, attr):
    val = 0.0
    for p in points:
        v = float(p.attribValue(attr))
        if v > val:
            val = v
    return val



branchNum = 0


def step(pointers, stepseed):
    """
    step through all head pointers (aka growing branches)
    """
    updatedPointers=[]
    for pointer in pointers:
        #print 'pointer is point w/ # ' + str(p.number())
        newP = duplicate_point(pointer)
        #print 'copy is point w/ # ' + str(newP.number())
        updatedPointers.append(newP)
        newP.setAttribValue('BranchPoint', newP.attribValue('BranchPoint')+1)
        # Should we branch?
        # TODO allow multiple branches at once
        rnd = random.random()
        rnd = hou.hmath.rand(stepseed + 11)
        # BRANCH OFF
        if rnd > 0.8:
            # Dirty, find a better way soon!
            global branchNum
            branchNum += 1
            global treeInfo
            #treeInfo['totalbranches']+=1
            #print 'branch'
            branchP = duplicate_point(pointer)
            branchP.setAttribValue('BranchPoint', 0)
            b = branchP.attribValue('Branch')
            branchP.setAttribValue('Branch', branchNum)
            #branchP.setAttribValue('Branch',b+1)
            g = branchP.attribValue('Generation')
            branchP.setAttribValue('Generation', g + 1)
            updatedPointers.append(branchP)
            stepseed += 1
        newP = rotate_normal(newP, 1.7, stepseed)
        # now move along normal
        move_along_normal(newP, 1)
        v = hou.Vector3(0, 1, 1)
        tropism(newP, TROPISM, TROPISMFACTOR)
        gravity(newP, GRAVITYFACTOR)
        newP.setAttribValue('Diameter', newP.attribValue('Diameter')*0.7)
        stepseed += 1
        print 'seed', stepseed
    return updatedPointers

    
x = step(pointers, 1)

seed = SEED
for arsch in range(GENERATIONS):
    #seed=1
    x = step(x, seed)
    seed += 1



# until there is a better way, 'postprocess' the tree to gain info that is not present at runtime.


def postprocess():
    tInfo = {}
    '''
    for n in range(0,branchNum):
        for p in geo.iterPoints():
            #print p
            print p.attribValue('BranchPoint')
            #if p.attribValue('Branch') == n:
            #    print p.attribValue('BranchPoint')
    '''
    for p in geo.iterPoints():
        curBranch = p.attribValue('Branch')
        tInfo['branch'] = curBranch


print '[[[[[[[[[[***]]]]]]]]]]'
totalBranches = int(get_highest_numeric_attr_val(geo.iterPoints(), 'Branch'))

for i in range(0,totalBranches):
    # go through all branches
    max = 0
    for p in get_points_by_attrib_value('Branch', i):
        v = p.attribValue('BranchPoint')
        if v > max:
            max = v
    for p in get_points_by_attrib_value('Branch', i):
        bp = p.attribValue('BranchPoint')
        #print max
        #print bp
        try:
            relpos = float(bp)/float(max)
        except:
            relpos = 0
        #p.setAttribValue('NormalizedPosition',1)
        p.setAttribValue('NormalizedPosition', float(relpos))

    #print 'branch {0} points {1}'.format(i,max)
    


# TODO and thoughts
# in addition to the branch 'id' we need a distance of each point of a branch towards it's parent.
# useful: how many 'levels' are we away from the main parent? > point needs to have parent id, increment on branch
# roots 
# DIAMETER should decrease by a step after each branch