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

luminarycloud/params/simulation/physics/heat_.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.heat.adjoint_controls_heat_ import AdjointControlsHeat
 from luminarycloud.params.simulation.physics.heat.adjoint_controls_heat_ import *

luminarycloud/params/simulation/physics/periodic_pair_.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)
@@ -28,6 +28,8 @@ class PeriodicPair(CodeRepr, ParamGroupWrapper[clientpb.PeriodicPair]):
     ""
     surfaces_side_b: list[str] = field(default_factory=list)
     ""
+    automatic_deduction: bool = False
+    "Deduce the sign and magnitude of the rotation or translation automatically. When this is enabled the translation and rotation vectors only provide the direction or axis. Note that for rotational periodicity this option works best if the full circle is evenly divisible by the angle. The center of rotation is not deduced."
     translation: Vector3 = field(default_factory=lambda: Vector3(0.0, 0.0, 0.0))
     "Holds the translation vector in the x,y,z directions from each point on the first periodic boundary to its matching point on the second periodic boundary."
     center: Vector3 = field(default_factory=lambda: Vector3(0.0, 0.0, 0.0))
@@ -43,6 +45,7 @@ class PeriodicPair(CodeRepr, ParamGroupWrapper[clientpb.PeriodicPair]):
             _proto.bound_a.extend(self.surfaces_side_a)
         if self.surfaces_side_b is not None:
             _proto.bound_b.extend(self.surfaces_side_b)
+        _proto.periodic_automatic_deduction = self.automatic_deduction
         _proto.periodic_translation.CopyFrom(self.translation._to_ad_proto())
         _proto.periodic_center_of_rotation.CopyFrom(self.center._to_ad_proto())
         _proto.periodic_rotation_angles.CopyFrom(self.rotation_vector._to_ad_proto())
@@ -52,6 +55,7 @@ class PeriodicPair(CodeRepr, ParamGroupWrapper[clientpb.PeriodicPair]):
         self.name = proto.periodic_pair_name
         self.surfaces_side_a.extend(proto.bound_a)
         self.surfaces_side_b.extend(proto.bound_b)
+        self.automatic_deduction = proto.periodic_automatic_deduction
         self.translation._from_ad_proto(proto.periodic_translation)
         self.center._from_ad_proto(proto.periodic_center_of_rotation)
         self.rotation_vector._from_ad_proto(proto.periodic_rotation_angles)

luminarycloud/params/simulation/physics/solution_controls/linear_solver_type/gauss_seidel_.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/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,