CODEWITHSUNDEEP
pythonpandas
Heathcliff
Heathcliff

Reputation: 3138

Create DataFrame from diagonal of other DataFrame

I have a covariance matrix for two stocks in a pandas DataFrame and would like to create a dataframe that has only the diagonal of that matrix.

The matrix looks like this:

        AAPL        GM
AAPL    0.079894    0.034113
GM      0.034113    0.060838

And I would like a DataFrame with indexes [AAPL, GM] and a column called covariance or something that would look like this:

        covariance      
AAPL    0.079894    
GM      0.060838

Is there a simple way to do this or will I end up writing for loops to fill other data structures to do it?

Thank you

Upvotes: 0

Views: 73

Answers (2)

Dani Mesejo
Dani Mesejo

Reputation: 61910

You could use lookup:

result = pd.DataFrame(index=df.index, data=df.lookup(df.index, df.columns), columns=['covariance'])
print(result)

Output

      covariance
AAPL    0.079894
GM      0.060838

Upvotes: 1

Erfan
Erfan

Reputation: 42886

Using np.diag to select the diagonal of the dataframe, then with with pd.DataFrame constructor to create your desired output:

df_diag = pd.DataFrame(data=np.diag(df), columns=['covariance'], index=df.index)

Output

      covariance
AAPL    0.079894
GM      0.060838

Upvotes: 3

Related Questions