pivtools 0.1.3__cp311-cp311-win_amd64.whl

pivtools_cli/piv/save_results.py

@@ -0,0 +1,342 @@
+"""
+Module for saving PIV results to .mat files compatible with post-processing code.
+"""
+import logging
+import sys
+from pathlib import Path
+from typing import List, Optional, Union
+
+import numpy as np
+import scipy.io
+from pivtools_core.config import Config
+from pivtools_core.paths import get_data_paths
+
+from pivtools_cli.piv.piv_result import PIVResult, PIVPassResult
+
+
+def save_piv_result_distributed(
+    piv_result: PIVResult,
+    output_path: Path,
+    frame_number: int,
+    runs_to_save: Optional[List[int]] = None,
+    vector_fmt: str = "B%05d.mat",
+) -> str:
+    """
+    Save a PIV result to disk. Designed to be submitted to Dask workers.
+    
+    This function can be called on Dask workers to save results in parallel,
+    avoiding the memory bottleneck of gathering all results to main.
+    Memory-efficient: uses direct serialization without unnecessary copies.
+    
+    Parameters
+    ----------
+    piv_result : PIVResult
+        The PIV result object containing one or more passes with complete data.
+    output_path : Path
+        Directory where the .mat file will be saved.
+    frame_number : int
+        Frame number (1-based) for the filename (e.g., 1 -> B00001.mat).
+    runs_to_save : Optional[List[int]]
+        List of pass indices (0-based) to save. If None, save all passes.
+        For passes not in this list, empty arrays will be saved.
+    vector_fmt : str
+        Format string for the filename, e.g., "B%05d.mat".
+        
+    Returns
+    -------
+    str
+        Path to the saved file (for verification/logging).
+    """
+    output_path = Path(output_path)
+    output_path.mkdir(parents=True, exist_ok=True)
+    
+    filename = output_path / (vector_fmt % frame_number)
+    
+    if len(piv_result.passes) == 0:
+        logging.warning(
+            f"PIVResult has no passes for frame {frame_number}. "
+            "Skipping save."
+        )
+        return str(filename)
+    
+    # Create single struct with arrays indexed by pass number
+    # All data is already in piv_result, no external lists needed
+    mat_data = _create_piv_struct_all_passes(piv_result, runs_to_save)
+    
+    # Save to .mat file with compression to reduce I/O
+    scipy.io.savemat(filename, {"piv_result": mat_data}, oned_as="row", do_compression=True)
+    logging.debug(f"Worker saved PIV result to {filename}")
+    
+    return str(filename)
+
+
+def save_coordinates_from_config_distributed(
+    config: Config,
+    output_path: Path,
+    correlator_cache: Optional[dict] = None,
+    runs_to_save: Optional[List[int]] = None,
+) -> str:
+    """
+    Generate and save coordinate grids. Designed for Dask workers.
+    
+    Parameters
+    ----------
+    config : Config
+        Configuration object containing window sizes and overlap.
+    output_path : Path
+        Directory where coordinates.mat will be saved.
+    correlator_cache : Optional[dict]
+        Precomputed correlator cache to avoid redundant computation.
+    runs_to_save : Optional[List[int]]
+        List of pass indices (0-based) to save with data. If None, save all passes.
+        For passes not in this list, empty coordinate grids will be saved.
+        
+    Returns
+    -------
+    str
+        Path to the saved coordinates file.
+    """
+    from pivtools_cli.piv.piv_backend.cpu_instantaneous import (
+        InstantaneousCorrelatorCPU
+    )
+    
+    # Create a temporary correlator with optional precomputed cache
+    correlator = InstantaneousCorrelatorCPU(config, precomputed_cache=correlator_cache)
+    
+    # Extract the cached window centers
+    win_ctrs_x_list = correlator.win_ctrs_x
+    win_ctrs_y_list = correlator.win_ctrs_y
+    
+    num_passes = len(config.window_sizes)
+    
+    if runs_to_save is None:
+        runs_to_save = list(range(num_passes))
+    
+
+    # Create MATLAB-style struct array with fields 'x' and 'y', shape (num_passes,)
+    dtype = [('x', object), ('y', object)]
+    coords_struct = np.empty((num_passes,), dtype=dtype)
+
+    for i in range(num_passes):
+        if i in runs_to_save:
+            x_centers = win_ctrs_x_list[i]
+            y_centers = win_ctrs_y_list[i]
+
+            # Create 2D coordinate grids with smallest y at the bottom
+            x_grid, y_grid = np.meshgrid(x_centers+1, y_centers[::-1]+1, indexing='xy')
+
+            # Convert to half precision for space saving
+            x_grid = _convert_to_half_precision(x_grid)
+            y_grid = _convert_to_half_precision(y_grid)
+
+            coords_struct['x'][i] = x_grid
+            coords_struct['y'][i] = y_grid
+        else:
+            # Empty arrays for non-selected passes
+            coords_struct['x'][i] = np.array([], dtype=np.float16)
+            coords_struct['y'][i] = np.array([], dtype=np.float16)
+
+    output_path = Path(output_path)
+    output_path.mkdir(parents=True, exist_ok=True)
+
+    filename = output_path / "coordinates.mat"
+    scipy.io.savemat(filename, {"coordinates": coords_struct}, oned_as="row", do_compression=True)
+    logging.info(f"Worker saved coordinates to {filename}")
+
+    return str(filename)
+
+
+def _create_piv_struct_all_passes(
+    piv_result: PIVResult,
+    runs_to_save: Optional[List[int]] = None,
+) -> np.ndarray:
+    """
+    Create a MATLAB-compatible struct with arrays indexed by pass number.
+    
+    This creates a single struct where each field (ux, uy, b_mask, etc.) is
+    an array with one element per pass, matching the expected format:
+        piv_result["ux"][pass_idx] = 2D array for that pass
+    
+    All required data (including window centers and masks) is extracted from
+    the PIVResult object, which contains all necessary information in each
+    PIVPassResult.
+    
+    Parameters
+    ----------
+    piv_result : PIVResult
+        PIV result object containing one or more passes with complete data.
+    runs_to_save : Optional[List[int]]
+        List of pass indices (0-based) to save with data. If None, save all passes.
+        For passes not in this list, empty arrays will be saved.
+        
+    Returns
+    -------
+    np.ndarray
+        Structured numpy array compatible with scipy.io.savemat.
+    """
+    n_passes = len(piv_result.passes)
+    
+    # Always save all passes, but empty arrays for non-selected passes
+    n_passes_to_save = n_passes
+    passes_to_save = list(range(n_passes))
+    
+    # If runs_to_save is specified, only fill data for those passes
+    if runs_to_save is None:
+        runs_to_save = passes_to_save
+    
+    # Create structured dtype with all fields
+    dtype = [
+        ('ux', object),
+        ('uy', object),
+        ('b_mask', object),
+        ('nan_mask', object),
+        ('win_ctrs_x', object),
+        ('win_ctrs_y', object),
+        ('peak_mag', object),
+        ('peak_choice', object),
+        ('n_windows', object),
+        ('predictor_field', object),
+        ('window_size', object),
+    ]
+    
+    # Create the struct with shape (n_passes_to_save,)
+    piv_struct = np.empty((n_passes_to_save,), dtype=dtype)
+    
+    # Get dtype from first pass for creating empty arrays
+    first_pass = piv_result.passes[0]
+    if first_pass.ux_mat is not None and first_pass.ux_mat.size > 0:
+        data_dtype = first_pass.ux_mat.dtype
+    else:
+        data_dtype = np.float64
+    
+    # Initialize all passes with empty arrays
+    empty = np.empty((0, 0), dtype=data_dtype)
+    for i in range(n_passes_to_save):
+        piv_struct['ux'][i] = empty
+        piv_struct['uy'][i] = empty
+        piv_struct['b_mask'][i] = empty
+        piv_struct['nan_mask'][i] = empty
+        piv_struct['win_ctrs_x'][i] = empty
+        piv_struct['win_ctrs_y'][i] = empty
+        piv_struct['peak_mag'][i] = empty
+        piv_struct['peak_choice'][i] = empty
+        piv_struct['n_windows'][i] = empty
+        piv_struct['predictor_field'][i] = empty
+        piv_struct['window_size'][i] = empty
+    
+    # Fill with actual data for selected passes
+    for local_idx, global_pass_idx in enumerate(passes_to_save):
+        if global_pass_idx not in runs_to_save:
+            continue  # Skip filling for non-selected passes
+        pass_result = piv_result.passes[global_pass_idx]
+        
+        # Save ux and uy directly without swapping - coordinate system is now correct
+        if pass_result.ux_mat is not None:
+            piv_struct['ux'][local_idx] = _convert_to_half_precision(pass_result.ux_mat)
+        if pass_result.uy_mat is not None:
+            piv_struct['uy'][local_idx] = _convert_to_half_precision(pass_result.uy_mat)
+
+        # Use b_mask from pass_result (already computed during PIV)
+        if pass_result.b_mask is not None:
+            piv_struct['b_mask'][local_idx] = pass_result.b_mask
+        elif pass_result.nan_mask is not None:
+            # Fallback to nan_mask if b_mask not available
+            piv_struct['b_mask'][local_idx] = pass_result.nan_mask
+
+        if pass_result.nan_mask is not None:
+            piv_struct['nan_mask'][local_idx] = pass_result.nan_mask
+
+        # Window centers are always stored in pass_result
+        if pass_result.win_ctrs_x is not None:
+            piv_struct['win_ctrs_x'][local_idx] = _convert_to_half_precision(pass_result.win_ctrs_x)
+        if pass_result.win_ctrs_y is not None:
+            piv_struct['win_ctrs_y'][local_idx] = _convert_to_half_precision(pass_result.win_ctrs_y)
+            
+        if pass_result.peak_mag is not None:
+            piv_struct['peak_mag'][local_idx] = _convert_to_half_precision(pass_result.peak_mag)
+        if pass_result.peak_choice is not None:
+            piv_struct['peak_choice'][local_idx] = pass_result.peak_choice
+        if pass_result.n_windows is not None:
+            piv_struct['n_windows'][local_idx] = pass_result.n_windows
+        if pass_result.predictor_field is not None:
+            piv_struct['predictor_field'][local_idx] = _convert_to_half_precision(pass_result.predictor_field)
+        if pass_result.window_size is not None:
+            piv_struct['window_size'][local_idx] = pass_result.window_size
+    
+    return piv_struct
+
+
+# Note: get_data_paths is imported from src/paths.py at the top of this file
+
+
+def get_output_path(
+    config: Config,
+    camera: Union[int, str],
+    create: bool = True,
+    use_uncalibrated: bool = True,
+) -> Path:
+    """
+    Get the output path for a specific camera's PIV results using the GUI path structure.
+    
+    Follows the standardized directory structure:
+    - Uncalibrated: base_path/uncalibrated_piv/{num_images}/Cam{camera}/instantaneous
+    - Calibrated: base_path/calibrated_piv/{num_images}/Cam{camera}/instantaneous
+    
+    Parameters
+    ----------
+    config : Config
+        Configuration object.
+    camera : Union[int, str]
+        Camera number (int) or camera folder name (str, e.g., "Cam1").
+    create : bool
+        If True, create the directory if it doesn't exist.
+    use_uncalibrated : bool
+        If True, save to uncalibrated_piv directory.
+        If False, save to calibrated_piv directory.
+        
+    Returns
+    -------
+    Path
+        Output path for PIV results.
+    """
+    base_path = config.base_paths[0]
+    
+    # Convert camera to int if it's a string
+    if isinstance(camera, str):
+        if camera.startswith("Cam"):
+            camera_num = int(camera[3:])
+        else:
+            camera_num = int(camera)
+    else:
+        camera_num = camera
+    
+    # Get PIV type - default to instantaneous
+    piv_type = "instantaneous" if config.data.get("processing", {}).get("instantaneous", True) else "ensemble"
+    
+    # Use get_data_paths from src/paths.py (positional args: base_dir, num_images, cam, type_name)
+    paths = get_data_paths(
+        base_path,
+        config.num_images,
+        camera_num,
+        piv_type,
+        endpoint="",
+        use_uncalibrated=use_uncalibrated
+    )
+    
+    output_path = paths["data_dir"]
+    
+    if create:
+        output_path.mkdir(parents=True, exist_ok=True)
+    
+    return output_path
+
+
+def _convert_to_half_precision(arr: np.ndarray) -> np.ndarray:
+    """
+    Convert float arrays to half precision (float16) for space saving.
+    """
+    if arr is None or arr.size == 0:
+        return arr
+    if arr.dtype.kind == 'f':
+        return arr.astype(np.float16)
+    return arr

