fitscube 2.0.0__tar.gz → 2.2.0__tar.gz

{fitscube-2.0.0 → fitscube-2.2.0}/.github/workflows/cd.yml

@@ -49,7 +49,7 @@ jobs:
           path: dist
 
       - name: Generate artifact attestation for sdist and wheel
-        uses: actions/attest-build-provenance@v2.2.2
+        uses: actions/attest-build-provenance@v2.4.0
         with:
           subject-path: "dist/*"
 

{fitscube-2.0.0 → fitscube-2.2.0}/.github/workflows/ci.yml

@@ -39,7 +39,7 @@ jobs:
       fail-fast: false
       matrix:
         python-version: ["3.10", "3.11", "3.12", "3.13"]
-        runs-on: [ubuntu-latest, windows-latest, macos-14]
+        runs-on: [ubuntu-latest]
 
     steps:
       - uses: actions/checkout@v4
@@ -62,6 +62,6 @@ jobs:
           --durations=20
 
       - name: Upload coverage report
-        uses: codecov/codecov-action@v5.4.0
+        uses: codecov/codecov-action@v5.4.3
         with:
           token: ${{ secrets.CODECOV_TOKEN }}

{fitscube-2.0.0 → fitscube-2.2.0}/.pre-commit-config.yaml

@@ -33,14 +33,14 @@ repos:
       - id: rst-inline-touching-normal
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: "v0.12.4"
+    rev: "v0.12.7"
     hooks:
       - id: ruff
         args: ["--fix", "--show-fixes"]
       - id: ruff-format
 
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: "v1.17.0"
+    rev: "v1.17.1"
     hooks:
       - id: mypy
         files: src|tests

