quantmllibrary 0.1.0__py3-none-any.whl

quantml/utils/reproducibility.py

@@ -0,0 +1,220 @@
+"""
+Reproducibility utilities for QuantML.
+
+Provides random seed management, version tracking, and experiment metadata.
+"""
+
+import random
+import os
+import sys
+from typing import Optional, Dict, Any
+from datetime import datetime
+import json
+
+# Try to import NumPy
+try:
+    import numpy as np
+    HAS_NUMPY = True
+except ImportError:
+    HAS_NUMPY = False
+    np = None
+
+# Try to get git info
+def get_git_info() -> Dict[str, Optional[str]]:
+    """Get git repository information."""
+    info = {
+        'commit_hash': None,
+        'branch': None,
+        'is_dirty': None
+    }
+    
+    try:
+        import subprocess
+        
+        # Get commit hash
+        try:
+            commit_hash = subprocess.check_output(
+                ['git', 'rev-parse', 'HEAD'],
+                stderr=subprocess.DEVNULL
+            ).decode().strip()
+            info['commit_hash'] = commit_hash
+        except Exception:
+            pass
+        
+        # Get branch
+        try:
+            branch = subprocess.check_output(
+                ['git', 'rev-parse', '--abbrev-ref', 'HEAD'],
+                stderr=subprocess.DEVNULL
+            ).decode().strip()
+            info['branch'] = branch
+        except Exception:
+            pass
+        
+        # Check if dirty
+        try:
+            result = subprocess.run(
+                ['git', 'diff', '--quiet'],
+                stderr=subprocess.DEVNULL
+            )
+            info['is_dirty'] = result.returncode != 0
+        except Exception:
+            pass
+    
+    except Exception:
+        pass
+    
+    return info
+
+
+def get_library_version() -> str:
+    """Get QuantML library version."""
+    try:
+        from quantml import __version__
+        return __version__
+    except ImportError:
+        return "unknown"
+
+
+def set_random_seed(seed: int, use_cuda: bool = False):
+    """
+    Set random seed for reproducibility.
+    
+    Args:
+        seed: Random seed value
+        use_cuda: Whether to set CUDA random seed (if available)
+    """
+    random.seed(seed)
+    os.environ['PYTHONHASHSEED'] = str(seed)
+    
+    if HAS_NUMPY:
+        np.random.seed(seed)
+    
+    # Try to set PyTorch seed if available
+    try:
+        import torch
+        torch.manual_seed(seed)
+        if use_cuda and torch.cuda.is_available():
+            torch.cuda.manual_seed_all(seed)
+            torch.backends.cudnn.deterministic = True
+            torch.backends.cudnn.benchmark = False
+    except ImportError:
+        pass
+    
+    # Try to set TensorFlow seed if available
+    try:
+        import tensorflow as tf
+        tf.random.set_seed(seed)
+    except ImportError:
+        pass
+
+
+def get_environment_info() -> Dict[str, Any]:
+    """
+    Get environment information for reproducibility.
+    
+    Returns:
+        Dictionary with environment details
+    """
+    info = {
+        'timestamp': datetime.now().isoformat(),
+        'python_version': sys.version,
+        'platform': sys.platform,
+        'quantml_version': get_library_version(),
+        'git_info': get_git_info()
+    }
+    
+    # Add NumPy version if available
+    if HAS_NUMPY:
+        info['numpy_version'] = np.__version__
+    
+    # Try to get other library versions
+    for lib_name in ['pandas', 'scipy', 'sklearn']:
+        try:
+            lib = __import__(lib_name)
+            if hasattr(lib, '__version__'):
+                info[f'{lib_name}_version'] = lib.__version__
+        except ImportError:
+            pass
+    
+    return info
+
+
+def save_experiment_metadata(
+    metadata: Dict[str, Any],
+    filepath: str
+):
+    """
+    Save experiment metadata to file.
+    
+    Args:
+        metadata: Metadata dictionary
+        filepath: Path to save metadata
+    """
+    os.makedirs(os.path.dirname(filepath) or '.', exist_ok=True)
+    
+    with open(filepath, 'w') as f:
+        json.dump(metadata, f, indent=2)
+
+
+def load_experiment_metadata(filepath: str) -> Dict[str, Any]:
+    """
+    Load experiment metadata from file.
+    
+    Args:
+        filepath: Path to metadata file
+    
+    Returns:
+        Metadata dictionary
+    """
+    with open(filepath, 'r') as f:
+        return json.load(f)
+
+
+def create_experiment_id(prefix: str = "exp") -> str:
+    """
+    Create a unique experiment ID.
+    
+    Args:
+        prefix: Prefix for experiment ID
+    
+    Returns:
+        Experiment ID string
+    """
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    random_suffix = random.randint(1000, 9999)
+    return f"{prefix}_{timestamp}_{random_suffix}"
+
+
+class ReproducibilityContext:
+    """Context manager for reproducible experiments."""
+    
+    def __init__(self, seed: int, experiment_name: Optional[str] = None):
+        """
+        Initialize reproducibility context.
+        
+        Args:
+            seed: Random seed
+            experiment_name: Optional experiment name
+        """
+        self.seed = seed
+        self.experiment_name = experiment_name
+        self.metadata = None
+    
+    def __enter__(self):
+        """Enter context and set random seed."""
+        set_random_seed(self.seed)
+        self.metadata = get_environment_info()
+        self.metadata['random_seed'] = self.seed
+        if self.experiment_name:
+            self.metadata['experiment_name'] = self.experiment_name
+        return self
+    
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Exit context."""
+        pass
+    
+    def get_metadata(self) -> Dict[str, Any]:
+        """Get experiment metadata."""
+        return self.metadata.copy() if self.metadata else {}
+

quantml/utils/serialization.py

@@ -0,0 +1,335 @@
+"""
+Model serialization utilities for saving and loading models.
+
+This module provides functions to save and load model weights,
+as well as full training checkpoints.
+"""
+
+import json
+import os
+from typing import Optional, Dict, Any, List, Union
+from quantml.tensor import Tensor
+
+# Try to import NumPy for efficient storage
+try:
+    import numpy as np
+    HAS_NUMPY = True
+except ImportError:
+    HAS_NUMPY = False
+    np = None
+
+
+def save_model(model, path: str, format: str = 'json') -> None:
+    """
+    Save model weights to file.
+    
+    Args:
+        model: Model with parameters() method
+        path: Path to save weights
+        format: 'json' or 'npz' (requires NumPy)
+    
+    Examples:
+        >>> from quantml.models import Linear
+        >>> model = Linear(10, 1)
+        >>> save_model(model, 'model_weights.json')
+    """
+    params = model.parameters()
+    
+    if format == 'json':
+        _save_json(params, path)
+    elif format == 'npz':
+        if not HAS_NUMPY:
+            raise RuntimeError("NumPy required for npz format")
+        _save_npz(params, path)
+    else:
+        raise ValueError(f"Unknown format: {format}. Use 'json' or 'npz'")
+
+
+def load_model(model, path: str, format: Optional[str] = None) -> None:
+    """
+    Load model weights from file.
+    
+    Args:
+        model: Model with parameters() method (weights will be loaded in-place)
+        path: Path to load weights from
+        format: 'json' or 'npz' (auto-detected from extension if None)
+    
+    Examples:
+        >>> from quantml.models import Linear
+        >>> model = Linear(10, 1)
+        >>> load_model(model, 'model_weights.json')
+    """
+    if format is None:
+        # Auto-detect from extension
+        if path.endswith('.json'):
+            format = 'json'
+        elif path.endswith('.npz'):
+            format = 'npz'
+        else:
+            format = 'json'  # Default
+    
+    params = model.parameters()
+    
+    if format == 'json':
+        loaded = _load_json(path)
+    elif format == 'npz':
+        if not HAS_NUMPY:
+            raise RuntimeError("NumPy required for npz format")
+        loaded = _load_npz(path)
+    else:
+        raise ValueError(f"Unknown format: {format}")
+    
+    # Update model parameters
+    if len(loaded) != len(params):
+        raise ValueError(f"Parameter count mismatch: model has {len(params)}, file has {len(loaded)}")
+    
+    for i, (param, loaded_data) in enumerate(zip(params, loaded)):
+        _update_tensor_data(param, loaded_data)
+
+
+def save_checkpoint(
+    model,
+    optimizer,
+    epoch: int,
+    path: str,
+    loss: Optional[float] = None,
+    metrics: Optional[Dict[str, float]] = None,
+    extra: Optional[Dict[str, Any]] = None
+) -> None:
+    """
+    Save a full training checkpoint.
+    
+    Includes model weights, optimizer state, epoch, and optional metrics.
+    
+    Args:
+        model: Model with parameters() method
+        optimizer: Optimizer with state to save
+        epoch: Current epoch number
+        path: Path to save checkpoint
+        loss: Optional current loss value
+        metrics: Optional dict of metric values
+        extra: Optional extra data to save
+    
+    Examples:
+        >>> save_checkpoint(model, optimizer, epoch=10, path='checkpoint.json')
+    """
+    checkpoint = {
+        'epoch': epoch,
+        'model_state': _params_to_list(model.parameters()),
+        'optimizer_state': _get_optimizer_state(optimizer),
+    }
+    
+    if loss is not None:
+        checkpoint['loss'] = loss
+    if metrics is not None:
+        checkpoint['metrics'] = metrics
+    if extra is not None:
+        checkpoint['extra'] = extra
+    
+    # Save as JSON
+    with open(path, 'w') as f:
+        json.dump(checkpoint, f, indent=2)
+
+
+def load_checkpoint(
+    model,
+    optimizer,
+    path: str
+) -> Dict[str, Any]:
+    """
+    Load a training checkpoint.
+    
+    Restores model weights and optimizer state.
+    
+    Args:
+        model: Model with parameters() method
+        optimizer: Optimizer to restore state to
+        path: Path to load checkpoint from
+    
+    Returns:
+        Dict with 'epoch', 'loss', 'metrics', and any extra data
+    
+    Examples:
+        >>> info = load_checkpoint(model, optimizer, 'checkpoint.json')
+        >>> print(f"Resuming from epoch {info['epoch']}")
+    """
+    with open(path, 'r') as f:
+        checkpoint = json.load(f)
+    
+    # Restore model weights
+    params = model.parameters()
+    loaded = checkpoint['model_state']
+    
+    if len(loaded) != len(params):
+        raise ValueError(f"Parameter count mismatch: model has {len(params)}, checkpoint has {len(loaded)}")
+    
+    for param, loaded_data in zip(params, loaded):
+        _update_tensor_data(param, loaded_data)
+    
+    # Restore optimizer state
+    if 'optimizer_state' in checkpoint:
+        _set_optimizer_state(optimizer, checkpoint['optimizer_state'])
+    
+    # Return info
+    return {
+        'epoch': checkpoint.get('epoch', 0),
+        'loss': checkpoint.get('loss'),
+        'metrics': checkpoint.get('metrics'),
+        'extra': checkpoint.get('extra')
+    }
+
+
+def get_model_state_dict(model) -> Dict[str, List]:
+    """
+    Get model state as a dictionary.
+    
+    Args:
+        model: Model with parameters() method
+    
+    Returns:
+        Dict mapping parameter names to their data
+    """
+    params = model.parameters()
+    state = {}
+    for i, param in enumerate(params):
+        name = f"param_{i}"
+        if hasattr(param, 'data'):
+            state[name] = _tensor_to_list(param)
+        else:
+            state[name] = param
+    return state
+
+
+def set_model_state_dict(model, state_dict: Dict[str, List]) -> None:
+    """
+    Set model state from a dictionary.
+    
+    Args:
+        model: Model with parameters() method
+        state_dict: Dict mapping parameter names to their data
+    """
+    params = model.parameters()
+    for i, param in enumerate(params):
+        name = f"param_{i}"
+        if name in state_dict:
+            _update_tensor_data(param, state_dict[name])
+
+
+# ============================================================================
+# Private helper functions
+# ============================================================================
+
+def _save_json(params: List[Tensor], path: str) -> None:
+    """Save parameters to JSON file."""
+    data = _params_to_list(params)
+    with open(path, 'w') as f:
+        json.dump(data, f, indent=2)
+
+
+def _load_json(path: str) -> List:
+    """Load parameters from JSON file."""
+    with open(path, 'r') as f:
+        return json.load(f)
+
+
+def _save_npz(params: List[Tensor], path: str) -> None:
+    """Save parameters to NumPy npz file."""
+    arrays = {}
+    for i, param in enumerate(params):
+        arr = _tensor_to_numpy(param)
+        arrays[f'param_{i}'] = arr
+    np.savez(path, **arrays)
+
+
+def _load_npz(path: str) -> List:
+    """Load parameters from NumPy npz file."""
+    data = np.load(path)
+    # Sort by parameter index
+    keys = sorted(data.files, key=lambda x: int(x.split('_')[1]))
+    return [data[k].tolist() for k in keys]
+
+
+def _params_to_list(params: List[Tensor]) -> List:
+    """Convert list of Tensors to list of nested lists."""
+    return [_tensor_to_list(p) for p in params]
+
+
+def _tensor_to_list(t: Tensor) -> List:
+    """Convert Tensor to nested list."""
+    data = t.data
+    if isinstance(data, list):
+        return data
+    # NumPy array
+    try:
+        return data.tolist()
+    except AttributeError:
+        return list(data)
+
+
+def _tensor_to_numpy(t: Tensor):
+    """Convert Tensor to NumPy array."""
+    if hasattr(t, 'numpy') and t.numpy is not None:
+        return t.numpy
+    data = t.data
+    if isinstance(data, list):
+        return np.array(data)
+    return data
+
+
+def _update_tensor_data(param: Tensor, data) -> None:
+    """Update tensor data in-place."""
+    if isinstance(data, list):
+        param._data = data
+    else:
+        # NumPy array
+        param._data = data.tolist() if hasattr(data, 'tolist') else list(data)
+    
+    # Clear cached numpy array if it exists
+    if hasattr(param, '_np_array'):
+        param._np_array = None
+
+
+def _get_optimizer_state(optimizer) -> Dict[str, Any]:
+    """Extract optimizer state for serialization."""
+    state = {
+        'lr': getattr(optimizer, 'lr', None),
+        'step': getattr(optimizer, '_step', 0),
+    }
+    
+    # Handle Adam-like optimizers with momentum buffers
+    if hasattr(optimizer, '_m'):
+        state['m'] = [_tensor_to_list(m) if hasattr(m, 'data') else m 
+                      for m in optimizer._m]
+    if hasattr(optimizer, '_v'):
+        state['v'] = [_tensor_to_list(v) if hasattr(v, 'data') else v 
+                      for v in optimizer._v]
+    
+    # Handle SGD momentum
+    if hasattr(optimizer, '_velocity'):
+        state['velocity'] = [_tensor_to_list(v) if hasattr(v, 'data') else v 
+                            for v in optimizer._velocity]
+    
+    return state
+
+
+def _set_optimizer_state(optimizer, state: Dict[str, Any]) -> None:
+    """Restore optimizer state from serialized data."""
+    if 'step' in state:
+        optimizer._step = state['step']
+    
+    # Handle Adam-like optimizers
+    if 'm' in state and hasattr(optimizer, '_m'):
+        for i, m_data in enumerate(state['m']):
+            if i < len(optimizer._m):
+                _update_tensor_data(optimizer._m[i], m_data)
+    
+    if 'v' in state and hasattr(optimizer, '_v'):
+        for i, v_data in enumerate(state['v']):
+            if i < len(optimizer._v):
+                _update_tensor_data(optimizer._v[i], v_data)
+    
+    # Handle SGD momentum
+    if 'velocity' in state and hasattr(optimizer, '_velocity'):
+        for i, v_data in enumerate(state['velocity']):
+            if i < len(optimizer._velocity):
+                _update_tensor_data(optimizer._velocity[i], v_data)