GitHub Gist for the stannatrane tripod test

FIRE.swift

//
//  FIRE.swift
//  MolecularRendererApp
//
//  Created by Philip Turner on 5/31/24.
//

import HDL
import Numerics

struct FIREMinimizationDescriptor {
  /// Optional. Indices of atoms to keep positionally constrained.
  var anchors: Set<UInt32>?
  
  /// Required. The tolerance for maximum force among the atoms.
  ///
  /// The default value is 10 pN. This value is sufficient for accurate
  /// structures. For accurate energies (when comparing two different
  /// structures), choose a value of 0.3 pN.
  var forceTolerance: Float = 10
  
  /// Required. The mass of each atom (in yoctograms).
  var masses: [Float]?
  
  /// Required. The position of each atom's nucleus.
  var positions: [SIMD3<Float>]?
  
  /// Required. The minimum timestep.
  var ΔtMin: Float = 0.25e-3
  
  /// Required. The initial timestep.
  var ΔtStart: Float = 0.001
  
  /// Required. The maximum timestep.
  var ΔtMax: Float = 0.010
}

/// A data structure for integration during an energy minimization.
///
/// This is not a full molecular dynamics engine. The current prototype is only
/// scoped to something capable of ONIOM simulation.
struct FIREMinimization {
  // Integrator settings.
  let anchors: Set<UInt32>
  let masses: [Float]
  let ΔtMin: Float
  let ΔtMax: Float
  let forceTolerance: Float
  
  // FIRE timestep trackers.
  private(set) var Δt: Float
  private(set) var NP0: Int
  
  // Dynamical variables.
  private(set) var time: Double
  var positions: [SIMD3<Float>] // must be mutable to enforce constraints
  var velocities: [SIMD3<Float>] // must be mutable to enforce constraints
  
  // Cached state.
  private(set) var oldTime: Double
  private(set) var oldPositions: [SIMD3<Float>]
  
  init(descriptor: FIREMinimizationDescriptor) {
    guard let masses = descriptor.masses,
          let positions = descriptor.positions else {
      fatalError("Descriptor was incomplete.")
    }
    guard masses.count == positions.count else {
      fatalError("Size of masses did not match size of positions.")
    }
    
    // Initialize the constraints.
    if let anchors = descriptor.anchors {
      self.anchors = anchors
    } else {
      self.anchors = []
    }
    self.masses = masses
    ΔtMin = descriptor.ΔtMin
    ΔtMax = descriptor.ΔtMax
    forceTolerance = descriptor.forceTolerance
    
    // Initialize the timestep trackers.
    Δt = descriptor.ΔtStart
    NP0 = 0
    
    // Initialize the dynamical variables.
    time = 0
    self.positions = positions
    velocities = [SIMD3<Float>](repeating: .zero, count: positions.count)
    
    // Initialize the cached state.
    oldTime = 0
    oldPositions = positions
  }
  
  // If the minimization has converged, return `true`. Otherwise, integrate for
  // one timestep.
  mutating func step(forces: [SIMD3<Float>]) -> Bool {
    guard forces.count == positions.count else {
      fatalError("Size of forces did not match size of positions.")
    }
    
    // Find the power (P) and maximum force.
    var P: Double = .zero
    var maxForce: Float = .zero
    for atomID in positions.indices {
      if anchors.contains(UInt32(atomID)) {
        continue
      }
      let force = forces[atomID]
      let velocity = velocities[atomID]
      P += Double((force * velocity).sum())
      
      let forceMagnitude = (force * force).sum().squareRoot()
      maxForce = max(maxForce, forceMagnitude)
    }
    
    // Return early if converged.
    if maxForce < forceTolerance {
      return true
    }
    
    // Prepare for integration.
    updateTimestep(P: P)
    
    // Integrate.
    integrate(forces: forces)
    oldTime = time + Double(0.5 * Δt)
    time = time + Double(Δt)
    
    // Return that you haven't converged.
    return false
  }
  
  // Adjust the timestep, according to the value of P.
  private mutating func updateTimestep(P: Double) {
    if time > 0, P < 0 {
      time = oldTime
      positions = oldPositions
      velocities = Array(repeating: .zero, count: positions.count)
      
      NP0 = 0
      Δt = max(0.5 * Δt, ΔtMin)
    } else {
      NP0 += 1
      if NP0 > 5 {
        Δt = min(1.1 * Δt, ΔtMax)
      }
    }
  }
  
  // Compute the force scale for the entire system.
  //
  // WARNING: This must be done after the velocities are reset when P < 0.
  private func createForceScale(forces: [SIMD3<Float>]) -> Float {
    var vAccumulator: Float = .zero
    var fAccumulator: Float = .zero
    for atomID in positions.indices {
      if anchors.contains(UInt32(atomID)) {
        continue
      }
      let velocity = velocities[atomID]
      let force = forces[atomID]
      vAccumulator += (velocity * velocity).sum()
      fAccumulator += (force * force).sum()
    }
    
    let vNorm = vAccumulator.squareRoot()
    let fNorm = fAccumulator.squareRoot()
    var forceScale = vNorm / fNorm
    if forceScale.isNaN || forceScale.isInfinite {
      forceScale = .zero
    }
    return forceScale
  }
  
