Improved support for mkv container format

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:

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

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):