from __future__ import absolute_import, division, print_function
import os
# Python imports
import random
import math
import numpy as nppython
# pymunk imports
import pymunkoptions
pymunkoptions.options["debug"] = False
import pymunk
PI = math.pi
DEBUG_SIZE = 600 # size of the pygame debug screen
[docs]def get_force_with_angle(force, angle):
x = force * math.cos(angle)
y = force * math.sin(angle)
return [x, y]
[docs]def front_from_corner(width, length, corner_position, angle):
half_width = width/2
dx = length * math.cos(angle) + half_width * math.cos(angle + PI/2) # PI/2 gives a half-rotation for the width component
dy = length * math.sin(angle) + half_width * math.sin(angle + PI/2)
front_position = [corner_position[0] + dx, corner_position[1] + dy]
return np.array([front_position[0], front_position[1], angle])
[docs]def corner_from_center(width, length, center_position, angle):
half_length = length/2
half_width = width/2
dx = half_length * math.cos(angle) + half_width * math.cos(angle + PI/2)
dy = half_length * math.sin(angle) + half_width * math.sin(angle + PI/2)
corner_position = [center_position[0] - dx, center_position[1] - dy]
return np.array([corner_position[0], corner_position[1], angle])
[docs]def random_body_position(body):
''' pick a random point along the boundary'''
width, length = body.dimensions
if random.randint(0, 1) == 0:
# force along ends
if random.randint(0, 1) == 0:
# force on the left end
location = (random.uniform(0, width), 0)
else:
# force on the right end
location = (random.uniform(0, width), length)
else:
# force along length
if random.randint(0, 1) == 0:
# force on the bottom end
location = (0, random.uniform(0, length))
else:
# force on the top end
location = (width, random.uniform(0, length))
return location
[docs]class NullScreen(object):
[docs] def update_screen(self):
pass
[docs]class PymunkMultibody(object):
"""
Multibody object for interfacing with pymunk
"""
defaults = {
'agent_shape': 'segment',
# hardcoded parameters
'elasticity': 0.9,
'damping': 0.5, # 1 is no damping, 0 is full damping
'angular_damping': 0.8,
'friction': 0.9, # does this do anything?
'physics_dt': 0.001,
'force_scaling': 1e2, # scales from pN
# configured parameters
'jitter_force': 1e-3, # pN
'bounds': [20, 20],
'barriers': False,
'initial_agents': {},
# for debugging
'screen': None,
}
def __init__(self, config):
# hardcoded parameters
self.elasticity = self.defaults['elasticity']
self.friction = self.defaults['friction']
self.damping = self.defaults['damping']
self.angular_damping = self.defaults['angular_damping']
self.physics_dt = config.get('physics_dt', self.defaults['physics_dt'])
self.force_scaling = self.defaults['force_scaling']
# configured parameters
self.agent_shape = config.get('agent_shape', self.defaults['agent_shape'])
self.jitter_force = config.get('jitter_force', self.defaults['jitter_force'])
self.bounds = config.get('bounds', self.defaults['bounds'])
barriers = config.get('barriers', self.defaults['barriers'])
# initialize pymunk space
self.space = pymunk.Space()
# debug screen
self.screen = config.get('screen')
if self.screen is None:
self.screen = NullScreen()
self.screen.configure({
'space': self.space,
'bounds': self.bounds})
# add static barriers
self.add_barriers(self.bounds, barriers)
# initialize agents
initial_agents = config.get('initial_agents', self.defaults['initial_agents'])
self.bodies = {}
for agent_id, specs in initial_agents.items():
self.add_body_from_center(agent_id, specs)
[docs] def run(self, timestep):
if self.physics_dt > timestep:
print('timestep skipped by pymunk_multibody: {}'.format(timestep))
return
time = 0
while time < timestep:
time += self.physics_dt
# apply forces
for body in self.space.bodies:
self.apply_jitter_force(body)
self.apply_motile_force(body)
self.apply_viscous_force(body)
# run for a physics timestep
self.space.step(self.physics_dt)
self.screen.update_screen()
[docs] def apply_motile_force(self, body):
width, length = body.dimensions
motile_location = (width / 2, 0) # apply force at back end of body
thrust = 0.0
torque = 0.0
motile_force = [thrust, torque]
if hasattr(body, 'thrust'):
thrust = body.thrust
torque = body.torque
motile_force = [thrust, 0.0]
# add to angular velocity
body.angular_velocity += torque
scaled_motile_force = [force * self.force_scaling for force in motile_force]
body.apply_impulse_at_local_point(scaled_motile_force, motile_location)
[docs] def apply_jitter_force(self, body):
jitter_location = random_body_position(body)
jitter_force = [
random.normalvariate(0, self.jitter_force),
random.normalvariate(0, self.jitter_force)]
scaled_jitter_force = [
force * self.force_scaling
for force in jitter_force]
body.apply_impulse_at_local_point(
scaled_jitter_force,
jitter_location)
[docs] def apply_viscous_force(self, body):
# dampen velocity
body.velocity = body.velocity * self.damping + (body.force / body.mass) * self.physics_dt
# dampen angular velocity
body.angular_velocity = body.angular_velocity * self.angular_damping + (body.torque / body.moment) * self.physics_dt
[docs] def add_barriers(self, bounds, barriers):
""" Create static barriers """
thickness = 5.0
offset = thickness
x_bound = bounds[0]
y_bound = bounds[1]
static_body = self.space.static_body
static_lines = [
pymunk.Segment(
static_body,
(0.0-offset, 0.0-offset),
(x_bound+offset, 0.0-offset),
thickness),
pymunk.Segment(
static_body,
(x_bound+offset, 0.0-offset),
(x_bound+offset, y_bound+offset),
thickness),
pymunk.Segment(
static_body,
(x_bound+offset, y_bound+offset),
(0.0-offset, y_bound+offset),
thickness),
pymunk.Segment(
static_body,
(0.0-offset, y_bound+offset),
(0.0-offset, 0.0-offset),
thickness),
]
if barriers:
spacer_thickness = barriers.get('spacer_thickness', 0.1)
channel_height = barriers.get('channel_height')
channel_space = barriers.get('channel_space')
n_lines = math.floor(x_bound/channel_space)
machine_lines = [
pymunk.Segment(
static_body,
(channel_space * line, 0),
(channel_space * line, channel_height), spacer_thickness)
for line in range(n_lines)]
static_lines += machine_lines
for line in static_lines:
line.elasticity = 0.0 # bounce
line.friction = 0.8
self.space.add(static_lines)
[docs] def get_shape(self, boundary):
'''
shape documentation at: https://pymunk-tutorial.readthedocs.io/en/latest/shape/shape.html
'''
if self.agent_shape == 'segment':
width = boundary['width']
length = boundary['length']
half_width = width / 2
half_length = length / 2 - half_width
shape = pymunk.Segment(
None,
(-half_length, 0),
(half_length, 0),
radius=half_width)
elif self.agent_shape == 'circle':
length = boundary['length']
half_length = length / 2
shape = pymunk.Circle(None, radius=half_length, offset=(0, 0))
elif self.agent_shape == 'rectangle':
width = boundary['width']
length = boundary['length']
half_length = length / 2
half_width = width / 2
shape = pymunk.Poly(None,
((-half_length, -half_width),
(half_length, -half_width),
(half_length, half_width),
(-half_length, half_width)))
return shape
[docs] def get_inertia(self, shape, mass):
if self.agent_shape == 'rectangle':
inertia = pymunk.moment_for_poly(mass, shape.get_vertices())
elif self.agent_shape == 'circle':
radius = shape.radius
inertia = pymunk.moment_for_circle(mass, radius, radius)
elif self.agent_shape == 'segment':
a = shape.a
b = shape.b
radius = shape.radius
inertia = pymunk.moment_for_segment(mass, a, b, radius)
return inertia
[docs] def add_body_from_center(self, body_id, specs):
boundary = specs['boundary']
mass = boundary['mass']
center_position = boundary['location']
angle = boundary['angle']
angular_velocity = boundary.get('angular_velocity', 0.0)
width = boundary['width']
length = boundary['length']
# get shape, inertia, make body, assign body to shape
shape = self.get_shape(boundary)
inertia = self.get_inertia(shape, mass)
body = pymunk.Body(mass, inertia)
shape.body = body
body.position = (
center_position[0],
center_position[1])
body.angle = angle
body.dimensions = (width, length)
body.angular_velocity = angular_velocity
shape.elasticity = self.elasticity
shape.friction = self.friction
# add body and shape to space
self.space.add(body, shape)
# add body to agents dictionary
self.bodies[body_id] = (body, shape)
[docs] def update_body(self, body_id, specs):
boundary = specs['boundary']
length = boundary['length']
width = boundary['width']
mass = boundary['mass']
thrust = boundary['thrust']
torque = boundary['torque']
body, shape = self.bodies[body_id]
position = body.position
angle = body.angle
# get shape, inertia, make body, assign body to shape
new_shape = self.get_shape(boundary)
inertia = self.get_inertia(new_shape, mass)
new_body = pymunk.Body(mass, inertia)
new_shape.body = new_body
new_body.position = position
new_body.angle = angle
new_body.velocity = body.velocity
new_body.angular_velocity = body.angular_velocity
new_body.dimensions = (width, length)
new_body.thrust = thrust
new_body.torque = torque
new_shape.elasticity = shape.elasticity
new_shape.friction = shape.friction
# swap bodies
self.space.remove(body, shape)
self.space.add(new_body, new_shape)
# update body
self.bodies[body_id] = (new_body, new_shape)
[docs] def update_bodies(self, bodies):
# if an agent has been removed from the agents store,
# remove it from space and bodies
removed_bodies = [
body_id for body_id in self.bodies.keys()
if body_id not in bodies.keys()]
for body_id in removed_bodies:
body, shape = self.bodies[body_id]
self.space.remove(body, shape)
del self.bodies[body_id]
# update agents, add new agents
for body_id, specs in bodies.items():
if body_id in self.bodies:
self.update_body(body_id, specs)
else:
self.add_body_from_center(body_id, specs)
[docs] def get_body_position(self, agent_id):
body, shape = self.bodies[agent_id]
return {
'location': [pos for pos in body.position],
'angle': body.angle,
}
[docs] def get_body_positions(self):
return {
body_id: {
'boundary': self.get_body_position(body_id)}
for body_id in self.bodies.keys()}
[docs]def test_multibody(
total_time=2,
agent_shape='rectangle',
n_agents=1,
jitter_force=1e1,
screen=None):
bounds = [500, 500]
center_location = [0.5*loc for loc in bounds]
agents = {
str(agent_idx): {
'boundary': {
'location': center_location,
'angle': random.uniform(0,2*PI),
'volume': 15,
'length': 30,
'width': 10,
'mass': 1,
'thrust': 1e3,
'torque': 0.0}}
for agent_idx in range(n_agents)
}
config = {
'agent_shape': agent_shape,
'jitter_force': jitter_force,
'bounds': bounds,
'barriers': False,
'initial_agents': agents,
'screen': screen
}
multibody = PymunkMultibody(config)
# run simulation
time = 0
time_step = 1
while time < total_time:
time += time_step
multibody.run(time_step)
if __name__ == '__main__':
test_multibody(10)