  // Perform an integration step.
  private mutating func integrate(forces: [SIMD3<Float>]) {
    // Find the force scale for redirecting atom velocities.
    let forceScale = createForceScale(forces: forces)
    
    // Iterate over the atoms.
    for atomID in positions.indices {
      var halfwayPosition: SIMD3<Float>
      var finalPosition: SIMD3<Float>
      var finalVelocity: SIMD3<Float>
      defer {
        oldPositions[atomID] = halfwayPosition
        positions[atomID] = finalPosition
        velocities[atomID] = finalVelocity
      }
      
      let initialPosition = positions[atomID]
      let initialVelocity = velocities[atomID]
      if anchors.contains(UInt32(atomID)) {
        // Prevent the anchors from moving.
        halfwayPosition = initialPosition
        finalPosition = initialPosition
        finalVelocity = .zero
      } else {
        let force = forces[atomID]
        let mass = masses[atomID]
        var velocity = initialVelocity
        
        // Semi-implicit Euler integration.
        // - This could be considered a valid form of molecular dynamics, just
        //   with a forcing term that accelerates the descent into the global
        //   minimum.
        let α: Float = 0.25
        velocity += Δt * force / mass
        velocity = (1 - α) * velocity + α * force * forceScale
        
        // Accelerated bias correction.
        // - This should be applied whenever the system's velocities are
        //   suddenly reset to zero. Ideally, such a situation will never
        //   happen in a molecular dynamics simulation.
        // - The correction can shrink the number of iterations by a factor of
        //   2x.
        if NP0 > 0 {
          var biasCorrection = 1 - α
          biasCorrection = Float.pow(biasCorrection, Float(NP0))
          biasCorrection = 1 / (1 - biasCorrection)
          velocity *= biasCorrection
        }
        
        // Clamp the velocity to 4000 m/s.
        let speed = (velocity * velocity).sum().squareRoot()
        if speed > 4.0 {
          velocity *= 4.0 / speed
        }
        finalVelocity = velocity
        
        // Integrate the position.
        halfwayPosition = initialPosition + 0.5 * Δt * velocity
        finalPosition = initialPosition + Δt * velocity
      }
    }
  }
}

Format.swift

// Utility for logging quantities to the console.
struct Format {
  static func pad(_ x: String, to size: Int) -> String {
    var output = x
    while output.count < size {
      output = " " + output
    }
    return output
  }
  static func time<T: BinaryFloatingPoint>(_ x: T) -> String {
    let xInFs = Float(x) * 1e3
    var repr = String(format: "%.2f", xInFs) + " fs"
    repr = pad(repr, to: 9)
    return repr
  }
  static func energy(_ x: Double) -> String {
    var repr = String(format: "%.2f", x / 160.218) + " eV"
    repr = pad(repr, to: 13)
    return repr
  }
  static func force(_ x: Float) -> String {
    var repr = String(format: "%.2f", x) + " pN"
    repr = pad(repr, to: 13)
    return repr
  }
  static func distance(_ x: Float) -> String {
    var repr = String(format: "%.2f", x) + " nm"
    repr = pad(repr, to: 9)
    return repr
  }
}

main.swift

import Foundation
import HDL
import MolecularRenderer
import QuaternionModule
import xTB

xTB_Environment.verbosity = .muted

// MARK: - Compile Structure

// Passivate only the carbon atoms.
func passivate(topology: inout Topology) {
  func createHydrogen(
    atomID: UInt32,
    orbital: SIMD3<Float>
  ) -> Atom {
    let atom = topology.atoms[Int(atomID)]
    
    var bondLength = atom.element.covalentRadius
    bondLength += Element.hydrogen.covalentRadius
    
    let position = atom.position + bondLength * orbital
    return Atom(position: position, element: .hydrogen)
  }
  
  let orbitalLists = topology.nonbondingOrbitals()
  
  var insertedAtoms: [Atom] = []
  var insertedBonds: [SIMD2<UInt32>] = []
  for atomID in topology.atoms.indices {
    let atom = topology.atoms[atomID]
    guard atom.atomicNumber == 6 else {
      continue
    }
    
    let orbitalList = orbitalLists[atomID]
    for orbital in orbitalList {
      let hydrogen = createHydrogen(
        atomID: UInt32(atomID),
        orbital: orbital)
      let hydrogenID = topology.atoms.count + insertedAtoms.count
      insertedAtoms.append(hydrogen)
      
      let bond = SIMD2(
        UInt32(atomID),
        UInt32(hydrogenID))
      insertedBonds.append(bond)
    }
  }
  topology.atoms += insertedAtoms
  topology.bonds += insertedBonds
}

