Improved typing and tests

droplets/droplets.py

@@ -32,11 +32,12 @@
 import warnings
 from abc import ABCMeta, abstractmethod
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Tuple, TypeVar
+from typing import Any, Callable, Dict, List, Optional, Tuple, Type, TypeVar, Union
 
 import numpy as np
 from numba.extending import register_jitable
 from numpy.lib.recfunctions import structured_to_unstructured
+from numpy.typing import DTypeLike
 from scipy import integrate
 
 from pde.fields import ScalarField
@@ -47,17 +48,11 @@
 
 from .tools import spherical
 
-# work-around to satisfy type checking in python 3.6
-if TYPE_CHECKING:
-    # TYPE_CHECKING will always be False in production and this circular import
-    # will thus be resolved.
-    from .emulsions import EmulsionTimeCourse  # @UnusedImport
-
-
 TDroplet = TypeVar("TDroplet", bound="DropletBase")
+DTypeList = List[Union[Tuple[str, Type[Any]], Tuple[str, Type[Any], Tuple[int, ...]]]]
 
 
-def get_dtype_field_size(dtype, field_name: str) -> int:
+def get_dtype_field_size(dtype: DTypeLike, field_name: str) -> int:
     """return the number of elements in a field of structured numpy array
 
     Args:
@@ -66,7 +61,7 @@ def get_dtype_field_size(dtype, field_name: str) -> int:
         field_name (str):
             The name of the field that needs to be checked
     """
-    shape = dtype.fields[field_name][0].shape
+    shape = dtype.fields[field_name][0].shape  # type: ignore
     return np.prod(shape) if shape else 1  # type: ignore
 
 
@@ -114,6 +109,9 @@ class DropletBase:
 
     data: np.recarray  # all information about the droplet in a record array
 
+    _merge_data: Callable[[np.ndarray, np.ndarray, np.ndarray], None]
+    """private method for merging droplet data, created by __init_subclass__"""
+
     @classmethod
     def from_data(cls, data: np.recarray) -> DropletBase:
         """create droplet class from a given data
@@ -147,14 +145,14 @@ def from_droplet(cls, droplet: DropletBase, **kwargs) -> DropletBase:
 
     @classmethod
     @abstractmethod
-    def get_dtype(cls, **kwargs):
+    def get_dtype(cls, **kwargs) -> DTypeList:
         """determine the dtype representing this droplet class
 
         Returns:
             :class:`numpy.dtype`:
                 The (structured) dtype associated with this class
         """
-        pass
+        ...
 
     def _init_data(self, **kwargs) -> None:
         """initializes the `data` attribute if it is not present
@@ -271,7 +269,7 @@ def check_data(self):
             raise ValueError("Radius must be positive")
 
     @classmethod
-    def get_dtype(cls, **kwargs):
+    def get_dtype(cls, **kwargs) -> DTypeList:
         """determine the dtype representing this droplet class
 
         Args:
@@ -320,7 +318,7 @@ def position(self) -> np.ndarray:
         return self.data["position"]
 
     @position.setter
-    def position(self, value: np.ndarray):
+    def position(self, value: np.ndarray) -> None:
         value = np.asanyarray(value)
         if len(value) != self.dim:
             raise ValueError(f"The dimension of the position must be {self.dim}")
@@ -332,7 +330,7 @@ def radius(self) -> float:
         return float(self.data["radius"])
 
     @radius.setter
-    def radius(self, value: float):
+    def radius(self, value: float) -> None:
         self.data["radius"] = value
         self.check_data()
 
@@ -342,7 +340,7 @@ def volume(self) -> float:
         return spherical.volume_from_radius(self.radius, self.dim)
 
     @volume.setter
-    def volume(self, volume: float):
+    def volume(self, volume: float) -> None:
         """set the radius from a supplied volume"""
         self.radius = spherical.radius_from_volume(volume, self.dim)
 
@@ -447,7 +445,7 @@ def interface_curvature(self) -> float:
         """float: the mean curvature of the interface of the droplet"""
         return 1 / self.radius
 
-    def _get_phase_field(self, grid: GridBase, dtype=np.double) -> np.ndarray:
+    def _get_phase_field(self, grid: GridBase, dtype: DTypeLike = float) -> np.ndarray:
         """Creates an image of the droplet on the `grid`
 
         Args:
@@ -626,7 +624,7 @@ def data_bounds(self) -> Tuple[np.ndarray, np.ndarray]:
         return l, h
 
     @classmethod
-    def get_dtype(cls, **kwargs):
+    def get_dtype(cls, **kwargs) -> DTypeList:
         """determine the dtype representing this droplet class
 
         Args:
@@ -661,7 +659,7 @@ def interface_width(self) -> Optional[float]:
             return float(self.data["interface_width"])
 
     @interface_width.setter
-    def interface_width(self, value: Optional[float]):
+    def interface_width(self, value: Optional[float]) -> None:
         if value is None:
             self.data["interface_width"] = math.nan
         elif value < 0:
@@ -670,7 +668,7 @@ def interface_width(self, value: Optional[float]):
             self.data["interface_width"] = value
         self.check_data()
 
-    def _get_phase_field(self, grid: GridBase, dtype=np.double) -> np.ndarray:
+    def _get_phase_field(self, grid: GridBase, dtype: DTypeLike = float) -> np.ndarray:
         """Creates an image of the droplet on the `grid`
 
         Args:
