From efba7ac306c5b3ddd27a7997e2c2ff7b446ab3d9 Mon Sep 17 00:00:00 2001 From: Jiarui Li Date: Wed, 8 Jul 2026 13:04:32 +0800 Subject: [PATCH] Write exposure windows sequentially --- dataset.py | 19 +++++++++++---- prepare_exposure_cache.py | 49 ++++++++++++++++++++++++++------------- 2 files changed, 48 insertions(+), 20 deletions(-) diff --git a/dataset.py b/dataset.py index 952e679..49cb0aa 100644 --- a/dataset.py +++ b/dataset.py @@ -64,6 +64,7 @@ class ExposureCache: token_path = cache_dir / "exposure_token.npy" age_path = cache_dir / "exposure_age_days.npy" onset_date_path = cache_dir / "exposure_onset_date.npy" + row_index_path = cache_dir / "exposure_row_index.npy" eid_index_path = cache_dir / "exposure_eid_index.npy" eid_start_path = cache_dir / "exposure_eid_start.npy" daily_path = cache_dir / "exposure_daily.npy" @@ -73,6 +74,7 @@ class ExposureCache: token_path, age_path, onset_date_path, + row_index_path, eid_index_path, eid_start_path, daily_path, @@ -88,6 +90,7 @@ class ExposureCache: self.raw_tokens = np.load(token_path, mmap_mode="r") self.age_days = np.load(age_path, mmap_mode="r") self.onset_dates = np.load(onset_date_path, mmap_mode="r") + self.row_index = np.load(row_index_path, mmap_mode="r") self.eid_index = np.load(eid_index_path, mmap_mode="r") self.eid_start = np.load(eid_start_path, mmap_mode="r") self.daily = np.load(daily_path, mmap_mode="r") @@ -110,10 +113,15 @@ class ExposureCache: len(self.raw_tokens) != n_rows or len(self.age_days) != n_rows or len(self.onset_dates) != n_rows - or self.daily.shape[0] != n_rows - or self.monthly.shape[0] != n_rows + or len(self.row_index) != n_rows ): - raise ValueError("Exposure cache metadata/daily/monthly row counts do not match") + raise ValueError("Exposure cache sequence metadata row counts do not match") + max_window_index = int(np.max(self.row_index)) if n_rows else -1 + if ( + max_window_index >= self.daily.shape[0] + or max_window_index >= self.monthly.shape[0] + ): + raise ValueError("Exposure row index points past daily/monthly window arrays") if len(self.eid_start) != len(self.eid_index) + 1: raise ValueError("exposure_eid_start.npy must have len(eid_index) + 1") if len(self.eid_start) and int(self.eid_start[-1]) != n_rows: @@ -156,7 +164,10 @@ class ExposureCache: if n_take == 0: return out - out[real_pos[:n_take]] = np.arange(start, start + n_take, dtype=np.int64) + out[real_pos[:n_take]] = np.asarray( + self.row_index[start:start + n_take], + dtype=np.int64, + ) expected_tokens = np.asarray(self.raw_tokens[start:start + n_take], dtype=np.int64) expected_age_days = np.asarray(self.age_days[start:start + n_take], dtype=np.int64) diff --git a/prepare_exposure_cache.py b/prepare_exposure_cache.py index 8614cc1..7d4dee7 100644 --- a/prepare_exposure_cache.py +++ b/prepare_exposure_cache.py @@ -12,13 +12,17 @@ The output directory contains: exposure_token.npy int32 raw disease token per real disease event exposure_age_days.npy int32 age in days per real disease event exposure_onset_date.npy datetime64[D] onset date per real disease event + exposure_row_index.npy int64 window row per real disease event, -1 if missing exposure_eid_index.npy int64 unique eids in cache order exposure_eid_start.npy int64 start offsets, length len(eid_index) + 1 - exposure_daily.npy float32 memmap, shape (N, 1826, 4) + exposure_daily.npy float32 memmap, capacity (N, 1826, 4); + first M rows are sequential matched windows channels: tmean, tmax, tmin, rhmean - exposure_monthly.npy float32 memmap, shape (N, 241, 2) + exposure_monthly.npy float32 memmap, capacity (N, 241, 2); + first M rows are sequential matched windows channels: tmean, rhmean - exposure_quality.npy float32 memmap, shape (N, 4) + exposure_quality.npy float32 memmap, capacity (N, 4); + first M rows are matched-window quality stats n_days, n_rh_days, n_months, n_rh_months exposure_manifest.json metadata @@ -265,6 +269,7 @@ def build_exposure_cache( output_dir / "exposure_token.npy", output_dir / "exposure_age_days.npy", output_dir / "exposure_onset_date.npy", + output_dir / "exposure_row_index.npy", output_dir / "exposure_eid_index.npy", output_dir / "exposure_eid_start.npy", output_dir / "exposure_daily.npy", @@ -305,6 +310,7 @@ def build_exposure_cache( token_path = output_dir / "exposure_token.npy" age_path = output_dir / "exposure_age_days.npy" onset_date_path = output_dir / "exposure_onset_date.npy" + row_index_path = output_dir / "exposure_row_index.npy" daily_path = output_dir / "exposure_daily.npy" monthly_path = output_dir / "exposure_monthly.npy" quality_path = output_dir / "exposure_quality.npy" @@ -318,6 +324,13 @@ def build_exposure_cache( sequence_rows["onset_date"].to_numpy(dtype="datetime64[D]"), ) _write_eid_offsets(sequence_rows, output_dir) + row_index_mm = np.lib.format.open_memmap( + row_index_path, + mode="w+", + dtype=np.int64, + shape=(n_rows,), + ) + row_index_mm[:] = -1 daily_mm = np.lib.format.open_memmap( daily_path, @@ -337,9 +350,6 @@ def build_exposure_cache( dtype=np.float32, shape=(n_rows, len(QUALITY_COLUMNS)), ) - daily_mm[:] = np.nan - monthly_mm[:] = np.nan - quality_mm[:] = np.nan daily_cols = _daily_columns() monthly_cols = _monthly_columns() @@ -347,7 +357,7 @@ def build_exposure_cache( int(token): frame.reset_index(drop=True) for token, frame in sequence_rows.groupby("token", sort=False) } - matched = np.zeros(n_rows, dtype=bool) + write_offset = 0 iterator = tqdm( summary.itertuples(index=False), @@ -393,24 +403,28 @@ def build_exposure_cache( daily_df = daily_df.set_index("position").loc[common_positions].reset_index() monthly_df = monthly_df.set_index("position").loc[common_positions].reset_index() positions = common_positions.astype(np.int64) - daily_mm[positions] = _reshape_window( + n_match = len(positions) + end_offset = write_offset + n_match + daily_mm[write_offset:end_offset] = _reshape_window( daily_df, daily_cols, DAILY_LENGTH, len(DAILY_CHANNELS), ) - monthly_mm[positions] = _reshape_window( + monthly_mm[write_offset:end_offset] = _reshape_window( monthly_df, monthly_cols, MONTHLY_LENGTH, len(MONTHLY_CHANNELS), ) - quality_mm[positions, 0] = daily_df.get("n_days_nonmissing", np.nan) - quality_mm[positions, 1] = daily_df.get("n_rh_days_nonmissing", np.nan) - quality_mm[positions, 2] = monthly_df.get("n_months_nonmissing", np.nan) - quality_mm[positions, 3] = monthly_df.get("n_rh_months_nonmissing", np.nan) - matched[positions] = True + quality_mm[write_offset:end_offset, 0] = daily_df.get("n_days_nonmissing", np.nan) + quality_mm[write_offset:end_offset, 1] = daily_df.get("n_rh_days_nonmissing", np.nan) + quality_mm[write_offset:end_offset, 2] = monthly_df.get("n_months_nonmissing", np.nan) + quality_mm[write_offset:end_offset, 3] = monthly_df.get("n_rh_months_nonmissing", np.nan) + row_index_mm[positions] = np.arange(write_offset, end_offset, dtype=np.int64) + write_offset = end_offset + row_index_mm.flush() daily_mm.flush() monthly_mm.flush() quality_mm.flush() @@ -421,13 +435,16 @@ def build_exposure_cache( "data_prefix": data_prefix, "labels_file": str(Path(labels_file).resolve()), "n_rows": int(n_rows), - "matched_rows": int(matched.sum()), - "missing_rows": int((~matched).sum()), + "window_capacity_rows": int(n_rows), + "matched_rows": int(write_offset), + "missing_rows": int(n_rows - write_offset), "alignment_key": "(eid, raw_token, date_of_birth + age_days)", "requires_basic_info_column": "date_of_birth", "daily_shape": [int(n_rows), DAILY_LENGTH, len(DAILY_CHANNELS)], + "active_daily_shape": [int(write_offset), DAILY_LENGTH, len(DAILY_CHANNELS)], "daily_channels": list(DAILY_CHANNELS), "monthly_shape": [int(n_rows), MONTHLY_LENGTH, len(MONTHLY_CHANNELS)], + "active_monthly_shape": [int(write_offset), MONTHLY_LENGTH, len(MONTHLY_CHANNELS)], "monthly_channels": list(MONTHLY_CHANNELS), "quality_columns": list(QUALITY_COLUMNS), "raw_token_convention": "padding=0, checkup=1, labels.csv first row token=2",