API Reference

Main Classes

class wl_stats_torch.WLStatistics(n_scales=5, device=None, pixel_arcmin=1.0, dtype=torch.float64)[source]

Bases: object

Complete weak lensing statistics calculator.

This class provides all the functionality from the original CosmoStat HOS_starlet_l1norm_peaks class, but implemented in PyTorch for GPU acceleration.

Parameters:

n_scales (int)
device (device | None)
pixel_arcmin (float)
dtype (dtype)

n_scales

Number of wavelet scales

Type:: int

device

Computation device (cpu or cuda)

Type:: torch.device

starlet

Starlet transform instance

Type:: Starlet2D

pixel_arcmin

Pixel resolution in arcminutes

Type:: float

Example

>>> stats = WLStatistics(n_scales=5, device='cuda')
>>> results = stats.compute_all_statistics(kappa_map, sigma_map)
>>> peak_counts = results['wavelet_peak_counts']
>>> l1_norms = results['wavelet_l1_norms']

__init__(n_scales=5, device=None, pixel_arcmin=1.0, dtype=torch.float64)[source]

Initialize weak lensing statistics calculator.

Parameters:

n_scales (int) – Number of wavelet scales (including coarse)
device (Optional[device]) – torch device for computation. If None, auto-detects.
pixel_arcmin (float) – Pixel resolution in arcminutes
dtype (dtype) – Data type for computations. Default: torch.float64 to match NumPy.

get_scale_resolutions()[source]

Get effective resolution of each scale in arcminutes.

Return type:: List[float]

compute_wavelet_transform(image, noise_sigma, mask=None, subtract_coarse_mean=True)[source]

Compute wavelet transform and SNR for an image or batch of images.

Parameters:

image (Tensor) – Input convergence map(s), shape (H, W) or (B, H, W)
noise_sigma (Union[float, Tensor]) – Noise standard deviation, scalar, (H, W), or (B, H, W)
mask (Optional[Tensor]) – Optional observation mask, None, (H, W), or (B, H, W)
subtract_coarse_mean (bool) – If True (default), subtract the spatial mean from the coarse scale before computing SNR. This removes the cosmologically uninformative DC component.

Returns:

‘wavelet_coeffs’: Wavelet coefficients (n_scales, H, W) or (B, n_scales, H, W)
’noise_levels’: Noise std for each coefficient (n_scales, H, W) or (B, n_scales, H, W)
’snr’: Signal-to-noise ratio (n_scales, H, W) or (B, n_scales, H, W)

Return type:

Dictionary with keys

compute_wavelet_peak_counts(min_snr=-2.0, max_snr=6.0, n_bins=31, mask=None, verbose=False, clamp_overflow=False)[source]

Compute histogram of peak counts at all wavelet scales.

Parameters:

min_snr (float) – Minimum SNR for histogram bins
max_snr (float) – Maximum SNR for histogram bins
n_bins (int) – Number of histogram bins
mask (Optional[Tensor]) – Optional mask to restrict peak detection (H, W) or (B, H, W)
verbose (bool) – Print min/max SNR at each scale
clamp_overflow (bool) – If True, peaks outside SNR range are included in edge bins. If False (default), peaks outside range are excluded. False matches cosmostat/pycs behavior.

Returns:

bin_centers: Tensor of shape (n_bins,)
peak_counts_list: List of tensors per scale, shape (n_bins,) or (B, n_bins)

Return type:

Tuple of (bin_centers, peak_counts_list) where

compute_wavelet_l1_norms(n_bins=40, mask=None, min_snr=None, max_snr=None, clamp_overflow=False)[source]

Compute L1-norm as a function of SNR threshold for each scale.

For each scale, we bin the SNR values and compute the sum of absolute wavelet coefficients in each bin.

Parameters:

n_bins (int) – Number of bins for SNR binning
mask (Optional[Tensor]) – Optional mask to restrict calculation (H, W) or (B, H, W)
min_snr (Optional[float]) – Minimum SNR (if None, uses data minimum)
max_snr (Optional[float]) – Maximum SNR (if None, uses data maximum)
clamp_overflow (bool) – If True, values outside SNR range are included in edge bins. If False (default), values outside range are excluded. False matches cosmostat/pycs behavior.

Returns:

