diff --git a/config_builder/README.md b/config_builder/README.md new file mode 100644 index 00000000..f36ab522 --- /dev/null +++ b/config_builder/README.md @@ -0,0 +1,7 @@ +# Datathon Project + +## How to use + +1. Run `python simulation_output_reviewer.py > available_data.txt` +2. Run `python zppy_config_generator.py`, after setting `ZPPY_CFG_LOCATION = "/home/ac.forsyth2/ez/zppy/config_builder/"` and a `UNIQUE_ID`. That will produce a starting point `zppy` `cfg` at `/home/ac.forsyth2/ez/zppy/config_builder/zppy_config_.cfg`. (Note: this actually reruns the `simulation_output_reviewer`. For a production use-case we'd only do single data read, and simply write the output as one step.) +3. Run `run_agent.py --base_config=zppy_config_.cfg --available_data=available_data.txt --query=`. diff --git a/config_builder/__init__.py b/config_builder/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/config_builder/knowledge_base.json b/config_builder/knowledge_base.json new file mode 100644 index 00000000..6d402281 --- /dev/null +++ b/config_builder/knowledge_base.json @@ -0,0 +1,67 @@ +{ + "Amazon dry bias": { + "description": "Amazon rainforest precipitation and soil moisture", + "variables": ["PRECT", "PRECL", "PRECC", "PRECSC", "PRECSL", "QSOIL", "QRUNOFF", + "SOILWATER_10CM", "SOILLIQ", "SOILICE", "H2OSNO"], + "regions": ["Amazon", "South America", "tropical"], + "diagnostics": ["lat_lon", "zonal_mean_xy", "annual_cycle_zonal_mean"], + "components": ["atm", "land"] + }, + "AMOC": { + "description": "Atlantic Meridional Overturning Circulation", + "variables": ["max_moc", "mocTimeSeries"], + "regions": ["Atlantic"], + "diagnostics": ["meridional_mean_2d"], + "components": ["ocn"], + "tasks": ["mpas_analysis", "global_time_series"] + }, + "precipitation": { + "description": "Precipitation diagnostics", + "variables": ["PRECT", "PRECL", "PRECC", "PRECSC", "PRECSL"], + "diagnostics": ["lat_lon", "zonal_mean_xy", "diurnal_cycle"], + "components": ["atm"] + }, + "temperature": { + "description": "Temperature diagnostics", + "variables": ["TREFHT", "TS", "TSA", "T"], + "diagnostics": ["lat_lon", "zonal_mean_xy", "zonal_mean_2d"], + "components": ["atm", "land"] + }, + "carbon cycle": { + "description": "Carbon cycle and ecosystem", + "variables": ["GPP", "NBP", "HR", "AR", "TOTVEGC", "TOTSOMC", "TOTLITC", + "SOIL1C", "SOIL2C", "SOIL3C", "SOIL4C", "NEP"], + "diagnostics": ["lat_lon", "annual_cycle_zonal_mean"], + "components": ["land"], + "tasks": ["ilamb"] + }, + "clouds": { + "description": "Cloud diagnostics", + "variables": ["CLDTOT", "CLDLOW", "CLDMED", "CLDHGH", "TGCLDLWP", "TGCLDIWP", + "CLOUD", "CLWMODIS", "CLIMODIS"], + "diagnostics": ["lat_lon", "zonal_mean_xy", "cosp_histogram"], + "components": ["atm"] + }, + "radiation": { + "description": "Radiative fluxes", + "variables": ["FSNT", "FLNT", "FSNS", "FLNS", "SWCF", "LWCF", "SOLIN", + "FSNTOA", "FSNTOAC"], + "diagnostics": ["lat_lon", "zonal_mean_xy", "annual_cycle_zonal_mean"], + "components": ["atm"] + }, + "aerosol": { + "description": "Aerosol optical depth and composition", + "variables": ["AODVIS", "AODDUST", "AODABS", "AODSO4", "AODBC", "AODPOM", "AODSOA"], + "diagnostics": ["lat_lon", "aerosol_aeronet"], + "components": ["atm"] + }, + "Labrador Sea cold bias": { + "description": "Labrador Sea temperature, convection, and circulation", + "variables": ["TS", "T", "TREFHT", "PSL", "U", "V", "SHFLX", "LHFLX", + "FLNS", "FSNS", "PRECL", "PRECC", "CLDTOT", "ICEFRAC"], + "regions": ["Labrador Sea", "North Atlantic", "Arctic"], + "diagnostics": ["lat_lon", "zonal_mean_xy", "polar", "meridional_mean_2d"], + "components": ["atm", "ocn"], + "tasks": ["mpas_analysis", "e3sm_diags"] + } +} diff --git a/config_builder/run_agent.py b/config_builder/run_agent.py new file mode 100644 index 00000000..6e4c2b68 --- /dev/null +++ b/config_builder/run_agent.py @@ -0,0 +1,364 @@ +#!/usr/bin/env python3 +""" +Agentic AI-powered zppy configuration modifier using ollama. +Allows scientists to request diagnostic configurations using natural language. +""" + +import argparse +import json +import re +import subprocess +import sys +from pathlib import Path +from typing import Any, Dict, List, Optional + + +class ZppyConfigAgent: + """Agent to modify zppy configurations based on natural language queries.""" + + def __init__( + self, + base_config_path: str, + available_data_path: str, + model: str, + knowledge_base_path: str, + ): + """ + Initialize the zppy config agent. + + Args: + base_config_path: Path to the base zppy configuration file + available_data_path: Path to simulation output reviewer results + model: Ollama model to use + knowledge_base_path: JSON file with topic->variable mappings + """ + self.model: str = model + self.base_config: str = self._load_file(base_config_path) + self.available_data: str = self._load_file(available_data_path) + + # Load or use default knowledge base + self.knowledge_base: Dict[str, Any] + if knowledge_base_path and Path(knowledge_base_path).exists(): + with open(knowledge_base_path) as f: + self.knowledge_base = json.load(f) + else: + self.knowledge_base = {} + + # Parse available variables from simulation output + self.available_vars: Dict[str, List[str]] = self._parse_available_variables() + + # Methods listed in order of first use #################################### + + def _load_file(self, path: str) -> str: + """Load file content as string.""" + with open(path) as f: + return f.read() + + def _parse_available_variables(self) -> Dict[str, List[str]]: + """Parse available variables from simulation output reviewer.""" + vars_by_component: Dict[str, List[str]] = {} + current_component: Optional[str] = None + + for line in self.available_data.split("\n"): + # Look for component headers like "component=atm" + comp_match = re.search(r"component=(\w+)", line) + if comp_match: + current_component = comp_match.group(1) + vars_by_component[current_component] = [] + + # Look for variable listings (numbered lines) + var_match = re.match(r"\s+\d+\.\s+(\w+)", line) + if var_match and current_component: + vars_by_component[current_component].append(var_match.group(1)) + + return vars_by_component + + def process_query(self, query: str, verbose: bool = False) -> str: + """ + Process a natural language query and return modified config. + + Args: + query: User's natural language request + verbose: Print intermediate steps + + Returns: + Modified zppy configuration as string + """ + if verbose: + print(f"\n{'=' * 70}") + print(f"Processing query: {query}") + print(f"{'=' * 70}\n") + + system_prompt = self._build_system_prompt() + user_prompt = self._build_user_prompt(query) + + if verbose: + print("Calling ollama...") + + response = self._call_ollama(user_prompt, system_prompt) + + # Clean up response if model added markdown fences + if response.startswith("```"): + lines = response.split("\n") + # Remove first and last line if they're fence markers + if lines[0].startswith("```") and lines[-1].startswith("```"): + response = "\n".join(lines[1:-1]) + + return response + + def _build_system_prompt(self) -> str: + """Build the system prompt with context about zppy configs.""" + + # Build a concise summary of available variables + vars_summary = [] + for comp, vars_list in self.available_vars.items(): + vars_summary.append( + f"{comp}: {len(vars_list)} variables including {', '.join(vars_list[:10])}" + ) + vars_text = "\n".join(vars_summary) + + return f"""You are an expert assistant for configuring E3SM post-processing diagnostics using zppy. + +ZPPY CONFIGURATION STRUCTURE: +- [default]: Global parameters (case name, paths, partition, walltime) +- [climo]: Climatology generation with subsections for different frequencies/grids +- [ts]: Time series generation with subsections for components (atm, land, ocean, etc.) +- [e3sm_to_cmip]: CMIP variable conversion +- [e3sm_diags]: Atmosphere diagnostics (sets: lat_lon, zonal_mean_xy, polar, cosp_histogram, etc.) +- [mpas_analysis]: Ocean/ice diagnostics +- [global_time_series]: Global time series plots +- [ilamb]: Land model benchmarking + +CONFIGURATION PARAMETERS: +- active = True/False: Enable/disable a task +- years = "start:end:interval": Time period (e.g., "1850:2050:20") +- vars = "VAR1,VAR2,...": Comma-separated variable list +- sets = "set1,set2,...": Diagnostic sets to run (for e3sm_diags) +- walltime = "HH:MM:SS": Job walltime +- frequency = "monthly"/"daily"/"diurnal_8xdaily": Data frequency + +AVAILABLE SIMULATION DATA: +{vars_text} + +SCIENTIFIC TOPIC MAPPINGS: +{json.dumps(self.knowledge_base, indent=2)} + +INSTRUCTIONS: +1. Analyze the user's request to identify: + - Relevant variables (check against available data) + - Time period (default to full range if not specified) + - Diagnostics/tasks to enable + - Components to include (atm, land, ocean, ice, rof) + +2. Modify the configuration to fulfill the request: + - Set active=True for relevant tasks, active=False for irrelevant ones + - Update variable lists to include only relevant variables + - Adjust time ranges if specified + - Update diagnostic sets for e3sm_diags + - Add comments explaining major changes + +3. Preserve configuration format: + - Keep proper indentation (2 spaces per level) + - Maintain section hierarchy + - Keep all required parameters + - Use proper quoting for strings + +4. Return ONLY the modified configuration file as plain text + - No markdown code fences + - No explanatory text before or after + - Just the valid zppy configuration + +EXAMPLES OF MODIFICATIONS: +- To focus on precipitation: Keep PRECT, PRECL, PRECC in vars; enable diurnal_cycle +- To focus on carbon: Enable ilamb task, select GPP, NBP, HR, AR variables +- For specific time period: Update years parameter (e.g., "1850:1950:10") +- To disable ocean: Set mpas_analysis active=False""" + + def _build_user_prompt(self, query: str) -> str: + """Build the user prompt for a specific query.""" + return f"""BASE CONFIGURATION: +{self.base_config} + +USER REQUEST: {query} + +Please modify the configuration to fulfill this request. Return only the modified zppy configuration file.""" + + def _call_ollama(self, prompt: str, system_prompt: str) -> str: + """Call ollama with the given prompt.""" + try: + result = subprocess.run( + ["ollama", "run", self.model], + input=f"{system_prompt}\n\n{prompt}", + capture_output=True, + text=True, + check=True, + timeout=120, # 2 minute timeout + ) + return result.stdout.strip() + except subprocess.CalledProcessError as e: + print(f"Error calling ollama: {e}", file=sys.stderr) + print(f"Stderr: {e.stderr}", file=sys.stderr) + raise + except subprocess.TimeoutExpired: + print("Ollama request timed out", file=sys.stderr) + raise + + def interactive_mode(self): + """Run in interactive mode for multiple queries.""" + print("=" * 70) + print("E3SM zppy Configuration Agent") + print("=" * 70) + print(f"Model: {self.model}") + print(f"Available components: {', '.join(self.available_vars.keys())}") + print("Type 'quit' to exit, 'help' for example queries\n") + + while True: + try: + query = input("Enter your request: ").strip() + + if query.lower() in ("quit", "exit", "q"): + print("\nExiting...") + break + + if query.lower() == "help": + self._print_help() + continue + + if not query: + continue + + modified_config = self.process_query(query, verbose=True) + + print("\n" + "=" * 70) + print("MODIFIED CONFIGURATION:") + print("=" * 70) + print(modified_config) + print("=" * 70) + + save = input("\nSave this configuration? (y/n): ").strip().lower() + if save == "y": + filename = input( + "Enter filename (default: modified_zppy.cfg): " + ).strip() + if not filename: + filename = "modified_zppy.cfg" + + with open(filename, "w") as f: + f.write(modified_config) + print(f"✓ Saved to {filename}") + + except KeyboardInterrupt: + print("\n\nExiting...") + break + except Exception as e: + print(f"✗ Error: {e}", file=sys.stderr) + + def _print_help(self): + """Print example queries.""" + print("\n" + "=" * 70) + print("EXAMPLE QUERIES:") + print("=" * 70) + print("• Show me diagnostics relevant to the Amazon dry bias") + print("• Show me diagnostics relevant to AMOC") + print("• Show me diagnostics on a hindcast 1850-1950") + print("• Focus on precipitation and clouds only") + print("• Disable ocean diagnostics") + print("• Enable carbon cycle analysis with ILAMB") + print("• Show aerosol diagnostics for 1980-2000") + print("• Generate only time series, no climatologies") + print("=" * 70 + "\n") + + +def main(): + """CLI entry point.""" + + parser = argparse.ArgumentParser( + description="Modify zppy configurations using natural language", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=""" +Examples: + # Interactive mode + %(prog)s base_config.cfg simulation_output.txt + + # Single query + %(prog)s base_config.cfg simulation_output.txt \\ + --query "show me diagnostics relevant to Amazon dry bias" \\ + --output amazon_config.cfg + + # Use larger model + %(prog)s base_config.cfg simulation_output.txt \\ + --model llama3.1:70b + """, + ) + + # Required parameters + parser.add_argument( + "base_config", + help="Path to starting point zppy configuration file (output of `zppy_config_generator.py`)", + ) + parser.add_argument( + "available_data", + help="Path to simulation output reviewer results (output of `simulation_output_reviewer.py`)", + ) + parser.add_argument( + "--query", help="Single query to process (otherwise runs interactively)" + ) + + # Optional parameters + parser.add_argument( + "--model", + default="llama3.1:8b", + help="Ollama model to use (default: llama3.1:8b)", + ) + parser.add_argument( + "--knowledge-base", + default="/home/ac.forsyth2/ez/zppy/config_builder/knowledge_base.json", + help="JSON file mapping topics to variables", + ) + parser.add_argument( + "--output", + default="tailored_zppy_config.cfg", + help="Output file for modified config (with --query)", + ) + parser.add_argument( + "--verbose", action="store_true", help="Print detailed processing steps" + ) + + args = parser.parse_args() + + # Verify ollama is available + try: + subprocess.run(["ollama", "list"], capture_output=True, check=True) + except (subprocess.CalledProcessError, FileNotFoundError): + print("Error: ollama not found. Please install ollama first:", file=sys.stderr) + print(" https://ollama.ai", file=sys.stderr) + print( + " Official instructions: https://ollama.com/download. For Linux, that shows `curl -fsSL https://ollama.com/install.sh | sh`. After installation, verify that Ollama is available by running `ollama --version`", + file=sys.stderr, + ) + sys.exit(1) + + agent = ZppyConfigAgent( + args.base_config, + args.available_data, + model=args.model, + knowledge_base_path=args.knowledge_base, + ) + + if args.query: + # Single query mode + modified_config = agent.process_query(args.query, verbose=args.verbose) + + if args.output: + with open(args.output, "w") as f: + f.write(modified_config) + print(f"✓ Modified configuration saved to {args.output}") + else: + print(modified_config) + else: + # Interactive mode + agent.interactive_mode() + + +if __name__ == "__main__": + main() diff --git a/config_builder/simulation_output_reviewer.py b/config_builder/simulation_output_reviewer.py new file mode 100644 index 00000000..28dd0d69 --- /dev/null +++ b/config_builder/simulation_output_reviewer.py @@ -0,0 +1,291 @@ +# TODO: Make mappings: +# These variables => run these tasks, sets +# These model components => run these tasks, sets + +import os +import re +import shlex +import subprocess +from typing import Dict, List, Optional, Set + +ARCHIVE_DIR: str = "/lcrc/group/e3sm2/ac.wlin/E3SMv3/v3.LR.historical_0051/" + + +class DirectoryMetadata(object): + def __init__(self, simulation_dir: str, component: str): + component_path: str = f"{simulation_dir}/archive/{component}/hist/" + component_files: List[str] = os.listdir(component_path) + metadata_dict: Dict[str, Set[str]] = get_available_filename_metadata( + component_files + ) + dates: Set[str] = metadata_dict["dates"] + + h_values: Set[str] = metadata_dict["h_values"] + h_value_to_first_file: Dict[str, str] = get_h_value_representative_files( + component_files, h_values + ) + variables_by_h_value: Dict[str, List[str]] = {} + for h_value in h_values: + representative_file: str = h_value_to_first_file[h_value] + variables_by_h_value[h_value] = get_netcdf_var_names( + f"{component_path}{representative_file}" + ) + + self.component: str = component + self.versions: Set[str] = metadata_dict["versions"] + self.resolutions: Set[str] = metadata_dict["resolutions"] + self.case_names: Set[str] = metadata_dict["case_names"] + self.components: Set[str] = metadata_dict["components"] + self.h_values: Set[str] = h_values + self.year_ranges: List[str] = collapse_year_range(dates) + self.file_extensions: Set[str] = metadata_dict["file_extensions"] + self.variables_by_h_value: Dict[str, List[str]] = variables_by_h_value + + def print_attributes(self, verbose=False): + print(f"\nPrinting directory metadata for component={self.component}") + print(f"versions={self.versions}") + print(f"resolutions={self.resolutions}") + print(f"case_names={self.case_names}") + print(f"components={self.components}") + print(f"h_values={self.h_values}") + for h_value in sorted(self.h_values): + vars: List[str] = self.variables_by_h_value[h_value] + print(f" {h_value} has {len(vars)} vars") + if verbose: + for i, var in enumerate(vars): + print(f" {i}. {var}") + print(f"year_ranges={self.year_ranges}") + print(f"file_extensions={self.file_extensions}") + + +def main(verbose: bool = False): + metadata_by_component: Dict[str, Optional[DirectoryMetadata]] = ( + get_metadata_by_component(ARCHIVE_DIR) + ) + for component in metadata_by_component: + m: Optional[DirectoryMetadata] = metadata_by_component[component] + if m: + m.print_attributes(verbose) + + +def get_metadata_by_component( + simulation_dir: str, +) -> Dict[str, Optional[DirectoryMetadata]]: + metadata_by_component: Dict[str, Optional[DirectoryMetadata]] = { + "atm": None, + "cpl": None, + "ice": None, + "lnd": None, + "ocn": None, + "rof": None, + } + for component in metadata_by_component.keys(): + metadata_by_component[component] = DirectoryMetadata(simulation_dir, component) + return metadata_by_component + + +def get_netcdf_var_names(filename: str) -> List[str]: + cmd = f"ncdump -h {shlex.quote(filename)}" + try: + result = subprocess.run( + cmd, + shell=True, + check=True, + capture_output=True, + text=True, + ) + except subprocess.CalledProcessError as e: + raise RuntimeError(f"ncdump failed: {e.stderr}") from e + + lines = result.stdout.splitlines() + + var_names = [] + in_variables = False + + # Regex: + # start of line (optional spaces) + # type (word) + # spaces + # name (word) + # optional spaces + # '(' ... ')' then ';' + # optionally followed by spaces and comments + decl_re = re.compile( + r"^\s*" # leading spaces + r"(\w+)\s+" # type, e.g., float, double, int + r"([A-Za-z0-9_]+)" # variable name + r"\s*\(" # opening parenthesis for dims + ) + + for raw_line in lines: + line = raw_line.rstrip() + + # Detect start of variables block + if line.strip().startswith("variables:"): + in_variables = True + continue + + if not in_variables: + continue + + # End of variables block + if line.strip().startswith("//") or line.strip() == "": + break + + # Skip attribute lines, which usually have ':' + # e.g.: " AODSO4:long_name = "Some name" ;" + stripped = line.strip() + if ":" in stripped or "=" in stripped: + # If it is exactly a declaration line with attributes folded later, + # it will still be matched by the regex above without ':' + # but attribute-only lines we want to skip. + # A conservative approach: only skip if it starts with ":" or has " = ". + if stripped.startswith(":") or " = " in stripped: + continue + + m = decl_re.match(line) + if m: + var_name = m.group(2) + var_names.append(var_name) + + return var_names + + +# Pure functions -- can be unit tested ######################################## +def get_h_value_representative_files( + component_files: List[str], available_h_values: Set[str] +) -> Dict[str, str]: + # Map each h-value to the first file containing it + h_value_to_first_file: Dict[str, str] = {} + for filename in component_files: + for h_value in available_h_values: + if (h_value in filename) and (h_value not in h_value_to_first_file): + h_value_to_first_file[h_value] = filename + break + return h_value_to_first_file + + +def get_available_filename_metadata(files: List[str]) -> Dict[str, Set[str]]: + metadata_dict: Dict[str, Set[str]] = { + "versions": set(), + "resolutions": set(), + "case_names": set(), + "components": set(), + "h_values": set(), + "dates": set(), + "file_extensions": set(), + } + for data_file in files: + parse_filename(data_file, metadata_dict) + return metadata_dict + + +def parse_filename(data_file: str, metadata_dict: Dict[str, Set[str]]): + parsed_string: List[str] = data_file.split(".") + num_entries: int = len(parsed_string) + if num_entries == 7: + # atm, cpl, lnd, rof + metadata_dict["versions"].add(parsed_string[0]) + metadata_dict["resolutions"].add(parsed_string[1]) + metadata_dict["case_names"].add(parsed_string[2]) + metadata_dict["components"].add(parsed_string[3]) + metadata_dict["h_values"].add(parsed_string[4]) + metadata_dict["dates"].add(parsed_string[5]) + metadata_dict["file_extensions"].add(parsed_string[6]) + elif num_entries == 9: + # ice, ocn + metadata_dict["versions"].add(parsed_string[0]) + metadata_dict["resolutions"].add(parsed_string[1]) + metadata_dict["case_names"].add(parsed_string[2]) + metadata_dict["components"].add(parsed_string[3]) + metadata_dict["h_values"].add( + f"{parsed_string[4]}.{parsed_string[5]}.{parsed_string[6]}" + ) + metadata_dict["dates"].add(parsed_string[7]) + metadata_dict["file_extensions"].add(parsed_string[8]) + elif num_entries == 10: + # ice, ocn + metadata_dict["versions"].add(parsed_string[0]) + metadata_dict["resolutions"].add(parsed_string[1]) + metadata_dict["case_names"].add(parsed_string[2]) + metadata_dict["components"].add(parsed_string[3]) + metadata_dict["h_values"].add( + f"{parsed_string[4]}.{parsed_string[5]}.{parsed_string[6]}" + ) + metadata_dict["dates"].add(f"{parsed_string[7]}-{parsed_string[8]}") + metadata_dict["file_extensions"].add(parsed_string[9]) + + +def is_valid_year_range_format(date_str: str) -> bool: + """ + Check if a string matches a format accepted by collapse_year_range. + + Accepted formats: + - YYYY-MM-DD-XXXXX (e.g., "1851-01-01-00000") + - YYYY-MM (e.g., "1985-03") + - YYYY-... (any string starting with a 4-digit year followed by hyphen) + + Args: + date_str: The string to validate + + Returns: + True if the string matches an accepted format, False otherwise + """ + if not date_str or not isinstance(date_str, str): + return False + + # Pattern: starts with 4 digits (year), followed by hyphen and anything + # Or just 4 digits (year only) + pattern = r"^\d{4}(-.*)?$" + + if not re.match(pattern, date_str): + return False + + # Extract and validate the year is reasonable (e.g., between 1000-9999) + year_str = date_str.split("-")[0] + try: + year = int(year_str) + return 1000 <= year <= 9999 + except ValueError: + return False + + +def collapse_year_range(year_ranges: Set[str]) -> List[str]: + if not year_ranges: + return [] + + # Extract years from the date strings + years = set() + for date_str in year_ranges: + year = int(date_str.split("-")[0]) + years.add(year) + + # Sort years + sorted_years = sorted(years) + + # Group consecutive years into ranges + result = [] + i = 0 + while i < len(sorted_years): + start = sorted_years[i] + end = start + + # Find consecutive years + while i + 1 < len(sorted_years) and sorted_years[i + 1] == sorted_years[i] + 1: + i += 1 + end = sorted_years[i] + + # Add range or single year + if start == end: + result.append(str(start)) + else: + result.append(f"{start}-{end}") + + i += 1 + + return result + + +# Run ######################################################################### +if __name__ == "__main__": + main(True) diff --git a/config_builder/test_simulation_output_reviewer.py b/config_builder/test_simulation_output_reviewer.py new file mode 100644 index 00000000..1e8e11d5 --- /dev/null +++ b/config_builder/test_simulation_output_reviewer.py @@ -0,0 +1,131 @@ +from typing import Dict, List, Set + +from .simulation_output_reviewer import ( + collapse_year_range, + get_available_filename_metadata, + get_h_value_representative_files, + is_valid_year_range_format, + parse_filename, +) + +# Run with: +# pytest test_*.py + + +def test_get_h_value_representative_files(): + files: List[str] = [ + "v3.LR.historical_0051.eam.h0.1851-01-01-00000.nc", + "v3.LR.historical_0051.eam.h0.1852-01-01-00000.nc", + "v3.LR.historical_0051.eam.h1.1851-01-01-00000.nc", + "v3.LR.historical_0051.eam.h1.1852-01-01-00000.nc", + "v3.LR.historical_0051.eam.h2.1851-01-01-00000.nc", + "v3.LR.historical_0051.eam.h2.1852-01-01-00000.nc", + ] + actual = get_h_value_representative_files(files, {"h0", "h1", "h2"}) + expected = { + "h0": "v3.LR.historical_0051.eam.h0.1851-01-01-00000.nc", + "h1": "v3.LR.historical_0051.eam.h1.1851-01-01-00000.nc", + "h2": "v3.LR.historical_0051.eam.h2.1851-01-01-00000.nc", + } + assert actual == expected + + +def test_get_available_filename_metadata(): + files: List[str] = [ + "v3.LR.historical_0051.cpl.hi.1851-01-01-00000.nc", + "v3.LR.historical_0051.cpl.hi.1852-01-01-00000.nc", + ] + actual: Dict[str, Set[str]] = get_available_filename_metadata(files) + expected: Dict[str, Set[str]] = { + "versions": {"v3"}, + "resolutions": {"LR"}, + "case_names": {"historical_0051"}, + "components": {"cpl"}, + "h_values": {"hi"}, + "dates": {"1851-01-01-00000", "1852-01-01-00000"}, + "file_extensions": {"nc"}, + } + assert actual == expected + + +def test_parse_filename(): + metadata_dict: Dict[str, Set[str]] = { + "versions": set(), + "resolutions": set(), + "case_names": set(), + "components": set(), + "h_values": set(), + "dates": set(), + "file_extensions": set(), + } + parse_filename("v3.LR.historical_0051.eam.h0.1850-01.nc", metadata_dict) + parse_filename("v3.LR.historical_0051.cpl.hi.1851-01-01-00000.nc", metadata_dict) + parse_filename( + "v3.LR.historical_0051.mpassi.hist.am.regionalStatistics.1850.01.nc", + metadata_dict, + ) + parse_filename( + "v3.LR.historical_0051.mpassi.hist.am.timeSeriesStatsMonthly.2050-12-01.nc", + metadata_dict, + ) + parse_filename("v3.LR.historical_0051.elm.h0.1850-01.nc", metadata_dict) + + assert metadata_dict["versions"] == {"v3"} + assert metadata_dict["resolutions"] == {"LR"} + assert metadata_dict["case_names"] == {"historical_0051"} + assert metadata_dict["components"] == {"eam", "cpl", "mpassi", "elm"} + assert metadata_dict["h_values"] == { + "h0", + "hi", + "hist.am.regionalStatistics", + "hist.am.timeSeriesStatsMonthly", + } + assert metadata_dict["dates"] == {"1850-01", "1851-01-01-00000", "2050-12-01"} + assert metadata_dict["file_extensions"] == {"nc"} + + +def test_string_split(): + string: str = "v3.LR.historical_0051.cpl.hi.1851-01-01-00000.nc" + actual: List[str] = string.split(".") + expected: List[str] = [ + "v3", + "LR", + "historical_0051", + "cpl", + "hi", + "1851-01-01-00000", + "nc", + ] + assert actual == expected + + +def test_is_valid_year_range(): + assert is_valid_year_range_format("1851-01-01-00000") == True + assert is_valid_year_range_format("1852-01-01-00000") == True + assert is_valid_year_range_format("2018-01-02-00000") == True + assert is_valid_year_range_format("1985-03") == True + assert is_valid_year_range_format("1907-06") == True + assert is_valid_year_range_format("2007") == True # Just year + + # Invalid formats + assert is_valid_year_range_format("") == False + assert is_valid_year_range_format("85-03") == False # Year too short + assert is_valid_year_range_format("abcd-01-01") == False + assert is_valid_year_range_format("2018/01/02") == False # Wrong separator + assert is_valid_year_range_format("0999-01") == False # Year too small + + +def test_collapse_year_range(): + year_ranges: Set[str] = { + "1851-01-01-00000", + "1852-01-01-00000", + "2007-01-01-00000", + } + actual: List[str] = collapse_year_range(year_ranges) + expected: List[str] = ["1851-1852", "2007"] + assert actual == expected + + year_ranges = {"2018-01-02-00000", "1985-03", "1907-06", "1907-05", "1908-01"} + actual = collapse_year_range(year_ranges) + expected = ["1907-1908", "1985", "2018"] + assert actual == expected diff --git a/config_builder/zppy_config_generator.py b/config_builder/zppy_config_generator.py new file mode 100644 index 00000000..689f155c --- /dev/null +++ b/config_builder/zppy_config_generator.py @@ -0,0 +1,893 @@ +#!/usr/bin/env python3 +""" +zppy Configuration Generator + +Generates zppy configuration files from simulation metadata. +""" + +from pathlib import Path +from typing import Dict, List, Optional, Set, Tuple + +from mache import MachineInfo +from simulation_output_reviewer import ( + ARCHIVE_DIR, + DirectoryMetadata, + get_metadata_by_component, +) + +UNIQUE_ID = "run16" +ZPPY_CFG_LOCATION = "/home/ac.forsyth2/ez/datathon/" +ZPPY_CFG_SUFFIX_OUTPUT = "datathon_placeholder/" +ZPPY_CFG_SUFFIX_WWW = "datathon_placeholder/" + + +class ZppyConfigGenerator: + def __init__( + self, + metadata_dict: Dict[str, Optional[DirectoryMetadata]], + use_only_standard_vars: bool = False, + ): + self.metadata_dict = metadata_dict + self.use_only_standard_vars: bool = use_only_standard_vars + + machine_specifics: Dict[str, str] = get_machine_specifics() + self.input_path: str = machine_specifics["input_path"] + self.output_path: str = machine_specifics["output_path"] + self.www_path: str = machine_specifics["output_path"] + self.diagnostics_base_path: str = machine_specifics["diagnostics_base_path"] + self.partition: str = machine_specifics["partition"] + self.qos: str = machine_specifics["qos"] + + self.case_name = get_case_name(metadata_dict) + + # Trackers of available climo & time-series data + self.dependency_tracker: Dict[str, Optional[Dependency]] = { + "climatology_atm_monthly": None, + "climatology_atm_monthly_diurnal": None, + "climatology_lnd_monthly": None, + "ts_atm_monthly": None, + "ts_atm_daily": None, + "ts_land_monthly": None, + "ts_rof_monthly": None, + "ts_atm_monthly_glb": None, + "ts_lnd_monthly_glb": None, + "e3sm_to_cmip_atm_monthly": None, + "e3sm_to_cmip_lnd_monthly": None, + "mpas_analysis": None, + } + + def generate(self) -> str: + sections: List[str] = [] + sections.append(self._generate_default_section()) + + climo_section = self._generate_climo_section() + sections.append(climo_section) + + ts_section = self._generate_ts_section() + sections.append(ts_section) + + if any(self.metadata_dict[k] for k in ("atm", "lnd")): + sections.append(self._generate_e3sm_to_cmip_section()) + + if self.metadata_dict["atm"]: + sections.append(self._generate_e3sm_diags_section()) + + if any(self.metadata_dict[k] for k in ("ocn", "ice")): + sections.append(self._generate_mpas_analysis_section()) + + if any(self.metadata_dict[k] for k in ("atm", "lnd", "ocn")): + sections.append(self._generate_global_time_series_section()) + + if any(self.metadata_dict[k] for k in ("atm", "lnd")): + sections.append(self._generate_ilamb_section()) + + # Skipping these tasks: + # tc_analysis, because of an issue with using the v3 dataset + # pcmdi_diags, because it's a beta version + + # Skipping model-vs-model for this initial prototype + + return "\n\n".join(sections) + + def _generate_default_section(self) -> str: + lines: List[str] = [ + "[default]", + f'case = "{self.case_name}"', + 'constraint = ""', + 'dry_run = "False"', + 'environment_commands = "" # Use Unified environment', + "fail_on_dependency_skip = True", + "infer_path_parameters = False", + "infer_section_parameters = False", + f'input = "{self.input_path}"', + "input_subdir = archive/atm/hist", + 'mapping_file = "map_ne30pg2_to_cmip6_180x360_aave.20200201.nc"', + f'output = "{self.output_path}"', + f'partition = "{self.partition}"', + f'qos = "{self.qos}"', + f'www = "{self.www_path}"', + ] + + return "\n".join(lines) + + def _generate_climo_section(self) -> str: + atm_h0_var_str: str = self._get_var_str("atm", "h0", set()) + atm_h3_var_str: str = self._get_var_str("atm", "h3", set(["PRECT"])) + lnd_h0_var_str: str = self._get_var_str("lnd", "h0", set()) + years_str: str + year_increment: int + years_str, _, _, year_increment = self._get_years_str("atm") + + lines: List[str] = [ + "[climo]", + "active = True", + 'walltime = "00:30:00"', + f'years = "{years_str}"', + "", + ] + + subtask_label: str + if atm_h0_var_str: + subtask_label = "atm_monthly_180x360_aave" + lines.extend( + [ + f" [[ {subtask_label} ]]", + ' frequency = "monthly"', + ' input_files = "eam.h0"', + ' input_subdir = "archive/atm/hist"', + f' vars = "{atm_h0_var_str}"', + "", + ] + ) + self.dependency_tracker["climatology_atm_monthly"] = Dependency( + subtask_label, atm_h0_var_str, years_str, year_increment + ) + + if atm_h3_var_str: + subtask_label = "atm_monthly_diurnal_8xdaily_180x360_aave" + lines.extend( + [ + f" [[ {subtask_label} ]]", + ' frequency = "diurnal_8xdaily"', + ' input_files = "eam.h3"', + ' input_subdir = "archive/atm/hist"', + f' vars = "{atm_h3_var_str}"', + "", + ] + ) + self.dependency_tracker["climatology_atm_monthly_diurnal"] = Dependency( + subtask_label, atm_h3_var_str, years_str, year_increment + ) + + if lnd_h0_var_str: + subtask_label = "land_monthly_climo" + lines.extend( + [ + f" [[ {subtask_label} ]]", + ' frequency = "monthly"', + ' input_files = "elm.h0"', + ' input_subdir = "archive/lnd/hist"', + ' mapping_file = "map_r05_to_cmip6_180x360_aave.20231110.nc"', + f' vars = "{lnd_h0_var_str}"', + "", + ] + ) + self.dependency_tracker["climatology_lnd_monthly"] = Dependency( + subtask_label, lnd_h0_var_str, years_str, year_increment + ) + + return "\n".join(lines) + + def _generate_ts_section(self) -> str: + atm_h0_standard: Set[str] = set( + "FSNTOA,FLUT,FSNT,FLNT,FSNS,FLNS,SHFLX,QFLX,TAUX,TAUY,PRECC,PRECL,PRECSC,PRECSL,TS,TREFHT,CLDTOT,CLDHGH,CLDMED,CLDLOW,U,PSL".split( + "," + ) + ) + lnd_h0_standard: Set[str] = set( + "FSH,RH2M,LAISHA,LAISUN,QINTR,QOVER,QRUNOFF,QSOIL,QVEGE,QVEGT,SOILICE,SOILLIQ,SOILWATER_10CM,TSA,TSOI,H2OSNO,TOTLITC,CWDC,SOIL1C,SOIL2C,SOIL3C,SOIL4C,WOOD_HARVESTC,TOTVEGC,NBP,GPP,AR,HR".split( + "," + ) + ) + + atm_h0_var_str: str = self._get_var_str("atm", "h0", atm_h0_standard) + atm_h1_var_str: str = self._get_var_str("atm", "h1", set(["PRECT"])) + lnd_h0_var_str: str = self._get_var_str("lnd", "h0", lnd_h0_standard) + lnd_h0_extra_vars_str: str = self._get_var_str("lnd", "h0", set(["landfrac"])) + rof_h0_var_str: str = self._get_var_str( + "rof", "h0", set(["RIVER_DISCHARGE_OVER_LAND_LIQ"]) + ) + rof_h0_extra_vars_str: str = self._get_var_str("rof", "h0", set(["areatotal2"])) + lnd_all_found_vars_str: str = self._get_var_str("lnd", "h0", set()) + years_str: str + year_increment: int + years_str, _, _, year_increment = self._get_years_str("atm") + + lines: List[str] = [ + "[ts]", + "active = True", + 'walltime = "00:30:00"', + f'years = "{years_str}"', + "", + ] + + subtask_label: str + if atm_h0_var_str: + subtask_label = "atm_monthly_180x360_aave" + lines.extend( + [ + f" [[ {subtask_label} ]]", + ' frequency = "monthly"', + ' input_files = "eam.h0"', + ' input_subdir = "archive/atm/hist"', + f' vars = "{atm_h0_var_str}"', + "", + ] + ) + self.dependency_tracker["ts_atm_monthly"] = Dependency( + subtask_label, atm_h0_var_str, years_str, year_increment + ) + + if atm_h1_var_str: + subtask_label = "atm_daily_180x360_aave" + lines.extend( + [ + f" [[ {subtask_label} ]]", + ' frequency = "daily"', + ' input_files = "eam.h1"', + ' input_subdir = "archive/atm/hist"', + f' vars = "{atm_h1_var_str}"', + "", + ] + ) + self.dependency_tracker["ts_atm_daily"] = Dependency( + subtask_label, atm_h1_var_str, years_str, year_increment + ) + + if lnd_h0_var_str: + if "landfrac" in lnd_h0_extra_vars_str: + subtask_label = "land_monthly" + lines.extend( + [ + f" [[ {subtask_label} ]]", + f' extra_vars = "{lnd_h0_extra_vars_str}"', + ' frequency = "monthly"', + ' input_files = "elm.h0"', + ' input_subdir = "archive/lnd/hist"', + f' vars = "{lnd_h0_var_str}"', + "", + ] + ) + self.dependency_tracker["ts_land_monthly"] = Dependency( + subtask_label, lnd_h0_var_str, years_str, year_increment + ) + else: + print( + f"Warning: land_monthly has valid variables, but extra_vars 'landfrac' is missing. Found variables={lnd_h0_var_str}" + ) + + if rof_h0_var_str: + if "areatotal2" in rof_h0_extra_vars_str: + subtask_label = "rof_monthly" + lines.extend( + [ + f" [[ {subtask_label} ]]", + f' extra_vars = "{rof_h0_extra_vars_str}"', + ' frequency = "monthly"', + ' input_files = "mosart.h0"', + ' input_subdir = "archive/rof/hist"', + f' vars = "{rof_h0_var_str}"', + "", + ] + ) + self.dependency_tracker["ts_rof_monthly"] = Dependency( + subtask_label, rof_h0_var_str, years_str, year_increment + ) + else: + print( + f"Warning: rof_monthly has valid variables, but extra_vars 'areatotal2' is missing. Found variables={rof_h0_var_str}" + ) + + # Use year increments of 5 for global time series + year_increment = 5 + years_str, _, _, year_increment = self._get_years_str("atm", year_increment) + + if atm_h0_var_str: + subtask_label = "atm_monthly_glb" + lines.extend( + [ + f" [[ {subtask_label} ]]", + ' frequency = "monthly"', + ' input_files = "eam.h0"', + ' input_subdir = "archive/atm/hist"', + ' mapping_file = "glb"', + f' vars = "{atm_h0_var_str}"', + "", + ] + ) + self.dependency_tracker["ts_atm_monthly_glb"] = Dependency( + subtask_label, atm_h0_var_str, years_str, year_increment + ) + + if lnd_all_found_vars_str: + subtask_label = "lnd_monthly_glb" + lines.extend( + [ + f" [[ {subtask_label} ]]", + ' frequency = "monthly"', + ' input_files = "elm.h0"', + ' input_subdir = "archive/lnd/hist"', + ' mapping_file = "glb"', + f' vars = "{lnd_all_found_vars_str}"', + "", + ] + ) + self.dependency_tracker["ts_lnd_monthly_glb"] = Dependency( + subtask_label, lnd_all_found_vars_str, years_str, year_increment + ) + + return "\n".join(lines) + + def _generate_e3sm_to_cmip_section(self) -> str: + standard_cmip_vars: Set[str] = set( + "ua, va, ta, wa, zg, hur, tas, ts, psl, ps, sfcWind, huss, pr, prc, prsn, evspsbl, tauu, tauv, hfls, clt, rlus, rsds, rsus, hfss, clivi, clwvi, rlut, rsdt, rsuscs, rsut, rtmt, abs550aer, od550aer, rsdscs, tasmax, tasmin".split( + ", " + ) + ) + standard_vars: Set[str] = set( + "ICEFRAC,LANDFRAC,OCNFRAC,PSL,FSNTC,FSNTOAC,SWCF,LWCF,FLUT,FSNT,FSNTOA,FLNT,FLNTC,FSNS,FLNS,FSNS,SHFLX,QFLX,LHFLX,TAUX,TAUY,PRECC,PRECL,PRECSC,PRECSL,TS,TREFHT,U10,QREFHT,TMQ,CLDTOT,CLDHGH,CLDMED,CLDLOW,FLDS,FSDS,TGCLDIWP,TGCLDCWP,TGCLDLWP,FLNSC,FLUTC,FSDSC,SOLIN,FSNSC,AODABS,AODVIS,AODDUST,AREL,TREFMNAV,TREFMXAV,PS,PHIS,U,V,T,Z3".split( + "," + ) + ) + + cmip_vars_str: str = self._get_var_str("atm", "h0", standard_cmip_vars) + vars_str: str = self._get_var_str("atm", "h0", standard_vars) + + lines: List[str] = [ + "[e3sm_to_cmip]", + "active = True", + 'frequency = "monthly"', + 'walltime = "00:30:00"', + "", + ] + + year_increment: int = 2 + subtask_label: str + + dependency_ts_atm_monthly: Optional[Dependency] = self.dependency_tracker[ + "ts_atm_monthly" + ] + if dependency_ts_atm_monthly: + subtask_label = "atm_monthly_180x360_aave" + lines.extend( + [ + f" [[{subtask_label}]]", + ' cmip_plevdata = "" # TODO: Set this', + f' cmip_vars = "{cmip_vars_str}"', + ' input_files = "eam.h0"', + f' ts_num_years = "{dependency_ts_atm_monthly.year_increment}"', + f' ts_subsection = "{dependency_ts_atm_monthly.subtask_label}"', + f' vars = "{vars_str}"', + f' years = "{dependency_ts_atm_monthly.years_str}"', + "", + ] + ) + self.dependency_tracker["e3sm_to_cmip_atm_monthly"] = Dependency( + subtask_label, + vars_str, + dependency_ts_atm_monthly.years_str, + year_increment, + ) + + dependency_ts_land_monthly: Optional[Dependency] = self.dependency_tracker[ + "ts_land_monthly" + ] + if dependency_ts_land_monthly: + subtask_label = "land_monthly" + lines.extend( + [ + f" [[{subtask_label}]]", + ' input_files = "elm.h0"', + f' ts_num_years = "{dependency_ts_land_monthly.year_increment}"', + f' ts_subsection = "{dependency_ts_land_monthly.subtask_label}"', + f' years = "{dependency_ts_land_monthly.years_str}"', + "", + ] + ) + self.dependency_tracker["e3sm_to_cmip_land_monthly"] = Dependency( + subtask_label, + vars_str, + dependency_ts_land_monthly.years_str, + year_increment, + ) + + return "\n".join(lines) + + def _generate_e3sm_diags_section(self) -> str: + sets_str: str = self._get_e3sm_diags_sets() + years_str: str + start_year: int + end_year: int + year_increment: int + years_str, start_year, end_year, year_increment = self._get_years_str("atm") + + lines: List[str] = [ + "[e3sm_diags]", + "active = True", + 'grid = "180x360_aave"', + "multiprocessing = True", + "num_workers = 8", + f"ref_final_yr = {end_year}", + f"ref_start_yr = {start_year}", + f'ref_years = "{start_year}-{end_year}",', + f'sets = "{sets_str}",', + f'short_name= "{self.case_name}"', + ] + + dependency_climo_atm_monthly: Optional[Dependency] = self.dependency_tracker[ + "climatology_atm_monthly" + ] + if dependency_climo_atm_monthly: + climo_subsection: str = dependency_climo_atm_monthly.subtask_label + lines.extend( + [ + f'climo_subsection = "{climo_subsection}"', + f'reference_data_path = "{self.diagnostics_base_path}observations/Atm/climatology/"', + ] + ) + + dependency_climo_atm_monthly_diurnal: Optional[Dependency] = ( + self.dependency_tracker["climatology_atm_monthly_diurnal"] + ) + if dependency_climo_atm_monthly_diurnal: + climo_diurnal_frequency: str = "diurnal_8xdaily" + climo_diurnal_subsection: str = ( + dependency_climo_atm_monthly_diurnal.subtask_label + ) + lines.extend( + [ + f'climo_diurnal_frequency = "{climo_diurnal_frequency}"', + f'climo_diurnal_subsection = "{climo_diurnal_subsection}"', + 'dc_obs_climo = "/lcrc/group/e3sm/public_html/e3sm_diags_test_data/unit_test_complete_run/obs/climatology"', + ] + ) + + dependency_ts_atm_monthly: Optional[Dependency] = self.dependency_tracker[ + "ts_atm_monthly" + ] + if dependency_ts_atm_monthly: + ts_subsection: str = dependency_ts_atm_monthly.subtask_label + lines.extend( + [ + f'ts_subsection = "{ts_subsection}"', + ] + ) + + dependency_ts_atm_daily: Optional[Dependency] = self.dependency_tracker[ + "ts_atm_daily" + ] + if dependency_ts_atm_daily: + ts_daily_subsection: str = dependency_ts_atm_daily.subtask_label + lines.extend( + [ + f'ts_daily_subsection = "{ts_daily_subsection}"', + ] + ) + + dependency_ts_rof_monthly: Optional[Dependency] = self.dependency_tracker[ + "ts_rof_monthly" + ] + if dependency_ts_rof_monthly: + lines.extend( + [ + f'streamflow_obs_ts = "{self.diagnostics_base_path}observations/Atm/time-series/"', + ] + ) + + if ( + dependency_ts_atm_monthly + or dependency_ts_atm_daily + or dependency_ts_rof_monthly + ): + lines.extend( + [ + f"ts_num_years = {year_increment}", + f'obs_ts = "{self.diagnostics_base_path}observations/Atm/time-series/"', + ] + ) + + # Add subtask + lines.extend( + [ + "", + " [[atm_monthly_180x360_aave]]", + f' years = "{years_str}"', + ] + ) + + return "\n".join(lines) + + def _get_e3sm_diags_sets(self) -> str: + # Eventually: Add "tc_analysis" back in after empty dat is resolved. + # Eventually: Add "aerosol_budget" back in once that's working for v3. + available_sets: List[str] = [ + "lat_lon", + "zonal_mean_xy", + "zonal_mean_2d", + "polar", + "cosp_histogram", + "meridional_mean_2d", + "annual_cycle_zonal_mean", + "zonal_mean_2d_stratosphere", + "enso_diags", + "qbo", + "diurnal_cycle", + "streamflow", + "tropical_subseasonal", + "aerosol_aeronet", + ] + + # Removing elements from the complete list allows us to retain the standard order of sets. + # Using `zppy/e3sm_diags.py` as a guide. + if not self.dependency_tracker["climatology_atm_monthly_diurnal"]: + available_sets.remove("diurnal_cycle") + if not self.dependency_tracker["ts_atm_monthly"]: + available_sets.remove("enso_diags") + available_sets.remove("qbo") + if not self.dependency_tracker["ts_atm_daily"]: + available_sets.remove("tropical_subseasonal") + if not self.dependency_tracker["ts_rof_monthly"]: + available_sets.remove("streamflow") + + return ",".join(available_sets) + + def _generate_mpas_analysis_section(self) -> str: + years_str: str + start_year: int + end_year: int + year_increment: int + years_str, start_year, end_year, year_increment = self._get_years_str("ocn") + if not (start_year and end_year): + # Use years from ice data as fall-back option + _, start_year, end_year, year_increment = self._get_years_str("ice") + + climo_years: str = f"{start_year}-{end_year}" + ts_years: str = f"{start_year}-{end_year}" + + self.dependency_tracker["mpas_analysis"] = Dependency( + "mpas_analysis", "", years_str, year_increment, climo_years, ts_years + ) + + return "\n".join( + [ + "[mpas_analysis]", + "active = True", + f"anomalyRefYear = {start_year}", + f'climo_years = "{climo_years}",', + f'enso_years = "{start_year}-{end_year}",', + 'mesh = "IcoswISC30E3r5"', + "parallelTaskCount = 6", + "shortTermArchive = True", + f'ts_years = "{ts_years}",', + ] + ) + + def _generate_global_time_series_section(self) -> str: + start_year: int + end_year: int + year_increment: int = 5 + _, start_year, end_year, year_increment = self._get_years_str( + "atm", year_increment + ) + if not (start_year and end_year): + # Use years from ocn data as fall-back option + _, start_year, end_year, year_increment = self._get_years_str( + "ocn", year_increment + ) + if not (start_year and end_year): + # Use years from lnd data as fall-back option + _, start_year, end_year, year_increment = self._get_years_str( + "lnd", year_increment + ) + + supported_plots: Set[str] = set() + + dependency_mpas_analysis: Optional[Dependency] = self.dependency_tracker[ + "mpas_analysis" + ] + climo_years: str = "" + ts_years: str = "" + moc_file_str: str = "" + if dependency_mpas_analysis: + climo_years = dependency_mpas_analysis.climo_years + ts_years = dependency_mpas_analysis.ts_years + if self.metadata_dict["ocn"]: + # We need ocn data specifically; ice data only won't work. + # Handle plots_original: + # As specified in `zppy_interfaces/global_time_series/utils.py` + supported_plots.add("change_ohc") + supported_plots.add("max_moc") + supported_plots.add("change_sea_level") + moc_file_str = f"mocTimeSeries_{start_year}-{end_year}.nc" + + dependency_ts_atm_monthly_glb: Optional[Dependency] = self.dependency_tracker[ + "ts_atm_monthly_glb" + ] + atm_h0_standard: Set[str] = set() + plots_atm: str = "" + if dependency_ts_atm_monthly_glb: + atm_h0_standard = set("TREFHT".split(",")) + plots_atm = self._get_var_str("atm", "h0", atm_h0_standard) + # Handle plots_original: + available_atm_vars: str = self._get_var_str("atm", "h0", set()) + # Set list of variables based on `get_vars_original()` in `zppy_interfaces/global_time_series/coupled_global/utils.py` + if "restom" in available_atm_vars: + supported_plots.add("net_toa_flux_restom") + if ("restom" in available_atm_vars) and ("ressurf" in available_atm_vars): + supported_plots.add("net_atm_energy_imbalance") + if "trefht" in available_atm_vars: + supported_plots.add("global_surface_air_temperature") + if ("fsntoa" in available_atm_vars) and ("flut" in available_atm_vars): + supported_plots.add("toa_radiation") + if ( + ("precc" in available_atm_vars) + and ("precl" in available_atm_vars) + and ("qflx" in available_atm_vars) + ): + supported_plots.add("net_atm_water_imbalance") + + plots_original: str = ",".join(supported_plots) + + dependency_ts_lnd_monthly_glb: Optional[Dependency] = self.dependency_tracker[ + "ts_lnd_monthly_glb" + ] + lnd_h0_standard: Set[str] = set() + plots_lnd: str = "" + if dependency_ts_lnd_monthly_glb: + lnd_h0_standard = set( + "FSH,RH2M,LAISHA,LAISUN,QINTR,QOVER,QRUNOFF,QSOIL,QVEGE,QVEGT,SOILWATER_10CM,TSA,H2OSNO,TOTLITC,CWDC,SOIL1C,SOIL2C,SOIL3C,SOIL4C,WOOD_HARVESTC,TOTVEGC,NBP,GPP,AR,HR".split( + "," + ) + ) + plots_lnd = self._get_var_str("lnd", "h0", lnd_h0_standard) + + return "\n".join( + [ + "[global_time_series]", + "active = True", + f'climo_years = "{climo_years}",', + f'experiment_name = "{self.case_name}"', + f'figstr = "{self.case_name}"', + f"ts_num_years = {year_increment}", + f'ts_years = "{ts_years}",', + f'years = "{start_year}-{end_year}",', + "", + " [[viewer]]", + " make_viewer = True", + f' plots_atm="{plots_atm}"', + f' plots_lnd="{plots_lnd}"', + f' plots_original="{plots_original}"', + f' moc_file = "{moc_file_str}"', + ] + ) + + def _generate_ilamb_section(self) -> str: + dependency_ts_atm_monthly: Optional[Dependency] = self.dependency_tracker[ + "ts_atm_monthly" + ] + ts_atm_subsection: str = "" + if dependency_ts_atm_monthly: + ts_atm_subsection = dependency_ts_atm_monthly.subtask_label + dependency_ts_land_monthly: Optional[Dependency] = self.dependency_tracker[ + "ts_land_monthly" + ] + ts_lnd_subsection: str = "" + if dependency_ts_land_monthly: + ts_lnd_subsection = dependency_ts_land_monthly.subtask_label + + dependency_e3sm_to_cmip_atm_monthly: Optional[Dependency] = ( + self.dependency_tracker["e3sm_to_cmip_atm_monthly"] + ) + e3sm_to_cmip_atm_subsection: str = "" + if dependency_e3sm_to_cmip_atm_monthly: + e3sm_to_cmip_atm_subsection = ( + dependency_e3sm_to_cmip_atm_monthly.subtask_label + ) + dependency_e3sm_to_cmip_lnd_monthly: Optional[Dependency] = ( + self.dependency_tracker["e3sm_to_cmip_lnd_monthly"] + ) + e3sm_to_cmip_lnd_subsection: str = "" + if dependency_e3sm_to_cmip_lnd_monthly: + e3sm_to_cmip_lnd_subsection = ( + dependency_e3sm_to_cmip_lnd_monthly.subtask_label + ) + + start_year: int + end_year: int + year_increment: int = 4 + _, start_year, end_year, year_increment = self._get_years_str( + "atm", year_increment + ) + if not (start_year and end_year): + # Use years from lnd data as fall-back option + _, start_year, end_year, year_increment = self._get_years_str( + "lnd", year_increment + ) + + return "\n".join( + [ + "[ilamb]", + "active = True", + f'e3sm_to_cmip_atm_subsection = "{e3sm_to_cmip_atm_subsection}"', + f'e3sm_to_cmip_land_subsection = "{e3sm_to_cmip_lnd_subsection}"', + f'ilamb_obs = "{self.diagnostics_base_path}ilamb_data"', + "nodes = 8", + f'short_name = "{self.case_name}"', + f'ts_atm_subsection = "{ts_atm_subsection}"', + f'ts_land_subsection = "{ts_lnd_subsection}"', + f"ts_num_years = {year_increment}", + f'years = "{start_year}:{end_year}:{year_increment}"', + ] + ) + + def _get_var_str( + self, component: str, h_value: str, standard_vars: Set[str] + ) -> str: + component_h_var_str: str = "" + + component_metadata: Optional[DirectoryMetadata] = self.metadata_dict[component] + if component_metadata: + vars_by_h_value: Dict[str, List[str]] = ( + component_metadata.variables_by_h_value + ) + if h_value in vars_by_h_value.keys(): + component_h_vars = set(vars_by_h_value[h_value]) + if self.use_only_standard_vars and standard_vars: + found_vars: Set[str] = component_h_vars & standard_vars + component_h_var_str = ",".join(found_vars) + else: + # Use ALL found h# variables + # (If `standard_vars=""`, the default is to use all available variables) + component_h_var_str = ",".join(component_h_vars) + # component_h_var_str = "" # Again, "" means use ALL + return component_h_var_str + + def _get_years_str( + self, component: str, year_increment: int = 0 + ) -> Tuple[str, int, int, int]: + component_metadata: Optional[DirectoryMetadata] = self.metadata_dict[component] + if component_metadata and component_metadata.year_ranges: + n: int = len(component_metadata.year_ranges) + if n > 1: + print(f"Warning: Found {n} year ranges for component {component}") + year_range: str = component_metadata.year_ranges[0] + if "-" in year_range: + parts = year_range.split("-") + start_year: int = int(parts[0]) + end_year: int = int(parts[1]) + year_diff: int = end_year - start_year + if (year_increment == 0) or (year_diff % year_increment != 0): + # Condition 1: we are explicitly being asked to find the increment. + # Condition 2: the increment is invalid and needs to be fixed. + year_increment = year_diff # Default to the entire period + potential_increments: List[int] = [20, 10, 5, 2] + for increment in potential_increments: + if year_diff % increment == 0: + year_increment = increment + break + return ( + f"{start_year}:{end_year}:{year_increment}", + start_year, + end_year, + year_increment, + ) + return "", 0, 0, 0 + + +class Dependency(object): + def __init__( + self, + subtask_label: str, + var_str: str, + years_str: str, + year_increment: int, + climo_years: str = "", + ts_years: str = "", + ): + self.subtask_label: str = subtask_label + self.var_str: str = var_str + self.years_str: str = years_str + self.year_increment: int = year_increment + self.climo_years: str = climo_years + self.ts_years: str = ts_years + + +# These functions are used to initialize the ZppyConfigGenerator +def get_machine_specifics() -> Dict[str, str]: + machine_info = MachineInfo() + machine: str = machine_info.machine + config = machine_info.config + username: str = config.get("web_portal", "username") + web_base_path: str = config.get("web_portal", "base_path") + diagnostics_base_path: str = config.get("diagnostics", "base_path") + + output_base_path: str = "" + partition: str = "" + qos: str = "" + if machine == "chrysalis": + output_base_path = "/lcrc/group/e3sm/" + partition = "compute" + qos = "regular" + elif machine == "perlmutter": + output_base_path = "/global/cfs/cdirs/e3sm/" + partition = "" + qos = "regular" + elif machine == "compy": + output_base_path = "/compyfs/" + partition = "slurm" + qos = "regular" + else: + print(f"Warning: invalid machine={machine}") + + return { + "input_path": ARCHIVE_DIR, + "output_path": f"{output_base_path}{username}/{ZPPY_CFG_SUFFIX_OUTPUT}", + "www_path": f"{web_base_path}/{username}/{ZPPY_CFG_SUFFIX_WWW}", + "diagnostics_base_path": f"{diagnostics_base_path}/", + "partition": partition, + "qos": qos, + } + + +def get_case_name(metadata: Dict[str, Optional[DirectoryMetadata]]) -> str: + found_case_names: Set[str] = set() + for component in metadata.keys(): + metadata_for_component: Optional[DirectoryMetadata] = metadata[component] + if metadata_for_component is None: + continue # No metadata to look through + case_names: Set[str] = metadata_for_component.case_names + if len(case_names) > 1: + print( + f"Warning: more than 1 case name found for component={component}. Cases: {case_names}" + ) + for case_name in case_names: + found_case_names.add(case_name) + if len(found_case_names) > 1: + print( + f"Warning: more than 1 case name found across components. Cases: {found_case_names}" + ) + return list(found_case_names)[0] + + +def main(verbose: bool = False): + print(f"Gathering metadata from input path={ARCHIVE_DIR}") + metadata: Dict[str, DirectoryMetadata] = get_metadata_by_component(ARCHIVE_DIR) + + # Filter out None values + metadata = {k: v for k, v in metadata.items() if v is not None} + if not metadata: + print("ERROR: No metadata found. Check ARCHIVE_DIR path.") + return + + # Create generator + generator = ZppyConfigGenerator(metadata, verbose) + + # Generate config + print("Generating zppy configuration...") + config = generator.generate() + + # Write to output file + output_file = Path(f"{ZPPY_CFG_LOCATION}zppy_config_{UNIQUE_ID}.cfg") + with open(output_file, "w") as f: + f.write(config) + + print(f"Configuration written to: {output_file.absolute()}") + print(f"Generated sections for components: {', '.join(metadata.keys())}") + + +if __name__ == "__main__": + # Eventually: we could add a zppy command-line flag `--build` that will simply generate a cfg, + # by calling this script. + # (Note we should never have zppy automatically run that generated cfg, + # because it should always be checked by a human first). + main(True)