func createTripod(
  isAzastannatrane: Bool
) -> Topology {
  var topology = Topology()
  topology.atoms += [
    Atom(position: SIMD3(0.00, 0.00, 0.00), element: .tin),
    Atom(position: SIMD3(0.00, -0.26, -0.00), element: .nitrogen),
  ]
  
  for legID in 0..<3 {
    let baseAtomID = topology.atoms.count
    var insertedAtoms: [Atom] = []
    var insertedBonds: [SIMD2<UInt32>] = []
    
    // Isolate the temporary variables for this part in a contained scope.
    do {
      let carbon0 = Atom(
        position: SIMD3(0.15, -0.28, 0.00), element: .carbon)
      insertedAtoms.append(carbon0)
      insertedBonds.append(
        SIMD2(UInt32(1), UInt32(baseAtomID + 0)))
      
      let carbon1 = Atom(
        position: SIMD3(0.23, -0.15, -0.05), element: .carbon)
      insertedAtoms.append(carbon1)
      insertedBonds.append(
        SIMD2(UInt32(baseAtomID + 0), UInt32(baseAtomID + 1)))
      
      let nitrogen2 = Atom(
        position: SIMD3(0.23, -0.00, 0.00),
        element: isAzastannatrane ? .nitrogen : .carbon)
      insertedAtoms.append(nitrogen2)
      insertedBonds.append(
        SIMD2(UInt32(baseAtomID + 1), UInt32(baseAtomID + 2)))
      insertedBonds.append(
        SIMD2(UInt32(baseAtomID + 2), UInt32(0)))
      
      let carbon3 = Atom(
        position: SIMD3(0.38, -0.20, -0.05), element: .carbon)
      insertedAtoms.append(carbon3)
      insertedBonds.append(
        SIMD2(UInt32(baseAtomID + 1), UInt32(baseAtomID + 3)))
      
      let sulfur4 = Atom(
        position: SIMD3(0.45, -0.36, -0.05), element: .sulfur)
      insertedAtoms.append(sulfur4)
      insertedBonds.append(
        SIMD2(UInt32(baseAtomID + 3), UInt32(baseAtomID + 4)))
      
      let hydrogen5 = Atom(
        position: SIMD3(0.57, -0.36, -0.05), element: .hydrogen)
      insertedAtoms.append(hydrogen5)
      insertedBonds.append(
        SIMD2(UInt32(baseAtomID + 4), UInt32(baseAtomID + 5)))
    }
    
    if isAzastannatrane {
      let carbon6 = Atom(
        position: SIMD3(0.32, 0.08, -0.05), element: .carbon)
      insertedAtoms.append(carbon6)
      insertedBonds.append(
        SIMD2(UInt32(baseAtomID + 2), UInt32(baseAtomID + 6)))
      
      let hydrogen7 = Atom(
        position: SIMD3(0.32, 0.20, -0.05), element: .hydrogen)
      insertedAtoms.append(hydrogen7)
      insertedBonds.append(
        SIMD2(UInt32(baseAtomID + 6), UInt32(baseAtomID + 7)))
    }
    
    // Apply the rotation transform to all atoms, just before inserting.
    let angleDegrees = Float(legID) * 120 - 30
    let rotation = Quaternion<Float>(
      angle: Float.pi / 180 * angleDegrees,
      axis: SIMD3(0, 1, 0))
    for relativeAtomID in insertedAtoms.indices {
      var atom = insertedAtoms[relativeAtomID]
      atom.position = rotation.act(on: atom.position)
      insertedAtoms[relativeAtomID] = atom
    }
    topology.atoms += insertedAtoms
    topology.bonds += insertedBonds
  }
  
  // Don't forget the feedstock.
  topology.atoms += [
    Atom(position: SIMD3(0.00, 0.19, 0.00), element: .hydrogen),
  ]
  
  passivate(topology: &topology)
  return topology
}

func createCarbatranePositions() -> [Atom] {
  let topology = createTripod(isAzastannatrane: false)
  let trajectory = loadCachedTrajectory(tripod: topology)
  guard trajectory.count > 0 else {
    fatalError("No starting structure to render.")
  }
  return trajectory.last!
}

func createAzatranePositions() -> [Atom] {
  let topology = createTripod(isAzastannatrane: true)
  let trajectory = loadCachedTrajectory(tripod: topology)
  guard trajectory.count > 0 else {
    fatalError("No starting structure to render.")
  }
  return trajectory.last!
}

let carbatranePositions = createCarbatranePositions()
let azatranePositions = createAzatranePositions()

// MARK: - Launch Application

@MainActor
func createApplication() -> Application {
  // Set up the device.
  var deviceDesc = DeviceDescriptor()
  deviceDesc.deviceID = Device.fastestDeviceID
  let device = Device(descriptor: deviceDesc)
  
  // Set up the display.
  var displayDesc = DisplayDescriptor()
  displayDesc.device = device
  displayDesc.frameBufferSize = SIMD2<Int>(1440, 1440)
  let display = Display(descriptor: displayDesc)
  
  // Set up the application.
  var applicationDesc = ApplicationDescriptor()
  applicationDesc.device = device
  applicationDesc.display = display
  applicationDesc.upscaleFactor = 1
  
  applicationDesc.addressSpaceSize = 4_000_000
  applicationDesc.voxelAllocationSize = 500_000_000
  applicationDesc.worldDimension = 64
  let application = Application(descriptor: applicationDesc)
  
  return application
}
let application = createApplication()

