def deltaC(size,center=None,magnitude=False):
"""
returns an array where every value is the distance from
the center point.
:param center: center point, if not specified, assume size/2
:param magnitude: return a single magnitude value instead of [dx,dy]
"""
size=(int(size[0]),int(size[1]))
if center is None:
center=size[0]/2.0,size[1]/2.0
if magnitude:
if False:
# old way, non-vector
img=np.ndarray((size[0],size[1]))
for x in range(size[0]):
for y in range(size[1]):
img[x,y]=math.sqrt((x-center[0])**2+(y-center[1])**2)
else:
img=np.sqrt((np.arange(size[0])-center[0])**2+(np.arange(size[1])[:,None]-center[1])**2)
else:
img=np.ndarray((size[0],size[1],2))
for x in range(size[0]):
for y in range(size[1]):
img[x,y]=(x-center[0]),(y-center[1])
return img
def monte_carlo_pi(samples, iterations):
s = np.zeros(1)
for i in xrange(0, iterations): # sample
x = np.random.random((samples), dtype=np.float64, bohrium=True)
y = np.random.random((samples), dtype=np.float64, bohrium=True)
m = np.sqrt(x*x + y*y) # model
c = np.add.reduce((m<1.0).astype('float')) # count
s += c*4.0 / samples # approximate
return s / (iterations)
def move(m, x, y, z, vx, vy, vz, dt, temporaries):
""" Move the bodies.
first find forces and change velocity and then move positions.
"""
start = time.time()
dx = x - x[:, None]
dy = numpy.subtract(y, y[:, None])
dz = numpy.subtract(z, z[:, None])
pm = numpy.multiply(m, m[:, None])
end = time.time()
print("Step 0:", end - start)
start = end
r = np.sqrt(dx**2 + dy**2 + dz**2)
tmp = G * pm / r**2
Fx = tmp * (dx / r)
Fy = tmp * (dy / r)
Fz = tmp * (dz / r)
end = time.time()
print("Step 1:", end - start)
start = end
fill_diagonal(Fx, 0.0)
fill_diagonal(Fy, 0.0)
fill_diagonal(Fz, 0.0)
end = time.time()
print("Step 2:", end - start)
start = end
mdt = m / dt
# Update state.
vx += np.add.reduce(Fx, axis=1) / mdt
vy += np.add.reduce(Fy, axis=1) / mdt
vz += np.add.reduce(Fz, axis=1) / mdt
x += vx * dt
y += vy * dt
z += vz * dt
end = time.time()
print("Step 3:", end - start)
start = end
return Fx, Fy, Fz
def norm(x, ord=None, axis=None):
"""
This version of norm is not fully compliant with the NumPy version,
it only supports computing 2-norm of a vector.
"""
if ord != None:
raise NotImplementedError("Unsupported value param ord=%s" % ord)
if axis != None:
raise NotImplementedError("Unsupported value of param ord=%s" % axis)
r = np.sum(x*x)
if issubclass(np.dtype(r.dtype).type, np.integer):
r_f32 = np.empty(r.shape, dtype=np.float32)
r_f32[:] = r
r = r_f32
return np.sqrt(r)
def norm(x, ord=None, axis=None):
"""
This version of norm is not fully compliant with the NumPy version,
it only supports computing 2-norm of a vector.
"""
if ord != None:
raise NotImplementedError("Unsupported value param ord=%s" % ord)
if axis != None:
raise NotImplementedError("Unsupported value of param ord=%s" % axis)
r = np.sum(x * x)
if issubclass(np.dtype(r.dtype).type, np.integer):
r_f32 = np.empty(r.shape, dtype=np.float32)
r_f32[:] = r
r = r_f32
return np.sqrt(r)
def logpolar2cartesian(polar_data):
"""
From:
https://stackoverflow.com/questions/2164570/reprojecting-polar-to-cartesian-grid
"""
from scipy.ndimage.interpolation import map_coordinates
theta_step=1
range_step=500
x=np.arange(-100000, 100000, 1000)
y=x
order=3
# "x" and "y" are numpy arrays with the desired cartesian coordinates
# we make a meshgrid with them
X, Y = np.meshgrid(x, y)
# Now that we have the X and Y coordinates of each point in the output plane
# we can calculate their corresponding theta and range
Tc = np.degrees(np.arctan2(Y, X)).ravel()
Rc = np.ln(np.sqrt(X**2 + Y**2)).ravel()
# Negative angles are corrected
Tc[Tc < 0] = 360 + Tc[Tc < 0]
# Using the known theta and range steps, the coordinates are mapped to
# those of the data grid
Tc = Tc / theta_step
Rc = Rc / range_step
# An array of polar coordinates is created stacking the previous arrays
#coords = np.vstack((Ac, Sc))
coords = np.vstack((Tc, Rc))
# To avoid holes in the 360º - 0º boundary, the last column of the data
# copied in the begining
polar_data = np.vstack((polar_data, polar_data[-1,:]))
# The data is mapped to the new coordinates
# Values outside range are substituted with nans
cart_data = map_coordinates(polar_data, coords, order=order, mode='constant', cval=np.nan)
# The data is reshaped and returned
return cart_data.reshape(len(Y), len(X)).T
def get(self, type='smooth', mode='reflect'):
""" type can be 'smooth' or 'fine'
mode can be 'reflect','constant','nearest','mirror', 'wrap' for handling borders """
gradfn = {'smooth': self.prewittgrad, 'fine': self.basicgrad}[type]
gradx, grady = gradfn()
# x, y and z below are now the gradient matrices,
# each entry from x,y,z is a gradient vector at an image point
x = filters.convolve(self.img, gradx, mode=mode)
y = filters.convolve(self.img, grady, mode=mode)
# norm is the magnitude of the x,y,z vectors,
# each entry is the magnitude of the gradient at an image point and z*z = 1
norm = np.sqrt(x * x + y * y + 1)
# divide by the magnitude to normalise
# as well scale to an image: negative 0-127, positive 127-255
x, y = [a / norm * 127.0 + 128.0 for a in (x, y)]
z = np.ones(
self.shape) / norm # generate z, matrix of ones, then normalise
z = z * 255.0 # all positive
# x, -y gives blender form
# convert to int, transpose to rgb and return the normal map
return np.array([x, -y, z]).transpose(1, 2, 0).astype(np.uint8)
def weighted_line(r0, c0, r1, c1, w, rmin=0, rmax=np.inf):
# The algorithm below works fine if c1 >= c0 and c1-c0 >= abs(r1-r0).
# If either of these cases are violated, do some switches.
if abs(c1-c0) < abs(r1-r0):
# Switch x and y, and switch again when returning.
xx, yy, val = weighted_line(c0, r0, c1, r1, w, rmin=rmin, rmax=rmax)
return (yy, xx, val)
# At this point we know that the distance in columns (x) is greater
# than that in rows (y). Possibly one more switch if c0 > c1.
if c0 > c1:
return weighted_line(r1, c1, r0, c0, w, rmin=rmin, rmax=rmax)
# The following is now always < 1 in abs
num=r1-r0
denom=c1-c0
slope = np.divide(num,denom,out=np.zeros_like(denom), where=denom!=0)
# Adjust weight by the slope
w *= np.sqrt(1+np.abs(slope)) / 2
# We write y as a function of x, because the slope is always <= 1
# (in absolute value)
x = np.arange(c0, c1+1, dtype=float)
y = x * slope + (c1*r0-c0*r1) / (c1-c0)
# Now instead of 2 values for y, we have 2*np.ceil(w/2).
# All values are 1 except the upmost and bottommost.
thickness = np.ceil(w/2)
yy = (np.floor(y).reshape(-1,1) + np.arange(-thickness-1,thickness+2).reshape(1,-1))
xx = np.repeat(x, yy.shape[1])
vals = trapez(yy, y.reshape(-1,1), w).flatten()
yy = yy.flatten()
# Exclude useless parts and those outside of the interval
# to avoid parts outside of the picture
mask = np.logical_and.reduce((yy >= rmin, yy < rmax, vals > 0))
return (yy[mask].astype(int), xx[mask].astype(int), vals[mask])
def move(galaxy, dt):
"""Move the bodies
first find forces and change velocity and then move positions
"""
n = len(galaxy['x'])
# Calculate all dictances component wise (with sign)
dx = galaxy['x'][np.newaxis, :].T - galaxy['x']
dy = galaxy['y'][np.newaxis, :].T - galaxy['y']
dz = galaxy['z'][np.newaxis, :].T - galaxy['z']
# Euclidian distances (all bodys)
r = np.sqrt(dx**2 + dy**2 + dz**2)
diagonal(r)[:] = 1.0
# prevent collition
mask = r < 1.0
r = r * ~mask + 1.0 * mask
m = galaxy['m'][np.newaxis, :].T
# Calculate the acceleration component wise
Fx = G * m * dx / r**3
Fy = G * m * dy / r**3
Fz = G * m * dz / r**3
# Set the force (acceleration) a body exerts on it self to zero
diagonal(Fx)[:] = 0.0
diagonal(Fy)[:] = 0.0
diagonal(Fz)[:] = 0.0
galaxy['vx'] += dt * np.sum(Fx, axis=0)
galaxy['vy'] += dt * np.sum(Fy, axis=0)
galaxy['vz'] += dt * np.sum(Fz, axis=0)
galaxy['x'] += dt * galaxy['vx']
galaxy['y'] += dt * galaxy['vy']
galaxy['z'] += dt * galaxy['vz']
plt.xticks([])
plt.yticks([])
k = 5 # number of plane waves
stripes = 37 # number of stripes per wave
N = 512 # image size in pixels
ite = 30 # iterations
phases = np.arange(0, 2 * pi, 2 * pi / ite)
image = np.empty((N, N))
d = np.arange(-N / 2, N / 2, dtype=np.float64)
xv, yv = np.meshgrid(d, d)
theta = np.arctan2(yv, xv)
r = np.log(np.sqrt(xv * xv + yv * yv))
r[np.isinf(r) == True] = 0
tcos = theta * np.cos(np.arange(0, pi, pi / k))[:, np.newaxis, np.newaxis]
rsin = r * np.sin(np.arange(0, pi, pi / k))[:, np.newaxis, np.newaxis]
inner = (tcos - rsin) * stripes
cinner = np.cos(inner)
sinner = np.sin(inner)
i = 0
for phase in phases:
image[:] = np.sum(cinner * np.cos(phase) - sinner * np.sin(phase),
axis=0) + k
plt.imshow(image.copy2numpy(), cmap="RdGy")
fig.savefig("quasi-{:03d}.png".format(i),
def integrate_tke(u, v, w, maskU, maskV, maskW, dxt, dxu, dyt, dyu, dzt, dzw, cost, cosu, kbot, kappaM, mxl, forc, forc_tke_surface, tke, dtke):
tau = 0
taup1 = 1
taum1 = 2
dt_tracer = 1
dt_mom = 1
AB_eps = 0.1
alpha_tke = 1.
c_eps = 0.7
K_h_tke = 2000.
flux_east = bh.zeros_like(maskU)
flux_north = bh.zeros_like(maskU)
flux_top = bh.zeros_like(maskU)
sqrttke = bh.sqrt(bh.maximum(0., tke[:, :, :, tau]))
"""
integrate Tke equation on W grid with surface flux boundary condition
"""
dt_tke = dt_mom # use momentum time step to prevent spurious oscillations
"""
vertical mixing and dissipation of TKE
"""
ks = kbot[2:-2, 2:-2] - 1
a_tri = bh.zeros_like(maskU[2:-2, 2:-2])
b_tri = bh.zeros_like(maskU[2:-2, 2:-2])
c_tri = bh.zeros_like(maskU[2:-2, 2:-2])
d_tri = bh.zeros_like(maskU[2:-2, 2:-2])
delta = bh.zeros_like(maskU[2:-2, 2:-2])
delta[:, :, :-1] = dt_tke / dzt[bh.newaxis, bh.newaxis, 1:] * alpha_tke * 0.5 \
* (kappaM[2:-2, 2:-2, :-1] + kappaM[2:-2, 2:-2, 1:])
a_tri[:, :, 1:-1] = -delta[:, :, :-2] / \
dzw[bh.newaxis, bh.newaxis, 1:-1]
a_tri[:, :, -1] = -delta[:, :, -2] / (0.5 * dzw[-1])
b_tri[:, :, 1:-1] = 1 + (delta[:, :, 1:-1] + delta[:, :, :-2]) / dzw[bh.newaxis, bh.newaxis, 1:-1] \
+ dt_tke * c_eps \
* sqrttke[2:-2, 2:-2, 1:-1] / mxl[2:-2, 2:-2, 1:-1]
b_tri[:, :, -1] = 1 + delta[:, :, -2] / (0.5 * dzw[-1]) \
+ dt_tke * c_eps / mxl[2:-2, 2:-
2, -1] * sqrttke[2:-2, 2:-2, -1]
b_tri_edge = 1 + delta / dzw[bh.newaxis, bh.newaxis, :] \
+ dt_tke * c_eps / mxl[2:-2, 2:-2, :] * sqrttke[2:-2, 2:-2, :]
c_tri[:, :, :-1] = -delta[:, :, :-1] / dzw[bh.newaxis, bh.newaxis, :-1]
d_tri[...] = tke[2:-2, 2:-2, :, tau] + dt_tke * forc[2:-2, 2:-2, :]
d_tri[:, :, -1] += dt_tke * forc_tke_surface[2:-2, 2:-2] / (0.5 * dzw[-1])
sol, water_mask = solve_implicit(ks, a_tri, b_tri, c_tri, d_tri, b_edge=b_tri_edge)
tke[2:-2, 2:-2, :, taup1] = where(water_mask, sol, tke[2:-2, 2:-2, :, taup1])
"""
Add TKE if surface density flux drains TKE in uppermost box
"""
tke_surf_corr = bh.zeros(maskU.shape[:2])
mask = tke[2:-2, 2:-2, -1, taup1] < 0.0
tke_surf_corr[2:-2, 2:-2] = where(
mask,
-tke[2:-2, 2:-2, -1, taup1] * 0.5 * dzw[-1] / dt_tke,
0.
)
tke[2:-2, 2:-2, -1, taup1] = bh.maximum(0., tke[2:-2, 2:-2, -1, taup1])
"""
add tendency due to lateral diffusion
"""
flux_east[:-1, :, :] = K_h_tke * (tke[1:, :, :, tau] - tke[:-1, :, :, tau]) \
/ (cost[bh.newaxis, :, bh.newaxis] * dxu[:-1, bh.newaxis, bh.newaxis]) * maskU[:-1, :, :]
flux_east[-1, :, :] = 0.
flux_north[:, :-1, :] = K_h_tke * (tke[:, 1:, :, tau] - tke[:, :-1, :, tau]) \
/ dyu[bh.newaxis, :-1, bh.newaxis] * maskV[:, :-1, :] * cosu[bh.newaxis, :-1, bh.newaxis]
flux_north[:, -1, :] = 0.
tke[2:-2, 2:-2, :, taup1] += dt_tke * maskW[2:-2, 2:-2, :] * \
((flux_east[2:-2, 2:-2, :] - flux_east[1:-3, 2:-2, :])
/ (cost[bh.newaxis, 2:-2, bh.newaxis] * dxt[2:-2, bh.newaxis, bh.newaxis])
+ (flux_north[2:-2, 2:-2, :] - flux_north[2:-2, 1:-3, :])
/ (cost[bh.newaxis, 2:-2, bh.newaxis] * dyt[bh.newaxis, 2:-2, bh.newaxis]))
"""
add tendency due to advection
"""
adv_flux_superbee_wgrid(
flux_east, flux_north, flux_top, tke[:, :, :, tau],
u[..., tau], v[..., tau], w[..., tau], maskW, dxt, dyt, dzw,
cost, cosu, dt_tracer
)
dtke[2:-2, 2:-2, :, tau] = maskW[2:-2, 2:-2, :] * (-(flux_east[2:-2, 2:-2, :] - flux_east[1:-3, 2:-2, :])
/ (cost[bh.newaxis, 2:-2, bh.newaxis] * dxt[2:-2, bh.newaxis, bh.newaxis])
- (flux_north[2:-2, 2:-2, :] - flux_north[2:-2, 1:-3, :])
/ (cost[bh.newaxis, 2:-2, bh.newaxis] * dyt[bh.newaxis, 2:-2, bh.newaxis]))
dtke[:, :, 0, tau] += -flux_top[:, :, 0] / dzw[0]
dtke[:, :, 1:-1, tau] += - \
(flux_top[:, :, 1:-1] - flux_top[:, :, :-2]) / dzw[1:-1]
dtke[:, :, -1, tau] += - \
(flux_top[:, :, -1] - flux_top[:, :, -2]) / \
(0.5 * dzw[-1])
"""
Adam Bashforth time stepping
"""
tke[:, :, :, taup1] += dt_tracer * ((1.5 + AB_eps) * dtke[:, :, :, tau] - (0.5 + AB_eps) * dtke[:, :, :, taum1])
return tke, dtke, tke_surf_corr
def gsw_dHdT(sa, ct, p):
"""
d/dT of dynamic enthalpy, analytical derivative
sa : Absolute Salinity [g/kg]
ct : Conservative Temperature [deg C]
p : sea pressure [dbar]
"""
t1 = v45 * ct
t2 = 0.2e1 * t1
t3 = v46 * sa
t4 = 0.5 * v12
t5 = v14 * ct
t7 = ct * (v13 + t5)
t8 = 0.5 * t7
t11 = sa * (v15 + v16 * ct)
t12 = 0.5 * t11
t13 = t4 + t8 + t12
t15 = v19 * ct
t19 = v17 + ct * (v18 + t15) + v20 * sa
t20 = 1.0 / t19
t24 = v47 + v48 * ct
t25 = 0.5 * v13
t26 = 1.0 * t5
t27 = sa * v16
t28 = 0.5 * t27
t29 = t25 + t26 + t28
t33 = t24 * t13
t34 = t19**2
t35 = 1.0 / t34
t37 = v18 + 2.0 * t15
t38 = t35 * t37
t48 = ct * (v44 + t1 + t3)
t57 = v40 * ct
t59 = ct * (v39 + t57)
t64 = t13**2
t68 = t20 * t29
t71 = t24 * t64
t74 = v04 * ct
t76 = ct * (v03 + t74)
t79 = v07 * ct
t82 = bh.sqrt(sa)
t83 = v11 * ct
t85 = ct * (v10 + t83)
t92 = v01 + ct * (v02 + t76) + sa * (v05 + ct * (v06 + t79) + t82 *
(v08 + ct * (v09 + t85)))
t93 = v48 * t92
t105 = v02 + t76 + ct * (v03 + 2.0 * t74) + sa * (v06 + 2.0 * t79 + t82 *
(v09 + t85 + ct *
(v10 + 2.0 * t83)))
t106 = t24 * t105
t107 = v44 + t2 + t3
t110 = v43 + t48
t117 = t24 * t92
t120 = 4.0 * t71 * t20 - t117 - 2.0 * t110 * t13
t123 = v38 + t59 + ct * (v39 + 2.0 * t57) + sa * v42 + (
4.0 * v48 * t64 * t20 + 8.0 * t33 * t68 - 4.0 * t71 * t38 - t93 -
t106 - 2.0 * t107 * t13 - 2.0 * t110 * t29) * t20 - t120 * t35 * t37
t128 = t19 * p
t130 = p * (1.0 * v12 + 1.0 * t7 + 1.0 * t11 + t128)
t131 = 1.0 / t92
t133 = 1.0 + t130 * t131
t134 = bh.log(t133)
t143 = v37 + ct * (v38 + t59) + sa * (v41 + v42 * ct) + t120 * t20
t152 = t37 * p
t156 = t92**2
t165 = v25 * ct
t167 = ct * (v24 + t165)
t169 = ct * (v23 + t167)
t175 = v30 * ct
t177 = ct * (v29 + t175)
t179 = ct * (v28 + t177)
t185 = v35 * ct
t187 = ct * (v34 + t185)
t189 = ct * (v33 + t187)
t199 = t13 * t20
t217 = 2.0 * t117 * t199 - t110 * t92
t234 = v21 + ct * (v22 + t169) + sa * (v26 + ct *
(v27 + t179) + v36 * sa + t82 *
(v31 + ct *
(v32 + t189))) + t217 * t20
t241 = t64 - t92 * t19
t242 = bh.sqrt(t241)
t243 = 1.0 / t242
t244 = t4 + t8 + t12 - t242
t245 = 1.0 / t244
t247 = t4 + t8 + t12 + t242 + t128
t248 = 1.0 / t247
t249 = t242 * t245 * t248
t252 = 1.0 + 2.0 * t128 * t249
t253 = bh.log(t252)
t254 = t243 * t253
t259 = t234 * t19 - t143 * t13
t264 = t259 * t20
t272 = 2.0 * t13 * t29 - t105 * t19 - t92 * t37
t282 = t128 * t242
t283 = t244**2
t287 = t243 * t272 / 2.0
t292 = t247**2
t305 = 0.1e5 * p * (v44 + t2 + t3 - 2.0 * v48 * t13 * t20
- 2.0 * t24 * t29 * t20 + 2.0 * t33 * t38 + 0.5 * v48 * p) * t20 \
- 0.1e5 * p * (v43 + t48 - 2.0 * t33 * t20 + 0.5 * t24 * p) * t38 \
+ 0.5e4 * t123 * t20 * t134 - 0.5e4 * t143 * t35 * t134 * t37 \
+ 0.5e4 * t143 * t20 * (p * (1.0 * v13 + 2.0 * t5 + 1.0 * t27 + t152) * t131
- t130 / t156 * t105) / t133 \
+ 0.5e4 * ((v22 + t169 + ct * (v23 + t167 + ct * (v24 + 2.0 * t165))
+ sa * (v27 + t179 + ct * (v28 + t177 + ct * (v29 + 2.0 * t175)) + t82 * (v32 + t189
+ ct * (v33 + t187 + ct * (v34 + 2.0 * t185)))) + (2.0 * t93 * t199 + 2.0 * t106 *
t199 + 2.0 * t117 * t68 - 2.0 * t117 * t13 * t35 * t37 - t107
* t92 - t110 * t105) * t20 - t217 * t35 * t37) * t19 + t234 * t37
- t123 * t13 - t143 * t29) * t20 * t254 - 0.5e4 * t259 * \
t35 * t254 * t37 - 0.25e4 * t264 / t242 / t241 * t253 * t272 \
+ 0.5e4 * t264 * t243 * (2.0 * t152 * t249 + t128 *
t243 * t245 * t248 * t272 - 2.0 * t282 /
t283 * t248 * (t25 + t26 + t28 - t287)
- 2.0 * t282 * t245 / t292 * (t25 + t26 + t28 + t287 + t152)) / t252
return t305
