MM4SpeedBenchmark.swift

import HDL
import MM4
import MolecularRenderer
import OpenMM
import QuaternionModule

import Foundation
import func QuartzCore.CACurrentMediaTime // profiling on macOS

// MARK: - Compile Structure

// 0.258 seconds to compile on M1 Max for 544079 atoms.
// 0.047 seconds to compile on M1 Max for 91891 atoms.
// 0.001 seconds to compile on M1 Max for 1509 atoms.
// 0.000 seconds to compile on M1 Max for 147 atoms.
let material: MaterialType = .elemental(.silicon)
let lattice = Lattice<Hexagonal> { h, k, l in
  let h2k = h + 2 * k
  
  // 554079-554080 atoms
  //
  // 40 * h + 40 * h2k + 40 * l
  //
  // 91891-91892 atoms
  //
  // 22 * h + 18 * h2k + 26 * l
  //
  // 1509-1510 atoms
  //
  // 5 * h + 5 * h2k + 5 * l
  //
  // 147-148 atoms
  //
  // 2 * h + 2 * h2k + 2 * l
  Bounds { 22 * h + 18 * h2k + 26 * l }
  Material { material }
  
  // Place one phosphorus atom to activate electrostatic forces with diamond.
//  Volume {
//    Concave {
//      Convex {
//        Origin { 0.1 * h }
//        Plane { -h }
//      }
//      Convex {
//        Origin { 0.2 * l }
//        Plane { -l }
//      }
//      Convex {
//        Origin { 0.5 * h2k }
//        Plane { -h2k }
//      }
//    }
//    
//    Replace { .atom(.phosphorus) }
//  }
}
var reconstruction = Reconstruction()
reconstruction.atoms = lattice.atoms
reconstruction.material = material
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 {
    // Don't H-passivate the phosphorus.
    let atom = topology.atoms[atomID]
    if atom.element == .phosphorus {
      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
}
passivate(topology: &topology)
topology.sort()

print()
print("atoms:", topology.atoms.count)
print("H:", topology.atoms.count(where: { $0.element == .hydrogen }))
print("C:", topology.atoms.count(where: { $0.element == .carbon }))
print("Si:", topology.atoms.count(where: { $0.element == .silicon }))
print("P:", topology.atoms.count(where: { $0.element == .phosphorus }))

// MARK: - Run Simulation Analysis

var parametersDesc = MM4ParametersDescriptor()
parametersDesc.atomicNumbers = topology.atoms.map(\.atomicNumber)
parametersDesc.bonds = topology.bonds

// 4.488 seconds to generate parameters on M1 Max for 91891 atoms.
// 0.685 seconds to generate parameters on M1 Max for 91891 atoms.
// 0.007 seconds to generate parameters on M1 Max for 1509 atoms.
// 0.001 seconds to generate parameters on M1 Max for 147 atoms.
//
// Compare that with 1.2-3.7 seconds projected for (61 ms) / (~3) execution
// time of NCFPart (1514 atoms) in the MM4 test suite
let parameters = try! MM4Parameters(descriptor: parametersDesc)
print("charged atoms:", parameters.atoms.parameters.count(where: { $0.charge != 0 }))

// # 554079-554080 atoms
//
// diamond (1st time): 10.939 s
// diamond (2nd time): 10.559 s
// diamond (3rd time): 10.406 s
//
// diamond + P (1st time): 15.819 s
// diamond + P (2nd time): 15.129 s
// diamond + P (3rd time): 15.078 s
//
// # 91891-91892 atoms
//
// diamond (1st time): 2.246 s
// diamond (2nd time): 1.544 s
// diamond (3rd time): 1.537 s
//
// diamond + P (1st time): 2.506 s
// diamond + P (2nd time): 2.198 s
// diamond + P (3rd time): 2.193 s
//
// silicon carbide (1st time): 2.228 s
// silicon carbide (2nd time): 2.174 s
// silicon carbide (3rd time): 2.162 s
//
// silicon (1st time): 1.578 s
// silicon (2nd time): 1.538 s
// silicon (3rd time): 1.546 s
//
// # 1509-1510 atoms
//
// diamond (1st time): 0.387 s
// diamond (2nd time): 0.075 s
// diamond (3rd time): 0.077 s
//
// diamond + P (1st time): 0.389 s
// diamond + P (2nd time): 0.092 s
// diamond + P (3rd time): 0.080 s
//
// # 147-148 atoms
//
// diamond (1st time): 0.369 s
// diamond (2nd time): 0.075 s
// diamond (3rd time): 0.061 s
//
// diamond + P (1st time): 0.361 s
// diamond + P (2nd time): 0.068 s
// diamond + P (3rd time): 0.063 s
var forceFieldDesc = MM4ForceFieldDescriptor()
forceFieldDesc.cutoffDistance = 1.5
 forceFieldDesc.integrator = .multipleTimeStep
forceFieldDesc.parameters = parameters
var forceField = try! MM4ForceField(descriptor: forceFieldDesc)
forceField.positions = topology.atoms.map(\.position)

// Give it a little push.
do {
  var rigidBodyDesc = MM4RigidBodyDescriptor()
  rigidBodyDesc.masses = parameters.atoms.masses
  rigidBodyDesc.positions = topology.atoms.map(\.position)
  var rigidBody = try! MM4RigidBody(descriptor: rigidBodyDesc)
  
  let linearSpeedInMetersPerS: Double = 5
  let linearSpeed = linearSpeedInMetersPerS / 1000
  let linearVelocity = linearSpeed * SIMD3<Double>(0, -1, 0)
  let linearMomentum = rigidBody.mass * linearVelocity
  rigidBody.linearMomentum = linearMomentum
  
  let angularFrequencyInGHz: Double = 3.2
  let angularFrequency = angularFrequencyInGHz / 1000
  let angularSpeed = angularFrequency * (2 * Double.pi)
  let angularVelocity = angularSpeed * SIMD3<Double>(1, 0, 0)
  let angularMomentum = rigidBody.momentOfInertia * angularVelocity
  rigidBody.angularMomentum = angularMomentum
  
  forceField.velocities = rigidBody.velocities
}

// # 554079-554080 atoms
//
// diamond (1st time): 0.586 s
// diamond (2nd time): 0.334 s
// diamond (3rd time): 0.330 s
//
// diamond + P (1st time): 0.534 s
// diamond + P (2nd time): 0.327 s
// diamond + P (3rd time): 0.321 s
//
// # 91891-91892 atoms
//
// diamond (1st time): 0.258 s
// diamond (2nd time): 0.074 s
// diamond (3rd time): 0.077 s
//
// diamond + P (1st time): 0.254 s
// diamond + P (2st time): 0.080 s
// diamond + P (3rd time): 0.084 s
//
// silicon carbide (1st time): 0.294 s
// silicon carbide (2nd time): 0.067 s
// silicon carbide (3rd time): 0.074 s
//
// silicon (1st time): 0.287 s
// silicon (2nd time): 0.064 s
// silicon (3rd time): 0.067 s
//
// # 1509-1510 atoms
//
// diamond (1st time): 0.212 s
// diamond (2nd time): 0.021 s
// diamond (3rd time): 0.025 s
//
// diamond + P (1st time): 0.203 s
// diamond + P (2nd time): 0.021 s
// diamond + P (3rd time): 0.024 s
//
// # 147-148 atoms
//
// diamond (1st time): 0.217 s
// diamond (2nd time): 0.024 s
// diamond (3rd time): 0.024 s
//
// diamond + P (1st time): 0.204 s
// diamond + P (2nd time): 0.024 s
// diamond + P (3rd time): 0.023 s
print("initial kinetic:", forceField.energy.kinetic)
print("initial potential:", forceField.energy.potential)

// Utility function for calculating temperature.
@MainActor
func temperature(kineticEnergy: Double) -> Float {
  let energyInJ = kineticEnergy * 1e-21
  let energyPerAtom = Float(energyInJ / Double(topology.atoms.count))
  return energyPerAtom * 2 / (3 * 1.380649e-23)
}

// Analyze the energy over a few timesteps.
let frameTimeInterval: Float = 0.3
var previousTimeStamp: Double?
for frameID in 0...10 {
  print()
  
//  _ = forceField.energy.kinetic
  
  // report statistics
  if false {
    print("frame ID = \(frameID)")
    
    let time = Float(frameID) * frameTimeInterval
    let timeRepr = String(format: "%.3f", time)
    print("time = \(timeRepr) ps")
    
    let kinetic = forceField.energy.kinetic
    let kineticRepr = String(format: "%.1f", kinetic)
    print("kinetic energy = \(kineticRepr) zJ")
    
    let potential = forceField.energy.potential
    let potentialRepr = String(format: "%.1f", potential)
    print("potential energy = \(potentialRepr) zJ")
    
    let totalEnergy = kinetic + potential
    let totalEnergyRepr = String(format: "%.1f", totalEnergy)
    print("total energy = \(totalEnergyRepr) zJ")
    
    let temperature = temperature(kineticEnergy: kinetic)
    let temperatureRepr = String(format: "%.1f", temperature)
    print("temperature = \(temperatureRepr) K")
    
    guard temperature < 1000 else {
      fatalError("Temperature exceeded 1000 K or was NaN.")
    }
  }
  
  // analyze rigid body statistics
  if false {
    var rigidBodyDesc = MM4RigidBodyDescriptor()
    rigidBodyDesc.masses = parameters.atoms.masses
    rigidBodyDesc.positions = forceField.positions
    rigidBodyDesc.velocities = forceField.velocities
    let rigidBody = try! MM4RigidBody(descriptor: rigidBodyDesc)
    
    // invert p = mv
    let linearMomentum = rigidBody.linearMomentum
    let linearVelocity = linearMomentum / rigidBody.mass
    let linearSpeed = (linearVelocity * linearVelocity).sum().squareRoot()
    let linearSpeedRepr = String(format: "%.3f", linearSpeed * 1000)
    print("speed = \(linearSpeedRepr) m/s")
    
    // invert L = Iω
    let angularMomentum = rigidBody.angularMomentum
    let angularVelocity = angularMomentum / rigidBody.momentOfInertia
    let angularSpeed = (angularVelocity * angularVelocity).sum().squareRoot()
    let angularFrequency = angularSpeed / (2 * Double.pi)
    let angularFrequencyRepr = String(format: "%.3f", angularFrequency * 1000)
    print("frequency = \(angularFrequencyRepr) GHz")
    
    let centerOfMass = rigidBody.centerOfMass
    let xRepr = String(format: "%.3f", centerOfMass.x)
    let yRepr = String(format: "%.3f", centerOfMass.y)
    let zRepr = String(format: "%.3f", centerOfMass.z)
    print("center of mass = (\(xRepr), \(yRepr), \(zRepr))")
  }
  
  // profile ns/day
  do {
    let currentTimeStamp = CACurrentMediaTime()
    if let previousTimeStamp {
      let elapsedTime = currentTimeStamp - previousTimeStamp
      let elapsedTimeInDays = elapsedTime / Double(86400)
      
      let simulationTime = Double(frameTimeInterval)
      let simulationTimeInNs = simulationTime / Double(1000)
      
      let speed = simulationTimeInNs / elapsedTimeInDays
      let speedRepr = String(format: "%.3f", speed)
      print(speedRepr, "ns/day")
    }
    previousTimeStamp = currentTimeStamp
  }
  
  if frameID < 10 {
    // perform time evolution
    forceField.simulate(time: Double(frameTimeInterval))
  }
}

// MARK: - Launch Application

#if false

@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.allocationSize = 10_000
  applicationDesc.device = device
  applicationDesc.display = display
  applicationDesc.upscaleFactor = 3
  let application = Application(descriptor: applicationDesc)
  
  return application
}
let application = createApplication()

@MainActor
func modifyAtoms() {
  let atoms = topology.atoms
  if atoms.count > 100 {
    fatalError("Tried to render a large lattice.")
  }
  
  for atomID in atoms.indices {
    let atom = atoms[atomID]
    application.atoms[atomID] = atom
  }
}

@MainActor
func modifyCamera() {
  let rotation = Quaternion<Float>(
    angle: Float.pi / 180 * -110,
    axis: SIMD3(0, 1, 0))
  
  // Place the camera 1.0 nm away from the origin.
  application.camera.position = rotation.act(on: SIMD3(0, 0, 2))
  
  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)
}

#endif