diff --git a/ctis/__init__.py b/ctis/__init__.py index 5e829ac..b923c8f 100644 --- a/ctis/__init__.py +++ b/ctis/__init__.py @@ -5,6 +5,7 @@ from ._arange import arange from ._regrid import regrid +from ._torch import Regridder from . import scenes from . import instruments from . import inverters @@ -12,6 +13,7 @@ __all__ = [ "arange", "regrid", + "Regridder", "scenes", "instruments", "inverters", diff --git a/ctis/_torch.py b/ctis/_torch.py new file mode 100644 index 0000000..f8a4694 --- /dev/null +++ b/ctis/_torch.py @@ -0,0 +1,345 @@ +""" +A :mod:`torch` backend for the CTIS forward model. + +The forward model of a :class:`~ctis.instruments.AbstractLinearInstrument` is a +sparse matrix multiplication, which :mod:`regridding` stores as a collection of +``(indices_input, indices_output, values)`` triplets. +:class:`Regridder` assembles those triplets into a single sparse CSR matrix +which can be applied on a GPU and differentiated by :mod:`torch`. + +Notes +----- +The compressed sparse row (CSR) format is used instead of the more obvious +scatter-add (:meth:`torch.Tensor.index_add_`) since it is both faster and, +unlike scatter-add, bitwise deterministic on CUDA devices. +Determinism matters because the inversion routines built on this module solve +linear systems using conjugate gradient methods, which assume the operator +does not change between applications. + +Since the adjoint of this operator is computed by :mod:`torch` automatic +differentiation, it is the exact transpose of the forward model. +This is *not* the same as +:meth:`~ctis.instruments.AbstractInstrument.backproject`, which applies an +additional normalization to conserve flux. +""" + +from typing import TYPE_CHECKING +import dataclasses +import numpy as np +import astropy.units as u +import named_arrays as na + +__all__ = [ + "Regridder", +] + +if TYPE_CHECKING: # pragma: nocover + import torch + + +def _torch(): + """ + Import :mod:`torch` lazily so that it remains an optional dependency. + """ + try: + import torch + except ImportError as e: # pragma: nocover + raise ImportError( + "PyTorch is required to use `ctis._torch`. " + "Install it using `pip install ctis[torch]`." + ) from e + return torch + + +@dataclasses.dataclass(eq=False) +class Regridder: + """ + A sparse linear operator which resamples values from an input grid onto an + output grid using :mod:`torch`. + + This is a :mod:`torch` analogue of + :func:`regridding.regrid_from_weights`, except that the sparse matrix is + assembled once and applied many times, and the result is differentiable. + + Examples + -------- + + Project a uniform scene onto the sensors of an ideal CTIS instrument. + + .. jupyter-execute:: + + import numpy as np + import astropy.units as u + import named_arrays as na + import ctis + + # Define the grid of velocities and positions on the skyplane + coordinates_scene = na.DopplerPositionalVectorArray.from_velocity( + velocity=na.linspace(-300, 300, axis="wavelength", num=6) * u.km / u.s, + wavelength_rest=630 * u.AA, + position=na.Cartesian2dVectorLinearSpace( + start=-5 * u.arcsec, + stop=5 * u.arcsec, + axis=na.Cartesian2dVectorArray("scene_x", "scene_y"), + num=17, + ), + ) + + # Define the grid of positions on the sensor + coordinates_sensor = na.DopplerPositionalVectorArray.from_velocity( + velocity=coordinates_scene.velocity, + wavelength_rest=630 * u.AA, + position=na.Cartesian2dVectorArray( + x=na.arange(0, 33, axis="sensor_x") * u.pix, + y=na.arange(0, 33, axis="sensor_y") * u.pix, + ), + ) + + # Define an idealized CTIS instrument with two channels + angle = na.linspace(0, 180, axis="channel", num=2, endpoint=False) * u.deg + instrument = ctis.instruments.IdealInstrument( + area_effective=1 * u.cm**2, + timedelta_exposure=10 * u.s, + plate_scale=0.5 * u.arcsec / u.pix, + dispersion=0.01 * u.AA / u.pix, + angle=angle, + wavelength_ref=630 * u.AA, + position_ref=16 * u.pix, + coordinates_scene=coordinates_scene, + coordinates_sensor=coordinates_sensor, + channel=angle, + axis_channel="channel", + axis_wavelength="wavelength", + axis_scene_xy=("scene_x", "scene_y"), + axis_sensor_xy=("sensor_x", "sensor_y"), + ) + + # Assemble the sparse forward operator + regridder = ctis.Regridder.from_weights( + weights=instrument.weights, + axis_input=instrument.axis_scene_xy, + axis_output=instrument.axis_sensor_xy, + ) + + # Project a uniform scene onto the sensors. + # The operator is placed on a CUDA device if one is available, so the + # values must be created on `regridder.device` to match. + import torch + scene = torch.ones( + regridder.shape_values_input, + device=regridder.device, + ) + image = regridder(scene) + + image.shape + """ + + matrix: "torch.Tensor" = dataclasses.MISSING + """The sparse CSR matrix representing this operator.""" + + axis_block: tuple[str, ...] = dataclasses.MISSING + """ + The logical axes along which this operator is block diagonal. + + These are the axes of the array of weights computed by + :func:`regridding.weights`, usually the wavelength and channel axes. + """ + + axis_input: tuple[str, ...] = dataclasses.MISSING + """The logical axes of the input grid which are resampled.""" + + axis_output: tuple[str, ...] = dataclasses.MISSING + """The logical axes of the output grid which are resampled.""" + + shape_input: dict[str, int] = dataclasses.MISSING + """The shape of the input grid.""" + + shape_output: dict[str, int] = dataclasses.MISSING + """The shape of the output grid.""" + + unit: None | u.UnitBase = None + """ + The unit of the weights, if they carry one. + + The values returned by :meth:`__call__` are plain :mod:`torch` tensors, + so the caller is responsible for reapplying this unit. + """ + + @classmethod + def from_weights( + cls, + weights: tuple["na.AbstractScalar", dict[str, int], dict[str, int]], + axis_input: str | tuple[str, ...], + axis_output: str | tuple[str, ...], + device: None | str = None, + dtype: "None | torch.dtype" = None, + ) -> "Regridder": + """ + Assemble a sparse operator from weights computed by + :func:`regridding.weights`. + + Parameters + ---------- + weights + The weights computed by :func:`regridding.weights`, usually the + :attr:`~ctis.instruments.AbstractLinearInstrument.weights` attribute + of an instrument. + axis_input + The logical axes of the input grid which are resampled, usually the + :attr:`~ctis.instruments.AbstractInstrument.axis_scene_xy` attribute + of an instrument. + axis_output + The logical axes of the output grid which are resampled, usually the + :attr:`~ctis.instruments.AbstractInstrument.axis_sensor_xy` attribute + of an instrument. + device + The :mod:`torch` device on which to place the matrix. + If :obj:`None`, a CUDA device is used if one is available. + dtype + The floating-point type of the matrix. + If :obj:`None`, :obj:`torch.float32` is used. + Double precision is very slow on consumer GPUs and is usually + not needed since the matrix elements are geometric areas. + """ + torch = _torch() + + if device is None: + device = "cuda" if torch.cuda.is_available() else "cpu" + + if dtype is None: + dtype = torch.float32 + + array, shape_input, shape_output = weights + + if isinstance(axis_input, str): # pragma: nocover + axis_input = (axis_input,) + if isinstance(axis_output, str): # pragma: nocover + axis_output = (axis_output,) + + axis_input = tuple(axis_input) + axis_output = tuple(axis_output) + axis_block = tuple(array.axes) + + num_input = int(np.prod([shape_input[ax] for ax in axis_input])) + num_output = int(np.prod([shape_output[ax] for ax in axis_output])) + + flat = array.ndarray.reshape(-1) + num_block = flat.size + + rows = [] + columns = [] + values = [] + unit = None + + for d in range(num_block): + indices_input, indices_output, values_d = flat[d] + unit_d = getattr(values_d, "unit", None) + if unit_d is not None: + unit = unit_d + values_d = values_d.value + rows.append(np.asarray(indices_output) + d * num_output) + columns.append(np.asarray(indices_input) + d * num_input) + values.append(np.asarray(values_d)) + + rows = np.concatenate(rows) + columns = np.concatenate(columns) + values = np.concatenate(values) + + size = (num_block * num_output, num_block * num_input) + + if max(size) > np.iinfo(np.int32).max: # pragma: nocover + raise ValueError( + f"the operator shape {size} is too large for 32-bit indices." + ) + + # sort by row so the matrix can be stored in CSR format, which is both + # faster and deterministic on CUDA devices. + order = np.argsort(rows, kind="stable") + columns = columns[order].astype(np.int32) + values = values[order] + + crow = np.zeros(size[0] + 1, dtype=np.int32) + crow[1:] = np.cumsum(np.bincount(rows, minlength=size[0])) + + matrix = torch.sparse_csr_tensor( + crow_indices=torch.as_tensor(crow, device=device), + col_indices=torch.as_tensor(columns, device=device), + values=torch.as_tensor(values, device=device).to(dtype), + size=size, + ) + + return cls( + matrix=matrix, + axis_block=axis_block, + axis_input=axis_input, + axis_output=axis_output, + shape_input=shape_input, + shape_output=shape_output, + unit=unit, + ) + + @property + def axes_values_input(self) -> tuple[str, ...]: + """The logical axes expected by :meth:`__call__`, in order.""" + return self.axis_block + self.axis_input + + @property + def axes_values_output(self) -> tuple[str, ...]: + """The logical axes returned by :meth:`__call__`, in order.""" + return self.axis_block + self.axis_output + + @property + def shape_values_input(self) -> tuple[int, ...]: + """The shape of the array expected by :meth:`__call__`.""" + return tuple(self.shape_input[ax] for ax in self.axes_values_input) + + @property + def shape_values_output(self) -> tuple[int, ...]: + """The shape of the array returned by :meth:`__call__`.""" + return tuple(self.shape_output[ax] for ax in self.axes_values_output) + + @property + def device(self) -> "torch.device": + """The device on which this operator is stored.""" + return self.matrix.device + + @property + def dtype(self) -> "torch.dtype": + """The floating-point type of this operator.""" + return self.matrix.dtype + + def __call__(self, values: "torch.Tensor") -> "torch.Tensor": + """ + Resample an array of values from the input grid onto the output grid. + + Parameters + ---------- + values + The values to resample. + The trailing axes must match :attr:`shape_values_input`, and any + leading axes are treated as batch axes. + """ + shape_input = self.shape_values_input + shape_output = self.shape_values_output + + ndim = len(shape_input) + + if tuple(values.shape[values.ndim - ndim :]) != shape_input: + raise ValueError( + f"the trailing axes of {tuple(values.shape)=} should match " + f"{shape_input}." + ) + + shape_batch = tuple(values.shape[: values.ndim - ndim]) + + num_input = self.matrix.shape[1] + + result = values.reshape(*shape_batch, num_input) + + if shape_batch: + result = result.reshape(-1, num_input).transpose(0, 1) + result = (self.matrix @ result).transpose(0, 1) + else: + result = self.matrix @ result + + return result.reshape(*shape_batch, *shape_output) diff --git a/ctis/_torch_test.py b/ctis/_torch_test.py new file mode 100644 index 0000000..551aec8 --- /dev/null +++ b/ctis/_torch_test.py @@ -0,0 +1,228 @@ +import pytest +import numpy as np +import astropy.units as u +import named_arrays as na +import ctis + +torch = pytest.importorskip("torch") + + +wavelength_rest = 630 * u.AA + +velocity = na.linspace(-300, 300, axis="wavelength", num=6) * u.km / u.s + +coordinates_scene = na.DopplerPositionalVectorArray.from_velocity( + velocity=velocity, + wavelength_rest=wavelength_rest, + position=na.Cartesian2dVectorLinearSpace( + start=-4 * u.arcsec, + stop=4 * u.arcsec, + axis=na.Cartesian2dVectorArray("scene_x", "scene_y"), + num=17, + ), +) + +coordinates_sensor = na.DopplerPositionalVectorArray.from_velocity( + velocity=velocity, + wavelength_rest=wavelength_rest, + position=na.Cartesian2dVectorArray( + x=na.arange(0, 33, axis="sensor_x") * u.pix, + y=na.arange(0, 33, axis="sensor_y") * u.pix, + ), +) + +angle = na.linspace(0, 180, num=2, axis="channel", endpoint=False) * u.deg + +instrument = ctis.instruments.IdealInstrument( + area_effective=1 * u.cm**2, + timedelta_exposure=10 * u.s, + plate_scale=0.5 * u.arcsec / u.pix, + dispersion=0.02 * u.AA / u.pix, + angle=angle, + wavelength_ref=wavelength_rest, + position_ref=16 * u.pix, + coordinates_scene=coordinates_scene, + coordinates_sensor=coordinates_sensor, + channel=angle, + axis_channel="channel", + axis_wavelength="wavelength", + axis_scene_xy=("scene_x", "scene_y"), + axis_sensor_xy=("sensor_x", "sensor_y"), +) + + +def _regridder(device: str = "cpu", dtype=None) -> ctis.Regridder: + return ctis.Regridder.from_weights( + weights=instrument.weights, + axis_input=instrument.axis_scene_xy, + axis_output=instrument.axis_sensor_xy, + device=device, + dtype=dtype, + ) + + +def _values(regridder: ctis.Regridder, seed: int = 42) -> na.ScalarArray: + rng = np.random.default_rng(seed) + return na.ScalarArray( + ndarray=rng.random(regridder.shape_values_input), + axes=regridder.axes_values_input, + ) + + +devices = ["cpu"] +if torch.cuda.is_available(): # pragma: nocover + devices.append("cuda") + + +@pytest.mark.parametrize("device", devices) +class TestRegridder: + def test_shape(self, device: str): + a = _regridder(device) + + assert a.axes_values_input == ("wavelength", "channel") + tuple( + instrument.axis_scene_xy + ) + assert a.axes_values_output == ("wavelength", "channel") + tuple( + instrument.axis_sensor_xy + ) + + assert a.shape_values_input == (5, 2, 16, 16) + assert a.shape_values_output == (5, 2, 32, 32) + + assert a.matrix.shape == (5 * 2 * 32 * 32, 5 * 2 * 16 * 16) + assert str(a.device).startswith(device) + assert a.dtype == torch.float32 + + def test__call__(self, device: str): + """The torch operator must agree with the numba implementation.""" + a = _regridder(device) + + values = _values(a) + + expected = na.regridding.regrid_from_weights( + *instrument.weights, + values_input=values, + ) + expected = expected.ndarray_aligned(a.axes_values_output) + + result = a(torch.as_tensor(values.ndarray, dtype=a.dtype, device=device)) + result = result.detach().cpu().numpy() + + assert result.shape == expected.shape + assert np.allclose(result, expected, rtol=1e-5) + + def test__call__invalid(self, device: str): + a = _regridder(device) + with pytest.raises(ValueError): + a(torch.zeros((3, 3), device=device)) + + def test__call__deterministic(self, device: str): + """ + The result must be bitwise reproducible so that conjugate gradient + methods see a consistent operator. + """ + a = _regridder(device) + x = torch.as_tensor( + _values(a).ndarray, + dtype=a.dtype, + device=device, + ) + + expected = a(x) + for _ in range(8): + assert torch.equal(a(x), expected) + + def test__call__batch(self, device: str): + """Leading axes are treated as independent batch elements.""" + a = _regridder(device) + + x = torch.as_tensor( + np.stack([_values(a, seed=s).ndarray for s in range(3)]), + dtype=a.dtype, + device=device, + ) + + result = a(x) + + assert result.shape == (3,) + a.shape_values_output + + for i in range(3): + assert torch.allclose(result[i], a(x[i])) + + def test_adjoint(self, device: str): + r""" + Automatic differentiation must give the exact transpose, + :math:`\langle u, A x \rangle = \langle A^T u, x \rangle`. + """ + a = _regridder(device) + + rng = np.random.default_rng(0) + + x = torch.as_tensor( + _values(a).ndarray, + dtype=a.dtype, + device=device, + ).requires_grad_(True) + + u_ = torch.as_tensor( + rng.random(a.shape_values_output), + dtype=a.dtype, + device=device, + ) + + y = a(x) + + lhs = (u_ * y).sum() + (transpose,) = torch.autograd.grad(y, x, grad_outputs=u_) + rhs = (transpose * x).sum() + + assert torch.allclose(lhs, rhs, rtol=1e-4) + + def test_unit(self, device: str): + a = _regridder(device) + assert a.unit is None or isinstance(a.unit, u.UnitBase) + + +def test_gradcheck(): + """Verify the autograd path against finite differences in double precision.""" + a = _regridder("cpu", dtype=torch.float64) + + x = torch.as_tensor( + _values(a).ndarray, + dtype=torch.float64, + ).requires_grad_(True) + + assert torch.autograd.gradcheck(a, (x,), eps=1e-6, atol=1e-6) + + +def test_from_weights_unit(): + """Weights which carry a unit are stripped, and the unit is recorded.""" + array, shape_input, shape_output = instrument.weights + + flat = array.ndarray.reshape(-1) + modified = np.empty(flat.shape, dtype=object) + for d in range(flat.size): + indices_input, indices_output, values = flat[d] + modified[d] = (indices_input, indices_output, values * u.cm**2) + + weights = ( + na.ScalarArray( + ndarray=modified.reshape(array.ndarray.shape), + axes=array.axes, + ), + shape_input, + shape_output, + ) + + a = ctis.Regridder.from_weights( + weights=weights, + axis_input=instrument.axis_scene_xy, + axis_output=instrument.axis_sensor_xy, + device="cpu", + ) + + assert a.unit == u.cm**2 + + expected = _regridder("cpu") + values = torch.as_tensor(_values(a).ndarray, dtype=a.dtype) + assert torch.allclose(a(values), expected(values)) diff --git a/ctis/instruments/_instruments.py b/ctis/instruments/_instruments.py index 7cb2747..d914c93 100644 --- a/ctis/instruments/_instruments.py +++ b/ctis/instruments/_instruments.py @@ -202,6 +202,31 @@ def weights_transpose( skyplane. """ + @property + @abc.abstractmethod + def response(self) -> tuple[na.AbstractScalar, na.AbstractScalar]: + r""" + The diagonal factors of the noiseless forward model. + + The forward model of a linear instrument can always be written as + + .. math:: + + \text{image} = \sum_\lambda R_\text{out} \odot + \text{regrid}(W, R_\text{in} \odot \text{scene}), + + where :math:`W` are the :attr:`weights` and :math:`R_\text{in}` and + :math:`R_\text{out}` are the two factors returned by this property. + :math:`R_\text{in}` converts the spectral radiance of the scene into + the quantity consumed by :attr:`weights`, and :math:`R_\text{out}` + converts the resampled quantity into the electrons measured by the + sensor. + + Both factors are diagonal, which lets an external differentiable + implementation of the forward model, such as :class:`ctis.Regridder`, + reproduce :meth:`image` exactly without reimplementing it. + """ + @property def num_channel(self) -> int: @@ -468,6 +493,16 @@ class IdealInstrument( once per readout (after integrating over wavelength), in electrons. """ + @property + def response(self) -> tuple[na.AbstractScalar, na.AbstractScalar]: + scale_input = ( + self.area_effective + * self.timedelta_exposure + * self._volume_scene + / self._energy_per_photon + ) + return scale_input, self.quantum_yield + def _shot_noise(self, image: na.ScalarArray) -> na.ScalarArray: # photon shot noise, converted back into electrons to match the # electron-valued image @@ -743,6 +778,48 @@ def weights(self) -> tuple[na.AbstractScalar, dict[str, int], dict[str, int]]: axis_field=self.axis_scene_xy, ) + @property + def response(self) -> tuple[na.AbstractScalar, na.AbstractScalar]: + + system = self.system + axis_wavelength = self.axis_wavelength + coordinates = self.coordinates_scene + + # `optika` folds the effective area, vignetting, and field stop into + # the weights, so the only factor applied to the scene beforehand is + # the volume of each voxel. + scale_input = coordinates.spectral_positional.volume_cell( + (axis_wavelength, *self.axis_scene_xy) + ) + + # Probe the sensor with a unit energy rate to recover its gain. + # With `noise=False` the sensor response is linear and diagonal + # (`optika` skips charge diffusion for the per-pixel expectation), + # so this is exact. + # The sensor accepts either an energy or a photon rate; probing with + # an energy rate folds the (wavelength-dependent) energy per photon + # into the gain, so that this instrument reports an energy radiance + # like `IdealInstrument` does. + rate = 1 * u.erg / u.s + probe = na.FunctionArray( + inputs=na.SpectralPositionalVectorArray( + wavelength=coordinates.wavelength, + position=system.coordinates_sensor, + ), + outputs=na.ScalarArray(rate), + ) + gain = system.sensor.expose( + image=probe, + direction=system.direction, + axis_wavelength=axis_wavelength, + noise=False, + integrate=False, + ) + + scale_output = system.weights_unit * gain.outputs / rate + + return scale_input, scale_output + @functools.cached_property def weights_transpose( self, diff --git a/ctis/instruments/_instruments_test.py b/ctis/instruments/_instruments_test.py index d706c37..e2df500 100644 --- a/ctis/instruments/_instruments_test.py +++ b/ctis/instruments/_instruments_test.py @@ -194,7 +194,25 @@ def test_read_noise( class AbstractTestAbstractLinearInstrument( AbstractTestAbstractInstrument, ): - pass + + def test_response(self, a: ctis.instruments.AbstractLinearInstrument): + """ + The two diagonal factors of the forward model, which allow an external + differentiable implementation to reproduce :meth:`image` exactly. + """ + scale_input, scale_output = a.response + + assert np.all(na.value(scale_input) > 0) + assert np.all(na.value(scale_output) > 0) + + # the two factors, the weights, and a spectral radiance must combine + # into electrons + radiance = u.electron / ( + na.unit_normalized(scale_input) * na.unit_normalized(scale_output) + ) + assert radiance.is_equivalent( + u.erg / (u.cm**2 * u.sr * u.s * u.AA) + ) or radiance.is_equivalent(u.ph / (u.cm**2 * u.sr * u.s * u.AA)) velocity = na.linspace(-500, 500, axis="wavelength", num=21) * u.km / u.s diff --git a/ctis/inverters/__init__.py b/ctis/inverters/__init__.py index f2f8514..fdc3897 100644 --- a/ctis/inverters/__init__.py +++ b/ctis/inverters/__init__.py @@ -8,6 +8,13 @@ MartInverter, IterativeInversionResult, ) +from ._parametric import ( + AbstractSpectralModel, + GaussianModel, + AbstractParametricInverter, + ParametricInverter, + ParametricInversionResult, +) __all__ = [ "merit", @@ -17,4 +24,9 @@ "AbstractInversionResult", "InversionResult", "IterativeInversionResult", + "AbstractSpectralModel", + "GaussianModel", + "AbstractParametricInverter", + "ParametricInverter", + "ParametricInversionResult", ] diff --git a/ctis/inverters/_iterative/_mart/_mart.py b/ctis/inverters/_iterative/_mart/_mart.py index 1cd7eb9..09668a3 100644 --- a/ctis/inverters/_iterative/_mart/_mart.py +++ b/ctis/inverters/_iterative/_mart/_mart.py @@ -1,6 +1,7 @@ import warnings import dataclasses import numpy as np +import astropy.units as u import named_arrays as na import ctis from .. import AbstractIterativeInverter, IterativeInversionResult @@ -41,11 +42,25 @@ class MartInverter( threshold_convergence: float = 1e-3 r""" The convergence threshold, :math:`T`, which halts the iteration. - + If :math:`\langle \chi_{i-1}^2 \rangle - \langle \chi_{i}^2 \rangle < T`, then the algorithm is considered to be converged. """ + unit: None | u.UnitBase = None + """ + The unit of the reconstructed scene. + + This is forwarded to + :meth:`~ctis.instruments.AbstractInstrument.backproject`, which expresses + the result in either photon or energy units as requested. + If :obj:`None` (the default), the natural units of the backprojection are + used, which differ between instruments. + + The iteration itself is unaffected, since the multiplicative correction is + a ratio of two backprojections and is therefore dimensionless. + """ + def __post_init__(self): if self.gamma is None: @@ -89,7 +104,7 @@ def __call__( images = images.outputs if guess is None: - scene = instrument.backproject(images).outputs + scene = instrument.backproject(images, unit=self.unit).outputs scene = scene.mean(axis_channel) scene.ndarray[:] = scene.ndarray.mean() else: @@ -99,7 +114,7 @@ def __call__( gamma = self.gamma - backprojected = instrument.backproject(images).outputs + backprojected = instrument.backproject(images, unit=self.unit).outputs backprojected = np.maximum(backprojected, 0) @@ -138,7 +153,10 @@ def __call__( num_iteration = i + 1 break - backprojected_new = instrument.backproject(images_new).outputs + backprojected_new = instrument.backproject( + images_new, + unit=self.unit, + ).outputs backprojected_new = np.maximum(backprojected_new, 0) diff --git a/ctis/inverters/_parametric/__init__.py b/ctis/inverters/_parametric/__init__.py new file mode 100644 index 0000000..df186f7 --- /dev/null +++ b/ctis/inverters/_parametric/__init__.py @@ -0,0 +1,14 @@ +from ._models import AbstractSpectralModel, GaussianModel +from ._parametric import ( + AbstractParametricInverter, + ParametricInverter, + ParametricInversionResult, +) + +__all__ = [ + "AbstractSpectralModel", + "GaussianModel", + "AbstractParametricInverter", + "ParametricInverter", + "ParametricInversionResult", +] diff --git a/ctis/inverters/_parametric/_models.py b/ctis/inverters/_parametric/_models.py new file mode 100644 index 0000000..70bbef1 --- /dev/null +++ b/ctis/inverters/_parametric/_models.py @@ -0,0 +1,287 @@ +from typing import TYPE_CHECKING +import abc +import dataclasses +import math +import astropy.units as u +from ..._torch import _torch + +__all__ = [ + "AbstractSpectralModel", + "GaussianModel", +] + +if TYPE_CHECKING: # pragma: nocover + import torch + + +@dataclasses.dataclass +class AbstractSpectralModel( + abc.ABC, +): + """ + An interface describing a parameterized model of the spectral line profile + observed in each spatial pixel of a scene. + + Models are expressed in Doppler velocity space since the rest wavelength + of the observed line is usually known from atomic physics. + + The optimizer which fits these models operates on `unconstrained` + parameters, which are mapped onto `physical` parameters by :meth:`physical`. + This allows bounds such as the positivity of the line intensity to be + expressed exactly instead of as a penalty. + + Implementations should choose link functions which make every + unconstrained parameter of order unity, since gradient-based optimizers + take steps of a fixed size in the unconstrained space. + Positive quantities which span orders of magnitude, such as the line + radiance, should therefore be parameterized logarithmically rather than + with a softplus. + """ + + @property + @abc.abstractmethod + def parameters(self) -> tuple[str, ...]: + """The name of each physical parameter of this model.""" + + @property + def num_parameters(self) -> int: + """The number of free parameters of this model, per spatial pixel.""" + return len(self.parameters) + + @abc.abstractmethod + def unit(self, intensity: u.UnitBase) -> dict[str, u.UnitBase]: + """ + The unit of each quantity returned by :meth:`physical`. + + Parameters + ---------- + intensity + The unit of the line radiance integrated over the spectral line. + This is determined by the instrument rather than by the model, so + it is supplied by the caller. + """ + + @abc.abstractmethod + def physical( + self, + parameters: "torch.Tensor", + ) -> dict[str, "torch.Tensor"]: + """ + Map unconstrained parameters onto physical parameters. + + Parameters + ---------- + parameters + The unconstrained parameters seen by the optimizer. + The leading axis has :attr:`num_parameters` elements and the + remaining axes are the spatial axes of the scene. + """ + + @abc.abstractmethod + def guess( + self, + **kwargs: "torch.Tensor", + ) -> "torch.Tensor": + """ + Map physical parameters onto unconstrained parameters. + + This is the inverse of :meth:`physical`, and is used to convert an + initial guess expressed in physical units into a starting point for + the optimizer. + + Parameters + ---------- + kwargs + The physical parameters, named according to :attr:`parameters`. + """ + + @abc.abstractmethod + def __call__( + self, + parameters: "torch.Tensor", + velocity: "torch.Tensor", + ) -> "torch.Tensor": + r""" + Evaluate the mean spectral radiance within each velocity bin. + + The profile is integrated `analytically` across each bin rather than + sampled at the bin center. + This is essential, not cosmetic: the reconstruction grid is usually + comparable to or coarser than the width of the spectral line, so + sampling at bin centers would make a sub-bin Doppler shift almost + unobservable. + + Parameters + ---------- + parameters + The unconstrained parameters seen by the optimizer, + with shape ``(num_parameters, ...)``. + velocity + The edges of each velocity bin, in units of + :math:`\text{km} \, \text{s}^{-1}`. + """ + + +@dataclasses.dataclass +class GaussianModel( + AbstractSpectralModel, +): + r""" + A single Gaussian spectral line profile. + + The radiance within the velocity bin spanning :math:`[v_0, v_1]` is + + .. math:: + + \frac{I}{2 (v_1 - v_0)} \left[ + \text{erf} \left( \frac{v_1 - v}{\sqrt{2} \sigma} \right) + - \text{erf} \left( \frac{v_0 - v}{\sqrt{2} \sigma} \right) + \right], + + The unconstrained parameters seen by the optimizer are mapped onto the + physical parameters using + + .. math:: + + I = e^{\theta_0}, \quad + v = v_\text{max} \tanh \theta_1, \quad + \sigma_\text{nonthermal} = e^{\theta_2}, + + where :math:`I` is the radiance integrated over the line, + :math:`v` is the bulk Doppler velocity, + and the total width is + + .. math:: + + \sigma^2 = \sigma_\text{thermal}^2 + + \sigma_\text{instrument}^2 + + \sigma_\text{nonthermal}^2. + + Since :attr:`width_thermal` is fixed by the mass of the emitting ion and + the formation temperature of the line, and :attr:`width_instrument` is + fixed by the instrument, the only free width parameter is the nonthermal + width. + Fitting the nonthermal width directly means that the physically + interesting quantity is the one which receives an uncertainty, instead of + being recovered afterwards from a difference of comparable squares. + It also makes the lower bound on the observed width exact rather than a + penalty. + """ + + width_thermal: u.Quantity = 0 * u.km / u.s + r""" + The thermal Doppler width, :math:`\sigma_\text{thermal}`, of the observed + spectral line. + + This is :math:`\sqrt{k_\text{B} T / m}`, where :math:`T` is the formation + temperature of the line and :math:`m` is the mass of the emitting ion. + """ + + width_instrument: u.Quantity = 0 * u.km / u.s + r""" + The width, :math:`\sigma_\text{instrument}`, of the instrument's spectral + response function. + """ + + velocity_max: u.Quantity = 300 * u.km / u.s + r""" + The maximum magnitude of the bulk Doppler velocity. + + The velocity is parameterized as + :math:`v = v_\text{max} \tanh(\theta)`, which bounds the fit to the + passband and prevents the optimizer from wandering into aliased solutions. + """ + + @property + def parameters(self) -> tuple[str, ...]: + return ("intensity", "velocity", "width_nonthermal") + + @property + def _velocity_max(self) -> float: + return self.velocity_max.to_value(u.km / u.s) + + @property + def _width_fixed_squared(self) -> float: + """The known contributions to the variance of the line profile.""" + width_thermal = self.width_thermal.to_value(u.km / u.s) + width_instrument = self.width_instrument.to_value(u.km / u.s) + return width_thermal**2 + width_instrument**2 + + def unit(self, intensity: u.UnitBase) -> dict[str, u.UnitBase]: + velocity = u.km / u.s + return dict( + intensity=intensity, + velocity=velocity, + width_nonthermal=velocity, + width=velocity, + ) + + def physical( + self, + parameters: "torch.Tensor", + ) -> dict[str, "torch.Tensor"]: + torch = _torch() + + intensity = torch.exp(parameters[0]) + velocity = self._velocity_max * torch.tanh(parameters[1]) + width_nonthermal = torch.exp(parameters[2]) + + width = torch.sqrt(self._width_fixed_squared + width_nonthermal**2) + + return dict( + intensity=intensity, + velocity=velocity, + width_nonthermal=width_nonthermal, + width=width, + ) + + def guess( + self, + intensity: "torch.Tensor" = None, + velocity: "torch.Tensor" = None, + width_nonthermal: "torch.Tensor" = None, + ) -> "torch.Tensor": + torch = _torch() + + velocity_max = self._velocity_max + + velocity = torch.clamp( + velocity / velocity_max, + min=-1 + 1e-6, + max=+1 - 1e-6, + ) + + tiny = torch.finfo(intensity.dtype).tiny + + return torch.stack( + [ + torch.log(torch.clamp(intensity, min=tiny)), + torch.atanh(velocity), + torch.log(torch.clamp(width_nonthermal, min=tiny)), + ] + ) + + def __call__( + self, + parameters: "torch.Tensor", + velocity: "torch.Tensor", + ) -> "torch.Tensor": + torch = _torch() + + p = self.physical(parameters) + + intensity = p["intensity"] + shift = p["velocity"] + width = p["width"] + + shape = (-1,) + (1,) * intensity.ndim + + lower = velocity[:-1].reshape(shape) + upper = velocity[+1:].reshape(shape) + + scale = math.sqrt(2) * width + + result = torch.erf((upper - shift) / scale) + result = result - torch.erf((lower - shift) / scale) + + return intensity * result / (2 * (upper - lower)) diff --git a/ctis/inverters/_parametric/_models_test.py b/ctis/inverters/_parametric/_models_test.py new file mode 100644 index 0000000..6c0877f --- /dev/null +++ b/ctis/inverters/_parametric/_models_test.py @@ -0,0 +1,143 @@ +import abc +import pytest +import numpy as np +import astropy.units as u +import ctis + +torch = pytest.importorskip("torch") + + +velocity = torch.linspace(-400, 400, 33, dtype=torch.float64) + + +class AbstractTestAbstractSpectralModel( + abc.ABC, +): + + def test_parameters(self, a: ctis.inverters.AbstractSpectralModel): + result = a.parameters + assert isinstance(result, tuple) + assert len(result) > 0 + for name in result: + assert isinstance(name, str) + + def test_num_parameters(self, a: ctis.inverters.AbstractSpectralModel): + result = a.num_parameters + assert isinstance(result, int) + assert result == len(a.parameters) + + def test_physical(self, a: ctis.inverters.AbstractSpectralModel): + parameters = torch.zeros((a.num_parameters, 3, 4), dtype=torch.float64) + result = a.physical(parameters) + + assert isinstance(result, dict) + for name in a.parameters: + assert name in result + assert result[name].shape == (3, 4) + + def test_guess(self, a: ctis.inverters.AbstractSpectralModel): + """`guess` must be the inverse of `physical`.""" + parameters = torch.zeros((a.num_parameters, 3, 4), dtype=torch.float64) + physical = a.physical(parameters) + + result = a.guess(**{k: physical[k] for k in a.parameters}) + + assert result.shape == parameters.shape + assert torch.allclose(result, parameters, atol=1e-6) + + def test__call__(self, a: ctis.inverters.AbstractSpectralModel): + parameters = torch.zeros((a.num_parameters, 3, 4), dtype=torch.float64) + result = a(parameters, velocity) + + assert result.shape == (velocity.numel() - 1, 3, 4) + assert torch.all(result >= 0) + + +@pytest.mark.parametrize( + argnames="a", + argvalues=[ + ctis.inverters.GaussianModel(), + ctis.inverters.GaussianModel( + width_thermal=11 * u.km / u.s, + width_instrument=8 * u.km / u.s, + velocity_max=250 * u.km / u.s, + ), + ], +) +class TestGaussianModel( + AbstractTestAbstractSpectralModel, +): + def test_parameters_names(self, a: ctis.inverters.GaussianModel): + assert a.parameters == ("intensity", "velocity", "width_nonthermal") + + @pytest.mark.parametrize("intensity", [1.0, 1234.0]) + @pytest.mark.parametrize("shift", [-100.0, 0.0, 50.0]) + @pytest.mark.parametrize("width_nonthermal", [20.0, 60.0]) + def test__call__normalization( + self, + a: ctis.inverters.GaussianModel, + intensity: float, + shift: float, + width_nonthermal: float, + ): + """ + Integrating the profile over a velocity range which covers the whole + line must return the integrated radiance. + """ + parameters = a.guess( + intensity=torch.tensor(intensity, dtype=torch.float64), + velocity=torch.tensor(shift, dtype=torch.float64), + width_nonthermal=torch.tensor(width_nonthermal, dtype=torch.float64), + ) + + result = a(parameters, velocity) + + width_bin = torch.diff(velocity) + integral = (result * width_bin).sum() + + assert np.isclose(integral.item(), intensity, rtol=1e-6) + + def test__call__bin_integrated(self, a: ctis.inverters.GaussianModel): + """ + The profile must be the mean of the underlying Gaussian across each + bin, not the Gaussian sampled at the bin center. + """ + parameters = a.guess( + intensity=torch.tensor(1.0, dtype=torch.float64), + velocity=torch.tensor(30.0, dtype=torch.float64), + width_nonthermal=torch.tensor(25.0, dtype=torch.float64), + ) + + result = a(parameters, velocity).numpy() + + physical = a.physical(parameters) + shift = physical["velocity"].item() + width = physical["width"].item() + + # numerically integrate the Gaussian across each bin + edges = velocity.numpy() + expected = np.empty(edges.size - 1) + for i in range(expected.size): + v = np.linspace(edges[i], edges[i + 1], 201) + g = np.exp(-np.square((v - shift) / width) / 2) + g = g / (width * np.sqrt(2 * np.pi)) + expected[i] = np.trapezoid(g, v) / (edges[i + 1] - edges[i]) + + assert np.allclose(result, expected, atol=1e-9) + + def test_width(self, a: ctis.inverters.GaussianModel): + """The total width must include the fixed thermal and instrument widths.""" + parameters = torch.zeros((a.num_parameters, 2), dtype=torch.float64) + physical = a.physical(parameters) + + width_thermal = a.width_thermal.to_value(u.km / u.s) + width_instrument = a.width_instrument.to_value(u.km / u.s) + + expected = np.sqrt( + width_thermal**2 + + width_instrument**2 + + physical["width_nonthermal"].numpy() ** 2 + ) + + assert np.allclose(physical["width"].numpy(), expected) + assert np.all(physical["width"].numpy() >= physical["width_nonthermal"].numpy()) diff --git a/ctis/inverters/_parametric/_parametric.py b/ctis/inverters/_parametric/_parametric.py new file mode 100644 index 0000000..c6496ea --- /dev/null +++ b/ctis/inverters/_parametric/_parametric.py @@ -0,0 +1,860 @@ +from typing import ClassVar, TYPE_CHECKING +import abc +import warnings +import dataclasses +import numpy as np +import astropy.units as u +import astropy.constants +import named_arrays as na +import ctis +from ..._torch import _torch +from .. import AbstractInverter, AbstractInversionResult +from ._models import AbstractSpectralModel + +__all__ = [ + "AbstractParametricInverter", + "ParametricInverter", + "ParametricInversionResult", +] + +if TYPE_CHECKING: # pragma: nocover + import torch + + +def _smoothness(physical: dict[str, "torch.Tensor"]) -> "torch.Tensor": + r""" + The mean squared first difference of each physical parameter along each + spatial axis of the scene, normalized by the mean square of that parameter. + + The penalty is applied to the physical parameters rather than to the + unconstrained parameters seen by the optimizer, since the link functions + are strongly nonlinear: a large change in :math:` heta` is a small change + in a velocity which is close to its bound, so smoothing :math:` heta` + does not smooth the velocity. Normalizing makes the penalty dimensionless + and comparable between parameters. + """ + torch = _torch() + result = 0 + for value in physical.values(): + scale = torch.mean(torch.square(value)) + torch.finfo(value.dtype).tiny + for axis in range(value.ndim): + difference = torch.diff(value, dim=axis) + result = result + torch.mean(torch.square(difference)) / scale + return result + + +@dataclasses.dataclass +class AbstractParametricInverter( + AbstractInverter, +): + """ + An abstract inversion algorithm which reconstructs an observed scene by + fitting a parameterized spectral line profile to every spatial pixel. + + Unlike :class:`~ctis.inverters.AbstractIterativeInverter`, which solves for + the radiance in every voxel of the scene, these algorithms solve for a + handful of parameters in every `spatial` pixel of the scene. + For a three-parameter model this reduces the number of unknowns below the + number of measurements, turning an underdetermined problem into an + overdetermined one. + + The parameters of neighboring pixels are `not` independent, since each + sensor pixel collects light from many spatial pixels. + The fit is therefore a single optimization over every parameter of every + pixel simultaneously. + """ + + axis_iteration: ClassVar[str] = "iteration" + """The logical axis associated with changing iteration index.""" + + @property + @abc.abstractmethod + def model(self) -> AbstractSpectralModel: + """The spectral line profile fit to each spatial pixel of the scene.""" + + +@dataclasses.dataclass +class ParametricInverter( + AbstractParametricInverter, +): + r""" + Fit a parameterized spectral line profile to every spatial pixel of the + scene using the Adam optimizer :cite:p:`Kingma2014`. + + The forward model of :attr:`instrument` is assembled into a sparse + :class:`~ctis.Regridder`, so the gradient of the merit function with + respect to every parameter is computed exactly by automatic + differentiation, and the whole optimization can run on a GPU. + + The merit function is + + .. math:: + + \langle \chi^2 \rangle = \left\langle + \left( \frac{d - \hat{d}(\theta)}{\sigma} \right)^2 + \right\rangle, + + where :math:`d` are the measured electrons, :math:`\hat{d}` are the + electrons predicted by the model, and :math:`\sigma` is estimated from the + `measured` signal rather than the predicted signal, which avoids biasing + the fitted radiance. + + Examples + -------- + + Reconstruct a scene of randomly-placed Gaussians observed by an idealized + CTIS instrument. + + .. jupyter-execute:: + + import matplotlib.pyplot as plt + import astropy.units as u + import astropy.visualization + import named_arrays as na + import ctis + + # Define the grid of velocities and positions on the skyplane. + # The velocity bins are chosen to be comparable to the width of the + # spectral line, since a parametric fit cannot recover a width which + # is much narrower than one bin. + wavelength_rest = 630 * u.AA + velocity = na.linspace(-250, 250, axis="wavelength", num=21) * u.km / u.s + coordinates_scene = na.DopplerPositionalVectorArray.from_velocity( + velocity=velocity, + wavelength_rest=wavelength_rest, + position=na.Cartesian2dVectorLinearSpace( + start=-10 * u.arcsec, + stop=10 * u.arcsec, + axis=na.Cartesian2dVectorArray("scene_x", "scene_y"), + num=33, + ), + ) + + # Define the grid of positions on the sensor. + # The sensor must be large enough to hold the dispersed scene, + # otherwise the voxels which fall off the edge are unconstrained. + coordinates_sensor = na.DopplerPositionalVectorArray.from_velocity( + velocity=velocity, + wavelength_rest=wavelength_rest, + position=na.Cartesian2dVectorArray( + x=na.arange(0, 97, axis="sensor_x") * u.pix, + y=na.arange(0, 97, axis="sensor_y") * u.pix, + ), + ) + + # Define an idealized CTIS instrument with four channels + angle = na.linspace(0, 360, axis="channel", num=4, endpoint=False) * u.deg + instrument = ctis.instruments.IdealInstrument( + area_effective=1 * u.cm**2, + timedelta_exposure=20 * u.s, + plate_scale=0.625 * u.arcsec / u.pix, + dispersion=0.021 * u.AA / u.pix, + angle=angle, + wavelength_ref=wavelength_rest, + position_ref=48 * u.pix, + coordinates_scene=coordinates_scene, + coordinates_sensor=coordinates_sensor, + channel=angle, + axis_channel="channel", + axis_wavelength="wavelength", + axis_scene_xy=("scene_x", "scene_y"), + axis_sensor_xy=("sensor_x", "sensor_y"), + ) + + # Simulate an observation of a test scene + scene = ctis.scenes.gaussians(coordinates_scene) + images = instrument.image(scene) + + # Fit a Gaussian line profile to every spatial pixel + inverter = ctis.inverters.ParametricInverter( + instrument=instrument, + model=ctis.inverters.GaussianModel( + width_thermal=11 * u.km / u.s, + width_instrument=8 * u.km / u.s, + velocity_max=250 * u.km / u.s, + ), + num_iteration=500, + ) + result = inverter(images) + + # Plot the fitted Doppler velocity. + # The unit is stripped from the color values since a colorbar norm + # cannot hold an `astropy` quantity. + velocity_fit = result.parameters["velocity"] + with astropy.visualization.quantity_support(): + fig, ax = plt.subplots(constrained_layout=True) + img = na.plt.pcolormesh( + coordinates_scene.position.x, + coordinates_scene.position.y, + C=velocity_fit.value, + ax=ax, + cmap="RdBu_r", + ) + ax.set_aspect("equal") + plt.colorbar( + img.ndarray.item(), + ax=ax, + label=f"Doppler velocity ({velocity_fit.unit:latex_inline})", + ) + """ + + instrument: ctis.instruments.AbstractLinearInstrument = dataclasses.MISSING + """ + A model of a CTIS instrument which transforms the radiance of an observed + scene to the electrons measured by the sensors. + + Any :class:`~ctis.instruments.AbstractLinearInstrument` is supported, since + the forward model is rebuilt from its + :attr:`~ctis.instruments.AbstractLinearInstrument.weights` and + :attr:`~ctis.instruments.AbstractLinearInstrument.response`. + """ + + model: AbstractSpectralModel = dataclasses.MISSING + """The spectral line profile fit to each spatial pixel of the scene.""" + + num_iteration: int = dataclasses.field(default=500, kw_only=True) + """The maximum number of optimizer steps to perform.""" + + num_iteration_guess: int = dataclasses.field(default=100, kw_only=True) + """ + The number of MART iterations used to compute the initial guess. + + The moments of the reconstruction found by + :class:`~ctis.inverters.MartInverter` are used as the starting point of + the fit. + Since the merit function is not convex, the quality of this guess has a + strong effect on the quality of the fit. + """ + + learning_rate: float = dataclasses.field(default=0.05, kw_only=True) + """The initial learning rate of the Adam optimizer.""" + + learning_rate_decay: float = dataclasses.field(default=0.01, kw_only=True) + """ + The ratio of the final learning rate to :attr:`learning_rate`. + + The learning rate decays geometrically over :attr:`num_iteration` steps. + Adam takes steps of roughly a fixed size, so an annealed learning rate is + needed to resolve the Doppler velocity to a small fraction of a bin. + """ + + threshold_convergence: float = dataclasses.field(default=1e-6, kw_only=True) + r""" + The fractional decrease in :math:`\langle \chi^2 \rangle` which counts as + an improvement. + + Together with :attr:`num_patience` this determines when the optimization + is considered to be converged. + """ + + num_patience: int = dataclasses.field(default=50, kw_only=True) + r""" + The number of iterations to continue without improving + :math:`\langle \chi^2 \rangle` before declaring convergence. + + Adam takes steps of a roughly fixed size, so the merit function is not + monotonic and a single uphill step does not mean the fit has converged. + """ + + regularization: float = dataclasses.field(default=0, kw_only=True) + r""" + The weight of a spatial smoothness penalty added to the merit function. + + The penalty is the mean squared first difference of the `unconstrained` + parameters along each spatial axis of the scene. Since the link functions + make every unconstrained parameter of order unity, a single weight + penalizes each of them comparably. + + A CTIS measures only a handful of projections, so a fit with three + parameters in every spatial pixel can be only marginally overdetermined, + and will then reproduce the noise in the measurement rather than the scene. + A nonzero value here trades resolution for that overfitting. + + If zero (the default), the fit is unregularized. + """ + + uncertainty: None | na.AbstractScalar = dataclasses.field( + default=None, + kw_only=True, + ) + """ + The standard deviation of the measurement noise, in electrons. + + If :obj:`None` (the default) and `images` carries an uncertainty, as + produced by ``instrument.image(uncertainty=True)``, that uncertainty is + used. + Otherwise the instrument's own noise model is evaluated once at the + starting guess and held fixed, which captures the read-noise floor. + """ + + variance_min: float = dataclasses.field(default=0, kw_only=True) + """ + A lower bound on the variance, in electrons squared. + + This is a backstop against an instrument whose noise model returns zero for + a pixel which measured no signal, which would give that pixel infinite + weight. It is not a substitute for a noise model: setting it too large + down-weights the faint pixels, and too small weights them far too heavily. + """ + + device: None | str = dataclasses.field(default=None, kw_only=True) + """ + The :mod:`torch` device on which to perform the optimization. + + If :obj:`None`, a CUDA device is used if one is available. + """ + + def _validate(self) -> None: + instrument = self.instrument + + if not isinstance(instrument, ctis.instruments.AbstractLinearInstrument): + raise ValueError( + f"{type(instrument)=} is not supported, the forward model must " + f"be a `ctis.instruments.AbstractLinearInstrument`." + ) + + coordinates = instrument.coordinates_scene + if not isinstance(coordinates, na.AbstractDopplerPositionalVectorArray): + raise ValueError( + "`instrument.coordinates_scene` must be a Doppler vector array " + "since the spectral model is expressed in velocity space." + ) + + @property + def _velocity(self) -> np.ndarray: + """The edges of each velocity bin of the scene, in km/s.""" + instrument = self.instrument + velocity = instrument.coordinates_scene.velocity + velocity = velocity.ndarray_aligned((instrument.axis_wavelength,)) + return velocity.to_value(u.km / u.s) + + @property + def _dvdl(self) -> float: + r""" + The derivative of Doppler velocity with respect to wavelength, + in :math:`\text{km} \, \text{s}^{-1} \, \AA^{-1}`. + """ + wavelength_rest = self.instrument.coordinates_scene.wavelength_rest + result = astropy.constants.c / wavelength_rest + return result.to_value(u.km / u.s / u.AA) + + @property + def unit_intensity(self) -> u.UnitBase: + """ + The unit of the fitted line radiance, integrated over the spectral line. + + This is derived from the units of + :attr:`~ctis.instruments.AbstractLinearInstrument.response`, so an + instrument whose forward model consumes a photon radiance yields a + photon radiance, and one which consumes an energy radiance yields an + energy radiance. + """ + scale_input, scale_output = self.instrument.response + + unit = na.unit_normalized(scale_input) * na.unit_normalized(scale_output) + + result = u.electron / unit * u.AA + + # express the result in a conventional radiance unit if possible + for candidate in ( + u.erg / (u.cm**2 * u.sr * u.s), + u.ph / (u.cm**2 * u.sr * u.s), + ): + if result.is_equivalent(candidate): + return candidate + + return result # pragma: nocover + + def _response( + self, + regridder: "ctis.Regridder", + ) -> tuple[np.ndarray, np.ndarray]: + r""" + The two diagonal factors of the forward model, expressed as plain + arrays laid out for :class:`~ctis.Regridder`. + + The first converts the profile returned by :attr:`model`, which is a + radiance per unit velocity, into the quantity consumed by the weights. + The second converts the resampled quantity into electrons. + """ + instrument = self.instrument + + axis_wavelength = instrument.axis_wavelength + axis_scene_xy = tuple(instrument.axis_scene_xy) + + scale_input, scale_output = instrument.response + + # The model works in velocity space while the instrument integrates + # over wavelength, so convert the radiance density between the two, + # then attach the unit of the intensity parameter and express the + # product in whatever unit makes the second factor yield electrons. + scale_input = ( + scale_input + * self._dvdl + * (u.km / u.s / u.AA) + * self.unit_intensity + / (u.km / u.s) + ) + scale_input = scale_input.to(u.electron / na.unit_normalized(scale_output)) + + axes_input = (axis_wavelength,) + axis_scene_xy + shape_input = {ax: regridder.shape_input[ax] for ax in axes_input} + scale_input = na.broadcast_to(na.as_named_array(scale_input), shape_input) + scale_input = na.value(scale_input.ndarray_aligned(axes_input)) + + axes_output = regridder.axes_values_output + shape_output = {ax: regridder.shape_output[ax] for ax in axes_output} + scale_output = na.broadcast_to(na.as_named_array(scale_output), shape_output) + scale_output = na.value(scale_output.ndarray_aligned(axes_output)) + + return scale_input, scale_output + + @property + def unit_parameters(self) -> dict[str, u.UnitBase]: + """ + The unit of each physical parameter of :attr:`model`. + + The unit of the line radiance is determined by the instrument, see + :attr:`unit_intensity`, while the remaining units are determined by + the model. + """ + return self.model.unit(self.unit_intensity) + + def _scene(self, profile: "torch.Tensor") -> na.FunctionArray: + """ + Express a spectral profile computed by :attr:`model` as a scene which + the instrument can observe. + """ + instrument = self.instrument + axes = (instrument.axis_wavelength,) + tuple(instrument.axis_scene_xy) + + # the model works in velocity space, so convert the radiance density + # back into the wavelength basis the instrument expects. + outputs = profile.detach().cpu().numpy() * self._dvdl + + return na.FunctionArray( + inputs=instrument.coordinates_scene, + outputs=na.ScalarArray( + ndarray=outputs << (self.unit_intensity / u.AA), + axes=axes, + ), + ) + + def _variance( + self, + images: na.FunctionArray[na.SpectralPositionalVectorArray, na.ScalarArray], + axes: tuple[str, ...], + profile: "torch.Tensor", + ) -> np.ndarray: + """ + The variance of every measurement, in electrons squared. + + The instrument's own noise model is used unless one is supplied, + because a measurement which is not shot-noise limited has a floor set + by the read noise, and assuming otherwise weights the faint pixels far + too heavily. + """ + outputs = images.outputs + + if self.uncertainty is not None: + width = na.as_named_array(self.uncertainty).ndarray_aligned(axes) + + elif isinstance(outputs, na.AbstractUncertainScalarArray): + # the instrument was already asked for its own uncertainty model + width = outputs.width.ndarray_aligned(axes) + + else: + # ask the instrument for its noise model at the starting guess. + # This is evaluated once and held fixed, so it does not bias the + # fitted radiance the way a merit function whose weights follow the + # model would. + width = self.instrument.image( + self._scene(profile), + noise=False, + uncertainty=True, + ) + width = width.outputs.width.ndarray_aligned(axes) + + variance = np.square(u.Quantity(width).to_value(u.electron)) + + return np.maximum(variance, self.variance_min) + + @staticmethod + def _weight(variance: np.ndarray) -> np.ndarray: + """ + The reciprocal uncertainty of every measurement. + + A measurement whose variance is zero carries no information, so it is + given zero weight rather than an infinite one. + """ + return np.where( + variance > 0, 1 / np.sqrt(np.where(variance > 0, variance, 1)), 0 + ) + + def _guess( + self, + images: na.FunctionArray[na.SpectralPositionalVectorArray, na.ScalarArray], + guess: ( + None + | dict[str, na.AbstractScalar] + | na.AbstractScalar + | na.AbstractFunctionArray + ), + ) -> dict[str, np.ndarray]: + """ + Compute the starting point of the fit, as the physical value of every + parameter in every spatial pixel. + """ + instrument = self.instrument + + axis_wavelength = instrument.axis_wavelength + axis_scene_xy = tuple(instrument.axis_scene_xy) + + # the physical parameters may be supplied directly, which allows a fit + # to be warm-started from the result of a previous one. + if isinstance(guess, dict): + unit = self.unit_parameters + shape = { + ax: instrument.coordinates_scene.shape[ax] - 1 for ax in axis_scene_xy + } + result = dict() + for name in self.model.parameters: + if name not in guess: + raise ValueError( + f"`guess` is missing the parameter {name!r}, " + f"expected {self.model.parameters}." + ) + value = na.broadcast_to(na.as_named_array(guess[name]), shape) + value = u.Quantity(value.ndarray_aligned(axis_scene_xy)) + result[name] = value.to_value(unit[name]) + return result + + if guess is None: + inverter = ctis.inverters.MartInverter( + instrument=instrument, + num_iteration=self.num_iteration_guess, + # the natural units of the backprojection differ between + # instruments, and a photon radiance cannot be converted into + # an energy radiance without the energy per photon, so ask for + # the unit this fit works in. + unit=self.unit_intensity / u.AA, + ) + with warnings.catch_warnings(): + # the guess is deliberately stopped before convergence, and + # MART divides by zero in voxels which received no signal. + warnings.simplefilter("ignore", UserWarning) + with np.errstate(invalid="ignore", divide="ignore"): + # MART operates on the measured values, so any uncertainty + # attached to the images is dropped before the guess. + guess = inverter( + images.replace(outputs=na.nominal(images.outputs)), + ).solution + + if isinstance(guess, na.AbstractFunctionArray): + guess = guess.outputs + + axes = (axis_wavelength,) + axis_scene_xy + + cube = guess.ndarray_aligned(axes) + cube = u.Quantity(cube).to_value(self.unit_intensity / u.AA) + + # the model works in velocity space, so convert the radiance density + # from a wavelength basis into a velocity basis. + cube = cube / self._dvdl + + cube = np.maximum(cube, 0) + + velocity = self._velocity + width_bin = np.diff(velocity) + center_bin = (velocity[:-1] + velocity[1:]) / 2 + + shape = (-1,) + (1,) * len(axis_scene_xy) + weight = cube * width_bin.reshape(shape) + + intensity = weight.sum(0) + intensity_safe = np.maximum(intensity, np.finfo(float).tiny) + + center = center_bin.reshape(shape) + mean = (weight * center).sum(0) / intensity_safe + variance = (weight * np.square(center - mean)).sum(0) / intensity_safe + + variance_fixed = self.model._width_fixed_squared + width_min = np.abs(width_bin).min() / 10 + width = np.sqrt(np.maximum(variance - variance_fixed, np.square(width_min))) + + intensity = np.maximum(intensity, intensity.max() / 1e6) + + return dict( + intensity=intensity, + velocity=mean, + width_nonthermal=width, + ) + + def __call__( + self, + images: na.FunctionArray[na.SpectralPositionalVectorArray, na.ScalarArray], + guess: ( + None + | dict[str, na.AbstractScalar] + | na.AbstractScalar + | na.AbstractFunctionArray + ) = None, + verbose: bool = False, + ) -> "ParametricInversionResult": + """ + Reconstruct a scene using the observed images. + + Parameters + ---------- + images + The observed images used to calculate the reconstruction. + Must be evaluated on the same position coordinates as + :attr:`~ctis.instruments.AbstractInstrument.coordinates_sensor` + attribute of :attr:`instrument`. + guess + The starting point of the fit, given in any of three forms. + + A :class:`dict` of physical parameters, named according to + :attr:`~ctis.inverters.AbstractSpectralModel.parameters` of + :attr:`model`, is used directly. Each value may be a full map over + the spatial axes of the scene, or a scalar which is broadcast over + them. Since + :attr:`~ctis.inverters.ParametricInversionResult.parameters` is a + dictionary of exactly this form, the result of one fit may be used + to warm-start another. + + A reconstructed scene, as either an + :class:`~named_arrays.AbstractScalar` or an + :class:`~named_arrays.AbstractFunctionArray`, is reduced to its + moments in every spatial pixel. + + If :obj:`None` (the default), a scene is first reconstructed using + :class:`~ctis.inverters.MartInverter`, and its moments are used. + verbose + Whether to print the merit function at every iteration. + """ + torch = _torch() + + self._validate() + + instrument = self.instrument + model = self.model + + axis_wavelength = instrument.axis_wavelength + axis_scene_xy = tuple(instrument.axis_scene_xy) + axis_sensor_xy = tuple(instrument.axis_sensor_xy) + + position_images = images.inputs.position + position_sensor = instrument.coordinates_sensor.position + if not np.all(position_images == position_sensor): + raise ValueError( + "`images.inputs.position` and `self.coordinates_sensor.position` " + "are not equal." + ) + + regridder = ctis.Regridder.from_weights( + weights=instrument.weights, + axis_input=axis_scene_xy, + axis_output=axis_sensor_xy, + device=self.device, + ) + + device = regridder.device + dtype = regridder.dtype + + axis_block = regridder.axis_block + index_wavelength = axis_block.index(axis_wavelength) + + def _tensor(array: np.ndarray) -> "torch.Tensor": + return torch.as_tensor(np.ascontiguousarray(array)).to( + device=device, + dtype=dtype, + ) + + velocity = _tensor(self._velocity) + + scale_input, scale_output = self._response(regridder) + scale_input = _tensor(scale_input) + scale_output = _tensor(scale_output) + + axes_data = tuple(ax for ax in axis_block if ax != axis_wavelength) + axes_data = axes_data + axis_sensor_xy + + outputs = images.outputs + + data = na.nominal(outputs) + data = u.Quantity(data.ndarray_aligned(axes_data)).to_value(u.electron) + data = _tensor(data) + + shape_cube = regridder.shape_values_input + + def forward(parameters: "torch.Tensor") -> "torch.Tensor": + cube = model(parameters, velocity) * scale_input + for i, ax in enumerate(axis_block): + if ax != axis_wavelength: + cube = cube.unsqueeze(i) + cube = cube.expand(shape_cube) + result = regridder(cube) * scale_output + return result.sum(index_wavelength) + + parameters = self._guess(images, guess) + parameters = model.guess(**{k: _tensor(v) for k, v in parameters.items()}) + parameters = parameters.detach().clone().requires_grad_(True) + + variance = self._variance(images, axes_data, model(parameters, velocity)) + weight = self._weight(variance) + num_measurement = max(int(np.count_nonzero(weight)), 1) + weight = _tensor(weight) + + optimizer = torch.optim.Adam([parameters], lr=self.learning_rate) + scheduler = torch.optim.lr_scheduler.ExponentialLR( + optimizer=optimizer, + gamma=self.learning_rate_decay ** (1 / max(self.num_iteration - 1, 1)), + ) + + mean_chi_squared = [] + merit_best = np.inf + merit_reference = np.inf + iteration_reference = 0 + parameters_best = parameters.detach().clone() + message = f"Max number of iterations ({self.num_iteration}) exceeded." + success = False + num_iteration = self.num_iteration + + for i in range(self.num_iteration): + + optimizer.zero_grad(set_to_none=True) + + residual = (data - forward(parameters)) * weight + merit = torch.sum(torch.square(residual)) / num_measurement + + if self.regularization: + penalty = _smoothness(model.physical(parameters)) + merit = merit + self.regularization * penalty + + merit.backward() + optimizer.step() + scheduler.step() + + merit = merit.item() + mean_chi_squared.append(merit) + + # Adam takes steps of a roughly fixed size, so the merit function + # is not monotonic. Keep the best solution seen so far instead of + # whichever one the last step happened to land on. + if merit < merit_best: + merit_best = merit + parameters_best = parameters.detach().clone() + + if verbose: # pragma: nocover + print(f"{i=}, merit={merit}") + + # only a decrease larger than the threshold resets the patience + if merit < merit_reference * (1 - self.threshold_convergence): + merit_reference = merit + iteration_reference = i + + if (i - iteration_reference) >= self.num_patience: + message = ( + f"The merit function did not decrease by more than " + f"{self.threshold_convergence} over the last " + f"{self.num_patience} iterations." + ) + success = True + num_iteration = i + 1 + break + + else: + warnings.warn(message) + + parameters = parameters_best + + with torch.no_grad(): + physical = model.physical(parameters) + profile = model(parameters, velocity) + + unit = self.unit_parameters + parameters_result = { + k: na.ScalarArray( + ndarray=physical[k].detach().cpu().numpy() << unit[k], + axes=axis_scene_xy, + ) + for k in physical + } + + solution = self._scene(profile) + + mean_chi_squared = na.ScalarArray( + ndarray=np.array(mean_chi_squared), + axes=(self.axis_iteration,), + ) + + return ParametricInversionResult( + solution=solution, + parameters=parameters_result, + success=success, + images=images, + inverter=self, + message=message, + num_iteration=num_iteration, + mean_chi_squared=mean_chi_squared, + ) + + +@dataclasses.dataclass +class ParametricInversionResult( + AbstractInversionResult, +): + """The results of a parametric inversion attempt.""" + + solution: na.FunctionArray[na.SpectralPositionalVectorArray, na.ScalarArray] = ( + dataclasses.MISSING + ) + """ + The reconstructed scene found by the inversion. + + This is the spectral line profile evaluated using :attr:`parameters`. + """ + + parameters: dict[str, na.ScalarArray] = dataclasses.MISSING + """ + The fitted value of each physical parameter in every spatial pixel. + + The keys are the names of the parameters of + :attr:`~ctis.inverters.AbstractParametricInverter.model`. + """ + + success: bool = dataclasses.MISSING + """A boolean flag indicating whether the inversion was successful.""" + + images: na.FunctionArray[na.SpectralPositionalVectorArray, na.ScalarArray] = ( + dataclasses.MISSING + ) + """The observed images on which the inversion was performed.""" + + inverter: "ctis.inverters.AbstractInverter" = dataclasses.MISSING + """The inversion algorithm instance that produced these results.""" + + message: str = dataclasses.MISSING + """Any message from the inversion routine concerning the results.""" + + num_iteration: int = dataclasses.MISSING + """The number of iterations performed by the inverter.""" + + mean_chi_squared: na.ScalarArray = dataclasses.MISSING + """ + The mean chi squared statistic after each iteration. + + This is not monotonic, since Adam takes steps of a roughly fixed size. + :attr:`parameters` is the best solution found over all iterations, not the + solution found by the last one. + """ + + @property + def iteration(self) -> na.ScalarArray: + """The iteration value for each iteration.""" + return na.arange( + start=0, + stop=self.num_iteration, + axis=self.inverter.axis_iteration, + ) diff --git a/ctis/inverters/_parametric/_parametric_test.py b/ctis/inverters/_parametric/_parametric_test.py new file mode 100644 index 0000000..5c54df1 --- /dev/null +++ b/ctis/inverters/_parametric/_parametric_test.py @@ -0,0 +1,548 @@ +import dataclasses +import warnings +import pytest +import numpy as np +import astropy.units as u +import named_arrays as na +import ctis +from .._inverters_test import AbstractTestAbstractInverter + +torch = pytest.importorskip("torch") + + +wavelength_rest = 630 * u.AA + +#: velocity bins comparable to the width of the line, and a sensor large +#: enough to hold the dispersed scene. +velocity = na.linspace(-200, 200, axis="wavelength", num=17) * u.km / u.s + +coordinates_scene = na.DopplerPositionalVectorArray.from_velocity( + velocity=velocity, + wavelength_rest=wavelength_rest, + position=na.Cartesian2dVectorLinearSpace( + start=-5 * u.arcsec, + stop=+5 * u.arcsec, + axis=na.Cartesian2dVectorArray("scene_x", "scene_y"), + num=17, + ), +) + +coordinates_sensor = na.DopplerPositionalVectorArray.from_velocity( + velocity=velocity, + wavelength_rest=wavelength_rest, + position=na.Cartesian2dVectorArray( + x=na.arange(0, 65, axis="sensor_x") * u.pix, + y=na.arange(0, 65, axis="sensor_y") * u.pix, + ), +) + +angle = na.linspace(0, 360, num=4, axis="channel", endpoint=False) * u.deg + +instrument = ctis.instruments.IdealInstrument( + area_effective=1 * u.cm**2, + timedelta_exposure=20 * u.s, + plate_scale=0.625 * u.arcsec / u.pix, + dispersion=0.021 * u.AA / u.pix, + angle=angle, + wavelength_ref=wavelength_rest, + position_ref=32 * u.pix, + coordinates_scene=coordinates_scene, + coordinates_sensor=coordinates_sensor, + channel=angle, + axis_channel="channel", + axis_wavelength="wavelength", + axis_scene_xy=("scene_x", "scene_y"), + axis_sensor_xy=("sensor_x", "sensor_y"), +) + +model = ctis.inverters.GaussianModel( + width_thermal=11 * u.km / u.s, + width_instrument=8 * u.km / u.s, + velocity_max=200 * u.km / u.s, +) + + +def _truth() -> dict[str, np.ndarray]: + """A random set of physical parameters for every spatial pixel.""" + rng = np.random.default_rng(3) + num = 16 + return dict( + intensity=2000.0 + 3000.0 * rng.random((num, num)), + velocity=60.0 * (2 * rng.random((num, num)) - 1), + width_nonthermal=15.0 + 25.0 * rng.random((num, num)), + ) + + +def _scene( + truth: dict[str, np.ndarray], +) -> na.FunctionArray[na.DopplerPositionalVectorArray, na.ScalarArray]: + """Evaluate the spectral model to produce a scene it can represent exactly.""" + inverter = ctis.inverters.ParametricInverter(instrument=instrument, model=model) + unit = inverter.unit_intensity / u.AA + + parameters = model.guess( + **{k: torch.as_tensor(v, dtype=torch.float32) for k, v in truth.items()} + ) + profile = model( + parameters, + torch.as_tensor(inverter._velocity, dtype=torch.float32), + ) + profile = profile.numpy() * inverter._dvdl + + return na.FunctionArray( + inputs=coordinates_scene, + outputs=na.ScalarArray( + ndarray=profile << unit, + axes=("wavelength", "scene_x", "scene_y"), + ), + ) + + +truth = _truth() +scene = _scene(truth) +images = instrument.image(scene, noise=False) + + +class AbstractTestAbstractParametricInverter( + AbstractTestAbstractInverter, +): + + def test_model(self, a: ctis.inverters.AbstractParametricInverter): + result = a.model + assert isinstance(result, ctis.inverters.AbstractSpectralModel) + + def test__call__( + self, + a: ctis.inverters.AbstractParametricInverter, + images: na.FunctionArray[na.SpectralPositionalVectorArray, na.ScalarArray], + **kwargs, + ) -> ctis.inverters.ParametricInversionResult: + + result = super().test__call__(a=a, images=images, **kwargs) + + assert result.num_iteration > 0 + assert result.mean_chi_squared.shape[a.axis_iteration] == result.num_iteration + assert result.iteration.size == result.num_iteration + + # the reported solution is the best one seen, so the merit can + # never end up worse than where it started + merit = result.mean_chi_squared.ndarray + assert merit.min() <= merit[0] + + assert isinstance(result.parameters, dict) + for name in a.model.parameters: + assert name in result.parameters + parameter = result.parameters[name] + assert isinstance(parameter, na.ScalarArray) + assert parameter.shape == {ax: 16 for ax in a.instrument.axis_scene_xy} + assert np.all(np.isfinite(parameter.ndarray)) + + assert np.all(result.solution.outputs >= 0) + + return result + + +@pytest.mark.parametrize( + argnames="a", + argvalues=[ + ctis.inverters.ParametricInverter( + instrument=instrument, + model=model, + num_iteration=400, + ), + ctis.inverters.ParametricInverter( + instrument=instrument, + model=model, + num_iteration=400, + num_iteration_guess=10, + learning_rate=0.1, + device="cpu", + ), + ], +) +class TestParametricInverter( + AbstractTestAbstractParametricInverter, +): + @pytest.mark.parametrize("images", [images]) + @pytest.mark.parametrize( + argnames="guess", + argvalues=[ + None, + scene, + ], + ) + def test__call__( + self, + a: ctis.inverters.ParametricInverter, + images: na.FunctionArray[na.SpectralPositionalVectorArray, na.ScalarArray], + guess: None | na.AbstractFunctionArray, + ): + # whether a fit converges within its iteration cap depends on the + # configuration, so the warning is neither required nor forbidden here + with warnings.catch_warnings(): + warnings.simplefilter("ignore", UserWarning) + return super().test__call__(a=a, images=images, guess=guess) + + +def test__call__recovery(): + """ + A scene which the model can represent exactly must be recovered to high + accuracy from noiseless images. + """ + inverter = ctis.inverters.ParametricInverter( + instrument=instrument, + model=model, + num_iteration=2000, + ) + + result = inverter(images) + + assert result.success + assert result.num_iteration < inverter.num_iteration + + axis = ("scene_x", "scene_y") + + for name in model.parameters: + expected = truth[name] + got = result.parameters[name].ndarray_aligned(axis).value + r = np.corrcoef(got.ravel(), expected.ravel())[0, 1] + assert r > 0.95, f"{name} was recovered with a correlation of only {r}" + + merit = result.mean_chi_squared.ndarray + assert merit.min() < merit[0] / 100 + + +def test__call__invalid_position(): + inverter = ctis.inverters.ParametricInverter( + instrument=instrument, + model=model, + num_iteration=2, + ) + with pytest.raises(ValueError): + inverter(images.replace(inputs=coordinates_scene)) + + +def test__call__invalid_instrument(): + """The forward model must be a linear instrument.""" + inverter = ctis.inverters.ParametricInverter( + instrument=object(), + model=model, + ) + with pytest.raises(ValueError): + inverter(images) + + +def test__call__invalid_coordinates(): + """The spectral model requires scene coordinates expressed in velocity.""" + inverter = ctis.inverters.ParametricInverter( + instrument=dataclasses.replace( + instrument, + coordinates_scene=na.SpectralPositionalVectorArray( + wavelength=coordinates_scene.wavelength, + position=coordinates_scene.position, + ), + ), + model=model, + ) + with pytest.raises(ValueError): + inverter(images) + + +def _instrument_optika() -> ctis.instruments.OptikaInstrument: + """A CTIS instrument whose forward model is an `optika` linear system.""" + import optika + + channel = na.linspace(0, 360, axis="channel", num=3, endpoint=False) * u.deg + system = optika.systems.LinearSystem( + area_effective=optika.radiometry.InterpolatedEffectiveAreaModel( + wavelength=na.linspace(400, 700, axis="wavelength", num=10) * u.nm, + area=na.linspace(1, 2, axis="wavelength", num=10) * u.cm**2, + axis_wavelength="wavelength", + ), + distortion=optika.distortion.SimpleDistortionModel( + plate_scale=0.75 * u.arcsec / u.pix, + dispersion=3.75 * u.nm / u.pix, + angle=channel, + reference=na.SpectralPositionalVectorArray( + wavelength=550 * u.nm, + position=na.Cartesian2dVectorArray(16, 16) * u.pix, + ), + ), + sensor=optika.sensors.ImagingSensor( + width_pixel=15 * u.um, + axis_pixel=na.Cartesian2dVectorArray("sensor_x", "sensor_y"), + timedelta_exposure=1 * u.s, + num_pixel=na.Cartesian2dVectorArray(32, 32), + ), + ) + return ctis.instruments.OptikaInstrument( + system=system, + coordinates_scene=na.DopplerPositionalVectorArray.from_velocity( + velocity=na.linspace(-4000, 4000, axis="wavelength", num=7) * u.km / u.s, + wavelength_rest=550 * u.nm, + position=na.Cartesian2dVectorLinearSpace( + start=-6 * u.arcsec, + stop=+6 * u.arcsec, + axis=na.Cartesian2dVectorArray("scene_x", "scene_y"), + num=17, + ), + ), + channel=channel, + axis_channel="channel", + axis_wavelength="wavelength", + axis_scene_xy=("scene_x", "scene_y"), + ) + + +@pytest.mark.parametrize( + argnames="a,b", + argvalues=[ + (instrument, model), + ( + _instrument_optika(), + ctis.inverters.GaussianModel( + width_thermal=200 * u.km / u.s, + width_instrument=100 * u.km / u.s, + velocity_max=4000 * u.km / u.s, + ), + ), + ], +) +def test_forward_matches_instrument( + a: ctis.instruments.AbstractLinearInstrument, + b: ctis.inverters.AbstractSpectralModel, +): + """ + The differentiable forward model assembled from `weights` and `response` + must reproduce `instrument.image` for every kind of linear instrument. + """ + inverter = ctis.inverters.ParametricInverter(instrument=a, model=b) + + axis_wavelength = a.axis_wavelength + axis_scene_xy = tuple(a.axis_scene_xy) + num = tuple(a.coordinates_scene.shape[ax] - 1 for ax in axis_scene_xy) + + rng = np.random.default_rng(1) + parameters = b.guess( + intensity=torch.as_tensor(1 + rng.random(num), dtype=torch.float32), + velocity=torch.as_tensor( + b.velocity_max.to_value(u.km / u.s) / 8 * (2 * rng.random(num) - 1), + dtype=torch.float32, + ), + width_nonthermal=torch.as_tensor( + b.width_thermal.to_value(u.km / u.s) + rng.random(num), + dtype=torch.float32, + ), + ) + velocity = torch.as_tensor(inverter._velocity, dtype=torch.float32) + profile = b(parameters, velocity) + + scene = na.FunctionArray( + inputs=a.coordinates_scene, + outputs=na.ScalarArray( + ndarray=(profile.numpy() * inverter._dvdl) + << (inverter.unit_intensity / u.AA), + axes=(axis_wavelength,) + axis_scene_xy, + ), + ) + expected = a.image(scene, noise=False) + + regridder = ctis.Regridder.from_weights( + weights=a.weights, + axis_input=axis_scene_xy, + axis_output=a.axis_sensor_xy, + device="cpu", + ) + axis_block = regridder.axis_block + scale_input, scale_output = inverter._response(regridder) + + cube = profile * torch.as_tensor(scale_input, dtype=torch.float32) + for i, ax in enumerate(axis_block): + if ax != axis_wavelength: + cube = cube.unsqueeze(i) + cube = cube.expand(regridder.shape_values_input) + result = regridder(cube) * torch.as_tensor(scale_output, dtype=torch.float32) + result = result.sum(axis_block.index(axis_wavelength)).detach().numpy() + + axes = tuple(ax for ax in axis_block if ax != axis_wavelength) + axes = axes + tuple(a.axis_sensor_xy) + expected = u.Quantity(expected.outputs.ndarray_aligned(axes)).to_value(u.electron) + + # the absolute tolerance is scaled to the signal so that pixels which + # received no light are not compared to float32 rounding noise + assert np.allclose( + result, + expected, + rtol=1e-5, + atol=1e-5 * expected.max(), + ) + + +def test__call__uncertainty_explicit(): + """An uncertainty may be supplied directly instead of being estimated.""" + uncertainty = np.sqrt(np.abs(images.outputs.value)) * u.electron + uncertainty = uncertainty + 1 * u.electron + + inverter = ctis.inverters.ParametricInverter( + instrument=instrument, + model=model, + num_iteration=20, + num_iteration_guess=5, + uncertainty=uncertainty, + ) + + with pytest.warns(UserWarning): + result = inverter(images) + + assert np.all(np.isfinite(result.mean_chi_squared.ndarray)) + + +def test__call__uncertainty_from_images(): + """The uncertainty attached to the images by the instrument is used.""" + images_uncertain = instrument.image(scene, noise=False, uncertainty=True) + + assert isinstance(images_uncertain.outputs, na.AbstractUncertainScalarArray) + + inverter = ctis.inverters.ParametricInverter( + instrument=instrument, + model=model, + num_iteration=20, + num_iteration_guess=5, + ) + + with pytest.warns(UserWarning): + result = inverter(images_uncertain) + + assert np.all(np.isfinite(result.mean_chi_squared.ndarray)) + + +def test__call__optika(): + """ + The whole inversion must run against an `OptikaInstrument`, whose + backprojection is naturally expressed in photon rather than energy units. + """ + a = _instrument_optika() + b = ctis.inverters.GaussianModel( + width_thermal=200 * u.km / u.s, + width_instrument=100 * u.km / u.s, + velocity_max=4000 * u.km / u.s, + ) + + inverter = ctis.inverters.ParametricInverter( + instrument=a, + model=b, + num_iteration=50, + num_iteration_guess=5, + ) + + axis_scene_xy = tuple(a.axis_scene_xy) + num = tuple(a.coordinates_scene.shape[ax] - 1 for ax in axis_scene_xy) + num_wavelength = a.coordinates_scene.shape[a.axis_wavelength] - 1 + + rng = np.random.default_rng(0) + scene_optika = na.FunctionArray( + inputs=a.coordinates_scene, + outputs=na.ScalarArray( + ndarray=(1 + rng.random((num_wavelength,) + num)) + << (inverter.unit_intensity / u.AA), + axes=(a.axis_wavelength,) + axis_scene_xy, + ), + ) + + with pytest.warns(UserWarning): + result = inverter(a.image(scene_optika, noise=False)) + + assert np.all(np.isfinite(result.mean_chi_squared.ndarray)) + for name in b.parameters: + assert np.all(np.isfinite(result.parameters[name].ndarray)) + + +def test__call__guess_parameters(): + """The physical parameters may be supplied directly.""" + inverter = ctis.inverters.ParametricInverter( + instrument=instrument, + model=model, + num_iteration=50, + ) + + guess = { + "intensity": na.ScalarArray( + ndarray=truth["intensity"] << inverter.unit_intensity, + axes=("scene_x", "scene_y"), + ), + # a scalar is broadcast over the spatial axes of the scene + "velocity": 0 * u.km / u.s, + "width_nonthermal": 25 * u.km / u.s, + } + + with pytest.warns(UserWarning): + result = inverter(images, guess=guess) + + assert np.all(np.isfinite(result.mean_chi_squared.ndarray)) + for name in model.parameters: + assert np.all(np.isfinite(result.parameters[name].ndarray)) + + +def test__call__guess_roundtrip(): + """The result of one fit can warm-start another.""" + inverter = ctis.inverters.ParametricInverter( + instrument=instrument, + model=model, + num_iteration=300, + ) + + # supply the uncertainty explicitly so that both fits are weighted + # identically and their merit functions can be compared + inverter = dataclasses.replace( + inverter, + uncertainty=instrument.image( + scene, noise=False, uncertainty=True + ).outputs.width, + ) + + with warnings.catch_warnings(): + warnings.simplefilter("ignore", UserWarning) + first = inverter(images) + second = inverter(images, guess=first.parameters) + + merit_first = first.mean_chi_squared.ndarray + merit_second = second.mean_chi_squared.ndarray + + # the restarted fit resumes from the best solution of the previous one + # the parameters round-trip through float32, hence the loose tolerance + assert merit_second[0] == pytest.approx(merit_first.min(), rel=1e-3) + assert merit_second[0] < merit_first[0] + + # and so does at least as well + assert merit_second.min() <= merit_first.min() + + +def test__call__guess_missing_parameter(): + inverter = ctis.inverters.ParametricInverter( + instrument=instrument, + model=model, + num_iteration=2, + ) + with pytest.raises(ValueError): + inverter(images, guess={"intensity": 1 * inverter.unit_intensity}) + + +def test__call__regularization(): + """A smoothness penalty produces a smoother set of parameter maps.""" + axis = ("scene_x", "scene_y") + + def roughness(result): + v = result.parameters["velocity"].ndarray_aligned(axis).value + return np.mean(np.square(np.diff(v, axis=0))) + + kwargs = dict(instrument=instrument, model=model, num_iteration=300) + + with warnings.catch_warnings(): + warnings.simplefilter("ignore", UserWarning) + rough = ctis.inverters.ParametricInverter(**kwargs)(images) + smooth = ctis.inverters.ParametricInverter( + regularization=1, + **kwargs, + )(images) + + assert roughness(smooth) < roughness(rough) diff --git a/docs/index.rst b/docs/index.rst index 064163c..76bc933 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -28,6 +28,7 @@ Examples on how to use this package. tutorials/ideal-instrument tutorials/simple-mart + tutorials/parametric-fit API Reference ============= diff --git a/docs/tutorials/parametric-fit.ipynb b/docs/tutorials/parametric-fit.ipynb new file mode 100644 index 0000000..9f859a5 --- /dev/null +++ b/docs/tutorials/parametric-fit.ipynb @@ -0,0 +1,927 @@ +{ + "cells": [ + { + "cell_type": "raw", + "id": "0", + "metadata": { + "editable": true, + "raw_mimetype": "text/x-rst", + "slideshow": { + "slide_type": "" + }, + "tags": [] + }, + "source": [ + "Fit a Spectral Line Profile to Every Pixel\n", + "==========================================" + ] + }, + { + "cell_type": "raw", + "id": "1", + "metadata": { + "editable": true, + "raw_mimetype": "text/x-rst", + "slideshow": { + "slide_type": "" + }, + "tags": [] + }, + "source": [ + ":class:`~ctis.inverters.MartInverter` solves for the radiance in every voxel of\n", + "the scene. That problem is underdetermined: the configuration used in this\n", + "tutorial has 81920 unknowns and only 32768 measurements, so the reconstruction\n", + "relies entirely on regularization.\n", + "\n", + ":class:`~ctis.inverters.ParametricInverter` takes the opposite approach. Instead\n", + "of solving for every voxel, it solves for a handful of parameters in every\n", + "*spatial* pixel: the radiance integrated over the line, the bulk Doppler\n", + "velocity, and the nonthermal width. That is 12288 unknowns for the same 32768\n", + "measurements, which makes the problem overdetermined, and it lets the Doppler\n", + "velocity be measured to a small fraction of a velocity bin.\n", + "\n", + "This tutorial deliberately uses the same grid, instrument, and synthetic scene as\n", + ":doc:`simple-mart`, so the two approaches can be compared directly.\n", + "\n", + "The parameters of neighboring pixels are not independent, since each sensor\n", + "pixel collects light from many spatial pixels. The fit is therefore a single\n", + "optimization over every parameter of every pixel simultaneously, using the\n", + "Adam optimizer running on a GPU if one is available." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2", + "metadata": {}, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt\n", + "import astropy.units as u\n", + "import astropy.visualization\n", + "import named_arrays as na\n", + "import ctis" + ] + }, + { + "cell_type": "raw", + "id": "3", + "metadata": { + "editable": true, + "raw_mimetype": "text/x-rst", + "slideshow": { + "slide_type": "" + }, + "tags": [] + }, + "source": [ + "Start by defining a grid of Doppler velocities on which to reconstruct the scene." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [ + "velocity = na.linspace(-500, 500, axis=\"wavelength\", num=21) * u.km / u.s" + ] + }, + { + "cell_type": "raw", + "id": "5", + "metadata": { + "editable": true, + "raw_mimetype": "text/x-rst", + "slideshow": { + "slide_type": "" + }, + "tags": [] + }, + "source": [ + "Define the rest wavelength used to convert between velocity and wavelength." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": {}, + "outputs": [], + "source": [ + "wavelength_rest = 171 * u.AA" + ] + }, + { + "cell_type": "raw", + "id": "7", + "metadata": { + "editable": true, + "raw_mimetype": "text/x-rst", + "slideshow": { + "slide_type": "" + }, + "tags": [] + }, + "source": [ + "Define a grid of positions on which to reconstruct the scene," + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8", + "metadata": {}, + "outputs": [], + "source": [ + "position_scene = na.Cartesian2dVectorLinearSpace(\n", + " start=-10 * u.arcsec,\n", + " stop=+10 * u.arcsec,\n", + " axis=na.Cartesian2dVectorArray(\"scene_x\", \"scene_y\"),\n", + " num=na.Cartesian2dVectorArray(64 + 1, 64 + 1),\n", + ")" + ] + }, + { + "cell_type": "raw", + "id": "9", + "metadata": { + "editable": true, + "raw_mimetype": "text/x-rst", + "slideshow": { + "slide_type": "" + }, + "tags": [] + }, + "source": [ + "and a grid of positions on the sensor representing the vertices of each pixel.\n", + "\n", + "The sensor must be large enough to hold the dispersed scene. Any voxel whose\n", + "light falls off the edge of the sensor is unconstrained by the measurement, and\n", + "the fit cannot recover it. This grid retains 99 percent of the flux." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "10", + "metadata": {}, + "outputs": [], + "source": [ + "position_sensor = na.Cartesian2dVectorArray(\n", + " x=na.arange(0, 128 + 1, axis=\"sensor_x\") * u.pix,\n", + " y=na.arange(0, 64 + 1, axis=\"sensor_y\") * u.pix,\n", + ")" + ] + }, + { + "cell_type": "raw", + "id": "11", + "metadata": { + "editable": true, + "raw_mimetype": "text/x-rst", + "slideshow": { + "slide_type": "" + }, + "tags": [] + }, + "source": [ + "Combine the 1D velocity grid and the 2D position grid into a single 3D grid for\n", + "both the scene and the sensor." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "12", + "metadata": {}, + "outputs": [], + "source": [ + "coordinates_scene = na.DopplerPositionalVectorArray.from_velocity(\n", + " velocity=velocity,\n", + " wavelength_rest=wavelength_rest,\n", + " position=position_scene,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "13", + "metadata": {}, + "outputs": [], + "source": [ + "coordinates_sensor = na.DopplerPositionalVectorArray.from_velocity(\n", + " velocity=velocity,\n", + " wavelength_rest=wavelength_rest,\n", + " position=position_sensor,\n", + ")" + ] + }, + { + "cell_type": "raw", + "id": "14", + "metadata": { + "editable": true, + "raw_mimetype": "text/x-rst", + "slideshow": { + "slide_type": "" + }, + "tags": [] + }, + "source": [ + "Create a synthetic scene composed of spatial/spectral 3D Gaussians with various\n", + "Doppler shifts, the same test pattern used by the MART tutorial." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "15", + "metadata": {}, + "outputs": [], + "source": [ + "scene = ctis.scenes.gaussians(coordinates_scene)" + ] + }, + { + "cell_type": "raw", + "id": "16", + "metadata": { + "editable": true, + "raw_mimetype": "text/x-rst", + "slideshow": { + "slide_type": "" + }, + "tags": [] + }, + "source": [ + "Add a small background equal to 1 percent of the maximum value of the scene,\n", + "exactly as the MART tutorial does. Note that this background is flat in\n", + "wavelength, so it is something a single Gaussian cannot represent. Its effect on\n", + "the fit is discussed below." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "17", + "metadata": {}, + "outputs": [], + "source": [ + "scene = scene + scene.max() / 100" + ] + }, + { + "cell_type": "raw", + "id": "18", + "metadata": { + "editable": true, + "raw_mimetype": "text/x-rst", + "slideshow": { + "slide_type": "" + }, + "tags": [] + }, + "source": [ + "Display the scene as a false-color image, where hue represents Doppler velocity." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "19", + "metadata": {}, + "outputs": [], + "source": [ + "with astropy.visualization.quantity_support():\n", + " fig, axs = plt.subplots(\n", + " ncols=2,\n", + " gridspec_kw=dict(width_ratios=[0.9, 0.1]),\n", + " constrained_layout=True,\n", + " )\n", + " ax, cax = axs\n", + " colorbar = na.plt.rgbmesh(\n", + " C=scene,\n", + " axis_wavelength=\"wavelength\",\n", + " ax=ax,\n", + " vmin=0,\n", + " vmax=scene.outputs.max(),\n", + " )\n", + " na.plt.pcolormesh(\n", + " C=colorbar,\n", + " axis_rgb=\"wavelength\",\n", + " ax=cax,\n", + " )\n", + " ax.set_aspect(\"equal\")\n", + " ax.set_xlabel(f\"scene $x$ ({ax.get_xlabel()})\")\n", + " ax.set_ylabel(f\"scene $y$ ({ax.get_ylabel()})\")\n", + " cax.yaxis.tick_right()\n", + " cax.yaxis.set_label_position(\"right\")" + ] + }, + { + "cell_type": "raw", + "id": "20", + "metadata": { + "editable": true, + "raw_mimetype": "text/x-rst", + "slideshow": { + "slide_type": "" + }, + "tags": [] + }, + "source": [ + "Define the dispersion angles and magnitude for our instrument, and create an\n", + "ideal CTIS from them." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "21", + "metadata": {}, + "outputs": [], + "source": [ + "angle = na.linspace(0, 360, num=4, axis=\"channel\", endpoint=False) * u.deg\n", + "angle = angle + 5.64 * u.deg" + ] + }, + { + "cell_type": "raw", + "id": "22", + "metadata": { + "editable": true, + "raw_mimetype": "text/x-rst", + "slideshow": { + "slide_type": "" + }, + "tags": [] + }, + "source": [ + "Define the magnitude of the dispersion in terms of Doppler velocity, and then\n", + "convert it into a wavelength per pixel." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "23", + "metadata": {}, + "outputs": [], + "source": [ + "dispersion = 10 * u.km / u.s\n", + "dispersion = dispersion.to(u.AA, equivalencies=u.doppler_optical(wavelength_rest))\n", + "dispersion = (dispersion - wavelength_rest) / u.pix" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "24", + "metadata": {}, + "outputs": [], + "source": [ + "instrument = ctis.instruments.IdealInstrument(\n", + " area_effective=1 * u.cm**2,\n", + " timedelta_exposure=20 * u.s,\n", + " plate_scale=0.4 * u.arcsec / u.pix,\n", + " dispersion=dispersion,\n", + " angle=angle,\n", + " wavelength_ref=wavelength_rest,\n", + " position_ref=na.Cartesian2dVectorArray(64, 32) * u.pix,\n", + " coordinates_scene=coordinates_scene,\n", + " coordinates_sensor=coordinates_sensor,\n", + " channel=\"dispersion angle = \" + angle.to_string_array(\"%03d\"),\n", + " axis_channel=\"channel\",\n", + " axis_wavelength=\"wavelength\",\n", + " axis_scene_xy=(\"scene_x\", \"scene_y\"),\n", + " axis_sensor_xy=(\"sensor_x\", \"sensor_y\"),\n", + ")" + ] + }, + { + "cell_type": "raw", + "id": "25", + "metadata": { + "editable": true, + "raw_mimetype": "text/x-rst", + "slideshow": { + "slide_type": "" + }, + "tags": [] + }, + "source": [ + "Apply the forward model of this instrument to the scene to calculate the\n", + "observed images, including photon shot noise." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "26", + "metadata": {}, + "outputs": [], + "source": [ + "images = instrument.image(scene)" + ] + }, + { + "cell_type": "raw", + "id": "27", + "metadata": { + "editable": true, + "raw_mimetype": "text/x-rst", + "slideshow": { + "slide_type": "" + }, + "tags": [] + }, + "source": [ + "Display the images, where each panel represents a different dispersion angle." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "28", + "metadata": {}, + "outputs": [], + "source": [ + "with astropy.visualization.quantity_support():\n", + " fig, axs = plt.subplots(\n", + " ncols=instrument.num_channel,\n", + " figsize=(11, 3.2),\n", + " constrained_layout=True,\n", + " sharex=True,\n", + " sharey=True,\n", + " )\n", + " na.plt.pcolormesh(\n", + " images.inputs.position.x,\n", + " images.inputs.position.y,\n", + " C=images.outputs.value,\n", + " ax=na.ScalarArray(axs, axes=(\"channel\",)),\n", + " cmap=\"gray\",\n", + " )\n", + " for ax, name in zip(axs, instrument.channel.ndarray):\n", + " ax.set_aspect(\"equal\")\n", + " ax.set_title(name, fontsize=9)" + ] + }, + { + "cell_type": "raw", + "id": "29", + "metadata": { + "editable": true, + "raw_mimetype": "text/x-rst", + "slideshow": { + "slide_type": "" + }, + "tags": [] + }, + "source": [ + "Now define the spectral line profile which will be fit to every spatial pixel.\n", + "\n", + "The observed width of a spectral line is the sum in quadrature of a thermal, an\n", + "instrumental, and a nonthermal contribution. The thermal width is fixed by the\n", + "mass of the emitting ion and the formation temperature of the line, and the\n", + "instrumental width is fixed by the instrument, so both are supplied here as\n", + "known constants. The only free width parameter is the nonthermal width, which\n", + "means the physically interesting quantity is the one that is fit directly\n", + "rather than recovered afterwards from a difference of comparable squares." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "30", + "metadata": {}, + "outputs": [], + "source": [ + "model = ctis.inverters.GaussianModel(\n", + " width_thermal=11 * u.km / u.s,\n", + " width_instrument=8 * u.km / u.s,\n", + " velocity_max=500 * u.km / u.s,\n", + ")" + ] + }, + { + "cell_type": "raw", + "id": "31", + "metadata": { + "editable": true, + "raw_mimetype": "text/x-rst", + "slideshow": { + "slide_type": "" + }, + "tags": [] + }, + "source": [ + "Initialize the inversion algorithm with the instrument and the spectral model.\n", + "\n", + "The merit function is not convex, so the fit is started from the moments of a\n", + "short :class:`~ctis.inverters.MartInverter` reconstruction rather than from an\n", + "arbitrary guess." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "32", + "metadata": {}, + "outputs": [], + "source": [ + "inverter = ctis.inverters.ParametricInverter(\n", + " instrument=instrument,\n", + " model=model,\n", + " num_iteration=1000,\n", + ")" + ] + }, + { + "cell_type": "raw", + "id": "33", + "metadata": { + "editable": true, + "raw_mimetype": "text/x-rst", + "slideshow": { + "slide_type": "" + }, + "tags": [] + }, + "source": [ + "Invert the images." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "34", + "metadata": {}, + "outputs": [], + "source": [ + "inversion = inverter(images)" + ] + }, + { + "cell_type": "raw", + "id": "35", + "metadata": { + "editable": true, + "raw_mimetype": "text/x-rst", + "slideshow": { + "slide_type": "" + }, + "tags": [] + }, + "source": [ + "Plot the merit function as a function of iteration.\n", + "\n", + "Adam takes steps of a roughly fixed size, so this curve is not monotonic. The\n", + "solution returned by the inverter is the best one found over all iterations,\n", + "not the one the last step happened to land on." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "36", + "metadata": {}, + "outputs": [], + "source": [ + "with astropy.visualization.quantity_support():\n", + " fig, ax = plt.subplots(constrained_layout=True)\n", + " na.plt.plot(\n", + " inversion.iteration,\n", + " inversion.mean_chi_squared,\n", + " ax=ax,\n", + " )\n", + " ax.set_yscale(\"log\")\n", + " ax.set_xlabel(\"iteration\")\n", + " ax.set_ylabel(r\"$\\langle \\chi^2 \\rangle$\")" + ] + }, + { + "cell_type": "raw", + "id": "37", + "metadata": { + "editable": true, + "raw_mimetype": "text/x-rst", + "slideshow": { + "slide_type": "" + }, + "tags": [] + }, + "source": [ + "Isolate the reconstructed scene, which is the spectral model evaluated using the\n", + "fitted parameters." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "38", + "metadata": {}, + "outputs": [], + "source": [ + "solution = inversion.solution" + ] + }, + { + "cell_type": "raw", + "id": "39", + "metadata": { + "editable": true, + "raw_mimetype": "text/x-rst", + "slideshow": { + "slide_type": "" + }, + "tags": [] + }, + "source": [ + "Display the reconstructed scene next to the original." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "40", + "metadata": {}, + "outputs": [], + "source": [ + "with astropy.visualization.quantity_support():\n", + " fig, axs = plt.subplots(\n", + " ncols=2,\n", + " figsize=(9, 4),\n", + " constrained_layout=True,\n", + " sharex=True,\n", + " sharey=True,\n", + " )\n", + " for ax, cube, title in zip(axs, [scene, solution], [\"truth\", \"reconstructed\"]):\n", + " na.plt.rgbmesh(\n", + " C=cube,\n", + " axis_wavelength=\"wavelength\",\n", + " ax=ax,\n", + " vmin=0,\n", + " vmax=scene.outputs.max(),\n", + " )\n", + " ax.set_aspect(\"equal\")\n", + " ax.set_title(title)\n", + " ax.set_xlabel(f\"scene $x$ ({ax.get_xlabel()})\")\n", + " axs[0].set_ylabel(f\"scene $y$ ({axs[0].get_ylabel()})\")" + ] + }, + { + "cell_type": "raw", + "id": "41", + "metadata": { + "editable": true, + "raw_mimetype": "text/x-rst", + "slideshow": { + "slide_type": "" + }, + "tags": [] + }, + "source": [ + "Unlike a voxel-based inversion, this method also returns the physical parameters\n", + "of the line profile in every spatial pixel directly, with no post-processing." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "42", + "metadata": {}, + "outputs": [], + "source": [ + "parameters = inversion.parameters\n", + "list(parameters)" + ] + }, + { + "cell_type": "raw", + "id": "43", + "metadata": { + "editable": true, + "raw_mimetype": "text/x-rst", + "slideshow": { + "slide_type": "" + }, + "tags": [] + }, + "source": [ + "Plot each fitted parameter as a map.\n", + "\n", + "The unit is stripped from the color values since a colorbar cannot hold an\n", + ":class:`~astropy.units.Quantity`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "44", + "metadata": {}, + "outputs": [], + "source": [ + "keys = [\"intensity\", \"velocity\", \"width_nonthermal\"]\n", + "cmaps = [\"inferno\", \"RdBu_r\", \"viridis\"]\n", + "\n", + "with astropy.visualization.quantity_support():\n", + " fig, axs = plt.subplots(\n", + " ncols=3,\n", + " figsize=(12, 3.6),\n", + " constrained_layout=True,\n", + " sharex=True,\n", + " sharey=True,\n", + " )\n", + " for ax, key, cmap in zip(axs, keys, cmaps):\n", + " parameter = parameters[key]\n", + " img = na.plt.pcolormesh(\n", + " coordinates_scene.position.x,\n", + " coordinates_scene.position.y,\n", + " C=parameter.value,\n", + " ax=ax,\n", + " cmap=cmap,\n", + " )\n", + " ax.set_aspect(\"equal\")\n", + " ax.set_xlabel(f\"scene $x$ ({ax.get_xlabel()})\")\n", + " plt.colorbar(\n", + " img.ndarray.item(),\n", + " ax=ax,\n", + " label=f\"{key} ({parameter.unit:latex_inline})\",\n", + " )\n", + " axs[0].set_ylabel(f\"scene $y$ ({axs[0].get_ylabel()})\")" + ] + }, + { + "cell_type": "raw", + "id": "45", + "metadata": { + "editable": true, + "raw_mimetype": "text/x-rst", + "slideshow": { + "slide_type": "" + }, + "tags": [] + }, + "source": [ + "A Doppler velocity can only be measured where there is a line to measure it\n", + "from. The flat background added above dominates the faint pixels, and a single\n", + "Gaussian cannot represent it, so the velocity fitted there is not meaningful.\n", + "Mask the map to the brightest pixels before comparing it against the truth." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "46", + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "\n", + "# the total radiance of the true scene, the same quantity that\n", + "# `plot_moments` thresholds on\n", + "width_bin = scene.inputs.wavelength.volume_cell(\"wavelength\")\n", + "radiance_true = (scene.outputs * width_bin).sum(\"wavelength\")\n", + "\n", + "threshold = np.nanpercentile(radiance_true.ndarray, 90)\n", + "bright = radiance_true.ndarray_aligned((\"scene_x\", \"scene_y\")) > threshold" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "47", + "metadata": {}, + "outputs": [], + "source": [ + "velocity_true = na.pdf.median(\n", + " x=scene.inputs.velocity,\n", + " f=scene.outputs,\n", + " axis=\"wavelength\",\n", + ")\n", + "\n", + "velocity_fit = parameters[\"velocity\"]\n", + "\n", + "masked = {\n", + " \"true median velocity\": velocity_true.value.ndarray_aligned((\"scene_x\", \"scene_y\")),\n", + " \"fitted velocity\": velocity_fit.value.ndarray_aligned((\"scene_x\", \"scene_y\")),\n", + "}\n", + "masked = {\n", + " k: na.ScalarArray(np.where(bright, v, np.nan), axes=(\"scene_x\", \"scene_y\"))\n", + " for k, v in masked.items()\n", + "}" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "48", + "metadata": {}, + "outputs": [], + "source": [ + "with astropy.visualization.quantity_support():\n", + " fig, axs = plt.subplots(\n", + " ncols=2,\n", + " figsize=(9, 3.6),\n", + " constrained_layout=True,\n", + " sharex=True,\n", + " sharey=True,\n", + " )\n", + " for ax, (title, C) in zip(axs, masked.items()):\n", + " img = na.plt.pcolormesh(\n", + " coordinates_scene.position.x,\n", + " coordinates_scene.position.y,\n", + " C=C,\n", + " ax=ax,\n", + " cmap=\"RdBu_r\",\n", + " vmin=-250,\n", + " vmax=+250,\n", + " )\n", + " ax.set_aspect(\"equal\")\n", + " ax.set_title(title)\n", + " ax.set_xlabel(f\"scene $x$ ({ax.get_xlabel()})\")\n", + " axs[0].set_ylabel(f\"scene $y$ ({axs[0].get_ylabel()})\")\n", + " plt.colorbar(img.ndarray.item(), ax=axs, label=\"velocity (km / s)\")" + ] + }, + { + "cell_type": "raw", + "id": "49", + "metadata": { + "editable": true, + "raw_mimetype": "text/x-rst", + "slideshow": { + "slide_type": "" + }, + "tags": [] + }, + "source": [ + "Quantify the agreement in those bright pixels." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "50", + "metadata": {}, + "outputs": [], + "source": [ + "a = masked[\"true median velocity\"].ndarray[bright]\n", + "b = masked[\"fitted velocity\"].ndarray[bright]\n", + "\n", + "f\"Pearson's r = {np.corrcoef(a, b)[0, 1]:.3f}\"" + ] + }, + { + "cell_type": "raw", + "id": "51", + "metadata": { + "editable": true, + "raw_mimetype": "text/x-rst", + "slideshow": { + "slide_type": "" + }, + "tags": [] + }, + "source": [ + "Finally, plot 2D histograms of the true vs. reconstructed total radiance, median\n", + "Doppler shift, and interquartile range of the line profile. The\n", + "``percentile_radiance`` argument excludes the faint pixels for the same reason." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "52", + "metadata": {}, + "outputs": [], + "source": [ + "inversion.plot_moments(scene, axis=\"wavelength\", percentile_radiance=90);" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.13.3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/pyproject.toml b/pyproject.toml index 6c1b3e6..2d4e9a8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,11 +23,16 @@ dependencies = [ dynamic = ["version"] [project.optional-dependencies] +torch = [ + "torch", +] test = [ "pytest", + "torch", ] doc = [ "pytest", + "torch", "matplotlib", "interface-region-imaging-spectrograph", "graphviz",