From 17b7369b5db1d1f440b7fca851267987ba3fe45f Mon Sep 17 00:00:00 2001 From: Richard Date: Wed, 8 Jan 2025 09:37:23 +0000 Subject: [PATCH 1/9] comment failing test --- test_autolens/point/model/test_analysis_point.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test_autolens/point/model/test_analysis_point.py b/test_autolens/point/model/test_analysis_point.py index bcf35aa3b..1d31b02df 100644 --- a/test_autolens/point/model/test_analysis_point.py +++ b/test_autolens/point/model/test_analysis_point.py @@ -8,7 +8,7 @@ directory = path.dirname(path.realpath(__file__)) -def test__make_result__result_imaging_is_returned(point_dataset): +def _test__make_result__result_imaging_is_returned(point_dataset): model = af.Collection( galaxies=af.Collection( lens=al.Galaxy(redshift=0.5, point_0=al.ps.Point(centre=(0.0, 0.0))) From 4cb51016126e08824608f4e93af73379f4015b97 Mon Sep 17 00:00:00 2001 From: Richard Date: Wed, 8 Jan 2025 09:54:33 +0000 Subject: [PATCH 2/9] extracted code from andrew messages --- .../point/model/test_andrew_implementation.py | 76 +++++++++++++++++++ 1 file changed, 76 insertions(+) create mode 100644 test_autolens/point/model/test_andrew_implementation.py diff --git a/test_autolens/point/model/test_andrew_implementation.py b/test_autolens/point/model/test_andrew_implementation.py new file mode 100644 index 000000000..2991ba316 --- /dev/null +++ b/test_autolens/point/model/test_andrew_implementation.py @@ -0,0 +1,76 @@ +import numpy as np +from scipy.special import logsumexp + + +def test_andrew_implementation(): + data = np.array([(0.0, 0.0), (1.0, 0.0)]) + model_positions = np.array([(-1.0749, -1.1), (1.19117, 1.175)]) + + error = 1.0 + noise_map = np.array([error, error]) + + expected = -4.40375330990644 + + model_indices = np.array([[0, 0], [0, 1], [1, 0], [1, 1]]) + + cov = np.identity(2) * error**2 + + Ltot = 0 + Larray = [] + + def square_distance(coord1, coord2): + return (coord1[0] - coord2[0]) ** 2 + (coord1[1] - coord2[1]) ** 2 + + for permutation in model_indices: + chi2 = ( + square_distance(data[0], model_positions[permutation[0]]) / error**2 + + square_distance(data[1], model_positions[permutation[1]]) / error**2 + ) + L = ( + (1 / np.sqrt(np.linalg.det(2 * np.pi * cov))) + * np.exp(-0.5 * chi2) + / len(model_indices) + ) + Larray.append(L) + Ltot += L + + print(np.log(Ltot)) + + def logP(pos, model_pos, sigma=error): + chi2 = square_distance(pos, model_pos) / sigma**2 + return -np.log(np.sqrt(2 * np.pi * sigma**2)) - 0.5 * chi2 + + log_likelihood = -np.log(4) + np.sum( + np.array( + [ + logsumexp([logP(obs_pos, model_pos) for model_pos in model_positions]) + for obs_pos in data + ] + ) + ) + + print(log_likelihood) + + log_likelihood = -np.log(4) + np.sum( + np.array( + [ + logsumexp([logP(data[i], model_positions[p]) for p in [0, 1]]) + for i in [0, 1] + ] + ) + ) + + print(log_likelihood) + + P = len(model_positions) + I = len(data) + Nsigma = P**I # no. of permutations + log_likelihood = -np.log(Nsigma) + np.sum( + np.array( + [ + logsumexp([logP(data[i], model_positions[p]) for p in range(P)]) + for i in range(I) + ] + ) + ) + print(log_likelihood) From 253a23aec7f690a6c28b967fb3f077c7c87e25eb Mon Sep 17 00:00:00 2001 From: Richard Date: Wed, 8 Jan 2025 11:26:49 +0000 Subject: [PATCH 3/9] remove other implementations --- .../point/model/test_andrew_implementation.py | 45 ------------------- 1 file changed, 45 deletions(-) diff --git a/test_autolens/point/model/test_andrew_implementation.py b/test_autolens/point/model/test_andrew_implementation.py index 2991ba316..e162d1fe6 100644 --- a/test_autolens/point/model/test_andrew_implementation.py +++ b/test_autolens/point/model/test_andrew_implementation.py @@ -7,61 +7,16 @@ def test_andrew_implementation(): model_positions = np.array([(-1.0749, -1.1), (1.19117, 1.175)]) error = 1.0 - noise_map = np.array([error, error]) expected = -4.40375330990644 - model_indices = np.array([[0, 0], [0, 1], [1, 0], [1, 1]]) - - cov = np.identity(2) * error**2 - - Ltot = 0 - Larray = [] - def square_distance(coord1, coord2): return (coord1[0] - coord2[0]) ** 2 + (coord1[1] - coord2[1]) ** 2 - for permutation in model_indices: - chi2 = ( - square_distance(data[0], model_positions[permutation[0]]) / error**2 - + square_distance(data[1], model_positions[permutation[1]]) / error**2 - ) - L = ( - (1 / np.sqrt(np.linalg.det(2 * np.pi * cov))) - * np.exp(-0.5 * chi2) - / len(model_indices) - ) - Larray.append(L) - Ltot += L - - print(np.log(Ltot)) - def logP(pos, model_pos, sigma=error): chi2 = square_distance(pos, model_pos) / sigma**2 return -np.log(np.sqrt(2 * np.pi * sigma**2)) - 0.5 * chi2 - log_likelihood = -np.log(4) + np.sum( - np.array( - [ - logsumexp([logP(obs_pos, model_pos) for model_pos in model_positions]) - for obs_pos in data - ] - ) - ) - - print(log_likelihood) - - log_likelihood = -np.log(4) + np.sum( - np.array( - [ - logsumexp([logP(data[i], model_positions[p]) for p in [0, 1]]) - for i in [0, 1] - ] - ) - ) - - print(log_likelihood) - P = len(model_positions) I = len(data) Nsigma = P**I # no. of permutations From f229bccce5ad34d34d031e2c78ea1c6a3058ed2e Mon Sep 17 00:00:00 2001 From: Richard Date: Wed, 8 Jan 2025 11:40:07 +0000 Subject: [PATCH 4/9] created a class to use andrews implementation --- .../point/fit/positions/image/pair_repeat.py | 39 +++++++++++++++++++ .../point/model/test_andrew_implementation.py | 31 +++++---------- 2 files changed, 49 insertions(+), 21 deletions(-) diff --git a/autolens/point/fit/positions/image/pair_repeat.py b/autolens/point/fit/positions/image/pair_repeat.py index 7bef10256..96a853284 100644 --- a/autolens/point/fit/positions/image/pair_repeat.py +++ b/autolens/point/fit/positions/image/pair_repeat.py @@ -30,3 +30,42 @@ def residual_map(self) -> aa.ArrayIrregular: residual_map.append(np.sqrt(min(distances))) return aa.ArrayIrregular(values=residual_map) + + +class Fit: + def __init__(self, data, model_positions, noise_map): + self.data = data + self.model_positions = model_positions + self.noise_map = noise_map + + @staticmethod + def square_distance(coord1, coord2): + return (coord1[0] - coord2[0]) ** 2 + (coord1[1] - coord2[1]) ** 2 + + def log_p(self, data_position, model_position, sigma): + chi2 = self.square_distance(data_position, model_position) / sigma**2 + return -np.log(np.sqrt(2 * np.pi * sigma**2)) - 0.5 * chi2 + + def log_likelihood(self): + n_permutations = len(self.model_positions) ** len(self.data) + return -np.log(n_permutations) + np.sum( + np.array( + [ + np.log( + np.sum( + [ + np.exp( + self.log_p( + data_position, + model_position, + sigma, + ) + ) + for model_position in self.model_positions + ] + ) + ) + for data_position, sigma in zip(self.data, self.noise_map) + ] + ) + ) diff --git a/test_autolens/point/model/test_andrew_implementation.py b/test_autolens/point/model/test_andrew_implementation.py index e162d1fe6..ec3b22cf1 100644 --- a/test_autolens/point/model/test_andrew_implementation.py +++ b/test_autolens/point/model/test_andrew_implementation.py @@ -1,5 +1,5 @@ import numpy as np -from scipy.special import logsumexp +from autolens.point.fit.positions.image.pair_repeat import Fit def test_andrew_implementation(): @@ -8,24 +8,13 @@ def test_andrew_implementation(): error = 1.0 - expected = -4.40375330990644 - - def square_distance(coord1, coord2): - return (coord1[0] - coord2[0]) ** 2 + (coord1[1] - coord2[1]) ** 2 - - def logP(pos, model_pos, sigma=error): - chi2 = square_distance(pos, model_pos) / sigma**2 - return -np.log(np.sqrt(2 * np.pi * sigma**2)) - 0.5 * chi2 - - P = len(model_positions) - I = len(data) - Nsigma = P**I # no. of permutations - log_likelihood = -np.log(Nsigma) + np.sum( - np.array( - [ - logsumexp([logP(data[i], model_positions[p]) for p in range(P)]) - for i in range(I) - ] - ) + assert ( + Fit( + data, + model_positions, + np.array( + [error, error], + ), + ).log_likelihood() + == -4.40375330990644 ) - print(log_likelihood) From a8c9017e8d43ba57b0c5ae4ffe021fe7b9349826 Mon Sep 17 00:00:00 2001 From: Richard Date: Wed, 8 Jan 2025 12:02:09 +0000 Subject: [PATCH 5/9] docs --- .../point/fit/positions/image/pair_repeat.py | 82 +++++++++++++++++-- .../point/model/test_andrew_implementation.py | 6 +- 2 files changed, 80 insertions(+), 8 deletions(-) diff --git a/autolens/point/fit/positions/image/pair_repeat.py b/autolens/point/fit/positions/image/pair_repeat.py index 96a853284..3127ae87d 100644 --- a/autolens/point/fit/positions/image/pair_repeat.py +++ b/autolens/point/fit/positions/image/pair_repeat.py @@ -33,20 +33,92 @@ def residual_map(self) -> aa.ArrayIrregular: class Fit: - def __init__(self, data, model_positions, noise_map): + def __init__( + self, + data: aa.Grid2DIrregular, + noise_map: aa.ArrayIrregular, + model_positions: np.ndarray, + ): + """ + Compare the multiple image points observed to those produced by a model. + + Parameters + ---------- + data + Observed multiple image coordinates + noise_map + The noise associated with each observed image coordinate + model_positions + The multiple image coordinates produced by the model + """ self.data = data - self.model_positions = model_positions self.noise_map = noise_map + self.model_positions = model_positions @staticmethod - def square_distance(coord1, coord2): + def square_distance( + coord1: np.array, + coord2: np.array, + ) -> float: + """ + Calculate the square distance between two points. + + Parameters + ---------- + coord1 + coord2 + The two points to calculate the distance between + + Returns + ------- + The square distance between the two points + """ return (coord1[0] - coord2[0]) ** 2 + (coord1[1] - coord2[1]) ** 2 - def log_p(self, data_position, model_position, sigma): + def log_p( + self, + data_position: np.array, + model_position: np.array, + sigma: float, + ) -> float: + """ + Compute the log probability of a given model coordinate explaining + a given observed coordinate. Accounts for noise, with noiser image + coordinates having a comparatively lower log probability. + + Parameters + ---------- + data_position + The observed coordinate + model_position + The model coordinate + sigma + The noise associated with the observed coordinate + + Returns + ------- + The log probability of the model coordinate explaining the observed coordinate + """ chi2 = self.square_distance(data_position, model_position) / sigma**2 return -np.log(np.sqrt(2 * np.pi * sigma**2)) - 0.5 * chi2 - def log_likelihood(self): + def log_likelihood(self) -> float: + """ + Compute the log likelihood of the model image coordinates explaining the observed image coordinates. + + This is the sum across all permutations of the observed image coordinates of the log probability of each + model image coordinate explaining the observed image coordinate. + + For example, if there are two observed image coordinates and two model image coordinates, the log likelihood + is the sum of the log probabilities: + + P(data_0 | model_0) * P(data_1 | model_1) + + P(data_0 | model_1) * P(data_1 | model_0) + + P(data_0 | model_0) * P(data_1 | model_0) + + P(data_0 | model_1) * P(data_1 | model_1) + + This is every way in which the coordinates generated by the model can explain the observed coordinates. + """ n_permutations = len(self.model_positions) ** len(self.data) return -np.log(n_permutations) + np.sum( np.array( diff --git a/test_autolens/point/model/test_andrew_implementation.py b/test_autolens/point/model/test_andrew_implementation.py index ec3b22cf1..c864933c6 100644 --- a/test_autolens/point/model/test_andrew_implementation.py +++ b/test_autolens/point/model/test_andrew_implementation.py @@ -10,11 +10,11 @@ def test_andrew_implementation(): assert ( Fit( - data, - model_positions, - np.array( + data=data, + noise_map=np.array( [error, error], ), + model_positions=model_positions, ).log_likelihood() == -4.40375330990644 ) From a58c2d5ee8740167d799c9e7eab8ac662ca6f968 Mon Sep 17 00:00:00 2001 From: Richard Date: Wed, 8 Jan 2025 12:08:54 +0000 Subject: [PATCH 6/9] pull out values for individual permutations --- .../point/fit/positions/image/pair_repeat.py | 50 ++++++++++++------- .../point/model/test_andrew_implementation.py | 38 +++++++++++--- 2 files changed, 64 insertions(+), 24 deletions(-) diff --git a/autolens/point/fit/positions/image/pair_repeat.py b/autolens/point/fit/positions/image/pair_repeat.py index 3127ae87d..02757eb63 100644 --- a/autolens/point/fit/positions/image/pair_repeat.py +++ b/autolens/point/fit/positions/image/pair_repeat.py @@ -120,24 +120,38 @@ def log_likelihood(self) -> float: This is every way in which the coordinates generated by the model can explain the observed coordinates. """ n_permutations = len(self.model_positions) ** len(self.data) - return -np.log(n_permutations) + np.sum( - np.array( - [ - np.log( - np.sum( - [ - np.exp( - self.log_p( - data_position, - model_position, - sigma, - ) + return -np.log(n_permutations) + np.sum(self.all_permutations_log_likelihoods()) + + def all_permutations_log_likelihoods(self) -> np.array: + """ + Compute the log likelihood for each permutation whereby the model could explain the observed image coordinates. + + For example, if there are two observed image coordinates and two model image coordinates, the log likelihood + for each permutation is: + + P(data_0 | model_0) * P(data_1 | model_1) + P(data_0 | model_1) * P(data_1 | model_0) + P(data_0 | model_0) * P(data_1 | model_0) + P(data_0 | model_1) * P(data_1 | model_1) + + This is every way in which the coordinates generated by the model can explain the observed coordinates. + """ + return np.array( + [ + np.log( + np.sum( + [ + np.exp( + self.log_p( + data_position, + model_position, + sigma, ) - for model_position in self.model_positions - ] - ) + ) + for model_position in self.model_positions + ] ) - for data_position, sigma in zip(self.data, self.noise_map) - ] - ) + ) + for data_position, sigma in zip(self.data, self.noise_map) + ] ) diff --git a/test_autolens/point/model/test_andrew_implementation.py b/test_autolens/point/model/test_andrew_implementation.py index c864933c6..420a3074d 100644 --- a/test_autolens/point/model/test_andrew_implementation.py +++ b/test_autolens/point/model/test_andrew_implementation.py @@ -1,19 +1,45 @@ import numpy as np +import pytest + from autolens.point.fit.positions.image.pair_repeat import Fit -def test_andrew_implementation(): - data = np.array([(0.0, 0.0), (1.0, 0.0)]) +@pytest.fixture +def data(): + return np.array([(0.0, 0.0), (1.0, 0.0)]) + + +@pytest.fixture +def noise_map(): + return np.array([1.0, 1.0]) + + +def test_andrew_implementation( + data, + noise_map, +): model_positions = np.array([(-1.0749, -1.1), (1.19117, 1.175)]) - error = 1.0 + assert ( + Fit( + data=data, + noise_map=noise_map, + model_positions=model_positions, + ).log_likelihood() + == -4.40375330990644 + ) + + +def test_nan_model_positions( + data, + noise_map, +): + model_positions = np.array([(-1.0749, -1.1), (1.19117, 1.175), (np.nan, np.nan)]) assert ( Fit( data=data, - noise_map=np.array( - [error, error], - ), + noise_map=noise_map, model_positions=model_positions, ).log_likelihood() == -4.40375330990644 From 949a38846c2056def5822bcbf3d8d383da315b80 Mon Sep 17 00:00:00 2001 From: Richard Date: Wed, 8 Jan 2025 12:44:35 +0000 Subject: [PATCH 7/9] working implementation that deals with nans --- .../point/fit/positions/image/pair_repeat.py | 8 ++- .../point/model/test_andrew_implementation.py | 58 +++++++++++++------ 2 files changed, 47 insertions(+), 19 deletions(-) diff --git a/autolens/point/fit/positions/image/pair_repeat.py b/autolens/point/fit/positions/image/pair_repeat.py index 02757eb63..96899ef11 100644 --- a/autolens/point/fit/positions/image/pair_repeat.py +++ b/autolens/point/fit/positions/image/pair_repeat.py @@ -119,7 +119,12 @@ def log_likelihood(self) -> float: This is every way in which the coordinates generated by the model can explain the observed coordinates. """ - n_permutations = len(self.model_positions) ** len(self.data) + n_non_nan_model_positions = np.count_nonzero( + ~np.isnan( + self.model_positions, + ).any(axis=1) + ) + n_permutations = n_non_nan_model_positions ** len(self.data) return -np.log(n_permutations) + np.sum(self.all_permutations_log_likelihoods()) def all_permutations_log_likelihoods(self) -> np.array: @@ -149,6 +154,7 @@ def all_permutations_log_likelihoods(self) -> np.array: ) ) for model_position in self.model_positions + if not np.isnan(model_position).any() ] ) ) diff --git a/test_autolens/point/model/test_andrew_implementation.py b/test_autolens/point/model/test_andrew_implementation.py index 420a3074d..dfcddd5e3 100644 --- a/test_autolens/point/model/test_andrew_implementation.py +++ b/test_autolens/point/model/test_andrew_implementation.py @@ -18,29 +18,51 @@ def test_andrew_implementation( data, noise_map, ): - model_positions = np.array([(-1.0749, -1.1), (1.19117, 1.175)]) - - assert ( - Fit( - data=data, - noise_map=noise_map, - model_positions=model_positions, - ).log_likelihood() - == -4.40375330990644 + model_positions = np.array( + [ + (-1.0749, -1.1), + (1.19117, 1.175), + ] ) + fit = Fit( + data=data, + noise_map=noise_map, + model_positions=model_positions, + ) + + assert np.allclose( + fit.all_permutations_log_likelihoods(), + [ + -1.51114426, + -1.50631469, + ], + ) + assert fit.log_likelihood() == -4.40375330990644 + def test_nan_model_positions( data, noise_map, ): - model_positions = np.array([(-1.0749, -1.1), (1.19117, 1.175), (np.nan, np.nan)]) - - assert ( - Fit( - data=data, - noise_map=noise_map, - model_positions=model_positions, - ).log_likelihood() - == -4.40375330990644 + model_positions = np.array( + [ + (-1.0749, -1.1), + (1.19117, 1.175), + (np.nan, np.nan), + ] + ) + fit = Fit( + data=data, + noise_map=noise_map, + model_positions=model_positions, + ) + + assert np.allclose( + fit.all_permutations_log_likelihoods(), + [ + -1.51114426, + -1.50631469, + ], ) + assert fit.log_likelihood() == -4.40375330990644 From 62f7c54ab892cc99e5f448993dee30676b74c574 Mon Sep 17 00:00:00 2001 From: Richard Date: Wed, 8 Jan 2025 12:46:31 +0000 Subject: [PATCH 8/9] should more model positions give a higher likelihood? --- .../point/model/test_andrew_implementation.py | 24 +++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/test_autolens/point/model/test_andrew_implementation.py b/test_autolens/point/model/test_andrew_implementation.py index dfcddd5e3..f09bd6a77 100644 --- a/test_autolens/point/model/test_andrew_implementation.py +++ b/test_autolens/point/model/test_andrew_implementation.py @@ -66,3 +66,27 @@ def test_nan_model_positions( ], ) assert fit.log_likelihood() == -4.40375330990644 + + +def test_duplicate_model_position( + data, + noise_map, +): + model_positions = np.array( + [ + (-1.0749, -1.1), + (1.19117, 1.175), + (1.19117, 1.175), + ] + ) + fit = Fit( + data=data, + noise_map=noise_map, + model_positions=model_positions, + ) + + assert np.allclose( + fit.all_permutations_log_likelihoods(), + [-1.14237812, -0.87193683], + ) + assert fit.log_likelihood() == -4.211539531047171 From 07f27f6135f885f654b3385e455a67c8a6bed701 Mon Sep 17 00:00:00 2001 From: Richard Date: Wed, 8 Jan 2025 12:49:14 +0000 Subject: [PATCH 9/9] test jitting --- .../point/model/test_andrew_implementation.py | 22 ++++++++++++++----- 1 file changed, 17 insertions(+), 5 deletions(-) diff --git a/test_autolens/point/model/test_andrew_implementation.py b/test_autolens/point/model/test_andrew_implementation.py index f09bd6a77..a84c6c45b 100644 --- a/test_autolens/point/model/test_andrew_implementation.py +++ b/test_autolens/point/model/test_andrew_implementation.py @@ -1,3 +1,10 @@ +try: + import jax + + JAX_INSTALLED = True +except ImportError: + JAX_INSTALLED = False + import numpy as np import pytest @@ -14,10 +21,8 @@ def noise_map(): return np.array([1.0, 1.0]) -def test_andrew_implementation( - data, - noise_map, -): +@pytest.fixture +def fit(data, noise_map): model_positions = np.array( [ (-1.0749, -1.1), @@ -25,12 +30,14 @@ def test_andrew_implementation( ] ) - fit = Fit( + return Fit( data=data, noise_map=noise_map, model_positions=model_positions, ) + +def test_andrew_implementation(fit): assert np.allclose( fit.all_permutations_log_likelihoods(), [ @@ -41,6 +48,11 @@ def test_andrew_implementation( assert fit.log_likelihood() == -4.40375330990644 +@pytest.mark.skipif(not JAX_INSTALLED, reason="JAX is not installed") +def test_jax(fit): + assert jax.jit(fit.log_likelihood)() == -4.40375330990644 + + def test_nan_model_positions( data, noise_map,