Unity Boid Simulation System using Job System and Burst Compiler

BoidSimulator.cs

using System;
using System.Collections;
using System.Collections.Generic;
using UnityEngine;
using Unity.Mathematics;
using Unity.Jobs;
using Unity.Burst;
using Unity.Collections;

public class BoidSimulator : IDisposable
{
    private bool _isDisposed;

    public struct BoidArgs
    {
        public int count;

        public float separateWeight;
        public float alignmentWeight;
        public float cohesionWeight;
        public float toCenterWeight; // Move to center point

        public float separateNeighborhoodRadiusSq;
        public float alignmentNeighborhoodRadiusSq;
        public float cohesionNeighborhoodRadiusSq;

        public float toCenterDistanceSq;
        public Vector3 toCenterScaleInv;

        public float steerForce;
        public float steerForceToCenter;
        public float maxSteerForce;
        public float maxSteerForceSq;
        public float forceMultiplier;

        public float dragIntensity;
        public float maxVelocity;
        public float maxVelocitySq;
    }

    public struct ExternalForce
    {
        public float3 targetPositionVelocity; // for each boid
        public float toTargetVelocity; // Instant velocity

        public float3 targetPositionForce; // for each boid
        public float toTargetForceWeight; // For steering

        public float3 velocity; // uniform velocity for all boids
        public float velocityIntensity;

        public float3 force; // uniform force for all boids (apply steer)
        public float forceWeight;

        public void ResetForceIfNecessary()
        {
            if (toTargetVelocity != 0f)
                toTargetVelocity = 0f;
            if (toTargetForceWeight != 0f)
                toTargetForceWeight = 0f;
            if (velocityIntensity != 0f)
                velocityIntensity = 0f;
            if (forceWeight != 0f)
                forceWeight = 0f;
        }
    }

    public struct BoidData
    {
        public float3 position;
        public float3 velocity;
    }

    private BoidArgs _boidArgs;
    private ExternalForce _externalForce;

    private NativeArray<BoidData> _boidDataArray;
    private NativeArray<float3> _boidForceArray; // Separate force buffer to avoid race condition
    public NativeArray<BoidData> data => _boidDataArray;

    // TODO: Find best value on Quest2
    // On M1 mac, 16: 6~7.5ms, 32: 3.5~4.5ms, 64: 3~4ms, 128: 5.5~6.5ms with Deep Profiling
    private const int INNER_LOOP_BATCH_COUNT = 64;
    private JobHandle _updateBoidJobHandle;

    public BoidSimulator(BoidArgs boidArgs, Vector3 centerPosition)
    {
        _externalForce = new ExternalForce();
        _boidArgs = boidArgs;
        int count = boidArgs.count;
        _boidDataArray = new NativeArray<BoidData>(count, Allocator.Persistent);
        _boidForceArray = new NativeArray<float3>(count, Allocator.Persistent);
        new InitializeBoidJob(_boidDataArray, centerPosition)
            .ScheduleParallel(count, INNER_LOOP_BATCH_COUNT, new JobHandle())
            .Complete();
    }

    #region IDisposable
    public void Dispose()
    {
        Dispose(true);
        GC.SuppressFinalize(this);
    }

    protected virtual void Dispose(bool disposing)
    {
        if (_isDisposed)
            return;

        if (disposing)
        {
            _updateBoidJobHandle.Complete();
            _boidDataArray.Dispose();
            _boidForceArray.Dispose();
        }

        _isDisposed = true;
    }

    ~BoidSimulator()
    {
        Dispose(false);
    }

    private void ThrowIfDisposed()
    {
        if (_isDisposed)
        {
            throw new ObjectDisposedException(nameof(BoidSimulator));
        }
    }
    #endregion

    public void UpdateArgs(BoidArgs boidArgs)
    {
        _boidArgs = boidArgs;
    }

