Skip to content

Instantly share code, notes, and snippets.

What would you like to do?
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'),
if in pdf_functions:
error_message = ("Using sympy.Abs or 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:
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:
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:
inverse_solutions = sym.solve(eq, y, rational=False)
# Sometimes, especially for exponents,
# there are garbage solutions with imaginary parts:
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:
inverse_cdf = recreate_piecewise(inverse_solutions)
# If inverse CDF will contain LambertW function,
# we must change its branch. For more details, see:
functions = map(operator.attrgetter('func'),
if sym.LambertW in functions:
inverse_cdf = replace_lambertw_branch(inverse_cdf)
# This is to prevent LambertW giving ComplexWarning after lambdifying
inverse_cdf =
max_value = cdf.args[-1][0]
# Warnings can happen with exponents in PDF:
lambda_function = sym.lambdify(args=x,
modules=[{'LambertW': lambertw}, 'numpy'])
return lambda_function(np.random.uniform(high=max_value,
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:
See also question on SO:
def remove_nans(expression_condition: ExprCondPair) -> ExprCondPair:
return expression_condition.args[0]
def right_hand_number(solution: ExprCondPair) -> sym.S:
return solution[1].args[1]
solutions = sorted(map(remove_nans, functions),
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[0], -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:
Solution is based on the 2nd example from:
:return: expression with replaced LambertW branch by a lower one
return expression.replace(sym.LambertW,
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.