cdo-toolkit 0.1.0__py3-none-any.whl

cdo_toolkit/__init__.py

@@ -0,0 +1,47 @@
+"""CDO-based NetCDF regridding toolkit."""
+
+from cdo_toolkit.api import (
+    extract_seafloor_single_file,
+    regrid_directory,
+    regrid_directory_both_levels,
+    regrid_large_files,
+    regrid_single_file,
+    regrid_single_file_both_levels,
+    regrid_single_file_extreme_levels,
+)
+from cdo_toolkit.cmip import (
+    filter_files_by_variables,
+    get_cmip_variable_name,
+    parse_variable_list,
+    pick_representative_file,
+    representative_files_by_directory,
+)
+from cdo_toolkit.errors import default_log_dir, init_regrid_error_log, log_regrid_error
+from cdo_toolkit.paths import is_intermediate_nc, is_weights_or_cache_file, weight_cache_dir_for_input
+from cdo_toolkit.pipeline import CDORegridPipeline
+from cdo_toolkit.workers import process_single_file_standalone
+
+__all__ = [
+    "CDORegridPipeline",
+    "default_log_dir",
+    "extract_seafloor_single_file",
+    "filter_files_by_variables",
+    "get_cmip_variable_name",
+    "init_regrid_error_log",
+    "is_intermediate_nc",
+    "is_weights_or_cache_file",
+    "log_regrid_error",
+    "parse_variable_list",
+    "pick_representative_file",
+    "process_single_file_standalone",
+    "representative_files_by_directory",
+    "regrid_directory",
+    "regrid_directory_both_levels",
+    "regrid_large_files",
+    "regrid_single_file",
+    "regrid_single_file_both_levels",
+    "regrid_single_file_extreme_levels",
+    "weight_cache_dir_for_input",
+]
+
+__version__ = "0.1.0"

cdo_toolkit/__main__.py

@@ -0,0 +1,6 @@
+"""Allow ``python -m cdo_toolkit``."""
+
+from cdo_toolkit.cli import main
+
+if __name__ == "__main__":
+    main()

cdo_toolkit/api.py

