Coverage for rhodent/typing.py: 94%

1from __future__ import annotations 

2from typing import Generic, Union, TypeVar 

3 

4import numpy as np 

5from numpy.typing import NDArray 

6 

7from gpaw import mpi 

8from gpaw.calculator import GPAW 

9 

10GPAWCalculator = GPAW 

11Communicator = mpi._Communicator 

12ArrayIndex = Union[NDArray[np.int_], list[int], int, slice] 

13 

14 

15DTypeT = TypeVar('DTypeT', bound=np.generic, covariant=True) 

16 

17 

18class ArrayIsOnRootRank(np.ndarray, Generic[DTypeT]): 

19 # numpy 2.5 turned NDArray into a typing alias that cannot be subclassed, 

20 # so the base class is the array type itself and the dtype parameter is 

21 # declared separately. 

22 def __new__(cls): 

23 """ Instances will act as empty numpy arrays """ 

24 return np.ndarray.__new__(cls, (0, )) 

25 

26 

27DistributedArray = Union[NDArray[DTypeT], ArrayIsOnRootRank] 

28Array1D = np.ndarray[tuple[int], np.dtype[DTypeT]] 

29Array2D = np.ndarray[tuple[int, int], np.dtype[DTypeT]] 

30Array3D = np.ndarray[tuple[int, int, int], np.dtype[DTypeT]]