A failure of a fluid simulation. I'll get back to this, but it'll be in an actual repo.

main.py

from vispy import app, gloo, io
import numpy as np
from scipy.spatial.distance import pdist, squareform

REFLECT_MATRIX_X = np.array([[1, 0],
                             [0, -1]])
REFLECT_MATRIX_Y = np.array([[-1, 0],
                             [0, 1]])
IDEAL_GAS_CONST = 5 # J/Kmol
REST_DENSITY = 0.8 # Experimentally determined
GRAVITY = np.array([0, -10]) # N/kg - TODO: make these values reasonable
DELTA_TIME = 0.003

def load_shader():
    """
    Loads the particle vertex and fragment shaders.
    Creates and returns a shader program.
    """
    with open('particle.vert') as vertex_file:
        vertex = '\n'.join(vertex_file.readlines())
    with open('particle.frag') as fragment_file:
        fragment = '\n'.join(fragment_file.readlines())
    return gloo.Program(vertex, fragment)

class Canvas(app.Canvas):
    """
    Canvas holds all vispy callbacks and setup up OpenGL
    """
    def __init__(self, save_images=True, *args, **kwargs):
        app.Canvas.__init__(self, *args, **kwargs)
        gloo.set_state(clear_color=(.1, .1, .1, 1.0),
                       blend=True,
                       blend_func=('src_alpha', 'one_minus_src_alpha'))
        self.program = load_shader()

        self.save_images = save_images

        self.count = 200
        self.particle_mass = 1/self.count

        self.left_boundary = -0.8
        self.right_boundary = 0.8
        self.top_boundary = -0.8
        self.bottom_boundary = 0.8

        self.particles = np.zeros(self.count, dtype=[('position', np.float32, 2),
                                                     ('velocity', np.float32, 2),
                                                     ('force', np.float32, 2),
                                                     ('density', np.float32, 1),
                                                     ('pressure', np.float32, 1)])


        self.particles['position'] = np.random.uniform(self.left_boundary, 0.0, (self.count, 2))
        self.particles['velocity'] = np.zeros((self.count, 2))


        self.vbo = gloo.VertexBuffer()
        self.update_vbos()
        self.program.bind(self.vbo)
        # self.program.bind(self.vbo_position)
        # self.program.bind(self.vbo_color)

        self.update_vbos()

        self._timer = app.Timer('auto', connect=self.on_timer, start=True)
        self.frame_count = 0

        self.show()

    def update_vbos(self):
        """
        Sets the data of the vbo based on the particles
        """
        self.vbo.set_data(self.particles[['position', 'density', 'pressure', 'force']])


    def on_timer(self, event):
        """
        Calls update on every timer tick
        """
        self.frame_count += 1
        self.update()

    def on_resize(self, event):
        """
        Callback on resize. Called at start of program.
        Sets openGL viewport size and shader resolutions
        """
        width, height = self.physical_size
        print('Resize: ', self.physical_size)
        self.program['u_resolution'] = self.physical_size
        gloo.set_viewport(0, 0, width, height)

    def on_draw(self, event):
        """
        Called every frame by Canvas.update.
        """
        gloo.clear((0.1, 0.1, 0.1, 1.0))
        self.iteration()

        self.update_vbos()

        self.program.draw('points')
        if self.save_images:
            io.imsave(f'anim/sph{self.frame_count:04}.png', gloo.util._screenshot())

    def compute_distance_matrix(self):
        """
        Returns a sparse matrix showing which particles are sufficently close to each other.
        Computes the distances between those that are. 
        """
        self.sqdistances = squareform(pdist(self.particles['position'], metric='sqeuclidean'))
        self.distances = np.sqrt(self.sqdistances)


    def compute_density_pressure(self, h):
        """
        Computes the densities and pressures of each particle
        """
        # Density at a particle r = sum for all particles r' (mj * W(r-r', h)
        poly6coeff = 315/(64*np.pi*h**9)
        kernel = poly6coeff*(np.subtract(h**2, self.sqdistances, where=self.distances <= h)**3)
        # Take the sum of all the kernel values * mass to find the value of the density field at each particle
        # Make sure that is greater than the rest density because we don't want negative pressures. 
        # TODO: Test with and without this.
        self.particles['density'] = np.maximum(np.sum(self.particle_mass*kernel, axis=0), REST_DENSITY)
        # Compute the pressure using the ideal gas law
        self.particles['pressure'] = (self.particles['density'] - REST_DENSITY)*IDEAL_GAS_CONST

    def compute_pressure_force(self, h):
        spiky_coeff =  -45/(np.pi*h**6)
        differences = (self.particles['position'][:, np.newaxis, :] - self.particles['position'])
        directions = np.divide(differences, self.distances[:, :, np.newaxis], 
                               where=np.logical_and(self.distances != 0, self.distances<h)[:, :, np.newaxis])
        W = spiky_coeff*directions*np.square((h - self.distances)[:, :, np.newaxis])

        pressure_sums = self.particles['pressure'][:, np.newaxis] + self.particles['pressure']
        pterm = pressure_sums/(2*self.particles['density'].T)

        self.particles['force'] += np.sum(self.particle_mass*W*pterm[:, :, np.newaxis], axis=0)

    def edges(self): 
        vert_bounce = np.logical_or(self.particles['position'][:, 1] <= self.top_boundary,
                                    self.particles['position'][:, 1] >= self.bottom_boundary)
        self.particles['velocity'][vert_bounce] = (
            np.matmul(self.particles['velocity'][vert_bounce], REFLECT_MATRIX_X))
        self.particles['position'][vert_bounce] += self.particles['velocity'][vert_bounce]*DELTA_TIME

        horz_bounce = np.logical_or(self.particles['position'][:, 0] <= self.left_boundary,
                                    self.particles['position'][:, 0] >= self.right_boundary)
        self.particles['velocity'][horz_bounce] = (
            np.matmul(self.particles['velocity'][horz_bounce], REFLECT_MATRIX_Y))
        self.particles['position'][horz_bounce] += self.particles['velocity'][horz_bounce]*DELTA_TIME



    def iteration(self):
        """
        One iteration of the fluid simulation algorithm
        """
        self.particles['force'] = np.zeros_like(self.particles['force'])

        self.compute_distance_matrix()
        self.compute_density_pressure(1)
        self.compute_pressure_force(1)

        self.particles['force'] += self.particles['density'][:, np.newaxis]*[0, -10]

        # Update velocity and position. Physics Step.
        self.particles['velocity'] += self.particles['force']/self.particles['density'][:, np.newaxis]*DELTA_TIME
        self.particles['position'] += self.particles['velocity']*DELTA_TIME
        # Bounce particles on edge
        self.edges()

