dkist-processing-dlnirsp 0.32.9__py3-none-any.whl → 0.33.0rc1__py3-none-any.whl

dkist_processing_dlnirsp/tasks/make_movie_frames.py

@@ -1,156 +0,0 @@
-"""Task for turning output science frames into movie frames."""
-
-import numpy as np
-from astropy.io import fits
-from astropy.visualization import ZScaleInterval
-from dkist_processing_common.codecs.fits import fits_array_encoder
-from dkist_processing_common.codecs.fits import fits_hdu_decoder
-from dkist_service_configuration.logging import logger
-
-from dkist_processing_dlnirsp.models.tags import DlnirspTag
-from dkist_processing_dlnirsp.tasks.dlnirsp_base import DlnirspTaskBase
-
-__all__ = ["MakeDlnirspMovieFrames"]
-
-
-class MakeDlnirspMovieFrames(DlnirspTaskBase):
-    """Task class for making individual movie frames that will be used for the browse movie."""
-
-    def run(self) -> None:
-        """
-        Construct movie frames from the set of output science frames.
-
-        For now just make a single movie frame for every output frame. If the data are polarimetric then all 4
-        Stokes images will be put together in a grid.
-        """
-        if self.constants.correct_for_polarization:
-            logger.info("Making polarimetric movie frames")
-        else:
-            logger.info("Making spectrographic movie frames")
-
-        self.make_movie_frames()
-
-    def make_movie_frames(self):
-        """Make a movie frame for every output frame."""
-        for mosaic_num in range(self.constants.num_mosaic_repeats):
-            for dither_step in range(self.constants.num_dither_steps):
-                for X_tile in range(1, self.constants.num_mosaic_tiles_x + 1):
-                    for Y_tile in range(1, self.constants.num_mosaic_tiles_y + 1):
-                        if self.constants.correct_for_polarization:
-                            data, header = self.make_single_polarimetric_frame(
-                                mosaic_num=mosaic_num,
-                                dither_step=dither_step,
-                                X_tile_num=X_tile,
-                                Y_tile_num=Y_tile,
-                            )
-                        else:
-                            data, header = self.make_single_spectrographic_frame(
-                                mosaic_num=mosaic_num,
-                                dither_step=dither_step,
-                                X_tile_num=X_tile,
-                                Y_tile_num=Y_tile,
-                            )
-
-                        self.write(
-                            data=data,
-                            header=header,
-                            tags=[
-                                DlnirspTag.movie_frame(),
-                                DlnirspTag.mosaic_num(mosaic_num),
-                                DlnirspTag.dither_step(dither_step),
-                                DlnirspTag.mosaic_tile_x(X_tile),
-                                DlnirspTag.mosaic_tile_y(Y_tile),
-                            ],
-                            encoder=fits_array_encoder,
-                        )
-
-    def make_single_polarimetric_frame(
-        self, mosaic_num: int, dither_step: int, X_tile_num: int, Y_tile_num: int
-    ) -> tuple[np.ndarray, fits.Header]:
-        """
-        Extract and grid together all 4 Stokes frames for a given instrument loop tuple.
-
-        The data are also scaled for visual appeal.
-        """
-        stokes_dict = dict()
-        for stokes in self.constants.stokes_params:
-            tags = [
-                DlnirspTag.frame(),
-                DlnirspTag.calibrated(),
-                DlnirspTag.mosaic_num(mosaic_num),
-                DlnirspTag.dither_step(dither_step),
-                DlnirspTag.mosaic_tile_x(X_tile_num),
-                DlnirspTag.mosaic_tile_y(Y_tile_num),
-            ]
-            hdu = next(self.read(tags=tags, decoder=fits_hdu_decoder))
-            data_2D = self.grab_wavelength_slice(hdu.data)
-            stokes_dict[stokes] = self.scale_for_rendering(data_2D)
-
-        # Don't care which hdu we get
-        header = hdu.header
-        movie_array = self.grid_movie_frame(
-            top_left=stokes_dict["I"],
-            top_right=stokes_dict["Q"],
-            bottom_left=stokes_dict["U"],
-            bottom_right=stokes_dict["V"],
-        )
-        return movie_array, header
-
-    def make_single_spectrographic_frame(
-        self, mosaic_num: int, dither_step: int, X_tile_num: int, Y_tile_num: int
-    ) -> tuple[np.ndarray, fits.Header]:
-        """Load a output spectrographic frame, extract a single wavelength, and scale for visual appeal."""
-        tags = [
-            DlnirspTag.frame(),
-            DlnirspTag.calibrated(),
-            DlnirspTag.mosaic_num(mosaic_num),
-            DlnirspTag.dither_step(dither_step),
-            DlnirspTag.mosaic_tile_x(X_tile_num),
-            DlnirspTag.mosaic_tile_y(Y_tile_num),
-        ]
-        hdu = next(self.read(tags=tags, decoder=fits_hdu_decoder))
-
-        data_2D = self.grab_wavelength_slice(hdu.data)
-        scaled_data = self.scale_for_rendering(data_2D)
-
-        movie_array = self.scale_for_rendering(scaled_data)
-
-        return movie_array, hdu.header
-
-    @staticmethod
-    def grab_wavelength_slice(data: np.ndarray):
-        """Convert a 3D IFU image into a 2D movie frame by extracting a certain wavelength region."""
-        wavelength_dim_size = data.shape[0]
-        middle_index = int(np.mean(range(wavelength_dim_size)))
-        return data[middle_index, :, :]
-
-    @staticmethod
-    def scale_for_rendering(data: np.ndarray):
-        """
-        Scale the output frame data using a normalization function to facilitate display as a movie frame.
-
-        Non-number pixels (nan, inf, -inf) are set to black.
-        """
-        bad_idx = ~np.isfinite(data)
-        data[bad_idx] = np.nanmedian(data)
-        zscale = ZScaleInterval()
-        scaled_data = zscale(data)
-        scaled_data[bad_idx] = 0.0
-        return scaled_data
-
-    @staticmethod
-    def grid_movie_frame(
-        top_left: np.ndarray,
-        top_right: np.ndarray,
-        bottom_left: np.ndarray,
-        bottom_right: np.ndarray,
-    ) -> np.ndarray:
-        """Combine multiple arrays into a 2x2 grid."""
-        result = np.concatenate(
-            (
-                np.concatenate((top_left, top_right), axis=1),
-                np.concatenate((bottom_left, bottom_right), axis=1),
-            ),
-            axis=0,
-        )
-        return result

