Updates to be consistent with the PICMI standard

1 files changed, 208 insertions, 165 deletions

diff --git a/Python/pywarpx/PICMI.py b/Python/pywarpx/PICMI.py
index b1711926f..271d363b9 100644
--- a/Python/pywarpx/PICMI.py
+++ b/Python/pywarpx/PICMI.py
@@ -1,117 +1,52 @@
 """Classes following the PICMI standard
 """
-from PICMI_Base import *
+import PICMI_Base
 import numpy as np
-from pywarpx import *
+import pywarpx
 
 codename = 'WarpX'
 
-class Grid(PICMI_Grid):
-    def init(self, **kw):
-
-        amr.n_cell = [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 = [self.xmin, self.ymin, self.zmin]  # physical domain
-        geometry.prob_hi = [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
+# --- Values from WarpXConst.H
+c   = 299792458.
+#ep0 = 8.854187817e-12
+#mu0 = 1.2566370614359173e-06
+#q_e = 1.6021764620000001e-19
+#m_e = 9.10938291e-31
+#m_p = 1.6726231000000001e-27
 
 
-class Gaussian_laser(PICMI_Gaussian_laser):
+class Species(PICMI_Base.PICMI_Species):
     def init(self, **kw):
 
-        warpx.use_laser = 1
-        laser.profile = "Gaussian"
-        laser.wavelength = self.wavelength  # The wavelength of the laser (in meters)
-        laser.e_max = self.E0  # Maximum amplitude of the laser field (in V/m)
-        laser.polarization = [np.cos(self.pol_angle), np.sin(self.pol_angle), 0.]  # The main polarization vector
-        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.focal_position - self.z0)/clight  # The time at which the laser reaches its peak (in seconds)
-
-
-class Laser_antenna(PICMI_Laser_antenna):
-    def init(self, **kw):
-
-        laser.position = [self.antenna_x0, self.antenna_y0, self.antenna_z0]  # This point is on the laser plane
-        laser.direction = [self.antenna_xvec, self.antenna_yvec, self.antenna_zvec]  # The plane normal direction
-        laser.profile_focal_distance = self.laser.focal_position - self.antenna_z0  # Focal distance from the antenna (in meters)
-
-
-class Species(PICMI_Species):
-    def init(self, **kw):
-
-        self.species_number = particles.nspecies
-        particles.nspecies = particles.nspecies + 1
-        if particles.species_names is None:
-            particles.species_names = self.name
+        if self.type == 'electron':
+            if self.charge is None: self.charge = '-q_e'
+            if self.mass is None: self.mass = 'm_e'
+        elif self.type == 'positron':
+            if self.charge is None: self.charge = 'q_e'
+            if self.mass is None: self.mass = 'm_e'
+        elif self.type == 'proton':
+            if self.charge is None: self.charge = 'q_e'
+            if self.mass is None: self.mass = 'm_p'
+        elif self.type == 'anti-proton':
+            if self.charge is None: self.charge = '-q_e'
+            if self.mass is None: self.mass = 'm_p'
+
+        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:
-            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)
+            pywarpx.particles.species_names += ' ' + self.name
 
-    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)
+        self.bucket = pywarpx.Bucket.Bucket(self.name, mass=self.mass, charge=self.charge, injection_style = 'python')
+        pywarpx.Particles.particles_list.append(self.bucket)
 
 
-class GaussianBeam(PICMI_GaussianBeam):
+class GaussianBeam(PICMI_Base.PICMI_GaussianBeam):
     def init(self, **kw):
 
         self.species.bucket.injection_style = "gaussian_beam"
@@ -122,41 +57,55 @@ class GaussianBeam(PICMI_GaussianBeam):
         self.species.bucket.y_rms = self.Yrms
         self.species.bucket.z_rms = self.Zrms
         self.species.bucket.npart = self.number_sim_particles
-        self.species.bucket.q_tot = self.number_sim_particles*self.species.charge
+        self.species.bucket.q_tot = self.number_real_particles*self.species.bucket.charge
 
-        # --- These are unused but need to be set (maybe)
-        self.species.bucket.profile = 'constant'
+        # --- These need to be defined even though they are not used
+        self.species.bucket.profile = "constant"
         self.species.bucket.density = 1
 
-        # --- Momentum distribution
-        if 'u_over_r' in kw:
-            # --- Radial expansion
+        # --- 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.
+        #self.species.bucket.xmin
+        #self.species.bucket.xmax
+        #self.species.bucket.ymin
+        #self.species.bucket.ymax
+        #self.species.bucket.zmin
+        #self.species.bucket.zmax
+
+        if self.UXdiv != 0. and UYdiv != 0. and UZdiv != 0.:
             self.species.bucket.momentum_distribution_type = "radial_expansion"
-            self.species.bucket.u_over_r = kw['u_over_r']
-
-        elif self.UXrms == 0. and self.UYrms == 0. and self.UZrms == 0.:
-            # --- Constant velocity
-            self.species.bucket.momentum_distribution_type = "constant"
-            self.species.bucket.ux = self.UXmean
-            self.species.bucket.uy = self.UYmean
-            self.species.bucket.uz = self.UZmean
-
-        else:
-            # --- Gaussian velocity distribution
+            self.species.bucket.u_over_r = self.UXdiv
+            #self.species.bucket.u_over_y = self.UYdiv
+            #self.species.bucket.u_over_z = self.UZdiv
+        elif self.UXrms != 0. or UYrms != 0. or UZrms != 0.:
             self.species.bucket.momentum_distribution_type = "gaussian"
             self.species.bucket.ux_m = self.UXmean
             self.species.bucket.uy_m = self.UYmean
             self.species.bucket.uz_m = self.UZmean
-            self.species.bucket.u_th = self.UZrms
-            # !!! UXrms and UYrms are unused. Only an isotropic distribution is supported
-            # !!! Maybe an error should be raised
+            self.species.bucket.ux_th = self.UXrms
+            self.species.bucket.uy_th = self.UYrms
+            self.species.bucket.uz_th = self.UZrms
+        else:
+            self.species.bucket.momentum_distribution_type = "constant"
+            self.species.bucket.ux = self.UXmean
+            self.species.bucket.uy = self.UYmean
+            self.species.bucket.uz = self.UZmean
 
 
-class Plasma(PICMI_Plasma):
+class Plasma(PICMI_Base.PICMI_Plasma):
     def init(self, **kw):
 
         for species in self.species:
-            species.bucket.injection_style = "NUniformPerCell"
+            if self.number_per_cell_each_dim is not None:
+                species.bucket.injection_style = "nuniformpercell"
+                species.bucket.num_particles_per_cell_each_dim = self.number_per_cell_each_dim
+            elif self.number_per_cell is not None:
+                species.bucket.injection_style = "nrandompercell"
+                species.bucket.num_particles_per_cell = self.number_per_cell
+            else:
+                raise Exception('Either nuniformpercell or nrandompercell must be specified')
+
             species.bucket.xmin = self.xmin
             species.bucket.xmax = self.xmax
             species.bucket.ymin = self.ymin
@@ -164,78 +113,150 @@ class Plasma(PICMI_Plasma):
             species.bucket.zmin = self.zmin
             species.bucket.zmax = self.zmax
 
-            species.bucket.profile = 'constant'
+            # --- Only constant density is supported at this time
+            species.bucket.profile = "constant"
             species.bucket.density = self.density
 
-            if self.number_per_cell is not None:
-                species.bucket.nrandompercell = self.number_per_cell
-            elif self.number_per_cell_each_dim is not None:
-                species.bucket.num_particles_per_cell_each_dim = self.number_per_cell_each_dim
-
-            # --- Momentum distribution
-            if 'u_over_r' in kw:
-                # --- Radial expansion
-                species.bucket.momentum_distribution_type = "radial_expansion"
-                species.bucket.u_over_r = kw['u_over_r']
-
-            elif self.vthx == 0. and self.vthy == 0. and self.vthz == 0.:
-                # --- Constant velocity
+            if self.vthx != 0. or self.vthy != 0. or self.vthz != 0.:
+                species.bucket.momentum_distribution_type = "gaussian"
+                species.bucket.ux_m = self.vxmean
+                species.bucket.uy_m = self.vymean
+                species.bucket.uz_m = self.vzmean
+                species.bucket.ux_th = self.vthx
+                species.bucket.uy_th = self.vthy
+                species.bucket.uz_th = self.vthz
+            else:
                 species.bucket.momentum_distribution_type = "constant"
                 species.bucket.ux = self.vxmean
                 species.bucket.uy = self.vymean
                 species.bucket.uz = self.vzmean
 
-            else:
-                # --- Gaussian velocity distribution
-                species.bucket.momentum_distribution_type = "gaussian"
-                species.bucket.ux_m = self.vxmean
-                species.bucket.uy_m = self.vymean
-                species.bucket.uz_m = self.vzmean
-                species.bucket.u_th = self.vthz
-                # !!! vthx and vthy are unused. Only an isotropic distribution is supported
-                # !!! Maybe an error should be raised
+            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.species_number)
+                pywarpx.warpx.num_injected_species = len(pywarpx.warpx.injected_plasma_species)
 
-class ParticleList(PICMI_ParticleList):
+
+class ParticleList(PICMI_Base.PICMI_ParticleList):
     def init(self, **kw):
 
-        assert len(self.x) == 1, "WarpX only supports initializing with a single particle"
+        if len(x) > 1:
+            raise Exception('Only a single particle can be loaded')
 
-        self.species.bucket.injection_style = "SingleParticle"
+        self.species.bucket.injection_style = "singleparticle"
         self.species.bucket.single_particle_pos = [self.x[0], self.y[0], self.z[0]]
-        self.species.bucket.single_particle_vel = [self.ux[0]/clight, self.uy[0]/clight, self.uz[0]/clight]
+        self.species.bucket.single_particle_vel = [self.ux[0]/c, self.uy[0]/c, self.uz[0]/c]
         self.species.bucket.single_particle_weight = self.weight
 
