nrlmsise00/scripts/generate-tables.py

#!/usr/bin/env python3
"""Generates src/tables.rs: the NRLMSISE-00 coefficient tables as Rust statics.

Fetches Dominik Brodowski's C translation of the NRL Fortran, parses the block data out of
nrlmsise-00_data.c, and writes the numbers as statics with the shapes the model reads them
at. The source is not kept. What is committed is the coefficients, which are the model
itself rather than anyone's code.

The decimal digits are emitted exactly as the reference writes them. Both C and Rust round
decimal to binary correctly, so the compiled doubles are bit-identical to the reference's.

The model itself is the work of Picone, Hedin and Drob at the US Naval Research Laboratory.
Needs network access; rerun it only to regenerate.

Usage: scripts/generate-tables.py
"""

import hashlib
import re
import subprocess
import urllib.request
from datetime import date, timezone, datetime
from pathlib import Path

BASE = "https://raw.githubusercontent.com/magnific0/nrlmsise-00/master"
ROOT = Path(__file__).resolve().parent.parent
OUT = ROOT / "src" / "tables.rs"

# The tables the model reads, with the shape it reads them at and what each one sets.
#
# ptm is declared double[50] in the data file but extern double[10] in the model, and only
# its first ten values are initialized. Ten is what is emitted: indices 0 to 8 are read and
# index 9 is the last one the data file gives.
TABLES = [
    ("pt", "PT", (150,), "Exospheric temperature."),
    ("pd", "PD", (9, 150), "Species densities at the lower boundary, and TLB."),
    ("ps", "PS", (150,), "The 120 km temperature gradient, S."),
    ("pdl", "PDL", (2, 25), "Turbopause heights, chemistry corrections and scale lengths."),
    ("ptm", "PTM", (10,), "Lower boundary temperatures, gradient and height."),
    ("pdm", "PDM", (8, 10), "Lower boundary densities, turbopause and correction heights."),
    ("ptl", "PTL", (4, 100), "Lower thermosphere temperature nodes, TN1."),
    ("pma", "PMA", (10, 100), "Mesosphere and stratosphere temperature nodes, TN2 and TN3."),
    ("pavgm", "PAVGM", (10,), "Global average temperatures the TN2 and TN3 nodes scale."),
]

# sam[100] is in the data file and no part of NRLMSISE-00 reads it, so it is not emitted.
UNUSED = ["sam"]


def fetch(name: str) -> str:
    with urllib.request.urlopen(f"{BASE}/{name}", timeout=120) as response:
        return response.read().decode("ascii")


def numbers(text: str) -> list[str]:
    """The numeric literals of a C initializer, as the digits the reference writes."""
    return re.findall(r"[-+]?\d*\.?\d+(?:[Ee][-+]?\d+)?", text)


def parse(source: str, name: str, shape: tuple[int, ...]) -> tuple[list[str], list[str]]:
    """The literals of one table, and the comment labelling each row of a 2D one."""
    pattern = rf"double\s+{name}\s*(?:\[\d+\])+\s*=\s*\{{(.*?)\}};"
    match = re.search(pattern, source, re.S)
    if match is None:
        raise SystemExit(f"{name}: not found in the reference data file")
    body = match.group(1)

    # Row labels, where the reference gives them: a comment between one row's closing brace
    # and the next row's opening brace labels the row that follows it.
    labels = ["" for _ in range(shape[0] if len(shape) > 1 else 1)]
    if len(shape) > 1:
        head = source[: match.start()]
        first = re.findall(r"/\*\s*([^*]*?)\s*\*/\s*$", head.strip())
        labels[0] = first[0] if first else ""
        for index, label in enumerate(re.findall(r"\}\s*,\s*/\*\s*(.*?)\s*\*/\s*\{", body, re.S)):
            if index + 1 < len(labels):
                labels[index + 1] = " ".join(label.split())

    literals = numbers(re.sub(r"/\*.*?\*/", "", body, flags=re.S))
    wanted = 1
    for dimension in shape:
        wanted *= dimension
    if len(literals) > wanted:
        raise SystemExit(f"{name}: {len(literals)} literals for a shape of {shape}")
    # ptm is initialized short of its declared length; C zero-fills the rest.
    literals += ["0.00000E+00"] * (wanted - len(literals))
    return literals, labels


def rust(literal: str) -> str:
    """One C literal as a Rust one, with the same digits."""
    text = literal.lower().replace("e+0", "e").replace("e-0", "e-")
    if text.endswith("e"):
        text = text[:-1]
    mantissa, _, exponent = text.partition("e")
    if "." not in mantissa:
        mantissa += ".0"
    return mantissa if exponent in ("", "0") else f"{mantissa}e{exponent}"


def rows(literals: list[str], per_line: int = 5) -> list[str]:
    out = []
    for start in range(0, len(literals), per_line):
        out.append(", ".join(rust(value) for value in literals[start : start + per_line]) + ",")
    return out


def emit(source: str, digest: str) -> str:
    lines = [
        "//! The NRLMSISE-00 coefficient tables.",
        "//!",
        "//! Generated by `scripts/generate-tables.py` from Dominik Brodowski's C translation",
        "//! of the NRL Fortran; do not edit. The coefficients are the model, which is the work",
        "//! of Picone, Hedin and Drob at the US Naval Research Laboratory.",
        "//!",
        f"//! Source: <{BASE}>",
        f"//! Source SHA-256 (nrlmsise-00.c and nrlmsise-00_data.c): {digest}",
        f"//! Generated: {date.today().isoformat()}",
        "//!",
        "//! The reference's `sam[100]` is not here: no part of NRLMSISE-00 reads it.",
        "",
        "#![allow(clippy::all)]",
        "",
    ]
    for name, konst, shape, what in TABLES:
        literals, labels = parse(source, name, shape)
        lines.append(f"/// {what}")
        lines.append(f"///")
        lines.append(f"/// `{name}` in the reference.")
        if len(shape) == 1:
            lines.append(f"pub(crate) static {konst}: [f64; {shape[0]}] = [")
            lines += [f"    {line}" for line in rows(literals)]
            lines.append("];")
        else:
            lines.append(f"pub(crate) static {konst}: [[f64; {shape[1]}]; {shape[0]}] = [")
            for row in range(shape[0]):
                chunk = literals[row * shape[1] : (row + 1) * shape[1]]
                label = f"  // {labels[row]}" if labels[row] else ""
                lines.append(f"    [{label}")
                lines += [f"        {line}" for line in rows(chunk)]
                lines.append("    ],")
            lines.append("];")
        lines.append("")
    return "\n".join(lines)


def main() -> None:
    files = {name: fetch(name) for name in ("nrlmsise-00.c", "nrlmsise-00_data.c")}
    digest = hashlib.sha256(
        (files["nrlmsise-00.c"] + files["nrlmsise-00_data.c"]).encode("ascii")
    ).hexdigest()
    print(f"source digest {digest}")
    for name in UNUSED:
        if f"double {name}" not in files["nrlmsise-00_data.c"]:
            raise SystemExit(f"{name} is gone from the reference; check what changed")
    OUT.write_text(emit(files["nrlmsise-00_data.c"], digest))
    subprocess.run(["rustfmt", "--edition", "2021", str(OUT)], check=True)
    count = len(re.findall(r"-?\d+\.\d+", OUT.read_text()))
    print(f"wrote {OUT.relative_to(ROOT)}: {count} coefficients")


if __name__ == "__main__":
    main()