In picmi, add support for all elements

2 files changed, 23 insertions, 4 deletions

diff --git a/Python/pywarpx/picmi.py b/Python/pywarpx/picmi.py
index 003ef523b..a0069ec8f 100644
--- a/Python/pywarpx/picmi.py
+++ b/Python/pywarpx/picmi.py
@@ -1,8 +1,10 @@
 """Classes following the PICMI standard
 """
+import re
 import picmistandard
 import numpy as np
 import pywarpx
+import periodictable
 
 codename = 'warpx'
 picmistandard.register_codename(codename)
@@ -33,9 +35,26 @@ class Species(picmistandard.PICMI_Species):
         elif self.particle_type == 'anti-proton':
             if self.charge is None: self.charge = '-q_e'
             if self.mass is None: self.mass = 'm_p'
-        elif self.particle_type == 'H' and self.charge_state == 1:
-            if self.charge is None: self.charge = 'q_e'
-            if self.mass is None: self.mass = 'm_p'
+        else:
+            if self.charge is None and self.charge_state is not None:
+                self.charge = self.charge_state*constants.q_e
+            # Match a string of the format '#nXx', with the '#n' optional isotope number.
+            m = re.match('(?P<iso>#[\d+])*(?P<sym>[A-Za-z]+)', self.particle_type)
+            if m is not None:
+                element = periodictable.elements.symbol(m['sym'])
+                if m['iso'] is not None:
+                    element = element[m['iso'][1:]]
+                if self.charge_state is not None:
+                    assert self.charge_state <= element.number, Exception('%s charge state not valid'%self.particle_type)
+                    try:
+                        element = element.ion[self.charge_state]
+                    except ValueError:
+                        # Note that not all valid charge states are defined in elements,
+                        # so this value error can be ignored.
+                        pass
+                self.element = element
+                if self.mass is None:
+                    self.mass = element.mass*periodictable.constants.atomic_mass_constant
 
     def initialize_inputs(self, layout):
         self.species_number = pywarpx.particles.nspecies
diff --git a/Python/setup.py b/Python/setup.py
index ecb87190c..8429cf95e 100644
--- a/Python/setup.py
+++ b/Python/setup.py
@@ -25,5 +25,5 @@ setup (name = 'pywarpx',
        package_dir = {'pywarpx':'pywarpx'},
        description = """Wrapper of WarpX""",
        package_data = package_data,
-       install_requires=['picmistandard']
+       install_requires=['picmistandard', 'periodictable']
        )