Coverage for fiqus/plotters/PlotPythonCCT.py: 13%

1import os 

2from operator import le, ge 

3 

4import numpy as np 

5import pandas as pd 

6import matplotlib.pyplot as plt 

7from matplotlib.pyplot import cm 

8from mpl_toolkits.mplot3d.art3d import Poly3DCollection 

9 

10from fiqus.geom_generators.GeometryCCT import Winding, FQPL 

11 

12 

13class PlotPythonCCT: 

14 def __init__(self, fdm, verbose=True): 

15 """ 

16 Class to cct models postprocessing 

17 :param fdm: FiQuS data model 

18 :param verbose: If True more information is printed in python console. 

19 """ 

20 self.cctdm = fdm.magnet 

21 self.model_folder = os.path.join(os.getcwd()) 

22 self.magnet_name = fdm.general.magnet_name 

23 

24 self.verbose = verbose 

25 # self.gu = GmshUtils(self.model_folder, self.verbose) 

26 # self.gu.initialize() 

27 

28 def plot_elements_file(self, lfes=None): 

29 field_map_3D_csv = 'field_map_3D.csv' 

30 df = pd.read_csv(field_map_3D_csv, delimiter=',', engine='python') 

31 sAvePositions = df['sAvePositions [m]'].to_numpy(dtype='float')[:-1] # [:-1] to remove last nan 

32 ph_order = df['nodeToPhysicalTurn [-]'].to_numpy(dtype='int')[:-1] # physical turns, [:-1] to remove last nan 

33 el_order = df['nodeToHalfTurn [-]'].to_numpy(dtype='int')[:-1] # electric order turns, [:-1] to remove last nan 

34 x3Dsurf_1 = df['x3Dsurf 1 [m]'].to_numpy(dtype='float') 

35 y3Dsurf_1 = df['y3Dsurf 1 [m]'].to_numpy(dtype='float') 

36 z3Dsurf_1 = df['z3Dsurf 1 [m]'].to_numpy(dtype='float') 

37 el_order_unique = np.unique(el_order) 

38 

39 per_turn = False 

40 

41 fig = plt.figure(figsize=(16, 10)) 

42 ax = fig.add_subplot(111, projection='3d', proj_type='ortho') 

43 

44 if per_turn: 

45 colors = cm.rainbow(np.linspace(0, 1, np.shape(el_order_unique)[0])) 

46 for idx_turn, el_u in enumerate(el_order_unique.tolist()): 

47 indexes = np.where(el_u == el_order)[0] 

48 min_idx=np.min(indexes) 

49 max_idx=np.max(indexes+2) 

50 ax.plot(x3Dsurf_1[min_idx:max_idx], y3Dsurf_1[min_idx:max_idx], z3Dsurf_1[min_idx:max_idx], color=colors[idx_turn]) 

51 else: 

52 ax.plot(x3Dsurf_1, y3Dsurf_1, z3Dsurf_1) 

53 # colors = cm.rainbow(np.linspace(0, 1, np.shape(x3Dsurf_1)[0])) 

54 # # ax.plot(x3Dsurf_1, y3Dsurf_1, z3Dsurf_1) 

55 

56 # for idx in range(np.shape(x3Dsurf_1)[0]-1): 

57 # ax.plot([x3Dsurf_1[idx], x3Dsurf_1[idx+1]], [y3Dsurf_1[idx], y3Dsurf_1[idx+1]], [z3Dsurf_1[idx], z3Dsurf_1[idx+1]], color=colors[idx]) 

58 # xmin, xmax, ymin, ymax, zmin, zmax = (0, 0, 0, 0, 0, 0) 

59 # #colors = cm.rainbow(np.linspace(0, 1, len(objects_list))) 

60 # ax.set_box_aspect((xmax - xmin, ymax - ymin, zmax - zmin)) # set axis to data aspect ratio 

61 # ax.set_xlim([xmin, xmax]) 

62 # ax.set_ylim([ymin, ymax]) 

63 # ax.set_zlim([zmin, zmax]) 

