GitHub Gist for the propargyl alcohol 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
      }
    }
  }
}

main.swift

import HDL
import MolecularRenderer
import QuaternionModule
import xTB

// MARK: - Compile Structure

let lattice = Lattice<Cubic> { h, k, l in
  Bounds { 1 * (h + k + l) }
  Material { .elemental(.carbon) }
  
  Volume {
    Concave {
      Convex {
        Origin { 0.3 * h }
        Plane { -h }
      }
      Convex {
        Origin { 0.3 * k }
        Plane { -k }
      }
      Convex {
        Origin { 0.3 * l }
        Plane { -l }
      }
    }
    
    Replace { .atom(.germanium) }
  }
}

var reconstruction = Reconstruction()
reconstruction.atoms = lattice.atoms
reconstruction.material = .elemental(.carbon)
var topology = reconstruction.compile()

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 orbitalList = orbitalLists[atomID]
    
    // Do not passivate atoms in the bridgehead position.
    if orbitalList.count == 1 {
      continue
    }
    
    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
}
passivate(topology: &topology)

// Move the germanium atom to the origin.
do {
  var germaniumPosition: SIMD3<Float>?
  for atom in topology.atoms {
    if atom.element == .germanium {
      germaniumPosition = atom.position
    }
  }
  guard let germaniumPosition else {
    fatalError("Could not find germanium position.")
  }
  
  let translation = SIMD3<Float>(0, 0, 0) - 1 * germaniumPosition
  for atomID in topology.atoms.indices {
    var atom = topology.atoms[atomID]
    atom.position += translation
    topology.atoms[atomID] = atom
  }
}

// Rotate so the germanium's free radical (at this state of molecule
// compilation) points toward +Y.
do {
  // All attempts to rotate with a quaternion failed, producing weird results.
  let basisVector1 = SIMD3<Float>(-1, -1, -1) / Float(3).squareRoot()
  let basisVector2 = SIMD3<Float>(-1, 0, 1) / Float(2).squareRoot()
  let basisVector3 = cross(basisVector1, basisVector2)
  
  for atomID in topology.atoms.indices {
    var atom = topology.atoms[atomID]
    let newX = (atom.position * basisVector2).sum()
    let newY = (atom.position * basisVector1).sum()
    let newZ = (atom.position * basisVector3).sum()
    
    atom.position = SIMD3(newX, newY, newZ)
    topology.atoms[atomID] = atom
  }
}

// Grow the ethynyl legs.
var carbonsToPassivate: [UInt32] = []
do {
  let orbitalLists = topology.nonbondingOrbitals()
  
  var insertedAtoms: [Atom] = []
  var insertedBonds: [SIMD2<UInt32>] = []
  for atomID in topology.atoms.indices {
    let atom = topology.atoms[atomID]
    let orbitalList = orbitalLists[atomID]
    guard orbitalList.count == 1 else {
      continue
    }
    
    var orbital = orbitalList.first!
    
    // To accelerate the minimization, the propargyl alcohol groups should be
    // angled in the direction the adamantane warps, when it expands from the
    // bulky germanium atom.
    if orbital.y < 0 {
      let downDirection = SIMD3<Float>(0, -1, 0)
      var newOrbital = 5 * orbital + 1 * downDirection
      newOrbital /= (newOrbital * newOrbital).sum().squareRoot()
      orbital = newOrbital
    }
    
    let bondLength1 = atom.element.covalentRadius + Element.carbon.covalentRadius
    let carbon1 = Atom(
      position: atom.position + orbital * bondLength1,
      element: .carbon)
    let carbonID1 = topology.atoms.count + insertedAtoms.count
    insertedAtoms.append(carbon1)
    
    let carbon2 = Atom(
      position: carbon1.position + orbital * 0.120,
      element: .carbon)
    let carbonID2 = topology.atoms.count + insertedAtoms.count
    insertedAtoms.append(carbon2)
    
    insertedBonds.append(SIMD2(
      UInt32(atomID),
      UInt32(carbonID1)))
    insertedBonds.append(SIMD2(
      UInt32(carbonID1),
      UInt32(carbonID2)))
    
    // Only form a propargyl alcohol on the bottom three bridgehead sites.
    // Leave the other one as an ethynyl group terminated in a free radical.
    guard orbital.y < 0 else {
      continue
    }
    
    let bondLength3 = 2 * Element.carbon.covalentRadius
    let carbon3 = Atom(
      position: carbon2.position + orbital * bondLength3,
      element: .carbon)
    let carbonID3 = topology.atoms.count + insertedAtoms.count
    insertedAtoms.append(carbon3)
    carbonsToPassivate.append(UInt32(carbonID3))
    
    let oxygenDirection = SIMD3<Float>(0, -1, 0)
    let bondLengthO = Element.carbon.covalentRadius + Element.oxygen.covalentRadius
    let oxygen = Atom(
      position: carbon3.position + oxygenDirection * bondLengthO,
      element: .oxygen)
    let oxygenID = topology.atoms.count + insertedAtoms.count
    insertedAtoms.append(oxygen)
    
    func createHydrogenVector() -> SIMD3<Float> {
      let vector1 = orbital
      let vector2 = oxygenDirection
      let perpendicular = cross(vector1, vector2)
      
      var output = 2 * perpendicular + 1 * vector2
      output /= (output * output).sum().squareRoot()
      return output
    }
    
    let bondLengthH = Element.oxygen.covalentRadius + Element.hydrogen.covalentRadius
    let hydrogen = Atom(
      position: oxygen.position + createHydrogenVector() * bondLengthH,
      element: .hydrogen)
    let hydrogenID = topology.atoms.count + insertedAtoms.count
    insertedAtoms.append(hydrogen)
    
    insertedBonds.append(SIMD2(
      UInt32(carbonID2),
      UInt32(carbonID3)))
    insertedBonds.append(SIMD2(
      UInt32(carbonID3),
      UInt32(oxygenID)))
    insertedBonds.append(SIMD2(
      UInt32(oxygenID),
      UInt32(hydrogenID)))
  }
  
  topology.atoms += insertedAtoms
  topology.bonds += insertedBonds
}
guard carbonsToPassivate.count == 3 else {
  fatalError("This should never happen.")
}

// Passivate the three secondary carbons that are currently carbene diradicals.
do {
  let orbitalLists = topology.nonbondingOrbitals(hybridization: .sp3)
  
  var insertedAtoms: [Atom] = []
  var insertedBonds: [SIMD2<UInt32>] = []
  for atomID in carbonsToPassivate {
    let atom = topology.atoms[Int(atomID)]
    let orbitalList = orbitalLists[Int(atomID)]
    guard orbitalList.count == 2 else {
      fatalError("This should never happen.")
    }
    
    for orbital in orbitalList {
      let bondLength = Element.carbon.covalentRadius + Element.hydrogen.covalentRadius
      let hydrogen = Atom(
        position: atom.position + orbital * bondLength,
        element: .hydrogen)
      let hydrogenID = topology.atoms.count + insertedAtoms.count
      insertedAtoms.append(hydrogen)
      
      insertedBonds.append(SIMD2(
        atomID,
        UInt32(hydrogenID)))
    }
  }
  
  topology.atoms += insertedAtoms
  topology.bonds += insertedBonds
}

// Should converge in 167 frames.
xTB_Environment.verbosity = .muted
let frames = minimize(atoms: topology.atoms)
print("Minimization completed in \(frames.count) frames.")

// Commented out line for debugging the static structure.
//let frames = [topology.atoms]

// MARK: - Launch Application

// Input: time in seconds
// Output: atoms
@MainActor
func interpolate(
  frames: [[Atom]],
  time: Float
) -> [Atom] {
  let multiple50Hz = time * 50
  var lowFrame = Int(multiple50Hz.rounded(.down))
  var highFrame = lowFrame + 1
  var lowInterpolationFactor = Float(highFrame) - multiple50Hz
  var highInterpolationFactor = multiple50Hz - Float(lowFrame)
  
  if lowFrame < -1 {
    fatalError("This should never happen.")
  }
  if lowFrame >= frames.count - 1 {
    lowFrame = frames.count - 1
    highFrame = frames.count - 1
    lowInterpolationFactor = 1
    highInterpolationFactor = 0
  }
  
  var output: [Atom] = []
  for atomID in topology.atoms.indices {
    let lowAtom = frames[lowFrame][atomID]
    let highAtom = frames[highFrame][atomID]
    
    var position: SIMD3<Float> = .zero
    position += lowAtom.position * lowInterpolationFactor
    position += highAtom.position * highInterpolationFactor
    
    let element = topology.atoms[atomID].element
    let atom = Atom(position: position, element: element)
    output.append(atom)
  }
  return output
}

@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>(1200, 1200)
  displayDesc.monitorID = device.fastestMonitorID
  let display = Display(descriptor: displayDesc)
  
  // Set up the application.
  var applicationDesc = ApplicationDescriptor()
  applicationDesc.device = device
  applicationDesc.display = display
  applicationDesc.upscaleFactor = 3
  
  applicationDesc.addressSpaceSize = 4_000_000
  applicationDesc.voxelAllocationSize = 500_000_000
  applicationDesc.worldDimension = 64
  let application = Application(descriptor: applicationDesc)
  
  return application
}
let application = createApplication()

@MainActor
func createTime() -> Float {
  let elapsedFrames = application.clock.frames
  let frameRate = application.display.frameRate
  let seconds = Float(elapsedFrames) / Float(frameRate)
  return seconds
}

@MainActor
func modifyAtoms() {
  let time = createTime()
  let delayTime: Float = 1
  let animationEndTime: Float = delayTime + Float(frames.count) / 50
  
  if time < delayTime {
    let atoms = topology.atoms
    for atomID in atoms.indices {
      let atom = atoms[atomID]
      application.atoms[atomID] = atom
    }
  } else if time < animationEndTime {
    let atoms = interpolate(
      frames: frames,
      time: time - delayTime)
    for atomID in atoms.indices {
      let atom = atoms[atomID]
      application.atoms[atomID] = atom
    }
  } else {
    let atoms = frames.last!
    
    // 0.1 Hz rotation rate
    let time = createTime()
    let angleDegrees = 0.1 * (time - animationEndTime) * 360
    let rotation = Quaternion<Float>(
      angle: Float.pi / 180 * angleDegrees,
      axis: SIMD3(0, 1, 0))
    
    for atomID in atoms.indices {
      var atom = atoms[atomID]
      atom.position = rotation.act(on: atom.position)
      application.atoms[atomID] = atom
    }
  }
}

@MainActor
func modifyCamera() {
  let rotation = Quaternion<Float>(
    angle: Float.pi / 180 * 0,
    axis: SIMD3(0, 1, 0))
  
  application.camera.position = rotation.act(on: SIMD3(0, -0.2, 1.5))
  
  application.camera.basis.0 = rotation.act(on: SIMD3(1, 0, 0))
  application.camera.basis.1 = rotation.act(on: SIMD3(0, 1, 0))
  application.camera.basis.2 = rotation.act(on: SIMD3(0, 0, 1))
  application.camera.fovAngleVertical = Float.pi / 180 * 60
}

// Enter the run loop.
application.run {
  modifyAtoms()
  modifyCamera()
  
  var image = application.render()
  image = application.upscale(image: image)
  application.present(image: image)
}

MathUtilities.swift

func cross(
  _ lhs: SIMD3<Float>,
  _ rhs: SIMD3<Float>
) -> SIMD3<Float> {
  let yzx = SIMD3<Int>(1, 2, 0)
  let zxy = SIMD3<Int>(2, 0, 1)
  return (lhs[yzx] * rhs[zxy]) - (lhs[zxy] * rhs[yzx])
}

Minimize.swift

//
//  Minimize.swift
//  MolecularRendererApp
//
//  Created by Philip Turner on 6/7/24.
//

import HDL
import var MM4.MM4YgPerAmu
import xTB

// Some quick and rough utilities for minimizing geometry.
// - We need some more sophisticated code to handle ONIOM models.
// - The code should handle both xTB-only and ONIOM minimizations.

// 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
  }
}

// Minimizes a structure with xTB.
func minimize(atoms: [Atom], anchorIDs: [UInt32] = []) -> [[Atom]] {
  var calculatorDesc = xTB_CalculatorDescriptor()
  calculatorDesc.atomicNumbers = atoms.map(\.atomicNumber)
  calculatorDesc.positions = atoms.map(\.position)
  calculatorDesc.hamiltonian = .tightBinding
  let calculator = xTB_Calculator(descriptor: calculatorDesc)
  
  var minimizationDesc = FIREMinimizationDescriptor()
  minimizationDesc.masses = atoms.map {
    if $0.atomicNumber == 1 {
      return Float(4.0 * MM4YgPerAmu)
    } else {
      return Float(12.011 * MM4YgPerAmu)
    }
  }
  minimizationDesc.positions = calculator.molecule.positions
  minimizationDesc.anchors = Set(anchorIDs)
  var minimization = FIREMinimization(descriptor: minimizationDesc)
  
  var frames: [[Atom]] = []
  func createFrame() -> [Atom] {
    var output: [Atom] = []
    for atomID in atoms.indices {
      var atom = 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

    let forces = calculator.molecule.forces
    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()
    
    if converged {
      frames.append(createFrame())
      break
    } else if trialID == 499 {
      print("failed to converge!")
    }
  }
  
  return frames
}