diff --git a/.changeset/optimization-constraints.md b/.changeset/optimization-constraints.md new file mode 100644 index 00000000000..927b68fdf6e --- /dev/null +++ b/.changeset/optimization-constraints.md @@ -0,0 +1,6 @@ +--- +"@hashintel/petrinaut": patch +"@hashintel/petrinaut-core": patch +--- + +Constraints are a concept of their own: boolean conditions over the parameter space or the simulation state, authored as TypeScript, lowered to serializable HIR, and validated against a runtime schema of the full HIR grammar. Optimization studies carry a list of them, authored in the create-optimization drawer and exposed through the describe protocol, where the Python binding reads them as callables with a boolean, a signed margin, a pydantic validator, and a SymPy view. Declarative only: nothing enforces them yet. diff --git a/apps/petrinaut-opt/uv.lock b/apps/petrinaut-opt/uv.lock index e0f898f09e6..a539590b85b 100644 --- a/apps/petrinaut-opt/uv.lock +++ b/apps/petrinaut-opt/uv.lock @@ -1061,7 +1061,11 @@ dependencies = [ ] [package.metadata] -requires-dist = [{ name = "pydantic", specifier = ">=2.13.4" }] +requires-dist = [ + { name = "pydantic", specifier = ">=2.13.4" }, + { name = "sympy", marker = "extra == 'sympy'", specifier = ">=1.13" }, +] +provides-extras = ["sympy"] [package.metadata.requires-dev] dev = [ @@ -1069,6 +1073,7 @@ dev = [ { name = "datamodel-code-generator", specifier = ">=0.74.0" }, { name = "pytest", specifier = ">=8.3" }, { name = "ruff", specifier = ">=0.16.4" }, + { name = "sympy", specifier = ">=1.13" }, ] [[package]] diff --git a/libs/@hashintel/petrinaut-cli/schemas/optimization-protocol.schema.json b/libs/@hashintel/petrinaut-cli/schemas/optimization-protocol.schema.json index 209d58a10cb..b0baee1b2a9 100644 --- a/libs/@hashintel/petrinaut-cli/schemas/optimization-protocol.schema.json +++ b/libs/@hashintel/petrinaut-cli/schemas/optimization-protocol.schema.json @@ -165,11 +165,18 @@ "items": { "$ref": "#/$defs/OptimizationDescribeParameter" } + }, + "constraints": { + "description": "The manifest's constraints, passed through verbatim so protocol clients (the Python binding) can evaluate their HIR. Absent means unconstrained.", + "type": "array", + "items": { + "$ref": "#/$defs/Constraint" + } } }, "required": ["direction", "study", "parameters"], "additionalProperties": false, - "description": "The `optimization.describe` result: direction, study settings, and the parameters that are not fixed." + "description": "The `optimization.describe` result: direction, study settings, the parameters that are not fixed, and the study's constraints." }, "OptimizationEvaluateResult": { "type": "object", @@ -187,6 +194,1067 @@ "required": ["objective"], "additionalProperties": false, "description": "The `optimization.evaluate` result. `objective` is the mean of the per-seed objectives (identical to the sole run's objective when the trial runs one seed); `replicates` reports the per-seed values whenever a trial runs more than one." + }, + "Constraint": { + "oneOf": [ + { + "$ref": "#/$defs/ParameterConstraint" + }, + { + "$ref": "#/$defs/StateConstraint" + } + ], + "description": "One boolean condition, discriminated by the space it ranges over.", + "discriminator": { + "propertyName": "space" + } + }, + "ParameterConstraint": { + "type": "object", + "properties": { + "space": { + "type": "string", + "const": "parameters" + }, + "id": { + "type": "string", + "minLength": 1 + }, + "name": { + "description": "Optional display name shown wherever the constraint is reported.", + "type": "string", + "minLength": 1 + }, + "code": { + "type": "string", + "minLength": 1, + "description": "The authored TypeScript source, the editable text of record. `hir` is its lowered form; regenerating `hir` from `code` must be a no-op." + }, + "hir": { + "$ref": "#/$defs/ParameterConstraintHir" + } + }, + "required": ["space", "id", "code", "hir"], + "additionalProperties": false, + "description": "One boolean condition over the parameter space: an expression over `scenario.*` and `parameters.*`, e.g. `scenario.min_load < scenario.max_load`. Checkable before a run starts." + }, + "ParameterConstraintHir": { + "type": "object", + "properties": { + "hirVersion": { + "type": "number", + "const": 1 + }, + "surface": { + "type": "string", + "const": "scenario-expression" + }, + "params": { + "type": "array", + "items": { + "$ref": "#/$defs/HirNamedSpan" + } + }, + "body": { + "$ref": "#/$defs/HirExpr" + }, + "span": { + "$ref": "#/$defs/HirSpan" + } + }, + "required": ["hirVersion", "surface", "params", "body", "span"], + "additionalProperties": false, + "description": "The lowered condition: a `scenario-expression` surface function with no declared parameters; `scenario.*` and `parameters.*` are ambient reads." + }, + "HirNamedSpan": { + "type": "object", + "properties": { + "name": { + "type": "string" + }, + "span": { + "$ref": "#/$defs/HirSpan" + } + }, + "required": ["name", "span"], + "additionalProperties": false, + "description": "A declared name (a parameter or a binding) and where it is spelled." + }, + "HirSpan": { + "type": "object", + "properties": { + "start": { + "type": "integer", + "minimum": 0, + "maximum": 9007199254740991 + }, + "length": { + "type": "integer", + "minimum": 0, + "maximum": 9007199254740991 + } + }, + "required": ["start", "length"], + "additionalProperties": false, + "description": "Half-open span into the user-visible source text, in UTF-16 code units." + }, + "HirExpr": { + "description": "One HIR expression node, discriminated by `kind`. Evaluators must reject kinds they do not know.", + "oneOf": [ + { + "$ref": "#/$defs/HirNumberLit" + }, + { + "$ref": "#/$defs/HirBoolLit" + }, + { + "$ref": "#/$defs/HirStringLit" + }, + { + "$ref": "#/$defs/HirStringCall" + }, + { + "$ref": "#/$defs/HirUuidGenerate" + }, + { + "$ref": "#/$defs/HirUuidFrom" + }, + { + "$ref": "#/$defs/HirConstant" + }, + { + "$ref": "#/$defs/HirLocalRef" + }, + { + "$ref": "#/$defs/HirParamRef" + }, + { + "$ref": "#/$defs/HirScenarioRef" + }, + { + "$ref": "#/$defs/HirRangeCall" + }, + { + "$ref": "#/$defs/HirFieldAccess" + }, + { + "$ref": "#/$defs/HirIndexAccess" + }, + { + "$ref": "#/$defs/HirLength" + }, + { + "$ref": "#/$defs/HirUnary" + }, + { + "$ref": "#/$defs/HirBinary" + }, + { + "$ref": "#/$defs/HirCond" + }, + { + "$ref": "#/$defs/HirLet" + }, + { + "$ref": "#/$defs/HirMathCall" + }, + { + "$ref": "#/$defs/HirRecordLit" + }, + { + "$ref": "#/$defs/HirArrayLit" + }, + { + "$ref": "#/$defs/HirArrayMap" + }, + { + "$ref": "#/$defs/HirArrayReduce" + }, + { + "$ref": "#/$defs/HirArrayConcat" + }, + { + "$ref": "#/$defs/HirDistribution" + }, + { + "$ref": "#/$defs/HirDistributionMap" + } + ], + "discriminator": { + "propertyName": "kind" + } + }, + "HirNumberLit": { + "type": "object", + "properties": { + "id": { + "type": "integer", + "minimum": 0, + "maximum": 9007199254740991 + }, + "span": { + "$ref": "#/$defs/HirSpan" + }, + "kind": { + "type": "string", + "const": "numberLit" + }, + "value": { + "type": "number" + }, + "raw": { + "type": "string" + } + }, + "required": ["id", "span", "kind", "value", "raw"], + "additionalProperties": false + }, + "HirBoolLit": { + "type": "object", + "properties": { + "id": { + "type": "integer", + "minimum": 0, + "maximum": 9007199254740991 + }, + "span": { + "$ref": "#/$defs/HirSpan" + }, + "kind": { + "type": "string", + "const": "boolLit" + }, + "value": { + "type": "boolean" + } + }, + "required": ["id", "span", "kind", "value"], + "additionalProperties": false + }, + "HirStringLit": { + "type": "object", + "properties": { + "id": { + "type": "integer", + "minimum": 0, + "maximum": 9007199254740991 + }, + "span": { + "$ref": "#/$defs/HirSpan" + }, + "kind": { + "type": "string", + "const": "stringLit" + }, + "value": { + "type": "string" + } + }, + "required": ["id", "span", "kind", "value"], + "additionalProperties": false + }, + "HirStringCall": { + "type": "object", + "properties": { + "id": { + "type": "integer", + "minimum": 0, + "maximum": 9007199254740991 + }, + "span": { + "$ref": "#/$defs/HirSpan" + }, + "kind": { + "type": "string", + "const": "stringCall" + }, + "fn": { + "$ref": "#/$defs/HirStringFn" + }, + "target": { + "$ref": "#/$defs/HirExpr" + }, + "argument": { + "$ref": "#/$defs/HirExpr" + } + }, + "required": ["id", "span", "kind", "fn", "target", "argument"], + "additionalProperties": false + }, + "HirStringFn": { + "type": "string", + "enum": ["startsWith", "endsWith", "includes"] + }, + "HirUuidGenerate": { + "type": "object", + "properties": { + "id": { + "type": "integer", + "minimum": 0, + "maximum": 9007199254740991 + }, + "span": { + "$ref": "#/$defs/HirSpan" + }, + "kind": { + "type": "string", + "const": "uuidGenerate" + } + }, + "required": ["id", "span", "kind"], + "additionalProperties": false + }, + "HirUuidFrom": { + "type": "object", + "properties": { + "id": { + "type": "integer", + "minimum": 0, + "maximum": 9007199254740991 + }, + "span": { + "$ref": "#/$defs/HirSpan" + }, + "kind": { + "type": "string", + "const": "uuidFrom" + }, + "operand": { + "$ref": "#/$defs/HirExpr" + } + }, + "required": ["id", "span", "kind", "operand"], + "additionalProperties": false + }, + "HirConstant": { + "type": "object", + "properties": { + "id": { + "type": "integer", + "minimum": 0, + "maximum": 9007199254740991 + }, + "span": { + "$ref": "#/$defs/HirSpan" + }, + "kind": { + "type": "string", + "const": "constant" + }, + "name": { + "$ref": "#/$defs/HirConstantName" + } + }, + "required": ["id", "span", "kind", "name"], + "additionalProperties": false + }, + "HirConstantName": { + "type": "string", + "enum": ["PI", "E", "Infinity", "NaN"] + }, + "HirLocalRef": { + "type": "object", + "properties": { + "id": { + "type": "integer", + "minimum": 0, + "maximum": 9007199254740991 + }, + "span": { + "$ref": "#/$defs/HirSpan" + }, + "kind": { + "type": "string", + "const": "localRef" + }, + "name": { + "type": "string" + } + }, + "required": ["id", "span", "kind", "name"], + "additionalProperties": false + }, + "HirParamRef": { + "type": "object", + "properties": { + "id": { + "type": "integer", + "minimum": 0, + "maximum": 9007199254740991 + }, + "span": { + "$ref": "#/$defs/HirSpan" + }, + "kind": { + "type": "string", + "const": "paramRef" + }, + "name": { + "type": "string" + } + }, + "required": ["id", "span", "kind", "name"], + "additionalProperties": false + }, + "HirScenarioRef": { + "type": "object", + "properties": { + "id": { + "type": "integer", + "minimum": 0, + "maximum": 9007199254740991 + }, + "span": { + "$ref": "#/$defs/HirSpan" + }, + "kind": { + "type": "string", + "const": "scenarioRef" + }, + "name": { + "type": "string" + } + }, + "required": ["id", "span", "kind", "name"], + "additionalProperties": false + }, + "HirRangeCall": { + "type": "object", + "properties": { + "id": { + "type": "integer", + "minimum": 0, + "maximum": 9007199254740991 + }, + "span": { + "$ref": "#/$defs/HirSpan" + }, + "kind": { + "type": "string", + "const": "rangeCall" + }, + "args": { + "type": "array", + "items": { + "$ref": "#/$defs/HirExpr" + } + } + }, + "required": ["id", "span", "kind", "args"], + "additionalProperties": false + }, + "HirFieldAccess": { + "type": "object", + "properties": { + "id": { + "type": "integer", + "minimum": 0, + "maximum": 9007199254740991 + }, + "span": { + "$ref": "#/$defs/HirSpan" + }, + "kind": { + "type": "string", + "const": "fieldAccess" + }, + "target": { + "$ref": "#/$defs/HirExpr" + }, + "field": { + "type": "string" + }, + "fieldSpan": { + "$ref": "#/$defs/HirSpan" + } + }, + "required": ["id", "span", "kind", "target", "field", "fieldSpan"], + "additionalProperties": false + }, + "HirIndexAccess": { + "type": "object", + "properties": { + "id": { + "type": "integer", + "minimum": 0, + "maximum": 9007199254740991 + }, + "span": { + "$ref": "#/$defs/HirSpan" + }, + "kind": { + "type": "string", + "const": "indexAccess" + }, + "target": { + "$ref": "#/$defs/HirExpr" + }, + "index": { + "$ref": "#/$defs/HirExpr" + } + }, + "required": ["id", "span", "kind", "target", "index"], + "additionalProperties": false + }, + "HirLength": { + "type": "object", + "properties": { + "id": { + "type": "integer", + "minimum": 0, + "maximum": 9007199254740991 + }, + "span": { + "$ref": "#/$defs/HirSpan" + }, + "kind": { + "type": "string", + "const": "length" + }, + "target": { + "$ref": "#/$defs/HirExpr" + } + }, + "required": ["id", "span", "kind", "target"], + "additionalProperties": false + }, + "HirUnary": { + "type": "object", + "properties": { + "id": { + "type": "integer", + "minimum": 0, + "maximum": 9007199254740991 + }, + "span": { + "$ref": "#/$defs/HirSpan" + }, + "kind": { + "type": "string", + "const": "unary" + }, + "op": { + "$ref": "#/$defs/HirUnaryOp" + }, + "operand": { + "$ref": "#/$defs/HirExpr" + } + }, + "required": ["id", "span", "kind", "op", "operand"], + "additionalProperties": false + }, + "HirUnaryOp": { + "type": "string", + "enum": ["-", "+", "!"] + }, + "HirBinary": { + "type": "object", + "properties": { + "id": { + "type": "integer", + "minimum": 0, + "maximum": 9007199254740991 + }, + "span": { + "$ref": "#/$defs/HirSpan" + }, + "kind": { + "type": "string", + "const": "binary" + }, + "op": { + "$ref": "#/$defs/HirBinaryOp" + }, + "left": { + "$ref": "#/$defs/HirExpr" + }, + "right": { + "$ref": "#/$defs/HirExpr" + } + }, + "required": ["id", "span", "kind", "op", "left", "right"], + "additionalProperties": false + }, + "HirBinaryOp": { + "type": "string", + "enum": [ + "+", + "-", + "*", + "/", + "%", + "**", + "<", + "<=", + ">", + ">=", + "==", + "!=", + "&&", + "||" + ] + }, + "HirCond": { + "type": "object", + "properties": { + "id": { + "type": "integer", + "minimum": 0, + "maximum": 9007199254740991 + }, + "span": { + "$ref": "#/$defs/HirSpan" + }, + "kind": { + "type": "string", + "const": "cond" + }, + "condition": { + "$ref": "#/$defs/HirExpr" + }, + "thenBranch": { + "$ref": "#/$defs/HirExpr" + }, + "elseBranch": { + "$ref": "#/$defs/HirExpr" + } + }, + "required": [ + "id", + "span", + "kind", + "condition", + "thenBranch", + "elseBranch" + ], + "additionalProperties": false + }, + "HirLet": { + "type": "object", + "properties": { + "id": { + "type": "integer", + "minimum": 0, + "maximum": 9007199254740991 + }, + "span": { + "$ref": "#/$defs/HirSpan" + }, + "kind": { + "type": "string", + "const": "let" + }, + "bindings": { + "type": "array", + "items": { + "$ref": "#/$defs/HirLetBinding" + } + }, + "body": { + "$ref": "#/$defs/HirExpr" + } + }, + "required": ["id", "span", "kind", "bindings", "body"], + "additionalProperties": false + }, + "HirLetBinding": { + "type": "object", + "properties": { + "name": { + "type": "string" + }, + "nameSpan": { + "$ref": "#/$defs/HirSpan" + }, + "value": { + "$ref": "#/$defs/HirExpr" + } + }, + "required": ["name", "nameSpan", "value"], + "additionalProperties": false + }, + "HirMathCall": { + "type": "object", + "properties": { + "id": { + "type": "integer", + "minimum": 0, + "maximum": 9007199254740991 + }, + "span": { + "$ref": "#/$defs/HirSpan" + }, + "kind": { + "type": "string", + "const": "mathCall" + }, + "fn": { + "$ref": "#/$defs/HirMathFn" + }, + "args": { + "type": "array", + "items": { + "$ref": "#/$defs/HirExpr" + } + } + }, + "required": ["id", "span", "kind", "fn", "args"], + "additionalProperties": false + }, + "HirMathFn": { + "type": "string", + "enum": [ + "abs", + "acos", + "asin", + "atan", + "atan2", + "cbrt", + "ceil", + "cos", + "cosh", + "exp", + "floor", + "hypot", + "log", + "log10", + "log2", + "max", + "min", + "pow", + "random", + "round", + "sign", + "sin", + "sinh", + "sqrt", + "tan", + "tanh", + "trunc" + ] + }, + "HirRecordLit": { + "type": "object", + "properties": { + "id": { + "type": "integer", + "minimum": 0, + "maximum": 9007199254740991 + }, + "span": { + "$ref": "#/$defs/HirSpan" + }, + "kind": { + "type": "string", + "const": "recordLit" + }, + "entries": { + "type": "array", + "items": { + "$ref": "#/$defs/HirRecordEntry" + } + } + }, + "required": ["id", "span", "kind", "entries"], + "additionalProperties": false + }, + "HirRecordEntry": { + "type": "object", + "properties": { + "key": { + "type": "string" + }, + "keySpan": { + "$ref": "#/$defs/HirSpan" + }, + "value": { + "$ref": "#/$defs/HirExpr" + } + }, + "required": ["key", "keySpan", "value"], + "additionalProperties": false + }, + "HirArrayLit": { + "type": "object", + "properties": { + "id": { + "type": "integer", + "minimum": 0, + "maximum": 9007199254740991 + }, + "span": { + "$ref": "#/$defs/HirSpan" + }, + "kind": { + "type": "string", + "const": "arrayLit" + }, + "elements": { + "type": "array", + "items": { + "$ref": "#/$defs/HirExpr" + } + } + }, + "required": ["id", "span", "kind", "elements"], + "additionalProperties": false + }, + "HirArrayMap": { + "type": "object", + "properties": { + "id": { + "type": "integer", + "minimum": 0, + "maximum": 9007199254740991 + }, + "span": { + "$ref": "#/$defs/HirSpan" + }, + "kind": { + "type": "string", + "const": "arrayMap" + }, + "target": { + "$ref": "#/$defs/HirExpr" + }, + "param": { + "$ref": "#/$defs/HirNamedSpan" + }, + "indexParam": { + "$ref": "#/$defs/HirNamedSpan" + }, + "body": { + "$ref": "#/$defs/HirExpr" + } + }, + "required": ["id", "span", "kind", "target", "param", "body"], + "additionalProperties": false + }, + "HirArrayReduce": { + "type": "object", + "properties": { + "id": { + "type": "integer", + "minimum": 0, + "maximum": 9007199254740991 + }, + "span": { + "$ref": "#/$defs/HirSpan" + }, + "kind": { + "type": "string", + "const": "arrayReduce" + }, + "target": { + "$ref": "#/$defs/HirExpr" + }, + "accParam": { + "$ref": "#/$defs/HirNamedSpan" + }, + "param": { + "$ref": "#/$defs/HirNamedSpan" + }, + "indexParam": { + "$ref": "#/$defs/HirNamedSpan" + }, + "body": { + "$ref": "#/$defs/HirExpr" + }, + "initial": { + "$ref": "#/$defs/HirExpr" + } + }, + "required": [ + "id", + "span", + "kind", + "target", + "accParam", + "param", + "body", + "initial" + ], + "additionalProperties": false + }, + "HirArrayConcat": { + "type": "object", + "properties": { + "id": { + "type": "integer", + "minimum": 0, + "maximum": 9007199254740991 + }, + "span": { + "$ref": "#/$defs/HirSpan" + }, + "kind": { + "type": "string", + "const": "arrayConcat" + }, + "left": { + "$ref": "#/$defs/HirExpr" + }, + "right": { + "$ref": "#/$defs/HirExpr" + } + }, + "required": ["id", "span", "kind", "left", "right"], + "additionalProperties": false + }, + "HirDistribution": { + "type": "object", + "properties": { + "id": { + "type": "integer", + "minimum": 0, + "maximum": 9007199254740991 + }, + "span": { + "$ref": "#/$defs/HirSpan" + }, + "kind": { + "type": "string", + "const": "distribution" + }, + "dist": { + "$ref": "#/$defs/HirDistributionKind" + }, + "args": { + "type": "array", + "items": { + "$ref": "#/$defs/HirExpr" + } + } + }, + "required": ["id", "span", "kind", "dist", "args"], + "additionalProperties": false + }, + "HirDistributionKind": { + "type": "string", + "enum": ["gaussian", "uniform", "lognormal"] + }, + "HirDistributionMap": { + "type": "object", + "properties": { + "id": { + "type": "integer", + "minimum": 0, + "maximum": 9007199254740991 + }, + "span": { + "$ref": "#/$defs/HirSpan" + }, + "kind": { + "type": "string", + "const": "distributionMap" + }, + "base": { + "$ref": "#/$defs/HirExpr" + }, + "param": { + "$ref": "#/$defs/HirNamedSpan" + }, + "body": { + "$ref": "#/$defs/HirExpr" + } + }, + "required": ["id", "span", "kind", "base", "param", "body"], + "additionalProperties": false + }, + "StateConstraint": { + "type": "object", + "properties": { + "space": { + "type": "string", + "const": "state" + }, + "id": { + "type": "string", + "minLength": 1 + }, + "name": { + "description": "Optional display name shown wherever the constraint is reported.", + "type": "string", + "minLength": 1 + }, + "code": { + "type": "string", + "minLength": 1, + "description": "The authored TypeScript source, the editable text of record. `hir` is its lowered form; regenerating `hir` from `code` must be a no-op." + }, + "hir": { + "$ref": "#/$defs/StateConstraintHir" + } + }, + "required": ["space", "id", "code", "hir"], + "additionalProperties": false, + "description": "One boolean condition over the simulation state, authored like a metric body but returning boolean, e.g. `return state.places.Queue.count <= 10;`. Observed while a run goes." + }, + "StateConstraintHir": { + "type": "object", + "properties": { + "hirVersion": { + "type": "number", + "const": 1 + }, + "surface": { + "type": "string", + "const": "metric" + }, + "params": { + "type": "array", + "items": { + "$ref": "#/$defs/HirNamedSpan" + } + }, + "body": { + "$ref": "#/$defs/HirExpr" + }, + "span": { + "$ref": "#/$defs/HirSpan" + } + }, + "required": ["hirVersion", "surface", "params", "body", "span"], + "additionalProperties": false, + "description": "The lowered condition: a `metric` surface function whose first declared parameter is the simulation `state`; `parameters.*` is ambient." + }, + "HirSurfaceKind": { + "type": "string", + "enum": [ + "dynamics", + "lambda", + "kernel", + "metric", + "scenario-expression", + "scenario-code" + ] + }, + "HirFunction": { + "type": "object", + "properties": { + "hirVersion": { + "type": "number", + "const": 1 + }, + "surface": { + "$ref": "#/$defs/HirSurfaceKind" + }, + "params": { + "type": "array", + "items": { + "$ref": "#/$defs/HirNamedSpan" + } + }, + "body": { + "$ref": "#/$defs/HirExpr" + }, + "span": { + "$ref": "#/$defs/HirSpan" + } + }, + "required": ["hirVersion", "surface", "params", "body", "span"], + "additionalProperties": false, + "description": "A lowered user function (see hir/hir.ts for the grammar). `params[0]` is the input parameter when the surface declares one; `parameters.*` and `scenario.*` reads are dedicated node kinds, never locals." } } } diff --git a/libs/@hashintel/petrinaut-cli/scripts/generate-protocol-schemas.ts b/libs/@hashintel/petrinaut-cli/scripts/generate-protocol-schemas.ts index f4936b6dee1..bcdf6777ffc 100644 --- a/libs/@hashintel/petrinaut-cli/scripts/generate-protocol-schemas.ts +++ b/libs/@hashintel/petrinaut-cli/scripts/generate-protocol-schemas.ts @@ -4,7 +4,8 @@ * * The document is checked in and regenerated by `codegen`, so CI fails when it * drifts from the Zod schemas in `@hashintel/petrinaut-core` that define the - * protocol. `@local/petrinaut-python` generates its pydantic models from it. + * protocol. `@local/petrinaut-python` generates its pydantic models from it, + * including the full HIR grammar the constraints carry. */ import { mkdir, writeFile } from "node:fs/promises"; @@ -14,26 +15,68 @@ import { fileURLToPath } from "node:url"; import { z } from "zod"; import { + constraintSchema, + hirFunctionSchema, + parameterConstraintSchema, petrinautOptimizationDescribeParameterSchema, petrinautOptimizationDescribeResultSchema, petrinautOptimizationEvaluateResultSchema, petrinautOptimizationReplicateSchema, + stateConstraintSchema, } from "@hashintel/petrinaut-core"; import type { ZodType } from "zod"; const packageRoot = fileURLToPath(new URL("..", import.meta.url)); +type JsonObject = Record; + +/** + * Definitions every conversion contributes to. Schemas carrying a `meta.id` + * (the HIR grammar, the constraint shapes) come out as `$ref`s into the + * converted schema's own `$defs`; those are hoisted here so every reference + * resolves against the document root, and one shared model is generated per + * definition. + */ +const sharedDefs: Record = {}; + +const hoistDefs = (json: JsonObject): void => { + const defs = json.$defs as Record | undefined; + delete json.$defs; + for (const [name, definition] of Object.entries(defs ?? {})) { + const existing = sharedDefs[name]; + if ( + existing !== undefined && + JSON.stringify(existing) !== JSON.stringify(definition) + ) { + throw new Error(`Two different schemas share the definition id ${name}`); + } + sharedDefs[name] = definition; + } +}; + /** `io: "output"` describes the serialized result the CLI writes to the wire. */ -const convert = (schema: ZodType): Record => { - const json = z.toJSONSchema(schema, { io: "output" }) as Record< - string, - unknown - >; +const convert = (schema: ZodType): JsonObject => { + const json = z.toJSONSchema(schema, { io: "output" }) as JsonObject; delete json.$schema; + hoistDefs(json); return json; }; +/** + * JSON Schema has no discriminator keyword; model generators read the + * OpenAPI one, and with it the generated union dispatches on the tag + * instead of trying every member. + */ +const discriminate = (definition: JsonObject, propertyName: string): void => { + if (!Array.isArray(definition.oneOf)) { + throw new Error( + `Expected a oneOf union to discriminate on ${propertyName}`, + ); + } + definition.discriminator = { propertyName }; +}; + const parameterUnion = petrinautOptimizationDescribeParameterSchema; const [floatParameter, intParameter, booleanParameter] = parameterUnion.options; @@ -41,16 +84,41 @@ const describeResult = convert(petrinautOptimizationDescribeResultSchema); // The parameter union is a named definition, so generators emit one shared // model instead of an anonymous copy inlined at each use. ( - (describeResult.properties as Record>) - .parameters as { items: unknown } + (describeResult.properties as Record).parameters as { + items: unknown; + } ).items = { $ref: "#/$defs/OptimizationDescribeParameter" }; const evaluateResult = convert(petrinautOptimizationEvaluateResultSchema); ( - (evaluateResult.properties as Record>) - .replicates as { items?: unknown } + (evaluateResult.properties as Record).replicates as { + items?: unknown; + } ).items = { $ref: "#/$defs/OptimizationReplicate" }; +// The constraint shapes and the HIR grammar arrive through the describe +// result's `constraints`; converting them on their own as well pins their +// definitions even when a future describe result stops carrying them. +convert(parameterConstraintSchema); +convert(stateConstraintSchema); +convert(constraintSchema); +// A root conversion inlines the schema instead of naming it, so the generic +// function shape is registered by hand next to its two pinned variants. +sharedDefs.HirFunction = convert(hirFunctionSchema); + +for (const [name, propertyName] of [ + ["Constraint", "space"], + ["HirExpr", "kind"], +] as const) { + const definition = sharedDefs[name]; + if (definition === undefined) { + throw new Error( + `The converted schemas did not produce a ${name} definition`, + ); + } + discriminate(definition, propertyName); +} + // No root schema and no title: the document only carries `$defs`, so model // generators emit one class per definition and nothing for the document itself. const document = { @@ -70,6 +138,7 @@ const document = { OptimizationReplicate: convert(petrinautOptimizationReplicateSchema), OptimizationDescribeResult: describeResult, OptimizationEvaluateResult: evaluateResult, + ...sharedDefs, }, }; diff --git a/libs/@hashintel/petrinaut-core/src/constraint/constraint.test.ts b/libs/@hashintel/petrinaut-core/src/constraint/constraint.test.ts new file mode 100644 index 00000000000..9f7533ea4d1 --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/constraint/constraint.test.ts @@ -0,0 +1,100 @@ +import { describe, expect, it } from "vitest"; + +import { lowerTypeScriptToHir } from "../hir/lower-typescript"; +import { + constraintListSchema, + constraintSchema, + constraintsInSpace, +} from "./constraint"; + +import type { Constraint } from "./constraint"; + +const lower = (code: string, surface: "metric" | "scenario-expression") => { + const lowered = lowerTypeScriptToHir(code, surface); + if (!lowered.ok) { + throw new Error(JSON.stringify(lowered.diagnostics)); + } + return JSON.parse(JSON.stringify(lowered.fn)) as unknown; +}; + +const parameterConstraint = { + space: "parameters", + id: "c-order", + name: "Load ordering", + code: "scenario.min_load < scenario.max_load", + hir: lower("scenario.min_load < scenario.max_load", "scenario-expression"), +}; + +const stateConstraint = { + space: "state", + id: "c-queue", + code: "return state.places.Queue.count <= 10;", + hir: lower("return state.places.Queue.count <= 10;", "metric"), +}; + +describe("constraintSchema", () => { + it("accepts each shape on its own surface", () => { + expect(constraintSchema.safeParse(parameterConstraint).success).toBe(true); + expect(constraintSchema.safeParse(stateConstraint).success).toBe(true); + }); + + it("pins the surface to the space", () => { + // A metric-surface function cannot pose as a parameter constraint, nor + // the reverse: the shape itself refuses, no cross-check needed. + const misfiledParameter = { + ...parameterConstraint, + hir: stateConstraint.hir, + }; + const misfiledState = { ...stateConstraint, hir: parameterConstraint.hir }; + expect(constraintSchema.safeParse(misfiledParameter).success).toBe(false); + expect(constraintSchema.safeParse(misfiledState).success).toBe(false); + }); + + it("rejects an unknown space", () => { + expect( + constraintSchema.safeParse({ ...parameterConstraint, space: "time" }) + .success, + ).toBe(false); + }); + + it("rejects an empty id or code", () => { + expect( + constraintSchema.safeParse({ ...parameterConstraint, id: "" }).success, + ).toBe(false); + expect( + constraintSchema.safeParse({ ...parameterConstraint, code: " " }) + .success, + ).toBe(false); + }); +}); + +describe("constraintListSchema", () => { + it("accepts mixed spaces and rejects a duplicate id across them", () => { + expect( + constraintListSchema.safeParse([parameterConstraint, stateConstraint]) + .success, + ).toBe(true); + const duplicated = constraintListSchema.safeParse([ + parameterConstraint, + { ...stateConstraint, id: parameterConstraint.id }, + ]); + expect(duplicated.success).toBe(false); + if (!duplicated.success) { + expect(duplicated.error.issues[0]?.path).toEqual([1, "id"]); + } + }); +}); + +describe("constraintsInSpace", () => { + it("narrows to one space", () => { + const constraints = constraintListSchema.parse([ + parameterConstraint, + stateConstraint, + ]) as Constraint[]; + const parameters = constraintsInSpace(constraints, "parameters"); + expect(parameters.map((constraint) => constraint.id)).toEqual(["c-order"]); + expect(parameters[0]?.hir.surface).toBe("scenario-expression"); + const state = constraintsInSpace(constraints, "state"); + expect(state[0]?.hir.surface).toBe("metric"); + }); +}); diff --git a/libs/@hashintel/petrinaut-core/src/constraint/constraint.ts b/libs/@hashintel/petrinaut-core/src/constraint/constraint.ts new file mode 100644 index 00000000000..8131d5459c5 --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/constraint/constraint.ts @@ -0,0 +1,139 @@ +/** + * Constraints: boolean conditions authored as TypeScript and carried as + * serialized HIR, so every consumer (the editors, the CLI, the Python + * binding) reads one shared expression representation without a TypeScript + * frontend of its own. + * + * Two shapes, discriminated by `space`: + * - a **parameter** constraint ranges over the parameter space: one + * expression over `scenario.*` and `parameters.*`, lowered on the + * `scenario-expression` surface with a boolean expected type. It is + * checkable before anything runs. + * - a **state** constraint ranges over the simulation state: a metric-shaped + * body over `state`, lowered on the `metric` surface and checked to return + * boolean. It is observed while a run goes. + * + * A constraint belongs to whatever carries it (today the optimization + * manifest); this module owns the shape, not the placement. Nothing enforces + * constraints yet: they are declared, validated, and evaluable. + * + * @layerRoot core.constraints + * @role Boolean conditions over the parameter space or the simulation state, authored as TypeScript, carried as HIR, and shaped once for every consumer + */ + +import { z } from "zod"; + +import { hirFunctionSchema } from "../hir/hir-schema"; + +import type { HirSurfaceKind } from "../hir/hir"; + +export const CONSTRAINT_SPACES = ["parameters", "state"] as const; + +export const constraintSpaceSchema = z.enum(CONSTRAINT_SPACES).meta({ + id: "ConstraintSpace", + description: + "What a constraint ranges over: the parameter space (`parameters`) or the simulation state (`state`).", +}); + +export type ConstraintSpace = (typeof CONSTRAINT_SPACES)[number]; + +/** The HIR surface each space lowers on. */ +export const CONSTRAINT_SURFACES = { + parameters: "scenario-expression", + state: "metric", +} as const satisfies Record; + +const constraintBaseShape = { + id: z.string().min(1), + name: z.string().trim().min(1).optional().meta({ + description: + "Optional display name shown wherever the constraint is reported.", + }), + code: z.string().trim().min(1).meta({ + description: + "The authored TypeScript source, the editable text of record. `hir` is its lowered form; regenerating `hir` from `code` must be a no-op.", + }), +}; + +export const parameterConstraintSchema = z + .strictObject({ + space: z.literal("parameters"), + ...constraintBaseShape, + hir: hirFunctionSchema + .extend({ surface: z.literal(CONSTRAINT_SURFACES.parameters) }) + .meta({ + id: "ParameterConstraintHir", + description: + "The lowered condition: a `scenario-expression` surface function with no declared parameters; `scenario.*` and `parameters.*` are ambient reads.", + }), + }) + .meta({ + id: "ParameterConstraint", + description: + "One boolean condition over the parameter space: an expression over `scenario.*` and `parameters.*`, e.g. `scenario.min_load < scenario.max_load`. Checkable before a run starts.", + }); + +export const stateConstraintSchema = z + .strictObject({ + space: z.literal("state"), + ...constraintBaseShape, + hir: hirFunctionSchema + .extend({ surface: z.literal(CONSTRAINT_SURFACES.state) }) + .meta({ + id: "StateConstraintHir", + description: + "The lowered condition: a `metric` surface function whose first declared parameter is the simulation `state`; `parameters.*` is ambient.", + }), + }) + .meta({ + id: "StateConstraint", + description: + "One boolean condition over the simulation state, authored like a metric body but returning boolean, e.g. `return state.places.Queue.count <= 10;`. Observed while a run goes.", + }); + +export const constraintSchema = z + .discriminatedUnion("space", [ + parameterConstraintSchema, + stateConstraintSchema, + ]) + .meta({ + id: "Constraint", + description: + "One boolean condition, discriminated by the space it ranges over.", + }); + +/** A set of constraints with unique ids across both spaces. */ +export const constraintListSchema = z + .array(constraintSchema) + .superRefine((constraints, context) => { + const seen = new Set(); + for (const [index, constraint] of constraints.entries()) { + if (seen.has(constraint.id)) { + context.addIssue({ + code: "custom", + path: [index, "id"], + message: `Duplicate constraint id "${constraint.id}"`, + }); + } + seen.add(constraint.id); + } + }) + .meta({ + description: + "Boolean conditions over the parameter space and the simulation state, ids unique across the list.", + }); + +export type ParameterConstraint = z.infer; +export type StateConstraint = z.infer; +export type Constraint = z.infer; + +/** The constraints of one space, typed by that space. */ +export function constraintsInSpace( + constraints: readonly Constraint[], + space: S, +): Extract[] { + return constraints.filter( + (constraint): constraint is Extract => + constraint.space === space, + ); +} diff --git a/libs/@hashintel/petrinaut-core/src/constraint/lower.test.ts b/libs/@hashintel/petrinaut-core/src/constraint/lower.test.ts new file mode 100644 index 00000000000..5f4b64c900a --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/constraint/lower.test.ts @@ -0,0 +1,153 @@ +import { describe, expect, it } from "vitest"; + +import { constraintListSchema } from "./constraint"; +import { lowerConstraint } from "./lower"; + +import type { SDCPN } from "../types/sdcpn"; +import type { LowerConstraintContext } from "./lower"; + +const sdcpn: SDCPN = { + places: [ + { + id: "place-queue", + name: "Queue", + colorId: null, + dynamicsEnabled: false, + differentialEquationId: null, + x: 0, + y: 0, + }, + ], + transitions: [], + types: [], + parameters: [ + { + id: "param-rate", + name: "Rate", + variableName: "rate", + type: "real", + defaultValue: "1.5", + }, + ], + differentialEquations: [], +}; + +const context: LowerConstraintContext = { + netParameters: sdcpn.parameters, + scenarioParameters: [ + { identifier: "min_load", type: "integer", default: 2 }, + { identifier: "max_load", type: "integer", default: 8 }, + ], + sdcpn, +}; + +describe("lowerConstraint", () => { + it("lowers a boolean parameter expression to a parameter constraint", () => { + const result = lowerConstraint( + { + space: "parameters", + id: "c-order", + name: "Load ordering", + code: "scenario.min_load < scenario.max_load && parameters.rate > 0", + }, + context, + ); + expect(result.ok).toBe(true); + if (!result.ok) { + return; + } + expect(result.constraint).toMatchObject({ + space: "parameters", + id: "c-order", + name: "Load ordering", + hir: { surface: "scenario-expression", params: [] }, + }); + }); + + it("omits the name when none was authored", () => { + const result = lowerConstraint( + { space: "parameters", id: "c", code: "scenario.min_load < 5" }, + context, + ); + expect(result.ok && "name" in result.constraint).toBe(false); + }); + + it("rejects a parameter expression that is not boolean", () => { + const result = lowerConstraint( + { space: "parameters", id: "c", code: "scenario.min_load + 1" }, + context, + ); + expect(result.ok).toBe(false); + if (result.ok) { + return; + } + expect(result.diagnostics[0]?.message).toContain("boolean"); + }); + + it("rejects a reference to an unknown scenario parameter", () => { + const result = lowerConstraint( + { space: "parameters", id: "c", code: "scenario.missing > 0" }, + context, + ); + expect(result.ok).toBe(false); + }); + + it("lowers a boolean state condition to a state constraint", () => { + const result = lowerConstraint( + { + space: "state", + id: "c-queue", + code: "return state.places.Queue.count <= 10;", + }, + context, + ); + expect(result.ok).toBe(true); + if (!result.ok) { + return; + } + expect(result.constraint).toMatchObject({ + space: "state", + hir: { surface: "metric", params: [{ name: "state" }] }, + }); + }); + + it("rejects a state condition returning a number", () => { + const result = lowerConstraint( + { space: "state", id: "c", code: "return state.places.Queue.count;" }, + context, + ); + expect(result.ok).toBe(false); + if (result.ok) { + return; + } + expect(result.diagnostics[0]?.message).toContain("boolean"); + }); + + it("produces constraints the list schema accepts verbatim", () => { + const results = [ + lowerConstraint( + { space: "parameters", id: "a", code: "scenario.min_load < 5" }, + context, + ), + lowerConstraint( + { + space: "state", + id: "b", + code: "return state.places.Queue.count > 0;", + }, + context, + ), + ]; + const constraints = results.map((result) => { + if (!result.ok) { + throw new Error(JSON.stringify(result.diagnostics)); + } + return result.constraint; + }); + const parsed = constraintListSchema.safeParse( + JSON.parse(JSON.stringify(constraints)), + ); + expect(parsed.success).toBe(true); + expect(parsed.data).toEqual(JSON.parse(JSON.stringify(constraints))); + }); +}); diff --git a/libs/@hashintel/petrinaut-core/src/constraint/lower.ts b/libs/@hashintel/petrinaut-core/src/constraint/lower.ts new file mode 100644 index 00000000000..4b1ee5f30a8 --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/constraint/lower.ts @@ -0,0 +1,98 @@ +/** + * Lowering of one constraint from its authored TypeScript to a typed + * `Constraint`: the source is lowered on the surface its space dictates, + * typechecked against what that space may read, and checked to produce a + * boolean. + * + * This module (transitively) imports `typescript`, so it stays out of + * browser main bundles: browser callers lower through the language worker + * (`sdcpn/lowerConstraint`), Node callers lower inline. + */ + +import { lowerTypeScriptToHir } from "../hir/lower-typescript"; +import { + buildMetricContext, + buildScenarioExpressionContext, +} from "../hir/surface-context"; +import { typecheckHir } from "../hir/typecheck"; +import { CONSTRAINT_SURFACES } from "./constraint"; + +import type { PetrinautExtensionSettings } from "../extensions"; +import type { HirDiagnostic } from "../hir/hir"; +import type { Parameter, ScenarioParameter, SDCPN } from "../types/sdcpn"; +import type { Constraint, ConstraintSpace } from "./constraint"; + +/** The authored side of a constraint: everything but its lowered form. */ +export type ConstraintSource = { + space: ConstraintSpace; + id: string; + name?: string; + code: string; +}; + +/** What a constraint's condition ranges over, per space. */ +export type LowerConstraintContext = { + /** Net parameters, ambient as `parameters.*` on both surfaces. */ + netParameters: readonly Parameter[]; + /** The scenario parameters (`scenario.*`); parameter space only. */ + scenarioParameters: readonly ScenarioParameter[]; + /** The net the simulation `state` exposes; state space only. */ + sdcpn: SDCPN; + extensions?: PetrinautExtensionSettings; +}; + +export type LowerConstraintResult = + | { ok: true; constraint: Constraint } + | { ok: false; diagnostics: HirDiagnostic[] }; + +/** + * Lowers one constraint's source and checks that it produces a boolean. + * Returns the constraint ready to carry, or the lowering and type + * diagnostics (spans relative to the user's source). + */ +export function lowerConstraint( + source: ConstraintSource, + context: LowerConstraintContext, +): LowerConstraintResult { + const lowered = lowerTypeScriptToHir( + source.code, + CONSTRAINT_SURFACES[source.space], + ); + if (!lowered.ok) { + return { ok: false, diagnostics: lowered.diagnostics }; + } + const surfaceContext = + source.space === "parameters" + ? buildScenarioExpressionContext( + [...context.netParameters], + [...context.scenarioParameters], + "boolean", + ) + : buildMetricContext(context.sdcpn, context.extensions, "boolean"); + const checked = typecheckHir(lowered.fn, surfaceContext); + const errors = checked.diagnostics.filter( + (diagnostic) => diagnostic.severity === "error", + ); + if (errors.length > 0) { + return { ok: false, diagnostics: errors }; + } + + const authored = { + id: source.id, + ...(source.name === undefined ? {} : { name: source.name }), + code: source.code, + }; + const constraint: Constraint = + source.space === "parameters" + ? { + space: "parameters", + ...authored, + hir: { ...lowered.fn, surface: CONSTRAINT_SURFACES.parameters }, + } + : { + space: "state", + ...authored, + hir: { ...lowered.fn, surface: CONSTRAINT_SURFACES.state }, + }; + return { ok: true, constraint }; +} diff --git a/libs/@hashintel/petrinaut-core/src/hir.ts b/libs/@hashintel/petrinaut-core/src/hir.ts index 089e50ff58a..78a594cdbfa 100644 --- a/libs/@hashintel/petrinaut-core/src/hir.ts +++ b/libs/@hashintel/petrinaut-core/src/hir.ts @@ -79,6 +79,18 @@ export { type HirInterpretBindings, type HirValue, } from "./hir/interpret"; +export { + lowerConstraint, + type ConstraintSource, + type LowerConstraintContext, + type LowerConstraintResult, +} from "./constraint/lower"; +export { + hirExprSchema, + hirFunctionSchema, + hirSurfaceKindSchema, + spanSchema, +} from "./hir/hir-schema"; export { lowerTypeScriptToHir, type LowerTypeScriptResult, diff --git a/libs/@hashintel/petrinaut-core/src/hir/hir-schema.test.ts b/libs/@hashintel/petrinaut-core/src/hir/hir-schema.test.ts new file mode 100644 index 00000000000..2d298d43dab --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/hir/hir-schema.test.ts @@ -0,0 +1,72 @@ +import { describe, expect, it } from "vitest"; + +import { hirExprSchema, hirFunctionSchema } from "./hir-schema"; +import { lowerTypeScriptToHir } from "./lower-typescript"; + +const lower = (code: string, surface: "metric" | "scenario-expression") => { + const lowered = lowerTypeScriptToHir(code, surface); + if (!lowered.ok) { + throw new Error(JSON.stringify(lowered.diagnostics)); + } + return lowered.fn; +}; + +describe("hirFunctionSchema", () => { + it("accepts what the lowering produces, node for node", () => { + const samples = [ + lower( + "scenario.min_load < scenario.max_load && parameters.rate > 0", + "scenario-expression", + ), + lower( + "Math.max(1, scenario.n) ** 2 % 3 !== 0 ? -1 : +Math.PI", + "scenario-expression", + ), + lower( + 'range(1, scenario.n).map((x) => x * 2).length > 0 && "a".startsWith("a")', + "scenario-expression", + ), + lower( + `const counts = state.places.Queue.tokens.map((token, index) => token.weight + index); + const total = counts.reduce((acc, value) => acc + value, 0); + return total <= 10 && [1, 2].concat([3])[0] === 1 && ({ a: 1 }).a === 1;`, + "metric", + ), + ]; + for (const fn of samples) { + const parsed = hirFunctionSchema.safeParse( + JSON.parse(JSON.stringify(fn)), + ); + expect(parsed.success, JSON.stringify(parsed.error?.issues)).toBe(true); + // Parsing is the identity on well-formed HIR: nothing is dropped or coerced. + expect(parsed.data).toEqual(JSON.parse(JSON.stringify(fn))); + } + }); + + it("rejects an unknown node kind", () => { + const fn = lower("scenario.n > 0", "scenario-expression"); + const forged = { ...fn, body: { ...fn.body, kind: "eval" } }; + expect(hirFunctionSchema.safeParse(forged).success).toBe(false); + }); + + it("rejects a node missing a field the grammar requires", () => { + const fn = lower("scenario.n > 0", "scenario-expression"); + const body = fn.body as { kind: "binary"; right: unknown }; + const { right: _dropped, ...withoutRight } = body; + expect(hirExprSchema.safeParse(withoutRight).success).toBe(false); + }); + + it("rejects fields the grammar does not declare", () => { + const fn = lower("scenario.n > 0", "scenario-expression"); + expect(hirExprSchema.safeParse({ ...fn.body, extra: 1 }).success).toBe( + false, + ); + }); + + it("rejects a foreign HIR version", () => { + const fn = lower("scenario.n > 0", "scenario-expression"); + expect(hirFunctionSchema.safeParse({ ...fn, hirVersion: 2 }).success).toBe( + false, + ); + }); +}); diff --git a/libs/@hashintel/petrinaut-core/src/hir/hir-schema.ts b/libs/@hashintel/petrinaut-core/src/hir/hir-schema.ts new file mode 100644 index 00000000000..0830d1ecef5 --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/hir/hir-schema.ts @@ -0,0 +1,468 @@ +/** + * Runtime schema for the serialized HIR grammar that `hir.ts` declares as + * types. Anything that carries HIR across a process boundary (an + * optimization manifest, the CLI's describe protocol) validates it here, and + * the CLI publishes this schema as JSON Schema so other languages generate a + * matching, fully typed AST from the one definition. + * + * The schema and the types are kept in lockstep by the `Equals` assertions + * at the bottom of this file: adding a node kind or a field to `hir.ts` + * without mirroring it here fails to compile. + */ + +import { z } from "zod"; + +import { HIR_MATH_FNS, HIR_STRING_FNS } from "./hir"; + +import type { + HirArrayConcat, + HirArrayLit, + HirArrayMap, + HirArrayReduce, + HirBinary, + HirBoolLit, + HirCond, + HirConstant, + HirDistribution, + HirDistributionMap, + HirExpr, + HirFieldAccess, + HirFunction, + HirIndexAccess, + HirLength, + HirLet, + HirLocalRef, + HirMathCall, + HirNumberLit, + HirParamRef, + HirRangeCall, + HirRecordLit, + HirScenarioRef, + HirStringCall, + HirStringLit, + HirUnary, + HirUuidFrom, + HirUuidGenerate, + Span, +} from "./hir"; + +export const spanSchema = z + .strictObject({ + start: z.int().nonnegative(), + length: z.int().nonnegative(), + }) + .meta({ + id: "HirSpan", + description: + "Half-open span into the user-visible source text, in UTF-16 code units.", + }); + +const namedSpanSchema = z + .strictObject({ + name: z.string(), + span: spanSchema, + }) + .meta({ + id: "HirNamedSpan", + description: + "A declared name (a parameter or a binding) and where it is spelled.", + }); + +export const hirStringFnSchema = z + .enum(HIR_STRING_FNS) + .meta({ id: "HirStringFn" }); +export const hirMathFnSchema = z.enum(HIR_MATH_FNS).meta({ id: "HirMathFn" }); +export const hirConstantNameSchema = z + .enum(["PI", "E", "Infinity", "NaN"]) + .meta({ id: "HirConstantName" }); +export const hirUnaryOpSchema = z + .enum(["-", "+", "!"]) + .meta({ id: "HirUnaryOp" }); +export const hirBinaryOpSchema = z + .enum([ + "+", + "-", + "*", + "/", + "%", + "**", + "<", + "<=", + ">", + ">=", + "==", + "!=", + "&&", + "||", + ]) + .meta({ id: "HirBinaryOp" }); +export const hirDistributionKindSchema = z + .enum(["gaussian", "uniform", "lognormal"]) + .meta({ id: "HirDistributionKind" }); + +const nodeBase = { + id: z.int().nonnegative(), + span: spanSchema, +}; + +/** + * The expression grammar. `z.lazy` breaks the recursion; the annotation + * pins the inferred output to `HirExpr` so every node schema below can + * reference it without widening. + */ +export const hirExprSchema: z.ZodType = z + // eslint-disable-next-line no-use-before-define -- recursive grammar: the node schemas below reference this union, and it is built from them once they exist + .lazy(() => z.discriminatedUnion("kind", hirExprOptions)) + .meta({ + id: "HirExpr", + description: + "One HIR expression node, discriminated by `kind`. Evaluators must reject kinds they do not know.", + }); + +const numberLitSchema = z + .strictObject({ + ...nodeBase, + kind: z.literal("numberLit"), + value: z.number(), + raw: z.string(), + }) + .meta({ id: "HirNumberLit" }); + +const boolLitSchema = z + .strictObject({ ...nodeBase, kind: z.literal("boolLit"), value: z.boolean() }) + .meta({ id: "HirBoolLit" }); + +const stringLitSchema = z + .strictObject({ + ...nodeBase, + kind: z.literal("stringLit"), + value: z.string(), + }) + .meta({ id: "HirStringLit" }); + +const stringCallSchema = z + .strictObject({ + ...nodeBase, + kind: z.literal("stringCall"), + fn: hirStringFnSchema, + target: hirExprSchema, + argument: hirExprSchema, + }) + .meta({ id: "HirStringCall" }); + +const uuidGenerateSchema = z + .strictObject({ ...nodeBase, kind: z.literal("uuidGenerate") }) + .meta({ id: "HirUuidGenerate" }); + +const uuidFromSchema = z + .strictObject({ + ...nodeBase, + kind: z.literal("uuidFrom"), + operand: hirExprSchema, + }) + .meta({ id: "HirUuidFrom" }); + +const constantSchema = z + .strictObject({ + ...nodeBase, + kind: z.literal("constant"), + name: hirConstantNameSchema, + }) + .meta({ id: "HirConstant" }); + +const localRefSchema = z + .strictObject({ ...nodeBase, kind: z.literal("localRef"), name: z.string() }) + .meta({ id: "HirLocalRef" }); + +const paramRefSchema = z + .strictObject({ ...nodeBase, kind: z.literal("paramRef"), name: z.string() }) + .meta({ id: "HirParamRef" }); + +const scenarioRefSchema = z + .strictObject({ + ...nodeBase, + kind: z.literal("scenarioRef"), + name: z.string(), + }) + .meta({ id: "HirScenarioRef" }); + +const rangeCallSchema = z + .strictObject({ + ...nodeBase, + kind: z.literal("rangeCall"), + args: z.array(hirExprSchema), + }) + .meta({ id: "HirRangeCall" }); + +const fieldAccessSchema = z + .strictObject({ + ...nodeBase, + kind: z.literal("fieldAccess"), + target: hirExprSchema, + field: z.string(), + fieldSpan: spanSchema, + }) + .meta({ id: "HirFieldAccess" }); + +const indexAccessSchema = z + .strictObject({ + ...nodeBase, + kind: z.literal("indexAccess"), + target: hirExprSchema, + index: hirExprSchema, + }) + .meta({ id: "HirIndexAccess" }); + +const lengthSchema = z + .strictObject({ + ...nodeBase, + kind: z.literal("length"), + target: hirExprSchema, + }) + .meta({ id: "HirLength" }); + +const unarySchema = z + .strictObject({ + ...nodeBase, + kind: z.literal("unary"), + op: hirUnaryOpSchema, + operand: hirExprSchema, + }) + .meta({ id: "HirUnary" }); + +const binarySchema = z + .strictObject({ + ...nodeBase, + kind: z.literal("binary"), + op: hirBinaryOpSchema, + left: hirExprSchema, + right: hirExprSchema, + }) + .meta({ id: "HirBinary" }); + +const condSchema = z + .strictObject({ + ...nodeBase, + kind: z.literal("cond"), + condition: hirExprSchema, + thenBranch: hirExprSchema, + elseBranch: hirExprSchema, + }) + .meta({ id: "HirCond" }); + +const letBindingSchema = z + .strictObject({ + name: z.string(), + nameSpan: spanSchema, + value: hirExprSchema, + }) + .meta({ id: "HirLetBinding" }); + +const letSchema = z + .strictObject({ + ...nodeBase, + kind: z.literal("let"), + bindings: z.array(letBindingSchema), + body: hirExprSchema, + }) + .meta({ id: "HirLet" }); + +const mathCallSchema = z + .strictObject({ + ...nodeBase, + kind: z.literal("mathCall"), + fn: hirMathFnSchema, + args: z.array(hirExprSchema), + }) + .meta({ id: "HirMathCall" }); + +const recordEntrySchema = z + .strictObject({ + key: z.string(), + keySpan: spanSchema, + value: hirExprSchema, + }) + .meta({ id: "HirRecordEntry" }); + +const recordLitSchema = z + .strictObject({ + ...nodeBase, + kind: z.literal("recordLit"), + entries: z.array(recordEntrySchema), + }) + .meta({ id: "HirRecordLit" }); + +const arrayLitSchema = z + .strictObject({ + ...nodeBase, + kind: z.literal("arrayLit"), + elements: z.array(hirExprSchema), + }) + .meta({ id: "HirArrayLit" }); + +const arrayMapSchema = z + .strictObject({ + ...nodeBase, + kind: z.literal("arrayMap"), + target: hirExprSchema, + param: namedSpanSchema, + indexParam: namedSpanSchema.optional(), + body: hirExprSchema, + }) + .meta({ id: "HirArrayMap" }); + +const arrayReduceSchema = z + .strictObject({ + ...nodeBase, + kind: z.literal("arrayReduce"), + target: hirExprSchema, + accParam: namedSpanSchema, + param: namedSpanSchema, + indexParam: namedSpanSchema.optional(), + body: hirExprSchema, + initial: hirExprSchema, + }) + .meta({ id: "HirArrayReduce" }); + +const arrayConcatSchema = z + .strictObject({ + ...nodeBase, + kind: z.literal("arrayConcat"), + left: hirExprSchema, + right: hirExprSchema, + }) + .meta({ id: "HirArrayConcat" }); + +const distributionSchema = z + .strictObject({ + ...nodeBase, + kind: z.literal("distribution"), + dist: hirDistributionKindSchema, + args: z.array(hirExprSchema), + }) + .meta({ id: "HirDistribution" }); + +const distributionMapSchema = z + .strictObject({ + ...nodeBase, + kind: z.literal("distributionMap"), + base: hirExprSchema, + param: namedSpanSchema, + body: hirExprSchema, + }) + .meta({ id: "HirDistributionMap" }); + +/** One schema per `HirExpr` member; the union above is built from this. */ +const hirExprOptions = [ + numberLitSchema, + boolLitSchema, + stringLitSchema, + stringCallSchema, + uuidGenerateSchema, + uuidFromSchema, + constantSchema, + localRefSchema, + paramRefSchema, + scenarioRefSchema, + rangeCallSchema, + fieldAccessSchema, + indexAccessSchema, + lengthSchema, + unarySchema, + binarySchema, + condSchema, + letSchema, + mathCallSchema, + recordLitSchema, + arrayLitSchema, + arrayMapSchema, + arrayReduceSchema, + arrayConcatSchema, + distributionSchema, + distributionMapSchema, +] as const; + +export const hirSurfaceKindSchema = z + .enum([ + "dynamics", + "lambda", + "kernel", + "metric", + "scenario-expression", + "scenario-code", + ]) + .meta({ id: "HirSurfaceKind" }); + +/** A lowered user function: the unit that crosses process boundaries. */ +export const hirFunctionSchema = z + .strictObject({ + hirVersion: z.literal(1), + surface: hirSurfaceKindSchema, + params: z.array(namedSpanSchema), + body: hirExprSchema, + span: spanSchema, + }) + .meta({ + id: "HirFunction", + description: + "A lowered user function (see hir/hir.ts for the grammar). `params[0]` is the input parameter when the surface declares one; `parameters.*` and `scenario.*` reads are dedicated node kinds, never locals.", + }); + +// -- Lockstep with hir.ts ----------------------------------------------------- +// +// Each assertion fails to compile when the schema's output type and the +// declared type diverge in either direction, so a new field or node kind in +// `hir.ts` has to be mirrored above before the package builds. + +/** + * The schema output and the declared type coincide: mutually assignable and + * with the same keys. Assignability alone would let an optional field go + * missing on either side unnoticed. + */ +type Equals = [A] extends [B] + ? [B] extends [A] + ? [keyof A] extends [keyof B] + ? [keyof B] extends [keyof A] + ? true + : false + : false + : false + : false; + +type SchemaKind = z.output<(typeof hirExprOptions)[number]>["kind"]; + +type _Lockstep = [ + Equals, Span>, + Equals, HirNumberLit>, + Equals, HirBoolLit>, + Equals, HirStringLit>, + Equals, HirStringCall>, + Equals, HirUuidGenerate>, + Equals, HirUuidFrom>, + Equals, HirConstant>, + Equals, HirLocalRef>, + Equals, HirParamRef>, + Equals, HirScenarioRef>, + Equals, HirRangeCall>, + Equals, HirFieldAccess>, + Equals, HirIndexAccess>, + Equals, HirLength>, + Equals, HirUnary>, + Equals, HirBinary>, + Equals, HirCond>, + Equals, HirLet>, + Equals, HirMathCall>, + Equals, HirRecordLit>, + Equals, HirArrayLit>, + Equals, HirArrayMap>, + Equals, HirArrayReduce>, + Equals, HirArrayConcat>, + Equals, HirDistribution>, + Equals, HirDistributionMap>, + Equals, HirFunction>, + // Every declared kind has a schema in the union above, and nothing more. + Equals, +]; + +const lockstep: _Lockstep extends true[] ? true : never = true; +void lockstep; diff --git a/libs/@hashintel/petrinaut-core/src/hir/surface-context.ts b/libs/@hashintel/petrinaut-core/src/hir/surface-context.ts index 6eb078584e8..4f3a4434565 100644 --- a/libs/@hashintel/petrinaut-core/src/hir/surface-context.ts +++ b/libs/@hashintel/petrinaut-core/src/hir/surface-context.ts @@ -132,6 +132,11 @@ export type HirMetricContext = { /** ALL places of the root net, keyed by display name (last name wins for * duplicates, matching the runtime object-key overwrite). */ places: HirMetricPlaceInfo[]; + /** + * What the body must return: a number (a metric, the default) or a + * boolean (a state constraint — a metric-shaped condition). + */ + expected?: "real" | "boolean"; }; /** One scenario parameter, read in scenario code as `scenario.`. */ @@ -371,6 +376,7 @@ export function buildLambdaContext( export function buildMetricContext( sdcpn: SDCPN, extensions: PetrinautExtensionSettings = DEFAULT_PETRINAUT_EXTENSIONS, + expected: "real" | "boolean" = "real", ): HirMetricContext { const colorById = collectColors(sdcpn, extensions); const placesByName = new Map(); @@ -388,6 +394,7 @@ export function buildMetricContext( surface: "metric", parameters: toParameterInfos(extensions.parameters ? sdcpn.parameters : []), places: [...placesByName.values()], + expected, }; } diff --git a/libs/@hashintel/petrinaut-core/src/hir/typecheck.ts b/libs/@hashintel/petrinaut-core/src/hir/typecheck.ts index 27a53161179..89d98643c30 100644 --- a/libs/@hashintel/petrinaut-core/src/hir/typecheck.ts +++ b/libs/@hashintel/petrinaut-core/src/hir/typecheck.ts @@ -848,7 +848,15 @@ class Typechecker { return; } case "metric": { - if (!isNumeric(returnType)) { + if (context.expected === "boolean") { + if (!isBoolish(returnType)) { + this.report( + bodySpan, + "hir:metric-return", + `This condition must return a boolean, got ${formatHirType(returnType)}.`, + ); + } + } else if (!isNumeric(returnType)) { this.report( bodySpan, "hir:metric-return", diff --git a/libs/@hashintel/petrinaut-core/src/index.ts b/libs/@hashintel/petrinaut-core/src/index.ts index 5196310c0ec..a14523b4162 100644 --- a/libs/@hashintel/petrinaut-core/src/index.ts +++ b/libs/@hashintel/petrinaut-core/src/index.ts @@ -438,6 +438,29 @@ export type { ScenarioHirItem, ScenarioLoweringInput, } from "./hir/scenario"; +export { + CONSTRAINT_SPACES, + CONSTRAINT_SURFACES, + constraintListSchema, + constraintSchema, + constraintSpaceSchema, + constraintsInSpace, + parameterConstraintSchema, + stateConstraintSchema, +} from "./constraint/constraint"; +export type { + Constraint, + ConstraintSpace, + ParameterConstraint, + StateConstraint, +} from "./constraint/constraint"; +// Type-only: lowering itself stays in ./hir (worker/Node). +export type { + ConstraintSource, + LowerConstraintContext, + LowerConstraintResult, +} from "./constraint/lower"; +export { hirFunctionSchema } from "./hir/hir-schema"; export { AD_HOC_DEFAULT_OPTIMIZE, AD_HOC_DEFAULT_COUNT_OPTIMIZE, diff --git a/libs/@hashintel/petrinaut-core/src/lsp/language-client.ts b/libs/@hashintel/petrinaut-core/src/lsp/language-client.ts index 5f0b3b93555..503ac6ba627 100644 --- a/libs/@hashintel/petrinaut-core/src/lsp/language-client.ts +++ b/libs/@hashintel/petrinaut-core/src/lsp/language-client.ts @@ -6,6 +6,11 @@ import { type LspWorkerFactory, } from "./transport"; +import type { + ConstraintSource, + LowerConstraintContext, + LowerConstraintResult, +} from "../constraint/lower"; import type { PetrinautExtensionSettings } from "../extensions"; // Type-only: must not pull the compiler (`typescript`) into client bundles. import type { HirCompileResult, ScenarioHir } from "../hir"; @@ -129,6 +134,17 @@ export interface LanguageClient { */ requestFormatExpression(this: void, code: string): Promise; + /** + * Lowers one constraint's TypeScript source to HIR (in the worker) and + * checks it produces a boolean. The result is the constraint ready to + * carry, e.g. in an optimization manifest. + */ + requestConstraint( + this: void, + source: ConstraintSource, + context: LowerConstraintContext, + ): Promise; + /** * Tear down the transport. Pending requests reject with "Worker terminated". * Idempotent. @@ -386,6 +402,12 @@ export function createLanguageClient( requestFormatExpression(code) { return sendRequest("sdcpn/formatExpression", { code }); }, + requestConstraint(source, context) { + return sendRequest("sdcpn/lowerConstraint", { + source, + context, + }); + }, dispose() { if (disposed) { diff --git a/libs/@hashintel/petrinaut-core/src/lsp/worker/language-server.worker.ts b/libs/@hashintel/petrinaut-core/src/lsp/worker/language-server.worker.ts index 543be8cd217..35585d5b4e2 100644 --- a/libs/@hashintel/petrinaut-core/src/lsp/worker/language-server.worker.ts +++ b/libs/@hashintel/petrinaut-core/src/lsp/worker/language-server.worker.ts @@ -30,6 +30,7 @@ import { buildScenarioExpressionContext, compileHirArtifacts, formatTypeScriptExpression, + lowerConstraint, lowerScenarioToHir, } from "../../hir"; import { getHirDiagnosticsForItem } from "../lib/check-hir"; @@ -592,6 +593,12 @@ workerRuntime.onMessage((data) => { break; } + case "sdcpn/lowerConstraint": { + const { id } = data; + respond(id, lowerConstraint(data.params.source, data.params.context)); + break; + } + case "textDocument/completion": { const { id } = data; if (!server) { diff --git a/libs/@hashintel/petrinaut-core/src/lsp/worker/protocol.ts b/libs/@hashintel/petrinaut-core/src/lsp/worker/protocol.ts index 277a9ac6ea0..cb5273bb67c 100644 --- a/libs/@hashintel/petrinaut-core/src/lsp/worker/protocol.ts +++ b/libs/@hashintel/petrinaut-core/src/lsp/worker/protocol.ts @@ -1,3 +1,7 @@ +import type { + ConstraintSource, + LowerConstraintContext, +} from "../../constraint/lower"; /** * LSP-inspired protocol types for the language server WebWorker. * @@ -194,6 +198,15 @@ type ClientRequest = /** A single scenario-expression to re-print canonically. */ code: string; }; + } + | { + jsonrpc: "2.0"; + id: number; + method: "sdcpn/lowerConstraint"; + params: { + source: ConstraintSource; + context: LowerConstraintContext; + }; }; /** Any message from the main thread to the worker. */ diff --git a/libs/@hashintel/petrinaut-core/src/optimization.ts b/libs/@hashintel/petrinaut-core/src/optimization.ts index db102ee09f1..31fbed5a296 100644 --- a/libs/@hashintel/petrinaut-core/src/optimization.ts +++ b/libs/@hashintel/petrinaut-core/src/optimization.ts @@ -1,5 +1,6 @@ import { z } from "zod"; +import { constraintListSchema } from "./constraint/constraint"; import { parseSDCPNFile } from "./file-format/parse-sdcpn-file"; import { sdcpnSchema } from "./file-format/types"; @@ -120,6 +121,14 @@ export const petrinautOptimizationParameterBindingSchema = z ]) .meta({ description: "The per-study treatment of one scenario parameter." }); +function addIssue( + context: z.core.$RefinementCtx, + path: PropertyKey[], + message: string, +): void { + context.addIssue({ code: "custom", path, message }); +} + export const petrinautOptimizationObjectiveSchema = z .strictObject({ metricId: z.string().min(1), @@ -185,14 +194,6 @@ const optimizationScenarioSchema = z "The sole scenario and the exhaustive, transient treatment of its parameters.", }); -function addIssue( - context: z.core.$RefinementCtx, - path: PropertyKey[], - message: string, -): void { - context.addIssue({ code: "custom", path, message }); -} - function validateScenarioParameterDefault( parameter: { identifier: string; @@ -238,6 +239,10 @@ export const petrinautOptimizationManifestSchema = z model: optimizationModelSchema, scenario: optimizationScenarioSchema, objective: petrinautOptimizationObjectiveSchema, + constraints: constraintListSchema.optional().meta({ + description: + 'Optional boolean conditions over the parameter space (`space: "parameters"`) and the simulation state (`space: "state"`). Absent means unconstrained; nothing enforces them yet.', + }), execution: petrinautOptimizationExecutionSchema, study: petrinautOptimizationStudySchema, }) @@ -521,10 +526,14 @@ export const petrinautOptimizationDescribeResultSchema = z "Study settings with the execution seed. `seedsPerTrial` is reported once the CLI runs seeded replicates; absent means 1.", }), parameters: z.array(petrinautOptimizationDescribeParameterSchema), + constraints: constraintListSchema.optional().meta({ + description: + "The manifest's constraints, passed through verbatim so protocol clients (the Python binding) can evaluate their HIR. Absent means unconstrained.", + }), }) .meta({ description: - "The `optimization.describe` result: direction, study settings, and the parameters that are not fixed.", + "The `optimization.describe` result: direction, study settings, the parameters that are not fixed, and the study's constraints.", }); export const petrinautOptimizationReplicateSchema = z diff --git a/libs/@hashintel/petrinaut-core/src/optimization/describe.ts b/libs/@hashintel/petrinaut-core/src/optimization/describe.ts index b1a5937766b..010993edda4 100644 --- a/libs/@hashintel/petrinaut-core/src/optimization/describe.ts +++ b/libs/@hashintel/petrinaut-core/src/optimization/describe.ts @@ -142,6 +142,9 @@ export const describeOptimization = ( ({ parameter, domain }) => describeOptimizationParameter(parameter, domain), ), + // Verbatim pass-through: protocol clients (the Python binding) evaluate + // the embedded constraint HIR themselves. + ...(manifest.constraints ? { constraints: manifest.constraints } : {}), }; }; diff --git a/libs/@hashintel/petrinaut/docs/optimization.md b/libs/@hashintel/petrinaut/docs/optimization.md index 174f81d78ad..e6024982547 100644 --- a/libs/@hashintel/petrinaut/docs/optimization.md +++ b/libs/@hashintel/petrinaut/docs/optimization.md @@ -86,6 +86,15 @@ The controls depend on the scenario parameter type: Parameters are fixed by default. Search ranges belong to this optimization run, not to the saved scenario. +## Constraints + +The **Constraints** section of the create-optimization drawer records boolean conditions with the study. They are carried in the study's manifest and readable by every consumer (the Python tooling included), but **nothing enforces them yet** -- they do not prune trials or stop runs. Two kinds: + +- **Parameter constraints** -- one-line expressions over the study's parameters (`scenario.*` for scenario parameters, `parameters.*` for net parameters) that must produce a boolean, for example `scenario.min_load < scenario.max_load`. In a later iteration these will let the optimizer avoid infeasible parameter combinations. +- **State constraints** -- small code bodies that read the simulation `state` exactly like a [metric](experiments.md#metrics) and `return` a boolean, for example `return state.places.Queue.count <= 10;`. In a later iteration these will measure how close a run comes to leaving the safe region, not just whether it did. + +Add a condition with its **Add ... constraint** button, edit it in place, and remove it with **Remove**. Conditions are checked when you press Run: one that does not compile, or does not produce a boolean, blocks the submission with its error message. Empty rows are ignored. + ## Watching results Open an optimization row to follow it while it runs. The drawer updates as diff --git a/libs/@hashintel/petrinaut/src/react/lsp/context.ts b/libs/@hashintel/petrinaut/src/react/lsp/context.ts index 9eb499ab5be..871522159cb 100644 --- a/libs/@hashintel/petrinaut/src/react/lsp/context.ts +++ b/libs/@hashintel/petrinaut/src/react/lsp/context.ts @@ -4,6 +4,9 @@ import type { AdHocSynthesisContext, CompileHirArtifactsOptions, CompletionList, + ConstraintSource, + LowerConstraintContext, + LowerConstraintResult, Diagnostic, DocumentUri, HirCompileResult, @@ -73,6 +76,14 @@ export interface LanguageClientContextValue { * code does not lower — keep the user's text in that case. */ requestFormatExpression: (code: string) => Promise; + /** + * Lower one constraint's source to HIR (in the language worker) and check + * it produces a boolean. + */ + requestConstraint: ( + source: ConstraintSource, + context: LowerConstraintContext, + ) => Promise; /** Initialize a temporary scenario editing session. */ initializeScenarioSession: (params: ScenarioSessionParams) => void; /** Update a scenario editing session. */ @@ -121,6 +132,18 @@ export const DEFAULT_LANGUAGE_CLIENT_CONTEXT: LanguageClientContextValue = { placeExpressions: {}, }), requestFormatExpression: () => Promise.resolve(null), + requestConstraint: () => + Promise.resolve({ + ok: false as const, + diagnostics: [ + { + code: "hir:no-language-client", + message: "No language client is wired; constraints cannot compile.", + severity: "error" as const, + span: { start: 0, length: 0 }, + }, + ], + }), initializeScenarioSession: () => {}, updateScenarioSession: () => {}, killScenarioSession: () => {}, diff --git a/libs/@hashintel/petrinaut/src/react/lsp/provider.tsx b/libs/@hashintel/petrinaut/src/react/lsp/provider.tsx index c8b362ca5bb..a61c244f326 100644 --- a/libs/@hashintel/petrinaut/src/react/lsp/provider.tsx +++ b/libs/@hashintel/petrinaut/src/react/lsp/provider.tsx @@ -112,6 +112,7 @@ export const LanguageClientProvider: React.FC<{ requestHirArtifacts: client.requestHirArtifacts, requestScenarioHir: client.requestScenarioHir, requestFormatExpression: client.requestFormatExpression, + requestConstraint: client.requestConstraint, initializeScenarioSession: client.initializeScenarioSession, updateScenarioSession: client.updateScenarioSession, killScenarioSession: client.killScenarioSession, diff --git a/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/experiments/create-experiment-drawer.test.tsx b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/experiments/create-experiment-drawer.test.tsx index 0dc0a97887b..e8b53861a38 100644 --- a/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/experiments/create-experiment-drawer.test.tsx +++ b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/experiments/create-experiment-drawer.test.tsx @@ -83,6 +83,12 @@ function makeLanguageClient(): LanguageClientContextValue { ), requestHover: vi.fn(() => Promise.resolve(null)), requestSignatureHelp: vi.fn(() => Promise.resolve(null)), + requestConstraint: vi.fn(() => + Promise.resolve({ + ok: false as const, + diagnostics: [], + }), + ), requestScenarioHir: vi.fn(() => Promise.resolve({ version: 1 as const, diff --git a/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/metrics/create-metric-drawer.test.tsx b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/metrics/create-metric-drawer.test.tsx index 45378f268bb..511ffdd618e 100644 --- a/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/metrics/create-metric-drawer.test.tsx +++ b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/metrics/create-metric-drawer.test.tsx @@ -39,6 +39,12 @@ function makeLanguageClientValue(): LanguageClientContextValue { ), requestHover: vi.fn(() => Promise.resolve(null)), requestSignatureHelp: vi.fn(() => Promise.resolve(null)), + requestConstraint: vi.fn(() => + Promise.resolve({ + ok: false as const, + diagnostics: [], + }), + ), requestScenarioHir: vi.fn(() => Promise.resolve({ version: 1 as const, diff --git a/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/create-optimization-drawer.test.tsx b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/create-optimization-drawer.test.tsx index 906d1a15b91..7561be31649 100644 --- a/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/create-optimization-drawer.test.tsx +++ b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/create-optimization-drawer.test.tsx @@ -45,6 +45,8 @@ import type { SDCPNContextValue } from "../../../../../../react/state/sdcpn-cont import type { OptimizationParameterDraft } from "./optimization-parameter-row"; import type { AdHocScenarioState, + ConstraintSource, + LowerConstraintResult, Metric, PetrinautOptimizationInput, Scenario, @@ -336,6 +338,30 @@ function makeSuccessfulLanguageClient(): LanguageClientContextValue { ), requestHover: vi.fn(() => Promise.resolve(null)), requestSignatureHelp: vi.fn(() => Promise.resolve(null)), + requestConstraint: vi.fn((source: ConstraintSource) => + Promise.resolve({ + ok: true, + constraint: { + ...source, + hir: { + hirVersion: 1, + surface: + source.space === "parameters" ? "scenario-expression" : "metric", + params: + source.space === "parameters" + ? [] + : [{ name: "state", span: { start: 0, length: 0 } }], + body: { + kind: "boolLit", + id: 0, + span: { start: 0, length: 0 }, + value: true, + }, + span: { start: 0, length: 0 }, + }, + }, + } as LowerConstraintResult), + ), requestScenarioHir: vi.fn(() => Promise.resolve({ version: 1 as const, @@ -686,6 +712,51 @@ describe("CreateOptimizationDrawer", () => { ).toBe(true); }); + it("lowers authored constraints and embeds them in the manifest", async () => { + const languageClient = makeSuccessfulLanguageClient(); + const createOptimization = vi.fn( + async (_input: PetrinautOptimizationInput) => "optimization-constrained", + ); + const savedMetric = sirSdcpnContextValue.petriNetDefinition.metrics?.[0]; + openConfiguration({ createOptimization, languageClient }); + + fireEvent.change( + screen.getByRole("combobox", { name: "Select a metric" }), + { + target: { value: `${MODEL_METRIC_VALUE_PREFIX}${savedMetric!.id}` }, + }, + ); + fireEvent.click( + screen.getByRole("checkbox", { name: "Optimize infected_ratio" }), + ); + fireEvent.click(screen.getByRole("button", { name: "Maximize" })); + + // Author one parameter constraint. + fireEvent.click( + screen.getByRole("button", { name: "Add parameter constraint" }), + ); + fireEvent.change(screen.getByRole("textbox", { name: "Metric code" }), { + target: { value: "scenario.infected_ratio < 0.9" }, + }); + + fireEvent.click(screen.getByRole("button", { name: /Run/ })); + await waitFor(() => expect(createOptimization).toHaveBeenCalledOnce()); + + expect( + vi.mocked(languageClient.requestConstraint).mock.calls[0]?.[0], + ).toMatchObject({ + space: "parameters", + code: "scenario.infected_ratio < 0.9", + }); + const submittedInput = createOptimization.mock.calls[0]![0]; + expect(submittedInput.constraints).toHaveLength(1); + expect(submittedInput.constraints?.[0]).toMatchObject({ + space: "parameters", + code: "scenario.infected_ratio < 0.9", + hir: { surface: "scenario-expression" }, + }); + }); + it("submits a transient custom metric without persisting it", async () => { const languageClient = makeSuccessfulLanguageClient(); const createOptimization = vi.fn( diff --git a/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/create-optimization-drawer.tsx b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/create-optimization-drawer.tsx index 4862560e7c2..b324411a07e 100644 --- a/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/create-optimization-drawer.tsx +++ b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/create-optimization-drawer.tsx @@ -73,6 +73,7 @@ import type { import type { AdHocScenarioState, AdHocSynthesisError, + Constraint, Metric, PetrinautOptimizationInput, PetrinautOptimizationParameterBinding, @@ -185,6 +186,9 @@ const errorsStyle = css({ }); type Direction = "maximize" | "minimize"; + +/** One constraint being authored: stable id + editable source. */ +type ConstraintDraft = { id: string; code: string }; type MetricSource = "saved" | "custom"; type ParameterDrafts = Record; @@ -227,6 +231,86 @@ const ScenarioSelectLabel = ({ ); }; +const constraintRowStyle = css({ + display: "flex", + alignItems: "flex-start", + gap: "2", + "& > :first-child": { + flex: "[1]", + minWidth: "0", + }, +}); + +const constraintListStyle = css({ + display: "flex", + flexDirection: "column", + gap: "2", +}); + +/** + * One editable list of constraints: an expression editor per row, a remove + * button, and a quiet add button. Parameter constraints edit as one-line + * expressions; state constraints as small code bodies. + */ +const ConstraintDraftList = ({ + drafts, + onChange, + multiline, + addLabel, + ariaPrefix, +}: { + drafts: ConstraintDraft[]; + onChange: (drafts: ConstraintDraft[]) => void; + multiline: boolean; + addLabel: string; + ariaPrefix: string; +}) => ( +
+ {drafts.map((draft, index) => ( +
+ + onChange( + drafts.map((candidate) => + candidate.id === draft.id + ? { ...candidate, code: code ?? "" } + : candidate, + ), + ) + } + /> + +
+ ))} +
+ +
+
+); + const InlineObjectiveMetricForm = ({ form }: { form: MetricFormInstance }) => { const values = useStore(form.store, (state) => state.values); const metricSessionId = useMetricLspSession(values.code); @@ -501,6 +585,7 @@ export function buildPetrinautOptimizationInput({ seed, dt, maxTime, + constraints, }: { name: string; title: string; @@ -514,6 +599,7 @@ export function buildPetrinautOptimizationInput({ seed: number; dt: number; maxTime: number; + constraints?: Constraint[]; }): PetrinautOptimizationInput { // Keyed by scenario parameter identifiers from the net definition: no // prototype. @@ -573,6 +659,7 @@ export function buildPetrinautOptimizationInput({ }, scenario: { id: scenario.id, parameterBindings }, objective: { metricId: metric.id, direction }, + ...(constraints ? { constraints } : {}), execution: { seed, dt, maxTime, seedsPerTrial }, study: { trials: optimizationSteps, sampler: OPTIMIZATION_SAMPLER }, }); @@ -597,6 +684,7 @@ export function buildAdHocPetrinautOptimizationInput({ seed, dt, maxTime, + constraints, }: { name: string; title: string; @@ -610,6 +698,7 @@ export function buildAdHocPetrinautOptimizationInput({ seed: number; dt: number; maxTime: number; + constraints?: Constraint[]; }): PetrinautOptimizationInput { return petrinautOptimizationInputSchema.parse({ kind: "petrinaut-optimization", @@ -625,6 +714,7 @@ export function buildAdHocPetrinautOptimizationInput({ }, scenario: { id: scenario.id, parameterBindings }, objective: { metricId: metric.id, direction }, + ...(constraints ? { constraints } : {}), execution: { seed, dt, maxTime, seedsPerTrial }, study: { trials: optimizationSteps, sampler: OPTIMIZATION_SAMPLER }, }); @@ -638,7 +728,7 @@ export const CreateOptimizationDrawer = ({ onClose: () => void; }) => { const { extensions, petriNetDefinition, title } = use(SDCPNContext); - const { requestHirArtifacts } = use(LanguageClientContext); + const { requestHirArtifacts, requestConstraint } = use(LanguageClientContext); const { createOptimization } = use(OptimizationsContext); const { enableAdHocScenarios, webGpuEnabled } = use(UserSettingsContext); const source = useOptimizationSource(); @@ -674,6 +764,15 @@ export const CreateOptimizationDrawer = ({ const [maxTime, setMaxTime] = useState(180); const [error, setError] = useState(null); const [isSubmitting, setIsSubmitting] = useState(false); + // Boolean conditions embedded in the manifest as source + lowered HIR. + // Declarative for now: carried and readable (Python included), not + // enforced by the study. + const [parameterConstraintDrafts, setParameterConstraintDrafts] = useState< + ConstraintDraft[] + >([]); + const [stateConstraintDrafts, setStateConstraintDrafts] = useState< + ConstraintDraft[] + >([]); const isAdHoc = enableAdHocScenarios && selectedScenarioId === AD_HOC_SCENARIO_VALUE; @@ -782,6 +881,8 @@ export const CreateOptimizationDrawer = ({ setMaxTime(180); setError(null); setIsSubmitting(false); + setParameterConstraintDrafts([]); + setStateConstraintDrafts([]); }; const resetState = () => { @@ -883,6 +984,42 @@ export const CreateOptimizationDrawer = ({ } } + // Lower each authored constraint against the study's parameters; a + // failing one blocks submission with its first diagnostic. + const constraintContext = { + netParameters: extensions.parameters + ? petriNetDefinition.parameters + : [], + scenarioParameters: scenarioForRun.scenarioParameters, + sdcpn: petriNetDefinition, + extensions, + }; + const constraints: Constraint[] = []; + for (const [space, drafts_] of [ + ["parameters", parameterConstraintDrafts], + ["state", stateConstraintDrafts], + ] as const) { + for (const [index, draft] of drafts_.entries()) { + if (draft.code.trim() === "") { + continue; + } + const lowered = await requestConstraint( + { space, id: draft.id, code: draft.code }, + constraintContext, + ); + if (!lowered.ok) { + setIsSubmitting(false); + setError( + `${space === "parameters" ? "Parameter" : "State"} constraint ${index + 1}: ${lowered.diagnostics[0]?.message ?? "does not compile"}`, + ); + return; + } + constraints.push(lowered.constraint); + } + } + const manifestConstraints = + constraints.length > 0 ? constraints : undefined; + const input = adHocBindings ? buildAdHocPetrinautOptimizationInput({ name, @@ -897,6 +1034,7 @@ export const CreateOptimizationDrawer = ({ seed, dt, maxTime, + constraints: manifestConstraints, }) : buildPetrinautOptimizationInput({ name, @@ -911,6 +1049,7 @@ export const CreateOptimizationDrawer = ({ seed, dt, maxTime, + constraints: manifestConstraints, }); await createOptimization(input, { computeBackend, parallelism }); resetState(); @@ -1276,6 +1415,38 @@ export const CreateOptimizationDrawer = ({ )} +
+ + Parameter constraints are expressions over the study's + parameters (scenario.*, parameters.*), e.g. scenario.min_load + < scenario.max_load. + + + + State constraints read the simulation state like a metric body + and must return a boolean, e.g. return + state.places.Queue.count <= 10; + + +
+
Choose a saved metric or write custom code for this run. diff --git a/libs/@local/petrinaut-arch-docs/content/cli/usage-manual.mdx b/libs/@local/petrinaut-arch-docs/content/cli/usage-manual.mdx index 3f9b556e94c..41acfe6fa1e 100644 --- a/libs/@local/petrinaut-arch-docs/content/cli/usage-manual.mdx +++ b/libs/@local/petrinaut-arch-docs/content/cli/usage-manual.mdx @@ -258,6 +258,18 @@ that are not fixed: A logarithmic integer domain requires `step: 1`. `default` is the scenario's default value; the optimizer picks each trial's value. +### Constraints + +A manifest may carry `constraints`: a list of boolean conditions, each +`{ "space", "id", "name"?, "code", "hir" }`. `space` is `"parameters"` for a +condition over the scenario and net parameters (checkable before a run) or +`"state"` for a condition over the simulation state (observed while a run +goes); `code` is the authored TypeScript and `hir` its lowered form, validated +against the HIR grammar. The describe response returns the list verbatim, so a +client evaluates them itself: the Python binding reads each one as a callable +with a boolean, a signed margin, and a pydantic validator. The CLI does not +enforce constraints yet. + ### Evaluate a trial Send every — and only — described parameter: diff --git a/libs/@local/petrinaut-python/README.md b/libs/@local/petrinaut-python/README.md index 40ba69776a1..803a541a4da 100644 --- a/libs/@local/petrinaut-python/README.md +++ b/libs/@local/petrinaut-python/README.md @@ -83,6 +83,55 @@ fixed-value injection, simulation, and metric evaluation. The manifest contract, protocol, and seeded-runs semantics are documented in the [CLI usage manual](../petrinaut-arch-docs/content/cli/usage-manual.mdx). +## Constraints + +A study's constraints come back from `describe()` as `{code, hir}` pairs, and +the binding evaluates the `hir` side itself: no TypeScript toolchain, and the +whole tree is validated node by node against the grammar's pydantic models +before anything runs. Two shapes, told apart by `space`: + +```python +from petrinaut import parse_constraints, violations + +constraints = parse_constraints(description.constraints) +for constraint in constraints: + print(constraint.space, constraint.code) + +ordering = constraints[0] # a ParameterConstraint +ordering({"min_load": 2, "max_load": 8}) # True +ordering.margin({"min_load": 2, "max_load": 8}) # 6.0, >= 0 iff satisfied +ordering.violation({"min_load": 8, "max_load": 2}) # 6.0, <= 0 iff satisfied +ordering.check({"min_load": 8, "max_load": 2}) # raises ConstraintViolation +``` + +- `ParameterConstraint` ranges over the parameter space and takes a `scenario` + mapping plus the net `parameters`. It is checkable before a run starts. +- `StateConstraint` ranges over the simulation state and takes a `state` + record (places keyed by name, each with `count` and `tokens`) plus the net + `parameters`. + +Both give four readings of the condition: the boolean (call it), the signed +`margin`, the `violation` in the sign Optuna's `constraints_func` expects, and +`check`, which raises. `violations(constraints, scenario=..., state=...)` +returns one violation per constraint for a sampler. `validator()` packages the +check for pydantic: + +```python +from typing import Annotated +from pydantic import AfterValidator, BaseModel + + +class Study(BaseModel): + scenario: Annotated[dict[str, float], AfterValidator(ordering.validator())] +``` + +A parameter constraint over plain arithmetic also has a symbolic reading with +the `sympy` extra (`petrinaut-python[sympy]`): `ordering.to_sympy()` returns a +`SymbolicConstraint` whose `.expression` is the relation over one real symbol +per parameter (`.scenario` and `.parameters` map names to symbols), ready for +`sympy.solve_univariate_inequality` or `simplify`. Arrays, records, strings +and `Math.random()` have no symbolic form and raise `NotSymbolicError`. + ## Timeouts and limits - Bootstrap (spawn to readiness) and each protocol response have deadlines, diff --git a/libs/@local/petrinaut-python/pyproject.toml b/libs/@local/petrinaut-python/pyproject.toml index 88e7e43b3d9..dd8daafff87 100644 --- a/libs/@local/petrinaut-python/pyproject.toml +++ b/libs/@local/petrinaut-python/pyproject.toml @@ -8,12 +8,18 @@ dependencies = [ "pydantic>=2.13.4", ] +[project.optional-dependencies] +sympy = [ + "sympy>=1.13", +] + [dependency-groups] dev = [ "basedpyright>=1.39.10", "datamodel-code-generator>=0.74.0", "pytest>=8.3", "ruff>=0.16.4", + "sympy>=1.13", ] [build-system] diff --git a/libs/@local/petrinaut-python/src/petrinaut/__init__.py b/libs/@local/petrinaut-python/src/petrinaut/__init__.py index 6cff6a23203..b71d2d16cbd 100644 --- a/libs/@local/petrinaut-python/src/petrinaut/__init__.py +++ b/libs/@local/petrinaut-python/src/petrinaut/__init__.py @@ -8,15 +8,29 @@ serializes requests to it, and shuts it down. "Session" rather than "client" because the object carries that lifecycle, not just the wire format. +Constraints a study carries are readable without a session: `parse_constraint` +turns the protocol's `{code, hir}` pairs into callables that evaluate the HIR +here, as a boolean, a signed margin, or a pydantic validator. + @layerRoot python-bindings @role Python sessions owning one CLI process each, translating protocol frames into methods and exceptions """ +from .constraint import ( + Constraint, + ConstraintViolation, + ParameterConstraint, + StateConstraint, + parse_constraint, + parse_constraints, + violations, +) from .errors import ( PetrinautClientError, PetrinautProtocolError, PetrinautRunError, ) +from .hir import HirEvaluationError, evaluate_hir, hir_margin from .models import ( OptimizationBooleanParameter, OptimizationDescribeResult, @@ -27,8 +41,13 @@ ) from .optimization import OptimizationSession from .session import PetrinautSession +from .symbolic import NotSymbolicError, SymbolicConstraint __all__ = [ + "Constraint", + "ConstraintViolation", + "HirEvaluationError", + "NotSymbolicError", "OptimizationBooleanParameter", "OptimizationDescribeResult", "OptimizationEvaluateResult", @@ -36,8 +55,16 @@ "OptimizationIntParameter", "OptimizationReplicate", "OptimizationSession", + "ParameterConstraint", "PetrinautClientError", "PetrinautProtocolError", "PetrinautRunError", "PetrinautSession", + "StateConstraint", + "SymbolicConstraint", + "evaluate_hir", + "hir_margin", + "parse_constraint", + "parse_constraints", + "violations", ] diff --git a/libs/@local/petrinaut-python/src/petrinaut/constraint.py b/libs/@local/petrinaut-python/src/petrinaut/constraint.py new file mode 100644 index 00000000000..39d9c0fd050 --- /dev/null +++ b/libs/@local/petrinaut-python/src/petrinaut/constraint.py @@ -0,0 +1,260 @@ +"""Constraints as callables. A constraint is a boolean condition authored in +Petrinaut and carried as ``{code, hir}``; it comes in two shapes, told apart +by ``space``: + +- :class:`ParameterConstraint` ranges over the parameter space and is + called with a ``scenario`` mapping (plus the net ``parameters``). It can + be checked before anything runs, so it doubles as a validator. +- :class:`StateConstraint` ranges over the simulation state and is called + with a ``state`` record (plus the net ``parameters``). + +Both are pydantic models: parsing one validates the whole HIR tree node by +node, and both expose the same four readings of a condition — the boolean +(:meth:`__call__`), the signed margin (:meth:`margin`, ``>= 0`` iff +satisfied), the violation Optuna-style constrained samplers consume +(:meth:`violation`, ``<= 0`` iff satisfied), and a check that raises +(:meth:`check`). :meth:`validator` packages the check for pydantic's +``AfterValidator``. A parameter constraint can also be read symbolically +through :meth:`ParameterConstraint.to_sympy`. + +>>> constraint = parse_constraint(described.constraints[0]) +>>> constraint(scenario={"min_load": 2, "max_load": 8}) +True +>>> constraint.margin(scenario={"min_load": 2, "max_load": 8}) +6.0 +""" + +from __future__ import annotations + +from collections.abc import Callable, Iterable, Mapping +from typing import TYPE_CHECKING, Annotated, Any, TypeAlias + +from pydantic import Field, TypeAdapter + +from . import models as m +from .hir import HirEvaluationError, Scalar, Value, evaluate_hir, hir_margin + +if TYPE_CHECKING: + from .symbolic import SymbolicConstraint + +__all__ = [ + "Constraint", + "ConstraintViolation", + "ParameterConstraint", + "StateConstraint", + "parse_constraint", + "parse_constraints", + "violations", +] + + +class ConstraintViolation(ValueError): + """A constraint check failed. ``margin`` says by how much (negative).""" + + def __init__( + self, constraint: ParameterConstraint | StateConstraint, margin: float + ) -> None: + self.constraint = constraint + self.margin = margin + label = constraint.name or constraint.id + super().__init__( + f'Constraint "{label}" is violated (margin {margin:g}): {constraint.code}' + ) + + +def _as_bool( + constraint: m.ParameterConstraint | m.StateConstraint, value: Value +) -> bool: + if not isinstance(value, bool): + raise HirEvaluationError( + f'Constraint "{constraint.id}" produced a {type(value).__name__}, expected a boolean' + ) + return value + + +class ParameterConstraint(m.ParameterConstraint): + """One boolean condition over the parameter space: an expression over + ``scenario.*`` and ``parameters.*``, checkable before a run starts.""" + + def __call__( + self, + scenario: Mapping[str, Scalar], + parameters: Mapping[str, Scalar] | None = None, + ) -> bool: + """Whether the constraint holds for these parameter values.""" + return _as_bool( + self, evaluate_hir(self.hir, scenario=scenario, parameters=parameters) + ) + + def margin( + self, + scenario: Mapping[str, Scalar], + parameters: Mapping[str, Scalar] | None = None, + ) -> float: + """Signed robustness margin: ``>= 0`` iff satisfied, its magnitude + the distance to the boundary.""" + return hir_margin(self.hir, scenario=scenario, parameters=parameters) + + def violation( + self, + scenario: Mapping[str, Scalar], + parameters: Mapping[str, Scalar] | None = None, + ) -> float: + """``-margin``: ``<= 0`` iff satisfied, the sign convention of + Optuna's ``constraints_func``.""" + return -self.margin(scenario, parameters) + + def check( + self, + scenario: Mapping[str, Scalar], + parameters: Mapping[str, Scalar] | None = None, + ) -> None: + """Raise :class:`ConstraintViolation` unless the constraint holds.""" + margin = self.margin(scenario, parameters) + if margin < 0: + raise ConstraintViolation(self, margin) + + def validator( + self, parameters: Mapping[str, Scalar] | None = None + ) -> Callable[[Mapping[str, Scalar]], Mapping[str, Scalar]]: + """The check as a pydantic ``AfterValidator`` body: returns the + scenario mapping when the constraint holds, raises otherwise (a + :class:`ConstraintViolation`, which pydantic reports as a + validation error). + + >>> Scenario = Annotated[dict[str, float], AfterValidator(constraint.validator())] + """ + + def validate(scenario: Mapping[str, Scalar]) -> Mapping[str, Scalar]: + self.check(scenario, parameters) + return scenario + + return validate + + def to_sympy(self) -> SymbolicConstraint: + """The condition as a SymPy relation over one symbol per parameter. + Needs the ``sympy`` extra; raises :class:`~petrinaut.symbolic.NotSymbolicError` + when the condition reads something SymPy cannot represent.""" + from .symbolic import to_sympy + + return to_sympy(self) + + +class StateConstraint(m.StateConstraint): + """One boolean condition over the simulation state, authored like a + metric body and observed while a run goes.""" + + def _locals(self, state: Mapping[str, Any]) -> dict[str, Value]: + # The annotation is a promise callers can break; the check is for them. + if not isinstance(state, Mapping): # pyright: ignore[reportUnnecessaryIsInstance] + raise HirEvaluationError( + f"A state constraint takes a state record, got {type(state).__name__}" + ) + # The metric surface declares the state record as its first parameter. + name = self.hir.params[0].name if self.hir.params else "state" + return {name: dict(state)} + + def __call__( + self, + state: Mapping[str, Any], + parameters: Mapping[str, Scalar] | None = None, + ) -> bool: + """Whether the constraint holds for this state record (places keyed + by name, each with ``count`` and ``tokens``).""" + return _as_bool( + self, + evaluate_hir(self.hir, parameters=parameters, locals_=self._locals(state)), + ) + + def margin( + self, + state: Mapping[str, Any], + parameters: Mapping[str, Scalar] | None = None, + ) -> float: + """Signed robustness margin: ``>= 0`` iff satisfied.""" + return hir_margin(self.hir, parameters=parameters, locals_=self._locals(state)) + + def violation( + self, + state: Mapping[str, Any], + parameters: Mapping[str, Scalar] | None = None, + ) -> float: + """``-margin``: ``<= 0`` iff satisfied.""" + return -self.margin(state, parameters) + + def check( + self, + state: Mapping[str, Any], + parameters: Mapping[str, Scalar] | None = None, + ) -> None: + """Raise :class:`ConstraintViolation` unless the constraint holds.""" + margin = self.margin(state, parameters) + if margin < 0: + raise ConstraintViolation(self, margin) + + def validator( + self, parameters: Mapping[str, Scalar] | None = None + ) -> Callable[[Mapping[str, Any]], Mapping[str, Any]]: + """The check as a pydantic ``AfterValidator`` body over a state record.""" + + def validate(state: Mapping[str, Any]) -> Mapping[str, Any]: + self.check(state, parameters) + return state + + return validate + + +Constraint: TypeAlias = ParameterConstraint | StateConstraint + +_CONSTRAINT_ADAPTER: TypeAdapter[Constraint] = TypeAdapter( + Annotated[ParameterConstraint | StateConstraint, Field(discriminator="space")] +) + + +def parse_constraint( + data: Mapping[str, Any] | m.ParameterConstraint | m.StateConstraint, +) -> Constraint: + """One constraint as a callable, from a protocol model or a mapping. + Validation covers the whole HIR tree and raises + :class:`pydantic.ValidationError` for anything outside the grammar.""" + if isinstance(data, (ParameterConstraint, StateConstraint)): + return data + if isinstance(data, (m.ParameterConstraint, m.StateConstraint)): + data = data.model_dump() + return _CONSTRAINT_ADAPTER.validate_python(data) + + +def parse_constraints( + items: Iterable[Mapping[str, Any] | m.ParameterConstraint | m.StateConstraint] + | None, +) -> list[Constraint]: + """Every constraint of a describe result (or any list) as callables; + ``None`` reads as no constraints.""" + return [parse_constraint(item) for item in items or ()] + + +def violations( + constraints: Iterable[Constraint], + *, + scenario: Mapping[str, Scalar] | None = None, + parameters: Mapping[str, Scalar] | None = None, + state: Mapping[str, Any] | None = None, +) -> list[float]: + """One signed violation per constraint, in order, ``<= 0`` iff + satisfied: the sequence a constrained sampler consumes. Parameter + constraints read ``scenario``; state constraints read ``state``, and + asking for one without a state is an error rather than a skipped entry, + so the sequence keeps one slot per constraint.""" + out: list[float] = [] + for constraint in constraints: + if isinstance(constraint, ParameterConstraint): + if scenario is None: + raise ValueError( + f'Parameter constraint "{constraint.id}" needs a scenario' + ) + out.append(constraint.violation(scenario, parameters)) + else: + if state is None: + raise ValueError(f'State constraint "{constraint.id}" needs a state') + out.append(constraint.violation(state, parameters)) + return out diff --git a/libs/@local/petrinaut-python/src/petrinaut/hir.py b/libs/@local/petrinaut-python/src/petrinaut/hir.py new file mode 100644 index 00000000000..ab667996384 --- /dev/null +++ b/libs/@local/petrinaut-python/src/petrinaut/hir.py @@ -0,0 +1,684 @@ +"""Evaluation of serialized HIR expressions — Petrinaut's shared expression +representation. ``hir/hir.ts`` in ``@hashintel/petrinaut-core`` owns the +grammar; :mod:`petrinaut.models` carries it as pydantic models generated from +the CLI's protocol schema, so a document is validated node by node before +anything here runs. Constraints travel as ``{code, hir}`` pairs; this module +evaluates the ``hir`` side so Python consumers need no TypeScript frontend. + +Two evaluation modes: + +- :func:`evaluate_hir` — the expression's value; for a constraint, ``True`` + means satisfied. +- :func:`hir_margin` — a signed robustness margin (``>= 0`` means + satisfied): comparisons yield signed slack, ``&&`` combines by ``min``, + ``||`` by ``max``, ``!`` negates. This is the learnable signal constrained + samplers (e.g. Optuna's ``constraints_func``, which expects violation + ``<= 0`` — i.e. ``-margin``) consume. The same walk evaluated over + intervals instead of scalars would bound a constraint over a whole + parameter box; that extension is deliberately not implemented yet. + +Malformed HIR — an unknown node kind, a missing field, a foreign version — +fails pydantic validation (:class:`pydantic.ValidationError`) before +evaluation starts. Nodes the grammar allows but a deterministic constraint +cannot evaluate (distributions, UUID generation, ``Math.random()``) raise +:class:`HirEvaluationError`, as do a value of the wrong shape, an index that +is not an integer, and a ``Math`` call with the wrong number of arguments. +""" + +from __future__ import annotations + +import math +from collections.abc import Callable, Mapping, Sequence +from typing import TypeAlias, TypeVar + +from pydantic import TypeAdapter + +from . import models as m + +__all__ = [ + "HirEvaluationError", + "HirExpr", + "HirFunction", + "Scalar", + "Value", + "evaluate_hir", + "hir_margin", + "validate_hir_function", +] + +Scalar: TypeAlias = float | int | bool + +#: What HIR evaluates to: scalars and strings, plus the arrays and records +#: the metric surface reads from the simulation state. +Value: TypeAlias = float | int | bool | str | list["Value"] | dict[str, "Value"] + +#: One expression node, any kind. The members are the generated models. +HirExpr: TypeAlias = ( + m.HirNumberLit + | m.HirBoolLit + | m.HirStringLit + | m.HirStringCall + | m.HirUuidGenerate + | m.HirUuidFrom + | m.HirConstant + | m.HirLocalRef + | m.HirParamRef + | m.HirScenarioRef + | m.HirRangeCall + | m.HirFieldAccess + | m.HirIndexAccess + | m.HirLength + | m.HirUnary + | m.HirBinary + | m.HirCond + | m.HirLet + | m.HirMathCall + | m.HirRecordLit + | m.HirArrayLit + | m.HirArrayMap + | m.HirArrayReduce + | m.HirArrayConcat + | m.HirDistribution + | m.HirDistributionMap +) + +#: A lowered function: the generic shape, or one of the two surface-pinned +#: shapes a constraint carries. +HirFunction: TypeAlias = m.HirFunction | m.ParameterConstraintHir | m.StateConstraintHir + +_FUNCTION_ADAPTER = TypeAdapter(m.HirFunction) + +_MAX_RANGE_LENGTH = 1_000_000 + + +class HirEvaluationError(Exception): + """A serialized HIR expression could not be evaluated.""" + + +def validate_hir_function(fn: HirFunction | Mapping[str, object]) -> HirFunction: + """The function as a model: a mapping is validated against the grammar + (raising :class:`pydantic.ValidationError` when it does not fit), a + model passes through.""" + if isinstance(fn, (m.HirFunction, m.ParameterConstraintHir, m.StateConstraintHir)): + return fn + return _FUNCTION_ADAPTER.validate_python(fn) + + +# -- ECMAScript arithmetic ------------------------------------------------------ + + +def _js_round(value: float) -> float: + """ECMAScript ``Math.round``: half-up toward positive infinity (Python's + ``round`` is banker's).""" + if not math.isfinite(value): + return value # JS: round(±Infinity) is ±Infinity, round(NaN) is NaN + return math.floor(value + 0.5) + + +def _js_sign(value: float) -> float: + if value > 0: + return 1 + if value < 0: + return -1 + return value # preserves 0 / -0 / NaN like JS + + +def _is_odd_integer(value: float) -> bool: + return math.isfinite(value) and value == int(value) and int(value) % 2 == 1 + + +def _js_pow(base: float, exponent: float) -> float: + """ECMAScript exponentiation (``**`` and ``Math.pow``): IEEE-754 via + ``math.pow``, never raising and never going complex, with the spec's + deviations from C ``pow`` restored.""" + # JS: any NaN exponent, and ±Infinity exponents on a |base| of exactly + # 1, yield NaN where C pow returns 1. + if math.isnan(exponent) or (math.isinf(exponent) and abs(base) == 1): + return math.nan + try: + return math.pow(base, exponent) + except OverflowError: + # Finite operands overflowing a double: ±Infinity, negative only + # for a negative base raised to an odd integer. + negative = base < 0 and _is_odd_integer(exponent) + return -math.inf if negative else math.inf + except ValueError: + if base == 0 and exponent < 0: + # JS: 0 ** negative is Infinity; -0 flips the sign through an + # odd integer exponent. + negative = math.copysign(1.0, base) < 0 and _is_odd_integer(exponent) + return -math.inf if negative else math.inf + # Negative base with a non-integer exponent, and other domain + # errors: NaN in JS. + return math.nan + + +_Unary = Callable[[float], float] + + +def _js_log(fn: _Unary) -> _Unary: + """JS ``Math.log`` family: 0 yields -Infinity and negatives yield NaN, + where Python raises for both.""" + + def wrapped(value: float) -> float: + if value == 0: + return -math.inf + if value < 0: + return math.nan + return fn(value) + + return wrapped + + +def _js_grows(fn: _Unary, *, odd: bool) -> _Unary: + """JS ``Math.exp``/``cosh``/``sinh``: a result too large for a double is + ±Infinity, where Python raises OverflowError. ``odd`` functions take the + argument's sign.""" + + def wrapped(value: float) -> float: + try: + return fn(value) + except OverflowError: + return math.copysign(math.inf, value) if odd else math.inf + + return wrapped + + +def _js_integral(fn: Callable[[float], int]) -> _Unary: + """JS ``Math.ceil``/``floor``/``trunc`` pass non-finite values through, + where Python raises.""" + + def wrapped(value: float) -> float: + if not math.isfinite(value): + return value + return fn(value) + + return wrapped + + +def _cbrt(value: float) -> float: + return math.copysign(abs(value) ** (1 / 3), value) + + +def _max(*values: float) -> float: + return max(values) if values else -math.inf # JS: Math.max() is -Infinity + + +def _min(*values: float) -> float: + return min(values) if values else math.inf # JS: Math.min() is Infinity + + +_MATH_FNS: dict[str, Callable[..., float]] = { + "abs": abs, + "acos": math.acos, + "asin": math.asin, + "atan": math.atan, + "atan2": math.atan2, + "cbrt": _cbrt, + "ceil": _js_integral(math.ceil), + "cos": math.cos, + "cosh": _js_grows(math.cosh, odd=False), + "exp": _js_grows(math.exp, odd=False), + "floor": _js_integral(math.floor), + "hypot": math.hypot, + "log": _js_log(math.log), + "log10": _js_log(math.log10), + "log2": _js_log(math.log2), + "max": _max, + "min": _min, + "pow": _js_pow, + "round": _js_round, + "sign": _js_sign, + "sin": math.sin, + "sinh": _js_grows(math.sinh, odd=True), + "sqrt": math.sqrt, + "tan": math.tan, + "tanh": math.tanh, + "trunc": _js_integral(math.trunc), +} + +_CONSTANTS: dict[str, float] = { + "PI": math.pi, + "E": math.e, + "Infinity": math.inf, + "NaN": math.nan, +} + +_COMPARISONS = ("<", "<=", ">", ">=") + + +def _strict_slack(slack: float) -> float: + """A strict comparison is violated at the boundary, so its margin must + go negative there: a zero slack becomes the smallest representable step + below zero, keeping "margin >= 0 iff satisfied" exact for ``<``, ``>`` + and ``!=`` while staying negligible for any consumer of magnitudes.""" + if slack != 0.0: + return slack + return -math.ulp(1.0) + + +def _strict_equal(left: Value, right: Value) -> bool: + """ECMAScript strict equality on the value kinds HIR produces: booleans + never equal numbers (`1 === true` is false in JS, unlike Python).""" + if isinstance(left, bool) != isinstance(right, bool): + return False + return left == right + + +def _truthy(value: Value) -> bool: + return bool(value) + + +def _number(value: Value, context: str) -> float: + """The value as a JS number: booleans coerce to 0/1, an int too large for + a double becomes ±Infinity; anything else is a type error the frontend's + typechecker would have refused.""" + if isinstance(value, bool): + return float(value) + if isinstance(value, (int, float)): + try: + return float(value) + except OverflowError: + return math.inf if value > 0 else -math.inf + raise HirEvaluationError(f"{context} expects a number, got {type(value).__name__}") + + +_Ordered = TypeVar("_Ordered", float, str) + + +def _compare(op: str, left: _Ordered, right: _Ordered) -> bool: + if op == "<": + return left < right + if op == "<=": + return left <= right + if op == ">": + return left > right + return left >= right + + +def _range(args: Sequence[float]) -> list[Value]: + """The scenario ``range(...)`` helper, matching the TypeScript + implementation (Python-style bounds, fractional steps allowed).""" + if not 1 <= len(args) <= 3: + raise HirEvaluationError(f"range() takes 1 to 3 arguments, got {len(args)}") + for argument in args: + if not math.isfinite(argument): + raise HirEvaluationError("range() arguments must be finite numbers.") + start = args[0] if len(args) > 1 else 0 + end = args[1] if len(args) > 1 else args[0] + step = args[2] if len(args) > 2 else 1 + if step == 0: + raise HirEvaluationError("range() step must not be zero.") + span = (end - start) / step + if span == -math.inf: + # A backward range whose bounds overflow a double: TypeScript ceilings + # the span to 0 and yields nothing. + return [] + if not math.isfinite(span): + raise HirEvaluationError( + f"range() would produce more than {_MAX_RANGE_LENGTH} elements." + ) + maximum_length = max(0, math.ceil(span)) + if maximum_length > _MAX_RANGE_LENGTH: + raise HirEvaluationError( + f"range() would produce {maximum_length} elements, exceeding the " + f"limit of {_MAX_RANGE_LENGTH}." + ) + values: list[Value] = [] + for i in range(maximum_length): + value = start + i * step + if value >= end if step > 0 else value <= end: + break + values.append(value) + return values + + +# -- The walker ----------------------------------------------------------------- + + +class _Evaluator: + def __init__( + self, + scenario: Mapping[str, Scalar], + parameters: Mapping[str, Scalar], + locals_: dict[str, Value], + ) -> None: + self.scenario = scenario + self.parameters = parameters + self.locals = locals_ + + def eval(self, node: HirExpr) -> Value: + match node: + case m.HirNumberLit() | m.HirBoolLit() | m.HirStringLit(): + return node.value + case m.HirConstant(): + return _CONSTANTS[node.name.value] + case m.HirLocalRef(): + if node.name not in self.locals: + raise HirEvaluationError(f'Unbound local "{node.name}"') + return self.locals[node.name] + case m.HirParamRef(): + if node.name not in self.parameters: + raise HirEvaluationError(f'Unknown net parameter "{node.name}"') + return self.parameters[node.name] + case m.HirScenarioRef(): + if node.name not in self.scenario: + raise HirEvaluationError( + f'Unknown scenario parameter "{node.name}"' + ) + return self.scenario[node.name] + case m.HirRangeCall(): + return _range( + [_number(self.eval(argument), "range()") for argument in node.args] + ) + case m.HirFieldAccess(): + target = self.eval(node.target) + if not isinstance(target, Mapping) or node.field not in target: + raise HirEvaluationError(f'No field "{node.field}" on {target!r}') + return target[node.field] + case m.HirIndexAccess(): + target = self.eval(node.target) + position = _number(self.eval(node.index), "An index") + if not math.isfinite(position) or position != int(position): + raise HirEvaluationError(f"Index {position!r} is not an integer") + index = int(position) + if not isinstance(target, list) or not 0 <= index < len(target): + raise HirEvaluationError(f"Index {index} out of range") + return target[index] + case m.HirLength(): + target = self.eval(node.target) + if not isinstance(target, (list, str)): + raise HirEvaluationError(".length target is not an array or string") + return len(target) + case m.HirStringCall(): + return self._string_call(node) + case m.HirUnary(): + operand = self.eval(node.operand) + op = node.op.value + if op == "!": + return not _truthy(operand) + number = _number(operand, f'Unary "{op}"') + return -number if op == "-" else number + case m.HirBinary(): + return self._binary(node) + case m.HirCond(): + taken = ( + node.thenBranch + if _truthy(self.eval(node.condition)) + else node.elseBranch + ) + return self.eval(taken) + case m.HirLet(): + saved = dict(self.locals) + try: + for binding in node.bindings: + self.locals[binding.name] = self.eval(binding.value) + return self.eval(node.body) + finally: + self.locals = saved + case m.HirMathCall(): + fn = node.fn.value + if fn == "random": + raise HirEvaluationError( + "Math.random() is not evaluable in a constraint" + ) + args = [ + _number(self.eval(argument), f"Math.{fn}()") + for argument in node.args + ] + try: + return _MATH_FNS[fn](*args) + except (ValueError, OverflowError): + return math.nan + except TypeError as error: + raise HirEvaluationError( + f"Math.{fn}() called with {len(args)} argument(s)" + ) from error + case m.HirRecordLit(): + return {entry.key: self.eval(entry.value) for entry in node.entries} + case m.HirArrayLit(): + return [self.eval(element) for element in node.elements] + case m.HirArrayMap(): + return self._array_map(node) + case m.HirArrayReduce(): + return self._array_reduce(node) + case m.HirArrayConcat(): + left = self.eval(node.left) + right = self.eval(node.right) + if not isinstance(left, list) or not isinstance(right, list): + raise HirEvaluationError(".concat operands must be arrays") + return [*left, *right] + case ( + m.HirDistribution() + | m.HirDistributionMap() + | m.HirUuidGenerate() + | m.HirUuidFrom() + ): + raise HirEvaluationError( + f'HIR node kind "{node.kind}" is not evaluable in a constraint' + ) + + def _string_call(self, node: m.HirStringCall) -> Value: + target = self.eval(node.target) + argument = self.eval(node.argument) + fn = node.fn.value + if not isinstance(target, str) or not isinstance(argument, str): + raise HirEvaluationError(f".{fn}(...) is only available on strings") + if fn == "startsWith": + return target.startswith(argument) + if fn == "endsWith": + return target.endswith(argument) + return argument in target + + def _binary(self, node: m.HirBinary) -> Value: + op = node.op.value + left = self.eval(node.left) + if op == "&&": + return self.eval(node.right) if _truthy(left) else left + if op == "||": + return left if _truthy(left) else self.eval(node.right) + right = self.eval(node.right) + if op == "==": + return _strict_equal(left, right) + if op == "!=": + return not _strict_equal(left, right) + if isinstance(left, str) and isinstance(right, str): + if op == "+": + return left + right + if op in _COMPARISONS: + return _compare(op, left, right) + left_number = _number(left, f'"{op}"') + right_number = _number(right, f'"{op}"') + if op in _COMPARISONS: + return _compare(op, left_number, right_number) + if op == "+": + return left_number + right_number + if op == "-": + return left_number - right_number + if op == "*": + return left_number * right_number + if op == "/": + if right_number == 0: + # ECMAScript division never raises. + if left_number == 0: + return math.nan + return math.copysign(math.inf, left_number) * math.copysign( + 1, right_number + ) + return left_number / right_number + if op == "%": + if right_number == 0 or math.isinf(left_number): + return math.nan + # ECMAScript remainder takes the dividend's sign (math.fmod). + return math.fmod(left_number, right_number) + # `**`, through the JS-faithful pow: Python's `**` raises on overflow + # and 0**negative, and goes complex for a negative base with a + # fractional exponent, where JS yields ±Infinity / NaN. + return _js_pow(left_number, right_number) + + def _array_map(self, node: m.HirArrayMap) -> list[Value]: + target = self.eval(node.target) + if not isinstance(target, list): + raise HirEvaluationError(".map target is not an array") + out: list[Value] = [] + saved = dict(self.locals) + try: + for index, element in enumerate(target): + self.locals[node.param.name] = element + if node.indexParam is not None: + self.locals[node.indexParam.name] = index + out.append(self.eval(node.body)) + finally: + self.locals = saved + return out + + def _array_reduce(self, node: m.HirArrayReduce) -> Value: + target = self.eval(node.target) + if not isinstance(target, list): + raise HirEvaluationError(".reduce target is not an array") + accumulator = self.eval(node.initial) + saved = dict(self.locals) + try: + for index, element in enumerate(target): + self.locals[node.accParam.name] = accumulator + self.locals[node.param.name] = element + if node.indexParam is not None: + self.locals[node.indexParam.name] = index + accumulator = self.eval(node.body) + finally: + self.locals = saved + return accumulator + + # -- Signed margins ---------------------------------------------------- + + def margin(self, node: HirExpr) -> float: + """Robustness of a boolean expression: ``>= 0`` iff it evaluates to + ``True``, with magnitude measuring the distance to the boundary. + Comparisons yield signed slack; ``&&`` = ``min``, ``||`` = ``max``, + ``!`` negates; a plain boolean is ``±inf`` (no boundary to measure). + + A slack that comes out NaN never leaves a comparison, because + ``min``/``max`` would drop it by argument order and let a compound + read satisfied while the boolean is false. The leaf resolves to + ``±inf`` by asking the comparison itself: NaN operands make every + JS comparison false, but ``inf`` against ``inf`` also cancels to a + NaN slack while ``<=``, ``>=`` and ``==`` still hold. + + ``&&`` and ``||`` short-circuit exactly as evaluation does, so an + arm guarded by the one before it — a count checked before the token + it indexes — is never walked when evaluation would not walk it.""" + match node: + case m.HirBinary(): + margin = self._binary_margin(node) + if margin is not None: + return margin + case m.HirUnary() if node.op.value == "!": + return -self.margin(node.operand) + case m.HirCond(): + taken = ( + node.thenBranch + if _truthy(self.eval(node.condition)) + else node.elseBranch + ) + return self.margin(taken) + case m.HirLet(): + saved = dict(self.locals) + try: + for binding in node.bindings: + self.locals[binding.name] = self.eval(binding.value) + return self.margin(node.body) + finally: + self.locals = saved + case _: + pass + # A boolean leaf (literal, parameter, field): no boundary to measure. + value = self.eval(node) + if not isinstance(value, bool): + raise HirEvaluationError( + f"margin() requires a boolean expression, got a {type(value).__name__}" + ) + return math.inf if value else -math.inf + + def _binary_margin(self, node: m.HirBinary) -> float | None: + """The margin of a logical or comparison node; ``None`` for an + arithmetic operator, which is a leaf for :meth:`margin`.""" + op = node.op.value + if op == "&&": + left_margin = self.margin(node.left) + if left_margin < 0: + return left_margin + return min(left_margin, self.margin(node.right)) + if op == "||": + left_margin = self.margin(node.left) + if left_margin >= 0: + return left_margin + return max(left_margin, self.margin(node.right)) + if op in _COMPARISONS: + left_value = _number(self.eval(node.left), f'"{op}"') + right_value = _number(self.eval(node.right), f'"{op}"') + slack = ( + right_value - left_value + if op in ("<", "<=") + else left_value - right_value + ) + if math.isnan(slack): + satisfied = _compare(op, left_value, right_value) + return math.inf if satisfied else -math.inf + return slack if op in ("<=", ">=") else _strict_slack(slack) + if op in ("==", "!="): + left, right = self.eval(node.left), self.eval(node.right) + equal = _strict_equal(left, right) + wanted = equal if op == "==" else not equal + if ( + isinstance(left, bool) + or isinstance(right, bool) + or not isinstance(left, (int, float)) + or not isinstance(right, (int, float)) + ): + # Booleans, strings, and composites have no distance to + # measure: the boolean's sign is the whole answer. + return math.inf if wanted else -math.inf + distance = abs(float(left) - float(right)) + if math.isnan(distance): + return math.inf if wanted else -math.inf + return -distance if op == "==" else _strict_slack(distance) + return None + + +# -- Entry points ----------------------------------------------------------------- + + +def evaluate_hir( + fn: HirFunction | Mapping[str, object], + *, + scenario: Mapping[str, Scalar] | None = None, + parameters: Mapping[str, Scalar] | None = None, + locals_: Mapping[str, Value] | None = None, +) -> Value: + """Evaluate one serialized HIR function body. + + ``scenario`` binds ``scenario.`` reads, ``parameters`` binds + ``parameters.`` reads, ``locals_`` binds the function's declared + parameters (e.g. a metric-surface ``state`` record, as plain dicts and + lists). A mapping is validated against the grammar first. + """ + function = validate_hir_function(fn) + return _Evaluator(scenario or {}, parameters or {}, dict(locals_ or {})).eval( + function.body + ) + + +def hir_margin( + fn: HirFunction | Mapping[str, object], + *, + scenario: Mapping[str, Scalar] | None = None, + parameters: Mapping[str, Scalar] | None = None, + locals_: Mapping[str, Value] | None = None, +) -> float: + """The signed robustness margin of one boolean HIR function: ``>= 0`` + iff :func:`evaluate_hir` would return ``True``. Bindings as for + :func:`evaluate_hir`.""" + function = validate_hir_function(fn) + return _Evaluator(scenario or {}, parameters or {}, dict(locals_ or {})).margin( + function.body + ) diff --git a/libs/@local/petrinaut-python/src/petrinaut/models.py b/libs/@local/petrinaut-python/src/petrinaut/models.py index cb95abfdcf4..1944b03dea7 100644 --- a/libs/@local/petrinaut-python/src/petrinaut/models.py +++ b/libs/@local/petrinaut-python/src/petrinaut/models.py @@ -4,7 +4,7 @@ from __future__ import annotations from enum import Enum -from typing import Literal +from typing import Annotated, Literal from pydantic import BaseModel, ConfigDict, Field @@ -76,6 +76,191 @@ class Study(BaseModel): seedsPerTrial: int | None = Field(None, ge=1, le=100) +class OptimizationEvaluateResult(BaseModel): + model_config = ConfigDict( + extra="forbid", + ) + objective: float + replicates: list[OptimizationReplicate] | None = None + + +class HirSpan(BaseModel): + model_config = ConfigDict( + extra="forbid", + ) + start: int = Field(..., ge=0, le=9007199254740991) + length: int = Field(..., ge=0, le=9007199254740991) + + +class HirNumberLit(BaseModel): + model_config = ConfigDict( + extra="forbid", + ) + id: int = Field(..., ge=0, le=9007199254740991) + span: HirSpan + kind: Literal["numberLit"] + value: float + raw: str + + +class HirBoolLit(BaseModel): + model_config = ConfigDict( + extra="forbid", + ) + id: int = Field(..., ge=0, le=9007199254740991) + span: HirSpan + kind: Literal["boolLit"] + value: bool + + +class HirStringLit(BaseModel): + model_config = ConfigDict( + extra="forbid", + ) + id: int = Field(..., ge=0, le=9007199254740991) + span: HirSpan + kind: Literal["stringLit"] + value: str + + +class HirStringFn(Enum): + startsWith = "startsWith" + endsWith = "endsWith" + includes = "includes" + + +class HirUuidGenerate(BaseModel): + model_config = ConfigDict( + extra="forbid", + ) + id: int = Field(..., ge=0, le=9007199254740991) + span: HirSpan + kind: Literal["uuidGenerate"] + + +class HirConstantName(Enum): + PI = "PI" + E = "E" + Infinity = "Infinity" + NaN = "NaN" + + +class HirLocalRef(BaseModel): + model_config = ConfigDict( + extra="forbid", + ) + id: int = Field(..., ge=0, le=9007199254740991) + span: HirSpan + kind: Literal["localRef"] + name: str + + +class HirParamRef(BaseModel): + model_config = ConfigDict( + extra="forbid", + ) + id: int = Field(..., ge=0, le=9007199254740991) + span: HirSpan + kind: Literal["paramRef"] + name: str + + +class HirScenarioRef(BaseModel): + model_config = ConfigDict( + extra="forbid", + ) + id: int = Field(..., ge=0, le=9007199254740991) + span: HirSpan + kind: Literal["scenarioRef"] + name: str + + +class HirUnaryOp(Enum): + field_ = "-" + field__1 = "+" + field__2 = "!" + + +class HirBinaryOp(Enum): + field_ = "+" + field__1 = "-" + field__2 = "*" + field__3 = "/" + field__4 = "%" + field__ = "**" + field__5 = "<" + field___1 = "<=" + field__6 = ">" + field___2 = ">=" + field___3 = "==" + field___4 = "!=" + field___5 = "&&" + field___6 = "||" + + +class HirMathFn(Enum): + abs = "abs" + acos = "acos" + asin = "asin" + atan = "atan" + atan2 = "atan2" + cbrt = "cbrt" + ceil = "ceil" + cos = "cos" + cosh = "cosh" + exp = "exp" + floor = "floor" + hypot = "hypot" + log = "log" + log10 = "log10" + log2 = "log2" + max = "max" + min = "min" + pow = "pow" + random = "random" + round = "round" + sign = "sign" + sin = "sin" + sinh = "sinh" + sqrt = "sqrt" + tan = "tan" + tanh = "tanh" + trunc = "trunc" + + +class HirDistributionKind(Enum): + gaussian = "gaussian" + uniform = "uniform" + lognormal = "lognormal" + + +class HirSurfaceKind(Enum): + dynamics = "dynamics" + lambda_ = "lambda" + kernel = "kernel" + metric = "metric" + scenario_expression = "scenario-expression" + scenario_code = "scenario-code" + + +class HirNamedSpan(BaseModel): + model_config = ConfigDict( + extra="forbid", + ) + name: str + span: HirSpan + + +class HirConstant(BaseModel): + model_config = ConfigDict( + extra="forbid", + ) + id: int = Field(..., ge=0, le=9007199254740991) + span: HirSpan + kind: Literal["constant"] + name: HirConstantName + + class OptimizationDescribeResult(BaseModel): model_config = ConfigDict( extra="forbid", @@ -90,11 +275,1321 @@ class OptimizationDescribeResult(BaseModel): | OptimizationIntParameter | OptimizationBooleanParameter ] + constraints: ( + list[ + Annotated[ + ParameterConstraint | StateConstraint, Field(discriminator="space") + ] + ] + | None + ) = Field( + None, + description="The manifest's constraints, passed through verbatim so protocol clients (the Python binding) can evaluate their HIR. Absent means unconstrained.", + ) -class OptimizationEvaluateResult(BaseModel): +class ParameterConstraint(BaseModel): model_config = ConfigDict( extra="forbid", ) - objective: float - replicates: list[OptimizationReplicate] | None = None + space: Literal["parameters"] + id: str = Field(..., min_length=1) + name: str | None = Field( + None, + description="Optional display name shown wherever the constraint is reported.", + min_length=1, + ) + code: str = Field( + ..., + description="The authored TypeScript source, the editable text of record. `hir` is its lowered form; regenerating `hir` from `code` must be a no-op.", + min_length=1, + ) + hir: ParameterConstraintHir + + +class ParameterConstraintHir(BaseModel): + model_config = ConfigDict( + extra="forbid", + ) + hirVersion: Literal[1] + surface: Literal["scenario-expression"] + params: list[HirNamedSpan] + body: ( + HirNumberLit + | HirBoolLit + | HirStringLit + | HirStringCall + | HirUuidGenerate + | HirUuidFrom + | HirConstant + | HirLocalRef + | HirParamRef + | HirScenarioRef + | HirRangeCall + | HirFieldAccess + | HirIndexAccess + | HirLength + | HirUnary + | HirBinary + | HirCond + | HirLet + | HirMathCall + | HirRecordLit + | HirArrayLit + | HirArrayMap + | HirArrayReduce + | HirArrayConcat + | HirDistribution + | HirDistributionMap + ) = Field( + ..., + description="One HIR expression node, discriminated by `kind`. Evaluators must reject kinds they do not know.", + discriminator="kind", + ) + span: HirSpan + + +class HirStringCall(BaseModel): + model_config = ConfigDict( + extra="forbid", + ) + id: int = Field(..., ge=0, le=9007199254740991) + span: HirSpan + kind: Literal["stringCall"] + fn: HirStringFn + target: ( + HirNumberLit + | HirBoolLit + | HirStringLit + | HirStringCall + | HirUuidGenerate + | HirUuidFrom + | HirConstant + | HirLocalRef + | HirParamRef + | HirScenarioRef + | HirRangeCall + | HirFieldAccess + | HirIndexAccess + | HirLength + | HirUnary + | HirBinary + | HirCond + | HirLet + | HirMathCall + | HirRecordLit + | HirArrayLit + | HirArrayMap + | HirArrayReduce + | HirArrayConcat + | HirDistribution + | HirDistributionMap + ) = Field( + ..., + description="One HIR expression node, discriminated by `kind`. Evaluators must reject kinds they do not know.", + discriminator="kind", + ) + argument: ( + HirNumberLit + | HirBoolLit + | HirStringLit + | HirStringCall + | HirUuidGenerate + | HirUuidFrom + | HirConstant + | HirLocalRef + | HirParamRef + | HirScenarioRef + | HirRangeCall + | HirFieldAccess + | HirIndexAccess + | HirLength + | HirUnary + | HirBinary + | HirCond + | HirLet + | HirMathCall + | HirRecordLit + | HirArrayLit + | HirArrayMap + | HirArrayReduce + | HirArrayConcat + | HirDistribution + | HirDistributionMap + ) = Field( + ..., + description="One HIR expression node, discriminated by `kind`. Evaluators must reject kinds they do not know.", + discriminator="kind", + ) + + +class HirUuidFrom(BaseModel): + model_config = ConfigDict( + extra="forbid", + ) + id: int = Field(..., ge=0, le=9007199254740991) + span: HirSpan + kind: Literal["uuidFrom"] + operand: ( + HirNumberLit + | HirBoolLit + | HirStringLit + | HirStringCall + | HirUuidGenerate + | HirUuidFrom + | HirConstant + | HirLocalRef + | HirParamRef + | HirScenarioRef + | HirRangeCall + | HirFieldAccess + | HirIndexAccess + | HirLength + | HirUnary + | HirBinary + | HirCond + | HirLet + | HirMathCall + | HirRecordLit + | HirArrayLit + | HirArrayMap + | HirArrayReduce + | HirArrayConcat + | HirDistribution + | HirDistributionMap + ) = Field( + ..., + description="One HIR expression node, discriminated by `kind`. Evaluators must reject kinds they do not know.", + discriminator="kind", + ) + + +class HirRangeCall(BaseModel): + model_config = ConfigDict( + extra="forbid", + ) + id: int = Field(..., ge=0, le=9007199254740991) + span: HirSpan + kind: Literal["rangeCall"] + args: list[ + Annotated[ + HirNumberLit + | HirBoolLit + | HirStringLit + | HirStringCall + | HirUuidGenerate + | HirUuidFrom + | HirConstant + | HirLocalRef + | HirParamRef + | HirScenarioRef + | HirRangeCall + | HirFieldAccess + | HirIndexAccess + | HirLength + | HirUnary + | HirBinary + | HirCond + | HirLet + | HirMathCall + | HirRecordLit + | HirArrayLit + | HirArrayMap + | HirArrayReduce + | HirArrayConcat + | HirDistribution + | HirDistributionMap, + Field(discriminator="kind"), + ] + ] + + +class HirFieldAccess(BaseModel): + model_config = ConfigDict( + extra="forbid", + ) + id: int = Field(..., ge=0, le=9007199254740991) + span: HirSpan + kind: Literal["fieldAccess"] + target: ( + HirNumberLit + | HirBoolLit + | HirStringLit + | HirStringCall + | HirUuidGenerate + | HirUuidFrom + | HirConstant + | HirLocalRef + | HirParamRef + | HirScenarioRef + | HirRangeCall + | HirFieldAccess + | HirIndexAccess + | HirLength + | HirUnary + | HirBinary + | HirCond + | HirLet + | HirMathCall + | HirRecordLit + | HirArrayLit + | HirArrayMap + | HirArrayReduce + | HirArrayConcat + | HirDistribution + | HirDistributionMap + ) = Field( + ..., + description="One HIR expression node, discriminated by `kind`. Evaluators must reject kinds they do not know.", + discriminator="kind", + ) + field: str + fieldSpan: HirSpan + + +class HirIndexAccess(BaseModel): + model_config = ConfigDict( + extra="forbid", + ) + id: int = Field(..., ge=0, le=9007199254740991) + span: HirSpan + kind: Literal["indexAccess"] + target: ( + HirNumberLit + | HirBoolLit + | HirStringLit + | HirStringCall + | HirUuidGenerate + | HirUuidFrom + | HirConstant + | HirLocalRef + | HirParamRef + | HirScenarioRef + | HirRangeCall + | HirFieldAccess + | HirIndexAccess + | HirLength + | HirUnary + | HirBinary + | HirCond + | HirLet + | HirMathCall + | HirRecordLit + | HirArrayLit + | HirArrayMap + | HirArrayReduce + | HirArrayConcat + | HirDistribution + | HirDistributionMap + ) = Field( + ..., + description="One HIR expression node, discriminated by `kind`. Evaluators must reject kinds they do not know.", + discriminator="kind", + ) + index: ( + HirNumberLit + | HirBoolLit + | HirStringLit + | HirStringCall + | HirUuidGenerate + | HirUuidFrom + | HirConstant + | HirLocalRef + | HirParamRef + | HirScenarioRef + | HirRangeCall + | HirFieldAccess + | HirIndexAccess + | HirLength + | HirUnary + | HirBinary + | HirCond + | HirLet + | HirMathCall + | HirRecordLit + | HirArrayLit + | HirArrayMap + | HirArrayReduce + | HirArrayConcat + | HirDistribution + | HirDistributionMap + ) = Field( + ..., + description="One HIR expression node, discriminated by `kind`. Evaluators must reject kinds they do not know.", + discriminator="kind", + ) + + +class HirLength(BaseModel): + model_config = ConfigDict( + extra="forbid", + ) + id: int = Field(..., ge=0, le=9007199254740991) + span: HirSpan + kind: Literal["length"] + target: ( + HirNumberLit + | HirBoolLit + | HirStringLit + | HirStringCall + | HirUuidGenerate + | HirUuidFrom + | HirConstant + | HirLocalRef + | HirParamRef + | HirScenarioRef + | HirRangeCall + | HirFieldAccess + | HirIndexAccess + | HirLength + | HirUnary + | HirBinary + | HirCond + | HirLet + | HirMathCall + | HirRecordLit + | HirArrayLit + | HirArrayMap + | HirArrayReduce + | HirArrayConcat + | HirDistribution + | HirDistributionMap + ) = Field( + ..., + description="One HIR expression node, discriminated by `kind`. Evaluators must reject kinds they do not know.", + discriminator="kind", + ) + + +class HirUnary(BaseModel): + model_config = ConfigDict( + extra="forbid", + ) + id: int = Field(..., ge=0, le=9007199254740991) + span: HirSpan + kind: Literal["unary"] + op: HirUnaryOp + operand: ( + HirNumberLit + | HirBoolLit + | HirStringLit + | HirStringCall + | HirUuidGenerate + | HirUuidFrom + | HirConstant + | HirLocalRef + | HirParamRef + | HirScenarioRef + | HirRangeCall + | HirFieldAccess + | HirIndexAccess + | HirLength + | HirUnary + | HirBinary + | HirCond + | HirLet + | HirMathCall + | HirRecordLit + | HirArrayLit + | HirArrayMap + | HirArrayReduce + | HirArrayConcat + | HirDistribution + | HirDistributionMap + ) = Field( + ..., + description="One HIR expression node, discriminated by `kind`. Evaluators must reject kinds they do not know.", + discriminator="kind", + ) + + +class HirBinary(BaseModel): + model_config = ConfigDict( + extra="forbid", + ) + id: int = Field(..., ge=0, le=9007199254740991) + span: HirSpan + kind: Literal["binary"] + op: HirBinaryOp + left: ( + HirNumberLit + | HirBoolLit + | HirStringLit + | HirStringCall + | HirUuidGenerate + | HirUuidFrom + | HirConstant + | HirLocalRef + | HirParamRef + | HirScenarioRef + | HirRangeCall + | HirFieldAccess + | HirIndexAccess + | HirLength + | HirUnary + | HirBinary + | HirCond + | HirLet + | HirMathCall + | HirRecordLit + | HirArrayLit + | HirArrayMap + | HirArrayReduce + | HirArrayConcat + | HirDistribution + | HirDistributionMap + ) = Field( + ..., + description="One HIR expression node, discriminated by `kind`. Evaluators must reject kinds they do not know.", + discriminator="kind", + ) + right: ( + HirNumberLit + | HirBoolLit + | HirStringLit + | HirStringCall + | HirUuidGenerate + | HirUuidFrom + | HirConstant + | HirLocalRef + | HirParamRef + | HirScenarioRef + | HirRangeCall + | HirFieldAccess + | HirIndexAccess + | HirLength + | HirUnary + | HirBinary + | HirCond + | HirLet + | HirMathCall + | HirRecordLit + | HirArrayLit + | HirArrayMap + | HirArrayReduce + | HirArrayConcat + | HirDistribution + | HirDistributionMap + ) = Field( + ..., + description="One HIR expression node, discriminated by `kind`. Evaluators must reject kinds they do not know.", + discriminator="kind", + ) + + +class HirCond(BaseModel): + model_config = ConfigDict( + extra="forbid", + ) + id: int = Field(..., ge=0, le=9007199254740991) + span: HirSpan + kind: Literal["cond"] + condition: ( + HirNumberLit + | HirBoolLit + | HirStringLit + | HirStringCall + | HirUuidGenerate + | HirUuidFrom + | HirConstant + | HirLocalRef + | HirParamRef + | HirScenarioRef + | HirRangeCall + | HirFieldAccess + | HirIndexAccess + | HirLength + | HirUnary + | HirBinary + | HirCond + | HirLet + | HirMathCall + | HirRecordLit + | HirArrayLit + | HirArrayMap + | HirArrayReduce + | HirArrayConcat + | HirDistribution + | HirDistributionMap + ) = Field( + ..., + description="One HIR expression node, discriminated by `kind`. Evaluators must reject kinds they do not know.", + discriminator="kind", + ) + thenBranch: ( + HirNumberLit + | HirBoolLit + | HirStringLit + | HirStringCall + | HirUuidGenerate + | HirUuidFrom + | HirConstant + | HirLocalRef + | HirParamRef + | HirScenarioRef + | HirRangeCall + | HirFieldAccess + | HirIndexAccess + | HirLength + | HirUnary + | HirBinary + | HirCond + | HirLet + | HirMathCall + | HirRecordLit + | HirArrayLit + | HirArrayMap + | HirArrayReduce + | HirArrayConcat + | HirDistribution + | HirDistributionMap + ) = Field( + ..., + description="One HIR expression node, discriminated by `kind`. Evaluators must reject kinds they do not know.", + discriminator="kind", + ) + elseBranch: ( + HirNumberLit + | HirBoolLit + | HirStringLit + | HirStringCall + | HirUuidGenerate + | HirUuidFrom + | HirConstant + | HirLocalRef + | HirParamRef + | HirScenarioRef + | HirRangeCall + | HirFieldAccess + | HirIndexAccess + | HirLength + | HirUnary + | HirBinary + | HirCond + | HirLet + | HirMathCall + | HirRecordLit + | HirArrayLit + | HirArrayMap + | HirArrayReduce + | HirArrayConcat + | HirDistribution + | HirDistributionMap + ) = Field( + ..., + description="One HIR expression node, discriminated by `kind`. Evaluators must reject kinds they do not know.", + discriminator="kind", + ) + + +class HirLet(BaseModel): + model_config = ConfigDict( + extra="forbid", + ) + id: int = Field(..., ge=0, le=9007199254740991) + span: HirSpan + kind: Literal["let"] + bindings: list[HirLetBinding] + body: ( + HirNumberLit + | HirBoolLit + | HirStringLit + | HirStringCall + | HirUuidGenerate + | HirUuidFrom + | HirConstant + | HirLocalRef + | HirParamRef + | HirScenarioRef + | HirRangeCall + | HirFieldAccess + | HirIndexAccess + | HirLength + | HirUnary + | HirBinary + | HirCond + | HirLet + | HirMathCall + | HirRecordLit + | HirArrayLit + | HirArrayMap + | HirArrayReduce + | HirArrayConcat + | HirDistribution + | HirDistributionMap + ) = Field( + ..., + description="One HIR expression node, discriminated by `kind`. Evaluators must reject kinds they do not know.", + discriminator="kind", + ) + + +class HirLetBinding(BaseModel): + model_config = ConfigDict( + extra="forbid", + ) + name: str + nameSpan: HirSpan + value: ( + HirNumberLit + | HirBoolLit + | HirStringLit + | HirStringCall + | HirUuidGenerate + | HirUuidFrom + | HirConstant + | HirLocalRef + | HirParamRef + | HirScenarioRef + | HirRangeCall + | HirFieldAccess + | HirIndexAccess + | HirLength + | HirUnary + | HirBinary + | HirCond + | HirLet + | HirMathCall + | HirRecordLit + | HirArrayLit + | HirArrayMap + | HirArrayReduce + | HirArrayConcat + | HirDistribution + | HirDistributionMap + ) = Field( + ..., + description="One HIR expression node, discriminated by `kind`. Evaluators must reject kinds they do not know.", + discriminator="kind", + ) + + +class HirMathCall(BaseModel): + model_config = ConfigDict( + extra="forbid", + ) + id: int = Field(..., ge=0, le=9007199254740991) + span: HirSpan + kind: Literal["mathCall"] + fn: HirMathFn + args: list[ + Annotated[ + HirNumberLit + | HirBoolLit + | HirStringLit + | HirStringCall + | HirUuidGenerate + | HirUuidFrom + | HirConstant + | HirLocalRef + | HirParamRef + | HirScenarioRef + | HirRangeCall + | HirFieldAccess + | HirIndexAccess + | HirLength + | HirUnary + | HirBinary + | HirCond + | HirLet + | HirMathCall + | HirRecordLit + | HirArrayLit + | HirArrayMap + | HirArrayReduce + | HirArrayConcat + | HirDistribution + | HirDistributionMap, + Field(discriminator="kind"), + ] + ] + + +class HirRecordLit(BaseModel): + model_config = ConfigDict( + extra="forbid", + ) + id: int = Field(..., ge=0, le=9007199254740991) + span: HirSpan + kind: Literal["recordLit"] + entries: list[HirRecordEntry] + + +class HirRecordEntry(BaseModel): + model_config = ConfigDict( + extra="forbid", + ) + key: str + keySpan: HirSpan + value: ( + HirNumberLit + | HirBoolLit + | HirStringLit + | HirStringCall + | HirUuidGenerate + | HirUuidFrom + | HirConstant + | HirLocalRef + | HirParamRef + | HirScenarioRef + | HirRangeCall + | HirFieldAccess + | HirIndexAccess + | HirLength + | HirUnary + | HirBinary + | HirCond + | HirLet + | HirMathCall + | HirRecordLit + | HirArrayLit + | HirArrayMap + | HirArrayReduce + | HirArrayConcat + | HirDistribution + | HirDistributionMap + ) = Field( + ..., + description="One HIR expression node, discriminated by `kind`. Evaluators must reject kinds they do not know.", + discriminator="kind", + ) + + +class HirArrayLit(BaseModel): + model_config = ConfigDict( + extra="forbid", + ) + id: int = Field(..., ge=0, le=9007199254740991) + span: HirSpan + kind: Literal["arrayLit"] + elements: list[ + Annotated[ + HirNumberLit + | HirBoolLit + | HirStringLit + | HirStringCall + | HirUuidGenerate + | HirUuidFrom + | HirConstant + | HirLocalRef + | HirParamRef + | HirScenarioRef + | HirRangeCall + | HirFieldAccess + | HirIndexAccess + | HirLength + | HirUnary + | HirBinary + | HirCond + | HirLet + | HirMathCall + | HirRecordLit + | HirArrayLit + | HirArrayMap + | HirArrayReduce + | HirArrayConcat + | HirDistribution + | HirDistributionMap, + Field(discriminator="kind"), + ] + ] + + +class HirArrayMap(BaseModel): + model_config = ConfigDict( + extra="forbid", + ) + id: int = Field(..., ge=0, le=9007199254740991) + span: HirSpan + kind: Literal["arrayMap"] + target: ( + HirNumberLit + | HirBoolLit + | HirStringLit + | HirStringCall + | HirUuidGenerate + | HirUuidFrom + | HirConstant + | HirLocalRef + | HirParamRef + | HirScenarioRef + | HirRangeCall + | HirFieldAccess + | HirIndexAccess + | HirLength + | HirUnary + | HirBinary + | HirCond + | HirLet + | HirMathCall + | HirRecordLit + | HirArrayLit + | HirArrayMap + | HirArrayReduce + | HirArrayConcat + | HirDistribution + | HirDistributionMap + ) = Field( + ..., + description="One HIR expression node, discriminated by `kind`. Evaluators must reject kinds they do not know.", + discriminator="kind", + ) + param: HirNamedSpan + indexParam: HirNamedSpan | None = None + body: ( + HirNumberLit + | HirBoolLit + | HirStringLit + | HirStringCall + | HirUuidGenerate + | HirUuidFrom + | HirConstant + | HirLocalRef + | HirParamRef + | HirScenarioRef + | HirRangeCall + | HirFieldAccess + | HirIndexAccess + | HirLength + | HirUnary + | HirBinary + | HirCond + | HirLet + | HirMathCall + | HirRecordLit + | HirArrayLit + | HirArrayMap + | HirArrayReduce + | HirArrayConcat + | HirDistribution + | HirDistributionMap + ) = Field( + ..., + description="One HIR expression node, discriminated by `kind`. Evaluators must reject kinds they do not know.", + discriminator="kind", + ) + + +class HirArrayReduce(BaseModel): + model_config = ConfigDict( + extra="forbid", + ) + id: int = Field(..., ge=0, le=9007199254740991) + span: HirSpan + kind: Literal["arrayReduce"] + target: ( + HirNumberLit + | HirBoolLit + | HirStringLit + | HirStringCall + | HirUuidGenerate + | HirUuidFrom + | HirConstant + | HirLocalRef + | HirParamRef + | HirScenarioRef + | HirRangeCall + | HirFieldAccess + | HirIndexAccess + | HirLength + | HirUnary + | HirBinary + | HirCond + | HirLet + | HirMathCall + | HirRecordLit + | HirArrayLit + | HirArrayMap + | HirArrayReduce + | HirArrayConcat + | HirDistribution + | HirDistributionMap + ) = Field( + ..., + description="One HIR expression node, discriminated by `kind`. Evaluators must reject kinds they do not know.", + discriminator="kind", + ) + accParam: HirNamedSpan + param: HirNamedSpan + indexParam: HirNamedSpan | None = None + body: ( + HirNumberLit + | HirBoolLit + | HirStringLit + | HirStringCall + | HirUuidGenerate + | HirUuidFrom + | HirConstant + | HirLocalRef + | HirParamRef + | HirScenarioRef + | HirRangeCall + | HirFieldAccess + | HirIndexAccess + | HirLength + | HirUnary + | HirBinary + | HirCond + | HirLet + | HirMathCall + | HirRecordLit + | HirArrayLit + | HirArrayMap + | HirArrayReduce + | HirArrayConcat + | HirDistribution + | HirDistributionMap + ) = Field( + ..., + description="One HIR expression node, discriminated by `kind`. Evaluators must reject kinds they do not know.", + discriminator="kind", + ) + initial: ( + HirNumberLit + | HirBoolLit + | HirStringLit + | HirStringCall + | HirUuidGenerate + | HirUuidFrom + | HirConstant + | HirLocalRef + | HirParamRef + | HirScenarioRef + | HirRangeCall + | HirFieldAccess + | HirIndexAccess + | HirLength + | HirUnary + | HirBinary + | HirCond + | HirLet + | HirMathCall + | HirRecordLit + | HirArrayLit + | HirArrayMap + | HirArrayReduce + | HirArrayConcat + | HirDistribution + | HirDistributionMap + ) = Field( + ..., + description="One HIR expression node, discriminated by `kind`. Evaluators must reject kinds they do not know.", + discriminator="kind", + ) + + +class HirArrayConcat(BaseModel): + model_config = ConfigDict( + extra="forbid", + ) + id: int = Field(..., ge=0, le=9007199254740991) + span: HirSpan + kind: Literal["arrayConcat"] + left: ( + HirNumberLit + | HirBoolLit + | HirStringLit + | HirStringCall + | HirUuidGenerate + | HirUuidFrom + | HirConstant + | HirLocalRef + | HirParamRef + | HirScenarioRef + | HirRangeCall + | HirFieldAccess + | HirIndexAccess + | HirLength + | HirUnary + | HirBinary + | HirCond + | HirLet + | HirMathCall + | HirRecordLit + | HirArrayLit + | HirArrayMap + | HirArrayReduce + | HirArrayConcat + | HirDistribution + | HirDistributionMap + ) = Field( + ..., + description="One HIR expression node, discriminated by `kind`. Evaluators must reject kinds they do not know.", + discriminator="kind", + ) + right: ( + HirNumberLit + | HirBoolLit + | HirStringLit + | HirStringCall + | HirUuidGenerate + | HirUuidFrom + | HirConstant + | HirLocalRef + | HirParamRef + | HirScenarioRef + | HirRangeCall + | HirFieldAccess + | HirIndexAccess + | HirLength + | HirUnary + | HirBinary + | HirCond + | HirLet + | HirMathCall + | HirRecordLit + | HirArrayLit + | HirArrayMap + | HirArrayReduce + | HirArrayConcat + | HirDistribution + | HirDistributionMap + ) = Field( + ..., + description="One HIR expression node, discriminated by `kind`. Evaluators must reject kinds they do not know.", + discriminator="kind", + ) + + +class HirDistribution(BaseModel): + model_config = ConfigDict( + extra="forbid", + ) + id: int = Field(..., ge=0, le=9007199254740991) + span: HirSpan + kind: Literal["distribution"] + dist: HirDistributionKind + args: list[ + Annotated[ + HirNumberLit + | HirBoolLit + | HirStringLit + | HirStringCall + | HirUuidGenerate + | HirUuidFrom + | HirConstant + | HirLocalRef + | HirParamRef + | HirScenarioRef + | HirRangeCall + | HirFieldAccess + | HirIndexAccess + | HirLength + | HirUnary + | HirBinary + | HirCond + | HirLet + | HirMathCall + | HirRecordLit + | HirArrayLit + | HirArrayMap + | HirArrayReduce + | HirArrayConcat + | HirDistribution + | HirDistributionMap, + Field(discriminator="kind"), + ] + ] + + +class HirDistributionMap(BaseModel): + model_config = ConfigDict( + extra="forbid", + ) + id: int = Field(..., ge=0, le=9007199254740991) + span: HirSpan + kind: Literal["distributionMap"] + base: ( + HirNumberLit + | HirBoolLit + | HirStringLit + | HirStringCall + | HirUuidGenerate + | HirUuidFrom + | HirConstant + | HirLocalRef + | HirParamRef + | HirScenarioRef + | HirRangeCall + | HirFieldAccess + | HirIndexAccess + | HirLength + | HirUnary + | HirBinary + | HirCond + | HirLet + | HirMathCall + | HirRecordLit + | HirArrayLit + | HirArrayMap + | HirArrayReduce + | HirArrayConcat + | HirDistribution + | HirDistributionMap + ) = Field( + ..., + description="One HIR expression node, discriminated by `kind`. Evaluators must reject kinds they do not know.", + discriminator="kind", + ) + param: HirNamedSpan + body: ( + HirNumberLit + | HirBoolLit + | HirStringLit + | HirStringCall + | HirUuidGenerate + | HirUuidFrom + | HirConstant + | HirLocalRef + | HirParamRef + | HirScenarioRef + | HirRangeCall + | HirFieldAccess + | HirIndexAccess + | HirLength + | HirUnary + | HirBinary + | HirCond + | HirLet + | HirMathCall + | HirRecordLit + | HirArrayLit + | HirArrayMap + | HirArrayReduce + | HirArrayConcat + | HirDistribution + | HirDistributionMap + ) = Field( + ..., + description="One HIR expression node, discriminated by `kind`. Evaluators must reject kinds they do not know.", + discriminator="kind", + ) + + +class StateConstraint(BaseModel): + model_config = ConfigDict( + extra="forbid", + ) + space: Literal["state"] + id: str = Field(..., min_length=1) + name: str | None = Field( + None, + description="Optional display name shown wherever the constraint is reported.", + min_length=1, + ) + code: str = Field( + ..., + description="The authored TypeScript source, the editable text of record. `hir` is its lowered form; regenerating `hir` from `code` must be a no-op.", + min_length=1, + ) + hir: StateConstraintHir + + +class StateConstraintHir(BaseModel): + model_config = ConfigDict( + extra="forbid", + ) + hirVersion: Literal[1] + surface: Literal["metric"] + params: list[HirNamedSpan] + body: ( + HirNumberLit + | HirBoolLit + | HirStringLit + | HirStringCall + | HirUuidGenerate + | HirUuidFrom + | HirConstant + | HirLocalRef + | HirParamRef + | HirScenarioRef + | HirRangeCall + | HirFieldAccess + | HirIndexAccess + | HirLength + | HirUnary + | HirBinary + | HirCond + | HirLet + | HirMathCall + | HirRecordLit + | HirArrayLit + | HirArrayMap + | HirArrayReduce + | HirArrayConcat + | HirDistribution + | HirDistributionMap + ) = Field( + ..., + description="One HIR expression node, discriminated by `kind`. Evaluators must reject kinds they do not know.", + discriminator="kind", + ) + span: HirSpan + + +class HirFunction(BaseModel): + model_config = ConfigDict( + extra="forbid", + ) + hirVersion: Literal[1] + surface: HirSurfaceKind + params: list[HirNamedSpan] + body: ( + HirNumberLit + | HirBoolLit + | HirStringLit + | HirStringCall + | HirUuidGenerate + | HirUuidFrom + | HirConstant + | HirLocalRef + | HirParamRef + | HirScenarioRef + | HirRangeCall + | HirFieldAccess + | HirIndexAccess + | HirLength + | HirUnary + | HirBinary + | HirCond + | HirLet + | HirMathCall + | HirRecordLit + | HirArrayLit + | HirArrayMap + | HirArrayReduce + | HirArrayConcat + | HirDistribution + | HirDistributionMap + ) = Field( + ..., + description="One HIR expression node, discriminated by `kind`. Evaluators must reject kinds they do not know.", + discriminator="kind", + ) + span: HirSpan + + +OptimizationDescribeResult.model_rebuild() +ParameterConstraint.model_rebuild() +ParameterConstraintHir.model_rebuild() +HirStringCall.model_rebuild() +HirUuidFrom.model_rebuild() +HirRangeCall.model_rebuild() +HirFieldAccess.model_rebuild() +HirIndexAccess.model_rebuild() +HirLength.model_rebuild() +HirUnary.model_rebuild() +HirBinary.model_rebuild() +HirCond.model_rebuild() +HirLet.model_rebuild() +HirLetBinding.model_rebuild() +HirMathCall.model_rebuild() +HirRecordLit.model_rebuild() +HirRecordEntry.model_rebuild() +HirArrayLit.model_rebuild() +HirArrayMap.model_rebuild() +HirArrayReduce.model_rebuild() +HirArrayConcat.model_rebuild() +HirDistribution.model_rebuild() +HirDistributionMap.model_rebuild() +StateConstraint.model_rebuild() diff --git a/libs/@local/petrinaut-python/src/petrinaut/symbolic.py b/libs/@local/petrinaut-python/src/petrinaut/symbolic.py new file mode 100644 index 00000000000..747a5c91c9b --- /dev/null +++ b/libs/@local/petrinaut-python/src/petrinaut/symbolic.py @@ -0,0 +1,309 @@ +"""A parameter constraint as a SymPy relation, for the things evaluation +cannot do: solve for the feasible interval of one parameter, simplify a +compound condition, or hand the feasible region to a symbolic tool. + +Only the arithmetic subset translates: numbers, the scenario and net +parameters (one real symbol each, named as authored), the ``Math`` functions +with a symbolic counterpart, comparisons, ``&&``/``||``/``!``, ternaries, and +``const`` bindings by substitution. Arrays, records, strings, ``range()``, +and ``Math.random`` raise :class:`NotSymbolicError`; state constraints are +never symbolic. + +Booleans and numbers are kept apart the way SymPy needs them: a ternary +whose arms are conditions becomes ``ITE``, one whose arms are numbers +becomes ``Piecewise``, and ``==``/``!=`` over conditions become +``Equivalent``/``Xor``. Arithmetic follows ECMAScript where SymPy's default +differs: ``%`` is the truncated remainder (the dividend's sign) and +``Math.cbrt`` is the real cube root. + +The translation is exact mathematics over the reals. It does not carry +ECMAScript's floating-point edges (NaN, signed zero, overflow), and a +comparison at an exact boundary of ``Math.log10``/``Math.log2`` may need +``simplify`` or ``nsimplify`` before SymPy decides it; ask the evaluator when +those matter. + +SymPy is an optional dependency: ``petrinaut-python[sympy]``. +""" + +# pyright: reportMissingTypeStubs=false, reportUnknownMemberType=false +# pyright: reportUnknownVariableType=false, reportUnknownArgumentType=false +# SymPy ships no type information; every value it hands back is Any. +from __future__ import annotations + +import math +from collections.abc import Callable +from dataclasses import dataclass, field +from typing import Any + +from . import models as m +from .hir import HirExpr + +__all__ = ["NotSymbolicError", "SymbolicConstraint", "to_sympy"] + +_BOOLEAN_BINARY_OPS = frozenset({"<", "<=", ">", ">=", "==", "!=", "&&", "||"}) + + +class NotSymbolicError(ValueError): + """The constraint reads something SymPy cannot represent.""" + + +@dataclass(frozen=True) +class SymbolicConstraint: + """A SymPy relation plus the symbols it was built over.""" + + expression: Any + """A SymPy ``Boolean`` (a relation, or ``And``/``Or``/``Not`` of them).""" + scenario: dict[str, Any] = field(default_factory=dict) + """Scenario parameter identifier → ``Symbol``.""" + parameters: dict[str, Any] = field(default_factory=dict) + """Net parameter name → ``Symbol``.""" + + @property + def symbols(self) -> list[Any]: + """Every symbol, scenario parameters first, in first-use order.""" + return [*self.scenario.values(), *self.parameters.values()] + + +def to_sympy(constraint: m.ParameterConstraint) -> SymbolicConstraint: + """Translate one parameter constraint. Raises :class:`ImportError` when + SymPy is not installed and :class:`NotSymbolicError` when the condition + leaves the arithmetic subset.""" + try: + import sympy + except ImportError as error: # pragma: no cover - exercised without the extra + raise ImportError( + "SymPy is not installed; install the `sympy` extra of petrinaut-python" + ) from error + translator = _Translator(sympy) + expression = translator.expr(constraint.hir.body) + return SymbolicConstraint(expression, translator.scenario, translator.parameters) + + +class _Translator: + def __init__(self, sympy: Any) -> None: + self.sympy = sympy + self.scenario: dict[str, Any] = {} + self.parameters: dict[str, Any] = {} + self.locals: dict[str, Any] = {} + #: Names of ``const`` bindings whose value is a condition. + self.boolean_locals: set[str] = set() + sp = sympy + + def hypot(*args: Any) -> Any: + return sp.sqrt(sum(argument**2 for argument in args)) + + def log10(value: Any) -> Any: + return sp.log(value, 10) + + def log2(value: Any) -> Any: + return sp.log(value, 2) + + def js_round(value: Any) -> Any: + # ECMAScript Math.round is floor(x + 1/2) over the reals. + return sp.floor(value + sp.Rational(1, 2)) + + def cbrt(value: Any) -> Any: + # The real cube root; sp.cbrt is the principal complex root. + return sp.real_root(value, 3) + + self.math_fns: dict[str, Callable[..., Any]] = { + "abs": sp.Abs, + "acos": sp.acos, + "asin": sp.asin, + "atan": sp.atan, + "atan2": sp.atan2, + "cbrt": cbrt, + "ceil": sp.ceiling, + "cos": sp.cos, + "cosh": sp.cosh, + "exp": sp.exp, + "floor": sp.floor, + "hypot": hypot, + "log": sp.log, + "log10": log10, + "log2": log2, + "max": sp.Max, + "min": sp.Min, + "pow": sp.Pow, + "round": js_round, + "sign": sp.sign, + "sin": sp.sin, + "sinh": sp.sinh, + "sqrt": sp.sqrt, + "tan": sp.tan, + "tanh": sp.tanh, + "trunc": self._trunc, + } + + def _trunc(self, value: Any) -> Any: + sp = self.sympy + return sp.sign(value) * sp.floor(sp.Abs(value)) + + def _symbol(self, table: dict[str, Any], other: dict[str, Any], name: str) -> Any: + if name not in table: + if name in other: + raise NotSymbolicError( + f'"{name}" names both a scenario parameter and a net parameter' + ) + table[name] = self.sympy.Symbol(name, real=True) + return table[name] + + def _number(self, value: float) -> Any: + sp = self.sympy + if math.isnan(value): + return sp.nan + if math.isinf(value): + return sp.oo if value > 0 else -sp.oo + if value == int(value): + return sp.Integer(int(value)) + return sp.Float(value) + + def is_boolean(self, node: HirExpr) -> bool: + """Whether the node is a condition rather than a number, read off the + structure: HIR carries no types, and a bare parameter reads as a + number unless the other side of an operator says otherwise.""" + match node: + case m.HirBoolLit(): + return True + case m.HirBinary(): + return node.op.value in _BOOLEAN_BINARY_OPS + case m.HirUnary(): + return node.op.value == "!" + case m.HirCond(): + return self.is_boolean(node.thenBranch) or self.is_boolean( + node.elseBranch + ) + case m.HirLet(): + return self.is_boolean(node.body) + case m.HirLocalRef(): + return node.name in self.boolean_locals + case _: + return False + + def expr(self, node: HirExpr) -> Any: + sp = self.sympy + match node: + case m.HirNumberLit(): + return self._number(node.value) + case m.HirBoolLit(): + return sp.true if node.value else sp.false + case m.HirConstant(): + return { + "PI": sp.pi, + "E": sp.E, + "Infinity": sp.oo, + "NaN": sp.nan, + }[node.name.value] + case m.HirScenarioRef(): + return self._symbol(self.scenario, self.parameters, node.name) + case m.HirParamRef(): + return self._symbol(self.parameters, self.scenario, node.name) + case m.HirLocalRef(): + if node.name not in self.locals: + raise NotSymbolicError(f'Unbound local "{node.name}"') + return self.locals[node.name] + case m.HirUnary(): + operand = self.expr(node.operand) + op = node.op.value + if op == "!": + return self._logic(sp.Not, operand) + return -operand if op == "-" else operand + case m.HirBinary(): + return self._binary(node) + case m.HirCond(): + condition = self.expr(node.condition) + then_branch = self.expr(node.thenBranch) + else_branch = self.expr(node.elseBranch) + if self.is_boolean(node): + # A condition-valued ternary must stay a Boolean: a + # Piecewise in a condition position is rewritten by SymPy + # and loses its own condition. + return self._logic(sp.ITE, condition, then_branch, else_branch) + return sp.Piecewise((then_branch, condition), (else_branch, True)) + case m.HirLet(): + saved_locals = dict(self.locals) + saved_booleans = set(self.boolean_locals) + try: + for binding in node.bindings: + self.locals[binding.name] = self.expr(binding.value) + if self.is_boolean(binding.value): + self.boolean_locals.add(binding.name) + else: + self.boolean_locals.discard(binding.name) + return self.expr(node.body) + finally: + self.locals = saved_locals + self.boolean_locals = saved_booleans + case m.HirMathCall(): + fn = node.fn.value + if fn not in self.math_fns: + raise NotSymbolicError(f"Math.{fn}() has no symbolic form") + return self.math_fns[fn]( + *(self.expr(argument) for argument in node.args) + ) + case _: + raise NotSymbolicError( + f'HIR node kind "{node.kind}" has no symbolic form' + ) + + def _condition(self, expression: Any) -> Any: + """A relation as a Boolean SymPy can put in a condition position. A + Piecewise operand folds into a Piecewise of relations, which SymPy + rewrites lossily under a condition, so it becomes a chain of ITEs; + an arm-less remainder reads as false.""" + sp = self.sympy + folded = sp.piecewise_fold(expression) + if not isinstance(folded, sp.Piecewise): + return folded + result: Any = None + for arm_expression, arm_condition in reversed(folded.args): + if result is None: + result = ( + arm_expression + if arm_condition == sp.true + else sp.ITE(arm_condition, arm_expression, sp.false) + ) + else: + result = sp.ITE(arm_condition, arm_expression, result) + return result + + def _logic(self, connective: Any, *operands: Any) -> Any: + """Apply a Boolean connective; SymPy's TypeError for a non-Boolean + operand is the subset's boundary, so it is reported as such.""" + try: + return connective(*operands) + except TypeError as error: + raise NotSymbolicError(str(error)) from error + + def _binary(self, node: m.HirBinary) -> Any: + sp = self.sympy + left = self.expr(node.left) + right = self.expr(node.right) + op = node.op.value + if op == "&&": + return self._logic(sp.And, left, right) + if op == "||": + return self._logic(sp.Or, left, right) + if op == "+": + return left + right + if op == "-": + return left - right + if op == "*": + return left * right + if op == "/": + return left / right + if op == "%": + # ECMAScript remainder takes the dividend's sign; sp.Mod takes + # the divisor's. The truncated remainder is p - q * trunc(p / q). + return left - right * self._trunc(left / right) + if op == "**": + return sp.Pow(left, right) + if op in ("==", "!="): + if self.is_boolean(node.left) or self.is_boolean(node.right): + connective = sp.Equivalent if op == "==" else sp.Xor + return self._logic(connective, left, right) + return self._condition( + sp.Eq(left, right) if op == "==" else sp.Ne(left, right) + ) + relation = {"<": sp.Lt, "<=": sp.Le, ">": sp.Gt, ">=": sp.Ge}[op] + return self._condition(relation(left, right)) diff --git a/libs/@local/petrinaut-python/tests/hir_fixtures.json b/libs/@local/petrinaut-python/tests/hir_fixtures.json new file mode 100644 index 00000000000..d25e0fea7b8 --- /dev/null +++ b/libs/@local/petrinaut-python/tests/hir_fixtures.json @@ -0,0 +1,683 @@ +{ + "ordering": { + "code": "scenario.min_load < scenario.max_load", + "space": "parameters", + "hir": { + "hirVersion": 1, + "surface": "scenario-expression", + "params": [], + "body": { + "kind": "binary", + "op": "<", + "left": { + "kind": "scenarioRef", + "name": "min_load", + "id": 0, + "span": { + "start": 0, + "length": 17 + } + }, + "right": { + "kind": "scenarioRef", + "name": "max_load", + "id": 1, + "span": { + "start": 20, + "length": 17 + } + }, + "id": 2, + "span": { + "start": 0, + "length": 37 + } + }, + "span": { + "start": 0, + "length": 37 + } + } + }, + "compound": { + "code": "scenario.min_load < scenario.max_load && (parameters.rate > 0 || scenario.turbo)", + "space": "parameters", + "hir": { + "hirVersion": 1, + "surface": "scenario-expression", + "params": [], + "body": { + "kind": "binary", + "op": "&&", + "left": { + "kind": "binary", + "op": "<", + "left": { + "kind": "scenarioRef", + "name": "min_load", + "id": 0, + "span": { + "start": 0, + "length": 17 + } + }, + "right": { + "kind": "scenarioRef", + "name": "max_load", + "id": 1, + "span": { + "start": 20, + "length": 17 + } + }, + "id": 2, + "span": { + "start": 0, + "length": 37 + } + }, + "right": { + "kind": "binary", + "op": "||", + "left": { + "kind": "binary", + "op": ">", + "left": { + "kind": "paramRef", + "name": "rate", + "id": 3, + "span": { + "start": 42, + "length": 15 + } + }, + "right": { + "kind": "numberLit", + "value": 0, + "raw": "0", + "id": 4, + "span": { + "start": 60, + "length": 1 + } + }, + "id": 5, + "span": { + "start": 42, + "length": 19 + } + }, + "right": { + "kind": "scenarioRef", + "name": "turbo", + "id": 6, + "span": { + "start": 65, + "length": 14 + } + }, + "id": 7, + "span": { + "start": 42, + "length": 37 + } + }, + "id": 8, + "span": { + "start": 0, + "length": 80 + } + }, + "span": { + "start": 0, + "length": 80 + } + } + }, + "math": { + "code": "Math.round(Math.abs(scenario.min_load - scenario.max_load)) >= 2", + "space": "parameters", + "hir": { + "hirVersion": 1, + "surface": "scenario-expression", + "params": [], + "body": { + "kind": "binary", + "op": ">=", + "left": { + "kind": "mathCall", + "fn": "round", + "args": [ + { + "kind": "mathCall", + "fn": "abs", + "args": [ + { + "kind": "binary", + "op": "-", + "left": { + "kind": "scenarioRef", + "name": "min_load", + "id": 0, + "span": { + "start": 20, + "length": 17 + } + }, + "right": { + "kind": "scenarioRef", + "name": "max_load", + "id": 1, + "span": { + "start": 40, + "length": 17 + } + }, + "id": 2, + "span": { + "start": 20, + "length": 37 + } + } + ], + "id": 3, + "span": { + "start": 11, + "length": 47 + } + } + ], + "id": 4, + "span": { + "start": 0, + "length": 59 + } + }, + "right": { + "kind": "numberLit", + "value": 2, + "raw": "2", + "id": 5, + "span": { + "start": 63, + "length": 1 + } + }, + "id": 6, + "span": { + "start": 0, + "length": 64 + } + }, + "span": { + "start": 0, + "length": 64 + } + } + }, + "ternary": { + "code": "scenario.turbo ? scenario.max_load <= 10 : scenario.max_load <= 6", + "space": "parameters", + "hir": { + "hirVersion": 1, + "surface": "scenario-expression", + "params": [], + "body": { + "kind": "cond", + "condition": { + "kind": "scenarioRef", + "name": "turbo", + "id": 0, + "span": { + "start": 0, + "length": 14 + } + }, + "thenBranch": { + "kind": "binary", + "op": "<=", + "left": { + "kind": "scenarioRef", + "name": "max_load", + "id": 1, + "span": { + "start": 17, + "length": 17 + } + }, + "right": { + "kind": "numberLit", + "value": 10, + "raw": "10", + "id": 2, + "span": { + "start": 38, + "length": 2 + } + }, + "id": 3, + "span": { + "start": 17, + "length": 23 + } + }, + "elseBranch": { + "kind": "binary", + "op": "<=", + "left": { + "kind": "scenarioRef", + "name": "max_load", + "id": 4, + "span": { + "start": 43, + "length": 17 + } + }, + "right": { + "kind": "numberLit", + "value": 6, + "raw": "6", + "id": 5, + "span": { + "start": 64, + "length": 1 + } + }, + "id": 6, + "span": { + "start": 43, + "length": 22 + } + }, + "id": 7, + "span": { + "start": 0, + "length": 65 + } + }, + "span": { + "start": 0, + "length": 65 + } + } + }, + "strictEquality": { + "code": "scenario.min_load == 1", + "space": "parameters", + "hir": { + "hirVersion": 1, + "surface": "scenario-expression", + "params": [], + "body": { + "kind": "binary", + "op": "==", + "left": { + "kind": "scenarioRef", + "name": "min_load", + "id": 0, + "span": { + "start": 0, + "length": 17 + } + }, + "right": { + "kind": "numberLit", + "value": 1, + "raw": "1", + "id": 1, + "span": { + "start": 21, + "length": 1 + } + }, + "id": 2, + "span": { + "start": 0, + "length": 22 + } + }, + "span": { + "start": 0, + "length": 22 + } + } + }, + "stateBound": { + "code": "return state.places.Queue.count <= 10;", + "space": "state", + "hir": { + "hirVersion": 1, + "surface": "metric", + "params": [ + { + "name": "state", + "span": { + "start": 0, + "length": 0 + } + } + ], + "body": { + "kind": "binary", + "op": "<=", + "left": { + "kind": "fieldAccess", + "target": { + "kind": "fieldAccess", + "target": { + "kind": "fieldAccess", + "target": { + "kind": "localRef", + "name": "state", + "id": 0, + "span": { + "start": 7, + "length": 5 + } + }, + "field": "places", + "fieldSpan": { + "start": 13, + "length": 6 + }, + "id": 1, + "span": { + "start": 7, + "length": 12 + } + }, + "field": "Queue", + "fieldSpan": { + "start": 20, + "length": 5 + }, + "id": 2, + "span": { + "start": 7, + "length": 18 + } + }, + "field": "count", + "fieldSpan": { + "start": 26, + "length": 5 + }, + "id": 3, + "span": { + "start": 7, + "length": 24 + } + }, + "right": { + "kind": "numberLit", + "value": 10, + "raw": "10", + "id": 4, + "span": { + "start": 35, + "length": 2 + } + }, + "id": 5, + "span": { + "start": 7, + "length": 30 + } + }, + "span": { + "start": 0, + "length": 38 + } + } + }, + "stateBlock": { + "code": "const total = state.places.Queue.tokens.reduce((acc, token) => acc + 1, 0);\nreturn total <= 5 && state.places.Queue.count >= 0;", + "space": "state", + "hir": { + "hirVersion": 1, + "surface": "metric", + "params": [ + { + "name": "state", + "span": { + "start": 0, + "length": 0 + } + } + ], + "body": { + "kind": "let", + "bindings": [ + { + "name": "total", + "nameSpan": { + "start": 6, + "length": 5 + }, + "value": { + "kind": "arrayReduce", + "target": { + "kind": "fieldAccess", + "target": { + "kind": "fieldAccess", + "target": { + "kind": "fieldAccess", + "target": { + "kind": "localRef", + "name": "state", + "id": 0, + "span": { + "start": 14, + "length": 5 + } + }, + "field": "places", + "fieldSpan": { + "start": 20, + "length": 6 + }, + "id": 1, + "span": { + "start": 14, + "length": 12 + } + }, + "field": "Queue", + "fieldSpan": { + "start": 27, + "length": 5 + }, + "id": 2, + "span": { + "start": 14, + "length": 18 + } + }, + "field": "tokens", + "fieldSpan": { + "start": 33, + "length": 6 + }, + "id": 3, + "span": { + "start": 14, + "length": 25 + } + }, + "accParam": { + "name": "acc", + "span": { + "start": 48, + "length": 3 + } + }, + "param": { + "name": "token", + "span": { + "start": 53, + "length": 5 + } + }, + "body": { + "kind": "binary", + "op": "+", + "left": { + "kind": "localRef", + "name": "acc", + "id": 5, + "span": { + "start": 63, + "length": 3 + } + }, + "right": { + "kind": "numberLit", + "value": 1, + "raw": "1", + "id": 6, + "span": { + "start": 69, + "length": 1 + } + }, + "id": 7, + "span": { + "start": 63, + "length": 7 + } + }, + "initial": { + "kind": "numberLit", + "value": 0, + "raw": "0", + "id": 4, + "span": { + "start": 72, + "length": 1 + } + }, + "id": 8, + "span": { + "start": 14, + "length": 60 + } + } + } + ], + "body": { + "kind": "binary", + "op": "&&", + "left": { + "kind": "binary", + "op": "<=", + "left": { + "kind": "localRef", + "name": "total", + "id": 9, + "span": { + "start": 83, + "length": 5 + } + }, + "right": { + "kind": "numberLit", + "value": 5, + "raw": "5", + "id": 10, + "span": { + "start": 92, + "length": 1 + } + }, + "id": 11, + "span": { + "start": 83, + "length": 10 + } + }, + "right": { + "kind": "binary", + "op": ">=", + "left": { + "kind": "fieldAccess", + "target": { + "kind": "fieldAccess", + "target": { + "kind": "fieldAccess", + "target": { + "kind": "localRef", + "name": "state", + "id": 12, + "span": { + "start": 97, + "length": 5 + } + }, + "field": "places", + "fieldSpan": { + "start": 103, + "length": 6 + }, + "id": 13, + "span": { + "start": 97, + "length": 12 + } + }, + "field": "Queue", + "fieldSpan": { + "start": 110, + "length": 5 + }, + "id": 14, + "span": { + "start": 97, + "length": 18 + } + }, + "field": "count", + "fieldSpan": { + "start": 116, + "length": 5 + }, + "id": 15, + "span": { + "start": 97, + "length": 24 + } + }, + "right": { + "kind": "numberLit", + "value": 0, + "raw": "0", + "id": 16, + "span": { + "start": 125, + "length": 1 + } + }, + "id": 17, + "span": { + "start": 97, + "length": 29 + } + }, + "id": 18, + "span": { + "start": 83, + "length": 43 + } + }, + "id": 19, + "span": { + "start": 0, + "length": 127 + } + }, + "span": { + "start": 0, + "length": 127 + } + } + } +} diff --git a/libs/@local/petrinaut-python/tests/test_constraint.py b/libs/@local/petrinaut-python/tests/test_constraint.py new file mode 100644 index 00000000000..29c99e0f24b --- /dev/null +++ b/libs/@local/petrinaut-python/tests/test_constraint.py @@ -0,0 +1,424 @@ +"""The two constraint shapes as callables: parsing with full HIR validation, +the binding each shape takes, the four readings of a condition, the pydantic +validator, and the symbolic view.""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Annotated, Any + +import pytest +import sympy +from pydantic import AfterValidator, BaseModel, ValidationError + +from petrinaut import ( + ConstraintViolation, + HirEvaluationError, + NotSymbolicError, + OptimizationDescribeResult, + ParameterConstraint, + StateConstraint, + parse_constraint, + parse_constraints, + violations, +) + +FIXTURES = json.loads( + (Path(__file__).parent / "hir_fixtures.json").read_text(encoding="utf-8") +) + + +def data(name: str, **overrides: Any) -> dict[str, Any]: + fixture = FIXTURES[name] + return { + "space": fixture["space"], + "id": name, + "code": fixture["code"], + "hir": fixture["hir"], + **overrides, + } + + +class TestParsing: + def test_discriminates_on_space(self) -> None: + assert isinstance(parse_constraint(data("ordering")), ParameterConstraint) + assert isinstance(parse_constraint(data("stateBound")), StateConstraint) + + def test_pins_the_surface_to_the_space(self) -> None: + # A metric-surface function cannot pose as a parameter constraint. + misfiled = data("ordering", hir=FIXTURES["stateBound"]["hir"]) + with pytest.raises(ValidationError, match="scenario-expression"): + parse_constraint(misfiled) + + def test_validates_every_node(self) -> None: + broken = data("ordering") + broken["hir"] = json.loads(json.dumps(broken["hir"])) + del broken["hir"]["body"]["left"]["span"] + with pytest.raises(ValidationError, match="span"): + parse_constraint(broken) + + def test_rejects_an_unknown_node_kind(self) -> None: + forged = data("ordering") + forged["hir"] = json.loads(json.dumps(forged["hir"])) + forged["hir"]["body"]["kind"] = "eval" + with pytest.raises(ValidationError, match="eval"): + parse_constraint(forged) + + def test_rejects_fields_outside_the_grammar(self) -> None: + extra = data("ordering") + extra["hir"] = {**extra["hir"], "compiled": True} + with pytest.raises(ValidationError, match="compiled"): + parse_constraint(extra) + + def test_reads_a_describe_result(self) -> None: + described = OptimizationDescribeResult.model_validate( + { + "direction": "maximize", + "study": {"trials": 3, "sampler": "random", "seed": 1}, + "parameters": [], + "constraints": [data("ordering"), data("stateBound")], + } + ) + constraints = parse_constraints(described.constraints) + assert [type(constraint).__name__ for constraint in constraints] == [ + "ParameterConstraint", + "StateConstraint", + ] + assert constraints[0].id == "ordering" + # A callable passes through untouched; None reads as no constraints. + assert parse_constraint(constraints[0]) is constraints[0] + assert parse_constraints(None) == [] + + +class TestParameterConstraint: + def test_takes_a_scenario(self) -> None: + ordering = parse_constraint(data("ordering")) + assert isinstance(ordering, ParameterConstraint) + assert ordering(scenario={"min_load": 2, "max_load": 8}) is True + assert ordering({"min_load": 8, "max_load": 2}) is False + + def test_reads_net_parameters(self) -> None: + compound = parse_constraint(data("compound")) + assert isinstance(compound, ParameterConstraint) + scenario = {"min_load": 1, "max_load": 4, "turbo": False} + assert compound(scenario, parameters={"rate": 1.5}) is True + with pytest.raises(HirEvaluationError, match="rate"): + compound(scenario) + + def test_four_readings_agree(self) -> None: + ordering = parse_constraint(data("ordering")) + assert isinstance(ordering, ParameterConstraint) + holds = {"min_load": 2, "max_load": 8} + fails = {"min_load": 8, "max_load": 2} + assert ordering.margin(holds) == 6.0 + assert ordering.violation(holds) == -6.0 + ordering.check(holds) + assert ordering.margin(fails) == -6.0 + with pytest.raises(ConstraintViolation, match="ordering") as raised: + ordering.check(fails) + assert raised.value.margin == -6.0 + assert raised.value.constraint is ordering + + def test_validator_plugs_into_pydantic(self) -> None: + ordering = parse_constraint(data("ordering")) + assert isinstance(ordering, ParameterConstraint) + + class Study(BaseModel): + scenario: Annotated[dict[str, float], AfterValidator(ordering.validator())] + + assert Study(scenario={"min_load": 1, "max_load": 2}).scenario == { + "min_load": 1, + "max_load": 2, + } + with pytest.raises(ValidationError, match="violated"): + Study(scenario={"min_load": 3, "max_load": 2}) + + +class TestStateConstraint: + def test_takes_a_state(self) -> None: + bound = parse_constraint(data("stateBound")) + assert isinstance(bound, StateConstraint) + assert bound(state={"places": {"Queue": {"count": 7}}}) is True + assert bound({"places": {"Queue": {"count": 11}}}) is False + assert bound.margin({"places": {"Queue": {"count": 7}}}) == 3.0 + assert bound.violation({"places": {"Queue": {"count": 11}}}) == 1.0 + + def test_check_and_validator(self) -> None: + bound = parse_constraint(data("stateBound")) + assert isinstance(bound, StateConstraint) + with pytest.raises(ConstraintViolation, match="stateBound"): + bound.check({"places": {"Queue": {"count": 11}}}) + + class Snapshot(BaseModel): + state: Annotated[dict[str, Any], AfterValidator(bound.validator())] + + Snapshot(state={"places": {"Queue": {"count": 1}}}) + with pytest.raises(ValidationError, match="violated"): + Snapshot(state={"places": {"Queue": {"count": 99}}}) + + +class TestViolations: + def test_one_slot_per_constraint(self) -> None: + constraints = parse_constraints([data("ordering"), data("stateBound")]) + out = violations( + constraints, + scenario={"min_load": 2, "max_load": 8}, + state={"places": {"Queue": {"count": 12}}}, + ) + assert out == [-6.0, 2.0] + + def test_a_missing_binding_is_an_error_not_a_gap(self) -> None: + constraints = parse_constraints([data("ordering"), data("stateBound")]) + with pytest.raises(ValueError, match="needs a state"): + violations(constraints, scenario={"min_load": 2, "max_load": 8}) + with pytest.raises(ValueError, match="needs a scenario"): + violations(constraints, state={}) + + +class TestSymbolic: + def test_ordering_becomes_a_relation(self) -> None: + ordering = parse_constraint(data("ordering")) + assert isinstance(ordering, ParameterConstraint) + symbolic = ordering.to_sympy() + min_load, max_load = ( + symbolic.scenario["min_load"], + symbolic.scenario["max_load"], + ) + assert symbolic.expression == sympy.Lt(min_load, max_load) + # The symbolic view answers what evaluation cannot: the feasible + # interval of one parameter given the others. + solved = sympy.solve_univariate_inequality( + symbolic.expression.subs(max_load, 8), min_load, relational=False + ) + assert solved == sympy.Interval.open(-sympy.oo, 8) + + def test_compound_keeps_both_parameter_kinds_apart(self) -> None: + compound = parse_constraint(data("compound")) + assert isinstance(compound, ParameterConstraint) + symbolic = compound.to_sympy() + assert set(symbolic.scenario) == {"min_load", "max_load", "turbo"} + assert set(symbolic.parameters) == {"rate"} + # Substituting a satisfying point evaluates the relation to true. + point = { + symbolic.scenario["min_load"]: 1, + symbolic.scenario["max_load"]: 4, + symbolic.scenario["turbo"]: sympy.false, + symbolic.parameters["rate"]: 2, + } + assert symbolic.expression.subs(point) == sympy.true + + def test_math_and_ternary_translate(self) -> None: + for name in ("math", "ternary"): + constraint = parse_constraint(data(name)) + assert isinstance(constraint, ParameterConstraint) + symbolic = constraint.to_sympy() + assert symbolic.expression is not None + + def test_state_reads_have_no_symbolic_form(self) -> None: + block = parse_constraint(data("stateBlock")) + assert isinstance(block, StateConstraint) + assert not hasattr(block, "to_sympy") + # A parameter constraint over an array is out of the subset too. + from petrinaut.symbolic import to_sympy + + with pytest.raises(NotSymbolicError): + to_sympy( + ParameterConstraint.model_validate( + data( + "ordering", + hir={ + "hirVersion": 1, + "surface": "scenario-expression", + "params": [], + "span": {"start": 0, "length": 1}, + "body": { + "kind": "binary", + "id": 2, + "span": {"start": 0, "length": 1}, + "op": ">", + "left": { + "kind": "length", + "id": 1, + "span": {"start": 0, "length": 1}, + "target": { + "kind": "arrayLit", + "id": 0, + "span": {"start": 0, "length": 1}, + "elements": [], + }, + }, + "right": { + "kind": "numberLit", + "id": 3, + "span": {"start": 0, "length": 1}, + "value": 0, + "raw": "0", + }, + }, + }, + ) + ) + ) + + +SPAN = {"start": 0, "length": 1} + + +def node(kind: str, **fields: Any) -> dict[str, Any]: + built: dict[str, Any] = {"kind": kind, "id": 0, "span": SPAN, **fields} + if kind == "fieldAccess": + built.setdefault("fieldSpan", SPAN) + return built + + +def num(value: float) -> dict[str, Any]: + return node("numberLit", value=value, raw=repr(value)) + + +def ref(name: str) -> dict[str, Any]: + return node("scenarioRef", name=name) + + +def parameter_constraint(body: dict[str, Any]) -> ParameterConstraint: + return ParameterConstraint.model_validate( + { + "space": "parameters", + "id": "inline", + "code": "", + "hir": { + "hirVersion": 1, + "surface": "scenario-expression", + "params": [], + "span": SPAN, + "body": body, + }, + } + ) + + +class TestSymbolicAgreesWithEvaluation: + """Every translated construct substitutes to the value the evaluator + computes; the cases are the ones SymPy gets wrong when handed naively + (Piecewise in condition positions, Mod's sign, the complex cube root).""" + + @staticmethod + def agree( + constraint: ParameterConstraint, assignments: list[dict[str, Any]] + ) -> None: + symbolic = constraint.to_sympy() + for scenario in assignments: + point = { + symbol: ( + sympy.true + if value is True + else sympy.false + if value is False + else value + ) + for name, symbol in symbolic.scenario.items() + for value in [scenario[name]] + } + expected = constraint(scenario) + actual = bool(symbolic.expression.subs(point)) + assert actual is expected, (scenario, symbolic.expression) + + def test_numeric_ternary_inside_a_condition(self) -> None: + # ((a > (flag ? b : c)) ? 1 : 2) == 2 + inner = node( + "cond", condition=ref("flag"), thenBranch=ref("b"), elseBranch=ref("c") + ) + outer = node( + "cond", + condition=node("binary", op=">", left=ref("a"), right=inner), + thenBranch=num(1), + elseBranch=num(2), + ) + constraint = parameter_constraint( + node("binary", op="==", left=outer, right=num(2)) + ) + self.agree( + constraint, + [ + {"a": 5, "b": 3, "c": 10, "flag": False}, + {"a": 5, "b": 3, "c": 10, "flag": True}, + {"a": 0, "b": 3, "c": -1, "flag": True}, + ], + ) + + def test_boolean_ternary_under_and(self) -> None: + # (flag ? a < 1 : a < 2) && b > 0 + ternary = node( + "cond", + condition=ref("flag"), + thenBranch=node("binary", op="<", left=ref("a"), right=num(1)), + elseBranch=node("binary", op="<", left=ref("a"), right=num(2)), + ) + constraint = parameter_constraint( + node( + "binary", + op="&&", + left=ternary, + right=node("binary", op=">", left=ref("b"), right=num(0)), + ) + ) + self.agree( + constraint, + [ + {"a": 1.5, "b": 1, "flag": False}, + {"a": 1.5, "b": 1, "flag": True}, + {"a": 0.5, "b": -1, "flag": True}, + ], + ) + + def test_equality_of_boolean_ternaries(self) -> None: + # (turbo ? flag : true) == (1.5 <= rate) and the same with != + for op in ("==", "!="): + left = node( + "cond", + condition=ref("turbo"), + thenBranch=ref("flag"), + elseBranch=node("boolLit", value=True), + ) + right = node("binary", op="<=", left=num(1.5), right=ref("rate")) + constraint = parameter_constraint( + node("binary", op=op, left=left, right=right) + ) + self.agree( + constraint, + [ + {"turbo": False, "flag": False, "rate": 100}, + {"turbo": True, "flag": False, "rate": 100}, + {"turbo": True, "flag": True, "rate": 1}, + ], + ) + + def test_remainder_takes_the_dividends_sign(self) -> None: + constraint = parameter_constraint( + node( + "binary", + op="<", + left=node("binary", op="%", left=ref("a"), right=num(3)), + right=num(0), + ) + ) + self.agree(constraint, [{"a": -7}, {"a": 7}, {"a": -4.5}, {"a": 6}]) + + def test_cube_root_stays_real(self) -> None: + constraint = parameter_constraint( + node( + "binary", + op="<", + left=node("mathCall", fn="cbrt", args=[ref("a")]), + right=num(0), + ) + ) + self.agree(constraint, [{"a": -8}, {"a": 8}, {"a": -0.749}]) + + +class TestStateBinding: + def test_a_state_must_be_a_record(self) -> None: + bound = parse_constraint(data("stateBound")) + assert isinstance(bound, StateConstraint) + with pytest.raises(HirEvaluationError, match="state record"): + bound([1, 2, 3]) # type: ignore[arg-type] diff --git a/libs/@local/petrinaut-python/tests/test_hir.py b/libs/@local/petrinaut-python/tests/test_hir.py new file mode 100644 index 00000000000..efce8c8367a --- /dev/null +++ b/libs/@local/petrinaut-python/tests/test_hir.py @@ -0,0 +1,583 @@ +"""Tests for the serialized-HIR evaluator against fixtures lowered by the +real TypeScript frontend (`hir_fixtures.json`, generated from +`lowerConstraint` in `@hashintel/petrinaut-core`).""" + +import json +import math +from pathlib import Path + +import pytest +from pydantic import ValidationError + +from petrinaut import ( + HirEvaluationError, + ParameterConstraint, + StateConstraint, + evaluate_hir, +) + +SPAN = {"start": 0, "length": 1} + +FIXTURES = json.loads( + (Path(__file__).parent / "hir_fixtures.json").read_text(encoding="utf-8") +) + + +def constraint(name: str) -> ParameterConstraint | StateConstraint: + fixture = FIXTURES[name] + data = { + "space": fixture["space"], + "id": name, + "code": fixture["code"], + "hir": fixture["hir"], + } + if fixture["space"] == "parameters": + return ParameterConstraint.model_validate(data) + return StateConstraint.model_validate(data) + + +def parameter_constraint( + id_: str, body: dict[str, object], code: str = "" +) -> ParameterConstraint: + return ParameterConstraint.model_validate( + { + "space": "parameters", + "id": id_, + "code": code, + "hir": { + "hirVersion": 1, + "surface": "scenario-expression", + "params": [], + "span": SPAN, + "body": body, + }, + } + ) + + +def state_constraint(id_: str, body: dict[str, object], param: str) -> StateConstraint: + return StateConstraint.model_validate( + { + "space": "state", + "id": id_, + "code": "", + "hir": { + "hirVersion": 1, + "surface": "metric", + "params": [{"name": param, "span": SPAN}], + "span": SPAN, + "body": body, + }, + } + ) + + +class TestParameterSpace: + def test_ordering(self) -> None: + ordering = constraint("ordering") + assert ordering(scenario={"min_load": 2, "max_load": 8}) is True + assert ordering(scenario={"min_load": 8, "max_load": 2}) is False + + def test_compound_short_circuit_and_boolean(self) -> None: + compound = constraint("compound") + assert ( + compound( + scenario={"min_load": 1, "max_load": 4, "turbo": False}, + parameters={"rate": 1.5}, + ) + is True + ) + assert ( + compound( + scenario={"min_load": 1, "max_load": 4, "turbo": False}, + parameters={"rate": 0.0}, + ) + is False + ) + # turbo rescues a non-positive rate through the `||`. + assert ( + compound( + scenario={"min_load": 1, "max_load": 4, "turbo": True}, + parameters={"rate": 0.0}, + ) + is True + ) + + def test_math_matches_ecmascript_rounding(self) -> None: + math_case = constraint("math") + # |2 - 4.5| = 2.5 → Math.round gives 3 in JS (half away from + # negative), so the constraint holds; Python's round(2.5) is 2. + assert math_case(scenario={"min_load": 2, "max_load": 4.5}) is True + assert math_case(scenario={"min_load": 2, "max_load": 3.4}) is False + + def test_ternary(self) -> None: + ternary = constraint("ternary") + assert ternary(scenario={"turbo": True, "max_load": 9}) is True + assert ternary(scenario={"turbo": False, "max_load": 9}) is False + + def test_strict_equality_keeps_booleans_apart(self) -> None: + strict = constraint("strictEquality") + assert strict(scenario={"min_load": 1}) is True + # JS: `true === 1` is false; Python's `True == 1` must not leak in. + assert strict(scenario={"min_load": True}) is False + + def test_unknown_scenario_parameter_raises(self) -> None: + ordering = constraint("ordering") + with pytest.raises(HirEvaluationError, match="min_load"): + ordering(scenario={"max_load": 8}) + + +class TestStateSpace: + def test_state_bound(self) -> None: + bound = constraint("stateBound") + assert bound(state={"places": {"Queue": {"count": 7}}}) is True + assert bound(state={"places": {"Queue": {"count": 11}}}) is False + + def test_state_block_with_reduce(self) -> None: + block = constraint("stateBlock") + state = { + "places": {"Queue": {"count": 3, "tokens": [{}, {}, {}]}}, + } + assert block(state=state) is True + state_over = { + "places": { + "Queue": {"count": 6, "tokens": [{}, {}, {}, {}, {}, {}]}, + }, + } + assert block(state=state_over) is False + + +class TestMargin: + def test_comparison_slack(self) -> None: + ordering = constraint("ordering") + assert ordering.margin(scenario={"min_load": 2, "max_load": 8}) == 6.0 + assert ordering.margin(scenario={"min_load": 8, "max_load": 2}) == -6.0 + + def test_and_takes_the_minimum(self) -> None: + compound = constraint("compound") + margin = compound.margin( + scenario={"min_load": 1, "max_load": 4, "turbo": False}, + parameters={"rate": 0.5}, + ) + # min(4 - 1, max(0.5 - 0, -inf)) = 0.5 + assert margin == 0.5 + + def test_boolean_leaf_is_infinite(self) -> None: + compound = constraint("compound") + margin = compound.margin( + scenario={"min_load": 1, "max_load": 9, "turbo": True}, + parameters={"rate": -1.0}, + ) + # The `|| turbo` arm is +inf, so the && is bounded by 9 - 1. + assert margin == 8.0 + + def test_strict_boundary_is_violated(self) -> None: + # `min_load < max_load` at equality is false, so the margin must go + # negative there rather than reporting a satisfied-looking zero. + ordering = constraint("ordering") + boundary = ordering.margin(scenario={"min_load": 5, "max_load": 5}) + assert boundary < 0 + assert not ordering(scenario={"min_load": 5, "max_load": 5}) + + def test_sign_agrees_with_the_boolean(self) -> None: + for name in ("ordering", "ternary", "strictEquality"): + case = constraint(name) + for scenario in ( + {"min_load": 2, "max_load": 8, "turbo": True}, + {"min_load": 8, "max_load": 2, "turbo": False}, + {"min_load": 1, "max_load": 6, "turbo": False}, + {"min_load": 4, "max_load": 4, "turbo": False}, + ): + satisfied = case(scenario=scenario) + margin = case.margin(scenario=scenario) + assert (margin >= 0) == satisfied, (name, scenario) + + +class TestStringNodes: + @staticmethod + def _node(kind: str, **fields: object) -> dict[str, object]: + return {"kind": kind, "id": 0, "span": {"start": 0, "length": 1}, **fields} + + def _constraint(self, body: dict[str, object]) -> ParameterConstraint: + return parameter_constraint("strings", body) + + def test_string_methods_evaluate(self) -> None: + def lit(value: str) -> dict[str, object]: + return self._node("stringLit", value=value) + + for fn, target, argument, expected in ( + ("startsWith", "pump-3", "pump", True), + ("endsWith", "pump-3", "-3", True), + ("includes", "pump-3", "mp", True), + ("includes", "pump-3", "xyz", False), + ): + case = self._constraint( + self._node( + "stringCall", fn=fn, target=lit(target), argument=lit(argument) + ) + ) + assert case(scenario={}) is expected, (fn, target, argument) + + def test_string_length(self) -> None: + body = self._node( + "binary", + op=">", + left=self._node("length", target=self._node("stringLit", value="abc")), + right=self._node("numberLit", value=2, raw="2"), + ) + assert self._constraint(body)(scenario={}) is True + + +class TestJsMathEdges: + """The evaluator's arithmetic must match ECMAScript at the edges Python + diverges: domain errors, overflow, and exponentiation.""" + + @staticmethod + def _num(value: float) -> dict[str, object]: + return { + "kind": "numberLit", + "id": 0, + "span": {"start": 0, "length": 1}, + "value": value, + "raw": repr(value), + } + + def _eval(self, body: dict[str, object]) -> object: + return evaluate_hir( + { + "hirVersion": 1, + "surface": "scenario-expression", + "params": [], + "span": SPAN, + "body": body, + } + ) + + def _math(self, fn: str, *args: float) -> object: + return self._eval( + { + "kind": "mathCall", + "id": 0, + "span": {"start": 0, "length": 1}, + "fn": fn, + "args": [self._num(argument) for argument in args], + } + ) + + def _pow(self, base: float, exponent: float) -> object: + return self._eval( + { + "kind": "binary", + "id": 0, + "span": {"start": 0, "length": 1}, + "op": "**", + "left": self._num(base), + "right": self._num(exponent), + } + ) + + def test_log_family_matches_js(self) -> None: + assert self._math("log", 0) == -math.inf + assert math.isnan(self._math("log", -1)) # type: ignore[arg-type] + assert self._math("log10", 0) == -math.inf + assert self._math("log2", 0) == -math.inf + + def test_overflow_grows_to_infinity(self) -> None: + assert self._math("exp", 1000) == math.inf + assert self._math("cosh", 1000) == math.inf + assert self._math("sinh", -1000) == -math.inf + + def test_pow_matches_js(self) -> None: + # Python raises ZeroDivisionError / OverflowError or goes complex + # for every one of these; JS defines them all. + assert self._pow(0, -1) == math.inf + assert self._pow(1e308, 2) == math.inf + assert self._pow(-1e308, 3) == -math.inf + assert math.isnan(self._pow(-8, 1 / 3)) # type: ignore[arg-type] + assert math.isnan(self._pow(1, math.inf)) # type: ignore[arg-type] + assert self._pow(-2, 3) == -8 + assert self._math("pow", 0, -1) == math.inf + + def test_integral_functions_pass_non_finite_through(self) -> None: + assert self._math("ceil", math.inf) == math.inf + assert self._math("floor", -math.inf) == -math.inf + assert math.isnan(self._math("round", math.nan)) # type: ignore[arg-type] + assert self._math("round", math.inf) == math.inf + + +class TestNanMargins: + """A NaN slack resolves at the comparison leaf with the boolean's sign, + so `min`/`max` in compounds can never drop it by argument order.""" + + @staticmethod + def _node(kind: str, **fields: object) -> dict[str, object]: + return {"kind": kind, "id": 0, "span": {"start": 0, "length": 1}, **fields} + + def _constraint(self, body: dict[str, object]) -> ParameterConstraint: + return parameter_constraint("nan-margins", body) + + def _num(self, value: float) -> dict[str, object]: + return self._node("numberLit", value=value, raw=repr(value)) + + def _nan_leaf(self) -> dict[str, object]: + # Math.sqrt(-1) > 2 — false in JS, its slack NaN in Python. + return self._node( + "binary", + op=">", + left=self._node("mathCall", fn="sqrt", args=[self._num(-1)]), + right=self._num(2), + ) + + def _sat_leaf(self) -> dict[str, object]: + return self._node("binary", op="<", left=self._num(5), right=self._num(9)) + + def test_and_with_nan_slack_reads_unsatisfied(self) -> None: + for left, right in ( + (self._nan_leaf(), self._sat_leaf()), + (self._sat_leaf(), self._nan_leaf()), + ): + case = self._constraint( + self._node("binary", op="&&", left=left, right=right) + ) + assert case(scenario={}) is False + assert case.margin(scenario={}) < 0 + + def test_or_with_nan_slack_follows_the_boolean(self) -> None: + rescued = self._constraint( + self._node("binary", op="||", left=self._nan_leaf(), right=self._sat_leaf()) + ) + assert rescued(scenario={}) is True + assert rescued.margin(scenario={}) >= 0 + + def test_negated_nan_comparison_reads_satisfied(self) -> None: + negated = self._constraint( + self._node("unary", op="!", operand=self._nan_leaf()) + ) + assert negated(scenario={}) is True + assert negated.margin(scenario={}) >= 0 + + def test_nan_equality_margins(self) -> None: + sqrt_neg = self._node("mathCall", fn="sqrt", args=[self._num(-1)]) + unequal = self._constraint( + self._node("binary", op="!=", left=sqrt_neg, right=self._num(5)) + ) + assert unequal(scenario={}) is True + assert unequal.margin(scenario={}) >= 0 + + +class TestMarginShortCircuit: + """`margin()` walks the same arms evaluation walks, and a slack that + cancels to NaN still reports the comparison's own answer.""" + + @staticmethod + def _node(kind: str, **fields: object) -> dict[str, object]: + node: dict[str, object] = {"kind": kind, "id": 0, "span": SPAN, **fields} + if kind == "fieldAccess": + node.setdefault("fieldSpan", SPAN) + return node + + def _constraint( + self, body: dict[str, object], params: list[str] | None = None + ) -> ParameterConstraint | StateConstraint: + if params: + return state_constraint("short-circuit", body, params[0]) + return parameter_constraint("short-circuit", body) + + def _num(self, value: float) -> dict[str, object]: + return self._node("numberLit", value=value, raw=repr(value)) + + def _queue(self) -> dict[str, object]: + state = self._node("localRef", name="state") + places = self._node("fieldAccess", target=state, field="places") + return self._node("fieldAccess", target=places, field="Queue") + + def _count(self) -> dict[str, object]: + return self._node("fieldAccess", target=self._queue(), field="count") + + def _first_token_x(self) -> dict[str, object]: + tokens = self._node("fieldAccess", target=self._queue(), field="tokens") + first = self._node("indexAccess", target=tokens, index=self._num(0)) + return self._node("fieldAccess", target=first, field="x") + + def test_guarded_arm_is_not_walked(self) -> None: + # `count > 0 && tokens[0].x < 5` over an empty place: evaluation + # stops at the guard, so the margin must stop there too instead of + # indexing a token that is not there. + guarded = self._constraint( + self._node( + "binary", + op="&&", + left=self._node( + "binary", op=">", left=self._count(), right=self._num(0) + ), + right=self._node( + "binary", op="<", left=self._first_token_x(), right=self._num(5) + ), + ), + params=["state"], + ) + empty = {"places": {"Queue": {"count": 0, "tokens": []}}} + assert guarded(state=empty) is False + assert guarded.margin(state=empty) < 0 + + def test_equal_infinities_keep_the_margin_sign(self) -> None: + # inf - inf is NaN, but JS says inf <= inf and inf == inf hold, so a + # cancelled slack must not read as a violation. + inf = self._node("mathCall", fn="exp", args=[self._num(1000)]) + for op, satisfied in ( + ("<=", True), + (">=", True), + ("==", True), + ("<", False), + ("!=", False), + ): + case = self._constraint(self._node("binary", op=op, left=inf, right=inf)) + assert case(scenario={}) is satisfied, op + assert (case.margin(scenario={}) >= 0) == satisfied, op + + +class TestRejections: + def test_unknown_node_kind_fails_validation(self) -> None: + with pytest.raises(ValidationError, match="mystery"): + evaluate_hir( + { + "hirVersion": 1, + "surface": "scenario-expression", + "params": [], + "span": SPAN, + "body": {"kind": "mystery", "id": 0, "span": SPAN}, + } + ) + + def test_distribution_rejected(self) -> None: + # Well-formed, so validation passes; the evaluator refuses it. + with pytest.raises(HirEvaluationError, match="distribution"): + evaluate_hir( + { + "hirVersion": 1, + "surface": "scenario-expression", + "params": [], + "span": SPAN, + "body": { + "kind": "distribution", + "id": 0, + "span": SPAN, + "dist": "gaussian", + "args": [], + }, + } + ) + + def test_wrong_version_fails_validation(self) -> None: + with pytest.raises(ValidationError, match="hirVersion"): + evaluate_hir( + { + "hirVersion": 2, + "surface": "scenario-expression", + "params": [], + "span": SPAN, + "body": {"kind": "boolLit", "id": 0, "span": SPAN, "value": True}, + } + ) + + def test_non_boolean_constraint_result_raises(self) -> None: + fixture = FIXTURES["ordering"] + # Evaluate the raw comparison fine, but a Constraint demanding a + # boolean rejects a numeric body. + numeric = parameter_constraint( + "numeric", + { + "kind": "numberLit", + "id": 0, + "span": fixture["hir"]["span"], + "value": 2, + "raw": "2", + }, + code="1 + 1", + ) + with pytest.raises(HirEvaluationError, match="boolean"): + numeric(scenario={}) + + +class TestRunTimeShapeErrors: + """Well-formed HIR whose values do not fit at run time raises + HirEvaluationError, never a bare Python exception.""" + + @staticmethod + def _node(kind: str, **fields: object) -> dict[str, object]: + return {"kind": kind, "id": 0, "span": SPAN, **fields} + + def _num(self, value: float) -> dict[str, object]: + return self._node("numberLit", value=value, raw=repr(value)) + + def _eval(self, body: dict[str, object], **scenario: float) -> object: + return evaluate_hir( + { + "hirVersion": 1, + "surface": "scenario-expression", + "params": [], + "span": SPAN, + "body": body, + }, + scenario=scenario, + ) + + def test_math_extrema_follow_ecmascript(self) -> None: + assert self._eval(self._node("mathCall", fn="max", args=[])) == -math.inf + assert self._eval(self._node("mathCall", fn="min", args=[])) == math.inf + assert self._eval(self._node("mathCall", fn="max", args=[self._num(5)])) == 5 + + def test_math_arity_is_an_evaluation_error(self) -> None: + with pytest.raises(HirEvaluationError, match=r"Math\.sqrt"): + self._eval( + self._node("mathCall", fn="sqrt", args=[self._num(1), self._num(2)]) + ) + with pytest.raises(HirEvaluationError, match=r"Math\.atan2"): + self._eval(self._node("mathCall", fn="atan2", args=[self._num(1)])) + + def test_non_integer_index_is_an_evaluation_error(self) -> None: + array = self._node("arrayLit", elements=[self._num(1), self._num(2)]) + for index in (math.nan, math.inf, 0.9): + with pytest.raises(HirEvaluationError, match="not an integer"): + self._eval( + self._node( + "indexAccess", + target=array, + index=self._node("scenarioRef", name="k"), + ), + k=index, + ) + + def test_range_argument_count(self) -> None: + with pytest.raises(HirEvaluationError, match="range"): + self._eval(self._node("rangeCall", args=[])) + with pytest.raises(HirEvaluationError, match="range"): + self._eval( + self._node("rangeCall", args=[self._num(-1e308), self._num(1e308)]) + ) + + def test_overflowing_backward_range_is_empty(self) -> None: + # The span underflows to -Infinity; TypeScript ceilings it to zero + # elements, and so must the evaluator instead of calling it oversized. + assert ( + self._eval( + self._node("rangeCall", args=[self._num(1e308), self._num(-1e308)]) + ) + == [] + ) + + def test_huge_integers_coerce_like_js_numbers(self) -> None: + # An int past the double range is a valid Scalar; JS would read it + # as Infinity, and so must the evaluator instead of overflowing. + body = self._node( + "binary", + op="/", + left=self._node("scenarioRef", name="a"), + right=self._num(2), + ) + assert self._eval(body, a=10**400) == math.inf + + def test_infinite_dividend_remainder_is_nan(self) -> None: + inf = self._node("mathCall", fn="exp", args=[self._num(1000)]) + result = self._eval(self._node("binary", op="%", left=inf, right=self._num(7))) + assert isinstance(result, float) and math.isnan(result) diff --git a/libs/@local/petrinaut-python/uv.lock b/libs/@local/petrinaut-python/uv.lock index 7949865be99..e62cb3e5609 100644 --- a/libs/@local/petrinaut-python/uv.lock +++ b/libs/@local/petrinaut-python/uv.lock @@ -279,6 +279,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e8/3d/1087453384dbde46a8c7f9356eead2c58be8a7bf156bca40243377c85715/more_itertools-11.1.0-py3-none-any.whl", hash = "sha256:4b65538ae22f6fed0ce4874efd317463a7489796a0939fa66824dd542125a192", size = 72226, upload-time = "2026-05-22T14:14:28.824Z" }, ] +[[package]] +name = "mpmath" +version = "1.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e0/47/dd32fa426cc72114383ac549964eecb20ecfd886d1e5ccf5340b55b02f57/mpmath-1.3.0.tar.gz", hash = "sha256:7a28eb2a9774d00c7bc92411c19a89209d5da7c4c9a9e227be8330a23a25b91f", size = 508106, upload-time = "2023-03-07T16:47:11.061Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/43/e3/7d92a15f894aa0c9c4b49b8ee9ac9850d6e63b03c9c32c0367a13ae62209/mpmath-1.3.0-py3-none-any.whl", hash = "sha256:a0b2b9fe80bbcd81a6647ff13108738cfb482d481d826cc0e02f5b35e5c88d2c", size = 536198, upload-time = "2023-03-07T16:47:09.197Z" }, +] + [[package]] name = "mypy-extensions" version = "1.1.0" @@ -330,16 +339,26 @@ dependencies = [ { name = "pydantic" }, ] +[package.optional-dependencies] +sympy = [ + { name = "sympy" }, +] + [package.dev-dependencies] dev = [ { name = "basedpyright" }, { name = "datamodel-code-generator" }, { name = "pytest" }, { name = "ruff" }, + { name = "sympy" }, ] [package.metadata] -requires-dist = [{ name = "pydantic", specifier = ">=2.13.4" }] +requires-dist = [ + { name = "pydantic", specifier = ">=2.13.4" }, + { name = "sympy", marker = "extra == 'sympy'", specifier = ">=1.13" }, +] +provides-extras = ["sympy"] [package.metadata.requires-dev] dev = [ @@ -347,6 +366,7 @@ dev = [ { name = "datamodel-code-generator", specifier = ">=0.74.0" }, { name = "pytest", specifier = ">=8.3" }, { name = "ruff", specifier = ">=0.16.4" }, + { name = "sympy", specifier = ">=1.13" }, ] [[package]] @@ -653,6 +673,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/0e/fe/da8b9e1347696bb22120b77280ec5ce25d500ca5cb39d5ad6e5c18de19c1/ruff-0.16.4-py3-none-win_arm64.whl", hash = "sha256:a3a61621c9b6f6a89573e938a080e648f1695baa3f58570a3a707bc51ff65a21", size = 10451579, upload-time = "2026-08-20T17:43:57.135Z" }, ] +[[package]] +name = "sympy" +version = "1.14.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "mpmath" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/83/d3/803453b36afefb7c2bb238361cd4ae6125a569b4db67cd9e79846ba2d68c/sympy-1.14.0.tar.gz", hash = "sha256:d3d3fe8df1e5a0b42f0e7bdf50541697dbe7d23746e894990c030e2b05e72517", size = 7793921, upload-time = "2025-04-27T18:05:01.611Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a2/09/77d55d46fd61b4a135c444fc97158ef34a095e5681d0a6c10b75bf356191/sympy-1.14.0-py3-none-any.whl", hash = "sha256:e091cc3e99d2141a0ba2847328f5479b05d94a6635cb96148ccb3f34671bd8f5", size = 6299353, upload-time = "2025-04-27T18:04:59.103Z" }, +] + [[package]] name = "tomli" version = "2.4.1"