Сюжет Sankey имеет собственные циклы из-за идентичных меток в нескольких столбцах - PullRequest
Новая
       22

Сюжет Sankey имеет собственные циклы из-за идентичных меток в нескольких столбцах

0 голосов
/ 24 апреля 2020

Рассмотрим фрейм данных: 

animal  animal2  size    count
dog     dog  small     2
dog     cat  large     3
cat     dog  small     1
dog     pig  large     5
cat     cat  large     3
pig     dog  small     9
pig     cat  large     2 
cat     pig  large     3

Я хочу построить график Санки с 3 столбцами, который показывает потоки между категориями (каждая строка в Санки будет числом раз, которое пара элементов, каждый из одного столбца, появляются вместе).

Этот код, который я нашел, кажется, работает, но он имеет много циклов, так как у меня есть похожие категории в нескольких столбцах: 

def genSankey(df,cat_cols=[],value_cols='',title='Sankey Diagram'):
    # maximum of 6 value cols -> 6 colors
    colorPalette = ['#4B8BBE','#306998','#FFE873','#FFD43B','#646464']
    labelList = []
    colorNumList = []
    for catCol in cat_cols:
        labelListTemp =  list(set(df[catCol].values))
        colorNumList.append(len(labelListTemp))
        labelList = labelList + labelListTemp

    # remove duplicates from labelList
    labelList = list(dict.fromkeys(labelList))

    # define colors based on number of levels
    colorList = []
    for idx, colorNum in enumerate(colorNumList):
        colorList = colorList + [colorPalette[idx]]*colorNum

    # transform df into a source-target pair
    for i in range(len(cat_cols)-1):
        if i==0:
            sourceTargetDf = df[[cat_cols[i],cat_cols[i+1],value_cols]]
            sourceTargetDf.columns = ['source','target','count']
        else:
            tempDf = df[[cat_cols[i],cat_cols[i+1],value_cols]]
            tempDf.columns = ['source','target','count']
            sourceTargetDf = pd.concat([sourceTargetDf,tempDf])
        sourceTargetDf = sourceTargetDf.groupby(['source','target']).agg({'count':'sum'}).reset_index()

    # add index for source-target pair
    sourceTargetDf['sourceID'] = sourceTargetDf['source'].apply(lambda x: labelList.index(x))
    sourceTargetDf['targetID'] = sourceTargetDf['target'].apply(lambda x: labelList.index(x))

    # creating the sankey diagram
    data = dict(
        type='sankey',
        node = dict(
          pad = 15,
          thickness = 20,
          line = dict(
            color = "black",
            width = 1.5   # was 0.5
          ),
          label = labelList,
          color = colorList
        ),
        link = dict(
          source = sourceTargetDf['sourceID'],
          target = sourceTargetDf['targetID'],
          value = sourceTargetDf['count']
        )
      )

    layout =  dict(
        title = title,
        font = dict(
          size = 20    # was 10
        )
    )

    fig = dict(data=[data], layout=layout)
    return fig

который может быть запущен как: 

import pandas as pd
import plotly
import chart_studio.plotly as py


fig = genSankey(df,cat_cols=['animal1','animal2','size'],value_cols='count',title='Animal List')
plotly.offline.plot(fig, validate=False)

Есть ли что-то, что я могу просто изменить в этой функции, чтобы перестать получать петли?

...