data_manager.py 29.8 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
tweaks    
BO ZHANG committed
9
    2022-09-13, Bo Zhang, added CsstMbiDataManager
BO ZHANG's avatar
BO ZHANG committed
10
    2022-09-29, Bo Zhang, favor CsstMsDataManager instead of CsstMbiDataManager
BO ZHANG's avatar
BO ZHANG committed
11
    2022-10-26, Bo Zhang, reconstruct CsstMsDataManager, deprecate CsstMbiDataManager
BO ZHANG's avatar
BO ZHANG committed
12
    2022-10-28, Bo Zhang, added CsstMsDataManager.query_rc(), dm.use_dfs, dm.node
BO ZHANG's avatar
BO ZHANG committed
13
    2022-11-06, Bo Zhang, deleted CsstMbiDataManager
BO ZHANG's avatar
BO ZHANG committed
14
    2022-11-20, Bo Zhang, added DFS APIs
BO ZHANG's avatar
BO ZHANG committed
15
"""
BO ZHANG's avatar
BO ZHANG committed
16
import _io
BO ZHANG's avatar
BO ZHANG committed
17
import glob
BO ZHANG's avatar
BO ZHANG committed
18
import os
BO ZHANG's avatar
BO ZHANG committed
19
import re
BO ZHANG's avatar
BO ZHANG committed
20
from typing import Union
BO ZHANG's avatar
BO ZHANG committed
21

BO ZHANG's avatar
BO ZHANG committed
22
import numpy as np
BO ZHANG's avatar
BO ZHANG committed
23
from astropy.io import fits
BO ZHANG's avatar
BO ZHANG committed
24
from astropy import time
BO ZHANG's avatar
BO ZHANG committed
25
from astropy.table import Table
BO ZHANG's avatar
BO ZHANG committed
26
27
from csst_dfs_api.facility.level0 import Level0DataApi
from csst_dfs_api.facility.level0prc import Level0PrcApi
BO ZHANG's avatar
BO ZHANG committed
28
29
from csst_dfs_api.facility.level1 import Level1DataApi
from csst_dfs_api.facility.level1prc import Level1PrcApi
BO ZHANG's avatar
BO ZHANG committed
30
from csst_dfs_api.mbi.level2 import Level2DataApi as MbiLevel2DataApi
BO ZHANG's avatar
BO ZHANG committed
31
from csst_dfs_api.sls.level2spectra import Level2SpectraApi as SlsLevel2DataApi
BO ZHANG's avatar
BO ZHANG committed
32
from csst_dfs_api.common.catalog import CatalogApi
BO ZHANG's avatar
BO ZHANG committed
33

BO ZHANG's avatar
BO ZHANG committed
34
from .logger import get_logger
35
from .params import CSST_PARAMS as CP
BO ZHANG's avatar
BO ZHANG committed
36
from .params import DFS_CONF
BO ZHANG's avatar
BO ZHANG committed
37
38


BO ZHANG's avatar
BO ZHANG committed
39
class CsstMsDataManager:
BO ZHANG's avatar
BO ZHANG committed
40
41
    """
    CSST MS data manager, including MBI and SLS.
BO ZHANG's avatar
tweaks    
BO ZHANG committed
42

BO ZHANG's avatar
tweaks    
BO ZHANG committed
43
44
45
46
    ``CsstMsDataManager`` provides an interface to switch between DFS and local file system.
    To initialize ``CsstMsDataManager`` from local directory, use ``CsstMsDataManager.from_dir()``
    To initialize ``CsstMsDataManager`` on ``dandelion`` or ``PM node``, ``CsstMsDataManager.quickstart()``.
    To initialize ``CsstMsDataManager`` from DFS, use ``CsstMsDataManager.from_dfs()``.
BO ZHANG's avatar
tweaks    
BO ZHANG committed
47
    To generate L0 and L1 file paths, use ``CsstMsDataManager.l0_detector()``, ``CsstMsDataManager.l1_detector()``, etc.
BO ZHANG's avatar
BO ZHANG committed
48
    Here are some examples for simulation with different versions.
BO ZHANG's avatar
BO ZHANG committed
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
    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
64
65
66
    Parameters
    ----------
    ver_sim : str
BO ZHANG's avatar
BO ZHANG committed
67
68
69
70
71
72
73
74
75
        The version of simulation data, see ``csst_common.params.CP``.
    datatype : str
        The options are {"mbi", "sls", "all"}.
        The "all" option is used for QC in particular.
        Note that in this case methods like ``get_bias`` are unavailable.
    available_detectors : list
        The list of available detector serial numbers of available images.
    target_detectors : list
        The list of target detector serial numbers of available images.
BO ZHANG's avatar
BO ZHANG committed
76
77
78
79
    dir_l0 : str
        The L0 directory.
    dir_l1 : str
        The L1 directory.
BO ZHANG's avatar
BO ZHANG committed
80
81
    path_aux : str
        The aux data directory (bias, flat, dark).
BO ZHANG's avatar
BO ZHANG committed
82
83
    use_dfs : bool
        If True, use DFS. In case some modules may have other options such as astroquery.
BO ZHANG's avatar
BO ZHANG committed
84
    dfs_node : str
BO ZHANG's avatar
BO ZHANG committed
85
86
        The environment in which the pipeline will run.
        Use "pml" for Purple Mountain Lab cluster and "local" for others.
BO ZHANG's avatar
BO ZHANG committed
87
88
    dfs_root : str
        The DFS root path.
BO ZHANG's avatar
BO ZHANG committed
89
    obs_id : str
BO ZHANG's avatar
BO ZHANG committed
90
        The exposure ID.
BO ZHANG's avatar
BO ZHANG committed
91
    exp_start : int
BO ZHANG's avatar
BO ZHANG committed
92
        The exposure start time in ``yyyymmddhhmmss`` format.
BO ZHANG's avatar
BO ZHANG committed
93
    exp_stop : int
