plotnine.geoms.geom_segment¶

data = pd.read_csv('data/survey-social-media.csv')
data = (
    data >>
    rename(country='COUNTRY',
           gender='Q145',
           age='Q146',
           use_internet='Q70',
           use_social_media='Q74') >>
    select('PSRAID', 'gender', 'use_internet', drop=True)
)

data >> sample_n(10, random_state=123)

def toint(s):
    res = []
    for x in s:
        try:
            x = int(x)
        except ValueError:
            x = -1
        res.append(x)
    return res

def percent(gdf):
    return 100 * np.sum(gdf['use_social_media'] == 'Yes') / len(gdf)

no_age = ['Refused', 'Don\'t know']
yes_no = ['Yes', 'No']
ages = ['18-34', '35-49', '50+']

def first(s):
    return s.iloc[0]

rdata = (
    data
    >> define(age_='toint(age)')
    >> define(age_group=case_when([
        ('age_ == -1', '""'),
        ('age_ < 35', '"18-34"'),
        ('age_ < 50', '"35-49"'),
        ('age_ >= 50', '"50+"'),
    ]))
    >> select('age_', drop=True)
    >> group_by('country')
    >> define(country_count='n()')
    >> query('age_group in @ages')
    >> query('use_social_media in @yes_no')
    >> group_by('country', 'age_group')
    >> do(
        # social media use percentage
        sm_use_percent=lambda gdf: 100 * sum(gdf['use_social_media'] == 'Yes') / len(gdf),

        # social media question response rate
        smq_response_rate=lambda gdf: len(gdf) / gdf['country_count'].iloc[0]
    )
)

rdata >> ungroup() >> head(9)

n = 14

top = (
    rdata
    >> group_by('country')
    >> summarize(r='sum(smq_response_rate)')
    >> arrange('-r')
    >> head(n)
)

top_countries = top['country']

point_data = (
    rdata
    >> query('country in @top_countries')

    # Format the floating point data that will be plotted into strings
    >> define(sm_use_percent_str=if_else(
        'country == "France"',
        'format_sequence(sm_use_percent, "{:.0f}%")',
        'format_sequence(sm_use_percent, "{:.0f}")')
    )
)

point_data >> ungroup() >> head(6)

segment_data = (
    point_data
    >> group_by('country')
    >> summarize(min='min(sm_use_percent)', max='max(sm_use_percent)')
    >> define(gap='max-min')
    >> arrange('-gap')
    # Format the floating point data that will be plotted into strings
    >> define(
        min_str='format_sequence(min, "{:.0f}")',
        max_str='format_sequence(max, "{:.0f}")',
        gap_str='format_sequence(gap, "{:.0f}")',
    )
)

segment_data >> head()

# The right column (youngest-oldest gap) location
xgap = 112

(ggplot()
 # Range strip
 + geom_segment(
     segment_data,
     aes(x='min', xend='max', y='country', yend='country'),
     size=6,
     color='#a7a9ac'
 )
 # Age group markers
 + geom_point(
     point_data,
     aes('sm_use_percent', 'country', color='age_group', fill='age_group'),
     size=5,
     stroke=0.7,
 )
 # Age group percentages
 + geom_text(
     point_data >> query('age_group=="50+"'),
     aes(x='sm_use_percent-2', y='country', label='sm_use_percent_str', color='age_group'),
     size=8,
     ha='right',
 )
 + geom_text(
     point_data >> query('age_group=="35-49"'),
     aes(x='sm_use_percent+2', y='country', label='sm_use_percent_str'),
     size=8,
     ha='left',
     va='center',
     color='white'
 )
 + geom_text(
     point_data >> query('age_group=="18-34"'),
     aes(x='sm_use_percent+2', y='country', label='sm_use_percent_str', color='age_group'),
     size=8,
     ha='left',
 )
 # gap difference
 + geom_text(
     segment_data,
     aes(x=xgap, y='country', label='gap_str'),
     size=9,
     fontweight='bold',
     format_string='+{}'
 )
)

