Source code for ecoli.library.pymunk_multibody

import random
import math
import numpy as np

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, rng): """pick a random point along the boundary""" width, length = body.dimensions if rng.integers(0, 2) == 0: # force along ends if rng.integers(0, 2) == 0: # force on the left end location = (rng.uniform(0, width), 0) else: # force on the right end location = (rng.uniform(0, width), length) else: # force along length if rng.integers(0, 2) == 0: # force on the bottom end location = (0, rng.uniform(0, length)) else: # force on the top end location = (width, rng.uniform(0, length)) return location
[docs] class NullScreen(object):
[docs] def update_screen(self): pass
[docs] def configure(self, config): 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": 5e4, # scales from pN # configured parameters "jitter_force": 1e-3, # pN "bounds": [20, 20], "barriers": False, "initial_agents": {}, # for debugging "screen": None, "seed": 0, } def __init__(self, config): # PRNG for jitter force self.rng = np.random.default_rng(seed=config.get("seed", self.defaults["seed"])) # 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, self.rng) jitter_force = [ self.rng.normal(0, self.jitter_force), self.rng.normal(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 = 50.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: assert isinstance(barriers, dict) spacer_thickness = barriers.get("spacer_thickness", 0.1) channel_height = barriers.get("channel_height", 0.7 * bounds[1]) channel_space = barriers.get("channel_space", 1.5) 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(line)
[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() }
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)