'''
Run as mpiexec -n 2 python mwe.py
for example
'''

import dedalus.public as d3
import numpy as np
from mpi4py import MPI
import logging

comm = MPI.COMM_WORLD
rank = comm.rank
ncpu = comm.size
Lmax = 1
Nmax = 1
c = d3.SphericalCoordinates('phi', 'theta', 'r')
d = d3.Distributor((c,),dtype=np.float64)
b = d3.ShellBasis(c, (2*(Lmax+1),Lmax+1,Nmax+1), radii=(1,2), dealias=1, dtype=np.float64)

u = d.VectorField(c,name='u',bases=b)
print('rank = ',rank, 'u[c].shape ',u['c'].shape)

print('Try to evaluate d3.trans(d3.grad(u))')
d3.trans(d3.grad(u)).evaluate()