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"]),
"dataset_index": int(row["dataset_index"]),
"landmark_age": float(row["landmark_age"]),
"t_query": float(row["t_query"]),
"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)
full_event = np.concatenate([seq_event, tgt_event[-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
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(
{
"patient_id": int(spec["patient_id"]),
"dataset_index": int(spec["dataset_index"]),
"sex": int(sample["sex"]),
"landmark_age": np.float32(float(spec["landmark_age"])),
"landmark_age": landmark_age,
"t_query": t_query,
"followup_end_time": np.float32(float(spec["followup_end_time"])),
"event_seq": full_event[prefix_mask].astype(np.int64, copy=False),
@@ -558,6 +562,7 @@ def _project_bi_rows(
out: list[dict[str, Any]] = []
for row_idx, row in enumerate(rows):
query_time = float(row["t_query"])
for item in projected:
matrix = item["matrix"]
historical = item["historical"][row_idx]
@@ -569,8 +574,8 @@ def _project_bi_rows(
"patient_id": row["patient_id"],
"dataset_index": row["dataset_index"],
"sex": row["sex"],
"landmark_age": float(row["landmark_age"]),
"t_query": float(row["t_query"]),
"landmark_age": query_time,
"t_query": query_time,
"followup_end_time": float(row["followup_end_time"]),
"horizon": float(horizon),
"formed_mode": formed_mode,