Repository: quantopian/zipline
File: tests/pipeline/test_engine.py (Python)

Code before:
                    'dv5': AverageDollarVolume(window_length=5),
                    'dv1_nan': AverageDollarVolume(
                        window_length=1,
                        inputs=[USEquityPricing.open, USEquityPricing.volume],
                    ),
                    'dv5_nan': AverageDollarVolume(
                        window_length=5,
                        inputs=[USEquityPricing.open, USEquityPricing.volume],
                    ),
                }
            ),
            self.dates[5],
            self.dates[-1],
        )

        expected_1 = (self.raw_data[5:] ** 2) * 2
        assert_frame_equal(results['dv1'].unstack(), expected_1)

        expected_5 = rolling_mean((self.raw_data ** 2) * 2, window=5)[5:]
        assert_frame_equal(results['dv5'].unstack(), expected_5)

        expected_1_nan = (self.raw_data_with_nans[5:]
                          * self.raw_data[5:] * 2).fillna(0)
        assert_frame_equal(results['dv1_nan'].unstack(), expected_1_nan)

        expected_5_nan = rolling_mean((self.raw_data_with_nans
                                       * self.raw_data * 2).fillna(0),
                                      window=5)[5:]
        assert_frame_equal(results['dv5_nan'].unstack(), expected_5_nan)

    @parameter_space(returns_length=[2, 3], correlation_length=[3, 4])
    def test_correlation_factors(self, returns_length, correlation_length):
        """
        Tests for the built-in factors `RollingPearsonOfReturns` and
        `RollingSpearmanOfReturns`.
        """
        my_asset_column = 0
        start_date_index = 14
        end_date_index = 18

        sids = self.sids
        dates = self.dates
        assets = self.asset_finder.retrieve_all(sids)
        my_asset = assets[my_asset_column]
        num_days = end_date_index - start_date_index + 1
        num_assets = len(assets)

        cascading_mask = \
            AssetIDPlusDay() < (sids[-1] + dates[start_date_index].day)
        expected_cascading_mask_result = make_cascading_boolean_array(
            shape=(num_days, num_assets),

Review comment (suggestion):
You can still just do `(self.raw_data_with_nans ** 2) * 2` here. Multiplying by `self.raw_data` will just fill with NaNs in the same places.


Code after:
                    'dv5': AverageDollarVolume(window_length=5),
                    'dv1_nan': AverageDollarVolume(
                        window_length=1,
                        inputs=[USEquityPricing.open, USEquityPricing.volume],
                    ),
                    'dv5_nan': AverageDollarVolume(
                        window_length=5,
                        inputs=[USEquityPricing.open, USEquityPricing.volume],
                    ),
                }
            ),
            self.dates[5],
            self.dates[-1],
        )

        expected_1 = (self.raw_data[5:] ** 2) * 2
        assert_frame_equal(results['dv1'].unstack(), expected_1)

        expected_5 = rolling_mean((self.raw_data ** 2) * 2, window=5)[5:]
        assert_frame_equal(results['dv5'].unstack(), expected_5)

        # The following two use USEquityPricing.open and .volume as inputs.
        # The former uses self.raw_data_with_nans, and the latter uses
        # .raw_data * 2.  Thus we multiply instead of squaring as above.
        expected_1_nan = (self.raw_data_with_nans[5:]
                          * self.raw_data[5:] * 2).fillna(0)
        assert_frame_equal(results['dv1_nan'].unstack(), expected_1_nan)

        expected_5_nan = rolling_mean((self.raw_data_with_nans
                                       * self.raw_data * 2).fillna(0),
                                      window=5)[5:]
        assert_frame_equal(results['dv5_nan'].unstack(), expected_5_nan)

    @parameter_space(returns_length=[2, 3], correlation_length=[3, 4])
    def test_correlation_factors(self, returns_length, correlation_length):
        """
        Tests for the built-in factors `RollingPearsonOfReturns` and
        `RollingSpearmanOfReturns`.
        """
        my_asset_column = 0
        start_date_index = 14
        end_date_index = 18

        sids = self.sids
        dates = self.dates
        assets = self.asset_finder.retrieve_all(sids)
        my_asset = assets[my_asset_column]
        num_days = end_date_index - start_date_index + 1
        num_assets = len(assets)

        cascading_mask = \