Particle Physics Simulation app for the Apple Vision Pro (visionOS 2.0)

ParticlePhysicsApp.swift

//
//  ParticlePhysicsApp.swift
//  ParticlePhysicsApp
//
//  Created by Fabio Dela Antonio on 15/09/2024.
//

import SwiftUI
import RealityKit
import ARKit
import OSLog

@main
struct ParticlePhysicsApp: App {
    @State private var model = EntityModel()
    @Environment(\.dismissImmersiveSpace) var dismissImmersiveSpace
    @Environment(\.openWindow) var openWindow

    var body: some SwiftUI.Scene {
        ImmersiveSpace {
            RealityView { content in
                content.add(await model.setupContentEntity())
            }
            .task {
                do {
                    if model.dataProvidersAreSupported && model.isReadyToRun {
                        try await model.session.run([model.handTracking])
                    } else {
                        await dismissImmersiveSpace()
                    }
                } catch {
                    logger.error("Failed to start session: \(error)")
                    await dismissImmersiveSpace()
                    openWindow(id: "error")
                }
            }
            .task {
                await model.processHandUpdates()
            }
            .task {
                await model.monitorSessionEvents()
            }
            .onChange(of: model.errorState) {
                openWindow(id: "error")
            }
        }
        .upperLimbVisibility(.hidden)
        .persistentSystemOverlays(.hidden)

        WindowGroup(id: "error") {
            Text("An error occurred; check the app's logs for details.")
        }
    }

    init() {
        ParticleComponent.registerComponent()
        ForcesContainerComponent.registerComponent()
        ParticlePhysicsSystem.registerSystem()
    }
}

typealias Vector3 = simd_float3

extension Vector3 {

    var magnitudeSquared: Float {
        x*x + y*y + z*z
    }

    var norm: Float {
        sqrtf(magnitudeSquared)
    }

    var normalized: Vector3 {
        normalize(self)
    }

    static func * (_ lhs: Vector3, _ rhs: Float) -> Vector3 {
        .init(x: lhs.x * rhs, y: lhs.y * rhs, z: lhs.z * rhs)
    }
}

typealias Vector4 = simd_float4

extension Vector4 {

    var xyz: Vector3 {
        .init(x: x, y: y, z: z)
    }
}

// MARK: - Particle Physics

protocol ParticleProtocol {
    var mass: Float { get }
    var velocity: Vector3 { get }
    var position: Vector3 { get }
}

protocol ForceProtocol {
    func result(with particle: ParticleProtocol) -> Vector3
}

final class GravityForce: ForceProtocol {
    func result(with particle: ParticleProtocol) -> Vector3 {
        let gravity: Float = 9.81 * 0.05 // m/s^2
        let gravityDirection = Vector3(0, -gravity, 0)
        return gravityDirection * particle.mass
    }
}

final class AttractionForce: ForceProtocol {
    var position: Vector3

    init(position: Vector3) {
        self.position = position
    }

    func result(with particle: ParticleProtocol) -> Vector3 {
        let bigG: Float = 6.674e-11 // N*(m/kg)^2
        let bigMass: Float = 1e10
        let intensity = (bigG * bigMass * particle.mass)/(position - particle.position).magnitudeSquared
        return (position - particle.position) * intensity
    }
}

final class DragForce: ForceProtocol {
    func result(with particle: ParticleProtocol) -> Vector3 {
        let dragCoefficient: Float = -0.4
        return particle.velocity * dragCoefficient
    }
}

final class ParticleComponent: Component, ParticleProtocol {
    let mass: Float
    private(set) var velocity: Vector3
    private(set) var position: Vector3

    init(mass: Float, velocity: Vector3, position: Vector3) {
        self.mass = mass
        self.velocity = velocity
        self.position = position
    }