BO ZHANG's avatar
BO ZHANG committed
94
95
96
97
98
99
100
        The exposure start time in ``yyyymmddhhmmss`` format.
    _telescope : str
        The telescope name. Defaults to ``CSST`` for C5.2 simulation.
    _instrument : str
        The instrument name. Defaults to ``MSC`` for C5.2 simulation.
    _survey : str
        The survey name. Defaults to ``MS`` for C5.2 simulation.
BO ZHANG's avatar
BO ZHANG committed
101
    obs_type : str
BO ZHANG's avatar
BO ZHANG committed
102
        The image type signature for science images. Defualts to ``SCI`` for C5.2 simulation.
BO ZHANG's avatar
BO ZHANG committed
103
    l0_post : str
BO ZHANG's avatar
BO ZHANG committed
104
        The postfix. Defaults to ``L0_1`` for C5.2 simulation.
BO ZHANG's avatar
BO ZHANG committed
105
106
107
108
    log_ppl : str
        The pipeline log file name.
    log_mod : str
        The module log file name.
109
110
    clear_dir : bool
        If True, clear ``dm.dir_l1`` directory.
BO ZHANG's avatar
BO ZHANG committed
111
112
    verbose : bool
        If True, print verbose info.
BO ZHANG's avatar
BO ZHANG committed
113
114
115
116
    n_jobs : int
        The number of jobs.
    backend : str
        The joblib backend.
BO ZHANG's avatar
BO ZHANG committed
117
118
    device : str
        The device for neural network. "CPU" or "GPU".
BO ZHANG's avatar
BO ZHANG committed
119
120
121

    Examples
    --------
BO ZHANG's avatar
BO ZHANG committed
122
    >>> dm_mbi = CsstMsDataManager(...)
BO ZHANG's avatar
BO ZHANG committed
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
    >>> # access L0 directory
    >>> dm_mbi.dir_l0
    >>> # access L1 directory
    >>> dm_mbi.dir_l1
    >>> # access path_aux
    >>> dm_mbi.path_aux
    >>> # access ver_sim
    >>> dm_mbi.ver_sim
    >>> # access target detectors
    >>> dm_mbi.target_detectors
    >>> # access available detectors
    >>> dm_mbi.available_detectors
    >>> # define an L1 file (detector-specified)
    >>> dm_mbi.l1_detector(detector=6)
    >>> # define an L1 file (non-detector-specified)
    >>> dm_mbi.l1_file("flipped_image.fits")
BO ZHANG's avatar
BO ZHANG committed
139
140
    """

BO ZHANG's avatar
BO ZHANG committed
141
    def __init__(self,
BO ZHANG's avatar
tweaks    
BO ZHANG committed
142
143
144
                 ver_sim: str = "C5.2",
                 datatype: str = "mbi",
                 available_detectors: Union[None, list] = None,
BO ZHANG's avatar
BO ZHANG committed
145
                 target_detectors: Union[None, list, int] = None,
BO ZHANG's avatar
tweaks    
BO ZHANG committed
146
147
                 dir_l0: str = ".",
                 dir_l1: str = ".",
BO ZHANG's avatar
BO ZHANG committed
148
                 path_aux: str = "/L1Pipeline/aux",  # aux dir
BO ZHANG's avatar
BO ZHANG committed
149
                 use_dfs: bool = False,
BO ZHANG's avatar
BO ZHANG committed
150
                 dfs_node: str = "kmust",
BO ZHANG's avatar
BO ZHANG committed
151
                 dfs_root: str = "/share/dfs",
BO ZHANG's avatar
BO ZHANG committed
152
153
154
                 obs_id: str = "100000100",
                 exp_start: int = "20270810081950",
                 exp_stop: int = "20270810082220",
BO ZHANG's avatar
tweaks    
BO ZHANG committed
155
156
157
                 _telescope: str = "CSST",
                 _instrument: str = "MSC",
                 _survey: str = "MS",
BO ZHANG's avatar
BO ZHANG committed
158
159
                 obs_type: str = "SCI",
                 l0_post: str = "L0_1",
BO ZHANG's avatar
BO ZHANG committed
160
161
162
163
164
165
166
167
                 log_ppl: str = "",
                 log_mod: str = "",
                 clear_dir: bool = False,
                 verbose: bool = True,
                 n_jobs: int = 18,
                 backend: str = "multiprocessing",
                 device: str = "CPU",
                 stamps: str = "",
BO ZHANG's avatar
BO ZHANG committed
168
169
                 ):

BO ZHANG's avatar
BO ZHANG committed
170
171
172
        # set DFS log dir
        os.environ["CSST_DFS_LOGS_DIR"] = "."

BO ZHANG's avatar
BO ZHANG committed
173
        # version
BO ZHANG's avatar
BO ZHANG committed
174
        assert ver_sim in CP["sim"]["versions"]
BO ZHANG's avatar
tweaks    
BO ZHANG committed
175
        self.ver_sim = ver_sim
BO ZHANG's avatar
BO ZHANG committed
176

BO ZHANG's avatar
BO ZHANG committed
177
        # datatype, valid_detectors, detector2filter
178
        assert datatype in ["mbi", "sls", "all"]
BO ZHANG's avatar
tweaks    
BO ZHANG committed
179
        self.datatype = datatype
BO ZHANG's avatar
BO ZHANG committed
180
181
182
183
        if datatype == "mbi":
            # MBI
            self.valid_detectors = CP["mbi"]["detectors"]
            self.detector2filter = CP["mbi"]["detector2filter"]
184
        elif datatype == "sls":
BO ZHANG's avatar
BO ZHANG committed
185
186
187
            # SLS
            self.valid_detectors = CP["sls"]["detectors"]
            self.detector2filter = CP["sls"]["detector2filter"]
188
189
        else:
            # ALL
BO ZHANG's avatar
BO ZHANG committed
190
191
            self.valid_detectors = CP["all"]["detectors"]
            self.detector2filter = CP["all"]["detector2filter"]
BO ZHANG's avatar
BO ZHANG committed
192
193
194
        if verbose:
            print("Data type is: ", self.datatype)
            print("Valid detectors are: ", self.valid_detectors)
BO ZHANG's avatar
BO ZHANG committed
195
196

        # available_detectors
197
        self.available_detectors = available_detectors if available_detectors is not None else list()
BO ZHANG's avatar
BO ZHANG committed
198
199
        if verbose:
            print("Available detectors are:", self.available_detectors)
BO ZHANG's avatar
BO ZHANG committed
200
        # set all available detectors by default
201
        self.target_detectors = target_detectors
BO ZHANG's avatar
BO ZHANG committed
202
203
        if verbose:
            print("Target detectors are: ", self._target_detectors)
BO ZHANG's avatar
BO ZHANG committed
204
205

        # exposure info
BO ZHANG's avatar
BO ZHANG committed
206
207
208
        self.obs_id = obs_id
        self.exp_start = exp_start
        self.exp_stop = exp_stop
BO ZHANG's avatar
BO ZHANG committed
209

BO ZHANG's avatar
BO ZHANG committed
210
211
212
213
        # file name components
        self._telescope = _telescope
        self._instrument = _instrument
        self._survey = _survey
BO ZHANG's avatar
BO ZHANG committed
214
215
        self.obs_type = obs_type
        self.l0_post = l0_post
BO ZHANG's avatar
BO ZHANG committed
216

BO ZHANG's avatar
BO ZHANG committed
217
218
        # DFS configuration
        self.use_dfs = use_dfs
219
        self.dfs_node = dfs_node
BO ZHANG's avatar
BO ZHANG committed
220
        self.dfs_root = dfs_root
BO ZHANG's avatar
BO ZHANG committed
221
222

        # data directory
BO ZHANG's avatar
BO ZHANG committed
223
224
225
        self.dir_l0 = dir_l0
        self.dir_l1 = dir_l1
        self.path_aux = path_aux
BO ZHANG's avatar
BO ZHANG committed
226
        self.ref_version = None
BO ZHANG's avatar
BO ZHANG committed
227

BO ZHANG's avatar
BO ZHANG committed
228
        # record hard code names in history
BO ZHANG's avatar
BO ZHANG committed
229
230
        self.hardcode_history = []

BO ZHANG's avatar
BO ZHANG committed
231
232
        self.n_jobs = n_jobs
        self.backend = backend
BO ZHANG's avatar
BO ZHANG committed
233
        self.device = device
BO ZHANG's avatar
BO ZHANG committed
234

BO ZHANG's avatar
BO ZHANG committed
235
        # aXe
BO ZHANG's avatar
BO ZHANG committed
236
        self.set_env()
BO ZHANG's avatar
BO ZHANG committed
237

BO ZHANG's avatar
BO ZHANG committed
238
239
        # change to working directory
        os.chdir(self.dir_l1)
240
241
242
        # clear dir_l1
        if clear_dir:
            self.clear_dir(self.dir_l1)
BO ZHANG's avatar
BO ZHANG committed
243

BO ZHANG's avatar
BO ZHANG committed
244
        # pipeline logger
BO ZHANG's avatar
BO ZHANG committed
245
246
247
        if log_ppl == "":
            self.logger_ppl = get_logger(name="CSST L1 Pipeline Logger", filename="")
        else:
BO ZHANG's avatar
BO ZHANG committed
248
            self.logger_ppl = get_logger(name="CSST L1 Pipeline Logger", filename=os.path.join(dir_l1, log_ppl))
BO ZHANG's avatar
BO ZHANG committed
249
        # module logger
BO ZHANG's avatar
BO ZHANG committed
250
251
252
        if log_mod == "":
            self.logger_mod = get_logger(name="CSST L1 Module Logger", filename="")
        else:
BO ZHANG's avatar
BO ZHANG committed
253
            self.logger_mod = get_logger(name="CSST L1 Module Logger", filename=os.path.join(dir_l1, log_mod))
BO ZHANG's avatar
BO ZHANG committed
254

255
256
257
258
        self.custom_bias = None
        self.custom_dark = None
        self.custom_flat = None

BO ZHANG's avatar
BO ZHANG committed
259
260
        self.stamps = stamps
        self.write_stamp()
BO ZHANG's avatar
BO ZHANG committed
261
262

    def write_stamp(self):
BO ZHANG's avatar
BO ZHANG committed
263
264
265
        if self.stamps is not None and not self.stamps == "":
            with open(self.stamps, "a+") as f:
                f.write(f"{time.Time.now().isot}\n")
BO ZHANG's avatar
BO ZHANG committed
266

267
268
269
270
271
272
    # DFS APIs
    @property
    def dfs_L0DataApi(self):
        return Level0DataApi()

    @property
BO ZHANG's avatar
BO ZHANG committed
273
274
    def dfs_L0PrcApi(self):
        return Level0PrcApi()
275
276

    @property
BO ZHANG's avatar
BO ZHANG committed
277
278
    def dfs_L1DataApi(self):
        return Level1DataApi()
279
280

    @property
BO ZHANG's avatar
BO ZHANG committed
281
282
    def dfs_L1PrcApi(self):
        return Level1PrcApi()
BO ZHANG's avatar
BO ZHANG committed
283
284
285
286

    @property
    def dfs_MbiL2DataApi(self):
        return MbiLevel2DataApi()
287

BO ZHANG's avatar
BO ZHANG committed
288
289
290
291
    @property
    def dfs_SlsL2DataApi(self):
        return SlsLevel2DataApi()

BO ZHANG's avatar
BO ZHANG committed
292
293
294
295
    @property
    def dfs_CatApi(self):
        return CatalogApi()

BO ZHANG's avatar
BO ZHANG committed
296
297
298
    def set_env(self):
        """ set environment variables """
        if os.uname()[1] == "dandelion":
BO ZHANG's avatar
BO ZHANG committed
299
300
            os.environ["LD_LIBRARY_PATH"] = "/home/csstpipeline/anaconda3/lib"
            os.environ["AXE_BINDIR"] = "/home/csstpipeline/PycharmProjects/axe/cextern/src"
BO ZHANG's avatar
BO ZHANG committed
301
302
        else:
            os.environ["LD_LIBRARY_PATH"] = ""
BO ZHANG's avatar
BO ZHANG committed
303
            os.environ["AXE_BINDIR"] = ""
BO ZHANG's avatar
BO ZHANG committed
304

305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
    @property
    def target_detectors(self):
        return self._target_detectors

    @target_detectors.setter
    def target_detectors(self, detectors: Union[None, list, int] = None):
        assert detectors is None or type(detectors) in [list, int]
        if detectors is None:
            self._target_detectors = list(set(self.available_detectors) & set(self.valid_detectors))
        elif isinstance(detectors, list):
            self._target_detectors = list(set(self.available_detectors) & set(self.valid_detectors) & set(detectors))
        elif isinstance(detectors, int):
            self._target_detectors = list(set(self.available_detectors) & set(self.valid_detectors) & {detectors})

    def set_detectors(self, detectors=None):
        raise DeprecationWarning("This method is deprecated, please directly use dm.target_detectors = detectors!")

BO ZHANG's avatar
BO ZHANG committed
322
    @staticmethod
BO ZHANG's avatar
BO ZHANG committed
323
324
325
326
327
328
329
330
    def from_dir(
            ver_sim="C5.2",
            datatype="mbi",
            dir_l0=".",
            dir_l1=".",
            path_aux="",
            use_dfs=False,
            dfs_node="kmust",
BO ZHANG's avatar
BO ZHANG committed
331
332
            log_ppl="csst-l1ppl.log",
            log_mod="csst-l1mod.log",
BO ZHANG's avatar
BO ZHANG committed
333
            n_jobs=18,
BO ZHANG's avatar
BO ZHANG committed
334
            backend="multiprocessing",
BO ZHANG's avatar
BO ZHANG committed
335
336
            device="CPU",
            **kwargs
BO ZHANG's avatar
BO ZHANG committed
337
    ):
BO ZHANG's avatar
BO ZHANG committed
338
        """ initialize the multi-band imaging data manager """
BO ZHANG's avatar
BO ZHANG committed
339

BO ZHANG's avatar
BO ZHANG committed
340
        assert ver_sim in CP["sim"]["versions"]
BO ZHANG's avatar
BO ZHANG committed
341
342
343
344
345
346
347
348
349
350
351

        # 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
        available_detectors = [int(re.split(r"[_.]", fp)[7]) for fp in fps_img]
        available_detectors.sort()

        # parse info
BO ZHANG's avatar
BO ZHANG committed
352
353
354
        (_telescope, _instrument, _survey, obs_type,
         exp_start, exp_stop, obs_id,
         _detector, *l0_post, _ext) = re.split(r"[_.]", fps_img[0])
BO ZHANG's avatar
BO ZHANG committed
355

BO ZHANG's avatar
BO ZHANG committed
356
357
358
        # exp_start = int(exp_start)
        # exp_stop = int(exp_stop)
        # obs_id = int(obs_id)
BO ZHANG's avatar
BO ZHANG committed
359

BO ZHANG's avatar
BO ZHANG committed
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
        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
            use_dfs=use_dfs,
            dfs_node=dfs_node,
            obs_id=obs_id,
            exp_start=exp_start,
            exp_stop=exp_stop,
            _telescope=_telescope,
            _instrument=_instrument,
            _survey=_survey,
            obs_type=obs_type,
            l0_post="_".join(l0_post),
BO ZHANG's avatar
BO ZHANG committed
378
379
            log_ppl=log_ppl,
            log_mod=log_mod,
BO ZHANG's avatar
BO ZHANG committed
380
            n_jobs=n_jobs,
BO ZHANG's avatar
BO ZHANG committed
381
            backend=backend,
BO ZHANG's avatar
BO ZHANG committed
382
383
            device=device,
            **kwargs
BO ZHANG's avatar
BO ZHANG committed
384
        )
BO ZHANG's avatar
BO ZHANG committed
385

BO ZHANG's avatar
BO ZHANG committed
386
    @staticmethod
BO ZHANG's avatar
BO ZHANG committed
387
    def glob_image(dir_l0, ver_sim="C5.2"):
BO ZHANG's avatar
BO ZHANG committed
388
        """ glob files in L0 data directory """
BO ZHANG's avatar
BO ZHANG committed
389
390
        assert ver_sim in CP["sim"]["versions"]
        pattern = os.path.join(dir_l0, "CSST_MSC_MS_SCI_*.fits")
BO ZHANG's avatar
BO ZHANG committed
391
392
393
394
        fps = glob.glob(pattern)
        fps = [os.path.basename(fp) for fp in fps]
        fps.sort()

395
        print("{} files found with pattern: {}".format(len(fps), pattern))
BO ZHANG's avatar
BO ZHANG committed
396
397
398
399
400
        return fps

    @staticmethod
    def glob_cat(dir_l0, ver_sim="C5"):
        """ glob input catalogs in L0 data directory """
BO ZHANG's avatar
BO ZHANG committed
401
        assert ver_sim in CP["sim"]["versions"]
BO ZHANG's avatar
BO ZHANG committed
402
        pattern = os.path.join(dir_l0, "MSC_*.cat")
BO ZHANG's avatar
BO ZHANG committed
403
404
405
406
407
408
409
        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

BO ZHANG's avatar
BO ZHANG committed
410
411
412
413
    def l0_id(self, detector=6):
        """ Level0 ID, consistent with DFS. """
        return f"{self.obs_id}{detector:02d}"

BO ZHANG's avatar
BO ZHANG committed
414
415
    def l0_cat(self, detector=6):
        """ the L0 cat file path"""
BO ZHANG's avatar
BO ZHANG committed
416
        assert self.ver_sim in CP["sim"]["versions"]
BO ZHANG's avatar
BO ZHANG committed
417
        fn = "{}_{}_chip_{:02d}_filt_{}.cat".format(
BO ZHANG's avatar
BO ZHANG committed
418
            self._instrument, self.obs_id, detector, self.detector2filter[detector])
BO ZHANG's avatar
BO ZHANG committed
419
420
421
422
        return os.path.join(self.dir_l0, fn)

    def l0_log(self, detector=6):
        """ L0 log file path """
BO ZHANG's avatar
BO ZHANG committed
423
        assert self.ver_sim in CP["sim"]["versions"]
BO ZHANG's avatar
BO ZHANG committed
424
        fn = "{}_{}_chip_{:02d}_filt_{}.log".format(
BO ZHANG's avatar
BO ZHANG committed
425
            self._instrument, self.obs_id, detector, self.detector2filter[detector])
BO ZHANG's avatar
BO ZHANG committed
426
427
428
429
        return os.path.join(self.dir_l0, fn)

    def l0_detector(self, detector=6):
        """ L0 detector-specific image file path """
BO ZHANG's avatar
BO ZHANG committed
430
        assert self.ver_sim in CP["sim"]["versions"]
BO ZHANG's avatar
BO ZHANG committed
431
432
        fn = "{}_{}_{}_SCI_{}_{}_{}_{:02d}_L0_1.fits".format(
            self._telescope, self._instrument, self._survey,
BO ZHANG's avatar
BO ZHANG committed
433
            self.exp_start, self.exp_stop, self.obs_id, detector)
BO ZHANG's avatar
BO ZHANG committed
434
435
436
437
        return os.path.join(self.dir_l0, fn)

    def l0_crs(self, detector=6):
        """ L0 cosmic ray file path """
BO ZHANG's avatar
BO ZHANG committed
438
        assert self.ver_sim in CP["sim"]["versions"]
BO ZHANG's avatar
BO ZHANG committed
439
440
        fn = "{}_{}_{}_CRS_{}_{}_{}_{:02d}_L0_1.fits".format(
            self._telescope, self._instrument, self._survey,
BO ZHANG's avatar
BO ZHANG committed
441
            self.exp_start, self.exp_stop, self.obs_id, detector)
BO ZHANG's avatar
BO ZHANG committed
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
        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

        """
