two_spouse.py

import random

p_friend_infected = 0.1
q_friend_transmission = 0.1
s_spouse_transmission = 0.3

DRAWS = 10000000

same_friend = {0: 0,
               1: 0,
               2: 0}
for i in range(DRAWS):
  friend_infected = random.random() < p_friend_infected
  if not friend_infected:
    same_friend[0] += 1
  else:
    A_infected = random.random() < q_friend_transmission
    B_infected = random.random() < q_friend_transmission

    A_conditionally_transmit = random.random() < s_spouse_transmission 
    B_conditionally_transmit = random.random() < s_spouse_transmission 

    A_infected = A_infected or (B_infected and B_conditionally_transmit)
    B_infected = B_infected or (A_infected and A_conditionally_transmit)

    if A_infected and B_infected:
      same_friend[2] += 1
    elif A_infected or B_infected:
      same_friend[1] += 1
    else:
      same_friend[0] += 1

different_friends = {0: 0,
                     1: 0,
                     2: 0}
for i in range(DRAWS):
  A_friend_infected = random.random() < p_friend_infected
  if not A_friend_infected:
    A_infected = False
  else:
    A_infected = random.random() < q_friend_transmission

  B_friend_infected = random.random() < p_friend_infected
  if not B_friend_infected:
    B_infected = False
  else:
    B_infected = random.random() < q_friend_transmission    

  A_conditionally_transmit = random.random() < s_spouse_transmission 
  B_conditionally_transmit = random.random() < s_spouse_transmission 

  A_infected = A_infected or (B_infected and B_conditionally_transmit)
  B_infected = B_infected or (A_infected and A_conditionally_transmit)

  if A_infected and B_infected:
    different_friends[2] += 1
  elif A_infected or B_infected:
    different_friends[1] += 1
  else:
    different_friends[0] += 1

Two-spouse correlated risks problem