@@ -699,7 +697,7 @@ def _get_phase_field(self, grid: GridBase, dtype=np.double) -> np.ndarray:
         dist: np.ndarray = grid.polar_coordinates_real(self.position, ret_angle=False)  # type: ignore
 
         # make the image
-        if interface_width == 0 or dtype == np.bool_:
+        if interface_width == 0 or np.issubdtype(dtype, bool):
             result = dist < self.radius
         else:
             result = 0.5 + 0.5 * np.tanh((self.radius - dist) / interface_width)  # type: ignore
@@ -745,7 +743,7 @@ def __init__(
             raise ValueError(f"Space dimension must be {self.__class__.dim}")
 
     @classmethod
-    def get_dtype(cls, **kwargs):
+    def get_dtype(cls, **kwargs) -> DTypeList:
         """determine the dtype representing this droplet class
 
         Args:
@@ -791,7 +789,7 @@ def amplitudes(self) -> np.ndarray:
         return np.atleast_1d(self.data["amplitudes"])  # type: ignore
 
     @amplitudes.setter
-    def amplitudes(self, value: Optional[np.ndarray] = None):
+    def amplitudes(self, value: Optional[np.ndarray] = None) -> None:
         if value is None:
             assert self.modes == 0
             self.data["amplitudes"] = np.broadcast_to(0.0, (0,))
@@ -812,14 +810,14 @@ def volume(self) -> float:
         raise NotImplementedError
 
     @volume.setter
-    def volume(self, volume: float):
+    def volume(self, volume: float) -> None:
         raise NotImplementedError
 
     @property
     def surface_area(self) -> float:
         raise NotImplementedError
 
-    def _get_phase_field(self, grid: GridBase, dtype=np.double) -> np.ndarray:
+    def _get_phase_field(self, grid: GridBase, dtype: DTypeLike = float) -> np.ndarray:
         """Creates a normalized image of the droplet on the `grid`
 
         Args:
@@ -848,7 +846,7 @@ def _get_phase_field(self, grid: GridBase, dtype=np.double) -> np.ndarray:
         interface = self.interface_distance(*angles)
 
         # make the image
-        if interface_width == 0 or dtype == np.bool_:
+        if interface_width == 0 or np.issubdtype(dtype, bool):
             result = dist < interface
         else:
             result = 0.5 + 0.5 * np.tanh((interface - dist) / interface_width)
@@ -927,7 +925,7 @@ def interface_distance(self, φ: np.ndarray) -> np.ndarray:  # type: ignore
         Returns:
             Array with distances of the interfacial points associated with each angle φ
         """
-        dist = np.ones(φ.shape, dtype=np.double)
+        dist = np.ones(φ.shape, dtype=float)
         for n, (a, b) in enumerate(iterate_in_pairs(self.amplitudes), 1):  # no 0th mode
             if a != 0:
                 dist += a * np.sin(n * φ)
@@ -964,7 +962,7 @@ def interface_curvature(self, φ: np.ndarray) -> np.ndarray:  # type: ignore
         Returns:
             Array with curvature at the interfacial points associated with each angle φ
         """
-        curv_radius = np.ones(φ.shape, dtype=np.double)
+        curv_radius = np.ones(φ.shape, dtype=float)
         for n, (a, b) in enumerate(iterate_in_pairs(self.amplitudes), 1):  # no 0th mode
             factor = n * n - 1
             if a != 0:
@@ -980,7 +978,7 @@ def volume(self) -> float:
         return np.pi * self.radius**2 * term  # type: ignore
 
     @volume.setter
-    def volume(self, volume: float):
+    def volume(self, volume: float) -> None:
         """set volume keeping relative perturbations"""
         term = 1 + np.sum(self.amplitudes**2) / 2
         self.radius = np.sqrt(volume / (np.pi * term))
@@ -991,8 +989,8 @@ def surface_area(self) -> float:
         # discretize surface for simple approximation to integral
         φs, dφ = np.linspace(0, 2 * np.pi, 256, endpoint=False, retstep=True)
 
-        dist = np.ones(φs.shape, dtype=np.double)
-        dist_dφ = np.zeros(φs.shape, dtype=np.double)
+        dist = np.ones(φs.shape, dtype=float)
+        dist_dφ = np.zeros(φs.shape, dtype=float)
         for n, (a, b) in enumerate(iterate_in_pairs(self.amplitudes), 1):  # no 0th mode
             if a != 0:
                 dist += a * np.sin(n * φs)
@@ -1116,7 +1114,7 @@ def interface_distance(  # type: ignore
             φ = np.zeros_like(θ)
         elif θ.shape != φ.shape:
             raise ValueError("Shape of θ and φ must agree")
-        dist = np.ones(θ.shape, dtype=np.double)
+        dist = np.ones(θ.shape, dtype=float)
         for k, a in enumerate(self.amplitudes, 1):  # skip zero-th mode!
             if a != 0:
                 dist += a * spherical.spherical_harmonic_real_k(k, θ, φ)  # type: ignore
@@ -1192,7 +1190,7 @@ def integrand(θ, φ):
         return volume  # type: ignore
 
     @volume.setter
-    def volume(self, volume: float):
+    def volume(self, volume: float) -> None:
         """set volume keeping relative perturbations"""
         raise NotImplementedError("Cannot set volume")
 
@@ -1242,7 +1240,7 @@ def interface_distance(self, θ: np.ndarray) -> np.ndarray:  # type: ignore
         Returns:
             Array with distances of the interfacial points associated with the angles
         """
-        dist = np.ones(θ.shape, dtype=np.double)
+        dist = np.ones(θ.shape, dtype=float)
         for order, a in enumerate(self.amplitudes, 1):  # skip zero-th mode!
             if a != 0:
                 dist += a * spherical.spherical_harmonic_symmetric(order, θ)  # type: ignore

droplets/emulsions.py

@@ -1,7 +1,6 @@
 """
 Classes describing collections of droplets, i.e. emulsions, and their temporal dynamics.
 
-
 .. autosummary::
    :nosignatures:
 

droplets/image_analysis.py

@@ -344,7 +344,7 @@ def locate_droplets_in_mask(mask: ScalarField) -> Emulsion:
 
 def locate_droplets(
     phase_field: ScalarField,
-    threshold: Union[float, str] = 0.5,
+    threshold: Union[float, Literal["auto", "extrema", "mean", "otsu"]] = 0.5,
     *,
     minimal_radius: float = 0,
     modes: int = 0,
@@ -533,7 +533,7 @@ def refine_droplet(
         droplet.interface_width = phase_field.grid.typical_discretization
 
     # enlarge the mask to also contain the shape change
-    mask = droplet._get_phase_field(phase_field.grid, dtype=np.bool_)
+    mask = droplet._get_phase_field(phase_field.grid, dtype=bool)
     dilation_iterations = 1 + int(2 * droplet.interface_width)
     mask = ndimage.binary_dilation(mask, iterations=dilation_iterations)
 
@@ -609,8 +609,8 @@ def _image_deviation(params):
 
 def get_structure_factor(
     scalar_field: ScalarField,
-    smoothing: Union[None, float, str] = "auto",
-    wave_numbers: Union[Sequence[float], str] = "auto",
+    smoothing: Union[None, float, Literal["auto", "none"]] = "auto",
+    wave_numbers: Union[Sequence[float], Literal["auto"]] = "auto",
     add_zero: bool = False,
 ) -> Tuple[np.ndarray, np.ndarray]:
     r"""Calculates the structure factor associated with a field

droplets/trackers.py

@@ -11,7 +11,7 @@
 """
 
 import math
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable, Dict, List, Literal, Optional, Union
 
 from pde.fields.base import FieldBase
 from pde.tools.docstrings import fill_in_docstring
@@ -36,7 +36,9 @@ def __init__(
         interval: IntervalData = 1,
         filename: Optional[str] = None,
         *,
-        method: str = "structure_factor_mean",
+        method: Literal[
+            "structure_factor_mean", "structure_factor_maximum", "droplet_detection"
+        ] = "structure_factor_mean",
         source: Union[None, int, Callable] = None,
         verbose: bool = False,
     ):
@@ -131,7 +133,7 @@ def __init__(
         *,
         emulsion_timecourse=None,
         source: Union[None, int, Callable] = None,
-        threshold: Union[float, str] = 0.5,
+        threshold: Union[float, Literal["auto", "extrema", "mean", "otsu"]] = 0.5,
         minimal_radius: float = 0,
         refine: bool = False,
         refine_args: Optional[Dict[str, Any]] = None,

tests/conftest.py

@@ -0,0 +1,35 @@
+"""
+This file is used to configure the test environment when running py.test
+
+.. codeauthor:: David Zwicker <[email protected]>
+"""
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pytest
+
+from pde.tools.numba import random_seed
+
+
+@pytest.fixture(scope="function", autouse=False, name="rng")
+def init_random_number_generators():
+    """get a random number generator and set the seed of the random number generator
+
+    The function returns an instance of :func:`~numpy.random.default_rng()` and
+    initializes the default generators of both :mod:`numpy` and :mod:`numba`.
+    """
+    random_seed()
+    return np.random.default_rng(0)
+
+
+@pytest.fixture(scope="function", autouse=True)
+def setup_and_teardown():
+    """helper function adjusting environment before and after tests"""
+    # raise all underflow errors
+    np.seterr(all="raise", under="ignore")
+
+    # run the actual test
+    yield
+
+    # clean up open matplotlib figures after the test
+    plt.close("all")