# The right column (youngest-oldest gap) location
xgap = 112

(ggplot()
 # Background Strips                                     # new
 + geom_segment(
     segment_data,
     aes(y='country', yend='country'),
     x=0, xend=100,
     size=8.5,
     color='#edece3'
 )
 # vertical grid lines along the strips                  # new
 + annotate(
     'segment',
     x=list(range(10, 100, 10)) * n,
     xend=list(range(10, 100, 10)) * n,
     y=np.tile(np.arange(1, n+1), 9)-.25,
     yend=np.tile(np.arange(1, n+1), 9) + .25,
     color='#CCCCCC'
 )
 # Range strip
 + geom_segment(
     segment_data,
     aes(x='min', xend='max', y='country', yend='country'),
     size=6,
     color='#a7a9ac'
 )
 # Age group markers
 + geom_point(
     point_data,
     aes('sm_use_percent', 'country', color='age_group', fill='age_group'),
     size=5,
     stroke=0.7,
 )
 # Age group percentages
 + geom_text(
     point_data >> query('age_group=="50+"'),
     aes(x='sm_use_percent-2', y='country', label='sm_use_percent_str', color='age_group'),
     size=8,
     ha='right',
 )
 + geom_text(
     point_data >> query('age_group=="35-49"'),
     aes(x='sm_use_percent+2', y='country', label='sm_use_percent_str'),
     size=8,
     ha='left',
     va='center',
     color='white'
 )
 + geom_text(
     point_data >> query('age_group=="18-34"'),
     aes(x='sm_use_percent+2', y='country', label='sm_use_percent_str', color='age_group'),
     size=8,
     ha='left',
 )
 # countries right-hand-size (instead of y-axis)         # new
 + geom_text(
     segment_data,
     aes(y='country', label='country'),
     x=-1,
     size=8,
     ha='right',
     fontweight='bold',
     color='#222222'
 )
 # gap difference
 + geom_vline(xintercept=xgap, color='#edece3', size=32)  # new
 + geom_text(
     segment_data,
     aes(x=xgap, y='country', label='gap_str'),
     size=9,
     fontweight='bold',
     format_string='+{}'
 )
 # Annotations                                            # new
 + annotate('text', x=31, y=n+1.1, label='50+', size=9, color='#ea9f2f', va='top')
 + annotate('text', x=56, y=n+1.1, label='35-49', size=9, color='#6d6e71', va='top')
 + annotate('text', x=85, y=n+1.1, label='18-34', size=9, color='#939c49', va='top')
 + annotate('text', x=xgap, y=n+.5, label='Youngest-\nOldest Gap', size=9, color='#444444', va='bottom', ha='center')
 + annotate('point', x=[31, 56, 85], y=n+.3, alpha=0.85, stroke=0)
 + annotate('segment', x=[31, 56, 85], xend=[31, 56, 85], y=n+.3, yend=n+.8, alpha=0.85)
 + annotate('hline', yintercept=[x+0.5 for x in range(2, n, 2)], alpha=.5, linetype='dotted', size=0.7)

 # Better spacing and color                              # new
 + scale_x_continuous(limits=(-18, xgap+2))
 + scale_y_discrete(expand=(0, 0, 0.1, 0))
 + scale_fill_manual(values=['#c3ca8c', '#d1d3d4', '#f2c480'])
 + scale_color_manual(values=['#939c49', '#6d6e71', '#ea9f2f'])
 + guides(color=None, fill=None)
 + theme_void()
 + theme(figure_size=(8, 7.5))
)

data = pd.read_csv(
    'data/API_SP.DYN.IMRT.IN_DS2_en_csv_v2/API_SP.DYN.IMRT.IN_DS2_en_csv_v2.csv',
    skiprows=[0, 1, 2])