bins_list: List of bin center tensors per scale
l1_norms_list: List of L1-norm tensors per scale, shape (n_bins,) or (B, n_bins)

Return type:

Tuple of (bins_list, l1_norms_list) where

compute_mono_scale_peaks(image, noise_sigma, smoothing_sigma=2.0, min_snr=-2.0, max_snr=6.0, n_bins=31, mask=None, clamp_overflow=False)[source]

Compute mono-scale peak counts with Gaussian smoothing.

Parameters:

image (Tensor) – Input convergence map (H, W) or (B, H, W)
noise_sigma (Union[float, Tensor]) – Standard deviation of noise, scalar or tensor
smoothing_sigma (float) – Gaussian smoothing scale in pixels
min_snr (float) – Minimum SNR for histogram
max_snr (float) – Maximum SNR for histogram
n_bins (int) – Number of histogram bins
mask (Optional[Tensor]) – Optional observation mask (H, W) or (B, H, W)
clamp_overflow (bool) – If True, peaks outside SNR range are included in edge bins. If False (default), peaks outside range are excluded. False matches cosmostat/pycs behavior.

Return type:

Tuple[Tensor, Tensor]

Returns:

Tuple of (bin_centers, peak_counts)

peak_counts shape: (n_bins,) or (B, n_bins)

compute_all_statistics(image, noise_sigma, mask=None, min_snr=-2.0, max_snr=6.0, n_bins=31, l1_nbins=40, l1_min_snr=None, l1_max_snr=None, compute_mono=True, mono_smoothing_sigma=2.0, verbose=False, clamp_overflow=False, subtract_coarse_mean=True)[source]

Compute all weak lensing statistics in one call.

Supports both single images (H, W) and batches (B, H, W).

Parameters:

image (Tensor) – Convergence map (H, W) or batch (B, H, W)
noise_sigma (Union[float, Tensor]) – Noise std, scalar, (H, W), or (B, H, W)
mask (Optional[Tensor]) – Optional observation mask, None, (H, W), or (B, H, W)
min_snr (float) – Minimum SNR for peak count histograms
max_snr (float) – Maximum SNR for peak count histograms
n_bins (int) – Number of bins for peak histograms
l1_nbins (int) – Number of bins for L1-norm
l1_min_snr (Optional[float]) – Minimum SNR for L1-norm (if None, uses min_snr)
l1_max_snr (Optional[float]) – Maximum SNR for L1-norm (if None, uses max_snr)
compute_mono (bool) – Whether to compute mono-scale peaks
mono_smoothing_sigma (float) – Smoothing scale for mono-scale peaks
verbose (bool) – Print progress information
clamp_overflow (bool) – If True, values outside SNR range are included in edge bins. If False (default), values outside range are excluded. False matches cosmostat/pycs behavior.
subtract_coarse_mean (bool) – If True (default), subtract the spatial mean from the coarse scale before computing SNR.

Return type:

Dict[str, any]

Returns:

Dictionary with all computed statistics. For batched input (B, H, W):

’wavelet_coeffs’: (B, n_scales, H, W)

’wavelet_peak_counts’: list of (B, n_bins) per scale

’wavelet_l1_norms’: list of (B, l1_nbins) per scale

’mono_peak_counts’: (B, n_bins) if compute_mono=True

For single input (H, W):

’wavelet_coeffs’: (n_scales, H, W)
’wavelet_peak_counts’: list of (n_bins,) per scale
’wavelet_l1_norms’: list of (l1_nbins,) per scale
’mono_peak_counts’: (n_bins,) if compute_mono=True

to(device)[source]

Move all components to specified device.

Parameters:: device (device)

Wavelet Transform

2D Starlet (à trous wavelet) Transform in PyTorch

This module implements the 2D starlet transform using PyTorch for GPU acceleration. The starlet transform is an isotropic undecimated wavelet transform using a B3-spline scaling function.

References

Starck, J.-L., Fadili, J., & Murtagh, F. (2007). “The Undecimated Wavelet Decomposition and its Reconstruction” IEEE Transactions on Image Processing, 16(2), 297-309.

wl_stats_torch.starlet.fft_convolve2d(signal, kernel)[source]

Perform 2D convolution using FFT (equivalent to scipy.signal.fftconvolve with mode=’same’).

