| """ |
| End-to-end raw -> processed dataset builder for FOXES. |
| |
| Runs, in order: clean_aia -> convert_aia -> combine_sxr -> align -> split -> |
| (optional) sxr_normalization. Each step can be skipped via the `steps` section |
| of the config if you already ran it. Everything is driven by one YAML config |
| — see build_dataset_config.yaml for the fields. |
| |
| Usage: |
| python data/build_dataset.py --config data/build_dataset_config.yaml |
| """ |
| import argparse |
| import os |
|
|
| import numpy as np |
| import yaml |
|
|
| from clean_aia import clean_aia_data |
| from convert_aia import process_aia_to_npy |
| from combine_sxr import SXRDataProcessor |
| from align_aia_sxr import align_aia_sxr |
| from split_train_val_test import split_train_val_test |
| from sxr_normalization import compute_sxr_norm |
|
|
|
|
| def load_config(config_path): |
| with open(config_path, 'r') as f: |
| return yaml.safe_load(f) |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description='Build a FOXES-ready dataset from raw AIA/GOES data.') |
| parser.add_argument('--config', type=str, default='data/build_dataset_config.yaml', |
| help='Path to build_dataset config YAML file') |
| args = parser.parse_args() |
|
|
| config = load_config(args.config) |
|
|
| wavelengths = config['wavelengths'] |
| aia = config['aia'] |
| sxr = config['sxr'] |
| output = config['output'] |
| processing = config.get('processing', {}) |
| steps = config.get('steps', {}) |
| split = config.get('split', {}) |
| norm = config.get('sxr_normalization', {}) |
|
|
| print("=== FOXES Dataset Build ===") |
|
|
| if steps.get('clean_aia', True): |
| print("\n--- Step 1/5: Clean AIA FITS files ---") |
| clean_aia_data(aia['raw_dir'], aia['bad_files_dir'], wavelengths) |
| else: |
| print("\n--- Step 1/5: Clean AIA FITS files (skipped) ---") |
|
|
| if steps.get('convert_aia', True): |
| print("\n--- Step 2/5: Convert AIA FITS -> .npy ---") |
| process_aia_to_npy(aia['raw_dir'], aia['processed_dir'], wavelengths) |
| else: |
| print("\n--- Step 2/5: Convert AIA FITS -> .npy (skipped) ---") |
|
|
| if steps.get('combine_sxr', True): |
| print("\n--- Step 3/5: Combine raw GOES satellite files ---") |
| SXRDataProcessor(data_dir=sxr['raw_dir'], output_dir=sxr['combined_dir']).combine_goes_data() |
| else: |
| print("\n--- Step 3/5: Combine raw GOES satellite files (skipped) ---") |
|
|
| if steps.get('align', True): |
| print("\n--- Step 4/5: Align AIA timestamps with GOES SXR data ---") |
| align_aia_sxr( |
| goes_data_dir=sxr['combined_dir'], |
| aia_processed_dir=aia['processed_dir'], |
| output_sxr_dir=output['sxr_dir'], |
| aia_missing_dir=output['aia_missing_dir'], |
| max_processes=processing.get('max_processes'), |
| batch_size_multiplier=processing.get('batch_size_multiplier', 4), |
| min_batch_size=processing.get('min_batch_size', 1), |
| ) |
| else: |
| print("\n--- Step 4/5: Align AIA timestamps with GOES SXR data (skipped) ---") |
|
|
| |
| |
| if steps.get('split', False): |
| print("\n--- Step 5/5: Split AIA/SXR into train/val/test ---") |
| split_kwargs = dict(train_range=split.get('train_range'), val_range=split.get('val_range'), |
| test_range=split.get('test_range'), copy_files=split.get('copy_files', False)) |
| split_train_val_test(aia['processed_dir'], aia['processed_dir'], **split_kwargs) |
| split_train_val_test(output['sxr_dir'], output['sxr_dir'], **split_kwargs) |
| else: |
| print("\n--- Step 5/5: Split AIA/SXR into train/val/test (skipped) ---") |
|
|
| if norm.get('compute', False): |
| print("\n--- Optional: Compute SXR normalization stats (from the train split) ---") |
| sxr_norm = compute_sxr_norm(os.path.join(output['sxr_dir'], 'train')) |
| np.save(norm['output_path'], sxr_norm) |
| print(f"Saved SXR normalization to {norm['output_path']}") |
|
|
| print("\nDataset build complete!") |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
