solarviewer 1.0.2__py3-none-any.whl

solar_radio_image_viewer/video_utils.py

@@ -0,0 +1,656 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+Video utilities for advanced video creation features.
+
+This module provides:
+- WCS coordinate detection and axes creation
+- Min/Max timeline plotting
+- Contour video processing
+"""
+
+import numpy as np
+import matplotlib.pyplot as plt
+from matplotlib.patches import Rectangle
+from astropy.io import fits
+from astropy.wcs import WCS
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+# =============================================================================
+# WCS Coordinate Detection and Axes
+# =============================================================================
+
+def detect_coordinate_system(fits_file):
+    """
+    Detect the coordinate system from a FITS file.
+    
+    Parameters
+    ----------
+    fits_file : str
+        Path to the FITS file
+        
+    Returns
+    -------
+    dict
+        Dictionary with:
+        - 'type': 'radec', 'helioprojective', or 'pixel'
+        - 'wcs': WCS object if available, None otherwise
+        - 'ctype1', 'ctype2': Coordinate type strings
+    """
+    try:
+        with fits.open(fits_file) as hdul:
+            header = hdul[0].header
+            
+            try:
+                wcs_obj = WCS(header, naxis=2)
+            except Exception:
+                wcs_obj = None
+            
+            ctype1 = header.get('CTYPE1', '').upper()
+            ctype2 = header.get('CTYPE2', '').upper()
+            
+            # Detect coordinate system
+            if 'HPLN' in ctype1 or 'HPLT' in ctype1:
+                coord_type = 'helioprojective'
+            elif 'RA' in ctype1 or 'DEC' in ctype1:
+                coord_type = 'radec'
+            elif 'GLON' in ctype1 or 'GLAT' in ctype1:
+                coord_type = 'galactic'
+            else:
+                coord_type = 'pixel'
+                
+            return {
+                'type': coord_type,
+                'wcs': wcs_obj,
+                'ctype1': ctype1,
+                'ctype2': ctype2,
+                'cunit1': header.get('CUNIT1', ''),
+                'cunit2': header.get('CUNIT2', ''),
+            }
+    except Exception as e:
+        logger.error(f"Error detecting coordinate system: {e}")
+        return {
+            'type': 'pixel',
+            'wcs': None,
+            'ctype1': '',
+            'ctype2': '',
+        }
+
+
+def get_wcs_axis_labels(coord_info):
+    """
+    Get appropriate axis labels based on coordinate system.
+    
+    Parameters
+    ----------
+    coord_info : dict
+        Dictionary from detect_coordinate_system()
+        
+    Returns
+    -------
+    tuple
+        (xlabel, ylabel)
+    """
+    coord_type = coord_info.get('type', 'pixel')
+    cunit1 = coord_info.get('cunit1', '')
+    cunit2 = coord_info.get('cunit2', '')
+    
+    if coord_type == 'helioprojective':
+        xlabel = f"Solar-X ({cunit1})" if cunit1 else "Solar-X (arcsec)"
+        ylabel = f"Solar-Y ({cunit2})" if cunit2 else "Solar-Y (arcsec)"
+    elif coord_type == 'radec':
+        xlabel = "Right Ascension"
+        ylabel = "Declination"
+    elif coord_type == 'galactic':
+        xlabel = "Galactic Longitude"
+        ylabel = "Galactic Latitude"
+    else:
+        xlabel = "X (pixels)"
+        ylabel = "Y (pixels)"
+        
+    return xlabel, ylabel
+
+
+def create_wcs_axes(fig, wcs_obj, subplot_spec=111):
+    """
+    Create matplotlib axes with WCS projection.
+    
+    Parameters
+    ----------
+    fig : matplotlib.figure.Figure
+        Figure to add axes to
+    wcs_obj : astropy.wcs.WCS
+        WCS object for projection
+    subplot_spec : int or SubplotSpec
+        Subplot specification
+        
+    Returns
+    -------
+    matplotlib.axes.Axes
+        Axes with WCS projection
+    """
+    if wcs_obj is not None:
+        ax = fig.add_subplot(subplot_spec, projection=wcs_obj)
+    else:
+        ax = fig.add_subplot(subplot_spec)
+    return ax
+
+
+# =============================================================================
+# Min/Max Timeline
+# =============================================================================
+
+class MinMaxTimeline:
+    """
+    Manages min/max timeline plotting for video frames.
+    
+    The timeline shows a continuous plot of min and max pixel values
+    across all frames processed so far.
+    """
+    
+    def __init__(self, total_frames, position='bottom-left', log_scale=False):
+        """
+        Initialize the timeline.
+        
+        Parameters
+        ----------
+        total_frames : int
+            Total number of frames in the video
+        position : str
+            Position of the timeline: 'bottom-left', 'bottom-right', 
+            'top-left', 'top-right'
+        log_scale : bool
+            If True, use logarithmic scale for y-axis
+        """
+        self.total_frames = total_frames
+        self.position = position
+        self.log_scale = log_scale
+        self.min_values = []
+        self.max_values = []
+        self.frame_numbers = []
+        
+        # Position configurations [left, bottom, width, height]
+        self.positions = {
+            'bottom-left': [0.05, 0.05, 0.20, 0.10],
+            'bottom-right': [0.75, 0.05, 0.20, 0.10],
+            'top-left': [0.05, 0.85, 0.20, 0.10],
+            'top-right': [0.75, 0.85, 0.20, 0.10],
+        }
+    
+    def add_frame_stats(self, frame_idx, vmin, vmax):
+        """
+        Add statistics for a frame.
+        
+        Parameters
+        ----------
+        frame_idx : int
+            Frame index (0-based)
+        vmin : float
+            Minimum pixel value for this frame
+        vmax : float
+            Maximum pixel value for this frame
+        """
+        self.frame_numbers.append(frame_idx)
+        self.min_values.append(vmin)
+        self.max_values.append(vmax)
+    
+    def precompute_stats(self, files, options, load_func):
+        """
+        Precompute min/max statistics for all frames.
+        
+        Parameters
+        ----------
+        files : list
+            List of file paths
+        options : dict
+            Video creation options
+        load_func : callable
+            Function to load FITS data: load_func(file, stokes) -> (data, header)
+            
+        Returns
+        -------
+        tuple
+            (all_mins, all_maxs) lists
+        """
+        all_mins = []
+        all_maxs = []
+        stokes = options.get('stokes', 'I')
+        
+        for file_path in files:
+            try:
+                data, _ = load_func(file_path, stokes=stokes)
+                if data is not None:
+                    # Apply region if enabled
+                    if options.get('region_enabled', False):
+                        x_min = options.get('x_min', 0)
+                        x_max = options.get('x_max', data.shape[1] - 1)
+                        y_min = options.get('y_min', 0)
+                        y_max = options.get('y_max', data.shape[0] - 1)
+                        data = data[y_min:y_max+1, x_min:x_max+1]
+                    
+                    all_mins.append(np.nanmin(data))
+                    all_maxs.append(np.nanmax(data))
+            except Exception as e:
+                logger.warning(f"Error computing stats for {file_path}: {e}")
+                all_mins.append(np.nan)
+                all_maxs.append(np.nan)
+        
+        self.min_values = all_mins
+        self.max_values = all_maxs
+        self.frame_numbers = list(range(len(files)))
+        
+        return all_mins, all_maxs
+    
+    def draw_timeline(self, fig, current_frame_idx, ax=None):
+        """
+        Draw the timeline on the figure.
+        
+        Parameters
+        ----------
+        fig : matplotlib.figure.Figure
+            Figure to draw on
+        current_frame_idx : int
+            Current frame index (0-based)
+        ax : matplotlib.axes.Axes, optional
+            If provided, use this axes instead of creating an overlay
+            
+        Returns
+        -------
+        matplotlib.axes.Axes
+            The timeline axes
+        """
+        # Use provided axes (dock panel mode) or create overlay (legacy mode)
+        if ax is None:
+            pos = self.positions.get(self.position, self.positions['bottom-left'])
+            ax = fig.add_axes(pos)
+        
+        # Plot data up to current frame
+        frames_to_show = self.frame_numbers[:current_frame_idx + 1]
+        mins_to_show = self.min_values[:current_frame_idx + 1]
+        maxs_to_show = self.max_values[:current_frame_idx + 1]
+        
+        # Plot filled area between min and max (shows range)
+        if len(frames_to_show) > 1:
+            ax.fill_between(frames_to_show, mins_to_show, maxs_to_show, 
+                           alpha=0.3, color='#4FC3F7', label='Range')
+        
+        # Plot min and max lines
+        ax.plot(frames_to_show, mins_to_show, color='#29B6F6', linewidth=1.5, label='Min')
+        ax.plot(frames_to_show, maxs_to_show, color='#FF7043', linewidth=1.5, label='Max')
+        
+        # Plot future data as faded lines
+        if current_frame_idx < len(self.min_values) - 1:
+            future_frames = self.frame_numbers[current_frame_idx:]
+            future_mins = self.min_values[current_frame_idx:]
+            future_maxs = self.max_values[current_frame_idx:]
+            ax.plot(future_frames, future_mins, color='#29B6F6', linewidth=0.8, alpha=0.3)
+            ax.plot(future_frames, future_maxs, color='#FF7043', linewidth=0.8, alpha=0.3)
+        
+        # Current frame marker - vertical line and points
+        if current_frame_idx < len(self.min_values):
+            # Use gray for current frame line to work on both light/dark backgrounds
+            ax.axvline(current_frame_idx, color='gray', linestyle='-', 
+                      linewidth=1.5, alpha=0.7)
+            ax.plot(current_frame_idx, self.min_values[current_frame_idx], 
+                   'o', color='#29B6F6', markersize=6, markeredgecolor='gray', markeredgewidth=1)
+            ax.plot(current_frame_idx, self.max_values[current_frame_idx], 
+                   'o', color='#FF7043', markersize=6, markeredgecolor='gray', markeredgewidth=1)
+        
+        # Set x limits to show full range
+        if self.total_frames > 1:
+            ax.set_xlim(-0.5, self.total_frames - 0.5)
+        
+        # Set y limits based on all data with margin
+        all_vals = self.min_values + self.max_values
+        valid_vals = [v for v in all_vals if not np.isnan(v)]
+        if valid_vals:
+            ymin, ymax = min(valid_vals), max(valid_vals)
+            margin = (ymax - ymin) * 0.15 if ymax != ymin else abs(ymax) * 0.1
+            ax.set_ylim(ymin - margin, ymax + margin)
+        
+        # Style for dock panel (dynamic theme)
+        # ax.set_facecolor('#1a1a2e') # Let theme decide facecolor
+        ax.tick_params(labelsize=10, direction='in', length=3)
+        
+        # Show only left and bottom spines (let theme decide color)
+        # ax.spines['bottom'].set_color('#555555')
+        # ax.spines['left'].set_color('#555555')
+        ax.spines['top'].set_visible(False)
+        ax.spines['right'].set_visible(False)
+        
+        # X-axis: show frame numbers
+        ax.set_xlabel('Frame', fontsize=11, labelpad=2)
+        
+        # Y-axis: show value range with scientific notation if needed
+        ax.set_ylabel('Value', fontsize=11, labelpad=2)
+        
+        # Apply log scale if enabled
+        if self.log_scale:
+            ax.set_yscale('log')
+        else:
+            ax.yaxis.set_major_locator(plt.MaxNLocator(3))  # Limit to 3 ticks
+        
+        # Format y-axis with scientific notation for large values
+        from matplotlib.ticker import ScalarFormatter
+        formatter = ScalarFormatter(useMathText=True)
+        formatter.set_scientific(True)
+        formatter.set_powerlimits((-2, 3))
+        ax.yaxis.set_major_formatter(formatter)
+        # ax.yaxis.get_offset_text().set_color('white')
+        ax.yaxis.get_offset_text().set_fontsize(8)
+        
+        # Add current values annotation
+        if current_frame_idx < len(self.min_values):
+            curr_min = self.min_values[current_frame_idx]
+            curr_max = self.max_values[current_frame_idx]
+            info_text = f"Frame {current_frame_idx + 1}/{self.total_frames}"
+            ax.text(0.98, 0.92, info_text, transform=ax.transAxes, fontsize=10,
+                   ha='right', va='top', alpha=0.9)
+        
+        return ax
+
+
+# =============================================================================
+# Contour Video Processor
+# =============================================================================
+
+class ContourVideoProcessor:
+    """
+    Handles contour processing for video creation.
+    
+    Supports three modes:
+    - Mode A: Fixed base image, evolving contours
+    - Mode B: Fixed contours, evolving colormap image
+    - Mode C: Both evolve
+    """
+    
+    def __init__(self, mode='A', contour_settings=None):
+        """
+        Initialize the contour processor.
+        
+        Parameters
+        ----------
+        mode : str
+            'A', 'B', or 'C'
+        contour_settings : dict
+            Contour configuration dictionary
+        """
+        self.mode = mode.upper()
+        self.contour_settings = contour_settings or self._default_settings()
+        
+        # Cache for base image / contour data
+        self.base_image_data = None
+        self.base_image_wcs = None
+        self.fixed_contour_data = None
+        self.fixed_contour_wcs = None
+        self.fixed_contour_levels = None
+    
+    def _default_settings(self):
+        """Return default contour settings."""
+        return {
+            'level_type': 'fraction',  # 'fraction', 'sigma', 'absolute'
+            'pos_levels': [0.1, 0.3, 0.5, 0.7, 0.9],
+            'neg_levels': [0.1, 0.3, 0.5, 0.7, 0.9],
+            'pos_color': 'white',
+            'neg_color': 'cyan',
+            'linewidth': 1.0,
+            'pos_linestyle': '-',
+            'neg_linestyle': '--',
+        }
+    
+    def load_base_image(self, file_path, stokes='I', load_func=None):
+        """Load the fixed base image (for mode A)."""
+        if load_func is None:
+            from .create_video import load_fits_data
+            load_func = load_fits_data
+            
+        self.base_image_data, header = load_func(file_path, stokes=stokes)
+        
+        # Create WCS from the same header used for data loading
+        if header:
+            try:
+                self.base_image_wcs = WCS(header, naxis=2)
+                logger.info(f"Loaded base_image_wcs: CRPIX={self.base_image_wcs.wcs.crpix}")
+            except Exception as e:
+                logger.warning(f"Failed to create base_image_wcs: {e}")
+                self.base_image_wcs = None
+        else:
+            self.base_image_wcs = None
+    
+    def load_fixed_contour(self, file_path, stokes='I', load_func=None):
+        """Load the fixed contour data (for mode B)."""
+        if load_func is None:
+            from .create_video import load_fits_data
+            load_func = load_fits_data
+            
+        self.fixed_contour_data, header = load_func(file_path, stokes=stokes)
+        
+        # Create WCS from the same header used for data loading
+        if header:
+            try:
+                self.fixed_contour_wcs = WCS(header, naxis=2)
+            except Exception:
+                self.fixed_contour_wcs = None
+        else:
+            self.fixed_contour_wcs = None
+        
+        # Pre-compute levels
+        self._compute_contour_levels(self.fixed_contour_data)
+    
+    def _compute_contour_levels(self, data):
+        """Compute contour levels based on settings."""
+        settings = self.contour_settings
+        level_type = settings.get('level_type', 'fraction')
+        pos_levels = settings.get('pos_levels', [])
+        neg_levels = settings.get('neg_levels', [])
+        
+        abs_max = np.nanmax(np.abs(data))
+        
+        if level_type == 'fraction':
+            pos = sorted([level * abs_max for level in pos_levels])
+            neg = sorted([-level * abs_max for level in neg_levels])
+        elif level_type == 'sigma':
+            # Use bottom 10% of image (full width, bottom 10% height) for RMS calculation
+            # This avoids including the sun in the noise estimate
+            height = data.shape[0]
+            bottom_10_pct = max(1, int(height * 0.1))  # At least 1 row
+            noise_region = data[:bottom_10_pct, :]  # Bottom rows (low y indices)
+            rms = np.nanstd(noise_region)
+            pos = sorted([level * rms for level in pos_levels])
+            neg = sorted([-level * rms for level in neg_levels])
+        else:  # absolute
+            pos = sorted(pos_levels)
+            neg = sorted([-level for level in neg_levels])
+        
+        self.fixed_contour_levels = {'pos': pos, 'neg': neg}
+        return self.fixed_contour_levels
+    
+    def compute_contour_levels(self, data):
+        """
+        Compute contour levels for given data.
+        
+        Parameters
+        ----------
+        data : ndarray
+            Image data
+            
+        Returns
+        -------
+        dict
+            {'pos': [...], 'neg': [...]}
+        """
+        return self._compute_contour_levels(data)
+    
+
+    def draw_contours(self, ax, contour_data, levels=None, target_wcs=None, 
+                     contour_wcs=None, target_shape=None, region_info=None):
+        """
+        Draw contours on axes with support for WCS reprojection.
+        
+        Parameters
+        ----------
+        region_info : dict, optional
+            If provided, contains 'x_min', 'y_min', 'x_max', 'y_max', 'full_shape'
+            Used to reproject to full shape first, then crop.
+        """
+        settings = self.contour_settings
+        
+        if levels is None:
+            levels = self.compute_contour_levels(contour_data)
+        
+        collections = []
+        # Reprojection logic
+        if target_wcs and contour_wcs and target_shape is not None:
+            try:
+                from reproject import reproject_interp
+                from astropy.wcs import WCS
+                
+                # Create axis-swapped WCS objects matching viewer.py's approach
+                def create_swapped_wcs(orig_wcs):
+                    """Create WCS with swapped axes for reprojection."""
+                    swapped = WCS(naxis=2)
+                    swapped.wcs.crpix = [orig_wcs.wcs.crpix[1], orig_wcs.wcs.crpix[0]]
+                    swapped.wcs.crval = [orig_wcs.wcs.crval[1], orig_wcs.wcs.crval[0]]
+                    try:
+                        swapped.wcs.cdelt = [orig_wcs.wcs.cdelt[1], orig_wcs.wcs.cdelt[0]]
+                    except Exception:
+                        pass
+                    if orig_wcs.wcs.ctype[0] and orig_wcs.wcs.ctype[1]:
+                        swapped.wcs.ctype = [orig_wcs.wcs.ctype[1], orig_wcs.wcs.ctype[0]]
+                    return swapped
+                
+                # Swap both WCS for consistent reprojection
+                target_wcs_swapped = create_swapped_wcs(target_wcs)
+                contour_wcs_swapped = create_swapped_wcs(contour_wcs)
+                
+                # If region_info provided, reproject to FULL shape first, then crop
+                if region_info:
+                    full_shape = region_info.get('full_shape', target_shape)
+                    
+                    # Reproject to full base image shape
+                    reprojected_data, footprint = reproject_interp(
+                        (contour_data, contour_wcs_swapped),
+                        target_wcs_swapped,
+                        shape_out=full_shape
+                    )
+                    
+                    # Then crop to the region
+                    x_min = region_info.get('x_min', 0)
+                    y_min = region_info.get('y_min', 0)
+                    x_max = region_info.get('x_max', full_shape[1] - 1)
+                    y_max = region_info.get('y_max', full_shape[0] - 1)
+                    reprojected_data = reprojected_data[y_min:y_max+1, x_min:x_max+1]
+                else:
+                    # No region - reproject directly to target shape
+                    reprojected_data, footprint = reproject_interp(
+                        (contour_data, contour_wcs_swapped),
+                        target_wcs_swapped,
+                        shape_out=target_shape
+                    )
+                
+                # Log reprojection result
+                nan_pct = 100 * np.isnan(reprojected_data).sum() / reprojected_data.size
+                logger.info(f"[REPROJ] Result: {nan_pct:.1f}% NaNs")
+                
+                # Check validity of reprojected data
+                is_all_nan = np.all(np.isnan(reprojected_data))
+                is_all_zero = np.all(np.nan_to_num(reprojected_data) == 0)
+                
+                if is_all_nan or is_all_zero:
+                    logger.warning("Reprojection yielded empty data. Falling back to pixel overlay.")
+                else:
+                    contour_data = reprojected_data
+
+            except Exception as e:
+                logger.warning(f"Contour reprojection failed: {e}. Falling back to pixel overlay.")
+        
+        # Draw positive contours
+        if levels.get('pos'):
+            try:
+                cs_pos = ax.contour(
+                    contour_data,
+                    levels=levels['pos'],
+                    colors=settings.get('pos_color', 'white'),
+                    linewidths=settings.get('linewidth', 1.0),
+                    linestyles=settings.get('pos_linestyle', '-'),
+                )
+                collections.append(cs_pos)
+            except Exception as e:
+                logger.warning(f"Error drawing positive contours: {e}")
+        
+        # Draw negative contours
+        if levels.get('neg'):
+            try:
+                cs_neg = ax.contour(
+                    contour_data,
+                    levels=levels['neg'],
+                    colors=settings.get('neg_color', 'cyan'),
+                    linewidths=settings.get('linewidth', 1.0),
+                    linestyles=settings.get('neg_linestyle', '--'),
+                )
+                collections.append(cs_neg)
+            except Exception as e:
+                logger.warning(f"Error drawing negative contours: {e}")
+        
+        return collections
+
+    
+    def get_frame_data(self, frame_idx, colormap_files=None, contour_files=None,
+                       stokes='I', load_func=None):
+        """
+        Get data for a specific frame based on mode.
+        
+        Parameters
+        ----------
+        frame_idx : int
+            Frame index
+        colormap_files : list
+            List of colormap image files
+        contour_files : list
+            List of contour image files
+        stokes : str
+            Stokes parameter
+        load_func : callable
+            Function to load FITS data
+            
+        Returns
+        -------
+        dict
+            {'colormap_data': ..., 'contour_data': ..., 'contour_levels': ...}
+        """
+        if load_func is None:
+            from .create_video import load_fits_data
+            load_func = load_fits_data
+        
+        result = {}
+        
+        if self.mode == 'A':
+            # Fixed base, evolving contours
+            result['colormap_data'] = self.base_image_data
+            if contour_files and frame_idx < len(contour_files):
+                data, _ = load_func(contour_files[frame_idx], stokes=stokes)
+                result['contour_data'] = data
+                result['contour_levels'] = self.compute_contour_levels(data)
+                
+        elif self.mode == 'B':
+            # Fixed contours, evolving colormap
+            result['contour_data'] = self.fixed_contour_data
+            result['contour_levels'] = self.fixed_contour_levels
+            if colormap_files and frame_idx < len(colormap_files):
+                data, _ = load_func(colormap_files[frame_idx], stokes=stokes)
+                result['colormap_data'] = data
+                
+        elif self.mode == 'C':
+            # Both evolve
+            if colormap_files and frame_idx < len(colormap_files):
+                data, _ = load_func(colormap_files[frame_idx], stokes=stokes)
+                result['colormap_data'] = data
+            if contour_files and frame_idx < len(contour_files):
+                data, _ = load_func(contour_files[frame_idx], stokes=stokes)
+                result['contour_data'] = data
+                result['contour_levels'] = self.compute_contour_levels(data)
+        
+        return result