Implementation of spatial lookup in Unity fluid simulation

Question

I'm trying to optimise my fluid simulation by implementing a spatial lookup algorithm I have read about online, but I am struggling to make it work with Unity's Job system.

I have very limited experience with Unity's job system and spatial lookup algorithms, so I am not quite sure what is causing some errors. I get lots of leaks in the console log

I believe I have made my SpatialLookup class compatible with Unity's Job system, but the main problem is the leaks that I'm unable to find and fix

Here is the code for my job system and my SpatialLookup and Cell class

Classes:

using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;
using Unity.Collections;
using Unity.Mathematics;


public class SpatialLookup
{
    public NativeArray<Cell> cells; // Array of cells for spatial indexing
    private float cellSize;
    public int gridHeight; // Number of rows
    public int gridWidth; // Number of columns

    public SpatialLookup(float smoothingRadius, int cameraWidth, int cameraHeight, int maxMoleculesPerCell, Allocator allocator)
    {
        this.cellSize = smoothingRadius;
        this.gridWidth = (int)Math.Ceiling(cameraWidth / cellSize); // Calculated number of columns
        this.gridHeight = (int)Math.Ceiling(cameraHeight / cellSize); // Calculated number of rows
        cells = new NativeArray<Cell>(gridWidth * gridHeight, allocator); // Create cells array

        // Create a cell for each space in the array
        for (int i = 0; i < cells.Length; i++)
        {
            cells[i] = new Cell(maxMoleculesPerCell, allocator);
        }
    }

    public (int, int) GetCellIndex(float2 position)
    {
        // Calculate the cell a molecule belongs in based on its position
        int cellX = (int)math.clamp(position.x / cellSize, 0, gridWidth - 1);
        int cellY = (int)math.clamp(position.y / cellSize, 0, gridHeight - 1);
        return (cellX, cellY);
    }

    public void AddMolecule(int index, float2 position)
    {
        var (cellX, cellY) = GetCellIndex(position); // Get cell index
        if (IsValidCell(cellX, cellY)) // If valid...
        {
            cells[cellX + cellY * gridWidth].Add(index); // Add molecule index to the cell
        }
    }

    public NativeArray<int> GetNeighbouringMolecules(float2 position, Allocator allocator)
    {
        NativeList<int> neighbours = new NativeList<int>(allocator);
        (int cellX, int cellY) = GetCellIndex(position);

        // Loop through the neighbouring cells (3x3 around current cell)
        for (int dx = -1; dx <= 1; dx++)
        {
            for (int dy = -1; dy <= 1; dy++)
            {
                int neighbourCellX = cellX + dx;
                int neighbourCellY = cellY + dy;

                if (IsValidCell(neighbourCellX, neighbourCellY))
                {
                    Cell cell = cells[neighbourCellX + neighbourCellY * gridWidth]; // Get the cell object
                    
                    // Add all molecules from that cell to the neighbours list
                    for (int i = 0; i < cell.MoleculeIndices.Length; i++)
                    {
                        neighbours.Add(cell.MoleculeIndices[i]);
                    }
                }
            }
        }

        // Convert to NativeArray<int> so I can use it for the job system
        NativeArray<int> neighbourArray = new NativeArray<int>(neighbours.Length, allocator);
        for (int i = 0; i < neighbours.Length; i++)
        {
            neighbourArray[i] = neighbours[i];
        }

        neighbours.Dispose();
        return neighbourArray;
    }

    public bool IsValidCell(int cellX, int cellY)
    {
        return cellX >= 0 && cellX < gridWidth && cellY >= 0 && cellY < gridHeight;
    }

    public void Dispose()
    {
        // Dispose resources for each cell so there are no leaks (bad) 
        for (int i = 0; i < cells.Length; i++)
        {
            cells[i].MoleculeIndices.Dispose();
        }
        if (cells.IsCreated)
            cells.Dispose(); // Dispose cell array
    }
}


