Repository: mattwigway/eqsormo

Issue:
endogenous attributes
Endogenous attributes should be supported. This doesn't matter in estimation, but in sorting, after each iteration, these endogenous attributes need to be updated.

Discussion:

Patch:
diff --git a/setup.cfg b/setup.cfg
index 79a16af..64f7da2 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -1,2 +1,4 @@
 [flake8]
-max-line-length = 120
\ No newline at end of file
+max-line-length = 120
+# disable errors about whitespace before func def parens
+ignore = E211
\ No newline at end of file
diff --git a/src/eqsormo/tra/clear_market.py b/src/eqsormo/tra/clear_market.py
index 80a69bd..f49e4f2 100644
--- a/src/eqsormo/tra/clear_market.py
+++ b/src/eqsormo/tra/clear_market.py
@@ -1,4 +1,4 @@
-#    Copyright 2019 Matthew Wigginton Conway
+#    Copyright 2019-2020 Matthew Wigginton Conway
 
 #    Licensed under the Apache License, Version 2.0 (the "License");
 #    you may not use this file except in compliance with the License.
@@ -12,21 +12,28 @@
 #    See the License for the specific language governing permissions and
 #    limitations under the License.
 
-# Author: Matthew Wigginton Conway <matt@indicatrix.org>, School of Geographical Sciences and Urban Planning, Arizona State University
+# Author: Matthew Wigginton Conway <matt@indicatrix.org>
+#         School of Geographical Sciences and Urban Planning,
+#         Arizona State University
 
-import numba
 import numpy as np
-from time import time
 
 from logging import getLogger
 LOG = getLogger(__name__)
 
-def clear_market (non_price_utilities, hhidx, choiceidx, supply, income, starting_price, price_income_transformation,
-        price_income_params, budget_coef, max_rent_to_income=None, convergence_criterion=1e-4, step=1e-2, maxiter=np.inf, weights=None):
+
+def clear_market_iter (non_price_utilities, hhidx, choiceidx, supply, income, starting_price,
+                       price_income_transformation, price_income_params, budget_coef,
+                       max_rent_to_income=None, convergence_criterion=1e-4, step=1e-2,
+                       weights=None):
     '''
-    Clear the market
+    Run one iteration of the market-clearing algorithm.
 
     income and price should be by household index/choice index, respectively.
+
+    returns price after next iteration, and flag for whether prices have converged. Note that once
+    prices converge, one more call is necessary to confirm convergence, to avoid additional computation
+    of market shares to check convergence after prices are updated.
     '''
     price = starting_price
 
