feat: Add parameter covariance estimation tool to model fitter

axiomatic_mcp/servers/axmodelfitter/README.md

@@ -114,6 +114,7 @@ Provides step-by-step guidance for setting up and executing model fitting workfl
 - **`cross_validate_model`** - Perform cross-validation to assess model generalization
 - **`calculate_information_criteria`** - Compute AIC/BIC for model comparison
 - **`compare_models`** - Statistical comparison of multiple models
+- **`compute_parameter_covariance`** - Provides estimates to quantify parameter uncertainty and correlations.
 
 ## Data Requirements
 

axiomatic_mcp/servers/axmodelfitter/server.py

@@ -2120,6 +2120,228 @@ async def compare_models(
         return ToolResult(content=[TextContent(type="text", text=error_text)])
 
 
+@mcp.tool(
+    name="compute_parameter_covariance",
+    description="""Compute parameter covariance matrices for fitted model parameters.
+
+    Provides uncertainty estimates using robust Huber-White sandwich estimator and
+    classical inverse Hessian approach. Use after fit_model to quantify parameter
+    uncertainty and correlations.
+
+    REQUIRED: Fitted parameters, model definition, same data used in fitting, variance estimate.
+    RETURNS: Covariance matrices, standard errors, correlation matrix.
+    """,
+    tags=["statistics", "uncertainty", "covariance", "parameter_estimation"],
+)
+async def compute_parameter_covariance(
+    model_name: Annotated[str, "Model name (e.g., 'ExponentialDecay', 'RingResonator')"],
+    function_source: Annotated[str, "JAX function source code. MUST use jnp operations: jnp.exp, jnp.sin, etc."],
+    function_name: Annotated[str, "Function name that computes the model output"],
+    parameters: Annotated[list, "Fitted parameter values: [{'name': 'a', 'value': {'magnitude': 2.0, 'unit': 'dimensionless'}}]"],
+    bounds: Annotated[
+        list,
+        "ALL parameter/input/output bounds: [{'name': 'a', 'lower': {'magnitude': 0, 'unit': 'dimensionless'}, 'upper': {'magnitude': 10, 'unit': 'dimensionless'}}]",  # noqa E501
+    ],
+    data_file: Annotated[str, "Path to data file (CSV, Excel, JSON, Parquet). All data must be provided via file."],
+    input_data: Annotated[
+        list, "Input column mappings: [{'column': 'time', 'name': 't', 'unit': 'second'}, {'column': 'x_col', 'name': 'x', 'unit': 'meter'}]"
+    ],
+    output_data: Annotated[
+        dict, "Output column mapping: {'columns': ['signal'], 'name': 'y', 'unit': 'volt'} OR {'columns': ['y1', 'y2'], 'name': 'y', 'unit': 'volt'}"
+    ],
+    file_format: Annotated[str | None, "File format: 'csv', 'excel', 'json', 'parquet' (auto-detect if None)"] = None,
+    variance: Annotated[float, "Noise variance (σ²) for uncertainty quantification. Estimate from residuals or domain knowledge."] = 1.0,
+    constants: Annotated[list | None, "Fixed constants: [{'name': 'c', 'value': {'magnitude': 3.0, 'unit': 'meter'}}]"] = None,
+    docstring: Annotated[str, "Brief description of the model"] = "",
+    cost_function_type: Annotated[str, "Cost function: 'mse' (default), 'mae'"] = "mse",
+    jit_compile: Annotated[bool, "Enable JIT compilation for performance"] = True,
+    scale_params: Annotated[bool, "Enable parameter scaling for numerical stability"] = False,
+) -> ToolResult:
+    """Compute parameter covariance matrix for fitted model parameters."""
+
+    try:
+        if data_file is None:
+            raise ValueError("data_file is required. All data must be provided via file.")
+        if input_data is None:
+            raise ValueError("input_data is required when using file-based input.")
+        if output_data is None:
+            raise ValueError("output_data is required when using file-based input.")
+
+        resolved_input_data, resolved_output_data = resolve_data_input(
+            data_file=data_file, input_data=input_data, output_data=output_data, file_format=file_format
+        )
+
+        input_names, const_names, param_names, bounds_names, n = validate_optimization_inputs(
+            resolved_input_data, resolved_output_data, parameters, bounds, constants
+        )
+
+        prepare_bounds_for_optimization(bounds, input_names, const_names, resolved_output_data["name"])
+
+        if variance <= 0:
+            raise ValueError("variance must be positive (σ² > 0). " "Estimate it from residuals (e.g., final_loss for MSE) or domain knowledge.")
+
+    except ValueError as e:
+        return ToolResult(content=[TextContent(type="text", text=str(e))])
+
+    if constants is None:
+        constants = []
+    request_data = {
+        "model_name": model_name,
+        "parameters": parameters,
+        "bounds": bounds,
+        "constants": constants,
+        "input": resolved_input_data,
+        "target": resolved_output_data,
+        "function_source": function_source,
+        "function_name": function_name,
+        "docstring": docstring,
+        "jit_compile": jit_compile,
+        "cost_function_type": cost_function_type,
+        "scale_params": scale_params,
+        "variance": variance,
+    }
+
+    try:
+        with AxiomaticAPIClient() as client:
+            response = client.post("/digital-twin/compute-parameter-covariance", data=request_data)
+
+        # New success criterion: treat as success if either covariance matrix is present
+        robust_cov = response.get("sandwich_covariance")
+        classical_cov = response.get("inverse_hessian_covariance")
+        has_cov = (isinstance(robust_cov, list) and len(robust_cov) > 0) or (isinstance(classical_cov, list) and len(classical_cov) > 0)
+
+        if not has_cov:
+            error_msg = response.get("error", "Unknown error occurred")
+            # Include full response for debugging if error message is generic
+            debug_info = ""
+            if error_msg == "Unknown error occurred":
+                import json
+
+                debug_info = f"\n\nFull API response:\n{json.dumps(response, indent=2)}"
+            return ToolResult(content=[TextContent(type="text", text=f"Parameter covariance computation failed: {error_msg}{debug_info}")])
+
+        result_text = f"""# Parameter Covariance Analysis: {model_name}
+
+Status: {"Success" if has_cov else "Failed"}
+
+## Fitted Parameters
+"""
+
+        for param in parameters:
+            name = param["name"]
+            value = param["value"]["magnitude"]
+            unit = param["value"]["unit"]
+            result_text += f"- **{name}:** {value:.6g} {unit}\n"
+
+        # Use provided param_names if present, else fall back to parameters list
+        parameter_names = response.get("param_names", param_names)
+
+        # Create lookup from parameter name to value/unit for robust matching
+        # (handles case where API returns param_names in different order than input parameters)
+        param_lookup = {p["name"]: p["value"] for p in parameters}
+
+        # Track std_errors for structured output
+        sandwich_std_errors = None
+        hessian_std_errors = None
+
+        if isinstance(robust_cov, list) and len(robust_cov) > 0:
+            result_text += "\n## Robust Covariance (Huber-White Sandwich)\n\n"
+            cov_array = np.array(robust_cov)
+            std_errors = np.sqrt(np.diag(cov_array))
+            sandwich_std_errors = std_errors.tolist()
+            # Use covariance matrix dimensions to avoid out-of-bounds access
+            n_cov_params = cov_array.shape[0]
+            cov_param_names = parameter_names[:n_cov_params]
+
+            result_text += "### Standard Errors\n"
+            for name, std_err in zip(cov_param_names, std_errors):
+                param_info = param_lookup.get(name, {"magnitude": float("nan"), "unit": "?"})
+                param_value = param_info["magnitude"]
+                param_unit = param_info["unit"]
+                relative_error = (std_err / abs(param_value) * 100) if param_value != 0 else float("inf")
+                result_text += f"- **{name}:** {std_err:.6g} {param_unit} ({relative_error:.2f}% relative)\n"
+
+            result_text += "\n### Correlation Matrix\n"
+            # Use len(cov_param_names) consistently to handle case where parameter_names is shorter than covariance matrix
+            n_display_params = len(cov_param_names)
+            result_text += "| Parameter | " + " | ".join(cov_param_names) + " |\n"
+            result_text += "|-----------|" + "|".join(["-------"] * n_display_params) + "|\n"
+
+            for i, name_i in enumerate(cov_param_names):
+                row_text = f"| **{name_i}** |"
+                for j in range(n_display_params):
+                    if std_errors[i] > 0 and std_errors[j] > 0:
+                        corr = cov_array[i, j] / (std_errors[i] * std_errors[j])
+                        row_text += f" {corr:+.3f} |"
+                    else:
+                        row_text += " N/A |"
+                result_text += row_text + "\n"
+
+        if isinstance(classical_cov, list) and len(classical_cov) > 0:
+            result_text += "\n## Classical Covariance (Inverse Hessian)\n\n"
+            cov_array_classical = np.array(classical_cov)
+            std_errors_classical = np.sqrt(np.diag(cov_array_classical))
+            hessian_std_errors = std_errors_classical.tolist()
+            # Use covariance matrix dimensions to avoid out-of-bounds access
+            n_cov_params_classical = cov_array_classical.shape[0]
+            cov_param_names_classical = parameter_names[:n_cov_params_classical]
+
+            result_text += "### Standard Errors\n"
+            for name, std_err in zip(cov_param_names_classical, std_errors_classical):
+                param_info = param_lookup.get(name, {"magnitude": float("nan"), "unit": "?"})
+                param_value = param_info["magnitude"]
+                param_unit = param_info["unit"]
+                relative_error = (std_err / abs(param_value) * 100) if param_value != 0 else float("inf")
+                result_text += f"- **{name}:** {std_err:.6g} {param_unit} ({relative_error:.2f}% relative)\n"
+
+        result_text += "\n## Notes\n"
+        result_text += "- Standard errors indicate parameter constraint quality\n"
+        result_text += "- 95% confidence interval: parameter ± 1.96 × std_error\n"
+        result_text += "- Correlation near ±1 shows parameters trade off\n"
+        result_text += "- Use robust when model may be misspecified\n"
+
+        return ToolResult(
+            content=[TextContent(type="text", text=result_text)],
+            structured_content={
+                "parameters": parameters,
+                "sandwich_covariance": robust_cov,
+                "inverse_hessian_covariance": classical_cov,
+                "parameter_names": parameter_names,
+                "sandwich_std_errors": sandwich_std_errors,
+                "inverse_hessian_std_errors": hessian_std_errors,
+            },
+        )
+
+    except Exception as e:
+        # Enhanced error handling for HTTP errors
+        import httpx
+
+        error_msg = str(e)
+        if isinstance(e, httpx.HTTPStatusError):
+            # Extract detailed error from HTTP response
+            try:
+                error_body = e.response.json()
+                if "detail" in error_body:
+                    error_msg = error_body["detail"]
+                else:
+                    error_msg = str(error_body)
+            except Exception:
+                error_msg = e.response.text if hasattr(e.response, "text") else str(e)
+
+        error_details = f"""Parameter covariance computation failed: {error_msg}
+
+Troubleshooting:
+- Ensure parameters match those from fit_model call
+- Use same data file and mappings as in fitting
+- Provide realistic noise variance from residuals
+- Verify all required fields are provided
+- Check that all units are valid pint units (e.g., "1/volt" not "dimensionless" for inverse volts)
+
+Variance can be estimated from final_loss: variance ≈ final_loss (for MSE)
+"""
+        return ToolResult(content=[TextContent(type="text", text=error_details)])
+
+
 def main():
     """Main entry point for the model fitting MCP server."""
     mcp.run()

examples/axmodelfitter/ring_modulator/prompt.md

@@ -1,11 +1,11 @@
 
 # Demo: Ring modulator
-Consider the following model for the power transmission a ring modulator as function of wavelength $\lambda$ and bias voltage $V$: 
+Consider the following model for the power transmission a ring modulator as function of wavelength $\lambda$ and bias voltage $V$:
 ```math
     T(\lambda, V) = \left \lvert \frac{t-a e^{2\pi j n(V) L/\lambda}}{1-t a e^{2\pi j n(V) L/\lambda}} \right\rvert^2.
 ```
-Transmission is determined by the factor $t = \sqrt{1-\kappa}$ governed by the power coupling coefficient $\kappa \in [0.1, 0.5]$. Power losses are accounted for by the decay factor $a= 10^{-\alpha L/20}$ with parameter $\alpha \in [0,0.6]$ $\frac{1}{\mu{\rm m}}$. $L$ denotes the ring length which is known to be quite tightly clustered around 5.0 $\mu{\rm m}$, i.e., $L \in [4.5, 5.5]$ $\mu {\rm m}$. 
+Transmission is determined by the factor $t = \sqrt{1-\kappa}$ governed by the power coupling coefficient $\kappa \in [0.1, 0.5]$. Power losses are accounted for by the decay factor $a= 10^{-\alpha L/20}$ with parameter $\alpha \in [0,0.6]$ $\frac{1}{\mu{\rm m}}$. $L$ denotes the ring length which is known to be quite tightly clustered around 5.0 $\mu{\rm m}$, i.e., $L \in [4.5, 5.5]$ $\mu {\rm m}$.
 
 We assume the effective refractive index depends linearly on the bias voltage $V$, i.e., $n(V) = n_0 + g_n V$, with parameters $n_0 \in [2.2, 2.4]$ and $g_n \in [-0.05, 0.05]$ $\frac{1}{{\rm V}}$.
 
-Using the AxModelFitter MCP, fit the parameters $(n_0, g_n, L, \alpha, \kappa)$ to the data provided in `data.csv` and compute the $R^2$-value. Finally, create a visualization of the fitting results, judge if the fit is good, if the parameter values are physically sensible and prepare a summary of your results, including visualizations, in a response.md file. 
+Using the AxModelFitter MCP, fit the parameters $(n_0, g_n, L, \alpha, \kappa)$ to the data provided in `data.csv` and compute the $R^2$-value. After fitting, compute the parameter covariance matrix to quantify uncertainty in the fitted parameters. Finally, create a visualization of the fitting results, judge if the fit is good, if the parameter values are physically sensible and prepare a summary of your results, including visualizations and parameter uncertainties, in a response.md file.

examples/axmodelfitter/ring_modulator/response.md

@@ -2,7 +2,7 @@
 
 ## Summary
 
-Successfully fitted the ring modulator transmission model to experimental data using AxModelFitter MCP. The model achieves an excellent fit with **R² = 0.964**, explaining 96.4% of the variance in the experimental data.
+Successfully fitted the ring modulator transmission model to experimental data using AxModelFitter MCP. The model achieves an excellent fit with **R² = 0.964374**, explaining 96.4% of the variance in the experimental data.
 
 ## Mathematical Model
 
@@ -13,6 +13,7 @@ T(\lambda, V) = \left \lvert \frac{t-a e^{2\pi j n(V) L/\lambda}}{1-t a e^{2\pi
 ```
 
 Where:
+
 - `t = √(1-κ)` - transmission coefficient determined by coupling coefficient κ
 - `a = 10^(-αL/20)` - decay factor accounting for ring losses with parameter α
 - `n(V) = n₀ + gₙV` - voltage-dependent effective refractive index
@@ -24,44 +25,50 @@ Where:
 
 | Parameter | Value | Unit | Physical Range | Status |
 |-----------|-------|------|---------------|---------|
-| n₀ | 2.277 | - | [2.2, 2.4] | ✅ Within range |
-| gₙ | 0.0196 | V⁻¹ | [-0.05, 0.05] | ✅ Within range |
-| L | 5.445 | μm | [4.5, 5.5] | ✅ Within range |
-| α | 0.403 | μm⁻¹ | [0, 0.6] | ✅ Within range |
-| κ | 0.214 | - | [0.1, 0.5] | ✅ Within range |
+| n₀ | 2.35504 | - | [2.2, 2.4] | ✅ Within range |
+| gₙ | 0.0202969 | V⁻¹ | [-0.05, 0.05] | ✅ Within range |
+| L | 5.26473 | μm | [4.5, 5.5] | ✅ Within range |
+| α | 0.41719 | μm⁻¹ | [0, 0.6] | ✅ Within range |
+| κ | 0.214359 | - | [0.1, 0.5] | ✅ Within range |
+
+Uncertainty (1σ, robust sandwich estimator): n₀ ≈ 2.41e-4, gₙ ≈ 2.31e-4 V⁻¹, L ≈ 1.73e-4 μm, α ≈ 7.06e-3 μm⁻¹, κ ≈ 2.14e-3.
+
+Correlation highlight (robust): n₀ and gₙ are strongly anti-correlated (ρ ≈ -0.945), indicating partial tradeoff between baseline index and voltage sensitivity in the fit.
 
 ## Fitting Quality Assessment
 
 ### Statistical Metrics
-- **R² coefficient**: 0.964 (excellent fit, >90% variance explained)
-- **Mean Squared Error (MSE)**: 0.002403
-- **Root Mean Square Error (RMSE)**: 0.049
-- **Mean Absolute Error (MAE)**: 0.039
+
+- **R² coefficient**: 0.964374 (excellent fit, >90% variance explained)
+- **Mean Squared Error (MSE)**: 0.002403405
+- **Root Mean Square Error (RMSE)**: 0.0490245
+- **Mean Absolute Error (MAE)**: 0.0391721
 
 ### Optimization Performance
-- **Execution time**: 3.75 seconds
-- **Function evaluations**: 7,020
+
+- **Execution time**: 15.01 seconds
+- **Function evaluations**: 55,645
 - **Optimizer**: nlopt (gradient-based)
 - **Convergence**: Successful
 
 ## Physical Parameter Analysis
 
 ### ✅ **All parameters are physically sensible:**
 
-1. **Baseline Refractive Index (n₀ = 2.277)**
+1. **Baseline Refractive Index (n₀ = 2.355)**
    - Within expected range for silicon photonics applications
    - Consistent with typical effective indices for ring resonators
 
-2. **Voltage Sensitivity (gₙ = 0.0196 V⁻¹)**
+2. **Voltage Sensitivity (gₙ = 0.0203 V⁻¹)**
    - Positive value indicates increasing refractive index with voltage
    - Magnitude is realistic for electro-optic modulation
    - Well within the expected range ±0.05 V⁻¹
 
-3. **Ring Length (L = 5.445 μm)**
+3. **Ring Length (L = 5.265 μm)**
    - Close to the expected value around 5.0 μm
    - Within the tight tolerance range [4.5, 5.5] μm
 
-4. **Loss Parameter (α = 0.403 μm⁻¹)**
+4. **Loss Parameter (α = 0.417 μm⁻¹)**
    - Moderate loss level, consistent with realistic ring resonators
    - Within the specified range [0, 0.6] μm⁻¹
 
@@ -99,18 +106,22 @@ The fitting results include comprehensive visualizations:
 
 All visualizations confirm the high quality of the fit and validate the physical model.
 
+Saved figure: `ring_modulator_fitting_results.png`.
+
 ## Conclusions
 
-1. **Successful Model Validation**: The ring modulator transmission model accurately describes the experimental data with R² = 0.964
+1. **Successful Model Validation**: The ring modulator transmission model accurately describes the experimental data with R² = 0.964374
 
 2. **Physical Parameter Consistency**: All fitted parameters fall within their expected physical ranges and are consistent with typical silicon photonic ring modulators
 
-3. **Robust Fitting Process**: The optimization converged efficiently (3.75s, 7,020 evaluations) using gradient-based methods
+3. **Robust Fitting Process**: The optimization converged successfully (15.01s, 55,645 evaluations) using gradient-based methods
 
 4. **Model Reliability**: The excellent fit quality and physical parameter values suggest the model is suitable for:
    - Device characterization
    - Performance optimization
    - Design parameter extraction
    - Predictive modeling for similar ring modulators
 
-The fitting demonstrates that the complex physics of ring modulator transmission can be accurately captured using the implemented model, providing valuable insights for silicon photonics device development and characterization.
+The fitting demonstrates that the complex physics of ring modulator transmission can be accurately captured using the implemented model, providing valuable insights for silicon photonics device development and characterization.
+
+