Python script for computing the probabilities of winning the New Mexico vaccination draw by "region"

Vax2TheMaxNM.py

"""
Simple script to compute the probability of winning
one or any of the New Mexico vaccination draws.

Uses the Python requests library to submit API
endpoint requests to the Wolfram|Alpha API, and
the Rich library for rendering to a table.

===============

The only variable that meaningfully needs to be
edited (besides WOLFRAM_APP_ID) is VACCINATION_RATE.
This was guessed by looking at Google time series
plots of vaccinations statewide and projecting to mid-July.

===============

REFERENCES
----------
NM Regions: https://www.nmhealth.org/about/phd/region/
NM Contests: https://vax2themaxnm.org/sweepstakes/sweepstakes-rules/
Wolfram|Alpha API: https://products.wolframalpha.com/api/documentation/

"""

import requests
import xml.etree.ElementTree as ET
import rich
from rich.console import Console
from rich.table import Table

# If True, prints extra information on API request status
DEBUG_MODE = False

# Wolfram|Alpha API key, which will need to be generated from:
#   https://products.wolframalpha.com/api/documentation/
WOLFRAM_APP_ID = "XXXXXX-XXXXXXXXXX"

# A static vaccination rate
# TODO: vaccination rates will tend upwards over time. This should reflect that.
VACCINATION_RATE = 0.65

# Number of unique drawings:
#   https://vax2themaxnm.org/sweepstakes/sweepstakes-rules/
N_SWEEPSTAKES = 5

# Regions available from:
#   https://www.nmhealth.org/about/phd/region/
REGIONS = {
    "northeast": [
        "Rio Arriba",
        "Los Alamos",
        "Taos",
        "Colfax",
        "Mora",
        "Santa Fe",
        "San Miguel",
        "Guadalupe",
        "Harding",
        "Union",
    ],
    "northwest": [
        "San Juan",
        "McKinley",
        "Sandoval",
        "Cibola",
        "Bernalillo",
        "Valencia",
        "Torrance",
    ],
    "southwest": [
        "Catron",
        "Socorro",
        "Grant",
        "Sierra",
        "Luna",
        "Hidalgo",
        "Luna",
        "Dona Ana",
        "Otero",
    ],
    "southeast": [
        "Lincoln",
        "De Baca",
        "Chaves",
        "Eddy",
        "Lea",
        "Quay",
        "Curry",
        "Roosevelt",
    ],
}


def log_debug(msg: str):
    """Logs a message to console only if DEBUG_MODE is True."""
    if DEBUG_MODE:
        rich.print(msg)


def wolfram_api_request(
    query: str,
    base_URL: str = "http://api.wolframalpha.com/v2/query",
    app_id: str = WOLFRAM_APP_ID,
    header: dict = {"user-agent": "Python"},
):
    """
    Submits a request to the Wolfram|Alpha API, and, if successful,
    returns the represented XML object.
    If not, fails with ValueError.
    """
    query = query.replace(" ", "%20")  # Replace spaces with web-compliant chars
    query = f"{base_URL}?input={query}&appid={app_id}"  # Encode query to the API specification

    log_debug(f"Submitting query: {query}")

    response = requests.get(query, headers=header)  # Submit request

    log_debug(f"{response}: {response.text}")

    # A successful response has the HTTP status code 200
    if response.status_code != 200:
        raise ValueError(
            f"Request failed with HTTP status code: {response.status_code}"
        )

    xml_root = ET.fromstring(response.text)

    if xml_root.attrib["success"] != "true":
        raise ValueError("Response claims to have failed")

    return xml_root


def get_region_population(counties: list) -> int:
    """
    Returns the total region population, and a dict with
    the schema: { "county": county_population, ... }
    """
    query = "population of %s county, New Mexico"
    population = {}

    for county in counties:
        query_result = wolfram_api_request(query % county)

        for elem in query_result.findall("pod"):
            if elem.get("title").lower() == "result":
                # Find the XML element with the population in plaintext
                # i.e., will be of the form: '38921 people (2019)'
                pop_txt = elem.find("subpod").find("plaintext").text

                log_debug(f"Found population text: {pop_txt}")

                # Convert that to an int and assign to the county dict entry
                population[county] = int(pop_txt[: pop_txt.find("people")])

    return sum([population[cnty] for cnty in counties]), population


def compute_single_probability(population: int) -> float:
    """Computes the probability of winning a single contest."""
    return 1.0 - ((population - 1) / population)


def compute_total_probability(
    population: int, num_entries: int = N_SWEEPSTAKES
) -> float:
    """Computes the probability of winning any one of `num_entries` contests."""
    return 1.0 - (((population - 1) / population) ** num_entries)


def float_to_percent(number: float, precision: int = 5) -> str:
    """Converts a float to a percent string. I.e., 0.95 -> "95%"."""
    return f"{round(100. * number, precision)}%"


if __name__ == "__main__":

    console = Console()
    table = Table(show_header=True, header_style="bold magenta")
    table.add_column("Region name", style="dim", width=12)
    table.add_column("Prob. of winning any drawing")
    table.add_column("Prob. of winning one drawing")
    table.add_column("Region population")
    table.add_column("Region counties")

    for region in REGIONS:
        counties = REGIONS[region]
        total_pop, county_pop = get_region_population(counties)
        adjusted_pop = total_pop * VACCINATION_RATE
        total_probability = float_to_percent(compute_total_probability(adjusted_pop))
        single_probability = float_to_percent(compute_single_probability(adjusted_pop))
        table.add_row(
            region.capitalize(),
            f"[bold]{total_probability}[/bold]",
            f"[bold]{single_probability}[/bold]",
            "{:,}".format(total_pop),
            ", ".join(counties),
        )

    console.print(
        f"[bold][underline]Note: results assuming an average vaccination rate of {VACCINATION_RATE}[/bold][/underline]"
    )
    console.print(table)