    func update(
        with deltaTime: TimeInterval,
        forces: [ForceProtocol],
        bounds: (Vector3, Vector3),
        maxVelocity: Float
    ) {
        position += velocity * Float(deltaTime)
        let (minPosition, maxPosition) = bounds

        if position.x > maxPosition.x {
            position.x = maxPosition.x
            velocity.x = 0
        }

        if position.y > maxPosition.y {
            position.y = maxPosition.y
            velocity.y = 0
        }

        if position.z > maxPosition.z {
            position.z = maxPosition.z
            velocity.z = 0
        }

        if position.x < minPosition.x {
            position.x = minPosition.x
            velocity.x = 0
        }

        if position.y < minPosition.y {
            position.y = minPosition.y
            velocity.y = 0
        }

        if position.z < minPosition.z {
            position.z = minPosition.z
            velocity.z = 0
        }

        var resultingForce = Vector3.zero

        for force in forces {
            resultingForce += force.result(with: self)
        }

        velocity += resultingForce * (Float(deltaTime)/mass)

        if abs(velocity.x) > maxVelocity {
            velocity.x = (velocity.x > 0) ? maxVelocity:-maxVelocity
        }

        if abs(velocity.y) > maxVelocity {
            velocity.y = (velocity.y > 0) ? maxVelocity:-maxVelocity
        }

        if abs(velocity.z) > maxVelocity {
            velocity.z = (velocity.z > 0) ? maxVelocity:-maxVelocity
        }
    }
}

final class ForcesContainerComponent: Component {
    var forces: [String: ForceProtocol]

    init(forces: [String: ForceProtocol]) {
        self.forces = forces
    }
}

enum PhysicsConstants {
    static let minPosition = Vector3(-0.5, -0.5, -0.5) // m
    static let maxPosition = Vector3(0.5, 0.5, 0.5) // m
    static let maxDimensionalVelocity: Float = 1.0 // m/s
    static let minMass: Float = 1.0 // kg
    static let maxMass: Float = 10.0 // kg
}

enum EntityBuilder {

    static let particleRadius: Float = 0.005

    static func buildParticle() -> Entity {
        let hue = Float.random(in: 0.0...1.0)
        let mass = (hue * (PhysicsConstants.maxMass - PhysicsConstants.minMass)) + PhysicsConstants.minMass

        let position = Vector3(
            Float.random(in: PhysicsConstants.minPosition.x...PhysicsConstants.maxPosition.x),
            Float.random(in: PhysicsConstants.minPosition.y...PhysicsConstants.maxPosition.y),
            Float.random(in: PhysicsConstants.minPosition.z...PhysicsConstants.maxPosition.z)
        )

        let velocity = Vector3(
            Float.random(in: -PhysicsConstants.maxDimensionalVelocity...PhysicsConstants.maxDimensionalVelocity),
            Float.random(in: -PhysicsConstants.maxDimensionalVelocity...PhysicsConstants.maxDimensionalVelocity),
            Float.random(in: -PhysicsConstants.maxDimensionalVelocity...PhysicsConstants.maxDimensionalVelocity)
        )

        let material = SimpleMaterial(
            color: .init(hue: CGFloat(hue), saturation: 0.8, brightness: 0.8, alpha: 1.0),
            roughness: .float(0.5),
            isMetallic: true
        )

        let modelEntity = ModelEntity(mesh: .generateSphere(radius: particleRadius), materials: [material])
        modelEntity.components.set(ParticleComponent(mass: mass, velocity: velocity, position: position))
        modelEntity.position = position
        return modelEntity
    }

    static let lineDiameter: Float = 0.01

    static func makeLine(from positionA: SIMD3<Float>, to positionB: SIMD3<Float>, reference: Entity) -> Entity {
        let vector = SIMD3<Float>(positionA.x - positionB.x, positionA.y - positionB.y, positionA.z - positionB.z)
        let distance = sqrt(vector.x * vector.x + vector.y * vector.y + vector.z * vector.z)

        let midPosition = SIMD3<Float>(
            (positionA.x + positionB.x)/2.0,
            (positionA.y + positionB.y)/2.0,
            (positionA.z + positionB.z)/2.0
        )

        let entity = reference.clone(recursive: false)
        entity.position = midPosition
        entity.look(at: positionB, from: midPosition, relativeTo: nil)

        entity.scale = .init(x: lineDiameter, y: lineDiameter, z: (lineDiameter + distance))
        return entity
    }

