Plot quiescent UVJ diagram

log10(age/Myr): 0.0 U-V: -0.63 V-J: -1.48
log10(age/Myr): 0.1 U-V: -0.61 V-J: -1.47
log10(age/Myr): 0.2 U-V: -0.60 V-J: -1.47
log10(age/Myr): 0.3 U-V: -0.58 V-J: -1.46
log10(age/Myr): 0.4 U-V: -0.52 V-J: -1.41
log10(age/Myr): 0.5 U-V: -0.41 V-J: -1.27
log10(age/Myr): 0.6 U-V: -0.42 V-J: -1.07
log10(age/Myr): 0.7 U-V: -0.34 V-J: -0.34
log10(age/Myr): 0.8 U-V: -0.27 V-J: 0.16
log10(age/Myr): 0.9 U-V: -0.14 V-J: 0.75
log10(age/Myr): 1.0 U-V: -0.06 V-J: 0.81
log10(age/Myr): 1.1 U-V: 0.02 V-J: 0.73
log10(age/Myr): 1.2 U-V: 0.04 V-J: 0.34
log10(age/Myr): 1.3 U-V: 0.13 V-J: 0.12
log10(age/Myr): 1.4 U-V: 0.16 V-J: -0.05
log10(age/Myr): 1.5 U-V: 0.21 V-J: -0.18
log10(age/Myr): 1.6 U-V: 0.30 V-J: -0.13
log10(age/Myr): 1.7 U-V: 0.35 V-J: -0.21
log10(age/Myr): 1.8 U-V: 0.46 V-J: 0.16
log10(age/Myr): 1.9 U-V: 0.48 V-J: -0.13
log10(age/Myr): 2.0 U-V: 0.56 V-J: 0.29
log10(age/Myr): 2.1 U-V: 0.67 V-J: 0.61
log10(age/Myr): 2.2 U-V: 0.64 V-J: 0.52
log10(age/Myr): 2.3 U-V: 0.78 V-J: 0.47
log10(age/Myr): 2.4 U-V: 0.89 V-J: 0.65
log10(age/Myr): 2.5 U-V: 0.87 V-J: 0.40
log10(age/Myr): 2.6 U-V: 1.13 V-J: 0.86
log10(age/Myr): 2.7 U-V: 1.19 V-J: 0.69
log10(age/Myr): 2.8 U-V: 1.42 V-J: 0.92
log10(age/Myr): 2.9 U-V: 1.54 V-J: 1.06
log10(age/Myr): 3.0 U-V: 1.49 V-J: 0.82
log10(age/Myr): 3.1 U-V: 1.37 V-J: 0.32
log10(age/Myr): 3.2 U-V: 1.59 V-J: 0.80
log10(age/Myr): 3.3 U-V: 1.67 V-J: 0.95
log10(age/Myr): 3.4 U-V: 1.78 V-J: 0.98
log10(age/Myr): 3.5 U-V: 1.86 V-J: 1.02
log10(age/Myr): 3.6 U-V: 1.94 V-J: 1.04
log10(age/Myr): 3.7 U-V: 2.02 V-J: 1.08
log10(age/Myr): 3.8 U-V: 2.04 V-J: 1.05
log10(age/Myr): 3.9 U-V: 1.91 V-J: 0.96
log10(age/Myr): 4.0 U-V: 2.17 V-J: 1.12
log10(age/Myr): 4.1 U-V: 2.05 V-J: 0.98
log10(age/Myr): 4.2 U-V: 2.30 V-J: 1.20
log10(age/Myr): 4.3 U-V: 2.41 V-J: 1.26
log10(age/Myr): 4.4 U-V: 2.46 V-J: 1.25
log10(age/Myr): 4.5 U-V: 2.51 V-J: 1.29
log10(age/Myr): 4.6 U-V: 2.34 V-J: 1.27
log10(age/Myr): 4.7 U-V: 2.16 V-J: 1.25
log10(age/Myr): 4.8 U-V: 2.56 V-J: 1.46
log10(age/Myr): 4.9 U-V: 2.82 V-J: 1.58
log10(age/Myr): 5.0 U-V: 2.63 V-J: 1.51

import cmasher as cmr
import matplotlib.pyplot as plt
import numpy as np
from astropy.table import Table
from synthesizer.filters import UVJ
from synthesizer.grid import Grid


def simple_UVJ(grid, target_metallicity=0.01):
    """Calculate UVJ colours as a function of age for single metallicity"""

    iZ = grid.get_nearest_index(target_metallicity, grid.metallicity)

    fc = UVJ(new_lam=grid.lam)
    # fc.plot_transmission_curves()

    for ia, log10age in enumerate(grid.log10age):
        sed = grid.get_spectra(
            (ia, iZ)
        )  # creates an SED object from a given grid point

        # --- now calculate the observed frame spectra

        sed.get_fnu0()  # generate dummy observed frame spectra.

        # --- measure broadband fluxes
        sed.get_photo_fluxes(fc)

        print(
            (
                f"log10(age/Myr): {log10age - 6:.1f} "
                f'U-V: {sed.measure_colour("U", "V"):.2f} '
                f'V-J: {sed.measure_colour("V", "J"):.2f}'
            )
        )


def UVJ_metallicity(grid):
    """
    Calculate UVJ as a function of metallicity and save as a .ecsv file
    and make a figure
    """

    fc = UVJ(new_lam=grid.lam)

    table = Table()
    table.meta["metallicities"] = list(grid.metallicity)
    table["log10ages"] = grid.log10age

    for iZ, Z in enumerate(grid.metallicity):
        for f in "UVJ":
            table[f"{Z}_{f}"] = np.zeros(len(grid.log10age))

        for ia, log10age in enumerate(grid.log10age):
            sed = grid.get_spectra(
                (ia, iZ)
            )  # creates an SED object from a given grid point

            # --- now calculate the observed frame spectra

            sed.get_fnu0()  # generate dummy observed frame spectra.

            # --- measure broadband fluxes
            sed.get_photo_fluxes(fc)

            for f in "UVJ":
                table[f"{Z}_{f}"][ia] = sed.photo_fluxes[f]

    # --- make plot

    fig, axes = plt.subplots(
        2, 1, figsize=(3.5, 4.5), sharex=True, sharey=True
    )
    plt.subplots_adjust(
        left=0.15, top=0.975, bottom=0.1, right=0.95, wspace=0.0, hspace=0.0
    )

    colors = cmr.take_cmap_colors("cmr.bubblegum", len(grid.metallicity))

    for Z, c in zip(grid.metallicity, colors):
        x = table["log10ages"] - 6.0

        for i, (f1, f2) in enumerate(["UV", "VJ"]):
            y = 2.5 * np.log10(table[f"{Z}_{f2}"] / table[f"{Z}_{f1}"])
            axes[i].plot(x, y, color=c, lw=1, label=f"Z={Z}")

    for i, (f1, f2) in enumerate(["UV", "VJ"]):
        axes[i].set_ylabel(rf"$\rm {f1}-{f2}$")

    axes[0].legend(fontsize=6, labelspacing=0.0)
    axes[1].set_xlabel(r"$\rm \log_{10}(age/Myr)$")

    return table, fig, axes


if __name__ == "__main__":
    grid_dir = "../../tests/test_grid"
    grid_name = "test_grid"

    grid = Grid(grid_name, grid_dir=grid_dir)

    simple_UVJ(grid)

    table, fig, axes = UVJ_metallicity(grid)
    plt.show()

    # table.write(f'data/{grid_name}_UVJ.ecsv', overwrite=True)