|
| 1 | +""" |
| 2 | +Tests to simulate dynamic spectrum WCSes (frequency x time). |
| 3 | +""" |
| 4 | +import pytest |
| 5 | +from numpy.testing import assert_allclose |
| 6 | + |
| 7 | +import astropy.units as u |
| 8 | + |
| 9 | +from ndcube.wcs.wrappers import ResampledLowLevelWCS |
| 10 | + |
| 11 | + |
| 12 | +def _world_at(cube, time_pixel, freq_pixel): |
| 13 | + return cube.wcs.low_level_wcs.pixel_to_world_values(time_pixel, freq_pixel) |
| 14 | + |
| 15 | + |
| 16 | +@pytest.mark.parametrize("ndc", [ |
| 17 | + "ndcube_gwcs_2d_t_f_linear", |
| 18 | + "ndcube_gwcs_2d_t_f_log", |
| 19 | +], indirect=True) |
| 20 | +def test_dynspec_array_axis_physical_types(ndc): |
| 21 | + types = ndc.array_axis_physical_types |
| 22 | + assert "em.freq" in types[0] |
| 23 | + assert "time" in types[1] |
| 24 | + |
| 25 | + |
| 26 | +def test_linear_dynspec_pixel_to_world(ndcube_gwcs_2d_t_f_linear): |
| 27 | + time, freq = ndcube_gwcs_2d_t_f_linear.wcs.low_level_wcs.pixel_to_world_values(3, 2) |
| 28 | + assert_allclose(time, 42.0) |
| 29 | + assert_allclose(freq, 2e6) |
| 30 | + |
| 31 | + |
| 32 | +def test_linear_dynspec_world_to_pixel(ndcube_gwcs_2d_t_f_linear): |
| 33 | + pix_t, pix_f = ndcube_gwcs_2d_t_f_linear.wcs.low_level_wcs.world_to_pixel_values(28.0, 4e6) |
| 34 | + assert_allclose(pix_t, 2.0) |
| 35 | + assert_allclose(pix_f, 4.0) |
| 36 | + |
| 37 | + |
| 38 | +@pytest.mark.parametrize(("bin_shape", "expected_shape", "expected_time", "expected_freq"), [ |
| 39 | + ((2, 1), (8, 10), 0.0, 0.5e6), |
| 40 | + ((1, 2), (16, 5), 7.0, 0.0), |
| 41 | +]) |
| 42 | +def test_linear_dynspec_rebin_wcs(ndcube_gwcs_2d_t_f_linear, bin_shape, |
| 43 | + expected_shape, expected_time, expected_freq): |
| 44 | + rebinned = ndcube_gwcs_2d_t_f_linear.rebin(bin_shape) |
| 45 | + time0, freq0 = rebinned.wcs.low_level_wcs.pixel_to_world_values(0, 0) |
| 46 | + |
| 47 | + assert rebinned.shape == expected_shape |
| 48 | + assert isinstance(rebinned.wcs.low_level_wcs, ResampledLowLevelWCS) |
| 49 | + assert_allclose(time0, expected_time) |
| 50 | + assert_allclose(freq0, expected_freq) |
| 51 | + |
| 52 | + |
| 53 | +@pytest.mark.parametrize(("lower_corner", "upper_corner", "expected_shape"), [ |
| 54 | + ([None, 3e6 * u.Hz], [None, 7e6 * u.Hz], (5, 10)), |
| 55 | + ([14 * u.s, None], [56 * u.s, None], (16, 4)), |
| 56 | +]) |
| 57 | +def test_linear_dynspec_crop_by_values_shape(ndcube_gwcs_2d_t_f_linear, |
| 58 | + lower_corner, upper_corner, |
| 59 | + expected_shape): |
| 60 | + cropped = ndcube_gwcs_2d_t_f_linear.crop_by_values(lower_corner, upper_corner) |
| 61 | + assert cropped.shape == expected_shape |
| 62 | + |
| 63 | + |
| 64 | +def test_log_dynspec_world_axis_units(ndcube_gwcs_2d_t_f_log): |
| 65 | + assert ndcube_gwcs_2d_t_f_log.wcs.world_axis_units == ("s", "Hz") |
| 66 | + |
| 67 | + |
| 68 | +@pytest.mark.parametrize(("time_pixel", "freq_pixel", "expected_time", "expected_freq"), [ |
| 69 | + (0, 0, 0.0, 3.992e6), |
| 70 | + (9, 15, 122.5, 978.572e6), |
| 71 | +]) |
| 72 | +def test_log_dynspec_pixel_to_world_endpoints(ndcube_gwcs_2d_t_f_log, |
| 73 | + time_pixel, freq_pixel, |
| 74 | + expected_time, expected_freq): |
| 75 | + time, freq = ndcube_gwcs_2d_t_f_log.wcs.low_level_wcs.pixel_to_world_values( |
| 76 | + time_pixel, freq_pixel) |
| 77 | + assert_allclose(time, expected_time) |
| 78 | + assert_allclose(freq, expected_freq, rtol=1e-6) |
| 79 | + |
| 80 | + |
| 81 | +def test_log_dynspec_world_to_pixel_roundtrip(ndcube_gwcs_2d_t_f_log): |
| 82 | + time, freq = _world_at(ndcube_gwcs_2d_t_f_log, 3, 7) |
| 83 | + pix_t, pix_f = ndcube_gwcs_2d_t_f_log.wcs.low_level_wcs.world_to_pixel_values( |
| 84 | + time, freq) |
| 85 | + assert_allclose(pix_t, 3.0, atol=1e-10) |
| 86 | + assert_allclose(pix_f, 7.0, atol=1e-10) |
| 87 | + |
| 88 | + |
| 89 | +@pytest.mark.parametrize(("bin_shape", "expected_shape", "axis"), [ |
| 90 | + ((2, 1), (8, 10), "freq"), |
| 91 | + ((1, 2), (16, 5), "time"), |
| 92 | +]) |
| 93 | +def test_log_dynspec_rebin_wcs_midpoint(ndcube_gwcs_2d_t_f_log, bin_shape, |
| 94 | + expected_shape, axis): |
| 95 | + rebinned = ndcube_gwcs_2d_t_f_log.rebin(bin_shape) |
| 96 | + time0, freq0 = rebinned.wcs.low_level_wcs.pixel_to_world_values(0, 0) |
| 97 | + |
| 98 | + assert rebinned.shape == expected_shape |
| 99 | + assert isinstance(rebinned.wcs.low_level_wcs, ResampledLowLevelWCS) |
| 100 | + if axis == "freq": |
| 101 | + _, freq_left = _world_at(ndcube_gwcs_2d_t_f_log, 0, 0) |
| 102 | + _, freq_right = _world_at(ndcube_gwcs_2d_t_f_log, 0, 1) |
| 103 | + assert_allclose(freq0, (freq_left + freq_right) / 2, rtol=1e-6) |
| 104 | + else: |
| 105 | + time_left, _ = _world_at(ndcube_gwcs_2d_t_f_log, 0, 0) |
| 106 | + time_right, _ = _world_at(ndcube_gwcs_2d_t_f_log, 1, 0) |
| 107 | + assert_allclose(time0, (time_left + time_right) / 2, rtol=1e-6) |
| 108 | + |
| 109 | + |
| 110 | +@pytest.mark.parametrize(("lower_corner", "upper_corner", "expected_shape", |
| 111 | + "axis", "bounds"), [ |
| 112 | + ([None, 10e6 * u.Hz], [None, 100e6 * u.Hz], (8, 10), "freq", (10e6, 100e6)), |
| 113 | + ([20 * u.s, None], [80 * u.s, None], (16, 6), "time", (20.0, 80.0)), |
| 114 | +]) |
| 115 | +def test_log_dynspec_crop_by_values_single_axis(ndcube_gwcs_2d_t_f_log, |
| 116 | + lower_corner, upper_corner, |
| 117 | + expected_shape, axis, bounds): |
| 118 | + cropped = ndcube_gwcs_2d_t_f_log.crop_by_values(lower_corner, upper_corner) |
| 119 | + assert cropped.shape == expected_shape |
| 120 | + |
| 121 | + if axis == "freq": |
| 122 | + values = [cropped.wcs.low_level_wcs.pixel_to_world_values(0, i)[1] |
| 123 | + for i in range(cropped.shape[0])] |
| 124 | + else: |
| 125 | + values = [cropped.wcs.low_level_wcs.pixel_to_world_values(i, 0)[0] |
| 126 | + for i in range(cropped.shape[1])] |
| 127 | + |
| 128 | + assert values[0] <= bounds[0] |
| 129 | + assert values[-1] >= bounds[1] |
| 130 | + |
| 131 | + |
| 132 | +def test_log_dynspec_crop_by_freq_and_time(ndcube_gwcs_2d_t_f_log): |
| 133 | + cropped = ndcube_gwcs_2d_t_f_log.crop_by_values( |
| 134 | + [20 * u.s, 10e6 * u.Hz], [80 * u.s, 100e6 * u.Hz]) |
| 135 | + assert cropped.shape == (8, 6) |
0 commit comments