BO ZHANG's avatar
BO ZHANG committed
460
        assert self.ver_sim in CP["sim"]["versions"]
BO ZHANG's avatar
BO ZHANG committed
461
462
        fn = "{}_{}_{}_SCI_{}_{}_{}_{:02d}_{}".format(
            self._telescope, self._instrument, self._survey,
BO ZHANG's avatar
BO ZHANG committed
463
            self.exp_start, self.exp_stop, self.obs_id, detector, post)
BO ZHANG's avatar
BO ZHANG committed
464
465
        return os.path.join(self.dir_l1, fn)

BO ZHANG's avatar
BO ZHANG committed
466
    def get_bias(self, detector=6):
BO ZHANG's avatar
tweaks    
BO ZHANG committed
467
        """ get bias data """
468
        if self.custom_bias is None:
BO ZHANG's avatar
BO ZHANG committed
469
470
471
472
473
474
475
            return os.path.join(
                self.path_aux,
                "C6.1_ref_crds",
                "csst_msc_{}_{:02d}_{:04d}.fits".format(
                    "bias", detector, 1 if self.ref_version is None else self.ref_version
                )
            )
BO ZHANG's avatar
BO ZHANG committed
476
        else:
477
            return self.custom_bias.format("bias", detector)
BO ZHANG's avatar
BO ZHANG committed
478

