Add more specific type annotations for Numpy arrays

[kevinsung]

So far we only use np.ndarray in type annotations but they could be made more specific to specify expected data types and shapes.