refactor: make partner correlation spec-driven instead of hardcoded

DeveshParagiri · DeveshParagiri · commit d45c8f681ae4 · 2026-02-16T17:05:03.000-05:00
- Add scope='partner_correlated' option to AttributeSpec
- Add correlation_rate field for explicit correlation probability
- Remove hardcoded PARTNER_CORRELATED_ATTRIBUTES list
- Update selector prompt to guide LLM on scope selection
- correlate_partner_attribute now takes attr_type and correlation_rate
- NPC partners always get correlated age regardless of scope
diff --git a/extropy/core/models/population.py b/extropy/core/models/population.py
@@ -400,9 +400,13 @@ class AttributeSpec(BaseModel):
         "universal", "population_specific", "context_specific", "personality"
     ] = Field(description="Category of attribute")
     description: str = Field(description="What this attribute represents")
-    scope: Literal["individual", "household"] = Field(
+    scope: Literal["individual", "household", "partner_correlated"] = Field(
         default="individual",
-        description="Whether this attribute is sampled per-individual or shared across a household",
+        description="individual: varies per person; household: shared across household members; partner_correlated: correlated between partners using assortative mating",
+    )
+    correlation_rate: float | None = Field(
+        default=None,
+        description="For partner_correlated scope: probability (0-1) that partner has same value. None uses type-specific defaults (age uses gaussian, race uses per-group rates).",
     )
     sampling: SamplingConfig
     grounding: GroundingInfo
@@ -670,16 +674,17 @@ class DiscoveredAttribute(BaseModel):
         default="independent",
         description="independent: sample directly; derived: zero-variance formula; conditional: probabilistic dependency",
     )
