<commit_before># -*- coding: utf-8 -*-
from __future__ import print_function

# plotting
import numpy as np
from matplotlib import pyplot as plt

from lifelines.utils import coalesce


def remove_spines(ax, sides):
    '''
    Remove spines of axis.

    Parameters:
      ax: axes to operate on
      sides: list of sides: top, left, bottom, right

    Examples:
    removespines(ax, ['top'])
    removespines(ax, ['top', 'bottom', 'right', 'left'])
    '''
    for side in sides:
        ax.spines[side].set_visible(False)


def move_spines(ax, sides, dists):
    '''
    Move the entire spine relative to the figure.

    Parameters:
      ax: axes to operate on
      sides: list of sides to move. Sides: top, left, bottom, right
      dists: list of float distances to move. Should match sides in length.

    Example:
    move_spines(ax, sides=['left', 'bottom'], dists=[-0.02, 0.1])
    '''
    for side, dist in zip(sides, dists):
        ax.spines[side].set_position(('axes', dist))


def remove_ticks(ax, x=False, y=False):
    '''
    Remove ticks from axis.

    Parameters:
      ax: axes to work on
      x: if True, remove xticks. Default False.
      y: if True, remove yticks. Default False.

    Examples:
    removeticks(ax, x=True)
    removeticks(ax, x=True, y=True)
    '''
    if x:
        ax.xaxis.set_ticks_position('none')
    if y:
        ax.yaxis.set_ticks_position('none')


def add_at_risk_counts(ax1, *kmfs):
    '''
    Add counts to plot
    '''
    # Create another axes where we can put size ticks
    ax2 = plt.twiny(ax=ax1)
    # Move the ticks below existing axes
    ax2_ypos = -0.15
    move_spines(ax2, ['bottom'], [ax2_ypos])
    # Hide all fluff
    remove_spines(ax2, ['top', 'right', 'bottom', 'left'])
    # Set ticks and labels on bottom
    ax2.xaxis.tick_bottom()
    # Match tick numbers and locations
    ax2.set_xlim(ax1.get_xlim())
    ax2.set_xticks(ax1.get_xticks())
    # Remove ticks, need to do this AFTER moving the ticks
    remove_ticks(ax2, x=True, y=True)
    # Add population size at times
    labels = []
    for tick in ax2.get_xticks():
        lbl = ""
        for kmf in kmfs:
            lbl += "\n{}".format(kmf.durations[kmf.durations >= tick].shape[0])
            # Last one add group name
            #if tick == ax2.get_xticks()[-1]:
            #    lbl += "  {}".format(kmf._label)

        labels.append(lbl.strip())

    # Align labels to the left so labels above don't change positions
    ax2.set_xticklabels(labels, ha = 'right')
    # Add a descriptive label, at a good position
    ax2.xaxis.set_label_coords(0, ax2_ypos)
    ax2.set_xlabel('At risk')


def plot_lifetimes(lifetimes, event_observed=None, birthtimes=None,
                   order=False):
    """
    Parameters:
      lifetimes: an (n,) numpy array of lifetimes.<commit_msg>Forgot tight layout call<commit_after># -*- coding: utf-8 -*-
from __future__ import print_function

# plotting
import numpy as np
from matplotlib import pyplot as plt

from lifelines.utils import coalesce


def remove_spines(ax, sides):
    '''
    Remove spines of axis.

    Parameters:
      ax: axes to operate on
      sides: list of sides: top, left, bottom, right

    Examples:
    removespines(ax, ['top'])
    removespines(ax, ['top', 'bottom', 'right', 'left'])
    '''
    for side in sides:
        ax.spines[side].set_visible(False)


def move_spines(ax, sides, dists):
    '''
    Move the entire spine relative to the figure.

    Parameters:
      ax: axes to operate on
      sides: list of sides to move. Sides: top, left, bottom, right
      dists: list of float distances to move. Should match sides in length.

    Example:
    move_spines(ax, sides=['left', 'bottom'], dists=[-0.02, 0.1])
    '''
    for side, dist in zip(sides, dists):
        ax.spines[side].set_position(('axes', dist))


def remove_ticks(ax, x=False, y=False):
    '''
    Remove ticks from axis.

    Parameters:
      ax: axes to work on
      x: if True, remove xticks. Default False.
      y: if True, remove yticks. Default False.

    Examples:
    removeticks(ax, x=True)
    removeticks(ax, x=True, y=True)
    '''
    if x:
        ax.xaxis.set_ticks_position('none')
    if y:
        ax.yaxis.set_ticks_position('none')


def add_at_risk_counts(ax1, *kmfs):
    '''
    Add counts to plot
    '''
    # Create another axes where we can put size ticks
    ax2 = plt.twiny(ax=ax1)
    # Move the ticks below existing axes
    ax2_ypos = -0.15
    move_spines(ax2, ['bottom'], [ax2_ypos])
    # Hide all fluff
    remove_spines(ax2, ['top', 'right', 'bottom', 'left'])
    # Set ticks and labels on bottom
    ax2.xaxis.tick_bottom()
    # Match tick numbers and locations
    ax2.set_xlim(ax1.get_xlim())
    ax2.set_xticks(ax1.get_xticks())
    # Remove ticks, need to do this AFTER moving the ticks
    remove_ticks(ax2, x=True, y=True)
    # Add population size at times
    labels = []
    for tick in ax2.get_xticks():
        lbl = ""
        for kmf in kmfs:
            lbl += "\n{}".format(kmf.durations[kmf.durations >= tick].shape[0])
            # Last one add group name
            #if tick == ax2.get_xticks()[-1]:
            #    lbl += "  {}".format(kmf._label)

        labels.append(lbl.strip())

    # Align labels to the left so labels above don't change positions
    ax2.set_xticklabels(labels, ha = 'right')
    # Add a descriptive label, at a good position
    ax2.xaxis.set_label_coords(0, ax2_ypos)
    ax2.set_xlabel('At risk')

    plt.tight_layout()


def plot_lifetimes(lifetimes, event_observed=None, birthtimes=None,
                   order=False):
    """
    Parameters:
      lifetimes: an (n,) numpy array of lifetimes.
{'fraction_non_alphanumeric': 0.12106835354067827, 'fraction_numerical': 0.010894394658232297, 'mean_word_length': 2.6463805253042922, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2810174', 'n_tokens_mistral': 1987, 'n_tokens_neox': 1909, 'n_words': 569}