luminarycloud 0.22.2__py3-none-any.whl → 0.23.0__py3-none-any.whl

luminarycloud/physics_ai/datasets.py

@@ -0,0 +1,301 @@
+# Copyright 2025 Luminary Cloud, Inc. All Rights Reserved.
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import TYPE_CHECKING, List, Optional, Union
+
+from google.protobuf.json_format import ParseDict
+from google.protobuf.struct_pb2 import Struct
+
+from .._client import get_default_client
+from .._proto.api.v0.luminarycloud.physics_ai import physics_ai_pb2 as physaipb
+from .._wrapper import ProtoWrapper, ProtoWrapperBase
+from ..enum import QuantityType
+
+if TYPE_CHECKING:
+    from ..solution import Solution
+
+
+@dataclass
+class ParameterDefinition:
+    """
+    Defines a parameter in a dataset's parameter schema.
+
+    .. warning:: This feature is experimental and may change or be removed without notice.
+
+    Parameters
+    ----------
+    name : str
+        Parameter name (e.g., "stream_velocity", "air_density").
+    type : str
+        Parameter type: "scalar" for single values, "vector" for multi-dimensional values.
+    dimension : int
+        Number of dimensions: 1 for scalar, N for N-dimensional vector.
+    default_reference : list[float], optional
+        Default reference value(s) for normalization. Length must match dimension if provided.
+    description : str, optional
+        Description of the parameter.
+
+    Examples
+    --------
+    >>> # Scalar parameter
+    >>> air_density = ParameterDefinition(name="air_density", type="scalar", dimension=1, default_reference=[1.225])
+    >>> # Vector parameter (3D velocity)
+    >>> velocity = ParameterDefinition(name="inlet_velocity", type="vector", dimension=3, default_reference=[30.0, 0.0, 0.0])
+    >>> # 1D velocity (stream velocity magnitude)
+    >>> stream_velocity = ParameterDefinition(name="stream_velocity", type="vector", dimension=1, default_reference=[148.28])
+    """
+
+    name: str
+    type: str  # "scalar" or "vector"
+    dimension: int = 1
+    default_reference: List[float] = field(default_factory=list)
+    description: str = ""
+
+
+@dataclass(kw_only=True)
+class ExportConfig:
+    """
+    Configuration for exporting Physics AI dataset cases.
+
+    .. warning:: This feature is experimental and may change or be removed without notice.
+
+    Parameters
+    ----------
+    exclude_surfaces : List[str], optional
+        Raw mesh boundary names to exclude from export.
+    exclude_tags : List[str], optional
+        Geometry tag names to exclude. These are resolved to mesh boundary names
+        server-side based on the geometry tags.
+    fill_holes : float, optional
+        Size threshold for filling holes in the surface mesh. Default is 0 (no filling).
+    single_precision : bool, optional
+        If True, export floating point fields in single precision. Default is True.
+    process_volume : bool, optional
+        If True, include volume mesh data in the export. Default is False.
+    surface_fields_to_keep : List[QuantityType], optional
+        Surface fields to include in the export. If empty, all fields are kept.
+    volume_fields_to_keep : List[QuantityType], optional
+        Volume fields to include in the export. If empty, all fields are kept.
+    """
+
+    exclude_surfaces: List[str] = field(default_factory=list)
+    exclude_tags: List[str] = field(default_factory=list)
+    fill_holes: float = 0.0
+    single_precision: bool = True
+    process_volume: bool = False
+    surface_fields_to_keep: List[QuantityType] = field(default_factory=list)
+    volume_fields_to_keep: List[QuantityType] = field(default_factory=list)
+
+
+@dataclass
+class DatasetCaseInput:
+    """
+    Input for a case in a Physics AI dataset.
+
+    .. warning:: This feature is experimental and may change or be removed without notice.
+
+    Parameters
+    ----------
+    solution : Solution
+        The solution to include in the dataset.
+    params : dict, optional
+        Physics parameters for this case (e.g., {"alpha": 5.0, "stream_velocity": 50.0}).
+        These are used as conditioning inputs during training.
+    """
+
+    solution: "Solution"
+    params: dict = field(default_factory=dict)
+
+
+@ProtoWrapper(physaipb.PhysicsAiDatasetCase)
+class PhysicsAiDatasetCase(ProtoWrapperBase):
+    """
+    Represents a case in a Physics AI dataset.
+
+    .. warning:: This feature is experimental and may change or be removed without notice.
+    """
+
+    case_id: str
+    solution_id: str
+    simulation_id: str
+    _proto: physaipb.PhysicsAiDatasetCase
+
+
+@ProtoWrapper(physaipb.PhysicsAiDataset)
+class PhysicsAiDataset(ProtoWrapperBase):
+    """
+    Represents a Physics AI dataset containing training cases.
+
+    .. warning:: This feature is experimental and may change or be removed without notice.
+    """
+
+    id: str
+    created_by: str
+    name: str
+    description: str
+    is_locked: bool
+    creation_time: datetime
+    update_time: datetime
+    _proto: physaipb.PhysicsAiDataset
+
+
+def create_dataset(
+    name: str,
+    cases: List[Union["Solution", DatasetCaseInput]],
+    export_config: Optional[ExportConfig] = None,
+    description: str = "",
+    parameter_schema: Optional[List[ParameterDefinition]] = None,
+) -> PhysicsAiDataset:
+    """
+    Create a Physics AI dataset from simulation solutions.
+
+    .. warning:: This feature is experimental and may change or be removed without notice.
+
+    Parameters
+    ----------
+    name : str
+        Name of the dataset
+    cases : List[Union[Solution, DatasetCaseInput]]
+        List of solutions or DatasetCaseInput objects. For simple cases, pass
+        Solution objects directly. For cases with physics parameters (e.g., alpha,
+        stream_velocity), wrap in DatasetCaseInput.
+    export_config : ExportConfig, optional
+        Export configuration. If not provided, uses defaults (single_precision=True).
+    description : str, optional
+        Description of the dataset
+    parameter_schema : List[ParameterDefinition], optional
+        List of parameter definitions for this dataset. Each parameter defines
+        the type, dimension, and optional default reference value.
+
+    Returns
+    -------
+    PhysicsAiDataset
+        The created dataset
+
+    Examples
+    --------
+    Simple usage with solutions (uses default export config):
+
+    >>> dataset = create_dataset(
+    ...     name="my-dataset",
+    ...     cases=[solution1, solution2, solution3],
+    ...     parameter_schema=[
+    ...         ParameterDefinition(name="stream_velocity", type="vector", dimension=1),
+    ...     ],
+    ... )
+
+    With custom export config:
+
+    >>> dataset = create_dataset(
+    ...     name="my-dataset",
+    ...     cases=[solution1, solution2, solution3],
+    ...     export_config=ExportConfig(
+    ...         exclude_tags=["Farfield", "Symmetry"],
+    ...         surface_fields_to_keep=[QuantityType.PRESSURE, QuantityType.WALL_SHEAR_STRESS],
+    ...     ),
+    ...     parameter_schema=[
+    ...         ParameterDefinition(name="stream_velocity", type="vector", dimension=1, default_reference=[148.28]),
+    ...         ParameterDefinition(name="air_density", type="scalar", dimension=1, default_reference=[0.38]),
+    ...     ],
+    ... )
+
+    With physics parameters:
+
+    >>> dataset = create_dataset(
+    ...     name="parametric-dataset",
+    ...     cases=[
+    ...         DatasetCaseInput(solution1, params={"alpha": 5.0}),
+    ...         DatasetCaseInput(solution2, params={"alpha": 10.0}),
+    ...     ],
+    ...     export_config=ExportConfig(process_volume=True),
+    ...     parameter_schema=[
+    ...         ParameterDefinition(name="alpha", type="scalar", dimension=1),
+    ...     ],
+    ... )
+    """
+    # Use default export config if not provided
+    if export_config is None:
+        export_config = ExportConfig()
+
+    # Build case inputs
+    case_inputs = []
+    for case in cases:
+        if isinstance(case, DatasetCaseInput):
+            solution = case.solution
+            params = case.params
+        else:
+            # It's a Solution object
+            solution = case
+            params = {}
+
+        case_input = physaipb.CreatePhysicsAiDatasetCaseInput(
+            solution_id=solution.id,
+            simulation_id=solution.simulation_id,
+        )
+        if params:
+            case_input.params.CopyFrom(ParseDict(params, Struct()))
+        case_inputs.append(case_input)
+
+    # Build export config proto from dataclass
+    export_config_proto = physaipb.GetSolutionDataPhysicsAIRequest(
+        exclude_surfaces=export_config.exclude_surfaces,
+        exclude_tags=export_config.exclude_tags,
+        fill_holes=export_config.fill_holes,
+        single_precision=export_config.single_precision,
+        process_volume=export_config.process_volume,
+    )
+    if export_config.surface_fields_to_keep:
+        export_config_proto.surface_fields_to_keep.extend(export_config.surface_fields_to_keep)
+    if export_config.volume_fields_to_keep:
+        export_config_proto.volume_fields_to_keep.extend(export_config.volume_fields_to_keep)
+
+    # Build parameter schema from list of ParameterDefinition
+    param_schema_proto = physaipb.ParameterSchema()
+    if parameter_schema:
+        for param_def in parameter_schema:
+            param_schema_proto.parameters.append(
+                physaipb.ParameterDefinition(
+                    name=param_def.name,
+                    type=param_def.type,
+                    dimension=param_def.dimension,
+                    default_reference=param_def.default_reference,
+                    description=param_def.description,
+                )
+            )
+
+    req = physaipb.CreateDatasetRequest(
+        name=name,
+        description=description,
+        cases=case_inputs,
+        parameter_schema=param_schema_proto,
+        export_config=export_config_proto,
+    )
+
+    response = get_default_client().CreateDataset(req)
+    return PhysicsAiDataset(response.dataset)
+
+
+def list_datasets() -> List[PhysicsAiDataset]:
+    """
+    List Physics AI datasets accessible to the current user.
+
+    Returns datasets created by the current user plus platform-curated datasets.
+
+    .. warning:: This feature is experimental and may change or be removed without notice.
+
+    Returns
+    -------
+    List[PhysicsAiDataset]
+        A list of accessible Physics AI datasets, ordered by creation time (newest first).
+
+    Examples
+    --------
+    List all datasets:
+
+    >>> datasets = list_datasets()
+    >>> for ds in datasets:
+    ...     print(f"{ds.name}: {ds.id} (locked={ds.is_locked})")
+    """
+    req = physaipb.ListDatasetsRequest()
+    response = get_default_client().ListDatasets(req)
+    return [PhysicsAiDataset(dataset) for dataset in response.datasets]

luminarycloud/physics_ai/training_jobs.py

@@ -1,20 +1,33 @@
 # Copyright 2025 Luminary Cloud, Inc. All Rights Reserved.
-from typing import List, Optional
+from typing import List
 from datetime import datetime
 
 from .._client import get_default_client
 from .._proto.api.v0.luminarycloud.physics_ai import physics_ai_pb2 as physaipb
 from .._proto.base import base_pb2 as basepb
 from .._wrapper import ProtoWrapper, ProtoWrapperBase
-from ..types.ids import PhysicsAiArchitectureVersionID
+from ..types.ids import PhysicsAiArchitectureVersionID, PhysicsAiTrainingJobID
+
+
+def get_training_job(training_job_id: PhysicsAiTrainingJobID) -> "PhysicsAiTrainingJob":
+    """
+    Get a Physics AI training job by ID.
+
+    Args:
+        training_job_id: The ID of the training job to retrieve.
+
+    Returns:
+        The training job.
+    """
+    req = physaipb.GetTrainingJobRequest(training_job_id=training_job_id)
+    resp = get_default_client().GetTrainingJob(req)
+    return PhysicsAiTrainingJob(resp.training_job)
 
 
 @ProtoWrapper(physaipb.PhysicsAiTrainingJob)
 class PhysicsAiTrainingJob(ProtoWrapperBase):
     """
     Represents a Physics AI training job.
-
-    .. warning:: This feature is experimental and may change or be removed without notice.
     """
 
     id: str
