import dedalus.public as d3
import matplotlib.pyplot as plt
import numpy as np
from scipy.special import sph_harm_y
Nθ, Nφ = 20, 16 # has to be this combination??
# Initialize coords/dist/basis/grid
coords = d3.CartesianCoordinates('θ', 'φ')
dist = d3.Distributor(coords, dtype=complex)
θbasis = d3.Chebyshev(coords['θ'], size=Nθ, bounds=(0, np.pi))
φbasis = d3.ComplexFourier(coords['φ'], size=Nφ, bounds=(0, 2*np.pi))
bases = (θbasis, φbasis)
θgrid, φgrid = dist.local_grids(θbasis, φbasis)
# Get 2D arrays of coordinates
θgrid_s, φgrid_s = dist.local_grids(θbasis, φbasis)
θgrid_s1d, φgrid_s1d = θgrid_s.T[0], φgrid_s[0] # extract 1D arrays out of these
θgrid_s2d, φgrid_s2d = np.meshgrid(θgrid_s1d, φgrid_s1d)
θgrid_s2d, φgrid_s2d = θgrid_s2d.T, φgrid_s2d.T
# Initialize field and try to set it with a high-l/|m| spherical harmonic
F = dist.Field(bases=(θbasis, φbasis), dtype=complex, name='F')
F['g'] = sph_harm_y(8, -8, θgrid_s2d, φgrid_s2d) # this function choice exaggerates the issue
plt.close() # plot once
plt.imshow(np.abs(F['g']))
plt.show()
plt.close() # plot again
plt.imshow(np.abs(F['g']))
plt.show()
plt.close() # plot again
plt.imshow(np.abs(F['g']))
plt.show()
print(F['c']) # print the spectral coefficients
plt.close() # suddenly the grid-space field is different
plt.imshow(np.abs(F['g']))
plt.show()