Write exposure windows sequentially

This commit is contained in:
2026-07-08 13:04:32 +08:00
parent 14eccc91f9
commit efba7ac306
2 changed files with 48 additions and 20 deletions

View File

@@ -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)