<commit_before># -*- coding: utf-8 -*-
"""
==============================================================================
Graphene structure generator (:mod:`sknano.generators._graphene_generator`)
==============================================================================

.. currentmodule:: sknano.generators._graphene_generator

"""
from __future__ import absolute_import, division, print_function
__docformat__ = 'restructuredtext en'

import numpy as np

from sknano.structures import Graphene
from ._base import GeneratorAtom as Atom, GeneratorAtoms as Atoms, \
    GeneratorMixin

__all__ = ['GrapheneGenerator']


class GrapheneGenerator(Graphene, GeneratorMixin):
    """Class for generating `N`-layer graphene nanostructures.

    Parameters
    ----------
    length : float
        Length of graphene sheet in **nanometers**
    width : float
        Width of graphene sheet in **nanometers**
    edge : {'AC', 'armchair', 'ZZ', 'zigzag'}, optional
        **A**\ rm\ **C**\ hair or **Z**\ ig\ **Z**\ ag edge along
        the `length` of the sheet.
    element1, element2 : {str, int}, optional
        Element symbol or atomic number of basis
        :class:`~sknano.core.atoms.Atom` 1 and 2
    bond : float, optional
        bond length between nearest-neighbor atoms in **Angstroms**.
    nlayers : int, optional
        Number of graphene layers.
    layer_spacing : float, optional
        Distance between layers in **Angstroms**.
    stacking_order : {'AA', 'AB'}, optional
        Stacking order of graphene layers
    autogen : bool, optional
        automatically generate unit cell and full structure
    verbose : bool, optional
        verbose output

    Notes
    -----
    The `GrapheneGenerator` class and its subclasses generate graphene with
    either an armchair or zigzag edge using a 4-atom conventional unit cell.
    If you want to generate graphene as an *unrolled nanotube*, see the
    :class:`~sknano.generators.UnrolledNanotubeGenerator` class.

    .. seealso:: :class:`~sknano.generators.UnrolledNanotubeGenerator`

    Examples
    --------

    Start an interactive python or ipython session, then import the
    `GrapheneGenerator` class.

    >>> from sknano.generators import GrapheneGenerator

    Now generate a **20 nm x 1 nm** armchair edge graphene nano-ribbon.

    >>> ACG = GrapheneGenerator(length=20, width=1, edge='AC')

    Save structure data in `xyz` format:

    >>> ACG.save_data()

    The rendered structure look like:

    .. image:: /images/20nmx1nm_AC_edge.png

    Now let's generate a **20 nm x 1 nm** zigzag edge graphene nano-ribbon.

    >>> ZZG = GrapheneGenerator(length=20, width=1, edge='ZZ')
    >>> ZZG.save_data()

    The rendered structure looks like:

    .. image:: /images/20nmx1nm_ZZ_edge.png

    Now generate **25 nm x 5 nm**, `armchair` edge,
    5 layer, `AB`-stacked graphene.

    >>> ACG_5layers = GrapheneGenerator(length=25, width=5, edge='AC',
    ...                                 nlayers=5)
    >>> ACG_5layers.save_data()

    The rendered structure looks like:

    .. image:: /images/25nmx5nm_5layer_AC_graphene.png

    Now generate single layer, **10 nm x 10 nm** sheet of BN Graphene.

    >>> BN_graphene = GrapheneGenerator(length=10, width=10, edge='AC',
    ...                                 element1='B', element2='N')
    >>> BN_graphene.save_data()

    The rendered structure looks like:

    .. image:: /images/10nmx10nm_single_layer_BN_graphene.png

    Now, just because we can, generate a **5 nm x 5 nm** sheet of
    Uranium-Einsteinium Graphene.

    >>> UEs_graphene = GrapheneGenerator(width=5, length=5, edge='zigzag',
    ...                                  element1='U', element2='Es')
    >>> UEs_graphene.save_data()

    The rendered structure looks like:

    .. image:: /images/5nmx5nm_single_layer_UEs_graphene.png

    """

    def __init__(self, autogen=True, **kwargs):

        super(GrapheneGenerator, self).__init__(**kwargs)

        if autogen:
            self.generate_unit_cell()
            self.generate_structure_data()

    def generate_unit_cell(self):
        """Generate the graphene unit cell.

        Called automatically if `autogen` is `True`.

        """
        bond = self._bond
        edge = self._edge
        e1 = self._element1
        e2 = self._element2

        self._unit_cell = Atoms()
        # Set up 4 atom basis
        # Leave atom 1 at the origin
        atom1 = Atom(element=e1)
        atom2 = Atom(element=e2)
        atom3 = Atom(element=e1)
        atom4 = Atom(element=e2)
        if edge == 'AC':
            # Move atom 2 to 2nd basis position
            atom2.x = -np.sqrt(3) / 2 * bond
            atom2.y = bond / 2

            # Move atom 3 from atom 1 along a1 primitive vector
            atom3.x = -np.sqrt(3) / 2 * bond
            atom3.y = 3 / 2 * bond

            # Move atom 4 from atom 2 along a2 primitive vector
            atom4.y = 2 * bond
        else:
            # Move atom 2 to 2nd basis position
            atom2.x = bond

            # Move atom 3 from atom 1 along a1 primitive vector
            atom3.x = 3 / 2 * bond
            atom3.y = np.sqrt(3) / 2 * bond

            # Move atom 4 from atom 2 along a2 primitive vector
            atom4.x = -bond / 2
            atom4.y = np.sqrt(3) / 2 * bond

        for atom in (atom1, atom2, atom3, atom4):
            self._unit_cell.append(atom)
        self._Natoms = self._unit_cell.Natoms

    def generate_structure_data(self):
        """Generate the full structure coordinates."""
        self._structure_atoms = Atoms()
        for nlayer in xrange(self._nlayers):
            layer = Atoms()
            for nx in xrange(self._Nx):
                for ny in xrange(self._Ny):
                    dr = np.array([nx * self._cell.x,
                                   ny * self._cell.y,
                                   nlayer * self._layer_spacing])
                    for atom in self._unit_cell:
                        layer_atom = Atom(element=atom.symbol)
                        layer_atom.r = atom.r + dr
                        layer.append(layer_atom)

            if self._Natoms_per_layer is None:
                self._Natoms_per_layer = layer.Natoms

            # translate layer to put its center of mass at the origin
            layer.center_CM(r_indices=[0, 1])
            if (nlayer % 2) != 0:
                layer.translate(self._layer_shift)

            self._structure_atoms.extend(layer.atoms)

    def save_data(self, fname=None, outpath=None, structure_format=None,
                  rotation_angle=-90, rot_axis='x', deg2rad=True,
                  center_CM=True, **kwargs):
        """Save structure data.

        See :meth:`~sknano.generators.StructureGenerator.save_data` method
        for documentation.<commit_msg>Updating translation of layer coordinates.<commit_after># -*- coding: utf-8 -*-
"""
==============================================================================
Graphene structure generator (:mod:`sknano.generators._graphene_generator`)
==============================================================================

