ionworks-api 0.1.0__py3-none-any.whl

ionworks/pipeline.py

@@ -0,0 +1,256 @@
+import os
+import re
+from typing import Any
+
+from pydantic import (
+    BaseModel,
+    Field,
+    ValidationError,
+    field_validator,
+    model_validator,
+)
+
+from .errors import IonworksError
+from .validators import run_validators_outbound
+
+
+def _prepare_payload(data: Any) -> Any:
+    """Prepare payload for API submission using outbound validators pipeline."""
+    return run_validators_outbound(data)
+
+
+class DataFitConfig(BaseModel):
+    objectives: dict[str, Any]
+    parameters: dict[str, Any]
+    cost: dict[str, Any] | None = None
+    optimizer: dict[str, Any] | None = None
+    existing_parameters: dict[str, Any] | None = None
+
+
+class EntryConfig(BaseModel):
+    values: dict[str, Any]
+
+
+class BuiltInEntryConfig(BaseModel):
+    name: str
+
+
+class CalculationConfig(BaseModel):
+    calculation: str
+    electrode: str | None = None
+    method: str | None = None
+    existing_parameters: dict[str, Any] | None = None
+
+
+class ValidationConfig(BaseModel):
+    objectives: dict[str, Any]
+    summary_stats: list[Any]
+    existing_parameters: dict[str, Any] | None = None
+
+
+class PipelineConfig(BaseModel):
+    project_id: str | None = Field(
+        default=None,
+        description="The project id to submit the pipeline to. "
+        "Can be found in the project settings page. "
+        "If not provided, will use PROJECT_ID environment variable.",
+    )
+    elements: dict[str, Any] = Field(
+        description="Dictionary of elements defining the pipeline. The key is the name "
+        "of the element and the value is the configuration of the element. "
+    )
+    name: str | None = Field(default=None, description="The name of the pipeline.")
+    description: str | None = Field(
+        default=None, description="The description of the pipeline."
+    )
+    options: dict[str, Any] | None = Field(
+        default=None,
+        description="Dictionary of options for the pipeline. "
+        "Options are used to configure the pipeline behavior. "
+        "Available options are: "
+        "live_progress_updates: bool ",
+    )
+
+    @field_validator("elements", mode="before")
+    @classmethod
+    def validate_elements_format(cls, v: Any) -> dict[str, Any]:
+        """Validate that elements is a dict, not the old list format."""
+        if isinstance(v, list):
+            raise ValueError(
+                "Pipeline elements must be provided as a dictionary, not a list. "
+                "The format has changed from:\n"
+                '  "elements": [{"name": {...}}, ...]\n'
+                "to:\n"
+                '  "elements": {"name": {...}, ...}\n'
+                "Please update your pipeline configuration."
+            )
+        if not isinstance(v, dict):
+            raise ValueError(
+                f"Pipeline elements must be a dictionary, got {type(v).__name__}"
+            )
+        return v
+
+    @model_validator(mode="after")
+    def set_defaults(self) -> "PipelineConfig":
+        """Set project_id from environment variable if not provided and defaults."""
+        if self.project_id is None:
+            env_project_id = os.getenv("PROJECT_ID")
+            if env_project_id is None:
+                raise ValueError(
+                    "project_id is required. Either provide it in the config "
+                    "or set the PROJECT_ID environment variable."
+                )
+            self.project_id = env_project_id
+        # Ensure options is never None to avoid 422 errors
+        if self.options is None:
+            self.options = {}
+        return self
+
+
+class DataFitResponse(BaseModel):
+    parameter_values: dict[str, Any]
+
+
+class CalculationResponse(BaseModel):
+    parameter_values: dict[str, Any]
+
+
+class ValidationResponse(BaseModel):
+    validation_results: dict[str, Any]
+    summary_stats: dict[str, list[Any]]
+
+
+class EntryResponse(BaseModel):
+    parameter_values: dict[str, Any]
+
+
+class PipelineSubmissionResponse(BaseModel):
+    id: str
+    name: str
+    description: str | None = None
+    status: str
+    error: str | None = None
+
+
+class PipelineResponse(BaseModel):
+    result: dict[str, Any]
+    element_results: dict[str, Any]
+
+
+class PipelineClient:
+    def __init__(self, client: Any):
+        self.client = client
+
+    def create(self, config: dict[str, Any]) -> PipelineSubmissionResponse:
+        """Run a complete pipeline with the given configuration.
+
+        Parameters
+        ----------
+        config : dict[str, Any]
+            Dictionary containing configuration for the pipeline.
+
+        Returns
+        -------
+        PipelineSubmissionResponse
+            The pipeline submission response.
+
+        Raises
+        ------
+        ValueError
+            If the configuration is invalid.
+        """
+        try:
+            validated_config = PipelineConfig(**config)
+            payload = _prepare_payload(validated_config.model_dump())
+            response_data = self.client.post("/pipelines", payload)
+            return PipelineSubmissionResponse(**response_data)
+        except ValidationError as e:
+            raise ValueError(f"Invalid pipeline configuration: {e}") from e
+        except IonworksError as e:
+            error_msg = str(e.message)
+            # Check for invalid UUID format in project_id
+            uuid_match = re.search(
+                r'invalid input syntax for type uuid: "([^"]*)"', error_msg
+            )
+            if uuid_match:
+                invalid_id = uuid_match.group(1)
+                raise ValueError(
+                    f"Invalid project_id format: '{invalid_id}' is not a valid UUID. "
+                    "Please provide a valid project ID from your project settings page."
+                ) from e
+            # Check for row-level security policy violation (project not accessible)
+            if "violates row-level security policy" in error_msg:
+                project_id = validated_config.project_id
+                raise ValueError(
+                    f"Access denied: The project '{project_id}' is not accessible "
+                    "with your API key. Please verify that your API key has access "
+                    "to this project."
+                ) from e
+            # Re-raise original error for other cases
+            raise
+
+    def list(self, project_id: str | None = None) -> list[PipelineSubmissionResponse]:
+        """List all pipelines.
+
+        Parameters
+        ----------
+        project_id : str, optional
+            The project id to filter pipelines. If not provided, will use
+            PROJECT_ID environment variable.
+
+        Returns
+        -------
+        list[PipelineSubmissionResponse]
+            List of pipeline submission responses.
+        """
+        if project_id is None:
+            project_id = os.getenv("PROJECT_ID")
+            if project_id is None:
+                raise ValueError(
+                    "project_id is required. Either provide it as an argument "
+                    "or set the PROJECT_ID environment variable."
+                )
+
+        endpoint = f"/pipelines?project_id={project_id}"
+        try:
+            response_data = self.client.get(endpoint)
+            if not isinstance(response_data, list):
+                raise ValueError(
+                    f"Unexpected response format from {endpoint}: expected a list, "
+                    f"got {type(response_data).__name__}"
+                )
+            return [PipelineSubmissionResponse(**item) for item in response_data]
+        except ValidationError as e:
+            raise ValueError(f"Invalid item format in list from {endpoint}: {e}") from e
+
+    def get(self, job_id: str) -> PipelineSubmissionResponse:
+        """Get the pipeline response for the given job id.
+
+        Parameters
+        ----------
+        job_id : str
+            The job id.
+
+        Returns
+        -------
+        PipelineSubmissionResponse
+            The pipeline submission response.
+        """
+        response_data = self.client.get(f"/pipelines/{job_id}")
+        return PipelineSubmissionResponse(**response_data)
+
+    def result(self, job_id: str) -> PipelineResponse:
+        """Get the result for the given job id.
+
+        Parameters
+        ----------
+        job_id : str
+            The job id.
+
+        Returns
+        -------
+        PipelineResponse
+            The pipeline results.
+        """
+        response_data = self.client.get(f"/pipelines/{job_id}/result")
+        return PipelineResponse(**response_data)

