run_csst_msc_mbi_photometry.py

"""_summary_
"""
import os
import yaml
from glob import glob
import mpi4py.MPI as MPI
from measurement_pipeline.L1_pipeline.csst_msc_mbi_photometry.csst_photometry import core_msc_l1_mbi_phot
from typing import Optional
from measurement_pipeline.config import parse_args

def run_csst_msc_mbi_photometry(image_path,
                                output_dir,
                                flag_weight_dir: str = None,
                                weight_path: str = None,
                                flag_path: str = None,
                                psf_ccds_path: Optional[str] = None,
                                plot_flag: bool = False,
                                getpsf_flag: bool = False,):

    input_dir = os.path.dirname(image_path)
    img_filename = os.path.basename(image_path)

    if flag_weight_dir is None:
        flag_weight_dir = input_dir

    if weight_path is None:
        weight_file_name = img_filename.replace("_injected", "")
        weight_path = os.path.join(
            flag_weight_dir, weight_file_name.replace("img", "wht"))
    if flag_path is None:
        flag_file_name = img_filename.replace("_injected", "")
        flag_path = os.path.join(
            flag_weight_dir, flag_file_name.replace("img", "flg"))

    psf_local_path = os.path.join(
        output_dir, img_filename.replace("img", "psf"))
    flux_path = os.path.join(
        output_dir, img_filename.replace("img", "flux"))
    cat_path = os.path.join(
        output_dir, img_filename.replace("img", "cat"))
    seg_path = os.path.join(
        output_dir, img_filename.replace("img", "seg"))
    sky_path = os.path.join(
        output_dir, img_filename.replace("img", "sky"))

    core_msc_l1_mbi_phot(
        image_path=image_path,
        weight_path=weight_path,
        flag_path=flag_path,
        psf_ccds_path=None,
        psf_local_path=psf_local_path,
        flux_path=flux_path,
        cat_path=cat_path,
        seg_path=seg_path,
        sky_path=sky_path,
        plot_flag=plot_flag,
        getpsf_flag=getpsf_flag)


def genearte_path_list_for_one_pointing(input_dir,
                                        pointing_label,
                                        chip_label_list=None):
    """_summary_

    Args:
        input_dir (_type_): _description_
        pointing_label (_type_): _description_
        chip_label_list (_type_, optional): _description_. Defaults to None.

    Returns:
        _type_: _description_
    """
    pointing_dir = os.path.join(input_dir, pointing_label)
    if chip_label_list is None:
        image_path_list = glob(
            pointing_dir + '/CSST_MSC_MS_SCIE_*_' + '*_img_*.fits')
    else:
        image_path_list = []
        for chip_label in chip_label_list:
            image_path = glob(pointing_dir + '/CSST_MSC_MS_SCIE_*_' +
                              chip_label + '_img_*.fits')[0]
            image_path_list.append(image_path)

    return image_path_list


def run_pointing_list(input_dir,
                      pointing_label_list,
                      output_dir,
                      chip_label_list=None,
                      flag_weight_dir=None):
    image_path_list = []
    output_path_list = []
    fw_dir_list = []
    try:
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
    except OSError:
        pass

    for pointing_label in pointing_label_list:
        output_pointing_dir = os.path.join(output_dir, pointing_label)
        fw_dir = os.path.join(flag_weight_dir, pointing_label)
        try:
            if not os.path.exists(output_pointing_dir):
                os.makedirs(output_pointing_dir)
        except OSError:
            pass
        temp_img_path_list = genearte_path_list_for_one_pointing(input_dir=input_dir,
                                                                 pointing_label=pointing_label,
                                                                 chip_label_list=chip_label_list)
        image_path_list = image_path_list + temp_img_path_list
        output_path_list = output_path_list + \
            [output_pointing_dir] * len(temp_img_path_list)
        fw_dir_list = fw_dir_list + [fw_dir] * len(temp_img_path_list)

    comm = MPI.COMM_WORLD
    ind_thread = comm.Get_rank()
    num_thread = comm.Get_size()

    for i in range(len(image_path_list)):
        if i % num_thread != ind_thread:
            continue
        image_path = image_path_list[i]
        output_path = output_path_list[i]
        fw_dir = fw_dir_list[i]
        run_csst_msc_mbi_photometry(image_path=image_path,
                                    output_dir=output_path,
                                    flag_weight_dir=fw_dir)

def main():
    args = parse_args()
    config_file = args.config_file
    with open(config_file, "r") as stream:
        try:
            config = yaml.safe_load(stream)
            for key, value in config.items():
                print(key + " : " + str(value))
        except yaml.YAMLError as exc:
            print(exc)
    
    run_pointing_list(input_dir=config["input_dir"],
                      flag_weight_dir=config["flag_weight_dir"],
                      pointing_label_list=config["pointing_label_list"],
                      output_dir=config["output_dir"],
                      chip_label_list=config["chip_label_list"])


if __name__ == "__main__":
    # input_dir = "/public/home/fangyuedong/project/injected_50sqDeg_L1_outputs"
    # pointing_label_list = ["MSC_0000000", "MSC_0000001",
    #                        "MSC_0000002", "MSC_0000003", "MSC_0000004", "MSC_0000005"]
    # chip_label_list = None
    # output_dir = "/public/home/fangyuedong/project/processed_injected_50sqDeg_L1_outputs"
    # flag_weight_dir = "/public/home/fangyuedong/project/50sqDeg_L1_outputs/50sqDeg_Photo_W2/"
    # flag_weight_dir = "/public/home/fangyuedong/project/50sqDeg_L1_outputs/50sqDeg_Photo_W3/"

    # # input_dir = "/public/home/fangyuedong/project/50sqDeg_L1_outputs/50sqDeg_Photo_W3/"
    # input_dir = "/public/home/fangyuedong/project/injected_50sqDeg_L1_outputs/50sqDeg_Photo_W3/"
    # pointing_label_list = ["MSC_0000000", "MSC_0000001",
    #                        "MSC_0000002", "MSC_0000003", "MSC_0000004", "MSC_0000005", "MSC_0000006", "MSC_0000007", "MSC_0000008", "MSC_0000009"]
    # chip_label_list = None
    # # output_dir = "/public/home/fangyuedong/project/50sqDeg_L1_outputs/50sqDeg_Photo_W3/"
    # output_dir = "/public/home/fangyuedong/project/processed_injected_50sqDeg_L1_outputs/50sqDeg_Photo_W3/"

    main()