.. currentmodule:: sknano.generators._graphene_generator

"""
from __future__ import absolute_import, division, print_function
__docformat__ = 'restructuredtext en'

import numpy as np

from sknano.structures import Graphene
from ._base import GeneratorAtom as Atom, GeneratorAtoms as Atoms, \
    GeneratorMixin

__all__ = ['GrapheneGenerator']


class GrapheneGenerator(Graphene, GeneratorMixin):
    """Class for generating `N`-layer graphene nanostructures.

    Parameters
    ----------
    length : float
        Length of graphene sheet in **nanometers**
    width : float
        Width of graphene sheet in **nanometers**
    edge : {'AC', 'armchair', 'ZZ', 'zigzag'}, optional
        **A**\ rm\ **C**\ hair or **Z**\ ig\ **Z**\ ag edge along
        the `length` of the sheet.
    element1, element2 : {str, int}, optional
        Element symbol or atomic number of basis
        :class:`~sknano.core.atoms.Atom` 1 and 2
    bond : float, optional
        bond length between nearest-neighbor atoms in **Angstroms**.
    nlayers : int, optional
        Number of graphene layers.
    layer_spacing : float, optional
        Distance between layers in **Angstroms**.
    stacking_order : {'AA', 'AB'}, optional
        Stacking order of graphene layers
    autogen : bool, optional
        automatically generate unit cell and full structure
    verbose : bool, optional
        verbose output

    Notes
    -----
    The `GrapheneGenerator` class and its subclasses generate graphene with
    either an armchair or zigzag edge using a 4-atom conventional unit cell.
    If you want to generate graphene as an *unrolled nanotube*, see the
    :class:`~sknano.generators.UnrolledNanotubeGenerator` class.

    .. seealso:: :class:`~sknano.generators.UnrolledNanotubeGenerator`

    Examples
    --------

    Start an interactive python or ipython session, then import the
    `GrapheneGenerator` class.

    >>> from sknano.generators import GrapheneGenerator

    Now generate a **20 nm x 1 nm** armchair edge graphene nano-ribbon.

    >>> ACG = GrapheneGenerator(length=20, width=1, edge='AC')

    Save structure data in `xyz` format:

    >>> ACG.save_data()

    The rendered structure look like:

    .. image:: /images/20nmx1nm_AC_edge.png

    Now let's generate a **20 nm x 1 nm** zigzag edge graphene nano-ribbon.

    >>> ZZG = GrapheneGenerator(length=20, width=1, edge='ZZ')
    >>> ZZG.save_data()

    The rendered structure looks like:

    .. image:: /images/20nmx1nm_ZZ_edge.png

    Now generate **25 nm x 5 nm**, `armchair` edge,
    5 layer, `AB`-stacked graphene.

    >>> ACG_5layers = GrapheneGenerator(length=25, width=5, edge='AC',
    ...                                 nlayers=5)
    >>> ACG_5layers.save_data()

    The rendered structure looks like:

    .. image:: /images/25nmx5nm_5layer_AC_graphene.png

    Now generate single layer, **10 nm x 10 nm** sheet of BN Graphene.

    >>> BN_graphene = GrapheneGenerator(length=10, width=10, edge='AC',
    ...                                 element1='B', element2='N')
    >>> BN_graphene.save_data()

    The rendered structure looks like:

    .. image:: /images/10nmx10nm_single_layer_BN_graphene.png

    Now, just because we can, generate a **5 nm x 5 nm** sheet of
    Uranium-Einsteinium Graphene.

    >>> UEs_graphene = GrapheneGenerator(width=5, length=5, edge='zigzag',
    ...                                  element1='U', element2='Es')
    >>> UEs_graphene.save_data()

    The rendered structure looks like:

    .. image:: /images/5nmx5nm_single_layer_UEs_graphene.png

    """

    def __init__(self, autogen=True, **kwargs):

        super(GrapheneGenerator, self).__init__(**kwargs)

        if autogen:
            self.generate_unit_cell()
            self.generate_structure_data()

    def generate_unit_cell(self):
        """Generate the graphene unit cell.

        Called automatically if `autogen` is `True`.

        """
        bond = self._bond
        edge = self._edge
        e1 = self._element1
        e2 = self._element2

        self._unit_cell = Atoms()
        # Set up 4 atom basis
        # Leave atom 1 at the origin
        atom1 = Atom(element=e1)
        atom2 = Atom(element=e2)
        atom3 = Atom(element=e1)
        atom4 = Atom(element=e2)
        if edge == 'AC':
            # Move atom 2 to 2nd basis position
            atom2.x = -np.sqrt(3) / 2 * bond
            atom2.y = bond / 2

            # Move atom 3 from atom 1 along a1 primitive vector
            atom3.x = -np.sqrt(3) / 2 * bond
            atom3.y = 3 / 2 * bond

            # Move atom 4 from atom 2 along a2 primitive vector
            atom4.y = 2 * bond
        else:
            # Move atom 2 to 2nd basis position
            atom2.x = bond

            # Move atom 3 from atom 1 along a1 primitive vector
            atom3.x = 3 / 2 * bond
            atom3.y = np.sqrt(3) / 2 * bond

            # Move atom 4 from atom 2 along a2 primitive vector
            atom4.x = -bond / 2
            atom4.y = np.sqrt(3) / 2 * bond

        for atom in (atom1, atom2, atom3, atom4):
            self._unit_cell.append(atom)
        self._Natoms = self._unit_cell.Natoms

    def generate_structure_data(self):
        """Generate the full structure coordinates."""
        self._structure_atoms = Atoms()
        for nlayer in xrange(self._nlayers):
            layer = Atoms()
            for nx in xrange(self._Nx):
                for ny in xrange(self._Ny):
                    dr = np.array([nx * self._cell.x,
                                   ny * self._cell.y,
                                   nlayer * self._layer_spacing])
                    for atom in self._unit_cell:
                        layer_atom = Atom(element=atom.symbol)
                        layer_atom.r = atom.r + dr
                        layer.append(layer_atom)

            if self._Natoms_per_layer is None:
                self._Natoms_per_layer = layer.Natoms

            # translate layer to put its center of mass at the origin
            dr = layer.CM
            dr.z = 0
            layer.translate(dr)
            if (nlayer % 2) != 0:
                layer.translate(self._layer_shift)

            self._structure_atoms.extend(layer)

    def save_data(self, fname=None, outpath=None, structure_format=None,
                  rotation_angle=-90, rot_axis='x', deg2rad=True,
                  center_CM=True, **kwargs):
        """Save structure data.

        See :meth:`~sknano.generators.StructureGenerator.save_data` method
        for documentation.
{'fraction_non_alphanumeric': 0.12255902069367532, 'fraction_numerical': 0.018289128533955114, 'mean_word_length': 2.162442396313364, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 69, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6633530', 'n_tokens_mistral': 4322, 'n_tokens_neox': 3877, 'n_words': 1170}