+        # --- These need to be defined even though they are not used
+        self.species.bucket.profile = "constant"
+        self.species.bucket.density = 1
+        self.species.bucket.momentum_distribution_type = 'constant'
+
+
+class Grid(PICMI_Base.PICMI_Grid):
+    def init(self, **kw):
+
+        pywarpx.amr.n_cell = [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.
+        pywarpx.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)
+        pywarpx.amr.max_level = kw.get('max_level', 0)
+
+        # Geometry
+        pywarpx.geometry.coord_sys = kw.get('coord_sys', 0)  # 0: Cartesian
+        pywarpx.geometry.is_periodic = '%d %d %d'%(self.bcxmin=='periodic', self.bcymin=='periodic', self.bczmin=='periodic')  # Is periodic?
+        pywarpx.geometry.prob_lo = [self.xmin, self.ymin, self.zmin]  # physical domain
+        pywarpx.geometry.prob_hi = [self.xmax, self.ymax, self.zmax]
+
+        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
+
+    def getmins(self, **kw):
+        return np.array([pywarpx.warpx.getProbLo(0), pywarpx.warpx.getProbLo(1), pywarpx.warpx.getProbLo(2)])
+
+    def getmaxs(self, **kw):
+        return np.array([pywarpx.warpx.getProbHi(0), pywarpx.warpx.getProbHi(1), pywarpx.warpx.getProbHi(2)])
+
+    def getxmin(self):
+        return pywarpx.warpx.getProbLo(0)
+
+    def getxmax(self):
+        return pywarpx.warpx.getProbHi(0)
+
+    def getymin(self):
+        return pywarpx.warpx.getProbLo(1)
+
+    def getymax(self):
+        return pywarpx.warpx.getProbHi(1)
+
+    def getzmin(self):
+        return pywarpx.warpx.getProbLo(2)
+
+    def getzmax(self):
+        return pywarpx.warpx.getProbHi(2)
+
+
+class EM_solver(PICMI_Base.PICMI_EM_solver):
+    def init(self, **kw):
+
+        if self.method is None:
+            self.method = 'Yee'
 
-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)
+        assert self.method in ['Yee'], Exception("Only 'Yee' FDTD is supported")
 
+        if 'current_deposition_algo' in kw:
+            pywarpx.algo.current_deposition = kw['current_deposition_algo']
+        if 'charge_deposition_algo' in kw:
+            pywarpx.algo.charge_deposition = kw['charge_deposition_algo']
+        if 'field_gathering_algo' in kw:
+            pywarpx.algo.field_gathering = kw['field_gathering_algo']
+        if 'particle_pusher_algo' in kw:
+            pywarpx.algo.particle_pusher = kw['particle_pusher_algo']
+
+
+class Simulation(PICMI_Base.PICMI_Simulation):
     def init(self, **kw):
 
-        warpx.verbose = self.verbose
-        warpx.cfl = self.timestep_over_cfl
+        pywarpx.warpx.verbose = self.verbose
+        pywarpx.warpx.cfl = self.timestep_over_cfl
         if self.timestep == 0.:
-            warpx.cfl = self.timestep_over_cfl
+            pywarpx.warpx.cfl = self.timestep_over_cfl
         else:
-            warpx.const_dt = self.timestep
-        amr.plot_int = self.plot_int
+            pywarpx.warpx.const_dt = self.timestep
+
+        if 'plot_int' in kw:
+            pywarpx.amr.plot_int = kw['plot_int']
 
         self.initialized = False
 
     def initialize(self, inputs_name=None):
         if not self.initialized:
             self.initialized = True
-            warpx.init()
+            pywarpx.warpx.init()
 
     def write_inputs(self, inputs_name='inputs'):
         kw = {}
-        if hasattr(self, 'max_step'):
+        if self.max_step is not None:
             kw['max_step'] = self.max_step
-        warpx.write_inputs(inputs_name, **kw)
+        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()
@@ -244,8 +265,30 @@ class Simulation(PICMI_Simulation):
                 nsteps = self.max_step
             else:
                 nsteps = -1
-        warpx.evolve(nsteps)
+        pywarpx.warpx.evolve(nsteps)
 
     def finalize(self):
-        warpx.finalize()
+        if self.initialized:
+            self.initialized = False
+            pywarpx.warpx.finalize()
+
+
+class Gaussian_laser(PICMI_Base.PICMI_Gaussian_laser):
+    def init(self, **kw):
+
+        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.pol_angle), np.sin(self.pol_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)
+        pywarpx.laser.profile_t_peak = (self.focal_position - self.z0)/c  # The time at which the laser reaches its peak (in seconds)
+
+
+class Laser_antenna(PICMI_Base.PICMI_Laser_antenna):
+    def init(self, **kw):
 
+        pywarpx.laser.position = [self.antenna_x0, self.antenna_y0, self.antenna_z0]  # This point is on the laser plane
+        pywarpx.laser.direction = [self.antenna_xvec, self.antenna_yvec, self.antenna_zvec]  # The plane normal direction
+        pywarpx.laser.profile_focal_distance = self.laser.focal_position - self.antenna_z0  # Focal distance from the antenna (in meters)