# Polya's Urn
#
# From Gregory F. Lawler, "Introduction to Stochastic Processes",
# second edition, ISBN 978-1-58488-651-8, Chapman and Hall/CRC, 2006.
# See Example 4 on p.109.
#
# Implementation for the R software system.
# 
# AUTHOR: M.V.Wickerhauser
# DATE: 2018-02-24
#
# Inputs:                               (default)
#    nb      Number of balls to add         (998)
#    red     Initial number of red balls      (1)
#    green   Initial number of green balls    (1)
#
# Output:
#    X      final number of red balls
#
# Notes:
#    The final number of balls is red+green+nb.
#    The final number of green balls is red+green+nb-X.
#    The fraction of red balls at the end is
#           X/(red+green+nb)
# 
PolyaUrn <- function(nb=998, red=1, green=1) {
  X <- red               # initial number of red balls
  start <- red+green     # initial number of balls 
  end <- start+nb        # final number of balls
  for(i in start:end) {  # sample and add nb more balls
    pred <- X/i          # probability that the ball is red
    pgreen <- 1-pred     # probability that the ball is green
    newball <- sample(c(1,0),1,prob=c(pred,pgreen))
    X <- X + newball     # increase X if the new ball is red
  }
  return(X)              # publish the final number of red balls
}


##################################
# Numerical experiments:
#
###
# Run 2000 trials of Polya's Urn up to 1000 final balls:
trials <- 2000		# number of trials
RedFraction <- 1:trials  # initialize a vector to hold the trial results
for(k in 1:trials)
  RedFraction[k] <- PolyaUrn(998,red=1,green=1)/1000
# Plot the fraction of red balls after each trial as a histogram with
# bins [0,0.05), [0.05,0.10), [0.10,0.15),...,[0.95,1)
hist(RedFraction,breaks=seq(0,1,by=0.05),right=F)

###
# Run 100 simulations of Polya's urn up to 1000 balls, then continue
# to 2000 balls, and scatterplot the red fractions (M998,M1998):
sims <- 100                        # perform 100 sumulations
M998 <- 1:sims; M1998 <- 1:sims;   # vectors for the red fractions
for(s in 1:sims) {
  red <- PolyaUrn(998,red=1,green=1) # 1 red and 1 green to start
  M998[s] <- red/1000;               # fraction of red balls at 1000 balls
  green <- 1000 - red;               # number of green balls at 1000 balls
  M1998[s] <- PolyaUrn(1000,red,green)/2000  # fraction of red at 2000 balls
}
plot(M998,M1998)   # scatterplot showing the correlation of M998 and M1998