some refactor of main
This commit is contained in:
parent
64e28327a1
commit
da72d26496
320
src/hf/hf_base.c
320
src/hf/hf_base.c
@ -197,8 +197,324 @@ hartree2eV(F64 energy_hartree) {
|
||||
return 27.2114*energy_hartree;
|
||||
}
|
||||
|
||||
|
||||
|
||||
function Problem problem_create(void) {
|
||||
Problem out = {0};
|
||||
out.arena = m_make_arena();
|
||||
out.atom.name = str8_lit("Hydrogen");
|
||||
out.atom.Z = 1;
|
||||
out.atom.occupancy[ECFG_1s] = 1;
|
||||
out.angular_momentum_l[ANGMOM_s] = 0.0;
|
||||
out.angular_momentum_l[ANGMOM_p] = 1.0;
|
||||
out.angular_momentum_l[ANGMOM_d] = 2.0;
|
||||
out.angular_momentum_l[ANGMOM_f] = 3.0;
|
||||
return out;
|
||||
}
|
||||
|
||||
/* Auxiliary routine: printing a matrix */
|
||||
function void print_eigenvalues(S32 l, S32 n, F64 *wr, F64 *wi) {
|
||||
ArenaTemp scratch = scratch_get(0, 0);
|
||||
S32 i, j;
|
||||
String8 newline = str8_lit("\n");
|
||||
String8 header = str8_pushf(scratch.arena, "\n ------- \n"
|
||||
"Sorted eigenvalues for l = %i\n", l);
|
||||
LOG(header.str);
|
||||
// printf("\n %s \n", desc);
|
||||
for (j = 0; j < n; j++) {
|
||||
String8 outstr =
|
||||
str8_pushf(scratch.arena, " (%4.5f, %4.5f) Hartree, %4.5f eV \n",
|
||||
wr[j], wi[j], hartree2eV(wr[j]));
|
||||
LOG(outstr.str);
|
||||
// printf(" (%6.2f,%6.2f)", a[i+j*lda].real, a[i+j*lda].imag );
|
||||
}
|
||||
LOG(newline.str);
|
||||
// printf("\n");
|
||||
|
||||
scratch_release(scratch);
|
||||
}
|
||||
|
||||
function void
|
||||
set_up_first_matrices(Problem *problem, 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
|
||||
{
|
||||
ArenaTemp scratch = scratch_get(0, 0);
|
||||
BSplineCtx *bspl_ctx = &problem->bspline_ctx;
|
||||
F64 *t = bspl_ctx->knotpoints;
|
||||
F64 Z = (F64)problem->atom.Z;
|
||||
U32 k = bspl_ctx->order;
|
||||
|
||||
// Skipping first bspline
|
||||
for (U32 i = 0; i < H->size1; i++) {
|
||||
for (U32 j = 0; j < H->size2; j++) {
|
||||
U32 bspl_index_i =
|
||||
i + 1; // The second Bspline has index 1 in our array etc.
|
||||
U32 bspl_index_j = j + 1;
|
||||
|
||||
// This logic assumes 1-indexed bsplines
|
||||
F64 abs_index_diff =
|
||||
fabs((F64)(bspl_index_i + 1) - (F64)(bspl_index_j + 1));
|
||||
if (!(abs_index_diff > ((F64)k - 1.0))) {
|
||||
// We do Gaussian quadrature between each knot point,
|
||||
// so we need to figure out where to start.
|
||||
// We start integration in the first shared knotpoint, which is the
|
||||
// one of the highest index.
|
||||
U32 start_knotpoint_index =
|
||||
bspl_index_i < bspl_index_j ? bspl_index_j : bspl_index_i;
|
||||
// And we integrate over the next k knotpoints.
|
||||
U32 end_knotpoint_index =
|
||||
bspl_index_i < bspl_index_j ? bspl_index_i + k : bspl_index_j + k;
|
||||
|
||||
F64 term1 = 0.0;
|
||||
F64 term2 = 0.0;
|
||||
F64 term3 = 0.0;
|
||||
F64 Bmat_term = 0.0;
|
||||
|
||||
for (U32 knotpoint_idx = start_knotpoint_index;
|
||||
knotpoint_idx < end_knotpoint_index; knotpoint_idx++) {
|
||||
F64 a = t[knotpoint_idx];
|
||||
F64 b = t[knotpoint_idx + 1];
|
||||
F64 prefac = 0.5 * (b - a);
|
||||
|
||||
// Only integrate non-zero intervals
|
||||
if (prefac > 1e-16) {
|
||||
for (U32 gq_i = 0; gq_i < g_gauss_legendre.order; gq_i++) {
|
||||
F64 w = g_gauss_legendre.weights[gq_i];
|
||||
F64 z = g_gauss_legendre.abscissae[gq_i];
|
||||
F64 r = (z * prefac) + ((a + b) * 0.5);
|
||||
F64 term_prefac = (prefac * w);
|
||||
F64 dB_i = compute_dBspline_F64(r, bspl_index_i);
|
||||
F64 dB_j = compute_dBspline_F64(r, bspl_index_j);
|
||||
F64 B_i = compute_bspline_F64(r, bspl_index_i);
|
||||
F64 B_j = compute_bspline_F64(r, bspl_index_j);
|
||||
term1 += term_prefac * dB_i * dB_j;
|
||||
term2 += term_prefac * B_i * B_j / (r * r);
|
||||
term3 += term_prefac * B_i * B_j / r;
|
||||
Bmat_term += term_prefac * B_i * B_j;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
F64 H_term_sum = 0.5 * term1 + (-Z) * term3;
|
||||
F64 H_l_term = 0.5 * term2;
|
||||
/* String8 debug = str8_pushf(scratch.arena,
|
||||
* "(i=%i,j=%i,t_i=%4.4f,t_i=%4.4f,term1=%.4e,term2=%.4e,term3=%.4e,term_sum=%.4e)
|
||||
* \n", */
|
||||
/* bspl_index_i, bspl_index_j,
|
||||
* t[bspl_index_i+k-1],t[bspl_index_j+k-1],term1,term2,term3,term_sum);
|
||||
*/
|
||||
/* LOG(debug.str); */
|
||||
mat_F64_set(H, i, j, H_term_sum);
|
||||
mat_F64_set(H_l, i, j, H_l_term);
|
||||
mat_F64_set(B_inv, i, j, Bmat_term);
|
||||
// mat_F64_set(&H, i, j, abs_index_diff);
|
||||
}
|
||||
// mat_F64_set(&H, i, j, abs_index_diff);
|
||||
}
|
||||
// LOG("\n");
|
||||
}
|
||||
|
||||
LOG(str8_pushf(scratch.arena, "H.size1=N-k-2=%i, last bspline index=%i \n",
|
||||
H->size1, bspl_ctx->num_bsplines - 1)
|
||||
.str);
|
||||
scratch_release(scratch);
|
||||
//print_mat_F64(H);
|
||||
LOG("\n");
|
||||
// print_mat_F64(B);
|
||||
}
|
||||
}
|
||||
|
||||
function void
|
||||
compute_wf_norm_F64(F64 *coeffs, U64 coeff_size, U64 n, U64 l) {
|
||||
ArenaTemp scratch = scratch_get(0, 0);
|
||||
|
||||
// Gauss legendre integration
|
||||
//
|
||||
F64 norm = 0.0;
|
||||
for (U64 i = 0; i < g_grid->num_steps - 1; i++) {
|
||||
F64 a = g_grid->points[i];
|
||||
F64 b = g_grid->points[i + 1];
|
||||
F64 prefac = 0.5 * (b - a);
|
||||
|
||||
// Only integrate non-zero intervals
|
||||
if (prefac > 1e-16) {
|
||||
for (U32 gq_i = 0; gq_i < g_gauss_legendre.order; gq_i++) {
|
||||
F64 w = g_gauss_legendre.weights[gq_i];
|
||||
F64 z = g_gauss_legendre.abscissae[gq_i];
|
||||
F64 r = (z * prefac) + ((a + b) * 0.5);
|
||||
F64 term_prefac = (prefac * w);
|
||||
F64 wf_at_r = 0.0;
|
||||
|
||||
for (U64 j = 0; j < coeff_size; j++) {
|
||||
wf_at_r += coeffs[j] * compute_bspline_F64(r, j + 1);
|
||||
}
|
||||
|
||||
norm += term_prefac * wf_at_r * wf_at_r;
|
||||
}
|
||||
}
|
||||
}
|
||||
String8 out =
|
||||
str8_pushf(scratch.arena, "n:%i, l:%i norm: %.2f \n", n, l, norm);
|
||||
LOG(out.str);
|
||||
scratch_release(scratch);
|
||||
}
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////
|
||||
//~
|
||||
// Main problem function, called from entry point.
|
||||
function void
|
||||
hf_main(void) {
|
||||
|
||||
Problem problem = problem_create();
|
||||
LOG(str8_pushf(problem.arena, "Created Problem-struct for %s \n", problem.atom.name.str).str);
|
||||
|
||||
set_up_gauss_legendre_points(problem.arena);
|
||||
|
||||
//- Set up grid and write to file.
|
||||
grid_assign(&problem.grid);
|
||||
set_up_grid(problem.arena);
|
||||
|
||||
write_array_binary_F64(str8_lit(grid_file_path_bin),
|
||||
problem.grid.points,
|
||||
problem.grid.num_steps);
|
||||
|
||||
write_array_F64(str8_lit(grid_file_path),
|
||||
problem.grid.points,
|
||||
problem.grid.num_steps,
|
||||
"%13.6e\n");
|
||||
|
||||
//- The BSpline context is the knotpoints and the BSpline order etc.
|
||||
bspline_ctx_assign(&problem.bspline_ctx);
|
||||
set_up_bspline_context(problem.arena);
|
||||
write_array_F64(str8_lit(knotpoints_file_path), problem.bspline_ctx.knotpoints,
|
||||
problem.bspline_ctx.num_knotpoints, "%13.6e\n");
|
||||
|
||||
//- Then we generate the BSplines and save them off for reference and
|
||||
// debugging.
|
||||
set_up_bsplines_at_points_and_write_matrix_F64(problem.arena);
|
||||
|
||||
U32 N = problem.bspline_ctx.num_knotpoints;
|
||||
U32 k = problem.bspline_ctx.order;
|
||||
U32 mat_size1 = N - k - 2;
|
||||
U32 mat_size2 = mat_size1;
|
||||
Mat_F64 H_base = mat_F64(problem.arena, mat_size1, mat_size2);
|
||||
Mat_F64 H_l_base = mat_F64(problem.arena, mat_size1, mat_size2);
|
||||
Mat_F64 H_l = mat_F64(problem.arena, mat_size1, mat_size2);
|
||||
Mat_F64 H = mat_F64(problem.arena, mat_size1, mat_size2);
|
||||
// This will be the inverse of B, but to start with we construct B.
|
||||
Mat_F64 B_inv = mat_F64(problem.arena, mat_size1, mat_size2);
|
||||
// A is the actual matrix for each eigenvalue problem.
|
||||
Mat_F64 A = mat_F64(problem.arena, H.size1, H.size2);
|
||||
set_up_first_matrices(&problem, &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
|
||||
// zgeev
|
||||
mat_invert_F64(&B_inv);
|
||||
|
||||
// This arena is used to push results from f. ex eigenvalue computations.
|
||||
// For each angular momentum
|
||||
for (U32 ang_mom_idx = 0; ang_mom_idx < MAX_NUM_ANGULAR_MOMENTA; ang_mom_idx++) {
|
||||
ArenaTemp scratch = scratch_get(0, 0);
|
||||
mat_F64_copy_to_dst(&H, &H_base);
|
||||
F64 l = problem.angular_momentum_l[ang_mom_idx];
|
||||
Eigensolution_F64 *eigsol = &problem.eigsols[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;
|
||||
mat_F64_copy_to_dst(&H_l, &H_l_base);
|
||||
// Multiply l(l+1)
|
||||
cblas_dscal(mat_size, l_factor, H_l.data, 1);
|
||||
// Add H = H_base + H_l
|
||||
cblas_daxpy(mat_size, 1.0, H_l.data, 1, H.data, 1);
|
||||
}
|
||||
|
||||
// Multiply to get A = B^-1 H
|
||||
{
|
||||
S32 n = A.size1;
|
||||
cblas_dgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, n, n, n, 1.0,
|
||||
B_inv.data, n, H.data, n, 0.0, A.data, n);
|
||||
//LOG("Matrix A: \n");
|
||||
//print_mat_F64(&A);
|
||||
}
|
||||
|
||||
// Solve generalised eigenvalue problem
|
||||
{
|
||||
S32 size1 = A.size1;
|
||||
S32 lda = size1;
|
||||
S32 ldvl = size1;
|
||||
S32 ldvr = size1;
|
||||
S32 info;
|
||||
S32 lwork;
|
||||
|
||||
F64 wkopt;
|
||||
F64 *work;
|
||||
|
||||
eigsol->eigvals_re= PushArray(problem.arena, F64, size1);
|
||||
F64 *wr = eigsol->eigvals_re;
|
||||
eigsol->eigvals_im = PushArray(problem.arena, F64, size1);
|
||||
F64 *wi = eigsol->eigvals_im;
|
||||
eigsol->left_eigvecs = mat_F64(problem.arena, ldvl, size1);
|
||||
F64 *vl = eigsol->left_eigvecs.data;
|
||||
eigsol->right_eigvecs = mat_F64(problem.arena, size1, ldvr);
|
||||
F64 *vr = eigsol->right_eigvecs.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 = PushArray(scratch.arena, F64, lwork);
|
||||
dgeev("Vectors", "Vectors", &size1, a, &lda, wr, wi, vl, &ldvl, vr, &ldvr,
|
||||
work, &lwork, &info);
|
||||
if (info > 0) {
|
||||
LOG("Failed to compute eigenvalues in dgeev\n");
|
||||
exit(1);
|
||||
}
|
||||
|
||||
// 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 );
|
||||
|
||||
U32 i = 0;
|
||||
F64 energy = -1000.0;
|
||||
U32 counter = 0;
|
||||
while (energy < 0.0) {
|
||||
energy = wr[i];
|
||||
U64 energy_index = sorted_indices[i];
|
||||
U64 n = 1 + i;
|
||||
if (ang_mom_idx > 0) {
|
||||
n = 2 + i;
|
||||
}
|
||||
|
||||
// compute_wf_norm_F64(eigensolution.right_eigenvectors.matrix[energy_index],
|
||||
// size1, n, ang_mom_idx);
|
||||
|
||||
U64 eigvec_idx = mat_get_col_major_idx(0, energy_index, size1);
|
||||
F64 *eigvecs = &eigsol->right_eigvecs.data[eigvec_idx];
|
||||
write_array_F64(
|
||||
get_eigenvector_filename(scratch.arena, n, ang_mom_idx),
|
||||
eigvecs, size1,
|
||||
"%13.6e\n");
|
||||
|
||||
i += 1;
|
||||
counter += 1;
|
||||
if (counter > 10) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
scratch_release(scratch);
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
|
||||
|
||||
|
||||
@ -26,6 +26,65 @@ struct GaussLegendre {
|
||||
F64 *abscissae;
|
||||
};
|
||||
|
||||
// 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 *eigvals_re;
|
||||
F64 *eigvals_im;
|
||||
Mat_F64 right_eigvecs;
|
||||
Mat_F64 left_eigvecs;
|
||||
};
|
||||
|
||||
typedef enum ElectronConfig ElectronConfig;
|
||||
enum ElectronConfig {
|
||||
ECFG_1s,
|
||||
ECFG_2s,
|
||||
ECFG_2p,
|
||||
ECFG_3s,
|
||||
ECFG_3p,
|
||||
ECFG_3d,
|
||||
ECFG_4s,
|
||||
ECFG_4p,
|
||||
ECFG_4d,
|
||||
ECFG_4f,
|
||||
ECFG_NUM_CONFIGS
|
||||
};
|
||||
|
||||
typedef enum AngularMomenta AngularMomenta;
|
||||
enum AngularMomenta {
|
||||
ANGMOM_s,
|
||||
ANGMOM_p,
|
||||
ANGMOM_d,
|
||||
ANGMOM_f,
|
||||
ANGMOM_NUM_MOMENTA
|
||||
};
|
||||
|
||||
typedef struct Atom Atom;
|
||||
struct Atom {
|
||||
String8 name;
|
||||
U32 Z;
|
||||
U32 occupancy[ECFG_NUM_CONFIGS];
|
||||
};
|
||||
|
||||
|
||||
// We use a "fat struct" approach where everything just exists here
|
||||
// in a single struct.
|
||||
#define MAX_NUM_ANGULAR_MOMENTA 3
|
||||
typedef struct Problem Problem;
|
||||
struct Problem {
|
||||
Arena *arena; // Just use a single arena to start with
|
||||
Grid grid;
|
||||
BSplineCtx bspline_ctx;
|
||||
Atom atom;
|
||||
F64 angular_momentum_l[ANGMOM_NUM_MOMENTA];
|
||||
Eigensolution_F64 eigsols[MAX_NUM_ANGULAR_MOMENTA];
|
||||
U32 num_eigsols;
|
||||
Mat_F64 H;
|
||||
Mat_F64 H_l;
|
||||
};
|
||||
|
||||
//~ Base math and utility functions
|
||||
|
||||
// Mat_F64 functions
|
||||
@ -35,7 +94,6 @@ function inline void mat_F64_set(Mat_F64 *mat, U32 i, U32 j, F64 val);
|
||||
function inline F64 mat_F64_get(Mat_F64 *mat, U32 i, U32 j);
|
||||
function Mat_F64 mat_F64_copy(Arena *arena, Mat_F64 *src);
|
||||
function void print_mat_F64(Mat_F64 *mat);
|
||||
|
||||
function void mat_invert_F64(Mat_F64 *mat);
|
||||
|
||||
// Gauss-Legendre
|
||||
@ -44,7 +102,17 @@ function void set_up_gauss_legendre_points(Arena *arena);
|
||||
// Random utility
|
||||
function void print_matrix_Z64(char *desc, int m, int n, Z64 *a, int lda);
|
||||
function void print_matrix_F64(char *desc, int m, int n, F64 *a, int lda);
|
||||
|
||||
function F64 hartree2eV(F64 energy_hartree);
|
||||
function void print_eigenvalues(S32 l, S32 n, F64 *wr, F64 *wi);
|
||||
function void compute_wf_norm_F64(F64 *coeffs, U64 coeff_size, U64 n, U64 l);
|
||||
|
||||
// Problem
|
||||
function Problem problem_create();
|
||||
function void set_up_first_matrices(Problem *problem, Mat_F64 *H,
|
||||
Mat_F64 *H_l, Mat_F64 *B_inv);
|
||||
function void hf_main();
|
||||
|
||||
|
||||
|
||||
|
||||
#endif /* HF_BASE_H */
|
||||
|
||||
@ -1,5 +1,5 @@
|
||||
#define DEBUG_LOG_WRITE_STRING_LIST_TO_FILE 0
|
||||
#define DEBUG_LOG_WRITE_ARRAY_BINARY 0
|
||||
#define DEBUG_LOG_WRITE_STRING_LIST_TO_FILE 1
|
||||
#define DEBUG_LOG_WRITE_ARRAY_BINARY 1
|
||||
|
||||
function void write_array_binary_F64(String8 path_to_file, F64 *values,
|
||||
U32 array_size) {
|
||||
373
src/main.c
373
src/main.c
@ -4,8 +4,8 @@
|
||||
#include "base/base_inc.h"
|
||||
#include "os/os_inc.h"
|
||||
|
||||
#include "hf/bsplines_and_grid.h"
|
||||
#include "hf/file_io.h"
|
||||
#include "hf/hf_bsplines_and_grid.h"
|
||||
#include "hf/hf_file_io.h"
|
||||
#include "hf/hf_base.h"
|
||||
|
||||
//----
|
||||
@ -14,8 +14,8 @@
|
||||
#include "os/os_entry_point.c"
|
||||
#include "os/os_inc.c"
|
||||
|
||||
#include "hf/bsplines_and_grid.c"
|
||||
#include "hf/file_io.c"
|
||||
#include "hf/hf_bsplines_and_grid.c"
|
||||
#include "hf/hf_file_io.c"
|
||||
#include "hf/hf_base.c"
|
||||
|
||||
// TODO make this a separate module that can be compiled instead
|
||||
@ -24,380 +24,21 @@
|
||||
//////
|
||||
//~
|
||||
|
||||
// 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 *eigvals_re;
|
||||
F64 *eigvals_im;
|
||||
Mat_F64 right_eigvecs;
|
||||
Mat_F64 left_eigvecs;
|
||||
};
|
||||
|
||||
typedef enum ElectronConfig ElectronConfig;
|
||||
enum ElectronConfig {
|
||||
ECFG_1s,
|
||||
ECFG_2s,
|
||||
ECFG_2p,
|
||||
ECFG_3s,
|
||||
ECFG_3p,
|
||||
ECFG_3d,
|
||||
ECFG_4s,
|
||||
ECFG_4p,
|
||||
ECFG_4d,
|
||||
ECFG_4f,
|
||||
ECFG_NUM_CONFIGS
|
||||
};
|
||||
|
||||
typedef enum AngularMomenta AngularMomenta;
|
||||
enum AngularMomenta {
|
||||
ANGMOM_s,
|
||||
ANGMOM_p,
|
||||
ANGMOM_d,
|
||||
ANGMOM_f,
|
||||
ANGMOM_NUM_MOMENTA
|
||||
};
|
||||
|
||||
typedef struct Atom Atom;
|
||||
struct Atom {
|
||||
String8 name;
|
||||
U32 Z;
|
||||
U32 occupancy[ECFG_NUM_CONFIGS];
|
||||
};
|
||||
|
||||
|
||||
// We use a "fat struct" approach where everything just exists here
|
||||
// in a single struct.
|
||||
#define MAX_NUM_ANGULAR_MOMENTA 3
|
||||
typedef struct Problem Problem;
|
||||
struct Problem {
|
||||
Arena *arena; // Just use a single arena to start with
|
||||
Grid grid;
|
||||
BSplineCtx bspline_ctx;
|
||||
Atom atom;
|
||||
F64 angular_momentum_l[ANGMOM_NUM_MOMENTA];
|
||||
Eigensolution_F64 eigsols[MAX_NUM_ANGULAR_MOMENTA];
|
||||
U32 num_eigsols;
|
||||
Mat_F64 H;
|
||||
Mat_F64 H_l;
|
||||
};
|
||||
|
||||
//////
|
||||
//~
|
||||
|
||||
function Problem problem_create() {
|
||||
Problem out = {0};
|
||||
out.arena = m_make_arena();
|
||||
out.atom.name = str8_lit("Hydrogen");
|
||||
out.atom.Z = 1;
|
||||
out.atom.occupancy[ECFG_1s] = 1;
|
||||
out.angular_momentum_l[ANGMOM_s] = 0.0;
|
||||
out.angular_momentum_l[ANGMOM_p] = 1.0;
|
||||
out.angular_momentum_l[ANGMOM_d] = 2.0;
|
||||
out.angular_momentum_l[ANGMOM_f] = 3.0;
|
||||
return out;
|
||||
}
|
||||
|
||||
/* Auxiliary routine: printing a matrix */
|
||||
function void print_eigenvalues(S32 l, S32 n, F64 *wr, F64 *wi) {
|
||||
ArenaTemp scratch = scratch_get(0, 0);
|
||||
S32 i, j;
|
||||
String8 newline = str8_lit("\n");
|
||||
String8 header = str8_pushf(scratch.arena, "\n ------- \n"
|
||||
"Sorted eigenvalues for l = %i\n", l);
|
||||
LOG(header.str);
|
||||
// printf("\n %s \n", desc);
|
||||
for (j = 0; j < n; j++) {
|
||||
String8 outstr =
|
||||
str8_pushf(scratch.arena, " (%4.5f, %4.5f) Hartree, %4.5f eV \n",
|
||||
wr[j], wi[j], hartree2eV(wr[j]));
|
||||
LOG(outstr.str);
|
||||
// printf(" (%6.2f,%6.2f)", a[i+j*lda].real, a[i+j*lda].imag );
|
||||
}
|
||||
LOG(newline.str);
|
||||
// printf("\n");
|
||||
|
||||
scratch_release(scratch);
|
||||
}
|
||||
|
||||
function void
|
||||
set_up_first_matrices(Problem *problem, 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
|
||||
{
|
||||
ArenaTemp scratch = scratch_get(0, 0);
|
||||
BSplineCtx *bspl_ctx = &problem->bspline_ctx;
|
||||
F64 *t = bspl_ctx->knotpoints;
|
||||
F64 Z = (F64)problem->atom.Z;
|
||||
U32 k = bspl_ctx->order;
|
||||
|
||||
// Skipping first bspline
|
||||
for (U32 i = 0; i < H->size1; i++) {
|
||||
for (U32 j = 0; j < H->size2; j++) {
|
||||
U32 bspl_index_i =
|
||||
i + 1; // The second Bspline has index 1 in our array etc.
|
||||
U32 bspl_index_j = j + 1;
|
||||
|
||||
// This logic assumes 1-indexed bsplines
|
||||
F64 abs_index_diff =
|
||||
fabs((F64)(bspl_index_i + 1) - (F64)(bspl_index_j + 1));
|
||||
if (!(abs_index_diff > ((F64)k - 1.0))) {
|
||||
// We do Gaussian quadrature between each knot point,
|
||||
// so we need to figure out where to start.
|
||||
// We start integration in the first shared knotpoint, which is the
|
||||
// one of the highest index.
|
||||
U32 start_knotpoint_index =
|
||||
bspl_index_i < bspl_index_j ? bspl_index_j : bspl_index_i;
|
||||
// And we integrate over the next k knotpoints.
|
||||
U32 end_knotpoint_index =
|
||||
bspl_index_i < bspl_index_j ? bspl_index_i + k : bspl_index_j + k;
|
||||
|
||||
F64 term1 = 0.0;
|
||||
F64 term2 = 0.0;
|
||||
F64 term3 = 0.0;
|
||||
F64 Bmat_term = 0.0;
|
||||
|
||||
for (U32 knotpoint_idx = start_knotpoint_index;
|
||||
knotpoint_idx < end_knotpoint_index; knotpoint_idx++) {
|
||||
F64 a = t[knotpoint_idx];
|
||||
F64 b = t[knotpoint_idx + 1];
|
||||
F64 prefac = 0.5 * (b - a);
|
||||
|
||||
// Only integrate non-zero intervals
|
||||
if (prefac > 1e-16) {
|
||||
for (U32 gq_i = 0; gq_i < g_gauss_legendre.order; gq_i++) {
|
||||
F64 w = g_gauss_legendre.weights[gq_i];
|
||||
F64 z = g_gauss_legendre.abscissae[gq_i];
|
||||
F64 r = (z * prefac) + ((a + b) * 0.5);
|
||||
F64 term_prefac = (prefac * w);
|
||||
F64 dB_i = compute_dBspline_F64(r, bspl_index_i);
|
||||
F64 dB_j = compute_dBspline_F64(r, bspl_index_j);
|
||||
F64 B_i = compute_bspline_F64(r, bspl_index_i);
|
||||
F64 B_j = compute_bspline_F64(r, bspl_index_j);
|
||||
term1 += term_prefac * dB_i * dB_j;
|
||||
term2 += term_prefac * B_i * B_j / (r * r);
|
||||
term3 += term_prefac * B_i * B_j / r;
|
||||
Bmat_term += term_prefac * B_i * B_j;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
F64 H_term_sum = 0.5 * term1 + (-Z) * term3;
|
||||
F64 H_l_term = 0.5 * term2;
|
||||
/* String8 debug = str8_pushf(scratch.arena,
|
||||
* "(i=%i,j=%i,t_i=%4.4f,t_i=%4.4f,term1=%.4e,term2=%.4e,term3=%.4e,term_sum=%.4e)
|
||||
* \n", */
|
||||
/* bspl_index_i, bspl_index_j,
|
||||
* t[bspl_index_i+k-1],t[bspl_index_j+k-1],term1,term2,term3,term_sum);
|
||||
*/
|
||||
/* LOG(debug.str); */
|
||||
mat_F64_set(H, i, j, H_term_sum);
|
||||
mat_F64_set(H_l, i, j, H_l_term);
|
||||
mat_F64_set(B_inv, i, j, Bmat_term);
|
||||
// mat_F64_set(&H, i, j, abs_index_diff);
|
||||
}
|
||||
// mat_F64_set(&H, i, j, abs_index_diff);
|
||||
}
|
||||
// LOG("\n");
|
||||
}
|
||||
|
||||
LOG(str8_pushf(scratch.arena, "H.size1=N-k-2=%i, last bspline index=%i \n",
|
||||
H->size1, bspl_ctx->num_bsplines - 1)
|
||||
.str);
|
||||
scratch_release(scratch);
|
||||
//print_mat_F64(H);
|
||||
LOG("\n");
|
||||
// print_mat_F64(B);
|
||||
}
|
||||
}
|
||||
|
||||
function void compute_wf_norm_F64(F64 *coeffs, U64 coeff_size, U64 n, U64 l) {
|
||||
ArenaTemp scratch = scratch_get(0, 0);
|
||||
|
||||
// Gauss legendre integration
|
||||
//
|
||||
F64 norm = 0.0;
|
||||
for (U64 i = 0; i < g_grid->num_steps - 1; i++) {
|
||||
F64 a = g_grid->points[i];
|
||||
F64 b = g_grid->points[i + 1];
|
||||
F64 prefac = 0.5 * (b - a);
|
||||
|
||||
// Only integrate non-zero intervals
|
||||
if (prefac > 1e-16) {
|
||||
for (U32 gq_i = 0; gq_i < g_gauss_legendre.order; gq_i++) {
|
||||
F64 w = g_gauss_legendre.weights[gq_i];
|
||||
F64 z = g_gauss_legendre.abscissae[gq_i];
|
||||
F64 r = (z * prefac) + ((a + b) * 0.5);
|
||||
F64 term_prefac = (prefac * w);
|
||||
F64 wf_at_r = 0.0;
|
||||
|
||||
for (U64 j = 0; j < coeff_size; j++) {
|
||||
wf_at_r += coeffs[j] * compute_bspline_F64(r, j + 1);
|
||||
}
|
||||
|
||||
norm += term_prefac * wf_at_r * wf_at_r;
|
||||
}
|
||||
}
|
||||
}
|
||||
String8 out =
|
||||
str8_pushf(scratch.arena, "n:%i, l:%i norm: %.2f \n", n, l, norm);
|
||||
LOG(out.str);
|
||||
scratch_release(scratch);
|
||||
}
|
||||
|
||||
/////////////////
|
||||
//~
|
||||
// Main entry point
|
||||
function void EntryPoint(void) {
|
||||
|
||||
// Init subsystems
|
||||
OS_InitReceipt os_receipt = OS_init();
|
||||
OS_InitGfxReceipt os_gfx_receipt = OS_gfx_init(os_receipt);
|
||||
|
||||
Problem problem = problem_create();
|
||||
LOG(str8_pushf(problem.arena, "Created Problem-struct for %s \n", problem.atom.name).str);
|
||||
|
||||
set_up_gauss_legendre_points(problem.arena);
|
||||
|
||||
//- Set up grid and write to file.
|
||||
grid_assign(&problem.grid);
|
||||
set_up_grid(problem.arena);
|
||||
|
||||
write_array_binary_F64(str8_lit(grid_file_path_bin),
|
||||
problem.grid.points,
|
||||
problem.grid.num_steps);
|
||||
write_array_F64(str8_lit(grid_file_path), problem.grid.points,problem.grid.num_steps,
|
||||
"%13.6e\n");
|
||||
|
||||
//- The BSpline context is the knotpoints and the BSpline order etc.
|
||||
bspline_ctx_assign(&problem.bspline_ctx);
|
||||
set_up_bspline_context(problem.arena);
|
||||
write_array_F64(str8_lit(knotpoints_file_path), problem.bspline_ctx.knotpoints,
|
||||
problem.bspline_ctx.num_knotpoints, "%13.6e\n");
|
||||
|
||||
//- Then we generate the BSplines and save them off for reference and
|
||||
// debugging.
|
||||
set_up_bsplines_at_points_and_write_matrix_F64(problem.arena);
|
||||
|
||||
U32 N = problem.bspline_ctx.num_knotpoints;
|
||||
U32 k = problem.bspline_ctx.order;
|
||||
U32 mat_size1 = N - k - 2;
|
||||
U32 mat_size2 = mat_size1;
|
||||
Mat_F64 H_base = mat_F64(problem.arena, mat_size1, mat_size2);
|
||||
Mat_F64 H_l_base = mat_F64(problem.arena, mat_size1, mat_size2);
|
||||
Mat_F64 H_l = mat_F64(problem.arena, mat_size1, mat_size2);
|
||||
Mat_F64 H = mat_F64(problem.arena, mat_size1, mat_size2);
|
||||
// This will be the inverse of B, but to start with we construct B.
|
||||
Mat_F64 B_inv = mat_F64(problem.arena, mat_size1, mat_size2);
|
||||
// A is the actual matrix for each eigenvalue problem.
|
||||
Mat_F64 A = mat_F64(problem.arena, H.size1, H.size2);
|
||||
set_up_first_matrices(&problem, &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
|
||||
// zgeev
|
||||
mat_invert_F64(&B_inv);
|
||||
|
||||
// This arena is used to push results from f. ex eigenvalue computations.
|
||||
// For each angular momentum
|
||||
for (U32 ang_mom_idx = 0; ang_mom_idx < MAX_NUM_ANGULAR_MOMENTA; ang_mom_idx++) {
|
||||
ArenaTemp scratch = scratch_get(0, 0);
|
||||
mat_F64_copy_to_dst(&H, &H_base);
|
||||
F64 l = problem.angular_momentum_l[ang_mom_idx];
|
||||
Eigensolution_F64 *eigsol = &problem.eigsols[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;
|
||||
mat_F64_copy_to_dst(&H_l, &H_l_base);
|
||||
// Multiply l(l+1)
|
||||
cblas_dscal(mat_size, l_factor, H_l.data, 1);
|
||||
// Add H = H_base + H_l
|
||||
cblas_daxpy(mat_size, 1.0, H_l.data, 1, H.data, 1);
|
||||
}
|
||||
|
||||
// Multiply to get A = B^-1 H
|
||||
{
|
||||
S32 n = A.size1;
|
||||
cblas_dgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, n, n, n, 1.0,
|
||||
B_inv.data, n, H.data, n, 0.0, A.data, n);
|
||||
//LOG("Matrix A: \n");
|
||||
//print_mat_F64(&A);
|
||||
}
|
||||
|
||||
// Solve generalised eigenvalue problem
|
||||
{
|
||||
S32 size1 = A.size1;
|
||||
S32 lda = size1;
|
||||
S32 ldvl = size1;
|
||||
S32 ldvr = size1;
|
||||
S32 info;
|
||||
S32 lwork;
|
||||
|
||||
F64 wkopt;
|
||||
F64 *work;
|
||||
|
||||
eigsol->eigvals_re= PushArray(problem.arena, F64, size1);
|
||||
F64 *wr = eigsol->eigvals_re;
|
||||
eigsol->eigvals_im = PushArray(problem.arena, F64, size1);
|
||||
F64 *wi = eigsol->eigvals_im;
|
||||
eigsol->left_eigvecs = mat_F64(problem.arena, ldvl, size1);
|
||||
F64 *vl = eigsol->left_eigvecs.data;
|
||||
eigsol->right_eigvecs = mat_F64(problem.arena, size1, ldvr);
|
||||
F64 *vr = eigsol->right_eigvecs.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 = PushArray(scratch.arena, F64, lwork);
|
||||
dgeev("Vectors", "Vectors", &size1, a, &lda, wr, wi, vl, &ldvl, vr, &ldvr,
|
||||
work, &lwork, &info);
|
||||
if (info > 0) {
|
||||
LOG("Failed to compute eigenvalues in dgeev\n");
|
||||
exit(1);
|
||||
}
|
||||
|
||||
// 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 );
|
||||
|
||||
U32 i = 0;
|
||||
F64 energy = -1000.0;
|
||||
U32 counter = 0;
|
||||
while (energy < 0.0) {
|
||||
energy = wr[i];
|
||||
U64 energy_index = sorted_indices[i];
|
||||
U64 n = 1 + i;
|
||||
if (ang_mom_idx > 0) {
|
||||
n = 2 + i;
|
||||
}
|
||||
|
||||
// compute_wf_norm_F64(eigensolution.right_eigenvectors.matrix[energy_index],
|
||||
// size1, n, ang_mom_idx);
|
||||
|
||||
U64 eigvec_idx = mat_get_col_major_idx(0, energy_index, size1);
|
||||
F64 *eigvecs = &eigsol->right_eigvecs.data[eigvec_idx];
|
||||
write_array_F64(
|
||||
get_eigenvector_filename(scratch.arena, n, ang_mom_idx),
|
||||
eigvecs, size1,
|
||||
"%13.6e\n");
|
||||
|
||||
i += 1;
|
||||
counter += 1;
|
||||
if (counter > 10) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
scratch_release(scratch);
|
||||
}
|
||||
// Main program
|
||||
hf_main();
|
||||
|
||||
}
|
||||
}
|
||||
|
||||
Loading…
Reference in New Issue
Block a user