From 642cd2b5a589196e2a7bd3100f8cf2bdb5c1ab82 Mon Sep 17 00:00:00 2001
From: David Zwicker <david.zwicker@ds.mpg.de>
Date: Tue, 6 Aug 2024 15:50:48 +0200
Subject: [PATCH] Improved plotting of emulsions on periodic grids

---
 droplets/emulsions.py  | 28 ++++++++++++++--------------
 tests/test_emulsion.py |  1 +
 2 files changed, 15 insertions(+), 14 deletions(-)

diff --git a/droplets/emulsions.py b/droplets/emulsions.py
index fc46bc4..7691130 100644
--- a/droplets/emulsions.py
+++ b/droplets/emulsions.py
@@ -16,7 +16,6 @@
 import json
 import logging
 import math
-import warnings
 from collections.abc import Iterable, Iterator, Sequence
 from typing import TYPE_CHECKING, Any, Callable, Literal, overload
 
@@ -52,7 +51,6 @@ def __init__(
         copy: bool = True,
         dtype: np.typing.DTypeLike | np.ndarray | SphericalDroplet = None,
         force_consistency: bool = False,
-        grid: GridBase | None = None,
     ):
         """
         Args:
@@ -71,10 +69,6 @@ def __init__(
         """
         super().__init__()
 
-        if grid is not None:
-            # deprecated on 2023-08-29
-            warnings.warn("`grid` argument is deprecated", DeprecationWarning)
-
         # store general information about droplets using a single dtype
         if isinstance(dtype, SphericalDroplet):
             dtype = dtype.data.dtype  # extract dtype from actual droplet
@@ -630,6 +624,7 @@ def plot(
         field: ScalarField | None = None,
         image_args: dict[str, Any] | None = None,
         repeat_periodically: bool = True,
+        grid: GridBase | None = None,
         color_value: Callable | None = None,
         cmap=None,
         norm=None,
@@ -652,6 +647,9 @@ def plot(
             repeat_periodically (bool):
                 flag determining whether droplets are shown on both sides of
                 periodic boundary conditions. This option can slow down plotting
+            grid (:class:`~pde.grids.base.GridBase`):
+                The grid on which the droplets are defined, which is necessary if
+                periodic boundary conditions should be respected for measuring distances
             color_value (callable):
                 Function used to determine the color of a droplet. The function is
                 called with individual droplet objects and must return a single scalar
@@ -694,16 +692,18 @@ def plot(
                 image_args = {}
             field.plot(kind="image", ax=ax, **image_args)
             ax.autoscale(False)  # fix image bounds to phase field
-            grid = field.grid
-            if isinstance(grid, CartesianGrid):
-                # determine the bounds from the (2d) grid
-                bounds = grid.axes_bounds
-            else:
-                # determine the bounds from the emulsion data itself
-                bounds = self.bbox.bounds
+
+            # extract grid for plotting
+            if grid is None:
+                grid = field.grid
+            elif grid != field.grid:
+                raise ValueError("Argument `grid` is incompatible with `field.grid`")
+
+        if grid is not None and isinstance(grid, CartesianGrid):
+            # determine the bounds from the (2d) grid
+            bounds = grid.axes_bounds
         else:
             # determine the bounds from the emulsion data itself
-            grid = None
             bounds = self.bbox.bounds
 
         ax.set_xlim(*bounds[0])
diff --git a/tests/test_emulsion.py b/tests/test_emulsion.py
index 75b057b..b2e0116 100644
--- a/tests/test_emulsion.py
+++ b/tests/test_emulsion.py
@@ -238,6 +238,7 @@ def test_emulsion_plotting():
     for e2 in es:
         e2.plot()
         e2.plot(field=field, repeat_periodically=True)
+        e2.plot(grid=field.grid, repeat_periodically=True)
         e2.plot(color_value=lambda droplet: droplet.radius)
 
     # 3d emulsion