data/aia_normalization.py

#!/home/ubuntu/miniconda3/envs/fox/bin/python
"""Compute AIA normalization parameters with bounded RAM.

Each finite pixel from every selected training file contributes to the global
per-wavelength statistics. Negative finite values are counted and clamped to
zero; NaN and infinite values are counted and excluded.

This script keeps only the statistics needed for normalization:

    Q90     -> asinh scale
    Q99.999 -> relaxed clipping threshold
    asinh_mean, asinh_std -> transformed normalization moments

Percentiles are approximated with a fine log-intensity histogram so that
hundreds of billions of pixels never need to be held in memory. Large batches
and parallel shards keep the CPUs busy, while atomic checkpoints make a long
run safe to resume after interruption.
"""

from __future__ import annotations

import argparse
import csv
import hashlib
import math
import os
import time
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

import numpy as np
from tqdm.auto import tqdm


# Channel order expected in each input array: (7, height, width).
WAVELENGTHS = (94, 131, 171, 193, 211, 304, 335)

# Normalization depends only on the 90th percentile and the relaxed clipping
# threshold at 99.999%.
PERCENTILES = (90.0, 99.999)


def available_cpu_count() -> int:
    """Return CPUs available to this process, respecting CPU affinity."""
    if hasattr(os, "sched_getaffinity"):
        return len(os.sched_getaffinity(0))
    return os.cpu_count() or 1


def parse_args() -> argparse.Namespace:
    cpu_count = available_cpu_count()
    parser = argparse.ArgumentParser(
        description=(
            "Stream AIA training NPY files and estimate the normalization "
            "parameters Q90, Q99.999, and asinh mean/std."
        )
    )
    parser.add_argument(
        "--train-dir", type=Path, default=Path("/data/AIA_processed/train")
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=Path("/data/aia_q90_asinh_q99999_norm_streamed.npz"),
    )
    parser.add_argument(
        "--checkpoint",
        type=Path,
        default=Path("/data/aia_streaming_stats_checkpoint.npz"),
    )
    parser.add_argument(
        "--summary-csv",
        type=Path,
        default=Path("/data/aia_q90_asinh_q99999_norm_streamed.csv"),
    )
    parser.add_argument(
        "--max-files",
        type=int,
        default=None,
        help="Randomly select this many files; default is every file.",
    )
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument(
        "--batch-size",
        type=int,
        default=4048,
        help=(
            "Images per batch. 2048 is a 14 GiB base buffer for 7x512x512 "
            "float32 images; parallel temporaries use substantially more RAM."
        ),
    )
    parser.add_argument("--load-workers", type=int, default=cpu_count)
    parser.add_argument("--reduce-workers", type=int, default=cpu_count)
    parser.add_argument(
        "--shards-per-channel",
        type=int,
        default=None,
        help="Default creates about twice as many reduction tasks as CPU slots.",
    )
    parser.add_argument("--histogram-bins", type=int, default=131_072)
    parser.add_argument(
        "--checkpoint-every",
        type=int,
        default=10,
        help="Checkpoint every N completed batches.",
    )
    parser.add_argument(
        "--no-resume",
        action="store_true",
        help="Ignore an existing checkpoint and start from the beginning.",
    )
    return parser.parse_args()


def validate_args(args: argparse.Namespace) -> None:
    positive_values = {
        "batch-size": args.batch_size,
        "load-workers": args.load_workers,
        "reduce-workers": args.reduce_workers,
        "histogram-bins": args.histogram_bins,
        "checkpoint-every": args.checkpoint_every,
    }
    if args.shards_per_channel is not None:
        positive_values["shards-per-channel"] = args.shards_per_channel
    if args.max_files is not None:
        positive_values["max-files"] = args.max_files
    invalid = [name for name, value in positive_values.items() if value <= 0]
    if invalid:
        raise ValueError(f"These arguments must be positive: {', '.join(invalid)}")


def atomic_savez(path: Path, **payload: object) -> None:
    """Write completely before replacing an existing checkpoint or result."""
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_name(f".{path.name}.{os.getpid()}.tmp.npz")
    try:
        np.savez(temporary, **payload)
        os.replace(temporary, path)
    finally:
        if temporary.exists():
            temporary.unlink()


def selection_digest(paths: list[Path]) -> str:
    """Identify the exact ordered file selection used by a checkpoint."""
    digest = hashlib.sha256()
    for path in paths:
        digest.update(path.name.encode("utf-8"))
        digest.update(b"\0")
    return digest.hexdigest()


def main() -> None:
    args = parse_args()
    validate_args(args)
    started = time.perf_counter()

    # 1. Discover all training files, optionally taking a reproducible sample.
    # With the default --max-files=None, no files or pixels are sampled.
    paths = sorted(args.train_dir.glob("*.npy"))
    if not paths:
        raise RuntimeError(f"No .npy files found in {args.train_dir}")
    if args.max_files is not None and args.max_files < len(paths):
        rng = np.random.default_rng(args.seed)
        selected = np.sort(
            rng.choice(len(paths), size=args.max_files, replace=False)
        )
        paths = [paths[index] for index in selected]

    # Inspect one memory-mapped file without loading all its pixels into RAM.
    probe = np.load(paths[0], mmap_mode="r", allow_pickle=False)
    if probe.ndim != 3 or probe.shape[0] != len(WAVELENGTHS):
        raise ValueError(
            f"Expected (7, H, W), got {probe.shape} in {paths[0]}"
        )
    image_shape = tuple(probe.shape)
    image_dtype = probe.dtype
    del probe
    if image_dtype != np.float32:
        print(
            f"Warning: input dtype is {image_dtype}; the batch buffer converts it "
            "to float32.",
            flush=True,
        )

    # More reduction tasks than workers helps prevent CPUs sitting idle when
    # individual shards finish at slightly different times.
    cpu_count = available_cpu_count()
    shards_per_channel = args.shards_per_channel or math.ceil(
        2 * args.reduce_workers / len(WAVELENGTHS)
    )
    file_bytes = paths[0].stat().st_size
    input_bytes = len(paths) * file_bytes
    buffer_bytes = args.batch_size * int(np.prod(image_shape)) * 4
    selected_digest = selection_digest(paths)

    print(f"Training directory: {args.train_dir}", flush=True)
    print(f"Files to stream: {len(paths):,}", flush=True)
    print(f"Input to read: {input_bytes / 2**40:.3f} TiB", flush=True)
    print(f"Batch buffer: {buffer_bytes / 2**30:.2f} GiB", flush=True)
    print(f"Available logical CPUs: {cpu_count}", flush=True)
    print(
        f"Parallelism: {args.load_workers} loaders, {args.reduce_workers} reducers, "
        f"up to {shards_per_channel * len(WAVELENGTHS)} reduction tasks/batch",
        flush=True,
    )

    # 2. Allocate the global accumulators needed for the histogram-based
    # normalization parameters.
    channel_count = len(WAVELENGTHS)

    # Uniform log1p bins provide roughly constant relative intensity precision
    # while covering every possible nonnegative float32 value.
    log_max = float(np.log1p(np.finfo(np.float32).max))
    log_width = log_max / args.histogram_bins
    histograms = np.zeros(
        (channel_count, args.histogram_bins), dtype=np.int64
    )
    counts = np.zeros(channel_count, dtype=np.int64)
    zero_counts = np.zeros(channel_count, dtype=np.int64)
    processed_files = 0

    def save_checkpoint() -> None:
        """Persist global accumulators at a completed-batch boundary."""
        atomic_savez(
            args.checkpoint,
            processed_files=np.asarray(processed_files),
            selected_file_count=np.asarray(len(paths)),
            selection_digest=np.asarray(selected_digest),
            first_file=np.asarray(paths[0].name),
            last_file=np.asarray(paths[-1].name),
            histogram_bins=np.asarray(args.histogram_bins),
            log_max=np.asarray(log_max),
            histograms=histograms,
            counts=counts,
            zero_counts=zero_counts,
        )

    # Resume only when both the exact file selection and histogram layout match.
    # Batch size and worker counts can safely differ from the previous process.
    if not args.no_resume and args.checkpoint.exists():
        with np.load(args.checkpoint, allow_pickle=False) as checkpoint:
            required = {
                "processed_files",
                "selected_file_count",
                "selection_digest",
                "histogram_bins",
                "log_max",
                "histograms",
                "counts",
                "zero_counts",
            }
            missing = required.difference(checkpoint.files)
            if missing:
                raise RuntimeError(
                    f"Checkpoint {args.checkpoint} is missing fields: "
                    f"{', '.join(sorted(missing))}. Use --no-resume or a new path."
                )
            compatible = (
                int(checkpoint["selected_file_count"]) == len(paths)
                and str(checkpoint["selection_digest"]) == selected_digest
                and int(checkpoint["histogram_bins"]) == args.histogram_bins
                and float(checkpoint["log_max"]) == log_max
            )
            if not compatible:
                raise RuntimeError(
                    f"Checkpoint {args.checkpoint} does not match this file "
                    "selection or histogram configuration. Use --no-resume or "
                    "choose a new checkpoint path."
                )
            processed_files = int(checkpoint["processed_files"])
            if not 0 <= processed_files <= len(paths):
                raise RuntimeError("Checkpoint has an invalid processed file count")
            histograms[:] = checkpoint["histograms"]
            counts[:] = checkpoint["counts"]
            zero_counts[:] = checkpoint["zero_counts"]
        print(f"Resuming after {processed_files:,} files", flush=True)

    # 3. Stream pixels in bounded batches. At the default size this persistent
    # buffer is ~14 GiB; masks and reducer temporaries raise peak RAM further.
    batch_buffer = np.empty(
        (args.batch_size,) + image_shape, dtype=np.float32
    )

    def load_item(item: tuple[int, Path]) -> None:
        """Load one complete (7, H, W) file into its batch-buffer slot."""
        slot, path = item
        array = np.load(path, mmap_mode="r", allow_pickle=False)
        if array.shape != image_shape:
            raise ValueError(
                f"Expected shape {image_shape}, got {array.shape} in {path}"
            )
        batch_buffer[slot] = array

    def reduce_shard(
        task: tuple[int, int, int],
    ) -> tuple[int, int, int, np.ndarray | None]:
        """Scan one channel slice and return the histogram contribution."""
        channel, shard_start, shard_stop = task
        values = batch_buffer[shard_start:shard_stop, channel].reshape(-1)
        finite = values[np.isfinite(values)]
        if finite.size == 0:
            return (channel, 0, 0, None)

        clean = np.maximum(finite, 0.0)
        zero_count = int(np.count_nonzero(clean == 0))
        shard_count = int(clean.size)

        log_values = np.log1p(clean)
        bin_indices = np.floor(log_values / log_width).astype(np.int64)
        np.clip(bin_indices, 0, args.histogram_bins - 1, out=bin_indices)
        histogram = np.bincount(
            bin_indices, minlength=args.histogram_bins
        )
        return (channel, shard_count, zero_count, histogram)

    def merge_shard(
        result: tuple[int, int, int, np.ndarray | None]
    ) -> None:
        """Merge one worker result into the per-wavelength global state."""
        channel, shard_count, zero_count, histogram = result
        if shard_count == 0:
            return

        counts[channel] += shard_count
        zero_counts[channel] += zero_count
        if histogram is None:
            raise RuntimeError("Missing histogram for a nonempty shard")
        histograms[channel] += histogram

    state_arrays = (histograms, counts, zero_counts)

    progress = tqdm(
        total=len(paths),
        initial=processed_files,
        desc="Streaming full AIA images",
        unit="file",
        dynamic_ncols=True,
    )
    batch_number = math.ceil(processed_files / args.batch_size)
    try:
        with (
            ThreadPoolExecutor(max_workers=args.load_workers) as load_executor,
            ThreadPoolExecutor(max_workers=args.reduce_workers) as reduce_executor,
        ):
            for batch_start in range(
                processed_files, len(paths), args.batch_size
            ):
                batch_paths = paths[
                    batch_start : batch_start + args.batch_size
                ]
                # Phase A: parallel disk reads into disjoint buffer slots.
                list(load_executor.map(load_item, enumerate(batch_paths)))

                # Phase B: divide every wavelength into enough independent
                # shards to occupy all reduction workers.
                shard_count = min(shards_per_channel, len(batch_paths))
                shard_bounds = np.linspace(
                    0, len(batch_paths), shard_count + 1, dtype=np.int64
                )
                reduction_tasks = [
                    (
                        channel,
                        int(shard_bounds[shard]),
                        int(shard_bounds[shard + 1]),
                    )
                    for channel in range(channel_count)
                    for shard in range(shard_count)
                    if shard_bounds[shard] < shard_bounds[shard + 1]
                ]

                # Restoring this small snapshot makes an interrupted partial batch
                # safe to checkpoint and process again on resume.
                stable_state = tuple(array.copy() for array in state_arrays)
                try:
                    for result in reduce_executor.map(
                        reduce_shard, reduction_tasks
                    ):
                        merge_shard(result)
                except BaseException:
                    for array, stable in zip(state_arrays, stable_state):
                        array[:] = stable
                    raise

                processed_files = batch_start + len(batch_paths)
                batch_number += 1
                progress.update(len(batch_paths))
                if batch_number % args.checkpoint_every == 0:
                    save_checkpoint()
    except BaseException:
        save_checkpoint()
        print(
            f"Saved a resumable checkpoint after {processed_files:,} complete "
            f"files: {args.checkpoint}",
            flush=True,
        )
        raise
    finally:
        progress.close()

    save_checkpoint()

    # 4. Convert histogram counts into the two percentile values used for
    # normalization: Q90 (the scale) and Q99.999 (the clipped tail cap).
    quantiles = np.empty(
        (channel_count, len(PERCENTILES)), dtype=np.float64
    )
    for channel in range(channel_count):
        histogram = histograms[channel]
        cumulative = np.cumsum(histogram)
        total = int(cumulative[-1])
        for percentile_index, percentile in enumerate(PERCENTILES):
            # Cumulative histogram entries are counts, so use a count target.
            target_count = percentile / 100.0 * total
            if target_count <= zero_counts[channel]:
                quantiles[channel, percentile_index] = 0.0
                continue
            bin_index = int(
                np.searchsorted(cumulative, target_count, side="left")
            )
            bin_index = min(bin_index, len(histogram) - 1)
            previous = 0 if bin_index == 0 else cumulative[bin_index - 1]
            fraction = (target_count - previous) / max(
                int(histogram[bin_index]), 1
            )
            log_value = (
                bin_index + np.clip(fraction, 0.0, 1.0)
            ) * log_width
            quantiles[channel, percentile_index] = float(
                np.expm1(log_value)
            )

    q90 = quantiles[:, PERCENTILES.index(90.0)]
    clip_q99999 = quantiles[:, PERCENTILES.index(99.999)]
    if np.any(q90 <= 0):
        bad_channels = np.asarray(WAVELENGTHS)[q90 <= 0].tolist()
        raise RuntimeError(f"Q90 is zero for wavelengths {bad_channels}")

    # 5. Estimate transformed mean/std using histogram-bin representatives.
    # This avoids rereading the entire ~1.75 TiB dataset. Unlike the exact raw
    # moments above, these transformed moments are fine approximations.
    asinh_mean = np.empty(channel_count, dtype=np.float64)
    asinh_std = np.empty(channel_count, dtype=np.float64)
    log_centers = (np.arange(args.histogram_bins) + 0.5) * log_width
    raw_representatives = np.expm1(log_centers)
    for channel in range(channel_count):
        weights = histograms[channel].astype(np.float64)
        zero_weight = float(zero_counts[channel])

        positive_weights = weights.copy()
        positive_weights[0] -= zero_weight

        transformed = np.arcsinh(
            np.minimum(raw_representatives, clip_q99999[channel])
            / q90[channel]
        )

        transformed_mean = np.dot(positive_weights, transformed) / counts[channel]
        positive_variance = np.dot(
            positive_weights,
            (transformed - transformed_mean) ** 2,
        )
        zero_variance = zero_weight * transformed_mean**2
        transformed_variance = (
            positive_variance + zero_variance
        ) / counts[channel]

        asinh_mean[channel] = transformed_mean
        asinh_std[channel] = np.sqrt(transformed_variance)

    # 6. Save both the machine-readable parameters and a readable CSV summary.
    atomic_savez(
        args.output,
        wavelengths=np.asarray(WAVELENGTHS, dtype=np.int32),
        percentiles=np.asarray(PERCENTILES),
        q90=q90,
        clip_q99999=clip_q99999,
        asinh_mean=asinh_mean,
        asinh_std=asinh_std,
    )

    args.summary_csv.parent.mkdir(parents=True, exist_ok=True)
    columns = [
        "wavelength_A",
        "q90",
        "q99.999",
        "asinh_mean",
        "asinh_std",
    ]
    with args.summary_csv.open("w", newline="") as csv_file:
        writer = csv.DictWriter(csv_file, fieldnames=columns)
        writer.writeheader()
        for channel, wavelength in enumerate(WAVELENGTHS):
            writer.writerow(
                {
                    "wavelength_A": wavelength,
                    "q90": q90[channel],
                    "q99.999": clip_q99999[channel],
                    "asinh_mean": asinh_mean[channel],
                    "asinh_std": asinh_std[channel],
                }
            )

    elapsed = time.perf_counter() - started
    print("\nNormalization parameters", flush=True)
    print("wave       Q90       Q99.999   asinh_mean  asinh_std", flush=True)
    for channel, wavelength in enumerate(WAVELENGTHS):
        print(
            f"{wavelength:>4}  {q90[channel]:>10.5g}  "
            f"{clip_q99999[channel]:>10.5g}  "
            f"{asinh_mean[channel]:>10.6f}  {asinh_std[channel]:>9.6f}",
            flush=True,
        )
    print(f"\nSaved parameters: {args.output}", flush=True)
    print(f"Saved summary: {args.summary_csv}", flush=True)
    print(f"Checkpoint: {args.checkpoint}", flush=True)
    print(
        f"Completed {len(paths):,} files in {elapsed / 3600:.2f} hours "
        f"({len(paths) / elapsed:.2f} files/s)",
        flush=True,
    )


if __name__ == "__main__":
    main()