Skip to content

Instantly share code, notes, and snippets.

@marl0ny

marl0ny/energy_eigenstates3d.cpp

Created November 29, 2025 15:42
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save marl0ny/ccfc6d5d2032e56d0e985f61fbf1ebd7 to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save marl0ny/ccfc6d5d2032e56d0e985f61fbf1ebd7 to your computer and use it in GitHub Desktop.
Download ZIP
Raw
energy_eigenstates3d.cpp
/* Numerically find the lowest energy bound states for a potential in 3D.
This uses a uniform cartesian grid finite difference approximation, where a
second-order accurate stencil is used for the Laplacian.
Dirichlet boundary conditions are imposed.
The eigenvalue problem is solved using Spectra (https://spectralib.org/),
and Eigen (https://eigen.tuxfamily.org/index.php?title=Main_Page).
If Spectra and Eigen are in the same directory as this file, compile with:
clang++ -std=c++17 -O3 -I${PWD} -c energy_eigenstates3d.cpp;
clang++ energy_eigenstates3d.o -o ees3d
Run the program with:
`./ees3d`
*/
#include <Eigen/Core>
#include <Eigen/Sparse>
#include <Spectra/SymEigsSolver.h>
#include <Spectra/SymEigsShiftSolver.h>
#include <Spectra/MatOp/SparseSymShiftSolve.h>
#include <Spectra/MatOp/SparseGenMatProd.h>
#include <iostream>
#define FPTYPE float
using complex = std::complex<FPTYPE>;
using namespace Eigen;
using namespace Spectra;
struct Settings {
enum {
LAPLACIAN_2ND_OR_7_PT=0,
};
FPTYPE hbar, m;
FPTYPE width, height, length;
FPTYPE global_shift;
int grid_width, grid_height, grid_length;
int n_states;
int laplacian_stencil;
};
static void construct_hamiltonian_2nd_order_7pt(
std::vector<Triplet<FPTYPE>> &triplets,
std::vector<FPTYPE> &potential,
Settings &settings
) {
triplets.reserve(
settings.grid_width*settings.grid_height*settings.grid_length*7);
FPTYPE dx = settings.width/FPTYPE(settings.grid_width);
FPTYPE dy = settings.height/FPTYPE(settings.grid_height);
FPTYPE dz = settings.length/FPTYPE(settings.grid_length);
FPTYPE hbar = settings.hbar;
FPTYPE m = settings.m;
FPTYPE global_shift = settings.global_shift;
int grid_width = settings.grid_width;
int grid_height = settings.grid_height;
int grid_length = settings.grid_length;
#define AT(i, j, k) (i)*grid_height*grid_width + (j)*grid_width + (k)
for (int i = 0; i < grid_length; i++) {
for (int j = 0; j < grid_height; j++) {
for (int k = 0; k < grid_width; k++) {
FPTYPE dx2 = dx*dx, dy2 = dy*dy, dz2 = dz*dz;
FPTYPE hbar2_2m = hbar*hbar/(2.0*m);
FPTYPE potential_val = potential[AT(i, j, k)] + global_shift;
triplets.push_back(
{
AT(i, j, k), AT(i, j, k),
-2.0*(-hbar2_2m/dx2 - hbar2_2m/dy2 - hbar2_2m/dz2)
+ potential_val
}
);
if (i < grid_length - 1)
triplets.push_back(
{AT(i, j, k), AT(i+1, j, k), -hbar2_2m/dz2});
if (i > 0)
triplets.push_back(
{AT(i, j, k), AT(i-1, j, k), -hbar2_2m/dz2});
if (j < grid_height - 1)
triplets.push_back(
{AT(i, j, k), AT(i, j+1, k), -hbar2_2m/dy2});
if (j > 0)
triplets.push_back(
{AT(i, j, k), AT(i, j-1, k), -hbar2_2m/dy2});
if (k < grid_height - 1)
triplets.push_back(
{AT(i, j, k), AT(i, j, k+1), -hbar2_2m/dx2});
if (k > 0)
triplets.push_back(
{AT(i, j, k), AT(i, j, k-1), -hbar2_2m/dx2});
}
}
}
#undef AT
}
int main(int argc, char **argv) {
Settings settings = {
.hbar=1.0, .m=1.0,
.width=10.0, .height=10.0, .length=10.0,
.global_shift=0.0,
.grid_width=32, .grid_height=32, .grid_length=32,
.n_states=4,
.laplacian_stencil=Settings::LAPLACIAN_2ND_OR_7_PT
};
// For a grid size of 50^3, expect to wait almost two minutes to find
// the lowest four energy eigenvalues.
if (argc > 1) {
settings.grid_width = std::stoi(argv[1]);
settings.grid_height = settings.grid_width;
settings.grid_length = settings.grid_width;
}
if (argc > 2)
settings.n_states = std::stoi(argv[2]);
size_t matrix_size
= settings.grid_width*settings.grid_height*settings.grid_length;
std::vector<Triplet<FPTYPE>> triplets;
std::vector<FPTYPE> potential(matrix_size, 0.0);
int grid_width = settings.grid_width, grid_height = settings.grid_height,
grid_length = settings.grid_length;
FPTYPE width = settings.width, height = settings.height,
length = settings.length;
#define AT(i, j, k) (i)*grid_height*grid_width + (j)*grid_width + (k)
for (int i = 0; i < settings.grid_length; i++) {
for (int j = 0; j < settings.grid_height; j++) {
for (int k = 0; k < settings.grid_width; k++) {
FPTYPE x = width*(k + 0.5)/grid_width - width/2.0;
FPTYPE y = height*(j + 0.5)/grid_height - height/2.0;
FPTYPE z = length*(i + 0.5)/grid_length - length/2.0;
potential[AT(i, j, k)] = x*x + y*y + z*z;
}
}
}
#undef AT
construct_hamiltonian_2nd_order_7pt(triplets, potential, settings);
size_t reserve_size = 7;
SparseMatrix<FPTYPE> mat(matrix_size, matrix_size);
mat.reserve(VectorXi::Constant(matrix_size, reserve_size));
mat.setFromTriplets(triplets.begin(), triplets.end());
mat.makeCompressed();
SparseSymShiftSolve<FPTYPE> op(mat);
SymEigsShiftSolver<SparseSymShiftSolve<FPTYPE>> eigs(
op, settings.n_states,
int(settings.n_states*1.25), 0.0);
eigs.init();
int nconv = eigs.compute(SortRule::LargestMagn);
VectorX<complex> eigenvalues;
if (eigs.info() == CompInfo::Successful) {
eigenvalues = eigs.eigenvalues();
} else {
fprintf(stderr, "Unable to compute eigenvalues and eigenvectors.\n");
}
std::cout <<
"Numerical computation of the lowest energies for "
"V(x, y, z) = x^2 + y^2 + z^2:"
<< std::endl;
for (auto &e: eigenvalues)
std::cout << e << std::endl;
std::cout << "nconv: " << nconv << std::endl;
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment