Skip to content

Enable colony simulation state saving and custom emit_step configuration - #442

Open
katha815 wants to merge 6 commits into
CovertLab:colony-sim-emitfrom
katha815:pr-colony_sim_emit-fix
Open

Enable colony simulation state saving and custom emit_step configuration#442
katha815 wants to merge 6 commits into
CovertLab:colony-sim-emitfrom
katha815:pr-colony_sim_emit-fix

Conversation

@katha815

Copy link
Copy Markdown

Summary

This PR enables proper colony simulation state saving and Parquet emission finalization for both dividing and non-dividing runs. It also ensures the emit_step configuration is correctly propagated to both outer and inner emitters.

All tests were run with the following command structure:

python ecoli/experiments/ecoli_engine_process.py --config configs/{customised_file}.json

Fix Details

Fix 0 — Convert Two-Component System to Vivarium PartitionedProcess

While investigating the colony state saving issue (Fix 1), I found that the Two-Component System process was still using the standalone interface and was not fully compatible with the current Vivarium framework. This may cause compatibility issues when the simulation attempted to save and restore colony states.

To address this, I migrated the process to PartitionedProcess:

  • Registered process name and topology with topology_registry.
  • Implemented ports_schema(), calculate_request(), and evolve_state() methods.
  • Replaced direct state access with Vivarium's bulk_name_to_idx() and counts() helpers.
  • Preserved the existing ODE, stoichiometry, and molecule-allocation logic.

This migration was necessary for the colony simulation to handle states correctly through the PartitionedProcess interface.

Fix 1 — Fix colony final-state JSON saving

  • Problem: The problem was identified using configs/colony_baseline_test.json. The colony's final state JSON was not being saved correctly. The Parquet emitter wraps agents in an agents/outer structure, but colony_save_states() expected agents/<agent_id>/... directly.
  • Solution: Detect emitter type and use correct internal path:
    • Parquet: agents/outer/agents/<agent_id>/cell_process
    • Other emitters: agents/<agent_id>/cell_process
  • Extract state_to_save["agents"]["outer"] before saving so colony JSON retains expected structure.
  • Preserve bulk_dtypes and unique_dtypes metadata.
  • Verified with: configs/colony_baseline_test.json (2-generation dividing run, 1-second smoke test)
  • Result: Final JSON loads successfully with expected agents/<agent_id>/... structure.

Fix 2 — Finalize outer Parquet emitter on shutdown

  • Problem: The problem was identified using configs/colony_baseline_test3.json. The emission step remained at 1 second regardless of custom configuration. The outer ParquetEmitter buffers writes in batches of 400 emits, and runs that don't reach this threshold leave buffered data unwritten.
  • Solution: After engine.end() in run_simulation(), mark the outer Parquet emitter as successful and explicitly call engine.emitter.finalize().
  • Verified with: configs/colony_baseline_test4.json (dividing run, 60-second emit step)
  • Result: Outer Parquet output now contains emission points at 0s, 60s, and 120s.

Fix 3 — Finalize inner Parquet emitters and propagate emit_step

  • Problem: The outer subfolder (/out/{project}/...) saved Parquet data, but the inner subfolder (/out/{project}/...__inner) either saved data every 1 second or did not save at all. This occurred because the inner ParquetEmitter was only finalized from division or exception paths inside EngineProcess.next_update(), so short non-dividing runs left inner Parquet data unwritten.
  • Solution:
    • Pass emit_step from top-level simulation configuration through EcoliEngineProcess into the inner EngineProcess.
    • Add finalize_parquet_emitters() helper that recursively traverses nested process dictionaries and calls finalize() on every ParquetEmitter while avoiding duplicate finalization.
    • After finalizing the outer emitter, recursively finalize inner emitters on normal shutdown.
  • Verified with: configs/colony_baseline_test4.json (short non-dividing run, 60-second emit step)
  • Result: Both outer and inner Parquet outputs contain emission points at 0s, 60s, and 120s, confirming emit_step propagation and proper finalization for non-dividing runs.

Test Configurations Used

Fix Config File Run Type Emit Step Purpose
Fix 1 configs/colony_baseline_test.json Dividing (2 generations) 1s Verify final colony JSON structure
Fix 2 configs/colony_baseline_test3.json Dividing (2 generations) 60s Verify emit_step propagation
Fix 3 configs/colony_baseline_test4.json Non-dividing (short) 60s Verify outer emitter finalization, inner emitter finalization, and emit_step propagation

All three test configurations produced the expected outputs. The final JSON from Fix 1 has the correct agents/<agent_id>/... structure. The Parquet outputs from Fix 2 and Fix 3 contain emissions at 0s, 60s, and 120s.

Files Changed

  • ecoli/experiments/ecoli_engine_process.py — Main fixes for state saving and emitter finalization
  • Two-Component System process file — Migration to PartitionedProcess
  • configs/colony_baseline_test.json, configs/colony_baseline_test3.json, configs/colony_baseline_test4.json — Test configurations

Notes

The current implementation prioritises correctness over code brevity. Additional simplification may be considered in a follow-up.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant