Skip to content

Commit

Permalink
Allowing Ekin array. Documenting axis option.
Browse files Browse the repository at this point in the history
  • Loading branch information
ejpaul committed Jun 11, 2024
1 parent af751a1 commit 8e225b3
Show file tree
Hide file tree
Showing 2 changed files with 10 additions and 34 deletions.
14 changes: 10 additions & 4 deletions src/simsopt/field/tracing.py
Original file line number Diff line number Diff line change
Expand Up @@ -261,7 +261,8 @@ def trace_particles_boozer(field: BoozerMagneticField, stz_inits: RealArray,
tmax: integration time
mass: particle mass in kg, defaults to the mass of an alpha particle
charge: charge in Coulomb, defaults to the charge of an alpha particle
Ekin: kinetic energy in Joule, defaults to 3.52MeV
Ekin: kinetic energy in Joule, defaults to 3.52MeV. Either a scalar (assumed
same energy for all particles), or ``(nparticles, )`` array.
tol: defualt tolerance for ode solver when solver-specific tolerances are not set
comm: MPI communicator to parallelize over
zetas: list of angles in [0, 2pi] for which intersection with the plane
Expand All @@ -283,7 +284,9 @@ def trace_particles_boozer(field: BoozerMagneticField, stz_inits: RealArray,
`zetas_stop`: whether to stop if hit provided zeta planes
`vpars_stop`: whether to stop if hit provided vpar planes
default solver (rk45) specific options are
`axis`: TODO
`axis`: If `True`, tracing is performed in coordinates :math:`x = \sqrt{s}\cos(\theta)`
and :math:`y = \sqrt{s}\sin(\theta)` to avoid a coordinate singularity on
the magnetic axis. If `False`, tracing is performed in :math:`(s,\theta)`.
`reltol`: relative tolerance for adaptive ode solver
`abstol`: absolute tolerance for adaptive ode solver
symplectic solver specific options are
Expand Down Expand Up @@ -341,7 +344,10 @@ def trace_particles_boozer(field: BoozerMagneticField, stz_inits: RealArray,
assert stz_inits.shape[0] == len(parallel_speeds)
speed_par = parallel_speeds
m = mass
speed_total = sqrt(2*Ekin/m) # Ekin = 0.5 * m * v^2 <=> v = sqrt(2*Ekin/m)
assert np.isscalar(Ekin) or (len(Ekin) == len(parallel_speeds))
if np.isscalar(Ekin):
Ekin = Ekin * np.ones((len(parallel_speeds),))
speed_total = np.sqrt(2*Ekin/m) # Ekin = 0.5 * m * v^2 <=> v = sqrt(2*Ekin/m)
mode = mode.lower()
assert mode in ['gc', 'gc_vac', 'gc_nok']

Expand All @@ -352,7 +358,7 @@ def trace_particles_boozer(field: BoozerMagneticField, stz_inits: RealArray,
for i in range(first, last):
res_ty, res_zeta_hit = sopp.particle_guiding_center_boozer_tracing(
field, stz_inits[i, :],
m, charge, speed_total, speed_par[i], tmax, vacuum=(mode == 'gc_vac'),
m, charge, speed_total[i], speed_par[i], tmax, vacuum=(mode == 'gc_vac'),
noK=(mode == 'gc_nok'), zetas=zetas, omegas=omegas, vpars=vpars, stopping_criteria=stopping_criteria,
**options)
if not forget_exact_path:
Expand Down
30 changes: 0 additions & 30 deletions src/simsoptpp/tracing.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -120,9 +120,6 @@ class GuidingCenterVacuumBoozerRHS {
if (axis==1) {
s = pow(ys[0],2)+pow(ys[1],2);
theta = atan2(ys[1],ys[0]);
} else if (axis==2) {
s = sqrt(pow(ys[0],2)+pow(ys[1],2));
theta = atan2(ys[1],ys[0]);
} else {
s = ys[0];
theta = ys[1];
Expand All @@ -149,9 +146,6 @@ class GuidingCenterVacuumBoozerRHS {
if (axis==1) {
dydt[0] = sdot*cos(theta)/(2*sqrt(s)) - sqrt(s) * sin(theta) * tdot;
dydt[1] = sdot*sin(theta)/(2*sqrt(s)) + sqrt(s) * cos(theta) * tdot;
} else if (axis==2) {
dydt[0] = sdot*cos(theta) - s * sin(theta) * tdot;
dydt[1] = sdot*sin(theta) + s * cos(theta) * tdot;
} else {
dydt[0] = sdot;
dydt[1] = tdot;
Expand Down Expand Up @@ -199,9 +193,6 @@ class GuidingCenterVacuumBoozerPerturbedRHS {
if (axis==1) {
s = pow(ys[0],2)+pow(ys[1],2);
theta = atan2(ys[1],ys[0]);
} else if (axis==2) {
s = sqrt(pow(ys[0],2)+pow(ys[1],2));
theta = atan2(ys[1],ys[0]);
} else {
s = ys[0];
theta = ys[1];
Expand Down Expand Up @@ -239,9 +230,6 @@ class GuidingCenterVacuumBoozerPerturbedRHS {
if (axis==1) {
dydt[0] = sdot*cos(theta)/(2*sqrt(s)) - sqrt(s) * sin(theta) * tdot;
dydt[1] = sdot*sin(theta)/(2*sqrt(s)) + sqrt(s) * cos(theta) * tdot;
} else if (axis==2) {
dydt[0] = sdot*cos(theta) - s * sin(theta) * tdot;
dydt[1] = sdot*sin(theta) + s * cos(theta) * tdot;
} else {
dydt[0] = sdot;
dydt[1] = tdot;
Expand Down Expand Up @@ -294,9 +282,6 @@ class GuidingCenterNoKBoozerPerturbedRHS {
if (axis==1) {
s = pow(ys[0],2)+pow(ys[1],2);
theta = atan2(ys[1],ys[0]);
} else if (axis==2) {
s = sqrt(pow(ys[0],2)+pow(ys[1],2));
theta = atan2(ys[1],ys[0]);
} else {
s = ys[0];
theta = ys[1];
Expand Down Expand Up @@ -340,9 +325,6 @@ class GuidingCenterNoKBoozerPerturbedRHS {
if (axis==1) {
dydt[0] = sdot*cos(theta)/(2*sqrt(s)) - sqrt(s) * sin(theta) * tdot;
dydt[1] = sdot*sin(theta)/(2*sqrt(s)) + sqrt(s) * cos(theta) * tdot;
} else if (axis==2) {
dydt[0] = sdot*cos(theta) - s * sin(theta) * tdot;
dydt[1] = sdot*sin(theta) + s * cos(theta) * tdot;
} else {
dydt[0] = sdot;
dydt[1] = tdot;
Expand Down Expand Up @@ -399,9 +381,6 @@ class GuidingCenterNoKBoozerRHS {
if (axis==1) {
s = pow(ys[0],2)+pow(ys[1],2);
theta = atan2(ys[1],ys[0]);
} else if (axis==2) {
s = sqrt(pow(ys[0],2)+pow(ys[1],2));
theta = atan2(ys[1],ys[0]);
} else {
s = ys[0];
theta = ys[1];
Expand Down Expand Up @@ -431,9 +410,6 @@ class GuidingCenterNoKBoozerRHS {
if (axis==1) {
dydt[0] = sdot*cos(theta)/(2*sqrt(s)) - sqrt(s) * sin(theta) * tdot;
dydt[1] = sdot*sin(theta)/(2*sqrt(s)) + sqrt(s) * cos(theta) * tdot;
} else if (axis==2) {
dydt[0] = sdot*cos(theta) - s * sin(theta) * tdot;
dydt[1] = sdot*sin(theta) + s * cos(theta) * tdot;
} else {
dydt[0] = sdot;
dydt[1] = tdot;
Expand Down Expand Up @@ -480,9 +456,6 @@ class GuidingCenterBoozerRHS {
if (axis==1) {
s = pow(ys[0],2)+pow(ys[1],2);
theta = atan2(ys[1],ys[0]);
} else if (axis==2) {
s = sqrt(pow(ys[0],2)+pow(ys[1],2));
theta = atan2(ys[1],ys[0]);
} else {
s = ys[0];
theta = ys[1];
Expand Down Expand Up @@ -517,9 +490,6 @@ class GuidingCenterBoozerRHS {
if (axis==1) {
dydt[0] = sdot*cos(theta)/(2*sqrt(s)) - sqrt(s) * sin(theta) * tdot;
dydt[1] = sdot*sin(theta)/(2*sqrt(s)) + sqrt(s) * cos(theta) * tdot;
} else if (axis==2) {
dydt[0] = sdot*cos(theta) - s * sin(theta) * tdot;
dydt[1] = sdot*sin(theta) + s * cos(theta) * tdot;
} else {
dydt[0] = sdot;
dydt[1] = tdot;
Expand Down

0 comments on commit 8e225b3

Please sign in to comment.