BO ZHANG's avatar
BO ZHANG committed
479
    def get_dark(self, detector=6):
BO ZHANG's avatar
tweaks    
BO ZHANG committed
480
        """ get dark data """
481
        if self.custom_dark is None:
BO ZHANG's avatar
BO ZHANG committed
482
483
484
485
486
487
488
            return os.path.join(
                self.path_aux,
                "C6.1_ref_crds",
                "csst_msc_{}_{:02d}_{:04d}.fits".format(
                    "dark", detector, 1 if self.ref_version is None else self.ref_version
                )
            )
BO ZHANG's avatar
BO ZHANG committed
489
        else:
490
            return self.custom_dark.format("dark", detector)
BO ZHANG's avatar
BO ZHANG committed
491

BO ZHANG's avatar
BO ZHANG committed
492
    def get_flat(self, detector=6):
BO ZHANG's avatar
tweaks    
BO ZHANG committed
493
        """ get flat data """
494
        if self.custom_flat is None:
BO ZHANG's avatar
BO ZHANG committed
495
496
497
498
499
500
501
            return os.path.join(
                self.path_aux,
                "C6.1_ref_crds",
                "csst_msc_{}_{:02d}_{:04d}.fits".format(
                    "flat", detector, 1 if self.ref_version is None else self.ref_version
                )
            )
BO ZHANG's avatar
BO ZHANG committed
502
        else:
503
            return self.custom_flat.format("flat", detector)
BO ZHANG's avatar
BO ZHANG committed
504

BO ZHANG's avatar
BO ZHANG committed
505
    def get_axeconf(self):
BO ZHANG's avatar
BO ZHANG committed
506
        return os.path.join(self.path_aux, "axeconf")  # "/home/csstpipeline/L1Pipeline/aux/axeconf"
BO ZHANG's avatar
BO ZHANG committed
507

BO ZHANG's avatar
BO ZHANG committed
508
509
510
511
512
    def l1_file(self, name="", comment=""):
        """ L1 file path

        Parameters
        ----------
BO ZHANG's avatar
BO ZHANG committed
513
        name : str
BO ZHANG's avatar
BO ZHANG committed
514
            file name
BO ZHANG's avatar
BO ZHANG committed
515
        comment : str
BO ZHANG's avatar
BO ZHANG committed
516
517
518
519
520
521
522
523
524
525
526
527
528
            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
tweaks    
BO ZHANG committed
529
530
    def get_sls_info(self):
        """ Get the target SLS image header info and return. """
BO ZHANG's avatar
BO ZHANG committed
531
        # if self.use_dfs:
BO ZHANG's avatar
BO ZHANG committed
532
533
534
535
        #     raise NotImplementedError()
        # else:
        assert len(self.target_detectors) == 1
        header = fits.getheader(self.l0_detector(self.target_detectors[0]), ext=1)
BO ZHANG's avatar
BO ZHANG committed
536
537
        return header

BO ZHANG's avatar
tweaks    
BO ZHANG committed
538
539
    def get_mbi_info(self):
        """ Get all MBI image header info and return as a table. """
BO ZHANG's avatar
BO ZHANG committed
540
        # if self.use_dfs:
BO ZHANG's avatar
BO ZHANG committed
541
542
543
        #     raise NotImplementedError()
        # else:
        info = Table.read("/nfsdata/share/csst_simulation_data/Cycle-5-SimuData/slitlessSpectroscopy/t_mbi_l1.fits")
BO ZHANG's avatar
BO ZHANG committed
544
545
        return info

BO ZHANG's avatar
BO ZHANG committed
546
    @staticmethod
547
    def quickstart(ver_sim="C5.2", datatype="mbi", dir_l1=".", exposure_id=100,
548
                   use_dfs=False, dfs_node="kmust", clear_l1=False, n_jobs=18, backend="multiprocessing"):
549
550
        """
        Quick dataset generator for tests on dandelion or PML