// Set the atoms of the compiled structure(s) here.
//
// Render static image showing both side and top view.
var baseAtomID: Int = .zero
for structureID in 0..<4 {
  let angleDegrees = Float(90)
  let rotation = Quaternion<Float>(
    angle: Float.pi / 180 * angleDegrees,
    axis: SIMD3(1, 0, 0))
  
  var positions: [Atom]
  if structureID < 2 {
    positions = carbatranePositions
  } else {
    positions = azatranePositions
  }
  
  for atomID in positions.indices {
    var atom = positions[atomID]
   
    // Apply the rotation before messing with any translations.
    if structureID % 2 == 1 {
      atom.position = rotation.act(on: atom.position)
    }
    
    let separationDistanceX: Float = 1.5
    let separationDistanceY: Float = 1.5
    if structureID < 2 {
      atom.x -= separationDistanceX / 2
    } else {
      atom.x += separationDistanceX / 2
    }
    if structureID % 2 == 0 {
      atom.y -= separationDistanceY / 2
    } else {
      atom.y += separationDistanceY / 2
    }
    
    application.atoms[baseAtomID + atomID] = atom
  }
  
  baseAtomID += positions.count
}

// Set up the camera statically here.
application.camera.position = SIMD3(0, 0, 9)
application.camera.fovAngleVertical = Float.pi / 180 * 20
application.camera.secondaryRayCount = 64

do {
  let image = application.render()
  let frameBufferSize = application.display.frameBufferSize
  let pixelCount = frameBufferSize[0] * frameBufferSize[1]
  guard image.pixels.count == pixelCount else {
    fatalError("Invalid pixel buffer size.")
  }
  
  // Create the header.
  let header = """
  P6
  \(frameBufferSize[0]) \(frameBufferSize[1])
  255
  
  """
  let headerData = header.data(using: .utf8)!
  
  // Convert the pixels from FP16 to UInt8.
  var output: [UInt8] = []
  for pixel in image.pixels {
    let scaled = pixel * 255
    var rounded = scaled.rounded(.toNearestOrEven)
    rounded.replace(
      with: SIMD4<Float16>(repeating: 0),
      where: rounded .< 0)
    rounded.replace(
      with: SIMD4<Float16>(repeating: 255),
      where: rounded .> 255)
    
    let integerValue = SIMD4<UInt8>(rounded)
    output.append(integerValue[0])
    output.append(integerValue[1])
    output.append(integerValue[2])
  }
  let outputData = output.withUnsafeBufferPointer { bufferPointer in
    Data(buffer: bufferPointer)
  }
  
  // Write to the file. The forward slash usage is safe on Windows.
  let ppmData = headerData + outputData
  let packagePath = FileManager.default.currentDirectoryPath
  let filePath = "\(packagePath)/.build/image.ppm"
  let succeeded = FileManager.default.createFile(
    atPath: filePath,
    contents: ppmData)
  guard succeeded else {
    fatalError("Could not write to file.")
  }
}

Minimize.swift

import HDL
import var MM4.MM4YgPerAmu
import xTB