    static func buildBox() -> Entity {
        let referenceEntity = ModelEntity(
            mesh: .generateBox(size: 1.0),
            materials: [SimpleMaterial(color: .init(white: 1.0, alpha: 0.5), isMetallic: false)]
        )

        let parentEntity = Entity()

        let minPosition = Vector3(
            PhysicsConstants.minPosition.x - particleRadius,
            PhysicsConstants.minPosition.y - particleRadius,
            PhysicsConstants.minPosition.z - particleRadius
        )

        let maxPosition = Vector3(
            PhysicsConstants.maxPosition.x + particleRadius,
            PhysicsConstants.maxPosition.y + particleRadius,
            PhysicsConstants.maxPosition.z + particleRadius
        )

        parentEntity.addChild(
            makeLine(
                from: Vector3(minPosition.x + lineDiameter, minPosition.y, minPosition.z),
                to: Vector3(maxPosition.x - lineDiameter, minPosition.y, minPosition.z),
                reference: referenceEntity
            )
        )

        parentEntity.addChild(
            makeLine(
                from: Vector3(maxPosition.x, minPosition.y + lineDiameter, minPosition.z),
                to: Vector3(maxPosition.x, maxPosition.y - lineDiameter, minPosition.z),
                reference: referenceEntity
            )
        )

        parentEntity.addChild(
            makeLine(
                from: Vector3(maxPosition.x - lineDiameter, maxPosition.y, minPosition.z),
                to: Vector3(minPosition.x + lineDiameter, maxPosition.y, minPosition.z),
                reference: referenceEntity
            )
        )

        parentEntity.addChild(
            makeLine(
                from: Vector3(minPosition.x, maxPosition.y - lineDiameter, minPosition.z),
                to: Vector3(minPosition.x, minPosition.y + lineDiameter, minPosition.z),
                reference: referenceEntity
            )
        )

        parentEntity.addChild(
            makeLine(
                from: Vector3(minPosition.x + lineDiameter, minPosition.y, maxPosition.z),
                to: Vector3(maxPosition.x - lineDiameter, minPosition.y, maxPosition.z),
                reference: referenceEntity
            )
        )

        parentEntity.addChild(
            makeLine(
                from: Vector3(maxPosition.x, minPosition.y + lineDiameter, maxPosition.z),
                to: Vector3(maxPosition.x, maxPosition.y - lineDiameter, maxPosition.z),
                reference: referenceEntity
            )
        )

        parentEntity.addChild(
            makeLine(
                from: Vector3(maxPosition.x - lineDiameter, maxPosition.y, maxPosition.z),
                to: Vector3(minPosition.x + lineDiameter, maxPosition.y, maxPosition.z),
                reference: referenceEntity
            )
        )

        parentEntity.addChild(
            makeLine(
                from: Vector3(minPosition.x, maxPosition.y - lineDiameter, maxPosition.z),
                to: Vector3(minPosition.x, minPosition.y + lineDiameter, maxPosition.z),
                reference: referenceEntity
            )
        )

        parentEntity.addChild(
            makeLine(
                from: Vector3(minPosition.x, minPosition.y, minPosition.z),
                to: Vector3(minPosition.x, minPosition.y, maxPosition.z),
                reference: referenceEntity
            )
        )

        parentEntity.addChild(
            makeLine(
                from: Vector3(maxPosition.x, minPosition.y, minPosition.z),
                to: Vector3(maxPosition.x, minPosition.y, maxPosition.z),
                reference: referenceEntity
            )
        )

        parentEntity.addChild(
            makeLine(
                from: Vector3(maxPosition.x, maxPosition.y, minPosition.z),
                to: Vector3(maxPosition.x, maxPosition.y, maxPosition.z),
                reference: referenceEntity
            )
        )

        parentEntity.addChild(
            makeLine(
                from: Vector3(minPosition.x, maxPosition.y, minPosition.z),
                to: Vector3(minPosition.x, maxPosition.y, maxPosition.z),
                reference: referenceEntity
            )
        )

        return parentEntity
    }
}

final class ParticlePhysicsSystem: System {

    init(scene: RealityKit.Scene) {
    }

    func update(context: SceneUpdateContext) {
        let forceContainers = context.entities(matching: .init(where: .has(ForcesContainerComponent.self)), updatingSystemWhen: .rendering)
        let particles = context.entities(matching: .init(where: .has(ParticleComponent.self)), updatingSystemWhen: .rendering)

        let forces = Array(
            forceContainers
                .compactMap({
                    $0.components[ForcesContainerComponent.self]?.forces
                })
                .reduce([:], {
                    $0.merging($1, uniquingKeysWith: { $1 })
                })
                .values
        )

        for particle in particles {
            guard let particleComponent = particle.components[ParticleComponent.self] else { continue }
            particleComponent.update(
                with: context.deltaTime,
                forces: forces,
                bounds: (PhysicsConstants.minPosition, PhysicsConstants.maxPosition),
                maxVelocity: PhysicsConstants.maxDimensionalVelocity
            )

            particle.position = particleComponent.position
        }
    }
}

// MARK: - Entity Model

@MainActor
let logger = Logger(subsystem: "ParticlePhysics", category: "general")

@Observable
@MainActor
final class EntityModel {
    let session = ARKitSession()
    let handTracking = HandTrackingProvider()

    private(set) var contentEntity = Entity()
    private let forcesContainer = ForcesContainerComponent(forces: [
        "drag": DragForce(),
        "gravity": GravityForce()
    ])

    var errorState = false

    var dataProvidersAreSupported: Bool {
        HandTrackingProvider.isSupported
    }

    var isReadyToRun: Bool {
        handTracking.state == .initialized
    }

    init() {
    }

    private let particlesOffset = Vector3(0, 1.5, -1.0)

    func setupContentEntity() async -> Entity {
        let particlesEntity = Entity()
        contentEntity.addChild(particlesEntity)
        particlesEntity.position = particlesOffset

        for _ in 0 ..< 100 {
            particlesEntity.addChild(EntityBuilder.buildParticle())
        }

        particlesEntity.addChild(EntityBuilder.buildBox())

        let forcesEntity = Entity()
        forcesEntity.components.set(forcesContainer)
        contentEntity.addChild(forcesEntity)

        contentEntity.transform.translation = .zero
        return contentEntity
    }

    func processHandUpdates() async {
        for await update in handTracking.anchorUpdates {
            let handAnchor = update.anchor

            guard handAnchor.isTracked,
                  let indexFingerTip = handAnchor.handSkeleton?.joint(.indexFingerTip),
                  let thumbTip = handAnchor.handSkeleton?.joint(.thumbTip)
            else {
                continue
            }

            let indexFingerTipPosition = (handAnchor.originFromAnchorTransform * indexFingerTip.anchorFromJointTransform).columns.3
            let thumbTipPosition = (handAnchor.originFromAnchorTransform * thumbTip.anchorFromJointTransform).columns.3

            let distance = (indexFingerTipPosition.xyz - thumbTipPosition.xyz).norm
            let midPoint = (indexFingerTipPosition.xyz + thumbTipPosition.xyz) * 0.5

            let key = switch handAnchor.chirality {
            case .right:
                "right-hand"
            case .left:
                "left-hand"
            }

            if distance < 0.02 {
                let pointInParticlesCoordinates = midPoint - particlesOffset
                forcesContainer.forces[key] = AttractionForce(position: pointInParticlesCoordinates)
            } else {
                forcesContainer.forces.removeValue(forKey: key)
            }
        }
    }

    func monitorSessionEvents() async {
        for await event in session.events {
            switch event {
            case .authorizationChanged(type: _, status: let status):
                logger.info("Authorization changed to: \(status)")

                if status == .denied {
                    errorState = true
                }
            case .dataProviderStateChanged(dataProviders: let providers, newState: let state, error: let error):
                logger.info("Data provider changed: \(providers), \(state)")
                if let error {
                    logger.error("Data provider reached an error state: \(error)")
                    errorState = true
                }
            @unknown default:
                fatalError("Unhandled new event type \(event)")
            }
        }
    }
}