from __future__ import print_function
from __future__ import absolute_import
from numpy.testing import assert_raises, assert_array_equal, assert_allclose
import numpy
import logging
from pmesh.pm import ParticleMesh
pm = ParticleMesh(BoxSize=1.0, Nmesh=(8, 8, 8), dtype='f8', resampler='cic')
from abopt.vmad2 import CodeSegment, logger
from cosmo4d.pmeshengine import ParticleMeshEngine, check_grad, Literal
logger.setLevel(level=logging.WARNING)
def test_compute():
def transfer(k): return 2.0
engine = ParticleMeshEngine(pm)
code = CodeSegment(engine)
code.r2c(real='r', complex='c')
code.transfer(complex='c', tf=transfer)
code.c2r(complex='c', real='r')
code.to_scalar(x='r', y='sum')
field = pm.generate_whitenoise(seed=1234).c2r()
norm = code.compute('sum', init={'r': field})
assert_allclose(norm, field.cnorm() * 4)
def test_vs():
from numpy.testing.decorators import skipif
import numpy
import logging
from pmesh.pm import ParticleMesh
pm = ParticleMesh(BoxSize=1.0, Nmesh=(8, 8, 8), dtype='f8', resampler='cic')
try:
from abopt.vmad2 import CodeSegment, logger
has_abopt = True
except:
has_abopt = False
if has_abopt:
from pmesh.abopt import ParticleMeshEngine, check_grad, Literal
logger.setLevel(level=logging.WARNING)
@skipif(not has_abopt)
def test_compute():
def transfer(k): return 2.0
engine = ParticleMeshEngine(pm)
code = CodeSegment(engine)
code.r2c(real='r', complex='c')
code.transfer(complex='c', tf=transfer)
code.c2r(complex='c', real='r')
code.to_scalar(x='r', y='sum')
field = pm.generate_whitenoise(seed=1234).c2r()
norm = code.compute('sum', init={'r': field})
assert_allclose(norm, field.cnorm() * 4)