From 8c3f3a2947f6b7588f8ddb2e48530f1695700c27 Mon Sep 17 00:00:00 2001 From: antonl Date: Tue, 24 Mar 2026 21:07:42 +0100 Subject: [PATCH] atom with eigensolution and arenas --- src/base/base_sort.c | 26 ++++++++++----------- src/main.c | 55 +++++++++++++++++++++++++++++++------------- 2 files changed, 52 insertions(+), 29 deletions(-) diff --git a/src/base/base_sort.c b/src/base/base_sort.c index 33ec3b8..1f92612 100644 --- a/src/base/base_sort.c +++ b/src/base/base_sort.c @@ -1,31 +1,31 @@ -U64 qsort_partition_F64(SortPair_F64 *array, U64 size, U64 low, U64 high) { - F64 pivot = array[high].value; +U64 qsort_partition_F64(SortPair_F64 *pairs, U64 size, U64 low, U64 high) { + F64 pivot = pairs[high].value; U64 i = low - 1; SortPair_F64 temp = {0}; for (U64 j = low; j < high; j++) { - if (array[j].value <= pivot) { + if (pairs[j].value <= pivot) { i += 1; - temp = array[i]; - array[i] = array[j]; - array[j] = temp; + temp = pairs[i]; + pairs[i] = pairs[j]; + pairs[j] = temp; } } U64 final_pivot_pos = i + 1; - temp = array[final_pivot_pos]; - array[final_pivot_pos] = array[high]; - array[high] = temp; + temp = pairs[final_pivot_pos]; + pairs[final_pivot_pos] = pairs[high]; + pairs[high] = temp; 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) { - U64 pivot_index = qsort_partition_F64(array, size, low, high); - qsort_F64(array, size, low, pivot_index - 1); - qsort_F64(array, size, pivot_index + 1, high); + U64 pivot_index = qsort_partition_F64(pairs, size, low, high); + qsort_F64(pairs, size, low, pivot_index - 1); + qsort_F64(pairs, size, pivot_index + 1, high); } } diff --git a/src/main.c b/src/main.c index 87f249c..824c05b 100644 --- a/src/main.c +++ b/src/main.c @@ -24,17 +24,31 @@ ////// //~ +// We have one set of eigenvalue problem solutions for each angular momentum +// quantum number typedef struct Eigensolution_F64 Eigensolution_F64; struct Eigensolution_F64 { + U32 l; F64 *eigenvalues_re; F64 *eigenvalues_im; Mat_F64 right_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_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}; ////// @@ -63,7 +77,7 @@ function void print_eigenvalues(S32 l, S32 n, F64 *wr, F64 *wi) { } 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 // 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 @@ -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); F64 *t = g_bspline_ctx.knotpoints; - F64 Z = 1.0; + F64 Z = (F64)atom->Z; U32 k = g_bspline_ctx.order; // Skipping first bspline @@ -198,6 +212,9 @@ function void EntryPoint(void) { g_filename_arena = m_make_arena_reserve(Megabytes(2)); 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); @@ -231,7 +248,7 @@ function void EntryPoint(void) { Mat_F64 B_inv = mat_F64(g_base_arena, mat_size1, mat_size2); // A is the actual matrix for each eigenvalue problem. 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, // 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); // This arena is used to push results from f. ex eigenvalue computations. - Arena *mkl_arena = m_make_arena(); // For each angular momentum 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); 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) { F64 l_factor = l * (l + 1.0); U64 mat_size = H_l.size1 * H_l.size2; @@ -263,7 +283,6 @@ function void EntryPoint(void) { //print_mat_F64(&A); } - Eigensolution_F64 eigensolution = {0}; // Solve generalised eigenvalue problem { S32 size1 = A.size1; @@ -276,17 +295,22 @@ function void EntryPoint(void) { F64 wkopt; F64 *work; - F64 *wr = PushArray(mkl_arena, F64, size1); - F64 *wi = PushArray(mkl_arena, F64, size1); - F64 *vl = PushArray(mkl_arena, F64, ldvl * size1); - F64 *vr = PushArray(mkl_arena, F64, ldvr * size1); + eigsol->eigenvalues_re = PushArray(g_atom.arena, F64, size1); + F64 *wr = eigsol->eigenvalues_re; + eigsol->eigenvalues_im = PushArray(g_atom.arena, F64, 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; F64 *a = A.data; dgeev("Vectors", "Vectors", &size1, a, &lda, wr, wi, vl, &ldvl, vr, &ldvr, &wkopt, &lwork, &info); 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, work, &lwork, &info); if (info > 0) { @@ -294,8 +318,8 @@ function void EntryPoint(void) { exit(1); } - F64 *right_eigenvectors = vr; - U64 *sorted_indices = PushArray(mkl_arena, U64, size1); + // Sort real and imaginary eigenvalues by real part + U64 *sorted_indices = PushArray(scratch.arena, U64, size1); sort_and_get_indices_F64(wr, sorted_indices, size1); sort_by_indices_F64(wi, sorted_indices, size1); print_eigenvalues((U32)l, size1, wr, wi ); @@ -315,7 +339,7 @@ function void EntryPoint(void) { // size1, n, ang_mom_idx); 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( get_eigenvector_filename(g_filename_arena, n, ang_mom_idx), eigvecs, size1, @@ -328,9 +352,8 @@ function void EntryPoint(void) { } } - free((void *)work); + scratch_release(scratch); } } - m_arena_release(mkl_arena); }