diff --git a/postreise/analyze/__init__.py b/postreise/analyze/__init__.py index 4fea14a5..9d84b87e 100644 --- a/postreise/analyze/__init__.py +++ b/postreise/analyze/__init__.py @@ -1 +1 @@ -__all__ = ['carbon', 'distance', 'mwmiles', 'transmission'] +__all__ = ['carbon', 'distance', 'mwmiles', 'statelines', 'transmission'] diff --git a/postreise/analyze/statelines.py b/postreise/analyze/statelines.py new file mode 100644 index 00000000..b1245928 --- /dev/null +++ b/postreise/analyze/statelines.py @@ -0,0 +1,135 @@ +zones_to_states = { + 1: 'ME', + 2: 'NH', + 3: 'VT', + 4: 'MA', + 5: 'RI', + 6: 'CT', + 7: 'NY', + 8: 'NY', + 9: 'NJ', + 10: 'PA', + 11: 'PA', + 12: 'DE', + 13: 'MD', + 14: 'VA', + 15: 'VA', + 16: 'NC', + 17: 'NC', + 18: 'SC', + 19: 'GA', + 20: 'GA', + 21: 'FL', + 22: 'FL', + 23: 'FL', + 24: 'AL', + 25: 'MS', + 26: 'TN', + 27: 'KY', + 28: 'WV', + 29: 'OH', + 30: 'OH', + 31: 'MI', + 32: 'MI', + 33: 'IN', + 34: 'IL', + 35: 'IL', + 36: 'WI', + 37: 'MN', + 38: 'MN', + 39: 'IA', + 40: 'MO', + 41: 'MO', + 42: 'AR', + 43: 'LA', + 44: 'TX', + 45: 'TX', + 46: 'NM', + 47: 'OK', + 48: 'KS', + 49: 'NE', + 50: 'SD', + 51: 'ND', + 52: 'MT', + 201: 'WA', + 202: 'OR', + 203: 'CA', + 204: 'CA', + 205: 'CA', + 206: 'CA', + 207: 'CA', + 208: 'NV', + 209: 'AZ', + 210: 'UT', + 211: 'NM', + 212: 'CO', + 213: 'WY', + 214: 'ID', + 215: 'MT', + 216: 'TX', + 301: 'TX', + 302: 'TX', + 303: 'TX', + 304: 'TX', + 305: 'TX', + 306: 'TX', + 307: 'TX', + 308: 'TX', + } + + +def classify_interstate_intrastate(scenario): + """Classifies branches in a change_table as interstate or intrastate. + + :param powersimdata.scenario.scenario.Scenario scenario: scenario instance. + :return: (*dict*) -- keys are *'interstate'*, *'intrastate'*. Values are + list of branch ids. + """ + + ct = scenario.state.get_ct() + grid = scenario.state.get_grid() + upgraded_branches = _classify_interstate_intrastate(ct, grid) + return upgraded_branches + + +def _classify_interstate_intrastate(ct, grid): + """Classifies branches in a change_table as interstate or intrastate. + This function is separate from classify_interstate_intrastate() for testing + purposes. + + :param dict ct: change_table dictionary. + :param powersimdata.input.grid.Grid grid: Grid instance. + :return: (*dict*) -- keys are *'interstate'*, *'intrastate'*. Values are + list of branch ids. + """ + + branch = grid.branch + upgraded_branches = {'interstate': [], 'intrastate': []} + + if 'branch' not in ct or 'branch_id' not in ct['branch']: + return upgraded_branches + + all_upgraded_branches = ct['branch']['branch_id'].keys() + for b in all_upgraded_branches: + # Alternatively: bus.loc[branch.loc[b, 'from_bus_id'], 'from_zone_id'] + try: + from_zone = branch.loc[b, 'from_zone_id'] + to_zone = branch.loc[b, 'to_zone_id'] + except KeyError: + raise ValueError('ct entry not found in branch: ' + str(b)) + try: + from_state = zones_to_states[from_zone] + except KeyError: + err_msg = 'zone not found in zones_to_states: ' + str(from_zone) + raise ValueError(err_msg) + try: + to_state = zones_to_states[to_zone] + except KeyError: + err_msg = 'zone not found in zones_to_states: ' + str(to_zone) + raise ValueError(err_msg) + if from_state == to_state: + upgraded_branches['intrastate'].append(b) + else: + upgraded_branches['interstate'].append(b) + + return upgraded_branches diff --git a/postreise/analyze/tests/__init__.py b/postreise/analyze/tests/__init__.py index e69de29b..0a1d6ac1 100644 --- a/postreise/analyze/tests/__init__.py +++ b/postreise/analyze/tests/__init__.py @@ -0,0 +1 @@ +__all__ = ['test_carbon', 'test_mwmiles', 'test_statelines'] diff --git a/postreise/analyze/tests/test_statelines.py b/postreise/analyze/tests/test_statelines.py new file mode 100644 index 00000000..91dc20e9 --- /dev/null +++ b/postreise/analyze/tests/test_statelines.py @@ -0,0 +1,83 @@ +import unittest + +from postreise.tests.mock_grid import MockGrid +from postreise.analyze.statelines import _classify_interstate_intrastate + +# branch_id is the index +mock_branch = { + 'branch_id': [1, 2, 3, 4, 5, 6, 7, 8], + 'from_zone_id': [10, 10, 10, 52, 216, 216, 216, 7], + 'to_zone_id': [10, 11, 12, 215, 301, 216, 301, 204], + } + +expected_keys = {'interstate', 'intrastate'} + + +class TestClassifyInterstateIntrastate(unittest.TestCase): + def setUp(self): + def check_expected(upgrades, expected_interstate, expected_intrastate): + err_msg = 'classify_interstate_intrastate should return a dict' + self.assertIsInstance(upgrades, dict, err_msg) + err_msg = "dict keys should be 'interstate' and 'intrastate'" + self.assertEqual(upgrades.keys(), expected_keys, err_msg) + for v in upgrades.values(): + self.assertIsInstance(v, list, "dict values should be lists") + for b in v: + self.assertIsInstance(b, int, "branch_ids should be ints") + err_msg = 'interstate values not as expected' + self.assertEqual( + set(upgrades['interstate']), expected_interstate, err_msg) + err_msg = 'intrastate values not as expected' + self.assertEqual( + set(upgrades['intrastate']), expected_intrastate, err_msg) + self.check_expected = check_expected + self.mock_grid = MockGrid(grid_attrs={'branch': mock_branch}) + + def test_classify_interstate_intrastate_empty_ct(self): + mock_ct = {} + expected_interstate = set() + expected_intrastate = set() + + upgrades = _classify_interstate_intrastate(mock_ct, self.mock_grid) + self.check_expected(upgrades, expected_interstate, expected_intrastate) + + def test_classify_interstate_intrastate_bad_ct(self): + mock_ct = {'branch': {'branch_id': {9: 1.5}}} + + with self.assertRaises(ValueError): + upgrades = _classify_interstate_intrastate(mock_ct, self.mock_grid) + + def test_classify_interstate_intrastate_none(self): + mock_ct = {'branch': {'branch_id': {}}} + expected_interstate = set() + expected_intrastate = set() + + upgrades = _classify_interstate_intrastate(mock_ct, self.mock_grid) + self.check_expected(upgrades, expected_interstate, expected_intrastate) + + def test_classify_interstate_intrastate_two(self): + mock_ct = {'branch': {'branch_id': {1: 2, 8: 10}}} + expected_interstate = {8} + expected_intrastate = {1} + + upgrades = _classify_interstate_intrastate(mock_ct, self.mock_grid) + self.check_expected(upgrades, expected_interstate, expected_intrastate) + + def test_classify_interstate_intrastate_several(self): + mock_ct = {'branch': {'branch_id': { + 1: 2, + 2: 3, + 3: 1.5, + 4: 4, + 5: 1.1, + 8: 10, + }}} + expected_interstate = {3, 8} + expected_intrastate = {1, 2, 4, 5} + + upgrades = _classify_interstate_intrastate(mock_ct, self.mock_grid) + self.check_expected(upgrades, expected_interstate, expected_intrastate) + + +if __name__ == '__main__': + unittest.main()