Skip to content

Instantly share code, notes, and snippets.

@marl0ny

marl0ny/eigenstates_1d.cpp

Created February 7, 2022 14:03
Show Gist options
  • Save marl0ny/4d2aa59e9a22f187430147a9624b73dc to your computer and use it in GitHub Desktop.
Save marl0ny/4d2aa59e9a22f187430147a9624b73dc to your computer and use it in GitHub Desktop.
Download ZIP
Raw
eigenstates_1d.cpp
/*
Finding the energy eigenstates for a single particle 1D system with
Dirichlet boundary conditions,
by discretizing the Hamiltonian using finite differences.
This requires that you have Eigen, Spectra, and SDL installed.
Eigen: https://eigen.tuxfamily.org/index.php?title=Main_Page
Spectra: https://spectralib.org/
SDL: https://www.libsdl.org/
Compile with:
clang++ -c eigenstates_1d.cpp -std=c++17 -O3 -Wall -g
clang++ c++ -std=c++17 -O3 -Wall -g -o program eigenstates_1d.o
You will need to also need to specify the path to Eigen, Spectra, and SDL as well.
*/
#include <complex>
#include <iostream>
#include <stdlib.h>
#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 <SDL2/SDL.h>
#define WIDTH 640
#define HEIGHT 480
using namespace Eigen;
using namespace Spectra;
using complex = std::complex<double>;
int main() {
const size_t size = 256;
float dx = 1.0/size;
SparseMatrix <double> mat(size, size);
mat.reserve(VectorXi::Constant(size, 3));
for (int i = 0; i < size; i++) {
double x = (double)((double)i - (double)size/2.0)/(double)size;
if (i < size-1) mat.insert(i, i+1) = -0.5/(dx*dx);
mat.insert(i, i) = 1.0/(dx*dx) + 15000.0*x*x;
if (i >= 1) mat.insert(i, i-1) = -0.5/(dx*dx);
}
mat.makeCompressed();
int total = 50;
// SparseGenMatProd<double> op(mat);
// SymEigsSolver<SparseGenMatProd<double>> eigs(op, total, total*10);
// eigs.init();
// int nconv = eigs.compute(SortRule::SmallestMagn);
SparseSymShiftSolve<double> op(mat);
SymEigsShiftSolver<SparseSymShiftSolve<double>> eigs(op, total, total + 10, 0.0);
eigs.init();
int nconv = eigs.compute(SortRule::LargestMagn);
std::cout << nconv << std::endl;
VectorXcd eigenvalues;
MatrixXcd eigenvectors;
if (eigs.info() == CompInfo::Successful) {
eigenvalues = eigs.eigenvalues();
eigenvectors = eigs.eigenvectors();
} else {
fprintf(stderr, "Computation failed\n");
exit(1);
}
std::cout << eigenvalues << std::endl;
SDL_Window* window = NULL;
if (SDL_Init(SDL_INIT_VIDEO) < 0) {
fprintf(stderr, "Unable to initialize SDL2: %s\n", SDL_GetError());
exit(1);
}
window = SDL_CreateWindow("Window",
SDL_WINDOWPOS_CENTERED, SDL_WINDOWPOS_CENTERED,
WIDTH, HEIGHT,
SDL_WINDOW_RESIZABLE);
if (window == NULL) {
fprintf(stderr, "Unable to create window: %s\n", SDL_GetError());
exit(1);
}
SDL_Renderer *renderer = SDL_CreateRenderer(window, -1, SDL_RENDERER_ACCELERATED);
if (renderer == NULL) {
fprintf(stderr, "Unable to create renderer: %s\n", SDL_GetError());
exit(1);
}
int window_width = WIDTH;
int window_height = HEIGHT;
SDL_Rect background_rect = {.x = 0, .y = 0, .w = window_width, .h = window_height};
SDL_SetRenderDrawColor(renderer, 100, 100, 100, 255);
SDL_RenderClear(renderer);
const Uint8 *keyboard_state;
// Point p;
int n = total - 1;
int returned = 0;
// int mouse_pressed = 0;
// int mouse_released = 0;
// int mouse_x = 0;
// int mouse_y = 0;
int delay_time = 15;
int key_state = 0;
do {
SDL_SetRenderDrawColor(renderer, 0, 0, 0, 255);
SDL_GetWindowSize(window, &window_width, &window_height);
background_rect.w = window_width;
background_rect.h = window_height;
SDL_RenderFillRect(renderer, &background_rect);
SDL_PumpEvents();
// mouse_released = 0;
keyboard_state = SDL_GetKeyboardState(NULL);
if (keyboard_state[SDL_SCANCODE_RETURN]) {
returned = 1;
}
int key_pressed = 0;
if (keyboard_state[SDL_SCANCODE_UP]) {
key_pressed = 1;
key_state = SDL_SCANCODE_UP;
}
if (keyboard_state[SDL_SCANCODE_DOWN]) {
key_pressed = 1;
key_state = SDL_SCANCODE_DOWN;
}
if (key_pressed == 0) {
if (key_state == SDL_SCANCODE_UP) {
n = (n - 1 < 0)? total - 1: n - 1;
} else if (key_state == SDL_SCANCODE_DOWN) {
n = (n + 1) % total;
}
key_state = 0;
}
SDL_SetRenderDrawColor(renderer, 255, 255, 255, 255);
for (int i = 0; i < size - 1; i++) {
double scale = 100.0*(double)size/201;
SDL_RenderDrawLine(renderer, (int)(window_width*i/size),
window_height/2 - (int)(scale*eigenvectors(i, n).real()),
(int)(window_width*(i + 1)/size),
window_height/2 - (int)(scale*eigenvectors(i+1, n).real()));
}
SDL_SetRenderDrawColor(renderer, 0, 0, 0, 255);
SDL_RenderPresent(renderer);
SDL_Delay(delay_time);
} while(returned == 0);
SDL_DestroyRenderer(renderer);
SDL_DestroyWindow(window);
SDL_Quit();
return 0;
}
Raw
eigenstates_2d.cpp
/*
Finding the energy eigenstates for a single particle 2D system with
Dirichlet boundary conditions,
by discretizing the Hamiltonian using finite differences.
This requires that you have Eigen, Spectra, and SDL installed.
Eigen: https://eigen.tuxfamily.org/index.php?title=Main_Page
Spectra: https://spectralib.org/
SDL: https://www.libsdl.org/
Compile with:
clang++ -c eigenstates_2d.cpp -std=c++17 -O3 -Wall -g
clang++ c++ -std=c++17 -O3 -Wall -g -o program eigenstates_2d.o
You will need to also need to specify the path
to Eigen, Spectra, and SDL as well.
*/
#include <complex>
#include <vector>
#include <iostream>
#include <stdlib.h>
#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 <SDL2/SDL.h>
using namespace Eigen;
using namespace Spectra;
using complex = std::complex<double>;
const size_t size = 256;
const size_t matrix_size = size*size;
const int WIDTH = 512;
const int HEIGHT = 512;
int main() {
std::vector<Triplet<double>> triplets;
triplets.reserve(size*size*5);
double dx = 1.0/size;
for (int i = 0; i < size; i++) {
for (int j = 0; j < size; j++) {
double y = (double)((double)i - (double)size/2.0)/(double)size;
double x = (double)((double)j - (double)size/2.0)/(double)size;
double u = 1.0/(dx*dx), d = 1.0/(dx*dx);
double l = 1.0/(dx*dx), r = 1.0/(dx*dx);
double c = -4.0/(dx*dx);
double potential = 10000.0*(x*x + y*y);
triplets.push_back(Triplet<double>(i*size + j, i*size + j,
-0.5*c + potential));
if (i < size-1) triplets.push_back(
Triplet<double>(i*size + j, (i + 1)*size + j, -0.5*u));
if (i > 0) triplets.push_back(
Triplet<double>(i*size + j, (i - 1)*size + j, -0.5*d));
if (j > 0) triplets.push_back(
Triplet<double>(i*size + j, i*size + (j - 1), -0.5*l));
if (j < size-1) triplets.push_back(
Triplet<double>(i*size + j, i*size + (j + 1), -0.5*r));
}
}
SparseMatrix<double> mat(matrix_size, matrix_size);
mat.reserve(VectorXi::Constant(matrix_size, 5));
mat.setFromTriplets(triplets.begin(), triplets.end());
mat.makeCompressed();
SparseSymShiftSolve<double> op(mat);
int total = 50;
SymEigsShiftSolver<SparseSymShiftSolve<double>> eigs(op, total,
total + 10, 0.0);
eigs.init();
int nconv = eigs.compute(SortRule::LargestMagn);
std::cout << nconv << std::endl;
VectorXcd eigenvalues;
MatrixXcd eigenvectors;
if (eigs.info() == CompInfo::Successful) {
eigenvalues = eigs.eigenvalues();
eigenvectors = eigs.eigenvectors();
} else {
fprintf(stderr, "Computation failed\n");
exit(1);
}
std::cout << eigenvalues << std::endl;
SDL_Window* window = NULL;
if (SDL_Init(SDL_INIT_VIDEO) < 0) {
fprintf(stderr, "Unable to initialize SDL2: %s\n", SDL_GetError());
exit(1);
}
window = SDL_CreateWindow("Window",
SDL_WINDOWPOS_CENTERED, SDL_WINDOWPOS_CENTERED,
WIDTH, HEIGHT,
SDL_WINDOW_RESIZABLE);
if (window == NULL) {
fprintf(stderr, "Unable to create window: %s\n", SDL_GetError());
exit(1);
}
SDL_Renderer *renderer = SDL_CreateRenderer(window, -1,
SDL_RENDERER_ACCELERATED);
if (renderer == NULL) {
fprintf(stderr, "Unable to create renderer: %s\n", SDL_GetError());
exit(1);
}
int window_width = WIDTH;
int window_height = HEIGHT;
SDL_Rect background_rect = {.x = 0, .y = 0,
.w = window_width, .h = window_height};
SDL_SetRenderDrawColor(renderer, 9, 0, 0, 255);
SDL_RenderClear(renderer);
const Uint8 *keyboard_state;
int n = 0;
int returned = 0;
int delay_time = 15;
int key_state = 0;
double scale = 1.0;
do {
SDL_SetRenderDrawColor(renderer, 0, 0, 0, 255);
SDL_GetWindowSize(window, &window_width, &window_height);
background_rect.w = window_width;
background_rect.h = window_height;
SDL_RenderFillRect(renderer, &background_rect);
SDL_PumpEvents();
keyboard_state = SDL_GetKeyboardState(NULL);
if (keyboard_state[SDL_SCANCODE_RETURN]) {
returned = 1;
}
int key_pressed = 0;
if (keyboard_state[SDL_SCANCODE_UP]) {
scale += 1.0;
key_pressed = 1;
key_state = SDL_SCANCODE_UP;
}
if (keyboard_state[SDL_SCANCODE_DOWN]) {
scale -= 1.0;
key_pressed = 1;
key_state = SDL_SCANCODE_DOWN;
}
if (keyboard_state[SDL_SCANCODE_LEFT]) {
key_pressed = 1;
key_state = SDL_SCANCODE_LEFT;
}
if (keyboard_state[SDL_SCANCODE_RIGHT]) {
key_pressed = 1;
key_state = SDL_SCANCODE_RIGHT;
}
if (key_pressed == 0) {
if (key_state == SDL_SCANCODE_RIGHT) {
n = (n - 1 < 0)? total - 1: n - 1;
} else if (key_state == SDL_SCANCODE_LEFT) {
n = (n + 1) % total;
}
key_state = 0;
}
for (int i = 0; i < size; i++) {
for (int j = 0; j < size; j++) {
double real = eigenvectors(i*size + j, n).real();
double imag = eigenvectors(i*size + j, n).imag();
double val = scale*250000.0*(real*real + imag*imag);
SDL_SetRenderDrawColor(renderer,
(int)val, (int)val, (int)val, 255);
SDL_Rect rect;
rect.x = i*2;
rect.y = j*2;
rect.w = 2;
rect.h = 2;
SDL_RenderFillRect(renderer, &rect);
}
}
SDL_RenderPresent(renderer);
SDL_Delay(delay_time);
} while (returned == 0);
SDL_DestroyRenderer(renderer);
SDL_DestroyWindow(window);
SDL_Quit();
return 0;
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment