Automate the def function in pandas for correlation

Question

def correlations(group):
    # Remove columns
    group = (
        group.reset_index()
        .drop(columns=["RNA1_Approved_Symbol"], level=0)
        .set_index("RNA1_ID")
    )
    # compute correlations
    correlation = group.T.corr()

    # rename pair column
    correlation = correlation.reset_index().rename(
        columns={"RNA1_ID": "RNA1_ID_new"}
    )

    # Remove NAs
    corr_result = (
        correlation.melt(value_name="Correlation", id_vars=["RNA1_ID_new"])
        .set_index(["RNA1_ID", "RNA1_ID_new"])
        .dropna()
    )
    return corr_result


def correlation_1(data):
    ## Create Pivot table
    pivot_library = data.pivot(
        index=["RNA1_Approved_Symbol", "RNA1_ID"], columns=["RNA2_ID"]
    )

    # Measure correlations between four id's
    return (
        pivot_library.groupby("RNA1_Approved_Symbol")
        .apply(correlations)
        .reset_index()
        .dropna()
    )

I am using the above function to calculate the correlation

data = df1
correlations = correlation_1(df1)

I want to automate the above function when I am using different data frame where there is RNA2_Approved_Symbol instead of RNA1_Approved_Symbol and RNA2_ID instead of RNA1_ID. Like using index 1 and index 2.

Answer 1

If I have understood the question is to generalise the column names to permit a function to accept dataframes where you don't have the same column name, or it's different from the previous dataframe column names. Here the column could be called RNA1_Approved_Symbol or RNA2_Approved_Symbol in addition a different column could be called RNA2_ID or RNA_ID. No problem providing the column order is identical for the columns in question. If the orders are not identical see the code below the ---- to correct that.

If the order is identical the simple answer is df.columns[i]. Thus it's obtaining the column name on its position within the dataframe (columns command), where i is the column location.

Example, to rename RNA1_ID to RNA1_ID_new, or RNA2_ID to RNA2_ID_new

import re

df1 = pd.DataFrame({"perch": [5,6,7], "RNA1_ID": [1,2,4], "roach": [30,42,51]})
df2 = pd.DataFrame({"perch": [5,6,7], "RNA2_ID": [1,2,4], "roach": [30,42,51]})

def renamecolumn(df):
    myname = df.columns[1]
    newname = re.sub(r'$','_new', myname)
    df = df.rename(columns={myname: newname})
    return df

print (renamecolumn(df1))
print (renamecolumn(df2))

Output,

-	perch	RNA1_ID_new	roach
0	5	1	30
1	6	2	42
2	7	4	51

-	perch	RNA2_ID_new	roach
0	5	1	30
1	6	2	42
2	7	4	51

I noted there was a fair number of set_index commands so the earlier the column name can be trapped inside a variable the better, because set_index will change the column order.

---

Alternatively, the columns can be re-arranged to be in a given order if the original order between df1 and df2 is different. Here I'll move the column to the front.

mycols = df1.columns.tolist()
mycols = mycols[-2:] + mycols[:-2]
df1 = df1[mycols]
print (df1)

-	RNA1_ID	roach	perch
0	1	30	5
1	2	42	6
2	4	51	7

Thus RNA1_ID has been moved from the second to the first column. If df2 was in the last column repeat as above use -1 instead of -2 and both columns will now be at the start of the dataframe. Modify the core dataframe generalisation renaming accordingly, viz.

myname = df.columns[0]

... sorted.