<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 The Shower Problem
This notebook is one of the examples in the second edition of Think Bayes.
Click here to run this notebook on Colab.
Recently I posted an article about an optimal (maybe) strategy for plugging in a USB connector.
I tweeted about the article and got a reply from Cameron Davidson-Pilon that pointed me to a related problem posted by Chris Said
Step1:  Starting with $t_1=0$ and $t_2=0$, when $x=1$, the time we wait turns out to be $\lambda$.
Step2:  Now suppose we've waited 50 seconds on the left and we switch to the right. How long should we stay?
Step3:  And now if we've spent left second on the left and switch to the right, how long should we stay.
Step4:  But all of that is based on the assumption that x=1. Let's see what happens to the average total time as we vary x.
Simulation
The following function simulates the strategy I just outlined.
It takes as parameters
Step5:  Here are a few test runs with x=0.5 and various values of correct and tau.
Step6:  Let's get the test data from Chris.
Step7:  The following function uses the samples to run simulations and returns the average total time.
Step8:  I'll train the algorithm with a different sample from the same distribution.
Step9:  Here's a test run with x=1.
Step10:  On average it takes about two tries and 116 seconds.
Step13:  Optimization
Now let's see how the average duration varies as we sweep through a range of values for the threshold probability, r
Step14:  Here's what the results look like.
Step15:  It looks like the optimal value of x is about 2.
Testing
Now let's run the simulations with the optimal value of x and see what the average total time is for the test data.
Step16:  The average total time is about 102 seconds, so it looks like this strategy is better than the current leader (assuming I haven't made a mistake in the simulation).
Step17:  Here's the distribution of total time, represented as a CDF.
Step18:  And here's the distribution of flips.
<ASSISTANT_TASK:>
Python Code:
def wait_time(x, other, lam, k):
    term = (other / lam) ** k
    total = lam * (x + term) ** (1/k)
    return total
Explanation: The Shower Problem
This notebook is one of the examples in the second edition of Think Bayes.
Click here to run this notebook on Colab.
Recently I posted an article about an optimal (maybe) strategy for plugging in a USB connector.
I tweeted about the article and got a reply from Cameron Davidson-Pilon that pointed me to a related problem posted by Chris Said:
The Shower Problem
Here’s the setup: You’re at your friend’s place and you need to take a shower. The shower knob is unlabeled. One direction is hot and the other direction is cold, and you don’t know which is which.
You turn it to the left. It’s cold. You wait.
At what point do you switch over to the right?
Let’s make this more explicit.
Your goal is to find a policy that minimizes the expected amount of time it takes to get hot water flowing out of the shower head. To simplify things, assume that the water coming out of the head is either hot or cold, and that the lukewarm transition time is effectively zero.
You know that the shower has a Time-To-Hot constant called $\tau$. This value is defined as the time it takes for hot water to arrive, assuming you have turned the knob to the hot direction and keep it there.
The constant is a fixed property of the shower and is sampled once from a known distribution. You have certain knowledge of the distribution, but you don’t know $\tau$.
The shower is memoryless, such that every time you turn the knob to the hot direction, it will take $\tau$ seconds until the hot water arrives, regardless of your prior actions. Every time you turn it to the cold direction, only cold water will come out.
I don’t know how to solve this problem. But as a starting point I realize it’s possible to keep track of the probability that the hot direction is to the left or to the right. In the animation above, the probability that the hot direction is to the right is just the unexplored white area under the right curve, divided by the total unexplored white area of both curves.
But how do you turn that into a policy for exploring the space? Does anybody know?
Chris provides a sample of 20,000 values for the time constant and the additional information that the distribution of $tau$ is Weibull with parameters $\lambda=50$ and $k=1.5$.
Also, Chris keeps track of the leader board for best strategy. At the moment, the leader is Cameron Davidson-Pilon himself, with an average duration of 111.365 seconds.
So let's see if we can do any better. As in the previous problem, I'm going to start with the conjecture that the optimal strategy is to stay on the first side until the probability that you are on the right side drops below some threshold, then switch to the second side and stay until the probability drops below the same threshold.
I am not confident that this strategy is actually optimal, but let's see if it does a decent job.
Suppose we try on the first side for $t_1$ seconds and then try on the second side for $t_2$ seconds. At that point, we can use the survival function of the Weibull distribution to compute the log odds that the second side is correct:
$$ LO = \left( \frac{t_1}{\lambda} \right)^k - \left( \frac{t_2}{\lambda} \right)^k$$
Now suppose we want to stay on the current side until the log odds we're on the wrong side exceeds some threshold, x, and the time we've spent on the other side is other. 
The following function computes how long we should stay.
End of explanation
right = 0
left = 0
x = 1
lam = 50
k = 1.5
right = wait_time(x, left, lam, k)
right
Explanation: Starting with $t_1=0$ and $t_2=0$, when $x=1$, the time we wait turns out to be $\lambda$.
End of explanation
left = wait_time(x, right, lam, k)
left
Explanation: Now suppose we've waited 50 seconds on the left and we switch to the right. How long should we stay?
End of explanation
right = wait_time(x, left, lam, k)
right
Explanation: And now if we've spent left second on the left and switch to the right, how long should we stay.
End of explanation
def simulate(correct, tau, this, that, lam, k, x, trace):
    if len(trace) > 20:
        return 'Too many recursions'
    
    # print(correct, tau, this, that, lam, k, x, trace)
    
    # figure out the maximum time we should try before flipping
    wait = wait_time(x, that, lam, k)
    
    # if we're on the correct side, see if we succeed before time's up
    if correct and tau < wait:
        # if so, update and return the trace
        return trace + [tau]
    
    # if time expired, add the wait time to the trace  
    # and make a recursive call to continue the simulation
    return simulate(not correct, tau, that, wait, lam, k, x, trace + [wait])
Explanation: But all of that is based on the assumption that x=1. Let's see what happens to the average total time as we vary x.
Simulation
The following function simulates the strategy I just outlined.
It takes as parameters:
correct: A Boolean indicating if we're on the right side.
tau: The actual time to get hot water.
this: The longest time we've tried on the current side.
that: The longest time we've tried on the other side..
lam: The scale parameter for the distribution of time until success.
k: The shape parameter for the distribution of time until success.
x: The threshold for the posterior log odds.
trace: A list that indicates how much time we have spent, so far, trying and flipping.
It runs the simulation and returns a sequence of waiting times. The sum of this sequence is the total time it took to connect.
End of explanation
x = 0.5
simulate(True, 20, 0, 0, lam, k, x, [])
simulate(False, 20, 0, 0, lam, k, x, [])
simulate(True, 40, 0, 0, lam, k, x, [])
simulate(False, 40, 0, 0, lam, k, x, [])
simulate(True, 60, 0, 0, lam, k, x, [])
simulate(False, 60, 0, 0, lam, k, x, [])
simulate(True, 120, 0, 0, lam, k, x, [])
simulate(False, 120, 0, 0, lam, k, x, [])
Explanation: Here are a few test runs with x=0.5 and various values of correct and tau.
End of explanation
import os
if not os.path.exists('shower_problem_tau_samples.csv'):
    !wget https://gist.github.com/csaid/a57c4ebaa1c7b0671cdc9692638ea4c4/raw/ad1709938834d7bc88b62ff0763733502eb6a329/shower_problem_tau_samples.csv
import pandas as pd
test_set = pd.read_csv('shower_problem_tau_samples.csv')
test_set.head()
test_set.describe()
Explanation: Let's get the test data from Chris.
End of explanation
import numpy as np
def run_simulations(samples, lam, k, x):
    res = []
    for i, (direction, tau) in samples.iterrows():
        # print(direction, tau)
        trace = simulate(direction, tau, 0, 0, lam, k, x, [])
        res.append((len(trace), sum(trace)))
        
    return np.transpose(res)
Explanation: The following function uses the samples to run simulations and returns the average total time.
End of explanation
from scipy.stats import weibull_min
train_set = test_set.copy()
train_set['tau'] = weibull_min.rvs(k, scale=lam, size=20000)
train_set.describe()
Explanation: I'll train the algorithm with a different sample from the same distribution.
End of explanation
x = 1
lengths, totals = run_simulations(train_set, lam, k, x)
Explanation: Here's a test run with x=1.
End of explanation
lengths.mean()
totals.mean()
Explanation: On average it takes about two tries and 116 seconds.
End of explanation
xs = np.linspace(1, 3, 20)
xs
res = []
for x in xs:
    lengths, totals = run_simulations(train_set, lam, k, x)
    res.append((x, totals.mean()))
from statsmodels.nonparametric.smoothers_lowess import lowess
def make_lowess(series):
    Use LOWESS to compute a smooth line.
    series: pd.Series
    returns: pd.Series
    
    endog = series.values
    exog = series.index.values
    smooth = lowess(endog, exog)
    index, data = np.transpose(smooth)
    return pd.Series(data, index=index)
def plot_series_lowess(series, color):
    Plots a series of data points and a smooth line.
    series: pd.Series
    color: string or tuple
    
    series.plot(lw=0, marker='o', color=color, alpha=0.5)
    smooth = make_lowess(series)
    smooth.plot(label='_', color=color)
Explanation: Optimization
Now let's see how the average duration varies as we sweep through a range of values for the threshold probability, r:
End of explanation
import matplotlib.pyplot as plt
rs, ts = np.transpose(res)
series = pd.Series(ts, rs)
plot_series_lowess(series, 'C1')
plt.xlabel("Threshold log odds where you flip (x)")
plt.ylabel("Average total duration (seconds)");
Explanation: Here's what the results look like.
End of explanation
lengths, totals = run_simulations(test_set, lam, k, 2)
Explanation: It looks like the optimal value of x is about 2.
Testing
Now let's run the simulations with the optimal value of x and see what the average total time is for the test data.
End of explanation
totals.mean()
Explanation: The average total time is about 102 seconds, so it looks like this strategy is better than the current leader (assuming I haven't made a mistake in the simulation).
End of explanation
try:
    import empiricaldist
except ImportError:
    !pip install empiricaldist
from empiricaldist import Cdf
Cdf.from_seq(totals).plot(lw=2)
plt.xlabel('Total time to connect (seconds)')
plt.ylabel('CDF')
plt.title('Distribution of total time to connect');
np.percentile(totals, 90)
Explanation: Here's the distribution of total time, represented as a CDF.
End of explanation
from empiricaldist import Pmf
pmf1 = Pmf.from_seq(lengths-1)
pmf1.bar(alpha=0.7)
plt.xlabel('How many times you have to flip')
plt.ylabel('PMF')
plt.title('Distribution of number of flips');
Explanation: And here's the distribution of flips.
End of explanation
<END_TASK>