Coverage for rhodent/typing.py: 94%

1from __future__ import annotations 

2from typing import 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(NDArray[DTypeT]): 

19 def __new__(cls): 

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

21 return NDArray.__new__(cls, (0, )) 

22 

23 

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

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

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

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