Computational Biology & Molecular Simulation

import numpy as np
import matplotlib.pyplot as plt

# Physical constants
kB = 1.381e-23  # Boltzmann constant (J/K)
NA = 6.022e23   # Avogadro's number

class MDSimulation:
    """Simple molecular dynamics simulation framework"""

    def __init__(self, n_atoms, box_size, dt=1e-15):
        """
        Initialize MD simulation

        Parameters:
        n_atoms: number of atoms
        box_size: simulation box size (nm)
        dt: timestep (seconds, default 1 fs)
        """
        self.n_atoms = n_atoms
        self.box_size = box_size * 1e-9  # Convert to meters
        self.dt = dt

        # Initialize positions randomly in box
        self.positions = np.random.rand(n_atoms, 3) * self.box_size

        # Initialize velocities (Maxwell-Boltzmann at 300K)
        self.velocities = self.maxwell_boltzmann_velocities(300)

        # Initialize forces
        self.forces = np.zeros((n_atoms, 3))

        # Energy tracking
        self.kinetic_energy = []
        self.potential_energy = []
        self.total_energy = []

    def maxwell_boltzmann_velocities(self, T, mass=1.67e-27):
        """Initialize velocities from Maxwell-Boltzmann distribution"""
        sigma = np.sqrt(kB * T / mass)
        velocities = np.random.normal(0, sigma, (self.n_atoms, 3))
        # Remove center of mass motion
        velocities -= np.mean(velocities, axis=0)
        return velocities

    def lennard_jones_force(self, epsilon=0.996e-21, sigma=3.4e-10):
        """
        Calculate Lennard-Jones forces between all atom pairs
        V(r) = 4ε[(σ/r)^12 - (σ/r)^6]
        """
        forces = np.zeros_like(self.positions)
        potential = 0.0

        for i in range(self.n_atoms):
            for j in range(i+1, self.n_atoms):
                # Distance vector with periodic boundary conditions
                rij = self.positions[i] - self.positions[j]
                rij -= self.box_size * np.round(rij / self.box_size)

                r = np.linalg.norm(rij)

                if r < 2.5 * sigma:  # Cutoff distance
                    # LJ potential and force
                    sr6 = (sigma / r)**6
                    sr12 = sr6 * sr6

                    potential += 4 * epsilon * (sr12 - sr6)

                    force_mag = 24 * epsilon * (2*sr12 - sr6) / r**2
                    force_vec = force_mag * rij

                    forces[i] += force_vec
                    forces[j] -= force_vec

        return forces, potential

    def velocity_verlet_step(self):
        """Velocity Verlet integration step"""
        # Update positions: r(t+dt) = r(t) + v(t)*dt + 0.5*a(t)*dt^2
        self.positions += self.velocities * self.dt + 0.5 * self.forces * self.dt**2 / 1.67e-27

        # Apply periodic boundary conditions
        self.positions = np.mod(self.positions, self.box_size)

        # Calculate new forces
        old_forces = self.forces.copy()
        self.forces, potential = self.lennard_jones_force()

        # Update velocities: v(t+dt) = v(t) + 0.5*[a(t) + a(t+dt)]*dt
        self.velocities += 0.5 * (old_forces + self.forces) * self.dt / 1.67e-27

        # Calculate energies
        kinetic = 0.5 * 1.67e-27 * np.sum(self.velocities**2)

        return kinetic, potential

    def run(self, n_steps, sample_interval=10):
        """Run MD simulation"""
        print(f"Starting MD simulation with {self.n_atoms} atoms")
        print(f"Box size: {self.box_size*1e9:.2f} nm")
        print(f"Timestep: {self.dt*1e15:.2f} fs")
        print(f"Total time: {n_steps*self.dt*1e12:.2f} ps")
        print("=" * 60)

        for step in range(n_steps):
            kinetic, potential = self.velocity_verlet_step()

            if step % sample_interval == 0:
                self.kinetic_energy.append(kinetic / 1.602e-19)  # Convert to eV
                self.potential_energy.append(potential / 1.602e-19)
                self.total_energy.append((kinetic + potential) / 1.602e-19)

                if step % (n_steps // 10) == 0:
                    T_inst = 2 * kinetic / (3 * self.n_atoms * kB)
                    print(f"Step {step:6d}: T = {T_inst:6.1f} K, "
                          f"E_total = {(kinetic+potential)/1.602e-19:8.3f} eV")

# Run simulation
md = MDSimulation(n_atoms=100, box_size=5.0, dt=1e-15)  # 5 nm box, 1 fs timestep
md.run(n_steps=5000, sample_interval=10)

# Plot results
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 8))

