unity and burst

Fluidsim.cs

using System;
using System.Collections.Generic;
using System.Runtime.CompilerServices;
using Unity.Burst;
using Unity.Collections;
using Unity.Mathematics;
using UnityEngine;

public unsafe class FluidSystem : IFluidSystem, IDisposable
{
    public float4 verts { get; set; }
    public float4 norms { get; set; }

    NativeArray<int> gridStarts;
    NativeArray<int> prevGridStarts;
    NativeArray<int> gridEnds;
    NativeArray<int> prevGridEnds;

    NativeArray<int> gridMakeStarts;
    NativeArray<int> gridMakeEnds;
    NativeArray<int> gridMakeParticleNexts;


    NativeArray<float> particleDensities;
    NativeArray<float4> particlePositions;
    NativeArray<float4> gridMakeParticlePositions;
    NativeArray<float4> particleVelocities;
    NativeArray<float4> gridMakeParticleVelocities;
    NativeArray<float4> particleForces;
    NativeArray<float4> gravDirection;



    // Physical constants...
    float MASS;
    float K;
    float HALF_K;
    float MU;
    float REST_DENSITY;
    float REST_DENSITY_TIMES_2;
    float SIGMA;
    float POINT_DAMPING;
    float H;
    float HH;
    float TIME_STEP;
    float GRAV_CONST;
    float SCALE;
    float COLLISION_ELASTICITY;

    float KERNEL_BASE;
    float PRESSURE_GRADIDENT_BASE;
    float VISCOSITY_LAPLACIAN_BASE;
    float GRADIENT_OR_LAPLACIAN_BASE;

    // Resolution related constants...
    float MC_RES;
    int MC_MARGIN;
    float MC_THRESHOLD;
    float MC_POLY_THRESHOLD;

    int PARTICLE_COUNT;
    int GRID_WIDTH;
    int GRID_HEIGHT;
    int GRID_WIDTH_HEIGHT;
    int GRID_DEPTH;
    int GRID_SIZE;

    float WIDTH;
    float HEIGHT;
    float DEPTH;

    bool sphereInited;
    int sphereId;
    bool centerDrawing;
    public bool drawParticlesMode;
    int drawParticlesTimer;
    int DRAW_PARTICLES_TIME_LIMIT;

    float4 drawingOffset;


    public FluidSystem(float width, float height, float depth) => internal_ctor(width, height, depth);