{fitscube-2.0.0 → fitscube-2.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: fitscube
-Version: 2.0.0
+Version: 2.2.0
 Summary: A package to produce produce FITS cubes.
 Project-URL: Homepage, https://github.com/AlecThomson/fitscube
 Project-URL: Bug Tracker, https://github.com/AlecThomson/fitscube/issues

{fitscube-2.0.0 → fitscube-2.2.0}/fitscube/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '2.0.0'
-__version_tuple__ = version_tuple = (2, 0, 0)
+__version__ = version = '2.2.0'
+__version_tuple__ = version_tuple = (2, 2, 0)

{fitscube-2.0.0 → fitscube-2.2.0}/fitscube/combine_fits.py

@@ -160,7 +160,13 @@ async def create_cube_from_scratch_coro(
     if output_file.exists() and overwrite:
         output_file.unlink()
 
-    output_wcs = WCS(output_header)
+    try:
+        output_wcs = WCS(output_header)
+    except Exception as e:
+        logger.error("Error creating new header")
+        for k in output_header:
+            logger.error(f"{k} = {output_header[k]}")
+        raise e
     output_shape = output_wcs.array_shape
     msg = f"Creating a new FITS file with shape {output_shape}"
     logger.info(msg)
@@ -191,6 +197,7 @@ async def create_cube_from_scratch_coro(
 
     for key, value in output_header.items():
         header[key] = value
+        logger.debug(f"{key}={value}")
 
     header.tofile(output_file, overwrite=overwrite)
 
@@ -253,11 +260,39 @@ async def create_output_cube_coro(
     ctype = "TIME" if time_domain_mode else "FREQ"
 
     old_data, old_header = fits.getdata(old_name, header=True, memmap=True)
-    even_spec = np.diff(specs).std() < (1e-4 * unit)
+    sorted_specs = np.sort(specs)
+    if time_domain_mode:
+        logger.info("Computing time-differences")
+
+        # This attempts to constrain the deviation away from 'asbolute' time
+        # as far as it can be. If some time steps are not regularly space
+        # (differencce between the time of adjacent scans) can be positive or negative,
+        # but so long as the assumulated total is close to 0 then we can say
+        # it is close. This approach is catering to some strangeness in
+        # ASKAP data. If the accumulated error is small enough we can assume
+        # the FITS header can encoude the times as regular steps.
+        diff_time = np.diff(sorted_specs)
+        diff_diff_time = np.diff(diff_time)
+        running_deviation_from_zero = np.abs(np.cumsum(diff_diff_time))
+        even_spec = np.max(running_deviation_from_zero) < (np.mean(diff_time) * 0.02)
+
+        # This is a simpler way where no attempt is made to ensure the total
+        # error on the irregular steps accumulates and violates the regular
+        # spacing we can encode in the fits header. I am less trustworthy of this,
+        # if all deviations are negative they would accumulate. Individually
+        # they may not fail but across the whole TIME dimension, as encoded by the
+        # C-type header fields that define a regularly spaced interval, may.
+        # even_spec = np.all(np.abs(np.diff(np.diff(sorted_specs))) < (0.15*u.s))
+    else:
+        even_spec = np.diff(sorted_specs).std() < (1e-4 * unit)
+
+    logger.debug(f"{np.diff(sorted_specs).std()=}")
     if not even_spec:
         spequency = "Times" if time_domain_mode else "Frequencies"
         msg = f"{spequency} are not evenly spaced"
         logger.warning(msg)
+        logger.debug(f"{np.max(np.diff(sorted_specs))=}")
+        logger.debug(f"{np.min(np.diff(sorted_specs))=}")
 
     n_chan = len(specs)
 
@@ -271,15 +306,15 @@ async def create_output_cube_coro(
         try:
             idx = wcs.axis_type_names[::-1].index(ctype)
 
-        except ValueError as e:
-            msg = f"No {ctype} axis found in WCS."
+            fits_idx = wcs.axis_type_names.index(ctype) + 1
+            logger.info(f"{ctype} axis found at index {idx} (NAXIS{fits_idx})")
 
-            raise ValueError(msg) from e
-        fits_idx = wcs.axis_type_names.index(ctype) + 1
-        logger.info(f"{ctype} axis found at index %s (NAXIS%s)", idx, fits_idx)
+        except ValueError:
+            msg = f"No {ctype} axis not found in WCS."
+            logger.info(msg)
+            fits_idx = len(old_data.shape) + 1
 
     new_header = old_header.copy()
-    new_header["NAXIS"] = 3 if is_2d else len(old_data.shape)
     new_header[f"NAXIS{fits_idx}"] = n_chan
     new_header[f"CRPIX{fits_idx}"] = 1
     new_header[f"CRVAL{fits_idx}"] = specs[0].value
@@ -287,7 +322,25 @@ async def create_output_cube_coro(
     new_header[f"CUNIT{fits_idx}"] = f"{unit:fits}"
     new_header[f"CTYPE{fits_idx}"] = ctype
 
+    # Figure out the correct number of dimensions to use
+    _no_of_naxis = [k for k in new_header if k.startswith("NAXIS") and k != "NAXIS"]
+    new_header["NAXIS"] = len(_no_of_naxis)
+
+    for k in ["CTYPE", "CUNIT", "CDELT"]:
+        key = f"{k}{fits_idx}"
+        logger.debug(f"{key}={new_header[key]}")
+
+    # Add extra transform fields for consistency
+    if ("CD1_1" in new_header or "PC1_1" in new_header) and fits_idx != 1:
+        transform_type = "CD" if "CD1_1" in new_header else "PC"
+        pv1 = f"{transform_type}{fits_idx}_{fits_idx}"
+        logger.info(f"Adding {pv1} to header")
+        new_header[pv1] = 1.0
+
     if ignore_spec or not even_spec:
+        logger.info(
+            f"Ignore the specrency information, {ignore_spec=} or {not even_spec=}"
+        )
         new_header[f"CDELT{fits_idx}"] = 1
         del new_header[f"CUNIT{fits_idx}"]
         new_header[f"CTYPE{fits_idx}"] = "CHAN"
@@ -305,9 +358,6 @@ async def create_output_cube_coro(
         )
         del new_header["BMAJ"], new_header["BMIN"], new_header["BPA"]
 
-    if time_domain_mode:
-        new_header["COMMENT"] = "The frequency/chan axis in this cube represents TIME."
-
     plane_shape = list(old_data.shape)
     cube_shape = plane_shape.copy()
     if is_2d:
@@ -393,6 +443,7 @@ async def parse_specs_coro(
         spec_file (str | None, optional): File containing frequencies/times. Defaults to None.
         spec_list (list[float] | None, optional): List of frequencies/times. Defaults to None.
         ignore_spec (bool | None, optional): Ignore frequency/time information. Defaults to False.
+        time_domain_mode (bool, optional): Whether these cubes dhould be formed over the time axis. Defaults to False.
 
     Raises:
         ValueError: If both spec_file and spec_list are specified

{fitscube-2.0.0 → fitscube-2.2.0}/fitscube/exceptions.py

@@ -5,8 +5,8 @@ class FITSCubeException(Exception):
     """Base container for FITSCube exceptions"""
 
 
-class FREQMissingException(FITSCubeException):
-    """Missing FREQ axis in fits cube"""
+class TargetAxisMissingException(FITSCubeException):
+    """The target axis is missing in fits cube"""
 
 
 class ChannelMissingException(FITSCubeException):

{fitscube-2.0.0 → fitscube-2.2.0}/fitscube/extract.py

@@ -11,7 +11,7 @@ import numpy as np
 from astropy.io import fits
 from radio_beam import Beam, Beams
 
-from fitscube.exceptions import ChannelMissingException, FREQMissingException
+from fitscube.exceptions import ChannelMissingException, TargetAxisMissingException
 from fitscube.logging import logger, set_verbosity
 
 
@@ -21,8 +21,10 @@ class ExtractOptions:
 
     hdu_index: int = 0
     """The HDU in the fits cube to access (e.g. for header and data)"""
-    channel_index: int = 0
-    """The channel of the cube to extract"""
+    channel_index: int | None = None
+    """The channel of the cube to extract. Defaults to None"""
+    time_index: int | None = None
+    """The timestep of the cube to extract. Defaults to None"""
     overwrite: bool = False
     """overwrite the output file, if it exists"""
     output_path: Path | None = None
@@ -30,8 +32,8 @@ class ExtractOptions:
 
 
 @dataclass
-class FreqWCS:
-    """Exxtract of frequency information in the WCS taken straight from
+class TargetWCS:
+    """Extract of target information in the WCS taken straight from
     the fits header."""
 
     axis: int
@@ -45,21 +47,90 @@ class FreqWCS:
     cdelt: float
     """The step between planes"""
     cunit: str
-    """the unit of the frequency information"""
+    """the unit of the target axis"""
 
 
-def get_output_path(input_path: Path, channel_index: int) -> Path:
+@dataclass
+class TargetIndex:
+    """Simple container to capture the mode to extract."""
+
+    axis_name: str
+    """The name of the axis to search for (e.g. TIME or FREQ)"""
+    axis_index: int
+    """The index (e.g. channel or timestep) to extract from the cube"""
+    output_name: str
+    """The output name to put in output files"""
+
+
+def create_target_index(
+    channel_index: int | None = None, time_index: int | None = None
+) -> TargetIndex:
+    """Define the properties of the target axis to subset based on the provided
+    index.
+
+    Args:
+        channel_index (int | None, optional): If not None, the frequency channel to extract. Defaults to None.
+        time_index (int | None, optional): If not None, the timestep to extract. Defaults to None.
+
+    Raises:
+        ValueError: If both channel and item indices are supplied
+        ValueError: If neither channel nor time index are set
+
+    Returns:
+        TargetIndex: Specified properties of target axis to subset
+    """
+    if time_index is not None and channel_index is not None:
+        msg = "Both time and channel index are set. Not allowed."
+        raise ValueError(msg)
+
+    if isinstance(time_index, int):
+        return TargetIndex(
+            axis_name="TIME", axis_index=time_index, output_name="timestep"
+        )
+    if isinstance(channel_index, int):
+        return TargetIndex(
+            axis_name="FREQ", axis_index=channel_index, output_name="channel"
+        )
+    msg = f"Something went wrong, target index could not be formed, {channel_index=} {time_index=}"
+    raise ValueError(msg)
+
+
+def _check_extract_mode(extract_options: ExtractOptions) -> None:
+    """Verify the operation of the extract options axis.
+
+    Args:
+        extract_options (ExtractOptions): The settings providedd to extract fitscube
+
+    Raises:
+        ValueError: Raised if both channel_index and time_index are unset
+        ValueError: Raise if both channel index and time index are set
+    """
+
+    if extract_options.channel_index is None and extract_options.time_index is None:
+        msg = "Both channel index and time index are None. One needs to be set."
+        raise ValueError(msg)
+    if (
+        extract_options.channel_index is not None
+        and extract_options.time_index is not None
+    ):
+        msg = "Both channel index and time index are set. Only one may be set. "
+        raise ValueError(msg)
+
+
+def get_output_path(input_path: Path, target_index: TargetIndex) -> Path:
     """Create the output path to write the plane to
 
     Args:
         input_path (Path): The base input path name
-        channel_index (int): The channel index to extract
+        target_index (TargetIndex): The target index to extract
 
     Returns:
         Path: New output path for the plane-only fits image
     """
+    index_suffix = f"{target_index.output_name.lower()}-{target_index.axis_index}"
+
     # The input_path suffix returns a string with a period
-    channel_suffix = f".channel-{channel_index}{input_path.suffix}"
+    channel_suffix = f".{index_suffix}{input_path.suffix}"
     output_path = input_path.with_suffix(channel_suffix)
 
     logger.debug(f"The formed {output_path=}")
@@ -86,14 +157,14 @@ def fits_file_contains_beam_table(header: fits.header.Header | Path) -> bool:
     return bool(loaded_header["CASAMBM"])
 
 
-def extract_beam_from_beam_table(fits_path: Path, channel_index: int) -> Beam:
+def extract_beam_from_beam_table(fits_path: Path, index: int) -> Beam:
     """Extract the beam that corresponds to the channel requested. The beam
     is drawn from a beam table that is inserted into the FITS cube. It is
     expected that the beam table exists.
 
     Args:
         fits_path (Path): The fits table to inspect for a beam table, and return the channel beam
-        channel_index (int): The channel to extract the beam for
+        index (int): The channel to extract the beam for
 
     Raises:
         ValueError: Raised when a beam table can not be found
@@ -116,26 +187,38 @@ def extract_beam_from_beam_table(fits_path: Path, channel_index: int) -> Beam:
 
         assert beams is not None, "beams is empty, which should not happen"
 
-    return beams[channel_index]
+    return beams[index]
 
 
-def find_freq_axis(header: fits.header.Header) -> FreqWCS:
-    """Attempt to find the axies of the channel in the data
-    cube that corresponds to frequency/channels.
+def find_target_axis(
+    header: fits.header.Header, target_index: TargetIndex | str = "FREQ"
+) -> TargetWCS:
+    """Attempt to find the axies of the target dimension in the data
+    cube that corresponds to the target type (e.g. time or frequency).
 
     Args:
         header (fits.header.Header): The header from the fits cube
+        target_index (TargetIndex | str): The name of axus to search for in the FITS header representing the axis to extract from. Defaults to FREQ.
 
     Returns:
-        FreqWCS: The information in the FITS header describing the frequency of the cube
+        TargetWCS: The information in the FITS header describing the target dimension of the cube
     """
 
+    axis_name = "FREQ"
+    if isinstance(target_index, str):
+        axis_name = target_index
+    elif isinstance(target_index, TargetIndex):
+        axis_name = target_index.axis_name
+
+    logger.info(f"Searching for {axis_name=} in header")
+
     naxis = header["NAXIS"]
-    # Remember that range upper limit is exclusive
+    # Remember that range upper limit is exclusive, and
+    # we start counting from 1
     for axis in range(1, naxis + 1):
-        if "FREQ" in header[f"CTYPE{axis}"]:
-            logger.info(f"Found FREQ at {axis=}")
-            return FreqWCS(
+        if axis_name in header[f"CTYPE{axis}"]:
+            logger.info(f"Found {axis_name} at {axis=}")
+            return TargetWCS(
                 axis=axis,
                 ctype=header[f"CTYPE{axis}"],
                 crpix=header[f"CRPIX{axis}"],
@@ -145,41 +228,56 @@ def find_freq_axis(header: fits.header.Header) -> FreqWCS:
             )
 
     msg = "Did not find the frequency axis"
-    raise FREQMissingException(msg)
+    raise TargetAxisMissingException(msg)
 
 
-def create_plane_freq_wcs(original_freq_wcs: FreqWCS, channel_index: int) -> FreqWCS:
-    """Create the frequency fields appropriate for a extracted channel
+def create_plane_target_wcs(
+    original_freq_wcs: TargetWCS, target_index: int | TargetIndex
+) -> TargetWCS:
+    """Create the target fields appropriate for a extracted channel/time index
 
     Args:
-        original_freq_wcs (FreqWCS): The frequency information describing the spectral axis
-        channel_index (int): The channel to extract
+        original_freq_wcs (TargetWCS): The frequency information describing the spectral axis
+        target_index (int | TargetIndex): The index to extract from the cube
 
     Returns:
-        FreqWCS: The frequency information for a channel
+        TargetWCS: The frequency information for a channel
     """
-    channel_freq = original_freq_wcs.crval + (channel_index * original_freq_wcs.cdelt)
-    return FreqWCS(
+    index = target_index if isinstance(target_index, int) else target_index.axis_index
+    update_index = original_freq_wcs.crval + (index * original_freq_wcs.cdelt)
+    return TargetWCS(
         axis=original_freq_wcs.axis,
         ctype=original_freq_wcs.ctype,
         crpix=1,
-        crval=channel_freq,
+        crval=update_index,
         cdelt=original_freq_wcs.cdelt,
         cunit=original_freq_wcs.cunit,
     )
 
 
-def update_header_for_frequency(
+def update_header_for_target_axis(
     header: fits.header.Header,
-    freq_wcs: FreqWCS,
-    channel_index: int,
+    target_wcs: TargetWCS,
+    target_index: TargetIndex,
     extract_beam_from_file: Path | None = None,
 ) -> fits.header.Header:
+    """Update the base header to indicate the new extracted characteristics
+    of the extracted plane.
+
+    Args:
+        header (fits.header.Header): The base header to examine and update
+        target_wcs (TargetWCS): The characteristics of the new extracted plane
+        channel_index (int): The extracted index
+        extract_beam_from_file (Path | None, optional): If not None, attempt to extract the beam table to update the stored beam inforionat. Defaults to None.
+
+    Returns:
+        fits.header.Header: The updated header
+    """
     # Get the new wcs items for the channels
-    plane_freq_wcs = create_plane_freq_wcs(
-        original_freq_wcs=freq_wcs, channel_index=channel_index
+    plane_freq_wcs = create_plane_target_wcs(
+        original_freq_wcs=target_wcs, target_index=target_index.axis_index
     )
-    _idx = freq_wcs.axis
+    _idx = target_wcs.axis
     out_header = header.copy()
     out_header[f"CTYPE{_idx}"] = plane_freq_wcs.ctype
     out_header[f"CRPIX{_idx}"] = plane_freq_wcs.crpix
@@ -192,7 +290,7 @@ def update_header_for_frequency(
     ):
         try:
             channel_beam: Beam = extract_beam_from_beam_table(
-                fits_path=extract_beam_from_file, channel_index=channel_index
+                fits_path=extract_beam_from_file, index=target_index.axis_index
             )
 
             out_header["BMAJ"] = channel_beam.major.to(u.deg).value
@@ -220,11 +318,20 @@ def extract_plane_from_cube(fits_cube: Path, extract_options: ExtractOptions) ->
     Returns:
         Path: The output file
     """
+    # Initial sanity check of the axis to extract
+    _check_extract_mode(extract_options=extract_options)
+
+    target_index = create_target_index(
+        channel_index=extract_options.channel_index,
+        time_index=extract_options.time_index,
+    )
+
     output_path: Path = (
         extract_options.output_path
         if extract_options.output_path
         else get_output_path(
-            input_path=fits_cube, channel_index=extract_options.channel_index
+            input_path=fits_cube,
+            target_index=target_index,
         )
     )
 
@@ -241,21 +348,21 @@ def extract_plane_from_cube(fits_cube: Path, extract_options: ExtractOptions) ->
     logger.info("Extracted header and data")
 
     logger.info(f"Data shape: {data.shape}")
-    freq_axis_wcs = find_freq_axis(header=header)
-    freq_cube_index = len(data.shape) - freq_axis_wcs.axis
+    target_axis_wcs = find_target_axis(header=header, target_index=target_index)
+    target_cube_index = len(data.shape) - target_axis_wcs.axis
 
-    if extract_options.channel_index > data.shape[freq_axis_wcs.axis] - 1:
-        msg = f"{extract_options.channel_index=} outside of channel cube {data.shape=}"
+    if target_index.axis_index > data.shape[target_cube_index]:
+        msg = f"{extract_options.channel_index=} outside of channel cube {data.shape=}, axis shape of {data.shape[target_axis_wcs.axis - 1]}"
         raise ChannelMissingException(msg)
 
     # Get the channel index requested
-    freq_plane_data = np.take(data, extract_options.channel_index, axis=freq_cube_index)
+    freq_plane_data = np.take(data, target_index.axis_index, axis=target_cube_index)
     # and pad it back so dimensions match
-    freq_plane_data = np.expand_dims(freq_plane_data, axis=freq_cube_index)
-    freq_plane_header = update_header_for_frequency(
+    freq_plane_data = np.expand_dims(freq_plane_data, axis=target_cube_index)
+    freq_plane_header = update_header_for_target_axis(
         header=header,
-        freq_wcs=freq_axis_wcs,
-        channel_index=extract_options.channel_index,
+        target_wcs=target_axis_wcs,
+        target_index=target_index,
         extract_beam_from_file=fits_cube
         if fits_file_contains_beam_table(header)  # replace with "BEAMS" in opened fits?
         else None,
@@ -281,7 +388,10 @@ def get_parser(parser: ArgumentParser | None = None) -> ArgumentParser:
 
     parser.add_argument("fits_cube", type=Path, help="The cube to extract a plane from")
     parser.add_argument(
-        "--channel-index", type=int, default=0, help="The channel to extract"
+        "--channel-index", type=int, default=None, help="The channel to extract"
+    )
+    parser.add_argument(
+        "--time-index", type=int, default=None, help="The channel to extract"
     )
     parser.add_argument(
         "--hdu-index",
@@ -315,6 +425,7 @@ def cli(args: Namespace | None = None) -> None:
     extract_options = ExtractOptions(
         hdu_index=args.hdu_index,
         channel_index=args.channel_index,
+        time_index=args.time_index,
         overwrite=args.overwrite,
         output_path=args.output_path,
     )

{fitscube-2.0.0 → fitscube-2.2.0}/tests/conftest.py

@@ -21,6 +21,16 @@ def headers() -> dict[str, str]:
     return {"base": EXAMPLE_HEADER, "beams": EXAMPLE_HEADER_WITH_BM}
 
 
+@pytest.fixture
+def timecube_path(tmpdir) -> Path:
+    tmp_dir = Path(tmpdir) / "timecube"
+    tmp_dir.mkdir(exist_ok=True, parents=True)
+    cube_zip = Path(__file__).parent / "data" / "timecube.zip"
+
+    unpack_archive(cube_zip, tmp_dir)
+    return tmp_dir / "test_timecube.fits"
+
+
 @pytest.fixture
 def cube_path(tmpdir) -> Path:
     tmp_dir = Path(tmpdir) / "cube"
@@ -42,3 +52,16 @@ def image_paths(tmpdir) -> list[Path]:
     image_paths.sort()
 
     return image_paths
+
+
+@pytest.fixture
+def time_image_paths(tmpdir) -> list[Path]:
+    tmp_dir = Path(tmpdir) / "time_images"
+    tmp_dir.mkdir(exist_ok=True, parents=True)
+    images_zip = Path(__file__).parent / "data" / "time_images.zip"
+
+    unpack_archive(images_zip, tmp_dir)
+    image_paths = list(tmp_dir.glob("*fits"))
+    image_paths.sort()
+
+    return image_paths

{fitscube-2.0.0 → fitscube-2.2.0}/tests/test_extract.py

@@ -7,27 +7,74 @@ from pathlib import Path
 import numpy as np
 import pytest
 from astropy.io import fits
-from fitscube.exceptions import ChannelMissingException
+from fitscube.exceptions import ChannelMissingException, TargetAxisMissingException
 from fitscube.extract import (
     ExtractOptions,
-    create_plane_freq_wcs,
+    TargetIndex,
+    _check_extract_mode,
+    create_plane_target_wcs,
+    create_target_index,
     extract_plane_from_cube,
-    find_freq_axis,
+    find_target_axis,
     fits_file_contains_beam_table,
     get_output_path,
-    update_header_for_frequency,
+    update_header_for_target_axis,
 )
 
 
+def test_target_index() -> None:
+    """Simple test of the target container"""
+    target_index = create_target_index(channel_index=2)
+    assert isinstance(target_index, TargetIndex)
+    assert target_index.axis_name == "FREQ"
+    assert target_index.axis_index == 2
+
+    target_index = create_target_index(time_index=20)
+    assert isinstance(target_index, TargetIndex)
+    assert target_index.axis_name == "TIME"
+    assert target_index.axis_index == 20
+
+    # int of 0 is False
+    target_index = create_target_index(time_index=0)
+    assert isinstance(target_index, TargetIndex)
+    assert target_index.axis_name == "TIME"
+    assert target_index.axis_index == 0
+
+    with pytest.raises(ValueError, match="index"):
+        create_target_index()
+
+    with pytest.raises(ValueError, match="index"):
+        create_target_index(channel_index=2, time_index=4)
+
+
+def test_check_mode_for_consistency() -> None:
+    """See if the basic consistency checks for the extraction mode makes sense"""
+    extract_options = ExtractOptions()
+    with pytest.raises(ValueError, match="index"):
+        _check_extract_mode(extract_options=extract_options)
+
+    extract_options = ExtractOptions(channel_index=3, time_index=4)
+    with pytest.raises(ValueError, match="index"):
+        _check_extract_mode(extract_options=extract_options)
+
+
 def test_get_output_path() -> None:
     """Make sure the output path generated is correct"""
 
+    target_index = create_target_index(channel_index=10)
     in_fits = Path("some.example.cube.fits")
-    channel_index = 10
     expected_fits = Path("some.example.cube.channel-10.fits")
 
     assert expected_fits == get_output_path(
-        input_path=in_fits, channel_index=channel_index
+        input_path=in_fits, target_index=target_index
+    )
+
+    target_index = create_target_index(time_index=10)
+    in_fits = Path("some.example.cube.fits")
+    expected_fits = Path("some.example.cube.timestep-10.fits")
+
+    assert expected_fits == get_output_path(
+        input_path=in_fits, target_index=target_index
     )
 
 
@@ -43,7 +90,7 @@ def test_find_freq_axis(example_header) -> None:
     """Find the components associated with frequency from the header"""
     header = fits.header.Header.fromstring(example_header)
 
-    freq_wcs = find_freq_axis(header=header)
+    freq_wcs = find_target_axis(header=header)
     assert freq_wcs.axis == 4
     assert freq_wcs.crpix == 1
     assert freq_wcs.crval == 801490740.740741
@@ -54,15 +101,15 @@ def test_create_plane_freq_wcs(example_header) -> None:
     """Update the freq wcs to indicate a plane"""
     header = fits.header.Header.fromstring(example_header)
 
-    freq_wcs = find_freq_axis(header=header)
-    plane_wcs = create_plane_freq_wcs(original_freq_wcs=freq_wcs, channel_index=1)
+    freq_wcs = find_target_axis(header=header)
+    plane_wcs = create_plane_target_wcs(original_freq_wcs=freq_wcs, target_index=1)
 
     assert plane_wcs.axis == freq_wcs.axis
     assert plane_wcs.crpix == 1
     assert plane_wcs.crval == 805490740.740741
     assert plane_wcs.cdelt == freq_wcs.cdelt
 
-    plane_wcs = create_plane_freq_wcs(original_freq_wcs=freq_wcs, channel_index=0)
+    plane_wcs = create_plane_target_wcs(original_freq_wcs=freq_wcs, target_index=0)
 
     assert plane_wcs.axis == freq_wcs.axis
     assert plane_wcs.crpix == 1
@@ -76,10 +123,10 @@ def test_update_header_for_frequency(example_header) -> None:
 
     header = fits.header.Header.fromstring(example_header)
 
-    freq_wcs = find_freq_axis(header=header)
-
-    new_header = update_header_for_frequency(
-        header=header, freq_wcs=freq_wcs, channel_index=1
+    freq_wcs = find_target_axis(header=header)
+    target_index = create_target_index(channel_index=1)
+    new_header = update_header_for_target_axis(
+        header=header, target_wcs=freq_wcs, target_index=target_index
     )
     assert new_header["CRPIX4"] == 1
     assert new_header["CRVAL4"] == 805490740.740741
@@ -167,3 +214,37 @@ def test_compare_extracted_to_image_bad_channel(cube_path, tmpdir) -> None:
     )
     with pytest.raises(ChannelMissingException):
         extract_plane_from_cube(fits_cube=cube_path, extract_options=extract_options)
+
+
+def test_extract_time_from_freq(cube_path, tmpdir) -> None:
+    """Attempt to extract a timestep from a cube without TIME. Should
+    raise an error
+    """
+    output_file = Path(tmpdir) / "extract" / "test.fits"
+    output_file.parent.mkdir(parents=True, exist_ok=True)
+    timestep = 0
+
+    extract_options = ExtractOptions(
+        hdu_index=0, time_index=timestep, output_path=output_file
+    )
+
+    with pytest.raises(TargetAxisMissingException):
+        extract_plane_from_cube(fits_cube=cube_path, extract_options=extract_options)
+
+
+def test_extract_time_cube(timecube_path, tmpdir) -> None:
+    """Attempt to extract a timestep from a cube without TIME. Should
+    raise an error
+    """
+    output_file = Path(tmpdir) / "extract" / "test.fits"
+    output_file.parent.mkdir(parents=True, exist_ok=True)
+    timestep = 200
+
+    extract_options = ExtractOptions(
+        hdu_index=0, time_index=timestep, output_path=output_file
+    )
+
+    timecube = extract_plane_from_cube(
+        fits_cube=timecube_path, extract_options=extract_options
+    )
+    assert timecube.exists()

{fitscube-2.0.0 → fitscube-2.2.0}/tests/test_times.py

@@ -126,3 +126,26 @@ def test_uneven_combine(
             assert chan in (1, 2)
             continue
         assert np.allclose(plane, image)
+
+
+@pytest.mark.filterwarnings("ignore:'datfix' made the change")
+def test_wsclean_images_create_axis(time_image_paths, tmpdir) -> None:
+    """Ensure that the combined cube conforms to the input data"""
+
+    tmpdir = Path(tmpdir) / "time_cube_combine"
+    tmpdir.mkdir(parents=True, exist_ok=True)
+    out_cube = tmpdir / "time_cube_mate.fits"
+
+    combine_fits(
+        file_list=time_image_paths,
+        out_cube=out_cube,
+        overwrite=True,
+        time_domain_mode=True,
+    )
+
+    cube_data = fits.getdata(out_cube)
+    for i, time_image_path in enumerate(time_image_paths):
+        image_data = fits.getdata(time_image_path)
+        # The TIME axis will be appended as a new dimension
+        cube_image_data = cube_data[i]
+        assert np.allclose(image_data.squeeze(), cube_image_data.squeeze())