Merged revision(s) 57979-58058 from trunk/blender into soc-2013-dingto
[blender-staging.git] / release / scripts / startup / bl_operators / rigidbody.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-80 compliant>
20
21 import bpy
22 from bpy.types import Operator
23 from bpy.props import IntProperty
24 from bpy.props import EnumProperty
25
26
27 class CopyRigidbodySettings(Operator):
28     '''Copy Rigid Body settings from active object to selected'''
29     bl_idname = "rigidbody.object_settings_copy"
30     bl_label = "Copy Rigid Body Settings"
31     bl_options = {'REGISTER', 'UNDO'}
32
33     _attrs = (
34         "type",
35         "kinematic",
36         "mass",
37         "collision_shape",
38         "use_margin",
39         "collision_margin",
40         "friction",
41         "restitution",
42         "use_deactivation",
43         "use_start_deactivated",
44         "deactivate_linear_velocity",
45         "deactivate_angular_velocity",
46         "linear_damping",
47         "angular_damping",
48         "collision_groups",
49         )
50
51     @classmethod
52     def poll(cls, context):
53         obj = context.object
54         return (obj and obj.rigid_body)
55
56     def execute(self, context):
57         obj_act = context.object
58         scene = context.scene
59
60         # deselect all but mesh objects
61         for o in context.selected_objects:
62             if o.type != 'MESH':
63                 o.select = False
64
65         objects = context.selected_objects
66         if objects:
67             # add selected objects to active one groups and recalculate
68             bpy.ops.group.objects_add_active()
69             scene.frame_set(scene.frame_current)
70             rb_from = obj_act.rigid_body
71             # copy settings
72             for o in objects:
73                 rb_to = o.rigid_body
74                 if (o == obj_act) or (rb_to is None):
75                     continue
76                 for attr in self._attrs:
77                     setattr(rb_to, attr, getattr(rb_from, attr))
78
79         return {'FINISHED'}
80
81
82 class BakeToKeyframes(Operator):
83     '''Bake rigid body transformations of selected objects to keyframes'''
84     bl_idname = "rigidbody.bake_to_keyframes"
85     bl_label = "Bake To Keyframes"
86     bl_options = {'REGISTER', 'UNDO'}
87
88     frame_start = IntProperty(
89             name="Start Frame",
90             description="Start frame for baking",
91             min=0, max=300000,
92             default=1,
93             )
94     frame_end = IntProperty(
95             name="End Frame",
96             description="End frame for baking",
97             min=1, max=300000,
98             default=250,
99             )
100     step = IntProperty(
101             name="Frame Step",
102             description="Frame Step",
103             min=1, max=120,
104             default=1,
105             )
106
107     @classmethod
108     def poll(cls, context):
109         obj = context.object
110         return (obj and obj.rigid_body)
111
112     def execute(self, context):
113         bake = []
114         objects = []
115         scene = context.scene
116         frame_orig = scene.frame_current
117         frames_step = range(self.frame_start, self.frame_end + 1, self.step)
118         frames_full = range(self.frame_start, self.frame_end + 1)
119
120         # filter objects selection
121         for obj in context.selected_objects:
122             if not obj.rigid_body or obj.rigid_body.type != 'ACTIVE':
123                 obj.select = False
124
125         objects = context.selected_objects
126
127         if objects:
128             # store transformation data
129             # need to start at scene start frame so simulation is run from the beginning
130             for f in frames_full:
131                 scene.frame_set(f)
132                 if f in frames_step:
133                     mat = {}
134                     for i, obj in enumerate(objects):
135                         mat[i] = obj.matrix_world.copy()
136                     bake.append(mat)
137
138             # apply transformations as keyframes
139             for i, f in enumerate(frames_step):
140                 scene.frame_set(f)
141                 obj_prev = objects[0]
142                 for j, obj in enumerate(objects):
143                     mat = bake[i][j]
144
145                     obj.location = mat.to_translation()
146
147                     rot_mode = obj.rotation_mode
148                     if rot_mode == 'QUATERNION':
149                         obj.rotation_quaternion = mat.to_quaternion()
150                     elif rot_mode == 'AXIS_ANGLE':
151                         # this is a little roundabout but there's no better way right now
152                         aa = mat.to_quaternion().to_axis_angle()
153                         obj.rotation_axis_angle = (aa[1], ) + aa[0][:]
154                     else:  # euler
155                         # make sure euler rotation is compatible to previous frame
156                         obj.rotation_euler = mat.to_euler(rot_mode, obj_prev.rotation_euler)
157
158                     obj_prev = obj
159
160                 bpy.ops.anim.keyframe_insert(type='BUILTIN_KSI_LocRot', confirm_success=False)
161
162             # remove baked objects from simulation
163             bpy.ops.rigidbody.objects_remove()
164
165             # clean up keyframes
166             for obj in objects:
167                 action = obj.animation_data.action
168                 for fcu in action.fcurves:
169                     keyframe_points = fcu.keyframe_points
170                     i = 1
171                     # remove unneeded keyframes
172                     while i < len(keyframe_points) - 1:
173                         val_prev = keyframe_points[i - 1].co[1]
174                         val_next = keyframe_points[i + 1].co[1]
175                         val = keyframe_points[i].co[1]
176
177                         if abs(val - val_prev) + abs(val - val_next) < 0.0001:
178                             keyframe_points.remove(keyframe_points[i])
179                         else:
180                             i += 1
181                     # use linear interpolation for better visual results
182                     for keyframe in keyframe_points:
183                         keyframe.interpolation = 'LINEAR'
184
185             # return to the frame we started on
186             scene.frame_set(frame_orig)
187
188         return {'FINISHED'}
189
190     def invoke(self, context, event):
191         scene = context.scene
192         self.frame_start = scene.frame_start
193         self.frame_end = scene.frame_end
194
195         wm = context.window_manager
196         return wm.invoke_props_dialog(self)
197
198
199 class ConnectRigidBodies(Operator):
200     '''Create rigid body constraints between selected rigid bodies'''
201     bl_idname = "rigidbody.connect"
202     bl_label = "Connect Rigid Bodies"
203     bl_options = {'REGISTER', 'UNDO'}
204
205     con_type = EnumProperty(
206             name="Type",
207             description="Type of generated constraint",
208             # XXX Would be nice to get icons too, but currently not possible ;)
209             items=tuple((e.identifier, e.name, e.description, e. value)
210                         for e in bpy.types.RigidBodyConstraint.bl_rna.properties["type"].enum_items),
211             default='FIXED',
212             )
213     pivot_type = EnumProperty(
214             name="Location",
215             description="Constraint pivot location",
216             items=(('CENTER', "Center", "Pivot location is between the constrained rigid bodies"),
217                    ('ACTIVE', "Active", "Pivot location is at the active object position"),
218                    ('SELECTED', "Selected", "Pivot location is at the selected object position")),
219             default='CENTER',
220             )
221     connection_pattern = EnumProperty(
222             name="Connection Pattern",
223             description="Pattern used to connect objects",
224             items=(('SELECTED_TO_ACTIVE', "Selected to Active", "Connect selected objects to the active object"),
225                    ('CHAIN_DISTANCE', "Chain by Distance", "Connect objects as a chain based on distance, starting at the active object")),
226             default='SELECTED_TO_ACTIVE',
227             )
228
229     @classmethod
230     def poll(cls, context):
231         obj = context.object
232         return (obj and obj.rigid_body)
233
234     def _add_constraint(self, context, object1, object2):
235         if object1 == object2:
236             return
237
238         if self.pivot_type == 'ACTIVE':
239             loc = object1.location
240         elif self.pivot_type == 'SELECTED':
241             loc = object2.location
242         else:
243             loc = (object1.location + object2.location) / 2.0
244
245         ob = bpy.data.objects.new("Constraint", object_data=None)
246         ob.location = loc
247         context.scene.objects.link(ob)
248         context.scene.objects.active = ob
249         ob.select = True
250
251         bpy.ops.rigidbody.constraint_add()
252         con_obj = context.active_object
253         con_obj.empty_draw_type = 'ARROWS'
254         con = con_obj.rigid_body_constraint
255         con.type = self.con_type
256
257         con.object1 = object1
258         con.object2 = object2
259
260     def execute(self, context):
261         scene = context.scene
262         objects = context.selected_objects
263         obj_act = context.active_object
264         change = False
265
266         if self.connection_pattern == 'CHAIN_DISTANCE':
267             objs_sorted = [obj_act]
268             objects_tmp = context.selected_objects
269             try:
270                 objects_tmp.remove(obj_act)
271             except ValueError:
272                 pass
273
274             last_obj = obj_act
275
276             while objects_tmp:
277                 objects_tmp.sort(key=lambda o: (last_obj.location - o.location).length)
278                 last_obj = objects_tmp.pop(0)
279                 objs_sorted.append(last_obj)
280
281             for i in range(1, len(objs_sorted)):
282                 self._add_constraint(context, objs_sorted[i - 1], objs_sorted[i])
283                 change = True
284
285         else:  # SELECTED_TO_ACTIVE
286             for obj in objects:
287                 self._add_constraint(context, obj_act, obj)
288                 change = True
289
290         if change:
291             # restore selection
292             bpy.ops.object.select_all(action='DESELECT')
293             for obj in objects:
294                 obj.select = True
295             scene.objects.active = obj_act
296             return {'FINISHED'}
297         else:
298             self.report({'WARNING'}, "No other objects selected")
299             return {'CANCELLED'}