Skip to content

Instantly share code, notes, and snippets.

Show Gist options
  • Save jmquintana79/c98bd0d0aeee8b2860ae1c9d6a481f59 to your computer and use it in GitHub Desktop.
Save jmquintana79/c98bd0d0aeee8b2860ae1c9d6a481f59 to your computer and use it in GitHub Desktop.
Esto permite un uso más flexible de la función *rolling*, tanto por poder aplicar una función custom como por poder usar más de un argumento para esta función (este ejemplo es para dos arguments pero es fácilmente modificable para otro número de argu
## custom function
def custom_function_example(x,y):
return x + y
## apply rolling on a custom function
def apply_rolling_function(df, window_size, func, column1, column2, verbose = False):
# validate arguments
assert column1 in df.columns.tolist()
assert column2 in df.columns.tolist()
# initialize
result = []
# loop
for i in range(len(df) - window_size + 1):
# collect data
window = df.iloc[i:i+window_size]
# apply function for this 2 columns
try:
result.append(func(window[column1], window[column2]))
except Exception as e:
result.append(np.nan)
if verbose:
print(f"Error in record {i}: {str(e)}")
#return pd.Series([np.nan] * (window_size - 1) + result).values
return pd.Series([result[0]] * (window_size - 1) + result).values
# windows size
window_size = 3
# usage
data[new_column] = apply_rolling_function(data, window_size, estimate_regression_r_value, column1, column2)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment