From 81c733660f8d7d6a449fd4beae81cfd0d6b5348c Mon Sep 17 00:00:00 2001
From: httse9 <62588559+httse9@users.noreply.github.com>
Date: Mon, 1 May 2023 18:29:41 -0400
Subject: [PATCH 1/6] bayesian seldonian

Set mode in spec
---
 .../candidate_selection.py                    |  4 +
 seldonian/parse_tree/mcmc/likelihood.py       | 45 +++++++++
 seldonian/parse_tree/mcmc/mcmc.py             | 96 +++++++++++++++++++
 seldonian/parse_tree/mcmc/prior.py            | 16 ++++
 seldonian/parse_tree/mcmc/proposal.py         | 11 +++
 seldonian/parse_tree/nodes.py                 | 23 ++++-
 seldonian/parse_tree/parse_tree.py            |  5 +-
 seldonian/safety_test/safety_test.py          |  4 +-
 seldonian/seldonian_algorithm.py              |  3 +
 seldonian/spec.py                             |  4 +
 10 files changed, 203 insertions(+), 8 deletions(-)
 create mode 100644 seldonian/parse_tree/mcmc/likelihood.py
 create mode 100644 seldonian/parse_tree/mcmc/mcmc.py
 create mode 100644 seldonian/parse_tree/mcmc/prior.py
 create mode 100644 seldonian/parse_tree/mcmc/proposal.py