This function stays entirely in PyTorch, avoiding CPU transfers and numpy conversions. It matches scipy’s behavior for ‘same’ mode convolution.

Parameters:

signal (Tensor) – Input signal, shape (H, W)
kernel (Tensor) – Convolution kernel, shape (H, W)

Return type:

Tensor

Returns:

Convolved result, shape (H, W), same size as input signal

class wl_stats_torch.starlet.Starlet2D(n_scales=5, device=None, dtype=torch.float32)[source]

Bases: Module

2D Starlet Transform (à trous algorithm) with B3-spline kernel.

This implements an undecimated wavelet transform where each scale is computed by taking the difference between successive smoothed versions of the image. The smoothing uses a B3-spline kernel with increasing dilation (holes).

The transform decomposes an image into multiple detail (wavelet) scales and a final coarse (approximation) scale.

Parameters:

n_scales (int)
device (device | None)
dtype (dtype)

n_scales

Total number of scales (including coarse scale)

Type:: int

kernel_2d

The 2D B3-spline convolution kernel

Type:: torch.Tensor

device

Device for computation (cpu or cuda)

Type:: torch.device

Example

>>> starlet = Starlet2D(n_scales=5)
>>> image = torch.randn(1, 1, 256, 256)  # (batch, channel, H, W)
>>> coeffs = starlet(image)  # (batch, n_scales, H, W)

__init__(n_scales=5, device=None, dtype=torch.float32)[source]

Initialize the Starlet2D transform.

Parameters:

n_scales (int) – Total number of scales (detail scales + 1 coarse scale). For example, n_scales=5 gives 4 detail scales and 1 coarse scale.
device (Optional[device]) – torch device for computation. If None, uses CPU.
dtype (dtype) – Data type for computations. Default: torch.float32 (PyTorch standard).

forward(x, return_coarse=True, return_dict=False)[source]

Apply the forward Starlet transform.

Parameters:

x (Tensor) – Input tensor of shape (B, C, H, W) or (H, W) or (C, H, W). For single-channel input, C should be 1.
return_coarse (bool) – If True, include the final coarse scale in output.
return_dict (bool) – If True, return a dictionary with additional info.

Returns:

Wavelet coefficients of shape (B, n_scales, H, W) if return_coarse=True,

or (B, n_scales-1, H, W) if return_coarse=False.

If return_dict is True:

Dictionary with keys:: ’coeffs’: wavelet coefficients ‘detail_scales’: list of detail scale tensors ‘coarse_scale’: final coarse scale tensor

Return type:

If return_dict is False

Raises:

ValueError – If input has more than 1 channel (excluding batch dimension)

reconstruct(wavelet_coeffs, gen2=True)[source]

Reconstruct image from wavelet coefficients.

Parameters:

wavelet_coeffs (Tensor) – Wavelet coefficients of shape (B, n_scales, H, W)
gen2 (bool) – If True, use second generation reconstruction (default). If False, use first generation (simple sum).

Return type:

Tensor

Returns:

Reconstructed image of shape (B, 1, H, W)

get_scale_resolution(pixel_size_arcmin=1.0)[source]

Get the effective resolution (FWHM) of each wavelet scale in arcminutes.

Parameters:: pixel_size_arcmin (float) – Size of one pixel in arcminutes
Return type:: list
Returns:: List of resolutions for each scale (including coarse)

get_noise_levels(noise_sigma, mask=None)[source]

Compute noise standard deviation for each wavelet coefficient.

This uses the same approach as CosmoStat: compute the wavelet impulse response (by transforming a delta function), square it, and convolve with the variance map.

Parameters:

noise_sigma (Tensor) – Noise std deviation map, shape (H, W) or (B, 1, H, W)
mask (Optional[Tensor]) – Optional mask where 1 = observed, 0 = not observed

Return type:

Tensor

Returns:

Noise levels for each coefficient, shape (B, n_scales, H, W)

get_snr(image, noise_sigma, mask=None, keep_sign=False)[source]

Compute Signal-to-Noise Ratio (SNR) for wavelet coefficients.

Parameters:

image (Tensor) – Input image, shape (H, W) or (B, 1, H, W)
noise_sigma (Tensor) – Noise std deviation map
mask (Optional[Tensor]) – Optional observation mask
keep_sign (bool) – If True, preserve sign of coefficients in SNR

Return type:

Tensor

Returns:

SNR map for each scale, shape (B, n_scales, H, W)

extra_repr()[source]

String representation for printing.

Return type:: str

wl_stats_torch.starlet.test_starlet()[source]: Quick test of Starlet2D implementation.

Peak Detection

Peak Detection for 2D Images in PyTorch

This module provides fast, GPU-accelerated peak detection for 2D images. A peak is defined as a local maximum - a pixel with a value greater than all of its 8 neighbors.

wl_stats_torch.peaks.find_peaks_2d(image, threshold=None, mask=None, include_border=False, ordered=True)[source]

Find local maxima (peaks) in a 2D image.

A peak is defined as a pixel with value strictly greater than all 8 neighbors.

Parameters:

image (Tensor) – 2D tensor (H, W) or (1, 1, H, W)
threshold (Optional[float]) – Minimum value to consider as peak. If None, uses image minimum.
mask (Optional[Tensor]) – Optional binary mask (H, W) where 1 = consider, 0 = ignore
include_border (bool) – If True, include peaks on the image border
ordered (bool) – If True, return peaks sorted by height (descending)

Returns:

positions: Tensor of shape (N, 2) with (row, col) coordinates
heights: Tensor of shape (N,) with peak values

Return type:

Tuple of (positions, heights)

Example

>>> image = torch.randn(128, 128)
>>> positions, heights = find_peaks_2d(image, threshold=2.0)
>>> print(f"Found {len(positions)} peaks above threshold")

wl_stats_torch.peaks.find_peaks_batch(images, threshold=None, masks=None, include_border=False, ordered=True)[source]

Find peaks in a batch of images (VECTORIZED VERSION).

This implementation processes all images in the batch in parallel, avoiding sequential loops for better GPU utilization.

Parameters:

images (Tensor) – Tensor of shape (B, 1, H, W) or (B, H, W)
threshold (Optional[float]) – Minimum value to consider as peak
masks (Optional[Tensor]) – Optional masks of shape (B, 1, H, W) or (B, H, W)
include_border (bool) – If True, include peaks on borders
ordered (bool) – If True, sort peaks by height

Return type:

List[Tuple[Tensor, Tensor]]

Returns:

List of (positions, heights) tuples, one per image in batch

wl_stats_torch.peaks.peaks_to_histogram(peak_heights, bins, digitize_mode=True, clamp_overflow=False)[source]

Compute histogram of peak heights.

This function mimics np.histogram behavior to match pycs output.

Parameters:

peak_heights (Tensor) – Tensor of peak values, shape (N,)
bins (Tensor) – Bin edges, shape (n_bins+1,)
digitize_mode (bool) – If True, use np.digitize-like behavior (default). If False, use torch.histogram behavior.
clamp_overflow (bool) – If True, values outside bin range are included in edge bins. If False (default), values outside range are excluded. False matches cosmostat/pycs behavior.

Return type:

Tensor

Returns:

Histogram counts, shape (n_bins,)

Note

To match pycs behavior with np.histogram: - Values x where bins[i] <= x < bins[i+1] go into bin i - The rightmost bin includes the right edge: bins[-2] <= x <= bins[-1] - When clamp_overflow=False: values outside [bins[0], bins[-1]] are excluded - When clamp_overflow=True: values < bins[0] go to first bin, > bins[-1] go to last bin

wl_stats_torch.peaks.mono_scale_peaks_smoothed(image, sigma_noise, smoothing_sigma=2.0, mask=None, bins=None, min_snr=-2.0, max_snr=6.0, n_bins=31, clamp_overflow=False)[source]

Compute mono-scale peak counts with Gaussian smoothing.

This applies Gaussian smoothing to the image, computes SNR, finds peaks, and returns histogram of peak counts.

NOTE: sigma_noise can now be a tensor for spatially-varying noise maps.

Parameters:

image (Tensor) – Input image (H, W)
sigma_noise (float) – Standard deviation of noise (scalar or H, W tensor)
smoothing_sigma (float) – Std dev for Gaussian smoothing (in pixels)
mask (Optional[Tensor]) – Optional observation mask
bins (Optional[Tensor]) – Optional custom bin edges for histogram
min_snr (float) – Minimum SNR for histogram (if bins not provided)
max_snr (float) – Maximum SNR for histogram (if bins not provided)
n_bins (int) – Number of bins for histogram (if bins not provided)
clamp_overflow (bool) – If True, peaks outside SNR range are included in edge bins. If False (default), peaks outside range are excluded. False matches cosmostat/pycs behavior.

