data_manager.py 26.3 KB
Newer Older
BO ZHANG's avatar
BO ZHANG committed
1
2
3
4
5
6
7
8
"""
Identifier:     KSC-SJ4-csst_common/data_manager.py
Name:           data_manager.py
Description:    file path generator
Author:         Bo Zhang
Created:        2022-09-13
Modified-History:
    2022-09-13, Bo Zhang, created
BO ZHANG's avatar
BO ZHANG committed
9
    2022-09-29, Bo Zhang, favor CsstMsDataManager instead of CsstMbiDataManager
BO ZHANG's avatar
BO ZHANG committed
10
"""
BO ZHANG's avatar
BO ZHANG committed
11
import os
BO ZHANG's avatar
BO ZHANG committed
12
13
14
15
16
import glob
import re

from astropy.io import fits

17
from .params import CSST_PARAMS as CP
BO ZHANG's avatar
BO ZHANG committed
18
19


BO ZHANG's avatar
BO ZHANG committed
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
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"):
BO ZHANG's avatar
BO ZHANG committed
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
        """ initialize the multi-band imaging data manager

        Parameters
        ----------
        ver_sim: str
            version of simulation data, see csst_common.params.CP
        dir_l0: str
            L0 directory
        dir_l1: str
            L1 directory
        dir_pcref: str
            position calibration data directory
        path_aux: str
            aux data directory (bias, flat, dark)
        force_all_detectors: bool
            if True, assert data for all detectors are available
BO ZHANG's avatar
BO ZHANG committed
60
61
        datatype: str
            {"mbi", "sls"}
BO ZHANG's avatar
BO ZHANG committed
62
63
64

        Examples
        --------
BO ZHANG's avatar
BO ZHANG committed
65
        >>> dm_mbi = CsstMbiDataManager(...)
BO ZHANG's avatar
BO ZHANG committed
66
        >>> # access L0 directory
BO ZHANG's avatar
BO ZHANG committed
67
        >>> dm_mbi.dir_l0
BO ZHANG's avatar
BO ZHANG committed
68
        >>> # access L1 directory
BO ZHANG's avatar
BO ZHANG committed
69
        >>> dm_mbi.dir_l1
BO ZHANG's avatar
BO ZHANG committed
70
        >>> # access dir_pcref
BO ZHANG's avatar
BO ZHANG committed
71
        >>> dm_mbi.dir_pcref
BO ZHANG's avatar
BO ZHANG committed
72
        >>> # access path_aux
BO ZHANG's avatar
BO ZHANG committed
73
        >>> dm_mbi.path_aux
BO ZHANG's avatar
BO ZHANG committed
74
        >>> # access ver_sim
BO ZHANG's avatar
BO ZHANG committed
75
        >>> dm_mbi.ver_sim
BO ZHANG's avatar
BO ZHANG committed
76
        >>> # access target detectors
BO ZHANG's avatar
BO ZHANG committed
77
        >>> dm_mbi.target_detectors
BO ZHANG's avatar
BO ZHANG committed
78
        >>> # access available detectors
BO ZHANG's avatar
BO ZHANG committed
79
        >>> dm_mbi.available_detectors
BO ZHANG's avatar
BO ZHANG committed
80
        >>> # define an L1 file (detector-specified)
BO ZHANG's avatar
BO ZHANG committed
81
        >>> dm_mbi.l1_detector(detector=6)
BO ZHANG's avatar
BO ZHANG committed
82
        >>> # define an L1 file (non-detector-specified)
BO ZHANG's avatar
BO ZHANG committed
83
        >>> dm_mbi.l1_file("flipped_image.fits")
BO ZHANG's avatar
BO ZHANG committed
84
85
86
        """
        assert ver_sim in CP["sim"]["versions"]

BO ZHANG's avatar
BO ZHANG committed
87
88
89
90
91
92
93
94
95
96
        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"]

BO ZHANG's avatar
BO ZHANG committed
97
98
99
100
101
102
103
104
105
106
107
108
109
        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:
BO ZHANG's avatar
BO ZHANG committed
110
            assert len(fps_img) == len(self.valid_detectors)
BO ZHANG's avatar
BO ZHANG committed
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
        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()

BO ZHANG's avatar
BO ZHANG committed
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
    @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)

BO ZHANG's avatar
BO ZHANG committed
356

BO ZHANG's avatar
BO ZHANG committed
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
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

    """

BO ZHANG's avatar
BO ZHANG committed
379
    def __init__(self, ver_sim="C5.2", dir_l0="", dir_l1="", dir_pcref="", path_aux="", force_all_detectors=False):
BO ZHANG's avatar
BO ZHANG committed
380
381
382
383
384
        """ initialize the multi-band imaging data manager

        Parameters
        ----------
        ver_sim: str
