Write exposure windows sequentially

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 len(self.row_index) != n_rows
-            or self.monthly.shape[0] != 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)

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[write_offset:end_offset, 0] = daily_df.get("n_days_nonmissing", np.nan)
-        quality_mm[positions, 1] = daily_df.get("n_rh_days_nonmissing", np.nan)
+        quality_mm[write_offset:end_offset, 1] = daily_df.get("n_rh_days_nonmissing", np.nan)
-        quality_mm[positions, 2] = monthly_df.get("n_months_nonmissing", np.nan)
+        quality_mm[write_offset:end_offset, 2] = monthly_df.get("n_months_nonmissing", np.nan)
-        quality_mm[positions, 3] = monthly_df.get("n_rh_months_nonmissing", np.nan)
+        quality_mm[write_offset:end_offset, 3] = monthly_df.get("n_rh_months_nonmissing", np.nan)
-        matched[positions] = True
+        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()),
+        "window_capacity_rows": int(n_rows),
-        "missing_rows": int((~matched).sum()),
+        "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",