misata 0.2.0b0__py3-none-any.whl → 0.3.1b0__py3-none-any.whl

misata/profiles.py

@@ -0,0 +1,332 @@
+"""
+Distribution Profiles for Realistic Data Generation.
+
+Pre-configured distribution parameters that match real-world patterns
+for common data types like age, salary, prices, etc.
+"""
+
+from typing import Any, Dict, List, Optional, Union
+import numpy as np
+
+
+class DistributionProfile:
+    """A named distribution configuration for realistic generation.
+    
+    Example:
+        profile = DistributionProfile(
+            name="age",
+            distribution="mixture",
+            params={
+                "components": [
+                    {"mean": 35, "std": 12, "weight": 0.6},  # Working age
+                    {"mean": 70, "std": 8, "weight": 0.2},   # Retirees
+                    {"mean": 12, "std": 4, "weight": 0.2},   # Children
+                ]
+            }
+        )
+        values = profile.generate(1000)
+    """
+    
+    def __init__(
+        self,
+        name: str,
+        distribution: str,
+        params: Dict[str, Any],
+        min_value: Optional[float] = None,
+        max_value: Optional[float] = None,
+        decimals: Optional[int] = None,
+    ):
+        self.name = name
+        self.distribution = distribution
+        self.params = params
+        self.min_value = min_value
+        self.max_value = max_value
+        self.decimals = decimals
+    
+    def generate(
+        self, 
+        size: int, 
+        rng: Optional[np.random.Generator] = None
+    ) -> np.ndarray:
+        """Generate values according to this profile."""
+        if rng is None:
+            rng = np.random.default_rng()
+        
+        if self.distribution == "normal":
+            mean = self.params.get("mean", 50)
+            std = self.params.get("std", 10)
+            values = rng.normal(mean, std, size)
+        
+        elif self.distribution == "lognormal":
+            mean = self.params.get("mean", 0)
+            sigma = self.params.get("sigma", 1)
+            values = rng.lognormal(mean, sigma, size)
+        
+        elif self.distribution == "exponential":
+            scale = self.params.get("scale", 1.0)
+            values = rng.exponential(scale, size)
+        
+        elif self.distribution == "pareto":
+            alpha = self.params.get("alpha", 2.0)
+            min_val = self.params.get("min", 1.0)
+            values = (rng.pareto(alpha, size) + 1) * min_val
+        
+        elif self.distribution == "beta":
+            a = self.params.get("a", 2)
+            b = self.params.get("b", 5)
+            scale = self.params.get("scale", 1.0)
+            values = rng.beta(a, b, size) * scale
+        
+        elif self.distribution == "mixture":
+            # Gaussian mixture model
+            components = self.params.get("components", [])
+            if not components:
+                values = rng.normal(0, 1, size)
+            else:
+                weights = np.array([c.get("weight", 1) for c in components])
+                weights = weights / weights.sum()
+                
+                # Sample component indices
+                component_indices = rng.choice(
+                    len(components), size=size, p=weights
+                )
+                
+                values = np.zeros(size)
+                for i, comp in enumerate(components):
+                    mask = component_indices == i
+                    n = mask.sum()
+                    if n > 0:
+                        values[mask] = rng.normal(
+                            comp.get("mean", 0),
+                            comp.get("std", 1),
+                            n
+                        )
+        
+        elif self.distribution == "zipf":
+            # Zipf distribution for long-tail data
+            a = self.params.get("alpha", 2.0)
+            values = rng.zipf(a, size).astype(float)
+        
+        elif self.distribution == "uniform":
+            low = self.params.get("min", 0)
+            high = self.params.get("max", 100)
+            values = rng.uniform(low, high, size)
+        
+        else:
+            # Default to uniform
+            values = rng.uniform(0, 100, size)
+        
+        # Apply constraints
+        if self.min_value is not None:
+            values = np.maximum(values, self.min_value)
+        if self.max_value is not None:
+            values = np.minimum(values, self.max_value)
+        if self.decimals is not None:
+            values = np.round(values, self.decimals)
+        
+        return values
+
+
+# ============ Pre-built Profiles ============
+
+PROFILES: Dict[str, DistributionProfile] = {}
+
+
+def _register_profile(profile: DistributionProfile) -> None:
+    """Register a profile by name."""
+    PROFILES[profile.name] = profile
+
+
+# Age distributions
+_register_profile(DistributionProfile(
+    name="age_adult",
+    distribution="mixture",
+    params={
+        "components": [
+            {"mean": 28, "std": 6, "weight": 0.3},   # Young adults
+            {"mean": 42, "std": 10, "weight": 0.45}, # Middle age
+            {"mean": 62, "std": 8, "weight": 0.25},  # Older adults
+        ]
+    },
+    min_value=18,
+    max_value=100,
+    decimals=0,
+))
+
+_register_profile(DistributionProfile(
+    name="age_population",
+    distribution="mixture",
+    params={
+        "components": [
+            {"mean": 8, "std": 4, "weight": 0.15},   # Children
+            {"mean": 25, "std": 8, "weight": 0.25},  # Young adults
+            {"mean": 42, "std": 12, "weight": 0.35}, # Middle age
+            {"mean": 68, "std": 10, "weight": 0.25}, # Seniors
+        ]
+    },
+    min_value=0,
+    max_value=105,
+    decimals=0,
+))
+
+# Salary distributions
+_register_profile(DistributionProfile(
+    name="salary_usd",
+    distribution="lognormal",
+    params={"mean": 11.0, "sigma": 0.5},  # Log of ~$60k median
+    min_value=25000,
+    max_value=500000,
+    decimals=0,
+))
+
+_register_profile(DistributionProfile(
+    name="salary_tech",
+    distribution="mixture",
+    params={
+        "components": [
+            {"mean": 75000, "std": 15000, "weight": 0.2},   # Junior
+            {"mean": 120000, "std": 25000, "weight": 0.4},  # Mid
+            {"mean": 180000, "std": 40000, "weight": 0.3},  # Senior
+            {"mean": 280000, "std": 60000, "weight": 0.1},  # Staff+
+        ]
+    },
+    min_value=50000,
+    max_value=600000,
+    decimals=0,
+))
+
+# Price distributions
+_register_profile(DistributionProfile(
+    name="price_retail",
+    distribution="lognormal",
+    params={"mean": 3.5, "sigma": 1.2},  # ~$30 median
+    min_value=0.99,
+    max_value=10000,
+    decimals=2,
+))
+
+_register_profile(DistributionProfile(
+    name="price_saas",
+    distribution="mixture",
+    params={
+        "components": [
+            {"mean": 15, "std": 5, "weight": 0.3},     # Basic tier
+            {"mean": 49, "std": 15, "weight": 0.4},    # Pro tier
+            {"mean": 199, "std": 50, "weight": 0.25},  # Enterprise
+            {"mean": 999, "std": 200, "weight": 0.05}, # Custom
+        ]
+    },
+    min_value=0,
+    max_value=5000,
+    decimals=0,
+))
+
+# Transaction amounts
+_register_profile(DistributionProfile(
+    name="transaction_amount",
+    distribution="pareto",
+    params={"alpha": 2.5, "min": 10},
+    min_value=1,
+    max_value=100000,
+    decimals=2,
+))
+
+# Counts / quantities
+_register_profile(DistributionProfile(
+    name="order_quantity",
+    distribution="zipf",
+    params={"alpha": 2.0},
+    min_value=1,
+    max_value=100,
+    decimals=0,
+))
+
+# Time-related
+_register_profile(DistributionProfile(
+    name="session_duration_seconds",
+    distribution="lognormal",
+    params={"mean": 5.5, "sigma": 1.5},  # ~4 min median
+    min_value=1,
+    max_value=7200,  # 2 hours max
+    decimals=0,
+))
+
+# Ratings and scores
+_register_profile(DistributionProfile(
+    name="rating_5star",
+    distribution="beta",
+    params={"a": 5, "b": 2, "scale": 5},  # Skewed towards higher ratings
+    min_value=1,
+    max_value=5,
+    decimals=1,
+))
+
+_register_profile(DistributionProfile(
+    name="nps_score",
+    distribution="mixture",
+    params={
+        "components": [
+            {"mean": 3, "std": 2, "weight": 0.15},   # Detractors
+            {"mean": 7, "std": 1, "weight": 0.25},   # Passives
+            {"mean": 9, "std": 0.8, "weight": 0.6},  # Promoters
+        ]
+    },
+    min_value=0,
+    max_value=10,
+    decimals=0,
+))
+
+# Percentages
+_register_profile(DistributionProfile(
+    name="conversion_rate",
+    distribution="beta",
+    params={"a": 2, "b": 50, "scale": 100},  # Low conversion (1-5%)
+    min_value=0,
+    max_value=100,
+    decimals=2,
+))
+
+_register_profile(DistributionProfile(
+    name="churn_rate",
+    distribution="beta",
+    params={"a": 1.5, "b": 30, "scale": 100},  # ~5% typical
+    min_value=0,
+    max_value=100,
+    decimals=2,
+))
+
+
+def get_profile(name: str) -> Optional[DistributionProfile]:
+    """Get a profile by name."""
+    return PROFILES.get(name)
+
+
+def list_profiles() -> List[str]:
+    """List all available profile names."""
+    return list(PROFILES.keys())
+
+
+def generate_with_profile(
+    profile_name: str,
+    size: int,
+    rng: Optional[np.random.Generator] = None
+) -> np.ndarray:
+    """Generate values using a named profile.
+    
+    Args:
+        profile_name: Name of the profile (e.g., "salary_tech")
+        size: Number of values to generate
+        rng: Random number generator
+        
+    Returns:
+        Array of generated values
+        
+    Raises:
+        ValueError: If profile not found
+    """
+    profile = get_profile(profile_name)
+    if profile is None:
+        available = ", ".join(list_profiles())
+        raise ValueError(f"Unknown profile: {profile_name}. Available: {available}")
+    
+    return profile.generate(size, rng)

misata/simulator.py

@@ -16,7 +16,9 @@ from typing import Any, Dict, List, Optional
 import numpy as np
 import pandas as pd
 
-from misata.generators import TextGenerator
+from misata.generators.base import TextGenerator as _FactoryTextGenerator  # Generator factory version
+# Use the original generators.py TextGenerator which supports seed
+from misata.generators_legacy import TextGenerator
 from misata.schema import Column, Relationship, ScenarioEvent, SchemaConfig
 
 
@@ -34,6 +36,10 @@ class DataSimulator:
         rng: NumPy random generator for reproducibility
     """
 
+    # Performance constants
+    MAX_CONTEXT_ROWS = 50000  # Cap context storage for memory efficiency
+    TEXT_POOL_SIZE = 10000    # Size of text value pools for vectorized sampling
+
     def __init__(self, config: SchemaConfig,
                  apply_semantic_fixes: bool = True, batch_size: int = 10_000,
                  smart_mode: bool = False, use_llm: bool = True):
@@ -57,6 +63,7 @@ class DataSimulator:
         self._unique_pools: Dict[str, np.ndarray] = {}  # Store pre-generated unique values
         self._unique_counters: Dict[str, int] = {}      # Track usage of unique pools
         self._smart_pools: Dict[str, np.ndarray] = {}   # Cache smart value pools
+        self._text_pools: Dict[str, np.ndarray] = {}    # Cache text pools for vectorized sampling
 
         # Apply semantic inference to fix column types
         if apply_semantic_fixes:
@@ -199,10 +206,24 @@ class DataSimulator:
         ctx_df = df[cols_to_store].copy()
 
         if table_name not in self.context:
+            # First batch: store up to MAX_CONTEXT_ROWS
+            if len(ctx_df) > self.MAX_CONTEXT_ROWS:
+                ctx_df = ctx_df.sample(n=self.MAX_CONTEXT_ROWS, random_state=self.config.seed)
             self.context[table_name] = ctx_df
         else:
-            # Append to existing context
-            self.context[table_name] = pd.concat([self.context[table_name], ctx_df], ignore_index=True)
+            # Append to existing context, but cap at MAX_CONTEXT_ROWS
+            current_len = len(self.context[table_name])
+            if current_len >= self.MAX_CONTEXT_ROWS:
+                # Already at capacity, use reservoir sampling for randomness
+                # Replace some existing rows with new ones (probability-based)
+                return  # Skip appending, we have enough IDs
+            
+            remaining_space = self.MAX_CONTEXT_ROWS - current_len
+            rows_to_add = ctx_df.iloc[:remaining_space]
+            self.context[table_name] = pd.concat(
+                [self.context[table_name], rows_to_add], 
+                ignore_index=True
+            )
 
     def generate_column(
         self,
@@ -225,6 +246,70 @@ class DataSimulator:
         """
         params = column.distribution_params
 
