Skip to content

Commit

Permalink
add options to specify excited states in QD
Browse files Browse the repository at this point in the history
  • Loading branch information
tjira committed Dec 26, 2024
1 parent bc7412d commit 337ee88
Show file tree
Hide file tree
Showing 11 changed files with 93 additions and 55 deletions.
2 changes: 1 addition & 1 deletion example/input/quantum_dynamics_1d-1s_imag.json
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
{
"quantum_dynamics" : {
"adiabatic" : false,
"imaginary" : true,
"iterations" : 1000,
"mode" : [1, 0],
"time_step" : 0.1,
"grid" : {
"limits" : [-8, 8],
Expand Down
2 changes: 1 addition & 1 deletion example/input/quantum_dynamics_1d-1s_real.json
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
{
"quantum_dynamics" : {
"adiabatic" : false,
"imaginary" : false,
"iterations" : 350,
"mode" : [0, 1],
"time_step" : 0.1,
"grid" : {
"limits" : [-8, 8],
Expand Down
2 changes: 1 addition & 1 deletion example/input/quantum_dynamics_1d-2s_real.json
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
{
"quantum_dynamics" : {
"adiabatic" : true,
"imaginary" : false,
"iterations" : 350,
"mode" : [0, 1],
"time_step" : 10,
"grid" : {
"limits" : [-16, 32],
Expand Down
2 changes: 1 addition & 1 deletion example/input/quantum_dynamics_2d-1s_imag.json
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
{
"quantum_dynamics" : {
"adiabatic" : false,
"imaginary" : true,
"iterations" : 1000,
"mode" : [1, 0],
"time_step" : 0.1,
"grid" : {
"limits" : [-8, 8],
Expand Down
2 changes: 1 addition & 1 deletion example/input/quantum_dynamics_2d-1s_real.json
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
{
"quantum_dynamics" : {
"adiabatic" : false,
"imaginary" : false,
"iterations" : 350,
"mode" : [0, 1],
"time_step" : 0.1,
"grid" : {
"limits" : [-8, 8],
Expand Down
2 changes: 1 addition & 1 deletion example/input/quantum_dynamics_3d-1s_imag.json
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
{
"quantum_dynamics" : {
"adiabatic" : false,
"imaginary" : true,
"iterations" : 1000,
"mode" : [1, 0],
"time_step" : 0.1,
"grid" : {
"limits" : [-8, 8],
Expand Down
2 changes: 1 addition & 1 deletion example/input/quantum_dynamics_3d-1s_real.json
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
{
"quantum_dynamics" : {
"adiabatic" : false,
"imaginary" : false,
"iterations" : 350,
"mode" : [0, 1],
"time_step" : 0.1,
"grid" : {
"limits" : [-8, 8],
Expand Down
2 changes: 1 addition & 1 deletion research/fssh_vs_lzsh/input.sh
Original file line number Diff line number Diff line change
Expand Up @@ -64,8 +64,8 @@ EOM
read -r -d '' TEMPLATE_EXACT_DYN <<- EOM
{
"adiabatic" : true,
"imaginary" : false,
"iterations" : 500,
"mode" : [0, 1],
"time_step" : 10,
"grid" : {
"limits" : [-48, 144],
Expand Down
8 changes: 4 additions & 4 deletions src/classicaldynamics.zig
Original file line number Diff line number Diff line change
Expand Up @@ -11,10 +11,10 @@ const asfloat = @import("helper.zig").asfloat;
pub fn ClassicalDynamicsOptions(comptime T: type) type {
return struct {
const InitialConditions = struct {
position_mean: []const T = &[_]f64{-10.0},
position_std: []const T = &[_]f64{ 0.5},
momentum_mean: []const T = &[_]f64{ 15.0},
momentum_std: []const T = &[_]f64{ 1.0},
position_mean: []const T = &[_]T{-10.0},
position_std: []const T = &[_]T{ 0.5},
momentum_mean: []const T = &[_]T{ 15.0},
momentum_std: []const T = &[_]T{ 1.0},
state: u32 = 1, mass: T = 2000
};
const FewestSwitches = struct {
Expand Down
102 changes: 70 additions & 32 deletions src/quantumdynamics.zig
Original file line number Diff line number Diff line change
Expand Up @@ -33,8 +33,8 @@ pub fn QuantumDynamicsOptions(comptime T: type) type {
};

adiabatic: bool = true,
imaginary: bool = false,
iterations: u32 = 300,
mode: []const u32 = &[_]u32{0, 1},
potential: []const u8 = "tully1D_1",
time_step: T = 10,

Expand Down Expand Up @@ -108,62 +108,100 @@ pub fn run(comptime T: type, opt: QuantumDynamicsOptions(T), print: bool, alloca
var W = try Wavefunction(T).init(ndim, nstate, opt.grid.points, allocator); defer W.deinit();
var WA = try Wavefunction(T).init(ndim, nstate, opt.grid.points, allocator); defer WA.deinit();

var WOPT = try std.ArrayList(Wavefunction(T)).initCapacity(allocator, if (opt.mode[0] > 0) opt.mode[0] - 1 else 0); defer WOPT.deinit();

for (0..WOPT.capacity) |_| {
try WOPT.append(try Wavefunction(T).init(ndim, nstate, opt.grid.points, allocator));
}

var P = try Matrix(T).init(nstate, nstate, allocator); defer P.deinit();

mpt.rgrid(T, &rvec, opt.grid.limits[0], opt.grid.limits[1], opt.grid.points);
mpt.kgrid(T, &kvec, opt.grid.limits[0], opt.grid.limits[1], opt.grid.points);

const VS = try rgridPotentials(T, opt.potential, rvec, allocator); var V = VS[0]; var VA = VS[1]; var VC = VS[2]; defer V.deinit(); defer VA.deinit(); defer VC.deinit();

const R = try rgridPropagators(T, VA, VC, rvec, opt.time_step, opt.imaginary, allocator); defer R.deinit();
const K = try kgridPropagators(T, W.nstate, kvec, opt.time_step, opt.initial_conditions.mass, opt.imaginary, allocator); defer K.deinit();
var R = try rgridPropagators(T, VA, VC, rvec, opt.time_step, opt.mode[0] > 0, allocator); defer R.deinit();
var K = try kgridPropagators(T, W.nstate, kvec, opt.time_step, opt.initial_conditions.mass, opt.mode[0] > 0, allocator); defer K.deinit();

if (opt.write.wavefunction != null) for (0..rdim) |i| for (0..ndim) |j| {
wavefunction.ptr(i, j).* = rvec.at(i, j);
};

wfn.guess(T, &W, rvec, opt.initial_conditions.position, opt.initial_conditions.momentum, opt.initial_conditions.gamma, opt.initial_conditions.state); wfn.normalize(T, &W, dr);
for (0..opt.mode[0] + opt.mode[1]) |i| {

if (print) try std.io.getStdOut().writer().print("\n{s:6} {s:12} {s:12} {s:12}", .{"ITER", "EKIN", "EPOT", "ETOT"});
if (print) try std.io.getStdOut().writer().print("\n{s} TIME DYNAMICS #{d}", .{if (i < opt.mode[0]) "IMAGINARY" else "REAL", if (i < opt.mode[0]) i + 1 else i - opt.mode[0] + 1});

if (print) {if (W.ndim > 1) for (0..W.ndim - 1) |_| {try std.io.getStdOut().writer().print(" " ** 11, .{});}; try std.io.getStdOut().writer().print(" {s:11}", .{"POSITION" });}
if (print) {if (W.ndim > 1) for (0..W.ndim - 1) |_| {try std.io.getStdOut().writer().print(" " ** 11, .{});}; try std.io.getStdOut().writer().print(" {s:11}", .{"MOMENTUM" });}
if (print) {if (W.nstate > 1) for (0..W.nstate - 1) |_| {try std.io.getStdOut().writer().print(" " ** 10, .{});}; try std.io.getStdOut().writer().print(" {s:10}\n", .{"POPULATION"});}
if (i == opt.mode[0]) {

for (0..opt.iterations) |i| {
for (R.items) |*e| {e.deinit();} for (K.items) |*e| {e.deinit();} R.deinit(); K.deinit();

R = try rgridPropagators(T, VA, VC, rvec, opt.time_step, false, allocator);
K = try kgridPropagators(T, W.nstate, kvec, opt.time_step, opt.initial_conditions.mass, false, allocator);
}

try wfn.propagate(T, &W, R, K, &T1); if (opt.imaginary) wfn.normalize(T, &W, dr);
if (i < opt.mode[0] or (opt.mode[0] == 0)) {
wfn.guess(T, &W, rvec, opt.initial_conditions.position, opt.initial_conditions.momentum, opt.initial_conditions.gamma, opt.initial_conditions.state); wfn.normalize(T, &W, dr);
}

if (opt.adiabatic) wfn.adiabatize(T, &WA, W, VC);
if (print) try std.io.getStdOut().writer().print("\n{s:6} {s:12} {s:12} {s:12}", .{"ITER", "EKIN", "EPOT", "ETOT"});

const Ekin = try wfn.ekin(T, W, kvec, opt.initial_conditions.mass, dr, &T1); const Epot: T = wfn.epot(T, W, V, dr);
if (print) {if (W.ndim > 1) for (0..W.ndim - 1) |_| {try std.io.getStdOut().writer().print(" " ** 11, .{});}; try std.io.getStdOut().writer().print(" {s:11}", .{"POSITION" });}
if (print) {if (W.ndim > 1) for (0..W.ndim - 1) |_| {try std.io.getStdOut().writer().print(" " ** 11, .{});}; try std.io.getStdOut().writer().print(" {s:11}", .{"MOMENTUM" });}
if (print) {if (W.nstate > 1) for (0..W.nstate - 1) |_| {try std.io.getStdOut().writer().print(" " ** 10, .{});}; try std.io.getStdOut().writer().print(" {s:10}\n", .{"POPULATION"});}

wfn.density(T, &P, if (opt.adiabatic) WA else W, dr); wfn.position(T, &r, W, rvec, dr); try wfn.momentum(T, &p, W, kvec, dr, &T1);
for (0..opt.iterations) |j| {

if (opt.write.population != null) for (0..nstate) |j| {pop.ptr(i, 1 + j).* = P.at(j, j);};
if (opt.write.position != null) for (0..ndim) |j| {position.ptr(i, 1 + j).* = r.at(j);};
if (opt.write.momentum != null) for (0..ndim) |j| {momentum.ptr(i, 1 + j).* = p.at(j);};
if (opt.write.kinetic_energy != null) ekin.ptr(i, 1 + 0).* = Ekin ;
if (opt.write.potential_energy != null) epot.ptr(i, 1 + 0).* = Epot ;
if (opt.write.total_energy != null) etot.ptr(i, 1 + 0).* = Ekin + Epot ;
try wfn.propagate(T, &W, R, K, &T1);

if (opt.write.wavefunction != null) for (0..rdim) |j| for (0..nstate) |k| {
wavefunction.ptr(j, ndim + 2 * i * nstate + 2 * k + 0).* = (if (opt.adiabatic) WA else W).data.at(j, k).re;
wavefunction.ptr(j, ndim + 2 * i * nstate + 2 * k + 1).* = (if (opt.adiabatic) WA else W).data.at(j, k).im;
};
if (i < opt.mode[0]) {

for (0..i) |k| {

if (i == opt.iterations - 1) {
@memcpy(output.P.data, P.data); @memcpy(output.r.data, r.data); @memcpy(output.p.data, p.data); output.Ekin = Ekin; output.Epot = Epot;
}
const overlap = wfn.overlap(T, WOPT.items[k], W, dr);

if (print and (i == 0 or (i + 1) % opt.log_intervals.iteration == 0)) try printIteration(T, @intCast(i), Ekin, Epot, r, p, P);
}
for (0..rdim) |l| for (0..nstate) |m| {
W.data.ptr(l, m).* = W.data.at(l, m).sub(WOPT.items[k].data.at(l, m).conjugate().mul(overlap));
};
}

for (R.items) |*e| {e.deinit();} for (K.items) |*e| {e.deinit();} for (V.items) |*e| {e.deinit();} for (VA.items) |*e| {e.deinit();} for (VC.items) |*e| {e.deinit();}
}
wfn.normalize(T, &W, dr);
}

if (opt.adiabatic) wfn.adiabatize(T, &WA, W, VC);

const Ekin = try wfn.ekin(T, W, kvec, opt.initial_conditions.mass, dr, &T1); const Epot: T = wfn.epot(T, W, V, dr);

wfn.density(T, &P, if (opt.adiabatic) WA else W, dr); wfn.position(T, &r, W, rvec, dr); try wfn.momentum(T, &p, W, kvec, dr, &T1);

if (i == opt.mode[0] + opt.mode[1] - 1 and opt.write.population != null) for (0..nstate) |k| {pop.ptr(j, 1 + k).* = P.at(k, k);};
if (i == opt.mode[0] + opt.mode[1] - 1 and opt.write.position != null) for (0..ndim) |k| {position.ptr(j, 1 + k).* = r.at(k);};
if (i == opt.mode[0] + opt.mode[1] - 1 and opt.write.momentum != null) for (0..ndim) |k| {momentum.ptr(j, 1 + k).* = p.at(k);};
if (i == opt.mode[0] + opt.mode[1] - 1 and opt.write.kinetic_energy != null) ekin.ptr(j, 1 + 0).* = Ekin ;
if (i == opt.mode[0] + opt.mode[1] - 1 and opt.write.potential_energy != null) epot.ptr(j, 1 + 0).* = Epot ;
if (i == opt.mode[0] + opt.mode[1] - 1 and opt.write.total_energy != null) etot.ptr(j, 1 + 0).* = Ekin + Epot ;

if (i == opt.mode[0] + opt.mode[1] - 1 and opt.write.wavefunction != null) for (0..rdim) |k| for (0..nstate) |l| {
wavefunction.ptr(k, ndim + 2 * j * nstate + 2 * l + 0).* = (if (opt.adiabatic) WA else W).data.at(k, l).re;
wavefunction.ptr(k, ndim + 2 * j * nstate + 2 * l + 1).* = (if (opt.adiabatic) WA else W).data.at(k, l).im;
};

if (print and (j == 0 or (j + 1) % opt.log_intervals.iteration == 0)) try printIteration(T, @intCast(j), Ekin, Epot, r, p, P);

if (j == opt.iterations - 1) {

if (opt.mode[0] > 0 and i < opt.mode[0] - 1) @memcpy(WOPT.items[i].data.data, W.data.data);

@memcpy(output.P.data, P.data); @memcpy(output.r.data, r.data); @memcpy(output.p.data, p.data); output.Ekin = Ekin; output.Epot = Epot;

for (0..nstate) |k| if (print) {
try std.io.getStdOut().writer().print("{s}FINAL POPULATION OF STATE {d:2}: {d:.6}\n", .{if (k == 0) "\n" else "", k, output.P.at(k, k)});
};
}
}
}

for (0..nstate) |i| {
if (print) {try std.io.getStdOut().writer().print("{s}FINAL POPULATION OF STATE {d:2}: {d:.6}\n", .{if (i == 0) "\n" else "", i, output.P.at(i, i)});}
for (R.items) |*e| {e.deinit();} for (K.items) |*e| {e.deinit();} for (V.items) |*e| {e.deinit();} for (VA.items) |*e| {e.deinit();} for (VC.items) |*e| {e.deinit();} for (WOPT.items) |*e| {e.deinit();}
}

try writeResults(T, opt, pop, ekin, epot, etot, position, momentum, wavefunction); return output;
Expand Down
22 changes: 11 additions & 11 deletions test/quantumdynamics.zig
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ test "qdyn_imag_harmonic1D_1" {

opt.initial_conditions.state = 0; opt.initial_conditions.position = &[_]f64{1}; opt.initial_conditions.momentum = &[_]f64{0}; opt.initial_conditions.mass = 1;

opt.adiabatic = false; opt.imaginary = true; opt.iterations = 300; opt.time_step = 0.1;
opt.adiabatic = false; opt.iterations = 300; opt.mode = &[_]u32{1, 0}; opt.time_step = 0.1;

const output = try qdn.run(f64, opt, false, std.testing.allocator); defer output.deinit();

Expand All @@ -26,7 +26,7 @@ test "qdyn_imag_harmonic2D_1" {

opt.initial_conditions.state = 0; opt.initial_conditions.position = &[_]f64{1, 1}; opt.initial_conditions.momentum = &[_]f64{0, 0}; opt.initial_conditions.mass = 1;

opt.adiabatic = false; opt.imaginary = true; opt.iterations = 300; opt.time_step = 0.1;
opt.adiabatic = false; opt.iterations = 300; opt.mode = &[_]u32{1, 0}; opt.time_step = 0.1;

const output = try qdn.run(f64, opt, false, std.testing.allocator); defer output.deinit();

Expand All @@ -40,7 +40,7 @@ test "qdyn_imag_harmonic3D_1" {

opt.initial_conditions.state = 0; opt.initial_conditions.position = &[_]f64{1, 1, 1}; opt.initial_conditions.momentum = &[_]f64{0, 0, 0}; opt.initial_conditions.mass = 1;

opt.adiabatic = false; opt.imaginary = true; opt.iterations = 300; opt.time_step = 0.1;
opt.adiabatic = false; opt.iterations = 300; opt.mode = &[_]u32{1, 0}; opt.time_step = 0.1;

const output = try qdn.run(f64, opt, false, std.testing.allocator); defer output.deinit();

Expand All @@ -59,7 +59,7 @@ test "qdyn_real_doubleState1D_1" {

opt.initial_conditions.state = 1; opt.initial_conditions.position = &[_]f64{-10}; opt.initial_conditions.momentum = &[_]f64{15}; opt.initial_conditions.mass = 2000;

opt.adiabatic = true; opt.imaginary = false; opt.iterations = 300; opt.time_step = 10;
opt.adiabatic = true; opt.iterations = 300; opt.mode = &[_]u32{0, 1}; opt.time_step = 10;

const output = try qdn.run(f64, opt, false, std.testing.allocator); defer output.deinit();

Expand All @@ -78,7 +78,7 @@ test "qdyn_real_doubleState1D_2" {

opt.initial_conditions.state = 1; opt.initial_conditions.position = &[_]f64{-10}; opt.initial_conditions.momentum = &[_]f64{15}; opt.initial_conditions.mass = 2000;

opt.adiabatic = true; opt.imaginary = false; opt.iterations = 300; opt.time_step = 10;
opt.adiabatic = true; opt.iterations = 300; opt.mode = &[_]u32{0, 1}; opt.time_step = 10;

const output = try qdn.run(f64, opt, false, std.testing.allocator); defer output.deinit();

Expand All @@ -98,7 +98,7 @@ test "qdyn_real_tripleState1D_1" {

opt.initial_conditions.state = 2; opt.initial_conditions.position = &[_]f64{-10}; opt.initial_conditions.momentum = &[_]f64{15}; opt.initial_conditions.mass = 2000;

opt.adiabatic = true; opt.imaginary = false; opt.iterations = 300; opt.time_step = 10;
opt.adiabatic = true; opt.iterations = 300; opt.mode = &[_]u32{0, 1}; opt.time_step = 10;

const output = try qdn.run(f64, opt, false, std.testing.allocator); defer output.deinit();

Expand All @@ -118,7 +118,7 @@ test "qdyn_real_tripleState1D_2" {

opt.initial_conditions.state = 2; opt.initial_conditions.position = &[_]f64{-10}; opt.initial_conditions.momentum = &[_]f64{15}; opt.initial_conditions.mass = 2000;

opt.adiabatic = true; opt.imaginary = false; opt.iterations = 300; opt.time_step = 10;
opt.adiabatic = true; opt.iterations = 300; opt.mode = &[_]u32{0, 1}; opt.time_step = 10;

const output = try qdn.run(f64, opt, false, std.testing.allocator); defer output.deinit();

Expand All @@ -138,7 +138,7 @@ test "qdyn_real_tripleState1D_3" {

opt.initial_conditions.state = 1; opt.initial_conditions.position = &[_]f64{-10}; opt.initial_conditions.momentum = &[_]f64{15}; opt.initial_conditions.mass = 2000;

opt.adiabatic = true; opt.imaginary = false; opt.iterations = 300; opt.time_step = 10;
opt.adiabatic = true; opt.iterations = 300; opt.mode = &[_]u32{0, 1}; opt.time_step = 10;

const output = try qdn.run(f64, opt, false, std.testing.allocator); defer output.deinit();

Expand All @@ -157,7 +157,7 @@ test "qdyn_real_tully1D_1" {

opt.initial_conditions.state = 1; opt.initial_conditions.position = &[_]f64{-10}; opt.initial_conditions.momentum = &[_]f64{15}; opt.initial_conditions.mass = 2000;

opt.adiabatic = true; opt.imaginary = false; opt.iterations = 300; opt.time_step = 10;
opt.adiabatic = true; opt.iterations = 300; opt.mode = &[_]u32{0, 1}; opt.time_step = 10;

const output = try qdn.run(f64, opt, false, std.testing.allocator); defer output.deinit();

Expand All @@ -176,7 +176,7 @@ test "qdyn_real_tully1D_2" {

opt.initial_conditions.state = 1; opt.initial_conditions.position = &[_]f64{-10}; opt.initial_conditions.momentum = &[_]f64{15}; opt.initial_conditions.mass = 2000;

opt.adiabatic = true; opt.imaginary = false; opt.iterations = 300; opt.time_step = 10;
opt.adiabatic = true; opt.iterations = 300; opt.mode = &[_]u32{0, 1}; opt.time_step = 10;

const output = try qdn.run(f64, opt, false, std.testing.allocator); defer output.deinit();

Expand All @@ -195,7 +195,7 @@ test "qdyn_real_tully1D_3" {

opt.initial_conditions.state = 1; opt.initial_conditions.position = &[_]f64{-10}; opt.initial_conditions.momentum = &[_]f64{15}; opt.initial_conditions.mass = 2000;

opt.adiabatic = true; opt.imaginary = false; opt.iterations = 300; opt.time_step = 10;
opt.adiabatic = true; opt.iterations = 300; opt.mode = &[_]u32{0, 1}; opt.time_step = 10;

const output = try qdn.run(f64, opt, false, std.testing.allocator); defer output.deinit();

Expand Down

0 comments on commit 337ee88

Please sign in to comment.