Fixes for Volume Profile

pandas_ta/volume/vp.py

@@ -1,6 +1,7 @@
 # -*- coding: utf-8 -*-
 from numpy import array_split
-from pandas import concat, DataFrame
+from numpy import mean
+from pandas import cut, concat, DataFrame
 from pandas_ta.utils import signed_series, verify_series
 
 
@@ -15,16 +16,17 @@ def vp(close, volume, width=None, **kwargs):
     if close is None or volume is None: return
 
     # Setup
-    signed_volume = signed_series(volume, initial=1)
-    pos_volume = signed_volume[signed_volume > 0] * volume
-    neg_volume = signed_volume[signed_volume < 0] * -volume
+    signed_price = signed_series(close, initial=1)
+    pos_volume = signed_price[signed_price > 0] * volume
+    pos_volume.name = volume.name
+    neg_volume = signed_price[signed_price < 0] * -volume
+    neg_volume.name = volume.name
     vp = concat([close, pos_volume, neg_volume], axis=1)
 
     close_col = f"{vp.columns[0]}"
     high_price_col = f"high_{close_col}"
     low_price_col = f"low_{close_col}"
     mean_price_col = f"mean_{close_col}"
-    mid_price_col = f"mid_{close_col}"
 
     volume_col = f"{vp.columns[1]}"
     pos_volume_col = f"pos_{volume_col}"
@@ -34,19 +36,28 @@ def vp(close, volume, width=None, **kwargs):
 
     # sort_close: Sort by close before splitting into ranges. Default: False
     # If False, it sorts by date index or chronological versus by price
+
     if sort_close:
-        vp.sort_values(by=[close_col], inplace=True)
-
-    # Calculate Result
-    vp_ranges = array_split(vp, width)
-    result = ({
-        low_price_col: r[close_col].min(),
-        mean_price_col: r[close_col].mean(),
-        high_price_col: r[close_col].max(),
-        pos_volume_col: r[pos_volume_col].sum(),
-        neg_volume_col: r[neg_volume_col].sum(),
-    } for r in vp_ranges)
-    vpdf = DataFrame(result)
+        vp[mean_price_col] = vp[close_col]
+        vpdf = vp.groupby(cut(vp[close_col], width, include_lowest=True, precision=2)).agg({
+            mean_price_col: mean,
+            pos_volume_col: sum,
+            neg_volume_col: sum,
+        })
+        vpdf[low_price_col] = [x.left for x in vpdf.index]
+        vpdf[high_price_col] = [x.right for x in vpdf.index]
+        vpdf = vpdf.reset_index(drop=True)
+        vpdf = vpdf[[low_price_col, mean_price_col, high_price_col, pos_volume_col, neg_volume_col]]
+    else:
+        vp_ranges = array_split(vp, width)
+        result = ({
+            low_price_col: r[close_col].min(),
+            mean_price_col: r[close_col].mean(),
+            high_price_col: r[close_col].max(),
+            pos_volume_col: r[pos_volume_col].sum(),
+            neg_volume_col: r[neg_volume_col].sum(),
+        } for r in vp_ranges)
+        vpdf = DataFrame(result)
     vpdf[total_volume_col] = vpdf[pos_volume_col] + vpdf[neg_volume_col]
 
     # Handle fills
@@ -79,8 +90,12 @@ def vp(close, volume, width=None, **kwargs):
         width=10
 
     vp = pd.concat([close, pos_volume, neg_volume], axis=1)
-    vp_ranges = np.array_split(vp, width)
-    result = ({high_close, low_close, mean_close, neg_volume, pos_volume} foreach range in vp_ranges)
+    if sort_close:
+        vp_ranges = cut(vp[close_col], width)
+        result = ({range_left, mean_close, range_right, pos_volume, neg_volume} foreach range in vp_ranges
+    else:
+        vp_ranges = np.array_split(vp, width)
+        result = ({low_close, mean_close, high_close, pos_volume, neg_volume} foreach range in vp_ranges
     vpdf = pd.DataFrame(result)
     vpdf['total_volume'] = vpdf['pos_volume'] + vpdf['neg_volume']