Write exposure windows sequentially
This commit is contained in:
19
dataset.py
19
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)
|
||||
|
||||
Reference in New Issue
Block a user