diff --git a/vivarium_readdy/__init__.py b/vivarium_readdy/__init__.py index 4808b39..8526c33 100644 --- a/vivarium_readdy/__init__.py +++ b/vivarium_readdy/__init__.py @@ -14,7 +14,11 @@ def get_module_version(): from .processes.readdy_process import ReaddyProcess # noqa: F401 -from vivarium.core.registry import emitter_registry # noqa: F401 -from .processes.simularium_emitter import SimulariumEmitter # noqa: F401 +from .util import monomer_ports_schema # noqa: F401 +from .util import create_monomer_update # noqa: F401 +from .util import agents_update # noqa: F401 -emitter_registry.register("simularium", SimulariumEmitter) +from .processes.simularium_monomer_emitter import SimulariumMonomerEmitter # noqa: F401 +from vivarium.core.registry import emitter_registry + +emitter_registry.register("simularium_monomers", SimulariumMonomerEmitter) diff --git a/vivarium_readdy/processes/readdy_process.py b/vivarium_readdy/processes/readdy_process.py index 1153803..dee4d6c 100644 --- a/vivarium_readdy/processes/readdy_process.py +++ b/vivarium_readdy/processes/readdy_process.py @@ -1,10 +1,13 @@ import numpy as np from vivarium.core.process import Process -from vivarium.core.engine import Engine, pf +from vivarium.core.engine import Engine from tqdm import tqdm import readdy +from pint import UnitRegistry + +from ..util import monomer_ports_schema, create_monomer_update NAME = "READDY" @@ -18,96 +21,37 @@ class ReaddyProcess(Process): name = NAME defaults = { - "internal_timestep": 0.1, # s - "box_size": 100.0, # m + "internal_timestep": 0.1, + "box_size": 100.0, "periodic_boundary": False, "temperature_C": 22.0, - "viscosity": 8.1, # cP, viscosity in cytoplasm + "viscosity": 1.0, # cP "force_constant": 250.0, "n_cpu": 4, "particle_radii": {}, + "topology_types": [], "topology_particles": [], + "bond_pairs": [], "reactions": [], + "time_units": "s", + "spatial_units": "m", } def __init__(self, parameters=None): super(ReaddyProcess, self).__init__(parameters) + self.create_readdy_system() def ports_schema(self): - return { - "box_size": { - "_default": 100.0, - "_updater": "set", - "_emit": True, - }, - "topologies": { - "*": { - "type_name": { - "_default": "", - "_updater": "set", - "_emit": True, - }, - "particle_ids": { - "_default": [], - "_updater": "set", - "_emit": True, - }, - } - }, - "particles": { - "*": { - "type_name": { - "_default": "", - "_updater": "set", - "_emit": True, - }, - "position": { - "_default": np.zeros(3), - "_updater": "set", - "_emit": True, - }, - "neighbor_ids": { - "_default": [], - "_updater": "set", - "_emit": True, - }, - } - }, - } + return monomer_ports_schema def next_update(self, timestep, states): - self.create_readdy_system(states) self.create_readdy_simulation() - self.add_particle_instances(states) + self.add_particle_instances(states["monomers"]) self.simulate_readdy(timestep) - current_state = self.get_current_state() - # update topologies - topologies_update = {"_add": [], "_delete": []} - for id, state in current_state["topologies"].items(): - if id in states["topologies"]: - topologies_update[id] = state - else: - topologies_update["_add"].append({"key": id, "state": state}) - for existing_id in states["topologies"].keys(): - if existing_id not in current_state["topologies"]: - topologies_update["_delete"].append(existing_id) - # update particles - particles_update = {"_add": [], "_delete": []} - for id, state in current_state["particles"].items(): - if id in states["particles"]: - particles_update[id] = state - else: - particles_update["_add"].append({"key": id, "state": state}) - for existing_id in states["particles"].keys(): - if existing_id not in current_state["particles"]: - particles_update["_delete"].append(existing_id) - return { - "box_size": states["box_size"], - "topologies": topologies_update, - "particles": particles_update, - } + new_monomers = self.get_current_monomers() + return create_monomer_update(states["monomers"], new_monomers) - def create_readdy_system(self, states): + def create_readdy_system(self): """ Create the ReaDDy system including particle species, constraints, and reactions @@ -120,30 +64,30 @@ def create_readdy_system(self, states): self.parameters["temperature_K"] = self.parameters["temperature_C"] + 273.15 self.system.temperature = self.parameters["temperature_K"] self.add_particle_species() - self.add_topology_types(states) - all_particle_types = set() - for particle_id in states["particles"]: - particle = states["particles"][particle_id] - all_particle_types.add(particle["type_name"]) - self.check_add_global_box_potential(states, all_particle_types) + self.add_topology_types() + all_particle_types = self.parameters["particle_radii"].keys() + self.check_add_global_box_potential(all_particle_types) self.add_repulsions(all_particle_types) - self.add_bonds(states) + self.add_bonds() self.add_reactions() @staticmethod - def calculate_diffusionCoefficient(radius, eta, T): + def calculate_diffusionCoefficient(radius, viscosity, temperature, spatial_units): """ calculate the theoretical diffusion constant of a spherical particle - with radius [m] - in a media with viscosity eta [cP] - at temperature T [Kelvin] + with radius [spatial_units] + in a media with viscosity [cP] + at temperature [Kelvin] - returns m^2/s + returns [spatial_units^2/s] """ + ureg = UnitRegistry() + convert_to_nm = ureg(spatial_units).to("m").magnitude return ( - (1.38065 * 10 ** (-23) * T) - / (6 * np.pi * eta * 10 ** (-3) * radius * 10 ** (-9)) + (1.38065 * 10 ** (-23) * temperature) + / (6 * np.pi * viscosity * 10 ** (-3) * radius * convert_to_nm) / 10**9 + / (convert_to_nm * convert_to_nm) ) def add_particle_species(self): @@ -158,6 +102,7 @@ def add_particle_species(self): self.parameters["particle_radii"][particle_name], self.parameters["viscosity"], self.parameters["temperature_K"], + self.parameters["spatial_units"], ) if particle_name in self.parameters["topology_particles"]: self.system.add_topology_species(particle_name, diffCoeff) @@ -165,17 +110,14 @@ def add_particle_species(self): self.system.add_species(particle_name, diffCoeff) added_particle_types.append(particle_name) - def add_topology_types(self, states): + def add_topology_types(self): """ Add all topology types """ - topology_types = set() - for topology_id in states["topologies"]: - topology_types.add(states["topologies"][topology_id]["type_name"]) - for topology_type in topology_types: + for topology_type in self.parameters["topology_types"]: self.system.topologies.add_type(topology_type) - def check_add_global_box_potential(self, states, all_particle_types): + def check_add_global_box_potential(self, all_particle_types): """ If the boundaries are not periodic, add a box potential for all particles @@ -199,51 +141,31 @@ def add_repulsions(self, all_particle_types): to enforce volume exclusion """ for type1 in all_particle_types: - if type1 not in self.parameters["particle_radii"]: - raise Exception( - "Please provide a radius for particle type " - f"{type1} in parameters['particle_radii']" - ) for type2 in all_particle_types: - if type2 not in self.parameters["particle_radii"]: - raise Exception( - "Please provide a radius for particle type " - f"{type1} in parameters['particle_radii']" - ) self.system.potentials.add_harmonic_repulsion( type1, type2, force_constant=self.parameters["force_constant"], interaction_distance=( - self.parameters["particle_radii"][type1] - + self.parameters["particle_radii"][type2] + self.parameters["particle_radii"].get(type1, 1.0) + + self.parameters["particle_radii"].get(type2, 1.0) ), ) - def add_bonds(self, states): + def add_bonds(self): """ Add harmonic bonds """ - added_bonds = [] - for particle_id in states["particles"]: - particle = states["particles"][particle_id] - for neighbor_id in particle["neighbor_ids"]: - neighbor = states["particles"][neighbor_id] - type1 = particle["type_name"] - type2 = neighbor["type_name"] - if (type1, type2) in added_bonds or (type2, type1) in added_bonds: - continue - self.system.topologies.configure_harmonic_bond( - type1, - type2, - force_constant=self.parameters["force_constant"], - length=( - self.parameters["particle_radii"][type1] - + self.parameters["particle_radii"][type2] - ), - ) - added_bonds.append((type1, type2)) - added_bonds.append((type2, type1)) + for bond_pair in self.parameters["bond_pairs"]: + self.system.topologies.configure_harmonic_bond( + bond_pair[0], + bond_pair[1], + force_constant=self.parameters["force_constant"], + length=( + self.parameters["particle_radii"].get(bond_pair[0], 1.0) + + self.parameters["particle_radii"].get(bond_pair[1], 1.0) + ), + ) def add_reactions(self): """ @@ -259,19 +181,19 @@ def create_readdy_simulation(self): self.simulation = self.system.simulation("CPU") self.simulation.kernel_configuration.n_threads = self.parameters["n_cpu"] - def add_particle_instances(self, states): + def add_particle_instances(self, monomers): """ Add particle instances to the simulation """ # add topology particles topology_particle_ids = [] - for topology_id in states["topologies"]: - topology = states["topologies"][topology_id] + for topology_id in monomers["topologies"]: + topology = monomers["topologies"][topology_id] topology_particle_ids += topology["particle_ids"] types = [] positions = [] for particle_id in topology["particle_ids"]: - particle = states["particles"][particle_id] + particle = monomers["particles"][particle_id] types.append(particle["type_name"]) positions.append(particle["position"]) top = self.simulation.add_topology( @@ -279,7 +201,7 @@ def add_particle_instances(self, states): ) added_edges = [] for index, particle_id in enumerate(topology["particle_ids"]): - for neighbor_id in states["particles"][particle_id]["neighbor_ids"]: + for neighbor_id in monomers["particles"][particle_id]["neighbor_ids"]: neighbor_index = topology["particle_ids"].index(neighbor_id) if (index, neighbor_index) not in added_edges and ( neighbor_index, @@ -290,10 +212,10 @@ def add_particle_instances(self, states): added_edges.append((index, neighbor_index)) added_edges.append((neighbor_index, index)) # add non-topology particles - for particle_id in states["particles"]: + for particle_id in monomers["particles"]: if particle_id in topology_particle_ids: continue - particle = states["particles"][particle_id] + particle = monomers["particles"][particle_id] self.simulation.add_particle( type=particle["type_name"], position=particle["position"] ) @@ -332,7 +254,7 @@ def loop(): calculate_forces() observe(t) - self.simulation._run_custom_loop(loop) + self.simulation._run_custom_loop(loop, show_summary=False) def current_particle_edges(self): """ @@ -349,7 +271,7 @@ def current_particle_edges(self): result.append((p1_id, p2_id)) return result - def get_current_state(self): + def get_current_monomers(self): """ Get data for topologies of particles from readdy.simulation.current_topologies @@ -363,18 +285,19 @@ def get_current_state(self): edges = self.current_particle_edges() for index, topology in enumerate(self.simulation.current_topologies): particle_ids = [] - for p in topology.particles: - particle_ids.append(p.id) + for p_ix, particle in enumerate(topology.particles): + particle_ids.append(particle.id) neighbor_ids = [] for edge in edges: - if p.id == edge[0]: + if particle.id == edge[0]: neighbor_ids.append(edge[1]) - elif p.id == edge[1]: + elif particle.id == edge[1]: neighbor_ids.append(edge[0]) - result["particles"][p.id] = { - "type_name": p.type, - "position": p.pos, + result["particles"][p_ix] = { + "type_name": particle.type, + "position": particle.pos, "neighbor_ids": neighbor_ids, + "radius": self.parameters["particle_radii"].get(particle.type, 1.0), } result["topologies"][index] = { "type_name": topology.type, @@ -382,10 +305,11 @@ def get_current_state(self): } # non-topology particles for index, particle in enumerate(self.simulation.current_particles): - result["particles"][p.id] = { - "type_name": p.type, - "position": p.pos, + result["particles"][index] = { + "type_name": particle.type, + "position": particle.pos, "neighbor_ids": [], + "radius": self.parameters["particle_radii"].get(particle.type, 1.0), } return result @@ -417,6 +341,7 @@ def initial_state(self, config=None): "type_name": "C", "position": position, "neighbor_ids": [], + "radius": 2.0, } last_id += 1 # inert chain particles @@ -433,22 +358,25 @@ def initial_state(self, config=None): "type_name": "D", "position": chain_position, "neighbor_ids": neighbor_ids, + "radius": 4.0, } chain_particle_ids.append(particle_id) chain_position += 2 * chain_particle_radius * np.random.uniform(3) return { - "box_size": box_size, - "topologies": { - 0: { - "type_name": "Chain", - "particle_ids": chain_particle_ids, - } + "monomers": { + "box_size": box_size, + "topologies": { + 0: { + "type_name": "Chain", + "particle_ids": chain_particle_ids, + } + }, + "particles": particles, }, - "particles": particles, } -def test_readdy_process(): +def run_readdy_process(): readdy_process = ReaddyProcess( { "particle_radii": { @@ -457,9 +385,15 @@ def test_readdy_process(): "C": 2.0, "D": 4.0, }, + "topology_types": [ + "Chain", + ], "topology_particles": [ "D", ], + "bond_pairs": [ + ["D", "D"], + ], "reactions": [ { "descriptor": "enz: A +(3) C -> B + C", @@ -468,22 +402,15 @@ def test_readdy_process(): ], } ) + composite = readdy_process.generate() engine = Engine( - processes={"readdy": readdy_process}, - topology={ - "readdy": { - "box_size": ("box_size",), - "topologies": ("topologies",), - "particles": ("particles",), - } - }, + composite=composite, initial_state=readdy_process.initial_state(), emitter="simularium", ) engine.update(1.0) # 10 steps - output = engine.emitter.get_data() - print(pf(output)) + engine.emitter.get_data() if __name__ == "__main__": - test_readdy_process() + run_readdy_process() diff --git a/vivarium_readdy/processes/simularium_emitter.py b/vivarium_readdy/processes/simularium_monomer_emitter.py similarity index 56% rename from vivarium_readdy/processes/simularium_emitter.py rename to vivarium_readdy/processes/simularium_monomer_emitter.py index e685519..2c9b746 100644 --- a/vivarium_readdy/processes/simularium_emitter.py +++ b/vivarium_readdy/processes/simularium_monomer_emitter.py @@ -1,4 +1,4 @@ -from typing import Any, Dict, Tuple +from typing import Any, Dict, List from vivarium.core.emitter import Emitter @@ -10,22 +10,16 @@ AgentData, MetaData, UnitData, + DisplayData, + DISPLAY_TYPE, ) -# TODO add viz type to state? -VIZ_FOR_CHOICE = { - "medyan_active": "fibers", - "cytosim_active": "fibers", - "readdy_active": "monomers", - "init": "fibers", -} - -class SimulariumEmitter(Emitter): +class SimulariumMonomerEmitter(Emitter): def __init__(self, config: Dict[str, str]) -> None: super().__init__(config) self.configuration_data = None - self.saved_data: Dict[float, Dict[str, Any]] = {} + self.saved_data: List[Dict[str, Any]] = [] def emit(self, data: Dict[str, Any]) -> None: """ @@ -37,45 +31,9 @@ def emit(self, data: Dict[str, Any]) -> None: self.configuration_data = data["data"] assert "processes" in self.configuration_data, "please emit processes" if data["table"] == "history": - emit_data = data["data"] - time = emit_data["time"] - self.saved_data[time] = { - key: value for key, value in emit_data.items() if key not in ["time"] - } - - def get_simularium_fibers( - self, time, fibers, actin_radius, position_offset, trajectory - ): - """ - Shape fiber state data into Simularium fiber agents - """ - print(f" {position_offset}") - n_agents = 0 - time_index = len(trajectory["times"]) - trajectory["times"].append(time) - trajectory["unique_ids"].append([]) - trajectory["type_names"].append([]) - trajectory["n_subpoints"].append([]) - trajectory["subpoints"].append([]) - for fiber_id in fibers: - fiber = fibers[fiber_id] - fiber_points = [ - np.array(point) + position_offset for point in fiber["points"] - ] - n_agents += 1 - trajectory["unique_ids"][time_index].append(int(fiber_id)) - trajectory["type_names"][time_index].append(fiber["type_name"]) - trajectory["n_subpoints"][time_index].append(len(fiber_points)) - trajectory["subpoints"][time_index].append(fiber_points) - trajectory["n_agents"].append(n_agents) - trajectory["viz_types"].append(n_agents * [1001.0]) - trajectory["positions"].append(n_agents * [[0.0, 0.0, 0.0]]) - trajectory["radii"].append(n_agents * [actin_radius]) - return trajectory + self.saved_data.append(data["data"]) - def get_simularium_monomers( - self, time, monomers, actin_radius, position_offset, trajectory - ): + def get_simularium_monomers(self, time, monomers, trajectory): """ Shape monomer state data into Simularium agents """ @@ -84,15 +42,21 @@ def get_simularium_monomers( trajectory["unique_ids"].append([]) trajectory["type_names"].append([]) trajectory["positions"].append([]) + trajectory["radii"].append([]) edge_ids = [] edge_positions = [] - for particle_id in monomers.get("particles", {}): + for particle_id in monomers["particles"]: particle = monomers["particles"][particle_id] trajectory["unique_ids"][time_index].append(int(particle_id)) trajectory["type_names"][time_index].append(particle["type_name"]) - trajectory["positions"][time_index].append( - np.array(particle["position"]) + position_offset - ) + # HACK needed until simulariumio default display data fix + if particle["type_name"] not in trajectory["display_data"]: + trajectory["display_data"][particle["type_name"]] = DisplayData( + name=particle["type_name"], + display_type=DISPLAY_TYPE.SPHERE, + ) + trajectory["positions"][time_index].append(np.array(particle["position"])) + trajectory["radii"][time_index].append(particle["radius"]) # visualize edges between particles for neighbor_id in particle["neighbor_ids"]: neighbor_id_str = str(neighbor_id) @@ -105,9 +69,10 @@ def get_simularium_monomers( edge_ids.append(edge) edge_positions.append( [ - np.array(particle["position"]) + position_offset, - np.array(monomers["particles"][neighbor_id_str]["position"]) - + position_offset, + np.array(particle["position"]), + np.array( + monomers["particles"][neighbor_id_str]["position"] + ), ] ) n_agents = len(trajectory["unique_ids"][time_index]) @@ -118,7 +83,6 @@ def get_simularium_monomers( trajectory["unique_ids"][time_index] += [1000 + i for i in range(n_edges)] trajectory["type_names"][time_index] += ["edge" for edge in range(n_edges)] trajectory["positions"][time_index] += n_edges * [[0.0, 0.0, 0.0]] - trajectory["radii"].append(n_agents * [actin_radius]) trajectory["radii"][time_index] += n_edges * [1.0] trajectory["n_subpoints"].append(n_agents * [0]) trajectory["n_subpoints"][time_index] += n_edges * [2] @@ -143,7 +107,9 @@ def jagged_3d_list_to_numpy_array(jagged_3d_list): """ Shape a jagged list with 3 dimensions to a numpy array """ - df = SimulariumEmitter.fill_df(pd.DataFrame(jagged_3d_list), [0.0, 0.0, 0.0]) + df = SimulariumMonomerEmitter.fill_df( + pd.DataFrame(jagged_3d_list), [0.0, 0.0, 0.0] + ) df_t = df.transpose() exploded = [df_t[col].explode() for col in list(df_t.columns)] return np.array(exploded).reshape((df.shape[0], df.shape[1], 3)) @@ -158,7 +124,7 @@ def get_subpoints_numpy_array(trajectory) -> np.ndarray: max_subpoints = 0 total_steps = len(trajectory["subpoints"]) for time_index in range(total_steps): - frame_array = SimulariumEmitter.jagged_3d_list_to_numpy_array( + frame_array = SimulariumMonomerEmitter.jagged_3d_list_to_numpy_array( trajectory["subpoints"][time_index] ) if frame_array.shape[0] > max_agents: @@ -184,70 +150,60 @@ def get_agent_data_from_jagged_lists(trajectory, scale_factor) -> AgentData: Shape a dictionary of jagged lists into a Simularium AgentData object """ return AgentData( - times=np.arange(len(trajectory["times"])), + times=trajectory["times"], n_agents=np.array(trajectory["n_agents"]), - viz_types=SimulariumEmitter.fill_df( + viz_types=SimulariumMonomerEmitter.fill_df( pd.DataFrame(trajectory["viz_types"]), 1000.0 ).to_numpy(), - unique_ids=SimulariumEmitter.fill_df( + unique_ids=SimulariumMonomerEmitter.fill_df( pd.DataFrame(trajectory["unique_ids"]), 0 ).to_numpy(dtype=int), types=trajectory["type_names"], positions=scale_factor - * SimulariumEmitter.jagged_3d_list_to_numpy_array(trajectory["positions"]), + * SimulariumMonomerEmitter.jagged_3d_list_to_numpy_array( + trajectory["positions"] + ), radii=scale_factor - * SimulariumEmitter.fill_df( + * SimulariumMonomerEmitter.fill_df( pd.DataFrame(trajectory["radii"]), 0.0 ).to_numpy(), - n_subpoints=SimulariumEmitter.fill_df( + n_subpoints=SimulariumMonomerEmitter.fill_df( pd.DataFrame(trajectory["n_subpoints"]), 0 ).to_numpy(dtype=int), subpoints=scale_factor - * SimulariumEmitter.get_subpoints_numpy_array(trajectory), + * SimulariumMonomerEmitter.get_subpoints_numpy_array(trajectory), + display_data=trajectory["display_data"], ) @staticmethod def get_simularium_converter( - trajectory, box_dimensions, scale_factor + trajectory, box_dimensions, scale_factor, time_units, spatial_units ) -> TrajectoryConverter: """ Shape a dictionary of jagged lists into a Simularium TrajectoryData object and provide it to a TrajectoryConverter for conversion """ - spatial_units = UnitData("nm") spatial_units.multiply(1 / scale_factor) return TrajectoryConverter( TrajectoryData( meta_data=MetaData( box_size=scale_factor * box_dimensions, ), - agent_data=SimulariumEmitter.get_agent_data_from_jagged_lists( + agent_data=SimulariumMonomerEmitter.get_agent_data_from_jagged_lists( trajectory, scale_factor ), - time_units=UnitData("count"), + time_units=time_units, spatial_units=spatial_units, ) ) - @staticmethod - def get_active_choice(choices) -> Tuple[str, bool]: - """ - Determine the active simulator - """ - for choice in choices: - if choices[choice]: - return choice - return None - def get_data(self) -> dict: """ Save the accumulated timeseries history of "emitted" data to file """ - if "readdy_actin" in self.configuration_data: - actin_radius = self.configuration_data["readdy_actin"]["actin_radius"] - else: - actin_radius = 3.0 # TODO add to MEDYAN/Cytosim configs - box_dimensions = None + box_size = None + time_units = None + spatial_units = None trajectory = { "times": [], "n_agents": [], @@ -258,37 +214,22 @@ def get_data(self) -> dict: "radii": [], "n_subpoints": [], "subpoints": [], + "display_data": {}, } - times = list(self.saved_data.keys()) - times.sort() - vizualize_time_index = 0 - for time, state in self.saved_data.items(): - index = times.index(time) - current_choice = SimulariumEmitter.get_active_choice(state["choices"]) - if box_dimensions is None and "fibers_box_extent" in state: - box_dimensions = np.array(state["fibers_box_extent"]) - if VIZ_FOR_CHOICE[current_choice] == "fibers": - center_fibers = ( - index == 0 or "readdy_active" in state["choices"] - ) # TODO generalize - trajectory = self.get_simularium_fibers( - vizualize_time_index, - state["fibers"], - actin_radius, - -0.5 * box_dimensions if center_fibers else np.zeros(3), - trajectory, - ) - vizualize_time_index += 1 - if VIZ_FOR_CHOICE[current_choice] == "monomers": - trajectory = self.get_simularium_monomers( - vizualize_time_index, - state["monomers"], - actin_radius, - (np.array(state["monomers"]["box_center"]) - 0.5 * box_dimensions), - trajectory, - ) - vizualize_time_index += 1 - simularium_converter = SimulariumEmitter.get_simularium_converter( - trajectory, box_dimensions, 0.1 + for state in self.saved_data: + if box_size is None: + box_size = state["monomers"]["box_size"] + if time_units is None: + time_units = state["time_units"] + if spatial_units is None: + spatial_units = state["spatial_units"] + trajectory = self.get_simularium_monomers( + state["time"], + state["monomers"], + trajectory, + ) + simularium_converter = SimulariumMonomerEmitter.get_simularium_converter( + trajectory, box_size, 0.1, UnitData(time_units), UnitData(spatial_units) ) - simularium_converter.write_JSON("out/actin_test") + simularium_converter.save("out/test") + return {"simularium_converter": simularium_converter} diff --git a/vivarium_readdy/util.py b/vivarium_readdy/util.py new file mode 100644 index 0000000..0395615 --- /dev/null +++ b/vivarium_readdy/util.py @@ -0,0 +1,94 @@ +import numpy as np + + +monomer_ports_schema = { + "time_units": { + "_default": "s", + "_updater": "set", + "_emit": True, + }, + "spatial_units": { + "_default": "m", + "_updater": "set", + "_emit": True, + }, + "monomers": { + "box_center": { + "_default": np.array([0.0, 0.0, 0.0]), + "_updater": "set", + "_emit": True, + }, + "box_size": { + "_default": 500.0, + "_updater": "set", + "_emit": True, + }, + "topologies": { + "*": { + "type_name": { + "_default": "", + "_updater": "set", + "_emit": True, + }, + "particle_ids": { + "_default": [], + "_updater": "set", + "_emit": True, + }, + } + }, + "particles": { + "*": { + "type_name": { + "_default": "", + "_updater": "set", + "_emit": True, + }, + "position": { + "_default": np.zeros(3), + "_updater": "set", + "_emit": True, + }, + "neighbor_ids": { + "_default": [], + "_updater": "set", + "_emit": True, + }, + "radius": { + "_default": 1.0, + "_updater": "set", + "_emit": True, + }, + } + }, + }, +} + + +def agents_update(existing, projected): + update = {"_add": [], "_delete": []} + + for id, state in projected.items(): + if id in existing: + update[id] = state + else: + update["_add"].append({"key": id, "state": state}) + + for existing_id in existing.keys(): + if existing_id not in projected: + update["_delete"].append(existing_id) + + return update + + +def create_monomer_update(previous_monomers, new_monomers): + return { + "monomers": { + "topologies": agents_update( + previous_monomers["topologies"], new_monomers["topologies"] + ), + "particles": agents_update( + previous_monomers["particles"], new_monomers["particles"] + ), + } + }