BO ZHANG's avatar
BO ZHANG committed
385
            version of simulation data, see csst_common.params.CP
BO ZHANG's avatar
BO ZHANG committed
386
387
388
389
390
391
392
393
394
395
396
397
398
        dir_l0: str
            L0 directory
        dir_l1: str
            L1 directory
        dir_pcref: str
            position calibration data directory
        path_aux: str
            aux data directory (bias, flat, dark)
        force_all_detectors: bool
            if True, assert data for all detectors are available

        Examples
        --------
BO ZHANG's avatar
BO ZHANG committed
399
        >>> dm_mbi = CsstMbiDataManager(...)
BO ZHANG's avatar
BO ZHANG committed
400
        >>> # access L0 directory
BO ZHANG's avatar
BO ZHANG committed
401
        >>> dm_mbi.dir_l0
BO ZHANG's avatar
BO ZHANG committed
402
        >>> # access L1 directory
BO ZHANG's avatar
BO ZHANG committed
403
        >>> dm_mbi.dir_l1
BO ZHANG's avatar
BO ZHANG committed
404
        >>> # access dir_pcref
BO ZHANG's avatar
BO ZHANG committed
405
        >>> dm_mbi.dir_pcref
BO ZHANG's avatar
BO ZHANG committed
406
        >>> # access path_aux
BO ZHANG's avatar
BO ZHANG committed
407
        >>> dm_mbi.path_aux
BO ZHANG's avatar
BO ZHANG committed
408
        >>> # access ver_sim
BO ZHANG's avatar
BO ZHANG committed
409
        >>> dm_mbi.ver_sim
BO ZHANG's avatar
BO ZHANG committed
410
        >>> # access target detectors
BO ZHANG's avatar
BO ZHANG committed
411
        >>> dm_mbi.target_detectors
BO ZHANG's avatar
BO ZHANG committed
412
        >>> # access available detectors
BO ZHANG's avatar
BO ZHANG committed
413
        >>> dm_mbi.available_detectors
BO ZHANG's avatar
BO ZHANG committed
414
        >>> # define an L1 file (detector-specified)
BO ZHANG's avatar
BO ZHANG committed
415
        >>> dm_mbi.l1_detector(detector=6)
BO ZHANG's avatar
BO ZHANG committed
416
        >>> # define an L1 file (non-detector-specified)
BO ZHANG's avatar
BO ZHANG committed
417
        >>> dm_mbi.l1_file("flipped_image.fits")
BO ZHANG's avatar
BO ZHANG committed
418
        """
BO ZHANG's avatar
BO ZHANG committed
419
420
        # raise DeprecationWarning(
        #     "CsstMbiDataManager will no longer be available in some days, please use CsstMsDataManager instead.")
421
        assert ver_sim in CP["sim"]["versions"]
BO ZHANG's avatar
BO ZHANG committed
422
423
424
425
426
427
428
429
430
431
432
433
434
435

        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:
436
            assert len(fps_img) == len(CP["mbi"]["detectors"])
BO ZHANG's avatar
BO ZHANG committed
437
438
439
440
441
442
443
        else:
            assert len(fps_img) > 0

        if ver_sim == "C3":
            # get info
            # print(re.split(r"[_.]", fps[0]))
            self._instrument, self._survey, \
BO ZHANG's avatar
BO ZHANG committed
444
445
                self._exp_start, self._exp_id, \
                _detector, self._l0_suffix, _ext = re.split(r"[_.]", fps_img[0])
BO ZHANG's avatar
BO ZHANG committed
446
447
448
449
450
451
452
453
454
455
456
457
458
            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, \
BO ZHANG's avatar
BO ZHANG committed
459
460
                self._exp_start, self._exp_stop, self._exp_id, \
                _detector, self._l0_suffix, self._version, _ext = re.split(r"[_.]", fps_img[0])
BO ZHANG's avatar
BO ZHANG committed
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
            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(
BO ZHANG's avatar
BO ZHANG committed
508
                self._instrument, self._exp_id - 90000000, detector, CP["mbi"]["detector2filter"][detector])
BO ZHANG's avatar
BO ZHANG committed
509
510
        elif self.ver_sim == "C5.2":
            fn = "{}_{}_chip_{:02d}_filt_{}.cat".format(
BO ZHANG's avatar
BO ZHANG committed
511
                self._instrument, self._exp_id, detector, CP["mbi"]["detector2filter"][detector])
BO ZHANG's avatar
BO ZHANG committed
512
513
514
515
516
517
        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(
BO ZHANG's avatar
BO ZHANG committed
518
                self._instrument, self._exp_id - 90000000, detector, CP["mbi"]["detector2filter"][detector])
BO ZHANG's avatar
BO ZHANG committed
519
520
        elif self.ver_sim == "C5.2":
            fn = "{}_{}_chip_{:02d}_filt_{}.log".format(
BO ZHANG's avatar
BO ZHANG committed
521
                self._instrument, self._exp_id, detector, CP["mbi"]["detector2filter"][detector])
BO ZHANG's avatar
BO ZHANG committed
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
        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))
BO ZHANG's avatar
BO ZHANG committed
584
                self.target_detectors = list(detectors)
BO ZHANG's avatar
BO ZHANG committed
585
            except AssertionError as ae:
BO ZHANG's avatar
BO ZHANG committed
586
587
588
589
                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))
BO ZHANG's avatar
BO ZHANG committed
590
                # raise ae
BO ZHANG's avatar
BO ZHANG committed
591
                self.target_detectors = set(detectors) & set(self.available_detectors)
BO ZHANG's avatar
BO ZHANG committed
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
        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


BO ZHANG's avatar
BO ZHANG committed
629
    @staticmethod
630
    def quickstart(ver_sim="C5.2", dir_l1=".", exposure_id=100):
BO ZHANG's avatar
BO ZHANG committed
631
        """ quick dataset generator for tests on dandelion or PMO