diff --git a/seldonian/candidate_selection/candidate_selection.py b/seldonian/candidate_selection/candidate_selection.py
index 70d5aed5..d4ff86ab 100644
--- a/seldonian/candidate_selection/candidate_selection.py
+++ b/seldonian/candidate_selection/candidate_selection.py
@@ -18,6 +18,7 @@ def __init__(
         n_safety,
         parse_trees,
         primary_objective,
+        mode,
         optimization_technique="barrier_function",
         optimizer="Powell",
         initial_solution=None,
@@ -68,6 +69,7 @@ def __init__(
 
         """
         self.regime = regime
+        self.mode = mode
         self.model = model
         self.candidate_dataset = candidate_dataset
         self.n_safety = n_safety
@@ -389,6 +391,7 @@ def objective_with_barrier(self, theta):
                 branch="candidate_selection",
                 n_safety=self.n_safety,
                 regime=self.regime,
+                mode=self.mode
             )
 
             pt.propagate_bounds(**bounds_kwargs)
@@ -479,6 +482,7 @@ def get_constraint_upper_bounds(self, theta):
                 branch="candidate_selection",
                 n_safety=self.n_safety,
                 regime=self.regime,
+                mode=self.mode
             )
 
             pt.propagate_bounds(**bounds_kwargs)
diff --git a/seldonian/parse_tree/mcmc/likelihood.py b/seldonian/parse_tree/mcmc/likelihood.py
new file mode 100644
index 00000000..5e85b54d
--- /dev/null
+++ b/seldonian/parse_tree/mcmc/likelihood.py
@@ -0,0 +1,45 @@
+import autograd.numpy as np
+
+##### Likelihhood Ratio Functions
+def Mean_Squared_Error_Likelihood_Ratio(proposal, original, zhat, std):
+    if std == 0:
+        std = 1e-15
+    return np.exp(- ((zhat - proposal) ** 2 - (zhat - original) ** 2).sum() / 2 / std ** 2)
+
+
+################################
+def get_likelihood_ratio(statistic_name, zhat, datasize):
+    """
+    Function to get likelihood ratio functions from
+    statistic name and zhat
+
+    :param datasize:
+        Size of safety data set.
+    :type dataset: int
+
+    :ivar power:
+        Ratio of size of safety data set
+        to size of zhat
+    :vartype power: float
+    """
+
+    power = (datasize / len(zhat))
+    # print("Power:", power)
+
+    if statistic_name == "Mean_Squared_Error":
+        # use std of zhat as std of normal assumption
+        std = np.std(zhat)
+
+        # Wrap likelihood ratio function such that
+        # it only takes proposal g and original (current) g
+        # as inputs. 
+        def wrapped_likelihood_ratio(proposal, original):
+            # If in candidate selection, scale likelihood ratio
+            # according to size of safety data set.
+            # If in safety test, power is 1.
+            return Mean_Squared_Error_Likelihood_Ratio(proposal, original, zhat, std) ** power
+        
+        return wrapped_likelihood_ratio
+    
+    raise NotImplementedError()
+
diff --git a/seldonian/parse_tree/mcmc/mcmc.py b/seldonian/parse_tree/mcmc/mcmc.py
new file mode 100644
index 00000000..8c23cb30
--- /dev/null
+++ b/seldonian/parse_tree/mcmc/mcmc.py
@@ -0,0 +1,96 @@
+import numpy as np
+from .proposal import ProposalDistribution
+from .prior import PriorDistribution
+from .likelihood import get_likelihood_ratio
+
+class MetropolisHastings:
+
+    def __init__(self, proposal_width, prior_type, prior_mean, prior_width, likelihood_ratio):
+        self.proposal_dist = ProposalDistribution(proposal_width)
+        self.prior_dist = PriorDistribution(prior_type, prior_mean, prior_width)
+        self.likelihood_ratio = likelihood_ratio
+
+    def run(self, N=200000, skip_interval=20, burn_in=5000):
+        verbose = False
+        if verbose:
+            print("Running MCMC:")
+        g = 0 # initial g
+
+        samples = []    # samples obtained
+        all_gs = [g]    # keep track of every g in the chain
+
+        t = 0
+        t_skip = 0
+        acceptance_count = 0
+
+        while True:
+
+            g_proposal = self.proposal_dist.propose(g)
+
+            prior_ratio = self.prior_dist.prior_ratio(g_proposal, g)
+            likelihood_ratio = self.likelihood_ratio(g_proposal, g)
+            accept_ratio = prior_ratio * likelihood_ratio
+            # print(accept_ratio)
+
+            if np.random.uniform() <= accept_ratio:
+                g = g_proposal
+                acceptance_count += 1
+
+            if t_skip == 0 and t >= burn_in :
+                samples.append(g)
+                # print(g)
+
+            all_gs.append(g)
+            t_skip = (t_skip + 1) % skip_interval
+            t += 1
+
+            if t % 10000 == 0:
+                if verbose:
+                    print(f"{t} Steps Completed")
+
+            if t == N:
+                break
+
+        if verbose:
+            print("Acceptance rate:", acceptance_count / N)
+
+        return samples, all_gs
+    
+
+def run_mcmc_default(statistic_name, zhat, datasize, **kwargs):
+    """
+    Default MCMC settings using jeffrey's pior.
+    """
+
+    # --TODO-- automatic tune proposal width
+    # such that the acceptance rate is between
+    # 0.2 and 0.8.
+
+    proposal_width = 0.2
+    if kwargs["branch"] == "candidate_selection":
+        prior_type = "jeffrey"
+        prior_width = None
+        prior_mean = None
+
+    elif kwargs["branch"] == "safety_test":
+        # use candidate zhat mean to construct prior for safety test
+        # prior_type = "normal"
+        # prior_mean = mean_zhat
+        prior_type = "jeffrey"
+        prior_width = None
+        prior_mean = None
+
+
+    likelihood_ratio = get_likelihood_ratio(statistic_name, zhat, datasize)
+
+    mh = MetropolisHastings(proposal_width, prior_type, prior_mean, prior_width, likelihood_ratio)
+    samples, _ = mh.run(N=100000, skip_interval=10, burn_in=3000)
+
+    # import matplotlib.pyplot as plt
+    # plt.hist(samples, bins=100, density=True)
+    # plt.show()
+
+    # print(np.quantile(samples, 0.95, method="inverted_cdf") - 2, )
+
+
+    return samples
diff --git a/seldonian/parse_tree/mcmc/prior.py b/seldonian/parse_tree/mcmc/prior.py
new file mode 100644
index 00000000..4108361d
--- /dev/null
+++ b/seldonian/parse_tree/mcmc/prior.py
@@ -0,0 +1,16 @@
+
+
+class PriorDistribution:
+
+    def __init__(self, type, mean, width):
+        self.type = type
+        self.mean = mean
+        self.width = width
+
+    def prior_ratio(self, proposal, original):
+        if  self.type == "jeffrey":
+            return 1
+        elif self.type == "normal":
+            raise NotImplementedError()
+        else:
+            raise NotImplementedError()
\ No newline at end of file
diff --git a/seldonian/parse_tree/mcmc/proposal.py b/seldonian/parse_tree/mcmc/proposal.py
new file mode 100644
index 00000000..c5ce0ad8
--- /dev/null
+++ b/seldonian/parse_tree/mcmc/proposal.py
@@ -0,0 +1,11 @@
+import numpy as np
+
+class ProposalDistribution:
+    """
+    Normal proposal distribution
+    """
+    def __init__(self, proposal_width):
+        self.proposal_width = proposal_width
+
+    def propose(self, x):
+        return np.random.normal(loc=x, scale=self.proposal_width)
\ No newline at end of file
diff --git a/seldonian/parse_tree/nodes.py b/seldonian/parse_tree/nodes.py
index b3361553..01a21fe9 100644
--- a/seldonian/parse_tree/nodes.py
+++ b/seldonian/parse_tree/nodes.py
@@ -5,7 +5,7 @@
 
 from seldonian.models.objectives import sample_from_statistic, evaluate_statistic
 from seldonian.utils.stats_utils import *
-
+from .mcmc.mcmc import run_mcmc_default
 
 class Node(object):
     def __init__(self, name, lower, upper):
@@ -273,17 +273,20 @@ def calculate_bounds(self, **kwargs):
                 return {"lower": lower, "upper": upper}
 
             else:
-                # Real confidence bound
+                # Real confidence bound / credible bound
 
                 # --TODO-- abstract away to support things like
                 # getting confidence intervals from bootstrap
                 # and RL cases
                 estimator_samples = self.zhat(**kwargs)
+                if kwargs["mode"] == "bayesian":
+                    posterior_samples = run_mcmc_default(self.measure_function_name, estimator_samples, **kwargs)
 
                 branch = kwargs["branch"]
                 data_dict = kwargs["data_dict"]
                 bound_kwargs = kwargs
-                bound_kwargs["data"] = estimator_samples
+                bound_kwargs["data"] = estimator_samples if kwargs["mode"] == "frequentist" else posterior_samples
+                bound_kwargs["bound_method"] = "ttest" if kwargs["mode"] == "frequentist" else "quantile"
                 bound_kwargs["delta"] = self.delta
 
                 # If lower and upper are both needed,
@@ -371,6 +374,8 @@ def predict_HC_lowerbound(self, data, datasize, delta, **kwargs):
                 lower = data.mean() - 2 * stddev(data) / np.sqrt(datasize) * tinv(
                     1.0 - delta, datasize - 1
                 )
+            elif bound_method == "quantile":
+                lower = np.quantile(data, delta, method="inverted_cdf")
             else:
                 raise NotImplementedError(
                     f"Bounding method {bound_method} is not supported"
@@ -401,6 +406,8 @@ def predict_HC_upperbound(self, data, datasize, delta, **kwargs):
                 lower = data.mean() + 2 * stddev(data) / np.sqrt(datasize) * tinv(
                     1.0 - delta, datasize - 1
                 )
+            elif bound_method == "quantile":
+                lower = np.quantile(data, 1 - delta, method="inverted_cdf")
             else:
                 raise NotImplementedError(
                     f"Bounding method {bound_method} is not supported"
@@ -442,6 +449,9 @@ def predict_HC_upper_and_lowerbound(self, data, datasize, delta, **kwargs):
 
             elif bound_method == "manual":
                 pass
+            elif bound_method == "quantile":
+                lower = np.quantile(data, delta / 2, method="inverted_cdf")
+                upper = np.quantile(data, 1 - delta / 2, method="inverted_cdf")
             else:
                 raise NotImplementedError(
                     f"Bounding method {bound_method}" " is not supported"
@@ -471,6 +481,8 @@ def compute_HC_lowerbound(self, data, datasize, delta, **kwargs):
                 lower = data.mean() - stddev(data) / np.sqrt(datasize) * tinv(
                     1.0 - delta, datasize - 1
                 )
+            elif bound_method == "quantile":
+                lower = np.quantile(data, delta, method="inverted_cdf")
             else:
                 raise NotImplementedError(
                     f"Bounding method {bound_method}" " is not supported"
@@ -499,6 +511,8 @@ def compute_HC_upperbound(self, data, datasize, delta, **kwargs):
                 upper = data.mean() + stddev(data) / np.sqrt(datasize) * tinv(
                     1.0 - delta, datasize - 1
                 )
+            elif bound_method == "quantile":
+                upper = np.quantile(data, 1 - delta, method="inverted_cdf")
             else:
                 raise NotImplementedError(
                     f"Bounding method {bound_method}" " is not supported"
@@ -539,6 +553,9 @@ def compute_HC_upper_and_lowerbound(self, data, datasize, delta, **kwargs):
 
             elif bound_method == "manual":
                 pass
+            elif bound_method == "quantile":
+                lower = np.quantile(data, delta / 2, method="inverted_cdf")
+                upper = np.quantile(data, 1 - delta / 2, method="inverted_cdf")
             else:
                 raise NotImplementedError(
                     f"Bounding method {bound_method}" " is not supported"
diff --git a/seldonian/parse_tree/parse_tree.py b/seldonian/parse_tree/parse_tree.py
index 3d21bc2f..82cac0a9 100644
--- a/seldonian/parse_tree/parse_tree.py
+++ b/seldonian/parse_tree/parse_tree.py
@@ -12,9 +12,6 @@
 from .nodes import *
 from .operators import *
 
-default_bound_method = "ttest"
-
-
 class ParseTree(object):
     def __init__(self, delta, regime, sub_regime, columns=[]):
         """
@@ -266,7 +263,7 @@ def _ast_tree_helper(self, ast_node):
             # then make a new entry
             if new_node.name not in self.base_node_dict:
                 self.base_node_dict[new_node.name] = {
-                    "bound_method": default_bound_method,
+                    "bound_method": "default",
                     "bound_computed": False,
                     "value_computed": False,
                     "lower": float("-inf"),
diff --git a/seldonian/safety_test/safety_test.py b/seldonian/safety_test/safety_test.py
index c84f6e5a..02252334 100644
--- a/seldonian/safety_test/safety_test.py
+++ b/seldonian/safety_test/safety_test.py
@@ -6,7 +6,7 @@
 
 class SafetyTest(object):
     def __init__(
-        self, safety_dataset, model, parse_trees, regime="supervised_learning", **kwargs
+        self, safety_dataset, model, parse_trees, mode, regime="supervised_learning", **kwargs
     ):
         """
         Object for running safety test
@@ -29,6 +29,7 @@ def __init__(
         self.model = model
         self.parse_trees = parse_trees
         self.regime = regime
+        self.mode = mode
         self.st_result = {}  # stores parse tree evaluated on safety test data
 
     def run(self, solution, batch_size_safety=None, **kwargs):
@@ -61,6 +62,7 @@ def run(self, solution, batch_size_safety=None, **kwargs):
                 branch="safety_test",
                 regime=self.regime,
                 batch_size_safety=batch_size_safety,
+                mode=self.mode,
                 **kwargs
             )
 
diff --git a/seldonian/seldonian_algorithm.py b/seldonian/seldonian_algorithm.py
index 21b288b2..f8913745 100644
--- a/seldonian/seldonian_algorithm.py
+++ b/seldonian/seldonian_algorithm.py
@@ -22,6 +22,7 @@ def __init__(self, spec):
         :type spec: :py:class:`.Spec` object
         """
         self.spec = spec
+        self.mode = spec.mode
         self.cs_has_been_run = False
         self.cs_result = None
         self.st_has_been_run = False
@@ -189,6 +190,7 @@ def candidate_selection(self, write_logfile=False):
             initial_solution=self.initial_solution,
             regime=self.regime,
             write_logfile=write_logfile,
+            mode=self.mode
         )
 
         cs = CandidateSelection(**cs_kwargs, **self.spec.regularization_hyperparams)
@@ -202,6 +204,7 @@ def safety_test(self):
             model=self.model,
             parse_trees=self.spec.parse_trees,
             regime=self.regime,
+            mode=self.mode
         )
 
         st = SafetyTest(**st_kwargs)
diff --git a/seldonian/spec.py b/seldonian/spec.py
index 177fbb24..84d888bc 100644
--- a/seldonian/spec.py
+++ b/seldonian/spec.py
@@ -81,6 +81,7 @@ def __init__(
         regularization_hyperparams={},
         batch_size_safety=None,
         verbose=False,
+        mode="frequentist",
     ):
         self.dataset = dataset
         self.model = model
@@ -97,6 +98,7 @@ def __init__(
         self.regularization_hyperparams = regularization_hyperparams
         self.batch_size_safety = batch_size_safety
         self.verbose = verbose
+        self.mode = mode
 
 
 class SupervisedSpec(Spec):
@@ -171,6 +173,7 @@ def __init__(
         regularization_hyperparams={},
         batch_size_safety=None,
         verbose=False,
+        mode="frequentist",
     ):
         super().__init__(
             dataset=dataset,
@@ -188,6 +191,7 @@ def __init__(
             regularization_hyperparams=regularization_hyperparams,
             batch_size_safety=batch_size_safety,
             verbose=verbose,
+            mode=mode,
         )
         self.sub_regime = sub_regime
 

From 670cd408713ee310ff6866b57df2ffdaa4c818c4 Mon Sep 17 00:00:00 2001
From: httse9 <62588559+httse9@users.noreply.github.com>
Date: Tue, 2 May 2023 12:33:46 -0400
Subject: [PATCH 2/6] hack like in frequentist seldonian

---
 seldonian/parse_tree/mcmc/likelihood.py |  9 +++++++--
 seldonian/parse_tree/mcmc/mcmc.py       | 10 ++++++----
 seldonian/parse_tree/nodes.py           |  8 ++++----
 3 files changed, 17 insertions(+), 10 deletions(-)

diff --git a/seldonian/parse_tree/mcmc/likelihood.py b/seldonian/parse_tree/mcmc/likelihood.py
index 5e85b54d..1d7e1ae8 100644
--- a/seldonian/parse_tree/mcmc/likelihood.py
+++ b/seldonian/parse_tree/mcmc/likelihood.py
@@ -27,8 +27,13 @@ def get_likelihood_ratio(statistic_name, zhat, datasize):
     # print("Power:", power)
 
     if statistic_name == "Mean_Squared_Error":
-        # use std of zhat as std of normal assumption
-        std = np.std(zhat)
+        # Use std of zhat as std of normal assumption
+        # Scale std of zhat to predict what would be
+        # the std for safety data (replace denominator
+        # of variance with size of safety set instead of
+        # size of candidate set)
+        std = np.std(zhat) * np.sqrt(len(zhat) / datasize)
+        # print(std)
 
         # Wrap likelihood ratio function such that
         # it only takes proposal g and original (current) g
diff --git a/seldonian/parse_tree/mcmc/mcmc.py b/seldonian/parse_tree/mcmc/mcmc.py
index 8c23cb30..7ff418f4 100644
--- a/seldonian/parse_tree/mcmc/mcmc.py
+++ b/seldonian/parse_tree/mcmc/mcmc.py
@@ -86,11 +86,13 @@ def run_mcmc_default(statistic_name, zhat, datasize, **kwargs):
     mh = MetropolisHastings(proposal_width, prior_type, prior_mean, prior_width, likelihood_ratio)
     samples, _ = mh.run(N=100000, skip_interval=10, burn_in=3000)
 
-    # import matplotlib.pyplot as plt
-    # plt.hist(samples, bins=100, density=True)
-    # plt.show()
+    # if kwargs["branch"] == "safety_test":
+    #     import matplotlib.pyplot as plt
+    #     plt.hist(samples, bins=100, density=True)
+    #     plt.axvline(np.quantile(samples, 0.9, method="inverted_cdf"))
+    #     plt.show()
 
-    # print(np.quantile(samples, 0.95, method="inverted_cdf") - 2, )
+    # print(np.quantile(samples, 0.9, method="inverted_cdf") )
 
 
     return samples
diff --git a/seldonian/parse_tree/nodes.py b/seldonian/parse_tree/nodes.py
index 01a21fe9..e9f609f7 100644
--- a/seldonian/parse_tree/nodes.py
+++ b/seldonian/parse_tree/nodes.py
@@ -375,7 +375,7 @@ def predict_HC_lowerbound(self, data, datasize, delta, **kwargs):
                     1.0 - delta, datasize - 1
                 )
             elif bound_method == "quantile":
-                lower = np.quantile(data, delta, method="inverted_cdf")
+                lower = np.quantile(data, delta / 2, method="inverted_cdf")
             else:
                 raise NotImplementedError(
                     f"Bounding method {bound_method} is not supported"
@@ -407,7 +407,7 @@ def predict_HC_upperbound(self, data, datasize, delta, **kwargs):
                     1.0 - delta, datasize - 1
                 )
             elif bound_method == "quantile":
-                lower = np.quantile(data, 1 - delta, method="inverted_cdf")
+                lower = np.quantile(data, 1 - delta / 2, method="inverted_cdf")
             else:
                 raise NotImplementedError(
                     f"Bounding method {bound_method} is not supported"
@@ -450,8 +450,8 @@ def predict_HC_upper_and_lowerbound(self, data, datasize, delta, **kwargs):
             elif bound_method == "manual":
                 pass
             elif bound_method == "quantile":
-                lower = np.quantile(data, delta / 2, method="inverted_cdf")
-                upper = np.quantile(data, 1 - delta / 2, method="inverted_cdf")
+                lower = np.quantile(data, delta / 4, method="inverted_cdf")
+                upper = np.quantile(data, 1 - delta / 4, method="inverted_cdf")
             else:
                 raise NotImplementedError(
                     f"Bounding method {bound_method}" " is not supported"

From 960b3c8b55919bab8d7e65ef8cafcd27d751940b Mon Sep 17 00:00:00 2001
From: httse9 <62588559+httse9@users.noreply.github.com>
Date: Tue, 2 May 2023 15:01:03 -0400
Subject: [PATCH 3/6] infer std, remove hack like frequentist

---
 seldonian/parse_tree/mcmc/likelihood.py | 67 +++++++++++++++++--------
 seldonian/parse_tree/mcmc/mcmc.py       | 39 +++++++++++---
 seldonian/parse_tree/mcmc/prior.py      |  9 +++-
 seldonian/parse_tree/mcmc/proposal.py   |  1 +
 seldonian/parse_tree/nodes.py           |  8 +--
 5 files changed, 89 insertions(+), 35 deletions(-)

diff --git a/seldonian/parse_tree/mcmc/likelihood.py b/seldonian/parse_tree/mcmc/likelihood.py
index 1d7e1ae8..203f0177 100644
--- a/seldonian/parse_tree/mcmc/likelihood.py
+++ b/seldonian/parse_tree/mcmc/likelihood.py
@@ -1,21 +1,41 @@
 import autograd.numpy as np
 
-##### Likelihhood Ratio Functions
+###### Likelihhood Ratio Functions
 def Mean_Squared_Error_Likelihood_Ratio(proposal, original, zhat, std):
     if std == 0:
         std = 1e-15
     return np.exp(- ((zhat - proposal) ** 2 - (zhat - original) ** 2).sum() / 2 / std ** 2)
 
+##### Likelihood Ratio Functions when also infer std
+def Mean_Squared_Error_Likelihood_Ratio_Infer_Std(proposal, original, zhat):
+    """
+    :param proposal:
+        Proposal parameters. [mean, std]
+    :type List(float)
+
+    :param original:
+        Current parameters. [mean, std]
+    :type List(float)
+    """
+    mean_prop, std_prop = proposal
+    mean_orig, std_orig = original
+
+    likelihood_ratio = np.exp(-(((zhat - mean_prop) / std_prop) ** 2 - ((zhat - mean_orig) / std_orig) ** 2).sum() / 2 + np.log(std_orig / std_prop) * len(zhat))
+    return likelihood_ratio
 
 ################################
-def get_likelihood_ratio(statistic_name, zhat, datasize):
+def get_likelihood_ratio(statistic_name, zhat, datasize, infer_std):
     """
     Function to get likelihood ratio functions from
     statistic name and zhat
 
     :param datasize:
         Size of safety data set.
-    :type dataset: int
+    :type datasize: int
+
+    :param infer_std:
+        Whether to infer std
+    :type infer_std: bool
 
     :ivar power:
         Ratio of size of safety data set
@@ -27,24 +47,29 @@ def get_likelihood_ratio(statistic_name, zhat, datasize):
     # print("Power:", power)
 
     if statistic_name == "Mean_Squared_Error":
-        # Use std of zhat as std of normal assumption
-        # Scale std of zhat to predict what would be
-        # the std for safety data (replace denominator
-        # of variance with size of safety set instead of
-        # size of candidate set)
-        std = np.std(zhat) * np.sqrt(len(zhat) / datasize)
-        # print(std)
-
-        # Wrap likelihood ratio function such that
-        # it only takes proposal g and original (current) g
-        # as inputs. 
-        def wrapped_likelihood_ratio(proposal, original):
-            # If in candidate selection, scale likelihood ratio
-            # according to size of safety data set.
-            # If in safety test, power is 1.
-            return Mean_Squared_Error_Likelihood_Ratio(proposal, original, zhat, std) ** power
-        
-        return wrapped_likelihood_ratio
+        if infer_std:
+            def wrapped_likelihood_ratio(proposal, original):
+                return Mean_Squared_Error_Likelihood_Ratio_Infer_Std(proposal, original, zhat)
+            return wrapped_likelihood_ratio
+        else:
+            # Use std of zhat as std of normal assumption
+            # Scale std of zhat to predict what would be
+            # the std for safety data (replace denominator
+            # of variance with size of safety set instead of
+            # size of candidate set)
+            std = np.std(zhat) * np.sqrt(len(zhat) / datasize)
+            # print(std)
+
+            # Wrap likelihood ratio function such that
+            # it only takes proposal g and original (current) g
+            # as inputs. 
+            def wrapped_likelihood_ratio(proposal, original):
+                # If in candidate selection, scale likelihood ratio
+                # according to size of safety data set.
+                # If in safety test, power is 1.
+                return Mean_Squared_Error_Likelihood_Ratio(proposal, original, zhat, std) ** power
+            
+            return wrapped_likelihood_ratio
     
     raise NotImplementedError()
 
diff --git a/seldonian/parse_tree/mcmc/mcmc.py b/seldonian/parse_tree/mcmc/mcmc.py
index 7ff418f4..f0538735 100644
--- a/seldonian/parse_tree/mcmc/mcmc.py
+++ b/seldonian/parse_tree/mcmc/mcmc.py
@@ -5,16 +5,22 @@
 
 class MetropolisHastings:
 
-    def __init__(self, proposal_width, prior_type, prior_mean, prior_width, likelihood_ratio):
+    def __init__(self, proposal_width, prior_type, prior_mean, prior_width, likelihood_ratio, infer_std):
         self.proposal_dist = ProposalDistribution(proposal_width)
-        self.prior_dist = PriorDistribution(prior_type, prior_mean, prior_width)
+        self.prior_dist = PriorDistribution(prior_type, prior_mean, prior_width, infer_std)
         self.likelihood_ratio = likelihood_ratio
+        self.infer_std = infer_std
 
     def run(self, N=200000, skip_interval=20, burn_in=5000):
         verbose = False
         if verbose:
             print("Running MCMC:")
-        g = 0 # initial g
+        
+        if self.infer_std:
+            # mean and std
+            g = np.array([0, 5])    # start with N(0, 1)
+        else:
+            g = 0 # initial g
 
         samples = []    # samples obtained
         all_gs = [g]    # keep track of every g in the chain
@@ -51,8 +57,13 @@ def run(self, N=200000, skip_interval=20, burn_in=5000):
             if t == N:
                 break
 
-        if verbose:
-            print("Acceptance rate:", acceptance_count / N)
+        # if verbose:
+        print("Acceptance rate:", acceptance_count / N)
+
+        if self.infer_std:
+            # keep only mean
+            # print(np.mean(np.array(samples)[:, 1]))
+            samples = np.array(samples)[:, 0]
 
         return samples, all_gs
     
@@ -60,13 +71,25 @@ def run(self, N=200000, skip_interval=20, burn_in=5000):
 def run_mcmc_default(statistic_name, zhat, datasize, **kwargs):
     """
     Default MCMC settings using jeffrey's pior.
+
+    :ivar infer_std:
+        If true, infer std using mcmc.
+        Else, use std derived from candidate
+        data.
+    :vartype infer_std: bool
     """
 
     # --TODO-- automatic tune proposal width
     # such that the acceptance rate is between
     # 0.2 and 0.8.
 
-    proposal_width = 0.2
+    infer_std = True
+
+    if infer_std:
+        proposal_width = 0.1
+    else:
+        proposal_width = 0.2
+
     if kwargs["branch"] == "candidate_selection":
         prior_type = "jeffrey"
         prior_width = None
@@ -81,9 +104,9 @@ def run_mcmc_default(statistic_name, zhat, datasize, **kwargs):
         prior_mean = None
 
 
-    likelihood_ratio = get_likelihood_ratio(statistic_name, zhat, datasize)
+    likelihood_ratio = get_likelihood_ratio(statistic_name, zhat, datasize, infer_std)
 
-    mh = MetropolisHastings(proposal_width, prior_type, prior_mean, prior_width, likelihood_ratio)
+    mh = MetropolisHastings(proposal_width, prior_type, prior_mean, prior_width, likelihood_ratio, infer_std)
     samples, _ = mh.run(N=100000, skip_interval=10, burn_in=3000)
 
     # if kwargs["branch"] == "safety_test":
diff --git a/seldonian/parse_tree/mcmc/prior.py b/seldonian/parse_tree/mcmc/prior.py
index 4108361d..78ecc82b 100644
--- a/seldonian/parse_tree/mcmc/prior.py
+++ b/seldonian/parse_tree/mcmc/prior.py
@@ -2,14 +2,19 @@
 
 class PriorDistribution:
 
-    def __init__(self, type, mean, width):
+    def __init__(self, type, mean, width, infer_std):
         self.type = type
         self.mean = mean
         self.width = width
+        self.infer_std = infer_std
 
     def prior_ratio(self, proposal, original):
         if  self.type == "jeffrey":
-            return 1
+            if self.infer_std:
+                # prior(mu, sigma) = 1 / sigma ** 2
+                return (original[1] / proposal[1]) ** 2
+            else:
+                return 1
         elif self.type == "normal":
             raise NotImplementedError()
         else:
diff --git a/seldonian/parse_tree/mcmc/proposal.py b/seldonian/parse_tree/mcmc/proposal.py
index c5ce0ad8..a4e2345a 100644
--- a/seldonian/parse_tree/mcmc/proposal.py
+++ b/seldonian/parse_tree/mcmc/proposal.py
@@ -8,4 +8,5 @@ def __init__(self, proposal_width):
         self.proposal_width = proposal_width
 
     def propose(self, x):
+        # supports multidimension x
         return np.random.normal(loc=x, scale=self.proposal_width)
\ No newline at end of file
diff --git a/seldonian/parse_tree/nodes.py b/seldonian/parse_tree/nodes.py
index e9f609f7..01a21fe9 100644
--- a/seldonian/parse_tree/nodes.py
+++ b/seldonian/parse_tree/nodes.py
@@ -375,7 +375,7 @@ def predict_HC_lowerbound(self, data, datasize, delta, **kwargs):
                     1.0 - delta, datasize - 1
                 )
             elif bound_method == "quantile":
-                lower = np.quantile(data, delta / 2, method="inverted_cdf")
+                lower = np.quantile(data, delta, method="inverted_cdf")
             else:
                 raise NotImplementedError(
                     f"Bounding method {bound_method} is not supported"
@@ -407,7 +407,7 @@ def predict_HC_upperbound(self, data, datasize, delta, **kwargs):
                     1.0 - delta, datasize - 1
                 )
             elif bound_method == "quantile":
-                lower = np.quantile(data, 1 - delta / 2, method="inverted_cdf")
+                lower = np.quantile(data, 1 - delta, method="inverted_cdf")
             else:
                 raise NotImplementedError(
                     f"Bounding method {bound_method} is not supported"
@@ -450,8 +450,8 @@ def predict_HC_upper_and_lowerbound(self, data, datasize, delta, **kwargs):
             elif bound_method == "manual":
                 pass
             elif bound_method == "quantile":
-                lower = np.quantile(data, delta / 4, method="inverted_cdf")
-                upper = np.quantile(data, 1 - delta / 4, method="inverted_cdf")
+                lower = np.quantile(data, delta / 2, method="inverted_cdf")
+                upper = np.quantile(data, 1 - delta / 2, method="inverted_cdf")
             else:
                 raise NotImplementedError(
                     f"Bounding method {bound_method}" " is not supported"

From 3bb2afdd591061999868e58b5c31db332f85aa66 Mon Sep 17 00:00:00 2001
From: httse9 <62588559+httse9@users.noreply.github.com>
Date: Thu, 4 May 2023 15:35:08 -0400
Subject: [PATCH 4/6] use candidate prior

use candidate data mean zhat to costruct prior for safety test
---
 .../candidate_selection.py                    |  3 +-
 seldonian/parse_tree/mcmc/mcmc.py             | 36 +++++++++++++------
 seldonian/parse_tree/mcmc/prior.py            |  7 ++--
 seldonian/parse_tree/nodes.py                 | 14 ++++++++
 seldonian/seldonian_algorithm.py              |  7 ++--
 seldonian/spec.py                             |  4 +++
 6 files changed, 55 insertions(+), 16 deletions(-)

diff --git a/seldonian/candidate_selection/candidate_selection.py b/seldonian/candidate_selection/candidate_selection.py
index d4ff86ab..b4824b01 100644
--- a/seldonian/candidate_selection/candidate_selection.py
+++ b/seldonian/candidate_selection/candidate_selection.py
@@ -391,7 +391,8 @@ def objective_with_barrier(self, theta):
                 branch="candidate_selection",
                 n_safety=self.n_safety,
                 regime=self.regime,
-                mode=self.mode
+                mode=self.mode,
+                use_candidate_prior=False
             )
 
             pt.propagate_bounds(**bounds_kwargs)
diff --git a/seldonian/parse_tree/mcmc/mcmc.py b/seldonian/parse_tree/mcmc/mcmc.py
index f0538735..a2b31be4 100644
--- a/seldonian/parse_tree/mcmc/mcmc.py
+++ b/seldonian/parse_tree/mcmc/mcmc.py
@@ -83,7 +83,7 @@ def run_mcmc_default(statistic_name, zhat, datasize, **kwargs):
     # such that the acceptance rate is between
     # 0.2 and 0.8.
 
-    infer_std = True
+    infer_std = False
 
     if infer_std:
         proposal_width = 0.1
@@ -97,11 +97,15 @@ def run_mcmc_default(statistic_name, zhat, datasize, **kwargs):
 
     elif kwargs["branch"] == "safety_test":
         # use candidate zhat mean to construct prior for safety test
-        # prior_type = "normal"
-        # prior_mean = mean_zhat
-        prior_type = "jeffrey"
-        prior_width = None
-        prior_mean = None
+        if kwargs["use_candidate_prior"]:
+            prior_type = "normal"
+            prior_mean = kwargs["zhat_mean"]
+            # print(prior_mean)
+            prior_width = 0.5
+        else:
+            prior_type = "jeffrey"
+            prior_width = None
+            prior_mean = None
 
 
     likelihood_ratio = get_likelihood_ratio(statistic_name, zhat, datasize, infer_std)
@@ -109,11 +113,23 @@ def run_mcmc_default(statistic_name, zhat, datasize, **kwargs):
     mh = MetropolisHastings(proposal_width, prior_type, prior_mean, prior_width, likelihood_ratio, infer_std)
     samples, _ = mh.run(N=100000, skip_interval=10, burn_in=3000)
 
+    
+
     # if kwargs["branch"] == "safety_test":
-    #     import matplotlib.pyplot as plt
-    #     plt.hist(samples, bins=100, density=True)
-    #     plt.axvline(np.quantile(samples, 0.9, method="inverted_cdf"))
-    #     plt.show()
+    # import matplotlib.pyplot as plt
+    # from scipy.stats import t
+    # plt.hist(samples, bins=100, density=True, label="posterior", alpha=0.5)
+    # plt.axvline(np.quantile(samples, 0.9, method="inverted_cdf"))
+
+    # print(np.quantile(samples, 0.95), np.mean(zhat) + np.std(zhat) / np.sqrt(datasize) * t.ppf(0.95, datasize - 1))
+
+    # normal_samples = np.random.normal(loc=np.mean(zhat), scale=np.std(zhat) / np.sqrt(datasize), size=(10000,))
+    # plt.hist(normal_samples, bins=100, density=True, label="gaussian", alpha=0.5)
+
+    # plt.hist(zhat, bins=100, density=True, label="zhat", alpha=0.5)
+    # plt.legend()
+    # plt.show()
+
 
     # print(np.quantile(samples, 0.9, method="inverted_cdf") )
 
diff --git a/seldonian/parse_tree/mcmc/prior.py b/seldonian/parse_tree/mcmc/prior.py
index 78ecc82b..5b17042d 100644
--- a/seldonian/parse_tree/mcmc/prior.py
+++ b/seldonian/parse_tree/mcmc/prior.py
@@ -1,4 +1,4 @@
-
+import numpy as np
 
 class PriorDistribution:
 
@@ -16,6 +16,9 @@ def prior_ratio(self, proposal, original):
             else:
                 return 1
         elif self.type == "normal":
-            raise NotImplementedError()
+            if self.infer_std:
+                raise NotImplementedError()
+            else:
+                return np.exp(- ((proposal - self.mean) ** 2 - (original - self.mean) ** 2) / 2 / self.width ** 2)
         else:
             raise NotImplementedError()
\ No newline at end of file
diff --git a/seldonian/parse_tree/nodes.py b/seldonian/parse_tree/nodes.py
index 01a21fe9..27db7ddc 100644
--- a/seldonian/parse_tree/nodes.py
+++ b/seldonian/parse_tree/nodes.py
@@ -6,6 +6,7 @@
 from seldonian.models.objectives import sample_from_statistic, evaluate_statistic
 from seldonian.utils.stats_utils import *
 from .mcmc.mcmc import run_mcmc_default
+import copy
 
 class Node(object):
     def __init__(self, name, lower, upper):
@@ -280,6 +281,19 @@ def calculate_bounds(self, **kwargs):
                 # and RL cases
                 estimator_samples = self.zhat(**kwargs)
                 if kwargs["mode"] == "bayesian":
+                    if kwargs["use_candidate_prior"]:
+                        # Use candidate data set to get a prior for MCMC
+                        candidate_kwargs = copy.deepcopy(kwargs)
+                        candidate_kwargs["dataset"] = kwargs["candidate_dataset"]
+                        candidate_kwargs["branch"] = "candidate_selection"
+                        candidate_kwargs["n_safety"] = kwargs["datasize"]
+                        candidate_data_dict, _ = self.calculate_data_forbound(**candidate_kwargs)
+                        candidate_kwargs["data_dict"] = candidate_data_dict
+                        candidate_kwargs["datasize"] = None     # doesn't matter for evaluate_statistics
+
+                        value = self.calculate_value(**candidate_kwargs)
+                        kwargs["zhat_mean"] = value
+
                     posterior_samples = run_mcmc_default(self.measure_function_name, estimator_samples, **kwargs)
 
                 branch = kwargs["branch"]
diff --git a/seldonian/seldonian_algorithm.py b/seldonian/seldonian_algorithm.py
index f8913745..fa3d1298 100644
--- a/seldonian/seldonian_algorithm.py
+++ b/seldonian/seldonian_algorithm.py
@@ -204,7 +204,7 @@ def safety_test(self):
             model=self.model,
             parse_trees=self.spec.parse_trees,
             regime=self.regime,
-            mode=self.mode
+            mode=self.mode,
         )
 
         st = SafetyTest(**st_kwargs)
@@ -293,7 +293,7 @@ def run_candidate_selection(self, write_logfile=False, debug=False):
             **self.spec.optimization_hyperparams,
             use_builtin_primary_gradient_fn=self.spec.use_builtin_primary_gradient_fn,
             custom_primary_gradient_fn=self.spec.custom_primary_gradient_fn,
-            debug=debug,
+            debug=debug
         )
 
         self.cs_has_been_run = True
@@ -316,7 +316,8 @@ def run_safety_test(self, candidate_solution, batch_size_safety=None, debug=Fals
         """
 
         st = self.safety_test()
-        passed_safety = st.run(candidate_solution, batch_size_safety=batch_size_safety)
+        passed_safety = st.run(candidate_solution, batch_size_safety=batch_size_safety, 
+                               candidate_dataset=self.candidate_dataset, use_candidate_prior=self.spec.use_candidate_prior)
         if not passed_safety:
             if debug:
                 print("Failed safety test")
diff --git a/seldonian/spec.py b/seldonian/spec.py
index 84d888bc..d8ef9fab 100644
--- a/seldonian/spec.py
+++ b/seldonian/spec.py
@@ -82,6 +82,7 @@ def __init__(
         batch_size_safety=None,
         verbose=False,
         mode="frequentist",
+        use_candidate_prior=False
     ):
         self.dataset = dataset
         self.model = model
@@ -99,6 +100,7 @@ def __init__(
         self.batch_size_safety = batch_size_safety
         self.verbose = verbose
         self.mode = mode
+        self.use_candidate_prior = use_candidate_prior
 
 
 class SupervisedSpec(Spec):
@@ -174,6 +176,7 @@ def __init__(
         batch_size_safety=None,
         verbose=False,
         mode="frequentist",
+        use_candidate_prior=False
     ):
         super().__init__(
             dataset=dataset,
@@ -192,6 +195,7 @@ def __init__(
             batch_size_safety=batch_size_safety,
             verbose=verbose,
             mode=mode,
+            use_candidate_prior=use_candidate_prior
         )
         self.sub_regime = sub_regime
 

From ba2db93c6bff28d5d0ee63bc28a22c8950aaddfb Mon Sep 17 00:00:00 2001
From: httse9 <62588559+httse9@users.noreply.github.com>
Date: Fri, 5 May 2023 05:17:31 -0400
Subject: [PATCH 5/6] infer std fix bug

make sure proposed std > 0
---
 seldonian/parse_tree/mcmc/likelihood.py |  2 ++
 seldonian/parse_tree/mcmc/mcmc.py       |  7 ++-----
 seldonian/parse_tree/mcmc/proposal.py   | 12 ++++++++++--
 3 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/seldonian/parse_tree/mcmc/likelihood.py b/seldonian/parse_tree/mcmc/likelihood.py
index 203f0177..b38e491c 100644
--- a/seldonian/parse_tree/mcmc/likelihood.py
+++ b/seldonian/parse_tree/mcmc/likelihood.py
@@ -4,6 +4,7 @@
 def Mean_Squared_Error_Likelihood_Ratio(proposal, original, zhat, std):
     if std == 0:
         std = 1e-15
+    # print(proposal, original, np.exp(- ((zhat - proposal) ** 2 - (zhat - original) ** 2).sum() / 2 / std ** 2))
     return np.exp(- ((zhat - proposal) ** 2 - (zhat - original) ** 2).sum() / 2 / std ** 2)
 
 ##### Likelihood Ratio Functions when also infer std
@@ -21,6 +22,7 @@ def Mean_Squared_Error_Likelihood_Ratio_Infer_Std(proposal, original, zhat):
     mean_orig, std_orig = original
 
     likelihood_ratio = np.exp(-(((zhat - mean_prop) / std_prop) ** 2 - ((zhat - mean_orig) / std_orig) ** 2).sum() / 2 + np.log(std_orig / std_prop) * len(zhat))
+    # print(likelihood_ratio)
     return likelihood_ratio
 
 ################################
diff --git a/seldonian/parse_tree/mcmc/mcmc.py b/seldonian/parse_tree/mcmc/mcmc.py
index a2b31be4..6823bbfa 100644
--- a/seldonian/parse_tree/mcmc/mcmc.py
+++ b/seldonian/parse_tree/mcmc/mcmc.py
@@ -6,7 +6,7 @@
 class MetropolisHastings:
 
     def __init__(self, proposal_width, prior_type, prior_mean, prior_width, likelihood_ratio, infer_std):
-        self.proposal_dist = ProposalDistribution(proposal_width)
+        self.proposal_dist = ProposalDistribution(proposal_width, infer_std)
         self.prior_dist = PriorDistribution(prior_type, prior_mean, prior_width, infer_std)
         self.likelihood_ratio = likelihood_ratio
         self.infer_std = infer_std
@@ -85,10 +85,7 @@ def run_mcmc_default(statistic_name, zhat, datasize, **kwargs):
 
     infer_std = False
 
-    if infer_std:
-        proposal_width = 0.1
-    else:
-        proposal_width = 0.2
+    proposal_width = 0.2
 
     if kwargs["branch"] == "candidate_selection":
         prior_type = "jeffrey"
diff --git a/seldonian/parse_tree/mcmc/proposal.py b/seldonian/parse_tree/mcmc/proposal.py
index a4e2345a..576db8f8 100644
--- a/seldonian/parse_tree/mcmc/proposal.py
+++ b/seldonian/parse_tree/mcmc/proposal.py
@@ -4,9 +4,17 @@ class ProposalDistribution:
     """
     Normal proposal distribution
     """
-    def __init__(self, proposal_width):
+    def __init__(self, proposal_width, infer_std):
         self.proposal_width = proposal_width
+        self.infer_std = infer_std
 
     def propose(self, x):
         # supports multidimension x
-        return np.random.normal(loc=x, scale=self.proposal_width)
\ No newline at end of file
+        proposal = np.random.normal(loc=x, scale=self.proposal_width)
+
+        if self.infer_std:
+            # make sure std > 0
+            if proposal[1] <= 0:
+                proposal[1] = 1e-10
+        
+        return proposal
\ No newline at end of file

From 8c51ec2f52012fd4f26bd306e4a7659af428e06f Mon Sep 17 00:00:00 2001
From: Hon Tik Tse <htse@kingkong.cs.umass.edu>
Date: Fri, 5 May 2023 05:17:47 -0400
Subject: [PATCH 6/6] infer std True

---
 seldonian/candidate_selection/candidate_selection.py | 1 -
 seldonian/parse_tree/mcmc/mcmc.py                    | 6 +++---
 seldonian/parse_tree/nodes.py                        | 1 -
 seldonian/seldonian_algorithm.py                     | 1 +
 4 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/seldonian/candidate_selection/candidate_selection.py b/seldonian/candidate_selection/candidate_selection.py
index b4824b01..2208e592 100644
--- a/seldonian/candidate_selection/candidate_selection.py
+++ b/seldonian/candidate_selection/candidate_selection.py
@@ -394,7 +394,6 @@ def objective_with_barrier(self, theta):
                 mode=self.mode,
                 use_candidate_prior=False
             )
-
             pt.propagate_bounds(**bounds_kwargs)
 
             # Check if the i-th behavioral constraint is satisfied
diff --git a/seldonian/parse_tree/mcmc/mcmc.py b/seldonian/parse_tree/mcmc/mcmc.py
index a2b31be4..7caafca5 100644
--- a/seldonian/parse_tree/mcmc/mcmc.py
+++ b/seldonian/parse_tree/mcmc/mcmc.py
@@ -83,7 +83,7 @@ def run_mcmc_default(statistic_name, zhat, datasize, **kwargs):
     # such that the acceptance rate is between
     # 0.2 and 0.8.
 
-    infer_std = False
+    infer_std = True
 
     if infer_std:
         proposal_width = 0.1
@@ -100,8 +100,8 @@ def run_mcmc_default(statistic_name, zhat, datasize, **kwargs):
         if kwargs["use_candidate_prior"]:
             prior_type = "normal"
             prior_mean = kwargs["zhat_mean"]
-            # print(prior_mean)
-            prior_width = 0.5
+            print("HERE", prior_mean)
+            prior_width = 0.1
         else:
             prior_type = "jeffrey"
             prior_width = None
diff --git a/seldonian/parse_tree/nodes.py b/seldonian/parse_tree/nodes.py
index 27db7ddc..69489bc7 100644
--- a/seldonian/parse_tree/nodes.py
+++ b/seldonian/parse_tree/nodes.py
@@ -293,7 +293,6 @@ def calculate_bounds(self, **kwargs):
 
                         value = self.calculate_value(**candidate_kwargs)
                         kwargs["zhat_mean"] = value
-
                     posterior_samples = run_mcmc_default(self.measure_function_name, estimator_samples, **kwargs)
 
                 branch = kwargs["branch"]
diff --git a/seldonian/seldonian_algorithm.py b/seldonian/seldonian_algorithm.py
index fa3d1298..46684710 100644
--- a/seldonian/seldonian_algorithm.py
+++ b/seldonian/seldonian_algorithm.py
@@ -28,6 +28,7 @@ def __init__(self, spec):
         self.st_has_been_run = False
         self.st_result = None
 
+
         self.parse_trees = self.spec.parse_trees
         # user can pass a dictionary that specifies
         # the bounding method for each base node