@prismatic-io/lux 0.0.1

package/src/orchestrator/experiment-selection.ts

@@ -0,0 +1,467 @@
+import { z } from "zod";
+import {
+  bootstrapMeanConfidenceInterval,
+  type Candidate,
+  CandidateIdSchema,
+  CandidateSchema,
+  type ExperimentCampaign,
+  type ExperimentConstraint,
+} from "../core/index.js";
+import {
+  type ExperimentEvaluation,
+  type ExperimentSelection,
+  ExperimentSelectionSchema,
+  type OpenedExperimentRun,
+} from "./experiment-runs.js";
+
+type ExperimentSplit = "train" | "validation" | "test";
+
+export const ExperimentDecisionCandidateSchema = CandidateSchema.pick({
+  id: true,
+  sourceBytes: true,
+});
+export type ExperimentDecisionCandidate = z.infer<typeof ExperimentDecisionCandidateSchema>;
+
+export const candidateEvaluations = (
+  experiment: OpenedExperimentRun["manifest"],
+  candidateId: string,
+): ExperimentEvaluation[] => {
+  const selected = new Map<string, ExperimentEvaluation>();
+  for (const evaluation of experiment.evaluations.filter(
+    (entry) => entry.candidateId === candidateId,
+  )) {
+    const key = `${evaluation.profileId}:${evaluation.split}`;
+    const current = selected.get(key);
+    if (!current || evaluation.repetitions > current.repetitions) selected.set(key, evaluation);
+  }
+  return [...selected.values()];
+};
+
+const caseRegressions = (
+  candidate: ExperimentDecisionCandidate,
+  profile: string | undefined,
+  split: ExperimentSplit,
+  experiment: OpenedExperimentRun["manifest"],
+): number => {
+  if (candidate.id === "baseline") return 0;
+  const baseline = candidateEvaluations(experiment, "baseline").filter(
+    (entry) => entry.split === split && (!profile || entry.profileId === profile),
+  );
+  const proposed = candidateEvaluations(experiment, candidate.id);
+  let regressions = 0;
+  for (const before of baseline) {
+    const after = proposed.find(
+      (entry) => entry.split === split && entry.profileId === before.profileId,
+    );
+    for (const [caseId, passRate] of Object.entries(before.casePassRates)) {
+      if (!after || (after.casePassRates[caseId] ?? 0) < passRate) regressions++;
+    }
+  }
+  return regressions;
+};
+
+const metricValues = (
+  candidate: ExperimentDecisionCandidate,
+  metric: string,
+  profile: string | undefined,
+  split: ExperimentSplit,
+  experiment: OpenedExperimentRun["manifest"],
+): number[] => {
+  if (metric === "quality.caseRegressions") {
+    return [caseRegressions(candidate, profile, split, experiment)];
+  }
+  if (metric === "subject.bytes") return [candidate.sourceBytes];
+  return candidateEvaluations(experiment, candidate.id)
+    .filter((entry) => entry.split === split && (!profile || entry.profileId === profile))
+    .flatMap((entry) => {
+      const value = metric === "quality.passRate" ? entry.passRate : entry.metrics[metric];
+      return value === undefined ? [] : [value];
+    });
+};
+
+const metricSampleGroups = (
+  candidate: ExperimentDecisionCandidate,
+  metric: string,
+  profile: string | undefined,
+  split: ExperimentSplit,
+  experiment: OpenedExperimentRun["manifest"],
+): number[][] => {
+  if (metric === "quality.caseRegressions" || metric === "subject.bytes") {
+    return metricValues(candidate, metric, profile, split, experiment).map((value) => [value]);
+  }
+  return candidateEvaluations(experiment, candidate.id)
+    .filter((entry) => entry.split === split && (!profile || entry.profileId === profile))
+    .flatMap((entry) => {
+      const samples = entry.samples?.[metric] ?? [];
+      return samples.length > 0 ? [samples] : [];
+    });
+};
+
+const satisfiesConstraint = (
+  campaign: ExperimentCampaign,
+  candidate: ExperimentDecisionCandidate,
+  constraint: ExperimentConstraint,
+  experiment: OpenedExperimentRun["manifest"],
+): boolean => {
+  let profiles: (string | undefined)[];
+  if (constraint.metric === "subject.bytes" || constraint.profile) {
+    profiles = [constraint.profile];
+  } else {
+    profiles = campaign.profiles;
+  }
+  return profiles.every((profile) => {
+    const values = metricValues(
+      candidate,
+      constraint.metric,
+      profile,
+      constraint.split,
+      experiment,
+    );
+    if (values.length === 0) return false;
+    if (constraint.confidence) {
+      const sampleGroups = metricSampleGroups(
+        candidate,
+        constraint.metric,
+        profile,
+        constraint.split,
+        experiment,
+      );
+      if (sampleGroups.length === 0) {
+        throw new Error(
+          `metric '${constraint.metric}' does not expose raw samples for confidence gating` +
+            `${profile ? ` in profile '${profile}'` : ""}`,
+        );
+      }
+      return sampleGroups.every((samples) => {
+        try {
+          const interval = bootstrapMeanConfidenceInterval(samples, constraint.confidence);
+          return (
+            (constraint.min === undefined || interval.lower >= constraint.min) &&
+            (constraint.max === undefined || interval.upper <= constraint.max)
+          );
+        } catch (error) {
+          if (error instanceof Error && error.message.includes("independent samples")) return false;
+          throw error;
+        }
+      });
+    }
+    return values.every(
+      (value) =>
+        (constraint.min === undefined || value >= constraint.min) &&
+        (constraint.max === undefined || value <= constraint.max),
+    );
+  });
+};
+
+export const average = (values: readonly number[]): number =>
+  values.reduce((sum, value) => sum + value, 0) / Math.max(values.length, 1);
+
+const normalizedConstraintViolation = (
+  campaign: ExperimentCampaign,
+  candidate: ExperimentDecisionCandidate,
+  constraint: ExperimentConstraint,
+  experiment: OpenedExperimentRun["manifest"],
+): number => {
+  let profiles: Array<string | undefined> = campaign.profiles;
+  if (constraint.metric === "subject.bytes" || constraint.profile) {
+    profiles = [constraint.profile];
+  }
+  const violation = (lower: number, upper: number): number =>
+    Math.max(
+      constraint.min === undefined
+        ? 0
+        : Math.max(0, constraint.min - lower) / Math.max(1, Math.abs(constraint.min)),
+      constraint.max === undefined
+        ? 0
+        : Math.max(0, upper - constraint.max) / Math.max(1, Math.abs(constraint.max)),
+    );
+  let worst = 0;
+  for (const profile of profiles) {
+    if (constraint.confidence) {
+      const groups = metricSampleGroups(
+        candidate,
+        constraint.metric,
+        profile,
+        constraint.split,
+        experiment,
+      );
+      if (groups.length === 0) return Number.POSITIVE_INFINITY;
+      for (const samples of groups) {
+        try {
+          const interval = bootstrapMeanConfidenceInterval(samples, constraint.confidence);
+          worst = Math.max(worst, violation(interval.lower, interval.upper));
+        } catch {
+          return Number.POSITIVE_INFINITY;
+        }
+      }
+      continue;
+    }
+    const values = metricValues(
+      candidate,
+      constraint.metric,
+      profile,
+      constraint.split,
+      experiment,
+    );
+    if (values.length === 0) return Number.POSITIVE_INFINITY;
+    for (const value of values) worst = Math.max(worst, violation(value, value));
+  }
+  return worst;
+};
+
+export const rankCandidates = <T extends ExperimentDecisionCandidate>(
+  campaign: ExperimentCampaign,
+  candidates: readonly T[],
+  experiment: OpenedExperimentRun["manifest"],
+  constraintSplits: readonly ExperimentSplit[] = ["train", "validation", "test"],
+): T[] =>
+  [...candidates].sort((a, b) => {
+    const activeConstraints = campaign.objectives.constraints.filter((constraint) =>
+      constraintSplits.includes(constraint.split),
+    );
+    const aEligible = activeConstraints.every((constraint) =>
+      satisfiesConstraint(campaign, a, constraint, experiment),
+    );
+    const bEligible = activeConstraints.every((constraint) =>
+      satisfiesConstraint(campaign, b, constraint, experiment),
+    );
+    if (aEligible !== bEligible) return aEligible ? -1 : 1;
+    if (!aEligible) {
+      const aViolation = Math.max(
+        0,
+        ...activeConstraints.map((constraint) =>
+          normalizedConstraintViolation(campaign, a, constraint, experiment),
+        ),
+      );
+      const bViolation = Math.max(
+        0,
+        ...activeConstraints.map((constraint) =>
+          normalizedConstraintViolation(campaign, b, constraint, experiment),
+        ),
+      );
+      if (aViolation !== bViolation) return aViolation < bViolation ? -1 : 1;
+    }
+
+    for (const objective of campaign.objectives.minimize.filter((entry) =>
+      constraintSplits.includes(entry.split),
+    )) {
+      const aValues = metricValues(
+        a,
+        objective.metric,
+        objective.profile,
+        objective.split,
+        experiment,
+      );
+      const bValues = metricValues(
+        b,
+        objective.metric,
+        objective.profile,
+        objective.split,
+        experiment,
+      );
+      if (aValues.length === 0 && bValues.length > 0) return 1;
+      if (bValues.length === 0 && aValues.length > 0) return -1;
+      if (aValues.length > 0 && bValues.length > 0) {
+        const delta = average(aValues) - average(bValues);
+        if (delta !== 0) return delta;
+      }
+    }
+    const quality = (candidate: ExperimentDecisionCandidate): number => {
+      const values = metricValues(
+        candidate,
+        "quality.passRate",
+        undefined,
+        "validation",
+        experiment,
+      );
+      return values.length > 0 ? Math.min(...values) : -1;
+    };
+    const qualityDelta = quality(b) - quality(a);
+    if (qualityDelta !== 0) return qualityDelta;
+    if (a.sourceBytes !== b.sourceBytes) return a.sourceBytes - b.sourceBytes;
+    return a.id.localeCompare(b.id);
+  });
+
+const candidateCaseVector = (
+  candidate: Candidate,
+  experiment: OpenedExperimentRun["manifest"],
+  preferredSplit: ExperimentSplit,
+): Record<string, number> => {
+  const evaluations = candidateEvaluations(experiment, candidate.id);
+  const split = evaluations.some((entry) => entry.split === preferredSplit)
+    ? preferredSplit
+    : "validation";
+  return Object.fromEntries(
+    evaluations
+      .filter((entry) => entry.split === split)
+      .flatMap((entry) =>
+        Object.entries(entry.casePassRates).map(([caseId, passRate]) => [
+          `${entry.profileId}:${caseId}`,
+          passRate,
+        ]),
+      ),
+  );
+};
+
+const dominates = (a: Record<string, number>, b: Record<string, number>): boolean => {
+  const keys = new Set([...Object.keys(a), ...Object.keys(b)]);
+  if (keys.size === 0) return false;
+  let strictlyBetter = false;
+  for (const key of keys) {
+    const av = a[key] ?? 0;
+    const bv = b[key] ?? 0;
+    if (av < bv) return false;
+    if (av > bv) strictlyBetter = true;
+  }
+  return strictlyBetter;
+};
+
+/**
+ * Per-case Pareto frontier used during reflective search. Quality constraints
+ * filter the population when at least one eligible candidate exists.
+ */
+export const paretoFrontier = (
+  campaign: ExperimentCampaign,
+  candidates: readonly Candidate[],
+  experiment: OpenedExperimentRun["manifest"],
+  split: ExperimentSplit = "train",
+  constraintSplits: readonly ExperimentSplit[] = ["train", "validation"],
+): Candidate[] => {
+  const activeConstraints = campaign.objectives.constraints.filter((constraint) =>
+    constraintSplits.includes(constraint.split),
+  );
+  const eligible = candidates.filter((candidate) =>
+    activeConstraints.every((constraint) =>
+      satisfiesConstraint(campaign, candidate, constraint, experiment),
+    ),
+  );
+  const population = eligible.length > 0 ? eligible : [...candidates];
+  const vectors = new Map(
+    population.map((candidate) => [
+      candidate.id,
+      candidateCaseVector(candidate, experiment, split),
+    ]),
+  );
+  return population
+    .filter((candidate) => {
+      const vector = vectors.get(candidate.id) ?? {};
+      return !population.some(
+        (other) => other.id !== candidate.id && dominates(vectors.get(other.id) ?? {}, vector),
+      );
+    })
+    .sort((a, b) => a.id.localeCompare(b.id));
+};
+
+const frontierCoverage = (
+  candidate: Candidate,
+  frontier: readonly Candidate[],
+  experiment: OpenedExperimentRun["manifest"],
+): number => {
+  const vectors = frontier.map((entry) => candidateCaseVector(entry, experiment, "train"));
+  const candidateVector = candidateCaseVector(candidate, experiment, "train");
+  const keys = new Set(vectors.flatMap((vector) => Object.keys(vector)));
+  return [...keys].filter((key) => {
+    const best = Math.max(...vectors.map((vector) => vector[key] ?? 0));
+    return (candidateVector[key] ?? 0) === best;
+  }).length;
+};
+
+export const searchFrontier = (
+  campaign: ExperimentCampaign,
+  candidates: readonly Candidate[],
+  experiment: OpenedExperimentRun["manifest"],
+): Candidate[] => {
+  const frontier = paretoFrontier(campaign, candidates, experiment);
+  return frontier.sort((a, b) => {
+    const coverage =
+      frontierCoverage(b, frontier, experiment) - frontierCoverage(a, frontier, experiment);
+    return coverage || a.id.localeCompare(b.id);
+  });
+};
+
+export const deriveExperimentSelection = (
+  campaign: ExperimentCampaign,
+  candidates: readonly ExperimentDecisionCandidate[],
+  experiment: OpenedExperimentRun["manifest"],
+  rawOptions: z.input<typeof ExperimentSelectionDerivationOptionsSchema> = {},
+): ExperimentSelection => {
+  const options = ExperimentSelectionDerivationOptionsSchema.parse(rawOptions);
+  const decisionCandidates = candidates.map((candidate) =>
+    ExperimentDecisionCandidateSchema.parse(candidate),
+  );
+  const ranked = rankCandidates(campaign, decisionCandidates, experiment);
+  const champion = ranked[0] ?? null;
+  const eligibleCandidates = ranked.filter((candidate) =>
+    campaign.objectives.constraints.every((constraint) =>
+      satisfiesConstraint(campaign, candidate, constraint, experiment),
+    ),
+  );
+  const eligible =
+    champion !== null && eligibleCandidates.some((candidate) => candidate.id === champion.id);
+  const constraintReports = decisionCandidates.flatMap((candidate) =>
+    campaign.objectives.constraints.flatMap((constraint) => {
+      const profiles: Array<string | undefined> =
+        constraint.profile || constraint.metric === "subject.bytes"
+          ? [constraint.profile]
+          : campaign.profiles;
+      return profiles.map((profile) => {
+        const values = metricValues(
+          candidate,
+          constraint.metric,
+          profile,
+          constraint.split,
+          experiment,
+        );
+        const scoped = { ...constraint, ...(profile ? { profile } : {}) };
+        let interval: ReturnType<typeof bootstrapMeanConfidenceInterval> | undefined;
+        let note: string | undefined;
+        if (constraint.confidence) {
+          const samples = metricSampleGroups(
+            candidate,
+            constraint.metric,
+            profile,
+            constraint.split,
+            experiment,
+          )[0];
+          if (samples) {
+            try {
+              interval = bootstrapMeanConfidenceInterval(samples, constraint.confidence);
+            } catch (error) {
+              note = error instanceof Error ? error.message : String(error);
+            }
+          } else {
+            note = "metric exposes no raw samples";
+          }
+        }
+        return {
+          candidateId: candidate.id,
+          metric: constraint.metric,
+          ...(profile ? { profile } : {}),
+          split: constraint.split,
+          values,
+          passed: satisfiesConstraint(campaign, candidate, scoped, experiment),
+          ...(interval ? { interval } : {}),
+          ...(note ? { note } : {}),
+        };
+      });
+    }),
+  );
+  return ExperimentSelectionSchema.parse({
+    decidedAt: options.decidedAt ?? new Date().toISOString(),
+    decisionCandidateIds: decisionCandidates.map((candidate) => candidate.id),
+    ...(options.promotionFinalistCandidateIds
+      ? { promotionFinalistCandidateIds: options.promotionFinalistCandidateIds }
+      : {}),
+    rankedCandidateIds: ranked.map((candidate) => candidate.id),
+    eligibleCandidateIds: eligibleCandidates.map((candidate) => candidate.id),
+    ...(eligible && champion ? { championCandidateId: champion.id } : {}),
+    constraints: constraintReports,
+  });
+};
+
+export const ExperimentSelectionDerivationOptionsSchema = z
+  .object({
+    decidedAt: z.string().optional(),
+    promotionFinalistCandidateIds: z.array(CandidateIdSchema).optional(),
+  })
+  .strict();