Improve the typing of pint wrapped method

roseau/load_flow/converters.py

@@ -23,7 +23,7 @@
         [1, ALPHA**2, ALPHA],
         [1, ALPHA, ALPHA**2],
     ],
-    dtype=complex,
+    dtype=np.complex128,
 )
 """numpy.ndarray[complex]: "A" matrix: transformation matrix from phasor to symmetrical components."""
 
@@ -32,15 +32,15 @@
 
 def phasor_to_sym(v_abc: Sequence[complex]) -> ComplexArray:
     """Compute the symmetrical components `(0, +, -)` from the phasor components `(a, b, c)`."""
-    v_abc_array = np.asarray(v_abc)
+    v_abc_array = np.array(v_abc)
     orig_shape = v_abc_array.shape
     v_012 = _A_INV @ v_abc_array.reshape((3, 1))
     return v_012.reshape(orig_shape)
 
 
 def sym_to_phasor(v_012: Sequence[complex]) -> ComplexArray:
     """Compute the phasor components `(a, b, c)` from the symmetrical components `(0, +, -)`."""
-    v_012_array = np.asarray(v_012)
+    v_012_array = np.array(v_012)
     orig_shape = v_012_array.shape
     v_abc = A @ v_012_array.reshape((3, 1))
     return v_abc.reshape(orig_shape)
@@ -124,13 +124,13 @@ def calculate_voltages(potentials: ComplexArray, phases: str) -> ComplexArray:
         Otherwise, the voltages are Phase-Phase.
 
     Example:
-        >>> potentials = 230 * np.array([1, np.exp(-2j*np.pi/3), np.exp(2j*np.pi/3), 0], dtype=complex)
+        >>> potentials = 230 * np.array([1, np.exp(-2j*np.pi/3), np.exp(2j*np.pi/3), 0], dtype=np.complex128)
         >>> calculate_voltages(potentials, "abcn")
         array([ 230.  +0.j        , -115.-199.18584287j, -115.+199.18584287j])
-        >>> potentials = np.array([230, 230 * np.exp(-2j*np.pi/3)], dtype=complex)
+        >>> potentials = np.array([230, 230 * np.exp(-2j*np.pi/3)], dtype=np.complex128)
         >>> calculate_voltages(potentials, "ab")
         array([345.+199.18584287j])
-        >>> calculate_voltages(np.array([230, 0], dtype=complex), "an")
+        >>> calculate_voltages(np.array([230, 0], dtype=np.complex128), "an")
         array([230.+0.j])
     """
     assert len(potentials) == len(phases), "Number of potentials must match number of phases."

roseau/load_flow/models/branches.py

@@ -140,8 +140,8 @@ def to_dict(self, *, _lf_only: bool = False) -> JsonDict:
         return res
 
     def results_from_dict(self, data: JsonDict) -> None:
-        currents1 = np.array([complex(i[0], i[1]) for i in data["currents1"]], dtype=complex)
-        currents2 = np.array([complex(i[0], i[1]) for i in data["currents2"]], dtype=complex)
+        currents1 = np.array([complex(i[0], i[1]) for i in data["currents1"]], dtype=np.complex128)
+        currents2 = np.array([complex(i[0], i[1]) for i in data["currents2"]], dtype=np.complex128)
         self._res_currents = (currents1, currents2)
 
     def _results_to_dict(self, warning: bool) -> JsonDict:

roseau/load_flow/models/buses.py

@@ -1,6 +1,6 @@
 import logging
-from collections.abc import Iterator, Sequence
-from typing import TYPE_CHECKING, Any, Optional
+from collections.abc import Iterator
+from typing import TYPE_CHECKING, Any, Optional, Union
 
 import numpy as np
 import pandas as pd
@@ -10,7 +10,7 @@
 from roseau.load_flow.converters import calculate_voltage_phases, calculate_voltages, phasor_to_sym
 from roseau.load_flow.exceptions import RoseauLoadFlowException, RoseauLoadFlowExceptionCode
 from roseau.load_flow.models.core import Element
-from roseau.load_flow.typing import ComplexArray, Id, JsonDict
+from roseau.load_flow.typing import ComplexArray, ComplexArrayLike1D, Id, JsonDict
 from roseau.load_flow.units import Q_, ureg_wraps
 
 logger = logging.getLogger(__name__)
@@ -36,7 +36,7 @@ def __init__(
         *,
         phases: str,
         geometry: Optional[Point] = None,
-        potentials: Optional[Sequence[complex]] = None,
+        potentials: Optional[ComplexArrayLike1D] = None,
         min_voltage: Optional[float] = None,
         max_voltage: Optional[float] = None,
         **kwargs: Any,
@@ -57,15 +57,20 @@ def __init__(
                 x-y coordinates of the bus.
 
             potentials:
-                An optional list of initial potentials of each phase of the bus.
+                An optional array-like of initial potentials of each phase of the bus. If given,
+                these potentials are used as the starting point of the load flow computation.
+                Either complex values (V) or a :data:`Quantity <roseau.load_flow.units.Q_>` of
+                complex values.
 
             min_voltage:
                 An optional minimum voltage of the bus (V). It is not used in the load flow.
                 It must be a phase-neutral voltage if the bus has a neutral, phase-phase otherwise.
+                Either a float (V) or a :data:`Quantity <roseau.load_flow.units.Q_>` of float.
 
             max_voltage:
                 An optional maximum voltage of the bus (V). It is not used in the load flow.
                 It must be a phase-neutral voltage if the bus has a neutral, phase-phase otherwise.
+                Either a float (V) or a :data:`Quantity <roseau.load_flow.units.Q_>` of float.
         """
         super().__init__(id, **kwargs)
         self._check_phases(id, phases=phases)
