Here is a Python script that generates synthetic data and uses scalar_spectrum.mojo
#!/usr/bin/env python3
from __future__ import annotations
import argparse
import importlib
from pathlib import Path
import sys
import time
import numpy as np
def gaussian_latitudes_deg(nlat: int) -> np.ndarray:
"""Return Gaussian latitude centers in degrees (south->north)."""
mu, _weights = np.polynomial.legendre.leggauss(nlat)
return np.degrees(np.arcsin(mu))
def make_field(
lat_deg: np.ndarray,
lon_deg: np.ndarray,
*,
seed: int,
noise_std: float,
) -> np.ndarray:
lat_rad = np.radians(lat_deg)[:, None]
lon_rad = np.radians(lon_deg)[None, :]
# Structured signal with several zonal wavenumbers.
mode_2 = 1.30 * np.cos(lat_rad) ** 2 * np.cos(2.0 * lon_rad)
mode_5 = 0.75 * np.sin(2.0 * lat_rad) * np.sin(5.0 * lon_rad + 0.7)
mode_9 = 0.40 * np.cos(3.0 * lat_rad) * np.cos(9.0 * lon_rad - 0.3)
# Broad large-scale meridional background.
background = 0.25 * (1.5 * np.sin(lat_rad) ** 2 - 0.5)
# Localized anomaly.
lat0 = np.radians(35.0)
lon0 = np.radians(120.0)
dist2 = (lat_rad - lat0) ** 2 + (np.angle(np.exp(1j * (lon_rad - lon0)))) ** 2
blob = 0.55 * np.exp(-dist2 / (2.0 * 0.28**2))
rng = np.random.default_rng(seed)
noise = rng.normal(loc=0.0, scale=noise_std, size=(lat_deg.size, lon_deg.size))
return (background + mode_2 + mode_5 + mode_9 + blob + noise).astype(np.float64)
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description=(
"Generate synthetic scalar Gaussian-grid data and run the Mojo "
"scalar_spectrum kernel."
)
)
parser.add_argument(
"--nlat",
type=int,
default=800,
help="Number of Gaussian latitudes (default: 800, stress-test size).",
)
parser.add_argument(
"--nlon",
type=int,
default=1600,
help="Number of longitudes (default: 1600, stress-test size).",
)
parser.add_argument("--seed", type=int, default=42, help="Random seed for reproducible noise.")
parser.add_argument("--noise-std", type=float, default=0.05, help="Std-dev of additive white noise.")
parser.add_argument(
"--max-degree",
type=int,
default=0,
help="Maximum harmonic degree (0 means full nlon/2 truncation used by the Mojo kernel).",
)
parser.add_argument(
"--mojo-dir",
type=Path,
default=None,
help="Directory containing scalar_spectrum.mojo (default: auto-detect).",
)
parser.add_argument(
"--output",
type=Path,
default=Path("tests/artifacts/synthetic_scalar_input.npz"),
help="Output .npz path for the generated synthetic field.",
)
return parser.parse_args()
def _resolve_mojo_dir(user_mojo_dir: Path | None) -> Path:
if user_mojo_dir is not None:
mojo_dir = user_mojo_dir.expanduser().resolve()
if not (mojo_dir / "scalar_spectrum.mojo").exists():
raise FileNotFoundError(
f"scalar_spectrum.mojo not found in --mojo-dir: {mojo_dir}"
)
return mojo_dir
script_dir = Path(__file__).resolve().parent
candidates = [
script_dir / "mojo",
script_dir.parent / "mojo",
Path.cwd() / "mojo",
Path.cwd(),
]
for cand in candidates:
if (cand / "scalar_spectrum.mojo").exists():
return cand.resolve()
raise FileNotFoundError(
"Could not locate scalar_spectrum.mojo automatically. "
"Pass --mojo-dir <path>."
)
def _load_scalar_kernel(mojo_dir: Path):
if str(mojo_dir) not in sys.path:
sys.path.insert(0, str(mojo_dir))
import mojo.importer # noqa: F401 # pyright: ignore[reportMissingImports]
return importlib.import_module("scalar_spectrum")
def main() -> None:
args = parse_args()
if args.nlat < 4:
raise ValueError("nlat must be >= 4.")
if args.nlon < 8:
raise ValueError("nlon must be >= 8.")
if args.nlon % 2 != 0:
raise ValueError("nlon must be even for typical spectral truncation use.")
if args.noise_std < 0.0:
raise ValueError("noise-std must be >= 0.")
if args.max_degree < 0:
raise ValueError("max-degree must be >= 0.")
lat_deg = gaussian_latitudes_deg(args.nlat)
lon_deg = np.linspace(0.0, 360.0, args.nlon, endpoint=False, dtype=np.float64)
field = make_field(lat_deg, lon_deg, seed=args.seed, noise_std=args.noise_std)
# Mirror the Mojo kernel behavior: if max_degree <= 0, use nlat then clamp to nlon/2.
# For the stress-test default (800x1600), this gives max_degree=800.
kernel_limit = args.nlon // 2
full_kernel_degree = min(args.nlat, kernel_limit)
suggested_max_degree = min(args.nlat - 1, kernel_limit)
effective_max_degree = full_kernel_degree if args.max_degree == 0 else min(args.max_degree, full_kernel_degree)
args.output.parent.mkdir(parents=True, exist_ok=True)
np.savez_compressed(
args.output,
field=field,
lat_deg=lat_deg,
lon_deg=lon_deg,
nlat=np.int64(args.nlat),
nlon=np.int64(args.nlon),
seed=np.int64(args.seed),
noise_std=np.float64(args.noise_std),
suggested_max_degree=np.int64(suggested_max_degree),
full_kernel_degree=np.int64(full_kernel_degree),
effective_max_degree=np.int64(effective_max_degree),
)
mojo_dir = _resolve_mojo_dir(args.mojo_dir)
print(f"Wrote synthetic payload: {args.output}")
print(f"field shape: {field.shape}, dtype={field.dtype}")
print(f"suggested_max_degree (API-safe): {suggested_max_degree}")
print(f"full_kernel_degree (Mojo max): {full_kernel_degree}")
print(f"loading scalar_spectrum.mojo from: {mojo_dir}")
print(f"running scalar_spectrum with max_degree={effective_max_degree} ...")
scalar_mod = _load_scalar_kernel(mojo_dir)
t0 = time.perf_counter()
power = np.asarray(
scalar_mod.scalar_spectrum(
np.ascontiguousarray(field, dtype=np.float64),
np.ascontiguousarray(lat_deg, dtype=np.float64),
int(effective_max_degree),
),
dtype=np.float64,
)
dt = time.perf_counter() - t0
print(f"kernel runtime: {dt:.3f} s")
print(f"output spectrum length: {power.size}")
print(f"power sum: {power.sum():.6e}")
print(f"power[0:5]: {power[:5]}")
if __name__ == "__main__":
main()
[✓][2026-03-26 15:09:46][USER@HOST][mojo-spectrum]-[36610748.HOST-01-ib] {module-build} - 22s
ρ pixi run python "scripts/generate_synthetic_scalar_input.py" --nlat 800 --nlon 1600 --max-degree 800
⠁
⠁ activating environment
⠁ activating environment
Wrote synthetic payload: tests/artifacts/synthetic_scalar_input.npz
field shape: (800, 1600), dtype=float64
suggested_max_degree (API-safe): 799
full_kernel_degree (Mojo max): 800
loading scalar_spectrum.mojo from: /fs/site8/eccc/cmd/cmds/yor000/gitlab.science.gc.ca/yor000/mojo-experiments/tke-mojo/mojo
running scalar_spectrum with max_degree=800 ...
[mojo] grid 800 x 1600 max_degree= 800
[mojo] Stage 0: extracting field to native buffers ...
[mojo] Stage 0 time: 0.163210318016354 s
[mojo] Stage 1: computing Gauss-Legendre weights ...
[mojo] Stage 1 time: 0.006960558996070176 s
[mojo] Stage 2: area-weighted mean ...
[mojo] Stage 2 time: 0.0018916879780590534 s
[mojo] Stage 3: weighted variance target ...
[mojo] Stage 3 time: 0.0017297330196015537 s
[mojo] Stage 4: Fourier precompute (twiddle table + nlat x n_m DFTs) ...
[mojo] Stage 4 time: 1.4660470120143145 s
[mojo] Stage 5: Legendre projection (streaming recurrence) ...
[mojo] Stage 5 time: 0.5123696019873023 s
[mojo] Stage 6: variance rescaling ...
[mojo] Stage 6 time: 1.5160185284912586e-06 s
[mojo] total kernel time: 2.152271553990431 s
[mojo] done.
kernel runtime: 2.153 s
output spectrum length: 800
power sum: 6.405006e-01
power[0:5]: [0.00076523 0.45140083 0.00095973 0.00074946 0.000495 ]
EDIT : By the way I keep trying to make vectorize work and Opus 4.6, Gemini 3.1 Pro and Composer 2 all give the same cute message that starts with “Ahhh the joys of writing code with a bleeding edge language…” and then proceed to ignore the Nightly documentation I give it 