Skip to content

⚡ Bolt: 벡터화된 행렬 곱을 통한 진단 집계 루프 성능 최적화 #204

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
seonghobae wants to merge 1 commit into from
Closed

⚡ Bolt: 벡터화된 행렬 곱을 통한 진단 집계 루프 성능 최적화 #204

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions .jules/bolt.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
169 changes: 112 additions & 57 deletions python/fast_mlsirm/diagnostics.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Comment thread
seonghobae marked this conversation as resolved.
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)


Expand All @@ -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
Comment thread
seonghobae marked this conversation as resolved.
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)


Expand Down Expand Up @@ -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)
Comment thread
seonghobae marked this conversation as resolved.
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)


Expand All @@ -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
Comment thread
seonghobae marked this conversation as resolved.
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)


Expand Down
Loading