Return type:

Tuple[Tensor, Tensor, Tuple[Tensor, Tensor]]

Returns:

Tuple of (bin_centers, counts, (peak_positions, peak_heights))

wl_stats_torch.peaks.test_peaks()[source]: Test peak detection functions.

Statistics

Weak Lensing Statistics Module

Main class for computing weak lensing summary statistics including: - Wavelet peak counts at multiple scales - Wavelet L1-norm statistics - Mono-scale peak counts with Gaussian smoothing

class wl_stats_torch.statistics.WLStatistics(n_scales=5, device=None, pixel_arcmin=1.0, dtype=torch.float64)[source]

Bases: object

Complete weak lensing statistics calculator.

This class provides all the functionality from the original CosmoStat HOS_starlet_l1norm_peaks class, but implemented in PyTorch for GPU acceleration.

Parameters:

n_scales (int)
device (device | None)
pixel_arcmin (float)
dtype (dtype)

n_scales

Number of wavelet scales

Type:: int

device

Computation device (cpu or cuda)

Type:: torch.device

starlet

Starlet transform instance

Type:: Starlet2D

pixel_arcmin

Pixel resolution in arcminutes

Type:: float

Example

>>> stats = WLStatistics(n_scales=5, device='cuda')
>>> results = stats.compute_all_statistics(kappa_map, sigma_map)
>>> peak_counts = results['wavelet_peak_counts']
>>> l1_norms = results['wavelet_l1_norms']

__init__(n_scales=5, device=None, pixel_arcmin=1.0, dtype=torch.float64)[source]

Initialize weak lensing statistics calculator.

Parameters:

n_scales (int) – Number of wavelet scales (including coarse)
device (Optional[device]) – torch device for computation. If None, auto-detects.
pixel_arcmin (float) – Pixel resolution in arcminutes
dtype (dtype) – Data type for computations. Default: torch.float64 to match NumPy.

get_scale_resolutions()[source]

Get effective resolution of each scale in arcminutes.

Return type:: List[float]

compute_wavelet_transform(image, noise_sigma, mask=None, subtract_coarse_mean=True)[source]

Compute wavelet transform and SNR for an image or batch of images.

Parameters:

image (Tensor) – Input convergence map(s), shape (H, W) or (B, H, W)
noise_sigma (Union[float, Tensor]) – Noise standard deviation, scalar, (H, W), or (B, H, W)
mask (Optional[Tensor]) – Optional observation mask, None, (H, W), or (B, H, W)
subtract_coarse_mean (bool) – If True (default), subtract the spatial mean from the coarse scale before computing SNR. This removes the cosmologically uninformative DC component.

Returns:

‘wavelet_coeffs’: Wavelet coefficients (n_scales, H, W) or (B, n_scales, H, W)
’noise_levels’: Noise std for each coefficient (n_scales, H, W) or (B, n_scales, H, W)
’snr’: Signal-to-noise ratio (n_scales, H, W) or (B, n_scales, H, W)

Return type:

Dictionary with keys

compute_wavelet_peak_counts(min_snr=-2.0, max_snr=6.0, n_bins=31, mask=None, verbose=False, clamp_overflow=False)[source]

Compute histogram of peak counts at all wavelet scales.

Parameters:

min_snr (float) – Minimum SNR for histogram bins
max_snr (float) – Maximum SNR for histogram bins
n_bins (int) – Number of histogram bins
mask (Optional[Tensor]) – Optional mask to restrict peak detection (H, W) or (B, H, W)
verbose (bool) – Print min/max SNR at each scale
clamp_overflow (bool) – If True, peaks outside SNR range are included in edge bins. If False (default), peaks outside range are excluded. False matches cosmostat/pycs behavior.

Returns:

bin_centers: Tensor of shape (n_bins,)
peak_counts_list: List of tensors per scale, shape (n_bins,) or (B, n_bins)

Return type:

Tuple of (bin_centers, peak_counts_list) where

compute_wavelet_l1_norms(n_bins=40, mask=None, min_snr=None, max_snr=None, clamp_overflow=False)[source]

