data_prep_create_loc_df.py

def create_location_dataframe(data):
    """Return a dataframe with the proper attributes needed for the plot."""
    origin_df = data[['OriginCity', 'OriginState', 'OriginCountry',
                      'OriginLat', 'OriginLong']].drop_duplicates()
    # If we have all the states, use them. Otherwise, use the countries
    state_or_country = (origin_df['OriginState'] if not origin_df['OriginState'].isnull().values.any()
                        else origin_df['OriginCountry'])
    origin_df['text'] = origin_df['OriginCity'] + '-' + state_or_country
    origin_df['size'] = 10
    origin_df['color'] = 'blue'
    origin_df['shape'] = 'square'
    origin_df.rename(columns={'OriginLat': 'lat', 'OriginLong': 'lon'}, inplace=True)
    # we'll concat origin_df with dest_df. Either rename and remove 'Origin' or drop 
    origin_df.drop(['OriginCity', 'OriginState', 'OriginCountry'], axis=1, inplace=True)

    dest_df = data[['DestinationCity', 'DestinationState', 'DestinationCountry',
                    'DestinationLat', 'DestinationLong']].drop_duplicates()
    state_or_country = (dest_df['DestinationState'] if not dest_df['DestinationState'].isnull().values.any()
                        else dest_df['DestinationCountry'])
    dest_df['text'] = dest_df['DestinationCity'] + '-' + state_or_country
    dest_df['size'] = 8
    dest_df['color'] = 'orange'
    dest_df['shape'] = 'circle'
    dest_df.rename(columns={'DestinationLat': 'lat', 'DestinationLong': 'lon'}, inplace=True)
    # we'll concat origin_df with dest_df. Either rename and remove 'Destination' or drop 
    dest_df.drop(['DestinationCity', 'DestinationState', 'DestinationCountry'], axis=1, inplace=True)

    df = pd.concat([origin_df, dest_df], sort=False)
    return df