Use customdata array storage in the ground object for storing sample
[blender-addons-contrib.git] / object_physics_meadow / blob.py
1 ### BEGIN GPL LICENSE BLOCK #####
2 #
3 #  This program is free software; you can redistribute it and/or
4 #  modify it under the terms of the GNU General Public License
5 #  as published by the Free Software Foundation; either version 2
6 #  of the License, or (at your option) any later version.
7 #
8 #  This program is distributed in the hope that it will be useful,
9 #  but WITHOUT ANY WARRANTY; without even the implied warranty of
10 #  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
11 #  GNU General Public License for more details.
12 #
13 #  You should have received a copy of the GNU General Public License
14 #  along with this program; if not, write to the Free Software Foundation,
15 #  Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
16 #
17 # ##### END GPL LICENSE BLOCK #####
18
19 # <pep8 compliant>
20
21 import bpy
22 from bpy_extras import object_utils
23 from math import *
24 from mathutils import *
25 from mathutils.kdtree import KDTree
26 from itertools import accumulate
27 import random
28
29 from object_physics_meadow import settings as _settings
30 from object_physics_meadow import duplimesh
31 from object_physics_meadow.duplimesh import project_on_ground, interp_weights_face
32 from object_physics_meadow.util import *
33
34 _blob_object_name = "__MeadowBlob__"
35 _blob_object_parent_name = "__MeadowBlobParent__"
36
37 def blob_objects(context):
38     settings = _settings.get(context)
39     blob_group = settings.blob_group(context)
40     # ignore objects with invalid blob index
41     return [ob for ob in blob_group.objects if ob.meadow.blob_index >= 0]
42
43 def blob_group_clear(context):
44     settings = _settings.get(context)
45     blob_group = settings.blob_group(context)
46     
47     if blob_group:
48         delete_objects(context, blob_group.objects)
49
50 def blob_group_assign(context, blobob, test=False):
51     settings = _settings.get(context)
52     blob_group = settings.blob_group(context)
53     
54     if test and blobob in blob_group.objects.values():
55         return
56     
57     blob_group.objects.link(blobob)
58     # NOTE: unsetting the type is important, otherwise gathering templates
59     # a second time will include deleted objects!
60     blobob.meadow.type = 'NONE'
61
62 def blob_group_remove(context, blobob):
63     settings = _settings.get(context)
64     blob_group = settings.blob_group(context)
65     
66     blob_group.objects.unlink(blobob)
67
68 def blob_apply_settings(ob, settings):
69     pass # TODO
70
71 #-----------------------------------------------------------------------
72
73 # 8-class qualitative Brewer color scheme for high-contrast colors
74 # http://colorbrewer2.org/
75 color_schemes = [
76     (228, 26, 28),
77     (55, 126, 184),
78     (77, 175, 74),
79     (152, 78, 163),
80     (255, 127, 0),
81     (255, 255, 51),
82     (166, 86, 40),
83     (247, 129, 191),
84     ]
85
86 def select_color(index):
87     base = color_schemes[hash(str(index)) % len(color_schemes)]
88     return (base[0]/255.0, base[1]/255.0, base[2]/255.0, 1.0)
89
90 def get_blob_material(context):
91     materials = context.blend_data.materials
92     if _blob_object_name in materials:
93         return materials[_blob_object_name]
94     
95     # setup new blob material
96     ma = materials.new(_blob_object_name)
97     ma.use_object_color = True
98     # make the material stand out a bit more using emission
99     ma.emit = 1.0
100     return ma
101
102 def get_blob_parent(context, obmat):
103     ob = context.blend_data.objects.get(_blob_object_parent_name, None)
104     if not ob:
105         ob = object_utils.object_data_add(bpy.context, None, name=_blob_object_parent_name).object
106     # put it in the blob group
107     blob_group_assign(context, ob, test=True)
108     
109     ob.matrix_world = obmat
110     
111     return ob
112
113 # assign sample to a blob, based on distance weighting
114 def assign_blob(blobtree, loc, nor):
115     num_nearest = 4 # number of blobs to consider
116     
117     nearest = blobtree.find_n(loc, num_nearest)
118     
119     totn = len(nearest)
120     if totn == 0:
121         return -1
122     if totn == 1:
123         return nearest[0][1]
124     totdist = fsum(dist for co, index, dist in nearest)
125     if totdist == 0.0:
126         return -1
127     
128     norm = 1.0 / (float(totn-1) * totdist)
129     accum = list(accumulate(((totdist - dist) * norm) ** 8 for co, index, dist in nearest))
130     
131     u = random.uniform(0.0, accum[-1])
132     for a, (co, index, dist) in zip(accum, nearest):
133         if u < a:
134             return index
135     return -1
136
137 def make_blob_object(context, index, loc, samples, display_radius):
138     settings = _settings.get(context)
139     
140     obmat = Matrix.Translation(loc)
141     
142     mesh = duplimesh.make_dupli_mesh(_blob_object_name, obmat, samples, display_radius)
143     mesh.materials.append(get_blob_material(context))
144     
145     ob = object_utils.object_data_add(bpy.context, mesh, operator=None).object
146     # assign the index for mapping
147     ob.meadow.blob_index = index
148     # objects get put at the cursor location by object_utils
149     ob.matrix_world = obmat
150     
151     blob_apply_settings(ob, settings)
152     
153     # assign color and material settings
154     ob.color = select_color(index)
155     ob.show_wire_color = True # XXX this is debatable, could make it an option
156     
157     return ob
158
159 class Blob():
160     def __init__(self, loc, nor, face_index):
161         self.loc = loc
162         self.nor = nor
163         self.face_index = face_index
164         self.samples = []
165
166 # store blobs list in ID datablock as customdata
167 def blobs_to_customprops(data, blobs):
168     data['loc'] = [x for b in blobs for x in b.loc]
169     data['nor'] = [x for b in blobs for x in b.nor]
170     data['face_index'] = [b.face_index for b in blobs]
171     data['samples_len'] = [len(b.samples) for b in blobs]
172     data['samples_loc'] = [x for b in blobs for loc, nor, idx in b.samples for x in loc]
173     data['samples_nor'] = [x for b in blobs for loc, nor, idx in b.samples for x in nor]
174     data['samples_idx'] = [idx for b in blobs for loc, nor, idx in b.samples]
175
176 # load blobs list from ID datablock customdata
177 def blobs_from_customprops(data):
178     blobs = []
179
180     iter_loc = iter(data['loc'])
181     iter_nor = iter(data['nor'])
182     iter_face_index = iter(data['face_index'])
183     iter_samples_len = iter(data['samples_len'])
184     iter_samples_loc = iter(data['samples_loc'])
185     iter_samples_nor = iter(data['samples_nor'])
186     iter_samples_idx = iter(data['samples_idx'])
187
188     while(True):
189         try:
190             loc = (next(iter_loc), next(iter_loc), next(iter_loc))
191             nor = (next(iter_nor), next(iter_nor), next(iter_nor))
192             face_index = next(iter_face_index)
193             
194             samples = []
195             num_samples = next(iter_samples_len)
196             for k in range(num_samples):
197                 sample_loc = (next(iter_samples_loc), next(iter_samples_loc), next(iter_samples_loc))
198                 sample_nor = (next(iter_samples_nor), next(iter_samples_nor), next(iter_samples_nor))
199                 sample_idx = next(iter_samples_idx)
200                 samples.append((sample_loc, sample_nor, sample_idx))
201
202             blob = Blob(loc, nor, face_index)
203             blob.samples = samples
204             blobs.append(blob)
205         except StopIteration:
206             break
207
208     return blobs
209
210 def make_blobs(context, gridob, groundob, samples, display_radius):
211     blob_group_clear(context)
212     blobs = []
213     
214     blobtree = KDTree(len(gridob.data.vertices))
215     for i, v in enumerate(gridob.data.vertices):
216         co = gridob.matrix_world * v.co
217         # note: only using 2D coordinates, otherwise weights get distorted by z offset
218         blobtree.insert((co[0], co[1], 0.0), i)
219     blobtree.balance()
220     
221     for v in gridob.data.vertices:
222         co = gridob.matrix_world * v.co
223         ok, loc, nor, face_index = project_on_ground(groundob, co)
224         blobs.append(Blob(loc, nor, face_index) if ok else None)
225     
226     for loc, nor, face_index in samples:
227         # note: use only 2D coordinates for weighting, z component should be 0
228         index = assign_blob(blobtree, (loc[0], loc[1], 0.0), nor)
229         if index >= 0:
230             blob = blobs[index]
231             if blob:
232                 blob.samples.append((loc, nor, face_index))
233     
234     # common parent empty for blobs
235     blob_parent = get_blob_parent(context, groundob.matrix_world)
236     
237     # preliminary display object
238     # XXX this could be removed eventually, but it's helpful as visual feedback to the user
239     # before creating the actual duplicator blob meshes
240     for index, blob in enumerate(blobs):
241         if blob:
242             samples = [(loc, nor) for loc, nor, _ in blob.samples]
243             ob = make_blob_object(context, index, blob.loc, samples, display_radius)
244             # put it in the blob group
245             blob_group_assign(context, ob)
246             # use parent to keep the outliner clean
247             set_object_parent(ob, blob_parent)
248     
249     blobs_to_customprops(groundob.meadow, blobs)
250
251 #-----------------------------------------------------------------------
252
253 from object_physics_meadow.patch import patch_objects, patch_group_assign
254
255 # select one patch object for each sample based on vertex groups
256 def assign_sample_patches(groundob, blob, patches):
257     vgroups = groundob.vertex_groups
258     faces = groundob.data.tessfaces
259     vertices = groundob.data.vertices
260     
261     used_vgroup_names = set(ob.meadow.density_vgroup_name for ob in patches)
262     
263     vgroup_samples = { vg.name : [] for vg in vgroups }
264     vgroup_samples[""] = [] # samples for unassigned patches
265     for loc, nor, face_index in blob.samples:
266         face = faces[face_index]
267         verts = [vertices[i] for i in face.vertices]
268         assert(len(verts) in {3, 4})
269         
270         # accumulate weights for each vertex group,
271         # by interpolating the face and 
272         fweight, findex = interp_weights_face(tuple(v.co for v in verts[0:4]), loc)
273         weights = [ 0.0 for vg in vgroups ]
274         for v, fac in zip(verts, fweight):
275             for vg in v.groups:
276                 weights[vg.group] += vg.weight * fac
277         
278         def select_vgroup():
279             if not weights:
280                 return None
281             used_vgroups = [(vg, w) for vg, w in zip(vgroups, weights) if vg.name in used_vgroup_names]
282             
283             totweight = sum(w for vg, w in used_vgroups)
284             # using 1.0 as the minimum total weight means we select
285             # the default "non-group" in uncovered areas:
286             # there is a 1.0-totweight chance of selecting no vgroup at all
287             u = random.uniform(0.0, max(totweight, 1.0))
288             for vg, w in used_vgroups:
289                 if u < w:
290                     return vg
291                 u -= w
292             return None
293         
294         vg = select_vgroup()
295         if vg:
296             vgroup_samples[vg.name].append((loc, nor))
297         else:
298             vgroup_samples[""].append((loc, nor))
299     
300     return vgroup_samples
301
302 def setup_blob_duplis(context, groundob, display_radius):
303     blobs = blobs_from_customprops(groundob.meadow)
304
305     groundob.data.calc_tessface()
306     patches = [ob for ob in patch_objects(context) if blobs[ob.meadow.blob_index] is not None]
307     
308     # common parent empty for blobs
309     blob_parent = get_blob_parent(context, groundob.matrix_world)
310     
311     del_patches = set() # patches to delete, keep this separate for iterator validity
312     for blob_index, blob in enumerate(blobs):
313         if blob is None:
314             continue
315         
316         vgroup_samples = assign_sample_patches(groundob, blob, patches)
317         
318         for ob in patches:
319             if ob.meadow.blob_index != blob_index:
320                 continue
321             
322             samples = vgroup_samples.get(ob.meadow.density_vgroup_name, [])
323             if not samples:
324                 del_patches.add(ob)
325                 continue
326             
327             if ob.meadow.use_as_dupli:
328                 # make a duplicator for the patch object
329                 dob = make_blob_object(context, blob_index, blob.loc, samples, display_radius)
330                 # put the duplicator in the patch group,
331                 # so it gets removed together with patch copies
332                 patch_group_assign(context, dob)
333                 # use parent to keep the outliner clean
334                 set_object_parent(dob, blob_parent)
335                 set_object_parent(ob, dob)
336                 
337                 dob.dupli_type = 'FACES'
338                 
339                 # make sure duplis are placed at the sample locations
340                 if ob.meadow.use_centered:
341                     # XXX centering is needed for particle instance modifier (this might be a bug!)
342                     ob.matrix_world = Matrix.Identity(4)
343                 else:
344                     ob.matrix_world = dob.matrix_world
345             else:
346                 # use parent to keep the outliner clean
347                 set_object_parent(ob, blob_parent)
348                 # move to the blob center
349                 ob.matrix_world = Matrix.Translation(blob.loc)
350         
351     # delete unused patch objects
352     delete_objects(context, del_patches)