From 5de75a402be706ea71ad92ce0ee379ab96b7a67c Mon Sep 17 00:00:00 2001
From: Ozzie Gooen <ozzieagooen@gmail.com>
Date: Sun, 16 Feb 2020 11:01:37 +0000
Subject: [PATCH] Improved handling of mixed discreteProbabilityMassFraction

---
 src/lib/Prop.re             | 13 +++++++++++--
 src/lib/Shape.re            | 20 +++++++++++++++-----
 src/utility/Guesstimator.re |  2 +-
 3 files changed, 27 insertions(+), 8 deletions(-)

diff --git a/src/lib/Prop.re b/src/lib/Prop.re
index 4221cb98..eac149ff 100644
--- a/src/lib/Prop.re
+++ b/src/lib/Prop.re
@@ -61,10 +61,19 @@ module Value = {
       let newDistribution =
         GenericDistribution.renderIfNeeded(~sampleCount=1000, r);
       switch (newDistribution) {
-      | Some({generationSource: Shape(Mixed({continuous: n, discrete: d}))}) =>
+      | Some({
+          generationSource:
+            Shape(
+              Mixed({
+                continuous: n,
+                discrete: d,
+                discreteProbabilityMassFraction: f,
+              }),
+            ),
+        }) =>
         <div>
           <Chart height=100 data={n |> Shape.Continuous.toJs} />
-          {Shape.Discrete.render(d)}
+          {d |> Shape.Discrete.scaleYToTotal(f) |> Shape.Discrete.render}
         </div>
       | None => "Something went wrong" |> ReasonReact.string
       | _ => <div />
diff --git a/src/lib/Shape.re b/src/lib/Shape.re
index 914f2f51..c98500f8 100644
--- a/src/lib/Shape.re
+++ b/src/lib/Shape.re
@@ -38,7 +38,7 @@ module Discrete = {
       Belt.Array.zip(p.xs, p.ys)
       ->Belt.Array.reduce([||], (items, (x, y)) =>
           switch (_lastElement(items)) {
-          | Some((_, yLast)) => [|(x, y +. yLast)|]
+          | Some((_, yLast)) => E.A.append(items, [|(x, y +. yLast)|])
           | None => [|(x, y)|]
           }
         )
@@ -46,6 +46,16 @@ module Discrete = {
     fromArrays(xs, ys);
   };
 
+  let ySum = (t: t) => {
+    E.A.fold_left((a, b) => a +. b, 0., t.ys);
+  };
+
+  let scaleYToTotal = (totalDesired, t: t): t => {
+    let currentSum = ySum(t);
+    let difference = totalDesired /. currentSum;
+    {xs: t.xs, ys: t.ys |> E.A.fmap(y => y *. difference)};
+  };
+
   let render = (t: t) =>
     Belt.Array.zip(t.xs, t.ys)
     |> E.A.fmap(((x, y)) =>
@@ -109,11 +119,11 @@ module Mixed = {
       | {
           continuous: ADDS_TO_1,
           discrete: ADDS_TO_CORRECT_PROBABILITY,
-          discreteProbabilityMass: Some(r),
+          discreteProbabilityMass: None,
         } =>
-        Some(
-          make(~continuous, ~discrete, ~discreteProbabilityMassFraction=r),
-        )
+        let discreteProbabilityMassFraction = Discrete.ySum(discrete);
+        let discrete = Discrete.scaleYToTotal(1.0, discrete);
+        Some(make(~continuous, ~discrete, ~discreteProbabilityMassFraction));
       | _ => None
       };
   };
diff --git a/src/utility/Guesstimator.re b/src/utility/Guesstimator.re
index 995616d9..560ca798 100644
--- a/src/utility/Guesstimator.re
+++ b/src/utility/Guesstimator.re
@@ -28,7 +28,7 @@ module Internals = {
     let assumptions: Shape.Mixed.Builder.assumptions = {
       continuous: ADDS_TO_1,
       discrete: ADDS_TO_CORRECT_PROBABILITY,
-      discreteProbabilityMass: Some(0.3),
+      discreteProbabilityMass: None,
     };
     Shape.Mixed.Builder.build(
       ~continuous=toContinous(r),