inequality.py

import psycopg2
import itertools
import pandas
import numpy as np
import matplotlib.pyplot as plt

def lorenz_curve(X):
    X = np.array(X, dtype='int64')
    X.sort()
    X_lorenz = X.cumsum() / X.sum()
    X_lorenz = np.insert(X_lorenz, 0, 0) 
    X_lorenz[0], X_lorenz[-1]
    fig, ax = plt.subplots(figsize=[6,6])
    ## scatter plot of Lorenz curve
    ax.scatter(np.arange(X_lorenz.size)/(X_lorenz.size-1), X_lorenz, 
               marker='x', color='darkgreen', s=10)
    ## line plot of equality
    ax.plot([0,1], [0,1], color='k')
    
def inequality_per_year(year, column):
    cur.execute("""
                SELECT id, play_count, share_count, comment_count, like_count, create_time
    	FROM public.tiktok_normalized
    where date_part('year', create_time) = (%s)
    and representative
                """, [year])
    result=cur.fetchall()
    data = [row[column] for row in result]
    lorenz_curve(data)
    
    return coefficient

conn=psycopg2.connect('dbname=postgres user=postgres')
cur = conn.cursor()
inequality_per_year(2020, 1)
cur.close()
conn.close()