bigz/margin-model

.gitignore

@@ -2,4 +2,5 @@ venv
 .mypy_cache
 /__pycache__/
 /src/__pycache__/
-/src/sections/__pycache__
+/src/sections/__pycache__
+/src/.DS_Store

.vscode/settings.json

@@ -0,0 +1,10 @@
+{
+    "python.linting.enabled": true,
+    "python.linting.mypyEnabled": true,
+    "python.linting.mypyArgs": [
+        "--ignore-missing-imports",
+        "--strict"
+    ],
+    "python.linting.pylintEnabled": true,  // Optional: Enables pylint
+    "python.linting.flake8Enabled": true   // Optional: Enables flake8
+}

mypy.ini

@@ -0,0 +1,2 @@
+[mypy]
+check_untyped_defs = True

pickles/userstats_0.pkl

@@ -0,0 +1 @@
+�]�.

requirements.txt

@@ -1,7 +1,6 @@
 aiodns==3.0.0
 aiohttp==3.8.3
 aiosignal==1.3.1
-altair==5.3.0
 anchorpy==0.20.1
 anchorpy-core==0.2.0
 anyio==3.6.2
@@ -11,7 +10,7 @@ attrs==22.1.0
 backoff==2.2.1
 base58==2.1.1
 based58==0.1.1
-blinker==1.8.2
+black==24.4.2
 borsh-construct==0.1.0
 cachetools==4.2.4
 certifi==2022.12.7
@@ -22,86 +21,68 @@ construct==2.10.68
 construct-typing==0.5.3
 Deprecated==1.2.14
 dnspython==2.2.1
-driftpy==0.7.59
+driftpy==0.7.73
 Events==0.5
 exceptiongroup==1.0.4
 flake8==6.0.0
 frozenlist==1.3.3
 ghp-import==2.1.0
-gitdb==4.0.11
-GitPython==3.1.43
-grpcio==1.64.1
+grpcio==1.65.1
 h11==0.14.0
 httpcore==0.16.3
 httpx==0.23.1
 idna==3.4
 iniconfig==1.1.1
-Jinja2==3.0.3
+Jinja2==3.1.4
 jito_searcher_client==0.1.4
 jsonalias==0.1.1
-jsonrpcclient==4.0.2
+jsonrpcclient==4.0.3
 jsonrpcserver==5.0.9
 jsonschema==4.17.3
 loguru==0.6.0
 Markdown==3.6
-markdown-it-py==3.0.0
 MarkupSafe==2.1.5
 mccabe==0.7.0
-mdurl==0.1.2
 mergedeep==1.3.4
 mkdocs==1.6.0
 mkdocs-get-deps==0.2.0
 more-itertools==8.14.0
 multidict==6.0.3
-mypy==1.10.0
+mypy==1.11.0
 mypy-extensions==1.0.0
 numpy==1.26.4
 OSlash==0.6.3
-packaging==22.0
-pandas==2.2.2
+packaging==23.1
 pathspec==0.12.1
-pillow==10.3.0
 platformdirs==4.2.2
-plotly==5.22.0
 pluggy==1.0.0
-protobuf==4.25.3
+protobuf==4.25.4
 psutil==5.9.4
 py==1.11.0
-pyarrow==16.1.0
 pycares==4.3.0
 pycodestyle==2.10.0
 pycparser==2.21
-pydeck==0.9.1
 pyflakes==3.0.1
-Pygments==2.18.0
 pyheck==0.1.5
 pyrsistent==0.19.2
 pythclient==0.1.4
 python-dateutil==2.9.0.post0
-pytz==2024.1
 PyYAML==6.0.1
 pyyaml_env_tag==0.1
 requests==2.32.3
 rfc3986==1.5.0
-rich==13.7.1
-rpds-py==0.18.1
 six==1.16.0
-smmap==5.0.1
 sniffio==1.3.0
 solana==0.34.0
 solders==0.21.0
-streamlit==1.35.0
 sumtypes==0.1a6
-tenacity==8.3.0
 toml==0.10.2
 tomli==2.0.1
 toolz==0.11.2
-tornado==6.4.1
 types-cachetools==4.2.10
 types-requests==2.31.0.6
 types-urllib3==1.26.25.14
 typing_extensions==4.12.2
-tzdata==2024.1
 urllib3==1.26.13
 watchdog==4.0.1
 websockets==10.4

script.py

