add linestyles based on mass limits and option for redshifts

parser = ArgumentParser()

parser.add_argument("savefig")

parser.add_argument("--plot_zou2024", action="store_true")

parser.add_argument("--redshifts", nargs="+", type=float, default=[0.0, 1.0, 2.0, 3.0, 4.0, 5.0])

parser.add_argument("--zfmt", default="%d")

parser.add_argument("--vmax", type=float, default=5.5)

args = parser.parse_args()

###############################################################################

lw3 = 1.5

alpha = 0.80

cmap = mpl.colormaps["gist_heat"]

norm = mpl.colors.Normalize(vmin=0.0, vmax=5.5)

norm = mpl.colors.Normalize(vmin=0.0, vmax=args.vmax)

cmap2 = mpl.colormaps["rainbow"]

norm2 = mpl.colors.Normalize(vmin=0.0, vmax=4.0)

###############################################################################

    # fig, axes = plt.subplots(1, 2, figsize=(1.3 * 6.4, 1.0 * 4.8), sharex=True, sharey=True, dpi=300)

    fig, axes = plt.subplots(1, 2, figsize=(8.3, 5.85 / 1.5), sharex=True, sharey=True, dpi=300)

    def get_mass_limit(z, which=99):

        from astropy.table import Table

        table = Table.read("notebooks/data/agile_dr1_redshift_mass_limit.fits")

        idx = np.argmin(np.abs(table["zcen"] - z))

        return table[idx][f"M{which}"]

    # Left panel: Plots the redshift evolution of the Con Eq

    for z in [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 5.5]:

    # for z in [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]:

    for z in args.redshifts:

        # log_mbh = mbh.get_log_mbh_continuity(log_mstar, z)

        # log_mbh = mbh.get_log_mbh_continuity_new(log_mstar, z)

        ls = "solid" if z < 4 else "dashed"

        if z < 5.5:

        # plot above and below the mass limit

        mlimit = get_mass_limit(z)

        select_below = log_mstar < mlimit

        select_above = log_mstar > mlimit

        print(f"{z=} {mlimit=}")

        log_mbh = mbh.get_log_mbh_continuity_new2(log_mstar, np.full_like(log_mstar, z))

            axes[0].plot(log_mstar, log_mbh, color=cmap(norm(z)), lw=lw1, label=f"$z = {z:.0f}$", ls=ls)

        else:

        axes[0].plot(

            log_mstar[select_below],

            log_mbh[select_below],

            color=cmap(norm(z)),

            lw=lw1,

            label=f"$z = {args.zfmt}$" % z,

            ls="solid",

        )

        axes[0].plot(

            log_mstar[select_above],

            log_mbh[select_above],

            color=cmap(norm(z)),

            lw=lw1,

            label=None,

            ls="dashed",

        )

    ## label mass limits?

    label_mlim = False

    if label_mlim:

        axes[0].plot([], [], lw=2, ls="solid", color="k", label="M<99%")

        axes[0].plot([], [], lw=2, ls="dashed", color="k", label="M>99%")

    # Plot IC

    plot_ic = True

    if plot_ic:

        log_mbh = mbh.get_log_mbh_reines_volonteri2015(log_mstar)

        axes[0].plot(log_mstar, log_mbh, color="C2", lw=lw1, label="IC (R&V15)", ls="dotted")

    fig.tight_layout()

    fig.savefig(args.savefig)

    print("wrote", args.savefig)

if __name__ == "__main__":