def main():
    """
    Main function
    """
    canvas = Canvas(title='Particle System Test',
                    size=(600, 600),
                    resizable=False,
                    vsync=False,
                    keys='interactive',
                    save_images=False)
    canvas.measure_fps(callback=lambda s : print(f'{canvas.frame_count}: {s:0.3} FPS'))
    app.run()



if __name__ == '__main__':
    main()
#!/usr/bin/python3
import requests
from PIL import Image
from io import BytesIO
import matplotlib.pyplot as plt
import matplotlib.animation as anim
import matplotlib as mpl
import numpy as np

width = height = 800
k = 12 
fig = plt.figure()

response = requests.get(f'https://picsum.photos/{width}/{height}/?random')
img = (Image.open(BytesIO(response.content)))
img = np.array(img)
flattened = np.reshape(img, (width*height, 3))

print('Image Recieved: ', response.headers)

def init_centroids(points, k):
    centroids = points.copy()
    np.random.shuffle(centroids)
    return centroids[:k]

centroids = init_centroids(flattened, k)
print('Centroids: ', centroids)


ax1 = fig.add_subplot(131)
ax1.imshow(img)
ax1.set_xticks([])
ax1.set_yticks([])

ax2 = fig.add_subplot(132)
im = ax2.imshow(img, animated=True)
ax2.set_xticks([])
ax2.set_yticks([])

# Create the subplot to show swatches of color
ax3 = fig.add_subplot(133)
ax3.set_xlim((0,k))
ax3.set_ylim((0, 1))
ax3.set_xticks([])
ax3.set_yticks([])
ax3.set_aspect('equal')


def closest_centroid(points, centroids):
    sq_dists = ((points - centroids[:, np.newaxis])**2).sum(axis=2)
    return np.argmin(sq_dists, axis=0)


def move_centroids(points, closest, centroids):
    return np.array([points[closest==c].mean(axis=0) for c in range(centroids.shape[0])])

def mean_squared_error(points, closest, centroids):
    # import ipdb; ipdb.set_trace()
    return np.sum([np.linalg.norm([points[closest==c] - centroids[c]]) for c in range(centroids.shape[0])])
    # return np.sum([(points[closest==c] - centroids[c])**2 for c in range(centroids.shape[0])])

def update(*args):
    global centroids

    closest = closest_centroid(flattened, centroids)
    segmented = centroids[np.reshape(closest, (width, height))]/255

    print('Error: ', mean_squared_error(flattened, closest, centroids))
    dirty=[]

    im.set_array(segmented)
    dirty.append(im)
    for x, color in enumerate(centroids):
        dirty.append(ax3.add_patch(mpl.patches.Rectangle((x, 0), 1, 1, facecolor=color/255)))

    centroids = move_centroids(flattened, closest, centroids)
    return dirty

ani = anim.FuncAnimation(fig, update, interval=20, blit=True)
plt.show()

particle.frag

precision mediump float;

varying vec3 v_color;

uniform vec2 u_resolution;

void main() {
    float x = 2.0*gl_PointCoord.x - 1.0;
    float y = 2.0*gl_PointCoord.y - 1.0;
    float a = x*x + y*y;
    a = a < 1.0 ? 1.0-a : 0.0;
    gl_FragColor = vec4(v_color, a);
}

particle.vert

attribute vec2 position;
attribute float pressure;
attribute float density;
attribute vec2 force;
//attribute vec2 a_velocity;

varying vec3 v_color;

void main() {
    gl_Position = vec4(position, 1.0, 1.0);
    v_color = vec3(pressure);
    gl_PointSize = 10.0;
}