    [BurstCompile]
    private void internal_ctor(float width, float height, float depth)
    {
        MASS = 1.0f;
        GRAV_CONST= 150.0f;
        SCALE= .2f;
        K= 1000.0f; // Gas constant
        HALF_K= K * .5f;
        MU= .1f; // Mu for Fluid
        REST_DENSITY= 1.2f;
        REST_DENSITY_TIMES_2= REST_DENSITY * 2.0f;
        SIGMA= 1.8f; // 1.5f
        POINT_DAMPING= 2.0f;
        H= 1.5f; // Kernel width
        TIME_STEP= 0.01f;
        HH= H * H;
        PARTICLE_COUNT = 12000;
        GRID_WIDTH = (int)(width / H) + 1;
        GRID_HEIGHT = (int)(height / H) + 1;
        GRID_DEPTH = (int)(depth / H) + 1;
        GRID_WIDTH_HEIGHT = GRID_WIDTH * GRID_HEIGHT;
        WIDTH = width;
        HEIGHT= height;
        DEPTH= depth;
        COLLISION_ELASTICITY= 0.8f; // 1.f
        KERNEL_BASE= 315.0f / (64.0f * math.PI * POW9(H));
        GRADIENT_OR_LAPLACIAN_BASE= -945.0f / (32.0f * math.PI * POW9(H));
        VISCOSITY_LAPLACIAN_BASE= 45.0f / (math.PI * POW6(H));
        PRESSURE_GRADIDENT_BASE= -VISCOSITY_LAPLACIAN_BASE;

        GRID_SIZE = GRID_WIDTH * GRID_HEIGHT * GRID_DEPTH;

        MC_RES = 0.5f;
        MC_MARGIN = 3;
        MC_THRESHOLD = 2.5f * MASS / KERNEL_BASE;
        MC_POLY_THRESHOLD = MASS * .6f / KERNEL_BASE;

        sphereInited = false;
        drawParticlesMode = false;
        drawParticlesTimer = 0;
        centerDrawing = true;
        DRAW_PARTICLES_TIME_LIMIT = 60 * 60;

        particleDensities = new NativeArray<float>(PARTICLE_COUNT, Allocator.Persistent);
        particlePositions = new NativeArray<float4>(PARTICLE_COUNT, Allocator.Persistent);
        particleVelocities = new NativeArray<float4>(PARTICLE_COUNT, Allocator.Persistent);
        particleForces = new NativeArray<float4>(PARTICLE_COUNT, Allocator.Persistent);
        gravDirection = new NativeArray<float4>(1, Allocator.Persistent);

        gridMakeParticlePositions = new NativeArray<float4>(PARTICLE_COUNT, Allocator.Persistent);
        gridMakeParticleVelocities = new NativeArray<float4>(PARTICLE_COUNT, Allocator.Persistent);
        gridStarts = new NativeArray<int>(GRID_SIZE + 2, Allocator.Persistent);
        gridEnds = new NativeArray<int>(GRID_SIZE + 2, Allocator.Persistent);
        gridMakeStarts = new NativeArray<int>(GRID_SIZE + 2, Allocator.Persistent);
        gridMakeEnds = new NativeArray<int>(GRID_SIZE + 2, Allocator.Persistent);
        prevGridStarts = new NativeArray<int>(GRID_SIZE + 2, Allocator.Persistent);
        prevGridEnds = new NativeArray<int>(GRID_SIZE + 2, Allocator.Persistent);
        gridMakeParticleNexts = new NativeArray<int>(PARTICLE_COUNT, Allocator.Persistent);

        for (var p = 0; p < PARTICLE_COUNT; ++p) gridMakeParticleNexts[p] = -1;

        for (var g = 0; g < GRID_SIZE; ++g)
            gridStarts[g] =
                gridEnds[g] =
                    gridMakeStarts[g] =
                        gridMakeEnds[g] =
                            prevGridStarts[g] =
                                prevGridEnds[g] = -1;

        for (var g = GRID_SIZE; g < GRID_SIZE + 2; ++g)
            gridStarts[g] =
                gridEnds[g] =
                    gridMakeStarts[g] =
                        gridMakeEnds[g] =
                            prevGridStarts[g] =
                                prevGridEnds[g] = PARTICLE_COUNT - 1;

        var r = 0;
        var loop = true;
        while (loop)
        {
            for (var j = 0; loop && j < height; ++j)
            {
                for (var k = 0; loop && k < depth; ++k)
                {
                    for (var i = 0; loop && i < width; ++i)
                    {
                        var pos = new float4(i, j, k, 0);

                        particlePositions[r] = pos;
                        particleVelocities[r] = new float4(0.0f);
                        particleForces[r] = new float4(0.0f);
                        particleDensities[r] = 0;

                        r++;

                        if (r == PARTICLE_COUNT)
                        {
                            loop = false;
                        }

                    }
                }
            }
        }

        updateGrid(); // Take note... must update grid here too.

        if (centerDrawing) drawingOffset = math.mul(new float4(WIDTH, HEIGHT, DEPTH, 0f), new float4(-0.5f, -0.5f, -0.5f, -0.5f));
        else drawingOffset = 0;
        Debug.Log("Success create fluid system");
    }

