diff --git a/changelog/207.feature.rst b/changelog/207.feature.rst new file mode 100644 index 00000000..44d76430 --- /dev/null +++ b/changelog/207.feature.rst @@ -0,0 +1 @@ +Added CHIANTI line-list and spectral response-function helpers to `sunkit_instruments.response`. diff --git a/pyproject.toml b/pyproject.toml index 619d5adb..40dab177 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,7 +28,17 @@ dependencies = [ dynamic = ["version"] [project.optional-dependencies] +chianti = [ + "ChiantiPy>=0.15.0", + "h5netcdf", + "numexpr", + "periodictable", +] tests = [ + "ChiantiPy>=0.15.0", + "h5netcdf", + "numexpr", + "periodictable", "pytest", "pytest-astropy", "pytest-doctestplus", diff --git a/sunkit_instruments/response/__init__.py b/sunkit_instruments/response/__init__.py index df185d83..d4732b02 100644 --- a/sunkit_instruments/response/__init__.py +++ b/sunkit_instruments/response/__init__.py @@ -2,4 +2,13 @@ A subpackage for computing instrument responses """ +from sunkit_instruments.response.abstractions import AbstractChannel, EmissionModel, LineEmissionModel +from sunkit_instruments.response.linelist import ( + LineListEmissionModel, + chianti_line_list, + get_line_list, + line_list_cache_path, + line_list_from_emission_model, +) +from sunkit_instruments.response.spectral import create_response_function, get_spectral_response from sunkit_instruments.response.thermal import SourceSpectra, get_temperature_response diff --git a/sunkit_instruments/response/abstractions.py b/sunkit_instruments/response/abstractions.py index 41a148d1..b9cecfc4 100644 --- a/sunkit_instruments/response/abstractions.py +++ b/sunkit_instruments/response/abstractions.py @@ -1,39 +1,130 @@ """This module defines abstractions for computing instrument response.""" import abc +import copy +from typing import Protocol, runtime_checkable import astropy.units as u -__all__ = ["AbstractChannel"] +__all__ = ["AbstractChannel", "EmissionModel", "LineEmissionModel"] + + +@runtime_checkable +class EmissionModel(Protocol): + """ + Protocol for emission models usable with + `~sunkit_instruments.response.get_temperature_response`. + + An emission model bundles the atomic physics: it knows how a plasma + emits as a function of temperature and can fold that emission through + an instrument channel. Any object providing this interface conforms — + no inheritance from this class is required. In particular + `synthesizAR.atomic.EmissionModel` satisfies it. + + Notes + ----- + Implementations should document whether the emissivity used by + `calculate_temperature_response` already includes the (equilibrium) + ionization fraction, since some models (e.g. synthesizAR) defer it + to support time-dependent ionization. + """ + + @property + def temperature(self) -> u.Quantity: + """Temperature grid of the model.""" + + def calculate_temperature_response(self, channel) -> u.Quantity: + """ + Temperature response of ``channel`` for this model's emission. + + ``channel`` is compatible with + `~sunkit_instruments.response.abstractions.AbstractChannel`; its + ``wavelength`` and ``wavelength_response()`` are the interface + implementations should rely on. + """ + + +@runtime_checkable +class LineEmissionModel(Protocol): + """ + Protocol for emission models that expose per-transition line emissivities. + + This is the surface consumed by + `~sunkit_instruments.response.line_list_from_emission_model` to build a + line list for spectrally-resolved response functions + (`~sunkit_instruments.response.create_response_function`). The model is + an iterable of ion objects (each providing ``atomic_number``, + ``ion_name_roman`` and ``ionization_fraction``), matching + `synthesizAR.atomic.EmissionModel` and `fiasco.IonCollection`. + """ + + @property + def temperature(self) -> u.Quantity: + """Temperature grid of the model.""" + + @property + def density(self) -> u.Quantity: + """Density grid of the model.""" + + def __iter__(self): + """Iterate over the ions in the model.""" + + def get_line_emissivity(self, ion, transition=None): + """ + Per-transition line emissivity for one ion. + + Returns a ``(wavelength, emissivity)`` pair where ``emissivity`` has + shape ``(temperature, density, transition)`` in photon units + (``cm3 ph s-1``) and excludes the ionization fraction, matching the + synthesizAR convention. + """ class AbstractChannel(abc.ABC): """ An abstract base class for defining instrument channels. + A channel instance is a snapshot of the instrument at one time: + time-dependent degradation enters exclusively through `at`, and every + consumer-facing method (`effective_area`, `wavelength_response`) takes + no time argument. Subclasses implement the time dependence in + `degradation`. + For all methods and properties defined here, see the topic guide on instrument response for more information. """ + _obstime = None + + def at(self, obstime): + """ + A copy of this channel evaluated at ``obstime``. + + The single way time enters a response calculation: bind the time + first, then pass the returned channel wherever a channel is + accepted:: + + get_temperature_response(channel.at("2020-01-01"), model) + + Parameters + ---------- + obstime: any format parsed by `sunpy.time.parse_time` + """ + bound = copy.copy(self) + bound._obstime = obstime + return bound + @u.quantity_input - def wavelength_response( - self, obstime=None - ) -> u.cm**2 * u.DN * u.steradian / (u.photon * u.pixel): + def wavelength_response(self) -> u.cm**2 * u.DN * u.steradian / (u.photon * u.pixel): """ Instrument response as a function of wavelength The wavelength response is the effective area with the conversion factors from photons to DN and steradians - to pixels. - - Parameters - ---------- - obstime: any format parsed by `~sunpy.time.parse_time`, optional - If specified, this is used to compute the time-dependent - instrument degradation. + to pixels. For time-dependent degradation, bind the time first + with `at`. """ - area_eff = self.effective_area(obstime=obstime) return ( - area_eff + self.effective_area() * self.energy_per_photon * self.pixel_solid_angle * self.camera_gain @@ -41,27 +132,25 @@ def wavelength_response( ) @u.quantity_input - def effective_area(self, obstime=None) -> u.cm**2: + def effective_area(self) -> u.cm**2: """ Effective area as a function of wavelength. The effective area is the geometrical collecting area weighted by the mirror reflectance, filter transmittance, - quantum efficiency, and instrument degradation. - - Parameters - ---------- - obstime: any format parsed by `sunpy.time.parse_time`, optional - If specified, this is used to compute the time-dependent - instrument degradation. + quantum efficiency, and instrument degradation evaluated at the + time bound with `at`. An unbound channel (including ``at(None)``) + never evaluates degradation — it is the pristine instrument. """ - return ( + area = ( self.geometrical_area * self.mirror_reflectance * self.filter_transmittance * self.effective_quantum_efficiency - * self.degradation(obstime=obstime) ) + if self._obstime is None: + return area + return area * self.degradation(obstime=self._obstime) @property @u.quantity_input diff --git a/sunkit_instruments/response/linelist.py b/sunkit_instruments/response/linelist.py new file mode 100644 index 00000000..a56243fc --- /dev/null +++ b/sunkit_instruments/response/linelist.py @@ -0,0 +1,463 @@ +""" +CHIANTI line lists with contribution functions (GOFNT). +""" + +import os +import logging +import pathlib +import warnings + +import numpy as np +import xarray as xr + +import astropy.constants as const +import astropy.units as u + +__all__ = [ + "LineListEmissionModel", + "chianti_line_list", + "get_line_list", + "line_list_cache_path", + "line_list_from_emission_model", +] + +log = logging.getLogger(__name__) + + +def chianti_line_list( + temperature: xr.DataArray, + density: xr.DataArray = None, + pressure: xr.DataArray = None, + abundance: str = None, + wavelength_range=None, + minimum_abundance: float = None, + element_list: list = None, + ion_list: list = None, +) -> xr.Dataset: + """ + Generate a line list with contribution functions (GOFNT) using ChiantiPy. + + Parameters + ---------- + temperature : `xarray.DataArray` + Temperature array in K with a ``logT`` dimension. + density : `xarray.DataArray`, optional + Electron density array in cm^-3. Mutually exclusive with ``pressure``. + pressure : `xarray.DataArray`, optional + Electron pressure array in K cm^-3. Mutually exclusive with ``density``. + abundance : `str`, optional + CHIANTI abundance name, e.g. ``"sun_coronal_2021_chianti"``. + wavelength_range : `tuple` + Two-element (min, max) wavelength range in Angstroms. + minimum_abundance : `float`, optional + Minimum elemental abundance to keep. + element_list : `list`, optional + List of elements, forwarded to `ChiantiPy.core.bunch`. + ion_list : `list`, optional + List of ions, forwarded to `ChiantiPy.core.bunch`. Ignored when + ``minimum_abundance`` is given. + + Returns + ------- + `xarray.Dataset` + Line list with GOFNT and per-transition metadata, restricted to + ``wavelength_range``. + + Notes + ----- + The ``XUVTOP`` environment variable must point at a local copy of the + CHIANTI database. ChiantiPy's ``gui`` default is forced to `False` for + the calling process, so no ``chiantirc`` file is needed on headless + systems. + """ + if (density is None) == (pressure is None): + msg = "Specify exactly one of density or pressure" + raise ValueError(msg) + if not isinstance(temperature, xr.DataArray): + msg = "temperature must be an xarray.DataArray with a logT dimension" + raise TypeError(msg) + wavelength_range = _validate_wavelength_range(wavelength_range) + if ion_list is not None and minimum_abundance is not None: + log.warning("minimum_abundance is set, the ion_list will be ignored") + + if "XUVTOP" not in os.environ: + msg = ( + "The XUVTOP environment variable is not set; ChiantiPy cannot locate the CHIANTI database. " + "Point it at a local copy of the database, e.g. `export XUVTOP=/path/to/chianti/dbase` " + "(available from https://www.chiantidatabase.org)." + ) + raise OSError(msg) + + with warnings.catch_warnings(): + # without a chiantirc file, ChiantiPy evaluates os.path.isfile(False) + # at import, which raises a RuntimeWarning on Python >= 3.14 + warnings.simplefilter("ignore", RuntimeWarning) + try: + import ChiantiPy.core as ch + import ChiantiPy.tools.data as chdata + except ImportError: + msg = "ChiantiPy is required for this function, install it with `pip install sunkit-instruments[chianti]`" + raise ImportError(msg) from None + + # never let ChiantiPy pop GUI selection dialogs (the rcfile default is + # gui=True, which hangs headless batch jobs) + chdata.Defaults["gui"] = False + + if density is not None: + temperature_bc, density_bc = xr.broadcast(temperature, density) + temperature_flat = temperature_bc.data.reshape(-1) + density_flat = density_bc.data.reshape(-1) + extra_coord_name = density.dims[0] + extra_coord = np.log10(density.data) + else: + density = pressure / temperature + temperature_bc = temperature.broadcast_like(density) + density_flat = density.data.reshape(-1) + temperature_flat = temperature_bc.data.reshape(-1) + extra_coord_name = pressure.dims[0] + extra_coord = pressure.data + + bunch = ch.bunch( + temperature_flat, + density_flat, + wavelength_range, + em=1.0, + abundance=abundance, + allLines=True, + keepIons=True, + minAbund=minimum_abundance, + ionList=ion_list, + elementList=element_list, + ) + + return _dataset_from_chianti_bunch( + bunch, + temperature, + temperature_bc, + {extra_coord_name: extra_coord}, + abundance, + wavelength_range, + ) + + +def line_list_cache_path( + output_dir: pathlib.Path, + abundance: str, + wavelength_range, + *, + density_dependent: bool = False, +) -> pathlib.Path: + """ + The canonical cache file path used by `get_line_list` for one abundance. + """ + wavelength_range = _validate_wavelength_range(wavelength_range) + prefix = "ll_wvl_eDens" if density_dependent else "ll_wvl" + lower, upper = (_format_wavelength_bound(bound) for bound in wavelength_range) + return pathlib.Path(output_dir) / f"{prefix}{lower}_{upper}_{abundance}.ncdf" + + +def get_line_list( + *, + output_dir: pathlib.Path, + abundance: str, + wavelength_range, + temperature: xr.DataArray, + pressure: xr.DataArray = None, + density: xr.DataArray = None, + minimum_abundance: float = None, + line_list_file: pathlib.Path = None, + compute_if_missing: bool = True, +) -> xr.Dataset: + """ + Load the cached CHIANTI line list for one abundance, computing and + caching it when absent. + + The cache file lives in ``output_dir`` and is named from the wavelength + range and abundance (``ll_wvl{lo}_{hi}_{abundance}.ncdf``, with an + ``_eDens`` prefix variant for density-dependent runs); the name is + re-derived per abundance so multiple abundances never share a cache. + Writes are atomic (tmp + replace) so concurrent job-array tasks racing + on the same missing cache cannot observe a half-written file. + + Parameters + ---------- + output_dir : `pathlib.Path` + Directory holding the cache files. + abundance : `str` + CHIANTI abundance name. + wavelength_range : array-like + Two-element (min, max) wavelength range in Angstroms. + temperature : `xarray.DataArray` + Temperature grid in K with a ``logT`` dimension. + pressure, density : `xarray.DataArray`, optional + Electron pressure grid, or density grid for density-dependent runs + (mutually exclusive; both forwarded to `chianti_line_list`). + minimum_abundance : `float`, optional + Minimum elemental abundance to keep. + line_list_file : `pathlib.Path`, optional + Explicit cache file to load, bypassing the derived name. + compute_if_missing : `bool`, optional + If `False`, raise instead of computing when the cache is absent + (useful to make job-array tasks fail fast when the preparation + step was skipped). + + Returns + ------- + `xarray.Dataset` + The line list. + """ + cache_path = ( + pathlib.Path(line_list_file) + if line_list_file is not None + else line_list_cache_path(output_dir, abundance, wavelength_range, density_dependent=density is not None) + ) + if line_list_file is not None or cache_path.exists(): + log.info(f"Loading line list from {cache_path}") + return _load_dataset(cache_path) + + if not compute_if_missing: + msg = f"line-list cache {cache_path} does not exist; run the line-list preparation step first" + raise FileNotFoundError(msg) + + log.info("Calculating line list") + line_list = chianti_line_list( + temperature=temperature, + pressure=pressure, + density=density, + abundance=abundance, + wavelength_range=wavelength_range, + minimum_abundance=minimum_abundance, + ) + + # write atomically so concurrent job-array tasks racing on the same + # missing cache cannot observe a half-written file + tmp_path = cache_path.with_name(f"{cache_path.name}.tmp{os.getpid()}") + _save_compressed_netcdf(tmp_path, line_list) + tmp_path.replace(cache_path) + return line_list + + +def _dataset_from_chianti_bunch(bunch, temperature, temperature_bc, extra_coords, abundance, wavelength_range): + import ChiantiPy + import ChiantiPy.tools.io as chio + + per_transition = { + "ion_name": bunch.Intensity["ionS"], + "wavelength": bunch.Intensity["wvl"], + "lower_level_label": bunch.Intensity["pretty1"], + "upper_level_label": bunch.Intensity["pretty2"], + "lower_level_index": bunch.Intensity["lvl1"], + "upper_level_index": bunch.Intensity["lvl2"], + "spectroscopic_name": np.array([bunch.IonInstances[ion].Spectroscopic for ion in bunch.Intensity["ionS"]]), + "atomic_number": np.array([bunch.IonInstances[ion].Z for ion in bunch.Intensity["ionS"]]), + "observed": bunch.Intensity["obs"] == "Y", + } + line_list = xr.Dataset({name: ("trans_index", values) for name, values in per_transition.items()}) + + gofnt_values = bunch.Intensity["intensity"].reshape((*temperature_bc.data.shape, 1, -1)) + line_list["gofnt"] = xr.DataArray( + gofnt_values, + dims=(*temperature_bc.dims, "abundance", "trans_index"), + coords={"logT": np.log10(temperature), **extra_coords, "abundance": np.array([abundance])}, + ) + line_list["logT_peak"] = np.log10(temperature[{"logT": line_list.gofnt.argmax(dim="logT")}]) + line_list["full_name"] = ( + line_list.spectroscopic_name.astype(object) + " " + line_list.wavelength.astype(str).astype(object) + ) + + line_list.attrs["Chiantipy"] = ChiantiPy.__version__ + line_list.attrs["Chianti"] = chio.versionRead() + line_list.gofnt.attrs["units"] = "erg cm3 / (s sr)" + + # cut down linelist to match wavelength range + in_range = (line_list.wavelength > wavelength_range[0]) & (line_list.wavelength < wavelength_range[1]) + return line_list.isel(trans_index=in_range) + + +def _save_compressed_netcdf(path: pathlib.Path, dataset: xr.Dataset) -> None: + encoding = dict.fromkeys(dataset.data_vars, {"zlib": True, "complevel": 5}) + # h5netcdf rather than netCDF4: the netCDF4 C bindings fail writing + # variable-length string variables under numpy >= 2.5 + dataset.to_netcdf(path, encoding=encoding, mode="w", engine="h5netcdf") + + +class LineListEmissionModel: + """ + A line-list dataset as an emission model for temperature responses. + + Wraps a line list (from `chianti_line_list`, `get_line_list` or + `line_list_from_emission_model`) so it conforms to the + `~sunkit_instruments.response.abstractions.EmissionModel` protocol and can + be passed to `~sunkit_instruments.response.get_temperature_response`. + + The response evaluates the channel's wavelength response at each line + centre (a delta-function approximation: no thermal broadening, accurate + to the variation of the wavelength response across a line profile) and + sums ``gofnt`` x response over all transitions. + + Parameters + ---------- + line_list : `xarray.Dataset` + Line list with ``wavelength``, ``gofnt`` and a ``logT`` coordinate. + Grid dimensions other than ``logT`` and ``trans_index`` (pressure, + density, abundance) must have size one — select a single entry first. + gofnt_unit : `str`, optional + Unit of ``gofnt``; by default read from its ``units`` attribute, + falling back to the ChiantiPy convention ``erg cm3 / (s sr)``. + """ + + def __init__(self, line_list: xr.Dataset, gofnt_unit: str = None): + for dim in set(line_list.gofnt.dims) - {"logT", "trans_index"}: + if line_list.sizes[dim] != 1: + msg = f"line_list must have a single {dim}; select one before building the model" + raise ValueError(msg) + line_list = line_list.isel({dim: 0}, drop=True) + self.line_list = line_list + self._gofnt_unit = u.Unit( + gofnt_unit if gofnt_unit is not None else line_list.gofnt.attrs.get("units", "erg cm3 / (s sr)") + ) + + @property + def temperature(self): + return 10 ** self.line_list.logT.data * u.K + + def calculate_temperature_response(self, channel): + """ + Temperature response of ``channel`` for this line list. + + ``channel`` is compatible with + `~sunkit_instruments.response.abstractions.AbstractChannel`. + """ + wave_response = channel.wavelength_response() + line_wavelength = u.Quantity(self.line_list.wavelength.data, u.AA) + response_at_lines = np.interp( + line_wavelength.to_value(channel.wavelength.unit), + channel.wavelength.value, + wave_response.value, + left=0.0, + right=0.0, + ) + # gofnt is an energy emissivity; the wavelength response is per photon + photons_per_erg = 1 / (const.h * const.c / line_wavelength).to_value(u.erg) + weight = xr.DataArray(response_at_lines * photons_per_erg, dims="trans_index") + response = (self.line_list.gofnt * weight).sum("trans_index") + unit = self._gofnt_unit * wave_response.unit * u.photon / u.erg + return u.Quantity(response.data, unit) + + +def line_list_from_emission_model( + model, + *, + include_ionization_fraction: bool = True, + abundance_label: str = "model", +) -> xr.Dataset: + """ + Build a line list from an emission model's per-transition emissivities. + + Converts the line emissivities of a + `~sunkit_instruments.response.abstractions.LineEmissionModel` (e.g. + `synthesizAR.atomic.EmissionModel`, backed by fiasco) into the line-list + dataset consumed by + `~sunkit_instruments.response.create_response_function`, matching the + `chianti_line_list` schema and its gofnt convention + (``erg cm3 s-1 sr-1``, equilibrium ionization fraction and abundance + included). + + Parameters + ---------- + model : `~sunkit_instruments.response.abstractions.LineEmissionModel` + Emission model exposing ``temperature``, ``density``, iteration over + ions, and ``get_line_emissivity(ion)`` returning per-transition + photon emissivities that exclude the ionization fraction (the + synthesizAR convention). + include_ionization_fraction : `bool`, optional + Multiply each ion's equilibrium ionization fraction into the gofnt, + by default `True`. Set to `False` to keep the model's convention of + deferring it (e.g. for time-dependent ionization). + abundance_label : `str`, optional + Label for the singleton ``abundance`` coordinate. + + Returns + ------- + `xarray.Dataset` + Line list with ``wavelength``, ``gofnt``, ``atomic_number``, + ``ion_name``, ``spectroscopic_name`` and ``full_name`` along + ``trans_index``, sorted by wavelength. + """ + gofnt_unit = u.Unit("erg cm3 s-1 sr-1") + wavelengths = [] + gofnts = [] + atomic_numbers = [] + ion_names = [] + spectroscopic_names = [] + for ion in model: + wavelength, emissivity = model.get_line_emissivity(ion) + if wavelength.size == 0 or not np.any(wavelength.value): + # placeholder entry for ions without level-population data + continue + if include_ionization_fraction: + emissivity = emissivity * ion.ionization_fraction[:, np.newaxis, np.newaxis] + # photon -> energy emissivity, per steradian (the ChiantiPy gofnt + # convention used throughout this subpackage) + photon_energy = (const.h * const.c / wavelength).to(u.erg) / u.photon + gofnt = (emissivity * photon_energy / (4 * np.pi * u.sr)).to_value(gofnt_unit) + n_transitions = wavelength.size + wavelengths.append(wavelength.to_value(u.AA)) + gofnts.append(gofnt) + atomic_numbers.append(np.full(n_transitions, ion.atomic_number)) + ion_names.append(np.full(n_transitions, getattr(ion, "ion_name", ""))) + spectroscopic_names.append(np.full(n_transitions, ion.ion_name_roman)) + + if not gofnts: + msg = "emission model produced no line emissivities" + raise ValueError(msg) + + wavelengths = np.concatenate(wavelengths) + # (logT, density, trans_index) -> (logT, density, abundance, trans_index) + gofnt = np.concatenate(gofnts, axis=-1)[..., np.newaxis, :] + atomic_numbers = np.concatenate(atomic_numbers) + ion_names = np.concatenate(ion_names) + spectroscopic_names = np.concatenate(spectroscopic_names) + order = np.argsort(wavelengths) + + line_list = xr.Dataset( + { + "wavelength": ("trans_index", wavelengths[order]), + "gofnt": (("logT", "density", "abundance", "trans_index"), gofnt[..., order]), + "atomic_number": ("trans_index", atomic_numbers[order]), + "ion_name": ("trans_index", ion_names[order]), + "spectroscopic_name": ("trans_index", spectroscopic_names[order]), + }, + coords={ + "logT": np.log10(model.temperature.to_value(u.K)), + "density": model.density.to_value(u.cm**-3), + "abundance": np.array([abundance_label]), + }, + ) + full_name = [ + f"{name} {wavelength:.3f}" + for name, wavelength in zip(line_list.spectroscopic_name.values, line_list.wavelength.values, strict=True) + ] + line_list["full_name"] = ("trans_index", np.array(full_name, dtype=object)) + line_list.gofnt.attrs["units"] = str(gofnt_unit) + return line_list + + +def _load_dataset(path: pathlib.Path) -> xr.Dataset: + # h5netcdf to match _save_compressed_netcdf (it also reads netCDF4-written + # HDF5 files); the netCDF4 C bindings misbehave under numpy >= 2.5 + with xr.open_dataset(path, engine="h5netcdf") as dataset: + return dataset.load() + + +def _validate_wavelength_range(wavelength_range): + try: + lower, upper = wavelength_range + except (TypeError, ValueError): + msg = "wavelength_range must contain exactly two values" + raise ValueError(msg) from None + return lower, upper + + +def _format_wavelength_bound(bound): + return f"{float(bound):g}" diff --git a/sunkit_instruments/response/spectral.py b/sunkit_instruments/response/spectral.py new file mode 100644 index 00000000..61fc0543 --- /dev/null +++ b/sunkit_instruments/response/spectral.py @@ -0,0 +1,465 @@ +""" +Spectral response functions built from CHIANTI line lists. +""" + +import numpy as np +import xarray as xr + +import astropy.constants as const +import astropy.units as u + +__all__ = ["create_response_function", "get_spectral_response"] + +_GAUSSIAN_EXPRESSION = "gofnt_scaled * exp(-0.5 * (shift / width)**2) / gaussian_norm / width" + + +def get_spectral_response(channel, model, *, vdop=None, wavelength_range=None, **kwargs): + """ + Spectrally-resolved response function of a channel for an emission model. + + The spectral counterpart of + `~sunkit_instruments.response.get_temperature_response`: instead of + integrating over wavelength, the emission is broadened by the thermal + (and instrumental) line width on a wavelength grid, optionally with a + Doppler-velocity axis — the form needed for spectrograph forward + modelling and inversions. + + Parameters + ---------- + channel : `~sunkit_instruments.response.abstractions.AbstractChannel` + Its ``wavelength`` and ``effective_area()`` define the instrument; + an ``instrumental_width`` attribute (Angstrom `~astropy.units.Quantity`) + is used as the instrumental line width when present. For channels + with time-dependent degradation, pass ``channel.at(obstime)``. + model : `~sunkit_instruments.response.abstractions.LineEmissionModel` or `xarray.Dataset` + Emission model converted via + `~sunkit_instruments.response.line_list_from_emission_model`, or an + already-built line-list dataset (e.g. from + `~sunkit_instruments.response.chianti_line_list`). + vdop : array-like, optional + Doppler-velocity axis in km/s, forwarded to + `create_response_function`. + wavelength_range : array-like, optional + Two-element (min, max) output wavelength range in Angstroms; by + default the channel's wavelength extent. + kwargs + Remaining keyword arguments forwarded to + `create_response_function` (``num_lines_keep``, ``window_sigma``, + ``normalization``, ...). + + Returns + ------- + `xarray.Dataset` + Response function from `create_response_function`. + """ + from sunkit_instruments.response.linelist import line_list_from_emission_model + + line_list = model if isinstance(model, xr.Dataset) else line_list_from_emission_model(model) + effective_area = xr.DataArray( + channel.effective_area().to_value(u.cm**2), + dims="wavelength", + coords={"wavelength": channel.wavelength.to_value(u.AA)}, + ) + if wavelength_range is None: + wavelength_range = [ + channel.wavelength.min().to_value(u.AA), + channel.wavelength.max().to_value(u.AA), + ] + instrumental_width = getattr(channel, "instrumental_width", 0 * u.AA) + return create_response_function( + line_list, + vdop=vdop, + instrumental_width=instrumental_width.to_value(u.AA), + effective_area=effective_area, + wavelength_range=wavelength_range, + **kwargs, + ) + + +def create_response_function( + line_list: xr.Dataset, + instrumental_width: float = 0, + normalization: float = 1e-27, + method: str = "linear", + vdop: u.Quantity | np.ndarray | list = None, + nonthermal_velocity: u.Quantity | np.ndarray | list = None, + wavelength_range: np.ndarray | list = None, + wavelength_step_mA: float = 4.9, + num_wavelength_bins: int = None, + effective_area: xr.DataArray = None, + num_lines_keep: int = 2, + band=None, + window_sigma: float = None, +) -> xr.Dataset: + """ + Computes a response function as a function of velocity and temperature. + + Parameters + ---------- + line_list : `xarray.Dataset` + Line list, can be created with `~sunkit_instruments.response.chianti_line_list`. + instrumental_width : `float`, optional + Instrumental width sigma in Angstroms, by default 0. May be a + `xarray.DataArray` carrying extra dims (e.g. a diffraction order + axis); the output inherits those dims. + normalization : `float`, optional + Normalization in the response function, by default 1e-27. + method : `str`, optional + Interpolation method for the effective area, by default "linear". + vdop : array-like, optional + Doppler axis array in km/s. + nonthermal_velocity : array-like, optional + Nonthermal velocity in km/s. + wavelength_range : array-like, optional + Two-element (min, max) wavelength range of the output grid in + Angstroms. Endpoints may be `xarray.DataArray` objects carrying + extra dims (requires ``num_wavelength_bins``). + wavelength_step_mA : `float`, optional + Wavelength bin size in milli-Angstroms, by default 4.9. + num_wavelength_bins : `int`, optional + Number of wavelength bins; overrides ``wavelength_step_mA``. + effective_area : `xarray.DataArray`, optional + Effective area with a ``wavelength`` dimension, interpolated onto + the output wavelength grid and multiplied into the response. + num_lines_keep : `int`, optional + Number of lines kept individually (in line-list order); the rest + are summed into a single "remaining lines" entry, by default 2. + band : `str`, optional + Bandpass label used for the summed-lines entry. If `None` it is + taken from a ``band`` coordinate of ``line_list`` when present. + window_sigma : `float`, optional + If set (e.g. 8-10), the summed contaminant response is built by + evaluating each line's Gaussian only within +/- ``window_sigma`` of + its (vdop-shifted) centre and scatter-adding it into a preallocated + full-grid accumulator, instead of computing and summing every line + across the entire wavelength grid. Same result to the chosen sigma + cutoff, dramatically faster for wide bands. `None` (default) keeps + the exact full-grid path. + + Returns + ------- + `xarray.Dataset` + Response function (data variable ``response``) with temperature, + velocity and wavelength axes. When the wavelength grid is + one-dimensional it is attached as the ``wavelength`` coordinate; + a grid carrying extra dims (e.g. order) is attached as the + ``wavelength_grid`` coordinate instead. + """ + if band is None and "band" in line_list.coords: + band = line_list.band.astype(str) + + try: + import periodictable as pt + except ImportError: + msg = "periodictable is required for this function, install it with `pip install sunkit-instruments[chianti]`" + raise ImportError(msg) from None + + try: + import numexpr as ne + except ImportError: + msg = "numexpr is required for this function, install it with `pip install sunkit-instruments[chianti]`" + raise ImportError(msg) from None + + CC_kms = const.c.to(u.km / u.s).value # Light speed (km/s) + CC_ms = const.c.to(u.m / u.s).value # Light speed (m/s) + KB = const.k_B.to(u.J / u.K).value # Boltzmann constant (J/K = kg m^2/K/s^2) + MP = const.m_p.to(u.kg).value # Proton mass (kg) + + # logic for single line response function + if "trans_index" not in line_list.dims: + line_list = line_list.expand_dims("trans_index") + _validate_line_list(line_list) + + wavelength_grid = _make_wavelength_grid(line_list, wavelength_range, wavelength_step_mA, num_wavelength_bins) + + if window_sigma is not None and wavelength_grid.ndim != 1: + msg = "window_sigma only supports one-dimensional wavelength grids" + raise ValueError(msg) + + if vdop is not None: + vdop = _velocity_axis(vdop, "vdop") + line_centers = line_list["wavelength"] * (1 + vdop / CC_kms) + else: + line_centers = line_list["wavelength"] + if nonthermal_velocity is not None: + nonthermal_velocity = _velocity_axis(nonthermal_velocity, "nonthermal_velocity") + + atomic_mass = _atomic_mass_from_atomic_number(line_list.atomic_number, pt.elements, MP) + + responses = [] + num_lines = line_list.sizes["trans_index"] + summed_lines_included = False + gaussian_norm = np.sqrt(2 * np.pi) + # windowed scatter-add state (only used when window_sigma is not None) + grid_values = np.asarray(wavelength_grid.values) + num_wavelengths = wavelength_grid.sizes["wavelength"] + contaminant_accumulator = None + contaminant_dims = None + contaminant_wavelength_axis = None + contaminant_coords = None + for i in range(num_lines): + line_center = line_centers.isel(trans_index=i) + + line_atomic_mass = atomic_mass.isel(trans_index=i) + thermal_velocity = np.sqrt(KB * 10**line_list.logT / line_atomic_mass) + thermal_line_width = line_center * thermal_velocity / CC_ms # (AA) sigma + + if nonthermal_velocity is not None: + doppler_width = np.sqrt( + thermal_line_width**2 + + instrumental_width**2 + + (line_list["wavelength"] * (nonthermal_velocity / CC_kms)) ** 2 + ) # (AA) + else: + doppler_width = np.sqrt(thermal_line_width**2 + instrumental_width**2) # (AA) + + if i < num_lines_keep: + # kept line, stored individually (few lines; full grid is cheap here) + line_response, gofnt_scaled = _evaluate_gaussian_response( + ne, + wavelength_grid, + line_center, + doppler_width, + line_list.gofnt.isel(trans_index=i), + normalization, + gaussian_norm, + ) + line_response = xr.DataArray(line_response, dims=gofnt_scaled.dims, coords=gofnt_scaled.coords) + line_response = line_response.expand_dims("line") + + line = line_list.full_name.isel(trans_index=i) + line_response = line_response.assign_coords({"line": line.expand_dims("line")}) + + line_response = line_response.assign_coords( + line_wavelength=line_list.isel(trans_index=i).wavelength.expand_dims("line") + ) + responses.append(line_response) + + elif window_sigma is not None: + # windowed scatter-add: evaluate this contaminant's Gaussian only + # within +/- window_sigma of its (vdop-shifted) centre and add it + # into a preallocated full-grid accumulator, skipping the ~zero + # wavelength bins that dominate the full-grid cost. + summed_lines_included = True + sigma_max = float(doppler_width.max()) + center_values = np.asarray(line_center.values, dtype=float) + i0 = max(0, int(np.searchsorted(grid_values, center_values.min() - window_sigma * sigma_max))) + i1 = min(num_wavelengths, int(np.searchsorted(grid_values, center_values.max() + window_sigma * sigma_max))) + if i0 < i1: + block, gofnt_scaled = _evaluate_gaussian_response( + ne, + wavelength_grid.isel(wavelength=slice(i0, i1)), + line_center, + doppler_width, + line_list.gofnt.isel(trans_index=i), + normalization, + gaussian_norm, + ) + if contaminant_accumulator is None: + contaminant_dims = gofnt_scaled.dims + contaminant_wavelength_axis = contaminant_dims.index("wavelength") + shape = list(block.shape) + shape[contaminant_wavelength_axis] = num_wavelengths + contaminant_accumulator = np.zeros(shape, dtype=block.dtype) + contaminant_coords = {d: gofnt_scaled.coords[d] for d in gofnt_scaled.coords} + elif gofnt_scaled.dims != contaminant_dims: + block = np.moveaxis( + block, [gofnt_scaled.dims.index(d) for d in contaminant_dims], range(len(contaminant_dims)) + ) + window = [slice(None)] * contaminant_accumulator.ndim + window[contaminant_wavelength_axis] = slice(i0, i1) + contaminant_accumulator[tuple(window)] += block + elif contaminant_accumulator is None: + gofnt_scaled = _broadcast_response_inputs( + wavelength_grid, + line_center, + doppler_width, + line_list.gofnt.isel(trans_index=i), + normalization, + )[0] + contaminant_dims = gofnt_scaled.dims + contaminant_wavelength_axis = contaminant_dims.index("wavelength") + contaminant_accumulator = np.zeros(gofnt_scaled.shape) + contaminant_coords = {d: gofnt_scaled.coords[d] for d in gofnt_scaled.coords} + + else: + # full-grid contaminant accumulation (original exact path) + contaminant_response, gofnt_scaled = _evaluate_gaussian_response( + ne, + wavelength_grid, + line_center, + doppler_width, + line_list.gofnt.isel(trans_index=i), + normalization, + gaussian_norm, + accumulator=contaminant_response if summed_lines_included else None, + ) + summed_lines_included = True + + if summed_lines_included: + if window_sigma is not None: + contaminant_response = xr.DataArray( + contaminant_accumulator, + dims=contaminant_dims, + coords=contaminant_coords, + ) + else: + contaminant_response = xr.DataArray( + contaminant_response, + dims=gofnt_scaled.dims, + coords=gofnt_scaled.coords, + ) + contaminant_response = contaminant_response.expand_dims("line") + # overwrite line coord with correct label + if band is None: + contaminant_label = xr.DataArray( + [f"Remaining {line_list.sizes['trans_index'] - num_lines_keep} lines"], dims="line" + ) + else: + contaminant_label = band + xr.DataArray( + [f" remaining {line_list.sizes['trans_index'] - num_lines_keep} lines"], dims="line" + ) + contaminant_response = contaminant_response.assign_coords(line=contaminant_label) + # assign the wvl of the first line in the list to summed lines. If this is the brightest line, it may well + # represent where peak intensity occurs + contaminant_response = contaminant_response.assign_coords( + {"line_wavelength": line_list.isel(trans_index=0).wavelength.expand_dims("line")} + ) + + # if all lines are summed, no need to broadcast against existing line label for concatenation + if num_lines_keep == 0: + contaminant_response["line"] = contaminant_label + else: + contaminant_response["line"] = contaminant_label.broadcast_like(line) + + if responses: + responses.append(contaminant_response) + else: + responses = contaminant_response + + responses = xr.concat(responses, dim="line", coords="different", compat="equals") + + ds = xr.Dataset() + ds["response"] = responses + response_attrs = ds.response.attrs + + if effective_area is not None: + interp = effective_area.interp(wavelength=wavelength_grid, method=method) + ds["response"] = ds.response * interp + ds.response.attrs.update(response_attrs) + ds.response.attrs["units"] = str(normalization * u.erg * u.cm**5 / u.s / u.sr / u.angstrom) + else: + ds.response.attrs["units"] = str(normalization * u.erg * u.cm**3 / u.s / u.sr / u.angstrom) + + # a 1-D grid becomes the wavelength index coordinate; a grid carrying + # extra dims (e.g. order) cannot share the dimension's name, so it is + # attached as wavelength_grid instead + if wavelength_grid.ndim == 1: + ds = ds.assign_coords(wavelength=("wavelength", wavelength_grid.data, {"units": str(u.AA)})) + else: + ds.coords["wavelength_grid"] = wavelength_grid + ds.coords["wavelength_grid"].attrs["units"] = str(u.AA) + + ds.attrs["normalization"] = normalization + + return ds + + +def _velocity_axis(values, dim): + if isinstance(values, u.Quantity): + values = values.to(u.km / u.s).value + values = np.atleast_1d(values) + return xr.DataArray(values, dims=dim, coords={dim: values}) + + +def _make_wavelength_grid(line_list, wavelength_range, wavelength_step_mA, num_wavelength_bins): + if wavelength_range is None: + # shortest and longest line wavelengths, with 1 Angstrom of padding + wavelength_range = [line_list.wavelength.min().data - 1, line_list.wavelength.max().data + 1] + + if num_wavelength_bins: + if isinstance(wavelength_range[0], xr.DataArray) or isinstance(wavelength_range[1], xr.DataArray): + # np.linspace cannot operate on xarray objects directly + # (https://github.com/pydata/xarray/issues/9043) + lower, upper = xr.broadcast(xr.DataArray(wavelength_range[0]), xr.DataArray(wavelength_range[1])) + return xr.DataArray( + np.linspace(lower.values, upper.values, num=num_wavelength_bins), + dims=("wavelength", *lower.dims), + coords=lower.coords, + ) + return xr.DataArray( + np.linspace(wavelength_range[0], wavelength_range[1], num=num_wavelength_bins), dims="wavelength" + ) + + if isinstance(wavelength_range[0], xr.DataArray) and wavelength_range[0].ndim > 0: + msg = ( + "wavelength_range endpoints with extra dims (e.g. order) require num_wavelength_bins; " + "np.arange only supports scalar endpoints" + ) + raise ValueError(msg) + return xr.DataArray( + np.arange(wavelength_range[0], wavelength_range[1] + wavelength_step_mA / 1e3, wavelength_step_mA / 1e3), # [A] + dims="wavelength", + ) + + +def _validate_line_list(line_list): + if not isinstance(line_list, xr.Dataset): + msg = "line_list must be an xarray.Dataset" + raise TypeError(msg) + missing = [name for name in ("wavelength", "atomic_number", "gofnt", "full_name") if name not in line_list] + if missing: + msg = f"line_list is missing required variables: {', '.join(missing)}" + raise ValueError(msg) + if "logT" not in line_list.coords: + msg = "line_list must include a logT coordinate" + raise ValueError(msg) + missing_trans_index = [ + name + for name in ("wavelength", "atomic_number", "gofnt", "full_name") + if "trans_index" not in line_list[name].dims + ] + if missing_trans_index: + msg = f"line_list variables must include a trans_index dimension: {', '.join(missing_trans_index)}" + raise ValueError(msg) + + +def _atomic_mass_from_atomic_number(atomic_number, elements, proton_mass): + atomic_mass = atomic_number.copy() + masses = np.array([elements[int(z)].mass for z in atomic_mass.data.reshape(-1)]) + atomic_mass.data = masses.reshape(atomic_mass.shape) * proton_mass + return atomic_mass + + +def _broadcast_response_inputs(wavelength_grid, line_center, doppler_width, gofnt, normalization): + shift = (wavelength_grid - line_center).broadcast_like(gofnt) + width, shift = xr.broadcast(doppler_width, shift) + gofnt_scaled = gofnt.broadcast_like(width) / normalization + return xr.broadcast(gofnt_scaled, width, shift) + + +def _evaluate_gaussian_response( + numexpr, + wavelength_grid, + line_center, + doppler_width, + gofnt, + normalization, + gaussian_norm, + *, + accumulator=None, +): + gofnt_scaled, width, shift = _broadcast_response_inputs( + wavelength_grid, line_center, doppler_width, gofnt, normalization + ) + local_dict = { + "gofnt_scaled": gofnt_scaled.data, + "shift": shift.data, + "width": width.data, + "gaussian_norm": gaussian_norm, + } + expression = _GAUSSIAN_EXPRESSION + if accumulator is not None: + local_dict["accumulator"] = accumulator + expression = f"{expression} + accumulator" + return numexpr.evaluate(expression, local_dict=local_dict), gofnt_scaled diff --git a/sunkit_instruments/response/tests/test_channel.py b/sunkit_instruments/response/tests/test_channel.py index 69a95f80..0c908928 100644 --- a/sunkit_instruments/response/tests/test_channel.py +++ b/sunkit_instruments/response/tests/test_channel.py @@ -3,8 +3,8 @@ import astropy.units as u -from sunkit_instruments.response import SourceSpectra -from sunkit_instruments.response.abstractions import AbstractChannel +from sunkit_instruments.response import SourceSpectra, get_temperature_response +from sunkit_instruments.response.abstractions import AbstractChannel, EmissionModel class TestChannel(AbstractChannel): @@ -87,6 +87,85 @@ def test_spectra_repr(fake_spectra): @pytest.mark.parametrize('obstime', [None, '2020-01-01']) def test_temperature_response(fake_channel, fake_spectra, obstime): - temp_response = fake_spectra.temperature_response(fake_channel, obstime=obstime) + channel = fake_channel if obstime is None else fake_channel.at(obstime) + temp_response = fake_spectra.temperature_response(channel) assert isinstance(temp_response, u.Quantity) assert temp_response.shape == fake_spectra.temperature.shape + + +class TimeDependentChannel(TestChannel): + @u.quantity_input + def degradation(self, obstime=None) -> u.dimensionless_unscaled: + return 2.0 if obstime is not None else 1.0 + + +class FakeEmissionModel: + temperature = np.array([1, 2]) * u.MK + + def calculate_temperature_response(self, channel): + return np.repeat(channel.wavelength_response().sum(), self.temperature.size) + + +def test_fake_emission_model_conforms_to_protocol(): + assert isinstance(FakeEmissionModel(), EmissionModel) + assert not isinstance(object(), EmissionModel) + + +def test_get_temperature_response_with_emission_model(): + temperature, temp_response = get_temperature_response( + TimeDependentChannel(), FakeEmissionModel() + ) + assert isinstance(temperature, u.Quantity) + assert isinstance(temp_response, u.Quantity) + assert temp_response.shape == temperature.shape + + +def test_get_temperature_response_with_emission_model_obstime(): + channel = TimeDependentChannel() + model = FakeEmissionModel() + _, response_unbound = get_temperature_response(channel, model) + _, response_bound = get_temperature_response(channel.at("2020-01-01"), model) + assert np.allclose(response_bound.value, 2 * response_unbound.value) + + +def test_get_temperature_response_with_unsupported_source(fake_channel): + with pytest.raises(TypeError, match="temperature_response"): + get_temperature_response(fake_channel, object()) + + +def test_at_binds_obstime(): + channel = TimeDependentChannel() + bound = channel.at("2020-01-01") + # TimeDependentChannel degrades by exactly 2x for any bound time (atol + # floor: the Gaussian filter tail underflows to subnormals where the + # factor-of-two identity no longer holds bit-exactly) + wave_response = channel.wavelength_response() + assert u.allclose(bound.wavelength_response(), 2 * wave_response, atol=1e-20 * wave_response.unit) + area = channel.effective_area() + assert u.allclose(bound.effective_area(), 2 * area, atol=1e-20 * area.unit) + # binding returns a copy; the original stays unbound + assert channel._obstime is None + assert u.allclose(bound.wavelength, channel.wavelength) + + +def test_at_view_reusable_with_emission_model(): + bound = TimeDependentChannel().at("2020-01-01") + _, first = get_temperature_response(bound, FakeEmissionModel()) + _, second = get_temperature_response(bound, FakeEmissionModel()) + assert u.allclose(first, second) + + +class DegradationRequiresTimeChannel(TestChannel): + @u.quantity_input + def degradation(self, obstime=None) -> u.dimensionless_unscaled: + if obstime is None: + raise ValueError("obstime required") + return 0.5 + + +def test_unbound_channel_never_evaluates_degradation(): + channel = DegradationRequiresTimeChannel() + # pristine instrument: degradation not called at all + pristine = channel.effective_area() + assert u.allclose(channel.at(None).effective_area(), pristine) + assert u.allclose(channel.at("2020-01-01").effective_area(), 0.5 * pristine) diff --git a/sunkit_instruments/response/tests/test_emission_model.py b/sunkit_instruments/response/tests/test_emission_model.py new file mode 100644 index 00000000..fb6f5ed5 --- /dev/null +++ b/sunkit_instruments/response/tests/test_emission_model.py @@ -0,0 +1,179 @@ +""" +Tests for the emission-model line-list adapter and get_spectral_response, +using a fake model that mimics the synthesizAR EmissionModel surface. +No fiasco or synthesizAR required. +""" + +import numpy as np +import pytest +import xarray as xr + +import astropy.constants as const +import astropy.units as u + +from sunkit_instruments.response import ( + LineListEmissionModel, + get_spectral_response, + get_temperature_response, + line_list_from_emission_model, +) +from sunkit_instruments.response.abstractions import EmissionModel, LineEmissionModel +from sunkit_instruments.response.tests.test_channel import TestChannel + +TEMPERATURE = np.array([5e5, 1e6, 2e6]) * u.K +DENSITY = np.array([1e9]) * u.cm**-3 + + +class FakeIon: + def __init__(self, ion_name, ion_name_roman, atomic_number, ionization_fraction): + self.ion_name = ion_name + self.ion_name_roman = ion_name_roman + self.atomic_number = atomic_number + self.ionization_fraction = ionization_fraction + + +class FakeLineEmissionModel: + temperature = TEMPERATURE + density = DENSITY + + def __init__(self): + self.fe9 = FakeIon("fe_9", "Fe IX", 26, np.array([0.1, 0.8, 0.2])) + self.missing = FakeIon("na_3", "Na III", 11, np.ones(3)) + self._ions = [self.fe9, self.missing] + self.emissivity = u.Quantity(np.ones((3, 1, 2)), "cm3 ph s-1") + self.wavelengths = np.array([174.5, 171.073]) * u.AA + + def __iter__(self): + return iter(self._ions) + + def get_line_emissivity(self, ion, transition=None): + if ion is self.missing: + # placeholder for ions without level-population data + return [0] * u.AA, u.Quantity(np.zeros((3, 1, 1)), "cm3 ph s-1") + return self.wavelengths, self.emissivity + + +def test_fake_model_conforms_to_protocol(): + assert isinstance(FakeLineEmissionModel(), LineEmissionModel) + + +def test_adapter_schema_and_sorting(): + line_list = line_list_from_emission_model(FakeLineEmissionModel()) + assert line_list.sizes["trans_index"] == 2 # missing-data ion skipped + np.testing.assert_allclose(line_list.wavelength.values, [171.073, 174.5]) # sorted + assert list(line_list.full_name.values) == ["Fe IX 171.073", "Fe IX 174.500"] + assert list(line_list.atomic_number.values) == [26, 26] + assert line_list.gofnt.dims == ("logT", "density", "abundance", "trans_index") + np.testing.assert_allclose(line_list.logT.values, np.log10(TEMPERATURE.to_value(u.K))) + assert u.Unit(line_list.gofnt.attrs["units"]) == u.Unit("erg cm3 / (s sr)") + + +def test_adapter_gofnt_values(): + """gofnt = emissivity x ionization fraction x photon energy / 4 pi sr.""" + model = FakeLineEmissionModel() + line_list = line_list_from_emission_model(model) + photon_energy = (const.h * const.c / (171.073 * u.AA)).to_value(u.erg) + expected = 1.0 * model.fe9.ionization_fraction * photon_energy / (4 * np.pi) + gofnt_171 = line_list.gofnt.sel(trans_index=line_list.wavelength == 171.073).squeeze() + np.testing.assert_allclose(gofnt_171.values, expected, rtol=1e-12) + + +def test_adapter_can_defer_ionization_fraction(): + model = FakeLineEmissionModel() + with_ioneq = line_list_from_emission_model(model) + without_ioneq = line_list_from_emission_model(model, include_ionization_fraction=False) + ratio = (with_ioneq.gofnt / without_ioneq.gofnt).isel(density=0, abundance=0, trans_index=0) + np.testing.assert_allclose(ratio.values, model.fe9.ionization_fraction) + + +def test_adapter_empty_model_raises(): + model = FakeLineEmissionModel() + model._ions = [model.missing] + with pytest.raises(ValueError, match="no line emissivities"): + line_list_from_emission_model(model) + + +def test_get_spectral_response_from_model(): + channel = TestChannel() + response = get_spectral_response( + channel, + FakeLineEmissionModel(), + vdop=np.array([-100.0, 0.0, 100.0]) * u.km / u.s, + wavelength_range=[170.0, 176.0], + num_lines_keep=1, + ) + assert {"line", "logT", "vdop", "density", "abundance", "wavelength"} <= set(response.response.dims) + lines = [str(line) for line in response.line.values] + assert lines[0] == "Fe IX 171.073" + assert any("remaining" in line.lower() for line in lines[1:]) + assert "erg" in response.response.attrs["units"] + + +def test_get_spectral_response_accepts_line_list_dataset(): + channel = TestChannel() + line_list = line_list_from_emission_model(FakeLineEmissionModel()) + response = get_spectral_response( + channel, + line_list, + wavelength_range=[170.0, 176.0], + num_lines_keep=2, + ) + assert response.sizes["line"] == 2 + + +@pytest.fixture +def chianti_style_line_list(): + logT = np.array([5.8, 6.0, 6.2]) + gofnt = np.array([1.0, 5.0, 2.0]) * 1e-25 + line_list = xr.Dataset( + { + "wavelength": ("trans_index", np.array([171.073])), + "atomic_number": ("trans_index", np.array([26])), + "gofnt": ( + ("logT", "pressure", "abundance", "trans_index"), + gofnt[:, np.newaxis, np.newaxis, np.newaxis], + ), + "full_name": ("trans_index", np.array(["Fe IX 171.073"], dtype=object)), + }, + coords={"logT": logT, "pressure": [3e15], "abundance": ["fake"]}, + ) + line_list.gofnt.attrs["units"] = "erg cm3 / (s sr)" + return line_list + + +def test_line_list_emission_model_conforms_to_protocol(chianti_style_line_list): + assert isinstance(LineListEmissionModel(chianti_style_line_list), EmissionModel) + + +def test_line_list_emission_model_temperature_response(chianti_style_line_list): + channel = TestChannel() + model = LineListEmissionModel(chianti_style_line_list) + temperature, response = get_temperature_response(channel, model) + + assert u.allclose(temperature, 10 ** chianti_style_line_list.logT.data * u.K) + assert response.unit.physical_type == u.Unit("cm5 DN / (pix s)").physical_type + # single line at 171.073: K(T) = gofnt(T) * R(171.073) * photons/erg + wave_response = channel.wavelength_response() + r_at_line = np.interp(171.073, channel.wavelength.to_value(u.AA), wave_response.value) + photons_per_erg = 1 / (const.h * const.c / (171.073 * u.AA)).to_value(u.erg) + expected = chianti_style_line_list.gofnt.values.squeeze() * r_at_line * photons_per_erg + np.testing.assert_allclose(response.to_value("cm5 DN / (pix s)"), expected, rtol=1e-12) + + +def test_line_list_emission_model_zero_outside_channel(chianti_style_line_list): + line_list = chianti_style_line_list.copy() + line_list["wavelength"] = ("trans_index", np.array([500.0])) # outside 100-200 A channel + _, response = get_temperature_response(TestChannel(), LineListEmissionModel(line_list)) + np.testing.assert_array_equal(response.value, 0.0) + + +def test_line_list_emission_model_requires_single_grid(): + line_list = xr.Dataset( + { + "wavelength": ("trans_index", np.array([171.073])), + "gofnt": (("logT", "pressure", "trans_index"), np.ones((2, 2, 1))), + }, + coords={"logT": [5.8, 6.0], "pressure": [3e15, 3e16]}, + ) + with pytest.raises(ValueError, match="single pressure"): + LineListEmissionModel(line_list) diff --git a/sunkit_instruments/response/tests/test_linelist.py b/sunkit_instruments/response/tests/test_linelist.py new file mode 100644 index 00000000..ea4cfc96 --- /dev/null +++ b/sunkit_instruments/response/tests/test_linelist.py @@ -0,0 +1,129 @@ +import os + +import numpy as np +import pytest +import xarray as xr + +from sunkit_instruments.response.linelist import ( + _save_compressed_netcdf, + chianti_line_list, + get_line_list, + line_list_cache_path, +) + + +@pytest.fixture +def fake_linelist(): + return xr.Dataset( + { + "wavelength": xr.DataArray([171.073, 174.531], dims="trans_index"), + "gofnt": xr.DataArray( + np.ones((3, 1, 1, 2)), + dims=("logT", "pressure", "abundance", "trans_index"), + coords={"logT": [5.8, 5.9, 6.0], "pressure": [3e15]}, + ), + "full_name": xr.DataArray( + np.array(["Fe IX 171.073", "Fe X 174.531"], dtype=object), dims="trans_index" + ), + } + ) + + +def test_cache_path_naming(tmp_path): + path = line_list_cache_path(tmp_path, "sun_coronal_2021_chianti", (100.5, 110.9)) + assert path == tmp_path / "ll_wvl100.5_110.9_sun_coronal_2021_chianti.ncdf" + path = line_list_cache_path(tmp_path, "abund", (100, 110), density_dependent=True) + assert path.name == "ll_wvl_eDens100_110_abund.ncdf" + + +def test_get_line_list_loads_cache(tmp_path, fake_linelist): + cache_path = line_list_cache_path(tmp_path, "abund", (170, 175)) + _save_compressed_netcdf(cache_path, fake_linelist) + loaded = get_line_list( + output_dir=tmp_path, + abundance="abund", + wavelength_range=(170, 175), + temperature=10 ** fake_linelist.logT, + pressure=fake_linelist.pressure, + compute_if_missing=False, + ) + xr.testing.assert_allclose(loaded.gofnt, fake_linelist.gofnt) + assert list(loaded.full_name.values) == list(fake_linelist.full_name.values) + + +def test_get_line_list_explicit_path(tmp_path, fake_linelist): + explicit = tmp_path / "custom_name.ncdf" + _save_compressed_netcdf(explicit, fake_linelist) + loaded = get_line_list( + output_dir=tmp_path, + abundance="ignored", + wavelength_range=(0, 1), + temperature=10 ** fake_linelist.logT, + line_list_file=explicit, + ) + xr.testing.assert_allclose(loaded.gofnt, fake_linelist.gofnt) + + +def test_get_line_list_missing_raises(tmp_path, fake_linelist): + with pytest.raises(FileNotFoundError, match="line-list cache"): + get_line_list( + output_dir=tmp_path, + abundance="abund", + wavelength_range=(170, 175), + temperature=10 ** fake_linelist.logT, + pressure=fake_linelist.pressure, + compute_if_missing=False, + ) + + +def test_rejects_both_density_and_pressure(): + temperature = xr.DataArray([1e6], dims="logT") + grid = xr.DataArray([1e9], dims="density") + with pytest.raises(ValueError, match="exactly one"): + chianti_line_list(temperature, density=grid, pressure=grid) + with pytest.raises(ValueError, match="exactly one"): + chianti_line_list(temperature) + + +def test_rejects_missing_wavelength_range(): + temperature = xr.DataArray([1e6], dims="logT") + pressure = xr.DataArray([3e15], dims="pressure") + with pytest.raises(ValueError, match="wavelength_range"): + chianti_line_list(temperature, pressure=pressure) + + +def test_missing_xuvtop_raises(monkeypatch): + monkeypatch.delenv("XUVTOP", raising=False) + temperature = xr.DataArray([1e6], dims="logT") + pressure = xr.DataArray([3e15], dims="pressure") + with pytest.raises(OSError, match="XUVTOP"): + chianti_line_list(temperature, pressure=pressure, wavelength_range=(170, 172)) + + +@pytest.mark.skipif(os.environ.get("XUVTOP") is None, reason="CHIANTI database (XUVTOP) not available") +def test_line_list_live(tmp_path): + pytest.importorskip("ChiantiPy") + temperature = 10 ** xr.DataArray(np.arange(5.6, 6.2, 0.2), dims="logT") + pressure = xr.DataArray([3e15], dims="pressure") + line_list = get_line_list( + output_dir=tmp_path, + abundance="sun_coronal_2021_chianti", + wavelength_range=(170, 172), + temperature=temperature, + pressure=pressure, + minimum_abundance=1e-5, + ) + assert line_list.sizes["trans_index"] > 0 + assert "Fe IX 171.073" in line_list.full_name.values + assert (line_list.wavelength > 170).all() and (line_list.wavelength < 172).all() + assert {"ion_name", "atomic_number", "spectroscopic_name", "logT_peak"} <= set(line_list.data_vars) + # cache round-trip: second call loads the file written by the first + cached = get_line_list( + output_dir=tmp_path, + abundance="sun_coronal_2021_chianti", + wavelength_range=(170, 172), + temperature=temperature, + pressure=pressure, + compute_if_missing=False, + ) + xr.testing.assert_allclose(cached.gofnt, line_list.gofnt) diff --git a/sunkit_instruments/response/tests/test_spectral.py b/sunkit_instruments/response/tests/test_spectral.py new file mode 100644 index 00000000..cc0ee88d --- /dev/null +++ b/sunkit_instruments/response/tests/test_spectral.py @@ -0,0 +1,226 @@ +""" +Tests for ``create_response_function``: response integrals, Doppler shifts, +xarray-broadcasting flexibility (parameters carrying extra dims such as a +diffraction order axis), and the windowed scatter-add contaminant path. +None of these tests require ChiantiPy or the CHIANTI database. +""" + +import numpy as np +import pytest +import xarray as xr + +import astropy.units as u + +from sunkit_instruments.response.spectral import create_response_function + +NORM = 1e-27 +ORDERS = xr.DataArray([1, 2], dims="order", coords={"order": [1, 2]}) +VDOP = np.array([-200.0, 0.0, 200.0]) * u.km / u.s + +# a CI-like spread of lines across a wide band, on a finer logT grid, to +# exercise the windowed scatter-add over many wavelength bins +WIDE_BAND_LINES = {"wavelength": np.linspace(105.0, 195.0, 30), "logT": np.linspace(4.6, 7.6, 15)} + + +def synthetic_line_list(n_lines=2, wavelength=None, logT=None): + """ + Minimal deterministic line list with the fields ``create_response_function`` + consumes: iron lines (near 171 A unless ``wavelength`` is given) with + Gaussian-in-logT contribution functions. + """ + wavelength = np.linspace(170.6, 171.4, n_lines) if wavelength is None else np.asarray(wavelength, dtype=float) + n_lines = wavelength.size + logT = np.array([5.8, 6.0, 6.2]) if logT is None else np.asarray(logT, dtype=float) + peaks = np.linspace(1.0, 0.5, n_lines) + gofnt = peaks[np.newaxis, :] * np.exp(-((logT[:, np.newaxis] - 6.0) ** 2) / 0.02) * 1e-25 + return xr.Dataset( + { + "wavelength": ("trans_index", wavelength), + "atomic_number": ("trans_index", np.full(n_lines, 26)), + "gofnt": (("logT", "trans_index"), gofnt), + "full_name": ("trans_index", [f"Fake Fe {i} {w:.3f}" for i, w in enumerate(wavelength)]), + }, + coords={"logT": logT}, + ) + + +class TestCreateResponseFunctionScalar: + def test_integral_matches_gofnt(self): + """ + Each per-line response is a normalized Gaussian in wavelength, so + integrating over wavelength at vdop=0 must recover gofnt/normalization. + """ + ll = synthetic_line_list(1) + resp = create_response_function( + ll, vdop=VDOP, instrumental_width=0.02, wavelength_range=[170.0, 172.0], wavelength_step_mA=2.0, num_lines_keep=1 + ) + dlam = float(resp.wavelength[1] - resp.wavelength[0]) + integral = (resp.response.sel(vdop=0).isel(line=0) * dlam).sum("wavelength") + expected = ll.gofnt.isel(trans_index=0) / NORM + np.testing.assert_allclose(integral.values, expected.values, rtol=1e-3) + + def test_peak_follows_doppler_shift(self): + ll = synthetic_line_list(1) + resp = create_response_function( + ll, vdop=VDOP, instrumental_width=0.02, wavelength_range=[170.0, 172.0], wavelength_step_mA=2.0, num_lines_keep=1 + ) + c_kms = 299792.458 + for v in (-200.0, 200.0): + peak_wvl = float(resp.wavelength[resp.response.sel(vdop=v).isel(line=0, logT=1).argmax("wavelength")]) + expected = float(ll.wavelength[0]) * (1 + v / c_kms) + assert abs(peak_wvl - expected) < 3e-3 + + def test_scalar_velocity_inputs(self): + ll = synthetic_line_list(1) + resp = create_response_function( + ll, + vdop=0.0 * u.km / u.s, + nonthermal_velocity=0.0, + instrumental_width=0.02, + wavelength_range=[170.0, 172.0], + wavelength_step_mA=2.0, + num_lines_keep=1, + ) + assert resp.response.sizes["vdop"] == 1 + assert resp.response.sizes["nonthermal_velocity"] == 1 + + def test_missing_line_list_fields_raises(self): + ll = synthetic_line_list(1).drop_vars("atomic_number") + with pytest.raises(ValueError, match="atomic_number"): + create_response_function(ll, wavelength_range=[170.0, 172.0], num_lines_keep=1) + + +class TestCreateResponseFunctionOrderDims: + def test_instrumental_width_order_dim(self): + ll = synthetic_line_list(2) + width = xr.DataArray([0.01, 0.05], dims="order", coords={"order": [1, 2]}) + resp = create_response_function( + ll, vdop=VDOP, instrumental_width=width, wavelength_range=[170.0, 172.0], wavelength_step_mA=2.0, num_lines_keep=2 + ) + assert "order" in resp.response.dims + assert resp.response.sizes["order"] == 2 + assert not np.allclose( + resp.response.isel(order=0).values, + resp.response.isel(order=1).values, + ) + + def test_equal_widths_give_identical_slices(self): + ll = synthetic_line_list(1) + width = xr.DataArray([0.02, 0.02], dims="order", coords={"order": [1, 2]}) + resp = create_response_function( + ll, vdop=VDOP, instrumental_width=width, wavelength_range=[170.0, 172.0], wavelength_step_mA=2.0, num_lines_keep=1 + ) + np.testing.assert_array_equal( + resp.response.isel(order=0).values, + resp.response.isel(order=1).values, + ) + + def test_wavelength_range_order_dependent(self): + ll = synthetic_line_list(2) + wavelength_range = [171.0 * 2 / ORDERS - 1.0, 171.0 * 2 / ORDERS + 1.0] + resp = create_response_function( + ll, vdop=VDOP, instrumental_width=0.02, wavelength_range=wavelength_range, num_wavelength_bins=64, num_lines_keep=2 + ) + assert set(resp.wavelength_grid.dims) == {"wavelength", "order"} + assert resp.response.sizes["wavelength"] == 64 + + def test_wavelength_range_order_dependent_requires_num_bins(self): + ll = synthetic_line_list(1) + wavelength_range = [171.0 * 2 / ORDERS - 1.0, 171.0 * 2 / ORDERS + 1.0] + with pytest.raises(ValueError, match="num_wavelength_bins"): + create_response_function( + ll, vdop=VDOP, instrumental_width=0.02, wavelength_range=wavelength_range, num_lines_keep=1 + ) + + def test_window_requires_one_dimensional_wavelength_grid(self): + ll = synthetic_line_list(3) + wavelength_range = [171.0 * 2 / ORDERS - 1.0, 171.0 * 2 / ORDERS + 1.0] + with pytest.raises(ValueError, match="one-dimensional"): + create_response_function( + ll, + wavelength_range=wavelength_range, + num_wavelength_bins=64, + num_lines_keep=0, + window_sigma=8.0, + ) + + def test_contam_sum_with_order_dim(self): + """num_lines_keep=0 sums every line; the order axis must survive.""" + ll = synthetic_line_list(3) + width = xr.DataArray([0.01, 0.05], dims="order", coords={"order": [1, 2]}) + resp = create_response_function( + ll, vdop=VDOP, instrumental_width=width, wavelength_range=[170.0, 172.0], wavelength_step_mA=2.0, num_lines_keep=0 + ) + assert "order" in resp.response.dims + assert resp.response.sizes["line"] == 1 + + def test_effective_area_band_order(self): + ll = synthetic_line_list(2) + wavelength = np.linspace(168.0, 174.0, 40) + ea = xr.DataArray( + np.full((40, 1), 10.0), + dims=("wavelength", "band"), + coords={"wavelength": wavelength, "band": [171]}, + ) + ea = ea * xr.DataArray([1.0, 0.5], dims="order", coords={"order": [1, 2]}) + resp = create_response_function( + ll, + vdop=VDOP, + instrumental_width=0.02, + wavelength_range=[170.0, 172.0], + wavelength_step_mA=2.0, + num_lines_keep=2, + effective_area=ea, + ) + assert {"band", "order", "line", "logT", "vdop", "wavelength"} <= set(resp.response.dims) + # the order axis came in through the effective area, so the two + # slices must differ by exactly its ratio + np.testing.assert_allclose( + resp.response.sel(order=2).values, + 0.5 * resp.response.sel(order=1).values, + rtol=1e-12, + ) + + +class TestWindowedContaminants: + """The windowed scatter-add path (``window_sigma``) must reproduce the + exact full-grid result.""" + + def test_window_matches_full_grid_all_contaminants(self): + """Wide-band case: num_lines_keep=0, every line summed.""" + ll = synthetic_line_list(**WIDE_BAND_LINES) + kw = { + "vdop": np.arange(-200.0, 210.0, 40.0), + "instrumental_width": 0.0, + "wavelength_range": [100.0, 200.0], + "num_lines_keep": 0, + } + full = create_response_function(ll, **kw) + win = create_response_function(ll, window_sigma=8.0, **kw) + peak = float(np.abs(full["response"]).max()) + xr.testing.assert_allclose(full["response"], win["response"], atol=peak * 1e-9) + + def test_window_with_kept_lines(self): + """num_lines_keep>0: kept lines stay full-grid, contaminants windowed.""" + ll = synthetic_line_list(**WIDE_BAND_LINES) + kw = { + "vdop": np.arange(-100.0, 110.0, 50.0), + "instrumental_width": 0.01, + "wavelength_range": [100.0, 200.0], + "num_lines_keep": 2, + } + full = create_response_function(ll, **kw) + win = create_response_function(ll, window_sigma=8.0, **kw) + peak = float(np.abs(full["response"]).max()) + xr.testing.assert_allclose(full["response"], win["response"], atol=peak * 1e-9) + + def test_window_off_grid_contaminants_are_zero(self): + ll = synthetic_line_list(1, wavelength=[500.0]) + resp = create_response_function( + ll, + wavelength_range=[100.0, 101.0], + num_lines_keep=0, + window_sigma=8.0, + ) + assert resp.response.dtype.kind != "O" + assert not resp.response.any() diff --git a/sunkit_instruments/response/thermal.py b/sunkit_instruments/response/thermal.py index 7de91ee5..56bf63a0 100644 --- a/sunkit_instruments/response/thermal.py +++ b/sunkit_instruments/response/thermal.py @@ -5,19 +5,29 @@ import astropy.units as u +from sunkit_instruments.response.abstractions import EmissionModel + __all__ = ["SourceSpectra", "get_temperature_response"] -def get_temperature_response(channel, spectra, obstime=None): +def get_temperature_response(channel, spectra): """ Calculate the temperature response function for a given instrument channel - and input spectra. + and input source. Parameters ---------- channel: `~sunkit_instruments.response.abstractions.AbstractChannel` - spectra: `~sunkit_instruments.response.SourceSpectra` - obstime: any format parsed by `sunpy.time.parse_time` , optional + For time-dependent instrument degradation, bind the time first: + ``get_temperature_response(channel.at(obstime), spectra)``. + spectra: `~sunkit_instruments.response.SourceSpectra` or emission model + The source used to calculate the temperature response. In addition to + `~sunkit_instruments.response.SourceSpectra`, this can be any + emission model conforming to + `~sunkit_instruments.response.abstractions.EmissionModel` (a + ``temperature`` attribute and a + ``calculate_temperature_response(channel)`` method), e.g. + `synthesizAR.atomic.EmissionModel`. Returns ------- @@ -27,8 +37,15 @@ def get_temperature_response(channel, spectra, obstime=None): See Also -------- sunkit_instruments.response.SourceSpectra.temperature_response + sunkit_instruments.response.abstractions.EmissionModel """ - return spectra.temperature, spectra.temperature_response(channel, obstime=obstime) + if hasattr(spectra, "temperature_response"): + return spectra.temperature, spectra.temperature_response(channel) + if isinstance(spectra, EmissionModel): + return spectra.temperature, spectra.calculate_temperature_response(channel) + raise TypeError( + "spectra must define either temperature_response(channel) or calculate_temperature_response(channel)" + ) class SourceSpectra: @@ -136,9 +153,7 @@ def data(self) -> u.photon * u.cm**3 / (u.s * u.Angstrom * u.steradian): return u.Quantity(self._da.data, self._da.attrs["unit"]) @u.quantity_input - def temperature_response( - self, channel, obstime=None - ) -> u.cm**5 * u.DN / (u.pixel * u.s): + def temperature_response(self, channel) -> u.cm**5 * u.DN / (u.pixel * u.s): """ Temperature response function for a given instrument channel. @@ -151,12 +166,10 @@ def temperature_response( ---------- channel: `~sunkit_instruments.response.abstractions.AbstractChannel` The relevant instrument channel object used to compute the wavelength - response function. - obstime: any format parsed by `sunpy.time.parse_time`, optional - A time of a particular observation. This is used to calculated any - time-dependent instrument degradation. + response function. For time-dependent instrument degradation, + bind the time first with ``channel.at(obstime)``. """ - wave_response = channel.wavelength_response(obstime=obstime) + wave_response = channel.wavelength_response() da_response = xarray.DataArray( wave_response.to_value(wave_response.unit), dims=['wavelength'],