Coverage for pyhiperta/utils/nanify.py: 100%

1import numpy as np 

2 

3 

4def nanify(a, condition): 

5 """Set a or values of a to np.nan where `condition` is true. 

6 

7 `a` can either be a number, then `nanify` returns np.nan if condition is true, or an array, 

8 in which case `nanify` returns the array with the values where condition is true set to np.nan. 

9 This function main purpose is to provide a convinient way to perform `a[condition] = np.nan` 

10 even when `a` is a single number and not an array. 

11 

12 Parameters 

13 ---------- 

14 a : np.scalar or np.ndarray 

15 The value to convert to np.nan is condition is true. 

16 b : bool or np.ndarray 

17 The boolean(s) controlling where to set values to np.nan. If an array it typically is a 

18 numpy condition array, see examples. 

19 

20 Returns 

21 ------- 

22 np.nan if a is a scalar, otherwise `a` after `a[condition] = np.nan`. 

23 

24 Examples 

25 -------- 

26 >>> a = 12 

27 >>> nanify(a, a < 13) 

28 np.nan 

29 >>> a = np.array([1, 2, 3]) 

30 >>> nanify(a, a < 2) 

31 array([np.nan, 2, 3]) 

32 """ 

33 if np.isscalar(a): 

34 if condition: 

35 return np.float32(np.nan) # this is subscriptable, while np.nan[..., np.newaxis] raises an exception! 

36 else: 

37 return a 

38 a[condition] = np.nan 

39 return a