1import numpy as np
2import time
3from sklearn.linear_model import LinearRegression
5########################################################################
6# 1) Functions to build the design matrix in float32
7########################################################################
9def build_design_matrix_naive(x: np.ndarray, M: int):
10 """
11 Naive Python approach to build design matrix A for Markov model.
12 x is int16, but we'll cast to float32 in A.
13 A has shape (N - M, M + 1).
14 Row j = [1, x_{j+M-1}, x_{j+M-2}, ..., x_{j+M-M}],
15 for j in 0..(N-M-1).
16 """
17 N = len(x)
18 A = np.ones((N - M, M + 1), dtype=np.float32) # store as float32
19 for row_idx in range(N - M):
20 for k in range(1, M + 1):
21 # cast x (int16) to float32 on assignment
22 A[row_idx, k] = x[M + row_idx - k]
23 return A
25def build_design_matrix_vectorized(x: np.ndarray, M: int):
26 """
27 Vectorized approach (using slicing) to build the same design matrix in float32.
28 x is int16; cast slices to float32.
29 """
30 N = len(x)
31 rows = N - M
32 # First column of 1's in float32
33 A_cols = [np.ones(rows, dtype=np.float32)]
34 # Next columns are shifted versions of x, cast to float32
35 for shift in range(1, M + 1):
36 A_cols.append(x[M - shift : N - shift].astype(np.float32))
37 A = np.column_stack(A_cols)
38 return A
40########################################################################
41# 2) Functions to solve for coefficients (stored as float32)
42########################################################################
44def solve_lstsq_numpy(A: np.ndarray, y: np.ndarray):
45 """
46 Solve for coefficients using NumPy's lstsq solver in float32.
47 Returns a 1D array of shape (M+1,) in float32.
48 """
49 # lstsq may do computations in float64 internally, but we can cast back
50 c, _, _, _ = np.linalg.lstsq(A, y, rcond=None)
51 return c.astype(np.float32)
53def solve_lstsq_sklearn(A: np.ndarray, y: np.ndarray):
54 """
55 Solve for coefficients using scikit-learn's LinearRegression.
56 (fit_intercept=False because we already have the 1-column in A).
57 We store final coefficients as float32.
58 """
59 reg = LinearRegression(fit_intercept=False)
60 reg.fit(A, y)
61 # sklearn might do double precision internally; cast to float32
62 return reg.coef_.astype(np.float32)
64########################################################################
65# 3) Functions to predict a new sequence
66########################################################################
68def predict_naive(coeffs: np.ndarray, seed: np.ndarray, N: int):
69 """
70 Naive Python prediction loop with float32 arithmetic.
71 - coeffs: float32 array of shape (M+1,) => [c0, c1, ..., cM].
72 - seed: last M known values (int16), but we'll cast to float32 internally.
73 - N: number of output points to produce.
74 The final predictions are rounded to int16.
75 """
76 M = len(coeffs) - 1
77 # We'll store predictions in float32, then round to int16 at the end
78 pred = np.zeros(N, dtype=np.float32)
79 # Copy seed (cast to float32)
80 pred[:M] = seed.astype(np.float32)
82 # Predict forward
83 for j in range(M, N):
84 val = coeffs[0]
85 for k in range(1, M + 1):
86 val += coeffs[k] * pred[j - k]
87 pred[j] = val
89 # Now round to int16
90 return np.round(pred).astype(np.int16)
92# Optional: Numba-accelerated version
93try:
94 from numba import njit
97 def _predict_markov_numba(pred, coeffs, M):
98 for j in range(M, len(pred)):
99 val = coeffs[0]
100 for k in range(1, M + 1):
101 val += coeffs[k] * pred[j - k]
102 pred[j] = val
104 def predict_numba(coeffs: np.ndarray, seed: np.ndarray, N: int):
105 """
106 Numba-accelerated version of predict_naive (float32 internal arithmetic).
107 Rounds final results to int16.
108 """
109 M = len(coeffs) - 1
110 pred = np.zeros(N, dtype=np.float32)
111 pred[:M] = seed.astype(np.float32)
112 _predict_markov_numba(pred, coeffs, M)
113 return np.round(pred).astype(np.int16)
115except ImportError:
116 def predict_numba(coeffs: np.ndarray, seed: np.ndarray, N: int):
117 """
118 Fallback if Numba is not installed: use naive approach.
119 """
120 print("Numba not installed. Falling back to naive prediction.")
121 return predict_naive(coeffs, seed, N)
123########################################################################
124# 4) Main benchmarking function
125########################################################################
127def benchmark_markov_model(N=1_000_000, M=5, seed=0):
128 """
129 1) Generate random data x of length N as int16.
130 2) Fit the Markov model in multiple ways:
131 - Naive (build + solve) in float32
132 - Vectorized+NumPy (float32)
133 - Vectorized+sklearn (float32)
134 Print total fitting time for each method, plus throughput in MB/s
135 (using 2*N bytes as the data size).
136 3) Predict a new sequence in multiple ways (naive vs numba-accelerated),
137 using float32 arithmetic internally, then rounding to int16.
138 Print total prediction time + throughput in MB/s.
139 """
140 np.random.seed(seed)
141 # Generate random integers in range [-32768, 32767], but let's just do typical range
142 x = np.random.randint(-30000, 30000, size=N, dtype=np.int16)
144 # We define "bytes processed" as 2*N because each x_i is an int16 (2 bytes).
145 bytes_processed = 2 * N
146 # Build target array (still in float32 for the solver)
147 # y corresponds to x[M:], but we cast to float32
148 y = x[M:].astype(np.float32)
150 print(f"\n=== Benchmarking Markov Model (N={N}, M={M}, data=int16, coeffs=float32) ===")
152 ########################################################################
153 # Fitting: Naive approach
154 ########################################################################
155 fit_start = time.perf_counter()
157 # a) build design matrix (naive)
158 t0 = time.perf_counter()
159 A_naive = build_design_matrix_naive(x, M)
160 dt_build_naive = time.perf_counter() - t0
162 # b) solve
163 t0 = time.perf_counter()
164 c_naive = solve_lstsq_numpy(A_naive, y)
165 dt_solve_naive = time.perf_counter() - t0
167 total_naive = time.perf_counter() - fit_start
169 throughput_build_naive = bytes_processed / (dt_build_naive * 1e6)
170 throughput_solve_naive = bytes_processed / (dt_solve_naive * 1e6)
171 throughput_total_naive = bytes_processed / (total_naive * 1e6)
173 print("\n--- Fitting (Naive) ---")
174 print(f" Build matrix: {dt_build_naive:.3f}s, ~{throughput_build_naive:.2f} MB/s")
175 print(f" Solve: {dt_solve_naive:.3f}s, ~{throughput_solve_naive:.2f} MB/s")
176 print(f" TOTAL: {total_naive:.3f}s, ~{throughput_total_naive:.2f} MB/s")
178 ########################################################################
179 # Fitting: Vectorized + NumPy
180 ########################################################################
181 fit_start = time.perf_counter()
183 # a) build design matrix (vectorized)
184 t0 = time.perf_counter()
185 A_vec = build_design_matrix_vectorized(x, M)
186 dt_build_vec = time.perf_counter() - t0
188 # b) solve
189 t0 = time.perf_counter()
190 c_vec = solve_lstsq_numpy(A_vec, y)
191 dt_solve_vec = time.perf_counter() - t0
193 total_vec = time.perf_counter() - fit_start
195 throughput_build_vec = bytes_processed / (dt_build_vec * 1e6)
196 throughput_solve_vec = bytes_processed / (dt_solve_vec * 1e6)
197 throughput_total_vec = bytes_processed / (total_vec * 1e6)
199 print("\n--- Fitting (Vectorized + NumPy) ---")
200 print(f" Build matrix: {dt_build_vec:.3f}s, ~{throughput_build_vec:.2f} MB/s")
201 print(f" Solve: {dt_solve_vec:.3f}s, ~{throughput_solve_vec:.2f} MB/s")
202 print(f" TOTAL: {total_vec:.3f}s, ~{throughput_total_vec:.2f} MB/s")
204 ########################################################################
205 # Fitting: Vectorized + scikit-learn
206 ########################################################################
207 fit_start = time.perf_counter()
209 # We can re-use A_vec from above, but to measure a true "total time"
210 # for the vectorized+sklearn path, let's rebuild it.
211 t0 = time.perf_counter()
212 A_vec_sklearn = build_design_matrix_vectorized(x, M)
213 dt_build_vec_sklearn = time.perf_counter() - t0
215 t0 = time.perf_counter()
216 c_sklearn = solve_lstsq_sklearn(A_vec_sklearn, y)
217 dt_solve_sklearn = time.perf_counter() - t0
219 total_sklearn = time.perf_counter() - fit_start
221 throughput_build_sklearn = bytes_processed / (dt_build_vec_sklearn * 1e6)
222 throughput_solve_sklearn = bytes_processed / (dt_solve_sklearn * 1e6)
223 throughput_total_sklearn = bytes_processed / (total_sklearn * 1e6)
225 print("\n--- Fitting (Vectorized + Sklearn) ---")
226 print(f" Build matrix: {dt_build_vec_sklearn:.3f}s, ~{throughput_build_sklearn:.2f} MB/s")
227 print(f" Solve: {dt_solve_sklearn:.3f}s, ~{throughput_solve_sklearn:.2f} MB/s")
228 print(f" TOTAL: {total_sklearn:.3f}s, ~{throughput_total_sklearn:.2f} MB/s")
230 ########################################################################
231 # Prediction
232 ########################################################################
233 # We can pick c_vec from above as "our" coefficients for prediction.
234 # We'll seed from the last M points of x (int16).
235 seed_vals = x[-M:]
237 # 1) Naive Prediction
238 t0 = time.perf_counter()
239 pred_naive = predict_naive(c_vec, seed_vals, N)
240 dt_pred_naive = time.perf_counter() - t0
241 throughput_pred_naive = bytes_processed / (dt_pred_naive * 1e6)
243 print("\n--- Prediction (Naive) ---")
244 print(f" Predict: {dt_pred_naive:.3f}s, ~{throughput_pred_naive:.2f} MB/s")
246 # Warm up Numba
247 _ = predict_numba(c_vec, seed_vals, len(seed_vals) + 1)
249 t0 = time.perf_counter()
250 pred_numba = predict_numba(c_vec, seed_vals, N)
251 dt_pred_numba = time.perf_counter() - t0
252 throughput_pred_numba = bytes_processed / (dt_pred_numba * 1e6)
254 print("\n--- Prediction (Numba) ---")
255 print(f" Predict: {dt_pred_numba:.3f}s, ~{throughput_pred_numba:.2f} MB/s")
257 print("\n=== Done ===\n")
259########################################################################
260# Script entry point
261########################################################################
262if __name__ == "__main__":
263 # Example usage
264 benchmark_markov_model(N=5_000_000, M=5, seed=42)
265 # Adjust N, M, and seed as needed.