diff --git a/workflow/config/config.common.yaml b/workflow/config/config.common.yaml index 8fafacbb8..1d8d0de58 100644 --- a/workflow/config/config.common.yaml +++ b/workflow/config/config.common.yaml @@ -1,3 +1,7 @@ +pudl_cache: + version: v2025.5.0 + base_url: https://s3.us-west-2.amazonaws.com/pudl.catalyst.coop + directory: repo_data/pudl pudl_path: s3://pudl.catalyst.coop/v2025.2.0 foresight: 'perfect' diff --git a/workflow/repo_data/config/config.common.yaml b/workflow/repo_data/config/config.common.yaml index a5900e7eb..8577cf638 100644 --- a/workflow/repo_data/config/config.common.yaml +++ b/workflow/repo_data/config/config.common.yaml @@ -1,5 +1,10 @@ pudl_path: s3://pudl.catalyst.coop/v2025.5.0 +pudl_cache: + version: v2025.5.0 + base_url: https://s3.us-west-2.amazonaws.com/pudl.catalyst.coop + directory: repo_data/pudl + # docs : renewable: EGS: diff --git a/workflow/rules/build_electricity.smk b/workflow/rules/build_electricity.smk index e45204d79..2c1c68d41 100644 --- a/workflow/rules/build_electricity.smk +++ b/workflow/rules/build_electricity.smk @@ -1,6 +1,107 @@ ################# ----------- Rules to Build Network ---------- ################# from itertools import chain +from pathlib import Path + + +PUDL_VERSION = config["pudl_cache"]["version"] +PUDL_BASE_URL = config["pudl_cache"]["base_url"].rstrip("/") +PUDL_DIRECTORY = Path(config["pudl_cache"]["directory"]) / PUDL_VERSION + +PUDL_POWERPLANT_FILES = [ + "out_eia__monthly_generators.parquet", + "out_eia__yearly_generators.parquet", + "core_eia860__scd_generators_energy_storage.parquet", + "core_eia860__scd_plants.parquet", +] + +PUDL_COST_FILES = [ + "core_nrelatb__yearly_projected_financial_cases_by_scenario.parquet", + "core_nrelatb__yearly_projected_cost_performance.parquet", + "core_eiaaeo__yearly_projected_fuel_cost_in_electric_sector_by_type.parquet", +] + +PUDL_DEMAND_FILES = [ + "censusdp1tract.state_2010census_dp1.parquet", + "core_eiaaeo__yearly_projected_generation_in_electric_sector_by_technology.parquet", +] + +PUDL_FUEL_PRICE_FILES = [ + "out_eia__monthly_generators.parquet", + "out_eia__yearly_generators.parquet", + "core_eia860__scd_plants.parquet", +] + +PUDL_FILES = sorted( + set( + PUDL_POWERPLANT_FILES + + PUDL_COST_FILES + + PUDL_DEMAND_FILES + + PUDL_FUEL_PRICE_FILES + ) +) + + +def pudl_paths(filenames): + return [str(PUDL_DIRECTORY / filename) for filename in filenames] + + +PUDL_POWERPLANT_PATHS = pudl_paths(PUDL_POWERPLANT_FILES) +PUDL_COST_PATHS = pudl_paths(PUDL_COST_FILES) +PUDL_FUEL_PRICE_PATHS = pudl_paths(PUDL_FUEL_PRICE_FILES) + + +def demand_pudl_paths(wildcards): + profile = config["electricity"]["demand"]["profile"] + + if profile == "ferc": + return pudl_paths(PUDL_DEMAND_FILES) + + if profile == "eia": + return pudl_paths( + [ + "core_eiaaeo__yearly_projected_generation_in_electric_sector_by_technology.parquet", + ] + ) + + return [] + + +rule retrieve_pudl_file: + output: + str(PUDL_DIRECTORY / "{pudl_file}.parquet"), + params: + url=lambda wildcards: ( + f"{PUDL_BASE_URL}/{PUDL_VERSION}/" f"{wildcards.pudl_file}.parquet" + ), + wildcard_constraints: + pudl_file="|".join(Path(filename).stem for filename in PUDL_FILES), + log: + "logs/retrieve_pudl/{pudl_file}.log", + resources: + mem_mb=1000, + walltime="00:30:00", + shell: + r""" + mkdir -p "$(dirname {output:q})" "$(dirname {log:q})" + + curl -fL \ + --retry 5 \ + --retry-all-errors \ + --connect-timeout 20 \ + --max-time 1800 \ + {params.url:q} \ + --output {output:q}.part \ + > {log:q} 2>&1 + + test -s {output:q}.part + + python -c \ + "import duckdb; duckdb.sql(\"SELECT * FROM read_parquet('{output}.part') LIMIT 1\").fetchall()" \ + >> {log:q} 2>&1 + + mv {output:q}.part {output:q} + """ rule build_shapes: @@ -101,8 +202,9 @@ rule build_bus_regions: rule build_cost_data: params: aeo=config_provider("costs", "aeo"), - pudl_path=config_provider("pudl_path"), + pudl_path=str(PUDL_DIRECTORY), input: + pudl=PUDL_COST_PATHS, efs_tech_costs="repo_data/costs/EFS_Technology_Data.xlsx", efs_icev_costs="repo_data/costs/efs_icev_costs.csv", eia_tech_costs="repo_data/costs/eia_tech_costs.csv", @@ -344,7 +446,6 @@ def demand_raw_data(wildcards): elif profile == "ferc": return [ DATA + "pudl/out_ferc714__hourly_estimated_state_demand.parquet", - DATA + "pudl/censusdp1tract.sqlite", ] elif profile == "eulp": return [ @@ -422,8 +523,9 @@ rule build_electrical_demand: planning_horizons=config["scenario"]["planning_horizons"], renewable_weather_years=config["renewable_weather_years"], snapshots=config["snapshots"], - pudl_path=config_provider("pudl_path"), + pudl_path=str(PUDL_DIRECTORY), input: + pudl=demand_pudl_paths, network=RESOURCES + "{interconnect}/elec_base_network.nc", demand_files=demand_raw_data, demand_scaling_file=demand_scaling_data, @@ -480,7 +582,6 @@ rule build_industry_demand: profile_year=pd.to_datetime(config["snapshots"]["start"]).year, eia_api=config_provider("api", "eia"), snapshots=config_provider("snapshots"), - pudl_path=config_provider("pudl_path"), input: network=RESOURCES + "{interconnect}/elec_base_network.nc", demand_files=demand_raw_data, @@ -629,8 +730,9 @@ rule build_fuel_prices: params: snapshots=config["snapshots"], api_eia=config["api"]["eia"], - pudl_path=config_provider("pudl_path"), + pudl_path=str(PUDL_DIRECTORY), input: + pudl=PUDL_FUEL_PRICE_PATHS, gas_balancing_area=ba_gas_dynamic_fuel_price_files, output: state_ng_fuel_prices=RESOURCES + "{interconnect}/state_ng_power_prices.csv", @@ -665,8 +767,10 @@ def dynamic_fuel_price_files(wildcards): rule build_powerplants: params: - pudl_path=config_provider("pudl_path"), + pudl_path=str(PUDL_DIRECTORY), + renewable_weather_year=config_provider("renewable_weather_years"), input: + pudl=PUDL_POWERPLANT_PATHS, wecc_ads="repo_data/WECC_ADS_public", eia_ads_generator_mapping="repo_data/WECC_ADS_public/eia_ads_generator_mapping_updated.csv", fuel_costs="repo_data/plants/fuelCost22.csv", diff --git a/workflow/rules/common.smk b/workflow/rules/common.smk index 3ddae2eb6..07295d0eb 100644 --- a/workflow/rules/common.smk +++ b/workflow/rules/common.smk @@ -7,8 +7,11 @@ from functools import partial, lru_cache import os, sys, glob -path = workflow.source_path("../scripts/_helpers.py") -sys.path.insert(0, os.path.dirname(path)) +# Cache local modules before importing them from Snakemake's source cache. +for source in ("../scripts/_helpers.py", "../scripts/constants.py"): + source_dir = os.path.dirname(workflow.source_path(source)) + if source_dir not in sys.path: + sys.path.insert(0, source_dir) from _helpers import validate_checksum, update_config_from_wildcards from constants import HOURS_PER_YEAR diff --git a/workflow/scripts/build_cost_data.py b/workflow/scripts/build_cost_data.py index 43a865df9..14a6a8ef2 100644 --- a/workflow/scripts/build_cost_data.py +++ b/workflow/scripts/build_cost_data.py @@ -63,17 +63,8 @@ ] # https://github.com/NREL/ReEDS-2.0/blob/e65ed5ed4ffff973071839481309f77d12d802cd/inputs/plant_characteristics/maxage.csv#L4 -def create_duckdb_instance(): - """Set up DuckDB to read parquet files directly.""" - duckdb.connect(database=":memory:", read_only=False) - # Install httpfs extension to access remote files if needed - duckdb.query("INSTALL httpfs;") - - def load_pudl_atb_data(parquet_path: str): """Loads ATB data directly from parquet files.""" - create_duckdb_instance() - query = f""" WITH finance_cte AS ( SELECT diff --git a/workflow/scripts/build_cutout.py b/workflow/scripts/build_cutout.py index 37e3a62f4..085068897 100644 --- a/workflow/scripts/build_cutout.py +++ b/workflow/scripts/build_cutout.py @@ -139,4 +139,8 @@ logging.info(f"Preparing cutout with parameters {cutout_params}.") features = cutout_params.pop("features", None) cutout = atlite.Cutout(snakemake.output[0], **cutout_params) - cutout.prepare(features=features) + cutout.prepare( + features=features, + monthly_requests=True, + concurrent_requests=False, + ) diff --git a/workflow/scripts/build_demand.py b/workflow/scripts/build_demand.py index 035c8766b..10f23d979 100644 --- a/workflow/scripts/build_demand.py +++ b/workflow/scripts/build_demand.py @@ -255,9 +255,6 @@ def _read_data(self) -> pd.DataFrame: def _read_census_data(self) -> pd.DataFrame: """Reads in census data for population weighting using parquet.""" - duckdb.connect(database=":memory:", read_only=False) - duckdb.query("INSTALL httpfs;") - parquet_path = snakemake.params.pudl_path sql = f""" @@ -699,7 +696,7 @@ class ReadCliu(ReadStrategy): - MECS is used to scale county level data Sources: - - https://data.nrel.gov/submissions/97 + - https://data.nlr.gov/submissions/97 - https://www.eia.gov/consumption/manufacturing/data/2014/#r3 (Table 3.2) - https://loadshape.epri.com/enduse - https://github.com/NREL/Industry-Energy-Tool/ @@ -1814,7 +1811,7 @@ class WriteIndustrial(WriteStrategy): """ Based on county level energy use from 2014. - https://data.nrel.gov/submissions/97 + https://data.nlr.gov/submissions/97 """ def __init__(self, n: pypsa.Network, filepath: str) -> None: @@ -2022,9 +2019,7 @@ def assign_scaler(self): # type DemandScaler assert self.api, "Must provide eia api key" return AeoEnergyScaler(self.api) elif self.scaling_method == "aeo_electricity": - assert self.filepath.startswith( - "s3://pudl.catalyst.coop/", - ), "Must provide pudl S3 URL (s3://pudl.catalyst.coop/...)" + assert self.filepath, "Must provide a local PUDL directory" return AeoElectricityScaler(self.filepath) elif self.scaling_method == "efs": assert self.filepath.endswith(".csv"), "Must provide EFS.csv data" @@ -2126,9 +2121,6 @@ def get_projections(self) -> pd.DataFrame: | 2049 | ### | ### | | 2050 | ### | ### | """ - duckdb.connect(database=":memory:", read_only=False) - duckdb.query("INSTALL httpfs;") - query = f""" SELECT projection_year, diff --git a/workflow/scripts/build_natural_gas.py b/workflow/scripts/build_natural_gas.py index 10cde0b04..f2669c1e4 100644 --- a/workflow/scripts/build_natural_gas.py +++ b/workflow/scripts/build_natural_gas.py @@ -853,7 +853,7 @@ def _expand_costs(self, n: pypsa.Network, costs: pd.DataFrame) -> pd.DataFrame: expanded_costs = [] for invesetment_period in n.investment_periods: # reindex to match any tsa - cost = costs.copy() + cost = costs.loc[costs.index.year == self.year].copy() cost.index = cost.index.map(lambda x: x.replace(year=invesetment_period)) cost = cost.reindex(n.snapshots.get_level_values(1), method="nearest") # set investment periods diff --git a/workflow/scripts/build_powerplants.py b/workflow/scripts/build_powerplants.py index 29c509554..aa64c46f7 100644 --- a/workflow/scripts/build_powerplants.py +++ b/workflow/scripts/build_powerplants.py @@ -12,11 +12,6 @@ logger = logging.getLogger(__name__) -def initialize_duckdb(): - duckdb.connect(database=":memory:", read_only=False) - duckdb.query("INSTALL httpfs;") - - def load_eia_operable_data(parquet_path: str): """Queries the parquet files directly for operable plant data.""" return duckdb.query( @@ -694,9 +689,9 @@ def apply_cems_heat_rates(plants, crosswalk_fn, cems_fn): start_date = f"{data_year}-01-01" end_date = f"{data_year + 1}-01-01" - initialize_duckdb() - eia_data_operable = load_eia_operable_data(snakemake.params.pudl_path) - heat_rates = load_heat_rates_data(snakemake.params.pudl_path, start_date, end_date) + pudl_path = snakemake.params.pudl_path + eia_data_operable = load_eia_operable_data(pudl_path) + heat_rates = load_heat_rates_data(pudl_path, start_date, end_date) eia_data_operable = merge_fc_hr_data( eia_data_operable, diff --git a/workflow/scripts/build_sector_costs.py b/workflow/scripts/build_sector_costs.py index 083931ca2..264f141cc 100644 --- a/workflow/scripts/build_sector_costs.py +++ b/workflow/scripts/build_sector_costs.py @@ -2,7 +2,7 @@ Builds costs data from EFS studies. https://www.nrel.gov/docs/fy18osti/70485.pdf -https://data.nrel.gov/submissions/78 +https://data.nlr.gov/submissions/78 """ from abc import ABC, abstractmethod @@ -22,7 +22,7 @@ # fixed maintenance costs in %/year # From https://www.nrel.gov/docs/fy18osti/71500.pdf -# Data in Sheet "13" at https://data.nrel.gov/submissions/93 +# Data in Sheet "13" at https://data.nlr.gov/submissions/93 # manually claculated as fixed/capex ICE_MAINTENANCE_COSTS = { # units of % / year "light_duty_cars": 11.34, @@ -43,7 +43,7 @@ # maintenance costs in $/mile. This gave very high maintenance costs # when using exogenous lifetime miles assumption. # From https://www.nrel.gov/docs/fy18osti/71500.pdf -# Data in Sheet "A2" at https://data.nrel.gov/submissions/93 +# Data in Sheet "A2" at https://data.nlr.gov/submissions/93 """ ICE_MAINTENANCE_COSTS = { # units of $ / miles "light_duty_cars": 0.033, @@ -96,7 +96,7 @@ class EfsTechnologyData: args: file_path: str - Path to this file https://data.nrel.gov/submissions/78 which is the + Path to this file https://data.nlr.gov/submissions/78 which is the data from this report https://www.nrel.gov/docs/fy18osti/70485.pdf Buildings sector will return both commercial and residential data @@ -254,7 +254,7 @@ def get_capex(self): # noqa: D102 df = df[ (df.Sector == "Transportation") & (df["EFS Case"] == self.efs_case) & (df.Metric == "Capital Cost") ].copy() - source = "NREL EFS at https://data.nrel.gov/submissions/78" + source = "NREL EFS at https://data.nlr.gov/submissions/78" description = "Supplemented with NREL ATB" df["Metric"] = "investment" df = self._format_data_structure(df, source=source, description=description) @@ -278,7 +278,7 @@ def get_efficiency(self): df = df[ (df.Sector == "Transportation") & (df["EFS Case"] == self.efs_case) & (df.Metric == "Main Efficiency") ].copy() - source = "NREL EFS at https://data.nrel.gov/submissions/78" + source = "NREL EFS at https://data.nlr.gov/submissions/78" df["Metric"] = "efficiency" df = self._format_data_structure(df, source=source, description="") df = self._correct_efficiency_units(df) @@ -329,7 +329,7 @@ class EfsIceTransportationData: Only contains ICE vehicles, as this data is manually scraped. - See Table 5 in https://www.nrel.gov/docs/fy18osti/70485.pdf for the sources - - See this file for the raw data https://data.nrel.gov/submissions/93 + - See this file for the raw data https://data.nlr.gov/submissions/93 Args: file_path: str @@ -511,7 +511,7 @@ def get_capex(self): # noqa: D102 & (df.Metric == "Installed Cost") & (df.Units.str.startswith("2016$/kBtu/hr")) ].copy() - source = "NREL EFS at https://data.nrel.gov/submissions/78" + source = "NREL EFS at https://data.nlr.gov/submissions/78" description = "" df["Metric"] = "investment" df = self._format_data_structure(df, source=source, description=description) @@ -529,7 +529,7 @@ def get_lifetime(self): # noqa: D102 def get_efficiency(self): # noqa: D102 df = self.data.copy() df = df[(df.Sector == "Buildings") & (df["EFS Case"] == self.efs_case) & (df.Metric == "Efficiency")].copy() - source = "NREL EFS at https://data.nrel.gov/submissions/78" + source = "NREL EFS at https://data.nlr.gov/submissions/78" df["Metric"] = "efficiency" df["Units"] = "per unit" df = self._format_data_structure(df, source=source, description="") diff --git a/workflow/scripts/build_stock_data.py b/workflow/scripts/build_stock_data.py index c97bb6f07..5a3b966f5 100644 --- a/workflow/scripts/build_stock_data.py +++ b/workflow/scripts/build_stock_data.py @@ -832,7 +832,7 @@ def add_brownfield_lpg( # existing stock efficiencies taken from 2016 EFS Technology data # This is consistent with where future efficiencies are taken from # Historical uses lowest EIA 2017 case where available - # https://data.nrel.gov/submissions/93 + # https://data.nlr.gov/submissions/93 # https://www.nrel.gov/docs/fy18osti/70485.pdf ## Efficiencies of existing stock are quite sensitive! ## diff --git a/workflow/scripts/eia.py b/workflow/scripts/eia.py index 5d27c9145..3a264e77b 100644 --- a/workflow/scripts/eia.py +++ b/workflow/scripts/eia.py @@ -1607,7 +1607,7 @@ def extract_state(description: str) -> str: @staticmethod def map_state_names(state: str) -> str: """Maps state name to code.""" - return "U.S." if state == "U.S. Total" else STATE_CODES[state] + return "U.S." if state in {"U.S.", "U.S. Total"} else STATE_CODES[state] class _GasProduction(DataExtractor): diff --git a/workflow/scripts/opts/sector.py b/workflow/scripts/opts/sector.py index 186ff19a7..1a1f05e9f 100644 --- a/workflow/scripts/opts/sector.py +++ b/workflow/scripts/opts/sector.py @@ -420,7 +420,8 @@ def add_export_limits(n, data, constraint, multiplier=None): # add domestic limits trade = Trade("gas", False, "exports", year, api).get_data() - trade = _format_data(trade, " trade") + trade = trade.loc[trade.index == year] + trade = _format_data(trade, " trade").groupby(level=0).sum() add_import_limits(n, trade, "min", import_min) add_export_limits(n, trade, "min", export_min) @@ -433,7 +434,8 @@ def add_export_limits(n, data, constraint, multiplier=None): # add international limits trade = Trade("gas", True, "exports", year, api).get_data() - trade = _format_data(trade, " trade") + trade = trade.loc[trade.index == year] + trade = _format_data(trade, " trade").groupby(level=0).sum() add_import_limits(n, trade, "min", import_min) add_export_limits(n, trade, "min", export_min) @@ -595,7 +597,7 @@ def add_ev_generation_constraint(n, config, snakemake): Default limits taken from: - (Fig ES2) https://www.nrel.gov/docs/fy18osti/71500.pdf - - (Sheet 6.3 - high case) https://data.nrel.gov/submissions/90 + - (Sheet 6.3 - high case) https://data.nlr.gov/submissions/90 """ mode_mapper = { "light_duty": "lgt", diff --git a/workflow/scripts/plot_network_maps.py b/workflow/scripts/plot_network_maps.py index 3e6c5c9a5..55b54bab9 100644 --- a/workflow/scripts/plot_network_maps.py +++ b/workflow/scripts/plot_network_maps.py @@ -466,9 +466,22 @@ def plot_capacity_map_by_horizon( bus_scale = get_bus_scale(interconnect) if interconnect else 1 line_scale = get_line_scale(interconnect) if interconnect else 1 - if bus_values.empty and kind == "new": + has_new_capacity = ( + bus_values.fillna(0).gt(0).any() or line_values.fillna(0).gt(0).any() or link_values.fillna(0).gt(0).any() + ) + + if kind == "new" and not has_new_capacity: + logger.info("No new capacity found; creating an empty new-capacity map.") fig, ax = plt.subplots(figsize=(10, 10)) - ax.text(0.5, 0.5, "No new capacity built", ha="center", va="center", fontsize=14) + ax.text( + 0.5, + 0.5, + "No new capacity built", + ha="center", + va="center", + fontsize=14, + transform=ax.transAxes, + ) ax.set_title(title, fontsize=TITLE_SIZE, pad=20) ax.axis("off") else: @@ -551,6 +564,25 @@ def plot_capacity_map_by_horizon( remove_sector_buses(bus_values).groupby(["bus", "carrier"]).sum() if not bus_values.empty else bus_values ) + has_new_capacity = ( + bus_values.fillna(0).gt(0).any() or line_values.fillna(0).gt(0).any() or link_values.fillna(0).gt(0).any() + ) + + if kind == "new" and not has_new_capacity: + logger.info("No new capacity found for horizon %s.", horizon) + ax.text( + 0.5, + 0.5, + "No new capacity built", + ha="center", + va="center", + fontsize=14, + transform=ax.transAxes, + ) + ax.set_title(f"{horizon}", fontsize=TITLE_SIZE) + ax.axis("off") + continue + artifacts = _plot_capacity_on_ax( n, bus_values, diff --git a/workflow/scripts/zenodo_downloader.py b/workflow/scripts/zenodo_downloader.py index 4d7df7d9f..1a73d4ad4 100644 --- a/workflow/scripts/zenodo_downloader.py +++ b/workflow/scripts/zenodo_downloader.py @@ -1,6 +1,8 @@ """Download scenarios from Zenodo.""" +import shutil from pathlib import Path +from zipfile import BadZipFile, ZipFile import requests @@ -42,7 +44,10 @@ def get_record_metadata(self, record_id): url = f"https://zenodo.org/api/records/{record_id}" try: - response = requests.get(url) + response = requests.get( + url, + timeout=(30, 120), + ) response.raise_for_status() metadata = response.json() self._metadata_cache[record_id] = metadata @@ -68,7 +73,10 @@ def download_scenario_file(self, scenario_final, scenario, filename, force_redow force_redownload : bool, optional If True, re-download the file even if it exists locally. Default is False. """ - (self.download_dir / "zenodo" / scenario).mkdir(exist_ok=True) + (self.download_dir / "zenodo" / scenario).mkdir( + parents=True, + exist_ok=True, + ) local_filepath = f"{self.download_dir}/zenodo/{scenario}/{filename}" # Check if file already exists locally and skip Zenodo @@ -82,12 +90,9 @@ def download_scenario_file(self, scenario_final, scenario, filename, force_redow record_id = self.scenario_records.get(scenario_final) if not record_id: - print(f"No record ID found for scenario: {scenario_final}") - print("Available scenarios with record IDs:") - for scenario, rec_id in self.scenario_records.items(): - if rec_id is not None: - print(f" - {scenario} (ID: {rec_id})") - return None + raise ValueError( + f"No Zenodo record ID configured for scenario: {scenario_final}", + ) return self._download_file(record_id, filename, Path(local_filepath), force_redownload) @@ -144,71 +149,126 @@ def download_by_record_id(self, record_id, filename, force_redownload=False): # Only proceed with download if needed return self._download_file(record_id, filename, Path(local_filepath), force_redownload) - def _download_file(self, record_id, filename, local_filepath, force_redownload=False): - """ - Internal method to download a file from Zenodo. - - This is only called after confirming the file doesn't exist locally. - """ - # Ensure directory exists + def _download_file( + self, + record_id, + filename, + local_filepath, + force_redownload=False, + ): + """Download a file directly or extract it from a Zenodo ZIP archive.""" local_filepath.parent.mkdir(parents=True, exist_ok=True) - # Get record metadata metadata = self.get_record_metadata(record_id) if not metadata: - return None - - # Find the specific file - target_file = None - for file_info in metadata.get("files", []): - if file_info["key"] == filename: - target_file = file_info - break - - if not target_file: - print(f"File '{filename}' not found in record {record_id}") - print("Available files:") - for file_info in metadata.get("files", []): - print(f" - {file_info['key']}") - return None - - # Download the file - download_url = target_file["links"]["self"] - file_size_mb = target_file["size"] / (1024 * 1024) + raise RuntimeError( + f"Could not retrieve metadata for Zenodo record {record_id}", + ) - print(f"Downloading {filename} from record {record_id}...") - print(f"Size: {file_size_mb:.1f} MB") - print(f"Saving to: {local_filepath}") + files = metadata.get("files", []) - try: - response = requests.get(download_url, stream=True) - response.raise_for_status() + # Some Zenodo records expose the requested NetCDF file directly. + target_file = next( + (file_info for file_info in files if file_info["key"] == filename), + None, + ) - total_size = int(response.headers.get("content-length", 0)) - downloaded_size = 0 + if target_file is not None: + self._download_url( + target_file["links"]["self"], + local_filepath, + target_file["size"], + ) + return str(local_filepath) - with open(local_filepath, "wb") as f: - for chunk in response.iter_content(chunk_size=8192): - if chunk: - f.write(chunk) - downloaded_size += len(chunk) + # Future renewable datasets are published as ZIP archives. + zip_file = next( + (file_info for file_info in files if file_info["key"].lower().endswith(".zip")), + None, + ) + + if zip_file is None: + available_files = ", ".join(file_info["key"] for file_info in files) + raise FileNotFoundError( + f"Neither '{filename}' nor a ZIP archive was found in " + f"Zenodo record {record_id}. Available files: {available_files}", + ) - # Show progress for large files - if total_size > 10 * 1024 * 1024: # Show progress for files > 10MB - progress = (downloaded_size / total_size) * 100 - print(f"\rProgress: {progress:.1f}%", end="", flush=True) + archive_path = local_filepath.parent / zip_file["key"] - if total_size > 10 * 1024 * 1024: - print() # New line after progress + if force_redownload or not archive_path.exists(): + self._download_url( + zip_file["links"]["self"], + archive_path, + zip_file["size"], + ) - print(f"Successfully downloaded {filename}") - return str(local_filepath) + try: + with ZipFile(archive_path) as archive: + matches = [member for member in archive.namelist() if Path(member).name == filename] + + if len(matches) != 1: + raise FileNotFoundError( + f"Expected exactly one '{filename}' in '{archive_path.name}', found {len(matches)}", + ) + + with ( + archive.open(matches[0]) as source, + local_filepath.open("wb") as destination, + ): + shutil.copyfileobj( + source, + destination, + length=1024 * 1024, + ) + + except BadZipFile: + local_filepath.unlink(missing_ok=True) + archive_path.unlink(missing_ok=True) + raise + + except OSError: + local_filepath.unlink(missing_ok=True) + raise + + return str(local_filepath) + + def _download_url(self, download_url, destination, expected_size): + """Download a URL to a local path using a temporary partial file.""" + partial_path = destination.with_suffix( + destination.suffix + ".part", + ) + + print( + f"Downloading {destination.name} ({expected_size / 1024**3:.1f} GiB)...", + ) - except requests.exceptions.RequestException as e: - print(f"Download failed: {e}") - if Path(local_filepath).exists(): - Path(local_filepath).unlink() # Remove partial file - return None + try: + with requests.get( + download_url, + stream=True, + timeout=(30, 300), + ) as response: + response.raise_for_status() + + with partial_path.open("wb") as output: + for chunk in response.iter_content( + chunk_size=1024 * 1024, + ): + if chunk: + output.write(chunk) + + actual_size = partial_path.stat().st_size + if expected_size and actual_size != expected_size: + raise OSError( + f"Incomplete download for {destination.name}: expected {expected_size} bytes, got {actual_size}", + ) + + partial_path.replace(destination) + + except Exception: + partial_path.unlink(missing_ok=True) + raise def list_available_files(self, scenario_name): """List all available files in a scenario dataset.""" @@ -252,16 +312,32 @@ def get_available_scenarios(self): return available -def download_scenario_file(scenario_final, scenario, filename, download_dir="./data/zenodo"): - """Quick function to download a single file from a scenario.""" +def download_scenario_file( + scenario_final, + scenario, + filename, + download_dir="./data", +): + """Download a single file from a configured scenario.""" downloader = ZenodoScenarioDownloader(download_dir) - return downloader.download_scenario_file(scenario_final, scenario, filename) - - -def download_by_record_id(record_id, filename, download_dir="./data/zenodo"): - """Quick function to download a file directly by record ID.""" + return downloader.download_scenario_file( + scenario_final, + scenario, + filename, + ) + + +def download_by_record_id( + record_id, + filename, + download_dir="./data", +): + """Download a file directly using its Zenodo record ID.""" downloader = ZenodoScenarioDownloader(download_dir) - return downloader.download_by_record_id(record_id, filename) + return downloader.download_by_record_id( + record_id, + filename, + ) def list_available_scenarios():