def multiple_timestep_ocean():
field1 = eastward_ocean()
field2 = equator_converging_ocean()
time = np.array([0, 100*3600]) # 0 hrs, 100 hrs
U = np.concatenate([field1.U, field2.U], axis=0)
V = np.concatenate([field1.V, field2.V], axis=0)
return Field2D(time, field1.x, field1.y, U, V)
def electric_dipole(nx=100):
"""electric field generation yoinked from https://scipython.com/blog/visualizing-a-vector-field-with-matplotlib/"""
def E(q, r0, x, y):
"""Return the electric field vector E=(Ex,Ey) due to charge q at r0."""
den = np.hypot(x - r0[0], y - r0[1]) ** 3
return q * (x - r0[0]) / den, q * (y - r0[1]) / den
# Grid of x, y points
ny = nx
x = np.linspace(-2, 2, nx)
y = np.linspace(-2, 2, ny)
X, Y = np.meshgrid(x, y)
# Create a multipole with nq charges of alternating sign, equally spaced
# on the unit circle.
nq = 2
charges = []
for i in range(nq):
q = i % 2 * 2 - 1
charges.append((q, (np.cos(2 * np.pi * i / nq), np.sin(2 * np.pi * i / nq))))
# Electric field vector, E=(Ex, Ey), as separate components
Ex, Ey = np.zeros((ny, nx)), np.zeros((ny, nx))
for charge in charges:
ex, ey = E(*charge, x=X, y=Y)
Ex += ex
Ey += ey
return Field2D(x=x, y=y, time=np.array([0]), U=Ex.T[np.newaxis], V=Ey.T[np.newaxis])
def eastward_ocean():
time = np.array([0]) # seconds
x = np.linspace(-180, 180, 1000)
y = np.linspace(-90, 90, 500)
U = 5 * np.ones([len(time), len(x), len(y)])
V = np.zeros([len(time), len(x), len(y)])
return Field2D(time, x, y, U, V)
def concentric_circles(nx=100):
x = np.linspace(-1, 1, nx)
y = np.linspace(-1, 1, nx)
X, Y = np.meshgrid(x, y)
mag = np.sqrt(X**2 + Y**2)
U = np.divide(-Y, mag, out=np.zeros_like(X), where=mag != 0)
V = np.divide(X, mag, out=np.zeros_like(X), where=mag != 0)
return Field2D(x=x, y=y, time=np.array([0]), U=U.T[np.newaxis], V=V.T[np.newaxis])
def equator_converging_ocean():
time = np.array([0]) # seconds
x = np.linspace(-180, 180, 1000)
y = np.linspace(-90, 90, 500)
[X, Y] = np.meshgrid(x, y)
U = np.zeros([len(time), len(x), len(y)])
V = np.zeros([len(time), len(x), len(y)])
V[0, :] = -np.sign(Y.T) * 5
return Field2D(time, x, y, U, V)
def converge():
time = np.array([0]) # seconds
x = np.linspace(-100, 100, 100)
y = np.linspace(-50, 50, 50)
[X, Y] = np.meshgrid(x, y)
U = np.zeros([len(time), len(x), len(y)])
V = np.zeros([len(time), len(x), len(y)])
U[0, :] = -X.T
V[0, :] = -Y.T
return Field2D(time, x, y, U * .01, V * .01)
def hycom_surface(months):
time = []
x = []
y = []
U = []
V = []
for month in months:
ds = xr.open_dataset(f'./data/uv_2015_{month}_3d.formatted.nc')
x.append(ds.x.data)
y.append(ds.y.data)
time.append(ds.time.data)
U.append(ds.water_u.sel(z=0).data.swapaxes(1, 2))
V.append(ds.water_v.sel(z=0).data.swapaxes(1, 2))
assert np.all(x[0] == x[i] for i in range(len(x)))
assert np.all(y[0] == y[i] for i in range(len(x)))
return Field2D(time=np.concatenate(time), x=x[0], y=y[0],
U=np.concatenate(U, axis=0), V=np.concatenate(V, axis=0))
def advect_particle(p0, field: Field2D, time):
"""advect a particle with initial position p0 through field along vector time"""
ntimesteps = len(time)
P = np.zeros([ntimesteps, 2]) # particle position in time, space
# initialize particle
P[0] = p0
for i in range(ntimesteps-1):
# find nearest field U and V indices
nearest_time_idx = np.argmin(np.abs(field.time - time[i]))
nearest_x_idx = np.argmin(np.abs(field.x - P[i, 0]))
nearest_y_idx = np.argmin(np.abs(field.y - P[i, 1]))
dt = time[i+1] - time[i]
dx, dy = 0, 0
if field.inbounds(P[i]):
dx = field.U[nearest_time_idx, nearest_x_idx, nearest_y_idx] * dt
dy = field.V[nearest_time_idx, nearest_x_idx, nearest_y_idx] * dt
if dy == 0:
print('huh')
P[i+1] = P[i] + np.array([dx, dy])
return P
def test_cartesian_electric_dipole():
field1 = generate_field.electric_dipole(nx=1000)
field2 = generate_field.electric_dipole(nx=1000)
field = Field2D(x=field1.x,
y=field1.y,
time=np.array([0, .5]),
U=np.concatenate([field1.U, -field2.U], axis=0),
V=np.concatenate([field1.V, -field2.V], axis=0))
# advect particles
np.random.seed(2)
p0 = np.random.rand(1000, 2) * [
field.x.max() - field.x.min(),
field.y.max() - field.y.min()
] + [field.x.min(), field.y.min()]
num_timesteps = 600
save_every = 10
dt = 1e-3
P, buffer_seconds, kernel_seconds = openCL_advect(
field=field,
p0=p0,
t0=0,
num_timesteps=num_timesteps,
save_every=save_every,
dt=dt,
device_index=2,
verbose=True)
fig, ax = plt.subplots(1, 1)
ax.set_aspect(1)
dot, = ax.plot([], [], '.', markersize=5, color='C0')
charge1 = ax.add_artist(Circle((-1, 0), 0.05, color='blue'))
charge2 = ax.add_artist(Circle((1, 0), 0.05, color='red'))
ax.streamplot(field.x,
field.y,
field.U[0].T,
field.V[0].T,
linewidth=1,
density=2,
arrowsize=0,
color='gray')
ax.set_xlim([min(field.x), max(field.x)])
ax.set_ylim([min(field.y), max(field.y)])
ax.axes.xaxis.set_visible(False)
ax.axes.yaxis.set_visible(False)
time = dt * np.arange(num_timesteps, step=save_every)
def update(i):
if i == 0:
charge1.set_color('blue')
charge2.set_color('red')
elif i == 24:
charge1.set_color('red')
charge2.set_color('blue')
dot.set_xdata(P[:, i, 0])
dot.set_ydata(P[:, i, 1])
ax.set_title(f't={i}')
ax.set_xlim([min(field.x), max(field.x)])
ax.set_ylim([min(field.y), max(field.y)])
return dot, charge1, charge2
ani = FuncAnimation(fig, update, frames=len(time))
plt.show()