Fixing doctest
adding function for calculating probability mass
testing the functionality