mitgcm
Analysis of MITgcm output using python
streamlines.py
Go to the documentation of this file.
1 """!
2 A collection of functions for streamlines and related shennanigans.
3 
4 Streamlines
5 ====================
6 
7 Functions for creating and analysing streamlines.
8 
9 These functions work on cartesian and spherical polar grids - other grids, such as cubed sphere, are not supported.
10 
11 Streamlines are defined to be the path that a parcel of fluid would follow when advected by an unchanging velocity field - the velocities are constant in time.
12 
13 Streaklines are defined as the path that a parcel of fluid would follow in the actual flow - the velocity fields change with time.
14 
15 """
16 
17 
18 import numpy as np
19 import netCDF4
20 import numba
21 import copy
22 import scipy.interpolate
23 from . import functions
24 
25 def stream2(u,v,
26  startx,starty,
27  grid_object,
28  t_max=2592000,delta_t=3600,
29  u_grid_loc='U',v_grid_loc='V'):
30  """!A two-dimensional streamline solver. The velocity fields *must* be two dimensional and not vary in time.
31  X_grid_loc variables specify where the field "X" is located on the C-grid. Possibles options are, U, V, T and Zeta.
32  """
33 
34  if u_grid_loc == 'U':
35  x_u = grid_object['Xp1'][:]
36  y_u = grid_object['Y'][:]
37  elif u_grid_loc == 'V':
38  x_u = grid_object['X'][:]
39  y_u = grid_object['Yp1'][:]
40  elif u_grid_loc == 'T':
41  x_u = grid_object['X'][:]
42  y_u = grid_object['Y'][:]
43  elif u_grid_loc == 'Zeta':
44  x_u = grid_object['Xp1'][:]
45  y_u = grid_object['Yp1'][:]
46  else:
47  print 'u_grid_loc not set correctly. Possible options are: U,V,T, Zeta'
48  return
49 
50  if v_grid_loc == 'U':
51  x_v = grid_object['Xp1'][:]
52  y_v = grid_object['Y'][:]
53  elif v_grid_loc == 'V':
54  x_v = grid_object['X'][:]
55  y_v = grid_object['Yp1'][:]
56  elif v_grid_loc == 'T':
57  x_v = grid_object['X'][:]
58  y_v = grid_object['Y'][:]
59  elif v_grid_loc == 'Zeta':
60  x_v = grid_object['Xp1'][:]
61  y_v = grid_object['Yp1'][:]
62  else:
63  print 'v_grid_loc not set correctly. Possible options are: U,V,T, Zeta'
64  return
65 
66 
67 
68 
69 
70  len_x_u = len(x_u)
71  len_y_u = len(y_u)
72 
73  len_x_v = len(x_v)
74  len_y_v = len(y_v)
75 
76  x_stream = np.ones((int(t_max/delta_t)+2))*startx
77  y_stream = np.ones((int(t_max/delta_t)+2))*starty
78  t_stream = np.zeros((int(t_max/delta_t)+2))
79 
80  t = 0 #set the initial time to be zero
81  i=0
82 
83  deg_per_m = np.array([1,1])
84 
85  # Runge-Kutta fourth order method to estimate next position.
86  while t < t_max:
87  if grid_object['grid_type']=='polar':
88  # use degrees per metre and convert all the velocities to degrees / second# calculate degrees per metre at current location - used to convert the m/s velocities in to degrees/s
89  deg_per_m = np.array([1./(1852.*60.),np.cos(starty*np.pi/180.)/(1852.*60.)])
90 
91  # Interpolate velocities to initial location
92  u_loc = bilinear_interp(startx,starty,u,x_u,y_u,len_x_u,len_y_u)
93  v_loc = bilinear_interp(startx,starty,v,x_v,y_v,len_x_v,len_y_v)
94  u_loc = u_loc*deg_per_m[1]
95  v_loc = v_loc*deg_per_m[0]
96  dx1 = delta_t*u_loc
97  dy1 = delta_t*v_loc
98 
99  u_loc1 = bilinear_interp(startx + 0.5*dx1,starty + 0.5*dy1,u,x_u,y_u,len_x_u,len_y_u)
100  v_loc1 = bilinear_interp(startx + 0.5*dx1,starty + 0.5*dy1,v,x_v,y_v,len_x_v,len_y_v)
101  u_loc1 = u_loc1*deg_per_m[1]
102  v_loc1 = v_loc1*deg_per_m[0]
103  dx2 = delta_t*u_loc1
104  dy2 = delta_t*v_loc1
105 
106  u_loc2 = bilinear_interp(startx + 0.5*dx2,starty + 0.5*dy2,u,x_u,y_u,len_x_u,len_y_u)
107  v_loc2 = bilinear_interp(startx + 0.5*dx2,starty + 0.5*dy2,v,x_v,y_v,len_x_v,len_y_v)
108  u_loc2 = u_loc2*deg_per_m[1]
109  v_loc2 = v_loc2*deg_per_m[0]
110  dx3 = delta_t*u_loc2
111  dy3 = delta_t*v_loc2
112 
113  u_loc3 = bilinear_interp(startx + dx3,starty + dy3,u,x_u,y_u,len_x_u,len_y_u)
114  v_loc3 = bilinear_interp(startx + dx3,starty + dy3,v,x_v,y_v,len_x_v,len_y_v)
115  u_loc3 = u_loc3*deg_per_m[1]
116  v_loc3 = v_loc3*deg_per_m[0]
117  dx4 = delta_t*u_loc3
118  dy4 = delta_t*v_loc3
119 
120  startx = startx + (dx1 + 2*dx2 + 2*dx3 + dx4)/6
121  starty = starty + (dy1 + 2*dy2 + 2*dy3 + dy4)/6
122 
123  t += delta_t
124  i += 1
125 
126  x_stream[i] = startx
127  y_stream[i] = starty
128  t_stream[i] = t
129 
130  # if x_stream[0] == x_stream[-1]:
131  # x_stream = np.delete(x_stream,-1)
132  # y_stream = np.delete(y_stream,-1)
133  # t_stream = np.delete(t_stream,-1)
134 
135  return x_stream,y_stream,t_stream
136 
137 
138 def stream3(u,v,w,
139  startx,starty,startz,
140  grid_object=None,
141  x_v='None',y_v='None',z_v='None',
142  x_w='None',y_w='None',z_w='None',
143  t_max=2592000,delta_t=3600,
144  u_grid_loc='U',v_grid_loc='V',w_grid_loc='W'):
145  """!A three-dimensional streamline solver. The velocity fields must be three dimensional and not vary in time.
146  X_grid_loc variables specify where the field "X" is located on the C-grid. Possibles options are, U, V, W, T and Zeta.
147 """
148  if grid_object:
149  if u_grid_loc == 'U':
150  x_u = copy.deepcopy(grid_object['Xp1'][:])
151  y_u = copy.deepcopy(grid_object['Y'][:])
152  z_u = copy.deepcopy(grid_object['Z'][:])
153  elif u_grid_loc == 'V':
154  x_u = copy.deepcopy(grid_object['X'][:])
155  y_u = copy.deepcopy(grid_object['Yp1'][:])
156  z_u = copy.deepcopy(grid_object['Z'][:])
157  elif u_grid_loc == 'W':
158  x_u = copy.deepcopy(grid_object['X'][:])
159  y_u = copy.deepcopy(grid_object['Y'][:])
160  z_u = copy.deepcopy(grid_object['Zl'][:])
161  elif u_grid_loc == 'T':
162  x_u = copy.deepcopy(grid_object['X'][:])
163  y_u = copy.deepcopy(grid_object['Y'][:])
164  z_u = copy.deepcopy(grid_object['Z'][:])
165  elif u_grid_loc == 'Zeta':
166  x_u = copy.deepcopy(grid_object['Xp1'][:])
167  y_u = copy.deepcopy(grid_object['Yp1'][:])
168  z_u = copy.deepcopy(grid_object['Z'][:])
169  else:
170  print 'u_grid_loc not set correctly. Possible options are: U,V,W,T, and Zeta'
171  return
172 
173  if v_grid_loc == 'U':
174  x_v = copy.deepcopy(grid_object['Xp1'][:])
175  y_v = copy.deepcopy(grid_object['Y'][:])
176  z_v = copy.deepcopy(grid_object['Z'][:])
177  elif v_grid_loc == 'V':
178  x_v = copy.deepcopy(grid_object['X'][:])
179  y_v = copy.deepcopy(grid_object['Yp1'][:])
180  z_v = copy.deepcopy(grid_object['Z'][:])
181  elif v_grid_loc == 'W':
182  x_v = copy.deepcopy(grid_object['X'][:])
183  y_v = copy.deepcopy(grid_object['Y'][:])
184  z_v = copy.deepcopy(grid_object['Zl'][:])
185  elif v_grid_loc == 'T':
186  x_v = copy.deepcopy(grid_object['X'][:])
187  y_v = copy.deepcopy(grid_object['Y'][:])
188  z_v = copy.deepcopy(grid_object['Z'][:])
189  elif v_grid_loc == 'Zeta':
190  x_v = copy.deepcopy(grid_object['Xp1'][:])
191  y_v = copy.deepcopy(grid_object['Yp1'][:])
192  z_v = copy.deepcopy(grid_object['Z'][:])
193  else:
194  print 'v_grid_loc not set correctly. Possible options are: U,V,W,T, and Zeta'
195  return
196 
197  if w_grid_loc == 'U':
198  x_w = copy.deepcopy(grid_object['Xp1'][:])
199  y_w = copy.deepcopy(grid_object['Y'][:])
200  z_w = copy.deepcopy(grid_object['Z'][:])
201  elif w_grid_loc == 'V':
202  x_w = copy.deepcopy(grid_object['X'][:])
203  y_w = copy.deepcopy(grid_object['Yp1'][:])
204  z_w = copy.deepcopy(grid_object['Z'][:])
205  elif w_grid_loc == 'W':
206  x_w = copy.deepcopy(grid_object['X'][:])
207  y_w = copy.deepcopy(grid_object['Y'][:])
208  z_w = copy.deepcopy(grid_object['Zl'][:])
209  elif w_grid_loc == 'T':
210  x_w = copy.deepcopy(grid_object['X'][:])
211  y_w = copy.deepcopy(grid_object['Y'][:])
212  z_w = copy.deepcopy(grid_object['Z'][:])
213  elif w_grid_loc == 'Zeta':
214  x_w = copy.deepcopy(grid_object['Xp1'][:])
215  y_w = copy.deepcopy(grid_object['Yp1'][:])
216  z_w = copy.deepcopy(grid_object['Z'][:])
217  else:
218  print 'w_grid_loc not set correctly. Possible options are: U,V,W,T, and Zeta'
219  return
220 
221 
222  len_x_u = len(x_u)
223  len_y_u = len(y_u)
224  len_z_u = len(z_u)
225 
226  len_x_v = len(x_v)
227  len_y_v = len(y_v)
228  len_z_v = len(z_v)
229 
230  len_x_w = len(x_w)
231  len_y_w = len(y_w)
232  len_z_w = len(z_w)
233 
234  x_stream = np.ones((int(np.fabs(t_max/delta_t))+2))*startx
235  y_stream = np.ones((int(np.fabs(t_max/delta_t))+2))*starty
236  z_stream = np.ones((int(np.fabs(t_max/delta_t))+2))*startz
237  t_stream = np.zeros((int(np.fabs(t_max/delta_t))+2))
238 
239  t = 0 #set the initial time to be zero
240  i=0
241 
242  deg_per_m = np.array([1,1])
243 
244 
245  # Runge-Kutta fourth order method to estimate next position.
246  while i < np.fabs(t_max/delta_t):
247  #t < t_max:
248  if grid_object['grid_type']=='polar':
249  # use degrees per metre and convert all the velocities to degrees / second# calculate degrees per metre at current location - used to convert the m/s velocities in to degrees/s
250  deg_per_m = np.array([1./(1852.*60.),np.cos(starty*np.pi/180.)/(1852.*60.)])
251 
252  # Interpolate velocities to initial location
253  u_loc = trilinear_interp(startx,starty,startz,u,x_u,y_u,z_u,len_x_u,len_y_u,len_z_u)
254  v_loc = trilinear_interp(startx,starty,startz,v,x_v,y_v,z_v,len_x_v,len_y_v,len_z_v)
255  w_loc = trilinear_interp(startx,starty,startz,w,x_w,y_w,z_w,len_x_w,len_y_w,len_z_w)
256  u_loc = u_loc*deg_per_m[1]
257  v_loc = v_loc*deg_per_m[0]
258  dx1 = delta_t*u_loc
259  dy1 = delta_t*v_loc
260  dz1 = delta_t*w_loc
261 
262  u_loc1 = trilinear_interp(startx + 0.5*dx1,starty + 0.5*dy1,startz + 0.5*dz1,u,x_u,y_u,z_u,len_x_u,len_y_u,len_z_u)
263  v_loc1 = trilinear_interp(startx + 0.5*dx1,starty + 0.5*dy1,startz + 0.5*dz1,v,x_v,y_v,z_v,len_x_v,len_y_v,len_z_v)
264  w_loc1 = trilinear_interp(startx + 0.5*dx1,starty + 0.5*dy1,startz + 0.5*dz1,w,x_w,y_w,z_w,len_x_w,len_y_w,len_z_w)
265  u_loc1 = u_loc1*deg_per_m[1]
266  v_loc1 = v_loc1*deg_per_m[0]
267  dx2 = delta_t*u_loc1
268  dy2 = delta_t*v_loc1
269  dz2 = delta_t*w_loc1
270 
271  u_loc2 = trilinear_interp(startx + 0.5*dx2,starty + 0.5*dy2,startz + 0.5*dz2,u,x_u,y_u,z_u,len_x_u,len_y_u,len_z_u)
272  v_loc2 = trilinear_interp(startx + 0.5*dx2,starty + 0.5*dy2,startz + 0.5*dz2,v,x_v,y_v,z_v,len_x_v,len_y_v,len_z_v)
273  w_loc2 = trilinear_interp(startx + 0.5*dx2,starty + 0.5*dy2,startz + 0.5*dz2,w,x_w,y_w,z_w,len_x_w,len_y_w,len_z_w)
274  u_loc2 = u_loc2*deg_per_m[1]
275  v_loc2 = v_loc2*deg_per_m[0]
276  dx3 = delta_t*u_loc2
277  dy3 = delta_t*v_loc2
278  dz3 = delta_t*w_loc2
279 
280  u_loc3 = trilinear_interp(startx + dx3,starty + dy3,startz + dz3,u,x_u,y_u,z_u,len_x_u,len_y_u,len_z_u)
281  v_loc3 = trilinear_interp(startx + dx3,starty + dy3,startz + dz3,v,x_v,y_v,z_v,len_x_v,len_y_v,len_z_v)
282  w_loc3 = trilinear_interp(startx + dx3,starty + dy3,startz + dz3,w,x_w,y_w,z_w,len_x_w,len_y_w,len_z_w)
283  u_loc3 = u_loc3*deg_per_m[1]
284  v_loc3 = v_loc3*deg_per_m[0]
285  dx4 = delta_t*u_loc3
286  dy4 = delta_t*v_loc3
287  dz4 = delta_t*w_loc3
288 
289  #recycle the "start_" variables to keep the code clean
290  startx = startx + (dx1 + 2*dx2 + 2*dx3 + dx4)/6
291  starty = starty + (dy1 + 2*dy2 + 2*dy3 + dy4)/6
292  startz = startz + (dz1 + 2*dz2 + 2*dz3 + dz4)/6
293  t += delta_t
294  i += 1
295 
296  x_stream[i] = startx
297  y_stream[i] = starty
298  z_stream[i] = startz
299  t_stream[i] = t
300 
301 
302  return x_stream,y_stream,z_stream,t_stream
303 
304 ###########################################################
305 def stream3_many(u,v,w,
306  startx,starty,startz,
307  grid_object=None,
308  x_v='None',y_v='None',z_v='None',
309  x_w='None',y_w='None',z_w='None',
310  t_max=2592000,delta_t=3600,
311  u_grid_loc='U',v_grid_loc='V',w_grid_loc='W'):
312  """!A three-dimensional streamline solver. The velocity fields must be three dimensional and not vary in time.
313  X_grid_loc variables specify where the field "X" is located on the C-grid. Possibles options are, U, V, W, T and Zeta.
314 
315  Returns:
316  * x_stream, y_stream, z_stream - all with dimensions (particle,time_level)
317  * t_stream - with dimensions (time_level)
318 """
319  if grid_object:
320  if u_grid_loc == 'U':
321  x_u = copy.deepcopy(grid_object['Xp1'][:])
322  y_u = copy.deepcopy(grid_object['Y'][:])
323  z_u = copy.deepcopy(grid_object['Z'][:])
324  elif u_grid_loc == 'V':
325  x_u = copy.deepcopy(grid_object['X'][:])
326  y_u = copy.deepcopy(grid_object['Yp1'][:])
327  z_u = copy.deepcopy(grid_object['Z'][:])
328  elif u_grid_loc == 'W':
329  x_u = copy.deepcopy(grid_object['X'][:])
330  y_u = copy.deepcopy(grid_object['Y'][:])
331  z_u = copy.deepcopy(grid_object['Zl'][:])
332  elif u_grid_loc == 'T':
333  x_u = copy.deepcopy(grid_object['X'][:])
334  y_u = copy.deepcopy(grid_object['Y'][:])
335  z_u = copy.deepcopy(grid_object['Z'][:])
336  elif u_grid_loc == 'Zeta':
337  x_u = copy.deepcopy(grid_object['Xp1'][:])
338  y_u = copy.deepcopy(grid_object['Yp1'][:])
339  z_u = copy.deepcopy(grid_object['Z'][:])
340  else:
341  print 'u_grid_loc not set correctly. Possible options are: U,V,W,T, and Zeta'
342  return
343 
344  if v_grid_loc == 'U':
345  x_v = copy.deepcopy(grid_object['Xp1'][:])
346  y_v = copy.deepcopy(grid_object['Y'][:])
347  z_v = copy.deepcopy(grid_object['Z'][:])
348  elif v_grid_loc == 'V':
349  x_v = copy.deepcopy(grid_object['X'][:])
350  y_v = copy.deepcopy(grid_object['Yp1'][:])
351  z_v = copy.deepcopy(grid_object['Z'][:])
352  elif v_grid_loc == 'W':
353  x_v = copy.deepcopy(grid_object['X'][:])
354  y_v = copy.deepcopy(grid_object['Y'][:])
355  z_v = copy.deepcopy(grid_object['Zl'][:])
356  elif v_grid_loc == 'T':
357  x_v = copy.deepcopy(grid_object['X'][:])
358  y_v = copy.deepcopy(grid_object['Y'][:])
359  z_v = copy.deepcopy(grid_object['Z'][:])
360  elif v_grid_loc == 'Zeta':
361  x_v = copy.deepcopy(grid_object['Xp1'][:])
362  y_v = copy.deepcopy(grid_object['Yp1'][:])
363  z_v = copy.deepcopy(grid_object['Z'][:])
364  else:
365  print 'v_grid_loc not set correctly. Possible options are: U,V,W,T, and Zeta'
366  return
367 
368  if w_grid_loc == 'U':
369  x_w = copy.deepcopy(grid_object['Xp1'][:])
370  y_w = copy.deepcopy(grid_object['Y'][:])
371  z_w = copy.deepcopy(grid_object['Z'][:])
372  elif w_grid_loc == 'V':
373  x_w = copy.deepcopy(grid_object['X'][:])
374  y_w = copy.deepcopy(grid_object['Yp1'][:])
375  z_w = copy.deepcopy(grid_object['Z'][:])
376  elif w_grid_loc == 'W':
377  x_w = copy.deepcopy(grid_object['X'][:])
378  y_w = copy.deepcopy(grid_object['Y'][:])
379  z_w = copy.deepcopy(grid_object['Zl'][:])
380  elif w_grid_loc == 'T':
381  x_w = copy.deepcopy(grid_object['X'][:])
382  y_w = copy.deepcopy(grid_object['Y'][:])
383  z_w = copy.deepcopy(grid_object['Z'][:])
384  elif w_grid_loc == 'Zeta':
385  x_w = copy.deepcopy(grid_object['Xp1'][:])
386  y_w = copy.deepcopy(grid_object['Yp1'][:])
387  z_w = copy.deepcopy(grid_object['Z'][:])
388  else:
389  print 'w_grid_loc not set correctly. Possible options are: U,V,W,T, and Zeta'
390  return
391 
392 
393  len_x_u = len(x_u)
394  len_y_u = len(y_u)
395  len_z_u = len(z_u)
396 
397  len_x_v = len(x_v)
398  len_y_v = len(y_v)
399  len_z_v = len(z_v)
400 
401  len_x_w = len(x_w)
402  len_y_w = len(y_w)
403  len_z_w = len(z_w)
404 
405  x_stream = np.ones((len(startx),int(np.fabs(t_max/delta_t))+2))*startx[:,np.newaxis]
406  y_stream = np.ones((len(startx),int(np.fabs(t_max/delta_t))+2))*starty[:,np.newaxis]
407  z_stream = np.ones((len(startx),int(np.fabs(t_max/delta_t))+2))*startz[:,np.newaxis]
408  t_stream = np.zeros((int(np.fabs(t_max/delta_t))+2))
409 
410  t = 0 #set the initial time to be zero
411  i=0
412 
413  # Prepare for spherical polar grids
414  deg_per_m = np.ones((len(startx),2),dtype=float)
415  if grid_object['grid_type']=='polar':
416  deg_per_m[:,0] = np.ones_like(startx)/(1852.*60.) # multiplier for v
417 
418  # Runge-Kutta fourth order method to estimate next position.
419  while i < np.fabs(t_max/delta_t):
420  #t < t_max:
421  if grid_object['grid_type']=='polar':
422  # use degrees per metre and convert all the velocities to degrees / second# calculate degrees per metre at current location - used to convert the m/s velocities in to degrees/s
423  deg_per_m[:,1] = np.cos(starty*np.pi/180.)/(1852.*60.)# multiplier for u
424 
425  # Interpolate velocities to initial location
426  u_loc = trilinear_interp_arrays(startx,starty,startz,u,x_u,y_u,z_u,len_x_u,len_y_u,len_z_u)
427  v_loc = trilinear_interp_arrays(startx,starty,startz,v,x_v,y_v,z_v,len_x_v,len_y_v,len_z_v)
428  w_loc = trilinear_interp_arrays(startx,starty,startz,w,x_w,y_w,z_w,len_x_w,len_y_w,len_z_w)
429  u_loc = u_loc*deg_per_m[:,1]
430  v_loc = v_loc*deg_per_m[:,0]
431  dx1 = delta_t*u_loc
432  dy1 = delta_t*v_loc
433  dz1 = delta_t*w_loc
434 
435  u_loc1 = trilinear_interp_arrays(startx + 0.5*dx1,starty + 0.5*dy1,startz + 0.5*dz1,u,x_u,y_u,z_u,len_x_u,len_y_u,len_z_u)
436  v_loc1 = trilinear_interp_arrays(startx + 0.5*dx1,starty + 0.5*dy1,startz + 0.5*dz1,v,x_v,y_v,z_v,len_x_v,len_y_v,len_z_v)
437  w_loc1 = trilinear_interp_arrays(startx + 0.5*dx1,starty + 0.5*dy1,startz + 0.5*dz1,w,x_w,y_w,z_w,len_x_w,len_y_w,len_z_w)
438  u_loc1 = u_loc1*deg_per_m[:,1]
439  v_loc1 = v_loc1*deg_per_m[:,0]
440  dx2 = delta_t*u_loc1
441  dy2 = delta_t*v_loc1
442  dz2 = delta_t*w_loc1
443 
444  u_loc2 = trilinear_interp_arrays(startx + 0.5*dx2,starty + 0.5*dy2,startz + 0.5*dz2,u,x_u,y_u,z_u,len_x_u,len_y_u,len_z_u)
445  v_loc2 = trilinear_interp_arrays(startx + 0.5*dx2,starty + 0.5*dy2,startz + 0.5*dz2,v,x_v,y_v,z_v,len_x_v,len_y_v,len_z_v)
446  w_loc2 = trilinear_interp_arrays(startx + 0.5*dx2,starty + 0.5*dy2,startz + 0.5*dz2,w,x_w,y_w,z_w,len_x_w,len_y_w,len_z_w)
447  u_loc2 = u_loc2*deg_per_m[:,1]
448  v_loc2 = v_loc2*deg_per_m[:,0]
449  dx3 = delta_t*u_loc2
450  dy3 = delta_t*v_loc2
451  dz3 = delta_t*w_loc2
452 
453  u_loc3 = trilinear_interp_arrays(startx + dx3,starty + dy3,startz + dz3,u,x_u,y_u,z_u,len_x_u,len_y_u,len_z_u)
454  v_loc3 = trilinear_interp_arrays(startx + dx3,starty + dy3,startz + dz3,v,x_v,y_v,z_v,len_x_v,len_y_v,len_z_v)
455  w_loc3 = trilinear_interp_arrays(startx + dx3,starty + dy3,startz + dz3,w,x_w,y_w,z_w,len_x_w,len_y_w,len_z_w)
456  u_loc3 = u_loc3*deg_per_m[:,1]
457  v_loc3 = v_loc3*deg_per_m[:,0]
458  dx4 = delta_t*u_loc3
459  dy4 = delta_t*v_loc3
460  dz4 = delta_t*w_loc3
461 
462  #recycle the "start_" variables to keep the code clean
463  startx = startx + (dx1 + 2*dx2 + 2*dx3 + dx4)/6
464  starty = starty + (dy1 + 2*dy2 + 2*dy3 + dy4)/6
465  startz = startz + (dz1 + 2*dz2 + 2*dz3 + dz4)/6
466  t += delta_t
467  i += 1
468 
469  x_stream[:,i] = startx
470  y_stream[:,i] = starty
471  z_stream[:,i] = startz
472  t_stream[i] = t
473 
474 
475  return x_stream,y_stream,z_stream,t_stream
476 #####################################################
477 
478 
479 def pathlines(u_netcdf_filename,v_netcdf_filename,w_netcdf_filename,
480  startx,starty,startz,startt,
481  t,
482  grid_object,
483  t_max,delta_t,
484  u_netcdf_variable='UVEL',
485  v_netcdf_variable='VVEL',
486  w_netcdf_variable='WVEL',
487  u_grid_loc='U',v_grid_loc='V',w_grid_loc='W',
488  u_bias_field=None,
489  v_bias_field=None,
490  w_bias_field=None):
491  """!A three-dimensional lagrangian particle tracker. The velocity fields must be four dimensional (three spatial, one temporal) and have units of m/s.
492  It should work to track particles forwards or backwards in time (set delta_t <0 for backwards in time). But, be warned, backwards in time hasn't been thoroughly tested yet.
493 
494  Because this is a very large amount of data, the fields are passed as netcdffile handles.
495 
496  The variables are:
497  * ?_netcdf_filename = name of the netcdf file with ?'s data in it.
498  * start? = intial value for x, y, z, or t.
499  * t = vector of time levels that are contained in the velocity data.
500  * grid_object is m.grid if you followed the standard naming conventions.
501  * ?_netcdf_variable = name of the "?" variable field in the netcdf file.
502  * t_max = length of time to track particles for, in seconds. This is always positive
503  * delta_t = timestep for particle tracking algorithm, in seconds. This can be positive or negative.
504  * ?_grid_loc = where the field "?" is located on the C-grid. Possibles options are, U, V, W, T and Zeta.
505  * ?_bias_field = bias to add to that velocity field omponent. If set to -mean(velocity component), then only the time varying portion of that field will be used.
506  """
507 
508  if u_grid_loc == 'U':
509  x_u = grid_object['Xp1'][:]
510  y_u = grid_object['Y'][:]
511  z_u = grid_object['Z'][:]
512  elif u_grid_loc == 'V':
513  x_u = grid_object['X'][:]
514  y_u = grid_object['Yp1'][:]
515  z_u = grid_object['Z'][:]
516  elif u_grid_loc == 'W':
517  x_u = grid_object['X'][:]
518  y_u = grid_object['Y'][:]
519  z_u = grid_object['Zl'][:]
520  elif u_grid_loc == 'T':
521  x_u = grid_object['X'][:]
522  y_u = grid_object['Y'][:]
523  z_u = grid_object['Z'][:]
524  elif u_grid_loc == 'Zeta':
525  x_u = grid_object['Xp1'][:]
526  y_u = grid_object['Yp1'][:]
527  z_u = grid_object['Z'][:]
528  else:
529  print 'u_grid_loc not set correctly. Possible options are: U,V,W,T, and Zeta'
530  return
531 
532  if v_grid_loc == 'U':
533  x_v = grid_object['Xp1'][:]
534  y_v = grid_object['Y'][:]
535  z_v = grid_object['Z'][:]
536  elif v_grid_loc == 'V':
537  x_v = grid_object['X'][:]
538  y_v = grid_object['Yp1'][:]
539  z_v = grid_object['Z'][:]
540  elif v_grid_loc == 'W':
541  x_v = grid_object['X'][:]
542  y_v = grid_object['Y'][:]
543  z_v = grid_object['Zl'][:]
544  elif v_grid_loc == 'T':
545  x_v = grid_object['X'][:]
546  y_v = grid_object['Y'][:]
547  z_v = grid_object['Z'][:]
548  elif v_grid_loc == 'Zeta':
549  x_v = grid_object['Xp1'][:]
550  y_v = grid_object['Yp1'][:]
551  z_v = grid_object['Z'][:]
552  else:
553  print 'v_grid_loc not set correctly. Possible options are: U,V,W,T, and Zeta'
554  return
555 
556  if w_grid_loc == 'U':
557  x_w = grid_object['Xp1'][:]
558  y_w = grid_object['Y'][:]
559  z_w = grid_object['Z'][:]
560  elif w_grid_loc == 'V':
561  x_w = grid_object['X'][:]
562  y_w = grid_object['Yp1'][:]
563  z_w = grid_object['Z'][:]
564  elif w_grid_loc == 'W':
565  x_w = grid_object['X'][:]
566  y_w = grid_object['Y'][:]
567  z_w = grid_object['Zl'][:]
568  elif w_grid_loc == 'T':
569  x_w = grid_object['X'][:]
570  y_w = grid_object['Y'][:]
571  z_w = grid_object['Z'][:]
572  elif w_grid_loc == 'Zeta':
573  x_w = grid_object['Xp1'][:]
574  y_w = grid_object['Yp1'][:]
575  z_w = grid_object['Z'][:]
576  else:
577  print 'w_grid_loc not set correctly. Possible options are: U,V,W,T, and Zeta'
578  return
579 
580  len_x_u = len(x_u)
581  len_y_u = len(y_u)
582  len_z_u = len(z_u)
583 
584  len_x_v = len(x_v)
585  len_y_v = len(y_v)
586  len_z_v = len(z_v)
587 
588  len_x_w = len(x_w)
589  len_y_w = len(y_w)
590  len_z_w = len(z_w)
591 
592  len_t = len(t)
593 
594  if u_bias_field is None:
595  u_bias_field = np.zeros_like(grid_object['wet_mask_U'][:])
596  if v_bias_field is None:
597  v_bias_field = np.zeros_like(grid_object['wet_mask_V'][:])
598  if w_bias_field is None:
599  w_bias_field = np.zeros_like(grid_object['wet_mask_W'][:])
600 
601  x_stream = np.ones((int(np.fabs(t_max/delta_t))+2))*startx
602  y_stream = np.ones((int(np.fabs(t_max/delta_t))+2))*starty
603  z_stream = np.ones((int(np.fabs(t_max/delta_t))+2))*startz
604  t_stream = np.ones((int(np.fabs(t_max/delta_t))+2))*startt
605 
606  t_RK = startt #set the initial time to be the given start time
607  z_RK = startz
608  y_RK = starty
609  x_RK = startx
610 
611  i=0
612 
613  u_netcdf_filehandle = netCDF4.Dataset(u_netcdf_filename)
614  v_netcdf_filehandle = netCDF4.Dataset(v_netcdf_filename)
615  w_netcdf_filehandle = netCDF4.Dataset(w_netcdf_filename)
616 
617  t_index = np.searchsorted(t,t_RK)
618  t_index_new = np.searchsorted(t,t_RK) # this is later used to test if new data needs to be read in.
619  if t_index == 0:
620  raise ValueError('Given time value is outside the given velocity fields - too small')
621  elif t_index == len_t:
622  raise ValueError('Given time value is outside the given velocity fields - too big')
623 
624 
625  # load fields in ready for the first run through the loop
626  # u
627  u_field,x_index_u,y_index_u,z_index_u = indices_and_field(x_u,y_u,z_u,
628  x_RK,y_RK,z_RK,t_index,
629  len_x_u,len_y_u,len_z_u,len_t,
630  u_netcdf_filehandle,u_netcdf_variable,u_bias_field)
631  u_field,x_index_u_new,y_index_u_new,z_index_u_new = indices_and_field(x_u,y_u,z_u,
632  x_RK,y_RK,z_RK,t_index,
633  len_x_u,len_y_u,len_z_u,len_t,
634  u_netcdf_filehandle,u_netcdf_variable,u_bias_field)
635  # v
636  v_field,x_index_v,y_index_v,z_index_v = indices_and_field(x_v,y_v,z_v,
637  x_RK,y_RK,z_RK,t_index,
638  len_x_v,len_y_v,len_z_v,len_t,
639  v_netcdf_filehandle,v_netcdf_variable,v_bias_field)
640  v_field,x_index_v_new,y_index_v_new,z_index_v_new = indices_and_field(x_v,y_v,z_v,
641  x_RK,y_RK,z_RK,t_index,
642  len_x_v,len_y_v,len_z_v,len_t,
643  v_netcdf_filehandle,v_netcdf_variable,v_bias_field)
644 
645  # w
646  w_field,x_index_w,y_index_w,z_index_w = indices_and_field(x_w,y_w,z_w,
647  x_RK,y_RK,z_RK,t_index,
648  len_x_w,len_y_w,len_z_w,len_t,
649  w_netcdf_filehandle,w_netcdf_variable,w_bias_field)
650  w_field,x_index_w_new,y_index_w_new,z_index_w_new = indices_and_field(x_w,y_w,z_w,
651  x_RK,y_RK,z_RK,t_index,
652  len_x_w,len_y_w,len_z_w,len_t,
653  w_netcdf_filehandle,w_netcdf_variable,w_bias_field)
654 
655 
656  # Prepare for spherical polar grids
657  deg_per_m = np.array([1,1])
658 
659  # Runge-Kutta fourth order method to estimate next position.
660  while i < np.fabs(t_max/delta_t):
661  #t_RK < t_max + startt:
662 
663  if grid_object['grid_type']=='polar':
664  # use degrees per metre and convert all the velocities to degrees / second# calculate degrees per metre at current location - used to convert the m/s velocities in to degrees/s
665  deg_per_m = np.array([1./(1852.*60.),np.cos(starty*np.pi/180.)/(1852.*60.)])
666  # Compute indices at location given
667 
668  if (y_index_u_new==y_index_u and
669  x_index_u_new==x_index_u and
670  z_index_u_new==z_index_u and
671 
672  y_index_v_new==y_index_v and
673  x_index_v_new==x_index_v and
674  z_index_v_new==z_index_v and
675 
676  y_index_w_new==y_index_w and
677  x_index_w_new==x_index_w and
678  z_index_w_new==z_index_w and
679 
680  t_index_new == t_index):
681  # the particle hasn't moved out of the grid cell it was in.
682  # So the loaded field is fine; there's no need to reload it.
683  pass
684  else:
685  t_index = np.searchsorted(t,t_RK)
686  if t_index == 0:
687  raise ValueError('Given time value is outside the given velocity fields - too small')
688  elif t_index == len_t:
689  raise ValueError('Given time value is outside the given velocity fields - too big')
690 
691 
692  # for u
693 
694  u_field,x_index_u,y_index_u,z_index_u = indices_and_field(x_u,y_u,z_u,
695  x_RK,y_RK,z_RK,t_index,
696  len_x_u,len_y_u,len_z_u,len_t,
697  u_netcdf_filehandle,u_netcdf_variable,u_bias_field)
698  # for v
699  v_field,x_index_v,y_index_v,z_index_v = indices_and_field(x_v,y_v,z_v,
700  x_RK,y_RK,z_RK,t_index,
701  len_x_v,len_y_v,len_z_v,len_t,
702  v_netcdf_filehandle,v_netcdf_variable,v_bias_field)
703 
704  # for w
705  w_field,x_index_w,y_index_w,z_index_w = indices_and_field(x_w,y_w,z_w,
706  x_RK,y_RK,z_RK,t_index,
707  len_x_w,len_y_w,len_z_w,len_t,
708  w_netcdf_filehandle,w_netcdf_variable,w_bias_field)
709 
710 
711 
712  # Interpolate velocities to initial location
713  u_loc = quadralinear_interp(x_RK,y_RK,z_RK,t_RK,
714  u_field,
715  x_u,y_u,z_u,t,
716  len_x_u,len_y_u,len_z_u,len_t,
717  x_index_u,y_index_u,z_index_u,t_index)
718  v_loc = quadralinear_interp(x_RK,y_RK,z_RK,t_RK,
719  v_field,
720  x_v,y_v,z_v,t,len_x_v,len_y_v,len_z_v,len_t,
721  x_index_v,y_index_v,z_index_v,t_index)
722  w_loc = quadralinear_interp(x_RK,y_RK,z_RK,t_RK,
723  w_field,
724  x_w,y_w,z_w,t,len_x_w,len_y_w,len_z_w,len_t,
725  x_index_w,y_index_w,z_index_w,t_index)
726  u_loc = u_loc*deg_per_m[1]
727  v_loc = v_loc*deg_per_m[0]
728  dx1 = delta_t*u_loc
729  dy1 = delta_t*v_loc
730  dz1 = delta_t*w_loc
731 
732  u_loc1 = quadralinear_interp(x_RK + 0.5*dx1,y_RK + 0.5*dy1,z_RK + 0.5*dz1,t_RK + 0.5*delta_t,
733  u_field,
734  x_u,y_u,z_u,t,len_x_u,len_y_u,len_z_u,len_t,
735  x_index_u,y_index_u,z_index_u,t_index)
736  v_loc1 = quadralinear_interp(x_RK + 0.5*dx1,y_RK + 0.5*dy1,z_RK + 0.5*dz1,t_RK + 0.5*delta_t,
737  v_field,
738  x_v,y_v,z_v,t,len_x_v,len_y_v,len_z_v,len_t,
739  x_index_v,y_index_v,z_index_v,t_index)
740  w_loc1 = quadralinear_interp(x_RK + 0.5*dx1,y_RK + 0.5*dy1,z_RK + 0.5*dz1,t_RK + 0.5*delta_t,
741  w_field,
742  x_w,y_w,z_w,t,len_x_w,len_y_w,len_z_w,len_t,
743  x_index_w,y_index_w,z_index_w,t_index)
744  u_loc1 = u_loc1*deg_per_m[1]
745  v_loc1 = v_loc1*deg_per_m[0]
746  dx2 = delta_t*u_loc1
747  dy2 = delta_t*v_loc1
748  dz2 = delta_t*w_loc1
749 
750  u_loc2 = quadralinear_interp(x_RK + 0.5*dx2,y_RK + 0.5*dy2,z_RK + 0.5*dz2,t_RK + 0.5*delta_t,
751  u_field,
752  x_u,y_u,z_u,t,len_x_u,len_y_u,len_z_u,len_t,
753  x_index_u,y_index_u,z_index_u,t_index)
754  v_loc2 = quadralinear_interp(x_RK + 0.5*dx2,y_RK + 0.5*dy2,z_RK + 0.5*dz2,t_RK + 0.5*delta_t,
755  v_field,
756  x_v,y_v,z_v,t,len_x_v,len_y_v,len_z_v,len_t,
757  x_index_v,y_index_v,z_index_v,t_index)
758  w_loc2 = quadralinear_interp(x_RK + 0.5*dx2,y_RK + 0.5*dy2,z_RK + 0.5*dz2,t_RK + 0.5*delta_t,
759  w_field,
760  x_w,y_w,z_w,t,len_x_w,len_y_w,len_z_w,len_t,
761  x_index_w,y_index_w,z_index_w,t_index)
762  u_loc2 = u_loc2*deg_per_m[1]
763  v_loc2 = v_loc2*deg_per_m[0]
764  dx3 = delta_t*u_loc2
765  dy3 = delta_t*v_loc2
766  dz3 = delta_t*w_loc2
767 
768  u_loc3 = quadralinear_interp(x_RK + dx3,y_RK + dy3,z_RK + dz3,t_RK + delta_t,
769  u_field,
770  x_u,y_u,z_u,t,len_x_u,len_y_u,len_z_u,len_t,
771  x_index_u,y_index_u,z_index_u,t_index)
772  v_loc3 = quadralinear_interp(x_RK + dx3,y_RK + dy3,z_RK + dz3,t_RK + delta_t,
773  v_field,
774  x_v,y_v,z_v,t,len_x_v,len_y_v,len_z_v,len_t,
775  x_index_v,y_index_v,z_index_v,t_index)
776  w_loc3 = quadralinear_interp(x_RK + dx3,y_RK + dy3,z_RK + dz3,t_RK + delta_t,
777  w_field,
778  x_w,y_w,z_w,t,len_x_w,len_y_w,len_z_w,len_t,
779  x_index_w,y_index_w,z_index_w,t_index)
780  u_loc3 = u_loc3*deg_per_m[1]
781  v_loc3 = v_loc3*deg_per_m[0]
782  dx4 = delta_t*u_loc3
783  dy4 = delta_t*v_loc3
784  dz4 = delta_t*w_loc3
785 
786  #recycle the variables to keep the code clean
787  x_RK = x_RK + (dx1 + 2*dx2 + 2*dx3 + dx4)/6
788  y_RK = y_RK + (dy1 + 2*dy2 + 2*dy3 + dy4)/6
789  z_RK = z_RK + (dz1 + 2*dz2 + 2*dz3 + dz4)/6
790  t_RK += delta_t
791  i += 1
792 
793  x_stream[i] = x_RK
794  y_stream[i] = y_RK
795  z_stream[i] = z_RK
796  t_stream[i] = t_RK
797 
798  t_index_new = np.searchsorted(t,t_RK)
799  x_index_w_new = np.searchsorted(x_w,x_RK)
800  y_index_w_new = np.searchsorted(y_w,y_RK)
801  if z_RK < 0:
802  z_index_w_new = np.searchsorted(-z_w,-z_RK)
803  else:
804  z_index_w_new = np.searchsorted(z_w,z_RK)
805 
806  x_index_v_new = np.searchsorted(x_v,x_RK)
807  y_index_v_new = np.searchsorted(y_v,y_RK)
808  if z_RK < 0:
809  z_index_v_new = np.searchsorted(-z_v,-z_RK)
810  else:
811  z_index_v_new = np.searchsorted(z_v,z_RK)
812 
813  x_index_u_new = np.searchsorted(x_u,x_RK)
814  y_index_u_new = np.searchsorted(y_u,y_RK)
815  if z_RK < 0:
816  z_index_u_new = np.searchsorted(-z_u,-z_RK)
817  else:
818  z_index_u_new = np.searchsorted(z_u,z_RK)
819 
820 
821  u_netcdf_filehandle.close()
822  v_netcdf_filehandle.close()
823  w_netcdf_filehandle.close()
824 
825  return x_stream,y_stream,z_stream,t_stream
826 
827 
828  ######################################################################
829 
830 def pathlines_many(u_netcdf_filename,v_netcdf_filename,w_netcdf_filename,
831  startx,starty,startz,startt,
832  t,
833  grid_object,
834  t_max,delta_t,
835  u_netcdf_variable='UVEL',
836  v_netcdf_variable='VVEL',
837  w_netcdf_variable='WVEL',
838  u_grid_loc='U',v_grid_loc='V',w_grid_loc='W',
839  u_bias_field=None,
840  v_bias_field=None,
841  w_bias_field=None):
842  """!A three-dimensional lagrangian particle tracker designed for tracking many particles at once. If you're tracking fewer than O(10) - use the pathlines function.
843 
844  The velocity fields must be four dimensional (three spatial, one temporal) and have units of m/s.
845  It should work to track particles forwards or backwards in time (set delta_t <0 for backwards in time). But, be warned, backwards in time hasn't been thoroughly tested yet.
846 
847  Because this is a very large amount of data, the fields are passed as netcdffile handles.
848 
849 
850  ## Returns:
851  * x_stream, y_stream, z_stream - all with dimensions (particle,time_level)
852  * t_stream - with dimensions (time_level)
853 
854  ## The variables are:
855  * ?_netcdf_filename = name of the netcdf file with ?'s data in it.
856  * start? = (nx1) arrays of initial values for x, y, or z.
857  * startt = start time
858  * t = vector of time levels that are contained in the velocity data.
859  * grid_object is m.grid if you followed the standard naming conventions.
860  * ?_netcdf_variable = name of the "?" variable field in the netcdf file.
861  * t_max = length of time to track particles for, in seconds. This is always positive
862  * delta_t = timestep for particle tracking algorithm, in seconds. This can be positive or negative.
863  * ?_grid_loc = where the field "?" is located on the C-grid. Possibles options are, U, V, W, T and Zeta.
864  * ?_bias_field = bias to add to that velocity field component. If set to -mean(velocity component), then only the time varying portion of that field will be used.
865  """
866 
867  if u_grid_loc == 'U':
868  x_u = grid_object['Xp1'][:]
869  y_u = grid_object['Y'][:]
870  z_u = grid_object['Z'][:]
871  elif u_grid_loc == 'V':
872  x_u = grid_object['X'][:]
873  y_u = grid_object['Yp1'][:]
874  z_u = grid_object['Z'][:]
875  elif u_grid_loc == 'W':
876  x_u = grid_object['X'][:]
877  y_u = grid_object['Y'][:]
878  z_u = grid_object['Zl'][:]
879  elif u_grid_loc == 'T':
880  x_u = grid_object['X'][:]
881  y_u = grid_object['Y'][:]
882  z_u = grid_object['Z'][:]
883  elif u_grid_loc == 'Zeta':
884  x_u = grid_object['Xp1'][:]
885  y_u = grid_object['Yp1'][:]
886  z_u = grid_object['Z'][:]
887  else:
888  print 'u_grid_loc not set correctly. Possible options are: U,V,W,T, and Zeta'
889  return
890 
891  if v_grid_loc == 'U':
892  x_v = grid_object['Xp1'][:]
893  y_v = grid_object['Y'][:]
894  z_v = grid_object['Z'][:]
895  elif v_grid_loc == 'V':
896  x_v = grid_object['X'][:]
897  y_v = grid_object['Yp1'][:]
898  z_v = grid_object['Z'][:]
899  elif v_grid_loc == 'W':
900  x_v = grid_object['X'][:]
901  y_v = grid_object['Y'][:]
902  z_v = grid_object['Zl'][:]
903  elif v_grid_loc == 'T':
904  x_v = grid_object['X'][:]
905  y_v = grid_object['Y'][:]
906  z_v = grid_object['Z'][:]
907  elif v_grid_loc == 'Zeta':
908  x_v = grid_object['Xp1'][:]
909  y_v = grid_object['Yp1'][:]
910  z_v = grid_object['Z'][:]
911  else:
912  print 'v_grid_loc not set correctly. Possible options are: U,V,W,T, and Zeta'
913  return
914 
915  if w_grid_loc == 'U':
916  x_w = grid_object['Xp1'][:]
917  y_w = grid_object['Y'][:]
918  z_w = grid_object['Z'][:]
919  elif w_grid_loc == 'V':
920  x_w = grid_object['X'][:]
921  y_w = grid_object['Yp1'][:]
922  z_w = grid_object['Z'][:]
923  elif w_grid_loc == 'W':
924  x_w = grid_object['X'][:]
925  y_w = grid_object['Y'][:]
926  z_w = grid_object['Zl'][:]
927  elif w_grid_loc == 'T':
928  x_w = grid_object['X'][:]
929  y_w = grid_object['Y'][:]
930  z_w = grid_object['Z'][:]
931  elif w_grid_loc == 'Zeta':
932  x_w = grid_object['Xp1'][:]
933  y_w = grid_object['Yp1'][:]
934  z_w = grid_object['Z'][:]
935  else:
936  print 'w_grid_loc not set correctly. Possible options are: U,V,W,T, and Zeta'
937  return
938 
939  len_x_u = len(x_u)
940  len_y_u = len(y_u)
941  len_z_u = len(z_u)
942 
943  len_x_v = len(x_v)
944  len_y_v = len(y_v)
945  len_z_v = len(z_v)
946 
947  len_x_w = len(x_w)
948  len_y_w = len(y_w)
949  len_z_w = len(z_w)
950 
951  len_t = len(t)
952 
953  if u_bias_field is None:
954  u_bias_field = np.zeros_like(grid_object['wet_mask_U'][:])
955  if v_bias_field is None:
956  v_bias_field = np.zeros_like(grid_object['wet_mask_V'][:])
957  if w_bias_field is None:
958  w_bias_field = np.zeros_like(grid_object['wet_mask_W'][1:,...])
959 
960  x_stream = np.ones((len(startx),int(np.fabs(t_max/delta_t))+2))*startx[:,np.newaxis]
961  y_stream = np.ones((len(startx),int(np.fabs(t_max/delta_t))+2))*starty[:,np.newaxis]
962  z_stream = np.ones((len(startx),int(np.fabs(t_max/delta_t))+2))*startz[:,np.newaxis]
963  t_stream = np.ones((int(np.fabs(t_max/delta_t))+2))*startt
964 
965  t_RK = startt #set the initial time to be the given start time
966  z_RK = startz
967  y_RK = starty
968  x_RK = startx
969 
970  i=0
971 
972  u_netcdf_filehandle = netCDF4.Dataset(u_netcdf_filename)
973  v_netcdf_filehandle = netCDF4.Dataset(v_netcdf_filename)
974  w_netcdf_filehandle = netCDF4.Dataset(w_netcdf_filename)
975 
976  t_index = np.searchsorted(t,t_RK)
977  t_index_new = np.searchsorted(t,t_RK) # this is later used to test if new data needs to be read in.
978  if t_index == 0:
979  raise ValueError('Given time value is outside the given velocity fields - too small')
980  elif t_index == len_t:
981  raise ValueError('Given time value is outside the given velocity fields - too big')
982 
983 
984  u_field_before = u_netcdf_filehandle.variables[u_netcdf_variable][t_index,...]
985  u_field_after = u_netcdf_filehandle.variables[u_netcdf_variable][t_index+1,...]
986  u_field = (u_field_before + ((u_field_before - u_field_after)*
987  (t_RK - t[t_index])/(t[t_index+1] - t[t_index])) -
988  u_bias_field)
989 
990  v_field_before = v_netcdf_filehandle.variables[v_netcdf_variable][t_index,...]
991  v_field_after = v_netcdf_filehandle.variables[v_netcdf_variable][t_index+1,...]
992  v_field = (v_field_before + ((v_field_before - v_field_after)*
993  (t_RK - t[t_index])/(t[t_index+1] - t[t_index])) -
994  v_bias_field)
995 
996  w_field_before = w_netcdf_filehandle.variables[w_netcdf_variable][t_index,...]
997  w_field_after = w_netcdf_filehandle.variables[w_netcdf_variable][t_index+1,...]
998  w_field = (w_field_before + ((w_field_before - w_field_after)*
999  (t_RK - t[t_index])/(t[t_index+1] - t[t_index])) -
1000  w_bias_field)
1001 
1002 
1003  # Prepare for spherical polar grids
1004  deg_per_m = np.ones((len(startx),2),dtype=float)
1005  if grid_object['grid_type']=='polar':
1006  deg_per_m[:,0] = np.ones_like(startx)/(1852.*60.) # multiplier for v
1007 
1008  # Runge-Kutta fourth order method to estimate next position.
1009  while i < np.fabs(t_max/delta_t):
1010  #t_RK < t_max + startt:
1011 
1012  if grid_object['grid_type']=='polar':
1013  # use degrees per metre and convert all the velocities to degrees / second# calculate degrees per metre at current location - used to convert the m/s velocities in to degrees/s
1014  deg_per_m[:,1] = np.cos(starty*np.pi/180.)/(1852.*60.)# multiplier for u
1015  # Compute indices at location given
1016 
1017  if (t_index_new == t_index):
1018  # time hasn't progressed beyond the loaded time slices
1019  pass
1020  else:
1021  t_index = np.searchsorted(t,t_RK)
1022  if t_index == 0:
1023  raise ValueError('Given time value is outside the given velocity fields - too small')
1024  elif t_index == len_t:
1025  raise ValueError('Given time value is outside the given velocity fields - too big')
1026 
1027 
1028  u_field_before = u_netcdf_filehandle.variables[u_netcdf_variable][t_index,...]
1029  u_field_after = u_netcdf_filehandle.variables[u_netcdf_variable][t_index+1,...]
1030  u_field = (u_field_before + ((u_field_before - u_field_after)*
1031  (t_RK - t[t_index])/(t[t_index+1] - t[t_index])) - u_bias_field)
1032 
1033  v_field_before = v_netcdf_filehandle.variables[v_netcdf_variable][t_index,...]
1034  v_field_after = v_netcdf_filehandle.variables[v_netcdf_variable][t_index+1,...]
1035  v_field = (v_field_before + ((v_field_before - v_field_after)*
1036  (t_RK - t[t_index])/(t[t_index+1] - t[t_index])) - v_bias_field)
1037 
1038  w_field_before = w_netcdf_filehandle.variables[w_netcdf_variable][t_index,...]
1039  w_field_after = w_netcdf_filehandle.variables[w_netcdf_variable][t_index+1,...]
1040  w_field = (w_field_before + ((w_field_before - w_field_after)*
1041  (t_RK - t[t_index])/(t[t_index+1] - t[t_index])) - w_bias_field)
1042 
1043  # Interpolate velocities at initial location
1044  u_loc = np.ones_like(startx)
1045  v_loc = np.ones_like(startx)
1046  w_loc = np.ones_like(startx)
1047 
1048  u_loc = trilinear_interp_arrays(x_RK,y_RK,z_RK,u_field,x_u,y_u,z_u,len_x_u,len_y_u,len_z_u)
1049  v_loc = trilinear_interp_arrays(x_RK,y_RK,z_RK,v_field,x_v,y_v,z_v,len_x_v,len_y_v,len_z_v)
1050  w_loc = trilinear_interp_arrays(x_RK,y_RK,z_RK,w_field,x_w,y_w,z_w,len_x_w,len_y_w,len_z_w)
1051 
1052  u_loc = u_loc*deg_per_m[:,1]
1053  v_loc = v_loc*deg_per_m[:,0]
1054 
1055  dx1 = delta_t*u_loc
1056  dy1 = delta_t*v_loc
1057  dz1 = delta_t*w_loc
1058 
1059  u_loc1 = np.ones_like(startx)
1060  v_loc1 = np.ones_like(startx)
1061  w_loc1 = np.ones_like(startx)
1062 
1063  u_loc1 = trilinear_interp_arrays(x_RK + 0.5*dx1,y_RK + 0.5*dy1,z_RK + 0.5*dz1,
1064  u_field,x_u,y_u,z_u,len_x_u,len_y_u,len_z_u)
1065  v_loc1 = trilinear_interp_arrays(x_RK + 0.5*dx1,y_RK + 0.5*dy1,z_RK + 0.5*dz1,
1066  v_field,x_v,y_v,z_v,len_x_v,len_y_v,len_z_v)
1067  w_loc1 = trilinear_interp_arrays(x_RK + 0.5*dx1,y_RK + 0.5*dy1,z_RK + 0.5*dz1,
1068  w_field,x_w,y_w,z_w,len_x_w,len_y_w,len_z_w)
1069  u_loc1 = u_loc1*deg_per_m[:,1]
1070  v_loc1 = v_loc1*deg_per_m[:,0]
1071  dx2 = delta_t*u_loc1
1072  dy2 = delta_t*v_loc1
1073  dz2 = delta_t*w_loc1
1074 
1075  u_loc2 = np.ones_like(startx)
1076  v_loc2 = np.ones_like(startx)
1077  w_loc2 = np.ones_like(startx)
1078  u_loc2 = trilinear_interp_arrays(x_RK + 0.5*dx2,y_RK + 0.5*dy2,z_RK + 0.5*dz2,
1079  u_field,x_u,y_u,z_u,len_x_u,len_y_u,len_z_u)
1080  v_loc2 = trilinear_interp_arrays(x_RK + 0.5*dx2,y_RK + 0.5*dy2,z_RK + 0.5*dz2,
1081  v_field,x_v,y_v,z_v,len_x_v,len_y_v,len_z_v)
1082  w_loc2 = trilinear_interp_arrays(x_RK + 0.5*dx2,y_RK + 0.5*dy2,z_RK + 0.5*dz2,
1083  w_field,x_w,y_w,z_w,len_x_w,len_y_w,len_z_w)
1084 
1085  u_loc2 = u_loc2*deg_per_m[:,1]
1086  v_loc2 = v_loc2*deg_per_m[:,0]
1087  dx3 = delta_t*u_loc2
1088  dy3 = delta_t*v_loc2
1089  dz3 = delta_t*w_loc2
1090 
1091  u_loc3 = np.ones_like(startx)
1092  v_loc3 = np.ones_like(startx)
1093  w_loc3 = np.ones_like(startx)
1094  u_loc3 = trilinear_interp_arrays(x_RK + dx3,y_RK + dy3,z_RK + dz3,
1095  u_field,x_u,y_u,z_u,len_x_u,len_y_u,len_z_u)
1096  v_loc3 = trilinear_interp_arrays(x_RK + dx3,y_RK + dy3,z_RK + dz3,
1097  v_field,x_v,y_v,z_v,len_x_v,len_y_v,len_z_v)
1098  w_loc3 = trilinear_interp_arrays(x_RK + dx3,y_RK + dy3,z_RK + dz3,
1099  w_field,x_w,y_w,z_w,len_x_w,len_y_w,len_z_w)
1100 
1101  u_loc3 = u_loc3*deg_per_m[:,1]
1102  v_loc3 = v_loc3*deg_per_m[:,0]
1103  dx4 = delta_t*u_loc3
1104  dy4 = delta_t*v_loc3
1105  dz4 = delta_t*w_loc3
1106 
1107  #recycle the variables to keep the code clean
1108  x_RK = x_RK + (dx1 + 2*dx2 + 2*dx3 + dx4)/6
1109  y_RK = y_RK + (dy1 + 2*dy2 + 2*dy3 + dy4)/6
1110  z_RK = z_RK + (dz1 + 2*dz2 + 2*dz3 + dz4)/6
1111  t_RK += delta_t
1112  i += 1
1113 
1114  x_stream[:,i] = x_RK
1115  y_stream[:,i] = y_RK
1116  z_stream[:,i] = z_RK
1117  t_stream[i] = t_RK
1118 
1119  t_index_new = np.searchsorted(t,t_RK)
1120 
1121 
1122  u_netcdf_filehandle.close()
1123  v_netcdf_filehandle.close()
1124  w_netcdf_filehandle.close()
1125 
1126  return x_stream,y_stream,z_stream,t_stream
1127 
1128 
1129  #################################################################################
1130 
1131 
1132 def pathlines_for_OLIC_xyzt_ani(u_netcdf_filename,v_netcdf_filename,w_netcdf_filename,
1133  n_particles,startt,
1134  t,
1135  grid_object,
1136  t_tracking,delta_t, trace_length,
1137  u_netcdf_variable='UVEL',
1138  v_netcdf_variable='VVEL',
1139  w_netcdf_variable='WVEL',
1140  u_grid_loc='U',v_grid_loc='V',w_grid_loc='W',
1141  u_bias_field=None,
1142  v_bias_field=None,
1143  w_bias_field=None):
1144  """!A three-dimensional lagrangian particle tracker designed for tracking many particles at once. If you're tracking fewer than O(10) - use the streaklines function.
1145 
1146  The velocity fields must be four dimensional (three spatial, one temporal) and have units of m/s.
1147  It should work to track particles forwards or backwards in time (set delta_t <0 for backwards in time). But, be warned, backwards in time hasn't been thoroughly tested yet.
1148 
1149  Because this is a very large amount of data, the fields are passed as netcdffile handles.
1150 
1151 
1152  ## Returns:
1153  * x_stream, y_stream, z_stream - all with dimensions (particle,time_level)
1154  * t_stream - with dimensions (time_level)
1155 
1156  ## The variables are:
1157  * ?_netcdf_filename = name of the netcdf file with ?'s data in it.
1158  * n_particles = number of particles to track
1159  * startt = start time
1160  * t = vector of time levels that are contained in the velocity data.
1161  * grid_object is m.grid if you followed the standard naming conventions.
1162  * ?_netcdf_variable = name of the "?" variable field in the netcdf file.
1163  * t_tracking = length of time to track particles for, in seconds. This is always positive
1164  * delta_t = timestep for particle tracking algorithm, in seconds. This can be positive or negative.
1165  * trace_length = length of time for each individual trace
1166  * ?_grid_loc = where the field "?" is located on the C-grid. Possibles options are, U, V, W, T and Zeta.
1167  * ?_bias_field = bias to add to that velocity field component. If set to -mean(velocity component), then only the time varying portion of that field will be used.
1168  """
1169 
1170  if u_grid_loc == 'U':
1171  x_u = grid_object['Xp1'][:]
1172  y_u = grid_object['Y'][:]
1173  z_u = grid_object['Z'][:]
1174  elif u_grid_loc == 'V':
1175  x_u = grid_object['X'][:]
1176  y_u = grid_object['Yp1'][:]
1177  z_u = grid_object['Z'][:]
1178  elif u_grid_loc == 'W':
1179  x_u = grid_object['X'][:]
1180  y_u = grid_object['Y'][:]
1181  z_u = grid_object['Zl'][:]
1182  elif u_grid_loc == 'T':
1183  x_u = grid_object['X'][:]
1184  y_u = grid_object['Y'][:]
1185  z_u = grid_object['Z'][:]
1186  elif u_grid_loc == 'Zeta':
1187  x_u = grid_object['Xp1'][:]
1188  y_u = grid_object['Yp1'][:]
1189  z_u = grid_object['Z'][:]
1190  else:
1191  print 'u_grid_loc not set correctly. Possible options are: U,V,W,T, and Zeta'
1192  return
1193 
1194  if v_grid_loc == 'U':
1195  x_v = grid_object['Xp1'][:]
1196  y_v = grid_object['Y'][:]
1197  z_v = grid_object['Z'][:]
1198  elif v_grid_loc == 'V':
1199  x_v = grid_object['X'][:]
1200  y_v = grid_object['Yp1'][:]
1201  z_v = grid_object['Z'][:]
1202  elif v_grid_loc == 'W':
1203  x_v = grid_object['X'][:]
1204  y_v = grid_object['Y'][:]
1205  z_v = grid_object['Zl'][:]
1206  elif v_grid_loc == 'T':
1207  x_v = grid_object['X'][:]
1208  y_v = grid_object['Y'][:]
1209  z_v = grid_object['Z'][:]
1210  elif v_grid_loc == 'Zeta':
1211  x_v = grid_object['Xp1'][:]
1212  y_v = grid_object['Yp1'][:]
1213  z_v = grid_object['Z'][:]
1214  else:
1215  print 'v_grid_loc not set correctly. Possible options are: U,V,W,T, and Zeta'
1216  return
1217 
1218  if w_grid_loc == 'U':
1219  x_w = grid_object['Xp1'][:]
1220  y_w = grid_object['Y'][:]
1221  z_w = grid_object['Z'][:]
1222  elif w_grid_loc == 'V':
1223  x_w = grid_object['X'][:]
1224  y_w = grid_object['Yp1'][:]
1225  z_w = grid_object['Z'][:]
1226  elif w_grid_loc == 'W':
1227  x_w = grid_object['X'][:]
1228  y_w = grid_object['Y'][:]
1229  z_w = grid_object['Zl'][:]
1230  elif w_grid_loc == 'T':
1231  x_w = grid_object['X'][:]
1232  y_w = grid_object['Y'][:]
1233  z_w = grid_object['Z'][:]
1234  elif w_grid_loc == 'Zeta':
1235  x_w = grid_object['Xp1'][:]
1236  y_w = grid_object['Yp1'][:]
1237  z_w = grid_object['Z'][:]
1238  else:
1239  print 'w_grid_loc not set correctly. Possible options are: U,V,W,T, and Zeta'
1240  return
1241 
1242  len_x_u = len(x_u)
1243  len_y_u = len(y_u)
1244  len_z_u = len(z_u)
1245 
1246  len_x_v = len(x_v)
1247  len_y_v = len(y_v)
1248  len_z_v = len(z_v)
1249 
1250  len_x_w = len(x_w)
1251  len_y_w = len(y_w)
1252  len_z_w = len(z_w)
1253 
1254  len_t = len(t)
1255 
1256  if u_bias_field is None:
1257  u_bias_field = np.zeros_like(grid_object['wet_mask_U'][:])
1258  if v_bias_field is None:
1259  v_bias_field = np.zeros_like(grid_object['wet_mask_V'][:])
1260  if w_bias_field is None:
1261  w_bias_field = np.zeros_like(grid_object['wet_mask_W'][1:,...])
1262 
1263 
1264  steps_per_trace = int(trace_length/delta_t)
1265  time_steps_until_jitter = np.random.randint(steps_per_trace, size=n_particles)
1266 
1267  startx = (np.random.rand(n_particles)*
1268  (np.max(grid_object['X'][:]) - np.min(grid_object['X'][:]))) + np.min(grid_object['X'][:])
1269  starty = (np.random.rand(n_particles)*
1270  (np.max(grid_object['Y'][:]) - np.min(grid_object['Y'][:]))) + np.min(grid_object['Y'][:])
1271  startz = (np.random.rand(n_particles)*grid_object['Z'][1])
1272  #(np.max(grid_object['Z'][:]) - np.min(grid_object['X'][:]))) + np.min(grid_object['X'][:])
1273 
1274  x_stream = np.ones((len(startx),int(np.fabs(t_tracking/delta_t))+2))*startx[:,np.newaxis]
1275  y_stream = np.ones((len(startx),int(np.fabs(t_tracking/delta_t))+2))*starty[:,np.newaxis]
1276  z_stream = np.ones((len(startx),int(np.fabs(t_tracking/delta_t))+2))*startz[:,np.newaxis]
1277  t_stream = np.ones((int(np.fabs(t_tracking/delta_t))+2))*startt
1278 
1279  t_RK = startt #set the initial time to be the given start time
1280  z_RK = startz
1281  y_RK = starty
1282  x_RK = startx
1283 
1284  i=0
1285 
1286  u_netcdf_filehandle = netCDF4.Dataset(u_netcdf_filename)
1287  v_netcdf_filehandle = netCDF4.Dataset(v_netcdf_filename)
1288  w_netcdf_filehandle = netCDF4.Dataset(w_netcdf_filename)
1289 
1290  t_index = np.searchsorted(t,t_RK)
1291  t_index_new = np.searchsorted(t,t_RK) # this is later used to test if new data needs to be read in.
1292  if t_index == 0:
1293  raise ValueError('Given time value is outside the given velocity fields - too small')
1294  elif t_index == len_t:
1295  raise ValueError('Given time value is outside the given velocity fields - too big')
1296 
1297 
1298  u_field_before = u_netcdf_filehandle.variables[u_netcdf_variable][t_index,...]
1299  u_field_after = u_netcdf_filehandle.variables[u_netcdf_variable][t_index+1,...]
1300  u_field = (u_field_before + ((u_field_before - u_field_after)*
1301  (t_RK - t[t_index])/(t[t_index+1] - t[t_index])) -
1302  u_bias_field)
1303 
1304  v_field_before = v_netcdf_filehandle.variables[v_netcdf_variable][t_index,...]
1305  v_field_after = v_netcdf_filehandle.variables[v_netcdf_variable][t_index+1,...]
1306  v_field = (v_field_before + ((v_field_before - v_field_after)*
1307  (t_RK - t[t_index])/(t[t_index+1] - t[t_index])) -
1308  v_bias_field)
1309 
1310  w_field_before = w_netcdf_filehandle.variables[w_netcdf_variable][t_index,...]
1311  w_field_after = w_netcdf_filehandle.variables[w_netcdf_variable][t_index+1,...]
1312  w_field = (w_field_before + ((w_field_before - w_field_after)*
1313  (t_RK - t[t_index])/(t[t_index+1] - t[t_index])) -
1314  w_bias_field)
1315 
1316 
1317  # Prepare for spherical polar grids
1318  deg_per_m = np.ones((len(startx),2),dtype=float)
1319  if grid_object['grid_type']=='polar':
1320  deg_per_m[:,0] = np.ones_like(startx)/(1852.*60.) # multiplier for v
1321 
1322  # Runge-Kutta fourth order method to estimate next position.
1323  while i < np.fabs(t_tracking/delta_t):
1324  #t_RK < t_tracking + startt:
1325 
1326  if grid_object['grid_type']=='polar':
1327  # use degrees per metre and convert all the velocities to degrees / second# calculate degrees per metre at current location - used to convert the m/s velocities in to degrees/s
1328  deg_per_m[:,1] = np.cos(starty*np.pi/180.)/(1852.*60.)# multiplier for u
1329  # Compute indices at location given
1330 
1331  if (t_index_new == t_index):
1332  # time hasn't progressed beyond the loaded time slices
1333  pass
1334  else:
1335  t_index = np.searchsorted(t,t_RK)
1336  if t_index == 0:
1337  raise ValueError('Given time value is outside the given velocity fields - too small')
1338  elif t_index == len_t:
1339  raise ValueError('Given time value is outside the given velocity fields - too big')
1340 
1341 
1342  u_field_before = u_netcdf_filehandle.variables[u_netcdf_variable][t_index,...]
1343  u_field_after = u_netcdf_filehandle.variables[u_netcdf_variable][t_index+1,...]
1344  u_field = (u_field_before + ((u_field_before - u_field_after)*
1345  (t_RK - t[t_index])/(t[t_index+1] - t[t_index])) - u_bias_field)
1346 
1347  v_field_before = v_netcdf_filehandle.variables[v_netcdf_variable][t_index,...]
1348  v_field_after = v_netcdf_filehandle.variables[v_netcdf_variable][t_index+1,...]
1349  v_field = (v_field_before + ((v_field_before - v_field_after)*
1350  (t_RK - t[t_index])/(t[t_index+1] - t[t_index])) - v_bias_field)
1351 
1352  w_field_before = w_netcdf_filehandle.variables[w_netcdf_variable][t_index,...]
1353  w_field_after = w_netcdf_filehandle.variables[w_netcdf_variable][t_index+1,...]
1354  w_field = (w_field_before + ((w_field_before - w_field_after)*
1355  (t_RK - t[t_index])/(t[t_index+1] - t[t_index])) - w_bias_field)
1356 
1357  # check if any particles need moving, and then move them
1358  for particle in xrange(0,n_particles):
1359  if time_steps_until_jitter[particle] <= 0:
1360  x_RK[particle] = (np.random.rand(1)*
1361  (np.max(grid_object['X'][:]) - np.min(grid_object['X'][:]))) + np.min(grid_object['X'][:])
1362  y_RK[particle] = (np.random.rand(1)*
1363  (np.max(grid_object['Y'][:]) - np.min(grid_object['Y'][:]))) + np.min(grid_object['Y'][:])
1364  z_RK[particle] = grid_object['Z'][1]
1365  #(np.random.rand(1)*
1366  #(np.max(grid_object['Z'][:]) - np.min(grid_object['Z'][:]))) + np.min(grid_object['Z'][:])
1367  time_steps_until_jitter[particle] = steps_per_trace
1368 
1369 
1370  # Interpolate velocities at initial location
1371  u_loc = np.ones_like(startx)
1372  v_loc = np.ones_like(startx)
1373  w_loc = np.ones_like(startx)
1374 
1375  u_loc = trilinear_interp_arrays(x_RK,y_RK,z_RK,u_field,x_u,y_u,z_u,len_x_u,len_y_u,len_z_u)
1376  v_loc = trilinear_interp_arrays(x_RK,y_RK,z_RK,v_field,x_v,y_v,z_v,len_x_v,len_y_v,len_z_v)
1377  w_loc = trilinear_interp_arrays(x_RK,y_RK,z_RK,w_field,x_w,y_w,z_w,len_x_w,len_y_w,len_z_w)
1378 
1379  u_loc = u_loc*deg_per_m[:,1]
1380  v_loc = v_loc*deg_per_m[:,0]
1381 
1382  dx1 = delta_t*u_loc
1383  dy1 = delta_t*v_loc
1384  dz1 = delta_t*w_loc
1385 
1386  u_loc1 = np.ones_like(startx)
1387  v_loc1 = np.ones_like(startx)
1388  w_loc1 = np.ones_like(startx)
1389 
1390  u_loc1 = trilinear_interp_arrays(x_RK + 0.5*dx1,y_RK + 0.5*dy1,z_RK + 0.5*dz1,
1391  u_field,x_u,y_u,z_u,len_x_u,len_y_u,len_z_u)
1392  v_loc1 = trilinear_interp_arrays(x_RK + 0.5*dx1,y_RK + 0.5*dy1,z_RK + 0.5*dz1,
1393  v_field,x_v,y_v,z_v,len_x_v,len_y_v,len_z_v)
1394  w_loc1 = trilinear_interp_arrays(x_RK + 0.5*dx1,y_RK + 0.5*dy1,z_RK + 0.5*dz1,
1395  w_field,x_w,y_w,z_w,len_x_w,len_y_w,len_z_w)
1396  u_loc1 = u_loc1*deg_per_m[:,1]
1397  v_loc1 = v_loc1*deg_per_m[:,0]
1398  dx2 = delta_t*u_loc1
1399  dy2 = delta_t*v_loc1
1400  dz2 = delta_t*w_loc1
1401 
1402  u_loc2 = np.ones_like(startx)
1403  v_loc2 = np.ones_like(startx)
1404  w_loc2 = np.ones_like(startx)
1405  u_loc2 = trilinear_interp_arrays(x_RK + 0.5*dx2,y_RK + 0.5*dy2,z_RK + 0.5*dz2,
1406  u_field,x_u,y_u,z_u,len_x_u,len_y_u,len_z_u)
1407  v_loc2 = trilinear_interp_arrays(x_RK + 0.5*dx2,y_RK + 0.5*dy2,z_RK + 0.5*dz2,
1408  v_field,x_v,y_v,z_v,len_x_v,len_y_v,len_z_v)
1409  w_loc2 = trilinear_interp_arrays(x_RK + 0.5*dx2,y_RK + 0.5*dy2,z_RK + 0.5*dz2,
1410  w_field,x_w,y_w,z_w,len_x_w,len_y_w,len_z_w)
1411 
1412  u_loc2 = u_loc2*deg_per_m[:,1]
1413  v_loc2 = v_loc2*deg_per_m[:,0]
1414  dx3 = delta_t*u_loc2
1415  dy3 = delta_t*v_loc2
1416  dz3 = delta_t*w_loc2
1417 
1418  u_loc3 = np.ones_like(startx)
1419  v_loc3 = np.ones_like(startx)
1420  w_loc3 = np.ones_like(startx)
1421  u_loc3 = trilinear_interp_arrays(x_RK + dx3,y_RK + dy3,z_RK + dz3,
1422  u_field,x_u,y_u,z_u,len_x_u,len_y_u,len_z_u)
1423  v_loc3 = trilinear_interp_arrays(x_RK + dx3,y_RK + dy3,z_RK + dz3,
1424  v_field,x_v,y_v,z_v,len_x_v,len_y_v,len_z_v)
1425  w_loc3 = trilinear_interp_arrays(x_RK + dx3,y_RK + dy3,z_RK + dz3,
1426  w_field,x_w,y_w,z_w,len_x_w,len_y_w,len_z_w)
1427 
1428  u_loc3 = u_loc3*deg_per_m[:,1]
1429  v_loc3 = v_loc3*deg_per_m[:,0]
1430  dx4 = delta_t*u_loc3
1431  dy4 = delta_t*v_loc3
1432  dz4 = delta_t*w_loc3
1433 
1434  #recycle the variables to keep the code clean
1435  x_RK = x_RK + (dx1 + 2*dx2 + 2*dx3 + dx4)/6
1436  y_RK = y_RK + (dy1 + 2*dy2 + 2*dy3 + dy4)/6
1437  z_RK = z_RK + (dz1 + 2*dz2 + 2*dz3 + dz4)/6
1438  t_RK += delta_t
1439  i += 1
1440  time_steps_until_jitter = time_steps_until_jitter - 1
1441 
1442  x_stream[:,i] = x_RK
1443  y_stream[:,i] = y_RK
1444  z_stream[:,i] = z_RK
1445  t_stream[i] = t_RK
1446 
1447  t_index_new = np.searchsorted(t,t_RK)
1448 
1449 
1450  u_netcdf_filehandle.close()
1451  v_netcdf_filehandle.close()
1452  w_netcdf_filehandle.close()
1453 
1454  return x_stream,y_stream,z_stream,t_stream
1455 
1456 ##########################
1457 
1458 # include xyz version
1459 
1460 def numeric_GLM_xyzt(u_netcdf_filename,v_netcdf_filename,w_netcdf_filename,
1461  n_particles,
1462  startx,starty,startz,startt,
1463  total_time,timestep,
1464  r_x,r_y,r_z,
1465  t,
1466  grid_object,
1467  r_cutoff_factor=3,
1468  u_netcdf_variable='UVEL',
1469  v_netcdf_variable='VVEL',
1470  w_netcdf_variable='WVEL',
1471  u_grid_loc='U',v_grid_loc='V',w_grid_loc='W',
1472  u_bias_field=None,
1473  v_bias_field=None,
1474  w_bias_field=None):
1475  """!A three-dimensional lagrangian particle tracker designed for tracking many particles at once as a method of estimating Generalised Lagrangian-Mean velocities.
1476 
1477  The algorithm is a little slow because it spends time checking that the particles are within a given distance of the centre of mass, that they are in the fluid not the bathymetry and that they are below the surface. This means that the cloud of particles should remain within the fluid and not get trapped by the bathymetry.
1478 
1479  The velocity fields must be four dimensional (three spatial, one temporal) and have units of m/s.
1480  It should work to track particles forwards or backwards in time (set timestep <0 for backwards in time). But, be warned, backwards in time hasn't been thoroughly tested yet.
1481 
1482  Because this is a very large amount of data, the fields are passed as netcdffile handles.
1483 
1484 
1485  ## Returns:
1486  * x_stream, y_stream, z_stream - all with dimensions (particle,time_level)
1487  * t_stream - with dimensions (time_level)
1488  * com_stream - with dimensions (time_level)
1489 
1490  ## The variables are:
1491  * ?_netcdf_filename = name of the netcdf file with ?'s data in it.
1492  * n_particles = number of particles to track
1493  * start? = starting value for x, y, z, or time
1494  * total_time = length of time to track particles for, in seconds. This is always positive
1495  * timestep = timestep for particle tracking algorithm, in seconds. This can be positive or negative.
1496  * r_? = radius of sedding ellipsoid in x, y, or z direction
1497  * t = vector of time levels that are contained in the velocity data.
1498  * grid_object is m.grid if you followed the standard naming conventions.
1499  * ?_netcdf_variable = name of the "?" variable field in the netcdf file.
1500  * ?_grid_loc = where the field "?" is located on the C-grid. Possibles options are, U, V, W, T and Zeta.
1501  * ?_bias_field = bias to add to that velocity field component. If set to -mean(velocity component), then only the time varying portion of that field will be used.
1502  """
1503 
1504  if u_grid_loc == 'U':
1505  x_u = grid_object['Xp1'][:]
1506  y_u = grid_object['Y'][:]
1507  z_u = grid_object['Z'][:]
1508  elif u_grid_loc == 'V':
1509  x_u = grid_object['X'][:]
1510  y_u = grid_object['Yp1'][:]
1511  z_u = grid_object['Z'][:]
1512  elif u_grid_loc == 'W':
1513  x_u = grid_object['X'][:]
1514  y_u = grid_object['Y'][:]
1515  z_u = grid_object['Zl'][:]
1516  elif u_grid_loc == 'T':
1517  x_u = grid_object['X'][:]
1518  y_u = grid_object['Y'][:]
1519  z_u = grid_object['Z'][:]
1520  elif u_grid_loc == 'Zeta':
1521  x_u = grid_object['Xp1'][:]
1522  y_u = grid_object['Yp1'][:]
1523  z_u = grid_object['Z'][:]
1524  else:
1525  print 'u_grid_loc not set correctly. Possible options are: U,V,W,T, and Zeta'
1526  return
1527 
1528  if v_grid_loc == 'U':
1529  x_v = grid_object['Xp1'][:]
1530  y_v = grid_object['Y'][:]
1531  z_v = grid_object['Z'][:]
1532  elif v_grid_loc == 'V':
1533  x_v = grid_object['X'][:]
1534  y_v = grid_object['Yp1'][:]
1535  z_v = grid_object['Z'][:]
1536  elif v_grid_loc == 'W':
1537  x_v = grid_object['X'][:]
1538  y_v = grid_object['Y'][:]
1539  z_v = grid_object['Zl'][:]
1540  elif v_grid_loc == 'T':
1541  x_v = grid_object['X'][:]
1542  y_v = grid_object['Y'][:]
1543  z_v = grid_object['Z'][:]
1544  elif v_grid_loc == 'Zeta':
1545  x_v = grid_object['Xp1'][:]
1546  y_v = grid_object['Yp1'][:]
1547  z_v = grid_object['Z'][:]
1548  else:
1549  print 'v_grid_loc not set correctly. Possible options are: U,V,W,T, and Zeta'
1550  return
1551 
1552  if w_grid_loc == 'U':
1553  x_w = grid_object['Xp1'][:]
1554  y_w = grid_object['Y'][:]
1555  z_w = grid_object['Z'][:]
1556  elif w_grid_loc == 'V':
1557  x_w = grid_object['X'][:]
1558  y_w = grid_object['Yp1'][:]
1559  z_w = grid_object['Z'][:]
1560  elif w_grid_loc == 'W':
1561  x_w = grid_object['X'][:]
1562  y_w = grid_object['Y'][:]
1563  z_w = grid_object['Zl'][:]
1564  elif w_grid_loc == 'T':
1565  x_w = grid_object['X'][:]
1566  y_w = grid_object['Y'][:]
1567  z_w = grid_object['Z'][:]
1568  elif w_grid_loc == 'Zeta':
1569  x_w = grid_object['Xp1'][:]
1570  y_w = grid_object['Yp1'][:]
1571  z_w = grid_object['Z'][:]
1572  else:
1573  print 'w_grid_loc not set correctly. Possible options are: U,V,W,T, and Zeta'
1574  return
1575 
1576  len_x_u = len(x_u)
1577  len_y_u = len(y_u)
1578  len_z_u = len(z_u)
1579 
1580  len_x_v = len(x_v)
1581  len_y_v = len(y_v)
1582  len_z_v = len(z_v)
1583 
1584  len_x_w = len(x_w)
1585  len_y_w = len(y_w)
1586  len_z_w = len(z_w)
1587 
1588  len_t = len(t)
1589 
1590  if u_bias_field is None:
1591  u_bias_field = np.zeros_like(grid_object['wet_mask_U'][:])
1592  if v_bias_field is None:
1593  v_bias_field = np.zeros_like(grid_object['wet_mask_V'][:])
1594  if w_bias_field is None:
1595  w_bias_field = np.zeros_like(grid_object['wet_mask_W'][1:,...])
1596 
1597 
1598  ini_x = ((np.random.rand(n_particles)-0.5)*2*r_x + startx)
1599  ini_y = ((np.random.rand(n_particles)-0.5)*2*r_y + starty)
1600  ini_z = ((np.random.rand(n_particles)-0.5)*2*r_z + startz)
1601 
1602  x_stream = np.ones((n_particles,int(np.fabs(total_time/timestep))+2))*startx
1603  y_stream = np.ones((n_particles,int(np.fabs(total_time/timestep))+2))*starty
1604  z_stream = np.ones((n_particles,int(np.fabs(total_time/timestep))+2))*startz
1605  t_stream = np.ones((int(np.fabs(total_time/timestep))+2))*startt
1606  com_stream = np.ones((3,int(np.fabs(total_time/timestep))+2))
1607 
1608  t_RK = startt #set the initial time to be the given start time
1609  z_RK = ini_z
1610  y_RK = ini_y
1611  x_RK = ini_x
1612 
1613  # define centre of mass variable
1614  x_com = startx
1615  y_com = starty
1616  z_com = startz
1617 
1618  normed_radius = ( ((ini_x - startx)**2)/(r_x**2) +
1619  ((ini_y - starty)**2)/(r_y**2) +
1620  ((ini_z - startz)**2)/(r_z**2) )
1621 
1622  wet_test = trilinear_interp_arrays(x_RK,y_RK,z_RK,grid_object['wet_mask_U'][:],x_u,y_u,z_u,len_x_u,len_y_u,len_z_u)
1623 
1624  for particle in xrange(n_particles): # check if any particles outside cutoffs
1625  if (normed_radius[particle] > r_cutoff_factor or wet_test[particle] < 0.5 or z_RK[particle]*grid_object['Z'][2] < 0): # find specific particles
1626  while (normed_radius[particle] > 1 or wet_test[particle] < 0.5 or z_RK[particle]*grid_object['Z'][2] < 0): # reseed inside ellipsoid
1627  x_RK[particle] = ((np.random.rand(1)-0.5)*r_x*2 + x_com)
1628  y_RK[particle] = ((np.random.rand(1)-0.5)*r_y*2 + y_com)
1629  z_RK[particle] = ((np.random.rand(1)-0.5)*r_z*2 + z_com)
1630  normed_radius[particle] = ( ((x_RK[particle] - x_com)**2)/(r_x**2) +
1631  ((y_RK[particle] - y_com)**2)/(r_y**2) +
1632  ((z_RK[particle] - z_com)**2)/(r_z**2) )
1633  wet_test[particle] = trilinear_interp(x_RK[particle],y_RK[particle],z_RK[particle],
1634  grid_object['wet_mask_U'][:],x_u,y_u,z_u,len_x_u,len_y_u,len_z_u)
1635 
1636 
1637 
1638 
1639 
1640  i=0
1641 
1642  u_netcdf_filehandle = netCDF4.Dataset(u_netcdf_filename)
1643  v_netcdf_filehandle = netCDF4.Dataset(v_netcdf_filename)
1644  w_netcdf_filehandle = netCDF4.Dataset(w_netcdf_filename)
1645 
1646  t_index = np.searchsorted(t,t_RK)
1647  t_index_new = np.searchsorted(t,t_RK) # this is later used to test if new data needs to be read in.
1648  if t_index == 0:
1649  raise ValueError('Given time value is outside the given velocity fields - too small')
1650  elif t_index == len_t:
1651  raise ValueError('Given time value is outside the given velocity fields - too big')
1652 
1653 
1654  u_field_before = u_netcdf_filehandle.variables[u_netcdf_variable][t_index,...]
1655  u_field_after = u_netcdf_filehandle.variables[u_netcdf_variable][t_index+1,...]
1656  u_field = (u_field_before + ((u_field_before - u_field_after)*
1657  (t_RK - t[t_index])/(t[t_index+1] - t[t_index])) -
1658  u_bias_field)
1659 
1660  v_field_before = v_netcdf_filehandle.variables[v_netcdf_variable][t_index,...]
1661  v_field_after = v_netcdf_filehandle.variables[v_netcdf_variable][t_index+1,...]
1662  v_field = (v_field_before + ((v_field_before - v_field_after)*
1663  (t_RK - t[t_index])/(t[t_index+1] - t[t_index])) -
1664  v_bias_field)
1665 
1666  w_field_before = w_netcdf_filehandle.variables[w_netcdf_variable][t_index,...]
1667  w_field_after = w_netcdf_filehandle.variables[w_netcdf_variable][t_index+1,...]
1668  w_field = (w_field_before + ((w_field_before - w_field_after)*
1669  (t_RK - t[t_index])/(t[t_index+1] - t[t_index])) -
1670  w_bias_field)
1671 
1672 
1673  # Prepare for spherical polar grids
1674  deg_per_m = np.ones((n_particles,2),dtype=float)
1675  if grid_object['grid_type']=='polar':
1676  deg_per_m[:,0] = np.ones((n_particles))/(1852.*60.) # multiplier for v
1677 
1678  # Runge-Kutta fourth order method to estimate next position.
1679  while i < np.fabs(total_time/timestep):
1680  #t_RK < total_time + startt:
1681 
1682  if grid_object['grid_type']=='polar':
1683  # use degrees per metre and convert all the velocities to degrees / second# calculate degrees per metre at current location - used to convert the m/s velocities in to degrees/s
1684  deg_per_m[:,1] = np.cos(np.ones((n_particles))*starty*np.pi/180.)/(1852.*60.)# multiplier for u
1685  # Compute indices at location given
1686 
1687  if (t_index_new == t_index):
1688  # time hasn't progressed beyond the loaded time slices
1689  pass
1690  else:
1691  t_index = np.searchsorted(t,t_RK)
1692  if t_index == 0:
1693  raise ValueError('Given time value is outside the given velocity fields - too small')
1694  elif t_index == len_t:
1695  raise ValueError('Given time value is outside the given velocity fields - too big')
1696 
1697 
1698  u_field_before = u_netcdf_filehandle.variables[u_netcdf_variable][t_index,...]
1699  u_field_after = u_netcdf_filehandle.variables[u_netcdf_variable][t_index+1,...]
1700  u_field = (u_field_before + ((u_field_before - u_field_after)*
1701  (t_RK - t[t_index])/(t[t_index+1] - t[t_index])) - u_bias_field)
1702 
1703  v_field_before = v_netcdf_filehandle.variables[v_netcdf_variable][t_index,...]
1704  v_field_after = v_netcdf_filehandle.variables[v_netcdf_variable][t_index+1,...]
1705  v_field = (v_field_before + ((v_field_before - v_field_after)*
1706  (t_RK - t[t_index])/(t[t_index+1] - t[t_index])) - v_bias_field)
1707 
1708  w_field_before = w_netcdf_filehandle.variables[w_netcdf_variable][t_index,...]
1709  w_field_after = w_netcdf_filehandle.variables[w_netcdf_variable][t_index+1,...]
1710  w_field = (w_field_before + ((w_field_before - w_field_after)*
1711  (t_RK - t[t_index])/(t[t_index+1] - t[t_index])) - w_bias_field)
1712 
1713 
1714  # check if any particles need moving, and then reseed them inside an ellipsoid around centre of mass with the original size
1715  for particle in xrange(n_particles): # check if any particles outside cutoff
1716  if (normed_radius[particle] > r_cutoff_factor or wet_test[particle] < 0.5 or z_RK[particle]*grid_object['Z'][2] < 0): # find specific particles
1717  while (normed_radius[particle] > 1 or wet_test[particle] < 0.5 or z_RK[particle]*grid_object['Z'][2] < 0): # reseed inside ellipsoid
1718  x_RK[particle] = ((np.random.rand(1)-0.5)*r_x*2 + x_com)
1719  y_RK[particle] = ((np.random.rand(1)-0.5)*r_y*2 + y_com)
1720  z_RK[particle] = ((np.random.rand(1)-0.5)*r_z*2 + z_com)
1721  normed_radius[particle] = ( ((x_RK[particle] - x_com)**2)/(r_x**2) +
1722  ((y_RK[particle] - y_com)**2)/(r_y**2) +
1723  ((z_RK[particle] - z_com)**2)/(r_z**2) )
1724  wet_test[particle] = trilinear_interp(x_RK[particle],y_RK[particle],z_RK[particle],
1725  grid_object['wet_mask_U'][:],x_u,y_u,z_u,len_x_u,len_y_u,len_z_u)
1726 
1727 
1728  # Interpolate velocities at initial location
1729 
1730  u_loc = trilinear_interp_arrays(x_RK,y_RK,z_RK,u_field,x_u,y_u,z_u,len_x_u,len_y_u,len_z_u)
1731  v_loc = trilinear_interp_arrays(x_RK,y_RK,z_RK,v_field,x_v,y_v,z_v,len_x_v,len_y_v,len_z_v)
1732  w_loc = trilinear_interp_arrays(x_RK,y_RK,z_RK,w_field,x_w,y_w,z_w,len_x_w,len_y_w,len_z_w)
1733  u_loc = u_loc*deg_per_m[:,1]
1734  v_loc = v_loc*deg_per_m[:,0]
1735  dx1 = timestep*u_loc
1736  dy1 = timestep*v_loc
1737  dz1 = timestep*w_loc
1738 
1739 
1740  u_loc1 = trilinear_interp_arrays(x_RK + 0.5*dx1,y_RK + 0.5*dy1,z_RK + 0.5*dz1,
1741  u_field,x_u,y_u,z_u,len_x_u,len_y_u,len_z_u)
1742  v_loc1 = trilinear_interp_arrays(x_RK + 0.5*dx1,y_RK + 0.5*dy1,z_RK + 0.5*dz1,
1743  v_field,x_v,y_v,z_v,len_x_v,len_y_v,len_z_v)
1744  w_loc1 = trilinear_interp_arrays(x_RK + 0.5*dx1,y_RK + 0.5*dy1,z_RK + 0.5*dz1,
1745  w_field,x_w,y_w,z_w,len_x_w,len_y_w,len_z_w)
1746  u_loc1 = u_loc1*deg_per_m[:,1]
1747  v_loc1 = v_loc1*deg_per_m[:,0]
1748  dx2 = timestep*u_loc1
1749  dy2 = timestep*v_loc1
1750  dz2 = timestep*w_loc1
1751 
1752 
1753  u_loc2 = trilinear_interp_arrays(x_RK + 0.5*dx2,y_RK + 0.5*dy2,z_RK + 0.5*dz2,
1754  u_field,x_u,y_u,z_u,len_x_u,len_y_u,len_z_u)
1755  v_loc2 = trilinear_interp_arrays(x_RK + 0.5*dx2,y_RK + 0.5*dy2,z_RK + 0.5*dz2,
1756  v_field,x_v,y_v,z_v,len_x_v,len_y_v,len_z_v)
1757  w_loc2 = trilinear_interp_arrays(x_RK + 0.5*dx2,y_RK + 0.5*dy2,z_RK + 0.5*dz2,
1758  w_field,x_w,y_w,z_w,len_x_w,len_y_w,len_z_w)
1759  u_loc2 = u_loc2*deg_per_m[:,1]
1760  v_loc2 = v_loc2*deg_per_m[:,0]
1761  dx3 = timestep*u_loc2
1762  dy3 = timestep*v_loc2
1763  dz3 = timestep*w_loc2
1764 
1765 
1766  u_loc3 = trilinear_interp_arrays(x_RK + dx3,y_RK + dy3,z_RK + dz3,
1767  u_field,x_u,y_u,z_u,len_x_u,len_y_u,len_z_u)
1768  v_loc3 = trilinear_interp_arrays(x_RK + dx3,y_RK + dy3,z_RK + dz3,
1769  v_field,x_v,y_v,z_v,len_x_v,len_y_v,len_z_v)
1770  w_loc3 = trilinear_interp_arrays(x_RK + dx3,y_RK + dy3,z_RK + dz3,
1771  w_field,x_w,y_w,z_w,len_x_w,len_y_w,len_z_w)
1772 
1773  u_loc3 = u_loc3*deg_per_m[:,1]
1774  v_loc3 = v_loc3*deg_per_m[:,0]
1775  dx4 = timestep*u_loc3
1776  dy4 = timestep*v_loc3
1777  dz4 = timestep*w_loc3
1778 
1779  #recycle the variables to keep the code clean
1780  x_RK = x_RK + (dx1 + 2*dx2 + 2*dx3 + dx4)/6
1781  y_RK = y_RK + (dy1 + 2*dy2 + 2*dy3 + dy4)/6
1782  z_RK = z_RK + (dz1 + 2*dz2 + 2*dz3 + dz4)/6
1783  t_RK += timestep
1784  i += 1
1785 
1786  x_com = np.mean(x_RK)
1787  y_com = np.mean(y_RK)
1788  z_com = np.mean(z_RK)
1789 
1790  x_stream[:,i] = x_RK
1791  y_stream[:,i] = y_RK
1792  z_stream[:,i] = z_RK
1793  t_stream[i] = t_RK
1794 
1795  com_stream[0,i] = x_com
1796  com_stream[1,i] = y_com
1797  com_stream[2,i] = z_com
1798 
1799  t_index_new = np.searchsorted(t,t_RK)
1800  normed_radius = ( ((x_RK - x_com)**2)/(r_x**2) +
1801  ((y_RK - y_com)**2)/(r_y**2) +
1802  ((z_RK - z_com)**2)/(r_z**2) )
1803  wet_test = trilinear_interp_arrays(x_RK,y_RK,z_RK,grid_object['wet_mask_U'][:],x_u,y_u,z_u,len_x_u,len_y_u,len_z_u)
1804 
1805  u_netcdf_filehandle.close()
1806  v_netcdf_filehandle.close()
1807  w_netcdf_filehandle.close()
1808 
1809  return x_stream,y_stream,z_stream,t_stream, com_stream
1810 ################################################
1811 
1812 def numeric_GLM_xyz(u,v,w,
1813  n_particles,
1814  startx,starty,startz,startt,
1815  total_time,timestep,
1816  r_x,r_y,r_z,
1817  grid_object,
1818  r_cutoff_factor=3,
1819  x_v='None',y_v='None',z_v='None',
1820  x_w='None',y_w='None',z_w='None',
1821  u_grid_loc='U',v_grid_loc='V',w_grid_loc='W'):
1822 
1823 
1824  """!A three-dimensional streamline solver. The velocity fields must be three dimensional and not vary in time.
1825  X_grid_loc variables specify where the field "X" is located on the C-grid. Possibles options are, U, V, W, T and Zeta.
1826 
1827  Returns:
1828  * x_stream, y_stream, z_stream - all with dimensions (particle,time_level)
1829  * t_stream - with dimensions (time_level)
1830  * com_stream - with dimensions (3(x,y,z),time_level)
1831 """
1832  if grid_object:
1833  if u_grid_loc == 'U':
1834  x_u = copy.deepcopy(grid_object['Xp1'][:])
1835  y_u = copy.deepcopy(grid_object['Y'][:])
1836  z_u = copy.deepcopy(grid_object['Z'][:])
1837  elif u_grid_loc == 'V':
1838  x_u = copy.deepcopy(grid_object['X'][:])
1839  y_u = copy.deepcopy(grid_object['Yp1'][:])
1840  z_u = copy.deepcopy(grid_object['Z'][:])
1841  elif u_grid_loc == 'W':
1842  x_u = copy.deepcopy(grid_object['X'][:])
1843  y_u = copy.deepcopy(grid_object['Y'][:])
1844  z_u = copy.deepcopy(grid_object['Zl'][:])
1845  elif u_grid_loc == 'T':
1846  x_u = copy.deepcopy(grid_object['X'][:])
1847  y_u = copy.deepcopy(grid_object['Y'][:])
1848  z_u = copy.deepcopy(grid_object['Z'][:])
1849  elif u_grid_loc == 'Zeta':
1850  x_u = copy.deepcopy(grid_object['Xp1'][:])
1851  y_u = copy.deepcopy(grid_object['Yp1'][:])
1852  z_u = copy.deepcopy(grid_object['Z'][:])
1853  else:
1854  print 'u_grid_loc not set correctly. Possible options are: U,V,W,T, and Zeta'
1855  return
1856 
1857  if v_grid_loc == 'U':
1858  x_v = copy.deepcopy(grid_object['Xp1'][:])
1859  y_v = copy.deepcopy(grid_object['Y'][:])
1860  z_v = copy.deepcopy(grid_object['Z'][:])
1861  elif v_grid_loc == 'V':
1862  x_v = copy.deepcopy(grid_object['X'][:])
1863  y_v = copy.deepcopy(grid_object['Yp1'][:])
1864  z_v = copy.deepcopy(grid_object['Z'][:])
1865  elif v_grid_loc == 'W':
1866  x_v = copy.deepcopy(grid_object['X'][:])
1867  y_v = copy.deepcopy(grid_object['Y'][:])
1868  z_v = copy.deepcopy(grid_object['Zl'][:])
1869  elif v_grid_loc == 'T':
1870  x_v = copy.deepcopy(grid_object['X'][:])
1871  y_v = copy.deepcopy(grid_object['Y'][:])
1872  z_v = copy.deepcopy(grid_object['Z'][:])
1873  elif v_grid_loc == 'Zeta':
1874  x_v = copy.deepcopy(grid_object['Xp1'][:])
1875  y_v = copy.deepcopy(grid_object['Yp1'][:])
1876  z_v = copy.deepcopy(grid_object['Z'][:])
1877  else:
1878  print 'v_grid_loc not set correctly. Possible options are: U,V,W,T, and Zeta'
1879  return
1880 
1881  if w_grid_loc == 'U':
1882  x_w = copy.deepcopy(grid_object['Xp1'][:])
1883  y_w = copy.deepcopy(grid_object['Y'][:])
1884  z_w = copy.deepcopy(grid_object['Z'][:])
1885  elif w_grid_loc == 'V':
1886  x_w = copy.deepcopy(grid_object['X'][:])
1887  y_w = copy.deepcopy(grid_object['Yp1'][:])
1888  z_w = copy.deepcopy(grid_object['Z'][:])
1889  elif w_grid_loc == 'W':
1890  x_w = copy.deepcopy(grid_object['X'][:])
1891  y_w = copy.deepcopy(grid_object['Y'][:])
1892  z_w = copy.deepcopy(grid_object['Zl'][:])
1893  elif w_grid_loc == 'T':
1894  x_w = copy.deepcopy(grid_object['X'][:])
1895  y_w = copy.deepcopy(grid_object['Y'][:])
1896  z_w = copy.deepcopy(grid_object['Z'][:])
1897  elif w_grid_loc == 'Zeta':
1898  x_w = copy.deepcopy(grid_object['Xp1'][:])
1899  y_w = copy.deepcopy(grid_object['Yp1'][:])
1900  z_w = copy.deepcopy(grid_object['Z'][:])
1901  else:
1902  print 'w_grid_loc not set correctly. Possible options are: U,V,W,T, and Zeta'
1903  return
1904 
1905 
1906  len_x_u = len(x_u)
1907  len_y_u = len(y_u)
1908  len_z_u = len(z_u)
1909 
1910  len_x_v = len(x_v)
1911  len_y_v = len(y_v)
1912  len_z_v = len(z_v)
1913 
1914  len_x_w = len(x_w)
1915  len_y_w = len(y_w)
1916  len_z_w = len(z_w)
1917 
1918 
1919  ini_x = ((np.random.rand(n_particles)-0.5)*2*r_x + startx)
1920  ini_y = ((np.random.rand(n_particles)-0.5)*2*r_y + starty)
1921  ini_z = ((np.random.rand(n_particles)-0.5)*2*r_z + startz)
1922 
1923  x_stream = np.ones((n_particles,int(np.fabs(total_time/timestep))+2))*startx
1924  y_stream = np.ones((n_particles,int(np.fabs(total_time/timestep))+2))*starty
1925  z_stream = np.ones((n_particles,int(np.fabs(total_time/timestep))+2))*startz
1926  t_stream = np.ones((int(np.fabs(total_time/timestep))+2))*startt
1927  com_stream = np.ones((3,int(np.fabs(total_time/timestep))+2))
1928 
1929  t_RK = startt #set the initial time to be the given start time
1930  z_RK = ini_z
1931  y_RK = ini_y
1932  x_RK = ini_x
1933 
1934  # define centre of mass variable
1935  x_com = startx
1936  y_com = starty
1937  z_com = startz
1938  t = startt
1939 
1940  normed_radius = ( ((ini_x - startx)**2)/(r_x**2) +
1941  ((ini_y - starty)**2)/(r_y**2) +
1942  ((ini_z - startz)**2)/(r_z**2) )
1943 
1944  wet_test = trilinear_interp_arrays(x_RK,y_RK,z_RK,grid_object['wet_mask_U'][:],x_u,y_u,z_u,len_x_u,len_y_u,len_z_u)
1945 
1946  for particle in xrange(n_particles): # check if any particles outside cutoff
1947  if (normed_radius[particle] > r_cutoff_factor or wet_test[particle] < 0.5 or z_RK[particle]*grid_object['Z'][2] < 0): # find specific particles
1948  while (normed_radius[particle] > 1 or wet_test[particle] < 0.5 or z_RK[particle]*grid_object['Z'][2] < 0): # reseed inside ellipsoid
1949  x_RK[particle] = ((np.random.rand(1)-0.5)*r_x*2 + x_com)
1950  y_RK[particle] = ((np.random.rand(1)-0.5)*r_y*2 + y_com)
1951  z_RK[particle] = ((np.random.rand(1)-0.5)*r_z*2 + z_com)
1952  normed_radius[particle] = ( ((x_RK[particle] - x_com)**2)/(r_x**2) +
1953  ((y_RK[particle] - y_com)**2)/(r_y**2) +
1954  ((z_RK[particle] - z_com)**2)/(r_z**2) )
1955  wet_test[particle] = trilinear_interp(x_RK[particle],y_RK[particle],z_RK[particle],
1956  grid_object['wet_mask_U'][:],x_u,y_u,z_u,len_x_u,len_y_u,len_z_u)
1957 
1958 
1959  i=0
1960 
1961  # Prepare for spherical polar grids
1962  deg_per_m = np.ones((n_particles,2),dtype=float)
1963  if grid_object['grid_type']=='polar':
1964  deg_per_m[:,0] = np.ones((n_particles))/(1852.*60.) # multiplier for v
1965 
1966  # Runge-Kutta fourth order method to estimate next position.
1967  while i < np.fabs(total_time/timestep):
1968  #t_RK < total_time:
1969  if grid_object['grid_type']=='polar':
1970  # use degrees per metre and convert all the velocities to degrees / second# calculate degrees per metre at current location - used to convert the m/s velocities in to degrees/s
1971  deg_per_m[:,1] = np.cos(np.ones((n_particles))*starty*np.pi/180.)/(1852.*60.)# multiplier for u
1972  # Compute indices at location given
1973 
1974  # check if any particles need moving, and then reseed them inside an ellipsoid around centre of mass with the original size
1975  for particle in xrange(n_particles): # check if any particles outside cutoff
1976  if (normed_radius[particle] > r_cutoff_factor or wet_test[particle] < 0.5 or z_RK[particle]*grid_object['Z'][2] < 0): # find specific particles
1977  while (normed_radius[particle] > 1 or wet_test[particle] < 0.5 or z_RK[particle]*grid_object['Z'][2] < 0): # reseed inside ellipsoid
1978  x_RK[particle] = ((np.random.rand(1)-0.5)*r_x*2 + x_com)
1979  y_RK[particle] = ((np.random.rand(1)-0.5)*r_y*2 + y_com)
1980  z_RK[particle] = ((np.random.rand(1)-0.5)*r_z*2 + z_com)
1981  normed_radius[particle] = ( ((x_RK[particle] - x_com)**2)/(r_x**2) +
1982  ((y_RK[particle] - y_com)**2)/(r_y**2) +
1983  ((z_RK[particle] - z_com)**2)/(r_z**2) )
1984  wet_test[particle] = trilinear_interp(x_RK[particle],y_RK[particle],z_RK[particle],
1985  grid_object['wet_mask_U'][:],x_u,y_u,z_u,len_x_u,len_y_u,len_z_u)
1986 
1987  # Interpolate velocities to initial location
1988  u_loc = trilinear_interp_arrays(x_RK,y_RK,z_RK,u,x_u,y_u,z_u,len_x_u,len_y_u,len_z_u)
1989  v_loc = trilinear_interp_arrays(x_RK,y_RK,z_RK,v,x_v,y_v,z_v,len_x_v,len_y_v,len_z_v)
1990  w_loc = trilinear_interp_arrays(x_RK,y_RK,z_RK,w,x_w,y_w,z_w,len_x_w,len_y_w,len_z_w)
1991  u_loc = u_loc*deg_per_m[:,1]
1992  v_loc = v_loc*deg_per_m[:,0]
1993  dx1 = timestep*u_loc
1994  dy1 = timestep*v_loc
1995  dz1 = timestep*w_loc
1996 
1997  u_loc1 = trilinear_interp_arrays(x_RK + 0.5*dx1,y_RK + 0.5*dy1,z_RK + 0.5*dz1,u,x_u,y_u,z_u,len_x_u,len_y_u,len_z_u)
1998  v_loc1 = trilinear_interp_arrays(x_RK + 0.5*dx1,y_RK + 0.5*dy1,z_RK + 0.5*dz1,v,x_v,y_v,z_v,len_x_v,len_y_v,len_z_v)
1999  w_loc1 = trilinear_interp_arrays(x_RK + 0.5*dx1,y_RK + 0.5*dy1,z_RK + 0.5*dz1,w,x_w,y_w,z_w,len_x_w,len_y_w,len_z_w)
2000  u_loc1 = u_loc1*deg_per_m[:,1]
2001  v_loc1 = v_loc1*deg_per_m[:,0]
2002  dx2 = timestep*u_loc1
2003  dy2 = timestep*v_loc1
2004  dz2 = timestep*w_loc1
2005 
2006  u_loc2 = trilinear_interp_arrays(x_RK + 0.5*dx2,y_RK + 0.5*dy2,z_RK + 0.5*dz2,u,x_u,y_u,z_u,len_x_u,len_y_u,len_z_u)
2007  v_loc2 = trilinear_interp_arrays(x_RK + 0.5*dx2,y_RK + 0.5*dy2,z_RK + 0.5*dz2,v,x_v,y_v,z_v,len_x_v,len_y_v,len_z_v)
2008  w_loc2 = trilinear_interp_arrays(x_RK + 0.5*dx2,y_RK + 0.5*dy2,z_RK + 0.5*dz2,w,x_w,y_w,z_w,len_x_w,len_y_w,len_z_w)
2009  u_loc2 = u_loc2*deg_per_m[:,1]
2010  v_loc2 = v_loc2*deg_per_m[:,0]
2011  dx3 = timestep*u_loc2
2012  dy3 = timestep*v_loc2
2013  dz3 = timestep*w_loc2
2014 
2015  u_loc3 = trilinear_interp_arrays(x_RK + dx3,y_RK + dy3,z_RK + dz3,u,x_u,y_u,z_u,len_x_u,len_y_u,len_z_u)
2016  v_loc3 = trilinear_interp_arrays(x_RK + dx3,y_RK + dy3,z_RK + dz3,v,x_v,y_v,z_v,len_x_v,len_y_v,len_z_v)
2017  w_loc3 = trilinear_interp_arrays(x_RK + dx3,y_RK + dy3,z_RK + dz3,w,x_w,y_w,z_w,len_x_w,len_y_w,len_z_w)
2018  u_loc3 = u_loc3*deg_per_m[:,1]
2019  v_loc3 = v_loc3*deg_per_m[:,0]
2020  dx4 = timestep*u_loc3
2021  dy4 = timestep*v_loc3
2022  dz4 = timestep*w_loc3
2023 
2024  #recycle the "start_" variables to keep the code clean
2025  x_RK = x_RK + (dx1 + 2*dx2 + 2*dx3 + dx4)/6
2026  y_RK = y_RK + (dy1 + 2*dy2 + 2*dy3 + dy4)/6
2027  z_RK = z_RK + (dz1 + 2*dz2 + 2*dz3 + dz4)/6
2028  t_RK += timestep
2029  i += 1
2030 
2031  x_com = np.mean(x_RK)
2032  y_com = np.mean(y_RK)
2033  z_com = np.mean(z_RK)
2034 
2035  x_stream[:,i] = x_RK
2036  y_stream[:,i] = y_RK
2037  z_stream[:,i] = z_RK
2038  t_stream[i] = t_RK
2039 
2040  com_stream[0,i] = x_com
2041  com_stream[1,i] = y_com
2042  com_stream[2,i] = z_com
2043 
2044  normed_radius = ( ((x_RK - x_com)**2)/(r_x**2) +
2045  ((y_RK - y_com)**2)/(r_y**2) +
2046  ((z_RK - z_com)**2)/(r_z**2) )
2047 
2048  wet_test = trilinear_interp_arrays(x_RK,y_RK,z_RK,grid_object['wet_mask_U'][:],x_u,y_u,z_u,len_x_u,len_y_u,len_z_u)
2049 
2050 
2051  return x_stream,y_stream,z_stream,t_stream, com_stream
2052 #####################################################
2053 
2054 
2055 
2056 def bilinear_interp(x0,y0,field,x,y,len_x,len_y):
2057  """!Do bilinear interpolation of a field. This function assumes that the grid can locally be regarded as cartesian, with everything at right angles.
2058 
2059  x0,y0 are the points to interpolate to.
2060  """
2061 
2062  # Compute indeces at location
2063  x_index = np.searchsorted(x,x0)
2064  if x_index == 0:
2065  x_index =1 # a dirty hack to deal with streamlines coming near the edge
2066  #raise ValueError('x location ', str(x0), ' is outside the model grid - too small')
2067  elif x_index == len_x:
2068  x_index =len_x - 1 # a dirty hack to deal with streamlines coming near the edge
2069  #raise ValueError('x location ', str(x0), ' is outside the model grid - too big')
2070 
2071  y_index = np.searchsorted(y,y0)
2072  if y_index == 0:
2073  y_index =1 # a dirty hack to deal with streamlines coming near the edge
2074  #raise ValueError('y location ', str(y0), ' is outside the model grid - too small')
2075  elif y_index == len_y:
2076  y_index =len_y - 1 # a dirty hack to deal with streamlines coming near the edge
2077  #raise ValueError('y location ', str(y0), ' is outside the model grid - too big')
2078 
2079  #print 'x index = ' + str(x_index)
2080  #print 'y index = ' + str(y_index)
2081 
2082  field_interp = actual_bilinear_interp(field,x0,y0,x,y,len_x,len_y,x_index,y_index)
2083  return field_interp
2084 
2085 @numba.jit
2086 def actual_bilinear_interp(field,x0,y0,x,y,len_x,len_y,x_index,y_index):
2087  """!This is a numba accelerated bilinear interpolation. The @numba.jit decorator just above this function causes it to be compiled just before it is run. This introduces a small, Order(1 second), overhead the first time, but not on subsequent calls.
2088  """
2089  field_interp = ((field[y_index-1,x_index-1]*(x[x_index] - x0)*(y[y_index] - y0) +
2090  field[y_index-1,x_index]*(x0 - x[x_index-1])*(y[y_index] - y0) +
2091  field[y_index,x_index]*(x[x_index] - x0)*(y0 - y[y_index-1]) +
2092  field[y_index,x_index]*(x0 - x[x_index-1])*(y0 - y[y_index-1]))/
2093  ((y[y_index] - y[y_index-1])*(x[x_index] - x[x_index-1])))
2094  return field_interp
2095 
2096 def trilinear_interp(x0,y0,z0,field,x,y,z,len_x,len_y,len_z):
2097  """!Do trilinear interpolation of the field in three spatial dimensions. This function assumes that the grid can locally be regarded as cartesian, with everything at right angles. It also requires that the entire field can be held in memory at the same time.
2098 
2099  x0,y0, and z0 represent the point to interpolate to.
2100  """
2101 
2102  # Compute indices at location given
2103  x_index = x.searchsorted(x0)
2104  if x_index == 0:
2105  x_index =1 # a dirty hack to deal with streamlines coming near the edge
2106  #raise ValueError('x location ', str(x0), ' is outside the model grid - too small')
2107  elif x_index == len_x:
2108  x_index =len_x - 1 # a dirty hack to deal with streamlines coming near the edge
2109  #raise ValueError('x location ', str(x0), ' is outside the model grid - too big')
2110 
2111  y_index = y.searchsorted(y0)
2112  if y_index == 0:
2113  y_index =1 # a dirty hack to deal with streamlines coming near the edge
2114  #raise ValueError('y location ', str(y0), ' is outside the model grid - too small')
2115  elif y_index == len_y:
2116  y_index =len_y - 1 # a dirty hack to deal with streamlines coming near the edge
2117  #raise ValueError('y location ', str(y0), ' is outside the model grid - too big')
2118 
2119  # np.searchsorted only works for positive arrays :/
2120  if z0 < 0:
2121  z0 = -z0
2122  z = -z
2123  z_index = z.searchsorted(z0)
2124  if z_index == 0:
2125  z_index =1 # a dirty hack to deal with streamlines coming near the surface
2126  #raise ValueError('z location ', str(z0), ' is outside the model grid - too small')
2127  elif z_index == len_z:
2128  z_index = len_z - 1 # a dirty hack to deal with streamlines coming near the bottom
2129  #raise ValueError('z location ', str(z0), ' is outside the model grid - too big')
2130 
2131  #print 'x index = ' + str(x_index)
2132  #print 'y index = ' + str(y_index)
2133 
2134  field_interp = actual_trilinear_interp(field,x0,y0,z0,x_index,y_index,z_index,x,y,z)
2135 
2136  return field_interp
2137 
2138 
2139 
2140 def trilinear_interp_arrays(x0,y0,z0,field,x,y,z,len_x,len_y,len_z):
2141  """!Do trilinear interpolation of the field in three spatial dimensions. This function assumes that the grid can locally be regarded as cartesian, with everything at right angles. It also requires that the entire field can be held in memory at the same time.
2142 
2143  x0,y0, and z0 represent the point to interpolate to.
2144  """
2145 
2146  # Compute indices at location given
2147  x0[x0<x[0]] = (x[0]+x[1])/2
2148  x0[x0>x[-1]] = (x[-1]+x[-2])/2
2149  x_index = x.searchsorted(x0)
2150  #x_index[x_index == 0] = 1 # a dirty hack to deal with streamlines coming near the edge
2151  #raise ValueError('x location ', str(x0), ' is outside the model grid - too small')
2152  #x_index[x_index == len_x] = len_x - 1 # a dirty hack to deal with streamlines coming near the edge
2153  #raise ValueError('x location ', str(x0), ' is outside the model grid - too big')
2154 
2155  y0[y0<y[0]] = (y[0]+y[1])/2
2156  y0[y0>y[-1]] = (y[-1]+y[-2])/2
2157  y_index = y.searchsorted(y0)
2158  #y_index[y_index == 0] = 1 # a dirty hack to deal with streamlines coming near the edge
2159  #raise ValueError('y location ', str(y0), ' is outside the model grid - too small')
2160  #y_index[y_index == len_y] = len_y - 1 # a dirty hack to deal with streamlines coming near the edge
2161  #raise ValueError('y location ', str(y0), ' is outside the model grid - too big')
2162 
2163  # np.searchsorted only works for positive arrays :/
2164  if any(z < 0):
2165  z0 = -z0
2166  z = -z
2167  z0[z0<z[0]] = (z[0]+z[1])/2
2168  z0[z0>z[-1]] = (z[-1]+z[-2])/2
2169  z_index = z.searchsorted(z0)
2170  #z_index[z_index == 0] = 1# a dirty hack to deal with streamlines coming near the surface
2171  #raise ValueError('z location ', str(z0), ' is outside the model grid - too small')
2172  #z_index[z_index == len_z] = len_z - 1 # a dirty hack to deal with streamlines coming near the bottom
2173  #raise ValueError('z location ', str(z0), ' is outside the model grid - too big')
2174 
2175  #print 'x index = ' + str(x_index)
2176  #print 'y index = ' + str(y_index)
2177 
2178  field_interp = actual_trilinear_interp(field,x0,y0,z0,x_index,y_index,z_index,x,y,z)
2179 
2180  return field_interp
2181 
2182 @numba.jit
2183 def actual_trilinear_interp(field,x0,y0,z0,x_index,y_index,z_index,x,y,z):
2184  """!This is a numba accelerated trilinear interpolation. The @numba.jit decorator just above this function causes it to be compiled just before it is run. This introduces a small, Order(1 second), overhead the first time, but not on subsequent calls.
2185  """
2186  field_interp = ((field[z_index-1,y_index-1,x_index-1]*
2187  (x[x_index] - x0)*(y[y_index] - y0)*(z[z_index] - z0) +
2188  field[z_index,y_index-1,x_index-1]*
2189  (x[x_index] - x0)*(y[y_index] - y0)*(z0 - z[z_index-1]) +
2190  field[z_index,y_index,x_index-1]*
2191  (x[x_index] - x0)*(y0 - y[y_index-1])*(z0 - z[z_index-1]) +
2192  field[z_index-1,y_index,x_index-1]*
2193  (x[x_index] - x0)*(y0 - y[y_index-1])*(z[z_index] - z0) +
2194  field[z_index-1,y_index-1,x_index]*
2195  (x0 - x[x_index-1])*(y[y_index] - y0)*(z[z_index] - z0) +
2196  field[z_index,y_index-1,x_index]*
2197  (x0 - x[x_index-1])*(y[y_index] - y0)*(z0 - z[z_index-1]) +
2198  field[z_index,y_index,x_index]*
2199  (x0 - x[x_index-1])*(y0 - y[y_index-1])*(z0 - z[z_index-1]) +
2200  field[z_index-1,y_index,x_index]*
2201  (x0 - x[x_index-1])*(y0 - y[y_index-1])*(z[z_index] - z0))/
2202  ((z[z_index] - z[z_index-1])*(y[y_index] - y[y_index-1])*(x[x_index] - x[x_index-1])))
2203  return field_interp
2204 
2205 def quadralinear_interp(x0,y0,z0,t0,
2206  field,
2207  x,y,z,t,
2208  len_x,len_y,len_z,len_t,
2209  x_index,y_index,z_index,t_index):
2210  """! Do quadralinear interpolation of the velocity field in three spatial dimensions and one temporal dimension to get nice accurate streaklines. This function assumes that the grid can locally be regarded as cartesian, with everything at right angles.
2211 
2212  x0,y0,z0, and t0 represent the point to interpolate to.
2213 
2214  The velocity field, "field", is passed as a truncated 4D field.
2215 
2216  x,y,z,t are vectors of these dimensions.
2217  """
2218 
2219  field_interp = actual_quadralinear_interp(field[1:3,1:3,1:3],
2220  x0,y0,z0,t0,
2221  x_index,y_index,z_index,t_index,
2222  x,y,z,t)
2223 
2224  return field_interp
2225 
2226 @numba.jit
2227 def actual_quadralinear_interp(field,x0,y0,z0,t0,
2228  x_index,y_index,z_index,t_index,
2229  x,y,z,t):
2230  """!This is a numba accelerated quadralinear interpolation. The @numba.jit decorator just above this function causes it to be compiled just before it is run. This introduces a small, Order(1 second), overhead the first time, but not on subsequent
2231  calls.
2232  """
2233  field_interp = ((field[0,0,0,0]*
2234  (x[x_index] - x0)*(y[y_index] - y0)*(z[z_index] - z0)*(t[t_index] - t0) +
2235  field[0,1,0,0]*
2236  (x[x_index] - x0)*(y[y_index] - y0)*(z0 - z[z_index-1])*(t[t_index] - t0) +
2237  field[0,1,1,0]*
2238  (x[x_index] - x0)*(y0 - y[y_index-1])*(z0 - z[z_index-1])*(t[t_index] - t0) +
2239  field[0,0,1,0]*
2240  (x[x_index] - x0)*(y0 - y[y_index-1])*(z[z_index] - z0)*(t[t_index] - t0) +
2241  field[0,0,0,1]*
2242  (x0 - x[x_index-1])*(y[y_index] - y0)*(z[z_index] - z0)*(t[t_index] - t0) +
2243  field[0,1,0,1]*
2244  (x0 - x[x_index-1])*(y[y_index] - y0)*(z0 - z[z_index-1])*(t[t_index] - t0) +
2245  field[0,1,1,1]*
2246  (x0 - x[x_index-1])*(y0 - y[y_index-1])*(z0 - z[z_index-1])*(t[t_index] - t0) +
2247  field[0,0,1,1]*
2248  (x0 - x[x_index-1])*(y0 - y[y_index-1])*(z[z_index] - z0)*(t[t_index] - t0) +
2249 
2250  field[1,0,0,0]*
2251  (x[x_index] - x0)*(y[y_index] - y0)*(z[z_index] - z0)*(t0 - t[t_index-1]) +
2252  field[1,1,0,0]*
2253  (x[x_index] - x0)*(y[y_index] - y0)*(z0 - z[z_index-1])*(t0 - t[t_index-1]) +
2254  field[1,1,1,0]*
2255  (x[x_index] - x0)*(y0 - y[y_index-1])*(z0 - z[z_index-1])*(t0 - t[t_index-1]) +
2256  field[1,0,1,0]*
2257  (x[x_index] - x0)*(y0 - y[y_index-1])*(z[z_index] - z0)*(t0 - t[t_index-1]) +
2258  field[1,0,0,1]*
2259  (x0 - x[x_index-1])*(y[y_index] - y0)*(z[z_index] - z0)*(t0 - t[t_index-1]) +
2260  field[1,1,0,1]*
2261  (x0 - x[x_index-1])*(y[y_index] - y0)*(z0 - z[z_index-1])*(t0 - t[t_index-1]) +
2262  field[1,1,1,1]*
2263  (x0 - x[x_index-1])*(y0 - y[y_index-1])*(z0 - z[z_index-1])*(t0 - t[t_index-1]) +
2264  field[1,0,1,1]*
2265  (x0 - x[x_index-1])*(y0 - y[y_index-1])*(z[z_index] - z0)*(t0 - t[t_index-1]))/
2266  ((t[t_index] - t[t_index-1])*(z[z_index] - z[z_index-1])*
2267  (y[y_index] - y[y_index-1])*(x[x_index] - x[x_index-1])))
2268 
2269  return field_interp
2270 
2271 
2272 
2273 
2274 
2275 def indices_and_field(x,y,z,
2276  x0,y0,z0,t_index,
2277  len_x,len_y,len_z,len_t,
2278  netcdf_filehandle,variable, bias_field):
2279  """!A helper function to extract a small 4D hypercube of data from a netCDF file. This isn't intended to be used on its own."""
2280 
2281  # Compute indices at location given
2282  x_index = np.searchsorted(x,x0)
2283  y_index = np.searchsorted(y,y0)
2284  # np.searchsorted only works for positive arrays :/
2285  if z0 < 0:
2286  z0 = -z0
2287  z = -z
2288  z_index = np.searchsorted(z,z0)
2289 
2290  # a bunch of dirty hacks to deal with streamlines coming near the edge
2291  x_index = max(x_index,1)
2292  x_index = min(x_index,len_x-3)
2293  y_index = max(y_index,1)
2294  y_index = min(y_index,len_y-3)
2295  z_index = max(z_index,1)
2296  z_index = min(z_index,len_z-3)
2297 
2298  field = netcdf_filehandle.variables[variable][t_index-1:t_index+2,
2299  z_index-1:z_index+3,
2300  y_index-1:y_index+3,
2301  x_index-1:x_index+3]
2302 
2303 
2304  field = (field +
2305  bias_field[np.newaxis,z_index-1:z_index+3,y_index-1:y_index+3,x_index-1:x_index+3])
2306 
2307  return field,x_index,y_index,z_index
2308 
2309 
2310 def extract_along_path4D(path_x,path_y,path_z,path_t,
2311  x_axis,y_axis,z_axis,t_axis,
2312  netcdf_filename='netcdf file with variable of interest',
2313  netcdf_variable='momVort3'):
2314  """!extract the value of a field along a path through a time varying, 3 dimensional field. The field must be in a NetCDF file, since it is assumed to be 4D and very large.
2315 
2316  This can also be used to pull out values at specific locations and times.
2317  """
2318 
2319  t_index = np.searchsorted(t_axis,path_t[0])
2320  t_index_new = np.searchsorted(t_axis,path_t[0]) # this is later used to test if new data needs to be read in.
2321 
2322  len_x = len(x_axis)
2323  len_y = len(y_axis)
2324  len_z = len(z_axis)
2325  len_t = len(t_axis)
2326 
2327 
2328  netcdf_filehandle = netCDF4.Dataset(netcdf_filename)
2329  field,x_index,y_index,z_index = indices_and_field(x_axis,y_axis,z_axis,
2330  path_x[0],path_y[0],path_z[0],t_index,
2331  len_x,len_y,len_z,len_t,
2332  netcdf_filehandle,netcdf_variable)
2333 
2334  field,x_index_new,y_index_new,z_index_new = indices_and_field(x_axis,y_axis,z_axis,
2335  path_x[0],path_y[0],path_z[0],t_index,
2336  len_x,len_y,len_z,len_t,
2337  netcdf_filehandle,netcdf_variable)
2338 
2339 
2340 
2341  path_variable = np.zeros((path_x.shape))
2342 
2343  i=0
2344 
2345 
2346  if t_index == 0:
2347  raise ValueError('Given start time is outside the given variable field - too small')
2348  elif t_index == len_t:
2349  raise ValueError('Given start time is outside the given variable field - too big')
2350 
2351 
2352  for i in xrange(0,len(path_t)):
2353 
2354  if (y_index_new==y_index and
2355  x_index_new==x_index and
2356  z_index_new==z_index and
2357 
2358  t_index_new == t_index):
2359  # the particle hasn't moved out of the grid cell it was in.
2360  # So the loaded field is fine; there's no need to reload it.
2361  pass
2362  else:
2363 
2364  t_index = np.searchsorted(t,path_t[i])
2365  if t_index == 0:
2366  raise ValueError('Time value is outside the given variable field - too small')
2367  elif t_index == len_t:
2368  raise ValueError('Time value is outside the given variable field - too big')
2369 
2370  field,x_index,y_index,z_index = indices_and_field(x_axis,y_axis,z_axis,
2371  path_x[i],path_y[i],path_z[i],t_index,
2372  len_x,len_y,len_z,len_t,
2373  netcdf_filehandle,netcdf_variable)
2374 
2375 
2376  # Interpolate field to location
2377  field_at_loc = quadralinear_interp(path_x[i],path_y[i],path_z[i],path_t[i],
2378  field,
2379  x_axis,y_axis,z_axis,t_axis,
2380  len_x,len_y,len_z,len_t,
2381  x_index,y_index,z_index,t_index)
2382 
2383 
2384  path_variable[i] = field_at_loc
2385 
2386  t_index_new = np.searchsorted(t_axis,path_t[i])
2387  x_index_new = np.searchsorted(x_axis,path_x[i])
2388  y_index_new = np.searchsorted(y_axis,path_y[i])
2389  if path_z[i] < 0:
2390  z_index_new = np.searchsorted(-z_axis,-path_z[i])
2391  else:
2392  z_index_new = np.searchsorted(z_axis,path_z[i])
2393 
2394 
2395 
2396  netcdf_filehandle.close()
2397 
2398  return path_variable
2399 
2400 
2401 
2402 def extract_along_path3D(path_x,path_y,path_z,
2403  x_axis,y_axis,z_axis,field):
2404  """!Extract the value of a field along a path through a 3 dimensional field. The field must be passed as an array. Time varying fields are not supported.
2405  """
2406 
2407  len_x = len(x_axis)
2408  len_y = len(y_axis)
2409  len_z = len(z_axis)
2410 
2411  path_variable = np.zeros((path_x.shape))
2412 
2413  #if np.min(path_z) == np.max(path_z):
2414  # the path is along a 2D depth surface
2415  # surf_loc_array = path_z[0]*np.ones((field.shape))
2416  # field_at_depth = functions.extract_on_surface(field,surf_loc_array,z_axis,direction='up')
2417 
2418 
2419  # FINISH THIS
2420 
2421  for i in xrange(0,len(path_x)):
2422 
2423  # Interpolate field to location
2424  path_variable[i] = trilinear_interp(path_x[i],path_y[i],path_z[i],field,x_axis,y_axis,z_axis,len_x,len_y,len_z)
2425 
2426  return path_variable
2427 
2428 
2429 
2430 def extract_along_path2D(path_x,path_y,
2431  x_axis,y_axis,field,order=3,dx=0,dy=0):
2432  """!Extract the value of a field along a path through a 2 dimensional field. The field must be passed as an array. Time varying fields are not supported.
2433 
2434  ------
2435  ##Parameters##
2436  * path_x - x-coordinate of the path.
2437  * path_y - y-coordinate of the path.
2438  * x_axis - vector of the x-coordinate of the input field.
2439  * y_axis - vector of the y-coordinate of the input field.
2440  * field - the input field, the values of which will be extracted along the path.
2441  * order - the order for the interpolation function, must be between 1 and 5 inclusive. 1 -> linear, 3 -> cubic.
2442  * dx, dy - order of the derivative to take in x or y. Optional arguments that are passed to the scipy cubic interpolation function.
2443  """
2444 
2445 
2446 
2447  path_variable = np.zeros((path_x.shape))
2448 
2449  if order == 1:
2450  len_x = len(x_axis)
2451  len_y = len(y_axis)
2452  # use linear interpolation function from this module
2453 
2454  for i in xrange(0,len(path_x)):
2455  # Interpolate field to location
2456  path_variable[i] = bilinear_interp(path_x[i],path_y[i],field,x_axis,y_axis,len_x,len_y)
2457 
2458  else:
2459  interp_field = scipy.interpolate.RectBivariateSpline(y_axis,x_axis,field,kx=order,ky=order)
2460 
2461  for i in xrange(0,len(path_x)):
2462 
2463  # Interpolate field to location
2464  path_variable[i] = interp_field(path_y[i],path_x[i],dx=dx,dy=dy)
2465 
2466  return path_variable
2467 
2468 
mitgcm.streamlines.quadralinear_interp
def quadralinear_interp(x0, y0, z0, t0, field, x, y, z, t, len_x, len_y, len_z, len_t, x_index, y_index, z_index, t_index)
Do quadralinear interpolation of the velocity field in three spatial dimensions and one temporal dime...
Definition: streamlines.py:2209
mitgcm.streamlines.pathlines
def pathlines
A three-dimensional lagrangian particle tracker.
Definition: streamlines.py:490
mitgcm.streamlines.pathlines_many
def pathlines_many
A three-dimensional lagrangian particle tracker designed for tracking many particles at once...
Definition: streamlines.py:841
mitgcm.streamlines.extract_along_path2D
def extract_along_path2D
Extract the value of a field along a path through a 2 dimensional field.
Definition: streamlines.py:2431
mitgcm.streamlines.numeric_GLM_xyzt
def numeric_GLM_xyzt
A three-dimensional lagrangian particle tracker designed for tracking many particles at once as a met...
Definition: streamlines.py:1474
mitgcm.streamlines.actual_trilinear_interp
def actual_trilinear_interp(field, x0, y0, z0, x_index, y_index, z_index, x, y, z)
This is a numba accelerated trilinear interpolation.
Definition: streamlines.py:2183
mitgcm.streamlines.stream3_many
def stream3_many
A three-dimensional streamline solver.
Definition: streamlines.py:311
mitgcm.streamlines.extract_along_path4D
def extract_along_path4D
extract the value of a field along a path through a time varying, 3 dimensional field.
Definition: streamlines.py:2313
mitgcm.streamlines.trilinear_interp
def trilinear_interp(x0, y0, z0, field, x, y, z, len_x, len_y, len_z)
Do trilinear interpolation of the field in three spatial dimensions.
Definition: streamlines.py:2096
mitgcm.streamlines.pathlines_for_OLIC_xyzt_ani
def pathlines_for_OLIC_xyzt_ani
A three-dimensional lagrangian particle tracker designed for tracking many particles at once...
Definition: streamlines.py:1143
mitgcm.functions.interp_field
def interp_field
Interpolate a given field onto a different grid.
Definition: functions.py:295
mitgcm.streamlines.stream3
def stream3
A three-dimensional streamline solver.
Definition: streamlines.py:144
mitgcm.streamlines.actual_quadralinear_interp
def actual_quadralinear_interp(field, x0, y0, z0, t0, x_index, y_index, z_index, t_index, x, y, z, t)
This is a numba accelerated quadralinear interpolation.
Definition: streamlines.py:2229
mitgcm.streamlines.extract_along_path3D
def extract_along_path3D(path_x, path_y, path_z, x_axis, y_axis, z_axis, field)
Extract the value of a field along a path through a 3 dimensional field.
Definition: streamlines.py:2403
mitgcm.streamlines.bilinear_interp
def bilinear_interp(x0, y0, field, x, y, len_x, len_y)
Do bilinear interpolation of a field.
Definition: streamlines.py:2056
mitgcm.streamlines.numeric_GLM_xyz
def numeric_GLM_xyz
A three-dimensional streamline solver.
Definition: streamlines.py:1821
mitgcm.streamlines.trilinear_interp_arrays
def trilinear_interp_arrays(x0, y0, z0, field, x, y, z, len_x, len_y, len_z)
Do trilinear interpolation of the field in three spatial dimensions.
Definition: streamlines.py:2140
mitgcm.streamlines.indices_and_field
def indices_and_field(x, y, z, x0, y0, z0, t_index, len_x, len_y, len_z, len_t, netcdf_filehandle, variable, bias_field)
A helper function to extract a small 4D hypercube of data from a netCDF file.
Definition: streamlines.py:2278
mitgcm.streamlines.actual_bilinear_interp
def actual_bilinear_interp(field, x0, y0, x, y, len_x, len_y, x_index, y_index)
This is a numba accelerated bilinear interpolation.
Definition: streamlines.py:2086
mitgcm.streamlines.stream2
def stream2
A two-dimensional streamline solver.
Definition: streamlines.py:29
Generated on Thu Nov 17 2016 10:22:54 for mitgcm by   doxygen 1.8.9.1