Esempio n. 1
0
    def _compute_Sq_with_reduced_lmax( self, new_lmax):
        """
        model: SphericalModel object
        pr: PhaseRetriver object
        """
        new_sh_obj = shtns.sht(new_lmax)

        old_sh_obj = self.sh
        
        new_S_q = np.zeros_like( self.S_q )
        
        old_all_slm = self.all_slm
        
        new_sh_obj.set_grid( self.n_theta, self.n_phi )
        n_q = self.n_q
        new_all_slm = np.zeros((self.n_q, int((new_lmax+2)*(new_lmax+1)/2) ) 
            , dtype=np.complex128)

        for qid in range(n_q):
            new_slm = new_sh_obj.spec_array()
            
            for lid in range(new_lmax+1):
                for mid in range(lid+1):
                    new_slm[new_sh_obj.idx(lid, mid)] = old_all_slm[qid][old_sh_obj.idx(lid,mid)]
            new_S_q[qid] = new_sh_obj.synth(new_slm)
            new_all_slm[qid, :] = new_slm

        return new_S_q, new_all_slm, new_sh_obj
Esempio n. 2
0
def get_Sq_with_reduced_lmax(new_lmax, model=None, pr=None):
    """
	model: SphericalModel object
	pr: PhaseRetriver object
	"""
    new_sh_obj = shtns.sht(new_lmax)

    if model is None:
        old_sh_obj = pr.sh
        new_S_q = np.zeros_like(pr.I_guess)
        old_all_slm = pr.all_slm_guess
        new_sh_obj.set_grid(pr.n_theta, pr.n_phi)
        n_q = pr.n_q
    else:
        old_sh_obj = model.sh
        new_S_q = np.zeros_like(model.S_q)
        old_all_slm = model.all_slm
        new_sh_obj.set_grid(model.n_theta, model.n_phi)
        n_q = model.n_q

    for qid in range(n_q):
        new_slm = new_sh_obj.spec_array()
        for lid in range(new_lmax + 1):
            for mid in range(lid + 1):
                new_slm[new_sh_obj.idx(lid,
                                       mid)] = old_all_slm[qid][old_sh_obj.idx(
                                           lid, mid)]
        new_S_q[qid] = new_sh_obj.synth(new_slm)

    return new_S_q
Esempio n. 3
0
def power_spectrum(Br, nphi, ntheta, l_trunc, r=ro):
    '''
    Obtain power spectrum of core surface field to degree l_trunc (default at the CMB)
    '''
    # SH transform from spatial to spectral space
    m_max = l_trunc
    # by default 'flag = sht_quick_init' uses gaussian grid
    sh = shtns.sht(l_trunc, m_max)
    nlat, nlon = sh.set_grid(nphi=nphi, nlat=ntheta)
    vr = Br.T.astype(
        'float64')  # NOTE: array has to be dtype='float64' and not 'float32'
    clm = sh.analys(vr)  # spatial to spectral
    # Construct Gauss coefficients
    glm = np.zeros(sh.nlm)
    hlm = np.zeros(sh.nlm)
    for l in range(l_trunc + 1):
        fac = 1 / (l + 1
                   )  # cmb_radius/(l+1) # from Br potential to full potential
        for m in range(l + 1):
            glm[sh.idx(l, m)] = fac * np.real(clm[sh.idx(l, m)])
            hlm[sh.idx(l, m)] = fac * np.imag(clm[sh.idx(l, m)])
    # Power spectrum
    degrees = np.arange(1, l_trunc + 1, 1)
    spectrum = np.zeros(l_trunc + 1)
    for l in range(l_trunc + 1):
        fac = (l + 1) * (earth_radius / cmb_radius)**(
            2 * l + 4)  # TODO: check factors using time average
        sum_m = 0
        for m in range(l + 1):
            sum_m += glm[sh.idx(l, m)]**2 + hlm[sh.idx(l, m)]**2
        spectrum[l] = fac * sum_m
    spectrum = spectrum[1:]  # ignore g00

    return degrees, spectrum, clm, glm, hlm
Esempio n. 4
0
    def setup(self, file_info, anti_aliasing=False):
        import shtns
        import numpy as np

        if file_info['modes_m_max'] != file_info['modes_m_max']:
            raise Exception("Only num_lon == num_lat supported")

        ntrunc = file_info['modes_n_max']
        self._shtns = shtns.sht(ntrunc, ntrunc, 1,
                                shtns.sht_orthonormal + shtns.SHT_NO_CS_PHASE)

        nlons = (ntrunc + 1) * 2
        nlats = (ntrunc + 1)

        if anti_aliasing:
            if nlons & 1:
                raise Exception(
                    "Only even numbers of longitudinal coordinates allowed for anti-aliasing"
                )
            if nlats & 1:
                raise Exception(
                    "Only even numbers of latitudinal coordinates allowed for anti-aliasing"
                )

            print("Anti-aliasing:")
            print(" + old lon/lat: ", nlons, nlats)

            nlons += nlons // 2
            nlats += nlats // 2

            print(" + new lon/lat: ", nlons, nlats)

        if file_info['grid_type'] == 'GAUSSIAN':
            #self._shtns.set_grid(nlats,nlons,shtns.sht_gauss_fly|shtns.SHT_PHI_CONTIGUOUS, 1.e-10)
            self._shtns.set_grid(
                nlats, nlons, shtns.sht_quick_init | shtns.SHT_PHI_CONTIGUOUS,
                0)

        elif file_info['grid_type'] == 'REGULAR':
            #self._shtns.set_grid(nlats,nlons,shtns.sht_reg_dct|shtns.SHT_PHI_CONTIGUOUS, 1.e-10)
            self._shtns.set_grid(nlats, nlons,
                                 shtns.sht_reg_dct | shtns.SHT_PHI_CONTIGUOUS,
                                 0)

        else:
            raise Exception("Grid type '" + file_info['grid_type'] +
                            "' not supported!")

        self.lats = np.arcsin(self._shtns.cos_theta)
        self.lons = (2. * np.pi / nlons) * np.arange(nlons)
        self.nlons = nlons
        self.nlats = nlats
        self.ntrunc = ntrunc
        self.nlm = self._shtns.nlm
        self.degree = self._shtns.l
        self.lap = -self.degree * (self.degree + 1.0).astype(np.complex)
        self.invlap = np.zeros(self.lap.shape, self.lap.dtype)
        self.invlap[1:] = 1. / self.lap[1:]
        self.lap = self.lap / self.rsphere**2
        self.invlap = self.invlap * self.rsphere**2
Esempio n. 5
0
 def __init__(self,nlons,nlats,ntrunc,rsphere,gridtype='gaussian'):
     """
     initialize
     nlons:  number of longitudes
     nlats:  number of latitudes
     """
     self._shtns = shtns.sht(ntrunc, ntrunc, 1, \
                             shtns.sht_orthonormal+shtns.SHT_NO_CS_PHASE)
     if gridtype == 'gaussian':
         #self._shtns.set_grid(nlats,nlons,shtns.sht_gauss_fly|shtns.SHT_PHI_CONTIGUOUS,1.e-10)
         self._shtns.set_grid(nlats,nlons,shtns.sht_quick_init|shtns.SHT_PHI_CONTIGUOUS,1.e-10)
     elif gridtype == 'regular':
         self._shtns.set_grid(nlats,nlons,shtns.sht_reg_dct|shtns.SHT_PHI_CONTIGUOUS,1.e-10)
     self.lats = np.arcsin(self._shtns.cos_theta)
     self.lons = (2.*np.pi/nlons)*np.arange(nlons)
     self.nlons = nlons
     self.nlats = nlats
     self.ntrunc = ntrunc
     self.nlm = self._shtns.nlm
     self.degree = self._shtns.l
     self.lap = -self.degree*(self.degree+1.0).astype(np.complex)
     self.invlap = np.zeros(self.lap.shape, self.lap.dtype)
     self.invlap[1:] = 1./self.lap[1:]
     self.rsphere = rsphere
     self.lap = self.lap/rsphere**2
     self.invlap = self.invlap*rsphere**2
Esempio n. 6
0
def write_model_h5(model, lmax, fname, varpath):
    # assumes the model is sampled on a regular grid where the poles are not
    # included

    # prepare the transform
    mmax = lmax
    sh = shtns.sht(lmax, mmax)

    npts_ph, npts_th = model.shape

    grid_typ = shtns.sht_reg_fast | shtns.SHT_THETA_CONTIGUOUS
    grid_typ = grid_typ | shtns.SHT_LOAD_SAVE_CFG

    polar_opt_threshold = 1.0e-10
    sh.set_grid(npts_th,
                npts_ph,
                flags=grid_typ,
                polar_opt=polar_opt_threshold)

    # apply transform
    ylm = sh.analys(model)

    with File(fname, 'r+') as f:
        path_l = os.path.split(varpath)[0] + '/l'
        if not path_l in f:
            f.create_dataset(path_l, data=sh.l)
        path_m = os.path.split(varpath)[0] + '/m'
        if not path_m in f:
            f.create_dataset(path_m, data=sh.m)
        f.create_dataset(varpath, data=ylm)

    return
