Files
Atomizer/optimization_engine/interview/study_blueprint.py

559 lines
20 KiB
Python
Raw Normal View History

"""
Study Blueprint
Data structures for the study blueprint - the validated configuration
ready for study generation.
"""
from dataclasses import dataclass, field, asdict
from typing import Dict, List, Any, Optional
import json
@dataclass
class DesignVariable:
"""Design variable specification."""
parameter: str
current_value: float
min_value: float
max_value: float
units: Optional[str] = None
is_integer: bool = False
def to_dict(self) -> Dict[str, Any]:
return asdict(self)
def to_config_format(self) -> Dict[str, Any]:
"""Convert to optimization_config.json format."""
return {
"expression_name": self.parameter,
"bounds": [self.min_value, self.max_value],
"units": self.units or "",
"is_integer": self.is_integer,
}
@dataclass
class Objective:
"""Optimization objective specification."""
name: str
goal: str # minimize, maximize, target
extractor: str # Extractor ID (e.g., "E1", "E4")
extractor_name: Optional[str] = None
extractor_params: Optional[Dict[str, Any]] = None
weight: float = 1.0
target_value: Optional[float] = None # For target objectives
def to_dict(self) -> Dict[str, Any]:
return asdict(self)
def to_config_format(self) -> Dict[str, Any]:
"""Convert to optimization_config.json format."""
config = {
"name": self.name,
"type": self.goal,
"extractor": self.extractor,
"weight": self.weight,
}
if self.extractor_params:
config["extractor_params"] = self.extractor_params
if self.target_value is not None:
config["target"] = self.target_value
return config
@dataclass
class Constraint:
"""Optimization constraint specification."""
name: str
constraint_type: str # max, min
threshold: float
extractor: str # Extractor ID
extractor_name: Optional[str] = None
extractor_params: Optional[Dict[str, Any]] = None
is_hard: bool = True
penalty_weight: float = 1000.0 # For soft constraints
def to_dict(self) -> Dict[str, Any]:
return asdict(self)
def to_config_format(self) -> Dict[str, Any]:
"""Convert to optimization_config.json format."""
config = {
"name": self.name,
"type": self.constraint_type,
"threshold": self.threshold,
"extractor": self.extractor,
"hard": self.is_hard,
}
if self.extractor_params:
config["extractor_params"] = self.extractor_params
if not self.is_hard:
config["penalty_weight"] = self.penalty_weight
return config
@dataclass
class StudyBlueprint:
"""
Complete study blueprint ready for generation.
This is the validated configuration that will be used to create
the study files (optimization_config.json, run_optimization.py, etc.)
"""
# Study metadata
study_name: str
study_description: str = ""
interview_session_id: str = ""
# Model paths
model_path: str = ""
sim_path: str = ""
fem_path: str = ""
# Design space
design_variables: List[DesignVariable] = field(default_factory=list)
# Optimization goals
objectives: List[Objective] = field(default_factory=list)
constraints: List[Constraint] = field(default_factory=list)
# Optimization settings
protocol: str = "protocol_10_single" # or "protocol_11_multi"
n_trials: int = 100
sampler: str = "TPE"
use_neural_acceleration: bool = False
# Solver settings
solver_config: Dict[str, Any] = field(default_factory=dict)
solve_all_solutions: bool = True
# Extractors configuration
extractors_config: Dict[str, Any] = field(default_factory=dict)
# Validation
warnings_acknowledged: List[str] = field(default_factory=list)
baseline_validated: bool = False
baseline_results: Optional[Dict[str, Any]] = None
def to_dict(self) -> Dict[str, Any]:
"""Convert to dictionary."""
return {
"study_name": self.study_name,
"study_description": self.study_description,
"interview_session_id": self.interview_session_id,
"model_path": self.model_path,
"sim_path": self.sim_path,
"fem_path": self.fem_path,
"design_variables": [dv.to_dict() for dv in self.design_variables],
"objectives": [obj.to_dict() for obj in self.objectives],
"constraints": [con.to_dict() for con in self.constraints],
"protocol": self.protocol,
"n_trials": self.n_trials,
"sampler": self.sampler,
"use_neural_acceleration": self.use_neural_acceleration,
"solver_config": self.solver_config,
"solve_all_solutions": self.solve_all_solutions,
"extractors_config": self.extractors_config,
"warnings_acknowledged": self.warnings_acknowledged,
"baseline_validated": self.baseline_validated,
"baseline_results": self.baseline_results,
}
def to_json(self) -> str:
"""Serialize to JSON string."""
return json.dumps(self.to_dict(), indent=2)
@classmethod
def from_dict(cls, data: Dict[str, Any]) -> "StudyBlueprint":
"""Create from dictionary."""
design_variables = [
DesignVariable(**dv) for dv in data.get("design_variables", [])
]
objectives = [
Objective(**obj) for obj in data.get("objectives", [])
]
constraints = [
Constraint(**con) for con in data.get("constraints", [])
]
return cls(
study_name=data.get("study_name", ""),
study_description=data.get("study_description", ""),
interview_session_id=data.get("interview_session_id", ""),
model_path=data.get("model_path", ""),
sim_path=data.get("sim_path", ""),
fem_path=data.get("fem_path", ""),
design_variables=design_variables,
objectives=objectives,
constraints=constraints,
protocol=data.get("protocol", "protocol_10_single"),
n_trials=data.get("n_trials", 100),
sampler=data.get("sampler", "TPE"),
use_neural_acceleration=data.get("use_neural_acceleration", False),
solver_config=data.get("solver_config", {}),
solve_all_solutions=data.get("solve_all_solutions", True),
extractors_config=data.get("extractors_config", {}),
warnings_acknowledged=data.get("warnings_acknowledged", []),
baseline_validated=data.get("baseline_validated", False),
baseline_results=data.get("baseline_results"),
)
def to_config_json(self) -> Dict[str, Any]:
"""
Convert to optimization_config.json format.
This is the format expected by the optimization runner.
"""
config = {
"study_name": self.study_name,
"description": self.study_description,
"version": "2.0",
"model": {
"part_file": self.model_path,
"sim_file": self.sim_path,
"fem_file": self.fem_path,
},
"design_variables": [
dv.to_config_format() for dv in self.design_variables
],
"objectives": [
obj.to_config_format() for obj in self.objectives
],
"constraints": [
con.to_config_format() for con in self.constraints
],
"optimization": {
"n_trials": self.n_trials,
"sampler": self.sampler,
"protocol": self.protocol,
"neural_acceleration": self.use_neural_acceleration,
},
"solver": {
"solve_all": self.solve_all_solutions,
**self.solver_config,
},
"extractors": self.extractors_config,
"_metadata": {
"interview_session_id": self.interview_session_id,
"warnings_acknowledged": self.warnings_acknowledged,
"baseline_validated": self.baseline_validated,
}
}
return config
def to_markdown(self) -> str:
"""Generate human-readable markdown summary."""
lines = []
lines.append(f"# Study Blueprint: {self.study_name}")
lines.append("")
if self.study_description:
lines.append(f"**Description**: {self.study_description}")
lines.append("")
# Design Variables
lines.append(f"## Design Variables ({len(self.design_variables)})")
lines.append("")
lines.append("| Parameter | Current | Min | Max | Units |")
lines.append("|-----------|---------|-----|-----|-------|")
for dv in self.design_variables:
lines.append(f"| {dv.parameter} | {dv.current_value} | {dv.min_value} | {dv.max_value} | {dv.units or '-'} |")
lines.append("")
# Objectives
lines.append(f"## Objectives ({len(self.objectives)})")
lines.append("")
lines.append("| Name | Goal | Extractor | Weight |")
lines.append("|------|------|-----------|--------|")
for obj in self.objectives:
lines.append(f"| {obj.name} | {obj.goal} | {obj.extractor} | {obj.weight} |")
lines.append("")
# Constraints
if self.constraints:
lines.append(f"## Constraints ({len(self.constraints)})")
lines.append("")
lines.append("| Name | Type | Threshold | Extractor | Hard? |")
lines.append("|------|------|-----------|-----------|-------|")
for con in self.constraints:
op = "<=" if con.constraint_type == "max" else ">="
lines.append(f"| {con.name} | {op} | {con.threshold} | {con.extractor} | {'Yes' if con.is_hard else 'No'} |")
lines.append("")
# Settings
lines.append("## Optimization Settings")
lines.append("")
lines.append(f"- **Protocol**: {self.protocol}")
lines.append(f"- **Trials**: {self.n_trials}")
lines.append(f"- **Sampler**: {self.sampler}")
lines.append(f"- **Neural Acceleration**: {'Enabled' if self.use_neural_acceleration else 'Disabled'}")
lines.append("")
# Validation
lines.append("## Validation")
lines.append("")
lines.append(f"- **Baseline Validated**: {'Yes' if self.baseline_validated else 'No'}")
if self.warnings_acknowledged:
lines.append(f"- **Warnings Acknowledged**: {len(self.warnings_acknowledged)}")
for w in self.warnings_acknowledged:
lines.append(f" - {w}")
lines.append("")
return "\n".join(lines)
def validate(self) -> List[str]:
"""
Validate blueprint completeness.
Returns:
List of validation errors (empty if valid)
"""
errors = []
if not self.study_name:
errors.append("Study name is required")
if not self.design_variables:
errors.append("At least one design variable is required")
if not self.objectives:
errors.append("At least one objective is required")
for dv in self.design_variables:
if dv.min_value >= dv.max_value:
errors.append(f"Invalid bounds for {dv.parameter}: min >= max")
return errors
def is_multi_objective(self) -> bool:
"""Check if this is a multi-objective study."""
return len(self.objectives) > 1
def get_objective_count(self) -> int:
"""Get number of objectives."""
return len(self.objectives)
def get_constraint_count(self) -> int:
"""Get number of constraints."""
return len(self.constraints)
def get_design_variable_count(self) -> int:
"""Get number of design variables."""
return len(self.design_variables)
class BlueprintBuilder:
"""
Helper class for building StudyBlueprint from interview state.
"""
def __init__(self):
"""Initialize builder."""
from .interview_intelligence import InterviewIntelligence
self.intelligence = InterviewIntelligence()
def from_interview_state(
self,
state: "InterviewState",
introspection: Optional[Dict[str, Any]] = None
) -> StudyBlueprint:
"""
Build StudyBlueprint from completed interview state.
Args:
state: Completed interview state
introspection: Optional introspection results
Returns:
StudyBlueprint ready for generation
"""
answers = state.answers
intro = introspection or state.introspection
# Build design variables
design_variables = []
for dv_data in answers.get("design_variables", []):
if isinstance(dv_data, dict):
dv = DesignVariable(
parameter=dv_data.get("parameter", ""),
current_value=dv_data.get("current_value", 0),
min_value=dv_data.get("min_value", 0),
max_value=dv_data.get("max_value", 1),
units=dv_data.get("units"),
is_integer=dv_data.get("is_integer", False),
)
design_variables.append(dv)
elif isinstance(dv_data, str):
# Just a parameter name - look up in introspection
expr = self._find_expression(dv_data, intro.get("expressions", []))
if expr:
value = expr.get("value", 0)
dv = DesignVariable(
parameter=dv_data,
current_value=value,
min_value=value * 0.5 if value > 0 else value * 1.5,
max_value=value * 1.5 if value > 0 else value * 0.5,
)
design_variables.append(dv)
# Build objectives
objectives = []
primary_goal = answers.get("objectives", [{}])
if isinstance(primary_goal, list) and primary_goal:
primary = primary_goal[0] if isinstance(primary_goal[0], dict) else {"goal": primary_goal[0]}
else:
primary = {"goal": str(primary_goal)}
# Map to extractor
extractor_sel = self.intelligence.extractor_mapper.map_goal_to_extractor(
primary.get("goal", ""),
intro
)
objectives.append(Objective(
name=primary.get("name", "primary_objective"),
goal=self._normalize_goal(primary.get("goal", "")),
extractor=extractor_sel.extractor_id,
extractor_name=extractor_sel.extractor_name,
extractor_params=extractor_sel.params,
weight=primary.get("weight", 1.0),
))
# Add secondary objectives
secondary = answers.get("objectives_secondary", [])
for sec_goal in secondary:
if sec_goal == "none" or not sec_goal:
continue
sec_sel = self.intelligence.extractor_mapper.map_goal_to_extractor(
sec_goal, intro
)
objectives.append(Objective(
name=f"secondary_{sec_goal}",
goal=self._normalize_goal(sec_goal),
extractor=sec_sel.extractor_id,
extractor_name=sec_sel.extractor_name,
extractor_params=sec_sel.params,
weight=0.5, # Default lower weight for secondary
))
# Build constraints
constraints = []
constraint_answers = answers.get("constraints", {})
constraint_handling = answers.get("constraint_handling", "hard")
if "max_stress" in constraint_answers and constraint_answers["max_stress"]:
stress_sel = self.intelligence.extractor_mapper.map_constraint_to_extractor("stress", intro)
constraints.append(Constraint(
name="max_stress",
constraint_type="max",
threshold=constraint_answers["max_stress"],
extractor=stress_sel.extractor_id,
extractor_name=stress_sel.extractor_name,
extractor_params=stress_sel.params,
is_hard=constraint_handling != "soft",
))
if "max_displacement" in constraint_answers and constraint_answers["max_displacement"]:
disp_sel = self.intelligence.extractor_mapper.map_constraint_to_extractor("displacement", intro)
constraints.append(Constraint(
name="max_displacement",
constraint_type="max",
threshold=constraint_answers["max_displacement"],
extractor=disp_sel.extractor_id,
extractor_name=disp_sel.extractor_name,
extractor_params=disp_sel.params,
is_hard=constraint_handling != "soft",
))
if "min_frequency" in constraint_answers and constraint_answers["min_frequency"]:
freq_sel = self.intelligence.extractor_mapper.map_constraint_to_extractor("frequency", intro)
constraints.append(Constraint(
name="min_frequency",
constraint_type="min",
threshold=constraint_answers["min_frequency"],
extractor=freq_sel.extractor_id,
extractor_name=freq_sel.extractor_name,
extractor_params=freq_sel.params,
is_hard=constraint_handling != "soft",
))
if "max_mass" in constraint_answers and constraint_answers["max_mass"]:
mass_sel = self.intelligence.extractor_mapper.map_constraint_to_extractor("mass", intro)
constraints.append(Constraint(
name="max_mass",
constraint_type="max",
threshold=constraint_answers["max_mass"],
extractor=mass_sel.extractor_id,
extractor_name=mass_sel.extractor_name,
is_hard=constraint_handling != "soft",
))
# Determine protocol
protocol = "protocol_11_multi" if len(objectives) > 1 else "protocol_10_single"
# Get settings
n_trials = answers.get("n_trials", 100)
if n_trials == "custom":
n_trials = 100 # Default
# Build blueprint
blueprint = StudyBlueprint(
study_name=state.study_name,
study_description=answers.get("problem_description", ""),
interview_session_id=state.session_id,
model_path=intro.get("part_file", ""),
sim_path=intro.get("sim_file", ""),
fem_path=intro.get("fem_file", ""),
design_variables=design_variables,
objectives=objectives,
constraints=constraints,
protocol=protocol,
n_trials=int(n_trials) if isinstance(n_trials, (int, float)) else 100,
sampler=self.intelligence.suggest_sampler(len(objectives), len(design_variables)),
use_neural_acceleration=answers.get("use_neural_acceleration", False),
solve_all_solutions=answers.get("solve_all_solutions", True),
warnings_acknowledged=state.warnings_acknowledged,
baseline_validated=answers.get("run_baseline_validation", False),
)
return blueprint
def _find_expression(self, name: str, expressions: List[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
"""Find expression by name."""
for expr in expressions:
if expr.get("name") == name:
return expr
return None
def _normalize_goal(self, goal: str) -> str:
"""Normalize goal string to standard format."""
goal_lower = goal.lower()
if "minimize" in goal_lower or "reduce" in goal_lower:
return "minimize"
elif "maximize" in goal_lower or "increase" in goal_lower:
return "maximize"
elif "target" in goal_lower:
return "target"
else:
return goal
# Import for type hints
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from .interview_state import InterviewState