Source code for abtem.parametrizations

"""Module for describing analytical potential parametrizations."""
from __future__ import annotations

import json
import os
from abc import ABCMeta, abstractmethod
from numbers import Number
from typing import Sequence

import numpy as np
from ase.data import chemical_symbols
from scipy.optimize import least_squares

from abtem.array import concatenate
from abtem.core.axes import OrdinalAxis
from abtem.core.config import config
from abtem.core.constants import kappa
from abtem.core.utils import EqualityMixin, get_dtype
from abtem.measurements import ReciprocalSpaceLineProfiles, RealSpaceLineProfiles
from abtem.parametrizations.functions import kirkland, ewald, peng, lobato


def _get_data_path():
    this_file = os.path.abspath(os.path.dirname(__file__))
    return os.path.join(this_file, "data")


[docs] def validate_sigmas(sigmas: float | dict) -> dict: if sigmas is None: sigmas = {} elif isinstance(sigmas, Number): sigmas = { chemical_symbol: sigmas for chemical_symbol in chemical_symbols.keys() } if not isinstance(sigmas, dict): raise ValueError() return sigmas
[docs] class Parametrization(EqualityMixin, metaclass=ABCMeta): """ Base class for potential parametrizations. Parameters ---------- parameters : str or dict A given string must be either the full path to a valid ``.json`` file or the name of a valid ``.json`` file in the ``abtem/parametrizations/data/`` folder. A given dictionary must be a mapping from chemical symbols to anything that can converted to a NumPy array with the correct shape for the given parametrization. sigmas : dict or float The standard deviation of isotropic displacements for each element as a mapping from chemical symbols to a number. If given as a float the standard deviation of the displacements is assumed to be identical for all atoms. """ _functions: dict _real_space_funcs = ( "potential", "projected_potential", "charge", "finite_projected_potential", ) _fourier_space_funcs = ( "scattering_factor", "projected_scattering_factor", "x_ray_scattering_factor", "finite_projected_scattering_factor", )
[docs] def __init__( self, parameters: dict[str, np.ndarray] | str, sigmas: dict[str, float] = None ): self._parameters = validate_parameters(parameters) self._sigmas = validate_sigmas(sigmas)
def to_json(self, file: str): with open(file, "w") as fp: data = { symbol: parameters.tolist() for symbol, parameters in self.parameters.items() } json.dump(data, fp) def from_json(self, file: str): with open(file, "r") as fp: self._parameters = { symbol: np.array(parameters) for symbol, parameters in json.load(fp).items() } @property def sigmas(self): """The standard deviation of isotropic displacements.""" return self._sigmas @property def parameters(self) -> dict[str, np.ndarray]: return self._parameters
[docs] @abstractmethod def scaled_parameters(self, symbol: str, name: str) -> np.ndarray: """ The parameters of the parametrization scaled to the abTEM units for a given function. Parameters ---------- symbol : str Chemical symbol to get the scaled parameters for. name : str Name of the function to get the scaled parameters for. Returns ------- scaled_parameters : np.ndarray """ pass
[docs] def potential(self, symbol: str, charge: float = 0.0) -> callable: """ Radial electrostatic potential for given chemical symbol and charge. Parameters ---------- symbol : str Chemical symbol of element. charge : float Charge of element. Given as elementary charges. Returns ------- electrostatic_potential : callable Function describing electrostatic potential parameterized by the radial distance to the core [Å]. """ return self.get_function("potential", symbol, charge)
[docs] def scattering_factor(self, symbol: str, charge: float = 0.0) -> callable: """ Radial scattering factor for given chemical symbol and charge. Parameters ---------- symbol : str Chemical symbol of element. charge : float Charge of element. Given as elementary charges. Returns ------- scattering_factor : callable Function describing scattering parameterized by the squared reciprocal radial distance to the core [1/Å^2]. """ return self.get_function("scattering_factor", symbol, charge)
[docs] def projected_potential(self, symbol: str, charge: float = 0.0) -> callable: """ Analytical infinite projection of radial electrostatic potential for given chemical symbol and charge. Parameters ---------- symbol : str Chemical symbol of element. charge : float Charge of element. Given as elementary charges. Returns ------- projected_potential : callable Function describing projected electrostatic potential parameterized by the radial distance to the core [Å]. """ return self.get_function("projected_potential", symbol, charge)
[docs] def projected_scattering_factor(self, symbol: str, charge: float = 0.0) -> callable: """ Analytical infinite projection of radial scattering factor for given chemical symbol and charge. Parameters ---------- symbol : str Chemical symbol of element. charge : float Charge of element. Given as elementary charges. Returns ------- projected_scattering_factor : callable Function describing projected scattering parameterized by the squared reciprocal radial distance to the core [1/Å^2]. """ return self.get_function("projected_scattering_factor", symbol, charge)
[docs] def charge(self, symbol: str, charge: float = 0.0) -> callable: """ Radial charge distribution for given chemical symbol and charge. Parameters ---------- symbol : str Chemical symbol of element. charge : float Charge of element. Given as elementary charges. Returns ------- charge : callable Function describing charge parameterized by the radial distance to the core [Å]. """ return self.get_function("charge", symbol, charge)
[docs] def x_ray_scattering_factor(self, symbol: str, charge: float = 0.0) -> callable: """ X-ray scattering factor for given chemical symbol and charge. Parameters ---------- symbol : str Chemical symbol of element. charge : float Charge of element. Given as elementary charges. Returns ------- x_ray_scattering_factor : callable """ return self.get_function("x_ray_scattering_factor", symbol, charge)
[docs] def finite_projected_potential(self, symbol: str, charge: float = 0.0) -> callable: """ X-ray scattering factor for given chemical symbol and charge. Parameters ---------- symbol : str Chemical symbol of element. charge : float Charge of element. Given as elementary charges. Returns ------- x_ray_scattering_factor : callable """ return self.get_function("finite_projected_potential", symbol, charge)
[docs] def finite_projected_scattering_factor( self, symbol: str, charge: float = 0.0 ) -> callable: """ Analytical infinite projection of radial scattering factor for given chemical symbol and charge. Parameters ---------- symbol : str Chemical symbol of element. charge : float Charge of element. Given as elementary charges. Returns ------- projected_scattering_factor : callable Function describing projected scattering parameterized by the squared reciprocal radial distance to the core [1/Å^2]. """ return self.get_function("finite_projected_scattering_factor", symbol, charge)
[docs] def get_function(self, name: str, symbol: str, charge: float = 0.0) -> callable: """ Returns the line profiles for a parameterized function for one or more element. Parameters ---------- name : str Name of the function to return. symbol : str Chemical symbol of element. charge : float Charge of element. Given as elementary charges. """ if isinstance(symbol, (int, np.int32, np.int64)): symbol = chemical_symbols[symbol] if charge > 0.0: raise NotImplementedError # if charge > 0.0: # raise RuntimeError( # f"charge not implemented for parametrization {self.__class__.__name__}" # ) if charge == 0.0: charge_symbol = "" elif charge > 0.0: charge_symbol = "+" * int(abs(charge)) else: charge_symbol = "-" * int(abs(charge)) try: func = self._functions[name] parameters = self.scaled_parameters(symbol + charge_symbol, name) dtype = get_dtype(complex=False) parameters = np.array(parameters, dtype=dtype) return lambda r, *args, **kwargs: func(r, parameters, *args, **kwargs) except KeyError: raise RuntimeError( f'parametrized function "{name}" does not exist for element {symbol} with charge {charge}' )
[docs] def line_profiles( self, symbol: str | Sequence[str], cutoff: float, sampling: float = 0.001, name: str = "potential", ) -> RealSpaceLineProfiles | ReciprocalSpaceLineProfiles: """ Returns the line profiles for a parameterized function for one or more element. Parameters ---------- symbol : str or list of str Chemical symbol(s) of atom(s). cutoff : float The outer radial cutoff distance of the line profiles in units of Ångstrom for potentials and units of reciprocal Ångstrom for scattering factors. sampling : float, optional The radial sampling of the line profiles in units of Ångstrom for potentials and units of reciprocal Ångstrom for scattering factors. Default is 0.001. name : str Name of the line profile to return. """ if not isinstance(symbol, str): return concatenate( [ self.line_profiles( s, cutoff=cutoff, sampling=sampling, name=name, ) for s in symbol ] ) func = self.get_function(name, symbol) ensemble_axes_metadata = [ OrdinalAxis(label="", values=(symbol,), _default_type="overlay") ] if name in self._real_space_funcs: r = np.arange(sampling, cutoff, sampling) metadata = {"label": "potential", "units": "eV/e"} return RealSpaceLineProfiles( func(r)[None], sampling=sampling, ensemble_axes_metadata=ensemble_axes_metadata, metadata=metadata, ) elif name in self._fourier_space_funcs: k2 = np.arange(0.0, cutoff, sampling) ** 2 metadata = {"label": "scattering factor", "units": "Å"} return ReciprocalSpaceLineProfiles( func(k2)[None], sampling=sampling, ensemble_axes_metadata=ensemble_axes_metadata, metadata=metadata, ) else: raise RuntimeError(f"function name {name} not recognized")
[docs] class KirklandParametrization(Parametrization): """ Potential parametrization by Earl J Kirkland. Parameters ---------- parameters : str or dict A given string must be either the full path to a valid ``.json`` file or the name of a valid ``.json`` file in the ``abtem/parametrizations/data/`` folder. A given dictionary must be a mapping from chemical symbols to anything that can converted to a NumPy array of shape (4, 3). sigmas : dict or float The standard deviation of isotropic displacements for each element as a mapping from chemical symbols to a number. If given as a float the standard deviation of the displacements is assumed to be identical for all atoms. References ---------- E.J. Kirkland. Advanced computing in electron microscopy. Springer, 2. edition, 2010. """ _functions = { "potential": kirkland.potential, "scattering_factor": kirkland.scattering_factor, "projected_potential": kirkland.projected_potential, "projected_scattering_factor": kirkland.projected_scattering_factor, }
[docs] def __init__( self, parameters: str | dict = "kirkland.json", sigmas: dict[str, float] = None ): super().__init__(parameters=parameters, sigmas=sigmas)
def fit(self, Z, k, f, guess=None): def reshape_parameters(p): p = p.reshape((4, 3)) return p def apply_constraint(p): p = p.copy() p = np.abs(p) return p def make_residuals_func(k2, target, func): def residuals_func(p): p = reshape_parameters(p) p = apply_constraint(p) return target - func(k2, p) return residuals_func if guess is None: if chemical_symbols[Z] in self.parameters: guess = self.scaled_parameters(chemical_symbols[Z], "scattering_factor") else: guess = np.array( validate_parameters("kirkland.json")[chemical_symbols[Z]] ) func = self._functions["scattering_factor"] residuals_func = make_residuals_func(k**2, f, func) result = least_squares(residuals_func, guess.ravel()) p_optimal = reshape_parameters(result.x) p_optimal = apply_constraint(p_optimal) self.parameters[chemical_symbols[Z]] = p_optimal return p_optimal
[docs] def scaled_parameters(self, symbol: str, name: str) -> np.ndarray: parameters = np.array(self.parameters[symbol]) a = np.pi * parameters[0] / kappa b = 2.0 * np.pi * np.sqrt(parameters[1]) c = np.pi ** (3.0 / 2.0) * parameters[2] / parameters[3] ** (3.0 / 2.0) / kappa d = np.pi**2 / parameters[3] scaled_parameters = np.vstack([a, b, c, d]) scaled_parameters = { "potential": scaled_parameters, "scattering_factor": parameters, "projected_potential": scaled_parameters, "projected_scattering_factor": scaled_parameters, } return scaled_parameters[name]
[docs] class LobatoParametrization(Parametrization): """ Potential parametrization by Ivan Lobato and Dirk Van Dyck. Parameters ---------- parameters : str or dict A given string must be either the full path to a valid ``.json`` file or the name of a valid ``.json`` file in the ``abtem/parametrizations/data/`` folder. A given dictionary must be a mapping from chemical symbols to anything that can converted to a NumPy array of shape (2, 5). sigmas : dict or float The standard deviation of isotropic displacements for each element as a mapping from chemical symbols to a number. If given as a float the standard deviation of the displacements is assumed to be identical for all atoms. References ---------- Ivan Lobato and Dirk Van Dyck. Acta Crystallographica Section A, 70:636–649, 2014. """ _functions = { "potential": lobato.potential, "scattering_factor": lobato.scattering_factor, "projected_potential": lobato.projected_potential, "projected_scattering_factor": lobato.projected_scattering_factor, "x_ray_scattering_factor": lobato.x_ray_scattering_factor, "charge": lobato.charge, }
[docs] def __init__( self, parameters: str | dict = "lobato.json", sigmas: dict[str, float] = None ): super().__init__(parameters=parameters, sigmas=sigmas)
def fit(self, Z, k, f, guess=None): def reshape_parameters(p): p = p.reshape((2, 5)) return p def apply_constraint(p): p = p.copy() p[1, :] = np.abs(p[1, :]) return p def make_residuals_func(k2, target, func): def residuals_func(p): p = reshape_parameters(p) p = apply_constraint(p) return target - func(k2, p) return residuals_func if guess is None: if chemical_symbols[Z] in self.parameters: guess = self.scaled_parameters(chemical_symbols[Z], "scattering_factor") else: guess = np.array( validate_parameters("lobato.json")[chemical_symbols[Z]] ) func = self._functions["scattering_factor"] residuals_func = make_residuals_func(k**2, f, func) result = least_squares(residuals_func, guess.ravel()) p_optimal = reshape_parameters(result.x) p_optimal = apply_constraint(p_optimal) self.parameters[chemical_symbols[Z]] = p_optimal return p_optimal
[docs] def scaled_parameters(self, symbol: str, name: str) -> np.ndarray: parameters = np.array(self.parameters[symbol]) a = np.pi**2 * parameters[0] / parameters[1] ** (3 / 2.0) / kappa b = 2 * np.pi / np.sqrt(parameters[1]) scaled_parameters = np.vstack((a, b)) scaled_parameters = { "potential": scaled_parameters, "scattering_factor": parameters, "projected_potential": scaled_parameters, "projected_scattering_factor": scaled_parameters, "x_ray_scattering_factor": parameters, "charge": parameters, } return scaled_parameters[name]
[docs] class PengParametrization(Parametrization): """ Potential parametrization by Lian-Mao Peng. Parameters ---------- parameters : str or dict A given string must be either the full path to a valid ``.json`` file or the name of a valid ``.json`` file in the ``abtem/parametrizations/data/`` folder. A given dictionary must be a mapping from chemical symbols to anything that can converted to a NumPy array of shape (2, 5). sigmas : dict or float The standard deviation of isotropic displacements for each element as a mapping from chemical symbols to a number. If given as a float the standard deviation of the displacements is assumed to be identical for all atoms. References ---------- L. Peng. Micron, 30(6):625–648, 1999. """ _functions = { "potential": peng.scattering_factor, "scattering_factor": peng.scattering_factor_k2, "projected_potential": peng.scattering_factor, "projected_scattering_factor": peng.scattering_factor_k2, "finite_projected_potential": peng.finite_projected_scattering_factor, "finite_projected_scattering_factor": peng.finite_projected_scattering_factor, }
[docs] def __init__( self, parameters: str | dict = "peng_high.json", sigmas: dict[str, float] = None ): super().__init__(parameters=parameters, sigmas=sigmas)
[docs] def scaled_parameters(self, symbol: str, name: str) -> np.ndarray: scattering_factor = np.array(self.parameters[symbol]) scattering_factor[1] /= 2**2 # convert scattering factor units potential = np.vstack( ( np.pi ** (3.0 / 2.0) * scattering_factor[0] / scattering_factor[1] ** (3 / 2.0) / kappa, np.pi**2 / scattering_factor[1], ) ) projected_potential = np.vstack( [ 1 / kappa * np.pi * scattering_factor[0] / scattering_factor[1], np.pi**2 / scattering_factor[1], ] ) projected_scattering_factor = np.vstack( [scattering_factor[0] / kappa, scattering_factor[1]] ) finite_projected_scattering_factor = np.concatenate( ( projected_scattering_factor, [np.pi / np.sqrt(projected_scattering_factor[1])], ) ) finite_projected_potential = np.concatenate( (projected_potential, [np.sqrt(projected_potential[1])]) ) scaled_parameters = { "potential": potential, "scattering_factor": scattering_factor, "finite_projected_potential": finite_projected_potential, "finite_projected_scattering_factor": finite_projected_scattering_factor, "projected_potential": projected_potential, "projected_scattering_factor": projected_scattering_factor, } return scaled_parameters[name]
[docs] class EwaldParametrization(Parametrization): _functions = {"potential": ewald.potential}
[docs] def __init__(self, width: float = 1.0): parameters = { symbol: [width, Z] for Z, symbol in enumerate(chemical_symbols[1:], 1) } super().__init__(parameters=parameters)
@property def width(self): return self.parameters["H"][0]
[docs] def scaled_parameters(self, symbol: str, name: str) -> np.ndarray: return {"potential": self.parameters[symbol]}[name]
[docs] def validate_parametrization(parametrization: str | Parametrization) -> Parametrization: named_parametrizations = { "ewald": EwaldParametrization, "lobato": LobatoParametrization, "peng": PengParametrization, "kirkland": KirklandParametrization, } if isinstance(parametrization, str): parametrization = named_parametrizations[parametrization]() return parametrization
[docs] def validate_parameters(parameters: str | dict) -> dict: if isinstance(parameters, str): if os.path.isabs(parameters): path = parameters else: path = os.path.join(_get_data_path(), parameters) with open(path, "r") as f: parameters = json.load(f) elif not isinstance(parameters, dict): raise ValueError() return parameters
# {'potential', 'projected_potential', 'charge', 'finite_projected_potential', 'scattering_factor', # 'projected_scattering_factor', 'x_ray_scattering_factor', 'finite_projected_scattering_factor'}