Coverage for tcvx21/grillix_post/lineouts/parametric_spline_m.py: 90%

1""" 

2Parametric spline interpolator 

3 

4Useful when you need to interpolate a curve of values, like points along a flux surface or a chord 

5""" 

6import warnings 

7import numpy as np 

8from scipy.interpolate import splprep, splev 

9 

10 

11class ParametricSpline: 

12 """ 

13 A wrapper class around slprep and splev from scipy.interpolate 

14 Uses cubic spline interpolation 

15 """ 

16 

17 def __init__(self, sample_points: list, t_points: np.array = None, 

18 smoothing: float = 0.0, periodic: bool = False, order: int = 3): 

19 """ 

20 Calculates the knots and coefficients for a parametric cubic spline interpolator 

21 

22 sample_points should be a list of 1D arrays, i.e. [x_points, y_points, ...] 

23 If t_points is given, then the arrays given in sample points should be parametrised by t 

24 

25 If smoothing if > 0.0, then the input points will be smoothed. (Typically smoothing required << 1, recommend 

26 to check the fit) 

27 

28 Order should be less than the number of sample points - 1 

29 """ 

30 

31 sample_length = None 

32 for sample_array in sample_points: 

33 if sample_length is None: 

34 # Use the first array to set the length 

35 sample_length = len(sample_array) 

36 else: 

37 assert len(sample_array) == sample_length 

38 

39 assert sample_length > order, f"Not enough sample points ({sample_length}) for an order ({order}) " \ 

40 "ParametricSpline" 

41 

42 if order < 3 and smoothing > 0.0: 

43 warnings.warn(UserWarning("Should not use smoothing for order < 3 in ParametricSpline")) 

44 

45 if t_points is None: 

46 tck, self.t_points = splprep(sample_points, s=smoothing, per=(1 if periodic else 0), k=order) 

47 else: 

48 tck, self.t_points = splprep(sample_points, u=t_points, s=smoothing, per=(1 if periodic else 0), k=order) 

49 

50 self.t_min, self.t_max = np.min(self.t_points), np.max(self.t_points) 

51 

52 self.knots, self.coeffs, self.order = tck 

53 

54 def __call__(self, t_evaluations: np.array): 

55 """ 

56 Returns the spline evaluations at points given by t_evaluations 

57 N.b. if no t_points are provided to init, then the given sample points are assumed to be parametrised between 0 and 1 

58 """ 

59 # Interpolation only! 

60 assert t_evaluations.min() >= self.t_min and t_evaluations.max() <= self.t_max, f"Requested points in the range\ 

61 {t_evaluations.min()}, {t_evaluations.max()}, which is outside the interval {self.t_min}, {self.t_max}" 

62 

63 return splev(t_evaluations, (self.knots, self.coeffs, self.order))