BO ZHANG's avatar
BO ZHANG committed
551
552
553

        Parameters
        ----------
BO ZHANG's avatar
BO ZHANG committed
554
        ver_sim : str
BO ZHANG's avatar
BO ZHANG committed
555
            {"C5.2"}
BO ZHANG's avatar
BO ZHANG committed
556
        datatype : str
BO ZHANG's avatar
BO ZHANG committed
557
            {"mbi", "sls"}
BO ZHANG's avatar
BO ZHANG committed
558
        dir_l1 : str
BO ZHANG's avatar
BO ZHANG committed
559
            output directory
BO ZHANG's avatar
BO ZHANG committed
560
        exposure_id : int
BO ZHANG's avatar
BO ZHANG committed
561
            The serial number of the exposure. 20-154 for C5.2.
BO ZHANG's avatar
BO ZHANG committed
562
        use_dfs : bool
BO ZHANG's avatar
BO ZHANG committed
563
            If True, use DFS.
BO ZHANG's avatar
BO ZHANG committed
564
565
        dfs_node : str
            The DFS node. Defaults to "kmust", could be "pml".
566
        clear_l1 : bool
567
            If True, clear dir_l1.
568
569
570
571
        n_jobs : int
            The number of jobs.
        backend : str
            The joblib backend.
BO ZHANG's avatar
BO ZHANG committed
572
573
574

        Returns
        -------
BO ZHANG's avatar
tweaks    
BO ZHANG committed
575
        CsstMsDataManager
BO ZHANG's avatar
BO ZHANG committed
576
            The Main Survey Data Manager instance.
