<commit_before># coding: utf-8
# Copyright (c) Pymatgen Development Team.
# Distributed under the terms of the MIT License.

"""
This module implements a TEM pattern calculator.
"""

from __future__ import division, print_function, unicode_literals
import json
import os
from collections import namedtuple
from fractions import Fraction
from typing import List, Dict, Tuple
import numpy as np  # type: ignore
import scipy.constants as sc  # type: ignore
import pandas as pd  # type: ignore
from pymatgen import Structure, Element  # type: ignore
from pymatgen.analysis.diffraction.core import AbstractDiffractionPatternCalculator  # type: ignore

with open(os.path.join(os.path.dirname(__file__),
                       "atomic_scattering_params.json")) as f:
    ATOMIC_SCATTERING_PARAMS = json.load(f)

# coding: utf-8
# Copyright (c) Pymatgen Development Team.
# Distributed under the terms of the MIT License.
# Credit to Dr. <NAME> for the template of the calculator

__author__ = "<NAME>, modified by <NAME>"
__copyright__ = "Copyright 2018, The Materials Project"
__version__ = "0.1"
__maintainer__ = "<NAME> respect for S.P.O"
__email__ = "<EMAIL>, <EMAIL>"
__date__ = "06/19/2019, updated 10/2019"


class TEMCalculator(AbstractDiffractionPatternCalculator):
    """
    Computes the TEM pattern of a crystal structure for multiple Laue zones.
    """

    def __init__(self, wavelength_cache: Dict[float, float] = {}, symprec: float = None, voltage: float = 200,
                 beam_direction: List[int] = [0, 0, 1], camera_length: int = 160,
                 debye_waller_factors: Dict[str, float] = None, cs: float = 1) -> None:
        """
        Initializes the TEM calculator with a given radiation.
        Args:
            symprec (float): Symmetry precision for structure refinement. If
            set to 0, no refinement is done. Otherwise, refinement is
            performed using spglib with provided precision.
            voltage (float): The wavelength is a function of the TEM microscope's
            voltage. By default, set to 200 kV. Units in kV.
            beam_direction: The direction of the electron beam fired onto the sample.
            By default, set to [0,0,1], which corresponds to the normal direction
            of the sample plane.
            camera_length (int): The distance from the sample to the projected diffraction pattern.
            By default, set to 160 cm. Units in cm.
            debye_waller_factors ({element symbol: float}): Allows the
            specification of Debye-Waller factors. Note that these
            factors are temperature dependent.
            cs (float): the chromatic aberration coefficient. set by default to 1 mm.
            later on: may want "number of iterations", "magnification", "critical value of beam",
            "twin direction" for certain materials, "sample thickness", and "excitation error s"
        """
        self.wavelength_cache = wavelength_cache
        self.symprec = symprec
        self.voltage = voltage
        self.beam_direction = beam_direction
        self.camera_length = camera_length
        self.debye_waller_factors = debye_waller_factors or {}
        self.cs = cs

    def wavelength_rel(self) -> float:
        """
        Calculates the wavelength of the electron beam with relativstic kinematic effects taken
        into account (electrons are way faster than X-rays, so you can't neglect these effects).
        Args:
            none
        Returns:
            relativisticWavelength (in meters)
        """
        if self.voltage in self.wavelength_cache:
            return self.wavelength_cache[self.voltage]
        wavelength_rel = sc.h / np.sqrt(2 * sc.m_e * sc.e * 1000 * self.voltage *
                                        (1 + (sc.e * 1000 * self.voltage) / (2 * sc.m_e * sc.c ** 2)))
        self.wavelength_cache[self.voltage] = wavelength_rel
        return wavelength_rel

    def generate_points(self, coord_left: int = -10, coord_right: int = 10) -> np.ndarray:
        """
        Generates a bunch of 3D points that span a cube.
        Args:
            coord_left (int): The minimum coordinate value.
            coord_right (int): The maximum coordinate value.
        Returns:
            Numpy 2d array
        """
        points = [0, 0, 0]
        coord_values = np.arange(coord_left, coord_right + 1)
        points[0], points[1], points[2] = np.meshgrid(coord_values, coord_values, coord_values)
        points_matrix = (np.ravel(points[i]) for i in range(0, 3))
        result = np.vstack(list(points_matrix)).transpose()
        return result

    def zone_axis_filter(self, points: np.ndarray, laue_zone: int = 0) -> List[Tuple[int, int, int]]:
        """
        Filters out all points that exist within the specified Laue zone according to the zone axis rule.
        Args:
            points (np.ndarray): The list of points to be filtered.
            laue_zone (int): The desired Laue zone.
        Returns:
            list of 3-tuples
        """
        if any(isinstance(n, tuple) for n in points):
            return points
        if len(points) == 0:
            return []
        filtered = np.where(np.dot(np.array(self.beam_direction), np.transpose(points)) == laue_zone)
        result = points[filtered]
        result_tuples = [tuple(x) for x in result.tolist()]
        return result_tuples  # type: ignore

    def get_interplanar_spacings(self, structure: Structure, points: List[Tuple[int, int, int]]) \
            -> Dict[Tuple[int, int, int], float]:
        """
        Args:
            structure (Structure): the input structure.
            points (tuple): the desired hkl indices.
        Returns:
            Dict of hkl to its interplanar spacing (float).
        """
        points_filtered = self.zone_axis_filter(points)
        if (0, 0, 0) in points_filtered:
            points_filtered.remove((0, 0, 0))
        interplanar_spacings_val = np.array(list(map(lambda x: structure.lattice.d_hkl(x), points_filtered)))
        interplanar_spacings = dict(zip(points_filtered, interplanar_spacings_val))
        return interplanar_spacings

    def bragg_angles(self, interplanar_spacings: Dict[Tuple[int, int, int], float]) \
            -> Dict[Tuple[int, int, int], float]:
        """
        Gets the Bragg angles for every hkl point passed in (where n = 1).
        Args:
            interplanar_spacings (dict): dictionary of hkl to interplanar spacing
        Returns:
            dict of hkl plane (3-tuple) to Bragg angle in radians (float)
        """
        plane = list(interplanar_spacings.keys())
        interplanar_spacings_val = np.array(list(interplanar_spacings.values()))
        bragg_angles_val = np.arcsin(self.wavelength_rel() / (2 * interplanar_spacings_val))
        bragg_angles = dict(zip(plane, bragg_angles_val))
        return bragg_angles

    def get_s2(self, bragg_angles: Dict[Tuple[int, int, int], float]) \
            -> Dict[Tuple[int, int, int], float]:
        """
        Calculates the s squared parameter (= square of sin theta over lambda) for each hkl plane.
        Args:
            bragg_angles (Dict): The bragg angles for each hkl plane.
        Returns:
            Dict of hkl plane to s2 parameter, calcs the s squared parameter (= square of sin theta over lambda).
        """
        plane = list(bragg_angles.keys())
        bragg_angles_val = np.array(list(bragg_angles.values()))
        s2_val = (np.sin(bragg_angles_val) / self.wavelength_rel()) ** 2
        s2 = dict(zip(plane, s2_val))
        return s2

    def x_ray_factors(self, structure: Structure, bragg_angles: Dict[Tuple[int, int, int], float]) \
            -> Dict[Element, Dict]:
        """
        Calculates x-ray factors, which are required to calculate atomic scattering factors. Method partially inspired
        by the equivalent process in the xrd module.
        Args:
            structure (Structure): The input structure.
            bragg_angles (Dict): Dictionary of hkl plane to Bragg angle.
        Returns:
            A Dict of atomic symbol to another dict of hkl plane to x-ray factor
        """
        x_ray_factors = {}
        s2 = self.get_s2(bragg_angles)
        atoms = structure.composition.elements
        scattering_factors_for_atom = {}
        for atom in atoms:
            coeffs = np.array(ATOMIC_SCATTERING_PARAMS[atom.symbol])
            for plane in bragg_angles:
                scattering_factor_curr = atom.Z - 41.78214 * s2[plane] * np.sum(coeffs[:, 0]
                                                                                * np.exp(-coeffs[:, 1] * s2[plane]),
                                                                                axis=None)
                scattering_factors_for_atom[plane] = scattering_factor_curr
            x_ray_factors[atom.symbol] = scattering_factors_for_atom
            scattering_factors_for_atom = {}
        return x_ray_factors

    def electron_scattering_factors(self, structure: Structure, bragg_angles: Dict[Tuple[int, int, int], float]) \
            -> Dict[Element, Dict]:
        """
        Calculates atomic scattering factors for electrons using the Mott-Bethe formula (1st order Born approximation).
        Args:
            structure (Structure): The input structure.
            bragg_angles (dict of 3-tuple to float): The Bragg angles for each hkl plane.<commit_msg>fixed type hints<commit_after># coding: utf-8
# Copyright (c) Pymatgen Development Team.
# Distributed under the terms of the MIT License.

"""
This module implements a TEM pattern calculator.
"""

from __future__ import division, print_function, unicode_literals
import json
import os
from collections import namedtuple
from fractions import Fraction
from typing import List, Dict, Tuple
import numpy as np  # type: ignore
import scipy.constants as sc  # type: ignore
import pandas as pd  # type: ignore
from pymatgen import Structure, Element  # type: ignore
from pymatgen.analysis.diffraction.core import AbstractDiffractionPatternCalculator  # type: ignore

with open(os.path.join(os.path.dirname(__file__),
                       "atomic_scattering_params.json")) as f:
    ATOMIC_SCATTERING_PARAMS = json.load(f)

# coding: utf-8
# Copyright (c) Pymatgen Development Team.
# Distributed under the terms of the MIT License.
# Credit to Dr. <NAME> for the template of the calculator

__author__ = "<NAME>, modified by <NAME>"
__copyright__ = "Copyright 2018, The Materials Project"
__version__ = "0.1"
__maintainer__ = "<NAME> respect for S.P.O"
__email__ = "<EMAIL>, <EMAIL>"
__date__ = "06/19/2019, updated 10/2019"


class TEMCalculator(AbstractDiffractionPatternCalculator):
    """
    Computes the TEM pattern of a crystal structure for multiple Laue zones.
    """

    def __init__(self, wavelength_cache: Dict[float, float] = {}, symprec: float = None, voltage: float = 200,
                 beam_direction: List[int] = [0, 0, 1], camera_length: int = 160,
                 debye_waller_factors: Dict[str, float] = None, cs: float = 1) -> None:
        """
        Initializes the TEM calculator with a given radiation.
        Args:
            symprec (float): Symmetry precision for structure refinement. If
            set to 0, no refinement is done. Otherwise, refinement is
            performed using spglib with provided precision.
            voltage (float): The wavelength is a function of the TEM microscope's
            voltage. By default, set to 200 kV. Units in kV.
            beam_direction: The direction of the electron beam fired onto the sample.
            By default, set to [0,0,1], which corresponds to the normal direction
            of the sample plane.
            camera_length (int): The distance from the sample to the projected diffraction pattern.
            By default, set to 160 cm. Units in cm.
            debye_waller_factors ({element symbol: float}): Allows the
            specification of Debye-Waller factors. Note that these
            factors are temperature dependent.
            cs (float): the chromatic aberration coefficient. set by default to 1 mm.
            later on: may want "number of iterations", "magnification", "critical value of beam",
            "twin direction" for certain materials, "sample thickness", and "excitation error s"
        """
        self.wavelength_cache = wavelength_cache
        self.symprec = symprec
        self.voltage = voltage
        self.beam_direction = beam_direction
        self.camera_length = camera_length
        self.debye_waller_factors = debye_waller_factors or {}
        self.cs = cs

    def wavelength_rel(self) -> float:
        """
        Calculates the wavelength of the electron beam with relativstic kinematic effects taken
        into account (electrons are way faster than X-rays, so you can't neglect these effects).
        Args:
            none
        Returns:
            relativisticWavelength (in meters)
        """
        if self.voltage in self.wavelength_cache:
            return self.wavelength_cache[self.voltage]
        wavelength_rel = sc.h / np.sqrt(2 * sc.m_e * sc.e * 1000 * self.voltage *
                                        (1 + (sc.e * 1000 * self.voltage) / (2 * sc.m_e * sc.c ** 2)))
        self.wavelength_cache[self.voltage] = wavelength_rel
        return wavelength_rel

    def generate_points(self, coord_left: int = -10, coord_right: int = 10) -> np.ndarray:
        """
        Generates a bunch of 3D points that span a cube.
        Args:
            coord_left (int): The minimum coordinate value.
            coord_right (int): The maximum coordinate value.
        Returns:
            Numpy 2d array
        """
        points = [0, 0, 0]
        coord_values = np.arange(coord_left, coord_right + 1)
        points[0], points[1], points[2] = np.meshgrid(coord_values, coord_values, coord_values)
        points_matrix = (np.ravel(points[i]) for i in range(0, 3))
        result = np.vstack(list(points_matrix)).transpose()
        return result

    def zone_axis_filter(self, points: np.ndarray, laue_zone: int = 0) -> List[Tuple[int, int, int]]:
        """
        Filters out all points that exist within the specified Laue zone according to the zone axis rule.
        Args:
            points (np.ndarray): The list of points to be filtered.
            laue_zone (int): The desired Laue zone.
        Returns:
            list of 3-tuples
        """
        if any(isinstance(n, tuple) for n in points):
            return points
        if len(points) == 0:
            return []
        filtered = np.where(np.dot(np.array(self.beam_direction), np.transpose(points)) == laue_zone)
        result = points[filtered]
        result_tuples = [tuple(x) for x in result.tolist()]
        return result_tuples  # type: ignore

    def get_interplanar_spacings(self, structure: Structure, points: List[Tuple[int, int, int]]) \
            -> Dict[Tuple[int, int, int], float]:
        """
        Args:
            structure (Structure): the input structure.
            points (tuple): the desired hkl indices.
        Returns:
            Dict of hkl to its interplanar spacing (float).
        """
        points_filtered = self.zone_axis_filter(points)
        if (0, 0, 0) in points_filtered:
            points_filtered.remove((0, 0, 0))
        interplanar_spacings_val = np.array(list(map(lambda x: structure.lattice.d_hkl(x), points_filtered)))
        interplanar_spacings = dict(zip(points_filtered, interplanar_spacings_val))
        return interplanar_spacings

    def bragg_angles(self, interplanar_spacings: Dict[Tuple[int, int, int], float]) \
            -> Dict[Tuple[int, int, int], float]:
        """
        Gets the Bragg angles for every hkl point passed in (where n = 1).
        Args:
            interplanar_spacings (dict): dictionary of hkl to interplanar spacing
        Returns:
            dict of hkl plane (3-tuple) to Bragg angle in radians (float)
        """
        plane = list(interplanar_spacings.keys())
        interplanar_spacings_val = np.array(list(interplanar_spacings.values()))
        bragg_angles_val = np.arcsin(self.wavelength_rel() / (2 * interplanar_spacings_val))
        bragg_angles = dict(zip(plane, bragg_angles_val))
        return bragg_angles

    def get_s2(self, bragg_angles: Dict[Tuple[int, int, int], float]) \
            -> Dict[Tuple[int, int, int], float]:
        """
        Calculates the s squared parameter (= square of sin theta over lambda) for each hkl plane.
        Args:
            bragg_angles (Dict): The bragg angles for each hkl plane.
        Returns:
            Dict of hkl plane to s2 parameter, calcs the s squared parameter (= square of sin theta over lambda).
        """
        plane = list(bragg_angles.keys())
        bragg_angles_val = np.array(list(bragg_angles.values()))
        s2_val = (np.sin(bragg_angles_val) / self.wavelength_rel()) ** 2
        s2 = dict(zip(plane, s2_val))
        return s2

    def x_ray_factors(self, structure: Structure, bragg_angles: Dict[Tuple[int, int, int], float]) \
            -> Dict[str, Dict[Tuple[int, int, int], float]]:
        """
        Calculates x-ray factors, which are required to calculate atomic scattering factors. Method partially inspired
        by the equivalent process in the xrd module.
        Args:
            structure (Structure): The input structure.
            bragg_angles (Dict): Dictionary of hkl plane to Bragg angle.
        Returns:
            A Dict of atomic symbol to another dict of hkl plane to x-ray factor
        """
        x_ray_factors = {}
        s2 = self.get_s2(bragg_angles)
        atoms = structure.composition.elements
        scattering_factors_for_atom = {}
        for atom in atoms:
            coeffs = np.array(ATOMIC_SCATTERING_PARAMS[atom.symbol])
            for plane in bragg_angles:
                scattering_factor_curr = atom.Z - 41.78214 * s2[plane] * np.sum(coeffs[:, 0]
                                                                                * np.exp(-coeffs[:, 1] * s2[plane]),
                                                                                axis=None)
                scattering_factors_for_atom[plane] = scattering_factor_curr
            x_ray_factors[atom.symbol] = scattering_factors_for_atom
            scattering_factors_for_atom = {}
        return x_ray_factors

    def electron_scattering_factors(self, structure: Structure, bragg_angles: Dict[Tuple[int, int, int], float]) \
            -> Dict[str, Dict[Tuple[int, int, int], float]]:
        """
        Calculates atomic scattering factors for electrons using the Mott-Bethe formula (1st order Born approximation).
        Args:
            structure (Structure): The input structure.
            bragg_angles (dict of 3-tuple to float): The Bragg angles for each hkl plane.
{'fraction_non_alphanumeric': 0.0888535875033095, 'fraction_numerical': 0.011014032300767805, 'mean_word_length': 2.3743076648204395, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 33, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 20, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '507663', 'n_tokens_mistral': 5527, 'n_tokens_neox': 5217, 'n_words': 1729}