Simplified weight calculation.
[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
32 from object_physics_meadow.util import *
33
34 _blob_object_name = "__MeadowBlob__"
35
36 def blob_objects(context):
37     settings = _settings.get(context)
38     blob_group = settings.blob_group(context)
39     # ignore objects with invalid blob index
40     return [ob for ob in blob_group.objects if ob.meadow.blob_index >= 0]
41
42 def blob_group_clear(context):
43     settings = _settings.get(context)
44     blob_group = settings.blob_group(context)
45     scene = context.scene
46     
47     if not blob_group:
48         return
49     
50     # local list copy to avoid messing up the iterator
51     objects = [ob for ob in blob_group.objects]
52     
53     for ob in objects:
54         scene.objects.unlink(ob)
55         blob_group.objects.unlink(ob)
56         
57         # note: this can fail if something still references the object
58         # we try to unlink own pointers above, but users might add own
59         if ob.users == 0:
60             bpy.data.objects.remove(ob)
61         else:
62             print("Warning: could not remove object %r" % ob.name)
63
64 def blob_group_assign(context, blobob):
65     settings = _settings.get(context)
66     blob_group = settings.blob_group(context)
67     
68     blob_group.objects.link(blobob)
69     # NOTE: unsetting the type is important, otherwise gathering templates
70     # a second time will include deleted objects!
71     blobob.meadow.type = 'NONE'
72
73 def blob_group_remove(context, blobob):
74     settings = _settings.get(context)
75     blob_group = settings.blob_group(context)
76     
77     blob_group.objects.unlink(blobob)
78
79 def blob_apply_settings(ob, settings):
80     pass # TODO
81
82 #-----------------------------------------------------------------------
83
84 # 8-class qualitative Brewer color scheme for high-contrast colors
85 # http://colorbrewer2.org/
86 color_schemes = [
87     (228, 26, 28),
88     (55, 126, 184),
89     (77, 175, 74),
90     (152, 78, 163),
91     (255, 127, 0),
92     (255, 255, 51),
93     (166, 86, 40),
94     (247, 129, 191),
95     ]
96
97 def select_color(index):
98     base = color_schemes[hash(str(index)) % len(color_schemes)]
99     return (base[0]/255.0, base[1]/255.0, base[2]/255.0, 1.0)
100
101 def get_blob_material(context):
102     materials = context.blend_data.materials
103     if _blob_object_name in materials:
104         return materials[_blob_object_name]
105     
106     # setup new blob material
107     ma = materials.new(_blob_object_name)
108     ma.use_object_color = True
109     # make the material stand out a bit more using emission
110     ma.emit = 1.0
111     return ma
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 blobs = []
167
168 def make_blobs(context, gridob, groundob, samples, display_radius):
169     global blobs
170     
171     blob_group_clear(context)
172     blobs = []
173     
174     blobtree = KDTree(len(gridob.data.vertices))
175     for i, v in enumerate(gridob.data.vertices):
176         # note: only using 2D coordinates, otherwise weights get distorted by z offset
177         blobtree.insert((v.co[0], v.co[1], 0.0), i)
178     blobtree.balance()
179     
180     for v in gridob.data.vertices:
181         ok, loc, nor, face_index = project_on_ground(groundob, v.co)
182         blobs.append(Blob(loc, nor, face_index) if ok else None)
183     
184     for loc, nor, face_index in samples:
185         # note: use only 2D coordinates for weighting, z component should be 0
186         index = assign_blob(blobtree, (loc[0], loc[1], 0.0), nor)
187         if index >= 0:
188             blob = blobs[index]
189             if blob:
190                 blob.samples.append((loc, nor, face_index))
191     
192     # preliminary display object
193     # XXX this could be removed eventually, but it's helpful as visual feedback to the user
194     # before creating the actual duplicator blob meshes
195     for index, blob in enumerate(blobs):
196         if blob:
197             samples = [(loc, nor) for loc, nor, _ in blob.samples]
198             ob = make_blob_object(context, index, blob.loc, samples, display_radius)
199             # put it in the blob group
200             blob_group_assign(context, ob)
201
202 #-----------------------------------------------------------------------
203
204 from object_physics_meadow.patch import patch_objects, patch_group_assign
205
206 # select one patch object for each sample based on vertex groups
207 def assign_sample_patches(groundob, blob):
208     vgroups = groundob.vertex_groups
209     faces = groundob.data.tessfaces
210     vertices = groundob.data.vertices
211     
212     vgroup_samples = { vg.name : [] for vg in vgroups }
213     vgroup_samples[""] = [] # samples for unassigned patches
214     for loc, nor, face_index in blob.samples:
215         face = faces[face_index]
216         # XXX this throws an exception if testing for vertex outside the group
217         # but there seems to be no nice way to test beforehand ...
218         #weights = [ [vg.weight(i) for i in face.vertices] for vg in vgroups ]
219         weights = [ 0.0 for vg in vgroups ]
220         for i in face.vertices:
221             v = vertices[i]
222             for vg in v.groups:
223                 weights[vg.group] += 0.25 * vg.weight # TODO
224         
225         # XXX testing
226         if weights:
227             select, (vg, w) = max(enumerate(zip(vgroups, weights)), key=lambda x: x[1][1])
228             if w > 0.0:
229                 vgroup_samples[vg.name].append((loc, nor))
230             else:
231                 vgroup_samples[""].append((loc, nor))
232         else:
233             vgroup_samples[""].append((loc, nor))
234     
235     return vgroup_samples
236
237 def setup_blob_duplis(context, groundob, display_radius):
238     global blobs
239     
240     assert(blobs)
241     
242     groundob.data.calc_tessface()
243     patches = [ob for ob in patch_objects(context) if blobs[ob.meadow.blob_index] is not None]
244     
245     for blob_index, blob in enumerate(blobs):
246         if blob is None:
247             continue
248         
249         vgroup_samples = assign_sample_patches(groundob, blob)
250         
251         for ob in patches:
252             if ob.meadow.blob_index != blob_index:
253                 continue
254             
255             samples = vgroup_samples.get(ob.meadow.density_vgroup_name, [])
256             if not samples:
257                 continue
258             
259             if ob.meadow.use_as_dupli:
260                 # make a duplicator for the patch object
261                 dob = make_blob_object(context, blob_index, blob.loc, samples, display_radius)
262                 # put the duplicator in the patch group,
263                 # so it gets removed together with patch copies
264                 patch_group_assign(context, dob)
265                 
266                 dob.dupli_type = 'FACES'
267                 
268                 ob.parent = dob
269                 # make sure duplis are placed at the sample locations
270                 if ob.meadow.use_centered:
271                     # XXX centering is needed for particle instance modifier (this might be a bug!)
272                     ob.matrix_world = Matrix.Identity(4)
273                 else:
274                     ob.matrix_world = dob.matrix_world
275             else:
276                 # move to the blob center
277                 ob.matrix_world = dob.matrix_world