"""Synthetic company generation using faker and numpy distributions.""" from __future__ import annotations import random from typing import Optional import numpy as np from faker import Faker from hr_env.server.company import DEPARTMENT_NAMES, Company, Department from hr_env.server.employee import Employee # Department-specific role templates ROLES = { "Engineering": [ ("Software Engineer", 1), ("Software Engineer", 2), ("Senior Engineer", 3), ("Staff Engineer", 4), ("Principal Engineer", 5), ], "Sales": [ ("Sales Rep", 1), ("Account Executive", 2), ("Senior AE", 3), ("Sales Manager", 4), ("VP Sales", 5), ], "Operations": [ ("Operations Analyst", 1), ("Operations Specialist", 2), ("Operations Manager", 3), ("Senior Ops Manager", 4), ("VP Operations", 5), ], "HR": [ ("HR Coordinator", 1), ("HR Specialist", 2), ("HR Manager", 3), ("Senior HR Manager", 4), ("VP People", 5), ], "Finance": [ ("Financial Analyst", 1), ("Senior Analyst", 2), ("Finance Manager", 3), ("Controller", 4), ("CFO", 5), ], } # Salary base by level (lognormal parameters) SALARY_BASE = {1: 55000, 2: 75000, 3: 100000, 4: 135000, 5: 180000} SALARY_SIGMA = 0.15 # lognormal spread # Department size distribution (fraction of total) DEPT_SIZE_FRACTIONS = { "Engineering": 0.35, "Sales": 0.25, "Operations": 0.20, "HR": 0.08, "Finance": 0.12, } # Skills by department DEPT_SKILLS = { "Engineering": ["python", "java", "cloud", "ml", "devops", "frontend", "backend", "data"], "Sales": ["negotiation", "crm", "pipeline", "presentation", "territory", "closing"], "Operations": ["logistics", "process", "lean", "supply_chain", "quality", "planning"], "HR": ["recruiting", "compliance", "training", "benefits", "employee_relations"], "Finance": ["accounting", "budgeting", "forecasting", "audit", "tax", "reporting"], } def generate_company( seed: int = 42, size: int = 300, name: str = "Simulated Corp", base_revenue: Optional[float] = None, hr_budget: Optional[float] = None, ) -> Company: """Generate a synthetic company with realistic employee distributions. Args: seed: Random seed for reproducibility. size: Total number of employees (200-500 recommended). name: Company name. base_revenue: Annual baseline revenue. Defaults to size * 150_000. hr_budget: Annual HR budget. Defaults to size * 6000. """ rng = np.random.default_rng(seed) fake = Faker() Faker.seed(seed) if base_revenue is None: base_revenue = size * 150_000 if hr_budget is None: hr_budget = size * 6000 company = Company( name=name, base_revenue=base_revenue, hr_budget=hr_budget, hr_budget_remaining=hr_budget / 4, # First quarter allocation rng=random.Random(seed), # seeded stream for deterministic turnover ) for dept_name in DEPARTMENT_NAMES: dept_size = max(5, int(size * DEPT_SIZE_FRACTIONS[dept_name])) dept = Department(name=dept_name) roles = ROLES[dept_name] skills_pool = DEPT_SKILLS[dept_name] for _ in range(dept_size): # Level distribution: pyramid shape level_probs = np.array([0.35, 0.30, 0.20, 0.10, 0.05]) level = int(rng.choice([1, 2, 3, 4, 5], p=level_probs)) # Find matching role role_name = next((r for r, l in roles if l == level), roles[0][0]) # Salary: lognormal around base for level base = SALARY_BASE[level] salary = float(rng.lognormal(np.log(base), SALARY_SIGMA)) salary = round(max(35000, salary), -2) # Round to nearest 100 # Performance: normal(3.2, 0.8) clipped [1, 5] perf = float(np.clip(rng.normal(3.2, 0.8), 1.0, 5.0)) # Tenure: exponential(lambda=24 months) tenure = int(rng.exponential(24)) # Engagement: beta(7, 3) * 100 engagement = float(rng.beta(7, 3) * 100) # Skills: 2-4 random from department pool n_skills = min(len(skills_pool), int(rng.integers(2, 5))) skills = list(rng.choice(skills_pool, size=n_skills, replace=False)) emp = Employee( id=f"{dept_name[:3].lower()}_{fake.unique.random_int(min=1000, max=9999)}", name=fake.name(), department=dept_name, role=role_name, level=level, salary=salary, tenure_months=tenure, performance_score=perf, engagement=engagement, skills=skills, training_hours=float(rng.exponential(10)), ) # Derived attributes emp.promotability = float(np.clip( 0.1 + perf * 0.12 + (tenure / 60) * 0.1 + rng.normal(0, 0.1), 0, 1 )) emp.transferability = float(np.clip( 0.3 + len(skills) * 0.08 + rng.normal(0, 0.1), 0, 1 )) emp.retainability = float(np.clip( 0.4 + engagement / 200 + (tenure / 48) * 0.1 + rng.normal(0, 0.1), 0, 1 )) emp.update_flight_risk() dept.employees.append(emp) company.departments[dept_name] = dept return company