"""
==========================
Multibody physics process
==========================
"""
import os
import random
import math
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches
# vivarium imports
from vivarium.library.units import units, remove_units
from vivarium.core.process import Process
from vivarium.core.composition import (
process_in_experiment,
simulate_experiment,
PROCESS_OUT_DIR,
)
# vivarium-cell imports
from ecoli.processes.environment.derive_globals import volume_from_length
from ecoli.library.pymunk_multibody import PymunkMultibody
from ecoli.analysis.colony.snapshots import (
plot_snapshots,
format_snapshot_data,
)
NAME = "multibody"
DEFAULT_BOUNDS = [40, 40] * units.um
DEFAULT_LENGTH_UNIT = units.um
DEFAULT_MASS_UNIT = units.fg
DEFAULT_VELOCITY_UNIT = units.um / units.s
DEFAULT_VOLUME_UNIT = DEFAULT_LENGTH_UNIT**3
# constants
PI = math.pi
[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]
def daughter_locations(value, state):
parent_length = state["length"]
parent_angle = state["angle"]
pos_ratios = [-0.25, 0.25]
daughter_locations = []
for daughter in range(2):
dx = parent_length * pos_ratios[daughter] * math.cos(parent_angle)
dy = parent_length * pos_ratios[daughter] * math.sin(parent_angle)
location = [value[0] + dx, value[1] + dy]
daughter_locations.append(location)
return daughter_locations
[docs]
class Multibody(Process):
"""Simulates collisions and forces between agent bodies with a multi-body physics engine.
:term:`Ports`:
* ``agents``: The store containing all agent sub-compartments. Each agent in
this store has values for location, angle, length, width, mass, thrust, and torque.
Arguments:
initial_parameters(dict): Accepts the following configuration keys:
jitter_force: force applied to random positions along agent
bodies to mimic thermal fluctuations. Produces Brownian motion.
agent_shape (:py:class:`str`): agents can take the shapes
``rectangle``, ``segment``, or ``circle``.
bounds (:py:class:`list`): size of the environment in
micrometers, with ``[x, y]``.
mother_machine (:py:class:`bool`): if set to ``True``, mother
machine barriers are introduced.
animate (:py:class:`bool`): interactive matplotlib option to
animate multibody. To run with animation turned on set True, and use
the TKAgg matplotlib backend:
.. code-block:: console
$ MPLBACKEND=TKAgg python vivarium/processes/snapshots.py
Notes:
* rotational diffusion in liquid medium with viscosity = 1 mPa.s: :math:`Dr = 3.5 \\pm0.3 rad^{2}/s`
(Saragosti, et al. 2012. Modeling E. coli tumbles by rotational diffusion.)
* translational diffusion in liquid medium with viscosity = 1 mPa.s: :math:`Dt = 100 um^{2}/s`
(Saragosti, et al. 2012. Modeling E. coli tumbles by rotational diffusion.)
"""
name = NAME
defaults = {
"jitter_force": 1e-4, # pN
"agent_shape": "segment",
"bounds": DEFAULT_BOUNDS,
"length_unit": DEFAULT_LENGTH_UNIT,
"mass_unit": DEFAULT_MASS_UNIT,
"velocity_unit": DEFAULT_VELOCITY_UNIT,
"boundary_key": "boundary",
"mother_machine": False,
"animate": False,
"seed": 0,
}
def __init__(self, parameters=None):
super().__init__(parameters)
# multibody parameters
jitter_force = self.parameters["jitter_force"]
self.agent_shape = self.parameters["agent_shape"]
self.bounds = self.parameters["bounds"]
self.mother_machine = self.parameters["mother_machine"]
# units
self.length_unit = self.parameters["length_unit"]
self.mass_unit = self.parameters["mass_unit"]
self.velocity_unit = self.parameters["velocity_unit"]
# make the multibody object
if self.mother_machine:
assert isinstance(self.mother_machine, dict), (
"mother_machine must be a dictionary with keys "
"spacer_thickness, channel_height, channel_space"
)
multibody_config = {
"agent_shape": self.agent_shape,
"jitter_force": jitter_force,
"bounds": remove_units(self.bounds),
"barriers": self.mother_machine,
"physics_dt": self.parameters["timestep"] / 10,
"seed": self.parameters["seed"],
}
self.physics = PymunkMultibody(multibody_config)
# interactive plot for visualization
self.animate = self.parameters["animate"]
if self.animate:
plt.ion()
self.ax = plt.gca()
self.ax.set_aspect("equal")
[docs]
def ports_schema(self):
glob_schema = {
"*": {
self.parameters["boundary_key"]: {
"location": {
"_emit": True,
"_default": [0.5 * bound for bound in self.bounds],
"_updater": "set",
"_divider": {
"divider": daughter_locations,
"topology": {
"length": (
"..",
"length",
),
"angle": (
"..",
"angle",
),
},
},
},
"length": {"_emit": True, "_default": 2.0 * units.um},
"width": {"_emit": True, "_default": 1.0 * units.um},
"angle": {
"_emit": True,
"_default": 0.0,
"_updater": "set",
},
"mass": {
"_emit": True,
"_default": 1339 * units.fg,
},
"thrust": {
"_default": 0.0,
"_updater": "set",
},
"torque": {
"_default": 0.0,
"_updater": "set",
},
}
}
}
schema = {"agents": glob_schema}
return schema
[docs]
def next_update(self, timestep, states):
agents = states["agents"]
# animate before update
agents = remove_units(agents)
if self.animate:
self.animate_frame(agents)
# update multibody with new agents
self.physics.update_bodies(agents)
# run simulation
self.physics.run(timestep)
# get new agent positions
agent_positions = self.physics.get_body_positions()
update = {"agents": agent_positions}
# for mother machine configurations, remove agents above the channel height
if self.mother_machine:
channel_height = self.mother_machine["channel_height"]
delete_agents = []
for agent_id, position in agent_positions.items():
location = position["boundary"]["location"]
y_loc = location[1]
if y_loc > channel_height:
# cell has moved past the channels
delete_agents.append(agent_id)
if delete_agents:
update["agents"] = {
agent_id: position
for agent_id, position in agent_positions.items()
if agent_id not in delete_agents
}
update["agents"]["_delete"] = [agent_id for agent_id in delete_agents]
for agent in update["agents"].values():
agent["boundary"]["location"] *= units.um
return update
## matplotlib interactive plot
[docs]
def animate_frame(self, agents):
plt.cla()
for agent_id, data in agents.items():
# location, orientation, length
data = data["boundary"]
x_center = data["location"][0]
y_center = data["location"][1]
angle = data["angle"] / PI * 180 + 90 # rotate 90 degrees to match field
length = data["length"]
width = data["width"]
# get bottom left position
x_offset = width / 2
y_offset = length / 2
theta_rad = math.radians(angle)
dx = x_offset * math.cos(theta_rad) - y_offset * math.sin(theta_rad)
dy = x_offset * math.sin(theta_rad) + y_offset * math.cos(theta_rad)
x = x_center - dx
y = y_center - dy
if self.agent_shape == "rectangle" or self.agent_shape == "segment":
# Create a rectangle
rect = patches.Rectangle(
(x, y), width, length, angle=angle, linewidth=1, edgecolor="b"
)
self.ax.add_patch(rect)
elif self.agent_shape == "circle":
# Create a circle
circle = patches.Circle((x, y), width, linewidth=1, edgecolor="b")
self.ax.add_patch(circle)
bounds = remove_units(self.bounds)
plt.xlim([0, bounds[0]])
plt.ylim([0, bounds[1]])
plt.draw()
plt.pause(0.01)
# configs
[docs]
def make_random_position(bounds):
unitless_bounds = [bound.to(units.um).magnitude for bound in bounds]
return [
np.random.uniform(0, unitless_bounds[0]),
np.random.uniform(0, unitless_bounds[1]),
] * units.um
[docs]
def single_agent_config(config):
# cell dimensions
width = 1.0 * units.um
length = 2.0 * units.um
volume = volume_from_length(length, width)
bounds = config.get("bounds", DEFAULT_BOUNDS)
location = config.get("location")
if location:
location = [loc * bounds[n] for n, loc in enumerate(location)]
else:
location = make_random_position(bounds)
return {
"boundary": {
"location": location,
"angle": np.random.uniform(0, 2 * PI),
"volume": volume,
"length": length,
"width": width,
"mass": 1339 * units.fg,
"thrust": 0,
"torque": 0,
}
}
[docs]
def agent_body_config(config):
agent_ids = config["agent_ids"]
agent_config = {agent_id: single_agent_config(config) for agent_id in agent_ids}
return {"agents": agent_config}
default_gd_config = {"bounds": DEFAULT_BOUNDS}
default_gd_config.update(
agent_body_config({"bounds": DEFAULT_BOUNDS, "agent_ids": ["1", "2"]})
)
[docs]
class InvokeUpdate(object):
def __init__(self, update):
self.update = update
[docs]
def get(self, timeout=0):
return self.update
# tests and simulations
def test_multibody(n_agents=1, time=10, return_data=False):
agent_ids = [str(agent_id) for agent_id in range(n_agents)]
multibody_config = {
"agents": agent_body_config({"bounds": DEFAULT_BOUNDS, "agent_ids": agent_ids})
}
multibody = Multibody(multibody_config)
# initialize agent's boundary state
initial_agents_state = multibody_config["agents"]
initial_state = {"agents": initial_agents_state}
experiment = process_in_experiment(multibody, initial_state=initial_state)
# run experiment
settings = {"timestep": 1, "total_time": time, "return_raw_data": True}
data = simulate_experiment(experiment, settings)
if return_data:
return data
def test_growth_division(
config=default_gd_config,
growth_rate=0.05,
growth_rate_noise=0.001,
division_volume=0.4**3 * units.fL,
total_time=10,
timestep=1,
experiment_settings=None,
return_data=False,
):
if not experiment_settings:
experiment_settings = {}
initial_agents_state = config["agents"]
# make the process
multibody = Multibody(config)
experiment = process_in_experiment(multibody, experiment_settings)
experiment.state._update_subschema(
("agents",),
{
"boundary": {
"mass": {"_updater": "set", "_divider": "split"},
"length": {"_updater": "set", "_divider": "split"},
"volume": {"_updater": "set", "_divider": "split"},
}
},
)
experiment.state._apply_subschemas()
# make initial agent state
experiment.state.set_value({"agents": initial_agents_state})
agents_store = experiment.state.get_path(["agents"])
# emit initial state
experiment._emit_store_data()
# run simulation
time = 0
while time < total_time:
experiment.update(timestep)
time += timestep
agents_state = agents_store.get_value()
invoked_update = []
for agent_id, state in agents_state.items():
state = state["boundary"]
length = state["length"]
width = state["width"]
mass = state["mass"] # .magnitude
# update
growth_rate2 = (
growth_rate + np.random.normal(0.0, growth_rate_noise)
) * timestep
new_mass = mass + mass * growth_rate2
new_length = length + length * growth_rate2
new_volume = volume_from_length(new_length, width)
if new_volume > division_volume:
daughter_ids = [str(agent_id) + "0", str(agent_id) + "1"]
daughter_updates = []
for daughter_id in daughter_ids:
daughter_updates.append(
{
"key": daughter_id,
"processes": {},
"topology": {},
"initial_state": {},
}
)
update = {
"_divide": {"mother": agent_id, "daughters": daughter_updates}
}
else:
update = {
agent_id: {
"boundary": {
"volume": new_volume,
"length": new_length,
"mass": new_mass,
}
}
} # * units.fg
invoked_update.append((InvokeUpdate({"agents": update}), None))
# update experiment
experiment._send_updates(invoked_update)
experiment.end()
if return_data:
return experiment.emitter.get_data_unitless()
[docs]
def run_growth_division(
out_dir="out",
animate=True,
):
n_agents = 2
agent_ids = [str(agent_id) for agent_id in range(n_agents)]
# configure the multibody process
bounds = DEFAULT_BOUNDS
multibody_config = {
"animate": animate,
# 'jitter_force': 1e0,
"bounds": bounds,
}
body_config = {"bounds": bounds, "agent_ids": agent_ids}
multibody_config.update(agent_body_config(body_config))
# experiment settings
experiment_settings = {"progress_bar": False, "display_info": False}
# run the test
gd_data = test_growth_division(
config=multibody_config,
growth_rate=0.05,
growth_rate_noise=0.001,
division_volume=volume_from_length(4, 1) * units.fL,
total_time=100,
experiment_settings=experiment_settings,
return_data=True,
)
agents, fields = format_snapshot_data(gd_data)
return plot_snapshots(bounds, agents=agents, fields=fields, out_dir=out_dir)
def test_daughter_locations():
locations = daughter_locations(
[2, 2],
{"length": 2 * math.sqrt(2), "angle": -math.pi / 4},
)
assert locations == [[1.5, 2.5], [2.5, 1.5]]
[docs]
def main():
out_dir = os.path.join(PROCESS_OUT_DIR, NAME)
if not os.path.exists(out_dir):
os.makedirs(out_dir)
run_growth_division(out_dir)
if __name__ == "__main__":
main()