cdo-toolkit 0.1.0__py3-none-any.whl

cdo_toolkit/cli.py

@@ -0,0 +1,166 @@
+"""Command-line interface for CDO regridding."""
+
+import argparse
+from pathlib import Path
+
+from rich.console import Console
+
+from cdo_toolkit.api import (
+    filter_files_by_variables,
+    get_cmip_variable_name,
+    parse_variable_list,
+    regrid_directory,
+    regrid_directory_both_levels,
+    regrid_single_file,
+    regrid_single_file_extreme_levels,
+)
+
+
+def main() -> None:
+        parser = argparse.ArgumentParser(description="CDO-based NetCDF regridding pipeline")
+        parser.add_argument("input", type=Path, help="Input file or directory")
+        parser.add_argument("-o", "--output", type=Path, help="Output file or directory")
+        parser.add_argument("-r", "--resolution", nargs=2, type=float, default=[1.0, 1.0],
+                           help="Target resolution (lon_res lat_res)")
+        parser.add_argument("-p", "--pattern", default="*.nc", help="File pattern (for directories)")
+        parser.add_argument(
+            "--variable", "-V", nargs="+", default=None,
+            help="CMIP variable prefix(es) to regrid (e.g. tos or tos uo). Comma-separated values allowed.",
+        )
+        parser.add_argument("--include-subdirectories", action="store_true", default=True, help="Include subdirectories")
+        parser.add_argument("--extract-surface", action="store_true", default=False,
+                            help="Extract top level only and regrid that (surface)")
+        parser.add_argument("--extract-seafloor", action="store_true", default=False,
+                            help="Extract seafloor values and regrid only that")
+        parser.add_argument("--extreme-levels", action="store_true", default=False,
+                            help="Extract and regrid both surface (top level) and seafloor for each file")
+        parser.add_argument("--no-regrid-cache", action="store_true", default=False,
+                            help="Do not reuse regrid weight cache (regenerate weights each time)")
+        parser.add_argument("--no-seafloor-cache", action="store_true", default=False,
+                            help="Do not reuse seafloor depth indices cache")
+        parser.add_argument("-v", "--verbose", action="store_true", default=True, help="Verbose output (progress UI)")
+        parser.add_argument("--verbose-max", action="store_true", default=False,
+                            help="Maximum verbosity: print Grid type, File size, Large file messages")
+        parser.add_argument("--quiet", action="store_true", help="Disable verbose output")
+        parser.add_argument("--max-workers", "-w", default=4, type=int, help="Maximum parallel workers")
+        parser.add_argument("--chunk-size-gb", type=float, default=2.0,
+                           help="Maximum chunk size in GB")
+        parser.add_argument("--max-memory-gb", default=8.0, type=float, help="Maximum memory usage in GB")
+        parser.add_argument("--no-parallel", action="store_true", default=False, help="Disable parallel processing")
+        parser.add_argument("--no-chunking", action="store_true", default=False, help="Disable chunked processing")
+        parser.add_argument("--use-ui", action="store_true", default=True, help="Use UI for processing")
+        # parser.add_argument("--cleanup", action="store_true", help="Clean up problematic files (*_top_level, *_chunk_*) before processing")
+        parser.add_argument("--unlink-unprocessed", action="store_true", default=False, help="Unlink unprocessed files after processing")
+        parser.add_argument("--overwrite", action="store_true", default=False, help="Overwrite existing output files")
+        
+        args = parser.parse_args()
+    
+        # handle verbose/quiet logic
+        verbose = args.verbose and not args.quiet
+        verbose_diagnostics = getattr(args, "verbose_max", False)
+        use_regrid_cache = not args.no_regrid_cache
+        use_seafloor_cache = not args.no_seafloor_cache
+    
+        # # handle cleanup if requested
+        # if args.cleanup:
+        #     if args.input.is_file():
+        #         # clean up in the same directory as the file
+        #         cleaned_count = cleanup_problematic_files(args.input.parent, verbose=verbose)
+        #     else:
+        #         # clean up in the directory
+        #         cleaned_count = cleanup_problematic_files(args.input, verbose=verbose)
+            
+        #     if cleaned_count == 0:
+        #         print("No problematic files found to clean up.")
+        #     else:
+        #         print(f"Cleaned up {cleaned_count} problematic files.")
+            
+        #     # exit after cleanup
+        #     exit(0)
+        
+        variables = args.variable
+    
+        # determine if input is file or directory
+        if args.input.is_file():
+            if variables and not filter_files_by_variables([args.input], variables):
+                allowed = sorted(parse_variable_list(variables) or [])
+                print(f"Skipping {args.input.name}: variable '{get_cmip_variable_name(args.input)}' not in {allowed}")
+                raise SystemExit(0)
+            # single file processing
+            if args.extreme_levels:
+                status = regrid_single_file_extreme_levels(
+                    input_path=args.input,
+                    output_dir=args.output,
+                    target_resolution=tuple(args.resolution),
+                    use_regrid_cache=use_regrid_cache,
+                    use_seafloor_cache=use_seafloor_cache,
+                    verbose=verbose,
+                    overwrite=args.overwrite,
+                )
+                success = status["top_level"] and status["seafloor"]
+            else:
+                out_path = getattr(args, "output", None)
+                out_dir = out_path if (out_path and out_path.is_dir()) else None
+                success = regrid_single_file(
+                    input_path=args.input,
+                    output_path=None if out_dir else out_path,
+                    output_dir=out_dir,
+                    target_resolution=tuple(args.resolution),
+                    extract_surface=args.extract_surface,
+                    extract_seafloor=args.extract_seafloor,
+                    use_regrid_cache=use_regrid_cache,
+                    use_seafloor_cache=use_seafloor_cache,
+                    verbose=verbose,
+                    verbose_diagnostics=verbose_diagnostics,
+                    use_ui=args.use_ui,
+                    overwrite=args.overwrite,
+                )
+    
+            if success:
+                print("Regridding successful!")
+            else:
+                print("Regridding failed!")
+        else:
+            # directory processing
+            if args.extreme_levels:
+                results = regrid_directory_both_levels(
+                    input_dir=args.input,
+                    output_dir=args.output,
+                    include_subdirectories=args.include_subdirectories,
+                    target_resolution=tuple(args.resolution),
+                    file_pattern=args.pattern,
+                    variables=variables,
+                    verbose=verbose,
+                    overwrite=args.overwrite,
+                    max_workers=args.max_workers,
+                    enable_parallel=not args.no_parallel,
+                )
+            else:
+                results = regrid_directory(
+                    input_dir=args.input,
+                    output_dir=args.output,
+                    include_subdirectories=args.include_subdirectories,
+                    target_resolution=tuple(args.resolution),
+                    file_pattern=args.pattern,
+                    variables=variables,
+                    extract_surface=args.extract_surface,
+                    extract_seafloor=args.extract_seafloor,
+                    use_regrid_cache=use_regrid_cache,
+                    use_seafloor_cache=use_seafloor_cache,
+                    verbose=verbose,
+                    verbose_diagnostics=verbose_diagnostics,
+                    max_workers=args.max_workers,
+                    enable_parallel=not args.no_parallel,
+                    use_ui=args.use_ui,
+                    overwrite=args.overwrite,
+                )
+            
+            # print results
+            console = Console()
+            console.print(f"\n[green]Successful: {len(results['successful'])}[/green]")
+            console.print(f"[red]Failed: {len(results['failed'])}[/red]")
+            console.print(f"[yellow]Skipped: {len(results['skipped'])}[/yellow]")
+
+
+if __name__ == "__main__":
+    main()