dkist_processing_dlnirsp/tests/test_assemble_movie.py

@@ -1,169 +0,0 @@
-import numpy as np
-import pytest
-from dkist_processing_common._util.scratch import WorkflowFileSystem
-from moviepy.video.io.VideoFileClip import VideoFileClip
-
-from dkist_processing_dlnirsp.models.tags import DlnirspTag
-from dkist_processing_dlnirsp.tasks.assemble_movie import AssembleDlnirspMovie
-from dkist_processing_dlnirsp.tests.conftest import DlnirspTestingConstants
-from dkist_processing_dlnirsp.tests.conftest import MovieFrameHeaders
-from dkist_processing_dlnirsp.tests.conftest import tag_on_mosaic_dither_loops
-from dkist_processing_dlnirsp.tests.conftest import write_frames_to_task
-
-
-def write_movie_frames_to_task(
-    task,
-    dither_mode_on: bool,
-    num_mosaics: int,
-    num_X_tiles: int,
-    num_Y_tiles: int,
-    array_shape: tuple[int, int],
-) -> int:
-    dataset = MovieFrameHeaders(
-        dither_mode_on=dither_mode_on,
-        num_mosaics=num_mosaics,
-        num_X_tiles=num_X_tiles,
-        num_Y_tiles=num_Y_tiles,
-        array_shape=array_shape,
-    )
-
-    num_written_frames = write_frames_to_task(
-        task=task,
-        frame_generator=dataset,
-        extra_tags=[
-            DlnirspTag.movie_frame(),
-        ],
-        tag_func=tag_on_mosaic_dither_loops,
-        data_func=make_movie_frame_data,
-    )
-    return num_written_frames
-
-
-def make_movie_frame_data(frame: MovieFrameHeaders) -> np.ndarray:
-    shape = frame.array_shape[1:]
-    data = np.zeros(shape)
-    total_frames = frame.dataset_shape[0]
-    current_index = frame.index + 1
-    horizontal_frac = int((current_index / total_frames) * shape[0])
-    data[:horizontal_frac, :] = 1.0
-
-    return data
-
-
-@pytest.fixture
-def assemble_movie_task(recipe_run_id, dither_mode_on, tmp_path, link_constants_db):
-    num_mosaics = 4
-    num_X_tiles = 3
-    num_Y_tiles = 2
-
-    constants = DlnirspTestingConstants(
-        NUM_MOSAIC_REPEATS=num_mosaics,
-        NUM_MOSAIC_TILES_X=num_X_tiles,
-        NUM_MOSAIC_TILES_Y=num_Y_tiles,
-        NUM_DITHER_STEPS=int(dither_mode_on) + 1,
-    )
-    link_constants_db(recipe_run_id, constants)
-
-    with AssembleDlnirspMovie(
-        recipe_run_id=recipe_run_id,
-        workflow_name="workflow_name",
-        workflow_version="workflow_version",
-    ) as task:
-        task.scratch = WorkflowFileSystem(scratch_base_path=tmp_path, recipe_run_id=recipe_run_id)
-
-        yield task, num_mosaics, num_X_tiles, num_Y_tiles
-        task._purge()
-
-
-@pytest.mark.parametrize("dither_mode_on", [pytest.param(True, id="dither_mode_on")])
-@pytest.mark.parametrize(
-    "is_polarimetric, requires_shrinking",
-    [
-        pytest.param(True, True, id="polarimetric_shrink"),
-        pytest.param(False, False, id="spectrographic_noshrink"),
-        pytest.param(True, False, id="polarimetric_noshrink"),
-    ],
-)
-def test_assemble_movie(
-    assemble_movie_task,
-    dither_mode_on,
-    is_polarimetric,
-    requires_shrinking,
-    mocker,
-    fake_gql_client,
-):
-    """
-    Given: An AssembleDlnirspMovie task and associated MOVIE_FRAME files
-    When: Running the task
-    Then: A movie is generated with the correct size
-    """
-    mocker.patch(
-        "dkist_processing_common.tasks.mixin.metadata_store.GraphQLClient", new=fake_gql_client
-    )
-    task, num_mosaics, num_X_tiles, num_Y_tiles = assemble_movie_task
-
-    if requires_shrinking:
-        movie_shape = (1080 * 2, 1920 * 2)  # Intentionally "backward" from normal
-        expected_shape = (1080, 1920)[::-1]
-    else:
-        movie_shape = (1920 // 2 - 1, 1080 // 2 - 1)  # Weird aspect ratio
-        expected_shape = (1920, 1080)[::-1]  # Because AssembleMovie will force even dimensions
-    write_movie_frames_to_task(
-        task,
-        dither_mode_on=dither_mode_on,
-        num_mosaics=num_mosaics,
-        num_X_tiles=num_X_tiles,
-        num_Y_tiles=num_Y_tiles,
-        array_shape=movie_shape,
-    )
-
-    task()
-
-    movie_file = list(task.read(tags=[DlnirspTag.movie()]))
-    assert len(movie_file) == 1
-    assert movie_file[0].exists()
-    clip = VideoFileClip(str(movie_file[0]))
-    assert tuple(clip.size) == expected_shape
-
-    # import os
-    # os.system(f"cp {movie_file[0]} foo.mp4")
-
-
-@pytest.mark.parametrize(
-    "dither_mode_on",
-    [pytest.param(True, id="dither_mode_on"), pytest.param(False, id="dither_mode_off")],
-)
-def test_tags_for_images(assemble_movie_task, dither_mode_on):
-    """
-    Given: A `AssembleDlnirspMovie` task with constants describing the number of L1 frames
-    When: Computing the tags for a certain movie image
-    Then: The correct tags are returned
-    """
-    task, num_mosaics, num_X_tiles, num_Y_tiles = assemble_movie_task
-    num_dither = int(dither_mode_on) + 1
-
-    assert task.num_images == num_mosaics * num_dither * num_X_tiles * num_Y_tiles
-
-    expected_mosaic_values = sum(
-        [[i for _ in range(num_X_tiles * num_Y_tiles * num_dither)] for i in range(num_mosaics)], []
-    )
-    expected_dither_values = (
-        sum([[d] * num_X_tiles * num_Y_tiles for d in range(num_dither)], []) * num_mosaics
-    )
-    expected_X_tile_values = (
-        sum([[i for _ in range(num_Y_tiles)] for i in range(1, num_X_tiles + 1)] * num_mosaics, [])
-        * num_dither
-    )
-    expected_Y_tile_values = (
-        list(range(1, num_Y_tiles + 1)) * num_X_tiles * num_mosaics * num_dither
-    )
-
-    for n in range(task.num_images):
-        assert sorted(task.tags_for_image_n(n)) == sorted(
-            [
-                DlnirspTag.mosaic_num(expected_mosaic_values[n]),
-                DlnirspTag.dither_step(expected_dither_values[n]),
-                DlnirspTag.mosaic_tile_x(expected_X_tile_values[n]),
-                DlnirspTag.mosaic_tile_y(expected_Y_tile_values[n]),
-            ]
-        )

dkist_processing_dlnirsp/tests/test_make_movie_frames.py

@@ -1,98 +0,0 @@
-import pytest
-from astropy.io import fits
-from dkist_processing_common._util.scratch import WorkflowFileSystem
-
-from dkist_processing_dlnirsp.models.tags import DlnirspTag
-from dkist_processing_dlnirsp.tasks.make_movie_frames import MakeDlnirspMovieFrames
-from dkist_processing_dlnirsp.tests.conftest import DlnirspTestingConstants
-from dkist_processing_dlnirsp.tests.conftest import write_calibrated_frames_to_task
-
-
-@pytest.fixture
-def make_movie_task(recipe_run_id, is_polarimetric, link_constants_db, tmp_path):
-    dither_mode_on = True
-    num_mosaics = 3
-    num_X_tiles = 2
-    num_Y_tiles = 1
-
-    constants_db = DlnirspTestingConstants(
-        NUM_MOSAIC_REPEATS=num_mosaics,
-        NUM_MOSAIC_TILES_X=num_X_tiles,
-        NUM_MOSAIC_TILES_Y=num_Y_tiles,
-        NUM_DITHER_STEPS=int(dither_mode_on) + 1,
-        POLARIMETER_MODE="Full Stokes" if is_polarimetric else "Stokes I",
-    )
-    link_constants_db(recipe_run_id, constants_db)
-
-    with MakeDlnirspMovieFrames(
-        recipe_run_id=recipe_run_id,
-        workflow_name="workflow_name",
-        workflow_version="workflow_version",
-    ) as task:
-        task.scratch = WorkflowFileSystem(scratch_base_path=tmp_path, recipe_run_id=recipe_run_id)
-
-        yield task, num_mosaics, dither_mode_on, num_X_tiles, num_Y_tiles
-        task._purge()
-
-
-@pytest.mark.parametrize(
-    "is_polarimetric",
-    [pytest.param(True, id="polarimetric"), pytest.param(False, id="spectrographic")],
-)
-def test_make_movie_frames(
-    make_movie_task, is_polarimetric, link_constants_db, mocker, fake_gql_client
-):
-    """
-    Given: A Make Movie Frames task and some calibrated Science frames
-    When: Running the task
-    Then: The correct number of movie frames are produced and they have the correct shape (depending on polarimetric or not)
-    """
-    mocker.patch(
-        "dkist_processing_common.tasks.mixin.metadata_store.GraphQLClient", new=fake_gql_client
-    )
-    task, num_mosaics, dither_mode_on, num_X_tiles, num_Y_tiles = make_movie_task
-
-    num_dither = int(dither_mode_on) + 1
-    wave_size = 3
-    X_size = 4
-    Y_size = 5
-    write_calibrated_frames_to_task(
-        task,
-        num_mosaics=num_mosaics,
-        num_X_tiles=num_X_tiles,
-        num_Y_tiles=num_Y_tiles,
-        is_polarimetric=is_polarimetric,
-        array_shape=(wave_size, X_size, Y_size),
-        dither_mode_on=dither_mode_on,
-    )
-
-    task()
-
-    if is_polarimetric:
-        expected_movie_frame_shape = (X_size * 2, Y_size * 2)
-    else:
-        expected_movie_frame_shape = (X_size, Y_size)
-
-    all_movie_frames = list(task.read([DlnirspTag.movie_frame()]))
-    assert len(all_movie_frames) == num_dither * num_mosaics * num_X_tiles * num_Y_tiles
-
-    for mosaic in range(num_mosaics):
-        for dither in range(num_dither):
-            for X_tile in range(1, num_X_tiles + 1):
-                for Y_tile in range(1, num_Y_tiles + 1):
-                    single_frame = list(
-                        task.read(
-                            [
-                                DlnirspTag.movie_frame(),
-                                DlnirspTag.mosaic_num(mosaic),
-                                DlnirspTag.dither_step(dither),
-                                DlnirspTag.mosaic_tile_x(X_tile),
-                                DlnirspTag.mosaic_tile_y(Y_tile),
-                            ]
-                        )
-                    )
-                    assert len(single_frame) == 1
-                    hdul = fits.open(single_frame[0])
-                    assert len(hdul) == 1
-                    data = hdul[0].data
-                    assert data.shape == expected_movie_frame_shape