Skip to content

Commit

Permalink
PSF size max for wfe drift and field positions coeffs
Browse files Browse the repository at this point in the history
  • Loading branch information
@JarronL
JarronL committed May 28, 2020
1 parent 1b6863f commit be11b6a
Show file tree
Hide file tree
Showing 2 changed files with 69 additions and 22 deletions.
46 changes: 27 additions & 19 deletions pynrc/nrc_utils.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1163,7 +1163,7 @@ def gen_psf_coeff(filter_or_bp, pupil=None, mask=None, module='A',


# How many processors to split into?
nproc = nproc_use(fov_pix, oversample, npsf, coron=coron_obs) #if poppy.conf.use_multiprocessing else 1
nproc = nproc_use(fov_pix, oversample, npsf)#, coron=coron_obs) #if poppy.conf.use_multiprocessing else 1
_log.debug('nprocessors: {}; npsf: {}'.format(nproc, npsf))

setup_logging('WARN', verbose=False)
Expand Down Expand Up @@ -1404,6 +1404,10 @@ def wfed_coeff(filter_or_bp, force=False, save=True, save_name=None, nsplit=None
Full path name of save file (.npy) to save/load.
If None, then a name is automatically generated,
matching the :func:`gen_psf_coeff` function.
nsplit : int
Number of processors to split over. There are checks to
make sure you're not requesting more processors than the
current machine has available.
Example
-------
Expand Down Expand Up @@ -1437,6 +1441,10 @@ def wfed_coeff(filter_or_bp, force=False, save=True, save_name=None, nsplit=None
bp = filter_or_bp
filter = bp.name

# defaults
fov_pix = kwargs['fov_pix'] if 'fov_pix' in list(kwargs.keys()) else 33
oversample = kwargs['oversample'] if 'oversample' in list(kwargs.keys()) else 4

# Final filename to save coeff
if save_name is None:
# Name to save array of oversampled coefficients
Expand All @@ -1459,15 +1467,14 @@ def wfed_coeff(filter_or_bp, force=False, save=True, save_name=None, nsplit=None

# Cycle through WFE drifts for fitting
wfe_list = np.array([0,1,2,5,10,20,40])
nwfe = len(wfe_list)
npos = len(wfe_list)

# Split over multiple processors?
nsplit_max = nproc_use(fov_pix, oversample, npos)#, coron=coron_obs)
if nsplit is None:
fov_pix = kwargs['fov_pix'] if 'fov_pix' in list(kwargs.keys()) else 33
oversample = kwargs['oversample'] if 'oversample' in list(kwargs.keys()) else 4
pupil = kwargs['pupil'] if 'pupil' in list(kwargs.keys()) else None
coron_obs = (pupil is not None) and ('LYOT' in pupil)
nsplit = nproc_use(fov_pix, oversample, nwfe, coron=coron_obs)
nsplit = nproc_use(fov_pix, oversample, npos)#, coron=coron_obs)

# Compare to number of PSFs
if ('quick' in list(kwargs.keys())) and (kwargs['quick']==True):
Expand All @@ -1478,10 +1485,14 @@ def wfed_coeff(filter_or_bp, force=False, save=True, save_name=None, nsplit=None
dw = 2.5
npsf = np.ceil(20 * dw)
npsf = 5 if npsf<5 else int(npsf)
nsplit_psf = nproc_use(fov_pix, oversample, npsf, coron=coron_obs)
nsplit_psf = nproc_use(fov_pix, oversample, npsf)#, coron=coron_obs)
if nsplit_psf > nsplit:
nsplit = 1

# Double check we're not requesting too many processors
nsplit = nsplit_max if nsplit > nsplit_max else nsplit
# nsplit = npos

# Create worker arguments with kwargs as an argument input
worker_args = []
args = [bp]
Expand Down Expand Up @@ -1520,7 +1531,7 @@ def wfed_coeff(filter_or_bp, force=False, save=True, save_name=None, nsplit=None

# Fit each pixel with a polynomial and save the coefficient
cf_shape = cf_wfe.shape[1:]
cf_wfe = cf_wfe.reshape([nwfe, -1])
cf_wfe = cf_wfe.reshape([npos, -1])
lxmap = np.array([np.min(wfe_list), np.max(wfe_list)])
cf_fit = jl_poly_fit(wfe_list, cf_wfe, deg=4, use_legendre=True, lxmap=lxmap)
cf_fit = cf_fit.reshape([-1, cf_shape[0], cf_shape[1], cf_shape[2]])
Expand Down Expand Up @@ -1567,6 +1578,10 @@ def field_coeff_resid(filter_or_bp, coeff0, force=False, save=True, save_name=No
Full path name of save file (.npy) to save/load.
If None, then a name is automatically generated,
matching the :func:`gen_psf_coeff` function.
nsplit : int
Number of processors to split over. There are checks to
make sure you're not requesting more processors than the
current machine has available.
Example
Expand Down Expand Up @@ -1610,17 +1625,6 @@ def field_coeff_resid(filter_or_bp, coeff0, force=False, save=True, save_name=No
fov_pix = kwargs['fov_pix'] if 'fov_pix' in list(kwargs.keys()) else 33
oversample = kwargs['oversample'] if 'oversample' in list(kwargs.keys()) else 4

fov_max = 128
if oversample>4: fov_max /= 2
if fov_pix>fov_max:
fov_pix = fov_max if (fov_pix % 2 == 0) else fov_max + 1
# Trim input coeff0
new_shape = fov_pix*oversample
cf0_new = np.array([pad_or_cut_to_size(im, new_shape) for im in coeff0])
coeff0 = cf0_new
kwargs['fov_pix'] = fov_pix
kwargs['oversample'] = oversample

# Final filename to save coeff
if save_name is None:
# Name to save array of oversampled coefficients
Expand Down Expand Up @@ -1680,10 +1684,11 @@ def field_coeff_resid(filter_or_bp, coeff0, force=False, save=True, save_name=No
kwargs['include_si_wfe'] = True

# Split over multiple processors?
nsplit_max = nproc_use(fov_pix, oversample, npos)#, coron=coron_obs)
if nsplit is None:
pupil = kwargs['pupil'] if 'pupil' in list(kwargs.keys()) else None
coron_obs = (pupil is not None) and ('LYOT' in pupil)
nsplit = nproc_use(fov_pix, oversample, npos, coron=coron_obs)
nsplit = nproc_use(fov_pix, oversample, npos)#, coron=coron_obs)

# Compare to number of PSFs
if ('quick' in list(kwargs.keys())) and (kwargs['quick']==True):
Expand All @@ -1698,6 +1703,9 @@ def field_coeff_resid(filter_or_bp, coeff0, force=False, save=True, save_name=No
if nsplit_psf > nsplit:
nsplit = 1

# Double check we're not requesting too many processors
nsplit = nsplit_max if nsplit > nsplit_max else nsplit

# Create worker arguments with kwargs as an input dict
worker_args = []
args = [filter]
Expand Down
45 changes: 42 additions & 3 deletions pynrc/pynrc_core.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -468,6 +468,10 @@ def __init__(self, filter='F210M', pupil=None, mask=None, module='A', ND_acq=Fal
self._wfe_drift = wfe_drift
# Field-dependent WFE
self._wfe_field = False # Don't calculate coefficients by default
# Max FoVs for calculationg drift and field dependnet coefficient residuals
# Any pixels beyond this size will be considered to have 0 residual difference
self._fovmax_wfedrift = 256
self._fovmax_wfefield = 128

self._psf_coeff_mod = {
'wfe_drift': None, 'wfe_drift_lxmap': None,
Expand Down Expand Up @@ -1442,7 +1446,6 @@ def update_psf_coeff(self, fov_pix=None, oversample=None,
if jitter=='none':
jitter = None

#print(opd)
self._psf_info={'fov_pix':fov_pix, 'oversample':oversample, 'quick':quick,
'offset_r':offset_r, 'offset_theta':offset_theta,
'tel_pupil':tel_pupil, 'save':save, 'force':force, 'use_legendre':use_legendre,
Expand All @@ -1461,6 +1464,11 @@ def update_psf_coeff(self, fov_pix=None, oversample=None,
wfe_kwargs['mask'] = self._mask
wfe_kwargs['module'] = self.module
wfe_kwargs['bar_offset'] = 0

# fov_pix should not be more than some size, otherwise memory issues
fov_max = self._fovmax_wfedrift if wfe_kwargs['oversample']<=4 else self._fovmax_wfedrift / 2
if wfe_kwargs['fov_pix']>fov_max:
wfe_kwargs['fov_pix'] = fov_max if (wfe_kwargs['fov_pix'] % 2 == 0) else fov_max + 1

res1, res2 = wfed_coeff(self.bandpass, **wfe_kwargs)
self._psf_coeff_mod['wfe_drift'] = res1
Expand All @@ -1479,7 +1487,17 @@ def update_psf_coeff(self, fov_pix=None, oversample=None,
field_kwargs['mask'] = self._mask
field_kwargs['module'] = self.module

res, v2grid, v3grid = field_coeff_resid(self.bandpass, self._psf_coeff, **field_kwargs)
# fov_pix should not be more than some size, otherwise memory issues
fov_max = self._fovmax_wfefield if field_kwargs['oversample']<=4 else self._fovmax_wfefield / 2
if field_kwargs['fov_pix']>fov_max:
field_kwargs['fov_pix'] = fov_max if (field_kwargs['fov_pix'] % 2 == 0) else fov_max + 1

new_shape = field_kwargs['fov_pix']*field_kwargs['oversample']
coeff0 = np.array([pad_or_cut_to_size(im, new_shape) for im in self._psf_coeff])
else:
coeff0 = self._psf_coeff

res, v2grid, v3grid = field_coeff_resid(self.bandpass, coeff0, **field_kwargs)
self._psf_coeff_mod['si_field'] = res
self._psf_coeff_mod['si_field_v2grid'] = v2grid
self._psf_coeff_mod['si_field_v3grid'] = v3grid
Expand Down Expand Up @@ -1520,6 +1538,11 @@ def update_psf_coeff(self, fov_pix=None, oversample=None,
wfe_kwargs['mask'] = None
wfe_kwargs['module'] = self.module

# fov_pix should not be more than some size, otherwise memory issues
fov_max = self._fovmax_wfedrift if wfe_kwargs['oversample']>4 else self._fovmax_wfedrift / 2
if wfe_kwargs['fov_pix']>fov_max:
wfe_kwargs['fov_pix'] = fov_max if (wfe_kwargs['fov_pix'] % 2 == 0) else fov_max + 1

res1, res2 = wfed_coeff(self.bandpass, **wfe_kwargs)
self._psf_coeff_bg_mod['wfe_drift'] = res1
self._psf_coeff_bg_mod['wfe_drift_lxmap'] = res2
Expand All @@ -1535,7 +1558,17 @@ def update_psf_coeff(self, fov_pix=None, oversample=None,
field_kwargs['mask'] = None
field_kwargs['module'] = self.module

res, v2grid, v3grid = field_coeff_resid(self.bandpass, self._psf_coeff, **field_kwargs)
# fov_pix should not be more than some size, otherwise memory issues
fov_max = self._fovmax_wfefield if field_kwargs['oversample']>4 else self._fovmax_wfefield / 2
if field_kwargs['fov_pix']>fov_max:
field_kwargs['fov_pix'] = fov_max if (field_kwargs['fov_pix'] % 2 == 0) else fov_max + 1

new_shape = field_kwargs['fov_pix']*field_kwargs['oversample']
coeff0 = np.array([pad_or_cut_to_size(im, new_shape) for im in self._psf_coeff])
else:
coeff0 = self._psf_coeff

res, v2grid, v3grid = field_coeff_resid(self.bandpass, coeff0, **field_kwargs)
self._psf_coeff_bg_mod['si_field'] = res
self._psf_coeff_bg_mod['si_field_v2grid'] = v2grid
self._psf_coeff_bg_mod['si_field_v3grid'] = v3grid
Expand Down Expand Up @@ -2020,6 +2053,11 @@ def gen_psf(self, sp=None, return_oversample=False, use_bg_psf=False,
cf_fit = cf_fit.reshape([cf_fit.shape[0], -1])
cf_mod = jl_poly(np.array([wfe_drift]), cf_fit, use_legendre=True, lxmap=lxmap)
cf_mod = cf_mod.reshape(psf_coeff.shape)
# Pad cf_mod array with 0s if undersized
if not np.allclose(psf_coeff.shape, cf_mod.shape):
new_shape = psf_coeff.shape[1:]
cf_mod_resize = np.array([pad_or_cut_to_size(im, new_shape) for im in cf_mod])
cf_mod = cf_mod_resize
psf_coeff_mod += cf_mod

# Modify PSF coefficients based on field-dependent
Expand Down Expand Up @@ -2070,6 +2108,7 @@ def gen_psf(self, sp=None, return_oversample=False, use_bg_psf=False,
# print(v2,v3)
nfield = np.size(v2)
cf_mod = field_coeff_func(v2grid, v3grid, cf_fit, v2, v3)
# Pad cf_mod array with 0s if undersized
if not np.allclose(psf_coeff.shape, cf_mod.shape):
new_shape = psf_coeff.shape[1:]
cf_mod_resize = np.array([pad_or_cut_to_size(im, new_shape) for im in cf_mod])
Expand Down

0 comments on commit be11b6a

Please sign in to comment.