public release 3.0.0 - see README and CHANGES for details

pmsco/cluster.py

@@ -17,22 +17,20 @@ pip install --user periodictable
 
 @author Matthias Muntwiler
 
-@copyright (c) 2015-20 by Paul Scherrer Institut @n
+@copyright (c) 2015-21 by Paul Scherrer Institut @n
 Licensed under the Apache License, Version 2.0 (the "License"); @n
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at
   http://www.apache.org/licenses/LICENSE-2.0
 """
 
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 import math
 import numpy as np
 import periodictable as pt
 import sys
 
+import pmsco.config as config
+
 ## default file format identifier
 FMT_DEFAULT = 0
 ## MSC file format identifier
@@ -227,13 +225,13 @@ class Cluster(object):
         """
         self.rmax = r
 
-    def build_element(self, index, element_number, x, y, z, emitter, charge=0., scatterer_class=0):
+    def build_element(self, index, element, x, y, z, emitter, charge=0., scatterer_class=0):
         """
         build a tuple in the format of the internal data array.
         
         @param index: (int) index
         
-        @param element_number: (int) chemical element number
+        @param element: chemical element number (int) or symbol (str)
         
         @param x, y, z: (float) atom coordinates in the cluster
         
@@ -243,7 +241,13 @@ class Cluster(object):
 
         @param scatterer_class: (int) scatterer class. default = 0.
         """
-        symbol = pt.elements[element_number].symbol
+        try:
+            element_number = int(element)
+            symbol = pt.elements[element_number].symbol
+        except ValueError:
+            symbol = element
+            element_number = pt.elements.symbol(symbol.strip()).number
+
         element = (index, element_number, symbol, scatterer_class, x, y, z, int(emitter), charge)
         return element
 
@@ -251,7 +255,7 @@ class Cluster(object):
         """
         add a single atom to the cluster.
         
-        @param atomtype: (int) chemical element number
+        @param atomtype: chemical element number (int) or symbol (str)
         
         @param v_pos: (numpy.ndarray, shape = (3)) position vector
         
@@ -274,7 +278,7 @@ class Cluster(object):
         self.rmax (maximum distance from the origin).
         all atoms are non-emitters.
         
-        @param atomtype: (int) chemical element number
+        @param atomtype: chemical element number (int) or symbol (str)
         
         @param v_pos: (numpy.ndarray, shape = (3))
             position vector of the first atom (basis vector)
@@ -307,7 +311,7 @@ class Cluster(object):
         and z_surf (position of the surface).
         all atoms are non-emitters.
 
-        @param atomtype: (int) chemical element number
+        @param atomtype: chemical element number (int) or symbol (str)
         
         @param v_pos: (numpy.ndarray, shape = (3))
             position vector of the first atom (basis vector)
@@ -1133,7 +1137,7 @@ class Cluster(object):
         np.savetxt(f, data, fmt=file_format, header=header, comments="")
 
 
-class ClusterGenerator(object):
+class ClusterGenerator(config.ConfigurableObject):
     """
     cluster generator class.
 
@@ -1151,6 +1155,7 @@ class ClusterGenerator(object):
         @param project: reference to the project object.
             cluster generators may need to look up project parameters.
         """
+        super().__init__()
         self.project = project
 
     def count_emitters(self, model, index):
@@ -1258,7 +1263,7 @@ class LegacyClusterGenerator(ClusterGenerator):
     """
 
     def __init__(self, project):
-        super(LegacyClusterGenerator, self).__init__(project)
+        super().__init__(project)
 
     def count_emitters(self, model, index):
         """