Source code for pmesh.transfer

"""
.. deprecated:: 0.1

"""

import warnings
warnings.warn("transfer.py is deprecated. Directly operate over slabiter.", DeprecationWarning)

import numpy
from mpi4py import MPI
try:
    from numba import jit
except ImportError:
    jit = None
    pass
[docs]class TransferFunction: """ these a a function of Window Transfer functions used by PM. they take the fourier-space field complex and the dimensionless circular frequency as inputs; complex is modified in-place. some functions are factories: they return an actually window with the given parameter. Working out the dimension of the input is important. the output of Poisson introduces a dimension to rho_k! the output of SuperLanzcos introduces a dimension to rho_k! w is a tuple of (w0, w1, w2, ...) w is in circular frequency units. The dimensionful k is w * Nmesh / BoxSize (nyquist is at about w = pi) they broadcast to the correct shape of complex. This is to reduce memory usage somewhat. complex is modified in place. """
[docs] @staticmethod def NormalizeDC(pm, complex): """ removes the DC amplitude. This effectively divides by the mean """ w = pm.w comm = pm.comm ind = [] value = 0.0 found = True for wi in w: if (wi != 0).all(): found = False break ind.append((wi == 0).nonzero()[0][0]) if found: ind = tuple(ind) value = numpy.abs(complex[ind]) value = comm.allreduce(value, MPI.SUM) complex[:] /= value
[docs] @staticmethod def RemoveDC(pm, complex): w = pm.w comm = pm.comm ind = [] for wi in w: if (wi != 0).all(): return ind.append((wi == 0).nonzero()[0][0]) ind = tuple(ind) complex[ind] = 0
[docs] @staticmethod def Trilinear(comm, complex, w): for wi in w: # convert to tmp = numpy.sinc(wi / (2 * numpy.pi)) ** 2 complex[:] /= tmp
[docs] @staticmethod def SuperLanzcos(dir, order=3): """ Super Lanzcos kernel of order 3. is complex * i * w in a fancier way. Notice that for differentiation, one actually wants complex * i * k which is complex * i * w * Nmesh / BoxSize """ def SuperLanzcosDir(pm, complex): w = pm.w comm = pm.comm wi = w[dir] * 1.0 # tmp = (1. / 594 * # (126 * numpy.sin(wi) + 193 * numpy.sin(2 * wi) + 142 * numpy.sin (3 * # wi) - 86 * numpy.sin(4 * wi))) tmp = 1 / 6.0 * (8 * numpy.sin (wi) - numpy.sin (2 * wi)) if order == 0: complex *= wi * 1j else: complex[:] *= tmp * 1j return SuperLanzcosDir
[docs] @staticmethod def Gaussian(smoothing): """ smoothing is in mesh units; Notice that this is different from the usual PM split convention. (used in Gadget2/3) The PM split is cut = sqrt(0.5) * smoothing """ sm2 = smoothing ** 2 def GaussianS(pm, complex): w = pm.w comm = pm.comm w2 = 0 for wi in w: wi2 = wi ** 2 complex *= numpy.exp(-0.5 * wi2 * sm2) return GaussianS
[docs] @staticmethod def Constant(C): def Constant(pm, complex): w = pm.w comm = pm.comm complex *= C return Constant
[docs] @staticmethod def Inspect(name, *indices): """ inspect the complex array at given indices mostly for debugging. """ def Inspect(pm, complex): w = pm.w comm = pm.comm V = ['%s = %s' %(str(i), str(complex[tuple(i)])) for i in indices] print(name, ','.join(V)) return Inspect
[docs] @staticmethod def PowerSpectrum(wout, psout): """ calculate the power spectrum. This shall be done after NormalizeDC and RemoveDC """ def PS(pm, complex): w = pm.w comm = pm.comm wedges = numpy.linspace(0, numpy.pi, wout.size + 1, endpoint=True) w2edges = wedges ** 2 P = numpy.zeros(len(psout)) N = numpy.zeros(len(psout)) for row in range(complex.shape[0]): # pickup the singular plane that is single counted (r2c transform) singular = w[0][-1] == 0 scratch = w[0][row] ** 2 for wi in w[1:]: scratch = scratch + wi[0] ** 2 # now scratch stores w ** 2 dig = numpy.digitize(scratch.flat, w2edges) # take the sum of w scratch **= 0.5 # the singular plane is down weighted by 0.5 scratch[singular] *= 0.5 wsum = numpy.bincount(dig, weights=scratch.flat, minlength=wout.size + 2)[1: -1] wsum = comm.allreduce(wsum, MPI.SUM) # take the sum of weights scratch[...] = 1.0 # the singular plane is down weighted by 0.5 scratch[singular] = 0.5 N1 = numpy.bincount(dig, weights=scratch.flat, minlength=wout.size + 2)[1: -1] N += comm.allreduce(N1, MPI.SUM) # take the sum of power numpy.abs(complex[row], out=scratch) scratch[...] **= 2.0 # the singular plane is down weighted by 0.5 scratch[singular] *= 0.5 P1 = numpy.bincount(dig, weights=scratch.flat, minlength=wout.size + 2)[1: -1] P += comm.allreduce(P1, MPI.SUM) psout[:] = P / N wout[:] = wsum / N return PS
[docs] @staticmethod def Laplace(pm, complex): """ Take the Laplacian k-space: complex *= -w2 where this function performs only the -w **-2 part. Note that k = w * Nmesh / BoxSize, thus the usual laplacian is - k ** 2 * complex = (Nmesh / BoxSize) ** 2 (-w**2) * complex """ w = pm.w comm = pm.comm for row in range(complex.shape[0]): w2 = w[0][row] ** 2 for wi in w[1:]: w2 = w2 + wi[0] ** 2 w2[w2 == 0] = numpy.inf w2 *= -1 complex[:] *= w2
[docs] @staticmethod def Poisson(pm, complex): """ Solve Poisson equation in k-space: complex /= -w2 Notes about gravity: gravity is .. math :: \phi_k = -4 \pi G \delta_k k^{-2} where :math:`k = \omega N_m / L`. hence .. math :: \phi_k = -4 \pi G \delta_k \omega^{-2} (N_m / L)^{-2} where this function performs only the :math:`- \omega **-2` part. """ w = pm.w comm = pm.comm for row in range(complex.shape[0]): w2 = w[0][row] ** 2 for wi in w[1:]: w2 = w2 + wi[0] ** 2 w2[w2 == 0] = numpy.inf w2 *= -1 complex[row] /= w2
if __name__ == '__main__': def test(): complex = numpy.ones((2, 3)) ndim = len(complex.shape) w = [] for d in range(ndim): s = numpy.ones(ndim, dtype='intp') s[d] = complex.shape[d] wi = numpy.arange(s[d], dtype='f8') w.append(wi.reshape(s)) pm = lambda : None pm.w = w pm.comm = None TransferFunction.Laplace(pm, complex) test()