diff --git a/activitysim/skim.py b/activitysim/skim.py new file mode 100644 index 0000000000..599c5eab71 --- /dev/null +++ b/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] diff --git a/activitysim/tests/__init__.py b/activitysim/tests/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/activitysim/tests/test_skim.py b/activitysim/tests/test_skim.py new file mode 100644 index 0000000000..2e8fce79f8 --- /dev/null +++ b/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])