Merge pull request #23 from ConorMacBride/rewrite-plot-map

Rewrite visualisation.plot_map

src/mcalf/tests/figure_hashes_mpl_334_ft_261.json

@@ -30,5 +30,19 @@
   "mcalf.tests.visualisation.test_spec.test_plot_spectrum": "79afa94c229e75646304b14bb388df30f23d6fd06ac82cfa4443ab05fc4c4a4a",
   "mcalf.tests.visualisation.test_spec.test_plot_spectrum_no_smooth": "c4bf62988715825cecaf75e03b66c05231854e5b84962fbd03f11e1c381e1273",
   "mcalf.tests.visualisation.test_spec.test_plot_spectrum_no_norm": "f4a52acdbb2a614e0a775ff7833a95151bf6188126286d1eb87fedef7f324711",
-  "mcalf.tests.visualisation.test_spec.test_plot_spectrum_no_smooth_norm": "96ce04857fe6fdc479c1620e72a0d26c981dd82082da2627e92e9716c46b2182"
+  "mcalf.tests.visualisation.test_spec.test_plot_spectrum_no_smooth_norm": "96ce04857fe6fdc479c1620e72a0d26c981dd82082da2627e92e9716c46b2182",
+  "mcalf.tests.visualisation.test_velocity.test_plot_map_basic": "7aafc54666ffd269480caecd10e0606fd0f59b1f43ed180c230421671df8bd63",
+  "mcalf.tests.visualisation.test_velocity.test_plot_map_mask": "a8d3af74f08a110ea6404aa782d5cf09820524b0cd9ec43ea22366757a922037",
+  "mcalf.tests.visualisation.test_velocity.test_plot_map_umbra_mask": "c8e305a8ab5da653387a585ea21a74d1d53020bb7d43657aecf4b47429b7b813",
+  "mcalf.tests.visualisation.test_velocity.test_plot_map_both_masks": "2e3313d66874786275d97589db95a106a75ba5328bcbf6626657722297445708",
+  "mcalf.tests.visualisation.test_velocity.test_plot_map_resolution": "62354784ccf0677f4837524f1887f6e15223940f7dca919b628bfcbba4d1d406",
+  "mcalf.tests.visualisation.test_velocity.test_plot_map_resolution_offset": "d1b73da96b58912d6e43e49885cfb263a633cd0a61a577ed07881dfde5ac6b5d",
+  "mcalf.tests.visualisation.test_velocity.test_plot_map_resolution_offset_masks": "6aa74ae92116d4da2cb201bf59f95a9fda71ebfd0ec847a9c1536248c7e35b9e",
+  "mcalf.tests.visualisation.test_velocity.test_plot_map_vmin": "9eb4e837594b4d515fefc5712ba559769c104ba2b390d69f74ddd687e1fc46ab",
+  "mcalf.tests.visualisation.test_velocity.test_plot_map_vmax": "9eb4e837594b4d515fefc5712ba559769c104ba2b390d69f74ddd687e1fc46ab",
+  "mcalf.tests.visualisation.test_velocity.test_plot_map_unit": "d9301400ee371e5af77f8f0d85e82000834bec0128a8afdfa675328546a3697c",
+  "mcalf.tests.visualisation.test_velocity.test_plot_map_unit_astropy": "9673f5206be729915f5843d48d9afecba1cde19c2638536b8fc6dbe9e8641c4f",
+  "mcalf.tests.visualisation.test_velocity.test_plot_map_unit_arr_astropy": "9673f5206be729915f5843d48d9afecba1cde19c2638536b8fc6dbe9e8641c4f",
+  "mcalf.tests.visualisation.test_velocity.test_plot_map_lw": "98013f267d1517c05f64e4b0d32960dc9b333ad3277152b4a8f33e3a479eb16b",
+  "mcalf.tests.visualisation.test_velocity.test_plot_map_colorbar": "c92b8186ce40312c307b5dceb491d2a47e6b46ce260a190d59fdaceb10970464"
 }

src/mcalf/tests/utils/test_misc.py

@@ -2,8 +2,9 @@
 import numpy as np
 import matplotlib.pyplot as plt
 from astropy.io import fits
+import astropy.units
 
-from mcalf.utils.misc import make_iter, load_parameter, merge_results, hide_existing_labels
+from mcalf.utils.misc import make_iter, load_parameter, merge_results, hide_existing_labels, calculate_extent
 
 from ..helpers import data_path_function
 data_path = data_path_function('utils')
@@ -129,3 +130,38 @@ def test_hide_existing_labels():
     plt.close()
     ret = [x['label'] for x in plot_settings.values()]
     assert ret == ['_A', '_B', 'C']
