Fix typing of numpy arrays (#144)

This PR fixes all annotations of numpy array and makes it explicit what
type of array a function expects/returns.

We were using invalid numpy array annotations in some places and didn't
precise the type of the array in other places. To properly type numpy
array one needs to write this notation `np.ndarray[Any,
np.dtype[np.complex_]]` where the first type is the shape of the array
and the second is its type. This is not practical and is why numpy
introduced `numpy.typing.NDArray` which can be simply used like
`NDArray[np.complex_]`. Because we use this type a lot, I added a type
alias.

roseau/load_flow/converters.py

@@ -11,6 +11,8 @@
 import numpy as np
 import pandas as pd
 
+from roseau.load_flow.typing import ComplexArray
+
 ALPHA = np.exp(2 / 3 * np.pi * 1j)
 """complex: Phasor rotation operator `alpha`, which rotates a phasor vector counterclockwise by 120
 degrees when multiplied by it."""
@@ -28,15 +30,15 @@
 _A_INV = np.linalg.inv(A)
 
 
-def phasor_to_sym(v_abc: Sequence[complex]) -> np.ndarray[complex]:
+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)
     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]) -> np.ndarray[complex]:
+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)
     orig_shape = v_012_array.shape
@@ -107,7 +109,7 @@ def series_phasor_to_sym(s_abc: pd.Series) -> pd.Series:
     return s_012
 
 
