def numerical_image_gradients(theta0, delta, scene=None, stamp=None):
dI_dp = []
for i, (p, dp) in enumerate(zip(theta0, delta)):
theta = theta0.copy()
imlo, _ = make_image(scene, stamp, Theta=theta)
theta[i] += dp
imhi, _ = make_image(scene, stamp, Theta=theta)
dI_dp.append((imhi - imlo) / (dp))
return np.array(dI_dp)
def display(stamp, scene, theta):
data = stamp.pixel_values
unc = 1 / (stamp.ierr.reshape(stamp.nx, stamp.ny))
mod, grad = make_image(scene, stamp, Theta=theta)
fig, axes = pl.subplots(1, 3)
axes[0].imshow(data.T, origin='lower')
axes[1].imshow(mod.T, origin='lower')
im = axes[2].imshow(((mod - data) / unc).T, origin='lower')
fig.colorbar(im)
return fig, axes, [data, unc, mod, grad]
def plot_model_images(pos,
scene,
stamps,
axes=None,
colorbars=True,
x=slice(None),
y=slice(None),
share=True,
scale_model=False,
scale_residuals=False):
vals = pos
same_scale = [False, scale_model, scale_residuals, False]
if axes is None:
rfig, raxes = pl.subplots(len(stamps),
4,
figsize=(14, 3.3 * len(stamps) + 0.5),
sharex=share,
sharey=share)
else:
rfig, raxes = None, axes
raxes = np.atleast_2d(raxes)
for i, stamp in enumerate(stamps):
data = stamp.pixel_values
im, grad = make_image(scene, stamp, Theta=vals)
resid = data - im
chi = resid * stamp.ierr.reshape(stamp.nx, stamp.ny)
ims = [data, im, resid, chi]
norm = None
for j, ii in enumerate(ims):
if same_scale[j] and colorbars:
vmin, vmax = cb.vmin, cb.vmax
else:
vmin = vmax = None
ci = raxes[i, j].imshow(ii[x, y].T,
origin='lower',
vmin=vmin,
vmax=vmax)
if (rfig is not None) & colorbars:
cb = rfig.colorbar(ci,
ax=raxes[i, j],
orientation='horizontal')
cb.ax.set_xticklabels(cb.ax.get_xticklabels(), rotation=-55)
text = "{}\n({}, {})".format(stamp.filtername, stamp.crval[0],
stamp.crval[1])
ax = raxes[i, 1]
ax.text(0.6, 0.1, text, transform=ax.transAxes, fontsize=10)
labels = ['Data', 'Model', 'Data-Model', "$\chi$"]
[ax.set_title(labels[i]) for i, ax in enumerate(raxes[0, :])]
return rfig, raxes
def setup_scene(sourceparams=[(1.0, 5., 5., 0.7, 30., 1.0, 0.05)],
splinedata=None,
perturb=0,
filters=['dummy'],
add_noise=False,
snr_max=100.,
stamp_kwargs=[]):
# --- Get Sources and a Scene -----
sources = []
for (flux, x, y, q, pa, n, rh) in sourceparams:
if splinedata is not None:
s = Galaxy(filters=filters, splinedata=splinedata)
s.sersic = n
s.rh = rh
else:
s = SimpleGalaxy(filters=filters)
s.flux = flux
s.ra = x
s.dec = y
s.q = q
s.pa = np.deg2rad(pa)
sources.append(s)
scene = Scene(sources)
theta = scene.get_all_source_params()
label = []
# get stamps
stamps = [make_stamp(**sk) for sk in stamp_kwargs]
if splinedata is not None:
for stamp in stamps:
stamp.scale = np.array([[32.0, 0.0], [0.0, 32.0]])
# --- Generate mock and add to stamp ---
ptrue = theta * np.random.normal(1.0, perturb, size=theta.shape)
for stamp in stamps:
true_image, _ = make_image(scene, stamp, Theta=ptrue)
stamp.pixel_values = true_image.copy()
err = stamp.pixel_values.max() / snr_max
#err = np.sqrt(err**2 + stamp.pixel_values.flatten())
err = np.ones(stamp.npix) * err
stamp.ierr = np.ones(stamp.npix) / err
if add_noise:
noise = np.random.normal(0, err)
stamp.pixel_values += noise.reshape(stamp.nx, stamp.ny)
return scene, stamps, ptrue, label
def setup_scene(galaxy=False, fwhm=1.0, offset=0.0,
size=(30, 30), peak_snr=1e1, add_noise=False):
stamp = make_stamp(size, fwhm, offset=offset)
# --- Get a Source and Scene -----
if galaxy:
ngauss = 1
source = SimpleGalaxy()
source.radii = np.arange(ngauss) * 0.5 + 1.0
source.q = 0.5
source.pa = np.deg2rad(30.)
theta = [100., 10., 10., 0.5, np.deg2rad(10.)]
label = ['$\psi$', '$x$', '$y$', '$q$', '$\\varphi$']
bounds = [(0, 1e4), (0., 30), (0, 30), (0, 1), (0, np.pi/2)]
else:
source = Star()
theta = [100., 10., 10.]
label = ['$\psi$', '$x$', '$y$']
bounds = [(-1e6, 1e6), (-1e5, 1e5), (-1e5, 1e5)]
scene = Scene([source])
# --- Generate mock and add to stamp ---
ptrue = np.array(theta)
true_image, partials = make_image(scene, stamp, Theta=ptrue)
stamp.pixel_values = true_image.copy()
err = stamp.pixel_values.max() / peak_snr
#err = np.sqrt(err**2 + stamp.pixel_values.flatten())
err = np.ones(stamp.npix) * err
stamp.ierr = np.ones(stamp.npix) / err
if add_noise:
noise = np.random.normal(0, err)
stamp.pixel_values += noise.reshape(stamp.nx, stamp.ny)
return scene, stamp, ptrue, label
# Now remove stamp astrometry (but keep stamp.scale)
stamp.dpix_dsky = np.eye(2)
stamp.crval = np.zeros([2])
stamp.crpix = np.zeros([2])
plans = [WorkPlan(stamp) for stamp in stamps]
scene, theta = prep_scene(sourcepars, filters=np.unique(filters).tolist(),
splinedata=paths.galmixture)
theta_init = theta.copy()
ptrue = theta.copy()
p0 = ptrue.copy()
ndim = len(theta)
nsource = len(sourcepars)
initial, _ = make_image(scene, stamps[0], Theta=theta_init)
#sys.exit()
# --------------------------------
# --- Show model and data ---
if True:
from phoplot import plot_model_images
fig, axes = plot_model_images(ptrue, scene, stamps)
pl.show()
# --------------------------------
# --- Priors ---
plate_scale = np.linalg.eigvals(np.linalg.inv(stamps[0].dpix_dsky))
plate_scale = np.abs(plate_scale).mean()
upper = [[12.0, s[1] + 3 * plate_scale, s[2] + 3 * plate_scale, 1.0, np.pi/2, 5.0, 0.12]
def plot_images(pos,
scene,
stamps,
source_idx=None,
x=slice(None),
y=slice(None),
dr=1,
dd=1,
axes=None,
colorbars=True,
share=True,
scale_model=False,
scale_residuals=False,
nchi=-1):
# --- Set up axes ---
vals = pos
same_scale = [False, scale_model, scale_residuals, nchi]
if axes is None:
figsize = (3.3 * len(stamps) + 0.5, 14)
rfig, raxes = pl.subplots(4,
len(stamps),
figsize=figsize,
sharex=share,
sharey=share)
raxes = raxes.T
else:
rfig, raxes = None, axes
raxes = np.atleast_2d(raxes)
# --- Restrict to pixels around source ---
scene.set_all_source_params(vals)
if source_idx is not None:
source = scene.sources[source_idx]
sky = np.array([source.ra, source.dec])
for i, stamp in enumerate(stamps):
if source_idx is not None:
xc, yc = stamp.sky_to_pix(sky)
dx, dy = np.abs(np.dot(stamp.scale, np.array([dr, dd])))
x = slice(int(xc - dx / 2), int(xc + dx / 2))
y = slice(int(yc - dy / 2), int(yc + dy / 2))
data = stamp.pixel_values
model, grad = make_image(scene, stamp, Theta=vals)
resid = data - model
chi = resid * stamp.ierr.reshape(stamp.nx, stamp.ny)
ims = [data, model, resid, chi]
for j, im in enumerate(ims):
if (same_scale[j] == 1) and colorbars:
vmin, vmax = cb.vmin, cb.vmax
elif (same_scale[j] > 0) and colorbars:
vmin, vmax = -same_scale[j], same_scale[j]
else:
vmin = vmax = None
ci = raxes[i, j].imshow(im[x, y].T,
origin='lower',
vmin=vmin,
vmax=vmax)
if colorbars:
cb = pl.colorbar(ci, ax=raxes[i, j], orientation='horizontal')
cb.ax.set_xticklabels(cb.ax.get_xticklabels(),
fontsize=8,
rotation=-50)
text = stamp.filtername
ax = raxes[i, 1]
ax.text(0.6, 0.1, text, transform=ax.transAxes, fontsize=10)
return rfig, raxes
nll = argfix(negative_lnlike_multi, scene=scene, plans=plans)
nll_nograd = argfix(negative_lnlike_multi,
scene=scene,
plans=plans,
grad=False)
# --- Initialize ---
theta_init = np.array([stamp.pixel_values.sum() * 1.0, ra_init, dec_init])
if inpixels:
world = np.array([ra_init, dec_init, 0])
hdr = fits.getheader(imname)
ast = apy_wcs.WCS(hdr)
center = ast.wcs_world2pix(world[None, :], 0)[0, :2] - stamp.lo
theta_init = np.array(
[stamp.pixel_values.sum() * 0.5, center[0], center[1]])
image_init, partials = make_image(scene, stamp, Theta=theta_init)
# --- Plot initial value ---
if True:
fig, axes = pl.subplots(3, 2, sharex=True, sharey=True)
ax = axes.flat[0]
i = ax.imshow(stamp.pixel_values.T, origin='lower')
ax.text(0.1, 0.9, 'Sim. Data', transform=ax.transAxes)
ax = axes.flat[1]
i = ax.imshow(image_init.T, origin='lower')
ax.text(0.1, 0.9, 'Initial Model', transform=ax.transAxes)
for i, ddtheta in enumerate(partials):
ax = axes.flat[i + 2]
ax.imshow(ddtheta.reshape(stamp.nx, stamp.ny).T, origin='lower')
ax.text(0.1,
0.9,
if __name__ == "__main__":
args = parser.parse_args()
plotname = os.path.basename(args.results_name) # filename for plots
galaxy = args.source_type == "galaxy"
# --- Setup Scene and Stamp ---
blob = setup_scene(galaxy=galaxy, fwhm=args.fwhm, size=args.size,
peak_snr=args.peak_snr, add_noise=args.add_noise)
scene, stamp, ptrue, label = blob
plans = [WorkPlan(stamp)]
plate_scale = np.abs(np.linalg.eigvals(np.linalg.inv(stamp.dpix_dsky)))
lower, upper = get_bounds(scene, plate_scale=plate_scale)
true_image, partials = make_image(scene, stamp, Theta=ptrue)
# ---- Plot mock image and gradients thereof -----
if args.show_grad:
fig, axes = pl.subplots(3, 2)
for i, ddtheta in enumerate(partials):
ax = axes.flat[i+1]
c = ax.imshow(ddtheta.reshape(stamp.nx, stamp.ny).T, origin='lower')
ax.text(0.1, 0.85, '$\partial I/\partial${}'.format(label[i]), transform=ax.transAxes)
fig.colorbar(c, ax=ax)
ax = axes.flat[0]
c = ax.imshow(true_image.T, origin='lower')
ax.text(0.1, 0.85, 'Mock (I)'.format(label[i]), transform=ax.transAxes)
fig.colorbar(c, ax=ax)
def plot_model_images(pos,
scene,
stamps,
axes=None,
colorbars=True,
x=slice(None),
y=slice(None),
share=True,
scale_model=False,
scale_residuals=False,
nchi=-1):
vals = pos
same_scale = [False, scale_model, scale_residuals, nchi]
if axes is None:
figsize = (3.3 * len(stamps) + 0.5, 14)
rfig, raxes = pl.subplots(4,
len(stamps),
figsize=figsize,
sharex=share,
sharey=share)
raxes = raxes.T
else:
rfig, raxes = None, axes
raxes = np.atleast_2d(raxes)
for i, stamp in enumerate(stamps):
data = stamp.pixel_values
im, grad = make_image(scene, stamp, Theta=vals)
resid = data - im
chi = resid * stamp.ierr.reshape(stamp.nx, stamp.ny)
ims = [data, im, resid, chi]
for j, ii in enumerate(ims):
ax = raxes[i, j]
if (same_scale[j] == 1) and colorbars:
vmin, vmax = cb.vmin, cb.vmax
elif (same_scale[j] > 0) and colorbars:
vmin, vmax = -same_scale[j], same_scale[j]
else:
vmin = vmax = None
ci = ax.imshow(ii[x, y].T, origin='lower', vmin=vmin, vmax=vmax)
_ = [
tick.label.set_fontsize(10)
for tick in ax.xaxis.get_major_ticks()
]
_ = [
tick.label.set_fontsize(10)
for tick in ax.yaxis.get_major_ticks()
]
if (rfig is not None) & colorbars:
cb = rfig.colorbar(ci,
ax=raxes[i, j],
orientation='horizontal')
cb.ax.set_xticklabels(cb.ax.get_xticklabels(),
rotation=-55,
fontsize=8)
text = stamp.filtername
ax = raxes[i, 1]
ax.text(0.6, 0.1, text, transform=ax.transAxes, fontsize=10)
raxes[i, 0].set_title(text, fontsize=14)
labels = ['Data', 'Model', 'Data-Model', r"$\chi$"]
_ = [
ax.set_ylabel(labels[i], rotation=60, labelpad=20)
for i, ax in enumerate(raxes[0, :])
]
return rfig, raxes
def residuals(self, z):
self.scene.set_all_source_params(z)
self.residuals = [make_image(self.scene, plan) for plan in self.plans]
return self.residuals
all_results.append(result)
ndim = len(vals)
if showfit:
sz = 6 * 3 + 2, len(stamps) * 2.5 # figure size
rfig, raxes = pl.subplots(len(stamps),
6,
sharex=True,
sharey=True,
figsize=sz)
raxes = np.atleast_2d(raxes)
for j, stamp in enumerate(stamps):
err = 1. / stamp.ierr.reshape(stamp.nx, stamp.ny)
snr = stamp.pixel_values * stamp.ierr.reshape(
stamp.nx, stamp.ny)
model, grad = make_image(scene, stamp, Theta=vals)
data = stamp.pixel_values
resid = (data - model)
chi = resid / err
imlist = [err, snr, data, model, resid, chi]
for ax, im in zip(raxes[j, :], imlist):
cc = ax.imshow(im[20:35, 20:35].T, origin='lower')
rfig.colorbar(cc, ax=raxes[j, :].tolist())
raxes[0,
0].text(0.6,
0.75,
"id={:3.0f}\nmag={:2.1f}".format(c["id"], c["mag"]),
transform=raxes[0, 0].transAxes)
pdf.savefig(rfig)
pl.close(rfig)
plans = [WorkPlan(stamp) for stamp in stamps]
# --- show model ----
if args.show_model:
rfig, raxes = pl.subplots(len(filters), len(offsets), sharex=True, sharey=True)
raxes = np.atleast_2d(raxes)
for i, stamp in enumerate(stamps):
raxes.flat[i].imshow(stamp.pixel_values.T, origin='lower')
pl.show()
# --- Gradient Check ---
if args.test_grad:
delta = np.ones_like(ptrue) * 1e-7
#numerical
grad_num = numerical_image_gradients(ptrue, delta, scene, stamp)
image, grad = make_image(scene, stamps[0], Theta=ptrue)
fig, axes = pl.subplots(len(ptrue), 3, sharex=True, sharey=True)
for i in range(len(ptrue)):
g = grad[i,:].reshape(stamp.nx, stamp.ny)
c = axes[i, 0].imshow(grad_num[i,:,:].T, origin='lower')
fig.colorbar(c, ax=axes[i, 0])
c = axes[i, 1].imshow(g.T, origin='lower')
fig.colorbar(c, ax=axes[i, 1])
c = axes[i, 2].imshow((grad_num[i,:,:] - g).T, origin='lower')
fig.colorbar(c, ax=axes[i, 2])
axes[0, 0].set_title('Numerical')
axes[0, 1].set_title('Analytic')
axes[0, 2].set_title('N - A')
# --- Optimization -----
def get_image(x, y, scene, stamp, **extras):
theta = np.array([1, x, y])
im, partials = make_image(scene, stamp, Theta=theta, **extras)
return im
#sys.exit()
# --- Set up posterior prob fns ----
plans = [WorkPlan(stamp) for stamp in stamps]
nll = argfix(negative_lnlike_multi, scene=scene, plans=plans)
#upper = np.zeros(5) + 1000
#lower = np.zeros(5) - 1000
model = Posterior(scene, plans, upper=upper, lower=lower)
# --- Gradient Check ---
if True:
delta = np.ones_like(ptrue) * 1e-7
#numerical
grad_num = numerical_image_gradients(ptrue, delta, scene, stamp)
image, grad = make_image(scene, stamps[0], Theta=ptrue)
fig, axes = pl.subplots(len(ptrue), 3, sharex=True, sharey=True)
for i in range(len(ptrue)):
g = grad[i, :].reshape(stamp.nx, stamp.ny)
c = axes[i, 0].imshow(grad_num[i, :, :].T, origin='lower')
fig.colorbar(c, ax=axes[i, 0])
c = axes[i, 1].imshow(g.T, origin='lower')
fig.colorbar(c, ax=axes[i, 1])
c = axes[i, 2].imshow((grad_num[i, :, :] - g).T, origin='lower')
fig.colorbar(c, ax=axes[i, 2])
axes[0, 0].set_title('Numerical')
axes[0, 1].set_title('Analytic')
axes[0, 2].set_title('N - A')
# --- Optimization -----
