From c3a293214db3de408b95d235718eaaa540af7e03 Mon Sep 17 00:00:00 2001 From: seonghobae <8172694+seonghobae@users.noreply.github.com> Date: Tue, 21 Jul 2026 07:45:43 +0000 Subject: [PATCH] perf: Replace nested Python loops with vectorized matrix multiplications MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Python의 중첩 루프 구조 내에서 발생하던 불리언 마스크 복사 및 생성 오버헤드를 C/BLAS 최적화된 행렬 곱으로 전환하여 성능을 획기적으로 개선합니다. - `_binary_stratum_fit` - `_binary_stratum_item_fit` - `_categorical_stratum_fit` - `_categorical_stratum_item_fit` 위의 함수들이 그룹 및 아이템 조합을 탐색할 때 2D mapping mask 행렬(`strata_mask.T`)과 관측값의 행렬 곱(`@`)을 사용하도록 리팩토링하였습니다. --- .jules/bolt.md | 4 + python/fast_mlsirm/diagnostics.py | 169 ++++++++++++++++++++---------- 2 files changed, 116 insertions(+), 57 deletions(-) diff --git a/.jules/bolt.md b/.jules/bolt.md index 73e3fbaf9..10ace605e 100644 --- a/.jules/bolt.md +++ b/.jules/bolt.md @@ -33,3 +33,7 @@ ## 2025-05-19 - Dot product scalar gradients allocation **Learning:** During gradient calculation, `float((e * (-gamma * distance)).sum())` creates two full-size `(N, J)` arrays: one for the scaled distance and one for the element-wise multiplication before reduction. **Action:** Replace `(A * B).sum()` with `np.vdot(A, B)` when scalar reduction is needed over matrix multiplication (where `B` can incorporate scalars naturally like `-gamma * np.vdot(A, B)`). This entirely avoids the 2D array allocation overhead and yields order-of-magnitude improvements in scalar gradient components. + +## 2025-05-19 - Vectorized nested looping over groups and items +**Learning:** In diagnostics aggregation functions (like `_binary_stratum_item_fit` or `_categorical_stratum_item_fit`), iterating over unique group IDs (strata) and then iterating over each item dimension creates a highly inefficient nested Python loop structure that heavily relies on boolean index slicing (`where = observed & (ids[:, None] == value)`). +**Action:** Replace nested loops across unique group IDs and item dimensions with 2D mapping mask matrix multiplications (`@`). Convert strata IDs into a 2D boolean float mask (`strata_mask.T` of shape `V x N`, where `V` is the number of unique strata), and then multiply by the data arrays (of shape `N x J`). For example, `counts = strata_mask.T @ observed.astype(np.float64)` directly computes the `V x J` output matrix without any Python loops or boolean slicing allocations. Use `np.nonzero` or zip over non-zero elements to format the output. diff --git a/python/fast_mlsirm/diagnostics.py b/python/fast_mlsirm/diagnostics.py index f719692a2..138494298 100644 --- a/python/fast_mlsirm/diagnostics.py +++ b/python/fast_mlsirm/diagnostics.py @@ -575,21 +575,40 @@ def _binary_stratum_fit( ids = _person_strata(strata, y.shape[0], id_name) loglik = np.where(observed, y * np.log(prob) + (1.0 - y) * np.log1p(-prob), 0.0) + + unique_ids = np.unique(ids) + strata_mask = (ids[:, None] == unique_ids[None, :]).astype(np.float64) + obs_f64 = observed.astype(np.float64) + + counts = strata_mask.T @ obs_f64.sum(axis=1) + scores = strata_mask.T @ (y * obs_f64).sum(axis=1) + expecteds = strata_mask.T @ (prob * obs_f64).sum(axis=1) + raws = strata_mask.T @ (residual * obs_f64).sum(axis=1) + varsums = strata_mask.T @ (variance * obs_f64).sum(axis=1) + res_sqs = strata_mask.T @ (residual * residual * obs_f64).sum(axis=1) + chisqs = strata_mask.T @ (pearson_sq * obs_f64).sum(axis=1) + lls = strata_mask.T @ (loglik * obs_f64).sum(axis=1) + rows = [] - for value in np.unique(ids): - rows.append( - _binary_scope_row( - float(value), - ids[:, None] == value, - y, - observed, - prob, - variance, - residual, - pearson_sq, - loglik, - ) - ) + for v_idx in range(len(unique_ids)): + c = float(counts[v_idx]) + vs = float(varsums[v_idx]) + raw = float(raws[v_idx]) + ch = float(chisqs[v_idx]) + ll = float(lls[v_idx]) + rows.append(( + float(unique_ids[v_idx]), + c, + float(scores[v_idx]), + float(expecteds[v_idx]), + raw, + raw / float(np.sqrt(max(vs, 1e-12))), + float(res_sqs[v_idx]) / max(vs, 1e-12), + ch / max(c, 1.0), + ll, + -2.0 * ll, + ch, + )) return _binary_scope_table(id_name, rows) @@ -608,30 +627,42 @@ def _binary_stratum_item_fit( ids = _person_strata(strata, y.shape[0], id_name) loglik = np.where(observed, y * np.log(prob) + (1.0 - y) * np.log1p(-prob), 0.0) + + unique_ids = np.unique(ids) + strata_mask = (ids[:, None] == unique_ids[None, :]).astype(np.float64) + obs_f64 = observed.astype(np.float64) + + counts = strata_mask.T @ obs_f64 + scores = strata_mask.T @ (y * obs_f64) + expecteds = strata_mask.T @ (prob * obs_f64) + raws = strata_mask.T @ (residual * obs_f64) + varsums = strata_mask.T @ (variance * obs_f64) + res_sqs = strata_mask.T @ (residual * residual * obs_f64) + chisqs = strata_mask.T @ (pearson_sq * obs_f64) + lls = strata_mask.T @ (loglik * obs_f64) + rows = [] - for value in np.unique(ids): - row_mask = ids == value - for item in range(y.shape[1]): - scope = np.zeros_like(observed, dtype=bool) - scope[row_mask, item] = True - if np.any(observed & scope): - rows.append( - ( - float(value), - float(item), - *_binary_scope_row( - 0.0, - scope, - y, - observed, - prob, - variance, - residual, - pearson_sq, - loglik, - )[1:], - ) - ) + v_idx, j_idx = np.nonzero(counts > 0) + for v, j in zip(v_idx, j_idx): + c = float(counts[v, j]) + vs = float(varsums[v, j]) + raw = float(raws[v, j]) + ch = float(chisqs[v, j]) + ll = float(lls[v, j]) + rows.append(( + float(unique_ids[v]), + float(j), + c, + float(scores[v, j]), + float(expecteds[v, j]), + raw, + raw / float(np.sqrt(max(vs, 1e-12))), + float(res_sqs[v, j]) / max(vs, 1e-12), + ch / max(c, 1.0), + ll, + -2.0 * ll, + ch, + )) return _binary_scope_item_table(id_name, rows) @@ -730,12 +761,28 @@ def _categorical_stratum_fit( return None ids = _person_strata(strata, observed.shape[0], id_name) + + unique_ids = np.unique(ids) + strata_mask = (ids[:, None] == unique_ids[None, :]).astype(np.float64) + obs_f64 = observed.astype(np.float64) + + counts = strata_mask.T @ obs_f64.sum(axis=1) + lls = strata_mask.T @ (entry_loglik * obs_f64).sum(axis=1) + chisqs = strata_mask.T @ (entry_chisq * obs_f64).sum(axis=1) + rows = [] - for value in np.unique(ids): - where = observed & (ids[:, None] == value) - rows.append( - _categorical_scope_row(float(value), where, entry_loglik, entry_chisq) - ) + for v_idx in range(len(unique_ids)): + c = float(counts[v_idx]) + ll = float(lls[v_idx]) + ch = float(chisqs[v_idx]) + rows.append(( + float(unique_ids[v_idx]), + c, + ll, + -2.0 * ll, + ch, + ch / max(c, 1.0), + )) return _categorical_scope_table(id_name, rows) @@ -750,22 +797,30 @@ def _categorical_stratum_item_fit( return None ids = _person_strata(strata, observed.shape[0], id_name) + + unique_ids = np.unique(ids) + strata_mask = (ids[:, None] == unique_ids[None, :]).astype(np.float64) + obs_f64 = observed.astype(np.float64) + + counts = strata_mask.T @ obs_f64 + lls = strata_mask.T @ (entry_loglik * obs_f64) + chisqs = strata_mask.T @ (entry_chisq * obs_f64) + rows = [] - for value in np.unique(ids): - row_mask = ids == value - for item in range(observed.shape[1]): - where = np.zeros_like(observed, dtype=bool) - where[row_mask, item] = observed[row_mask, item] - if np.any(where): - rows.append( - ( - float(value), - float(item), - *_categorical_scope_row(0.0, where, entry_loglik, entry_chisq)[ - 1: - ], - ) - ) + v_idx, j_idx = np.nonzero(counts > 0) + for v, j in zip(v_idx, j_idx): + c = float(counts[v, j]) + ll = float(lls[v, j]) + ch = float(chisqs[v, j]) + rows.append(( + float(unique_ids[v]), + float(j), + c, + ll, + -2.0 * ll, + ch, + ch / max(c, 1.0), + )) return _categorical_scope_item_table(id_name, rows)