File size: 3,595 Bytes
c2a61b6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 | """
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 has an extension, assume it's a file and use parent
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))
|