def test_sort_img_index():
ximg, yimg = np.array([1, 2, 3, 4]), np.array([0, 0, 1, 2])
xref, yref = np.array([2, 3, 1, 4]), np.array([0, 1, 0, 2])
indexes = util.sort_image_index(ximg, yimg, xref, yref)
npt.assert_allclose(ximg[indexes], xref)
npt.assert_allclose(yimg[indexes], yref)
xref, yref = np.array([2, 3, 1, 4]), np.array([0, 1, 0, 2])
indexes = util.sort_image_index(xref, yref, xref, yref)
npt.assert_allclose(np.array(indexes), [0, 1, 2, 3])
def _image_position_penalty(self, lens_args_tovary, chi_mode):
self.betax,self.betay = self.lensing._ray_shooting_fast(lens_args_tovary)
dx_source = ((self.betax[0] - self.betax[1]) ** 2 + (self.betax[0] - self.betax[2]) ** 2 + (
self.betax[0] - self.betax[3]) ** 2 + (
self.betax[1] - self.betax[2]) ** 2 +
(self.betax[1] - self.betax[3]) ** 2 + (self.betax[2] - self.betax[3]) ** 2)
dy_source = ((self.betay[0] - self.betay[1]) ** 2 + (self.betay[0] - self.betay[2]) ** 2 + (
self.betay[0] - self.betay[3]) ** 2 + (
self.betay[1] - self.betay[2]) ** 2 +
(self.betay[1] - self.betay[3]) ** 2 + (self.betay[2] - self.betay[3]) ** 2)
src_plane_pen = 0.5 * (dx_source + dy_source) * self.tol_source ** -2
if chi_mode == 0:
return src_plane_pen
elif src_plane_pen > self._pso_convergence_mean:
return src_plane_pen
else:
# compute the positions in the image plane
kwargs_lens_final = lens_args_tovary + self.param_class.argsfixed_todictionary()
x_image, y_image = self._solver.findBrightImage(np.mean(self.betax), np.mean(self.betay),
kwargs_lens_final, arrival_time_sort=False)
# if we have the wrong number of images, use the source plane chi^2
if len(x_image) != len(self._ximage) or len(y_image) != len(self._yimage):
return src_plane_pen
# compute the image plane chi^2
else:
# try to match image positions, compute chi^2
inds = sort_image_index(x_image, y_image, self._ximage, self._yimage)
dx = (x_image[inds] - self._ximage) ** 2
dy = (y_image[inds] - self._yimage) ** 2
image_plane_penalty = 0
if isinstance(self.tol_image, list) or isinstance(self.tol_image, np.ndarray):
for i, dx_dy in enumerate(self.tol_image):
image_plane_penalty += 0.5 * (dx[i] + dy[i]) * self.tol_image[i] ** -2
else:
image_plane_penalty = np.sum(0.5 * (dx + dy) * self.tol_image ** -2)
return image_plane_penalty