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;
|
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;
|
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
|
//~ Base math and utility functions
|
||||||
|
|
||||||
// Mat_F64 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 inline F64 mat_F64_get(Mat_F64 *mat, U32 i, U32 j);
|
||||||
function Mat_F64 mat_F64_copy(Arena *arena, Mat_F64 *src);
|
function Mat_F64 mat_F64_copy(Arena *arena, Mat_F64 *src);
|
||||||
function void print_mat_F64(Mat_F64 *mat);
|
function void print_mat_F64(Mat_F64 *mat);
|
||||||
|
|
||||||
function void mat_invert_F64(Mat_F64 *mat);
|
function void mat_invert_F64(Mat_F64 *mat);
|
||||||
|
|
||||||
// Gauss-Legendre
|
// Gauss-Legendre
|
||||||
@ -44,7 +102,17 @@ function void set_up_gauss_legendre_points(Arena *arena);
|
|||||||
// Random utility
|
// Random utility
|
||||||
function void print_matrix_Z64(char *desc, int m, int n, Z64 *a, int lda);
|
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 void print_matrix_F64(char *desc, int m, int n, F64 *a, int lda);
|
||||||
|
|
||||||
function F64 hartree2eV(F64 energy_hartree);
|
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 */
|
#endif /* HF_BASE_H */
|
||||||
|
|||||||
@ -1,5 +1,5 @@
|
|||||||
#define DEBUG_LOG_WRITE_STRING_LIST_TO_FILE 0
|
#define DEBUG_LOG_WRITE_STRING_LIST_TO_FILE 1
|
||||||
#define DEBUG_LOG_WRITE_ARRAY_BINARY 0
|
#define DEBUG_LOG_WRITE_ARRAY_BINARY 1
|
||||||
|
|
||||||
function void write_array_binary_F64(String8 path_to_file, F64 *values,
|
function void write_array_binary_F64(String8 path_to_file, F64 *values,
|
||||||
U32 array_size) {
|
U32 array_size) {
|
||||||
373
src/main.c
373
src/main.c
@ -4,8 +4,8 @@
|
|||||||
#include "base/base_inc.h"
|
#include "base/base_inc.h"
|
||||||
#include "os/os_inc.h"
|
#include "os/os_inc.h"
|
||||||
|
|
||||||
#include "hf/bsplines_and_grid.h"
|
#include "hf/hf_bsplines_and_grid.h"
|
||||||
#include "hf/file_io.h"
|
#include "hf/hf_file_io.h"
|
||||||
#include "hf/hf_base.h"
|
#include "hf/hf_base.h"
|
||||||
|
|
||||||
//----
|
//----
|
||||||
@ -14,8 +14,8 @@
|
|||||||
#include "os/os_entry_point.c"
|
#include "os/os_entry_point.c"
|
||||||
#include "os/os_inc.c"
|
#include "os/os_inc.c"
|
||||||
|
|
||||||
#include "hf/bsplines_and_grid.c"
|
#include "hf/hf_bsplines_and_grid.c"
|
||||||
#include "hf/file_io.c"
|
#include "hf/hf_file_io.c"
|
||||||
#include "hf/hf_base.c"
|
#include "hf/hf_base.c"
|
||||||
|
|
||||||
// TODO make this a separate module that can be compiled instead
|
// 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
|
// Main entry point
|
||||||
function void EntryPoint(void) {
|
function void EntryPoint(void) {
|
||||||
|
|
||||||
|
// Init subsystems
|
||||||
OS_InitReceipt os_receipt = OS_init();
|
OS_InitReceipt os_receipt = OS_init();
|
||||||
OS_InitGfxReceipt os_gfx_receipt = OS_gfx_init(os_receipt);
|
OS_InitGfxReceipt os_gfx_receipt = OS_gfx_init(os_receipt);
|
||||||
|
|
||||||
Problem problem = problem_create();
|
// Main program
|
||||||
LOG(str8_pushf(problem.arena, "Created Problem-struct for %s \n", problem.atom.name).str);
|
hf_main();
|
||||||
|
|
||||||
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);
|
|
||||||
}
|
|
||||||
|
|
||||||
}
|
}
|
||||||
}
|
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user