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]
Like this:
Atoms = {i:v for i,v in enumerate(L_XYZ)}