Selecting correct value in Numpy matrix inside dataframe-CodePudding

I have a dataframe with mixed types. I have the columns, x, y, z, matrix, where matrix is a numpy matrix

10 10 10 [[1,2,3],[4,5,6],[7,8,9]]
12 12 12 [[1,2,3],[4,5,6],[7,8,9]]
14 14 14 [[1,2,3],[4,5,6],[7,8,9]]

I want to take the value of the matrix in index row 1 column index 1 (5) for plotting together with x, y, z

I have tried the following, but it does not work

print(df["matrix"][:][1,1])

CodePudding user response：

new = df_max_timestep
tangential_stress = []
for value in df_max_timestep["rotated stress tensor"]:
    tangential_stress.append(value[1,1])
new["Tangential stress"] = tangential_stress

CodePudding user response：

Try to get the matrix elements by string accessor .str[], as follows:

df['matrix'].str[1].str[1]

Result:

0    5
1    5
2    5
Name: matrix, dtype: int64

Data Setup

data = {'x': [10, 12, 14],
 'y': [10, 12, 14],
 'z': [10, 12, 14],
 'matrix': [np.array([[1,2,3],[4,5,6],[7,8,9]]),
  np.array([[1,2,3],[4,5,6],[7,8,9]]),
  np.array([[1,2,3],[4,5,6],[7,8,9]])]}

df = pd.DataFrame(data)

print(df)

    x   y   z                             matrix
0  10  10  10  [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
1  12  12  12  [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
2  14  14  14  [[1, 2, 3], [4, 5, 6], [7, 8, 9]]

CodePudding user response：

You can first stack the arrays. Then you will be able to then handle them as numpy arrays:

import pandas as pd
import numpy as np

df = pd.DataFrame([[10, 10, 10, np.array([[1,2,3],[4,5,6],[7,8,9]])],
                   [12, 12, 12, np.array([[1,2,3],[4,5,6],[7,8,9]])],
                   [14, 14, 14, np.array([[1,2,3],[4,5,6],[7,8,9]])],],
                 columns=['x', 'y', 'z', 'matrix']
                 )

np.stack(df['matrix'])[:, 1,1]

output:

array([5, 5, 5])

example of assignment as new column:

>>> df['new'] = np.stack(df['matrix'])[:, 1,1]
>>> df
    x   y   z                             matrix  new
0  10  10  10  [[1, 2, 3], [4, 5, 6], [7, 8, 9]]    5
1  12  12  12  [[1, 2, 3], [4, 5, 6], [7, 8, 9]]    5
2  14  14  14  [[1, 2, 3], [4, 5, 6], [7, 8, 9]]    5

plotting

import plotly.express as px

df['matrix_center'] = np.stack(df['matrix'])[:, 1,1]
px.scatter_3d(df, x='x', y='y', z='z', color='matrix_center')