/// <summary>
/// Constructor. Initializes a Visscher wavefunction from an ordinary wavefunction.
/// </summary>
public VisscherWf(WaveFunction inputWf, float[][][] V, float mass, float dt, bool multiThread = true)
{
int sx = inputWf.GridSpec.SizeX;
int sy = inputWf.GridSpec.SizeY;
int sz = inputWf.GridSpec.SizeZ;
LatticeSpacing = inputWf.LatticeSpacing;
// Allocate the arrays
RealPart = TdseUtils.Misc.Allocate3DArray(sz, sy, sx);
ImagPartM = TdseUtils.Misc.Allocate3DArray(sz, sy, sx);
ImagPartP = TdseUtils.Misc.Allocate3DArray(sz, sy, sx);
// For the real part, just copy the values from the input wavefunction.
for (int z = 0; z < sz; z++)
{
for (int y = 0; y < sy; y++)
{
float[] realPartZY = RealPart[z][y];
float[] inputWfZY = inputWf.Data[z][y];
for (int x = 0; x < sx; x++)
{
realPartZY[x] = inputWfZY[2 * x];
}
}
}
// For the imaginary parts, we need to compute the time evolutions.
// We use a power series expansion of the time-evolution operator, accurate to 2nd order in H*dt
WaveFunction H_PsiIn = inputWf.ApplyH(V, mass, multiThread);
WaveFunction H2_PsiIn = H_PsiIn.ApplyH(V, mass, multiThread);
float halfDt = dt / 2;
float eighthDt2 = dt * dt / 8;
TdseUtils.Misc.ForLoop(0, sz, z =>
{
for (int y = 0; y < sy; y++)
{
float[] imPZY = this.ImagPartP[z][y];
float[] imMZY = this.ImagPartM[z][y];
float[] inputWfZY = inputWf.Data[z][y];
float[] H_PsiInZY = H_PsiIn.Data[z][y];
float[] H2_PsiInZY = H2_PsiIn.Data[z][y];
for (int x = 0; x < sx; x++)
{
int x2 = 2 * x;
int x2p1 = x2 + 1;
float dt0Term = inputWfZY[x2p1];
float dt1Term = halfDt * H_PsiInZY[x2];
float dt2Term = eighthDt2 * H2_PsiInZY[x2p1];
imPZY[x] = dt0Term - dt1Term - dt2Term; // [1 - i*H*(dt/2) - (1/2)H^2*(dt/2)^2] * Psi
imMZY[x] = dt0Term + dt1Term - dt2Term; // [1 + i*H*(dt/2) - (1/2)H^2*(dt/2)^2] * Psi
}
}
}, multiThread);
}
/// <summary>
/// Creates a Gaussian wavepacket with given properties.
/// </summary>
public static WaveFunction CreateGaussianWavePacket(
GridSpec gridSpec, float latticeSpacing, bool originAtLatticeCenter, float mass,
Vec3 packetCenter, Vec3 packetWidth, Vec3 avgMomentum, bool multiThread = true)
{
WaveFunction wf = new WaveFunction(gridSpec, latticeSpacing);
int sx = gridSpec.SizeX;
int sy = gridSpec.SizeY;
int sz = gridSpec.SizeZ;
float[][][] wfData = wf.Data;
Complex I = Complex.I;
float rootPi = (float)Math.Sqrt(Math.PI);
float sigmaXSq = packetWidth.X * packetWidth.X;
float sigmaYSq = packetWidth.Y * packetWidth.Y;
float sigmaZSq = packetWidth.Z * packetWidth.Z;
float norm = (float)Math.Sqrt(1.0 / (rootPi * packetWidth.X * rootPi * packetWidth.Y * rootPi * packetWidth.Z));
int halfGridSizeX = (sx - 1) / 2;
int halfGridSizeY = (sy - 1) / 2;
int halfGridSizeZ = (sz - 1) / 2;
TdseUtils.Misc.ForLoop(0, sz, (z) =>
{
float zf = (originAtLatticeCenter) ? (z - halfGridSizeZ) * latticeSpacing : (z * latticeSpacing);
Complex expArgZ = I * zf * avgMomentum.Z - (zf - packetCenter.Z) * (zf - packetCenter.Z) / (2 * sigmaZSq);
for (int y = 0; y < sy; y++)
{
float yf = (originAtLatticeCenter) ? (y - halfGridSizeY) * latticeSpacing : (y * latticeSpacing);
Complex expArgZY = expArgZ + I * yf * avgMomentum.Y - (yf - packetCenter.Y) * (yf - packetCenter.Y) / (2 * sigmaYSq);
float[] wfDataZY = wf.Data[z][y];
for (int x = 0; x < sx; x++)
{
float xf = (originAtLatticeCenter) ? (x - halfGridSizeX) * latticeSpacing : (x * latticeSpacing);
Complex expArgZYX = expArgZY + I * xf * avgMomentum.X - (xf - packetCenter.X) * (xf - packetCenter.X) / (2 * sigmaXSq);
Complex wfVal = norm * Complex.Exp(expArgZYX);
wfDataZY[2 * x] = wfVal.Re;
wfDataZY[2 * x + 1] = wfVal.Im;
}
}
}, multiThread);
wf.Normalize();
return(wf);
}
/// <summary>
/// Converts a Visscher wavefunction to a regular wavefunction.
/// </summary>
public WaveFunction ToRegularWavefunction(bool multiThread = true)
{
int sx = GridSpec.SizeX;
int sy = GridSpec.SizeY;
int sz = GridSpec.SizeZ;
WaveFunction result = new WaveFunction(GridSpec, LatticeSpacing);
TdseUtils.Misc.ForLoop(0, sz, z =>
{
for (int y = 0; y < sy; y++)
{
float[] thisRzy = RealPart[z][y];
float[] thisIMzy = ImagPartM[z][y];
float[] thisIPzy = ImagPartP[z][y];
float[] outWfzy = result.Data[z][y];
for (int x = 0; x < sx; x++)
{
outWfzy[2 * x] = thisRzy[x];
// Take the square root of Im(t-dt/2) * I(t+dt/2), as this is the quantity that gives a conserved probability.
float IM = thisIMzy[x];
float IP = thisIPzy[x];
float Iproduct = IM * IP;
float Iout;
if (Iproduct > 0.0f)
{
Iout = (float)Math.Sqrt(Iproduct);
if (IM < 0.0f)
{
Iout = -Iout;
}
}
else
{
Iout = 0.0f;
}
outWfzy[2 * x + 1] = Iout;
}
}
}, multiThread);
return(result);
}
/// <summary>
/// Worker method.
/// </summary>
protected override void WorkerMethod()
{
// Reset counters
m_currentTimeStepIndex = 0;
m_lastSavedFrame = -1;
// Precompute the potential everywhere on the grid
float[][][] V = PrecomputeV();
// Create the initial relative wavefunction
Vec3 r0 = m_initialPosition1 - m_initialPosition2;
Vec3 p0 = (m_mass2 / m_totalMass) * m_initialMomentum1 - (m_mass1 / m_totalMass) * m_initialMomentum2;
int sx = 2 * m_gridSizeX - 1; // The range of r = r1-r2 is twice the range of r1, r2
int sy = 2 * m_gridSizeY - 1;
int sz = 2 * m_gridSizeZ - 1;
WaveFunction initialWf = WaveFunctionUtils.CreateGaussianWavePacket(
new GridSpec(sx, sy, sz), m_latticeSpacing, true, m_reducedMass, r0, m_sigmaRel, p0, m_multiThread
);
m_visscherWf = new VisscherWf(initialWf, V, m_reducedMass, m_deltaT, m_multiThread);
initialWf = null; // Allow initialWf to be garbage collected
TimeStepCompleted();
if (IsCancelled)
{
return;
}
// Main loop: Evolve the relative wavefunction
while (m_currentTimeStepIndex < m_totalNumTimeSteps)
{
// Evolve the wavefunction by one timestep
EvolveByOneTimeStep(m_visscherWf, V);
m_currentTimeStepIndex++;
// Report progress to the client
TimeStepCompleted();
if (IsCancelled)
{
return;
}
}
}
/// <summary>
/// Computes the result of applying a given Hamiltonian operator to this wavefunction.
/// </summary>
public WaveFunction ApplyH(float[][][] V, float mass, bool multiThread = true)
{
// Initialize locals
int sx = GridSpec.SizeX;
int sy = GridSpec.SizeY;
int sz = GridSpec.SizeZ;
int sx2 = 2 * sx;
int sx2m2 = sx2 - 2;
int sym1 = sy - 1;
int szm1 = sz - 1;
float keFactor = (1.0f / 12.0f) / (2 * mass * m_latticeSpacing * m_latticeSpacing);
WaveFunction outWf = new WaveFunction(GridSpec, m_latticeSpacing);
// Compute H * Wf
TdseUtils.Misc.ForLoop(0, sz, z =>
{
int zp = (z < szm1) ? z + 1 : 0;
int zpp = (zp < szm1) ? zp + 1 : 0;
int zm = (z > 0) ? z - 1 : szm1;
int zmm = (zm > 0) ? zm - 1 : szm1;
for (int y = 0; y < sy; y++)
{
int yp = (y < sym1) ? y + 1 : 0;
int ypp = (yp < sym1) ? yp + 1 : 0;
int ym = (y > 0) ? y - 1 : sym1;
int ymm = (ym > 0) ? ym - 1 : sym1;
float[] inWf_z_y = m_data[z][y];
float[] inWf_z_ym = m_data[z][ym];
float[] inWf_z_yp = m_data[z][yp];
float[] inWf_z_ymm = m_data[z][ymm];
float[] inWf_z_ypp = m_data[z][ypp];
float[] inWf_zm_y = m_data[zm][y];
float[] inWf_zp_y = m_data[zp][y];
float[] inWf_zmm_y = m_data[zmm][y];
float[] inWf_zpp_y = m_data[zpp][y];
float[] outWf_z_y = outWf.m_data[z][y];
float[] V_z_y = V[z][y];
for (int rx = 0; rx < sx2; rx += 2)
{
int rxp = (rx < sx2m2) ? rx + 2 : 0;
int rxpp = (rxp < sx2m2) ? rxp + 2 : 0;
int rxm = (rx > 0) ? rx - 2 : sx2m2;
int rxmm = (rxm > 0) ? rxm - 2 : sx2m2;
int x = rx / 2;
// Kinetic energy terms.
float kR = keFactor * (
90.0f * inWf_z_y[rx] -
16.0f * (inWf_zm_y[rx] + inWf_zp_y[rx] + inWf_z_yp[rx] + inWf_z_ym[rx] + inWf_z_y[rxm] + inWf_z_y[rxp]) +
(inWf_zmm_y[rx] + inWf_zpp_y[rx] + inWf_z_ypp[rx] + inWf_z_ymm[rx] + inWf_z_y[rxmm] + inWf_z_y[rxpp])
);
int ix = rx + 1;
float kI = keFactor * (
90.0f * inWf_z_y[ix] -
16.0f * (inWf_zm_y[ix] + inWf_zp_y[ix] + inWf_z_yp[ix] + inWf_z_ym[ix] + inWf_z_y[rxm + 1] + inWf_z_y[rxp + 1]) +
(inWf_zmm_y[ix] + inWf_zpp_y[ix] + inWf_z_ypp[ix] + inWf_z_ymm[ix] + inWf_z_y[rxmm + 1] + inWf_z_y[rxpp + 1])
);
// Potential energy terms
float vR = V_z_y[x] * inWf_z_y[rx];
float vI = V_z_y[x] * inWf_z_y[ix];
outWf_z_y[rx] = kR + vR;
outWf_z_y[ix] = kI + vI;
}
}
}, multiThread);
return(outWf);
}