Switched from math to numpy