Skip to content

Commit 95554eb

Browse files
gornalexandergornalexander
committed
New diagnostics
1 parent 71f81db commit 95554eb

File tree

6 files changed

+466
-45
lines changed

6 files changed

+466
-45
lines changed

‎plasma_checker.ipynb renamed to ‎diag-test.ipynb

+72-9
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@
1111
"from lcode2dPy.simulation.interface import Simulation\n",
1212
"from lcode2dPy.diagnostics.targets import BeamDiagnostics, PlasmaDiagnostics\n",
1313
"from lcode2dPy.config.default_config import default_config\n",
14-
"from lcode2dPy.beam.beam_generator import make_beam, Gauss, rGauss\n",
14+
"from lcode2dPy.beam_generator.beam_generator import make_beam, Gauss, rGauss\n",
1515
"import numpy as np\n",
1616
"import matplotlib.pyplot as plt\n",
1717
"import subprocess"
@@ -43,7 +43,47 @@
4343
},
4444
"outputs": [],
4545
"source": [
46-
"!rm ../c_code/*"
46+
"import numpy as np\n",
47+
"#import openpmd_api as io\n",
48+
"import shutil\n",
49+
"from pathlib import Path\n",
50+
"import matplotlib.pyplot as plt\n",
51+
"\n",
52+
"class Diagnostics:\n",
53+
" def __init__(self, diag_list):\n",
54+
" self.diag_list = diag_list\n",
55+
"\n",
56+
"class TDiagnostics():\n",
57+
" def __init__(self, t_start=0, t_end=None, period=100):\n",
58+
" self.t_start = t_start\n",
59+
" self.t_end = t_end \n",
60+
" self.period = period \n",
61+
" self.data = {}\n",
62+
" \n",
63+
" def process(self, config, t, layer_idx, steps, plasma_particles, plasma_fields, rho_beam, beam): \n",
64+
" if self.t_end and (t<self.t_start or t>self.t_end):\n",
65+
" return False\n",
66+
" elif t<self.t_start:\n",
67+
" return False\n",
68+
" if (t-self.t_start)%self.period != 0:\n",
69+
" return False\n",
70+
" return True\n",
71+
"\n",
72+
"class BeamDiagnostics(TDiagnostics):\n",
73+
" def __init__(self, t_start=0, t_end=None, period=100):\n",
74+
" super().__init__(t_start, t_end, period)\n",
75+
" def process(self, config, t, layer_idx, steps, plasma_particles, plasma_fields, rho_beam, beam):\n",
76+
" if super().process(config, t, layer_idx, steps, plasma_particles, plasma_fields, rho_beam, beam):\n",
77+
" beam_slice = beam\n",
78+
" #lost_slice = beam[1]\n",
79+
" #self.lost[t] = np.array([],dtype=particle_dtype)\n",
80+
" if layer_idx == 0:\n",
81+
" particle_dtype = np.dtype([('xi', 'f8'), ('r', 'f8'), ('p_z', 'f8'), ('p_r', 'f8'), ('M', 'f8'), ('q_m', 'f8'),\n",
82+
" ('q_norm', 'f8'), ('id', 'i8')])\n",
83+
" self.data[t] = np.array([],dtype=particle_dtype)\n",
84+
" self.data[t] = np.append(self.data[t], beam_slice.particles)\n",
85+
" #self.lost[t] = np.append(self.data[t], lost_slice.particles)\n",
86+
" #self.test[t].append(rho_beam.tolist())"
4787
]
4888
},
4989
{
@@ -53,6 +93,25 @@
5393
"Collapsed": "false"
5494
},
5595
"outputs": [],
96+
"source": [
97+
"class EveryxiDiagnostics(TDiagnostics):\n",
98+
" def __init__(self, t_start=0, t_end=None, period=1, xi_period=100):\n",
99+
" super().__init__(t_start, t_end, period)\n",
100+
" self.xi_period = xi_period\n",
101+
" def process(self, config, t, layer_idx, steps, plasma_particles, plasma_fields, rho_beam, beam):\n",
102+
" super().process(config, t, layer_idx, steps, plasma_particles, plasma_fields, rho_beam, beam)\n",
103+
" if layer_idx % self.xi_period == 0:\n",
104+
" xi_step_p = config.getfloat('time-step')\n",
105+
" print('xi={xi:.6f} Ez={Ez:e} N={N}'.format(xi=layer_idx * -xi_step_p, Ez=plasma_fields.E_z[0], N=steps))"
106+
]
107+
},
108+
{
109+
"cell_type": "code",
110+
"execution_count": 5,
111+
"metadata": {
112+
"Collapsed": "false"
113+
},
114+
"outputs": [],
56115
"source": [
57116
"# Config\n",
58117
"config = default_config\n",
@@ -74,19 +133,19 @@
74133
" pz_distr=Gauss(gamma*m_proton, gamma*m_proton*1e-4, vmin=None, vmax=None),\n",
75134
" Ipeak_kA=2*40/1000,\n",
76135
" q_m=1/m_proton,\n",
77-
" saveto=\"../c_code\")\n",
136+
" saveto=None)\n",
78137
"\n",
79-
"diagnostics = [\n",
138+
"diagnostics = [EveryxiDiagnostics(),\n",
80139
" BeamDiagnostics(period=time_limit//time_step * time_step),\n",
81-
" PlasmaDiagnostics(period=time_limit//time_step * time_step)\n",
140+
"# PlasmaDiagnostics(period=time_limit//time_step * time_step)\n",
82141
"]\n",
83142
"\n",
84143
"sim = Simulation(beam_pars=beam_pars, diagnostics=diagnostics, config=config)"
85144
]
86145
},
87146
{
88147
"cell_type": "code",
89-
"execution_count": 5,
148+
"execution_count": 6,
90149
"metadata": {
91150
"Collapsed": "false"
92151
},
@@ -96,7 +155,8 @@
96155
"output_type": "stream",
97156
"text": [
98157
"Number of particles: 10093\n",
99-
"Number of particles in the middle layer: 1261\n"
158+
"Number of particles in the middle layer: 1261\n",
159+
"xi=-0.000000 Ez=-3.149575e-09 N=1\n"
100160
]
101161
}
102162
],
@@ -200,7 +260,9 @@
200260
{
201261
"cell_type": "code",
202262
"execution_count": null,
203-
"metadata": {},
263+
"metadata": {
264+
"Collapsed": "false"
265+
},
204266
"outputs": [
205267
{
206268
"ename": "Error",
@@ -232,7 +294,8 @@
232294
"hash": "40d3a090f54c6569ab1632332b64b2c03c39dcf918b08424e98f38b5ae0af88f"
233295
},
234296
"kernelspec": {
235-
"display_name": "Python 3.7.4 64-bit ('base': conda)",
297+
"display_name": "Python 3",
298+
"language": "python",
236299
"name": "python3"
237300
},
238301
"language_info": {

‎lcode2dPy/diagnostics/diagnostics_class.py

+11
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
class Diagnostics2d:
2+
def __init__(self, dt_diag, dxi_diag):
3+
self.config = None
4+
self.dt_diag = dt_diag
5+
self.dxi_diag = dxi_diag
6+
7+
def process(self, t, layer_idx, plasma_particles, plasma_fields, rho_beam, beam_slice):
8+
for diag_name in self.dxi_diag.keys():
9+
diag, pars = self.dxi_diag[diag_name]
10+
diag(self, t, layer_idx, plasma_particles, plasma_fields, rho_beam, beam_slice, **pars)
11+
return None

‎lcode2dPy/diagnostics/diagnostics_list.py

+54
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
import numpy as np
2+
3+
# def E_z_diag(diagnostics, t, layer_idx, plasma_particles, plasma_fields, rho_beam, beam_slice,
4+
# t_start, t_end, period, r_selected):
5+
# time_step = diagnostics.config.getfloat('time-step')
6+
# if t < t_start or t > t_end:
7+
# return
8+
# if t <= t_start + time_step and layer_idx == 0:
9+
# diagnostics.Ez = {}
10+
# if (t - t_start) % period != 0:
11+
# return
12+
# if layer_idx == 0:
13+
# diagnostics.Ez[t] = []
14+
# r_grid_steps = plasma_fields.E_z.size
15+
# rs = np.linspace(0, diagnostics.config.getfloat('window-width'), r_grid_steps)
16+
# E_z_selected = plasma_fields.E_z[rs == r_selected][0]
17+
# diagnostics.Ez[t].append(E_z_selected)
18+
# return E_z_selected
19+
20+
class Diagnostics:
21+
def __init__(self, diag_list):
22+
self.diag_list = diag_list
23+
24+
class TDiagnostics():
25+
def __init__(self, t_start=0, t_end=None, period=100):
26+
self.t_start = t_start
27+
self.t_end = t_end
28+
self.period = period
29+
self.data = {}
30+
31+
def process(self, config, t, layer_idx, steps, plasma_particles, plasma_fields, rho_beam, beam):
32+
if self.t_end and (t<self.t_start or t>self.t_end):
33+
return False
34+
elif t<self.t_start:
35+
return False
36+
if (t-self.t_start)%self.period != 0:
37+
return False
38+
return True
39+
40+
class BeamDiagnostics(TDiagnostics):
41+
def __init__(self, t_start=0, t_end=None, period=100):
42+
super().__init__(t_start, t_end, period)
43+
def process(self, config, t, layer_idx, steps, plasma_particles, plasma_fields, rho_beam, beam):
44+
if super().process(config, t, layer_idx, steps, plasma_particles, plasma_fields, rho_beam, beam):
45+
beam_slice = beam
46+
#lost_slice = beam[1]
47+
#self.lost[t] = np.array([],dtype=particle_dtype)
48+
if layer_idx == 0:
49+
particle_dtype = np.dtype([('xi', 'f8'), ('r', 'f8'), ('p_z', 'f8'), ('p_r', 'f8'), ('M', 'f8'), ('q_m', 'f8'),
50+
('q_norm', 'f8'), ('id', 'i8')])
51+
self.data[t] = np.array([],dtype=particle_dtype)
52+
self.data[t] = np.append(self.data[t], beam_slice.particles)
53+
#self.lost[t] = np.append(self.data[t], lost_slice.particles)
54+
#self.test[t].append(rho_beam.tolist())

‎lcode2dPy/push_solver.py

+2-5
Original file line numberDiff line numberDiff line change
@@ -109,9 +109,6 @@ def step_dt(self, plasma_particles, plasma_fields, beam_source, beam_drain, t, d
109109

110110
# Every xi step diagnostics
111111
if diagnostics:
112-
diagnostics.dxi(t + self.config.getfloat('time-step'), layer_idx, plasma_particles, plasma_fields, rho_beam, stable_slice)
113-
else:
114-
if layer_idx % 100 == 0:
115-
print('xi={xi:.6f} Ez={Ez:e} N={N}'.format(xi=layer_idx * -self.xi_step_p, Ez=plasma_fields.E_z[0], N=steps))
116-
print('xi={xi:.6f} Ez={Ez:e} N={N}'.format(xi=layer_idx * -self.xi_step_p, Ez=plasma_fields.E_z[0], N=steps))
112+
for diagnostic in diagnostics:
113+
diagnostic.process(self.config, t + self.config.getfloat('time-step'), layer_idx, steps, plasma_particles, plasma_fields, rho_beam, stable_slice)
117114
return plasma_particles, plasma_fields

‎lcode2dPy/simulation/interface.py

+16-24
Original file line numberDiff line numberDiff line change
@@ -11,10 +11,6 @@
1111
from lcode2dPy.plasma.initialization import init_plasma
1212
from lcode2dPy.diagnostics.targets import MyDiagnostics
1313

14-
# Imports for 3d simulation
15-
from lcode2dPy.push_solver_3d import PushAndSolver3d
16-
from lcode2dPy.beam3d.beam import particle_dtype as beam_particle_dtype_3d
17-
1814

1915
class Simulation:
2016
def __init__(self, config=default_config, beam_generator=make_beam,
@@ -23,10 +19,7 @@ def __init__(self, config=default_config, beam_generator=make_beam,
2319
self.t_step = config.getfloat('time-step')
2420

2521
geometry = config.get('geometry')
26-
if geometry == '3d' or geometry == '3D':
27-
self.push_solver = PushAndSolver3d(self.config) # 3d
28-
self.beam_particle_dtype = beam_particle_dtype_3d
29-
elif geometry == 'circ' or geometry == 'c':
22+
if geometry == 'circ' or geometry == 'c':
3023
self.push_solver = PusherAndSolver(self.config) # circ
3124
self.beam_particle_dtype = beam_particle_dtype_2d
3225
elif geometry == 'plane':
@@ -40,14 +33,13 @@ def __init__(self, config=default_config, beam_generator=make_beam,
4033
self.beam_source = None
4134
self.beam_drain = None
4235

43-
self.diagnostics = MyDiagnostics(config, diagnostics)
36+
self.diagnostics = diagnostics
4437

4538
def step(self, N_steps):
4639
"""Compute N time steps."""
4740
# t step function, makes N_steps time steps.
4841

4942
# Beam generation
50-
self.diagnostics.config()
5143
if self.beam_source is None:
5244
beam_particles = self.beam_generator(self.config, **self.beam_pars)
5345
beam_particles = np.array(list(map(tuple, beam_particles.to_numpy())),
@@ -56,8 +48,8 @@ def step(self, N_steps):
5648
beam_slice = BeamSlice(beam_particles.size, beam_particles)
5749
self.beam_source = MemoryBeamSource(beam_slice) #TODO mpi_beam_source
5850
self.beam_drain = MemoryBeamDrain()
59-
if self.diagnostics:
60-
self.diagnostics.config = self.config
51+
# if self.diagnostics:
52+
# self.diagnostics.sim = self
6153
# Time loop
6254
for t_i in range(N_steps):
6355
fields, plasma_particles = init_plasma(self.config)
@@ -76,17 +68,17 @@ def step(self, N_steps):
7668
# self.diagnostics.every_dt()
7769

7870

79-
class Diagnostics2d:
80-
def __init__(self, dt_diag, dxi_diag):
81-
self.config = None
82-
self.dt_diag = dt_diag
83-
self.dxi_diag = dxi_diag
71+
# class Diagnostics2d:
72+
# def __init__(self, dt_diag, dxi_diag):
73+
# self.config = None
74+
# self.dt_diag = dt_diag
75+
# self.dxi_diag = dxi_diag
8476

85-
def every_dt(self):
86-
pass # TODO
77+
# def every_dt(self):
78+
# pass # TODO
8779

88-
def every_dxi(self, t, layer_idx, plasma_particles, plasma_fields, rho_beam, beam_slice):
89-
for diag_name in self.dxi_diag.keys():
90-
diag, pars = self.dxi_diag[diag_name]
91-
diag(self, t, layer_idx, plasma_particles, plasma_fields, rho_beam, beam_slice, **pars)
92-
return None
80+
# def every_dxi(self, t, layer_idx, plasma_particles, plasma_fields, rho_beam, beam_slice):
81+
# for diag_name in self.dxi_diag.keys():
82+
# diag, pars = self.dxi_diag[diag_name]
83+
# diag(self, t, layer_idx, plasma_particles, plasma_fields, rho_beam, beam_slice, **pars)
84+
# return None

0 commit comments

Comments
 (0)