From 6c46be2ca48d6cda935b9d2213855bc4efdc7c17 Mon Sep 17 00:00:00 2001 From: Shane Maloney Date: Sun, 26 May 2024 17:46:30 +0100 Subject: [PATCH] Tighten up API, tidy docs, (#58) * Tighten and tidy API, docs page per module and better tests * Optonal args are now keyword only * Use of qantity aware type annotations * Doc refactor - page per submodule * Fix bug in transform code * Add fft equivalence test * Fix WCS v array indexing errors * Update mem and clean modules * Fix type error for VisMeta --- .pre-commit-config.yaml | 5 + changelog/58.breaking.rst | 2 + changelog/58.bugfix.rst | 1 + changelog/58.docs.rst | 1 + docs/conf.py | 8 +- docs/reference/clean.rst | 8 + docs/reference/imaging.rst | 8 + docs/reference/index.rst | 24 +- docs/reference/mem.rst | 8 + docs/reference/transform.rst | 8 + docs/reference/utils.rst | 8 + docs/reference/visibility.rst | 8 + docs/tutorials/introduction.rst | 13 +- examples/rhessi.py | 12 +- examples/stix.py | 13 +- setup.cfg | 3 + xrayvision/__init__.py | 8 +- xrayvision/clean.py | 219 +++++++---- xrayvision/imaging.py | 306 ++++++++------- xrayvision/mem.py | 110 +++--- xrayvision/tests/test_clean.py | 45 ++- xrayvision/tests/test_imaging.py | 200 ++++++---- xrayvision/tests/test_mem.py | 125 +------ xrayvision/tests/test_transform.py | 551 ++++++++++++++++------------ xrayvision/tests/test_visibility.py | 15 +- xrayvision/transform.py | 167 ++++++--- xrayvision/utils.py | 7 +- xrayvision/visibility.py | 207 +++++------ 28 files changed, 1203 insertions(+), 887 deletions(-) create mode 100644 changelog/58.breaking.rst create mode 100644 changelog/58.bugfix.rst create mode 100644 changelog/58.docs.rst create mode 100644 docs/reference/clean.rst create mode 100644 docs/reference/imaging.rst create mode 100644 docs/reference/mem.rst create mode 100644 docs/reference/transform.rst create mode 100644 docs/reference/utils.rst create mode 100644 docs/reference/visibility.rst diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index d16b96e..6feceb9 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -27,6 +27,11 @@ repos: exclude: ".*(.fits|.fts|.fit|.header|.txt|tca.*|.json)$|^CITATION.rst$" - id: mixed-line-ending exclude: ".*(.fits|.fts|.fit|.header|.txt|tca.*)$" + - repo: https://github.com/pre-commit/mirrors-mypy + rev: "v1.10.0" + hooks: + - id: mypy + additional_dependencies: [ "types-setuptools" ] - repo: https://github.com/codespell-project/codespell rev: v2.2.6 hooks: diff --git a/changelog/58.breaking.rst b/changelog/58.breaking.rst new file mode 100644 index 0000000..4840028 --- /dev/null +++ b/changelog/58.breaking.rst @@ -0,0 +1,2 @@ +Optional parameters are now keyword only for the :mod:`xrayvision.transform`, :mod:`xrayvision.imaging` and :mod:`xrayvision.visibility` modules. +Remove ``natural`` keyword in favour of ``scheme`` keyword which can be either 'natural' or 'uniform'. diff --git a/changelog/58.bugfix.rst b/changelog/58.bugfix.rst new file mode 100644 index 0000000..325b4ce --- /dev/null +++ b/changelog/58.bugfix.rst @@ -0,0 +1 @@ +Fix a bug where the x, y dimensions were not being treated consistently in :mod:`xrayvision.transform`. diff --git a/changelog/58.docs.rst b/changelog/58.docs.rst new file mode 100644 index 0000000..3499690 --- /dev/null +++ b/changelog/58.docs.rst @@ -0,0 +1 @@ +Add per module reference pages, switch to documenting types using type annotations. diff --git a/docs/conf.py b/docs/conf.py index 6f20ad1..28d1c4e 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -46,6 +46,9 @@ # Add any paths that contain templates here, relative to this directory. # templates_path = ['_templates'] +# Set automodapi to generate files inside the generated directory +automodapi_toctreedirnm = "generated/api" + # List of patterns, relative to source directory, that match files and # directories to ignore when looking for source files. # This pattern also affects html_static_path and html_extra_path. @@ -63,7 +66,7 @@ default_role = "obj" # Disable having a separate return type row -napoleon_use_rtype = False +napoleon_use_rtype = True # Disable google style docstrings napoleon_google_docstring = False @@ -71,6 +74,9 @@ # until sphinx-gallery / sphinx is fixed https://github.com/sphinx-doc/sphinx/issues/12300 suppress_warnings = ["config.cache"] +autodoc_typehints = "description" +autoclass_content = "init" + # -- Options for intersphinx extension --------------------------------------- # Example configuration for intersphinx: refer to the Python standard library. diff --git a/docs/reference/clean.rst b/docs/reference/clean.rst new file mode 100644 index 0000000..a3911f2 --- /dev/null +++ b/docs/reference/clean.rst @@ -0,0 +1,8 @@ +.. _clean: + +Clean ('xrayvision.clean') +************************** + +The ``clean`` submodule contains clean imaging methods. + +.. automodapi:: xrayvision.clean diff --git a/docs/reference/imaging.rst b/docs/reference/imaging.rst new file mode 100644 index 0000000..959cb52 --- /dev/null +++ b/docs/reference/imaging.rst @@ -0,0 +1,8 @@ +.. _imaging: + +Imaging ('xrayvision.imaging') +****************************** + +The ``imaging`` submodule contains functions to make map and images from visibilities. + +.. automodapi:: xrayvision.imaging diff --git a/docs/reference/index.rst b/docs/reference/index.rst index cd0ca4e..389d57b 100644 --- a/docs/reference/index.rst +++ b/docs/reference/index.rst @@ -1,23 +1,15 @@ -.. _reference: - ********* Reference ********* -.. automodapi:: xrayvision.visibility - -.. automodapi:: xrayvision.imaging - -.. automodapi:: xrayvision.transform - -.. automodapi:: xrayvision.clean - -.. automodapi:: xrayvision.mem - -.. automodapi:: xrayvision.utils - - .. toctree:: - :maxdepth: 2 + :maxdepth: 1 + + clean + imaging + mem + transform + utils + visibility ../whatsnew/index diff --git a/docs/reference/mem.rst b/docs/reference/mem.rst new file mode 100644 index 0000000..c6141ca --- /dev/null +++ b/docs/reference/mem.rst @@ -0,0 +1,8 @@ +.. _mem: + +MEM ('xrayvision.mem') +********************** + +The ``mem`` submodule contains the Maximum Entropy methods. + +.. automodapi:: xrayvision.mem diff --git a/docs/reference/transform.rst b/docs/reference/transform.rst new file mode 100644 index 0000000..fa0a3bd --- /dev/null +++ b/docs/reference/transform.rst @@ -0,0 +1,8 @@ +.. _transform: + +Transform ('xrayvision.transform') +********************************** + +The ``transform`` submodule forward and reverse transforms. + +.. automodapi:: xrayvision.transform diff --git a/docs/reference/utils.rst b/docs/reference/utils.rst new file mode 100644 index 0000000..08f8df0 --- /dev/null +++ b/docs/reference/utils.rst @@ -0,0 +1,8 @@ +.. _utils: + +Utils ('xrayvision.utils') +************************** + +The ``utils`` submodule contains utility functions. + +.. automodapi:: xrayvision.utils diff --git a/docs/reference/visibility.rst b/docs/reference/visibility.rst new file mode 100644 index 0000000..c3cf537 --- /dev/null +++ b/docs/reference/visibility.rst @@ -0,0 +1,8 @@ +.. _visibility: + +Visibility ('xrayvision.visibility') +************************************ + +The ``visibility`` submodule contains generic visibility classes. + +.. automodapi:: xrayvision.visibility diff --git a/docs/tutorials/introduction.rst b/docs/tutorials/introduction.rst index 5747c40..09ff152 100644 --- a/docs/tutorials/introduction.rst +++ b/docs/tutorials/introduction.rst @@ -41,10 +41,10 @@ measured visibilities in f can the original map a be recovered? data = make_data() - full_uv = transform.generate_uv(65) + full_uv = transform.generate_uv(65*u.pix) - uu = transform.generate_uv(33) - vv = transform.generate_uv(33) + uu = transform.generate_uv(33*u.pix) + vv = transform.generate_uv(33*u.pix) uu, vv = np.meshgrid(uu, vv) @@ -53,7 +53,7 @@ measured visibilities in f can the original map a be recovered? full_vis = transform.dft_map(data, u=uv[0,:], v=uv[1,:]) - res = transform.idft_map(full_vis, u=uv[0,:], v=uv[1,:], shape=(33, 33)) + res = transform.idft_map(full_vis, u=uv[0,:], v=uv[1,:], weights=1/33**2, shape=(33, 33)*u.pix) # assert np.allclose(data, res) # Generate log spaced radial u, v sampeling @@ -77,9 +77,10 @@ measured visibilities in f can the original map a be recovered? sub_vis = transform.dft_map(data, u=sub_uv[0,:], v=sub_uv[1,:]) psf1 = transform.idft_map(np.full(sub_vis.size, 1), u=sub_uv[0,:], v=sub_uv[1,:], - shape=(65, 65)) + weights=1/sub_vis.size, shape=(65, 65)*u.pix) - sub_res = transform.idft_map(sub_vis, u=sub_uv[0,:], v=sub_uv[1,:], shape=(65, 65)) + sub_res = transform.idft_map(sub_vis, u=sub_uv[0,:], v=sub_uv[1,:], + weights=1/sub_vis.size, shape=(65, 65)*u.pix) xp = np.round(x * 33 + 33/2 - 0.5 + 16).astype(int) yp = np.round(y * 33 + 33/2 - 0.5 + 16).astype(int) diff --git a/examples/rhessi.py b/examples/rhessi.py index 79f113c..bd75884 100644 --- a/examples/rhessi.py +++ b/examples/rhessi.py @@ -64,13 +64,13 @@ ############################################################################### # Lets have a look at the point spread function (PSF) or dirty beam -psf_map = vis_psf_map(vis, shape=(101, 101) * apu.pixel, pixel_size=1.5 * apu.arcsec, natural=False) +psf_map = vis_psf_map(vis, shape=(101, 101) * apu.pixel, pixel_size=1.5 * apu.arcsec / apu.pixel, scheme="uniform") ############################################################################### # We can now make an image using the back projection algorithm essentially and # inverse Fourier transform of the visibilities. -backproj_map = vis_to_map(vis, shape=[101, 101] * apu.pixel, pixel_size=1.5 * apu.arcsec) +backproj_map = vis_to_map(vis, shape=[101, 101] * apu.pixel, pixel_size=1.5 * apu.arcsec / apu.pix) ############################################################################### # Back projection contain many artifact due to the incomplete sampling of the u-v @@ -83,13 +83,17 @@ # v=vis_data_59['v']/apu.arcsec, offset=vis_data_59['xyoffset'][0]*apu.arcsec) clean_map, model_map, residual_map = vis_clean( - vis, shape=[101, 101] * apu.pixel, pixel=[1.5, 1.5] * apu.arcsec, clean_beam_width=10 * apu.arcsec, niter=100 + vis, + shape=[101, 101] * apu.pixel, + pixel_size=[1.5, 1.5] * apu.arcsec / apu.pix, + clean_beam_width=10 * apu.arcsec, + niter=100, ) ############################################################################### # MEM -mem_map = mem(vis, shape=[129, 129] * apu.pixel, pixel=[2, 2] * apu.arcsec) +mem_map = mem(vis, shape=[129, 129] * apu.pixel, pixel_size=[2, 2] * apu.arcsec / apu.pix) mem_map.plot() diff --git a/examples/stix.py b/examples/stix.py index 4762e67..c8511c3 100644 --- a/examples/stix.py +++ b/examples/stix.py @@ -26,32 +26,37 @@ stix_data = pickle.load(urllib.request.urlopen("https://pub099.cs.technik.fhnw.ch/demo/stix_vis.pkl")) time_range, energy_range, offset, stix_vis = stix_data +stix_vis.phase_centre = [0, 0] * apu.arcsec stix_vis.offset = offset ############################################################################### # Lets have a look at the point spread function (PSF) or dirty beam -psf_map = vis_psf_map(stix_vis, shape=(129, 129) * apu.pixel, pixel_size=2 * apu.arcsec, natural=False) +psf_map = vis_psf_map(stix_vis, shape=(129, 129) * apu.pixel, pixel_size=2 * apu.arcsec / apu.pix, scheme="uniform") psf_map.plot() ############################################################################### # Back projection -backproj_map = vis_to_map(stix_vis, shape=(129, 129) * apu.pixel, pixel_size=2 * apu.arcsec, natural=False) +backproj_map = vis_to_map(stix_vis, shape=(129, 129) * apu.pixel, pixel_size=2 * apu.arcsec / apu.pix, scheme="uniform") backproj_map.plot() ############################################################################### # Clean clean_map, model_map, resid_map = vis_clean( - stix_vis, shape=[129, 129] * apu.pixel, pixel=[2, 2] * apu.arcsec, niter=100, clean_beam_width=20 * apu.arcsec + stix_vis, + shape=[129, 129] * apu.pixel, + pixel_size=[2, 2] * apu.arcsec / apu.pix, + clean_beam_width=20 * apu.arcsec, + niter=100, ) clean_map.plot() ############################################################################### # MEM -mem_map = mem(stix_vis, shape=[129, 129] * apu.pixel, pixel=[2, 2] * apu.arcsec) +mem_map = mem(stix_vis, shape=[129, 129] * apu.pixel, pixel_size=[2, 2] * apu.arcsec / apu.pix) mem_map.plot() ############################################################################### diff --git a/setup.cfg b/setup.cfg index a5a5c9f..28af907 100644 --- a/setup.cfg +++ b/setup.cfg @@ -97,6 +97,9 @@ ignore-words-list = process, technik +[mypy] +disable_error_code = import-untyped + [coverage:run] omit = xrayvision/_sunpy_init* diff --git a/xrayvision/__init__.py b/xrayvision/__init__.py index 1623e23..612afc7 100644 --- a/xrayvision/__init__.py +++ b/xrayvision/__init__.py @@ -1,11 +1,7 @@ # Licensed under a 3-clause BSD style license - see LICENSE.rst try: - from xrayvision.version import __version__ + from xrayvision.version import __version__ # type: ignore except ImportError: __version__ = "unknown" -__all__ = [] - -from pkg_resources import resource_filename - -SAMPLE_RHESSI_VISIBILITIES = resource_filename("xrayvision", "data/hsi_visibili_20131028_0156_20131028_0200_6_12.fits") +__all__: list[str] = [] diff --git a/xrayvision/clean.py b/xrayvision/clean.py index f10ead0..32327a3 100644 --- a/xrayvision/clean.py +++ b/xrayvision/clean.py @@ -8,34 +8,41 @@ """ import logging +from typing import Union, Optional +from collections.abc import Iterable +import astropy.units as u import numpy as np from astropy.convolution import Gaussian2DKernel +from astropy.units import Quantity +from numpy.typing import NDArray from scipy import signal -from scipy.ndimage.interpolation import shift +from scipy.ndimage import shift from sunpy.map.map_factory import Map from xrayvision.imaging import vis_psf_image, vis_to_map +from xrayvision.visibility import Visibility __all__ = ["clean", "vis_clean", "ms_clean", "vis_ms_clean"] + logger = logging.getLogger(__name__) logger.setLevel("DEBUG") __common_clean_doc__ = r""" - clean_beam_width : `float` + clean_beam_width : The width of the gaussian to convolve the model with. If set to 0.0 \ the gaussian to convolution is disabled - gain : `float` + gain : The gain per loop or loop gain - thres : `float` + thres : Terminates clean when ``residual.max() <= thres`` - niter : `int` + niter : Maximum number of iterations to perform Returns ------- - `numpy.ndarray` + : The CLEAN image 2D Notes @@ -59,30 +66,39 @@ """ -def clean(dirty_map, dirty_beam, pixel=None, clean_beam_width=4.0, gain=0.1, thres=0.01, niter=5000): +@u.quantity_input +def clean( + dirty_map: Quantity, + dirty_beam: Quantity, + pixel_size: Quantity[u.arcsec / u.pix] = None, + clean_beam_width: Optional[Quantity[u.arcsec]] = 4.0 * u.arcsec, + gain: Optional[float] = 0.1, + thres: Optional[float] = None, + niter: int = 5000, +) -> Union[Quantity, NDArray[np.float64]]: r""" Clean the image using Hogbom's original method. CLEAN iteratively subtracts the PSF or dirty beam from the dirty map to create the residual. - At each iteration the location of the maximum residual is found and a shifted dirty beam is - subtracted that location updating the residual. This process continues until either `niter` - iterations is reached or the maximum residual <= `thres`. + At each iteration, the location of the maximum residual is found and a shifted dirty beam is + subtracted at that location. This process continues until either `niter` iterations are reached or + the maximum residual <= `thres`. Parameters ---------- - dirty_map : `numpy.ndarray` + dirty_map : The dirty map to be cleaned 2D - dirty_beam : `numpy.ndarray` + dirty_beam : The dirty beam or point spread function (PSF) 2D must - pixel : - Size of a pixel + pixel_size : + The pixel size in arcsec """ # Ensure both beam and map are even/odd on same axes # if not [x % 2 == 0 for x in dirty_map.shape] == [x % 2 == 0 for x in dirty_beam.shape]: # raise ValueError('') pad = [0 if x % 2 == 0 else 1 for x in dirty_map.shape] - # Assume beam, map center is in middle + # Assume beam, map phase_centre is in middle beam_center = (dirty_beam.shape[0] - 1) / 2.0, (dirty_beam.shape[1] - 1) / 2.0 map_center = (dirty_map.shape[0] - 1) / 2.0, (dirty_map.shape[1] - 1) / 2.0 @@ -104,7 +120,8 @@ def clean(dirty_map, dirty_beam, pixel=None, clean_beam_width=4.0, gain=0.1, thr # imax = imax * max_beam model[mx, my] += gain * imax - logger.debug(f"Iter: {i}, strength: {imax}, location: {mx, my}") + if i % 25 == 0: + logger.debug(f"Iter: {i}, strength: {imax}, location: {mx, my}") offset = map_center[0] - mx, map_center[1] - my shifted_beam_center = int(beam_center[0] + offset[0]), int(beam_center[1] + offset[1]) @@ -119,10 +136,11 @@ def clean(dirty_map, dirty_beam, pixel=None, clean_beam_width=4.0, gain=0.1, thr dirty_map = np.subtract(dirty_map, comp) - # if dirty_map.max() <= thres: - # logger.info("Threshold reached") - # break - # # el + if thres: + if np.abs(dirty_map).max() <= thres: + logger.info("Threshold reached") + break + if np.abs(dirty_map.min()) > dirty_map.max(): logger.info("Largest residual negative") break @@ -130,28 +148,40 @@ def clean(dirty_map, dirty_beam, pixel=None, clean_beam_width=4.0, gain=0.1, thr else: print("Max iterations reached") - if clean_beam_width != 0.0: + if clean_beam_width is not None: # Convert from FWHM to StDev FWHM = sigma*(8ln2)**0.5 = 2.3548200450309493 - x_stdev = (clean_beam_width / pixel[0] / 2.3548200450309493).value - y_stdev = (clean_beam_width / pixel[1] / 2.3548200450309493).value + x_stdev = (clean_beam_width / pixel_size[1]).to_value(u.pix) / 2.3548200450309493 + y_stdev = (clean_beam_width / pixel_size[0]).to_value(u.pix) / 2.3548200450309493 clean_beam = Gaussian2DKernel(x_stdev, y_stdev, x_size=dirty_beam.shape[1], y_size=dirty_beam.shape[0]).array # Normalise beam clean_beam = clean_beam / clean_beam.max() # Convolve clean beam with model and scale - clean_map = signal.convolve2d(model, clean_beam / clean_beam.sum(), mode="same") / (pixel[0] * pixel[1]) + clean_map = signal.convolve2d(model, clean_beam / clean_beam.sum(), mode="same") / ( + pixel_size[0].value * pixel_size[1].value + ) # Scale residual map with model and scale - dirty_map = dirty_map / clean_beam.sum() / (pixel[0] * pixel[1]) + dirty_map = dirty_map / clean_beam.sum() / (pixel_size[0].value * pixel_size[1].value) return clean_map + dirty_map, model, dirty_map return model + dirty_map, model, dirty_map -clean.__doc__ += __common_clean_doc__ +clean.__doc__ += __common_clean_doc__ # type: ignore -def vis_clean(vis, shape, pixel, clean_beam_width=4.0, niter=5000, map=True, gain=0.1, **kwargs): +@u.quantity_input +def vis_clean( + vis: Visibility, + shape: Quantity[u.pix], + pixel_size: Quantity[u.arcsec / u.pix], + clean_beam_width: Optional[Quantity[u.arcsec]] = 4.0, + niter: Optional[int] = 5000, + map: Optional[bool] = True, + gain: Optional[float] = 0.1, + **kwargs, +): r""" Clean the visibilities using Hogbom's original method. @@ -159,21 +189,26 @@ def vis_clean(vis, shape, pixel, clean_beam_width=4.0, niter=5000, map=True, gai Parameters ---------- - vis : `xrayvision.visibility.Visibly` + vis : The visibilities to clean shape : Size of map - pixel : - Size of pixel - map : `boolean` optional - Return an `sunpy.map.Map` by default or data only if `False` + pixel_size : + The pixel size in arcsec + map : + Return a `sunpy.map.Map` by default or array only if `False` """ - dirty_map = vis_to_map(vis, shape=shape, pixel_size=pixel, **kwargs) + dirty_map = vis_to_map(vis, shape=shape, pixel_size=pixel_size, **kwargs) dirty_beam_shape = [x.value * 3 + 1 if x.value * 3 % 2 == 0 else x.value * 3 for x in shape] * shape.unit - dirty_beam = vis_psf_image(vis, shape=dirty_beam_shape, pixel_size=pixel, **kwargs) + dirty_beam = vis_psf_image(vis, shape=dirty_beam_shape, pixel_size=pixel_size, **kwargs) clean_map, model, residual = clean( - dirty_map.data, dirty_beam.value, pixel=pixel, gain=gain, clean_beam_width=clean_beam_width, niter=niter + dirty_map.data, + dirty_beam.value, + pixel_size=pixel_size, + clean_beam_width=clean_beam_width, + gain=gain, + niter=niter, ) if not map: return clean_map, model, residual @@ -181,25 +216,24 @@ def vis_clean(vis, shape, pixel, clean_beam_width=4.0, niter=5000, map=True, gai return [Map((data, dirty_map.meta)) for data in (clean_map, model, residual)] -vis_clean.__doc__ += __common_clean_doc__ - +vis_clean.__doc__ += __common_clean_doc__ # type: ignore __common_ms_clean_doc__ = r""" scales : array-like, optional, optional The scales to use eg ``[1, 2, 4, 8]`` - clean_beam_width : `float` + clean_beam_width : The width of the gaussian to convolve the model with. If set to 0.0 the gaussian \ convolution is disabled - gain : `float` + gain : The gain per loop or loop gain - thres : `float` + thres : Terminates clean when `residuals.max() <= thres`` - niter : `int` + niter : Maximum number of iterations to perform Returns ------- - `numpy.ndarray` + : Cleaned image Notes @@ -219,20 +253,32 @@ def vis_clean(vis, shape, pixel, clean_beam_width=4.0, niter=5000, map=True, gai """ -def ms_clean(dirty_map, dirty_beam, pixel, scales=None, clean_beam_width=4.0, gain=0.1, thres=0.01, niter=5000): +@u.quantity_input +def ms_clean( + dirty_map: Quantity, + dirty_beam: Quantity, + pixel_size: Quantity[u.arcsec / u.pix], + scales: Union[Iterable, NDArray, None] = None, + clean_beam_width: Quantity = 4.0 * u.arcsec, + gain: float = 0.1, + thres: float = 0.01, + niter: int = 5000, +) -> Union[Quantity, NDArray[np.float64]]: r""" Clean the map using a multiscale clean algorithm. Parameters ---------- - dirty_map : `numpy.ndarray` + dirty_map : The 2D dirty map to be cleaned - dirty_beam : `numpy.ndarray` + dirty_beam : The 2D dirty beam should have the same dimensions as `dirty_map` + pixel_size : + The pixel size in arcsec """ # Compute the number of dyadic scales, their sizes and scale biases - number_of_scales = np.floor(np.log2(min(dirty_map.shape))).astype(int) - scale_sizes = 2 ** np.arange(number_of_scales) + number_of_scales: int = np.floor(np.log2(min(dirty_map.shape))).astype(int) + scale_sizes: NDArray[np.int_] = 2 ** np.arange(number_of_scales) if scales: scales = np.array(scales) @@ -258,7 +304,7 @@ def ms_clean(dirty_map, dirty_beam, pixel, scales=None, clean_beam_width=4.0, ga cross_terms = {} for i, scale in enumerate(scale_sizes): - scales[:, :, i] = component(scale=scale, shape=dirty_map.shape) + scales[:, :, i] = _component(scale=scale, shape=dirty_map.shape) scaled_residuals[:, :, i] = signal.convolve(dirty_map, scales[:, :, i], mode="same") scaled_dirty_beams[:, :, i] = signal.convolve(dirty_beam, scales[:, :, i], mode="same") max_scaled_dirty_beams[i] = scaled_dirty_beams[:, :, i].max() @@ -272,8 +318,11 @@ def ms_clean(dirty_map, dirty_beam, pixel, scales=None, clean_beam_width=4.0, ga # print(f'Clean loop {i}') # For each scale find the strength and location of max residual # Chose scale with has maximum strength - max_index = np.argmax(scaled_residuals) - max_x, max_y, max_scale = np.unravel_index(max_index, scaled_residuals.shape) + max_index: int = np.argmax(scaled_residuals).astype(int) + max_x: int + max_y: int + max_scale: int + max_x, max_y, max_scale = (int(x) for x in np.unravel_index(max_index, scaled_residuals.shape)) strength = scaled_residuals[max_x, max_y, max_scale] @@ -319,9 +368,9 @@ def ms_clean(dirty_map, dirty_beam, pixel, scales=None, clean_beam_width=4.0, ga scaled_residuals[:, :, j] = np.subtract(scaled_residuals[:, :, j], comp) # End max(res(a)) or niter - if scaled_residuals[:, :, max_scale].max() <= thres: + if np.abs(scaled_residuals[:, :, max_scale].max()) <= thres: logger.info("Threshold reached") - # break + break # Largest scales largest residual is negative if np.abs(scaled_residuals[:, :, 0].min()) > scaled_residuals[:, :, 0].max(): @@ -331,29 +380,39 @@ def ms_clean(dirty_map, dirty_beam, pixel, scales=None, clean_beam_width=4.0, ga else: logger.info("Max iterations reached") - # Convolve model with clean beam B_G * I^M - if clean_beam_width != 0.0: - x_stdev = (clean_beam_width / pixel[0] / (2.0 * np.sqrt(2.0 * np.log(2.0)))).value - y_stdev = (clean_beam_width / pixel[1] / (2.0 * np.sqrt(2.0 * np.log(2.0)))).value + # Convolve model with clean beam B_G * I^M + if clean_beam_width is not None: + x_stdev = ((clean_beam_width / pixel_size[0]).to_value(u.pix) / (2.0 * np.sqrt(2.0 * np.log(2.0)))).value + y_stdev = ((clean_beam_width / pixel_size[1]).to_value(u.pix) / (2.0 * np.sqrt(2.0 * np.log(2.0)))).value clean_beam = Gaussian2DKernel(x_stdev, y_stdev, x_size=dirty_beam.shape[1], y_size=dirty_beam.shape[0]).array # Normalise beam clean_beam = clean_beam / clean_beam.max() - clean_map = signal.convolve2d(model, clean_beam, mode="same") / (pixel[0] * pixel[1]) + clean_map = signal.convolve2d(model, clean_beam, mode="same") / (pixel_size[0] * pixel_size[1]) # Scale residual map with model and scale - dirty_map = (scaled_residuals / clean_beam.sum() / (pixel[0] * pixel[1])).sum(axis=2) + dirty_map = (scaled_residuals / clean_beam.sum() / (pixel_size[0] * pixel_size[1])).sum(axis=2) return clean_map + dirty_map, model, dirty_map # Add residuals B_G * I^M + I^R return model, scaled_residuals.sum(axis=2) -ms_clean.__doc__ += __common_ms_clean_doc__ +ms_clean.__doc__ += __common_ms_clean_doc__ # type: ignore -def vis_ms_clean(vis, shape, pixel, scales=None, clean_beam_width=4.0, gain=0.1, thres=0.01, niter=5000, map=True): +def vis_ms_clean( + vis: Visibility, + shape: Quantity[u.pix], + pixel_size: Quantity[u.arcsec / u.pix], + scales: Optional[Iterable], + clean_beam_width: Optional[Quantity[u.arcsec]] = 4.0, + niter: Optional[int] = 5000, + map: Optional[bool] = True, + gain: Optional[float] = 0.1, + thres: Optional[float] = 0.01, +) -> Union[Quantity, NDArray[np.float64]]: r""" Clean the visibilities using a multiscale clean method. @@ -361,20 +420,38 @@ def vis_ms_clean(vis, shape, pixel, scales=None, clean_beam_width=4.0, gain=0.1, Parameters ---------- - vis : `xrayvision.visibility.Visibly` + vis : The visibilities to clean shape : Size of map - pixel : - Size of pixel + pixel_size : + The pixel size in arcsec + scales : array-like, optional, optional + The scales to use eg ``[1, 2, 4, 8]`` + clean_beam_width : + The width of the gaussian to convolve the model with. If set to 0.0 the gaussian \ + convolution is disabled + gain : + The gain per loop or loop gain + thres : + Terminates clean when `residuals.max() <= thres`` + niter : + Maximum number of iterations to perform + map : + Return a `sunpy.map.Map` by default or array only if `False` + + Returns + ------- + : + Cleaned image """ - dirty_map = vis_to_map(vis, shape=shape, pixel_size=pixel) - dirty_beam = vis_psf_image(vis, shape=shape * 3, pixel_size=pixel) + dirty_map = vis_to_map(vis, shape=shape, pixel_size=pixel_size) + dirty_beam = vis_psf_image(vis, shape=shape * 3, pixel_size=pixel_size) clean_map, model, residual = ms_clean( dirty_map.data, dirty_beam, - pixel, + pixel_size=pixel_size, scales=scales, clean_beam_width=clean_beam_width, gain=gain, @@ -390,7 +467,7 @@ def vis_ms_clean(vis, shape, pixel, scales=None, clean_beam_width=4.0, gain=0.1, # vis_ms_clean.__doc__ += __common_ms_clean_doc__ -def radial_prolate_sphereoidal(nu): +def _radial_prolate_sphereoidal(nu): r""" Calculate prolate spheroidal wave function approximation. @@ -467,7 +544,7 @@ def radial_prolate_sphereoidal(nu): return 0 -def vec_radial_prolate_sphereoidal(nu): +def _vec_radial_prolate_sphereoidal(nu): r""" Calculate prolate spheroidal wave function approximation. @@ -546,7 +623,7 @@ def vec_radial_prolate_sphereoidal(nu): return out -def component(scale, shape): +def _component(scale, shape): r""" Parameters @@ -579,9 +656,9 @@ def component(scale, shape): rad_zeros_indices = radii_squared <= 0.0 amp_zero_indices = radii_squared >= 1.0 - wave_amp = vec_radial_prolate_sphereoidal(np.sqrt(radii_squared.reshape(radii_squared.size))) + wave_amp = _vec_radial_prolate_sphereoidal(np.sqrt(radii_squared.reshape(radii_squared.size))) wave_amp = wave_amp.reshape(shape) - wave_amp[rad_zeros_indices] = vec_radial_prolate_sphereoidal([0])[0] + wave_amp[rad_zeros_indices] = _vec_radial_prolate_sphereoidal([0])[0] wave_amp = wave_amp * (1 - radii_squared) diff --git a/xrayvision/imaging.py b/xrayvision/imaging.py index d1559d7..c304266 100644 --- a/xrayvision/imaging.py +++ b/xrayvision/imaging.py @@ -1,6 +1,9 @@ +from typing import Optional + import astropy.units as apu import numpy as np -from sunpy.map import Map +from astropy.units import Quantity +from sunpy.map import GenericMap, Map from xrayvision.transform import dft_map, idft_map from xrayvision.visibility import Visibility @@ -18,30 +21,32 @@ ] ANGLE = apu.get_physical_type(apu.deg) +WEIGHT_SCHEMES = ("natural", "uniform") -def get_weights(vis, natural=True, norm=True): +def get_weights(vis: Visibility, scheme: str = "natural", norm: bool = True) -> np.ndarray: r""" - Return natural spatial frequency weight factor for each visibility. + Return spatial frequency weight factors for each visibility. Defaults to use natural weighting scheme given by `(vis.u**2 + vis.v**2)^{1/2}`. Parameters ---------- - vis : `xrayvision.visibility.Visibility` + vis : Input visibilities - natural : `boolean` optional - Use natural or uniform weighting scheme - norm : `boolean` - Normalise the weighs before returning + scheme : + Weighting scheme to use, defaults to natural + norm : + Normalise the weighs before returning, defaults to True. Returns ------- - `weights` """ + if scheme not in WEIGHT_SCHEMES: + raise ValueError(f"Invalid weighting scheme {scheme}, must be one of: {WEIGHT_SCHEMES}") weights = np.sqrt(vis.u**2 + vis.v**2).value - if natural: + if scheme == "natural": weights = np.ones_like(vis.vis, dtype=float) if norm: @@ -50,7 +55,8 @@ def get_weights(vis, natural=True, norm=True): return weights -def validate_and_expand_kwarg(q, name=""): +@apu.quantity_input() +def validate_and_expand_kwarg(q: Quantity, name: Optional[str] = "") -> Quantity: r""" Expand a scalar or array of size one to size two by repeating. @@ -85,223 +91,243 @@ def validate_and_expand_kwarg(q, name=""): return q -@apu.quantity_input(shape=apu.pixel, pixel_size="angle") -def vis_psf_image(vis, *, shape=(65, 65) * apu.pixel, pixel_size=2 * apu.arcsec, natural=True): +@apu.quantity_input +def image_to_vis( + image: Quantity, + *, + u: Quantity[apu.arcsec**-1], + v: Quantity[apu.arcsec**-1], + phase_centre: Optional[Quantity[apu.arcsec]] = (0.0, 0.0) * apu.arcsec, + pixel_size: Optional[Quantity[apu.arcsec / apu.pix]] = 1.0 * apu.arcsec / apu.pix, +) -> Visibility: + r""" + Return a Visibility created from the image and u, v sampling. + + Parameters + ---------- + image : + The 2D input image + u : + Array of u coordinates where the visibilities will be evaluated + v : + Array of v coordinates where the visibilities will be evaluated + phase_centre : + The coordinates the phase_centre. + pixel_size : + Size of pixels, if only one value is passed, assume square pixels (repeating the value). + + Returns + ------- + : + The new visibility object + """ - Create the point spread function for given u, v point of the visibilities. + pixel_size = validate_and_expand_kwarg(pixel_size, "pixel_size") + if not (apu.get_physical_type((1 / u).unit) == ANGLE and apu.get_physical_type((1 / v).unit) == ANGLE): + raise ValueError("u and v must be inverse angle (e.g. 1/deg or 1/arcsec") + vis = dft_map(image, u=u, v=v, phase_centre=phase_centre, pixel_size=pixel_size) + return Visibility(vis, u=u, v=v, offset=phase_centre) + + +@apu.quantity_input() +def vis_to_image( + vis: Visibility, + shape: Quantity[apu.pix] = (65, 65) * apu.pixel, + pixel_size: Optional[Quantity[apu.arcsec / apu.pix]] = 1 * apu.arcsec / apu.pix, + scheme: str = "natural", +) -> Quantity: + """ + Create an image by 'back projecting' the given visibilities onto the sky. Parameters ---------- - vis : `xrayvision.visibility.Visibility` + vis : Input visibilities - shape : `astropy.units.quantity.Quantity`, optional - Shape of the image, if only one value is given assume square (repeating the value). - pixel_size : `astropy.units.quantity.Quantity`, optional - Size of pixels, if only one value is given assume square pixels (repeating the value). - natural : `boolean` optional - Weight scheme use natural by default, uniform if `False` + shape : + Shape of the image, if only one value is passed assume square (repeating the value). + pixel_size : + Size of pixels, if only one value is passed assume square pixels (repeating the value). + scheme : + Weight scheme natural by default. Returns ------- - `astropy.units.quantity.Quantity` - Point spread function + `~astropy.units.Quantity` + Back projection image """ shape = validate_and_expand_kwarg(shape, "shape") pixel_size = validate_and_expand_kwarg(pixel_size, "pixel_size") - shape = shape.to_value(apu.pixel) - weights = get_weights(vis, natural=natural) - - # Make sure psf is always odd so power is in exactly one pixel - m, n = (s // 2 * 2 + 1 for s in shape) - psf_arr = idft_map( - np.ones(vis.vis.shape) * vis.vis.unit, u=vis.u, v=vis.v, shape=(m, n), weights=weights, pixel_size=pixel_size + shape = shape.to(apu.pixel) + weights = get_weights(vis, scheme=scheme) + bp_arr = idft_map( + vis.vis, u=vis.u, v=vis.v, shape=shape, weights=weights, pixel_size=pixel_size, phase_centre=vis.phase_centre ) - return psf_arr + + return bp_arr -@apu.quantity_input(shape=apu.pixel, pixel_size="angle") -def vis_psf_map(vis, *, shape=(65, 65) * apu.pixel, pixel_size=2 * apu.arcsec, natural=True): +@apu.quantity_input +def vis_psf_map( + vis: Visibility, + *, + shape: Quantity[apu.pix] = (65, 65) * apu.pixel, + pixel_size: Optional[Quantity[apu.arcsec / apu.pix]] = 1 * apu.arcsec / apu.pix, + scheme: Optional[str] = "natural", +) -> GenericMap: r""" Create a map of the point spread function for given the visibilities. Parameters ---------- - vis : `xrayvision.visibility.Visibility` + vis : Input visibilities - shape : `astropy.units.quantity.Quantity`, optional - Shape of the image, if only one value is given assume square (repeating the value). - pixel_size : `Tuple[~astropy.units.Quantity]`, optional - Size of pixels, if only one value is given assume square pixels (repeating the value). - natural : `boolean` optional - Weight scheme use natural by default, uniform if `False` + shape : + Shape of the image, if only one value is passed, assume square (repeating the value). + pixel_size : + Size of pixels, if only one value is passed, assume square pixels (repeating the value). + scheme : + Weight scheme to use natural by default. Returns ------- - + : + Map of the point spread function """ shape = validate_and_expand_kwarg(shape, "shape") pixel_size = validate_and_expand_kwarg(pixel_size, "pixel_size") header = generate_header(vis, shape=shape, pixel_size=pixel_size) - psf = vis_psf_image(vis, shape=shape, pixel_size=pixel_size, natural=natural) + psf = vis_psf_image(vis, shape=shape, pixel_size=pixel_size, scheme=scheme) return Map((psf, header)) -@apu.quantity_input(shape=apu.pixel, pixel_size="angle") -def vis_to_image(vis, shape=(65, 65) * apu.pixel, pixel_size=2 * apu.arcsec, natural=True): +@apu.quantity_input() +def vis_psf_image( + vis: Visibility, + *, + shape: Quantity[apu.pix] = (65, 65) * apu.pixel, + pixel_size: Quantity[apu.arcsec / apu.pix] = 1 * apu.arcsec / apu.pix, + scheme: str = "natural", +) -> Quantity: """ - Create an image by 'back projecting' the given visibilities onto the sky. + Create the point spread function for given u, v point of the visibilities. Parameters ---------- - vis : `xrayvision.visibility.Visibility` + vis : Input visibilities - shape : `~astropy.units.Quantity` - Shape of the image, if only one value is given assume square (repeating the value). - pixel_size : `~astropy.units.Quantity` - Size of pixels, if only one value is given assume square pixels (repeating the value). - natural : `boolean` optional - Weight scheme use natural by default, uniform if `False` + shape : + Shape of the image, if only one value passed, assume square (repeating the value). + pixel_size : + Size of pixels, if only one value is passed, assume square pixels (repeating the value). + scheme : + Weight scheme, natural by default. Returns ------- - `~astropy.units.Quantity` - Back projection image + : + Point spread function """ shape = validate_and_expand_kwarg(shape, "shape") pixel_size = validate_and_expand_kwarg(pixel_size, "pixel_size") - shape = shape.to_value(apu.pixel) - weights = get_weights(vis, natural=natural) - bp_arr = idft_map(vis.vis, u=vis.u, v=vis.v, shape=shape, weights=weights, pixel_size=pixel_size, center=vis.center) + shape = shape.to(apu.pixel) + weights = get_weights(vis, scheme=scheme) - return bp_arr + # Make sure psf is always odd so power is in exactly one pixel + shape = [s // 2 * 2 + 1 for s in shape.to_value(apu.pix)] * shape.unit + psf_arr = idft_map( + np.ones(vis.vis.shape) * vis.vis.unit, u=vis.u, v=vis.v, shape=shape, weights=weights, pixel_size=pixel_size + ) + return psf_arr -@apu.quantity_input(shape=apu.pixel, pixel_size="angle") -def vis_to_map(vis, shape=(65, 65) * apu.pixel, pixel_size=2 * apu.arcsec, natural=True): +@apu.quantity_input() +def vis_to_map( + vis: Visibility, + shape: Quantity[apu.pix] = (65, 65) * apu.pixel, + pixel_size: Optional[Quantity[apu.arcsec / apu.pix]] = 1 * apu.arcsec / apu.pixel, + scheme: Optional[str] = "natural", +) -> GenericMap: r""" Create a map by performing a back projection of inverse transform on the visibilities. Parameters ---------- - vis : `xrayvision.visibility.Visibility` + vis : Input visibilities - shape : `int` (m, n) - Shape of the image, if only one value is given assume square (repeating the value). - pixel_size : `float` (dx, dy), optional - Size of pixels, if only one value is given assume square pixels (repeating the value). - natural : `boolean` optional - Weight scheme use natural by default, uniform if `False`. + shape : + Shape of the image, if only one value is passed, assume square (repeating the value). + pixel_size : + Size of pixels, if only one value is passed, assume square pixels (repeating the value). + scheme : + Weight scheme natural by default. Returns ------- - `sunpy.map.Map` - Map object with the map created from the visibilities and the metadata will contain the - offset and the pixel size + : + Map object with the map created from the visibilities and the metadata will contain the offset and the pixel size """ shape = validate_and_expand_kwarg(shape, "shape") pixel_size = validate_and_expand_kwarg(pixel_size, "pixel_size") header = generate_header(vis, shape=shape, pixel_size=pixel_size) - image = vis_to_image(vis, shape=shape, pixel_size=pixel_size, natural=natural) + image = vis_to_image(vis, shape=shape, pixel_size=pixel_size, scheme=scheme) return Map((image, header)) -def generate_header(vis, *, shape, pixel_size): +@apu.quantity_input() +def generate_header(vis: Visibility, *, shape: Quantity[apu.pix], pixel_size: Quantity[apu.arcsec / apu.pix]) -> dict: r""" Generate a map head given the visibilities, pixel size and shape Parameters ---------- vis : - shape : `~astropy.units.Quantity` - Shape of the image, if only one value is given assume square (repeating the value). - pixel_size : `~astropy.units.Quantity` - Size of pixels, if only one value is given assume square pixels (repeating the value) - offset + Input visibilities. + shape : + Shape of the image, if only one value is passed assume square (repeating the value). + pixel_size : + Size of pixels, if only one value is passed assume square pixels (repeating the value) Returns ------- - + : """ header = { - "crval1": (vis.offset[0]).to_value(apu.arcsec), - "crval2": (vis.offset[1]).to_value(apu.arcsec), - "cdelt1": pixel_size[0].to_value(apu.arcsec), - "cdelt2": pixel_size[1].to_value(apu.arcsec), + "crval1": (vis.offset[1]).to_value(apu.arcsec), + "crval2": (vis.offset[0]).to_value(apu.arcsec), + "cdelt1": (pixel_size[1] * apu.pix).to_value(apu.arcsec), + "cdelt2": (pixel_size[0] * apu.pix).to_value(apu.arcsec), "ctype1": "HPLN-TAN", "ctype2": "HPLT-TAN", "naxis": 2, - "naxis1": shape[0].value, - "naxis2": shape[1].value, + "naxis1": shape[1].value, + "naxis2": shape[0].value, "cunit1": "arcsec", "cunit2": "arcsec", } - # if vis.center is not None: - # header['crval1'] = vis.center[0].value - # header['crval2'] = vis.center[1].value - if pixel_size is not None: - if pixel_size.ndim == 0: - header["cdelt1"] = pixel_size.value - header["cdelt2"] = pixel_size.value - elif pixel_size.ndim == 1 and pixel_size.size == 2: - header["cdelt1"] = pixel_size[0].value - header["cdelt2"] = pixel_size[1].value - else: - raise ValueError(f"pixel_size can have a length of 1 or 2 not {pixel_size.shape}") return header -@apu.quantity_input(center="angle", pixel_size="angle") -def image_to_vis(image, *, u, v, center=(0.0, 0.0) * apu.arcsec, pixel_size=2.0 * apu.arcsec): - r""" - Return a Visibility created from the image and u, v sampling. - - Parameters - ---------- - image : `numpy.ndarray` - The 2D input image - u : `astropy.units.Quantity` - Array of u coordinates where the visibilities will be evaluated - v : `astropy.units.Quantity` - Array of v coordinates where the visibilities will be evaluated - center : `~astropy.units.Quantity` (x, y) - The coordinates of the center of the image. - pixel_size : `~astropy.units.Quantity` - Size of pixels, if only one value is given assume square pixels (repeating the value). - - Returns - ------- - `Visibility` - - The new visibility object - - """ - pixel_size = validate_and_expand_kwarg(pixel_size, "pixel_size") - if not (apu.get_physical_type((1 / u).unit) == ANGLE and apu.get_physical_type((1 / v).unit) == ANGLE): - raise ValueError("u and v must be inverse angle (e.g. 1/deg or 1/arcsec") - vis = dft_map(image, u=u, v=v, center=center, pixel_size=pixel_size) - return Visibility(vis, u=u, v=v, center=center) - - -def map_to_vis(amap, *, u, v): +@apu.quantity_input() +def map_to_vis(amap: GenericMap, *, u: Quantity[1 / apu.arcsec], v: Quantity[1 / apu.arcsec]) -> Visibility: r""" - Return a Visibility object created from the map and u, v sampling. + Return a Visibility object created from the map, sampling it at give `u`, `v` coordinates. Parameters ---------- - amap : `sunpy.map.Map` - The input map - u : `numpy.ndarray` + amap : + Input map. + u : Array of u coordinates where the visibilities will be evaluated - v : `numpy.ndarray` + v : Array of v coordinates where the visibilities will be evaluated Returns ------- - `Visibility` + : The new visibility object """ @@ -310,16 +336,16 @@ def map_to_vis(amap, *, u, v): meta = amap.meta new_pos = np.array([0.0, 0.0]) if "crval1" in meta: - new_pos[0] = float(meta["crval1"]) + new_pos[1] = float(meta["crval1"]) if "crval2" in meta: - new_pos[1] = float(meta["crval2"]) + new_pos[0] = float(meta["crval2"]) new_psize = np.array([1.0, 1.0]) if "cdelt1" in meta: - new_psize[0] = float(meta["cdelt1"]) + new_psize[1] = float(meta["cdelt1"]) if "cdelt2" in meta: - new_psize[1] = float(meta["cdelt2"]) + new_psize[0] = float(meta["cdelt2"]) - vis = image_to_vis(amap.data, u=u, v=v, pixel_size=new_psize * apu.arcsec) + vis = image_to_vis(amap.data, u=u, v=v, pixel_size=new_psize * apu.arcsec / apu.pix) vis.offset = new_pos * apu.arcsec return vis diff --git a/xrayvision/mem.py b/xrayvision/mem.py index 1ee4916..485e86f 100644 --- a/xrayvision/mem.py +++ b/xrayvision/mem.py @@ -3,10 +3,13 @@ """ from types import SimpleNamespace +from typing import Union, Optional import astropy.units as apu import numpy as np +from astropy.units import Quantity from numpy.linalg import norm +from numpy.typing import NDArray from sunpy.map import Map from xrayvision.imaging import generate_header @@ -14,21 +17,22 @@ from xrayvision.utils import get_logger __all__ = [ - "get_entropy", - "get_fourier_matrix", - "estimate_flux", - "get_mean_visibilities", - "proximal_entropy", - "proximal_operator", - "optimise_fb", + "_get_entropy", + "_get_fourier_matrix", + "_estimate_flux", + "_get_mean_visibilities", + "_proximal_entropy", + "_proximal_operator", + "_optimise_fb", "mem", ] +from xrayvision.visibility import Visibility logger = get_logger(__name__, "DEBUG") -def get_entropy(image, flux): +def _get_entropy(image, flux): r""" Return the entropy of an image. @@ -55,7 +59,7 @@ def get_entropy(image, flux): return np.sum(image * np.log(image / (flux * np.e))) -def get_fourier_matrix(vis, shape=[64, 64] * apu.pix, pixel_size=[4.0312500, 4.0312500] * apu.arcsec): +def _get_fourier_matrix(vis, shape=(64, 64) * apu.pix, pixel_size=(4.0312500, 4.0312500) * apu.arcsec): r""" Return the complex Fourier matrix used to compute the value of the visibilities. @@ -76,10 +80,10 @@ def get_fourier_matrix(vis, shape=[64, 64] * apu.pix, pixel_size=[4.0312500, 4.0 ------- The complex Fourier matrix """ - m, n = shape.to_value("pix") - y = generate_xy(m, 0 * apu.arcsec, pixel_size[1]) - x = generate_xy(n, 0 * apu.arcsec, pixel_size[0]) - x, y = np.meshgrid(x, y) + m, n = shape.to("pix") + y = generate_xy(m, phase_centre=0 * apu.arcsec, pixel_size=pixel_size[1]) + x = generate_xy(n, phase_centre=0 * apu.arcsec, pixel_size=pixel_size[0]) + x, y = np.meshgrid(x, y, indexing="ij") uv = np.vstack([vis.u, vis.v]) # Check apu are correct for exp need to be dimensionless and then remove apu for speed if (vis.u * x[0, 0]).unit == apu.dimensionless_unscaled and (vis.v * y[0, 0]).unit == apu.dimensionless_unscaled: @@ -91,10 +95,10 @@ def get_fourier_matrix(vis, shape=[64, 64] * apu.pix, pixel_size=[4.0312500, 4.0 1j * 2 * np.pi * (x[..., np.newaxis] * uv[np.newaxis, 0, :] + y[..., np.newaxis] * uv[np.newaxis, 1, :]) ) - return Hv * pixel_size[0] * pixel_size[1] + return Hv * pixel_size[0].value * pixel_size[1].value -def estimate_flux(vis, shape, pixel, maxiter=1000, tol=1e-3): +def _estimate_flux(vis, shape, pixel, maxiter=1000, tol=1e-3): r""" Estimate the total flux in the image by solving an optimisation problem. @@ -151,7 +155,9 @@ def estimate_flux(vis, shape, pixel, maxiter=1000, tol=1e-3): diff_V = Hvx - Visib chi2 = (diff_V**2.0).sum() - logger.info(f"Iter: {i}, Chi2: {chi2}") + if i % 25 == 0: + logger.info(f"Iter: {i}, Chi2: {chi2}") + if np.sqrt(((x - x_old) ** 2.0).sum()) < tol * np.sqrt((x_old**2.0).sum()): break @@ -179,7 +185,7 @@ def _prepare_for_optimise(pixel, shape, vis): ------- """ - Hv = get_fourier_matrix(vis, shape, pixel) + Hv = _get_fourier_matrix(vis, shape, pixel) # Division of real and imaginary part of the matrix 'Hv' ReHv = np.real(Hv) ImHv = np.imag(Hv) @@ -210,7 +216,7 @@ def _prepare_for_optimise(pixel, shape, vis): return Hv, Lip, Visib -def get_mean_visibilities(vis, shape, pixel): +def _get_mean_visibilities(vis, shape, pixel): r""" Return the mean visibilities sampling the same call in the discretisation of the (u,v) plane. @@ -237,9 +243,9 @@ def get_mean_visibilities(vis, shape, pixel): rv = np.around(iv) # index of the u coordinates of the sampling frequencies in the discretisation of the u axis - ru = ru + imsize2 + ru = ru * apu.pix + imsize2.to(apu.pix) # index of the v coordinates of the sampling frequencies in the discretisation of the v axis - rv = rv + imsize2 + rv = rv * apu.pix + imsize2.to(apu.pix) # matrix that represents the discretization of the (u,v)-plane iuarr = np.zeros(shape.to_value("pixel").astype(int)) @@ -293,7 +299,7 @@ def get_mean_visibilities(vis, shape, pixel): return SimpleNamespace(u=u, v=v, vis=visib, amplitude_error=weights) -def proximal_entropy(y, m, lamba, Lip, tol=10**-10): +def _proximal_entropy(y, m, lamba, Lip, tol=10**-10): r""" This function computes the value of the proximity operator of the entropy function subject to positivity constraint, i.e. it solves the problem @@ -335,7 +341,7 @@ def proximal_entropy(y, m, lamba, Lip, tol=10**-10): return c -def proximal_operator(z, f, m, lamb, Lip, niter=250): +def _proximal_operator(z, f, m, lamb, Lip, niter=250): r""" Computes the value of the proximity operator of the entropy function subject to positivity constraint and flux constraint by means of a Dykstra-like proximal algorithm @@ -371,7 +377,7 @@ def proximal_operator(z, f, m, lamb, Lip, niter=250): y = tmp + (f - tmp.sum()) / tmp.size p = x + p - y - x = proximal_entropy(y + q, m, lamb, Lip) + x = _proximal_entropy(y + q, m, lamb, Lip) if np.abs(x.sum() - f) <= 0.01 * f: break @@ -381,7 +387,7 @@ def proximal_operator(z, f, m, lamb, Lip, niter=250): return x, i -def optimise_fb(Hv, Visib, Lip, flux, lambd, shape, pixel, maxiter, tol): +def _optimise_fb(Hv, Visib, Lip, flux, lambd, shape, pixel, maxiter, tol): r""" Solve the optimization problem using a forward-backward splitting algorithm @@ -441,7 +447,7 @@ def optimise_fb(Hv, Visib, Lip, flux, lambd, shape, pixel, maxiter, tol): tmp = x.flatten()[:] Hvx = np.matmul(Hv, tmp) - f_R = get_entropy(x, m) + f_R = _get_entropy(x, m) diff_V = Hvx - Visib f_0 = (diff_V**2).sum() @@ -458,12 +464,12 @@ def optimise_fb(Hv, Visib, Lip, flux, lambd, shape, pixel, maxiter, tol): y = z - 1 / Lip * grad # PROXIMAL STEP - p, pi = proximal_operator(y, f, m, lambd, Lip) + p, pi = _proximal_operator(y, f, m, lambd, Lip) # COMPUTATION OF THE OBJECTIVE FUNCTION 'Jp' IN 'p' tmp = p.flatten() Hvp = np.matmul(Hv, tmp) - f_Rp = get_entropy(p, m) + f_Rp = _get_entropy(p, m) diff_Vp = Hvp - Visib f_0 = (diff_Vp**2).sum() @@ -491,7 +497,8 @@ def optimise_fb(Hv, Visib, Lip, flux, lambd, shape, pixel, maxiter, tol): tau = (t_old - 1.0) / t z = x + tau * (x - x_old) + (t_old / t) * (p - x) - logger.info(f"Iter: {i}, Obj function: {J}") + if i % 25 == 0: + logger.info(f"Iter: {i}, Obj function: {J}") if check and (np.sqrt(((x - x_old) ** 2.0).sum()) < tol * np.sqrt((x_old**2).sum())): break @@ -558,7 +565,7 @@ def resistant_mean(data, sigma_cut): return mean, sigma -def get_percent_lambda(vis): +def _get_percent_lambda(vis): r""" Return 'percent_lambda' use with MEM @@ -602,52 +609,63 @@ def get_percent_lambda(vis): return percent_lambda -def mem(vis, percent_lambda=None, shape=None, pixel=None, maxiter=1000, tol=1e-3, map=True): +@apu.quantity_input +def mem( + vis: Visibility, + shape: Quantity[apu.pix], + pixel_size: Quantity[apu.arcsec / apu.pix], + *, + percent_lambda: Optional[Quantity[apu.percent]] = None, + maxiter: int = 1000, + tol: float = 1e-3, + map: bool = True, +) -> Union[Quantity, NDArray[np.float64]]: r""" - Maximum Entropy Method for visibility based image reconstruction + Maximum Entropy Method visibility based image reconstruction Parameters ---------- vis : Input Visibilities - percent_lambda - value used to compute the regularization parameter as a percentage of a maximum value - automatically overestimated by the algorithm. Must be in the range [0.0001,0.2] - shape + shape : Image size - pixel + pixel_size : Pixel size - maxiter : int - Maximum number of iterations of the optimization loop + percent_lambda + Value used to compute the regularization parameter as a percentage of a maximum value + automatically overestimated by the algorithm. Must be in the range [0.0001,0.2] + maxiter : + Maximum number of iterations of the optimisation loop tol : float tolerance value used in the stopping rule ( || x - x_old || <= tol || x_old ||) - + map : + Return a sunpy map or bare array Returns ------- """ - total_flux = estimate_flux(vis, shape, pixel) + total_flux = _estimate_flux(vis, shape, pixel_size) if percent_lambda is None: - percent_lambda = get_percent_lambda(vis) + percent_lambda = _get_percent_lambda(vis) - mean_vis = get_mean_visibilities(vis, shape, pixel) - Hv, Lip, Visib = _prepare_for_optimise(pixel, shape, mean_vis) + mean_vis = _get_mean_visibilities(vis, shape, pixel_size) + Hv, Lip, Visib = _prepare_for_optimise(pixel_size, shape, mean_vis) # should have same apu as image ct/ keV cm s arcsec**2 x = np.zeros(shape.to_value("pix").astype(int)).flatten() # COMPUTATION OF THE; OBJECTIVE; FUNCTION; 'chi2' - x = x + total_flux / (shape[0] * shape[1] * pixel[0] * pixel[1]).value + x = x + total_flux / (shape[0] * shape[1] * pixel_size[0] * pixel_size[1]).value Hvx = np.matmul(Hv, x) lambd = 2 * np.abs(np.matmul((Hvx.value - Visib), Hv)).max() * percent_lambda - im = optimise_fb(Hv, Visib, Lip, total_flux, lambd, shape, pixel, maxiter, tol) + im = _optimise_fb(Hv, Visib, Lip, total_flux, lambd, shape, pixel_size, maxiter, tol) # This is needed to match IDL output # im = np.rot90(im, -1) if map: - header = generate_header(vis, shape=shape, pixel_size=pixel) + header = generate_header(vis, shape=shape, pixel_size=pixel_size) return Map((im, header)) return im diff --git a/xrayvision/tests/test_clean.py b/xrayvision/tests/test_clean.py index 31e781c..0e9673c 100644 --- a/xrayvision/tests/test_clean.py +++ b/xrayvision/tests/test_clean.py @@ -1,14 +1,15 @@ import astropy.units as u import numpy as np from astropy.convolution.kernels import Gaussian2DKernel +from numpy.testing import assert_allclose from scipy import signal from xrayvision.clean import ( + _component, + _radial_prolate_sphereoidal, + _vec_radial_prolate_sphereoidal, clean, - component, ms_clean, - radial_prolate_sphereoidal, - vec_radial_prolate_sphereoidal, vis_clean, ) from xrayvision.imaging import image_to_vis @@ -23,6 +24,7 @@ def test_clean_ideal(): clean_map = np.zeros((n, m)) clean_map[pos1[0], pos1[1]] = 10.0 clean_map[pos2[0], pos2[1]] = 7.0 + clean_map = clean_map # << 1 / u.arcsec**2 dirty_beam = np.zeros((n, m)) dirty_beam[(n - 1) // 4 : (n - 1) // 4 + (n - 1) // 2, (m - 1) // 2] = 0.75 @@ -36,35 +38,35 @@ def test_clean_ideal(): dirty_map = signal.convolve(clean_map, dirty_beam, mode="same") # Disable convolution of model with gaussian for testing - out_map, model, resid = clean(dirty_map, dirty_beam, clean_beam_width=0.0) + out_map, model, resid = clean(dirty_map, dirty_beam, clean_beam_width=None) # Within threshold default threshold of 0.1 - assert np.allclose(clean_map, out_map, atol=dirty_beam.max() * 0.1) + assert_allclose(out_map, clean_map, atol=dirty_beam.max() * 1e-13) def test_component(): comp = np.zeros((3, 3)) comp[1, 1] = 1.0 - res = component(scale=0, shape=(3, 3)) + res = _component(scale=0, shape=(3, 3)) assert np.array_equal(res, comp) - res = component(scale=1, shape=(3, 3)) + res = _component(scale=1, shape=(3, 3)) assert np.array_equal(res, comp) - res = component(scale=2, shape=(6, 6)) + res = _component(scale=2, shape=(6, 6)) assert np.all(res[0, :] == 0.0) assert np.all(res[:, 0] == 0.0) assert np.all(res[2:4, 2:4] == res.max()) - res = component(scale=3, shape=(7, 7)) + res = _component(scale=3, shape=(7, 7)) assert np.all(res[0, :] == 0.0) assert np.all(res[:, 0] == 0.0) assert res[3, 3] == 1 def test_radial_prolate_spheroidal(): - amps = [radial_prolate_sphereoidal(r) for r in [-1.0, 0.0, 0.5, 1.0, 2.0]] + amps = [_radial_prolate_sphereoidal(r) for r in [-1.0, 0.0, 0.5, 1.0, 2.0]] assert amps[0] == 1.0 assert amps[1] == 1.0 assert amps[2] == 0.36106538453111797 @@ -74,8 +76,8 @@ def test_radial_prolate_spheroidal(): def test_vec_radial_prolate_spheroidal(): radii = np.linspace(-0.5, 1.5, 1000) - amps1 = [radial_prolate_sphereoidal(r) for r in radii] - amps2 = vec_radial_prolate_sphereoidal(radii) + amps1 = [_radial_prolate_sphereoidal(r) for r in radii] + amps2 = _vec_radial_prolate_sphereoidal(radii) assert np.allclose(amps1, amps2) @@ -100,7 +102,9 @@ def test_ms_clean_ideal(): dirty_map = signal.convolve2d(clean_map, dirty_beam, mode="same") # Disable convolution of model with gaussian for testing - model, res = ms_clean(dirty_map, dirty_beam, [1, 1] * u.arcsec, scales=[1], clean_beam_width=0.0) + model, res = ms_clean( + dirty_map, dirty_beam, pixel_size=[1, 1] * u.arcsec / u.pixel, scales=[1], clean_beam_width=None + ) recovered = model + res # Within threshold default threshold @@ -129,16 +133,16 @@ def test_clean_sim(): sub_uv = np.hstack([sub_uv, np.zeros((2, 1))]) / u.arcsec # Factor of 9 is compensate for the factor of 3 * 3 increase in size - dirty_beam = idft_map(np.ones(321) / 321, u=sub_uv[0, :], v=sub_uv[1, :], shape=(n * 3 + 1, m * 3 + 1)) + dirty_beam = idft_map(np.ones(321) / 321, u=sub_uv[0, :], v=sub_uv[1, :], shape=(n * 3 + 1, m * 3 + 1) * u.pix) vis = dft_map(data, u=sub_uv[0, :], v=sub_uv[1, :]) - dirty_map = idft_map(vis, u=sub_uv[0, :], v=sub_uv[1, :], shape=(n, m)) + dirty_map = idft_map(vis, weights=1 / 321, u=sub_uv[0, :], v=sub_uv[1, :], shape=(n, m) * u.pix) clean_map, model, res = clean( - dirty_map, dirty_beam, clean_beam_width=0.1 * u.arcsec, pixel=[2, 2] * u.arcsec, niter=500 + dirty_map, dirty_beam, pixel_size=[2, 2] * u.arcsec / u.pix, clean_beam_width=0.1 * u.arcsec, niter=500 ) - np.allclose(data, clean_map.data, rtol=dirty_beam.max() * 0.1) + assert_allclose(clean_map, data, atol=dirty_beam.max() * 0.1) def test_vis_clean_sim(): @@ -164,6 +168,11 @@ def test_vis_clean_sim(): vis = image_to_vis(data * u.dimensionless_unscaled, u=sub_uv[0, :], v=sub_uv[1, :]) clean_map, model, res = vis_clean( - vis, shape=(m, n) * u.pix, clean_beam_width=0, pixel=[2, 2] * u.arcsec, natural=False, niter=100 + vis, + shape=(m, n) * u.pix, + pixel_size=[2, 2] * u.arcsec / u.pix, + clean_beam_width=None, + niter=100, + scheme="uniform", ) np.allclose(data, clean_map.data, atol=0.1) diff --git a/xrayvision/tests/test_imaging.py b/xrayvision/tests/test_imaging.py index 6ea4b98..82819c4 100644 --- a/xrayvision/tests/test_imaging.py +++ b/xrayvision/tests/test_imaging.py @@ -2,6 +2,7 @@ import numpy as np import pytest from astropy.convolution.kernels import Gaussian2DKernel +from numpy.testing import assert_allclose from sunpy.map import Map from xrayvision.imaging import image_to_vis, map_to_vis, vis_psf_image, vis_to_image, vis_to_map @@ -27,30 +28,88 @@ def uv(): return u.flatten() / apu.arcsec, v.flatten() / apu.arcsec +@pytest.mark.parametrize("shape", [(10, 10), (5, 5), (2, 2)]) +def test_image_to_vis_vis_to_image(shape): + # Images with the same total flux + extent = [10, 10] * apu.arcsec + arr1 = np.full((shape), 1) * apu.ct / np.prod(shape) + ps1 = extent / (shape * apu.pix) + img1 = arr1 / (ps1[0] * ps1[1] * apu.pix**2) + + # flux in image is 1 + assert_allclose(img1.sum() * ps1[0] * apu.pix * ps1[1] * apu.pix, 1 * apu.ct) + + # Calculate full u, v coverage so will be equivalent to a discrete Fourier transform (DFT) + u = generate_uv((shape[0]) * apu.pix, pixel_size=ps1[0]) + v = generate_uv((shape[1]) * apu.pix, pixel_size=ps1[1]) + u, v = np.meshgrid(u, v, indexing="ij") + u, v = np.array([u, v]).reshape(2, np.prod(shape)) / apu.arcsec + + vis1 = image_to_vis(img1, u=u, v=v, pixel_size=ps1) + res1 = vis_to_image(vis1, shape=shape * apu.pix, pixel_size=ps1) + assert_allclose(res1, img1) + + +@pytest.mark.skip(reason="WIP") +def test_vis_to_image_conserve(): + shape = (3, 3) + shape1 = (5, 5) + shape2 = (7, 7) + # Images with the same total flux + extent = [10, 10] * apu.arcsec + arr1 = np.zeros(shape) + arr1[shape[0] // 2, shape[1] // 2] = 1 / np.prod(shape) + ps = extent / (shape * apu.pix) + img1 = arr1 / (ps[0] * ps[1] * apu.pix**2) + + # flux in image is same + assert_allclose(img1.sum() * ps[0] * apu.pix * ps[1] * apu.pix, arr1.sum()) + + u = generate_uv(shape[0] * apu.pix, pixel_size=ps[0]) + v = generate_uv(shape[1] * apu.pix, pixel_size=ps[1]) + u, v = np.meshgrid(u, v, indexing="ij") + u, v = np.array([u, v]).reshape(2, np.prod(shape)) / apu.arcsec + + vis1 = image_to_vis(img1, u=u, v=v, pixel_size=ps) + + ps2 = extent / (shape1 * apu.pix) + ps3 = extent / (shape2 * apu.pix) + res = vis_to_image(vis1, shape=shape * apu.pix, pixel_size=ps) + res2 = vis_to_image(vis1, shape=shape1 * apu.pix, pixel_size=ps2) + res3 = vis_to_image(vis1, shape=shape2 * apu.pix, pixel_size=ps3) + + assert_allclose(res, img1) + assert_allclose(res2, img1) + assert_allclose(res3, img1) + + @pytest.mark.parametrize("pixel_size", [(0.5), (1), (2)]) -def test_psf_to_image(pixel_size, uv): +def test_vis_to_psf(pixel_size, uv): u, v = uv - img = np.zeros((65, 65)) * (apu.ph / apu.arcsec**2) - img[32, 32] = 1.0 * (apu.ph / apu.arcsec**2) - obs_vis = dft_map(img, u=u, v=v, pixel_size=[2.0, 2.0] * apu.arcsec) + ps = [pixel_size, pixel_size] * apu.arcsec / apu.pix + img = np.zeros((65, 65)) * (apu.ph / apu.cm**2) + img[32, 32] = 1 * (apu.ph / apu.cm**2) # pixel size of 4 -> 1/4 + obs_vis = dft_map(img, u=u, v=v, pixel_size=ps) weights = np.sqrt(u**2 + v**2).value weights /= weights.sum() - psf_calc = idft_map(obs_vis, u=u, v=v, shape=[65, 65], pixel_size=[2, 2] * apu.arcsec, weights=weights) + psf_calc = idft_map(obs_vis, u=u, v=v, shape=[65, 65] * apu.pix, pixel_size=ps, weights=weights) vis = Visibility(obs_vis, u=u, v=v) - res = vis_psf_image(vis, shape=[65, 65] * apu.pixel, pixel_size=2 * apu.arcsec, natural=False) - assert np.allclose(psf_calc, res) + res = vis_psf_image(vis, shape=[65, 65] * apu.pixel, pixel_size=ps, scheme="uniform") + assert_allclose(res, psf_calc) -def test_vis_to_image(uv): +def test_vis_to_image_against_idft(uv): u, v = uv - img = np.zeros((65, 65)) * (apu.ph / apu.arcsec**2) - img[32, 32] = 1.0 * (apu.ph / apu.arcsec**2) - obs_vis = dft_map(img, u=u, v=v, pixel_size=[2.0, 2.0] * apu.arcsec) + img = np.zeros((65, 65)) * (apu.ph / apu.cm**2) + img[32, 32] = 1.0 * (apu.ph / apu.cm**2) + obs_vis = dft_map(img, u=u, v=v, pixel_size=[2.0, 2.0] * apu.arcsec / apu.pix) weights = np.sqrt(u**2 + v**2).value weights /= weights.sum() - bp_calc = idft_map(obs_vis, u=u, v=v, shape=[65, 65], pixel_size=[2, 2] * apu.arcsec, weights=weights) + bp_calc = idft_map( + obs_vis, u=u, v=v, shape=[65, 65] * apu.pix, pixel_size=[2, 2] * apu.arcsec / apu.pix, weights=weights + ) vis = Visibility(obs_vis, u=u, v=v) - res = vis_to_image(vis, shape=[65, 65] * apu.pixel, pixel_size=2 * apu.arcsec, natural=False) + res = vis_to_image(vis, shape=[65, 65] * apu.pixel, pixel_size=2 * apu.arcsec / apu.pix, scheme="uniform") assert np.allclose(bp_calc, res) @@ -61,14 +120,14 @@ def test_image_to_vis(): image = np.zeros((m, n)) # Calculate full u, v coverage so will be equivalent to a discrete Fourier transform (DFT) - u = generate_uv(m) - v = generate_uv(n) - u, v = np.meshgrid(u, v) + u = generate_uv(m * apu.pix) + v = generate_uv(n * apu.pix) + u, v = np.meshgrid(u, v, indexing="ij") u, v = np.array([u, v]).reshape(2, size) / apu.arcsec # For an empty map visibilities should all be zero (0+0j) empty_vis = image_to_vis(image, u=v, v=v) - assert np.array_equal(empty_vis.center, (0.0, 0.0) * apu.arcsec) + assert np.array_equal(empty_vis.phase_centre, (0.0, 0.0) * apu.arcsec) assert np.array_equal(empty_vis.vis, np.zeros(n * m, dtype=complex)) @@ -79,14 +138,14 @@ def test_image_to_vis_center(): image = np.zeros((m, n)) # Calculate full u, v coverage so will be equivalent to a discrete Fourier transform (DFT) - u = generate_uv(m) - v = generate_uv(n) - u, v = np.meshgrid(u, v) + u = generate_uv(m * apu.pix) + v = generate_uv(n * apu.pix) + u, v = np.meshgrid(u, v, indexing="ij") u, v = np.array([u, v]).reshape(2, size) / apu.arcsec # For an empty map visibilities should all be zero (0+0j) - empty_vis = image_to_vis(image, u=u, v=v, center=(2.0, -3.0) * apu.arcsec) - assert np.array_equal(empty_vis.center, (2.0, -3.0) * apu.arcsec) + empty_vis = image_to_vis(image, u=u, v=v, phase_centre=(2.0, -3.0) * apu.arcsec) + assert np.array_equal(empty_vis.offset, (2.0, -3.0) * apu.arcsec) assert np.array_equal(empty_vis.vis, np.zeros(n * m, dtype=complex)) @@ -99,19 +158,19 @@ def test_map_to_vis(pos, pixel): size = m * n pos = pos * apu.arcsec - pixel = pixel * apu.arcsec + pixel = pixel * apu.arcsec / apu.pix # Calculate full u, v coverage so will be equivalent to a discrete Fourier transform (DFT) - u = generate_uv(m, pos[0]) - v = generate_uv(n, pos[1]) - u, v = np.meshgrid(u, v) + u = generate_uv(m * apu.pix, phase_centre=pos[0], pixel_size=pixel[0]) + v = generate_uv(n * apu.pix, phase_centre=pos[1], pixel_size=pixel[1]) + u, v = np.meshgrid(u, v, indexing="ij") u, v = np.array([u, v]).reshape(2, size) / apu.arcsec header = { - "crval1": pos[0].value, - "crval2": pos[1].value, - "cdelt1": pixel[0].value, - "cdelt2": pixel[1].value, + "crval1": pos[1].value, + "crval2": pos[0].value, + "cdelt1": pixel[1].value, + "cdelt2": pixel[0].value, "cunit1": "arcsec", "cunit2": "arcsec", } @@ -127,19 +186,21 @@ def test_map_to_vis(pos, pixel): assert np.allclose(res, data) -def test_vis_to_image1(): - m = n = 33 +def test_vis_to_image(): + m = n = 30 size = m * n # Calculate full u, v coverage so will be equivalent to a discrete Fourier transform (DFT) - u = generate_uv(m) - v = generate_uv(n) + u = generate_uv(m * apu.pix) + v = generate_uv(n * apu.pix) u, v = np.meshgrid(u, v) uv = np.array([u, v]).reshape(2, size) / apu.arcsec u, v = uv # Astropy index order is opposite to that of numpy, is 1st dim is across second down data = Gaussian2DKernel(6, x_size=n, y_size=m).array + data = np.zeros((m, n)) + data[m // 2, n // 2] = 1 vis = image_to_vis(data, u=u, v=v) res = vis_to_image(vis, shape=(m, n) * apu.pixel) @@ -152,16 +213,16 @@ def test_vis_to_image_single_pixel_size(): size = m * n # Calculate full u, v coverage so will be equivalent to a discrete Fourier transform (DFT) - u = generate_uv(m, pixel_size=2.0 * apu.arcsec) - v = generate_uv(n, pixel_size=2.0 * apu.arcsec) - u, v = np.meshgrid(u, v) + u = generate_uv(m * apu.pix, pixel_size=2.0 * apu.arcsec / apu.pix) + v = generate_uv(n * apu.pix, pixel_size=2.0 * apu.arcsec / apu.pix) + u, v = np.meshgrid(u, v, indexing="ij") uv = np.array([u, v]).reshape(2, size) / apu.arcsec u, v = uv # Astropy index order is opposite to that of numpy, is 1st dim is across second down data = Gaussian2DKernel(6, x_size=n, y_size=m).array - vis = image_to_vis(data, u=u, v=v, pixel_size=(2.0, 2.0) * apu.arcsec) - res = vis_to_image(vis, shape=(m, n) * apu.pixel, pixel_size=2.0 * apu.arcsec) + vis = image_to_vis(data, u=u, v=v, pixel_size=(2.0, 2.0) * apu.arcsec / apu.pix) + res = vis_to_image(vis, shape=(m, n) * apu.pixel, pixel_size=2.0 * apu.arcsec / apu.pix) assert res.shape == (m, n) assert np.allclose(data, res) @@ -171,9 +232,9 @@ def test_vis_to_image_invalid_pixel_size(): size = m * n # Calculate full u, v coverage so will be equivalent to a discrete Fourier transform (DFT) - u = generate_uv(m) - v = generate_uv(n) - u, v = np.meshgrid(u, v) + u = generate_uv(m * apu.pix) + v = generate_uv(n * apu.pix) + u, v = np.meshgrid(u, v, indexing="ij") uv = np.array([u, v]).reshape(2, size) / apu.arcsec u, v = uv @@ -182,24 +243,30 @@ def test_vis_to_image_invalid_pixel_size(): vis = image_to_vis(data, u=u, v=v) with pytest.raises(ValueError): - vis_to_image(vis, shape=(m, n) * apu.pixel, pixel_size=[1, 2, 2] * apu.arcsec) + vis_to_image(vis, shape=(m, n) * apu.pixel, pixel_size=[1, 2, 2] * apu.arcsec / apu.pix) -@pytest.mark.parametrize("m,n,pos,pixel", [(33, 33, (10.0, -5.0), (2.0, 3.0)), (32, 32, (-12, -19), (1.0, 5.0))]) +@pytest.mark.parametrize( + "m,n,pos,pixel", + [ # (33, 33, (0, 0), (1.5, 1.5)), + (33, 33, (10.0, -5.0), (2.0, 3.0)), + (32, 32, (-12, -19), (1.0, 5.0)), + ], +) def test_vis_to_map(m, n, pos, pixel): pos = pos * apu.arcsec - pixel = pixel * apu.arcsec - u = generate_uv(m, pixel[0]) - v = generate_uv(n, pixel[1]) - u, v = np.meshgrid(u, v) + pixel = pixel * apu.arcsec / apu.pix + u = generate_uv(m * apu.pix, phase_centre=pos[0], pixel_size=pixel[0]) + v = generate_uv(n * apu.pix, phase_centre=pos[1], pixel_size=pixel[1]) + u, v = np.meshgrid(u, v, indexing="ij") uv = np.array([u, v]).reshape(2, m * n) / apu.arcsec u, v = uv header = { - "crval1": pos[0].value, - "crval2": pos[1].value, - "cdelt1": pixel[0].value, - "cdelt2": pixel[1].value, + "crval1": pos[1].value, + "crval2": pos[0].value, + "cdelt1": pixel[1].value, + "cdelt2": pixel[0].value, "cunit1": "arcsec", "cunit2": "arcsec", } @@ -208,22 +275,23 @@ def test_vis_to_map(m, n, pos, pixel): vis = map_to_vis(mp, u=u, v=v) - res = vis_to_map(vis, shape=(m, n) * apu.pixel, pixel_size=pixel, natural=False) - # assert np.allclose(res.data, data) + res = vis_to_map(vis, shape=(m, n) * apu.pixel, pixel_size=pixel, scheme="natural") + # TODO: figure out why atol is needed here + assert_allclose(data, res.data, atol=1e-15) - assert res.reference_coordinate.Tx == pos[0] - assert res.reference_coordinate.Ty == pos[1] - assert res.scale.axis1 == pixel[0] / apu.pix - assert res.scale.axis2 == pixel[1] / apu.pix + assert res.reference_coordinate.Tx == pos[1] + assert res.reference_coordinate.Ty == pos[0] + assert res.scale.axis1 == pixel[1] + assert res.scale.axis2 == pixel[0] assert res.dimensions.x == m * apu.pix assert res.dimensions.y == n * apu.pix def test_to_sunpy_single_pixel_size(): m = n = 32 - u = generate_uv(m, pixel_size=2.0 * apu.arcsec) - v = generate_uv(m, pixel_size=2.0 * apu.arcsec) - u, v = np.meshgrid(u, v) + u = generate_uv(m * apu.pix, pixel_size=2.0 * apu.arcsec / apu.pix) + v = generate_uv(m * apu.pix, pixel_size=2.0 * apu.arcsec / apu.pix) + u, v = np.meshgrid(u, v, indexing="ij") uv = np.array([u, v]).reshape(2, m * n) / apu.arcsec u, v = uv @@ -232,7 +300,7 @@ def test_to_sunpy_single_pixel_size(): mp = Map((data, header)) vis = map_to_vis(mp, u=u, v=v) - res = vis_to_map(vis, shape=(m, n) * apu.pixel, pixel_size=2 * apu.arcsec) + res = vis_to_map(vis, shape=(m, n) * apu.pixel, pixel_size=2 * apu.arcsec / apu.pix) assert res.meta["cdelt1"] == 2.0 assert res.meta["cdelt1"] == 2.0 assert np.allclose(data, res.data) @@ -240,9 +308,9 @@ def test_to_sunpy_single_pixel_size(): def test_to_sunpy_map_invalid_pixel_size(): m = n = 32 - u = generate_uv(m) - v = generate_uv(m) - u, v = np.meshgrid(u, v) + u = generate_uv(m * apu.pix) + v = generate_uv(n * apu.pix) + u, v = np.meshgrid(u, v, indexing="ij") uv = np.array([u, v]).reshape(2, m * n) / apu.arcsec u, v = uv @@ -253,4 +321,4 @@ def test_to_sunpy_map_invalid_pixel_size(): vis = map_to_vis(mp, u=u, v=v) with pytest.raises(ValueError): - vis_to_map(vis, shape=(m, n) * apu.pixel, pixel_size=[1, 2, 3] * apu.arcsec) + vis_to_map(vis, shape=(m, n) * apu.pixel, pixel_size=[1, 2, 3] * apu.arcsec / apu.pix) diff --git a/xrayvision/tests/test_mem.py b/xrayvision/tests/test_mem.py index e58471b..f2568a5 100644 --- a/xrayvision/tests/test_mem.py +++ b/xrayvision/tests/test_mem.py @@ -1,116 +1,29 @@ import astropy.units as u import numpy as np from astropy.convolution.kernels import Gaussian2DKernel +from numpy.testing import assert_allclose from xrayvision.imaging import image_to_vis from xrayvision.mem import mem, resistant_mean def test_resistant_mean(): + # fmt: off x = np.array( - [ - -0.61566182, - -0.3362035, - -0.73589487, - -0.13014149, - -0.91786147, - -0.02883366, - 1.37688696, - -0.22389746, - -1.37010835, - 0.21693273, - -0.82412723, - 0.56856134, - -0.36824209, - -0.65785301, - 0.10807291, - -0.16730037, - -0.30308355, - 0.12605051, - -1.14017418, - -1.36991364, - -0.45275646, - 0.31636448, - -0.1316372, - -1.19194891, - 0.58557167, - 1.18596057, - 0.01104923, - 0.61225975, - 0.133178, - -0.37117463, - -0.14882964, - -0.59864443, - -0.49450362, - 0.83710754, - -0.51270022, - 0.87380341, - 0.52904637, - 0.37973441, - 0.34061863, - -0.10493781, - 0.86224542, - -0.84283578, - -0.35069229, - -1.29952663, - 0.39678443, - -0.94903521, - 0.008999, - 0.52126955, - 0.24956007, - -2.18242896, - 0.15877853, - -0.99771069, - -0.14426402, - 1.48957386, - -0.68078661, - 1.07328579, - -2.57370777, - 2.26526367, - 1.20964125, - -1.26854227, - 0.20587741, - -0.12066582, - -0.78971368, - -0.36412889, - 0.12816666, - -0.50670234, - 1.19987195, - 1.29270178, - 1.03906824, - 1.41541003, - 0.15639539, - 0.92386204, - -2.62109581, - 0.68909503, - -0.34981646, - 0.93647975, - -0.32268089, - 1.67739013, - -0.29216895, - -2.30162528, - -0.89649313, - 0.37481482, - -2.18661087, - 0.2190169, - -0.30387412, - 0.84001059, - -0.25604226, - 0.50463104, - -1.25674997, - 0.66314647, - -0.0073464, - -0.59719724, - 0.64527612, - 1.25720287, - 0.5032624, - 0.07773715, - 0.29573575, - -0.80567372, - -1.92934995, - -2.0358119, - ] + [-0.61566182, -0.3362035, -0.73589487, -0.13014149, -0.91786147, -0.02883366, 1.37688696, -0.22389746, + -1.37010835, 0.21693273, -0.82412723, 0.56856134, -0.36824209, -0.65785301, 0.10807291, -0.16730037, + -0.30308355, 0.12605051, -1.14017418, -1.36991364, -0.45275646, 0.31636448, -0.1316372, -1.19194891, + 0.58557167, 1.18596057, 0.01104923, 0.61225975, 0.133178, -0.37117463, -0.14882964, -0.59864443, -0.49450362, + 0.83710754, -0.51270022, 0.87380341, 0.52904637, 0.37973441, 0.34061863, -0.10493781, 0.86224542, -0.84283578, + -0.35069229, -1.29952663, 0.39678443, -0.94903521, 0.008999, 0.52126955, 0.24956007, -2.18242896, 0.15877853, + -0.99771069, -0.14426402, 1.48957386, -0.68078661, 1.07328579, -2.57370777, 2.26526367, 1.20964125, + -1.26854227, 0.20587741, -0.12066582, -0.78971368, -0.36412889, 0.12816666, -0.50670234, 1.19987195, + 1.29270178, 1.03906824, 1.41541003, 0.15639539, 0.92386204, -2.62109581, 0.68909503, -0.34981646, 0.93647975, + -0.32268089, 1.67739013, -0.29216895, -2.30162528, -0.89649313, 0.37481482, -2.18661087, 0.2190169, + -0.30387412, 0.84001059, -0.25604226, 0.50463104, -1.25674997, 0.66314647, -0.0073464, -0.59719724, 0.64527612, + 1.25720287, 0.5032624, 0.07773715, 0.29573575, -0.80567372, -1.92934995, -2.0358119,] ) + # fmt: on rmean, rsigma = resistant_mean(x, 3) # values for IDL routine "RESISTANT_Mean, xx, 3, rmean, rsig, num" assert np.allclose(rmean, -0.10947954) @@ -143,9 +56,11 @@ def test_mem(): sub_uv = np.vstack([x.flatten(), y.flatten()]) / u.arcsec # sub_uv = np.hstack([sub_uv, np.zeros((2, 1))]) / u.arcsec - vis = image_to_vis(data, u=sub_uv[0, :], v=sub_uv[1, :], pixel_size=2 * u.arcsec) + vis = image_to_vis(data * u.dimensionless_unscaled, u=sub_uv[0, :], v=sub_uv[1, :], pixel_size=2 * u.arcsec / u.pix) setattr(vis, "amplitude_error", np.sqrt(np.abs(vis.vis))) setattr(vis, "label", [str(x) for x in np.sqrt(x**2 + y**2).flatten()]) - res = mem(vis, percent_lambda=None, shape=(m, n) * u.pix, pixel=[2, 2] * u.arcsec, maxiter=1000, tol=1e-3) - np.testing.assert_allclose(res.data - data / 4, 0.0, atol=1e-3) + res = mem( + vis, percent_lambda=None, shape=(m, n) * u.pix, pixel_size=[2, 2] * u.arcsec / u.pix, maxiter=1000, tol=1e-3 + ) + assert_allclose(res.data, data, atol=1e-1) diff --git a/xrayvision/tests/test_transform.py b/xrayvision/tests/test_transform.py index 97fe9c9..4fcb6d8 100644 --- a/xrayvision/tests/test_transform.py +++ b/xrayvision/tests/test_transform.py @@ -1,8 +1,8 @@ import astropy.units as apu import numpy as np import pytest -from astropy.convolution import Gaussian2DKernel -from numpy.testing import assert_allclose, assert_array_equal +from numpy.fft import fft2, fftshift, ifft2, ifftshift +from numpy.testing import assert_allclose from scipy import signal from xrayvision.transform import dft_map, generate_uv, generate_xy, idft_map @@ -10,268 +10,295 @@ @pytest.mark.parametrize("pixel_size", [0.5, 1, 2, 3]) def test_generate_xy_pixel_size(pixel_size): - # Odd - pixel_size = pixel_size * apu.arcsec - odd = np.linspace(-16 * pixel_size, 16 * pixel_size, 33) - assert np.array_equal(odd, generate_xy(33, pixel_size=pixel_size)) + odd = np.linspace(-16 * pixel_size, 16 * pixel_size, 33) * apu.arcsec + even = np.linspace(-15.5 * pixel_size, 15.5 * pixel_size, 32) * apu.arcsec + pixel_size = pixel_size * apu.arcsec / apu.pix - # Even - even = np.linspace(-15.5 * pixel_size, 15.5 * pixel_size, 32) - assert np.array_equal(even, generate_xy(32, pixel_size=pixel_size)) + assert np.array_equal(odd, generate_xy(33 * apu.pix, pixel_size=pixel_size)) + assert np.array_equal(even, generate_xy(32 * apu.pix, pixel_size=pixel_size)) -@pytest.mark.parametrize("center", [0, +5.5, -5.5]) -def test_generate_xy_offset(center): - # Odd - center = center * apu.arcsec - odd = np.linspace(-16, 16, 33) * apu.arcsec + center - assert np.array_equal(odd, generate_xy(33, center=center)) +@pytest.mark.parametrize("phase_centre", [0, +5.5, -5.5]) +def test_generate_xy_offset(phase_centre): + phase_centre = phase_centre * apu.arcsec + even = np.linspace(-15.5, 15.5, 32) * apu.arcsec + phase_centre + odd = np.linspace(-16, 16, 33) * apu.arcsec + phase_centre - # Even - even = np.linspace(-15.5, 15.5, 32) * apu.arcsec + center - assert np.array_equal(even, generate_xy(32, center=center)) + assert np.array_equal(even, generate_xy(32 * apu.pix, phase_centre=phase_centre)) + assert np.array_equal(odd, generate_xy(33 * apu.pix, phase_centre=phase_centre)) -@pytest.mark.parametrize("center, pixel_size", [(0, (0.5, 1, 2, 3)), (+5.5, (0.5, 1, 2, 3)), (-5.5, (0.5, 1, 2, 3))]) -def test_generate_xy_offset_size(center, pixel_size): - center = center * apu.arcsec - pixel_size = pixel_size * apu.arcsec +@pytest.mark.parametrize( + "phase_centre, pixel_size", [(0, (0.5, 1, 2, 3)), (+5.5, (0.5, 1, 2, 3)), (-5.5, (0.5, 1, 2, 3))] +) +def test_generate_xy_offset_size(phase_centre, pixel_size): + phase_centre = phase_centre * apu.arcsec + pixel_size = pixel_size * apu.arcsec / apu.pix # No sure if this is a good idea could be hard to debug for size in pixel_size: # Odd - odd = np.linspace(-16 * size, 16 * size, 33) + center - assert np.array_equal(odd, generate_xy(33, center=center, pixel_size=size)) + odd = np.linspace(-16 * size, 16 * size, 33) * apu.pix + phase_centre + assert np.array_equal(odd, generate_xy(33 * apu.pix, phase_centre=phase_centre, pixel_size=size)) # Even - even = np.linspace(-15.5 * size, 15.5 * size, 32) + center - assert np.array_equal(even, generate_xy(32, center=center, pixel_size=size)) + even = np.linspace(-15.5 * size, 15.5 * size, 32) * apu.pix + phase_centre + assert np.array_equal(even, generate_xy(32 * apu.pix, phase_centre=phase_centre, pixel_size=size)) @pytest.mark.parametrize("pixel_size", [0.5, 1, 2, 3]) def test_generate_uv_pixel_size(pixel_size): - pixel_size = pixel_size * apu.arcsec + pixel_size = pixel_size * apu.arcsec / apu.pix m = 33 n = 32 # Odd - odd = np.linspace(-((m - 1) / 2) * (1 / (m * pixel_size)), ((m - 1) / 2) * (1 / (m * pixel_size)), m) + odd = np.linspace(-((m - 1) / 2) * (1 / (m * pixel_size)), ((m - 1) / 2) * (1 / (m * pixel_size)), m) / apu.pix # Know issue with quantities and isclose/allclose need to add unit to atol default value - assert np.allclose(odd, generate_uv(m, pixel_size=pixel_size), atol=1e-8 / apu.arcsec) + assert np.allclose(odd, generate_uv(m * apu.pix, pixel_size=pixel_size), atol=1e-8 / apu.arcsec) # Even - even = (np.arange(n) - n / 2 + 0.5) * (1 / (pixel_size * n)) - assert np.allclose(even, generate_uv(32, pixel_size=pixel_size), atol=1e-8 / apu.arcsec) + even = (np.arange(n) - n / 2 + 0.5) * (1 / (pixel_size * n)) / apu.pix + assert np.allclose(even, generate_uv(32 * apu.pix, pixel_size=pixel_size), atol=1e-8 / apu.arcsec) -@pytest.mark.parametrize("center", [0.0, -5.5, 5.5]) -def test_generate_uv_pixel_offset(center): - center = center * apu.arcsec +@pytest.mark.parametrize("phase_centre", [0.0, -5.5, 5.5]) +def test_generate_uv_pixel_offset(phase_centre): + phase_centre = phase_centre * apu.arcsec m = 33 n = 32 # Odd odd = np.linspace(-((m - 1) / 2) * (1 / m), ((m - 1) / 2) * (1 / m), m) / apu.arcsec - if center.value != 0: - odd += 1 / center - assert np.allclose(odd, generate_uv(m, center=center), atol=1e-8 / apu.arcsec) + if phase_centre.value != 0: + odd += 1 / phase_centre + assert np.allclose(odd, generate_uv(m * apu.pix, phase_centre=phase_centre), atol=1e-8 / apu.arcsec) # Even even = (np.arange(n) - n / 2 + 0.5) * (1 / n) / apu.arcsec - if center.value != 0: - even += 1 / center - assert np.allclose(even, generate_uv(32, center=center), atol=1e-8 / apu.arcsec) + if phase_centre.value != 0: + even += 1 / phase_centre + assert np.allclose(even, generate_uv(n * apu.pix, phase_centre=phase_centre), atol=1e-8 / apu.arcsec) -@pytest.mark.parametrize("center, pixel_size", [(0, (0.5, 1, 2, 3)), (+5.5, (0.5, 1, 2, 3)), (-5.5, (0.5, 1, 2, 3))]) -def test_generate_uv_offset_size(center, pixel_size): - center = center * apu.arcsec - pixel_size = pixel_size * apu.arcsec +@pytest.mark.parametrize( + "phase_centre, pixel_size", [(0, (0.5, 1, 2, 3)), (+5.5, (0.5, 1, 2, 3)), (-5.5, (0.5, 1, 2, 3))] +) +def test_generate_uv_offset_size(phase_centre, pixel_size): + phase_centre = phase_centre * apu.arcsec + pixel_size = pixel_size * apu.arcsec / apu.pix m = 33 n = 32 for size in pixel_size: # Odd - odd = np.linspace(-((m - 1) / 2) * (1 / (size * m)), ((m - 1) / 2) * (1 / (size * m)), m) - if center != 0: - odd += 1 / center - assert np.allclose(odd, generate_uv(m, center=center, pixel_size=size), atol=1e-8 / apu.arcsec) + odd = np.linspace(-((m - 1) / 2) * (1 / (size * m)), ((m - 1) / 2) * (1 / (size * m)), m) / apu.pix + if phase_centre != 0: + odd += 1 / phase_centre + assert np.allclose( + odd, generate_uv(m * apu.pix, phase_centre=phase_centre, pixel_size=size), atol=1e-8 / apu.arcsec + ) # Even - even = (np.arange(n) - n / 2 + 0.5) * (1 / (size * n)) - if center != 0: - even += 1 / center - assert np.allclose(even, generate_uv(32, center=center, pixel_size=size), atol=1e-8 / apu.arcsec) - - -@pytest.mark.parametrize("shape", [(33, 33), (32, 32), (33, 32), (32, 33)]) -def test_dft_idft_map(shape): - m, n = shape - m * n - uu = generate_uv(n) - vv = generate_uv(m) - u, v = np.meshgrid(uu, vv) + even = (np.arange(n) - n / 2 + 0.5) * (1 / (size * n)) / apu.pix + if phase_centre != 0: + even += 1 / phase_centre + assert np.allclose( + even, generate_uv(n * apu.pix, phase_centre=phase_centre, pixel_size=size), atol=1e-8 / apu.arcsec + ) + + +@pytest.mark.parametrize("shape", [(3, 3), (2, 3), (3, 2), (2, 2)]) +@pytest.mark.parametrize("pixel_size", [(1, 1), (1, 2), (2, 1), (2, 2)]) +@pytest.mark.parametrize("center", [(0, 0), (-1, 1), (1, -1), (1, 1)]) +def test_dft_idft(shape, pixel_size, center): + data = np.arange(np.prod(shape)).reshape(shape) + uu = generate_uv( + shape[0] * apu.pix, phase_centre=center[0] * apu.arcsec, pixel_size=pixel_size[0] * apu.arcsec / apu.pix + ) + vv = generate_uv( + shape[1] * apu.pix, phase_centre=center[1] * apu.arcsec, pixel_size=pixel_size[1] * apu.arcsec / apu.pix + ) + u, v = np.meshgrid(uu, vv, indexing="ij") u = u.flatten() v = v.flatten() - # All zeros - zeros = np.zeros((m, n)) - vis = dft_map(zeros, u=u, v=v) - # All visibilities should be zero - # assert_array_equal(np.zeros(np.prod((m, n)), dtype=complex), vis) - out_map = idft_map(vis, u=u, v=v, shape=shape) - # Should get back original map after dft(idft()) - assert_allclose(zeros, out_map / np.prod((m, n))) - - # All ones - ones = np.ones((m, n)) - vis = dft_map(ones, u=u, v=v) - # All visibilities should be 1 - # assert_allclose(np.ones(np.prod((m, n)), dtype=complex), vis) - out_map = idft_map(vis, u=u, v=v, shape=(m, n)) - assert_allclose(ones, out_map / np.prod((m, n))) - - # Delta - delta = zeros[:, :] - delta[m // 2, n // 2] = 1.0 - vis = dft_map(delta, u=u, v=v) - out_map = idft_map(vis, u=u, v=v, shape=shape) - assert_allclose(out_map / np.prod(shape), delta, atol=1e-14) - - # Gaussian - astropy has axis in reverse order compared to numpy - gaussian = Gaussian2DKernel(5, x_size=n, y_size=m).array - vis = dft_map(gaussian, u=u, v=v) - out_map = idft_map(vis, u=u, v=v, shape=shape) - assert_allclose(out_map / np.prod(shape), gaussian) - - -# @pytest.mark.skip('UV coordinate generation off somewhere') -@pytest.mark.parametrize("pixel_size", [(1.0, 1.0), (0.5, 0.5), (1.5, 1.5), (2.0, 0.5), (0.5, 2.0)]) -def test_dft_idft_map_pixel_size(pixel_size): - pixel_size = pixel_size * apu.arcsec - shape = (32, 32) - m, n = shape - size = m * n - uu = generate_uv(m, pixel_size=pixel_size[0]) - vv = generate_uv(n, pixel_size=pixel_size[1]) - u, v = np.meshgrid(uu, vv) - u = u.flatten() - v = v.flatten() - - # All zeros - zeros = np.zeros(shape) - vis = dft_map(zeros, u=u, v=v, pixel_size=pixel_size) - # All visibilities should be zero - assert_array_equal(np.zeros(size, dtype=complex), vis) - out_map = idft_map(vis, u=u, v=v, shape=shape, pixel_size=pixel_size) - # Should get back original map after dft(idft()) - assert_array_equal(out_map / np.prod(shape), zeros) - - # All ones - ones = np.ones(shape) - vis = dft_map(ones, u=u, v=v, pixel_size=pixel_size) - out_map = idft_map(vis, u=u, v=v, shape=shape, pixel_size=pixel_size) - assert_allclose(out_map / np.prod(shape), ones) - - # Delta - delta = zeros[:, :] - delta[m // 2, n // 2] = 1.0 - vis = dft_map(delta, u=u, v=v, pixel_size=pixel_size) - out_map = idft_map(vis, u=u, v=v, shape=shape, pixel_size=pixel_size) - assert_allclose(out_map / np.prod(shape), delta, atol=1e-14) - - # Gaussian - astropy has axis in reverse order compared to numpy - gaussian = Gaussian2DKernel(5, x_size=n, y_size=m).array - vis = dft_map(gaussian, u=u, v=v, pixel_size=pixel_size) - out_map = idft_map(vis, u=u, v=v, shape=shape, pixel_size=pixel_size) - assert_allclose(out_map / np.prod(shape), gaussian) - - -# @pytest.mark.skip('UV coordinate generation off somewhere') -@pytest.mark.parametrize("center", [(0, 0), (2.1, 2.1), (5.4, -5.4), (-5.6, 5.6)]) -def test_dft_idft_map_center(center): - center = center * apu.arcsec - shape = (33, 33) - m, n = shape - size = m * n - uu = generate_uv(n, center=center[1]) - vv = generate_uv(m, center=center[0]) - u, v = np.meshgrid(uu, vv) - u = u.flatten() - v = v.flatten() - - # All zeros - zeros = np.zeros(shape) - vis = dft_map(zeros, u=u, v=v, center=center) - # All visibilities should be zero - assert_array_equal(vis, np.zeros(size, dtype=complex)) - out_map = idft_map(vis, u=u, v=v, shape=shape, center=center) - # Should get back original map after dft(idft()) - assert_array_equal(out_map / np.prod(shape), zeros) - - # All ones - ones = np.ones(shape) - vis = dft_map(ones, u=u, v=v, center=center) - out_map = idft_map(vis, u=u, v=v, shape=shape, center=center) - assert_allclose(out_map / np.prod(shape), ones) - - # Delta - delta = zeros[:, :] - delta[m // 2, n // 2] = 1.0 - vis = dft_map(delta, u=u, v=v, center=center) - out_map = idft_map(vis, u=u, v=v, shape=shape, center=center) - assert_allclose(out_map / np.prod(shape), delta, atol=1e-14) - - # Gaussian - astropy has axis in reverse order compared to numpy - gaussian = Gaussian2DKernel(5, x_size=n, y_size=m).array - vis = dft_map(gaussian, u=u, v=v, center=center) - out_map = idft_map(vis, u=u, v=v, shape=shape, center=center) - assert_allclose(out_map / np.prod(shape), gaussian) - - -# @pytest.mark.skip('UV coordinate generation off somewhere') -@pytest.mark.parametrize( - "shape, pixel_size", - [((11, 10), (0.5, 0.5)), ((11, 10), (1.0, 1.0)), ((11, 10), (2.0, 2.0)), ((11, 10), (3.0, 3.0))], -) -def test_dft_idft_map_shape_pixel_size(shape, pixel_size): - pixel_size = pixel_size * apu.arcsec - m, n = shape - size = m * n - - vv = generate_uv(m, pixel_size=pixel_size[0]) - uu = generate_uv(n, pixel_size=pixel_size[1]) - u, v = np.meshgrid(uu, vv) - u = u.flatten() - v = v.flatten() - - # All zeros - zeros = np.zeros(shape) - vis = dft_map(zeros, u=u, v=v, pixel_size=pixel_size) - # All visibilities should be zero - assert_array_equal(np.zeros(size, dtype=complex), vis) - out_map = idft_map(vis, u=u, v=v, shape=shape, pixel_size=pixel_size) - # Should get back original map after dft(idft()) - assert_array_equal(out_map / np.prod(shape), zeros) - - # All ones - ones = np.ones(shape) - vis = dft_map(ones, u=u, v=v, pixel_size=pixel_size) - out_map = idft_map(vis, u=u, v=v, shape=shape, pixel_size=pixel_size) - assert_allclose(out_map / np.prod(shape), ones) - - # Delta - delta = zeros[:, :] - delta[m // 2, n // 2] = 1.0 - vis = dft_map(delta, u=u, v=v, pixel_size=pixel_size) - out_map = idft_map(vis, u=u, v=v, shape=shape, pixel_size=pixel_size) - assert_allclose(out_map / np.prod(shape), delta, atol=1e-14) - - # Gaussian - astropy has axis in reverse order compared to numpy - gaussian = Gaussian2DKernel(5, x_size=n, y_size=m).array - vis = dft_map(gaussian, u=u, v=v, pixel_size=pixel_size) - out_map = idft_map(vis, u=u, v=v, shape=shape, pixel_size=pixel_size) - assert_allclose(out_map / np.prod(shape), gaussian) + vis = dft_map(data, u=u, v=v, pixel_size=pixel_size * apu.arcsec / apu.pix) + img = idft_map( + vis, u=u, v=v, weights=1 / np.prod(shape), shape=shape * apu.pix, pixel_size=pixel_size * apu.arcsec / apu.pix + ) + assert_allclose(img, data, atol=1e-13) + + +# @pytest.mark.parametrize("shape", [(33, 33), (32, 32), (33, 32), (32, 33)]) +# def test_dft_idft_map(shape): +# m, n = shape +# uu = generate_uv(m * apu.pix) +# vv = generate_uv(n * apu.pix) +# u, v = np.meshgrid(uu, vv, indexing="ij") +# u = u.flatten() +# v = v.flatten() +# +# # All zeros +# zeros = np.zeros((m, n)) +# vis = dft_map(zeros, u=u, v=v) +# # All visibilities should be zero +# # assert_array_equal(np.zeros(np.prod((m, n)), dtype=complex), vis) +# out_map = idft_map(vis, u=u, v=v, shape=shape * apu.pix) +# # Should get back the original map after dft(idft()) +# assert_allclose(zeros, out_map / np.prod((m, n))) +# +# # All ones +# ones = np.ones((m, n)) +# vis = dft_map(ones, u=u, v=v) +# # All visibilities should be 1 +# # assert_allclose(np.ones(np.prod((m, n)), dtype=complex), vis) +# out_map = idft_map(vis, u=u, v=v, shape=shape * apu.pix) +# assert_allclose(ones, out_map / np.prod((m, n))) +# +# # Delta +# delta = zeros[:, :] +# delta[m // 2, n // 2] = 1.0 +# vis = dft_map(delta, u=u, v=v) +# out_map = idft_map(vis, u=u, v=v, shape=shape * apu.pix) +# assert_allclose(out_map / np.prod(shape), delta, atol=1e-14) +# +# # Gaussian - astropy has axis in reverse order compared to numpy +# gaussian = Gaussian2DKernel(5, x_size=n, y_size=m).array +# vis = dft_map(gaussian, u=u, v=v) +# out_map = idft_map(vis, u=u, v=v, shape=shape * apu.pix) +# assert_allclose(out_map / np.prod(shape), gaussian) +# +# +# # @pytest.mark.skip('UV coordinate generation off somewhere') +# @pytest.mark.parametrize("pixel_size", [(1.0, 1.0), (0.5, 0.5), (1.5, 1.5), (2.0, 0.5), (0.5, 2.0)]) +# def test_dft_idft_map_pixel_size(pixel_size): +# pixel_size = pixel_size * apu.arcsec / apu.pix +# shape = (32, 32) * apu.pix +# m, n = shape.value.astype(int) +# size = m * n +# uu = generate_uv(m * apu.pix, pixel_size=pixel_size[0]) +# vv = generate_uv(n * apu.pix, pixel_size=pixel_size[1]) +# u, v = np.meshgrid(uu, vv) +# u = u.flatten() +# v = v.flatten() +# +# # All zeros +# zeros = np.zeros(shape.value.astype(int)) +# vis = dft_map(zeros, u=u, v=v, pixel_size=pixel_size) +# # All visibilities should be zero +# assert_array_equal(np.zeros(size, dtype=complex), vis) +# out_map = idft_map(vis, u=u, v=v, shape=shape, pixel_size=pixel_size) +# # Should get back the original map after dft(idft()) +# assert_array_equal(out_map / np.prod(shape.value), zeros) +# +# # All ones +# ones = np.ones(shape.value.astype(int)) +# vis = dft_map(ones, u=u, v=v, pixel_size=pixel_size) +# out_map = idft_map(vis, u=u, v=v, shape=shape, pixel_size=pixel_size) +# assert_allclose(out_map / np.prod(shape.value), ones) +# +# # Delta +# delta = zeros[:, :] +# delta[m // 2, n // 2] = 1.0 +# vis = dft_map(delta, u=u, v=v, pixel_size=pixel_size) +# out_map = idft_map(vis, u=u, v=v, shape=shape, pixel_size=pixel_size) +# assert_allclose(out_map / np.prod(shape.value), delta, atol=1e-14) +# +# # Gaussian - astropy has axis in reverse order compared to numpy +# gaussian = Gaussian2DKernel(5, x_size=n, y_size=m).array +# vis = dft_map(gaussian, u=u, v=v, pixel_size=pixel_size) +# out_map = idft_map(vis, u=u, v=v, shape=shape, pixel_size=pixel_size) +# assert_allclose(out_map / np.prod(shape.value), gaussian) +# +# +# # @pytest.mark.skip('UV coordinate generation off somewhere') +# @pytest.mark.parametrize("phase_centre", [(0, 0), (2.1, 1.1), (5.4, -4.5), (-5.6, 5.6)]) +# def test_dft_idft_map_center(phase_centre): +# phase_centre = phase_centre * apu.arcsec +# shape = (33, 33) +# m, n = shape +# size = m * n +# shape = shape * apu.pix +# uu = generate_uv(n * apu.pix, phase_centre=phase_centre[0]) +# vv = generate_uv(m * apu.pix, phase_centre=phase_centre[1]) +# u, v = np.meshgrid(uu, vv) +# u = u.flatten() +# v = v.flatten() +# +# # All zeros +# zeros = np.zeros((m, n)) +# vis = dft_map(zeros, u=u, v=v, phase_centre=phase_centre) +# # All visibilities should be zero +# assert_array_equal(vis, np.zeros(size, dtype=complex)) +# out_map = idft_map(vis, u=u, v=v, shape=shape, phase_centre=phase_centre) +# # Should get back the original map after dft(idft()) +# assert_array_equal(out_map / np.prod((m, n)), zeros) +# +# # All ones +# ones = np.ones((m, n)) +# vis = dft_map(ones, u=u, v=v, phase_centre=phase_centre) +# out_map = idft_map(vis, u=u, v=v, shape=shape, phase_centre=phase_centre) +# assert_allclose(out_map / np.prod((m, n)), ones) +# +# # Delta +# delta = zeros[:, :] +# delta[m // 2, n // 2] = 1.0 +# vis = dft_map(delta, u=u, v=v, phase_centre=phase_centre) +# out_map = idft_map(vis, u=u, v=v, shape=shape, phase_centre=phase_centre) +# assert_allclose(out_map / np.prod((m, n)), delta, atol=1e-14) +# +# # Gaussian - astropy has axis in reverse order compared to numpy +# gaussian = Gaussian2DKernel(5, x_size=n, y_size=m).array +# vis = dft_map(gaussian, u=u, v=v, phase_centre=phase_centre) +# out_map = idft_map(vis, u=u, v=v, shape=shape, phase_centre=phase_centre) +# assert_allclose(out_map / np.prod((m, n)), gaussian) +# +# +# # @pytest.mark.skip('UV coordinate generation off somewhere') +# @pytest.mark.parametrize( +# "shape, pixel_size", +# [((11, 10), (0.5, 0.5)), ((11, 10), (1.0, 1.0)), ((11, 10), (2.0, 2.0)), ((11, 10), (3.0, 3.0))], +# ) +# def test_dft_idft_map_shape_pixel_size(shape, pixel_size): +# pixel_size = pixel_size * apu.arcsec / apu.pix +# m, n = shape +# size = m * n +# shape = shape * apu.pix +# +# uu = generate_uv(m * apu.pix, pixel_size=pixel_size[0]) +# vv = generate_uv(n * apu.pix, pixel_size=pixel_size[1]) +# u, v = np.meshgrid(uu, vv, indexing="ij") +# u = u.flatten() +# v = v.flatten() +# +# # All zeros +# zeros = np.zeros((m, n)) +# vis = dft_map(zeros, u=u, v=v, pixel_size=pixel_size) +# # All visibilities should be zero +# assert_array_equal(np.zeros(size, dtype=complex), vis) +# out_map = idft_map(vis, u=u, v=v, shape=shape, pixel_size=pixel_size) +# # Should get back the original map after dft(idft()) +# assert_array_equal(out_map / np.prod((m, n)), zeros) +# +# # All ones +# ones = np.ones((m, n)) +# vis = dft_map(ones, u=u, v=v, pixel_size=pixel_size) +# out_map = idft_map(vis, u=u, v=v, shape=shape, pixel_size=pixel_size) +# assert_allclose(out_map / np.prod((m, n)), ones) +# +# # Delta +# delta = zeros[:, :] +# delta[m // 2, n // 2] = 1.0 +# vis = dft_map(delta, u=u, v=v, pixel_size=pixel_size) +# out_map = idft_map(vis, u=u, v=v, shape=shape, pixel_size=pixel_size) +# assert_allclose(out_map / np.prod((m, n)), delta, atol=1e-14) +# +# # Gaussian - astropy has axis in reverse order compared to numpy +# gaussian = Gaussian2DKernel(5, x_size=n, y_size=m).array +# vis = dft_map(gaussian, u=u, v=v, pixel_size=pixel_size) +# out_map = idft_map(vis, u=u, v=v, shape=shape, pixel_size=pixel_size) +# assert_allclose(out_map / np.prod((m, n)), gaussian) def test_equivalence_of_convolve(): @@ -280,10 +307,9 @@ def test_equivalence_of_convolve(): data[3:6, 3:6] = 5.0 m, n = (33, 33) - m * n - vv = generate_uv(m) - uu = generate_uv(n) + vv = generate_uv(m * apu.pix) + uu = generate_uv(n * apu.pix) u, v = np.meshgrid(uu, vv) u = u.flatten() v = v.flatten() @@ -298,18 +324,87 @@ def test_equivalence_of_convolve(): # bp1 = idft_map(full_vis, u=u, v=v, shape=(33, 33)) - bp2 = idft_map(sub_vis[non_zero], u=u[non_zero], v=v[non_zero], shape=(33, 33)) + bp2 = idft_map(sub_vis[non_zero], u=u[non_zero], v=v[non_zero], shape=(33, 33) * apu.pix) - # Need to make the psf large enough to slide over entire data window - psf1 = idft_map(sampling[non_zero], u=u[non_zero], v=v[non_zero], shape=(33 * 3, 33 * 3)) + # Need to make the psf is large enough to slide over the entire data window + psf1 = idft_map(sampling[non_zero], u=u[non_zero], v=v[non_zero], shape=(33 * 3, 33 * 3) * apu.pix) conv = signal.convolve(data, psf1, mode="same", method="fft") - psf2 = idft_map(sampling[non_zero], u=u[non_zero], v=v[non_zero], shape=(33, 33)) + psf2 = idft_map(sampling[non_zero], u=u[non_zero], v=v[non_zero], shape=(33, 33) * apu.pix) - bp3 = idft_map(full_vis[non_zero], u=u[non_zero], v=v[non_zero], shape=(33, 33)) + bp3 = idft_map(full_vis[non_zero], u=u[non_zero], v=v[non_zero], shape=(33, 33) * apu.pix) assert np.allclose(bp2, conv) assert np.allclose(bp2, bp3) - # Due to the enlarged psf need to only use center portion + # Due to the enlarged psf need to only use the centre portion assert np.allclose(psf1[33:66, 33:66], psf2) + + +def test_phase_centre_equivalence(): + # Calculate full u, v coverage so will be equivalent to a discrete Fourier transform (DFT) + data = np.random.randn(8, 8) / apu.arcsec**2 + u = generate_uv(8 * apu.pix, phase_centre=0 * apu.arcsec, pixel_size=1 * apu.arcsec / apu.pix) + v = generate_uv(8 * apu.pix, phase_centre=0 * apu.arcsec, pixel_size=1 * apu.arcsec / apu.pix) + u, v = np.meshgrid(u, v) + u, v = np.array([u, v]).reshape(2, u.size) / apu.arcsec + + # calculate normal 0,0 + vis = dft_map(data, u=u, v=v, pixel_size=[1, 1] * apu.arcsec / apu.pix) + img1 = idft_map( + vis, u=u, v=v, weights=1 / u.size, pixel_size=[1, 1] * apu.arcsec / apu.pix, shape=data.shape * apu.pix + ) + assert_allclose(data, img1) + + # extract phase and amp + phase = np.arctan2(np.imag(vis), np.real(vis)) + amp = np.abs(vis) + + # change vis to a phase centre of 5, 5 + phase_shift = 2 * np.pi * (5 * apu.arcsec * u + 5 * apu.arcsec * v) * apu.rad + vis_shifted = (np.cos(phase + phase_shift) + np.sin(phase + phase_shift) * 1j) * amp + + # make image with centre of 5, 5 with shifted vis + img2 = idft_map( + vis_shifted, + u=u, + v=v, + weights=1 / u.size, + pixel_size=[1, 1] * apu.arcsec / apu.pix, + shape=data.shape * apu.pix, + phase_centre=[5, 5] * apu.arcsec, + ) + assert np.allclose(data, img2) + + +def test_fft_equivalence(): + # Odd (3, 3) so symmetric and chose shape and pixel size so xy values will run from 0 to 2 the same as in fft + # TODO: add same kind of test for even for fft2 then A[n/2] has both pos and negative nyquist frequencies + # e.g shape (2, 2), (3, 2), (2, 3) + shape = (3, 3) + pixel = (1, 1) + center = (1, 1) + + data = np.arange(np.prod(shape)).reshape(shape) + uu = generate_uv( + shape[0] * apu.pix, phase_centre=center[0] * apu.arcsec, pixel_size=pixel[0] * apu.arcsec / apu.pix + ) + vv = generate_uv( + shape[1] * apu.pix, phase_centre=center[1] * apu.arcsec, pixel_size=pixel[1] * apu.arcsec / apu.pix + ) + u, v = np.meshgrid(uu, vv, indexing="ij") + u = u.flatten() + v = v.flatten() + + vis = dft_map(data, u=u, v=v, pixel_size=pixel * apu.arcsec / apu.pix, phase_centre=center * apu.arcsec) + + ft = fft2(data) + fts = fftshift(ft) + vis = vis.reshape(shape) + # Convention in xrayvison has the minus sign on the forward transform but numpy has it on reverse + vis_conj = np.conjugate(vis) + assert_allclose(fts, vis_conj, atol=1e-13) + + vis_ft = ifftshift(vis_conj) + img = ifft2(vis_ft) + assert_allclose(np.real(img), data, atol=1e-14) diff --git a/xrayvision/tests/test_visibility.py b/xrayvision/tests/test_visibility.py index b47ab04..cfa1d0e 100644 --- a/xrayvision/tests/test_visibility.py +++ b/xrayvision/tests/test_visibility.py @@ -1,16 +1,19 @@ -from pathlib import Path - import astropy.units as apu import pytest from astropy.tests.helper import assert_quantity_allclose +from numpy.testing import assert_array_equal import xrayvision.visibility as vm +from xrayvision.visibility import Visibility -@pytest.fixture -def test_data_dir(): - path = Path(__file__).parent.parent / "data" - return path +def test_visibility(): + vis = Visibility(vis=1 * apu.ct, u=1 / apu.deg, v=1 / apu.deg) + assert vis.vis == 1 * apu.ct + assert vis.u == 1 / apu.deg + assert vis.v == 1 / apu.deg + assert_array_equal([0, 0] * apu.arcsec, vis.phase_centre) + assert_array_equal([0, 0] * apu.arcsec, vis.offset) @pytest.fixture diff --git a/xrayvision/transform.py b/xrayvision/transform.py index 4ceebcb..d1ed6b7 100644 --- a/xrayvision/transform.py +++ b/xrayvision/transform.py @@ -7,128 +7,163 @@ """ +from typing import Union, Optional + import astropy.units as apu import numpy as np +import numpy.typing as npt +from astropy.units import Quantity from astropy.units.core import UnitsError +__all__ = ["generate_xy", "generate_uv", "dft_map", "idft_map"] + -@apu.quantity_input(center="angle", pixel_size="angle") -def generate_xy(number_pixels, center=0.0 * apu.arcsec, pixel_size=1.0 * apu.arcsec): +@apu.quantity_input() +def generate_xy( + number_pixels: Quantity[apu.pix], + *, + phase_centre: Optional[Quantity[apu.arcsec]] = 0.0 * apu.arcsec, + pixel_size: Optional[Quantity[apu.arcsec / apu.pix]] = 1.0 * apu.arcsec / apu.pix, +) -> Quantity[apu.arcsec]: """ - Generate the x or y coordinates given the number of pixels, center and pixel size. + Generate the x or y coordinates given the number of pixels, phase_centre and pixel size. Parameters ---------- number_pixels : `int` Number of pixels - center : `float`, optional + phase_centre : `float`, optional Center coordinates pixel_size : `float`, optional Size of pixel in physical units (e.g. arcsecs, meters) Returns ------- - `numpy.array` + : The generated x, y coordinates See Also -------- - `generate_uv` : Generates corresponding coordinates but un u, v space + generate_uv: + Generates corresponding coordinates but in Fourier or u, v space. Examples -------- - >>> generate_xy(9) + >>> import astropy.units as apu + >>> generate_xy(9*apu.pix) - >>> generate_xy(9, pixel_size=2.5 * apu.arcsec) - + >>> generate_xy(9*apu.pix, pixel_size=2.5 * apu.arcsec/apu.pix) + - >>> generate_xy(9, center=10 * apu.arcsec, pixel_size=2.5 * apu.arcsec) + >>> generate_xy(9*apu.pix, phase_centre=10 * apu.arcsec, pixel_size=2.5 * apu.arcsec/apu.pix) """ - x = (np.arange(number_pixels) - number_pixels / 2 + 0.5) * pixel_size + center + x = ( + np.arange(number_pixels.to_value(apu.pixel)) - (number_pixels.to_value(apu.pix) / 2) + 0.5 + ) * apu.pix * pixel_size + phase_centre return x -@apu.quantity_input(center="angle", pixel_size="angle") -def generate_uv(number_pixels, center=0.0 * apu.arcsec, pixel_size=1.0 * apu.arcsec): +@apu.quantity_input() +def generate_uv( + number_pixels: Quantity[apu.pix], + *, + phase_centre: Optional[Quantity[apu.arcsec]] = 0.0 * apu.arcsec, + pixel_size: Optional[Quantity[apu.arcsec / apu.pix]] = 1.0 * apu.arcsec / apu.pix, +) -> Quantity[1 / apu.arcsec]: """ - Generate the u or v coordinates given the number of pixels, center and pixel size. + Generate the u or v coordinates given the number of pixels, phase_centre and pixel size. Parameters ---------- number_pixels : `int` Number of pixels - center : `float`, optional + phase_centre : `float`, optional Center coordinates pixel_size : `float`, optional Size of pixel in physical units (e.g. arcsecs, meters) Returns ------- - `numpy.array` - The generated u, v coordinates + : + The generated u, v coordinates. See Also -------- - `generate_xy` : Generates corresponding coordinate but un x, y space + generate_xy: + Generates corresponding coordinates but in Fourier or u, v space. Examples -------- - >>> generate_uv(9) + >>> import astropy.units as apu + >>> generate_uv(9*apu.pix) - >>> generate_uv(9, pixel_size=2.5 * apu.arcsec) + >>> generate_uv(9*apu.pix, pixel_size=2.5 * apu.arcsec/apu.pix) - >>> generate_uv(9, center=10 * apu.arcsec, pixel_size=2.5 * apu.arcsec) + >>> generate_uv(9*apu.pix, phase_centre=10 * apu.arcsec, pixel_size=2.5 * apu.arcsec/apu.pix) """ - x = (np.arange(number_pixels) - number_pixels / 2 + 0.5) / (pixel_size * number_pixels) - if center.value != 0.0: - x += 1 / center + # x = (np.arange(number_pixels) - number_pixels / 2 + 0.5) / (pixel_size * number_pixels) + + x = (np.arange(number_pixels.to_value(apu.pixel)) - (number_pixels.to_value(apu.pix) / 2) + 0.5) / ( + pixel_size * number_pixels + ) + + if phase_centre.value != 0.0: # type: ignore + x += 1 / phase_centre # type: ignore return x -@apu.quantity_input(center="angle", pixel_size="angle") -def dft_map(input_array, *, u, v, center=(0.0, 0.0) * apu.arcsec, pixel_size=(1.0, 1.0) * apu.arcsec): +@apu.quantity_input() +def dft_map( + input_array: Union[Quantity, npt.NDArray], + *, + u: Quantity[1 / apu.arcsec], + v: Quantity[1 / apu.arcsec], + phase_centre: Quantity[apu.arcsec] = (0.0, 0.0) * apu.arcsec, + pixel_size: Quantity[apu.arcsec / apu.pix] = (1.0, 1.0) * apu.arcsec / apu.pix, +) -> Union[Quantity, npt.NDArray]: r""" Discrete Fourier transform in terms of coordinates returning 1-D array complex visibilities. Parameters ---------- - input_array : `numpy.ndarray` + input_array : Input array to be transformed should be 2D (m, n) - uv : `numpy.array` - Array of 2xN u, v coordinates where the visibilities will be evaluated - center : `float` (x, y), optional - Coordinates of the center of the map e.g. ``(0,0)`` or ``[5.0, -2.0]`` + u : + Array of 2xN u coordinates where the visibilities are evaluated. + v : + Array of 2xN v coordinates where the visibilities are evaluated. + phase_centre : + Coordinates of the phase_centre of the map e.g. ``(0,0)`` or ``[5.0, -2.0]``. pixel_size : `float` (dx, dy), optional - The pixel size in x and y directions, need not be square e.g. ``(1, 3)`` + The pixel size need not be square e.g. ``(1, 3)``. Returns ------- - `numpy.ndarray` - Array of N `complex` visibilities evaluated at the u, v coordinates given by `uv` + : + Array of N `complex` visibilities evaluated at the given `u`, `v` coordinates. """ - m, n = input_array.shape + m, n = input_array.shape * apu.pix + x = generate_xy(m, phase_centre=phase_centre[0], pixel_size=pixel_size[0]) # type: ignore + y = generate_xy(n, phase_centre=phase_centre[1], pixel_size=pixel_size[1]) # type: ignore - y = generate_xy(m, center[1], pixel_size[1]) - x = generate_xy(n, center[0], pixel_size[0]) - - x, y = np.meshgrid(x, y) + x, y = np.meshgrid(x, y, indexing="ij") uv = np.vstack([u, v]) # Check units are correct for exp need to be dimensionless and then remove units for speed if (uv[0, :] * x[0, 0]).unit == apu.dimensionless_unscaled and ( uv[1, :] * y[0, 0] ).unit == apu.dimensionless_unscaled: - uv = uv.value + uv = uv.value # type: ignore x = x.value y = y.value @@ -147,38 +182,48 @@ def dft_map(input_array, *, u, v, center=(0.0, 0.0) * apu.arcsec, pixel_size=(1. ) -@apu.quantity_input(center="angle", pixel_size="angle") +@apu.quantity_input def idft_map( - input_vis, *, u, v, shape, weights=None, center=(0.0, 0.0) * apu.arcsec, pixel_size=(1.0, 1.0) * apu.arcsec -): + input_vis: Union[Quantity, npt.NDArray], + *, + u: Quantity[1 / apu.arcsec], + v: Quantity[1 / apu.arcsec], + shape: Quantity[apu.pix], + weights: Optional[npt.NDArray] = None, + phase_centre: Quantity[apu.arcsec] = (0.0, 0.0) * apu.arcsec, + pixel_size: Quantity[apu.arcsec / apu.pix] = (1.0, 1.0) * apu.arcsec / apu.pix, +) -> Union[Quantity, npt.NDArray]: r""" Inverse discrete Fourier transform in terms of coordinates returning a 2D real array or image. Parameters ---------- - uv : `numpy.ndarray` - Array of 2xN u, v coordinates corresponding to the input visibilities in `input_vis` - input_vis : `numpy.ndarray` - Array of N `complex` input visibilities - shape : `float` (m,n) - The shape of the output array to create - weights : `numpy.ndarray` - Array of weights for visibilities - center : `float` (x, y), optional - Coordinates of the center of the map e.g. ``(0,0)`` or ``[5.0, -2.0]`` - pixel_size : `float` (dx, dy), optional - The pixel size in x and y directions, need not be square e.g. ``(1, 3)`` + input_vis : + Input array of N complex visibilities to be transformed to a 2D array. + u : + Array of N u coordinates corresponding to the input visibilities in `input_vis` + v : + Array of N v coordinates corresponding to the input visibilities in `input_vis` + shape : + The shape of the output array to create. + weights : + Array of weights for visibilities. + phase_centre : + Coordinates of the phase_centre of the map e.g. ``(0,0)`` or ``[5.0, -2.0]``. + pixel_size : + The pixel size this need not be square e.g. ``(1, 3)``. Returns ------- - `numpy.ndarray` - The complex visibilities evaluated at the u, v coordinates + : + 2D image obtained from the visibilities evaluated at the given `u`, `v` coordinates. """ m, n = shape - y = generate_xy(m, center[1], pixel_size[1]) - x = generate_xy(n, center[0], pixel_size[0]) - x, y = np.meshgrid(x, y) + x = generate_xy(m, phase_centre=phase_centre[0], pixel_size=pixel_size[0]) # type: ignore + y = generate_xy(n, phase_centre=phase_centre[1], pixel_size=pixel_size[1]) # type: ignore + + x, y = np.meshgrid(x, y, indexing="ij") if weights is None: weights = np.ones(input_vis.shape) @@ -187,7 +232,7 @@ def idft_map( if (uv[0, :] * x[0, 0]).unit == apu.dimensionless_unscaled and ( uv[1, :] * y[0, 0] ).unit == apu.dimensionless_unscaled: - uv = uv.value + uv = uv.value # type: ignore x = x.value y = y.value diff --git a/xrayvision/utils.py b/xrayvision/utils.py index 2935b24..bb9789f 100644 --- a/xrayvision/utils.py +++ b/xrayvision/utils.py @@ -1,7 +1,8 @@ import logging +from typing import Union, Optional -def get_logger(name, level=logging.WARNING): +def get_logger(name: str, level: Optional[Union[int, str]] = logging.WARNING) -> logging.Logger: """ Return a configured logger instance. @@ -18,9 +19,9 @@ def get_logger(name, level=logging.WARNING): Configured logger """ logger = logging.getLogger(name) - logger.setLevel(level) + logger.setLevel(level) # type: ignore handler = logging.StreamHandler() - handler.setLevel(level) + handler.setLevel(level) # type: ignore formatter = logging.Formatter( "%(asctime)s %(levelname)s %(name)s %(lineno)s: %(message)s", datefmt="%Y-%m-%dT%H:%M:%SZ" ) diff --git a/xrayvision/visibility.py b/xrayvision/visibility.py index 2d4fb75..66a5398 100644 --- a/xrayvision/visibility.py +++ b/xrayvision/visibility.py @@ -1,5 +1,5 @@ """ -Modules contains visibility related classes. +Module contains visibility related classes. This contains classes to hold general visibilities and specialised classes hold visibilities from certain spacecraft or instruments @@ -7,7 +7,7 @@ import abc from typing import Union, Optional -from collections.abc import Iterable +from collections.abc import Iterable, Sequence import astropy.units as apu import numpy as np @@ -17,6 +17,8 @@ __all__ = ["Visibility", "Visibilities", "VisMeta", "VisibilitiesABC", "VisMetaABC"] +from astropy.units import Quantity + _E_RANGE_KEY = "spectral_range" _T_RANGE_KEY = "time_range" _OBS_COORD_KEY = "observer_coordinate" @@ -34,21 +36,21 @@ def observer_coordinate(self) -> Union[SkyCoord, None]: @property @abc.abstractmethod - def spectral_range(self) -> Union[Iterable[apu.Quantity], None]: + def spectral_range(self) -> Optional[Iterable[apu.Quantity]]: """ Spectral range over which the visibilities are computed. """ @property @abc.abstractmethod - def time_range(self) -> Union[Iterable[Time], None]: + def time_range(self) -> Optional[Iterable[Time]]: """ Time range over which the visibilities are computed. """ @property @abc.abstractmethod - def vis_labels(self) -> Union[Iterable[str], None]: + def vis_labels(self) -> Sequence[Iterable[str]]: """ Labels of each visibility. """ @@ -133,6 +135,60 @@ def phase_uncertainty(self) -> Union[Iterable[apu.Quantity[apu.deg]], None]: """ +class VisMeta(VisMetaABC, dict): + """ + A class for holding Visibility-specific metadata. + + Parameters + ---------- + meta: `dict` + A dictionary of the metadata + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # Check controlled/expected inputs are of correct type and units. + controled_args = ( + (_OBS_COORD_KEY, SkyCoord), + (_E_RANGE_KEY, apu.Quantity, apu.keV, apu.spectral()), + (_T_RANGE_KEY, Time), + ) + for args in controled_args: + self._check_input_type_and_unit(*args) + + def _check_input_type_and_unit(self, key, key_type, unit=None, equivalencies=None): + value = self.get(key, None) + if not isinstance(value, (key_type, type(None))): + raise KeyError(f"Inputs must include a key, '{key}', that gives a {key_type}.") + if unit is not None and value is not None and not value.unit.is_equivalent(unit, equivalencies=equivalencies): + raise ValueError(f"'{key}' must have angular units.") + + @property + def observer_coordinate(self): + return self.get(_OBS_COORD_KEY, None) + + @property + def spectral_range(self): + return self.get(_E_RANGE_KEY, None) + + @property + def time_range(self): + return self.get(_T_RANGE_KEY, None) + + @property + def vis_labels(self): + return self.get(_VIS_LABELS_KEY, None) + + @property + def instrument(self): + instr = None + i, n = 0, len(_INSTR_KEYS) + while not instr and i < n: + instr = self.get(_INSTR_KEYS[i], None) + i += 1 + return instr + + class Visibilities(VisibilitiesABC): @apu.quantity_input() def __init__( @@ -224,6 +280,10 @@ def __init__( self._uv_key = "uv" self._units_key = "units" + # Build meta. Make sure that phase center is included. + if not isinstance(meta, VisMetaABC): + meta = VisMeta(meta) + # Construct underlying data object. # In case visibilities is multi-dimensional, assume last axis is the uv-axis. # and give other axes arbitrary names. @@ -244,12 +304,12 @@ def __init__( data[self._phase_uncert_key] = (dims, phase_uncertainty.to_value(phase.unit)) if meta is None: meta = VisMeta(dict()) - vis_labels = meta.vis_labels + vis_labels = getattr(meta, "vis_labels", None) if vis_labels is not None: if len(vis_labels) != nvis: raise ValueError( "meta.vis_labels must be same length as number of visibilites. " - f"Number of labels = {len(meta.vis_labels)}; " + f"Number of labels = {len(vis_labels)}; " f"Number of visibilities = {nvis}" ) coords[_VIS_LABELS_KEY] = ([self._uv_key], vis_labels) @@ -353,113 +413,48 @@ def __repr__(self): return f"{self.__class__.__name__}< {self.u.size}, {self.visibilities}>" -class VisMeta(VisMetaABC, dict): - """ - A class for holding Visibility-specific metadata. - - Parameters - ---------- - meta: `dict` - A dictionary of the metadata - """ - - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - # Check controlled/expected inputs are of correct type and units. - controled_args = ( - (_OBS_COORD_KEY, SkyCoord), - (_E_RANGE_KEY, apu.Quantity, apu.keV, apu.spectral()), - (_T_RANGE_KEY, Time), - ) - for args in controled_args: - self._check_input_type_and_unit(*args) - - def _check_input_type_and_unit(self, key, key_type, unit=None, equivalencies=None): - value = self.get(key, None) - if not isinstance(value, (key_type, type(None))): - raise KeyError(f"Inputs must include a key, '{key}', that gives a {key_type}.") - if unit is not None and value is not None and not value.unit.is_equivalent(unit, equivalencies=equivalencies): - raise ValueError(f"'{key}' must have angular units.") - - @property - def observer_coordinate(self): - return self.get(_OBS_COORD_KEY, None) - - @property - def spectral_range(self): - return self.get(_E_RANGE_KEY, None) - - @property - def time_range(self): - return self.get(_T_RANGE_KEY, None) - - @property - def vis_labels(self): - return self.get(_VIS_LABELS_KEY, None) - - @property - def instrument(self): - instr = None - i, n = 0, len(_INSTR_KEYS) - while not instr and i < n: - instr = self.get(_INSTR_KEYS[i], None) - i += 1 - return instr - - -class BaseVisibility: - r""" - Base visibility containing bare essential fields, u, v, and complex vis - """ - - @apu.quantity_input(u=1 / apu.arcsec, v=1 / apu.arcsec, center=apu.arcsec) - def __int__(self, u, v, vis, center=(0, 0) * apu.arcsec): - self.u = u - self.v = v - self.vis = vis - self.center = center - - class Visibility: r""" Hold a set of related visibilities and information. - Attributes - ---------- - vis : `numpy.ndarray` - Array of N complex visibilities at coordinates in `uv` - u : `numpy.ndarray` - Array of `u` coordinates where visibilities will be evaluated - v : `numpy.ndarray` - Array of `v` coordinates where visibilities will be evaluated - center : `float` (x, y), optional - The x, y offset of phase center """ - @apu.quantity_input(uv=1 / apu.arcsec, offset=apu.arcsec, center=apu.arcsec, pixel_size=apu.arcsec) - def __init__(self, vis, *, u, v, offset=(0.0, 0.0) * apu.arcsec, center=(0.0, 0.0) * apu.arcsec): + @apu.quantity_input + def __init__( + self, + vis, + *, + u: Quantity[1 / apu.arcsec], + v: Quantity[1 / apu.arcsec], + offset: Optional[Quantity[apu.arcsec]] = (0.0, 0.0) * apu.arcsec, + phase_centre: Optional[Quantity[apu.arcsec]] = (0.0, 0.0) * apu.arcsec, + ) -> None: r""" - Initialise a new Visibility object. + Generic Visibility object. Parameters ---------- - vis : `numpy.ndarray` - Array of N complex visibilities at coordinates in `uv`. - u : `numpy.ndarray` - Array of `u` coordinates where visibilities will be evaluated. - v : `numpy.ndarray` - Array of `v` coordinates where visibilities will be evaluated. - center : - Phase centre - """ - self.u = u - self.v = v - self.vis = vis - self.center = center - self.offset = offset - - def __repr__(self): + vis: + Array of N complex visibilities sampled at the `u`, `v` coordinates. + u: + Array of `u` coordinates where visibilities are sampled. + v: + Array of `v` coordinates where visibilities are sampled. + phase_centre: + Phase centre of the visibility, defaults to (0,0). + offset: + Offset of the phase_centre visibility, defaults to (0,0). + + + """ + self.u: Quantity[1 / apu.arcsec] = u + self.v: Quantity[1 / apu.arcsec] = v + self.vis: Quantity = vis + self.phase_centre: Quantity[apu.arcsec] = phase_centre + self.offset: Quantity[apu.arcsec] = offset + + def __repr__(self) -> str: r""" Return a printable representation of the visibility. @@ -470,7 +465,7 @@ def __repr__(self): """ return f"{self.__class__.__name__}< {self.u.size}, {self.vis}>" - def __eq__(self, other): + def __eq__(self, other) -> bool: r""" Equality for Visibility class @@ -490,5 +485,5 @@ def __eq__(self, other): if all(props_equal): return True - else: - return False + + return False