using Unity.Collections;
using Unity.Collections.LowLevel.Unsafe;
using System.Collections.Generic;

public struct Cell
{
    public NativeArray<int> MoleculeIndices; // Holds indicies of molecules (from molecules list) in this cell
    public int Count; // Holds the number of molecules currently in the cell

    public Cell(int maxSize, Allocator allocator)
    {
        MoleculeIndices = new NativeArray<int>(maxSize, allocator);
        Count = 0;
    }
    
    // woah we learnt this in lesson
    public void Add(int index)
    {
        if (Count < MoleculeIndices.Length)
        {
            MoleculeIndices[Count++] = index;
        }
    }

    public void Clear()
    {
        Count = 0; // Clear cell
    }
}

Job System

    int maxNeighboursPerCell = 1000;
    public float smoothingRadius = 2f;
    public float restDensity = 1f;
    public float pressureMultiplier = 1f;
    public float deltaTime;

    /// <summary>
    /// Calculates the repulsive force between the molecules
    /// </summary>
    public void CalculateForces()
    {
        SpatialLookup spatialLookup = new SpatialLookup(smoothingRadius, 18, 10, maxNeighboursPerCell, Allocator.Temp);
        int moleculeCount = molecules.Count;
        if (molecules.Count == 0) { return; }

        deltaTime = Time.deltaTime;

        // Create the native arrays for the job. I use native arrays to take advantage of the multithreading used in Jobs
        NativeArray<float2> positions = new(moleculeCount, Allocator.TempJob);
        NativeArray<float2> velocities = new(moleculeCount, Allocator.TempJob);
        NativeArray<float> densities = new(moleculeCount, Allocator.TempJob);
        NativeArray<float2> pressures = new(moleculeCount, Allocator.TempJob);
        NativeArray<float2> neighbourMoleculesPos = new(moleculeCount, Allocator.TempJob);
        NativeArray<int> neighbourArray = new NativeArray<int>(moleculeCount * maxNeighboursPerCell, Allocator.TempJob); // Or other appropriate allocator


        // Copy data to the arrays
        for (int i = 0; i < moleculeCount; i++)
        {
            positions[i] = molecules[i].position;
            velocities[i] = molecules[i].velocity;
        }

        GetNeighbourMoleculesJob neighbourJob = new()
        {
            positions = positions,
            neighbouringMoleculesPos = neighbourMoleculesPos,
            spatialLookup = spatialLookup,
            neighbouringMoleculesInt = neighbourArray,
            maxNeighboursPerCell = maxNeighboursPerCell,
        };
        JobHandle neighbourHandle = neighbourJob.Schedule(moleculeCount, 64);

        CalculateDensitiesJob densityJob = new()
        {
            positions = positions,
            velocities = velocities,
            densities = densities,
            smoothingRadius = smoothingRadius,
            deltaTime = deltaTime
        };
        JobHandle densityHandle = densityJob.Schedule(moleculeCount, 64, neighbourHandle);

        CalculatePressureForcesJob pressureJob = new()
        {
            positions = positions,
            velocities = velocities,
            densities = densities,
            pressures = pressures,
            smoothingRadius = smoothingRadius,
            restDensity = restDensity,
            pressureMultiplier = pressureMultiplier,
            deltaTime = deltaTime
        };
        JobHandle pressureHandle = pressureJob.Schedule(moleculeCount, 64, densityHandle);

        pressureHandle.Complete();


        for (int i = 0; i < moleculeCount; i++)
        {
            molecules[i].ApplyForce(pressures[i]);
        }


        positions.Dispose();
        velocities.Dispose();
        densities.Dispose();
        pressures.Dispose();
        neighbourMoleculesPos.Dispose();
        neighbourArray.Dispose();
    }

    public static float KernelSmoother(float x, float r)
    {
        if (x >= r) return 0;

        float volume = (math.PI * math.pow(r, 4) / 6);
        return (r - x) * (r - x) / volume;
    }

    public static float KernelSmootherGradient(float x, float r)
    {
        if (x >= r) return 0;

        float scale = 12 / (math.pow(r, 4) * math.PI);
        return (x - r) * scale;
    }
}

