HexagonDraft3.swift

import Foundation
import HDL
import MM4
import Numerics
import OpenMM

func createGeometry() -> [Entity] {
  // Create the compiled structure.
  var hexagon = Hexagon()
  hexagon.minimize()
  hexagon.center(position: [0, 0, -2])
  
  // Render the minimized structure.
  //
  // Minimization is a necessary step before running an MD simulation. Every
  // part needs to be minimized. Equilibriation, on the other hand, can either
  // be very cheap or neglected (if working at ~0 K). I've found a way to get
  // useful equilibriation at 273 K with only 1 ps of simulation.
  return hexagon.topology.atoms
}

struct Hexagon {
  var topology: Topology
  
  init() {
    let lattice = Self.createLattice()
    topology = Self.createTopology(lattice: lattice)
  }
  
  static func createLattice() -> Lattice<Hexagonal> {
    Lattice<Hexagonal> { h, k, l in
      let h2k = h + 2 * k
      Bounds { 16 * h + 10 * h2k + 2 * l }
      Material { .elemental(.carbon) }
      
      Volume {
        Origin { 8 * h + 5 * h2k }
        
        var directions: [SIMD3<Float>] = []
        directions.append(k)
        directions.append(h)
        directions.append(-k)
        directions.append(-h)
        directions.append(k + h)
        directions.append(-k - h)
        
        for direction in directions {
          Convex {
            Origin { 5 * direction }
            Plane { direction }
          }
        }
        
        Replace { .empty }
      }
    }
  }
  
  static func createTopology(lattice: Lattice<Hexagonal>) -> Topology {
    var topology = Topology()
    topology.insert(atoms: lattice.atoms)
    guard topology.bonds.count == .zero else {
      fatalError("The topology should not have any bonds yet.")
    }
    
    // Encapsulate 'insertedBonds' in a scope. This way, you cannot reference
    // the variable name (by accident) later in the code.
    do {
      let matches = topology.match(
        topology.atoms, algorithm: .covalentBondLength(1.5))
      
      var insertedBonds: [SIMD2<UInt32>] = []
      for i in topology.atoms.indices {
        for j in matches[i] where i < j {
          let bond = SIMD2(UInt32(i), UInt32(j))
          insertedBonds.append(bond)
        }
      }
      topology.insert(bonds: insertedBonds)
    }
    guard topology.bonds.count > .zero else {
      fatalError("No bonds were added.")
    }
    
    // Encapsulate 'nonbondingOrbitals', 'insertedAtoms', and 'insertedBonds'
    // within their own scope, so they will be deleted as soon as they're not
    // relevant.
    do {
      // 'nonbondingOrbitals' is a list of lists.
      let nonbondingOrbitals = topology.nonbondingOrbitals(hybridization: .sp3)
      
      var insertedAtoms: [Entity] = []
      var insertedBonds: [SIMD2<UInt32>] = []
      for atomID in topology.atoms.indices {
        let carbon = topology.atoms[atomID]
        let carbonPosition = carbon.position
        
        // Here, you have fetched a list from the 'list of lists'.
        let orbitals = nonbondingOrbitals[atomID]
        for orbital in orbitals {
          // Get a rough estimate of the bond length. I often substitute this
          // with precise values from the literature, or MM3/MM4 equilibrium
          // bond lengths. Usually it doesn't matter too much, because
          // minimization will relax it to the exact value.
          let chBondLength = Element.carbon.covalentRadius +
          Element.hydrogen.covalentRadius
          
          // Position + Length * (Direction Vector)
          //
          // The direction vector is normalized (unit length), and
          // 'chBondLength' effectively gives it a length.
          let hydrogenPosition = carbonPosition + chBondLength * orbital
          let hydrogen = Entity(
            position: hydrogenPosition, type: .atom(.hydrogen))
          
          // Predict the passivator's index in the topology, after the current
          // atoms are merged with the 'inserted atoms'.
          let hydrogenID = topology.atoms.count + insertedAtoms.count
          
          let bond = SIMD2(UInt32(atomID), UInt32(hydrogenID))
          insertedAtoms.append(hydrogen)
          insertedBonds.append(bond)
        }
      }
      
      topology.insert(atoms: insertedAtoms)
      topology.insert(bonds: insertedBonds)
    }
    
    return topology
  }
  
  // Runs an energy minimization to relax the structure into the equilibrium
  // position.
  mutating func minimize() {
    var paramsDesc = MM4ParametersDescriptor()
    paramsDesc.atomicNumbers = topology.atoms.map(\.atomicNumber)
    paramsDesc.bonds = topology.bonds
    let parameters = try! MM4Parameters(descriptor: paramsDesc)
    
    var forceFieldDesc = MM4ForceFieldDescriptor()
    forceFieldDesc.parameters = parameters
    let forceField = try! MM4ForceField(descriptor: forceFieldDesc)
    
    // When initialized, the forcefield knows nothing about the atom positions.
    forceField.positions = topology.atoms.map(\.position)
    forceField.minimize()
    
    // The atoms shouldn't change much, because hexagonal diamond surfaces are
    // already compiled in the ~equilibrium position. For cubic diamond (100)
    // surfaces, you would see a much more noticeable change.
    for atomID in parameters.atoms.indices {
      var atom = topology.atoms[atomID]
      let position = forceField.positions[atomID]
      atom.position = position
      topology.atoms[atomID] = atom
    }
  }
  
  // Moves the center of mass to the specified position.
  mutating func center(position centerPosition: SIMD3<Float>) {
    // Accumulate in double precision to minimize rounding error.
    var weightedCenterAccumulator: SIMD3<Double> = .zero
    var massAccumulator: Double = .zero
    
    // Iterate over the atoms.
    for atomID in topology.atoms.indices {
      let atom = topology.atoms[atomID]
      
      // Approximate the atomic mass by using the atomic number (larger atoms
      // are heavier). We don't have access to atomic masses this early on in
      // compilation.
      let mass = Float(atom.atomicNumber)
      let position: SIMD3<Float> = atom.position
      
      weightedCenterAccumulator += SIMD3<Double>(position * mass)
      massAccumulator += Double(mass)
    }
    
    // Cast the final sum to single precision, because we don't need the extra
    // precision anymore.
    let weightedCenter = SIMD3<Float>(weightedCenterAccumulator)
    let mass = Float(massAccumulator)
    let centerOfMass = weightedCenter / mass
    
    // Shift the center of mass to the origin.
    for atomID in topology.atoms.indices {
      var atom = topology.atoms[atomID]
      atom.position -= centerOfMass
      topology.atoms[atomID] = atom
    }
    
    // Shift the center of mass to the specified position.
    for atomID in topology.atoms.indices {
      var atom = topology.atoms[atomID]
      atom.position += centerPosition
      topology.atoms[atomID] = atom
    }
  }
}