diff --git a/src/ellphi/__init__.py b/src/ellphi/__init__.py index 9a58db7..fe3e6f0 100644 --- a/src/ellphi/__init__.py +++ b/src/ellphi/__init__.py @@ -28,7 +28,7 @@ coef_from_axes, coef_from_cov, ) -from .ellcloud import ellipse_cloud, EllipseCloud, LocalCov +from .ellcloud import ellipse_cloud, EllipseCloud, LocalCov, RescaleDiagnostics # solver from .solver import ( @@ -60,6 +60,7 @@ "ellipse_cloud", "EllipseCloud", "LocalCov", + "RescaleDiagnostics", # solver "quad_eval", "pencil", diff --git a/src/ellphi/__init__.pyi b/src/ellphi/__init__.pyi index d98f65c..e608f45 100644 --- a/src/ellphi/__init__.pyi +++ b/src/ellphi/__init__.pyi @@ -6,6 +6,7 @@ FloatArray = NDArray[np.float64] from .ellcloud import ( EllipseCloud as EllipseCloud, LocalCov as LocalCov, + RescaleDiagnostics as RescaleDiagnostics, ellipse_cloud as ellipse_cloud, ) from .geometry import ( @@ -44,6 +45,7 @@ __all__ = [ "ellipse_cloud", "EllipseCloud", "LocalCov", + "RescaleDiagnostics", "quad_eval", "pencil", "tangency", diff --git a/src/ellphi/ellcloud.py b/src/ellphi/ellcloud.py index 01df549..23848c9 100644 --- a/src/ellphi/ellcloud.py +++ b/src/ellphi/ellcloud.py @@ -19,7 +19,22 @@ from .geometry import axes_from_cov, coef_from_cov, infer_dim_from_coef_length from .solver import pdist_tangency -__all__ = ["ellipse_cloud", "EllipseCloud", "LocalCov"] +__all__ = ["ellipse_cloud", "EllipseCloud", "LocalCov", "RescaleDiagnostics"] + + +@dataclass(frozen=True) +class RescaleDiagnostics: + """Diagnostics returned by :meth:`EllipseCloud.rescale`. + + Attributes: + scale: The scaling factor applied to the cloud. + pre_summary: Per-axis aggregate semi-axis lengths before rescaling. + post_summary: Per-axis aggregate semi-axis lengths after rescaling. + """ + + scale: float + pre_summary: numpy.ndarray + post_summary: numpy.ndarray @dataclass @@ -312,7 +327,7 @@ def from_local_cov( """ return LocalCov(k=k)(X) - def rescale(self, *, method="median") -> float: + def rescale(self, *, method="median", return_diagnostics=False): """Applies rescaling to all the ellipses in the cloud. This method rescales all the ellipses in the cloud based on the @@ -320,9 +335,13 @@ def rescale(self, *, method="median") -> float: Args: method: The rescaling method to use. Can be "median" or "average". + return_diagnostics: If ``True``, return a + :class:`RescaleDiagnostics` object instead of just the scale + factor. Default is ``False`` for backward compatibility. Returns: - The scaling factor used to rescale the ellipses. + The scaling factor (float) when *return_diagnostics* is ``False``, + or a :class:`RescaleDiagnostics` object when it is ``True``. """ if self.n_dim != 2: raise NotImplementedError( @@ -342,6 +361,12 @@ def rescale(self, *, method="median") -> float: ell_scale = ell_scales[1] ** 2 / ell_scales[0] self.cov /= ell_scale**2 self.coef *= ell_scale**2 + if return_diagnostics: + return RescaleDiagnostics( + scale=float(ell_scale), + pre_summary=ell_scales, + post_summary=ell_scales / ell_scale, + ) return float(ell_scale) diff --git a/src/ellphi/ellcloud.pyi b/src/ellphi/ellcloud.pyi index 108c98d..bb4ec8e 100644 --- a/src/ellphi/ellcloud.pyi +++ b/src/ellphi/ellcloud.pyi @@ -1,13 +1,19 @@ from __future__ import annotations from __future__ import annotations -from typing import Any, Iterator, Sequence +from typing import Any, Iterator, Literal, Sequence, overload from dataclasses import dataclass, field import matplotlib.pyplot as plt import numpy as np from numpy.typing import NDArray from ellphi import FloatArray -__all__ = ["ellipse_cloud", "EllipseCloud", "LocalCov"] +__all__ = ["ellipse_cloud", "EllipseCloud", "LocalCov", "RescaleDiagnostics"] + +@dataclass(frozen=True) +class RescaleDiagnostics: + scale: float + pre_summary: FloatArray + post_summary: FloatArray @dataclass class EllipseCloud: @@ -54,7 +60,18 @@ class EllipseCloud: ) -> EllipseCloud: ... @classmethod def from_local_cov(cls, X: FloatArray, *, k: int = 5) -> EllipseCloud: ... - def rescale(self, *, method: str = "median") -> float: ... + @overload + def rescale( + self, *, method: str = ..., return_diagnostics: Literal[False] = ... + ) -> float: ... + @overload + def rescale( + self, *, method: str = ..., return_diagnostics: Literal[True] + ) -> RescaleDiagnostics: ... + @overload + def rescale( + self, *, method: str = ..., return_diagnostics: bool + ) -> float | RescaleDiagnostics: ... def ellipse_cloud( X: FloatArray, diff --git a/tests/test_ellcloud_rescale.py b/tests/test_ellcloud_rescale.py index 2c48c85..54d49dd 100644 --- a/tests/test_ellcloud_rescale.py +++ b/tests/test_ellcloud_rescale.py @@ -4,6 +4,8 @@ import numpy as np +from ellphi import RescaleDiagnostics + from .factories import random_cloud @@ -21,3 +23,43 @@ def test_rescale_returns_float_and_updates_arrays(): assert isinstance(scale, float) np.testing.assert_allclose(cloud.cov, cov_before / scale**2) np.testing.assert_allclose(cloud.coef, coef_before * scale**2) + + +def test_rescale_diagnostics_type_and_scale_consistency(): + """``rescale(return_diagnostics=True)`` returns a RescaleDiagnostics.""" + + rng = np.random.default_rng(2025) + cloud = random_cloud(rng, n_ellipses=8) + + diag = cloud.rescale(return_diagnostics=True) + + assert isinstance(diag, RescaleDiagnostics) + assert isinstance(diag.scale, float) + assert diag.pre_summary.shape == (2,) + assert diag.post_summary.shape == (2,) + + +def test_rescale_diagnostics_pre_post_relationship(): + """post_summary equals pre_summary / scale.""" + + rng = np.random.default_rng(2026) + cloud = random_cloud(rng, n_ellipses=10) + + diag = cloud.rescale(return_diagnostics=True) + + np.testing.assert_allclose(diag.post_summary, diag.pre_summary / diag.scale) + + +def test_rescale_default_matches_diagnostics_scale(): + """Default (no diagnostics) scale matches diagnostics.scale.""" + + rng = np.random.default_rng(2027) + cloud_a = random_cloud(rng, n_ellipses=6) + + rng2 = np.random.default_rng(2027) + cloud_b = random_cloud(rng2, n_ellipses=6) + + scale = cloud_a.rescale() + diag = cloud_b.rescale(return_diagnostics=True) + + assert scale == diag.scale