Slope rotation feature: Dupli rotation gets aligned to the ground
[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 mathutils.interpolate import poly_3d_calc
27 from itertools import accumulate
28 import random
29
30 from object_physics_meadow import settings as _settings
31 from object_physics_meadow import duplimesh
32 from object_physics_meadow.duplimesh import project_on_ground
33 from object_physics_meadow.util import *
34 from object_physics_meadow import progress
35
36 _blob_object_name = "__MeadowBlob__"
37 _sampleviz_parent_name = "Meadow_SampleViz_Root"
38 _duplicator_parent_name = "Meadow_Duplicators_Root"
39
40 def blob_objects(context):
41     settings = _settings.get(context)
42     blob_group = settings.blob_group(context)
43     # ignore objects with invalid blob index
44     return [ob for ob in blob_group.objects if ob.meadow.blob_index >= 0]
45
46 def blob_group_clear(context):
47     settings = _settings.get(context)
48     blob_group = settings.blob_group(context)
49     
50     if blob_group:
51         delete_objects(context, blob_group.objects)
52
53 def blob_group_assign(context, blobob, test=False):
54     settings = _settings.get(context)
55     blob_group = settings.blob_group(context)
56     
57     if test and blobob in blob_group.objects.values():
58         return
59     
60     blob_group.objects.link(blobob)
61     # NOTE: unsetting the type is important, otherwise gathering templates
62     # a second time will include deleted objects!
63     blobob.meadow.type = 'NONE'
64
65 def blob_group_remove(context, blobob):
66     settings = _settings.get(context)
67     blob_group = settings.blob_group(context)
68     
69     blob_group.objects.unlink(blobob)
70
71 def blob_apply_settings(ob, settings):
72     pass # TODO
73
74 #-----------------------------------------------------------------------
75
76 # 8-class qualitative Brewer color scheme for high-contrast colors
77 # http://colorbrewer2.org/
78 color_schemes = [
79     (228, 26, 28),
80     (55, 126, 184),
81     (77, 175, 74),
82     (152, 78, 163),
83     (255, 127, 0),
84     (255, 255, 51),
85     (166, 86, 40),
86     (247, 129, 191),
87     ]
88
89 def select_color(index):
90     base = color_schemes[hash(str(index)) % len(color_schemes)]
91     return (base[0]/255.0, base[1]/255.0, base[2]/255.0, 1.0)
92
93 def get_blob_material(context):
94     materials = context.blend_data.materials
95     if _blob_object_name in materials:
96         return materials[_blob_object_name]
97     
98     # setup new blob material
99     ma = materials.new(_blob_object_name)
100     ma.use_object_color = True
101     # make the material stand out a bit more using emission
102     ma.emit = 1.0
103     return ma
104
105 def get_blob_parent(context, obmat, name):
106     ob = context.blend_data.objects.get(name, None)
107     if not ob:
108         ob = object_utils.object_data_add(bpy.context, None, name=name).object
109     # put it in the blob group
110     blob_group_assign(context, ob, test=True)
111     
112     ob.matrix_world = obmat
113     
114     return ob
115
116 # assign sample to a blob, based on distance weighting
117 def assign_blob(blobtree, loc, nor):
118     num_nearest = 4 # number of blobs to consider
119     
120     nearest = blobtree.find_n(loc, num_nearest)
121     
122     totn = len(nearest)
123     if totn == 0:
124         return -1
125     if totn == 1:
126         return nearest[0][1]
127     totdist = fsum(dist for co, index, dist in nearest)
128     if totdist == 0.0:
129         return -1
130     
131     norm = 1.0 / (float(totn-1) * totdist)
132     accum = list(accumulate(((totdist - dist) * norm) ** 8 for co, index, dist in nearest))
133     
134     u = random.uniform(0.0, accum[-1])
135     for a, (co, index, dist) in zip(accum, nearest):
136         if u < a:
137             return index
138     return -1
139
140 def make_blob_object(context, index, loc, samples, display_radius):
141     settings = _settings.get(context)
142     
143     obmat = Matrix.Translation(loc)
144     
145     mesh = duplimesh.make_dupli_mesh(_blob_object_name, obmat, samples, display_radius)
146     mesh.materials.append(get_blob_material(context))
147     
148     ob = object_utils.object_data_add(bpy.context, mesh, operator=None).object
149     # assign the index for mapping
150     ob.meadow.blob_index = index
151     # objects get put at the cursor location by object_utils
152     ob.matrix_world = obmat
153     
154     blob_apply_settings(ob, settings)
155     
156     # assign color and material settings
157     ob.color = select_color(index)
158     ob.show_wire_color = True # XXX this is debatable, could make it an option
159     
160     return ob
161
162 class Blob():
163     def __init__(self, loc, nor, poly_index):
164         self.loc = loc
165         self.nor = nor
166         self.poly_index = poly_index
167         self.samples = []
168     
169     def add_sample(self, loc, nor, poly, verts, weights):
170         self.samples.append((loc, nor, poly, verts, weights))
171     
172     # note: Vector instances cannot be pickled directly,
173     # therefore define own pickle methods here
174     def __getstate__(self):
175         return self.loc[:], self.nor[:], self.poly_index, [(sloc[:], snor[:], spoly, sverts, sweights) for sloc, snor, spoly, sverts, sweights in self.samples]
176     
177     def __setstate__(self, state):
178         self.loc = Vector(state[0])
179         self.nor = Vector(state[1])
180         self.poly_index = state[2]
181         self.samples = [(Vector(sloc), Vector(snor), spoly, sverts, sweights) for sloc, snor, spoly, sverts, sweights in state[3]]
182
183 # store blobs list in ID datablock as customdata
184 def blobs_to_customprops(data, blobs):
185     import pickle, array
186     B = pickle.dumps(blobs)
187     pad = (4 - len(B)) % 4
188     A = array.array('i', B + b'\x00' * pad)
189     data['blobs'] = A.tolist()
190
191 # load blobs list from ID datablock customdata
192 def blobs_from_customprops(data):
193     import pickle, array
194     A = array.array('i', data['blobs'])
195     blobs = pickle.loads(A.tobytes())
196     return blobs
197
198 def object_has_blob_data(ob):
199     return ob.meadow.get('blobs') is not None
200
201 def object_free_blob_data(ob):
202     if ob.meadow.get('blobs') is not None:
203         del ob.meadow['blobs']
204
205 def make_blob_visualizer(context, groundob, blobs, display_radius, hide, hide_render=True):
206     slope_factor = 1.0
207
208     # common parent empty for blobs
209     blob_parent = get_blob_parent(context, groundob.matrix_world, _sampleviz_parent_name)
210     blob_parent.hide = hide
211     blob_parent.hide_render = hide_render
212
213     # preliminary display object
214     # XXX this could be removed eventually, but it's helpful as visual feedback to the user
215     # before creating the actual duplicator blob meshes
216     with progress.ProgressContext("Sample Visualization", 0, len(blobs)):
217         for index, blob in enumerate(blobs):
218             progress.progress_add(1)
219             if blob:
220                 # generator for duplimesh, yielding (loc, rot) pairs
221                 def mesh_samples():
222                     up = Vector((0,0,1))
223                     for loc, nor, _, _, _ in blob.samples:
224                         mat = (slope_factor * up.rotation_difference(nor)).to_matrix()
225                         mat.resize_4x4()
226                         yield loc, mat
227                 ob = make_blob_object(context, index, blob.loc, mesh_samples(), display_radius)
228                 # put it in the blob group
229                 blob_group_assign(context, ob)
230                 # use parent to keep the outliner clean
231                 set_object_parent(ob, blob_parent)
232                 # apply layers
233                 if groundob.meadow.use_layers:
234                     ob.layers = groundob.meadow.layers
235                 
236                 ob.hide = hide
237                 ob.hide_render = hide_render
238                 # make children unselectable by default
239                 ob.hide_select = True
240
241 def make_blobs(context, gridob, groundob, samples2D, display_radius):
242     blob_group_clear(context)
243     blobs = []
244     
245     imat = groundob.matrix_world.inverted()
246
247     blobtree = KDTree(len(gridob.data.vertices))
248     for i, v in enumerate(gridob.data.vertices):
249         co = gridob.matrix_world * v.co
250         # note: only using 2D coordinates, otherwise weights get distorted by z offset
251         blobtree.insert((co[0], co[1], 0.0), i)
252     blobtree.balance()
253     
254     for v in gridob.data.vertices:
255         co = gridob.matrix_world * v.co
256         ok, loc, nor, poly_index = project_on_ground(groundob, co)
257         blobs.append(Blob(loc, nor, poly_index) if ok else None)
258     
259     with progress.ProgressContext("Grouping Samples", 0, len(samples2D)):
260         mpolys = groundob.data.polygons
261         mverts = groundob.data.vertices
262         for xy in samples2D:
263             progress.progress_add(1)
264
265             # note: use only 2D coordinates for weighting, z component should be 0
266             index = assign_blob(blobtree, (xy[0], xy[1], 0.0), nor)
267             if index < 0:
268                 continue
269             blob = blobs[index]
270             if blob is None:
271                 continue
272             
273             # project samples onto the ground object
274             ok, sloc, snor, spoly = project_on_ground(groundob, xy[0:2]+(0,))
275             if not ok:
276                 continue
277             
278             # calculate barycentric vertex weights on the poly
279             poly = mpolys[spoly]
280             sverts = list(poly.vertices)
281             # note: coordinate space has to be consistent, use sloc in object space
282             sweights = poly_3d_calc(tuple(mverts[i].co for i in sverts), imat * sloc)
283
284             blob.add_sample(sloc, snor, spoly, sverts, sweights)
285     
286     blobs_to_customprops(groundob.meadow, blobs)
287
288     make_blob_visualizer(context, groundob, blobs, display_radius, hide=True)
289
290 #-----------------------------------------------------------------------
291
292 from object_physics_meadow.patch import patch_objects, patch_group_assign
293
294 # select one patch object for each sample based on vertex groups
295 def assign_sample_patches(groundob, blob, patches):
296     vgroups = groundob.vertex_groups
297     mverts = groundob.data.vertices
298     
299     used_vgroup_names = set(ob.meadow.density_vgroup_name for ob in patches)
300     
301     vgroup_samples = { vg.name : [] for vg in vgroups }
302     vgroup_samples[""] = [] # samples for unassigned patches
303     for sloc, snor, spoly, sverts, sweights in blob.samples:
304         verts = [mverts[i] for i in sverts]
305         # accumulate weights for each vertex group by interpolating the poly
306         weights = [ 0.0 for vg in vgroups ]
307         for v, fac in zip(verts, sweights):
308             for vg in v.groups:
309                 weights[vg.group] += vg.weight * fac
310         
311         def select_vgroup():
312             if not weights:
313                 return None
314             used_vgroups = [(vg, w) for vg, w in zip(vgroups, weights) if vg.name in used_vgroup_names]
315             
316             totweight = sum(w for vg, w in used_vgroups)
317             # using 1.0 as the minimum total weight means we select
318             # the default "non-group" in uncovered areas:
319             # there is a 1.0-totweight chance of selecting no vgroup at all
320             u = random.uniform(0.0, max(totweight, 1.0))
321             for vg, w in used_vgroups:
322                 if u < w:
323                     return vg
324                 u -= w
325             return None
326         
327         vg = select_vgroup()
328         if vg:
329             vgroup_samples[vg.name].append((sloc, snor))
330         else:
331             vgroup_samples[""].append((sloc, snor))
332     
333     return vgroup_samples
334
335 def setup_blob_duplis(context, groundob, display_radius):
336     slope_rotation = groundob.meadow.slope_rotation
337
338     blobs = blobs_from_customprops(groundob.meadow)
339
340     patches = [ob for ob in patch_objects(context) if blobs[ob.meadow.blob_index] is not None]
341     for ob in patches:
342         # make unselectable by default
343         ob.hide_select = True
344     
345     # common parent empty for blobs
346     blob_parent = get_blob_parent(context, groundob.matrix_world, _duplicator_parent_name)
347
348     del_patches = set() # patches to delete, keep this separate for iterator validity
349     for blob_index, blob in enumerate(blobs):
350         if blob is None:
351             continue
352         
353         vgroup_samples = assign_sample_patches(groundob, blob, patches)
354         
355         for ob in patches:
356             if ob.meadow.blob_index != blob_index:
357                 continue
358             
359             samples = vgroup_samples.get(ob.meadow.density_vgroup_name, [])
360             if not samples:
361                 del_patches.add(ob)
362                 continue
363             
364             # generator for duplimesh, yielding (loc, rot) pairs
365             def mesh_samples():
366                 up = Vector((0,0,1))
367                 for loc, nor in samples:
368                     mat = (slope_rotation * up.rotation_difference(nor)).to_matrix()
369                     mat.resize_4x4()
370                     yield loc, mat
371
372             if ob.meadow.use_as_dupli:
373                 # make a duplicator for the patch object
374                 dob = make_blob_object(context, blob_index, blob.loc, mesh_samples(), display_radius)
375                 # put the duplicator in the patch group,
376                 # so it gets removed together with patch copies
377                 patch_group_assign(context, dob)
378                 # use parent to keep the outliner clean
379                 set_object_parent(dob, blob_parent)
380                 # make unselectable by default
381                 dob.hide_select = True
382                 # apply layers
383                 if groundob.meadow.use_layers:
384                     dob.layers = groundob.meadow.layers
385
386                 set_object_parent(ob, dob)
387                 
388                 dob.dupli_type = 'FACES'
389                 
390                 # make sure duplis are placed at the sample locations
391                 if ob.meadow.use_centered:
392                     # XXX centering is needed for particle instance modifier (this might be a bug!)
393                     ob.matrix_world = Matrix.Identity(4)
394                 else:
395                     ob.matrix_world = dob.matrix_world
396             else:
397                 # use parent to keep the outliner clean
398                 set_object_parent(ob, blob_parent)
399                 # move to the blob center
400                 ob.matrix_world = Matrix.Translation(blob.loc)
401         
402     # delete unused patch objects
403     delete_objects(context, del_patches)