feat: add country correlation for multi-country household sampling

- Add country to PARTNER_CORRELATED_ATTRIBUTES (95% same-country rate)
- Add same_country_rate config field to HouseholdConfig
- Auto-inject country attribute for multi-country geographies (world, continents, regions)
- Detect patterns like 'east asia', 'europe', 'global', 'apac', etc.

Author: DeveshParagiri

extropy/core/models/population.py

@@ -106,6 +106,8 @@ class HouseholdConfig(BaseModel):
         }
     )
     default_same_group_rate: float = 0.85
+    # Partner correlation: same-country rate for international populations
+    same_country_rate: float = 0.95
     assortative_mating: dict[str, float] = Field(
         default_factory=lambda: {
             "education_level": 0.6,

extropy/population/sampler/households.py

@@ -29,6 +29,7 @@
 # Attributes correlated between partners (not copied, but biased)
 PARTNER_CORRELATED_ATTRIBUTES = [
     "age",
+    "country",
     "race_ethnicity",
     "education_level",
     "religious_affiliation",
@@ -115,6 +116,15 @@ def correlate_partner_attribute(
         )
         return max(config.min_adult_age, partner_age)
 
+    if attr_name == "country":
+        if rng.random() < config.same_country_rate:
+            return primary_value
+        if available_options:
+            others = [o for o in available_options if o != primary_value]
+            if others:
+                return rng.choice(others)
+        return primary_value
+
     if attr_name == "race_ethnicity":
         same_rate = config.same_group_rates.get(
             str(primary_value).lower(), config.default_same_group_rate

extropy/population/spec_builder/selector.py

@@ -11,6 +11,65 @@
 from ...core.models import AttributeSpec, DiscoveredAttribute
 
 
+# Multi-country geography patterns that should trigger country attribute injection
+MULTI_COUNTRY_PATTERNS = [
+    # Global
+    "world",
+    "global",
+    "globe",
+    "international",
+    "worldwide",
+    # Continents
+    "africa",
+    "asia",
+    "europe",
+    "north america",
+    "south america",
+    "latin america",
+    "oceania",
+    "antarctica",
+    # Regions spanning multiple countries
+    "east asia",
+    "southeast asia",
+    "south asia",
+    "central asia",
+    "middle east",
+    "western europe",
+    "eastern europe",
+    "central europe",
+    "northern europe",
+    "southern europe",
+    "sub-saharan africa",
+    "north africa",
+    "central america",
+    "caribbean",
+    "pacific",
+    "nordic",
+    "scandinavian",
+    "balkan",
+    "mediterranean",
+    "gulf",
+    "apac",
+    "emea",
+    "latam",
+    "amer",
+]
+
+
+def _is_multi_country_geography(geography: str | None, description: str) -> bool:
+    """Detect if the population spans multiple countries."""
+    if not geography and not description:
+        return False
+
+    text = f"{geography or ''} {description}".lower()
+
+    for pattern in MULTI_COUNTRY_PATTERNS:
+        if pattern in text:
+            return True
+
+    return False
+
+
 # JSON schema for attribute selection response
 ATTRIBUTE_SELECTION_SCHEMA = {
     "type": "object",
@@ -334,4 +393,22 @@ def trait_already_exists(trait: str, existing: set[str]) -> bool:
                     )
                 )
 
+    # Inject country attribute for multi-country geographies
+    if _is_multi_country_geography(geography, description):
+        existing_names = {a.name for a in attributes}
+        if "country" not in existing_names:
+            # Insert at the beginning (universal demographic)
+            attributes.insert(
+                0,
+                DiscoveredAttribute(
+                    name="country",
+                    type="categorical",
+                    category="universal",
+                    description="Country of residence",
+                    strategy="independent",
+                    scope="household",
+                    depends_on=[],
+                ),
+            )
+
     return attributes

tests/test_household_sampling.py

@@ -422,3 +422,23 @@ def test_partner_age_correlation(self):
             avg_diff = sum(age_diffs) / len(age_diffs)
             # Average age gap should be small (< 10 years)
             assert avg_diff < 10, f"Average age gap {avg_diff:.1f} too large"
+
+    def test_country_correlation(self):
+        """Country should be correlated with high same-country rate (~95%)."""
+        rng = random.Random(42)
+        config = HouseholdConfig(same_country_rate=0.95)
+        countries = ["USA", "India", "UK", "Japan", "Brazil"]
+
+        same_country = 0
+        total = 1000
+        for _ in range(total):
+            primary_country = rng.choice(countries)
+            partner_country = correlate_partner_attribute(
+                "country", primary_country, rng, config, countries
+            )
+            if partner_country == primary_country:
+                same_country += 1
+
+        rate = same_country / total
+        # Should be close to 0.95 (within statistical margin)
+        assert 0.90 < rate < 0.99, f"Same-country rate {rate:.2%} out of expected range"