#!/usr/bin/env python
# coding: utf-8

# In[ ]:


import k3d
from math import sin, cos, pi

width = height = length = 20
model_matrix = [
    7.0,  5.0, -5.0, 0.0,
    0.0,  7.0,  7.0, 5.0,
    7.0, -5.0,  5.0, 0.0,
    0.0,  0.0,  0.0, 1.0
]

def f(x, y, z):
    return sin(float(x) / width * pi * 2.0), cos(float(y) / height * pi * 2.0), sin(float(z) / length * pi * 2.0)

colors = (0xFF0000, 0x00FF00) * width * height * length
vectors = [[[f(x, y, z) for x in range(width)] for y in range(height)] for z in range(length)]

plot = k3d.plot()
plot += k3d.vector_field(vectors, colors, model_matrix=model_matrix, use_head=False)
plot.display()


# In[ ]:


plot.objects[0].scale = 10.0


# In[ ]:


plot.objects[0].use_head = False


# In[ ]:


#credits:https://github.com/Davide-sd

import numpy as np 

fig = k3d.plot()
n = 20
x = np.linspace(-5, 5, n)
y = np.linspace(-5, 5, n)
z = np.linspace(-5, 5, n)
xx, yy, zz = np.meshgrid(x, y, z)
uu, vv, ww = zz, yy, xx
xx, yy, zz, uu, vv, ww = [t.flatten().astype(np.float32) for t in [xx, yy, zz, uu, vv, ww]]
scale = 0.5
magnitude = np.sqrt(uu**2 + vv**2 + ww**2)
vectors = np.array((uu, vv, ww)).T * scale
origins = np.array((xx, yy, zz)).T
colors = k3d.helpers.map_colors(magnitude, k3d.matplotlib_color_maps.Plasma, [])
vec_colors = np.zeros(2 * len(colors))
for i, c in enumerate(colors):
    vec_colors[2 * i] = c
    vec_colors[2 * i + 1] = c
vec_colors = vec_colors.astype(np.uint32)
vec = k3d.vectors(
    origins = origins - vectors / 2,
    vectors = vectors,
    colors = vec_colors,
)
fig += vec
fig.display()


# In[ ]: