Plot Discrete Distribution using np.linspace()

Question

I'm trying to plot a simple discrete distribution using matplotlib:

If -1<=x<0, p=0.3;
If 0<=x<1, p=0.5;
If 1<=x<=2, p=0.2.

How can I start from x = np.linspace(-1, 2)?

What I tried so far is:

def mapDiscProb(x):
    if np.any(x < 0):
        return 0.3 + x * 0
    elif np.any(x >= 1):
        return 0.2 + x * 0
    else:
        return 0.5 + x * 0

x = np.linspace(-1, 2)
y = mapDiscProb(x4)

ax.plot(x, y, clip_on = False)

And the outcome is just a whole line at 0.3 from -1 to 2 as if the elif and else are not excuted.

My expected output is three disconnected horizontal lines, as is standard for a discrete pmf.

What have you tried so far? What issue are you having with it? Are you just trying to plot the PDF (like a histogram), or something else? Your title says "mapping", but your question says "plot". — Engineero, Dec 05 '17 at 20:24
Are you trying to use this function independently on all of the numbers in `linespace`? otherwise it makes sense that only the first`if` will trigger, as you will always have a number in the array less than 0 — DJK, Dec 05 '17 at 20:55
@djk47463 Yes, that's what I thought it would be: having all numbers in `linspace` fall into the right category and make a plot. — PacmanKX, Dec 05 '17 at 21:03
@PaulH You mean use "<" and ">" instead of "<=" and ">="? That's not working, still connected. — PacmanKX, Dec 05 '17 at 21:57

ImportanceOfBeingErnest · Answer 1 · 2017-12-05T21:53:11.273

You may use

`numpy.piecewise`

numpy.piecewise allows to define a function dependent on some conditions. Here you have three conditions [x<0, x>=1, (x>=0) & (x<1)], and you may define a function to use for each of them.

import matplotlib.pyplot as plt
import numpy as np

l1 = lambda x: 0.3 + x * 0
l2 = lambda x: 0.2 + x * 0
l3 = lambda x: 0.5 + x * 0

mapDiscProb=lambda x: np.piecewise(x, [x<0, x>=1, (x>=0) & (x<1)],[l1,l2,l3])

x = np.linspace(-1, 2)
y = mapDiscProb(x)

fig, ax = plt.subplots()
ax.plot(x, y, clip_on = False)

plt.show()

`numpy.vectorize`

numpy.vectorize vectorizes a function which is meant to be called with scalars, such that is evaluated for each element in an array. This allows if/else statements to be used as expected.

import matplotlib.pyplot as plt
import numpy as np

def mapDiscProb(x):
    if x < 0:
        return 0.3
    elif x >= 1:
        return 0.2
    else:
        return 0.5

x = np.linspace(-1, 2)
y = np.vectorize(mapDiscProb)(x)

fig, ax = plt.subplots()
ax.plot(x, y, clip_on = False)

plt.show()

`numpy.select`

(credit to PaulH for this idea) numpy.select can choose select values from different arrays based on a condition. For piecewise constant functions this is an easy tool, because it does not require to build any additional functions (one-liner).

import matplotlib.pyplot as plt
import numpy as np

x = np.linspace(-1, 2)
y = np.select([x<0, x<1, x>1], [0.3, 0.5, 0.2])

fig, ax = plt.subplots()
ax.plot(x, y, clip_on = False)

plt.show()

Output in all cases:

No vertical lines

In case you don't want any vertical lines to appear, it makes sense to plot as many plots as you have conditions.

import matplotlib.pyplot as plt
import numpy as np

l1 = lambda x: 0.3 + x * 0
l2 = lambda x: 0.2 + x * 0
l3 = lambda x: 0.5 + x * 0

x = np.linspace(-1, 2)
func = [l1,l2,l3]
cond = [x<0, x>=1, (x>=0) & (x<1)]

fig, ax = plt.subplots()
for f,c in zip(func,cond):
    xi = x[c]
    ax.plot(xi, f(xi), color="C0")

plt.show()

Alternatively, using numpy.select, you may modify the x array to surely include the values [0,1], which lie on the edge between conditions. Choosing the conditions to exclude those values explicitely, [x<0, (x>0) & (x<1), x>1] (note the lack of any equal sign) will allow to set those values to nan. Nan values are not shown, hence a gap appears.

import matplotlib.pyplot as plt
import numpy as np

x = np.sort(np.append(np.linspace(-1, 2),[0,1]))
y = np.select([x<0, (x>0) & (x<1), x>1], [0.3, 0.5, 0.2], np.nan)

fig, ax = plt.subplots()
ax.plot(x, y)

plt.show()

Thanks, but is there a way to eliminate the "nearly vertical" lines in the graph? — PacmanKX, Dec 05 '17 at 21:21
Do you not want any vertical lines or do you want them to be vertical? — ImportanceOfBeingErnest, Dec 05 '17 at 21:25
you could use `numpy.select` to set the values at the boundaries to `numpy.nan` ... e.g., `y = numpy.select([x<0, x<1, x>1], [0.3, 0.5, 0.2], numpy.nan)` — Paul H, Dec 05 '17 at 21:37
@PaulH nice one. However, in order to get rid of the lines one would first need to find out the edge values and modify the conditions to exclude those values. In any case I updated the answer. — ImportanceOfBeingErnest, Dec 05 '17 at 21:54
So basically you plot 3 different functions as if one single function. Way to go. — PacmanKX, Dec 05 '17 at 21:58
@PaulH No, unfortunately it doesn't. But the idea goes of course in the right direction. First, 0 and 1 are not in `x`, so we need to add them. Then, 0 is <1, such that the second condition is fulfilled, value 0.5 is taken and a vertical line appears. I think the way I modified it in the answer is correct though. — ImportanceOfBeingErnest, Dec 05 '17 at 22:00
@PaulH A rule of thumb: `linspace` never hits any integer if you put in integer values, unless you modify the `num` argument to an odd number or set `endpoint=False`. (necessary, not sufficient condition). — ImportanceOfBeingErnest, Dec 05 '17 at 22:09

Harald Heitmann · Accepted Answer · 2017-12-06T13:11:36.230

0

If you just want to do it for that example, this should work

import numpy as np
from matplotlib import pyplot as plt
x=np.linspace(-1,2)
plt.plot(x[x < 0], 0.3 + x[x < 0]*0, color='blue')
plt.plot(x[(x >= 0) & (x <= 1)], 0.5 + x[(x>=0) & (x<=1)]*0, color='blue')
plt.plot(x[x > 1],0.2 + x[x > 1]*0, color='blue')
plt.show()

masking is my way to go when it comes to discrete functions.

edited Dec 06 '17 at 13:11

answered Dec 05 '17 at 21:35

Harald Heitmann

129
2
10

Like the previous answer you basically plotted 3 different uniform distributions as 1 discrete distribution. I think this is a smart alternative and you used much less code to do the same thing. I'm just wondering for discrete distribution, it seems like there is just no way to build a function that y=f(x) first and use one single `plot(x, y)` to draw, is that true? – PacmanKX Dec 05 '17 at 22:09
Just to defend my answer: The last option in it uses much less code than this. ;-) – ImportanceOfBeingErnest Dec 05 '17 at 22:10

Plot Discrete Distribution using np.linspace()

2 Answers2

`numpy.piecewise`

`numpy.vectorize`

`numpy.select`

No vertical lines

Linked