Enhance price shock cache generation and API functionality:

- Added support for filtering by pool ID in the price shock API, allowing users to specify which pool's data to retrieve.
- Enhanced user metrics calculations to filter users based on the specified pool ID, improving data accuracy.
- Updated the frontend to allow users to select a pool ID, with corresponding changes to cache key generation for better data management.
- Modified the `generate_ucache.py` script to accept and process the new `pool_id` parameter, ensuring it is included in the API requests.
- Updated `gen.sh` to include comprehensive logging for price shock cache generation, detailing asset groups, scenarios, and pool IDs.

These changes improve the flexibility and usability of the price shock features, enabling more targeted analysis and reporting.

Author: goldhaxx

backend/api/price_shock.py

@@ -15,6 +15,7 @@ async def _get_price_shock(
     oracle_distortion: float = 0.1,
     asset_group: str = PriceShockAssetGroup.IGNORE_STABLES.value,
     n_scenarios: int = 5,
+    pool_id: int = None,
 ) -> dict:
     asset_group = asset_group.replace("+", " ")
     price_shock_asset_group = PriceShockAssetGroup(asset_group)
@@ -26,6 +27,7 @@ async def _get_price_shock(
         oracle_distortion=oracle_distortion,
         asset_group=price_shock_asset_group,
         n_scenarios=n_scenarios,
+        pool_id=pool_id,
     )
 
 
@@ -35,6 +37,7 @@ async def get_price_shock(
     oracle_distortion: float = 0.1,
     asset_group: str = PriceShockAssetGroup.IGNORE_STABLES.value,
     n_scenarios: int = 5,
+    pool_id: int = None,
 ):
     return await _get_price_shock(
         request.state.backend_state.last_oracle_slot,
@@ -43,4 +46,5 @@ async def get_price_shock(
         oracle_distortion,
         asset_group,
         n_scenarios,
+        pool_id,
     )

backend/scripts/generate_ucache.py

@@ -50,23 +50,33 @@ class PriceShockEndpoint(Endpoint):
     asset_group: PriceShockAssetGroup
     oracle_distortion: float
     n_scenarios: int
+    pool_id: int
 
     def __init__(
         self,
         asset_group: PriceShockAssetGroup,
         oracle_distortion: float,
         n_scenarios: int,
+        pool_id: int = None,
     ):
         self.asset_group = asset_group
         self.oracle_distortion = oracle_distortion
         self.n_scenarios = n_scenarios
+        self.pool_id = pool_id
+        
+        params = {
+            "asset_group": asset_group,
+            "oracle_distortion": oracle_distortion,
+            "n_scenarios": n_scenarios,
+        }
+        
+        # Only add pool_id if it's not None
+        if pool_id is not None:
+            params["pool_id"] = pool_id
+            
         super().__init__(
             "price-shock/usermap",
-            {
-                "asset_group": asset_group,
-                "oracle_distortion": oracle_distortion,
-                "n_scenarios": n_scenarios,
-            },
+            params,
         )
 
 
@@ -99,6 +109,7 @@ async def mock_call_next(request):
                     oracle_distortion=query_params["oracle_distortion"],
                     asset_group=query_params["asset_group"],
                     n_scenarios=query_params["n_scenarios"],
+                    pool_id=query_params.get("pool_id"),
                 )
 
             if endpoint == "asset-liability/matrix":
@@ -167,6 +178,7 @@ async def main():
     ps_parser.add_argument("--asset-group", type=str, required=True)
     ps_parser.add_argument("--oracle-distortion", type=float, required=True)
     ps_parser.add_argument("--n-scenarios", type=int, required=True)
+    ps_parser.add_argument("--pool-id", type=int, help="Filter by pool ID (optional)")
 
     args = parser.parse_args()
 
@@ -191,6 +203,7 @@ async def main():
                 asset_group=args.asset_group,
                 oracle_distortion=args.oracle_distortion,
                 n_scenarios=args.n_scenarios,
