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

luminarycloud/params/simulation/physics/solution_controls/linear_solver_type/krylov_amg_.py

@@ -15,7 +15,7 @@ from luminarycloud._proto.client.entity_pb2 import EntityIdentifier
 from luminarycloud._helpers import CodeRepr
 import luminarycloud.params.enum._enum_wrappers as enum
 
-from luminarycloud.params.simulation._lib import ParamGroupWrapper, create_unique_id
+from luminarycloud.params._lib import ParamGroupWrapper, create_unique_id
 
 from luminarycloud.params.simulation.physics.solution_controls.linear_solver_type_ import (
     LinearSolverType,

luminarycloud/params/simulation/physics/solution_controls/linear_solver_type_.py

@@ -15,7 +15,7 @@ from luminarycloud._proto.client.entity_pb2 import EntityIdentifier
 from luminarycloud._helpers import CodeRepr
 import luminarycloud.params.enum._enum_wrappers as enum
 
-from luminarycloud.params.simulation._lib import ParamGroupWrapper, create_unique_id
+from luminarycloud.params._lib import ParamGroupWrapper, create_unique_id
 
 
 @dataclass(kw_only=True)

luminarycloud/params/simulation/physics_.py

@@ -15,7 +15,7 @@ from luminarycloud._proto.client.entity_pb2 import EntityIdentifier
 from luminarycloud._helpers import CodeRepr
 import luminarycloud.params.enum._enum_wrappers as enum
 
-from luminarycloud.params.simulation._lib import ParamGroupWrapper, create_unique_id
+from luminarycloud.params._lib import ParamGroupWrapper, create_unique_id
 
 from luminarycloud.params.simulation.physics.fluid_ import Fluid
 from luminarycloud.params.simulation.physics.fluid_ import *

luminarycloud/params/simulation/simulation_param_.py

@@ -15,7 +15,7 @@ from luminarycloud._proto.client.entity_pb2 import EntityIdentifier
 from luminarycloud._helpers import CodeRepr
 import luminarycloud.params.enum._enum_wrappers as enum
 
-from luminarycloud.params.simulation._lib import ParamGroupWrapper, create_unique_id
+from luminarycloud.params._lib import ParamGroupWrapper, create_unique_id
 
 from luminarycloud.params.simulation.adaptive_mesh_refinement_ import AdaptiveMeshRefinement
 from luminarycloud.params.simulation.adaptive_mesh_refinement_ import *
@@ -37,6 +37,8 @@ from luminarycloud.params.simulation.multi_physics_coupling_options_ import (
     MultiPhysicsCouplingOptions,
 )
 from luminarycloud.params.simulation.multi_physics_coupling_options_ import *
+from luminarycloud.params.simulation.nonlinear_control_ import NonlinearControl
+from luminarycloud.params.simulation.nonlinear_control_ import *
 from luminarycloud.params.simulation.output_ import Output
 from luminarycloud.params.simulation.output_ import *
 from luminarycloud.params.simulation.particle_group_ import ParticleGroup
@@ -65,6 +67,8 @@ class SimulationParam(CodeRepr, ParamGroupWrapper[clientpb.SimulationParam]):
     "Basic parameters used by a solution."
     adjoint: Adjoint = field(default_factory=Adjoint)
     "Settings for adjoint sensitivity analysis."
+    nonlinear_control: NonlinearControl = field(default_factory=NonlinearControl)
+    "Determine input values such that target output values are met."
     time: Time = field(default_factory=Time)
     "Time parameters used by a transient solution (required if flow_behavior = transient)."
     motion_data: list[MotionData] = field(default_factory=list)
@@ -104,6 +108,8 @@ class SimulationParam(CodeRepr, ParamGroupWrapper[clientpb.SimulationParam]):
             _proto.general.CopyFrom(self.basic._to_proto())
         if self.adjoint is not None:
             _proto.adjoint.CopyFrom(self.adjoint._to_proto())
+        if self.nonlinear_control is not None:
+            _proto.nonlinear_control.CopyFrom(self.nonlinear_control._to_proto())
         if self.time is not None:
             _proto.time.CopyFrom(self.time._to_proto())
         if self.motion_data is not None:
@@ -139,6 +145,7 @@ class SimulationParam(CodeRepr, ParamGroupWrapper[clientpb.SimulationParam]):
         self.volume_entity = [VolumeEntity.from_proto(v) for v in proto.volume_entity]
         self.basic = General.from_proto(proto.general)
         self.adjoint = Adjoint.from_proto(proto.adjoint)
+        self.nonlinear_control = NonlinearControl.from_proto(proto.nonlinear_control)
         self.time = Time.from_proto(proto.time)
         self.motion_data = [MotionData.from_proto(v) for v in proto.motion_data]
         self.particle_group = [ParticleGroup.from_proto(v) for v in proto.particle_group]

luminarycloud/params/simulation/sliding_interfaces_.py

@@ -15,7 +15,7 @@ from luminarycloud._proto.client.entity_pb2 import EntityIdentifier
 from luminarycloud._helpers import CodeRepr
 import luminarycloud.params.enum._enum_wrappers as enum
 
-from luminarycloud.params.simulation._lib import ParamGroupWrapper, create_unique_id
+from luminarycloud.params._lib import ParamGroupWrapper, create_unique_id
 
 
 @dataclass(kw_only=True)

luminarycloud/params/simulation/surface_name_.py

@@ -15,7 +15,7 @@ from luminarycloud._proto.client.entity_pb2 import EntityIdentifier
 from luminarycloud._helpers import CodeRepr
 import luminarycloud.params.enum._enum_wrappers as enum
 
-from luminarycloud.params.simulation._lib import ParamGroupWrapper, create_unique_id
+from luminarycloud.params._lib import ParamGroupWrapper, create_unique_id
 
 
 @dataclass(kw_only=True)

luminarycloud/params/simulation/time/compute_statistics/compute_statistics_off_.py

@@ -15,7 +15,7 @@ from luminarycloud._proto.client.entity_pb2 import EntityIdentifier
 from luminarycloud._helpers import CodeRepr
 import luminarycloud.params.enum._enum_wrappers as enum
 
-from luminarycloud.params.simulation._lib import ParamGroupWrapper, create_unique_id
+from luminarycloud.params._lib import ParamGroupWrapper, create_unique_id
 
 from luminarycloud.params.simulation.time.compute_statistics_ import ComputeStatistics
 from luminarycloud.params.simulation.time.compute_statistics_ import *

luminarycloud/params/simulation/time/compute_statistics/compute_statistics_on_.py

@@ -15,7 +15,7 @@ from luminarycloud._proto.client.entity_pb2 import EntityIdentifier
 from luminarycloud._helpers import CodeRepr
 import luminarycloud.params.enum._enum_wrappers as enum
 
-from luminarycloud.params.simulation._lib import ParamGroupWrapper, create_unique_id
+from luminarycloud.params._lib import ParamGroupWrapper, create_unique_id
 
 from luminarycloud.params.simulation.time.compute_statistics_ import ComputeStatistics
 from luminarycloud.params.simulation.time.compute_statistics_ import *

luminarycloud/params/simulation/time/compute_statistics_.py

@@ -15,7 +15,7 @@ from luminarycloud._proto.client.entity_pb2 import EntityIdentifier
 from luminarycloud._helpers import CodeRepr
 import luminarycloud.params.enum._enum_wrappers as enum
 
-from luminarycloud.params.simulation._lib import ParamGroupWrapper, create_unique_id
+from luminarycloud.params._lib import ParamGroupWrapper, create_unique_id
 
 
 @dataclass(kw_only=True)

luminarycloud/params/simulation/time/time_marching/time_explicit_.py

@@ -15,7 +15,7 @@ from luminarycloud._proto.client.entity_pb2 import EntityIdentifier
 from luminarycloud._helpers import CodeRepr
 import luminarycloud.params.enum._enum_wrappers as enum
 
-from luminarycloud.params.simulation._lib import ParamGroupWrapper, create_unique_id
+from luminarycloud.params._lib import ParamGroupWrapper, create_unique_id
 
 from luminarycloud.params.simulation.time.time_marching_ import TimeMarching
 from luminarycloud.params.simulation.time.time_marching_ import *

luminarycloud/params/simulation/time/time_marching/time_implicit_.py

@@ -15,7 +15,7 @@ from luminarycloud._proto.client.entity_pb2 import EntityIdentifier
 from luminarycloud._helpers import CodeRepr
 import luminarycloud.params.enum._enum_wrappers as enum
 
-from luminarycloud.params.simulation._lib import ParamGroupWrapper, create_unique_id
+from luminarycloud.params._lib import ParamGroupWrapper, create_unique_id
 
 from luminarycloud.params.simulation.time.time_marching_ import TimeMarching
 from luminarycloud.params.simulation.time.time_marching_ import *

luminarycloud/params/simulation/time/time_marching_.py

@@ -15,7 +15,7 @@ from luminarycloud._proto.client.entity_pb2 import EntityIdentifier
 from luminarycloud._helpers import CodeRepr
 import luminarycloud.params.enum._enum_wrappers as enum
 
-from luminarycloud.params.simulation._lib import ParamGroupWrapper, create_unique_id
+from luminarycloud.params._lib import ParamGroupWrapper, create_unique_id
 
 
 @dataclass(kw_only=True)

luminarycloud/params/simulation/time/time_step_ramp/time_step_ramp_off_.py

@@ -15,7 +15,7 @@ from luminarycloud._proto.client.entity_pb2 import EntityIdentifier
 from luminarycloud._helpers import CodeRepr
 import luminarycloud.params.enum._enum_wrappers as enum
 
-from luminarycloud.params.simulation._lib import ParamGroupWrapper, create_unique_id
+from luminarycloud.params._lib import ParamGroupWrapper, create_unique_id
 
 from luminarycloud.params.simulation.time.time_step_ramp_ import TimeStepRamp
 from luminarycloud.params.simulation.time.time_step_ramp_ import *

luminarycloud/params/simulation/time/time_step_ramp/time_step_ramp_on_.py

@@ -15,7 +15,7 @@ from luminarycloud._proto.client.entity_pb2 import EntityIdentifier
 from luminarycloud._helpers import CodeRepr
 import luminarycloud.params.enum._enum_wrappers as enum
 
-from luminarycloud.params.simulation._lib import ParamGroupWrapper, create_unique_id
+from luminarycloud.params._lib import ParamGroupWrapper, create_unique_id
 
 from luminarycloud.params.simulation.time.time_step_ramp_ import TimeStepRamp
 from luminarycloud.params.simulation.time.time_step_ramp_ import *

luminarycloud/params/simulation/time/time_step_ramp_.py

@@ -15,7 +15,7 @@ from luminarycloud._proto.client.entity_pb2 import EntityIdentifier
 from luminarycloud._helpers import CodeRepr
 import luminarycloud.params.enum._enum_wrappers as enum
 
-from luminarycloud.params.simulation._lib import ParamGroupWrapper, create_unique_id
+from luminarycloud.params._lib import ParamGroupWrapper, create_unique_id
 
 
 @dataclass(kw_only=True)

luminarycloud/params/simulation/time_.py

@@ -15,7 +15,7 @@ from luminarycloud._proto.client.entity_pb2 import EntityIdentifier
 from luminarycloud._helpers import CodeRepr
 import luminarycloud.params.enum._enum_wrappers as enum
 
-from luminarycloud.params.simulation._lib import ParamGroupWrapper, create_unique_id
+from luminarycloud.params._lib import ParamGroupWrapper, create_unique_id
 
 from luminarycloud.params.simulation.time.compute_statistics_ import ComputeStatistics
 from luminarycloud.params.simulation.time.compute_statistics_ import *

luminarycloud/params/simulation/volume_entity_.py

@@ -15,7 +15,7 @@ from luminarycloud._proto.client.entity_pb2 import EntityIdentifier
 from luminarycloud._helpers import CodeRepr
 import luminarycloud.params.enum._enum_wrappers as enum
 
-from luminarycloud.params.simulation._lib import ParamGroupWrapper, create_unique_id
+from luminarycloud.params._lib import ParamGroupWrapper, create_unique_id
 
 
 @dataclass(kw_only=True)

luminarycloud/physics_ai/__init__.py

@@ -2,6 +2,8 @@ from .architectures import (
     list_architectures as list_architectures,
     PhysicsAiArchitecture as PhysicsAiArchitecture,
     PhysicsAiArchitectureVersion as PhysicsAiArchitectureVersion,
+    ModelOutputs as ModelOutputs,
+    InputParameter as InputParameter,
 )
 from .models import (
     list_pretrained_models as list_pretrained_models,
@@ -26,4 +28,11 @@ from .datasets import (
     DatasetCaseInput as DatasetCaseInput,
     ExportConfig as ExportConfig,
     PhysicsAiDataset as PhysicsAiDataset,
+    ParameterDefinition as ParameterDefinition,
+)
+
+from .training_jobs import (
+    get_training_job as get_training_job,
+    list_training_jobs as list_training_jobs,
+    PhysicsAiTrainingJob as PhysicsAiTrainingJob,
 )

luminarycloud/physics_ai/architectures.py

@@ -1,16 +1,68 @@
 # Copyright 2025 Luminary Cloud, Inc. All Rights Reserved.
-from typing import List, Optional
+import copy
+from dataclasses import dataclass, field
 import json
 
 from .._client import get_default_client
 from .._proto.api.v0.luminarycloud.physics_ai import physics_ai_pb2 as physaipb
+from .._proto.quantity import quantity_pb2 as quantitypb
 from .._wrapper import ProtoWrapper, ProtoWrapperBase
 from ..types.ids import PhysicsAiArchitectureID, PhysicsAiArchitectureVersionID
 from ..enum.physics_ai_lifecycle_state import PhysicsAiLifecycleState
+from ..enum.quantity_type import QuantityType
 
 from .training_jobs import PhysicsAiTrainingJob
 
 
+@dataclass
+class ModelOutputs:
+    """
+    Defines the output variables for the model.
+
+    Parameters
+    ----------
+    surface : list[QuantityType], optional
+        Surface outputs to predict (e.g., PRESSURE, WALL_SHEAR_STRESS)
+    volume : list[QuantityType], optional
+        Volume outputs to predict (e.g., VELOCITY, PRESSURE)
+    """
+
+    surface: list[QuantityType] | None = None
+    volume: list[QuantityType] | None = None
+
+
+@dataclass
+class InputParameter:
+    """
+    Defines an input parameter for the model with reference values for normalization.
+
+    Type is inherited from the dataset's ParameterSchema, so only the name and
+    optional reference values need to be specified here.
+
+    Parameters
+    ----------
+    name : str
+        Parameter name (must exist in the dataset's ParameterSchema)
+    reference : list[float], optional
+        Reference value(s) for normalization. If not provided, falls back to
+        the default_reference from the dataset's ParameterSchema. For scalars,
+        a single-element list. For vectors, the number of elements must match
+        the dimension from the schema.
+
+    Examples
+    --------
+    >>> # Use schema defaults (no reference specified)
+    >>> air_density = InputParameter(name="air_density")
+    >>> # Override with custom reference values
+    >>> velocity = InputParameter(name="inlet_velocity", reference=[30.0, 0.0, 0.0])
+    >>> # Scalar parameter with explicit reference
+    >>> stream_velocity = InputParameter(name="stream_velocity", reference=[148.28])
+    """
+
+    name: str
+    reference: list[float] = field(default_factory=list)
+
+
 @ProtoWrapper(physaipb.PhysicsAiArchitectureVersion)
 class PhysicsAiArchitectureVersion(ProtoWrapperBase):
     """
@@ -26,30 +78,87 @@ class PhysicsAiArchitectureVersion(ProtoWrapperBase):
     lifecycle_state: PhysicsAiLifecycleState
     _proto: physaipb.PhysicsAiArchitectureVersion
 
-    def get_training_description(self, config: dict) -> str:
-        if config.get("description"):
-            return config["description"]
+    def get_training_description(self, _internal: dict | None) -> str:
+        if _internal and _internal.get("description"):
+            return _internal["description"]
 
         desc = f"Training job for architecture {self.name}"
-        if config.get("custom_args"):
-            desc += f" with custom args: {config['custom_args']}"
+        if _internal and _internal.get("custom_args"):
+            desc += f" with custom args: {_internal['custom_args']}"
         return desc
 
-    def train(self, config: dict) -> PhysicsAiTrainingJob:
+    def train(
+        self,
+        dataset_id: str | None = None,
+        outputs: ModelOutputs | None = None,
+        inputs: list[InputParameter] | None = None,
+        _internal: dict | None = None,
+    ) -> PhysicsAiTrainingJob:
         """
         Submit a training job for this architecture version.
 
         Parameters
         ----------
-        config : dict, optional
-            Training configuration dictionary
+        dataset_id : str, optional
+            ID of a managed dataset to train on. If provided, uses session-based
+            GCS authentication.
+        outputs : ModelOutputs, optional
+            Output variables for the model (surface and/or volume quantities).
+        inputs : list[InputParameter], optional
+            Input parameters for the model with reference values for normalization.
+            Parameter names must exist in the dataset's ParameterSchema.
+        _internal : dict, optional
+            Internal configuration (resources, priority_class, custom_args, etc.).
+            This is for internal use only and may change without notice.
 
         Returns
         -------
         PhysicsAiTrainingJob
             The submitted training job object
-        """
 
+        Examples
+        --------
+        >>> from luminarycloud.enum.quantity_type import QuantityType
+        >>> from luminarycloud.physics_ai.architectures import ModelOutputs, InputParameter
+        >>>
+        >>> job = arch_version.train(
+        ...     dataset_id="dataset-abc123",
+        ...     outputs=ModelOutputs(
+        ...         surface=[QuantityType.PRESSURE_TIME_AVERAGE, QuantityType.WALL_SHEAR_STRESS_TIME_AVERAGE],
+        ...         volume=[QuantityType.VELOCITY_TIME_AVERAGE],
+        ...     ),
+        ...     inputs=[
+        ...         InputParameter(name="stream_velocity", reference=[148.28]),
+        ...         InputParameter(name="air_density", reference=[0.38]),
+        ...     ],
+        ... )
+        """
+        # Build typed proto fields for outputs
+        surface_outputs: list[quantitypb.QuantityType] = []
+        volume_outputs: list[quantitypb.QuantityType] = []
+
+        if outputs:
+            if outputs.surface:
+                surface_outputs = [quantitypb.QuantityType.Value(q.name) for q in outputs.surface]
+            if outputs.volume:
+                volume_outputs = [quantitypb.QuantityType.Value(q.name) for q in outputs.volume]
+
+        # Build typed proto fields for input parameters
+        # Type is inherited from the dataset's ParameterSchema at HydraConfig generation time
+        input_parameters: list[physaipb.TrainingInputParameter] = []
+        if inputs:
+            for param in inputs:
+                input_parameters.append(
+                    physaipb.TrainingInputParameter(
+                        name=param.name,
+                        reference=param.reference,
+                    )
+                )
+
+        # Build training_config JSON for internal fields only (resources, priority_class, etc.)
+        config: dict = copy.deepcopy(_internal) if _internal else {}
+
+        # Apply defaults for _internal fields
         if "custom_args" not in config:
             config["custom_args"] = ""
         if "priority_class" not in config:
@@ -66,14 +175,17 @@ class PhysicsAiArchitectureVersion(ProtoWrapperBase):
             config["mode"] = "full"
 
         training_config_json = json.dumps(config, indent=2)
-        external_dataset_uri = f"gs://training-data/architecture-{self.id}"
+
         req = physaipb.SubmitTrainingJobRequest(
             architecture_version_id=self.id,
-            training_description=self.get_training_description(config),
-            external_dataset_uri=external_dataset_uri,
+            training_description=self.get_training_description(_internal),
             training_config=training_config_json,
             initialization_type=physaipb.MODEL_INITIALIZATION_TYPE_RANDOM,
             base_model_version_id="",
+            dataset_id=dataset_id if dataset_id else "",
+            surface_outputs=surface_outputs,
+            volume_outputs=volume_outputs,
+            input_parameters=input_parameters,
         )
 
         response = get_default_client().SubmitTrainingJob(req)
@@ -92,10 +204,10 @@ class PhysicsAiArchitecture(ProtoWrapperBase):
     id: PhysicsAiArchitectureID
     name: str
     description: str
-    versions: List[PhysicsAiArchitectureVersion]
+    versions: list[PhysicsAiArchitectureVersion]
     _proto: physaipb.PhysicsAiArchitecture
 
-    def get_latest_version(self) -> Optional[PhysicsAiArchitectureVersion]:
+    def get_latest_version(self) -> PhysicsAiArchitectureVersion | None:
         """
         Get the latest version of this architecture based on name.
 
@@ -107,10 +219,10 @@ class PhysicsAiArchitecture(ProtoWrapperBase):
         """
         if not self.versions:
             return None
-        return self.versions[0] if self.versions else None
+        return self.versions[0]
 
 
-def list_architectures() -> List[PhysicsAiArchitecture]:
+def list_architectures() -> list[PhysicsAiArchitecture]:
     """
     List available Physics AI architectures for model training.
 

luminarycloud/physics_ai/datasets.py

@@ -1,20 +1,57 @@
 # Copyright 2025 Luminary Cloud, Inc. All Rights Reserved.
 from dataclasses import dataclass, field
-from typing import List, Optional, Union, TYPE_CHECKING
 from datetime import datetime
+from typing import TYPE_CHECKING, List, Optional, Union
 
-from google.protobuf.struct_pb2 import Struct
 from google.protobuf.json_format import ParseDict
+from google.protobuf.struct_pb2 import Struct
 
-from ..enum import QuantityType
 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:
     """
@@ -107,7 +144,7 @@ def create_dataset(
     cases: List[Union["Solution", DatasetCaseInput]],
     export_config: Optional[ExportConfig] = None,
     description: str = "",
-    parameter_schema: Optional[dict] = None,
+    parameter_schema: Optional[List[ParameterDefinition]] = None,
 ) -> PhysicsAiDataset:
     """
     Create a Physics AI dataset from simulation solutions.
@@ -126,8 +163,9 @@ def create_dataset(
         Export configuration. If not provided, uses defaults (single_precision=True).
     description : str, optional
         Description of the dataset
-    parameter_schema : dict, optional
-        JSON schema for case parameters
+    parameter_schema : List[ParameterDefinition], optional
+        List of parameter definitions for this dataset. Each parameter defines
+        the type, dimension, and optional default reference value.
 
     Returns
     -------
@@ -141,6 +179,9 @@ def create_dataset(
     >>> dataset = create_dataset(
     ...     name="my-dataset",
     ...     cases=[solution1, solution2, solution3],
+    ...     parameter_schema=[
+    ...         ParameterDefinition(name="stream_velocity", type="vector", dimension=1),
+    ...     ],
     ... )
 
     With custom export config:
@@ -152,6 +193,10 @@ def create_dataset(
     ...         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:
@@ -163,6 +208,9 @@ def create_dataset(
     ...         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
@@ -201,12 +249,19 @@ def create_dataset(
     if export_config.volume_fields_to_keep:
         export_config_proto.volume_fields_to_keep.extend(export_config.volume_fields_to_keep)
 
-    # Build parameter schema
+    # Build parameter schema from list of ParameterDefinition
+    param_schema_proto = physaipb.ParameterSchema()
     if parameter_schema:
-        param_schema_proto = ParseDict(parameter_schema, Struct())
-    else:
-        # Default empty schema
-        param_schema_proto = Struct()
+        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,

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]