+
+
+def test_calculate_extent():
+    # Invalid `res`
+    with pytest.raises(TypeError) as e:
+        calculate_extent('1 m', 1000)
+    assert '`resolution` values must be' in str(e.value)
+
+    # Invalid `px`
+    with pytest.raises(TypeError) as e:
+        calculate_extent(250.0, 1000.5)
+    assert '`px` must be an integer' in str(e.value)
+
+    # Invalid `offset`
+    with pytest.raises(TypeError) as e:
+        calculate_extent(250.0, 1000, [10])
+    assert '`offset` must be an float or integer' in str(e.value)
+
+    # Default `unit` used
+    _, _, un = calculate_extent(250.0, 1000, 10)
+    assert un == 'Mm'  # Will fail if default value is changed
+    _, _, un = calculate_extent(250.0, 1000, 10, 'testunit')
+    assert un == 'testunit'  # Will fail if default value is changed
+
+    # Unit extracted
+    _, _, un = calculate_extent(250.0 * astropy.units.kg, 1000, 10, 'testunit')
+    assert un == '$\\mathrm{kg}$'
+
+    # Test good values
+    f, l, _ = calculate_extent(3., 7, -3.5)
+    assert -10.5 == pytest.approx(f)
+    assert 10.5 == pytest.approx(l)
+    f, l, _ = calculate_extent(2., 2, 1)
+    assert 2. == pytest.approx(f)
+    assert 6. == pytest.approx(l)

src/mcalf/tests/visualisation/test_velocity.py

@@ -0,0 +1,148 @@
+import pytest
+import numpy as np
+import astropy.units as u
+
+from mcalf.visualisation import plot_map
+from ..helpers import figure_test
+
+
+@pytest.fixture
+def arr():
+    np.random.seed(1234)
+    a = np.random.rand(8, 6)
+    a[3, 3] = np.nan
+    a[5, 1] = np.nan
+    return (a * 16) - 8
+
+
+@pytest.fixture
+def mask(arr):
+    mask = np.full_like(arr, True, dtype=bool)
+    mask[1:7, 2:] = False
+    return mask
+
+
+@pytest.fixture
+def umbra_mask(arr):
+    mask = np.full_like(arr, True, dtype=bool)
+    mask[2:6, 2:5] = False
+    mask[4, 2] = True
+    mask[2, 4] = True
+    return mask
+
+
+@figure_test
+def test_plot_map_basic(pytestconfig, arr):
+    plot_map(arr)
+
+
+@figure_test
+def test_plot_map_mask(pytestconfig, arr, mask):
+    plot_map(arr, mask)
+
+
+@figure_test
+def test_plot_map_umbra_mask(pytestconfig, arr, umbra_mask):
+    plot_map(arr, umbra_mask=umbra_mask)
+
+
+@figure_test
+def test_plot_map_both_masks(pytestconfig, arr, mask, umbra_mask):
+    plot_map(arr, mask, umbra_mask)
+
+
+@figure_test
+def test_plot_map_resolution(pytestconfig, arr):
+    plot_map(arr, resolution=(2.5, 3 * u.kg / u.Hz))
+
+
+@figure_test
+def test_plot_map_resolution_offset(pytestconfig, arr):
+    plot_map(arr, resolution=(2.5, 3 * u.kg / u.Hz), offset=(-3, -4))
+
+
+@figure_test
+def test_plot_map_resolution_offset_masks(pytestconfig, arr, mask, umbra_mask):
+    plot_map(arr, mask, umbra_mask, resolution=(2.5, 3 * u.kg / u.Hz), offset=(-3, -4))
+
+
+@figure_test
+def test_plot_map_vmin(pytestconfig, arr):
+    plot_map(arr, vmin=-4.5)
+
+
+@figure_test
+def test_plot_map_vmax(pytestconfig, arr):
+    plot_map(arr, vmax=4.5)
+
+
+@figure_test
+def test_plot_map_unit(pytestconfig, arr):
+    plot_map(arr, unit='test/unit')
+
+
+@figure_test
+def test_plot_map_unit_astropy(pytestconfig, arr):
+    plot_map(arr, unit=(u.m / u.s))
+
+
+@figure_test
+def test_plot_map_unit_arr_astropy(pytestconfig, arr):
+    plot_map(arr * u.m / u.s, unit='test/unit')
+
+
+@figure_test
+def test_plot_map_lw(pytestconfig, arr, umbra_mask):
+    plot_map(arr, umbra_mask=umbra_mask, lw=5.)
+
+
+@figure_test
+def test_plot_map_colorbar(pytestconfig, arr):
+    plot_map(arr, show_colorbar=False)
+
+
+def test_plot_map_validate_arr(arr):
+    for a in ([[1, 2, 3], [4, 5, 6], [7, 8, 9]], np.array([arr, arr])):
+        with pytest.raises(TypeError) as e:
+            plot_map(a)
+        assert '`arr` must be a numpy.ndarray with 2 dimensions' in str(e.value)
+
+
+def test_plot_map_validate_masks(arr, mask, umbra_mask):
+
+    # Mask is 2D array
+    for a in ([[1, 2, 3], [4, 5, 6], [7, 8, 9]], np.array([mask, mask])):
+        with pytest.raises(TypeError) as e:
+            plot_map(arr, a, umbra_mask)
+        assert '`mask` must be a numpy.ndarray with 2 dimensions' in str(e.value)
+
+    # Umbra mask is 2D array
+    for a in ([[1, 2, 3], [4, 5, 6], [7, 8, 9]], np.array([umbra_mask, umbra_mask])):
+        with pytest.raises(TypeError) as e:
+            plot_map(arr, mask, a)
+        assert '`umbra_mask` must be a numpy.ndarray with 2 dimensions' in str(e.value)
+
+    # Mask wrong shape
+    with pytest.raises(ValueError) as e:
+        plot_map(arr, np.vstack([mask, mask]), umbra_mask)
+    assert '`mask` must be the same shape as `arr`' in str(e.value)
+
+    # Umbra mask wrong shape
+    with pytest.raises(ValueError) as e:
+        plot_map(arr, mask, np.vstack([umbra_mask, umbra_mask]))
+    assert '`umbra_mask` must be the same shape as `arr`' in str(e.value)
+
+
+def test_plot_map_validate_resolution(arr):
+
+    # Resolution is 2-tuple
+    for a in ((1, 1, 1), np.array([1, 1])):
+        with pytest.raises(TypeError) as e:
+            plot_map(arr, resolution=a, offset=(1, 1))
+        assert '`resolution` must be a tuple of length 2' in str(e.value)
+
+    # Offset is 2-tuple
+    for a in ((1, 1, 1), np.array([1, 1])):
+        with pytest.raises(TypeError) as e:
+            plot_map(arr, resolution=(1, 1), offset=a)
+        assert '`offset` must be a tuple of length 2' in str(e.value)

