RealityView - organic, wave-like sphere morphing using LowLevelMesh and metal

MorphingSphereMetalView.swift

import SwiftUI
import RealityKit
import Metal

struct MorphingSphereMetalView: View {
    let rootEntity: Entity = Entity()
    let latitudeBands = 120
    let longitudeBands = 120
    var vertexCapacity: Int {
        return (latitudeBands + 1) * (longitudeBands + 1)
    }
    var indexCount: Int {
        return latitudeBands * longitudeBands * 6
    }

    @State var mesh: LowLevelMesh?
    @State var isMorphForward: Bool = true
    @State var morphAmount: Float = 0.0
    @State var morphPhase: Float = 0.0
    @State var timer: Timer?
    @State var frameDuration: TimeInterval = 0.0
    @State var lastUpdateTime = CACurrentMediaTime()
    @State var rotationAngles: SIMD3<Float> = [0, 0, 0]
    @State var time: Double = 0.0
    @State var lastRotationUpdateTime = CACurrentMediaTime()
    @State var radius: Float = 0.1
    
    let device: MTLDevice
    let commandQueue: MTLCommandQueue
    let computePipeline: MTLComputePipelineState
    
    init() {
        self.device = MTLCreateSystemDefaultDevice()!
        self.commandQueue = device.makeCommandQueue()!
        
        let library = device.makeDefaultLibrary()!
        let updateFunction = library.makeFunction(name: "updateMorphingSphere")!
        self.computePipeline = try! device.makeComputePipelineState(function: updateFunction)
    }
    
    var body: some View {
        GeometryReader3D { proxy in
            RealityView { content in
                let size = content.convert(proxy.frame(in: .local), from: .local, to: .scene).extents
                let radius = Float(0.5 * size.x)
                let mesh = try! createMesh()
                let modelComponent = try! getModelComponent(mesh: mesh)
                rootEntity.components.set(modelComponent)
                rootEntity.scale *= scalePreviewFactor
                content.add(rootEntity)
                self.radius = radius
                self.mesh = mesh
            }
            .onAppear { startTimer() }
            .onDisappear { stopTimer() }
        }
    }
    
    func startTimer() {
        timer = Timer.scheduledTimer(withTimeInterval: 1/120.0, repeats: true) { _ in
            let currentTime = CACurrentMediaTime()
            frameDuration = currentTime - lastUpdateTime
            lastUpdateTime = currentTime
            morphPhase = morphPhase + Float(frameDuration * 0.1)
            updateMesh()
            stepMorphAmount()
            stepRotationAndScale()
        }
    }
    
    func stopTimer() {
        timer?.invalidate()
        timer = nil
    }
    
    func stepMorphAmount() {
        if isMorphForward {
            morphAmount += 0.125
            if morphAmount >= 5.0 {
                isMorphForward = false
            }
        } else {
            morphAmount -= 0.01
            if morphAmount <= 0.0 {
                isMorphForward = true
            }
        }
    }
    
    func stepRotationAndScale() {
        let currentTime = CACurrentMediaTime()
        let frameDuration = currentTime - lastRotationUpdateTime
        
        rotationAngles.x += Float(frameDuration * 1.0)
        rotationAngles.y += Float(frameDuration * 0.0)
        rotationAngles.z += Float(frameDuration * 0.25)
        
        let rotationX = simd_quatf(angle: rotationAngles.x, axis: [1, 0, 0])
        let rotationY = simd_quatf(angle: rotationAngles.y, axis: [0, 1, 0])
        let rotationZ = simd_quatf(angle: rotationAngles.z, axis: [0, 0, 1])
        rootEntity.transform.rotation = rotationX * rotationY * rotationZ
        
        lastRotationUpdateTime = currentTime
    }
    
    func getModelComponent(mesh: LowLevelMesh) throws -> ModelComponent {
        let resource = try MeshResource(from: mesh)
        
        var material = PhysicallyBasedMaterial()
        material.baseColor.tint = .orange //.init(white: 0.05, alpha: 1.0)
        material.roughness.scale = 0.5
        material.metallic.scale = 1.0
        material.blending = .transparent(opacity: 1.0)
        material.faceCulling = .none

        return ModelComponent(mesh: resource, materials: [material])
    }
    
    func createMesh() throws -> LowLevelMesh {
        var desc = VertexData.descriptor
        desc.vertexCapacity = vertexCapacity
        desc.indexCapacity = indexCount
        
        let mesh = try LowLevelMesh(descriptor: desc)
        
        mesh.withUnsafeMutableIndices { rawIndices in
            let indices = rawIndices.bindMemory(to: UInt32.self)
            var index = 0
            
            for latNumber in 0..<latitudeBands {
                for longNumber in 0..<longitudeBands {
                    let first = (latNumber * (longitudeBands + 1)) + longNumber
                    let second = first + longitudeBands + 1
                    
                    indices[index] = UInt32(first)
                    indices[index + 1] = UInt32(second)
                    indices[index + 2] = UInt32(first + 1)
                    
                    indices[index + 3] = UInt32(second)
                    indices[index + 4] = UInt32(second + 1)
                    indices[index + 5] = UInt32(first + 1)
                    
                    index += 6
                }
            }
        }
        
        return mesh
    }
    
    func updateMesh() {
        guard let mesh = mesh,
              let commandBuffer = commandQueue.makeCommandBuffer(),
              let computeEncoder = commandBuffer.makeComputeCommandEncoder() else { return }
        
        let vertexBuffer = mesh.replace(bufferIndex: 0, using: commandBuffer)
        
        computeEncoder.setComputePipelineState(computePipeline)
        computeEncoder.setBuffer(vertexBuffer, offset: 0, index: 0)
        
        var params = MorphingSphereParams(
            latitudeBands: Int32(latitudeBands),
            longitudeBands: Int32(longitudeBands),
            radius: radius,
            morphAmount: morphAmount,
            morphPhase: morphPhase
        )
        computeEncoder.setBytes(&params, length: MemoryLayout<MorphingSphereParams>.size, index: 1)
        
        let threadsPerGrid = MTLSize(width: vertexCapacity, height: 1, depth: 1)
        let threadsPerThreadgroup = MTLSize(width: 64, height: 1, depth: 1)
        computeEncoder.dispatchThreads(threadsPerGrid, threadsPerThreadgroup: threadsPerThreadgroup)
        
        computeEncoder.endEncoding()
        commandBuffer.commit()
        
        let meshBounds = BoundingBox(min: [-radius, -radius, -radius], max: [radius, radius, radius])
        mesh.parts.replaceAll([
            LowLevelMesh.Part(
                indexCount: indexCount,
                topology: .triangle,
                bounds: meshBounds
            )
        ])
    }
    
    struct VertexData {
        var position: SIMD3<Float> = .zero
        var normal: SIMD3<Float> = .zero
        var uv: SIMD2<Float> = .zero
        
        static var vertexAttributes: [LowLevelMesh.Attribute] = [
            .init(semantic: .position, format: .float3, offset: MemoryLayout<Self>.offset(of: \.position)!),
            .init(semantic: .normal, format: .float3, offset: MemoryLayout<Self>.offset(of: \.normal)!),
            .init(semantic: .uv0, format: .float2, offset: MemoryLayout<Self>.offset(of: \.uv)!)
        ]

        static var vertexLayouts: [LowLevelMesh.Layout] = [
            .init(bufferIndex: 0, bufferStride: MemoryLayout<Self>.stride)
        ]

        static var descriptor: LowLevelMesh.Descriptor {
            var desc = LowLevelMesh.Descriptor()
            desc.vertexAttributes = VertexData.vertexAttributes
            desc.vertexLayouts = VertexData.vertexLayouts
            desc.indexType = .uint32
            return desc
        }
    }
    
    struct MorphingSphereParams {
        var latitudeBands: Int32
        var longitudeBands: Int32
        var radius: Float
        var morphAmount: Float
        var morphPhase: Float
    }
}

#Preview {
    MorphingSphereMetalView()
}

MorphSphere.metal

#include <metal_stdlib>
using namespace metal;

struct VertexData {
    float3 position;
    float3 normal;
    float2 uv;
};

struct MorphingSphereParams {
    int32_t latitudeBands;
    int32_t longitudeBands;
    float radius;
    float morphAmount;
    float morphPhase;
};

// Smooth noise function
float smoothNoise(float2 st) {
    float2 i = floor(st);
    float2 f = fract(st);
    
    // Generate pseudo-random values for the four corners of the containing unit square
    float a = sin(dot(i, float2(12.9898, 78.233)) * 43758.5453);
    float b = sin(dot(i + float2(1.0, 0.0), float2(12.9898, 78.233)) * 43758.5453);
    float c = sin(dot(i + float2(0.0, 1.0), float2(12.9898, 78.233)) * 43758.5453);
    float d = sin(dot(i + float2(1.0, 1.0), float2(12.9898, 78.233)) * 43758.5453);

    // Compute smooth interpolation weights
    float2 u = f * f * (3.0 - 2.0 * f);

    // Interpolate between corner values
    return mix(a, b, u.x) +
            (c - a)* u.y * (1.0 - u.x) +
            (d - b) * u.x * u.y;
}

// Fractal Brownian Motion
float fbm(float2 st) {
    float value = 0.0;
    float amplitude = 0.5;
    float frequency = 1.0;
    // Reduce octaves for smoother effect
    for (int i = 0; i < 4; ++i) {
        value += amplitude * smoothNoise(st * frequency);
        st = st * 2.0 + float2(3.14, 2.71);
        amplitude *= 0.5;
        frequency *= 2.0;
    }
    return value;
}

kernel void updateMorphingSphere(device VertexData* vertices [[buffer(0)]],
                                 constant MorphingSphereParams& params [[buffer(1)]],
                                 uint id [[thread_position_in_grid]])
{
    int x = id % (params.longitudeBands + 1);
    int y = id / (params.longitudeBands + 1);
    
    if (x > params.longitudeBands || y > params.latitudeBands) return;
    
    float lat = float(y) / float(params.latitudeBands);
    float lon = float(x) / float(params.longitudeBands);
    float theta = (1.0 - lat) * M_PI_F;
    float phi = lon * 2 * M_PI_F;
    
    float sinTheta = sin(theta);
    float cosTheta = cos(theta);
    float sinPhi = sin(phi);
    float cosPhi = cos(phi);
    
    float3 basePosition = float3(cosPhi * sinTheta, cosTheta, sinPhi * sinTheta);
    
    // Use a very low frequency for the noise
    float2 noiseCoord = float2(theta, phi) + params.morphPhase * 10000;
    float noiseValue = fbm(noiseCoord) * 2.0 - 1.0;
    
    // Apply a sine wave to create a more organic, wave-like motion
    float waveEffect = sin(params.morphPhase + theta * 6.0 + phi * 2.0) * 0.5 + 0.5;
    
    // Combine noise and wave effect
    float combinedEffect = mix(noiseValue, waveEffect, 0.5);
    
    // Reduce the overall morphing amount for subtler effect
    float morphScale = 0.15 * params.morphAmount;
    
    float3 offset = basePosition * (combinedEffect * morphScale);
    
    float3 position = (basePosition + offset) * params.radius;
    float3 normal = normalize(position);
    
    vertices[id].position = position;
    vertices[id].normal = normal;
    vertices[id].uv = float2(lon, lat);
}