Equal size buckets on discrete color scale regardless of range

Hi there,

I was able to create the discrete colorbar shown here

but Iโ€™m wondering if it is possible to separate the values that are equal to zero into a separate, equal length bucket in the scale. Essentially, Iโ€™d have the 6 color ranges on the right-hand side, plus a 7th for values of zero.

Here is my existing code.

DISCRETE = 6

def gen_colorscale(obs, color="viridis"):
    color = px.colors.sample_colorscale(color, obs)
    p1 = tuple(zip(np.linspace(0, 1, obs+1)[:-1], color))
    p2 = tuple(zip(np.linspace(0, 1, obs+1)[1:], color))
    cs = []
    for a, b in zip(p1, p2):
        cs.append(a)
        cs.append(b)
    return cs

cs = gen_colorscale(DISCRETE)

# color range
cr = [0, 20000]
# tick vals
v = np.linspace(*cr, DISCRETE)
vt = (
    pd.DataFrame(v, columns=["v"])
    .apply(lambda v: (v / 10 ** 3).round(0))
    .apply(lambda v: v.round(0).astype(int).astype(str) + "k to " + v.shift(-1).fillna(0).round(0).astype(int).astype(str) + "k")
    .values
)
vt[0] = v[0].round(0).astype(int).astype(str) + " to " + (v[1] / 10 ** 3).round(0).astype(int).astype(str) + "k"
vt[-1] = ">" + (v[-1] / 10 ** 3).round(0).astype(int).astype(str) + "k"

# =============================================================================
# Create interactive chart with all data
# =============================================================================
fig = px.choropleth_mapbox(
    df,
    geojson=counties,
    locations="fips",
    color="migration",
    range_color=[cr[0], cr[1] + cr[1]/(DISCRETE-1)],
    color_continuous_scale=cs,
    labels={"migration": "Migration (k)"},
    center={"lat": 37.0902, "lon": -95.7129},
    zoom=4.2,
    opacity=1.0,
    mapbox_style="white-bg",
)
fig.update_layout(
    title=dict(
        text="The Great California Migration",
        font=dict(family="Rockwell", size=40, color="rgba(255,255,255,1)"),
        x=0.5,
        y=0.97,
    ),
    coloraxis_colorbar=dict(
        tickvals=np.linspace(cr[0]+cr[1]/(DISCRETE-1)/2,cr[1] + cr[1]/(DISCRETE-1)/2,DISCRETE),
        ticktext=vt,
        len=0.8,
        thickness=50,
        xanchor="right",
        x=1.0,
        bgcolor="rgba(22,33,49,1)",
        title="Migration Inflow<sup>1</sup>",
        titlefont=dict(
            family="Rockwell",
            color="rgba(255,255,255,1)"
        ),
        tickfont=dict(color="rgba(255,255,255,1)"),
    ),
    margin=dict(l=0, r=0, b=50, t=75, pad=4),
    paper_bgcolor="rgba(8,18,23,1)",
    plot_bgcolor="rgba(8,18,23,1)",
    showlegend=True,
    annotations=[
        dict(
            x=0.005,
            y=0.005,
            xref="paper",
            yref="paper",
            align="left",
            text="1. Migration flow estimated by changes to IRS filing (returns and exceptions) location from 2014 through 2019, only includes migration out of California<br>@econgraphix | Source: 2014-2019 IRS Migration Data",
            showarrow=False,
            font=dict(family="Rockwell", color="rgba(255,255,255,1)"),
            bgcolor="rgba(8,18,23,1)",
        )
    ],
)