Coverage for suppy\utils\_func_wrapper.py: 88%

1from typing import Callable 

2import numpy as np 

3 

4try: 

5 import cupy as cp 

6 

7 NO_GPU = False 

8except ImportError: 

9 cp = np 

10 NO_GPU = True 

11 

12 

13class FuncWrapper: 

14 """ 

15 A callable class for a function that keeps track of the number of times 

16 it is called. 

17 

18 Parameters 

19 ---------- 

20 func : Callable 

21 The function to be wrapped. 

22 args : list 

23 The arguments to be passed to the function. 

24 

25 Attributes 

26 ---------- 

27 func : Callable 

28 The function to be wrapped. 

29 args : list 

30 The arguments to be passed to the function. 

31 fcount : int 

32 The number of times the function has been called. 

33 """ 

34 

35 def __init__(self, func: Callable, args=[]): 

36 self.func = func 

37 self.args = args 

38 self.fcount = 0 

39 self._intermediate_x = None 

40 self._intermediate_value = 0.0 

41 

42 def __call__(self, x): 

43 xp = cp if isinstance(x, cp.ndarray) else np 

44 self.fcount += 1 

45 if ( 

46 self._intermediate_x is not None 

47 and x is not None 

48 and xp.array_equal(x, self._intermediate_x) 

49 ): 

50 return self._intermediate_value 

51 else: 

52 if x is None: # should mainly happen when evaluating perturbation scheme for 

53 self._intermediate_x = None 

54 else: 

55 self._intermediate_x = x.copy() 

56 self._intermediate_value = self.func(x, *self.args) 

57 return self._intermediate_value