# Columns as valid python variables
year_columns = {'y{}'.format(c): c for c in data.columns if c[:2] in {'19', '20'}}

data = (
    data
    >> rename({'country': 'Country Name', 'code': 'Country Code'})
    >> rename(year_columns)
    >> select('Indicator Name', 'Indicator Code', 'Unnamed: 61', drop=True)
)

data >> head()

metadata = pd.read_csv(
    'data/API_SP.DYN.IMRT.IN_DS2_en_csv_v2/Metadata_Country_API_SP.DYN.IMRT.IN_DS2_en_csv_v2.csv'
)

metadata = (
    metadata
    >> rename({'code': 'Country Code',
               'region': 'Region',
               'income_group': 'IncomeGroup'
               })
    >> select('code', 'region', 'income_group')
)

cat_order = ['High income', 'Upper middle income', 'Lower middle income', 'Low income']
metadata['income_group'] = pd.Categorical(metadata['income_group'], categories=cat_order, ordered=True)

metadata >> head(10)

# Drop the aggregates from
country_metadata = metadata >> call('.dropna', subset=['region'])
country_metadata >> head(10)

country_data = inner_join(data, country_metadata, on='code')
country_data >> head()

sampled_data = (
    country_data
    >> call('.dropna', subset=['y1960', 'y2015'])
    >> sample_n(35, random_state=123)
    >> define(
        y1960_rank='y1960.rank(method="min").astype(int)',
        y2015_rank='y2015.rank(method="min").astype(int)'
    )
)

sampled_data >> head()

(ggplot(sampled_data)
 + geom_text(aes(1, 'y1960_rank', label='country'), ha='right', size=9)
 + geom_text(aes(2, 'y2015_rank', label='country'), ha='left', size=9)
 + geom_point(aes(1, 'y1960_rank', color='income_group'), size=2.5)
 + geom_point(aes(2, 'y2015_rank', color='income_group'), size=2.5)
 + geom_segment(aes(x=1, y='y1960_rank', xend=2, yend='y2015_rank', color='income_group'))
)

# Text colors
black1 = '#252525'
black2 = '#222222'

(ggplot(sampled_data)
 # Slight modifications for the original lines,
 # 1. Nudge the text to either sides of the points
 # 2. Alter the color and alpha values
 + geom_text(aes(1, 'y1960_rank', label='country'), nudge_x=-0.05, ha='right', size=9, color=black1)
 + geom_text(aes(2, 'y2015_rank', label='country'), nudge_x=0.05, ha='left', size=9, color=black1)
 + geom_point(aes(1, 'y1960_rank', color='income_group'), size=2.5, alpha=.7)
 + geom_point(aes(2, 'y2015_rank', color='income_group'), size=2.5, alpha=.7)
 + geom_segment(aes(x=1, y='y1960_rank', xend=2, yend='y2015_rank', color='income_group'), alpha=.7)

 # Text Annotations
 + annotate('text', x=1, y=0, label='Rank in 1960', fontweight='bold', ha='right', size=10, color=black2)
 + annotate('text', x=2, y=0, label='Rank in 2015', fontweight='bold', ha='left', size=10, color=black2)
 + annotate('text', x=1.5, y=0, label='Lines show change in rank', size=9, color=black1)
 + annotate('label', x=1.5, y=3, label='Lower infant\ndeath rates', size=9, color=black1,
            label_size=0, fontstyle='italic')
 + annotate('label', x=1.5, y=33, label='Higher infant\ndeath rates', size=9, color=black1,
            label_size=0, fontstyle='italic')

 # Prevent country names from being chopped off
 + lims(x=(0.35, 2.65))
 + labs(color='Income Group')
 # Countries with lower rates on top
 + scale_y_reverse()
 # Change colors
 + scale_color_brewer(type='qual', palette=2)
 # Removes all decorations
 + theme_void()
 # Changing the figure size prevents the country names from squishing up
 + theme(figure_size=(8, 11))
)