ionworks/simulation.py

@@ -0,0 +1,371 @@
+from __future__ import annotations
+
+from datetime import datetime, timedelta
+import time
+from typing import Any, List, cast
+
+from pydantic import BaseModel, Field, ValidationError
+
+
+class QuickModelConfig(BaseModel):
+    """Quick model configuration for protocol-based simulations."""
+
+    capacity: float = Field(default=1.0, description="Cell capacity in Ah")
+    chemistry: str = Field(default="NMC/Graphite", description="Chemistry name")
+
+
+class ProtocolExperimentConfig(BaseModel):
+    """Protocol experiment configuration."""
+
+    protocol: str = Field(description="YAML protocol string (UCP format)")
+    name: str = Field(description="Protocol name for template naming")
+
+
+class ProtocolSimulationRequest(BaseModel):
+    """Request model for single protocol-based simulation."""
+
+    parameterized_model: Any = Field(
+        description=(
+            "Model can be: quick_model dict, full model dict, or model ID string"
+        )
+    )
+    protocol_experiment: ProtocolExperimentConfig = Field(
+        description="Protocol experiment configuration"
+    )
+    experiment_parameters: dict[str, float] | None = Field(
+        default=None,
+        description=("Experiment parameters for any inputs in the protocol."),
+    )
+    design_parameters: dict[str, float] | None = Field(
+        default=None, description="Design parameters for the simulation"
+    )
+    max_backward_jumps: int | None = Field(
+        default=None,
+        description="Maximum backward jumps allowed (for goto statements)",
+    )
+    study_id: str | None = Field(default=None, description="Optional study UUID")
+    extra_variables: list[str] | None = Field(
+        default=None,
+        description=(
+            "Optional list of extra variables to include in simulation output "
+            "(e.g., ['Negative electrode potential [V]', 'Positive electrode "
+            "potential [V]']). If provided, these override any extra variables "
+            "defined in the experiment template."
+        ),
+    )
+
+
+class DOERow(BaseModel):
+    """Design of experiments row configuration."""
+
+    type: str = Field(description="Type: 'range', 'discrete', or 'normal'")
+    name: str = Field(description="Parameter name")
+    # For range type
+    min: float | None = Field(default=None, description="Minimum value")
+    max: float | None = Field(default=None, description="Maximum value")
+    count: int | None = Field(
+        default=None, description="Number of samples (for grid/random)"
+    )
+    # For discrete type
+    values: list[float] | None = Field(default=None, description="Discrete values")
+    # For normal type
+    mean: float | None = Field(default=None, description="Mean value")
+    std: float | None = Field(default=None, description="Standard deviation")
+
+
+class DesignParametersDOE(BaseModel):
+    """Design of experiments configuration."""
+
+    sampling: str = Field(
+        description="Sampling strategy: 'grid', 'random', or 'latin_hypercube'"
+    )
+    rows: list[DOERow] = Field(description="DOE row configurations")
+    count: int | None = Field(
+        default=None, description="Total count for non-grid sampling"
+    )
+
+
+class ProtocolSimulationBatchRequest(BaseModel):
+    """Request model for batch protocol-based simulation."""
+
+    parameterized_model: Any = Field(
+        description=(
+            "Model can be: quick_model dict, full model dict, or model ID string"
+        )
+    )
+    protocol_experiment: ProtocolExperimentConfig = Field(
+        description="Protocol experiment configuration"
+    )
+    design_parameters_doe: DesignParametersDOE = Field(
+        description="Design of experiments configuration"
+    )
+    experiment_parameters: dict[str, float] | None = Field(
+        default=None,
+        description=("Experiment parameters for any inputs in the protocol."),
+    )
+    max_backward_jumps: int | None = Field(
+        default=None,
+        description="Maximum backward jumps allowed (for goto statements)",
+    )
+    study_id: str | None = Field(default=None, description="Optional study UUID")
+    extra_variables: list[str] | None = Field(
+        default=None,
+        description=(
+            "Optional list of extra variables to include in simulation output "
+            "(e.g., ['Negative electrode potential [V]', 'Positive electrode "
+            "potential [V]']). If provided, these override any extra variables "
+            "defined in the experiment template."
+        ),
+    )
+
+
+class SimulationResponse(BaseModel):
+    """Response model for simulation creation."""
+
+    simulation_id: str = Field(description="Simulation UUID")
+    job_id: str = Field(description="Job UUID")
+
+
+class SimulationClient:
+    """Client for running simulations."""
+
+    def __init__(self, client: Any):
+        self.client = client
+
+    def protocol(self, config: dict[str, Any]) -> SimulationResponse:
+        """Create a single protocol-based simulation.
+
+        Parameters
+        ----------
+        config : dict[str, Any]
+            Configuration dictionary containing:
+            - parameterized_model: quick_model dict, full model dict, or model ID string
+            - protocol_experiment: ProtocolExperimentConfig dict with protocol
+              and name fields
+            - experiment_parameters: Optional ProtocolExperimentParameters dict
+              with initial_soc and initial_temperature
+            - design_parameters: Optional dict[str, float]
+            - max_backward_jumps: Optional int
+            - study_id: Optional str
+            - extra_variables: Optional list[str] - Extra variables to include
+              in simulation output
+
+        Returns
+        -------
+        SimulationResponse
+            Response containing simulation_id and job_id.
+
+        Raises
+        ------
+        ValueError
+            If the configuration is invalid.
+        """
+        endpoint = "/simulations/protocol"
+        try:
+            validated_config = ProtocolSimulationRequest(**config)
+            response_data = self.client.post(
+                endpoint, json_payload=validated_config.model_dump(exclude_none=True)
+            )
+            return SimulationResponse(**response_data)
+        except ValidationError as e:
+            raise ValueError(f"Invalid protocol simulation configuration: {e}") from e
+
+    def protocol_batch(self, config: dict[str, Any]) -> list[SimulationResponse]:
+        """Create multiple protocol-based simulations using DOE.
+
+        Parameters
+        ----------
+        config : dict[str, Any]
+            Configuration dictionary containing:
+            - parameterized_model: quick_model dict, full model dict, or model ID string
+            - protocol_experiment: ProtocolExperimentConfig dict with protocol
+              and name fields
+            - design_parameters_doe: DesignParametersDOE dict
+            - experiment_parameters: Optional ProtocolExperimentParameters dict
+            - max_backward_jumps: Optional int
+            - study_id: Optional str
+            - extra_variables: Optional list[str] - Extra variables to include
+              in simulation output
+
+        Returns
+        -------
+        list[SimulationResponse]
+            List of responses, each containing simulation_id and job_id.
+
+        Raises
+        ------
+        ValueError
+            If the configuration is invalid.
+        """
+        endpoint = "/simulations/protocol/batch"
+        try:
+            validated_config = ProtocolSimulationBatchRequest(**config)
+            response_data = self.client.post(
+                endpoint, json_payload=validated_config.model_dump(exclude_none=True)
+            )
+            if not isinstance(response_data, list):
+                raise ValueError(
+                    f"Unexpected response format from {endpoint}: expected a "
+                    f"list, got {type(response_data).__name__}"
+                )
+            return [SimulationResponse(**item) for item in response_data]
+        except ValidationError as e:
+            raise ValueError(
+                f"Invalid batch protocol simulation configuration: {e}"
+            ) from e
+
+    def list(self) -> list[dict[str, Any]]:
+        """List all simulations for the current user.
+
+        Returns
+        -------
+        list[dict[str, Any]]
+            List of simulation objects with joined model and experiment data.
+        """
+        endpoint = "/simulations"
+        response_data = self.client.get(endpoint)
+        if not isinstance(response_data, list):
+            raise ValueError(
+                f"Unexpected response format from {endpoint}: expected a list, "
+                f"got {type(response_data).__name__}"
+            )
+        return response_data
+
+    def get(self, simulation_id: str) -> dict[str, Any]:
+        """Get a specific simulation by ID.
+
+        Parameters
+        ----------
+        simulation_id : str
+            The UUID of the simulation to retrieve.
+
+        Returns
+        -------
+        dict[str, Any]
+            Simulation object with full joined data including model, experiment,
+            and simulation_data (null if not completed).
+        """
+        endpoint = f"/simulations/{simulation_id}"
+        response_data = self.client.get(endpoint)
+        return cast(dict[str, Any], response_data)
+
+    def get_result(self, simulation_id: str) -> dict[str, Any]:
+        """Get simulation data/result for a completed simulation.
+
+        Parameters
+        ----------
+        simulation_id : str
+            The UUID of the simulation.
+
+        Returns
+        -------
+        dict[str, Any]
+            Simulation data object containing time_series, steps, and metrics.
+            Returns 404 if simulation hasn't completed yet.
+
+        Raises
+        ------
+        Exception
+            If simulation data not found (simulation may not be completed yet).
+            The client will raise an appropriate error for 404 responses.
+        """
+        endpoint = f"/simulations/{simulation_id}/result"
+        response_data = self.client.get(endpoint)
+        return cast(dict[str, Any], response_data)
+
+    def wait_for_completion(
+        self,
+        simulation_id: str | List[str],  # List to avoid conflict with list() method
+        timeout: int = 60,
+        poll_interval: int = 2,
+        verbose: bool = True,
+    ) -> dict[str, Any] | List[dict[str, Any]]:
+        """Wait for simulation(s) to complete by polling until done or timeout.
+
+        Parameters
+        ----------
+        simulation_id : str | list[str]
+            Single simulation ID or list of simulation IDs to wait for.
+        timeout : int, optional
+            Maximum time to wait in seconds (default: 60).
+        poll_interval : int, optional
+            Time between polls in seconds (default: 2).
+        verbose : bool, optional
+            Whether to print status updates (default: True).
+
+        Returns
+        -------
+        dict[str, Any] | list[dict[str, Any]]
+            Completed simulation(s). Returns single dict if single ID provided,
+            list of dicts if list of IDs provided. Only returns completed
+            simulations if timeout is reached.
+
+        Raises
+        ------
+        TimeoutError
+            If timeout is reached before all simulations complete.
+        """
+        is_single = isinstance(simulation_id, str)
+        if is_single:
+            simulation_ids: List[str] = [simulation_id]  # type: ignore[list-item]
+        else:
+            simulation_ids = simulation_id  # type: ignore[assignment]
+        timeout_delta = timedelta(seconds=timeout)
+        start_time = datetime.now()
+        completed_simulations: dict[str, dict[str, Any]] = {}
+
+        if verbose:
+            print(f"Polling for {len(simulation_ids)} simulation(s) completion...")
+
+        while datetime.now() - start_time < timeout_delta:
+            completed_count = len(completed_simulations)
+            for sim_id in simulation_ids:
+                if sim_id in completed_simulations:
+                    continue
+                try:
+                    simulation = self.get(sim_id)
+                    if simulation.get("simulation_data") is not None:
+                        completed_simulations[sim_id] = simulation
+                        completed_count += 1
+                except Exception:
+                    # Continue polling if there's an error
+                    pass
+
+            elapsed = (datetime.now() - start_time).seconds
+            if verbose:
+                print(
+                    f"  Status: {completed_count}/{len(simulation_ids)} completed "
+                    f"(elapsed: {elapsed}s)"
+                )
+
+            if len(completed_simulations) == len(simulation_ids):
+                if verbose:
+                    print("All simulations completed!")
+                break
+
+            time.sleep(poll_interval)
+        else:
+            # Timeout reached
+            if verbose:
+                print(
+                    f"Timeout: Only {len(completed_simulations)}/{len(simulation_ids)} "
+                    f"simulations completed within {timeout} seconds"
+                )
+            if not completed_simulations:
+                raise TimeoutError(f"No simulations completed within {timeout} seconds")
+
+        # Return results in the same format as input
+        if is_single:
+            if simulation_ids[0] not in completed_simulations:
+                raise TimeoutError(
+                    f"Simulation {simulation_ids[0]} did not complete within "
+                    f"{timeout} seconds"
+                )
+            return completed_simulations[simulation_ids[0]]
+        else:
+            # Return list of completed simulations in order
+            return [
+                completed_simulations[sim_id]
+                for sim_id in simulation_ids
+                if sim_id in completed_simulations
+            ]