2
\$\begingroup\$

I defined a function that generates tables and creates countplots and barplots based on the arguments that it receives. But the snippets of code inside the function are repeated several times, which looks bad. Is there a way to reduce the repetition?

The code was run in a Jupyter Notebook. Here is the Titanic dataset link

data

import os
print(os.listdir("../input"))
from IPython.display import display  # to use display
import numpy as np
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
%matplotlib inline  

# because I am using Jupyter Notebook
# data frame I am using is the Titanic dataset on kaggle.com
# use the train.csv file, it has a column called 'Survived'

# Read in the dataset
titanic_train = pd.read_csv('../input/train.csv')

# Define the count_n_plot function
def count_n_plot(df, col_name, countsplit = None, bar = False, barsplit = None):

    """
    Creates countplots and barplots of the specified feature 
    (with options to split the columns) and generates the 
    corresponding table of counts and percentages.

    Parameters
    ----------
    df : DataFrame
        Dataset for plotting.
    col_name : string
        Name of column/feature in "data".
    countsplit : string
        Use countsplit to specify the "hue" argument of the countplot.
    bar : Boolean
        If True, a barplot of the column col_name is created, showing
        the fraction of survivors on the y-axis.
    barsplit: string
        Use barsplit to specify the "hue" argument of the barplot.
    """

    if (countsplit != None) & bar & (barsplit != None):
        col_count1 = df[[col_name]].groupby(by = col_name).size()
        col_perc1 = col_count1.apply(lambda x: x / sum(col_count1) * 100).round(1)
        tcount1 = pd.DataFrame({'Count': col_count1, 'Percentage': col_perc1})

        col_count2 = df[[col_name,countsplit]].groupby(by = [col_name,countsplit]).size()
        col_perc2 = col_count2.apply(lambda x: x / sum(col_count2) * 100).round(1)
        tcount2 = pd.DataFrame({'Count': col_count2, 'Percentage': col_perc2})


        display(tcount1, tcount2) #, tbar1, tbar2)

        figc, axc = plt.subplots(1, 2, figsize = (10,4))
        sns.countplot(data = df, x = col_name, hue = None, ax = axc[0])
        sns.countplot(data = df, x = col_name, hue = countsplit, ax = axc[1])

        figb, axb = plt.subplots(1, 2, figsize = (10,4))
        sns.barplot(data = df, x = col_name, y = 'Survived', hue = None, ax = axb[0])
        sns.barplot(data = df, x = col_name, y = 'Survived', hue = barsplit, ax = axb[1])

    elif (countsplit != None) & bar:
        col_count1 = df[[col_name]].groupby(by = col_name).size()
        col_perc1 = col_count1.apply(lambda x: x / sum(col_count1) * 100).round(1)
        tcount1 = pd.DataFrame({'Count': col_count1, 'Percentage': col_perc1})

        col_count2 = df[[col_name,countsplit]].groupby(by = [col_name,countsplit]).size()
        col_perc2 = col_count2.apply(lambda x: x / sum(col_count2) * 100).round(1)
        tcount2 = pd.DataFrame({'Count': col_count2, 'Percentage': col_perc2})


        display(tcount1, tcount2) #, tbar1)

        fig, axes = plt.subplots(1, 3, figsize = (15,4))
        sns.countplot(data = df, x = col_name, hue = None, ax = axes[0])
        sns.countplot(data = df, x = col_name, hue = countsplit, ax = axes[1])
        sns.barplot(data = df, x = col_name, y = 'Survived', hue = None, ax = axes[2])

    elif countsplit != None:
        col_count1 = df[[col_name]].groupby(by = col_name).size()
        col_perc1 = col_count1.apply(lambda x: x / sum(col_count1) * 100).round(1)
        tcount1 = pd.DataFrame({'Count': col_count1, 'Percentage': col_perc1})

        col_count2 = df[[col_name,countsplit]].groupby(by = [col_name,countsplit]).size()
        col_perc2 = col_count2.apply(lambda x: x / sum(col_count2) * 100).round(1)
        tcount2 = pd.DataFrame({'Count': col_count2, 'Percentage': col_perc2})
        display(tcount1, tcount2)

        fig, axes = plt.subplots(1, 2, figsize = (10,4))
        sns.countplot(data = df, x = col_name, hue = None, ax = axes[0])
        sns.countplot(data = df, x = col_name, hue = countsplit, ax = axes[1])

    else:
        col_count = df[[col_name]].groupby(by = col_name).size()
        col_perc = col_count.apply(lambda x: x / sum(col_count) * 100).round(1)
        tcount1 = pd.DataFrame({'Count': col_count, 'Percentage': col_perc})
        display(tcountl)

        sns.countplot(data = df, x = col_name)
\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.