# Energy evolution
time = np.arange(len(md.total_energy)) * 10 * 1e-15 * 1e12  # ps
ax1.plot(time, md.kinetic_energy, 'r-', label='Kinetic', alpha=0.7)
ax1.plot(time, md.potential_energy, 'b-', label='Potential', alpha=0.7)
ax1.plot(time, md.total_energy, 'k-', label='Total', linewidth=2)
ax1.set_xlabel('Time (ps)')
ax1.set_ylabel('Energy (eV)')
ax1.set_title('Energy Conservation in MD Simulation')
ax1.legend()
ax1.grid(True, alpha=0.3)

# Temperature evolution
temperatures = 2 * np.array(md.kinetic_energy) * 1.602e-19 / (3 * md.n_atoms * kB)
ax2.plot(time, temperatures, 'g-', linewidth=1.5)
ax2.axhline(y=300, color='r', linestyle='--', label='Target (300 K)')
ax2.set_xlabel('Time (ps)')
ax2.set_ylabel('Temperature (K)')
ax2.set_title('Temperature Fluctuations')
ax2.legend()
ax2.grid(True, alpha=0.3)

plt.tight_layout()
plt.savefig('md_simulation.png', dpi=150, bbox_inches='tight')
print("\nSimulation complete! Plot saved as 'md_simulation.png'")

program md_simulation
  implicit none

  ! Physical constants
  real(8), parameter :: kB = 1.381d-23      ! Boltzmann constant (J/K)
  real(8), parameter :: mass = 1.67d-27     ! Particle mass (kg)
  real(8), parameter :: epsilon = 0.996d-21 ! LJ epsilon (J)
  real(8), parameter :: sigma = 3.4d-10     ! LJ sigma (m)

  ! Simulation parameters
  integer, parameter :: n_atoms = 100
  integer, parameter :: n_steps = 5000
  integer, parameter :: sample_interval = 10
  real(8), parameter :: dt = 1.0d-15        ! Timestep (1 fs)
  real(8), parameter :: box_size = 5.0d-9   ! Box size (5 nm)
  real(8), parameter :: cutoff = 2.5d0 * sigma
  real(8), parameter :: target_temp = 300.0d0

  ! Arrays
  real(8), dimension(3, n_atoms) :: positions, velocities, forces, old_forces
  real(8), dimension(n_steps/sample_interval) :: kinetic_E, potential_E, total_E
  real(8) :: kinetic, potential, temperature
  integer :: i, step, sample_idx

  ! Initialize simulation
  call initialize_system()

  ! Main MD loop
  print *, 'Starting Fortran MD simulation'
  print *, 'Atoms:', n_atoms, '  Box size:', box_size*1d9, 'nm'
  print *, 'Timestep:', dt*1d15, 'fs  Total time:', n_steps*dt*1d12, 'ps'
  print *, '============================================================'

  sample_idx = 1
  do step = 0, n_steps-1
    ! Velocity Verlet integration
    call velocity_verlet_step(kinetic, potential)

    ! Sample and output
    if (mod(step, sample_interval) == 0) then
      kinetic_E(sample_idx) = kinetic
      potential_E(sample_idx) = potential
      total_E(sample_idx) = kinetic + potential
      sample_idx = sample_idx + 1

      if (mod(step, n_steps/10) == 0) then
        temperature = 2.0d0 * kinetic / (3.0d0 * n_atoms * kB)
        print '(A,I6,A,F7.1,A,F10.3,A)', &
          'Step', step, ': T =', temperature, ' K, E_total =', &
          (kinetic+potential)/1.602d-19, ' eV'
      end if
    end if
  end do

  print *, ''
  print *, 'Simulation complete!'
  call print_statistics()