+        # ========== CORRELATED COLUMN GENERATION ==========
+        # If this column depends on another column's value, use conditional distribution
+        if "depends_on" in params and table_data is not None:
+            parent_col = params["depends_on"]
+            mapping = params.get("mapping", {})
+            
+            if parent_col in table_data.columns and mapping:
+                parent_values = table_data[parent_col].values
+                
+                # Check if it's numeric or categorical mapping
+                first_val = next(iter(mapping.values()))
+                if isinstance(first_val, dict) and "mean" in first_val:
+                    # Numeric conditional distribution (e.g., salary based on job_title)
+                    # mapping = {"Intern": {"mean": 40000, "std": 5000}, "CTO": {"mean": 200000, "std": 30000}}
+                    values = np.zeros(size)
+                    for key, dist in mapping.items():
+                        mask = parent_values == key
+                        count = mask.sum()
+                        if count > 0:
+                            mean = dist.get("mean", 50000)
+                            std = dist.get("std", mean * 0.1)
+                            values[mask] = self.rng.normal(mean, std, count)
+                    
+                    # Handle values that didn't match any key (use default)
+                    default = params.get("default", {"mean": 50000, "std": 10000})
+                    unmatched = ~np.isin(parent_values, list(mapping.keys()))
+                    if unmatched.sum() > 0:
+                        values[unmatched] = self.rng.normal(
+                            default.get("mean", 50000), 
+                            default.get("std", 10000), 
+                            unmatched.sum()
+                        )
+                    return values
+                    
+                elif isinstance(first_val, list):
+                    # Categorical conditional (e.g., state based on country)
+                    # mapping = {"USA": ["CA", "TX", "NY"], "UK": ["England", "Scotland"]}
+                    values = np.empty(size, dtype=object)
+                    for key, choices in mapping.items():
+                        mask = parent_values == key
+                        count = mask.sum()
+                        if count > 0:
+                            values[mask] = self.rng.choice(choices, count)
+                    
+                    # Default for unmatched
+                    default_choices = params.get("default", ["Unknown"])
+                    unmatched = values == None  # noqa
+                    if unmatched.sum() > 0:
+                        values[unmatched] = self.rng.choice(default_choices, unmatched.sum())
+                    return values
+                    
+                elif isinstance(first_val, (int, float)):
+                    # Probability-based boolean (e.g., churn probability based on plan)
+                    # mapping = {"free": 0.3, "pro": 0.1, "enterprise": 0.05}
+                    values = np.zeros(size, dtype=bool)
+                    for key, prob in mapping.items():
+                        mask = parent_values == key
+                        count = mask.sum()
+                        if count > 0:
+                            values[mask] = self.rng.random(count) < prob
+                    return values
+
+        # ========== STANDARD COLUMN GENERATION ==========
+        
         # CATEGORICAL
         if column.type == "categorical":
             choices = params.get("choices", ["A", "B", "C"])
