It compiles

Value: [1e-3 to 4e-1]

packages/squiggle-lang/__tests__/Distributions/Scoring/KlDivergence_test.res

@@ -3,6 +3,7 @@ open Expect
 open TestHelpers
 open GenericDist_Fixtures
 
+let klDivergence = DistributionOperation.Constructors.logScore_DistEstimateDistAnswer(~env)
 // integral from low to high of 1 / (high - low) log(normal(mean, stdev)(x) / (1 / (high - low))) dx
 let klNormalUniform = (mean, stdev, low, high): float =>
   -.Js.Math.log((high -. low) /. Js.Math.sqrt(2.0 *. MagicNumbers.Math.pi *. stdev ** 2.0)) +.
@@ -11,8 +12,6 @@ let klNormalUniform = (mean, stdev, low, high): float =>
   (mean ** 2.0 -. (high +. low) *. mean +. (low ** 2.0 +. high *. low +. high ** 2.0) /. 3.0)
 
 describe("klDivergence: continuous -> continuous -> float", () => {
-  let klDivergence = DistributionOperation.Constructors.klDivergence(~env)
-
   let testUniform = (lowAnswer, highAnswer, lowPrediction, highPrediction) => {
     test("of two uniforms is equal to the analytic expression", () => {
       let answer =
@@ -82,7 +81,6 @@ describe("klDivergence: continuous -> continuous -> float", () => {
 })
 
 describe("klDivergence: discrete -> discrete -> float", () => {
-  let klDivergence = DistributionOperation.Constructors.klDivergence(~env)
   let mixture = a => DistributionTypes.DistributionOperation.Mixture(a)
   let a' = [(point1, 1e0), (point2, 1e0)]->mixture->run
   let b' = [(point1, 1e0), (point2, 1e0), (point3, 1e0)]->mixture->run
@@ -117,7 +115,6 @@ describe("klDivergence: discrete -> discrete -> float", () => {
 })
 
 describe("klDivergence: mixed -> mixed -> float", () => {
-  let klDivergence = DistributionOperation.Constructors.klDivergence(~env)
   let mixture' = a => DistributionTypes.DistributionOperation.Mixture(a)
   let mixture = a => {
     let dist' = a->mixture'->run
@@ -193,7 +190,7 @@ describe("combineAlongSupportOfSecondArgument0", () => {
     let predictionWrapped = E.R.fmap(a => run(FromDist(ToDist(ToPointSet), a)), prediction)
 
     let interpolator = XYShape.XtoY.continuousInterpolator(#Stepwise, #UseZero)
-    let integrand = PointSetDist_Scoring.KLDivergence.integrand
+    let integrand = PointSetDist_Scoring.WithDistAnswer.integrand
 
     let result = switch (answerWrapped, predictionWrapped) {
     | (Ok(Dist(PointSet(Continuous(a)))), Ok(Dist(PointSet(Continuous(b))))) =>

packages/squiggle-lang/src/rescript/Distributions/DistributionOperation.res

@@ -145,18 +145,9 @@ let rec run = (~env, functionCallInfo: functionCallInfo): outputType => {
         Dist(dist)
       }
     | ToDist(Normalize) => dist->GenericDist.normalize->Dist
-    | ToScore(KLDivergence(t2)) =>
-      GenericDist.Score.klDivergence(dist, t2, ~toPointSetFn)
-      ->E.R2.fmap(r => Float(r))
-      ->OutputLocal.fromResult
     | ToScore(LogScore(answer, prior)) =>
-      GenericDist.Score.logScoreWithPointResolution(
-        ~prediction=dist,
-        ~answer,
-        ~prior,
-        ~toPointSetFn,
-      )
-      ->E.R2.fmap(r => Float(r))
+      GenericDist.Score.logScore(~estimate=Score_Dist(dist), ~answer, ~prior)
+      ->E.R2.fmap(s => Float(s))
       ->OutputLocal.fromResult
     | ToBool(IsNormalized) => dist->GenericDist.isNormalized->Bool
     | ToDist(Truncate(leftCutoff, rightCutoff)) =>
@@ -271,13 +262,22 @@ module Constructors = {
   let pdf = (~env, dist, f) => C.pdf(dist, f)->run(~env)->toFloatR
   let normalize = (~env, dist) => C.normalize(dist)->run(~env)->toDistR
   let isNormalized = (~env, dist) => C.isNormalized(dist)->run(~env)->toBoolR
-  let klDivergence = (~env, dist1, dist2) => C.klDivergence(dist1, dist2)->run(~env)->toFloatR
-  let logScoreWithPointResolution = (
-    ~env,
-    ~prediction: DistributionTypes.genericDist,
-    ~answer: float,
-    ~prior: option<DistributionTypes.genericDist>,
-  ) => C.logScoreWithPointResolution(~prediction, ~answer, ~prior)->run(~env)->toFloatR
+  let logScore_DistEstimateDistAnswer = (~env, estimate, answer) =>
+    C.logScore_DistEstimateDistAnswer(estimate, answer)->run(~env)->toFloatR
+  let logScore_DistEstimateDistAnswerWithPrior = (~env, estimate, answer, prior) =>
+    C.logScore_DistEstimateDistAnswerWithPrior(estimate, answer, prior)->run(~env)->toFloatR
+  let logScore_DistEstimateScalarAnswer = (~env, estimate, answer) =>
+    C.logScore_DistEstimateScalarAnswer(estimate, answer)->run(~env)->toFloatR
+  let logScore_DistEstimateScalarAnswerWithPrior = (~env, estimate, answer, prior) =>
+    C.logScore_DistEstimateScalarAnswerWithPrior(estimate, answer, prior)->run(~env)->toFloatR
+  let logScore_ScalarEstimateDistAnswer = (~env, estimate, answer) =>
+    C.logScore_ScalarEstimateDistAnswer(estimate, answer)->run(~env)->toFloatR
+  let logScore_ScalarEstimateDistAnswerWithPrior = (~env, estimate, answer, prior) =>
+    C.logScore_ScalarEstimateDistAnswerWithPrior(estimate, answer, prior)->run(~env)->toFloatR
+  let logScore_ScalarEstimateScalarAnswer = (~env, estimate, answer) =>
+    C.logScore_ScalarEstimateScalarAnswer(estimate, answer)->run(~env)->toFloatR
+  let logScore_ScalarEstimateScalarAnswerWithPrior = (~env, estimate, answer, prior) =>
+    C.logScore_ScalarEstimateScalarAnswerWithPrior(estimate, answer, prior)->run(~env)->toFloatR
   let toPointSet = (~env, dist) => C.toPointSet(dist)->run(~env)->toDistR
   let toSampleSet = (~env, dist, n) => C.toSampleSet(dist, n)->run(~env)->toDistR
   let fromSamples = (~env, xs) => C.fromSamples(xs)->run(~env)->toDistR

packages/squiggle-lang/src/rescript/Distributions/DistributionOperation.resi

@@ -61,13 +61,40 @@ module Constructors: {
   @genType
   let isNormalized: (~env: env, genericDist) => result<bool, error>
   @genType
-  let klDivergence: (~env: env, genericDist, genericDist) => result<float, error>
+  let logScore_DistEstimateDistAnswer: (~env: env, genericDist, genericDist) => result<float, error>
   @genType
-  let logScoreWithPointResolution: (
+  let logScore_DistEstimateDistAnswerWithPrior: (
     ~env: env,
-    ~prediction: genericDist,
-    ~answer: float,
-    ~prior: option<genericDist>,
+    genericDist,
+    genericDist,
+    DistributionTypes.DistributionOperation.scoreDistOrScalar,
+  ) => result<float, error>
+  @genType
+  let logScore_DistEstimateScalarAnswer: (~env: env, genericDist, float) => result<float, error>
+  @genType
+  let logScore_DistEstimateScalarAnswerWithPrior: (
+    ~env: env,
+    genericDist,
+    float,
+    DistributionTypes.DistributionOperation.scoreDistOrScalar,
+  ) => result<float, error>
+  @genType
+  let logScore_ScalarEstimateDistAnswer: (~env: env, float, genericDist) => result<float, error>
+  @genType
+  let logScore_ScalarEstimateDistAnswerWithPrior: (
+    ~env: env,
+    float,
+    genericDist,
+    DistributionTypes.DistributionOperation.scoreDistOrScalar,
+  ) => result<float, error>
+  @genType
+  let logScore_ScalarEstimateScalarAnswer: (~env: env, float, float) => result<float, error>
+  @genType
+  let logScore_ScalarEstimateScalarAnswerWithPrior: (
+    ~env: env,
+    float,
+    float,
+    DistributionTypes.DistributionOperation.scoreDistOrScalar,
   ) => result<float, error>
   @genType
   let toPointSet: (~env: env, genericDist) => result<genericDist, error>

packages/squiggle-lang/src/rescript/Distributions/DistributionTypes.res

@@ -92,7 +92,9 @@ module DistributionOperation = {
     | ToString
     | ToSparkline(int)
 
-  type toScore = KLDivergence(genericDist) | LogScore(float, option<genericDist>)
+  type scoreDistOrScalar = Score_Dist(genericDist) | Score_Scalar(float)
+
+  type toScore = LogScore(scoreDistOrScalar, option<scoreDistOrScalar>)
 
   type fromDist =
     | ToFloat(toFloat)
@@ -120,8 +122,7 @@ module DistributionOperation = {
     | ToFloat(#Pdf(r)) => `pdf(${E.Float.toFixed(r)})`
     | ToFloat(#Sample) => `sample`
     | ToFloat(#IntegralSum) => `integralSum`
-    | ToScore(KLDivergence(_)) => `klDivergence`
-    | ToScore(LogScore(x, _)) => `logScore against ${E.Float.toFixed(x)}`
+    | ToScore(_) => `logScore`
     | ToDist(Normalize) => `normalize`
     | ToDist(ToPointSet) => `toPointSet`
     | ToDist(ToSampleSet(r)) => `toSampleSet(${E.I.toString(r)})`
@@ -162,10 +163,37 @@ module Constructors = {
     let fromSamples = (xs): t => FromSamples(xs)
     let truncate = (dist, left, right): t => FromDist(ToDist(Truncate(left, right)), dist)
     let inspect = (dist): t => FromDist(ToDist(Inspect), dist)
-    let klDivergence = (dist1, dist2): t => FromDist(ToScore(KLDivergence(dist2)), dist1)
-    let logScoreWithPointResolution = (~prediction, ~answer, ~prior): t => FromDist(
-      ToScore(LogScore(answer, prior)),
-      prediction,
+    let logScore_DistEstimateDistAnswer = (estimate, answer): t => FromDist(
+      ToScore(LogScore(Score_Dist(answer), None)),
+      estimate,
+    )
+    let logScore_DistEstimateDistAnswerWithPrior = (estimate, answer, prior): t => FromDist(
+      ToScore(LogScore(Score_Dist(answer), Some(prior))),
+      estimate,
+    )
+    let logScore_DistEstimateScalarAnswer = (estimate, answer): t => FromDist(
+      ToScore(LogScore(Score_Scalar(answer), None)),
+      estimate,
+    )
+    let logScore_DistEstimateScalarAnswerWithPrior = (estimate, answer, prior): t => FromDist(
+      ToScore(LogScore(Score_Scalar(answer), Some(prior))),
+      estimate,
+    )
+    let logScore_ScalarEstimateDistAnswer = (estimate, answer): t => FromFloat(
+      ToScore(LogScore(Score_Dist(answer), None)),
+      estimate,
+    )
+    let logScore_ScalarEstimateDistAnswerWithPrior = (estimate, answer, prior): t => FromFloat(
+      ToScore(LogScore(Score_Dist(answer), Some(prior))),
+      estimate,
+    )
+    let logScore_ScalarEstimateScalarAnswer = (estimate, answer): t => FromFloat(
+      ToScore(LogScore(Score_Scalar(answer), None)),
+      estimate,
+    )
+    let logScore_ScalarEstimateScalarAnswerWithPrior = (estimate, answer, prior): t => FromFloat(
+      ToScore(LogScore(Score_Scalar(answer), Some(prior))),
+      estimate,
     )
     let scalePower = (dist, n): t => FromDist(ToDist(Scale(#Power, n)), dist)
     let scaleLogarithm = (dist, n): t => FromDist(ToDist(Scale(#Logarithm, n)), dist)

packages/squiggle-lang/src/rescript/Distributions/GenericDist.res

@@ -61,46 +61,6 @@ let integralEndY = (t: t): float =>
 
 let isNormalized = (t: t): bool => Js.Math.abs_float(integralEndY(t) -. 1.0) < 1e-7
 
-module Score = {
-  let klDivergence = (prediction, answer, ~toPointSetFn: toPointSetFn): result<float, error> => {
-    let pointSets = E.R.merge(toPointSetFn(prediction), toPointSetFn(answer))
-    pointSets |> E.R2.bind(((predi, ans)) =>
-      PointSetDist.T.klDivergence(predi, ans)->E.R2.errMap(x => DistributionTypes.OperationError(x))
-    )
-  }
-
-  let logScoreWithPointResolution = (
-    ~prediction: DistributionTypes.genericDist,
-    ~answer: float,
-    ~prior: option<DistributionTypes.genericDist>,
-    ~toPointSetFn: toPointSetFn,
-  ): result<float, error> => {
-    switch prior {
-    | Some(prior') =>
-      E.R.merge(toPointSetFn(prior'), toPointSetFn(prediction))->E.R.bind(((
-        prior'',
-        prediction'',
-      )) =>
-        PointSetDist.T.logScoreWithPointResolution(
-          ~prediction=prediction'',
-          ~answer,
-          ~prior=prior''->Some,
-        )->E.R2.errMap(x => DistributionTypes.OperationError(x))
-      )
-    | None =>
-      prediction
-      ->toPointSetFn
-      ->E.R.bind(x =>
-        PointSetDist.T.logScoreWithPointResolution(
-          ~prediction=x,
-          ~answer,
-          ~prior=None,
-        )->E.R2.errMap(x => DistributionTypes.OperationError(x))
-      )
-    }
-  }
-}
-
 let toFloatOperation = (
   t,
   ~toPointSetFn: toPointSetFn,
@@ -159,6 +119,125 @@ let toPointSet = (
   }
 }
 
+module Score = {
+  type scoreDistOrScalar = DistributionTypes.DistributionOperation.scoreDistOrScalar
+
+  type pointSet_ScoreDistOrScalar = D(PointSetTypes.pointSetDist) | S(float)
+  let argsMake = (
+    ~esti: scoreDistOrScalar,
+    ~answ: scoreDistOrScalar,
+    ~prior: option<scoreDistOrScalar>,
+  ): result<PointSetDist_Scoring.scoreArgs, error> => {
+    let toPointSetFn = toPointSet(
+      ~xyPointLength=MagicNumbers.Environment.defaultXYPointLength,
+      ~sampleCount=MagicNumbers.Environment.defaultSampleCount,
+      ~xSelection=#ByWeight,
+    )
+    let prior': option<result<pointSet_ScoreDistOrScalar, error>> = switch prior {
+    | None => None
+    | Some(Score_Dist(d)) => toPointSetFn(d, ())->E.R.bind(x => x->D->Ok)->Some
+    | Some(Score_Scalar(s)) => s->S->Ok->Some
+    }
+    let twoDists = (esti': t, answ': t): result<
+      (PointSetTypes.pointSetDist, PointSetTypes.pointSetDist),
+      error,
+    > => E.R.merge(toPointSetFn(esti', ()), toPointSetFn(answ', ()))
+    switch (esti, answ, prior') {
+    | (Score_Dist(esti'), Score_Dist(answ'), None) =>
+      twoDists(esti', answ')->E.R.bind(((esti'', answ'')) =>
+        {estimate: esti'', answer: answ'', prior: None}
+        ->PointSetDist_Scoring.DistEstimateDistAnswer
+        ->Ok
+      )
+    | (Score_Dist(esti'), Score_Dist(answ'), Some(Ok(D(prior'')))) =>
+      twoDists(esti', answ')->E.R.bind(((esti'', answ'')) =>
+        {estimate: esti'', answer: answ'', prior: Some(prior'')}
+        ->PointSetDist_Scoring.DistEstimateDistAnswer
+        ->Ok
+      )
+    | (Score_Dist(esti'), Score_Scalar(answ'), None) =>
+      toPointSetFn(esti', ())->E.R.bind(esti'' =>
+        {estimate: esti'', answer: answ', prior: None}
+        ->PointSetDist_Scoring.DistEstimateScalarAnswer
+        ->Ok
+      )
+    | (Score_Dist(esti'), Score_Scalar(answ'), Some(Ok(D(prior'')))) =>
+      toPointSetFn(esti', ())->E.R.bind(esti'' =>
+        {estimate: esti'', answer: answ', prior: Some(prior'')}
+        ->PointSetDist_Scoring.DistEstimateScalarAnswer
+        ->Ok
+      )
+    | (Score_Scalar(esti'), Score_Dist(answ'), None) =>
+      toPointSetFn(answ', ())->E.R.bind(answ'' =>
+        {estimate: esti', answer: answ'', prior: None}
+        ->PointSetDist_Scoring.ScalarEstimateDistAnswer
+        ->Ok
+      )
+    | (Score_Scalar(esti'), Score_Dist(answ'), Some(Ok(S(prior'')))) =>
+      toPointSetFn(answ', ())->E.R.bind(answ'' =>
+        {estimate: esti', answer: answ'', prior: Some(prior'')}
+        ->PointSetDist_Scoring.ScalarEstimateDistAnswer
+        ->Ok
+      )
+    | (Score_Scalar(esti'), Score_Scalar(answ'), None) =>
+      {estimate: esti', answer: answ', prior: None}
+      ->PointSetDist_Scoring.ScalarEstimateScalarAnswer
+      ->Ok
+    | (Score_Scalar(esti'), Score_Scalar(answ'), Some(Ok(S(prior'')))) =>
+      {estimate: esti', answer: answ', prior: prior''->Some}
+      ->PointSetDist_Scoring.ScalarEstimateScalarAnswer
+      ->Ok
+    | (_, _, Some(Error(err))) => err->Error
+    }
+  }
+
+  let logScore = (
+    ~estimate: scoreDistOrScalar,
+    ~answer: scoreDistOrScalar,
+    ~prior: option<scoreDistOrScalar>,
+  ): result<float, error> =>
+    argsMake(~esti=estimate, ~answ=answer, ~prior)->E.R.bind(x =>
+      x->PointSetDist.logScore->E.R2.errMap(y => DistributionTypes.OperationError(y))
+    )
+
+  //  let klDivergence = (prediction, answer, ~toPointSetFn: toPointSetFn): result<float, error> => {
+  //    let pointSets = E.R.merge(toPointSetFn(prediction), toPointSetFn(answer))
+  //    pointSets |> E.R2.bind(((predi, ans)) =>
+  //      PointSetDist.T.klDivergence(predi, ans)->E.R2.errMap(x => DistributionTypes.OperationError(x))
+  //    )
+  //  }
+  //
+  //  let logScoreWithPointResolution = (
+  //    ~prediction: DistributionTypes.genericDist,
+  //    ~answer: float,
+  //    ~prior: option<DistributionTypes.genericDist>,
+  //    ~toPointSetFn: toPointSetFn,
+  //  ): result<float, error> => {
+  //    switch prior {
+  //    | Some(prior') =>
+  //      E.R.merge(toPointSetFn(prior'), toPointSetFn(prediction))->E.R.bind(((
+  //        prior'',
+  //        prediction'',
+  //      )) =>
+  //        PointSetDist.T.logScoreWithPointResolution(
+  //          ~prediction=prediction'',
+  //          ~answer,
+  //          ~prior=prior''->Some,
+  //        )->E.R2.errMap(x => DistributionTypes.OperationError(x))
+  //      )
+  //    | None =>
+  //      prediction
+  //      ->toPointSetFn
+  //      ->E.R.bind(x =>
+  //        PointSetDist.T.logScoreWithPointResolution(
+  //          ~prediction=x,
+  //          ~answer,
+  //          ~prior=None,
+  //        )->E.R2.errMap(x => DistributionTypes.OperationError(x))
+  //      )
+  //    }
+  //  }
+}
 /*
   PointSetDist.toSparkline calls "downsampleEquallyOverX", which downsamples it to n=bucketCount.
   It first needs a pointSetDist, so we convert to a pointSetDist. In this process we want the 

packages/squiggle-lang/src/rescript/Distributions/GenericDist.resi

@@ -25,12 +25,10 @@ let toFloatOperation: (
 ) => result<float, error>
 
 module Score: {
-  let klDivergence: (t, t, ~toPointSetFn: toPointSetFn) => result<float, error>
-  let logScoreWithPointResolution: (
-    ~prediction: t,
-    ~answer: float,
-    ~prior: option<t>,
-    ~toPointSetFn: toPointSetFn,
+  let logScore: (
+    ~estimate: DistributionTypes.DistributionOperation.scoreDistOrScalar,
+    ~answer: DistributionTypes.DistributionOperation.scoreDistOrScalar,
+    ~prior: option<DistributionTypes.DistributionOperation.scoreDistOrScalar>,
   ) => result<float, error>
 }
 

packages/squiggle-lang/src/rescript/Distributions/PointSetDist/PointSetDist.res

@@ -215,7 +215,7 @@ module T = Dist({
 })
 
 let logScore = (args: PointSetDist_Scoring.scoreArgs): result<float, Operation.Error.t> =>
-  PointSetDist_Scoring.logScore(args)
+  PointSetDist_Scoring.logScore(args, ~combineFn=combinePointwise, ~integrateFn=T.integralEndY)
 
 let pdf = (f: float, t: t) => {
   let mixedPoint: PointSetTypes.mixedPoint = T.xToY(f, t)

packages/squiggle-lang/src/rescript/Distributions/PointSetDist/PointSetDist_Scoring.res

@@ -2,6 +2,7 @@ type t = PointSetTypes.pointSetDist
 type continuousShape = PointSetTypes.continuousShape
 type discreteShape = PointSetTypes.discreteShape
 type mixedShape = PointSetTypes.mixedShape
+
 type scalar = float
 type abstractScoreArgs<'a, 'b> = {estimate: 'a, answer: 'b, prior: option<'a>}
 type scoreArgs =
@@ -9,6 +10,7 @@ type scoreArgs =
   | DistEstimateScalarAnswer(abstractScoreArgs<t, scalar>)
   | ScalarEstimateDistAnswer(abstractScoreArgs<scalar, t>)
   | ScalarEstimateScalarAnswer(abstractScoreArgs<scalar, scalar>)
+
 let logFn = Js.Math.log // base e
 let minusScaledLogOfQuot = (~esti, ~answ): result<float, Operation.Error.t> => {
   let quot = esti /. answ
@@ -56,7 +58,7 @@ module WithScalarAnswer = {
   }
   let scoreWithPrior' = (
     ~estimatePdf: float => float,
-    ~answer: float,
+    ~answer: scalar,
     ~priorPdf: float => float,
   ): result<float, Operation.Error.t> => {
     let numerator = answer->estimatePdf
@@ -69,23 +71,44 @@ module WithScalarAnswer = {
       minusScaledLogOfQuot(~esti=numerator, ~answ=priorDensityOfAnswer)
     }
   }
-  let score = (~estimate: t, ~answer: t, ~mapper): result<float, Operation.Error.t> => {
-    let pdf = (shape, ~x) => XYShape.XtoY.linear(x, shape.xyShape)
-    let estimatePdf = mapper((x => pdf(~x), x => pdf(~x), x => pdf(~x)))
+  let score = (~estimate: t, ~answer: scalar): result<float, Operation.Error.t> => {
+    let estimatePdf = x =>
+      switch estimate {
+      | Continuous(esti) => XYShape.XtoY.linear(x, esti.xyShape)
+      | Discrete(esti) => XYShape.XtoY.linear(x, esti.xyShape)
+      | Mixed(esti) =>
+        XYShape.XtoY.linear(x, esti.continuous.xyShape) +.
+        XYShape.XtoY.linear(x, esti.discrete.xyShape)
+      }
+
     score'(~estimatePdf, ~answer)
   }
-  let scoreWithPrior = (~estimate: t, ~answer: t, ~prior: t, ~mapper): result<
+  let scoreWithPrior = (~estimate: t, ~answer: scalar, ~prior: t): result<
     float,
     Operation.Error.t,
   > => {
-    let estimatePdf = x => XYShape.XtoY.linear(x, estimate.xyShape)
-    let priorPdf = x => XYShape.XtoY.linear(x, prior.xyShape)
+    let estimatePdf = x =>
+      switch estimate {
+      | Continuous(esti) => XYShape.XtoY.linear(x, esti.xyShape)
+      | Discrete(esti) => XYShape.XtoY.linear(x, esti.xyShape)
+      | Mixed(esti) =>
+        XYShape.XtoY.linear(x, esti.continuous.xyShape) +.
+        XYShape.XtoY.linear(x, esti.discrete.xyShape)
+      }
+    let priorPdf = x =>
+      switch prior {
+      | Continuous(prio) => XYShape.XtoY.linear(x, prio.xyShape)
+      | Discrete(prio) => XYShape.XtoY.linear(x, prio.xyShape)
+      | Mixed(prio) =>
+        XYShape.XtoY.linear(x, prio.continuous.xyShape) +.
+        XYShape.XtoY.linear(x, prio.discrete.xyShape)
+      }
     scoreWithPrior'(~estimatePdf, ~answer, ~priorPdf)
   }
 }
 
 module TwoScalars = {
-  let score = (~estimate: float, ~answer: float) =>
+  let score = (~estimate: scalar, ~answer: scalar) =>
     if answer == 0.0 {
       0.0->Ok
     } else if estimate == 0.0 {
@@ -107,9 +130,9 @@ module TwoScalars = {
 let logScore = (args: scoreArgs, ~combineFn, ~integrateFn): result<float, Operation.Error.t> =>
   switch args {
   | DistEstimateDistAnswer({estimate, answer, prior: None}) =>
-    WithDistAnswer.sum(~estimate, ~answer, ~integrateFn)
+    WithDistAnswer.sum(~estimate, ~answer, ~integrateFn, ~combineFn)
   | DistEstimateDistAnswer({estimate, answer, prior: Some(prior)}) =>
-    WithDistAnswer.sumWithPrior(~estimate, ~answer, ~prior, ~integrateFn)
+    WithDistAnswer.sumWithPrior(~estimate, ~answer, ~prior, ~integrateFn, ~combineFn)
   | DistEstimateScalarAnswer({estimate, answer, prior: None}) =>
     WithScalarAnswer.score(~estimate, ~answer)
   | DistEstimateScalarAnswer({estimate, answer, prior: Some(prior)}) =>

packages/squiggle-lang/src/rescript/ReducerInterface/ReducerInterface_GenericDistribution.res

@@ -162,20 +162,6 @@ module Helpers = {
       }
     }
   }
-
-  let klDivergenceWithPrior = (
-    prediction: DistributionTypes.genericDist,
-    answer: DistributionTypes.genericDist,
-    prior: DistributionTypes.genericDist,
-    env: DistributionOperation.env,
-  ) => {
-    let term1 = DistributionOperation.Constructors.klDivergence(~env, prediction, answer)
-    let term2 = DistributionOperation.Constructors.klDivergence(~env, prior, answer)
-    switch E.R.merge(term1, term2)->E.R2.fmap(((a, b)) => a -. b) {
-    | Ok(x) => x->DistributionOperation.Float->Some
-    | Error(_) => None
-    }
-  }
 }
 
 module SymbolicConstructors = {
@@ -226,28 +212,20 @@ let dispatchToGenericOutput = (
       ~env,
     )->Some
   | ("normalize", [EvDistribution(dist)]) => Helpers.toDistFn(Normalize, dist, ~env)
-  | ("klDivergence", [EvDistribution(prediction), EvDistribution(answer)]) =>
-    Some(DistributionOperation.run(FromDist(ToScore(KLDivergence(answer)), prediction), ~env))
-  | ("klDivergence", [EvDistribution(prediction), EvDistribution(answer), EvDistribution(prior)]) =>
-    Helpers.klDivergenceWithPrior(prediction, answer, prior, env)
-  | (
-    "logScoreWithPointAnswer",
-    [EvDistribution(prediction), EvNumber(answer), EvDistribution(prior)],
-  )
-  | (
-    "logScoreWithPointAnswer",
-    [EvDistribution(prediction), EvDistribution(Symbolic(#Float(answer))), EvDistribution(prior)],
-  ) =>
-    DistributionOperation.run(
-      FromDist(ToScore(LogScore(answer, prior->Some)), prediction),
-      ~env,
-    )->Some
-  | ("logScoreWithPointAnswer", [EvDistribution(prediction), EvNumber(answer)])
-  | (
-    "logScoreWithPointAnswer",
-    [EvDistribution(prediction), EvDistribution(Symbolic(#Float(answer)))],
-  ) =>
-    DistributionOperation.run(FromDist(ToScore(LogScore(answer, None)), prediction), ~env)->Some
+  | ("klDivergence", [EvDistribution(estimate), EvDistribution(answer)]) =>
+    Some(
+      DistributionOperation.run(
+        FromDist(ToScore(LogScore(Score_Dist(answer), None)), estimate),
+        ~env,
+      ),
+    )
+  | ("klDivergence", [EvDistribution(estimate), EvDistribution(answer), EvDistribution(prior)]) =>
+    Some(
+      DistributionOperation.run(
+        FromDist(ToScore(LogScore(Score_Dist(answer), Some(Score_Dist(prior)))), estimate),
+        ~env,
+      ),
+    )
   | ("isNormalized", [EvDistribution(dist)]) => Helpers.toBoolFn(IsNormalized, dist, ~env)
   | ("toPointSet", [EvDistribution(dist)]) => Helpers.toDistFn(ToPointSet, dist, ~env)
   | ("scaleLog", [EvDistribution(dist)]) =>