src/mcalf/utils/misc.py

@@ -4,10 +4,11 @@
 import numpy as np
 import matplotlib.pyplot as plt
 from astropy.io import fits
+import astropy.units
 from scipy.io import readsav
 
 
-__all__ = ['make_iter', 'load_parameter', 'merge_results', 'hide_existing_labels']
+__all__ = ['make_iter', 'load_parameter', 'merge_results', 'hide_existing_labels', 'calculate_extent']
 
 
 def make_iter(*args):
@@ -436,3 +437,50 @@ def hide_existing_labels(plot_settings, axes=None, fig=None):
     for name in plot_settings:
         if plot_settings[name]['label'] in existing:
             plot_settings[name]['label'] = '_' + plot_settings[name]['label']
+
+
+def calculate_extent(res, px, offset=0, unit="Mm"):
+    """Calculate the extent from a resolution value.
+
+    Parameters
+    ----------
+    res : float or astropy.units.quantity.Quantity
+        Length of each pixel. Unit defaults to `unit` is not an astropy quantity.
+    px : int
+        Number of pixels extent is being calculated for.
+    offset : int or float, default=0
+        Number of pixels from the 0 pixel to the first pixel. Defaults to the first
+        pixel being at 0 length units. For example, in a 1000 pixel wide dataset,
+        setting offset to -500 would place the 0 Mm location at the centre.
+    unit : str, default="Mm"
+        Default unit string to use if `res` is not an astropy quantity.
+
+    Returns
+    -------
+    first : float
+        First extent value.
+    last : float
+        Last extent value.
+    unit : str
+        Unit of extent values.
+    """
+
+    # Ensure a valid spatial and pixel resolution is provided
+    if not isinstance(res, (float, astropy.units.quantity.Quantity)):
+        raise TypeError('`resolution` values must be either floats or astropy quantities'
+                        f', got {type(res)}.')
+    if not isinstance(px, (int, np.integer)):
+        raise TypeError(f'`px` must be an integer, got {type(px)}.')
+    if not isinstance(offset, (float, int, np.integer)):
+        raise TypeError(f'`offset` must be an float or integer, got {type(offset)}.')
+
+    # Update the default unit if a quantity is provided
+    if isinstance(res, astropy.units.quantity.Quantity):
+        unit = res.unit.to_string(astropy.units.format.LatexInline)
+        res = float(res.value)  # Remove the unit
+
+    # Calculate the extent values
+    first = offset * res
+    last = (px + offset) * res
+
+    return first, last, unit