comparison pipeline_sources.py @ 0:a35056104c2c draft default tip

planemo upload for repository https://github.com/esg-epfl-apc/tools-astro/tree/main/tools commit da42ae0d18f550dec7f6d7e29d297e7cf1909df2

0:a35056104c2c

import pandas as pd
import numpy as np
import re
import time
import io
import requests
from astroquery.simbad import Simbad
from astropy.coordinates import SkyCoord
from astropy import units as u

from query_tns_aux import parse_data, query_tns_main_name, query_tns_survey_name


def query_simbad(name):
    Simbad.add_votable_fields("otypes")

    table_ids = Simbad.query_objectids(name)
    table_obj = Simbad.query_object(name)

    dict_data = {}

    list_ids = None
    main_id = None
    otype_ = None
    ra = None
    dec = None
    if table_ids:
        df_ids = table_ids.to_pandas()
        df_ids.columns = df_ids.columns.str.upper()
        list_ids = list(df_ids["ID"].str.decode("utf-8").values)
    if table_obj:
        df_obj = table_obj.to_pandas()
        df_obj.columns = df_obj.columns.str.upper()
        main_id = df_obj["MAIN_ID"].values[0]
        otype_ = '|'.join(list(set(df_obj["OTYPES.OTYPE"].values)))
        str_ra_dec = str(df_obj["RA"].values[0]) + " " + str(df_obj["DEC"].values[0])
        ra_dec = SkyCoord(str_ra_dec, unit=(u.hourangle, u.deg))
        ra = ra_dec.ra.value
        dec = ra_dec.dec.value

    dict_data[name] = {"IDs": list_ids, "MAIN_ID": main_id, "OTYPES": otype_, "RA": ra, "DEC": dec, "DISCOVERY_TIME": None}

    return dict_data


def query_tns(name):
    dict_data = {}

    if ((name[0:3] == "at2") or (name[0:2] == "sn")):
        my_text = query_tns_main_name(name[2:])
    else:
        my_text = query_tns_survey_name(name)

    time.sleep(5)

    main_id, list_ids, list_otypes, ra, dec, discovery_time = parse_data(my_text)
    if ra and dec:
        str_ra_dec = ra + ":" + dec
        ra_dec = SkyCoord(str_ra_dec.replace(":", " "), unit=(u.hourangle, u.deg))
        ra = ra_dec.ra.value
        dec = ra_dec.dec.value

    if list_otypes is None:
        otype_ = None
    elif len(list_otypes) == 1:
        otype_ = list_otypes[0]
    else:
        otype_ = '|'.join(list_otypes)

    dict_data[name] = {"IDs": list_ids, "MAIN_ID": main_id, "OTYPES": otype_, "RA": ra, "DEC": dec, "DISCOVERY_TIME": discovery_time}

    return dict_data


def query_fink(name):
    dict_data = {}
    list_ids = None
    main_id = None
    otype_ = None
    ra = None
    dec = None
    discovery_time = None

    query_name = name.replace("ztf", "ZTF")
    r = requests.post('https://api.fink-portal.org/api/v1/objects', json={'objectId': query_name, 'output-format': 'json', 'columns': 'i:ra,i:dec,i:objectId'})
    # Format output in a DataFrame
    df = pd.read_json(io.BytesIO(r.content))
    if not df.empty:
        object_id = df["i:objectId"].values[0]
        main_id = object_id
        ra_tmp = df["i:ra"][df["i:objectId"] == object_id].values
        dec_tmp = df["i:dec"][df["i:objectId"] == object_id].values
        ra = np.mean(np.float32(ra_tmp))
        dec = np.mean(np.float32(dec_tmp))
        list_ids = '|'.join(list(set(df["i:objectId"].values)))

    dict_data[name] = {"IDs": list_ids, "MAIN_ID": main_id, "OTYPES": otype_, "RA": ra, "DEC": dec, "DISCOVERY_TIME": discovery_time}

    return dict_data


def create_pattern_list():
    pattern_list_low = []
    pattern_list_low += ["\\b(ngc) *?([0-9]{1,4})\\b", "\\b(m) *?([0-9]{1,3})\\b"]
    pattern_list_low += ["\\b(ugc) *?([0-9]{1,5})\\b"]

    pattern_list_low += ["\\b(icecube|grb|frb|pks|mrk|hawc|maxi|gw)([ -]?)([0-9\\.\\-\\+]{2,}[a-z]?)\\b"]
    pattern_list_low += ["\\b(ic)([ -]?)([0-9]{1,4})\\b"]
    pattern_list_low += ["\\b(ztf) *?([0-9]{2}[a-z]{7})\\b"]

    pattern_list_low += ["\\b(at|sn) *?([1-2]{1}[0-9]{3}[a-z]{1,4})\\b"]
    pattern_list_low += ["\\b(asas)([ -]?)(sn)([ -]?)([0-9]{2}[a-z]{2,3})\\b"]

    pattern_list_low += ["\\b(ps|gaia) *?([1-2]{1}[0-9]{1}[a-z]{1,4})\\b"]

    pattern_list_low += ["\\b(m31n) *?([0-9]{4})-([0-9]{2}[a-z]{1})\\b"]
    pattern_list_low += ["\\b(ptf|atlas) *?([0-9]{2}[a-z]{1,4})\\b"]

    pattern_list_low += ["\\b(4c) *?((\\+|-)[0-9]{2}\\.[0-9]{2}\\.[0-9]{1})\\b"]
    pattern_list_low += ["\\b(4c) *?((\\+|-)[0-9]{2}\\.[0-9]{2})\\b"]

    # new from lm-astronomy
    pattern_list_low += ["\\b(lsq12) *?([a-z]{3})\\b"]
    pattern_list_low += ["\\b(des14) *?([a-z]{1}[0-9]{1})\\b"]

    pattern_list = []

    HH = "(2[0-4]|[0-1][0-9])"
    MM = "(90|[0-8][0-9])"
    MMm = "(90|[0-8][0-9])([0-9])"
    SS = "(90|[0-8][0-9])"
    DD = "(90|[0-8][0-9])"
    pDDd = "(\\+|-)(90|[0-8][0-9])([0-9])"

    LLL = "([0-9][0-9][0-9])"
    BB = "([0-9][0-9])"
    VVV = "([0-9][0-9][0-9])"

    NAAA = "(([0-9][A-Z]{1,4})|([A-Z]{2,4}))"

    pattern_list += [f"\\b{NAAA} *?([0-9]{{6}})(?!-)\\b"]

    pattern_list += [f"\\b{NAAA} *?(J[0-9]{{6}}(\\+|-)[0-9]{{5}})(?!-)\\b"]
    pattern_list += [f"\\b{NAAA} *?(J[0-9]{{6}}(\\.?)[0-9]{{1}}(\\+|-)[0-9]{{6}})(?!-)\\b"]

    pattern_list += [f"\\b{NAAA} *?(J[0-9]{{4}}(\\.?)[0-9]{{1}}(\\+|-)[0-9]{{4}})[a-z]{{0,1}}(?!-)\\b"]
    pattern_list += [f"\\b{NAAA} *?{HH}{MM}(\\+|-){DD}(?!-)\\b"]
    pattern_list += [f"\\b{NAAA} *?(J?){HH}{MM}{pDDd}(?!-)\\b"]
    pattern_list += [f"\\b{NAAA} *?B{HH}{MM}(\\+|-){DD}(?!-)\\b"]
    pattern_list += [f"\\b{NAAA} *?B{HH}{MM}{pDDd}(?!-)\\b"]
    pattern_list += [f"\\b{NAAA} *?J{HH}{MM}(\\+|-){DD}{MM}(?!-)\\b"]
    pattern_list += [f"\\b{NAAA} *?J{HH}{MM}([0-9])(\\+|-){DD}{MM}(?!-)\\b"]

    pattern_list += [f"\\b{NAAA} *?{HH}{MMm}(\\+|-){DD}{MM}(?!-)\\b"]
    pattern_list += [f"\\b{NAAA} *?{HH}{MM}(\\.?)([0-9])(\\+|-){DD}{MM}(?!-)\\b"]
    pattern_list += [f"\\b{NAAA} *?{HH}{MM}{SS}(\\+|-){DD}{MM}{SS}(?!-)\\b"]
    pattern_list += [f"\\b{NAAA} *?{HH}{MM}{SS}(\\.?)([0-9])(\\+|-){DD}{MM}{SS}(?!-)\\b"]
    pattern_list += [f"\\b{NAAA} *?B{HH}{MM}{SS}(\\.?)([0-9])(\\+|-){DD}{MM}{SS}(?!-)\\b"]
    pattern_list += [f"\\b{NAAA} *?J{HH}{MM}{SS}(\\.?)([0-9][0-9])(\\+|-){DD}{MM}{SS}(\\.?)([0-9])(?!-)\\b"]

    pattern_list += [f"\\b{NAAA} *?{LLL}(\\.?)([0-9])(\\+|-){BB}(\\.?)([0-9])(\\+|-){VVV}(?!-)\\b"]
    pattern_list += [f"\\b{NAAA} *?G{LLL}(\\.?)([0-9])(\\+|-){BB}(\\.?)([0-9])(\\+|-){VVV}(?!-)\\b"]

    FFF = "([0-9][0-9][0-9])"
    FFFF = "([0-9][0-9][0-9][0-9])"
    NNNN = FFFF
    NNNNN = "([0-9][0-9][0-9][0-9][0-9])"

    pattern_list += [f"\\b{NAAA} *?{FFF}-{NNNN}(?!-)\\b"]
    pattern_list += [f"\\b{NAAA} *?{FFFF}-{NNNNN}(?!-)\\b"]

    YY = "[0-9][0-9]"
    MM = "[0-1][0-9]"
    DD = "[0-3][0-9]"
    pattern_list += [f"\\b{NAAA} *?{YY}{MM}{DD}A{{0,1}}(?!-)\\b"]

    return pattern_list, pattern_list_low


def rule_based_source_detector(text_id, text_id_text):
    pattern_list, pattern_list_low = create_pattern_list()

    regex_sources = []

    for pattern in pattern_list_low:
        for m in re.finditer(pattern, text_id_text.lower()):
            source_ = m.group(0).replace(" ", "")

            if "asas" in source_:
                source_ = source_.replace("-", "")
                source_ = source_.replace("sn", "sn-")

            regex_sources.append(source_)

    for pattern in pattern_list:
        for m in re.finditer(pattern, text_id_text):
            source_ = m.group(0)
            if source_.replace(" ", "").lower() not in regex_sources:
                regex_sources.append(source_)

    return list(set(regex_sources))


def query_info_sources(text_id, sources):
    if len(sources) != 0:
        source_list = []
        otype_list = []
        mainid_list = []
        ra_list = []
        dec_list = []
        pattern_string = re.compile("[a-z]")
        dict_unknown = {}
        dict_unknown = {"Raw Source Name": []}

        for source_name_0 in sources:
            # TODO:  remove all non a-z0-9 characters from the beginning and ending of the source name
            # Now: remove "," from the beginning and ending of the source name
            if source_name_0[0] == ",":
                source_name = source_name_0[1:]
            elif source_name_0[-1] == ",":
                source_name = source_name_0[:-1]
            else:
                source_name = source_name_0

            if pattern_string.findall(source_name.lower()):
                dict_otype = query_simbad(source_name)
                if dict_otype[source_name]["MAIN_ID"] is None:
                    dict_otype = query_tns(source_name)
                if dict_otype[source_name]["MAIN_ID"] is None:
                    dict_otype = query_fink(source_name)

                if dict_otype[source_name]["MAIN_ID"] is None:
                    dict_unknown["Raw Source Name"].append(source_name)
                else:
                    mainid_list.append(dict_otype[source_name]["MAIN_ID"])
                    otype_list.append(dict_otype[source_name]["OTYPES"])
                    ra_list.append(dict_otype[source_name]["RA"])
                    dec_list.append(dict_otype[source_name]["DEC"])
                    source_list.append(source_name)
            else:
                dict_unknown["Raw Source Name"].append(source_name_0)

        dict_data = {"TEXT_ID": [text_id] * len(source_list), "Raw Source Name": source_list, "Main ID Name": mainid_list, "OTYPE": otype_list, "RA": ra_list, "Dec": dec_list}
        df_save = pd.DataFrame(dict_data)
        df_save.replace({None: "NotKnown"}, inplace=True)
        return df_save.drop_duplicates(subset=['Main ID Name']), pd.DataFrame(dict_unknown)

    return pd.DataFrame(), pd.DataFrame()