Inverse transform sampling

Inverse transform sampling is a method to generate random values that follow an arbitrary distribution. For some reason this method was never implemented in any popular scientific libraries. And as I often need to use it, instead of calculating it by hands every time, I decided to implement the function doing it for me.

What I want reviewed:

• Am I reinventing the wheel? I searched thoroughly but couldn't find anything similar.
• Is this a correct approach for this problem? I've seen the following code on SO: link. There the PDF is approximated by a discrete distribution. Maybe, that way is better. I don't know.
• There are so many issues within that my function looks like a bunch of patches and workarounds in order to make it work. Maybe there are more elegant and correct ways to deal with those drawbacks.
• Missed cases for inputs. There are an infinite number of PDFs. I could miss some.
• Type hints. Did I write them correctly? With the types of objects are quite confusing.
• Code style.

Code:
inverse_transform.py

import operator
from typing import Iterator

import numpy as np
import sympy as sym
from scipy.special import lambertw
from sympy.functions.elementary.piecewise import ExprCondPair

def sample(pdf: sym.Function,
*,
size: int) -> np.array:
"""
Generates random values following the given distribution
:param pdf: input Probability Density Function (PDF)
:param size: number of generated values
"""
if not isinstance(pdf, sym.Piecewise):
raise ValueError("PDF must be constructed by sympy.Piecewise")

pdf_functions = map(operator.attrgetter('func'),
pdf.atoms(sym.Function))
if sym.re in pdf_functions:
error_message = ("Using sympy.Abs or sympy.re is not supported "
"due to not implemented computing of their integrals "
"within SymPy. Split the relevant expression.")
raise NotImplementedError(error_message)

# The following is used in order to prevent an error
# when using PDF in a form of, for example, x**-2.5.
# For more details see:
# https://stackoverflow.com/questions/50543587/integrating-piecewise-with-irrational-exponent-gives-error
pdf = sym.nsimplify(pdf)

x = pdf.free_symbols.pop()
y = sym.Dummy('y')

cdf = sym.integrate(pdf, (x, -sym.oo, y))
# The following is used in order to prevent
# long erroneous polynomials
# when calculating PDF in a form of, for example,  x**-2.5
# Beware that this will add too much precision. Bug.
# Issue submitted: https://github.com/sympy/sympy/issues/14787
cdf = cdf.evalf()

eq = sym.Eq(x, cdf)

# TODO: Use solveset when it will be able to deal with LambertW
# With default rational == True, there will be an error
# as 'solve' doesn't play along with Piecewise.
# Related issue: https://github.com/sympy/sympy/issues/12024
inverse_solutions = sym.solve(eq, y, rational=False)
# Sometimes, especially for exponents,
# there are garbage solutions with imaginary parts:
# https://github.com/sympy/sympy/issues/9973
inverse_solutions = filter(is_real, inverse_solutions)

# As, for some reason, 'solve' returns a list of Piecewise's,
# it's necessary to collect them back together.
# Related issue: https://github.com/sympy/sympy/issues/14733
inverse_cdf = recreate_piecewise(inverse_solutions)
# If inverse CDF will contain LambertW function,
# we must change its branch. For more details, see:
# https://stackoverflow.com/questions/49817984/sympy-solve-doesnt-give-one-of-the-solutions-with-lambertw
functions = map(operator.attrgetter('func'),
inverse_cdf.atoms(sym.Function))
if sym.LambertW in functions:
inverse_cdf = replace_lambertw_branch(inverse_cdf)
# This is to prevent LambertW giving ComplexWarning after lambdifying
inverse_cdf = sym.re(inverse_cdf)

max_value = cdf.args[-1]

# Warnings can happen with exponents in PDF:
# https://github.com/sympy/sympy/issues/14789
lambda_function = sym.lambdify(args=x,
expr=inverse_cdf,
modules=[{'LambertW': lambertw}, 'numpy'])
return lambda_function(np.random.uniform(high=max_value,
size=size))

