aboutsummaryrefslogtreecommitdiff
path: root/Python
diff options
context:
space:
mode:
Diffstat (limited to 'Python')
-rw-r--r--Python/pywarpx/AMReX.py42
-rw-r--r--Python/pywarpx/Bucket.py11
-rw-r--r--Python/pywarpx/PGroup.py5
-rw-r--r--Python/pywarpx/PICMI.py133
-rw-r--r--Python/pywarpx/Particles.py2
-rw-r--r--Python/pywarpx/WarpX.py50
-rw-r--r--Python/pywarpx/WarpXPIC.py50
-rw-r--r--Python/pywarpx/__init__.py6
-rw-r--r--Python/pywarpx/callbacks.py481
-rw-r--r--Python/pywarpx/picmi.py388
-rw-r--r--Python/pywarpx/timestepper.py329
-rw-r--r--Python/setup.py7
12 files changed, 1004 insertions, 500 deletions
diff --git a/Python/pywarpx/AMReX.py b/Python/pywarpx/AMReX.py
deleted file mode 100644
index 4ce7ac61c..000000000
--- a/Python/pywarpx/AMReX.py
+++ /dev/null
@@ -1,42 +0,0 @@
-from .Bucket import Bucket
-
-from .WarpX import warpx
-from .Amr import amr
-from .Geometry import geometry
-from .Algo import algo
-from .Langmuirwave import langmuirwave
-from .Interpolation import interpolation
-from .Laser import laser
-from . import Particles
-from .Particles import particles, particles_list
-
-import ctypes
-from ._libwarpx import libwarpx
-from ._libwarpx import amrex_init
-
-class AMReX(object):
-
- def init(self):
- argv = ['warpx']
- argv += warpx.attrlist()
- argv += amr.attrlist()
- argv += geometry.attrlist()
- argv += algo.attrlist()
- argv += langmuirwave.attrlist()
- argv += interpolation.attrlist()
- argv += particles.attrlist()
- argv += laser.attrlist()
-
- if not particles_list:
- # --- This is needed in case only species_names has been set,
- # --- assuming that only the built in particle types are being used.
- for pstring in particles.species_names.split(' '):
- particles_list.append(getattr(Particles, pstring))
-
- for particle in particles_list:
- argv += particle.attrlist()
-
- amrex_init(argv)
-
- def finalize(self, finalize_mpi=1):
- libwarpx.amrex_finalize(finalize_mpi)
diff --git a/Python/pywarpx/Bucket.py b/Python/pywarpx/Bucket.py
index 5c429f530..cd5c81793 100644
--- a/Python/pywarpx/Bucket.py
+++ b/Python/pywarpx/Bucket.py
@@ -26,9 +26,18 @@ class Bucket(object):
"Concatenate the attributes into a string"
result = []
for attr, value in iteritems(self.argvattrs):
+ if value is None:
+ continue
# --- repr is applied to value so that for floats, all of the digits are included.
# --- The strip is then needed when value is a string.
- attrstring = '{0}.{1}={2} '.format(self.instancename, attr, repr(value).strip("'\""))
+ if isinstance(value, str):
+ rhs = value
+ elif hasattr(value, '__iter__'):
+ # --- For lists, tuples, and arrays make a space delimited string of the values
+ rhs = ' '.join(map(repr, value))
+ else:
+ rhs = value
+ attrstring = '{0}.{1} = {2}'.format(self.instancename, attr, repr(rhs).strip("'\""))
result += [attrstring]
return result
diff --git a/Python/pywarpx/PGroup.py b/Python/pywarpx/PGroup.py
index fa8572c0a..64e396d38 100644
--- a/Python/pywarpx/PGroup.py
+++ b/Python/pywarpx/PGroup.py
@@ -148,7 +148,7 @@ class PGroups(object):
self._pgroups = []
for igroup in range(self.ngroups):
self._pgroups.append(PGroup(igroup, self.ispecie))
-
+
def __iter__(self):
self.setuppgroups()
for igroup in range(self.ngroups):
@@ -158,3 +158,6 @@ class PGroups(object):
self.setuppgroups()
return self._pgroups[key]
+ def __len__(self):
+ self.setuppgroups()
+ return len(self._pgroups)
diff --git a/Python/pywarpx/PICMI.py b/Python/pywarpx/PICMI.py
deleted file mode 100644
index 1bc36e5ca..000000000
--- a/Python/pywarpx/PICMI.py
+++ /dev/null
@@ -1,133 +0,0 @@
-"""Classes following the PICMI standard
-"""
-from PICMI_Base import *
-import numpy as np
-from pywarpx import *
-
-codename = 'WarpX'
-
-def _args_to_string(*args):
- # --- Converts of sequence of number to a string that is appropriate for input.
- return ' '.join(map(repr, args))
-
-class Grid(PICMI_Grid):
- def init(self, **kw):
-
- amr.n_cell = _args_to_string(self.nx, self.ny, self.nz)
-
- # Maximum allowable size of each subdomain in the problem domain;
- # this is used to decompose the domain for parallel calculations.
- amr.max_grid_size = kw.get('max_grid_size', 32)
-
- # Maximum level in hierarchy (for now must be 0, i.e., one level in total)
- amr.max_level = kw.get('max_level', 0)
-
- # Geometry
- geometry.coord_sys = kw.get('coord_sys', 0) # 0: Cartesian
- geometry.is_periodic = '%d %d %d'%(self.bcxmin=='periodic', self.bcymin=='periodic', self.bczmin=='periodic') # Is periodic?
- geometry.prob_lo = _args_to_string(self.xmin, self.ymin, self.zmin) # physical domain
- geometry.prob_hi = _args_to_string(self.xmax, self.ymax, self.zmax)
-
- if self.moving_window_velocity is not None and np.any(self.moving_window_velocity != 0):
- warpx.do_moving_window = 1
- if self.moving_window_velocity[0] != 0.:
- warpx.moving_window_dir = 'x'
- warpx.moving_window_v = self.moving_window_velocity[0]/clight # in units of the speed of light
- if self.moving_window_velocity[1] != 0.:
- warpx.moving_window_dir = 'y'
- warpx.moving_window_v = self.moving_window_velocity[1]/clight # in units of the speed of light
- if self.moving_window_velocity[2] != 0.:
- warpx.moving_window_dir = 'z'
- warpx.moving_window_v = self.moving_window_velocity[2]/clight # in units of the speed of light
-
- def getmins(self, **kw):
- return np.array([warpx.getProbLo(0), warpx.getProbLo(1), warpx.getProbLo(2)])
-
- def getmaxs(self, **kw):
- return np.array([warpx.getProbHi(0), warpx.getProbHi(1), warpx.getProbHi(2)])
-
- def getxmin(self):
- return warpx.getProbLo(0)
- def getxmax(self):
- return warpx.getProbHi(0)
- def getymin(self):
- return warpx.getProbLo(1)
- def getymax(self):
- return warpx.getProbHi(1)
- def getzmin(self):
- return warpx.getProbLo(2)
- def getzmax(self):
- return warpx.getProbHi(2)
-
-
-class EM_solver(PICMI_EM_solver):
- def init(self, **kw):
-
- if self.current_deposition_algo is not None:
- algo.current_deposition = self.current_deposition_algo
- if self.charge_deposition_algo is not None:
- algo.charge_deposition = self.charge_deposition_algo
- if self.field_gathering_algo is not None:
- algo.field_gathering = self.field_gathering_algo
- if self.particle_pusher_algo is not None:
- algo.particle_pusher = self.particle_pusher_algo
-
-
-class Gaussian_laser(PICMI_Gaussian_laser):
- def init(self, **kw):
-
- warpx.use_laser = 1
- laser.profile = "Gaussian"
- laser.position = _args_to_string(self.antenna_x0, self.antenna_y0, self.antenna_z0) # This point is on the laser plane
- laser.direction = _args_to_string(self.antenna_xvec, self.antenna_yvec, self.antenna_zvec) # The plane normal direction
- laser.polarization = _args_to_string(np.cos(self.pol_angle), np.sin(self.pol_angle), 0.) # The main polarization vector
- laser.e_max = self.E0 # Maximum amplitude of the laser field (in V/m)
- laser.profile_waist = self.waist # The waist of the laser (in meters)
- laser.profile_duration = self.duration # The duration of the laser (in seconds)
- laser.profile_t_peak = self.t_peak # The time at which the laser reaches its peak (in seconds)
- laser.profile_focal_distance = self.focal_position - self.antenna_z0 # Focal distance from the antenna (in meters)
- laser.wavelength = self.wavelength # The wavelength of the laser (in meters)
-
-
-class Species(PICMI_Species):
- def init(self, **kw):
-
- self.species_number = particles.nspecies
- particles.nspecies = particles.nspecies + 1
- particles.species_names = particles.species_names + ' ' + self.name
-
- self.bucket = Bucket.Bucket(self.name, mass=self.mass, charge=self.charge, injection_style = 'python')
- Particles.particles_list.append(self.bucket)
-
- def add_particles(self, n=None,
- x=None, y=None, z=None,
- ux=None, uy=None, uz=None, w=None,
- unique_particles=None, **kw):
- pid = np.array([w]).T
- add_particles(self.species_number, x, y, z, ux, uy, uz, pid, unique_particles)
-
-
-class Simulation(PICMI_Simulation):
- def set_warpx_attr(self, warpx_obj, attr, kw):
- value = kw.get(attr, None)
- if value is not None:
- setattr(warpx_obj, attr, value)
- setattr(self, attr, value)
-
- def init(self, **kw):
-
- warpx.verbose = self.verbose
- warpx.cfl = self.cfl
- amr.plot_int = self.plot_int
-
- self.amrex = AMReX()
- self.amrex.init()
- warpx.init()
-
- def step(self, nsteps=-1):
- warpx.evolve(nsteps)
-
- def finalize(self):
- warpx.finalize()
- self.amrex.finalize()
-
diff --git a/Python/pywarpx/Particles.py b/Python/pywarpx/Particles.py
index e7bac7e10..395a888d7 100644
--- a/Python/pywarpx/Particles.py
+++ b/Python/pywarpx/Particles.py
@@ -1,6 +1,6 @@
from .Bucket import Bucket
-particles = Bucket('particles', nspecies=0, species_names='')
+particles = Bucket('particles', nspecies=0, species_names=None)
particles_list = []
electrons = Bucket('electrons')
diff --git a/Python/pywarpx/WarpX.py b/Python/pywarpx/WarpX.py
index f0a31f3f9..7a05c40fb 100644
--- a/Python/pywarpx/WarpX.py
+++ b/Python/pywarpx/WarpX.py
@@ -1,19 +1,56 @@
from .Bucket import Bucket
+
+from .Amr import amr
+from .Geometry import geometry
+from .Algo import algo
+from .Langmuirwave import langmuirwave
+from .Interpolation import interpolation
+from .Laser import laser
+from . import Particles
+from .Particles import particles, particles_list
+
+import ctypes
from ._libwarpx import libwarpx
+from ._libwarpx import amrex_init
class WarpX(Bucket):
"""
A Python wrapper for the WarpX C++ class
"""
+ def create_argv_list(self):
+ argv = []
+ argv += warpx.attrlist()
+ argv += amr.attrlist()
+ argv += geometry.attrlist()
+ argv += algo.attrlist()
+ argv += langmuirwave.attrlist()
+ argv += interpolation.attrlist()
+ argv += particles.attrlist()
+ argv += laser.attrlist()
+
+ if not particles_list:
+ # --- This is needed in case only species_names has been set,
+ # --- assuming that only the built in particle types are being used.
+ for pstring in particles.species_names.split(' '):
+ particles_list.append(getattr(Particles, pstring))
+
+ for particle in particles_list:
+ argv += particle.attrlist()
+
+ return argv
+
def init(self):
+ argv = ['warpx'] + self.create_argv_list()
+ amrex_init(argv)
libwarpx.warpx_init()
def evolve(self, nsteps=-1):
libwarpx.warpx_evolve(nsteps)
- def finalize(self):
+ def finalize(self, finalize_mpi=1):
libwarpx.warpx_finalize()
+ libwarpx.amrex_finalize(finalize_mpi)
def getProbLo(self, direction):
return libwarpx.warpx_getProbLo(direction)
@@ -21,4 +58,15 @@ class WarpX(Bucket):
def getProbHi(self, direction):
return libwarpx.warpx_getProbHi(direction)
+ def write_inputs(self, filename='inputs', **kw):
+ argv = self.create_argv_list()
+ with open(filename, 'w') as ff:
+
+ for k, v in kw.iteritems():
+ ff.write('{0} = {1}\n'.format(k, v))
+
+ for arg in argv:
+ ff.write('{0}\n'.format(arg))
+
warpx = WarpX('warpx')
+
diff --git a/Python/pywarpx/WarpXPIC.py b/Python/pywarpx/WarpXPIC.py
new file mode 100644
index 000000000..77ab8464f
--- /dev/null
+++ b/Python/pywarpx/WarpXPIC.py
@@ -0,0 +1,50 @@
+from warp.run_modes.timestepper import PICAPI
+from ._libwarpx import libwarpx
+
+class WarpXPIC(PICAPI):
+
+ def get_time(self):
+ return libwarpx.warpx_gett_new(0)
+
+ def set_time(self, time):
+ for i in range(libwarpx.warpx_finestLevel()+1):
+ libwarpx.warpx_sett_new(i, time)
+
+ def get_step_size(self):
+ libwarpx.warpx_ComputeDt()
+ return libwarpx.warpx_getdt(0)
+
+ def get_step_number(self):
+ return libwarpx.warpx_getistep(0)
+
+ def set_step_number(self, it):
+ for i in range(libwarpx.warpx_finestLevel()+1):
+ libwarpx.warpx_setistep(i, it)
+
+ def push_positions(self, dt):
+ libwarpx.warpx_PushX(0, dt)
+
+ def push_velocities_withE(self, dt):
+ libwarpx.warpx_EPushV(0, dt)
+
+ def push_velocities_withB(self, dt):
+ libwarpx.warpx_BPushV(0, dt)
+
+ def get_self_fields(self):
+ libwarpx.warpx_FieldGather(0)
+
+ def calculate_source(self):
+ libwarpx.warpx_CurrentDeposition(0)
+
+ def push_Efields(self, dt):
+ libwarpx.warpx_EvolveE(0, dt)
+ libwarpx.warpx_FillBoundaryE(0, True)
+
+ def push_Bfields(self, dt):
+ libwarpx.warpx_EvolveB(0, dt)
+ libwarpx.warpx_FillBoundaryB(0, True)
+
+ def apply_particle_boundary_conditions(self):
+ libwarpx.mypc_Redistribute() # Redistribute particles
+ libwarpx.warpx_MoveWindow() # !!! not the correct place yet
+
diff --git a/Python/pywarpx/__init__.py b/Python/pywarpx/__init__.py
index bddbb78b7..14d028a45 100644
--- a/Python/pywarpx/__init__.py
+++ b/Python/pywarpx/__init__.py
@@ -5,12 +5,14 @@ from .Geometry import geometry
from .Algo import algo
from .Langmuirwave import langmuirwave
from .Interpolation import interpolation
-from .Particles import particles
+from .Particles import particles, electrons, positrons, protons
from .Laser import laser
-from .AMReX import AMReX
#from .timestepper import TimeStepper
from .PGroup import PGroup
from .PGroup import PGroups
+#from .WarpXPIC import WarpXPIC
from ._libwarpx import add_particles
+
+from .callbacks import *
diff --git a/Python/pywarpx/callbacks.py b/Python/pywarpx/callbacks.py
new file mode 100644
index 000000000..68629d164
--- /dev/null
+++ b/Python/pywarpx/callbacks.py
@@ -0,0 +1,481 @@
+"""call back operations
+=====================
+
+These are the functions which allow installing user created functions so that
+they are called at various places along the time step.
+
+For each call back type, the following three functions are defined.
+ - install___: Installs a function to be called at that specified time
+ - uninstall___: Uninstalls the function (so it won't be called anymore)
+ - isinstalled___: Checks if the function is installed
+
+These functions all take a function or instance method as an argument. Note that
+if an instance method is used, an extra reference to the method's object is saved.
+
+The install can be done using a decorator, which has the prefix "callfrom". See example below.
+
+Functions can be called at the following times:
+ - afterinit <installafterinit>: immediately after the init is complete
+ - beforeEsolve <installbeforeEsolve>: before the solve for E fields
+ - afterEsolve <installafterEsolve>: after the solve for E fields
+ - beforedeposition <installbeforedeposition>: before the particle deposition (for charge and/or current)
+ - afterdeposition <installafterdeposition>: after particle deposition (for charge and/or current)
+ - beforestep <installbeforestep>: before the time step
+ - afterstep <installafterstep>: after the time step
+ - particlescraper <installparticlescraper>: just after the particle boundary conditions are applied
+ but before lost particles are processed
+ - particleloader <installparticleloader>: at the time that the standard particle loader is called
+ - particleinjection <installparticleinjection>: called when particle injection happens, after the position
+ advance and before deposition is called, allowing a user defined
+ particle distribution to be injected each time step
+ - appliedfields <installappliedfields>: allows directly specifying any fields to be applied to the particles
+ during the advance
+
+To use a decorator, the syntax is as follows. This will install the function myplots to be called after each step.
+
+@callfromafterstep
+def myplots():
+ ppzx()
+
+This is equivalent to the following:
+
+def myplots():
+ ppzx()
+installafterstep(myplots)
+
+"""
+from __future__ import generators
+
+import types
+import copy
+import time
+import ctypes
+import sys
+import numpy
+
+from ._libwarpx import libwarpx
+
+class CallbackFunctions(object):
+ """
+ Class to handle call back function lists.
+
+ Note that for functions passed in that are methods of a class instance,
+ a full reference of the instance is saved. This extra reference means
+ that the object will not actually deleted if the user deletes the
+ original reference. This is good since the user does not need to keep
+ the reference to the object (for example it can be created using a local
+ variable in a function). It may be bad if the user thinks an object was
+ deleted, but it actually isn't since it had (unkown to the user)
+ installed a method in one of the call back lists
+ """
+
+ def __init__(self,name=None,lcallonce=0):
+ self.funcs = []
+ self.time = 0.
+ self.timers = {}
+ self.name = name
+ self.lcallonce = lcallonce
+
+ def __call__(self,*args,**kw):
+ "Call all of the functions in the list"
+ tt = self.callfuncsinlist(*args,**kw)
+ self.time = self.time + tt
+ if self.lcallonce: self.funcs = []
+
+ def clearlist(self):
+ self.funcs = []
+
+ def __nonzero__(self):
+ "Returns True if functions are installed, otherwise False"
+ return self.hasfuncsinstalled()
+
+ def __len__(self):
+ "Returns number of functions installed"
+ return len(self.funcs)
+
+ def hasfuncsinstalled(self):
+ "Checks if there are any functions installed"
+ return len(self.funcs) > 0
+
+ def _getmethodobject(self,func):
+ "For call backs that are methods, returns the method's instance"
+ return func[0]
+
+ def callbackfunclist(self):
+ "Generator returning callable functions from the list"
+ funclistcopy = copy.copy(self.funcs)
+ for f in funclistcopy:
+ if isinstance(f,list):
+ object = self._getmethodobject(f)
+ if object is None:
+ self.funcs.remove(f)
+ continue
+ result = getattr(object,f[1])
+ elif isinstance(f,basestring):
+ import __main__
+ if f in __main__.__dict__:
+ result = __main__.__dict__[f]
+ # --- If the function with the name is found, then replace the
+ # --- name in the list with the function.
+ self.funcs[self.funcs.index(f)] = result
+ else:
+ continue
+ else:
+ result = f
+ if not callable(result):
+ print("\n\nWarning: a call back was found that is not callable.")
+ if self.name is not None:
+ print("For %s"%self.name)
+ print("Only callable objects can be installed.")
+ print("It is possible that the callable's name has been overwritten")
+ print("by something not callable. This can happen during restart")
+ print("if a function name had later been used as a variable name.")
+ print(self.name)
+ if isinstance(f,basestring):
+ print("The name of the call back is %s"%f)
+ print("\n\n")
+ continue
+ yield result
+
+ def installfuncinlist(self,f):
+ "Check if the specified function is installed"
+ if isinstance(f,types.MethodType):
+ # --- If the function is a method of a class instance, then save a full
+ # --- reference to that instance and the method name.
+ finstance = f.im_self
+ fname = f.__name__
+ self.funcs.append([finstance,fname])
+ elif callable(f):
+ # --- If a function had already been installed by name, then skip the install.
+ # --- This is problematic, since no warning message is given, but it is unlikely
+ # --- to arise under normal circumstances.
+ # --- The purpose of this check is to avoid redundant installation of functions
+ # --- during a restore from a dump file. Without the check, functions that had been
+ # --- installed via a decorator would be installed an extra time since the source
+ # --- of the function contains the decoration (which is activated when the source
+ # --- is exec'd).
+ if f.__name__ not in self.funcs:
+ self.funcs.append(f)
+ else:
+ self.funcs.append(f)
+
+ def uninstallfuncinlist(self,f):
+ "Uninstall the specified function"
+ # --- An element by element search is needed
+ # --- f can be a function or method object, or a name (string).
+ # --- Note that method objects can not be removed by name.
+ funclistcopy = copy.copy(self.funcs)
+ for func in funclistcopy:
+ if f == func:
+ self.funcs.remove(f)
+ return
+ elif isinstance(func,list) and isinstance(f,types.MethodType):
+ object = self._getmethodobject(func)
+ if f.im_self is object and f.__name__ == func[1]:
+ self.funcs.remove(func)
+ return
+ elif isinstance(func,basestring):
+ if f.__name__ == func:
+ self.funcs.remove(func)
+ return
+ elif isinstance(f,basestring):
+ if isinstance(func,basestring): funcname = func
+ elif isinstance(func,list): funcname = None
+ else: funcname = func.__name__
+ if f == funcname:
+ self.funcs.remove(func)
+ return
+ raise Exception('Warning: no such function had been installed')
+
+ def isinstalledfuncinlist(self,f):
+ "Checks if the specified function is installed"
+ # --- An element by element search is needed
+ funclistcopy = copy.copy(self.funcs)
+ for func in funclistcopy:
+ if f == func:
+ return 1
+ elif isinstance(func,list) and isinstance(f,types.MethodType):
+ object = self._getmethodobject(func)
+ if f.im_self is object and f.__name__ == func[1]:
+ return 1
+ elif isinstance(func,basestring):
+ if f.__name__ == func:
+ return 1
+ return 0
+
+ def callfuncsinlist(self,*args,**kw):
+ "Call the functions in the list"
+ bb = time.time()
+ for f in self.callbackfunclist():
+ #barrier()
+ t1 = time.time()
+ f(*args,**kw)
+ #barrier()
+ t2 = time.time()
+ # --- For the timers, use the function (or method) name as the key.
+ self.timers[f.__name__] = self.timers.get(f.__name__,0.) + (t2 - t1)
+ aa = time.time()
+ return aa - bb
+
+#=============================================================================
+
+# --- Now create the actual instances.
+_afterinit = CallbackFunctions('afterinit')
+_beforeEsolve = CallbackFunctions('beforeEsolve')
+_afterEsolve = CallbackFunctions('afterEsolve')
+_beforedeposition = CallbackFunctions('beforedeposition')
+_afterdeposition = CallbackFunctions('afterdeposition')
+_particlescraper = CallbackFunctions('particlescraper')
+_particleloader = CallbackFunctions('particleloader')
+_beforestep = CallbackFunctions('beforestep')
+_afterstep = CallbackFunctions('afterstep')
+_afterrestart = CallbackFunctions('afterrestart',lcallonce=1)
+_particleinjection = CallbackFunctions('particleinjection')
+_appliedfields = CallbackFunctions('appliedfields')
+
+# --- Create the objects that can be called from C.
+# --- Note that each of the CFUNCTYPE instances need to be saved
+_CALLBACK_FUNC_0 = ctypes.CFUNCTYPE(None)
+_c_afterinit = _CALLBACK_FUNC_0(_afterinit)
+libwarpx.warpx_set_callback_py_afterinit(_c_afterinit)
+_c_beforeEsolve = _CALLBACK_FUNC_0(_beforeEsolve)
+libwarpx.warpx_set_callback_py_beforeEsolve(_c_beforeEsolve)
+_c_afterEsolve = _CALLBACK_FUNC_0(_afterEsolve)
+libwarpx.warpx_set_callback_py_afterEsolve(_c_afterEsolve)
+_c_beforedeposition = _CALLBACK_FUNC_0(_beforedeposition)
+libwarpx.warpx_set_callback_py_beforedeposition(_c_beforedeposition)
+_c_afterdeposition = _CALLBACK_FUNC_0(_afterdeposition)
+libwarpx.warpx_set_callback_py_afterdeposition(_c_afterdeposition)
+_c_particlescraper = _CALLBACK_FUNC_0(_particlescraper)
+libwarpx.warpx_set_callback_py_particlescraper(_c_particlescraper)
+_c_particleloader = _CALLBACK_FUNC_0(_particleloader)
+libwarpx.warpx_set_callback_py_particleloader(_c_particleloader)
+_c_beforestep = _CALLBACK_FUNC_0(_beforestep)
+libwarpx.warpx_set_callback_py_beforestep(_c_beforestep)
+_c_afterstep = _CALLBACK_FUNC_0(_afterstep)
+libwarpx.warpx_set_callback_py_afterstep(_c_afterstep)
+_c_afterrestart = _CALLBACK_FUNC_0(_afterrestart)
+libwarpx.warpx_set_callback_py_afterrestart(_c_afterrestart)
+_c_particleinjection = _CALLBACK_FUNC_0(_particleinjection)
+libwarpx.warpx_set_callback_py_particleinjection(_c_particleinjection)
+_c_appliedfields = _CALLBACK_FUNC_0(_appliedfields)
+libwarpx.warpx_set_callback_py_appliedfields(_c_appliedfields)
+
+#=============================================================================
+def printcallbacktimers(tmin=1.,lminmax=False,ff=None):
+ """Prints timings of installed functions.
+ - tmin=1.: only functions with time greater than tmin will be printed
+ - lminmax=False: If True, prints the min and max times over all processors
+ - ff=None: If given, timings will be written to the file object instead of stdout
+ """
+ if ff is None: ff = sys.stdout
+ for c in [_afterinit,_beforeEsolve,_afterEsolve,
+ _beforedeposition,_afterdeposition,
+ _particlescraper,
+ _particleloader,
+ _beforestep,_afterstep,
+ _afterrestart,
+ _particleinjection,
+ _appliedfields]:
+ for fname, time in c.timers.items():
+ #vlist = numpy.array(gather(time))
+ vlist = numpy.array([time])
+ #if me > 0: continue
+ vsum = numpy.sum(vlist)
+ if vsum <= tmin: continue
+ vrms = numpy.sqrt(max(0.,numpy.sum(vlist**2)/len(vlist) - (numpy.sum(vlist)/len(vlist))**2))
+ npes = 1. # Only works for one processor
+ ff.write('%20s %s %10.4f %10.4f %10.4f'%(c.name,fname,vsum,vsum/npes,vrms))
+ if lminmax:
+ vmin = numpy.min(vlist)
+ vmax = numpy.max(vlist)
+ ff.write(' %10.4f %10.4f'%(vmin,vmax))
+ it = libwarpx.warpx_getistep(0)
+ if it > 0:
+ ff.write(' %10.4f'%(vsum/npes/(it)))
+ ff.write('\n')
+
+#=============================================================================
+# ----------------------------------------------------------------------------
+def callfromafterinit(f):
+ installafterinit(f)
+ return f
+def installafterinit(f):
+ "Adds a function to the list of functions called after the init"
+ _afterinit.installfuncinlist(f)
+def uninstallafterinit(f):
+ "Removes the function from the list of functions called after the init"
+ _afterinit.uninstallfuncinlist(f)
+def isinstalledafterinit(f):
+ "Checks if the function is called after a init"
+ return _afterinit.isinstalledfuncinlist(f)
+
+# ----------------------------------------------------------------------------
+def callfrombeforeEsolve(f):
+ installbeforeEsolve(f)
+ return f
+def installbeforeEsolve(f):
+ "Adds a function to the list of functions called before an E solve"
+ _beforeEsolve.installfuncinlist(f)
+def uninstallbeforeEsolve(f):
+ "Removes the function from the list of functions called before an E solve"
+ _beforeEsolve.uninstallfuncinlist(f)
+def isinstalledbeforeEsolve(f):
+ "Checks if the function is called before an E solve"
+ return _beforeEsolve.isinstalledfuncinlist(f)
+
+# ----------------------------------------------------------------------------
+def callfromafterEsolve(f):
+ installafterEsolve(f)
+ return f
+def installafterEsolve(f):
+ "Adds a function to the list of functions called after an E solve"
+ _afterEsolve.installfuncinlist(f)
+def uninstallafterEsolve(f):
+ "Removes the function from the list of functions called after an E solve"
+ _afterEsolve.uninstallfuncinlist(f)
+def isinstalledafterEsolve(f):
+ "Checks if the function is called after an E solve"
+ return _afterEsolve.isinstalledfuncinlist(f)
+
+# ----------------------------------------------------------------------------
+def callfrombeforedeposition(f):
+ installbeforedeposition(f)
+ return f
+def installbeforedeposition(f):
+ "Adds a function to the list of functions called before a particle deposition"
+ _beforedeposition.installfuncinlist(f)
+def uninstallbeforedeposition(f):
+ "Removes the function from the list of functions called before a particle deposition"
+ _beforedeposition.uninstallfuncinlist(f)
+def isinstalledbeforedeposition(f):
+ "Checks if the function is called before a particle deposition"
+ return _beforedeposition.isinstalledfuncinlist(f)
+
+# ----------------------------------------------------------------------------
+def callfromafterdeposition(f):
+ installafterdeposition(f)
+ return f
+def installafterdeposition(f):
+ "Adds a function to the list of functions called after a particle deposition"
+ _afterdeposition.installfuncinlist(f)
+def uninstallafterdeposition(f):
+ "Removes the function from the list of functions called after a particle deposition"
+ _afterdeposition.uninstallfuncinlist(f)
+def isinstalledafterdeposition(f):
+ "Checks if the function is called after a particle deposition"
+ return _afterdeposition.isinstalledfuncinlist(f)
+
+# ----------------------------------------------------------------------------
+def callfromparticlescraper(f):
+ installparticlescraper(f)
+ return f
+def installparticlescraper(f):
+ "Adds a function to the list of functions called to scrape particles"
+ _particlescraper.installfuncinlist(f)
+def uninstallparticlescraper(f):
+ "Removes the function from the list of functions called to scrape particles"
+ _particlescraper.uninstallfuncinlist(f)
+def isinstalledparticlescraper(f):
+ "Checks if the function is called to scrape particles"
+ return _particlescraper.isinstalledfuncinlist(f)
+
+# ----------------------------------------------------------------------------
+def callfromparticleloader(f):
+ installparticleloader(f)
+ return f
+def installparticleloader(f):
+ "Adds a function to the list of functions called to load particles"
+ _particleloader.installfuncinlist(f)
+def uninstallparticleloader(f):
+ "Removes the function from the list of functions called to load particles"
+ _particleloader.uninstallfuncinlist(f)
+def isinstalledparticleloader(f):
+ "Checks if the function is called to load particles"
+ return _particleloader.isinstalledfuncinlist(f)
+
+# ----------------------------------------------------------------------------
+def callfrombeforestep(f):
+ installbeforestep(f)
+ return f
+def installbeforestep(f):
+ "Adds a function to the list of functions called before a step"
+ _beforestep.installfuncinlist(f)
+def uninstallbeforestep(f):
+ "Removes the function from the list of functions called before a step"
+ _beforestep.uninstallfuncinlist(f)
+def isinstalledbeforestep(f):
+ "Checks if the function is called before a step"
+ return _beforestep.isinstalledfuncinlist(f)
+
+# ----------------------------------------------------------------------------
+def callfromafterstep(f):
+ installafterstep(f)
+ return f
+def installafterstep(f):
+ "Adds a function to the list of functions called after a step"
+ _afterstep.installfuncinlist(f)
+def uninstallafterstep(f):
+ "Removes the function from the list of functions called after a step"
+ _afterstep.uninstallfuncinlist(f)
+def isinstalledafterstep(f):
+ "Checks if the function is called after a step"
+ return _afterstep.isinstalledfuncinlist(f)
+
+# ----------------------------------------------------------------------------
+def callfromafterrestart(f):
+ raise Exception('restart call back not implemented yet')
+ installafterrestart(f)
+ return f
+def installafterrestart(f):
+ "Adds a function to the list of functions called immediately after a restart"
+ raise Exception('restart call back not implemented yet')
+ _afterrestart.installfuncinlist(f)
+def uninstallafterrestart(f):
+ "Removes the function from the list of functions called immediately after a restart"
+ raise Exception('restart call back not implemented yet')
+ _afterrestart.uninstallfuncinlist(f)
+def isinstalledafterrestart(f):
+ "Checks if the function is called immediately after a restart"
+ raise Exception('restart call back not implemented yet')
+ return _afterrestart.isinstalledfuncinlist(f)
+
+# ----------------------------------------------------------------------------
+def callfromparticleinjection(f):
+ installparticleinjection(f)
+ return f
+def installparticleinjection(f):
+ """
+ Adds a user defined function that is to be called when particle
+ injection happens, after the position advance and before deposition is
+ called, allowing a user defined particle distribution to be injected
+ each time step"""
+ _particleinjection.installfuncinlist(f)
+def uninstallparticleinjection(f):
+ "Removes the function installed by installparticleinjection"
+ _particleinjection.uninstallfuncinlist(f)
+def isinstalledparticleinjection(f):
+ "Checks if the function is called when particles injection happens"
+ return _particleinjection.isinstalledfuncinlist(f)
+
+# ----------------------------------------------------------------------------
+def callfromappliedfields(f):
+ raise Exception('applied fields call back not implemented yet')
+ installappliedfields(f)
+ return f
+def installappliedfields(f):
+ """
+ Adds a user defined function which can specify E and B fields which are applied
+ to the particles during the particle advance.
+ """
+ raise Exception('applied fields call back not implemented yet')
+ _appliedfields.installfuncinlist(f)
+def uninstallappliedfields(f):
+ "Removes the function installed by installappliedfields"
+ raise Exception('applied fields call back not implemented yet')
+ _appliedfields.uninstallfuncinlist(f)
+def isinstalledappliedfields(f):
+ "Checks if the function is called when which applies fields"
+ raise Exception('applied fields call back not implemented yet')
+ return _appliedfields.isinstalledfuncinlist(f)
+
diff --git a/Python/pywarpx/picmi.py b/Python/pywarpx/picmi.py
new file mode 100644
index 000000000..771edf85f
--- /dev/null
+++ b/Python/pywarpx/picmi.py
@@ -0,0 +1,388 @@
+"""Classes following the PICMI standard
+"""
+import PICMI_Base
+import numpy as np
+import pywarpx
+
+codename = 'WarpX'
+
+# --- Values from WarpXConst.H
+c = 299792458.
+ep0 = 8.854187817e-12
+mu0 = 1.2566370614359173e-06
+q_e = 1.602176462e-19
+m_e = 9.10938291e-31
+m_p = 1.6726231e-27
+
+
+class Species(PICMI_Base.PICMI_Species):
+ def init(self, **kw):
+
+ if self.particle_type == 'electron':
+ if self.charge is None: self.charge = '-q_e'
+ if self.mass is None: self.mass = 'm_e'
+ elif self.particle_type == 'positron':
+ if self.charge is None: self.charge = 'q_e'
+ if self.mass is None: self.mass = 'm_e'
+ elif self.particle_type == 'proton':
+ if self.charge is None: self.charge = 'q_e'
+ if self.mass is None: self.mass = 'm_p'
+ elif self.particle_type == 'anti-proton':
+ if self.charge is None: self.charge = '-q_e'
+ if self.mass is None: self.mass = 'm_p'
+
+ def initialize_inputs(self, layout):
+ self.species_number = pywarpx.particles.nspecies
+ pywarpx.particles.nspecies += 1
+
+ if self.name is None:
+ self.name = 'species{}'.format(self.species_number)
+
+ if pywarpx.particles.species_names is None:
+ pywarpx.particles.species_names = self.name
+ else:
+ pywarpx.particles.species_names += ' ' + self.name
+
+ self.species = pywarpx.Bucket.Bucket(self.name, mass=self.mass, charge=self.charge, injection_style = 'python')
+ pywarpx.Particles.particles_list.append(self.species)
+
+ if self.initial_distribution is not None:
+ self.initial_distribution.initialize_inputs(self.species_number, layout, self.species)
+
+
+PICMI_Base.PICMI_MultiSpecies.Species_class = Species
+class MultiSpecies(PICMI_Base.PICMI_MultiSpecies):
+ pass
+
+
+class GaussianBunchDistribution(PICMI_Base.PICMI_GaussianBunchDistribution):
+ def init(self, **kw):
+ if self.seed is not None:
+ print('Warning: WarpX does not support specifying the random number seed')
+
+ def initialize_inputs(self, species_number, layout, species):
+ species.injection_style = "gaussian_beam"
+ species.x_m = self.centroid_position[0]
+ species.y_m = self.centroid_position[1]
+ species.z_m = self.centroid_position[2]
+ species.x_rms = self.rms_bunch_size[0]
+ species.y_rms = self.rms_bunch_size[1]
+ species.z_rms = self.rms_bunch_size[2]
+
+ # --- Only PseudoRandomLayout is supported
+ species.npart = layout.n_macroparticles
+
+ # --- Calculate the total charge. Note that charge might be a string instead of a number.
+ charge = species.charge
+ if charge == 'q_e' or charge == '+q_e':
+ charge = q_e
+ elif charge == '-q_e':
+ charge = -q_e
+ species.q_tot = self.number_real_particles*charge
+
+ # --- These need to be defined even though they are not used
+ species.profile = "constant"
+ species.density = 1
+
+ # --- The PICMI standard doesn't yet have a way of specifying these values.
+ # --- They should default to the size of the domain. They are not typically
+ # --- necessary though since any particles outside the domain are rejected.
+ #species.xmin
+ #species.xmax
+ #species.ymin
+ #species.ymax
+ #species.zmin
+ #species.zmax
+
+ if np.any(np.not_equal(self.velocity_divergence, 0.)):
+ species.momentum_distribution_type = "radial_expansion"
+ species.u_over_r = self.velocity_divergence[0]
+ #species.u_over_y = self.velocity_divergence[1]
+ #species.u_over_z = self.velocity_divergence[2]
+ elif np.any(np.not_equal(self.rms_velocity, 0.)):
+ species.momentum_distribution_type = "gaussian"
+ species.ux_m = self.centroid_velocity[0]
+ species.uy_m = self.centroid_velocity[1]
+ species.uz_m = self.centroid_velocity[2]
+ species.ux_th = self.rms_velocity[0]
+ species.uy_th = self.rms_velocity[1]
+ species.uz_th = self.rms_velocity[2]
+ else:
+ species.momentum_distribution_type = "constant"
+ species.ux = self.centroid_velocity[0]
+ species.uy = self.centroid_velocity[1]
+ species.uz = self.centroid_velocity[2]
+
+
+class UniformDistribution(PICMI_Base.PICMI_UniformDistribution):
+
+ def initialize_inputs(self, species_number, layout, species):
+
+ if isinstance(layout, GriddedLayout):
+ species.injection_style = "nuniformpercell"
+ species.num_particles_per_cell_each_dim = layout.n_macroparticle_per_cell
+ elif isinstance(layout, PseudoRandomLayout):
+ assert (layout.n_macroparticles_per_cell is not None), Exception('WarpX only supports n_macroparticles_per_cell for the GriddedLayout with UniformDistribution')
+ species.injection_style = "nrandompercell"
+ species.num_particles_per_cell = layout.n_macroparticles_per_cell
+ else:
+ raise Exception('WarpX does not support the specified layout for UniformDistribution')
+
+ species.xmin = self.lower_bound[0]
+ species.xmax = self.upper_bound[0]
+ species.ymin = self.lower_bound[1]
+ species.ymax = self.upper_bound[1]
+ species.zmin = self.lower_bound[2]
+ species.zmax = self.upper_bound[2]
+
+ # --- Only constant density is supported at this time
+ species.profile = "constant"
+ species.density = self.density
+
+ if np.any(np.not_equal(self.rms_velocity, 0.)):
+ species.momentum_distribution_type = "gaussian"
+ species.ux_m = self.directed_velocity[0]
+ species.uy_m = self.directed_velocity[1]
+ species.uz_m = self.directed_velocity[2]
+ species.ux_th = self.rms_velocity[0]
+ species.uy_th = self.rms_velocity[1]
+ species.uz_th = self.rms_velocity[2]
+ else:
+ species.momentum_distribution_type = "constant"
+ species.ux = self.directed_velocity[0]
+ species.uy = self.directed_velocity[1]
+ species.uz = self.directed_velocity[2]
+
+ if self.fill_in:
+ pywarpx.warpx.do_plasma_injection = 1
+ if not hasattr(pywarpx.warpx, 'injected_plasma_species'):
+ pywarpx.warpx.injected_plasma_species = []
+
+ pywarpx.warpx.injected_plasma_species.append(species_number)
+ pywarpx.warpx.num_injected_species = len(pywarpx.warpx.injected_plasma_species)
+
+
+class AnalyticDistribution(PICMI_Base.PICMI_AnalyticDistribution):
+
+ def initialize_inputs(self, species_number, layout, species):
+ raise Exception('WarpX does not support AnalyticDistribution')
+
+
+class ParticleList(PICMI_Base.PICMI_ParticleList):
+ def init(self, **kw):
+
+ if len(x) > 1:
+ raise Exception('Only a single particle can be loaded')
+
+ def initialize_inputs(self, species_number, layout, species):
+
+ species.injection_style = "singleparticle"
+ species.single_particle_pos = [self.x[0], self.y[0], self.z[0]]
+ species.single_particle_vel = [self.ux[0]/c, self.uy[0]/c, self.uz[0]/c]
+ species.single_particle_weight = self.weight
+
+ # --- These need to be defined even though they are not used
+ species.profile = "constant"
+ species.density = 1
+ species.momentum_distribution_type = 'constant'
+
+
+class ParticleDistributionPlanarInjector(PICMI_Base.PICMI_ParticleDistributionPlanarInjector):
+ pass
+
+
+class GriddedLayout(PICMI_Base.PICMI_GriddedLayout):
+ pass
+
+
+class PseudoRandomLayout(PICMI_Base.PICMI_PseudoRandomLayout):
+ pass
+
+
+class BinomialSmoother(PICMI_Base.PICMI_BinomialSmoother):
+ pass
+
+
+class CylindricalGrid(PICMI_Base.PICMI_CylindricalGrid):
+ def init(self, **kw):
+ raise Exception('WarpX does not support CylindricalGrid')
+
+
+class Cartesian2DGrid(PICMI_Base.PICMI_Cartesian2DGrid):
+ def init(self, **kw):
+ self.max_grid_size = kw.get('max_grid_size', 32)
+ self.max_level = kw.get('max_level', 0)
+ self.coord_sys = kw.get('coord_sys', 0)
+
+ def initialize_inputs(self):
+ pywarpx.amr.n_cell = self.number_of_cells
+
+ # Maximum allowable size of each subdomain in the problem domain;
+ # this is used to decompose the domain for parallel calculations.
+ pywarpx.amr.max_grid_size = self.max_grid_size
+
+ # Maximum level in hierarchy (for now must be 0, i.e., one level in total)
+ pywarpx.amr.max_level = self.max_level
+
+ # Geometry
+ pywarpx.geometry.coord_sys = self.coord_sys
+ pywarpx.geometry.is_periodic = '%d %d %d'%(self.bc_xmin=='periodic', self.bc_ymin=='periodic') # Is periodic?
+ pywarpx.geometry.prob_lo = self.lower_bound # physical domain
+ pywarpx.geometry.prob_hi = self.upper_bound
+
+ if self.moving_window_velocity is not None and np.any(np.not_equal(self.moving_window_velocity, 0.)):
+ pywarpx.warpx.do_moving_window = 1
+ if self.moving_window_velocity[0] != 0.:
+ pywarpx.warpx.moving_window_dir = 'x'
+ pywarpx.warpx.moving_window_v = self.moving_window_velocity[0]/c # in units of the speed of light
+ if self.moving_window_velocity[1] != 0.:
+ pywarpx.warpx.moving_window_dir = 'y'
+ pywarpx.warpx.moving_window_v = self.moving_window_velocity[1]/c # in units of the speed of light
+
+
+class Cartesian3DGrid(PICMI_Base.PICMI_Cartesian3DGrid):
+ def init(self, **kw):
+ self.max_grid_size = kw.get('max_grid_size', 32)
+ self.max_level = kw.get('max_level', 0)
+ self.coord_sys = kw.get('coord_sys', 0)
+
+ def initialize_inputs(self):
+ pywarpx.amr.n_cell = self.number_of_cells
+
+ # Maximum allowable size of each subdomain in the problem domain;
+ # this is used to decompose the domain for parallel calculations.
+ pywarpx.amr.max_grid_size = self.max_grid_size
+
+ # Maximum level in hierarchy (for now must be 0, i.e., one level in total)
+ pywarpx.amr.max_level = self.max_level
+
+ # Geometry
+ pywarpx.geometry.coord_sys = self.coord_sys
+ pywarpx.geometry.is_periodic = '%d %d %d'%(self.bc_xmin=='periodic', self.bc_ymin=='periodic', self.bc_zmin=='periodic') # Is periodic?
+ pywarpx.geometry.prob_lo = self.lower_bound # physical domain
+ pywarpx.geometry.prob_hi = self.upper_bound
+
+ if self.moving_window_velocity is not None and np.any(np.not_equal(self.moving_window_velocity, 0.)):
+ pywarpx.warpx.do_moving_window = 1
+ if self.moving_window_velocity[0] != 0.:
+ pywarpx.warpx.moving_window_dir = 'x'
+ pywarpx.warpx.moving_window_v = self.moving_window_velocity[0]/c # in units of the speed of light
+ if self.moving_window_velocity[1] != 0.:
+ pywarpx.warpx.moving_window_dir = 'y'
+ pywarpx.warpx.moving_window_v = self.moving_window_velocity[1]/c # in units of the speed of light
+ if self.moving_window_velocity[2] != 0.:
+ pywarpx.warpx.moving_window_dir = 'z'
+ pywarpx.warpx.moving_window_v = self.moving_window_velocity[2]/c # in units of the speed of light
+
+
+class ElectromagneticSolver(PICMI_Base.PICMI_ElectromagneticSolver):
+ def init(self, **kw):
+ assert self.method is None or self.method in ['Yee'], Exception("Only 'Yee' FDTD is supported")
+
+ def initialize_inputs(self):
+
+ self.grid.initialize_inputs()
+
+ if self.cfl is not None:
+ pywarpx.warpx.cfl = self.cfl
+
+ if self.stencil_order is not None:
+ pywarpx.interpolation.nox = self.stencil_order[0]
+ pywarpx.interpolation.noy = self.stencil_order[1]
+ pywarpx.interpolation.noz = self.stencil_order[2]
+
+
+class Electrostatic_solver(PICMI_Base.PICMI_Electrostatic_solver):
+ def initialize_inputs(self):
+ pass
+
+
+class GaussianLaser(PICMI_Base.PICMI_GaussianLaser):
+
+ def initialize_inputs(self):
+ pywarpx.warpx.use_laser = 1
+ pywarpx.laser.profile = "Gaussian"
+ pywarpx.laser.wavelength = self.wavelength # The wavelength of the laser (in meters)
+ pywarpx.laser.e_max = self.E0 # Maximum amplitude of the laser field (in V/m)
+ pywarpx.laser.polarization = [np.cos(self.polarization_angle), np.sin(self.polarization_angle), 0.] # The main polarization vector
+ pywarpx.laser.profile_waist = self.waist # The waist of the laser (in meters)
+ pywarpx.laser.profile_duration = self.duration # The duration of the laser (in seconds)
+
+
+class LaserAntenna(PICMI_Base.PICMI_LaserAntenna):
+
+ def initialize_inputs(self, laser):
+ pywarpx.laser.position = self.position # This point is on the laser plane
+ pywarpx.laser.direction = self.normal_vector # The plane normal direction
+ pywarpx.laser.profile_focal_distance = laser.focal_position[2] - self.position[2] # Focal distance from the antenna (in meters)
+ pywarpx.laser.profile_t_peak = (self.position[2] - laser.centroid_position[2])/c # The time at which the laser reaches its peak (in seconds)
+
+
+class Simulation(PICMI_Base.PICMI_Simulation):
+ def init(self, **kw):
+
+ self.plot_int = kw.get('plot_int', None)
+ self.current_deposition_algo = kw.get('current_deposition_algo', None)
+ self.charge_deposition_algo = kw.get('charge_deposition_algo', None)
+ self.field_gathering_algo = kw.get('field_gathering_algo', None)
+ self.particle_pusher_algo = kw.get('particle_pusher_algo', None)
+
+ self.inputs_initialized = False
+ self.warpx_initialized = False
+
+ def initialize_inputs(self):
+ if self.inputs_initialized:
+ return
+
+ self.inputs_initialized = True
+
+ pywarpx.warpx.verbose = self.verbose
+ if self.time_step_size is not None:
+ pywarpx.warpx.const_dt = self.timestep
+
+ pywarpx.amr.plot_int = self.plot_int
+ pywarpx.algo.current_deposition = self.current_deposition_algo
+ pywarpx.algo.charge_deposition = self.charge_deposition_algo
+ pywarpx.algo.field_gathering = self.field_gathering_algo
+ pywarpx.algo.particle_pusher = self.particle_pusher_algo
+
+ self.solver.initialize_inputs()
+
+ for i in range(len(self.species)):
+ assert self.calculate_self_fields[i], Exception('WarpX does not support species without self fields')
+ self.species[i].initialize_inputs(self.layouts[i])
+
+ for i in range(len(self.lasers)):
+ self.lasers[i].initialize_inputs()
+ self.laser_injection_methods[i].initialize_inputs(self.lasers[i])
+
+ def initialize_warpx(self, inputs_name=None):
+ if self.warpx_initialized:
+ return
+
+ self.warpx_initialized = True
+ pywarpx.warpx.init()
+
+ def write_input_file(self, inputs_name='inputs'):
+ self.initialize_inputs()
+ kw = {}
+ if self.max_steps is not None:
+ kw['max_step'] = self.max_steps
+ if self.max_time is not None:
+ kw['stop_time'] = self.max_time
+ pywarpx.warpx.write_inputs(inputs_name, **kw)
+
+ def step(self, nsteps=None):
+ self.initialize_inputs()
+ self.initialize_warpx()
+ if nsteps is None:
+ if self.max_steps is not None:
+ nsteps = self.max_steps
+ else:
+ nsteps = -1
+ pywarpx.warpx.evolve(nsteps)
+
+ def finalize(self):
+ if self.warpx_initialized:
+ self.warpx_initialized = False
+ pywarpx.warpx.finalize()
diff --git a/Python/pywarpx/timestepper.py b/Python/pywarpx/timestepper.py
index f72b1669f..3068cbc82 100644
--- a/Python/pywarpx/timestepper.py
+++ b/Python/pywarpx/timestepper.py
@@ -1,8 +1,5 @@
-import warp
-from warp import top
-from warp import w3d
-
from ._libwarpx import libwarpx
+from . import callbacks
class TimeStepper(object):
@@ -12,6 +9,8 @@ class TimeStepper(object):
def onestep(self):
+ callbacks._beforestep()
+
self.cur_time = libwarpx.warpx_gett_new(0)
self.istep = libwarpx.warpx_getistep(0)
@@ -33,7 +32,11 @@ class TimeStepper(object):
# --- Evolve particles to p^{n+1/2} and x^{n+1}
# --- Depose current, j^{n+1/2}
+ callbacks._particleinjection()
+ callbacks._particlescraper()
+ callbacks._beforedeposition()
libwarpx.warpx_PushParticlesandDepose(self.cur_time)
+ callbacks._afterdeposition()
libwarpx.mypc_Redistribute() # Redistribute particles
@@ -43,7 +46,9 @@ class TimeStepper(object):
libwarpx.warpx_SyncCurrent()
libwarpx.warpx_FillBoundaryB()
+ callbacks._beforeEsolve()
libwarpx.warpx_EvolveE(dt,0) # We now have E^{n+1}
+ callbacks._afterEsolve()
self.istep += 1
@@ -67,318 +72,4 @@ class TimeStepper(object):
if libwarpx.warpx_checkInt() > 0 and (self.istep+1)%libwarpx.warpx_plotInt() == 0 or max_time_reached:
libwarpx.warpx_WriteCheckPointFile()
-
-
-
-
-
-
-
-
-
-
-
-# --- This is not used
-class TimeStepperFromPICSAR(object):
-
- def __init__(self, package=None, solver=None, l_debug=False):
- self.package = package
- self.solver = solver
- self.l_debug = l_debug
-
- def setpackagestepnumber(self, it):
- if self.package is not None:
- self.package.setstepnumber(it)
-
- def step(self, n=1, freq_print=10, lallspecl=0):
- """
- This function performs a range of Particle-In-Cell iterations
-
- Inputs:
- - n: number of iterations
- - freq_print: print frequency
- """
-
- if (self.l_debug): print("Call step")
-
- for i in range(n):
- if(me == 0):
- if top.it%freq_print == 0:
- print('it = %g time = %g'%(top.it, top.time))
-
- l_first = (lallspecl or (i == 0))
- l_last = (lallspecl or (i == n-1))
-
- self.onestep(l_first, l_last)
-
- if (self.l_debug): print("End step")
-
- def onestep(self, l_first, l_last):
- """
- Perform a single particle-in-cell step
- """
-
- if (self.l_debug): print("Call onestep")
-
- # --- Iteration number
- self.setpackagestepnumber(top.it)
-
- # --- call beforestep functions
- if (self.l_debug): print("Call beforestep functions")
- warp.callbeforestepfuncs.callfuncsinlist()
-
- # --- gather fields from grid to particles
- if (self.l_debug): print("Call Field gathering and particle push")
-
- # --- push
- if l_first:
- if self.package is None:
- # --- Standard Warp advance
- for specie in warp.listofallspecies:
- for pg in specie.flatten(specie.pgroups):
- for js in range(pg.ns):
- self.push_velocity_second_part(js, pg)
- self.push_positions(js, pg)
- warp.particleboundaries3d(pg, -1, False)
-
- else:
- # --- Particle pusher
- if (self.l_debug): print("Call package particle push")
- pxr.pxrpush_particles_part2()
-
- # --- Particle boundary consitions
- if (self.l_debug): print("Call package particle boundary conditions")
- pxr.particle_bcs()
-
- if (self.l_debug): print("Call aliasparticlearrays()")
- self.aliasparticlearrays()
-
- else:
- if self.package is None:
- # --- Standard Warp advance
-
- for specie in warp.listofallspecies:
- for pg in specie.flatten(specie.pgroups):
- for js in range(pg.ns):
- self.push_velocity_full(js, pg)
- self.push_positions(js, pg)
-
- warp.particleboundaries3d(pg, -1, False)
-
- else:
- # --- Particle pusher
- if (self.l_debug): print("Call package particle pusher")
- pxr.field_gathering_plus_particle_pusher()
-
- # --- Particle boundary conditions
- if (self.l_debug): print("Call package particle boundary conditions")
- pxr.particle_bcs()
-
- self.aliasparticlearrays()
-
- # --- Particle sorting
- if (self.l_debug): print("Call Particle Sorting")
- if self.package is not None:
- # --- This should be a function installed before load rho
- pxr.particle_sorting_sub()
-
- # --- call beforeloadrho functions
- if (self.l_debug): print("Call beforeloadrho functions")
- warp.beforeloadrho.callfuncsinlist()
-
- pgroups = []
- for specie in warp.listofallspecies:
- pgroups += specie.flatten(specie.pgroups)
-
- self.pgroups = pgroups
-
- # --- Call user-defined injection routines
- if (self.l_debug): print("Call user-defined injection routines")
- warp.userinjection.callfuncsinlist()
-
- xgriddiff = w3d.xmmin - pxr.xmin
- ygriddiff = w3d.ymmin - pxr.ymin
- zgriddiff = w3d.zmmin - pxr.zmin
-
- if (xgriddiff != 0 or ygriddiff != 0 or zgriddiff != 0):
- pxr.pxr_move_sim_boundaries(xgriddiff, ygriddiff, zgriddiff)
-
-
- if (self.l_debug): print("Call loadrho")
- self.solver.loadrho(pgroups = pgroups)
- if (self.l_debug): print("Call loadj")
- self.solver.loadj(pgroups = pgroups)
-
- if (self.l_debug): print("Call dosolve")
- self.solver.dosolve()
-
- if self.package is None:
- for specie in warp.listofallspecies:
- for pg in specie.flatten(specie.pgroups):
- for js in range(pg.ns):
- self.fetcheb(js, pg)
- if l_last:
- self.push_velocity_first_part(js, pg)
-
- else:
- if l_last:
- if (self.l_debug): print("Call package push particles 1")
- pxr.pxrpush_particles_part1()
-
- # --- update time, time counter
- top.time += top.dt
- if top.it%top.nhist == 0:
-# zmmnt()
- minidiag(top.it, top.time, top.lspecial)
- top.it += 1
-
- # --- Load balance function should be installed after step
-
- # --- call afterstep functions
- if (self.l_debug): print("Call callafterstepfuncs.callfuncsinlist()")
- warp.callafterstepfuncs.callfuncsinlist()
-
-
- # --- The following methods use the standard Warp routines
- def fetcheb(self, js, pg=None):
- if self.l_verbose:print(me, 'enter fetcheb')
- if pg is None:
- pg = top.pgroup
- np = pg.nps[js]
- if np == 0: return
- il = pg.ins[js] - 1
- iu = il + pg.nps[js]
- w3d.pgroupfsapi = pg
- w3d.ipminfsapi = pg.ins[js]
- w3d.npfsapi = pg.nps[js]
- pg.ex[il:iu] = 0.
- pg.ey[il:iu] = 0.
- pg.ez[il:iu] = 0.
- pg.bx[il:iu] = 0.
- pg.by[il:iu] = 0.
- pg.bz[il:iu] = 0.
- self.fetche()
- self.fetchb()
- w3d.pgroupfsapi = None
-
- def push_velocity_full(self, js, pg=None):
- if self.l_verbose:print(me, 'enter push_velocity_full')
- if pg is None:
- pg = top.pgroup
- np = pg.nps[js]
- if np == 0: return
- il = pg.ins[js] - 1
- iu = il + pg.nps[js]
- if pg.lebcancel_pusher:
- warp.ebcancelpush3d(np, pg.uxp[il:iu], pg.uyp[il:iu], pg.uzp[il:iu], pg.gaminv[il:iu],
- pg.ex[il:iu], pg.ey[il:iu], pg.ez[il:iu],
- pg.bx[il:iu], pg.by[il:iu], pg.bz[il:iu],
- pg.sq[js], pg.sm[js], top.dt, 0)
- else:
- # --- push velocity from electric field (half step)
- warp.epush3d(np, pg.uxp[il:iu], pg.uyp[il:iu], pg.uzp[il:iu],
- pg.ex[il:iu], pg.ey[il:iu], pg.ez[il:iu],
- pg.sq[js], pg.sm[js], 0.5*top.dt)
- # --- update gamma
- self.set_gamma(js, pg)
- # --- push velocity from magnetic field
- warp.bpush3d(np, pg.uxp[il:iu], pg.uyp[il:iu], pg.uzp[il:iu], pg.gaminv[il:iu],
- pg.bx[il:iu], pg.by[il:iu], pg.bz[il:iu],
- pg.sq[js], pg.sm[js], top.dt, top.ibpush)
- # --- push velocity from electric field (half step)
- warp.epush3d(np, pg.uxp[il:iu], pg.uyp[il:iu], pg.uzp[il:iu],
- pg.ex[il:iu], pg.ey[il:iu], pg.ez[il:iu],
- pg.sq[js], pg.sm[js], 0.5*top.dt)
- # --- update gamma
- self.set_gamma(js, pg)
-
- if self.l_verbose:print(me, 'exit push_velocity_first_part')
-
- def push_velocity_first_part(self, js, pg=None):
- if self.l_verbose:print(me, 'enter push_velocity_first_part')
- if pg is None:
- pg = top.pgroup
- np = pg.nps[js]
- if np == 0: return
- il = pg.ins[js] - 1
- iu = il + pg.nps[js]
- if pg.lebcancel_pusher:
- warp.ebcancelpush3d(np, pg.uxp[il:iu], pg.uyp[il:iu], pg.uzp[il:iu], pg.gaminv[il:iu],
- pg.ex[il:iu], pg.ey[il:iu], pg.ez[il:iu],
- pg.bx[il:iu], pg.by[il:iu], pg.bz[il:iu],
- pg.sq[js], pg.sm[js], top.dt, 1)
- else:
- # --- push velocity from electric field (half step)
- warp.epush3d(np, pg.uxp[il:iu], pg.uyp[il:iu], pg.uzp[il:iu],
- pg.ex[il:iu], pg.ey[il:iu], pg.ez[il:iu],
- pg.sq[js], pg.sm[js], 0.5*top.dt)
- # --- update gamma
- self.set_gamma(js,pg)
- # --- push velocity from magnetic field
- warp.bpush3d(np, pg.uxp[il:iu], pg.uyp[il:iu], pg.uzp[il:iu], pg.gaminv[il:iu],
- pg.bx[il:iu], pg.by[il:iu], pg.bz[il:iu],
- pg.sq[js], pg.sm[js], 0.5*top.dt, top.ibpush)
-
- if self.l_verbose:print(me, 'exit push_velocity_first_part')
-
- def push_velocity_second_part(self, js, pg=None):
- if self.l_verbose:print(me, 'enter push_velocity_second_part')
- if pg is None:
- pg = top.pgroup
- np = pg.nps[js]
- if np == 0: return
- il = pg.ins[js] - 1
- iu = il + pg.nps[js]
- if pg.lebcancel_pusher:
- warp.ebcancelpush3d(np, pg.uxp[il:iu], pg.uyp[il:iu], pg.uzp[il:iu], pg.gaminv[il:iu],
- pg.ex[il:iu], pg.ey[il:iu], pg.ez[il:iu],
- pg.bx[il:iu], pg.by[il:iu], pg.bz[il:iu],
- pg.sq[js], pg.sm[js], top.dt, 2)
- else:
- # --- push velocity from magnetic field
- warp.bpush3d(np, pg.uxp[il:iu], pg.uyp[il:iu], pg.uzp[il:iu], pg.gaminv[il:iu],
- pg.bx[il:iu], pg.by[il:iu], pg.bz[il:iu],
- pg.sq[js], pg.sm[js], 0.5*top.dt, top.ibpush)
- # --- push velocity from electric field (half step)
- warp.epush3d(np, pg.uxp[il:iu], pg.uyp[il:iu], pg.uzp[il:iu],
- pg.ex[il:iu], pg.ey[il:iu], pg.ez[il:iu],
- pg.sq[js], pg.sm[js], 0.5*top.dt)
- # --- update gamma
- self.set_gamma(js, pg)
-
- if self.l_verbose:print(me, 'exit push_velocity_second_part')
-
- def set_gamma(self, js, pg=None):
- if self.l_verbose:print(me, 'enter set_gamma')
- if pg is None:
- pg = top.pgroup
- np = pg.nps[js]
- if np == 0: return
- il = pg.ins[js] - 1
- iu = il + pg.nps[js]
- # --- update gamma
- warp.gammaadv(np, pg.gaminv[il:iu], pg.uxp[il:iu], pg.uyp[il:iu], pg.uzp[il:iu], top.gamadv, top.lrelativ)
-
- if self.l_verbose:print(me, 'exit set_gamma')
-
- def push_positions(self, js, pg=None):
- if self.l_verbose:print(me, 'enter push_positions')
- if pg is None:
- pg = top.pgroup
- np = pg.nps[js]
- if np == 0: return
- il = pg.ins[js] - 1
- iu = il + pg.nps[js]
- warp.xpush3d(np, pg.xp[il:iu], pg.yp[il:iu], pg.zp[il:iu],
- pg.uxp[il:iu], pg.uyp[il:iu], pg.uzp[il:iu],
- pg.gaminv[il:iu], top.dt)
-
- if self.l_verbose:print(me, 'exit push_positions')
-
-
-
-
-
-
-
-
+ callbacks._afterstep()
diff --git a/Python/setup.py b/Python/setup.py
index 20e7c5d9f..71f82dbc7 100644
--- a/Python/setup.py
+++ b/Python/setup.py
@@ -6,9 +6,16 @@ setup.py file for WarpX
from distutils.core import setup
+try:
+ from distutils.command.build_py import build_py_2to3 as build_py
+except ImportError:
+ from distutils.command.build_py import build_py
+
+
setup (name = 'pywarpx',
packages = ['pywarpx'],
package_dir = {'pywarpx':'pywarpx'},
description = """Wrapper of WarpX""",
package_data = {'pywarpx' : ['libwarpx.so']},
+ cmdclass={'build_py': build_py}
)
generated by cgit v1.2.3 (git 2.46.0) at 2025-08-12 17:26:57 +0000