Updated moon fetch script to use less memory

This commit is contained in:
connor committed 2026-10-15 10:08:41 -04:00
1 parent f8b515e198
commit 6888b586dc
2 files changed
+59 -25

No files matched your search

+1
View File
@@ -1,3 +1,4 @@
/target /target
/examples/models/nasa/ /examples/models/nasa/
/data/ /data/
__pycache__/
+58 -25
View File
@@ -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 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. 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 import hashlib
@@ -67,6 +68,7 @@ LINES = {"N": 5321, "S": 5320}
ROWS = 70 * PPD # 5320 lines from the equator to 70 degrees ROWS = 70 * PPD # 5320 lines from the equator to 70 degrees
OUT_W, OUT_H = 16384, 8192 OUT_W, OUT_H = 16384, 8192
SCALE = 100_000 # stored value per unit I/F SCALE = 100_000 # stored value per unit I/F
STRIP = 512 # rows or columns resampled at a time
def url(band, tile): def url(band, tile):
@@ -106,18 +108,24 @@ def fetch(band, tile, sha, directory):
return path return path
def read_tile(path, hemisphere): def open_tile(path, hemisphere):
"""The tile's 5320 lines from the equator outward, as rows from north to south.""" """The tile's 5320 lines from the equator outward, as rows from north to south,
data = np.fromfile(path, dtype="<f4", offset=RECORD) 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) data = data.reshape(LINES[hemisphere], WIDTH)
if hemisphere == "N": if hemisphere == "N":
data = data[1:] # the line centered just above 70 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 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): def row_weights(rows, north):
"""cos(latitude) of each row's center, for area weighting.""" """cos(latitude) of each row's center, for area weighting."""
k = np.arange(rows) + 0.5 k = np.arange(rows) + 0.5
@@ -143,30 +151,55 @@ def resample(image, axis, size):
def build(root): 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 = root / "source"
source.mkdir(parents=True, exist_ok=True) source.mkdir(parents=True, exist_ok=True)
# Columns from -180: the tiles starting at 180 E and 270 E come first. # Columns from -180: the tiles starting at 180 E and 270 E come first.
order = sorted(TILES, key=lambda t: (t[1] + 180) % 360) 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"] halves = [
south = [read_tile(fetch(643, t, s, source), "S") for t, w, h, s in order if h == "S"] ([open_tile(fetch(643, t, s, source), h) for t, w, h, s in order if h == hemisphere],
mosaic = np.vstack([np.hstack(north), np.hstack(south)]) row_weights(ROWS, hemisphere == "N"))
weights = np.concatenate([row_weights(ROWS, True), row_weights(ROWS, False)]) for hemisphere in "NS"
valid = ~np.isnan(mosaic) ]
w = weights[:, None] * valid # The mean of the valid pixels, by area.
mean = float(np.nansum(mosaic * weights[:, None]) / w.sum()) total = weight = 0.0
missing = int((~valid).sum()) missing = 0
print(f"mosaic {mosaic.shape[1]} x {mosaic.shape[0]}: mean I/F {mean:.5f}, " for tiles, weights in halves:
f"{missing} missing pixels ({100 * missing / mosaic.size:.4f} %)") for start in range(0, ROWS, STRIP):
mosaic[~valid] = mean rows = mosaic_rows(tiles, start, start + STRIP)
cap = np.full((20 * PPD, mosaic.shape[1]), mean, dtype=np.float32) w = weights[start:start + STRIP, None]
full = np.vstack([cap, mosaic, cap]) # 90 N to 90 S valid = ~np.isnan(rows)
print(f"resampling {full.shape[1]} x {full.shape[0]} to {OUT_W} x {OUT_H}", flush=True) total += float(np.nansum(rows * w))
full = resample(full, 1, OUT_W) weight += float((w * valid).sum())
full = resample(full, 0, OUT_H) missing += int((~valid).sum())
values = np.clip(np.round(full * SCALE), 0, 65535).astype(np.uint16) 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" out = root / "wac-hapke-643.png"
Image.fromarray(values).save(out) 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): def band_mean(band):