Search code examples
pythonpython-3.xdictionarydictionary-comprehension

Dictionary Comprehension with multiple statements


How would I simplify this using dictionary comprehension? I am working on a basic script to produce crystal supercells. Any criticism is welcome as I am a chemist and a rather poor coder. prim_cell is a user-defined dictionary of the following format:

for i in range(len(L_XYZ)):
    if Atoms == {}:
        Atoms.update({i: L_XYZ[i]})
    else:
        Atoms.update({max(Atoms, key=int)+1: L_XYZ[i]})

For context, the Atoms dictionary is initially empty

X = 2
Y = 2
Z = 2

prim_cell = {

1 : [ 'Pb', (0.000, 0.000, 0.000) ], 2 : [ 'Pb', (0.500, 0.500, 0.500) ], 
3 : [ 'O', (0.305, 0.305, 0.000) ], 4 : [ 'O', (-0.305, -0.305, 0.000) ], 
5 : [ 'O', (0.195, 0.805, 0.500) ], 6 : [ 'O', (0.805, 0.195, 0.500) ] }


Atoms = {}
Layers = {}


def Expand(prim_cell,Atoms,prim):

    Atom = [ prim[1][0]/X, prim[1][1]/Y, prim[1][2]/Z ]

    L_X = [ [ Atom[0] + i/X, Atom[1], Atom[2] ] for i in range(X) ] 

    L_XY = [ [ Coord[0], Coord[1] + i/Y, Coord[2] ] for Coord in L_X for i in range(Y) ]

    L_XYZ = [ [ prim[0], round(Coord[0],9), round(Coord[1],9), round(Coord[2] +i/Z, 9) ] for Coord in L_XY for i in range(Z) ]

    for i in range(len(L_XYZ)):

        if Atoms == { }:

            Atoms.update( { i : L_XYZ[i] } )

        else:

            Atoms.update( { max( Atoms, key=int ) + 1 : L_XYZ[i] } )

    return Atoms


def Supercell(prim_cell,Atoms):

    for i in range(len(prim_cell)):

        Atoms = Expand(prim_cell,Atoms,prim_cell[i+1])

    Layered_Atoms = sorted(Atoms.items(), key = lambda x: x[1][3])

    for i in range(len(Layered_Atoms)):

        Layers.update( { i + 1 : Layered_Atoms[i][1] } )

    for key, value in Layers.items():

        print(key,value)


Supercell(prim_cell,Atoms)

Here is an example of the output. It's basically a cartesian expansion of each point in x y and z depending upon the values of X Y and Z

1 ['Pb', 0.0, 0.0, 0.0]
2 ['O', 0.305, 0.305, 0.0]
3 ['O', -0.305, -0.305, 0.0]
4 ['Pb', 0.5, 0.5, 0.25]
5 ['O', 0.195, 0.805, 0.25]
6 ['O', 0.805, 0.195, 0.25]
7 ['Pb', 0.0, 0.0, 0.5]
8 ['O', 0.305, 0.305, 0.5]
9 ['O', -0.305, -0.305, 0.5]
10 ['Pb', 0.5, 0.5, 0.75]
11 ['O', 0.195, 0.805, 0.75]
12 ['O', 0.805, 0.195, 0.75]

Solution

  • Like this:

    Atoms = {i:v for i,v in enumerate(L_XYZ)}