152 lines
5.0 KiB
Python
152 lines
5.0 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
|
|
# (setq python-shell-interpreter (shell-command-to-string "nix-shell --command \"which python3 | tr -d '\n'\""))
|
|
|
|
import functools
|
|
import math
|
|
|
|
import numpy as np
|
|
import vispy
|
|
import vispy.scene
|
|
from vispy.scene import visuals
|
|
|
|
import opensimplex
|
|
|
|
def gradient(fn, eps=1e-3):
|
|
eps_inv = 1.0 / eps
|
|
# Numerical gradient by finite differences
|
|
def _grad(x, y, z):
|
|
p = fn(x, y, z)
|
|
p_dx = fn(x+eps, y, z)
|
|
p_dy = fn(x, y+eps, z)
|
|
p_dz = fn(x, y, z+eps)
|
|
return (p_dx - p)*eps_inv, (p_dy - p)*eps_inv, (p_dz - p)*eps_inv
|
|
return _grad
|
|
|
|
def curl_3d(grads):
|
|
# 'grads' should be of shape (..., 3, 3). Each 3x3 matrix should
|
|
# be the Jacobian of the potential.
|
|
cx = grads[..., 2, 1] - grads[..., 1, 2]
|
|
cy = grads[..., 0, 2] - grads[..., 2, 0]
|
|
cz = grads[..., 1, 0] - grads[..., 0, 1]
|
|
return np.stack((cx, cy, cz), axis=-1)
|
|
|
|
class Potential(object):
|
|
def __init__(self):
|
|
self.x_spx = opensimplex.OpenSimplex(seed=0)
|
|
self.y_spx = opensimplex.OpenSimplex(seed=12345)
|
|
self.z_spx = opensimplex.OpenSimplex(seed=45678)
|
|
def eval(self, x, y, z):
|
|
y2 = y# + 0.5*math.sin(3*x) + 0.5*math.sin(2.5*z)
|
|
x2 = x
|
|
z2 = z
|
|
f1 = np.array([
|
|
self.x_spx.noise3d(x2, y2, z2),
|
|
self.y_spx.noise3d(x2, y2, z2),
|
|
self.z_spx.noise3d(x2, y2, z2),
|
|
])
|
|
f2 = np.array([z*0.5, 0, 0])
|
|
return f1# + f2
|
|
def grad(self, x, y, z, eps=1e-3):
|
|
# Returns gradients in matrix form.
|
|
#
|
|
# Element (i, j) is gradient of i'th component with respect to
|
|
# j'th component, where components are (x,y,z) - i.e. row 0 is
|
|
# (d/dx Px, d/dy Px, d/dz Px). I'm fairly sure this is just
|
|
# the Jacobian.
|
|
p = self.eval(x, y, z)
|
|
p_dx = self.eval(x+eps, y, z)
|
|
p_dy = self.eval(x, y+eps, z)
|
|
p_dz = self.eval(x, y, z+eps)
|
|
return (np.stack((p_dx, p_dy, p_dz)) - p).T / eps
|
|
|
|
def generate(grid):
|
|
p = Potential()
|
|
grads = np.array([p.grad(*pt) for pt in grid])
|
|
curl = curl_3d(grads)
|
|
return curl
|
|
|
|
class Data(object):
|
|
def __init__(self, view):
|
|
self.use_tubes = False
|
|
self.view = view
|
|
s = 0.1
|
|
self.s = s
|
|
count = 8
|
|
xs = zs = np.linspace(-s, s, count)
|
|
ys = np.array([0])
|
|
self.points = np.array([i.flatten() for i in np.meshgrid(xs,ys,zs)]).T
|
|
self.points_old = None
|
|
if self.use_tubes:
|
|
self.visual = vispy.scene.visuals.Line(
|
|
color=(1,1,1,0.75),
|
|
connect='segments',
|
|
)
|
|
self.view.add(self.visual)
|
|
self.view.camera = 'turntable'
|
|
self.update()
|
|
else:
|
|
p = self.points
|
|
points = [p]
|
|
for i in range(200):
|
|
print(i)
|
|
curl = generate(p)
|
|
p2 = p + curl*0.1*s
|
|
p = p2
|
|
points.append(p)
|
|
points = np.stack(points, axis=1) / s
|
|
# points = (count*count, N, 3) where first dimension chooses which
|
|
# trajectory, and second dimension proceeds along time/iterations
|
|
# of that trajectory.
|
|
for traj in points:
|
|
tube = vispy.scene.visuals.Tube(points=traj, radius=0.4*s)
|
|
self.view.add(tube)
|
|
self.view.camera = 'turntable'
|
|
def update(self, ev=None):
|
|
if not self.use_tubes:
|
|
return
|
|
t = 0 if ev is None else ev.elapsed
|
|
# Get velocity for current points:
|
|
curl = generate(self.points)
|
|
#
|
|
a1 = self.points
|
|
a2 = self.points + curl*0.1*self.s
|
|
lines = np.hstack((a1, a2)).reshape(self.points.shape[0]*2, -1) / self.s
|
|
self.points_old = np.vstack((self.points_old, lines))
|
|
maxpoints = self.points.shape[0] * 200
|
|
extra = self.points_old.shape[0] - maxpoints
|
|
if extra > 0:
|
|
self.points_old = self.points_old[extra:]
|
|
self.visual.set_data(
|
|
pos=self.points_old,
|
|
#arrows=np.hstack((a1, a2)),
|
|
)
|
|
self.points = a2
|
|
# self.scatter = visuals.Markers()
|
|
# self.scatter.set_data(self.d, edge_color=None, face_color=(1, 0.5, 1, .5), size=4)
|
|
# view.add(self.scatter)
|
|
#m = np.array([[np.cos(t), np.sin(t*1.01), np.cos(t*1.02)]])
|
|
#d2 = self.d*m
|
|
#self.scatter.set_data(d2, edge_color=None, face_color=(1, 0.5, 1, .5), size=4)
|
|
|
|
def main():
|
|
#
|
|
# Make a canvas and add simple view
|
|
#
|
|
canvas = vispy.scene.SceneCanvas(keys='interactive', show=True)
|
|
view = canvas.central_widget.add_view()
|
|
import sys
|
|
# add a colored 3D axis for orientation
|
|
axis = visuals.XYZAxis(parent=view.scene)
|
|
timer = vispy.app.Timer()
|
|
da = Data(view)
|
|
# Problem is how slow update() is:
|
|
timer.connect(da.update)
|
|
timer.start(0.05)
|
|
if sys.flags.interactive != 1:
|
|
vispy.app.run()
|
|
|
|
if __name__ == '__main__':
|
|
main()
|