contains

  subroutine initialize_system()
    ! Initialize positions randomly in box
    call random_number(positions)
    positions = positions * box_size

    ! Initialize velocities (Maxwell-Boltzmann distribution)
    call initialize_velocities(target_temp)

    ! Calculate initial forces
    call calculate_forces(potential)
  end subroutine initialize_system

  subroutine initialize_velocities(T)
    real(8), intent(in) :: T
    real(8) :: sigma_v, rand_normal
    integer :: i, j

    sigma_v = sqrt(kB * T / mass)

    ! Generate Gaussian random velocities
    do i = 1, n_atoms
      do j = 1, 3
        call random_normal(rand_normal)
        velocities(j, i) = sigma_v * rand_normal
      end do
    end do

    ! Remove center of mass motion
    do j = 1, 3
      velocities(j, :) = velocities(j, :) - sum(velocities(j, :)) / n_atoms
    end do
  end subroutine initialize_velocities

  subroutine random_normal(x)
    ! Box-Muller transform for Gaussian random numbers
    real(8), intent(out) :: x
    real(8) :: u1, u2
    call random_number(u1)
    call random_number(u2)
    x = sqrt(-2.0d0 * log(u1)) * cos(2.0d0 * 3.14159265359d0 * u2)
  end subroutine random_normal

  subroutine calculate_forces(pot_energy)
    real(8), intent(out) :: pot_energy
    real(8), dimension(3) :: rij, force_vec
    real(8) :: r, r2, sr6, sr12, force_mag
    integer :: i, j, k

    forces = 0.0d0
    pot_energy = 0.0d0

    ! Loop over all atom pairs
    do i = 1, n_atoms-1
      do j = i+1, n_atoms
        ! Distance vector with periodic boundary conditions
        rij = positions(:, i) - positions(:, j)

        ! Apply minimum image convention
        do k = 1, 3
          if (rij(k) > box_size/2.0d0) rij(k) = rij(k) - box_size
          if (rij(k) < -box_size/2.0d0) rij(k) = rij(k) + box_size
        end do

        r2 = dot_product(rij, rij)
        r = sqrt(r2)

        ! Only calculate if within cutoff
        if (r < cutoff) then
          sr6 = (sigma / r)**6
          sr12 = sr6 * sr6

          ! Lennard-Jones potential
          pot_energy = pot_energy + 4.0d0 * epsilon * (sr12 - sr6)

          ! Lennard-Jones force
          force_mag = 24.0d0 * epsilon * (2.0d0*sr12 - sr6) / r2
          force_vec = force_mag * rij

          forces(:, i) = forces(:, i) + force_vec
          forces(:, j) = forces(:, j) - force_vec
        end if
      end do
    end do
  end subroutine calculate_forces

  subroutine velocity_verlet_step(kin_energy, pot_energy)
    real(8), intent(out) :: kin_energy, pot_energy
    integer :: i

    ! Update positions: r(t+dt) = r(t) + v(t)*dt + 0.5*a(t)*dt^2
    positions = positions + velocities * dt + 0.5d0 * forces * dt**2 / mass

    ! Apply periodic boundary conditions
    do i = 1, 3
      where (positions(i, :) < 0.0d0) positions(i, :) = positions(i, :) + box_size
      where (positions(i, :) >= box_size) positions(i, :) = positions(i, :) - box_size
    end do

    ! Save old forces
    old_forces = forces

    ! Calculate new forces
    call calculate_forces(pot_energy)

    ! Update velocities: v(t+dt) = v(t) + 0.5*[a(t) + a(t+dt)]*dt
    velocities = velocities + 0.5d0 * (old_forces + forces) * dt / mass

    ! Calculate kinetic energy
    kin_energy = 0.5d0 * mass * sum(velocities**2)
  end subroutine velocity_verlet_step

  subroutine print_statistics()
    real(8) :: avg_temp, avg_total_E, std_total_E
    integer :: n_samples

    n_samples = size(total_E)

    avg_total_E = sum(total_E) / n_samples
    std_total_E = sqrt(sum((total_E - avg_total_E)**2) / n_samples)

    avg_temp = 2.0d0 * sum(kinetic_E) / (3.0d0 * n_atoms * kB * n_samples)

    print *, ''
    print *, 'Performance Statistics:'
    print *, '---------------------'
    print '(A,F7.1,A)', '  Average Temperature: ', avg_temp, ' K'
    print '(A,F10.3,A)', '  Average Total Energy: ', avg_total_E/1.602d-19, ' eV'
    print '(A,F10.5,A)', '  Energy Conservation: ', std_total_E/abs(avg_total_E)*100, ' %'
    print *, ''
    print *, 'Fortran Performance Advantages:'
    print *, '  • Native array operations (SIMD vectorization)'
    print *, '  • Efficient memory layout (column-major)'
    print *, '  • Aggressive compiler optimizations (-O3, -march=native)'
    print *, '  • 5-10× faster than Python for large systems'
  end subroutine print_statistics

end program md_simulation