From 4e246bcdc588a8cb22918d9d10467b1457bef792 Mon Sep 17 00:00:00 2001
From: David Zwicker <david.zwicker@ds.mpg.de>
Date: Sat, 16 Mar 2024 18:18:49 +0100
Subject: [PATCH 01/13] A first version of the MovieStorage

---
 pde/storage/__init__.py       |   1 +
 pde/storage/file.py           |   2 +-
 pde/storage/movie.py          | 411 ++++++++++++++++++++++++++++++++++
 pde/visualization/plotting.py |   2 +-
 4 files changed, 414 insertions(+), 2 deletions(-)
 create mode 100644 pde/storage/movie.py

diff --git a/pde/storage/__init__.py b/pde/storage/__init__.py
index d0cf9afb..bfbb423b 100644
--- a/pde/storage/__init__.py
+++ b/pde/storage/__init__.py
@@ -7,6 +7,7 @@
    ~memory.get_memory_storage
    ~memory.MemoryStorage
    ~file.FileStorage
+   ~movie.MovieStorage
 
 .. codeauthor:: David Zwicker <david.zwicker@ds.mpg.de> 
 """
diff --git a/pde/storage/file.py b/pde/storage/file.py
index ff34e58a..82f2358d 100644
--- a/pde/storage/file.py
+++ b/pde/storage/file.py
@@ -25,7 +25,7 @@ class FileStorage(StorageBase):
 
     def __init__(
         self,
-        filename: str,
+        filename: str | Path,
         *,
         info: InfoDict | None = None,
         write_mode: WriteModeType = "truncate_once",
diff --git a/pde/storage/movie.py b/pde/storage/movie.py
new file mode 100644
index 00000000..9c7a7c13
--- /dev/null
+++ b/pde/storage/movie.py
@@ -0,0 +1,411 @@
+"""
+Defines a class storing data on the file system as a compressed movie
+
+This package requires the optional :mod:`ffmpeg-python` package to use FFmpeg for
+reading and writing movies.
+
+.. codeauthor:: David Zwicker <david.zwicker@ds.mpg.de> 
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import shlex
+from collections.abc import Iterator
+from pathlib import Path
+from typing import Any, Callable, Literal
+
+import numpy as np
+from matplotlib.colors import Normalize
+from numpy.typing import ArrayLike
+
+from ..fields import FieldCollection, ScalarField
+from ..fields.base import FieldBase
+from ..tools.docstrings import fill_in_docstring
+from ..tools.misc import module_available
+from ..trackers.interrupts import ConstantInterrupts
+from .base import InfoDict, StorageBase, StorageTracker, WriteModeType
+
+
+def _get_limits(value: float | ArrayLike, dim: int) -> np.ndarray:
+    """helper function creating sequence of length `dim` from input"""
+    if np.isscalar(value):
+        return np.full(dim, value)
+    else:
+        return np.atleast_1d(value)[:dim]
+
+
+class MovieStorage(StorageBase):
+    def __init__(
+        self,
+        filename: str | Path,
+        *,
+        vmin: float | ArrayLike = 0,
+        vmax: float | ArrayLike = 1,
+        info: InfoDict | None = None,
+        write_mode: WriteModeType = "truncate_once",
+    ):
+        """
+        Args:
+            filename (str):
+                The path to the hdf5-file where the data is stored
+            vmin (float or array):
+                Lowest values that are encoded (per field)
+            vmax (float or array):
+                Highest values that are encoded (per field)
+            info (dict):
+                Supplies extra information that is stored in the storage
+            write_mode (str):
+                Determines how new data is added to already existing data. Possible
+                values are: 'append' (data is always appended), 'truncate' (data is
+                cleared every time this storage is used for writing), or 'truncate_once'
+                (data is cleared for the first writing, but appended subsequently).
+                Alternatively, specifying 'readonly' will disable writing completely.
+
+
+        TODO:
+            - allow more bits for colorchannels
+            - allow choosing bitrate for video
+            - support different write_mode
+            - track whether times roughly work (checking for frame drops)
+            - we could embedd extra information (like time, and maybe colorscaling) in
+              the individual frames if we extended the shape
+        """
+        if not module_available("ffmpeg"):
+            raise ModuleNotFoundError("`MovieStorage` needs `ffmpeg-python` package")
+
+        super().__init__(info=info, write_mode=write_mode)
+        self.filename = Path(filename)
+        self.vmin = vmin
+        self.vmax = vmax
+        self.dt = None
+        self.t_start = None
+
+        self._logger = logging.getLogger(self.__class__.__name__)
+        self._ffmpeg: Any = None
+        self._state: Literal["closed", "reading", "writing"] = "closed"
+        self._norms: list[Normalize] | None = None
+        self._is_writing = False
+
+    def __del__(self):
+        self.close()  # ensure open files are closed when the FileStorage is deleted
+
+    def close(self) -> None:
+        """close the currently opened file"""
+        if self._ffmpeg is not None:
+            self._logger.info(f"Close movie file `{self.filename}`")
+            if self._state == "writing":
+                self._ffmpeg.stdin.close()
+                self._ffmpeg.wait()
+            elif self._state == "reading":
+                self._ffmpeg.stdout.close()
+                self._ffmpeg.wait()
+            self._ffmpeg = None
+            self._state = "closed"
+
+    def __enter__(self) -> MovieStorage:
+        return self
+
+    def __exit__(self, exc_type, exc_value, exc_traceback):
+        self.close()
+
+    def clear(self):
+        """truncate the storage by removing all stored data."""
+        if self.filename.exists():
+            self.filename.unlink()
+
+    def _get_metadata(self) -> str:
+        """obtain metadata stored in the video"""
+        info = self.info.copy()
+        info["version"] = 1
+        info["vmin"] = self.vmin
+        info["vmax"] = self.vmax
+        return json.dumps(info)
+
+    def _read_metadata(self) -> None:
+        """read metadata from video and store it in :attr:`info`"""
+        import ffmpeg
+
+        info = ffmpeg.probe(self.filename)
+
+        # sanity checks on the video
+        nb_streams = info["format"]["nb_streams"]
+        if nb_streams != 1:
+            self._logger.warning(f"Only using first of {nb_streams} streams")
+
+        raw_comment = info["format"].get("tags", {}).get("comment", "{}")
+        metadata = json.loads(shlex.split(raw_comment)[0])
+
+        version = metadata.pop("version", 1)
+        if version != 1:
+            self._logger.warning(f"Unknown metadata version `{version}`")
+        self.vmin = metadata.pop("vmin", 0)
+        self.vmax = metadata.pop("vmax", 1)
+        self.info.update(metadata)
+
+        stream = info["streams"][0]
+        self.info["num_frames"] = int(stream["nb_frames"])
+        self.info["width"] = stream["coded_width"]
+        self.info["height"] = stream["coded_height"]
+        if stream["pix_fmt"] == "gray":
+            self.info["channels"] = 1
+        elif stream["pix_fmt"] in {"rgb24", "yuv420p"}:
+            self.info["channels"] = 3
+        else:
+            self._logger.warning(f"Unknown pixel format {stream['pix_fmt']}")
+
+    def _init_normalization(
+        self, field: FieldBase, *, inverse: bool = False
+    ) -> list[Callable]:
+        """initialize the normalizations of the color information
+
+        Args:
+            field (:class:`~pde.fields.base.FieldBase):
+                Example field to obtain information about grid and data rank
+            inverse (bool):
+                Whether inverse normalization function should be returned
+        """
+        self._norms = []
+        fields = field if isinstance(field, FieldCollection) else [field]
+        vmin = _get_limits(self.vmin, len(fields))
+        vmax = _get_limits(self.vmax, len(fields))
+        for f_id, f in enumerate(fields):
+            if inverse:
+                norm = lambda data: vmin[f_id] + (vmax[f_id] - vmin[f_id]) * data
+            else:
+                norm = Normalize(vmin[f_id], vmax[f_id], clip=True)
+            num = self.grid.dim**f.rank  # independent components in the field
+            self._norms.extend([norm] * num)
+
+    def _reshape_data(self, data: np.ndarray) -> np.ndarray:
+        """reshape data such that it has exactly three dimensions"""
+        data = np.copy(data)
+        if data.ndim < 3:
+            data = data.reshape(data.shape + (1,) * (3 - data.ndim))  # set ndim=3
+        elif data.ndim > 3:
+            data = data.reshape(data.shape[:2] + (-1,))  # collapse last axes
+        return data
+
+    def start_writing(self, field: FieldBase, info: InfoDict | None = None) -> None:
+        """initialize the storage for writing data
+
+        Args:
+            field (:class:`~pde.fields.FieldBase`):
+                An example of the data that will be written to extract the grid and the
+                data_shape
+            info (dict):
+                Supplies extra information that is stored in the storage
+        """
+        import ffmpeg
+
+        if self._is_writing:
+            raise RuntimeError(f"{self.__class__.__name__} is already in writing mode")
+        if self._ffmpeg is not None:
+            raise RuntimeError("ffmpeg process already started")
+
+        # delete data if truncation is requested. This is for instance necessary
+        # to remove older data with incompatible data_shape
+        if self.write_mode == "truncate" or self.write_mode == "truncate_once":
+            self.clear()
+
+        # initialize the writing, setting current data shape
+        super().start_writing(field, info=info)
+        if info:
+            self.info.update(info)
+
+        # get spatial dimension of the video
+        grid = field.grid
+        if grid.num_axes == 1:
+            width, height = field.grid.shape[0], 1
+        elif grid.num_axes == 2:
+            width, height = field.grid.shape
+        else:
+            raise RuntimeError("Cannot use grid with more than two axes")
+
+        # get color channel information
+        if isinstance(field, ScalarField):
+            pix_fmt = "gray"  # 1 color channel
+            self._data_shape = (width, height, 1)
+        else:
+            pix_fmt = "rgb24"  # 3 color channels
+            self._data_shape = (width, height, 3)
+        self.info["frame_shape"] = self._data_shape
+        self.info["field_shape"] = field.data.shape
+        if field.is_complex:
+            raise NotImplementedError("MovieStorage does not support complex values")
+
+        # set up the normalization
+        self._init_normalization(field)
+
+        # set input
+        self._logger.debug(f"Start ffmpeg process for `{self.filename}`")
+        f_input = ffmpeg.input(
+            "pipe:", format="rawvideo", pix_fmt=pix_fmt, s=f"{width}x{height}"
+        )
+        # set output format
+        f_output = f_input.output(
+            filename=self.filename,
+            pix_fmt=pix_fmt,
+            metadata="comment=" + shlex.quote(self._get_metadata()),
+        )
+        f_output = f_output.overwrite_output()  # allow overwriting file
+        self._ffmpeg = f_output.run_async(pipe_stdin=True)  # start process
+        self._state = "writing"
+
+    def _append_data(self, data: np.ndarray, time: float) -> None:
+        """append a new data set
+
+        Args:
+            data (:class:`~numpy.ndarray`): The actual data
+            time (float): The time point associated with the data (currently not used)
+        """
+        if self._state != "writing" or self._ffmpeg is None:
+            RuntimeError(
+                "Writing not initialized. Call "
+                f"`{self.__class__.__name__}.start_writing`"
+            )
+
+        # ensure the data has the shape width x height x depth
+        data = self._reshape_data(data)
+        assert len(self._data_shape) == data.ndim == 3
+        if self._data_shape[2] == 1:
+            # single color channel
+            data = data.reshape(self._data_shape)
+        elif self._data_shape[2] == 3:
+            # three color channels
+            if data.shape[2] == 1:
+                z = np.zeros(data.shape[:2] + (1,), dtype=self._dtype)
+                data = np.dstack((data, z, z))
+            elif data.shape[2] == 2:
+                z = np.zeros(data.shape[:2] + (1,), dtype=self._dtype)
+                data = np.dstack((data[..., 0], data[..., 1], z))
+        else:
+            raise RuntimeError
+
+        assert data.shape == self._data_shape
+        assert data.shape[2] == len(self._norms)
+
+        # map values to [0, 1] float values
+        for i in range(data.shape[2]):
+            data[..., i] = self._norms[i](data[..., i])
+        data *= 256
+
+        # write the new data
+        self._ffmpeg.stdin.write(data.astype(np.uint8).tobytes())
+
+    def end_writing(self) -> None:
+        """finalize the storage after writing"""
+        if not self._state == "writing":
+            return  # writing mode was already ended
+        self._logger.debug("End writing")
+        self.close()
+
+    def __len__(self):
+        """return the number of stored items, i.e., time steps"""
+        if "num_frames" not in self.info:
+            self._read_metadata()
+        return self.info["num_frames"]
+
+    @property
+    def times(self):
+        """:class:`~numpy.ndarray`: The times at which data is available"""
+        return np.arange(self.t_start, len(self), self.dt)
+
+    @property
+    def data(self):
+        """:class:`~numpy.ndarray`: The actual data for all time"""
+        raise NotImplementedError
+
+    def __iter__(self) -> Iterator[FieldBase]:
+        """iterate over all stored fields"""
+        import ffmpeg
+
+        if "width" not in self.info:
+            self._read_metadata()
+        if self._field is None:
+            self._init_field()
+        self._init_normalization(self._field, inverse=True)
+        # frame_shape = (self.info["width"], self.info["height"], self.info["channels"])
+        data_shape = (self.info["width"], self.info["height"], len(self._norms))
+        data = np.empty(data_shape, dtype=self._dtype)
+
+        # iterate over entire movie
+        f_input = ffmpeg.input(self.filename)
+        f_output = f_input.output(
+            "pipe:", format="rawvideo", pix_fmt="rgb24", vframes=8
+        )
+        proc = f_output.run_async(pipe_stdout=True)
+        while True:
+            read_bytes = proc.stdout.read(np.prod(data_shape))
+            if not read_bytes:
+                break
+            frame = np.frombuffer(read_bytes, np.uint8).reshape(data_shape)
+
+            for i, norm in enumerate(self._norms):
+                data[..., i] = norm(frame[:, :, i] / 256)
+
+            # create the field with the data of the given index
+            assert self._field is not None
+            field = self._field.copy()
+            field.data = data.reshape(self.info["field_shape"])
+            yield field
+
+    def items(self) -> Iterator[tuple[float, FieldBase]]:
+        """iterate over all times and stored fields, returning pairs"""
+        # iterate over entire movie
+        t = self.info.get("t_start", 0)
+        dt = self.info.get("dt", 1)
+        for field in self:
+            yield t, field
+            t += dt
+
+    @fill_in_docstring
+    def tracker(
+        self,
+        interrupts: ConstantInterrupts | float = 1,
+        *,
+        transformation: Callable[[FieldBase, float], FieldBase] | None = None,
+    ) -> StorageTracker:
+        """create object that can be used as a tracker to fill this storage
+
+        Args:
+            interrupts:
+                {ARG_TRACKER_INTERRUPT}
+            transformation (callable, optional):
+                A function that transforms the current state into a new field or field
+                collection, which is then stored. This allows to store derived
+                quantities of the field during calculations. The argument needs to be a
+                callable function taking 1 or 2 arguments. The first argument always is
+                the current field, while the optional second argument is the associated
+                time.
+
+        Returns:
+            :class:`StorageTracker`: The tracker that fills the current storage
+
+        Example:
+            The `transformation` argument allows storing additional fields:
+
+            .. code-block:: python
+
+                def add_to_state(state):
+                    transformed_field = state.smooth(1)
+                    return field.append(transformed_field)
+
+                storage = pde.MemoryStorage()
+                tracker = storage.tracker(1, transformation=add_to_state)
+                eq.solve(..., tracker=tracker)
+
+            In this example, :obj:`storage` will contain a trajectory of the fields of
+            the simulation as well as the smoothed fields. Other transformations are
+            possible by defining appropriate :func:`add_to_state`
+        """
+        if np.isscalar(interrupts):
+            interrupts = ConstantInterrupts(interrupts)
+        if not isinstance(interrupts, ConstantInterrupts):
+            self._logger.warning("`VideoTracker` can only use `ConstantInterrupts`")
+        self.info["dt"] = interrupts.dt
+        self.info["t_start"] = interrupts.t_start
+        return StorageTracker(
+            storage=self, interrupts=interrupts, transformation=transformation
+        )
diff --git a/pde/visualization/plotting.py b/pde/visualization/plotting.py
index 965317f4..e4e5f9b9 100644
--- a/pde/visualization/plotting.py
+++ b/pde/visualization/plotting.py
@@ -19,7 +19,7 @@
 import math
 import time
 import warnings
-from typing import Any, Callable, Literal, Tuple, Union
+from typing import Any, Callable, Literal, Union
 
 import numpy as np
 

From a177321ec02f61e302b23283585198839c576b29 Mon Sep 17 00:00:00 2001
From: David Zwicker <david.zwicker@ds.mpg.de>
Date: Sat, 16 Mar 2024 20:50:38 +0100
Subject: [PATCH 02/13] Major improvements

---
 docs/source/_static/requirements_optional.csv |   1 +
 pde/storage/__init__.py                       |   1 +
 pde/storage/movie.py                          | 137 ++++++++++++------
 pde/tools/resources/requirements_full.txt     |   1 +
 pyproject.toml                                |   1 +
 scripts/_templates/pyproject.toml             |   1 +
 scripts/create_requirements.py                |   6 +
 tests/requirements_full.txt                   |   1 +
 tests/storage/test_movie_storages.py          |  50 +++++++
 9 files changed, 153 insertions(+), 46 deletions(-)
 create mode 100644 tests/storage/test_movie_storages.py

diff --git a/docs/source/_static/requirements_optional.csv b/docs/source/_static/requirements_optional.csv
index c06293af..e8006041 100644
--- a/docs/source/_static/requirements_optional.csv
+++ b/docs/source/_static/requirements_optional.csv
@@ -1,4 +1,5 @@
 Package,Minimal version,Usage
+ffmpeg-python,0.2,Reading and writing videos
 h5py,2.10,Storing data in the hierarchical file format
 ipywidgets,8,Jupyter notebook support
 mpi4py,3,Parallel processing using MPI
diff --git a/pde/storage/__init__.py b/pde/storage/__init__.py
index bfbb423b..17d08908 100644
--- a/pde/storage/__init__.py
+++ b/pde/storage/__init__.py
@@ -14,3 +14,4 @@
 
 from .file import FileStorage
 from .memory import MemoryStorage, get_memory_storage
+from .movie import MovieStorage
diff --git a/pde/storage/movie.py b/pde/storage/movie.py
index 9c7a7c13..a583b4fd 100644
--- a/pde/storage/movie.py
+++ b/pde/storage/movie.py
@@ -13,12 +13,13 @@
 import logging
 import shlex
 from collections.abc import Iterator
+from dataclasses import dataclass
 from pathlib import Path
 from typing import Any, Callable, Literal
 
 import numpy as np
 from matplotlib.colors import Normalize
-from numpy.typing import ArrayLike
+from numpy.typing import ArrayLike, DTypeLike
 
 from ..fields import FieldCollection, ScalarField
 from ..fields.base import FieldBase
@@ -36,6 +37,30 @@ def _get_limits(value: float | ArrayLike, dim: int) -> np.ndarray:
         return np.atleast_1d(value)[:dim]
 
 
