Skip to content

Instantly share code, notes, and snippets.

@TonyBogomoloff
Last active January 20, 2020 15:28
Show Gist options
  • Save TonyBogomoloff/1a1ae9a0dd718d198c312adfa36966eb to your computer and use it in GitHub Desktop.
Save TonyBogomoloff/1a1ae9a0dd718d198c312adfa36966eb to your computer and use it in GitHub Desktop.
A clean example of markBlocks()
import numpy as np, pandas as pd
def markBlocks(data, condition, blockName):
df = data.copy()
### count the length of periods where the condition is satisfied
df['flag'] = np.where(condition, 1, 0)
## give a consequetive number to each step, where the the condition is satisfied
df['counter_' + blockName] = df.groupby((df['flag'] != df['flag'].shift(1)).cumsum()).cumcount()+1
## drop the counter to zero when the condition is not satisfied
df['counter_' + blockName] = np.where(df.flag == 0, 0, df['counter_' + blockName])
## get the length of each of the parts, where the condition is satisfied
df['length_' + blockName] = np.where(df.flag.diff().shift(-1) != 0, df['counter_' + blockName], np.nan)
## propagate it to the whole parts
df['length_' + blockName] = df['length_' + blockName].bfill().astype(int)
## number each block
df[blockName + '_block'] = (df['flag'] != df['flag'].shift(1)).astype(int).cumsum()
df[blockName + '_block'] = np.where(df.flag == 0, 0, df[blockName + '_block'])
df[blockName + '_block'] = np.where(df[blockName + '_block'] != 0, (df[blockName + '_block'] - 1)//2 + 1, 0)
# df = df[df.flag == 1]
# df = df.drop(columns=['flag'])
return df
data = dict(y=[1, 1, 1, 0, 0, 1, 0, 0, 0, 0, 1, 1, 0, 0, 0, 1, 1, 1, 1, 0, 1])
df = pd.DataFrame(data)
df = markBlocks(df, df.y == 1, 'ones')
df.head(50)
@TonyBogomoloff
Copy link
Author

Output:

index   y	flag	counter_ones	length_ones	ones_block
0	1	1	1	3	1
1	1	1	2	3	1
2	1	1	3	3	1
3	0	0	0	0	0
4	0	0	0	0	0
5	1	1	1	1	2
6	0	0	0	0	0
7	0	0	0	0	0
8	0	0	0	0	0
9	0	0	0	0	0
10	1	1	1	2	3
11	1	1	2	2	3
12	0	0	0	0	0
13	0	0	0	0	0
14	0	0	0	0	0
15	1	1	1	4	4
16	1	1	2	4	4
17	1	1	3	4	4
18	1	1	4	4	4
19	0	0	0	0	0
20	1	1	1	1	5

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment