def floor_plot()

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

• 96 lines of code
• 3 McCabe index (conditional complexity)

def floor_plot(df):
    # Create the plot.
    radon_floor_plot = figure(
        plot_width=800,
        plot_height=500,
        title="log(radon) measurement vs floor",
        y_range=[-6, 6],
        x_range=[-0.5, 1.5],
        x_axis_label="Floor",
        y_axis_label="log(radon)",
    )

    # Prepare data for the plot.
    basement_floor_data = df[df["floor"] == 0]["log_activity"].values
    basement_floor_kde = sm.nonparametric.KDEUnivariate(basement_floor_data)
    basement_floor_kde.fit()
    ground_floor_data = df[df["floor"] == 1]["log_activity"].values
    ground_floor_kde = sm.nonparametric.KDEUnivariate(ground_floor_data)
    ground_floor_kde.fit()
    radon_floor_source = ColumnDataSource(
        {
            "x": (df["floor"].values + np.random.normal(scale=0.02, size=df.shape[0])),
            "y": df["log_activity"].values,
            "county": df["county"].values,
            "color": df["floor"].apply(
                lambda floor: "orange" if floor == 1 else "steelblue"
            ),
        }
    )

    # Bind data to the plot.
    radon_floor_glyph = radon_floor_plot.circle(
        x="x",
        y="y",
        source=radon_floor_source,
        size=5,
        fill_color="color",
        line_color="white",
        alpha=0.7,
    )
    radon_floor_tooltips = HoverTool(
        renderers=[radon_floor_glyph],
        tooltips=[
            ("County", "@county"),
            ("log(radon)", "@y{0.000}"),
        ],
    )
    radon_floor_plot.add_tools(radon_floor_tooltips)
    x = [-0.25, 0.25]
    y = [basement_floor_kde.support[np.argmax(basement_floor_kde.density)]] * 2
    radon_floor_plot.line(
        y=y,
        x=x,
        line_color="steelblue",
        line_dash="dashed",
        line_width=4.0,
        alpha=0.5,
    )
    x = 0.25 * (basement_floor_kde.density / basement_floor_kde.density.max())
    radon_floor_plot.line(
        y=basement_floor_kde.support,
        x=x,
        line_color="steelblue",
        alpha=0.7,
    )
    radon_floor_plot.line(
        y=basement_floor_kde.support,
        x=-x,
        line_color="steelblue",
        alpha=0.7,
    )
    x = 0.25 * (ground_floor_kde.density / ground_floor_kde.density.max())
    radon_floor_plot.line(
        y=ground_floor_kde.support,
        x=1 + x,
        line_color="orange",
        alpha=0.7,
    )
    radon_floor_plot.line(
        y=ground_floor_kde.support,
        x=1 - x,
        line_color="orange",
        alpha=0.7,
    )
    radon_floor_plot.line(
        y=[ground_floor_kde.support[np.argmax(ground_floor_kde.density)]] * 2,
        x=[0.75, 1.25],
        line_color="orange",
        line_dash="dashed",
        line_width=4.0,
        alpha=0.5,
    )

    # Style the plot.
    radon_floor_plot.xaxis.ticker = FixedTicker(ticks=[0, 1])
    radon_floor_plot.xaxis.major_label_overrides = {
        0: "Basement",
        1: "First",
    }
    radon_floor_plot.grid.grid_line_color = None
    radon_floor_plot.outline_line_color = "black"
    radon_floor_plot.output_backend = "svg"

    return radon_floor_plot