[BurstCompile]
struct CalculateDensitiesJob : IJobParallelFor
{
    [ReadOnly] public NativeArray<float2> positions;
    [ReadOnly] public NativeArray<float2> velocities;
    [WriteOnly] public NativeArray<float> densities;

    public float smoothingRadius;
    public float deltaTime;
       
    public void Execute(int i)
    {
        float density = 0f;
        const float mass = 1f;
        float2 m1Pos = positions[i] + velocities[i] * deltaTime;

        for (int j = 0; j < positions.Length; j++)
        {
            float dst = math.length(positions[j] - m1Pos);
            float influence = MoleculeManager.KernelSmoother(dst, smoothingRadius);
            density += mass * influence;
        }
        densities[i] = density;
    }
}

[BurstCompile]
struct CalculatePressureForcesJob : IJobParallelFor
{
    [ReadOnly] public NativeArray<float2> positions;
    [ReadOnly] public NativeArray<float2> velocities;
    [ReadOnly] public NativeArray<float> densities;
    public NativeArray<float2> pressures;

    public float smoothingRadius;
    const float mass = 1f;
    public float restDensity;
    public float pressureMultiplier;
    public float deltaTime;


    public void Execute(int i)
    {
        float2 pressureForce = float2.zero;
        float2 acceleration = pressures[i] / densities[i]; // Approximate acceleration
        float2 predictedVelocity = velocities[i] + acceleration * deltaTime;
        float2 m1Pos = positions[i] + predictedVelocity * deltaTime;

        for (int j = 0; j < positions.Length; j++)
        {
            if (i == j) continue;

            float dst = math.length(m1Pos - positions[j]);
            if (dst == 0) continue;
            float2 dir = (m1Pos - positions[j]) / dst;

            float slope = MoleculeManager.KernelSmootherGradient(dst, smoothingRadius);
            float density = densities[i];
            float sharedPressure = CalculateSharedPressure(density, densities[j]);
            pressureForce += sharedPressure * mass * slope * dir / density;
        }

        pressures[i] = pressureForce;
    }

    private readonly float ConvertDensityToPressure(float density, float restDensity, float pressureMultiplier)
    {
        float densityOffset = density - restDensity;
        float pressure = densityOffset * pressureMultiplier;
        return pressure;
    }

    private readonly float CalculateSharedPressure(float density_i, float density_j)
    {
        float pressure_i = -ConvertDensityToPressure(density_i, restDensity, pressureMultiplier);
        float pressure_j = -ConvertDensityToPressure(density_j, restDensity, pressureMultiplier);
        return (pressure_i + pressure_j) / 2;
    }
}

[BurstCompile]
struct GetNeighbourMoleculesJob : IJobParallelFor
{
    [ReadOnly] public NativeArray<float2> positions;
    [WriteOnly] public NativeArray<float2> neighbouringMoleculesPos;
    [ReadOnly] public SpatialLookup spatialLookup;

    public NativeArray<int> neighbouringMoleculesInt;
    public int maxNeighboursPerCell;

    public void Execute(int i)
    {
        // Get neighbouring molecules for the position
        int count = 0; // Track how many neighbouring molecules are found
        NativeArray<int> tempNeighbours = spatialLookup.GetNeighbouringMolecules(positions[i], Allocator.Temp);
        count = tempNeighbours.Length;

        // Store neighbouring molecule positions
        for (int j = 0; j < count; j++)
        {
            neighbouringMoleculesPos[i * maxNeighboursPerCell + j] = positions[tempNeighbours[j]];
        }
    }
}

If there is a better way of handling this then that would be greatly appreciated. Thank you!