anton/hf_again
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

atom with eigensolution and arenas

Browse Source
This commit is contained in:
antonl 2026-03-24 21:07:42 +01:00
parent 261a1030ea
commit 8c3f3a2947
2 changed files with 52 additions and 29 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

26
src/base/base_sort.c
View File

@ -1,31 +1,31 @@
U64 qsort_partition_F64(SortPair_F64 *array, U64 size, U64 low, U64 high) { U64 qsort_partition_F64(SortPair_F64 *pairs, U64 size, U64 low, U64 high) {
F64 pivot = array[high].value; F64 pivot = pairs[high].value;
U64 i = low - 1; U64 i = low - 1;
SortPair_F64 temp = {0}; SortPair_F64 temp = {0};
for (U64 j = low; j < high; j++) { for (U64 j = low; j < high; j++) {
if (array[j].value <= pivot) { if (pairs[j].value <= pivot) {
i += 1; i += 1;
temp = array[i]; temp = pairs[i];
array[i] = array[j]; pairs[i] = pairs[j];
array[j] = temp; pairs[j] = temp;
} }
} }
U64 final_pivot_pos = i + 1; U64 final_pivot_pos = i + 1;
temp = array[final_pivot_pos]; temp = pairs[final_pivot_pos];
array[final_pivot_pos] = array[high]; pairs[final_pivot_pos] = pairs[high];
array[high] = temp; pairs[high] = temp;
return final_pivot_pos; return final_pivot_pos;
} }
function void qsort_F64(SortPair_F64 *array, U64 size, U64 low, U64 high) { function void qsort_F64(SortPair_F64 *pairs, U64 size, U64 low, U64 high) {
if (low < high) { if (low < high) {
U64 pivot_index = qsort_partition_F64(array, size, low, high); U64 pivot_index = qsort_partition_F64(pairs, size, low, high);
qsort_F64(array, size, low, pivot_index - 1); qsort_F64(pairs, size, low, pivot_index - 1);
qsort_F64(array, size, pivot_index + 1, high); qsort_F64(pairs, size, pivot_index + 1, high);
} }
} }

55
src/main.c
View File