+                pool_id=getattr(args, 'pool_id', None),
             )
         )
 
@@ -203,3 +216,4 @@ async def main():
 # Usage example:
 # python -m backend.scripts.generate_ucache --use-snapshot asset-liability --mode 0 --perp-market-index 0
 # python -m backend.scripts.generate_ucache --use-snapshot price-shock --asset-group "ignore+stables" --oracle-distortion 0.05 --n-scenarios 5
+# python -m backend.scripts.generate_ucache --use-snapshot price-shock --asset-group "ignore+stables" --oracle-distortion 0.05 --n-scenarios 5 --pool-id 0

backend/utils/price_shock.py

@@ -21,13 +21,19 @@ def create_dataframes(leverages):
 
 
 def calculate_spot_bankruptcies(df):
+    if df.empty or 'spot_asset' not in df.columns or 'spot_liability' not in df.columns or 'net_usd_value' not in df.columns:
+        return 0.0
+    
     spot_bankrupt = df[
         (df["spot_asset"] < df["spot_liability"]) & (df["net_usd_value"] < 0)
     ]
     return (spot_bankrupt["spot_liability"] - spot_bankrupt["spot_asset"]).sum()
 
 
 def calculate_total_bankruptcies(df):
+    if df.empty or 'net_usd_value' not in df.columns:
+        return 0.0
+    
     return -df[df["net_usd_value"] < 0]["net_usd_value"].sum()
 
 
