-
Notifications
You must be signed in to change notification settings - Fork 2
/
particle_func.py
118 lines (97 loc) · 3.62 KB
/
particle_func.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
# particles_func.py: ti.kernels for calculating particles
# Author: Bowen Zhao ([email protected])
from parameters import *
from constants import *
from helper_func import *
import taichi as ti
@ti.kernel
def initial_push(q: float,
mass: float,
p_field: ti.template(),
u_field: ti.template(),
Eg: ti.template(),
Bg: ti.template()):
"""
Calculate u(0 - dt/2)
"""
for idx_ptc in p_field:
pos_t = p_field[idx_ptc] - pos_ori
u_t = u_field[idx_ptc]
if ti.static(DIM == 3):
E_ptc = trilerp(Eg, pos_t)
B_ptc = trilerp(Bg, pos_t)
# Doing "half" of a boris push
gam = ti.sqrt(1. + u_t.norm_sqr())
t = q * e * B_ptc * dt / (4. * gam * mass * c)
s = 2. * t / (1. + t.norm_sqr())
u0 = -u_t + u_t.cross(t)
up = u_t + u0.cross(s)
# uL at t = -dt/2
uL = up - q * e * E_ptc * dt / (2. * mass * c)
# Write final results back to u-field
u_field[idx_ptc] = uL
else:
E_ptc = bilerp(Eg, pos_t[0:2])
B_ptc = bilerp(Bg, pos_t[0:2])
# Doing "half" of a boris push
gam = ti.sqrt(1. + u_t.norm_sqr())
t = q * e * B_ptc * dt / (4. * gam * mass * c)
s = 2. * t / (1. + t.norm_sqr())
u0 = -u_t + u_t.cross(t)
up = u_t + u0.cross(s)
# uL at t = -dt/2
uL = up - q * e * E_ptc * dt / (2. * mass * c)
# Write final results back to u-field
u_field[idx_ptc] = uL
@ti.kernel
def boris_push(q: float,
mass: float,
p_field: ti.template(),
u_field: ti.template(),
Eg: ti.template(),
Bg: ti.template()
):
"""
Calculate u(t + dt/2) based on u(t - dt/2)
"""
for idx_ptc in p_field:
pos_t = p_field[idx_ptc] - pos_ori
u_tmhalf = u_field[idx_ptc]
# Transfer fields from grid to particle
if ti.static(DIM == 3):
E_ptc = trilerp(Eg, pos_t)
B_ptc = trilerp(Bg, pos_t)
# Doing "half" of a boris push
gam = ti.sqrt(1. + u_tmhalf.norm_sqr())
um = u_tmhalf + q * e * E_ptc * dt / (2. * mass * c)
t = q * e * B_ptc * dt / (2. * gam * mass * c)
s = 2. * t / (1. + t.norm_sqr())
u0 = um + um.cross(t)
up = u_tmhalf + u0.cross(s)
# u
u_tphalf = up + q * e * E_ptc * dt / (2. * mass * c)
# Write final results back to u-field
u_field[idx_ptc] = u_tphalf
else:
E_ptc = bilerp(Eg, pos_t[0:2])
B_ptc = bilerp(Bg, pos_t[0:2])
# Doing "half" of a boris push
gam = ti.sqrt(1. + u_tmhalf.norm_sqr())
um = u_tmhalf + q * e * E_ptc * dt / (2. * mass * c)
t = q * e * B_ptc * dt / (2. * gam * mass * c)
s = 2. * t / (1. + t.norm_sqr())
u0 = um + um.cross(t)
up = u_tmhalf + u0.cross(s)
# u
u_tphalf = up + q * e * E_ptc * dt / (2. * mass * c)
# Write final results back to u-field
u_field[idx_ptc] = u_tphalf
@ti.kernel
def push_particles(p_field: ti.template(),
u_field: ti.template()):
"""
Update the positions of the particles using the 4-velocity known at t+dt/2
"""
for idx_ptc in p_field:
gam = ti.sqrt(1. + u_field[idx_ptc].norm_sqr())
p_field[idx_ptc] = p_field[idx_ptc] + (c * dt) / gam * u_field[idx_ptc]