@@ -88,17 +93,17 @@ def __repr__(self) -> str:
     @property
     @ureg_wraps("V", (None,), strict=False)
     def potentials(self) -> Q_[ComplexArray]:
-        """The potentials of the bus (V)."""
+        """An array of initial potentials of the bus (V)."""
         return self._potentials
 
     @potentials.setter
     @ureg_wraps(None, (None, "V"), strict=False)
-    def potentials(self, value: Sequence[complex]) -> None:
+    def potentials(self, value: ComplexArrayLike1D) -> None:
         if len(value) != len(self.phases):
             msg = f"Incorrect number of potentials: {len(value)} instead of {len(self.phases)}"
             logger.error(msg)
             raise RoseauLoadFlowException(msg=msg, code=RoseauLoadFlowExceptionCode.BAD_POTENTIALS_SIZE)
-        self._potentials = np.asarray(value, dtype=complex)
+        self._potentials = np.array(value, dtype=np.complex128)
         self._invalidate_network_results()
 
     def _res_potentials_getter(self, warning: bool) -> ComplexArray:
@@ -111,7 +116,7 @@ def res_potentials(self) -> Q_[ComplexArray]:
         return self._res_potentials_getter(warning=True)
 
     def _res_voltages_getter(self, warning: bool) -> ComplexArray:
-        potentials = np.asarray(self._res_potentials_getter(warning=warning))
+        potentials = np.array(self._res_potentials_getter(warning=warning))
         return calculate_voltages(potentials, self.phases)
 
     @property
@@ -142,7 +147,7 @@ def min_voltage(self) -> Optional[Q_[float]]:
 
     @min_voltage.setter
     @ureg_wraps(None, (None, "V"), strict=False)
-    def min_voltage(self, value: Optional[float]) -> None:
+    def min_voltage(self, value: Optional[Union[float, Q_[float]]]) -> None:
         if value is not None and self._max_voltage is not None and value > self._max_voltage:
             msg = (
                 f"Cannot set min voltage of bus {self.id!r} to {value} V as it is higher than its "
@@ -161,7 +166,7 @@ def max_voltage(self) -> Optional[Q_[float]]:
 
     @max_voltage.setter
     @ureg_wraps(None, (None, "V"), strict=False)
-    def max_voltage(self, value: Optional[float]) -> None:
+    def max_voltage(self, value: Optional[Union[float, Q_[float]]]) -> None:
         if value is not None and self._min_voltage is not None and value < self._min_voltage:
             msg = (
                 f"Cannot set max voltage of bus {self.id!r} to {value} V as it is lower than its "
@@ -334,7 +339,7 @@ def to_dict(self, *, _lf_only: bool = False) -> JsonDict:
         return res
 
     def results_from_dict(self, data: JsonDict) -> None:
-        self._res_potentials = np.array([complex(v[0], v[1]) for v in data["potentials"]], dtype=complex)
+        self._res_potentials = np.array([complex(v[0], v[1]) for v in data["potentials"]], dtype=np.complex128)
 
     def _results_to_dict(self, warning: bool) -> JsonDict:
         return {

roseau/load_flow/models/lines/lines.py

@@ -1,6 +1,6 @@
 import logging
 import warnings
-from typing import Any, Optional
+from typing import Any, Optional, Union
 
 import numpy as np
 from shapely import LineString, Point
@@ -144,7 +144,7 @@ def __init__(
         bus2: Bus,
         *,
         parameters: LineParameters,
-        length: float,
+        length: Union[float, Q_[float]],
         phases: Optional[str] = None,
         ground: Optional[Ground] = None,
         geometry: Optional[LineString] = None,
@@ -231,7 +231,7 @@ def length(self) -> Q_[float]:
 
     @length.setter
     @ureg_wraps(None, (None, "km"), strict=False)
-    def length(self, value: float) -> None:
+    def length(self, value: Union[float, Q_[float]]) -> None:
         if value <= 0:
             msg = f"A line length must be greater than 0. {value:.2f} km provided."
             logger.error(msg)
@@ -335,7 +335,7 @@ def _res_shunt_values_getter(self, warning: bool) -> tuple[ComplexArray, Complex
 
     def _res_shunt_currents_getter(self, warning: bool) -> tuple[ComplexArray, ComplexArray]:
         if not self.with_shunt:
-            zeros = np.zeros(len(self.phases), dtype=complex)
+            zeros = np.zeros(len(self.phases), dtype=np.complex128)
             return zeros[:], zeros[:]
         _, _, cur1, cur2 = self._res_shunt_values_getter(warning)
         return cur1, cur2
@@ -348,7 +348,7 @@ def res_shunt_currents(self) -> tuple[Q_[ComplexArray], Q_[ComplexArray]]:
 
     def _res_shunt_power_losses_getter(self, warning: bool) -> ComplexArray:
         if not self.with_shunt:
-            return np.zeros(len(self.phases), dtype=complex)
+            return np.zeros(len(self.phases), dtype=np.complex128)
         pot1, pot2, cur1, cur2 = self._res_shunt_values_getter(warning)
         return pot1 * cur1.conj() + pot2 * cur2.conj()