@@ -46,6 +52,7 @@ def get_price_shock_df(
     oracle_distortion: float,
     asset_group: PriceShockAssetGroup,
     n_scenarios: int,
+    pool_id: int = None,
 ):
     user_leverages = get_user_leverages_for_price_shock(
         slot,
@@ -54,8 +61,35 @@ def get_price_shock_df(
         oracle_distortion,
         asset_group,
         n_scenarios,
+        pool_id,
     )
     levs = user_leverages
+    
+    # Handle case where no users are found after filtering
+    if not levs["user_keys"] or len(levs["leverages_none"]) == 0:
+        print(f"No users found for the specified filters. Returning empty result.")
+        oracle_moves = generate_oracle_moves(n_scenarios, oracle_distortion)
+        
+        # Create empty DataFrame with all zeros
+        df_plot = pd.DataFrame(
+            {
+                "Oracle Move (%)": oracle_moves,
+                "Total Bankruptcy ($)": [0.0] * len(oracle_moves),
+                "Spot Bankruptcy ($)": [0.0] * len(oracle_moves),
+                "Perpetual Bankruptcy ($)": [0.0] * len(oracle_moves),
+            }
+        )
+        
+        df_plot = df_plot.sort_values("Oracle Move (%)")
+        
+        return {
+            "slot": slot,
+            "result": df_plot.to_json(),
+            "distorted_oracles": levs["distorted_oracles"],
+            "oracle_down_max": pd.DataFrame().to_json(),
+            "oracle_up_max": pd.DataFrame().to_json(),
+        }
+    
     dfs = (
         create_dataframes(levs["leverages_down"])
         + [pd.DataFrame(levs["leverages_none"])]

backend/utils/user_metrics.py

@@ -260,9 +260,39 @@ def get_user_leverages_for_price_shock(
     oracle_distortion: float = 0.1,
     asset_group: PriceShockAssetGroup = PriceShockAssetGroup.IGNORE_STABLES,
     scenarios: int = 5,
+    pool_id: int = None,
 ):
     user_keys = list(user_map.user_map.keys())
     user_vals = list(user_map.values())
+    
+    # Filter users by pool_id if specified
+    if pool_id is not None:
+        filtered_user_vals = []
+        filtered_user_keys = []
+        pool_id_counts = {}  # Track distribution of pool_ids
+        
+        for user_key, user in zip(user_keys, user_vals):
+            try:
+                user_account = user.get_user_account()
+                actual_pool_id = user_account.pool_id
+                
+                # Track pool_id distribution for debugging
+                pool_id_counts[actual_pool_id] = pool_id_counts.get(actual_pool_id, 0) + 1
+                
+                if actual_pool_id == pool_id:
+                    filtered_user_vals.append(user)
+                    filtered_user_keys.append(user_key)
+            except Exception as e:
+                print(f"Error checking pool_id for user {user_key}: {e}")
+                continue
+        
+        print(f"Pool ID distribution: {dict(sorted(pool_id_counts.items()))}")
+        print(f"Requested pool_id: {pool_id}")
+        print(f"Filtered to {len(filtered_user_vals)} users with pool_id {pool_id} (out of {len(user_vals)} total users)")
+        
+        user_vals = filtered_user_vals
+        user_keys = filtered_user_keys
+    
     all_configs = mainnet_spot_market_configs + mainnet_perp_market_configs
 
     print(f"User keys : {len(user_keys)}")

gen.sh

@@ -1,47 +1,184 @@
 #!/bin/bash
 set -e
 
-
 # Run the first one sync, this will generate a fresh pickle
 python -m backend.scripts.generate_ucache \
     asset-liability \
     --mode 0 \
     --perp-market-index 0
 
-
 # The next ones will use the --use-snapshot flag, so they will reuse the pickle
 # We can run all commands in parallel by adding & at the end
+
+echo "Generating comprehensive price shock cache matrix..."
+echo "Asset groups: ignore+stables, jlp+only"
+echo "Scenarios: 5 (0.05 distortion), 10 (0.1 distortion)" 
+echo "Pool IDs: all, 0, 1, 3"
+echo "Total combinations: 16"
+
+# =============================================================================
+# ALL POOLS (no pool filter) - 4 combinations
+# =============================================================================
+
+# ignore+stables, 5 scenarios, 0.05 distortion, all pools
 python -m backend.scripts.generate_ucache \
     --use-snapshot \
     price-shock \
     --asset-group "ignore+stables" \
     --oracle-distortion 0.05 \
     --n-scenarios 5 &
 
+# ignore+stables, 10 scenarios, 0.1 distortion, all pools  
+python -m backend.scripts.generate_ucache \
+    --use-snapshot \
+    price-shock \
+    --asset-group "ignore+stables" \
+    --oracle-distortion 0.1 \
+    --n-scenarios 10 &
+
+# jlp+only, 5 scenarios, 0.05 distortion, all pools
 python -m backend.scripts.generate_ucache \
     --use-snapshot \
     price-shock \
     --asset-group "jlp+only" \
     --oracle-distortion 0.05 \
     --n-scenarios 5 &
 
+# jlp+only, 10 scenarios, 0.1 distortion, all pools
 python -m backend.scripts.generate_ucache \
     --use-snapshot \
     price-shock \
     --asset-group "jlp+only" \
     --oracle-distortion 0.1 \
     --n-scenarios 10 &
 
+# =============================================================================
+# MAIN POOL (pool_id = 0) - 4 combinations  
+# =============================================================================
+
+# ignore+stables, 5 scenarios, 0.05 distortion, pool 0
+python -m backend.scripts.generate_ucache \
+    --use-snapshot \
+    price-shock \
+    --asset-group "ignore+stables" \
+    --oracle-distortion 0.05 \
+    --n-scenarios 5 \
+    --pool-id 0 &
 
+# ignore+stables, 10 scenarios, 0.1 distortion, pool 0
 python -m backend.scripts.generate_ucache \
     --use-snapshot \
     price-shock \
     --asset-group "ignore+stables" \
     --oracle-distortion 0.1 \
-    --n-scenarios 10 &
+    --n-scenarios 10 \
+    --pool-id 0 &
+
+# jlp+only, 5 scenarios, 0.05 distortion, pool 0  
+python -m backend.scripts.generate_ucache \
+    --use-snapshot \
+    price-shock \
+    --asset-group "jlp+only" \
+    --oracle-distortion 0.05 \
+    --n-scenarios 5 \
+    --pool-id 0 &
+
+# jlp+only, 10 scenarios, 0.1 distortion, pool 0
+python -m backend.scripts.generate_ucache \
+    --use-snapshot \
+    price-shock \
+    --asset-group "jlp+only" \
+    --oracle-distortion 0.1 \
+    --n-scenarios 10 \
+    --pool-id 0 &
+
+# =============================================================================
+# ISOLATED POOL 1 (pool_id = 1) - 4 combinations
+# =============================================================================
+
+# ignore+stables, 5 scenarios, 0.05 distortion, pool 1
+python -m backend.scripts.generate_ucache \
+    --use-snapshot \
+    price-shock \
+    --asset-group "ignore+stables" \
+    --oracle-distortion 0.05 \
+    --n-scenarios 5 \
+    --pool-id 1 &
+
+# ignore+stables, 10 scenarios, 0.1 distortion, pool 1
+python -m backend.scripts.generate_ucache \
+    --use-snapshot \
+    price-shock \
+    --asset-group "ignore+stables" \
+    --oracle-distortion 0.1 \
+    --n-scenarios 10 \
+    --pool-id 1 &
 
+# jlp+only, 5 scenarios, 0.05 distortion, pool 1
+python -m backend.scripts.generate_ucache \
+    --use-snapshot \
+    price-shock \
+    --asset-group "jlp+only" \
+    --oracle-distortion 0.05 \
+    --n-scenarios 5 \
+    --pool-id 1 &
+
+# jlp+only, 10 scenarios, 0.1 distortion, pool 1  
+python -m backend.scripts.generate_ucache \
+    --use-snapshot \
+    price-shock \
+    --asset-group "jlp+only" \
+    --oracle-distortion 0.1 \
+    --n-scenarios 10 \
+    --pool-id 1 &
+
+# =============================================================================
+# ISOLATED POOL 3 (pool_id = 3) - 4 combinations
+# =============================================================================
+
+# ignore+stables, 5 scenarios, 0.05 distortion, pool 3
+python -m backend.scripts.generate_ucache \
+    --use-snapshot \
+    price-shock \
+    --asset-group "ignore+stables" \
+    --oracle-distortion 0.05 \
+    --n-scenarios 5 \
+    --pool-id 3 &
+
+# ignore+stables, 10 scenarios, 0.1 distortion, pool 3
+python -m backend.scripts.generate_ucache \
+    --use-snapshot \
+    price-shock \
+    --asset-group "ignore+stables" \
+    --oracle-distortion 0.1 \
+    --n-scenarios 10 \
+    --pool-id 3 &
+
+# jlp+only, 5 scenarios, 0.05 distortion, pool 3
+python -m backend.scripts.generate_ucache \
+    --use-snapshot \
+    price-shock \
+    --asset-group "jlp+only" \
+    --oracle-distortion 0.05 \
+    --n-scenarios 5 \
+    --pool-id 3 &
+
+# jlp+only, 10 scenarios, 0.1 distortion, pool 3
+python -m backend.scripts.generate_ucache \
+    --use-snapshot \
+    price-shock \
+    --asset-group "jlp+only" \
+    --oracle-distortion 0.1 \
+    --n-scenarios 10 \
+    --pool-id 3 &
+
+# =============================================================================
 # Wait for all background processes to complete
+# =============================================================================
+echo "Waiting for all 16 price shock cache generation processes to complete..."
 wait
 
+echo "✅ All cache generation completed successfully!"
+
 # Delete old pickles
 cd pickles && ls -t | tail -n +4 | xargs rm -rf

requirements.txt

@@ -12,7 +12,7 @@ uvicorn==0.34.0
 requests==2.32.3
 plotly==6.0.0
 anchorpy==0.21.0
-driftpy>=0.8.40
+driftpy>=0.8.56
 ccxt==4.2.17
 rich>=10.14.0
 aiofiles==24.1.0