def is_real(expression: sym.Expr) -> bool:
"""Checks if expression doesn't contain imaginary part with sympy.I"""
return sym.I not in expression.atoms()

def recreate_piecewise(functions: Iterator[ExprCondPair]) -> sym.Piecewise:
"""
Collects Piecewise from list of unsorted Piecewise's,
ignoring parts with NaNs.
Solution for the issue: https://github.com/sympy/sympy/issues/14733
https://stackoverflow.com/questions/50428912/how-to-get-sorted-exprcondpairs-in-a-piecewise-function-that-was-obtained-from
"""
def remove_nans(expression_condition: ExprCondPair) -> ExprCondPair:
return expression_condition.args

def right_hand_number(solution: ExprCondPair) -> sym.S:
return solution.args

solutions = sorted(map(remove_nans, functions),
key=right_hand_number)
return sym.Piecewise(*solutions)

def to_lower_lambertw_branch(*args) -> sym.Function:
"""
Wraps the first argument from a given list of arguments
as a lower branch of LambertW function.
:return: lower LambertW branch
"""
return sym.LambertW(args, -1)

def replace_lambertw_branch(expression: sym.Expr) -> sym.Expr:
"""
Replaces upper branch of LambertW function with the lower one.
For details of the bug see:
https://stackoverflow.com/questions/49817984/sympy-solve-doesnt-give-one-of-the-solutions-with-lambertw
Solution is based on the 2nd example from:
http://docs.sympy.org/latest/modules/core.html?highlight=replace#sympy.core.basic.Basic.replace
:return: expression with replaced LambertW branch by a lower one
"""
return expression.replace(sym.LambertW,
to_lower_lambertw_branch)

Examples of usage:
I will plot results in order to give a better idea:

import matplotlib.pyplot as plt
import sympy as sym

import inverse_transform

x = sym.Symbol('x')
f = sym.Piecewise((0, x < 0.),
(1, x <= 1.),
(0, True))
plt.hist(inverse_transform.sample(f, size=10**6),
bins=100)
plt.show() f = sym.Piecewise((0, x < 4.3),
(1, x < 12.9),
(5, x <= 13.5),
(0, True))
plt.hist(inverse_transform.sample(f, size=10**6),
bins=100)
plt.show() shift = 1.5
f = sym.Piecewise((0., x <= shift),
((x - shift) * sym.exp(-(x - shift)), x <= 13.5),
(0., True))
plt.hist(inverse_transform.sample(f, size=10**6),
bins=100)
plt.show() f = sym.Piecewise((0, x < 6.5),
(97.25 / (25 + x**2) , x < 10.5),
(0, True))
plt.hist(inverse_transform.sample(f, size=10**6),
bins=100)
plt.show() f = sym.Piecewise((0, x < 0.4),
(x ** -2.35, x < 50),
(0, True))
plt.hist(inverse_transform.sample(f, size=10**6),
bins=100)
plt.show() f = sym.Piecewise((0, x < 6.5),
(sym.exp(-x/3.5) , x < 10.5),
(0, True))
plt.hist(inverse_transform.sample(f, size=10**6),
bins=100)
plt.show() f = sym.Piecewise((0, x < -2),
(sym.exp(x/0.25) , x < 0),
(sym.exp(-x/0.25) , x < 2),
(0, True))
plt.hist(inverse_transform.sample(f, size=10**6),
bins=100)
plt.show() • This is a nice code, have you done any performance tests comparing it with the answers here stackoverflow.com/q/21100716/1391441? Also, if sympy causes so many issues (does is still?), couldn't it be replaced by some other package or vanilla Python? – Gabriel Sep 15 at 13:38
• Thanks! I haven't performed any performance tests as I was very much satisfied with how fast my code ran. About SymPy, from all the described bugs in the code I know only about one that was fixed, no idea about others but looks like nothing has changed. As my approach uses symbolic math, then, I'm afraid, there are no other similar libraries for Python. Maybe some other languages have better functionality for this task like MATLAB, Mathematica or R, and we could call them from Python? Worth investigating I think. – Georgy Sep 15 at 14:57