bokeh-autoscale

autoscale_problem.py

from __future__ import division

import time
import numpy as np
import bokeh.plotting as bkhplt
from bokeh.objects import Range1d

RADIUS = 1
FRAMERATE = 1 / 24

class Physics(object):

    def __init__(self):
        N_SQRT = 60
        N = N_SQRT ** 2

        step = 4 * RADIUS
        grid = np.mgrid[-N_SQRT / 2 * step: N_SQRT / 2 * step: step,
                        -N_SQRT / 2 * step: N_SQRT / 2 * step: step]

        self.posX = grid[0].flatten().astype(np.float32)
        self.posY = grid[1].flatten().astype(np.float32)

        self.n = N
        self.radii = [RADIUS] * N
        self.colors = ["#%02x%02x%02x" % (r, g, 150)
                       for r, g in zip(np.floor(
                                        np.minimum(50+2*abs(self.posX), 255)),
                                       np.floor(
                                        np.minimum(30+2*abs(self.posY), 255)))]

        self.t = 0
        self.dt = 1

    def update(self):
        rate = 1.01
        if self.dt < 0:
            rate = 1/rate
        self.posX *= rate
        self.posY *= rate
        self.t += self.dt

        if self.dt > 0 and self.t > 100:
            self.dt = -1
        elif self.dt < 0 and self.t <= 0:
            self.dt = +1




def main():
    bkhplt.output_server("autoscale_problem.py example")

    bkhplt.hold()

    phy = Physics()


    bkhplt.scatter(phy.posX, phy.posY, radius=phy.radii, radius_units="data",
                   fill_color=phy.colors, fill_alpha=0.6,
                   line_color=None, tools="pan,zoom,resize", name="color_scatter_example")

    # TODO: pending fix in Bokeh
    bkhplt.curplot().x_range = Range1d(start=-250, end=250)
    bkhplt.curplot().y_range = Range1d(start=250, end=250)

    renderer = [r for r in bkhplt.curplot().renderers
                if isinstance(r, bkhplt.Glyph)][0]
    ds = renderer.data_source

    time.sleep(1)
    while True:
        phy.update()
        ds.data["x"] = phy.posX
        ds.data["y"] = phy.posY
        ds._dirty = True
        bkhplt.session().store_obj(ds)
        time.sleep(FRAMERATE)


if __name__ == '__main__':
    main()

particles.py

from __future__ import division

import time
import math
import numpy as np
from numbapro import cuda, float32, vectorize
import bokeh.plotting as bkhplt
from bokeh.objects import Range1d

MASS = 1
DAMPING = 0.001
GRAVITY = -0.05
RADIUS = 1
RIGHT_MARGIN = 150
FRAMERATE = 1 / 24
ACCURACY = 4

class Physics(object):
    class Device(object):
        pass

    def __init__(self):
        N_SQRT = 30
        N = N_SQRT ** 2

        step = 4 * RADIUS
        grid = np.mgrid[step:(N_SQRT * step + step):step,
                        step:(N_SQRT * step + step):step]

        noise = np.random.random(N) * RADIUS

        self.posX = cuda.pinned_array(shape=N, dtype=np.float32)
        self.posY = cuda.pinned_array(shape=N, dtype=np.float32)

        self.posX[:] = grid[0].flatten().astype(np.float32) + noise
        self.posY[:] = grid[1].flatten().astype(np.float32)

        self.n = N
        self.radii = [RADIUS] * N
        self.colors = ["#%02x%02x%02x" % (r, g, 150)
                       for r, g in zip(np.floor(50+2*self.posX),
                                       np.floor(30+2*self.posY))]

        self.dt = 1 / ACCURACY

        self.gpu = self.Device()
        self.gpu.stream = cuda.stream()

        # current positions
        self.gpu.cx = cuda.to_device(self.posX)
        self.gpu.cy = cuda.to_device(self.posY)

        # previous positions
        self.gpu.px = cuda.device_array_like(self.gpu.cx)
        self.gpu.py = cuda.device_array_like(self.gpu.cy)

        # next positions
        self.gpu.nx = cuda.device_array_like(self.gpu.cx)
        self.gpu.ny = cuda.device_array_like(self.gpu.cy)

        # forces
        self.gpu.fx = cuda.device_array_like(self.gpu.cx)
        self.gpu.fy = cuda.device_array_like(self.gpu.cy)

        # init previous positions
        device_copy(self.gpu.cx, out=self.gpu.px, stream=self.gpu.stream)
        device_copy(self.gpu.cy, out=self.gpu.py, stream=self.gpu.stream)

        # init gpu kernels

        self.gpu.init_force = device_init_force.configure(N_SQRT, N_SQRT,
                                                          self.gpu.stream)
        self.gpu.integrate = device_integrate.configure(N_SQRT, N_SQRT,
                                                        self.gpu.stream)
        self.gpu.wall_collide = device_wall_collide.configure(N_SQRT, N_SQRT,
                                                              self.gpu.stream)
        self.gpu.collide = device_collide.configure(N_SQRT, N_SQRT,
                                                    self.gpu.stream)

    def init_force(self):
        self.gpu.init_force(self.gpu.fx, self.gpu.fy)

    def integrate(self):
        self.gpu.integrate(self.gpu.cx, self.gpu.cy, self.gpu.px,
                           self.gpu.py, self.gpu.nx, self.gpu.ny,
                           self.gpu.fx, self.gpu.fy, self.dt)

    def swap(self):
        nx, ny = self.gpu.px, self.gpu.py
        self.gpu.px, self.gpu.py = self.gpu.cx, self.gpu.cy
        self.gpu.cx, self.gpu.cy = self.gpu.nx, self.gpu.ny
        self.gpu.nx, self.gpu.ny = nx, ny

    def collision(self):
        self.gpu.wall_collide(self.gpu.cx, self.gpu.cy, self.gpu.px,
                              self.gpu.py, self.gpu.fx, self.gpu.fy, self.dt)
        self.gpu.collide(self.gpu.cx, self.gpu.cy, self.gpu.fx, self.gpu.fy,
                         self.dt)

    def update(self):
        for i in range(ACCURACY):
            self.init_force()
            self.collision()
            self.integrate()
            self.swap()
        # copy to host
        self.gpu.cx.copy_to_host(self.posX, stream=self.gpu.stream)
        self.gpu.cy.copy_to_host(self.posY, stream=self.gpu.stream)
        # sync gpu
        self.gpu.stream.synchronize()


