additory 0.1.0a2__py3-none-any.whl → 0.1.0a3__py3-none-any.whl

additory/synthetic/engines.py

@@ -1,529 +0,0 @@
-"""
-Distribution strategy engines for synthetic data generation.
-
-Implements the six essential distribution strategies: equal, custom, categorical,
-high_cardinality, numeric_range, and skewed. All engines use polars for
-high-performance data generation.
-"""
-
-from abc import ABC, abstractmethod
-from typing import Dict, Any, List, Optional, Union
-import polars as pl
-import re
-from dataclasses import dataclass
-
-from .models import DistributionType, DistributionStrategy, ValidationResult
-from .exceptions import SyntheticDataError, ValidationError
-
-
-@dataclass
-class DistributionConfig:
-    """Configuration for distribution strategy application."""
-    strategy: DistributionStrategy
-    base_values: List[str]
-    target_rows: int
-    seed: Optional[int] = None
-
-
-class DistributionEngine(ABC):
-    """Abstract base class for distribution strategy engines."""
-    
-    def __init__(self, seed: Optional[int] = None):
-        """Initialize the distribution engine with optional seed."""
-        self.seed = seed
-    
-    @abstractmethod
-    def supports_strategy(self, strategy_type: DistributionType) -> bool:
-        """Check if this engine supports the given strategy type."""
-        pass
-    
-    @abstractmethod
-    def validate_config(self, config: DistributionConfig) -> ValidationResult:
-        """Validate the distribution configuration."""
-        pass
-    
-    @abstractmethod
-    def apply_distribution(self, config: DistributionConfig) -> pl.Series:
-        """Apply the distribution strategy to generate data."""
-        pass
-    
-    def _validate_base_requirements(self, config: DistributionConfig) -> ValidationResult:
-        """Validate basic requirements common to all distributions."""
-        result = ValidationResult(is_valid=True)
-        
-        if config.target_rows <= 0:
-            result.add_error("Target rows must be positive", "Set target_rows to a positive integer")
-        
-        if not config.base_values:
-            result.add_error("Base values cannot be empty", "Provide at least one base value")
-        
-        return result
-
-
-class EqualDistributionEngine(DistributionEngine):
-    """Engine for equal (uniform) distribution strategy."""
-    
-    def supports_strategy(self, strategy_type: DistributionType) -> bool:
-        """Check if this engine supports equal distribution."""
-        return strategy_type == DistributionType.EQUAL
-    
-    def validate_config(self, config: DistributionConfig) -> ValidationResult:
-        """Validate equal distribution configuration."""
-        result = self._validate_base_requirements(config)
-        
-        # Equal distribution has no additional requirements
-        return result
-    
-    def apply_distribution(self, config: DistributionConfig) -> pl.Series:
-        """Apply equal distribution - uniform selection from base values."""
-        if not self.validate_config(config).is_valid:
-            raise ValidationError("Invalid configuration for equal distribution")
-        
-        # Create uniform distribution using polars
-        return (
-            pl.Series("values", config.base_values)
-            .sample(n=config.target_rows, with_replacement=True, seed=self.seed)
-        )
-
-
-class CustomDistributionEngine(DistributionEngine):
-    """Engine for custom weighted distribution strategy."""
-    
-    def supports_strategy(self, strategy_type: DistributionType) -> bool:
-        """Check if this engine supports custom distribution."""
-        return strategy_type == DistributionType.CUSTOM
-    
-    def validate_config(self, config: DistributionConfig) -> ValidationResult:
-        """Validate custom distribution configuration."""
-        result = self._validate_base_requirements(config)
-        
-        weights = config.strategy.parameters.get('weights', {})
-        if not weights:
-            result.add_error("Custom distribution requires weights", "Add 'weights' parameter with value:percentage pairs")
-            return result
-        
-        # Validate weights format and sum
-        total_weight = 0
-        for value, weight in weights.items():
-            if not isinstance(weight, (int, float)):
-                result.add_error(f"Weight for '{value}' must be numeric", f"Change weight to a number")
-            else:
-                total_weight += weight
-        
-        if abs(total_weight - 100) > 0.01:  # Allow small floating point errors
-            result.add_error(f"Weights must sum to 100%, got {total_weight}%", "Adjust weights to sum to 100%")
-        
-        return result
-    
-    def apply_distribution(self, config: DistributionConfig) -> pl.Series:
-        """Apply custom weighted distribution."""
-        validation = self.validate_config(config)
-        if not validation.is_valid:
-            raise ValidationError(f"Invalid configuration for custom distribution: {validation.errors}")
-        
-        weights = config.strategy.parameters['weights']
-        
-        # Convert percentages to probabilities and create weighted samples
-        values = []
-        probabilities = []
-        
-        for value, weight in weights.items():
-            values.append(value)
-            probabilities.append(weight / 100.0)
-        
-        # Create weighted distribution manually since polars sample doesn't support weights
-        import random
-        if self.seed is not None:
-            random.seed(self.seed)
-        
-        # Generate weighted samples
-        result_values = []
-        for _ in range(config.target_rows):
-            rand_val = random.random()
-            cumulative = 0.0
-            for i, prob in enumerate(probabilities):
-                cumulative += prob
-                if rand_val <= cumulative:
-                    result_values.append(values[i])
-                    break
-        
-        return pl.Series("values", result_values)
-
-
-class CategoricalDistributionEngine(DistributionEngine):
-    """Engine for categorical distribution strategy."""
-    
-    def supports_strategy(self, strategy_type: DistributionType) -> bool:
-        """Check if this engine supports categorical distribution."""
-        return strategy_type == DistributionType.CATEGORICAL
-    
-    def validate_config(self, config: DistributionConfig) -> ValidationResult:
-        """Validate categorical distribution configuration."""
-        result = self._validate_base_requirements(config)
-        
-        categories = config.strategy.parameters.get('categories', [])
-        if not categories:
-            result.add_error("Categorical distribution requires categories", "Add 'categories' parameter with list of allowed values")
-            return result
-        
-        if not isinstance(categories, list):
-            result.add_error("Categories must be a list", "Change categories to a list format")
-        
-        # Check that all categories exist in base values
-        base_set = set(config.base_values)
-        for category in categories:
-            if category not in base_set:
-                result.add_warning(f"Category '{category}' not found in base values")
-        
-        return result
-    
-    def apply_distribution(self, config: DistributionConfig) -> pl.Series:
-        """Apply categorical distribution - select only from specified categories."""
-        validation = self.validate_config(config)
-        if not validation.is_valid:
-            raise ValidationError(f"Invalid configuration for categorical distribution: {validation.errors}")
-        
-        categories = config.strategy.parameters['categories']
-        
-        # Filter base values to only include specified categories
-        available_categories = [val for val in config.base_values if val in categories]
-        
-        if not available_categories:
-            raise SyntheticDataError("No valid categories found in base values")
-        
-        return (
-            pl.Series("values", available_categories)
-            .sample(n=config.target_rows, with_replacement=True, seed=self.seed)
-        )
-
-
-class HighCardinalityDistributionEngine(DistributionEngine):
-    """Engine for high cardinality distribution strategy."""
-    
-    def supports_strategy(self, strategy_type: DistributionType) -> bool:
-        """Check if this engine supports high cardinality distribution."""
-        return strategy_type == DistributionType.HIGH_CARDINALITY
-    
-    def validate_config(self, config: DistributionConfig) -> ValidationResult:
-        """Validate high cardinality distribution configuration."""
-        result = self._validate_base_requirements(config)
-        
-        # Check if we have enough base values for high cardinality
-        if len(config.base_values) < config.target_rows * 0.8:
-            result.add_warning(
-                f"Base values ({len(config.base_values)}) may be insufficient for high cardinality with {config.target_rows} rows"
-            )
-        
-        return result
-    
-    def apply_distribution(self, config: DistributionConfig) -> pl.Series:
-        """Apply high cardinality distribution - mostly unique values with controlled duplicates."""
-        validation = self.validate_config(config)
-        if not validation.is_valid:
-            raise ValidationError(f"Invalid configuration for high cardinality distribution: {validation.errors}")
-        
-        # Strategy: Use most values once, then fill remainder with random selection
-        unique_count = min(len(config.base_values), int(config.target_rows * 0.9))
-        duplicate_count = config.target_rows - unique_count
-        
-        # Create base series for sampling
-        base_series = pl.Series("values", config.base_values)
-        
-        # Get unique values first
-        unique_values = base_series.head(unique_count)
-        
-        # Fill remainder with random selection if needed
-        if duplicate_count > 0:
-            duplicate_values = base_series.sample(n=duplicate_count, with_replacement=True, seed=self.seed)
-            # Both series now have the same dtype since they come from the same base series
-            result = pl.concat([unique_values, duplicate_values])
-        else:
-            result = unique_values
-        
-        # Shuffle the final result - handle case where we have fewer values than requested
-        if result.len() < config.target_rows:
-            # Need to sample with replacement to reach target
-            return result.sample(n=config.target_rows, with_replacement=True, seed=self.seed)
-        else:
-            return result.sample(n=config.target_rows, with_replacement=False, seed=self.seed)
-
-
-class NumericRangeDistributionEngine(DistributionEngine):
-    """Engine for numeric range distribution strategy."""
-    
-    def supports_strategy(self, strategy_type: DistributionType) -> bool:
-        """Check if this engine supports numeric range distribution."""
-        return strategy_type == DistributionType.NUMERIC_RANGE
-    
-    def validate_config(self, config: DistributionConfig) -> ValidationResult:
-        """Validate numeric range distribution configuration."""
-        result = self._validate_base_requirements(config)
-        
-        min_val = config.strategy.parameters.get('min')
-        max_val = config.strategy.parameters.get('max')
-        
-        if min_val is None:
-            result.add_error("Numeric range requires 'min' parameter", "Add 'min' parameter with minimum value")
-        
-        if max_val is None:
-            result.add_error("Numeric range requires 'max' parameter", "Add 'max' parameter with maximum value")
-        
-        if min_val is not None and max_val is not None:
-            if not isinstance(min_val, (int, float)):
-                result.add_error("Min value must be numeric", "Change 'min' to a number")
-            
-            if not isinstance(max_val, (int, float)):
-                result.add_error("Max value must be numeric", "Change 'max' to a number")
-            
-            if isinstance(min_val, (int, float)) and isinstance(max_val, (int, float)) and min_val >= max_val:
-                result.add_error("Min value must be less than max value", "Ensure min < max")
-        
-        return result
-    
-    def apply_distribution(self, config: DistributionConfig) -> pl.Series:
-        """Apply numeric range distribution - filter values within specified range."""
-        validation = self.validate_config(config)
-        if not validation.is_valid:
-            raise ValidationError(f"Invalid configuration for numeric range distribution: {validation.errors}")
-        
-        min_val = config.strategy.parameters['min']
-        max_val = config.strategy.parameters['max']
-        
-        # Filter base values to numeric values within range
-        numeric_values = []
-        for val in config.base_values:
-            try:
-                num_val = float(val)
-                if min_val <= num_val <= max_val:
-                    numeric_values.append(val)
-            except (ValueError, TypeError):
-                continue
-        
-        if not numeric_values:
-            raise SyntheticDataError(f"No numeric values found in range [{min_val}, {max_val}]")
-        
-        return (
-            pl.Series("values", numeric_values)
-            .sample(n=config.target_rows, with_replacement=True, seed=self.seed)
-        )
-
-
-class SkewedDistributionEngine(DistributionEngine):
-    """Engine for skewed distribution strategy (80/20 rule)."""
-    
-    def supports_strategy(self, strategy_type: DistributionType) -> bool:
-        """Check if this engine supports skewed distribution."""
-        return strategy_type == DistributionType.SKEWED
-    
-    def validate_config(self, config: DistributionConfig) -> ValidationResult:
-        """Validate skewed distribution configuration."""
-        result = self._validate_base_requirements(config)
-        
-        # Need at least 2 values for skewed distribution
-        if len(config.base_values) < 2:
-            result.add_error("Skewed distribution requires at least 2 base values", "Provide more base values")
-        
-        return result
-    
-    def apply_distribution(self, config: DistributionConfig) -> pl.Series:
-        """Apply skewed distribution - 80% from 20% of values."""
-        validation = self.validate_config(config)
-        if not validation.is_valid:
-            raise ValidationError(f"Invalid configuration for skewed distribution: {validation.errors}")
-        
-        # Calculate 20% of values (at least 1)
-        popular_count = max(1, len(config.base_values) // 5)
-        popular_values = config.base_values[:popular_count]
-        remaining_values = config.base_values[popular_count:]
-        
-        # 80% of rows from popular values, 20% from remaining
-        popular_rows = int(config.target_rows * 0.8)
-        remaining_rows = config.target_rows - popular_rows
-        
-        # Generate popular values
-        popular_series = (
-            pl.Series("values", popular_values)
-            .sample(n=popular_rows, with_replacement=True, seed=self.seed)
-        )
-        
-        # Generate remaining values
-        if remaining_values and remaining_rows > 0:
-            remaining_series = (
-                pl.Series("values", remaining_values)
-                .sample(n=remaining_rows, with_replacement=True, seed=self.seed)
-            )
-            result = pl.concat([popular_series, remaining_series])
-        else:
-            # If no remaining values, use all from popular
-            result = popular_series
-        
-        # Shuffle the final result
-        return result.sample(n=config.target_rows, with_replacement=False, seed=self.seed)
-
-
-class DistributionEngineFactory:
-    """Factory for creating distribution engines with plugin support."""
-    
-    # Class-level registry for custom engines (shared across all instances)
-    _custom_engines: List[type] = []
-    
-    def __init__(self):
-        """Initialize the factory with all available engines."""
-        self._engines = [
-            EqualDistributionEngine,
-            CustomDistributionEngine,
-            CategoricalDistributionEngine,
-            HighCardinalityDistributionEngine,
-            NumericRangeDistributionEngine,
-            SkewedDistributionEngine,
-        ]
-    
-    @classmethod
-    def register_custom_engine(cls, engine_class: type) -> None:
-        """
-        Register a custom distribution engine.
-        
-        Args:
-            engine_class: A class that inherits from DistributionEngine
-            
-        Raises:
-            ValidationError: If the engine class is invalid
-            
-        Example:
-            >>> class MyCustomEngine(DistributionEngine):
-            ...     def supports_strategy(self, strategy_type):
-            ...         return strategy_type == DistributionType.CUSTOM
-            ...     # ... implement other methods
-            >>> DistributionEngineFactory.register_custom_engine(MyCustomEngine)
-        """
-        # Validate that the class inherits from DistributionEngine
-        if not issubclass(engine_class, DistributionEngine):
-            raise ValidationError(
-                f"Custom engine must inherit from DistributionEngine, got {engine_class.__name__}"
-            )
-        
-        # Validate that required methods are implemented
-        required_methods = ['supports_strategy', 'validate_config', 'apply_distribution']
-        for method_name in required_methods:
-            if not hasattr(engine_class, method_name):
-                raise ValidationError(
-                    f"Custom engine {engine_class.__name__} must implement {method_name} method"
-                )
-        
-        # Check if engine is already registered
-        if engine_class in cls._custom_engines:
-            raise ValidationError(
-                f"Engine {engine_class.__name__} is already registered"
-            )
-        
-        # Add to registry
-        cls._custom_engines.append(engine_class)
-    
-    @classmethod
-    def unregister_custom_engine(cls, engine_class: type) -> None:
-        """
-        Unregister a custom distribution engine.
-        
-        Args:
-            engine_class: The engine class to unregister
-            
-        Raises:
-            ValidationError: If the engine is not registered
-        """
-        if engine_class not in cls._custom_engines:
-            raise ValidationError(
-                f"Engine {engine_class.__name__} is not registered"
-            )
-        cls._custom_engines.remove(engine_class)
-    
-    @classmethod
-    def list_custom_engines(cls) -> List[type]:
-        """Get list of all registered custom engines."""
-        return cls._custom_engines.copy()
-    
-    @classmethod
-    def clear_custom_engines(cls) -> None:
-        """Clear all registered custom engines (useful for testing)."""
-        cls._custom_engines.clear()
-    
-    def create_engine(self, strategy_type: DistributionType, seed: Optional[int] = None) -> DistributionEngine:
-        """
-        Create an appropriate engine for the given strategy type.
-        
-        Custom engines are checked first, then built-in engines.
-        """
-        # Check custom engines first (allows overriding built-in engines)
-        for engine_class in self._custom_engines:
-            engine = engine_class(seed=seed)
-            if engine.supports_strategy(strategy_type):
-                return engine
-        
-        # Check built-in engines
-        for engine_class in self._engines:
-            engine = engine_class(seed=seed)
-            if engine.supports_strategy(strategy_type):
-                return engine
-        
-        raise SyntheticDataError(f"No engine available for distribution strategy: {strategy_type}")
-    
-    def get_supported_strategies(self) -> List[DistributionType]:
-        """Get list of all supported distribution strategies (built-in and custom)."""
-        strategies = []
-        
-        # Check custom engines
-        for engine_class in self._custom_engines:
-            engine = engine_class()
-            for strategy_type in DistributionType:
-                if engine.supports_strategy(strategy_type):
-                    strategies.append(strategy_type)
-        
-        # Check built-in engines
-        for engine_class in self._engines:
-            engine = engine_class()
-            for strategy_type in DistributionType:
-                if engine.supports_strategy(strategy_type):
-                    strategies.append(strategy_type)
-        
-        return list(set(strategies))
-
-
-class DistributionManager:
-    """High-level manager for distribution strategy operations."""
-    
-    def __init__(self, seed: Optional[int] = None):
-        """Initialize the distribution manager."""
-        self.factory = DistributionEngineFactory()
-        self.seed = seed
-    
-    def apply_distribution(self, strategy: DistributionStrategy, base_values: List[str], target_rows: int) -> pl.Series:
-        """Apply a distribution strategy to base values."""
-        config = DistributionConfig(
-            strategy=strategy,
-            base_values=base_values,
-            target_rows=target_rows,
-            seed=self.seed
-        )
-        
-        engine = self.factory.create_engine(strategy.strategy_type, self.seed)
-        return engine.apply_distribution(config)
-    
-    def validate_strategy(self, strategy: DistributionStrategy, base_values: List[str], target_rows: int) -> ValidationResult:
-        """Validate a distribution strategy configuration."""
-        config = DistributionConfig(
-            strategy=strategy,
-            base_values=base_values,
-            target_rows=target_rows,
-            seed=self.seed
-        )
-        
-        try:
-            engine = self.factory.create_engine(strategy.strategy_type, self.seed)
-            return engine.validate_config(config)
-        except SyntheticDataError as e:
-            result = ValidationResult(is_valid=False)
-            result.add_error(str(e))
-            return result
-    
-    def get_supported_strategies(self) -> List[DistributionType]:
-        """Get list of supported distribution strategies."""
-        return self.factory.get_supported_strategies()