# Import the basic libraries

# ASE system
import ase
from ase import Atom, Atoms
from ase import io
from ase.lattice.spacegroup import crystal

# Spacegroup/symmetry library
from pyspglib import spglib

# The qe-util package
from qeutil import QuantumEspresso

# iPython utility function
from IPython.core.display import Image

# Configure qe-util for local execution of the Quantum Espresso on four processors
QuantumEspresso.pw_cmd='mpiexec -n 4 pw.x < %(infile)s > %(outfile)s'

a=4.3596                                         # Lattice constant in Angstrom
cryst = crystal(['Si', 'C'],                     # Atoms in the crystal
                [(0, 0, 0), (0.25, 0.25, 0.25)], # Atomic positions (fractional coordinates)
                spacegroup=216,                  # International number of the spacegroup of the crystal
                cellpar=[a, a, a, 90, 90, 90])   # Unit cell (a, b, c, alpha, beta, gamma) in Angstrom, Degrees

# Write the image to disk file
ase.io.write('crystal.png',       # The file where the picture get stored
             cryst,               # The object holding the crystal definition
             format='png',        # Format of the file
             show_unit_cell=2,    # Draw the unit cell boundaries
             rotation='115y,15x', # Rotate the scene by 115deg around Y axis and 15deg around X axis
             scale=35)            # Scale of the picture

# Display the image
Image(filename='crystal.png')

# Find a primitive cell of the structure
puc=spglib.find_primitive(cryst)

# Parameters of the unit cell
print 'Primitive unit cell (Angstrom, carthesian coordinates):\n', puc[0]
print '\nAtomic positions (fractional, crystalographic):\n', puc[1]

# Create the primitive cell crystal
cryst_prim=Atoms(                         # Create a generic atomic structure
                 cell=puc[0],             # Unit cell (Angstrom)
                 scaled_positions=puc[1], # Atomic positions (fractional)
                 numbers=puc[2],          # Atomic numbers
                 pbc=True)                # Use Periodic Boundary Conditions

# Check if the symmetry is still the same
print spglib.get_spacegroup(cryst_prim)

# define the new unit cell
new_uc=a*array([[1,0,1],[1,1,0],[0,1,1]])/2

print "New primitive unit cell (A)\n", new_uc

# Create a new crystal with a new primitive unit cell
new_crystal=Atoms(cell=new_uc,                               # Define new unit cell
                    positions=cryst_prim.get_positions(),    # Put atoms in cartesian positions from the cryst_prim
                    numbers=cryst_prim.get_atomic_numbers(),         # Set atomic numbers
                    pbc=True)                                # Make the structure periodic

print "Atomic positions (fractional)\n", new_crystal.get_scaled_positions()

# Take a look, same orientation as before
ase.io.write('crystal-prim.png',new_crystal,format='png',show_unit_cell=2, rotation='115y,15x', scale=40)
Image(filename='crystal-prim.png')