Create a 3D lattice of points for some lattice parameters

lattice_points.py

# -*- coding: utf-8 -*-
"""
Create a 3D lattice of points for some lattice parameters

@author: Nicolas Guarin-Zapata
"""
from __future__ import division, print_function
import numpy as np
from numpy import sin, cos, sqrt



def gen_lattice(repetitions=(5, 5, 5), basis=None, lattice_cons=None,
                origin=(0, 0, 0)):
    """
    Generate a set of coordinates arranged in a lattice

    Parameters
    ----------
    repetitions : tuple, int
        Repetitions in each direction.
    basis : tuple, ndarray
        Tuple of basis vectors.
    lattice_cons : tuple, float
        Lattice constants.
    origin : tuple, float
        Origin for the first point.

    Returns
    -------
    pts : ndarray, (nx * ny * nz, 3)
        Array with coordinates for the points.

    """
    nx, ny, nz = repetitions
    origin = np.asarray(origin)
    if basis == None:
        v1 = np.array([1, 0, 0])
        v2 = np.array([0, 1, 0])
        v3 = np.array([0, 0, 1])
    else:
        v1 = np.asarray(basis[0])
        v1 = np.asarray(basis[1])
        v1 = np.asarray(basis[2])
    if lattice_cons== None:
        a, b, c = 1, 1, 1
    else:
        a, b, c, = lattice_cons

    pts = np.zeros((nx * ny * nz, 3))
    cont = 0
    for contx in range(nx):
        for conty in range(ny):
            for contz in range(nz):
                pts[cont, :] = origin + contx * a * v1 +\
                                conty * b * v2 + contz * c * v3
                cont += 1
    return pts


def trans_array_pts(coords, delta):
    """Translate an array of coordinates (coords) by delta"""
    return coords + delta


def rotate_array_pts(coords, angle):
    """
    Rotate points around its center mass an angle
    in the z axis
    """
    mean = np.mean(coords, axis=0)
    mat = np.array([
            [cos(angle), -sin(angle), 0],
            [sin(angle), cos(angle), 0],
            [0, 0, 1]])
    return (coords - mean).dot(mat) + mean


# Example
if __name__ == "__main__":
    import matplotlib.pyplot as plt
    from mpl_toolkits.mplot3d import Axes3D

    Axes3D


    # Lattice
    pts = gen_lattice()
    max_coords = pts.min(axis=0)
    max_coords = pts.max(axis=0)
    mean_coords = pts.mean(axis=0)
    offset = 1

    # Molecule
    molec = 0.25*np.array([
            [0, 0, 0],
            [sqrt(8/9), 0, 4/3],
            [-sqrt(2/9), sqrt(2/3), 4/3],
            [-sqrt(2/9), -sqrt(2/3), 4/3],])
    size_atoms = [80, 60, 60, 60]
    color_atoms = ["blue", "red", "red", "red"]
    delta = [mean_coords[0], mean_coords[1], max_coords[-1] + offset]
    molec = trans_array_pts(molec, delta)
    angle = np.pi/3
    molec = rotate_array_pts(molec, angle)

    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    ax.scatter(pts[:, 0], pts[:, 1], pts[:, 2], s=100, c="#d1937b", alpha=0.6)
    ax.scatter(molec[:, 0], molec[:, 1], molec[:, 2],
               s=size_atoms, c=color_atoms, alpha=1.0)
    plt.show()