diff --git a/.changeset/connected-optimizer-source.md b/.changeset/connected-optimizer-source.md new file mode 100644 index 00000000000..342df51f6d8 --- /dev/null +++ b/.changeset/connected-optimizer-source.md @@ -0,0 +1,6 @@ +--- +"@hashintel/petrinaut": patch +"@hashintel/ds-components": patch +--- + +A connected optimization source runs studies in this browser behind the experimental In-browser optimization setting. The optimization form gains Runs per step and the experiments' Backend switch, which stays on the CPU because the GPU backend cannot compute an expression objective. A connected study's drawer streams the objective's metrics for the step being evaluated, and for whichever point the navigator or the surface picks once the study is over. The connected study's Surface draws only the study's steps — each a dot the field interpolates between, the best emphasized, pruned steps hollow — and fills in as the step in flight streams; it becomes navigable once the study is over or Follow steps is off, as do the Parameters band sliders. `Slider` accepts `disabled`. A connected study can be stopped and continued with more steps on the same sampler, settles its controls on the best step when it ends, stops refining a picked point that cannot beat the best after its first runs, evaluates up to four steps at once with a Parallel steps field, and lists every batch computing under the summary's progress bar. The connected study's drawer lays everything out in view at once: one summary strip (status, steps, best, backend), the Parameters band, the Surface beside the objective's chart, and the steps table filling the rest, the best step starred. The create form seeds each study with a fresh random **Seed**, editable for reproducibility. diff --git a/libs/@hashintel/ds-components/src/components/Slider/slider.tsx b/libs/@hashintel/ds-components/src/components/Slider/slider.tsx index 4930cbb14f1..e5d2e1bf2b6 100644 --- a/libs/@hashintel/ds-components/src/components/Slider/slider.tsx +++ b/libs/@hashintel/ds-components/src/components/Slider/slider.tsx @@ -28,6 +28,8 @@ export interface SliderProps { defaultValue?: number; label?: string; showValueText?: boolean; + /** Shows the value without letting the pointer or keyboard move it. */ + disabled?: boolean; onChange?: (value: number) => void; /** Fires once when a drag or keyboard interaction settles. */ onChangeEnd?: (value: number) => void; @@ -43,6 +45,7 @@ export const Slider: React.FC = ({ defaultValue, label, showValueText = false, + disabled, onChange, onChangeEnd, }) => { @@ -50,12 +53,17 @@ export const Slider: React.FC = ({ = { "compilation-output": "The Compilation bottom-panel tab: enabling it, the GPU verdict line, structural blockers, shader emission failures, per-item GPU/CPU/untested/no-HIR/unused status, and HIR node counts.", examples: - "Walkthroughs of the built-in examples and the scenarios/metrics each ships with: SIR, Supply Chain with Disruption, Supply Chain Profit, Deployment Pipeline, Production with Machine Failure, Probabilistic Satellite Launcher, Café Queue, Drone Patrol.", + "Walkthroughs of the built-in examples and the scenarios/metrics each ships with: SIR, Vaccination Campaign, Supply Chain with Disruption, Supply Chain Profit, Deployment Pipeline, Production with Machine Failure, Probabilistic Satellite Launcher, Café Queue, Drone Patrol.", }; const getLatestNetDefinitionToolInputSchema = z diff --git a/libs/@hashintel/petrinaut-core/src/examples/examples.test.ts b/libs/@hashintel/petrinaut-core/src/examples/examples.test.ts index 7fe4bb387c4..327ba705c0a 100644 --- a/libs/@hashintel/petrinaut-core/src/examples/examples.test.ts +++ b/libs/@hashintel/petrinaut-core/src/examples/examples.test.ts @@ -10,6 +10,7 @@ import { sirModel, supplyChainProfit, supplyChainWithDisruption, + vaccinationCampaign, } from "./index"; const EXAMPLES = [ @@ -19,6 +20,7 @@ const EXAMPLES = [ sirModel, supplyChainProfit, supplyChainWithDisruption, + vaccinationCampaign, ]; describe.each(EXAMPLES.map((example) => [example.title, example] as const))( diff --git a/libs/@hashintel/petrinaut-core/src/examples/index.ts b/libs/@hashintel/petrinaut-core/src/examples/index.ts index 8ae051588b8..4fe133bc8e8 100644 --- a/libs/@hashintel/petrinaut-core/src/examples/index.ts +++ b/libs/@hashintel/petrinaut-core/src/examples/index.ts @@ -11,3 +11,4 @@ export { cafeQueue } from "./cafe-queue"; export { dronePatrol } from "./drone-patrol"; export { supplyChainWithDisruption } from "./supply-chain-with-disruption"; export { supplyChainProfit } from "./supply-chain-profit"; +export { vaccinationCampaign } from "./vaccination-campaign"; diff --git a/libs/@hashintel/petrinaut-core/src/examples/vaccination-campaign.test.ts b/libs/@hashintel/petrinaut-core/src/examples/vaccination-campaign.test.ts new file mode 100644 index 00000000000..6da04d132ad --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/examples/vaccination-campaign.test.ts @@ -0,0 +1,165 @@ +import { describe, expect, it } from "vitest"; + +import { compileHirArtifacts } from "../hir"; +import { lowerScenarioToHir } from "../hir/scenario"; +import { compileScenario } from "../simulation/authoring/scenario/compile-scenario"; +import { + createMonteCarloExperiment, + runExperimentToCompletion, +} from "../simulation/monte-carlo"; +import { analyzeCompilation } from "../webgpu/compilation-report"; +import { assessGpuEligibility } from "../webgpu/eligibility"; +import { vaccinationCampaign } from "./vaccination-campaign"; + +import type { CompiledScenarioResult } from "../simulation/authoring/scenario/compile-scenario"; + +const { petriNetDefinition } = vaccinationCampaign; + +const winterWave = petriNetDefinition.scenarios!.find( + (scenario) => scenario.id === "scenario__winter_wave", +)!; +const totalCost = petriNetDefinition.metrics!.find( + (metric) => metric.id === "metric__total_cost", +)!; +const infectedPlace = petriNetDefinition.places.find( + (place) => place.name === "Infected", +)!; + +const { artifacts } = compileHirArtifacts(petriNetDefinition, undefined, { + includeHir: true, +}); +const winterWaveHir = lowerScenarioToHir(winterWave); + +/** The two levers the optimization stories range over. */ +type Levers = { vaccination_coverage: number; contact_reduction: number }; + +const compile = (levers?: Levers): CompiledScenarioResult => { + const outcome = compileScenario( + winterWave, + winterWaveHir, + petriNetDefinition.parameters, + petriNetDefinition.places, + petriNetDefinition.types, + levers ? { scenarioParameterValues: levers } : undefined, + ); + if (!outcome.ok) { + throw new Error( + `scenario failed to compile: ${outcome.errors + .map((error) => error.message) + .join("; ")}`, + ); + } + return outcome.result; +}; + +/** Mean Total cost on the final state over eight seeded 60-day runs. */ +const meanTotalCost = async (levers: Levers): Promise => { + const compiled = compile(levers); + const runCount = 8; + const handle = await createMonteCarloExperiment({ + sdcpn: petriNetDefinition, + hirArtifacts: artifacts, + initialMarking: compiled.initialState, + parameterValues: compiled.parameterValues, + seed: 1, + dt: 0.1, + maxTime: 60, + runCount, + runs: Array.from({ length: runCount }, (_, index) => ({ + seed: 1000 + index, + })), + metricSpecs: [ + { + kind: "expression", + id: totalCost.id, + label: totalCost.name, + sampleRuns: "all", + code: totalCost.code, + artifact: artifacts.metrics[totalCost.id]!, + }, + ], + }); + const completion = await runExperimentToCompletion(handle); + if (completion.event.type !== "complete") { + throw new Error(`experiment ended with ${completion.event.type}`); + } + let sum = 0; + for (const result of completion.runResults.values()) { + sum += result[totalCost.id]!; + } + return sum / completion.runResults.size; +}; + +describe("Vaccination Campaign", () => { + it("is GPU-eligible as an uncoloured net", () => { + const result = assessGpuEligibility(petriNetDefinition); + + expect(result.eligible).toBe(true); + if (!result.eligible) return; + expect(result.profile.uncolouredOnly).toBe(true); + // Four counts, two firing counts, rng, status = 8 words. + expect(result.profile.bytesPerRun).toBe(32); + }); + + it("compiles to a GPU shader with a place-count objective", () => { + const report = analyzeCompilation({ + sdcpn: petriNetDefinition, + artifacts, + metricSpecs: [ + { + id: "infected", + label: "Infected", + kind: "placeTokenCountMean", + placeId: infectedPlace.id, + }, + ], + }); + + expect(report.gpuReady).toBe(true); + expect(report.eligibilityReasons).toStrictEqual([]); + expect(report.shaderFailure).toBeNull(); + expect(report.metricFailure).toBeNull(); + expect( + report.items + .filter((item) => item.kind === "lambda") + .map((item) => item.status), + ).toStrictEqual(["gpu-ready", "gpu-ready"]); + }); + + it("seeds the Winter wave from the coverage and the initial cases", () => { + const result = compile(); + + expect(result.initialState).toEqual({ + place__susceptible: 686, + place__infected: 20, + place__recovered: 0, + place__vaccinated: 294, + }); + expect(Number(result.parameterValues.vaccination_coverage)).toBeCloseTo( + 0.3, + ); + expect(Number(result.parameterValues.contact_reduction)).toBeCloseTo(0.2); + }); + + it("prices the cheapest response inside the levers' domain", async () => { + const floor = await meanTotalCost({ + vaccination_coverage: 0.45, + contact_reduction: 0.4, + }); + const boundary: Levers[] = [ + { vaccination_coverage: 0, contact_reduction: 0 }, + { vaccination_coverage: 0, contact_reduction: 0.4 }, + { vaccination_coverage: 0, contact_reduction: 0.8 }, + { vaccination_coverage: 0.45, contact_reduction: 0 }, + { vaccination_coverage: 0.9, contact_reduction: 0 }, + { vaccination_coverage: 0.9, contact_reduction: 0.8 }, + ]; + + for (const point of boundary) { + expect( + await meanTotalCost(point), + `coverage ${point.vaccination_coverage}, contact reduction ${point.contact_reduction} costs more than the valley floor`, + ).toBeGreaterThan(floor); + } + }); +}); diff --git a/libs/@hashintel/petrinaut-core/src/examples/vaccination-campaign.ts b/libs/@hashintel/petrinaut-core/src/examples/vaccination-campaign.ts new file mode 100644 index 00000000000..6978ee0fbe3 --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/examples/vaccination-campaign.ts @@ -0,0 +1,263 @@ +import { GRID_SIZE } from "../grid-size"; + +import type { SDCPN } from "../types/sdcpn"; + +/** + * Vaccination campaign — an SIR wave with two policy levers and a cost account, + * built as the model to optimize. + * + * Susceptible, Infected, Recovered and Vaccinated are plain counts. Infection + * (`S + I -> 2I`) fires at `infection_rate` scaled down by the two levers: + * `contact_reduction` (distancing) and `vaccine_efficacy × vaccination_coverage` + * (the share of contacts that land on a protected person). Recovery moves + * Infected to Recovered at `recovery_rate`. The wave persists for the whole + * horizon while the scaled infection rate exceeds the recovery rate and dies + * out below it, so the case count bends sharply around that threshold. + * + * The Winter wave scenario seeds `Vaccinated` from `vaccination_coverage` and + * exposes both levers as scenario parameters. The Total cost metric charges + * every case at `case_cost` and each lever at a price quadratic in its + * intensity (`campaign_cost` and `distancing_cost`, per head at full + * intensity), so both levers have diminishing returns against a rising price + * and the minimum lies inside the domain rather than on a bound: a shallow + * valley along the epidemic threshold with its floor near a coverage of 0.45 + * and a contact reduction of 0.4 for the default costs, at about 960 against + * 1,280–2,220 in the corners over a 60-day horizon. + * + * GPU-ready as shipped: uncoloured places, rates that read only parameters, + * and place counts as the experiment observables. + */ +export const vaccinationCampaign: { title: string; petriNetDefinition: SDCPN } = + { + title: "Vaccination Campaign", + petriNetDefinition: { + places: [ + { + id: "place__susceptible", + name: "Susceptible", + colorId: null, + dynamicsEnabled: false, + differentialEquationId: null, + showAsInitialState: true, + x: -29 * GRID_SIZE, + y: 10 * GRID_SIZE, + }, + { + id: "place__infected", + name: "Infected", + colorId: null, + dynamicsEnabled: false, + differentialEquationId: null, + showAsInitialState: true, + x: -13 * GRID_SIZE, + y: 19 * GRID_SIZE, + }, + { + id: "place__recovered", + name: "Recovered", + colorId: null, + dynamicsEnabled: false, + differentialEquationId: null, + x: 25 * GRID_SIZE, + y: 13 * GRID_SIZE, + }, + { + id: "place__vaccinated", + name: "Vaccinated", + colorId: null, + dynamicsEnabled: false, + differentialEquationId: null, + showAsInitialState: true, + x: -29 * GRID_SIZE, + y: -4 * GRID_SIZE, + }, + ], + transitions: [ + { + id: "transition__infection", + name: "Infection", + inputArcs: [ + { + placeId: "place__susceptible", + weight: 1, + type: "standard", + }, + { + placeId: "place__infected", + weight: 1, + type: "standard", + }, + ], + outputArcs: [ + { + placeId: "place__infected", + weight: 2, + }, + ], + lambdaType: "stochastic", + lambdaCode: `// Infectious contacts per day, cut by distancing and by the share of +// contacts that land on a protected (vaccinated and immune) person. +const distancing = 1 - parameters.contact_reduction; +const protection = 1 - parameters.vaccine_efficacy * parameters.vaccination_coverage; +return parameters.infection_rate * distancing * protection;`, + transitionKernelCode: `// Consumes 1 Susceptible + 1 Infected and produces 2 Infected (the output +// arc has weight 2): the susceptible has become newly infected. +return { + Infected: [{}, {}], +};`, + x: -10 * GRID_SIZE, + y: 5 * GRID_SIZE, + }, + { + id: "transition__recovery", + name: "Recovery", + inputArcs: [ + { + placeId: "place__infected", + weight: 1, + type: "standard", + }, + ], + outputArcs: [ + { + placeId: "place__recovered", + weight: 1, + }, + ], + lambdaType: "stochastic", + lambdaCode: `// Recoveries per day. The wave dies out once the scaled infection rate +// falls below this rate, and persists for the whole horizon above it. +return parameters.recovery_rate;`, + transitionKernelCode: `// Move one Infected to Recovered (1-to-1). +return { + Recovered: [{}], +};`, + x: 6 * GRID_SIZE, + y: 16 * GRID_SIZE, + }, + ], + types: [], + differentialEquations: [], + parameters: [ + { + id: "param__infection_rate", + name: "Infection Rate", + variableName: "infection_rate", + type: "real", + defaultValue: "3", + }, + { + id: "param__recovery_rate", + name: "Recovery Rate", + variableName: "recovery_rate", + type: "real", + defaultValue: "2", + }, + { + id: "param__vaccine_efficacy", + name: "Vaccine Efficacy", + variableName: "vaccine_efficacy", + type: "real", + defaultValue: "0.9", + }, + { + id: "param__vaccination_coverage", + name: "Vaccination Coverage", + variableName: "vaccination_coverage", + type: "real", + defaultValue: "0", + }, + { + id: "param__contact_reduction", + name: "Contact Reduction", + variableName: "contact_reduction", + type: "real", + defaultValue: "0", + }, + { + id: "param__case_cost", + name: "Cost per Case", + variableName: "case_cost", + type: "real", + defaultValue: "10", + }, + { + id: "param__campaign_cost", + name: "Campaign Cost per Head at Full Coverage", + variableName: "campaign_cost", + type: "real", + defaultValue: "1.2", + }, + { + id: "param__distancing_cost", + name: "Distancing Cost per Head at Full Reduction", + variableName: "distancing_cost", + type: "real", + defaultValue: "1.6", + }, + ], + scenarios: [ + { + id: "scenario__winter_wave", + name: "Winter wave", + description: + "A town of 1,000 seeded with 20 cases. Vaccination coverage is set before the wave and contact reduction holds for its whole run; optimize both against Total cost to find the cheapest mix.", + scenarioParameters: [ + { type: "integer", identifier: "population", default: 1000 }, + { type: "integer", identifier: "initial_infected", default: 20 }, + { type: "ratio", identifier: "vaccination_coverage", default: 0.3 }, + { type: "ratio", identifier: "contact_reduction", default: 0.2 }, + ], + parameterOverrides: { + param__vaccination_coverage: "scenario.vaccination_coverage", + param__contact_reduction: "scenario.contact_reduction", + }, + initialState: { + type: "per_place", + content: { + place__susceptible: + "scenario.population - scenario.initial_infected - Math.round((scenario.population - scenario.initial_infected) * scenario.vaccination_coverage)", + place__infected: "scenario.initial_infected", + place__recovered: "0", + place__vaccinated: + "Math.round((scenario.population - scenario.initial_infected) * scenario.vaccination_coverage)", + }, + }, + }, + ], + metrics: [ + { + id: "metric__total_cost", + name: "Total cost", + description: + "Cases, the vaccination campaign and distancing priced together: the objective to minimize on the final state.", + code: `const cases = state.places.Infected.count + state.places.Recovered.count; +const population = + state.places.Susceptible.count + state.places.Vaccinated.count + cases; +const coverage = parameters.vaccination_coverage; +const reduction = parameters.contact_reduction; +return ( + cases * parameters.case_cost + + coverage * coverage * parameters.campaign_cost * population + + reduction * reduction * parameters.distancing_cost * population +);`, + }, + { + id: "metric__infected", + name: "Infected", + description: + "People currently infected: the wave's curve, dying out or growing.", + code: `return state.places.Infected.count;`, + }, + { + id: "metric__attack_rate", + name: "Attack rate", + description: "Share of the population infected so far.", + code: `const cases = state.places.Infected.count + state.places.Recovered.count; +const population = + state.places.Susceptible.count + state.places.Vaccinated.count + cases; +return population === 0 ? 0 : cases / population;`, + }, + ], + }, + }; diff --git a/libs/@hashintel/petrinaut-core/src/webgpu/compilation-report.test.ts b/libs/@hashintel/petrinaut-core/src/webgpu/compilation-report.test.ts index ac4a7187b3d..2a97266d8ab 100644 --- a/libs/@hashintel/petrinaut-core/src/webgpu/compilation-report.test.ts +++ b/libs/@hashintel/petrinaut-core/src/webgpu/compilation-report.test.ts @@ -92,6 +92,7 @@ describe("analyzeCompilation", () => { dronePatrol: true, supplyChainWithDisruption: true, supplyChainProfit: true, + vaccinationCampaign: true, }); const production = analyze( allExamples.productionMachines.petriNetDefinition, diff --git a/libs/@hashintel/petrinaut/docs/examples.md b/libs/@hashintel/petrinaut/docs/examples.md index 6f75c584000..fd61b9ef5b8 100644 --- a/libs/@hashintel/petrinaut/docs/examples.md +++ b/libs/@hashintel/petrinaut/docs/examples.md @@ -22,6 +22,22 @@ The classic Susceptible-Infected-Recovered compartmental model from epidemiology SIR +## Vaccination Campaign + +The SIR model with two policy levers and a cost account, built as the model to optimize: a town of 1,000 seeded with 20 cases, a vaccination campaign whose coverage is set before the wave, and distancing that holds for its whole run. + +**Demonstrates:** + +- **Parameter-driven rates** -- Infection fires at `infection_rate` scaled by `(1 - contact_reduction)` and by `(1 - vaccine_efficacy × vaccination_coverage)`, the share of contacts that land on an unprotected person; Recovery fires at `recovery_rate`. The wave persists while the scaled infection rate exceeds the recovery rate and dies out below it. +- **Scenario parameters wired to the initial state and the rates** -- the _Winter wave_ scenario seeds `Vaccinated` from `vaccination_coverage` and overrides both lever parameters, so an optimization or a sweep over the levers changes the initial marking and the rates together. +- **An objective with an interior optimum** -- the **Total cost** [metric](simulation.md) charges every case at `case_cost` and each lever at a price quadratic in its intensity (`campaign_cost`, `distancing_cost`), so both levers have diminishing returns against a rising price. Over a 60-day horizon the cost is about 960 near a coverage of 0.45 and a contact reduction of 0.4, against 1,280 to 2,220 in the corners of the domain. +- **GPU-ready modelling** -- untyped places and rates that read only parameters, so an experiment measuring the **Infected** place's token count (**Built-in › Place tokens**) runs on the GPU backend as shipped. The model metric of the same name is an expression, which keeps an experiment on the CPU. +- Two further metrics -- **Infected** (the wave's curve, dying out or growing) and **Attack rate** (share of the population infected so far). + +**Suggested initial state:** pick **Winter wave** and, in the Optimizations tab, minimize **Total cost** over `vaccination_coverage` (0 to 0.9) and `contact_reduction` (0 to 0.8) with a max time of 60: the surface shows a valley along the epidemic threshold and the steps settle around a coverage of 0.45 and a contact reduction of 0.4. To watch a single run instead, press Play and select the **Infected** metric in the timeline. + +**Key concepts:** [stochastic firing](petri-net-extensions.md#stochastic-rate), [parameters](petri-net-extensions.md#global-parameters), [scenarios](scenarios.md), [optimization objectives](useful-patterns.md#optimization-objectives-metrics-that-read-parameters), [optimization](optimization.md). + ## Café Queue A small service system: customers arrive, wait, are served by a limited staff pool, and leave. Built to run on the **GPU compute backend out of the box** — every place is untyped, and the interesting measurements are plain token counts. diff --git a/libs/@hashintel/petrinaut/src/main.ts b/libs/@hashintel/petrinaut/src/main.ts index 97bac2f2784..71c3f06398c 100644 --- a/libs/@hashintel/petrinaut/src/main.ts +++ b/libs/@hashintel/petrinaut/src/main.ts @@ -12,7 +12,14 @@ export type { ErrorTracker } from "./react/error-tracker-context"; export { ErrorTrackerContext } from "./react/error-tracker-context"; -export type { PetrinautOptimization } from "./react/optimization-context"; +export type { + PetrinautConnectedOptimization, + PetrinautConnectedOptimizationCapability, + PetrinautConnectedRunOptions, + PetrinautOptimization, + PetrinautOptimizationChannel, + PetrinautOptimizationSource, +} from "./react/optimization-context"; export { PetrinautOptimizationContext } from "./react/optimization-context"; export type { PetrinautSlots } from "./ui/types/petrinaut-slots"; diff --git a/libs/@hashintel/petrinaut/src/react/experiments/context.ts b/libs/@hashintel/petrinaut/src/react/experiments/context.ts index df25a7d4133..707b2e50973 100644 --- a/libs/@hashintel/petrinaut/src/react/experiments/context.ts +++ b/libs/@hashintel/petrinaut/src/react/experiments/context.ts @@ -18,6 +18,7 @@ import type { MonteCarloMetricSpec, MonteCarloUserDefinedMetricFrame, MonteCarloWorkerProgress, + ReadableStore, } from "@hashintel/petrinaut-core"; export type ExperimentStatus = @@ -227,6 +228,17 @@ export type ExperimentsContextValue = { sampleDetachedObjective: ( request: DetachedObjectiveRequest, ) => Promise; + /** + * Streams one batch of a study's objective at one parameter point on the + * requested backend: the in-browser optimizer's trials and the study + * drawer's selected-point refinement. Batches queue per `queueKey` (the + * `cacheKey` by default); different keys run side by side. The + * returned run never rejects — refusal, failure and cancellation all + * settle `completion` with a failed outcome naming the reason. + */ + runDetachedObjective: ( + request: DetachedObjectiveRunRequest, + ) => DetachedObjectiveRun; }; /** One local compute batch for an optimization study's objective. */ @@ -246,6 +258,63 @@ export type DetachedObjectiveRequest = { maxTime: number; }; +export type DetachedObjectiveRunRequest = DetachedObjectiveRequest & { + /** + * Pinned per-run seeds, `runCount` long; CPU only. Absent (and always on + * the GPU, which derives every run's seed from `seed`), runs derive their + * seeds from `seed`. + */ + runSeeds?: readonly number[]; + /** + * Runs sharing a queue key run one at a time, in order; runs with + * different keys overlap. Defaults to `cacheKey`, so a study's batches + * queue unless the caller gives each its own key. + */ + queueKey?: string; + computeBackend: ExperimentComputeBackend; + signal?: AbortSignal; +}; + +export type DetachedObjectiveRunResult = { + runsCompleted: number; + metricFrames: readonly MonteCarloUserDefinedMetricFrame[]; + /** Per-run final metric values; empty on the GPU, which reports no run axis. */ + runResults: ReadonlyMap>>; + /** Where the batch ran. */ + computeBackend: ExperimentComputeBackend; + /** Why the requested backend declined, when the batch ran elsewhere. */ + computeBackendFallbackReason: string | null; +}; + +/** + * How a batch ended. A failure carries a reason the user can act on: the + * diagnostics of a metric that did not compile, each backend that declined + * and why, how many runs errored. `cancelled` marks a batch stopped through + * `cancel` or the request's signal, which nobody needs to act on. + */ +export type DetachedObjectiveRunOutcome = + | ({ readonly ok: true } & DetachedObjectiveRunResult) + | { + readonly ok: false; + readonly reason: string; + readonly cancelled: boolean; + }; + +/** One streaming batch for a study's objective at one parameter point. */ +export type DetachedObjectiveRun = { + /** Frames so far; replaced as the batch streams, at most every 100 ms. */ + readonly frames: ReadableStore; + readonly progress: ReadableStore; + /** Settles on the terminal event; never rejects. */ + readonly completion: Promise; + cancel(this: void): void; +}; + +const constantStore = (value: T): ReadableStore => ({ + get: () => value, + subscribe: () => () => {}, +}); + const DEFAULT_CONTEXT_VALUE: ExperimentsContextValue = { experiments: [], selectedExperimentId: null, @@ -257,6 +326,16 @@ const DEFAULT_CONTEXT_VALUE: ExperimentsContextValue = { setSweepSelection: () => {}, sampleSurfaceCells: () => Promise.resolve(null), sampleDetachedObjective: () => Promise.resolve(null), + runDetachedObjective: () => ({ + frames: constantStore([]), + progress: constantStore(null), + completion: Promise.resolve({ + ok: false, + cancelled: false, + reason: "Experiments are unavailable", + }), + cancel: () => {}, + }), }; export const ExperimentsContext = createContext( @@ -278,6 +357,7 @@ export type ExperimentsActionsValue = Pick< | "setSweepSelection" | "sampleSurfaceCells" | "sampleDetachedObjective" + | "runDetachedObjective" >; export const ExperimentsActionsContext = createContext( diff --git a/libs/@hashintel/petrinaut/src/react/experiments/provider.tsx b/libs/@hashintel/petrinaut/src/react/experiments/provider.tsx index 070231f352b..dfa658853c9 100644 --- a/libs/@hashintel/petrinaut/src/react/experiments/provider.tsx +++ b/libs/@hashintel/petrinaut/src/react/experiments/provider.tsx @@ -175,7 +175,10 @@ export const ExperimentsProvider: React.FC = ({ const pendingRegistrations = pendingRegistrationsRef.current; const sweepSessions = sweepSessionsRef.current; const chosenBackends = backendsRef.current; + const detachedObjectiveSampler = detachedObjectiveSamplerRef; return () => { + detachedObjectiveSampler.current?.dispose(); + detachedObjectiveSampler.current = null; for (const registration of pendingRegistrations.values()) { registration.abortController.abort(); } @@ -658,23 +661,25 @@ export const ExperimentsProvider: React.FC = ({ const stableRemoveExperiment = useStableCallback(removeExperiment); const stableSetSweepSelection = useStableCallback(setSweepSelection); const stableSampleSurfaceCells = useStableCallback(sampleSurfaceCells); + // Built on first use: a session that never opens an optimization surface + // or runs a study in the browser spawns no extra worker lane. + const getDetachedObjectiveSampler = (): DetachedObjectiveSampler => { + detachedObjectiveSamplerRef.current ??= createDetachedObjectiveSampler({ + languageClient: languageClientRef, + createWorker: reusableWorkerFactory, + shardCount: shardCountRef.current ?? getDefaultMonteCarloShardCount(), + }); + return detachedObjectiveSamplerRef.current; + }; const sampleDetachedObjective: ExperimentsContextValue["sampleDetachedObjective"] = - (request) => { - // Built on first use: a session that never opens an optimization - // surface spawns no extra worker lane. - const sampler = - detachedObjectiveSamplerRef.current ?? - createDetachedObjectiveSampler({ - languageClient: languageClientRef, - createWorker: reusableWorkerFactory, - }); - detachedObjectiveSamplerRef.current = sampler; - return sampler.sample(request); - }; + (request) => getDetachedObjectiveSampler().sample(request); + const runDetachedObjective: ExperimentsContextValue["runDetachedObjective"] = + (request) => getDetachedObjectiveSampler().run(request); const stableSampleDetachedObjective = useStableCallback( sampleDetachedObjective, ); + const stableRunDetachedObjective = useStableCallback(runDetachedObjective); // Every callback is identity-stable, so this object never changes and // actions-only consumers sit out the per-publish re-render storm. @@ -686,6 +691,7 @@ export const ExperimentsProvider: React.FC = ({ setSweepSelection: stableSetSweepSelection, sampleSurfaceCells: stableSampleSurfaceCells, sampleDetachedObjective: stableSampleDetachedObjective, + runDetachedObjective: stableRunDetachedObjective, })); const contextValue: ExperimentsContextValue = { diff --git a/libs/@hashintel/petrinaut/src/react/experiments/provider/detached-objective.test.ts b/libs/@hashintel/petrinaut/src/react/experiments/provider/detached-objective.test.ts new file mode 100644 index 00000000000..0119f5c8193 --- /dev/null +++ b/libs/@hashintel/petrinaut/src/react/experiments/provider/detached-objective.test.ts @@ -0,0 +1,601 @@ +import { describe, expect, it, vi } from "vitest"; + +import { sirModel } from "@hashintel/petrinaut-core/examples"; +import { WORKER_POOL_BACKEND_ID } from "@hashintel/petrinaut-core/experiments"; +import { + compileHirArtifacts, + lowerScenarioToHir, +} from "@hashintel/petrinaut-core/hir"; + +import { createDetachedObjectiveSampler } from "./detached-objective"; +import { createWritableStore } from "./detached-objective/writable-store"; + +import type { LanguageClientContextValue } from "../../lsp/context"; +import type { DetachedObjectiveRunRequest } from "../context"; +import type { experimentBackendRegistrations } from "./create-experiment"; +import type { + AbortSignalLike, + MonteCarloExperiment, + MonteCarloExperimentEvent, + MonteCarloExperimentMetrics, + MonteCarloExperimentState, + MonteCarloUserDefinedMetricFrame, + MonteCarloWorkerProgress, +} from "@hashintel/petrinaut-core"; +import type { + ExperimentAssessment, + ExperimentBackend, + ExperimentRequest, + ReusableWorkerFactory, +} from "@hashintel/petrinaut-core/experiments"; + +const scenario = sirModel.petriNetDefinition.scenarios?.find( + (candidate) => candidate.id === "scenario__seasonal_flu", +); +const metric = sirModel.petriNetDefinition.metrics?.find( + (candidate) => candidate.id === "metric__infected_fraction", +); +if (!scenario || !metric) { + throw new Error("The SIR fixtures are incomplete"); +} +const definition = { + ...sirModel.petriNetDefinition, + scenarios: [scenario], + metrics: [metric], +}; + +const runRequest = ( + overrides: Partial = {}, +): DetachedObjectiveRunRequest => ({ + cacheKey: "study", + definition, + scenarioId: scenario.id, + scenarioParameterValues: { population: 1_000, infected_ratio: 0.05 }, + metric: { id: metric.id, label: metric.name, code: metric.code }, + seed: 7, + runCount: 3, + runSeeds: [7, 11, 13], + dt: 1, + maxTime: 180, + computeBackend: "cpu", + ...overrides, +}); + +const progressOf = ( + completedRuns: number, + erroredRuns = 0, +): MonteCarloWorkerProgress => ({ + activeRuns: 0, + advancedRuns: completedRuns, + allFinished: completedRuns + erroredRuns >= 3, + completedRuns, + erroredRuns, + frameNumber: 180, + runCount: 3, + time: 180, +}); + +const frameOf = (value: number): MonteCarloUserDefinedMetricFrame => ({ + metricId: metric.id, + label: metric.name, + outputType: "distribution", + frameNumber: 1, + time: 1, + bins: [[value, 3]], + value: null, + frameValue: null, + timeValue: null, + runSampleCount: 3, + timeSampleCount: 0, +}); + +type FakeHandle = { + handle: MonteCarloExperiment; + metrics: ReturnType>; + progress: ReturnType< + typeof createWritableStore + >; + runResults: ReturnType< + typeof createWritableStore< + ReadonlyMap>> + > + >; + emit: (event: MonteCarloExperimentEvent) => void; +}; + +/** + * A handle shaped like the worker pool's: a cancel is answered by the shards + * a tick later, and a handle whose instantiation signal fires drops its shard + * listeners at once, so a cancel after that is never answered. + */ +const createFakeHandle = (signal?: AbortSignalLike): FakeHandle => { + let tornDown = false; + signal?.addEventListener( + "abort", + () => { + tornDown = true; + }, + { once: true }, + ); + const status = createWritableStore("Ready"); + const progress = createWritableStore(null); + const metrics = createWritableStore({ + frames: [], + latestByMetricId: {}, + }); + const runResults = createWritableStore< + ReadonlyMap>> + >(new Map()); + const listeners = new Set<(event: MonteCarloExperimentEvent) => void>(); + const emit = (event: MonteCarloExperimentEvent) => { + for (const listener of listeners) { + listener(event); + } + }; + const handle: MonteCarloExperiment = { + status, + progress, + metrics, + runResults, + events: { + subscribe: (listener) => { + listeners.add(listener); + return () => { + listeners.delete(listener); + }; + }, + }, + start: vi.fn(), + cancel: vi.fn(() => { + queueMicrotask(() => { + if (!tornDown) { + emit({ type: "cancelled", progress: progress.get() }); + } + }); + }), + dispose: vi.fn(), + }; + return { handle, metrics, progress, runResults, emit }; +}; + +type FakeBackend = { + backend: ExperimentBackend; + requests: ExperimentRequest[]; + handles: FakeHandle[]; +}; + +const createFakeBackend = ( + id: string, + options: { + refuse?: string; + /** Index of the first request refused; earlier ones are accepted. */ + refuseFrom?: number; + needsHirTrees?: boolean; + } = {}, +): FakeBackend => { + const requests: ExperimentRequest[] = []; + const handles: FakeHandle[] = []; + const backend: ExperimentBackend = { + id, + label: id, + needsHirTrees: options.needsHirTrees ?? false, + isAvailable: () => true, + assess: (request) => { + const requestIndex = requests.length; + requests.push(request); + const assessment: ExperimentAssessment = + options.refuse === undefined || requestIndex < (options.refuseFrom ?? 0) + ? { + eligible: true, + notes: [], + instantiate: (instantiateOptions) => { + const fake = createFakeHandle(instantiateOptions?.signal); + handles.push(fake); + return Promise.resolve({ ok: true, handle: fake.handle }); + }, + } + : { + eligible: false, + blockers: [ + { code: "refused", message: options.refuse, origin: "model" }, + ], + }; + return Promise.resolve(assessment); + }, + dispose: vi.fn(), + }; + return { backend, requests, handles }; +}; + +/** Compiles inline what the language worker compiles in the app. */ +const languageClient: Pick< + LanguageClientContextValue, + "requestHirArtifacts" | "requestScenarioHir" +> = { + requestHirArtifacts: (sdcpn, extensions, options) => + Promise.resolve(compileHirArtifacts(sdcpn, extensions, options)), + requestScenarioHir: (candidate, adHocContext) => + Promise.resolve(lowerScenarioToHir(candidate, { adHocContext })), +}; + +const unusedWorkerFactory = Object.assign( + () => Promise.reject(new Error("The fake backends lease no workers")), + { drain: () => {}, dispose: () => {} }, +) as ReusableWorkerFactory; + +const createSampler = (backends: { cpu: FakeBackend; gpu?: FakeBackend }) => { + const registrations = vi.fn( + ({ + computeBackend, + }: Parameters[0]) => [ + ...(computeBackend === "webgpu" && backends.gpu + ? [ + { + id: "webgpu", + label: "GPU", + load: () => Promise.resolve(backends.gpu!.backend), + }, + ] + : []), + { + id: WORKER_POOL_BACKEND_ID, + label: "CPU", + load: () => Promise.resolve(backends.cpu.backend), + }, + ], + ); + const sampler = createDetachedObjectiveSampler({ + languageClient: { current: languageClient }, + createWorker: unusedWorkerFactory, + shardCount: 6, + backendRegistrations: registrations, + }); + return { sampler, registrations }; +}; + +const completeWith = (fake: FakeHandle, value: number) => { + fake.metrics.set({ + frames: [frameOf(value)], + latestByMetricId: { [metric.id]: frameOf(value) }, + }); + fake.runResults.set( + new Map([ + [0, { [metric.id]: value }], + [1, { [metric.id]: value }], + [2, { [metric.id]: value }], + ]), + ); + fake.emit({ type: "complete", progress: progressOf(3) }); +}; + +describe("createDetachedObjectiveSampler().run", () => { + it("streams frames and progress, then settles the result with the seeds pinned on the CPU pool", async () => { + const cpu = createFakeBackend(WORKER_POOL_BACKEND_ID); + const { sampler } = createSampler({ cpu }); + + const run = sampler.run(runRequest()); + await vi.waitFor(() => expect(cpu.handles).toHaveLength(1)); + const fake = cpu.handles[0]!; + expect(fake.handle.start).toHaveBeenCalledOnce(); + expect(cpu.requests[0]).toMatchObject({ + seed: 7, + runCount: 3, + runs: [{ seed: 7 }, { seed: 11 }, { seed: 13 }], + metricSpecs: [ + { + kind: "expression", + id: metric.id, + sampleRuns: "all", + runOutput: { type: "distribution" }, + }, + ], + }); + + fake.progress.set(progressOf(1)); + fake.metrics.set({ + frames: [frameOf(0.2)], + latestByMetricId: { [metric.id]: frameOf(0.2) }, + }); + await vi.waitFor(() => expect(run.frames.get()).toEqual([frameOf(0.2)])); + expect(run.progress.get()).toEqual(progressOf(1)); + + completeWith(fake, 0.25); + const outcome = await run.completion; + expect(outcome).toMatchObject({ + ok: true, + runsCompleted: 3, + metricFrames: [frameOf(0.25)], + computeBackend: "cpu", + computeBackendFallbackReason: null, + }); + expect(outcome.ok && outcome.runResults.get(2)).toEqual({ + [metric.id]: 0.25, + }); + expect(run.frames.get()).toEqual([frameOf(0.25)]); + expect(run.progress.get()).toEqual(progressOf(3)); + expect(fake.handle.dispose).toHaveBeenCalled(); + }); + + it("names why a batch failed: errored runs, a terminal error, every backend refusing, a study that does not compile", async () => { + const cpu = createFakeBackend(WORKER_POOL_BACKEND_ID); + const { sampler } = createSampler({ cpu }); + + const errored = sampler.run(runRequest()); + await vi.waitFor(() => expect(cpu.handles).toHaveLength(1)); + cpu.handles[0]!.emit({ type: "complete", progress: progressOf(2, 1) }); + await expect(errored.completion).resolves.toEqual({ + ok: false, + cancelled: false, + reason: "1 of 3 runs failed", + }); + + const crashed = sampler.run(runRequest()); + await vi.waitFor(() => expect(cpu.handles).toHaveLength(2)); + cpu.handles[1]!.emit({ + type: "error", + message: "worker crashed", + itemId: null, + }); + await expect(crashed.completion).resolves.toEqual({ + ok: false, + cancelled: false, + reason: "worker crashed", + }); + + const refusing = createFakeBackend(WORKER_POOL_BACKEND_ID, { + refuse: "no", + }); + const refused = createSampler({ cpu: refusing }).sampler.run(runRequest()); + await expect(refused.completion).resolves.toEqual({ + ok: false, + cancelled: false, + reason: "cpu: no", + }); + expect(refusing.handles).toHaveLength(0); + + const uncompilable = sampler.run( + runRequest({ cacheKey: "missing-scenario", scenarioId: "missing" }), + ); + await expect(uncompilable.completion).resolves.toEqual({ + ok: false, + cancelled: false, + reason: "Scenario missing is not in the model snapshot", + }); + + const broken = sampler.run( + runRequest({ + cacheKey: "broken-metric", + definition: { + ...definition, + metrics: [{ ...metric, code: "return (" }], + }, + }), + ); + const outcome = await broken.completion; + expect(outcome).toMatchObject({ ok: false, cancelled: false }); + expect(outcome.ok ? "" : outcome.reason).toMatch( + new RegExp(`^${metric.id}: .+`), + ); + expect(cpu.handles).toHaveLength(2); + }); + + it("names why the scenario does not compile at a point", async () => { + const cpu = createFakeBackend(WORKER_POOL_BACKEND_ID); + const { sampler } = createSampler({ cpu }); + + const run = sampler.run( + runRequest({ + scenarioParameterValues: { + population: Number.NaN, + infected_ratio: 0.05, + }, + }), + ); + const outcome = await run.completion; + expect(outcome).toMatchObject({ ok: false, cancelled: false }); + expect(outcome.ok ? "" : outcome.reason).toMatch( + /^Scenario parameter "population" must be a finite number\./, + ); + expect(cpu.requests).toHaveLength(0); + }); + + it("names the kept backend when it refuses a later batch", async () => { + const cpu = createFakeBackend(WORKER_POOL_BACKEND_ID, { + refuse: "pool drained", + refuseFrom: 1, + }); + const { sampler } = createSampler({ cpu }); + + const first = sampler.run(runRequest()); + await vi.waitFor(() => expect(cpu.handles).toHaveLength(1)); + completeWith(cpu.handles[0]!, 0.1); + await expect(first.completion).resolves.toMatchObject({ ok: true }); + + const second = sampler.run(runRequest({ seed: 8, runSeeds: [8, 9, 10] })); + await expect(second.completion).resolves.toEqual({ + ok: false, + cancelled: false, + reason: "cpu: pool drained", + }); + expect(cpu.requests).toHaveLength(2); + expect(cpu.handles).toHaveLength(1); + }); + + it("passes no pinned seeds to the GPU and records where the batch ran", async () => { + const cpu = createFakeBackend(WORKER_POOL_BACKEND_ID); + const gpu = createFakeBackend("webgpu", { needsHirTrees: true }); + const { sampler } = createSampler({ cpu, gpu }); + + const run = sampler.run(runRequest({ computeBackend: "webgpu" })); + await vi.waitFor(() => expect(gpu.handles).toHaveLength(1)); + expect(gpu.requests[0]?.runs).toBeUndefined(); + expect(gpu.requests[0]?.seed).toBe(7); + expect(gpu.requests[0]?.hirArtifacts).toBeDefined(); + expect(cpu.requests).toHaveLength(0); + + completeWith(gpu.handles[0]!, 0.1); + await expect(run.completion).resolves.toMatchObject({ + computeBackend: "webgpu", + computeBackendFallbackReason: null, + }); + }); + + it("pins the seeds after all when a GPU request falls back to the CPU pool", async () => { + const cpu = createFakeBackend(WORKER_POOL_BACKEND_ID); + const gpu = createFakeBackend("webgpu", { refuse: "unsupported net" }); + const { sampler } = createSampler({ cpu, gpu }); + + const run = sampler.run(runRequest({ computeBackend: "webgpu" })); + await vi.waitFor(() => expect(cpu.handles).toHaveLength(2)); + // The walk's request carried no seeds (the GPU would have refused them); + // the handle it produced is replaced by one that pins them. + expect(cpu.requests[0]?.runs).toBeUndefined(); + expect(cpu.requests[1]?.runs).toEqual([ + { seed: 7 }, + { seed: 11 }, + { seed: 13 }, + ]); + expect(cpu.handles[0]!.handle.dispose).toHaveBeenCalled(); + expect(cpu.handles[0]!.handle.start).not.toHaveBeenCalled(); + + completeWith(cpu.handles[1]!, 0.3); + await expect(run.completion).resolves.toMatchObject({ + computeBackend: "cpu", + computeBackendFallbackReason: "unsupported net", + }); + }); + + it("walks the registrations once per study and backend, reusing the chosen backend", async () => { + const cpu = createFakeBackend(WORKER_POOL_BACKEND_ID); + const { sampler, registrations } = createSampler({ cpu }); + + const first = sampler.run(runRequest()); + await vi.waitFor(() => expect(cpu.handles).toHaveLength(1)); + completeWith(cpu.handles[0]!, 0.1); + await first.completion; + + const second = sampler.run(runRequest({ seed: 8, runSeeds: [8, 9, 10] })); + await vi.waitFor(() => expect(cpu.handles).toHaveLength(2)); + expect(cpu.requests[1]?.runs).toEqual([ + { seed: 8 }, + { seed: 9 }, + { seed: 10 }, + ]); + completeWith(cpu.handles[1]!, 0.1); + await second.completion; + expect(registrations).toHaveBeenCalledOnce(); + expect(registrations).toHaveBeenCalledWith( + expect.objectContaining({ computeBackend: "cpu", shardCount: 2 }), + ); + }); + + it("lets a cancelled batch settle and release the queue to the next batch of the same study", async () => { + const cpu = createFakeBackend(WORKER_POOL_BACKEND_ID); + const { sampler } = createSampler({ cpu }); + + const first = sampler.run(runRequest()); + await vi.waitFor(() => expect(cpu.handles).toHaveLength(1)); + const second = sampler.run(runRequest({ seed: 8, runSeeds: [8, 9, 10] })); + first.cancel(); + + await expect(first.completion).resolves.toEqual({ + ok: false, + cancelled: true, + reason: "cancelled", + }); + await vi.waitFor(() => + expect(cpu.handles[1]?.handle.start).toHaveBeenCalledOnce(), + ); + completeWith(cpu.handles[1]!, 0.4); + await expect(second.completion).resolves.toMatchObject({ + ok: true, + metricFrames: [frameOf(0.4)], + }); + }); + + it("runs batches with distinct queue keys side by side while compiling their study once", async () => { + const cpu = createFakeBackend(WORKER_POOL_BACKEND_ID); + const { sampler, registrations } = createSampler({ cpu }); + + const firstTrial = sampler.run(runRequest({ queueKey: "study:trial:0" })); + const secondTrial = sampler.run( + runRequest({ queueKey: "study:trial:1", seed: 8, runSeeds: [8, 9, 10] }), + ); + await vi.waitFor(() => { + expect(cpu.handles).toHaveLength(2); + for (const { handle } of cpu.handles) { + expect(handle.start).toHaveBeenCalledOnce(); + } + }); + + completeWith(cpu.handles[1]!, 0.2); + completeWith(cpu.handles[0]!, 0.1); + await expect(secondTrial.completion).resolves.toMatchObject({ + metricFrames: [frameOf(0.2)], + }); + await expect(firstTrial.completion).resolves.toMatchObject({ + metricFrames: [frameOf(0.1)], + }); + expect(registrations).toHaveBeenCalledOnce(); + }); + + it("queues one study's runs in order and runs studies side by side", async () => { + const cpu = createFakeBackend(WORKER_POOL_BACKEND_ID); + const { sampler } = createSampler({ cpu }); + + const firstOfA = sampler.run(runRequest({ cacheKey: "a" })); + const secondOfA = sampler.run(runRequest({ cacheKey: "a" })); + const onlyOfB = sampler.run(runRequest({ cacheKey: "b" })); + await vi.waitFor(() => expect(cpu.handles).toHaveLength(2)); + await new Promise((resolve) => { + setTimeout(resolve, 0); + }); + expect(cpu.handles).toHaveLength(2); + + completeWith(cpu.handles[0]!, 0.1); + completeWith(cpu.handles[1]!, 0.2); + await Promise.all([firstOfA.completion, onlyOfB.completion]); + await vi.waitFor(() => expect(cpu.handles).toHaveLength(3)); + completeWith(cpu.handles[2]!, 0.3); + await expect(secondOfA.completion).resolves.toMatchObject({ + metricFrames: [frameOf(0.3)], + }); + }); + + it("settles as cancelled on cancel, whether the batch is running or still queued", async () => { + const cpu = createFakeBackend(WORKER_POOL_BACKEND_ID); + const { sampler } = createSampler({ cpu }); + + const running = sampler.run(runRequest()); + const queued = sampler.run(runRequest()); + await vi.waitFor(() => expect(cpu.handles).toHaveLength(1)); + queued.cancel(); + running.cancel(); + expect(cpu.handles[0]!.handle.cancel).toHaveBeenCalledOnce(); + const cancelled = { ok: false, cancelled: true, reason: "cancelled" }; + await expect(running.completion).resolves.toEqual(cancelled); + await expect(queued.completion).resolves.toEqual(cancelled); + expect(cpu.handles).toHaveLength(1); + }); + + it("cancels through the request's signal and releases chosen backends on dispose", async () => { + const cpu = createFakeBackend(WORKER_POOL_BACKEND_ID); + const gpu = createFakeBackend("webgpu", { needsHirTrees: true }); + const { sampler } = createSampler({ cpu, gpu }); + const controller = new AbortController(); + + const run = sampler.run( + runRequest({ computeBackend: "webgpu", signal: controller.signal }), + ); + await vi.waitFor(() => expect(gpu.handles).toHaveLength(1)); + controller.abort(); + await expect(run.completion).resolves.toEqual({ + ok: false, + cancelled: true, + reason: "cancelled", + }); + + sampler.dispose(); + expect(gpu.backend.dispose).toHaveBeenCalledOnce(); + }); +}); diff --git a/libs/@hashintel/petrinaut/src/react/experiments/provider/detached-objective.ts b/libs/@hashintel/petrinaut/src/react/experiments/provider/detached-objective.ts index fccfe087f40..dc558f3fe20 100644 --- a/libs/@hashintel/petrinaut/src/react/experiments/provider/detached-objective.ts +++ b/libs/@hashintel/petrinaut/src/react/experiments/provider/detached-objective.ts @@ -3,17 +3,35 @@ import { DEFAULT_PETRINAUT_EXTENSIONS, getOwn, runExperimentToCompletion, + type MonteCarloExperiment, + type MonteCarloUserDefinedMetricFrame, + type MonteCarloWorkerProgress, type Scenario, } from "@hashintel/petrinaut-core"; -import { createWorkerPoolExperimentBackend } from "@hashintel/petrinaut-core/experiments"; +import { + createWorkerPoolExperimentBackend, + selectExperimentBackend, + WORKER_POOL_BACKEND_ID, +} from "@hashintel/petrinaut-core/experiments"; +import { createThrottle } from "../shared/throttle"; +import { experimentBackendRegistrations } from "./create-experiment"; +import { createWritableStore } from "./detached-objective/writable-store"; import { instantiateOnBackend } from "./shared/instantiate-on-backend"; import type { LanguageClientContextValue } from "../../lsp/context"; -import type { DetachedObjectiveRequest } from "../context"; +import type { + DetachedObjectiveRequest, + DetachedObjectiveRun, + DetachedObjectiveRunOutcome, + DetachedObjectiveRunRequest, + ExperimentComputeBackend, +} from "../context"; import type { SweepCellSnapshot } from "../sweep-session"; +import type { WritableStore } from "./detached-objective/writable-store"; import type { ExperimentBackend, + ExperimentRequest, ReusableWorkerFactory, } from "@hashintel/petrinaut-core/experiments"; @@ -31,6 +49,13 @@ type CompiledStudy = { metricArtifact: NonNullable; }; +/** The backend a study's runs settled on for one requested backend. */ +type ChosenBackend = { + backend: ExperimentBackend; + backendId: ExperimentComputeBackend; + fallbackReason: string | null; +}; + export type DetachedObjectiveSampler = { /** * Computes one objective sample against a study's frozen model snapshot. @@ -41,16 +66,47 @@ export type DetachedObjectiveSampler = { sample: ( request: DetachedObjectiveRequest, ) => Promise; + /** + * Streams one batch on the requested backend. The first run of a study on + * a backend walks the registrations and keeps the winner for the study's + * later runs. Runs queue per `queueKey` (the `cacheKey` when unset); + * studies run side by side. A batch that cannot run settles with the + * reason: the compile diagnostics, each backend's refusal, the terminal + * error, or the count of errored runs. + */ + run: (request: DetachedObjectiveRunRequest) => DetachedObjectiveRun; + /** Cancels every run in flight and releases the backends runs chose. */ + dispose: () => void; +}; + +/** How often a run republishes its frames and progress while streaming. */ +const RUN_PUBLISH_WINDOW_MS = 100; + +const cancelledOutcome: DetachedObjectiveRunOutcome = { + ok: false, + cancelled: true, + reason: "cancelled", }; +const failedOutcome = (reason: string): DetachedObjectiveRunOutcome => ({ + ok: false, + cancelled: false, + reason, +}); + +const errorMessage = (error: unknown): string => + error instanceof Error ? error.message : String(error); + /** * The frozen definition, its scenario HIR and its HIR artifacts never change - * for a given `cacheKey`, so they compile once per study. A failed compile - * is retried on the next sample rather than cached. + * for a given `cacheKey`, so they compile once per study and per artifact + * shape (with or without the HIR trees the GPU backend reads). A failed + * compile is retried on the next batch rather than cached. */ const compileStudy = async ( languageClient: LanguageClient, request: DetachedObjectiveRequest, + includeHir: boolean, ): Promise => { const scenario = (request.definition.scenarios ?? []).find( (candidate: Scenario) => candidate.id === request.scenarioId, @@ -65,14 +121,17 @@ const compileStudy = async ( const { artifacts, failures } = await languageClient.requestHirArtifacts( request.definition, DEFAULT_PETRINAUT_EXTENSIONS, - { includeHir: false }, + { includeHir }, ); const metricArtifact = getOwn(artifacts.metrics, request.metric.id); if (!metricArtifact) { throw new Error( failures - .map((failure) => failure.diagnostics[0]?.message) - .filter(Boolean) + .flatMap((failure) => + failure.diagnostics.map( + (diagnostic) => `${failure.itemId}: ${diagnostic.message}`, + ), + ) .join("; ") || "The objective metric did not compile", ); } @@ -80,89 +139,131 @@ const compileStudy = async ( return { scenario, scenarioHir, artifacts, metricArtifact }; }; +/** + * Scenario compilation is numeric; boolean bindings arrive as their 0/1 + * encoding, matching how the engine stores them. + */ +const numericScenarioValues = ( + values: DetachedObjectiveRequest["scenarioParameterValues"], +): Record => + Object.fromEntries( + Object.entries(values).map(([identifier, value]) => [ + identifier, + typeof value === "boolean" ? (value ? 1 : 0) : value, + ]), + ); + export const createDetachedObjectiveSampler = ({ languageClient, createWorker, + shardCount, + backendRegistrations = experimentBackendRegistrations, }: { /** Read per call, so a replaced language client is picked up. */ languageClient: { readonly current: LanguageClient }; createWorker: ReusableWorkerFactory; + /** The full pool's width; runs take a third of it. */ + shardCount: number; + backendRegistrations?: typeof experimentBackendRegistrations; }): DetachedObjectiveSampler => { const compileCache = new Map>(); - let backend: ExperimentBackend | null = null; - let chain: Promise = Promise.resolve(); + const chosenBackends = new Map(); + /** Walks in progress, so runs that overlap wait for one choice. */ + const pendingChoices = new Map>(); + const runQueues = new Map>(); + const runsInFlight = new Set(); + let sampleBackend: ExperimentBackend | null = null; + let sampleChain: Promise = Promise.resolve(); + // The wide CPU lane of a sweep: a third of the pool, so a study's runs + // leave room for the surface walk and the user's own experiments. + const runShards = Math.max(1, Math.floor(shardCount / 3)); const compiledFor = ( request: DetachedObjectiveRequest, + includeHir: boolean, ): Promise => { - let compiled = compileCache.get(request.cacheKey); + const key = `${request.cacheKey}|${includeHir ? "hir" : "flat"}`; + let compiled = compileCache.get(key); if (!compiled) { - compiled = compileStudy(languageClient.current, request); - compileCache.set(request.cacheKey, compiled); + compiled = compileStudy(languageClient.current, request, includeHir); + compileCache.set(key, compiled); compiled.catch(() => { - compileCache.delete(request.cacheKey); + compileCache.delete(key); }); } return compiled; }; - const runBatch = async ( + /** + * The request for one batch: the compiled snapshot with its scenario + * compiled at the batch's parameter point. Throws when the scenario does + * not compile there. + */ + const buildRequest = async ( request: DetachedObjectiveRequest, - ): Promise => { - try { - const { scenario, scenarioHir, artifacts, metricArtifact } = - await compiledFor(request); - const compiledScenario = compileScenario( - scenario, - scenarioHir, - request.definition.parameters, - request.definition.places, - request.definition.types, + options: { includeHir: boolean; runSeeds?: readonly number[] }, + ): Promise => { + const { scenario, scenarioHir, artifacts, metricArtifact } = + await compiledFor(request, options.includeHir); + const compiledScenario = compileScenario( + scenario, + scenarioHir, + request.definition.parameters, + request.definition.places, + request.definition.types, + { + scenarioParameterValues: numericScenarioValues( + request.scenarioParameterValues, + ), + }, + ); + if (!compiledScenario.ok) { + throw new Error( + compiledScenario.errors.map((error) => error.message).join("; ") || + `Scenario "${scenario.name}" did not compile at this point`, + ); + } + return { + sdcpn: request.definition, + extensions: DEFAULT_PETRINAUT_EXTENSIONS, + initialMarking: compiledScenario.result.initialState, + parameterValues: compiledScenario.result.parameterValues, + seed: request.seed, + dt: request.dt, + maxTime: request.maxTime, + runCount: request.runCount, + metricSpecs: [ { - // Scenario compilation is numeric; boolean bindings arrive as - // their 0/1 encoding, matching how the engine stores them. - scenarioParameterValues: Object.fromEntries( - Object.entries(request.scenarioParameterValues).map( - ([identifier, value]) => [ - identifier, - typeof value === "boolean" ? (value ? 1 : 0) : value, - ], - ), - ), + kind: "expression", + id: request.metric.id, + label: request.metric.label, + code: request.metric.code, + sampleRuns: "all", + runOutput: { type: "distribution" }, + artifact: metricArtifact, }, - ); - if (!compiledScenario.ok) { - return null; - } + ], + hirArtifacts: artifacts, + ...(options.runSeeds === undefined + ? {} + : { runs: options.runSeeds.map((seed) => ({ seed })) }), + }; + }; - backend ??= createWorkerPoolExperimentBackend({ + const sampleBatch = async ( + request: DetachedObjectiveRequest, + ): Promise => { + try { + const experimentRequest = await buildRequest(request, { + includeHir: false, + }); + sampleBackend ??= createWorkerPoolExperimentBackend({ createWorker, shardCount: 1, }); const handle = await instantiateOnBackend( - backend, - { - sdcpn: request.definition, - extensions: DEFAULT_PETRINAUT_EXTENSIONS, - initialMarking: compiledScenario.result.initialState, - parameterValues: compiledScenario.result.parameterValues, - seed: request.seed, - dt: request.dt, - maxTime: request.maxTime, - runCount: request.runCount, - metricSpecs: [ - { - kind: "expression", - id: request.metric.id, - label: request.metric.label, - code: request.metric.code, - sampleRuns: "all", - runOutput: { type: "distribution" }, - artifact: metricArtifact, - }, - ], - hirArtifacts: artifacts, - }, + sampleBackend, + experimentRequest, {}, ); const { event, frames } = await runExperimentToCompletion(handle); @@ -175,11 +276,281 @@ export const createDetachedObjectiveSampler = ({ } }; + /** A run's handle on the backend its study settled on. */ + const instantiateOnChosen = async ( + request: DetachedObjectiveRunRequest, + chosen: ChosenBackend, + signal: AbortSignal, + ): Promise => { + const experimentRequest = await buildRequest(request, { + includeHir: chosen.backend.needsHirTrees, + runSeeds: + chosen.backendId === WORKER_POOL_BACKEND_ID + ? request.runSeeds + : undefined, + }); + try { + return await instantiateOnBackend(chosen.backend, experimentRequest, { + signal, + }); + } catch (error) { + // A refusal reads as the walk's declines do: the backend, then why. + throw new Error(`${chosen.backendId}: ${errorMessage(error)}`); + } + }; + + /** Walks the registrations for a study's first run on a requested backend. */ + const walkBackends = async ( + request: DetachedObjectiveRunRequest, + signal: AbortSignal, + ): Promise<{ + handle: MonteCarloExperiment; + chosen: ChosenBackend; + }> => { + // The walk reports a request it cannot build as the first candidate's + // refusal; building it here first keeps a compile failure's diagnostics + // as the reason. The compile is cached for the candidate that needs it. + await buildRequest(request, { + includeHir: request.computeBackend === "webgpu", + }); + // The GPU backend refuses pinned seeds, and a refusal on the walk would + // read as a fallback. The seeds ride along only when every candidate is + // the CPU pool. + const pinSeedsOnWalk = request.computeBackend === "cpu"; + const selection = await selectExperimentBackend({ + registrations: backendRegistrations({ + computeBackend: request.computeBackend, + createWorker, + shardCount: runShards, + }), + buildRequest: ({ needsHirTrees }) => + buildRequest(request, { + includeHir: needsHirTrees, + runSeeds: pinSeedsOnWalk ? request.runSeeds : undefined, + }), + instantiateOptions: { signal }, + }); + if (!selection.ok) { + throw new Error( + selection.declined + .map((entry) => `${entry.backendId}: ${entry.reason}`) + .join("; ") || "Every backend declined the batch", + ); + } + const won: ChosenBackend = { + backend: selection.backend, + backendId: selection.backendId as ExperimentComputeBackend, + fallbackReason: selection.declined[0]?.reason ?? null, + }; + if ( + pinSeedsOnWalk || + won.backendId !== WORKER_POOL_BACKEND_ID || + request.runSeeds === undefined + ) { + return { handle: selection.handle, chosen: won }; + } + // The walk fell back to the CPU pool without the seeds. The pool takes + // them, so its handle is replaced by one that pins them. + selection.handle.dispose(); + const handle = await instantiateOnBackend( + won.backend, + await buildRequest(request, { + includeHir: false, + runSeeds: request.runSeeds, + }), + { signal }, + ); + return { handle, chosen: won }; + }; + + /** + * The handle for one run. The first run of a study on a requested backend + * walks the registrations and keeps the winner; later runs instantiate on + * it directly, and runs that begin while the walk is out wait for its + * choice. Throws when the kept backend or every candidate refuses, naming + * each and why. + */ + const acquireHandle = async ( + request: DetachedObjectiveRunRequest, + signal: AbortSignal, + ): Promise<{ + handle: MonteCarloExperiment; + chosen: ChosenBackend; + }> => { + const key = `${request.cacheKey}|${request.computeBackend}`; + const chosen = chosenBackends.get(key); + if (chosen) { + return { + handle: await instantiateOnChosen(request, chosen, signal), + chosen, + }; + } + const pending = pendingChoices.get(key); + if (pending) { + // A walk that failed leaves this run to walk for itself, so its own + // refusal, or its own cancellation, is what it reports. + const settled = await pending.catch(() => null); + if (settled) { + return { + handle: await instantiateOnChosen(request, settled, signal), + chosen: settled, + }; + } + } + const walk = walkBackends(request, signal); + const choice = walk.then(({ chosen: won }) => won); + choice.catch(() => undefined); + pendingChoices.set(key, choice); + try { + const result = await walk; + chosenBackends.set(key, result.chosen); + return result; + } finally { + if (pendingChoices.get(key) === choice) { + pendingChoices.delete(key); + } + } + }; + + const streamRun = async ( + request: DetachedObjectiveRunRequest, + signal: AbortSignal, + frames: WritableStore, + progress: WritableStore, + ): Promise => { + let handle: MonteCarloExperiment | null = null; + const cancelHandle = () => handle?.cancel(); + signal.addEventListener("abort", cancelHandle, { once: true }); + // The backend keeps listening to the signal it was instantiated with and + // tears the handle down when it fires, before the shards can answer the + // cancel with the terminal event this run waits for. Instantiation gets a + // signal of its own that dies once the handle is ready. + const instantiation = new AbortController(); + const abortInstantiation = () => instantiation.abort(); + signal.addEventListener("abort", abortInstantiation, { once: true }); + // Read through a call so the abort flag is re-checked after the await (a + // plain property read would be control-flow-narrowed to `false`). + const isCancelled = () => signal.aborted; + try { + if (isCancelled()) { + return cancelledOutcome; + } + let acquired: Awaited>; + try { + acquired = await acquireHandle(request, instantiation.signal); + } finally { + signal.removeEventListener("abort", abortInstantiation); + } + if (isCancelled()) { + acquired.handle.dispose(); + return cancelledOutcome; + } + handle = acquired.handle; + const live = acquired.handle; + const publish = createThrottle(() => { + frames.set(live.metrics.get().frames); + progress.set(live.progress.get()); + }, RUN_PUBLISH_WINDOW_MS); + const offMetrics = live.metrics.subscribe(publish.call); + const offProgress = live.progress.subscribe(publish.call); + let completion: Awaited>; + try { + completion = await runExperimentToCompletion(live); + } finally { + offMetrics(); + offProgress(); + publish.cancel(); + } + const { event, frames: finalFrames, runResults } = completion; + frames.set(finalFrames); + if (event.type === "error") { + return failedOutcome(event.message); + } + if (event.progress !== null) { + progress.set(event.progress); + } + if (event.type === "cancelled") { + return cancelledOutcome; + } + const { erroredRuns, runCount } = event.progress; + if (erroredRuns > 0) { + return failedOutcome(`${erroredRuns} of ${runCount} runs failed`); + } + return { + ok: true, + runsCompleted: event.progress.completedRuns, + metricFrames: finalFrames, + runResults, + computeBackend: acquired.chosen.backendId, + computeBackendFallbackReason: acquired.chosen.fallbackReason, + }; + } catch (error) { + return isCancelled() + ? cancelledOutcome + : failedOutcome(errorMessage(error)); + } finally { + signal.removeEventListener("abort", cancelHandle); + } + }; + + const run: DetachedObjectiveSampler["run"] = (request) => { + const frames = createWritableStore< + readonly MonteCarloUserDefinedMetricFrame[] + >([]); + const progress = createWritableStore(null); + const controller = new AbortController(); + const forwardAbort = () => controller.abort(); + if (request.signal?.aborted) { + controller.abort(); + } else { + request.signal?.addEventListener("abort", forwardAbort, { once: true }); + } + runsInFlight.add(controller); + + const queueKey = request.queueKey ?? request.cacheKey; + const previous = runQueues.get(queueKey) ?? Promise.resolve(); + const completion = previous.then(() => + streamRun(request, controller.signal, frames, progress), + ); + const settled = completion.then( + () => undefined, + () => undefined, + ); + runQueues.set(queueKey, settled); + void settled.then(() => { + runsInFlight.delete(controller); + request.signal?.removeEventListener("abort", forwardAbort); + if (runQueues.get(queueKey) === settled) { + runQueues.delete(queueKey); + } + }); + + return { + frames, + progress, + completion, + cancel: () => controller.abort(), + }; + }; + return { sample: (request) => { - const next = chain.then(() => runBatch(request)); - chain = next.catch(() => null); + const next = sampleChain.then(() => sampleBatch(request)); + sampleChain = next.catch(() => null); return next; }, + run, + dispose: () => { + for (const controller of runsInFlight) { + controller.abort(); + } + runsInFlight.clear(); + for (const chosen of chosenBackends.values()) { + chosen.backend.dispose?.(); + } + chosenBackends.clear(); + sampleBackend?.dispose?.(); + sampleBackend = null; + }, }; }; diff --git a/libs/@hashintel/petrinaut/src/react/experiments/provider/detached-objective/writable-store.ts b/libs/@hashintel/petrinaut/src/react/experiments/provider/detached-objective/writable-store.ts new file mode 100644 index 00000000000..b9472cf5a97 --- /dev/null +++ b/libs/@hashintel/petrinaut/src/react/experiments/provider/detached-objective/writable-store.ts @@ -0,0 +1,29 @@ +import type { ReadableStore } from "@hashintel/petrinaut-core"; + +export type WritableStore = ReadableStore & { + set(this: void, value: T): void; +}; + +/** A readable store with a setter. Setting an identical value notifies nobody. */ +export const createWritableStore = (initial: T): WritableStore => { + let current = initial; + const listeners = new Set<(value: T) => void>(); + return { + get: () => current, + set: (value) => { + if (Object.is(current, value)) { + return; + } + current = value; + for (const listener of listeners) { + listener(value); + } + }, + subscribe: (listener) => { + listeners.add(listener); + return () => { + listeners.delete(listener); + }; + }, + }; +}; diff --git a/libs/@hashintel/petrinaut/src/react/experiments/sweep-session/throttle.ts b/libs/@hashintel/petrinaut/src/react/experiments/shared/throttle.ts similarity index 100% rename from libs/@hashintel/petrinaut/src/react/experiments/sweep-session/throttle.ts rename to libs/@hashintel/petrinaut/src/react/experiments/shared/throttle.ts diff --git a/libs/@hashintel/petrinaut/src/react/experiments/sweep-session.ts b/libs/@hashintel/petrinaut/src/react/experiments/sweep-session.ts index a1c999a9ccc..f55cdbf9c2d 100644 --- a/libs/@hashintel/petrinaut/src/react/experiments/sweep-session.ts +++ b/libs/@hashintel/petrinaut/src/react/experiments/sweep-session.ts @@ -32,6 +32,7 @@ import { mergeMetricFramesAcrossCells, normalizeSweepSelection, } from "./parameter-grid"; +import { createThrottle } from "./shared/throttle"; import { sweepCellObjective } from "./sweep-cell-objective"; import { createBatchRegistry } from "./sweep-session/batch-registry"; import { @@ -46,7 +47,6 @@ import { sweepRangeDraws, sweepSelectionKey, } from "./sweep-session/selection-draws"; -import { createThrottle } from "./sweep-session/throttle"; import type { ExperimentParameterAxis, SweepSelection } from "./parameter-grid"; import type { SweepBatchStatus } from "./sweep-session/batch-registry"; diff --git a/libs/@hashintel/petrinaut/src/react/experiments/sweep-session/README.md b/libs/@hashintel/petrinaut/src/react/experiments/sweep-session/README.md index d04ec65ae0d..313d7fd82b1 100644 --- a/libs/@hashintel/petrinaut/src/react/experiments/sweep-session/README.md +++ b/libs/@hashintel/petrinaut/src/react/experiments/sweep-session/README.md @@ -3,4 +3,4 @@ layer: react.experiments.sweep role: The sweep session's private pieces (selection keys and range draws, the batch registry, cell batching, the publish throttle) --- -`sweep-session.ts` in the parent folder is the orchestrator (the refine ladder with pipelined rungs, the per-selection cache, the streamed gate). These modules are its private pieces: `selection-draws.ts` names selections and draws per-run values for a range, `batch-registry.ts` tracks every computing batch for the activity list, `cell-batch.ts` turns a chunk of surface cells into one experiment and regroups per-run values into cell means, `throttle.ts` is the leading-edge, trailing-coalesce timer both the publish and the batch refresh use. +`sweep-session.ts` in the parent folder is the orchestrator (the refine ladder with pipelined rungs, the per-selection cache, the streamed gate). These modules are its private pieces: `selection-draws.ts` names selections and draws per-run values for a range, `batch-registry.ts` tracks every computing batch for the activity list, `cell-batch.ts` turns a chunk of surface cells into one experiment and regroups per-run values into cell means. The leading-edge, trailing-coalesce timer both the publish and the batch refresh use lives in `../shared/throttle.ts`, shared with the detached objective runs. diff --git a/libs/@hashintel/petrinaut/src/react/experiments/sweep-session/batch-registry.ts b/libs/@hashintel/petrinaut/src/react/experiments/sweep-session/batch-registry.ts index 92b0874e9a4..f476926dc0d 100644 --- a/libs/@hashintel/petrinaut/src/react/experiments/sweep-session/batch-registry.ts +++ b/libs/@hashintel/petrinaut/src/react/experiments/sweep-session/batch-registry.ts @@ -1,4 +1,4 @@ -import { createThrottle } from "./throttle"; +import { createThrottle } from "../shared/throttle"; import type { MonteCarloExperiment } from "@hashintel/petrinaut-core"; diff --git a/libs/@hashintel/petrinaut/src/react/index.ts b/libs/@hashintel/petrinaut/src/react/index.ts index 7a63fe159aa..a9624a5043a 100644 --- a/libs/@hashintel/petrinaut/src/react/index.ts +++ b/libs/@hashintel/petrinaut/src/react/index.ts @@ -69,18 +69,29 @@ export { type NetManagement, } from "./net-management-context"; export { PetrinautOptimizationContext } from "./optimization-context"; -export type { PetrinautOptimization } from "./optimization-context"; +export type { + PetrinautConnectedOptimization, + PetrinautConnectedOptimizationCapability, + PetrinautConnectedRunOptions, + PetrinautOptimization, + PetrinautOptimizationChannel, + PetrinautOptimizationSource, +} from "./optimization-context"; export { isOptimizationActive, OptimizationsContext, } from "./optimizations/context"; export type { + CreateOptimizationOptions, OptimizationBest, OptimizationConnectionState, + OptimizationNavigation, OptimizationRecord, + OptimizationSelectionStream, OptimizationStatus, OptimizationsContextValue, } from "./optimizations/context"; +export { useOptimizationSource } from "./optimizations/use-optimization-source"; export { ExperimentsActionsContext, ExperimentsContext, @@ -88,6 +99,12 @@ export { } from "./experiments/context"; export type { CreateExperimentInput, + DetachedObjectiveRequest, + DetachedObjectiveRun, + DetachedObjectiveRunOutcome, + DetachedObjectiveRunRequest, + DetachedObjectiveRunResult, + ExperimentComputeBackend, ExperimentRecord, ExperimentsActionsValue, ExperimentStatus, diff --git a/libs/@hashintel/petrinaut/src/react/optimization-context.ts b/libs/@hashintel/petrinaut/src/react/optimization-context.ts index 81af913dea4..dcedf14465f 100644 --- a/libs/@hashintel/petrinaut/src/react/optimization-context.ts +++ b/libs/@hashintel/petrinaut/src/react/optimization-context.ts @@ -1,13 +1,28 @@ import { createContext } from "react"; import type { PetrinautOptimization } from "@hashintel/petrinaut-core"; +import type { + PetrinautConnectedOptimization, + PetrinautConnectedOptimizationCapability, + PetrinautConnectedRunOptions, + PetrinautOptimizationChannel, + PetrinautOptimizationSource, +} from "@hashintel/petrinaut-core/optimization"; /** - * Optional host-provided optimization capability. + * Optional host-provided optimization source: a remote capability, or a + * connected optimizer that runs its trials through the host's own compute. * * A `null` value means that optimization is unavailable and its UI is hidden. */ export const PetrinautOptimizationContext = - createContext(null); + createContext(null); -export type { PetrinautOptimization }; +export type { + PetrinautConnectedOptimization, + PetrinautConnectedOptimizationCapability, + PetrinautConnectedRunOptions, + PetrinautOptimization, + PetrinautOptimizationChannel, + PetrinautOptimizationSource, +}; diff --git a/libs/@hashintel/petrinaut/src/react/optimizations/channel/create-optimization-channel.test.ts b/libs/@hashintel/petrinaut/src/react/optimizations/channel/create-optimization-channel.test.ts new file mode 100644 index 00000000000..cfe0a51c454 --- /dev/null +++ b/libs/@hashintel/petrinaut/src/react/optimizations/channel/create-optimization-channel.test.ts @@ -0,0 +1,207 @@ +import { describe, expect, it, vi } from "vitest"; + +import { + resolveTrialScenarioParameterValues, + type PetrinautOptimizationTrialRequest, +} from "@hashintel/petrinaut-core/optimization"; + +import { + completedRunResult, + createFakeDetachedObjectiveRuns, + distributionFrame, + failedRunOutcome, +} from "../fake-detached-objective-runs.fixtures"; +import { + sirOptimizationInput, + sirOptimizationMetric, +} from "../sir-optimization-input.fixtures"; +import { + createOptimizationChannel, + type OptimizationChannelStudy, +} from "./create-optimization-channel"; + +const metricId = sirOptimizationMetric.id; + +const trialRequest = ( + overrides: Partial = {}, +): PetrinautOptimizationTrialRequest => { + const suggestedValues = { infected_ratio: 0.05 }; + return { + runId: "run-1", + trial: 0, + manifest: sirOptimizationInput, + suggestedValues, + scenarioParameterValues: resolveTrialScenarioParameterValues( + sirOptimizationInput, + suggestedValues, + ), + seeds: [1, 2, 3], + signal: new AbortController().signal, + ...overrides, + }; +}; + +const setup = () => { + const fake = createFakeDetachedObjectiveRuns(); + const study: OptimizationChannelStudy = { + computeBackend: "webgpu", + trialStarted: vi.fn(), + trialSettled: vi.fn(), + }; + const channel = createOptimizationChannel({ + runDetachedObjective: fake.runDetachedObjective, + resolveStudy: (runId) => (runId === "run-1" ? study : null), + }); + return { fake, study, channel }; +}; + +describe("createOptimizationChannel", () => { + it("runs a trial on the study's backend with its seeds pinned, and reports the mean of the per-seed finals", async () => { + const { fake, study, channel } = setup(); + + const outcome = channel.evaluateTrial(trialRequest()); + expect(fake.runs[0]?.request).toMatchObject({ + cacheKey: "run-1", + scenarioId: sirOptimizationInput.scenario.id, + scenarioParameterValues: { population: 1_000, infected_ratio: 0.05 }, + metric: { id: metricId, label: sirOptimizationMetric.name }, + seed: 1, + runCount: 3, + runSeeds: [1, 2, 3], + dt: 1, + maxTime: 180, + computeBackend: "webgpu", + }); + expect(study.trialStarted).toHaveBeenCalledWith( + 0, + { infected_ratio: 0.05 }, + fake.runs[0]!.run, + 3, + ); + expect(fake.runs[0]?.request.queueKey).toBe("run-1:trial:0"); + + const result = completedRunResult({ + metricId, + frames: [distributionFrame(metricId, 180, [[0.25, 3]])], + runValues: [0.5, 0.25, 0], + }); + fake.runs[0]!.settle(result); + await expect(outcome).resolves.toEqual({ + kind: "objective", + objective: 0.25, + replicates: [ + { seed: 1, objective: 0.5 }, + { seed: 2, objective: 0.25 }, + { seed: 3, objective: 0 }, + ], + }); + expect(study.trialSettled).toHaveBeenCalledWith(0, result); + }); + + it("reads the objective off the last sampled frame when the backend reports no run axis", async () => { + const { fake, channel } = setup(); + + const outcome = channel.evaluateTrial(trialRequest()); + fake.runs[0]!.settle( + completedRunResult({ + metricId, + frames: [ + distributionFrame(metricId, 1, [[0.9, 3]]), + distributionFrame(metricId, 180, [ + [0.1, 1], + [0.3, 1], + ]), + ], + runsCompleted: 3, + computeBackend: "webgpu", + }), + ); + await expect(outcome).resolves.toEqual({ + kind: "objective", + objective: 0.2, + }); + }); + + it("prunes a batch that did not complete with the batch's own reason, cancellation included", async () => { + const { fake, channel } = setup(); + + const failed = channel.evaluateTrial(trialRequest()); + fake.runs[0]!.settle(failedRunOutcome("2 of 3 runs failed")); + await expect(failed).resolves.toEqual({ + kind: "pruned", + reason: "2 of 3 runs failed", + }); + + const controller = new AbortController(); + const cancelled = channel.evaluateTrial( + trialRequest({ trial: 1, signal: controller.signal }), + ); + controller.abort(); + await expect(cancelled).resolves.toEqual({ + kind: "pruned", + reason: "cancelled", + }); + expect(fake.runs[1]!.cancelled).toBe(true); + + const aborted = new AbortController(); + aborted.abort(); + await expect( + channel.evaluateTrial(trialRequest({ trial: 2, signal: aborted.signal })), + ).resolves.toEqual({ kind: "pruned", reason: "cancelled" }); + expect(fake.runs).toHaveLength(2); + }); + + it("prunes a trial whose objective is not finite", async () => { + const { fake, channel } = setup(); + + const outcome = channel.evaluateTrial(trialRequest()); + fake.runs[0]!.settle( + completedRunResult({ metricId, frames: [], runsCompleted: 3 }), + ); + await expect(outcome).resolves.toEqual({ + kind: "pruned", + reason: `The objective metric "${metricId}" did not produce a finite value`, + }); + }); + + it("evaluates a run the provider does not know on the CPU, unwatched", async () => { + const { fake, study, channel } = setup(); + + const outcome = channel.evaluateTrial(trialRequest({ runId: "unknown" })); + expect(fake.runs[0]?.request.computeBackend).toBe("cpu"); + fake.runs[0]!.settle( + completedRunResult({ + metricId, + frames: [distributionFrame(metricId, 180, [[0.3, 1]])], + runValues: [0.3], + }), + ); + await expect(outcome).resolves.toMatchObject({ + kind: "objective", + objective: 0.3, + }); + expect(study.trialStarted).not.toHaveBeenCalled(); + }); + + it("never throws: a failing run request becomes a pruned trial, and dispose cancels runs in flight", async () => { + const throwing = createOptimizationChannel({ + runDetachedObjective: () => { + throw new Error("no compute"); + }, + resolveStudy: () => null, + }); + await expect(throwing.evaluateTrial(trialRequest())).resolves.toEqual({ + kind: "pruned", + reason: "no compute", + }); + + const { fake, channel } = setup(); + const outcome = channel.evaluateTrial(trialRequest()); + channel.dispose(); + expect(fake.runs[0]!.cancelled).toBe(true); + await expect(outcome).resolves.toEqual({ + kind: "pruned", + reason: "cancelled", + }); + }); +}); diff --git a/libs/@hashintel/petrinaut/src/react/optimizations/channel/create-optimization-channel.ts b/libs/@hashintel/petrinaut/src/react/optimizations/channel/create-optimization-channel.ts new file mode 100644 index 00000000000..987128579a6 --- /dev/null +++ b/libs/@hashintel/petrinaut/src/react/optimizations/channel/create-optimization-channel.ts @@ -0,0 +1,138 @@ +/** + * @layerRoot react.optimizations.channel + * @role Evaluates optimizer trials as detached objective runs on the experiments backend + */ +import { + prunedTrialOutcome, + trialOutcome, +} from "./create-optimization-channel/trial-outcome"; + +import type { + DetachedObjectiveRun, + DetachedObjectiveRunOutcome, + ExperimentComputeBackend, + ExperimentsActionsValue, +} from "../../experiments/context"; +import type { + OptimizationScalar, + PetrinautOptimizationChannel, +} from "@hashintel/petrinaut-core/optimization"; + +/** + * The study a run belongs to, as the channel needs it: which backend to ask + * for, and who watches the trials as they evaluate. + */ +export type OptimizationChannelStudy = { + computeBackend: ExperimentComputeBackend; + trialStarted: ( + trial: number, + values: Readonly>, + run: DetachedObjectiveRun, + runCount: number, + ) => void; + trialSettled: (trial: number, outcome: DetachedObjectiveRunOutcome) => void; +}; + +export type OptimizationChannel = PetrinautOptimizationChannel & { + dispose(this: void): void; +}; + +const errorMessage = (error: unknown): string => + error instanceof Error ? error.message : String(error); + +/** + * The channel a connected optimizer evaluates its trials through. Each trial + * becomes one detached objective run compiled once per optimizer run id and + * queued on its own, so trials the optimizer keeps in flight together + * overlap. The channel never throws: whatever stops a trial reaches Optuna + * as a pruned trial carrying the reason. + */ +export const createOptimizationChannel = ({ + runDetachedObjective, + resolveStudy, +}: { + runDetachedObjective: ExperimentsActionsValue["runDetachedObjective"]; + /** + * The study behind a run id, or null for a run the provider does not + * know, whose trials run on the CPU with nobody watching. + */ + resolveStudy: (runId: string) => OptimizationChannelStudy | null; +}): OptimizationChannel => { + const runsInFlight = new Set(); + + const evaluateTrial: PetrinautOptimizationChannel["evaluateTrial"] = async ( + request, + ) => { + // Read through a call so the abort flag is re-checked after an await (a + // plain property read would be control-flow-narrowed to `false`). + const isCancelled = () => request.signal.aborted; + const metric = request.manifest.model.definition.metrics?.find( + (candidate) => candidate.id === request.manifest.objective.metricId, + ); + const [firstSeed] = request.seeds; + if (!metric) { + return prunedTrialOutcome( + `The study has no metric "${request.manifest.objective.metricId}" to optimize`, + ); + } + if (firstSeed === undefined) { + return prunedTrialOutcome("The trial has no seed to run with"); + } + if (isCancelled()) { + return prunedTrialOutcome("cancelled"); + } + + const controller = new AbortController(); + const forwardAbort = () => controller.abort(); + request.signal.addEventListener("abort", forwardAbort, { once: true }); + let run: DetachedObjectiveRun | null = null; + let outcome: DetachedObjectiveRunOutcome; + try { + const study = resolveStudy(request.runId); + run = runDetachedObjective({ + cacheKey: request.runId, + // Trials in flight at once each take a queue of their own; the + // compiled study is shared through the cache key. + queueKey: `${request.runId}:trial:${request.trial}`, + definition: request.manifest.model.definition, + scenarioId: request.manifest.scenario.id, + scenarioParameterValues: request.scenarioParameterValues, + metric: { id: metric.id, label: metric.name, code: metric.code }, + seed: firstSeed, + runCount: request.seeds.length, + runSeeds: request.seeds, + dt: request.manifest.execution.dt, + maxTime: request.manifest.execution.maxTime, + computeBackend: study?.computeBackend ?? "cpu", + signal: controller.signal, + }); + runsInFlight.add(run); + study?.trialStarted( + request.trial, + request.suggestedValues, + run, + request.seeds.length, + ); + outcome = await run.completion; + study?.trialSettled(request.trial, outcome); + } catch (error) { + outcome = { ok: false, cancelled: false, reason: errorMessage(error) }; + } finally { + if (run) { + runsInFlight.delete(run); + } + request.signal.removeEventListener("abort", forwardAbort); + } + return trialOutcome(outcome, metric.id, request.seeds); + }; + + return { + evaluateTrial, + dispose: () => { + for (const run of runsInFlight) { + run.cancel(); + } + runsInFlight.clear(); + }, + }; +}; diff --git a/libs/@hashintel/petrinaut/src/react/optimizations/channel/create-optimization-channel/trial-outcome.ts b/libs/@hashintel/petrinaut/src/react/optimizations/channel/create-optimization-channel/trial-outcome.ts new file mode 100644 index 00000000000..588dc72ac5a --- /dev/null +++ b/libs/@hashintel/petrinaut/src/react/optimizations/channel/create-optimization-channel/trial-outcome.ts @@ -0,0 +1,80 @@ +import { getOwn } from "@hashintel/petrinaut-core"; + +import { sweepCellObjective } from "../../../experiments/sweep-cell-objective"; + +import type { + DetachedObjectiveRunOutcome, + DetachedObjectiveRunResult, +} from "../../../experiments/context"; +import type { PetrinautOptimizationTrialOutcome } from "@hashintel/petrinaut-core/optimization"; + +export const prunedTrialOutcome = ( + reason: string, +): PetrinautOptimizationTrialOutcome => ({ kind: "pruned", reason }); + +type TrialReplicate = { seed: number; objective: number }; + +/** + * The per-seed objectives, read off the per-run finals the CPU backend + * reports. Run `i` ran `seeds[i]`, which is how the request pinned them. + * Undefined when the backend reports no run axis, or a run's value is + * missing or not finite. + */ +const trialReplicates = ( + result: DetachedObjectiveRunResult, + metricId: string, + seeds: readonly number[], +): TrialReplicate[] | undefined => { + if (result.runResults.size === 0) { + return undefined; + } + const replicates: TrialReplicate[] = []; + const byRunIndex = [...result.runResults].sort( + ([left], [right]) => left - right, + ); + for (const [runIndex, values] of byRunIndex) { + const seed = seeds[runIndex]; + const objective = getOwn(values, metricId); + if ( + seed === undefined || + objective === undefined || + !Number.isFinite(objective) + ) { + return undefined; + } + replicates.push({ seed, objective }); + } + return replicates; +}; + +/** + * A settled trial batch as Optuna receives it. A batch that did not complete + * prunes the trial with the batch's own reason. The objective is the mean of + * the per-seed objectives, as the optimizer service reports it; where the + * backend reports no run axis it is the metric's last sampled frame, which + * a distribution frame reduces to the mean of its bins. + */ +export const trialOutcome = ( + outcome: DetachedObjectiveRunOutcome, + metricId: string, + seeds: readonly number[], +): PetrinautOptimizationTrialOutcome => { + if (!outcome.ok) { + return prunedTrialOutcome(outcome.reason); + } + const replicates = trialReplicates(outcome, metricId, seeds); + const objective = replicates + ? replicates.reduce((sum, replicate) => sum + replicate.objective, 0) / + replicates.length + : sweepCellObjective(outcome.metricFrames, metricId); + if (objective === null || !Number.isFinite(objective)) { + return prunedTrialOutcome( + `The objective metric "${metricId}" did not produce a finite value`, + ); + } + return { + kind: "objective", + objective, + ...(replicates === undefined ? {} : { replicates }), + }; +}; diff --git a/libs/@hashintel/petrinaut/src/react/optimizations/context.ts b/libs/@hashintel/petrinaut/src/react/optimizations/context.ts index ee0474595a5..e8e66fa531a 100644 --- a/libs/@hashintel/petrinaut/src/react/optimizations/context.ts +++ b/libs/@hashintel/petrinaut/src/react/optimizations/context.ts @@ -1,10 +1,14 @@ import { createContext } from "react"; +import type { ExperimentComputeBackend } from "../experiments/context"; +import type { OptimizationSurfaceAxis } from "./surface-grid"; import type { + MonteCarloUserDefinedMetricFrame, PetrinautOptimizationEvent, PetrinautOptimizationInput, PetrinautOptimizationTrialEvent, } from "@hashintel/petrinaut-core"; +import type { OptimizationScalar } from "@hashintel/petrinaut-core/optimization"; export type OptimizationStatus = | "initializing" @@ -39,6 +43,67 @@ export type OptimizationBest = NonNullable< Extract["best"] >; +/** Where a connected study's drawer points: one parameter point. */ +export type OptimizationNavigation = { + /** Axis position (0..stepCount) per optimized numeric parameter identifier. */ + positions: Readonly>; + /** Value per optimized boolean parameter identifier. */ + booleans: Readonly>; + /** + * While true, the navigation follows each trial as it is evaluated. On at + * creation; cleared by a user move. + */ + followTrials: boolean; +}; + +/** The objective's live metric stream at the navigation, or at the followed trial. */ +export type OptimizationSelectionStream = { + /** + * `trial:` while following a trial; otherwise the navigation key + * (positions in axis order, then booleans). + */ + key: string; + metricFrames: readonly MonteCarloUserDefinedMetricFrame[]; + runsCompleted: number; + /** + * Ladder target the in-flight batch climbs to; null when saturated or + * while following a trial. + */ + runTarget: number | null; + computing: boolean; + /** + * Why the last batch at this key failed — the metric's compile + * diagnostics, the backend's refusal, the count of errored runs — so the + * drawer can say what to fix. Null while computing and once a batch has + * succeeded; a cancellation records nothing. + */ + error: string | null; + /** + * Why the ladder stopped short of its top rung — "8 runs · cannot beat the + * best" — or null while it climbs, once it reaches the top, or on a trial. + */ + note: string | null; +}; + +/** One batch a connected study is computing, for the drawer's activity list. */ +export type OptimizationBatchStatus = { + id: string; + /** A step's runs, or one rung of the navigated point's refinement ladder. */ + kind: "step" | "refine"; + /** "Step 4", or "Refining population 1850 · infected_ratio 0.36". */ + label: string; + runCount: number; + completedRuns: number; +}; + +/** A step the optimizer is evaluating, with its objective so far. */ +export type OptimizationInFlightStep = { + trial: number; + parameters: Readonly>; + /** The running objective, null before the first frame with samples. */ + objective: number | null; +}; + export type OptimizationRecord = { id: string; input: PetrinautOptimizationInput; @@ -65,24 +130,116 @@ export type OptimizationRecord = { failedTrials: number; trials: readonly PetrinautOptimizationTrialEvent[]; best: OptimizationBest | null; + /** + * Whether more steps can be run on the study: a connected study keeps its + * sampler's history until it is removed, so it is resumable once a segment + * ends — by completion, or by a stop once its steps in flight are pruned. + * False for a remote study, and for one that failed. + */ + resumable: boolean; + /** Steps a connected study keeps in flight at once; 1 for a remote study. */ + parallelism: number; + /** + * The backend the study's trials run on: the one asked for, until the + * first trial that ran elsewhere reports where. `cpu` for a remote study. + */ + computeBackend: ExperimentComputeBackend; + /** + * Why the requested backend declined, from the first trial that ran + * elsewhere; null while every trial ran where asked. + */ + computeBackendFallbackReason: string | null; + /** The study's navigable axes: its optimized numeric parameters. */ + axes: readonly OptimizationSurfaceAxis[]; + /** + * Where the drawer points; null for a remote study, which computes nothing + * locally. + */ + navigation: OptimizationNavigation | null; + /** + * The objective's live stream at the navigation or the followed trial; + * null for a remote study. + */ + selection: OptimizationSelectionStream | null; + /** + * Every batch a connected study computes right now — the steps in flight + * and the navigated point's refinement rung. Empty when idle, and always + * for a remote study. + */ + activity: readonly OptimizationBatchStatus[]; + /** + * The steps a connected study is evaluating, most recently started last, + * each with its running objective. Empty when none is, and always for a + * remote study. + */ + inFlight: readonly OptimizationInFlightStep[]; }; +const TRIAL_SELECTION_KEY_PREFIX = "trial:"; + +/** The trial a selection stream follows, or null when the stream is a point's. */ +export function followedTrial(selectionKey: string): number | null { + if (!selectionKey.startsWith(TRIAL_SELECTION_KEY_PREFIX)) { + return null; + } + const trial = Number(selectionKey.slice(TRIAL_SELECTION_KEY_PREFIX.length)); + return Number.isInteger(trial) ? trial : null; +} + export function isOptimizationActive( - optimization: OptimizationRecord, + optimization: Pick, ): boolean { return ( optimization.status === "initializing" || optimization.status === "running" ); } +export type CreateOptimizationOptions = { + /** + * Backend a connected study's trials and refinement try first; a remote + * study ignores it. Defaults to `cpu`. + */ + computeBackend?: ExperimentComputeBackend; + /** + * Steps a connected study keeps in flight at once, 1 to + * `PETRINAUT_OPTIMIZATION_MAX_PARALLELISM`; a remote study ignores it. + * Defaults to 1. + */ + parallelism?: number; +}; + export type OptimizationsContextValue = { optimizations: readonly OptimizationRecord[]; selectedOptimizationId: string | null; selectedOptimization: OptimizationRecord | null; setSelectedOptimizationId: (optimizationId: string | null) => void; - createOptimization: (input: PetrinautOptimizationInput) => Promise; + createOptimization: ( + input: PetrinautOptimizationInput, + options?: CreateOptimizationOptions, + ) => Promise; + /** + * Stops the study. A remote run is cancelled server-side; a connected + * study prunes the steps in flight and keeps its sampler's history, so it + * can be continued. + */ cancelOptimization: (optimizationId: string) => void; removeOptimization: (optimizationId: string) => void; + /** + * Runs `trials` more steps on a resumable connected study, following them + * as they are evaluated. Rejects for a study that is running, was removed, + * failed, or would exceed the trial cap; the record's `error` carries the + * reason as well. + */ + extendOptimization: (optimizationId: string, trials: number) => Promise; + /** + * Moves a connected study's navigation. A position or boolean change stops + * following trials, and the selection refines at the new point; a remote + * study has no navigation and ignores the call. + */ + setOptimizationNavigation: ( + optimizationId: string, + patch: Partial, + ) => void; /** * Start a fresh optimization from a prior one's input (e.g. after a * transport failure). Returns the new id, or null if the record is gone. @@ -99,6 +256,9 @@ const DEFAULT_CONTEXT_VALUE: OptimizationsContextValue = { Promise.reject(new Error("Optimization is unavailable")), cancelOptimization: () => {}, removeOptimization: () => {}, + extendOptimization: () => + Promise.reject(new Error("Optimization is unavailable")), + setOptimizationNavigation: () => {}, retryOptimization: () => Promise.resolve(null), }; diff --git a/libs/@hashintel/petrinaut/src/react/optimizations/fake-detached-objective-runs.fixtures.ts b/libs/@hashintel/petrinaut/src/react/optimizations/fake-detached-objective-runs.fixtures.ts new file mode 100644 index 00000000000..bc8e17869db --- /dev/null +++ b/libs/@hashintel/petrinaut/src/react/optimizations/fake-detached-objective-runs.fixtures.ts @@ -0,0 +1,134 @@ +import type { + DetachedObjectiveRun, + DetachedObjectiveRunOutcome, + DetachedObjectiveRunRequest, + ExperimentComputeBackend, +} from "../experiments/context"; +import type { + MonteCarloUserDefinedMetricFrame, + MonteCarloWorkerProgress, +} from "@hashintel/petrinaut-core"; + +type Store = { + get(): T; + set(value: T): void; + subscribe(listener: (value: T) => void): () => void; +}; + +const createStore = (initial: T): Store => { + let current = initial; + const listeners = new Set<(value: T) => void>(); + return { + get: () => current, + set: (value) => { + current = value; + for (const listener of listeners) { + listener(value); + } + }, + subscribe: (listener) => { + listeners.add(listener); + return () => { + listeners.delete(listener); + }; + }, + }; +}; + +export type FakeDetachedObjectiveRun = { + request: DetachedObjectiveRunRequest; + frames: Store; + progress: Store; + run: DetachedObjectiveRun; + cancelled: boolean; + settle: (outcome: DetachedObjectiveRunOutcome) => void; +}; + +export const cancelledRunOutcome: DetachedObjectiveRunOutcome = { + ok: false, + cancelled: true, + reason: "cancelled", +}; + +export const failedRunOutcome = ( + reason: string, +): DetachedObjectiveRunOutcome => ({ ok: false, cancelled: false, reason }); + +/** Records every requested run and lets the test stream into and settle each one. */ +export const createFakeDetachedObjectiveRuns = () => { + const runs: FakeDetachedObjectiveRun[] = []; + const runDetachedObjective = ( + request: DetachedObjectiveRunRequest, + ): DetachedObjectiveRun => { + const frames = createStore([]); + const progress = createStore(null); + const { promise, resolve } = + Promise.withResolvers(); + const entry: FakeDetachedObjectiveRun = { + request, + frames, + progress, + cancelled: false, + settle: resolve, + run: { + frames, + progress, + completion: promise, + cancel: () => { + entry.cancelled = true; + resolve(cancelledRunOutcome); + }, + }, + }; + request.signal?.addEventListener("abort", entry.run.cancel, { + once: true, + }); + runs.push(entry); + return entry.run; + }; + return { runs, runDetachedObjective }; +}; + +export const distributionFrame = ( + metricId: string, + frameNumber: number, + bins: readonly (readonly [number, number])[], +): MonteCarloUserDefinedMetricFrame => ({ + metricId, + label: metricId, + outputType: "distribution", + frameNumber, + time: frameNumber, + bins, + value: null, + frameValue: null, + timeValue: null, + runSampleCount: bins.reduce((sum, [, frequency]) => sum + frequency, 0), + timeSampleCount: 0, +}); + +/** A finished batch: `runValues` are the per-run finals the CPU pool reports; none for the GPU. */ +export const completedRunResult = ({ + metricId, + frames, + runValues = [], + runsCompleted = runValues.length, + computeBackend = "cpu", + fallbackReason = null, +}: { + metricId: string; + frames: readonly MonteCarloUserDefinedMetricFrame[]; + runValues?: readonly number[]; + runsCompleted?: number; + computeBackend?: ExperimentComputeBackend; + fallbackReason?: string | null; +}): Extract => ({ + ok: true, + runsCompleted, + metricFrames: frames, + runResults: new Map( + runValues.map((value, runIndex) => [runIndex, { [metricId]: value }]), + ), + computeBackend, + computeBackendFallbackReason: fallbackReason, +}); diff --git a/libs/@hashintel/petrinaut/src/react/optimizations/provider.test.tsx b/libs/@hashintel/petrinaut/src/react/optimizations/provider.test.tsx index d8bef5545e3..fdae2d134be 100644 --- a/libs/@hashintel/petrinaut/src/react/optimizations/provider.test.tsx +++ b/libs/@hashintel/petrinaut/src/react/optimizations/provider.test.tsx @@ -7,68 +7,57 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { PETRINAUT_OPTIMIZATION_CANCELLED_ERROR_CODE, - petrinautOptimizationInputSchema, type PetrinautOptimization, + type PetrinautOptimizationEvent, } from "@hashintel/petrinaut-core"; -import { sirModel } from "@hashintel/petrinaut-core/examples"; +import { + type PetrinautConnectedOptimization, + resolveTrialScenarioParameterValues, +} from "@hashintel/petrinaut-core/optimization"; +import { + ExperimentsActionsContext, + type ExperimentsActionsValue, +} from "../experiments/context"; import { PetrinautNavigationProvider, usePetrinautNavigation, } from "../navigation"; import { PetrinautOptimizationContext } from "../optimization-context"; +import { UserSettingsContext } from "../state/user-settings-context"; import { OptimizationsContext, type OptimizationsContextValue, } from "./context"; +import { + completedRunResult, + createFakeDetachedObjectiveRuns, + distributionFrame, +} from "./fake-detached-objective-runs.fixtures"; import { OptimizationsProvider } from "./provider"; +import { + sirOptimizationInput, + sirOptimizationMetric, +} from "./sir-optimization-input.fixtures"; +import { + buildOptimizationSurfaceAxes, + optimizationAxisPositionFor, + optimizationAxisValueAt, +} from "./surface-grid"; import type { PetrinautNavigationState } from "../navigation"; +import type { PropsWithChildren } from "react"; -const scenario = sirModel.petriNetDefinition.scenarios?.find( - (candidate) => candidate.id === "scenario__seasonal_flu", -); -const metric = sirModel.petriNetDefinition.metrics?.find( - (candidate) => candidate.id === "metric__infected_fraction", -); -if (!scenario || !metric) { - throw new Error("The SIR optimization fixtures are incomplete"); -} +const input = sirOptimizationInput; +const metricId = sirOptimizationMetric.id; +const infectedRatioAxis = buildOptimizationSurfaceAxes(input)[0]!; -const input = petrinautOptimizationInputSchema.parse({ - kind: "petrinaut-optimization", - version: 1, - name: "SIR optimization", - model: { - title: sirModel.title, - definition: { - ...sirModel.petriNetDefinition, - scenarios: [scenario], - metrics: [metric], - }, - }, - scenario: { - id: scenario.id, - parameterBindings: { - population: { kind: "fixed", value: 1_000 }, - infected_ratio: { - kind: "optimize", - domain: { - kind: "continuous", - minimum: 0.001, - maximum: 0.2, - scale: "log", - }, - }, - }, - }, - objective: { - metricId: "metric__infected_fraction", - direction: "minimize", - }, - execution: { seed: 1, dt: 1, maxTime: 180 }, - study: { trials: 2, sampler: "tpe" }, -}); +/** An event before a fake log stamps its `seq`, each variant on its own. */ +type UnsequencedEvent = PetrinautOptimizationEvent extends infer Event + ? Event extends unknown + ? Omit + : never + : never; const CaptureContext = ({ onValue, @@ -88,6 +77,173 @@ const CaptureNavigation = ({ return null; }; +/** Overrides the In-browser optimization setting below the default context. */ +const InBrowserOptimizationSetting = ({ + enabled, + children, +}: PropsWithChildren<{ enabled: boolean }>) => { + const value = use(UserSettingsContext); + return ( + + {children} + + ); +}; + +/** Routes the provider's detached objective runs to a fake. */ +const ExperimentsActionsOverride = ({ + runDetachedObjective, + children, +}: PropsWithChildren<{ + runDetachedObjective: ExperimentsActionsValue["runDetachedObjective"]; +}>) => { + const value = use(ExperimentsActionsContext); + return ( + + {children} + + ); +}; + +/** + * A connected source whose runs stay quiet until aborted, counting connections + * and disposals so tests can observe what the setting gates. + */ +const createQuietConnectedSource = () => { + const calls = { connect: 0, dispose: 0 }; + const source: PetrinautConnectedOptimization = { + kind: "connected", + connect: () => { + calls.connect += 1; + return { + createOptimizationRun: () => + Promise.resolve({ runId: "run-quiet-connected" }), + // eslint-disable-next-line require-yield -- the run stays quiet until aborted + async *attachOptimizationRun(_runId, options) { + options?.onAttached?.(); + await new Promise((resolve) => { + options?.signal?.addEventListener("abort", resolve, { + once: true, + }); + }); + }, + cancelOptimizationRun: () => Promise.resolve(), + extendOptimizationRun: () => Promise.resolve(), + releaseOptimizationRun: () => Promise.resolve(), + dispose: () => { + calls.dispose += 1; + }, + }; + }, + }; + return { source, calls }; +}; + +/** + * A connected source whose study evaluates one trial per value through the + * channel, in order, then completes — the shape of the in-browser optimizer. + */ +const createEvaluatingSource = (infectedRatios: readonly number[]) => { + const calls = { connect: 0, dispose: 0 }; + const source: PetrinautConnectedOptimization = { + kind: "connected", + connect: (channel) => { + calls.connect += 1; + return { + createOptimizationRun: () => + Promise.resolve({ runId: "run-connected" }), + async *attachOptimizationRun(runId, options) { + options?.onAttached?.(); + let seq = 0; + for (const [trial, infectedRatio] of infectedRatios.entries()) { + const suggestedValues = { infected_ratio: infectedRatio }; + const outcome = await channel.evaluateTrial({ + runId, + trial, + manifest: input, + suggestedValues, + scenarioParameterValues: resolveTrialScenarioParameterValues( + input, + suggestedValues, + ), + seeds: [1, 2, 3], + signal: options?.signal ?? new AbortController().signal, + }); + seq += 1; + yield { + type: "trial", + trial, + parameters: suggestedValues, + objective: + outcome.kind === "objective" ? outcome.objective : null, + state: outcome.kind === "objective" ? "complete" : "pruned", + best: null, + seq, + }; + } + seq += 1; + yield { + type: "complete", + requestedTrials: infectedRatios.length, + completedTrials: infectedRatios.length, + prunedTrials: 0, + failedTrials: 0, + best: null, + seq, + }; + }, + cancelOptimizationRun: () => Promise.resolve(), + extendOptimizationRun: () => Promise.resolve(), + releaseOptimizationRun: () => Promise.resolve(), + dispose: () => { + calls.dispose += 1; + }, + }; + }, + }; + return { source, calls }; +}; + +const renderConnectedProvider = ({ + source, + runDetachedObjective, + enabled = true, +}: { + source: PetrinautConnectedOptimization; + runDetachedObjective: ExperimentsActionsValue["runDetachedObjective"]; + enabled?: boolean; +}) => { + let latest: OptimizationsContextValue | null = null; + const tree = (isEnabled: boolean) => ( + + + + + { + latest = value; + }} + /> + + + + + ); + const { rerender, unmount } = render(tree(enabled)); + return { + getValue: () => { + if (!latest) { + throw new Error("Optimization context was not captured"); + } + return latest; + }, + setEnabled: (isEnabled: boolean) => rerender(tree(isEnabled)), + unmount, + }; +}; + function renderProvider(capability: PetrinautOptimization) { let latest: OptimizationsContextValue | null = null; render( @@ -952,4 +1108,754 @@ describe("OptimizationsProvider", () => { expect(optimization.status).toBe("running"); expect(optimization.error).toBeNull(); }); + + it("treats a connected source as absent while In-browser optimization is off", async () => { + const { source, calls } = createQuietConnectedSource(); + const fake = createFakeDetachedObjectiveRuns(); + const { getValue } = renderConnectedProvider({ + source, + runDetachedObjective: fake.runDetachedObjective, + enabled: false, + }); + + await expect(getValue().createOptimization(input)).rejects.toThrow( + "Optimization is unavailable", + ); + expect(calls.connect).toBe(0); + expect(getValue().optimizations).toHaveLength(0); + }); + + it("connects and disposes a connected source as In-browser optimization is toggled", async () => { + const { source, calls } = createQuietConnectedSource(); + const fake = createFakeDetachedObjectiveRuns(); + const { getValue, setEnabled } = renderConnectedProvider({ + source, + runDetachedObjective: fake.runDetachedObjective, + }); + + await act(async () => { + await getValue().createOptimization(input); + }); + await waitFor(() => + expect(getValue().optimizations[0]?.status).toBe("running"), + ); + expect(calls).toEqual({ connect: 1, dispose: 0 }); + expect(getValue().optimizations[0]?.navigation).toEqual({ + positions: { infected_ratio: 25 }, + booleans: {}, + followTrials: true, + }); + expect( + sessionStorage.getItem("petrinaut:active-optimization-runs"), + "a run in this page cannot be re-attached to after a reload", + ).toBeNull(); + + setEnabled(false); + expect(calls).toEqual({ connect: 1, dispose: 1 }); + await waitFor(() => + expect(getValue().optimizations[0]?.status).toBe("cancelled"), + ); + await expect(getValue().createOptimization(input)).rejects.toThrow( + "Optimization is unavailable", + ); + + setEnabled(true); + await act(async () => { + await getValue().createOptimization(input); + }); + expect(calls).toEqual({ connect: 2, dispose: 1 }); + }); + + it("does not re-attach stored runs through a connected source", async () => { + sessionStorage.setItem( + "petrinaut:active-optimization-runs", + JSON.stringify({ "run-stale": { input, createdAt: 1 } }), + ); + const { source, calls } = createQuietConnectedSource(); + const fake = createFakeDetachedObjectiveRuns(); + const { getValue } = renderConnectedProvider({ + source, + runDetachedObjective: fake.runDetachedObjective, + }); + + await act(async () => { + await Promise.resolve(); + }); + expect(calls.connect).toBe(0); + expect(getValue().optimizations).toHaveLength(0); + expect( + sessionStorage.getItem("petrinaut:active-optimization-runs"), + ).not.toBeNull(); + }); + + it("uses a remote capability regardless of the In-browser optimization setting", async () => { + const capability: PetrinautOptimization = { + createOptimizationRun: () => Promise.resolve({ runId: "run-remote" }), + async *attachOptimizationRun(_runId, options) { + options?.onAttached?.(); + yield { type: "started", requestedTrials: 2, seq: 1 }; + }, + cancelOptimizationRun: () => Promise.resolve(), + }; + let latest: OptimizationsContextValue | null = null; + render( + + + + { + latest = value; + }} + /> + + + , + ); + const getValue = () => { + if (!latest) { + throw new Error("Optimization context was not captured"); + } + return latest; + }; + + await act(async () => { + await getValue().createOptimization(input, { computeBackend: "webgpu" }); + }); + await waitFor(() => + expect(getValue().optimizations[0]?.runId).toBe("run-remote"), + ); + // A remote study computes nothing locally: no backend choice, no navigation. + expect(getValue().optimizations[0]).toMatchObject({ + computeBackend: "cpu", + navigation: null, + selection: null, + axes: [expect.objectContaining({ identifier: "infected_ratio" })], + }); + }); + + it("evaluates a connected study's trials through runDetachedObjective, following each step, then refines the selection", async () => { + const { source, calls } = createEvaluatingSource([0.05, 0.02]); + const fake = createFakeDetachedObjectiveRuns(); + const { getValue, unmount } = renderConnectedProvider({ + source, + runDetachedObjective: fake.runDetachedObjective, + }); + + let optimizationId = ""; + await act(async () => { + optimizationId = await getValue().createOptimization(input, { + computeBackend: "webgpu", + }); + }); + + // Trial 0 runs on the study's backend with its seeds pinned, and the + // navigation follows it while its batch streams as the selection. + await waitFor(() => expect(fake.runs).toHaveLength(1)); + expect(fake.runs[0]!.request).toMatchObject({ + cacheKey: "run-connected", + seed: 1, + runCount: 3, + runSeeds: [1, 2, 3], + computeBackend: "webgpu", + scenarioParameterValues: { population: 1_000, infected_ratio: 0.05 }, + }); + const followedPosition = optimizationAxisPositionFor( + infectedRatioAxis, + 0.05, + ); + await waitFor(() => + expect(getValue().optimizations[0]?.selection?.key).toBe("trial:0"), + ); + expect(getValue().optimizations[0]).toMatchObject({ + computeBackend: "webgpu", + computeBackendFallbackReason: null, + navigation: { + positions: { infected_ratio: followedPosition }, + followTrials: true, + }, + selection: { key: "trial:0", runTarget: null, computing: true }, + }); + const streamed = distributionFrame(metricId, 1, [[0.2, 3]]); + fake.runs[0]!.frames.set([streamed]); + await waitFor(() => + expect(getValue().optimizations[0]?.selection?.metricFrames).toEqual([ + streamed, + ]), + ); + + // Its outcome reaches Optuna; the first fallback reason lands on the record. + fake.runs[0]!.settle( + completedRunResult({ + metricId, + frames: [distributionFrame(metricId, 180, [[0.25, 3]])], + runValues: [0.25, 0.25, 0.25], + computeBackend: "cpu", + fallbackReason: "no adapter", + }), + ); + await waitFor(() => expect(fake.runs).toHaveLength(2)); + await waitFor(() => + expect(getValue().optimizations[0]?.trials).toEqual([ + expect.objectContaining({ + trial: 0, + objective: 0.25, + state: "complete", + }), + ]), + ); + // The record names the backend the trials ran on, not the one asked for. + expect(getValue().optimizations[0]).toMatchObject({ + computeBackend: "cpu", + computeBackendFallbackReason: "no adapter", + }); + expect(getValue().optimizations[0]?.selection?.key).toBe("trial:1"); + expect( + getValue().optimizations[0]?.navigation?.positions.infected_ratio, + ).toBe(optimizationAxisPositionFor(infectedRatioAxis, 0.02)); + + fake.runs[1]!.settle( + completedRunResult({ + metricId, + frames: [distributionFrame(metricId, 180, [[0.125, 3]])], + runValues: [0.125, 0.125, 0.125], + }), + ); + await waitFor(() => + expect(getValue().optimizations[0]?.status).toBe("complete"), + ); + expect(getValue().optimizations[0]?.best).toMatchObject({ + trial: 1, + objective: 0.125, + }); + + // Complete: the selection refines at the followed point, up the ladder. + const lastPosition = optimizationAxisPositionFor(infectedRatioAxis, 0.02); + await waitFor(() => expect(fake.runs).toHaveLength(3)); + expect(fake.runs[2]!.request).toMatchObject({ + cacheKey: optimizationId, + computeBackend: "webgpu", + seed: 1, + runCount: 8, + scenarioParameterValues: { + population: 1_000, + infected_ratio: optimizationAxisValueAt( + infectedRatioAxis, + lastPosition, + ), + }, + }); + expect(fake.runs[2]!.request.runSeeds).toBeUndefined(); + await waitFor(() => + expect(getValue().optimizations[0]?.selection).toMatchObject({ + key: `infected_ratio=${lastPosition}`, + runsCompleted: 0, + runTarget: 8, + computing: true, + }), + ); + fake.runs[2]!.settle( + completedRunResult({ + metricId, + frames: [distributionFrame(metricId, 180, [[0.1, 8]])], + runsCompleted: 8, + }), + ); + await waitFor(() => expect(fake.runs).toHaveLength(4)); + expect(fake.runs[3]!.request.runCount).toBe(17); + await waitFor(() => + expect(getValue().optimizations[0]?.selection).toMatchObject({ + runsCompleted: 8, + runTarget: 25, + computing: true, + }), + ); + + // A navigation change cancels the batch in flight and refines the new point. + act(() => { + getValue().setOptimizationNavigation(optimizationId, { + positions: { infected_ratio: 3 }, + }); + }); + expect(fake.runs[3]!.cancelled).toBe(true); + await waitFor(() => expect(fake.runs).toHaveLength(5)); + expect(fake.runs[4]!.request).toMatchObject({ + runCount: 8, + scenarioParameterValues: { + infected_ratio: optimizationAxisValueAt(infectedRatioAxis, 3), + }, + }); + expect(getValue().optimizations[0]?.selection?.key).toBe( + "infected_ratio=3", + ); + + // Removing the study cancels its batches; unmounting disposes the source. + act(() => getValue().removeOptimization(optimizationId)); + expect(fake.runs[4]!.cancelled).toBe(true); + expect(getValue().optimizations).toHaveLength(0); + expect(calls).toEqual({ connect: 1, dispose: 0 }); + unmount(); + expect(calls).toEqual({ connect: 1, dispose: 1 }); + }); + + it("lets a user move stop following while the study runs, refining beside the trials", async () => { + const { source } = createEvaluatingSource([0.05, 0.02]); + const fake = createFakeDetachedObjectiveRuns(); + const { getValue } = renderConnectedProvider({ + source, + runDetachedObjective: fake.runDetachedObjective, + }); + + let optimizationId = ""; + await act(async () => { + optimizationId = await getValue().createOptimization(input); + }); + await waitFor(() => + expect(getValue().optimizations[0]?.selection?.key).toBe("trial:0"), + ); + expect(fake.runs[0]!.request.computeBackend).toBe("cpu"); + + act(() => { + getValue().setOptimizationNavigation(optimizationId, { + positions: { infected_ratio: 40 }, + }); + }); + await waitFor(() => + expect(getValue().optimizations[0]?.navigation).toEqual({ + positions: { infected_ratio: 40 }, + booleans: {}, + followTrials: false, + }), + ); + expect(fake.runs[1]!.request).toMatchObject({ + cacheKey: optimizationId, + computeBackend: "cpu", + runCount: 8, + scenarioParameterValues: { + infected_ratio: optimizationAxisValueAt(infectedRatioAxis, 40), + }, + }); + expect(getValue().optimizations[0]?.selection?.key).toBe( + "infected_ratio=40", + ); + + // The next trial starts without moving the navigation or the selection. + fake.runs[0]!.settle( + completedRunResult({ + metricId, + frames: [distributionFrame(metricId, 180, [[0.25, 3]])], + runValues: [0.25, 0.25, 0.25], + }), + ); + await waitFor(() => expect(fake.runs).toHaveLength(3)); + expect(fake.runs[2]!.request.cacheKey).toBe("run-connected"); + expect(getValue().optimizations[0]?.navigation?.positions).toEqual({ + infected_ratio: 40, + }); + expect(getValue().optimizations[0]?.selection?.key).toBe( + "infected_ratio=40", + ); + expect(fake.runs[1]!.cancelled).toBe(false); + }); +}); + +/** + * A connected source shaped like the in-browser optimizer's lifecycle: a run + * log in segments, each begun by `started` and ended by a terminal event, + * which a settled study continues with more trials; a stop ends the segment + * once the trial in flight has settled. Segment `n` evaluates + * `ratiosBySegment[n]`, one trial per value, through the channel. + */ +const createResumableSource = ( + ratiosBySegment: readonly (readonly number[])[], + { rejectExtension }: { rejectExtension?: string } = {}, +) => { + const calls = { extend: [] as number[], release: [] as string[], cancel: 0 }; + // The cancelled terminal is the worker's own message, sent once the pruned + // steps in flight have reported; a test decides when it arrives. + let closeStoppedSegment: () => void = () => {}; + const source: PetrinautConnectedOptimization = { + kind: "connected", + connect: (channel) => { + const events: PetrinautOptimizationEvent[] = []; + const listeners = new Set<() => void>(); + let controller = new AbortController(); + let segment = 0; + let trial = 0; + let requested = 0; + let running = false; + let cancelled = false; + // Read through a call so the flag is re-checked after each await (a + // plain property read would be control-flow-narrowed to `false`). + const isCancelled = () => cancelled; + const append = (event: UnsequencedEvent) => { + events.push({ + ...event, + seq: events.length + 1, + } as PetrinautOptimizationEvent); + for (const listener of listeners) { + listener(); + } + }; + const runSegment = async (ratios: readonly number[]) => { + running = true; + cancelled = false; + for (const ratio of ratios) { + if (isCancelled()) { + break; + } + const suggestedValues = { infected_ratio: ratio }; + const outcome = await channel.evaluateTrial({ + runId: "run-resumable", + trial, + manifest: input, + suggestedValues, + scenarioParameterValues: resolveTrialScenarioParameterValues( + input, + suggestedValues, + ), + seeds: [1, 2, 3], + signal: controller.signal, + }); + append({ + type: "trial", + trial, + parameters: suggestedValues, + objective: outcome.kind === "objective" ? outcome.objective : null, + state: outcome.kind === "objective" ? "complete" : "pruned", + best: null, + }); + trial += 1; + } + if (isCancelled()) { + await new Promise((resolve) => { + closeStoppedSegment = resolve; + }); + } + running = false; + append( + isCancelled() + ? { + type: "error", + code: PETRINAUT_OPTIMIZATION_CANCELLED_ERROR_CODE, + message: "optimization cancelled", + retryable: false, + } + : { + type: "complete", + requestedTrials: requested, + completedTrials: trial, + prunedTrials: 0, + failedTrials: 0, + best: null, + }, + ); + }; + return { + createOptimizationRun: () => { + const ratios = ratiosBySegment[0] ?? []; + requested = ratios.length; + append({ type: "started", requestedTrials: requested }); + // The worker asks for its first evaluation a task after the run + // is created, once the provider knows the run id. + setTimeout(() => void runSegment(ratios), 0); + return Promise.resolve({ runId: "run-resumable" }); + }, + extendOptimizationRun: (_runId, trials) => { + if (rejectExtension !== undefined) { + return Promise.reject(new Error(rejectExtension)); + } + if (running) { + return Promise.reject(new Error("still running")); + } + calls.extend.push(trials); + segment += 1; + requested += trials; + controller = new AbortController(); + append({ type: "started", requestedTrials: requested }); + const ratios = ratiosBySegment[segment] ?? []; + setTimeout(() => void runSegment(ratios), 0); + return Promise.resolve(); + }, + async *attachOptimizationRun(_runId, options) { + options?.onAttached?.(); + let index = options?.cursor ?? 0; + for (;;) { + const event = events[index]; + if (event) { + index += 1; + yield event; + if (event.type === "complete" || event.type === "error") { + return; + } + continue; + } + await new Promise((resolve) => { + const wake = () => { + listeners.delete(wake); + resolve(); + }; + listeners.add(wake); + options?.signal?.addEventListener("abort", wake, { once: true }); + }); + if (options?.signal?.aborted) { + return; + } + } + }, + cancelOptimizationRun: () => { + calls.cancel += 1; + cancelled = true; + controller.abort(); + return Promise.resolve(); + }, + releaseOptimizationRun: (runId) => { + calls.release.push(runId); + return Promise.resolve(); + }, + dispose: () => {}, + }; + }, + }; + return { source, calls, closeStoppedSegment: () => closeStoppedSegment() }; +}; + +describe("OptimizationsProvider lifecycle of a connected study", () => { + it("settles on the best, then continues from its cursor, following the new steps", async () => { + const { source, calls } = createResumableSource([[0.05], [0.02]]); + const fake = createFakeDetachedObjectiveRuns(); + const { getValue } = renderConnectedProvider({ + source, + runDetachedObjective: fake.runDetachedObjective, + }); + + let optimizationId = ""; + await act(async () => { + optimizationId = await getValue().createOptimization(input, { + parallelism: 2, + }); + }); + await waitFor(() => expect(fake.runs).toHaveLength(1)); + expect(fake.runs[0]!.request.queueKey).toBe("run-resumable:trial:0"); + await waitFor(() => + expect(getValue().optimizations[0]).toMatchObject({ + parallelism: 2, + resumable: false, + inFlight: [ + { trial: 0, parameters: { infected_ratio: 0.05 }, objective: null }, + ], + activity: [expect.objectContaining({ label: "Step 1", runCount: 3 })], + }), + ); + + fake.runs[0]!.settle( + completedRunResult({ + metricId, + frames: [distributionFrame(metricId, 180, [[0.25, 3]])], + runValues: [0.25, 0.25, 0.25], + }), + ); + await waitFor(() => + expect(getValue().optimizations[0]?.status).toBe("complete"), + ); + // Following ended where the study did best, and that point refines. + const bestPosition = optimizationAxisPositionFor(infectedRatioAxis, 0.05); + expect(getValue().optimizations[0]).toMatchObject({ + resumable: true, + requestedTrials: 1, + navigation: { + positions: { infected_ratio: bestPosition }, + followTrials: false, + }, + inFlight: [], + }); + await waitFor(() => expect(fake.runs).toHaveLength(2)); + expect(fake.runs[1]!.request).toMatchObject({ + cacheKey: optimizationId, + scenarioParameterValues: { + infected_ratio: optimizationAxisValueAt( + infectedRatioAxis, + bestPosition, + ), + }, + }); + + await act(async () => { + await getValue().extendOptimization(optimizationId, 1); + }); + expect(calls.extend).toEqual([1]); + expect(fake.runs[1]!.cancelled).toBe(true); + await waitFor(() => + expect(getValue().optimizations[0]).toMatchObject({ + status: "running", + resumable: false, + requestedTrials: 2, + navigation: { followTrials: true }, + }), + ); + await waitFor(() => expect(fake.runs).toHaveLength(3)); + expect(fake.runs[2]!.request.queueKey).toBe("run-resumable:trial:1"); + await waitFor(() => + expect(getValue().optimizations[0]?.selection?.key).toBe("trial:1"), + ); + + fake.runs[2]!.settle( + completedRunResult({ + metricId, + frames: [distributionFrame(metricId, 180, [[0.125, 3]])], + runValues: [0.125, 0.125, 0.125], + }), + ); + await waitFor(() => + expect(getValue().optimizations[0]?.status).toBe("complete"), + ); + expect(getValue().optimizations[0]).toMatchObject({ + resumable: true, + requestedTrials: 2, + completedTrials: 2, + trials: [ + expect.objectContaining({ trial: 0 }), + expect.objectContaining({ trial: 1 }), + ], + }); + expect(getValue().optimizations[0]?.best).toMatchObject({ + trial: 1, + objective: 0.125, + }); + + act(() => getValue().removeOptimization(optimizationId)); + expect(calls.release).toEqual(["run-resumable"]); + expect(getValue().optimizations).toHaveLength(0); + }); + + it("stops a study without dropping its attachment, so the segment's terminal event lands before a continuation", async () => { + const { source, calls, closeStoppedSegment } = createResumableSource([ + [0.05, 0.02], + [0.01], + ]); + const fake = createFakeDetachedObjectiveRuns(); + const { getValue } = renderConnectedProvider({ + source, + runDetachedObjective: fake.runDetachedObjective, + }); + + let optimizationId = ""; + await act(async () => { + optimizationId = await getValue().createOptimization(input); + }); + await waitFor(() => expect(fake.runs).toHaveLength(1)); + + act(() => getValue().cancelOptimization(optimizationId)); + expect(calls.cancel).toBe(1); + expect(getValue().optimizations[0]).toMatchObject({ + status: "cancelled", + resumable: false, + }); + // The trial in flight is pruned as cancelled and reports before the + // worker acknowledges the stop. + await waitFor(() => expect(getValue().optimizations[0]?.lastSeq).toBe(2)); + expect(getValue().optimizations[0]).toMatchObject({ + status: "cancelled", + resumable: false, + prunedTrials: 1, + }); + closeStoppedSegment(); + await waitFor(() => expect(getValue().optimizations[0]?.lastSeq).toBe(3)); + expect(getValue().optimizations[0]).toMatchObject({ + status: "cancelled", + resumable: true, + prunedTrials: 1, + }); + + await act(async () => { + await getValue().extendOptimization(optimizationId, 1); + }); + await waitFor(() => + expect(getValue().optimizations[0]).toMatchObject({ + status: "running", + requestedTrials: 3, + lastSeq: 4, + }), + ); + // The stop settled the study on a point, which began refining (the + // second run); the continuation cancels that and runs the new trial. + await waitFor(() => expect(fake.runs).toHaveLength(3)); + expect(fake.runs[1]!.request.cacheKey).toBe(optimizationId); + expect(fake.runs[1]!.cancelled).toBe(true); + // The stopped segment asked for its second trial never, so numbering + // continues from the pruned one. + expect(fake.runs[2]!.request).toMatchObject({ + queueKey: "run-resumable:trial:1", + scenarioParameterValues: { infected_ratio: 0.01 }, + }); + }); + + it("puts a refused continuation on the record and leaves the study resumable", async () => { + const { source } = createResumableSource([[0.05]], { + rejectExtension: "An optimization may run at most 1,000 trials in total", + }); + const fake = createFakeDetachedObjectiveRuns(); + const { getValue } = renderConnectedProvider({ + source, + runDetachedObjective: fake.runDetachedObjective, + }); + + let optimizationId = ""; + await act(async () => { + optimizationId = await getValue().createOptimization(input); + }); + await waitFor(() => expect(fake.runs).toHaveLength(1)); + fake.runs[0]!.settle( + completedRunResult({ + metricId, + frames: [distributionFrame(metricId, 180, [[0.25, 3]])], + runValues: [0.25, 0.25, 0.25], + }), + ); + await waitFor(() => + expect(getValue().optimizations[0]?.resumable).toBe(true), + ); + + await expect( + getValue().extendOptimization(optimizationId, 999), + ).rejects.toThrow("at most 1,000 trials"); + // The refusal's state update landed outside an act scope; flush it. + await act(async () => { + await Promise.resolve(); + }); + expect(getValue().optimizations[0]).toMatchObject({ + status: "complete", + resumable: true, + error: "An optimization may run at most 1,000 trials in total", + }); + }); + + it("never marks a remote run resumable", async () => { + const capability: PetrinautOptimization = { + createOptimizationRun: () => Promise.resolve({ runId: "run-remote-2" }), + async *attachOptimizationRun() { + yield { + type: "complete", + requestedTrials: 2, + completedTrials: 0, + prunedTrials: 0, + failedTrials: 0, + best: null, + seq: 1, + }; + }, + cancelOptimizationRun: () => Promise.resolve(), + }; + const getValue = renderProvider(capability); + + await act(async () => { + await getValue().createOptimization(input); + }); + await waitFor(() => + expect(getValue().optimizations[0]?.status).toBe("complete"), + ); + expect(getValue().optimizations[0]?.resumable).toBe(false); + await expect( + getValue().extendOptimization(getValue().optimizations[0]!.id, 1), + ).rejects.toThrow("cannot be continued"); + }); }); diff --git a/libs/@hashintel/petrinaut/src/react/optimizations/provider.tsx b/libs/@hashintel/petrinaut/src/react/optimizations/provider.tsx index 277236662d4..a505675d305 100644 --- a/libs/@hashintel/petrinaut/src/react/optimizations/provider.tsx +++ b/libs/@hashintel/petrinaut/src/react/optimizations/provider.tsx @@ -1,3 +1,7 @@ +/** + * @layerRoot react.optimizations + * @role Tracks optimization runs, folds their event streams into records, and drives a connected study's navigation and live selection + */ import { use, useCallback, useEffect, useRef, useState } from "react"; import { @@ -7,13 +11,26 @@ import { type PetrinautOptimizationEvent, type PetrinautOptimizationInput, } from "@hashintel/petrinaut-core"; +import { + isConnectedOptimization, + type PetrinautConnectedOptimization, + type PetrinautConnectedOptimizationCapability, +} from "@hashintel/petrinaut-core/optimization"; +import { + ExperimentsActionsContext, + type ExperimentsActionsValue, +} from "../experiments/context"; import { useBlockWindowClose } from "../hooks/use-block-window-close"; +import { useLatest } from "../hooks/use-latest"; import { openPetrinautSimulationResource, usePetrinautNavigation, } from "../navigation"; -import { PetrinautOptimizationContext } from "../optimization-context"; +import { + createOptimizationChannel, + type OptimizationChannelStudy, +} from "./channel/create-optimization-channel"; import { type OptimizationBest, type OptimizationErrorCategory, @@ -23,6 +40,13 @@ import { OptimizationsContext, type OptimizationsContextValue, } from "./context"; +import { + type ConnectedStudy, + type ConnectedStudyOutcome, + createConnectedStudy, +} from "./provider/connected-study"; +import { buildOptimizationSurfaceAxes } from "./surface-grid"; +import { useOptimizationSource } from "./use-optimization-source"; import type { PropsWithChildren } from "react"; @@ -293,15 +317,59 @@ const createOptimizationRecord = ( failedTrials: 0, trials: [], best: null, + resumable: false, + parallelism: 1, + computeBackend: "cpu", + computeBackendFallbackReason: null, + axes: buildOptimizationSurfaceAxes(input), + navigation: null, + selection: null, + activity: [], + inFlight: [], ...overrides, }); +/** + * A connected source's capability together with the channel it evaluates + * trials through. Both die with the connection. + */ +type OptimizationConnection = { + source: PetrinautConnectedOptimization; + capability: PetrinautConnectedOptimizationCapability; + dispose: () => void; +}; + +const connectOptimizationSource = ( + source: PetrinautConnectedOptimization, + experimentsActions: React.RefObject, + resolveStudy: (runId: string) => OptimizationChannelStudy | null, +): OptimizationConnection => { + const channel = createOptimizationChannel({ + runDetachedObjective: (request) => + experimentsActions.current.runDetachedObjective(request), + resolveStudy, + }); + const capability = source.connect(channel); + return { + source, + capability, + dispose: () => { + capability.dispose(); + channel.dispose(); + }, + }; +}; + export const OptimizationsProvider = ({ children }: PropsWithChildren) => { - const capability = use(PetrinautOptimizationContext); + const source = useOptimizationSource(); + const experimentsActionsRef = useLatest(use(ExperimentsActionsContext)); + const connectionRef = useRef(null); const navigation = usePetrinautNavigation(); const abortControllersRef = useRef(new Map()); /** Server run ids of active detached runs, keyed by record id. */ const runIdsRef = useRef(new Map()); + /** The local machinery behind each connected study, keyed by record id. */ + const studiesRef = useRef(new Map()); const [optimizations, setOptimizations] = useState([]); const selectedOptimizationId = navigation.state.simulateResource?.type === "optimization" @@ -326,11 +394,16 @@ export const OptimizationsProvider = ({ children }: PropsWithChildren) => { useEffect(() => { const abortControllers = abortControllersRef.current; + const studies = studiesRef.current; return () => { for (const controller of abortControllers.values()) { controller.abort(); } abortControllers.clear(); + for (const study of studies.values()) { + study.dispose(); + } + studies.clear(); }; }, []); @@ -368,16 +441,45 @@ export const OptimizationsProvider = ({ children }: PropsWithChildren) => { } }, [navigation, optimizations, selectedOptimizationId]); + const settleStudy = ( + optimizationId: string, + outcome: ConnectedStudyOutcome, + best?: OptimizationBest | null, + ) => { + studiesRef.current.get(optimizationId)?.settle(outcome, best); + }; + + const disposeStudy = (optimizationId: string) => { + studiesRef.current.get(optimizationId)?.dispose(); + studiesRef.current.delete(optimizationId); + }; + + /** + * Whether a settled record can run more steps: a connected study whose + * local machinery is still here. The machinery goes when the study is + * removed or its connection is disposed, and with it the kept sampler. + */ + const resumableAfterSettling = ( + optimizationId: string, + current: OptimizationRecord, + ): boolean => + current.navigation !== null && studiesRef.current.has(optimizationId); + const markOptimizationCancelled = useCallback( (optimizationId: string) => { patchOptimization(optimizationId, (current) => ({ ...current, status: "cancelled", + // The segment's terminal event, not this mark, makes a connected + // study resumable: a stop lands here while its steps in flight are + // still being pruned, and the core refuses to extend it until then. + resumable: false, error: null, errorCategory: null, errorDiagnostics: null, connectionState: null, })); + settleStudy(optimizationId, "cancelled"); }, [patchOptimization], ); @@ -391,6 +493,7 @@ export const OptimizationsProvider = ({ children }: PropsWithChildren) => { patchOptimization(optimizationId, (current) => ({ ...current, status: "error", + resumable: false, connectionState: null, // A classified transport failure yields a safe, actionable message // and correlation ids; anything else keeps its message. @@ -402,6 +505,7 @@ export const OptimizationsProvider = ({ children }: PropsWithChildren) => { errorCategory: classified?.category ?? null, errorDiagnostics: classified?.diagnostics ?? null, })); + settleStudy(optimizationId, "error"); }, [patchOptimization], ); @@ -426,6 +530,7 @@ export const OptimizationsProvider = ({ children }: PropsWithChildren) => { ...current, ...extra, status: "running", + resumable: false, requestedTrials: event.requestedTrials, })); break; @@ -433,7 +538,9 @@ export const OptimizationsProvider = ({ children }: PropsWithChildren) => { patchOptimization(optimizationId, (current) => ({ ...current, ...extra, - status: "running", + // A stopped study's pruned steps still report; they do not + // revive it. + status: current.status === "cancelled" ? "cancelled" : "running", completedTrials: current.completedTrials + (event.state === "complete" ? 1 : 0), prunedTrials: @@ -443,12 +550,14 @@ export const OptimizationsProvider = ({ children }: PropsWithChildren) => { trials: [...current.trials, event], best: event.best ?? computeRunningBest(current, event), })); + studiesRef.current.get(optimizationId)?.trialReported(event); break; case "complete": patchOptimization(optimizationId, (current) => ({ ...current, ...extra, status: "complete", + resumable: resumableAfterSettling(optimizationId, current), connectionState: null, // The complete event's requested-trial count is the true total, // but its completed/pruned/failed counts only cover the frames @@ -458,6 +567,7 @@ export const OptimizationsProvider = ({ children }: PropsWithChildren) => { requestedTrials: event.requestedTrials, best: event.best ?? current.best, })); + settleStudy(optimizationId, "complete", event.best); break; case "error": patchOptimization(optimizationId, (current) => ({ @@ -474,12 +584,23 @@ export const OptimizationsProvider = ({ children }: PropsWithChildren) => { ...(event.code === PETRINAUT_OPTIMIZATION_CANCELLED_ERROR_CODE ? { status: "cancelled" as const, + resumable: resumableAfterSettling(optimizationId, current), error: null, errorCategory: null, errorDiagnostics: null, } - : { status: "error" as const, error: event.message }), + : { + status: "error" as const, + resumable: false, + error: event.message, + }), })); + settleStudy( + optimizationId, + event.code === PETRINAUT_OPTIMIZATION_CANCELLED_ERROR_CODE + ? "cancelled" + : "error", + ); break; } }, @@ -517,6 +638,7 @@ export const OptimizationsProvider = ({ children }: PropsWithChildren) => { attach, cancel, abortController, + cursor = 0, dropRecordOnNotFound = false, }: { optimizationId: string; @@ -524,6 +646,8 @@ export const OptimizationsProvider = ({ children }: PropsWithChildren) => { attach: PetrinautOptimization["attachOptimizationRun"]; cancel: PetrinautOptimization["cancelOptimizationRun"]; abortController: AbortController; + /** The record's last applied `seq`, when it already holds earlier events. */ + cursor?: number; /** * Silently drop the record when the very first attachment 404s — used * when re-attaching to a stored run that may have expired server-side. @@ -534,7 +658,7 @@ export const OptimizationsProvider = ({ children }: PropsWithChildren) => { // Read through a call so the abort flag is re-checked after each await // (a plain property read would be control-flow-narrowed to `false`). const isCancelled = () => signal.aborted; - let lastSeq = 0; + let lastSeq = cursor; let sawTerminalEvent = false; let consecutiveFailures = 0; let receivedAnyEvent = false; @@ -683,8 +807,90 @@ export const OptimizationsProvider = ({ children }: PropsWithChildren) => { ], ); + /** + * The study behind an optimizer run id, for the channel: its requested + * backend, and the hooks that follow its trials. The first trial that ran + * elsewhere than asked records where, and why, on the record. + */ + const resolveChannelStudy = useCallback( + (runId: string): OptimizationChannelStudy | null => { + const entry = [...runIdsRef.current].find( + ([, knownRunId]) => knownRunId === runId, + ); + const study = entry ? studiesRef.current.get(entry[0]) : undefined; + if (!entry || !study) { + return null; + } + const [optimizationId] = entry; + return { + computeBackend: study.computeBackend, + trialStarted: study.trialStarted, + trialSettled: (trial, outcome) => { + study.trialSettled(trial, outcome); + if (outcome.ok && outcome.computeBackendFallbackReason !== null) { + const { computeBackend, computeBackendFallbackReason } = outcome; + patchOptimization(optimizationId, (current) => + current.computeBackendFallbackReason === null + ? { ...current, computeBackend, computeBackendFallbackReason } + : current, + ); + } + }, + }; + }, + [patchOptimization], + ); + + /** + * The capability behind the source: the remote one as given, or a connected + * one wired to the experiments backend on first use and kept while the + * source stays the same. Connecting happens on demand rather than in render + * so a source never connects twice, and the cleanup below tears the + * connection down, with the runs made through it, when the source changes + * or the provider unmounts. + */ + const resolveCapability = useCallback((): PetrinautOptimization | null => { + if (source === null || !isConnectedOptimization(source)) { + return source; + } + const current = connectionRef.current; + if (current?.source === source) { + return current.capability; + } + current?.dispose(); + const connection = connectOptimizationSource( + source, + experimentsActionsRef, + resolveChannelStudy, + ); + connectionRef.current = connection; + return connection.capability; + }, [experimentsActionsRef, resolveChannelStudy, source]); + + useEffect( + () => () => { + const connection = connectionRef.current; + if (connection?.source === source) { + connection.dispose(); + connectionRef.current = null; + // A connected capability's runs end with its connection: aborting + // their attach loops settles each record as cancelled, and the + // studies' own batches stop with them. + for (const controller of abortControllersRef.current.values()) { + controller.abort(); + } + for (const study of studiesRef.current.values()) { + study.dispose(); + } + studiesRef.current.clear(); + } + }, + [source], + ); + const createOptimization: OptimizationsContextValue["createOptimization"] = - async (rawInput) => { + async (rawInput, options) => { + const capability = resolveCapability(); if (!capability) { throw new Error("Optimization is unavailable"); } @@ -692,10 +898,42 @@ export const OptimizationsProvider = ({ children }: PropsWithChildren) => { const input = petrinautOptimizationInputSchema.parse(rawInput); const optimizationId = crypto.randomUUID(); const abortController = new AbortController(); + const connection = + connectionRef.current?.capability === capability + ? connectionRef.current + : null; + const connected = connection !== null; + const computeBackend = connected + ? (options?.computeBackend ?? "cpu") + : "cpu"; + const parallelism = connected ? (options?.parallelism ?? 1) : 1; + const study = connected + ? createConnectedStudy({ + optimizationId, + input, + axes: buildOptimizationSurfaceAxes(input), + computeBackend, + runDetachedObjective: (request) => + experimentsActionsRef.current.runDetachedObjective(request), + onUpdate: (update) => { + patchOptimization(optimizationId, (current) => ({ + ...current, + ...update, + })); + }, + }) + : null; + if (study) { + studiesRef.current.set(optimizationId, study); + } abortControllersRef.current.set(optimizationId, abortController); setOptimizations((current) => [ - createOptimizationRecord(optimizationId, input), + createOptimizationRecord(optimizationId, input, { + computeBackend, + parallelism, + navigation: study?.initialNavigation ?? null, + }), ...current, ]); setSelectedOptimizationId(optimizationId); @@ -703,9 +941,14 @@ export const OptimizationsProvider = ({ children }: PropsWithChildren) => { const consumeRun = async () => { let runId: string; try { - ({ runId } = await capability.createOptimizationRun(input, { - signal: abortController.signal, - })); + ({ runId } = await (connection + ? connection.capability.createOptimizationRun(input, { + signal: abortController.signal, + parallelism, + }) + : capability.createOptimizationRun(input, { + signal: abortController.signal, + }))); } catch (error) { const classified = classifyError(error); if ( @@ -729,7 +972,11 @@ export const OptimizationsProvider = ({ children }: PropsWithChildren) => { } runIdsRef.current.set(optimizationId, runId); - storeActiveRun(runId, input); + if (!connected) { + // A connected study's run lives in this page; a reload cannot + // re-attach to it. + storeActiveRun(runId, input); + } patchOptimization(optimizationId, (current) => ({ ...current, runId, @@ -750,8 +997,13 @@ export const OptimizationsProvider = ({ children }: PropsWithChildren) => { }; void consumeRun().finally(() => { - abortControllersRef.current.delete(optimizationId); - runIdsRef.current.delete(optimizationId); + // A continuation may have taken the entries over by now. + if ( + abortControllersRef.current.get(optimizationId) === abortController + ) { + abortControllersRef.current.delete(optimizationId); + runIdsRef.current.delete(optimizationId); + } }); return optimizationId; @@ -769,6 +1021,14 @@ export const OptimizationsProvider = ({ children }: PropsWithChildren) => { * again. */ useEffect(() => { + if (source !== null && isConnectedOptimization(source)) { + return; + } + const storedRuns = Object.entries(readStoredActiveRuns()); + if (storedRuns.length === 0) { + return; + } + const capability = resolveCapability(); if (!capability) { return; } @@ -779,7 +1039,7 @@ export const OptimizationsProvider = ({ children }: PropsWithChildren) => { const runIds = runIdsRef.current; const startedIds: string[] = []; - for (const [runId, storedRun] of Object.entries(readStoredActiveRuns())) { + for (const [runId, storedRun] of storedRuns) { const parsedInput = petrinautOptimizationInputSchema.safeParse( storedRun.input, ); @@ -826,7 +1086,7 @@ export const OptimizationsProvider = ({ children }: PropsWithChildren) => { current.filter((optimization) => !startedIds.includes(optimization.id)), ); }; - }, [capability, runAttachLoop]); + }, [resolveCapability, runAttachLoop, source]); /** * The run id of a detached record: from the live-loop map while its attach @@ -844,12 +1104,26 @@ export const OptimizationsProvider = ({ children }: PropsWithChildren) => { optimizationId, ) => { const runId = resolveRunId(optimizationId); + const connected = studiesRef.current.has(optimizationId); if (runId !== undefined) { - runIdsRef.current.delete(optimizationId); removeStoredActiveRun(runId); // Stop the detached run server-side; aborting the local attachment // below only drops this tab's connection to it. - void capability?.cancelOptimizationRun(runId).catch(() => undefined); + void resolveCapability() + ?.cancelOptimizationRun(runId) + .catch(() => undefined); + } + if (connected) { + // The study's segment ends with a terminal event once its steps in + // flight are pruned. The attachment stays to apply it, so the record's + // cursor covers the whole segment and a continuation resumes right + // after it; the status settles here without waiting, and the terminal + // event offers the continuation. + markOptimizationCancelled(optimizationId); + return; + } + if (runId !== undefined) { + runIdsRef.current.delete(optimizationId); } abortControllersRef.current.get(optimizationId)?.abort(); abortControllersRef.current.delete(optimizationId); @@ -863,13 +1137,89 @@ export const OptimizationsProvider = ({ children }: PropsWithChildren) => { if (runId !== undefined) { runIdsRef.current.delete(optimizationId); removeStoredActiveRun(runId); - void capability?.cancelOptimizationRun(runId).catch(() => undefined); + const connection = connectionRef.current; + // A connected study keeps its sampler until it is released; a remote + // run is stopped server-side. + void ( + connection && studiesRef.current.has(optimizationId) + ? connection.capability.releaseOptimizationRun(runId) + : (resolveCapability()?.cancelOptimizationRun(runId) ?? + Promise.resolve()) + ).catch(() => undefined); } abortControllersRef.current.get(optimizationId)?.abort(); abortControllersRef.current.delete(optimizationId); + disposeStudy(optimizationId); dropOptimizationRecord(optimizationId); }; + const extendOptimization: OptimizationsContextValue["extendOptimization"] = + async (optimizationId, trials) => { + const existing = optimizations.find( + (optimization) => optimization.id === optimizationId, + ); + const connection = connectionRef.current; + const study = studiesRef.current.get(optimizationId); + if ( + !existing?.resumable || + existing.runId === null || + !connection || + !study + ) { + throw new Error("This optimization cannot be continued"); + } + const { runId } = existing; + try { + await connection.capability.extendOptimizationRun(runId, trials, { + parallelism: existing.parallelism, + }); + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + patchOptimization(optimizationId, (current) => ({ + ...current, + error: message, + })); + throw error; + } + const abortController = new AbortController(); + abortControllersRef.current.set(optimizationId, abortController); + runIdsRef.current.set(optimizationId, runId); + study.resume(); + patchOptimization(optimizationId, (current) => ({ + ...current, + status: "running", + resumable: false, + error: null, + errorCategory: null, + errorDiagnostics: null, + connectionState: "streaming", + })); + void runAttachLoop({ + optimizationId, + runId, + attach: connection.capability.attachOptimizationRun.bind( + connection.capability, + ), + cancel: connection.capability.cancelOptimizationRun.bind( + connection.capability, + ), + abortController, + cursor: existing.lastSeq, + }).finally(() => { + if ( + abortControllersRef.current.get(optimizationId) === abortController + ) { + abortControllersRef.current.delete(optimizationId); + runIdsRef.current.delete(optimizationId); + } + }); + }; + + const setOptimizationNavigation: OptimizationsContextValue["setOptimizationNavigation"] = + (optimizationId, patch) => { + studiesRef.current.get(optimizationId)?.setNavigation(patch); + }; + const retryOptimization: OptimizationsContextValue["retryOptimization"] = async (optimizationId) => { const existing = optimizations.find( @@ -878,7 +1228,10 @@ export const OptimizationsProvider = ({ children }: PropsWithChildren) => { if (!existing) { return null; } - return createOptimization(existing.input); + return createOptimization(existing.input, { + computeBackend: existing.computeBackend, + parallelism: existing.parallelism, + }); }; const selectedOptimization = @@ -894,6 +1247,8 @@ export const OptimizationsProvider = ({ children }: PropsWithChildren) => { createOptimization, cancelOptimization, removeOptimization, + extendOptimization, + setOptimizationNavigation, retryOptimization, }; diff --git a/libs/@hashintel/petrinaut/src/react/optimizations/provider/activity-registry.test.ts b/libs/@hashintel/petrinaut/src/react/optimizations/provider/activity-registry.test.ts new file mode 100644 index 00000000000..33b5706d53c --- /dev/null +++ b/libs/@hashintel/petrinaut/src/react/optimizations/provider/activity-registry.test.ts @@ -0,0 +1,151 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; + +import { createActivityRegistry } from "./activity-registry"; + +import type { OptimizationBatchStatus } from "../context"; +import type { + MonteCarloWorkerProgress, + ReadableStore, +} from "@hashintel/petrinaut-core"; + +/** A progress store the test ticks by hand. */ +const fakeProgress = () => { + const listeners = new Set<(value: MonteCarloWorkerProgress | null) => void>(); + let completedRuns = 0; + const progress: ReadableStore = { + get: () => ({ completedRuns }) as MonteCarloWorkerProgress, + subscribe: (listener) => { + listeners.add(listener); + return () => { + listeners.delete(listener); + }; + }, + }; + return { + progress, + listenerCount: () => listeners.size, + tick: (runs: number) => { + completedRuns = runs; + for (const listener of listeners) { + listener(progress.get()); + } + }, + }; +}; + +beforeEach(() => { + vi.useFakeTimers(); +}); + +afterEach(() => { + vi.useRealTimers(); +}); + +describe("createActivityRegistry", () => { + it("publishes on register and unregister, and throttles progress ticks", () => { + const published: number[][] = []; + const registry = createActivityRegistry((activity) => { + published.push(activity.map((batch) => batch.completedRuns)); + }); + const step = fakeProgress(); + const unregister = registry.register({ + kind: "step", + label: "Step 1", + runCount: 8, + progress: step.progress, + }); + expect(published).toEqual([[0]]); + + step.tick(3); + expect(published).toEqual([[0], [3]]); + step.tick(5); + vi.advanceTimersByTime(100); + expect(published).toEqual([[0], [3], [5]]); + + unregister(); + expect(published.at(-1)).toEqual([]); + expect(step.listenerCount()).toBe(0); + }); + + it("lists steps before refinement, each in the order it began, with its label and budget", () => { + let latest: readonly OptimizationBatchStatus[] = []; + const registry = createActivityRegistry((activity) => { + latest = activity; + }); + registry.register({ + kind: "refine", + label: "Refining infected_ratio 0.05", + runCount: 17, + progress: fakeProgress().progress, + }); + registry.register({ + kind: "step", + label: "Step 3", + runCount: 3, + progress: fakeProgress().progress, + }); + registry.register({ + kind: "step", + label: "Step 4", + runCount: 3, + progress: fakeProgress().progress, + }); + + expect(latest).toEqual([ + { + id: "step-2", + kind: "step", + label: "Step 3", + runCount: 3, + completedRuns: 0, + }, + { + id: "step-3", + kind: "step", + label: "Step 4", + runCount: 3, + completedRuns: 0, + }, + { + id: "refine-1", + kind: "refine", + label: "Refining infected_ratio 0.05", + runCount: 17, + completedRuns: 0, + }, + ]); + }); + + it("stops listening to every batch on clear, so later ticks publish nothing", () => { + const published: number[][] = []; + const registry = createActivityRegistry((activity) => { + published.push(activity.map((batch) => batch.completedRuns)); + }); + const step = fakeProgress(); + const rung = fakeProgress(); + registry.register({ + kind: "step", + label: "Step 1", + runCount: 3, + progress: step.progress, + }); + const unregisterRung = registry.register({ + kind: "refine", + label: "Refining", + runCount: 8, + progress: rung.progress, + }); + + registry.clear(); + expect(published.at(-1)).toEqual([]); + expect(step.listenerCount()).toBe(0); + expect(rung.listenerCount()).toBe(0); + + const publishedAfterClear = published.length; + step.tick(2); + rung.tick(4); + vi.advanceTimersByTime(500); + unregisterRung(); + expect(published.length).toBe(publishedAfterClear); + }); +}); diff --git a/libs/@hashintel/petrinaut/src/react/optimizations/provider/activity-registry.ts b/libs/@hashintel/petrinaut/src/react/optimizations/provider/activity-registry.ts new file mode 100644 index 00000000000..f73ce077834 --- /dev/null +++ b/libs/@hashintel/petrinaut/src/react/optimizations/provider/activity-registry.ts @@ -0,0 +1,99 @@ +import { createThrottle } from "../../experiments/shared/throttle"; + +import type { OptimizationBatchStatus } from "../context"; +import type { + MonteCarloWorkerProgress, + ReadableStore, +} from "@hashintel/petrinaut-core"; + +export type ActivityRegistry = { + /** Lists the batch until the returned function is called. */ + register: (batch: { + kind: OptimizationBatchStatus["kind"]; + label: string; + runCount: number; + progress: ReadableStore; + }) => () => void; + /** Drops every batch and publishes the empty list. */ + clear: () => void; +}; + +const KIND_ORDER: Record = { + step: 0, + refine: 1, +}; + +/** + * Progress ticks republish on a 100 ms throttle: the list feeds a small + * activity display, not the charts. + */ +const PROGRESS_TICK_MS = 100; + +/** + * Tracks every batch a connected study computes and publishes the sorted + * list on each change — steps first, then refinement, each in the order it + * began. A batch appearing or leaving publishes at once; its progress ticks + * are throttled. + */ +export const createActivityRegistry = ( + onActivity: (activity: readonly OptimizationBatchStatus[]) => void, +): ActivityRegistry => { + let sequence = 0; + const active = new Map< + number, + { + kind: OptimizationBatchStatus["kind"]; + label: string; + runCount: number; + progress: ReadableStore; + offProgress: () => void; + } + >(); + + const publish = () => { + onActivity( + [...active.entries()] + .map(([sequenceNumber, batch]) => ({ + id: `${batch.kind}-${sequenceNumber}`, + kind: batch.kind, + label: batch.label, + runCount: batch.runCount, + completedRuns: batch.progress.get()?.completedRuns ?? 0, + sequenceNumber, + })) + .sort( + (left, right) => + KIND_ORDER[left.kind] - KIND_ORDER[right.kind] || + left.sequenceNumber - right.sequenceNumber, + ) + .map(({ sequenceNumber: _sequenceNumber, ...batch }) => batch), + ); + }; + const progressTick = createThrottle(publish, PROGRESS_TICK_MS); + + return { + register: ({ kind, label, runCount, progress }) => { + const id = ++sequence; + const offProgress = progress.subscribe(progressTick.call); + active.set(id, { kind, label, runCount, progress, offProgress }); + publish(); + return () => { + const batch = active.get(id); + if (batch === undefined) { + return; + } + active.delete(id); + batch.offProgress(); + publish(); + }; + }, + clear: () => { + for (const batch of active.values()) { + batch.offProgress(); + } + active.clear(); + progressTick.cancel(); + publish(); + }, + }; +}; diff --git a/libs/@hashintel/petrinaut/src/react/optimizations/provider/connected-study.test.ts b/libs/@hashintel/petrinaut/src/react/optimizations/provider/connected-study.test.ts new file mode 100644 index 00000000000..06362e84ae3 --- /dev/null +++ b/libs/@hashintel/petrinaut/src/react/optimizations/provider/connected-study.test.ts @@ -0,0 +1,384 @@ +import { describe, expect, it } from "vitest"; + +import { + cancelledRunOutcome, + completedRunResult, + createFakeDetachedObjectiveRuns, + distributionFrame, + failedRunOutcome, +} from "../fake-detached-objective-runs.fixtures"; +import { + sirOptimizationInput, + sirOptimizationMetric, +} from "../sir-optimization-input.fixtures"; +import { + buildOptimizationSurfaceAxes, + optimizationAxisPositionFor, + optimizationAxisValueAt, +} from "../surface-grid"; +import { + createConnectedStudy, + type ConnectedStudyUpdate, +} from "./connected-study"; + +import type { PetrinautOptimizationTrialEvent } from "@hashintel/petrinaut-core"; + +const metricId = sirOptimizationMetric.id; +const axes = buildOptimizationSurfaceAxes(sirOptimizationInput); +const axis = axes[0]!; + +const trialEvent = ( + trial: number, + infectedRatio: number, + objective: number | null, +): PetrinautOptimizationTrialEvent => ({ + type: "trial", + trial, + parameters: { infected_ratio: infectedRatio }, + objective, + state: objective === null ? "pruned" : "complete", + best: null, + seq: trial + 2, +}); + +const setup = () => { + const refinementRuns = createFakeDetachedObjectiveRuns(); + const trialRuns = createFakeDetachedObjectiveRuns(); + const updates: ConnectedStudyUpdate[] = []; + const study = createConnectedStudy({ + optimizationId: "optimization-1", + input: sirOptimizationInput, + axes, + computeBackend: "webgpu", + runDetachedObjective: refinementRuns.runDetachedObjective, + onUpdate: (update) => { + updates.push(update); + }, + }); + /** A trial's batch as the channel would hand it over. */ + const startTrial = (trial: number, infectedRatio: number) => { + const entry = trialRuns.runDetachedObjective({ + cacheKey: "run-1", + definition: sirOptimizationInput.model.definition, + scenarioId: sirOptimizationInput.scenario.id, + scenarioParameterValues: { + population: 1_000, + infected_ratio: infectedRatio, + }, + metric: { id: metricId, label: "m", code: "" }, + seed: 1, + runCount: 3, + dt: 1, + maxTime: 180, + computeBackend: "webgpu", + }); + study.trialStarted(trial, { infected_ratio: infectedRatio }, entry, 3); + return trialRuns.runs.at(-1)!; + }; + return { + refinementRuns, + updates, + study, + startTrial, + latest: () => updates.at(-1), + }; +}; + +describe("createConnectedStudy", () => { + it("starts at the axis midpoints, following trials, computing nothing", () => { + const { study, refinementRuns, updates } = setup(); + expect(study.initialNavigation).toEqual({ + positions: { infected_ratio: 25 }, + booleans: {}, + followTrials: true, + }); + expect(study.computeBackend).toBe("webgpu"); + expect(refinementRuns.runs).toHaveLength(0); + expect(updates).toHaveLength(0); + }); + + it("follows a trial: the navigation moves to its values and its batch streams as the selection", () => { + const { study, startTrial, latest, refinementRuns } = setup(); + + const trial = startTrial(0, 0.05); + expect(latest()?.navigation).toEqual({ + positions: { infected_ratio: optimizationAxisPositionFor(axis, 0.05) }, + booleans: {}, + followTrials: true, + }); + expect(latest()?.selection).toEqual({ + key: "trial:0", + metricFrames: [], + runsCompleted: 0, + runTarget: null, + computing: true, + error: null, + note: null, + }); + expect(latest()?.activity).toEqual([ + { + id: "step-1", + kind: "step", + label: "Step 1", + runCount: 3, + completedRuns: 0, + }, + ]); + expect(latest()?.inFlight).toEqual([ + { trial: 0, parameters: { infected_ratio: 0.05 }, objective: null }, + ]); + + const frame = distributionFrame(metricId, 1, [[0.2, 2]]); + trial.frames.set([frame]); + expect(latest()?.selection).toMatchObject({ + key: "trial:0", + metricFrames: [frame], + computing: true, + }); + expect(latest()?.inFlight[0]?.objective).toBeCloseTo(0.2); + + const result = completedRunResult({ + metricId, + frames: [frame], + runValues: [0.2, 0.2, 0.2], + }); + study.trialSettled(0, result); + expect(latest()?.selection).toEqual({ + key: "trial:0", + metricFrames: [frame], + runsCompleted: 3, + runTarget: null, + computing: false, + error: null, + note: null, + }); + expect(latest()?.activity).toEqual([]); + expect(latest()?.inFlight).toEqual([]); + expect(refinementRuns.runs).toHaveLength(0); + }); + + it("follows the most recently started of several trials in flight, then the next when it settles", () => { + const { study, startTrial, latest } = setup(); + + startTrial(0, 0.05); + const second = startTrial(1, 0.02); + expect(latest()?.selection?.key).toBe("trial:1"); + expect(latest()?.navigation?.positions).toEqual({ + infected_ratio: optimizationAxisPositionFor(axis, 0.02), + }); + expect(latest()?.activity.map((batch) => batch.label)).toEqual([ + "Step 1", + "Step 2", + ]); + expect(latest()?.inFlight.map((step) => step.trial)).toEqual([0, 1]); + + // The unfollowed trial's frames still reach the record as its running value. + const frame = distributionFrame(metricId, 1, [[0.4, 3]]); + second.frames.set([frame]); + expect(latest()?.inFlight[1]?.objective).toBeCloseTo(0.4); + + study.trialSettled( + 1, + completedRunResult({ metricId, frames: [frame], runValues: [0.4] }), + ); + expect(latest()?.selection).toMatchObject({ + key: "trial:0", + computing: true, + }); + expect(latest()?.navigation?.positions).toEqual({ + infected_ratio: optimizationAxisPositionFor(axis, 0.05), + }); + expect(latest()?.inFlight.map((step) => step.trial)).toEqual([0]); + }); + + it("a followed trial's failure lands on the selection with its reason", () => { + const { study, startTrial, latest, refinementRuns } = setup(); + startTrial(0, 0.05); + + study.trialSettled(0, failedRunOutcome(`${metricId}: Unexpected token`)); + expect(latest()?.selection).toEqual({ + key: "trial:0", + metricFrames: [], + runsCompleted: 0, + runTarget: null, + computing: false, + error: `${metricId}: Unexpected token`, + note: null, + }); + expect(refinementRuns.runs).toHaveLength(0); + }); + + it("a user move stops following and refines the new point on the study's backend, listing the rung", () => { + const { study, startTrial, latest, refinementRuns } = setup(); + startTrial(0, 0.05); + + study.setNavigation({ positions: { infected_ratio: 10 } }); + expect(latest()?.navigation).toEqual({ + positions: { infected_ratio: 10 }, + booleans: {}, + followTrials: false, + }); + expect(refinementRuns.runs[0]?.request).toMatchObject({ + cacheKey: "optimization-1", + computeBackend: "webgpu", + seed: 1, + runCount: 8, + scenarioParameterValues: { + population: 1_000, + infected_ratio: optimizationAxisValueAt(axis, 10), + }, + }); + expect(latest()?.selection).toMatchObject({ + key: "infected_ratio=10", + runTarget: 8, + computing: true, + }); + expect(latest()?.activity).toEqual([ + expect.objectContaining({ kind: "step", label: "Step 1", runCount: 3 }), + expect.objectContaining({ + kind: "refine", + label: `Refining infected_ratio ${optimizationAxisValueAt(axis, 10) + .toPrecision(3) + .replace(/\.?0+$/, "")}`, + runCount: 8, + }), + ]); + + // Later trials no longer move the navigation or replace the selection. + startTrial(1, 0.02); + expect(latest()?.navigation?.positions).toEqual({ infected_ratio: 10 }); + expect(latest()?.selection?.key).toBe("infected_ratio=10"); + }); + + it("settles on the best trial's point and refines it there, once the followed trial has settled", () => { + const { study, startTrial, latest, refinementRuns } = setup(); + study.trialReported(trialEvent(0, 0.05, 0.3)); + study.trialReported(trialEvent(1, 0.02, 0.1)); + const trial = startTrial(2, 0.15); + + study.settle("complete"); + expect(refinementRuns.runs).toHaveLength(0); + + const failed = failedRunOutcome("1 of 3 runs failed"); + study.trialSettled(2, failed); + trial.settle(failed); + const bestPosition = optimizationAxisPositionFor(axis, 0.02); + expect(latest()?.navigation).toEqual({ + positions: { infected_ratio: bestPosition }, + booleans: {}, + followTrials: false, + }); + expect(refinementRuns.runs[0]?.request).toMatchObject({ + scenarioParameterValues: { + infected_ratio: optimizationAxisValueAt(axis, bestPosition), + }, + }); + expect(latest()?.selection?.key).toBe(`infected_ratio=${bestPosition}`); + }); + + it("takes the best the terminal event carries, and stays at the midpoint without any", () => { + const { study, latest, refinementRuns } = setup(); + + study.settle("complete", { + trial: 4, + parameters: { infected_ratio: 0.01 }, + objective: 0.05, + }); + const bestPosition = optimizationAxisPositionFor(axis, 0.01); + expect(latest()?.navigation?.positions).toEqual({ + infected_ratio: bestPosition, + }); + expect(refinementRuns.runs[0]?.request.scenarioParameterValues).toEqual({ + population: 1_000, + infected_ratio: optimizationAxisValueAt(axis, bestPosition), + }); + + const empty = setup(); + empty.study.settle("cancelled"); + expect(empty.latest()?.navigation).toEqual({ + positions: { infected_ratio: 25 }, + booleans: {}, + followTrials: false, + }); + expect(empty.refinementRuns.runs).toHaveLength(1); + }); + + it("a stop settles on the best too; a navigation the user moved earlier stays where it is", () => { + const { study, startTrial, latest, refinementRuns } = setup(); + study.trialReported(trialEvent(0, 0.05, 0.3)); + const trial = startTrial(1, 0.02); + const frame = distributionFrame(metricId, 1, [[0.2, 1]]); + trial.frames.set([frame]); + + study.settle("cancelled"); + trial.run.cancel(); + study.trialSettled(1, cancelledRunOutcome); + const bestPosition = optimizationAxisPositionFor(axis, 0.05); + expect(latest()?.navigation?.positions).toEqual({ + infected_ratio: bestPosition, + }); + expect(refinementRuns.runs).toHaveLength(1); + expect(latest()?.selection?.key).toBe(`infected_ratio=${bestPosition}`); + + const moved = setup(); + moved.startTrial(0, 0.05); + moved.study.setNavigation({ positions: { infected_ratio: 10 } }); + moved.study.settle("complete", { + trial: 0, + parameters: { infected_ratio: 0.05 }, + objective: 0.3, + }); + expect(moved.latest()?.navigation?.positions).toEqual({ + infected_ratio: 10, + }); + expect(moved.refinementRuns.runs).toHaveLength(1); + }); + + it("turning following back on attaches to the trial being evaluated", () => { + const { study, startTrial, latest, refinementRuns } = setup(); + study.setNavigation({ positions: { infected_ratio: 10 } }); + startTrial(1, 0.02); + expect(latest()?.selection?.key).toBe("infected_ratio=10"); + + study.setNavigation({ followTrials: true }); + expect(refinementRuns.runs[0]!.cancelled).toBe(true); + expect(latest()?.navigation).toEqual({ + positions: { infected_ratio: optimizationAxisPositionFor(axis, 0.02) }, + booleans: {}, + followTrials: true, + }); + expect(latest()?.selection?.key).toBe("trial:1"); + }); + + it("resuming a settled study stops the refinement and follows the next trial", () => { + const { study, startTrial, latest, refinementRuns } = setup(); + study.settle("complete", { + trial: 0, + parameters: { infected_ratio: 0.05 }, + objective: 0.3, + }); + expect(refinementRuns.runs).toHaveLength(1); + + study.resume(); + expect(refinementRuns.runs[0]!.cancelled).toBe(true); + expect(latest()?.navigation?.followTrials).toBe(true); + + startTrial(1, 0.02); + expect(latest()?.selection?.key).toBe("trial:1"); + expect(latest()?.navigation?.positions).toEqual({ + infected_ratio: optimizationAxisPositionFor(axis, 0.02), + }); + }); + + it("dispose cancels the refinement, clears the activity and publishes nothing further", () => { + const { study, latest, refinementRuns, updates } = setup(); + study.setNavigation({ positions: { infected_ratio: 3 } }); + const published = updates.length; + + study.dispose(); + expect(refinementRuns.runs[0]!.cancelled).toBe(true); + study.setNavigation({ positions: { infected_ratio: 4 } }); + expect(updates).toHaveLength(published); + expect(latest()?.navigation?.positions).toEqual({ infected_ratio: 3 }); + }); +}); diff --git a/libs/@hashintel/petrinaut/src/react/optimizations/provider/connected-study.ts b/libs/@hashintel/petrinaut/src/react/optimizations/provider/connected-study.ts new file mode 100644 index 00000000000..4333138ba6d --- /dev/null +++ b/libs/@hashintel/petrinaut/src/react/optimizations/provider/connected-study.ts @@ -0,0 +1,485 @@ +import { sweepCellObjective } from "../../experiments/sweep-cell-objective"; +import { + optimizationAxisMidpoint, + optimizationAxisPositionFor, + optimizationBooleanIdentifiers, + optimizationNavigationKey, + optimizationNavigationValues, +} from "../surface-grid"; +import { createActivityRegistry } from "./activity-registry"; +import { createPointRefinement } from "./point-refinement"; + +import type { + DetachedObjectiveRun, + DetachedObjectiveRunOutcome, + ExperimentComputeBackend, + ExperimentsActionsValue, +} from "../../experiments/context"; +import type { + OptimizationBatchStatus, + OptimizationBest, + OptimizationInFlightStep, + OptimizationNavigation, + OptimizationRecord, + OptimizationSelectionStream, + OptimizationStatus, +} from "../context"; +import type { OptimizationSurfaceAxis } from "../surface-grid"; +import type { + PetrinautOptimizationInput, + PetrinautOptimizationTrialEvent, +} from "@hashintel/petrinaut-core"; +import type { OptimizationScalar } from "@hashintel/petrinaut-core/optimization"; + +/** What a connected study publishes into its record. */ +export type ConnectedStudyUpdate = Pick< + OptimizationRecord, + "navigation" | "selection" | "activity" | "inFlight" +>; + +/** The status a study settles with. */ +export type ConnectedStudyOutcome = Extract< + OptimizationStatus, + "complete" | "error" | "cancelled" +>; + +/** A trial as the channel reports it: the optimizer's values and the batch evaluating them. */ +type EvaluatingTrial = { + trial: number; + values: Readonly>; + run: DetachedObjectiveRun; + /** Stops listing the trial in the activity and watching its frames. */ + release: () => void; +}; + +export type ConnectedStudy = { + readonly computeBackend: ExperimentComputeBackend; + /** The navigation at creation, for the record's first render. */ + readonly initialNavigation: OptimizationNavigation; + setNavigation(this: void, patch: Partial): void; + /** + * A trial began evaluating. While following, the navigation moves to the + * trial and its stream becomes the selection; with several trials in + * flight the most recently started one is followed. + */ + trialStarted( + this: void, + trial: number, + values: Readonly>, + run: DetachedObjectiveRun, + runCount: number, + ): void; + /** + * The trial's batch settled; a followed trial's selection stops computing + * and, when the batch failed, carries its reason. + */ + trialSettled( + this: void, + trial: number, + outcome: DetachedObjectiveRunOutcome, + ): void; + /** A trial event landed on the record; the study keeps the best from it. */ + trialReported(this: void, event: PetrinautOptimizationTrialEvent): void; + /** + * The study reached a terminal status, `best` overriding the best kept from + * the trials when given. While following, the navigation settles on the + * best trial's point, following ends, and the point refines; a navigation + * the user moved earlier stays where it is. + */ + settle( + this: void, + outcome: ConnectedStudyOutcome, + best?: OptimizationBest | null, + ): void; + /** + * More steps were asked of a settled study: following turns back on so the + * next step is followed, and the point refining stops. + */ + resume(this: void): void; + dispose(this: void): void; +}; + +const labelValue = new Intl.NumberFormat("en-US", { + maximumSignificantDigits: 3, +}); + +/** + * The local machinery behind one connected study: where its drawer points, + * whether that follows the trials as they are evaluated, the objective's + * live stream there — the followed trial's batch while following, the point + * refinement ladder once the study is terminal or the user has moved away — + * and the list of every batch computing for it. + */ +export const createConnectedStudy = ({ + optimizationId, + input, + axes, + computeBackend, + runDetachedObjective, + onUpdate, +}: { + optimizationId: string; + input: PetrinautOptimizationInput; + axes: readonly OptimizationSurfaceAxis[]; + computeBackend: ExperimentComputeBackend; + runDetachedObjective: ExperimentsActionsValue["runDetachedObjective"]; + onUpdate: (update: ConnectedStudyUpdate) => void; +}): ConnectedStudy => { + const booleanIdentifiers = optimizationBooleanIdentifiers(input); + const optimizedIdentifiers = [ + ...axes.map((axis) => axis.identifier), + ...booleanIdentifiers, + ]; + const { direction } = input.objective; + const scenario = input.model.definition.scenarios?.find( + (candidate) => candidate.id === input.scenario.id, + ); + const metric = input.model.definition.metrics?.find( + (candidate) => candidate.id === input.objective.metricId, + ); + if (!metric) { + throw new Error( + `The study has no metric "${input.objective.metricId}" to optimize`, + ); + } + + let navigation: OptimizationNavigation = { + positions: Object.fromEntries( + axes.map((axis) => [axis.identifier, optimizationAxisMidpoint(axis)]), + ), + booleans: Object.fromEntries( + booleanIdentifiers.map((identifier) => [ + identifier, + (scenario?.scenarioParameters.find( + (parameter) => parameter.identifier === identifier, + )?.default ?? 0) !== 0, + ]), + ), + followTrials: true, + }; + let selection: OptimizationSelectionStream | null = null; + let activity: readonly OptimizationBatchStatus[] = []; + let best: OptimizationBest | null = null; + let terminal: ConnectedStudyOutcome | null = null; + let disposed = false; + /** Trials being evaluated, in the order they started. */ + const evaluating = new Map(); + let followed: { trial: number; off: () => void } | null = null; + + const inFlight = (): readonly OptimizationInFlightStep[] => + [...evaluating.values()].map((entry) => ({ + trial: entry.trial, + parameters: entry.values, + objective: sweepCellObjective(entry.run.frames.get(), metric.id), + })); + + const publish = () => { + if (!disposed) { + onUpdate({ navigation, selection, activity, inFlight: inFlight() }); + } + }; + + const activityRegistry = createActivityRegistry((next) => { + activity = next; + publish(); + }); + + /** The navigation at a trial's values; unset axes keep their position. */ + const navigationAt = ( + values: Readonly>, + followTrials: boolean, + ): OptimizationNavigation => ({ + positions: Object.fromEntries( + axes.map((axis) => { + const value = values[axis.identifier]; + return [ + axis.identifier, + typeof value === "number" + ? optimizationAxisPositionFor(axis, value) + : (navigation.positions[axis.identifier] ?? + optimizationAxisMidpoint(axis)), + ]; + }), + ), + booleans: Object.fromEntries( + booleanIdentifiers.map((identifier) => { + const value = values[identifier]; + return [ + identifier, + typeof value === "boolean" + ? value + : (navigation.booleans[identifier] ?? false), + ]; + }), + ), + followTrials, + }); + + const keyOf = (target: OptimizationNavigation): string => + optimizationNavigationKey(axes, booleanIdentifiers, target); + + const refineLabel = ( + values: Readonly>, + ): string => + `Refining ${optimizedIdentifiers + .map((identifier) => { + const value = values[identifier]; + return `${identifier} ${ + typeof value === "number" ? labelValue.format(value) : String(value) + }`; + }) + .join(" · ")}`; + + const refinement = createPointRefinement({ + runDetachedObjective: (request) => { + const run = runDetachedObjective(request); + const off = activityRegistry.register({ + kind: "refine", + label: refineLabel(request.scenarioParameterValues), + runCount: request.runCount, + progress: run.progress, + }); + void run.completion.then(off, off); + return run; + }, + study: { + cacheKey: optimizationId, + definition: input.model.definition, + scenarioId: input.scenario.id, + metric: { id: metric.id, label: metric.name, code: metric.code }, + seed: input.execution.seed, + dt: input.execution.dt, + maxTime: input.execution.maxTime, + computeBackend, + direction, + }, + bestObjective: () => best?.objective ?? null, + onUpdate: (next) => { + selection = next; + publish(); + }, + }); + + const refineHere = () => { + const key = keyOf(navigation); + refinement.refine({ + key, + scenarioParameterValues: optimizationNavigationValues( + input, + axes, + booleanIdentifiers, + navigation, + ), + isBest: + best !== null && key === keyOf(navigationAt(best.parameters, false)), + }); + }; + + const stopFollowing = () => { + followed?.off(); + followed = null; + }; + + const follow = ({ trial, values, run }: EvaluatingTrial) => { + stopFollowing(); + navigation = navigationAt(values, true); + const key = `trial:${trial}`; + const mirror = () => { + selection = { + key, + metricFrames: run.frames.get(), + runsCompleted: run.progress.get()?.completedRuns ?? 0, + runTarget: null, + computing: true, + error: null, + note: null, + }; + publish(); + }; + const offFrames = run.frames.subscribe(mirror); + const offProgress = run.progress.subscribe(mirror); + followed = { + trial, + off: () => { + offFrames(); + offProgress(); + }, + }; + mirror(); + }; + + const mostRecentlyStarted = (): EvaluatingTrial | undefined => + [...evaluating.values()].at(-1); + + /** Following ends where the study did best, and that point refines. */ + const settleOnBest = () => { + stopFollowing(); + navigation = best + ? navigationAt(best.parameters, false) + : { ...navigation, followTrials: false }; + refineHere(); + publish(); + }; + + const foldBest = (event: PetrinautOptimizationTrialEvent) => { + if (event.best) { + best = event.best; + return; + } + if (event.state !== "complete" || event.objective === null) { + return; + } + const isBetter = + best === null || + (direction === "maximize" + ? event.objective > best.objective + : event.objective < best.objective); + if (isBetter) { + best = { + trial: event.trial, + parameters: event.parameters, + objective: event.objective, + }; + } + }; + + return { + computeBackend, + initialNavigation: navigation, + setNavigation: (patch) => { + if (disposed) { + return; + } + const moved = + patch.positions !== undefined || patch.booleans !== undefined; + navigation = { + positions: { ...navigation.positions, ...patch.positions }, + booleans: { ...navigation.booleans, ...patch.booleans }, + followTrials: + patch.followTrials ?? (moved ? false : navigation.followTrials), + }; + if (terminal !== null || !navigation.followTrials) { + stopFollowing(); + refineHere(); + } else { + refinement.stop(); + const latest = mostRecentlyStarted(); + if (latest) { + follow(latest); + } + } + publish(); + }, + trialStarted: (trial, values, run, runCount) => { + if (disposed) { + return; + } + const offActivity = activityRegistry.register({ + kind: "step", + label: `Step ${trial + 1}`, + runCount, + progress: run.progress, + }); + // The followed trial's own mirror publishes its frames. + const offFrames = run.frames.subscribe(() => { + if (followed?.trial !== trial) { + publish(); + } + }); + const entry: EvaluatingTrial = { + trial, + values, + run, + release: () => { + offActivity(); + offFrames(); + }, + }; + evaluating.set(trial, entry); + if (terminal !== null || !navigation.followTrials) { + publish(); + return; + } + refinement.stop(); + follow(entry); + }, + trialSettled: (trial, outcome) => { + if (disposed) { + return; + } + const entry = evaluating.get(trial); + entry?.release(); + evaluating.delete(trial); + if (followed?.trial !== trial) { + publish(); + return; + } + stopFollowing(); + selection = outcome.ok + ? { + key: `trial:${trial}`, + metricFrames: outcome.metricFrames, + runsCompleted: outcome.runsCompleted, + runTarget: null, + computing: false, + error: null, + note: null, + } + : { + key: `trial:${trial}`, + metricFrames: selection?.metricFrames ?? [], + runsCompleted: selection?.runsCompleted ?? 0, + runTarget: null, + computing: false, + error: outcome.cancelled ? null : outcome.reason, + note: null, + }; + if (terminal !== null) { + settleOnBest(); + return; + } + const latest = mostRecentlyStarted(); + if (latest) { + follow(latest); + return; + } + publish(); + }, + trialReported: (event) => { + if (!disposed) { + foldBest(event); + } + }, + settle: (outcome, settledBest) => { + if (disposed || terminal !== null) { + return; + } + terminal = outcome; + if (settledBest !== undefined && settledBest !== null) { + best = settledBest; + } + if (navigation.followTrials && !followed) { + settleOnBest(); + } + }, + resume: () => { + if (disposed || terminal === null) { + return; + } + terminal = null; + refinement.stop(); + navigation = { ...navigation, followTrials: true }; + publish(); + }, + dispose: () => { + disposed = true; + stopFollowing(); + for (const entry of evaluating.values()) { + entry.release(); + } + evaluating.clear(); + activityRegistry.clear(); + refinement.dispose(); + }, + }; +}; diff --git a/libs/@hashintel/petrinaut/src/react/optimizations/provider/point-refinement.test.ts b/libs/@hashintel/petrinaut/src/react/optimizations/provider/point-refinement.test.ts new file mode 100644 index 00000000000..347eb0fd18f --- /dev/null +++ b/libs/@hashintel/petrinaut/src/react/optimizations/provider/point-refinement.test.ts @@ -0,0 +1,350 @@ +import { describe, expect, it } from "vitest"; + +import { deriveRunSeed } from "@hashintel/petrinaut-core"; + +import { distributionStats } from "../../experiments/distribution-stats"; +import { + completedRunResult, + createFakeDetachedObjectiveRuns, + distributionFrame, + failedRunOutcome, +} from "../fake-detached-objective-runs.fixtures"; +import { + sirOptimizationInput, + sirOptimizationMetric, +} from "../sir-optimization-input.fixtures"; +import { + cannotBeatBestNote, + createPointRefinement, + type PointRefinementStudy, +} from "./point-refinement"; + +import type { OptimizationSelectionStream } from "../context"; + +const metricId = sirOptimizationMetric.id; + +const study: PointRefinementStudy = { + cacheKey: "study", + definition: sirOptimizationInput.model.definition, + scenarioId: sirOptimizationInput.scenario.id, + metric: { + id: metricId, + label: sirOptimizationMetric.name, + code: sirOptimizationMetric.code, + }, + seed: 42, + dt: 1, + maxTime: 180, + computeBackend: "cpu", + direction: "minimize", +}; + +const target = (key: string, infectedRatio: number, isBest = false) => ({ + key, + scenarioParameterValues: { population: 1_000, infected_ratio: infectedRatio }, + isBest, +}); + +const setup = ({ + maxRuns = 25, + best = null, +}: { maxRuns?: number; best?: number | null } = {}) => { + const fake = createFakeDetachedObjectiveRuns(); + const updates: OptimizationSelectionStream[] = []; + const refinement = createPointRefinement({ + runDetachedObjective: fake.runDetachedObjective, + study, + bestObjective: () => best, + maxRuns, + onUpdate: (update) => { + updates.push(update); + }, + }); + return { fake, updates, refinement, latest: () => updates.at(-1) }; +}; + +const settled = async () => { + await new Promise((resolve) => { + setTimeout(resolve, 0); + }); +}; + +/** Eight runs whose values sit `spread` either side of `mean`. */ +const spreadFrame = (mean: number, spread = 0.01) => + distributionFrame(metricId, 180, [ + [mean - spread, 4], + [mean + spread, 4], + ]); + +describe("createPointRefinement", () => { + it("climbs the ladder from the point's first rung, seeding each batch from its first run index", async () => { + const { fake, refinement, latest } = setup(); + + refinement.refine(target("a", 0.05)); + expect(latest()).toEqual({ + key: "a", + metricFrames: [], + runsCompleted: 0, + runTarget: 8, + computing: true, + error: null, + note: null, + }); + expect(fake.runs[0]?.request).toMatchObject({ + cacheKey: "study", + seed: 42, + runCount: 8, + computeBackend: "cpu", + scenarioParameterValues: { population: 1_000, infected_ratio: 0.05 }, + }); + + const first = distributionFrame(metricId, 1, [[0.1, 8]]); + fake.runs[0]!.settle( + completedRunResult({ metricId, frames: [first], runsCompleted: 8 }), + ); + await settled(); + expect(latest()).toEqual({ + key: "a", + metricFrames: [first], + runsCompleted: 8, + runTarget: 25, + computing: true, + error: null, + note: null, + }); + expect(fake.runs[1]?.request).toMatchObject({ + seed: deriveRunSeed(42, 8), + runCount: 17, + }); + + // The in-flight batch streams merged with the finished rungs. + const second = distributionFrame(metricId, 1, [[0.3, 17]]); + fake.runs[1]!.frames.set([second]); + expect(latest()).toMatchObject({ + runsCompleted: 8, + runTarget: 25, + computing: true, + }); + expect(distributionStats(latest()!.metricFrames, metricId)).toMatchObject({ + runs: 25, + mean: (0.1 * 8 + 0.3 * 17) / 25, + }); + + fake.runs[1]!.settle( + completedRunResult({ metricId, frames: [second], runsCompleted: 17 }), + ); + await settled(); + expect(latest()).toMatchObject({ + runsCompleted: 25, + runTarget: null, + computing: false, + note: null, + }); + expect(fake.runs).toHaveLength(2); + }); + + it("a new key cancels the batch in flight, and a refined key resumes from its cached rungs", async () => { + const { fake, refinement, latest } = setup(); + + refinement.refine(target("a", 0.05)); + fake.runs[0]!.settle( + completedRunResult({ + metricId, + frames: [distributionFrame(metricId, 1, [[0.1, 8]])], + runsCompleted: 8, + }), + ); + await settled(); + expect(fake.runs).toHaveLength(2); + + refinement.refine(target("b", 0.01)); + expect(fake.runs[1]!.cancelled).toBe(true); + expect(fake.runs[2]?.request).toMatchObject({ + seed: 42, + runCount: 8, + scenarioParameterValues: { infected_ratio: 0.01 }, + }); + expect(latest()).toMatchObject({ + key: "b", + runsCompleted: 0, + runTarget: 8, + }); + + refinement.refine(target("a", 0.05)); + expect(fake.runs[2]!.cancelled).toBe(true); + expect(latest()).toMatchObject({ + key: "a", + runsCompleted: 8, + runTarget: 25, + }); + expect(fake.runs[3]?.request).toMatchObject({ + seed: deriveRunSeed(42, 8), + runCount: 17, + scenarioParameterValues: { infected_ratio: 0.05 }, + }); + }); + + it("refining the active key again changes nothing; stop cancels and keeps the cache", async () => { + const { fake, refinement, latest } = setup(); + + refinement.refine(target("a", 0.05)); + refinement.refine(target("a", 0.05)); + expect(fake.runs).toHaveLength(1); + + fake.runs[0]!.settle( + completedRunResult({ + metricId, + frames: [distributionFrame(metricId, 1, [[0.1, 8]])], + runsCompleted: 8, + }), + ); + await settled(); + refinement.stop(); + expect(fake.runs[1]!.cancelled).toBe(true); + + refinement.refine(target("a", 0.05)); + expect(latest()).toMatchObject({ + key: "a", + runsCompleted: 8, + runTarget: 25, + }); + expect(fake.runs[2]?.request).toMatchObject({ runCount: 17 }); + }); + + it("a failed rung stops the ladder with its reason, and refining the key again retries it", async () => { + const { fake, refinement, latest } = setup(); + + refinement.refine(target("a", 0.05)); + fake.runs[0]!.settle(failedRunOutcome("cpu: unsupported net")); + await settled(); + expect(latest()).toEqual({ + key: "a", + metricFrames: [], + runsCompleted: 0, + runTarget: null, + computing: false, + error: "cpu: unsupported net", + note: null, + }); + expect(fake.runs).toHaveLength(1); + + refinement.refine(target("a", 0.05)); + expect(fake.runs).toHaveLength(2); + expect(latest()).toMatchObject({ + key: "a", + runTarget: 8, + computing: true, + error: null, + note: null, + }); + }); + + it("a batch cancelled from beneath stops the ladder without an error", async () => { + const { fake, refinement, latest } = setup(); + + refinement.refine(target("a", 0.05)); + fake.runs[0]!.run.cancel(); + await settled(); + expect(latest()).toEqual({ + key: "a", + metricFrames: [], + runsCompleted: 0, + runTarget: null, + computing: false, + error: null, + note: null, + }); + expect(fake.runs).toHaveLength(1); + }); + + it("stops after the first rung, with a note, at a point that cannot beat the best", async () => { + // Minimizing, with a best of 0.1: a point around 0.3 is hopeless. + const { fake, refinement, latest } = setup({ best: 0.1 }); + + refinement.refine(target("a", 0.05)); + fake.runs[0]!.settle( + completedRunResult({ + metricId, + frames: [spreadFrame(0.3)], + runsCompleted: 8, + }), + ); + await settled(); + + expect(latest()).toMatchObject({ + key: "a", + runsCompleted: 8, + runTarget: null, + computing: false, + error: null, + note: cannotBeatBestNote(8), + }); + expect(latest()?.note).toBe("8 runs · cannot beat the best"); + expect(fake.runs).toHaveLength(1); + + // Returning to the point later changes nothing: the verdict stands. + refinement.refine(target("b", 0.01)); + refinement.refine(target("a", 0.05)); + expect(fake.runs).toHaveLength(2); + expect(latest()).toMatchObject({ key: "a", note: cannotBeatBestNote(8) }); + }); + + it("keeps climbing at a point that might beat the best, and at the best trial's own point", async () => { + const { fake, refinement, latest } = setup({ best: 0.1 }); + + // Within reach: a mean of 0.11 whose eight runs spread 0.05 either side, + // so 2.5 standard errors reach below the best. + refinement.refine(target("a", 0.05)); + fake.runs[0]!.settle( + completedRunResult({ + metricId, + frames: [spreadFrame(0.11, 0.05)], + runsCompleted: 8, + }), + ); + await settled(); + expect(latest()).toMatchObject({ + runTarget: 25, + computing: true, + note: null, + }); + expect(fake.runs).toHaveLength(2); + + // The best trial's point: hopeless by its estimate, refined regardless. + refinement.refine(target("best", 0.02, true)); + fake.runs[2]!.settle( + completedRunResult({ + metricId, + frames: [spreadFrame(0.3)], + runsCompleted: 8, + }), + ); + await settled(); + expect(latest()).toMatchObject({ + key: "best", + runTarget: 25, + computing: true, + note: null, + }); + expect(fake.runs).toHaveLength(4); + }); + + it("refines as before while the study has no best", async () => { + const { fake, refinement, latest } = setup(); + + refinement.refine(target("a", 0.05)); + fake.runs[0]!.settle( + completedRunResult({ + metricId, + frames: [spreadFrame(0.3)], + runsCompleted: 8, + }), + ); + await settled(); + expect(latest()).toMatchObject({ + runTarget: 25, + computing: true, + note: null, + }); + }); +}); diff --git a/libs/@hashintel/petrinaut/src/react/optimizations/provider/point-refinement.ts b/libs/@hashintel/petrinaut/src/react/optimizations/provider/point-refinement.ts new file mode 100644 index 00000000000..eb1479c4b93 --- /dev/null +++ b/libs/@hashintel/petrinaut/src/react/optimizations/provider/point-refinement.ts @@ -0,0 +1,232 @@ +import { + getNextRunTarget, + mergeMetricFramesAcrossCells, +} from "../../experiments/parameter-grid"; +import { sweepBatchSeed } from "../../experiments/sweep-session"; +import { + estimateObjective, + shouldStopRefining, +} from "./point-refinement/objective-estimate"; + +import type { + DetachedObjectiveRun, + DetachedObjectiveRunRequest, + ExperimentsActionsValue, +} from "../../experiments/context"; +import type { SweepCellSnapshot } from "../../experiments/sweep-session"; +import type { OptimizationSelectionStream } from "../context"; +import type { MonteCarloUserDefinedMetricFrame } from "@hashintel/petrinaut-core"; + +export { shouldStopRefining } from "./point-refinement/objective-estimate"; + +/** The most runs the selected point is refined to. */ +export const POINT_REFINEMENT_MAX_RUNS = 100; + +/** The study fields every refinement batch shares. */ +export type PointRefinementStudy = Pick< + DetachedObjectiveRunRequest, + | "cacheKey" + | "definition" + | "scenarioId" + | "metric" + | "dt" + | "maxTime" + | "computeBackend" +> & { seed: number; direction: "maximize" | "minimize" }; + +export type PointRefinementTarget = { + key: string; + scenarioParameterValues: DetachedObjectiveRunRequest["scenarioParameterValues"]; + /** + * The point is the best trial's: it climbs to the top rung whatever its + * estimate says, since it is the value the study reports. + */ + isBest: boolean; +}; + +export type PointRefinement = { + /** + * Climbs the run ladder at `target`, streaming into `onUpdate`. A new key + * cancels the batch in flight and resumes from the key's cached rungs; the + * key already refining, or settled, changes nothing. A failed rung stops + * the ladder and records the reason; refining the key again retries it. + * Between rungs, a point whose mean sits too far from the study's best to + * ever beat it stops with a note saying so. + */ + refine(this: void, target: PointRefinementTarget): void; + /** Cancels the batch in flight. Finished rungs stay cached. */ + stop(this: void): void; + dispose(this: void): void; +}; + +/** The point being refined, and how to stop it. */ +type RefinementSession = { + key: string; + cancel: () => void; +}; + +const mergeFrames = ( + base: readonly MonteCarloUserDefinedMetricFrame[], + streamed: readonly MonteCarloUserDefinedMetricFrame[], +): readonly MonteCarloUserDefinedMetricFrame[] => + base.length === 0 ? streamed : mergeMetricFramesAcrossCells([base, streamed]); + +/** The note a ladder stops with when the point cannot beat the best. */ +export const cannotBeatBestNote = (runs: number): string => + `${runs} runs · cannot beat the best`; + +/** + * Refines one parameter point of a study, as the sweep session refines the + * navigator's selection: cumulative batches up the run ladder, each batch + * seeded from its first global run index so a rung repeats exactly, merged + * into a cache keyed by the point. + */ +export const createPointRefinement = ({ + runDetachedObjective, + study, + bestObjective, + maxRuns = POINT_REFINEMENT_MAX_RUNS, + onUpdate, +}: { + runDetachedObjective: ExperimentsActionsValue["runDetachedObjective"]; + study: PointRefinementStudy; + /** The study's best objective so far, read before each rung. */ + bestObjective: () => number | null; + maxRuns?: number; + onUpdate: (selection: OptimizationSelectionStream) => void; +}): PointRefinement => { + const cache = new Map(); + let active: RefinementSession | null = null; + + const stop = () => { + active?.cancel(); + active = null; + }; + + /** Whether a point's finished rungs already rule it out against the best. */ + const cannotBeatBest = (snapshot: SweepCellSnapshot): boolean => { + if (snapshot.runsCompleted === 0) { + return false; + } + const estimate = estimateObjective(snapshot.metricFrames, study.metric.id); + return ( + estimate !== null && + shouldStopRefining({ + direction: study.direction, + best: bestObjective(), + mean: estimate.mean, + standardError: estimate.standardError, + }) + ); + }; + + const refine = (target: PointRefinementTarget) => { + if (active?.key === target.key) { + return; + } + stop(); + let cancelled = false; + let inFlight: DetachedObjectiveRun | null = null; + // Read through a call so the flag is re-checked after each await (a plain + // property read would be control-flow-narrowed to `false`). + const isCancelled = () => cancelled; + active = { + key: target.key, + cancel: () => { + cancelled = true; + inFlight?.cancel(); + }, + }; + + const publish = ( + snapshot: SweepCellSnapshot, + runTarget: number | null, + note: string | null, + ) => { + onUpdate({ + key: target.key, + metricFrames: snapshot.metricFrames, + runsCompleted: snapshot.runsCompleted, + runTarget, + computing: runTarget !== null, + error: null, + note, + }); + }; + + const climb = async (): Promise => { + let snapshot: SweepCellSnapshot = cache.get(target.key) ?? { + runsCompleted: 0, + metricFrames: [], + }; + + while (!isCancelled()) { + const runTarget = getNextRunTarget(snapshot.runsCompleted, maxRuns); + if (runTarget === null) { + publish(snapshot, null, null); + return; + } + if (!target.isBest && cannotBeatBest(snapshot)) { + publish(snapshot, null, cannotBeatBestNote(snapshot.runsCompleted)); + return; + } + publish(snapshot, runTarget, null); + + const base = snapshot; + const run = runDetachedObjective({ + ...study, + scenarioParameterValues: target.scenarioParameterValues, + seed: sweepBatchSeed(study.seed, base.runsCompleted), + runCount: runTarget - base.runsCompleted, + }); + inFlight = run; + const offFrames = run.frames.subscribe((frames) => { + if (!isCancelled()) { + publish( + { + runsCompleted: base.runsCompleted, + metricFrames: mergeFrames(base.metricFrames, frames), + }, + runTarget, + null, + ); + } + }); + const outcome = await run.completion; + offFrames(); + inFlight = null; + if (isCancelled()) { + return; + } + if (!outcome.ok) { + active = null; + onUpdate({ + key: target.key, + metricFrames: base.metricFrames, + runsCompleted: base.runsCompleted, + runTarget: null, + computing: false, + error: outcome.cancelled ? null : outcome.reason, + note: null, + }); + return; + } + snapshot = { + runsCompleted: base.runsCompleted + outcome.runsCompleted, + metricFrames: mergeFrames(base.metricFrames, outcome.metricFrames), + }; + cache.set(target.key, snapshot); + } + }; + void climb(); + }; + + return { + refine, + stop, + dispose: () => { + stop(); + cache.clear(); + }, + }; +}; diff --git a/libs/@hashintel/petrinaut/src/react/optimizations/provider/point-refinement/objective-estimate.test.ts b/libs/@hashintel/petrinaut/src/react/optimizations/provider/point-refinement/objective-estimate.test.ts new file mode 100644 index 00000000000..418fb00737d --- /dev/null +++ b/libs/@hashintel/petrinaut/src/react/optimizations/provider/point-refinement/objective-estimate.test.ts @@ -0,0 +1,122 @@ +import { describe, expect, it } from "vitest"; + +import { distributionFrame } from "../../fake-detached-objective-runs.fixtures"; +import { estimateObjective, shouldStopRefining } from "./objective-estimate"; + +const metricId = "metric"; + +describe("estimateObjective", () => { + it("reads the mean and its standard error off the last frame with samples", () => { + const estimate = estimateObjective( + [ + distributionFrame(metricId, 1, [[0.9, 4]]), + distributionFrame(metricId, 2, [ + [0.1, 2], + [0.3, 2], + ]), + distributionFrame(metricId, 3, []), + ], + metricId, + ); + + // Values 0.1, 0.1, 0.3, 0.3: mean 0.2, sample variance 0.04/3. + expect(estimate).toEqual({ + runs: 4, + mean: expect.closeTo(0.2, 12) as number, + standardError: expect.closeTo(Math.sqrt(0.04 / 3 / 4), 12) as number, + }); + }); + + it("leaves the error unbounded with one run, and estimates nothing without a distribution", () => { + expect( + estimateObjective([distributionFrame(metricId, 1, [[0.5, 1]])], metricId), + ).toEqual({ runs: 1, mean: 0.5, standardError: Number.POSITIVE_INFINITY }); + expect(estimateObjective([], metricId)).toBeNull(); + expect( + estimateObjective([distributionFrame("other", 1, [[0.5, 3]])], metricId), + ).toBeNull(); + }); +}); + +describe("shouldStopRefining", () => { + it("stops a maximized point whose mean plus 2.5 errors falls short of the best, and not at the boundary", () => { + expect( + shouldStopRefining({ + direction: "maximize", + best: 10, + mean: 7, + standardError: 1, + }), + ).toBe(true); + expect( + shouldStopRefining({ + direction: "maximize", + best: 10, + mean: 7.5, + standardError: 1, + }), + ).toBe(false); + expect( + shouldStopRefining({ + direction: "maximize", + best: 10, + mean: 7.4, + standardError: 1, + }), + ).toBe(true); + expect( + shouldStopRefining({ + direction: "maximize", + best: 10, + mean: 12, + standardError: 1, + }), + ).toBe(false); + }); + + it("stops a minimized point whose mean minus 2.5 errors exceeds the best", () => { + expect( + shouldStopRefining({ + direction: "minimize", + best: 0.1, + mean: 0.4, + standardError: 0.1, + }), + ).toBe(true); + expect( + shouldStopRefining({ + direction: "minimize", + best: 0.1, + mean: 0.35, + standardError: 0.1, + }), + ).toBe(false); + expect( + shouldStopRefining({ + direction: "minimize", + best: 0.1, + mean: 0.05, + standardError: 0.1, + }), + ).toBe(false); + }); + + it("never stops without a best, or while a single run leaves the error unbounded", () => { + expect( + shouldStopRefining({ + direction: "maximize", + best: null, + mean: 0, + standardError: 0, + }), + ).toBe(false); + expect( + shouldStopRefining({ + direction: "maximize", + best: 10, + mean: 0, + standardError: Number.POSITIVE_INFINITY, + }), + ).toBe(false); + }); +}); diff --git a/libs/@hashintel/petrinaut/src/react/optimizations/provider/point-refinement/objective-estimate.ts b/libs/@hashintel/petrinaut/src/react/optimizations/provider/point-refinement/objective-estimate.ts new file mode 100644 index 00000000000..880defe2075 --- /dev/null +++ b/libs/@hashintel/petrinaut/src/react/optimizations/provider/point-refinement/objective-estimate.ts @@ -0,0 +1,79 @@ +import type { MonteCarloUserDefinedMetricFrame } from "@hashintel/petrinaut-core"; + +/** The objective's mean over a point's runs and how sure that mean is. */ +export type ObjectiveEstimate = { + runs: number; + mean: number; + /** Standard error of the mean; infinite with a single run. */ + standardError: number; +}; + +/** + * Standard errors the point's mean must fall short of the best by before the + * ladder stops: 2.5 leaves under a 1% chance of giving up on a point that + * could in fact beat it. + */ +const STOP_MARGIN_STANDARD_ERRORS = 2.5; + +/** + * The estimate from `metricId`'s last distribution frame with samples among + * `frames` — the frame the objective is read from — or null without one. + */ +export const estimateObjective = ( + frames: readonly MonteCarloUserDefinedMetricFrame[], + metricId: string, +): ObjectiveEstimate | null => { + for (let index = frames.length - 1; index >= 0; index--) { + const frame = frames[index]!; + if (frame.metricId !== metricId || frame.outputType !== "distribution") { + continue; + } + let runs = 0; + let sum = 0; + for (const [value, frequency] of frame.bins) { + runs += frequency; + sum += value * frequency; + } + if (runs === 0) { + continue; + } + const mean = sum / runs; + if (runs < 2) { + return { runs, mean, standardError: Number.POSITIVE_INFINITY }; + } + let squares = 0; + for (const [value, frequency] of frame.bins) { + squares += frequency * (value - mean) ** 2; + } + return { + runs, + mean, + standardError: Math.sqrt(squares / (runs - 1) / runs), + }; + } + return null; +}; + +/** + * Whether refining a point further is pointless: its mean sits more than + * the margin below (maximizing) or above (minimizing) the study's best, so + * more runs would only sharpen a value that cannot win. Never true without a + * best, or while a single run leaves the error unbounded. + */ +export const shouldStopRefining = ({ + direction, + best, + mean, + standardError, +}: { + direction: "maximize" | "minimize"; + best: number | null; + mean: number; + standardError: number; +}): boolean => { + if (best === null || !Number.isFinite(standardError)) { + return false; + } + const margin = STOP_MARGIN_STANDARD_ERRORS * standardError; + return direction === "maximize" ? mean + margin < best : mean - margin > best; +}; diff --git a/libs/@hashintel/petrinaut/src/react/optimizations/sir-optimization-input.fixtures.ts b/libs/@hashintel/petrinaut/src/react/optimizations/sir-optimization-input.fixtures.ts new file mode 100644 index 00000000000..3ec4f7a1da8 --- /dev/null +++ b/libs/@hashintel/petrinaut/src/react/optimizations/sir-optimization-input.fixtures.ts @@ -0,0 +1,51 @@ +import { petrinautOptimizationInputSchema } from "@hashintel/petrinaut-core"; +import { sirModel } from "@hashintel/petrinaut-core/examples"; + +const scenario = sirModel.petriNetDefinition.scenarios?.find( + (candidate) => candidate.id === "scenario__seasonal_flu", +); +const metric = sirModel.petriNetDefinition.metrics?.find( + (candidate) => candidate.id === "metric__infected_fraction", +); +if (!scenario || !metric) { + throw new Error("The SIR optimization fixtures are incomplete"); +} + +export const sirOptimizationScenario = scenario; +export const sirOptimizationMetric = metric; + +/** A two-trial study minimizing the SIR model's infected fraction. */ +export const sirOptimizationInput = petrinautOptimizationInputSchema.parse({ + kind: "petrinaut-optimization", + version: 1, + name: "SIR optimization", + model: { + title: sirModel.title, + definition: { + ...sirModel.petriNetDefinition, + scenarios: [scenario], + metrics: [metric], + }, + }, + scenario: { + id: scenario.id, + parameterBindings: { + population: { kind: "fixed", value: 1_000 }, + infected_ratio: { + kind: "optimize", + domain: { + kind: "continuous", + minimum: 0.001, + maximum: 0.2, + scale: "log", + }, + }, + }, + }, + objective: { + metricId: "metric__infected_fraction", + direction: "minimize", + }, + execution: { seed: 1, dt: 1, maxTime: 180 }, + study: { trials: 2, sampler: "tpe" }, +}); diff --git a/libs/@hashintel/petrinaut/src/react/optimizations/surface-grid.ts b/libs/@hashintel/petrinaut/src/react/optimizations/surface-grid.ts index c7bdde9b6d9..3674612371b 100644 --- a/libs/@hashintel/petrinaut/src/react/optimizations/surface-grid.ts +++ b/libs/@hashintel/petrinaut/src/react/optimizations/surface-grid.ts @@ -119,3 +119,66 @@ export function optimizationAxisMidpoint( ): number { return Math.round(axis.stepCount / 2); } + +/** The optimized boolean parameters, in binding order; they toggle rather than slide. */ +export function optimizationBooleanIdentifiers( + input: PetrinautOptimizationInput, +): string[] { + return Object.entries(input.scenario.parameterBindings) + .filter( + ([, binding]) => + binding.kind === "optimize" && binding.domain.kind === "boolean", + ) + .map(([identifier]) => identifier); +} + +type NavigationPoint = { + positions: Readonly>; + booleans: Readonly>; +}; + +/** One point as a cache key: positions in axis order, then booleans in binding order. */ +export function optimizationNavigationKey( + axes: readonly OptimizationSurfaceAxis[], + booleanIdentifiers: readonly string[], + point: NavigationPoint, +): string { + return [ + ...axes.map( + (axis) => `${axis.identifier}=${point.positions[axis.identifier] ?? 0}`, + ), + ...booleanIdentifiers.map( + (identifier) => `${identifier}=${point.booleans[identifier] ?? false}`, + ), + ].join("|"); +} + +/** + * Every scenario parameter's value at one point: the fixed bindings, each + * axis's value at its position, and each boolean as toggled. + */ +export function optimizationNavigationValues( + input: PetrinautOptimizationInput, + axes: readonly OptimizationSurfaceAxis[], + booleanIdentifiers: readonly string[], + point: NavigationPoint, +): Record { + const values: Record = {}; + for (const [identifier, binding] of Object.entries( + input.scenario.parameterBindings, + )) { + if (binding.kind === "fixed") { + values[identifier] = binding.value; + } + } + for (const axis of axes) { + values[axis.identifier] = optimizationAxisValueAt( + axis, + point.positions[axis.identifier] ?? 0, + ); + } + for (const identifier of booleanIdentifiers) { + values[identifier] = point.booleans[identifier] ?? false; + } + return values; +} diff --git a/libs/@hashintel/petrinaut/src/react/optimizations/use-optimization-source.ts b/libs/@hashintel/petrinaut/src/react/optimizations/use-optimization-source.ts new file mode 100644 index 00000000000..1bbb7a55d44 --- /dev/null +++ b/libs/@hashintel/petrinaut/src/react/optimizations/use-optimization-source.ts @@ -0,0 +1,27 @@ +import { use } from "react"; + +import { + isConnectedOptimization, + type PetrinautOptimizationSource, +} from "@hashintel/petrinaut-core/optimization"; + +import { PetrinautOptimizationContext } from "../optimization-context"; +import { UserSettingsContext } from "../state/user-settings-context"; + +/** + * The host's optimization source as the UI may use it. A remote capability + * passes through unchanged; a connected one counts only while the experimental + * In-browser optimization setting is on. `null` keeps the Optimizations + * surfaces hidden and nothing connects. + */ +export const useOptimizationSource = (): PetrinautOptimizationSource | null => { + const source = use(PetrinautOptimizationContext); + const { enableInBrowserOptimization } = use(UserSettingsContext); + if (source === null) { + return null; + } + if (isConnectedOptimization(source) && !enableInBrowserOptimization) { + return null; + } + return source; +}; diff --git a/libs/@hashintel/petrinaut/src/react/state/user-settings-context.ts b/libs/@hashintel/petrinaut/src/react/state/user-settings-context.ts index a8cdf2d7cf6..d04f6a93ba4 100644 --- a/libs/@hashintel/petrinaut/src/react/state/user-settings-context.ts +++ b/libs/@hashintel/petrinaut/src/react/state/user-settings-context.ts @@ -94,6 +94,14 @@ export type UserSettings = { * study drawer runs no compute of its own. */ enableOptimizationSurface: boolean; + /** + * Experimental: connect a host-supplied in-browser optimizer, which runs + * studies through the experiments backend and streams each step's metrics + * as it is evaluated. Off, a connected optimizer counts as none at all and + * the Optimizations surfaces stay hidden. A remote optimization capability + * is unaffected either way. + */ + enableInBrowserOptimization: boolean; subViewPanels: SubViewPanelsSettings; /** Where each document's canvas was last left, keyed by document id. */ canvasViewports: Record; @@ -124,6 +132,7 @@ export type UserSettingsActions = { setShowCompilationOutput: (value: boolean) => void; setEnableParameterSweeps: (value: boolean) => void; setEnableOptimizationSurface: (value: boolean) => void; + setEnableInBrowserOptimization: (value: boolean) => void; updateSubViewSection: ( containerName: string, sectionId: string, @@ -159,6 +168,7 @@ export const defaultUserSettings: UserSettings = { showCompilationOutput: false, enableParameterSweeps: false, enableOptimizationSurface: false, + enableInBrowserOptimization: false, subViewPanels: {}, canvasViewports: {}, }; @@ -189,6 +199,7 @@ const DEFAULT_CONTEXT_VALUE: UserSettingsContextValue = { setShowCompilationOutput: () => {}, setEnableParameterSweeps: () => {}, setEnableOptimizationSurface: () => {}, + setEnableInBrowserOptimization: () => {}, updateSubViewSection: () => {}, setCanvasViewport: () => {}, }; diff --git a/libs/@hashintel/petrinaut/src/react/state/user-settings-provider.tsx b/libs/@hashintel/petrinaut/src/react/state/user-settings-provider.tsx index 711b3ff67a9..0b9e467606e 100644 --- a/libs/@hashintel/petrinaut/src/react/state/user-settings-provider.tsx +++ b/libs/@hashintel/petrinaut/src/react/state/user-settings-provider.tsx @@ -127,6 +127,8 @@ export const UserSettingsProvider: React.FC = ({ ), })); }, + setEnableInBrowserOptimization: (value: boolean) => + setState((prev) => ({ ...prev, enableInBrowserOptimization: value })), updateSubViewSection: ( containerName: string, sectionId: string, diff --git a/libs/@hashintel/petrinaut/src/ui/components/contour-surface.tsx b/libs/@hashintel/petrinaut/src/ui/components/contour-surface.tsx index 28ca9bcf63c..0a1cf4c4f9a 100644 --- a/libs/@hashintel/petrinaut/src/ui/components/contour-surface.tsx +++ b/libs/@hashintel/petrinaut/src/ui/components/contour-surface.tsx @@ -1,13 +1,17 @@ /** * A filled contour plot over a sparse grid of sampled values, Optuna-style: * inverse-distance-weighted interpolation, marching-squares iso-lines, a - * Blues ramp, dots where data exists, and rings for external markers. + * Blues ramp, dots where data exists, and markers for external points — amber + * rings or filled dots, a hollow grey ring for a point without a value, and + * the navigation mark. * * Purely presentational: callers hand it grid-indexed values (`"x,y"` keys - * into an `nx × ny` index space, y up) and receive picks back as fractions - * of the plot area. The plot repaints as `values` stream in, one paint per - * animation frame; a caller clearing `values` for a new slice keeps the - * previous picture up, dimmed, until the new samples can replace it. + * into an `nx × ny` index space, y up — fractional coordinates included) and + * receive picks back as fractions of the plot area. The plot repaints as + * `values` stream in, one paint per animation frame; a caller clearing + * `values` for a new slice keeps the previous picture up, dimmed, until the + * new samples can replace it. Without `onPickFraction` the plot is + * display-only: no crosshair cursor, and a drag never arms. */ import { useEffect, useRef } from "react"; @@ -16,6 +20,7 @@ import { css } from "@hashintel/ds-helpers/css"; import { useElementSize } from "../../react/hooks/use-element-size"; import { type ContourSurfaceMarker, + type ContourSurfaceSampleMarks, type ContourSurfaceValues, createPaintState, paintField, @@ -29,6 +34,7 @@ import { export type { ContourSurfaceMarker, + ContourSurfaceSampleMarks, ContourSurfaceValues, } from "./contour-surface/paint-field"; export type { ContourSurfaceFraction } from "./contour-surface/use-surface-drag"; @@ -50,7 +56,7 @@ const frameStyle = css({ const canvasStyle = css({ display: "block", width: "[100%]", - cursor: "crosshair", + "&[data-interactive]": { cursor: "crosshair" }, // Horizontal touch drags navigate; vertical swipes stay the browser's to // scroll the drawer (it fires pointercancel, which aborts the drag). touchAction: "pan-y", @@ -62,6 +68,7 @@ export const ContourSurface = ({ ny, values, markers = [], + sampleMarks = "dot", height = 280, contentKey, onPickFraction, @@ -73,6 +80,11 @@ export const ContourSurface = ({ ny: number; values: ContourSurfaceValues; markers?: readonly ContourSurfaceMarker[]; + /** + * Whether every sampled cell gets a dot. `none` suits a plot whose samples + * are already drawn as markers. + */ + sampleMarks?: ContourSurfaceSampleMarks; /** Plot height in pixels; the width follows the container. */ height?: number; /** @@ -115,11 +127,12 @@ export const ContourSurface = ({ ny, values, markers, + sampleMarks, contentKey, }); }); return () => cancelAnimationFrame(frame); - }, [contentKey, height, markers, nx, ny, size, values]); + }, [contentKey, height, markers, nx, ny, sampleMarks, size, values]); return (
@@ -128,6 +141,7 @@ export const ContourSurface = ({ className={canvasStyle} style={{ height }} aria-label={ariaLabel} + data-interactive={onPickFraction ? "" : undefined} {...handlers} /> {preview ? : null} diff --git a/libs/@hashintel/petrinaut/src/ui/components/contour-surface/contour-field.test.ts b/libs/@hashintel/petrinaut/src/ui/components/contour-surface/contour-field.test.ts index 2076fa8849b..6d09504a31e 100644 --- a/libs/@hashintel/petrinaut/src/ui/components/contour-surface/contour-field.test.ts +++ b/libs/@hashintel/petrinaut/src/ui/components/contour-surface/contour-field.test.ts @@ -36,6 +36,31 @@ describe("idwRaster", () => { // Raster rows are top-down; grid y is up, so grid (0,1)=3 is top-left. expect([...raster]).toEqual([3, 4, 1, 2]); }); + + it("stays inside the sample range and reads each sample's value beside it with sparse, irregular samples", () => { + // Three fractional positions on an 11×11 grid, as a study's first steps. + const samples = [ + { x: 1.3, y: 7.6, value: 2 }, + { x: 8.1, y: 2.2, value: 10 }, + { x: 4.9, y: 4.4, value: 5 }, + ]; + const size = { nx: 11, ny: 11, width: 81, height: 81 }; + const raster = idwRaster({ samples, ...size }); + + expect(Math.min(...raster)).toBeGreaterThanOrEqual(2); + expect(Math.max(...raster)).toBeLessThanOrEqual(10); + for (const sample of samples) { + const px = Math.round((sample.x / 10) * 80); + const py = Math.round((1 - sample.y / 10) * 80); + expect(raster[py * 81 + px]).toBeCloseTo(sample.value, 0); + } + // Midway between two samples the field blends them rather than snapping. + const midX = Math.round(((1.3 + 8.1) / 2 / 10) * 80); + const midY = Math.round((1 - (7.6 + 2.2) / 2 / 10) * 80); + const midway = raster[midY * 81 + midX]!; + expect(midway).toBeGreaterThan(3); + expect(midway).toBeLessThan(9); + }); }); describe("createIdwAccumulator", () => { diff --git a/libs/@hashintel/petrinaut/src/ui/components/contour-surface/paint-field.ts b/libs/@hashintel/petrinaut/src/ui/components/contour-surface/paint-field.ts index d8b069d92a1..d807eba8577 100644 --- a/libs/@hashintel/petrinaut/src/ui/components/contour-surface/paint-field.ts +++ b/libs/@hashintel/petrinaut/src/ui/components/contour-surface/paint-field.ts @@ -1,11 +1,13 @@ /** * The imperative paint of a contour plot: the filled field blitted as one * raster-resolution image, iso-lines on top, dots where samples exist, and - * rings for external markers. + * markers for external points — amber rings or filled dots, a hollow grey ring + * for a point without a value, and the navigation mark. * - * A restart (the caller clearing `values` for a new slice) keeps the previous - * field up, dimmed, until the new samples can say something: two samples - * interpolate to a near-uniform wash that says less than the old picture. + * A field needs three samples to say anything: below that, two samples + * interpolate to a near-uniform wash, so the plot shows only its dots and + * markers — or, after a restart (the caller clearing `values` for a new + * slice), the previous field dimmed until the new samples can replace it. */ import { BLUES_STOPS, @@ -30,13 +32,25 @@ export type ContourSurfaceMarker = { y: number; /** Draw larger and stronger — e.g. a study's best trial. */ emphasis?: boolean; + /** + * `point` is an amber ring over a field computed elsewhere; `dot` is the + * same point filled, for a plot whose markers are the field's own samples. + * `navigation` marks where the viewer's controls sit rather than a data + * point: a dark ring with a centre dot, distinct from the amber data marks. + * `muted` is a point that carries no value, such as a pruned trial: a faint + * grey ring. + */ + kind?: "point" | "dot" | "navigation" | "muted"; }; +/** Whether the plot dots every sampled cell. */ +export type ContourSurfaceSampleMarks = "dot" | "none"; + /** Interpolation lattice points per grid cell. */ const RASTER_SUBDIVISION = 8; -/** Samples a fresh walk needs before its field replaces the ghost. */ -const GHOST_MIN_SAMPLES = 3; +/** Samples a field needs before it is painted, or replaces the ghost. */ +const FIELD_MIN_SAMPLES = 3; const ISO_LINE_COUNT = 10; @@ -118,6 +132,38 @@ const drawMarkers = ( ): void => { for (const marker of markers) { const [x, y] = toPixel(marker.x, marker.y); + if (marker.kind === "navigation") { + context.beginPath(); + context.arc(x, y, 6, 0, Math.PI * 2); + context.strokeStyle = "rgba(15, 23, 42, 0.9)"; + context.lineWidth = 1.5; + context.stroke(); + context.beginPath(); + context.arc(x, y, 1.5, 0, Math.PI * 2); + context.fillStyle = "rgba(15, 23, 42, 0.9)"; + context.fill(); + continue; + } + if (marker.kind === "muted") { + context.beginPath(); + context.arc(x, y, 3.5, 0, Math.PI * 2); + context.strokeStyle = "rgba(100, 116, 139, 0.6)"; + context.lineWidth = 1; + context.stroke(); + continue; + } + if (marker.kind === "dot") { + context.beginPath(); + context.arc(x, y, marker.emphasis ? 5.5 : 3.5, 0, Math.PI * 2); + context.fillStyle = marker.emphasis + ? "rgba(217, 119, 6, 0.95)" + : "rgba(217, 119, 6, 0.8)"; + context.fill(); + context.strokeStyle = "rgba(255, 255, 255, 0.9)"; + context.lineWidth = marker.emphasis ? 1.5 : 1; + context.stroke(); + continue; + } context.beginPath(); context.arc(x, y, marker.emphasis ? 5 : 3.5, 0, Math.PI * 2); context.strokeStyle = marker.emphasis @@ -205,7 +251,7 @@ const updateField = ( }; // The live field's canvas is reused across versions, so the ghost copies // it rather than aliasing it. - if (samples.length >= GHOST_MIN_SAMPLES) { + if (samples.length >= FIELD_MIN_SAMPLES) { const ghostImage = state.ghost?.image ?? document.createElement("canvas"); ghostImage.width = image.width; ghostImage.height = image.height; @@ -224,11 +270,22 @@ export const paintField = (options: { ny: number; values: ContourSurfaceValues; markers: readonly ContourSurfaceMarker[]; + sampleMarks: ContourSurfaceSampleMarks; /** Identity of the plotted quantity; a change drops the ghost. */ contentKey: string | undefined; }): void => { - const { canvas, state, width, height, nx, ny, values, markers, contentKey } = - options; + const { + canvas, + state, + width, + height, + nx, + ny, + values, + markers, + sampleMarks, + contentKey, + } = options; if (state.contentKey !== contentKey) { state.contentKey = contentKey; state.ghost = null; @@ -259,13 +316,14 @@ export const paintField = (options: { height - (y / Math.max(ny - 1, 1)) * height, ]; - if (samples.length < GHOST_MIN_SAMPLES && state.ghost !== null) { + if (samples.length >= FIELD_MIN_SAMPLES) { + drawField(context, updateField(state, samples, nx, ny), width, height); + } else if (state.ghost !== null) { context.globalAlpha = 0.45; drawField(context, state.ghost, width, height); context.globalAlpha = 1; - drawSamples(context, samples, toPixel); - } else if (samples.length > 0) { - drawField(context, updateField(state, samples, nx, ny), width, height); + } + if (sampleMarks === "dot") { drawSamples(context, samples, toPixel); } diff --git a/libs/@hashintel/petrinaut/src/ui/views/Editor/editor-view.tsx b/libs/@hashintel/petrinaut/src/ui/views/Editor/editor-view.tsx index d791b121d47..2d3b1e08289 100644 --- a/libs/@hashintel/petrinaut/src/ui/views/Editor/editor-view.tsx +++ b/libs/@hashintel/petrinaut/src/ui/views/Editor/editor-view.tsx @@ -22,6 +22,7 @@ import { sirModel, supplyChainWithDisruption, supplyChainProfit, + vaccinationCampaign, } from "@hashintel/petrinaut-core/examples"; import { usePetrinautCommands } from "../../../react"; @@ -436,6 +437,14 @@ export const EditorView = ({ clearSelection(); }, }, + { + id: "load-example-vaccination-campaign", + text: "Vaccination Campaign", + onClick: () => { + createNewNet(vaccinationCampaign); + clearSelection(); + }, + }, ], }, ] 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 6bd9572f759..0dc0a97887b 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 @@ -147,6 +147,7 @@ const TestProviders = ({ setEnableParameterSweeps: () => {}, setEnableOptimizationSurface: () => {}, setCanvasViewport: () => {}, + setEnableInBrowserOptimization: () => {}, updateSubViewSection: () => {}, }; @@ -162,6 +163,16 @@ const TestProviders = ({ setSweepSelection: () => {}, sampleSurfaceCells: () => Promise.resolve(null), sampleDetachedObjective: () => Promise.resolve(null), + runDetachedObjective: () => ({ + frames: { get: () => [], subscribe: () => () => {} }, + progress: { get: () => null, subscribe: () => () => {} }, + completion: Promise.resolve({ + ok: false, + cancelled: false, + reason: "unused", + }), + cancel: () => {}, + }), }} > diff --git a/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/experiments/create-experiment-drawer.tsx b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/experiments/create-experiment-drawer.tsx index 291032b5070..9d5e3c8a212 100644 --- a/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/experiments/create-experiment-drawer.tsx +++ b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/experiments/create-experiment-drawer.tsx @@ -10,7 +10,6 @@ import { Select, TextInput, Toggle, - Tooltip, type SelectItem, } from "@hashintel/ds-components"; import { css, cx } from "@hashintel/ds-helpers/css"; @@ -18,11 +17,6 @@ import { EMPTY_AD_HOC_STATE, isWebGpuAvailable, } from "@hashintel/petrinaut-core"; -import { - analyzeCompilation, - summarizeGpuUnavailability, - toGpuMetricSpecs, -} from "@hashintel/petrinaut-core/webgpu"; import { ExperimentsActionsContext, @@ -50,6 +44,8 @@ import { MODEL_METRIC_VALUE_PREFIX, type MetricKindGroup, } from "../metrics/metric-picker-options"; +import { ComputeBackendToggle } from "../shared/compute-backend-toggle"; +import { useGpuAvailability } from "../shared/use-gpu-availability"; import { areMetricLspDiagnosticSummariesEqual, EMPTY_METRIC_LSP_DIAGNOSTICS, @@ -61,7 +57,6 @@ import { ExperimentScenarioRun } from "./experiment-scenario-run"; import type { AdHocScenarioState, MonteCarloMetricSpec, - PetrinautExtensionSettings, Scenario, ScenarioParameter, SDCPN, @@ -81,64 +76,6 @@ const labelStyle = css({ color: "neutral.s120", }); -const backendControlStyle = css({ - display: "inline-flex", - alignItems: "center", - gap: "1.5", - flexShrink: "[0]", - // Matches the height the sibling inputs occupy, so the grid row's baselines - // line up rather than the control floating in a shorter cell. - minHeight: "[34px]", -}); - -const backendSideLabelStyle = css({ - fontSize: "sm", - fontWeight: "medium", - lineHeight: "[1]", - // Muted until selected, so the toggle's position reads as a choice between two - // named backends rather than an unlabelled on/off. - color: "neutral.s100", - transition: "[color 0.15s ease]", - "&[data-selected=true]": { - color: "neutral.s120", - }, -}); - -/** - * The GPU side is purple rather than neutral, so the accelerated path is visibly - * a different thing and not merely the toggle in its other position. - */ -const gpuSideLabelStyle = css({ - "&[data-selected=true]": { - color: "purple.s90", - }, -}); - -/* - * The design system's toggle has no purple tone, and adding one there would change - * a shared component for one screen's sake. These reach into its parts from - * outside instead: `&[data-state='checked'] [data-part='control']` is one - * selector more specific than the recipe's own `&[data-state='checked']`, so it - * wins without `!important`. - */ -const gpuToggleStyle = css({ - "&[data-state='checked'] [data-part='control']": { - backgroundColor: "purple.s80", - }, - "&[data-state='checked']:hover:not([data-disabled]) [data-part='control']": { - backgroundColor: "purple.s70", - }, -}); - -const gpuToggleGlowStyle = css({ - "&[data-state='checked'] [data-part='control']": { - animationName: "[petrinautGpuGlow]", - animationDuration: "[2.4s]", - animationIterationCount: "[infinite]", - animationTimingFunction: "ease-in-out", - }, -}); - const gridStyle = css({ display: "grid", gridTemplateColumns: "[repeat(3, minmax(0, 1fr))]", @@ -959,109 +896,6 @@ interface CreateExperimentDrawerProps { onClose: () => void; } -/** - * Whether the GPU backend could run this experiment, and the reason when it - * could not. - * - * The net is analysed asynchronously (lowering user code happens in the language - * worker) but the metric gate is evaluated synchronously from the drafts, so - * editing a metric updates the answer without another round-trip. - */ -function useGpuAvailability({ - enabled, - sdcpn, - extensions, - metricSpecs, -}: { - enabled: boolean; - sdcpn: SDCPN; - extensions: PetrinautExtensionSettings; - metricSpecs: readonly ExperimentMetricSpecInput[] | null; -}): { available: boolean; reason: string | null; pending: boolean } { - const { requestHirArtifacts } = use(LanguageClientContext); - const [netReason, setNetReason] = useState(null); - const [pending, setPending] = useState(false); - - useEffect(() => { - if (!enabled) { - return; - } - - let cancelled = false; - setPending(true); - - const analyze = async () => { - try { - const { artifacts } = await requestHirArtifacts(sdcpn, extensions, { - includeHir: true, - }); - if (cancelled) { - return; - } - setNetReason( - summarizeGpuUnavailability( - analyzeCompilation({ sdcpn, artifacts, extensions }), - ), - ); - } catch (caught) { - if (!cancelled) { - setNetReason( - caught instanceof Error - ? `The net could not be compiled: ${caught.message}` - : "The net could not be compiled.", - ); - } - } finally { - if (!cancelled) { - setPending(false); - } - } - }; - - void analyze(); - - return () => { - cancelled = true; - }; - }, [enabled, sdcpn, extensions, requestHirArtifacts]); - - if (!enabled) { - return { available: false, reason: null, pending: false }; - } - if (pending) { - return { available: false, reason: null, pending: true }; - } - if (netReason !== null) { - return { available: false, reason: netReason, pending: false }; - } - - // Expression metrics are computed from full simulation state, which the GPU - // path never materialises on the host, so they rule the backend out before the - // histogram gate is worth consulting. Narrowing as we go also gives - // `toGpuMetricSpecs` the compiled-spec type it wants without a cast: only - // expression specs lack an `artifact`. - const histogramSpecs: MonteCarloMetricSpec[] = []; - for (const spec of metricSpecs ?? []) { - if (spec.kind === "expression") { - return { - available: false, - reason: `Metric "${spec.label}" is an expression metric, which the GPU backend cannot compute. Use place token-count metrics to run on the GPU.`, - pending: false, - }; - } - histogramSpecs.push(spec); - } - - if (histogramSpecs.length > 0) { - const gpuMetrics = toGpuMetricSpecs(histogramSpecs); - if (!gpuMetrics.ok) { - return { available: false, reason: gpuMetrics.reason, pending: false }; - } - } - - return { available: true, reason: null, pending: false }; -} - export const CreateExperimentDrawer = ({ open, onClose, @@ -1069,12 +903,8 @@ export const CreateExperimentDrawer = ({ const { petriNetDefinition, extensions } = use(SDCPNContext); // Read here, not in ExperimentsProvider: that provider is mounted outside // UserSettingsProvider and so cannot see these settings. - const { - webGpuEnabled, - showAnimations, - enableAdHocScenarios, - enableParameterSweeps, - } = use(UserSettingsContext); + const { webGpuEnabled, enableAdHocScenarios, enableParameterSweeps } = + use(UserSettingsContext); const { createExperiment } = use(ExperimentsActionsContext); const scenarios = petriNetDefinition.scenarios ?? EMPTY_SCENARIOS; const [name, setName] = useState(DEFAULT_EXPERIMENT_NAME); @@ -1382,51 +1212,13 @@ export const CreateExperimentDrawer = ({ the backend is a property of the experiment like the rest, and a bare control below the grid read as an orphan. */} {webGpuEnabled && webGpuAvailable && ( -
+
Backend - - {/* Wrapped so the tooltip still opens while the control is - disabled — a disabled control fires no pointer events. */} - - - CPU - - - - GPU - - - +
)}
diff --git a/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/experiments/experiment-metric-timeline.tsx b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/experiments/experiment-metric-timeline.tsx index 5146b10596d..a70ae68238a 100644 --- a/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/experiments/experiment-metric-timeline.tsx +++ b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/experiments/experiment-metric-timeline.tsx @@ -58,9 +58,10 @@ const titleStyle = css({ whiteSpace: "nowrap", }); +/** The plot's height unless the owner sizes it. */ +const DEFAULT_PLOT_HEIGHT = 260; + const chartStyle = css({ - height: "[260px]", - minHeight: "[260px]", width: "full", minWidth: "[0]", _empty: { @@ -128,8 +129,6 @@ const aggregateNumberStyle = css({ display: "flex", alignItems: "center", justifyContent: "center", - height: "[260px]", - minHeight: "[260px]", width: "full", fontSize: "[44px]", fontWeight: "semibold", @@ -161,10 +160,12 @@ export const ExperimentMetricTimeline = ({ expectedOutputType, timeDomain, contentEpoch, + plotHeight = DEFAULT_PLOT_HEIGHT, }: { frames: readonly MetricFrame[]; displaySize: MetricSize; - onDisplaySizeChange: (size: MetricSize) => void; + /** Toggles between the two sizes; absent when the chart's slot is fixed. */ + onDisplaySizeChange?: (size: MetricSize) => void; /** * Title shown before any frame arrives. With it, the component keeps its * full shell — header, fixed-height plot area, footer — while empty, so @@ -187,8 +188,11 @@ export const ExperimentMetricTimeline = ({ * change crossfades the previous picture out instead of cutting. */ contentEpoch?: string; + /** The plot area's height in pixels; the header and controls add to it. */ + plotHeight?: number; }) => { const chartRootRef = useRef(null); + const plotSizeStyle = { height: plotHeight, minHeight: plotHeight }; const size = useElementSize(chartRootRef, { debounce: 50 }); const [settings, setSettings] = useState(DEFAULT_METRIC_VIEW_SETTINGS); const [selection, setSelection] = useState(null); @@ -245,28 +249,34 @@ export const ExperimentMetricTimeline = ({
{latestFrame?.label ?? label} -
-
+ {onDisplaySizeChange ? ( +
+
+ ) : null}
{view.displayMode === "number" ? ( -
+
{view.aggregateNumber === null ? "n/a" : formatNumber(view.aggregateNumber)}
) : (
-
+
{view.hasPlotData || lastOutputType !== null ? null : (
Waiting for metric data
)} diff --git a/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/experiments/experiments-story-fixtures.tsx b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/experiments/experiments-story-fixtures.tsx index fbca7653be2..36bb6313d99 100644 --- a/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/experiments/experiments-story-fixtures.tsx +++ b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/experiments/experiments-story-fixtures.tsx @@ -5,6 +5,7 @@ import { sirModel } from "@hashintel/petrinaut-core/examples"; import { type CreateExperimentInput, + type DetachedObjectiveRunOutcome, ExperimentsActionsContext, type ExperimentsActionsValue, ExperimentsContext, @@ -21,6 +22,11 @@ import { } from "../../../../../../react/state/editor-context"; import type { SDCPNContextValue } from "../../../../../../react/state/sdcpn-context"; +import type { + MonteCarloUserDefinedMetricFrame, + MonteCarloWorkerProgress, + ReadableStore, +} from "@hashintel/petrinaut-core"; export const sirSdcpnContextValue: SDCPNContextValue = { createNewNet: () => {}, @@ -397,6 +403,106 @@ export function makeFakeSurfaceSampler( }); } +/** A store the fake compute writes and the UI subscribes to. */ +function createFakeStore( + initial: T, +): ReadableStore & { set(next: T): void } { + let current = initial; + const listeners = new Set<(value: T) => void>(); + return { + get: () => current, + subscribe: (listener) => { + listeners.add(listener); + return () => listeners.delete(listener); + }, + set: (next) => { + current = next; + for (const listener of listeners) { + listener(next); + } + }, + }; +} + +/** + * The fake of a streaming objective batch: ten frames of the synthetic bump + * at the request's parameter values, one every 60 ms, then the result. + */ +export const fakeRunDetachedObjective: ExperimentsActionsValue["runDetachedObjective"] = + (request) => { + const values = Object.values(request.scenarioParameterValues).filter( + (entry): entry is number => typeof entry === "number", + ); + const objective = syntheticSweepObjective(values[0] ?? 0, values[1] ?? 0); + const frames = createFakeStore( + [], + ); + const progress = createFakeStore(null); + let cancelled = false; + const completion = new Promise((resolve) => { + const totalTicks = 10; + let tick = 0; + const step = () => { + if (cancelled) { + resolve({ ok: false, cancelled: true, reason: "cancelled" }); + return; + } + tick += 1; + const fraction = tick / totalTicks; + const time = request.maxTime * fraction; + frames.set([ + ...frames.get(), + { + metricId: request.metric.id, + label: request.metric.label, + outputType: "distribution", + frameNumber: Math.round(time / request.dt), + time, + bins: [ + [Math.round(objective * fraction * 100) / 100, request.runCount], + ], + value: null, + frameValue: null, + timeValue: null, + runSampleCount: request.runCount, + timeSampleCount: request.runCount, + }, + ]); + progress.set({ + activeRuns: tick < totalTicks ? request.runCount : 0, + advancedRuns: request.runCount, + allFinished: tick >= totalTicks, + completedRuns: tick < totalTicks ? 0 : request.runCount, + erroredRuns: 0, + frameNumber: Math.round(time / request.dt), + runCount: request.runCount, + time, + }); + if (tick < totalTicks) { + setTimeout(step, 60); + return; + } + resolve({ + ok: true, + runsCompleted: request.runCount, + metricFrames: frames.get(), + runResults: new Map(), + computeBackend: request.computeBackend, + computeBackendFallbackReason: null, + }); + }; + setTimeout(step, 60); + }); + return { + frames, + progress, + completion, + cancel: () => { + cancelled = true; + }, + }; + }; + export function FakeExperimentsProvider({ children, initialExperiments, @@ -413,7 +519,7 @@ export function FakeExperimentsProvider({ overrides?: Partial< Pick< ExperimentsContextValue, - "sampleSurfaceCells" | "sampleDetachedObjective" + "sampleSurfaceCells" | "sampleDetachedObjective" | "runDetachedObjective" > >; /** @@ -607,6 +713,7 @@ export function FakeExperimentsProvider({ ); }); }, + runDetachedObjective: fakeRunDetachedObjective, ...overrides, })); diff --git a/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/experiments/sweep-surface.tsx b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/experiments/sweep-surface.tsx index b74de0b4f17..63ac3d2cf94 100644 --- a/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/experiments/sweep-surface.tsx +++ b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/experiments/sweep-surface.tsx @@ -23,6 +23,7 @@ import { import { ContourSurface } from "../../../../../components/contour-surface"; import { formatAxisValue } from "../shared/format-axis-value"; import { + describeSurfaceSampling, SurfaceAxisControls, SurfaceCaption, SurfaceControlLabel, @@ -226,7 +227,7 @@ export const SweepSurface = ({ { x: nearestGridIndex(xAxis, sweepSelection?.[xAxis.identifier]), y: nearestGridIndex(yAxis, sweepSelection?.[yAxis.identifier]), - emphasis: true, + kind: "navigation", }, ]} onPickFraction={handlePickFraction} @@ -240,9 +241,11 @@ export const SweepSurface = ({ ? { x: readoutAt(xAxis, preview.x), y: readoutAt(yAxis, preview.y) } : null } - sampledCount={cellValues.size} - totalCells={totalCells} - runsPerCell={SURFACE_CELL_RUNS} + text={describeSurfaceSampling({ + sampledCount: cellValues.size, + totalCells, + runsPerCell: SURFACE_CELL_RUNS, + })} /> ); diff --git a/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/experiments/view-experiment-drawer.tsx b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/experiments/view-experiment-drawer.tsx index 1d9e561686b..0f45662e661 100644 --- a/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/experiments/view-experiment-drawer.tsx +++ b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/experiments/view-experiment-drawer.tsx @@ -1,6 +1,6 @@ import { use } from "react"; -import { Button, Drawer, Icon, Tooltip } from "@hashintel/ds-components"; +import { Button, Drawer, Icon } from "@hashintel/ds-components"; import { css } from "@hashintel/ds-helpers/css"; import { @@ -8,31 +8,12 @@ import { type ExperimentRecord, } from "../../../../../../react/experiments/context"; import { Section, SectionList } from "../../../../../components/section"; +import { ComputeBackendBadge } from "../shared/compute-backend-badge"; import { SweepNavigator } from "./sweep-navigator"; import { SweepSurface } from "./sweep-surface"; import { ExperimentMetrics } from "./view-experiment-drawer/experiment-metrics"; import { ExperimentSummary } from "./view-experiment-drawer/experiment-summary"; -// Local rather than the design system's `Badge`, whose `brand` scheme puts -// #5EB1EF on a near-white #FBFDFF — about 2.3:1, below the 4.5:1 WCAG AA -// needs for text this size. -const backendBadgeStyle = css({ - display: "inline-flex", - alignItems: "center", - gap: "1", - paddingX: "1.5", - paddingY: "[2px]", - borderRadius: "sm", - fontSize: "xs", - fontWeight: "medium", - color: "neutral.s110", - backgroundColor: "neutral.s10", - "&[data-tone=active]": { - color: "blue.s100", - backgroundColor: "blue.s10", - }, -}); - // The drawer body is a column: the summary, the navigator, and the surface // hold still at the top, and the metric charts alone scroll below them. const drawerBodyStyle = css({ @@ -54,18 +35,6 @@ const metricsScrollStyle = css({ scrollbarWidth: "[thin]", }); -const describeComputeBackend = (experiment: ExperimentRecord): string => { - if (experiment.computeBackend === "webgpu") { - return "Stepped on the GPU through WebGPU. Distributions match the CPU backend statistically; individual trajectories differ (different random generators)."; - } - if (experiment.computeBackendFallbackReason !== null) { - // The notification that carried this is gone by the time anyone wonders - // why the results are not GPU-backed. - return `The GPU backend was requested but could not run this net: ${experiment.computeBackendFallbackReason}`; - } - return "Stepped on the CPU, across worker threads."; -}; - // Keeps its footprint when a run can no longer be cancelled, so Remove and // Close do not slide when a run finishes. const cancelSlotStyle = css({ @@ -73,26 +42,6 @@ const cancelSlotStyle = css({ "&[data-hidden=true]": { visibility: "hidden" }, }); -const ComputeBackendBadge = ({ - experiment, -}: { - experiment: ExperimentRecord; -}) => { - const isGpu = experiment.computeBackend === "webgpu"; - - return ( - - - {isGpu ? : null} - {isGpu ? "GPU" : "CPU"} - - - ); -}; - export const ViewExperimentDrawer = ({ open, onClose, @@ -134,7 +83,7 @@ export const ViewExperimentDrawer = ({ // In the header rather than the strip below, so which backend ran // stays visible when the section is collapsed. renderHeaderAction={() => ( - + )} > diff --git a/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/experiments/view-experiment-drawer/experiment-metrics.tsx b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/experiments/view-experiment-drawer/experiment-metrics.tsx index 78907dd8150..c4d58e393d5 100644 --- a/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/experiments/view-experiment-drawer/experiment-metrics.tsx +++ b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/experiments/view-experiment-drawer/experiment-metrics.tsx @@ -1,55 +1,13 @@ /** - * The drawer's metric charts: one tile per metric, each resizable between a - * half-width and a full-width slot. Before any frame has arrived the tiles - * are stable shells per configured metric, so the first data causes no - * layout shift. + * The experiment drawer's metric charts: one tile per configured metric, + * fed the record's frames. */ -import { useState } from "react"; - -import { css, cx } from "@hashintel/ds-helpers/css"; - -import { - ExperimentMetricTimeline, - type MetricSize, -} from "../experiment-metric-timeline"; +import { MetricTiles, type MetricTile } from "../../shared/metric-tiles"; import type { ExperimentRecord } from "../../../../../../../react/experiments/context"; -const gridStyle = css({ - display: "grid", - gridTemplateColumns: "repeat(2, minmax(0, 1fr))", - alignItems: "start", - gap: "3", -}); - -const tileStyle = css({ - display: "flex", - flexDirection: "column", - gap: "1", - minWidth: "[0]", - padding: "3", - borderWidth: "[1px]", - borderStyle: "solid", - borderColor: "neutral.bd.subtle", - borderRadius: "md", - backgroundColor: "neutral.s00", -}); - -const largeTileStyle = css({ - gridColumn: "[1 / -1]", -}); - -type MetricFrame = ExperimentRecord["metricFrames"][number]; - -const metricTiles = ( - experiment: ExperimentRecord, -): { - id: string; - label: string; - frames: MetricFrame[]; - outputType: MetricFrame["outputType"]; -}[] => { - const framesById = new Map(); +const metricTiles = (experiment: ExperimentRecord): MetricTile[] => { + const framesById = new Map(); for (const frame of experiment.metricFrames) { const frames = framesById.get(frame.metricId) ?? []; frames.push(frame); @@ -68,35 +26,12 @@ export const ExperimentMetrics = ({ experiment, }: { experiment: ExperimentRecord; -}) => { - const [sizes, setSizes] = useState>({}); - // What the frames represent: a selection change fades the previous picture - // out inside each plot instead of cutting to the sparse new stream. - const contentEpoch = JSON.stringify(experiment.sweep?.selection ?? null); - - return ( -
- {metricTiles(experiment).map((tile) => { - const size = sizes[tile.id] ?? "small"; - return ( -
- - setSizes((previous) => ({ ...previous, [tile.id]: nextSize })) - } - /> -
- ); - })} -
- ); -}; +}) => ( + +); diff --git a/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/experiments/view-experiment-drawer/experiment-summary.tsx b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/experiments/view-experiment-drawer/experiment-summary.tsx index e14ae5b3d71..2dd376971f1 100644 --- a/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/experiments/view-experiment-drawer/experiment-summary.tsx +++ b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/experiments/view-experiment-drawer/experiment-summary.tsx @@ -2,7 +2,7 @@ * The drawer's summary: a strip of stats with a status dot, the compute * activity underneath, and the error text when the experiment failed. */ -import { type ReactNode, useEffect, useState } from "react"; +import { useEffect, useState } from "react"; import { css } from "@hashintel/ds-helpers/css"; @@ -10,78 +10,27 @@ import { type ExperimentRecord, getExperimentElapsedMs, isExperimentActive, + type SweepBatchStatus, } from "../../../../../../../react/experiments/context"; +import { experimentProgressPercent } from "../../../../shared/experiment-progress"; +import { + ComputeActivity, + type ComputeActivityBatch, +} from "../../shared/compute-activity"; +import { + SummaryStat, + SummaryStatusDot, + type SummaryStatusTone, + SummaryStrip, +} from "../../shared/summary-strip"; import { formatDurationMs } from "../format-duration"; import { formatNumber } from "../shared/format-number"; -import { ComputeActivity } from "./experiment-summary/compute-activity"; const summaryStyle = css({ marginTop: "-1", marginBottom: "3", }); -// Every stat carries its own leading hairline, and the strip shifts left by -// exactly one divider-plus-gap so each row's first divider lands outside the -// clipping wrapper — wrapped rows therefore start flush, not with a floating -// rule (a sibling selector cannot see flex line breaks). -const stripClipStyle = css({ - overflow: "hidden", -}); - -const stripStyle = css({ - display: "flex", - flexWrap: "wrap", - alignItems: "center", - rowGap: "2", - marginLeft: "[-17px]", -}); - -const statStyle = css({ - display: "flex", - flexDirection: "column", - gap: "[1px]", - minWidth: "[0]", - paddingLeft: "4", - marginLeft: "[1px]", - borderLeftWidth: "[1px]", - borderLeftStyle: "solid", - borderLeftColor: "neutral.bd.subtle", - paddingRight: "4", -}); - -const statLabelStyle = css({ - fontSize: "[10px]", - fontWeight: "medium", - letterSpacing: "[0.04em]", - textTransform: "uppercase", - color: "neutral.s70", -}); - -const statValueStyle = css({ - fontSize: "sm", - fontWeight: "medium", - color: "neutral.s120", - fontVariantNumeric: "tabular-nums", - overflow: "hidden", - textOverflow: "ellipsis", - whiteSpace: "nowrap", -}); - -// Inline-block inside the value span, so a long value still ellipsizes (a -// flex value container turns its text into an item ellipsis cannot reach). -const statusDotStyle = css({ - display: "inline-block", - width: "[7px]", - height: "[7px]", - borderRadius: "full", - marginRight: "1.5", - verticalAlign: "[1px]", - backgroundColor: "neutral.s60", - "&[data-tone=active]": { backgroundColor: "blue.s100" }, - "&[data-tone=done]": { backgroundColor: "green.s90" }, - "&[data-tone=error]": { backgroundColor: "red.s100" }, -}); - const activityStyle = css({ marginTop: "2", }); @@ -94,7 +43,7 @@ const errorStyle = css({ const STATUS_DISPLAY: Record< ExperimentRecord["status"], - { label: string; tone: "active" | "done" | "error" | "neutral" } + { label: string; tone: SummaryStatusTone } > = { initializing: { label: "Initializing", tone: "active" }, running: { label: "Running", tone: "active" }, @@ -124,27 +73,6 @@ const useNow = (active: boolean): number => { return now; }; -const Stat = ({ - label, - minChars, - children, -}: { - label: string; - /** Reserve this many characters so a changing value never reflows the strip. */ - minChars?: number; - children: ReactNode; -}) => ( -
- {label} - - {children} - -
-); - /** Longest status label, so the strip never reflows as the status changes. */ /** * Simulated time to show when no batch is publishing progress: an idle sweep @@ -160,6 +88,38 @@ const STATUS_CHARS = ...Object.values(STATUS_DISPLAY).map((entry) => entry.label.length), ) + 2; +/** + * "selection" is the navigator's own ladder — the priority work; "surface" + * is a contour chunk; "refine" is a single cell brought up to depth. + */ +const BATCH_KIND_META: Record< + SweepBatchStatus["kind"], + Pick +> = { + selection: { label: "Selection", tone: "priority" }, + surface: { label: "Surface", tone: "background" }, + refine: { label: "Refine", tone: "background" }, +}; + +/** The sweep's batches as the activity list shows them. */ +const activityBatches = ( + sweepBatches: readonly SweepBatchStatus[], +): ComputeActivityBatch[] => + sweepBatches.map((batch) => ({ + id: String(batch.id), + ...BATCH_KIND_META[batch.kind], + runCount: batch.runCount, + completedRuns: batch.completedRuns, + })); + +/** The bar under the stats: the selection's runs for a sweep, simulated time otherwise. */ +const activityBar = (experiment: ExperimentRecord) => ({ + percent: experimentProgressPercent(experiment), + label: experiment.sweep + ? `Selection · ${experiment.sweep.runsSampled.toLocaleString("en-US")} / ${experiment.runCount.toLocaleString("en-US")} runs` + : `Time · ${(experiment.progress?.time ?? 0).toLocaleString("en-US")} / ${experiment.maxTime.toLocaleString("en-US")}`, +}); + export const ExperimentSummary = ({ experiment, }: { @@ -172,54 +132,54 @@ export const ExperimentSummary = ({ return (
-
-
- - - {status.label} - - {experiment.scenarioName ?? "Default"} - - {progress - ? `${progress.activeRuns} active, ${progress.completedRuns} complete` - : experiment.runCount} - - - {progress?.erroredRuns ?? 0} - - - {formatNumber(progress?.time ?? settledTime(experiment))} /{" "} - {formatNumber(experiment.maxTime)} - - {/* Wall-clock, as distinct from the simulated time; dashed out + + + + {status.label} + + + {experiment.scenarioName ?? "Default"} + + + {progress + ? `${progress.activeRuns} active, ${progress.completedRuns} complete` + : experiment.runCount} + + + {progress?.erroredRuns ?? 0} + + + {formatNumber(progress?.time ?? settledTime(experiment))} /{" "} + {formatNumber(experiment.maxTime)} + + {/* Wall-clock, as distinct from the simulated time; dashed out when stepping never began. */} - - {elapsedMs === null ? "—" : formatDurationMs(elapsedMs)} - -
-
+ + {elapsedMs === null ? "—" : formatDurationMs(elapsedMs)} + +
{experiment.error ? ( diff --git a/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/browser-optimizer.stories.tsx b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/browser-optimizer.stories.tsx new file mode 100644 index 00000000000..b27766e6657 --- /dev/null +++ b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/browser-optimizer.stories.tsx @@ -0,0 +1,224 @@ +import { createBrowserOptimization } from "@hashintel/petrinaut-core/browser-optimization"; +import { + sirModel, + supplyChainProfit, + vaccinationCampaign, +} from "@hashintel/petrinaut-core/examples"; + +import { + AutoStudy, + type AutoStudyDescription, + RunnableSimulateViewStory, + type StoryExample, +} from "../simulate-view-story-harness"; + +import type { ExperimentComputeBackend } from "../../../../../../react/experiments/context"; +import type { UserSettings } from "../../../../../../react/state/user-settings-context"; +import type { Meta, StoryObj } from "@storybook/react-vite"; + +/** One optimizer for the whole Storybook session, as the website keeps one per page. */ +const browserOptimization = createBrowserOptimization(); + +type BrowserOptimizerArgs = { + steps: number; + runsPerStep: number; + maxTime: number; + computeBackend: ExperimentComputeBackend; + autoStart: boolean; +}; + +const meta = { + title: "Simulate / Browser optimizer (real)", + parameters: { layout: "fullscreen" }, + args: { + steps: 4, + runsPerStep: 3, + maxTime: 60, + computeBackend: "cpu", + autoStart: true, + }, + argTypes: { + steps: { control: { type: "range", min: 1, max: 20, step: 1 } }, + runsPerStep: { control: { type: "range", min: 1, max: 10, step: 1 } }, + maxTime: { control: { type: "number", min: 1 } }, + computeBackend: { control: "inline-radio", options: ["cpu", "webgpu"] }, + autoStart: { control: "boolean" }, + }, +} satisfies Meta; + +export default meta; + +type Story = StoryObj; + +/** A study's fixed part; the args supply steps, runs per step and max time. */ +type StudyPreset = Omit< + AutoStudyDescription, + "steps" | "runsPerStep" | "maxTime" +>; + +const seasonalFluStudy: StudyPreset = { + scenarioName: "Seasonal Flu", + name: "Peak infection", + dt: 0.1, + optimize: { + population: { minimum: 500, maximum: 5_000 }, + infected_ratio: { minimum: 0, maximum: 1 }, + }, + objective: { metricName: "Infected Fraction", direction: "maximize" }, +}; + +const richStockStudy: StudyPreset = { + scenarioName: "Rich stock", + name: "Adjusted profit", + dt: 1, + optimize: { + production_rate: { minimum: 50, maximum: 400 }, + selling_price: { minimum: 20, maximum: 60 }, + }, + objective: { metricName: "Adjusted profit", direction: "maximize" }, +}; + +const winterWaveStudy: StudyPreset = { + scenarioName: "Winter wave", + name: "Cheapest response", + dt: 0.1, + optimize: { + vaccination_coverage: { minimum: 0, maximum: 0.9 }, + contact_reduction: { minimum: 0, maximum: 0.8 }, + }, + objective: { metricName: "Total cost", direction: "minimize" }, +}; + +const BrowserOptimizerStory = ({ + example, + study, + settings, + steps, + runsPerStep, + maxTime, + computeBackend, + autoStart, +}: BrowserOptimizerArgs & { + example: StoryExample; + study: StudyPreset; + settings?: Partial; +}) => ( + + {autoStart ? ( + + ) : null} + +); + +const firstRunNote = + "The first study in a browser downloads the Python runtime and the optimizer packages from jsDelivr and PyPI (about 10 MB, a few seconds); the record shows Running with no steps until then, and later studies reuse the browser's cache. The whole study runs in this tab: Optuna in a worker, each step as seeded simulations on the experiments backend."; + +const watchForNote = + "Everything is in view at once: the summary strip (status, steps, best, backend, progress bars and the computing chip), the Parameters band, the Surface beside the objective's chart, and the steps table filling the rest. Watch the band follow each step, the Surface gain a dot per step — the best emphasized, the field filling in between them, the ringed dot on the step in flight streaming its running value — and the chart beside it stream the objective over the step's runs. While the study runs the sliders are disabled and a drag on the Surface does nothing; turn Follow steps off to take over early. Once complete, click the Surface or move a slider: the point refines in escalating batches, its value enters the field, and the chart streams again."; + +const gpuNote = + "With WebGPU on in settings, the create form's Backend switch appears but stays disabled for an expression objective by design: the GPU backend cannot compute expression metrics, so steps run on the CPU."; + +export const SirCpu: Story = { + name: "SIR CPU", + parameters: { + docs: { + description: { + story: `The SIR model's Seasonal Flu scenario, maximizing Infected Fraction over population and infected ratio on the CPU. ${firstRunNote} ${watchForNote} ${gpuNote}`, + }, + }, + }, + render: (args) => ( + + ), +}; + +export const SirGpuRequested: Story = { + name: "SIR GPU requested", + args: { computeBackend: "webgpu" }, + parameters: { + docs: { + description: { + story: `The SIR study with WebGPU enabled and the GPU requested for its steps. The GPU backend declines the expression objective, so the record's badge reads CPU and its tooltip carries the reason: the real fallback. ${firstRunNote} ${watchForNote} ${gpuNote}`, + }, + }, + }, + render: (args) => ( + + ), +}; + +export const SupplyChain: Story = { + name: "Supply Chain", + parameters: { + docs: { + description: { + story: `The supply chain example's Rich stock scenario, maximizing Adjusted profit over production rate and selling price on the CPU; two numeric parameters, so the Surface shows. ${firstRunNote} ${watchForNote} ${gpuNote}`, + }, + }, + }, + render: (args) => ( + + ), +}; + +export const VaccinationCampaign: Story = { + name: "Vaccination Campaign", + args: { steps: 6 }, + parameters: { + docs: { + description: { + story: `The Vaccination Campaign example's Winter wave scenario, minimizing Total cost over vaccination coverage (0 to 0.9) and contact reduction (0 to 0.8) on the CPU: the model built for this drawer. Cases are priced against a campaign and distancing whose prices rise quadratically, so the Surface shows a valley along the epidemic threshold with its floor near a coverage of 0.45 and a contact reduction of 0.4 (about 960 against 1,280 to 2,220 in the corners). Six steps are still the sampler's random start-up, so expect scattered dots with the best step landing in the valley and the Surface field dipping there. The net is GPU-eligible, so an experiment on it runs on the GPU when one is available, while the study's expression objective keeps its steps on the CPU. ${firstRunNote} ${watchForNote} ${gpuNote}`, + }, + }, + }, + render: (args) => ( + + ), +}; + +export const Manual: Story = { + args: { autoStart: false }, + parameters: { + docs: { + description: { + story: `The real optimizer with the In-browser optimization setting on and the Optimizations tab open, and no study: the entry point for hand-testing the create form. ${firstRunNote} ${watchForNote} ${gpuNote}`, + }, + }, + }, + render: (args) => ( + + ), +}; 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 d65eb0f6045..906d1a15b91 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 @@ -18,7 +18,12 @@ import { } from "@hashintel/petrinaut-core"; import { LanguageClientContext } from "../../../../../../react/lsp/context"; -import { OptimizationsContext } from "../../../../../../react/optimizations/context"; +import { PetrinautOptimizationContext } from "../../../../../../react/optimization-context"; +import { + type CreateOptimizationOptions, + OptimizationsContext, + type OptimizationsContextValue, +} from "../../../../../../react/optimizations/context"; import { SDCPNContext } from "../../../../../../react/state/sdcpn-context"; import { UserSettingsContext } from "../../../../../../react/state/user-settings-context"; import { UserSettingsProvider } from "../../../../../../react/state/user-settings-provider"; @@ -36,7 +41,6 @@ import { import { createOptimizationParameterDraft } from "./optimization-parameter-row"; import type { LanguageClientContextValue } from "../../../../../../react/lsp/context"; -import type { OptimizationsContextValue } from "../../../../../../react/optimizations/context"; import type { SDCPNContextValue } from "../../../../../../react/state/sdcpn-context"; import type { OptimizationParameterDraft } from "./optimization-parameter-row"; import type { @@ -46,6 +50,7 @@ import type { Scenario, SDCPN, } from "@hashintel/petrinaut-core"; +import type { PetrinautConnectedOptimization } from "@hashintel/petrinaut-core/optimization"; import type { ReactNode } from "react"; const { addMetricMock } = vi.hoisted(() => ({ addMetricMock: vi.fn() })); @@ -137,10 +142,12 @@ vi.mock("@hashintel/ds-components", async (importOriginal) => { const Toggle = ({ "aria-label": ariaLabel, + disabled, onChange, value, }: { "aria-label": string; + disabled?: boolean; onChange: (value: boolean) => void; value: boolean; }) => ( @@ -148,6 +155,7 @@ vi.mock("@hashintel/ds-components", async (importOriginal) => { aria-label={ariaLabel} type="checkbox" checked={value} + disabled={disabled} onChange={(event) => onChange(event.target.checked)} /> ); @@ -200,20 +208,53 @@ type TestProviderProps = { sdcpnContextValue?: SDCPNContextValue; /** Turns the Ad-hoc scenarios user setting on for this render. */ enableAdHocScenarios?: boolean; + /** Turns the WebGPU user setting on for this render. */ + webGpuEnabled?: boolean; + /** + * Supplies a connected optimizer (with the In-browser optimization setting + * on), so the form offers a backend choice. + */ + connectedSource?: boolean; +}; + +/** A connected source that never runs: the form only asks what kind it is. */ +const connectedSource: PetrinautConnectedOptimization = { + kind: "connected", + connect: () => ({ + createOptimizationRun: () => Promise.resolve({ runId: "run-test" }), + async *attachOptimizationRun() { + yield { type: "started", requestedTrials: 1, seq: 1 }; + }, + cancelOptimizationRun: () => Promise.resolve(), + extendOptimizationRun: () => Promise.resolve(), + releaseOptimizationRun: () => Promise.resolve(), + dispose: () => {}, + }), }; -/** Overrides one user setting below the provider (localStorage is not +/** Overrides user settings below the provider (localStorage is not * writable in this environment). */ -const AdHocSettingOverride = ({ - enabled, +const SettingsOverride = ({ + enableAdHocScenarios, + webGpuEnabled, + enableInBrowserOptimization, children, }: { - enabled: boolean; + enableAdHocScenarios: boolean; + webGpuEnabled: boolean; + enableInBrowserOptimization: boolean; children: ReactNode; }) => { const value = use(UserSettingsContext); return ( - + {children} ); @@ -224,6 +265,8 @@ const TestProviders = ({ languageClient, sdcpnContextValue = sirSdcpnContextValue, enableAdHocScenarios = false, + webGpuEnabled = false, + connectedSource: withConnectedSource = false, }: TestProviderProps) => { const portalContainerRef = useRef(null); const optimizations: OptimizationsContextValue = { @@ -234,19 +277,29 @@ const TestProviders = ({ createOptimization, cancelOptimization: () => {}, removeOptimization: () => {}, + extendOptimization: () => Promise.resolve(), + setOptimizationNavigation: () => {}, retryOptimization: () => Promise.resolve(null), }; const drawer = ( - - - - -
- {}} /> - - - - + + + + + +
+ {}} /> + + + + + ); return ( @@ -264,6 +317,7 @@ const TestProviders = ({ afterEach(() => { cleanup(); + vi.unstubAllGlobals(); vi.clearAllMocks(); }); @@ -478,7 +532,10 @@ describe("CreateOptimizationDrawer", () => { it("submits a successfully validated saved metric", async () => { const languageClient = makeSuccessfulLanguageClient(); const createOptimization = vi.fn( - async (_input: PetrinautOptimizationInput) => "optimization-saved", + async ( + _input: PetrinautOptimizationInput, + _options?: CreateOptimizationOptions, + ) => "optimization-saved", ); const savedMetric = sirSdcpnContextValue.petriNetDefinition.metrics?.[0]; expect(savedMetric).toBeDefined(); @@ -505,11 +562,128 @@ describe("CreateOptimizationDrawer", () => { const submittedInput = createOptimization.mock.calls[0]![0]; expect(submittedInput.model.definition.metrics).toEqual([savedMetric]); expect(submittedInput.objective.metricId).toBe(savedMetric!.id); - expect(submittedInput.execution).toEqual({ - seed: 1234, - dt: 0.1, - maxTime: 180, + const { seed: submittedSeed, ...execution } = submittedInput.execution; + expect(Number.isInteger(submittedSeed)).toBe(true); + expect(execution).toEqual({ dt: 0.1, maxTime: 180, seedsPerTrial: 1 }); + expect(createOptimization.mock.calls[0]![1]).toEqual({ + computeBackend: "cpu", + parallelism: 1, + }); + expect(screen.queryByLabelText("Parallel steps")).toBeNull(); + }); + + it("sends runs per step as the manifest's seeds per trial", async () => { + const languageClient = makeSuccessfulLanguageClient(); + const createOptimization = vi.fn( + async (_input: PetrinautOptimizationInput) => "optimization-seeded", + ); + const savedMetric = sirSdcpnContextValue.petriNetDefinition.metrics?.[0]; + expect(savedMetric).toBeDefined(); + openConfiguration({ createOptimization, languageClient }); + + fireEvent.change(screen.getByLabelText("Runs per step"), { + target: { value: "3" }, + }); + 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" })); + fireEvent.click(screen.getByRole("button", { name: /Run/ })); + + await waitFor(() => expect(createOptimization).toHaveBeenCalledOnce()); + expect(createOptimization.mock.calls[0]![0].execution.seedsPerTrial).toBe( + 3, + ); + }); + + it("sends the typed seed with the manifest", async () => { + const languageClient = makeSuccessfulLanguageClient(); + const createOptimization = vi.fn( + async (_input: PetrinautOptimizationInput) => "optimization-seed", + ); + const savedMetric = sirSdcpnContextValue.petriNetDefinition.metrics?.[0]; + expect(savedMetric).toBeDefined(); + openConfiguration({ createOptimization, languageClient }); + + fireEvent.change(screen.getByLabelText("Seed"), { + target: { value: "4242" }, + }); + 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" })); + fireEvent.click(screen.getByRole("button", { name: /Run/ })); + + await waitFor(() => expect(createOptimization).toHaveBeenCalledOnce()); + expect(createOptimization.mock.calls[0]![0].execution.seed).toBe(4242); + }); + + it("rejects a seed above the limit before submitting", () => { + openConfiguration(); + fireEvent.change( + screen.getByRole("combobox", { name: "Select a metric" }), + { + target: { + value: `${MODEL_METRIC_VALUE_PREFIX}metric__infected_fraction`, + }, + }, + ); + fireEvent.click( + screen.getByRole("checkbox", { name: "Optimize infected_ratio" }), + ); + fireEvent.click(screen.getByRole("button", { name: "Maximize" })); + + fireEvent.change(screen.getByLabelText("Seed"), { + target: { value: "2147483648" }, + }); + + expect( + screen.getByText("Seed must be an integer between 0 and 2,147,483,647"), + ).toBeTruthy(); + expect(screen.getByRole("button", { name: /Run/ })).toHaveProperty( + "disabled", + true, + ); + }); + + it("rejects runs per step outside 1..100 before submitting", () => { + openConfiguration(); + fireEvent.change( + screen.getByRole("combobox", { name: "Select a metric" }), + { + target: { + value: `${MODEL_METRIC_VALUE_PREFIX}metric__infected_fraction`, + }, + }, + ); + fireEvent.click( + screen.getByRole("checkbox", { name: "Optimize infected_ratio" }), + ); + fireEvent.click(screen.getByRole("button", { name: "Maximize" })); + + fireEvent.change(screen.getByLabelText("Runs per step"), { + target: { value: "101" }, }); + + expect( + screen.getByText("Runs per step must be an integer between 1 and 100"), + ).toBeTruthy(); + expect( + (screen.getByRole("button", { name: /Run/ }) as HTMLButtonElement) + .disabled, + ).toBe(true); }); it("submits a transient custom metric without persisting it", async () => { @@ -618,6 +792,8 @@ describe("CreateOptimizationDrawer", () => { metric, direction: "minimize", optimizationSteps: 20, + seedsPerTrial: 4, + seed: 99, dt: 0.5, maxTime: 100, }); @@ -660,7 +836,12 @@ describe("CreateOptimizationDrawer", () => { metricId: metric.id, direction: "minimize", }); - expect(input.execution).toEqual({ seed: 1234, dt: 0.5, maxTime: 100 }); + expect(input.execution).toEqual({ + seed: 99, + dt: 0.5, + maxTime: 100, + seedsPerTrial: 4, + }); expect(input.study).toEqual({ trials: 20, sampler: "tpe" }); }); @@ -735,6 +916,8 @@ describe("CreateOptimizationDrawer", () => { metric, direction: "maximize", optimizationSteps: 10, + seedsPerTrial: 1, + seed: 7, dt: 0.5, maxTime: 50, }); @@ -801,3 +984,134 @@ describe("CreateOptimizationDrawer", () => { ); }); }); + +describe("CreateOptimizationDrawer backend choice", () => { + const openWithWebGpu = (props: TestProviderProps) => { + // `isWebGpuAvailable()` only reads `navigator.gpu`, so a bare object is + // enough — and spreading the real Navigator would drop its prototype. + vi.stubGlobal("navigator", { gpu: {} }); + openConfiguration(props); + }; + + it("offers no backend cell while WebGPU is off in settings", () => { + openWithWebGpu({ connectedSource: true, webGpuEnabled: false }); + + expect(document.querySelector("[data-backend-state]")).toBeNull(); + expect(screen.queryByText("Backend")).toBeNull(); + }); + + it("offers no backend cell for a remote optimizer, which runs elsewhere", () => { + openWithWebGpu({ connectedSource: false, webGpuEnabled: true }); + + expect(document.querySelector("[data-backend-state]")).toBeNull(); + }); + + it("offers the cell for a connected optimizer and rules the GPU out for the expression objective", async () => { + openWithWebGpu({ + connectedSource: true, + webGpuEnabled: true, + languageClient: makeSuccessfulLanguageClient(), + }); + + expect(screen.getByText("Backend")).toBeTruthy(); + const savedMetric = sirSdcpnContextValue.petriNetDefinition.metrics?.[0]; + fireEvent.change( + screen.getByRole("combobox", { name: "Select a metric" }), + { + target: { value: `${MODEL_METRIC_VALUE_PREFIX}${savedMetric!.id}` }, + }, + ); + + await waitFor(() => { + expect( + document + .querySelector("[data-backend-state]") + ?.getAttribute("data-backend-state"), + ).toBe("unavailable"); + }); + expect( + document.querySelector( + "[data-backend-state] input[type='checkbox']", + )!.disabled, + ).toBe(true); + }); + + it("passes the backend as a creation option", async () => { + const languageClient = makeSuccessfulLanguageClient(); + const createOptimization = vi.fn( + async ( + _input: PetrinautOptimizationInput, + _options?: CreateOptimizationOptions, + ) => "optimization-backend", + ); + const savedMetric = sirSdcpnContextValue.petriNetDefinition.metrics?.[0]; + openWithWebGpu({ + connectedSource: true, + webGpuEnabled: true, + languageClient, + createOptimization, + }); + + 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" })); + fireEvent.click(screen.getByRole("button", { name: /Run/ })); + + await waitFor(() => expect(createOptimization).toHaveBeenCalledOnce()); + // The switch never left the CPU side: the objective is an expression + // metric, which the GPU backend cannot compute. + expect(createOptimization.mock.calls[0]![1]).toEqual({ + computeBackend: "cpu", + parallelism: 1, + }); + }); + + it("offers parallel steps to a connected optimizer and passes the count as a creation option", async () => { + const createOptimization = vi.fn( + async ( + _input: PetrinautOptimizationInput, + _options?: CreateOptimizationOptions, + ) => "optimization-parallel", + ); + const savedMetric = sirSdcpnContextValue.petriNetDefinition.metrics?.[0]; + openConfiguration({ + connectedSource: true, + languageClient: makeSuccessfulLanguageClient(), + createOptimization, + }); + + 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" })); + fireEvent.change(screen.getByLabelText("Parallel steps"), { + target: { value: "5" }, + }); + expect( + screen.getByText("Parallel steps must be an integer between 1 and 4"), + ).toBeTruthy(); + fireEvent.change(screen.getByLabelText("Parallel steps"), { + target: { value: "3" }, + }); + fireEvent.click(screen.getByRole("button", { name: /Run/ })); + + await waitFor(() => expect(createOptimization).toHaveBeenCalledOnce()); + expect(createOptimization.mock.calls[0]![1]).toEqual({ + computeBackend: "cpu", + parallelism: 3, + }); + }); +}); 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 5fc3073afce..4862560e7c2 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 @@ -14,20 +14,27 @@ import { } from "@hashintel/ds-components"; import { css } from "@hashintel/ds-helpers/css"; import { - PETRINAUT_DEFAULT_SEED, + PETRINAUT_OPTIMIZATION_MAX_SEEDS_PER_TRIAL, PETRINAUT_OPTIMIZATION_MAX_STEPS_PER_TRIAL, PETRINAUT_OPTIMIZATION_MAX_TOTAL_STEPS, PETRINAUT_OPTIMIZATION_MAX_TRIALS, createUserKeyedRecord, EMPTY_AD_HOC_STATE, + isWebGpuAvailable, metricSchema, petrinautOptimizationInputSchema, adHocOptimizationBindings, synthesizeAdHocOptimization, } from "@hashintel/petrinaut-core"; +import { + isConnectedOptimization, + PETRINAUT_OPTIMIZATION_MAX_PARALLELISM, + PETRINAUT_OPTIMIZATION_MAX_SEED, +} from "@hashintel/petrinaut-core/optimization"; import { LanguageClientContext } from "../../../../../../react/lsp/context"; import { OptimizationsContext } from "../../../../../../react/optimizations/context"; +import { useOptimizationSource } from "../../../../../../react/optimizations/use-optimization-source"; import { SDCPNContext } from "../../../../../../react/state/sdcpn-context"; import { UserSettingsContext } from "../../../../../../react/state/user-settings-context"; import { AdHocScenarioForm } from "../../../../../components/ad-hoc-scenario-form/ad-hoc-scenario-form"; @@ -47,12 +54,22 @@ import { getMetricKindIcon, MODEL_METRIC_VALUE_PREFIX, } from "../metrics/metric-picker-options"; +import { ComputeBackendToggle } from "../shared/compute-backend-toggle"; +import { useGpuAvailability } from "../shared/use-gpu-availability"; import { createOptimizationParameterDraft, type OptimizationParameterDraft, OptimizationParameterRow, } from "./optimization-parameter-row"; +import { + isValidOptimizationSeed, + randomOptimizationSeed, +} from "./optimization-seed"; +import type { + ExperimentComputeBackend, + ExperimentMetricSpecInput, +} from "../../../../../../react/experiments/context"; import type { AdHocScenarioState, AdHocSynthesisError, @@ -177,6 +194,8 @@ const directionOptions = [ ]; const OPTIMIZATION_SAMPLER = "tpe" as const; +const DEFAULT_SEEDS_PER_TRIAL = 1; +const DEFAULT_PARALLELISM = 1; const AD_HOC_SCENARIO_VALUE = "__adhoc__"; const AD_HOC_SCENARIO_LABEL = "No scenario"; const DEFAULT_DT = 0.1; @@ -363,6 +382,9 @@ function getConfigurationError({ missingObjectiveMessage, direction, optimizationSteps, + seedsPerTrial, + parallelism, + seed, dt, maxTime, }: { @@ -375,6 +397,9 @@ function getConfigurationError({ missingObjectiveMessage: string; direction: Direction | null; optimizationSteps: number | null; + seedsPerTrial: number | null; + parallelism: number | null; + seed: number | null; dt: number | null; maxTime: number | null; }): string | null { @@ -421,6 +446,25 @@ function getConfigurationError({ ) { return `Optimization steps must be an integer between 1 and ${PETRINAUT_OPTIMIZATION_MAX_TRIALS.toLocaleString()}`; } + if ( + seedsPerTrial === null || + !Number.isInteger(seedsPerTrial) || + seedsPerTrial < 1 || + seedsPerTrial > PETRINAUT_OPTIMIZATION_MAX_SEEDS_PER_TRIAL + ) { + return `Runs per step must be an integer between 1 and ${PETRINAUT_OPTIMIZATION_MAX_SEEDS_PER_TRIAL.toLocaleString()}`; + } + if ( + parallelism === null || + !Number.isInteger(parallelism) || + parallelism < 1 || + parallelism > PETRINAUT_OPTIMIZATION_MAX_PARALLELISM + ) { + return `Parallel steps must be an integer between 1 and ${PETRINAUT_OPTIMIZATION_MAX_PARALLELISM}`; + } + if (!isValidOptimizationSeed(seed)) { + return `Seed must be an integer between 0 and ${PETRINAUT_OPTIMIZATION_MAX_SEED.toLocaleString()}`; + } if (dt === null || !Number.isFinite(dt) || dt <= 0) { return "Time step must be a positive number"; } @@ -435,10 +479,10 @@ function getConfigurationError({ return `Use at most ${PETRINAUT_OPTIMIZATION_MAX_STEPS_PER_TRIAL.toLocaleString()} simulation steps per optimization step`; } if ( - simulationStepsPerOptimization * optimizationSteps > + simulationStepsPerOptimization * seedsPerTrial * optimizationSteps > PETRINAUT_OPTIMIZATION_MAX_TOTAL_STEPS ) { - return `Use at most ${PETRINAUT_OPTIMIZATION_MAX_TOTAL_STEPS.toLocaleString()} simulation steps across the optimization`; + return `Use at most ${PETRINAUT_OPTIMIZATION_MAX_TOTAL_STEPS.toLocaleString()} simulation steps across the optimization (time steps × runs per step × optimization steps)`; } return null; } @@ -453,6 +497,8 @@ export function buildPetrinautOptimizationInput({ metric, direction, optimizationSteps, + seedsPerTrial, + seed, dt, maxTime, }: { @@ -464,6 +510,8 @@ export function buildPetrinautOptimizationInput({ metric: Metric; direction: Direction; optimizationSteps: number; + seedsPerTrial: number; + seed: number; dt: number; maxTime: number; }): PetrinautOptimizationInput { @@ -525,7 +573,7 @@ export function buildPetrinautOptimizationInput({ }, scenario: { id: scenario.id, parameterBindings }, objective: { metricId: metric.id, direction }, - execution: { seed: PETRINAUT_DEFAULT_SEED, dt, maxTime }, + execution: { seed, dt, maxTime, seedsPerTrial }, study: { trials: optimizationSteps, sampler: OPTIMIZATION_SAMPLER }, }); } @@ -545,6 +593,8 @@ export function buildAdHocPetrinautOptimizationInput({ metric, direction, optimizationSteps, + seedsPerTrial, + seed, dt, maxTime, }: { @@ -556,6 +606,8 @@ export function buildAdHocPetrinautOptimizationInput({ metric: Metric; direction: Direction; optimizationSteps: number; + seedsPerTrial: number; + seed: number; dt: number; maxTime: number; }): PetrinautOptimizationInput { @@ -573,7 +625,7 @@ export function buildAdHocPetrinautOptimizationInput({ }, scenario: { id: scenario.id, parameterBindings }, objective: { metricId: metric.id, direction }, - execution: { seed: PETRINAUT_DEFAULT_SEED, dt, maxTime }, + execution: { seed, dt, maxTime, seedsPerTrial }, study: { trials: optimizationSteps, sampler: OPTIMIZATION_SAMPLER }, }); } @@ -588,7 +640,11 @@ export const CreateOptimizationDrawer = ({ const { extensions, petriNetDefinition, title } = use(SDCPNContext); const { requestHirArtifacts } = use(LanguageClientContext); const { createOptimization } = use(OptimizationsContext); - const { enableAdHocScenarios } = use(UserSettingsContext); + const { enableAdHocScenarios, webGpuEnabled } = use(UserSettingsContext); + const source = useOptimizationSource(); + // A remote study runs wherever the service runs, so only a connected + // source — trials evaluated in this browser — gets a backend choice. + const backendSelectable = source !== null && isConnectedOptimization(source); const scenarios = petriNetDefinition.scenarios ?? []; const metrics = petriNetDefinition.metrics ?? []; const [selectedScenarioId, setSelectedScenarioId] = useState( @@ -606,6 +662,14 @@ export const CreateOptimizationDrawer = ({ const [optimizationSteps, setOptimizationSteps] = useState( 100, ); + const [seedsPerTrial, setSeedsPerTrial] = useState( + DEFAULT_SEEDS_PER_TRIAL, + ); + const [seed, setSeed] = useState(randomOptimizationSeed); + const [parallelism, setParallelism] = useState( + DEFAULT_PARALLELISM, + ); + const [gpuRequested, setGpuRequested] = useState(false); const [dt, setDt] = useState(DEFAULT_DT); const [maxTime, setMaxTime] = useState(180); const [error, setError] = useState(null); @@ -666,6 +730,42 @@ export const CreateOptimizationDrawer = ({ ); }; + // The objective is an expression metric whichever way it is authored, which + // the GPU backend cannot compute, so the switch stays disabled with that + // reason; the net analysis still runs so the reason names the first + // blocker. The gate reads the metric's kind, not its code, so the custom + // objective counts before any code is typed. + const objectiveMetricForGpu = + metricSource === "saved" + ? selectedSavedMetric + : { id: customMetricId, name: CUSTOM_OBJECTIVE_METRIC_NAME, code: "" }; + const objectiveMetricSpecs: ExperimentMetricSpecInput[] | null = + objectiveMetricForGpu + ? [ + { + kind: "expression", + id: objectiveMetricForGpu.id, + label: objectiveMetricForGpu.name, + code: objectiveMetricForGpu.code, + sampleRuns: "all", + runOutput: { type: "distribution" }, + }, + ] + : null; + const webGpuAvailable = isWebGpuAvailable(); + const gpu = useGpuAvailability({ + enabled: open && backendSelectable && webGpuEnabled && webGpuAvailable, + sdcpn: petriNetDefinition, + extensions, + metricSpecs: objectiveMetricSpecs, + }); + // Derived rather than stored, so a net edited into ineligibility after the + // switch was flipped neither shows as on nor submits a GPU study. + const gpuSelected = gpuRequested && gpu.available; + const computeBackend: ExperimentComputeBackend = gpuSelected + ? "webgpu" + : "cpu"; + const resetConfigurationState = (scenario?: Scenario) => { setName("Optimization"); setDrafts(scenario ? createParameterDrafts(scenario) : {}); @@ -674,6 +774,10 @@ export const CreateOptimizationDrawer = ({ setCustomMetricId(crypto.randomUUID()); setDirection(null); setOptimizationSteps(100); + setSeedsPerTrial(DEFAULT_SEEDS_PER_TRIAL); + setParallelism(DEFAULT_PARALLELISM); + setSeed(randomOptimizationSeed()); + setGpuRequested(false); setDt(DEFAULT_DT); setMaxTime(180); setError(null); @@ -703,6 +807,9 @@ export const CreateOptimizationDrawer = ({ missingObjectiveMessage: "Select an objective metric", direction, optimizationSteps, + seedsPerTrial, + parallelism, + seed, dt, maxTime, }) @@ -713,6 +820,9 @@ export const CreateOptimizationDrawer = ({ validationError || direction === null || optimizationSteps === null || + seedsPerTrial === null || + parallelism === null || + !isValidOptimizationSeed(seed) || dt === null || maxTime === null ) { @@ -783,6 +893,8 @@ export const CreateOptimizationDrawer = ({ metric, direction, optimizationSteps, + seedsPerTrial, + seed, dt, maxTime, }) @@ -795,10 +907,12 @@ export const CreateOptimizationDrawer = ({ metric, direction, optimizationSteps, + seedsPerTrial, + seed, dt, maxTime, }); - await createOptimization(input); + await createOptimization(input, { computeBackend, parallelism }); resetState(); resetMetricForm(); } catch (submitError) { @@ -881,6 +995,9 @@ export const CreateOptimizationDrawer = ({ : "Define the custom objective metric", direction, optimizationSteps, + seedsPerTrial, + parallelism, + seed, dt, maxTime, }) @@ -1000,16 +1117,73 @@ export const CreateOptimizationDrawer = ({ - - - + > + + , + + + , + // Steps overlap only where this browser evaluates them; + // a remote study's service decides its own pace. + ...(backendSelectable + ? [ + + + , + ] + : []), + // Only offered where the choice exists: a connected source + // with WebGPU switched on in settings. + ...(backendSelectable && webGpuEnabled && webGpuAvailable + ? [ + + + , + ] + : []), + ]} + + + + + diff --git a/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/optimization-seed.ts b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/optimization-seed.ts new file mode 100644 index 00000000000..5cea13b1a14 --- /dev/null +++ b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/optimization-seed.ts @@ -0,0 +1,16 @@ +import { PETRINAUT_OPTIMIZATION_MAX_SEED } from "@hashintel/petrinaut-core/optimization"; + +/** + * A fresh study seed. Every study used to share one fixed seed, so two studies + * over different models drew the same normalized positions for their random + * start-up steps and painted the same surface. A draw per form keeps a study + * reproducible through the field while making studies differ by default. + */ +export const randomOptimizationSeed = (): number => + Math.floor(Math.random() * (PETRINAUT_OPTIMIZATION_MAX_SEED + 1)); + +export const isValidOptimizationSeed = (seed: number | null): seed is number => + seed !== null && + Number.isInteger(seed) && + seed >= 0 && + seed <= PETRINAUT_OPTIMIZATION_MAX_SEED; diff --git a/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/optimization-status.ts b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/optimization-status.ts new file mode 100644 index 00000000000..7f0338dedf0 --- /dev/null +++ b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/optimization-status.ts @@ -0,0 +1,22 @@ +import type { OptimizationRecord } from "../../../../../../react/optimizations/context"; + +/** + * The status as the drawer and the list name it. A connected study is + * stopped rather than cancelled: its sampler stays, and it can be continued. + */ +export const describeOptimizationStatus = ( + optimization: Pick, +): string => { + switch (optimization.status) { + case "initializing": + return "Initializing"; + case "running": + return "Running"; + case "complete": + return "Complete"; + case "error": + return "Error"; + case "cancelled": + return optimization.navigation === null ? "Cancelled" : "Stopped"; + } +}; diff --git a/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/optimization-surface.stories.tsx b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/optimization-surface.stories.tsx index 697515cf0bb..f5358b8a13d 100644 --- a/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/optimization-surface.stories.tsx +++ b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/optimization-surface.stories.tsx @@ -1,17 +1,25 @@ import { useEffect, useState } from "react"; import { FakeExperimentsProvider } from "../experiments/experiments-story-fixtures"; -import { OptimizationSurface } from "./optimization-surface"; +import { + NavigatedOptimizationSurface, + OptimizationSurface, +} from "./optimization-surface"; import { makeOptimizationInput, makeOptimizationRecord, + makeSelectionStream, + makeSyntheticObjectiveSampler, makeTrials, + navigationAtTrial, optimizedBindingSets, - syntheticObjective, + useFakeStudyClock, } from "./optimizations-story-fixtures"; -import type { DetachedObjectiveRequest } from "../../../../../../react/experiments/context"; -import type { OptimizationRecord } from "../../../../../../react/optimizations/context"; +import type { + OptimizationNavigation, + OptimizationRecord, +} from "../../../../../../react/optimizations/context"; import type { Meta, StoryObj } from "@storybook/react-vite"; const meta = { @@ -23,43 +31,6 @@ export default meta; type Story = StoryObj; -/** - * The stories' local compute: the same synthetic objective the fake trials - * used, returned as a single-bin distribution frame after `delayFor` the - * batch — so the contour fills in progressively and the trial rings land on - * it, at whatever pace the story simulates. - */ -const makeSyntheticObjectiveSampler = - (delayFor: (runCount: number) => number) => - (request: DetachedObjectiveRequest) => { - const objective = syntheticObjective(request.scenarioParameterValues); - const frame = { - metricId: request.metric.id, - label: request.metric.label, - outputType: "distribution" as const, - frameNumber: 365, - time: 365, - bins: [ - [Math.round(objective * 100) / 100, request.runCount], - ] as (readonly [number, number])[], - value: null, - frameValue: null, - timeValue: null, - runSampleCount: request.runCount, - timeSampleCount: request.runCount, - }; - return new Promise<{ - runsCompleted: number; - metricFrames: [typeof frame]; - }>((resolve) => { - setTimeout( - () => - resolve({ runsCompleted: request.runCount, metricFrames: [frame] }), - delayFor(request.runCount), - ); - }); - }; - const sampleSyntheticObjective = makeSyntheticObjectiveSampler(() => 80); /** Batches cost real simulation time on the CPU lane, scaling with runs. */ @@ -219,3 +190,113 @@ export const ManyParameters: Story = { /> ), }; + +/** + * A connected study's surface computes nothing: its steps are the field's + * samples — a dot each, the best emphasized, pruned steps hollow — and the + * field is interpolated between them. The navigation lives in the drawer's + * navigator, so the plot has no sliders of its own; once the study is over, + * clicking the plot moves the navigation and the picked point's value enters + * the field from the selection stream. + */ +const ConnectedSurfaceStory = ({ stepCount }: { stepCount: number }) => { + const study = makeTrials(baseInput, stepCount); + const [navigation, setNavigation] = useState(() => + navigationAtTrial(baseInput, study.trials[study.best?.trial ?? 0]!, false), + ); + const selection = makeSelectionStream({ + input: baseInput, + navigation, + runsCompleted: 100, + }); + const optimization = makeOptimizationRecord({ + input: baseInput, + trials: study.trials, + best: study.best, + status: "complete", + navigation, + selection, + }); + + return ( +
+ + setNavigation((previous) => ({ ...previous, ...patch })) + } + /> +
+ ); +}; + +export const ConnectedTwoSteps: Story = { + name: "Connected study, two steps", + render: () => , +}; + +export const ConnectedTwelveSteps: Story = { + name: "Connected study, twelve steps", + render: () => , +}; + +/** + * A connected study mid-run, following its steps: one lands every 1.5 s, and + * the step in flight streams its running objective into the field at the + * ringed dot before its own dot lands. The plot only displays until the last + * step lands, then a click picks a point. + */ +const ConnectedMidRunStory = () => { + const study = makeTrials(baseInput, 12); + const { landed, progress } = useFakeStudyClock({ + steps: study.trials.length, + ticksPerStep: 10, + tickMs: 150, + }); + const trials = study.trials.slice(0, landed); + const inFlight = study.trials[landed]; + const [chosen, setChosen] = useState(() => + navigationAtTrial(baseInput, study.trials[0]!, true), + ); + // While following, the navigation is wherever the optimizer is evaluating; + // once every step has landed it holds at the last one. + const navigation = chosen.followTrials + ? navigationAtTrial(baseInput, inFlight ?? study.trials.at(-1)!, true) + : chosen; + const selection = inFlight + ? makeSelectionStream({ + input: baseInput, + navigation, + followedTrial: inFlight.trial, + runsCompleted: 1, + computing: true, + progress, + }) + : makeSelectionStream({ input: baseInput, navigation, runsCompleted: 100 }); + const optimization = makeOptimizationRecord({ + input: baseInput, + trials, + best: trials.at(-1)?.best ?? null, + status: inFlight ? "running" : "complete", + navigation, + selection, + }); + + return ( +
+ setChosen({ ...navigation, ...patch })} + /> +
+ ); +}; + +export const ConnectedMidRun: Story = { + name: "Connected study mid-run, streaming a step", + render: () => , +}; diff --git a/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/optimization-surface.tsx b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/optimization-surface.tsx index 9c50326322c..6f398ce2c6b 100644 --- a/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/optimization-surface.tsx +++ b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/optimization-surface.tsx @@ -1,62 +1,72 @@ /** * The optimization surface: an Optuna-style filled contour of the study's - * objective over two optimized parameters, computed locally. + * objective over two optimized parameters. The study's trials arrive with + * parameter and objective values and are projected onto the two shown axes. * - * The study's trials arrive with parameter and objective values and are - * drawn as markers projected onto the two shown axes. The interpolated fill - * comes from points this view computes itself: it walks an X×Y sub-grid of - * the shown parameters in quad-tree order, running the study's frozen model - * with its objective metric on a background worker, holding every other - * optimized parameter at its slider position (initially the best trial's - * value). The selected point refines with escalating batches, and the - * readout streams the objective's mean and median as runs accumulate. + * Two variants share the plot and differ in where the field comes from. + * `OptimizationSurface`, for a study run elsewhere, computes the fill itself: + * it walks an X×Y sub-grid of the shown parameters in quad-tree order, + * running the study's frozen model with its objective metric on a background + * worker, holding every other optimized parameter at its slider position, and + * refines the selected point with escalating batches. The trials sit on that + * fill as rings. `NavigatedOptimizationSurface`, for a study evaluated in this + * browser, computes nothing: the trials are the samples, the field is + * interpolated between them, and the point being evaluated — or, once the + * study is over, the point the navigation refines — streams its running value + * into the field as it lands. It follows the record's navigation and the + * provider's selection stream; the drawer's navigator holds the controls. */ -import { use, useEffect, useRef, useState } from "react"; +import { type ReactNode, use, useEffect, useRef, useState } from "react"; import { Slider } from "@hashintel/ds-components"; import { css } from "@hashintel/ds-helpers/css"; -import { createUserKeyedRecord, getOwn } from "@hashintel/petrinaut-core"; import { ExperimentsActionsContext } from "../../../../../../react/experiments/context"; import { distributionStats } from "../../../../../../react/experiments/distribution-stats"; -import { - EXPERIMENT_RUN_LADDER, - mergeMetricFramesAcrossCells, -} from "../../../../../../react/experiments/parameter-grid"; +import { EXPERIMENT_RUN_LADDER } from "../../../../../../react/experiments/parameter-grid"; import { sweepCellObjective } from "../../../../../../react/experiments/sweep-cell-objective"; -import { sweepBatchSeed } from "../../../../../../react/experiments/sweep-session"; import { buildOptimizationSurfaceAxes, - optimizationAxisMidpoint, - optimizationAxisPositionFor, optimizationAxisValueAt, + optimizationBooleanIdentifiers, } from "../../../../../../react/optimizations/surface-grid"; -import { - ContourSurface, - contourSurfaceKey, -} from "../../../../../components/contour-surface"; import { formatAxisValue } from "../shared/format-axis-value"; +import { describeSurfaceSampling } from "../shared/surface-frame"; import { - SurfaceAxisControls, - SurfaceCaption, - SurfaceFrame, -} from "../shared/surface-frame"; -import { - quadTreeLevels, SURFACE_CELL_RUNS, surfacePositions, } from "../shared/surface-sampling"; -import { useSurfaceWalk } from "../shared/use-surface-walk"; +import { + type OptimizationSurfaceView, + resolveSurfaceBooleans, + resolveSurfacePositions, + surfaceSliceKey, + surfaceWalkKey, +} from "./optimization-surface/navigation-slice"; +import { + sampleStudyCell, + type StudyCellCache, +} from "./optimization-surface/sample-study-cell"; +import { + describeSurfaceState, + inFlightSurfaceField, + mergeSurfaceFields, + navigatedSurfaceSample, + OptimizationSurfacePlot, + surfaceCellKeyAt, + surfaceInteraction, + trialSurfaceField, + withNavigatedSample, +} from "./optimization-surface/surface-plot"; +import { useStudySurfaceWalk } from "./optimization-surface/use-study-surface-walk"; -import type { ExperimentsContextValue } from "../../../../../../react/experiments/context"; import type { DistributionStats } from "../../../../../../react/experiments/distribution-stats"; -import type { SweepCellSnapshot } from "../../../../../../react/experiments/sweep-session"; -import type { OptimizationRecord } from "../../../../../../react/optimizations/context"; -import type { OptimizationSurfaceAxis } from "../../../../../../react/optimizations/surface-grid"; import type { - ContourSurfaceFraction, - ContourSurfaceMarker, -} from "../../../../../components/contour-surface"; + OptimizationNavigation, + OptimizationRecord, + OptimizationSelectionStream, +} from "../../../../../../react/optimizations/context"; +import type { OptimizationSurfaceAxis } from "../../../../../../react/optimizations/surface-grid"; /** Ladder cap for the selected point's local refinement. */ const SELECTED_POINT_MAX_RUNS = 100; @@ -97,111 +107,12 @@ const readoutStyle = css({ fontVariantNumeric: "tabular-nums", }); -/** - * Brings one cell up to at least `minRuns` locally computed runs, merging - * batches into `cache`. A cell's entry is the promise of its deepest result, - * so the walk and the selected point's refinement queue behind each other - * instead of both sampling from the same run index. - */ -const sampleStudyCell = async (options: { - sampleDetachedObjective: ExperimentsContextValue["sampleDetachedObjective"]; - cache: Map>; - optimization: Pick; - axes: readonly OptimizationSurfaceAxis[]; - xAxisId: string; - yAxisId: string; - /** Slider position per off-surface axis, plus boolean fallbacks. */ - slice: string; - xPosition: number; - yPosition: number; - minRuns: number; -}): Promise => { - const { - sampleDetachedObjective, - cache, - optimization, - axes, - xAxisId, - yAxisId, - slice, - xPosition, - yPosition, - minRuns, - } = options; - const input = optimization.input; - const objectiveMetric = input.model.definition.metrics?.find( - (metric) => metric.id === input.objective.metricId, - ); - if (!objectiveMetric) { - return null; - } - - const sliceEntries = new Map( - slice - .split("|") - .filter((entry) => entry !== "") - .map((entry) => entry.split("=") as [string, string]), - ); - - const values = createUserKeyedRecord(); - for (const [identifier, binding] of Object.entries( - input.scenario.parameterBindings, - )) { - if (binding.kind === "fixed") { - values[identifier] = binding.value; - } else if (binding.domain.kind === "boolean") { - values[identifier] = sliceEntries.get(identifier) === "true"; - } - } - for (const axis of axes) { - const position = - axis.identifier === xAxisId - ? xPosition - : axis.identifier === yAxisId - ? yPosition - : Number(sliceEntries.get(axis.identifier) ?? 0); - values[axis.identifier] = optimizationAxisValueAt(axis, position); - } - - const key = `${slice}|x=${xPosition}|y=${yPosition}`; - const pending = cache.get(key); - const settled = (async (): Promise => { - const cached = await pending; - if (cached && cached.runsCompleted >= minRuns) { - return cached; - } - const from = cached?.runsCompleted ?? 0; - const snapshot = await sampleDetachedObjective({ - cacheKey: optimization.id, - definition: input.model.definition, - scenarioId: input.scenario.id, - scenarioParameterValues: values, - metric: { - id: objectiveMetric.id, - label: objectiveMetric.name, - code: objectiveMetric.code, - }, - seed: sweepBatchSeed(input.execution.seed, from), - runCount: minRuns - from, - dt: input.execution.dt, - maxTime: input.execution.maxTime, - }); - if (!snapshot) { - return cached ?? null; - } - return { - runsCompleted: minRuns, - metricFrames: cached - ? mergeMetricFramesAcrossCells([ - cached.metricFrames, - snapshot.metricFrames, - ]) - : snapshot.metricFrames, - }; - })(); - cache.set(key, settled); - return await settled; -}; +const initialView = ( + axes: readonly OptimizationSurfaceAxis[], +): OptimizationSurfaceView => ({ + xAxisId: axes[0]?.identifier ?? "", + yAxisId: axes[1]?.identifier ?? "", +}); export const OptimizationSurface = ({ optimization, @@ -216,10 +127,10 @@ export const OptimizationSurface = ({ (metric) => metric.id === metricId, ); - const [xAxisId, setXAxisId] = useState(axes[0]?.identifier ?? ""); - const [yAxisId, setYAxisId] = useState(axes[1]?.identifier ?? ""); - const [positions, setPositions] = useState>({}); - const [preview, setPreview] = useState(null); + const [view, setView] = useState(() => initialView(axes)); + const [chosenPositions, setChosenPositions] = useState< + Record + >({}); /** * The selected point's refinement so far, tagged with its walk: the current * point's stats, and every grid cell a selection has refined within this @@ -230,93 +141,36 @@ export const OptimizationSurface = ({ stats: DistributionStats | null; cells: ReadonlyMap; } | null>(null); - /** Per position tuple, the promise of its deepest merged result. */ - const cellCacheRef = useRef( - new Map>(), - ); - - const xAxis = axes.find((axis) => axis.identifier === xAxisId); - const yAxis = axes.find((axis) => axis.identifier === yAxisId); + const cellCacheRef = useRef(new Map()); - /** Slider position per axis: explicit, else best trial, else midpoint. */ - const positionOf = (axis: OptimizationSurfaceAxis): number => { - const explicit = getOwn(positions, axis.identifier); - if (explicit !== undefined) { - return explicit; - } - const best = optimization.best?.parameters[axis.identifier]; - return typeof best === "number" - ? optimizationAxisPositionFor(axis, best) - : optimizationAxisMidpoint(axis); - }; - - // The off-surface coordinates: slider positions of the other axes, plus - // boolean bindings at the best trial's values. Part of the walk key, so a - // best-trial change that moves a boolean restarts the walk rather than - // mixing slices. - const booleanSlice = Object.entries(input.scenario.parameterBindings) - .filter( - ( - entry, - ): entry is [string, { kind: "optimize"; domain: { kind: "boolean" } }] => - entry[1].kind === "optimize" && entry[1].domain.kind === "boolean", - ) - .map(([identifier]) => { - const best = optimization.best?.parameters[identifier]; - return `${identifier}=${typeof best === "boolean" ? best : false}`; - }); - const slice = [ - ...axes - .filter( - (axis) => axis.identifier !== xAxisId && axis.identifier !== yAxisId, - ) - .map((axis) => `${axis.identifier}=${positionOf(axis)}`), - ...booleanSlice, - ].join("|"); - const walkKey = `${optimization.id}|${xAxisId}|${yAxisId}|${slice}`; - - const xSelected = xAxis ? positionOf(xAxis) : 0; - const ySelected = yAxis ? positionOf(yAxis) : 0; + const positions = resolveSurfacePositions( + axes, + chosenPositions, + optimization.best, + ); + const booleans = resolveSurfaceBooleans( + optimizationBooleanIdentifiers(input), + {}, + optimization.best, + ); + const slice = surfaceSliceKey({ axes, view, positions, booleans }); + const walkKey = surfaceWalkKey(optimization.id, view, slice); + const xSelected = positions[view.xAxisId] ?? 0; + const ySelected = positions[view.yAxisId] ?? 0; + const xAxis = axes.find((axis) => axis.identifier === view.xAxisId); + const yAxis = axes.find((axis) => axis.identifier === view.yAxisId); - // The sampler is serialised, so one lane of single-cell chunks. - const walkValues = useSurfaceWalk({ - walkKey, - lanes: 1, - buildWalk: () => { - if (!xAxis || !yAxis || xAxis === yAxis) { - return null; - } - const xPositions = surfacePositions(xAxis); - const yPositions = surfacePositions(yAxis); - return { - chunks: quadTreeLevels(xPositions.length, yPositions.length) - .flat() - .map((cell) => [cell]), - sample: (chunk) => - Promise.all( - chunk.map(async (cell) => { - const snapshot = await sampleStudyCell({ - sampleDetachedObjective, - cache: cellCacheRef.current, - optimization, - axes, - xAxisId, - yAxisId, - slice, - xPosition: xPositions[cell.x]!, - yPosition: yPositions[cell.y]!, - minRuns: SURFACE_CELL_RUNS, - }); - return snapshot - ? sweepCellObjective(snapshot.metricFrames, metricId) - : null; - }), - ), - }; - }, + const walkValues = useStudySurfaceWalk({ + sampleDetachedObjective, + cellCache: cellCacheRef, + optimization, + axes, + view, + slice, }); const optimizationId = optimization.id; + const { xAxisId, yAxisId } = view; // The selected point's refinement: escalating batches, streaming the // objective's mean/median into the readout and refreshing its grid cell. useEffect(() => { @@ -332,10 +186,12 @@ export const OptimizationSurface = ({ } let stale = false; const isStale = () => stale; - const xIndex = surfacePositions(walkXAxis).indexOf(xSelected); - const yIndex = surfacePositions(walkYAxis).indexOf(ySelected); - const cellKey = - xIndex === -1 || yIndex === -1 ? null : contourSurfaceKey(xIndex, yIndex); + const cellKey = surfaceCellKeyAt( + walkXAxis, + walkYAxis, + xSelected, + ySelected, + ); const run = async () => { for (const target of EXPERIMENT_RUN_LADDER) { @@ -387,77 +243,55 @@ export const OptimizationSurface = ({ sampleDetachedObjective, ]); + if (axes.length < 2 || !objectiveMetric) { + return null; + } + const currentRefined = refined?.walkKey === walkKey ? refined : null; - // A selected point is usually also a grid cell: its refined value wins. + const stats = currentRefined?.stats; + const direction = + input.objective.direction === "maximize" ? "Maximize" : "Minimize"; + + // A refined point is usually also a grid cell: its deeper value wins. const cellValues = currentRefined && currentRefined.cells.size > 0 ? new Map([...walkValues, ...currentRefined.cells]) : walkValues; - - /** Completed trials projected onto the shown axes, as ring markers. */ - const trialMarkers: ContourSurfaceMarker[] = + const markers = xAxis && yAxis - ? optimization.trials - .filter( - (trial) => trial.state === "complete" && trial.objective !== null, - ) - .map((trial) => { - const xValue = trial.parameters[xAxis.identifier]; - const yValue = trial.parameters[yAxis.identifier]; - if (typeof xValue !== "number" || typeof yValue !== "number") { - return null; - } - return { - x: - (optimizationAxisPositionFor(xAxis, xValue) / xAxis.stepCount) * - (surfacePositions(xAxis).length - 1), - y: - (optimizationAxisPositionFor(yAxis, yValue) / yAxis.stepCount) * - (surfacePositions(yAxis).length - 1), - emphasis: optimization.best?.trial === trial.trial, - }; - }) - .filter((marker) => marker !== null) + ? trialSurfaceField({ + trials: optimization.trials, + best: optimization.best, + xAxis, + yAxis, + mark: "ring", + }).markers : []; - - if (axes.length < 2 || !objectiveMetric) { - return null; - } - - const handlePickFraction = (fraction: ContourSurfaceFraction) => { - if (!xAxis || !yAxis) { - return; - } - setPositions((previous) => ({ - ...previous, - [xAxis.identifier]: Math.round(fraction.x * xAxis.stepCount), - [yAxis.identifier]: Math.round(fraction.y * yAxis.stepCount), - })); - }; - - /** The axis readout a plot fraction lands on. */ - const readoutAt = (axis: OptimizationSurfaceAxis, fraction: number): string => - `${axis.identifier} = ${formatAxisValue( - optimizationAxisValueAt(axis, Math.round(fraction * axis.stepCount)), - )}`; - - const direction = - input.objective.direction === "maximize" ? "Maximize" : "Minimize"; - const stats = currentRefined?.stats; const totalCells = xAxis && yAxis ? surfacePositions(xAxis).length * surfacePositions(yAxis).length : 0; return ( - - + + setChosenPositions((previous) => ({ ...previous, ...picked })) + } + caption={describeSurfaceSampling({ + sampledCount: cellValues.size, + totalCells, + runsPerCell: SURFACE_CELL_RUNS, + note: "rings are the study's trials (best highlighted), the ringed dot the current parameters", + })} + > {axes.map((axis) => (
@@ -468,16 +302,18 @@ export const OptimizationSurface = ({ min={0} max={axis.stepCount} step={1} - value={positionOf(axis)} + value={positions[axis.identifier]} onChangeEnd={(position) => - setPositions((previous) => ({ + setChosenPositions((previous) => ({ ...previous, [axis.identifier]: position, })) } /> - {formatAxisValue(optimizationAxisValueAt(axis, positionOf(axis)))} + {formatAxisValue( + optimizationAxisValueAt(axis, positions[axis.identifier] ?? 0), + )}
))} @@ -489,29 +325,96 @@ export const OptimizationSurface = ({ )} median · ${stats.runs} runs` : "computing…"}
- {xAxis && yAxis ? ( - - ) : null} - - + + ); +}; + +export const NavigatedOptimizationSurface = ({ + optimization, + navigation, + selection, + onNavigationChange, + controls, +}: { + optimization: OptimizationRecord; + navigation: OptimizationNavigation; + /** The provider's stream at the navigated point (or the followed step). */ + selection: OptimizationSelectionStream | null; + onNavigationChange: (patch: Partial) => void; + /** Further controls at the end of the axis row, e.g. a help tooltip. */ + controls?: ReactNode; +}) => { + const input = optimization.input; + const axes = optimization.axes; + const [view, setView] = useState(() => initialView(axes)); + + const positions = resolveSurfacePositions( + axes, + navigation.positions, + optimization.best, + ); + const booleans = resolveSurfaceBooleans( + optimizationBooleanIdentifiers(input), + navigation.booleans, + optimization.best, + ); + const xAxis = axes.find((axis) => axis.identifier === view.xAxisId); + const yAxis = axes.find((axis) => axis.identifier === view.yAxisId); + + if (axes.length < 2 || !xAxis || !yAxis) { + return null; + } + + const field = mergeSurfaceFields( + trialSurfaceField({ + trials: optimization.trials, + best: optimization.best, + xAxis, + yAxis, + mark: "dot", + }), + inFlightSurfaceField({ inFlight: optimization.inFlight, xAxis, yAxis }), + ); + const values = withNavigatedSample( + field.values, + navigatedSurfaceSample({ + selection, + trials: optimization.trials, + metricId: input.objective.metricId, + xAxis, + yAxis, + positions, + }), + ); + const interaction = surfaceInteraction(optimization, navigation); + const slice = surfaceSliceKey({ axes, view, positions, booleans }); + + return ( + + onNavigationChange({ + positions: { ...navigation.positions, ...picked }, + followTrials: false, + }) + : undefined + } + caption={describeSurfaceState({ + trials: optimization.trials, + best: optimization.best, + interaction, + selection, + })} + controls={controls} + /> ); }; diff --git a/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/optimization-surface/navigation-slice.ts b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/optimization-surface/navigation-slice.ts new file mode 100644 index 00000000000..e515f5feee1 --- /dev/null +++ b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/optimization-surface/navigation-slice.ts @@ -0,0 +1,98 @@ +/** + * Where a study's surface is looked at: a position per numeric axis and a + * value per boolean parameter, with gaps in an explicit navigation filled + * from the best trial and then the domain midpoint; and the keys naming one + * X/Y view's off-surface coordinates. + */ +import { getOwn } from "@hashintel/petrinaut-core"; + +import { + optimizationAxisMidpoint, + optimizationAxisPositionFor, +} from "../../../../../../../react/optimizations/surface-grid"; + +import type { OptimizationBest } from "../../../../../../../react/optimizations/context"; +import type { OptimizationSurfaceAxis } from "../../../../../../../react/optimizations/surface-grid"; + +/** The two axes a surface shows. */ +export type OptimizationSurfaceView = { xAxisId: string; yAxisId: string }; + +/** A position per axis: explicit, else the best trial's value, else the midpoint. */ +export const resolveSurfacePositions = ( + axes: readonly OptimizationSurfaceAxis[], + explicit: Readonly>, + best: OptimizationBest | null, +): Record => { + const positions: Record = {}; + for (const axis of axes) { + const chosen = getOwn(explicit, axis.identifier); + if (chosen !== undefined) { + positions[axis.identifier] = chosen; + continue; + } + const bestValue = best?.parameters[axis.identifier]; + positions[axis.identifier] = + typeof bestValue === "number" + ? optimizationAxisPositionFor(axis, bestValue) + : optimizationAxisMidpoint(axis); + } + return positions; +}; + +/** A value per boolean parameter: explicit, else the best trial's, else false. */ +export const resolveSurfaceBooleans = ( + identifiers: readonly string[], + explicit: Readonly>, + best: OptimizationBest | null, +): Record => { + const booleans: Record = {}; + for (const identifier of identifiers) { + const chosen = getOwn(explicit, identifier); + if (chosen !== undefined) { + booleans[identifier] = chosen; + continue; + } + const bestValue = best?.parameters[identifier]; + booleans[identifier] = typeof bestValue === "boolean" ? bestValue : false; + } + return booleans; +}; + +/** + * The off-surface coordinates of one view: `identifier=position` per hidden + * axis, then `identifier=true|false` per boolean, joined with `|`. Part of + * the walk key, so a move on a hidden axis or a boolean restarts the walk + * rather than mixing slices. + */ +export const surfaceSliceKey = ({ + axes, + view, + positions, + booleans, +}: { + axes: readonly OptimizationSurfaceAxis[]; + view: OptimizationSurfaceView; + positions: Readonly>; + booleans: Readonly>; +}): string => + [ + ...axes + .filter( + (axis) => + axis.identifier !== view.xAxisId && axis.identifier !== view.yAxisId, + ) + .map( + (axis) => + `${axis.identifier}=${positions[axis.identifier] ?? optimizationAxisMidpoint(axis)}`, + ), + ...Object.entries(booleans).map( + ([identifier, value]) => `${identifier}=${value}`, + ), + ].join("|"); + +/** Identity of one sampled slice of one study's surface. */ +export const surfaceWalkKey = ( + optimizationId: string, + view: OptimizationSurfaceView, + slice: string, +): string => `${optimizationId}|${view.xAxisId}|${view.yAxisId}|${slice}`; diff --git a/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/optimization-surface/sample-study-cell.ts b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/optimization-surface/sample-study-cell.ts new file mode 100644 index 00000000000..f3f3359bcbc --- /dev/null +++ b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/optimization-surface/sample-study-cell.ts @@ -0,0 +1,121 @@ +import { createUserKeyedRecord } from "@hashintel/petrinaut-core"; + +import { mergeMetricFramesAcrossCells } from "../../../../../../../react/experiments/parameter-grid"; +import { sweepBatchSeed } from "../../../../../../../react/experiments/sweep-session"; +import { optimizationAxisValueAt } from "../../../../../../../react/optimizations/surface-grid"; + +import type { ExperimentsContextValue } from "../../../../../../../react/experiments/context"; +import type { SweepCellSnapshot } from "../../../../../../../react/experiments/sweep-session"; +import type { OptimizationRecord } from "../../../../../../../react/optimizations/context"; +import type { OptimizationSurfaceAxis } from "../../../../../../../react/optimizations/surface-grid"; + +/** + * Per position tuple, the promise of a cell's deepest merged result. A cell's + * entry is a promise so the walk and a selected point's refinement queue + * behind each other instead of both sampling from the same run index. + */ +export type StudyCellCache = Map>; + +/** + * Brings one cell up to at least `minRuns` locally computed runs, merging + * batches into `cache`. + */ +export const sampleStudyCell = async (options: { + sampleDetachedObjective: ExperimentsContextValue["sampleDetachedObjective"]; + cache: StudyCellCache; + optimization: Pick; + axes: readonly OptimizationSurfaceAxis[]; + xAxisId: string; + yAxisId: string; + /** Position per off-surface axis and value per boolean, as a slice key. */ + slice: string; + xPosition: number; + yPosition: number; + minRuns: number; +}): Promise => { + const { + sampleDetachedObjective, + cache, + optimization, + axes, + xAxisId, + yAxisId, + slice, + xPosition, + yPosition, + minRuns, + } = options; + const input = optimization.input; + const objectiveMetric = input.model.definition.metrics?.find( + (metric) => metric.id === input.objective.metricId, + ); + if (!objectiveMetric) { + return null; + } + + const sliceEntries = new Map( + slice + .split("|") + .filter((entry) => entry !== "") + .map((entry) => entry.split("=") as [string, string]), + ); + + const values = createUserKeyedRecord(); + for (const [identifier, binding] of Object.entries( + input.scenario.parameterBindings, + )) { + if (binding.kind === "fixed") { + values[identifier] = binding.value; + } else if (binding.domain.kind === "boolean") { + values[identifier] = sliceEntries.get(identifier) === "true"; + } + } + for (const axis of axes) { + const position = + axis.identifier === xAxisId + ? xPosition + : axis.identifier === yAxisId + ? yPosition + : Number(sliceEntries.get(axis.identifier) ?? 0); + values[axis.identifier] = optimizationAxisValueAt(axis, position); + } + + const key = `${slice}|x=${xPosition}|y=${yPosition}`; + const pending = cache.get(key); + const settled = (async (): Promise => { + const cached = await pending; + if (cached && cached.runsCompleted >= minRuns) { + return cached; + } + const from = cached?.runsCompleted ?? 0; + const snapshot = await sampleDetachedObjective({ + cacheKey: optimization.id, + definition: input.model.definition, + scenarioId: input.scenario.id, + scenarioParameterValues: values, + metric: { + id: objectiveMetric.id, + label: objectiveMetric.name, + code: objectiveMetric.code, + }, + seed: sweepBatchSeed(input.execution.seed, from), + runCount: minRuns - from, + dt: input.execution.dt, + maxTime: input.execution.maxTime, + }); + if (!snapshot) { + return cached ?? null; + } + return { + runsCompleted: minRuns, + metricFrames: cached + ? mergeMetricFramesAcrossCells([ + cached.metricFrames, + snapshot.metricFrames, + ]) + : snapshot.metricFrames, + }; + })(); + cache.set(key, settled); + return await settled; +}; diff --git a/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/optimization-surface/surface-plot.test.tsx b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/optimization-surface/surface-plot.test.tsx new file mode 100644 index 00000000000..9b3c89e052d --- /dev/null +++ b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/optimization-surface/surface-plot.test.tsx @@ -0,0 +1,435 @@ +/** + * @vitest-environment jsdom + */ +import { cleanup, render, screen } from "@testing-library/react"; +import { afterEach, describe, expect, it, vi } from "vitest"; + +import { buildOptimizationSurfaceAxes } from "../../../../../../../react/optimizations/surface-grid"; +import { contourSurfaceKey } from "../../../../../../components/contour-surface"; +import { + makeOptimizationInput, + optimizedBindingSets, +} from "../optimizations-story-fixtures"; +import { + describeSurfaceState, + inFlightSurfaceField, + mergeSurfaceFields, + navigatedSurfaceSample, + OptimizationSurfacePlot, + surfaceInteraction, + trialSurfaceField, + withNavigatedSample, +} from "./surface-plot"; + +import type { OptimizationSelectionStream } from "../../../../../../../react/optimizations/context"; +import type { + MonteCarloUserDefinedMetricFrame, + PetrinautOptimizationTrialEvent, +} from "@hashintel/petrinaut-core"; + +vi.mock("@hashintel/ds-components", async (importOriginal) => { + const actual = + await importOriginal(); + const Select = ({ + items, + onChange, + value, + "aria-label": ariaLabel, + }: { + items: readonly { value: string; text: string }[]; + onChange: (value: string | null) => void; + value: string | null; + "aria-label"?: string; + }) => ( + + ); + return { ...actual, Select }; +}); + +vi.mock( + "../../../../../../components/contour-surface", + async (importOriginal) => { + const actual = + await importOriginal< + typeof import("../../../../../../components/contour-surface") + >(); + const ContourSurface = ({ + onPickFraction, + onPreviewFraction, + "aria-label": ariaLabel, + }: { + onPickFraction?: unknown; + onPreviewFraction?: unknown; + "aria-label"?: string; + }) => ( + + ); + return { ...actual, ContourSurface }; + }, +); + +afterEach(cleanup); + +const input = makeOptimizationInput(optimizedBindingSets.base); +const metricId = input.objective.metricId; +const axes = buildOptimizationSurfaceAxes(input); +const [xAxis, yAxis] = axes as [(typeof axes)[number], (typeof axes)[number]]; + +// production_rate spans 50..400 and selling_price 20..60, both over 50 +// positions drawn on an 11-point grid: 225 and 40 sit at the grid's centre. +const trial = ( + index: number, + parameters: Record, + objective: number | null, +): PetrinautOptimizationTrialEvent => ({ + type: "trial", + trial: index, + parameters, + objective, + state: objective === null ? "pruned" : "complete", + best: null, + seq: index + 2, +}); + +const trials = [ + trial(0, { production_rate: 225, selling_price: 40 }, 4), + trial(1, { production_rate: 50, selling_price: 20 }, 1), + trial(2, { production_rate: 400, selling_price: 60 }, null), +]; +const best = { trial: 0, parameters: trials[0]!.parameters, objective: 4 }; + +const distributionFrame = ( + bins: readonly (readonly [number, number])[], +): MonteCarloUserDefinedMetricFrame => ({ + metricId, + label: "Profit", + outputType: "distribution", + frameNumber: 1, + time: 1, + bins, + value: null, + frameValue: null, + timeValue: null, + runSampleCount: 4, + timeSampleCount: 4, +}); + +const stream = ( + overrides: Partial, +): OptimizationSelectionStream => ({ + key: "production_rate=25|selling_price=25", + metricFrames: [ + distributionFrame([ + [10, 2], + [20, 2], + ]), + ], + runsCompleted: 4, + runTarget: null, + computing: true, + error: null, + note: null, + ...overrides, +}); + +const centre = { production_rate: 25, selling_price: 25 }; + +describe("trialSurfaceField", () => { + it("samples the field at each trial with an objective and marks pruned trials hollow", () => { + const field = trialSurfaceField({ + trials, + best, + xAxis, + yAxis, + mark: "dot", + }); + + expect([...field.values]).toEqual([ + [contourSurfaceKey(5, 5), 4], + [contourSurfaceKey(0, 0), 1], + ]); + expect(field.markers).toEqual([ + { x: 5, y: 5, kind: "dot", emphasis: true }, + { x: 0, y: 0, kind: "dot", emphasis: false }, + { x: 10, y: 10, kind: "muted" }, + ]); + }); + + it("draws rings instead of dots for a field computed elsewhere, leaving pruned trials out", () => { + const field = trialSurfaceField({ + trials, + best: null, + xAxis, + yAxis, + mark: "ring", + }); + + expect(field.markers.map((marker) => marker.kind)).toEqual([ + "point", + "point", + ]); + expect(field.markers.every((marker) => marker.emphasis !== true)).toBe( + true, + ); + }); + + it("skips a trial without a numeric value on a shown axis", () => { + const field = trialSurfaceField({ + trials: [trial(0, { production_rate: 225 }, 4)], + best: null, + xAxis, + yAxis, + mark: "dot", + }); + + expect(field.values.size).toBe(0); + expect(field.markers).toHaveLength(0); + }); +}); + +describe("inFlightSurfaceField", () => { + it("rings every step being evaluated and samples the field where one has a running value", () => { + const field = inFlightSurfaceField({ + inFlight: [ + { + trial: 3, + parameters: { production_rate: 225, selling_price: 40 }, + objective: 2.5, + }, + { + trial: 4, + parameters: { production_rate: 50, selling_price: 60 }, + objective: null, + }, + ], + xAxis, + yAxis, + }); + + expect([...field.values]).toEqual([[contourSurfaceKey(5, 5), 2.5]]); + expect(field.markers).toEqual([ + { x: 5, y: 5, kind: "point" }, + { x: 0, y: 10, kind: "point" }, + ]); + }); + + it("merges beneath the trials' field, the later value winning at a shared point", () => { + const merged = mergeSurfaceFields( + { values: new Map([[contourSurfaceKey(5, 5), 4]]), markers: [] }, + { + values: new Map([[contourSurfaceKey(5, 5), 2.5]]), + markers: [{ x: 5, y: 5, kind: "point" }], + }, + ); + + expect(merged.values.get(contourSurfaceKey(5, 5))).toBe(2.5); + expect(merged.markers).toHaveLength(1); + }); +}); + +describe("navigatedSurfaceSample", () => { + it("streams the followed step's running objective at the navigation until its event lands", () => { + const following = stream({ key: "trial:3" }); + + expect( + navigatedSurfaceSample({ + selection: following, + trials, + metricId, + xAxis, + yAxis, + positions: centre, + }), + ).toEqual({ x: 5, y: 5, value: 15 }); + + expect( + navigatedSurfaceSample({ + selection: following, + trials: [ + ...trials, + trial(3, { production_rate: 225, selling_price: 40 }, 15.2), + ], + metricId, + xAxis, + yAxis, + positions: centre, + }), + ).toBeNull(); + }); + + it("places the refined value at the picked point once the study has settled", () => { + expect( + navigatedSurfaceSample({ + selection: stream({ computing: false, runsCompleted: 100 }), + trials, + metricId, + xAxis, + yAxis, + positions: { production_rate: 50, selling_price: 0 }, + }), + ).toEqual({ x: 10, y: 0, value: 15 }); + }); + + it("has no value before frames arrive, on a failed point, or without a stream", () => { + const arguments_ = { trials, metricId, xAxis, yAxis, positions: centre }; + + expect( + navigatedSurfaceSample({ ...arguments_, selection: null }), + ).toBeNull(); + expect( + navigatedSurfaceSample({ + ...arguments_, + selection: stream({ metricFrames: [] }), + }), + ).toBeNull(); + expect( + navigatedSurfaceSample({ + ...arguments_, + selection: stream({ error: "cpu: unsupported net" }), + }), + ).toBeNull(); + }); +}); + +describe("withNavigatedSample", () => { + it("lays the live sample over the trials' field, replacing a value at the same point", () => { + const values = new Map([ + [contourSurfaceKey(5, 5), 4], + [contourSurfaceKey(0, 0), 1], + ]); + + expect(withNavigatedSample(values, null)).toBe(values); + expect([...withNavigatedSample(values, { x: 2, y: 3, value: 9 })]).toEqual([ + [contourSurfaceKey(5, 5), 4], + [contourSurfaceKey(0, 0), 1], + [contourSurfaceKey(2, 3), 9], + ]); + expect( + withNavigatedSample(values, { x: 5, y: 5, value: 4.5 }).get( + contourSurfaceKey(5, 5), + ), + ).toBe(4.5); + }); +}); + +describe("surfaceInteraction", () => { + it("only displays while a running study is followed, and navigates otherwise", () => { + expect( + surfaceInteraction({ status: "running" }, { followTrials: true }), + ).toBe("following"); + expect( + surfaceInteraction({ status: "initializing" }, { followTrials: true }), + ).toBe("following"); + expect( + surfaceInteraction({ status: "running" }, { followTrials: false }), + ).toBe("navigable"); + expect( + surfaceInteraction({ status: "complete" }, { followTrials: true }), + ).toBe("navigable"); + expect( + surfaceInteraction({ status: "cancelled" }, { followTrials: true }), + ).toBe("navigable"); + }); +}); + +describe("describeSurfaceState", () => { + it("counts the placed steps and the best while the optimizer chooses", () => { + expect( + describeSurfaceState({ + trials: [], + best: null, + interaction: "following", + selection: null, + }), + ).toBe("no steps placed yet · the optimizer is choosing the next point"); + expect( + describeSurfaceState({ + trials, + best, + interaction: "following", + selection: stream({ key: "trial:3" }), + }), + ).toBe( + "3 steps placed · best 4 · the optimizer is choosing the next point", + ); + }); + + it("reports the picked point's refinement, then invites a pick", () => { + expect( + describeSurfaceState({ + trials, + best, + interaction: "navigable", + selection: stream({ runsCompleted: 8, runTarget: 25 }), + }), + ).toBe( + "3 steps · refining the picked point: 8 of 25 runs · drag or click to refine a point", + ); + expect( + describeSurfaceState({ + trials, + best, + interaction: "navigable", + selection: stream({ runsCompleted: 8 }), + }), + ).toBe( + "3 steps · refining the picked point: 8 runs · drag or click to refine a point", + ); + expect( + describeSurfaceState({ + trials: trials.slice(0, 1), + best, + interaction: "navigable", + selection: stream({ computing: false, runsCompleted: 100 }), + }), + ).toBe("1 step · drag or click to refine a point"); + }); +}); + +describe("OptimizationSurfacePlot", () => { + const renderPlot = ( + onPick: ((picked: Record) => void) | undefined, + ) => + render( + {}} + positions={centre} + values={new Map()} + markers={[]} + sampleMarks="none" + contentKey="study" + onPick={onPick} + caption="3 steps" + />, + ); + + it("is display-only without a pick handler and arms picks and previews with one", () => { + const { unmount } = renderPlot(undefined); + const passive = screen.getByLabelText("Optimization surface"); + expect(passive.hasAttribute("data-interactive")).toBe(false); + expect(passive.hasAttribute("data-previews")).toBe(false); + expect(screen.getByText("3 steps")).toBeTruthy(); + unmount(); + + renderPlot(() => {}); + const active = screen.getByLabelText("Optimization surface"); + expect(active.hasAttribute("data-interactive")).toBe(true); + expect(active.hasAttribute("data-previews")).toBe(true); + }); +}); diff --git a/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/optimization-surface/surface-plot.tsx b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/optimization-surface/surface-plot.tsx new file mode 100644 index 00000000000..036e416bc24 --- /dev/null +++ b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/optimization-surface/surface-plot.tsx @@ -0,0 +1,387 @@ +/** + * The plot of a study's surface: the X/Y axis selects, a contour over the + * field the owner hands in, the study's trials as markers, the navigation + * marker where the parameters are, and a caption. The pure helpers beside it + * turn a study's trials and its live selection stream into that field, and + * decide whether the plot navigates or only displays. + */ +import { type ReactNode, useState } from "react"; + +import { sweepCellObjective } from "../../../../../../../react/experiments/sweep-cell-objective"; +import { + followedTrial, + isOptimizationActive, +} from "../../../../../../../react/optimizations/context"; +import { + optimizationAxisPositionFor, + optimizationAxisValueAt, +} from "../../../../../../../react/optimizations/surface-grid"; +import { + ContourSurface, + contourSurfaceKey, +} from "../../../../../../components/contour-surface"; +import { formatAxisValue } from "../../shared/format-axis-value"; +import { + SurfaceAxisControls, + SurfaceCaption, + SurfaceFrame, +} from "../../shared/surface-frame"; +import { surfacePositions } from "../../shared/surface-sampling"; + +import type { + OptimizationBest, + OptimizationInFlightStep, + OptimizationNavigation, + OptimizationRecord, + OptimizationSelectionStream, +} from "../../../../../../../react/optimizations/context"; +import type { OptimizationSurfaceAxis } from "../../../../../../../react/optimizations/surface-grid"; +import type { + ContourSurfaceFraction, + ContourSurfaceMarker, + ContourSurfaceSampleMarks, + ContourSurfaceValues, +} from "../../../../../../components/contour-surface"; +import type { OptimizationSurfaceView } from "./navigation-slice"; +import type { PetrinautOptimizationTrialEvent } from "@hashintel/petrinaut-core"; + +/** Grid-index coordinate of an axis position, fractional between samples. */ +export const surfaceGridCoordinate = ( + axis: OptimizationSurfaceAxis, + position: number, +): number => (position / axis.stepCount) * (surfacePositions(axis).length - 1); + +/** The sampled cell an axis position pair lands on, or null between cells. */ +export const surfaceCellKeyAt = ( + xAxis: OptimizationSurfaceAxis, + yAxis: OptimizationSurfaceAxis, + xPosition: number, + yPosition: number, +): string | null => { + const xIndex = surfacePositions(xAxis).indexOf(xPosition); + const yIndex = surfacePositions(yAxis).indexOf(yPosition); + return xIndex === -1 || yIndex === -1 + ? null + : contourSurfaceKey(xIndex, yIndex); +}; + +/** One point of the field, in grid-index space. */ +export type SurfaceSample = { x: number; y: number; value: number }; + +/** A study's trials as a field: a sample per objective, a marker per trial. */ +export type TrialSurfaceField = { + values: ReadonlyMap; + markers: readonly ContourSurfaceMarker[]; +}; + +/** How a trial with an objective is drawn. */ +export type TrialSurfaceMark = "ring" | "dot"; + +/** + * Projects the trials onto the shown axes. A trial with an objective is a + * sample of the field and a mark — a ring over a field computed elsewhere, a + * filled dot when the trials are the field's only samples — the best + * emphasized. A trial without one — pruned or failed — is no sample; among + * dots it is a muted ring, among rings it is absent. + */ +export const trialSurfaceField = ({ + trials, + best, + xAxis, + yAxis, + mark, +}: { + trials: readonly PetrinautOptimizationTrialEvent[]; + best: OptimizationBest | null; + xAxis: OptimizationSurfaceAxis; + yAxis: OptimizationSurfaceAxis; + mark: TrialSurfaceMark; +}): TrialSurfaceField => { + const values = new Map(); + const markers: ContourSurfaceMarker[] = []; + for (const trial of trials) { + const xValue = trial.parameters[xAxis.identifier]; + const yValue = trial.parameters[yAxis.identifier]; + if (typeof xValue !== "number" || typeof yValue !== "number") { + continue; + } + const x = surfaceGridCoordinate( + xAxis, + optimizationAxisPositionFor(xAxis, xValue), + ); + const y = surfaceGridCoordinate( + yAxis, + optimizationAxisPositionFor(yAxis, yValue), + ); + if (trial.objective === null) { + if (mark === "dot") { + markers.push({ x, y, kind: "muted" }); + } + continue; + } + values.set(contourSurfaceKey(x, y), trial.objective); + markers.push({ + x, + y, + kind: mark === "dot" ? "dot" : "point", + emphasis: best?.trial === trial.trial, + }); + } + return { values, markers }; +}; + +/** + * The steps being evaluated, projected onto the shown axes: each is a ring + * where the optimizer is looking, and one with a running objective is a + * sample of the field too, so the surface fills in as its runs complete. + */ +export const inFlightSurfaceField = ({ + inFlight, + xAxis, + yAxis, +}: { + inFlight: readonly OptimizationInFlightStep[]; + xAxis: OptimizationSurfaceAxis; + yAxis: OptimizationSurfaceAxis; +}): TrialSurfaceField => { + const values = new Map(); + const markers: ContourSurfaceMarker[] = []; + for (const step of inFlight) { + const xValue = step.parameters[xAxis.identifier]; + const yValue = step.parameters[yAxis.identifier]; + if (typeof xValue !== "number" || typeof yValue !== "number") { + continue; + } + const x = surfaceGridCoordinate( + xAxis, + optimizationAxisPositionFor(xAxis, xValue), + ); + const y = surfaceGridCoordinate( + yAxis, + optimizationAxisPositionFor(yAxis, yValue), + ); + if (step.objective !== null) { + values.set(contourSurfaceKey(x, y), step.objective); + } + markers.push({ x, y, kind: "point" }); + } + return { values, markers }; +}; + +/** The objective's running value on a selection stream; null before it has one. */ +const selectionSurfaceValue = ( + selection: OptimizationSelectionStream | null, + metricId: string, +): number | null => + selection === null || selection.error !== null + ? null + : sweepCellObjective(selection.metricFrames, metricId); + +/** + * The navigated point's live sample: the followed trial's running objective + * until its own event lands, then whatever point the navigation refines. + * Null while the stream has no value. + */ +export const navigatedSurfaceSample = ({ + selection, + trials, + metricId, + xAxis, + yAxis, + positions, +}: { + selection: OptimizationSelectionStream | null; + trials: readonly PetrinautOptimizationTrialEvent[]; + metricId: string; + xAxis: OptimizationSurfaceAxis; + yAxis: OptimizationSurfaceAxis; + positions: Readonly>; +}): SurfaceSample | null => { + if (selection === null) { + return null; + } + const value = selectionSurfaceValue(selection, metricId); + if (value === null) { + return null; + } + const trial = followedTrial(selection.key); + if (trial !== null && trials.some((event) => event.trial === trial)) { + return null; + } + return { + x: surfaceGridCoordinate(xAxis, positions[xAxis.identifier] ?? 0), + y: surfaceGridCoordinate(yAxis, positions[yAxis.identifier] ?? 0), + value, + }; +}; + +/** The trials' field with the live sample laid over it. */ +export const withNavigatedSample = ( + values: ReadonlyMap, + sample: SurfaceSample | null, +): ReadonlyMap => + sample === null + ? values + : new Map([ + ...values, + [contourSurfaceKey(sample.x, sample.y), sample.value], + ]); + +/** Fields laid over one another; a later field's value wins at a shared point. */ +export const mergeSurfaceFields = ( + ...fields: readonly TrialSurfaceField[] +): TrialSurfaceField => ({ + values: new Map(fields.flatMap((field) => [...field.values])), + markers: fields.flatMap((field) => field.markers), +}); + +/** + * Whether the plot navigates. While the study runs and the navigation follows + * its steps, the optimizer chooses the points and the plot only displays; + * otherwise a click or drag picks a point. + */ +export type SurfaceInteraction = "following" | "navigable"; + +export const surfaceInteraction = ( + optimization: Pick, + navigation: Pick, +): SurfaceInteraction => + isOptimizationActive(optimization) && navigation.followTrials + ? "following" + : "navigable"; + +/** The caption's state line for a connected study's surface. */ +export const describeSurfaceState = ({ + trials, + best, + interaction, + selection, +}: { + trials: readonly PetrinautOptimizationTrialEvent[]; + best: OptimizationBest | null; + interaction: SurfaceInteraction; + selection: OptimizationSelectionStream | null; +}): string => { + const count = trials.length; + const steps = `${count} ${count === 1 ? "step" : "steps"}`; + if (interaction === "following") { + return [ + count === 0 ? "no steps placed yet" : `${steps} placed`, + ...(best === null ? [] : [`best ${formatAxisValue(best.objective)}`]), + "the optimizer is choosing the next point", + ].join(" · "); + } + const refining = + selection !== null && + selection.computing && + followedTrial(selection.key) === null + ? selection.runTarget === null + ? `refining the picked point: ${selection.runsCompleted} runs` + : `refining the picked point: ${selection.runsCompleted} of ${selection.runTarget} runs` + : null; + return [ + steps, + ...(refining === null ? [] : [refining]), + "drag or click to refine a point", + ].join(" · "); +}; + +export const OptimizationSurfacePlot = ({ + axes, + view, + onViewChange, + positions, + values, + markers, + sampleMarks, + contentKey, + onPick, + caption, + controls, + children, +}: { + axes: readonly OptimizationSurfaceAxis[]; + view: OptimizationSurfaceView; + onViewChange: (view: OptimizationSurfaceView) => void; + /** A position per axis; the navigation marker sits at the shown pair. */ + positions: Readonly>; + values: ContourSurfaceValues; + /** The data markers; the navigation marker is added here. */ + markers: readonly ContourSurfaceMarker[]; + sampleMarks: ContourSurfaceSampleMarks; + /** Identity of the plotted field; a change drops the dimmed previous picture. */ + contentKey: string; + /** + * The X and Y positions a click or drag on the plot picked. Undefined + * makes the plot display-only. + */ + onPick: ((positions: Record) => void) | undefined; + /** The state line under the plot, outside a drag. */ + caption: string; + /** Further controls at the end of the axis row. */ + controls?: ReactNode; + /** Rows between the axis selects and the plot. */ + children?: ReactNode; +}) => { + const [preview, setPreview] = useState(null); + const xAxis = axes.find((axis) => axis.identifier === view.xAxisId); + const yAxis = axes.find((axis) => axis.identifier === view.yAxisId); + + const handlePickFraction = + onPick && xAxis && yAxis + ? (fraction: ContourSurfaceFraction) => + onPick({ + [xAxis.identifier]: Math.round(fraction.x * xAxis.stepCount), + [yAxis.identifier]: Math.round(fraction.y * yAxis.stepCount), + }) + : undefined; + + /** The axis readout a plot fraction lands on. */ + const readoutAt = (axis: OptimizationSurfaceAxis, fraction: number): string => + `${axis.identifier} = ${formatAxisValue( + optimizationAxisValueAt(axis, Math.round(fraction * axis.stepCount)), + )}`; + + return ( + + onViewChange({ ...view, xAxisId })} + onYAxisIdChange={(yAxisId) => onViewChange({ ...view, yAxisId })} + > + {controls} + + {children} + {xAxis && yAxis ? ( + + ) : null} + + + ); +}; diff --git a/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/optimization-surface/use-study-surface-walk.ts b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/optimization-surface/use-study-surface-walk.ts new file mode 100644 index 00000000000..6dc5ccc5d7a --- /dev/null +++ b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/optimization-surface/use-study-surface-walk.ts @@ -0,0 +1,81 @@ +/** + * The fill of a remote study's surface: an X×Y sub-grid of the shown axes + * walked in quad-tree order, each cell brought to `SURFACE_CELL_RUNS` runs of + * the study's frozen model through the detached sampler, which is serialised + * — so one lane of single-cell chunks. + */ +import { sweepCellObjective } from "../../../../../../../react/experiments/sweep-cell-objective"; +import { + quadTreeLevels, + SURFACE_CELL_RUNS, + surfacePositions, +} from "../../shared/surface-sampling"; +import { useSurfaceWalk } from "../../shared/use-surface-walk"; +import { + type OptimizationSurfaceView, + surfaceWalkKey, +} from "./navigation-slice"; +import { sampleStudyCell, type StudyCellCache } from "./sample-study-cell"; + +import type { ExperimentsContextValue } from "../../../../../../../react/experiments/context"; +import type { OptimizationRecord } from "../../../../../../../react/optimizations/context"; +import type { OptimizationSurfaceAxis } from "../../../../../../../react/optimizations/surface-grid"; +import type { RefObject } from "react"; + +export const useStudySurfaceWalk = ({ + sampleDetachedObjective, + cellCache, + optimization, + axes, + view, + slice, +}: { + sampleDetachedObjective: ExperimentsContextValue["sampleDetachedObjective"]; + cellCache: RefObject; + optimization: Pick; + axes: readonly OptimizationSurfaceAxis[]; + view: OptimizationSurfaceView; + /** Position per off-surface axis and value per boolean, as a slice key. */ + slice: string; +}): ReadonlyMap => { + const metricId = optimization.input.objective.metricId; + const xAxis = axes.find((axis) => axis.identifier === view.xAxisId); + const yAxis = axes.find((axis) => axis.identifier === view.yAxisId); + + return useSurfaceWalk({ + walkKey: surfaceWalkKey(optimization.id, view, slice), + lanes: 1, + buildWalk: () => { + if (!xAxis || !yAxis || xAxis === yAxis) { + return null; + } + const xPositions = surfacePositions(xAxis); + const yPositions = surfacePositions(yAxis); + return { + chunks: quadTreeLevels(xPositions.length, yPositions.length) + .flat() + .map((cell) => [cell]), + sample: (chunk) => + Promise.all( + chunk.map(async (cell) => { + const snapshot = await sampleStudyCell({ + sampleDetachedObjective, + cache: cellCache.current, + optimization, + axes, + xAxisId: view.xAxisId, + yAxisId: view.yAxisId, + slice, + xPosition: xPositions[cell.x]!, + yPosition: yPositions[cell.y]!, + minRuns: SURFACE_CELL_RUNS, + }); + return snapshot + ? sweepCellObjective(snapshot.metricFrames, metricId) + : null; + }), + ), + }; + }, + }); +}; diff --git a/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/optimizations-story-fixtures.ts b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/optimizations-story-fixtures.ts index 88f0dc90713..a3c9af679a6 100644 --- a/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/optimizations-story-fixtures.ts +++ b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/optimizations-story-fixtures.ts @@ -1,18 +1,41 @@ /** * Fixtures for the optimization stories: a real study manifest over the * supply-chain example, deterministic fake trials, and the synthetic - * objective both the trials and the stories' fake local compute share — so - * trial rings land on the contour they would on a real study. + * objective the trials, the selection streams and the remote surface's fake + * local compute all share — so a step's mark lands on the contour a real + * study would give. For a connected study, a navigation at a trial's point, + * the selection stream the provider would publish there, and a clock that + * lands one step after another. */ +import { useEffect, useState } from "react"; + import { petrinautOptimizationInputSchema } from "@hashintel/petrinaut-core"; import { supplyChainProfit } from "@hashintel/petrinaut-core/examples"; +import { + buildOptimizationSurfaceAxes, + optimizationAxisPositionFor, + optimizationBooleanIdentifiers, + optimizationNavigationKey, + optimizationNavigationValues, +} from "../../../../../../react/optimizations/surface-grid"; + +import type { + DetachedObjectiveRequest, + ExperimentComputeBackend, +} from "../../../../../../react/experiments/context"; +import type { SweepCellSnapshot } from "../../../../../../react/experiments/sweep-session"; import type { + OptimizationBatchStatus, OptimizationBest, + OptimizationInFlightStep, + OptimizationNavigation, OptimizationRecord, + OptimizationSelectionStream, OptimizationStatus, } from "../../../../../../react/optimizations/context"; import type { + MonteCarloUserDefinedMetricFrame, PetrinautOptimizationInput, PetrinautOptimizationParameterBinding, PetrinautOptimizationTrialEvent, @@ -231,8 +254,32 @@ export function makeOptimizationRecord(options: { trials?: readonly PetrinautOptimizationTrialEvent[]; best?: OptimizationBest | null; status?: OptimizationStatus; + computeBackend?: ExperimentComputeBackend; + computeBackendFallbackReason?: string | null; + /** Set for a connected study; a remote study has neither. */ + navigation?: OptimizationNavigation | null; + selection?: OptimizationSelectionStream | null; + /** Whether the study can be continued; a settled connected study by default. */ + resumable?: boolean; + parallelism?: number; + activity?: readonly OptimizationBatchStatus[]; + inFlight?: readonly OptimizationInFlightStep[]; }): OptimizationRecord { - const { input, trials = [], best = null, status = "running" } = options; + const { + input, + trials = [], + best = null, + status = "running", + computeBackend = "cpu", + computeBackendFallbackReason = null, + navigation = null, + selection = null, + resumable = navigation !== null && + (status === "complete" || status === "cancelled"), + parallelism = 1, + activity = [], + inFlight = [], + } = options; return { id: "optimization-story-1", input, @@ -251,5 +298,231 @@ export function makeOptimizationRecord(options: { failedTrials: trials.filter((trial) => trial.state === "failed").length, trials, best, + resumable, + parallelism, + computeBackend, + computeBackendFallbackReason, + axes: buildOptimizationSurfaceAxes(input), + navigation, + selection, + activity, + inFlight, + }; +} + +/** The navigation at a trial's parameters, following steps while running. */ +export function navigationAtTrial( + input: PetrinautOptimizationInput, + trial: PetrinautOptimizationTrialEvent, + followTrials = true, +): OptimizationNavigation { + const positions: Record = {}; + for (const axis of buildOptimizationSurfaceAxes(input)) { + const value = trial.parameters[axis.identifier]; + positions[axis.identifier] = + typeof value === "number" + ? optimizationAxisPositionFor(axis, value) + : Math.round(axis.stepCount / 2); + } + const booleans: Record = {}; + for (const identifier of optimizationBooleanIdentifiers(input)) { + booleans[identifier] = trial.parameters[identifier] === true; + } + return { positions, booleans, followTrials }; +} + +/** The provider's key for a navigated point. */ +export function navigationKey( + input: PetrinautOptimizationInput, + navigation: OptimizationNavigation, +): string { + return optimizationNavigationKey( + buildOptimizationSurfaceAxes(input), + optimizationBooleanIdentifiers(input), + navigation, + ); +} + +/** + * Distribution frames of the objective at one point, streamed up to + * `frameCount` of the study's time steps: the synthetic profit accrues + * linearly over the year, spread across `runs` runs with a jitter that + * shrinks as runs accumulate — so a refinement visibly sharpens the band. + */ +export function makeObjectiveFrames( + input: PetrinautOptimizationInput, + values: Readonly>, + runs: number, + frameCount = 40, +): MonteCarloUserDefinedMetricFrame[] { + const metric = input.model.definition.metrics?.[0]; + if (!metric) { + throw new Error("The study manifest carries no objective metric"); + } + const final = syntheticObjective(values); + const { maxTime } = input.execution; + const frames: MonteCarloUserDefinedMetricFrame[] = []; + for (let index = 0; index <= frameCount; index++) { + const fraction = index / frameCount; + const time = maxTime * fraction; + const mean = final * fraction; + const spread = Math.max(1, Math.abs(final) * 0.08 * (0.3 + fraction)); + const binCount = Math.min(9, 2 + Math.floor(Math.sqrt(runs))); + const bins: (readonly [number, number])[] = []; + let assigned = 0; + for (let bin = 0; bin < binCount; bin++) { + const offset = ((bin - (binCount - 1) / 2) / (binCount - 1)) * 2; + const weight = Math.exp(-(offset ** 2) * 1.5); + const frequency = + bin === binCount - 1 + ? runs - assigned + : Math.max(0, Math.round((weight * runs) / binCount)); + assigned += frequency; + if (frequency > 0) { + bins.push([ + Math.round((mean + offset * spread) * 100) / 100, + frequency, + ]); + } + } + frames.push({ + metricId: metric.id, + label: metric.name, + outputType: "distribution", + frameNumber: Math.round(time / input.execution.dt), + time, + bins, + value: null, + frameValue: null, + timeValue: null, + runSampleCount: runs, + timeSampleCount: runs, + }); + } + return frames; +} + +/** The selection stream a connected study publishes at a navigated point. */ +export function makeSelectionStream(options: { + input: PetrinautOptimizationInput; + navigation: OptimizationNavigation; + /** Set while following that step: the key becomes the trial's. */ + followedTrial?: number; + runsCompleted: number; + runTarget?: number | null; + computing?: boolean; + frameCount?: number; + /** + * How far through the simulated time the frames have streamed, 0..1: the + * frames stop there, so the running objective reads part-way to its final + * value. Complete when omitted. + */ + progress?: number; + /** Why the point could not compute; the stream then stops at `runsCompleted`. */ + error?: string | null; + /** Why the ladder stopped short, e.g. "8 runs · cannot beat the best". */ + note?: string | null; +}): OptimizationSelectionStream { + const { + input, + navigation, + followedTrial, + runsCompleted, + runTarget = null, + computing = false, + frameCount, + progress, + error = null, + note = null, + } = options; + const axes = buildOptimizationSurfaceAxes(input); + const booleanIdentifiers = optimizationBooleanIdentifiers(input); + const values = optimizationNavigationValues( + input, + axes, + booleanIdentifiers, + navigation, + ); + const frames = makeObjectiveFrames( + input, + values, + Math.max(1, runsCompleted), + frameCount, + ); + return { + key: + followedTrial === undefined + ? optimizationNavigationKey(axes, booleanIdentifiers, navigation) + : `trial:${followedTrial}`, + metricFrames: + progress === undefined + ? frames + : frames.slice(0, Math.max(1, Math.ceil(frames.length * progress))), + runsCompleted, + runTarget, + computing, + error, + note, }; } + +/** + * The stories' clock for a study in flight: `landed` steps have reported and + * the next one is `progress` of the way through its runs. Advances every + * `tickMs`, `ticksPerStep` ticks per step, until all `steps` have landed. + */ +export function useFakeStudyClock({ + steps, + ticksPerStep, + tickMs, +}: { + steps: number; + ticksPerStep: number; + tickMs: number; +}): { landed: number; progress: number } { + const [tick, setTick] = useState(0); + const total = steps * ticksPerStep; + useEffect(() => { + if (tick >= total) { + return; + } + const timer = setTimeout(() => setTick((previous) => previous + 1), tickMs); + return () => clearTimeout(timer); + }, [tick, tickMs, total]); + return { + landed: Math.min(steps, Math.floor(tick / ticksPerStep)), + progress: (tick % ticksPerStep) / ticksPerStep, + }; +} + +/** + * The remote surface stories' local compute: the same synthetic objective + * the fake trials used, returned as a single-bin distribution frame after + * `delayFor` the batch — so the walked contour fills in progressively and + * the trial rings land on it, at whatever pace the story simulates. + */ +export const makeSyntheticObjectiveSampler = + (delayFor: (runCount: number) => number) => + (request: DetachedObjectiveRequest): Promise => { + const objective = syntheticObjective(request.scenarioParameterValues); + const frame: MonteCarloUserDefinedMetricFrame = { + metricId: request.metric.id, + label: request.metric.label, + outputType: "distribution", + frameNumber: 365, + time: 365, + bins: [[Math.round(objective * 100) / 100, request.runCount]], + value: null, + frameValue: null, + timeValue: null, + runSampleCount: request.runCount, + timeSampleCount: request.runCount, + }; + return new Promise((resolve) => { + setTimeout( + () => + resolve({ runsCompleted: request.runCount, metricFrames: [frame] }), + delayFor(request.runCount), + ); + }); + }; diff --git a/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/optimizations-view.tsx b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/optimizations-view.tsx index 30b3742c96f..ae9f836316a 100644 --- a/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/optimizations-view.tsx +++ b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/optimizations-view.tsx @@ -9,23 +9,9 @@ import { import { EditorContext } from "../../../../../../react/state/editor-context"; import { Table, type TableColumn } from "../../../../../components/table"; import { SimulateSubviewFrame } from "../simulate-subview-frame"; +import { describeOptimizationStatus } from "./optimization-status"; import { ViewOptimizationDrawer } from "./view-optimization-drawer"; -function formatStatus(status: OptimizationRecord["status"]): string { - switch (status) { - case "initializing": - return "Initializing"; - case "running": - return "Running"; - case "complete": - return "Complete"; - case "error": - return "Error"; - case "cancelled": - return "Cancelled"; - } -} - const OptimizationStatusBadge = ({ optimization, }: { @@ -51,7 +37,7 @@ const OptimizationStatusBadge = ({ : undefined } > - {formatStatus(optimization.status)} + {describeOptimizationStatus(optimization)} ); }; diff --git a/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/view-optimization-drawer.stories.tsx b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/view-optimization-drawer.stories.tsx new file mode 100644 index 00000000000..09ebda5fbbf --- /dev/null +++ b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/view-optimization-drawer.stories.tsx @@ -0,0 +1,221 @@ +/** + * The study drawer against fake records. For a connected study the navigator + * and the surface follow each step while the study runs — the step in flight + * streams its running objective into the surface before its dot lands — then + * the surface, the controls and the chart move together when the parameters + * are picked by hand; the selection stream is faked from the synthetic + * objective and refines in three batches after every move. + */ +import { useEffect, useState } from "react"; + +import { + type OptimizationNavigation, + OptimizationsContext, + type OptimizationsContextValue, +} from "../../../../../../react/optimizations/context"; +import { FakeExperimentsProvider } from "../experiments/experiments-story-fixtures"; +import { + makeOptimizationInput, + makeOptimizationRecord, + makeSelectionStream, + makeTrials, + navigationAtTrial, + navigationKey, + optimizedBindingSets, + useFakeStudyClock, +} from "./optimizations-story-fixtures"; +import { ViewOptimizationDrawer } from "./view-optimization-drawer"; + +import type { Meta, StoryObj } from "@storybook/react-vite"; + +const meta = { + title: "Simulate / ViewOptimizationDrawer", + parameters: { layout: "fullscreen" }, +} satisfies Meta; + +export default meta; + +type Story = StoryObj; + +const input = makeOptimizationInput(optimizedBindingSets.logScale); +const allTrials = makeTrials(input, 30); + +/** The refinement ladder a navigated point climbs, one rung per 900 ms. */ +const REFINEMENT_LADDER = [8, 25, 100]; + +const FakeConnectedStudy = ({ + running, + fallbackReason = null, + refinementError = null, +}: { + /** Lands one step every 1.2 s and follows the next; else shows the complete study. */ + running: boolean; + fallbackReason?: string | null; + /** Set to have every navigated point fail with this reason instead of refining. */ + refinementError?: string | null; +}) => { + const clock = useFakeStudyClock({ + steps: running ? allTrials.trials.length : 0, + ticksPerStep: 8, + tickMs: 150, + }); + const landed = running ? clock.landed : allTrials.trials.length; + const trials = allTrials.trials.slice(0, landed); + const inFlight = running ? allTrials.trials[landed] : undefined; + + const [chosen, setChosen] = useState(() => + navigationAtTrial(input, allTrials.trials[0]!, true), + ); + // While following, the navigation is wherever the optimizer is evaluating; + // once every step has landed it holds at the last one. + const navigation = chosen.followTrials + ? navigationAtTrial(input, inFlight ?? allTrials.trials.at(-1)!, true) + : chosen; + const key = navigationKey(input, navigation); + + const [refinement, setRefinement] = useState({ key, rung: 0 }); + if (refinement.key !== key) { + setRefinement({ key, rung: 0 }); + } + useEffect(() => { + if (refinement.rung >= REFINEMENT_LADDER.length - 1) { + return; + } + const timer = setTimeout( + () => + setRefinement((previous) => + previous.key === key ? { key, rung: previous.rung + 1 } : previous, + ), + 900, + ); + return () => clearTimeout(timer); + }, [key, refinement.rung]); + + const rung = refinement.key === key ? refinement.rung : 0; + const selection = + inFlight && navigation.followTrials + ? makeSelectionStream({ + input, + navigation, + followedTrial: inFlight.trial, + runsCompleted: 1, + computing: true, + progress: clock.progress, + }) + : refinementError !== null + ? makeSelectionStream({ + input, + navigation, + runsCompleted: 0, + error: refinementError, + }) + : makeSelectionStream({ + input, + navigation, + runsCompleted: REFINEMENT_LADDER[rung]!, + runTarget: REFINEMENT_LADDER[rung + 1] ?? null, + computing: rung < REFINEMENT_LADDER.length - 1, + }); + + const optimization = makeOptimizationRecord({ + input, + trials, + best: trials.at(-1)?.best ?? null, + status: inFlight ? "running" : "complete", + computeBackendFallbackReason: fallbackReason, + navigation, + selection, + }); + + const value: OptimizationsContextValue = { + optimizations: [optimization], + selectedOptimizationId: optimization.id, + selectedOptimization: optimization, + setSelectedOptimizationId: () => {}, + createOptimization: () => Promise.resolve(optimization.id), + cancelOptimization: () => {}, + removeOptimization: () => {}, + extendOptimization: () => Promise.resolve(), + setOptimizationNavigation: (_optimizationId, patch) => + setChosen({ ...navigation, ...patch }), + retryOptimization: () => Promise.resolve(null), + }; + + return ( + + {}} + optimization={optimization} + /> + + ); +}; + +export const ConnectedRunning: Story = { + name: "Connected study, following steps", + render: () => , +}; + +export const ConnectedComplete: Story = { + name: "Connected study, complete", + render: () => , +}; + +export const ConnectedRefinementFailed: Story = { + name: "Connected study whose point could not compute", + render: () => ( + + ), +}; + +export const ConnectedAfterGpuFallback: Story = { + name: "Connected study after GPU fallback", + render: () => ( + + ), +}; + +const RemoteStudy = () => { + const optimization = makeOptimizationRecord({ + input, + trials: allTrials.trials, + best: allTrials.best, + status: "complete", + }); + const value: OptimizationsContextValue = { + optimizations: [optimization], + selectedOptimizationId: optimization.id, + selectedOptimization: optimization, + setSelectedOptimizationId: () => {}, + createOptimization: () => Promise.resolve(optimization.id), + cancelOptimization: () => {}, + removeOptimization: () => {}, + extendOptimization: () => Promise.resolve(), + setOptimizationNavigation: () => {}, + retryOptimization: () => Promise.resolve(null), + }; + + return ( + + + {}} + optimization={optimization} + /> + + + ); +}; + +export const Remote: Story = { + name: "Remote study", + render: () => , +}; diff --git a/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/view-optimization-drawer.test.tsx b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/view-optimization-drawer.test.tsx new file mode 100644 index 00000000000..354a74a7e08 --- /dev/null +++ b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/view-optimization-drawer.test.tsx @@ -0,0 +1,460 @@ +/** + * @vitest-environment jsdom + */ +import { cleanup, fireEvent, render, screen } from "@testing-library/react"; +import { use } from "react"; +import { afterEach, describe, expect, it, vi } from "vitest"; + +import { + type OptimizationRecord, + OptimizationsContext, + type OptimizationsContextValue, +} from "../../../../../../react/optimizations/context"; +import { UserSettingsContext } from "../../../../../../react/state/user-settings-context"; +import { + makeOptimizationInput, + makeOptimizationRecord, + makeSelectionStream, + makeTrials, + navigationAtTrial, + optimizedBindingSets, +} from "./optimizations-story-fixtures"; +import { ViewOptimizationDrawer } from "./view-optimization-drawer"; + +import type { ReactNode } from "react"; + +vi.mock("@hashintel/ds-components", async (importOriginal) => { + const actual = + await importOriginal(); + const Drawer = Object.assign( + ({ children }: { children: ReactNode }) =>
{children}
, + { + Header: ({ + title, + description, + }: { + title: ReactNode; + description?: ReactNode; + }) => ( +
+ {title} +

{description}

+
+ ), + Body: ({ children }: { children: ReactNode }) =>
{children}
, + Footer: ({ actions }: { actions: ReactNode }) => ( +
{actions}
+ ), + }, + ); + const Slider = ({ + value, + disabled, + onChange, + }: { + value: number; + disabled?: boolean; + onChange?: (value: number) => void; + }) => ( + onChange?.(Number(event.target.value))} + /> + ); + const Tooltip = ({ children }: { children: ReactNode }) => <>{children}; + + return { ...actual, Drawer, Slider, Tooltip }; +}); + +vi.mock("./optimization-surface", () => ({ + OptimizationSurface: () =>
, + NavigatedOptimizationSurface: ({ + navigation, + onNavigationChange, + }: { + navigation: { positions: Record }; + onNavigationChange: (patch: { + positions: Record; + followTrials: boolean; + }) => void; + }) => ( +