From 737ea7318ba41091f05190d4c1d83553f95f3907 Mon Sep 17 00:00:00 2001 From: Tammo Jan Dijkema Date: Fri, 24 Apr 2020 16:44:25 +0200 Subject: [PATCH] Add function for searching HDF5 --- lofarimaging/hdf5util.py | 40 ++++++++++++++++++++++++++++++++++++++-- 1 file changed, 38 insertions(+), 2 deletions(-) diff --git a/lofarimaging/hdf5util.py b/lofarimaging/hdf5util.py index b575a0a..ccb80f7 100644 --- a/lofarimaging/hdf5util.py +++ b/lofarimaging/hdf5util.py @@ -32,7 +32,7 @@ from typing import List import numpy as np import h5py -__all__ = ["get_new_obsname", "write_hdf5", "merge_hdf5"] +__all__ = ["get_new_obsname", "write_hdf5", "merge_hdf5", "get_obsnums"] def get_new_obsname(h5file: h5py.File): @@ -143,10 +143,46 @@ def merge_hdf5(src_filename: str, dest_filename: str, obslist: List[str] = None) """ with h5py.File(dest_filename) as dest_file: with h5py.File(src_filename, 'r') as src_file: - if obslist == None: + if obslist is None: obslist = src_file for src_obsname in obslist: dest_obsname = get_new_obsname(dest_file) h5py.h5o.copy(src_file.id, bytes(src_obsname, 'utf-8'), dest_file.id, bytes(dest_obsname, 'utf-8')) dest_file.flush() + +def get_obsnums(h5file: h5py.File, + start_date: datetime.datetime = None, + end_date: datetime.datetime = None, + rcu_modes: List[int] = None, + station_name: str = None, + subband: int = None, + extent: List[int] = None) -> List[str]: + """ + Find observations in an HDF5 file with many observations + """ + matching_obs = [] + for obs in h5file: + if start_date is not None or end_date is not None: + obsdate = datetime.datetime.strptime(h5file[obs].attrs["obstime"], "%Y-%m-%d %H:%M:%S") + if start_date is not None and obsdate < start_date: + continue + if end_date is not None and obsdate > end_date: + continue + if rcu_modes is not None and h5file[obs].attrs["rcu_mode"] not in rcu_modes: + continue + if station_name is not None and h5file[obs].attrs["station_name"] != station_name: + continue + if subband is not None and h5file[obs].attrs["subband"] != subband: + continue + if extent is not None: + has_ground_image_with_extent = False + for groundimg in h5file[obs]["ground_images"]: + if list(h5file[obs]["ground_images"][groundimg].attrs["extent"]) == extent: + has_ground_image_with_extent = True + break + if not has_ground_image_with_extent: + continue + matching_obs.append(obs) + + return matching_obs