func runMinimization(tripod: Topology) -> [[SIMD4<Float>]] {
  var calculatorDesc = xTB_CalculatorDescriptor()
  calculatorDesc.atomicNumbers = tripod.atoms.map(\.atomicNumber)
  let calculator = xTB_Calculator(descriptor: calculatorDesc)
  
  var minimizationDesc = FIREMinimizationDescriptor()
  minimizationDesc.anchors = [UInt32(0)] // apex atom
  minimizationDesc.masses = tripod.atoms.map {
    if $0.atomicNumber == 1 {
      return Float(4.0 * 1.660539)
    } else {
      return Float(12.011 * 1.660539)
    }
  }
  minimizationDesc.positions = tripod.atoms.map(\.position)
  var minimization = FIREMinimization(descriptor: minimizationDesc)
  
  // Find the sulfur indices.
  var sulfurIDs: [UInt32] = []
  for atomID in tripod.atoms.indices {
    let atom = tripod.atoms[atomID]
    if atom.element == .sulfur {
      sulfurIDs.append(UInt32(atomID))
    }
  }
  
  guard sulfurIDs.count == 3 else {
    fatalError("Failed to locate all the sulfurs on the legs.")
  }
  
  var frames: [[Atom]] = []
  func createFrame() -> [Atom] {
    var output: [Atom] = []
    for atomID in tripod.atoms.indices {
      var atom = tripod.atoms[atomID]
      let position = minimization.positions[atomID]
      atom.position = position
      output.append(atom)
    }
    return output
  }
  
  print()
  for trialID in 0..<500 {
    frames.append(createFrame())
    calculator.molecule.positions = minimization.positions
    
    // Enforce the constraints on leg sulfurs.
    var forces = calculator.molecule.forces
    do {
      var forceAccumulator: Float = .zero
      for atomID in sulfurIDs {
        let force = forces[Int(atomID)]
        forceAccumulator += force.y
      }
      forceAccumulator /= 3
      for atomID in sulfurIDs {
        var force = forces[Int(atomID)]
        force.y = forceAccumulator
        forces[Int(atomID)] = force
      }
    }
    
    var maximumForce: Float = .zero
    for atomID in calculator.molecule.atomicNumbers.indices {
      if minimization.anchors.contains(UInt32(atomID)) {
        continue
      }
      let force = forces[atomID]
      let forceMagnitude = (force * force).sum().squareRoot()
      maximumForce = max(maximumForce, forceMagnitude)
    }
    
    print("time: \(Format.time(minimization.time))", terminator: " | ")
    print("energy: \(Format.energy(calculator.energy))", terminator: " | ")
    print("max force: \(Format.force(maximumForce))", terminator: " | ")
    
    let converged = minimization.step(forces: forces)
    if !converged {
      print("Δt: \(Format.time(minimization.Δt))", terminator: " | ")
    }
    print()
    
    // Enforce the constraints on leg sulfurs.
    do {
      var positions = minimization.positions
      var positionAccumulator: Float = .zero
      for atomID in sulfurIDs {
        let position = positions[Int(atomID)]
        positionAccumulator += position.y
      }
      positionAccumulator /= 3
      for atomID in sulfurIDs {
        var position = positions[Int(atomID)]
        position.y = positionAccumulator
        positions[Int(atomID)] = position
      }
      minimization.positions = positions
    }
    
    if converged {
      print("converged at trial \(trialID)")
      frames.append(createFrame())
      break
    } else if trialID == 499 {
      print("failed to converge!")
    }
  }
  
  return frames
}

Serialization.swift

//
//  Serialization.swift
//  MolecularRendererApp
//
//  Created by Philip Turner on 4/12/24.
//

import Foundation
import HDL

struct Serialization {
  
}

// MARK: - Atoms

extension Serialization {
  // Generates a unique hash from a list of atoms.
  static func hash(atoms: [SIMD4<Float>]) -> Data {
    var xorHash: SIMD4<UInt32> = .zero
    var rotateHash: SIMD4<UInt32> = .one
    var floatSum: SIMD4<UInt32> = .zero
    
    for atom in atoms {
      let storageCasted = unsafeBitCast(atom, to: SIMD4<UInt32>.self)
      
      xorHash ^= storageCasted
      xorHash = (xorHash &<< 3) | (xorHash &>> (32 - 3))
      
      rotateHash &*= storageCasted
      rotateHash &+= 1
      rotateHash = (rotateHash &<< 9) | (rotateHash &>> (32 - 9))
      
      var quantized = atom
      quantized *= 1024
      quantized.round(.toNearestOrEven)
      floatSum &+= unsafeBitCast(quantized, to: SIMD4<UInt32>.self)
    }
    
    var data = Data()
    func addVector(_ vector: SIMD4<UInt32>) {
      let castedVector = unsafeBitCast(vector, to: SIMD16<UInt8>.self)
      for laneID in 0..<16 {
        let byte = castedVector[laneID]
        data.append(byte)
      }
    }
    addVector(xorHash)
    addVector(rotateHash)
    addVector(floatSum)
    return data
  }
  
  // Compresses a list of atoms into binary data.
  //
  // Segments the volume into 64x64x64 nm blocks. The block-sparse volume must
  // not exceed 67 million cubic nanometers.
  //
  // The compression is a factor of ~2x from the information-theoretic limit.
  static func serialize(atoms: [SIMD4<Float>]) -> Data {
    // Allocate a dictionary for the chunk indices.
    var chunkIDs: [SIMD3<UInt16>: UInt8] = [:]
    var chunkCount: UInt8 = .zero
    
    // Allocate an array for the compressed atoms.
    var compressedAtoms: [SIMD4<UInt16>] = []
    
    // Loop over the atoms.
    for atom in atoms {
      // Quantize the position to ~1 picometer precision.
      var position = atom.position
      position *= 1024
      position.round(.toNearestOrEven)
      
      // Convert to a 32-bit integer.
      let integerValue = SIMD3<Int32>(position)
      let bitPattern = SIMD3<UInt32>(truncatingIfNeeded: integerValue)
      
      // Search for a matching sector.
      let upperBits = SIMD3<UInt16>(truncatingIfNeeded: bitPattern &>> 16)
      var chunkID: UInt8
      if let matchingID = chunkIDs[upperBits] {
        chunkID = matchingID
      } else if chunkCount < UInt8.max {
        chunkIDs[upperBits] = chunkCount
        chunkID = chunkCount
        chunkCount += 1
      } else {
        fatalError("Exceeded available vocabulary of 256 chunks.")
      }
      
      // Encode the lower bits.
      let lowerBits = SIMD3<UInt16>(truncatingIfNeeded: bitPattern)
      
      // Pack the atomic number and chunk ID into 16 bits.
      let idPair = SIMD2<UInt8>(atom.atomicNumber, chunkID)
      
      // Create a 4-lane vector with the data.
      let vector = SIMD4(lowerBits, unsafeBitCast(idPair, to: UInt16.self))
      compressedAtoms.append(vector)
    }
    
    // Allocate an array for the output data.
    var output = [SIMD4<UInt16>](repeating: .zero, count: 257)
    
    // Create a header for the serialized data.
    let atomCount = UInt64(atoms.count)
    output[0] = unsafeBitCast(atomCount, to: SIMD4<UInt16>.self)
    
    // Encode the chunks.
    for (key, value) in chunkIDs {
      let vector = SIMD4(key, 0)
      output[Int(1 + value)] = vector
    }
    
    // Encode the atoms.
    output += compressedAtoms
    
    // Return the data for further processing in the caller.
    let byteCount = output.count * MemoryLayout<SIMD4<UInt16>>.stride
    return Data(bytes: output, count: byteCount)
  }
  
