# Ticket #10775: streamlines.py

File streamlines.py, 6.6 KB (added by , 11 years ago) |
---|

Line | |
---|---|

1 | """ |

2 | Copyright (c) 2011 Raymond Speth. |

3 | |

4 | Permission is hereby granted, free of charge, to any person obtaining a copy |

5 | of this software and associated documentation files (the "Software"), to deal |

6 | in the Software without restriction, including without limitation the rights |

7 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell |

8 | copies of the Software, and to permit persons to whom the Software is |

9 | furnished to do so, subject to the following conditions: |

10 | |

11 | The above copyright notice and this permission notice shall be included in |

12 | all copies or substantial portions of the Software. |

13 | |

14 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR |

15 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, |

16 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE |

17 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER |

18 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, |

19 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN |

20 | THE SOFTWARE. |

21 | """ |

22 | |

23 | import numpy as np |

24 | import matplotlib.pyplot as plt |

25 | import matplotlib as mpl |

26 | |

27 | |

28 | class Streamlines(object): |

29 | def __init__(self, X, Y, U, V, res=0.125, spacing=4, maxLen=50, detectLoops=True): |

30 | """ |

31 | Compute a set of streamlines covering the given velocity field. |

32 | |

33 | X and Y - 1D or 2D (e.g. generated by np.meshgrid) arrays of the |

34 | grid points. The mesh spacing is assumed to be uniform |

35 | in each dimension. |

36 | U and V - 2D arrays of the velocity field. |

37 | res - Sets the distance between successive points in each |

38 | streamline (same units as X and Y) |

39 | spacing - Sets the minimum density of streamlines, in grid points. |

40 | maxLen - The maximum length of an individual streamline segment. |

41 | detectLoops - Determines whether an attempt is made to stop extending |

42 | a given streamline before reaching maxLen points if |

43 | it forms a closed loop or reaches a velocity node. |

44 | |

45 | Plots are generated with the 'plot' or 'plotArrows' methods. |

46 | """ |

47 | |

48 | self.spacing = spacing |

49 | self.detectLoops = detectLoops |

50 | self.maxLen = maxLen |

51 | self.res = res |

52 | |

53 | xa = np.asanyarray(X) |

54 | ya = np.asanyarray(Y) |

55 | self.x = xa if xa.ndim == 1 else xa[0] |

56 | self.y = ya if ya.ndim == 1 else ya[:,0] |

57 | self.u = U |

58 | self.v = V |

59 | self.dx = (self.x[-1]-self.x[0])/(self.x.size-1) # assume a regular grid |

60 | self.dy = (self.y[-1]-self.y[0])/(self.y.size-1) # assume a regular grid |

61 | self.dr = self.res * np.sqrt(self.dx * self.dy) |

62 | |

63 | # marker for which regions have contours |

64 | self.used = np.zeros(self.u.shape, dtype=bool) |

65 | self.used[0] = True |

66 | self.used[-1] = True |

67 | self.used[:,0] = True |

68 | self.used[:,-1] = True |

69 | |

70 | # Don't try to compute streamlines in regions where there is no velocity data |

71 | for i in range(self.x.size): |

72 | for j in range(self.y.size): |

73 | if self.u[j,i] == 0.0 and self.v[j,i] == 0.0: |

74 | self.used[j,i] = True |

75 | |

76 | # Make the streamlines |

77 | self.streamlines = [] |

78 | while not self.used.all(): |

79 | nz = np.transpose(np.logical_not(self.used).nonzero()) |

80 | # Make a streamline starting at the first unrepresented grid point |

81 | self.streamlines.append(self._makeStreamline(self.x[nz[0][1]], |

82 | self.y[nz[0][0]])) |

83 | |

84 | def plot(self, lw=1, ax=None): |

85 | """ |

86 | Draw the computed streamline segments. |

87 | |

88 | Optional keyword arguments: |

89 | lw - line width |

90 | ax - axes to use for plotting |

91 | |

92 | """ |

93 | if ax is None: |

94 | ax = plt.gca() |

95 | |

96 | for streamline in self.streamlines: |

97 | ax.plot(streamline[0], streamline[1], 'k', lw=lw) |

98 | |

99 | ax.axis('tight') |

100 | |

101 | return ax |

102 | |

103 | def plotArrows(self, lw=1, ax=None, size=16): |

104 | """ |

105 | Draw the computed streamline segments with arrows indicating flow direction. |

106 | |

107 | Optional keyword arguments: |

108 | lw - line width |

109 | size - size of the arrow head |

110 | ax - axes to use for plotting |

111 | |

112 | """ |

113 | if ax is None: |

114 | ax = plt.gca() |

115 | |

116 | for streamline in self.streamlines: |

117 | path = mpl.path.Path(np.asarray((streamline[0], streamline[1])).T) |

118 | patch = mpl.patches.FancyArrowPatch(path=path, arrowstyle='->', |

119 | mutation_scale=arrowSize, lw=lw) |

120 | ax.add_patch(patch) |

121 | |

122 | ax.axis('tight') |

123 | |

124 | return ax |

125 | |

126 | def _interp(self, x, y): |

127 | """ Compute the velocity at point (x,y) """ |

128 | i = (x-self.x[0])/self.dx |

129 | ai = i % 1 |

130 | |

131 | j = (y-self.y[0])/self.dy |

132 | aj = j % 1 |

133 | |

134 | # Bilinear interpolation |

135 | u = (self.u[j,i]*(1-ai)*(1-aj) + |

136 | self.u[j,i+1]*ai*(1-aj) + |

137 | self.u[j+1,i]*(1-ai)*aj + |

138 | self.u[j+1,i+1]*ai*aj) |

139 | |

140 | v = (self.v[j,i]*(1-ai)*(1-aj) + |

141 | self.v[j,i+1]*ai*(1-aj) + |

142 | self.v[j+1,i]*(1-ai)*aj + |

143 | self.v[j+1,i+1]*ai*aj) |

144 | |

145 | self.used[j:j+self.spacing,i:i+self.spacing] = True |

146 | |

147 | return u,v |

148 | |

149 | def _makeStreamline(self, x0, y0): |

150 | """ |

151 | Compute a streamline extending in both directions from the given point. |

152 | """ |

153 | |

154 | sx, sy = self._makeHalfStreamline(x0, y0, 1) # forwards |

155 | rx, ry = self._makeHalfStreamline(x0, y0, -1) # backwards |

156 | |

157 | rx.reverse() |

158 | ry.reverse() |

159 | |

160 | return rx+[x0]+sx, ry+[y0]+sy |

161 | |

162 | def _makeHalfStreamline(self, x0, y0, sign): |

163 | """ |

164 | Compute a streamline extending in one direction from the given point. |

165 | """ |

166 | |

167 | xmin = self.x[0] |

168 | xmax = self.x[-1] |

169 | ymin = self.y[0] |

170 | ymax = self.y[-1] |

171 | |

172 | sx = [] |

173 | sy = [] |

174 | |

175 | x = x0 |

176 | y = y0 |

177 | i = 0 |

178 | while xmin < x < xmax and ymin < y < ymax: |

179 | u, v = self._interp(x, y) |

180 | theta = np.arctan2(v,u) |

181 | |

182 | x += sign * self.dr * np.cos(theta) |

183 | y += sign * self.dr * np.sin(theta) |

184 | sx.append(x) |

185 | sy.append(y) |

186 | |

187 | i += 1 |

188 | |

189 | if self.detectLoops and i % 10 == 0 and self._detectLoop(sx, sy): |

190 | break |

191 | |

192 | if i > self.maxLen / 2: |

193 | break |

194 | |

195 | return sx, sy |

196 | |

197 | def _detectLoop(self, xVals, yVals): |

198 | """ Detect closed loops and nodes in a streamline. """ |

199 | x = xVals[-1] |

200 | y = yVals[-1] |

201 | D = np.array([np.hypot(x-xj, y-yj) |

202 | for xj,yj in zip(xVals[:-1],yVals[:-1])]) |

203 | return (D < 0.9 * self.dr).any() |