def uranium()

in src/beanmachine/tutorials/utils/radon.py [0:0]

• 86 lines of code
• 1 McCabe index (conditional complexity)

def uranium(summary_df: DataFrame, df: DataFrame) -> Figure:
    """
    Plot uranium linear regression.

    :param summary_df: The dataframe output from arviz.
    :param df: The original dataframe data.
    :returns plot: A bokeh Figure object.
    """
    alpha_hat_df = (
        summary_df[["mean", "sd"]]
        .loc[summary_df.index.astype(str).str.startswith("alpha_hat"), :]
        .reset_index(drop=True)
        .copy()
    )
    alpha_hat_df["log_Uppm"] = df["log_Uppm"].values
    alpha_hat_df["county"] = df["county"].values
    alpha_hat_df = alpha_hat_df.drop_duplicates().reset_index(drop=True)
    alpha_hat_df["lower"] = alpha_hat_df["mean"] - alpha_hat_df["sd"]
    alpha_hat_df["upper"] = alpha_hat_df["mean"] + alpha_hat_df["sd"]

    intercepts_source = ColumnDataSource(
        {
            "x": alpha_hat_df["log_Uppm"].values,
            "y": alpha_hat_df["mean"].values,
            "lower": alpha_hat_df["lower"].values,
            "upper": alpha_hat_df["upper"].values,
            "county": alpha_hat_df["county"].values,
        }
    )

    plot = figure(
        plot_width=800,
        plot_height=500,
        title="Partial-pooling with individual and group level predictors",
        x_axis_label="log(uranium)",
        y_axis_label="Intercept estimate (log(radon activity))",
        y_range=[0.5, 2.2],
        x_range=[-1, 0.6],
    )
    markers = plot.circle(
        x="x",
        y="y",
        source=intercepts_source,
        size=10,
        fill_color="steelblue",
        line_color="white",
        fill_alpha=0.7,
        line_alpha=0.7,
        hover_fill_color="orange",
        hover_line_color="black",
        hover_fill_alpha=1.0,
        legend_label="County",
    )
    tooltips = HoverTool(
        renderers=[markers],
        tooltips=[
            ("County", "@county"),
            ("Estimated α", "@y{0.000}"),
        ],
    )
    plot.add_tools(tooltips)
    whiskers = Whisker(
        base="x",
        upper="upper",
        lower="lower",
        source=intercepts_source,
        line_color="steelblue",
    )
    whiskers.upper_head.line_color = "steelblue"
    whiskers.lower_head.line_color = "steelblue"
    plot.add_layout(whiskers)

    x = np.array([-1, 1])
    a = summary_df.loc[
        summary_df.index.astype(str).str.startswith("mu_alpha"), "mean"
    ].values
    g = summary_df.loc[
        summary_df.index.astype(str).str.startswith("gamma"), "mean"
    ].values
    y = a + g * x
    plot.line(
        x=x,
        y=y,
        line_color="black",
        line_alpha=0.3,
        line_width=3,
        legend_label="Estimated linear regression",
        level="underlay",
    )

    plot.outline_line_color = "black"
    plot.grid.grid_line_alpha = 0.2
    plot.grid.grid_line_color = "grey"
    plot.grid.grid_line_width = 0.2
    plot.legend.location = "top_left"
    plot.output_backend = "svg"

    return plot