added dm.query_rc() (7a6611cb) · Commits · csst-pipeline / csst_common

csst_common/data_manager.py

+42 −14

Original line number	Diff line number	Diff line
		@@ -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

+0 −2

Original line number	Diff line number	Diff line
		@@ -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)