Files
Atomizer/optimization_engine/config/spec_models.py
Anto01 a26914bbe8 feat: Add Studio UI, intake system, and extractor improvements
Dashboard:
- Add Studio page with drag-drop model upload and Claude chat
- Add intake system for study creation workflow
- Improve session manager and context builder
- Add intake API routes and frontend components

Optimization Engine:
- Add CLI module for command-line operations
- Add intake module for study preprocessing
- Add validation module with gate checks
- Improve Zernike extractor documentation
- Update spec models with better validation
- Enhance solve_simulation robustness

Documentation:
- Add ATOMIZER_STUDIO.md planning doc
- Add ATOMIZER_UX_SYSTEM.md for UX patterns
- Update extractor library docs
- Add study-readme-generator skill

Tools:
- Add test scripts for extraction validation
- Add Zernike recentering test

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
2026-01-27 12:02:30 -05:00

811 lines
26 KiB
Python

"""
AtomizerSpec v2.0 Pydantic Models
These models match the JSON Schema at optimization_engine/schemas/atomizer_spec_v2.json
They provide validation and type safety for the unified configuration system.
"""
from datetime import datetime
from enum import Enum
from typing import Any, Dict, List, Literal, Optional, Tuple, Union
from pydantic import BaseModel, Field, field_validator, model_validator
import re
# ============================================================================
# Enums
# ============================================================================
class SpecCreatedBy(str, Enum):
"""Who/what created the spec."""
CANVAS = "canvas"
CLAUDE = "claude"
API = "api"
MIGRATION = "migration"
MANUAL = "manual"
DASHBOARD_INTAKE = "dashboard_intake"
class SpecStatus(str, Enum):
"""Study lifecycle status."""
DRAFT = "draft"
INTROSPECTED = "introspected"
CONFIGURED = "configured"
VALIDATED = "validated"
READY = "ready"
RUNNING = "running"
COMPLETED = "completed"
FAILED = "failed"
class SolverType(str, Enum):
"""Supported solver types."""
NASTRAN = "nastran"
NX_NASTRAN = "NX_Nastran"
ABAQUS = "abaqus"
class SubcaseType(str, Enum):
"""Subcase analysis types."""
STATIC = "static"
MODAL = "modal"
THERMAL = "thermal"
BUCKLING = "buckling"
class DesignVariableType(str, Enum):
"""Design variable types."""
CONTINUOUS = "continuous"
INTEGER = "integer"
CATEGORICAL = "categorical"
class ExtractorType(str, Enum):
"""Physics extractor types."""
DISPLACEMENT = "displacement"
FREQUENCY = "frequency"
STRESS = "stress"
MASS = "mass"
MASS_EXPRESSION = "mass_expression"
ZERNIKE_OPD = "zernike_opd"
ZERNIKE_CSV = "zernike_csv"
TEMPERATURE = "temperature"
CUSTOM_FUNCTION = "custom_function"
class OptimizationDirection(str, Enum):
"""Optimization direction."""
MINIMIZE = "minimize"
MAXIMIZE = "maximize"
class ConstraintType(str, Enum):
"""Constraint types."""
HARD = "hard"
SOFT = "soft"
class ConstraintOperator(str, Enum):
"""Constraint comparison operators."""
LE = "<="
GE = ">="
LT = "<"
GT = ">"
EQ = "=="
class PenaltyMethod(str, Enum):
"""Penalty methods for constraints."""
LINEAR = "linear"
QUADRATIC = "quadratic"
EXPONENTIAL = "exponential"
class AlgorithmType(str, Enum):
"""Optimization algorithm types."""
TPE = "TPE"
CMA_ES = "CMA-ES"
NSGA_II = "NSGA-II"
RANDOM_SEARCH = "RandomSearch"
SAT_V3 = "SAT_v3"
GP_BO = "GP-BO"
class SurrogateType(str, Enum):
"""Surrogate model types."""
MLP = "MLP"
GNN = "GNN"
ENSEMBLE = "ensemble"
# ============================================================================
# Position Model
# ============================================================================
class CanvasPosition(BaseModel):
"""Canvas position for nodes."""
x: float = 0
y: float = 0
# ============================================================================
# Introspection Models (for intake workflow)
# ============================================================================
class ExpressionInfo(BaseModel):
"""Information about an NX expression from introspection."""
name: str = Field(..., description="Expression name in NX")
value: Optional[float] = Field(default=None, description="Current value")
units: Optional[str] = Field(default=None, description="Physical units")
formula: Optional[str] = Field(default=None, description="Expression formula if any")
is_candidate: bool = Field(
default=False, description="Whether this is a design variable candidate"
)
confidence: float = Field(
default=0.0, ge=0.0, le=1.0, description="Confidence that this is a DV"
)
class BaselineData(BaseModel):
"""Results from baseline FEA solve."""
timestamp: datetime = Field(..., description="When baseline was run")
solve_time_seconds: float = Field(..., description="How long the solve took")
mass_kg: Optional[float] = Field(default=None, description="Computed mass from BDF/FEM")
max_displacement_mm: Optional[float] = Field(
default=None, description="Max displacement result"
)
max_stress_mpa: Optional[float] = Field(default=None, description="Max von Mises stress")
success: bool = Field(default=True, description="Whether baseline solve succeeded")
error: Optional[str] = Field(default=None, description="Error message if failed")
class IntrospectionData(BaseModel):
"""Model introspection results stored in the spec."""
timestamp: datetime = Field(..., description="When introspection was run")
solver_type: Optional[str] = Field(default=None, description="Detected solver type")
mass_kg: Optional[float] = Field(
default=None, description="Mass from expressions or properties"
)
volume_mm3: Optional[float] = Field(default=None, description="Volume from mass properties")
expressions: List[ExpressionInfo] = Field(
default_factory=list, description="Discovered expressions"
)
baseline: Optional[BaselineData] = Field(default=None, description="Baseline solve results")
warnings: List[str] = Field(default_factory=list, description="Warnings from introspection")
def get_design_candidates(self) -> List[ExpressionInfo]:
"""Return expressions marked as design variable candidates."""
return [e for e in self.expressions if e.is_candidate]
# ============================================================================
# Meta Models
# ============================================================================
class SpecMeta(BaseModel):
"""Metadata about the spec."""
version: str = Field(..., pattern=r"^2\.\d+$", description="Schema version (e.g., '2.0')")
created: Optional[datetime] = Field(default=None, description="When the spec was created")
modified: Optional[datetime] = Field(
default=None, description="When the spec was last modified"
)
created_by: Optional[SpecCreatedBy] = Field(
default=None, description="Who/what created the spec"
)
modified_by: Optional[str] = Field(default=None, description="Who/what last modified the spec")
study_name: str = Field(
...,
min_length=3,
max_length=100,
pattern=r"^[a-z0-9_]+$",
description="Unique study identifier (snake_case)",
)
description: Optional[str] = Field(
default=None, max_length=1000, description="Human-readable description"
)
tags: Optional[List[str]] = Field(default=None, description="Tags for categorization")
engineering_context: Optional[str] = Field(
default=None, description="Real-world engineering context"
)
status: SpecStatus = Field(default=SpecStatus.DRAFT, description="Study lifecycle status")
topic: Optional[str] = Field(
default=None,
pattern=r"^[A-Za-z0-9_]+$",
description="Topic folder for grouping related studies",
)
# ============================================================================
# Model Configuration Models
# ============================================================================
class NxPartConfig(BaseModel):
"""NX geometry part file configuration."""
path: Optional[str] = Field(default=None, description="Path to .prt file")
hash: Optional[str] = Field(default=None, description="File hash for change detection")
idealized_part: Optional[str] = Field(
default=None, description="Idealized part filename (_i.prt)"
)
class FemConfig(BaseModel):
"""FEM mesh file configuration."""
path: Optional[str] = Field(default=None, description="Path to .fem file")
element_count: Optional[int] = Field(default=None, description="Number of elements")
node_count: Optional[int] = Field(default=None, description="Number of nodes")
class Subcase(BaseModel):
"""Simulation subcase definition."""
id: int
name: Optional[str] = None
type: Optional[SubcaseType] = None
class SimConfig(BaseModel):
"""Simulation file configuration."""
path: str = Field(..., description="Path to .sim file")
solver: SolverType = Field(..., description="Solver type")
solution_type: Optional[str] = Field(
default=None, pattern=r"^SOL\d+$", description="Solution type (e.g., SOL101)"
)
subcases: Optional[List[Subcase]] = Field(default=None, description="Defined subcases")
class NxSettings(BaseModel):
"""NX runtime settings."""
nx_install_path: Optional[str] = None
simulation_timeout_s: Optional[int] = Field(default=None, ge=60, le=7200)
auto_start_nx: Optional[bool] = None
class ModelConfig(BaseModel):
"""NX model files and configuration."""
nx_part: Optional[NxPartConfig] = None
fem: Optional[FemConfig] = None
sim: Optional[SimConfig] = Field(
default=None, description="Simulation file config (required for optimization)"
)
nx_settings: Optional[NxSettings] = None
introspection: Optional[IntrospectionData] = Field(
default=None, description="Model introspection results from intake"
)
# ============================================================================
# Design Variable Models
# ============================================================================
class DesignVariableBounds(BaseModel):
"""Design variable bounds."""
min: float
max: float
@model_validator(mode="after")
def validate_bounds(self) -> "DesignVariableBounds":
if self.min >= self.max:
raise ValueError(f"min ({self.min}) must be less than max ({self.max})")
return self
class DesignVariable(BaseModel):
"""A design variable to optimize."""
id: str = Field(..., pattern=r"^dv_\d{3}$", description="Unique identifier (pattern: dv_XXX)")
name: str = Field(..., description="Human-readable name")
expression_name: str = Field(
...,
pattern=r"^[a-zA-Z_][a-zA-Z0-9_]*$",
description="NX expression name (must match model)",
)
type: DesignVariableType = Field(..., description="Variable type")
bounds: DesignVariableBounds = Field(..., description="Value bounds")
baseline: Optional[float] = Field(default=None, description="Current/initial value")
units: Optional[str] = Field(default=None, description="Physical units (mm, deg, etc.)")
step: Optional[float] = Field(default=None, description="Step size for integer/discrete")
enabled: bool = Field(default=True, description="Whether to include in optimization")
description: Optional[str] = None
canvas_position: Optional[CanvasPosition] = None
# ============================================================================
# Extractor Models
# ============================================================================
class ExtractorConfig(BaseModel):
"""Type-specific extractor configuration."""
inner_radius_mm: Optional[float] = None
outer_radius_mm: Optional[float] = None
n_modes: Optional[int] = None
filter_low_orders: Optional[int] = None
displacement_unit: Optional[str] = None
reference_subcase: Optional[int] = None
expression_name: Optional[str] = None
mode_number: Optional[int] = None
element_type: Optional[str] = None
result_type: Optional[str] = None
metric: Optional[str] = None
class Config:
extra = "allow" # Allow additional fields for flexibility
class CustomFunction(BaseModel):
"""Custom function definition for custom_function extractors."""
name: Optional[str] = Field(default=None, description="Function name")
module: Optional[str] = Field(default=None, description="Python module path")
signature: Optional[str] = Field(default=None, description="Function signature")
source_code: Optional[str] = Field(default=None, description="Python source code")
class ExtractorOutput(BaseModel):
"""Output definition for an extractor."""
name: str = Field(..., description="Output name (used by objectives/constraints)")
metric: Optional[str] = Field(
default=None, description="Specific metric (max, total, rms, etc.)"
)
subcase: Optional[int] = Field(default=None, description="Subcase ID for this output")
units: Optional[str] = None
class Extractor(BaseModel):
"""Physics extractor that computes outputs from FEA."""
id: str = Field(..., pattern=r"^ext_\d{3}$", description="Unique identifier (pattern: ext_XXX)")
name: str = Field(..., description="Human-readable name")
type: ExtractorType = Field(..., description="Extractor type")
builtin: bool = Field(default=True, description="Whether this is a built-in extractor")
config: Optional[ExtractorConfig] = Field(
default=None, description="Type-specific configuration"
)
function: Optional[CustomFunction] = Field(
default=None, description="Custom function definition (for custom_function type)"
)
outputs: List[ExtractorOutput] = Field(..., min_length=1, description="Output values")
canvas_position: Optional[CanvasPosition] = None
@model_validator(mode="after")
def validate_custom_function(self) -> "Extractor":
if self.type == ExtractorType.CUSTOM_FUNCTION and self.function is None:
raise ValueError("custom_function extractor requires function definition")
return self
# ============================================================================
# Objective Models
# ============================================================================
class ObjectiveSource(BaseModel):
"""Source reference for objective value."""
extractor_id: str = Field(..., description="Reference to extractor")
output_name: str = Field(..., description="Which output from the extractor")
class Objective(BaseModel):
"""Optimization objective."""
id: str = Field(..., pattern=r"^obj_\d{3}$", description="Unique identifier (pattern: obj_XXX)")
name: str = Field(..., description="Human-readable name")
direction: OptimizationDirection = Field(..., description="Optimization direction")
weight: float = Field(default=1.0, ge=0, description="Weight for weighted sum")
source: ObjectiveSource = Field(..., description="Where the value comes from")
target: Optional[float] = Field(default=None, description="Target value (for goal programming)")
units: Optional[str] = None
description: Optional[str] = None
canvas_position: Optional[CanvasPosition] = None
# ============================================================================
# Constraint Models
# ============================================================================
class ConstraintSource(BaseModel):
"""Source reference for constraint value."""
extractor_id: str
output_name: str
class PenaltyConfig(BaseModel):
"""Penalty method configuration for constraints."""
method: Optional[PenaltyMethod] = None
weight: Optional[float] = None
margin: Optional[float] = Field(default=None, description="Soft margin before penalty kicks in")
class Constraint(BaseModel):
"""Hard or soft constraint."""
id: str = Field(..., pattern=r"^con_\d{3}$", description="Unique identifier (pattern: con_XXX)")
name: str
type: ConstraintType = Field(..., description="Constraint type")
operator: ConstraintOperator = Field(..., description="Comparison operator")
threshold: float = Field(..., description="Constraint threshold value")
source: ConstraintSource = Field(..., description="Where the value comes from")
penalty_config: Optional[PenaltyConfig] = None
description: Optional[str] = None
canvas_position: Optional[CanvasPosition] = None
# ============================================================================
# Optimization Models
# ============================================================================
class AlgorithmConfig(BaseModel):
"""Algorithm-specific settings."""
population_size: Optional[int] = None
n_generations: Optional[int] = None
mutation_prob: Optional[float] = None
crossover_prob: Optional[float] = None
seed: Optional[int] = None
n_startup_trials: Optional[int] = None
sigma0: Optional[float] = None
class Config:
extra = "allow" # Allow additional algorithm-specific fields
class Algorithm(BaseModel):
"""Optimization algorithm configuration."""
type: AlgorithmType
config: Optional[AlgorithmConfig] = None
class OptimizationBudget(BaseModel):
"""Computational budget for optimization."""
max_trials: Optional[int] = Field(default=None, ge=1, le=10000)
max_time_hours: Optional[float] = None
convergence_patience: Optional[int] = Field(
default=None, description="Stop if no improvement for N trials"
)
class SurrogateConfig(BaseModel):
"""Neural surrogate model configuration."""
n_models: Optional[int] = None
architecture: Optional[List[int]] = None
train_every_n_trials: Optional[int] = None
min_training_samples: Optional[int] = None
acquisition_candidates: Optional[int] = None
fea_validations_per_round: Optional[int] = None
class Surrogate(BaseModel):
"""Surrogate model settings."""
enabled: Optional[bool] = None
type: Optional[SurrogateType] = None
config: Optional[SurrogateConfig] = None
class OptimizationConfig(BaseModel):
"""Optimization algorithm configuration."""
algorithm: Algorithm
budget: OptimizationBudget
surrogate: Optional[Surrogate] = None
canvas_position: Optional[CanvasPosition] = None
# ============================================================================
# Workflow Models
# ============================================================================
class WorkflowStage(BaseModel):
"""A stage in a multi-stage optimization workflow."""
id: str
name: str
algorithm: Optional[str] = None
trials: Optional[int] = None
purpose: Optional[str] = None
class WorkflowTransition(BaseModel):
"""Transition between workflow stages."""
from_: str = Field(..., alias="from")
to: str
condition: Optional[str] = None
class Config:
populate_by_name = True
class Workflow(BaseModel):
"""Multi-stage optimization workflow."""
stages: Optional[List[WorkflowStage]] = None
transitions: Optional[List[WorkflowTransition]] = None
# ============================================================================
# Reporting Models
# ============================================================================
class InsightConfig(BaseModel):
"""Insight-specific configuration."""
include_html: Optional[bool] = None
show_pareto_evolution: Optional[bool] = None
class Config:
extra = "allow"
class Insight(BaseModel):
"""Reporting insight definition."""
type: Optional[str] = None
for_trials: Optional[str] = None
config: Optional[InsightConfig] = None
class ReportingConfig(BaseModel):
"""Reporting configuration."""
auto_report: Optional[bool] = None
report_triggers: Optional[List[str]] = None
insights: Optional[List[Insight]] = None
# ============================================================================
# Canvas Models
# ============================================================================
class CanvasViewport(BaseModel):
"""Canvas viewport settings."""
x: float = 0
y: float = 0
zoom: float = 1.0
class CanvasEdge(BaseModel):
"""Connection between canvas nodes."""
source: str
target: str
sourceHandle: Optional[str] = None
targetHandle: Optional[str] = None
class CanvasGroup(BaseModel):
"""Grouping of canvas nodes."""
id: str
name: str
node_ids: List[str]
class CanvasConfig(BaseModel):
"""Canvas UI state (persisted for reconstruction)."""
layout_version: Optional[str] = None
viewport: Optional[CanvasViewport] = None
edges: Optional[List[CanvasEdge]] = None
groups: Optional[List[CanvasGroup]] = None
# ============================================================================
# Main AtomizerSpec Model
# ============================================================================
class AtomizerSpec(BaseModel):
"""
AtomizerSpec v2.0 - The unified configuration schema for Atomizer optimization studies.
This is the single source of truth used by:
- Canvas UI (rendering and editing)
- Backend API (validation and storage)
- Claude Assistant (reading and modifying)
- Optimization Engine (execution)
"""
meta: SpecMeta = Field(..., description="Metadata about the spec")
model: ModelConfig = Field(..., description="NX model files and configuration")
design_variables: List[DesignVariable] = Field(
default_factory=list,
max_length=50,
description="Design variables to optimize (required for running)",
)
extractors: List[Extractor] = Field(
default_factory=list, description="Physics extractors (required for running)"
)
objectives: List[Objective] = Field(
default_factory=list,
max_length=5,
description="Optimization objectives (required for running)",
)
constraints: Optional[List[Constraint]] = Field(
default=None, description="Hard and soft constraints"
)
optimization: OptimizationConfig = Field(..., description="Algorithm configuration")
workflow: Optional[Workflow] = Field(default=None, description="Multi-stage workflow")
reporting: Optional[ReportingConfig] = Field(default=None, description="Reporting config")
canvas: Optional[CanvasConfig] = Field(default=None, description="Canvas UI state")
@model_validator(mode="after")
def validate_references(self) -> "AtomizerSpec":
"""Validate that all references are valid."""
# Collect valid extractor IDs and their outputs
extractor_outputs: Dict[str, set] = {}
for ext in self.extractors:
extractor_outputs[ext.id] = {o.name for o in ext.outputs}
# Validate objective sources
for obj in self.objectives:
if obj.source.extractor_id not in extractor_outputs:
raise ValueError(
f"Objective '{obj.name}' references unknown extractor: {obj.source.extractor_id}"
)
if obj.source.output_name not in extractor_outputs[obj.source.extractor_id]:
raise ValueError(
f"Objective '{obj.name}' references unknown output: {obj.source.output_name}"
)
# Validate constraint sources
if self.constraints:
for con in self.constraints:
if con.source.extractor_id not in extractor_outputs:
raise ValueError(
f"Constraint '{con.name}' references unknown extractor: {con.source.extractor_id}"
)
if con.source.output_name not in extractor_outputs[con.source.extractor_id]:
raise ValueError(
f"Constraint '{con.name}' references unknown output: {con.source.output_name}"
)
return self
def get_enabled_design_variables(self) -> List[DesignVariable]:
"""Return only enabled design variables."""
return [dv for dv in self.design_variables if dv.enabled]
def get_extractor_by_id(self, extractor_id: str) -> Optional[Extractor]:
"""Find an extractor by ID."""
for ext in self.extractors:
if ext.id == extractor_id:
return ext
return None
def get_objective_by_id(self, objective_id: str) -> Optional[Objective]:
"""Find an objective by ID."""
for obj in self.objectives:
if obj.id == objective_id:
return obj
return None
def get_constraint_by_id(self, constraint_id: str) -> Optional[Constraint]:
"""Find a constraint by ID."""
if not self.constraints:
return None
for con in self.constraints:
if con.id == constraint_id:
return con
return None
def has_custom_extractors(self) -> bool:
"""Check if spec has any custom function extractors."""
return any(ext.type == ExtractorType.CUSTOM_FUNCTION for ext in self.extractors)
def is_multi_objective(self) -> bool:
"""Check if this is a multi-objective optimization."""
return len(self.objectives) > 1
def is_ready_for_optimization(self) -> Tuple[bool, List[str]]:
"""
Check if spec is complete enough to run optimization.
Returns:
Tuple of (is_ready, list of missing requirements)
"""
missing = []
# Check required fields for optimization
if not self.model.sim:
missing.append("No simulation file (.sim) configured")
if not self.design_variables:
missing.append("No design variables defined")
if not self.extractors:
missing.append("No extractors defined")
if not self.objectives:
missing.append("No objectives defined")
# Check that enabled DVs have valid bounds
for dv in self.get_enabled_design_variables():
if dv.bounds.min >= dv.bounds.max:
missing.append(f"Design variable '{dv.name}' has invalid bounds")
return len(missing) == 0, missing
# ============================================================================
# Validation Response Models
# ============================================================================
class ValidationError(BaseModel):
"""A validation error."""
type: str # 'schema', 'semantic', 'reference'
path: List[str]
message: str
class ValidationWarning(BaseModel):
"""A validation warning."""
type: str
path: List[str]
message: str
class ValidationSummary(BaseModel):
"""Summary of spec contents."""
design_variables: int
extractors: int
objectives: int
constraints: int
custom_functions: int
class ValidationReport(BaseModel):
"""Full validation report."""
valid: bool
errors: List[ValidationError]
warnings: List[ValidationWarning]
summary: ValidationSummary