83 lines
2.6 KiB
Python
83 lines
2.6 KiB
Python
|
"""Process indicators for many symbols in parallel.
|
||
|
|
|
||
|
|
Python's GIL makes pure-Python parallelism a poor fit, but Wickra's heavy
|
||
|
|
lifting happens inside the Rust extension — which releases the GIL during
|
||
|
|
batch computation. That means a `ThreadPoolExecutor` actually delivers
|
||
|
|
multi-core speedup here.
|
||
|
|
|
||
|
|
Run with::
|
||
|
|
|
||
|
|
python -m examples.python.parallel_assets --assets 1000 --bars 5000
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
import concurrent.futures
|
||
|
|
import time
|
||
|
|
from typing import Tuple
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
|
||
|
|
import wickra as ta
|
||
|
|
|
||
|
|
|
||
|
|
def synthesize_panel(n_assets: int, n_bars: int, seed: int = 0xBADC0FFEE0DDF00D) -> np.ndarray:
|
||
|
|
"""Build an `(n_assets, n_bars)` matrix of synthetic prices."""
|
||
|
|
rng = np.random.default_rng(seed & 0xFFFF_FFFF)
|
||
|
|
drift = rng.standard_normal((n_assets, n_bars)) * 0.4
|
||
|
|
return 100.0 + np.cumsum(drift, axis=1)
|
||
|
|
|
||
|
|
|
||
|
|
def compute_one(prices: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
|
||
|
|
"""Return (RSI, EMA, MACD-line) for one asset."""
|
||
|
|
rsi = ta.RSI(14).batch(prices)
|
||
|
|
ema = ta.EMA(20).batch(prices)
|
||
|
|
macd = ta.MACD().batch(prices)
|
||
|
|
return rsi, ema, macd[:, 0]
|
||
|
|
|
||
|
|
|
||
|
|
def main() -> int:
|
||
|
|
p = argparse.ArgumentParser(description=__doc__.splitlines()[0] if __doc__ else None)
|
||
|
|
p.add_argument("--assets", type=int, default=200)
|
||
|
|
p.add_argument("--bars", type=int, default=5000)
|
||
|
|
p.add_argument(
|
||
|
|
"--workers",
|
||
|
|
type=int,
|
||
|
|
default=0,
|
||
|
|
help="thread pool size; 0 means use os.cpu_count()",
|
||
|
|
)
|
||
|
|
args = p.parse_args()
|
||
|
|
|
||
|
|
print(f"Generating {args.assets}x{args.bars} synthetic panel…")
|
||
|
|
panel = synthesize_panel(args.assets, args.bars)
|
||
|
|
|
||
|
|
# Serial baseline.
|
||
|
|
t0 = time.perf_counter()
|
||
|
|
serial_results = [compute_one(panel[i]) for i in range(args.assets)]
|
||
|
|
t_serial = time.perf_counter() - t0
|
||
|
|
print(f"Serial: {t_serial:.3f} s ({args.assets} assets)")
|
||
|
|
|
||
|
|
# Parallel.
|
||
|
|
workers = args.workers or None
|
||
|
|
t0 = time.perf_counter()
|
||
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as pool:
|
||
|
|
parallel_results = list(pool.map(compute_one, panel))
|
||
|
|
t_parallel = time.perf_counter() - t0
|
||
|
|
used = workers or 0
|
||
|
|
print(
|
||
|
|
f"Parallel: {t_parallel:.3f} s (workers={used or 'os.cpu_count()'}, "
|
||
|
|
f"speedup ~{t_serial / max(t_parallel, 1e-9):.2f}x)"
|
||
|
|
)
|
||
|
|
|
||
|
|
# Sanity-check that the parallel results match the serial baseline.
|
||
|
|
for i in range(args.assets):
|
||
|
|
for s, p_ in zip(serial_results[i], parallel_results[i]):
|
||
|
|
np.testing.assert_array_equal(s, p_)
|
||
|
|
print("Parallel results match serial results — OK.")
|
||
|
|
return 0
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
raise SystemExit(main())
|