cdo_toolkit/cmip.py

@@ -0,0 +1,61 @@
+"""Optional CMIP6 filename conventions and filters."""
+
+from pathlib import Path
+from typing import Optional
+
+import xarray as xa
+
+def parse_variable_list(variables: Optional[list[str] | str]) -> Optional[set[str]]:
+    """Parse CLI/YAML variable filter into a lowercase set, or None if unset."""
+    if not variables:
+        return None
+    tokens: list[str] = []
+    if isinstance(variables, str):
+        tokens = variables.split(",")
+    else:
+        for item in variables:
+            tokens.extend(item.split(","))
+    parsed = {token.strip().lower() for token in tokens if token.strip()}
+    return parsed or None
+
+
+def get_cmip_variable_name(file_path: Path) -> str:
+    """CMIP6 filename prefix before the first underscore (e.g. tos from tos_Omon_...)."""
+    return file_path.stem.split("_")[0].lower()
+
+
+def filter_files_by_variables(
+    files: list[Path],
+    variables: Optional[list[str] | str],
+) -> list[Path]:
+    """Keep only files whose CMIP variable prefix is in *variables*."""
+    allowed = parse_variable_list(variables)
+    if not allowed:
+        return files
+    return [f for f in files if get_cmip_variable_name(f) in allowed]
+
+
+def pick_representative_file(input_files: list[Path]) -> Optional[Path]:
+    """Pick a file with nominal_resolution metadata, else the first file in the group."""
+    if not input_files:
+        return None
+    for file_path in input_files:
+        try:
+            with xa.open_dataset(file_path, decode_times=False) as ds:
+                if "nominal_resolution" in ds.attrs:
+                    return file_path
+        except Exception:
+            continue
+    return input_files[0]
+
+
+def representative_files_by_directory(files: list[Path]) -> dict[Path, Path]:
+    """Map each file's parent directory to a representative file for resolution."""
+    by_dir: dict[Path, list[Path]] = {}
+    for file_path in files:
+        by_dir.setdefault(file_path.parent, []).append(file_path)
+    return {
+        parent: rep
+        for parent, group_files in by_dir.items()
+        if (rep := pick_representative_file(group_files)) is not None
+    }

cdo_toolkit/constants.py

@@ -0,0 +1,9 @@
+"""Shared constants for the CDO regrid package."""
+
+REGRID_ERROR_LOGGER_NAME = "cdo_toolkit.errors"
+
+NC4_ENCODING_KEYS = frozenset({
+    "szip_pixels_per_block", "contiguous", "quantize_mode", "_FillValue", "fletcher32",
+    "endian", "chunksizes", "least_significant_digit", "complevel", "szip_coding",
+    "significant_digits", "dtype", "shuffle", "zlib", "blosc_shuffle", "compression",
+})

cdo_toolkit/errors.py

@@ -0,0 +1,79 @@
+"""Error logging and file locking for regridding."""
+
+import logging
+import os
+import traceback
+from contextlib import contextmanager
+from datetime import datetime
+from pathlib import Path
+from typing import Optional
+
+try:
+    import fcntl
+except ImportError:
+    fcntl = None
+
+from cdo_toolkit.constants import REGRID_ERROR_LOGGER_NAME
+
+_regrid_error_log_path: Optional[Path] = None
+
+
+def default_log_dir() -> Path:
+    """Default directory for regrid error logs."""
+    return Path.cwd() / "logs"
+
+
+def init_regrid_error_log(path: Optional[Path] = None) -> Path:
+    """One regrid error log per run; safe for parallel workers to append."""
+    global _regrid_error_log_path
+    if path is not None:
+        _regrid_error_log_path = Path(path)
+    elif _regrid_error_log_path is None:
+        log_dir = default_log_dir()
+        log_dir.mkdir(parents=True, exist_ok=True)
+        timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+        _regrid_error_log_path = log_dir / f"regrid_errors_{timestamp}.log"
+
+    logger = logging.getLogger(REGRID_ERROR_LOGGER_NAME)
+    logger.setLevel(logging.WARNING)
+    logger.propagate = False
+    if not any(
+        isinstance(h, logging.FileHandler)
+        and getattr(h, "baseFilename", None) == str(_regrid_error_log_path.resolve())
+        for h in logger.handlers
+    ):
+        handler = logging.FileHandler(_regrid_error_log_path)
+        handler.setFormatter(
+            logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
+        )
+        logger.addHandler(handler)
+    return _regrid_error_log_path
+
+
+def log_regrid_error(message: str, exc: Optional[BaseException] = None) -> None:
+    """Append a regrid error to the shared session log."""
+    init_regrid_error_log()
+    logger = logging.getLogger(REGRID_ERROR_LOGGER_NAME)
+    if exc is not None:
+        logger.error("%s\n%s", message, traceback.format_exc())
+    else:
+        logger.error(message)
+    for handler in logger.handlers:
+        handler.flush()
+
+
+@contextmanager
+def weight_file_lock(weight_path: Path):
+    """Serialize weight generation and validation across parallel workers."""
+    if fcntl is None:
+        yield
+        return
+    lock_path = weight_path.parent / f"{weight_path.name}.lock"
+    lock_path.parent.mkdir(parents=True, exist_ok=True)
+    fd = os.open(str(lock_path), os.O_CREAT | os.O_RDWR, 0o644)
+    try:
+        fcntl.flock(fd, fcntl.LOCK_EX)
+        yield
+    finally:
+        fcntl.flock(fd, fcntl.LOCK_UN)
+        os.close(fd)

cdo_toolkit/memory.py

@@ -0,0 +1,22 @@
+"""Memory usage monitoring."""
+
+import psutil
+
+
+class MemoryMonitor:
+    """Track peak memory usage during processing."""
+
+    def __init__(self):
+        self.peak_memory_gb = 0.0
+
+    def get_memory_usage_gb(self) -> float:
+        process = psutil.Process()
+        return process.memory_info().rss / (1024**3)
+
+    def update_peak(self):
+        current = self.get_memory_usage_gb()
+        if current > self.peak_memory_gb:
+            self.peak_memory_gb = current
+
+    def get_peak_memory_gb(self) -> float:
+        return self.peak_memory_gb

cdo_toolkit/paths.py

@@ -0,0 +1,30 @@
+"""Path helpers for weight caches and intermediate files."""
+
+from pathlib import Path
+
+
+def weight_cache_dir_for_input(input_path: Path) -> Path:
+    """Per-leaf-directory CDO weight cache (avoids cross-model collisions)."""
+    return input_path.parent / "cdo_weights"
+
+
+def is_weights_or_cache_file(path: Path) -> bool:
+    """Return True if path is a weight file or under a weight cache directory (exclude from regrid list)."""
+    if path.stem.lower().startswith("weights_"):
+        return True
+    if "cdo_weights" in path.parts:
+        return True
+    return False
+
+
+def is_intermediate_nc(path: Path) -> bool:
+    """Return True if path is an intermediate/product we should not regrid (top_level, regridded, seafloor, chunk, weights)."""
+    if is_weights_or_cache_file(path):
+        return True
+    stem = path.stem.lower()
+    return (
+        "_top_level" in stem
+        or "_regridded" in stem
+        or "_seafloor" in stem
+        or "_chunk_" in stem
+    )