-    scope: Literal["individual", "household"] = Field(
+    scope: Literal["individual", "household", "partner_correlated"] = Field(
         default="individual",
-        description="individual: varies per person; household: shared across household members",
+        description="individual: varies per person; household: shared across household members; partner_correlated: correlated between partners",
+    )
+    correlation_rate: float | None = Field(
+        default=None,
+        description="For partner_correlated scope: probability (0-1) that partner has same value",
     )
     depends_on: list[str] = Field(default_factory=list)
 
 
-# hydrated attribute seems to be an extension of discovered attribute.
-
-
 class HydratedAttribute(BaseModel):
     """An attribute with distribution data from research (Step 2).
 
@@ -696,9 +701,13 @@ class HydratedAttribute(BaseModel):
     strategy: Literal["independent", "derived", "conditional"] = Field(
         default="independent", description="Sampling strategy determined in Step 1"
     )
-    scope: Literal["individual", "household"] = Field(
+    scope: Literal["individual", "household", "partner_correlated"] = Field(
         default="individual",
-        description="individual: varies per person; household: shared across household members",
+        description="individual: varies per person; household: shared across household members; partner_correlated: correlated between partners",
+    )
+    correlation_rate: float | None = Field(
+        default=None,
+        description="For partner_correlated scope: probability (0-1) that partner has same value",
     )
     depends_on: list[str] = Field(default_factory=list)
     sampling: SamplingConfig
diff --git a/extropy/population/sampler/core.py b/extropy/population/sampler/core.py
@@ -33,7 +33,6 @@
     correlate_partner_attribute,
     generate_dependents,
     estimate_household_count,
-    PARTNER_CORRELATED_ATTRIBUTES,
 )
 from .modifiers import apply_modifiers_and_sample
 from ...utils.eval_safe import eval_formula, FormulaError
@@ -214,7 +213,7 @@ def _sample_population_independent(
 
 def _generate_npc_partner(
     primary: dict[str, Any],
-    household_attrs: set[str],
+    attr_map: dict[str, AttributeSpec],
     categorical_options: dict[str, list[str]],
     rng: random.Random,
     config: HouseholdConfig,
@@ -223,37 +222,52 @@ def _generate_npc_partner(
     """Generate a lightweight NPC partner profile for context.
 
     Not a full agent — just enough for persona prompts and conversations.
+    Uses attr.scope from the spec to determine which attributes to include.
     """
     partner: dict[str, Any] = {}
 
-    if "age" in primary:
-        partner["age"] = correlate_partner_attribute("age", primary["age"], rng, config)
+    # Always include gender
     partner["gender"] = rng.choice(["male", "female"])
 
-    for attr in (
-        "race_ethnicity",
-        "education_level",
-        "religious_affiliation",
-        "political_orientation",
-    ):
-        if attr in primary:
-            correlated = correlate_partner_attribute(
-                attr, primary[attr], rng, config, categorical_options.get(attr)
-            )
-            if correlated is not None:
-                partner[attr] = correlated
+    # Always correlate age if present (essential for NPC identity, regardless of scope)
+    if "age" in primary:
+        partner["age"] = correlate_partner_attribute(
+            "age",
+            "int",
+            primary["age"],
+            None,  # Uses gaussian offset
+            rng,
+            config,
+        )
 
-    # Shared household attrs
-    for attr in household_attrs:
-        if attr in primary:
-            partner[attr] = primary[attr]
+    # Process attributes based on their scope
+    for attr_name, attr in attr_map.items():
+        if attr_name not in primary or attr_name == "age":
+            continue
+
+        if attr.scope == "household":
+            # Shared: copy from primary
+            partner[attr_name] = primary[attr_name]
+        elif attr.scope == "partner_correlated":
+            # Correlated: use assortative mating
+            partner[attr_name] = correlate_partner_attribute(
+                attr_name,
+                attr.type,
+                primary[attr_name],
+                attr.correlation_rate,
+                rng,
+                config,
+                available_options=categorical_options.get(attr_name),
+            )
+        # Individual scope: skip for NPC (not enough data to sample fully)
 
     # Generate name for partner
     partner_age = partner.get("age")
     birth_decade = age_to_birth_decade(partner_age) if partner_age is not None else None
+    ethnicity = partner.get("race_ethnicity") or partner.get("ethnicity") or partner.get("race")
     first_name, _ = generate_name(
         gender=partner["gender"],
-        ethnicity=partner.get("race_ethnicity"),
+        ethnicity=str(ethnicity) if ethnicity else None,
         birth_decade=birth_decade,
         seed=rng.randint(0, 2**31),
         name_config=name_config,
@@ -406,7 +420,7 @@ def _sample_population_households(
                 # Partner is NPC context on the primary agent
                 npc_partner = _generate_npc_partner(
                     adult1,
-                    household_attrs,
+                    attr_map,
                     categorical_options,
                     rng,
                     config,
@@ -502,10 +516,10 @@ def _sample_partner_agent(
 ) -> dict[str, Any]:
     """Sample a partner agent with correlated demographics.
 
-    - Household-scoped attributes are copied from the primary.
-    - Correlated attributes (age, race, education, religion, politics)
-      use assortative mating tables.
-    - Everything else is sampled independently.
+    Uses attr.scope from the spec to determine sampling behavior:
+    - scope="household": copy from primary
+    - scope="partner_correlated": use assortative mating correlation
+    - scope="individual": sample independently
     """
     if config is None:
         config = HouseholdConfig()
@@ -517,29 +531,21 @@ def _sample_partner_agent(
             continue
 
         # Household-scoped: copy from primary
-        if attr_name in household_attrs and attr_name in primary:
+        if attr.scope == "household" and attr_name in primary:
             value = primary[attr_name]
-        # Correlated: use partner correlation
-        elif attr_name in PARTNER_CORRELATED_ATTRIBUTES and attr_name in primary:
-            correlated = correlate_partner_attribute(
+        # Partner-correlated: use assortative mating
+        elif attr.scope == "partner_correlated" and attr_name in primary:
+            value = correlate_partner_attribute(
                 attr_name,
+                attr.type,
                 primary[attr_name],
+                attr.correlation_rate,
                 rng,
                 config,
                 available_options=categorical_options.get(attr_name),
             )
-            if correlated is not None:
-                value = correlated
-            else:
-                # Fallback: sample independently
-                try:
-                    value = _sample_attribute(attr, rng, agent, stats)
-                except FormulaError as e:
-                    raise SamplingError(
-                        f"Agent {index}: Failed to sample '{attr_name}': {e}"
-                    ) from e
         else:
-            # Independent sampling
+            # Individual scope: sample independently
             try:
                 value = _sample_attribute(attr, rng, agent, stats)
             except FormulaError as e:
diff --git a/extropy/population/sampler/households.py b/extropy/population/sampler/households.py
@@ -3,6 +3,11 @@
 All household composition rates and correlation tables are read from
 a HouseholdConfig instance (populated by LLM research at spec time,
 with US Census defaults as the safety net).
+
+Attribute scopes are now spec-driven:
+- scope="household": shared across all household members
+- scope="partner_correlated": correlated between partners using assortative mating
+- scope="individual": sampled independently for each person
 """
 
 from __future__ import annotations
@@ -15,32 +20,7 @@
 from ..names.generator import generate_name
 
 if TYPE_CHECKING:
-    from ...core.models.population import NameConfig
-
-
-# Attributes that are always shared within a household
-HOUSEHOLD_SHARED_ATTRIBUTES = [
-    "state",
-    "urban_rural",
-    "household_income",
-    "household_size",
-]
-
-# Attributes correlated between partners (not copied, but biased)
-PARTNER_CORRELATED_ATTRIBUTES = [
-    "age",
-    "country",
-    "race_ethnicity",
-    "education_level",
-    "religious_affiliation",
-    "political_orientation",
-]
-
-# Attributes sampled independently for each partner
-PARTNER_INDEPENDENT_ATTRIBUTES = [
-    "personality",
-    "occupation_sector",
-]
+    from ...core.models.population import AttributeSpec, NameConfig
 
 
 def _age_bracket(age: int, config: HouseholdConfig) -> str:
@@ -92,20 +72,34 @@ def household_needs_kids(htype: HouseholdType) -> bool:
 
 def correlate_partner_attribute(
     attr_name: str,
+    attr_type: str,
     primary_value: Any,
+    correlation_rate: float | None,
     rng: random.Random,
     config: HouseholdConfig,
     available_options: list[str] | None = None,
 ) -> Any:
     """Produce a correlated value for a partner based on the primary's value.
 
-    For categorical attributes, uses assortative mating rates to decide
-    whether to copy or re-sample.  For age, applies a Gaussian offset.
-
-    Returns the correlated value, or None if the attribute isn't in the
-    correlation tables (caller should sample independently).
+    Uses the correlation_rate from the attribute spec. Special handling:
+    - age (int/float): Gaussian offset using config.partner_age_gap_mean/std
+    - race_ethnicity-like attrs: Per-group rates from config.same_group_rates
+    - Other categorical/boolean: Simple probability of same value
+
+    Args:
+        attr_name: Name of the attribute
+        attr_type: Type of the attribute (int, float, categorical, boolean)
+        primary_value: The primary partner's value
+        correlation_rate: Probability (0-1) that partner has same value, or None for defaults
+        rng: Random number generator
+        config: HouseholdConfig with default rates
+        available_options: For categorical attrs, list of valid options to sample from
+
+    Returns:
+        The correlated value for the partner.
     """
-    if attr_name == "age" and isinstance(primary_value, (int, float)):
+    # Age uses gaussian offset, not simple correlation
+    if attr_name == "age" and attr_type in ("int", "float"):
         partner_age = int(
             round(
                 rng.gauss(
@@ -116,16 +110,8 @@ 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":
+    # Race/ethnicity uses per-group rates from config
+    if attr_name in ("race_ethnicity", "ethnicity", "race"):
         same_rate = config.same_group_rates.get(
             str(primary_value).lower(), config.default_same_group_rate
         )
@@ -137,17 +123,29 @@ def correlate_partner_attribute(
                 return rng.choice(others)
         return primary_value
 
-    if attr_name in config.assortative_mating:
-        correlation = config.assortative_mating[attr_name]
-        if rng.random() < correlation:
+    # Country uses same_country_rate from config if no explicit rate
+    if attr_name == "country":
+        rate = correlation_rate if correlation_rate is not None else config.same_country_rate
+        if rng.random() < 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
 
-    return None  # Not a correlated attribute
+    # For all other attributes, use the explicit correlation_rate or a default
+    rate = correlation_rate if correlation_rate is not None else config.default_same_group_rate
+    if rng.random() < rate:
+        return primary_value
+
+    # Sample a different value
+    if available_options:
+        others = [o for o in available_options if o != primary_value]
+        if others:
+            return rng.choice(others)
+
+    return primary_value
 
 
 def generate_dependents(
diff --git a/extropy/population/spec_builder/binder.py b/extropy/population/spec_builder/binder.py
@@ -140,6 +140,7 @@ def bind_constraints(
             category=attr.category,
             description=attr.description,
             scope=attr.scope,
+            correlation_rate=attr.correlation_rate,
             sampling=filtered_sampling,
             grounding=attr.grounding,
             constraints=attr.constraints,
diff --git a/extropy/population/spec_builder/selector.py b/extropy/population/spec_builder/selector.py
diff --git a/tests/test_household_sampling.py b/tests/test_household_sampling.py
