Keep burden landmark query times consistent

This commit is contained in:
2026-06-26 12:24:51 +08:00
parent 7ee29329f5
commit fb873326ee

View File

@@ -235,7 +235,6 @@ def _row_to_worker_spec(row: dict[str, Any]) -> dict[str, Any]:
"patient_id": int(row["patient_id"]), "patient_id": int(row["patient_id"]),
"dataset_index": int(row["dataset_index"]), "dataset_index": int(row["dataset_index"]),
"landmark_age": float(row["landmark_age"]), "landmark_age": float(row["landmark_age"]),
"t_query": float(row["t_query"]),
"followup_end_time": float(row["followup_end_time"]), "followup_end_time": float(row["followup_end_time"]),
} }
@@ -253,14 +252,19 @@ def _materialize_worker_rows(
tgt_time = np.asarray(sample["target_time_seq"], dtype=np.float32) tgt_time = np.asarray(sample["target_time_seq"], dtype=np.float32)
full_event = np.concatenate([seq_event, tgt_event[-1:]]) full_event = np.concatenate([seq_event, tgt_event[-1:]])
full_time = np.concatenate([seq_time, tgt_time[-1:]]) full_time = np.concatenate([seq_time, tgt_time[-1:]])
t_query = np.float32(float(spec["t_query"])) t_query = np.float32(float(spec["landmark_age"]))
prefix_mask = full_time <= t_query prefix_mask = full_time <= t_query
landmark_age = np.float32(float(spec["landmark_age"]))
if landmark_age != t_query:
raise RuntimeError(
f"landmark_age and t_query diverged for dataset_index={spec['dataset_index']}"
)
rows.append( rows.append(
{ {
"patient_id": int(spec["patient_id"]), "patient_id": int(spec["patient_id"]),
"dataset_index": int(spec["dataset_index"]), "dataset_index": int(spec["dataset_index"]),
"sex": int(sample["sex"]), "sex": int(sample["sex"]),
"landmark_age": np.float32(float(spec["landmark_age"])), "landmark_age": landmark_age,
"t_query": t_query, "t_query": t_query,
"followup_end_time": np.float32(float(spec["followup_end_time"])), "followup_end_time": np.float32(float(spec["followup_end_time"])),
"event_seq": full_event[prefix_mask].astype(np.int64, copy=False), "event_seq": full_event[prefix_mask].astype(np.int64, copy=False),
@@ -558,6 +562,7 @@ def _project_bi_rows(
out: list[dict[str, Any]] = [] out: list[dict[str, Any]] = []
for row_idx, row in enumerate(rows): for row_idx, row in enumerate(rows):
query_time = float(row["t_query"])
for item in projected: for item in projected:
matrix = item["matrix"] matrix = item["matrix"]
historical = item["historical"][row_idx] historical = item["historical"][row_idx]
@@ -569,8 +574,8 @@ def _project_bi_rows(
"patient_id": row["patient_id"], "patient_id": row["patient_id"],
"dataset_index": row["dataset_index"], "dataset_index": row["dataset_index"],
"sex": row["sex"], "sex": row["sex"],
"landmark_age": float(row["landmark_age"]), "landmark_age": query_time,
"t_query": float(row["t_query"]), "t_query": query_time,
"followup_end_time": float(row["followup_end_time"]), "followup_end_time": float(row["followup_end_time"]),
"horizon": float(horizon), "horizon": float(horizon),
"formed_mode": formed_mode, "formed_mode": formed_mode,