BO ZHANG's avatar
BO ZHANG committed
577
578
579
580
581
582
583
584
585
586
        """
        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)
BO ZHANG's avatar
BO ZHANG committed
587
            path_aux = "/nfsdata/users/csstpipeline/L1Pipeline/aux"
BO ZHANG's avatar
BO ZHANG committed
588
        elif hostname == "dandelion" and datatype == "sls":
BO ZHANG's avatar
BO ZHANG committed
589
            dir_l0 = "/nfsdata/share/csst_simulation_data/Cycle-5-SimuData/slitlessSpectroscopy/" \
BO ZHANG's avatar
BO ZHANG committed
590
                     "NGP_AstrometryON_shearOFF_Spec/MSC_{:07d}/".format(exposure_id)
BO ZHANG's avatar
BO ZHANG committed
591
            path_aux = "/nfsdata/users/csstpipeline/L1Pipeline/aux"
BO ZHANG's avatar
BO ZHANG committed
592
593
594
595
596

        # PMO
        elif hostname == "ubuntu" and datatype == "mbi":
            dir_l0 = "/share/simudata/CSSOSDataProductsSims/data/CSSTSimImage_C5/" \
                     "NGP_AstrometryON_shearOFF/MSC_{:07d}/".format(exposure_id)
BO ZHANG's avatar
BO ZHANG committed
597
            path_aux = "/L1Pipeline/aux"
BO ZHANG's avatar
BO ZHANG committed
598
599
600
        elif hostname == "ubuntu" and datatype == "sls":
            dir_l0 = "/share/simudata/CSSOSDataProductsSims/data/CSSTSimImage_C5/" \
                     "NGP_AstrometryON_shearOFF_Spec/MSC_{:07d}/".format(exposure_id)
BO ZHANG's avatar
BO ZHANG committed
601
            path_aux = "/L1Pipeline/aux"
BO ZHANG's avatar
BO ZHANG committed
602
        else:
BO ZHANG's avatar
tweaks    
BO ZHANG committed
603
            raise ValueError("@DM: invalid hostname {} or datatype {}!".format(hostname, datatype))
BO ZHANG's avatar
BO ZHANG committed
604

BO ZHANG's avatar
BO ZHANG committed
605
        return CsstMsDataManager.from_dir(
BO ZHANG's avatar
BO ZHANG committed
606
            ver_sim=ver_sim, datatype=datatype, dir_l0=dir_l0, dir_l1=dir_l1, path_aux=path_aux,
BO ZHANG's avatar
BO ZHANG committed
607
608
            use_dfs=use_dfs, dfs_node=dfs_node, n_jobs=n_jobs, backend=backend
        )
BO ZHANG's avatar
BO ZHANG committed
609

610
611
612
    def __repr__(self):
        lines = ""
        lines += "<CsstMsDataManager>\n"
613
614
615
616
617
618
        lines += f"- Data type = {self.datatype}\n"
        lines += f"- Valid detectors = {self.valid_detectors}\n"
        lines += f"- Available detectors = {self.available_detectors}\n"
        lines += f"- Target detectors = {self.target_detectors}\n"
        lines += f"- dir_l0 = {self.dir_l0}\n"
        lines += f"- dir_l1 = {self.dir_l1}\n"
BO ZHANG's avatar
BO ZHANG committed
619
        lines += f"- use_dfs = {self.use_dfs}\n"
620
621
        lines += f"- dfs_node = {self.dfs_node}\n"
        lines += f"- CSST_DFS_GATEWAY = " + os.getenv("CSST_DFS_GATEWAY") + "\n"
622
623
        return lines

624
    def remove_files(self, fmt="*.fits"):
BO ZHANG's avatar
BO ZHANG committed
625
        """ Remove L1 files conforming the format. """
626
        os.system(f"rm -rf {os.path.join(self.dir_l1, fmt)}")
627

BO ZHANG's avatar
BO ZHANG committed
628
629
630
631
    def remove_dir(self, dir_name):
        """ Remove L1 (sub-)directory. """
        os.system(f"rm -rf {os.path.join(self.dir_l1, dir_name)}")

632
633
634
635
    @staticmethod
    def clear_dir(dir_path):
        os.system(f"rm -rf {dir_path}/*")

BO ZHANG's avatar
BO ZHANG committed
636
637
638
639
640
641
642
643
644
645
646
647
    # DFS interfaces
    @property
    def dfs_node(self):
        return self._dfs_node

    @dfs_node.setter
    def dfs_node(self, dfs_node):
        # for DFS configuration, defaults to "local", could be "pml"
        assert dfs_node in DFS_CONF.keys()
        self._dfs_node = dfs_node
        for k, v in DFS_CONF[dfs_node].items():
            os.environ[k] = v
648

649
650
    @staticmethod
    def dfs_is_available():
BO ZHANG's avatar
BO ZHANG committed
651
652
        """ Test if DFS works. """
        try:
BO ZHANG's avatar
BO ZHANG committed
653
            tbl = CatalogApi().catalog_query(
BO ZHANG's avatar
tweaks    
BO ZHANG committed
654
655
656
657
658
659
660
661
662
                catalog_name="gaia3",
                ra=180,
                dec=0,
                radius=.1,
                columns=("ra", "dec"),
                min_mag=0,
                max_mag=30,
                obstime=-1,
                limit=-1
BO ZHANG's avatar
BO ZHANG committed
663
664
            )
            return len(tbl) > 0
BO ZHANG's avatar
BO ZHANG committed
665
666
667
        except:
            return False

668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
    def get_coord(self):
        """ Get pointing coordinate. """
        header = fits.getheader(self.l0_detector(detector=self.target_detectors[0]))
        pointing_ra = header["RA_OBJ"]
        pointing_dec = header["DEC_OBJ"]
        return pointing_ra, pointing_dec

    def dfs_rc_auto(self):
        """ Download RC. """
        assert self.dfs_is_available()
        assert len(self.target_detectors) >= 1
        pointing_ra, pointing_dec = self.get_coord()
        refcat = self.dfs_rc_query(
            ra=pointing_ra,
            dec=pointing_dec,
            columns=(
                    'ref_epoch',
                    'ra',
                    'ra_error',
                    'dec',
                    'dec_error',
                    'parallax',
                    'parallax_error',
                    'pmra',
                    'pmra_error',
                    'pmdec',
                    'pmdec_error',
                    'phot_g_mean_mag',
                    'source_id'
            ),
            radius=2,
            min_mag=0,
            max_mag=30,
            obstime=-1,
            limit=-1
        )
        print(
            f"RA, Dec = [{pointing_ra}, {pointing_dec}], "
            f"Status = {isinstance(refcat, Table)}, "
            f"N_entries = {len(refcat)}"
        )
        return refcat

BO ZHANG's avatar
BO ZHANG committed
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
    def dfs_rc_query(
            self,
            ra=180,
            dec=0,
            radius=2,
            columns=(
                    'ref_epoch',
                    'ra',
                    'ra_error',
                    'dec',
                    'dec_error',
                    'parallax',
                    'parallax_error',
                    'pmra',
                    'pmra_error',
                    'pmdec',
                    'pmdec_error',
                    'phot_g_mean_mag',
                    'source_id'
            ),
            min_mag=0,
            max_mag=30,
            obstime=-1,
            limit=-1
    ):
BO ZHANG's avatar
BO ZHANG committed
736
737
738
739
740
        """ Query Reference Catalog (RC) from DFS.
        Ref.
        https://gea.esac.esa.int/archive/documentation/GDR3/Gaia_archive/chap_datamodel/
        sec_dm_main_source_catalogue/ssec_dm_gaia_source.html
        """
741
742
743
744
745
746
747
748
749
750
751
752
753
        cat = self.dfs_CatApi.catalog_query(
            catalog_name="gaia3",
            ra=ra,
            dec=dec,
            columns=columns,
            radius=radius,
            min_mag=min_mag,
            max_mag=max_mag,
            obstime=obstime,
            limit=limit
        )
        self.logger_ppl.info(f"Results from DFS CATAPI: code = {cat['code']}, message = {cat['message']}")
        if cat["code"] == 0:
754
            self.logger_ppl.info(f"totalCount = {cat['totalCount']}, columns = {cat['columns']}")
755
756
757
            return self.dfs_CatApi.to_table(cat)
        else:
            raise ValueError("Bad results!")
BO ZHANG's avatar
BO ZHANG committed
758
759
760
761
762

    def dfs_l1_push(self):
        """ Push MBI/SLS L1 data to DFS. """
        # l1api = get_l1api()
        # l1api.write()
763
        return
BO ZHANG's avatar
BO ZHANG committed
764
765
766
767
768

    def dfs_l2_push(self):
        """ Push SLS spectra to DFS. """
        pass

BO ZHANG's avatar
BO ZHANG committed
769
770
771
772
773
774
775
776
777
778
779
780
781
    @staticmethod
    def from_l1id(
            l1_id="1000000001",
            datatype="sls",
            dir_l0="/L1Pipeline/L0",
            dir_l1="/L1Pipeline/L1",
            use_dfs=True,
            dfs_node="pml",
            clear_l1=False,
            dfs_root="/share/dfs"
    ):
        pass

BO ZHANG's avatar
BO ZHANG committed
782
    @staticmethod
BO ZHANG's avatar
BO ZHANG committed
783
784
    def from_dfs(
            obs_id="100000100",
BO ZHANG's avatar
BO ZHANG committed
785
            ver_sim="C5.2",
BO ZHANG's avatar
BO ZHANG committed
786
787
788
            datatype="mbi",
            dir_l0="/L1Pipeline/L0",
            dir_l1="/L1Pipeline/L1",
BO ZHANG's avatar
tweaks    
BO ZHANG committed
789
            path_aux="/L1Pipeline/aux",
BO ZHANG's avatar
BO ZHANG committed
790
            use_dfs=True,
BO ZHANG's avatar
BO ZHANG committed
791
792
            dfs_node="pml",
            clear_l1=False,
BO ZHANG's avatar
BO ZHANG committed
793
794
            dfs_root="/share/dfs",
            n_jobs=18,
BO ZHANG's avatar
BO ZHANG committed
795
            backend="multiprocessing",
BO ZHANG's avatar
BO ZHANG committed
796
797
            device="CPU",
            **kwargs
BO ZHANG's avatar
BO ZHANG committed
798
    ):
BO ZHANG's avatar
BO ZHANG committed
799
        """ Initialize CsstMsDataManager from DFS. """
800
801
802
803
804
        # (clear and) make directories
        if os.path.exists(dir_l0):
            os.system(f"rm -rf {dir_l0}/*")
        else:
            os.mkdir(dir_l0)
BO ZHANG's avatar
BO ZHANG committed
805
806
807
808
        # if os.path.exists(dir_l1):
        #     os.system(f"rm -rf {dir_l1}/*")
        # else:
        #     os.mkdir(dir_l1)
BO ZHANG's avatar
BO ZHANG committed
809
810
811
        # os.chdir(dir_l1)
        if not os.path.exists(dir_l1):
            os.mkdir(dir_l1)
BO ZHANG's avatar
BO ZHANG committed
812
813
        elif clear_l1:
            os.system(f"rm -rf {dir_l1}/*")
814
815
        os.chdir(dir_l1)

BO ZHANG's avatar
BO ZHANG committed
816
        print(f"Query obs_id={obs_id} ...", end="")
BO ZHANG's avatar
BO ZHANG committed
817
        records = CsstMsDataManager(dfs_node=dfs_node, verbose=False).dfs_L0DataApi.find(obs_id=obs_id)
BO ZHANG's avatar
tweaks    
BO ZHANG committed
818
        print(f"{records['totalCount']} records obtained!")
BO ZHANG's avatar
BO ZHANG committed
819
820
821
        tbl = Table([_.__dict__ for _ in records["data"]])
        tbl.sort(["detector_no", "obs_type"])

BO ZHANG's avatar
tweaks    
BO ZHANG committed
822
823
        print("Making symbolic links ...")
        for i_rec in range(len(tbl)):
BO ZHANG's avatar
BO ZHANG committed
824
            os.symlink(
BO ZHANG's avatar
tweaks    
BO ZHANG committed
825
826
                src=os.path.join(dfs_root, tbl["file_path"][i_rec]),
                dst=os.path.join(dir_l0, os.path.basename(tbl["file_path"][i_rec])),
BO ZHANG's avatar
BO ZHANG committed
827
828
829
830
            )

        # initialize dm
        dm = CsstMsDataManager.from_dir(
BO ZHANG's avatar
BO ZHANG committed
831
            ver_sim=ver_sim,
BO ZHANG's avatar
BO ZHANG committed
832
833
834
835
836
837
838
            datatype=datatype,
            dir_l0=dir_l0,
            dir_l1=dir_l1,
            path_aux=path_aux,
            use_dfs=use_dfs,
            dfs_node=dfs_node,
            n_jobs=n_jobs,
BO ZHANG's avatar
BO ZHANG committed
839
            backend=backend,
BO ZHANG's avatar
BO ZHANG committed
840
841
            device=device,
            **kwargs
BO ZHANG's avatar
BO ZHANG committed
842
        )
BO ZHANG's avatar
BO ZHANG committed
843
        assert dm.obs_id == obs_id
BO ZHANG's avatar
BO ZHANG committed
844
845
846

        return dm

BO ZHANG's avatar
BO ZHANG committed
847
    def dfs_l0_query(self, obs_id: str = "100000100"):
BO ZHANG's avatar
BO ZHANG committed
848
        """ Query L0 data from DFS. """
BO ZHANG's avatar
BO ZHANG committed
849
        result = self.dfs_L0DataApi.find(obs_id=str(obs_id))
BO ZHANG's avatar
BO ZHANG committed
850
        print(f"{result['totalCount']} records returned from DFS.")
BO ZHANG's avatar
tweaks    
BO ZHANG committed
851
852
853
        if not result["code"] == 0:
            raise ValueError(f"DFS returns non-zero code! ({result['code']})")
        tbl = Table([_.__dict__ for _ in result["data"]])
BO ZHANG's avatar
BO ZHANG committed
854
855
856
857
858
859
860
        tbl.sort(["detector_no", "obs_type"])
        # Check if all 30 detectors are available
        for detector in CP["all"]["detectors"]:
            for obs_type in ["sci", "cosmic_ray"]:
                if np.sum((tbl["detector_no"] == f"{detector:02d}") & (tbl["obs_type"] == obs_type)) == 0:
                    self.logger_ppl.warning(f"Record not found for detector {detector:02d} and obs_type {obs_type}")
        return tbl
BO ZHANG's avatar
BO ZHANG committed
861

BO ZHANG's avatar
BO ZHANG committed
862
863
864
865
866
867
868
869
870
871
872
873
    def dfs_l0_check_all(self):
        """ Check all C5.2 L0 data is available in DFS. """
        is_good = True
        for obs_id in range(100000020, 100000155):
            tbl = self.dfs_l0_query(obs_id=f"{obs_id}")
            if len(tbl) == 60:
                self.logger_ppl.info(f"DFS returns {len(tbl)} records for obs_id={obs_id}")
            else:
                is_good = False
                self.logger_ppl.warning(f"DFS returns {len(tbl)} records for obs_id={obs_id}")
        return is_good

874
    def dfs_l1_query(self, obs_id, detector):
BO ZHANG's avatar
BO ZHANG committed
875
876
        """ Query L1 data from DFS. """
        pass
877
878
879
880


# temporarily compatible with old interface
CsstMbiDataManager = CsstMsDataManager