diff options
Diffstat (limited to 'Python/pywarpx/picmi.py')
| -rw-r--r-- | Python/pywarpx/picmi.py | 21 |
1 files changed, 18 insertions, 3 deletions
diff --git a/Python/pywarpx/picmi.py b/Python/pywarpx/picmi.py index 3cbefc920..b89c3fc43 100644 --- a/Python/pywarpx/picmi.py +++ b/Python/pywarpx/picmi.py @@ -180,7 +180,6 @@ class AnalyticDistribution(picmistandard.PICMI_AnalyticDistribution): species.zmin = self.lower_bound[2] species.zmax = self.upper_bound[2] - # --- Only constant density is supported at this time species.profile = "parse_density_function" species.__setattr__('density_function(x,y,z)', self.density_expression) @@ -188,7 +187,10 @@ class AnalyticDistribution(picmistandard.PICMI_AnalyticDistribution): setattr(pywarpx.my_constants, k, v) # --- Note that WarpX takes gamma*beta as input - if np.any(np.not_equal(self.rms_velocity, 0.)): + if np.any(np.not_equal(self.momentum_expressions, None)): + species.momentum_distribution_type = 'parse_momentum_function' + self.setup_parse_momentum_functions(species) + elif np.any(np.not_equal(self.rms_velocity, 0.)): species.momentum_distribution_type = "gaussian" species.ux_m = self.directed_velocity[0]/c species.uy_m = self.directed_velocity[1]/c @@ -205,6 +207,19 @@ class AnalyticDistribution(picmistandard.PICMI_AnalyticDistribution): if self.fill_in: species.do_continuous_injection = 1 + def setup_parse_momentum_functions(self, species): + if self.momentum_expressions[0] is not None: + species.__setattr__('momentum_function_ux(x,y,z)', '({0})/{1}'.format(self.momentum_expressions[0], c)) + else: + species.__setattr__('momentum_function_ux(x,y,z)', '({0})/{1}'.format(self.directed_velocity[0], c)) + if self.momentum_expressions[1] is not None: + species.__setattr__('momentum_function_uy(x,y,z)', '({0})/{1}'.format(self.momentum_expressions[1], c)) + else: + species.__setattr__('momentum_function_uy(x,y,z)', '({0})/{1}'.format(self.directed_velocity[1], c)) + if self.momentum_expressions[2] is not None: + species.__setattr__('momentum_function_uz(x,y,z)', '({0})/{1}'.format(self.momentum_expressions[2], c)) + else: + species.__setattr__('momentum_function_uz(x,y,z)', '({0})/{1}'.format(self.directed_velocity[2], c)) class ParticleListDistribution(picmistandard.PICMI_ParticleListDistribution): def init(self, kw): @@ -259,13 +274,13 @@ class CylindricalGrid(picmistandard.PICMI_CylindricalGrid): assert self.lower_bound[0] >= 0., Exception('Lower radial boundary must be >= 0.') assert self.bc_rmin != 'periodic' and self.bc_rmax != 'periodic', Exception('Radial boundaries can not be periodic') - assert self.n_azimuthal_modes is None or self.n_azimuthal_modes == 1, Exception('Only one azimuthal mode supported') # Geometry pywarpx.geometry.coord_sys = 1 # RZ pywarpx.geometry.is_periodic = '0 %d'%(self.bc_zmin=='periodic') # Is periodic? pywarpx.geometry.prob_lo = self.lower_bound # physical domain pywarpx.geometry.prob_hi = self.upper_bound + pywarpx.warpx.nmodes = self.n_azimuthal_modes 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 |