Update script, format with ruff

Author: SinaKhalili

.gitignore

@@ -6,3 +6,4 @@ venv
 pickles/*
 cache
 ignore
+ucache

backend/scripts/generate_ucache.py

@@ -1,3 +1,4 @@
+import argparse
 import asyncio
 import glob
 import json
@@ -29,6 +30,35 @@ class Endpoint:
     params: dict
 
 
+@dataclass
+class AssetLiabilityEndpoint(Endpoint):
+    mode: int
+    perp_market_index: int
+
+    def __init__(self, mode: int, perp_market_index: int):
+        super().__init__(
+            "asset-liability/matrix",
+            {"mode": mode, "perp_market_index": perp_market_index},
+        )
+
+
+@dataclass
+class PriceShockEndpoint(Endpoint):
+    asset_group: str
+    oracle_distortion: float
+    n_scenarios: int
+
+    def __init__(self, asset_group: str, oracle_distortion: float, n_scenarios: int):
+        super().__init__(
+            "price-shock/usermap",
+            {
+                "asset_group": asset_group,
+                "oracle_distortion": oracle_distortion,
+                "n_scenarios": n_scenarios,
+            },
+        )
+
+
 async def process_multiple_endpoints(state_pickle_path: str, endpoints: list[Endpoint]):
     """Process a single endpoint in its own process"""
     state = BackendState()
@@ -102,7 +132,6 @@ async def run_request():
 async def generate_ucache(endpoints: list[Endpoint]):
     """Generate ucache files by splitting endpoints across processes"""
     ucache_dir = "ucache"
-    use_snapshot = os.getenv("USE_SNAPSHOT", "false").lower() == "true"
 
     print("Generating ucache")
     if not os.path.exists(ucache_dir):
@@ -112,33 +141,49 @@ async def generate_ucache(endpoints: list[Endpoint]):
 
 
 async def main():
-    load_dotenv()
+    parser = argparse.ArgumentParser(description="Generate ucache files")
+    parser.add_argument(
+        "--use-snapshot", action="store_true", help="Use existing snapshot"
+    )
+
+    subparsers = parser.add_subparsers(dest="command", required=True)
 
-    use_snapshot = os.getenv("USE_SNAPSHOT", "false").lower() == "true"
-    print(f"use_snapshot: {use_snapshot}")
+    # Asset Liability parser
+    al_parser = subparsers.add_parser("asset-liability")
+    al_parser.add_argument("--mode", type=int, required=True)
+    al_parser.add_argument("--perp-market-index", type=int, required=True)
 
-    if not use_snapshot:
+    # Price Shock parser
+    ps_parser = subparsers.add_parser("price-shock")
+    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)
+
+    args = parser.parse_args()
+
+    if not args.use_snapshot:
         state = BackendState()
         state.initialize(os.getenv("RPC_URL") or "")
         print("Taking snapshot")
         await state.bootstrap()
         await state.take_pickle_snapshot()
         await state.close()
 
-    endpoints = [
-        Endpoint(
-            endpoint="asset-liability/matrix",
-            params={"mode": 0, "perp_market_index": 0},
-        ),
-        Endpoint(
-            endpoint="price-shock/usermap",
-            params={
-                "asset_group": "ignore+stables",
-                "oracle_distortion": 0.05,
-                "n_scenarios": 5,
-            },
-        ),
-    ]
+    endpoints = []
+    if args.command == "asset-liability":
+        endpoints.append(
+            AssetLiabilityEndpoint(
+                mode=args.mode, perp_market_index=args.perp_market_index
+            )
+        )
+    elif args.command == "price-shock":
+        endpoints.append(
+            PriceShockEndpoint(
+                asset_group=args.asset_group,
+                oracle_distortion=args.oracle_distortion,
+                n_scenarios=args.n_scenarios,
+            )
+        )
 
     await generate_ucache(endpoints)
 

src/lib/user_metrics.py

@@ -1,8 +1,7 @@
 import copy
 from typing import List, Optional
 
-from driftpy.constants.numeric_constants import MARGIN_PRECISION
-from driftpy.constants.numeric_constants import QUOTE_PRECISION
+from driftpy.constants.numeric_constants import MARGIN_PRECISION, QUOTE_PRECISION
 from driftpy.constants.perp_markets import mainnet_perp_market_configs
 from driftpy.constants.spot_markets import mainnet_spot_market_configs
 from driftpy.drift_client import DriftClient

src/main.py

@@ -4,14 +4,13 @@
 from dotenv import load_dotenv
 
 from lib.page import header, needs_backend, sidebar
-from page.asset_liability_cached import asset_liab_matrix_cached_page
+from page.asset_liability import asset_liab_matrix_cached_page
 from page.backend import backend_page
 from page.health import health_page
-from page.health_cached import health_cached_page
 from page.liquidation_curves import liquidation_curves_page
 from page.orderbook import orderbook_page
-from page.price_shock_cached import price_shock_cached_page
-from sections.welcome import welcome_page
+from page.price_shock import price_shock_cached_page
+from page.welcome import welcome_page
 
 load_dotenv()
 
@@ -51,27 +50,15 @@ def apply_custom_css(css):
             title="Health",
             icon="🏥",
         ),
-        # st.Page(
-        #     needs_backend(price_shock_page),
-        #     url_path="price-shock",
-        #     title="Price Shock",
-        #     icon="💸",
-        # ),
-        # st.Page(
-        #     needs_backend(asset_liab_matrix_page),
-        #     url_path="asset-liability-matrix",
-        #     title="Asset-Liability Matrix",
-        #     icon="📊",
-        # ),
         st.Page(
             price_shock_cached_page,
-            url_path="price-shock-cached",
+            url_path="price-shock",
             title="Price Shock",
             icon="💸",
         ),
         st.Page(
             asset_liab_matrix_cached_page,
-            url_path="asset-liability-matrix-cached",
+            url_path="asset-liability-matrix",
             title="Asset-Liability Matrix",
             icon="📊",
         ),
@@ -82,14 +69,7 @@ def apply_custom_css(css):
             icon="🌊",
         ),
     ]
-    cached_pages = [
-        st.Page(
-            health_cached_page,
-            url_path="health-cached",
-            title="Health (Cached)",
-            icon="🏥",
-        ),
-    ]
+
     if os.getenv("DEV"):
         main_pages.append(
             st.Page(
@@ -103,7 +83,6 @@ def apply_custom_css(css):
     pg = st.navigation(
         {
             "Main": main_pages,
-            # "Cached": cached_pages,
         }
     )
     pg.run()

src/page/asset_liability.py

@@ -1,12 +1,10 @@
-import json
-
-from driftpy.constants.perp_markets import mainnet_perp_market_configs
-from driftpy.constants.spot_markets import mainnet_spot_market_configs
-from lib.api import api
 import pandas as pd
-from requests.exceptions import JSONDecodeError
 import streamlit as st
+from driftpy.constants.perp_markets import mainnet_perp_market_configs
+from driftpy.constants.spot_markets import mainnet_spot_market_configs
 
+from lib.api import api2
+from utils import get_current_slot
 
 options = [0, 1, 2, 3]
 labels = [
@@ -17,77 +15,133 @@
 ]
 
 
-def asset_liab_matrix_page():
+def calculate_effective_leverage(assets: float, liabilities: float) -> float:
+    return liabilities / assets if assets != 0 else 0
+
+
+def format_metric(
+    value: float, should_highlight: bool, mode: int, financial: bool = False
+) -> str:
+    formatted = f"{value:,.2f}" if financial else f"{value:.2f}"
+    return f"{formatted} ✅" if should_highlight and mode > 0 else formatted
+
+
+def generate_summary_data(
+    df: pd.DataFrame, mode: int, perp_market_index: int
+) -> pd.DataFrame:
+    summary_data = {}
+    for i in range(len(mainnet_spot_market_configs)):
+        prefix = f"spot_{i}"
+        assets = df[f"{prefix}_all_assets"].sum()
+        liabilities = df[f"{prefix}_all"].sum()
+
+        summary_data[f"spot{i}"] = {
+            "all_assets": assets,
+            "all_liabilities": format_metric(
+                liabilities, 0 < liabilities < 1_000_000, mode, financial=True
+            ),
+            "effective_leverage": format_metric(
+                calculate_effective_leverage(assets, liabilities),
+                0 < calculate_effective_leverage(assets, liabilities) < 2,
+                mode,
+            ),
+            "all_spot": df[f"{prefix}_all_spot"].sum(),
+            "all_perp": df[f"{prefix}_all_perp"].sum(),
+            f"perp_{perp_market_index}_long": df[
+                f"{prefix}_perp_{perp_market_index}_long"
+            ].sum(),
+            f"perp_{perp_market_index}_short": df[
+                f"{prefix}_perp_{perp_market_index}_short"
+            ].sum(),
+        }
+    return pd.DataFrame(summary_data).T
+
+
+def asset_liab_matrix_cached_page():
+    if "min_leverage" not in st.session_state:
+        st.session_state.min_leverage = 0.0
+    if "only_high_leverage_mode_users" not in st.session_state:
+        st.session_state.only_high_leverage_mode_users = False
+
     params = st.query_params
     mode = int(params.get("mode", 0))
     perp_market_index = int(params.get("perp_market_index", 0))
 
     mode = st.selectbox(
         "Options", options, format_func=lambda x: labels[x], index=options.index(mode)
     )
-    st.query_params.update({"mode": mode})
+    st.query_params.update({"mode": str(mode)})
 
     perp_market_index = st.selectbox(
         "Market index",
         [x.market_index for x in mainnet_perp_market_configs],
         index=[x.market_index for x in mainnet_perp_market_configs].index(
             perp_market_index
         ),
+        format_func=lambda x: f"{x} ({mainnet_perp_market_configs[int(x)].symbol})",
     )
-    st.query_params.update({"perp_market_index": perp_market_index})
+    st.query_params.update({"perp_market_index": str(perp_market_index)})
 
-    min_leverage = st.slider(
-        "Filter by minimum leverage", 0.0, 50.0, 0.0, 0.5, key="min_leverage"
+    result = api2(
+        "asset-liability/matrix",
+        _params={"mode": mode, "perp_market_index": perp_market_index},
+        key=f"asset-liability/matrix_{mode}_{perp_market_index}",
     )
+    df = pd.DataFrame(result["df"])
 
-    try:
-        result = api(
-            "asset-liability",
-            "matrix",
-            params=params,
-            as_json=True,
-        )
-        if "result" in result and result["result"] == "miss":
-            st.write("Fetching data for the first time...")
-            st.image(
-                "https://i.gifer.com/origin/8a/8a47f769c400b0b7d81a8f6f8e09a44a_w200.gif"
-            )
-            st.write("Check again in one minute!")
-            st.stop()
-
-    except Exception as e:
-        if type(e) == JSONDecodeError:
-            print("HIT A JSONDecodeError...", e)
-            st.write("Fetching data for the first time...")
-            st.image(
-                "https://i.gifer.com/origin/8a/8a47f769c400b0b7d81a8f6f8e09a44a_w200.gif"
-            )
-            st.write("Check again in one minute!")
-            st.stop()
-        else:
-            st.write(e)
-            st.stop()
+    if st.session_state.only_high_leverage_mode_users:
+        df = df[df["is_high_leverage"]]
 
-    res = pd.DataFrame(result["res"])
-    df = pd.DataFrame(result["df"])
+    filtered_df = df[df["leverage"] >= st.session_state.min_leverage].sort_values(
+        "leverage", ascending=False
+    )
 
-    st.write(f"{df.shape[0]} users for scenario")
-    st.write(res)
+    summary_df = generate_summary_data(filtered_df, mode, perp_market_index)
+    slot = result["slot"]
+    current_slot = get_current_slot()
+
+    st.info(
+        f"This data is for slot {slot}, which is now {int(current_slot) - int(slot)} slots old"
+    )
+    st.write(f"{df.shape[0]} users")
+    st.checkbox(
+        "Only show high leverage mode users", key="only_high_leverage_mode_users"
+    )
+    st.slider(
+        "Filter by minimum leverage",
+        0.0,
+        110.0,
+        0.0,
+        key="min_leverage",
+    )
+    st.write(summary_df)
 
     tabs = st.tabs(["FULL"] + [x.symbol for x in mainnet_spot_market_configs])
 
-    # Add leverage filter to FULL tab
     with tabs[0]:
-        filtered_df = df[df["leverage"] >= min_leverage].sort_values(
-            "leverage", ascending=False
-        )
+        if st.session_state.only_high_leverage_mode_users:
+            st.write(
+                f"There are **{len(filtered_df)}** users with high leverage mode and {st.session_state.min_leverage}x leverage or more"
+            )
+        else:
+            st.write(
+                f"There are **{len(filtered_df)}** users with this **{st.session_state.min_leverage}x** leverage or more"
+            )
+        st.write(f"Total USD value: **{filtered_df['net_usd_value'].sum():,.2f}**")
+        st.write(f"Total collateral: **{filtered_df['spot_asset'].sum():,.2f}**")
+        st.write(f"Total liabilities: **{filtered_df['spot_liability'].sum():,.2f}**")
         st.dataframe(filtered_df, hide_index=True)
 
     for idx, tab in enumerate(tabs[1:]):
-        important_cols = [x for x in df.columns if "spot_" + str(idx) in x]
-        toshow = df[["spot_asset", "net_usd_value"] + important_cols]
+        important_cols = [x for x in filtered_df.columns if "spot_" + str(idx) in x]
+
+        toshow = filtered_df[
+            ["user_key", "spot_asset", "net_usd_value"] + important_cols
+        ]
         toshow = toshow[toshow[important_cols].abs().sum(axis=1) != 0].sort_values(
             by="spot_" + str(idx) + "_all", ascending=False
         )
-        tab.write(f"{ len(toshow)} users with this asset to cover liabilities")
+        tab.write(
+            f"{len(toshow)} users with this asset to cover liabilities (with {st.session_state.min_leverage}x leverage or more)"
+        )
         tab.dataframe(toshow, hide_index=True)