+@dataclass
+class FFmpegPixelFormat:
+    pix_fmt: str
+    channels: int
+    value_max: int
+    dtype: DTypeLike
+
+    def data_to_frame(self, normalized_data: np.ndarray) -> np.ndarray:
+        return (normalized_data * self.value_max).astype(self.dtype)
+
+    def data_from_frame(self, frame_data: np.ndarray):
+        return frame_data.astype(float) / self.value_max
+
+
+FFmpegPixelFormats = {
+    "gray": FFmpegPixelFormat(
+        pix_fmt="gray", channels=1, value_max=255, dtype=np.uint8
+    ),
+    "rgb24": FFmpegPixelFormat(
+        pix_fmt="rgb24", channels=3, value_max=255, dtype=np.uint8
+    ),
+}
+
+
 class MovieStorage(StorageBase):
     def __init__(
         self,
@@ -43,6 +68,7 @@ def __init__(
         *,
         vmin: float | ArrayLike = 0,
         vmax: float | ArrayLike = 1,
+        bitrate: int = -1,
         info: InfoDict | None = None,
         write_mode: WriteModeType = "truncate_once",
     ):
@@ -54,6 +80,9 @@ def __init__(
                 Lowest values that are encoded (per field)
             vmax (float or array):
                 Highest values that are encoded (per field)
+            bitrate (float):
+                The bitrate of the movie (in kilobits per second). The default value of
+                -1 let's the encode choose an appropriate bit rate.
             info (dict):
                 Supplies extra information that is stored in the storage
             write_mode (str):
@@ -71,6 +100,9 @@ def __init__(
             - track whether times roughly work (checking for frame drops)
             - we could embedd extra information (like time, and maybe colorscaling) in
               the individual frames if we extended the shape
+            - introduce encode/decode method for images
+              (provide abstract classes that manage bit depth, color channels and ffmpeg
+              codes)
         """
         if not module_available("ffmpeg"):
             raise ModuleNotFoundError("`MovieStorage` needs `ffmpeg-python` package")
@@ -79,8 +111,7 @@ def __init__(
         self.filename = Path(filename)
         self.vmin = vmin
         self.vmax = vmax
-        self.dt = None
-        self.t_start = None
+        self.bitrate = bitrate
 
         self._logger = logging.getLogger(self.__class__.__name__)
         self._ffmpeg: Any = None
@@ -155,9 +186,7 @@ def _read_metadata(self) -> None:
         else:
             self._logger.warning(f"Unknown pixel format {stream['pix_fmt']}")
 
-    def _init_normalization(
-        self, field: FieldBase, *, inverse: bool = False
-    ) -> list[Callable]:
+    def _init_normalization(self, field: FieldBase, *, inverse: bool = False) -> None:
         """initialize the normalizations of the color information
 
         Args:
@@ -167,7 +196,7 @@ def _init_normalization(
                 Whether inverse normalization function should be returned
         """
         self._norms = []
-        fields = field if isinstance(field, FieldCollection) else [field]
+        fields = field.fields if isinstance(field, FieldCollection) else [field]
         vmin = _get_limits(self.vmin, len(fields))
         vmax = _get_limits(self.vmax, len(fields))
         for f_id, f in enumerate(fields):
@@ -179,13 +208,27 @@ def _init_normalization(
             self._norms.extend([norm] * num)
 
     def _reshape_data(self, data: np.ndarray) -> np.ndarray:
-        """reshape data such that it has exactly three dimensions"""
-        data = np.copy(data)
-        if data.ndim < 3:
-            data = data.reshape(data.shape + (1,) * (3 - data.ndim))  # set ndim=3
+        """reshape data such that it has exactly three dimensions
+
+        Assumes that that the input data has the spatial dimension last. The returned
+        data will have two spatial dimensions and all other dimensions (tensorial and
+        field collections) collapsed to the last dimension
+        """
+        data = np.copy(data)  # safety, so we don't change the shape of the original
+
+        # make sure there are two spatial dimensions
+        grid_dim = self._grid.num_axes
+        if grid_dim > 2:
+            raise NotImplementedError
+        if grid_dim == 1:
+            data = data.reshape(data.shape + (1,))
+        assert data.ndim >= 2
+        if data.ndim == 2:
+            return data.reshape(data.shape + (1,))  # explicitly scalar data
         elif data.ndim > 3:
-            data = data.reshape(data.shape[:2] + (-1,))  # collapse last axes
-        return data
+            # collapse first dimensions, so we have three in total
+            data = data.reshape((-1,) + data.shape[-2:])
+        return np.moveaxis(data, 0, -1)
 
     def start_writing(self, field: FieldBase, info: InfoDict | None = None) -> None:
         """initialize the storage for writing data
@@ -215,21 +258,21 @@ def start_writing(self, field: FieldBase, info: InfoDict | None = None) -> None:
             self.info.update(info)
 
         # get spatial dimension of the video
-        grid = field.grid
-        if grid.num_axes == 1:
+        self._grid = field.grid
+        if self._grid.num_axes == 1:
             width, height = field.grid.shape[0], 1
-        elif grid.num_axes == 2:
+        elif self._grid.num_axes == 2:
             width, height = field.grid.shape
         else:
             raise RuntimeError("Cannot use grid with more than two axes")
 
         # get color channel information
         if isinstance(field, ScalarField):
-            pix_fmt = "gray"  # 1 color channel
-            self._data_shape = (width, height, 1)
+            self.info["pixel_format"] = "gray"
         else:
-            pix_fmt = "rgb24"  # 3 color channels
-            self._data_shape = (width, height, 3)
+            self.info["pixel_format"] = "rgb24"
+        self._pix_fmt = FFmpegPixelFormats[self.info["pixel_format"]]
+        self._data_shape = (width, height, self._pix_fmt.channels)
         self.info["frame_shape"] = self._data_shape
         self.info["field_shape"] = field.data.shape
         if field.is_complex:
@@ -240,16 +283,23 @@ def start_writing(self, field: FieldBase, info: InfoDict | None = None) -> None:
 
         # set input
         self._logger.debug(f"Start ffmpeg process for `{self.filename}`")
-        f_input = ffmpeg.input(
-            "pipe:", format="rawvideo", pix_fmt=pix_fmt, s=f"{width}x{height}"
-        )
+        args = {
+            "format": "rawvideo",
+            "s": f"{width}x{height}",
+            "pix_fmt": self._pix_fmt.pix_fmt,
+            "loglevel": "warning",
+        }
+        f_input = ffmpeg.input("pipe:", **args)
         # set output format
-        f_output = f_input.output(
-            filename=self.filename,
-            pix_fmt=pix_fmt,
-            metadata="comment=" + shlex.quote(self._get_metadata()),
-        )
-        f_output = f_output.overwrite_output()  # allow overwriting file
+        args = {
+            "pix_fmt": self._pix_fmt.pix_fmt,
+            "metadata": "comment=" + shlex.quote(self._get_metadata()),
+        }
+        if self.bitrate > 0:
+            args["video_bitrate"] = self.bitrate
+        f_output = f_input.output(filename=self.filename, **args)
+        if self.write_mode in {"truncate", "truncate_once"}:
+            f_output = f_output.overwrite_output()  # allow overwriting file
         self._ffmpeg = f_output.run_async(pipe_stdin=True)  # start process
         self._state = "writing"
 
@@ -265,34 +315,28 @@ def _append_data(self, data: np.ndarray, time: float) -> None:
                 "Writing not initialized. Call "
                 f"`{self.__class__.__name__}.start_writing`"
             )
-
         # ensure the data has the shape width x height x depth
         data = self._reshape_data(data)
         assert len(self._data_shape) == data.ndim == 3
+        assert data.shape[2] == len(self._norms)
         if self._data_shape[2] == 1:
             # single color channel
             data = data.reshape(self._data_shape)
         elif self._data_shape[2] == 3:
             # three color channels
-            if data.shape[2] == 1:
-                z = np.zeros(data.shape[:2] + (1,), dtype=self._dtype)
-                data = np.dstack((data, z, z))
-            elif data.shape[2] == 2:
-                z = np.zeros(data.shape[:2] + (1,), dtype=self._dtype)
-                data = np.dstack((data[..., 0], data[..., 1], z))
+            if data.shape[2] < 3:
+                zero_shape = data.shape[:2] + (3 - data.shape[2],)
+                data = np.dstack((data, np.zeros(zero_shape, dtype=self._dtype)))
         else:
             raise RuntimeError
-
         assert data.shape == self._data_shape
-        assert data.shape[2] == len(self._norms)
 
         # map values to [0, 1] float values
-        for i in range(data.shape[2]):
-            data[..., i] = self._norms[i](data[..., i])
-        data *= 256
+        for i, norm in enumerate(self._norms):
+            data[..., i] = norm(data[..., i])
 
         # write the new data
-        self._ffmpeg.stdin.write(data.astype(np.uint8).tobytes())
+        self._ffmpeg.stdin.write(self._pix_fmt.data_to_frame(data).tobytes())
 
     def end_writing(self) -> None:
         """finalize the storage after writing"""
@@ -329,21 +373,22 @@ def __iter__(self) -> Iterator[FieldBase]:
         # frame_shape = (self.info["width"], self.info["height"], self.info["channels"])
         data_shape = (self.info["width"], self.info["height"], len(self._norms))
         data = np.empty(data_shape, dtype=self._dtype)
+        pix_fmt = FFmpegPixelFormats[self.info["pixel_format"]]
 
         # iterate over entire movie
-        f_input = ffmpeg.input(self.filename)
+        f_input = ffmpeg.input(self.filename, loglevel="warning")
         f_output = f_input.output(
-            "pipe:", format="rawvideo", pix_fmt="rgb24", vframes=8
+            "pipe:", format="rawvideo", pix_fmt=pix_fmt.pix_fmt, vframes=8
         )
         proc = f_output.run_async(pipe_stdout=True)
         while True:
             read_bytes = proc.stdout.read(np.prod(data_shape))
             if not read_bytes:
                 break
-            frame = np.frombuffer(read_bytes, np.uint8).reshape(data_shape)
+            frame = np.frombuffer(read_bytes, pix_fmt.dtype).reshape(data_shape)
 
             for i, norm in enumerate(self._norms):
-                data[..., i] = norm(frame[:, :, i] / 256)
+                data[..., i] = norm(pix_fmt.data_from_frame(frame[:, :, i]))
 
             # create the field with the data of the given index
             assert self._field is not None
diff --git a/pde/tools/resources/requirements_full.txt b/pde/tools/resources/requirements_full.txt
index b901132d..8730d2f6 100644
--- a/pde/tools/resources/requirements_full.txt
+++ b/pde/tools/resources/requirements_full.txt
@@ -1,4 +1,5 @@
 # These are the full requirements used to test all functions
+ffmpeg-python>=0.2
 h5py>=2.10
 matplotlib>=3.1
 numba>=0.59
diff --git a/pyproject.toml b/pyproject.toml
index c6cd0d2c..e8b579d3 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -92,6 +92,7 @@ follow_imports_for_stubs = true
 
 [[tool.mypy.overrides]]
 module = [
+    "ffmpeg.*",
     "h5py.*",
     "IPython.*",
     "ipywidgets.*",
diff --git a/scripts/_templates/pyproject.toml b/scripts/_templates/pyproject.toml
index ef572028..d008b82f 100644
--- a/scripts/_templates/pyproject.toml
+++ b/scripts/_templates/pyproject.toml
@@ -89,6 +89,7 @@ follow_imports_for_stubs = true
 
 [[tool.mypy.overrides]]
 module = [
+    "ffmpeg.*",
     "h5py.*",
     "IPython.*",
     "ipywidgets.*",
diff --git a/scripts/create_requirements.py b/scripts/create_requirements.py
index 2d6da2cb..0677c164 100755
--- a/scripts/create_requirements.py
+++ b/scripts/create_requirements.py
@@ -95,6 +95,12 @@ def line(self, relation: str = ">=") -> str:
         essential=True,
     ),
     # general, optional requirements
+    Requirement(
+        name="ffmpeg-python",
+        version_min="0.2",
+        usage="Reading and writing videos",
+        collections={"full"},
+    ),
     Requirement(
         name="h5py",
         version_min="2.10",
diff --git a/tests/requirements_full.txt b/tests/requirements_full.txt
index b901132d..8730d2f6 100644
--- a/tests/requirements_full.txt
+++ b/tests/requirements_full.txt
@@ -1,4 +1,5 @@
 # These are the full requirements used to test all functions
+ffmpeg-python>=0.2
 h5py>=2.10
 matplotlib>=3.1
 numba>=0.59
diff --git a/tests/storage/test_movie_storages.py b/tests/storage/test_movie_storages.py
new file mode 100644
index 00000000..7f6cef07
--- /dev/null
+++ b/tests/storage/test_movie_storages.py
@@ -0,0 +1,50 @@
+"""
+.. codeauthor:: David Zwicker <david.zwicker@ds.mpg.de>
+"""
+
+import numpy as np
+import pytest
+
+import pde
+from pde import MovieStorage
+from pde.tools.misc import module_available
+
+
+@pytest.mark.skipif(not module_available("ffmpeg"), reason="requires `ffmpeg` module")
+@pytest.mark.parametrize("dim", [1, 2])
+def test_movie_storage_scalar_field(dim, tmp_path):
+    """test writing scalar field"""
+    path = tmp_path / f"test_movie_storage_scalar_{dim}.mp4"
+
+    grid = pde.UnitGrid([8] * dim)
+    field = pde.ScalarField(grid)
+    eq = pde.PDE({"a": "1"})
+    writer = MovieStorage(path, vmax=5)
+    eq.solve(field, t_range=3.5, dt=0.1, backend="numpy", tracker=writer.tracker(1))
+
+    reader = MovieStorage(path)
+    assert len(reader) == 4
+    for i, field in enumerate(reader):
+        assert field.grid == grid
+        np.testing.assert_allclose(field.data, i, rtol=0.2)
+
+
+@pytest.mark.skipif(not module_available("ffmpeg"), reason="requires `ffmpeg` module")
+@pytest.mark.parametrize("dim", [1, 2])
+@pytest.mark.parametrize("num_fields", [1, 2, 3])
+def test_movie_storage_scalar_fields(dim, num_fields, tmp_path):
+    """test writing three scalar field"""
+    path = tmp_path / f"test_movie_storage_scalar_{dim}_{num_fields}.mp4"
+
+    grid = pde.UnitGrid([8] * dim)
+    field = pde.FieldCollection([pde.ScalarField(grid)] * num_fields, copy_fields=True)
+    eq = pde.PDE(dict([("a", "1"), ("b", "2"), ("c", "3")][:num_fields]))
+    writer = MovieStorage(path, vmax=[5, 8, 12])
+    eq.solve(field, t_range=3.5, dt=0.1, backend="numpy", tracker=writer.tracker(1))
+
+    reader = MovieStorage(path)
+    assert len(reader) == 4
+    for i, fields in enumerate(reader):
+        assert fields.grid == grid
+        for j, field in enumerate(fields, 1):
+            np.testing.assert_allclose(field.data, j * i)

From 7b8e938a4772dc04d2f7c280d144a91b7b2dd3ac Mon Sep 17 00:00:00 2001
From: David Zwicker <david.zwicker@ds.mpg.de>
Date: Sat, 16 Mar 2024 23:36:35 +0100
Subject: [PATCH 03/13] Minor improvement

---
 pde/storage/_ffmpeg.py               | 60 ++++++++++++++++++++++++
 pde/storage/movie.py                 | 69 +++++++++++++---------------
 tests/storage/test_movie_storages.py | 40 +++++++++-------
 3 files changed, 114 insertions(+), 55 deletions(-)
 create mode 100644 pde/storage/_ffmpeg.py

diff --git a/pde/storage/_ffmpeg.py b/pde/storage/_ffmpeg.py
new file mode 100644
index 00000000..bd843d5e
--- /dev/null
+++ b/pde/storage/_ffmpeg.py
@@ -0,0 +1,60 @@
+"""
+Functions for interacting with ffmpeg
+
+.. codeauthor:: David Zwicker <david.zwicker@ds.mpg.de>
+"""
+
+import subprocess as sp
+from dataclasses import dataclass
+
+import numpy as np
+from numpy.typing import DTypeLike
+
+
+@dataclass
+class FFmpegPixelFormat:
+    pix_fmt: str
+    channels: int
+    value_max: int
+    dtype: DTypeLike
+
+    def data_to_frame(self, normalized_data: np.ndarray) -> np.ndarray:
+        return (normalized_data * self.value_max).astype(self.dtype)
+
+    def data_from_frame(self, frame_data: np.ndarray):
+        return frame_data.astype(float) / self.value_max
+
+
+pixel_formats = {
+    "gray": FFmpegPixelFormat(
+        pix_fmt="gray", channels=1, value_max=255, dtype=np.uint8
+    ),
+    "rgb24": FFmpegPixelFormat(
+        pix_fmt="rgb24", channels=3, value_max=255, dtype=np.uint8
+    ),
+    "bgr24": FFmpegPixelFormat(
+        pix_fmt="bgr24", channels=3, value_max=255, dtype=np.uint8
+    ),
+    "rgb32": FFmpegPixelFormat(
+        pix_fmt="rgb32", channels=4, value_max=255, dtype=np.uint8
+    ),
+    "gbrp": FFmpegPixelFormat(
+        pix_fmt="gbrp", channels=4, value_max=255, dtype=np.uint8
+    ),
+}
+
+
+def _run_ffmpeg(args: list[str]):
+    return sp.check_output(["ffmpeg"] + args)
+
+
+def codecs() -> list[str]:
+    """list: all supported ffmpeg codecs"""
+    res = _run_ffmpeg(["-codecs"])
+
+
+def get_pixel_formats(encoder=None):
+    if encoder is None:
+        res = _run_ffmpeg(["-pix_fmts"])
+    else:
+        res = _run_ffmpeg(["-h", f"encoder={encoder}"])
diff --git a/pde/storage/movie.py b/pde/storage/movie.py
index a583b4fd..1b989fe8 100644
--- a/pde/storage/movie.py
+++ b/pde/storage/movie.py
@@ -13,19 +13,19 @@
 import logging
 import shlex
 from collections.abc import Iterator
-from dataclasses import dataclass
 from pathlib import Path
 from typing import Any, Callable, Literal
 
 import numpy as np
 from matplotlib.colors import Normalize
-from numpy.typing import ArrayLike, DTypeLike
+from numpy.typing import ArrayLike
 
 from ..fields import FieldCollection, ScalarField
 from ..fields.base import FieldBase
 from ..tools.docstrings import fill_in_docstring
 from ..tools.misc import module_available
 from ..trackers.interrupts import ConstantInterrupts
+from . import _ffmpeg as FFmpeg
 from .base import InfoDict, StorageBase, StorageTracker, WriteModeType
 
 
@@ -37,31 +37,11 @@ def _get_limits(value: float | ArrayLike, dim: int) -> np.ndarray:
         return np.atleast_1d(value)[:dim]
 
 
-@dataclass
-class FFmpegPixelFormat:
-    pix_fmt: str
-    channels: int
-    value_max: int
-    dtype: DTypeLike
-
-    def data_to_frame(self, normalized_data: np.ndarray) -> np.ndarray:
-        return (normalized_data * self.value_max).astype(self.dtype)
-
-    def data_from_frame(self, frame_data: np.ndarray):
-        return frame_data.astype(float) / self.value_max
-
-
-FFmpegPixelFormats = {
-    "gray": FFmpegPixelFormat(
-        pix_fmt="gray", channels=1, value_max=255, dtype=np.uint8
-    ),
-    "rgb24": FFmpegPixelFormat(
-        pix_fmt="rgb24", channels=3, value_max=255, dtype=np.uint8
-    ),
-}
+class MovieStorage(StorageBase):
 
+    codecs: list[str] = ["libx264", "libx264rgb", "libx265"]
+    """list: List of prefered codecs"""
 
-class MovieStorage(StorageBase):
     def __init__(
         self,
         filename: str | Path,
@@ -97,6 +77,9 @@ def __init__(
             - allow more bits for colorchannels
             - allow choosing bitrate for video
             - support different write_mode
+            - support different video codecs (from a list of priorities)
+               Check in order which one supports the chosen pixel format (and store that
+               information)
             - track whether times roughly work (checking for frame drops)
             - we could embedd extra information (like time, and maybe colorscaling) in
               the individual frames if we extended the shape
@@ -159,6 +142,7 @@ def _read_metadata(self) -> None:
         import ffmpeg
 
         info = ffmpeg.probe(self.filename)
+        # TODO: Check stored pixel format against the one
 
         # sanity checks on the video
         nb_streams = info["format"]["nb_streams"]
@@ -179,12 +163,13 @@ def _read_metadata(self) -> None:
         self.info["num_frames"] = int(stream["nb_frames"])
         self.info["width"] = stream["coded_width"]
         self.info["height"] = stream["coded_height"]
-        if stream["pix_fmt"] == "gray":
-            self.info["channels"] = 1
-        elif stream["pix_fmt"] in {"rgb24", "yuv420p"}:
-            self.info["channels"] = 3
-        else:
-            self._logger.warning(f"Unknown pixel format {stream['pix_fmt']}")
+        if self.info["pixel_format"] != stream["pix_fmt"]:
+            self._logger.warning(
+                f"Inconsistent pixel format {self.info['pixel_format']} != "
+                f"{stream['pix_fmt']}"
+            )
+        if self.info["pixel_format"] not in FFmpeg.pixel_formats:
+            self._logger.warning(f"Unknown pixel format {self.info['pixel_format']}")
 
     def _init_normalization(self, field: FieldBase, *, inverse: bool = False) -> None:
         """initialize the normalizations of the color information
@@ -268,10 +253,12 @@ def start_writing(self, field: FieldBase, info: InfoDict | None = None) -> None:
 
         # get color channel information
         if isinstance(field, ScalarField):
+            codec = "libx264"
             self.info["pixel_format"] = "gray"
         else:
+            codec = "libx264rgb"
             self.info["pixel_format"] = "rgb24"
-        self._pix_fmt = FFmpegPixelFormats[self.info["pixel_format"]]
+        self._pix_fmt = FFmpeg.pixel_formats[self.info["pixel_format"]]
         self._data_shape = (width, height, self._pix_fmt.channels)
         self.info["frame_shape"] = self._data_shape
         self.info["field_shape"] = field.data.shape
@@ -292,6 +279,8 @@ def start_writing(self, field: FieldBase, info: InfoDict | None = None) -> None:
         f_input = ffmpeg.input("pipe:", **args)
         # set output format
         args = {
+            "vcodec": codec,
+            "crf": "0",  # Constant Rate Factor - aim for lossless compression
             "pix_fmt": self._pix_fmt.pix_fmt,
             "metadata": "comment=" + shlex.quote(self._get_metadata()),
         }
@@ -354,7 +343,11 @@ def __len__(self):
     @property
     def times(self):
         """:class:`~numpy.ndarray`: The times at which data is available"""
-        return np.arange(self.t_start, len(self), self.dt)
+        t_start = self.info.get("t_start")
+        if t_start is None:
+            t_start = 0
+        dt = self.info.get("dt", 1)
+        return t_start + dt * np.arange(len(self))
 
     @property
     def data(self):
@@ -370,10 +363,10 @@ def __iter__(self) -> Iterator[FieldBase]:
         if self._field is None:
             self._init_field()
         self._init_normalization(self._field, inverse=True)
-        # frame_shape = (self.info["width"], self.info["height"], self.info["channels"])
+        pix_fmt = FFmpeg.pixel_formats[self.info["pixel_format"]]
+        frame_shape = (self.info["width"], self.info["height"], pix_fmt.channels)
         data_shape = (self.info["width"], self.info["height"], len(self._norms))
         data = np.empty(data_shape, dtype=self._dtype)
-        pix_fmt = FFmpegPixelFormats[self.info["pixel_format"]]
 
         # iterate over entire movie
         f_input = ffmpeg.input(self.filename, loglevel="warning")
@@ -385,7 +378,7 @@ def __iter__(self) -> Iterator[FieldBase]:
             read_bytes = proc.stdout.read(np.prod(data_shape))
             if not read_bytes:
                 break
-            frame = np.frombuffer(read_bytes, pix_fmt.dtype).reshape(data_shape)
+            frame = np.frombuffer(read_bytes, pix_fmt.dtype).reshape(frame_shape)
 
             for i, norm in enumerate(self._norms):
                 data[..., i] = norm(pix_fmt.data_from_frame(frame[:, :, i]))
@@ -415,8 +408,8 @@ def tracker(
         """create object that can be used as a tracker to fill this storage
 
         Args:
-            interrupts:
-                {ARG_TRACKER_INTERRUPT}
+            interrupts (:class:`~pde.tracker.ConstantInterrupts` or float):
+                Time interval with which the tracker is being called
             transformation (callable, optional):
                 A function that transforms the current state into a new field or field
                 collection, which is then stored. This allows to store derived
diff --git a/tests/storage/test_movie_storages.py b/tests/storage/test_movie_storages.py
index 7f6cef07..b53d7c02 100644
--- a/tests/storage/test_movie_storages.py
+++ b/tests/storage/test_movie_storages.py
@@ -11,29 +11,35 @@
 
 
 @pytest.mark.skipif(not module_available("ffmpeg"), reason="requires `ffmpeg` module")
-@pytest.mark.parametrize("dim", [1, 2])
-def test_movie_storage_scalar_field(dim, tmp_path):
-    """test writing scalar field"""
-    path = tmp_path / f"test_movie_storage_scalar_{dim}.mp4"
-
-    grid = pde.UnitGrid([8] * dim)
-    field = pde.ScalarField(grid)
-    eq = pde.PDE({"a": "1"})
-    writer = MovieStorage(path, vmax=5)
-    eq.solve(field, t_range=3.5, dt=0.1, backend="numpy", tracker=writer.tracker(1))
+def test_movie_storage_simple(tmp_path, rng):
+    """test storing field as movie"""
+    path = tmp_path / f"test_movie_storage_simple.mov"
+
+    grid = pde.UnitGrid([16, 16])
+    field = pde.ScalarField.random_uniform(grid, rng=rng)
+    eq = pde.DiffusionPDE()
+    writer = MovieStorage(path)
+    storage = pde.MemoryStorage()
+    eq.solve(
+        field,
+        t_range=3.5,
+        dt=0.1,
+        backend="numpy",
+        tracker=[storage.tracker(2), writer.tracker(2)],
+    )
 
     reader = MovieStorage(path)
-    assert len(reader) == 4
+    assert len(reader) == 2
+    np.testing.assert_allclose(reader.times, [0, 2])
     for i, field in enumerate(reader):
-        assert field.grid == grid
-        np.testing.assert_allclose(field.data, i, rtol=0.2)
+        np.testing.assert_allclose(field.data, storage[i].data, atol=0.1)
 
 
 @pytest.mark.skipif(not module_available("ffmpeg"), reason="requires `ffmpeg` module")
 @pytest.mark.parametrize("dim", [1, 2])
 @pytest.mark.parametrize("num_fields", [1, 2, 3])
 def test_movie_storage_scalar_fields(dim, num_fields, tmp_path):
-    """test writing three scalar field"""
+    """test storing field as movie"""
     path = tmp_path / f"test_movie_storage_scalar_{dim}_{num_fields}.mp4"
 
     grid = pde.UnitGrid([8] * dim)
@@ -44,7 +50,7 @@ def test_movie_storage_scalar_fields(dim, num_fields, tmp_path):
 
     reader = MovieStorage(path)
     assert len(reader) == 4
-    for i, fields in enumerate(reader):
+    for t, fields in enumerate(reader):
         assert fields.grid == grid
-        for j, field in enumerate(fields, 1):
-            np.testing.assert_allclose(field.data, j * i)
+        for i, field in enumerate(fields, 1):
+            np.testing.assert_allclose(field.data, i * t, atol=0.1, rtol=0.1)

From 64d3288faac32bc40ca1ca64441d49ef02588478 Mon Sep 17 00:00:00 2001
From: David Zwicker <david.zwicker@ds.mpg.de>
Date: Sun, 17 Mar 2024 11:26:45 +0100
Subject: [PATCH 04/13] Fixed a nasty bug involving the normalization

---
 pde/storage/_ffmpeg.py               |  83 ++++++----
 pde/storage/movie.py                 | 233 +++++++++++++++------------
 tests/storage/test_movie_storages.py |  71 ++++++--
 3 files changed, 235 insertions(+), 152 deletions(-)

diff --git a/pde/storage/_ffmpeg.py b/pde/storage/_ffmpeg.py
index bd843d5e..54cb51f8 100644
--- a/pde/storage/_ffmpeg.py
+++ b/pde/storage/_ffmpeg.py
@@ -1,23 +1,45 @@
 """
-Functions for interacting with ffmpeg
+Functions for interacting with FFmpeg
 
 .. codeauthor:: David Zwicker <david.zwicker@ds.mpg.de>
 """
 
-import subprocess as sp
+# import subprocess as sp
 from dataclasses import dataclass
 
 import numpy as np
 from numpy.typing import DTypeLike
 
+# def _run_ffmpeg(args: list[str]):
+#     return sp.check_output(["ffmpeg"] + args)
+#
+#
+# def codecs() -> list[str]:
+#     """list: all supported ffmpeg codecs"""
+#     res = _run_ffmpeg(["-codecs"])
+#
+#
+# def get_pixel_formats(encoder=None):
+#     if encoder is None:
+#         res = _run_ffmpeg(["-pix_fmts"])
+#     else:
+#         res = _run_ffmpeg(["-h", f"encoder={encoder}"])
+
 
 @dataclass
-class FFmpegPixelFormat:
-    pix_fmt: str
+class FFmpegFormat:
+    """defines a FFmpeg format used for storing field data in a video"""
+
+    pix_fmt_file: str
+    codec: str
     channels: int
     value_max: int
     dtype: DTypeLike
 
+    @property
+    def pix_fmt_data(self) -> str:
+        return {1: "gray", 3: "rgb24", 4: "rgba"}[self.channels]
+
     def data_to_frame(self, normalized_data: np.ndarray) -> np.ndarray:
         return (normalized_data * self.value_max).astype(self.dtype)
 
@@ -25,36 +47,33 @@ def data_from_frame(self, frame_data: np.ndarray):
         return frame_data.astype(float) / self.value_max
 
 
-pixel_formats = {
-    "gray": FFmpegPixelFormat(
-        pix_fmt="gray", channels=1, value_max=255, dtype=np.uint8
+formats = {
+    "gray": FFmpegFormat(
+        pix_fmt_file="gray",
+        codec="libx264",
+        channels=1,
+        value_max=255,
+        dtype=np.uint8,
     ),
-    "rgb24": FFmpegPixelFormat(
-        pix_fmt="rgb24", channels=3, value_max=255, dtype=np.uint8
+    "rgb24": FFmpegFormat(
+        pix_fmt_file="rgb24",
+        codec="libx264rgb",
+        channels=3,
+        value_max=255,
+        dtype=np.uint8,
     ),
-    "bgr24": FFmpegPixelFormat(
-        pix_fmt="bgr24", channels=3, value_max=255, dtype=np.uint8
+    "bgr24": FFmpegFormat(
+        pix_fmt_file="bgr24",
+        codec="libx264rgb",
+        channels=3,
+        value_max=255,
+        dtype=np.uint8,
     ),
-    "rgb32": FFmpegPixelFormat(
-        pix_fmt="rgb32", channels=4, value_max=255, dtype=np.uint8
-    ),
-    "gbrp": FFmpegPixelFormat(
-        pix_fmt="gbrp", channels=4, value_max=255, dtype=np.uint8
+    "rgb32": FFmpegFormat(
+        pix_fmt_file="rgb32",
+        codec="libx264rgb",
+        channels=4,
+        value_max=255,
+        dtype=np.uint8,
     ),
 }
-
-
-def _run_ffmpeg(args: list[str]):
-    return sp.check_output(["ffmpeg"] + args)
-
-
-def codecs() -> list[str]:
-    """list: all supported ffmpeg codecs"""
-    res = _run_ffmpeg(["-codecs"])
-
-
-def get_pixel_formats(encoder=None):
-    if encoder is None:
-        res = _run_ffmpeg(["-pix_fmts"])
-    else:
-        res = _run_ffmpeg(["-h", f"encoder={encoder}"])
diff --git a/pde/storage/movie.py b/pde/storage/movie.py
index 1b989fe8..affd738a 100644
--- a/pde/storage/movie.py
+++ b/pde/storage/movie.py
@@ -7,12 +7,20 @@
 .. codeauthor:: David Zwicker <david.zwicker@ds.mpg.de> 
 """
 
+# TODO:
+#     - allow more bits for colorchannels
+#     - allow reading single frames
+#     - support different write_mode
+#     - track whether times roughly work (checking for frame drops)
+#     - we could embedd extra information (like time, and maybe colorscaling) in
+#       the individual frames if we extended the shape
+
 from __future__ import annotations
 
 import json
 import logging
 import shlex
-from collections.abc import Iterator
+from collections.abc import Iterator, Sequence
 from pathlib import Path
 from typing import Any, Callable, Literal
 
@@ -21,7 +29,7 @@
 from numpy.typing import ArrayLike
 
 from ..fields import FieldCollection, ScalarField
-from ..fields.base import FieldBase
+from ..fields.base import DataFieldBase, FieldBase
 from ..tools.docstrings import fill_in_docstring
 from ..tools.misc import module_available
 from ..trackers.interrupts import ConstantInterrupts
@@ -32,15 +40,20 @@
 def _get_limits(value: float | ArrayLike, dim: int) -> np.ndarray:
     """helper function creating sequence of length `dim` from input"""
     if np.isscalar(value):
-        return np.full(dim, value)
+        return np.full(dim, value, dtype=float)
     else:
-        return np.atleast_1d(value)[:dim]
+        return np.asarray(value)[:dim].astype(float)
 
 
 class MovieStorage(StorageBase):
+    """store discretized fields in a movie file
 
-    codecs: list[str] = ["libx264", "libx264rgb", "libx265"]
-    """list: List of prefered codecs"""
+    Warning:
+        This storage potentially compresses data and can thus lead to loss of some
+        information. The data quality depends on many parameters, but most important are
+        the bit depth of the video format, the range that is encoded (determined by
+        `vmin` and `vmax`), and the target bitrate.
+    """
 
     def __init__(
         self,
@@ -48,18 +61,28 @@ def __init__(
         *,
         vmin: float | ArrayLike = 0,
         vmax: float | ArrayLike = 1,
+        video_format: str = "auto",
         bitrate: int = -1,
         info: InfoDict | None = None,
         write_mode: WriteModeType = "truncate_once",
+        loglevel: str = "warning",
     ):
         """
         Args:
             filename (str):
-                The path to the hdf5-file where the data is stored
+                The path where the movie is stored. The file extension determines the
+                container format of the movie.
             vmin (float or array):
-                Lowest values that are encoded (per field)
+                Lowest values that are encoded (per field). Lower values are clipped to
+                this value.
             vmax (float or array):
-                Highest values that are encoded (per field)
+                Highest values that are encoded (per field). Larger values are clipped
+                to this value.
+            video_format (str):
+                How to write data to the movie. This determines the number of color
+                channels and the bit depth of individual colors. Available options are
+                listed in :func:`~pde.storage._ffmpeg.formats`. The special value
+                `auto` tries to find a suitable format automatically.
             bitrate (float):
                 The bitrate of the movie (in kilobits per second). The default value of
                 -1 let's the encode choose an appropriate bit rate.
@@ -71,21 +94,8 @@ def __init__(
                 cleared every time this storage is used for writing), or 'truncate_once'
                 (data is cleared for the first writing, but appended subsequently).
                 Alternatively, specifying 'readonly' will disable writing completely.
-
-
-        TODO:
-            - allow more bits for colorchannels
-            - allow choosing bitrate for video
-            - support different write_mode
-            - support different video codecs (from a list of priorities)
-               Check in order which one supports the chosen pixel format (and store that
-               information)
-            - track whether times roughly work (checking for frame drops)
-            - we could embedd extra information (like time, and maybe colorscaling) in
-              the individual frames if we extended the shape
-            - introduce encode/decode method for images
-              (provide abstract classes that manage bit depth, color channels and ffmpeg
-              codes)
+            loglevel (str):
+                FFmpeg log level
         """
         if not module_available("ffmpeg"):
             raise ModuleNotFoundError("`MovieStorage` needs `ffmpeg-python` package")
@@ -94,7 +104,9 @@ def __init__(
         self.filename = Path(filename)
         self.vmin = vmin
         self.vmax = vmax
+        self.video_format = video_format
         self.bitrate = bitrate
+        self.loglevel = loglevel
 
         self._logger = logging.getLogger(self.__class__.__name__)
         self._ffmpeg: Any = None
@@ -142,7 +154,6 @@ def _read_metadata(self) -> None:
         import ffmpeg
 
         info = ffmpeg.probe(self.filename)
-        # TODO: Check stored pixel format against the one
 
         # sanity checks on the video
         nb_streams = info["format"]["nb_streams"]
@@ -163,13 +174,16 @@ def _read_metadata(self) -> None:
         self.info["num_frames"] = int(stream["nb_frames"])
         self.info["width"] = stream["coded_width"]
         self.info["height"] = stream["coded_height"]
-        if self.info["pixel_format"] != stream["pix_fmt"]:
-            self._logger.warning(
-                f"Inconsistent pixel format {self.info['pixel_format']} != "
-                f"{stream['pix_fmt']}"
-            )
-        if self.info["pixel_format"] not in FFmpeg.pixel_formats:
-            self._logger.warning(f"Unknown pixel format {self.info['pixel_format']}")
+        try:
+            self._format = FFmpeg.formats[self.info["video_format"]]
+        except KeyError:
+            self._logger.warning(f"Unknown pixel format `{self.info['pixel_format']}`")
+        else:
+            if self._format.pix_fmt_file != stream["pix_fmt"]:
+                self._logger.info(
+                    "Pixel format differs from requested one: "
+                    f"{self._format.pix_fmt_file} != {stream['pix_fmt']}"
+                )
 
     def _init_normalization(self, field: FieldBase, *, inverse: bool = False) -> None:
         """initialize the normalizations of the color information
@@ -179,42 +193,21 @@ def _init_normalization(self, field: FieldBase, *, inverse: bool = False) -> Non
                 Example field to obtain information about grid and data rank
             inverse (bool):
                 Whether inverse normalization function should be returned
+
+        The resulting normalization functions are stored in `self._norms`
         """
         self._norms = []
-        fields = field.fields if isinstance(field, FieldCollection) else [field]
+        if isinstance(field, FieldCollection):
+            fields: Sequence[DataFieldBase] = field.fields
+        else:
+            fields = [field]  # type: ignore
         vmin = _get_limits(self.vmin, len(fields))
         vmax = _get_limits(self.vmax, len(fields))
         for f_id, f in enumerate(fields):
-            if inverse:
-                norm = lambda data: vmin[f_id] + (vmax[f_id] - vmin[f_id]) * data
-            else:
-                norm = Normalize(vmin[f_id], vmax[f_id], clip=True)
-            num = self.grid.dim**f.rank  # independent components in the field
+            norm = Normalize(vmin[f_id], vmax[f_id], clip=True)
+            num = f.grid.dim**f.rank  # independent components in the field
             self._norms.extend([norm] * num)
 
-    def _reshape_data(self, data: np.ndarray) -> np.ndarray:
-        """reshape data such that it has exactly three dimensions
-
-        Assumes that that the input data has the spatial dimension last. The returned
-        data will have two spatial dimensions and all other dimensions (tensorial and
-        field collections) collapsed to the last dimension
-        """
-        data = np.copy(data)  # safety, so we don't change the shape of the original
-
-        # make sure there are two spatial dimensions
-        grid_dim = self._grid.num_axes
-        if grid_dim > 2:
-            raise NotImplementedError
-        if grid_dim == 1:
-            data = data.reshape(data.shape + (1,))
-        assert data.ndim >= 2
-        if data.ndim == 2:
-            return data.reshape(data.shape + (1,))  # explicitly scalar data
-        elif data.ndim > 3:
-            # collapse first dimensions, so we have three in total
-            data = data.reshape((-1,) + data.shape[-2:])
-        return np.moveaxis(data, 0, -1)
-
     def start_writing(self, field: FieldBase, info: InfoDict | None = None) -> None:
         """initialize the storage for writing data
 
@@ -252,40 +245,39 @@ def start_writing(self, field: FieldBase, info: InfoDict | None = None) -> None:
             raise RuntimeError("Cannot use grid with more than two axes")
 
         # get color channel information
-        if isinstance(field, ScalarField):
-            codec = "libx264"
-            self.info["pixel_format"] = "gray"
+        if self.video_format == "auto":
+            if isinstance(field, ScalarField):
+                self.info["video_format"] = "gray"
+            else:
+                self.info["video_format"] = "rgb24"
         else:
-            codec = "libx264rgb"
-            self.info["pixel_format"] = "rgb24"
-        self._pix_fmt = FFmpeg.pixel_formats[self.info["pixel_format"]]
-        self._data_shape = (width, height, self._pix_fmt.channels)
-        self.info["frame_shape"] = self._data_shape
-        self.info["field_shape"] = field.data.shape
+            self.info["video_format"] = self.video_format
+        self._format = FFmpeg.formats[self.info["video_format"]]
+        self._data_shape = (width, height, self._format.channels)
         if field.is_complex:
             raise NotImplementedError("MovieStorage does not support complex values")
 
         # set up the normalization
-        self._init_normalization(field)
+        self._init_normalization(field, inverse=False)
 
         # set input
         self._logger.debug(f"Start ffmpeg process for `{self.filename}`")
         args = {
             "format": "rawvideo",
             "s": f"{width}x{height}",
-            "pix_fmt": self._pix_fmt.pix_fmt,
-            "loglevel": "warning",
+            "pix_fmt": self._format.pix_fmt_data,
+            "loglevel": self.loglevel,
         }
         f_input = ffmpeg.input("pipe:", **args)
         # set output format
         args = {
-            "vcodec": codec,
-            "crf": "0",  # Constant Rate Factor - aim for lossless compression
-            "pix_fmt": self._pix_fmt.pix_fmt,
+            "vcodec": self._format.codec,
+            "crf": "10",  # Constant Rate Factor - lower values for less compression
+            "pix_fmt": self._format.pix_fmt_file,
             "metadata": "comment=" + shlex.quote(self._get_metadata()),
         }
         if self.bitrate > 0:
-            args["video_bitrate"] = self.bitrate
+            args["video_bitrate"] = str(self.bitrate)
         f_output = f_input.output(filename=self.filename, **args)
         if self.write_mode in {"truncate", "truncate_once"}:
             f_output = f_output.overwrite_output()  # allow overwriting file
@@ -304,32 +296,40 @@ def _append_data(self, data: np.ndarray, time: float) -> None:
                 "Writing not initialized. Call "
                 f"`{self.__class__.__name__}.start_writing`"
             )
-        # ensure the data has the shape width x height x depth
-        data = self._reshape_data(data)
+        assert self._norms is not None  # normalization has been initialized
+
+        # make sure there are two spatial dimensions
+        grid_dim = self._grid.num_axes
+        if grid_dim > 2:
+            raise NotImplementedError
+        if grid_dim == 1:
+            data = data.reshape(data.shape + (1,))
+        assert data.ndim >= 2
+        assert data.shape[-2:] == self._data_shape[:2]  # check spatial dimensions
+
+        # ensure the data has the shape extra_dim x width x height
+        # where `extra_dim` are separated fields or capture the rank of the field
+        if data.ndim == 2:
+            data = data.reshape((1,) + data.shape)  # explicitly scalar data
+        elif data.ndim > 3:
+            # collapse first dimensions, so we have three in total
+            data = data.reshape((-1,) + data.shape[-2:])
         assert len(self._data_shape) == data.ndim == 3
-        assert data.shape[2] == len(self._norms)
-        if self._data_shape[2] == 1:
-            # single color channel
-            data = data.reshape(self._data_shape)
-        elif self._data_shape[2] == 3:
-            # three color channels
-            if data.shape[2] < 3:
-                zero_shape = data.shape[:2] + (3 - data.shape[2],)
-                data = np.dstack((data, np.zeros(zero_shape, dtype=self._dtype)))
-        else:
-            raise RuntimeError
-        assert data.shape == self._data_shape
+        assert len(self._norms) == data.shape[0]  # same non-spatial dimension
 
-        # map values to [0, 1] float values
+        # map data values to frame values
+        frame_data = np.zeros(self._data_shape, dtype=self._format.dtype)
         for i, norm in enumerate(self._norms):
-            data[..., i] = norm(data[..., i])
+            data_norm = norm(data[i, ...])
+            frame_data[..., i] = self._format.data_to_frame(data_norm)
 
         # write the new data
-        self._ffmpeg.stdin.write(self._pix_fmt.data_to_frame(data).tobytes())
+        self._ffmpeg.stdin.write(frame_data.tobytes())
 
     def end_writing(self) -> None:
         """finalize the storage after writing"""
         if not self._state == "writing":
+            self._logger.warning("Writing was already terminated")
             return  # writing mode was already ended
         self._logger.debug("End writing")
         self.close()
@@ -354,6 +354,21 @@ def data(self):
         """:class:`~numpy.ndarray`: The actual data for all time"""
         raise NotImplementedError
 
+    def _get_field(self, t_index: int) -> FieldBase:
+        """return the field corresponding to the given time index
+
+        Load the data given an index, i.e., the data at time `self.times[t_index]`.
+
+        Args:
+            t_index (int):
+                The index of the data to load
+
+        Returns:
+            :class:`~pde.fields.FieldBase`:
+            The field class containing the grid and data
+        """
+        raise NotImplementedError
+
     def __iter__(self) -> Iterator[FieldBase]:
         """iterate over all stored fields"""
         import ffmpeg
@@ -362,31 +377,35 @@ def __iter__(self) -> Iterator[FieldBase]:
             self._read_metadata()
         if self._field is None:
             self._init_field()
+        assert self._field is not None
         self._init_normalization(self._field, inverse=True)
-        pix_fmt = FFmpeg.pixel_formats[self.info["pixel_format"]]
-        frame_shape = (self.info["width"], self.info["height"], pix_fmt.channels)
-        data_shape = (self.info["width"], self.info["height"], len(self._norms))
+        assert self._norms is not None
+        frame_shape = (self.info["width"], self.info["height"], self._format.channels)
+        data_shape = (len(self._norms), self.info["width"], self.info["height"])
         data = np.empty(data_shape, dtype=self._dtype)
 
         # iterate over entire movie
-        f_input = ffmpeg.input(self.filename, loglevel="warning")
+        f_input = ffmpeg.input(self.filename, loglevel=self.loglevel)
         f_output = f_input.output(
-            "pipe:", format="rawvideo", pix_fmt=pix_fmt.pix_fmt, vframes=8
+            "pipe:", format="rawvideo", pix_fmt=self._format.pix_fmt_data
         )
         proc = f_output.run_async(pipe_stdout=True)
+        f_id = 0
         while True:
-            read_bytes = proc.stdout.read(np.prod(data_shape))
+            f_id += 1
+            read_bytes = proc.stdout.read(np.prod(frame_shape))
             if not read_bytes:
                 break
-            frame = np.frombuffer(read_bytes, pix_fmt.dtype).reshape(frame_shape)
+            frame = np.frombuffer(read_bytes, self._format.dtype).reshape(frame_shape)
 
             for i, norm in enumerate(self._norms):
-                data[..., i] = norm(pix_fmt.data_from_frame(frame[:, :, i]))
+                frame_data = self._format.data_from_frame(frame[:, :, i])
+                data[i, :, :] = norm.inverse(frame_data)
 
             # create the field with the data of the given index
             assert self._field is not None
             field = self._field.copy()
-            field.data = data.reshape(self.info["field_shape"])
+            field.data = data.reshape(field.data.shape)
             yield field
 
     def items(self) -> Iterator[tuple[float, FieldBase]]:
@@ -399,7 +418,7 @@ def items(self) -> Iterator[tuple[float, FieldBase]]:
             t += dt
 
     @fill_in_docstring
-    def tracker(
+    def tracker(  # type: ignore
         self,
         interrupts: ConstantInterrupts | float = 1,
         *,
@@ -439,11 +458,11 @@ def add_to_state(state):
             possible by defining appropriate :func:`add_to_state`
         """
         if np.isscalar(interrupts):
-            interrupts = ConstantInterrupts(interrupts)
+            interrupts = ConstantInterrupts(interrupts)  # type: ignore
         if not isinstance(interrupts, ConstantInterrupts):
             self._logger.warning("`VideoTracker` can only use `ConstantInterrupts`")
-        self.info["dt"] = interrupts.dt
-        self.info["t_start"] = interrupts.t_start
+        self.info["dt"] = interrupts.dt  # type: ignore
+        self.info["t_start"] = interrupts.t_start  # type: ignore
         return StorageTracker(
             storage=self, interrupts=interrupts, transformation=transformation
         )
diff --git a/tests/storage/test_movie_storages.py b/tests/storage/test_movie_storages.py
index b53d7c02..3b238f96 100644
--- a/tests/storage/test_movie_storages.py
+++ b/tests/storage/test_movie_storages.py
@@ -10,12 +10,13 @@
 from pde.tools.misc import module_available
 
 
-@pytest.mark.skipif(not module_available("ffmpeg"), reason="requires `ffmpeg` module")
-def test_movie_storage_simple(tmp_path, rng):
-    """test storing field as movie"""
-    path = tmp_path / f"test_movie_storage_simple.mov"
+@pytest.mark.skipif(not module_available("ffmpeg"), reason="requires `ffmpeg-python`")
+@pytest.mark.parametrize("dim", [1, 2])
+def test_movie_storage_scalar(dim, tmp_path, rng):
+    """test storing scalar field as movie"""
+    path = tmp_path / f"test_movie_storage_scalar.mov"
 
-    grid = pde.UnitGrid([16, 16])
+    grid = pde.UnitGrid([16] * dim)
     field = pde.ScalarField.random_uniform(grid, rng=rng)
     eq = pde.DiffusionPDE()
     writer = MovieStorage(path)
@@ -32,25 +33,69 @@ def test_movie_storage_simple(tmp_path, rng):
     assert len(reader) == 2
     np.testing.assert_allclose(reader.times, [0, 2])
     for i, field in enumerate(reader):
+        assert field.grid == grid
         np.testing.assert_allclose(field.data, storage[i].data, atol=0.1)
 
 
-@pytest.mark.skipif(not module_available("ffmpeg"), reason="requires `ffmpeg` module")
+@pytest.mark.skipif(not module_available("ffmpeg"), reason="requires `ffmpeg-python`")
 @pytest.mark.parametrize("dim", [1, 2])
 @pytest.mark.parametrize("num_fields", [1, 2, 3])
-def test_movie_storage_scalar_fields(dim, num_fields, tmp_path):
-    """test storing field as movie"""
-    path = tmp_path / f"test_movie_storage_scalar_{dim}_{num_fields}.mp4"
+def test_movie_storage_collection(dim, num_fields, tmp_path):
+    """test storing field collection as movie"""
+    path = tmp_path / f"test_movie_storage_collection_{dim}_{num_fields}.mp4"
 
     grid = pde.UnitGrid([8] * dim)
     field = pde.FieldCollection([pde.ScalarField(grid)] * num_fields, copy_fields=True)
     eq = pde.PDE(dict([("a", "1"), ("b", "2"), ("c", "3")][:num_fields]))
-    writer = MovieStorage(path, vmax=[5, 8, 12])
-    eq.solve(field, t_range=3.5, dt=0.1, backend="numpy", tracker=writer.tracker(1))
+    writer = MovieStorage(path, vmax=[5, 10, 15])
+    storage = pde.MemoryStorage()
+    eq.solve(
+        field,
+        t_range=3.5,
+        dt=0.1,
+        backend="numpy",
+        tracker=[storage.tracker(1), writer.tracker(1)],
+    )
 
     reader = MovieStorage(path)
     assert len(reader) == 4
+    assert reader.vmax == [5, 10, 15]
     for t, fields in enumerate(reader):
         assert fields.grid == grid
-        for i, field in enumerate(fields, 1):
-            np.testing.assert_allclose(field.data, i * t, atol=0.1, rtol=0.1)
+        for i, field in enumerate(fields):
+            np.testing.assert_allclose(field.data, (i + 1) * t, atol=0.1, rtol=0.1)
+            # np.testing.assert_allclose(field.data, storage[t][i].data)
+
+
+@pytest.mark.skipif(not module_available("ffmpeg"), reason="requires `ffmpeg-python`")
+@pytest.mark.parametrize("dim", [1, 2])
+def test_movie_storage_vector(dim, tmp_path, rng):
+    """test storing scalar field as movie"""
+    path = tmp_path / f"test_movie_storage_vector.mov"
+
+    grid = pde.UnitGrid([16] * dim)
+    field = pde.VectorField.random_uniform(grid, rng=rng)
+    eq = pde.PDE({"a": "-0.5 * a"})
+    writer = MovieStorage(path)
+    storage = pde.MemoryStorage()
+    eq.solve(
+        field,
+        t_range=3.5,
+        dt=0.1,
+        backend="numpy",
+        tracker=[storage.tracker(2), writer.tracker(2)],
+    )
+    assert len(writer._norms) == dim
+
+    reader = MovieStorage(path)
+    assert len(reader) == 2
+    np.testing.assert_allclose(reader.times, [0, 2])
+    for i, field in enumerate(reader):
+        assert field.grid == grid
+        assert isinstance(field, pde.VectorField)
+        assert field.data.shape == (dim,) + grid.shape
+        np.testing.assert_allclose(field.data, storage[i].data, atol=0.1)
+
+
+# test all pixel formats
+# test failure of complex data

From 3126b11ba1c20835ac70fce272d79cdbe5eb672f Mon Sep 17 00:00:00 2001
From: David Zwicker <david.zwicker@ds.mpg.de>
Date: Sun, 17 Mar 2024 11:41:59 +0100
Subject: [PATCH 05/13] Install FFmpeg as a requirement

---
 .github/workflows/tests_all.yml | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/.github/workflows/tests_all.yml b/.github/workflows/tests_all.yml
index 82161432..40d7a108 100644
--- a/.github/workflows/tests_all.yml
+++ b/.github/workflows/tests_all.yml
@@ -12,11 +12,19 @@ jobs:
 
     steps:
     - uses: actions/checkout@v3
+
     - name: Set up Python
       uses: actions/setup-python@v4
       with:
         python-version: '3.12'
 
+    - uses: FedericoCarboni/setup-ffmpeg@v3
+      # install ffmpeg as special requirement
+      id: setup-ffmpeg
+      with:
+        ffmpeg-version: release
+        github-token: ${{ github.server_url == 'https://github.com' && github.token || '' }}
+
     - name: Install dependencies
       # install all requirements
       run: |

From e12d6adafa56d28f908b100897cf210783019b92 Mon Sep 17 00:00:00 2001
From: David Zwicker <david.zwicker@ds.mpg.de>
Date: Sun, 17 Mar 2024 14:08:35 +0100
Subject: [PATCH 06/13] Fixed color space problem

Minor additional fixed and improved tests
---
 pde/storage/base.py                    |   6 +-
 pde/storage/movie.py                   |  90 ++++++++++++++++----
 tests/storage/test_generic_storages.py | 109 ++++++++++++++-----------
 tests/storage/test_movie_storages.py   |  18 ++--
 4 files changed, 144 insertions(+), 79 deletions(-)

diff --git a/pde/storage/base.py b/pde/storage/base.py
index 6ff28a1c..7f946441 100644
--- a/pde/storage/base.py
+++ b/pde/storage/base.py
@@ -599,9 +599,9 @@ def initialize(self, field: FieldBase, info: InfoDict | None = None) -> float:
         Returns:
             float: The first time the tracker needs to handle data
         """
-        result = super().initialize(field, info)
-        self.storage.start_writing(self._transform(field, 0), info)
-        return result
+        t_first = super().initialize(field, info)
+        self.storage.start_writing(self._transform(field, t_first), info)
+        return t_first
 
     def handle(self, field: FieldBase, t: float) -> None:
         """handle data supplied to this tracker
diff --git a/pde/storage/movie.py b/pde/storage/movie.py
index affd738a..06f29f85 100644
--- a/pde/storage/movie.py
+++ b/pde/storage/movie.py
@@ -48,11 +48,18 @@ def _get_limits(value: float | ArrayLike, dim: int) -> np.ndarray:
 class MovieStorage(StorageBase):
     """store discretized fields in a movie file
 
+    This storage only works when the `ffmpeg` program and :mod:`ffmpeg` is installed.
+
     Warning:
         This storage potentially compresses data and can thus lead to loss of some
         information. The data quality depends on many parameters, but most important are
         the bit depth of the video format, the range that is encoded (determined by
         `vmin` and `vmax`), and the target bitrate.
+
+        Note also that selecting individual time points might be quite slow since the
+        video needs to be read from the beginning each time. Instead, it is much more
+        efficient to process entire videos (by iterating over them or using
+        :func:`~pde.storage.movie.MovieStorage.items()`).
     """
 
     def __init__(
@@ -227,8 +234,13 @@ def start_writing(self, field: FieldBase, info: InfoDict | None = None) -> None:
 
         # delete data if truncation is requested. This is for instance necessary
         # to remove older data with incompatible data_shape
-        if self.write_mode == "truncate" or self.write_mode == "truncate_once":
+        if self.write_mode == "truncate":
             self.clear()
+        elif self.write_mode == "truncate_once":
+            self.clear()
+            self.write_mode = "append"  # do not truncate in subsequent calls
+        elif self.write_mode == "append":
+            raise NotImplementedError("Appending to movies is not possible")
 
         # initialize the writing, setting current data shape
         super().start_writing(field, info=info)
@@ -253,35 +265,39 @@ def start_writing(self, field: FieldBase, info: InfoDict | None = None) -> None:
         else:
             self.info["video_format"] = self.video_format
         self._format = FFmpeg.formats[self.info["video_format"]]
-        self._data_shape = (width, height, self._format.channels)
         if field.is_complex:
             raise NotImplementedError("MovieStorage does not support complex values")
+        self._frame_shape = (width, height, self._format.channels)
 
         # set up the normalization
         self._init_normalization(field, inverse=False)
 
         # set input
         self._logger.debug(f"Start ffmpeg process for `{self.filename}`")
-        args = {
+        input_args = {
             "format": "rawvideo",
             "s": f"{width}x{height}",
             "pix_fmt": self._format.pix_fmt_data,
             "loglevel": self.loglevel,
         }
-        f_input = ffmpeg.input("pipe:", **args)
+        f_input = ffmpeg.input("pipe:", **input_args)
         # set output format
-        args = {
+        output_args = {
             "vcodec": self._format.codec,
-            "crf": "10",  # Constant Rate Factor - lower values for less compression
+            "crf": "0",  # Constant Rate Factor - lower values for less compression
             "pix_fmt": self._format.pix_fmt_file,
             "metadata": "comment=" + shlex.quote(self._get_metadata()),
         }
+        if "264" in self._format.codec:
+            # make the H.264 codec use the full color range
+            output_args["bsf"] = "h264_metadata=video_full_range_flag=1"
         if self.bitrate > 0:
-            args["video_bitrate"] = str(self.bitrate)
-        f_output = f_input.output(filename=self.filename, **args)
-        if self.write_mode in {"truncate", "truncate_once"}:
-            f_output = f_output.overwrite_output()  # allow overwriting file
+            # set the specified bitrate
+            output_args["video_bitrate"] = str(self.bitrate)
+        f_output = f_input.output(filename=self.filename, **output_args)
         self._ffmpeg = f_output.run_async(pipe_stdin=True)  # start process
+
+        self.info["num_frames"] = 0
         self._state = "writing"
 
     def _append_data(self, data: np.ndarray, time: float) -> None:
@@ -304,8 +320,9 @@ def _append_data(self, data: np.ndarray, time: float) -> None:
             raise NotImplementedError
         if grid_dim == 1:
             data = data.reshape(data.shape + (1,))
+        # check spatial dimensions
         assert data.ndim >= 2
-        assert data.shape[-2:] == self._data_shape[:2]  # check spatial dimensions
+        assert data.shape[-2:] == self._frame_shape[:2]
 
         # ensure the data has the shape extra_dim x width x height
         # where `extra_dim` are separated fields or capture the rank of the field
@@ -314,17 +331,18 @@ def _append_data(self, data: np.ndarray, time: float) -> None:
         elif data.ndim > 3:
             # collapse first dimensions, so we have three in total
             data = data.reshape((-1,) + data.shape[-2:])
-        assert len(self._data_shape) == data.ndim == 3
+        assert len(self._frame_shape) == data.ndim == 3
         assert len(self._norms) == data.shape[0]  # same non-spatial dimension
 
         # map data values to frame values
-        frame_data = np.zeros(self._data_shape, dtype=self._format.dtype)
+        frame_data = np.zeros(self._frame_shape, dtype=self._format.dtype)
         for i, norm in enumerate(self._norms):
             data_norm = norm(data[i, ...])
             frame_data[..., i] = self._format.data_to_frame(data_norm)
 
         # write the new data
         self._ffmpeg.stdin.write(frame_data.tobytes())
+        self.info["num_frames"] += 1
 
     def end_writing(self) -> None:
         """finalize the storage after writing"""
@@ -367,7 +385,45 @@ def _get_field(self, t_index: int) -> FieldBase:
             :class:`~pde.fields.FieldBase`:
             The field class containing the grid and data
         """
-        raise NotImplementedError
+        import ffmpeg
+
+        if t_index < 0:
+            t_index += len(self)
+
+        if not 0 <= t_index < len(self):
+            raise IndexError("Time index out of range")
+
+        if "width" not in self.info:
+            self._read_metadata()
+        if self._field is None:
+            self._init_field()
+        assert self._field is not None
+        self._init_normalization(self._field, inverse=True)
+        assert self._norms is not None
+        frame_shape = (self.info["width"], self.info["height"], self._format.channels)
+        data_shape = (len(self._norms), self.info["width"], self.info["height"])
+        data = np.empty(data_shape, dtype=self._dtype)
+
+        # iterate over entire movie
+        f_input = ffmpeg.input(self.filename, loglevel=self.loglevel)
+        f_input = f_input.filter("select", f"gte(n,{t_index})")
+        f_output = f_input.output(
+            "pipe:", vframes=1, format="rawvideo", pix_fmt=self._format.pix_fmt_data
+        )
+        read_bytes, _ = f_output.run(capture_stdout=True)
+        if not read_bytes:
+            raise OSError("Could not read any data")
+        frame = np.frombuffer(read_bytes, self._format.dtype).reshape(frame_shape)
+
+        for i, norm in enumerate(self._norms):
+            frame_data = self._format.data_from_frame(frame[:, :, i])
+            data[i, :, :] = norm.inverse(frame_data)
+
+        # create the field with the data of the given index
+        assert self._field is not None
+        field = self._field.copy()
+        field.data = data.reshape(field.data.shape)
+        return field
 
     def __iter__(self) -> Iterator[FieldBase]:
         """iterate over all stored fields"""
@@ -390,9 +446,7 @@ def __iter__(self) -> Iterator[FieldBase]:
             "pipe:", format="rawvideo", pix_fmt=self._format.pix_fmt_data
         )
         proc = f_output.run_async(pipe_stdout=True)
-        f_id = 0
         while True:
-            f_id += 1
             read_bytes = proc.stdout.read(np.prod(frame_shape))
             if not read_bytes:
                 break
@@ -411,7 +465,9 @@ def __iter__(self) -> Iterator[FieldBase]:
     def items(self) -> Iterator[tuple[float, FieldBase]]:
         """iterate over all times and stored fields, returning pairs"""
         # iterate over entire movie
-        t = self.info.get("t_start", 0)
+        t = self.info.get("t_start")
+        if t is None:
+            t = 0
         dt = self.info.get("dt", 1)
         for field in self:
             yield t, field
diff --git a/tests/storage/test_generic_storages.py b/tests/storage/test_generic_storages.py
index 10b912ea..efa8862f 100644
--- a/tests/storage/test_generic_storages.py
+++ b/tests/storage/test_generic_storages.py
@@ -8,13 +8,14 @@
 import numpy as np
 import pytest
 
-from pde import DiffusionPDE, FileStorage, MemoryStorage, UnitGrid
+from pde import DiffusionPDE, FileStorage, MemoryStorage, MovieStorage, UnitGrid
 from pde.fields import FieldCollection, ScalarField, Tensor2Field, VectorField
 from pde.storage.base import StorageView
 from pde.tools import mpi
 from pde.tools.misc import module_available
 
 STORAGE_CLASSES = [MemoryStorage, FileStorage]
+STORAGE_CLASSES_ALL = [(0, MemoryStorage), (0, FileStorage), (0.1, MovieStorage)]
 
 
 @pytest.fixture
@@ -27,12 +28,19 @@ def storage_factory(tmp_path, storage_class):
         file_path = tmp_path / "test_storage_write.hdf5"
         return functools.partial(FileStorage, file_path)
 
+    if storage_class is MovieStorage:
+        # provide factory that initializes a MovieStorage with a file
+        if not module_available("ffmpeg"):
+            pytest.skip("No module `ffmpeg-python`")
+        file_path = tmp_path / "test_storage_write.mp4"
+        return functools.partial(MovieStorage, file_path, vmax=5)
+
     # simply return the storage class assuming it is a factory function already
     return storage_class
 
 
-@pytest.mark.parametrize("storage_class", STORAGE_CLASSES)
-def test_storage_write(storage_factory):
+@pytest.mark.parametrize("atol,storage_class", STORAGE_CLASSES_ALL)
+def test_storage_write(atol, storage_factory):
     """test simple memory storage"""
     dim = 5
     grid = UnitGrid([dim])
@@ -50,20 +58,22 @@ def test_storage_write(storage_factory):
 
     np.testing.assert_allclose(storage.times, np.arange(2))
     for f in storage:
-        np.testing.assert_array_equal(f.data, np.arange(dim))
+        np.testing.assert_allclose(f.data, np.arange(dim), atol=atol)
     for i in range(2):
-        np.testing.assert_array_equal(storage[i].data, np.arange(dim))
+        np.testing.assert_allclose(storage[i].data, np.arange(dim), atol=atol)
     assert {"a": 1, "b": 2}.items() <= storage.info.items()
 
-    storage = storage_factory()
-    storage.clear()
-    for i in range(3):
-        storage.start_writing(field)
-        field.data = np.arange(dim) + i
-        storage.append(field, i)
-        storage.end_writing()
+    if not isinstance(storage, MovieStorage):
+        storage = storage_factory()
+        storage.clear()
+        for i in range(3):
+            storage.start_writing(field)
+            field.data = np.arange(dim) + i
+            storage.append(field, i)
+            storage.end_writing()
 
-    np.testing.assert_allclose(storage.times, np.arange(3))
+        assert len(storage) == 3
+        np.testing.assert_allclose(storage.times, np.arange(3))
 
 
 def test_storage_truncation(tmp_path, rng):
@@ -100,8 +110,8 @@ def test_storage_truncation(tmp_path, rng):
         assert not storage.has_collection
 
 
-@pytest.mark.parametrize("storage_class", STORAGE_CLASSES)
-def test_storing_extract_range(storage_factory):
+@pytest.mark.parametrize("atol,storage_class", STORAGE_CLASSES_ALL)
+def test_storing_extract_range(atol, storage_factory):
     """test methods specific to FieldCollections in memory storage"""
     sf = ScalarField(UnitGrid([1]))
 
@@ -125,15 +135,16 @@ def test_storing_extract_range(storage_factory):
         s1[-3]
 
     # test extraction
-    s2 = s1.extract_time_range()
-    assert s2.times == list(s1.times)
-    np.testing.assert_allclose(s2.data, s1.data)
-    s3 = s1.extract_time_range(0.5)
-    assert s3.times == s1.times[:1]
-    np.testing.assert_allclose(s3.data, s1.data[:1])
-    s4 = s1.extract_time_range((0.5, 1.5))
-    assert s4.times == s1.times[1:]
-    np.testing.assert_allclose(s4.data, s1.data[1:])
+    if not isinstance(s1, MovieStorage):
+        s2 = s1.extract_time_range()
+        assert s2.times == list(s1.times)
+        np.testing.assert_allclose(s2.data, s1.data)
+        s3 = s1.extract_time_range(0.5)
+        assert s3.times == s1.times[:1]
+        np.testing.assert_allclose(s3.data, s1.data[:1])
+        s4 = s1.extract_time_range((0.5, 1.5))
+        assert s4.times == s1.times[1:]
+        np.testing.assert_allclose(s4.data, s1.data[1:])
 
 
 @pytest.mark.parametrize("storage_class", STORAGE_CLASSES)
@@ -175,8 +186,8 @@ def test_storing_collection(storage_factory, rng):
         storage.view_field("nonsense")
 
 
-@pytest.mark.parametrize("storage_class", STORAGE_CLASSES + [None])
-def test_storage_apply(storage_factory):
+@pytest.mark.parametrize("atol,storage_class", STORAGE_CLASSES_ALL + [(0, None)])
+def test_storage_apply(atol, storage_factory):
     """test the apply function of StorageBase"""
     grid = UnitGrid([2])
     field = ScalarField(grid)
@@ -203,8 +214,8 @@ def test_storage_apply(storage_factory):
     assert len(s2) == 0
 
 
-@pytest.mark.parametrize("storage_class", STORAGE_CLASSES + [None])
-def test_storage_copy(storage_factory):
+@pytest.mark.parametrize("atol,storage_class", STORAGE_CLASSES_ALL + [(0, None)])
+def test_storage_copy(atol, storage_factory):
     """test the copy function of StorageBase"""
     grid = UnitGrid([2])
     field = ScalarField(grid)
@@ -253,8 +264,8 @@ def test_storage_types(storage_factory, dtype, rng):
 
 
 @pytest.mark.multiprocessing
-@pytest.mark.parametrize("storage_class", STORAGE_CLASSES)
-def test_storage_mpi(storage_factory, rng):
+@pytest.mark.parametrize("atol,storage_class", STORAGE_CLASSES_ALL)
+def test_storage_mpi(atol, storage_factory, rng):
     """test writing data using MPI"""
     eq = DiffusionPDE()
     grid = UnitGrid([8])
@@ -270,8 +281,8 @@ def test_storage_mpi(storage_factory, rng):
         assert len(storage) == 11
 
 
-@pytest.mark.parametrize("storage_class", STORAGE_CLASSES)
-def test_storing_transformation_collection(storage_factory, rng):
+@pytest.mark.parametrize("atol,storage_class", STORAGE_CLASSES_ALL)
+def test_storing_transformation_collection(atol, storage_factory, rng):
     """test transformation yielding field collections in storage classes"""
     grid = UnitGrid([8])
     field = ScalarField.random_normal(grid, rng=rng).smooth(1)
@@ -293,14 +304,14 @@ def trans1(field, t):
     assert storage.has_collection
     for t, sol in storage.items():
         a, a2 = sol
-        np.testing.assert_allclose(a2.data, 2 * a.data + t)
+        np.testing.assert_allclose(a2.data, 2 * a.data + t, atol=atol)
 
 
-@pytest.mark.parametrize("storage_class", STORAGE_CLASSES)
-def test_storing_transformation_scalar(storage_factory, rng):
+@pytest.mark.parametrize("atol,storage_class", STORAGE_CLASSES_ALL)
+def test_storing_transformation_scalar(atol, storage_factory, rng):
     """test transformations yielding scalar fields in storage classes"""
     grid = UnitGrid([8])
-    field = ScalarField.random_normal(grid, rng=rng).smooth(1)
+    field = ScalarField.random_uniform(grid, rng=rng).smooth(1)
 
     storage = storage_factory()
     eq = DiffusionPDE(diffusivity=0)
@@ -309,11 +320,14 @@ def test_storing_transformation_scalar(storage_factory, rng):
 
     assert not storage.has_collection
     for sol in storage:
-        np.testing.assert_allclose(sol.data, field.data**2)
+        if isinstance(storage, MovieStorage):
+            np.testing.assert_allclose(sol.data, field.data**2, atol=0.1)
+        else:
+            np.testing.assert_allclose(sol.data, field.data**2)
 
 
-@pytest.mark.parametrize("storage_class", STORAGE_CLASSES)
-def test_storage_view(storage_factory, rng):
+@pytest.mark.parametrize("atol,storage_class", STORAGE_CLASSES_ALL)
+def test_storage_view(atol, storage_factory, rng):
     """test StorageView"""
     grid = UnitGrid([2, 2])
     f1 = ScalarField.random_uniform(grid, 0.1, 0.4, label="a", rng=rng)
@@ -332,13 +346,16 @@ def test_storage_view(storage_factory, rng):
     assert not view.has_collection
     np.testing.assert_allclose(view.times, range(3))
     assert len(view) == 3
-    assert view[0] == f1
+    np.testing.assert_allclose(view[0].data, f1.data, atol=atol)
     for field in view:
-        assert field == f1
+        np.testing.assert_allclose(field.data, f1.data, atol=atol)
     for i, (j, field) in enumerate(view.items()):
         assert i == j
-        assert field == f1
-
-    assert StorageView(storage, field="a")[0] == f1
-    assert StorageView(storage, field="b")[0] == f2
-    assert StorageView(storage, field=1)[0] == f2
+        np.testing.assert_allclose(field.data, f1.data, atol=atol)
+
+    field = StorageView(storage, field="a")[0]
+    np.testing.assert_allclose(field.data, f1.data, atol=atol)
+    field = StorageView(storage, field="b")[0]
+    np.testing.assert_allclose(field.data, f2.data, atol=atol)
+    field = StorageView(storage, field=1)[0]
+    np.testing.assert_allclose(field.data, f2.data, atol=atol)
diff --git a/tests/storage/test_movie_storages.py b/tests/storage/test_movie_storages.py
index 3b238f96..762fe36c 100644
--- a/tests/storage/test_movie_storages.py
+++ b/tests/storage/test_movie_storages.py
@@ -19,7 +19,7 @@ def test_movie_storage_scalar(dim, tmp_path, rng):
     grid = pde.UnitGrid([16] * dim)
     field = pde.ScalarField.random_uniform(grid, rng=rng)
     eq = pde.DiffusionPDE()
-    writer = MovieStorage(path)
+    writer = MovieStorage(path, loglevel="info")
     storage = pde.MemoryStorage()
     eq.solve(
         field,
@@ -34,7 +34,7 @@ def test_movie_storage_scalar(dim, tmp_path, rng):
     np.testing.assert_allclose(reader.times, [0, 2])
     for i, field in enumerate(reader):
         assert field.grid == grid
-        np.testing.assert_allclose(field.data, storage[i].data, atol=0.1)
+        np.testing.assert_allclose(field.data, storage[i].data, atol=0.01)
 
 
 @pytest.mark.skipif(not module_available("ffmpeg"), reason="requires `ffmpeg-python`")
@@ -48,14 +48,7 @@ def test_movie_storage_collection(dim, num_fields, tmp_path):
     field = pde.FieldCollection([pde.ScalarField(grid)] * num_fields, copy_fields=True)
     eq = pde.PDE(dict([("a", "1"), ("b", "2"), ("c", "3")][:num_fields]))
     writer = MovieStorage(path, vmax=[5, 10, 15])
-    storage = pde.MemoryStorage()
-    eq.solve(
-        field,
-        t_range=3.5,
-        dt=0.1,
-        backend="numpy",
-        tracker=[storage.tracker(1), writer.tracker(1)],
-    )
+    eq.solve(field, t_range=3.5, dt=0.1, backend="numpy", tracker=writer.tracker(1))
 
     reader = MovieStorage(path)
     assert len(reader) == 4
@@ -63,8 +56,7 @@ def test_movie_storage_collection(dim, num_fields, tmp_path):
     for t, fields in enumerate(reader):
         assert fields.grid == grid
         for i, field in enumerate(fields):
-            np.testing.assert_allclose(field.data, (i + 1) * t, atol=0.1, rtol=0.1)
-            # np.testing.assert_allclose(field.data, storage[t][i].data)
+            np.testing.assert_allclose(field.data, (i + 1) * t, atol=0.02, rtol=0.02)
 
 
 @pytest.mark.skipif(not module_available("ffmpeg"), reason="requires `ffmpeg-python`")
@@ -94,7 +86,7 @@ def test_movie_storage_vector(dim, tmp_path, rng):
         assert field.grid == grid
         assert isinstance(field, pde.VectorField)
         assert field.data.shape == (dim,) + grid.shape
-        np.testing.assert_allclose(field.data, storage[i].data, atol=0.1)
+        np.testing.assert_allclose(field.data, storage[i].data, atol=0.01)
 
 
 # test all pixel formats

From f5e4eff2a7354928d71ff068d7317c02070e2a03 Mon Sep 17 00:00:00 2001
From: David Zwicker <david.zwicker@ds.mpg.de>
Date: Sun, 17 Mar 2024 14:26:10 +0100
Subject: [PATCH 07/13] Minor improvements

---
 pde/storage/_ffmpeg.py | 10 +++++++++-
 pde/storage/movie.py   |  7 +++----
 2 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/pde/storage/_ffmpeg.py b/pde/storage/_ffmpeg.py
index 54cb51f8..5e27edb5 100644
--- a/pde/storage/_ffmpeg.py
+++ b/pde/storage/_ffmpeg.py
@@ -38,7 +38,15 @@ class FFmpegFormat:
 
     @property
     def pix_fmt_data(self) -> str:
-        return {1: "gray", 3: "rgb24", 4: "rgba"}[self.channels]
+        """return a suitable pixel format for the field data"""
+        if self.channels == 1:
+            return "gray"
+        elif self.channels == 3:
+            return "rgb24"
+        elif self.channels == 4:
+            return "rgba"
+        else:
+            raise NotImplementedError(f"Cannot deal with {self.channels} channels")
 
     def data_to_frame(self, normalized_data: np.ndarray) -> np.ndarray:
         return (normalized_data * self.value_max).astype(self.dtype)
diff --git a/pde/storage/movie.py b/pde/storage/movie.py
index 06f29f85..539f89b0 100644
--- a/pde/storage/movie.py
+++ b/pde/storage/movie.py
@@ -9,11 +9,10 @@
 
 # TODO:
 #     - allow more bits for colorchannels
-#     - allow reading single frames
-#     - support different write_mode
 #     - track whether times roughly work (checking for frame drops)
 #     - we could embedd extra information (like time, and maybe colorscaling) in
-#       the individual frames if we extended the shape
+#       the individual frames if we extended the shape (or we could potentially use
+#       subtitles?)
 
 from __future__ import annotations
 
@@ -483,7 +482,7 @@ def tracker(  # type: ignore
         """create object that can be used as a tracker to fill this storage
 
         Args:
-            interrupts (:class:`~pde.tracker.ConstantInterrupts` or float):
+            interrupts (:class:`~pde.trackers.interrupts.ConstantInterrupts` or float):
                 Time interval with which the tracker is being called
             transformation (callable, optional):
                 A function that transforms the current state into a new field or field

From b3c58ff34c3838370040d134f5e5c7946cb98f22 Mon Sep 17 00:00:00 2001
From: David Zwicker <david.zwicker@ds.mpg.de>
Date: Mon, 18 Mar 2024 12:48:58 +0100
Subject: [PATCH 08/13] Support videos with more bits per channel

Multiple other improvements
---
 pde/storage/_ffmpeg.py                 |  87 ------------------
 pde/storage/movie.py                   |  67 ++++++++------
 pde/tools/ffmpeg.py                    | 122 +++++++++++++++++++++++++
 tests/storage/test_generic_storages.py |   2 +-
 tests/storage/test_movie_storages.py   |  56 ++++++++++--
 tests/tools/test_ffmpeg.py             |  16 ++++
 6 files changed, 228 insertions(+), 122 deletions(-)
 delete mode 100644 pde/storage/_ffmpeg.py
 create mode 100644 pde/tools/ffmpeg.py
 create mode 100644 tests/tools/test_ffmpeg.py

diff --git a/pde/storage/_ffmpeg.py b/pde/storage/_ffmpeg.py
deleted file mode 100644
index 5e27edb5..00000000
--- a/pde/storage/_ffmpeg.py
+++ /dev/null
@@ -1,87 +0,0 @@
-"""
-Functions for interacting with FFmpeg
-
-.. codeauthor:: David Zwicker <david.zwicker@ds.mpg.de>
-"""
-
-# import subprocess as sp
-from dataclasses import dataclass
-
-import numpy as np
-from numpy.typing import DTypeLike
-
-# def _run_ffmpeg(args: list[str]):
-#     return sp.check_output(["ffmpeg"] + args)
-#
-#
-# def codecs() -> list[str]:
-#     """list: all supported ffmpeg codecs"""
-#     res = _run_ffmpeg(["-codecs"])
-#
-#
-# def get_pixel_formats(encoder=None):
-#     if encoder is None:
-#         res = _run_ffmpeg(["-pix_fmts"])
-#     else:
-#         res = _run_ffmpeg(["-h", f"encoder={encoder}"])
-
-
-@dataclass
-class FFmpegFormat:
-    """defines a FFmpeg format used for storing field data in a video"""
-
-    pix_fmt_file: str
-    codec: str
-    channels: int
-    value_max: int
-    dtype: DTypeLike
-
-    @property
-    def pix_fmt_data(self) -> str:
-        """return a suitable pixel format for the field data"""
-        if self.channels == 1:
-            return "gray"
-        elif self.channels == 3:
-            return "rgb24"
-        elif self.channels == 4:
-            return "rgba"
-        else:
-            raise NotImplementedError(f"Cannot deal with {self.channels} channels")
-
-    def data_to_frame(self, normalized_data: np.ndarray) -> np.ndarray:
-        return (normalized_data * self.value_max).astype(self.dtype)
-
-    def data_from_frame(self, frame_data: np.ndarray):
-        return frame_data.astype(float) / self.value_max
-
-
-formats = {
-    "gray": FFmpegFormat(
-        pix_fmt_file="gray",
-        codec="libx264",
-        channels=1,
-        value_max=255,
-        dtype=np.uint8,
-    ),
-    "rgb24": FFmpegFormat(
-        pix_fmt_file="rgb24",
-        codec="libx264rgb",
-        channels=3,
-        value_max=255,
-        dtype=np.uint8,
-    ),
-    "bgr24": FFmpegFormat(
-        pix_fmt_file="bgr24",
-        codec="libx264rgb",
-        channels=3,
-        value_max=255,
-        dtype=np.uint8,
-    ),
-    "rgb32": FFmpegFormat(
-        pix_fmt_file="rgb32",
-        codec="libx264rgb",
-        channels=4,
-        value_max=255,
-        dtype=np.uint8,
-    ),
-}
diff --git a/pde/storage/movie.py b/pde/storage/movie.py
index 539f89b0..fe618e89 100644
--- a/pde/storage/movie.py
+++ b/pde/storage/movie.py
@@ -7,13 +7,9 @@
 .. codeauthor:: David Zwicker <david.zwicker@ds.mpg.de> 
 """
 
-# TODO:
-#     - allow more bits for colorchannels
-#     - track whether times roughly work (checking for frame drops)
-#     - we could embedd extra information (like time, and maybe colorscaling) in
-#       the individual frames if we extended the shape (or we could potentially use
-#       subtitles?)
-
+# TODO: write time as the time stamps (potentially using a factor to convert simulation
+#       time to real time); this might not be possible with rawvideo. An alternative
+#       might be to store the time stamps and apply them later, e.g., using `mkvmerge`
 from __future__ import annotations
 
 import json
@@ -27,12 +23,12 @@
 from matplotlib.colors import Normalize
 from numpy.typing import ArrayLike
 
-from ..fields import FieldCollection, ScalarField
+from ..fields import FieldCollection
 from ..fields.base import DataFieldBase, FieldBase
+from ..tools import ffmpeg as FFmpeg
 from ..tools.docstrings import fill_in_docstring
 from ..tools.misc import module_available
 from ..trackers.interrupts import ConstantInterrupts
-from . import _ffmpeg as FFmpeg
 from .base import InfoDict, StorageBase, StorageTracker, WriteModeType
 
 
@@ -52,8 +48,8 @@ class MovieStorage(StorageBase):
     Warning:
         This storage potentially compresses data and can thus lead to loss of some
         information. The data quality depends on many parameters, but most important are
-        the bit depth of the video format, the range that is encoded (determined by
-        `vmin` and `vmax`), and the target bitrate.
+        the bits per channel of the video format, the range that is encoded (determined
+        by `vmin` and `vmax`), and the target bitrate.
 
         Note also that selecting individual time points might be quite slow since the
         video needs to be read from the beginning each time. Instead, it is much more
@@ -67,6 +63,7 @@ def __init__(
         *,
         vmin: float | ArrayLike = 0,
         vmax: float | ArrayLike = 1,
+        bits_per_channel: int = 8,
         video_format: str = "auto",
         bitrate: int = -1,
         info: InfoDict | None = None,
@@ -77,23 +74,28 @@ def __init__(
         Args:
             filename (str):
                 The path where the movie is stored. The file extension determines the
-                container format of the movie.
+                container format of the movie. The standard codec FFV1 plays well with
+                the ".avi" and ".mkv" container format.
             vmin (float or array):
-                Lowest values that are encoded (per field). Lower values are clipped to
-                this value.
+                Lowest values that are encoded (per field). Smaller values are clipped
+                to this value.
             vmax (float or array):
                 Highest values that are encoded (per field). Larger values are clipped
                 to this value.
+            bits_per_channel (int):
+                The number of bits used per color channel. Typical values are 8 and 16.
+                The relative accuracy of stored values is 0.01 and 0.0001, respectively.
             video_format (str):
-                How to write data to the movie. This determines the number of color
-                channels and the bit depth of individual colors. Available options are
-                listed in :func:`~pde.storage._ffmpeg.formats`. The special value
-                `auto` tries to find a suitable format automatically.
+                Identifier for a video format from :data:`~pde.tools.ffmpeg.formats`,
+                which determines the number of channels, the bit depth of individual
+                colors, and the codec. The special value `auto` tries to find a suitable
+                format automatically, taking `bits_per_channel` into account.
             bitrate (float):
                 The bitrate of the movie (in kilobits per second). The default value of
-                -1 let's the encode choose an appropriate bit rate.
+                -1 let's the encoder choose an appropriate bit rate.
             info (dict):
-                Supplies extra information that is stored in the storage
+                Supplies extra information that is stored in the storage alongside
+                additional information necessary to reconstruct fields and grids.
             write_mode (str):
                 Determines how new data is added to already existing data. Possible
                 values are: 'append' (data is always appended), 'truncate' (data is
@@ -101,7 +103,9 @@ def __init__(
                 (data is cleared for the first writing, but appended subsequently).
                 Alternatively, specifying 'readonly' will disable writing completely.
             loglevel (str):
-                FFmpeg log level
+                FFmpeg log level determining the amount of data sent to stdout. The
+                default only emits warnings and errors, but setting this to `"info"` can
+                be useful to get additioanl information about the encoding.
         """
         if not module_available("ffmpeg"):
             raise ModuleNotFoundError("`MovieStorage` needs `ffmpeg-python` package")
@@ -110,6 +114,7 @@ def __init__(
         self.filename = Path(filename)
         self.vmin = vmin
         self.vmax = vmax
+        self.bits_per_channel = bits_per_channel
         self.video_format = video_format
         self.bitrate = bitrate
         self.loglevel = loglevel
@@ -257,10 +262,14 @@ def start_writing(self, field: FieldBase, info: InfoDict | None = None) -> None:
 
         # get color channel information
         if self.video_format == "auto":
-            if isinstance(field, ScalarField):
-                self.info["video_format"] = "gray"
-            else:
-                self.info["video_format"] = "rgb24"
+            channels = field._data_flat.shape[0]
+            video_format = FFmpeg.find_format(channels, self.bits_per_channel)
+            if video_format is None:
+                raise RuntimeError(
+                    f"Could not find a video format with {channels} channels and "
+                    f"{self.bits_per_channel} bits per channel."
+                )
+            self.info["video_format"] = video_format
         else:
             self.info["video_format"] = self.video_format
         self._format = FFmpeg.formats[self.info["video_format"]]
@@ -283,12 +292,13 @@ def start_writing(self, field: FieldBase, info: InfoDict | None = None) -> None:
         # set output format
         output_args = {
             "vcodec": self._format.codec,
-            "crf": "0",  # Constant Rate Factor - lower values for less compression
             "pix_fmt": self._format.pix_fmt_file,
             "metadata": "comment=" + shlex.quote(self._get_metadata()),
         }
         if "264" in self._format.codec:
-            # make the H.264 codec use the full color range
+            # set extra options for the H.264 codec
+            output_args["crf"] = "0"  # Constant Rate Factor (lower = less compression)
+            # make the H.264 codec use the full color range:
             output_args["bsf"] = "h264_metadata=video_full_range_flag=1"
         if self.bitrate > 0:
             # set the specified bitrate
@@ -438,6 +448,7 @@ def __iter__(self) -> Iterator[FieldBase]:
         frame_shape = (self.info["width"], self.info["height"], self._format.channels)
         data_shape = (len(self._norms), self.info["width"], self.info["height"])
         data = np.empty(data_shape, dtype=self._dtype)
+        frame_bytes = np.prod(frame_shape) * self._format.bytes_per_channel
 
         # iterate over entire movie
         f_input = ffmpeg.input(self.filename, loglevel=self.loglevel)
@@ -446,7 +457,7 @@ def __iter__(self) -> Iterator[FieldBase]:
         )
         proc = f_output.run_async(pipe_stdout=True)
         while True:
-            read_bytes = proc.stdout.read(np.prod(frame_shape))
+            read_bytes = proc.stdout.read(frame_bytes)
             if not read_bytes:
                 break
             frame = np.frombuffer(read_bytes, self._format.dtype).reshape(frame_shape)
diff --git a/pde/tools/ffmpeg.py b/pde/tools/ffmpeg.py
new file mode 100644
index 00000000..7d3478dd
--- /dev/null
+++ b/pde/tools/ffmpeg.py
@@ -0,0 +1,122 @@
+"""
+Functions for interacting with FFmpeg
+
+.. codeauthor:: David Zwicker <david.zwicker@ds.mpg.de>
+"""
+
+from dataclasses import dataclass
+
+# import subprocess as sp
+from typing import Optional
+
+import numpy as np
+from numpy.typing import DTypeLike
+
+# def _run_ffmpeg(args: list[str]):
+#     return sp.check_output(["ffmpeg"] + args)
+#
+#
+# def codecs() -> list[str]:
+#     """list: all supported ffmpeg codecs"""
+#     res = _run_ffmpeg(["-codecs"])
+#
+#
+# def get_pixel_formats(encoder=None):
+#     if encoder is None:
+#         res = _run_ffmpeg(["-pix_fmts"])
+#     else:
+#         res = _run_ffmpeg(["-h", f"encoder={encoder}"])
+
+
+@dataclass
+class FFmpegFormat:
+    """defines a FFmpeg format used for storing field data in a video"""
+
+    pix_fmt_file: str
+    channels: int
+    bits_per_channel: int
+    codec: str = "ffv1"
+
+    @property
+    def pix_fmt_data(self) -> str:
+        """return a suitable pixel format for the field data"""
+        if self.bits_per_channel == 8:
+            if self.channels == 1:
+                return "gray"
+            elif self.channels == 3:
+                return "rgb24"
+            elif self.channels == 4:
+                return "rgba"
+            else:
+                raise NotImplementedError(f"Cannot deal with {self.channels} channels")
+        elif self.bits_per_channel == 16:
+            if self.channels == 1:
+                return "gray16le"
+            elif self.channels == 3:
+                return "gbrp16le"
+            elif self.channels == 4:
+                return "rgba64le"
+            else:
+                raise NotImplementedError(f"Cannot deal with {self.channels} channels")
+        else:
+            raise NotImplementedError(f"Cannot use {self.bits_per_channel} bits")
+
+    @property
+    def bytes_per_channel(self) -> int:
+        return self.bits_per_channel // 8
+
+    @property
+    def dtype(self) -> DTypeLike:
+        if self.bits_per_channel == 8:
+            return np.uint8
+        elif self.bits_per_channel == 16:
+            return np.uint16
+        else:
+            raise NotImplementedError(f"Cannot use {self.bits_per_channel} bits")
+
+    @property
+    def value_max(self) -> int:
+        return 2**self.bits_per_channel - 1  # type: ignore
+
+    def data_to_frame(self, normalized_data: np.ndarray) -> np.ndarray:
+        return (normalized_data * self.value_max).astype(self.dtype)
+
+    def data_from_frame(self, frame_data: np.ndarray):
+        return frame_data.astype(float) / self.value_max
+
+
+formats = {
+    "gray": FFmpegFormat(pix_fmt_file="gray", channels=1, bits_per_channel=8),
+    "rgb24": FFmpegFormat(pix_fmt_file="rgb24", channels=3, bits_per_channel=8),
+    "rgb32": FFmpegFormat(pix_fmt_file="rgb32", channels=4, bits_per_channel=8),
+    "gray16le": FFmpegFormat(pix_fmt_file="gray16le", channels=1, bits_per_channel=16),
+    "gbrp16le": FFmpegFormat(pix_fmt_file="gbrp16le", channels=3, bits_per_channel=16),
+    "rgba64le": FFmpegFormat(pix_fmt_file="rgba64le", channels=4, bits_per_channel=16),
+}
+
+
+def find_format(channels: int, bits_per_channel: int = 8) -> Optional[str]:
+    """find a defined FFmpegFormat that satisifies the requirements
+
+    Args:
+        channels (int):
+            Minimal number of color channels
+        bits_per_channel (int):
+            Minimal number of bits per channel
+
+    Returns:
+        str: Identifier for a format that satisifies the requirements (but might have
+        more channels or more bits per channel then requested. `None` is returned if no
+        format can be identified.
+    """
+    n_best, f_best = None, None
+    for n, f in formats.items():  # iterate through all defined formats
+        if f.channels >= channels and f.bits_per_channel >= bits_per_channel:
+            # this format satisfies the requirements
+            if (
+                f_best is None
+                or f.bits_per_channel < f_best.bits_per_channel
+                or f.channels < f_best.channels
+            ):
+                n_best, f_best = n, f
+    return n_best
diff --git a/tests/storage/test_generic_storages.py b/tests/storage/test_generic_storages.py
index efa8862f..8b4f1a0c 100644
--- a/tests/storage/test_generic_storages.py
+++ b/tests/storage/test_generic_storages.py
@@ -32,7 +32,7 @@ def storage_factory(tmp_path, storage_class):
         # provide factory that initializes a MovieStorage with a file
         if not module_available("ffmpeg"):
             pytest.skip("No module `ffmpeg-python`")
-        file_path = tmp_path / "test_storage_write.mp4"
+        file_path = tmp_path / "test_storage_write.avi"
         return functools.partial(MovieStorage, file_path, vmax=5)
 
     # simply return the storage class assuming it is a factory function already
diff --git a/tests/storage/test_movie_storages.py b/tests/storage/test_movie_storages.py
index 762fe36c..936a4a82 100644
--- a/tests/storage/test_movie_storages.py
+++ b/tests/storage/test_movie_storages.py
@@ -7,6 +7,7 @@
 
 import pde
 from pde import MovieStorage
+from pde.tools.ffmpeg import formats
 from pde.tools.misc import module_available
 
 
@@ -14,12 +15,12 @@
 @pytest.mark.parametrize("dim", [1, 2])
 def test_movie_storage_scalar(dim, tmp_path, rng):
     """test storing scalar field as movie"""
-    path = tmp_path / f"test_movie_storage_scalar.mov"
+    path = tmp_path / f"test_movie_storage_scalar.avi"
 
     grid = pde.UnitGrid([16] * dim)
     field = pde.ScalarField.random_uniform(grid, rng=rng)
     eq = pde.DiffusionPDE()
-    writer = MovieStorage(path, loglevel="info")
+    writer = MovieStorage(path)
     storage = pde.MemoryStorage()
     eq.solve(
         field,
@@ -42,7 +43,7 @@ def test_movie_storage_scalar(dim, tmp_path, rng):
 @pytest.mark.parametrize("num_fields", [1, 2, 3])
 def test_movie_storage_collection(dim, num_fields, tmp_path):
     """test storing field collection as movie"""
-    path = tmp_path / f"test_movie_storage_collection_{dim}_{num_fields}.mp4"
+    path = tmp_path / f"test_movie_storage_collection_{dim}_{num_fields}.avi"
 
     grid = pde.UnitGrid([8] * dim)
     field = pde.FieldCollection([pde.ScalarField(grid)] * num_fields, copy_fields=True)
@@ -63,7 +64,7 @@ def test_movie_storage_collection(dim, num_fields, tmp_path):
 @pytest.mark.parametrize("dim", [1, 2])
 def test_movie_storage_vector(dim, tmp_path, rng):
     """test storing scalar field as movie"""
-    path = tmp_path / f"test_movie_storage_vector.mov"
+    path = tmp_path / f"test_movie_storage_vector.avi"
 
     grid = pde.UnitGrid([16] * dim)
     field = pde.VectorField.random_uniform(grid, rng=rng)
@@ -89,5 +90,48 @@ def test_movie_storage_vector(dim, tmp_path, rng):
         np.testing.assert_allclose(field.data, storage[i].data, atol=0.01)
 
 
-# test all pixel formats
-# test failure of complex data
+@pytest.mark.parametrize("name,video_format", formats.items())
+def test_video_format(name, video_format, tmp_path, rng):
+    """test all video_formats"""
+    path = tmp_path / f"test_movie_storage_{name}.avi"
+
+    field = pde.ScalarField.random_uniform(pde.UnitGrid([16]), rng=rng)
+    writer = MovieStorage(path, video_format=name)
+    storage = pde.MemoryStorage()
+    pde.DiffusionPDE().solve(
+        field,
+        t_range=3.5,
+        dt=0.1,
+        backend="numpy",
+        tracker=[storage.tracker(2), writer.tracker(2)],
+    )
+
+    atol = 0.01 if video_format.bits_per_channel == 8 else 0.0001
+    reader = MovieStorage(path)
+    assert len(reader) == 2
+    np.testing.assert_allclose(reader.times, [0, 2])
+    for i, field in enumerate(reader):
+        np.testing.assert_allclose(field.data, storage[i].data, atol=atol)
+
+
+def test_too_many_channels(tmp_path, rng):
+    """test that data with too many channels throws an error"""
+    path = tmp_path / f"test_movie_complex.avi"
+
+    field = pde.FieldCollection.scalar_random_uniform(5, pde.UnitGrid([16]), rng=rng)
+    writer = MovieStorage(path)
+    eq = pde.PDE({s: "0" for s in "abcde"})
+    with pytest.raises(RuntimeError):
+        eq.solve(field, t_range=3.5, backend="numpy", tracker=writer.tracker(2))
+
+
+def test_complex_data(tmp_path, rng):
+    """test that complex data throws an error"""
+    path = tmp_path / f"test_movie_complex.avi"
+
+    field = pde.ScalarField.random_uniform(pde.UnitGrid([16]), dtype=complex, rng=rng)
+    writer = MovieStorage(path)
+    with pytest.raises(NotImplementedError):
+        pde.DiffusionPDE().solve(
+            field, t_range=3.5, backend="numpy", tracker=writer.tracker(2)
+        )
diff --git a/tests/tools/test_ffmpeg.py b/tests/tools/test_ffmpeg.py
new file mode 100644
index 00000000..4077ce37
--- /dev/null
+++ b/tests/tools/test_ffmpeg.py
@@ -0,0 +1,16 @@
+"""
+.. codeauthor:: David Zwicker <david.zwicker@ds.mpg.de>
+"""
+
+import pytest
+
+from pde.tools.ffmpeg import find_format
+
+
+@pytest.mark.parametrize(
+    "channels,bits_per_channel,result",
+    [(1, 8, "gray"), (2, 7, "rgb24"), (3, 9, "gbrp16le"), (5, 8, None), (1, 17, None)],
+)
+def test_find_format(channels, bits_per_channel, result):
+    """test_find_format function"""
+    assert find_format(channels, bits_per_channel) == result

From 9f18ccb11b2b1c80fe3185fa390857f7b8f94394 Mon Sep 17 00:00:00 2001
From: David Zwicker <david.zwicker@ds.mpg.de>
Date: Mon, 18 Mar 2024 14:05:35 +0100
Subject: [PATCH 09/13] Multiple smaller improvements

---
 pde/tools/__init__.py                |  1 +
 pde/tools/ffmpeg.py                  | 43 ++++++++++++++++------------
 tests/storage/test_movie_storages.py |  3 ++
 3 files changed, 29 insertions(+), 18 deletions(-)

diff --git a/pde/tools/__init__.py b/pde/tools/__init__.py
index 8bc758ca..33e08a37 100644
--- a/pde/tools/__init__.py
+++ b/pde/tools/__init__.py
@@ -9,6 +9,7 @@
    cuboid
    docstrings
    expressions
+   ffmpeg
    math
    misc
    mpi
diff --git a/pde/tools/ffmpeg.py b/pde/tools/ffmpeg.py
index 7d3478dd..4b83b694 100644
--- a/pde/tools/ffmpeg.py
+++ b/pde/tools/ffmpeg.py
@@ -1,41 +1,42 @@
 """
 Functions for interacting with FFmpeg
 
+.. autosummary::
+   :nosignatures:
+
+   FFmpegFormat
+   formats
+   find_format
+
 .. codeauthor:: David Zwicker <david.zwicker@ds.mpg.de>
 """
 
 from dataclasses import dataclass
-
-# import subprocess as sp
 from typing import Optional
 
 import numpy as np
 from numpy.typing import DTypeLike
 
-# def _run_ffmpeg(args: list[str]):
-#     return sp.check_output(["ffmpeg"] + args)
-#
-#
-# def codecs() -> list[str]:
-#     """list: all supported ffmpeg codecs"""
-#     res = _run_ffmpeg(["-codecs"])
-#
-#
-# def get_pixel_formats(encoder=None):
-#     if encoder is None:
-#         res = _run_ffmpeg(["-pix_fmts"])
-#     else:
-#         res = _run_ffmpeg(["-h", f"encoder={encoder}"])
-
 
 @dataclass
 class FFmpegFormat:
-    """defines a FFmpeg format used for storing field data in a video"""
+    """defines a FFmpeg format used for storing field data in a video
+
+    All pixel formats supported by FFmpeg can be obtained by running
+    :code:`ffmpeg -pix_fmts`. However, not all pixel formats are supported by all
+    codecs. Supported pixel formats are listed in the output of
+    :code:`ffmpeg -h encoder=<ENCODER>`, where `<ENCODER>` is one of the encoders listed
+    in :code:`ffmpeg -codecs`.
+    """
 
     pix_fmt_file: str
+    """str: name of the pixel format used in the codec"""
     channels: int
+    """int: number of color channels in this pixel format"""
     bits_per_channel: int
+    """int: number of bits per color channel in this pixel format"""
     codec: str = "ffv1"
+    """str: name of the codec that supports this pixel format"""
 
     @property
     def pix_fmt_data(self) -> str:
@@ -63,10 +64,12 @@ def pix_fmt_data(self) -> str:
 
     @property
     def bytes_per_channel(self) -> int:
+        """int:number of bytes per color channel"""
         return self.bits_per_channel // 8
 
     @property
     def dtype(self) -> DTypeLike:
+        """numpy dtype corresponding to the data of a single channel"""
         if self.bits_per_channel == 8:
             return np.uint8
         elif self.bits_per_channel == 16:
@@ -76,12 +79,15 @@ def dtype(self) -> DTypeLike:
 
     @property
     def value_max(self) -> int:
+        """maximal value stored in a color channel"""
         return 2**self.bits_per_channel - 1  # type: ignore
 
     def data_to_frame(self, normalized_data: np.ndarray) -> np.ndarray:
+        """converts normalized data to data being stored in a color channel"""
         return (normalized_data * self.value_max).astype(self.dtype)
 
     def data_from_frame(self, frame_data: np.ndarray):
+        """converts data stored in a color channel to normalized data"""
         return frame_data.astype(float) / self.value_max
 
 
@@ -93,6 +99,7 @@ def data_from_frame(self, frame_data: np.ndarray):
     "gbrp16le": FFmpegFormat(pix_fmt_file="gbrp16le", channels=3, bits_per_channel=16),
     "rgba64le": FFmpegFormat(pix_fmt_file="rgba64le", channels=4, bits_per_channel=16),
 }
+"""dict of :class:`FFmpegFormat` formats"""
 
 
 def find_format(channels: int, bits_per_channel: int = 8) -> Optional[str]:
diff --git a/tests/storage/test_movie_storages.py b/tests/storage/test_movie_storages.py
index 936a4a82..1574bc7d 100644
--- a/tests/storage/test_movie_storages.py
+++ b/tests/storage/test_movie_storages.py
@@ -90,6 +90,7 @@ def test_movie_storage_vector(dim, tmp_path, rng):
         np.testing.assert_allclose(field.data, storage[i].data, atol=0.01)
 
 
+@pytest.mark.skipif(not module_available("ffmpeg"), reason="requires `ffmpeg-python`")
 @pytest.mark.parametrize("name,video_format", formats.items())
 def test_video_format(name, video_format, tmp_path, rng):
     """test all video_formats"""
@@ -114,6 +115,7 @@ def test_video_format(name, video_format, tmp_path, rng):
         np.testing.assert_allclose(field.data, storage[i].data, atol=atol)
 
 
+@pytest.mark.skipif(not module_available("ffmpeg"), reason="requires `ffmpeg-python`")
 def test_too_many_channels(tmp_path, rng):
     """test that data with too many channels throws an error"""
     path = tmp_path / f"test_movie_complex.avi"
@@ -125,6 +127,7 @@ def test_too_many_channels(tmp_path, rng):
         eq.solve(field, t_range=3.5, backend="numpy", tracker=writer.tracker(2))
 
 
+@pytest.mark.skipif(not module_available("ffmpeg"), reason="requires `ffmpeg-python`")
 def test_complex_data(tmp_path, rng):
     """test that complex data throws an error"""
     path = tmp_path / f"test_movie_complex.avi"

From 7a9d2095edc4b407aca577b5b0541c132048fcbb Mon Sep 17 00:00:00 2001
From: David Zwicker <david.zwicker@ds.mpg.de>
Date: Mon, 18 Mar 2024 16:13:22 +0100
Subject: [PATCH 10/13] Improved ffmpeg format definition

---
 pde/storage/movie.py                 |  30 +++++--
 pde/tools/ffmpeg.py                  | 125 +++++++++++++++++----------
 tests/storage/test_movie_storages.py |   8 +-
 3 files changed, 107 insertions(+), 56 deletions(-)

diff --git a/pde/storage/movie.py b/pde/storage/movie.py
index fe618e89..0b359f17 100644
--- a/pde/storage/movie.py
+++ b/pde/storage/movie.py
@@ -75,7 +75,7 @@ def __init__(
             filename (str):
                 The path where the movie is stored. The file extension determines the
                 container format of the movie. The standard codec FFV1 plays well with
-                the ".avi" and ".mkv" container format.
+                the ".avi" container format.
             vmin (float or array):
                 Lowest values that are encoded (per field). Smaller values are clipped
                 to this value.
@@ -181,14 +181,26 @@ def _read_metadata(self) -> None:
         self.vmax = metadata.pop("vmax", 1)
         self.info.update(metadata)
 
+        # read information from the first stream
         stream = info["streams"][0]
-        self.info["num_frames"] = int(stream["nb_frames"])
-        self.info["width"] = stream["coded_width"]
-        self.info["height"] = stream["coded_height"]
         try:
-            self._format = FFmpeg.formats[self.info["video_format"]]
+            self.info["num_frames"] = int(stream["nb_frames"])
         except KeyError:
-            self._logger.warning(f"Unknown pixel format `{self.info['pixel_format']}`")
+            self.info["num_frames"] = None  # number of frames was not stored
+        self.info["width"] = stream["width"]
+        self.info["height"] = stream["height"]
+        video_format = self.info.get("video_format")
+        if video_format is None:
+            video_format = stream.get("pix_fmt")
+        if video_format is None:
+            if self.video_format == "auto":
+                raise RuntimeError("Could not determine video format from file")
+            else:
+                video_format = self.video_format
+        try:
+            self._format = FFmpeg.formats[video_format]
+        except KeyError:
+            self._logger.warning(f"Unknown pixel format `{video_format}`")
         else:
             if self._format.pix_fmt_file != stream["pix_fmt"]:
                 self._logger.info(
@@ -285,7 +297,7 @@ def start_writing(self, field: FieldBase, info: InfoDict | None = None) -> None:
         input_args = {
             "format": "rawvideo",
             "s": f"{width}x{height}",
-            "pix_fmt": self._format.pix_fmt_data,
+            "pixel_format": self._format.pix_fmt_data,
             "loglevel": self.loglevel,
         }
         f_input = ffmpeg.input("pipe:", **input_args)
@@ -321,7 +333,9 @@ def _append_data(self, data: np.ndarray, time: float) -> None:
                 "Writing not initialized. Call "
                 f"`{self.__class__.__name__}.start_writing`"
             )
-        assert self._norms is not None  # normalization has been initialized
+        # normalization and video format have been initialized
+        assert self._norms is not None
+        assert self._format is not None
 
         # make sure there are two spatial dimensions
         grid_dim = self._grid.num_axes
diff --git a/pde/tools/ffmpeg.py b/pde/tools/ffmpeg.py
index 4b83b694..1807ab70 100644
--- a/pde/tools/ffmpeg.py
+++ b/pde/tools/ffmpeg.py
@@ -12,7 +12,7 @@
 """
 
 from dataclasses import dataclass
-from typing import Optional
+from typing import Optional, Union
 
 import numpy as np
 from numpy.typing import DTypeLike
@@ -31,73 +31,108 @@ class FFmpegFormat:
 
     pix_fmt_file: str
     """str: name of the pixel format used in the codec"""
+    pix_fmt_data: str
+    """str: name of the pixel format used in the frame data"""
     channels: int
     """int: number of color channels in this pixel format"""
     bits_per_channel: int
     """int: number of bits per color channel in this pixel format"""
+    dtype: DTypeLike
+    """numpy dtype corresponding to the data of a single channel"""
     codec: str = "ffv1"
     """str: name of the codec that supports this pixel format"""
 
-    @property
-    def pix_fmt_data(self) -> str:
-        """return a suitable pixel format for the field data"""
-        if self.bits_per_channel == 8:
-            if self.channels == 1:
-                return "gray"
-            elif self.channels == 3:
-                return "rgb24"
-            elif self.channels == 4:
-                return "rgba"
-            else:
-                raise NotImplementedError(f"Cannot deal with {self.channels} channels")
-        elif self.bits_per_channel == 16:
-            if self.channels == 1:
-                return "gray16le"
-            elif self.channels == 3:
-                return "gbrp16le"
-            elif self.channels == 4:
-                return "rgba64le"
-            else:
-                raise NotImplementedError(f"Cannot deal with {self.channels} channels")
-        else:
-            raise NotImplementedError(f"Cannot use {self.bits_per_channel} bits")
-
     @property
     def bytes_per_channel(self) -> int:
         """int:number of bytes per color channel"""
         return self.bits_per_channel // 8
 
     @property
-    def dtype(self) -> DTypeLike:
-        """numpy dtype corresponding to the data of a single channel"""
-        if self.bits_per_channel == 8:
-            return np.uint8
-        elif self.bits_per_channel == 16:
-            return np.uint16
-        else:
-            raise NotImplementedError(f"Cannot use {self.bits_per_channel} bits")
-
-    @property
-    def value_max(self) -> int:
+    def max_value(self) -> Union[float, int]:
         """maximal value stored in a color channel"""
-        return 2**self.bits_per_channel - 1  # type: ignore
+        if np.issubdtype(self.dtype, np.integer):
+            return 2**self.bits_per_channel - 1  # type: ignore
+        else:
+            return 1.0
 
     def data_to_frame(self, normalized_data: np.ndarray) -> np.ndarray:
         """converts normalized data to data being stored in a color channel"""
-        return (normalized_data * self.value_max).astype(self.dtype)
+        return np.ascontiguousarray(normalized_data * self.max_value, dtype=self.dtype)
 
     def data_from_frame(self, frame_data: np.ndarray):
         """converts data stored in a color channel to normalized data"""
-        return frame_data.astype(float) / self.value_max
+        return frame_data.astype(float) / self.max_value
 
 
 formats = {
-    "gray": FFmpegFormat(pix_fmt_file="gray", channels=1, bits_per_channel=8),
-    "rgb24": FFmpegFormat(pix_fmt_file="rgb24", channels=3, bits_per_channel=8),
-    "rgb32": FFmpegFormat(pix_fmt_file="rgb32", channels=4, bits_per_channel=8),
-    "gray16le": FFmpegFormat(pix_fmt_file="gray16le", channels=1, bits_per_channel=16),
-    "gbrp16le": FFmpegFormat(pix_fmt_file="gbrp16le", channels=3, bits_per_channel=16),
-    "rgba64le": FFmpegFormat(pix_fmt_file="rgba64le", channels=4, bits_per_channel=16),
+    # 8 bit formats
+    "gray": FFmpegFormat(
+        pix_fmt_file="gray",
+        pix_fmt_data="gray",
+        channels=1,
+        bits_per_channel=8,
+        dtype=np.uint8,
+    ),
+    "rgb24": FFmpegFormat(
+        pix_fmt_file="rgb24",
+        pix_fmt_data="rgb24",
+        channels=3,
+        bits_per_channel=8,
+        dtype=np.uint8,
+    ),
+    "rgb32": FFmpegFormat(
+        pix_fmt_file="rgb32",
+        pix_fmt_data="rgb32",
+        channels=4,
+        bits_per_channel=8,
+        dtype=np.uint8,
+    ),
+    # 16 bit formats
+    "gray16le": FFmpegFormat(
+        pix_fmt_file="gray16le",
+        pix_fmt_data="gray16le",
+        channels=1,
+        bits_per_channel=16,
+        dtype=np.dtype("<u2"),
+    ),
+    "gbrp16le": FFmpegFormat(
+        pix_fmt_file="gbrp16le",
+        pix_fmt_data="gbrp16le",
+        channels=3,
+        bits_per_channel=16,
+        dtype=np.dtype("<u2"),
+    ),
+    "gbrap16le": FFmpegFormat(
+        pix_fmt_file="gbrap16le",
+        pix_fmt_data="gbrap16le",
+        channels=4,
+        bits_per_channel=16,
+        dtype=np.dtype("<u2"),
+    ),
+    # 32 bit formats
+    # "grayf32le": FFmpegFormat(
+    #     pix_fmt_file="grayf32le",
+    #     pix_fmt_data="grayf32le",
+    #     codec="pfm",  # pfm or dpx
+    #     channels=1,
+    #     bits_per_channel=32,
+    #     dtype=np.dtype("<f4"),
+    # ),
+    # "gbrpf32le": FFmpegFormat(
+    #     pix_fmt_file="gbrpf32le",
+    #     pix_fmt_data="gbrpf32le",
+    #     codec="pfm",  # or dpx?
+    #     channels=3,
+    #     bits_per_channel=32,
+    #     dtype=np.dtype("<f4"),
+    # ),
+    # "gbrapf32le": FFmpegFormat(
+    #     pix_fmt_file="gbrapf32le",
+    #     pix_fmt_data="gbrapf32le",
+    #     channels=4,
+    #     bits_per_channel=32,
+    #     dtype=np.dtype("<f4"),
+    # ),
 }
 """dict of :class:`FFmpegFormat` formats"""
 
diff --git a/tests/storage/test_movie_storages.py b/tests/storage/test_movie_storages.py
index 1574bc7d..4df24d43 100644
--- a/tests/storage/test_movie_storages.py
+++ b/tests/storage/test_movie_storages.py
@@ -94,6 +94,10 @@ def test_movie_storage_vector(dim, tmp_path, rng):
 @pytest.mark.parametrize("name,video_format", formats.items())
 def test_video_format(name, video_format, tmp_path, rng):
     """test all video_formats"""
+    if np.issubdtype(video_format.dtype, np.integer):
+        assert video_format.max_value == np.iinfo(video_format.dtype).max
+    assert np.dtype(video_format.dtype).itemsize == video_format.bytes_per_channel
+
     path = tmp_path / f"test_movie_storage_{name}.avi"
 
     field = pde.ScalarField.random_uniform(pde.UnitGrid([16]), rng=rng)
@@ -107,10 +111,8 @@ def test_video_format(name, video_format, tmp_path, rng):
         tracker=[storage.tracker(2), writer.tracker(2)],
     )
 
-    atol = 0.01 if video_format.bits_per_channel == 8 else 0.0001
+    atol = {8: 1e-2, 16: 1e-4, 32: 1e-8, 64: 1e-15}[video_format.bits_per_channel]
     reader = MovieStorage(path)
-    assert len(reader) == 2
-    np.testing.assert_allclose(reader.times, [0, 2])
     for i, field in enumerate(reader):
         np.testing.assert_allclose(field.data, storage[i].data, atol=atol)
 

From b4e942a2aef5a754eec7bc7b19f0619b58205171 Mon Sep 17 00:00:00 2001
From: David Zwicker <david.zwicker@ds.mpg.de>
Date: Tue, 19 Mar 2024 09:16:24 +0100
Subject: [PATCH 11/13] Minor improvements

---
 pde/storage/file.py  |  2 --
 pde/storage/movie.py | 31 +++++++++++++------------------
 2 files changed, 13 insertions(+), 20 deletions(-)

diff --git a/pde/storage/file.py b/pde/storage/file.py
index 82f2358d..a2fd7dc4 100644
--- a/pde/storage/file.py
+++ b/pde/storage/file.py
@@ -7,7 +7,6 @@
 from __future__ import annotations
 
 import json
-import logging
 from pathlib import Path
 from typing import Any, Literal
 
@@ -69,7 +68,6 @@ def __init__(
         self.keep_opened = keep_opened
         self.check_mpi = check_mpi
 
-        self._logger = logging.getLogger(self.__class__.__name__)
         self._file: Any = None
         self._is_writing = False
         self._data_length: int = None  # type: ignore
diff --git a/pde/storage/movie.py b/pde/storage/movie.py
index 0b359f17..f3cdee2d 100644
--- a/pde/storage/movie.py
+++ b/pde/storage/movie.py
@@ -7,13 +7,9 @@
 .. codeauthor:: David Zwicker <david.zwicker@ds.mpg.de> 
 """
 
-# TODO: write time as the time stamps (potentially using a factor to convert simulation
-#       time to real time); this might not be possible with rawvideo. An alternative
-#       might be to store the time stamps and apply them later, e.g., using `mkvmerge`
 from __future__ import annotations
 
 import json
-import logging
 import shlex
 from collections.abc import Iterator, Sequence
 from pathlib import Path
@@ -119,7 +115,6 @@ def __init__(
         self.bitrate = bitrate
         self.loglevel = loglevel
 
-        self._logger = logging.getLogger(self.__class__.__name__)
         self._ffmpeg: Any = None
         self._state: Literal["closed", "reading", "writing"] = "closed"
         self._norms: list[Normalize] | None = None
@@ -189,26 +184,26 @@ def _read_metadata(self) -> None:
             self.info["num_frames"] = None  # number of frames was not stored
         self.info["width"] = stream["width"]
         self.info["height"] = stream["height"]
-        video_format = self.info.get("video_format")
-        if video_format is None:
-            video_format = stream.get("pix_fmt")
-        if video_format is None:
-            if self.video_format == "auto":
+        if self.video_format == "auto":
+            video_format = self.info.get("video_format")
+            if video_format is None:
+                video_format = stream.get("pix_fmt")
+            if video_format is None:
                 raise RuntimeError("Could not determine video format from file")
-            else:
-                video_format = self.video_format
+        else:
+            video_format = self.video_format
         try:
             self._format = FFmpeg.formats[video_format]
         except KeyError:
             self._logger.warning(f"Unknown pixel format `{video_format}`")
         else:
-            if self._format.pix_fmt_file != stream["pix_fmt"]:
+            if self._format.pix_fmt_file != stream.get("pix_fmt"):
                 self._logger.info(
                     "Pixel format differs from requested one: "
-                    f"{self._format.pix_fmt_file} != {stream['pix_fmt']}"
+                    f"{self._format.pix_fmt_file} != {stream.get('pix_fmt')}"
                 )
 
-    def _init_normalization(self, field: FieldBase, *, inverse: bool = False) -> None:
+    def _init_normalization(self, field: FieldBase) -> None:
         """initialize the normalizations of the color information
 
         Args:
@@ -290,7 +285,7 @@ def start_writing(self, field: FieldBase, info: InfoDict | None = None) -> None:
         self._frame_shape = (width, height, self._format.channels)
 
         # set up the normalization
-        self._init_normalization(field, inverse=False)
+        self._init_normalization(field)
 
         # set input
         self._logger.debug(f"Start ffmpeg process for `{self.filename}`")
@@ -421,7 +416,7 @@ def _get_field(self, t_index: int) -> FieldBase:
         if self._field is None:
             self._init_field()
         assert self._field is not None
-        self._init_normalization(self._field, inverse=True)
+        self._init_normalization(self._field)
         assert self._norms is not None
         frame_shape = (self.info["width"], self.info["height"], self._format.channels)
         data_shape = (len(self._norms), self.info["width"], self.info["height"])
@@ -457,7 +452,7 @@ def __iter__(self) -> Iterator[FieldBase]:
         if self._field is None:
             self._init_field()
         assert self._field is not None
-        self._init_normalization(self._field, inverse=True)
+        self._init_normalization(self._field)
         assert self._norms is not None
         frame_shape = (self.info["width"], self.info["height"], self._format.channels)
         data_shape = (len(self._norms), self.info["width"], self.info["height"])

From d755511868a7b17e3bc22f7e04fc5866ef86dba2 Mon Sep 17 00:00:00 2001
From: David Zwicker <david.zwicker@ds.mpg.de>
Date: Tue, 19 Mar 2024 14:46:52 +0100
Subject: [PATCH 12/13] Improved support for mkv container format

---
 pde/storage/movie.py                 | 38 ++++++++++++++++++++++++----
 pde/tools/ffmpeg.py                  | 14 +++++-----
 tests/storage/test_movie_storages.py | 27 ++++++++++++++++++++
 3 files changed, 68 insertions(+), 11 deletions(-)

diff --git a/pde/storage/movie.py b/pde/storage/movie.py
index f3cdee2d..0edaad07 100644
--- a/pde/storage/movie.py
+++ b/pde/storage/movie.py
@@ -12,6 +12,7 @@
 import json
 import shlex
 from collections.abc import Iterator, Sequence
+from fractions import Fraction
 from pathlib import Path
 from typing import Any, Callable, Literal
 
@@ -24,6 +25,7 @@
 from ..tools import ffmpeg as FFmpeg
 from ..tools.docstrings import fill_in_docstring
 from ..tools.misc import module_available
+from ..tools.parse_duration import parse_duration
 from ..trackers.interrupts import ConstantInterrupts
 from .base import InfoDict, StorageBase, StorageTracker, WriteModeType
 
@@ -40,12 +42,15 @@ class MovieStorage(StorageBase):
     """store discretized fields in a movie file
 
     This storage only works when the `ffmpeg` program and :mod:`ffmpeg` is installed.
+    The default codec is `FFV1 <https://en.m.wikipedia.org/wiki/FFV1>`_, which supports
+    lossless compression for various configurations. Not all video players support this
+    codec, but `VLC <https://www.videolan.org>`_ usually works quite well.
 
     Warning:
         This storage potentially compresses data and can thus lead to loss of some
         information. The data quality depends on many parameters, but most important are
-        the bits per channel of the video format, the range that is encoded (determined
-        by `vmin` and `vmax`), and the target bitrate.
+        the bits per channel of the video format and the range that is encoded
+        (determined by `vmin` and `vmax`).
 
         Note also that selecting individual time points might be quite slow since the
         video needs to be read from the beginning each time. Instead, it is much more
@@ -71,7 +76,7 @@ def __init__(
             filename (str):
                 The path where the movie is stored. The file extension determines the
                 container format of the movie. The standard codec FFV1 plays well with
-                the ".avi" container format.
+                the ".avi", ".mkv", and ".mov" container format.
             vmin (float or array):
                 Lowest values that are encoded (per field). Smaller values are clipped
                 to this value.
@@ -166,7 +171,11 @@ def _read_metadata(self) -> None:
         if nb_streams != 1:
             self._logger.warning(f"Only using first of {nb_streams} streams")
 
-        raw_comment = info["format"].get("tags", {}).get("comment", "{}")
+        tags = info["format"].get("tags", {})
+        # read comment field, which can be either lower case or upper case
+        raw_comment = tags.get("comment", tags.get("COMMENT", "{}"))
+        if raw_comment == "{}":
+            self._logger.warning("Could not find metadata written by `py-pde`")
         metadata = json.loads(shlex.split(raw_comment)[0])
 
         version = metadata.pop("version", 1)
@@ -181,7 +190,15 @@ def _read_metadata(self) -> None:
         try:
             self.info["num_frames"] = int(stream["nb_frames"])
         except KeyError:
-            self.info["num_frames"] = None  # number of frames was not stored
+            # frame count is not directly in the video
+            # => try determining it from the duration
+            try:
+                fps = Fraction(stream.get("avg_frame_rate", None))
+                duration = parse_duration(stream.get("tags", {}).get("DURATION"))
+            except TypeError:
+                raise RuntimeError("Frame count could not be read from video")
+            else:
+                self.info["num_frames"] = int(duration.total_seconds() * float(fps))
         self.info["width"] = stream["width"]
         self.info["height"] = stream["height"]
         if self.video_format == "auto":
@@ -332,6 +349,17 @@ def _append_data(self, data: np.ndarray, time: float) -> None:
         assert self._norms is not None
         assert self._format is not None
 
+        # check time
+        t_start = self.info.get("t_start")
+        if t_start is None:
+            t_start = 0
+        dt = self.info.get("dt", 1)
+        time_expect = t_start + dt * self.info["num_frames"]
+        if not np.isclose(time, time_expect):
+            if self.info.get("time_mismatch", False):
+                self._logger.warning(f"Detected time mismatch: {time} != {time_expect}")
+                self.info["time_mismatch"] = True
+
         # make sure there are two spatial dimensions
         grid_dim = self._grid.num_axes
         if grid_dim > 2:
diff --git a/pde/tools/ffmpeg.py b/pde/tools/ffmpeg.py
index 1807ab70..1fcf36cd 100644
--- a/pde/tools/ffmpeg.py
+++ b/pde/tools/ffmpeg.py
@@ -22,11 +22,12 @@
 class FFmpegFormat:
     """defines a FFmpeg format used for storing field data in a video
 
-    All pixel formats supported by FFmpeg can be obtained by running
-    :code:`ffmpeg -pix_fmts`. However, not all pixel formats are supported by all
-    codecs. Supported pixel formats are listed in the output of
-    :code:`ffmpeg -h encoder=<ENCODER>`, where `<ENCODER>` is one of the encoders listed
-    in :code:`ffmpeg -codecs`.
+    Note:
+        All pixel formats supported by FFmpeg can be obtained by running
+        :code:`ffmpeg -pix_fmts`. However, not all pixel formats are supported by all
+        codecs. Supported pixel formats are listed in the output of
+        :code:`ffmpeg -h encoder=<ENCODER>`, where `<ENCODER>` is one of the encoders
+        listed in :code:`ffmpeg -codecs`.
     """
 
     pix_fmt_file: str
@@ -134,7 +135,7 @@ def data_from_frame(self, frame_data: np.ndarray):
     #     dtype=np.dtype("<f4"),
     # ),
 }
-"""dict of :class:`FFmpegFormat` formats"""
+"""dict of pre-defined :class:`FFmpegFormat` formats"""
 
 
 def find_format(channels: int, bits_per_channel: int = 8) -> Optional[str]:
@@ -160,5 +161,6 @@ def find_format(channels: int, bits_per_channel: int = 8) -> Optional[str]:
                 or f.bits_per_channel < f_best.bits_per_channel
                 or f.channels < f_best.channels
             ):
+                # the current format is better than the previous one
                 n_best, f_best = n, f
     return n_best
diff --git a/tests/storage/test_movie_storages.py b/tests/storage/test_movie_storages.py
index 4df24d43..96f30e22 100644
--- a/tests/storage/test_movie_storages.py
+++ b/tests/storage/test_movie_storages.py
@@ -90,6 +90,33 @@ def test_movie_storage_vector(dim, tmp_path, rng):
         np.testing.assert_allclose(field.data, storage[i].data, atol=0.01)
 
 
+@pytest.mark.skipif(not module_available("ffmpeg"), reason="requires `ffmpeg-python`")
+@pytest.mark.parametrize("ext", [".mov", ".avi", ".mkv"])
+def test_movie_storage_containers(ext, tmp_path, rng):
+    """test storing scalar field as movie with different extensions"""
+    path = tmp_path / f"test_movie_storage_scalar{ext}"
+
+    grid = pde.UnitGrid([16])
+    field = pde.ScalarField.random_uniform(grid, rng=rng)
+    eq = pde.DiffusionPDE()
+    writer = MovieStorage(path)
+    storage = pde.MemoryStorage()
+    eq.solve(
+        field,
+        t_range=3.5,
+        dt=0.1,
+        backend="numpy",
+        tracker=[storage.tracker(2), writer.tracker(2)],
+    )
+
+    reader = MovieStorage(path)
+    assert len(reader) == 2
+    np.testing.assert_allclose(reader.times, [0, 2])
+    for i, field in enumerate(reader):
+        assert field.grid == grid
+        np.testing.assert_allclose(field.data, storage[i].data, atol=0.01)
+
+
 @pytest.mark.skipif(not module_available("ffmpeg"), reason="requires `ffmpeg-python`")
 @pytest.mark.parametrize("name,video_format", formats.items())
 def test_video_format(name, video_format, tmp_path, rng):

From 404ae2247da01f39f82293e87e73071d2782a5f1 Mon Sep 17 00:00:00 2001
From: David Zwicker <david.zwicker@ds.mpg.de>
Date: Tue, 19 Mar 2024 15:41:04 +0100
Subject: [PATCH 13/13] Improved corner interpolation of fields with boundary
 conditions

---
 pde/fields/base.py                 | 10 ++++++---
 pde/grids/boundaries/axes.py       | 34 +++++++++++++++++++++++++++++-
 tests/fields/test_scalar_fields.py | 24 +++++++++++++++++++++
 3 files changed, 64 insertions(+), 4 deletions(-)

diff --git a/pde/fields/base.py b/pde/fields/base.py
index 148d9913..3a0cd253 100644
--- a/pde/fields/base.py
+++ b/pde/fields/base.py
@@ -1372,7 +1372,7 @@ def interpolate(
         """
         if bc is not None:
             # impose boundary conditions and then interpolate using ghost cells
-            self.set_ghost_cells(bc)
+            self.set_ghost_cells(bc, set_corners=True)
             interpolator = self.make_interpolator(fill=fill, with_ghost_cells=True)
 
         else:
@@ -1502,19 +1502,23 @@ def get_boundary_values(
         return (self._data_full[i_wall] + self._data_full[i_ghost]) / 2  # type: ignore
 
     @fill_in_docstring
-    def set_ghost_cells(self, bc: BoundariesData, *, args=None) -> None:
+    def set_ghost_cells(
+        self, bc: BoundariesData, *, set_corners: bool = False, args=None
+    ) -> None:
         """set the boundary values on virtual points for all boundaries
 
         Args:
             bc (str or list or tuple or dict):
                 The boundary conditions applied to the field.
                 {ARG_BOUNDARIES}
+            set_corners (bool):
+                Determines whether the corner cells are set using interpolation
             args:
                 Additional arguments that might be supported by special boundary
                 conditions.
         """
         bcs = self.grid.get_boundary_conditions(bc, rank=self.rank)
-        bcs.set_ghost_cells(self._data_full, args=args)
+        bcs.set_ghost_cells(self._data_full, args=args, set_corners=set_corners)
 
     @property
     @abstractmethod
diff --git a/pde/grids/boundaries/axes.py b/pde/grids/boundaries/axes.py
index 59ef458c..56933c4c 100644
--- a/pde/grids/boundaries/axes.py
+++ b/pde/grids/boundaries/axes.py
@@ -8,6 +8,8 @@
 
 from __future__ import annotations
 
+import itertools
+import logging
 from collections.abc import Iterator, Sequence
 from typing import Union
 
@@ -277,12 +279,16 @@ def get_mathematical_representation(self, field_name: str = "C") -> str:
 
         return "\n".join(result)
 
-    def set_ghost_cells(self, data_full: np.ndarray, *, args=None) -> None:
+    def set_ghost_cells(
+        self, data_full: np.ndarray, *, set_corners: bool = False, args=None
+    ) -> None:
         """set the ghost cells for all boundaries
 
         Args:
             data_full (:class:`~numpy.ndarray`):
                 The full field data including ghost points
+            set_corners (bool):
+                Determines whether the corner cells are set using interpolation
             args:
                 Additional arguments that might be supported by special boundary
                 conditions.
@@ -290,6 +296,32 @@ def set_ghost_cells(self, data_full: np.ndarray, *, args=None) -> None:
         for b in self:
             b.set_ghost_cells(data_full, args=args)
 
+        if set_corners and self.grid.num_axes >= 2:
+            d = data_full  # abbreviation
+            nxt = [1, -2]  # maps 0 to 1 and -1 to -2 to obtain neighboring cells
+            if self.grid.num_axes == 2:
+                # iterate all corners
+                for i, j in itertools.product([0, -1], [0, -1]):
+                    d[..., i, j] = (d[..., nxt[i], j] + d[..., i, nxt[j]]) / 2
+
+            elif self.grid.num_axes >= 3:
+                # iterate all edges
+                for i, j in itertools.product([0, -1], [0, -1]):
+                    d[..., :, i, j] = (+d[..., :, nxt[i], j] + d[..., :, i, nxt[j]]) / 2
+                    d[..., i, :, j] = (+d[..., nxt[i], :, j] + d[..., i, :, nxt[j]]) / 2
+                    d[..., i, j, :] = (+d[..., nxt[i], j, :] + d[..., i, nxt[j], :]) / 2
+                # iterate all corners
+                for i, j, k in itertools.product(*[[0, -1]] * 3):
+                    d[..., i, j, k] = (
+                        d[..., nxt[i], j, k]
+                        + d[..., i, nxt[j], k]
+                        + d[..., i, j, nxt[k]]
+                    ) / 3
+            else:
+                logging.getLogger(self.__class__.__name__).warning(
+                    f"Can't interpolate corners for grid with {self.grid.num_axes} axes"
+                )
+
     def make_ghost_cell_setter(self) -> GhostCellSetter:
         """return function that sets the ghost cells on a full array"""
         ghost_cell_setters = tuple(b.make_ghost_cell_setter() for b in self)
diff --git a/tests/fields/test_scalar_fields.py b/tests/fields/test_scalar_fields.py
index a6cf810a..e16152e3 100644
--- a/tests/fields/test_scalar_fields.py
+++ b/tests/fields/test_scalar_fields.py
@@ -558,3 +558,27 @@ def test_field_split(decomp, rng):
         np.testing.assert_allclose(field.data, field_data)
     else:
         assert field_data is None
+
+
+def test_field_corner_interpolation_2d():
+    """test corner interpolation for a 2d field"""
+    f = ScalarField(UnitGrid([1, 1]), 0)
+    bc_x = [{"value": -1}, {"value": 2}]
+    bc_y = [{"value": -2}, {"value": 1}]
+    f.set_ghost_cells(bc=[bc_x, bc_y], set_corners=True)
+    expect = np.array([[-1.5, -2, 0], [-1, 0, 2], [0, 1, 1.5]])
+    np.testing.assert_allclose(f._data_full, 2 * expect.T)
+
+
+def test_field_corner_interpolation_3d():
+    """test corner interpolation for a 3d field"""
+    f = ScalarField(UnitGrid([1, 1, 1]), 0)
+    f.set_ghost_cells(bc=[[{"value": -3}, {"value": 3}]] * 3, set_corners=True)
+    expect = np.array(
+        [
+            [[-6, -6, -2], [-6, -6, 0], [-2, 0, 2]],
+            [[-6, -6, 0], [-6, 0, 6], [0, 6, 6]],
+            [[-2, 0, 2], [0, 6, 6], [2, 6, 6]],
+        ]
+    )
+    np.testing.assert_allclose(f._data_full, expect)