[Edit Linked Library] Fixes based on code review from Campbell
[blender-addons-contrib.git] / wetted_mesh.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 bl_info = {
20     "name": "Add Wetted Mesh",
21     "author": "freejack",
22     "version": (0, 2, 1),
23     "blender": (2, 58, 0),
24     "location": "View3D > Tool Shelf > Wetted Mesh Panel",
25     "description": "Adds separated fluid, dry and wetted mesh for selected pair.",
26     "warning": "",
27     "wiki_url": "http://wiki.blender.org/index.php/Extensions:2.6/Py/"\
28         "Scripts/Mesh/Wetted_Mesh",
29     "tracker_url": "http://projects.blender.org/tracker/index.php?"\
30         "func=detail&aid=27156",
31     "category": "Mesh"}
32
33 import bpy
34 import collections
35 import math
36
37 ### Tool Panel ###
38 class VIEW3D_PT_tools_WettedMesh(bpy.types.Panel):
39     """Wetted Mesh Tool Panel"""
40     bl_space_type = 'VIEW_3D'
41     bl_region_type = 'TOOLS'
42     bl_label = 'Wetted Mesh'
43     bl_context = 'objectmode'
44     bl_options = {'DEFAULT_CLOSED'}
45
46     def draw(self, context):
47         layout = self.layout
48         col = layout.column(align=True)
49         slcnt = len(context.selected_objects)
50
51         if slcnt != 2:
52             col.label(text = 'Select two mesh objects')
53             col.label(text = 'to generate separated')
54             col.label(text = 'fluid, dry and wetted')
55             col.label(text = 'meshes.')
56         else:
57             (solid, fluid) = getSelectedPair(context)
58             col.label(text = 'solid = '+solid.name)
59             col.label(text = 'fluid = '+fluid.name)
60             col.operator('mesh.primitive_wetted_mesh_add', text='Generate Meshes')
61
62 ### Operator ###
63 class AddWettedMesh(bpy.types.Operator):
64     """Add wetted mesh for selected mesh pair"""
65     bl_idname = "mesh.primitive_wetted_mesh_add"
66     bl_label = "Add Wetted Mesh"
67     bl_options = {'REGISTER', 'UNDO'}
68     statusMessage = ''
69
70     def draw(self, context):
71         layout = self.layout
72         col = layout.column(align=True)
73         col.label(text = self.statusMessage)
74
75     def execute(self, context):
76         # make sure a pair of objects is selected
77         if len(context.selected_objects) != 2:
78             # should not happen if called from tool panel
79             self.report({'WARNING'}, "no mesh pair selected, operation cancelled")
80             return {'CANCELLED'}
81
82         print("add_wetted_mesh begin")
83         
84         # super-selected object is solid, other object is fluid
85         (solid, fluid) = getSelectedPair(context)
86         print("   solid = "+solid.name)
87         print("   fluid = "+fluid.name)
88             
89         # make a copy of fluid object, convert to mesh if required
90         print("   copy fluid")
91         bpy.ops.object.select_all(action='DESELECT')
92         fluid.select = True
93         context.scene.objects.active = fluid
94         bpy.ops.object.duplicate()
95         bpy.ops.object.convert(target='MESH', keep_original=False)
96         bpy.ops.object.transform_apply(location=True, rotation=True, scale=True)
97         fluidCopy = context.object
98         
99         # substract solid from fluidCopy
100         print("   bool: fluidCopy DIFFERENCE solid")
101         bpy.ops.object.modifier_add(type='BOOLEAN')
102         bop = fluidCopy.modifiers.items()[0]
103         bop[1].operation = 'DIFFERENCE'
104         bop[1].object = solid
105         bpy.ops.object.modifier_apply(apply_as='DATA', modifier=bop[0])
106         fluidMinusSolid = fluidCopy
107         fluidMinusSolid.name = "fluidMinusSolid"
108         
109         # make a second copy of fluid object
110         print("   copy fluid")
111         bpy.ops.object.select_all(action='DESELECT')
112         fluid.select = True
113         context.scene.objects.active = fluid
114         bpy.ops.object.duplicate()
115         bpy.ops.object.convert(target='MESH', keep_original=False)
116         bpy.ops.object.transform_apply(location=True, rotation=True, scale=True)
117         fluidCopy = context.object
118         
119         # make union from fluidCopy and solid
120         print("   bool: fluidCopy UNION solid")
121         bpy.ops.object.modifier_add(type='BOOLEAN')
122         bop = fluidCopy.modifiers.items()[0]
123         bop[1].operation = 'UNION'
124         bop[1].object = solid
125         bpy.ops.object.modifier_apply(apply_as='DATA', modifier=bop[0])
126         fluidUnionSolid = fluidCopy
127         fluidUnionSolid.name = "fluidUnionSolid"
128         
129         # index meshes
130         print("   KDTree index fluidMinusSolid")
131         fluidMinusSolidKDT = KDTree(3, fluidMinusSolid.data.vertices)
132         print("   KDTree index fluidUnionSolid")
133         fluidUnionSolidKDT = KDTree(3, fluidUnionSolid.data.vertices)
134         kdtrees = (fluidMinusSolidKDT, fluidUnionSolidKDT)
135         
136         # build mesh face sets
137         faceDict = { }
138         vertDict = { }
139         
140         print("   processing fluidMinusSolid faces")
141         cacheDict = { }
142         setFMSfaces = set()
143         numFaces = len(fluidUnionSolid.data.faces)
144         i = 0
145         for f in fluidMinusSolid.data.faces:
146             if i % 500 == 0:
147                 print("      ", i, " / ", numFaces)
148             i += 1
149             fuid = unifiedFaceId(kdtrees, f, fluidMinusSolid.data.vertices, \
150                                  faceDict, vertDict, cacheDict)
151             setFMSfaces.add(fuid)
152         
153         print("   processing fluidUnionSolid faces")
154         cacheDict = { }
155         setFUSfaces = set()
156         numFaces = len(fluidUnionSolid.data.faces)
157         i = 0
158         for f in fluidUnionSolid.data.faces:
159             if i % 500 == 0:
160                 print("      ", i, " / ", numFaces)
161             i += 1
162             fuid = unifiedFaceId(kdtrees, f, fluidUnionSolid.data.vertices, \
163                                  faceDict, vertDict, cacheDict)
164             setFUSfaces.add(fuid)
165         
166         # remove boolean helpers
167         print("   delete helper objects")
168         bpy.ops.object.select_all(action='DESELECT')
169         fluidUnionSolid.select = True
170         fluidMinusSolid.select = True
171         bpy.ops.object.delete()
172
173         # wetted = FMS - FUS
174         print("   set operation FMS diff FUS")
175         setWetFaces = setFMSfaces.difference(setFUSfaces)
176         print("   build wetted mesh")
177         verts, faces = buildMesh(setWetFaces, faceDict, vertDict)
178         print("   create wetted mesh")
179         wetted = createMesh("Wetted", verts, faces)
180
181         # fluid = FMS x FUS
182         print("   set operation FMS intersect FUS")
183         setFluidFaces = setFMSfaces.intersection(setFUSfaces)
184         print("   build fluid mesh")
185         verts, faces = buildMesh(setFluidFaces, faceDict, vertDict)
186         print("   create fluid mesh")
187         fluid = createMesh("Fluid", verts, faces)
188         
189         # solid = FUS - FMS
190         print("   set operation FUS diff FMS")
191         setSolidFaces = setFUSfaces.difference(setFMSfaces)
192         print("   build solid mesh")
193         verts, faces = buildMesh(setSolidFaces, faceDict, vertDict)
194         print("   create solid mesh")
195         solid = createMesh("Solid", verts, faces)
196         
197         # parent wetted mesh
198         print("   parent mesh")
199         bpy.ops.object.add(type='EMPTY')
200         wettedMesh = context.object
201         solid.select = True
202         fluid.select = True
203         wetted.select = True
204         wettedMesh.select = True
205         bpy.ops.object.parent_set(type='OBJECT')
206         wettedMesh.name = 'WettedMesh'
207         
208         print("add_wetted_mesh done")
209         self.statusMessage = 'created '+wettedMesh.name
210
211         return {'FINISHED'}
212
213
214 ### Registration ###
215 def register():
216     bpy.utils.register_class(VIEW3D_PT_tools_WettedMesh)
217     bpy.utils.register_class(AddWettedMesh)
218
219
220 def unregister():
221     bpy.utils.unregister_class(VIEW3D_PT_tools_WettedMesh)
222     bpy.utils.unregister_class(AddWettedMesh)
223
224 if __name__ == "__main__":
225     register()
226
227
228 #
229 # KD tree (used to create a geometric index of mesh vertices)
230 #
231
232 def distance(a, b):
233     return (a-b).length
234
235 Node = collections.namedtuple("Node", 'point axis label left right')
236
237 class KDTree(object):
238     """A tree for nearest neighbor search in a k-dimensional space.
239
240     For information about the implementation, see
241     http://en.wikipedia.org/wiki/Kd-tree
242
243     Usage:
244     objects is an iterable of (co, index) tuples (so MeshVertex is useable)
245     k is the number of dimensions (=3)
246     
247     t = KDTree(k, objects)
248     point, label, distance = t.nearest_neighbor(destination)
249     """
250
251     def __init__(self, k, objects=[]):
252
253         def build_tree(objects, axis=0):
254
255             if not objects:
256                 return None
257
258             objects.sort(key=lambda o: o.co[axis])
259             median_idx = len(objects) // 2
260             median_point = objects[median_idx].co
261             median_label = objects[median_idx].index
262
263             next_axis = (axis + 1) % k
264             return Node(median_point, axis, median_label,
265                         build_tree(objects[:median_idx], next_axis),
266                         build_tree(objects[median_idx + 1:], next_axis))
267
268         self.root = build_tree(list(objects))
269         self.size = len(objects)
270
271
272     def nearest_neighbor(self, destination):
273
274         best = [None, None, float('inf')]
275         # state of search: best point found, its label,
276         # lowest distance
277
278         def recursive_search(here):
279
280             if here is None:
281                 return
282             point, axis, label, left, right = here
283
284             here_sd = distance(point, destination)
285             if here_sd < best[2]:
286                 best[:] = point, label, here_sd
287
288             diff = destination[axis] - point[axis]
289             close, away = (left, right) if diff <= 0 else (right, left)
290
291             recursive_search(close)
292             if math.fabs(diff) < best[2]:
293                 recursive_search(away)
294
295         recursive_search(self.root)
296         return best[0], best[1], best[2]
297
298
299 #
300 # helper functions
301 #
302
303 # get super-selected object and other object from selected pair
304 def getSelectedPair(context):
305     objA = context.object
306     objB = context.selected_objects[0]
307     if objA == objB:
308         objB = context.selected_objects[1]
309     return (objA, objB)
310
311 # get a unified vertex id for given coordinates
312 def unifiedVertexId(kdtrees, location, vertDict):
313     eps = 0.0001
314     offset = 0
315     for t in kdtrees:
316         co, index, d = t.nearest_neighbor(location)
317         if d < eps:
318             uvid = offset + index
319             if uvid not in vertDict:
320                 vertDict[uvid] = co
321             return uvid
322         offset += t.size
323     return -1
324
325 # get a unified face id tuple
326 #    Stores the ordered face id tuple in faceDict
327 #    and the used coordinates for vertex id in vertDict.
328 #    cacheDict caches the unified vertex id (lookup in kdtree is expensive).
329 #    For each mesh (where the face belongs to) a separate cacheDict is expected.
330 def unifiedFaceId(kdtrees, face, vertices, faceDict, vertDict, cacheDict):
331     fids = [ ]
332     for v in face.vertices:
333         uvid = cacheDict.get(v)
334         if uvid == None:
335             uvid = unifiedVertexId(kdtrees, vertices[v].co, vertDict)
336             cacheDict[v] = uvid
337         fids.append(uvid)
338     ofids = tuple(fids)
339     fids.sort()
340     fuid = tuple(fids)
341     if fuid not in faceDict:
342         faceDict[fuid] = ofids
343     return fuid
344
345 # build vertex and face array from unified face sets
346 def buildMesh(unifiedFaceSet, faceDict, vertDict):
347     verts = [ ]
348     nextV = 0
349     myV = { }
350     faces = [ ]
351     for uf in unifiedFaceSet:
352         of = faceDict[uf]
353         myf = [ ]
354         for uV in of:
355             v = myV.get(uV)
356             if v == None:
357                 v = nextV
358                 myV[uV] = nextV
359                 verts.append(vertDict[uV])
360                 nextV += 1
361             myf.append(v)
362         faces.append(myf)
363     return verts, faces
364
365 # create mesh object and link to scene
366 def createMesh(name, verts, faces):
367     me = bpy.data.meshes.new(name+"Mesh")
368     ob = bpy.data.objects.new(name, me)
369     ob.show_name = True
370     bpy.context.scene.objects.link(ob)
371     me.from_pydata(verts, [], faces)
372     me.update(calc_edges=True)
373     return ob