Skip to content

Commit

Permalink
fix stricter array in electron funcs
Browse files Browse the repository at this point in the history
  • Loading branch information
@johnomotani
johnomotani committed Jan 27, 2025
1 parent dcdcf15 commit 5e97463
Show file tree
Hide file tree
Showing 2 changed files with 42 additions and 42 deletions.
28 changes: 14 additions & 14 deletions moment_kinetics/src/electron_vpa_advection.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -154,31 +154,31 @@ function add_electron_vpa_advection_to_Jacobian!(jacobian_matrix, f, dens, upar,
# + w_∥*1/2*source_density_amplitude/n) * dg/dw_∥
if include ∈ (:all, :explicit_v)
if ielement_vpa == 1 && igrid_vpa == 1
jacobian_matrix[row,(iz-1)*v_size+(ivperp-1)*vpa.n+icolumn_min_vpa:(iz-1)*v_size+(ivperp-1)*vpa.n+icolumn_max_vpa] .+=
dt * vpa_speed * vpa_Dmat[1,:] ./ vpa_element_scale[ielement_vpa]
@. jacobian_matrix[row,(iz-1)*v_size+(ivperp-1)*vpa.n+icolumn_min_vpa:(iz-1)*v_size+(ivperp-1)*vpa.n+icolumn_max_vpa] +=
dt * vpa_speed * vpa_Dmat[1,:] / vpa_element_scale[ielement_vpa]
elseif ielement_vpa == vpa.nelement_local && igrid_vpa == vpa.ngrid
jacobian_matrix[row,(iz-1)*v_size+(ivperp-1)*vpa.n+icolumn_min_vpa:(iz-1)*v_size+(ivperp-1)*vpa.n+icolumn_max_vpa] .+=
dt * vpa_speed * vpa_Dmat[end,:] ./ vpa_element_scale[ielement_vpa]
@. jacobian_matrix[row,(iz-1)*v_size+(ivperp-1)*vpa.n+icolumn_min_vpa:(iz-1)*v_size+(ivperp-1)*vpa.n+icolumn_max_vpa] +=
dt * vpa_speed * vpa_Dmat[end,:] / vpa_element_scale[ielement_vpa]
elseif igrid_vpa == vpa.ngrid
# Note igrid_vpa is only ever 1 when ielement_vpa==1, because
# of the way element boundaries are counted.
icolumn_min_vpa_next = vpa.imin[ielement_vpa+1] - 1
icolumn_max_vpa_next = vpa.imax[ielement_vpa+1]
if vpa_speed < 0.0
jacobian_matrix[row,(iz-1)*v_size+(ivperp-1)*vpa.n+icolumn_min_vpa_next:(iz-1)*v_size+(ivperp-1)*vpa.n+icolumn_max_vpa_next] .+=
dt * vpa_speed * vpa_Dmat[1,:] ./ vpa_element_scale[ielement_vpa+1]
@. jacobian_matrix[row,(iz-1)*v_size+(ivperp-1)*vpa.n+icolumn_min_vpa_next:(iz-1)*v_size+(ivperp-1)*vpa.n+icolumn_max_vpa_next] +=
dt * vpa_speed * vpa_Dmat[1,:] / vpa_element_scale[ielement_vpa+1]
elseif vpa_speed > 0.0
jacobian_matrix[row,(iz-1)*v_size+(ivperp-1)*vpa.n+icolumn_min_vpa:(iz-1)*v_size+(ivperp-1)*vpa.n+icolumn_max_vpa] .+=
dt * vpa_speed * vpa_Dmat[end,:] ./ vpa_element_scale[ielement_vpa]
@. jacobian_matrix[row,(iz-1)*v_size+(ivperp-1)*vpa.n+icolumn_min_vpa:(iz-1)*v_size+(ivperp-1)*vpa.n+icolumn_max_vpa] +=
dt * vpa_speed * vpa_Dmat[end,:] / vpa_element_scale[ielement_vpa]
else
jacobian_matrix[row,(iz-1)*v_size+(ivperp-1)*vpa.n+icolumn_min_vpa:(iz-1)*v_size+(ivperp-1)*vpa.n+icolumn_max_vpa] .+=
dt * vpa_speed * 0.5 * vpa_Dmat[end,:] ./ vpa_element_scale[ielement_vpa]
jacobian_matrix[row,(iz-1)*v_size+(ivperp-1)*vpa.n+icolumn_min_vpa_next:(iz-1)*v_size+(ivperp-1)*vpa.n+icolumn_max_vpa_next] .+=
dt * vpa_speed * 0.5 * vpa_Dmat[1,:] ./ vpa_element_scale[ielement_vpa+1]
@. jacobian_matrix[row,(iz-1)*v_size+(ivperp-1)*vpa.n+icolumn_min_vpa:(iz-1)*v_size+(ivperp-1)*vpa.n+icolumn_max_vpa] +=
dt * vpa_speed * 0.5 * vpa_Dmat[end,:] / vpa_element_scale[ielement_vpa]
@. jacobian_matrix[row,(iz-1)*v_size+(ivperp-1)*vpa.n+icolumn_min_vpa_next:(iz-1)*v_size+(ivperp-1)*vpa.n+icolumn_max_vpa_next] +=
dt * vpa_speed * 0.5 * vpa_Dmat[1,:] / vpa_element_scale[ielement_vpa+1]
end
else
jacobian_matrix[row,(iz-1)*v_size+(ivperp-1)*vpa.n+icolumn_min_vpa:(iz-1)*v_size+(ivperp-1)*vpa.n+icolumn_max_vpa] .+=
dt * vpa_speed * vpa_Dmat[igrid_vpa,:] ./ vpa_element_scale[ielement_vpa]
@. jacobian_matrix[row,(iz-1)*v_size+(ivperp-1)*vpa.n+icolumn_min_vpa:(iz-1)*v_size+(ivperp-1)*vpa.n+icolumn_max_vpa] +=
dt * vpa_speed * vpa_Dmat[igrid_vpa,:] / vpa_element_scale[ielement_vpa]
end
end
# q = 2*p*vth*∫dw_∥ w_∥^3 g
Expand Down
56 changes: 28 additions & 28 deletions moment_kinetics/src/electron_z_advection.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -121,31 +121,31 @@ function add_electron_z_advection_to_Jacobian!(jacobian_matrix, f, dens, upar, p
# Contributions from (w_∥*vth + upar)*dg/dz
if include ∈ (:all, :explicit_z)
if ielement_z == 1 && igrid_z == 1
jacobian_matrix[row,(icolumn_min_z-1)*v_size+v_remainder:v_size:(icolumn_max_z-1)*v_size+v_remainder] .+=
dt * this_z_speed * z_Dmat[1,:] ./ z_element_scale[ielement_z]
@. jacobian_matrix[row,(icolumn_min_z-1)*v_size+v_remainder:v_size:(icolumn_max_z-1)*v_size+v_remainder] +=
dt * this_z_speed * z_Dmat[1,:] / z_element_scale[ielement_z]
elseif ielement_z == z.nelement_local && igrid_z == z.ngrid
jacobian_matrix[row,(icolumn_min_z-1)*v_size+v_remainder:v_size:(icolumn_max_z-1)*v_size+v_remainder] .+=
dt * this_z_speed * z_Dmat[end,:] ./ z_element_scale[ielement_z]
@. jacobian_matrix[row,(icolumn_min_z-1)*v_size+v_remainder:v_size:(icolumn_max_z-1)*v_size+v_remainder] +=
dt * this_z_speed * z_Dmat[end,:] / z_element_scale[ielement_z]
elseif igrid_z == z.ngrid
# Note igrid_z is only ever 1 when ielement_z==1, because
# of the way element boundaries are counted.
icolumn_min_z_next = z.imin[ielement_z+1] - 1
icolumn_max_z_next = z.imax[ielement_z+1]
if this_z_speed < 0.0
jacobian_matrix[row,(icolumn_min_z_next-1)*v_size+v_remainder:v_size:(icolumn_max_z_next-1)*v_size+v_remainder] .+=
dt * this_z_speed * z_Dmat[1,:] ./ z_element_scale[ielement_z+1]
@. jacobian_matrix[row,(icolumn_min_z_next-1)*v_size+v_remainder:v_size:(icolumn_max_z_next-1)*v_size+v_remainder] +=
dt * this_z_speed * z_Dmat[1,:] / z_element_scale[ielement_z+1]
elseif this_z_speed > 0.0
jacobian_matrix[row,(icolumn_min_z-1)*v_size+v_remainder:v_size:(icolumn_max_z-1)*v_size+v_remainder] .+=
dt * this_z_speed * z_Dmat[end,:] ./ z_element_scale[ielement_z]
@. jacobian_matrix[row,(icolumn_min_z-1)*v_size+v_remainder:v_size:(icolumn_max_z-1)*v_size+v_remainder] +=
dt * this_z_speed * z_Dmat[end,:] / z_element_scale[ielement_z]
else
jacobian_matrix[row,(icolumn_min_z-1)*v_size+v_remainder:v_size:(icolumn_max_z-1)*v_size+v_remainder] .+=
dt * this_z_speed * 0.5 * z_Dmat[end,:] ./ z_element_scale[ielement_z]
jacobian_matrix[row,(icolumn_min_z_next-1)*v_size+v_remainder:v_size:(icolumn_max_z_next-1)*v_size+v_remainder] .+=
dt * this_z_speed * 0.5 * z_Dmat[1,:] ./ z_element_scale[ielement_z+1]
@. jacobian_matrix[row,(icolumn_min_z-1)*v_size+v_remainder:v_size:(icolumn_max_z-1)*v_size+v_remainder] +=
dt * this_z_speed * 0.5 * z_Dmat[end,:] / z_element_scale[ielement_z]
@. jacobian_matrix[row,(icolumn_min_z_next-1)*v_size+v_remainder:v_size:(icolumn_max_z_next-1)*v_size+v_remainder] +=
dt * this_z_speed * 0.5 * z_Dmat[1,:] / z_element_scale[ielement_z+1]
end
else
jacobian_matrix[row,(icolumn_min_z-1)*v_size+v_remainder:v_size:(icolumn_max_z-1)*v_size+v_remainder] .+=
dt * this_z_speed * z_Dmat[igrid_z,:] ./ z_element_scale[ielement_z]
@. jacobian_matrix[row,(icolumn_min_z-1)*v_size+v_remainder:v_size:(icolumn_max_z-1)*v_size+v_remainder] +=
dt * this_z_speed * z_Dmat[igrid_z,:] / z_element_scale[ielement_z]
end
end
# vth = sqrt(2*p/n/me)
Expand Down Expand Up @@ -192,31 +192,31 @@ function add_electron_z_advection_to_z_only_Jacobian!(

# Contributions from (w_∥*vth + upar)*dg/dz
if ielement_z == 1 && igrid_z == 1
jacobian_matrix[row,icolumn_min_z:icolumn_max_z] .+=
dt * this_z_speed * z_Dmat[1,:] ./ z_element_scale[ielement_z]
@. jacobian_matrix[row,icolumn_min_z:icolumn_max_z] +=
dt * this_z_speed * z_Dmat[1,:] / z_element_scale[ielement_z]
elseif ielement_z == z.nelement_local && igrid_z == z.ngrid
jacobian_matrix[row,icolumn_min_z:icolumn_max_z] .+=
dt * this_z_speed * z_Dmat[end,:] ./ z_element_scale[ielement_z]
@. jacobian_matrix[row,icolumn_min_z:icolumn_max_z] +=
dt * this_z_speed * z_Dmat[end,:] / z_element_scale[ielement_z]
elseif igrid_z == z.ngrid
# Note igrid_z is only ever 1 when ielement_z==1, because
# of the way element boundaries are counted.
icolumn_min_z_next = z.imin[ielement_z+1] - 1
icolumn_max_z_next = z.imax[ielement_z+1]
if this_z_speed < 0.0
jacobian_matrix[row,icolumn_min_z_next:icolumn_max_z_next] .+=
dt * this_z_speed * z_Dmat[1,:] ./ z_element_scale[ielement_z+1]
@. jacobian_matrix[row,icolumn_min_z_next:icolumn_max_z_next] +=
dt * this_z_speed * z_Dmat[1,:] / z_element_scale[ielement_z+1]
elseif this_z_speed > 0.0
jacobian_matrix[row,icolumn_min_z:icolumn_max_z] .+=
dt * this_z_speed * z_Dmat[end,:] ./ z_element_scale[ielement_z]
@. jacobian_matrix[row,icolumn_min_z:icolumn_max_z] +=
dt * this_z_speed * z_Dmat[end,:] / z_element_scale[ielement_z]
else
jacobian_matrix[row,icolumn_min_z:icolumn_max_z] .+=
dt * this_z_speed * 0.5 * z_Dmat[end,:] ./ z_element_scale[ielement_z]
jacobian_matrix[row,icolumn_min_z_next:icolumn_max_z_next] .+=
dt * this_z_speed * 0.5 * z_Dmat[1,:] ./ z_element_scale[ielement_z+1]
@. jacobian_matrix[row,icolumn_min_z:icolumn_max_z] +=
dt * this_z_speed * 0.5 * z_Dmat[end,:] / z_element_scale[ielement_z]
@. jacobian_matrix[row,icolumn_min_z_next:icolumn_max_z_next] +=
dt * this_z_speed * 0.5 * z_Dmat[1,:] / z_element_scale[ielement_z+1]
end
else
jacobian_matrix[row,icolumn_min_z:icolumn_max_z] .+=
dt * this_z_speed * z_Dmat[igrid_z,:] ./ z_element_scale[ielement_z]
@. jacobian_matrix[row,icolumn_min_z:icolumn_max_z] +=
dt * this_z_speed * z_Dmat[igrid_z,:] / z_element_scale[ielement_z]
end
end

Expand Down

0 comments on commit 5e97463

Please sign in to comment.