f9d2c6baf6cfdfe16c12a4a0ac256dab38abff8a
[blender-addons-contrib.git] / add_mesh_space_tree / sca.py
1 # ##### BEGIN GPL LICENSE BLOCK #####
2 #
3 #  SCA Tree Generator, a Blender addon
4 #  (c) 2013 Michel J. Anders (varkenvarken)
5 #
6 #  This program is free software; you can redistribute it and/or
7 #  modify it under the terms of the GNU General Public License
8 #  as published by the Free Software Foundation; either version 2
9 #  of the License, or (at your option) any later version.
10 #
11 #  This program is distributed in the hope that it will be useful,
12 #  but WITHOUT ANY WARRANTY; without even the implied warranty of
13 #  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
14 #  GNU General Public License for more details.
15 #
16 #  You should have received a copy of the GNU General Public License
17 #  along with this program; if not, write to the Free Software Foundation,
18 #  Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
19 #
20 # ##### END GPL LICENSE BLOCK #####
21
22 # <pep8 compliant>
23
24 from collections import defaultdict as dd
25 from random import random, seed, expovariate
26 from math import sqrt, pow, sin, cos
27 from functools import partial
28
29 from mathutils import Vector
30
31 from .kdtree import Tree
32
33
34 class Branchpoint:
35     def __init__(self, p, parent):
36         self.v = Vector(p)
37         self.parent = parent
38         self.connections = 0
39         self.apex = None
40         self.shoot = None
41         self.index = None
42
43
44 def sphere(r, p):
45     r2 = r * r
46     while True:
47         x = (random() * 2 - 1) * r
48         y = (random() * 2 - 1) * r
49         z = (random() * 2 - 1) * r
50         if x * x + y * y + z * z <= r2:
51             yield p + Vector((x, y, z))
52
53
54 class SCA:
55     def __init__(self, NENDPOINTS=100, d=0.3, NBP=2000, KILLDIST=5, INFLUENCE=15, SEED=42, volume=partial(sphere, 5, Vector((0, 0, 8))), TROPISM=0.0, exclude=lambda p: False,
56         startingpoints=[]):
57         seed(SEED)
58         self.d = d
59         self.NBP = NBP
60         self.KILLDIST = KILLDIST * KILLDIST * d * d
61         self.INFLUENCE = INFLUENCE * INFLUENCE * d * d
62         self.TROPISM = TROPISM
63         self.exclude = exclude
64         self.endpoints = []
65
66         self.volumepoint = volume()
67         for i in range(NENDPOINTS):
68             self.endpoints.append(next(self.volumepoint))
69
70         self.branchpoints = [Branchpoint((0, 0, 0), None)] if len(startingpoints) == 0 else startingpoints
71
72     def iterate(self, newendpointsper1000=0, maxtime=0.0):  # maxtime still ignored for now
73
74         endpointsadded = 0.0
75         niterations = 0.0
76         newendpointsper1000 /= 1000.0
77         t = expovariate(newendpointsper1000) if newendpointsper1000 > 0.0 else 1  # time to the first new 'endpoint add event'
78
79         while self.NBP > 0 and (len(self.endpoints) > 0):
80             self.NBP -= 1
81             closestsendpoints = dd(list)
82
83             kill = set()
84
85             for ei, e in enumerate(self.endpoints):
86                 distance = None
87                 closestbp = None
88                 for bi, b in enumerate(self.branchpoints):
89                     ddd = b.v - e
90                     ddd = ddd.dot(ddd)
91                     if ddd < self.KILLDIST:
92                         kill.add(ei)
93                     elif (ddd < self.INFLUENCE and b.shoot is None) and ((distance is None) or (ddd < distance)):
94                         closestbp = bi
95                         distance = ddd
96                 if not (closestbp is None):
97                     closestsendpoints[closestbp].append(ei)
98
99             if len(closestsendpoints) < 1:
100                 break
101
102             for bi in closestsendpoints:
103                 sd = Vector((0, 0, 0))
104                 n = 0
105                 for ei in closestsendpoints[bi]:
106                     dv = self.branchpoints[bi].v - self.endpoints[ei]
107                     ll = sqrt(dv.dot(dv))
108                     sd -= dv / ll
109                     n += 1
110                 sd /= n
111                 ll = sqrt(sd.dot(sd))
112                 # don't know if this assumption is true:
113                 # if the unnormalised direction is very small, the endpoints are nearly coplanar/colinear and at roughly the same distance
114                 # so no endpoints will be killed and we might end up adding the same branch again and again
115                 if ll < 1e-3:
116                     #print('SD very small')
117                     continue
118
119                 sd /= ll
120                 sd[2] += self.TROPISM
121                 ll = sqrt(sd.dot(sd))
122                 sd /= ll
123
124                 newp = self.branchpoints[bi].v + sd * self.d
125                 # the assumption we made earlier is not suffucient to prevent adding the same branch so we need an expensive check:
126                 tooclose = False
127                 for dbi in self.branchpoints:
128                     dddd = newp - dbi.v
129                     if dddd.dot(dddd) < 1e-3:
130                         #print('BP to close to another')
131                         tooclose = True
132                 if tooclose:
133                     continue
134
135                 if not self.exclude(newp):
136                     bp = Branchpoint(newp, bi)
137                     self.branchpoints.append(bp)
138                     nbpi = len(self.branchpoints) - 1
139                     bp = self.branchpoints[bi]
140                     bp.connections += 1
141                     if bp.apex is None:
142                         bp.apex = nbpi
143                     else:
144                         bp.shoot = nbpi
145                     while not (bp.parent is None):
146                         bp = self.branchpoints[bp.parent]
147                         bp.connections += 1
148
149             self.endpoints = [ep for ei, ep in enumerate(self.endpoints) if not(ei in kill)]
150
151             if newendpointsper1000 > 0.0:
152                 # generate new endpoints with a poisson process
153                 # when we first arrive here, t already hold the time to the first event
154                 niterations += 1
155                 while t < niterations:  # we keep on adding endpoints as long as the next event still happens within this iteration
156                     self.endpoints.append(next(self.volumepoint))
157                     endpointsadded += 1
158                     t += expovariate(newendpointsper1000)  # time to new 'endpoint add event'
159
160         #if newendpointsper1000 > 0.0:
161             #print("newendpoints/iteration %.3f, actual %.3f in %5.f iterations"%(newendpointsper1000,endpointsadded/niterations,niterations))
162
163     def iterate2(self, newendpointsper1000=0, maxtime=0.0):  # maxtime still ignored for now
164         """iterate using a kdtree fr the branchpoints"""
165         endpointsadded = 0.0
166         niterations = 0.0
167         newendpointsper1000 /= 1000.0
168         t = expovariate(newendpointsper1000) if newendpointsper1000 > 0.0 else 1  # time to the first new 'endpoint add event'
169
170         tree = Tree(3)
171         for bpi, bp in enumerate(self.branchpoints):
172             bp.index = bpi
173             tree.insert(bp.v, bp)
174
175         while self.NBP > 0 and (len(self.endpoints) > 0):
176             self.NBP -= 1
177             closestsendpoints = dd(list)
178
179             kill = set()
180
181             for ei, e in enumerate(self.endpoints):
182                 distance = None
183                 closestbp, distance = tree.nearest(e, checkempty=True)  # ignore forks, they have .data set to None
184                 if closestbp is not None:
185                     if distance < self.KILLDIST:
186                         kill.add(ei)
187                     elif distance < self.INFLUENCE:
188                         closestsendpoints[closestbp.data.index].append(ei)
189
190             if len(closestsendpoints) < 1:
191                 break
192
193             for bi in closestsendpoints:
194                 sd = Vector((0, 0, 0))
195                 n = 0
196                 for ei in closestsendpoints[bi]:
197                     dv = self.branchpoints[bi].v - self.endpoints[ei]
198                     dv.normalize()
199                     sd -= dv
200                     n += 1
201                 sd /= n
202                 ll = sd.length_squared
203                 # don't know if this assumption is true:
204                 # if the unnormalised direction is very small, the endpoints are nearly coplanar/colinear and at roughly the same distance
205                 # so no endpoints will be killed and we might end up adding the same branch again and again
206                 if ll < 1e-3:
207                     #print('SD very small')
208                     continue
209
210                 sd /= ll
211                 sd[2] += self.TROPISM
212                 sd.normalize()
213
214                 newp = self.branchpoints[bi].v + sd * self.d
215                 # the assumption we made earlier is not suffucient to prevent adding the same branch so we need an expensive check:
216                 _, dddd = tree.nearest(newp)  # no checkempty here, we want to check for forks as well
217                 if dddd < 1e-3:
218                         #print('BP to close to another')
219                         continue
220
221                 if not self.exclude(newp):
222                     bp = Branchpoint(newp, bi)
223                     self.branchpoints.append(bp)
224                     nbpi = len(self.branchpoints) - 1
225                     bp.index = nbpi
226                     tree.insert(bp.v, bp)
227                     bp = self.branchpoints[bi]
228                     bp.connections += 1
229                     if bp.apex is None:
230                         bp.apex = nbpi
231                     else:
232                         bp.shoot = nbpi
233                         node, _ = tree.nearest(bp.v)
234                         node.data = None  # signal that this is a fork so that we might ignore it when searching for nearest neighbors
235                     while not (bp.parent is None):
236                         bp = self.branchpoints[bp.parent]
237                         bp.connections += 1
238
239             self.endpoints = [ep for ei, ep in enumerate(self.endpoints) if not(ei in kill)]
240
241             if newendpointsper1000 > 0.0:
242                 # generate new endpoints with a poisson process
243                 # when we first arrive here, t already hold the time to the first event
244                 niterations += 1
245                 while t < niterations:  # we keep on adding endpoints as long as the next event still happens within this iteration
246                     self.endpoints.append(next(self.volumepoint))
247                     endpointsadded += 1
248                     t += expovariate(newendpointsper1000)  # time to new 'endpoint add event'