Esempio n. 7
0
    def __init__(self, nlons, nlats, ntrunc, rsphere, gridtype='gaussian'):
        """initialize
        nlons:  number of longitudes
        nlats:  number of latitudes"""
        self._shtns = shtns.sht(ntrunc, ntrunc, 1,
                                shtns.sht_orthonormal + shtns.SHT_NO_CS_PHASE)
        if gridtype == 'gaussian':
            self._shtns.set_grid(
                nlats, nlons, shtns.sht_quick_init | shtns.SHT_PHI_CONTIGUOUS,
                0)
        elif gridtype == 'regular':
            self._shtns.set_grid(nlats, nlons,
                                 shtns.sht_reg_dct | shtns.SHT_PHI_CONTIGUOUS,
                                 0)
        self.lats = np.arcsin(self._shtns.cos_theta)
        self.lons = (2. * np.pi / nlons) * np.arange(nlons)
        self.nlons = nlons
        self.nlats = nlats
        self.ntrunc = ntrunc
        self.nlm = self._shtns.nlm
        self.degree = self._shtns.l
        self.lap = -self.degree * (self.degree + 1.0).astype(complex)
        self.invlap = np.zeros(self.lap.shape, self.lap.dtype)
        self.invlap[1:] = 1. / self.lap[1:]
        self.rsphere = rsphere
        self.lap = self.lap / rsphere**2
        self.invlap = self.invlap * rsphere**2

        print("N: " + str(self.nlons) + ", " + str(self.nlats))
        print("Mtrunc: " + str(self.ntrunc))
        print("Nlm: " + str(self.nlm))
Esempio n. 8
0
    def __init__(self, Lmax, Mmax, fmetric=None, fricci=None):
        self.grid = shtns.sht(Lmax,Mmax, norm=shtns.SHT_REAL_NORM, nthreads=4)
        self.Lmax = Lmax
        self.nTheta, self.nPhi = self.grid.set_grid()
        self.nlm = self.grid.nlm
        self.numTerms = (self.Lmax+1)*(self.Lmax+1)
        self.extents = np.array([self.nTheta,self.nPhi])
        self.l, self.m = YlmIndex(np.arange(self.numTerms))
        self.index = np.array([self.grid.idx(int(self.l[i]),int(abs(self.m[i]))) for i in xrange(self.numTerms)])

#grid coordinates and properties:
        self.theta, self.phi = np.meshgrid(
            np.arccos(np.polynomial.legendre.leggauss(self.nTheta)[0])[::-1], 
            np.linspace(0,2*pi*(1-1./self.nPhi),self.nPhi),
            sparse=False)
        self.theta, self.phi = self.theta.T, self.phi.T
        self.gaussian_weights = np.polynomial.legendre.leggauss(self.nTheta)[1]
        self.costheta = np.cos(self.theta)
        self.sintheta = np.sin(self.theta)

        #Metric on grid points:
        if fmetric == None:
            self.gthth = np.zeros(self.extents)
            self.gphph = np.zeros(self.extents)
            self.gthph = np.zeros(self.extents)
        else:
            self.gthth, self.gphph, self.gthph = fmetric(self.theta, self.phi)
            self.UpdateMetric()

        # Compute ricci scalar, or not...
        if fricci==None:
            self.ricci = None
        else:
            self.ricci = fricci(self.theta, self.phi)
Esempio n. 9
0
def write_model_ylm(model, lmax, fname):
    # assumes the model is sampled on a regular grid where the poles are not
    # included

    # prepare the transform
    mmax = lmax
    sh = shtns.sht(lmax, mmax)

    npts_ph, npts_th = model.shape

    grid_typ = shtns.sht_reg_fast | shtns.SHT_THETA_CONTIGUOUS
    grid_typ = grid_typ | shtns.SHT_LOAD_SAVE_CFG

    polar_opt_threshold = 1.0e-10
    sh.set_grid(npts_th,
                npts_ph,
                flags=grid_typ,
                polar_opt=polar_opt_threshold)

    # apply transform
    ylm = sh.analys(model)

    with open(fname, 'w') as f:
        for l in zip(sh.l, sh.m, np.real(ylm), np.imag(ylm)):
            f.write('%5d %5d %14e %14e\n' % l)
    return
Esempio n. 10
0
 def __init__(self,nlons,nlats,ntrunc,rsphere,gridtype='gaussian'):
     """initialize
     nlons:  number of longitudes
     nlats:  number of latitudes
     ntrunc: spectral truncation
     rsphere: sphere radius (m)
     gridtype: 'gaussian' (default) or 'regular'"""
     self._shtns = shtns.sht(ntrunc, ntrunc, 1,\
             shtns.sht_fourpi|shtns.SHT_NO_CS_PHASE)
     if gridtype == 'gaussian':
         self._shtns.set_grid(nlats,nlons,shtns.sht_quick_init|shtns.SHT_PHI_CONTIGUOUS,1.e-8)
     elif gridtype == 'regular':
         self._shtns.set_grid(nlats,nlons,shtns.sht_reg_dct|shtns.SHT_PHI_CONTIGUOUS,1.e-8)
     #self._shtns.print_info()
     self.lats = np.arcsin(self._shtns.cos_theta)
     self.lons = (2.*np.pi/nlons)*np.arange(nlons)
     self.nlons = nlons
     self.nlats = nlats
     self.ntrunc = ntrunc
     self.nlm = self._shtns.nlm
     self.degree = self._shtns.l
     self.order = self._shtns.m
     if gridtype == 'gaussian':
         self.gauwts =\
         np.concatenate((self._shtns.gauss_wts(),self._shtns.gauss_wts()[::-1]))
     else:
         self.gauwts = None
     self.gridtype = gridtype
     self.lap = -self.degree*(self.degree+1.0).astype(np.complex)
     self.invlap = np.zeros(self.lap.shape, self.lap.dtype)
     self.invlap[1:] = 1./self.lap[1:]
     self.rsphere = rsphere
     self.lap = self.lap/rsphere**2
     self.invlap = self.invlap*rsphere**2
Esempio n. 11
0
 def __init__(self, l_max, force_python=False):
     """initialize a spherical harmonic transform object"""
     self._l_max = l_max
     if not force_python:
         try:
             import shtns
             self._shtns = shtns.sht(l_max, l_max)
         except ImportError:
             LOG.debug('Could not import shtns')
Esempio n. 12
0
 def __init__(self, l_max, force_python=False):
     """initialize a spherical harmonic transform object"""
     self._l_max = l_max
     if not force_python:
         try:
             import shtns
             self._shtns = shtns.sht(l_max, l_max)
         except ImportError:
             LOG.debug('Could not import shtns')
Esempio n. 13
0
def fcf(run_ID, directory, l_trunc=8):
    '''
    Flux concentration factor used in measure of semblance (Christensen et al. 2010). Note that we truncate at degree and order 8
    '''
    St_file = list_St_files(run_ID, directory)  # Find all St_no in folder
    n = len(St_file)

    # Loop over time
    FCF_s = []
    i = 0
    for file in St_file:
        print('Loading {} ({}/{})'.format(file, i + 1, n))
        filename = '{}/{}'.format(directory, file)

        (_, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, ntheta, nphi, _, _,
         theta, phi, _, _, Br, _) = surfaceload(filename)

        # Truncate degree 8
        m_max = l_trunc
        sh = shtns.sht(l_trunc, m_max)
        sh = shtns.sht(l_trunc, m_max)
        nlat, nlon = sh.set_grid(nphi=nphi, nlat=ntheta)
        # NOTE: array has to be dtype='float64' and not 'float32'
        vr = Br.T.astype('float64')
        clm = sh.analys(vr)  # spatial to spectral
        Br_f = sh.synth(clm)  # spectral to spatial
        Br_4 = (Br_f**4).mean()
        Br_2 = (Br_f**2).mean()
        FCF = (Br_4 - Br_2**2) / Br_2**2

        # Append
        FCF_s.append(np.mean(FCF))
        i += 1
    # Time average (should really divide by dt but n is good enough)
    FCF_out = sum(FCF_s) / n

    # Save
    with h5py.File('{}/fcf'.format(directory), 'w') as f:
        f.create_dataset('FCF_out', data=FCF_out)
    print('{}/fcf saved'.format(directory))

    return FCF_out
Esempio n. 14
0
 def __init__(self,nlons,nlats,ntrunc,rsphere,gridtype='gaussian'):
     """initialize
     nlons:  number of longitudes
     nlats:  number of latitudes
     """
     self._shtns = shtns.sht(ntrunc, ntrunc, 1, shtns.sht_orthonormal+shtns.SHT_NO_CS_PHASE)
     if gridtype == 'gaussian':
         self._shtns.set_grid(nlats,nlons,shtns.sht_quick_init|shtns.SHT_PHI_CONTIGUOUS,0)
     elif gridtype == 'regular':
         self._shtns.set_grid(nlats,nlons,shtns.sht_reg_dct|shtns.SHT_PHI_CONTIGUOUS,0)
     self.lats = np.arcsin(self._shtns.cos_theta)
     self.lons = (2.*np.pi/nlons)*np.arange(nlons)
Esempio n. 15
0
def filter_field(Br, nphi, ntheta, l_trunc):
    '''Use SHTns to truncate field to degree l_trunc'''
    m_max = l_trunc
    sh = shtns.sht(
        l_trunc,
        m_max)  # by default 'flag = sht_quick_init' uses gaussian grid
    nlat, nlon = sh.set_grid(
        nphi=nphi,
        nlat=ntheta)  # NOTE: array has to be dtype='float64' and not 'float32'
    vr = Br.T.astype('float64')
    ylm = sh.analys(vr)  # spatial to spectral
    Br_f = sh.synth(ylm)  # spectral to spatial
    return Br_f
Esempio n. 16
0
def extrapot(lmax,rcmb,brcmb,rout):

    nphi, ntheta = brcmb.shape
    nrout = len(rout)

    polar_opt = 1e-10

    lmax = int(nphi/3)
    mmax = lmax

    try:
        import shtns
    except ImportError:
        print("Potential extrapolation requires the SHTns library")
        print("It can be obtained here: https://bitbucket.org/nschaeff/shtns")

    norm=shtns.sht_orthonormal | shtns.SHT_NO_CS_PHASE

    sh = shtns.sht(lmax,mmax=mmax,norm=norm)
    ntheta, nphi = sh.set_grid(ntheta, nphi, polar_opt=polar_opt)


    L = sh.l * (sh.l + 1)

    brlm = sh.analys(brcmb.T)
    bpolcmb = np.zeros_like(brlm)
    bpolcmb[1:] = rcmb**2 * brlm[1:]/L[1:]
    btor = np.zeros_like(brlm)

    brout = np.zeros([ntheta,nphi,nrout])
    btout = np.zeros([ntheta,nphi,nrout])
    bpout = np.zeros([ntheta,nphi,nrout])

    for k,radius in enumerate(rout):
        print(("%d/%d" %(k,nrout)))

        radratio = rcmb/radius
        bpol = bpolcmb * radratio**(sh.l)
        brlm = bpol * L/radius**2
        brout[...,k] = sh.synth(brlm)

        dbpoldr = -sh.l/radius * bpol
        slm = dbpoldr

        btout[...,k], bpout[...,k] = sh.synth(slm,btor)

    brout = np.transpose(brout,(1,0,2))
    btout = np.transpose(btout,(1,0,2))
    bpout = np.transpose(bpout,(1,0,2))

    return brout, btout, bpout
Esempio n. 17
0
        def setup(self, file_info, data):
            import shtns
            import numpy as np

            if file_info['modes_m_max'] != file_info['modes_m_max']:
                raise Exception("Only num_lon == num_lat supported")

            ntrunc = file_info['modes_n_max']
            self._shtns = shtns.sht(
                ntrunc, ntrunc, 1,
                shtns.sht_orthonormal + shtns.SHT_NO_CS_PHASE)

            nlons = (ntrunc + 1) * 2
            nlats = (ntrunc + 1)

            if file_info['grid_type'] == 'GAUSSIAN':
                #self._shtns.set_grid(nlats,nlons,shtns.sht_gauss_fly|shtns.SHT_PHI_CONTIGUOUS, 1.e-10)
                self._shtns.set_grid(
                    nlats, nlons,
                    shtns.sht_quick_init | shtns.SHT_PHI_CONTIGUOUS, 0)
            elif file_info['grid_type'] == 'REGULAR':
                #self._shtns.set_grid(nlats,nlons,shtns.sht_reg_dct|shtns.SHT_PHI_CONTIGUOUS, 1.e-10)
                self._shtns.set_grid(
                    nlats, nlons, shtns.sht_reg_dct | shtns.SHT_PHI_CONTIGUOUS,
                    0)
            else:
                raise Exception("Grid type '" + file_info['grid_type'] +
                                "' not supported!")

            self.lats = np.arcsin(self._shtns.cos_theta)
            self.lons = (2. * np.pi / nlons) * np.arange(nlons)
            self.nlons = nlons
            self.nlats = nlats
            self.ntrunc = ntrunc
            self.nlm = self._shtns.nlm
            self.degree = self._shtns.l
            self.lap = -self.degree * (self.degree + 1.0).astype(np.complex)
            self.invlap = np.zeros(self.lap.shape, self.lap.dtype)
            self.invlap[1:] = 1. / self.lap[1:]
            self.lap = self.lap / self.rsphere**2
            self.invlap = self.invlap * self.rsphere**2

            sh_data = np.frombuffer(data, dtype=np.complex128)

            return sh_data
Esempio n. 18
0
    def simulate( self, q_values, 
        n_theta, n_phi, 
        lmax, n_psi, 
        dont_rotate = True,
        make_positive = True):
        """
        Simulates S(q), projects into spherical harmonics, and computes correlation  
        
        Parameters
        ----------
        q_values : np.array, q magnitudes to be simulated
        n_theta : int, number of thetas in q space, range = [0, pi]
        This is not the same as Bragg theta; this theta = Bragg theta + pi/2
        n_phi: int, number of phis, range = [0, 2pi] 
        lmax: int, maximum order of spherical harmonics projection
        dont_rotate: default True, do not rotate the traj when simulating S(q)
        
        """
        self.n_theta = n_theta
        self.n_phi = n_phi
        self.q_values = q_values
        self.n_q = q_values.shape[0]
        self.lmax = lmax
        
        self.n_psi = n_psi
        
        self.sh = shtns.sht(lmax)
    
        self.sh.set_grid( n_theta, n_phi )
        
        # simulate S_q
        self._compute_Sq( self.model, q_values, dont_rotate)

         # make sure that the inverse spherical transformation always gets a positive intensity map
        if make_positive:
            self._make_Sq_positive()
    
        # project into spherical harmonics
        self._sph_harm_project()
        
        # sum slm to get cl
        self._leg_coefs_from_sph_harm()
        
        # compute correlations
        self._sph_coefs_to_corr()
Esempio n. 19
0
    def __init__(self,
                 R,
                 lmax,
                 nphi=None,
                 ntheta=None,
                 delta_phi=1.,
                 verbose=False):
        self.R = R
        self.lmax = lmax

        # prepare the spherical harmonic transform
        self.nphi = nphi or self.lmax * 4
        self.ntheta = ntheta or self.lmax * 2
        self.mmax = self.lmax
        self.sh = shtns.sht(self.lmax, self.mmax)

        grid_typ = shtns.sht_gauss | shtns.SHT_THETA_CONTIGUOUS | \
            shtns.SHT_LOAD_SAVE_CFG

        polar_opt_threshold = 1.0e-10
        self.sh.set_grid(self.ntheta,
                         self.nphi,
                         flags=grid_typ,
                         polar_opt=polar_opt_threshold)

        self.theta = np.arccos(self.sh.cos_theta)
        self.phi = np.linspace(0, 2 * np.pi, self.nphi)

        self.workspace = self.sh.spec_array()

        self.delta_phi = np.deg2rad(delta_phi)

        self.group_velocity = None
        self.phase_velocity = None
        self.topo = None
        self.group_velocity_ylm = None
        self.phase_velocity_ylm = None
        self.topo_ylm = None
        self.group_velocity_notrot = None
        self.phase_velocity_notrot = None
        self.topo_notrot = None
        self.group_velocity_ylm_notrot = None
        self.phase_velocity_ylm_notrot = None
        self.topo_ylm_notrot = None
Esempio n. 20
0
 def __init__(self, nlons, nlats, ntrunc, rsphere, gridtype='gaussian'):
     """initialize
     nlons:  number of longitudes
     nlats:  number of latitudes
     ntrunc: spectral truncation
     rsphere: sphere radius (m)
     gridtype: 'gaussian' (default) or 'regular'"""
     self._shtns = shtns.sht(ntrunc, ntrunc, 1,\
             shtns.sht_fourpi|shtns.SHT_NO_CS_PHASE)
     if gridtype == 'gaussian':
         self._shtns.set_grid(
             nlats, nlons, shtns.sht_quick_init | shtns.SHT_PHI_CONTIGUOUS,
             1.e-8)
     elif gridtype == 'regular':
         self._shtns.set_grid(nlats, nlons,
                              shtns.sht_reg_dct | shtns.SHT_PHI_CONTIGUOUS,
                              1.e-8)
     #self._shtns.print_info()
     self.lats = np.arcsin(self._shtns.cos_theta)
     self.lons = (2. * np.pi / nlons) * np.arange(nlons)
     self.nlons = nlons
     self.nlats = nlats
     self.ntrunc = ntrunc
     self.nlm = self._shtns.nlm
     self.degree = self._shtns.l
     self.order = self._shtns.m
     if gridtype == 'gaussian':
         self.gauwts =\
         np.concatenate((self._shtns.gauss_wts(),self._shtns.gauss_wts()[::-1]))
     else:
         self.gauwts = None
     self.gridtype = gridtype
     self.lap = -self.degree * (self.degree + 1.0).astype(np.complex)
     self.invlap = np.zeros(self.lap.shape, self.lap.dtype)
     self.invlap[1:] = 1. / self.lap[1:]
     self.rsphere = rsphere
     self.lap = self.lap / rsphere**2
     self.invlap = self.invlap * rsphere**2
Esempio n. 21
0
    def setup(self, file_info, anti_aliasing):
        import shtns
        import numpy as np

        if file_info['modes_m_max'] != file_info['modes_m_max']:
            raise Exception("Only num_lon == num_lat supported")

        ntrunc = file_info['modes_n_max']
        self._shtns = shtns.sht(ntrunc, ntrunc, 1,
                                shtns.sht_orthonormal + shtns.SHT_NO_CS_PHASE)

        nlons = (ntrunc + 1) * 2
        nlats = (ntrunc + 1)

        if anti_aliasing:
            if nlons & 1:
                raise Exception(
                    "Only even numbers of longitudes coordinates allowed for anti-aliasing"
                )
            if nlats & 1:
                raise Exception(
                    "Only even numbers of latitudinal coordinates allowed for anti-aliasing"
                )

            print("Anti-aliasing:")
            print(" + old lon/lat: ", nlons, nlats)

            nlons += nlons // 2
            nlats += nlats // 2

            print(" + new lon/lat: ", nlons, nlats)

        self._shtns.set_grid(nlats, nlons,
                             shtns.sht_quick_init | shtns.SHT_PHI_CONTIGUOUS,
                             0)

        self.lats = np.arcsin(self._shtns.cos_theta)
        self.lons = (2. * np.pi / nlons) * np.arange(nlons)