-def calculate_voltages(potentials: np.ndarray, phases: str) -> np.ndarray:
+def calculate_voltages(potentials: ComplexArray, phases: str) -> ComplexArray:
     """Calculate the voltages between phases given the potentials of each phase.
 
     Args:

roseau/load_flow/models/branches.py

@@ -8,7 +8,7 @@
 from roseau.load_flow.converters import calculate_voltages
 from roseau.load_flow.models.buses import Bus
 from roseau.load_flow.models.core import Element
-from roseau.load_flow.typing import Id, JsonDict
+from roseau.load_flow.typing import ComplexArray, Id, JsonDict
 from roseau.load_flow.units import Q_, ureg_wraps
 
 logger = logging.getLogger(__name__)
@@ -66,7 +66,7 @@ def __init__(
         self.bus2 = bus2
         self.geometry = geometry
         self._connect(bus1, bus2)
-        self._res_currents: Optional[tuple[np.ndarray, np.ndarray]] = None
+        self._res_currents: Optional[tuple[ComplexArray, ComplexArray]] = None
 
     def __repr__(self) -> str:
         s = f"{type(self).__name__}(id={self.id!r}, phases1={self.phases1!r}, phases2={self.phases2!r}"
@@ -76,16 +76,16 @@ def __repr__(self) -> str:
         s += ")"
         return s
 
-    def _res_currents_getter(self, warning: bool) -> tuple[np.ndarray, np.ndarray]:
+    def _res_currents_getter(self, warning: bool) -> tuple[ComplexArray, ComplexArray]:
         return self._res_getter(value=self._res_currents, warning=warning)
 
     @property
     @ureg_wraps(("A", "A"), (None,), strict=False)
-    def res_currents(self) -> tuple[Q_[np.ndarray], Q_[np.ndarray]]:
+    def res_currents(self) -> tuple[Q_[ComplexArray], Q_[ComplexArray]]:
         """The load flow result of the branch currents (A)."""
         return self._res_currents_getter(warning=True)
 
-    def _res_powers_getter(self, warning: bool) -> tuple[np.ndarray, np.ndarray]:
+    def _res_powers_getter(self, warning: bool) -> tuple[ComplexArray, ComplexArray]:
         cur1, cur2 = self._res_currents_getter(warning)
         pot1, pot2 = self._res_potentials_getter(warning=False)  # we warn on the previous line
         powers1 = pot1 * cur1.conj()
@@ -94,28 +94,28 @@ def _res_powers_getter(self, warning: bool) -> tuple[np.ndarray, np.ndarray]:
 
     @property
     @ureg_wraps(("VA", "VA"), (None,), strict=False)
-    def res_powers(self) -> tuple[Q_[np.ndarray], Q_[np.ndarray]]:
+    def res_powers(self) -> tuple[Q_[ComplexArray], Q_[ComplexArray]]:
         """The load flow result of the branch powers (VA)."""
         return self._res_powers_getter(warning=True)
 
-    def _res_potentials_getter(self, warning: bool) -> tuple[np.ndarray, np.ndarray]:
+    def _res_potentials_getter(self, warning: bool) -> tuple[ComplexArray, ComplexArray]:
         pot1 = self.bus1._get_potentials_of(self.phases1, warning)
         pot2 = self.bus2._get_potentials_of(self.phases2, warning=False)  # we warn on the previous line
         return pot1, pot2
 
     @property
     @ureg_wraps(("V", "V"), (None,), strict=False)
-    def res_potentials(self) -> tuple[Q_[np.ndarray], Q_[np.ndarray]]:
+    def res_potentials(self) -> tuple[Q_[ComplexArray], Q_[ComplexArray]]:
         """The load flow result of the branch potentials (V)."""
         return self._res_potentials_getter(warning=True)
 
-    def _res_voltages_getter(self, warning: bool) -> tuple[np.ndarray, np.ndarray]:
+    def _res_voltages_getter(self, warning: bool) -> tuple[ComplexArray, ComplexArray]:
         pot1, pot2 = self._res_potentials_getter(warning)
         return calculate_voltages(pot1, self.phases1), calculate_voltages(pot2, self.phases2)
 
     @property
     @ureg_wraps(("V", "V"), (None,), strict=False)
-    def res_voltages(self) -> tuple[Q_[np.ndarray], Q_[np.ndarray]]:
+    def res_voltages(self) -> tuple[Q_[ComplexArray], Q_[ComplexArray]]:
         """The load flow result of the branch voltages (V)."""
         return self._res_voltages_getter(warning=True)
 

roseau/load_flow/models/buses.py

@@ -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 Id, JsonDict
+from roseau.load_flow.typing import ComplexArray, Id, JsonDict
 from roseau.load_flow.units import Q_, ureg_wraps
 
 logger = logging.getLogger(__name__)
@@ -79,15 +79,15 @@ def __init__(
         self.min_voltage = min_voltage
         self.max_voltage = max_voltage
 
-        self._res_potentials: Optional[np.ndarray] = None
+        self._res_potentials: Optional[ComplexArray] = None
         self._short_circuits: list[dict[str, Any]] = []
 
     def __repr__(self) -> str:
         return f"{type(self).__name__}(id={self.id!r}, phases={self.phases!r})"
 
     @property
     @ureg_wraps("V", (None,), strict=False)
-    def potentials(self) -> Q_[np.ndarray]:
+    def potentials(self) -> Q_[ComplexArray]:
         """The potentials of the bus (V)."""
         return self._potentials
 
@@ -101,22 +101,22 @@ def potentials(self, value: Sequence[complex]) -> None:
         self._potentials = np.asarray(value, dtype=complex)
         self._invalidate_network_results()
 
-    def _res_potentials_getter(self, warning: bool) -> np.ndarray:
+    def _res_potentials_getter(self, warning: bool) -> ComplexArray:
         return self._res_getter(value=self._res_potentials, warning=warning)
 
     @property
     @ureg_wraps("V", (None,), strict=False)
-    def res_potentials(self) -> Q_[np.ndarray]:
+    def res_potentials(self) -> Q_[ComplexArray]:
         """The load flow result of the bus potentials (V)."""
         return self._res_potentials_getter(warning=True)
 
-    def _res_voltages_getter(self, warning: bool) -> np.ndarray:
+    def _res_voltages_getter(self, warning: bool) -> ComplexArray:
         potentials = np.asarray(self._res_potentials_getter(warning=warning))
         return calculate_voltages(potentials, self.phases)
 
     @property
     @ureg_wraps("V", (None,), strict=False)
-    def res_voltages(self) -> Q_[np.ndarray]:
+    def res_voltages(self) -> Q_[ComplexArray]:
         """The load flow result of the bus voltages (V).
 
         If the bus has a neutral, the voltages are phase-neutral voltages for existing phases in
@@ -130,7 +130,7 @@ def voltage_phases(self) -> list[str]:
         """The phases of the voltages."""
         return calculate_voltage_phases(self.phases)
 
-    def _get_potentials_of(self, phases: str, warning: bool) -> np.ndarray:
+    def _get_potentials_of(self, phases: str, warning: bool) -> ComplexArray:
         """Get the potentials of the given phases."""
         potentials = self._res_potentials_getter(warning)
         return np.array([potentials[self.phases.index(p)] for p in phases])

