Refactor code structure and remove redundant code blocks for improved readability and maintainability. Added Sphinx API documentation
4672e76 | """ | |
| Compute log-space mean/std normalization stats over processed SXR .npy files. | |
| Only needed for training, not for running inference against a released | |
| checkpoint. Called by build_dataset.py, or run directly to regenerate only the | |
| normalization artifact after the dataset has already been built. | |
| """ | |
| import argparse | |
| import numpy as np | |
| from pathlib import Path | |
| import glob | |
| import os | |
| SXR_LOG_OFFSET = 1e-8 | |
| from tqdm import tqdm | |
| def compute_sxr_norm(sxr_dir): | |
| """ | |
| Compute the mean and standard deviation of log10-transformed Soft X-Ray (SXR) flux values. | |
| This function scans a given directory containing `.npy` SXR data files, | |
| loads each file, filters out invalid or non-finite values, applies a logarithmic | |
| transformation (`log10(SXR + 1e-8)`), and computes the mean and standard deviation | |
| for normalization purposes. These normalization statistics are typically used | |
| during model training and inference to ensure consistent SXR scaling. | |
| Parameters | |
| ---------- | |
| sxr_dir : str or Path | |
| Path to the directory containing `.npy` SXR flux files. | |
| Returns | |
| ------- | |
| tuple of (float, float) | |
| - mean : Mean of log10(SXR + 1e-8) flux values. | |
| - std : Standard deviation of log10(SXR + 1e-8) flux values. | |
| Raises | |
| ------ | |
| FileNotFoundError | |
| If the specified SXR directory does not exist. | |
| ValueError | |
| If no valid `.npy` files or no valid SXR values are found. | |
| Notes | |
| ----- | |
| - Files are expected to contain scalar SXR flux values in W/m². | |
| - Invalid (non-finite or nonpositive) values are automatically skipped. | |
| - The logarithmic transform helps stabilize the variance and normalize scale differences. | |
| """ | |
| sxr_dir = Path(sxr_dir).resolve() | |
| print(f"Checking SXR directory: {sxr_dir}") | |
| if not sxr_dir.is_dir(): | |
| raise FileNotFoundError(f"SXR directory does not exist or is not a directory: {sxr_dir}") | |
| # Use glob for case-insensitive matching | |
| sxr_files = sorted(glob.glob(os.path.join(sxr_dir, "*.npy"))) | |
| print(f"Found {len(sxr_files)} SXR files in {sxr_dir}") | |
| if len(sxr_files) == 0: | |
| print("No files matching '*.npy' found. Listing directory contents:") | |
| print(os.listdir(sxr_dir)[:10]) # Show first 10 files | |
| raise ValueError(f"No SXR files found in {sxr_dir}") | |
| sxr_values = [] | |
| for f in tqdm(sxr_files, desc="Computing SXR normalization", unit="file"): | |
| try: | |
| sxr = np.load(f) | |
| sxr = np.atleast_1d(sxr).flatten()[0] | |
| if not np.isfinite(sxr) or sxr <= 0: | |
| print(f"Skipping invalid SXR value in {f}: {sxr}") | |
| continue | |
| sxr_values.append(np.log10(sxr + SXR_LOG_OFFSET)) | |
| except Exception as e: | |
| print(f"Failed to load SXR file {f}: {e}") | |
| continue | |
| sxr_values = np.array(sxr_values) | |
| if len(sxr_values) == 0: | |
| raise ValueError(f"No valid SXR values found in {sxr_dir}. All files failed to load or contained invalid data.") | |
| mean = np.mean(sxr_values) | |
| std = np.std(sxr_values) | |
| print(f"Computed SXR normalization: mean={mean}, std={std}") | |
| return mean, std | |
| def main(): | |
| parser = argparse.ArgumentParser(description=__doc__) | |
| parser.add_argument( | |
| "--sxr-dir", type=Path, required=True, | |
| help="Training-split directory containing scalar SXR .npy files.", | |
| ) | |
| parser.add_argument( | |
| "--output", type=Path, required=True, | |
| help="Destination .npy normalization artifact.", | |
| ) | |
| args = parser.parse_args() | |
| mean, std = compute_sxr_norm(args.sxr_dir) | |
| args.output.parent.mkdir(parents=True, exist_ok=True) | |
| np.save(args.output, np.asarray([mean, std], dtype=np.float32)) | |
| print(f"Saved SXR normalization to {args.output.resolve()}") | |
| if __name__ == "__main__": | |
| main() | |