feat: add optimization (#5)

* feat: add optimization

* chore: add optimization docs regarding integration with design flow

* v.01.11

Author: a96lex

axiomatic_mcp/servers/pic/README.md

@@ -1,10 +1,10 @@
 # AxPhotonicsPreview
 
-An MCP server for designing, and simulating Photonic Integrated Circuits (PICs) using the Axiomatic AI platform.
+An MCP server for designing, optimizing, and simulating Photonic Integrated Circuits (PICs) using the Axiomatic AI platform.
 
 ## Overview
 
-The AxPhotonicsPreview Server enables AI assistants to create photonic integrated circuit designs using natural language descriptions. It leverages the Axiomatic AI platform to generate GDSFactory-compatible Python code for photonic components and circuits, and provides tools to simulate them.
+The AxPhotonicsPreview Server enables AI assistants to create photonic integrated circuit designs using natural language descriptions. It leverages the Axiomatic AI platform to generate GDSFactory-compatible Python code for photonic components and circuits, and provides tools to simulate and optimize them.
 
 ## Tools Available
 
@@ -47,6 +47,101 @@ Simulates a previously generated circuit and returns both the wavelength range (
 
 ---
 
+### `optimize_circuit`
+
+Optimizes a photonic circuit by refining the generated code using its formalized statements.
+
+**Parameters:**
+
+- `code_path` (Path, required): Path to the Python file containing the circuit code.
+- `statements_path` (Path, required): Path to the JSON file containing the circuit statements.
+
+**Outputs:**
+
+- **Optimized Python code (`*_optimized.py`)**: Refined version of the input circuit code.
+- **Structured Output**: Returns the optimized file path and the optimized code.
+
+**Example Usage:**
+
+Suppose you have the following circuit in `circuit.py`:
+
+```python
+import gdsfactory as gf
+import cspdk.si220.cband
+
+pdk = cspdk.si220.cband.get_pdk()
+pdk.activate()
+
+c = gf.Component()
+
+# Add MZI component
+mzi = c << pdk.get_component("mzi", delta_length=10, cross_section="strip")
+mzi.move((-34.55, 2.14))
+mzi.name = "mzi"
+
+# Add ports
+c.add_port("in0", port=mzi.ports["o1"])
+c.add_port("in1", port=mzi.ports["o2"])
+c.add_port("out0", port=mzi.ports["o3"])
+c.add_port("out1", port=mzi.ports["o4"])
+
+c
+```
+
+And the following `statements.json`:
+
+```json
+{
+  "statements": [
+    {
+      "type": "PARAMETER_CONSTRAINT",
+      "text": "The circuit has one input port named 'in0' (corresponding to 'o1') and one output port named 'out0' (corresponding to 'o2').",
+      "formalization": {
+        "code": "n_ports == 2",
+        "default_tolerance": 0.05,
+        "mapping": {
+          "n_ports": {
+            "name": "number_of_ports",
+            "arguments": { "port_type": "all", "component": "netlist" }
+          }
+        }
+      },
+      "validation": null
+    },
+    {
+      "type": "PARAMETER_CONSTRAINT",
+      "text": "The circuit has exactly one input port.",
+      "formalization": {
+        "code": "n_ports_in == 1",
+        "default_tolerance": 0.05,
+        "mapping": {
+          "n_ports_in": {
+            "name": "number_of_ports",
+            "arguments": { "port_type": "in", "component": "netlist" }
+          }
+        }
+      },
+      "validation": null
+    }
+  ]
+}
+```
+
+Running `optimize_circuit` with these inputs will produce a new file:
+
+- `circuit_optimized.py` – optimized code satisfying the given constraints.
+
+**Sample Output (structured):**
+
+```json
+{
+  "optimized_file_path": "/path/to/circuit_optimized.py",
+  "optimized_code": "import gdsfactory as gf ..."
+}
+```
+
+---
+
 ### `list_available_pdks`
 
 Gets a list of all available Process Design Kits (PDKs), with its granted status.
@@ -184,6 +279,7 @@ This server requires extra dependencies, please see [Check system requirements](
 1. **Design Generation**: Use the MCP server to create initial designs
 2. **Local Execution**: Run generated code with GDSFactory
 3. **Simulation**: Use `simulate_circuit` to generate wavelength sweeps and notebook results
+4. **Optimization**: Use `optimize_circuit` to refine the circuit code and parameters based on formalized design constraints
 
 ---
 

axiomatic_mcp/servers/pic/server.py

@@ -15,6 +15,7 @@
 from ...shared.utils.prompt_utils import get_feedback_prompt
 from .services.circuit_service import CircuitService
 from .services.notebook_service import NotebookService
+from .services.optimization_service import OptimizationService
 from .services.pdk_service import PdkService
 from .services.simulation_service import SimulationService
 from .services.statements_service import StatementsService
@@ -38,6 +39,7 @@
 notebook_service = NotebookService()
 pdk_service = PdkService()
 statements_service = StatementsService()
+optimization_service = OptimizationService()
 
 
 @mcp.tool(
@@ -259,6 +261,56 @@ async def validate_statements(
     )
 
 
+@mcp.tool(
+    name="optimize_circuit",
+    description="Optimizes a photonic circuit by refining the generated code using its statements.",
+    tags=["optimize", "gfsfactory"],
+)
+async def optimize_circuit(
+    code_path: Annotated[Path, "Path to the Python file containing the circuit code"],
+    statements_path: Annotated[Path, "Path to the JSON file containing the circuit statements"],
+) -> ToolResult:
+    """Optimize a photonic integrated circuit."""
+    if not code_path.exists():
+        raise FileNotFoundError(f"Circuit code not found: {code_path}")
+    if not statements_path.exists():
+        raise FileNotFoundError(f"Statements file not found: {statements_path}")
+
+    code = code_path.read_text(encoding="utf-8")
+
+    with statements_path.open("r", encoding="utf-8") as f:
+        statements_json = json.load(f)
+
+    statements_list = statements_json.get("statements", [])
+
+    request_body = {
+        "code": code,
+        "statements": statements_list,
+    }
+
+    response = await optimization_service.optimize_code(request_body)
+
+    optimized_code = response.get("optimized_code", "")
+
+    optimized_file_path = code_path.parent / f"{code_path.stem}_optimized.py"
+
+    with optimized_file_path.open("w", encoding="utf-8") as f:
+        f.write(optimized_code)
+
+    return ToolResult(
+        content=[
+            TextContent(
+                type="text",
+                text=f"Optimized circuit saved at {optimized_file_path}",
+            )
+        ],
+        structured_content={
+            "optimized_file_path": str(optimized_file_path),
+            "optimized_code": optimized_code,
+        },
+    )
+
+
 @mcp.tool(
     name="list_available_pdks",
     description="Get a list of all available PDKs that the user has access to.",

axiomatic_mcp/servers/pic/services/optimization_service.py

@@ -0,0 +1,23 @@
+from typing import Any
+
+from ....shared import AxiomaticAPIClient
+from ....shared.models.singleton_base import SingletonBase
+from ...constants.api_constants import ApiRoutes
+
+
+class OptimizationService(SingletonBase):
+    async def optimize_code(self, query: dict[str, Any]) -> dict[str, Any]:
+        """
+        Call the GET_OPTIMIZED_CODE API endpoint with optimization request.
+        query: {
+            "code": str,          # Python code (gdsfactory circuit)
+            "statements": list    # Parsed statements from statements.json
+        }
+        """
+
+        response = AxiomaticAPIClient().post(ApiRoutes.GET_OPTIMIZED_CODE, data=query)
+
+        if not response:
+            raise RuntimeError("No response from get_optimized_code API")
+
+        return response

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "axiomatic-mcp"
-version = "0.1.10"
+version = "0.1.11"
 description = "Modular MCP servers for Axiomatic_AI"
 authors = [{name = "Axiomatic Team", email = "[email protected]"}]
 readme = "README.md"