diff --git a/psfex/psfex_lib.py b/psfex/psfex_lib.py index 4b908a9..25770c9 100644 --- a/psfex/psfex_lib.py +++ b/psfex/psfex_lib.py @@ -2,6 +2,8 @@ import numpy from . import _psfex_pywrap +DEFAULT_EXT='psf_data' + DOUBLE_DTYPE = 'f8' LONG_DTYPE = 'i8' POLY_DIM = 2 @@ -19,8 +21,19 @@ def __str__(self): class PSFEx(dict): - def __init__(self, filename): - self._load(filename) + def __init__(self, *args, **kw): + """ + Load data into a PSFEx object + + parameters + ---------- + filename or FITS object: string, optional + Send filename or open fitsio FITS object + ext: string or int, optional + The extension from which to read + """ + + self.load(*args, **kw) def get_rec(self, row, col): """ @@ -53,19 +66,34 @@ def get_sigma(self): sigma = fwhm/fac return sigma - def _load(self, filename): + def load(self, arg, **kw): """ - Load the PSF information from the fits file - - Also load the C PSFEx object + Load data into a PSFEx object + + parameters + ---------- + filename or FITS object: string, optional + Send filename or open fitsio FITS object + ext: string or int, optional + The extension from which to read """ - self['filename'] = filename - with fitsio.FITS(filename) as fits: - psf_mask=fits['psf_data']['psf_mask'][:] - psf_mask=numpy.array(psf_mask, dtype=DOUBLE_DTYPE) + if isinstance(arg, fitsio.FITS): + self._load_from_fits(arg, **kw) + else: + self['filename'] = arg + with fitsio.FITS(arg) as fits: + self._load_from_fits(fits, **kw) + + def _load_from_fits(self, fits, ext=DEFAULT_EXT): + + hdu = fits[ext] + + psf_mask=hdu['psf_mask'][:] + + psf_mask=numpy.array(psf_mask, dtype=DOUBLE_DTYPE) - h=fits['psf_data'].read_header() + h=hdu.read_header() self._psf_mask=psf_mask @@ -100,4 +128,5 @@ def _load(self, filename): self['contextscale'], self['psf_samp'], self._psf_mask) - + +