From 6543a119e8e2cd5d83b12c3d048b806e4d01e27b Mon Sep 17 00:00:00 2001 From: chaitanya chawak Date: Wed, 9 Sep 2026 18:46:07 +0200 Subject: [PATCH 01/15] abstract out estimate_noise_sigma method from calculate_sample_weights + unit tests --- .../tests/test_training/train_utils_test.py | 53 ++++++++++ src/wf_psf/training/train_utils.py | 96 +++++++++++++++---- 2 files changed, 130 insertions(+), 19 deletions(-) diff --git a/src/wf_psf/tests/test_training/train_utils_test.py b/src/wf_psf/tests/test_training/train_utils_test.py index db05d8f3..09921799 100644 --- a/src/wf_psf/tests/test_training/train_utils_test.py +++ b/src/wf_psf/tests/test_training/train_utils_test.py @@ -228,6 +228,59 @@ def test_calculate_sample_weights_high_variance(mock_noise_estimator, loss): assert np.all(weights >= 0), f"Weights are negative: {weights}" +@pytest.mark.parametrize( + "loss", [None, "mean_squared_error", "masked_mean_squared_error"] +) +def test_estimate_noise_sigma_shape_and_positivity(loss): + """estimate_noise_sigma returns one positive sigma per observation.""" + batch_size, height, width = 5, 32, 32 + + if loss == "masked_mean_squared_error": + outputs = np.random.rand(batch_size, height, width, 2) + outputs[..., 1] = np.random.randint(0, 2, size=(batch_size, height, width)) + else: + outputs = np.random.rand(batch_size, height, width) + + sigma = train_utils.estimate_noise_sigma(outputs, loss) + + assert isinstance(sigma, np.ndarray) + assert sigma.shape == (batch_size,) + assert np.all(sigma > 0) + + +@pytest.mark.parametrize( + "loss", [None, "mean_squared_error", "masked_mean_squared_error"] +) +def test_calculate_sample_weights_uses_estimate_noise_sigma(loss): + """calculate_sample_weights delegates noise estimation to estimate_noise_sigma.""" + batch_size, height, width = 5, 32, 32 + + if loss == "masked_mean_squared_error": + outputs = np.random.rand(batch_size, height, width, 2) + outputs[..., 1] = np.random.randint(0, 2, size=(batch_size, height, width)) + else: + outputs = np.random.rand(batch_size, height, width) + + fake_sigma = np.linspace(1.0, 2.0, batch_size) + with patch.object( + train_utils, "estimate_noise_sigma", return_value=fake_sigma + ) as mock_sigma: + result = train_utils.calculate_sample_weights( + outputs, use_sample_weights=True, loss=loss + ) + + mock_sigma.assert_called_once() + # Passed the same outputs and loss through + called_outputs, called_loss = mock_sigma.call_args[0][:2] + assert called_outputs is outputs + assert called_loss == loss + + # Weights are the median-normalised inverse variance of the returned sigma + expected = 1 / fake_sigma**2 + expected /= np.median(expected) + np.testing.assert_allclose(result, expected) + + @pytest.mark.parametrize( "loss_function, metric_function", [ diff --git a/src/wf_psf/training/train_utils.py b/src/wf_psf/training/train_utils.py index 16deca0d..c41ef393 100644 --- a/src/wf_psf/training/train_utils.py +++ b/src/wf_psf/training/train_utils.py @@ -367,6 +367,81 @@ def configure_optimizer_and_loss( return optimizer, loss, metrics +def _is_masked_loss(loss: Union[str, Callable, None]) -> bool: + """ + Check whether the given loss corresponds to the masked mean squared error. + + Parameters + ---------- + loss: str, callable, optional + The loss function (or its name) used for training. + + Returns + ------- + bool + True if the loss is the ``masked_mean_squared_error``, False otherwise. + """ + if loss is None: + return False + return (isinstance(loss, str) and loss == "masked_mean_squared_error") or ( + hasattr(loss, "name") and loss.name == "masked_mean_squared_error" + ) + + +def estimate_noise_sigma( + outputs: tf.Tensor, + loss: Union[str, Callable, None] = None, +) -> np.ndarray: + """ + Estimate the per-observation noise standard deviation. + + For each image in ``outputs``, the noise standard deviation is estimated with + :class:`wf_psf.utils.utils.NoiseEstimator`, which applies a robust MAD-based + estimator on the unmasked background pixels (i.e. pixels outside a central + exclusion window, and outside the per-image mask when a masked loss is used). + + This helper isolates the noise-estimation step so that the resulting sigma can + be reused independently of the training-specific sample-weight logic (e.g. for + goodness-of-fit metrics such as reduced chi-squared in validation workflows). + + Parameters + ---------- + outputs: tf.Tensor + Image data. When ``loss`` is ``"masked_mean_squared_error"``, a 4D tensor + of shape ``(batch_size, height, width, 2)`` is expected, where the last + dimension holds ``[image, mask]``. Otherwise, a 3D tensor of shape + ``(batch_size, height, width)`` is expected. + loss: str, callable, optional + The loss function (or its name) used for training. If the loss is + ``"masked_mean_squared_error"``, the per-image mask stored in + ``outputs[..., 1]`` is combined with the exclusion window before + estimating the noise. Default is None. + + Returns + ------- + np.ndarray + A 1D array of shape ``(batch_size,)`` containing the estimated noise + standard deviation for each image. + """ + img_dim = (outputs.shape[1], outputs.shape[2]) + win_rad = np.ceil(outputs.shape[1] / 3.33) + std_est = NoiseEstimator(img_dim=img_dim, win_rad=win_rad) + + if _is_masked_loss(loss): + logger.info("Estimating noise standard deviation for masked images..") + images = outputs[..., 0] + masks = np.array(1 - outputs[..., 1], dtype=bool) + imgs_std = np.array( + [std_est.estimate_noise(_im, _win) for _im, _win in zip(images, masks)] + ) + else: + logger.info("Estimating noise standard deviation for images..") + # Estimate noise standard deviation + imgs_std = np.array([std_est.estimate_noise(_im) for _im in outputs]) + + return imgs_std + + def calculate_sample_weights( outputs: tf.Tensor, use_sample_weights: bool, @@ -405,25 +480,8 @@ def calculate_sample_weights( An array of sample weights, or None if `use_sample_weights` is False. """ if use_sample_weights: - img_dim = (outputs.shape[1], outputs.shape[2]) - win_rad = np.ceil(outputs.shape[1] / 3.33) - std_est = NoiseEstimator(img_dim=img_dim, win_rad=win_rad) - - if loss is not None and ( - (isinstance(loss, str) and loss == "masked_mean_squared_error") - or (hasattr(loss, "name") and loss.name == "masked_mean_squared_error") - ): - - logger.info("Estimating noise standard deviation for masked images..") - images = outputs[..., 0] - masks = np.array(1 - outputs[..., 1], dtype=bool) - imgs_std = np.array( - [std_est.estimate_noise(_im, _win) for _im, _win in zip(images, masks)] - ) - else: - logger.info("Estimating noise standard deviation for images..") - # Estimate noise standard deviation - imgs_std = np.array([std_est.estimate_noise(_im) for _im in outputs]) + # Estimate the per-observation noise standard deviation + imgs_std = estimate_noise_sigma(outputs, loss) # Calculate variances variances = imgs_std**2 From 4e79ee4e8a8aa33f2b94045a12f2b5fc637a7bf9 Mon Sep 17 00:00:00 2001 From: chaitanya chawak Date: Wed, 16 Sep 2026 15:18:01 +0200 Subject: [PATCH 02/15] move NoiseEstimator into separate noise.py --- src/wf_psf/utils/noise.py | 151 ++++++++++++++++++++++++++++++++++++++ src/wf_psf/utils/utils.py | 143 ------------------------------------ 2 files changed, 151 insertions(+), 143 deletions(-) create mode 100644 src/wf_psf/utils/noise.py diff --git a/src/wf_psf/utils/noise.py b/src/wf_psf/utils/noise.py new file mode 100644 index 00000000..c8066e91 --- /dev/null +++ b/src/wf_psf/utils/noise.py @@ -0,0 +1,151 @@ +"""Noise estimation utilities for image data. + +:Authors: Tobias Liaudat + +""" + +import numpy as np + + +class NoiseEstimator: + """ + A class for estimating noise levels in an image. + + Parameters + ---------- + img_dim : tuple of int + The dimensions of the image as (height, width). + win_rad : int + The radius of the exclusion window (in pixels). + """ + + def __init__(self, img_dim: tuple[int, int], win_rad: int) -> None: + """ + Initialize a NoiseEstimator instance. + + This constructor sets up the noise estimator by storing the image dimensions + and exclusion window radius, then initializes the exclusion window mask. + + Parameters + ---------- + img_dim : tuple of int + The dimensions of the image as (height, width). + win_rad : int + The radius of the exclusion window in pixels. Pixels within this radius + of the image center are excluded from noise estimation. + + Notes + ----- + The exclusion window is initialized automatically via _init_window(), creating + a boolean mask where pixels inside the exclusion radius are marked False + (excluded) and pixels outside are marked True (included). + """ + self.img_dim: tuple[int, int] = img_dim + self.win_rad: int = win_rad + + self._init_window() # Initialize self.window + + def _init_window(self): + """ + Initialize the exclusion window mask stored in self.window. + + The mask is a boolean array of shape `self.img_dim` (rows, cols). Pixels + whose Euclidean distance from the image center is less than or equal to + `self.win_rad` are marked False (excluded); all other pixels are True + (included). The mask dtype is `bool`. + + Notes + ----- + - The image center is computed as (rows / 2, cols / 2). This yields a + floating-point center so the distance is computed with sub-pixel + precision; for even dimensions the center lies between pixels. + - The comparison uses "<=" so pixels exactly at distance `win_rad` are + excluded. Change to "<" if you prefer a strict interior exclusion. + - Time complexity is O(rows * cols) for mask construction. + - No return value; the constructed mask is assigned to `self.window`. + """ + self.window = np.ones(self.img_dim, dtype=bool) + + mid_x = self.img_dim[0] / 2 + mid_y = self.img_dim[1] / 2 + + for _x in range(self.img_dim[0]): + for _y in range(self.img_dim[1]): + # If pixel is within the exclusion radius, set it to False + if np.sqrt((_x - mid_x) ** 2 + (_y - mid_y) ** 2) <= self.win_rad: + self.window[_x, _y] = False + + def apply_mask(self, mask: np.ndarray = None) -> np.ndarray: + """ + Apply a given mask to the exclusion window. + + Parameters + ---------- + mask : np.ndarray, optional + A boolean mask to apply to the exclusion window. If None, the exclusion + window is returned without any modification. + + Returns + ------- + np.ndarray + The resulting boolean array after applying the mask to the exclusion window. + """ + if mask is None: + return self.window # Return just the window if no mask is provided + return self.window & mask # Otherwise, apply the mask as usual + + @staticmethod + def sigma_mad(x): + """ + Robustly estimate the standard deviation using the Median Absolute Deviation (MAD). + + Computes MAD = ``median(|x - median(x)|)`` and scales it by 1.4826 to make the + estimator consistent with the standard deviation for a Gaussian distribution: + + sigma ≈ 1.4826 * MAD + + Parameters + ---------- + x : array-like + Input data. The values are flattened before computation. NaNs are not + specially handled and will propagate; remove or mask them prior to + calling if needed. + + Returns + ------- + float + Robust estimate of the standard deviation of the input data. + + Notes + ----- + - The MAD-based estimator is much less sensitive to outliers than the + sample standard deviation, making it appropriate for noisy data with + occasional large deviations. + - The constant 1.4826 is the scaling factor for consistency with the + standard deviation of a normal distribution. + """ + return 1.4826 * np.median(np.abs(x - np.median(x))) + + def estimate_noise(self, image: np.ndarray, mask: np.ndarray = None) -> float: + """ + Estimates the noise level of an image using the MAD estimator. + + Parameters + ---------- + image : np.ndarray + The input image for noise estimation. + mask : np.ndarray, optional + A boolean mask specifying which pixels to include in the noise estimation. + If None, the default exclusion window is used. The mask should have the same shape as `image`. + + Returns + ------- + float + The estimated noise standard deviation (MAD of the image pixels within the window or mask). + """ + if mask is not None: + return self.sigma_mad(image[self.apply_mask(mask)]) + + # Use the default window if no mask is provided + return self.sigma_mad(image[self.window]) + diff --git a/src/wf_psf/utils/utils.py b/src/wf_psf/utils/utils.py index 17219ad9..401cc6d8 100644 --- a/src/wf_psf/utils/utils.py +++ b/src/wf_psf/utils/utils.py @@ -437,149 +437,6 @@ def add_noise(image, desired_SNR): return noisy_image -class NoiseEstimator: - """ - A class for estimating noise levels in an image. - - Parameters - ---------- - img_dim : tuple of int - The dimensions of the image as (height, width). - win_rad : int - The radius of the exclusion window (in pixels). - """ - - def __init__(self, img_dim: tuple[int, int], win_rad: int) -> None: - """ - Initialize a NoiseEstimator instance. - - This constructor sets up the noise estimator by storing the image dimensions - and exclusion window radius, then initializes the exclusion window mask. - - Parameters - ---------- - img_dim : tuple of int - The dimensions of the image as (height, width). - win_rad : int - The radius of the exclusion window in pixels. Pixels within this radius - of the image center are excluded from noise estimation. - - Notes - ----- - The exclusion window is initialized automatically via _init_window(), creating - a boolean mask where pixels inside the exclusion radius are marked False - (excluded) and pixels outside are marked True (included). - """ - self.img_dim: tuple[int, int] = img_dim - self.win_rad: int = win_rad - - self._init_window() # Initialize self.window - - def _init_window(self): - """ - Initialize the exclusion window mask stored in self.window. - - The mask is a boolean array of shape `self.img_dim` (rows, cols). Pixels - whose Euclidean distance from the image center is less than or equal to - `self.win_rad` are marked False (excluded); all other pixels are True - (included). The mask dtype is `bool`. - - Notes - ----- - - The image center is computed as (rows / 2, cols / 2). This yields a - floating-point center so the distance is computed with sub-pixel - precision; for even dimensions the center lies between pixels. - - The comparison uses "<=" so pixels exactly at distance `win_rad` are - excluded. Change to "<" if you prefer a strict interior exclusion. - - Time complexity is O(rows * cols) for mask construction. - - No return value; the constructed mask is assigned to `self.window`. - """ - self.window = np.ones(self.img_dim, dtype=bool) - - mid_x = self.img_dim[0] / 2 - mid_y = self.img_dim[1] / 2 - - for _x in range(self.img_dim[0]): - for _y in range(self.img_dim[1]): - # If pixel is within the exclusion radius, set it to False - if np.sqrt((_x - mid_x) ** 2 + (_y - mid_y) ** 2) <= self.win_rad: - self.window[_x, _y] = False - - def apply_mask(self, mask: np.ndarray = None) -> np.ndarray: - """ - Apply a given mask to the exclusion window. - - Parameters - ---------- - mask : np.ndarray, optional - A boolean mask to apply to the exclusion window. If None, the exclusion - window is returned without any modification. - - Returns - ------- - np.ndarray - The resulting boolean array after applying the mask to the exclusion window. - """ - if mask is None: - return self.window # Return just the window if no mask is provided - return self.window & mask # Otherwise, apply the mask as usual - - @staticmethod - def sigma_mad(x): - """ - Robustly estimate the standard deviation using the Median Absolute Deviation (MAD). - - Computes MAD = ``median(|x - median(x)|)`` and scales it by 1.4826 to make the - estimator consistent with the standard deviation for a Gaussian distribution: - - sigma ≈ 1.4826 * MAD - - Parameters - ---------- - x : array-like - Input data. The values are flattened before computation. NaNs are not - specially handled and will propagate; remove or mask them prior to - calling if needed. - - Returns - ------- - float - Robust estimate of the standard deviation of the input data. - - Notes - ----- - - The MAD-based estimator is much less sensitive to outliers than the - sample standard deviation, making it appropriate for noisy data with - occasional large deviations. - - The constant 1.4826 is the scaling factor for consistency with the - standard deviation of a normal distribution. - """ - return 1.4826 * np.median(np.abs(x - np.median(x))) - - def estimate_noise(self, image: np.ndarray, mask: np.ndarray = None) -> float: - """ - Estimates the noise level of an image using the MAD estimator. - - Parameters - ---------- - image : np.ndarray - The input image for noise estimation. - mask : np.ndarray, optional - A boolean mask specifying which pixels to include in the noise estimation. - If None, the default exclusion window is used. The mask should have the same shape as `image`. - - Returns - ------- - float - The estimated noise standard deviation (MAD of the image pixels within the window or mask). - """ - if mask is not None: - return self.sigma_mad(image[self.apply_mask(mask)]) - - # Use the default window if no mask is provided - return self.sigma_mad(image[self.window]) - - class ZernikeInterpolation: """Interpolate Zernike coefficients using K-nearest RBF splines. From 74cd5411592b20874ca2f74ff145db1ed8da4a41 Mon Sep 17 00:00:00 2001 From: chaitanya chawak Date: Wed, 16 Sep 2026 15:22:09 +0200 Subject: [PATCH 03/15] add method to estimate noise for image batches --- src/wf_psf/training/train_utils.py | 26 +++++++++++--------------- src/wf_psf/utils/noise.py | 30 ++++++++++++++++++++++++++++++ 2 files changed, 41 insertions(+), 15 deletions(-) diff --git a/src/wf_psf/training/train_utils.py b/src/wf_psf/training/train_utils.py index c41ef393..9e724d57 100644 --- a/src/wf_psf/training/train_utils.py +++ b/src/wf_psf/training/train_utils.py @@ -13,7 +13,8 @@ import tensorflow as tf from typing import Optional, Callable, Union from wf_psf.psf_models.psf_models import compile_PSF_model -from wf_psf.utils.utils import NoiseEstimator, generalised_sigmoid +from wf_psf.utils.noise import NoiseEstimator +from wf_psf.utils.utils import generalised_sigmoid import logging logger = logging.getLogger(__name__) @@ -395,14 +396,12 @@ def estimate_noise_sigma( """ Estimate the per-observation noise standard deviation. - For each image in ``outputs``, the noise standard deviation is estimated with - :class:`wf_psf.utils.utils.NoiseEstimator`, which applies a robust MAD-based - estimator on the unmasked background pixels (i.e. pixels outside a central - exclusion window, and outside the per-image mask when a masked loss is used). - - This helper isolates the noise-estimation step so that the resulting sigma can - be reused independently of the training-specific sample-weight logic (e.g. for - goodness-of-fit metrics such as reduced chi-squared in validation workflows). + This resolves the image and, when applicable, the per-image mask from the + training-specific ``outputs``/``loss`` representation, then delegates the + actual batch estimation to :meth:`wf_psf.utils.noise.NoiseEstimator.estimate_noise_batch`, + which applies a robust MAD-based estimator on the unmasked background pixels + (i.e. pixels outside a central exclusion window, and outside the per-image + mask when one is provided). Parameters ---------- @@ -431,15 +430,12 @@ def estimate_noise_sigma( logger.info("Estimating noise standard deviation for masked images..") images = outputs[..., 0] masks = np.array(1 - outputs[..., 1], dtype=bool) - imgs_std = np.array( - [std_est.estimate_noise(_im, _win) for _im, _win in zip(images, masks)] - ) else: logger.info("Estimating noise standard deviation for images..") - # Estimate noise standard deviation - imgs_std = np.array([std_est.estimate_noise(_im) for _im in outputs]) + images = outputs + masks = None - return imgs_std + return std_est.estimate_noise_batch(images, masks) def calculate_sample_weights( diff --git a/src/wf_psf/utils/noise.py b/src/wf_psf/utils/noise.py index c8066e91..889f0c78 100644 --- a/src/wf_psf/utils/noise.py +++ b/src/wf_psf/utils/noise.py @@ -149,3 +149,33 @@ def estimate_noise(self, image: np.ndarray, mask: np.ndarray = None) -> float: # Use the default window if no mask is provided return self.sigma_mad(image[self.window]) + def estimate_noise_batch( + self, images: np.ndarray, masks: np.ndarray = None + ) -> np.ndarray: + """ + Estimate the noise standard deviation for a batch of images. + + Parameters + ---------- + images : np.ndarray + A batch of images with shape ``(batch_size, height, width)``. + masks : np.ndarray, optional + A batch of boolean masks with the same shape as `images`, specifying + which pixels to include in the noise estimation for each image. If + None, only the exclusion window is used for every image. + + Returns + ------- + np.ndarray + A 1D array of shape ``(batch_size,)`` containing the estimated noise + standard deviation for each image. + """ + if masks is None: + return np.array([self.estimate_noise(_im) for _im in images]) + + return np.array( + [ + self.estimate_noise(_im, _mask) + for _im, _mask in zip(images, masks) + ] + ) From 7560d27e3fbebb6c96ea38b34d8ea38e7b8f548d Mon Sep 17 00:00:00 2001 From: chaitanya chawak Date: Wed, 16 Sep 2026 15:47:18 +0200 Subject: [PATCH 04/15] update tests --- .../tests/test_training/train_utils_test.py | 3 + src/wf_psf/tests/test_utils/noise_test.py | 156 ++++++++++++++++++ src/wf_psf/tests/test_utils/utils_test.py | 138 ---------------- 3 files changed, 159 insertions(+), 138 deletions(-) create mode 100644 src/wf_psf/tests/test_utils/noise_test.py diff --git a/src/wf_psf/tests/test_training/train_utils_test.py b/src/wf_psf/tests/test_training/train_utils_test.py index 09921799..108d7777 100644 --- a/src/wf_psf/tests/test_training/train_utils_test.py +++ b/src/wf_psf/tests/test_training/train_utils_test.py @@ -34,6 +34,9 @@ def mock_noise_estimator(): mock_instance.estimate_noise.side_effect = lambda img, mask=None: np.std( img ) # Mock behavior + mock_instance.estimate_noise_batch.side_effect = ( + lambda images, masks=None: np.array([np.std(img) for img in images]) + ) yield mock_instance diff --git a/src/wf_psf/tests/test_utils/noise_test.py b/src/wf_psf/tests/test_utils/noise_test.py new file mode 100644 index 00000000..be458e63 --- /dev/null +++ b/src/wf_psf/tests/test_utils/noise_test.py @@ -0,0 +1,156 @@ +import numpy as np +from wf_psf.utils.noise import NoiseEstimator + + +def test_initialization(): + """Test if NoiseEstimator initializes correctly.""" + img_dim = (50, 50) + win_rad = 10 + estimator = NoiseEstimator(img_dim, win_rad) + + assert estimator.img_dim == img_dim + assert estimator.win_rad == win_rad + assert isinstance(estimator.window, np.ndarray) + assert estimator.window.shape == img_dim + + +def test_init_window(): + """Test that the exclusion window is correctly initialized.""" + img_dim = (50, 50) + win_rad = 10 + estimator = NoiseEstimator(img_dim, win_rad) + + mid_x, mid_y = img_dim[0] / 2, img_dim[1] / 2 + + for x in range(img_dim[0]): + for y in range(img_dim[1]): + # Pixel inside the exclusion radius should be False, others True + inside_radius = np.sqrt((x - mid_x) ** 2 + (y - mid_y) ** 2) <= win_rad + assert estimator.window[x, y] == (not inside_radius) + + +def test_sigma_mad(): + """Test the MAD-based standard deviation estimation.""" + data = np.array( + [1, 1, 2, 2, 3, 3, 4, 4, 100] + ) # Outlier should not heavily influence MAD + expected_sigma = 1.4826 * np.median(np.abs(data - np.median(data))) + + assert np.isclose(NoiseEstimator.sigma_mad(data), expected_sigma, atol=1e-4) + + +def test_estimate_noise_without_default_window(): + """Test noise estimation with the default exclusion window (no custom mask).""" + img_dim = (50, 50) + win_rad = 5 + estimator = NoiseEstimator(img_dim, win_rad) + + # Create a synthetic noisy image (Gaussian noise with mean=0, std=10) + np.random.seed(42) + image = np.random.normal(0, 10, img_dim) + + noise_estimation = estimator.estimate_noise(image) + + # The estimated noise should be close to 10 (the true std) + assert np.isclose(noise_estimation, 10, atol=2) + + +def test_estimate_noise_with_custom_mask(): + """Test noise estimation with a custom mask applied outside the exclusion radius.""" + img_dim = (50, 50) + estimator = NoiseEstimator(img_dim, win_rad=5) + + # Create synthetic noise with std=5 + np.random.seed(42) + image = np.random.normal(0, 5, img_dim) + + # Custom mask excluding top-left corner + custom_mask = np.ones(img_dim, dtype=bool) + custom_mask[:10, :10] = False # Mask out top-left 10x10 pixels + + noise_estimation = estimator.estimate_noise(image, mask=custom_mask) + + assert np.isclose(noise_estimation, 5, atol=1) + + +def test_apply_mask_with_none_mask(): + """Test apply_mask when mask is None.""" + img_dim = (10, 10) + estimator = NoiseEstimator(img_dim, win_rad=3) + + result = estimator.apply_mask(None) # Pass None as the mask + + # It should return the window itself when no mask is provided + assert np.array_equal( + result, estimator.window + ), "apply_mask should return the window when mask is None." + + +def test_apply_mask_with_valid_mask(): + """Test apply_mask when a valid mask is provided.""" + img_dim = (10, 10) + estimator = NoiseEstimator(img_dim, win_rad=3) + + # Create a custom mask + custom_mask = np.ones(img_dim, dtype=bool) + custom_mask[5, 5] = False # Set a pixel to False to exclude it from the window + + result = estimator.apply_mask(custom_mask) + + # Check that the mask was applied correctly: pixel (5, 5) should be False, others True + expected_result = estimator.window & custom_mask + assert np.array_equal( + result, expected_result + ), "apply_mask did not apply the mask correctly." + + +def test_apply_mask_with_zeroed_mask(): + """Test apply_mask when a zeroed mask is provided.""" + img_dim = (10, 10) + estimator = NoiseEstimator(img_dim, win_rad=3) + + # Create a mask where all pixels are excluded (False) + zeroed_mask = np.zeros(img_dim, dtype=bool) + + result = estimator.apply_mask(zeroed_mask) + + # The result should be an array of False values, as the mask excludes all pixels + expected_result = np.zeros(img_dim, dtype=bool) + assert np.array_equal( + result, expected_result + ), "apply_mask did not handle the zeroed mask correctly." + + +def test_estimate_noise_batch_without_masks(): + """estimate_noise_batch with no masks matches per-image estimate_noise.""" + img_dim = (30, 30) + win_rad = 5 + estimator = NoiseEstimator(img_dim, win_rad) + + np.random.seed(0) + images = np.random.normal(0, 3, size=(4, *img_dim)) + + batch_result = estimator.estimate_noise_batch(images) + expected = np.array([estimator.estimate_noise(img) for img in images]) + + assert batch_result.shape == (4,) + np.testing.assert_allclose(batch_result, expected) + + +def test_estimate_noise_batch_with_masks(): + """estimate_noise_batch with per-image masks matches per-image estimate_noise.""" + img_dim = (30, 30) + win_rad = 5 + estimator = NoiseEstimator(img_dim, win_rad) + + np.random.seed(1) + images = np.random.normal(0, 3, size=(4, *img_dim)) + masks = np.random.randint(0, 2, size=(4, *img_dim)).astype(bool) + + batch_result = estimator.estimate_noise_batch(images, masks) + expected = np.array( + [estimator.estimate_noise(img, mask) for img, mask in zip(images, masks)] + ) + + assert batch_result.shape == (4,) + np.testing.assert_allclose(batch_result, expected) diff --git a/src/wf_psf/tests/test_utils/utils_test.py b/src/wf_psf/tests/test_utils/utils_test.py index cc7f2a2b..c849c692 100644 --- a/src/wf_psf/tests/test_utils/utils_test.py +++ b/src/wf_psf/tests/test_utils/utils_test.py @@ -11,7 +11,6 @@ import numpy as np from wf_psf.utils.utils import ( downsample_im, - NoiseEstimator, zernike_generator, compute_unobscured_zernike_projection, decompose_tf_obscured_opd_basis, @@ -19,143 +18,6 @@ from wf_psf.sims.psf_simulator import PSFSimulator from unittest import mock - -def test_downsample_basic(): - """Test apply_mask when a zeroed mask is provided.""" - img_dim = (10, 10) - estimator = NoiseEstimator(img_dim, win_rad=3) - - # Create a mask where all pixels are excluded (False) - zeroed_mask = np.zeros(img_dim, dtype=bool) - - result = estimator.apply_mask(zeroed_mask) - - # The result should be an array of False values, as the mask excludes all pixels - expected_result = np.zeros(img_dim, dtype=bool) - assert np.array_equal( - result, expected_result - ), "apply_mask did not handle the zeroed mask correctly." - - -def test_initialization(): - """Test if NoiseEstimator initializes correctly.""" - img_dim = (50, 50) - win_rad = 10 - estimator = NoiseEstimator(img_dim, win_rad) - - assert estimator.img_dim == img_dim - assert estimator.win_rad == win_rad - assert isinstance(estimator.window, np.ndarray) - assert estimator.window.shape == img_dim - - -def test_init_window(): - """Test that the exclusion window is correctly initialized.""" - img_dim = (50, 50) - win_rad = 10 - estimator = NoiseEstimator(img_dim, win_rad) - - mid_x, mid_y = img_dim[0] / 2, img_dim[1] / 2 - - for x in range(img_dim[0]): - for y in range(img_dim[1]): - # Pixel inside the exclusion radius should be False, others True - inside_radius = np.sqrt((x - mid_x) ** 2 + (y - mid_y) ** 2) <= win_rad - assert estimator.window[x, y] == (not inside_radius) - - -def test_sigma_mad(): - """Test the MAD-based standard deviation estimation.""" - data = np.array( - [1, 1, 2, 2, 3, 3, 4, 4, 100] - ) # Outlier should not heavily influence MAD - expected_sigma = 1.4826 * np.median(np.abs(data - np.median(data))) - - assert np.isclose(NoiseEstimator.sigma_mad(data), expected_sigma, atol=1e-4) - - -def test_estimate_noise_without_default_window(): - """Test noise estimation with the default exclusion window (no custom mask).""" - img_dim = (50, 50) - win_rad = 5 - estimator = NoiseEstimator(img_dim, win_rad) - - # Create a synthetic noisy image (Gaussian noise with mean=0, std=10) - np.random.seed(42) - image = np.random.normal(0, 10, img_dim) - - noise_estimation = estimator.estimate_noise(image) - - # The estimated noise should be close to 10 (the true std) - assert np.isclose(noise_estimation, 10, atol=2) - - -def test_estimate_noise_with_custom_mask(): - """Test noise estimation with a custom mask applied outside the exclusion radius.""" - img_dim = (50, 50) - estimator = NoiseEstimator(img_dim, win_rad=5) - - # Create synthetic noise with std=5 - np.random.seed(42) - image = np.random.normal(0, 5, img_dim) - - # Custom mask excluding top-left corner - custom_mask = np.ones(img_dim, dtype=bool) - custom_mask[:10, :10] = False # Mask out top-left 10x10 pixels - - noise_estimation = estimator.estimate_noise(image, mask=custom_mask) - - assert np.isclose(noise_estimation, 5, atol=1) - - -def test_apply_mask_with_none_mask(): - """Test apply_mask when mask is None.""" - img_dim = (10, 10) - estimator = NoiseEstimator(img_dim, win_rad=3) - - result = estimator.apply_mask(None) # Pass None as the mask - - # It should return the window itself when no mask is provided - assert np.array_equal( - result, estimator.window - ), "apply_mask should return the window when mask is None." - - -def test_apply_mask_with_valid_mask(): - """Test apply_mask when a valid mask is provided.""" - img_dim = (10, 10) - estimator = NoiseEstimator(img_dim, win_rad=3) - - # Create a custom mask - custom_mask = np.ones(img_dim, dtype=bool) - custom_mask[5, 5] = False # Set a pixel to False to exclude it from the window - - result = estimator.apply_mask(custom_mask) - - # Check that the mask was applied correctly: pixel (5, 5) should be False, others True - expected_result = estimator.window & custom_mask - assert np.array_equal( - result, expected_result - ), "apply_mask did not apply the mask correctly." - - -def test_apply_mask_with_zeroed_mask(): - """Test apply_mask when a zeroed mask is provided.""" - img_dim = (10, 10) - estimator = NoiseEstimator(img_dim, win_rad=3) - - # Create a mask where all pixels are excluded (False) - zeroed_mask = np.zeros(img_dim, dtype=bool) - - result = estimator.apply_mask(zeroed_mask) - - # The result should be an array of False values, as the mask excludes all pixels - expected_result = np.zeros(img_dim, dtype=bool) - assert np.array_equal( - result, expected_result - ), "apply_mask did not handle the zeroed mask correctly." - - def test_unobscured_zernike_projection(): n_zernikes = 20 wfe_dim = 256 From 5278721d8d113caa79522f8325de8280a9b949c3 Mon Sep 17 00:00:00 2001 From: chaitanya chawak Date: Thu, 17 Sep 2026 14:00:45 +0200 Subject: [PATCH 05/15] decouple loss from calculate_sample_weights and move it into general_train_cycle --- src/wf_psf/training/train_utils.py | 133 ++++++++++------------------- src/wf_psf/utils/noise.py | 34 ++++++-- 2 files changed, 75 insertions(+), 92 deletions(-) diff --git a/src/wf_psf/training/train_utils.py b/src/wf_psf/training/train_utils.py index 9e724d57..bb948ea5 100644 --- a/src/wf_psf/training/train_utils.py +++ b/src/wf_psf/training/train_utils.py @@ -389,80 +389,31 @@ def _is_masked_loss(loss: Union[str, Callable, None]) -> bool: ) -def estimate_noise_sigma( - outputs: tf.Tensor, - loss: Union[str, Callable, None] = None, -) -> np.ndarray: - """ - Estimate the per-observation noise standard deviation. - - This resolves the image and, when applicable, the per-image mask from the - training-specific ``outputs``/``loss`` representation, then delegates the - actual batch estimation to :meth:`wf_psf.utils.noise.NoiseEstimator.estimate_noise_batch`, - which applies a robust MAD-based estimator on the unmasked background pixels - (i.e. pixels outside a central exclusion window, and outside the per-image - mask when one is provided). - - Parameters - ---------- - outputs: tf.Tensor - Image data. When ``loss`` is ``"masked_mean_squared_error"``, a 4D tensor - of shape ``(batch_size, height, width, 2)`` is expected, where the last - dimension holds ``[image, mask]``. Otherwise, a 3D tensor of shape - ``(batch_size, height, width)`` is expected. - loss: str, callable, optional - The loss function (or its name) used for training. If the loss is - ``"masked_mean_squared_error"``, the per-image mask stored in - ``outputs[..., 1]`` is combined with the exclusion window before - estimating the noise. Default is None. - - Returns - ------- - np.ndarray - A 1D array of shape ``(batch_size,)`` containing the estimated noise - standard deviation for each image. - """ - img_dim = (outputs.shape[1], outputs.shape[2]) - win_rad = np.ceil(outputs.shape[1] / 3.33) - std_est = NoiseEstimator(img_dim=img_dim, win_rad=win_rad) - - if _is_masked_loss(loss): - logger.info("Estimating noise standard deviation for masked images..") - images = outputs[..., 0] - masks = np.array(1 - outputs[..., 1], dtype=bool) - else: - logger.info("Estimating noise standard deviation for images..") - images = outputs - masks = None - - return std_est.estimate_noise_batch(images, masks) - - def calculate_sample_weights( - outputs: tf.Tensor, - use_sample_weights: bool, - loss: Union[str, Callable, None], + images: tf.Tensor, + masks: Optional[np.ndarray] = None, apply_sigmoid: bool = False, sigmoid_max_val: float = 5.0, sigmoid_power_k: float = 1.0, -) -> Optional[np.ndarray]: +) -> np.ndarray: """ Calculate sample weights based on image noise standard deviation. - The function computes sample weights by estimating the noise standard deviation for each image, calculating the inverse variance, - and then normalizing the weights by dividing by the median. + The function estimates the per-observation noise standard deviation with + :class:`wf_psf.utils.noise.NoiseEstimator`, computes sample weights from the + inverse variance, and normalizes the weights by dividing by the median. Parameters ---------- - outputs: tf.Tensor + images: tf.Tensor A 3D tensor of shape (batch_size, height, width) representing images, where the first dimension is the batch size and the next two dimensions are the image height and width. - use_sample_weights: bool - Flag indicating whether to compute sample weights. If True, sample weights will be computed based on the image noise. - loss: str, callable, optional - The loss function used for training. If the loss name is "masked_mean_squared_error", the function will calculate the noise standard deviation for masked images. + masks: np.ndarray, optional + A batch of boolean masks with the same shape as `images`, specifying + which pixels to include in the noise estimation for each image. If + None, only the exclusion window is used for every image. Default is None. apply_sigmoid: bool, optional - Flag indicating whether to apply a generalized sigmoid function to the sample weights. Default is True. + Flag indicating whether to apply a generalized sigmoid function to the sample weights. Default is False. sigmoid_max_val: float, optional The maximum value for the sigmoid function. Default is 5.0. sigmoid_power_k: float, optional @@ -472,28 +423,27 @@ def calculate_sample_weights( Returns ------- - np.ndarray or None - An array of sample weights, or None if `use_sample_weights` is False. + np.ndarray + An array of sample weights. """ - if use_sample_weights: - # Estimate the per-observation noise standard deviation - imgs_std = estimate_noise_sigma(outputs, loss) + img_dim = (images.shape[1], images.shape[2]) + std_est = NoiseEstimator(img_dim=img_dim) - # Calculate variances - variances = imgs_std**2 + # Estimate the per-observation noise standard deviation + imgs_std = std_est.estimate_noise_batch(images, masks) - # Use inverse variance for weights and scale by median - sample_weight = 1 / variances - sample_weight /= np.median(sample_weight) + # Calculate variances + variances = imgs_std**2 - # Apply a generalized sigmoid function to the sample weights - if apply_sigmoid: - sample_weight = generalised_sigmoid( - sample_weight, max_val=sigmoid_max_val, power_k=sigmoid_power_k - ) + # Use inverse variance for weights and scale by median + sample_weight = 1 / variances + sample_weight /= np.median(sample_weight) - else: - sample_weight = None + # Apply a generalized sigmoid function to the sample weights + if apply_sigmoid: + sample_weight = generalised_sigmoid( + sample_weight, max_val=sigmoid_max_val, power_k=sigmoid_power_k + ) return sample_weight @@ -776,15 +726,26 @@ def general_train_cycle( learning_rate_param, param_optim, param_loss, param_metrics ) + # Resolve the training-specific outputs/loss representation into a plain + # image (and, for a masked loss, a per-image mask) before computing weights + if _is_masked_loss(loss): + images = outputs[..., 0] + masks = np.array(1 - outputs[..., 1], dtype=bool) + else: + images = outputs + masks = None + # Calculate sample weights - sample_weight = calculate_sample_weights( - outputs, - use_sample_weights, - loss, - apply_sigmoid, - sigmoid_max_val, - sigmoid_power_k, - ) + if use_sample_weights: + sample_weight = calculate_sample_weights( + images, + masks, + apply_sigmoid, + sigmoid_max_val, + sigmoid_power_k, + ) + else: + sample_weight = None # Define the training cycle if cycle_def in ("parametric", "complete", "only-parametric"): diff --git a/src/wf_psf/utils/noise.py b/src/wf_psf/utils/noise.py index 889f0c78..b1d8b36c 100644 --- a/src/wf_psf/utils/noise.py +++ b/src/wf_psf/utils/noise.py @@ -15,11 +15,12 @@ class NoiseEstimator: ---------- img_dim : tuple of int The dimensions of the image as (height, width). - win_rad : int - The radius of the exclusion window (in pixels). + win_rad : int, optional + The radius of the exclusion window (in pixels). If None, a default + radius is derived from `img_dim` (see :meth:`default_win_rad`). """ - def __init__(self, img_dim: tuple[int, int], win_rad: int) -> None: + def __init__(self, img_dim: tuple[int, int], win_rad: int = None) -> None: """ Initialize a NoiseEstimator instance. @@ -30,9 +31,10 @@ def __init__(self, img_dim: tuple[int, int], win_rad: int) -> None: ---------- img_dim : tuple of int The dimensions of the image as (height, width). - win_rad : int + win_rad : int, optional The radius of the exclusion window in pixels. Pixels within this radius - of the image center are excluded from noise estimation. + of the image center are excluded from noise estimation. If None, the + radius is derived from `img_dim` via :meth:`default_win_rad`. Notes ----- @@ -41,10 +43,30 @@ def __init__(self, img_dim: tuple[int, int], win_rad: int) -> None: (excluded) and pixels outside are marked True (included). """ self.img_dim: tuple[int, int] = img_dim - self.win_rad: int = win_rad + self.win_rad: int = ( + win_rad if win_rad is not None else self.default_win_rad(img_dim) + ) self._init_window() # Initialize self.window + @staticmethod + def default_win_rad(img_dim: tuple[int, int]) -> int: + """ + Compute a default exclusion-window radius from the image dimensions. + + Parameters + ---------- + img_dim : tuple of int + The dimensions of the image as (height, width). + + Returns + ------- + int + The default exclusion-window radius, in pixels, for an image of + this size. + """ + return int(np.ceil(img_dim[0] / 3.33)) + def _init_window(self): """ Initialize the exclusion window mask stored in self.window. From 5e0a11f7af5390d5d3a9bb3764652e83aece6ac0 Mon Sep 17 00:00:00 2001 From: chaitanya chawak Date: Thu, 17 Sep 2026 14:13:56 +0200 Subject: [PATCH 06/15] update tests --- .../tests/test_training/train_utils_test.py | 218 +++++++++--------- 1 file changed, 110 insertions(+), 108 deletions(-) diff --git a/src/wf_psf/tests/test_training/train_utils_test.py b/src/wf_psf/tests/test_training/train_utils_test.py index 108d7777..51156994 100644 --- a/src/wf_psf/tests/test_training/train_utils_test.py +++ b/src/wf_psf/tests/test_training/train_utils_test.py @@ -121,72 +121,31 @@ def test_learning_rate_affects_default_optimizer(): @pytest.mark.parametrize( - "use_sample_weights, expected_output", + "masks", [ - (False, None), - (True, np.ndarray), # When enabled, it should return an array + None, + np.random.randint(0, 2, size=(5, 32, 32)).astype(bool), # Per-image mask ], ) -def test_calculate_sample_weights_no_weights( - mock_noise_estimator, use_sample_weights, expected_output -): - """Test when sample weights are disabled and ensure correct return type.""" - outputs = np.random.rand(5, 32, 32) # (batch_size=5, height=32, width=32) - result = train_utils.calculate_sample_weights( - outputs, use_sample_weights, loss=None - ) - if not use_sample_weights: - assert result is None - else: - assert isinstance(result, expected_output) - - -@pytest.mark.parametrize( - "use_sample_weights, loss, expected_output_type", - [ - (False, None, None), # Null test: when disabled, should return None - (True, "mean_squared_error", np.ndarray), # Standard MSE case - (True, "masked_mean_squared_error", np.ndarray), # Masked MSE case - ], -) -def test_calculate_sample_weights_integration( - use_sample_weights, loss, expected_output_type -): - """Test different cases for sample weight computation.""" - # Generate dummy image data +def test_calculate_sample_weights_integration(masks): + """Test sample weight computation with real (unmocked) noise estimation.""" batch_size, height, width = 5, 32, 32 + images = np.random.rand(batch_size, height, width) - if loss == "masked_mean_squared_error": - # Create image-mask pairs: last dimension has [image, mask] - outputs = np.random.rand(batch_size, height, width, 2) - outputs[..., 1] = np.random.randint( - 0, 2, size=(batch_size, height, width) - ) # Binary mask - else: - outputs = np.random.rand(batch_size, height, width) + result = train_utils.calculate_sample_weights(images, masks=masks) - result = train_utils.calculate_sample_weights(outputs, use_sample_weights, loss) - - if not use_sample_weights: - assert result is None - else: - assert isinstance(result, expected_output_type) - assert result.shape == (batch_size,) # Expecting one weight per sample - assert np.all(result > 0) # Weights should be positive - assert np.isclose( - np.median(result), 1, atol=0.1 - ) # Weights should be scaled around 1 + assert isinstance(result, np.ndarray) + assert result.shape == (batch_size,) # Expecting one weight per sample + assert np.all(result > 0) # Weights should be positive + assert np.isclose( + np.median(result), 1, atol=0.1 + ) # Weights should be scaled around 1 -@pytest.mark.parametrize( - "loss", [None, "mean_squared_error", "masked_mean_squared_error"] -) -def test_calculate_sample_weights_unit(mock_noise_estimator, loss): +def test_calculate_sample_weights_unit(mock_noise_estimator): """Test sample weighting strategy with random images.""" - outputs = np.random.rand(10, 32, 32) # 10 images of size 32x32 - result = train_utils.calculate_sample_weights( - outputs, use_sample_weights=True, loss=loss - ) + images = np.random.rand(10, 32, 32) # 10 images of size 32x32 + result = train_utils.calculate_sample_weights(images) # Check the type and shape of the result assert isinstance(result, np.ndarray) @@ -199,27 +158,19 @@ def test_calculate_sample_weights_unit(mock_noise_estimator, loss): assert np.isclose(np.median(result), 1, atol=0.1) -@pytest.mark.parametrize( - "loss", [None, "mean_squared_error", "masked_mean_squared_error"] -) -def test_calculate_sample_weights_high_variance(mock_noise_estimator, loss): +def test_calculate_sample_weights_high_variance(mock_noise_estimator): """Test case for high variance (noisy images).""" # Create high variance images with more noise - outputs = np.random.normal(loc=0.0, scale=10.0, size=(5, 32, 32)) # Larger noise + images = np.random.normal(loc=0.0, scale=10.0, size=(5, 32, 32)) # Larger noise # Calculate sample weights - weights = train_utils.calculate_sample_weights( - outputs, use_sample_weights=True, loss=loss - ) + weights = train_utils.calculate_sample_weights(images) - # Check that weights are relatively lower for high variance images # Compare with weights from a small variance case - small_variance_outputs = np.random.normal( + small_variance_images = np.random.normal( loc=0.0, scale=1e-6, size=(5, 32, 32) ) # Small noise - small_variance_weights = train_utils.calculate_sample_weights( - small_variance_outputs, use_sample_weights=True, loss=None - ) + small_variance_weights = train_utils.calculate_sample_weights(small_variance_images) # Check if the median of the high variance weights is smaller than the median of the low variance weights assert np.median(weights) <= np.median(small_variance_weights), ( @@ -232,51 +183,37 @@ def test_calculate_sample_weights_high_variance(mock_noise_estimator, loss): @pytest.mark.parametrize( - "loss", [None, "mean_squared_error", "masked_mean_squared_error"] + "loss, expected", + [ + (None, False), + ("mean_squared_error", False), + ("masked_mean_squared_error", True), + (train_utils.MaskedMeanSquaredError(), True), + ], ) -def test_estimate_noise_sigma_shape_and_positivity(loss): - """estimate_noise_sigma returns one positive sigma per observation.""" - batch_size, height, width = 5, 32, 32 - - if loss == "masked_mean_squared_error": - outputs = np.random.rand(batch_size, height, width, 2) - outputs[..., 1] = np.random.randint(0, 2, size=(batch_size, height, width)) - else: - outputs = np.random.rand(batch_size, height, width) - - sigma = train_utils.estimate_noise_sigma(outputs, loss) +def test_is_masked_loss(loss, expected): + """_is_masked_loss identifies the masked_mean_squared_error loss by name.""" + assert train_utils._is_masked_loss(loss) is expected - assert isinstance(sigma, np.ndarray) - assert sigma.shape == (batch_size,) - assert np.all(sigma > 0) - -@pytest.mark.parametrize( - "loss", [None, "mean_squared_error", "masked_mean_squared_error"] -) -def test_calculate_sample_weights_uses_estimate_noise_sigma(loss): - """calculate_sample_weights delegates noise estimation to estimate_noise_sigma.""" +def test_calculate_sample_weights_uses_noise_estimator_batch(): + """calculate_sample_weights delegates noise estimation to NoiseEstimator.estimate_noise_batch.""" batch_size, height, width = 5, 32, 32 - - if loss == "masked_mean_squared_error": - outputs = np.random.rand(batch_size, height, width, 2) - outputs[..., 1] = np.random.randint(0, 2, size=(batch_size, height, width)) - else: - outputs = np.random.rand(batch_size, height, width) + images = np.random.rand(batch_size, height, width) + masks = np.random.randint(0, 2, size=(batch_size, height, width)).astype(bool) fake_sigma = np.linspace(1.0, 2.0, batch_size) - with patch.object( - train_utils, "estimate_noise_sigma", return_value=fake_sigma - ) as mock_sigma: - result = train_utils.calculate_sample_weights( - outputs, use_sample_weights=True, loss=loss - ) + with patch("wf_psf.training.train_utils.NoiseEstimator") as MockNoiseEstimator: + mock_instance = MockNoiseEstimator.return_value + mock_instance.estimate_noise_batch.return_value = fake_sigma + + result = train_utils.calculate_sample_weights(images, masks) - mock_sigma.assert_called_once() - # Passed the same outputs and loss through - called_outputs, called_loss = mock_sigma.call_args[0][:2] - assert called_outputs is outputs - assert called_loss == loss + mock_instance.estimate_noise_batch.assert_called_once() + # Passed the same images and masks through, with no knowledge of any loss + called_images, called_masks = mock_instance.estimate_noise_batch.call_args[0][:2] + assert called_images is images + assert called_masks is masks # Weights are the median-normalised inverse variance of the returned sigma expected = 1 / fake_sigma**2 @@ -424,6 +361,71 @@ def test_general_train_cycle_handles_cycle_def( mock_model.set_trainable_layers.assert_has_calls(expected_calls, any_order=False) +def test_general_train_cycle_resolves_images_and_masks_for_masked_loss( + mock_test_setup, +): + """general_train_cycle unpacks image/mask from outputs when the loss is masked, + before calculate_sample_weights ever sees the data.""" + mock_model = mock_test_setup["mock_model"] + inputs = mock_test_setup["inputs"] + + batch_size, height, width = 4, 8, 8 + images = np.random.rand(batch_size, height, width) + raw_masks = np.random.randint(0, 2, size=(batch_size, height, width)) + outputs = np.stack([images, raw_masks], axis=-1) + + with patch("wf_psf.training.train_utils.calculate_sample_weights") as mock_calc: + mock_calc.return_value = np.ones(batch_size) + train_utils.general_train_cycle( + psf_model=mock_model, + inputs=inputs, + outputs=outputs, + validation_data=mock_test_setup["validation_data"], + batch_size=mock_test_setup["batch_size"], + learning_rate_param=mock_test_setup["learning_rate_param"], + learning_rate_non_param=mock_test_setup["learning_rate_non_param"], + n_epochs_param=mock_test_setup["n_epochs_param"], + n_epochs_non_param=mock_test_setup["n_epochs_non_param"], + param_loss=train_utils.MaskedMeanSquaredError(), + use_sample_weights=True, + cycle_def="parametric", + ) + + mock_calc.assert_called_once() + called_images, called_masks = mock_calc.call_args[0][:2] + np.testing.assert_array_equal(called_images, images) + np.testing.assert_array_equal(called_masks, np.array(1 - raw_masks, dtype=bool)) + + +def test_general_train_cycle_skips_sample_weights_when_disabled(mock_test_setup): + """general_train_cycle leaves sample_weight as None and never calls + calculate_sample_weights when use_sample_weights is False.""" + mock_model = mock_test_setup["mock_model"] + + with ( + patch("wf_psf.training.train_utils.calculate_sample_weights") as mock_calc, + patch("wf_psf.training.train_utils.train_cycle_part") as mock_train_cycle_part, + ): + mock_train_cycle_part.return_value = MagicMock() + train_utils.general_train_cycle( + psf_model=mock_model, + inputs=mock_test_setup["inputs"], + outputs=mock_test_setup["outputs"], + validation_data=mock_test_setup["validation_data"], + batch_size=mock_test_setup["batch_size"], + learning_rate_param=mock_test_setup["learning_rate_param"], + learning_rate_non_param=mock_test_setup["learning_rate_non_param"], + n_epochs_param=mock_test_setup["n_epochs_param"], + n_epochs_non_param=mock_test_setup["n_epochs_non_param"], + use_sample_weights=False, + cycle_def="parametric", + ) + + mock_calc.assert_not_called() + _, train_cycle_part_kwargs = mock_train_cycle_part.call_args + assert train_cycle_part_kwargs["sample_weight"] is None + + def test_get_callbacks(): # Test when both are None assert train_utils.get_callbacks(None, None) is None From d12dc5b9387917f35d4448e8e7a292e96e2af3b4 Mon Sep 17 00:00:00 2001 From: chaitanya chawak Date: Thu, 17 Sep 2026 16:32:46 +0200 Subject: [PATCH 07/15] resolve training outputs inside the use_sample_weights conditional --- src/wf_psf/training/train_utils.py | 43 +++++++++++++++++++++++------- 1 file changed, 34 insertions(+), 9 deletions(-) diff --git a/src/wf_psf/training/train_utils.py b/src/wf_psf/training/train_utils.py index bb948ea5..41661d0d 100644 --- a/src/wf_psf/training/train_utils.py +++ b/src/wf_psf/training/train_utils.py @@ -389,6 +389,39 @@ def _is_masked_loss(loss: Union[str, Callable, None]) -> bool: ) +def _resolve_training_outputs( + outputs: tf.Tensor, loss: Union[str, Callable, None] +) -> tuple[tf.Tensor, Optional[np.ndarray]]: + """ + Resolve the training-specific ``outputs``/``loss`` representation into a + plain image and, for a masked loss, a per-image mask. + + Parameters + ---------- + outputs: tf.Tensor + Image data. When ``loss`` is ``"masked_mean_squared_error"``, a 4D tensor + of shape ``(batch_size, height, width, 2)`` is expected, where the last + dimension holds ``[image, mask]``. Otherwise, a 3D tensor of shape + ``(batch_size, height, width)`` is expected. + loss: str, callable, optional + The loss function (or its name) used for training. + + Returns + ------- + tuple of (tf.Tensor, np.ndarray or None) + The resolved images and, for a masked loss, the corresponding boolean + masks (``outputs[..., 1]`` inverted); otherwise ``None``. + """ + if _is_masked_loss(loss): + images = outputs[..., 0] + masks = np.array(1 - outputs[..., 1], dtype=bool) + else: + images = outputs + masks = None + + return images, masks + + def calculate_sample_weights( images: tf.Tensor, masks: Optional[np.ndarray] = None, @@ -726,17 +759,9 @@ def general_train_cycle( learning_rate_param, param_optim, param_loss, param_metrics ) - # Resolve the training-specific outputs/loss representation into a plain - # image (and, for a masked loss, a per-image mask) before computing weights - if _is_masked_loss(loss): - images = outputs[..., 0] - masks = np.array(1 - outputs[..., 1], dtype=bool) - else: - images = outputs - masks = None - # Calculate sample weights if use_sample_weights: + images, masks = _resolve_training_outputs(outputs, loss) sample_weight = calculate_sample_weights( images, masks, From 8ee2232e125eda0def6f79a66bef26af94405e1d Mon Sep 17 00:00:00 2001 From: chaitanya chawak Date: Thu, 17 Sep 2026 17:43:51 +0200 Subject: [PATCH 08/15] update tests - add new and minor fixes in existing --- .../tests/test_training/train_utils_test.py | 55 ++++++++++++++----- src/wf_psf/tests/test_utils/noise_test.py | 19 +++++++ 2 files changed, 60 insertions(+), 14 deletions(-) diff --git a/src/wf_psf/tests/test_training/train_utils_test.py b/src/wf_psf/tests/test_training/train_utils_test.py index 51156994..440c6e0d 100644 --- a/src/wf_psf/tests/test_training/train_utils_test.py +++ b/src/wf_psf/tests/test_training/train_utils_test.py @@ -31,9 +31,6 @@ def mock_noise_estimator(): """Mock the NoiseEstimator class to return a fixed noise estimate.""" with patch("wf_psf.training.train_utils.NoiseEstimator") as MockNoiseEstimator: mock_instance = MockNoiseEstimator.return_value - mock_instance.estimate_noise.side_effect = lambda img, mask=None: np.std( - img - ) # Mock behavior mock_instance.estimate_noise_batch.side_effect = ( lambda images, masks=None: np.array([np.std(img) for img in images]) ) @@ -192,12 +189,38 @@ def test_calculate_sample_weights_high_variance(mock_noise_estimator): ], ) def test_is_masked_loss(loss, expected): - """_is_masked_loss identifies the masked_mean_squared_error loss by name.""" + """Test that _is_masked_loss identifies the masked_mean_squared_error loss by name.""" assert train_utils._is_masked_loss(loss) is expected +def test_resolve_training_outputs_unmasked_loss(): + """Test that for a non-masked loss, outputs pass through unchanged and masks is None.""" + outputs = np.random.rand(4, 8, 8) + images, masks = train_utils._resolve_training_outputs( + outputs, "mean_squared_error" + ) + + assert images is outputs + assert masks is None + + +def test_resolve_training_outputs_masked_loss(): + """Test that for a masked loss, outputs are split into images and an inverted boolean mask.""" + batch_size, height, width = 4, 8, 8 + images = np.random.rand(batch_size, height, width) + raw_masks = np.random.randint(0, 2, size=(batch_size, height, width)) + outputs = np.stack([images, raw_masks], axis=-1) + + resolved_images, resolved_masks = train_utils._resolve_training_outputs( + outputs, "masked_mean_squared_error" + ) + + np.testing.assert_array_equal(resolved_images, images) + np.testing.assert_array_equal(resolved_masks, np.array(1 - raw_masks, dtype=bool)) + + def test_calculate_sample_weights_uses_noise_estimator_batch(): - """calculate_sample_weights delegates noise estimation to NoiseEstimator.estimate_noise_batch.""" + """Test that calculate_sample_weights delegates noise estimation to NoiseEstimator.estimate_noise_batch.""" batch_size, height, width = 5, 32, 32 images = np.random.rand(batch_size, height, width) masks = np.random.randint(0, 2, size=(batch_size, height, width)).astype(bool) @@ -358,13 +381,13 @@ def test_general_train_cycle_handles_cycle_def( if non_param_train: expected_calls.append(call(param_bool=False, nonparam_bool=True)) - mock_model.set_trainable_layers.assert_has_calls(expected_calls, any_order=False) + mock_model.set_trainable_layers.asserFootert_has_calls(expected_calls, any_order=False) def test_general_train_cycle_resolves_images_and_masks_for_masked_loss( mock_test_setup, ): - """general_train_cycle unpacks image/mask from outputs when the loss is masked, + """Test that general_train_cycle unpacks image/mask from outputs when the loss is masked, before calculate_sample_weights ever sees the data.""" mock_model = mock_test_setup["mock_model"] inputs = mock_test_setup["inputs"] @@ -374,8 +397,10 @@ def test_general_train_cycle_resolves_images_and_masks_for_masked_loss( raw_masks = np.random.randint(0, 2, size=(batch_size, height, width)) outputs = np.stack([images, raw_masks], axis=-1) - with patch("wf_psf.training.train_utils.calculate_sample_weights") as mock_calc: - mock_calc.return_value = np.ones(batch_size) + with patch( + "wf_psf.training.train_utils.calculate_sample_weights" + ) as mock_calculate_sample_weights: + mock_calculate_sample_weights.return_value = np.ones(batch_size) train_utils.general_train_cycle( psf_model=mock_model, inputs=inputs, @@ -391,19 +416,21 @@ def test_general_train_cycle_resolves_images_and_masks_for_masked_loss( cycle_def="parametric", ) - mock_calc.assert_called_once() - called_images, called_masks = mock_calc.call_args[0][:2] + mock_calculate_sample_weights.assert_called_once() + called_images, called_masks = mock_calculate_sample_weights.call_args[0][:2] np.testing.assert_array_equal(called_images, images) np.testing.assert_array_equal(called_masks, np.array(1 - raw_masks, dtype=bool)) def test_general_train_cycle_skips_sample_weights_when_disabled(mock_test_setup): - """general_train_cycle leaves sample_weight as None and never calls + """Test that general_train_cycle leaves sample_weight as None and never calls calculate_sample_weights when use_sample_weights is False.""" mock_model = mock_test_setup["mock_model"] with ( - patch("wf_psf.training.train_utils.calculate_sample_weights") as mock_calc, + patch( + "wf_psf.training.train_utils.calculate_sample_weights" + ) as mock_calculate_sample_weights, patch("wf_psf.training.train_utils.train_cycle_part") as mock_train_cycle_part, ): mock_train_cycle_part.return_value = MagicMock() @@ -421,7 +448,7 @@ def test_general_train_cycle_skips_sample_weights_when_disabled(mock_test_setup) cycle_def="parametric", ) - mock_calc.assert_not_called() + mock_calculate_sample_weights.assert_not_called() _, train_cycle_part_kwargs = mock_train_cycle_part.call_args assert train_cycle_part_kwargs["sample_weight"] is None diff --git a/src/wf_psf/tests/test_utils/noise_test.py b/src/wf_psf/tests/test_utils/noise_test.py index be458e63..6a20302d 100644 --- a/src/wf_psf/tests/test_utils/noise_test.py +++ b/src/wf_psf/tests/test_utils/noise_test.py @@ -14,6 +14,25 @@ def test_initialization(): assert estimator.window.shape == img_dim +def test_default_win_rad(): + """Test that the default_win_rad derives an int radius from the image dimensions.""" + assert NoiseEstimator.default_win_rad((32, 32)) == 10 + assert NoiseEstimator.default_win_rad((40, 40)) == 13 + assert NoiseEstimator.default_win_rad((100, 76)) == 31 + assert isinstance(NoiseEstimator.default_win_rad((32, 32)), int) + + +def test_initialization_uses_default_win_rad_when_omitted(): + """Test that omitting win_rad falls back to default_win_rad and produces the same window.""" + img_dim = (50, 50) + estimator = NoiseEstimator(img_dim) + explicit_estimator = NoiseEstimator(img_dim, NoiseEstimator.default_win_rad(img_dim)) + + assert estimator.win_rad == NoiseEstimator.default_win_rad(img_dim) + assert isinstance(estimator.win_rad, int) + np.testing.assert_array_equal(estimator.window, explicit_estimator.window) + + def test_init_window(): """Test that the exclusion window is correctly initialized.""" img_dim = (50, 50) From e8cf878c16cdf071a4134d2452dcc8c46df19607 Mon Sep 17 00:00:00 2001 From: chaitanya chawak Date: Thu, 24 Sep 2026 08:09:15 +0200 Subject: [PATCH 09/15] simplify _resolve_training_outputs by passing loss_name --- src/wf_psf/training/train_utils.py | 44 ++++++++++-------------------- 1 file changed, 14 insertions(+), 30 deletions(-) diff --git a/src/wf_psf/training/train_utils.py b/src/wf_psf/training/train_utils.py index 41661d0d..7365ade7 100644 --- a/src/wf_psf/training/train_utils.py +++ b/src/wf_psf/training/train_utils.py @@ -11,7 +11,7 @@ import numpy as np import tensorflow as tf -from typing import Optional, Callable, Union +from typing import Optional, Callable from wf_psf.psf_models.psf_models import compile_PSF_model from wf_psf.utils.noise import NoiseEstimator from wf_psf.utils.utils import generalised_sigmoid @@ -368,43 +368,25 @@ def configure_optimizer_and_loss( return optimizer, loss, metrics -def _is_masked_loss(loss: Union[str, Callable, None]) -> bool: - """ - Check whether the given loss corresponds to the masked mean squared error. - - Parameters - ---------- - loss: str, callable, optional - The loss function (or its name) used for training. - - Returns - ------- - bool - True if the loss is the ``masked_mean_squared_error``, False otherwise. - """ - if loss is None: - return False - return (isinstance(loss, str) and loss == "masked_mean_squared_error") or ( - hasattr(loss, "name") and loss.name == "masked_mean_squared_error" - ) - - def _resolve_training_outputs( - outputs: tf.Tensor, loss: Union[str, Callable, None] + outputs: tf.Tensor, loss_name: Optional[str] ) -> tuple[tf.Tensor, Optional[np.ndarray]]: """ - Resolve the training-specific ``outputs``/``loss`` representation into a - plain image and, for a masked loss, a per-image mask. + Resolve the training-specific ``outputs`` representation into a plain + image, and in case of a masked loss, a per-image mask. Parameters ---------- outputs: tf.Tensor - Image data. When ``loss`` is ``"masked_mean_squared_error"``, a 4D tensor + Image data. When ``loss_name`` starts with ``"masked_"``, a 4D tensor of shape ``(batch_size, height, width, 2)`` is expected, where the last dimension holds ``[image, mask]``. Otherwise, a 3D tensor of shape ``(batch_size, height, width)`` is expected. - loss: str, callable, optional - The loss function (or its name) used for training. + loss_name: str, optional + The name of the loss function used for training (e.g. ``loss.name``). + Any name starting with ``"masked_"`` (e.g. ``"masked_mean_squared_error"``) + is treated as a masked loss, so new masked losses work without updating + this function. Returns ------- @@ -412,7 +394,7 @@ def _resolve_training_outputs( The resolved images and, for a masked loss, the corresponding boolean masks (``outputs[..., 1]`` inverted); otherwise ``None``. """ - if _is_masked_loss(loss): + if loss_name is not None and loss_name.startswith("masked_"): images = outputs[..., 0] masks = np.array(1 - outputs[..., 1], dtype=bool) else: @@ -761,7 +743,9 @@ def general_train_cycle( # Calculate sample weights if use_sample_weights: - images, masks = _resolve_training_outputs(outputs, loss) + images, masks = _resolve_training_outputs( + outputs, getattr(loss, "name", loss) + ) sample_weight = calculate_sample_weights( images, masks, From 0d7fb3c8055964e57ebd83223c9b3b09a5b66ae6 Mon Sep 17 00:00:00 2001 From: chaitanya chawak Date: Thu, 24 Sep 2026 08:13:00 +0200 Subject: [PATCH 10/15] update unit tests --- .../tests/test_training/train_utils_test.py | 31 +++++++------------ 1 file changed, 11 insertions(+), 20 deletions(-) diff --git a/src/wf_psf/tests/test_training/train_utils_test.py b/src/wf_psf/tests/test_training/train_utils_test.py index 440c6e0d..fdf21862 100644 --- a/src/wf_psf/tests/test_training/train_utils_test.py +++ b/src/wf_psf/tests/test_training/train_utils_test.py @@ -180,39 +180,30 @@ def test_calculate_sample_weights_high_variance(mock_noise_estimator): @pytest.mark.parametrize( - "loss, expected", - [ - (None, False), - ("mean_squared_error", False), - ("masked_mean_squared_error", True), - (train_utils.MaskedMeanSquaredError(), True), - ], + "loss_name", [None, "mean_squared_error", "some_other_loss"] ) -def test_is_masked_loss(loss, expected): - """Test that _is_masked_loss identifies the masked_mean_squared_error loss by name.""" - assert train_utils._is_masked_loss(loss) is expected - - -def test_resolve_training_outputs_unmasked_loss(): - """Test that for a non-masked loss, outputs pass through unchanged and masks is None.""" +def test_resolve_training_outputs_unmasked_loss(loss_name): + """Test that for a non-"masked_"-prefixed loss name, outputs pass through unchanged and masks is None.""" outputs = np.random.rand(4, 8, 8) - images, masks = train_utils._resolve_training_outputs( - outputs, "mean_squared_error" - ) + images, masks = train_utils._resolve_training_outputs(outputs, loss_name) assert images is outputs assert masks is None -def test_resolve_training_outputs_masked_loss(): - """Test that for a masked loss, outputs are split into images and an inverted boolean mask.""" +@pytest.mark.parametrize( + "loss_name", ["masked_mean_squared_error", "masked_huber_error"] +) +def test_resolve_training_outputs_masked_loss(loss_name): + """Test that any loss name starting with "masked_" splits outputs into images + and an inverted boolean mask, so new masked losses work without code changes.""" batch_size, height, width = 4, 8, 8 images = np.random.rand(batch_size, height, width) raw_masks = np.random.randint(0, 2, size=(batch_size, height, width)) outputs = np.stack([images, raw_masks], axis=-1) resolved_images, resolved_masks = train_utils._resolve_training_outputs( - outputs, "masked_mean_squared_error" + outputs, loss_name ) np.testing.assert_array_equal(resolved_images, images) From 1f76342f97b46b1da1192dbb8470d223ed166636 Mon Sep 17 00:00:00 2001 From: chaitanya chawak Date: Thu, 24 Sep 2026 08:14:56 +0200 Subject: [PATCH 11/15] add changelog --- ..._estimation_from_sample_weight_computation.md | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) create mode 100644 changelog.d/20260909_155518_chaitanyachawak001_225_separate_noise_estimation_from_sample_weight_computation.md diff --git a/changelog.d/20260909_155518_chaitanyachawak001_225_separate_noise_estimation_from_sample_weight_computation.md b/changelog.d/20260909_155518_chaitanyachawak001_225_separate_noise_estimation_from_sample_weight_computation.md new file mode 100644 index 00000000..3dcec0ef --- /dev/null +++ b/changelog.d/20260909_155518_chaitanyachawak001_225_separate_noise_estimation_from_sample_weight_computation.md @@ -0,0 +1,16 @@ + + +### New features + +- Added `NoiseEstimator.estimate_noise_batch` for batch noise estimation, which is reusable outside training (e.g. SHE validation) since it no longer depends on the training loss/output representation. + +### Internal changes + +- Moved `NoiseEstimator` from `utils.utils` to a dedicated `utils.noise` module, and initialised it with a default exclusion-window radius (`NoiseEstimator.default_win_rad`). +- Simplified `train_utils.calculate_sample_weights` to `calculate_sample_weights(images, masks=None, ...)`. It no longer needs to know about `loss` or `use_sample_weights`. `general_train_cycle` now resolves the training output representation. +- Removed `train_utils.estimate_noise_sigma`, which is now redundant after the refactors. From 44e9aeb244bc3c4f4dd3a65106dbbc38ac339824 Mon Sep 17 00:00:00 2001 From: chaitanya chawak Date: Thu, 24 Sep 2026 08:50:44 +0200 Subject: [PATCH 12/15] update changelog --- ...5_separate_noise_estimation_from_sample_weight_computation.md | 1 - 1 file changed, 1 deletion(-) diff --git a/changelog.d/20260909_155518_chaitanyachawak001_225_separate_noise_estimation_from_sample_weight_computation.md b/changelog.d/20260909_155518_chaitanyachawak001_225_separate_noise_estimation_from_sample_weight_computation.md index 3dcec0ef..96504038 100644 --- a/changelog.d/20260909_155518_chaitanyachawak001_225_separate_noise_estimation_from_sample_weight_computation.md +++ b/changelog.d/20260909_155518_chaitanyachawak001_225_separate_noise_estimation_from_sample_weight_computation.md @@ -13,4 +13,3 @@ For top level release notes, leave all the headers commented out. - Moved `NoiseEstimator` from `utils.utils` to a dedicated `utils.noise` module, and initialised it with a default exclusion-window radius (`NoiseEstimator.default_win_rad`). - Simplified `train_utils.calculate_sample_weights` to `calculate_sample_weights(images, masks=None, ...)`. It no longer needs to know about `loss` or `use_sample_weights`. `general_train_cycle` now resolves the training output representation. -- Removed `train_utils.estimate_noise_sigma`, which is now redundant after the refactors. From ea78bedff5f0c753e639d5c58936d57f2331e550 Mon Sep 17 00:00:00 2001 From: chaitanya chawak Date: Thu, 24 Sep 2026 08:53:05 +0200 Subject: [PATCH 13/15] update test_resolve_training_outputs_unmasked_loss for null test fuzziness --- src/wf_psf/tests/test_training/train_utils_test.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/src/wf_psf/tests/test_training/train_utils_test.py b/src/wf_psf/tests/test_training/train_utils_test.py index fdf21862..a624640c 100644 --- a/src/wf_psf/tests/test_training/train_utils_test.py +++ b/src/wf_psf/tests/test_training/train_utils_test.py @@ -180,7 +180,14 @@ def test_calculate_sample_weights_high_variance(mock_noise_estimator): @pytest.mark.parametrize( - "loss_name", [None, "mean_squared_error", "some_other_loss"] + "loss_name", + [ + None, + "mean_squared_error", + "some_other_loss", + "non_masked_loss", + "unmasked_mse_loss", + ], ) def test_resolve_training_outputs_unmasked_loss(loss_name): """Test that for a non-"masked_"-prefixed loss name, outputs pass through unchanged and masks is None.""" From fa9d3e15d4269e584a6f2756f813f84ef8004cda Mon Sep 17 00:00:00 2001 From: chaitanya chawak Date: Thu, 24 Sep 2026 08:59:42 +0200 Subject: [PATCH 14/15] update test_resolve_training_outputs_unmasked_loss for null test fuzziness --- src/wf_psf/tests/test_training/train_utils_test.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/wf_psf/tests/test_training/train_utils_test.py b/src/wf_psf/tests/test_training/train_utils_test.py index a624640c..b412c946 100644 --- a/src/wf_psf/tests/test_training/train_utils_test.py +++ b/src/wf_psf/tests/test_training/train_utils_test.py @@ -187,6 +187,7 @@ def test_calculate_sample_weights_high_variance(mock_noise_estimator): "some_other_loss", "non_masked_loss", "unmasked_mse_loss", + "masked", ], ) def test_resolve_training_outputs_unmasked_loss(loss_name): From a2298ec9308b0048070b2f960f87870be968a196 Mon Sep 17 00:00:00 2001 From: chaitanya chawak Date: Thu, 24 Sep 2026 10:14:38 +0200 Subject: [PATCH 15/15] update changelog --- ..._separate_noise_estimation_from_sample_weight_computation.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/changelog.d/20260909_155518_chaitanyachawak001_225_separate_noise_estimation_from_sample_weight_computation.md b/changelog.d/20260909_155518_chaitanyachawak001_225_separate_noise_estimation_from_sample_weight_computation.md index 96504038..82b6f9c6 100644 --- a/changelog.d/20260909_155518_chaitanyachawak001_225_separate_noise_estimation_from_sample_weight_computation.md +++ b/changelog.d/20260909_155518_chaitanyachawak001_225_separate_noise_estimation_from_sample_weight_computation.md @@ -12,4 +12,4 @@ For top level release notes, leave all the headers commented out. ### Internal changes - Moved `NoiseEstimator` from `utils.utils` to a dedicated `utils.noise` module, and initialised it with a default exclusion-window radius (`NoiseEstimator.default_win_rad`). -- Simplified `train_utils.calculate_sample_weights` to `calculate_sample_weights(images, masks=None, ...)`. It no longer needs to know about `loss` or `use_sample_weights`. `general_train_cycle` now resolves the training output representation. +- Simplified `train_utils.calculate_sample_weights` to `calculate_sample_weights(images, masks=None, ...)`. It no longer needs to know about `loss` or `use_sample_weights`. `general_train_cycle` now resolves the training output representation based on `loss.name`: any name starting with the `masked_` prefix (e.g. `masked_mean_squared_error`) is treated as a masked loss, so new masked loss functions must follow this naming convention to be handled correctly.