@@ -37,99 +44,105 @@ def clear_market (non_price_utilities, hhidx, choiceidx, supply, income, startin
 
     prev_price = np.copy(price)
 
-    itr = 0
-    while True:
-        itr += 1
-        if itr > maxiter:
-            LOG.error(f'Prices FAILED TO CONVERGE in {maxiter} iterations!')
-            return price
-
-        if max_rent_to_income is not None:
-            # could cause oscillation if price keeps going above max income
-            # Also, this will affect sorting equilibrium, because the price may effectively get fixed to the
-            # 90th percentile income because it always moves off, and everything else equilibrate around it
-            if np.any(price > np.max(income) * max_rent_to_income):
-                LOG.error('Some prices have exceeded max income - setting to be affordable to 90th percentile household to keep process going.')
-                price[price > np.max(income) * max_rent_to_income] = np.percentile(income, 0.9) * max_rent_to_income
-
-        shares = compute_shares(
-            price=price,
-            supply=supply,
-            alt_income=alt_income,
-            choiceidx=choiceidx,
-            hhidx=hhidx,
-            non_price_utilities=non_price_utilities,
-            price_income_transformation=price_income_transformation,
-            price_income_params=price_income_params,
-            budget_coef=budget_coef,
-            max_rent_to_income=max_rent_to_income,
-            weights=weights
-        )
-
-        if np.any(shares == 0):
-            raise ValueError('Some shares are zero.')
-
-        excess_demand = shares - supply
-
-        #maxdiff = np.max(np.abs(shares - supply) / supply)
-        maxdiff = np.max(np.abs(excess_demand / supply))
-        if maxdiff < convergence_criterion:
-            LOG.info(f'Prices converged after {itr} iterations')
-            return price
-
-        # probably will need to remove if using numba
-        maxpricediff = np.max(price - prev_price)
-        minpricediff = np.min(price - prev_price)
-        prev_price[:] = price
-        if max_rent_to_income is not None:
-            deltaNExcluded = np.sum(alt_income * max_rent_to_income <= price[choiceidx]) - origNExcluded
-        else:
-            deltaNExcluded = 'n/a'
-
-        # Use the approach defined in Tra (2007), page 108, eq. 7.7/7.7a, which is copied from Anas (1982)
-        # first, compute derivative. Since the budget transformation is an arbitrary Python function, first compute its derivative.
-        # since the budgets are independent between houses and between choosers, this is a fast numpy vectorized operation. We need a
-        # loop to compute derivatives of budget. That can be done as a numpy vectorized operation
-        alt_price = price[choiceidx]
-
-        budget = np.zeros_like(alt_income)
-        if max_rent_to_income is None:
-            feasible_alts = np.full_like(alt_price, True)
-        else:
-            feasible_alts = alt_income * max_rent_to_income > alt_price
-
-        budget = np.full_like(alt_income, np.nan)
-        budget[feasible_alts] = price_income_transformation.apply(alt_income[feasible_alts], alt_price[feasible_alts], *price_income_params)
-        price_step = 1e-8
-        budget_step = np.full_like(alt_income, np.nan)
-
-        if max_rent_to_income is None:
-            feasible_alts_step = np.full_like(alt_price, True)
-        else:
-            feasible_alts_step = alt_income * max_rent_to_income > alt_price
-        budget_step[feasible_alts] = price_income_transformation.apply(alt_income[feasible_alts_step], alt_price[feasible_alts_step] + price_step, *price_income_params)
-
-        # fix one price from changing
-        # TODO pick price in a smarter way (PUMA with least change?)
-        # Appears to be causing convergence problems.
-        #excess_demand[0] = 0
-
-        # NB derivatives are ill-behaved at the boundary where a choice enters a household's choice set due to the price dropping below their
-        # income. We assume that what is in each household's choice set is nonchanging for the purpose of calculating derivatives.
-        # Nope, that won't work.
-        deriv = compute_derivatives(price, alt_income, choiceidx, hhidx, non_price_utilities, budget, budget_step, price_step, budget_coef,
-            shares, max_rent_to_income, feasible_alts, feasible_alts_step, weights)
+    if max_rent_to_income is not None:
+        # could cause oscillation if price keeps going above max income
+        # Also, this will affect sorting equilibrium, because the price may effectively get fixed to the
+        # 90th percentile income because it always moves off, and everything else equilibrate around it
+        if np.any(price > np.max(income) * max_rent_to_income):
+            LOG.error('Some prices have exceeded max income - setting to be affordable to 90th'
+                        + 'percentile household to keep process going.')
+            price[price > np.max(income) * max_rent_to_income] = np.percentile(income, 0.9) * max_rent_to_income
+
+    shares = compute_shares(
+        price=price,
+        supply=supply,
+        alt_income=alt_income,
+        choiceidx=choiceidx,
+        hhidx=hhidx,
+        non_price_utilities=non_price_utilities,
+        price_income_transformation=price_income_transformation,
+        price_income_params=price_income_params,
+        budget_coef=budget_coef,
+        max_rent_to_income=max_rent_to_income,
+        weights=weights
+    )
+
+    if np.any(shares == 0):
+        raise ValueError('Some shares are zero.')
+
+    excess_demand = shares - supply
+
+    # maxdiff = np.max(np.abs(shares - supply) / supply)
+    maxdiff = np.max(np.abs(excess_demand / supply))
+    if maxdiff < convergence_criterion:
+        return price, False
+
+    # probably will need to remove if using numba
+    maxpricediff = np.max(price - prev_price)
+    minpricediff = np.min(price - prev_price)
+    prev_price[:] = price
+    if max_rent_to_income is not None:
+        deltaNExcluded = np.sum(alt_income * max_rent_to_income <= price[choiceidx]) - origNExcluded
+    else:
+        deltaNExcluded = 'n/a'
+
+    # Use the approach defined in Tra (2007), page 108, eq. 7.7/7.7a, which is copied from Anas (1982)
+    # first, compute derivative. Since the budget transformation is an arbitrary Python function,
+    # first compute its derivative.
+    # since the budgets are independent between houses and between choosers, this is a fast numpy
+    # vectorized operation. We need a loop to compute derivatives of budget. That can be done as a
+    # numpy vectorized operation. TODO this doesn't make sense - maybe using vectorize to mean diff
+    # things, i.e. actually vectorized vs np.vectorize?
+    alt_price = price[choiceidx]
+
+    budget = np.zeros_like(alt_income)
+    if max_rent_to_income is None:
+        feasible_alts = np.full_like(alt_price, True)
+    else:
+        feasible_alts = alt_income * max_rent_to_income > alt_price
+
+    budget = np.full_like(alt_income, np.nan)
+    budget[feasible_alts] =\
+        price_income_transformation.apply(alt_income[feasible_alts], alt_price[feasible_alts], *price_income_params)
+    price_step = 1e-8
+    budget_step = np.full_like(alt_income, np.nan)
+
+    if max_rent_to_income is None:
+        feasible_alts_step = np.full_like(alt_price, True)
+    else:
+        feasible_alts_step = alt_income * max_rent_to_income > alt_price
+    budget_step[feasible_alts] =\
+        price_income_transformation.apply(alt_income[feasible_alts_step],
+                                          alt_price[feasible_alts_step] + price_step, *price_income_params)
+
+    # fix one price from changing
+    # TODO pick price in a smarter way (PUMA with least change?)
+    # Appears to be causing convergence problems.
+    # excess_demand[0] = 0
+
+    # NB derivatives are ill-behaved at the boundary where a choice enters a household's choice set due to
+    # price dropping below their income. We assume that what is in each household's choice set is
+    # nonchanging for the purpose of calculating derivatives.
+    # Nope, that won't work.
+    # TODO is the above comment correct? how is this handled now?
+    deriv = compute_derivatives(price, alt_income, choiceidx, hhidx, non_price_utilities, budget, budget_step,
+                                price_step, budget_coef, shares, max_rent_to_income, feasible_alts,
+                                feasible_alts_step, weights)
+
+    if not np.all(deriv < 0):
+        raise ValueError('some derivatives of price are nonnegative')
+
+    # this is 7.7a from Tra's dissertation
+    price = price - (excess_demand / deriv)
+    LOG.info(f'Max unit diff: {maxdiff}, max price diff: {maxpricediff}, '
+             + f'min price diff: {minpricediff}, additional excluded due to rent/income ration {deltaNExcluded}')
     
-        if not np.all(deriv < 0):
-            raise ValueError('some derivatives of price are nonnegative')
+    return price, False  # not converged yet (or we don't know anyhow)
 
-        # this is 7.7a from Tra's dissertation
-        price = price - (excess_demand / deriv)
-        LOG.info(f'After iteration {itr}, max unit diff: {maxdiff}, max price diff: {maxpricediff}, min price diff: {minpricediff}, additional excluded due to rent/income ration {deltaNExcluded}')
 
-#@numba.jit(nopython=True) seems to slow things down but I may just not be letting it warm up
-def compute_derivatives (price, alt_income, choiceidx, hhidx, non_price_utilities, budget, budget_step, price_step, budget_coef, base_shares, max_rent_to_income,
-        feasible_alts, feasible_alts_step, weights):
+# @numba.jit(nopython=True) seems to slow things down but I may just not be letting it warm up
+def compute_derivatives (price, alt_income, choiceidx, hhidx, non_price_utilities, budget, budget_step, price_step,
+                         budget_coef, base_shares, max_rent_to_income, feasible_alts, feasible_alts_step, weights):
     deriv = np.zeros_like(price)
 
     base_exp_utilities = np.exp(non_price_utilities + budget_coef * budget)
@@ -158,9 +171,10 @@ def compute_derivatives (price, alt_income, choiceidx, hhidx, non_price_utilitie
         exp_utilities[choicemask] = base_exp_utilities[choicemask]
 
     return deriv
-    
-def compute_shares (price, supply, alt_income, choiceidx, hhidx, non_price_utilities, price_income_transformation, price_income_params, budget_coef,
-        max_rent_to_income, weights):
+
+
+def compute_shares (price, supply, alt_income, choiceidx, hhidx, non_price_utilities, price_income_transformation,
+                    price_income_params, budget_coef, max_rent_to_income, weights):
     alt_price = price[choiceidx]
 
     # unfortunately feasible alts does need to be recalculated on each iter as prices may change
@@ -170,11 +184,12 @@ def compute_shares (price, supply, alt_income, choiceidx, hhidx, non_price_utili
         feasible_alts = alt_income * max_rent_to_income > alt_price
 
     budgets = np.zeros_like(alt_price)
-    budgets[feasible_alts] = price_income_transformation.apply(alt_income[feasible_alts], alt_price[feasible_alts], *price_income_params)
+    budgets[feasible_alts] = price_income_transformation.apply(alt_income[feasible_alts],
+                                                               alt_price[feasible_alts], *price_income_params)
     full_utilities = non_price_utilities + budget_coef * budgets
     exp_utility = np.exp(full_utilities)
-    exp_utility[~feasible_alts] = 0 # will force choice probability to zero for infeasible alts
-    
+    exp_utility[~feasible_alts] = 0  # will force choice probability to zero for infeasible alts
+
     if not np.all(np.isfinite(exp_utility)):
         raise FloatingPointError('Not all exp(utilities) are finite (scaling?)')
 
diff --git a/src/eqsormo/tra/price_income.py b/src/eqsormo/tra/price_income.py
index cd8e898..d64a738 100644
--- a/src/eqsormo/tra/price_income.py
+++ b/src/eqsormo/tra/price_income.py
@@ -21,6 +21,7 @@ Functional forms for price and income
 import numpy as np
 import scipy.stats
 
+
 class FunctionalForm(object):
     def __init__ (self, func, name, n_params=0, starting_values=None, param_names=None):
         '''
@@ -70,7 +71,7 @@ sqrtdiff = FunctionalForm(lambda income, price: np.sqrt(income - price), 'sqrt(i
 # NB any price > income should be taken care of by a max_rent_to_income parameter
 # TODO can we remove the max rent to income param here?
 boxcoxdiff = FunctionalForm(lambda income, price, lmbda: scipy.stats.boxcox(np.maximum(income - price, 1), lmbda=lmbda), 'boxcox(income - price)',
-     n_params=1, param_names=['boxcox_lambda'])
+                            n_params=1, param_names=['boxcox_lambda'])
 
 # Form suggested by Stephane Hess in his advanced discrete choice modelling class, London, UK, July 2019
 # Note that this will not behave right if called on a subset of the data, since it uses the mean of income
diff --git a/src/eqsormo/tra/tra_sorting_model.py b/src/eqsormo/tra/tra_sorting_model.py
index 52f7ee6..de4e967 100644
--- a/src/eqsormo/tra/tra_sorting_model.py
+++ b/src/eqsormo/tra/tra_sorting_model.py
@@ -12,86 +12,126 @@
 #    See the License for the specific language governing permissions and
 #    limitations under the License.
 
-# Author: Matthew Wigginton Conway <matt@indicatrix.org>, School of Geographical Sciences and Urban Planning, Arizona State University
+# Author: Matthew Wigginton Conway <matt@indicatrix.org>
+#         School of Geographical Sciences and Urban Planning
+#         Arizona State University
 
 from logging import getLogger
 import time
 import statsmodels.api as sm
 import pandas as pd
 import numpy as np
-import scipy.optimize, scipy.stats
-from tqdm import tqdm
-import pickle
 import datetime
 
 from . import price_income
-from .clear_market import clear_market
+from .clear_market import clear_market_iter
 from eqsormo.common import BaseSortingModel, MNLFullASC
 from eqsormo.common.util import human_bytes, human_time, human_shape
 from eqsormo.common.compute_ascs import compute_ascs
 import eqsormo
 from . import save_load
 
+from typing import Dict, Callable, Any
+
 LOG = getLogger(__name__)
 
+
 class TraSortingModel(BaseSortingModel):
     '''
-    The variety of sorting model described in Tra CI (2013) Measuring the General Equilibrium Benefits of Air Quality Regulation in Small Urban Areas. Land Economics 89(2): 291–307.
-
-    The key difference between this and the Bayer et al. model is that price is included in the first stage rather than the second stage. This obviates the need to instrument
-    for price, because there are housing type fixed effects in the first stage. But we can't include price directly, due to identification
-    challenges since the prices could also be captured by the ASCs. So instead we include a budget constraint of the form f(income, price) - Tra uses
-    ln(income - price) but the exact functional form is not important, as long as income (or another household attribute) is somehow interacted with price
-    so that there is variation in price between households 
+    The variety of sorting model described in Tra CI (2013) Measuring the General Equilibrium Benefits of Air Quality
+    Regulation in Small Urban Areas. Land Economics 89(2): 291–307.
+
+    The key difference between this and the Bayer et al. model is that price is included in the first stage rather than
+    the second stage. This obviates the need to instrument for price, because there are housing type fixed effects in
+    the first stage. But we can't include price directly, due to identification challenges since the prices could also
+    be captured by the ASCs. So instead we include a budget constraint of the form f(income, price) - Tra uses
+    ln(income - price) but the exact functional form is not important, as long as income (or another household
+    attribute) is somehow interacted with price so that there is variation in price between households within a single
+    housing type.
+
+    The other key difference from both the Bayer and the Tra formulations is that this model supports "unequilibrated
+    choice," that is it can model a joint choice between a good for which prices are adjusted to bring the market to
+    equilibrium and one for which they are not. In my dissertation I use this to model the joint choice of housing
+    (equilibrated good) and vehicle ownership (unequilibrated good).
     '''
 
-    def __init__ (self, housing_attributes, household_attributes,  interactions, unequilibrated_hh_params, unequilibrated_hsg_params, second_stage_params, price, income, choice, unequilibrated_choice, price_income_transformation=price_income.logdiff,
-            price_income_starting_values=[], sample_alternatives=None, method='L-BFGS-B', max_rent_to_income=None, household_housing_attributes=None, weights=None, max_chunk_bytes=2e9, est_first_stage_ses=True, seed=None):
+    def __init__ (self, housing_attributes, household_attributes, interactions, unequilibrated_hh_params,
+                  unequilibrated_hsg_params, second_stage_params, price, income, choice, unequilibrated_choice,
+                  price_income_transformation=price_income.logdiff, price_income_starting_values=[],
+                  endogenous_variable_defs: Dict[str, Callable[[Any, Any], float]] = None,
+                  neighborhoods=None, sample_alternatives=None, method='L-BFGS-B',
+                  max_rent_to_income=None, household_housing_attributes=None, weights=None, max_chunk_bytes=2e9,
+                  est_first_stage_ses=True, seed=None):
         '''
         Initialize a Tra sorting model
 
         :param housing_attributes: Attributes of housing choices. Price should not be an attribute here.
         :type housing_attributes: Pandas dataframe.
 
-        :param household_attributes: Attributes of households. It is okay if income is an attribute here as long as it is also passed in separately - since we assume it doesn't change when a household moves.
+        :param household_attributes: Attributes of households. It is okay if income is an attribute here as long as it
+            is also passed in separately - since we assume it doesn't change when a household moves.
         :type household_attributes: Pandas dataframe
 
-        :param household_housing_attributes: Attributes of the combination of households and housing choices, multiindexed with household ID and choice ID. Must contain all possible combinations.
+        :param household_housing_attributes: Attributes of the combination of households and housing choices,
+            multiindexed with household ID and choice ID. Must contain all possible combinations.
         :type household_housing_attributes: Pandas dataframe
 
-        :param interactions: Which household attributes should be interacted with which housing attributes in the first stage
+        :param interactions: Which household attributes should be interacted with which housing attributes in the
+            first stage. Endogenous variables may be included as housing attributes here.
         :type interactions: iterable of tuple (household_attributes column, housing_attributes column)
 
-        :param second_stage_params: Which housing attributes to include in the second stage fit (decomposition of mean parameters). If None, don't fit a second stage (this is not necessary for simulation).
+        :param second_stage_params: Which housing attributes to include in the second stage fit (decomposition of mean
+            parameters). If None, don't fit a second stage (this stage is not necessary for simulation). Note that
+            currently endogenous variables cannot be used in this stage.
         :type second_stage_params: iterable of string
 
-        :param price: Observed price of each housing choice. This price should be over the same time period as income (for example, annualized prices with annual incomes)
+        :param price: Observed price of each housing choice. This price should be over the same time period as income
+            (for example, annualized prices with annual incomes)
         :type price: Pandas series of numbers, indexed like housing_attributes
 
         :param income: Household income
         :type income: Pandas series of numbers, indexed like household_attributes
-        
+
         :param choice: The choice made by the household. Should be index values in household_attributes
         :param type: Pandas series, indexed like household_attributes
 
-        :param sample_alternatives: If a positive integer, take a random sample of this size for available alternatives to estimate the model. Recommended in models with a large number of alternatives. If 0 or None, no sampling will be performed. Use np.set_seed() before calling fit for model reproducibility when using this option.
+        :param endogenous_variable_defs: This is a dictionary mapping variable names to functions that calculate them,
+            for endogenous variables that need to be updated in the sorting stage of the model, as household re-sort
+            themselves across the region. The function should accept a data frame of housing attributes, a series with
+            household income, and a numpy array of weights for how likely that household is to live in a particular
+            neighborhood. Should return a float scalar for that neighborhood.
+        :type endogenous_variable_defs: Dict of string -> function
+
+        :param neighborhoods: Pandas series of which neighborhood each housing choice is in, for calculation of
+            endogenous variables. Indexed like housing_attributes.
+        :type neighborhoods: pd.Series
+
+        :param sample_alternatives: If a positive integer, take a random sample of this size for available alternatives
+            to estimate the model. Recommended in models with a large number of alternatives. If 0 or None, no sampling
+                will be performed.
         :type sample_alternatives: int or None
 
-        :param method: A scipy.optimize.minimize method used for maximizing log-likelihood of first stage, default 'bfgs'
+        :param method: A scipy.optimize.minimize method used for maximizing log-likelihood of first stage, default
+            'bfgs'
+        :type method: str
 
-        :param max_rent_to_income: Maximum proportion of income rent can be for the alternative to be included in the choice set - in (0, 1] if the price/income transformation can't handle income less than rent
+        :param max_rent_to_income: Maximum proportion of income rent can be for the alternative to be included in the
+            choice set - in (0, 1] if the price/income transformation can't handle income less than rent
         :type max_rent_to_income: float or None
 
-        :param max_chunk_bytes: Maximum number of bytes to use for a single chunk of the full alternatives array. The most memory-constrained part of the alogorithm
-        is clearing the market, which requires computing the dot product of a n_households x n_housing_alternatives * n_unequilibrated_alternatives array and a coefficients
-        array. However, since the coefficients array is a column vector, we can compute the product of chunks of the full alternatives array and the column vector, then concatenate.
-        This tuning parameter controls how large these chunks are in bytes. Default 2GB.
+        :param max_chunk_bytes: Maximum number of bytes to use for a single chunk of the full alternatives array.
+            The most memory-constrained part of the alogorithm is clearing the market, which requires computing the dot
+            product of a n_households x n_housing_alternatives * n_unequilibrated_alternatives array and a coefficients
+            array. However, since the coefficients array is a column vector, we can compute the product of chunks of the
+            full alternatives array and the column vector, then concatenate.
+            This tuning parameter controls how large these chunks are in bytes. Default 2GB.
         :type max_chunk_bytes: int
 
-        :param est_first_stage_ses: If False, do not estimate standard errors for the first stage. This can significantly speed convergence and is recommended during model development.
+        :param est_first_stage_ses: If False, do not estimate standard errors for the first stage. This can
+            significantly speed convergence and is recommended during model development.
         :type est_first_stage_ses: bool
 
-        :param seed: seed for the RNG
+        :param seed: seed for the random number generator used in sampling alternatives
         :type seed: int
         '''
         self.housing_attributes = housing_attributes
@@ -108,7 +148,7 @@ class TraSortingModel(BaseSortingModel):
         self.unequilibrated_hsg_params = unequilibrated_hsg_params
         self.sample_alternatives = sample_alternatives
         self.alternatives_stds = None
-        
+
         if weights is None:
             self.weights = None
             self.weighted_supply = self.choice.value_counts()
@@ -116,6 +156,12 @@ class TraSortingModel(BaseSortingModel):
             self.weights = weights.reindex(household_attributes.index)
             self.weighted_supply = self.weights.groupby(self.choice).sum()
 
+        self.endogenous_variable_defs = endogenous_variable_defs
+        if neighborhoods is not None:
+            self.neighborhoods = neighborhoods.reindex(self.housing_attributes.index)
+        else:
+            self.neighborhoods
+
         self.price_income_transformation = price_income_transformation
         self.price_income_starting_values = price_income_starting_values
         self.method = method
@@ -131,6 +177,7 @@ class TraSortingModel(BaseSortingModel):
         self.creation_time = datetime.datetime.today()
 
     def validate (self):
+        # TODO this could easily be moved to a new file
         allPassed = True
 
         choiceCount = self.choice.value_counts().reindex(self.housing_attributes.index, fill_value=0)
@@ -138,8 +185,9 @@ class TraSortingModel(BaseSortingModel):
             choiceList = ' - ' + '\n - '.join(choiceCount.index[choiceCount == 0])
             LOG.error(f'Some housing alternatives are not chosen by any households!\n{choiceList}')
             allPassed = False
-        
-        if self.max_rent_to_income is not None and not np.all(self.income.values * self.max_rent_to_income > self.price.loc[self.choice].values):
+
+        if (self.max_rent_to_income is not None
+                and not np.all(self.income.values * self.max_rent_to_income > self.price.loc[self.choice].values)):
             LOG.error('Some households pay more in rent than the max rent to income ratio')
             allPassed = False
 
@@ -153,10 +201,16 @@ class TraSortingModel(BaseSortingModel):
                 allPassed = False
 
             if hsgattr not in self.housing_attributes.columns:
-                LOG.error(f'Attribute {hsgattr} is used in interactions but is not in housing_attributes')
+                if self.endogenous_variable_defs is None or hsgattr not in self.endogenous_variable_defs:
+                    LOG.error(f'Attribute {hsgattr} is used in interactions but is not in housing_attributes')
+                    allPassed = False
+
+            if hsgattr in self.housing_attributes.columns and self.endogenous_variable_defs is not None\
+                    and hsgattr in self.endogenous_variable_defs:
+                LOG.error(f'{hsgattr} in both housing_attributes and endogenous variables')
                 allPassed = False
 
-            if self.housing_attributes[hsgattr].isnull().any():
+            if hsgattr in self.housing_attributes.columns and self.housing_attributes[hsgattr].isnull().any():
                 LOG.error(f'Attribute {hsgattr} contains NaNs')
                 allPassed = False
 
@@ -170,6 +224,14 @@ class TraSortingModel(BaseSortingModel):
                     LOG.error(f'Attribute {hsgattr} contains NaNs')
                     allPassed = False
 
+        if self.endogenous_variable_defs is not None:
+            if self.neighborhoods is None:
+                LOG.error('Neighborhoods required when endogenous variables in use')
+                allPassed = False
+            if self.neighborhoods.isnull().any():
+                LOG.error('All choices must have a neighborhood defined')
+                allPassed = False
+
         # TODO more checks
         if allPassed:
             LOG.info('All validation checks passed!')
@@ -198,7 +260,7 @@ class TraSortingModel(BaseSortingModel):
             price_income_params = self.price_income_starting_values
 
         assert len(hhidx) == len(choiceidx) and len(choiceidx) == len(uneqchoiceidx)
-        
+
         if self.household_housing_attributes is not None:
             assert hh_hsgidx is not None and len(hhidx) == len(hh_hsgidx)
 
@@ -214,7 +276,7 @@ class TraSortingModel(BaseSortingModel):
 
         if self.household_housing_attributes is not None:
             ncols += len(self.household_housing_attributes.columns)
-        
+
         LOG.info(f'Allocating alternatives array of size {human_bytes(len(hhidx) * ncols * 8)}')
         alternatives = np.zeros((len(hhidx), ncols))
 
@@ -232,7 +294,7 @@ class TraSortingModel(BaseSortingModel):
             feasible_alts = np.full(len(alt_income), True)
 
         LOG.info(f'{np.sum(feasible_alts)} options appear in choice sets')
-        
+
         budget = np.full(len(hhidx), np.nan)
         budget[feasible_alts] = self.price_income_transformation.apply(alt_income[feasible_alts], alt_price[feasible_alts], *price_income_params)
         assert not np.any(np.isnan(budget[feasible_alts])) # should be no nans left
@@ -241,11 +303,23 @@ class TraSortingModel(BaseSortingModel):
         del alt_income, alt_price, budget # save memory
 
         for hh_attr, hsg_attr in self.interactions:
-            alternatives[:,current_col] =\
-                self.household_attributes[hh_attr].astype('float64').values[hhidx] * self.housing_attributes[hsg_attr].astype('float64').values[choiceidx]
+            if hsg_attr in self.housing_attributes.columns:
+                # TODO lots of type conversion happening here. Could maybe refactor to do less.
+                alternatives[:, current_col] = (
+                    self.household_attributes[hh_attr].astype('float64').values[hhidx]
+                    * self.housing_attributes[hsg_attr].astype('float64').values[choiceidx]
+                )
+            elif hsg_attr in self.endogenous_varnames:
+                endogenous_col = self.endogenous_varnames.index(hsg_attr)
+                alternatives[:, current_col] = (
+                    self.household_attributes[hh_attr].astype('float64').values[hhidx]
+                    * self.endogenous_variables[self.nbhd_for_choice[choiceidx], endogenous_col]
+                )
+            else:
+                raise KeyError(f'{hsg_attr} is not a housing attribute, exogenous or endogenous')
             colnames.append(f'{hh_attr}:{hsg_attr}')
             current_col += 1
-        
+
         # now add the attributes for the unequilibrated choice
         for param in self.unequilibrated_hh_params:
             vals = self.household_attributes[param].astype('float64').values[hhidx]
@@ -266,8 +340,10 @@ class TraSortingModel(BaseSortingModel):
         if self.household_housing_attributes is not None:
             for c in self.household_housing_attributes.columns:
                 colnames.append(c)
-            
-            alternatives[:,current_col:] = self.household_housing_attributes.values[hh_hsgidx,:]
+
+            alternatives[:,current_col:current_col + len(self.household_housing_attributes.columns)] =\
+                self.household_housing_attributes.values[hh_hsgidx, :]
+            current_col += len(self.household_housing_attributes.columns)
 
         total_time = time.perf_counter() - start_time
         LOG.info(f'Materialized alternatives into {human_shape(alternatives.shape)} array using {human_bytes(alternatives.nbytes)} in {human_time(total_time)}')
@@ -275,7 +351,7 @@ class TraSortingModel(BaseSortingModel):
 
         if self.alternatives_stds is None:
             self.alternatives_stds = np.std(alternatives, axis=0)
-        
+
         alternatives /= self.alternatives_stds
 
         return alternatives
@@ -291,12 +367,20 @@ class TraSortingModel(BaseSortingModel):
         # good ol' mononomial logit model
         unequilibrated_choice = self.unequilibrated_choice.copy() if self.unequilibrated_choice is not None else pd.Series(np.zeros(len(self.choice), index=self.choice.index))
         unique_unequilibrated_choices = unequilibrated_choice.unique()
-        self.unequilibrated_choice_xwalk = pd.Series(np.arange(len(unique_unequilibrated_choices)), index=unique_unequilibrated_choices)
+        self.unequilibrated_choice_xwalk = pd.Series(np.arange(len(unique_unequilibrated_choices)),
+                                                     index=unique_unequilibrated_choices)
         self.hh_xwalk = pd.Series(np.arange(len(self.household_attributes)), index=self.household_attributes.index)
 
         self.hh_hsg_choice = self.housing_xwalk.loc[self.choice.loc[self.hh_xwalk.index]].values
         self.hh_unequilibrated_choice = self.unequilibrated_choice_xwalk.loc[self.unequilibrated_choice.loc[self.hh_xwalk.index]].values
 
+        if self.endogenous_variable_defs is not None:
+            # neighborhood indices
+            unique_neighborhoods = self.neighborhoods.unique()
+            self.nbhd_xwalk = pd.Series(np.arange(len(unique_neighborhoods)), index=unique_neighborhoods)
+            # neighborhood for each housing choice, neighborhood index for each neighborhood
+            self.nbhd_for_choice = self.nbhd_xwalk.loc[self.neighborhoods.loc[self.housing_xwalk.index]].values  # index is sorted
+
         # index of each household in the full alternatives dataset
         # repeated for each alternative (housing choice)
         # so if there are three households, three housing choices, and three unequilibrated choices:
@@ -309,7 +393,7 @@ class TraSortingModel(BaseSortingModel):
         self.full_hhidx = np.repeat(np.arange(len(self.household_attributes)), len(self.housing_attributes) * len(unique_unequilibrated_choices))
         self.full_choiceidx = np.repeat(np.tile(np.arange(len(self.housing_attributes)), len(self.household_attributes)), len(unique_unequilibrated_choices))
         self.full_uneqchoiceidx = np.tile(np.arange(len(unique_unequilibrated_choices)), len(self.housing_attributes) * len(self.household_attributes))
-        self.full_hsgchosen = (self.hh_hsg_choice[self.full_hhidx] == self.full_choiceidx) 
+        self.full_hsgchosen = (self.hh_hsg_choice[self.full_hhidx] == self.full_choiceidx)
         self.full_uneqchosen = (self.hh_unequilibrated_choice[self.full_hhidx] == self.full_uneqchoiceidx)
         self.full_chosen = self.full_hsgchosen & self.full_uneqchosen
 
@@ -328,7 +412,10 @@ class TraSortingModel(BaseSortingModel):
                     self.housing_xwalk.index[self.full_choiceidx[chunk_start:chunk_end]]))].values
             del hh_hsg_loc # save memory
 
-        if self.sample_alternatives <= 0 or self.sample_alternatives is None:
+        # calculate endogenous variables if necessary (conditional is inside the function)
+        self.initialize_or_update_endogenous_variables(initial=True)
+
+        if self.sample_alternatives is None or self.sample_alternatives <= 0:
             if self.max_rent_to_income is None:
                 self.alternatives = self.materialize_alternatives(self.full_hhidx, self.full_choiceidx, self.full_uneqchoiceidx,
                     self.full_hh_hsgidx if self.household_housing_attributes is not None else None)
@@ -419,7 +506,7 @@ class TraSortingModel(BaseSortingModel):
         # unequilibrated alternatives do not
         # TODO could move this calculation into MNLFullASC
         utils += self._log_supply
-        
+
         return utils
 
     def chunked_full_alternatives (self, price_income_params):
@@ -582,7 +669,8 @@ class TraSortingModel(BaseSortingModel):
         TODO document which data structures can be changed and still have this function return correct results
         '''
         LOG.info('Clearing the market and sorting households')
-        LOG.info("There's nothing hidden in your head / the Sorting Hat can't see / so try me on and I will tell you / where you ought to be.\n" +\
+        LOG.info("There's nothing hidden in your head / the Sorting Hat can't see / so try me on and I will tell you /'\
+                 ' where you ought to be.\n" +\
             "    -JK Rowling, Harry Potter and the Sorcerer's Stone")
 
         # convert supply to an np array
@@ -596,7 +684,7 @@ class TraSortingModel(BaseSortingModel):
         if self.second_stage_params is not None:
             LOG.info('updating second stage')
             pred_ascs = self.second_stage_fit.predict(sm.add_constant(self.housing_attributes[self.second_stage_params])) + self.type_shock
-            
+
             maxabsdiff = np.max(np.abs(pred_ascs - self.first_stage_ascs))
             LOG.info(f'Second stage updated with changes to first-stage ASCs of up to {maxabsdiff:.2f}')
             self.first_stage_ascs = pred_ascs
@@ -604,43 +692,114 @@ class TraSortingModel(BaseSortingModel):
             LOG.info('No second stage fit, not updating')
             pred_ascs = self.first_stage_ascs[0]
 
-        LOG.info('finding non-price utilites')
-        non_price_utilities = self.full_utility(include_budget=False)
-        assert not np.any(np.isnan(non_price_utilities))
-
-        # Note that I am intentionally _not_ dropping hh/choice combinations here that do not meet rent to income criteria, because which households those are
-        # might change in the sorting phase. The filtering happens there. This does not matter for the calculation of utility below, since utilities for
-        # different alternatives are independent of each other and the budget is set to zero anyhow.
-
         if self.price_income_transformation.n_params > 0:
             price_income_params = self.first_stage_fit.params.values[-self.price_income_params:]
         else:
             price_income_params = np.zeros(0)
 
-        startTimeClear = time.perf_counter()
-        new_prices = clear_market(
-            non_price_utilities=non_price_utilities,
-            hhidx=self.full_hhidx,
-            choiceidx=self.full_choiceidx,
-            supply=supply,
-            income=self.income.loc[self.hh_xwalk.index].values,
-            starting_price=self.price.loc[self.housing_xwalk.index].values,
-            price_income_transformation=self.price_income_transformation,
-            price_income_params=price_income_params,
-            budget_coef=self.first_stage_fit.params['budget'],
-            max_rent_to_income=self.max_rent_to_income,
-            maxiter=maxiter,
-            weights=self.weights.loc[self.hh_xwalk.index].values if self.weights is not None else None
-        )
-        endTimeClear = time.perf_counter()
+        itr = 0
+        while True:
+            if itr > maxiter:
+                LOG.error(f'Prices FAILED TO CONVERGE after {itr} iterations')
+                break
+            itr += 1
+            LOG.info(f'sorting: begin iteration {itr}')
+            self.initialize_or_update_endogenous_variables(initial=False)
+
+            LOG.info('finding non-price utilites')
+            non_price_utilities = self.full_utility(include_budget=False)
+            assert not np.any(np.isnan(non_price_utilities))
+
+            # Note that I am intentionally _not_ dropping hh/choice combinations here that do not meet rent to income criteria, because which households those are
+            # might change in the sorting phase. The filtering happens there. This does not matter for the calculation of utility below, since utilities for
+            # different alternatives are independent of each other and the budget is set to zero anyhow.
+
+            startTimeClear = time.perf_counter()
+            new_prices = clear_market_iter(
+                non_price_utilities=non_price_utilities,
+                hhidx=self.full_hhidx,
+                choiceidx=self.full_choiceidx,
+                supply=supply,
+                income=self.income.loc[self.hh_xwalk.index].values,
+                starting_price=self.price.loc[self.housing_xwalk.index].values,
+                price_income_transformation=self.price_income_transformation,
+                price_income_params=price_income_params,
+                budget_coef=self.first_stage_fit.params['budget'],
+                max_rent_to_income=self.max_rent_to_income,
+                maxiter=maxiter,
+                weights=self.weights.loc[self.hh_xwalk.index].values if self.weights is not None else None
+            )
+            endTimeClear = time.perf_counter()
+
+            new_prices = pd.Series(new_prices, index=self.housing_xwalk.index)
 
-        new_prices = pd.Series(new_prices, index=self.housing_xwalk.index)
+            maxPriceChange = np.max(np.abs(new_prices - self.price))
 
-        maxPriceChange = np.max(np.abs(new_prices - self.price))
+            LOG.info(f'Market cleared in {human_time(endTimeClear - startTimeClear)}, with price changes up to {maxPriceChange:.4f}')
 
-        LOG.info(f'Market cleared in {human_time(endTimeClear - startTimeClear)}, with price changes up to {maxPriceChange:.4f}')
+            self.price = new_prices
+
+    def initialize_or_update_endogenous_variables (self, initial: bool) -> None:
+        '''
+        Initialize endogenous variables, or update them based on predicted probabilities after the model has been fit.
+        :param initial: if True, will compute endogenous variables based on observed choices. If False, will compute
+            based on predicted probabilities
+        :type initial: bool
+        '''
+        if self.endogenous_variable_defs is None:
+            LOG.info('not updating endogenous variables as none are defined')
+        else:
+            if initial:
+                # initial fit: use actual chosen values to compute endogenous variables
+                LOG.info('creating endogenous variables')
+                weighted_probs = self.full_hsgchosen
+            else:
+                # if model is already fitted, update endogenous variables with fitted values
+                LOG.info('updating endogenous variables')
+                weighted_probs = self._probabilities()
+
+            if self.weights is not None:
+                # if weights are present, use them
+                weighted_probs = weighted_probs * self.weights.loc[self.hh_xwalk.index].values[self.full_hhidx]
+
+            # if this is the initial run, initialize endogenous_varnames
+            # need to materialize and save because dict.keys() order is not guaranteed, though in recent versions of
+            # Python it is stable (in fact, may have always been stable?)
+            if initial:
+                self.endogenous_varnames = list(self.endogenous_variable_defs.keys())
+
+            # Neighborhood for every element of weighted_probs
+            nbhdidx = self.nbhd_for_choice[self.full_choiceidx]
+
+            # reset endogenous variables. this creates a new array. if we instead cleared the old array, we would need
+            # to create a copy in the if block below to ensure that the original values were preserved for comparison
+            # purposes
+            self.endogenous_variables = np.full(
+                (len(self.endogenous_varnames), np.max(self.nbhd_for_choice) + 1),
+                np.nan,
+                'float64'
+            )
 
-        self.price = new_prices
+            # compute endogenous variables
+            for neighborhood in range(np.max(self.nbhd_for_choice) + 1):
+                nbhdmask = nbhdidx == neighborhood
+                # all probabilities for the neighborhood
+                nbhdweights = weighted_probs[nbhdmask]
+                # how likely every household is to choose this neighborhood
+                hhweights = np.bincount(self.full_hhidx[nbhdmask], weights=nbhdweights)
+                # use the precomputed list to avoid order changing, as iteration order in python dicts is not guaranteed
+                # endogenous_varnames is created by initialize_endogenous_variables
+                for i, varname in enumerate(self.endogenous_varnames):
+                    func = self.endogenous_variable_defs[varname]
+                    self.endogenous_variables[neighborhood, i] = func(self.household_attributes, self.income, hhweights)
+
+            # check our work
+            assert not np.any(np.isnan(self.endogenous_variables)), 'some endogenous variables are nan!'
+
+            if initial:
+                # save the initial variables for later comparison
+                # this array is never modified, an entirely new array is created when sorting, so no need to copy
+                self.orig_endogenous_variables = self.endogenous_variables
 
     def _probabilities (self):
         'Compute probabilities and return as numpy array, use .probabilities() for a Pandas data frame'
@@ -676,7 +835,7 @@ class TraSortingModel(BaseSortingModel):
 
     def mkt_shares (self):
         return pd.Series(self._mkt_shares(), index=self.housing_xwalk.index)
-    
+
     def _uneq_mkt_shares (self):
         probs = self._probabilities() * self.weights.loc[self.hh_xwalk.index].values[self.full_hhidx]
         return np.bincount(self.full_uneqchoiceidx, probs)
@@ -710,7 +869,7 @@ Unequilibrated ASCs:
 Second stage (OLS parameters):
 {second_stage_summary}
 
-Fit with EqSorMo version {version}, https://github.com/mattwigway/eqsormo            
+Fit with EqSorMo version {version}, https://github.com/mattwigway/eqsormo
         '''.format(
             first_stage_summary=self.first_stage_fit.summary(),
             unequilibrated_ascs=pd.DataFrame(self.first_stage_uneq_ascs).to_string(),
@@ -736,10 +895,3 @@ Fit with EqSorMo version {version}, https://github.com/mattwigway/eqsormo
         tra = super().from_pickle(fn)
         tra.create_full_alternatives()
         return tra
-
-
-    
-
-
-
-


Test patch:
diff --git a/tests/tra/test_endogenous_variables.py b/tests/tra/test_endogenous_variables.py
new file mode 100644
index 0000000..8ec702d
--- /dev/null
+++ b/tests/tra/test_endogenous_variables.py
@@ -0,0 +1,188 @@
+#    Copyright 2020 Matthew Wigginton Conway
+
+#    Licensed under the Apache License, Version 2.0 (the "License");
+#    you may not use this file except in compliance with the License.
+#    You may obtain a copy of the License at
+
+#        http://www.apache.org/licenses/LICENSE-2.0
+
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS,
+#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#    See the License for the specific language governing permissions and
+#    limitations under the License.
+
+# Author: Matthew Wigginton Conway <matt@indicatrix.org>,
+#         School of Geographical Sciences and Urban Planning
+#         Arizona State University
+
+import pandas as pd
+import numpy as np
+
+import eqsormo
+from eqsormo.tra import TraSortingModel
+eqsormo.enable_logging()
+
+
+def test_endogenous_variables (monkeypatch):
+    model = TraSortingModel(
+        housing_attributes=pd.DataFrame({
+            'age': [10, 20, 30, 40, 50]
+        }),
+        household_attributes=pd.DataFrame({
+            'white': [1, 1, 0, 0, 1]
+        }),
+        interactions=(('white', 'prop_white'),),
+        unequilibrated_hh_params=tuple(),
+        unequilibrated_hsg_params=tuple(),
+        second_stage_params=tuple(),
+        # avoid price-income issues by making all prices less than all incomes - or maybe this should be tested
+        # explicitly?
+        price=pd.Series([5, 5, 5, 5, 5]),
+        income=pd.Series([10, 20, 30, 40, 50]),
+        choice=pd.Series([0, 1, 2, 3, 4]),
+        unequilibrated_choice=pd.Series([0, 0, 0, 0, 0]),
+        endogenous_variable_defs={
+            'prop_white': lambda hh, inc, weights: np.average(hh.white, weights=weights),
+            'mean_income': lambda hh, inc, weights: np.average(inc, weights=weights)
+        },
+        neighborhoods=pd.Series(['a', 'a', 'b', 'b', 'b'])
+    )
+
+    model.create_alternatives()
+
+    model.initialize_or_update_endogenous_variables(initial=True)
+
+    # as of Python 3.something, dict iteration order is the same as definition order. This test relies on this;
+    # otherwise, we will not compare prop_white to expected prop_white. Can be fixed by actually reading the values in
+    # endogenous_varnames
+    assert model.endogenous_varnames == ['prop_white', 'mean_income'],\
+        'endogenous varnames incorrect or not in expected order'
+
+    # as before, this could be harmless - but if our assumption in testing here was not met, later asserts would fail
+    # in weird and wonderful ways
+    assert np.all(model.nbhd_xwalk == pd.Series([0, 1], index=['a', 'b'])), 'nbhd_xwalk not in expected order'
+
+    # again, check that order is as expected - changes could be harmless as long as they are consistent across all
+    # indices
+    assert np.all(model.housing_xwalk == pd.Series([0, 1, 2, 3, 4], index=[0, 1, 2, 3, 4]))
+
+    # make sure nbhd_for_choice is correct
+    assert np.all(model.nbhd_for_choice == np.array([0, 0, 1, 1, 1]))
+
+    # prop_white
+    # First neighborhood is neighborhood 'a', should be 100% white b/c it contains households 0/1
+    # second neighborhood is 'b', should be 33% white b/c it contains households 2/3/4
+    assert np.allclose(model.endogenous_variables[:, 0], np.array([1, 1 / 3]))
+
+    # mean income
+    # First neighborhood is neighborhood 'a', should be 10 + 20 / 2 = 15
+    # Second should be 30 + 40 + 50 / 3 = 40.
+    assert np.allclose(model.endogenous_variables[:, 1], np.array([(10 + 20) / 2, (30 + 40 + 50) / 3]))
+
+    # test that the interactions with the household variables are properly computed
+    # construct the interaction term outside the model
+    # first two choices are in 100% white nbhd, next three are 1/3
+    intterm = (model.household_attributes.white.to_numpy()[model.full_hhidx]
+               * np.array([1, 1, 1 / 3, 1 / 3, 1 / 3])[model.full_choiceidx])
+
+    assert model.alternatives_colnames[1] == 'white:prop_white'
+
+    # need to destandardize alternatives here, as they are roughly standardized to promote convergence, and
+    # coefficients are automatically unstandardized after estimation.
+    assert np.allclose((model.alternatives * model.alternatives_stds)[:, 1], intterm)
+
+    # test that variables are updated correctly - this makes every household equally likely to choose each location
+    faked_probabilities = np.full(model.alternatives.shape[0], 0.2)
+    monkeypatch.setattr(model, '_probabilities', lambda: np.copy(faked_probabilities))
+
+    # do the update
+    model.initialize_or_update_endogenous_variables(initial=False)
+
+    # every neighborhood should have 3 out of five households white because every neighborhood has same composition now
+    assert np.allclose(model.endogenous_variables[:, 0], np.array([3 / 5, 3 / 5]))
+    # and have a mean income of 30
+    assert np.allclose(model.endogenous_variables[:, 1], np.array([30, 30]))
+
+
+def test_endogenous_variables_with_weights (monkeypatch):
+    model = TraSortingModel(
+        housing_attributes=pd.DataFrame({
+            'age': [10, 20, 30, 40, 50]
+        }),
+        household_attributes=pd.DataFrame({
+            'white': [1, 1, 0, 0, 1]
+        }),
+        interactions=(('white', 'prop_white'),),
+        unequilibrated_hh_params=tuple(),
+        unequilibrated_hsg_params=tuple(),
+        second_stage_params=tuple(),
+        # avoid price-income issues by making all prices less than all incomes - or maybe this should be tested
+        # explicitly?
+        price=pd.Series([5, 5, 5, 5, 5]),
+        income=pd.Series([20, 20, 30, 40, 50]),
+        choice=pd.Series([0, 1, 2, 3, 4]),
+        weights=pd.Series([1, 1, 2, 1, 1]),
+        unequilibrated_choice=pd.Series([0, 0, 0, 0, 0]),
+        endogenous_variable_defs={
+            'prop_white': lambda hh, inc, weights: np.average(hh.white, weights=weights),
+            'mean_income': lambda hh, inc, weights: np.average(inc, weights=weights)
+        },
+        neighborhoods=pd.Series(['a', 'a', 'b', 'b', 'b'])
+    )
+
+    model.create_alternatives()
+
+    model.initialize_or_update_endogenous_variables(initial=True)
+
+    # as of Python 3.something, dict iteration order is the same as definition order. This test relies on this;
+    # otherwise, we will not compare prop_white to expected prop_white. Can be fixed by actually reading the values in
+    # endogenous_varnames
+    assert model.endogenous_varnames == ['prop_white', 'mean_income'],\
+        'endogenous varnames incorrect or not in expected order'
+
+    # as before, this could be harmless - but if our assumption in testing here was not met, later asserts would fail
+    # in weird and wonderful ways
+    assert np.all(model.nbhd_xwalk == pd.Series([0, 1], index=['a', 'b'])), 'nbhd_xwalk not in expected order'
+
+    # again, check that order is as expected - changes could be harmless as long as they are consistent across all
+    # indices
+    assert np.all(model.housing_xwalk == pd.Series([0, 1, 2, 3, 4], index=[0, 1, 2, 3, 4]))
+
+    # make sure nbhd_for_choice is correct
+    assert np.all(model.nbhd_for_choice == np.array([0, 0, 1, 1, 1]))
+
+    # prop_white
+    # First neighborhood is neighborhood 'a', should be 100% white b/c it contains households 0/1
+    # second neighborhood is 'b', should be 25% white b/c it contains households 2/3/4, and hh 2 is weighted double
+    assert np.allclose(model.endogenous_variables[:, 0], np.array([1, 1 / 4]))
+
+    # mean income
+    # First neighborhood is neighborhood 'a', should be 20 as both hhs are homogenous
+    # Second should be 30 + 30 + 40 + 50 / 4 bc hh 4 is weighted twice
+    assert np.allclose(model.endogenous_variables[:, 1], np.array([20, (30 + 30 + 40 + 50) / 4]))
+
+    # test that the interactions with the household variables are properly computed
+    # construct the interaction term outside the model
+    # first two choices are in 100% white nbhd, next three are 1/3
+    intterm = (model.household_attributes.white.to_numpy()[model.full_hhidx]
+               * np.array([1, 1, 1 / 4, 1 / 4, 1 / 4])[model.full_choiceidx])
+
+    assert model.alternatives_colnames[1] == 'white:prop_white'
+
+    # need to destandardize alternatives here, as they are roughly standardized to promote convergence, and
+    # coefficients are automatically unstandardized after estimation.
+    assert np.allclose((model.alternatives * model.alternatives_stds)[:, 1], intterm)
+
+    # test that variables are updated correctly - this makes every household equally likely to choose each location
+    faked_probabilities = np.full(model.alternatives.shape[0], 0.2)
+    monkeypatch.setattr(model, '_probabilities', lambda: np.copy(faked_probabilities))
+
+    # do the update
+    model.initialize_or_update_endogenous_variables(initial=False)
+
+    # every neighborhood should have 3 out of 6 households white because every neighborhood has same composition now,
+    # and one nonwhite hh is weighted double
+    assert np.allclose(model.endogenous_variables[:, 0], np.array([3 / 6, 3 / 6]))
+    # and have a mean income of 20, 20, 30, 30, 40, 50 b/c hh 2 is (30) is weighted double
+    assert np.allclose(model.endogenous_variables[:, 1], np.full(2, np.mean([20, 20, 30, 30, 40, 50])))