pmsco-public/projects/demo/molecule.py

385 lines
13 KiB
Python

"""
@package pmsco.projects.demo.molecule
scattering calculation project for single molecules
the atomic positions are read from a molecule file.
cluster file, emitter (by chemical symbol), initial state and kinetic energy are specified on the command line.
there are no structural parameters.
@author Matthias Muntwiler, matthias.muntwiler@psi.ch
@copyright (c) 2015-20 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
"""
import math
import numpy as np
import os.path
from pathlib import Path
import periodictable as pt
import argparse
import logging
# noinspection PyUnresolvedReferences
from pmsco.calculators.calculator import InternalAtomicCalculator
# noinspection PyUnresolvedReferences
from pmsco.calculators.edac import EdacCalculator
# noinspection PyUnresolvedReferences
from pmsco.calculators.phagen.runner import PhagenCalculator
import pmsco.cluster as cluster
from pmsco.data import calc_modfunc_loess
# noinspection PyUnresolvedReferences
import pmsco.elements.bindingenergy
from pmsco.helpers import BraceMessage as BMsg
import pmsco.project as project
logger = logging.getLogger(__name__)
class MoleculeFileCluster(cluster.ClusterGenerator):
"""
cluster generator based on external file.
work in progress.
"""
def __init__(self, project):
super(MoleculeFileCluster, self).__init__(project)
self.base_cluster = None
def load_base_cluster(self):
"""
load and cache the project-defined coordinate file.
the file path is set in self.project.cluster_file.
the file must be in XYZ (.xyz) or PMSCO cluster (.clu) format (cf. pmsco.cluster module).
@return: Cluster object (also referenced by self.base_cluster)
"""
if self.base_cluster is None:
clu = cluster.Cluster()
clu.set_rmax(120.0)
p = Path(self.project.cluster_file)
ext = p.suffix
if ext == ".xyz":
fmt = cluster.FMT_XYZ
elif ext == ".clu":
fmt = cluster.FMT_PMSCO
else:
raise ValueError(f"unknown cluster file extension {ext}")
clu.load_from_file(self.project.cluster_file, fmt=fmt)
self.base_cluster = clu
return self.base_cluster
def count_emitters(self, model, index):
"""
count the number of emitter configurations.
the method creates the full cluster and counts the emitters.
@param model: model parameters.
@param index: scan and domain are used by the create_cluster() method,
emit decides whether the method returns the number of configurations (-1),
or the number of emitters in the specified configuration (>= 0).
@return: number of emitter configurations.
"""
clu = self.create_cluster(model, index)
return clu.get_emitter_count()
def create_cluster(self, model, index):
"""
import a cluster from a coordinate file (XYZ format).
the method does the following:
- load the cluster file specified by self.cluster_file.
- trim the cluster according to model['rmax'].
- mark the 6 nitrogen atoms at the center of the trimer as emitters.
@param model: rmax is the trim radius of the cluster in units of the surface lattice constant.
@param index (named tuple CalcID) calculation index.
this method uses the domain index to look up domain parameters in
`pmsco.project.Project.domains`.
`index.emit` selects whether a single-emitter (>= 0) or all-emitter cluster (== -1) is returned.
@return pmsco.cluster.Cluster object
"""
self.load_base_cluster()
clu = cluster.Cluster()
clu.copy_from(self.base_cluster)
clu.comment = f"{self.__class__}, {index}"
dom = self.project.domains[index.domain]
# trim
clu.set_rmax(model['rmax'])
clu.trim_sphere(clu.rmax)
# emitter selection
idx_emit = np.where(clu.data['s'] == self.project.scans[index.scan].emitter)
assert isinstance(idx_emit, tuple)
idx_emit = idx_emit[0]
if index.emit >= 0:
idx_emit = idx_emit[index.emit]
clu.data['e'][idx_emit] = 1
# rotation
if 'xrot' in model:
clu.rotate_z(model['xrot'])
elif 'xrot' in dom:
clu.rotate_z(dom['xrot'])
if 'yrot' in model:
clu.rotate_z(model['yrot'])
elif 'yrot' in dom:
clu.rotate_z(dom['yrot'])
if 'zrot' in model:
clu.rotate_z(model['zrot'])
elif 'zrot' in dom:
clu.rotate_z(dom['zrot'])
logger.info(f"cluster for calculation {index}: "
f"{clu.get_atom_count()} atoms, {clu.get_emitter_count()} emitters")
return clu
class MoleculeProject(project.Project):
"""
general molecule project.
the following model parameters are used:
@arg `model['zsurf']` : position of surface above molecule (angstrom)
@arg `model['Texp']` : experimental temperature (K)
@arg `model['Tdeb']` : debye temperature (K)
@arg `model['V0']` : inner potential (eV)
@arg `model['rmax']` : cluster radius (angstrom)
@arg `model['ares']` : angular resolution (degrees, FWHM)
@arg `model['distm']` : dmax for EDAC (angstrom)
the following domain parameters are used.
they can also be specified as model parameters.
@arg `'xrot'` : rotation about x-axis (applied first) (deg)
@arg `'yrot'` : rotation about y-axis (applied after x) (deg)
@arg `'zrot'` : rotation about z-axis (applied after x and y) (deg)
the project parameters are:
@arg `cluster_file` : name of cluster file of template molecule.
default: "dpdi-trimer.xyz"
"""
def __init__(self):
"""
initialize a project instance
"""
super(MoleculeProject, self).__init__()
self.model_space = project.ModelSpace()
self.scan_dict = {}
self.cluster_file = "demo-cluster.xyz"
self.cluster_generator = MoleculeFileCluster(self)
self.atomic_scattering_factory = PhagenCalculator
self.multiple_scattering_factory = EdacCalculator
self.phase_files = {}
self.rme_files = {}
self.modf_smth_ei = 0.5
def create_params(self, model, index):
"""
set a specific set of parameters given the optimizable parameters.
@param model: (dict) optimization parameters
this method requires zsurf, V0, Texp, Tdeb, ares and distm.
@param index (named tuple CalcID) calculation index.
this method formats the index into the comment line.
"""
params = project.CalculatorParams()
params.title = "molecule demo"
params.comment = f"{self.__class__} {index}"
params.cluster_file = ""
params.output_file = ""
initial_state = self.scans[index.scan].initial_state
params.initial_state = initial_state
emitter = self.scans[index.scan].emitter
params.binding_energy = pt.elements.symbol(emitter).binding_energy[initial_state]
params.polarization = "H"
params.z_surface = model['zsurf']
params.inner_potential = model['V0']
params.work_function = 4.5
params.polar_incidence_angle = 60.0
params.azimuthal_incidence_angle = 0.0
params.angular_resolution = model['ares']
params.experiment_temperature = model['Texp']
params.debye_temperature = model['Tdeb']
params.phase_files = self.phase_files
params.rme_files = self.rme_files
# edac_interface only
params.emitters = []
params.lmax = 15
params.dmax = model['distm']
params.orders = [20]
return params
def create_model_space(self):
"""
define the range of model parameters.
see the class description for a list of parameters.
"""
return self.model_space
# noinspection PyUnusedLocal
def calc_modulation(self, data, model):
"""
calculate the modulation function with project-specific smoothing factor
see @ref pmsco.pmsco.project.calc_modulation.
@param data: (numpy.ndarray) experimental data in ETPI, or ETPAI format.
@param model: (dict) model parameters of the calculation task. not used.
@return copy of the data array with the modulation function in the 'i' column.
"""
return calc_modfunc_loess(data, smth=self.modf_smth_ei)
def create_model_space(mode):
"""
define the model space.
"""
dom = project.ModelSpace()
if mode == "single":
dom.add_param('zsurf', 1.20)
dom.add_param('Texp', 300.00)
dom.add_param('Tdeb', 100.00)
dom.add_param('V0', 10.00)
dom.add_param('rmax', 50.00)
dom.add_param('ares', 5.00)
dom.add_param('distm', 5.00)
dom.add_param('wdom1', 1.0)
dom.add_param('wdom2', 1.0)
dom.add_param('wdom3', 1.0)
dom.add_param('wdom4', 1.0)
dom.add_param('wdom5', 1.0)
else:
raise ValueError(f"undefined model space for {mode} optimization")
return dom
def create_project():
"""
create the project instance.
"""
proj = MoleculeProject()
proj_dir = os.path.dirname(os.path.abspath(__file__))
proj.project_dir = proj_dir
# scan dictionary
# to select any number of scans, add their dictionary keys as scans option on the command line
proj.scan_dict['empty'] = {'filename': os.path.join(proj_dir, "../common/empty-hemiscan.etpi"),
'emitter': "N", 'initial_state': "1s"}
proj.mode = 'single'
proj.model_space = create_model_space(proj.mode)
proj.job_name = 'molecule0000'
proj.description = 'molecule demo'
return proj
def set_project_args(project, project_args):
"""
set the project arguments.
@param project: project instance
@param project_args: (Namespace object) project arguments.
"""
scans = []
try:
if project_args.scans:
scans = project_args.scans
else:
logger.error("missing scan argument")
exit(1)
except AttributeError:
logger.error("missing scan argument")
exit(1)
for scan_key in scans:
scan_spec = project.scan_dict[scan_key]
project.add_scan(**scan_spec)
try:
project.cluster_file = os.path.abspath(project_args.cluster_file)
project.cluster_generator = MoleculeFileCluster(project)
except (AttributeError, TypeError):
logger.error("missing cluster-file argument")
exit(1)
try:
if project_args.emitter:
for scan in project.scans:
scan.emitter = project_args.emitter
logger.warning(f"override emitters of all scans to {project_args.emitter}")
except AttributeError:
pass
try:
if project_args.initial_state:
for scan in project.scans:
scan.initial_state = project_args.initial_state
logger.warning(f"override initial states of all scans to {project_args.initial_state}")
except AttributeError:
pass
try:
if project_args.energy:
for scan in project.scans:
scan.energies = np.asarray((project_args.energy, ))
logger.warning(f"override scan energy of all scans to {project_args.energy}")
except AttributeError:
pass
try:
if project_args.symmetry:
for angle in np.linspace(0, 360, num=project_args.symmetry, endpoint=False):
project.add_domain({'xrot': 0., 'yrot': 0., 'zrot': angle})
logger.warning(f"override rotation symmetry to {project_args.symmetry}")
except AttributeError:
pass
def parse_project_args(_args):
parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
# main arguments
parser.add_argument('--scans', nargs="*",
help="nick names of scans to use in calculation (see create_project function)")
parser.add_argument('--cluster-file',
help="path name of molecule file (xyz format).")
# conditional arguments
parser.add_argument('--emitter',
help="emitter: chemical symbol")
parser.add_argument('--initial-state',
help="initial state term: e.g. 2p1/2")
parser.add_argument('--energy', type=float,
help="kinetic energy (eV)")
parser.add_argument('--symmetry', type=int, default=1,
help="n-fold rotational symmetry")
parsed_args = parser.parse_args(_args)
return parsed_args