    public void UpdateBoid(Vector3 centerPosition)
    {
        int count = _boidArgs.count;

        _updateBoidJobHandle =
            new CalculateBoidForceJob(_boidDataArray, _boidForceArray, _boidArgs)
                .ScheduleParallel(count, INNER_LOOP_BATCH_COUNT, _updateBoidJobHandle);
        _updateBoidJobHandle.Complete();

        _updateBoidJobHandle =
            new ApplyForceToBoidDataJob(_boidDataArray, _boidForceArray, _boidArgs, _externalForce, centerPosition, Time.deltaTime)
                .ScheduleParallel(count, INNER_LOOP_BATCH_COUNT, _updateBoidJobHandle);
        _updateBoidJobHandle.Complete();

        _externalForce.ResetForceIfNecessary();
    }

    public void AddToTargetVelocity(float3 targetPosition, float velocity)
    {
        _externalForce.targetPositionVelocity = targetPosition;
        _externalForce.toTargetVelocity = velocity;
    }

    public void AddToTargetForce(float3 targetPosition, float force)
    {
        _externalForce.targetPositionForce = targetPosition;
        _externalForce.toTargetForceWeight = force;
    }

    public void AddVelocity(float3 velocity, float velocityIntensity)
    {
        _externalForce.force = velocity;
        _externalForce.forceWeight = velocityIntensity;
    }

    public void AddVelocity(Vector3 velocity)
    {
        _externalForce.force = velocity;
        _externalForce.forceWeight = 1f;
    }

    public void AddForce(float3 force, float forceWeight)
    {
        _externalForce.force = force;
        _externalForce.forceWeight = forceWeight;
    }

    [BurstCompile(FloatPrecision.Low, FloatMode.Fast)]
    public struct InitializeBoidJob : IJobFor
    {
        [NativeDisableParallelForRestriction]
        [WriteOnly] private NativeArray<BoidData> _boidData;
        private float3 _swarmCenterPosition;

        public InitializeBoidJob(NativeArray<BoidData> boidData, float3 swarmCenterPosition)
        {
            this._boidData = boidData;
            this._swarmCenterPosition = swarmCenterPosition;
        }

        public void Execute(int index)
        {
            Unity.Mathematics.Random random = new Unity.Mathematics.Random((uint)index * 3498 + 1);

            // float3 randomPosition = new float3(random.NextFloat(-10, 10), random.NextFloat(-10, 10), random.NextFloat(-10, 10));
            // float3 randomVelocity = new float3(random.NextFloat(-1f, 1f), random.NextFloat(-1f, 1f), random.NextFloat(-1f, 1f));
            float3 randomPosition = math.normalize(random.NextFloat3() - 0.5f) * 10.0f;
            float3 randomVelocity = math.normalize(random.NextFloat3() - 0.5f) * 1f;

            _boidData[index] = new BoidData
            {
                position = _swarmCenterPosition + randomPosition,
                velocity = randomVelocity
            };
        }
    }

    [BurstCompile(FloatPrecision.Low, FloatMode.Fast)]
    public struct CalculateBoidForceJob : IJobFor
    {
        [ReadOnly] private NativeArray<BoidData> _boidData;

        [NativeDisableParallelForRestriction]
        [WriteOnly] private NativeArray<float3> _boidForce;
        private BoidArgs _boidArgs;

        public CalculateBoidForceJob(NativeArray<BoidData> boidData, NativeArray<float3> boidForce, BoidArgs boidArgs)
        {
            this._boidData = boidData;
            this._boidForce = boidForce;
            this._boidArgs = boidArgs;
        }

        public float3 CalculateSteer(float3 desired, float3 velocity, float steerForce, float maxForce, float maxForceSq)
        {
            float lengthSq = math.lengthsq(desired);
            if (lengthSq < 1e-5f)
                return float3.zero;

            // Normalize and scale in one step
            float3 desiredDir = desired * (math.rsqrt(lengthSq) * steerForce);

            // Calculate steering force
            float3 steer = desiredDir - velocity;

            LimitVector(ref steer, maxForce, maxForceSq);
            return steer;
        }

        public void LimitVector(ref float3 vector, float maxLength, float maxLengthSq)
        {
            float lengthSq = math.lengthsq(vector);
            if (lengthSq > maxLengthSq)
            {
                vector = vector * (maxLength * math.rsqrt(lengthSq));
            }
        }

        public void Execute(int index)
        {
            BoidData boidData = _boidData[index];

            float3 position = boidData.position;
            float3 velocity = boidData.velocity;

            float3 sepDirSum = float3.zero; // Separation
            int sepCount = 0;
            float3 aliVelSum = float3.zero; // Alignment
            int aliCount = 0;
            float3 cohPosSum = float3.zero; // Cohesion
            int cohCount = 0;

            // Calculate force for each boid
            for (int i = 0; i < _boidArgs.count; i++)
            {
                if (i == index)
                    continue;

                float3 neighborPosition = _boidData[i].position;

                float3 towardsThisBoid = position - neighborPosition;
                float distSquared = math.lengthsq(towardsThisBoid);

                // Use max radius to early exit
                float maxRadiusSq = math.max(math.max(
                    _boidArgs.separateNeighborhoodRadiusSq,
                    _boidArgs.alignmentNeighborhoodRadiusSq),
                    _boidArgs.cohesionNeighborhoodRadiusSq);

                if (distSquared > maxRadiusSq)
                    continue;

                // Separation
                if (distSquared <= _boidArgs.separateNeighborhoodRadiusSq)
                {
                    // float dist = math.sqrt(distSquared);
                    // float3 repulse = towardsThisBoid / dist; // normalize
                    // repulse /= dist; // Reduce the force by distance
                    float invDist = math.rsqrt(distSquared);  // Optimized 1/sqrt()
                    float3 repulse = towardsThisBoid * invDist * invDist; // Normalize and scale in one step
                    sepDirSum += repulse;
                    sepCount++; // To calculate average
                }

                // Alignment
                if (distSquared <= _boidArgs.alignmentNeighborhoodRadiusSq)
                {
                    float3 neighborVelocity = _boidData[i].velocity;  // Only fetch if needed
                    aliVelSum += neighborVelocity;
                    aliCount++;
                }

                // Cohesion
                if (distSquared <= _boidArgs.cohesionNeighborhoodRadiusSq)
                {
                    cohPosSum += neighborPosition;
                    cohCount++;
                }
            }

            float steeringForce = _boidArgs.steerForce;
            float maxSteeringForce = _boidArgs.maxSteerForce;
            float maxSteeringForceSq = _boidArgs.maxSteerForceSq;

            // Separation
            float3 sepSteer = (sepCount > 0)
                ? CalculateSteer(sepDirSum, velocity, steeringForce, maxSteeringForce, maxSteeringForceSq)
                : float3.zero;

            // Alignment
            float3 aliSteer = (aliCount > 0)
                ? CalculateSteer(aliVelSum, velocity, steeringForce, maxSteeringForce, maxSteeringForceSq)
                : float3.zero;

            // Cohesion
            float3 cohSteer = float3.zero;
            if (cohCount > 0)
            {
                cohPosSum = cohPosSum / (float)cohCount; // average position
                cohSteer = cohPosSum - position; // towards average position
                cohSteer = CalculateSteer(cohSteer, velocity, steeringForce, maxSteeringForce, maxSteeringForceSq);
            }

            // Calculate force
            float3 force = float3.zero;
            force += sepSteer * _boidArgs.separateWeight;
            force += aliSteer * _boidArgs.alignmentWeight;
            force += cohSteer * _boidArgs.cohesionWeight;

            _boidForce[index] = force * _boidArgs.forceMultiplier;
        }
    }

    [BurstCompile(FloatPrecision.Low, FloatMode.Fast)]
    public struct ApplyForceToBoidDataJob : IJobFor
    {
        [NativeDisableParallelForRestriction]
        private NativeArray<BoidData> _boidData;
        [ReadOnly] private NativeArray<float3> _boidForce;
        private BoidArgs _boidArgs;
        private ExternalForce _externalForce;
        private float3 _centerPosition;
        private float _deltaTime;

        public ApplyForceToBoidDataJob(NativeArray<BoidData> boidData, NativeArray<float3> boidForce, BoidArgs boidArgs, ExternalForce externalForce, Vector3 centerPosition, float deltaTime)
        {
            this._boidData = boidData;
            this._boidForce = boidForce;
            this._boidArgs = boidArgs;
            this._externalForce = externalForce;
            this._centerPosition = centerPosition;
            this._deltaTime = deltaTime;
        }

        public float3 CalculateSteer(float3 desired, float3 velocity, float steerForce, float maxForce, float maxForceSq)
        {
            float lengthSq = math.lengthsq(desired);
            if (lengthSq < 1e-5f)
                return float3.zero;

            // Normalize and scale in one step
            float3 desiredDir = desired * (math.rsqrt(lengthSq) * steerForce);

            // Calculate steering force
            float3 steer = desiredDir - velocity;

            LimitVector(ref steer, maxForce, maxForceSq);
            return steer;
        }

        public void LimitVector(ref float3 vector, float maxLength, float maxLengthSq)
        {
            float lengthSq = math.lengthsq(vector);
            if (lengthSq > maxLengthSq)
            {
                vector = vector * (maxLength * math.rsqrt(lengthSq));
            }
        }

        public void Execute(int index)
        {
            BoidData boidData = _boidData[index];
            float3 position = boidData.position;
            float3 velocity = boidData.velocity;

            // Apply force to velocity
            float3 force = _boidForce[index];

            // Check and apply center force if needed
            float3 positionScaled = position * _boidArgs.toCenterScaleInv;
            float towardsCenterLengthSq = math.lengthsq(_centerPosition - positionScaled);
            if (towardsCenterLengthSq > _boidArgs.toCenterDistanceSq)
            {
                float3 towardsCenter = _centerPosition - position;
                float3 toCenterSteer = CalculateSteer(towardsCenter, velocity, _boidArgs.steerForceToCenter, _boidArgs.maxSteerForce, _boidArgs.maxSteerForceSq);
                force += toCenterSteer * _boidArgs.toCenterWeight;
            }

            // External - To Target Force
            if (_externalForce.toTargetForceWeight != 0f)
            {
                float3 toTargetDir = _externalForce.targetPositionForce - position;
                float3 toTargetSteer = CalculateSteer(toTargetDir, velocity, _boidArgs.steerForce, _boidArgs.maxSteerForce, _boidArgs.maxSteerForceSq);
                force += toTargetSteer * _externalForce.toTargetForceWeight;
            }

            // External - Uniform Force (steering)
            if (_externalForce.forceWeight != 0f)
            {
                float3 externalVelocitySteer = CalculateSteer(_externalForce.force, velocity, _boidArgs.steerForce, _boidArgs.maxSteerForce, _boidArgs.maxSteerForceSq);
                force += externalVelocitySteer * _externalForce.forceWeight;
            }

            velocity += force * _deltaTime; // Apply force to velocity

            // External - To Target Velocity
            if (_externalForce.toTargetVelocity != 0f)
            {
                float3 towardsTargetDir = math.normalize(_externalForce.targetPositionVelocity - position);
                velocity += towardsTargetDir * _externalForce.toTargetVelocity;
            }

            // External - Uniform Velocity
            if (_externalForce.velocityIntensity != 0f)
            {
                velocity += _externalForce.velocity * _externalForce.velocityIntensity;
            }

            LimitVector(ref velocity, _boidArgs.maxVelocity, _boidArgs.maxVelocitySq);

            position += velocity * _deltaTime;

            // Apply drag
            velocity *= 1.0f - (_boidArgs.dragIntensity * _deltaTime);

            boidData.position = position;
            boidData.velocity = velocity;
            _boidData[index] = boidData;
        }
    }
}