Load multiple images separately or as a polarization hypercube

VERSION.txt

@@ -1 +1 @@
-1.1.10
+1.1.11

carta/constants.py

@@ -153,8 +153,33 @@ class ContourDashMode(StrEnum):
     NEGATIVE_ONLY = "NegativeOnly"
 
 
+PROTO_POLARIZATION = {
+    "I": 1,
+    "Q": 2,
+    "U": 3,
+    "V": 4,
+    "RR": 5,
+    "LL": 6,
+    "RL": 7,
+    "LR": 8,
+    "XX": 9,
+    "YY": 10,
+    "XY": 11,
+    "YX": 12,
+    "PTOTAL": 13,
+    "PLINEAR": 14,
+    "PFTOTAL": 15,
+    "PFLINEAR": 16,
+    "PANGLE": 17,
+}
+
+
 class Polarization(IntEnum):
     """Polarizations, corresponding to the POLARIZATIONS enum in the frontend."""
+
+    def __init__(self, value):
+        self.proto_index = PROTO_POLARIZATION[self.name]
+
     YX = -8
     XY = -7
     YY = -6

carta/session.py

@@ -10,11 +10,11 @@
 import posixpath
 
 from .image import Image
-from .constants import CoordinateSystem, LabelType, BeamType, PaletteColor, Overlay, PanelMode, GridMode, ComplexComponent, NumberFormat
+from .constants import CoordinateSystem, LabelType, BeamType, PaletteColor, Overlay, PanelMode, GridMode, ComplexComponent, NumberFormat, Polarization
 from .backend import Backend
 from .protocol import Protocol
 from .util import logger, Macro, split_action_path, CartaBadID, CartaBadSession, CartaBadUrl
-from .validation import validate, String, Number, Color, Constant, Boolean, NoneOr, OneOf
+from .validation import validate, String, Number, Color, Constant, Boolean, NoneOr, OneOf, IterableOf, MapOf, Union
 
 
 class Session:
