Coverage for src/evutils/transforms/functional/_spatial.py: 93%

1"""Spatial functional transforms (flips, jitter).""" 

2import math 

3 

4import numpy as np 

5from evutils.jit import lazy_njit 

6 

7from ._common import apply_kernel 

8 

9@lazy_njit 

10def _flip_lr_jit(t, x, y, p, width: int): 

11 """Mirror x about the sensor width: ``x' = width - 1 - x``.""" 

12 # Promote through int64 so the uint16 subtraction can't wrap around. 

13 new_x = (np.int64(width) - 1 - x.astype(np.int64)).astype(np.uint16) 

14 return t, new_x, y, p 

15 

16def flip_lr(events, sensor_size): 

17 """Flip events horizontally: ``x' = width - 1 - x``. 

18 

19 Parameters 

20 ---------- 

21 events : np.ndarray or EventArray 

22 Events to flip. 

23 sensor_size : tuple 

24 ``(W, H, P)`` sensor size; only the width ``W`` is used. 

25 

26 Returns 

27 ------- 

28 np.ndarray or EventArray 

29 Horizontally flipped events, in their original container type. 

30 """ 

31 return apply_kernel(events, _flip_lr_jit, int(sensor_size[0])) 

32 

33@lazy_njit 

34def _spatial_jitter_jit(t, x, y, p, width: int, height: int, 

35 var_x: float, var_y: float, sigma_xy: float, 

36 clip_outliers: bool): 

37 """Add correlated Gaussian noise to ``x``/``y``. 

38 

39 Numba has no ``multivariate_normal``, so we sample two independent standard 

40 normals and correlate them with the Cholesky factor of 

41 ``[[var_x, sigma_xy], [sigma_xy, var_y]]``. 

42 """ 

43 n = len(x) 

44 l11 = math.sqrt(var_x) 

45 l21 = sigma_xy / l11 if l11 > 0.0 else 0.0 

46 under = var_y - l21 * l21 

47 l22 = math.sqrt(under) if under > 0.0 else 0.0 

48 

49 z1 = np.random.normal(0.0, 1.0, n) 

50 z2 = np.random.normal(0.0, 1.0, n) 

51 new_x = x.astype(np.float64) + l11 * z1 

52 new_y = y.astype(np.float64) + (l21 * z1 + l22 * z2) 

53 

54 if clip_outliers: 

55 keep = (new_x >= 0) & (new_x < width) & (new_y >= 0) & (new_y < height) 

56 t, p = t[keep], p[keep] 

57 new_x, new_y = new_x[keep], new_y[keep] 

58 

59 return t, new_x.astype(np.uint16), new_y.astype(np.uint16), p 

60 

61def spatial_jitter(events, sensor_size, var_x=1.0, var_y=1.0, sigma_xy=0.0, 

62 clip_outliers=False): 

63 """Jitter event coordinates with a 2D Gaussian. 

64 

65 Parameters 

66 ---------- 

67 events : np.ndarray or EventArray 

68 Events to jitter. 

69 sensor_size : tuple 

70 ``(W, H, P)`` sensor size, used for clipping. 

71 var_x, var_y : float, optional 

72 Variances of the jitter in x and y. Default ``1.0``. 

73 sigma_xy : float, optional 

74 Off-diagonal covariance (diagonal skew). Default ``0.0``. 

75 clip_outliers : bool, optional 

76 Drop events jittered outside the sensor instead of casting them back 

77 (which would wrap the unsigned coordinate). Default ``False``. 

78 

79 Returns 

80 ------- 

81 np.ndarray or EventArray 

82 Jittered events, in their original container type. 

83 """ 

84 return apply_kernel(events, _spatial_jitter_jit, int(sensor_size[0]), 

85 int(sensor_size[1]), float(var_x), float(var_y), 

86 float(sigma_xy), bool(clip_outliers))