diff --git a/.github/workflows/workflow_actions.yml b/.github/workflows/workflow_actions.yml index 500be92..d3f439f 100644 --- a/.github/workflows/workflow_actions.yml +++ b/.github/workflows/workflow_actions.yml @@ -10,6 +10,9 @@ on: env: FORCE_JAVASCRIPT_ACTIONS_TO_NODE24: true + # Use non-interactive matplotlib backend, to prevent the workflow from + # trying to make interactive windows. + MPLBACKEND: Agg concurrency: group: ${{ github.workflow }}-${{ github.ref }} @@ -17,7 +20,7 @@ concurrency: jobs: # --- JOB 1: RUN TESTS --- - test: + pytest: name: Test (Python ${{ matrix.python-version }}) runs-on: ubuntu-latest strategy: @@ -41,13 +44,9 @@ jobs: run: poetry install --with dev - name: Run Pytest - env: - MPLBACKEND: Agg run: poetry run pytest - name: Smoke test package installation - env: - MPLBACKEND: Agg run: | poetry build python -m venv test_env @@ -55,12 +54,38 @@ jobs: pip install dist/*.whl radas_config -o ./new_config.yaml radas -s hydrogen -c ./new_config.yaml + + # --- JOB 2: RUN RADAS COMMAND, CHECK IT WORKS WITH DEFAULT CONFIG --- + test_radas: + name: Test radas command + runs-on: ubuntu-latest + if: (github.event_name == 'release' || github.event_name == 'workflow_dispatch') + permissions: + contents: write + + steps: + - uses: actions/checkout@v5 + + - name: Install Poetry + run: pipx install "poetry>=2,<3" + + - name: Set up Python + uses: actions/setup-python@v6 + with: + python-version: '3.12' + cache: 'poetry' + + - name: Install dependencies + run: poetry install - # --- JOB 2: BUILD RELEASE ARTIFACTS --- + - name: Run radas command + run: poetry run radas -vvv + + # --- JOB 3: BUILD RELEASE ARTIFACTS --- build_release: name: Build Release - needs: test - if: startsWith(github.ref, 'refs/tags') + needs: [pytest, test_radas] + if: github.event_name == 'release' runs-on: ubuntu-latest steps: - uses: actions/checkout@v5 @@ -78,11 +103,11 @@ jobs: path: dist/ retention-days: 1 - # --- JOB 3: PUBLISH TO PYPI --- + # --- JOB 4: PUBLISH TO PYPI --- publish: name: Publish to PyPI needs: build_release - if: startsWith(github.ref, 'refs/tags') + if: github.event_name == 'release' runs-on: ubuntu-latest environment: name: pypi-publish diff --git a/pyproject.toml b/pyproject.toml index 297509a..eb39aa1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,7 +5,7 @@ name = "radas" # used a date-based format (e.g., 2025.6.2), package managers will treat # Semantic Versioning (i.e. 1.0.0) as an "older" release. The epoch ensures that # the new versioning takes precedence over all legacy YYYY.MM.version releases. -version = "1!1.0.0" +version = "1!1.0.1" description = "Plasma radiated power calculated using OpenADAS" authors = ["Commonwealth Fusion Systems"] readme = "README.md" diff --git a/radas/adas_interface/download_adas_datasets.py b/radas/adas_interface/download_adas_datasets.py index a8ec666..c9df09b 100644 --- a/radas/adas_interface/download_adas_datasets.py +++ b/radas/adas_interface/download_adas_datasets.py @@ -32,7 +32,7 @@ def download_species_data( year_key = f"{year}"[-2:] dataset_prefix = dataset_config["prefix"].lower() - output_filename = data_file_dir / f"{species_name}_{dataset_type}.dat" + output_filename = data_file_dir / f"{species_name}_{dataset_type}_{year_key}.dat" query_path = f"{url_base}/download/{reader_class}/{dataset_prefix}{year_key}/{dataset_prefix}{year_key}_{species_key}.dat" if not output_filename.exists(): diff --git a/radas/adas_interface/read_adf11_file.py b/radas/adas_interface/read_adf11_file.py index 3e06fcc..b8075e3 100644 --- a/radas/adas_interface/read_adf11_file.py +++ b/radas/adas_interface/read_adf11_file.py @@ -14,15 +14,15 @@ def load_library(library_name: str, filepath: Path): def read_adf11_file( - data_file_dir, species_name, dataset_type + data_file_dir, species_name, year, dataset_type ) -> dict: """Open and read an ADF11 OpenADAS file. Uses the format specification from https://www.adas.ac.uk/man/appxa-11.pdf """ - - filename = data_file_dir / f"{species_name}_{dataset_type}.dat" + year_key = f"{year}"[-2:] + filename = data_file_dir / f"{species_name}_{dataset_type}_{year_key}.dat" if not filename.exists(): raise FileNotFoundError(f"{filename} does not exist.") diff --git a/radas/cli.py b/radas/cli.py index 95a0437..e05f349 100644 --- a/radas/cli.py +++ b/radas/cli.py @@ -4,6 +4,7 @@ from pathlib import Path from functools import partial from typing import Optional +import contextlib from .shared import open_yaml_file, default_config_file from .adas_interface.download_adas_datasets import download_species_data @@ -14,7 +15,6 @@ from .time_evolution import calculate_time_evolution from .unit_handling import convert_units, ureg from .mavrin_reference import compare_radas_to_mavrin -from .interpolate_rates import interpolate_dataset @click.command() @@ -69,14 +69,22 @@ def run_radas_cli( verbose=verbose, debug=debug, ) - try: - from ipdb import launch_ipdb_on_exception - - with launch_ipdb_on_exception(): + + if debug: + with _post_mortem_debugger(): run_radas(**kwargs) - except ModuleNotFoundError: + else: run_radas(**kwargs) +def _post_mortem_debugger(): + """Context manager that drops into ipdb on unhandled exceptions, or a no-op if ipdb is not installed.""" + try: + from ipdb import launch_ipdb_on_exception + except ModuleNotFoundError: + print("Warning: --debug set but ipdb is not installed; " + "install with `pip install ipdb` for post-mortem debugging.") + return contextlib.nullcontext() + return launch_ipdb_on_exception() def run_radas( directory: Path, @@ -126,22 +134,9 @@ def run_radas( (species_name in species) or (species == ("all",)) ): datasets[species_name] = read_rate_coeff( - data_file_dir, species_name, configuration + data_file_dir, species_name, configuration, verbose=verbose, ) - if ("electron_density_resolution" in configuration["globals"]) or ("electron_temp_resolution") in configuration["globals"]: - if verbose: - print("Interpolating rate coefficients") - - new_datasets = dict() - for species_name, dataset in datasets.items(): - electron_density_resolution = configuration["globals"].get("electron_density_resolution", dataset.sizes["dim_electron_density"]) - electron_temp_resolution = configuration["globals"].get("electron_temp_resolution", dataset.sizes["dim_electron_temp"]) - new_datasets[species_name] = interpolate_dataset(dataset, - electron_density_resolution = electron_density_resolution, - electron_temp_resolution = electron_temp_resolution) - datasets = new_datasets - output_dir.mkdir(exist_ok=True, parents=True) if not debug: with mp.Pool() as pool: @@ -150,11 +145,15 @@ def run_radas( species_name: datasets[species_name] for species_name in species } + # Sort by atomic number, to process heavier elements first since they take longer + # N.b. there will be a race condition, so it might not appear these start first + sorted_datasets = dict(sorted(datasets.items(), key=lambda item: item[1].atomic_number, reverse=True)) + pool.map( partial( run_radas_computation, output_dir=output_dir, verbose=verbose ), - [(ds) for ds in datasets.values()], + [(ds) for ds in sorted_datasets.values()], ) else: for ds in datasets.values(): @@ -199,6 +198,9 @@ def run_radas_computation(dataset: xr.Dataset, output_dir: Path, verbose: int): output_dir.mkdir(exist_ok=True) dataset.pint.dequantify().to_netcdf(output_dir / f"{dataset.species_name}.nc") + if verbose: + print(f"Finished computation for {dataset.species_name}") + @click.command() @click.option( diff --git a/radas/config.yaml b/radas/config.yaml index 68734d2..c777ff5 100644 --- a/radas/config.yaml +++ b/radas/config.yaml @@ -10,14 +10,13 @@ globals: # electron density (ne) * residence time (tau) (in m^-3 s) ne_tau: - value: [0.5e+17] + value: [0.5e+16, 0.5e+17, 0.5e+18] units: "m^-3 s" # Number of log-spaced points to use for electron density and electron temp. # Interpolation is used to map from raw data to points. - # To disable interpolation, do not provide these inputs. - electron_density_resolution: 50 - electron_temp_resolution: 100 + electron_density_resolution: 20 + electron_temp_resolution: 80 data_file_config: adf11: @@ -459,4 +458,4 @@ species: atomic_number: 81 lead: atomic_symbol: "Pb" - atomic_number: 82 + atomic_number: 82 \ No newline at end of file diff --git a/radas/interpolate_rates.py b/radas/interpolate_rates.py index 73b0f57..88629cb 100644 --- a/radas/interpolate_rates.py +++ b/radas/interpolate_rates.py @@ -1,78 +1,86 @@ -"""Routines to interpolate a dataset of rate coefficients to higher resolution.""" +"""Routines for log-log interpolation of rate coefficients with boundary clipping.""" import xarray as xr import numpy as np from scipy.interpolate import RectBivariateSpline from numpy.typing import NDArray +import warnings -def interpolate_array(array: xr.DataArray, new_electron_density: NDArray[np.floating], new_electron_temp: NDArray[np.floating]) -> xr.DataArray: - """Interpolate array onto new values for the electron density and electron temp. +def is_significantly_below(requested, limit): + return requested < limit and not np.isclose(requested, limit) + +def is_significantly_above(requested, limit): + return requested > limit and not np.isclose(requested, limit) + +def interpolate_array( + array: xr.DataArray, + new_electron_density: NDArray[np.floating], + new_electron_temp: NDArray[np.floating] +) -> xr.DataArray: + """ + Interpolate rate coefficients onto a new density/temperature grid in log-log space. - The interpolation is performed for logarithmic values. + Uses nearest-neighbor extrapolation by clipping out-of-bounds coordinates to + the original grid edges. """ units = array.pint.units array = array.pint.dequantify().squeeze() + # Handle zero-value edge cases (log of zero is undefined) if np.allclose(array, 0.0, atol=0.0, rtol=1e-6): - # If all values of the array are zero, return a zero array. - return xr.DataArray(np.zeros((np.size(new_electron_temp), np.size(new_electron_density))), + return xr.DataArray( + np.zeros((np.size(new_electron_temp), np.size(new_electron_density))), coords=dict(dim_electron_temp=new_electron_temp, dim_electron_density=new_electron_density) ) * units - elif np.any(array <= 0.0): - # If only some of the values of the array are zero, raise an error. - raise NotImplementedError("Cannot handle zero-valued entries in non-zero rate coefficients.") + if np.any(array <= 0.0): + raise NotImplementedError("Cannot log-interpolate rate coefficients containing zeros.") + + # Check if extrapolation is needed and raise a warning if this is the case. + out_of_bounds_msg = [] + + req_dens_min, req_dens_max = new_electron_density.min(), new_electron_density.max() + grid_dens_min, grid_dens_max = array.dim_electron_density.min(), array.dim_electron_density.max() + + if is_significantly_below(req_dens_min, grid_dens_min) or is_significantly_above(req_dens_max, grid_dens_max): + out_of_bounds_msg.append( + f"Density requested [{req_dens_min:.2e}, {req_dens_max:.2e}] " + f"exceeds grid [{grid_dens_min:.2e}, {grid_dens_max:.2e}]." + ) + + # Check Temperature Bounds + req_temp_min, req_temp_max = new_electron_temp.min(), new_electron_temp.max() + grid_temp_min, grid_temp_max = array.dim_electron_temp.min(), array.dim_electron_temp.max() + + if is_significantly_below(req_temp_min, grid_temp_min) or is_significantly_above(req_temp_max, grid_temp_max): + out_of_bounds_msg.append( + f"Temperature requested [{req_temp_min:.2e}, {req_temp_max:.2e}] " + f"exceeds grid [{grid_temp_min:.2e}, {grid_temp_max:.2e}]." + ) + + if out_of_bounds_msg: + full_msg = "Nearest-neighbour extrapolation used for off-grid values: " + " ".join(out_of_bounds_msg) + warnings.warn(full_msg, RuntimeWarning) + + # ------------------------------------ + + # Prepare original grid and data in log10 space x = np.log10(array.dim_electron_density) y = np.log10(array.dim_electron_temp) z = np.log10(array.transpose("dim_electron_density", "dim_electron_temp").pint.magnitude) + # Transform target coordinates to log10 x_interp = np.log10(new_electron_density) y_interp = np.log10(new_electron_temp) - z_interp = np.power(10, RectBivariateSpline(x, y, z)(x_interp, y_interp, grid=True).T) - - return xr.DataArray(z_interp, - coords=dict(dim_electron_temp=new_electron_temp, dim_electron_density=new_electron_density) - ) * units -def interpolate_dataset(dataset: xr.Dataset, electron_density_resolution: int, electron_temp_resolution: int) -> xr.Dataset: - """Interpolate all rate coefficients in a dataset.""" - new_electron_density = np.logspace( - np.log10(dataset["dim_electron_density"].min().item()), - np.log10(dataset["dim_electron_density"].max().item()), - num = electron_density_resolution - ) + # Force nearest-neighbor extrapolation by clipping points to the grid domain + x_clipped = np.clip(x_interp, x.min().item(), x.max().item()) + y_clipped = np.clip(y_interp, y.min().item(), y.max().item()) - new_electron_temp = np.logspace( - np.log10(dataset["dim_electron_temp"].min().item()), - np.log10(dataset["dim_electron_temp"].max().item()), - num = electron_temp_resolution - ) + # Perform spline interpolation and revert from log space + z_interp_log = RectBivariateSpline(x, y, z)(x_clipped, y_clipped, grid=True) + z_interp = np.power(10, z_interp_log.T) - new_dataset = xr.Dataset().assign_attrs(dataset.attrs) - - for key, array in dataset.items(): - if key in [ - "electron_density", - "electron_temp", - ]: - # Don't copy in the coordinate arrays which we'll interpolate - continue - elif key in [ - "ne_tau" - ]: - # Directly copy in the coordinate arrays which we'll leave unchanged - new_dataset[key] = array - elif array.ndim == 0: - # Directly copy in scalar arrays - new_dataset[key] = array - elif (("dim_electron_density" in array.coords) - and ("dim_electron_temp" in array.coords) - and ("dim_charge_state" in array.coords)): - # For each charge state, interpolate the rate coefficient - new_dataset[key] = array.groupby("dim_charge_state").map(interpolate_array, args=(new_electron_density, new_electron_temp)) - else: - raise NotImplementedError(f"Could not process array '{key}' with coords {array.coords}") - - new_dataset["electron_density"] = xr.DataArray(new_electron_density, dims="dim_electron_density") * dataset["electron_density"].pint.units - new_dataset["electron_temp"] = xr.DataArray(new_electron_temp, dims="dim_electron_temp") * dataset["electron_temp"].pint.units - - return new_dataset \ No newline at end of file + return xr.DataArray( + z_interp, + coords=dict(dim_electron_temp=new_electron_temp, dim_electron_density=new_electron_density) + ) * units \ No newline at end of file diff --git a/radas/mavrin_reference/compare_to_mavrin.py b/radas/mavrin_reference/compare_to_mavrin.py index 4f16e88..7b8eda9 100644 --- a/radas/mavrin_reference/compare_to_mavrin.py +++ b/radas/mavrin_reference/compare_to_mavrin.py @@ -1,12 +1,11 @@ import xarray as xr -import numpy as np from pathlib import Path import matplotlib.pyplot as plt from .read_mavrin_data import ( read_mavrin_data, compute_Mavrin_polynomial_fit, ) -from ..unit_handling import ureg, magnitude +from ..unit_handling import ureg, magnitude_in_units def compare_radas_to_mavrin(output_dir: Path): @@ -17,7 +16,7 @@ def compare_radas_to_mavrin(output_dir: Path): compare_radas_to_mavrin_per_species(output_dir, species) -def compare_radas_to_mavrin_per_species(output_dir: Path, species: str, max_decades: int = 6): +def compare_radas_to_mavrin_per_species(output_dir: Path, species: str, max_decades: int = 4, show: bool=False): mavrin_data = read_mavrin_data() ds = xr.open_dataset(output_dir / f"{species}.nc").pint.quantify() @@ -46,41 +45,56 @@ def compare_radas_to_mavrin_per_species(output_dir: Path, species: str, max_deca Lz_radas = ds["equilibrium_Lz"].pint.to(ureg.W * ureg.m**3) mean_charge_radas = ds["equilibrium_mean_charge_state"].pint.to(ureg.dimensionless) - fig, axs = plt.subplots(ncols=2) + fig, axs = plt.subplots(ncols=2, nrows=2, sharex="all", sharey="row") for i in range(ds.sizes["dim_ne_tau"]): ne_tau = ds.ne_tau.isel(dim_ne_tau=i).item() - Lz_radas.isel(dim_ne_tau=i).plot(ax=axs[0], label=f"{ne_tau:~P}", color=f"C{i}") + Lz_radas.isel(dim_ne_tau=i).plot(ax=axs[0][0], color=f"C{i}") if Lz_mavrin is not None: - Lz_mavrin.isel(dim_ne_tau=i).plot(ax=axs[0], color=f"C{i}", linestyle="--") + Lz_mavrin.isel(dim_ne_tau=i).plot(ax=axs[0][1], color=f"C{i}", label=f"{ne_tau:~.1P}") + else: + axs[0][1].plot([], [], color=f"C{i}", label=f"{ne_tau:~.1P}") - mean_charge_radas.isel(dim_ne_tau=i).plot(ax=axs[1], color=f"C{i}") + mean_charge_radas.isel(dim_ne_tau=i).plot(ax=axs[1][0], color=f"C{i}") if mean_charge_mavrin is not None: - mean_charge_mavrin.isel(dim_ne_tau=i).plot( - ax=axs[1], color=f"C{i}", linestyle="--" - ) + mean_charge_mavrin.isel(dim_ne_tau=i).plot(ax=axs[1][1], color=f"C{i}") - ds["coronal_Lz"].pint.to(ureg.W * ureg.m**3).plot(ax=axs[0], label="coronal", color="k", linestyle="--") - ds["coronal_mean_charge_state"].pint.to(ureg.dimensionless).plot(ax=axs[1], label="coronal", color="k", linestyle="--") + ds["coronal_Lz"].pint.to(ureg.W * ureg.m**3).plot(ax=axs[0][0], color="k") + axs[0][1].plot([], [], color="k", label="coronal") + ds["coronal_mean_charge_state"].pint.to(ureg.dimensionless).plot(ax=axs[1][0], color="k") - axs[0].legend() - axs[0].set_yscale("log") + axs[0][1].legend() - Lz_radas_mag = magnitude(Lz_radas) - mean_charge_radas_mag = magnitude(mean_charge_radas) + Lz_radas_mag = magnitude_in_units(Lz_radas, ureg.W * ureg.m**3) + Lz_coronal_mag = magnitude_in_units(ds["coronal_Lz"], ureg.W * ureg.m**3) + Lz_min = min(Lz_radas_mag.min(), Lz_coronal_mag.min()) + Lz_max = max(Lz_radas_mag.max(), Lz_coronal_mag.max()) - axs[0].set_ylim(max(np.min(Lz_radas_mag), np.max(Lz_radas_mag) / 10**max_decades) / 2, np.max(Lz_radas_mag) * 2) - axs[1].set_ylim(0, np.max(mean_charge_radas_mag) * 1.2) + for ax in axs.flatten(): + ax.set_title("") + ax.set_ylabel("") + + axs[0][0].set_yscale("log") + axs[0][0].set_ylim(max(Lz_min, Lz_max / 10**max_decades) / 2, Lz_max * 2) + axs[1][0].set_ylim(0, ds.atomic_number * 1.2) + + axs[0][0].set_ylabel("$L_z$ $[W m^3]$") + axs[1][0].set_ylabel("$$") - axs[0].set_title("$L_z$ $[W m^3]$") - axs[1].set_title("$$") + axs[0][0].set_title("radas") + axs[0][1].set_title("Mavrin") for ax in axs.flatten(): ax.set_xscale("log") + ax.set_xlabel("") + + for ax in axs[-1][:].flatten(): ax.set_xlabel("$T_e$ [$eV$]") - ax.set_ylabel("") plt.suptitle(species) - plt.savefig(output_dir / f"{species}.png") + if show: + plt.show() + + plt.savefig(output_dir / f"{species}.png", dpi=300) diff --git a/radas/read_rate_coeffs.py b/radas/read_rate_coeffs.py index 3e895fc..a3e279a 100644 --- a/radas/read_rate_coeffs.py +++ b/radas/read_rate_coeffs.py @@ -2,60 +2,136 @@ from .adas_interface.determine_adas_dataset_type import ( determine_reader_class_and_config, ) -from .shared import get_git_revision_short_hash from importlib.metadata import version, PackageNotFoundError import datetime import xarray as xr import numpy as np +import warnings +from .interpolate_rates import interpolate_array +# Reference units for non-dimensionalizing coordinates +reference_electron_density = Quantity(1.0, ureg.m**-3) +reference_electron_temp = Quantity(1.0, ureg.eV) -def read_rate_coeff(data_file_dir, species_name, config): - """Builds a rate_dataset combining all of the raw data available for a given species.""" - config_for_species = config["species"][species_name] - +def read_rate_coeff(data_file_dir, species_name, config, verbose=0): + """ + Main pipeline to assemble an atomic rate dataset for a specific species. + + Reads raw ADAS files, standardizes their grids, aligns charge states, + and attaches metadata. + """ try: - radas_version=version("radas") + radas_version = version("radas") except PackageNotFoundError: - radas_version="UNDEFINED" + radas_version = "UNDEFINED" + + # 1. Collect and sort data by year + rate_coefficients = build_sorted_dictionary_of_rate_coefficients(config, species_name, data_file_dir) + + # 2. Resample all datasets to a common resolution + rate_coefficients = interpolate_rates_onto_matching_grids(config, species_name, rate_coefficients, verbose=verbose) - dataset = xr.Dataset().assign_attrs( - atomic_number=config_for_species["atomic_number"], + # 3. Merge individual datasets (e.g., recombination, ionization) into one + try: + dataset = xr.merge([v.rename(k) for k, v in rate_coefficients.items()], join="exact") + except xr.AlignmentError as e: + raise xr.AlignmentError(f"Alignment failed for {species_name}: {e}") + + if dataset.sizes["dim_electron_density"] <= 2 or dataset.sizes["dim_electron_temp"] <= 2: + raise xr.AlignmentError(f"Alignment resulted in dataset sizes {dataset.sizes} for {species_name}.") + + # Convert dimensionless coordinates back to physical quantities + dataset["electron_density"] = dataset["dim_electron_density"] * reference_electron_density + dataset["electron_temp"] = dataset["dim_electron_temp"] * reference_electron_temp + dataset["reference_electron_density"] = reference_electron_density + dataset["reference_electron_temp"] = reference_electron_temp + + # 4. Standardize charge state indexing and attach global attributes + dataset = align_rates_on_charge_states(dataset) + dataset = dataset.assign_attrs( + atomic_number=config["species"][species_name]["atomic_number"], species_name=species_name, - git_hash=get_git_revision_short_hash(), radas_version=radas_version, created=datetime.date.today().strftime("%Y-%b-%d"), ) - dataset = write_global_attributes(dataset, config["globals"]) + return write_global_attributes(dataset, config["globals"]) + +def build_sorted_dictionary_of_rate_coefficients(config, species_name, data_file_dir): + """Make a dictionary of rate coefficient datasets, ordered most-recent first.""" + rate_coefficients = dict() + years = dict() - for dataset_type in config_for_species["data_files"].keys(): + for dataset_type, file_to_read in config["species"][species_name]["data_files"].items(): reader_key, dataset_config = determine_reader_class_and_config( config["data_file_config"], dataset_type ) + # Extract year for sorting; entries can be a simple int or [file, year] list + if isinstance(file_to_read, int): + years[dataset_type] = file_to_read + elif isinstance(file_to_read, list) and len(file_to_read) == 2: + years[dataset_type] = file_to_read[1] + else: + raise NotImplementedError(f"Unsupported config format for {species_name} {dataset_type}") + if reader_key == "adf11": rate_dataset = build_adf11_rate_dataset( - data_file_dir, - species_name, - dataset_type, - dataset_config, + data_file_dir, species_name, years[dataset_type], dataset_type, dataset_config, ) else: - raise NotImplementedError( - f"No implementation for reading {reader_key} files." - ) + raise NotImplementedError(f"No implementation for reader: {reader_key}") - determine_coordinates(dataset, rate_dataset) - dataset[dataset_type] = rate_dataset.rate_coefficient + rate_coefficients[dataset_type] = rate_dataset.rate_coefficient + + # Sort keys by year descending + sorted_keys = sorted(years, key=years.get, reverse=True) + return {k: rate_coefficients[k] for k in sorted_keys} - dataset = align_rates_on_charge_states(dataset) +def interpolate_rates_onto_matching_grids(config, species_name, rate_coefficients, verbose): + """Resample all rate coefficients to a uniform log-grid defined by the newest dataset.""" + + # Use the range of the most recent dataset to define the master grid + most_recent_rate_coeff = list(rate_coefficients.values())[0] - return dataset + new_electron_density = np.logspace( + np.log10(most_recent_rate_coeff["dim_electron_density"].min().item()), + np.log10(most_recent_rate_coeff["dim_electron_density"].max().item()), + num = config["globals"]["electron_density_resolution"] + ) + + new_electron_temp = np.logspace( + np.log10(most_recent_rate_coeff["dim_electron_temp"].min().item()), + np.log10(most_recent_rate_coeff["dim_electron_temp"].max().item()), + num = config["globals"]["electron_temp_resolution"] + ) + + interpolated_rate_coefficients = dict() + for key, value in rate_coefficients.items(): + with warnings.catch_warnings(record=True) as captured_warnings: + warnings.simplefilter("always") + # Map interpolation across charge states + interpolated_rate_coefficients[key] = value.groupby("dim_charge_state").map( + interpolate_array, args=(new_electron_density, new_electron_temp) + ) + + if verbose: + for w in captured_warnings: + warnings.warn_explicit( + message=f"when interpolating {key} for {species_name}. {w.message}", + category=w.category, + filename=w.filename, + lineno=w.lineno, + ) + + return interpolated_rate_coefficients def write_global_attributes(dataset: xr.Dataset, globals: dict) -> xr.Dataset: + """Attach global configuration parameters to the dataset as attributes or DataArrays.""" for attribute, value in globals.items(): if isinstance(value, dict): + # Complex attributes (with units) are added as coordinates/DataArrays if np.ndim(value["value"]) >= 1: dataset[attribute] = xr.DataArray( Quantity(value["value"], value["units"]), @@ -64,101 +140,58 @@ def write_global_attributes(dataset: xr.Dataset, globals: dict) -> xr.Dataset: else: dataset[attribute] = Quantity(value["value"], value["units"]) else: + # Simple metadata (strings/ints) added as attributes dataset[attribute] = value - return dataset - -def determine_coordinates(dataset: xr.Dataset, rate_dataset: xr.Dataset): - - for key in [ - "electron_density", - "electron_temp", - "reference_electron_density", - "reference_electron_temp", - ]: - if key not in dataset: - dataset[key] = rate_dataset[key] - - for key in ["electron_density", "electron_temp"]: - np.testing.assert_allclose( - dimensionless_magnitude( - (dataset[key] - rate_dataset[key]) / dataset[f"reference_{key}"] - ), - 0.0, - ) - - -def build_adf11_rate_dataset( - data_file_dir, species_name, dataset_type, dataset_config -): +def build_adf11_rate_dataset(data_file_dir, species_name, year, dataset_type, dataset_config): + """Read a specific ADF11 file and format it as a quantified xarray Dataset.""" from .adas_interface.read_adf11_file import read_adf11_file - data = read_adf11_file(data_file_dir, species_name, dataset_type) - + data = read_adf11_file(data_file_dir, species_name, year, dataset_type) ds = xr.Dataset() - ds["species"] = species_name - ds["dataset"] = dataset_type - ds["charge"] = data["IZMAX"] - - electron_density = convert_units( - Quantity(10 ** data["DDENSD"][: data["IDMAXD"]], ureg.cm**-3), ureg.m**-3 - ) - electron_temp = Quantity(10 ** data["DTEVD"][: data["ITMAXD"]], ureg.eV) - - # Use logarithmic quantities to define the coordinates, so that we can interpolate over logarithmic quantities. - ds["electron_density"] = xr.DataArray( - electron_density, coords=dict(dim_electron_density=electron_density.magnitude) - ) - ds["electron_temp"] = xr.DataArray( - electron_temp, coords=dict(dim_electron_temp=electron_temp.magnitude) - ) - - ds["reference_electron_density"] = Quantity(1.0, ureg.m**-3) - ds["reference_electron_temp"] = Quantity(1.0, ureg.eV) + # Log values stored in ADAS files are converted to linear scale if required + electron_density = convert_units(Quantity(10**data["DDENSD"][:data["IDMAXD"]], ureg.cm**-3), ureg.m**-3) + electron_temp = Quantity(10**data["DTEVD"][:data["ITMAXD"]], ureg.eV) - ds["number_of_charge_states"] = data["IZMAX"] - charge_state = np.arange(data["IZMAX"]) - ds["charge_state"] = xr.DataArray( - charge_state, coords=dict(dim_charge_state=charge_state) - ) - - coefficient = data["DRCOFD"][: data["IZMAX"], : data["ITMAXD"], : data["IDMAXD"]] + coefficient = data["DRCOFD"][:data["IZMAX"], :data["ITMAXD"], :data["IDMAXD"]] if dataset_config["code"] <= 9: coefficient = 10**coefficient - input_units = dataset_config["stored_units"] - output_units = dataset_config["desired_units"] - ds["rate_coefficient"] = convert_units( - xr.DataArray( - coefficient, - dims=("dim_charge_state", "dim_electron_temp", "dim_electron_density"), - ).pint.quantify(input_units), - output_units, - ) + # Create dimensionless indices for internal processing + dim_electron_density = dimensionless_magnitude(electron_density / reference_electron_density) + dim_electron_temp = dimensionless_magnitude(electron_temp / reference_electron_temp) + + rate_coefficient = xr.DataArray(coefficient, coords=dict( + dim_charge_state = np.arange(data["IZMAX"]), + dim_electron_temp = dim_electron_temp, + dim_electron_density = dim_electron_density, + )).pint.quantify(dataset_config["stored_units"]) + ds["rate_coefficient"] = convert_units(rate_coefficient, dataset_config["desired_units"]) return ds - def align_rates_on_charge_states(dataset: xr.Dataset) -> xr.Dataset: - """For rates which are for k+1->k reactions, we shift these by one position - in the dim_charge_state dimension so that the kth position of a rate always - refers to the reactant species.""" - - dataset = dataset.pad( - pad_width=dict(dim_charge_state=(0, 1)), mode="constant", constant_values=0.0 - ) - dataset = dataset.assign_coords( - dim_charge_state=np.arange(dataset.sizes["dim_charge_state"]) - ) - - for key in [ + """ + Standardize charge state mapping so index 'k' always refers to the reactant. + + For k+1 -> k reactions (e.g. recombination), the rates are shifted so that + index k represents the species being recombined. + """ + # Pad to accommodate the N+1 charge state after shifting + dataset = dataset.pad(pad_width=dict(dim_charge_state=(0, 1)), mode="constant", constant_values=0.0) + dataset = dataset.assign_coords(dim_charge_state=np.arange(dataset.sizes["dim_charge_state"])) + + # Shift k+1 -> k processes + keys_to_shift = [ "effective_recombination", "charge_exchange_cross_coupling", "recombination_and_bremsstrahlung", "charge_exchange_emission", - ]: + ] + + for key in [k for k in keys_to_shift if k in dataset]: dataset[key] = dataset[key].roll(dim_charge_state=+1) - return dataset + return dataset \ No newline at end of file diff --git a/radas/shared.py b/radas/shared.py index f8768b2..87ffa66 100644 --- a/radas/shared.py +++ b/radas/shared.py @@ -2,7 +2,6 @@ from pathlib import Path from importlib.resources import files -import subprocess import yaml default_config_file = files("radas").joinpath("config.yaml") @@ -11,18 +10,6 @@ library_extensions = [".a", ".so"] -def get_git_revision_short_hash() -> str: - try: - return ( - subprocess.check_output(["git", "rev-parse", "--short", "HEAD"]) - .decode("ascii") - .strip() - ) - except: # noqa:E722 - # If git isn't available (sometimes the case in tests), return a blank - return "UNDEFINED" - - def open_yaml_file(yaml_file: Path) -> dict: with open(yaml_file, "r") as file: return yaml.load(file, Loader=yaml.FullLoader)