#!/usr/bin/env python3 """ CM1 Precipitation Data Extraction Script Extracts rain, prate, xh, yh, and time from large CM1 files and saves to smaller NetCDF files. Output: /global/cfs/projectdirs/m4486/Haochen/Extracted_Data/Hourly_Precip/CM1/ """ import argparse import os import warnings from datetime import datetime, timedelta warnings.filterwarnings('ignore', category=FutureWarning) import numpy as np import xarray as xr # ============================================================================ # PATH CONFIGURATION # ============================================================================ CM1_BASE_PATH = "/global/cfs/projectdirs/m4486/Haochen/TRACER_Model/CM1_idealized" OUTPUT_BASE = "/global/cfs/projectdirs/m4486/Haochen/Extracted_Data/Hourly_Precip/CM1" # Scenario configurations # Folder name -> output experiment name SCENARIO_CONFIG = { 'control': 'control', 'x3': 'high_ccn', '03': 'low_ccn' } CASE_CONFIG = { 'june17': { 'case_folder': 'June17', 'base_datetime': datetime(2022, 6, 17, 6, 0), # 06 UTC June 17 'file1': 'cm1out_0600-1800UTC.nc', 'file2': 'cm1out_1800-3000UTC.nc', }, 'aug7': { 'case_folder': 'Aug7', 'base_datetime': datetime(2022, 8, 7, 6, 0), # 06 UTC August 7 'file1': 'cm1out_0600-2100UTC.nc', 'file2': 'cm1out_2100-3000UTC.nc', } } # ============================================================================ # EXTRACTION FUNCTION # ============================================================================ def extract_cm1_precipitation(case_name, scenario_folder, output_experiment): """ Extract precipitation variables from CM1 files. Parameters: ----------- case_name : str Case name (june17, aug7) scenario_folder : str CM1 folder name (control, x3, 03) output_experiment : str Output experiment name (control, high_ccn, low_ccn) """ config = CASE_CONFIG[case_name] # Build file paths case_path = os.path.join(CM1_BASE_PATH, config['case_folder'], scenario_folder) file1 = os.path.join(case_path, config['file1']) file2 = os.path.join(case_path, config['file2']) # Output file output_file = os.path.join(OUTPUT_BASE, f"CM1_{case_name}_{output_experiment}.nc") # Check if output exists if os.path.exists(output_file): print(f" Output file already exists: {output_file}") print(f" Skipping... (delete file to re-extract)") return True print(f"\n Input path: {case_path}") print(f" Output file: {output_file}") # Check input files if not os.path.exists(file1): print(f" Warning: File not found: {file1}") return False # Process files one at a time to save memory ds_list = [] time_offset = 0 # Track time offset for second file for i, fpath in enumerate([file1, file2]): if not os.path.exists(fpath): print(f" Warning: File not found: {fpath}") continue print(f" Reading file {i+1}: {os.path.basename(fpath)}...") # Open dataset and extract only needed variables ds = xr.open_dataset(fpath) # Extract only the variables we need vars_to_extract = [] if 'rain' in ds: vars_to_extract.append('rain') if 'prate' in ds: vars_to_extract.append('prate') if not vars_to_extract: print(f" Error: No precipitation variables found in {fpath}") ds.close() continue # Create subset with only needed variables ds_subset = ds[vars_to_extract].copy() # Keep coordinate variables if 'xh' in ds.coords: ds_subset['xh'] = ds['xh'] if 'yh' in ds.coords: ds_subset['yh'] = ds['yh'] # Get number of timesteps n_times = len(ds['time']) print(f" Extracted {len(vars_to_extract)} variables, {n_times} timesteps") # Store time offset for later if i == 0: time_offset = n_times ds_list.append(ds_subset) ds.close() if not ds_list: print(" Error: No data extracted") return False # Concatenate along time print(" Concatenating files...") combined_ds = xr.concat(ds_list, dim='time') # Create proper datetime array print(" Creating time coordinates...") base_dt = config['base_datetime'] total_timesteps = len(combined_ds['time']) # CM1 outputs at 10-minute intervals (600 seconds) # Create datetime array starting from base_datetime time_datetime = [] for i in range(total_timesteps): dt = base_dt + timedelta(seconds=i * 600) # 10 min = 600 seconds time_datetime.append(dt) print(f" Total timesteps: {total_timesteps}") print(f" Time range: {time_datetime[0]} to {time_datetime[-1]}") # Assign new time coordinate combined_ds = combined_ds.assign_coords(time=np.array(time_datetime, dtype='datetime64[ns]')) # Add attributes combined_ds.attrs = { 'description': 'Extracted CM1 precipitation data', 'source': f'CM1 idealized simulation - {scenario_folder}', 'case': case_name, 'experiment': output_experiment, 'time_interval': '10 minutes', 'created': datetime.now().strftime('%Y-%m-%d %H:%M:%S'), } if 'rain' in combined_ds: combined_ds['rain'].attrs = { 'description': 'Accumulated precipitation', 'units': 'mm', } if 'prate' in combined_ds: combined_ds['prate'].attrs = { 'description': 'Precipitation rate', 'units': 'mm/s', } # Save to NetCDF with compression print(f" Saving to {output_file}...") encoding = {} for var in combined_ds.data_vars: encoding[var] = {'zlib': True, 'complevel': 4} combined_ds.to_netcdf(output_file, encoding=encoding) # Print summary file_size = os.path.getsize(output_file) / 1e6 print(f" Done! File size: {file_size:.1f} MB") print(f" Time range: {time_datetime[0]} to {time_datetime[-1]}") print(f" Timesteps: {len(time_datetime)}") combined_ds.close() return True # ============================================================================ # MAIN PROCESSING # ============================================================================ def process_case(case_name, scenarios='all'): """Process a single case.""" config = CASE_CONFIG[case_name] print("\n" + "="*60) print(f"Processing: {case_name}") print("="*60) # Create output directory os.makedirs(OUTPUT_BASE, exist_ok=True) # Determine which scenarios to process if scenarios == 'all': scenarios_to_process = list(SCENARIO_CONFIG.keys()) else: scenarios_to_process = [scenarios] for scenario_folder in scenarios_to_process: output_experiment = SCENARIO_CONFIG[scenario_folder] print(f"\n--- {scenario_folder} -> {output_experiment} ---") extract_cm1_precipitation(case_name, scenario_folder, output_experiment) print(f"\nCase {case_name} complete!") # ============================================================================ # ENTRY POINT # ============================================================================ if __name__ == '__main__': parser = argparse.ArgumentParser( description='Extract precipitation data from CM1 files' ) parser.add_argument('--case', choices=['june17', 'aug7', 'both'], default='both', help='Case to process (default: both)') parser.add_argument('--scenario', choices=['control', 'x3', '03', 'all'], default='all', help='Scenario to process (default: all)') args = parser.parse_args() cases = ['june17', 'aug7'] if args.case == 'both' else [args.case] for case in cases: process_case(case, args.scenario) print("\n" + "="*60) print("All extraction complete!") print("="*60) print(f"\nExtracted files saved to: {OUTPUT_BASE}")