diff --git a/zipline/testing/__init__.py b/zipline/testing/__init__.py index 16bc73d7..34bf9db8 100644 --- a/zipline/testing/__init__.py +++ b/zipline/testing/__init__.py @@ -21,6 +21,7 @@ from .core import ( # noqa make_trade_panel_for_asset_info, num_days_in_range, parameter_space, + permute_rows, powerset, product_upper_triangle, seconds_to_timestamp, diff --git a/zipline/testing/core.py b/zipline/testing/core.py index 8bc05bc0..c759480b 100644 --- a/zipline/testing/core.py +++ b/zipline/testing/core.py @@ -31,6 +31,7 @@ from zipline.finance.order import ORDER_STATUS from zipline.pipeline.engine import SimplePipelineEngine from zipline.pipeline.loaders.testing import make_seeded_random_loader from zipline.utils import security_list +from zipline.utils.input_validation import expect_dimensions from zipline.utils.tradingcalendar import trading_days @@ -891,6 +892,22 @@ def parameter_space(**params): return decorator +@expect_dimensions(array=2) +def permute_rows(seed, array): + """ + Shuffle each row in ``array`` based on permutations generated by ``seed``. + + Parameters + ---------- + seed : int + Seed for numpy.RandomState + array : np.ndarray[ndim=2] + Array over which to apply permutations. + """ + rand = np.random.RandomState(seed) + return np.apply_along_axis(rand.permutation, 1, array) + + @nottest def make_test_handler(testcase, *args, **kwargs): """