BO ZHANG's avatar
BO ZHANG committed
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646

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

        Returns
        -------
        CsstMbiDataManager
            the MBI data manager instance
        """
BO ZHANG's avatar
BO ZHANG committed
647
648
649
        # auto identify node name
        hostname = os.uname()[1]

BO ZHANG's avatar
BO ZHANG committed
650
651
        if hostname == "dandelion":
            # dandelion node
BO ZHANG's avatar
BO ZHANG committed
652
653
            dir_l0 = "/nfsdata/share/csst_simulation_data/Cycle-5-SimuData/multipleBandsImaging/" \
                     "NGP_AstrometryON_shearOFF/MSC_{:07d}/".format(exposure_id)
BO ZHANG's avatar
BO ZHANG committed
654
655
            path_aux = "/nfsdata/users/cham/L1Test/ref_C5.2/MSC_{}_*_{:02d}_combine.fits"
            dir_pcref = "/nfsdata/users/csstpipeline/L1Pipeline/msc/gaia_dr3/"
BO ZHANG's avatar
BO ZHANG committed
656
657
        elif hostname == "ubuntu":
            # PMO node
BO ZHANG's avatar
BO ZHANG committed
658
659
            dir_l0 = "/share/simudata/CSSOSDataProductsSims/data/CSSTSimImage_C5/" \
                     "NGP_AstrometryON_shearOFF/MSC_{:07d}/".format(exposure_id)
BO ZHANG's avatar
BO ZHANG committed
660
661
662
            path_aux = "/data/sim_data/MSC_0000100/ref/MSC_{}_*_{:02d}_combine.fits"
            dir_pcref = "/home/user/L1Pipeline/msc/gaia_dr3/"
        else:
BO ZHANG's avatar
BO ZHANG committed
663
            raise ValueError("@DM: invalid hostname {}!".format(hostname))
BO ZHANG's avatar
BO ZHANG committed
664
665
666

        return CsstMbiDataManager(ver_sim=ver_sim, dir_l0=dir_l0, dir_l1=dir_l1, dir_pcref=dir_pcref, path_aux=path_aux)

BO ZHANG's avatar
BO ZHANG committed
667

BO ZHANG's avatar
BO ZHANG committed
668
669
670
671
672
673
# 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}/",
# }
BO ZHANG's avatar
BO ZHANG committed
674

BO ZHANG's avatar
BO ZHANG committed
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
if __name__ == "__main__":
    # test C3
    dm = CsstMbiDataManager(
        ver_sim="C3", dir_l0="/data/L1Pipeline/msc/MSC_0000020", dir_l1="/data/L1Pipeline/msc/work")
    print("----- L0 images -----")
    print(dm.l0_detector(detector=6))
    print(os.path.exists(dm.l0_detector(detector=6)))
    print("----- L0 crs -----")
    print(dm.l0_crs(detector=6))
    print(os.path.exists(dm.l0_detector(detector=8)))
    print("----- L0 input cat -----")
    print(dm.l0_cat(8))
    print(os.path.exists(dm.l0_cat(detector=8)))
    print("----- available detectors -----")
    print(dm.available_detectors)
    print("----- L1 images -----")
BO ZHANG's avatar
BO ZHANG committed
691
    print(dm.l1_detector(25, "img.fits"))
BO ZHANG's avatar
BO ZHANG committed
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712

    # test C5.1
    dm = CsstMbiDataManager(
        ver_sim="C5.1", dir_l0="/data/sim_data/MSC_0000100", dir_l1="/home/user/L1Pipeline/msc/work")
    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"))