let dx = 1; // Grid spacing

let dt = 0.05; // Time step to ensure stability (CFL condition)

let V = []; // Potential energy (includes the slits)

let psiRe = []; // Real part of the wave function

let psiIm = []; // Imaginary part of the wave function

let probability = []; // Probability density

let slitWidth = 5; // Width of the slits

let slitDistance = 40; // Distance between the two slits (vertical spacing)

let slitHeight = 20; // Height of each slit (how large the slits are)

​

function setup() {

  createCanvas(gridSize, gridSize);

  pixelDensity(1);

​

  // Initialize wave function and potential

  for (let x = 0; x < gridSize; x++) {

    psiRe[x] = [];

    psiIm[x] = [];

    probability[x] = [];

    V[x] = [];

    for (let y = 0; y < gridSize; y++) {

      psiRe[x][y] = 0;

      psiIm[x][y] = 0;

      V[x][y] = 0; // No potential initially

​

      // Create two slits in the center of the screen (vertically)

      if (

        x == floor(gridSize / 2) &&

        !(

          (y > gridSize / 2 - slitDistance / 2 - slitHeight && y < gridSize / 2 - slitDistance / 2) ||

          (y > gridSize / 2 + slitDistance / 2 && y < gridSize / 2 + slitDistance / 2 + slitHeight)

        )

      ) {

        V[x][y] = 1000; // Set high potential for the barrier except for slit positions

      }

    }

  }

​

  // Initialize a Gaussian wave packet on the left side

  let x0 = 20; // Initial position of the wave packet

  let y0 = gridSize / 2;

  let sigma = 10; // Width of the wave packet

  let kx = 1.5; // Wavenumber

​

  for (let x = 0; x < gridSize; x++) {

    for (let y = 0; y < gridSize; y++) {

      let r = exp(-((x - x0) ** 2 + (y - y0) ** 2) / (2 * sigma ** 2));

      psiRe[x][y] = r * cos(kx * (x - x0));

      psiIm[x][y] = r * sin(kx * (x - x0));

    }

  }

​

  // Check wave packet initialization by printing a sample value

  console.log("Sample psiRe[20][gridSize/2]:", psiRe[20][floor(gridSize / 2)]);

}

​

function draw() {

  background(0);

  updateWaveFunction();

  normalizeWaveFunction(); // Ensuring the wavefunction stays normalized

​

  // Display the probability density

  loadPixels();

  let maxP = getMaxProbability(); // Find the max probability for dynamic scaling

  if (maxP > 0) {

    for (let x = 0; x < gridSize; x++) {

      for (let y = 0; y < gridSize; y++) {

        let p = probability[x][y]; // Probability density

        let index = (x + y * gridSize) * 4;

        // Adjust scaling based on max probability to avoid black screen

        let scaledP = map(p, 0, maxP, 0, 255); // Dynamic scaling based on max probability

        pixels[index] = scaledP; // Red

        pixels[index + 1] = scaledP; // Green

        pixels[index + 2] = scaledP; // Blue

        pixels[index + 3] = 255; // Alpha

      }

    }

  }

  updatePixels();

​

  // Debugging: display max probability value to understand the scale

  console.log("Max Probability:", maxP);

}

​

function updateWaveFunction() {

  // Time evolution of the wave function using finite difference method

  for (let x = 1; x < gridSize - 1; x++) {

    for (let y = 1; y < gridSize - 1; y++) {

      if (V[x][y] != 1000) {

        // Discretized Schrödinger equation

        let laplacianRe =

          psiRe[x + 1][y] +

          psiRe[x - 1][y] +

          psiRe[x][y + 1] +

          psiRe[x][y - 1] -

          4 * psiRe[x][y];

​

        let laplacianIm =

          psiIm[x + 1][y] +

          psiIm[x - 1][y] +

          psiIm[x][y + 1] +

          psiIm[x][y - 1] -

          4 * psiIm[x][y];

​

        // Update real and imaginary parts

        psiRe[x][y] += (-dt * laplacianIm / (dx * dx) + dt * V[x][y] * psiIm[x][y]);

        psiIm[x][y] += (dt * laplacianRe / (dx * dx) - dt * V[x][y] * psiRe[x][y]);

​

        // Calculate the probability density

        let prob = psiRe[x][y] ** 2 + psiIm[x][y] ** 2;

        probability[x][y] = isNaN(prob) ? 0 : prob; // Check for NaN

      }

    }

  }

}

​

function normalizeWaveFunction() {

  let sum = 0;

  for (let x = 0; x < gridSize; x++) {

    for (let y = 0; y < gridSize; y++) {

      sum += psiRe[x][y] ** 2 + psiIm[x][y] ** 2; // Sum of probabilities

    }

  }

​

  let normFactor = sqrt(sum);

​

  for (let x = 0; x < gridSize; x++) {

    for (let y = 0; y < gridSize; y++) {

      psiRe[x][y] /= normFactor;

      psiIm[x][y] /= normFactor;

    }

  }

}

​

// Debug function to find the maximum probability value

function getMaxProbability() {

  let maxP = 0;

  for (let x = 0; x < gridSize; x++) {

    for (let y = 0; y < gridSize; y++) {

      if (probability[x][y] > maxP) {

        maxP = probability[x][y];

      }

    }

  }

  return maxP;

}

​