update public distribution

based on internal repository c9a2ac8 2019-01-03 16:04:57 +0100
tagged rev-master-2.0.0

projects/twoatom/twoatom.py

@@ -1,5 +1,3 @@
-#!/usr/bin/env python
-
 """
 @package projects.twoatom
 Two-atom demo scattering calculation project
@@ -8,30 +6,129 @@ this file is specific to the project and the state of the data analysis,
 as it contains particular parameter values.
 """
 
+from __future__ import absolute_import
 from __future__ import division
-import sys
-import os
-import math
-import numpy as np
-import periodictable as pt
+from __future__ import print_function
+
 import argparse
 import logging
+import math
+import numpy as np
+import os.path
+import periodictable as pt
 
-# adjust the system path so that the main PMSCO code is found
-base_dir = os.path.dirname(__file__) or '.'
-package_dir = os.path.join(base_dir, '../..')
-package_dir = os.path.abspath(package_dir)
-sys.path.append(package_dir)
-
-import pmsco.pmsco
 import pmsco.cluster as mc
 import pmsco.project as mp
-import pmsco.data as md
 from pmsco.helpers import BraceMessage as BMsg
 
 logger = logging.getLogger(__name__)
 
 
+class TwoatomCluster(mc.ClusterGenerator):
+    """
+    cluster of two atoms.
+
+    atom A (top) is set at position (0, 0, 0), atom B (bottom) at (-dx, -dy, -dz)
+    where dx, dy and dz are calculated from model parameters.
+    the type of the atoms is set upon construction.
+
+    the model parameters are:
+    @arg @c model['dAB']   : distance between the two atoms in Angstrom.
+    @arg @c model['th']    : polar angle of the connection line, 0 = on top geometry.
+    @arg @c model['ph']    : azimuthal angle of the connection line, 0 = polar angle affects X coordinate.
+
+    the class is designed to be reusable in various projects.
+    object attributes refine the atom types and the mapping of project-specific model parameters.
+    """
+
+    ## @var atom_types (dict)
+    # chemical element numbers of the cluster atoms.
+    #
+    # atom 'A' is the top atom, 'B' the bottom one.
+    # upon construction both atoms are set to oxygen.
+    # to customize, call @ref set_atom_type.
+
+    ## @var model_dict (dict)
+    # mapping of model parameters to cluster parameters
+    #
+    # the default model parameters used by the cluster are 'dAB', 'th' and 'ph'.
+    # if the project uses other parameter names, e.g. 'dCO' instead of 'dAB',
+    # the project-specific names can be declared here.
+    # in the example, set model_dict['dAB'] = 'dCO'.
+
+    def __init__(self, project):
+        """
+        initialize the cluster generator.
+
+        the atoms and model dictionary are given default values.
+        see @ref set_atom_type and @ref model_dict for customization.
+
+        @param project: project instance.
+        """
+        super(TwoatomCluster, self).__init__(project)
+
+        self.atom_types = {'A': pt.O.number, 'B': pt.O.number}
+        self.model_dict = {'dAB': 'dAB', 'th': 'th', 'ph': 'ph'}
+
+    def set_atom_type(self, atom, element):
+        """
+        set the type (chemical element) of an atom.
+
+        @param atom: atom key, 'A' (top) or 'B' (bottom).
+        @param element: chemical element number or symbol.
+        """
+        try:
+            self.atom_types[atom] = int(element)
+        except ValueError:
+            self.atom_types[atom] = pt.elements.symbol(element.strip()).number
+
+    def count_emitters(self, model, index):
+        """
+        return the number of emitter configurations.
+
+        this cluster supports only one configuration.
+
+        @param model:
+        @param index:
+        @return 1
+        """
+        return 1
+
+    def create_cluster(self, model, index):
+        """
+        create a cluster given the model parameters and index.
+
+        @param model:
+        @param index:
+        @return a pmsco.cluster.Cluster object containing the atomic coordinates.
+        """
+        r = model[self.model_dict['dAB']]
+        try:
+            th = math.radians(model[self.model_dict['th']])
+        except KeyError:
+            th = 0.
+        try:
+            ph = math.radians(model[self.model_dict['ph']])
+        except KeyError:
+            ph = 0.
+
+        dx = r * math.sin(th) * math.cos(ph)
+        dy = r * math.sin(th) * math.sin(ph)
+        dz = r * math.cos(th)
+
+        clu = mc.Cluster()
+        clu.comment = "{0} {1}".format(self.__class__, index)
+        clu.set_rmax(r * 2.0)
+
+        a_top = np.array((0.0, 0.0, 0.0))
+        a_bot = np.array((-dx, -dy, -dz))
+
+        clu.add_atom(self.atom_types['A'], a_top, 1)
+        clu.add_atom(self.atom_types['B'], a_bot, 0)
+
+        return clu
+
+
 class TwoatomProject(mp.Project):
     """
     two-atom calculation project class.
@@ -49,31 +146,12 @@ class TwoatomProject(mp.Project):
     def __init__(self):
         super(TwoatomProject, self).__init__()
         self.scan_dict = {}
-
-    def create_cluster(self, model, index):
-        """
-        calculate a specific set of atom positions given the optimizable parameters.
-
-        the cluster contains a nitrogen in the top layer,
-        and a nickel atom in the second layer.
-        The layer distance and the angle can be adjusted by parameters.
-
-        @param model: (dict) optimizable parameters
-        """
-        clu = mc.Cluster()
-        clu.comment = "{0} {1}".format(self.__class__, index)
-        clu.set_rmax(10.0)
-
-        a_N = np.array((0.0, 0.0, 0.0))
-        rad_pNNi = math.radians(model['pNNi'])
-        a_Ni1 = np.array((0.0,
-            -model['dNNi'] * math.sin(rad_pNNi),
-            -model['dNNi'] * math.cos(rad_pNNi)))
-
-        clu.add_atom(pt.N.number, a_N, 1)
-        clu.add_atom(pt.Ni.number, a_Ni1, 0)
-
-        return clu
+        self.cluster_generator = TwoatomCluster(self)
+        self.cluster_generator.set_atom_type('A', 'N')
+        self.cluster_generator.set_atom_type('B', 'Ni')
+        self.cluster_generator.model_dict['dAB'] = 'dNNi'
+        self.cluster_generator.model_dict['th'] = 'pNNi'
+        self.cluster_generator.model_dict['ph'] = 'aNNi'
 
     def create_params(self, model, index):
         """
