CODEWITHSUNDEEP
pythonpandasdataframe
Andrew Drake
Andrew Drake

Reputation: 665

Dataframe: Mask with Previous Value in Same Column that is not NaN

I have the following dataframe:

Trajectory Direction Resulting_Direction
STRAIGHT   NORTH     NORTH
STRAIGHT   NaN       NORTH
LEFT       NaN       WEST
LEFT       NaN       WEST
LEFT       NaN       WEST
STRAIGHT   NaN       WEST
STRAIGHT   NaN       WEST
RIGHT      NaN       NORTH
RIGHT      NaN       NORTH
RIGHT      NaN       NORTH

My goal is to have direction change whenever I encounter three straight trajectories. So in this example, my new column would be Resulting_Direction (assume it is not initially in the df).

Currently I am doing this by doing row-by-row if-statements. However this is painfully slow and inefficient. I wish to use a mask to set the resulting direction in rows where it turns, then use fillna(method="ffill"). This is my attempt:

df.loc[:,'direction'] = np.NaN
df.loc[df.index == 0, "direction"] = "WEST"
# mask is for finding when a signal hasnt changed in three seconds, but now has
mask = (df.trajectory != df.trajectory.shift(1)) & (df.trajectory == df.trajectory.shift(-1)) & (df.trajectory == df.trajectory.shift(-2))
df.loc[(mask) & (df['trajectory'] == 'LEFT') & (df['direction'].dropna().shift() == "WEST"),'direction'] = 'SOUTH'
df.loc[(mask) & (df['trajectory'] == 'LEFT') & (df['direction'].dropna().shift() == "SOUTH"),'direction'] = 'EAST'
df.loc[(mask) & (df['trajectory'] == 'LEFT') & (df['direction'].dropna().shift() == "EAST"),'direction'] = 'NORTH'
df.loc[(mask) & (df['trajectory'] == 'LEFT') & (df['direction'].dropna().shift() == "NORTH"),'direction'] = 'WEST'
df.loc[(mask) & (df['trajectory'] == 'RIGHT') & (df['direction'].dropna().shift() == "WEST"),'direction'] = 'NORTH'
df.loc[(mask) & (df['trajectory'] == 'RIGHT') & (df['direction'].dropna().shift() == "SOUTH"),'direction'] = 'WEST'
df.loc[(mask) & (df['trajectory'] == 'RIGHT') & (df['direction'].dropna().shift() == "EAST"),'direction'] = 'SOUTH'
df.loc[(mask) & (df['trajectory'] == 'RIGHT') & (df['direction'].dropna().shift() == "NORTH"),'direction'] = 'EAST'
df.loc[:,'direction'] = df.direction.fillna(method="ffill")
print(df[['trajectory','direction']])

I believe my issue is in df['direction'].dropna().shift(). How do I find the previous value in the same column that is not NaN?

Upvotes: 2

Views: 335

Answers (1)

Quang Hoang
Quang Hoang

Reputation: 150735

IIUC, the problem is to detect where the direction change, supposedly at the beginning of 3 consecutive change commands:

thresh = 3
# mark the consecutive direction commands
blocks = df.Trajectory.ne(df.Trajectory.shift()).cumsum()


# group by blocks
groups = df.groupby(blocks)

# enumerate each block
df['mask'] = groups.cumcount()

# shift up to mark the beginning
# mod thresh to divide each block into small block of thresh
df['mask'] = groups['mask'].shift(1-thresh) % thresh

# for conversion of direction to letters:
changes = {'LEFT': -1,'RIGHT':1}

# all the directions
directions = ['NORTH', 'EAST', 'SOUTH', 'WEST']

# update directions according to the start direction
start = df['Direction'].iloc[0]
start_idx = directions.index(start)
directions = {k%4: v for k,v in enumerate(directions, start=start_idx)}


# update direction changes
direction_changes = (df.Trajectory
                     .where(df['mask'].eq(2))   # where the changes happends
                     .map(changes)              # replace the changes with number
                     .fillna(0)                 # where no direction change is 0
                    )
# mod 4 for the 4 direction
# and map
df['Resulting_Direction'] = (direction_changes.cumsum() % 4).map(directions)

Output:

  Trajectory Direction Resulting_Direction  mask
0   STRAIGHT     NORTH               NORTH   NaN
1   STRAIGHT       NaN               NORTH   NaN
2       LEFT       NaN                WEST   2.0
3       LEFT       NaN                WEST   NaN
4       LEFT       NaN                WEST   NaN
5   STRAIGHT       NaN                WEST   NaN
6   STRAIGHT       NaN                WEST   NaN
7      RIGHT       NaN               NORTH   2.0
8      RIGHT       NaN               NORTH   NaN
9      RIGHT       NaN               NORTH   NaN

Upvotes: 1

Related Questions