    [BurstCompile]
    public void updateGrid()
    {
        for (int g = 0; g < GRID_SIZE; ++g)
        {
#if OPENCL_SPH
            prevGridStarts[g] = gridStarts[g];
            prevGridEnds[g] = gridEnds[g];
#endif
            gridMakeStarts[g] = gridMakeEnds[g] = -1;
        }

        for (int p = 0; p < PARTICLE_COUNT; ++p)
        {
            float4 positionIndexes = particlePositions[p] / H;

            int g = POSITION_GRID_INDEX(positionIndexes);

            if (gridMakeStarts[g] < 0)
            {
                gridMakeStarts[g] = p;
                gridMakeEnds[g] = p;
            }
            else
            {
                gridMakeParticleNexts[gridMakeEnds[g]] = p;
                gridMakeEnds[g] = p;
            }
            gridMakeParticleNexts[p] = -1;
        }
        int c = 0;

        for (int g = 0; g < GRID_SIZE; ++g)
        {
            gridStarts[g] = c;
            for (int p = gridMakeStarts[g]; p > -1; p = gridMakeParticleNexts[p])
            {
                gridMakeParticlePositions[c] = particlePositions[p];
                gridMakeParticleVelocities[c] = particleVelocities[p];
                ++c;
            }
            gridEnds[g] = c;
        }
        for (int p = 0; p < PARTICLE_COUNT; ++p)
        {
            particleVelocities[p] = gridMakeParticleVelocities[p];
#if OPENCL_SPH
            swap(particlePositions[p], gridMakeParticlePositions[p]);
#else
            particlePositions[p] = gridMakeParticlePositions[p];
#endif
        }
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    private int POSITION_GRID_INDEX(float4 pos) =>
        GRID_INDEX((int)pos[0],
            (int)pos[1],
            (int)pos[2]);

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    private int GRID_INDEX(int i, int j, int k)
        => (GRID_WIDTH_HEIGHT * k + GRID_WIDTH * j + i);

    public void setGravDirection(float4 gravDirection)
        => this.gravDirection[0] = gravDirection;
    [BurstCompile]
    public void updateDensities()
    {
        int gw = GRID_WIDTH;
        int gwh = GRID_WIDTH_HEIGHT;
        int gwmm = GRID_WIDTH - 1, ghmm = GRID_HEIGHT - 1, gdmm = GRID_DEPTH - 1;
        int p, pEnd, i, j, k, gridK, gridKEnd, gridJK, gridJKEnd, gridIJK, gridIJKEnd;
        for (k = 0, gridK = 0, gridKEnd = GRID_DEPTH * gwh; gridK < gridKEnd; ++k, gridK += gwh)
            for (j = 0, gridJK = gridK, gridJKEnd = gridK + gwh; gridJK < gridJKEnd; ++j, gridJK += gw)
                for (i = 0, gridIJK = gridJK, gridIJKEnd = gridJK + gw; gridIJK < gridIJKEnd; ++i, ++gridIJK)
                    for (p = gridStarts[gridIJK], pEnd = gridEnds[gridIJK]; p < pEnd; ++p)

                    {
                        float pDensity = 0.0f;
                        float4 particlePosition = particlePositions[p];
                        int zStart = math.max(k - 1, 0), zEnd = math.min(k + 1, gwmm);
                        int yStart = math.max(j - 1, 0), yEnd = math.min(j + 1, ghmm);
                        int xStart = math.max(i - 1, 0), xEnd = math.min(i + 1, gdmm);
                        int np, npEnd, gridZ, gridYZ, gridXYZ, gridZEnd, gridYZEnd, gridXYZEnd;
                        int gridYStart = gw * yStart;
                        int gridYEnd = gw * yEnd;
                        for (gridZ = gwh * zStart, gridZEnd = gwh * zEnd; gridZ <= gridZEnd; gridZ += gwh)
                            for (gridYZ = gridZ + gridYStart, gridYZEnd = gridZ + gridYEnd; gridYZ <= gridYZEnd; gridYZ += gw)
                                for (gridXYZ = gridYZ + xStart, gridXYZEnd = gridYZ + xEnd; gridXYZ <= gridXYZEnd; ++gridXYZ)
                                    for (np = gridStarts[gridXYZ], npEnd = gridEnds[gridXYZ]; np < npEnd; ++np)
                                        if (p <= np)
                                        {
                                            //float4 d = particlePosition - particlePositions[np];
                                            //float rSq = math.abs(math.sqrt(d));

                                            float rSq = math.distance(particlePosition, particlePositions[np]);
                                            if (rSq <= HH)
                                            {
                                                float kernelDiff = HH - rSq;
                                                float common = KERNEL(kernelDiff) * MASS;
                                                pDensity += common;
                                                particleDensities[np] += common;
                                            }
                                        }
                        particleDensities[p] += pDensity;
                    }
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    private float KERNEL(float kernelDiff) => KERNEL_BASE * CUBE(kernelDiff);

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    private float CUBE(float x) => ((x) * (x) * (x));

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    private float POW6(float x) => (CUBE(x) * CUBE(x));
    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    private float POW9(float x) => (POW6(x) * CUBE(x));

    [BurstCompile]
    public void updateForces()
    {
        float4 gravCommon = MASS * GRAV_CONST * gravDirection[0];

        int gw = GRID_WIDTH;
        int gwh = GRID_WIDTH_HEIGHT;
        int gwmm = GRID_WIDTH - 1, ghmm = GRID_HEIGHT - 1, gdmm = GRID_DEPTH - 1;
        int p, pEnd, i, j, k, gridK, gridKEnd, gridJK, gridJKEnd, gridIJK, gridIJKEnd;
        for (k = 0, gridK = 0, gridKEnd = GRID_DEPTH * gwh; gridK < gridKEnd; ++k, gridK += gwh)
            for (j = 0, gridJK = gridK, gridJKEnd = gridK + gwh; gridJK < gridJKEnd; ++j, gridJK += gw)
                for (i = 0, gridIJK = gridJK, gridIJKEnd = gridJK + gw; gridIJK < gridIJKEnd; ++i, ++gridIJK)
                    for (p = gridStarts[gridIJK], pEnd = gridEnds[gridIJK]; p < pEnd; ++p)
                    {
                        float particleDensity = particleDensities[p];
                        float4 pForce = particleDensity * gravCommon;
                        float4 particlePosition = particlePositions[p];
                        float particleDensityRecip = MASS / particleDensity;

                        float pDensity = 0.0f;
                        int zStart = math.max(k - 1, 0), zEnd = math.min(k + 1, gwmm);
                        int yStart = math.max(j - 1, 0), yEnd = math.min(j + 1, ghmm);
                        int xStart = math.max(i - 1, 0), xEnd = math.min(i + 1, gdmm);
                        int np, npEnd, gridZ, gridYZ, gridXYZ, gridZEnd, gridYZEnd, gridXYZEnd;
                        int gridYStart = gw * yStart;
                        int gridYEnd = gw * yEnd;
                        for (gridZ = gwh * zStart, gridZEnd = gwh * zEnd; gridZ <= gridZEnd; gridZ += gwh)
                            for (gridYZ = gridZ + gridYStart, gridYZEnd = gridZ + gridYEnd; gridYZ <= gridYZEnd; gridYZ += gw)
                                for (gridXYZ = gridYZ + xStart, gridXYZEnd = gridYZ + xEnd; gridXYZ <= gridXYZEnd; ++gridXYZ)
                                    for (np = gridStarts[gridXYZ], npEnd = gridEnds[gridXYZ]; np < npEnd; ++np)
                                        if (p <= np)
                                        {
                                            float4 d = particlePosition - particlePositions[np];
                                            float rSq = math.distance(particlePosition, particlePositions[np]);
                                            if (rSq <= HH)
                                            {
                                                float r = math.sqrt(rSq);
                                                float neighborDensity = particleDensities[np];
                                                float neighborDensityRecip = MASS / neighborDensity;
                                                float kernelDiff = HH - rSq;
                                                float kernelDiff2 = H - r;
                                                float4 npForce = float4.zero;

                                                // Compute the pressure force.
                                                float4 common = HALF_K
                                                                * (particleDensity + neighborDensity - REST_DENSITY_TIMES_2)
                                                                * PRESSURE_GRADIENT(d, r, kernelDiff2);
                                                pForce -= neighborDensityRecip * common;
                                                npForce += particleDensityRecip * common;

                                                // Compute the viscosity force.
                                                common = MU * (particleVelocities[np] - particleVelocities[p])
                                                            * VISCOSITY_LAPLACIAN(kernelDiff2);
                                                pForce += neighborDensityRecip * common;
                                                npForce -= particleDensityRecip * common;

                                                // Compute the gradient of the color field.
                                                common = GRADIENT(d, kernelDiff);
                                                float4 particleColorGradient = neighborDensityRecip * common;
                                                float4 neighbourColorGradient = particleDensityRecip * common;

                                                // Compute the laplacian of the color field.
                                                float value = SIGMA * LAPLACIAN(kernelDiff, rSq);

                                                float particleColorGradLen = math.abs(math.length(particleColorGradient));
                                                if (particleColorGradLen > 0.001f)
                                                {
                                                    pForce -= neighborDensityRecip * value
                                                                                   * (particleColorGradient / particleColorGradLen);
                                                }

                                                float neighborColorGradLen =
                                                    math.abs(math.length(neighbourColorGradient));
                                                if (neighborColorGradLen > 0.001f)
                                                {
                                                    npForce -= particleDensityRecip * value * (neighbourColorGradient / neighborColorGradLen);
                                                }
                                                particleForces[np] += npForce;

                                            }
                                        }
                        particleForces[p] += pForce;

                    }
    }

    [BurstCompile]
    public void updateParticles()
    {
        for (int p = 0; p < PARTICLE_COUNT; ++p)
        {
            float4 vel = particleVelocities[p];
            float4 pos = particlePositions[p];
            vel += TIME_STEP * (particleForces[p] / particleDensities[p] - (POINT_DAMPING * vel) / MASS); // step vel
            pos += TIME_STEP * vel; // step pos

            // boundaries
            if (pos[0] < 0.0f)
            {
                pos[0] = 0.0f;
                vel[0] *= -COLLISION_ELASTICITY;
            }
            else if (pos[0] > WIDTH)
            {
                pos[0] = WIDTH;
                vel[0] *= -COLLISION_ELASTICITY;
            }

            if (pos[1] < 0.0f)
            {
                pos[1] = 0.0f;
                vel[1] *= -COLLISION_ELASTICITY;
            }
            else if (pos[1] > HEIGHT)
            {
                pos[1] = HEIGHT;
                vel[1] *= -COLLISION_ELASTICITY;
            }

            if (pos[2] < 0.0f)
            {
                pos[2] = 0.0f;
                vel[2] *= -COLLISION_ELASTICITY;
            }
            else if (pos[2] > DEPTH)
            {
                pos[2] = DEPTH;
                vel[2] *= -COLLISION_ELASTICITY;
            }

            particlePositions[p] = pos; // write back pos
            particleVelocities[p] = vel; // write back vel
        }
    }


    public void draw()
    {
        updateDensities();
        updateForces();
        updateParticles();

        //if (!sphereInited)
        //{
        //    GL.Begin(1);
        //    sphereId = glGenLists(1);
        //    glNewList(sphereId, GL_COMPILE);
        //    glutSolidSphere(0.3 * SCALE, 5, 5);
        //    glEndList();
        //    sphereInited = true;
        //}
        if (drawParticlesMode)
        {
            for (int p = 0; p < PARTICLE_COUNT; ++p)
                drawParticle(p);
            if (++drawParticlesTimer > DRAW_PARTICLES_TIME_LIMIT) drawParticlesMode = false;
        }
        else
        {
            //verts.clear();
            //norms.clear();
            //faces.clear();

            //generateIsoSurface(verts, norms, faces);
            //drawFaces(verts, norms, faces);

        }
    }

    public void generateIsoSurface(float4 verts, float4 norms, List<float3> faces)
    {
        throw new NotImplementedException();
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public float4 PRESSURE_GRADIENT(float4 d, float r, float kernelDiff2)
        => ((r < 0.001f) ? 0.0f : (PRESSURE_GRADIDENT_BASE * (kernelDiff2 * kernelDiff2) * d / r));


    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public float4 GRADIENT(float4 d, float kernelDiff) 
        => (-GRADIENT_OR_LAPLACIAN_BASE * (kernelDiff * kernelDiff) * d);

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public float4 VISCOSITY_LAPLACIAN(float kernelDiff2)
        => (VISCOSITY_LAPLACIAN_BASE * kernelDiff2);
    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public float LAPLACIAN(float kernelDiff, float rSq)
        => (GRADIENT_OR_LAPLACIAN_BASE * kernelDiff * (7.0f * rSq - 3.0f * HH));

    public void drawParticle(int particleId)
    {
        float4 p = SCALE * (particlePositions[particleId] + drawingOffset);
        Gizmos.DrawSphere(p.xyz, 0.2f);
        //GL.Begin(0);
        //GL.MultMatrix(glTranslatef(+p[0], +p[1], +p[2]));
        ////glCallList(sphereId);
        //GL.MultMatrix(glTranslatef(-p[0], -p[1], -p[2]));
        //GL.End();
    }

    public void Dispose()
    {
        gridStarts.Dispose();
        prevGridStarts.Dispose();
        gridEnds.Dispose();
        prevGridEnds.Dispose();
        gridMakeStarts.Dispose();
        gridMakeEnds.Dispose();
        gridMakeParticleNexts.Dispose();
        particleDensities.Dispose();
        particlePositions.Dispose();
        gridMakeParticlePositions.Dispose();
        particleVelocities.Dispose();
        gridMakeParticleVelocities.Dispose();
        particleForces.Dispose();
        gravDirection.Dispose();
    }

    Matrix4x4 glTranslatef(float x, float y, float z)
    {
        var m = new Matrix4x4();
        m[0, 0] = 1;
        m[0, 1] = 0;
        m[0, 2] = 0;
        m[0, 3] = x;
        m[1, 0] = 0;
        m[1, 1] = 1;
        m[1, 2] = 0;
        m[1, 3] = y;
        m[2, 0] = 0;
        m[2, 1] = 0;
        m[2, 2] = 1;
        m[2, 3] = z;
        m[3, 0] = 0;
        m[3, 1] = 0;
        m[3, 2] = 0;
        m[3, 3] = 1;
        return m;
    }
}

public unsafe interface IFluidSystem
{
    float4 verts { get; set; }
    float4 norms { get; set; }


    void updateGrid(); // To be done in CPU, as non-trivially parallel

    void updateDensities(); // To be done in OpenCL

    void updateForces();  // To be done in OpenCL

    void updateParticles();  // To be done in OpenCL

    void generateIsoSurface(float4 verts, float4 norms, List<float3> faces);

    void drawParticle(int p);
}