add script for adding columns to the dr1 catalog

# Date: 2026-02-25 13:50:08

"""

Update the DR1 catalog and database with respect to new columns

"""

Update the DR1 catalog and database with respect to new columns.

The following script was run on 2026 03 16 in order to udpate the truth catalog

FITS file as well as master.db.

The columns added are `is_agn_focc` and `*_tau` and `*_sf_inf` (for the ugrizy

bands).

After running the script, the dr1_new_new catalog.fits was replaced by the

catalog generated by this script.

Then, the master.db was updated by dropping the original truth catalog and

joining truth.db generated by this script.

20260316

av

"""

import glob

import os

import astropy.coordinates as c

from lsst_inaf_agile.catalog_combined import CatalogCombined

dirname = "data/catalog/dr1_new_new"

columns = ["ID", "Z", "M", "is_agn", "magabs_lsst-i_point", "MBH", "is_optical_type2"]

for band in "ugrizy":

    columns.append(f"lsst-{band}_point")

#columns = ["ID", "Z", "M", "is_agn", "magabs_lsst-i_point", "MBH", "is_optical_type2"]

#for band in "ugrizy":

#    columns.append(f"lsst-{band}_point")

columns = None

catalog = CatalogCombined(dirname, columns=columns)

tau_sf_infs: list[list[float]] = []

tau_sf_infs: dict[tuple[int, str], list[float]] = {}

mjd = util.get_mjd_vec()

for b in "ugrizy":

    band = f"lsst-{b}"

    i = catalog[catalog["is_agn"]]["ID"][0]

bands = [f"lsst-{b}" for b in "ugrizy"]

#uids = catalog[catalog["is_agn"]]["ID"]

uids = list(

    map(

        lambda f: int(os.path.basename(f)),

        glob.glob("data/catalog/dr1_new_new/lightcurves/*")

    )

)

print(uids)

def get_one(uid, band, return_tau_sf_inf):

    lckw = {

            "mjd0": 0.0,

            "mjd": mjd,

            "band": band,

        "flux": catalog[f"{band}_point"][i],

        "z": catalog["Z"][i],

        "mag_i": catalog["magabs_lsst-i_point"][i],

        "logMbh": catalog["MBH"][i],

        "type2": catalog["is_optical_type2"][i],

            "flux": catalog[f"{band}_point"][uid],

            "z": catalog["Z"][uid],

            "mag_i": catalog["magabs_lsst-i_point"][uid],

            "logMbh": catalog["MBH"][uid],

            "type2": catalog["is_optical_type2"][uid],

            "T": mjd.max() + 1,

            "deltatc": 1,

        "seed": 9999,

        "return_tau_sf_inf": True,

            "seed": uid + 1,

            "return_tau_sf_inf": return_tau_sf_inf,

            "divide_sf_inf_by_z1": True,

            }

    lc = lightcurve.get_lightcurve_agn(**lckw)

    print(lc)

    return lightcurve.get_lightcurve_agn(**lckw)

    if b == "u":

        tau_sf_infs.append([])

    tau_sf_infs[-1].append(lc)

def get_tau_sf_inf_all(band="lsst-r"):

    import astropy.units as u

    from lsst_inaf_agile.my_lsst import filter_lam_fwhm

    from AGN_lightcurves_simulations.AGN_sim.DRW_sim_new import get_tau_sf_inf

    lambda_of = filter_lam_fwhm[band.replace("lsst-", "")][0] * u.um.to(u.angstrom)

    tau, sf_inf = get_tau_sf_inf(

            lambda_of / (1 + catalog["Z"]),

            catalog["magabs_lsst-i_point"],

            catalog["MBH"],

            catalog["Z"],

            catalog["is_optical_type2"],

            "observed",

            True

    )

    return tau, sf_inf

print(np.array(tau_sf_infs).shape)

quit()

# add the columns

print(c["tau"])

print(c["sf_inf"])

quit()

for band in bands:

    for nproc, uid in enumerate(uids[:10]):

table = "Truth"

filename = "test.db"

        # 1: test 'get_one' against disk light curves

        lc = get_one(uid, band, return_tau_sf_inf=False)

        filename = f"data/catalog/dr1_new_new/lightcurves/{uid}/{band}.npy"

        lcref = np.load(filename)

        assert np.allclose(lc, lcref)

        # 2: get vectorized tau, sf_inf

        tau1, sf_inf1 = get_one(uid, band, return_tau_sf_inf=True)

        select = catalog["ID"] == uid

        tau2, sf_inf2 = map(lambda a: a[select][0], get_tau_sf_inf_all(band))

        assert np.allclose([tau1, sf_inf1], [tau2, sf_inf2])

from astropy.table import Table

catalog.catalog_combined = Table(catalog.catalog_combined)

from lsst_inaf_agile.mbh import get_occupation_fraction

focc = get_occupation_fraction(np.where(catalog["M"] > 0.0, catalog["M"], 7.0))

U = np.random.uniform(0.0, 1.0, catalog["M"].size)

has_bh = focc > U

is_agn_focc = has_bh & catalog["is_agn"]

catalog["is_agn_focc"] = is_agn_focc

print(catalog["is_agn"].sum())

print(is_agn_focc.sum())

for band in bands:

    tau, sf_inf = get_tau_sf_inf_all(band)

    tau = np.where(catalog["is_agn"], tau, np.nan)

    sf_inf = np.where(catalog["is_agn"], sf_inf, np.nan)

    catalog[f"{band}_tau"] = tau

    catalog[f"{band}_sf_inf"] = sf_inf

filename = "truth.fits"

if os.path.exists(filename):

    os.remove(filename)

c.write_database(table, filename)

catalog.write(filename)

print("wrote", filename)

print(c)

table = "Truth"

filename = "truth.db"

if os.path.exists(filename):

    os.remove(filename)

catalog.write_database(table, filename)

print("wrote", filename)