|
|
@ -1,5 +1,6 @@ |
|
|
|
import glob |
|
|
|
from datetime import date |
|
|
|
import copy |
|
|
|
|
|
|
|
import numpy as np |
|
|
|
from uncertainties import unumpy as unp |
|
|
@ -8,14 +9,37 @@ import xarray as xr |
|
|
|
|
|
|
|
|
|
|
|
def get_mask(dataArray): |
|
|
|
"""generate a bool mask array for given dataArray |
|
|
|
|
|
|
|
:param dataArray: The given dataArray |
|
|
|
:type dataArray: xarray DataArray |
|
|
|
:return: the mask array |
|
|
|
:rtype: numpy array of bool elements |
|
|
|
""" |
|
|
|
return np.ones(dataArray.shape, dtype=bool) |
|
|
|
|
|
|
|
|
|
|
|
def remove_bad_shots(dataArray, **kwargs): |
|
|
|
"""copy and remove bad shots from the dataArray |
|
|
|
|
|
|
|
:param dataArray: The given dataArray |
|
|
|
:type dataArray: xarray DataArray |
|
|
|
:return: The dataArray after removement |
|
|
|
:rtype: xarray DataArray |
|
|
|
""" |
|
|
|
dataArray = copy.deepcopy(dataArray) |
|
|
|
dataArray.loc[dict(kwargs)] = np.nan |
|
|
|
return dataArray |
|
|
|
|
|
|
|
|
|
|
|
def auto_rechunk(dataSet): |
|
|
|
"""rechunk the dataSet or dataArray using auto rechunk function |
|
|
|
|
|
|
|
:param dataSet: The given dataArray or dataSet |
|
|
|
:type dataSet: xarray DataArray or xarray DataSet |
|
|
|
:return: The chuncked dataArray or dataSet |
|
|
|
:rtype: xarray DataArray or xarray DataSet |
|
|
|
""" |
|
|
|
kwargs = { |
|
|
|
key: "auto" |
|
|
|
for key in dataSet.dims |
|
|
@ -24,6 +48,15 @@ def auto_rechunk(dataSet): |
|
|
|
|
|
|
|
|
|
|
|
def copy_chunk(dataSet, dataChunk): |
|
|
|
"""copy the chunk and apply to another dataArray or dataSet |
|
|
|
|
|
|
|
:param dataSet: The dataArray or dataSet will be chunked |
|
|
|
:type dataSet: xarray DataArray or xarray DataSet |
|
|
|
:param dataChunk: The dataArray or dataSet giving the chunk |
|
|
|
:type dataChunk: xarray DataArray or xarray DataSet |
|
|
|
:return: The chuncked dataArray or dataSet |
|
|
|
:rtype: xarray DataArray or xarray DataSet |
|
|
|
""" |
|
|
|
kwargs = { |
|
|
|
key: dataChunk.chunksizes[key] |
|
|
|
for key in dataChunk.chunksizes |
|
|
@ -33,6 +66,17 @@ def copy_chunk(dataSet, dataChunk): |
|
|
|
|
|
|
|
|
|
|
|
def get_h5_file_path(folderpath, maxFileNum=None, filename='*.h5',): |
|
|
|
"""_summary_ |
|
|
|
|
|
|
|
:param folderpath: _description_ |
|
|
|
:type folderpath: _type_ |
|
|
|
:param maxFileNum: _description_, defaults to None |
|
|
|
:type maxFileNum: _type_, optional |
|
|
|
:param filename: _description_, defaults to '*.h5' |
|
|
|
:type filename: str, optional |
|
|
|
:return: _description_ |
|
|
|
:rtype: _type_ |
|
|
|
""" |
|
|
|
filepath = np.sort(glob.glob(folderpath + filename)) |
|
|
|
if maxFileNum is None: |
|
|
|
return filepath |
|
|
|