diff --git a/.github/workflows/CI.yaml b/.github/workflows/CI.yaml index d774b32..a5bfe0a 100644 --- a/.github/workflows/CI.yaml +++ b/.github/workflows/CI.yaml @@ -13,23 +13,19 @@ jobs: - name: Checkout repository uses: actions/checkout@v4 - - name: Set up Miniforge - uses: conda-incubator/setup-miniconda@v2 + - name: Install uv + uses: astral-sh/setup-uv@v8.2.0 with: - miniforge-version: latest - activate-environment: test-env - python-version: 3.12 + python-version: "3.12" + + - name: Install ffmpeg + run: sudo apt-get update && sudo apt-get install -y --no-install-recommends ffmpeg - name: Install dependencies - shell: bash -l {0} - run: | - mamba install -c conda-forge xarray geopandas pooch matplotlib tqdm pytest ffmpeg netcdf4 rioxarray - pip install -e . + run: uv sync --extra test - name: Run pytest - shell: bash -l {0} - run: | - pytest tests/ -vv -s + run: uv run pytest tests/ -vv -s format: runs-on: ubuntu-latest @@ -39,11 +35,11 @@ jobs: with: ref: ${{ github.head_ref }} - - name: Install Ruff - run: pip install ruff + - name: Install uv + uses: astral-sh/setup-uv@v8.2.0 - name: Run ruff format - run: ruff format . + run: uvx ruff format . - name: Commit changes if: success() diff --git a/mapflow/_classic.py b/mapflow/_classic.py index c3d7d28..56a7999 100644 --- a/mapflow/_classic.py +++ b/mapflow/_classic.py @@ -113,6 +113,15 @@ def _log_norm(data, vmin, vmax, qmin, qmax): raise ValueError(f"Normalization range for log scale must be positive. Got vmin={vmin}, vmax={vmax}") return LogNorm(vmin=vmin, vmax=vmax) + @staticmethod + def _validate_quantiles(qmin, qmax): + if not (0 <= qmin <= 100): + raise ValueError(f"qmin must be between 0 and 100, got {qmin}") + if not (0 <= qmax <= 100): + raise ValueError(f"qmax must be between 0 and 100, got {qmax}") + if qmin >= qmax: + raise ValueError(f"qmin must be less than qmax, got {qmin} and {qmax}") + @staticmethod def _norm(data, vmin, vmax, qmin, qmax, norm, log, diff=False): """Generates a normalization based on the specified parameters. @@ -133,13 +142,7 @@ def _norm(data, vmin, vmax, qmin, qmax, norm, log, diff=False): Raises: ValueError: If qmin/qmax are not between 0-100 or if log=True with no positive values. """ - # Validate quantile ranges - if not (0 <= qmin <= 100): - raise ValueError(f"qmin must be between 0 and 100, got {qmin}") - if not (0 <= qmax <= 100): - raise ValueError(f"qmax must be between 0 and 100, got {qmax}") - if qmin >= qmax: - raise ValueError(f"qmin must be less than qmax, got {qmin} and {qmax}") + PlotModel._validate_quantiles(qmin, qmax) if norm is not None: return norm @@ -157,6 +160,77 @@ def _norm(data, vmin, vmax, qmin, qmax, norm, log, diff=False): vmax = np.nanpercentile(data, q=qmax) if vmax is None else vmax return Normalize(vmin=vmin, vmax=vmax) + @staticmethod + def _norm_streaming(frames, vmin, vmax, qmin, qmax, norm, log, diff=False): + """Streaming counterpart of :meth:`_norm` operating on an iterable of 2D frames. + + Frames are consumed one at a time, so the full 3D array is never held + in memory. Quantile-based bounds are aggregated per frame (minimum of + the per-frame lower quantiles, maximum of the per-frame upper + quantiles), which yields a range at least as wide as the exact global + quantiles. The frames iterable is not consumed at all when explicit + bounds (``norm``, or ``vmin`` and ``vmax``) make the data pass + unnecessary. + + Args: + frames (Iterable[np.ndarray]): Iterable of 2D frames. + vmin (float): Minimum value for normalization. + vmax (float): Maximum value for normalization. + qmin (float): Minimum quantile for normalization (0-100). + qmax (float): Maximum quantile for normalization (0-100). + norm (matplotlib.colors.Normalize): Custom normalization object. + log (bool): Indicates if a logarithmic scale should be used. + diff (bool): Indicates if a divergent colormap should be used. + + Returns: + matplotlib.colors.Normalize: Normalization object. + """ + PlotModel._validate_quantiles(qmin, qmax) + + if norm is not None: + return norm + + if diff: + if vmax is None: + high = None + for frame in frames: + value = np.nanpercentile(np.abs(frame), q=qmax) + if not np.isnan(value): + high = value if high is None else max(high, value) + vmax = high + vmin = None if vmax is None else -vmax + return Normalize(vmin=vmin, vmax=vmax) + + if log: + if vmin is None or vmax is None: + low = high = None + for frame in frames: + positive = frame[frame > 0] + if positive.size == 0: + continue + frame_min, frame_max = np.nanpercentile(positive, q=[qmin, qmax]) + low = frame_min if low is None else min(low, frame_min) + high = frame_max if high is None else max(high, frame_max) + if low is None or high is None: + return Normalize(vmin=1e-1, vmax=1e0) + vmin = low if vmin is None else vmin + vmax = high if vmax is None else vmax + if vmin <= 0 or vmax <= 0: + raise ValueError(f"Normalization range for log scale must be positive. Got vmin={vmin}, vmax={vmax}") + return LogNorm(vmin=vmin, vmax=vmax) + + if vmin is None or vmax is None: + low = high = None + for frame in frames: + frame_min, frame_max = np.nanpercentile(frame, q=[qmin, qmax]) + if not np.isnan(frame_min): + low = frame_min if low is None else min(low, frame_min) + if not np.isnan(frame_max): + high = frame_max if high is None else max(high, frame_max) + vmin = low if vmin is None else vmin + vmax = high if vmax is None else vmax + return Normalize(vmin=vmin, vmax=vmax) + def _process_data(self, data): if not isinstance(data, np.ndarray): data = np.asarray(data) @@ -387,6 +461,42 @@ def upsample(data, ratio=5): ret[k::ratio] = ret[::ratio][:-1] + k * delta / ratio return ret + @staticmethod + def _frame_values(data, k): + """Extracts frame ``k`` as a numpy array, loading it lazily if possible.""" + frame = data[k] + return np.asarray(getattr(frame, "values", frame)) + + @classmethod + def _iter_raw_frames(cls, data): + """Yields raw frames one at a time as numpy arrays.""" + for k in range(len(data)): + yield cls._frame_values(data, k) + + @classmethod + def _iter_upsampled_frames(cls, data, ratio=1): + """Yields temporally interpolated frames one at a time. + + Streaming equivalent of :meth:`upsample`: at most two source frames + are held in memory at any time, so lazily-backed data (e.g. an xarray + DataArray opened from a netCDF/zarr store or backed by dask) is never + fully loaded, and the interpolated frames are never all materialized + at once. + """ + n_raw = len(data) + if n_raw == 0: + raise ValueError("data must contain at least one frame.") + previous = cls._frame_values(data, 0) + for k in range(1, n_raw): + yield previous + current = cls._frame_values(data, k) + if ratio > 1: + delta = (current - previous) / ratio + for j in range(1, ratio): + yield (previous + j * delta).astype(previous.dtype, copy=False) + previous = current + yield previous + @staticmethod def _process_title(title, upsample_ratio): if title is None: @@ -475,8 +585,10 @@ def __call__( these frames into a video file using FFmpeg. Args: - data (np.ndarray): A 3D numpy array where the first dimension is time - (or frame sequence) and the next two are spatial (y, x). + data (np.ndarray | xr.DataArray): A 3D array where the first dimension + is time (or frame sequence) and the next two are spatial (y, x). + Lazily-backed xarray DataArrays (netCDF/zarr stores, dask) are + streamed one time step at a time and never fully loaded in memory. path (str | Path): The output path for the generated video file. Supported formats are avi, mkv, mov, and mp4. figsize (tuple[float, float], optional): Figure size (width, height) @@ -525,7 +637,10 @@ def __call__( self.plot.aspect, ) - norm = self.plot._norm(data, vmin, vmax, qmin, qmax, norm, log, diff) + if isinstance(data, np.ndarray): + norm = self.plot._norm(data, vmin, vmax, qmin, qmax, norm, log, diff) + else: + norm = self.plot._norm_streaming(self._iter_raw_frames(data), vmin, vmax, qmin, qmax, norm, log, diff) self._animate( data=data, path=path, @@ -569,17 +684,16 @@ def _animate( fixed_frame: bool = False, ): titles = self._process_title(title, upsample_ratio) - data = self.upsample(data, ratio=upsample_ratio) - data_len = len(data) + n_raw = len(data) + data_len = (n_raw - 1) * upsample_ratio + 1 if n_raw > 1 else 1 with TemporaryDirectory() as tempdir: - frame_paths = [Path(tempdir) / f"frame_{k:08d}.png" for k in range(data_len)] - args = [] - for k in range(data_len): - frame_data = data[k] - arg_tuple = ( - frame_data, - frame_paths[k], + # Generator consumed lazily by pool.imap: frames are interpolated + # and dispatched to workers on the fly, never all held in memory. + args = ( + ( + frame, + Path(tempdir) / f"frame_{k:08d}.png", figsize, titles[k] if titles and k < len(titles) else None, cmap, @@ -590,7 +704,8 @@ def _animate( fixed_frame, {"diff": diff}, ) - args.append(arg_tuple) + for k, frame in enumerate(self._iter_upsampled_frames(data, ratio=upsample_ratio)) + ) cpu_total = cpu_count() or 1 default_jobs = max(1, int((2 * cpu_total) / 3)) @@ -735,7 +850,10 @@ def animate( Args: da (xr.DataArray): Input DataArray with time as the animation dimension and x/y spatial dimensions. - 2D inputs are not supported. + 2D inputs are not supported. Lazily-backed DataArrays (netCDF/zarr stores, dask) are streamed + one time step at a time: the full dataset is never loaded in memory. Note that when relying on + quantile-based color scaling (no ``vmin``/``vmax``/``norm`` given), an extra streaming pass over + the data is required to compute the color range. path (str): Output path for the video file. Supported formats are avi, mkv, mov, and mp4. time_name (str, optional): Name of the time coordinate in `da`. If None, @@ -806,7 +924,7 @@ def animate( field = da.name or da.attrs.get("long_name") titles = [f"{field} - {t}" for t in time] animation( - data=da.values, + data=da, path=output_path, title=titles, label=unit, diff --git a/mapflow/_quiver.py b/mapflow/_quiver.py index 20407e1..7f6fa23 100644 --- a/mapflow/_quiver.py +++ b/mapflow/_quiver.py @@ -153,7 +153,9 @@ def quiver( Args: u (xr.DataArray): 3D DataArray for the U-component of the vector field. + Lazily-backed DataArrays are streamed one time step at a time. v (xr.DataArray): 3D DataArray for the V-component of the vector field. + Lazily-backed DataArrays are streamed one time step at a time. path (str | Path): The output path for the generated video file. subsample (int, optional): Subsampling factor for quiver arrows. Defaults to 1. figsize (tuple, optional): Figure size (width, height) in inches. @@ -176,9 +178,11 @@ def quiver( timeout (int | str, optional): Timeout for the ffmpeg command. Defaults to "auto". **kwargs: Additional keyword arguments. """ - magnitude = np.sqrt(u**2 + v**2) - norm = self.plot._norm( - magnitude.values, + magnitude_frames = ( + np.sqrt(self._frame_values(u, k) ** 2 + self._frame_values(v, k) ** 2) for k in range(len(u)) + ) + norm = self.plot._norm_streaming( + magnitude_frames, vmin=vmin, vmax=vmax, qmin=qmin, @@ -240,19 +244,20 @@ def _animate( titles = self._process_title(title, upsample_ratio) u, v = data - u_upsampled = self.upsample(u.values, ratio=upsample_ratio) - v_upsampled = self.upsample(v.values, ratio=upsample_ratio) - data = (u_upsampled, v_upsampled) - data_len = len(u_upsampled) + n_raw = len(u) + data_len = (n_raw - 1) * upsample_ratio + 1 if n_raw > 1 else 1 + frames = zip( + self._iter_upsampled_frames(u, ratio=upsample_ratio), + self._iter_upsampled_frames(v, ratio=upsample_ratio), + ) with TemporaryDirectory() as tempdir: - frame_paths = [Path(tempdir) / f"frame_{k:08d}.png" for k in range(data_len)] - args = [] - for k in range(data_len): - frame_data = (data[0][k], data[1][k]) - arg_tuple = ( + # Generator consumed lazily by pool.imap: frames are interpolated + # and dispatched to workers on the fly, never all held in memory. + args = ( + ( frame_data, - frame_paths[k], + Path(tempdir) / f"frame_{k:08d}.png", figsize, titles[k] if titles and k < len(titles) else None, cmap, @@ -263,7 +268,8 @@ def _animate( fixed_frame, kwargs, ) - args.append(arg_tuple) + for k, frame_data in enumerate(frames) + ) cpu_total = cpu_count() or 1 default_jobs = max(1, int((2 * cpu_total) / 3)) diff --git a/pyproject.toml b/pyproject.toml index cb34c38..c2522fa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,7 +15,7 @@ requires-python = ">=3.10" dependencies = ["xarray", "matplotlib", "geopandas", "tqdm"] [project.optional-dependencies] -test = ["netcdf4", "pytest", "rioxarray"] +test = ["netcdf4", "pytest", "rioxarray", "pooch"] [project.urls] Homepage = "https://github.com/CyrilJl/mapflow" diff --git a/tests/test_lazy.py b/tests/test_lazy.py new file mode 100644 index 0000000..28d93e5 --- /dev/null +++ b/tests/test_lazy.py @@ -0,0 +1,141 @@ +import numpy as np +import pytest +import xarray as xr +from matplotlib.colors import LogNorm, Normalize + +from mapflow import Animation, QuiverAnimation, animate +from mapflow._classic import PlotModel + + +class LoadCounter: + """Sequence wrapper that records which frames are materialized.""" + + def __init__(self, arr): + self.arr = arr + self.loads = [] + + def __len__(self): + return len(self.arr) + + def __getitem__(self, k): + self.loads.append(k) + return self.arr[k] + + +@pytest.mark.parametrize("ratio", [1, 2, 5]) +@pytest.mark.parametrize("dtype", ["float32", "float64"]) +def test_iter_upsampled_frames_matches_upsample(ratio, dtype): + rng = np.random.default_rng(0) + data = rng.random((5, 4, 3)).astype(dtype) + expected = Animation.upsample(data, ratio=ratio) + frames = list(Animation._iter_upsampled_frames(data, ratio=ratio)) + got = np.stack(frames) + assert got.dtype == expected.dtype + np.testing.assert_allclose(got, expected, rtol=1e-6) + + +def test_iter_upsampled_frames_single_frame(): + data = np.random.default_rng(0).random((1, 4, 3)) + frames = list(Animation._iter_upsampled_frames(data, ratio=5)) + assert len(frames) == 1 + np.testing.assert_array_equal(frames[0], data[0]) + + +def test_iter_upsampled_frames_empty(): + with pytest.raises(ValueError): + list(Animation._iter_upsampled_frames(np.empty((0, 4, 3)), ratio=2)) + + +def test_animation_streams_each_frame_once(tmp_path): + rng = np.random.default_rng(0) + arr = rng.random((6, 16, 16)) + counter = LoadCounter(arr) + animation = Animation(x=np.linspace(0, 15, 16), y=np.linspace(40, 55, 16)) + animation(counter, tmp_path / "out.mp4", vmin=0.0, vmax=1.0, upsample_ratio=3, dpi=60) + # vmin/vmax provided: no pass over the data for the norm, each raw frame + # is loaded exactly once, in order, by the streaming interpolation. + assert counter.loads == list(range(6)) + assert (tmp_path / "out.mp4").exists() + + +def test_quiver_streams_each_frame_once(tmp_path): + rng = np.random.default_rng(0) + u = LoadCounter(rng.random((5, 12, 12))) + v = LoadCounter(rng.random((5, 12, 12))) + animation = QuiverAnimation(x=np.linspace(0, 11, 12), y=np.linspace(40, 51, 12)) + animation.quiver( + u, + v, + tmp_path / "out.mp4", + vmin=0.0, + vmax=2.0, + dpi=60, + x_name="lon", + y_name="lat", + ) + assert u.loads == list(range(5)) + assert v.loads == list(range(5)) + assert (tmp_path / "out.mp4").exists() + + +def test_animate_lazy_netcdf(tmp_path): + rng = np.random.default_rng(0) + da = xr.DataArray( + rng.random((5, 12, 14)).astype("float32"), + dims=("time", "lat", "lon"), + coords={ + "time": np.arange(np.datetime64("2020-01-01"), np.datetime64("2020-01-06")), + "lat": np.linspace(40, 50, 12), + "lon": np.linspace(-5, 5, 14), + }, + name="field", + ) + nc_path = tmp_path / "data.nc" + da.to_netcdf(nc_path) + out = tmp_path / "out.mp4" + with xr.open_dataarray(nc_path) as lazy: + # Default quantile-based scaling: exercises the streaming norm pass. + animate(da=lazy, path=str(out), dpi=60) + assert out.exists() + + +def test_norm_streaming_default_bounds_widen_exact(): + rng = np.random.default_rng(0) + data = rng.normal(size=(4, 20, 20)) + exact = PlotModel._norm(data, None, None, 1, 99, None, log=False) + streamed = PlotModel._norm_streaming(iter(data), None, None, 1, 99, None, log=False) + assert type(streamed) is Normalize + assert streamed.vmin <= exact.vmin + assert streamed.vmax >= exact.vmax + + +def test_norm_streaming_log(): + rng = np.random.default_rng(0) + data = rng.random((4, 10, 10)) + 0.1 + streamed = PlotModel._norm_streaming(iter(data), None, None, 0.01, 99.9, None, log=True) + assert isinstance(streamed, LogNorm) + assert streamed.vmin > 0 + assert streamed.vmax > streamed.vmin + + +def test_norm_streaming_log_no_positive_values(): + data = np.full((3, 5, 5), -1.0) + streamed = PlotModel._norm_streaming(iter(data), None, None, 0.01, 99.9, None, log=True) + assert streamed.vmin == pytest.approx(1e-1) + assert streamed.vmax == pytest.approx(1e0) + + +def test_norm_streaming_diff_is_symmetric(): + rng = np.random.default_rng(0) + data = rng.normal(size=(4, 10, 10)) + streamed = PlotModel._norm_streaming(iter(data), None, None, 0.01, 99.9, None, log=False, diff=True) + assert streamed.vmin == -streamed.vmax + + +def test_norm_streaming_explicit_bounds_skip_data(): + def frames(): + raise AssertionError("frames should not be consumed when vmin/vmax are given") + yield + + norm = PlotModel._norm_streaming(frames(), 0.0, 1.0, 0.01, 99.9, None, log=False) + assert (norm.vmin, norm.vmax) == (0.0, 1.0)