Coverage for src / sdynpy / signal_processing / sdynpy_complex.py: 18%

1# -*- coding: utf-8 -*- 

2""" 

3Functions for dealing with complex numbers 

4""" 

5""" 

6Copyright 2022 National Technology & Engineering Solutions of Sandia, 

7LLC (NTESS). Under the terms of Contract DE-NA0003525 with NTESS, the U.S. 

8Government retains certain rights in this software. 

9 

10This program is free software: you can redistribute it and/or modify 

11it under the terms of the GNU General Public License as published by 

12the Free Software Foundation, either version 3 of the License, or 

13(at your option) any later version. 

14 

15This program is distributed in the hope that it will be useful, 

16but WITHOUT ANY WARRANTY; without even the implied warranty of 

17MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

18GNU General Public License for more details. 

19 

20You should have received a copy of the GNU General Public License 

21along with this program. If not, see <https://www.gnu.org/licenses/>. 

22""" 

23 

24import numpy as np 

25import matplotlib.pyplot as plt 

26 

27 

28def collapse_complex_to_real(vector, axis=-1, preserve_magnitude=False, plot=False, 

29 force_angle=None): 

30 x = np.real(vector) 

31 y = np.imag(vector) 

32 slope = np.sum(x * y, axis=axis, keepdims=True) / np.sum(x * x, axis=axis, keepdims=True) 

33 angle = np.arctan(slope) 

34 if force_angle is not None: 

35 angle[...] = force_angle 

36 rotated_vector = vector * np.exp(-1j * angle) 

37 if plot: 

38 plt.figure('Complex Rotation') 

39 shape = list(vector.shape) 

40 shape[axis] = 1 

41 for key in np.ndindex(*shape): 

42 print(key) 

43 index = list(key) 

44 index[axis] = slice(None) 

45 this_vector = vector[index] 

46 this_rotated_vector = rotated_vector[index] 

47 plt.plot(np.real(this_vector), np.imag(this_vector), 'x') 

48 plt.plot(np.real(this_rotated_vector), np.imag(this_rotated_vector), 'o') 

49 if preserve_magnitude: 

50 return np.sign(np.real(rotated_vector)) * np.abs(rotated_vector) 

51 else: 

52 return np.real(rotated_vector) 

53 

54def fit_complex_angle(vector,axis=-1): 

55 x = np.real(vector) 

56 y = np.imag(vector) 

57 slope = np.sum(x * y, axis=axis, keepdims=True) / np.sum(x * x, axis=axis, keepdims=True) 

58 angle = np.arctan(slope) 

59 return angle 

60 

61def rotate_vector(vector,angle): 

62 return vector * np.exp(1j * angle)