@@ -0,0 +1,59 @@
+import asyncio
+import os 
+import shutil
+import glob
+import time
+import traceback
+
+from anchorpy import Wallet
+
+from solana.rpc.async_api import AsyncClient
+
+from driftpy.drift_client import DriftClient
+from driftpy.account_subscription_config import AccountSubscriptionConfig
+from driftpy.user_map.user_map import UserMap
+from driftpy.user_map.user_map_config import UserMapConfig, WebsocketConfig as Cfg1
+from driftpy.user_map.userstats_map import UserStatsMap, UserStatsMapConfig
+from driftpy.market_map.market_map_config import MarketMapConfig, WebsocketConfig as Cfg2
+from driftpy.market_map.market_map import MarketMap
+from driftpy.pickle.vat import Vat
+from driftpy.types import MarketType
+
+
+async def main():
+    rpc = str(os.getenv("MAINNET_RPC_ENDPOINT"))
+    wallet = Wallet.dummy()
+
+    dc = DriftClient(
+        AsyncClient(rpc),
+        wallet,
+        "mainnet",
+        account_subscription=AccountSubscriptionConfig("cached")
+    )
+
+    await dc.subscribe()
+
+    users = UserMap(UserMapConfig(dc, Cfg1))
+    stats = UserStatsMap(UserStatsMapConfig(dc))
+    spot = MarketMap(MarketMapConfig(dc.program, MarketType.Spot(), Cfg2, dc.connection))
+    perp = MarketMap(MarketMapConfig(dc.program, MarketType.Perp(), Cfg2, dc.connection))
+
+    vat = Vat(dc, users, stats, spot, perp)
+
+    start = time.time()
+    try:
+        await vat.pickle()
+    except:
+        print("an error occurred during pickling:")
+        traceback.print_exc()
+        os._exit(0)
+    print("pickled in:", time.time() - start)
+
+    pkl_files = glob.glob("*.pkl")
+
+    for file in pkl_files:
+        shutil.move(file, "./pickles")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

src/cache.py

@@ -0,0 +1,41 @@
+from asyncio import AbstractEventLoop
+from scenario import get_usermap_df
+import time
+import streamlit as st
+from driftpy.drift_client import DriftClient
+from driftpy.pickle.vat import Vat
+import pandas as pd
+
+
+def _load_asset_liab_dfs(
+    dc: DriftClient, vat: Vat, loop: AbstractEventLoop
+) -> tuple[tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame], list[str]]:
+    start = time.time()
+    oracle_distort = 0
+    (levs_none, levs_init, levs_maint), user_keys = loop.run_until_complete(
+        get_usermap_df(
+            dc,
+            vat.users,
+            "margins",
+            oracle_distort,
+            None,
+            "ignore stables",
+            n_scenarios=0,
+            all_fields=True,
+        )
+    )
+    print(f"Loaded asset/liability data in {time.time() - start:.2f} seconds")
+    return (levs_none, levs_init, levs_maint), user_keys
+
+
+@st.cache_data
+def _cache_dataframes(
+    dfs: tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame], user_keys: list[str]
+) -> tuple[tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame], list[str]]:
+    return dfs, user_keys
+
+
+def get_cached_asset_liab_dfs(
+    dc: DriftClient, vat: Vat, loop: AbstractEventLoop
+) -> tuple[tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame], list[str]]:
+    return _cache_dataframes(*_load_asset_liab_dfs(dc, vat, loop))

src/health_utils.py

@@ -24,11 +24,18 @@
 )
 from driftpy.types import is_variant
 from driftpy.pickle.vat import Vat
-from driftpy.constants.spot_markets import mainnet_spot_market_configs, devnet_spot_market_configs
-from driftpy.constants.perp_markets import mainnet_perp_market_configs, devnet_perp_market_configs
+from driftpy.constants.spot_markets import (
+    mainnet_spot_market_configs,
+    devnet_spot_market_configs,
+)
+from driftpy.constants.perp_markets import (
+    mainnet_perp_market_configs,
+    devnet_perp_market_configs,
+)
 
 from utils import load_newest_files, load_vat, to_financial
 
+
 def get_largest_perp_positions(vat: Vat):
     top_positions: list[Any] = []
 
@@ -293,4 +300,4 @@ def get_most_levered_spot_borrows_above_1m(vat: Vat):
         "Public Key": [pos[1] for pos in borrows],
     }
 
-    return data
+    return data