cohesion.py

"""
in dummy s
out verts v
"""

def setup():

    """
    braap
    """

    from math import pi
    from math import sqrt as sq
    from random import random as rnd
    from collections import defaultdict

    from mathutils import Vector


    TWO_PI = pi * 2

    # PARAMETERS
    _maxForce = 0.2  # Maximum steering force
    _maxSpeed = 2  # Maximum speed
    _desiredSeparation = 10
    _separationCohesionRation = 1.1
    _maxEdgeLen = 5
    width = 399
    height = 200


    class PVector():
        """
        more repetition than I like... but let's see if the rewrite works before
        making it pretty / efficient
        """

        def __init__(self, x=0, y=0):
            self.value = Vector((x, y))

        @staticmethod
        def random():
            return Vector((rnd(), rnd()))

        @property
        def x(self): return self.value.x

        @property
        def y(self): return self.value.y

        def add(self, other, other2=None):
            if other2:
                return other.value + other2.value
            else:
                self.value = self.value + other.value
                return self.value

        def sub(self, other, other2=None):
            if other2:
                return other.value - other2.value
            else:
                self.value = self.value - other.value
                return self.value

        def limit(self, scalar):
            # is this the right analogue?
            self.value.magnitude = scalar
            return self.value

        def mag(self): return self.value.magnitude

        def setMag(self, scalar):
            self.value.magnitude = scalar
            return self.value

        def mult(self, scalar):
            self.value = self.value * scalar
            return self.value

        def div(self, scalar):
            self.value = self.value / scalar
            return self.value

        def dist(self, other, other2=None):
            if other2:
                return (other.value - other2.value).length
            else:
                return (self.value - other.value).length

        def __repr__(self):
            return repr(self.value)


    class Node():
        # position  $ Vec
        # velocity   $ Vec
        # acceleration  $ vec

        # float maxForce
        # float maxSpeed

        def __init__(self, x, y, mF, mS):
            self.acceleration = PVector()
            self.velocity = PVector.random()
            self.position = PVector(x, y)
            self.maxSpeed = mF
            self.maxForce = mS

        def applyForce(self, force):
            self.acceleration.add(force)

        def update(self):
            velocity.add(acceleration)
            velocity.limit(maxSpeed)
            position.add(velocity)
            acceleration.mult(0)

        def seek(self, target):
            desired = PVector.sub(target, self.position)
            desired.setMag(self.maxSpeed)
            steer = PVector.sub(desired, self.velocity)
            steer.limit(self.maxForce)
            return steer


    class DifferentialLine():

        def __init__(self, mF, mS, dS, sCr, eL):
            self.nodes = [] # NodeArray()
            self.maxSpeed = mF
            self.maxForce = mS
            self.desiredSeparation = dS
            self.sq_desiredSeparation = sq(self.desiredSeparation)
            self.separationCohesionRation = sCr
            self.maxEdgeLen = eL

        def addNode(self, n):
            self.nodes.append(n)

        def addNodeAt(self, n, index):
            self.nodes.insert(index, n)   # maybe list is not the appropriate structure

        def run(self):
            self.differentiate()
            self.growth()

        def growth(self):
            for i in range(len(self.nodes)-1):
                n1 = self.nodes[i]
                n2 = self.nodes[i+1]
                d = PVector.dist(n1.position, n2.position)

                # Can add more rules for inserting nodes
                if (d > self.maxEdgeLen): 
                    index = self.nodes.index(n2) # indexof? .. is not i+1 ?
                    middleNode = PVector.add(n1.position, n2.position).div(2)
                    self.addNodeAt(Node(middleNode.x, middleNode.y, self.maxForce, self.maxSpeed), index)


        def differentiate(self):
            separationForces = self.getSeparationForces()
            cohesionForces = self.getEdgeCohesionForces()

            for i in range(len(self.nodes)):
                separation = separationForces[i]  # a PVector
                cohesion = cohesionForces[i]  # a PVector

                separation.mult(self.separationCohesionRation)

                self.nodes[i].applyForce(separation)
                self.nodes[i].applyForce(cohesion)
                self.nodes[i].update()


        def getSeparationForces(self):  # PVector[] 
            n = len(self.nodes)
            separateForces = []  # prefil with n elements ?!
            nearNodes = defaultdict(int) # int[] nearNodes = new int[n];  (prefil with n elements?)

            separateForces = [PVector() for i in range(n)]

            for i in range(n):
                nodei = self.nodes[i]
                for j in range(i+1, n):
                    nodej = self.nodes[j]

                    forceij = self.getSeparationForce(nodei, nodej)
                    if forceij.mag() > 0:
                        separateForces[i].add(forceij)        
                        separateForces[j].sub(forceij)
                        nearNodes[i] += 1
                        nearNodes[j] += 1

                if nearNodes[i] > 0:
                    separateForces[i].div(nearNodes[i])

                if separateForces[i].mag() > 0:
                    separateForces[i].setMag(self.maxSpeed)
                    separateForces[i].sub(self.nodes[i].velocity)
                    separateForces[i].limit(self.maxForce)

            return separateForces

        def getSeparationForce(self, n1, n2):
            steer = PVector()
            sq_d = sq(n2.position.x-n1.position.x) + sq(n2.position.y-n1.position.y)
            if (sq_d > 0 and sq_d < sq_desiredSeparation):
                diff = PVector.sub(n1.position, n2.position)
                diff.normalize()
                diff.div(sq(sq_d)) #  //Weight by distacne
                steer.add(diff)

            return steer


        def getEdgeCohesionForces(self):
            n = len(self.nodes)
            nodes = self.nodes
            cohesionForces = list(range(n))

            for i in range(n):
                _sum = PVector()
                if (i != 0 and i != n-1):
                    _sum.add(nodes[i-1].position).add(nodes[i+1].position)
                elif i == 0:
                    _sum.add(nodes[n-1].position).add(nodes[i+1].position)
                elif i == n-1:
                    _sum.add(nodes[i-1].position).add(nodes[0].position)

                _sum.div(2)
                cohesionForces.append(nodes[i].seek(_sum))

            return cohesionForces


        def renderShape(self):
            return [(n.position.x, n.position.y, 0.0) for n in self.nodes]


    _diff_line = DifferentialLine(_maxForce, _maxSpeed, _desiredSeparation, _separationCohesionRation, _maxEdgeLen);

    nodesStart = 20
    angInc = TWO_PI / nodesStart
    rayStart = 10

    num_nodes = int(TWP_PI/ angInc)
    for a in range(num_nodes):
        theta = a * angInc

        x = width/2 + cos(theta) * rayStart
        y = height/2 + sin(theta) * rayStart
    
        _diff_line.addNode(Node(x, y, _maxForce, _maxSpeed))

# def draw() {
for i in range(40):

    _diff_line.run()
    _diff_line.renderLine()
    verts = _diff_line.renderShape()

# draw verts