64 ax.set_xlabel('x (m)') 

65 ax.set_ylabel('y (m)') 

66 ax.set_zlabel('z (m)') 

67 ax.view_init(elev=90, azim=0, vertical_axis='x') 

68 plt.show() 

69 pass 

70 

71 def plot_elements_computed(self, lfes=None): 

72 """ 

73 Makes a plot oh hexahedra by calculating 'on the fly' what is defined in the yaml data model for cct magnet. 

74 :param lfes: list with numbers of elements to put text labels on. For two windings an example would be [[3], [9]] to plot 3rd hex in the first winding and 9th in the second 

75 :return: Nothing. A matplotlib graph pops up on the screen. 

76 """ 

77 

78 windings = [] 

79 fqpls = [] 

80 for ww, _ in enumerate(self.cctdm.geometry.windings.names): 

81 winding_obj = Winding(self.cctdm, ww, post_proc=True) 

82 windings.append(winding_obj) 

83 for ff, _ in enumerate(self.cctdm.geometry.fqpls.names): 

84 fqpl_obj = FQPL(self.cctdm, ff) 

85 fqpls.append(fqpl_obj) 

86 objects_list = windings + fqpls 

87 if not lfes: 

88 lfes = [[] for _ in objects_list] # make it work if lfes is not supplied as input. 

89 

90 def val_check(var, array, oper, func): 

91 new = func(array) 

92 if oper(new, var): 

93 return new 

94 else: 

95 return var 

96 

97 fig = plt.figure(figsize=(16, 10)) 

98 ax = fig.add_subplot(111, projection='3d', proj_type='ortho') 

99 xmin, xmax, ymin, ymax, zmin, zmax = (0, 0, 0, 0, 0, 0) 

100 colors = cm.rainbow(np.linspace(0, 1, len(objects_list))) 

101 for obj, lfe, fc in zip(objects_list, lfes, colors): 

102 xs = [] 

103 ys = [] 

104 zs = [] 

105 hexes = obj.hexes 

106 vts = obj.vertices_to_surf 

107 for elem_num, points_dict in hexes.items(): 

108 list_of_coor_tuples = [] 

109 points_dict.pop('ct', None) # removing turn from dict 

110 for p_num, p_coords in points_dict.items(): 

111 

112 list_of_coor_tuples.append((p_coords['x'], p_coords['y'], p_coords['z'])) 

113 if elem_num in lfe: 

114 ax.scatter(p_coords['x'], p_coords['y'], p_coords['z']) 

115 ax.text(p_coords['x'], p_coords['y'], p_coords['z'], f'{p_num}', zdir='z') 

116 for coor_acc, coor_key in zip([xs, ys, zs], ['x', 'y', 'z']): 

117 coor_acc.append(p_coords[coor_key]) 

118 poly3d = [[list_of_coor_tuples[vts[ix][iy]] for iy in range(len(vts[0]))] for ix in range(len(vts))] 

119 ax.add_collection3d(Poly3DCollection(poly3d, edgecolors='k', facecolors=fc, linewidths=0.5, alpha=0.5)) 

120 xmin = val_check(xmin, xs, le, np.min) 

121 xmax = val_check(xmax, xs, ge, np.max) 

122 ymin = val_check(ymin, ys, le, np.min) 

123 ymax = val_check(ymax, ys, ge, np.max) 

124 zmin = val_check(zmin, zs, le, np.min) 

125 zmax = val_check(zmax, zs, ge, np.max) 

126 ax.set_box_aspect((xmax - xmin, ymax - ymin, zmax - zmin)) # set axis to data aspect ratio 

127 ax.set_xlim([xmin, xmax]) 

128 ax.set_ylim([ymin, ymax]) 

129 ax.set_zlim([zmin, zmax]) 

130 ax.set_xlabel('x (m)') 

131 ax.set_ylabel('y (m)') 

132 ax.set_zlabel('z (m)') 

133 ax.view_init(elev=90, azim=0, vertical_axis='x') 

134 plt.show()