fitscube 2.7.0__tar.gz → 2.8.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (40) hide show
  1. {fitscube-2.7.0 → fitscube-2.8.0}/PKG-INFO +2 -2
  2. {fitscube-2.7.0 → fitscube-2.8.0}/fitscube/_version.py +2 -2
  3. {fitscube-2.7.0 → fitscube-2.8.0}/fitscube/combine_fits.py +203 -53
  4. {fitscube-2.7.0 → fitscube-2.8.0}/fitscube/exceptions.py +8 -0
  5. {fitscube-2.7.0 → fitscube-2.8.0}/pyproject.toml +1 -1
  6. {fitscube-2.7.0 → fitscube-2.8.0}/tests/test_combine.py +37 -12
  7. {fitscube-2.7.0 → fitscube-2.8.0}/tests/test_combine_guards.py +31 -1
  8. fitscube-2.8.0/tests/test_combine_layouts.py +328 -0
  9. {fitscube-2.7.0 → fitscube-2.8.0}/tests/test_frequencies.py +1 -1
  10. {fitscube-2.7.0 → fitscube-2.8.0}/tests/test_times.py +3 -1
  11. {fitscube-2.7.0 → fitscube-2.8.0}/uv.lock +1 -1
  12. {fitscube-2.7.0 → fitscube-2.8.0}/.github/CONTRIBUTING.md +0 -0
  13. {fitscube-2.7.0 → fitscube-2.8.0}/.github/dependabot.yml +0 -0
  14. {fitscube-2.7.0 → fitscube-2.8.0}/.github/release.yml +0 -0
  15. {fitscube-2.7.0 → fitscube-2.8.0}/.github/workflows/cd.yml +0 -0
  16. {fitscube-2.7.0 → fitscube-2.8.0}/.github/workflows/ci.yml +0 -0
  17. {fitscube-2.7.0 → fitscube-2.8.0}/.gitignore +0 -0
  18. {fitscube-2.7.0 → fitscube-2.8.0}/.pre-commit-config.yaml +0 -0
  19. {fitscube-2.7.0 → fitscube-2.8.0}/CHANGELOG.md +0 -0
  20. {fitscube-2.7.0 → fitscube-2.8.0}/LICENSE +0 -0
  21. {fitscube-2.7.0 → fitscube-2.8.0}/README.md +0 -0
  22. {fitscube-2.7.0 → fitscube-2.8.0}/fitscube/__init__.py +0 -0
  23. {fitscube-2.7.0 → fitscube-2.8.0}/fitscube/asyncio.py +0 -0
  24. {fitscube-2.7.0 → fitscube-2.8.0}/fitscube/bounding_box.py +0 -0
  25. {fitscube-2.7.0 → fitscube-2.8.0}/fitscube/cli.py +0 -0
  26. {fitscube-2.7.0 → fitscube-2.8.0}/fitscube/extract.py +0 -0
  27. {fitscube-2.7.0 → fitscube-2.8.0}/fitscube/logging.py +0 -0
  28. {fitscube-2.7.0 → fitscube-2.8.0}/fitscube/version.pyi +0 -0
  29. {fitscube-2.7.0 → fitscube-2.8.0}/noxfile.py +0 -0
  30. {fitscube-2.7.0 → fitscube-2.8.0}/tests/__init__.py +0 -0
  31. {fitscube-2.7.0 → fitscube-2.8.0}/tests/conftest.py +0 -0
  32. {fitscube-2.7.0 → fitscube-2.8.0}/tests/data/cube.zip +0 -0
  33. {fitscube-2.7.0 → fitscube-2.8.0}/tests/data/images.zip +0 -0
  34. {fitscube-2.7.0 → fitscube-2.8.0}/tests/data/time_images.zip +0 -0
  35. {fitscube-2.7.0 → fitscube-2.8.0}/tests/data/timecube.zip +0 -0
  36. {fitscube-2.7.0 → fitscube-2.8.0}/tests/test_bb.py +0 -0
  37. {fitscube-2.7.0 → fitscube-2.8.0}/tests/test_compress.py +0 -0
  38. {fitscube-2.7.0 → fitscube-2.8.0}/tests/test_extract.py +0 -0
  39. {fitscube-2.7.0 → fitscube-2.8.0}/tests/test_package.py +0 -0
  40. {fitscube-2.7.0 → fitscube-2.8.0}/tests/test_zero_beams.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.5
2
2
  Name: fitscube
3
- Version: 2.7.0
3
+ Version: 2.8.0
4
4
  Summary: A package to produce produce FITS cubes.
5
5
  Project-URL: Homepage, https://github.com/AlecThomson/fitscube
6
6
  Project-URL: Bug Tracker, https://github.com/AlecThomson/fitscube/issues
@@ -45,7 +45,7 @@ Classifier: Programming Language :: Python :: 3.14
45
45
  Classifier: Topic :: Scientific/Engineering
46
46
  Classifier: Typing :: Typed
47
47
  Requires-Python: >=3.10
48
- Requires-Dist: astropy
48
+ Requires-Dist: astropy>=6.0
49
49
  Requires-Dist: numpy>=2
50
50
  Requires-Dist: radio-beam
51
51
  Requires-Dist: tqdm
@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
18
18
  commit_id: str | None
19
19
  __commit_id__: str | None
20
20
 
21
- __version__ = version = '2.7.0'
22
- __version_tuple__ = version_tuple = (2, 7, 0)
21
+ __version__ = version = '2.8.0'
22
+ __version_tuple__ = version_tuple = (2, 8, 0)
23
23
 
24
24
  __commit_id__ = commit_id = None
@@ -4,6 +4,8 @@
4
4
  Assumes:
5
5
  - All files have the same WCS
6
6
  - All files have the same shape / pixel grid (checked, see `check_matching_shapes`)
7
+ - All files have the same axes, in the same order, and the same Stokes parameter
8
+ (checked, see `check_matching_axes`)
7
9
  - All the relevant information is in the first header of the first image
8
10
  - Frequency is either a WCS axis or in the REFFREQ header keyword OR
9
11
  - Time is present in the DATE-OBS header keyword for time-domain-mode
@@ -24,11 +26,12 @@ from typing import IO, Any, Literal, NamedTuple, TypeVar, cast
24
26
 
25
27
  import astropy.units as u
26
28
  import numpy as np
29
+ from astropy.coordinates import StokesCoord
27
30
  from astropy.io import fits
28
31
  from astropy.io.fits.verify import VerifyWarning
29
32
  from astropy.table import Table
30
33
  from astropy.time import Time
31
- from astropy.wcs import WCS
34
+ from astropy.wcs import WCS, WCSSUB_STOKES
32
35
  from numpy.typing import ArrayLike, NDArray
33
36
  from radio_beam import Beam, Beams
34
37
  from radio_beam.beam import NoBeamException
@@ -40,9 +43,11 @@ from fitscube.bounding_box import (
40
43
  get_common_bounding_box_coro,
41
44
  )
42
45
  from fitscube.exceptions import (
46
+ AxisMismatchException,
43
47
  AxisOrderException,
44
48
  IrregularSpacingException,
45
49
  ShapeMismatchException,
50
+ StokesMismatchException,
46
51
  )
47
52
  from fitscube.logging import TQDM_OUT, logger, set_verbosity
48
53
 
@@ -134,56 +139,141 @@ async def write_channel_to_cube_coro(
134
139
  write_channel_to_cube = sync_wrapper(write_channel_to_cube_coro)
135
140
 
136
141
 
137
- async def check_matching_shapes_coro(
142
+ async def read_headers_coro(
138
143
  file_list: list[Path],
139
144
  max_workers: int | None = None,
140
- ) -> tuple[int, int]:
141
- """Confirm every input image shares the same NAXIS1/NAXIS2 pixel grid.
142
-
143
- Planes are written to the cube at a fixed byte offset per channel, so
144
- inputs of differing shape would silently slide against each other and
145
- produce a scrambled cube with a self-consistent header.
145
+ ) -> list[fits.Header]:
146
+ """Read the primary header of every input image, without its data.
146
147
 
147
148
  Args:
148
- file_list (list[Path]): The FITS images to check
149
+ file_list (list[Path]): The FITS images to read
149
150
  max_workers (int | None, optional): Maximum number of concurrent header reads. Defaults to None.
150
151
 
151
- Raises:
152
- ShapeMismatchException: If any image differs in shape from the first
153
-
154
152
  Returns:
155
- tuple[int, int]: The common (NAXIS1, NAXIS2)
153
+ list[fits.Header]: One header per image, in ``file_list`` order
156
154
  """
157
- headers = await gather_with_limit(
155
+ return await gather_with_limit(
158
156
  max_workers,
159
157
  *(asyncio.to_thread(fits.getheader, fits_path) for fits_path in file_list),
160
- desc="Checking shapes",
158
+ desc="Reading headers",
161
159
  )
162
- shapes = [(header["NAXIS1"], header["NAXIS2"]) for header in headers]
163
- expected = shapes[0]
160
+
161
+
162
+ def check_values_match(
163
+ file_list: list[Path],
164
+ values: list[T],
165
+ describe: Callable[[T], str],
166
+ requirement: str,
167
+ exception: type[Exception],
168
+ ) -> T:
169
+ """Check that every input has the same value as the first.
170
+
171
+ Args:
172
+ file_list (list[Path]): The FITS images the values came from
173
+ values (list[T]): One value per image
174
+ describe (Callable[[T], str]): Describes a value for the error message
175
+ requirement (str): What the inputs must share, for the error message
176
+ exception (type[Exception]): Raised when any value differs from the first
177
+
178
+ Returns:
179
+ T: The common value
180
+ """
181
+ expected = values[0]
164
182
  offenders = [
165
- f"{fits_path} has (NAXIS1, NAXIS2)={shape}"
166
- for fits_path, shape in zip(file_list, shapes, strict=False)
167
- if shape != expected
183
+ f"{fits_path} has {describe(value)}"
184
+ for fits_path, value in zip(file_list, values, strict=True)
185
+ if value != expected
168
186
  ]
169
187
  if offenders:
170
188
  # Keep the message readable when a whole run is mismatched
171
- shown = offenders[:10]
189
+ max_shown = 10
190
+ shown = offenders[:max_shown]
172
191
  if len(offenders) > len(shown):
173
192
  shown.append(f"...and {len(offenders) - len(shown)} more")
174
193
  listing = "\n".join(shown)
175
194
  msg = (
176
- "All input images must share the same pixel grid. Expected "
177
- f"(NAXIS1, NAXIS2)={expected} from {file_list[0]}, but found:\n{listing}"
195
+ f"All input images must share {requirement}. Expected "
196
+ f"{describe(expected)} from {file_list[0]}, but found:\n{listing}"
178
197
  )
179
- raise ShapeMismatchException(msg)
180
-
198
+ raise exception(msg)
181
199
  return expected
182
200
 
183
201
 
202
+ async def check_matching_shapes_coro(
203
+ file_list: list[Path],
204
+ max_workers: int | None = None,
205
+ ) -> tuple[int, int]:
206
+ """Confirm every input image shares the same NAXIS1/NAXIS2 pixel grid.
207
+
208
+ Planes are written to the cube at a fixed byte offset per channel, so
209
+ inputs of differing shape would silently slide against each other and
210
+ produce a scrambled cube with a self-consistent header.
211
+
212
+ Args:
213
+ file_list (list[Path]): The FITS images to check
214
+ max_workers (int | None, optional): Maximum number of concurrent header reads. Defaults to None.
215
+
216
+ Raises:
217
+ ShapeMismatchException: If any image differs in shape from the first
218
+
219
+ Returns:
220
+ tuple[int, int]: The common (NAXIS1, NAXIS2)
221
+ """
222
+ headers = await read_headers_coro(file_list, max_workers=max_workers)
223
+ return check_values_match(
224
+ file_list,
225
+ [(header["NAXIS1"], header["NAXIS2"]) for header in headers],
226
+ describe=lambda shape: f"(NAXIS1, NAXIS2)={shape}",
227
+ requirement="the same pixel grid",
228
+ exception=ShapeMismatchException,
229
+ )
230
+
231
+
184
232
  check_matching_shapes = sync_wrapper(check_matching_shapes_coro)
185
233
 
186
234
 
235
+ async def check_matching_axes_coro(
236
+ file_list: list[Path],
237
+ max_workers: int | None = None,
238
+ ) -> None:
239
+ """Confirm every input image has the same axes and Stokes parameter.
240
+
241
+ The cube's header is taken from the first input, so an input with other
242
+ axes (e.g. RA/DEC swapped, or an extra axis) or another Stokes parameter
243
+ would be silently mislabelled.
244
+
245
+ Args:
246
+ file_list (list[Path]): The FITS images to check
247
+ max_workers (int | None, optional): Maximum number of concurrent header reads. Defaults to None.
248
+
249
+ Raises:
250
+ AxisMismatchException: If any image's axis types or their order differ from the first
251
+ StokesMismatchException: If any image differs in Stokes from the first
252
+ """
253
+ headers = await read_headers_coro(file_list, max_workers=max_workers)
254
+ check_values_match(
255
+ file_list,
256
+ [get_axis_types(header) for header in headers],
257
+ describe=lambda ctypes: f"axes {ctypes}",
258
+ requirement="the same axes, in the same order",
259
+ exception=AxisMismatchException,
260
+ )
261
+ check_values_match(
262
+ file_list,
263
+ [get_stokes_codes(header) for header in headers],
264
+ describe=lambda codes: (
265
+ "no Stokes axis"
266
+ if codes is None
267
+ else f"Stokes {StokesCoord(list(codes)).symbol.tolist()} (codes {codes})"
268
+ ),
269
+ requirement="the same Stokes parameter",
270
+ exception=StokesMismatchException,
271
+ )
272
+
273
+
274
+ check_matching_axes = sync_wrapper(check_matching_axes_coro)
275
+
276
+
187
277
  # https://stackoverflow.com/a/66082278
188
278
  def np_arange_fix(start: float, stop: float, step: float) -> NDArray[np.float64]:
189
279
  n = (stop - start) / step + 1
@@ -337,7 +427,7 @@ async def create_cube_from_scratch_coro(
337
427
 
338
428
  header.tofile(output_file, overwrite=overwrite)
339
429
 
340
- bytes_per_value = BIT_DICT.get(abs(output_header["BITPIX"]), None)
430
+ bytes_per_value = BIT_DICT.get(abs(output_header["BITPIX"]))
341
431
  msg = f"Header BITPIX={output_header['BITPIX']}, bytes_per_value={bytes_per_value}"
342
432
  logger.info(msg)
343
433
  if bytes_per_value is None:
@@ -418,7 +508,11 @@ async def create_output_cube_coro(
418
508
  diff_time = np.diff(sorted_specs)
419
509
  diff_diff_time = np.diff(diff_time)
420
510
  running_deviation_from_zero = np.abs(np.cumsum(diff_diff_time))
421
- even_spec = np.max(running_deviation_from_zero) < (np.mean(diff_time) * 0.02)
511
+ # Fewer than three times have no second difference, and are always
512
+ # evenly spaced
513
+ even_spec = len(diff_diff_time) == 0 or bool(
514
+ np.max(running_deviation_from_zero) < (np.mean(diff_time) * 0.02)
515
+ )
422
516
 
423
517
  # This is a simpler way where no attempt is made to ensure the total
424
518
  # error on the irregular steps accumulates and violates the regular
@@ -814,48 +908,99 @@ def nan_zero_beams(beams: Beams, zero_beam_idx: NDArray[np.bool_]) -> Beams:
814
908
  )
815
909
 
816
910
 
817
- def get_polarisation(header: fits.Header) -> int:
818
- """Get the polarisation axis.
911
+ def _stokes_wcs(header: fits.Header) -> WCS | None:
912
+ """Get the Stokes axis of an image as a 1D WCS.
819
913
 
820
914
  Args:
821
- header (fits.Header): Primary header
915
+ header (fits.Header): Primary header of an image
822
916
 
823
917
  Returns:
824
- int: Polarisation axis (in FITS)
918
+ WCS | None: The Stokes axis alone, or None if there is no Stokes axis
825
919
  """
826
920
  wcs = WCS(header)
827
- array_shape = wcs.array_shape
828
- if array_shape is None:
921
+ if wcs.array_shape is None:
829
922
  msg = "WCS does not have an array shape"
830
923
  raise ValueError(msg)
831
924
 
832
- for _, (ctype, naxis, crval) in enumerate(
833
- zip(wcs.axis_type_names, array_shape[::-1], wcs.wcs.crval, strict=False)
834
- ):
835
- if ctype == "STOKES":
836
- assert naxis <= 1, (
837
- f"Only one polarisation axis is supported - found {naxis}"
838
- )
839
- # FITS Stokes codes are 1=I, 2=Q, 3=U, 4=V; the BEAMS table POL
840
- # column is 0-indexed, so subtract 1.
841
- return int(crval - 1)
842
- return 0
925
+ stokes_wcs = wcs.sub([WCSSUB_STOKES])
926
+ return stokes_wcs if stokes_wcs.naxis else None
843
927
 
844
928
 
845
- def make_beam_table(beams: Beams, old_header: fits.Header) -> fits.BinTableHDU:
929
+ def get_polarisation(header: fits.Header) -> NDArray[np.int_]:
930
+ """Get the 0-based plane indices along the Stokes axis.
931
+
932
+ These are the values the CASA beam table expects in its POL column: indices
933
+ into the cube's Stokes axis, not FITS Stokes codes (see `get_stokes_codes`).
934
+
935
+ Args:
936
+ header (fits.Header): Primary header of the cube
937
+
938
+ Returns:
939
+ NDArray[np.int_]: ``arange(n_stokes)``, or ``[0]`` if there is no Stokes axis
940
+ """
941
+ stokes_wcs = _stokes_wcs(header)
942
+ # CASA counts an image without a Stokes axis as one Stokes plane in its beam set
943
+ n_stokes = 1 if stokes_wcs is None else stokes_wcs.pixel_shape[0]
944
+ return np.arange(n_stokes)
945
+
946
+
947
+ def get_axis_types(header: fits.Header) -> tuple[str, ...]:
948
+ """Get the CTYPE of every axis of an image, in FITS axis order.
949
+
950
+ Args:
951
+ header (fits.Header): Primary header of an image
952
+
953
+ Returns:
954
+ tuple[str, ...]: One CTYPE per axis, "" where an axis has none
955
+ """
956
+ return tuple(
957
+ str(header.get(f"CTYPE{axis}", "")).strip()
958
+ for axis in range(1, header["NAXIS"] + 1)
959
+ )
960
+
961
+
962
+ def get_stokes_codes(header: fits.Header) -> tuple[int, ...] | None:
963
+ """Get the FITS Stokes code of each plane along the Stokes axis.
964
+
965
+ The codes are those of `astropy.coordinates.StokesCoord`, e.g. 1..4 for
966
+ I, Q, U, V (see ``astropy.coordinates.polarization.FITS_STOKES_VALUE_SYMBOL_MAP``).
967
+
968
+ Args:
969
+ header (fits.Header): Primary header of an image
970
+
971
+ Returns:
972
+ tuple[int, ...] | None: One code per Stokes plane, or None if there is no Stokes axis
973
+ """
974
+ stokes_wcs = _stokes_wcs(header)
975
+ if stokes_wcs is None:
976
+ return None
977
+ stokes: StokesCoord = stokes_wcs.pixel_to_world(
978
+ np.arange(stokes_wcs.pixel_shape[0])
979
+ )
980
+ return tuple(int(code) for code in np.round(stokes.value))
981
+
982
+
983
+ def make_beam_table(beams: Beams, cube_header: fits.Header) -> fits.BinTableHDU:
846
984
  """Make a beam table.
847
985
 
986
+ CHAN and POL are 0-based indices into the cube's frequency and Stokes axes
987
+ (CASA convention), not FITS Stokes codes.
988
+
848
989
  Args:
849
- beams (Beams): Beams object
850
- header (fits.Header): Old header to infer polarisation
990
+ beams (Beams): One beam per output channel
991
+ cube_header (fits.Header): Header of the output cube, used to find the
992
+ Stokes axis
851
993
 
852
994
  Returns:
853
995
  fits.BinTableHDU: Beam table
854
996
  """
997
+ stokes_idx = get_polarisation(cube_header)
998
+ if len(stokes_idx) != 1:
999
+ msg = f"Only single-Stokes cubes are supported - found {len(stokes_idx)} Stokes planes"
1000
+ raise NotImplementedError(msg)
855
1001
  nchan = len(beams.major)
856
1002
  chans = np.arange(nchan)
857
- pol = get_polarisation(old_header)
858
- pols = np.ones(nchan, dtype=int) * pol
1003
+ pols = np.full(nchan, stokes_idx[0])
859
1004
  tiny = np.finfo(np.float32).tiny
860
1005
  # A zero-sized beam is not a valid PSF, and a literal zero is exactly what
861
1006
  # the sentinel below exists to keep out of the table. NaN them first so they
@@ -880,7 +1025,7 @@ def make_beam_table(beams: Beams, old_header: fits.Header) -> fits.BinTableHDU:
880
1025
  tab_header = tab_hdu.header
881
1026
  tab_header["EXTNAME"] = "BEAMS"
882
1027
  tab_header["NCHAN"] = nchan
883
- tab_header["NPOL"] = 1 # Only one pol for now
1028
+ tab_header["NPOL"] = len(stokes_idx)
884
1029
 
885
1030
  return tab_hdu
886
1031
 
@@ -1007,6 +1152,7 @@ async def combine_fits_coro(
1007
1152
  # TODO: Check that all files have the same WCS
1008
1153
 
1009
1154
  await check_matching_shapes_coro(file_list=file_list, max_workers=max_workers)
1155
+ await check_matching_axes_coro(file_list=file_list, max_workers=max_workers)
1010
1156
 
1011
1157
  file_specs, specs, missing_chan_idx = await parse_specs_coro(
1012
1158
  spec_file=spec_file,
@@ -1104,6 +1250,12 @@ async def combine_fits_coro(
1104
1250
  )
1105
1251
  raise ValueError(msg)
1106
1252
 
1253
+ # Build the beam table before writing any data, so an unsupported cube
1254
+ # (e.g. multi-Stokes) fails before the planes are copied in
1255
+ beam_table_hdu = (
1256
+ make_beam_table(beams, new_header) if has_beams and not single_beam else None
1257
+ )
1258
+
1107
1259
  new_channels = np.arange(len(specs))
1108
1260
  old_channels = np.arange(len(file_specs))
1109
1261
 
@@ -1145,9 +1297,7 @@ async def combine_fits_coro(
1145
1297
  await gather_with_limit(max_workers, *coros, desc="Writing channels")
1146
1298
 
1147
1299
  # Handle beams
1148
- if has_beams and not single_beam:
1149
- old_header = fits.getheader(file_list[0])
1150
- beam_table_hdu = make_beam_table(beams, old_header)
1300
+ if beam_table_hdu is not None:
1151
1301
  msg = f"Appending beam table to {out_cube}"
1152
1302
  logger.info(msg)
1153
1303
  fits.append(
@@ -17,6 +17,14 @@ class ShapeMismatchException(FITSCubeException):
17
17
  """Input images do not share a common pixel grid"""
18
18
 
19
19
 
20
+ class AxisMismatchException(FITSCubeException):
21
+ """Input images do not share the same axes, in the same order"""
22
+
23
+
24
+ class StokesMismatchException(FITSCubeException):
25
+ """Input images do not share the same Stokes parameter"""
26
+
27
+
20
28
  class AxisOrderException(FITSCubeException):
21
29
  """The spectral axis is not the slowest-varying axis of the cube"""
22
30
 
@@ -31,7 +31,7 @@ classifiers = [
31
31
  ]
32
32
  dynamic = ["version"]
33
33
  dependencies = [
34
- "astropy",
34
+ "astropy>=6.0", # StokesCoord
35
35
  "numpy>=2",
36
36
  "tqdm",
37
37
  "radio_beam",
@@ -9,17 +9,40 @@ import numpy as np
9
9
  import pytest
10
10
  from astropy.io import fits
11
11
  from astropy.time import Time
12
- from fitscube.combine_fits import check_for_any_beam, combine_fits, get_polarisation
12
+ from fitscube.combine_fits import (
13
+ check_for_any_beam,
14
+ combine_fits,
15
+ get_polarisation,
16
+ make_beam_table,
17
+ )
18
+ from radio_beam import Beams
13
19
 
14
20
 
15
21
  @pytest.mark.filterwarnings("ignore:'datfix' made the change")
16
- def test_get_polarisation_uses_crval(headers: dict[str, str]) -> None:
17
- """POL must track CRVAL4 (the actual Stokes code)"""
22
+ def test_get_polarisation_is_axis_index(headers: dict[str, str]) -> None:
23
+ """POL values are 0-based indices along the Stokes axis, not Stokes codes"""
18
24
  header = fits.Header.fromstring(headers["beams"])
19
- assert get_polarisation(header) == 0 # CRVAL4 == 1.0 -> Stokes I
25
+ header["NAXIS4"] = 1
26
+ header["CRVAL4"] = 3.0 # Stokes U, but a single plane
27
+ assert get_polarisation(header).tolist() == [0]
28
+
29
+ header["NAXIS4"] = 3
30
+ assert get_polarisation(header).tolist() == [0, 1, 2]
31
+
32
+
33
+ @pytest.mark.filterwarnings("ignore:'datfix' made the change")
34
+ def test_make_beam_table_rejects_multi_stokes(headers: dict[str, str]) -> None:
35
+ """Multi-Stokes beam tables are not supported yet"""
36
+ header = fits.Header.fromstring(headers["beams"])
37
+ header["NAXIS4"] = 3
38
+ beams = Beams(
39
+ major=np.ones(4) * u.arcsec,
40
+ minor=np.ones(4) * u.arcsec,
41
+ pa=np.zeros(4) * u.deg,
42
+ )
20
43
 
21
- header["CRVAL4"] = 2.0 # Stokes Q
22
- assert get_polarisation(header) == 1
44
+ with pytest.raises(NotImplementedError):
45
+ make_beam_table(beams, header)
23
46
 
24
47
 
25
48
  def test_check_for_any_beams_no_beams(file_list) -> None:
@@ -58,16 +81,15 @@ def test_combine_beam_not_in_first_file(
58
81
 
59
82
 
60
83
  @pytest.mark.parametrize(
61
- ("stokes_code", "expected_pol"),
62
- [(1, 0), (2, 1), (3, 2), (4, 3)], # I, Q, U, V
84
+ "stokes_code",
85
+ [1, 2, 3, 4], # I, Q, U, V
63
86
  )
64
- def test_combine_beam_polarisation_matches_stokes(
87
+ def test_combine_beam_pol_is_zero_index_for_single_stokes(
65
88
  tmp_path: Path,
66
89
  even_specs: u.Quantity,
67
90
  stokes_code: int,
68
- expected_pol: int,
69
91
  ) -> None:
70
- """Each Stokes plane's own beam table must carry its own POL, not always 0."""
92
+ """POL/CHAN are 0-based axis indices, not FITS Stokes codes."""
71
93
  image = np.ones((1, 1, 10, 10))
72
94
  file_list = []
73
95
  for i, spec in enumerate(even_specs):
@@ -103,4 +125,7 @@ def test_combine_beam_polarisation_matches_stokes(
103
125
  with fits.open(out_cube) as hdul:
104
126
  assert hdul[0].header["CRVAL4"] == stokes_code
105
127
  assert hdul[1].name == "BEAMS"
106
- assert set(hdul[1].data["POL"].tolist()) == {expected_pol}
128
+ assert hdul[0].header["NAXIS3"] == len(even_specs)
129
+ assert hdul[0].header["NAXIS4"] == 1
130
+ assert np.all(hdul[1].data["POL"] == 0)
131
+ assert np.array_equal(hdul[1].data["CHAN"], np.arange(len(even_specs)))
@@ -9,11 +9,16 @@ import numpy as np
9
9
  import pytest
10
10
  from astropy.io import fits
11
11
  from fitscube.bounding_box import get_common_bounding_box
12
- from fitscube.combine_fits import check_matching_shapes, combine_fits
12
+ from fitscube.combine_fits import (
13
+ check_matching_axes,
14
+ check_matching_shapes,
15
+ combine_fits,
16
+ )
13
17
  from fitscube.exceptions import (
14
18
  AxisOrderException,
15
19
  IrregularSpacingException,
16
20
  ShapeMismatchException,
21
+ StokesMismatchException,
17
22
  )
18
23
 
19
24
 
@@ -96,6 +101,31 @@ def test_check_matching_shapes(file_list: list[Path]) -> None:
96
101
  assert check_matching_shapes(file_list=file_list) == (8, 8)
97
102
 
98
103
 
104
+ def test_check_matching_axes(file_list: list[Path]) -> None:
105
+ for path in file_list:
106
+ fits.setval(path, "CRVAL3", value=3.0) # Stokes U
107
+ check_matching_axes(file_list=file_list)
108
+
109
+
110
+ def test_check_matching_axes_compares_stokes_codes_not_keywords(
111
+ file_list: list[Path],
112
+ ) -> None:
113
+ """CRVAL3=2 at CRPIX3=2 is still Stokes I at the (only) first pixel"""
114
+ fits.setval(file_list[-1], "CRVAL3", value=2.0)
115
+ fits.setval(file_list[-1], "CRPIX3", value=2.0)
116
+ check_matching_axes(file_list=file_list)
117
+
118
+
119
+ def test_mismatched_stokes_raise(tmp_path: Path, file_list: list[Path]) -> None:
120
+ """A Q plane among I planes would be mislabelled as I in the cube"""
121
+ fits.setval(file_list[2], "CRVAL3", value=2.0) # Stokes Q
122
+ out_cube = tmp_path / "cube.fits"
123
+
124
+ with pytest.raises(StokesMismatchException, match=file_list[2].name):
125
+ combine_fits(file_list=file_list, out_cube=out_cube, overwrite=True)
126
+ assert not out_cube.exists()
127
+
128
+
99
129
  def test_spectral_axis_must_be_slowest(tmp_path: Path, specs: u.Quantity) -> None:
100
130
  """A non-degenerate axis above FREQ breaks the per-plane seek."""
101
131
  file_list = [
@@ -0,0 +1,328 @@
1
+ """End-to-end combines of image pairs across input axis layouts
2
+
3
+ Each case combines two single-plane images and states the cube it expects in
4
+ the notation ``(RA, DEC, FREQ) + (RA, DEC, FREQ) -> (RA, DEC, FREQ)``, or the
5
+ error it expects. Axes are named by CTYPE. A Stokes symbol (e.g. ``I`` or
6
+ ``XX``) is a STOKES axis holding that parameter, and ``FOO`` is an axis type
7
+ fitscube knows nothing about.
8
+ """
9
+
10
+ from __future__ import annotations
11
+
12
+ from collections.abc import Callable
13
+ from pathlib import Path
14
+
15
+ import numpy as np
16
+ import pytest
17
+ from astropy.coordinates import StokesCoord
18
+ from astropy.coordinates.polarization import FITS_STOKES_VALUE_SYMBOL_MAP
19
+ from astropy.io import fits
20
+ from astropy.time import Time
21
+ from fitscube.combine_fits import combine_fits
22
+ from fitscube.exceptions import AxisMismatchException, StokesMismatchException
23
+
24
+ Axes = tuple[str, ...]
25
+ WritePlane = Callable[[Axes, int], Path]
26
+
27
+
28
+ def _is_stokes(axis: str) -> bool:
29
+ return axis in {stokes.symbol for stokes in FITS_STOKES_VALUE_SYMBOL_MAP.values()}
30
+
31
+
32
+ def _axes_id(value: object) -> str:
33
+ if isinstance(value, bool):
34
+ return "time" if value else "freq"
35
+ if isinstance(value, tuple):
36
+ return f"({','.join(value)})"
37
+ if isinstance(value, type):
38
+ return value.__name__
39
+ return str(value)
40
+
41
+
42
+ @pytest.fixture
43
+ def write_plane(tmp_path: Path) -> WritePlane:
44
+ """Write the ``index``-th single-plane image with the given axes.
45
+
46
+ Plane ``index`` holds the value ``index`` and its own beam, is one step
47
+ later in frequency and time than the plane before, and each non-celestial
48
+ axis has length one.
49
+ """
50
+
51
+ def _write_plane(axes: Axes, index: int) -> Path:
52
+ time = Time(60000.0 + index * 10 / 86400, format="mjd")
53
+ cards = {
54
+ "RA": {"CTYPE": "RA---SIN", "CRVAL": 0.0, "CDELT": -1e-3, "CRPIX": 5.0},
55
+ "DEC": {"CTYPE": "DEC--SIN", "CRVAL": 0.0, "CDELT": 1e-3, "CRPIX": 5.0},
56
+ "FREQ": {
57
+ "CTYPE": "FREQ",
58
+ "CRVAL": 1e9 + index * 1e6,
59
+ "CDELT": 1e6,
60
+ "CRPIX": 1.0,
61
+ "CUNIT": "Hz",
62
+ },
63
+ "TIME": {
64
+ "CTYPE": "TIME",
65
+ "CRVAL": time.mjd * 86400,
66
+ "CDELT": 10.0,
67
+ "CRPIX": 1.0,
68
+ "CUNIT": "s",
69
+ },
70
+ "FOO": {"CTYPE": "FOO", "CRVAL": 1.0, "CDELT": 1.0, "CRPIX": 1.0},
71
+ }
72
+ header = fits.Header()
73
+ for fits_axis, axis in enumerate(axes, start=1):
74
+ axis_cards = (
75
+ {
76
+ "CTYPE": "STOKES",
77
+ "CRVAL": float(StokesCoord(axis).value),
78
+ "CDELT": 1.0,
79
+ "CRPIX": 1.0,
80
+ }
81
+ if _is_stokes(axis)
82
+ else cards[axis]
83
+ )
84
+ for key, value in axis_cards.items():
85
+ header[f"{key}{fits_axis}"] = value
86
+ header["DATE-OBS"] = time.isot
87
+ header["MJD-OBS"] = time.mjd
88
+ if len(axes) == 2:
89
+ header["REFFREQ"] = 1e9 + index * 1e6
90
+ # Vary the beam per plane so a beam table is written
91
+ header["BMAJ"] = 1e-3 * (1 + index)
92
+ header["BMIN"] = 1e-3
93
+ header["BPA"] = 0.0
94
+
95
+ shape = (1,) * (len(axes) - 2) + (10, 10)
96
+ path = tmp_path / f"plane_{index}.fits"
97
+ fits.PrimaryHDU(np.full(shape, float(index)), header=header).writeto(path)
98
+ return path
99
+
100
+ return _write_plane
101
+
102
+
103
+ @pytest.fixture
104
+ def out_cube(tmp_path: Path) -> Path:
105
+ return tmp_path / "cube.fits"
106
+
107
+
108
+ @pytest.mark.parametrize(
109
+ ("time_domain_mode", "first", "second", "expected"),
110
+ [
111
+ # Frequency cubes
112
+ (False, ("RA", "DEC"), ("RA", "DEC"), ("RA", "DEC", "FREQ")),
113
+ (False, ("RA", "DEC", "FREQ"), ("RA", "DEC", "FREQ"), ("RA", "DEC", "FREQ")),
114
+ (False, ("DEC", "RA", "FREQ"), ("DEC", "RA", "FREQ"), ("DEC", "RA", "FREQ")),
115
+ (
116
+ False,
117
+ ("RA", "DEC", "FOO", "FREQ"),
118
+ ("RA", "DEC", "FOO", "FREQ"),
119
+ ("RA", "DEC", "FOO", "FREQ"),
120
+ ),
121
+ (
122
+ False,
123
+ ("RA", "DEC", "FREQ", "FOO"),
124
+ ("RA", "DEC", "FREQ", "FOO"),
125
+ ("RA", "DEC", "FREQ", "FOO"),
126
+ ),
127
+ # wsclean axis order
128
+ (
129
+ False,
130
+ ("RA", "DEC", "FREQ", "I"),
131
+ ("RA", "DEC", "FREQ", "I"),
132
+ ("RA", "DEC", "FREQ", "STOKES"),
133
+ ),
134
+ # CASA axis order
135
+ (
136
+ False,
137
+ ("RA", "DEC", "U", "FREQ"),
138
+ ("RA", "DEC", "U", "FREQ"),
139
+ ("RA", "DEC", "STOKES", "FREQ"),
140
+ ),
141
+ (
142
+ False,
143
+ ("RA", "DEC", "FREQ", "XX"),
144
+ ("RA", "DEC", "FREQ", "XX"),
145
+ ("RA", "DEC", "FREQ", "STOKES"),
146
+ ),
147
+ (
148
+ False,
149
+ ("RA", "DEC", "TIME", "FREQ"),
150
+ ("RA", "DEC", "TIME", "FREQ"),
151
+ ("RA", "DEC", "TIME", "FREQ"),
152
+ ),
153
+ (
154
+ False,
155
+ ("RA", "DEC", "FREQ", "TIME", "V"),
156
+ ("RA", "DEC", "FREQ", "TIME", "V"),
157
+ ("RA", "DEC", "FREQ", "TIME", "STOKES"),
158
+ ),
159
+ # Time cubes. A missing TIME axis is appended.
160
+ (True, ("RA", "DEC"), ("RA", "DEC"), ("RA", "DEC", "TIME")),
161
+ (True, ("RA", "DEC", "TIME"), ("RA", "DEC", "TIME"), ("RA", "DEC", "TIME")),
162
+ (
163
+ True,
164
+ ("RA", "DEC", "FREQ"),
165
+ ("RA", "DEC", "FREQ"),
166
+ ("RA", "DEC", "FREQ", "TIME"),
167
+ ),
168
+ (
169
+ True,
170
+ ("RA", "DEC", "FREQ", "I"),
171
+ ("RA", "DEC", "FREQ", "I"),
172
+ ("RA", "DEC", "FREQ", "STOKES", "TIME"),
173
+ ),
174
+ (
175
+ True,
176
+ ("RA", "DEC", "Q", "TIME"),
177
+ ("RA", "DEC", "Q", "TIME"),
178
+ ("RA", "DEC", "STOKES", "TIME"),
179
+ ),
180
+ (
181
+ True,
182
+ ("RA", "DEC", "TIME", "FOO"),
183
+ ("RA", "DEC", "TIME", "FOO"),
184
+ ("RA", "DEC", "TIME", "FOO"),
185
+ ),
186
+ ],
187
+ ids=_axes_id,
188
+ )
189
+ def test_combine_pair(
190
+ write_plane: WritePlane,
191
+ out_cube: Path,
192
+ time_domain_mode: bool,
193
+ first: Axes,
194
+ second: Axes,
195
+ expected: Axes,
196
+ ) -> None:
197
+ """first + second -> expected, with the planes in order and a valid BEAMS table"""
198
+ combine_fits(
199
+ file_list=[write_plane(first, 0), write_plane(second, 1)],
200
+ out_cube=out_cube,
201
+ time_domain_mode=time_domain_mode,
202
+ overwrite=True,
203
+ )
204
+
205
+ with fits.open(out_cube) as hdul:
206
+ header = hdul[0].header
207
+ n_axes = header["NAXIS"]
208
+ ctypes = tuple(
209
+ header[f"CTYPE{axis}"].split("-")[0] for axis in range(1, n_axes + 1)
210
+ )
211
+ assert ctypes == expected
212
+
213
+ combine_axis = ctypes.index("TIME" if time_domain_mode else "FREQ") + 1
214
+ naxes = tuple(header[f"NAXIS{axis}"] for axis in range(1, n_axes + 1))
215
+ assert naxes == tuple(
216
+ 10 if ctype in ("RA", "DEC") else 2 if axis == combine_axis else 1
217
+ for axis, ctype in enumerate(ctypes, start=1)
218
+ )
219
+ stokes = [axis for axis in first if _is_stokes(axis)]
220
+ if stokes:
221
+ stokes_axis = ctypes.index("STOKES") + 1
222
+ assert header[f"CRVAL{stokes_axis}"] == StokesCoord(stokes[0]).value
223
+
224
+ # Plane i along the combine axis holds input i
225
+ planes = np.moveaxis(hdul[0].data, n_axes - combine_axis, 0)
226
+ assert np.all(planes[0] == 0)
227
+ assert np.all(planes[1] == 1)
228
+
229
+ beam_table = hdul["BEAMS"]
230
+ assert beam_table.header["NCHAN"] == 2
231
+ assert beam_table.header["NPOL"] == 1
232
+ assert beam_table.data["CHAN"].tolist() == [0, 1]
233
+ assert beam_table.data["POL"].tolist() == [0, 0]
234
+
235
+
236
+ @pytest.mark.parametrize(
237
+ ("time_domain_mode", "first", "second", "error"),
238
+ [
239
+ (
240
+ False,
241
+ ("RA", "DEC", "FREQ"),
242
+ ("DEC", "RA", "FREQ"),
243
+ AxisMismatchException,
244
+ ),
245
+ (False, ("RA", "DEC"), ("RA", "DEC", "FREQ"), AxisMismatchException),
246
+ (
247
+ False,
248
+ ("RA", "DEC", "FREQ"),
249
+ ("RA", "DEC", "FOO", "FREQ"),
250
+ AxisMismatchException,
251
+ ),
252
+ (
253
+ False,
254
+ ("RA", "DEC", "FREQ", "I"),
255
+ ("RA", "DEC", "I", "FREQ"),
256
+ AxisMismatchException,
257
+ ),
258
+ (
259
+ False,
260
+ ("RA", "DEC", "FREQ", "I"),
261
+ ("RA", "DEC", "FREQ"),
262
+ AxisMismatchException,
263
+ ),
264
+ (
265
+ False,
266
+ ("RA", "DEC", "FREQ", "I"),
267
+ ("RA", "DEC", "FREQ", "Q"),
268
+ StokesMismatchException,
269
+ ),
270
+ (
271
+ False,
272
+ ("RA", "DEC", "FREQ", "XX"),
273
+ ("RA", "DEC", "FREQ", "YY"),
274
+ StokesMismatchException,
275
+ ),
276
+ # Frequency mode needs a FREQ axis, or REFFREQ on a 2D image
277
+ (False, ("RA", "DEC", "FOO"), ("RA", "DEC", "FOO"), ValueError),
278
+ (True, ("RA", "DEC", "TIME"), ("RA", "DEC", "FREQ"), AxisMismatchException),
279
+ (
280
+ True,
281
+ ("RA", "DEC", "TIME", "I"),
282
+ ("RA", "DEC", "TIME", "Q"),
283
+ StokesMismatchException,
284
+ ),
285
+ ],
286
+ ids=_axes_id,
287
+ )
288
+ def test_combine_pair_raises(
289
+ write_plane: WritePlane,
290
+ out_cube: Path,
291
+ time_domain_mode: bool,
292
+ first: Axes,
293
+ second: Axes,
294
+ error: type[Exception],
295
+ ) -> None:
296
+ """first + second -> error, before any output is written"""
297
+ with pytest.raises(error):
298
+ combine_fits(
299
+ file_list=[write_plane(first, 0), write_plane(second, 1)],
300
+ out_cube=out_cube,
301
+ time_domain_mode=time_domain_mode,
302
+ overwrite=True,
303
+ )
304
+ assert not out_cube.exists()
305
+
306
+
307
+ def test_combine_rejects_multi_stokes_beams(tmp_path: Path, out_cube: Path) -> None:
308
+ """A multi-Stokes cube with varying beams fails before any plane is written"""
309
+ file_list = []
310
+ for index in range(2):
311
+ header = fits.Header()
312
+ header["CTYPE3"], header["CRVAL3"] = "STOKES", 1.0
313
+ header["CDELT3"], header["CRPIX3"] = 1.0, 1.0
314
+ header["CTYPE4"], header["CRVAL4"] = "FREQ", 1e9 + index * 1e6
315
+ header["CDELT4"], header["CRPIX4"], header["CUNIT4"] = 1e6, 1.0, "Hz"
316
+ header["BMAJ"] = 1e-3 * (1 + index)
317
+ header["BMIN"], header["BPA"] = 1e-3, 0.0
318
+ path = tmp_path / f"plane_{index}.fits"
319
+ fits.PrimaryHDU(np.full((1, 2, 10, 10), 1.0 + index), header=header).writeto(
320
+ path
321
+ )
322
+ file_list.append(path)
323
+
324
+ with pytest.raises(NotImplementedError, match="single-Stokes"):
325
+ combine_fits(file_list=file_list, out_cube=out_cube, overwrite=True)
326
+
327
+ # The data were never written: every plane is still the blank fill
328
+ assert np.all(fits.getdata(out_cube) == 0)
@@ -44,7 +44,7 @@ def file_list(even_specs: u.Quantity):
44
44
 
45
45
 
46
46
  def test_parse_specs(file_list: list[Path], even_specs: u.Quantity):
47
- file_specs, specs, missing_chan_idx = parse_specs(file_list)
47
+ file_specs, _specs, _missing_chan_idx = parse_specs(file_list)
48
48
  assert np.array_equal(file_specs, even_specs)
49
49
 
50
50
 
@@ -10,7 +10,9 @@ from fitscube.combine_fits import combine_fits, parse_specs
10
10
 
11
11
 
12
12
  def test_parse_specs(file_list: list[Path], even_specs: u.Quantity):
13
- file_specs, specs, missing_chan_idx = parse_specs(file_list, time_domain_mode=True)
13
+ file_specs, _specs, _missing_chan_idx = parse_specs(
14
+ file_list, time_domain_mode=True
15
+ )
14
16
  # assert np.array_equal(file_specs, even_specs)
15
17
  assert np.allclose(file_specs, even_specs, rtol=1e-10, atol=1e-9)
16
18
 
@@ -446,7 +446,7 @@ uvloop = [
446
446
 
447
447
  [package.metadata]
448
448
  requires-dist = [
449
- { name = "astropy" },
449
+ { name = "astropy", specifier = ">=6.0" },
450
450
  { name = "furo", marker = "extra == 'docs'", specifier = ">=2023.8.17" },
451
451
  { name = "myst-parser", marker = "extra == 'docs'", specifier = ">=0.13" },
452
452
  { name = "numpy", specifier = ">=2" },
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes