CODEWITHSUNDEEP
pythonnumpymatplotliboptional-parametersnormal-distribution
sully11
sully11

Reputation: 37

Python function to draw random numbers based on chosen distribution

I am writing a function that will draw random numbers from a given distribution like Poisson, Normal or Binomial. It takes one argument for number of samples and a second argument for type of distribution. It would accept additional parameters based on the distribution chosen. So if I take Normal samples then its mean and sd.

Is there an optimal way of writing this?

My code 

import matplotlib.pyplot as plt
import numpy as np

def randNumberDistribution(samples, distribution,*optional):
    if distribution.capitalize() == 'Normal':
        if len(optional) == 2:
            mean, sd, = optional
            s = np.random.normal(mean, sd, samples)
            print(s)
            count, bins, ignored = plt.hist(s, 20, density=True)
            plt.plot(bins, 1/(sigma * np.sqrt(2 * np.pi)) * np.exp( - (bins - mu)**2 / (2 * sigma**2) ),linewidth=3, color='y')
            plt.show()
        else:
            print("Invalid number of arguments")
    if distribution.capitalize() == 'Binomial':
        if len(optional) == 2:
            numOfTrials, probSuccess = optional # number of trials, probability of success(each trial)
            s = np.random.binomial(n, p, samples)
            count, bins, ignored = plt.hist(s, 14, density=True)
        else:
            print("Invalid number of arguments")
    if distribution.capitalize() == 'Poisson':
        if len(optional) == 1:
            exp = optional
            s = np.random.poisson(exp, samples) #Expectation of interval(should be >= 0)
            count, bins, ignored = plt.hist(s, 14, density=True)
        else:
            print("Invalid number of arguments")


print(randNumberDistribution(5,'Poisson',5))

Upvotes: 0

Views: 522

Answers (1)

Sheldore
Sheldore

Reputation: 39072

You can make use of things like dictionaries or Classes but IMHO, I find your code simple and easy to understand. I would just make use of elif to avoid checking all the if statements once a condition is satisfied. Below is a slightly short version. Your else statements were specific for a given distribution type. Mine simply checks for the 3 cases and says invalid for any other input.

def randNumberDistribution(samples, distribution,*optional):
    if distribution.capitalize() == 'Normal' and len(optional) == 2:
            mean, sd, = optional
            s = np.random.normal(mean, sd, samples)
            count, bins, ignored = plt.hist(s, 20, density=True)
            plt.plot(bins, 1/(sigma * np.sqrt(2 * np.pi)) * np.exp( - (bins - mu)**2 / (2 * sigma**2) ),linewidth=3, color='y')
            plt.show()
    elif distribution.capitalize() == 'Binomial' and len(optional) == 2:
            n, p = optional # number of trials, probability of success(each trial)
            s = np.random.binomial(n, p, samples)
            count, bins, ignored = plt.hist(s, 14, density=True)
    elif distribution.capitalize() == 'Poisson' and len(optional) == 1:
            exp = optional
            s = np.random.poisson(exp, samples) #Expectation of interval(should be >= 0)
            count, bins, ignored = plt.hist(s, 14, density=True)
    else:
        print("Invalid number of arguments")

Upvotes: 1

Related Questions