  static func deserialize(atoms: Data) -> [SIMD4<Float>] {
    let wordCount = atoms.count / MemoryLayout<SIMD4<UInt16>>.stride
    var rawData = [SIMD4<UInt16>](repeating: .zero, count: wordCount)
    rawData.withUnsafeMutableBytes {
      let copiedCount = atoms.copyBytes(to: $0)
      guard copiedCount == atoms.count else {
        fatalError("Incorrect number of bytes was copied.")
      }
    }
    
    // Read the header.
    let header = rawData[0]
    let atomCount = unsafeBitCast(header, to: UInt64.self)
    guard 1 + 256 + atomCount == wordCount else {
      fatalError("Atom count was incorrect.")
    }
    
    // Loop over the atoms.
    var output: [SIMD4<Float>] = []
    for atomID in 0..<atomCount {
      let vector = rawData[Int(257 + atomID)]
      
      // Decode the lower bits.
      let lowerBits = SIMD3(vector[0], vector[1], vector[2])
      
      // Decode the atomic number and chunk ID.
      let idPair = unsafeBitCast(vector[3], to: SIMD2<UInt8>.self)
      
      // Locate the matching sector.
      let chunkID = idPair[1]
      let chunkVector = rawData[Int(1) + Int(chunkID)]
      let upperBits = unsafeBitCast(chunkVector, to: SIMD3<UInt16>.self)
      
      // Convert to a 32-bit integer.
      var bitPattern: SIMD3<UInt32> = .zero
      bitPattern |= SIMD3(truncatingIfNeeded: lowerBits)
      bitPattern |= SIMD3(truncatingIfNeeded: upperBits) &<< 16
      let integerValue = SIMD3<Int32>(truncatingIfNeeded: bitPattern)
      
      // Dequantize the position, from picometers to nanometers.
      var position = SIMD3<Float>(integerValue)
      position /= 1024
      
      // Decode the atomic number and convert to an entity.
      let atomicNumber = UInt8(idPair[0])
      let atom = SIMD4(position, Float(atomicNumber))
      output.append(atom)
    }
    
    return output
  }
}

// MARK: - Bonds

extension Serialization {
  // Compresses a bond topology into binary data.
  //
  // The typical compression ratio is ~2.0x. This is close to the
  // information-theoretic limit of ~3.5x.
  static func serialize(bonds: [SIMD2<UInt32>]) -> Data {
    var lastAtomID: UInt32 = .zero
    
    // Sort the bonds into compressable and non-compressable groups.
    var compressedBonds: [SIMD2<UInt16>] = []
    var decompressedBonds: [SIMD2<UInt32>] = []
    for bond in bonds {
      // We rely on there always being an 8-compressed atom nearby.
      guard bond[0] >= lastAtomID,
            bond[0] - lastAtomID < 65536 else {
        print(compressedBonds.count, decompressedBonds.count, bonds.count)
        fatalError("Bond could not be compressed.")
      }
      guard bond[1] > bond[0] else {
        fatalError("Bond was not sorted.")
      }
      
      let difference = bond[1] - bond[0]
      if difference < 65536 {
        let startDelta = UInt16(bond[0] - lastAtomID)
        let lengthDelta = UInt16(bond[1] - bond[0])
        let compressedBond = SIMD2(startDelta, lengthDelta)
        compressedBonds.append(compressedBond)
        
        // Update the atom cursor.
        lastAtomID = bond[0]
      } else {
        // Do not update the atom cursor.
        decompressedBonds.append(bond)
      }
    }
    
    // Allocate an array for the raw data.
    var rawData: [SIMD2<UInt32>] = []
    
    // Write to the header.
    var header: SIMD2<UInt32> = .zero
    header[0] = UInt32(compressedBonds.count)
    header[1] = UInt32(decompressedBonds.count)
    rawData.append(header)
    
    // Pad the compressable bonds to a multiple of two.
    while compressedBonds.count % 2 != 0 {
      compressedBonds.append(SIMD2<UInt16>.zero)
    }
    
    // Write the compressable bonds.
    for groupID in 0..<compressedBonds.count / 2 {
      var vector: SIMD4<UInt16> = .zero
      for laneID in 0..<2 {
        let compressedBond = compressedBonds[groupID * 2 + laneID]
        vector[laneID * 2 + 0] = compressedBond[0]
        vector[laneID * 2 + 1] = compressedBond[1]
      }
      let castedVector = unsafeBitCast(vector, to: SIMD2<UInt32>.self)
      rawData.append(castedVector)
    }
    
    // Write the non-compressable bonds.
    rawData.append(contentsOf: decompressedBonds)
    
    // Return the data for further processing in the caller.
    return Data(bytes: rawData, count: rawData.count * 8)
  }
  
