| """
|
| I/O Utility Functions
|
|
|
| Provides file system operations and tensor serialization utilities.
|
| All operations are designed for safety and reproducibility.
|
|
|
| Design Decisions:
|
| - Atomic file operations where possible
|
| - Type-safe tensor serialization
|
| - Directory creation on demand
|
| - No silent failures
|
|
|
| Time Complexity: O(n) for file operations where n = file size
|
| Space Complexity: O(n) for in-memory tensor operations
|
| """
|
|
|
| from pathlib import Path
|
| from typing import Any, Dict, Optional, Union
|
| import numpy as np
|
|
|
|
|
| def ensure_dir(path: Union[str, Path]) -> Path:
|
| """
|
| Ensure a directory exists, creating it if necessary.
|
|
|
| Args:
|
| path: Directory path (can be file path, will use parent)
|
|
|
| Returns:
|
| Path object for the directory
|
| """
|
| path = Path(path)
|
|
|
|
|
| if path.suffix:
|
| path = path.parent
|
|
|
| path.mkdir(parents=True, exist_ok=True)
|
| return path
|
|
|
|
|
| def save_tensor(
|
| data: np.ndarray,
|
| path: Union[str, Path],
|
| metadata: Optional[Dict[str, Any]] = None
|
| ) -> None:
|
| """
|
| Save a NumPy tensor to disk with optional metadata.
|
|
|
| Uses NumPy's compressed format for storage efficiency.
|
| Metadata is stored alongside for reproducibility.
|
|
|
| Args:
|
| data: NumPy array to save
|
| path: Save path (will add .npz extension if not present)
|
| metadata: Optional metadata dictionary
|
|
|
| Example:
|
| >>> save_tensor(arr, "data/processed/train.npz", {"mean": 0.5})
|
| """
|
| path = Path(path)
|
| if path.suffix != ".npz":
|
| path = path.with_suffix(".npz")
|
|
|
| ensure_dir(path)
|
|
|
| save_dict = {"data": data}
|
| if metadata is not None:
|
| save_dict["metadata"] = np.array([metadata], dtype=object)
|
|
|
| np.savez_compressed(path, **save_dict)
|
|
|
|
|
| def load_tensor(path: Union[str, Path]) -> tuple:
|
| """
|
| Load a NumPy tensor from disk with metadata.
|
|
|
| Args:
|
| path: Path to .npz file
|
|
|
| Returns:
|
| Tuple of (data array, metadata dict or None)
|
|
|
| Raises:
|
| FileNotFoundError: If file doesn't exist
|
| """
|
| path = Path(path)
|
| if path.suffix != ".npz":
|
| path = path.with_suffix(".npz")
|
|
|
| if not path.exists():
|
| raise FileNotFoundError(f"Tensor file not found: {path}")
|
|
|
| loaded = np.load(path, allow_pickle=True)
|
| data = loaded["data"]
|
|
|
| metadata = None
|
| if "metadata" in loaded:
|
| metadata = loaded["metadata"].item()
|
|
|
| return data, metadata
|
|
|
|
|
| def get_file_size_mb(path: Union[str, Path]) -> float:
|
| """
|
| Get file size in megabytes.
|
|
|
| Args:
|
| path: Path to file
|
|
|
| Returns:
|
| File size in MB
|
| """
|
| path = Path(path)
|
| if not path.exists():
|
| return 0.0
|
| return path.stat().st_size / (1024 * 1024)
|
|
|
|
|
| def list_files(
|
| directory: Union[str, Path],
|
| pattern: str = "*",
|
| recursive: bool = False
|
| ) -> list:
|
| """
|
| List files in a directory matching a pattern.
|
|
|
| Args:
|
| directory: Directory to search
|
| pattern: Glob pattern (e.g., "*.nc" for NetCDF files)
|
| recursive: If True, search subdirectories
|
|
|
| Returns:
|
| List of Path objects
|
| """
|
| directory = Path(directory)
|
| if not directory.exists():
|
| return []
|
|
|
| if recursive:
|
| return sorted(directory.rglob(pattern))
|
| return sorted(directory.glob(pattern))
|
|
|
