11"""Estimate command for predicting simulation costs before running."""
22
3+ import re
4+
35import typer
46
57from ..app import app , console , get_study_path
68from ..study import StudyContext , detect_study_folder , parse_version_ref
79from ..utils import Output , ExitCode
810
911
12+ def _parse_base_population_ref (ref : str ) -> tuple [str , int | None ]:
13+ """Parse base_population reference like 'population.v2'."""
14+ match = re .match (r"^(.+)\.v(\d+)$" , ref )
15+ if match :
16+ return match .group (1 ), int (match .group (2 ))
17+ return ref , None
18+
19+
1020@app .command ("estimate" )
1121def estimate_command (
1222 scenario : str = typer .Option (
@@ -85,8 +95,14 @@ def estimate_command(
8595 out .error (f"Failed to load scenario: { e } " )
8696 raise typer .Exit (1 )
8797
88- # Load population spec (resolve relative path from scenario meta)
89- pop_path = study_ctx .root / scenario_spec .meta .population_spec
98+ # Load population spec (resolve from base_population or population_spec)
99+ base_pop_ref = scenario_spec .meta .base_population or scenario_spec .meta .population_spec
100+ if not base_pop_ref :
101+ out .error ("Scenario has no base_population or population_spec reference" )
102+ raise typer .Exit (1 )
103+
104+ pop_name , pop_version = _parse_base_population_ref (base_pop_ref )
105+ pop_path = study_ctx .get_population_path (pop_name , pop_version )
90106 if not pop_path .exists ():
91107 out .error (f"Population spec not found: { pop_path } " )
92108 raise typer .Exit (1 )
@@ -96,22 +112,20 @@ def estimate_command(
96112 out .error (f"Failed to load population spec: { e } " )
97113 raise typer .Exit (1 )
98114
99- # Load agents and network from study DB using scenario meta IDs
100- pop_id = scenario_spec .meta .population_id
101- net_id = scenario_spec .meta .network_id
115+ # Load agents and network from study DB using scenario ID
102116 with open_study_db (study_db ) as db :
103- agents = db .get_agents ( pop_id )
104- network = db .get_network (net_id )
117+ agents = db .get_agents_by_scenario ( scenario_name )
118+ network = db .get_network (scenario_name )
105119
106120 if not agents :
107121 out .error (
108- f"No agents found for population '{ pop_id } '. "
122+ f"No agents found for scenario '{ scenario_name } '. "
109123 f"Run 'extropy sample -s { scenario_name } ' first." ,
110124 )
111125 raise typer .Exit (1 )
112126 if not network .get ("edges" ):
113127 out .error (
114- f"No network found for network '{ net_id } '. "
128+ f"No network found for scenario '{ scenario_name } '. "
115129 f"Run 'extropy network -s { scenario_name } ' first." ,
116130 )
117131 raise typer .Exit (1 )
0 commit comments