  static func deserialize(bonds: Data) -> [SIMD2<UInt32>] {
    var rawData = [SIMD2<UInt32>](repeating: .zero, count: bonds.count / 8)
    rawData.withUnsafeMutableBytes {
      let copiedCount = bonds.copyBytes(to: $0)
      guard copiedCount == bonds.count else {
        fatalError("Incorrect number of bytes was copied.")
      }
    }
    
    // Read the header.
    let header = rawData[0]
    let compressedBondCount = UInt32(header[0])
    let decompressedBondCount = UInt32(header[1])
    
    // Pad the compressable bonds to a multiple of two.
    var paddedCompressedBondCount = compressedBondCount
    while paddedCompressedBondCount % 2 != 0 {
      paddedCompressedBondCount += 1
    }
    
    // Read the compressable bonds.
    var compressedBonds: [SIMD2<UInt16>] = []
    for groupID in 0..<paddedCompressedBondCount / 2 {
      let castedVector = rawData[Int(1 + groupID)]
      let vector = unsafeBitCast(castedVector, to: SIMD4<UInt16>.self)
      for laneID in 0..<2 {
        var compressedBond: SIMD2<UInt16> = .zero
        compressedBond[0] = vector[laneID * 2 + 0]
        compressedBond[1] = vector[laneID * 2 + 1]
        compressedBonds.append(compressedBond)
      }
    }
    
    // Restore the padded count to the original count.
    while compressedBonds.count > compressedBondCount {
      compressedBonds.removeLast()
    }
    
    // Read the non-compressable bonds.
    let decompressedRange = Int(1 + paddedCompressedBondCount / 2)...
    let decompressedBonds = Array(rawData[decompressedRange])
    guard decompressedBonds.count == decompressedBondCount else {
      fatalError("Could not decode decompressed bonds.")
    }
    
    // Merge the bonds into a common array.
    var lastAtomID: UInt32 = .zero
    var bonds: [SIMD2<UInt32>] = []
    for compressedBond in compressedBonds {
      let startDelta = UInt32(compressedBond[0])
      let lengthDelta = UInt32(compressedBond[1])
      let bond = SIMD2(
        lastAtomID + startDelta,
        lastAtomID + startDelta + lengthDelta)
      bonds.append(bond)
      
      // Update the atom cursor.
      lastAtomID += startDelta
    }
    bonds.append(contentsOf: decompressedBonds)
    
    // Sort the array of bonds.
    bonds.sort(by: { lhs, rhs in
      if lhs[0] > rhs[0] {
        return false
      } else if lhs[0] < rhs[0] {
        return true
      }
      if lhs[1] > rhs[1] {
        return false
      } else if lhs[1] < rhs[1] {
        return true
      }
      return true
    })
    return bonds
  }
}

// MARK: - Frames

extension Serialization {
  // Procedure for saving a simulation trajectory.
  //
  // Truncates the atom coordinates to picometer precision. This does not
  // make the data robust to infinitesimal rounding errors.
  static func encode(frames: [[SIMD4<Float>]]) -> Data {
    var mergedFrames: [SIMD4<Float>] = []
    for frame in frames {
      mergedFrames.append(contentsOf: frame)
    }
    let compressedTrajectory = Serialization.serialize(atoms: mergedFrames)
    
    var header: [UInt32] = []
    header.append(UInt32(compressedTrajectory.count))
    header.append(UInt32(frames.count))
    for frame in frames {
      header.append(UInt32(frame.count))
    }
    
    var serialized = Data()
    header.withUnsafeBytes {
      let chunk = Data(bytes: $0.baseAddress!, count: $0.count)
      serialized.append(chunk)
    }
    serialized.append(compressedTrajectory)
    return serialized
  }
  
  // Procedure for loading a simulation trajectory.
  static func decode(frames: Data) -> [[SIMD4<Float>]] {
    guard frames.count >= 4 else {
      fatalError("Not enough room for header.")
    }
    var dataCursor: Int = .zero
    
    // Decode the header.
    var header: [UInt32] = []
    frames.withUnsafeBytes { buffer in
      let rawPointer = buffer.baseAddress!
      do {
        let wordPointer = (rawPointer + dataCursor)
          .assumingMemoryBound(to: UInt32.self)
        header.append(wordPointer.pointee)
        dataCursor += 4
      }
      do {
        let wordPointer = (rawPointer + dataCursor)
          .assumingMemoryBound(to: UInt32.self)
        header.append(wordPointer.pointee)
        dataCursor += 4
      }
      
      let frameCount = Int(header[1])
      guard (frames.count - dataCursor) >= 4 * frameCount else {
        fatalError("Not enough room for frames.")
      }
      for _ in 0..<frameCount {
        let wordPointer = (rawPointer + dataCursor)
          .assumingMemoryBound(to: UInt32.self)
        header.append(wordPointer.pointee)
        dataCursor += 4
      }
    }
    
    // Decompress the trajectory data.
    guard (frames.count - dataCursor) == header[0] else {
      fatalError("Not enough room for compressed data.")
    }
    let compressedTrajectory: Data = frames[dataCursor...]
    let decompressedTrajectory = Serialization
      .deserialize(atoms: compressedTrajectory)
    
    // Check that the atom count matches the expected value.
    var totalAtomCount: Int = .zero
    for frameID in 0..<header[1] {
      let atomCount = header[2 + Int(frameID)]
      totalAtomCount += Int(atomCount)
    }
    guard decompressedTrajectory.count == totalAtomCount else {
      fatalError("Decompressed trajectory had unexpected size.")
    }
    
    // Decode the frames.
    var atomCursor: Int = .zero
    var output: [[SIMD4<Float>]] = []
    for frameID in 0..<header[1] {
      let atomCount = header[2 + Int(frameID)]
      let nextAtomCursor = atomCursor + Int(atomCount)
      let atoms = decompressedTrajectory[atomCursor..<nextAtomCursor]
      output.append(Array(atoms))
      
      atomCursor = nextAtomCursor
    }
    guard atomCursor == decompressedTrajectory.count else {
      fatalError("Failed to parse entire trajectory.")
    }
    
    return output
  }
  
  // Converts unsafe characters to underscores, to make the file name
  // compatible with the file system.
  static func fileSafeString(_ input: String) -> String {
    // Fetch the null-terminated C string.
    var cString = input.utf8CString
    for characterID in cString.indices {
      let byte = cString[characterID]
      let scalar = UnicodeScalar(UInt32(byte))!
      var character = Character(scalar)
      
      if character == "/" {
        character = "_"
      } else if character == ":" {
        character = "_"
      } else if character == ";" {
        character = "_"
      }
      cString[characterID] = CChar(character.asciiValue!)
    }
    return cString.withUnsafeBufferPointer {
      return String(cString: $0.baseAddress!)
    }
  }
}

SSDCaching.swift

//
//  SSDCaching.swift
//  molecular-renderer
//
//  Created by Philip Turner on 10/23/25.
//

import Foundation
import HDL

// Procedure for generating a unique identifier for the current state.
func createKey(tripod: Topology) -> String {
  let key = Serialization.hash(atoms: tripod.atoms)
  
  // RFC 3548 encoding: https://www.rfc-editor.org/rfc/rfc3548#page-6
  // "/" -> "_"
  // "+" -> "-"
  var base64Key = key.base64EncodedString()
  do {
    // Fetch the null-terminated C string.
    var cString = base64Key.utf8CString
    for characterID in cString.indices {
      let byte = cString[characterID]
      let scalar = UnicodeScalar(UInt32(byte))!
      var character = Character(scalar)
      
      if character == "/" {
        character = "_"
      } else if character == "+" {
        character = "-"
      }
      cString[characterID] = CChar(character.asciiValue!)
    }
    base64Key = cString.withUnsafeBufferPointer {
      return String(cString: $0.baseAddress!)
    }
  }
  return base64Key
}

func loadCachedTrajectory(
  tripod: Topology
) -> [[SIMD4<Float>]] {
  // Find the path.
  let packagePath = FileManager.default.currentDirectoryPath
  let key = createKey(tripod: tripod)
  
  // Load the cached trajectory.
  var frames: [[SIMD4<Float>]] = []
  do {
    let url = URL(
      filePath: "\(packagePath)/.build/tooltips/\(key)")
    let data = try Data(contentsOf: url)
    frames = Serialization.decode(frames: data)
  } catch {
    frames = runMinimization(tripod: tripod)
    
    let directoryURL = URL(
      filePath: "\(packagePath)/.build/tooltips")
    try! FileManager.default.createDirectory(
      at: directoryURL,
      withIntermediateDirectories: true)
    
    let data = Serialization.encode(frames: frames)
    let succeeded = FileManager.default.createFile(
      atPath: "\(packagePath)/.build/tooltips/\(key)",
      contents: data)
    guard succeeded else {
      fatalError("Could not create file.")
    }
  }
  
  return frames
}