@@ -469,23 +554,59 @@ class DataSimulator:
                         return values
 
             if text_type == "name":
-                values = np.array([self.text_gen.name() for _ in range(size)])
+                pool_key = "text_name"
+                if pool_key not in self._text_pools:
+                    pool_size = min(size, self.TEXT_POOL_SIZE)
+                    self._text_pools[pool_key] = np.array([self.text_gen.name() for _ in range(pool_size)])
+                values = self.rng.choice(self._text_pools[pool_key], size=size)
             elif text_type == "email":
-                values = np.array([self.text_gen.email() for _ in range(size)])
+                pool_key = "text_email"
+                if pool_key not in self._text_pools:
+                    pool_size = min(size, self.TEXT_POOL_SIZE)
+                    self._text_pools[pool_key] = np.array([self.text_gen.email() for _ in range(pool_size)])
+                values = self.rng.choice(self._text_pools[pool_key], size=size)
             elif text_type == "company":
-                values = np.array([self.text_gen.company() for _ in range(size)])
+                pool_key = "text_company"
+                if pool_key not in self._text_pools:
+                    pool_size = min(size, self.TEXT_POOL_SIZE)
+                    self._text_pools[pool_key] = np.array([self.text_gen.company() for _ in range(pool_size)])
+                values = self.rng.choice(self._text_pools[pool_key], size=size)
             elif text_type == "sentence":