roseau/load_flow/models/lines/lines.py

@@ -12,7 +12,7 @@
 from roseau.load_flow.models.grounds import Ground
 from roseau.load_flow.models.lines.parameters import LineParameters
 from roseau.load_flow.models.sources import VoltageSource
-from roseau.load_flow.typing import Id, JsonDict
+from roseau.load_flow.typing import ComplexArray, Id, JsonDict
 from roseau.load_flow.units import Q_, ureg_wraps
 
 logger = logging.getLogger(__name__)
@@ -275,13 +275,13 @@ def parameters(self, value: LineParameters) -> None:
 
     @property
     @ureg_wraps("ohm", (None,), strict=False)
-    def z_line(self) -> Q_[np.ndarray]:
+    def z_line(self) -> Q_[ComplexArray]:
         """Impedance of the line in Ohm"""
         return self.parameters._z_line * self._length
 
     @property
     @ureg_wraps("S", (None,), strict=False)
-    def y_shunt(self) -> Q_[np.ndarray]:
+    def y_shunt(self) -> Q_[ComplexArray]:
         """Shunt admittance of the line in Siemens"""
         return self.parameters._y_shunt * self._length
 
@@ -296,33 +296,33 @@ def max_current(self) -> Optional[Q_[float]]:
     def with_shunt(self) -> bool:
         return self.parameters.with_shunt
 
-    def _res_series_values_getter(self, warning: bool) -> tuple[np.ndarray, np.ndarray]:
+    def _res_series_values_getter(self, warning: bool) -> tuple[ComplexArray, ComplexArray]:
         pot1, pot2 = self._res_potentials_getter(warning)  # V
         du_line = pot1 - pot2
         i_line = np.linalg.inv(self.z_line.m_as("ohm")) @ du_line  # Zₗ x Iₗ = ΔU -> I = Zₗ⁻¹ x ΔU
         return du_line, i_line
 
-    def _res_series_currents_getter(self, warning: bool) -> np.ndarray:
+    def _res_series_currents_getter(self, warning: bool) -> ComplexArray:
         _, i_line = self._res_series_values_getter(warning)
         return i_line
 
     @property
     @ureg_wraps("A", (None,), strict=False)
-    def res_series_currents(self) -> Q_[np.ndarray]:
+    def res_series_currents(self) -> Q_[ComplexArray]:
         """Get the current in the series elements of the line (A)."""
         return self._res_series_currents_getter(warning=True)
 
-    def _res_series_power_losses_getter(self, warning: bool) -> np.ndarray:
+    def _res_series_power_losses_getter(self, warning: bool) -> ComplexArray:
         du_line, i_line = self._res_series_values_getter(warning)
         return du_line * i_line.conj()  # Sₗ = ΔU.Iₗ*
 
     @property
     @ureg_wraps("VA", (None,), strict=False)
-    def res_series_power_losses(self) -> Q_[np.ndarray]:
+    def res_series_power_losses(self) -> Q_[ComplexArray]:
         """Get the power losses in the series elements of the line (VA)."""
         return self._res_series_power_losses_getter(warning=True)
 
-    def _res_shunt_values_getter(self, warning: bool) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+    def _res_shunt_values_getter(self, warning: bool) -> tuple[ComplexArray, ComplexArray, ComplexArray, ComplexArray]:
         assert self.with_shunt, "This method only works when there is a shunt"
         assert self.ground is not None
         pot1, pot2 = self._res_potentials_getter(warning)
@@ -333,7 +333,7 @@ def _res_shunt_values_getter(self, warning: bool) -> tuple[np.ndarray, np.ndarra
         i2_shunt = (y_shunt @ pot2 - yg * vg) / 2
         return pot1, pot2, i1_shunt, i2_shunt
 
-    def _res_shunt_currents_getter(self, warning: bool) -> tuple[np.ndarray, np.ndarray]:
+    def _res_shunt_currents_getter(self, warning: bool) -> tuple[ComplexArray, ComplexArray]:
         if not self.with_shunt:
             zeros = np.zeros(len(self.phases), dtype=complex)
             return zeros[:], zeros[:]
@@ -342,30 +342,30 @@ def _res_shunt_currents_getter(self, warning: bool) -> tuple[np.ndarray, np.ndar
 
     @property
     @ureg_wraps(("A", "A"), (None,), strict=False)
-    def res_shunt_currents(self) -> tuple[Q_[np.ndarray], Q_[np.ndarray]]:
+    def res_shunt_currents(self) -> tuple[Q_[ComplexArray], Q_[ComplexArray]]:
         """Get the currents in the shunt elements of the line (A)."""
         return self._res_shunt_currents_getter(warning=True)
 
-    def _res_shunt_power_losses_getter(self, warning: bool) -> np.ndarray:
+    def _res_shunt_power_losses_getter(self, warning: bool) -> ComplexArray:
         if not self.with_shunt:
             return np.zeros(len(self.phases), dtype=complex)
         pot1, pot2, cur1, cur2 = self._res_shunt_values_getter(warning)
         return pot1 * cur1.conj() + pot2 * cur2.conj()
 
     @property
     @ureg_wraps("VA", (None,), strict=False)
-    def res_shunt_power_losses(self) -> Q_[np.ndarray]:
+    def res_shunt_power_losses(self) -> Q_[ComplexArray]:
         """Get the power losses in the shunt elements of the line (VA)."""
         return self._res_shunt_power_losses_getter(warning=True)
 
-    def _res_power_losses_getter(self, warning: bool) -> np.ndarray:
+    def _res_power_losses_getter(self, warning: bool) -> ComplexArray:
         series_losses = self._res_series_power_losses_getter(warning)
         shunt_losses = self._res_shunt_power_losses_getter(warning=False)  # we warn on the previous line
         return series_losses + shunt_losses
 
     @property
     @ureg_wraps("VA", (None,), strict=False)
-    def res_power_losses(self) -> Q_[np.ndarray]:
+    def res_power_losses(self) -> Q_[ComplexArray]:
         """Get the power losses in the line (VA)."""
         return self._res_power_losses_getter(warning=True)
 

roseau/load_flow/models/lines/parameters.py

@@ -8,7 +8,7 @@
 from typing_extensions import Self
 
 from roseau.load_flow.exceptions import RoseauLoadFlowException, RoseauLoadFlowExceptionCode
-from roseau.load_flow.typing import Id, JsonDict
+from roseau.load_flow.typing import ComplexArray, Id, JsonDict
 from roseau.load_flow.units import Q_, ureg_wraps
 from roseau.load_flow.utils import (
     CX,
@@ -41,7 +41,11 @@ class LineParameters(Identifiable, JsonMixin):
 
     @ureg_wraps(None, (None, None, "ohm/km", "S/km", "A"), strict=False)
     def __init__(
-        self, id: Id, z_line: np.ndarray, y_shunt: Optional[np.ndarray] = None, max_current: Optional[float] = None
+        self,
+        id: Id,
+        z_line: ComplexArray,
+        y_shunt: Optional[ComplexArray] = None,
+        max_current: Optional[float] = None,
     ) -> None:
         """LineParameters constructor.
 
@@ -89,12 +93,12 @@ def __eq__(self, other: object) -> bool:
 
     @property
     @ureg_wraps("ohm/km", (None,), strict=False)
-    def z_line(self) -> Q_[np.ndarray]:
+    def z_line(self) -> Q_[ComplexArray]:
         return self._z_line
 
     @property
     @ureg_wraps("S/km", (None,), strict=False)
-    def y_shunt(self) -> Q_[np.ndarray]:
+    def y_shunt(self) -> Q_[ComplexArray]:
         return self._y_shunt
 
     @property
@@ -183,7 +187,7 @@ def _sym_to_zy(
         xpn: Optional[float] = None,
         bn: Optional[float] = None,
         bpn: Optional[float] = None,
-    ) -> tuple[np.ndarray, np.ndarray]:
+    ) -> tuple[ComplexArray, ComplexArray]:
         """Create impedance and admittance matrix from a symmetrical model.
 
         Args:
@@ -364,7 +368,7 @@ def _geometry_to_zy(
         section_neutral: float,
         height: float,
         external_diameter: float,
-    ) -> tuple[np.ndarray, np.ndarray]:
+    ) -> tuple[ComplexArray, ComplexArray]:
         """Create impedance and admittance matrix using a geometric model.
 
         Args: