Merge b298c0e into e099bf9

urbansim/models/tests/test_transition.py

@@ -301,3 +301,22 @@ def test_transition_model(basic_df, grow_targets_filters, totals_col, year):
     assert added.isin(new.index).all()
     assert not added.isin(basic_df.index).any()
     npt.assert_array_equal(added.values, [basic_df.index.values.max() + 1])
+
+
+def test_tabular_transition_add_and_remove():
+    data = pd.DataFrame(
+        {'a': ['x', 'x', 'y', 'y', 'y', 'y', 'y', 'y', 'z', 'z']})
+
+    totals = pd.DataFrame(
+        {'a': ['x', 'y', 'z'],
+         'total': [3, 1, 10]},
+        index=[2112, 2112, 2112])
+
+    tran = transition.TabularTotalsTransition(totals, 'total')
+    model = transition.TransitionModel(tran)
+
+    new, added, _ = model.transition(data, 2112)
+
+    assert len(new) == totals.total.sum()
+    assert added.is_unique is True
+    assert new.index.is_unique is True

urbansim/models/transition.py

@@ -5,8 +5,6 @@
 """
 from __future__ import division
 
-import itertools
-
 import numpy as np
 import pandas as pd
 
@@ -264,7 +262,7 @@ def transition(self, data, year):
         copied_indexes = []
         removed_indexes = []
 
-        # since we're looping over descrete segments we need to track
+        # since we're looping over discrete segments we need to track
         # out here where their new indexes will begin
         starting_index = data.index.values.max() + 1
 
@@ -273,7 +271,9 @@ def transition(self, data, year):
             nrows = self._calc_nrows(len(subset), row[self._config_column])
             updated, added, copied, removed = \
                 add_or_remove_rows(subset, nrows, starting_index)
-            starting_index = starting_index + nrows + 1
+            if nrows > 0:
+                # only update the starting index if rows were added
+                starting_index = starting_index + nrows
             segments.append(updated)
             added_indexes.append(added)
             copied_indexes.append(copied)