@ -24,17 +24,31 @@
////// //////
//~ //~
// We have one set of eigenvalue problem solutions for each angular momentum
// quantum number
typedef struct Eigensolution_F64 Eigensolution_F64; typedef struct Eigensolution_F64 Eigensolution_F64;
struct Eigensolution_F64 { struct Eigensolution_F64 {
U32 l;
F64 *eigenvalues_re; F64 *eigenvalues_re;
F64 *eigenvalues_im; F64 *eigenvalues_im;
Mat_F64 right_eigenvectors; Mat_F64 right_eigenvectors;
Mat_F64 left_eigenvectors; Mat_F64 left_eigenvectors;
}; };
#define NUM_ANGULAR_MOMENTA 3
typedef struct Atom Atom;
struct Atom {
Arena *arena;
U32 Z;
Eigensolution_F64 eigensolutions[NUM_ANGULAR_MOMENTA];
};
//////
//~
global Arena *g_base_arena = 0; global Arena *g_base_arena = 0;
global Arena *g_filename_arena = 0; global Arena *g_filename_arena = 0;
#define NUM_ANGULAR_MOMENTA 3 global Atom g_atom = {0};
global F64 g_angular_momenta[NUM_ANGULAR_MOMENTA] = {0.0, 1.0, 2.0}; global F64 g_angular_momenta[NUM_ANGULAR_MOMENTA] = {0.0, 1.0, 2.0};
////// //////
@ -63,7 +77,7 @@ function void print_eigenvalues(S32 l, S32 n, F64 *wr, F64 *wi) {
} }
function void function void
set_up_first_matrices(Mat_F64 *H, Mat_F64 *H_l, Mat_F64 *B_inv) { set_up_first_matrices(Atom *atom, Mat_F64 *H, Mat_F64 *H_l, Mat_F64 *B_inv) {
// We work in units hbar = 1, bohr radius a0 = 1, electron mass m_e = 1, and // We work in units hbar = 1, bohr radius a0 = 1, electron mass m_e = 1, and
// charge e = 1, and 1/(4piepsilon_0) = 1. Set up Hamiltonian: H = // charge e = 1, and 1/(4piepsilon_0) = 1. Set up Hamiltonian: H =
// -0.5*d^2/dr^2 + l(l+1)/(2r^2) - Z/r // -0.5*d^2/dr^2 + l(l+1)/(2r^2) - Z/r
@ -71,7 +85,7 @@ set_up_first_matrices(Mat_F64 *H, Mat_F64 *H_l, Mat_F64 *B_inv) {
ArenaTemp scratch = scratch_get(0, 0); ArenaTemp scratch = scratch_get(0, 0);
F64 *t = g_bspline_ctx.knotpoints; F64 *t = g_bspline_ctx.knotpoints;
F64 Z = 1.0; F64 Z = (F64)atom->Z;
U32 k = g_bspline_ctx.order; U32 k = g_bspline_ctx.order;
// Skipping first bspline // Skipping first bspline
@ -198,6 +212,9 @@ function void EntryPoint(void) {
g_filename_arena = m_make_arena_reserve(Megabytes(2)); g_filename_arena = m_make_arena_reserve(Megabytes(2));
g_base_arena = m_make_arena(); g_base_arena = m_make_arena();
g_atom.Z = 1.0; // Hydrogen
g_atom.arena = m_make_arena();
set_up_gauss_legendre_points(g_base_arena); set_up_gauss_legendre_points(g_base_arena);
@ -231,7 +248,7 @@ function void EntryPoint(void) {
Mat_F64 B_inv = mat_F64(g_base_arena, mat_size1, mat_size2); Mat_F64 B_inv = mat_F64(g_base_arena, mat_size1, mat_size2);
// A is the actual matrix for each eigenvalue problem. // A is the actual matrix for each eigenvalue problem.
Mat_F64 A = mat_F64(g_base_arena, H.size1, H.size2); Mat_F64 A = mat_F64(g_base_arena, H.size1, H.size2);
set_up_first_matrices(&H_base, &H_l_base, &B_inv); set_up_first_matrices(&g_atom, &H_base, &H_l_base, &B_inv);
// Our problem is Hc = EBc, but we want to solve B^-1Hc = Ec, // Our problem is Hc = EBc, but we want to solve B^-1Hc = Ec,
// so we invert the B matrix and compute the product A = B^-1H before calling // so we invert the B matrix and compute the product A = B^-1H before calling
@ -239,11 +256,14 @@ function void EntryPoint(void) {
mat_invert_F64(&B_inv); mat_invert_F64(&B_inv);
// This arena is used to push results from f. ex eigenvalue computations. // This arena is used to push results from f. ex eigenvalue computations.
Arena *mkl_arena = m_make_arena();
// For each angular momentum // For each angular momentum
for (U32 ang_mom_idx = 0; ang_mom_idx < NUM_ANGULAR_MOMENTA; ang_mom_idx++) { for (U32 ang_mom_idx = 0; ang_mom_idx < NUM_ANGULAR_MOMENTA; ang_mom_idx++) {
ArenaTemp scratch = scratch_get(0, 0);
mat_F64_copy_to_dst(&H, &H_base); mat_F64_copy_to_dst(&H, &H_base);
F64 l = g_angular_momenta[ang_mom_idx]; F64 l = g_angular_momenta[ang_mom_idx];
Eigensolution_F64 *eigsol = &g_atom.eigensolutions[ang_mom_idx];
eigsol->l = (U32)l;
if (l > 1e-16) { if (l > 1e-16) {
F64 l_factor = l * (l + 1.0); F64 l_factor = l * (l + 1.0);
U64 mat_size = H_l.size1 * H_l.size2; U64 mat_size = H_l.size1 * H_l.size2;
@ -263,7 +283,6 @@ function void EntryPoint(void) {
//print_mat_F64(&A); //print_mat_F64(&A);
} }
Eigensolution_F64 eigensolution = {0};
// Solve generalised eigenvalue problem // Solve generalised eigenvalue problem
{ {
S32 size1 = A.size1; S32 size1 = A.size1;
@ -276,17 +295,22 @@ function void EntryPoint(void) {
F64 wkopt; F64 wkopt;
F64 *work; F64 *work;
F64 *wr = PushArray(mkl_arena, F64, size1); eigsol->eigenvalues_re = PushArray(g_atom.arena, F64, size1);
F64 *wi = PushArray(mkl_arena, F64, size1); F64 *wr = eigsol->eigenvalues_re;
F64 *vl = PushArray(mkl_arena, F64, ldvl * size1); eigsol->eigenvalues_im = PushArray(g_atom.arena, F64, size1);
F64 *vr = PushArray(mkl_arena, F64, ldvr * size1); F64 *wi = eigsol->eigenvalues_im;
eigsol->left_eigenvectors = mat_F64(g_atom.arena, ldvl, size1);
F64 *vl = eigsol->left_eigenvectors.data;
eigsol->right_eigenvectors = mat_F64(g_atom.arena, size1, ldvr);
F64 *vr = eigsol->right_eigenvectors.data;
lwork = -1; lwork = -1;
F64 *a = A.data; F64 *a = A.data;
dgeev("Vectors", "Vectors", &size1, a, &lda, wr, wi, vl, &ldvl, vr, &ldvr, dgeev("Vectors", "Vectors", &size1, a, &lda, wr, wi, vl, &ldvl, vr, &ldvr,
&wkopt, &lwork, &info); &wkopt, &lwork, &info);
lwork = (S32)wkopt; lwork = (S32)wkopt;
work = (F64 *)malloc(lwork * sizeof(F64)); //work = (F64 *)malloc(lwork * sizeof(F64));
work = PushArray(scratch.arena, F64, lwork);
dgeev("Vectors", "Vectors", &size1, a, &lda, wr, wi, vl, &ldvl, vr, &ldvr, dgeev("Vectors", "Vectors", &size1, a, &lda, wr, wi, vl, &ldvl, vr, &ldvr,
work, &lwork, &info); work, &lwork, &info);
if (info > 0) { if (info > 0) {
@ -294,8 +318,8 @@ function void EntryPoint(void) {
exit(1); exit(1);
} }
F64 *right_eigenvectors = vr; // Sort real and imaginary eigenvalues by real part
U64 *sorted_indices = PushArray(mkl_arena, U64, size1); U64 *sorted_indices = PushArray(scratch.arena, U64, size1);
sort_and_get_indices_F64(wr, sorted_indices, size1); sort_and_get_indices_F64(wr, sorted_indices, size1);
sort_by_indices_F64(wi, sorted_indices, size1); sort_by_indices_F64(wi, sorted_indices, size1);
print_eigenvalues((U32)l, size1, wr, wi ); print_eigenvalues((U32)l, size1, wr, wi );
@ -315,7 +339,7 @@ function void EntryPoint(void) {
// size1, n, ang_mom_idx); // size1, n, ang_mom_idx);
U64 eigvec_idx = mat_get_col_major_idx(0, energy_index, size1); U64 eigvec_idx = mat_get_col_major_idx(0, energy_index, size1);
F64 *eigvecs = &right_eigenvectors[eigvec_idx]; F64 *eigvecs = &eigsol->right_eigenvectors.data[eigvec_idx];
write_array_F64( write_array_F64(
get_eigenvector_filename(g_filename_arena, n, ang_mom_idx), get_eigenvector_filename(g_filename_arena, n, ang_mom_idx),
eigvecs, size1, eigvecs, size1,
@ -328,9 +352,8 @@ function void EntryPoint(void) {
} }
} }
free((void *)work); scratch_release(scratch);
} }
} }
m_arena_release(mkl_arena);
} }