update readfos for addition and copy, query for fos

package/lib/query.py

@@ -25,11 +25,11 @@ def divide_proposal(products):
     """
     for pid in np.unique(products["Proposal ID"]):
         obs = products[products["Proposal ID"] == pid].copy()
-        same_filt = np.unique(np.array(np.sum([obs["Filters"][:, 1:] == filt[1:] for filt in obs["Filters"]], axis=2) < 3, dtype=bool), axis=0)
+        same_filt = np.unique(np.array(np.sum([obs["Filters"] == filt for filt in obs["Filters"]], axis=2) >= len(obs["Filters"][0]), dtype=bool), axis=0)
         if len(same_filt) > 1:
             for filt in same_filt:
                 products["Proposal ID"][np.any([products["Dataset"] == dataset for dataset in obs["Dataset"][filt]], axis=0)] = "_".join(
-                    [obs["Proposal ID"][filt][0], "_".join([fi for fi in obs["Filters"][filt][0][1:] if fi[:-1] != "CLEAR"])]
+                    [obs["Proposal ID"][filt][0], "_".join([fi for fi in obs["Filters"][filt][0] if fi[:-1] != "CLEAR"])]
                 )
     for pid in np.unique(products["Proposal ID"]):
         obs = products[products["Proposal ID"] == pid].copy()
@@ -44,38 +44,69 @@ def divide_proposal(products):
     return products
 
 
-def get_product_list(target=None, proposal_id=None):
+def get_product_list(target=None, proposal_id=None, instrument="foc"):
     """
     Retrieve products list for a given target from the MAST archive
     """
     mission = MastMissions(mission="hst")
     radius = "3"
     select_cols = [
-        "sci_data_set_name",
-        "sci_spec_1234",
-        "sci_actual_duration",
-        "sci_start_time",
-        "sci_stop_time",
-        "sci_central_wavelength",
-        "sci_instrume",
-        "sci_aper_1234",
-        "sci_targname",
         "sci_pep_id",
         "sci_pi_last_name",
+        "sci_targname",
+        "sci_aper_1234",
+        "sci_spec_1234",
+        "sci_central_wavelength",
+        "sci_actual_duration",
+        "sci_instrume",
+        "sci_operating_mode",
+        "sci_data_set_name",
+        "sci_start_time",
+        "sci_stop_time",
+        "sci_refnum",
     ]
 
-    cols = ["Dataset", "Filters", "Exptime", "Start", "Stop", "Central wavelength", "Instrument", "Size", "Target name", "Proposal ID", "PI last name"]
+    cols = [
+        "Proposal ID",
+        "PI last name",
+        "Target name",
+        "Aperture",
+        "Filters",
+        "Central wavelength",
+        "Exptime",
+        "Instrument",
+        "Operating Mode",
+        "Dataset",
+        "Start",
+        "Stop",
+        "References",
+    ]
 
     if target is None:
         target = input("Target name:\n>")
 
     # Use query_object method to resolve the object name into coordinates
-    results = mission.query_object(target, radius=radius, select_cols=select_cols, sci_spec_1234="POL*", sci_obs_type="image", sci_aec="S", sci_instrume="foc")
+    if instrument == "foc":
+        results = mission.query_object(
+            target, radius=radius, select_cols=select_cols, sci_spec_1234="POL*", sci_obs_type="image", sci_aec="S", sci_instrume="foc"
+        )
+        dataproduct_type = "image"
+        description = "DADS C0F file - Calibrated exposure WFPC/WFPC2/FOC/FOS/GHRS/HSP"
+    elif instrument == "fos":
+        results = mission.query_object(
+            target, radius=radius, select_cols=select_cols, sci_operating_mode="SPECTROPOLARIMETRY", sci_obs_type="spectrum", sci_aec="S", sci_instrume="fos"
+        )
+        dataproduct_type = "spectrum"
+        description = ["DADS C0F file - Calibrated exposure WFPC/WFPC2/FOC/FOS/GHRS/HSP", "DADS C3F file - Calibrated exposure GHRS/FOS/HSP"]
 
     for c, n_c in zip(select_cols, cols):
         results.rename_column(c, n_c)
     results["Proposal ID"] = Column(results["Proposal ID"], dtype="U35")
-    results["Filters"] = Column(np.array([filt.split(";") for filt in results["Filters"]], dtype=str))
+    if instrument == "foc":
+        results["POLFilters"] = Column(np.array([filt.split(";")[0] for filt in results["Filters"]], dtype=str))
+        results["Filters"] = Column(np.array([filt.split(";")[1:] for filt in results["Filters"]], dtype=str))
+    else:
+        results["Filters"] = Column(np.array([filt.split(";") for filt in results["Filters"]], dtype=str))
     results["Start"] = Column(Time(results["Start"]))
     results["Stop"] = Column(Time(results["Stop"]))
 
@@ -89,20 +120,34 @@ def get_product_list(target=None, proposal_id=None):
             to_remove.append(i)
     obs.remove_rows(to_remove)
     # Remove observations for which a polarization filter is missing
-    polfilt = {"POL0": 0, "POL60": 1, "POL120": 2}
-    for pid in np.unique(obs["Proposal ID"]):
-        used_pol = np.zeros(3)
-        for dataset in obs[obs["Proposal ID"] == pid]:
-            used_pol[polfilt[dataset["Filters"][0]]] += 1
-        if np.any(used_pol < 1):
-            obs.remove_rows(np.arange(len(obs))[obs["Proposal ID"] == pid])
+    if instrument == "foc":
+        polfilt = {"POL0": 0, "POL60": 1, "POL120": 2}
+        for pid in np.unique(obs["Proposal ID"]):
+            used_pol = np.zeros(3)
+            for dataset in obs[obs["Proposal ID"] == pid]:
+                used_pol[polfilt[dataset["POLFilters"]]] += 1
+            if np.any(used_pol < 1):
+                obs.remove_rows(np.arange(len(obs))[obs["Proposal ID"] == pid])
+    # Remove observations for which a spectropolarization has not been reduced
+    if instrument == "fos":
+        for pid in np.unique(obs["Proposal ID"]):
+            observations = Observations.query_criteria(proposal_id=pid.split("_")[0])
+            c3prod = Observations.filter_products(
+                Observations.get_product_list(observations),
+                productType=["SCIENCE"],
+                dataproduct_type=dataproduct_type,
+                calib_level=[2],
+                description=description[1],
+            )
+            if len(c3prod) < 1:
+                obs.remove_rows(np.arange(len(obs))[obs["Proposal ID"] == pid])
 
     tab = unique(obs, ["Target name", "Proposal ID"])
     obs["Obs"] = [np.argmax(np.logical_and(tab["Proposal ID"] == data["Proposal ID"], tab["Target name"] == data["Target name"])) + 1 for data in obs]
     try:
-        n_obs = unique(obs[["Obs", "Filters", "Start", "Central wavelength", "Instrument", "Size", "Target name", "Proposal ID", "PI last name"]], "Obs")
+        n_obs = unique(obs[["Obs", "Filters", "Start", "Central wavelength", "Instrument", "Aperture", "Target name", "Proposal ID", "PI last name"]], "Obs")
     except IndexError:
-        raise ValueError("There is no observation with POL0, POL60 and POL120 for {0:s} in HST/FOC Legacy Archive".format(target))
+        raise ValueError("There is no observation with polarimetry for {0:s} in HST/{1:s} Legacy Archive".format(target, instrument.upper()))
 
     b = np.zeros(len(results), dtype=bool)
     if proposal_id is not None and str(proposal_id) in obs["Proposal ID"]:
@@ -130,29 +175,27 @@ def get_product_list(target=None, proposal_id=None):
     products = Observations.filter_products(
         Observations.get_product_list(observations),
         productType=["SCIENCE"],
-        dataproduct_type=["image"],
+        dataproduct_type=dataproduct_type,
         calib_level=[2],
-        description="DADS C0F file - Calibrated exposure WFPC/WFPC2/FOC/FOS/GHRS/HSP",
+        description=description,
     )
+
     products["proposal_id"] = Column(products["proposal_id"], dtype="U35")
-    products["target_name"] = Column(observations["target_name"])
 
     for prod in products:
         prod["proposal_id"] = results["Proposal ID"][results["Dataset"] == prod["productFilename"][: len(results["Dataset"][0])].upper()][0]
 
-    for prod in products:
-        prod["target_name"] = observations["target_name"][observations["obsid"] == prod["obsID"]][0]
-    tab = unique(products, ["target_name", "proposal_id"])
+    tab = unique(products, "proposal_id")
 
-    products["Obs"] = [np.argmax(np.logical_and(tab["proposal_id"] == data["proposal_id"], tab["target_name"] == data["target_name"])) + 1 for data in products]
+    products["Obs"] = [np.argmax(tab["proposal_id"] == data["proposal_id"]) + 1 for data in products]
     return target, products
 
 
-def retrieve_products(target=None, proposal_id=None, output_dir="./data"):
+def retrieve_products(target=None, proposal_id=None, instrument="foc", output_dir="./data"):
     """
     Given a target name and a proposal_id, create the local directories and retrieve the fits files from the MAST Archive
     """
-    target, products = get_product_list(target=target, proposal_id=proposal_id)
+    target, products = get_product_list(target=target, proposal_id=proposal_id, instrument=instrument)
     prodpaths = []
     # data_dir = path_join(output_dir, target)
     out = ""
@@ -183,10 +226,11 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Query MAST for target products")
     parser.add_argument("-t", "--target", metavar="targetname", required=False, help="the name of the target", type=str, default=None)
     parser.add_argument("-p", "--proposal_id", metavar="proposal_id", required=False, help="the proposal id of the data products", type=int, default=None)
+    parser.add_argument("-i", "--instrum", metavar="instrum", required=False, help="the instrument used for observation", type=str, default="foc")
     parser.add_argument(
         "-o", "--output_dir", metavar="directory_path", required=False, help="output directory path for the data products", type=str, default="./data"
     )
     args = parser.parse_args()
     print(args.target)
-    prodpaths = retrieve_products(target=args.target, proposal_id=args.proposal_id)
+    prodpaths = retrieve_products(target=args.target, proposal_id=args.proposal_id, instrument=args.instrum.lower(), output_dir=args.output_dir)
     print(prodpaths)