@@ -93,13 +171,13 @@ class TwoatomProject(mp.Project):
         params.scattering_level = 5
         params.fcut = 15.0
         params.cut = 15.0
-        params.angular_broadening = 0.0
+        params.angular_resolution = 0.0
         params.lattice_constant = 1.0
         params.z_surface = model['Zsurf']
-        params.atom_types = 3
-        params.atomic_number = [7, 28]
-        params.phase_file = ["hbn_n.pha", "ni.pha"]
-        params.msq_displacement = [0.01, 0.01, 0.00]
+        params.phase_files = {self.cluster_generator.atom_types['A']: "",
+                              self.cluster_generator.atom_types['B']: ""}
+        params.msq_displacement = {self.cluster_generator.atom_types['A']: 0.01,
+                                   self.cluster_generator.atom_types['B']: 0.0}
         params.planewave_attenuation = 1.0
         params.inner_potential = model['V0']
         params.work_function = 3.6
@@ -153,6 +231,27 @@ class TwoatomProject(mp.Project):
         return dom
 
 
+def example_intensity(e, t, p, a):
+    """
+    arbitrary intensity pattern for example data
+
+    this function can be used to calculate the intensity in example scan files.
+    the function implements an arbitrary modulation function
+
+    @param e: energy
+    @param t: theta
+    @param p: phi
+    @param a: alpha
+    @return intensity
+    """
+    i = np.random.random() * 1e6 * \
+        np.cos(np.radians(t)) ** 2 * \
+        np.cos(np.radians(a)) ** 2 * \
+        np.cos(np.radians(p)) ** 2 * \
+        np.sin(e / 1000. * np.pi * 0.1 / np.sqrt(e)) ** 2
+    return i
+
+
 def create_project():
     """
     create a new TwoatomProject calculation project.
@@ -230,20 +329,3 @@ def parse_project_args(_args):
     parsed_args = parser.parse_args(_args)
 
     return parsed_args
-
-
-def main():
-    args, unknown_args = pmsco.pmsco.parse_cli()
-    if unknown_args:
-        project_args = parse_project_args(unknown_args)
-    else:
-        project_args = None
-
-    project = create_project()
-    pmsco.pmsco.set_common_args(project, args)
-    set_project_args(project, project_args)
-    pmsco.pmsco.run_project(project)
-
-if __name__ == '__main__':
-    main()
-    sys.exit(0)