Compute L1-norm as a function of SNR threshold for each scale.

For each scale, we bin the SNR values and compute the sum of absolute wavelet coefficients in each bin.

Parameters:

n_bins (int) – Number of bins for SNR binning
mask (Optional[Tensor]) – Optional mask to restrict calculation (H, W) or (B, H, W)
min_snr (Optional[float]) – Minimum SNR (if None, uses data minimum)
max_snr (Optional[float]) – Maximum SNR (if None, uses data maximum)
clamp_overflow (bool) – If True, values outside SNR range are included in edge bins. If False (default), values outside range are excluded. False matches cosmostat/pycs behavior.

Returns:

bins_list: List of bin center tensors per scale
l1_norms_list: List of L1-norm tensors per scale, shape (n_bins,) or (B, n_bins)

Return type:

Tuple of (bins_list, l1_norms_list) where

compute_mono_scale_peaks(image, noise_sigma, smoothing_sigma=2.0, min_snr=-2.0, max_snr=6.0, n_bins=31, mask=None, clamp_overflow=False)[source]

Compute mono-scale peak counts with Gaussian smoothing.

Parameters:

image (Tensor) – Input convergence map (H, W) or (B, H, W)
noise_sigma (Union[float, Tensor]) – Standard deviation of noise, scalar or tensor
smoothing_sigma (float) – Gaussian smoothing scale in pixels
min_snr (float) – Minimum SNR for histogram
max_snr (float) – Maximum SNR for histogram
n_bins (int) – Number of histogram bins
mask (Optional[Tensor]) – Optional observation mask (H, W) or (B, H, W)
clamp_overflow (bool) – If True, peaks outside SNR range are included in edge bins. If False (default), peaks outside range are excluded. False matches cosmostat/pycs behavior.

Return type:

Tuple[Tensor, Tensor]

Returns:

Tuple of (bin_centers, peak_counts)

peak_counts shape: (n_bins,) or (B, n_bins)

compute_all_statistics(image, noise_sigma, mask=None, min_snr=-2.0, max_snr=6.0, n_bins=31, l1_nbins=40, l1_min_snr=None, l1_max_snr=None, compute_mono=True, mono_smoothing_sigma=2.0, verbose=False, clamp_overflow=False, subtract_coarse_mean=True)[source]

Compute all weak lensing statistics in one call.

Supports both single images (H, W) and batches (B, H, W).

Parameters:

image (Tensor) – Convergence map (H, W) or batch (B, H, W)
noise_sigma (Union[float, Tensor]) – Noise std, scalar, (H, W), or (B, H, W)
mask (Optional[Tensor]) – Optional observation mask, None, (H, W), or (B, H, W)
min_snr (float) – Minimum SNR for peak count histograms
max_snr (float) – Maximum SNR for peak count histograms
n_bins (int) – Number of bins for peak histograms
l1_nbins (int) – Number of bins for L1-norm
l1_min_snr (Optional[float]) – Minimum SNR for L1-norm (if None, uses min_snr)
l1_max_snr (Optional[float]) – Maximum SNR for L1-norm (if None, uses max_snr)
compute_mono (bool) – Whether to compute mono-scale peaks
mono_smoothing_sigma (float) – Smoothing scale for mono-scale peaks
verbose (bool) – Print progress information
clamp_overflow (bool) – If True, values outside SNR range are included in edge bins. If False (default), values outside range are excluded. False matches cosmostat/pycs behavior.
subtract_coarse_mean (bool) – If True (default), subtract the spatial mean from the coarse scale before computing SNR.

Return type:

Dict[str, any]

Returns:

Dictionary with all computed statistics. For batched input (B, H, W):

’wavelet_coeffs’: (B, n_scales, H, W)

’wavelet_peak_counts’: list of (B, n_bins) per scale

’wavelet_l1_norms’: list of (B, l1_nbins) per scale

’mono_peak_counts’: (B, n_bins) if compute_mono=True

For single input (H, W):

’wavelet_coeffs’: (n_scales, H, W)
’wavelet_peak_counts’: list of (n_bins,) per scale
’wavelet_l1_norms’: list of (l1_nbins,) per scale
’mono_peak_counts’: (n_bins,) if compute_mono=True

to(device)[source]

