Fix up initial velocities, deployment script and add rattling option for optimisation

Author: peterspackman

.github/workflows/website.yml

@@ -64,10 +64,6 @@ jobs:
         working-directory: website
         run: npm run build
 
-      - name: Convert PET-MAD model
-        run: |
-          uv run scripts/convert_pet_mad.py --output website/dist/pet-mad.gguf
-
       - name: Setup Pages
         uses: actions/configure-pages@v4
 

website/src/components/MolecularDynamics.tsx

@@ -297,6 +297,7 @@ export default function MolecularDynamics() {
   const [mode, setMode] = useState<'md' | 'optimize'>('md')
   const [maxOptSteps, setMaxOptSteps] = useState(100)
   const [forceThreshold, setForceThreshold] = useState(0.05)
+  const [rattleAmount, setRattleAmount] = useState(0.1)  // Angstroms
 
   // Initialize NGL Stage
   useEffect(() => {
@@ -408,6 +409,11 @@ export default function MolecularDynamics() {
           setState(s => ({ ...s, isRunning: false }))
           break
 
+        case 'rattled':
+          // Update visualization with rattled positions
+          updateVisualization(msg.positions)
+          break
+
         case 'error':
           setState(s => ({ ...s, error: msg.message, isRunning: false }))
           break
@@ -657,13 +663,21 @@ export default function MolecularDynamics() {
     setEnergyHistory([])
     setState(s => ({ ...s, step: 0, optimizationConverged: false }))
     lastStepTimeRef.current = 0
-    workerRef.current?.postMessage({ type: 'start', stepsPerFrame: 1, mode })
+    workerRef.current?.postMessage({
+      type: 'start',
+      stepsPerFrame: 1,
+      mode,
+    })
   }
 
   const stopMD = () => {
     workerRef.current?.postMessage({ type: 'stop' })
   }
 
+  const rattleStructure = () => {
+    workerRef.current?.postMessage({ type: 'rattle', amount: rattleAmount })
+  }
+
   return (
     <div className="md-simulation">
       {/* Left panel - Structure and parameters */}
@@ -789,8 +803,32 @@ export default function MolecularDynamics() {
                   />
                 </div>
               </div>
+              <div className="params-row">
+                <div className="control-group">
+                  <label>Rattle (Å)</label>
+                  <input
+                    type="number"
+                    value={rattleAmount}
+                    onChange={e => setRattleAmount(Number(e.target.value))}
+                    min={0}
+                    max={1.0}
+                    step={0.05}
+                    className="number-input"
+                  />
+                </div>
+                <div className="control-group">
+                  <label>&nbsp;</label>
+                  <button
+                    onClick={rattleStructure}
+                    className="control-button"
+                    disabled={!state.isModelLoaded || atomicNumbersRef.current.length === 0 || state.isRunning}
+                  >
+                    Rattle
+                  </button>
+                </div>
+              </div>
               <p className="nc-note">
-                FIRE geometry optimization using gradient forces.
+                FIRE optimization. Rattle perturbs atom positions.
               </p>
             </>
           )}
@@ -867,15 +905,16 @@ export default function MolecularDynamics() {
                         stroke="rgba(59, 130, 246, 0.9)"
                         strokeWidth="2.5"
                       />
-                      {/* Show dot at current position */}
-                      {dots.length > 0 && (
+                      {/* Show small dots at each data point */}
+                      {dots.map((dot, i) => (
                         <circle
-                          cx={dots[dots.length - 1].x}
-                          cy={dots[dots.length - 1].y}
-                          r="8"
+                          key={i}
+                          cx={dot.x}
+                          cy={dot.y}
+                          r={i === dots.length - 1 ? 6 : 3}
                           fill="rgba(59, 130, 246, 1)"
                         />
-                      )}
+                      ))}
                       <text x="10" y="18" className="energy-label">
                         {mode === 'md' ? 'Total E' : 'E'} (eV)
                       </text>

website/src/workers/mdWorker.ts

@@ -389,15 +389,42 @@ const FIRE_F_DEC = 0.5
 const FIRE_N_MIN = 5
 const FIRE_DT_MAX = 1.0  // fs
 
-// Reset FIRE optimizer state
+// Reset FIRE optimizer state and initialize velocities along force direction
 function resetFIRE(): void {
   state.fireAlpha = FIRE_ALPHA_START
   state.fireNpos = 0
   state.fireDt = 0.1  // Start with small timestep
   state.optStep = 0
-  // Zero velocities
-  if (state.velocities) {
-    state.velocities.fill(0)
+
+  // Initialize velocities along force direction for faster startup
+  if (state.module && state.model && state.positions && state.velocities && state.masses) {
+    // Get initial forces
+    const system = state.module.AtomicSystem.create(
+      state.positions,
+      state.atomicNumbers!,
+      state.cell,
+      state.isPeriodic
+    )
+    const result = state.model.predictWithOptions(system, false)
+    const forces = new Float64Array(result.forces)
+
+    // Calculate force magnitude
+    let fNorm = 0
+    for (let i = 0; i < forces.length; i++) {
+      fNorm += forces[i] * forces[i]
+    }
+    fNorm = Math.sqrt(fNorm)
+
+    // Set initial velocity along force direction with small magnitude
+    // v = dt * F / |F| gives unit velocity in force direction scaled by timestep
+    if (fNorm > 1e-10) {
+      const vScale = state.fireDt * 0.1  // Small initial velocity
+      for (let i = 0; i < state.velocities.length; i++) {
+        state.velocities[i] = vScale * forces[i] / fNorm
+      }
+    } else {
+      state.velocities.fill(0)
+    }
   }
 }
 
@@ -607,7 +634,17 @@ function runMDStep(): void {
   }
 }
 
-function handleStart(data: { stepsPerFrame?: number, mode?: 'md' | 'optimize' }): void {
+// Apply random perturbations to positions
+function rattlePositions(amount: number): void {
+  if (!state.positions || amount <= 0) return
+
+  for (let i = 0; i < state.positions.length; i++) {
+    // Uniform random in [-amount, +amount]
+    state.positions[i] += (Math.random() * 2 - 1) * amount
+  }
+}
+
+function handleStart(data: { stepsPerFrame?: number, mode?: 'md' | 'optimize', rattleAmount?: number }): void {
   if (state.isRunning) return
 
   // Update mode if provided
@@ -621,6 +658,11 @@ function handleStart(data: { stepsPerFrame?: number, mode?: 'md' | 'optimize' })
     // Reset FIRE state for new optimization
     resetFIRE()
 
+    // Apply rattle if requested
+    if (data.rattleAmount && data.rattleAmount > 0) {
+      rattlePositions(data.rattleAmount)
+    }
+
     // Run optimization steps at ~30 fps
     mdInterval = setInterval(() => {
       const done = runFIREStep()
@@ -656,6 +698,21 @@ function handleStep(): void {
   runMDStep()
 }
 
+function handleRattle(data: { amount: number }): void {
+  if (!state.positions) {
+    self.postMessage({ type: 'error', message: 'No system loaded' })
+    return
+  }
+
+  rattlePositions(data.amount)
+
+  // Send back the new positions so visualization can update
+  self.postMessage({
+    type: 'rattled',
+    positions: Array.from(state.positions),
+  })
+}
+
 // Message router
 self.onmessage = async (e: MessageEvent) => {
   const { type, ...data } = e.data
@@ -685,6 +742,9 @@ self.onmessage = async (e: MessageEvent) => {
     case 'step':
       handleStep()
       break
+    case 'rattle':
+      handleRattle(data)
+      break
     default:
       self.postMessage({ type: 'error', message: `Unknown message type: ${type}` })
   }