@@ -0,0 +1,573 @@
+"""High-level regridding API."""
+
+import multiprocessing as mp
+import os
+import tempfile
+from concurrent.futures import ProcessPoolExecutor
+from pathlib import Path
+from typing import Optional
+
+import xarray as xa
+
+from cdo_toolkit.cmip import (
+    filter_files_by_variables,
+    get_cmip_variable_name,
+    parse_variable_list,
+    pick_representative_file,
+    representative_files_by_directory,
+)
+from cdo_toolkit.errors import init_regrid_error_log
+from cdo_toolkit.paths import is_intermediate_nc
+from cdo_toolkit.pipeline import CDORegridPipeline
+from cdo_toolkit.timing import format_processing_time, get_processing_time, print_timestamp
+from cdo_toolkit.ui import RegridProgressUI
+
+def extract_seafloor_single_file(
+    input_path: Path,
+    output_path: Optional[Path] = None,
+    verbose: bool = True,
+    overwrite: bool = False,
+) -> Path:
+    """
+    Convenience function to extract seafloor values from a single file.
+    
+    Args:
+    - input_path (Path): Path to input file
+    - output_path (Path): Path to output file (if None, auto-generates <filename>_seafloor.nc)
+    - verbose (bool): Enable verbose output
+    - overwrite (bool): If True, overwrite existing output files
+    
+    Returns (Path): Path to the seafloor-extracted file
+    
+    Raises:
+    - RuntimeError: If seafloor extraction fails
+    """
+    pipeline = CDORegridPipeline(
+        extract_seafloor=True,
+        verbose=verbose,
+    )
+    
+    if output_path is None:
+        output_path = input_path.parent / f"{input_path.stem}_seafloor{input_path.suffix}"
+    
+    if output_path.exists() and not overwrite:
+        if verbose:
+            pipeline.console.print(f"[yellow]Seafloor file already exists: {output_path.name}[/yellow]")
+        return output_path
+    
+    try:
+        seafloor_path = pipeline._extract_seafloor_values(input_path)
+        if verbose:
+            pipeline.console.print(f"[green]Seafloor file created: {seafloor_path}[/green]")
+        return seafloor_path
+    except Exception as e:
+        if verbose:
+            pipeline.console.print(f"[red]Failed to extract seafloor: {e}[/red]")
+        raise
+
+
+def regrid_single_file(
+    input_path: Path,
+    output_path: Optional[Path] = None,
+    output_dir: Optional[Path] = None,
+    target_resolution: tuple[float, float] = (1.0, 1.0),
+    extract_surface: bool = False,
+    extract_seafloor: bool = False,
+    use_regrid_cache: bool = True,
+    use_seafloor_cache: bool = True,
+    verbose: bool = True,
+    verbose_diagnostics: bool = False,
+    cleanup_weights: bool = False,
+    overwrite: bool = False,
+    use_ui: bool = True,
+) -> bool:
+    """
+    Convenience function to regrid a single file.
+
+    Pipeline: if neither extract_surface nor extract_seafloor, regrid whole file;
+    if extract_surface, extract top level and regrid; if extract_seafloor, extract
+    seafloor and regrid. For both surface and seafloor use regrid_single_file_extreme_levels.
+
+    Args:
+    - input_path (Path): Path to input file
+    - output_path (Path): Path to output file
+    - target_resolution (tuple): Target resolution as (lon_res, lat_res)
+    - extract_surface (bool): Extract top level only and regrid that
+    - extract_seafloor (bool): Extract seafloor values and regrid only that
+    - use_regrid_cache (bool): Reuse existing regrid weight files when present
+    - use_seafloor_cache (bool): Reuse seafloor depth indices cache
+    - verbose (bool): Enable verbose output (progress UI)
+    - verbose_diagnostics (bool): If True, print Grid type, File size, Large file (max verbosity)
+    - cleanup_weights (bool): Clean up weights after processing
+    - overwrite (bool): If True, overwrite existing output files
+    - use_ui (bool): Use rich progress UI
+
+    Returns (bool): True if successful, False otherwise
+    """
+    pipeline = CDORegridPipeline(
+        target_resolution=target_resolution,
+        extract_surface=extract_surface,
+        extract_seafloor=extract_seafloor,
+        use_regrid_cache=use_regrid_cache,
+        use_seafloor_cache=use_seafloor_cache,
+        verbose=verbose,
+        verbose_diagnostics=verbose_diagnostics,
+        cleanup_weights=cleanup_weights,
+    )
+    if output_dir is not None and output_path is None:
+        has_level = pipeline._has_level_lightweight(input_path)
+        filename = pipeline._generate_output_filename(
+            input_path, has_level, extract_surface, extract_seafloor
+        )
+        output_path = output_dir / filename
+    error_log_path = init_regrid_error_log()
+    pipeline.set_error_log_path(error_log_path)
+    # Initialize UI if requested
+    ui = None
+    if use_ui and verbose:
+        ui = RegridProgressUI(
+            [input_path],
+            verbose=verbose,
+            verbose_diagnostics=pipeline.verbose_diagnostics,
+            log_file=error_log_path,
+        )
+        ui.__enter__()
+    # Track processing time
+    import time
+    start_time = print_timestamp(pipeline.console, "START") if verbose else time.localtime()
+    try:
+        result = pipeline.regrid_file(input_path, output_path, overwrite=overwrite, ui=ui)
+        return result
+    finally:
+        if ui:
+            # Add timing to stats
+            end_time = print_timestamp(pipeline.console, "END") if verbose else time.localtime()
+            pipeline.stats['processing_time'] = format_processing_time(
+                get_processing_time(start_time, end_time)
+            )
+            ui._update_stats(pipeline.stats)
+            ui.print_summary()
+            ui.__exit__(None, None, None)
+
+
+def regrid_single_file_both_levels(
+    input_path: Path,
+    output_dir: Optional[Path] = None,
+    target_resolution: tuple[float, float] = (1.0, 1.0),
+    use_regrid_cache: bool = True,
+    use_seafloor_cache: bool = True,
+    verbose: bool = True,
+    cleanup_weights: bool = False,
+    overwrite: bool = False,
+    weight_cache_dir: Optional[Path] = None,
+) -> dict[str, bool]:
+    """
+    Regrid both the top level (surface) and the seafloor values for a single file.
+
+    Performs steps 2 and 3 sequentially: seafloor extraction+regrid, then surface extraction+regrid.
+    Outputs: ``<name>_seafloor_regridded.nc`` and ``<name>_top_level_regridded.nc`` (if multi-level).
+
+    Parameters
+    ----------
+    input_path : Path
+        Input NetCDF file.
+    output_dir : Path, optional
+        Directory for regridded outputs. If ``None``, files are written next to ``input_path``.
+    target_resolution : tuple[float, float]
+        Target grid resolution.
+    use_regrid_cache : bool
+        Reuse existing regrid weight files when present.
+    use_seafloor_cache : bool
+        Reuse seafloor depth indices cache.
+    verbose : bool
+        Enable verbose logging.
+    cleanup_weights : bool
+        Clean up weights after processing.
+    overwrite : bool
+        Overwrite existing outputs.
+    weight_cache_dir : Path, optional
+        Directory for regrid weight cache. If None, each input file uses ``<input_dir>/cdo_weights``.
+
+    Returns
+    -------
+    dict[str, bool]
+        Mapping ``{'top_level': bool, 'seafloor': bool}`` indicating success for each stream.
+    """
+    # 1) Seafloor: extract seafloor indices (from file or cache), extract values, regrid
+    seafloor_pipeline = CDORegridPipeline(
+        target_resolution=target_resolution,
+        extract_surface=False,
+        extract_seafloor=True,
+        use_regrid_cache=use_regrid_cache,
+        use_seafloor_cache=use_seafloor_cache,
+        verbose=verbose,
+        cleanup_weights=cleanup_weights,
+        weight_cache_dir=weight_cache_dir,
+    )
+    if overwrite:
+        seafloor_intermediate = input_path.parent / f"{input_path.stem}_seafloor{input_path.suffix}"
+        if seafloor_intermediate.exists():
+            seafloor_intermediate.unlink()
+    seafloor_success = seafloor_pipeline.regrid_file(
+        input_path=input_path,
+        output_path=output_dir,
+        overwrite=overwrite,
+        ui=None,
+    )
+
+    # 2) Surface: extract top level and regrid
+    top_pipeline = CDORegridPipeline(
+        target_resolution=target_resolution,
+        extract_surface=True,
+        extract_seafloor=False,
+        use_regrid_cache=use_regrid_cache,
+        use_seafloor_cache=use_seafloor_cache,
+        verbose=verbose,
+        cleanup_weights=cleanup_weights,
+        weight_cache_dir=weight_cache_dir,
+    )
+    top_success = top_pipeline.regrid_file(
+        input_path=input_path,
+        output_path=output_dir,
+        overwrite=overwrite,
+        ui=None,
+    )
+
+    return {"top_level": top_success, "seafloor": seafloor_success}
+
+
+# Alias for CLI --extreme-levels
+regrid_single_file_extreme_levels = regrid_single_file_both_levels
+
+
+def process_directory_both_levels(
+    dir_path: Path,
+    file_list: list[Path],
+    output_dir: Optional[Path] = None,
+    target_resolution: tuple[float, float] = (1.0, 1.0),
+    use_regrid_cache: bool = True,
+    use_seafloor_cache: bool = True,
+    verbose: bool = True,
+    overwrite: bool = False,
+) -> list[tuple[Path, dict[str, bool]]]:
+    """
+    Process all files in one directory with shared pipelines (seafloor cache + weight cache).
+    Per file: seafloor then surface (extreme levels). Call from a single worker per directory.
+    """
+    weight_cache_dir = dir_path / "cdo_weights"
+    seafloor_pipeline = CDORegridPipeline(
+        target_resolution=target_resolution,
+        extract_surface=False,
+        extract_seafloor=True,
+        use_regrid_cache=use_regrid_cache,
+        use_seafloor_cache=use_seafloor_cache,
+        verbose=verbose,
+        cleanup_weights=False,
+        weight_cache_dir=weight_cache_dir,
+    )
+    top_pipeline = CDORegridPipeline(
+        target_resolution=target_resolution,
+        extract_surface=True,
+        extract_seafloor=False,
+        use_regrid_cache=use_regrid_cache,
+        use_seafloor_cache=use_seafloor_cache,
+        verbose=verbose,
+        cleanup_weights=False,
+        weight_cache_dir=weight_cache_dir,
+    )
+    results: list[tuple[Path, dict[str, bool]]] = []
+    for fp in file_list:
+        seafloor_ok = seafloor_pipeline.regrid_file(
+            input_path=fp,
+            output_path=output_dir,
+            overwrite=overwrite,
+            ui=None,
+        )
+        top_ok = top_pipeline.regrid_file(
+            input_path=fp,
+            output_path=output_dir,
+            overwrite=overwrite,
+            ui=None,
+        )
+        results.append((fp, {"top_level": top_ok, "seafloor": seafloor_ok}))
+    return results
+
+
+def _worker_process_directory_both_levels(
+    args: tuple,
+) -> list[tuple[Path, dict[str, bool]]]:
+    """Picklable worker: process one directory (all files sequentially) with shared pipelines."""
+    dir_path, file_list, output_dir, target_resolution, verbose, overwrite = args
+    return process_directory_both_levels(
+        dir_path=dir_path,
+        file_list=file_list,
+        output_dir=output_dir,
+        target_resolution=target_resolution,
+        verbose=verbose,
+        overwrite=overwrite,
+    )
+
+
+def _worker_both_levels(
+    args: tuple,
+) -> tuple[Path, dict[str, bool]]:
+    """Picklable worker for regrid_single_file_both_levels. Used by ProcessPoolExecutor.
+    Uses a per-process weight cache dir to avoid races on weight generation/reuse.
+    Seafloor depth cache is in-memory per pipeline, so no cross-process conflict.
+    """
+    input_path, output_dir, target_resolution, verbose, overwrite = args
+    # Per-process weight dir so parallel workers don't share weight files (avoid races)
+    weight_cache_dir = Path(tempfile.gettempdir()) / f"cdo_weights_{os.getpid()}"
+    status = regrid_single_file_both_levels(
+        input_path=input_path,
+        output_dir=output_dir,
+        target_resolution=target_resolution,
+        verbose=verbose,
+        cleanup_weights=False,
+        overwrite=overwrite,
+        weight_cache_dir=weight_cache_dir,
+    )
+    return (input_path, status)
+
+
+def regrid_directory(
+    input_dir: Path,
+    include_subdirectories: bool = False,
+    output_dir: Optional[Path] = None,
+    target_resolution: tuple[float, float] = None,
+    file_pattern: str = "*.nc",
+    variables: Optional[list[str] | str] = None,
+    extract_surface: bool = False,
+    extract_seafloor: bool = False,
+    use_regrid_cache: bool = True,
+    use_seafloor_cache: bool = True,
+    verbose: bool = True,
+    verbose_diagnostics: bool = False,
+    max_workers: Optional[int] = 4,
+    enable_parallel: bool = True,
+    overwrite: bool = False,
+    use_ui: bool = True,
+) -> dict[str, list[Path]]:
+    """
+    Convenience function to regrid all files in a directory.
+
+    Args:
+    - input_dir (Path): Input directory containing NetCDF files
+    - output_dir (Path): Output directory for regridded files
+    - target_resolution (tuple): Target resolution as (lon_res, lat_res)
+    - file_pattern (str): File pattern to match (e.g., "*.nc", "*.nc4")
+    - variables (list[str] | str, optional): CMIP variable prefix(es) to process (e.g. ``tos`` or ``tos,uo``)
+    - extract_surface (bool): Extract top level only and regrid that
+    - extract_seafloor (bool): Extract seafloor values and regrid only that
+    - use_regrid_cache (bool): Reuse existing regrid weight files
+    - use_seafloor_cache (bool): Reuse seafloor depth indices cache
+    - verbose (bool): Enable verbose output (progress UI)
+    - verbose_diagnostics (bool): If True, print Grid type, File size, Large file (max verbosity)
+    - max_workers (int): Maximum number of parallel workers
+    - enable_parallel (bool): Enable parallel processing
+    - overwrite (bool): If True, overwrite existing output files
+    - use_ui (bool): Use rich progress UI
+
+    Returns (dict[str, list[Path]]): Dictionary mapping status to list of file paths
+    """
+    # find all matching files (exclude intermediates: _top_level, _regridded, _seafloor, _chunk_)
+    if include_subdirectories:
+        raw = list(input_dir.rglob(file_pattern))
+    else:
+        raw = list(input_dir.glob(file_pattern))
+    input_files = [p for p in raw if not is_intermediate_nc(p)]
+    input_files = filter_files_by_variables(input_files, variables)
+
+    if not input_files:
+        var_msg = f" and variables {sorted(parse_variable_list(variables) or [])}" if variables else ""
+        print(f"No files found matching pattern '{file_pattern}'{var_msg} in {input_dir}")
+        return {"successful": [], "failed": [], "skipped": []}
+
+    error_log_path = init_regrid_error_log()
+    if verbose:
+        print(f"Regrid errors log: {error_log_path}")
+
+    # create pipeline
+    pipeline = CDORegridPipeline(
+        target_resolution=target_resolution,
+        extract_surface=extract_surface,
+        extract_seafloor=extract_seafloor,
+        use_regrid_cache=use_regrid_cache,
+        use_seafloor_cache=use_seafloor_cache,
+        verbose=verbose,
+        verbose_diagnostics=verbose_diagnostics,
+        max_workers=max_workers,
+        enable_parallel=enable_parallel,
+    )
+    pipeline.set_error_log_path(error_log_path)
+
+    # process files
+    results = pipeline.regrid_batch(input_files, output_dir, overwrite=overwrite, use_ui=use_ui)
+    # print statistics
+    pipeline.print_statistics()
+    
+    return results
+
+
+def regrid_directory_both_levels(
+    input_dir: Path,
+    include_subdirectories: bool = False,
+    output_dir: Optional[Path] = None,
+    target_resolution: tuple[float, float] = (1.0, 1.0),
+    file_pattern: str = "*.nc",
+    variables: Optional[list[str] | str] = None,
+    verbose: bool = True,
+    overwrite: bool = False,
+    max_workers: Optional[int] = 4,
+    enable_parallel: bool = True,
+) -> dict[str, list[Path]]:
+    """
+    Regrid both the top level and the seafloor values for all files
+    in a directory.
+
+    This is a higher-level orchestrator that calls
+    :func:`regrid_single_file_both_levels` for each matching file and
+    aggregates the results into the same status dictionary structure
+    as :func:`regrid_directory`.     When ``enable_parallel`` is True and there are multiple directories, directories
+    are processed in parallel (one worker per directory). Within each directory
+    all files are processed sequentially with shared pipelines, so per-directory
+    seafloor depth cache and weight cache (``<dir>/cdo_weights``) are reused and
+    no workers contend for the same directory.
+
+    Parameters
+    ----------
+    input_dir : Path
+        Directory containing input NetCDF files.
+    include_subdirectories : bool
+        Recurse into subdirectories.
+    output_dir : Path, optional
+        Directory for outputs (defaults to alongside inputs when None).
+    target_resolution : tuple[float, float]
+        Target grid resolution.
+    file_pattern : str
+        Glob pattern for selecting input files.
+    verbose : bool
+        Enable verbose logging.
+    overwrite : bool
+        Overwrite existing outputs.
+    max_workers : int, optional
+        Maximum parallel workers (default 4). Used only if enable_parallel is True.
+    enable_parallel : bool
+        Process multiple files in parallel (default True).
+
+    Returns
+    -------
+    dict[str, list[Path]]
+        Dictionary with ``'successful'``, ``'failed'`` and ``'skipped'``
+        keys mapping to lists of input file paths.
+    """
+    # Only process source files; exclude intermediates (_top_level, _regridded, _seafloor, _chunk_)
+    if include_subdirectories:
+        raw = list(input_dir.rglob(file_pattern))
+    else:
+        raw = list(input_dir.glob(file_pattern))
+    input_files = [p for p in raw if not is_intermediate_nc(p)]
+    input_files = filter_files_by_variables(input_files, variables)
+
+    if not input_files:
+        var_msg = f" and variables {sorted(parse_variable_list(variables) or [])}" if variables else ""
+        print(f"No files found matching pattern '{file_pattern}'{var_msg} in {input_dir}")
+        return {"successful": [], "failed": [], "skipped": []}
+
+    results: dict[str, list[Path]] = {
+        "successful": [],
+        "failed": [],
+        "skipped": [],
+    }
+
+    # Group files by directory so one worker owns each directory (reuses seafloor cache + weight cache)
+    by_dir: dict[Path, list[Path]] = {}
+    for fp in input_files:
+        by_dir.setdefault(fp.parent, []).append(fp)
+    dir_jobs = [(d, sorted(fs)) for d, fs in by_dir.items()]
+
+    use_parallel = (
+        enable_parallel
+        and len(dir_jobs) >= 2
+        and max_workers is not None
+        and max_workers > 1
+    )
+
+    if use_parallel:
+        worker_args = [
+            (dir_path, file_list, output_dir, target_resolution, verbose, overwrite)
+            for dir_path, file_list in dir_jobs
+        ]
+        n_workers = min(max_workers, len(dir_jobs), mp.cpu_count())
+        if verbose:
+            print(f"Processing {len(dir_jobs)} directories in parallel with {n_workers} workers (files within each directory processed sequentially with shared cache).")
+        with ProcessPoolExecutor(max_workers=n_workers) as executor:
+            for file_results in executor.map(_worker_process_directory_both_levels, worker_args):
+                for input_path, status in file_results:
+                    if status["top_level"] and status["seafloor"]:
+                        results["successful"].append(input_path)
+                    elif status["top_level"] or status["seafloor"]:
+                        results["failed"].append(input_path)
+                    else:
+                        results["failed"].append(input_path)
+    else:
+        for dir_path, file_list in dir_jobs:
+            try:
+                file_results = process_directory_both_levels(
+                    dir_path=dir_path,
+                    file_list=file_list,
+                    output_dir=output_dir,
+                    target_resolution=target_resolution,
+                    verbose=verbose,
+                    overwrite=overwrite,
+                )
+                for fp, status in file_results:
+                    if status["top_level"] and status["seafloor"]:
+                        results["successful"].append(fp)
+                    elif status["top_level"] or status["seafloor"]:
+                        results["failed"].append(fp)
+                    else:
+                        results["failed"].append(fp)
+            except Exception:
+                for fp in file_list:
+                    results["failed"].append(fp)
+
+    return results
+
+def regrid_large_files(
+    input_files: list[Path],
+    output_dir: Optional[Path] = None,
+    target_resolution: tuple[float, float] = (1.0, 1.0),
+    chunk_size_gb: float = 2.0,
+    max_memory_gb: float = 8.0,
+    verbose: bool = True,
+    overwrite: bool = False,
+) -> dict[str, list[Path]]:
+    """
+    Convenience function for regridding large files with memory optimization.
+    
+    Args:
+    - input_files (list[Path]): List of input files to regrid
+    - output_dir (Path): Output directory for regridded files
+    - target_resolution (tuple): Target resolution as (lon_res, lat_res)
+    - chunk_size_gb (float): Maximum chunk size in GB
+    - max_memory_gb (float): Maximum memory usage in GB
+    - verbose (bool): Enable verbose output
+    - overwrite (bool): If True, overwrite existing output files
+    
+    Returns (dict[str, list[Path]]): Dictionary mapping status to list of file paths
+    """
+    pipeline = CDORegridPipeline(
+        target_resolution=target_resolution,
+        chunk_size_gb=chunk_size_gb,
+        max_memory_gb=max_memory_gb,
+        verbose=verbose,
+        enable_chunking=True,
+        memory_monitoring=True,
+    )
+    
+    results = pipeline.regrid_batch(input_files, output_dir, overwrite=overwrite)
+    pipeline.print_statistics()
+    
+    return results
+