Move all components to specified device.

Parameters:: device (device)

wl_stats_torch.statistics.test_statistics()[source]: Test the complete statistics pipeline.

Visualization

Visualization utilities for weak lensing statistics.

wl_stats_torch.visualization.plot_peak_histograms(bin_centers, peak_counts, scale_labels=None, title='Wavelet Peak Counts', xlabel='SNR', ylabel='Peak Counts', log_scale=True, figsize=(10, 6), save_path=None)[source]

Plot peak count histograms for multiple scales.

Parameters:

bin_centers (Tensor) – Bin centers, shape (n_bins,)
peak_counts (List[Tensor]) – List of peak counts per scale
scale_labels (Optional[List[str]]) – Optional labels for each scale
title (str) – Plot title
xlabel (str) – X-axis label
ylabel (str) – Y-axis label
log_scale (bool) – Use logarithmic y-axis
figsize (Tuple[int, int]) – Figure size (width, height)
save_path (Optional[str]) – If provided, save figure to this path

wl_stats_torch.visualization.plot_l1_norms(l1_bins, l1_norms, scale_labels=None, title='Wavelet L1-Norms', xlabel='SNR', ylabel='L1-Norm', log_scale=False, xlim=None, figsize=(10, 6), save_path=None)[source]

Plot L1-norm as a function of SNR for multiple scales.

Parameters:

l1_bins (List[Tensor]) – List of bin centers per scale
l1_norms (List[Tensor]) – List of L1-norms per scale
scale_labels (Optional[List[str]]) – Optional labels for each scale
title (str) – Plot title
xlabel (str) – X-axis label
ylabel (str) – Y-axis label
log_scale (bool) – Use logarithmic y-axis
xlim (Optional[Tuple[float, float]]) – X-axis limits (min, max)
figsize (Tuple[int, int]) – Figure size
save_path (Optional[str]) – If provided, save figure to this path

wl_stats_torch.visualization.plot_wavelet_scales(wavelet_coeffs, peak_positions=None, titles=None, cmap='viridis', vmin=None, vmax=None, figsize=(15, 10), mark_peaks=True, save_path=None)[source]

Visualize wavelet scales with optional peak markers.

Parameters:

wavelet_coeffs (Tensor) – Wavelet coefficients (n_scales, H, W)
peak_positions (Optional[List[Tensor]]) – Optional list of peak positions per scale
titles (Optional[List[str]]) – Optional titles for each scale
cmap (str) – Colormap name
vmin (Optional[float]) – Minimum value for colorscale
vmax (Optional[float]) – Maximum value for colorscale
figsize (Tuple[int, int]) – Figure size
mark_peaks (bool) – Whether to mark peak positions
save_path (Optional[str]) – If provided, save figure to this path

wl_stats_torch.visualization.plot_snr_map(snr_coeffs, scale_idx=0, peak_positions=None, title=None, cmap='RdBu_r', vmin=-5, vmax=5, figsize=(10, 8), save_path=None)[source]

Plot SNR map for a specific scale with optional peak markers.

Parameters:

snr_coeffs (Tensor) – SNR coefficients (n_scales, H, W)
scale_idx (int) – Which scale to plot
peak_positions (Optional[Tensor]) – Optional peak positions (N, 2)
title (Optional[str]) – Plot title
cmap (str) – Colormap name
vmin (float) – Minimum SNR for colorscale
vmax (float) – Maximum SNR for colorscale
figsize (Tuple[int, int]) – Figure size
save_path (Optional[str]) – If provided, save figure to this path

wl_stats_torch.visualization.plot_comparison(results_list, labels, statistic='wavelet_peak_counts', scale_idx=0, title=None, log_scale=True, figsize=(10, 6), save_path=None)[source]

Compare the same statistic across multiple result sets.

Parameters:

results_list (List[dict]) – List of result dictionaries from compute_all_statistics
labels (List[str]) – Labels for each result set
statistic (str) – Which statistic to compare (‘wavelet_peak_counts’ or ‘wavelet_l1_norms’)
scale_idx (int) – Which scale to plot (for multi-scale statistics)
title (Optional[str]) – Plot title
log_scale (bool) – Use logarithmic y-axis
figsize (Tuple[int, int]) – Figure size
save_path (Optional[str]) – If provided, save figure to this path