pivtools_cli/piv_cluster/cluster.py

@@ -0,0 +1,108 @@
+import logging
+import os
+import sys
+from collections import defaultdict
+from pathlib import Path
+from typing import List, Tuple
+
+from dask.distributed import Client, LocalCluster
+
+from pivtools_core.config import Config
+
+
+def make_cluster(
+    threads_per_worker: int = 1,  # None,
+    n_workers_per_node: int = 2,
+    memory_limit: str = "auto",
+) -> Tuple[LocalCluster, Client]:
+    cluster = LocalCluster(
+        n_workers=n_workers_per_node,
+        threads_per_worker=threads_per_worker,
+        memory_limit=memory_limit,
+        nanny=True,
+    )
+    client = Client(cluster)
+    return cluster, client
+
+
+def group_workers_by_host(client: Client) -> dict[str, List[str]]:
+    workers = client.scheduler_info()["workers"]
+    grouped = defaultdict(list)
+    for addr, info in workers.items():
+        grouped[info["host"]].append(addr)
+    return dict(grouped)
+
+
+def select_workers_per_node(client: Client, n_workers_per_node: int = 1) -> List[str]:
+    grouped = group_workers_by_host(client)
+    selected = []
+    for node_workers in grouped.values():
+        selected.extend(node_workers[:n_workers_per_node])
+    return selected
+
+
+def start_cluster(
+    n_workers_per_node: int = 1,
+    threads_per_worker: int = None,
+    memory_limit: str = "auto",
+    config: Config = Config(),
+) -> tuple[LocalCluster, Client]:
+    """
+    Start a local Dask cluster.
+
+    Returns:
+        client: Dask Client
+        piv_workers: list of workers to use for PIV
+    """
+    cluster = None
+    client = None
+
+    try:
+        cluster, client = make_cluster(
+            n_workers_per_node=n_workers_per_node,
+            threads_per_worker=threads_per_worker,
+            memory_limit=memory_limit,
+        )
+        client.run(
+            setup_worker_logging,
+            log_level=getattr(logging, config.log_level, logging.INFO),
+            log_file=config.log_file if hasattr(config, "log_file") else None,
+            log_console=True,
+        )
+
+        return cluster, client
+
+    except Exception as e:
+        print(f"Error starting Dask cluster: {e}")
+        if client is not None:
+            client.close()
+        if cluster is not None:
+            cluster.close()
+        raise
+
+
+def setup_worker_logging(log_level=logging.INFO, log_file=None, log_console=True):
+    """
+    Configure logging inside a Dask worker process.
+    """
+    logger = logging.getLogger()
+    logger.setLevel(log_level)
+
+    for handler in logger.handlers[:]:
+        logger.removeHandler(handler)
+
+    if log_file:
+        file_handler = logging.FileHandler(log_file)
+        file_handler.setLevel(log_level)
+        file_formatter = logging.Formatter("%(asctime)s [%(levelname)s] %(message)s")
+        file_handler.setFormatter(file_formatter)
+        logger.addHandler(file_handler)
+
+    if log_console:
+        console_handler = logging.StreamHandler()
+        console_handler.setLevel(log_level)
+        console_formatter = logging.Formatter("%(asctime)s [%(levelname)s] %(message)s")
+        console_handler.setFormatter(console_formatter)
+        logger.addHandler(console_handler)
+
+    logger.info("Worker logging configured successfully")