Simple container class for skim arrays

activitysim/skim.py

@@ -0,0 +1,43 @@
+import numpy as np
+
+
+class Skim(object):
+    """
+    Container for skim arrays.
+
+    Parameters
+    ----------
+    data : 2D array
+    offset : int, optional
+        An optional offset that will be added to origin/destination
+        values to turn them into array indices.
+        For example, if zone IDs are 1-based, an offset of -1
+        would turn them into 0-based array indices.
+
+    """
+    def __init__(self, data, offset=None):
+        self.data = data
+        self.offset = offset
+
+    def get(self, orig, dest):
+        """
+        Get impedence values for a set of origin, destination pairs.
+
+        Parameters
+        ----------
+        orig : 1D array
+        dest : 1D array
+
+        Returns
+        -------
+        values : 1D array
+
+        """
+        orig = np.asanyarray(orig)
+        dest = np.asanyarray(dest)
+
+        if self.offset:
+            orig = orig + self.offset
+            dest = dest + self.offset
+
+        return self.data[orig, dest]

activitysim/tests/test_skim.py

@@ -0,0 +1,32 @@
+import numpy as np
+import numpy.testing as npt
+import pytest
+
+from .. import skim
+
+
+@pytest.fixture
+def data():
+    return np.arange(100, dtype='int').reshape((10, 10))
+
+
+def test_basic(data):
+    sk = skim.Skim(data)
+
+    orig = [5, 9, 1]
+    dest = [2, 9, 6]
+
+    npt.assert_array_equal(
+        sk.get(orig, dest),
+        [52, 99, 16])
+
+
+def test_offset(data):
+    sk = skim.Skim(data, offset=-1)
+
+    orig = [6, 10, 2]
+    dest = [3, 10, 7]
+
+    npt.assert_array_equal(
+        sk.get(orig, dest),
+        [52, 99, 16])