使用蒙版设置numpy ndarray的值

Question

I want to calculate business days between two times, both of which contain null values, following this question related to calculating business days. I've identified that the way I'm setting values using a mask does not behave as expected.

I'm using python 2.7.11, pandas 0.18.1 and numpy 1.11.0. My slightly modified code:

import datetime
import numpy as np
import pandas as pd

def business_date_diff(start, end):
    mask = pd.notnull(start) & pd.notnull(end)
    start = start[mask]
    end = end[mask]
    start = start.values.astype('datetime64[D]')
    end = end.values.astype('datetime64[D]')
    result = np.empty(len(mask), dtype=float)
    result[mask] = np.busday_count(start, end)
    result[~mask] = np.nan
    return result

Unfortunately, this doesn't return the expected business day differences (instead I get a number of very near 0 floats). When I check np.busday_count(start, end) the results look correct.

print start[0:5]
print end[0:5]
print np.busday_count(start, end)[0:5]

# ['2016-07-04' '2016-07-04' '2016-07-04' '2016-07-04' '2016-07-04']
# ['2016-07-05' '2016-07-05' '2016-07-05' '2016-07-06' '2016-07-06']
# [1 1 1 2 2]

But when I check the values for results the results do not make sense:

...
result = np.empty(len(mask), dtype=float)
result[mask] = np.busday_count(start, end)
result[~mask] = np.nan
print result

# [           nan               nan   1.43700866e-210   1.45159738e-210
# 1.45159738e-210   1.45159738e-210   1.45159738e-210   1.46618609e-210
# 1.45159738e-210   1.64491834e-210   1.45159738e-210   1.43700866e-210
# 1.43700866e-210   1.43700866e-210   1.43700866e-210   1.45159738e-210
# 1.43700866e-210   1.43700866e-210   1.43700866e-210   1.43700866e-210

What am I doing wrong?

Answer 1

Your problem is that with your version of numpy, you can't use a boolean array as an index to an array. Just use np.where(mask==True) instead of mask and np.where(mask==False) instead of ~mask, and it will work as desired.