run_csst_msc_instrument.py

"""_summary_
"""
import os
from glob import glob
import mpi4py.MPI as MPI
# from .L1_pipeline.csst_msc_instrument.csst_msc_mbi_instrument import core_msc_l1_mbi_instrument
from L1_pipeline.csst_msc_instrument.csst_msc_mbi_instrument import core_msc_l1_mbi_instrument


def run_csst_msc_instrument(image_path,
                            output_dir,
                            calib_data_path="/public/home/fangyuedong/project/calib_data/"):
    """_summary_

    Args:
        image_path (_type_): _description_
        output_dir (_type_): _description_
        calib_data_path (str, optional): _description_. Defaults to "/public/home/fangyuedong/project/calib_data/".
    """

    img_filename = os.path.splitext(os.path.basename(image_path))[0]
    chip_label = img_filename.split("_")[-3]

    # Get paths to corresponding calibration files
    bias_path = os.path.join(calib_data_path, "bias_%s.fits" % (chip_label))
    dark_path = os.path.join(calib_data_path, "dark_%s.fits" % (chip_label))
    flat_path = os.path.join(calib_data_path, "flat_%s.fits" % (chip_label))

    output_img_types = ['img', 'wht', 'flg']
    output_paths = []
    for i in range(len(output_img_types)):
        output_paths.append(os.path.join(output_dir,
                                         img_filename[:-6] + output_img_types[i] + "_L1.fits"))

    core_msc_l1_mbi_instrument(
        image_path=image_path,
        bias_path=bias_path,
        dark_path=dark_path,
        flat_path=flat_path,
        shutter_path="/dummy/csst_msc_ms_shutter_09_000001.fits",  # Didn't use at this moment
        image_output_path=output_paths[0],
        weight_output_path=output_paths[1],
        flag_output_path=output_paths[2],
    )


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_*_' + '*_*')
    else:
        image_path_list = []
        for chip_label in chip_label_list:
            image_path = glob(pointing_dir + '/CSST_MSC_MS_SCIE_*_' +
                              chip_label + '_*')[0]
            image_path_list.append(image_path)

    return image_path_list


def run_pointing_list(input_dir,
                      pointing_label_list,
                      output_dir,
                      calib_data_path="/public/home/fangyuedong/project/calib_data/",
                      chip_label_list=None):
    """_summary_

    Args:
        input_dir (_type_): _description_
        pointing_label_list (_type_): _description_
        output_dir (_type_): _description_
        calib_data_path (str, optional): _description_. Defaults to "/public/home/fangyuedong/project/calib_data/".
        chip_label_list (_type_, optional): _description_. Defaults to None.
    """
    image_path_list = []
    output_path_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)
        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)

    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]
        run_csst_msc_instrument(image_path=image_path,
                                output_dir=output_path,
                                calib_data_path=calib_data_path)


if __name__ == "__main__":
    # input_dir = "/public/share/yangxuliu/CSSOSDataProductsSims/outputs_50sqDeg/50sqDeg_Photo_W1/"
    # pointing_label_list = ["MSC_0000000"]
    # chip_label_list = ["08"]
    # output_dir = "/public/home/fangyuedong/project/test_deepcr"
    # input_dir = "/public/share/yangxuliu/CSSOSDataProductsSims/outputs_50sqDeg/50sqDeg_Photo_W2/"
    # output_dir = "/public/home/fangyuedong/project/50sqDeg_L1_outputs/50sqDeg_Photo_W2/"
    input_dir = "/public/share/yangxuliu/CSSOSDataProductsSims/outputs_50sqDeg/50sqDeg_Photo_W3/"
    output_dir = "/public/home/fangyuedong/project/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

    run_pointing_list(input_dir=input_dir,
                      pointing_label_list=pointing_label_list,
                      output_dir=output_dir,
                      chip_label_list=chip_label_list)