docs: add type hints for emod viscosity

dclab/features/emodulus/viscosity.py

@@ -5,6 +5,7 @@
 import warnings
 
 import numpy as np
+import numpy.typing as npt
 
 from ...warn import PipelineWarning
 
@@ -45,10 +46,17 @@ class TemperatureOutOfRangeWarning(PipelineWarning):
 
 
 def check_temperature(model: str,
-                      temperature: float | np.array,
+                      temperature: float | npt.NDArray,
                       tmin: float,
-                      tmax: float):
-    """Raise a TemperatureOutOfRangeWarning if applicable"""
+                      tmax: float) -> None:
+    """Raise a TemperatureOutOfRangeWarning if applicable
+
+    Raises
+    ------
+    TemperatureOutOfRangeWarning
+        If the given temperature is out of the given range.
+
+    """
     if np.min(temperature) < tmin or np.max(temperature) > tmax:
         warnings.warn(
             f"For the {model} model, the temperature should be "
@@ -60,11 +68,12 @@ def check_temperature(model: str,
 def get_viscosity(medium: str = "0.49% MC-PBS",
                   channel_width: float = 20.0,
                   flow_rate: float = 0.16,
-                  temperature: float | np.ndarray = 23.0,
+                  temperature: float | npt.NDArray = 23.0,
                   model: Literal['herold-2017',
                                  'herold-2017-fallback',
                                  'buyukurganci-2022',
-                                 'kestin-1978'] = 'herold-2017-fallback'):
+                                 'kestin-1978'] = 'herold-2017-fallback'
+                  ) -> float | npt.NDArray:
     """Returns the viscosity for RT-DC-specific media
 
     Media that are not pure (e.g. ketchup or polymer solutions)
@@ -83,16 +92,16 @@ def get_viscosity(medium: str = "0.49% MC-PBS",
 
     Parameters
     ----------
-    medium: str
+    medium
         The medium to compute the viscosity for; Valid values
         are defined in :const:`KNOWN_MEDIA`.
-    channel_width: float
+    channel_width
         The channel width in µm
-    flow_rate: float
+    flow_rate
         Flow rate in µL/s
-    temperature: float or ndarray
+    temperature
         Temperature in °C
-    model: str
+    model
         The model name to use for computing the medium viscosity.
         For water, this value is ignored, as there is only the
         'kestin-1978' model :cite:`Kestin_1978`. For MC-PBS media,
@@ -101,7 +110,7 @@ def get_viscosity(medium: str = "0.49% MC-PBS",
 
     Returns
     -------
-    viscosity: float or ndarray
+    viscosity
         Viscosity in mPa*s
 
     Notes
@@ -148,21 +157,23 @@ def get_viscosity(medium: str = "0.49% MC-PBS",
     return eta
 
 
-def shear_rate_square_channel(flow_rate, channel_width, flow_index):
+def shear_rate_square_channel(flow_rate: float,
+                              channel_width: float,
+                              flow_index: float) -> float:
     """Returns The wall shear rate of a power law liquid in a squared channel.
 
     Parameters
     ----------
-    flow_rate: float
+    flow_rate
         Flow rate in µL/s
-    channel_width: float
+    channel_width
         The channel width in µm
-    flow_index: float
+    flow_index
         Flow behavior index aka the power law exponent of the shear thinning
 
     Returns
     -------
-    shear_rate: float
+    shear_rate
         Shear rate in 1/s.
     """
     # convert channel width to mm
@@ -176,7 +187,7 @@ def get_viscosity_mc_pbs_buyukurganci_2022(
                         "0.83% MC-PBS"] = "0.49% MC-PBS",
         channel_width: float = 20.0,
         flow_rate: float = 0.16,
-        temperature: float = 23.0):
+        temperature: float = 23.0) -> float | npt.NDArray:
     """Compute viscosity of MC-PBS according to :cite:`Buyukurganci2022`
 
     This viscosity model was derived in :cite:`Buyukurganci2022`
@@ -211,7 +222,7 @@ def get_viscosity_mc_pbs_herold_2017(
         medium: Literal["0.49% MC-PBS", "0.59% MC-PBS"] = "0.49% MC-PBS",
         channel_width: float = 20.0,
         flow_rate: float = 0.16,
-        temperature: float = 23.0):
+        temperature: float = 23.0) -> float | npt.NDArray:
     r"""Compute viscosity of MC-PBS according to :cite:`Herold2017`
 
     Note that all the factors in equation 5.2 in :cite:`Herold2017`
@@ -242,7 +253,8 @@ def get_viscosity_mc_pbs_herold_2017(
     return eta
 
 
-def get_viscosity_water_kestin_1978(temperature: float = 23.0):
+def get_viscosity_water_kestin_1978(
+        temperature: float = 23.0) -> float | npt.NDArray:
     """Compute the viscosity of water according to :cite:`Kestin_1978`"""
     # see equation (15) in Kestin et al, J. Phys. Chem. 7(3) 1978
     check_temperature("'kestin-1978' water", temperature, 0, 40)