@vectorize(['float32(float32)'], target='gpu')
def device_copy(src):
    return src


@cuda.jit('void(float32[:], float32[:])')
def device_init_force(ax, ay):
    i = cuda.grid(1)
    if i < ax.shape[0]:
        ax[i] = 0
        ay[i] = GRAVITY


@cuda.jit(argtypes=[float32[:]] * 8 + [float32])
def device_integrate(cx, cy, px, py, nx, ny, fx, fy, dt):
    i = cuda.grid(1)
    if i < cx.shape[0]:
        ax = fx[i] / MASS
        ay = fy[i] / MASS
        dm = float32(DAMPING)
        nx[i] = (2 - dm) * cx[i] - (1 - dm) * px[i] + ax * (dt * dt)
        ny[i] = (2 - dm) * cy[i] - (1 - dm) * py[i] + ay * (dt * dt)


@cuda.jit(argtypes=[float32[:]] * 6 + [float32])
def device_wall_collide(cx, cy, px, py, fx, fy, dt):
    i = cuda.grid(1)
    if i < cx.shape[0]:
        dtdt = (dt * dt)
        ax = (cx[i] - px[i]) / dtdt
        ay = (cy[i] - py[i]) / dtdt

        # left
        if cy[i] <= RADIUS and ay < 0:          # ground collision
            cy[i] = -cy[i] + 2 * RADIUS         # correct location
            fy[i] *= -(1 - DAMPING)             # reverse direction

        ## right
        if cx[i] <= RADIUS and ax < 0:
            cx[i] = -cx[i] + 2 * RADIUS         # correct location
            fx[i] *= -(1 - DAMPING)             # reverse direction
        elif cx[i] >= RIGHT_MARGIN - RADIUS and ax > 0:
            cx[i] = - cx[i] + 2 * (RIGHT_MARGIN - RADIUS)   # correct location
            fx[i] *= -(1 - DAMPING)             # reverse direction


@cuda.jit('float32(float32, float32, float32, float32)', device=True)
def distance_square(ax, ay, bx, by):
    dx = ax - bx
    dy = ay - by
    return math.sqrt(dx * dx + dy * dy)


@cuda.jit(argtypes=[float32[:]] * 4 + [float32])
def device_collide(cx, cy, fx, fy, dt):
    i = cuda.grid(1)
    n = cx.shape[0]
    if i < n:
        x = cx[i]
        y = cy[i]
        for j in range(n):
            if i != j:
                jx = cx[j]
                jy = cy[j]
                dist = distance_square(x, y, jx, jy)

                if dist <= 2 * RADIUS:
                    dtdt = dt * dt
                    overlapped = 2 * RADIUS - dist
                    dx = x - jx
                    dy = y - jy

                    total = overlapped / 2
                    ddx = total * (dx / dist)
                    ddy = total * (dy / dist)

                    nx = x + ddx
                    ny = y + ddy

                    fx[i] += (nx - x) / dtdt * MASS
                    fy[i] += (ny - y) / dtdt * MASS

#-----------------------------------------------------------------------------


def main():
    bkhplt.output_server("particles.py example")

    bkhplt.hold()

    phy = Physics()


    bkhplt.scatter(phy.posX, phy.posY, radius=phy.radii, radius_units="data",
                   fill_color=phy.colors, fill_alpha=0.6,
                   line_color=None, tools="pan,zoom,resize", name="color_scatter_example")

    # TODO: pending fix in Bokeh
    bkhplt.curplot().x_range = Range1d(start=0, end=500)
    bkhplt.curplot().y_range = Range1d(start=0, end=500)

    renderer = [r for r in bkhplt.curplot().renderers
                if isinstance(r, bkhplt.Glyph)][0]
    ds = renderer.data_source

    time.sleep(2)
    while True:
        phy.update()
        ds.data["x"] = phy.posX
        ds.data["y"] = phy.posY
        ds._dirty = True
        bkhplt.session().store_obj(ds)
        time.sleep(FRAMERATE)


if __name__ == '__main__':
    main()