added dm.query_rc()

csst_common/data_manager.py

@@ -9,16 +9,19 @@ Modified-History:
     2022-09-13, Bo Zhang, added CsstMbiDataManager
     2022-09-29, Bo Zhang, favor CsstMsDataManager instead of CsstMbiDataManager
     2022-10-26, Bo Zhang, reconstruct CsstMsDataManager, deprecate CsstMbiDataManager
+    2022-10-28, Bo Zhang, added CsstMsDataManager.query_rc(), dm.dfs_mode, dm.node
 """
-import os
 import glob
+import os
 import re
+from typing import Union
 
 from astropy.io import fits
 from astropy.table import Table
-from typing import Union, Tuple
 
+from .dfs import gaia_query_from_dfs
 from .params import CSST_PARAMS as CP
+from .params import DFS_CONF
 
 
 class CsstMsDataManager:
@@ -66,6 +69,9 @@ class CsstMsDataManager:
         The aux data directory (bias, flat, dark).
     dfs_mode : bool
         DFS mode. If True, ``CsstMsDataManager`` will use DFS.
+    node : str
+        The environment in which the pipeline will run.
+        Use "pml" for Purple Mountain Lab cluster and "local" for others.
     _exp_id : int
         The exposure ID.
     _exp_start : int
@@ -113,6 +119,7 @@ class CsstMsDataManager:
                  dir_l1: str = ".",
                  path_aux: str = "",  # bias dark flat
                  dfs_mode: bool = False,
+                 node: str = "local",
                  _exp_id: int = 100000100,
                  _exp_start: int = 20270810081950,
                  _exp_stop: int = 20270810082220,
@@ -160,7 +167,14 @@ class CsstMsDataManager:
         self._survey = _survey
         self._imagetype = _imagetype
         self._l0_post = _l0_post
+
+        # for catalog query
         self.dfs_mode = dfs_mode
+        # for DFS configuration, defaults to "local", could be "pml"
+        assert node in DFS_CONF.keys()
+        self.node = node
+        for k, v in DFS_CONF[node]:
+            os.putenv(k, v)
 
         # data directory
         self.dir_l0 = dir_l0
@@ -175,7 +189,7 @@ class CsstMsDataManager:
         raise NotImplementedError("from_dfs is currently not available!")
 
     @staticmethod
-    def from_dir(ver_sim="C5.2", datatype="mbi", dir_l0=".", dir_l1=".", dir_pcref="", path_aux="", dfs_mode=False):
+    def from_dir(ver_sim="C5.2", datatype="mbi", dir_l0=".", dir_l1=".", path_aux="", dfs_mode=False, node="local"):
         """ initialize the multi-band imaging data manager """
 
         assert ver_sim in ["C5.2", ]
@@ -206,6 +220,7 @@ class CsstMsDataManager:
                                  dir_l1=dir_l1,
                                  path_aux=path_aux,  # bias dark flat
                                  dfs_mode=dfs_mode,
+                                 node=node,
                                  _exp_id=_exp_id,
                                  _exp_start=_exp_start,
                                  _exp_stop=_exp_stop,
@@ -362,23 +377,23 @@ class CsstMsDataManager:
 
     def get_sls_info(self):
         """ Get the target SLS image header info and return. """
-        if self.dfs_mode:
-            raise NotImplementedError()
-        else:
+        # if self.dfs_mode:
+        #     raise NotImplementedError()
+        # else:
         assert len(self.target_detectors) == 1
         header = fits.getheader(self.l0_detector(self.target_detectors[0]), ext=1)
         return header
 
     def get_mbi_info(self):
         """ Get all MBI image header info and return as a table. """
-        if self.dfs_mode:
-            raise NotImplementedError()
-        else:
+        # if self.dfs_mode:
+        #     raise NotImplementedError()
+        # else:
         info = Table.read("/nfsdata/share/csst_simulation_data/Cycle-5-SimuData/slitlessSpectroscopy/t_mbi_l1.fits")
         return info
 
     @staticmethod
-    def quickstart(ver_sim="C5.2", datatype="mbi", dir_l1=".", exposure_id=100, dfs_mode=False):
+    def quickstart(ver_sim="C5.2", datatype="mbi", dir_l1=".", exposure_id=100, dfs_mode=False, node="local"):
         """Quick dataset generator for tests on dandelion or PMO
 
         Parameters
@@ -393,6 +408,8 @@ class CsstMsDataManager:
             The serial number of the exposure. 20-154 for C5.2.
         dfs_mode : bool
             If True, use DFS.
+        node : str
+            The node in which this program runs. Defaults to "local", could be "pml".
 
         Returns
         -------
@@ -428,7 +445,18 @@ class CsstMsDataManager:
             raise ValueError("@DM: invalid hostname {} or datatype {}!".format(hostname, datatype))
 
         return CsstMsDataManager.from_dir(
-            ver_sim=ver_sim, datatype=datatype, dir_l0=dir_l0, dir_l1=dir_l1, path_aux=path_aux, dfs_mode=dfs_mode)
+            ver_sim=ver_sim, datatype=datatype, dir_l0=dir_l0, dir_l1=dir_l1, path_aux=path_aux,
+            dfs_mode=dfs_mode, node=node)
+
+    def query_rc(self, ra=180, dec=0, radius=2, min_mag=0, max_mag=20, obstime=-1, limit=-1):
+        """ Query Reference Catalog (RC) from DFS """
+        try:
+            tbl = gaia_query_from_dfs(
+                ra=ra, dec=dec, radius=radius, min_mag=min_mag, max_mag=max_mag, obstime=obstime, limit=limit)
+            assert len(tbl) > 0
+            return tbl
+        except Exception as e:
+            return None
 
 
 class CsstMbiDataManager:

csst_common/params.py

@@ -17,5 +17,3 @@ with open(PACKAGE_PATH + "/data/csst_params.yml") as f:
 
 with open(PACKAGE_PATH + "/data/dfs_conf.yml") as f:
     DFS_CONF = yaml.safe_load(f)
-
-