-                values = np.array([self.text_gen.sentence() for _ in range(size)])
+                pool_key = "text_sentence"
+                if pool_key not in self._text_pools:
+                    pool_size = min(size, self.TEXT_POOL_SIZE)
+                    self._text_pools[pool_key] = np.array([self.text_gen.sentence() for _ in range(pool_size)])
+                values = self.rng.choice(self._text_pools[pool_key], size=size)
             elif text_type == "word":
-                values = np.array([self.text_gen.word() for _ in range(size)])
+                pool_key = "text_word"
+                if pool_key not in self._text_pools:
+                    pool_size = min(size, self.TEXT_POOL_SIZE)
+                    self._text_pools[pool_key] = np.array([self.text_gen.word() for _ in range(pool_size)])
+                values = self.rng.choice(self._text_pools[pool_key], size=size)
             elif text_type == "address":
-                values = np.array([self.text_gen.full_address() for _ in range(size)])
+                pool_key = "text_address"
+                if pool_key not in self._text_pools:
+                    pool_size = min(size, self.TEXT_POOL_SIZE)
+                    self._text_pools[pool_key] = np.array([self.text_gen.full_address() for _ in range(pool_size)])
+                values = self.rng.choice(self._text_pools[pool_key], size=size)
             elif text_type == "phone":
-                values = np.array([self.text_gen.phone_number() for _ in range(size)])
+                pool_key = "text_phone"
+                if pool_key not in self._text_pools:
+                    pool_size = min(size, self.TEXT_POOL_SIZE)
+                    self._text_pools[pool_key] = np.array([self.text_gen.phone_number() for _ in range(pool_size)])
+                values = self.rng.choice(self._text_pools[pool_key], size=size)
             elif text_type == "url":
-                values = np.array([self.text_gen.url() for _ in range(size)])
+                pool_key = "text_url"
+                if pool_key not in self._text_pools:
+                    pool_size = min(size, self.TEXT_POOL_SIZE)
+                    self._text_pools[pool_key] = np.array([self.text_gen.url() for _ in range(pool_size)])
+                values = self.rng.choice(self._text_pools[pool_key], size=size)
             else:
-                values = np.array([self.text_gen.sentence() for _ in range(size)])
+                pool_key = "text_sentence"
+                if pool_key not in self._text_pools:
+                    pool_size = min(size, self.TEXT_POOL_SIZE)
+                    self._text_pools[pool_key] = np.array([self.text_gen.sentence() for _ in range(pool_size)])
+                values = self.rng.choice(self._text_pools[pool_key], size=size)
 
             return values