Inverse transform sampling

Question

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 sympy 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 sympy 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][0]

    # 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
    See also question on SO:
    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[0]

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

    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[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:
    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()

Answer 1

1. If you do not need custom piecewise probabilities (which look very strange to me) and instead want to use one of the many SciPy distributions, you can use their predefined percent point functions (which are inverse CDFs). Only if you actually need to create a new distribution, you might find my next observations useful. Also, as agtoever mentioned, you might want to take advantage of rv_continuous in scipy.
2. Using SymPy is a two-edged sword. Defining and computing CDFs and inverse samples symbolically might be exact, but it may cost more in terms of syntax or CPU time. Still, the CDF is costly to compute numerically also, so your approach might be right for such piecewise functions.
3. For each call of sample, the CDF should not be recalculated. You should split the computation of the CDF from a PDF into a different function (and have def sample(cdf: [...]). In general, my advice is the same as uncle Bob's, namely, to have functions as short as you can make them. (for instance, you could also extract lines from if not isinstance(pdf, sym.Piecewise): to raise NotImplementedError(error_message) in a method which could be called validate.
4. After extracting the CDF and/or the validation methods, they become testable, and you might want to add tests for them. Perhaps you can also extract and test the pointwise solution of the inverse. Sadly, testing the random sampling itself might be impossible (would fail randomly).
5. I am not very skilled in optimization, but looking at another solution, it seems to me you've got the right idea in deriving the CDF from the given PDF, then solve the pointwise inverse. The other solution finds the CDF numerically, and has the option of approximating the CDF using Chebyshev polynomials for performance reasons.

While I can't find more significant problems with the solution, I can't be sure it's correct either (you can never be with software in general). But tests might increase one's confidence in the code.

Answer 2

Given that there hasn't been an answer on this question for almost 2 years, I thought I'd throw my 2c in and add something.

Am I reinventing the wheel? I searched thoroughly but couldn't find anything similar

Is this a correct approach for this problem?

I think these questions would be better asked on Stack Overflow/Math than Code Review. Validating proofs are more their areas, we're more about fixing your code problems (which is probably why this question hasn't had an answer yet has over 3K views and almost 2 years?).

Looking at your code, it seems fine, however there are some things I'd point out (in sample)-

   if not isinstance(pdf, sym.Piecewise):
        raise ValueError("PDF must be constructed by sympy.Piecewise")

If you raise this error - does sympy catch it and deal with it? Or does it just die? If it does not catch it properly, it would be a better approach for you to just print the message and return None. That's a more standard pattern IMO. Again for the raise NotImplementedError(error_message) - if you reach an impassible spot in your code - there's no point raising a message unless it is handled above it. Just log/print the error (no point continuing, right?) and end the program.

Regarding the functions replace_lambertw_branch and to_lower_lambertw_branch - they are only used once. I understand the need to introduce clarity but don't create a whole function with all the scaffolding when it's only used once (and is a single line). It's unnecessary code and forces readers of your code to jump around trying to figure out what is going on.

I could also say the same for the inner functions of remove_nans and right_hand_number, even though they're closer, it still requires the code reader to stop the flow and go looking elsewhere for what is happening. They too are only used once - there's no need to create a whole scaffold when a lambda will do.

Adding further comment - if users/readers of the code are unclear what the lambdas do, you can improve the naming of the variables in the lambda to make it clear, or add in a comment above the lambda to explain "the why" - but these should be kept to an absolute minimum.

Continuing the point about comments - they're scattered everywhere throughout the code. This is bad for two reasons. One - it interrupts the comprehension of the code flow - and comments lie. If they're all very important - they should be moved into a separate document, explaining your approach, and the reasons why you did what you did in the way you did it.

Those sorts of "the why" comments don't belong in code - because the minute the code changes for whatever reason - the comment then lies about the code - and it will confuse everyone who looks at your code. Rarely, you have well-meaning coders that attempt to fix the code to match the comment (rare, but I have seen that).

Otherwise, the rest of the code seems fine. Perhaps others can add to what I've written. I hope this helps in some small way.