reconstructed CsstMsDataManager (500a78ec) · Commits · csst-pipeline / csst_common

csst_common/data/csst_params.yml

+1 −1

Original line number	Diff line number	Diff line
		@@ -81,7 +81,7 @@ sim: # simulation
		- C5.1
		- C3

		basic:
		all:
		detectors:
		- 1
		- 2

csst_common/data_manager.py

+114 −366

Original line number	Diff line number	Diff line
		@@ -8,12 +8,12 @@ Modified-History:
		2022-09-13, Bo Zhang, created
		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
		"""
		import os
		import glob
		import re

		import numpy as np
		from astropy.io import fits

		from .params import CSST_PARAMS as CP
		@@ -53,8 +53,6 @@ class CsstMsDataManager:
		Will be removed in the next version.
		path_aux : str
		The aux data directory (bias, flat, dark).
		assert_all_detectors : bool
		If True, assert data for all detectors are available.
		datatype : str
		The options are {"mbi", "sls", "all"}.
		The "all" option is used for QC in particular.
		@@ -83,15 +81,32 @@ class CsstMsDataManager:
		>>> dm_mbi.l1_file("flipped_image.fits")
		"""

		def __init__(self, ver_sim="C5.2", dir_l0="", dir_l1="", dir_pcref="", path_aux="", assert_all_detectors=False,
		datatype="mbi"):
		""" initialize the multi-band imaging data manager """
		def __init__(self,
		ver_sim="C5.2",
		datatype="mbi",
		available_detectors=None,
		target_detectors=None,
		dir_l0=".",
		dir_l1=".",
		path_aux="", # bias dark flat
		dir_pcref="", # deprecated
		_exp_id=100000100,
		_exp_start=20270810081950,
		_exp_stop=20270810082220,
		_telescope="CSST",
		_instrument="MSC",
		_survey="MS",
		_imagetype="SCI",
		_l0_post="L0_1",
		):

		# version
		assert ver_sim in CP["sim"]["versions"]
		self.ver_sim = ver_sim

		# datatype, valid_detectors, detector2filter
		assert datatype in ["mbi", "sls", "all"]
		self.datatype = datatype

		if datatype == "mbi":
		# MBI
		self.valid_detectors = CP["mbi"]["detectors"]
		@@ -102,59 +117,81 @@ class CsstMsDataManager:
		self.detector2filter = CP["sls"]["detector2filter"]
		else:
		# ALL
		self.valid_detectors = CP["basic"]["detectors"]
		self.detector2filter = CP["basic"]["detector2filter"]
		self.valid_detectors = CP["all"]["detectors"]
		self.detector2filter = CP["all"]["detector2filter"]

		# available_detectors
		self.available_detectors = available_detectors
		# set all available detectors by default
		self.target_detectors = []
		self.set_detectors(target_detectors)

		# exposure info
		self._exp_id = _exp_id
		self._exp_start = _exp_start
		self._exp_stop = _exp_stop

		# file name components
		self._telescope = _telescope
		self._instrument = _instrument
		self._survey = _survey
		self._imagetype = _imagetype
		self._l0_post = _l0_post

		# data directory
		self.dir_l0 = dir_l0
		self.dir_l1 = dir_l1
		self.dir_pcref = dir_pcref
		self.path_aux = path_aux
		self.ver_sim = ver_sim
		self.target_detectors = []

		# record hard code names in history
		self.hardcode_history = []

		fps_img = self.glob_image(dir_l0, ver_sim=ver_sim)
		fps_cat = self.glob_cat(dir_l0, ver_sim=ver_sim)

		if assert_all_detectors:
		assert len(fps_img) == len(self.valid_detectors)
		else:
		assert len(fps_img) > 0

		if ver_sim == "C3":
		# get info
		# print(re.split(r"[_.]", fps[0]))
		self._instrument, self._survey, \
		self._exp_start, self._exp_id, \
		_detector, self._l0_suffix, _ext = re.split(r"[_.]", fps_img[0])
		self._cat_id = re.split(r"[_.]", fps_cat[0])[1]

		self._exp_start = int(self._exp_start)
		self._exp_id = int(self._exp_id)
		@staticmethod
		def from_dfs(ver_sim="C5.2", data_type="mbi", exp_id=10000100, dir_l1="."):
		raise NotImplementedError("from_dfs is currently not available!")

		# available detectors
		self.available_detectors = [int(re.split(r"[_.]", fp)[4]) for fp in fps_img]
		self.available_detectors.sort()
		@staticmethod
		def from_dir(ver_sim="C5.2", datatype="mbi", dir_l0=".", dir_l1=".", dir_pcref="", path_aux=""):
		""" initialize the multi-band imaging data manager """

		elif ver_sim in ["C5.1", "C5.2"]:
		# get info
		# print(re.split(r"[_.]", fps[0]))
		self._telescope, self._instrument, self._survey, self._imagetype, \
		self._exp_start, self._exp_stop, self._exp_id, \
		_detector, self._l0_suffix, self._version, _ext = re.split(r"[_.]", fps_img[0])
		self._cat_id = re.split(r"[_.]", fps_cat[0])[1]
		assert ver_sim in ["C5.2", ]

		self._exp_start = int(self._exp_start)
		self._exp_stop = int(self._exp_stop)
		self._exp_id = int(self._exp_id)
		# glob files
		fps_img = CsstMsDataManager.glob_image(dir_l0, ver_sim=ver_sim)
		if len(fps_img) == 0:
		raise FileNotFoundError(f"No file found in dir_l0: {dir_l0}")

		# available detectors
		self.available_detectors = [int(re.split(r"[_.]", fp)[7]) for fp in fps_img]
		self.available_detectors.sort()

		# set all available detectors by default
		self.set_detectors(None)
		available_detectors = [int(re.split(r"[_.]", fp)[7]) for fp in fps_img]
		available_detectors.sort()

		# parse info
		_telescope, _instrument, _survey, _imagetype, \
		_exp_start, _exp_stop, _exp_id, \
		_detector, *_l0_post, _ext = re.split(r"[_.]", fps_img[0])

		_exp_start = int(_exp_start)
		_exp_stop = int(_exp_stop)
		_exp_id = int(_exp_id)

		return CsstMsDataManager(ver_sim=ver_sim,
		datatype=datatype,
		available_detectors=available_detectors,
		target_detectors=None,
		dir_l0=dir_l0,
		dir_l1=dir_l1,
		path_aux=path_aux, # bias dark flat
		dir_pcref=dir_pcref, # deprecated
		_exp_id=_exp_id,
		_exp_start=_exp_start,
		_exp_stop=_exp_stop,
		_telescope=_telescope,
		_instrument=_instrument,
		_survey=_survey,
		_imagetype=_imagetype,
		_l0_post="_".join(_l0_post),
		)

		@staticmethod
		def glob_image(dir_l0, ver_sim="C5"):
		@@ -174,9 +211,6 @@ class CsstMsDataManager:
		@staticmethod
		def glob_cat(dir_l0, ver_sim="C5"):
		""" glob input catalogs in L0 data directory """
		if ver_sim == "C3":
		pattern = os.path.join(dir_l0, "MSC_*.cat")
		else:
		assert ver_sim in ["C5.1", "C5.2"]
		pattern = os.path.join(dir_l0, "MSC_*.cat")
		fps = glob.glob(pattern)
		@@ -188,33 +222,20 @@ class CsstMsDataManager:

		def l0_cat(self, detector=6):
		""" the L0 cat file path"""
		if self.ver_sim == "C3":
		fn = "{}_{}_{:07d}_{:02d}.cat".format(
		self._instrument, self._cat_id, self._exp_id - 100000000, detector)
		elif self.ver_sim == "C5.1":
		fn = "{}_{}_chip_{:02d}_filt_{}.cat".format(
		self._instrument, self._exp_id - 90000000, detector, self.detector2filter[detector])
		elif self.ver_sim == "C5.2":
		assert self.ver_sim == "C5.2"
		fn = "{}_{}_chip_{:02d}_filt_{}.cat".format(
		self._instrument, self._exp_id, detector, self.detector2filter[detector])
		return os.path.join(self.dir_l0, fn)

		def l0_log(self, detector=6):
		""" L0 log file path """
		if self.ver_sim == "C5.1":
		fn = "{}_{}_chip_{:02d}_filt_{}.log".format(
		self._instrument, self._exp_id - 90000000, detector, self.detector2filter[detector])
		elif self.ver_sim == "C5.2":
		assert self.ver_sim == "C5.2"
		fn = "{}_{}_chip_{:02d}_filt_{}.log".format(
		self._instrument, self._exp_id, detector, self.detector2filter[detector])
		return os.path.join(self.dir_l0, fn)

		def l0_detector(self, detector=6):
		""" L0 detector-specific image file path """
		if self.ver_sim == "C3":
		fn = "{}_{}_{}_{}_{:02d}_raw.fits".format(
		self._instrument, self._survey, self._exp_start, self._exp_id, detector)
		else:
		assert self.ver_sim in ["C5.1", "C5.2"]
		fn = "{}_{}_{}_SCI_{}_{}_{}_{:02d}_L0_1.fits".format(
		self._telescope, self._instrument, self._survey,
		@@ -223,10 +244,6 @@ class CsstMsDataManager:

		def l0_crs(self, detector=6):
		""" L0 cosmic ray file path """
		if self.ver_sim == "C3":
		fn = "{}_CRS_{}_{}_{:02d}_raw.fits".format(
		self._instrument, self._exp_start, self._exp_id, detector)
		else:
		assert self.ver_sim in ["C5.1", "C5.2"]
		fn = "{}_{}_{}_CRS_{}_{}_{}_{:02d}_L0_1.fits".format(
		self._telescope, self._instrument, self._survey,
		@@ -249,11 +266,6 @@ class CsstMsDataManager:
		L1 file path

		"""
		if self.ver_sim == "C3":
		fn = "{}_{}_{}_{}_{:02d}_{}".format(
		self._instrument, self._survey,
		self._exp_start, self._exp_id, detector, post)
		else:
		assert self.ver_sim in ["C5.1", "C5.2"]
		fn = "{}_{}_{}_SCI_{}_{}_{}_{:02d}_{}".format(
		self._telescope, self._instrument, self._survey,
		@@ -338,12 +350,12 @@ class CsstMsDataManager:
		dir_l1 : str
		output directory
		exposure_id : int
		21-154 for C5.2
		The serial number of the exposure. 20-154 for C5.2.

		Returns
		-------
		CsstMsDataManager
		the MBI data manager instance
		The Main Survey Data Manager instance.
		"""
		assert datatype in ["mbi", "sls"]
		# auto identify node name
		@@ -377,31 +389,11 @@ class CsstMsDataManager:
		else:
		raise ValueError("@DM: invalid hostname {} or datatype {}!".format(hostname, datatype))

		return CsstMsDataManager(
		return CsstMsDataManager.from_dir(
		ver_sim=ver_sim, datatype=datatype, dir_l0=dir_l0, dir_l1=dir_l1, dir_pcref=dir_pcref, path_aux=path_aux)


		class CsstMbiDataManager:
		""" this class defines the file format of the input & output of CSST MSC L1 pipeline

		C3:
		MSC_MS_210525220000_100000020_06_raw.fits
		MSC_CRS_210525220000_100000020_06_raw.fits
		MSC_210525120000_0000020_06.cat

		C5.1:
		CSST_MSC_MS_SCI_20270810081950_20270810082220_100000100_06_L0_1.fits
		CSST_MSC_MS_CRS_20270810081950_20270810082220_100000100_06_L0_1.fits
		MSC_10000100_chip_06_filt_y.cat
		MSC_10000100_chip_06_filt_y.log

		C5.2
		CSST_MSC_MS_SCI_20270810081950_20270810082220_100000100_06_L0_1.fits
		CSST_MSC_MS_CRS_20270810081950_20270810082220_100000100_06_L0_1.fits
		MSC_100000100_chip_06_filt_y.cat
		MSC_100000100_chip_06_filt_y.log

		"""

		def __init__(self, ver_sim="C5.2", dir_l0="", dir_l1="", dir_pcref="", path_aux="", force_all_detectors=False):
		""" initialize the multi-band imaging data manager
		@@ -443,249 +435,5 @@ class CsstMbiDataManager:
		>>> # define an L1 file (non-detector-specified)
		>>> dm_mbi.l1_file("flipped_image.fits")
		"""
		# raise DeprecationWarning(
		# "CsstMbiDataManager will no longer be available in some days, please use CsstMsDataManager instead.")
		assert ver_sim in CP["sim"]["versions"]

		self.dir_l0 = dir_l0
		self.dir_l1 = dir_l1
		self.dir_pcref = dir_pcref
		self.path_aux = path_aux
		self.ver_sim = ver_sim
		self.target_detectors = []

		self.hardcode_history = []

		fps_img = self.glob_image(dir_l0, ver_sim=ver_sim)
		fps_cat = self.glob_cat(dir_l0, ver_sim=ver_sim)

		if force_all_detectors:
		assert len(fps_img) == len(CP["mbi"]["detectors"])
		else:
		assert len(fps_img) > 0

		if ver_sim == "C3":
		# get info
		# print(re.split(r"[_.]", fps[0]))
		self._instrument, self._survey, \
		self._exp_start, self._exp_id, \
		_detector, self._l0_suffix, _ext = re.split(r"[_.]", fps_img[0])
		self._cat_id = re.split(r"[_.]", fps_cat[0])[1]

		self._exp_start = int(self._exp_start)
		self._exp_id = int(self._exp_id)

		# available detectors
		self.available_detectors = [int(re.split(r"[_.]", fp)[4]) for fp in fps_img]
		self.available_detectors.sort()

		elif ver_sim in ["C5.1", "C5.2"]:
		# get info
		# print(re.split(r"[_.]", fps[0]))
		self._telescope, self._instrument, self._survey, self._imagetype, \
		self._exp_start, self._exp_stop, self._exp_id, \
		_detector, self._l0_suffix, self._version, _ext = re.split(r"[_.]", fps_img[0])
		self._cat_id = re.split(r"[_.]", fps_cat[0])[1]

		self._exp_start = int(self._exp_start)
		self._exp_stop = int(self._exp_stop)
		self._exp_id = int(self._exp_id)

		# available detectors
		self.available_detectors = [int(re.split(r"[_.]", fp)[7]) for fp in fps_img]
		self.available_detectors.sort()

		@staticmethod
		def glob_image(dir_l0, ver_sim="C5"):
		""" glob files in L0 data directory """
		if ver_sim == "C3":
		pattern = os.path.join(dir_l0, "MSC_MS_*_raw.fits")
		else:
		assert ver_sim in ["C5.1", "C5.2"]
		pattern = os.path.join(dir_l0, "CSST_MSC_MS_SCI_*.fits")
		fps = glob.glob(pattern)
		fps = [os.path.basename(fp) for fp in fps]
		fps.sort()

		print("@DM.glob_dir: {} files found with pattern: {}".format(len(fps), pattern))
		return fps

		@staticmethod
		def glob_cat(dir_l0, ver_sim="C5"):
		""" glob input catalogs in L0 data directory """
		if ver_sim == "C3":
		pattern = os.path.join(dir_l0, "MSC_*.cat")
		else:
		assert ver_sim in ["C5.1", "C5.2"]
		pattern = os.path.join(dir_l0, "MSC_*.cat")
		fps = glob.glob(pattern)
		fps = [os.path.basename(fp) for fp in fps]
		fps.sort()

		print("@DM.glob_dir: {} files found with pattern: {}".format(len(fps), pattern))
		return fps

		def l0_cat(self, detector=6):
		""" the L0 cat file path"""
		if self.ver_sim == "C3":
		fn = "{}_{}_{:07d}_{:02d}.cat".format(
		self._instrument, self._cat_id, self._exp_id - 100000000, detector)
		elif self.ver_sim == "C5.1":
		fn = "{}_{}_chip_{:02d}_filt_{}.cat".format(
		self._instrument, self._exp_id - 90000000, detector, CP["mbi"]["detector2filter"][detector])
		elif self.ver_sim == "C5.2":
		fn = "{}_{}_chip_{:02d}_filt_{}.cat".format(
		self._instrument, self._exp_id, detector, CP["mbi"]["detector2filter"][detector])
		return os.path.join(self.dir_l0, fn)

		def l0_log(self, detector=6):
		""" L0 log file path """
		if self.ver_sim == "C5.1":
		fn = "{}_{}_chip_{:02d}_filt_{}.log".format(
		self._instrument, self._exp_id - 90000000, detector, CP["mbi"]["detector2filter"][detector])
		elif self.ver_sim == "C5.2":
		fn = "{}_{}_chip_{:02d}_filt_{}.log".format(
		self._instrument, self._exp_id, detector, CP["mbi"]["detector2filter"][detector])
		return os.path.join(self.dir_l0, fn)

		def l0_detector(self, detector=6):
		""" L0 detector-specific image file path """
		if self.ver_sim == "C3":
		fn = "{}_{}_{}_{}_{:02d}_raw.fits".format(
		self._instrument, self._survey, self._exp_start, self._exp_id, detector)
		else:
		assert self.ver_sim in ["C5.1", "C5.2"]
		fn = "{}_{}_{}_SCI_{}_{}_{}_{:02d}_L0_1.fits".format(
		self._telescope, self._instrument, self._survey,
		self._exp_start, self._exp_stop, self._exp_id, detector)
		return os.path.join(self.dir_l0, fn)

		def l0_crs(self, detector=6):
		""" L0 cosmic ray file path """
		if self.ver_sim == "C3":
		fn = "{}_CRS_{}_{}_{:02d}_raw.fits".format(
		self._instrument, self._exp_start, self._exp_id, detector)
		else:
		assert self.ver_sim in ["C5.1", "C5.2"]
		fn = "{}_{}_{}_CRS_{}_{}_{}_{:02d}_L0_1.fits".format(
		self._telescope, self._instrument, self._survey,
		self._exp_start, self._exp_stop, self._exp_id, detector)
		return os.path.join(self.dir_l0, fn)

		def l1_detector(self, detector=6, post="img.fits"):
		""" generate L1 file path

		Parameters
		----------
		detector:
		detector ID
		post:
		postfix
		{"img.fits", "wht.fits", "flg.fits", "img_L1.fits", "wht_L1.fits", "flg_L1.fits"}

		Returns
		-------
		L1 file path

		"""
		if self.ver_sim == "C3":
		fn = "{}_{}_{}_{}_{:02d}_{}".format(
		self._instrument, self._survey,
		self._exp_start, self._exp_id, detector, post)
		else:
		assert self.ver_sim in ["C5.1", "C5.2"]
		fn = "{}_{}_{}_SCI_{}_{}_{}_{:02d}_{}".format(
		self._telescope, self._instrument, self._survey,
		self._exp_start, self._exp_stop, self._exp_id, detector, post)
		return os.path.join(self.dir_l1, fn)

		def set_detectors(self, detectors=None):
		""" set target detector """
		if detectors is None:
		# default detectors
		self.target_detectors = self.available_detectors
		else:
		try:
		# assert detectors is a subset of available detectors
		assert set(detectors).issubset(set(self.available_detectors))
		self.target_detectors = list(detectors)
		except AssertionError as ae:
		print("@DM: available detector are ", self.available_detectors)
		print("@DM: target detector are ", detectors)

		print("@DM: final target detectors are ", set(detectors) & set(self.available_detectors))
		# raise ae
		self.target_detectors = set(detectors) & set(self.available_detectors)
		print("final target detector IDs are ", self.target_detectors)
		return

		def get_bias(self, detector=6):
		fp = glob.glob(self.path_aux.format("CLB", detector))[0]
		return fits.getdata(fp)

		def get_dark(self, detector=6):
		fp = glob.glob(self.path_aux.format("CLD", detector))[0]
		return fits.getdata(fp)

		def get_flat(self, detector=6):
		fp = glob.glob(self.path_aux.format("CLF", detector))[0]
		return fits.getdata(fp)

		def l1_file(self, name="", comment=""):
		"""

		Parameters
		----------
		name: str
		file name
		comment: str
		use the function name plz

		Returns
		-------
		fp: str
		the synthetic file path

		"""
		fp = os.path.join(self.dir_l1, name)
		# record hardcode history
		self.hardcode_history.append(dict(hdcd=fp, comment=comment))
		return fp

		@staticmethod
		def quickstart(ver_sim="C5.2", dir_l1=".", exposure_id=100):
		""" quick dataset generator for tests on dandelion or PMO

		Parameters
		----------
		ver_sim: str
		{"C5.2"}
		dir_l1: str
		output directory
		exposure_id: int
		21-154 for C5.2

		Returns
		-------
		CsstMbiDataManager
		the MBI data manager instance
		"""
		# auto identify node name
		hostname = os.uname()[1]

		if hostname == "dandelion":
		# dandelion node
		dir_l0 = "/nfsdata/share/csst_simulation_data/Cycle-5-SimuData/multipleBandsImaging/" \
		"NGP_AstrometryON_shearOFF/MSC_{:07d}/".format(exposure_id)
		path_aux = "/nfsdata/users/cham/L1Test/ref_C5.2/MSC_{}_*_{:02d}_combine.fits"
		dir_pcref = "/nfsdata/users/csstpipeline/L1Pipeline/msc/gaia_dr3/"
		elif hostname == "ubuntu":
		# PMO node
		dir_l0 = "/share/simudata/CSSOSDataProductsSims/data/CSSTSimImage_C5/" \
		"NGP_AstrometryON_shearOFF/MSC_{:07d}/".format(exposure_id)
		path_aux = "/data/sim_data/MSC_0000100/ref/MSC_{}_*_{:02d}_combine.fits"
		dir_pcref = "/home/user/L1Pipeline/msc/gaia_dr3/"
		else:
		raise ValueError("@DM: invalid hostname {}!".format(hostname))

		return CsstMbiDataManager(ver_sim=ver_sim, dir_l0=dir_l0, dir_l1=dir_l1, dir_pcref=dir_pcref, path_aux=path_aux)
		raise DeprecationWarning(
		"CsstMbiDataManager will no longer be available in some days, please use CsstMsDataManager instead.")

tests/test_data_manager.py

+0 −15

Original line number	Diff line number	Diff line
		@@ -14,21 +14,6 @@ from csst_common.data_manager import CsstMbiDataManager, CsstMsDataManager
		from csst_common.params import CSST_PARAMS as CP


		# deprecated in some days
		class TestCsstMbiDataManager(unittest.TestCase):
		def setUp(self) -> None:
		self.dm = CsstMbiDataManager.quickstart(ver_sim="C5.2", dir_l1=".", exposure_id=100)

		def test_l0data_existence(self):
		self.assertTrue(self.dm.available_detectors == CP["mbi"]["detectors"])
		self.assertTrue(os.path.exists(self.dm.l0_detector(6)))
		self.assertTrue(os.path.exists(self.dm.l0_log(6)))
		self.assertTrue(os.path.exists(self.dm.l0_cat(6)))
		self.assertTrue(os.path.exists(self.dm.l0_crs(6)))
		self.assertTrue(isinstance(self.dm.l1_detector(6, post="img.fits"), str))
		self.assertTrue(isinstance(self.dm.l1_file(name="some_file.ext", comment="a demo file"), str))


		class TestCsstMsDataManager(unittest.TestCase):
		def setUp(self) -> None:
		self.dm_mbi = CsstMsDataManager.quickstart(ver_sim="C5.2", datatype="mbi", dir_l1=".", exposure_id=100)