@@ -422,6 +422,83 @@ def open_LEL_image(self, expression, directory=".", append=False, make_active=Tr
         """
         return Image.new(self, directory, expression, "", append, True, make_active=make_active, update_directory=update_directory)
 
+    @validate(IterableOf(String()), Boolean())
+    def open_images(self, image_paths, append=False):
+        """Open multiple images
+
+        This is a utility function for adding multiple images in a single command. It assumes that the images are not complex-valued or LEL expressions, and that the default HDU can be used for each image. For more complicated use cases, the methods for opening individual images should be used.
+
+        Parameters
+        ----------
+        image_paths : {0}
+            The image paths, either relative to the session's current directory or absolute paths relative to the CARTA backend's root directory.
+        append : {1}
+            Whether the images should be appended to existing images. By default this is ``False`` and any existing open images are closed.
+
+        Returns
+        -------
+        list of :obj:`carta.image.Image` objects
+            The list of opened images.
+        """
+        images = []
+        for path in image_paths[:1]:
+            images.append(self.open_image(path, append=append))
+        for path in image_paths[1:]:
+            images.append(self.open_image(path, append=True))
+        return images
+
+    @validate(Union(IterableOf(String(), min_size=2), MapOf(Constant(Polarization), String(), min_size=2)), Boolean())
+    def open_hypercube(self, image_paths, append=False):
+        """Open multiple images merged into a polarization hypercube.
+
+        Parameters
+        ----------
+        image_paths : {0}
+            The image paths, either relative to the session's current directory or absolute paths relative to the CARTA backend's root directory. If this is a list of paths, the polarizations will be deduced from the image headers or names. If this is a dictionary, the polarizations must be used as keys.
+        append : {1}
+            Whether the hypercube should be appended to existing images. By default this is ``False`` and any existing open images are closed.
+
+        Returns
+        -------
+        :obj:`carta.image.Image`
+            The opened hypercube.
+
+        Raises
+        ------
+        ValueError
+            If explicit polarizations are not provided, and cannot be deduced from the image headers or names.
+        """
+        stokes_images = []
+
+        if isinstance(image_paths, dict):
+            for stokes, path in image_paths.items():
+                directory, file_name = posixpath.split(path)
+                directory = self.resolve_file_path(directory)
+                stokes_images.append({"directory": directory, "file": file_name, "hdu": "", "polarizationType": stokes.proto_index})
+        else:
+            stokes_guesses = set()
+
+            for path in image_paths:
+                directory, file_name = posixpath.split(path)
+                directory = self.resolve_file_path(directory)
+
+                stokes_guess = self.call_action("fileBrowserStore.getStokesFile", directory, file_name, "")
+
+                if not stokes_guess:
+                    raise ValueError(f"Could not deduce polarization for {path}. Please use a dictionary to specify the polarization mapping explicitly.")
+
+                stokes_guesses.add(stokes_guess["polarizationType"])
+                stokes_images.append(stokes_guess)
+
+            if len(stokes_guesses) < len(stokes_images):
+                raise ValueError("Duplicate polarizations deduced for provided images. Please use a dictionary to specify the polarization mapping explicitly.")
+
+        output_directory = self.pwd()
+        output_hdu = ""
+        command = "appendConcatFile" if append else "openConcatFile"
+        image_id = self.call_action(command, stokes_images, output_directory, output_hdu)
+        return Image(self, image_id)
+
     def image_list(self):
         """Return the list of currently open images.
 

tests/test_session.py

@@ -4,7 +4,7 @@
 from carta.session import Session
 from carta.image import Image
 from carta.util import CartaValidationFailed, Macro
-from carta.constants import CoordinateSystem, NumberFormat as NF, ComplexComponent as CC
+from carta.constants import CoordinateSystem, NumberFormat as NF, ComplexComponent as CC, Polarization as Pol
 
 # FIXTURES
 
@@ -134,8 +134,6 @@ def test_open_image(mocker, session, args, kwargs, expected_args, expected_kwarg
     session.open_image(*args, **kwargs)
     mock_image_new.assert_called_with(session, *expected_args, **expected_kwargs)
 
-# TODO this should be merged with the test above when this separate function is removed
-
 
 @pytest.mark.parametrize("args,kwargs,expected_args,expected_kwargs", [
     # Open complex image with default component
@@ -180,6 +178,117 @@ def test_open_LEL_image(mocker, session, args, kwargs, expected_args, expected_k
     mock_image_new.assert_called_with(session, *expected_args, **expected_kwargs)
 
 
+@pytest.mark.parametrize("append", [True, False])
+def test_open_images(mocker, session, mock_method, append):
+    mock_open_image = mock_method("open_image", ["1", "2", "3"])
+    images = session.open_images(["foo.fits", "bar.fits", "baz.fits"], append)
+    mock_open_image.assert_has_calls([
+        mocker.call("foo.fits", append=append),
+        mocker.call("bar.fits", append=True),
+        mocker.call("baz.fits", append=True),
+    ])
+    assert images == ["1", "2", "3"]
+
+
+@pytest.mark.parametrize("paths,expected_args", [
+    (["foo.fits", "bar.fits", "baz.fits"], [
+        [
+            {"directory": "/resolved/path", "file": "foo.fits", "hdu": "", "polarizationType": 1},
+            {"directory": "/resolved/path", "file": "bar.fits", "hdu": "", "polarizationType": 2},
+            {"directory": "/resolved/path", "file": "baz.fits", "hdu": "", "polarizationType": 3},
+        ], "/current/dir", ""]),
+])
+@pytest.mark.parametrize("append,expected_command", [
+    (True, "appendConcatFile"),
+    (False, "openConcatFile"),
+])
+def test_open_hypercube_guess_polarization(mocker, session, mock_call_action, mock_method, paths, expected_args, append, expected_command):
+    mock_method("pwd", ["/current/dir"])
+    mock_method("resolve_file_path", ["/resolved/path"] * 3)
+    mock_call_action.side_effect = [*expected_args[0], 123]
+
+    hypercube = session.open_hypercube(paths, append)
+
+    mock_call_action.assert_has_calls([
+        mocker.call("fileBrowserStore.getStokesFile", "/resolved/path", "foo.fits", ""),
+        mocker.call("fileBrowserStore.getStokesFile", "/resolved/path", "bar.fits", ""),
+        mocker.call("fileBrowserStore.getStokesFile", "/resolved/path", "baz.fits", ""),
+        mocker.call(expected_command, *expected_args),
+    ])
+
+    assert type(hypercube) is Image
+    assert hypercube.session == session
+    assert hypercube.image_id == 123
+
+
+@pytest.mark.parametrize("paths,expected_calls,mocked_side_effect,expected_error", [
+    (["foo.fits", "bar.fits"], [
+        ("fileBrowserStore.getStokesFile", "/resolved/path", "foo.fits", ""),
+    ], [
+        None,
+    ], "Could not deduce polarization for"),
+    (["foo.fits", "bar.fits"], [
+        ("fileBrowserStore.getStokesFile", "/resolved/path", "foo.fits", ""),
+        ("fileBrowserStore.getStokesFile", "/resolved/path", "bar.fits", ""),
+    ], [
+        {"directory": "/resolved/path", "file": "foo.fits", "hdu": "", "polarizationType": 1},
+        {"directory": "/resolved/path", "file": "bar.fits", "hdu": "", "polarizationType": 1},
+    ], "Duplicate polarizations deduced"),
+])
+def test_open_hypercube_guess_polarization_bad(mocker, session, mock_call_action, mock_method, paths, expected_calls, mocked_side_effect, expected_error):
+    mock_method("pwd", ["/current/dir"])
+    mock_method("resolve_file_path", ["/resolved/path"] * 3)
+    mock_call_action.side_effect = mocked_side_effect
+
+    with pytest.raises(ValueError) as e:
+        session.open_hypercube(paths)
+    assert expected_error in str(e.value)
+
+    mock_call_action.assert_has_calls([mocker.call(*args) for args in expected_calls])
+
+
+@pytest.mark.parametrize("paths,expected_args", [
+    ({Pol.I: "foo.fits", Pol.Q: "bar.fits", Pol.U: "baz.fits"}, [
+        [
+            {"directory": "/resolved/path", "file": "foo.fits", "hdu": "", "polarizationType": 1},
+            {"directory": "/resolved/path", "file": "bar.fits", "hdu": "", "polarizationType": 2},
+            {"directory": "/resolved/path", "file": "baz.fits", "hdu": "", "polarizationType": 3},
+        ], "/current/dir", ""]),
+])
+@pytest.mark.parametrize("append,expected_command", [
+    (True, "appendConcatFile"),
+    (False, "openConcatFile"),
+])
+def test_open_hypercube_explicit_polarization(mocker, session, mock_call_action, mock_method, paths, expected_args, append, expected_command):
+    mock_method("pwd", ["/current/dir"])
+    mock_method("resolve_file_path", ["/resolved/path"] * 3)
+    mock_call_action.side_effect = [123]
+
+    hypercube = session.open_hypercube(paths, append)
+
+    mock_call_action.assert_has_calls([
+        mocker.call(expected_command, *expected_args),
+    ])
+
+    assert type(hypercube) is Image
+    assert hypercube.session == session
+    assert hypercube.image_id == 123
+
+
+@pytest.mark.parametrize("paths,expected_error", [
+    ({Pol.I: "foo.fits"}, "at least 2"),
+    (["foo.fits"], "at least 2"),
+])
+@pytest.mark.parametrize("append", [True, False])
+def test_open_hypercube_bad(mocker, session, mock_call_action, mock_method, paths, expected_error, append):
+    mock_method("pwd", ["/current/dir"])
+    mock_method("resolve_file_path", ["/resolved/path"] * 3)
+
+    with pytest.raises(Exception) as e:
+        session.open_hypercube(paths, append)
+    assert expected_error in str(e.value)
+
+
 # OVERLAY