Esempio n. 22
0
def gauss_coeffs(Br, nphi, ntheta, l_trunc):
    '''
    Obtain Gauss coefficients from core surface field
    '''
    # SH transform from spatial to spectral space
    m_max = l_trunc
    # by default 'flag = sht_quick_init' uses gaussian grid
    sh = shtns.sht(l_trunc, m_max)
    nlat, nlon = sh.set_grid(nphi=nphi, nlat=ntheta)
    vr = Br.T.astype(
        'float64')  # NOTE: array has to be dtype='float64' and not 'float32'
    clm = sh.analys(vr)  # spatial to spectral
    # Construct Gauss coefficients
    glm = np.zeros(sh.nlm)
    hlm = np.zeros(sh.nlm)
    for l in range(l_trunc + 1):
        fac = 1 / (l + 1
                   )  # cmb_radius/(l+1) # from Br potential to full potential
        for m in range(l + 1):
            glm[sh.idx(l, m)] = fac * np.real(clm[sh.idx(l, m)])
            hlm[sh.idx(l, m)] = fac * np.imag(clm[sh.idx(l, m)])

    return glm, hlm
Esempio n. 23
0
def filter_model_shtns(model,
                       lmax,
                       nphi_out=None,
                       ntheta_out=None,
                       order=2,
                       lmax_transform=None):
    # assumes the model is sampled on a regular grid where the poles are not
    # included

    if lmax_transform is None:
        lmax_transform = lmax * 2

    nphi, ntheta = model.shape
    lmax_transform = min(min(ntheta - 2, nphi / 2 - 1), lmax_transform)

    # prepare the transform
    mmax = lmax_transform
    sh = shtns.sht(lmax_transform, mmax)

    grid_typ = shtns.sht_reg_fast | shtns.SHT_THETA_CONTIGUOUS
    grid_typ = grid_typ | shtns.SHT_LOAD_SAVE_CFG

    polar_opt_threshold = 1.0e-10
    sh.set_grid(ntheta, nphi, flags=grid_typ, polar_opt=polar_opt_threshold)

    # apply transform
    ylm = sh.analys(model)

    ylm *= 1. / (1. + (sh.l * 1. / lmax)**(2 * order))

    if nphi_out and ntheta_out:
        sh.set_grid(ntheta_out,
                    nphi_out,
                    flags=grid_typ,
                    polar_opt=polar_opt_threshold)

    return sh.synth(ylm)
Esempio n. 24
0
    def load(self, filename):
        ff = h5py.File(filename, 'r')
        self.S_q = ff['intensity'].value
        self.n_theta = self.S_q.shape[1]
        self.n_phi = self.S_q.shape[2]

        self.lmax = ff['lmax'].value
        self.sh = shtns.sht(self.lmax)
        # generate the shtns object
        self.sh.set_grid( self.n_theta, self.n_phi )

        self.Sq_original = ff['intensity_original'].value
        
        self.all_slm =ff['slm'].value
        self.corr = ff['corr'].value
        self.cl = ff['leg_coefs'].value
        self.cospsi = ff['cospsi'].value
        self.n_psi = self.cospsi.size

        self.q_values = ff['qvalues'].value
        self.n_q = self.q_values.size


        ff.close()
Esempio n. 25
0
#zsurf = nfile.variables['zsurf']             #              lat, lon
#pres_half   = nfile.variables['pres_half']   # time, phalf, lat, lon
#height_half = nfile.variables['height_half'] # time, phalf, lat, lon
#pres_full   = nfile.variables['pres_full']   # time, pfull, lat, lon
#height_full = nfile.variables['height']      # time, pfull, lat, lon
field = nfile.variables[fieldname]  # time, pfull, lat, lon

#field_window = field[-w:,:,:,:]
field_state = field[-1, :, :, :]  # snapshot for now

nlon = field.shape[3]
nlat = field.shape[2]

mmax = min(tnum, (nlon - 1) // 2)
print tnum, mmax
sh = shtns.sht(tnum, mmax)  # reduce 2nd argument here as needed
#nlm = sh.nlm	# size of combined lm index
sh.set_grid(nlat, nlon)  # build default grid (gauss grid, phi-contiguous)

#print np.sctypeDict.keys()
field_k_l_m = np.empty((field.shape[1], tnum + 1, mmax + 1), dtype='complex64')
for k in range(0, field.shape[1]):
    field_lm = sh.analys(
        field_state[k, :, :].astype('float64'))  # sh.spec_array implicit?
    for l in range(0, tnum + 1):
        for m in range(0, mmax + 1):
            field_k_l_m[k, l, m] = field_lm[sh.idx(l, m)] if m <= l else np.nan

uspec = np.abs(field_k_l_m)**2  # pfull, l, m
uspec[:, :,
      1:] *= 2  # weight factor for missing m<0 (reality condition); 'optional' if not thinking of as summing displayed components
Esempio n. 26
0
    def __init__(self,
                 lmax=15,
                 mmax=None,
                 mres=1,
                 norm=shtns.sht_fourpi,
                 nlat=None,
                 nlon=None,
                 flags=(shtns.sht_quick_init | shtns.SHT_PHI_CONTIGUOUS
                        | shtns.SHT_SOUTH_POLE_FIRST),
                 polar_opt=1.0e-8,
                 nl_order=2,
                 radius=radius_earth):

        #print(lmax,mmax,mres,nlat,nlon,flags,polar_opt,nl_order,radius)
        #print('flags',flags,shtns.sht_quick_init,shtns.SHT_PHI_CONTIGUOUS,shtns.SHT_SOUTH_POLE_FIRST)

        if lmax is None and nlat is None:
            raise ValueError('lmax or nlat should be given.')
        elif lmax is None:
            lmax = compute_lmax(nlat)

        #print('lmax',lmax)

        self.lmax = int(lmax)
        self.radius = float(radius)
        #print(self.radius,radius_earth)

        if mmax is None and mres == 1:
            # triangular truncation
            self.mmax = self.lmax
            self.mres = 1
        else:
            self.mmax = int(float(lmax) / mres)
            self.mres = mres

        #print('norm',norm,self.lmax,self.mmax,self.mres)

        self.sh = shtns.sht(self.lmax, self.mmax, self.mres, norm=norm)

        bin_flags = bin_int(flags, 14)
        #print(bin_flags,bin_flags[-14])
        if bin_flags[-14] == '1':
            self.order_lat = 'south_to_north'
        else:
            self.order_lat = 'north_to_south'

        # in shtns, 0 means None
        if nlat is None:
            nlat = 0
        if nlon is None:
            nlon = 0

        self.nlat, self.nlon = self.sh.set_grid(nlat=nlat,
                                                nphi=nlon,
                                                flags=flags,
                                                polar_opt=polar_opt,
                                                nl_order=nl_order)
        #print('flags',flags)

        self.sin_lats = self.sh.cos_theta
        self.lons = np.arange(self.nlon) * 360. / (self.nlon * self.mres)
        self.lats = np.arcsin(self.sin_lats) / np.pi * 180.

        self.l_idx = self.sh.l
        #print('l_idx',self.l_idx)
        self.l2_idx = self.l_idx * (self.l_idx + 1)
        #print('l2_idx',self.l2_idx)
        self.m_idx = self.sh.m
        self.nlm = self.sh.nlm

        self.delta = 360. / (self.nlon * self.mres)
        # create arrays 2D lats et lons
        self.LONS, self.LATS = np.meshgrid(self.lons, self.lats)
        self.cosLATS = np.cos(self.LATS / 180 * np.pi)

        self.lrange = np.arange(self.lmax + 1)
        self.l2_l = self.lrange * (self.lrange + 1)
        self.kh_l = np.sqrt(self.l2_l) / self.radius

        self._complex64_save_netCFD = np.dtype([('real', np.float32),
                                                ('imag', np.float32)])
Esempio n. 27
0
    return vx, vy, vz


m = 20
lmax = 20

lUsr = int(sys.argv[1])
mUsr = int(sys.argv[2])

ntheta, nphi = [128, 256]

polar_opt = 1e-10

norm = shtns.sht_orthonormal | shtns.SHT_NO_CS_PHASE

sh = shtns.sht(lmax, mmax=m, norm=norm, nthreads=1)

ntheta, nphi = sh.set_grid(ntheta, nphi, polar_opt=polar_opt)

S = sh.spec_array()
T = sh.spec_array()

T[sh.idx(lUsr, mUsr)] = 1.

utheta, uphi = sh.synth(S, T)
psi = sh.synth(T)

x, y, z, th2D, p2D = get_grid(np.arccos(sh.cos_theta),
                              np.linspace(0., 2 * np.pi, nphi))

ux, uy, uz = get_cart(utheta, uphi, th2D, p2D)
Esempio n. 28
0
    def timeLongitude(self,
                      removeMean=True,
                      lat0=0.,
                      levels=12,
                      cm='RdYlBu_r',
                      deminc=True,
                      shtns_lib='shtns'):
        """
        Plot the time-longitude diagram of Br (input latitude can be chosen)

        .. warning:: the python bindings of `SHTns <https://bitbucket.org/bputigny/shtns-magic>`_ are mandatory to use this plotting function!

        :param lat0: value of the latitude
        :type lat0: float
        :param levels: number of contour levels
        :type levels: int
        :param cm: name of the colormap
        :type cm: str
        :param deminc: a logical to indicate if one wants do get rid of the
                       possible azimuthal symmetry
        :type deminc: bool
        :param shtns_lib: version of shtns library used: can be either 'shtns'
                          or 'shtns-magic'
        :type shtns_lib: char
        :param removeMean: remove the time-averaged part when set to True
        :type removeMean: bool
        """
        # The python bindings of shtns are mandatory to use this function !!!
        import shtns

        if removeMean:
            blmCut = self.blm - self.blm.mean(axis=0)
        else:
            blmCut = self.blm

        # Define shtns setup
        sh = shtns.sht(int(self.l_max_cmb),
                       int(self.m_max_cmb / self.minc),
                       mres=int(self.minc),
                       norm=shtns.sht_orthonormal | shtns.SHT_NO_CS_PHASE)

        polar_opt_threshold = 1e-10
        nlat = max(int(self.l_max_cmb * (3. / 2. / 2.) * 2.), 192)
        nphi = 2 * nlat / self.minc
        nlat, nphi = sh.set_grid(nlat, nphi, polar_opt=polar_opt_threshold)

        th = np.linspace(np.pi / 2., -np.pi / 2., nlat)
        lat0 *= np.pi / 180.
        mask = np.where(abs(th - lat0) == abs(th - lat0).min(), 1, 0)
        idx = np.nonzero(mask)[0][0]

        # Transform data on grid space
        BrCMB = np.zeros((self.nstep, nphi, nlat), 'Float64')
        if deminc:
            dat = np.zeros((self.nstep, self.minc * nphi + 1), 'Float64')
        else:
            dat = np.zeros((self.nstep, nphi), 'Float64')
        for k in range(self.nstep):
            tmp = sh.synth(blmCut[k, :] * sh.l * (sh.l + 1) / self.rcmb**2)
            tmp = tmp.T  # Longitude, Latitude

            if shtns_lib == 'shtns-magic':
                BrCMB[k, ...] = rearangeLat(tmp)
            else:
                BrCMB[k, ...] = tmp

            if deminc:
                dat[k, :] = symmetrize(BrCMB[k, :, idx], self.minc)
            else:
                dat[k, :] = BrCMB[k, :, idx]

        th = np.linspace(np.pi / 2., -np.pi / 2., nlat)
        if deminc:
            phi = np.linspace(-np.pi, np.pi, self.minc * nphi + 1)
        else:
            phi = np.linspace(-np.pi / self.minc, np.pi / self.minc, nphi)

        fig = plt.figure()
        ax = fig.add_subplot(111)
        vmin = -max(abs(dat.max()), abs(dat.min()))
        vmax = -vmin
        cs = np.linspace(vmin, vmax, levels)
        ax.contourf(phi, self.time, dat, cs, cmap=plt.get_cmap(cm))

        ax.set_xlabel('Longitude')
        ax.set_ylabel('Time')

        w2 = np.fft.fft2(dat)
        w2 = abs(w2[1:self.nstep / 2 + 1, 0:self.m_max_cmb + 1])

        dw = 2. * np.pi / (self.time[-1] - self.time[0])
        omega = dw * np.arange(self.nstep)
        omega = omega[1:self.nstep / 2 + 1]
        ms = np.arange(self.m_max_cmb + 1)

        fig1 = plt.figure()
        ax1 = fig1.add_subplot(111)
        ax1.contourf(ms, omega, w2, 17, cmap=plt.get_cmap('jet'))
        ax1.set_yscale('log')
        ax1.set_xlim(0, 13)
        ax1.set_xlabel(r'Azimuthal wavenumber')
        ax1.set_ylabel(r'Frequency')
Esempio n. 29
0
    def __init__(self, 
        q_values,
        lmax, n_theta, n_phi,
        corr = None, cospsi = None,
        ref_SphModel = None, auto_only = False,
        bark=False):
        
        # q values correponding to correlations, in ascending order!!!
        self.q_values = q_values
        # can pass a reference Spherical Model, use it for guessing, and comparing
        self.ref_SphModel = ref_SphModel

        self.lmax = lmax

        self.n_q = len(q_values)

        self.n_theta = n_theta
        self.n_phi = n_phi

        self.sh = shtns.sht(lmax)
        self.sh.set_grid( n_theta, n_phi )
        
        if corr is None:
            self.corr = self.ref_SphModel.corr
            self.cospsi = self.ref_SphModel.cospsi

            self.auto_only = auto_only
            if auto_only:
                self.cl = np.zeros_like(self.ref_SphModel.cl)
                self.cl[:,range(self.n_q),range(self.n_q)] = self.ref_SphModel.cl[:,range(self.n_q),range(self.n_q)]
            else:
                self.cl = self.ref_SphModel.cl
        else:

            if len(corr.shape) == 2:
                # that means there are only autocorrelator
                self.auto_only = True
                if len(cospsi.shape) == 1:
                    # that means all the q values have the same cospsi, reshape
                    cospsi = np.array( [cospsi] * self.n_q )
                    
            elif len(corr.shape) == 3:
                self.auto_only = False
                # auto and cross-correlations
                if len(cospsi.shape) == 1:
                    # that means all the q values have the same cospsi, reshape
                    cospsi = np.array( [[cospsi] * self.n_q] * self.n_q )
            else:
                print("Error: correlations should be either a 2D or 3D array")
        
            self.corr = corr
            self.cospsi = cospsi
            
            try:
                assert( self.cospsi.shape[0] == self.n_q)
                assert( self.corr.shape[-1] == self.cospsi.shape[-1])
                assert( len(self.corr.shape) == len(self.cospsi.shape) )
            except:
                print ("ERROR: Mismatch in cosines and correlations provided!")
            
            # project leg poly
            self._compute_legendre_projection()
        
        
        if bark:
            print ("\
                 .:##:::.\n \
              .:::::/;;\:.\n\
        ()::::::@::/;;#;|:.\n\
        ::::##::::|;;##;|::\n\
         \':::::::::\;;;/::\'\n\
              ':::::::::::\n\
               |O|O|O|O|O|O\n\
               :#:::::::##::.\n\
              .:###:::::#:::::.\n\
              :::##:::::::::::#:.\n\
               ::::;:::::::::###::.\n\
               \':::;::###::;::#:::::\n\
                ::::;::#::;::::::::::\n\
                :##:;::::::;::::###:::     .\n\
              .:::::; .:::##::::::::::::::::\n\
              ::::::; :::::::::::::::::##::\n\
              The Phase Retrieving Golden Retriever Puppy!")

        
        def save(self, save_filename):
            ff = h5py.File(save_filename,'a')
            ff.create_dataset('intensity', data = self.I_guess)
            ff.create_dataset('slm',data = self.all_slm_guess)

            ff.create_dataset('qvalues', data = self.q_values)
            ff.create_dataset('leg_coefs', data = self.cl)
            ff.create_dataset('corr', data = self.corr)

            ff.create_dataset('deltas', data  = self.deltas)
            ff.create_dataset('cospsi', data = self.cospsi)
            ff.create_dataset('lmax', data = self.lmax)

            ff.close()
Esempio n. 30
0
u2d = u3d[2]
v2d = v3d[2]

hdiv_lm, hrot_lm = ds.oper.hdivrotsh_from_uv(u2d.astype(np.float64),
                                             v2d.astype(np.float64))
uuu, vvv = ds.oper.uv_from_hdivrotsh(hdiv_lm, hrot_lm)

print 'uuu', uuu
print 'u2d', u2d
print 'vvv', vvv
print 'v2d', v2d

print 'shtns'

#does not work :(:(:(:(
sh1 = shtns.sht(127, 127, 1, norm=1)
nla, nlo = sh1.set_grid(nlat=192,
                        nphi=384,
                        flags=8708,
                        polar_opt=1e-08,
                        nl_order=2)

radius = 6367470.0
l2 = sh1.l * (sh1.l + 1)

hd = ds.oper.create_array_sh()
hr = ds.oper.create_array_sh()

vx1 = -v2d

sh1.spat_to_SHsphtor(vx1, u2d, hd, hr)
Esempio n. 31
0
def filt(dat_r, dat_t, dat_p, lmax=None, lCut=10, filt_type="high"):

    dat_r = float64(dat_r)
    dat_t = float64(dat_t)
    dat_p = float64(dat_p)

    nphi = dat_r.shape[0]
    nlat = dat_r.shape[1]
    nr = dat_r.shape[2]

    # Convert from MagIC to shtns format: nphi,ntheta,nr -> ntheta,nphi,nr

    dat_r = transpose(dat_r, (1, 0, 2))
    dat_t = transpose(dat_t, (1, 0, 2))
    dat_p = transpose(dat_p, (1, 0, 2))

    if lmax is None:
        lmax = nphi / 3

    sh = shtns.sht(lmax=lmax,
                   mmax=lmax,
                   norm=shtns.sht_schmidt | shtns.SHT_NO_CS_PHASE)
    polar_opt = 1e-10
    nlat, nphi = sh.set_grid(nlat=nlat, nphi=nphi, polar_opt=polar_opt)

    Q = zeros([nr, sh.nlm], dtype="complex128")
    S = zeros([nr, sh.nlm], dtype="complex128")
    T = zeros([nr, sh.nlm], dtype="complex128")

    Q1 = zeros([nr, sh.nlm], dtype="complex128")
    S1 = zeros([nr, sh.nlm], dtype="complex128")
    T1 = zeros([nr, sh.nlm], dtype="complex128")

    # Forward transform: data -> SH

    for i in range(nr):
        Q[i, :], S[i, :], T[i, :] = sh.analys(dat_r[..., i], dat_t[..., i],
                                              dat_p[..., i])

    # Apply filter

    if filt_type == "high":
        mask = sh.l >= lCut
    elif filt_type == "low":
        mask = sh.l <= lCut
    else:
        print('Filter type must be "high" or "low"')
        exit()

    Q1[:, mask] = Q[:, mask]
    S1[:, mask] = S[:, mask]
    T1[:, mask] = T[:, mask]

    # Inverse transform: SH -> data

    for k in range(nr):
        dat_r[..., k], dat_t[...,
                             k], dat_p[...,
                                       k] = sh.synth(Q1[k, :], S1[k, :],
                                                     T1[k, :])

    # Convert from shtns to MagIC format: ntheta,nphi,nr -> nphi,ntheta,nr

    dat_r = transpose(dat_r, (1, 0, 2))
    dat_t = transpose(dat_t, (1, 0, 2))
    dat_p = transpose(dat_p, (1, 0, 2))

    return dat_r, dat_t, dat_p
Esempio n. 32
0
def plot_radial_cmplx_wrapper(dir_NM, i_mode, n_lat_grid, i_radius_str = None, show = True, fmt = 'png', transparent = True):

    # Define directories.
    dir_processed = os.path.join(dir_NM, 'processed')
    dir_spectral = os.path.join(dir_processed, 'spectral')
    dir_plot = os.path.join(dir_processed, 'plots')
    mkdir_if_not_exist(dir_plot)

    # Determine if the plot is 'quick' mode or 'full' mode.
    if i_radius_str is None:

        option = 'quick'

    else:

        option = 'full'

    fig = plt.figure(figsize = (7.0, 5.0))

    # Reconstruct the coordinate grid.
    n_lon_grid = (2*n_lat_grid) - 1
    lon_grid = np.linspace(0.0, 2.0*np.pi, n_lon_grid + 1, endpoint = True)[:-1]
    lat_grid = np.linspace(-np.pi/2.0, np.pi/2.0, n_lat_grid, endpoint=True)

    # List of radii to sample.
    if i_radius_str == 'all':

        _, _, _, _, _, header_info = load_vsh_coefficients(dir_NM, i_mode, 0)
        n_samples = len(header_info['r_sample'])

        i_radius_list = list(range(n_samples))

    elif i_radius_str is None:

        i_radius_list = [None]

    else:

        i_radius_list = [int(i_radius_str)]

    first_iteration = True
    for j_radius in i_radius_list:
        
        # Load the VSH coefficients for the specified mode.
        coeffs, header_info, r_sample, i_sample = load_vsh_coefficients(dir_NM, i_mode, j_radius)
        Ulm_real, Vlm_real, Wlm_real, Ulm_imag, Vlm_imag, Wlm_imag = coeffs

        title = region_int_to_title(i_sample, r_sample, shell_name_path = os.path.join(dir_processed, 'shell_names.txt'))

        if first_iteration:

            # Infer the maximum l-value used.
            #n_coeffs = len(Ulm)
            n_coeffs = coeffs.shape[-1]
            l_max = (int((np.round(np.sqrt(8*n_coeffs + 1)) - 1))//2) - 1

            # Re-construct the shtns calculator.
            m_max = l_max
            sh_calculator   = shtns.sht(l_max, m_max)
            grid_type       =       shtns.sht_reg_fast          \
                                |   shtns.SHT_SOUTH_POLE_FIRST  \
                                |   shtns.SHT_PHI_CONTIGUOUS
            sh_calculator.set_grid(n_lat_grid, n_lon_grid, flags = grid_type)

            first_iteration = False

        n_c_levels_default = 20 

        # Expand the VSH into spatial components and plot.
        U_real_r, V_real_e, V_real_n, W_real_e, W_real_n = project_from_spherical_harmonics(sh_calculator, Ulm_real, Vlm_real, Wlm_real)
        U_imag_r, V_imag_e, V_imag_n, W_imag_e, W_imag_n = project_from_spherical_harmonics(sh_calculator, Ulm_imag, Vlm_imag, Wlm_imag)

        U_cplx_r = U_real_r + 1.0j*U_imag_r
        #U_abs_r = np.abs(U_cplx_r)
        #U_phase_r = np.angle(U_cplx_r)

        plot_radial_cplx(lon_grid, lat_grid, U_cplx_r, fig = fig, ax_arr = None, show = False, title = title, n_c_levels = n_c_levels_default)

        # Save the plot.
        fig_base_name = 'disp_radial_cplx'
        if option == 'quick':

            name_fig = '{:}_{:}_{:>05d}.{:}'.format(fig_base_name, option, i_mode, fmt)

        else:

            name_fig = '{:}_{:}_{:>05d}_{:>03d}.{:}'.format(fig_base_name, option, i_mode, j_radius, fmt)

        path_fig = os.path.join(dir_plot, name_fig)
        print('Saving figure to {:}'.format(path_fig))
        save_figure(path_fig, fmt, transparent = transparent)

        # Show the plot.
        if show:

            plt.show()

        # Close the plot.
        plt.close()

    return
Esempio n. 33
0
def plot_sh_disp_wrapper(dir_NM, i_mode, n_lat_grid, mode_real_or_complex, show = True, fmt = 'pdf', transparent = False, i_radius_str = None, close = True, c_bar_label = 'Default', figsize = None, path_outline = None):
    '''
    Plots a vector field on the surface of a sphere in terms of the radial, consoidal and toroidal components.
    This is a wrapper for plot_sh_disp() which first loads the necessary arrays. 

    Input:

    dir_NM, i_mode, n_lat_grid, fmt
        See 'Definitions of variables' in NMPostProcess/process.py.

    Output:

    None
    '''

    # Set transparency.
    transparent = True

    # Define directories.
    dir_processed = os.path.join(dir_NM, 'processed')
    dir_spectral = os.path.join(dir_processed, 'spectral')
    dir_plot = os.path.join(dir_processed, 'plots')
    mkdir_if_not_exist(dir_plot)

    # Load outline data if specified.
    if path_outline is not None:

        outline_data = np.loadtxt(path_outline)

    else:

        outline_data = None

    # Determine if the plot is 'quick' mode or 'full' mode.
    if i_radius_str is None:

        option = 'quick'

    else:

        option = 'full'

    if mode_real_or_complex in ['real', 'complex_imag_only', 'complex_real_only']:
        
        if figsize is None:

            figsize = (3.5, 6.0)

        two_panel = False

    elif mode_real_or_complex == 'complex':

        if figsize is None:

            figsize = (7.0, 6.0)

        two_panel = True

    else:

        raise ValueError

    # Reconstruct the coordinate grid.
    n_lon_grid = (2*n_lat_grid) - 1
    lon_grid = np.linspace(0.0, 2.0*np.pi, n_lon_grid + 1, endpoint = True)[:-1]
    lat_grid = np.linspace(-np.pi/2.0, np.pi/2.0, n_lat_grid, endpoint=True)

    if i_radius_str == 'all':

        _, header_info, _, _ = load_vsh_coefficients(dir_NM, i_mode, 0)
        n_samples = len(header_info['r_sample'])

        i_radius_list = list(range(n_samples))

    elif i_radius_str is None:

        i_radius_list = [None]

    else:

        i_radius_list = [int(i_radius_str)]

    # Identify the eigenvalue list file.
    # Load the frequency of the mode.
    path_eigenvalues = os.path.join(dir_processed, 'eigenvalue_list.txt')
    freq_mHz = read_eigenvalues(path_eigenvalues, i_mode = i_mode)

    first_iteration = True
    for j_radius in i_radius_list:
        
        fig = plt.figure(figsize = figsize)

        # Load the VSH coefficients for the specified mode.
        coeffs, header_info, r_sample, i_sample = load_vsh_coefficients(dir_NM, i_mode, j_radius)

        if mode_real_or_complex == 'real':

            Ulm, Vlm, Wlm = coeffs

        else:

            Ulm_real, Vlm_real, Wlm_real, Ulm_imag, Vlm_imag, Wlm_imag = coeffs

        title_str_mode = 'Mode {:>4d}'.format(i_mode)
        title_str_freq = 'freq. {:>7.4f} mHz'.format(freq_mHz)
        title_str_sample = region_int_to_title(i_sample, r_sample, shell_name_path = os.path.join(dir_processed, 'shell_names.txt'))
        title_str_sample = title_str_sample[0].lower() + title_str_sample[1:]
        if two_panel:

            title = '{:} ({:}), sampled at {:}'.format(title_str_mode, title_str_freq, title_str_sample)

        else:

            title = '{:} ({:}),\nsampled at {:}'.format(title_str_mode, title_str_freq, title_str_sample)

        if first_iteration:

            # Infer the maximum l-value used.
            #n_coeffs = len(Ulm)
            n_coeffs = coeffs.shape[-1]
            l_max = (int((np.round(np.sqrt(8*n_coeffs + 1)) - 1))//2) - 1

            # Re-construct the shtns calculator.
            m_max = l_max
            sh_calculator   = shtns.sht(l_max, m_max)
            grid_type       =       shtns.sht_reg_fast          \
                                |   shtns.SHT_SOUTH_POLE_FIRST  \
                                |   shtns.SHT_PHI_CONTIGUOUS
            sh_calculator.set_grid(n_lat_grid, n_lon_grid, flags = grid_type)

            first_iteration = False

        n_c_levels_default = 20 

        # Expand the VSH into spatial components and plot.
        if mode_real_or_complex == 'real':

            U_r, V_e, V_n, W_e, W_n = project_from_spherical_harmonics(sh_calculator, Ulm, Vlm, Wlm)
            fig, ax_arr, handles, contour_info = plot_sh_disp_3_comp(lon_grid, lat_grid, U_r, V_e, V_n, W_e, W_n, fig = fig, ax_arr = None, show = False, title = title, n_c_levels = n_c_levels_default)

            fields = {  'U' : {'r' : U_r},
                        'V' : {'e' : V_e, 'n' : V_n},
                        'W' : {'e' : W_e, 'n' : W_n}}

        else:

            U_real_r, V_real_e, V_real_n, W_real_e, W_real_n = project_from_spherical_harmonics(sh_calculator, Ulm_real, Vlm_real, Wlm_real)
            U_imag_r, V_imag_e, V_imag_n, W_imag_e, W_imag_n = project_from_spherical_harmonics(sh_calculator, Ulm_imag, Vlm_imag, Wlm_imag)

            fields = {
                'U' :   {   'real' : {  'r' : U_real_r},
                            'imag' : {  'r' : U_imag_r}},
                'V' :   {   'real' : {  'n' : V_real_n,
                                        'e' : V_real_e},
                            'imag' : {  'n' : V_imag_n,
                                        'e' : V_imag_e}},
                'W' :   {   'real' : {  'n' : W_real_n,
                                        'e' : W_real_e},
                            'imag' : {  'n' : W_imag_n,
                                        'e' : W_imag_e}}}

            #U_abs_r = np.sqrt((U_real_r**2.0) + (U_imag_r**2.0))
            #V_abs_e = np.sqrt((V_real_e**2.0) + (V_imag_e**2.0))
            #V_abs_n = np.sqrt((V_real_n**2.0) + (V_imag_n**2.0))
            #W_abs_e = np.sqrt((W_real_e**2.0) + (W_imag_e**2.0))
            #W_abs_n = np.sqrt((W_real_n**2.0) + (W_imag_n**2.0))

            if c_bar_label == 'Default':

                c_bar_label_real = 'Real part'
                c_bar_label_imag = 'Imaginary part'

            else:

                c_bar_label_real = c_bar_label
                c_bar_label_imag = c_bar_label

            if mode_real_or_complex == 'complex':

                #plt.imshow(W_real_e)
                #plt.show()

                fig, ax_arr_real, handles_real, contour_info_real = plot_sh_disp_3_comp(lon_grid, lat_grid, U_real_r, V_real_e, V_real_n, W_real_e, W_real_n, fig = fig, axes_grid_int = 121, ax_arr = None, show = False, title = title, n_c_levels = n_c_levels_default, c_bar_label = c_bar_label_real, outline_data = outline_data)
                fig, ax_arr_imag, handles_imag, contour_info_imag = plot_sh_disp_3_comp(lon_grid, lat_grid, U_imag_r, V_imag_e, V_imag_n, W_imag_e, W_imag_n, fig = fig, axes_grid_int = 122, ax_arr = None, show = False, title = None, n_c_levels = n_c_levels_default, c_bar_label = c_bar_label_imag, outline_data = outline_data)

                contour_info = [contour_info_real, contour_info_imag]

                ax_arr = [ax_arr_real, ax_arr_imag]
                handles = [handles_real, handles_imag]

            elif mode_real_or_complex == 'complex_real_only':

                fig, ax_arr, handles, contour_info = plot_sh_disp_3_comp(lon_grid, lat_grid, U_real_r, V_real_e, V_real_n, W_real_e, W_real_n, fig = fig, ax_arr = None, show = False, title = title, n_c_levels = n_c_levels_default, c_bar_label = c_bar_label_real, outline_data = outline_data) 

            elif mode_real_or_complex == 'complex_imag_only':

                fig, ax_arr, handles, contour_info = plot_sh_disp_3_comp(lon_grid, lat_grid, U_real_r, V_real_e, V_real_n, W_real_e, W_real_n, fig = fig, ax_arr = None, show = False, title = title, n_c_levels = n_c_levels_default, c_bar_label = c_bar_label_imag, outline_data = outline_data)

            else:
                
                raise ValueError

        # Save the plot.
        if fmt is not None:

            real_or_complex_str_dict = {
                'real'              : 'real',
                'complex'              : 'cplx',
                'complex_real_only'     : 'real_only',
                'complex_imag_only'    : 'imag_only'}

            real_or_complex_str = real_or_complex_str_dict[mode_real_or_complex]

            if option == 'quick':

                name_fig = 'displacement_{:}_{:}_{:>05d}.{:}'.format(option, real_or_complex_str, i_mode, fmt)

            else:

                name_fig = 'displacement_{:}_{:}_{:>05d}_{:>03d}.{:}'.format(option, real_or_complex_str, i_mode, j_radius, fmt)
            
            #plt.draw()
            path_fig = os.path.join(dir_plot, name_fig)
            print('Saving figure to {:}'.format(path_fig))
            save_figure(path_fig, fmt, transparent = transparent)

        # Show the plot.
        if show:

            plt.show()

        # Close the plot.
        if close:
            
            plt.close()

    return fig, ax_arr, handles, fields, contour_info
Esempio n. 34
0
    def timeLongitude(self, removeMean=True, lat0=0., levels=12, cm='RdYlBu_r',
                      deminc=True, shtns_lib='shtns'):
        """
        Plot the time-longitude diagram of Br (input latitude can be chosen)

        .. warning:: the python bindings of `SHTns <https://bitbucket.org/bputigny/shtns-magic>`_ are mandatory to use this plotting function!

        :param lat0: value of the latitude
        :type lat0: float
        :param levels: number of contour levels
        :type levels: int
        :param cm: name of the colormap
        :type cm: str
        :param deminc: a logical to indicate if one wants do get rid of the
                       possible azimuthal symmetry
        :type deminc: bool
        :param shtns_lib: version of shtns library used: can be either 'shtns'
                          or 'shtns-magic'
        :type shtns_lib: char
        :param removeMean: remove the time-averaged part when set to True
        :type removeMean: bool
        """
        # The python bindings of shtns are mandatory to use this function !!!
        import shtns

        if removeMean:
            blmCut = self.blm-self.blm.mean(axis=0)
        else:
            blmCut = self.blm

        # Define shtns setup
        sh = shtns.sht(int(self.l_max_cmb), int(self.m_max_cmb/self.minc), 
                       mres=int(self.minc), 
                       norm=shtns.sht_orthonormal | shtns.SHT_NO_CS_PHASE)

        polar_opt_threshold = 1e-10
        nlat = max((self.l_max_cmb*(3/2/2)*2),192)
        nphi = 2*nlat/self.minc
        nlat, nphi = sh.set_grid(nlat, nphi, polar_opt=polar_opt_threshold)

        th = np.linspace(np.pi/2., -np.pi/2., nlat)
        lat0 *= np.pi/180.
        mask = np.where(abs(th-lat0) == abs(th-lat0).min(), 1, 0)
        idx = np.nonzero(mask)[0][0]

        # Transform data on grid space
        BrCMB = np.zeros((self.nstep, nphi, nlat), 'Float64')
        if deminc:
            dat = np.zeros((self.nstep, self.minc*nphi+1), 'Float64')
        else:
            dat = np.zeros((self.nstep, nphi), 'Float64')
        for k in range(self.nstep):
            tmp = sh.synth(blmCut[k, :]*sh.l*(sh.l+1)/self.rcmb**2)
            tmp = tmp.T # Longitude, Latitude

            if shtns_lib == 'shtns-magic':
                BrCMB[k, ...] = rearangeLat(tmp)
            else:
                BrCMB[k, ...] = tmp

            if deminc:
                dat[k, :] = symmetrize(BrCMB[k, :, idx], self.minc)
            else:
                dat[k, :] = BrCMB[k, :, idx]


        th = np.linspace(np.pi/2., -np.pi/2., nlat)
        if deminc:
            phi = np.linspace(-np.pi, np.pi, self.minc*nphi+1)
        else:
            phi = np.linspace(-np.pi/self.minc, np.pi/self.minc, nphi)

        fig = plt.figure()
        ax = fig.add_subplot(111)
        vmin = -max(abs(dat.max()), abs(dat.min()))
        vmax = -vmin
        cs = np.linspace(vmin, vmax, levels)
        ax.contourf(phi, self.time, dat, cs, cmap=plt.get_cmap(cm))

        ax.set_xlabel('Longitude')
        ax.set_ylabel('Time')

        w2 = np.fft.fft2(dat)
        w2 = abs(w2[1:self.nstep/2+1, 0:self.m_max_cmb+1])

        dw = 2.*np.pi/(self.time[-1]-self.time[0])
        omega = dw*np.arange(self.nstep)
        omega = omega[1:self.nstep/2+1]
        ms = np.arange(self.m_max_cmb+1)

        fig1 = plt.figure()
        ax1 = fig1.add_subplot(111)
        ax1.contourf(ms, omega, w2, 17, cmap=plt.get_cmap('jet'))
        ax1.set_yscale('log')
        ax1.set_xlim(0,13)
        ax1.set_xlabel(r'Azimuthal wavenumber')
        ax1.set_ylabel(r'Frequency')
Esempio n. 35
0
    def movieRad(self, cut=0.5, levels=12, cm='RdYlBu_r', png=False, step=1,
                 normed=False, dpi=80, bgcolor=None, deminc=True, removeMean=False,
                 precision='Float64', shtns_lib='shtns', contour=False, mer=False):
        """
        Plotting function (it can also write the png files)

        .. warning:: the python bindings of `SHTns <https://bitbucket.org/bputigny/shtns-magic>`_ are mandatory to use this plotting function!

        :param levels: number of contour levels
        :type levels: int
        :param cm: name of the colormap
        :type cm: str
        :param cut: adjust the contour extrema to max(abs(data))*cut
        :type cut: float
        :param png: save the movie as a series of png files when
                    set to True
        :type png: bool
        :param dpi: dot per inch when saving PNGs
        :type dpi: int
        :param bgcolor: background color of the figure
        :type bgcolor: str
        :param normed: the colormap is rescaled every timestep when set to True,
                       otherwise it is calculated from the global extrema
        :type normed: bool
        :param step: the stepping between two timesteps
        :type step: int
        :param deminc: a logical to indicate if one wants do get rid of the
                       possible azimuthal symmetry
        :type deminc: bool
        :param precision: single or double precision
        :type precision: char
        :param shtns_lib: version of shtns library used: can be either 'shtns'
                          or 'shtns-magic'
        :type shtns_lib: char
        :param contour: also display the solid contour levels when set to True
        :type contour: bool
        :param mer: display meridians and circles when set to True
        :type mer: bool
        :param removeMean: remove the time-averaged part when set to True
        :type removeMean: bool
        """

        # The python bindings of shtns are mandatory to use this function !!!
        import shtns

        if removeMean:
            dataCut = self.wlm-self.wlm.mean(axis=0)
        else:
            dataCut = self.wlm

        # Define shtns setup
        sh = shtns.sht(int(self.l_max_r), int(self.m_max_r/self.minc), 
                       mres=int(self.minc), 
                       norm=shtns.sht_orthonormal | shtns.SHT_NO_CS_PHASE)

        polar_opt_threshold = 1e-10
        nlat = max((self.l_max_r*(3/2/2)*2),192)
        nphi = 2*nlat/self.minc
        nlat, nphi = sh.set_grid(nlat, nphi, polar_opt=polar_opt_threshold)

        # Transform data on grid space
        data = np.zeros((self.nstep, nphi, nlat), precision)
        for k in range(self.nstep):
            tmp = sh.synth(dataCut[k, :]*sh.l*(sh.l+1)/self.radius**2)
            tmp = tmp.T # Longitude, Latitude

            if shtns_lib == 'shtns-magic':
                data[k, ...] = rearangeLat(tmp)
            else:
                data[k, ...] = tmp

        if png:
            plt.ioff()
            if not os.path.exists('movie'):
                os.mkdir('movie')
        else:
            plt.ion()

        if not normed:
            vmin = - max(abs(data.max()), abs(data.min()))
            vmin = cut * vmin
            vmax = -vmin
            cs = np.linspace(vmin, vmax, levels)

        th = np.linspace(np.pi/2., -np.pi/2., nlat)
        if deminc:
            phi = np.linspace(-np.pi, np.pi, self.minc*nphi+1)
            xxout, yyout = hammer2cart(th, -np.pi)
            xxin, yyin = hammer2cart(th, np.pi)
        else:
            phi = np.linspace(-np.pi/self.minc, np.pi/self.minc, nphi)
            xxout, yyout = hammer2cart(th, -np.pi/self.minc)
            xxin, yyin = hammer2cart(th, np.pi/self.minc)
        ttheta, pphi = np.meshgrid(th, phi)
        xx, yy = hammer2cart(ttheta, pphi)
        if deminc:
            fig = plt.figure(figsize=(8, 4))
        else:
            fig = plt.figure(figsize=(8/self.minc, 4))
        fig.subplots_adjust(top=0.99, right=0.99, bottom=0.01, left=0.01)
        ax = fig.add_subplot(111, frameon=False)

        if mer:
            theta = np.linspace(np.pi/2, -np.pi/2, nlat)
            meridians = np.r_[-120, -60, 0, 60, 120]
            circles = np.r_[ 60, 30, 0, -30, -60]

        for k in range(self.nstep):
            if k == 0:
                if normed:
                    vmin = - max(abs(data[k, ...].max()), abs(data[k, ...].min()))
                    vmin = cut * vmin
                    vmax = -vmin
                    cs = np.linspace(vmin, vmax, levels)
                if deminc:
                    dat = symmetrize(data[k, ...], self.minc)
                else:
                    dat = data[k, ...]
                im = ax.contourf(xx, yy, dat, cs, cmap=plt.get_cmap(cm), extend='both')
                if contour:
                    ax.contour(xx, yy, dat, cs, linestyles=['-', '-'],
                               colors=['k', 'k'], linewidths=[0.7, 0.7])
                ax.plot(xxout, yyout, 'k-', lw=1.5)
                ax.plot(xxin, yyin, 'k-', lw=1.5)

                if mer:
                    for lat0 in circles:
                        x0, y0 = hammer2cart(lat0*np.pi/180., phi)
                        ax.plot(x0, y0, 'k:', linewidth=0.7)
                    for lon0 in meridians:
                        x0, y0 = hammer2cart(theta, lon0*np.pi/180.)
                        ax.plot(x0, y0, 'k:', linewidth=0.7)

                ax.axis('off')
                man = plt.get_current_fig_manager()
                man.canvas.draw()

                if png:
                    filename = 'movie/img%05d.png' % k
                    print('write %s' % filename)
                    #st = 'echo %i' % ivar + ' > movie/imgmax'
                    if bgcolor is not None:
                        fig.savefig(filename, facecolor=bgcolor, dpi=dpi)
                    else:
                        fig.savefig(filename, dpi=dpi)

            elif k != 0 and k % step == 0:
                if not png:
                    print(k)
                plt.cla()
                if normed:
                    vmin = - max(abs(data[k, ...].max()), abs(data[k, ...].min()))
                    vmin = cut * vmin
                    vmax = -vmin
                    cs = np.linspace(vmin, vmax, levels)
                if deminc:
                    dat = symmetrize(data[k, ...], self.minc)
                else:
                    dat = data[k, ...]
                im = ax.contourf(xx, yy, dat, cs, cmap=plt.get_cmap(cm), extend='both')
                if contour:
                    ax.contour(xx, yy, dat, cs, colors=['k'],
                               linestyles=['-', '-'], linewidths=[0.7, 0.7])
                ax.plot(xxout, yyout, 'k-', lw=1.5)
                ax.plot(xxin, yyin, 'k-', lw=1.5)

                if mer:
                    for lat0 in circles:
                        x0, y0 = hammer2cart(lat0*np.pi/180., phi)
                        ax.plot(x0, y0, 'k:', linewidth=0.7)
                    for lon0 in meridians:
                        x0, y0 = hammer2cart(theta, lon0*np.pi/180.)
                        ax.plot(x0, y0, 'k:', linewidth=0.7)

                ax.axis('off')
                man.canvas.draw()
                if png:
                    filename = 'movie/img%05d.png' % k
                    print('write %s' % filename)
                    #st = 'echo %i' % ivar + ' > movie/imgmax'
                    if bgcolor is not None:
                        fig.savefig(filename, facecolor=bgcolor, dpi=dpi)
                    else:
                        fig.savefig(filename, dpi=dpi)
Esempio n. 36
0
#  The fact that you are presently reading this means that you have had
#  knowledge of the CeCILL license and that you accept its terms.
#  

###################################
# SHTns Python interface example  #
###################################

import numpy		# numpy for arrays
import shtns		# shtns module
					#   compiled and installe using ./configure --enable-python && make && make install

lmax = 7			# maximum degree of spherical harmonic representation.
mmax = 3			# maximum order of spherical harmonic representation.

sh = shtns.sht(lmax, mmax)		# create sht object with given lmax and mmax (orthonormalized)
# sh = shtns.sht(lmax, mmax, mres=2, norm=shtns.sht_schmidt | shtns.SHT_NO_CS_PHASE)	# use schmidt semi-normalized harmonics

nlat, nphi = sh.set_grid()		# build default grid (gauss grid, phi-contiguous)
print(sh.nlat, sh.nphi)			# displays the latitudinal and longitudinal grid sizes.

cost = sh.cos_theta			# latitudinal coordinates of the grid as cos(theta)
el = sh.l					# array of size sh.nlm giving the spherical harmonic degree l for any sh coefficient
l2 = el*(el+1)				# array l(l+1) that is useful for computing laplacian

### use advanced options to create a regular grid, theta-contiguous, and with south-pole comming first.
# nlat = lmax*2
# nphi = mmax*3
# grid_typ = shtns.sht_reg_fast | shtns.SHT_THETA_CONTIGUOUS | shtns.SHT_SOUTH_POLE_FIRST
# polar_opt_threshold = 1.0e-10
# sh.set_grid(nlat, nphi, flags=grid_typ, polar_opt=polar_opt_threshold)
Esempio n. 37
0
#  The fact that you are presently reading this means that you have had
#  knowledge of the CeCILL license and that you accept its terms.
#

###################################
# SHTns Python interface example  #
###################################

import numpy  # numpy for arrays
import shtns  # shtns module
#   compiled and installe using ./configure --enable-python && make && make install

lmax = 7  # maximum degree of spherical harmonic representation.
mmax = 3  # maximum order of spherical harmonic representation.

sh = shtns.sht(
    lmax, mmax)  # create sht object with given lmax and mmax (orthonormalized)
# sh = shtns.sht(lmax, mmax, mres=2, norm=shtns.sht_schmidt | shtns.SHT_NO_CS_PHASE)	# use schmidt semi-normalized harmonics

nlat, nphi = sh.set_grid()  # build default grid (gauss grid, phi-contiguous)
print(sh.nlat,
      sh.nphi)  # displays the latitudinal and longitudinal grid sizes.

cost = sh.cos_theta  # latitudinal coordinates of the grid as cos(theta)
el = sh.l  # array of size sh.nlm giving the spherical harmonic degree l for any sh coefficient
l2 = el * (el + 1)  # array l(l+1) that is useful for computing laplacian

### use advanced options to create a regular grid, theta-contiguous, and with south-pole comming first.
# nlat = lmax*2
# nphi = mmax*3
# grid_typ = shtns.sht_reg_fast | shtns.SHT_THETA_CONTIGUOUS | shtns.SHT_SOUTH_POLE_FIRST
# polar_opt_threshold = 1.0e-10