added CsstMsDataManager (f83bc722) · Commits · Wu Maochun / csst_common

csst_common/data_manager.py

+277 −58

Original line number	Diff line number	Diff line
		@@ -6,7 +6,7 @@ Author: Bo Zhang
		Created: 2022-09-13
		Modified-History:
		2022-09-13, Bo Zhang, created
		2022-09-13, Bo Zhang, fixed a bug
		2022-09-29, Bo Zhang, favor CsstMsDataManager instead of CsstMbiDataManager
		"""
		import os
		import glob
		@@ -17,8 +17,30 @@ from astropy.io import fits
		from .params import CSST_PARAMS as CP


		class CsstSlsDataManager:
		def __init__(self, ver_sim="C5.2", dir_l0="", dir_l1="", dir_pcref="", path_aux="", force_all_detectors=False):
		class CsstMsDataManager:
		""" CSST MS data manager, including MBI and SLS

		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,
		datatype="mbi"):
		""" initialize the multi-band imaging data manager

		Parameters
		@@ -35,31 +57,43 @@ class CsstSlsDataManager:
		aux data directory (bias, flat, dark)
		force_all_detectors: bool
		if True, assert data for all detectors are available
		datatype: str
		{"mbi", "sls"}

		Examples
		--------
		>>> dm = CsstSlsDataManager(...)
		>>> dm_mbi = CsstMbiDataManager(...)
		>>> # access L0 directory
		>>> dm.dir_l0
		>>> dm_mbi.dir_l0
		>>> # access L1 directory
		>>> dm.dir_l1
		>>> dm_mbi.dir_l1
		>>> # access dir_pcref
		>>> dm.dir_pcref
		>>> dm_mbi.dir_pcref
		>>> # access path_aux
		>>> dm.path_aux
		>>> dm_mbi.path_aux
		>>> # access ver_sim
		>>> dm.ver_sim
		>>> dm_mbi.ver_sim
		>>> # access target detectors
		>>> dm.target_detectors
		>>> dm_mbi.target_detectors
		>>> # access available detectors
		>>> dm.available_detectors
		>>> dm_mbi.available_detectors
		>>> # define an L1 file (detector-specified)
		>>> dm.l1_detector(detector=6)
		>>> dm_mbi.l1_detector(detector=6)
		>>> # define an L1 file (non-detector-specified)
		>>> dm.l1_file("flipped_image.fits")
		>>> dm_mbi.l1_file("flipped_image.fits")
		"""
		assert ver_sim in CP["sim"]["versions"]

		assert datatype in ["mbi", "sls"]
		if datatype == "mbi":
		# MBI
		self.valid_detectors = CP["mbi"]["detectors"]
		self.detector2filter = CP["mbi"]["detector2filter"]
		else:
		# SLS
		self.valid_detectors = CP["sls"]["detectors"]
		self.detector2filter = CP["sls"]["detector2filter"]

		self.dir_l0 = dir_l0
		self.dir_l1 = dir_l1
		self.dir_pcref = dir_pcref
		@@ -73,7 +107,7 @@ class CsstSlsDataManager:
		fps_cat = self.glob_cat(dir_l0, ver_sim=ver_sim)

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

		@@ -108,6 +142,217 @@ class CsstSlsDataManager:
		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, self.detector2filter[detector])
		elif 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":
		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,
		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
		e.g, {"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=""):
		""" L1 file path

		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", datatype="mbi", dir_l1=".", exposure_id=100):
		""" quick dataset generator for tests on dandelion or PMO

		Parameters
		----------
		ver_sim:
		{"C5.2"}
		datatype: str
		{"mbi", "sls"}
		dir_l1:
		output directory
		exposure_id:
		21-154 for C5.2

		Returns
		-------
		CsstMbiDataManager
		the MBI data manager instance
		"""
		assert datatype in ["mbi", "sls"]
		# auto identify node name
		hostname = os.uname()[1]
		assert hostname in ["dandelion", "ubuntu"]

		# dandelion
		if hostname == "dandelion" and datatype == "mbi":
		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 == "dandelion" and datatype == "sls":
		dir_l0 = "/nfsdata/share/csst_simulation_data/Cycle-5-SimuData/multipleBandsImaging/" \
		"NGP_AstrometryON_shearOFF_Spec/MSC_{:07d}/".format(exposure_id)
		path_aux = ""
		dir_pcref = ""

		# PMO
		elif hostname == "ubuntu" and datatype == "mbi":
		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/"
		elif hostname == "ubuntu" and datatype == "sls":
		dir_l0 = "/share/simudata/CSSOSDataProductsSims/data/CSSTSimImage_C5/" \
		"NGP_AstrometryON_shearOFF_Spec/MSC_{:07d}/".format(exposure_id)
		path_aux = ""
		dir_pcref = ""

		else:
		raise ValueError("@DM: invalid hostname {}!".format(hostname))

		return CsstMsDataManager(ver_sim=ver_sim, 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
		@@ -151,26 +396,28 @@ class CsstMbiDataManager:

		Examples
		--------
		>>> dm = CsstMbiDataManager(...)
		>>> dm_mbi = CsstMbiDataManager(...)
		>>> # access L0 directory
		>>> dm.dir_l0
		>>> dm_mbi.dir_l0
		>>> # access L1 directory
		>>> dm.dir_l1
		>>> dm_mbi.dir_l1
		>>> # access dir_pcref
		>>> dm.dir_pcref
		>>> dm_mbi.dir_pcref
		>>> # access path_aux
		>>> dm.path_aux
		>>> dm_mbi.path_aux
		>>> # access ver_sim
		>>> dm.ver_sim
		>>> dm_mbi.ver_sim
		>>> # access target detectors
		>>> dm.target_detectors
		>>> dm_mbi.target_detectors
		>>> # access available detectors
		>>> dm.available_detectors
		>>> dm_mbi.available_detectors
		>>> # define an L1 file (detector-specified)
		>>> dm.l1_detector(detector=6)
		>>> dm_mbi.l1_detector(detector=6)
		>>> # define an L1 file (non-detector-specified)
		>>> dm.l1_file("flipped_image.fits")
		>>> 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
		@@ -378,34 +625,6 @@ class CsstMbiDataManager:
		self.hardcode_history.append(dict(hdcd=fp, comment=comment))
		return fp

		@staticmethod
		def __testddl__():
		""" test function """
		dm = CsstMbiDataManager(
		ver_sim="C5.2",
		dir_l0="/nfsdata/share/csst_simulation_data/Cycle-5-SimuData/multipleBandsImaging"
		"/NGP_AstrometryON_shearOFF/MSC_0000100",
		dir_l1=".",
		force_all_detectors=True,
		)
		print("----- available detectors -----")
		print(dm.available_detectors)
		for detector in dm.available_detectors[:2]:
		print("----- L0 images -----")
		print(dm.l0_detector(detector=detector))
		print(os.path.exists(dm.l0_detector(detector=detector)))
		print("----- L0 crs -----")
		print(dm.l0_crs(detector=detector))
		print(os.path.exists(dm.l0_detector(detector=detector)))
		print("----- L0 input cat -----")
		print(dm.l0_cat(detector=detector))
		print(os.path.exists(dm.l0_cat(detector=detector)))
		print("----- L0 input log -----")
		print(dm.l0_log(detector=detector))
		print(os.path.exists(dm.l0_log(detector=detector)))
		print("----- L1 images -----")
		print(dm.l1_detector(detector, post="img.fits"))
		return dm

		@staticmethod
		def quickstart(ver_sim="C5.2", dir_l1=".", exposure_id=100):
		@@ -446,12 +665,12 @@ class CsstMbiDataManager:
		return CsstMbiDataManager(ver_sim=ver_sim, dir_l0=dir_l0, dir_l1=dir_l1, dir_pcref=dir_pcref, path_aux=path_aux)


		DIR_SIM = {
		("C5.2", "dandelion"): "/nfsdata/share/csst_simulation_data/Cycle-5-SimuData/multipleBandsImaging/"
		"NGP_AstrometryON_shearOFF/MSC_{:07d}/",
		("C5.2", "ubuntu"): "/nfsdata/share/csst_simulation_data/Cycle-5-SimuData/multipleBandsImaging/"
		"NGP_AstrometryON_shearOFF/MSC_{:07d}/",
		}
		# DIR_SIM = {
		# ("C5.2", "dandelion"): "/nfsdata/share/csst_simulation_data/Cycle-5-SimuData/multipleBandsImaging/"
		# "NGP_AstrometryON_shearOFF/MSC_{:07d}/",
		# ("C5.2", "ubuntu"): "/nfsdata/share/csst_simulation_data/Cycle-5-SimuData/multipleBandsImaging/"
		# "NGP_AstrometryON_shearOFF/MSC_{:07d}/",
		# }

		if __name__ == "__main__":
		# test C3

tests/test_data_manager.py

+28 −3

Original line number	Diff line number	Diff line
		@@ -10,16 +10,17 @@ Modified-History:
		"""
		import os
		import unittest
		from csst_common.data_manager import CsstMbiDataManager
		from csst_common.data_manager import CsstMbiDataManager, CsstMsDataManager
		from csst_common.params import CSST_PARAMS as CP


		class TestDataManager(unittest.TestCase):
		# 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(len(self.dm.available_detectors) == len(CP["mbi"]["detectors"]))
		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)))
		@@ -28,6 +29,30 @@ class TestDataManager(unittest.TestCase):
		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)
		self.dm_sls = CsstMsDataManager.quickstart(ver_sim="C5.2", datatype="sls", dir_l1=".", exposure_id=100)

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

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


		# class TestParams(unittest.TestCase):
		# def test_params(self):
		# self.assertTrue(True)