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) {
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);
}
}

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;
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);
}