natural language toolkit dispersion plots c 20012023 nltk project steven bird stevenbird1gmail com url https www nltk org for license information see license txt a utility for displaying lexical dispersion generate a lexical dispersion plot param text the source text type text liststr or iterstr param words the target words type words list of str param ignorecase flag to set if case should be ignored when searching text type ignorecase bool return a matplotlib axes object that may still be modified before plotting rtype axes natural language toolkit dispersion plots c 2001 2023 nltk project steven bird stevenbird1 gmail com url https www nltk org for license information see license txt a utility for displaying lexical dispersion generate a lexical dispersion plot param text the source text type text list str or iter str param words the target words type words list of str param ignore_case flag to set if case should be ignored when searching text type ignore_case bool return a matplotlib axes object that may still be modified before plotting rtype axes
def dispersion_plot(text, words, ignore_case=False, title="Lexical Dispersion Plot"):
    

    try:
        import matplotlib.pyplot as plt
    except ImportError as e:
        raise ImportError(
            "The plot function requires matplotlib to be installed. "
            "See https://matplotlib.org/"
        ) from e

    word2y = {
        word.casefold() if ignore_case else word: y
        for y, word in enumerate(reversed(words))
    }
    xs, ys = [], []
    for x, token in enumerate(text):
        token = token.casefold() if ignore_case else token
        y = word2y.get(token)
        if y is not None:
            xs.append(x)
            ys.append(y)

    words = words[::-1]

    _, ax = plt.subplots()
    ax.plot(xs, ys, "|")
    ax.set_yticks(list(range(len(words))), words, color="C0")
    ax.set_ylim(-1, len(words))
    ax.set_title(title)
    ax.set_xlabel("Word Offset")
    return ax


if __name__ == "__main__":
    import matplotlib.pyplot as plt

    from nltk.corpus import gutenberg

    words = ["Elinor", "Marianne", "Edward", "Willoughby"]
    dispersion_plot(gutenberg.words("austen-sense.txt"), words)
    plt.show()