Updated moon fetch script to use less memory

scripts/fetch-moon.py

@@ -22,7 +22,8 @@ bands relative to 643 nm over the whole mosaic, weighted by area: it downloads e
 tile of every band in turn (56 files, 8 GB) and keeps none of them. That is how the spectrum in crates/resolved/src/moon.rs
 was found.
 
-Needs numpy and Pillow.
+Needs numpy and Pillow. Building the image takes about 20 s and 2.6 GB of memory once
+the tiles are downloaded (about a minute at 50 MB/s, with a pause between files).
 """
 
 import hashlib
@@ -67,6 +68,7 @@ LINES = {"N": 5321, "S": 5320}
 ROWS = 70 * PPD  # 5320 lines from the equator to 70 degrees
 OUT_W, OUT_H = 16384, 8192
 SCALE = 100_000  # stored value per unit I/F
+STRIP = 512  # rows or columns resampled at a time
 
 
 def url(band, tile):
@@ -106,18 +108,24 @@ def fetch(band, tile, sha, directory):
     return path
 
 
-def read_tile(path, hemisphere):
-    """The tile's 5320 lines from the equator outward, as rows from north to south."""
-    data = np.fromfile(path, dtype="<f4", offset=RECORD)
+def open_tile(path, hemisphere):
+    """The tile's 5320 lines from the equator outward, as rows from north to south,
+    mapped from the file rather than read into memory."""
+    data = np.memmap(path, dtype="<f4", mode="r", offset=RECORD)
     data = data.reshape(LINES[hemisphere], WIDTH)
     if hemisphere == "N":
         data = data[1:]  # the line centered just above 70 N
-    data = data.copy()
-    # PDS special values (null, saturation) are all below -3e38.
-    data[data < -1e30] = np.nan
     return data
 
 
+def mosaic_rows(tiles, start, stop):
+    """Rows `start` to `stop` of the hemisphere `tiles` make side by side, as float32
+    with the PDS special values (null, saturation; all below -3e38) as NaN."""
+    rows = np.hstack([t[start:stop] for t in tiles]).astype(np.float32)
+    rows[rows < -1e30] = np.nan
+    return rows
+
+
 def row_weights(rows, north):
     """cos(latitude) of each row's center, for area weighting."""
     k = np.arange(rows) + 0.5
@@ -143,30 +151,55 @@ def resample(image, axis, size):
 
 
 def build(root):
+    """Works through the mosaic STRIP rows (or columns) at a time, so memory stays near
+    the 0.9 GB of the half-resampled image instead of the many GB a whole-mosaic
+    resampling in float64 takes."""
     source = root / "source"
     source.mkdir(parents=True, exist_ok=True)
     # Columns from -180: the tiles starting at 180 E and 270 E come first.
     order = sorted(TILES, key=lambda t: (t[1] + 180) % 360)
-    north = [read_tile(fetch(643, t, s, source), "N") for t, w, h, s in order if h == "N"]
-    south = [read_tile(fetch(643, t, s, source), "S") for t, w, h, s in order if h == "S"]
-    mosaic = np.vstack([np.hstack(north), np.hstack(south)])
-    weights = np.concatenate([row_weights(ROWS, True), row_weights(ROWS, False)])
-    valid = ~np.isnan(mosaic)
-    w = weights[:, None] * valid
-    mean = float(np.nansum(mosaic * weights[:, None]) / w.sum())
-    missing = int((~valid).sum())
-    print(f"mosaic {mosaic.shape[1]} x {mosaic.shape[0]}: mean I/F {mean:.5f}, "
-          f"{missing} missing pixels ({100 * missing / mosaic.size:.4f} %)")
-    mosaic[~valid] = mean
-    cap = np.full((20 * PPD, mosaic.shape[1]), mean, dtype=np.float32)
-    full = np.vstack([cap, mosaic, cap])  # 90 N to 90 S
-    print(f"resampling {full.shape[1]} x {full.shape[0]} to {OUT_W} x {OUT_H}", flush=True)
-    full = resample(full, 1, OUT_W)
-    full = resample(full, 0, OUT_H)
-    values = np.clip(np.round(full * SCALE), 0, 65535).astype(np.uint16)
+    halves = [
+        ([open_tile(fetch(643, t, s, source), h) for t, w, h, s in order if h == hemisphere],
+         row_weights(ROWS, hemisphere == "N"))
+        for hemisphere in "NS"
+    ]
+    # The mean of the valid pixels, by area.
+    total = weight = 0.0
+    missing = 0
+    for tiles, weights in halves:
+        for start in range(0, ROWS, STRIP):
+            rows = mosaic_rows(tiles, start, start + STRIP)
+            w = weights[start:start + STRIP, None]
+            valid = ~np.isnan(rows)
+            total += float(np.nansum(rows * w))
+            weight += float((w * valid).sum())
+            missing += int((~valid).sum())
+    mean = total / weight
+    pixels = 2 * ROWS * 4 * WIDTH
+    print(f"mosaic {4 * WIDTH} x {2 * ROWS}: mean I/F {mean:.5f}, "
+          f"{missing} missing pixels ({100 * missing / pixels:.4f} %)")
+    # Missing pixels and the polar caps get the mean. Resampled across, then down.
+    cap = 20 * PPD
+    height = 2 * cap + 2 * ROWS  # 90 N to 90 S
+    print(f"resampling {4 * WIDTH} x {height} to {OUT_W} x {OUT_H}", flush=True)
+    across = np.empty((height, OUT_W), dtype=np.float32)
+    across[:cap] = mean
+    across[height - cap:] = mean
+    for k, (tiles, _) in enumerate(halves):
+        for start in range(0, ROWS, STRIP):
+            rows = mosaic_rows(tiles, start, start + STRIP)
+            rows[np.isnan(rows)] = mean
+            top = cap + k * ROWS + start
+            across[top:top + len(rows)] = resample(rows, 1, OUT_W)
+    values = np.empty((OUT_H, OUT_W), dtype=np.uint16)
+    low, high = np.inf, -np.inf
+    for start in range(0, OUT_W, STRIP):
+        down = resample(across[:, start:start + STRIP], 0, OUT_H)
+        low, high = min(low, float(down.min())), max(high, float(down.max()))
+        values[:, start:start + STRIP] = np.clip(np.round(down * SCALE), 0, 65535)
     out = root / "wac-hapke-643.png"
     Image.fromarray(values).save(out)
-    print(f"wrote {out}: I/F from {full.min():.4f} to {full.max():.4f}")
+    print(f"wrote {out}: I/F from {low:.4f} to {high:.4f}")
 
 
 def band_mean(band):