@@ -24,6 +37,7 @@ class PhysicsAiTrainingJob(ProtoWrapperBase):
     training_data_source_type: physaipb.TrainingDataSourceType
     training_description: str
     external_dataset_uri: str
+    dataset_id: str
     initialization_type: physaipb.ModelInitializationType
     base_model_version_id: str
     status: basepb.JobStatus
@@ -35,3 +49,35 @@ class PhysicsAiTrainingJob(ProtoWrapperBase):
 
     def get_status(self) -> str:
         return basepb.JobStatusType.Name(self.status.typ)
+
+    def cancel(self) -> None:
+        """
+        Cancel this training job.
+
+        This method requests cancellation of the training job. If the job is already
+        suspended or suspending, this is a no-op.
+        """
+        req = physaipb.CancelTrainingJobRequest(training_job_id=self.id)
+        get_default_client().CancelTrainingJob(req)
+
+
+def list_training_jobs() -> List[PhysicsAiTrainingJob]:
+    """
+    List Physics AI training jobs for the current user.
+
+    Returns
+    -------
+    List[PhysicsAiTrainingJob]
+        A list of training jobs, ordered by creation time (newest first).
+
+    Examples
+    --------
+    List all training jobs:
+
+    >>> jobs = list_training_jobs()
+    >>> for job in jobs:
+    ...     print(f"{job.id}: {job.get_status()}")
+    """
+    req = physaipb.ListTrainingJobsRequest()
+    response = get_default_client().ListTrainingJobs(req)
+    return [PhysicsAiTrainingJob(job) for job in response.training_jobs]

luminarycloud/pipelines/__init__.py

@@ -2,8 +2,12 @@
 from .core import (
     Pipeline as Pipeline,
     PipelineParameter as PipelineParameter,
+    stage as stage,
     # Stage base class, mainly exported for testing
     Stage as Stage,
+    # RunScript stage lives in core because it's a special snowflake
+    RunScript as RunScript,
+    StopRun as StopRun,
 )
 
 from .parameters import (
@@ -23,28 +27,36 @@ from .stages import (
     ReadGeometryOutputs as ReadGeometryOutputs,
     ReadMesh as ReadMesh,
     ReadMeshOutputs as ReadMeshOutputs,
-    ModifyGeometry as ModifyGeometry,
-    ModifyGeometryOutputs as ModifyGeometryOutputs,
     Mesh as Mesh,
     MeshOutputs as MeshOutputs,
+    WaitForMesh as WaitForMesh,
+    WaitForMeshOutputs as WaitForMeshOutputs,
     Simulate as Simulate,
     SimulateOutputs as SimulateOutputs,
 )
 
 from .arguments import (
     PipelineArgs as PipelineArgs,
+    PP_NAMED_VARIABLE_SET_ID as PP_NAMED_VARIABLE_SET_ID,
     ArgNamedVariableSet as ArgNamedVariableSet,
 )
 
+from .flowables import (
+    FlowableType as FlowableType,
+    FlowableIOSchema as FlowableIOSchema,
+)
+
 from .api import (
     create_pipeline as create_pipeline,
-    list_pipelines as list_pipelines,
+    iterate_pipelines as iterate_pipelines,
     get_pipeline as get_pipeline,
     create_pipeline_job as create_pipeline_job,
     get_pipeline_job as get_pipeline_job,
-    list_pipeline_jobs as list_pipeline_jobs,
+    iterate_pipeline_jobs as iterate_pipeline_jobs,
     PipelineJobRecord as PipelineJobRecord,
     PipelineRecord as PipelineRecord,
     PipelineJobRunRecord as PipelineJobRunRecord,
+    PipelineTaskRecord as PipelineTaskRecord,
+    StageDefinition as StageDefinition,
     LogLine as LogLine,
 )