Unverified Commit b9093301 authored by Stephan Kuschel's avatar Stephan Kuschel Committed by GitHub
Browse files

Merge pull request #256 from skuschel/fbpic

Fbpic modeexpansion upgrades. pyfftw is limited to 1 thread to be compatible with multiprocessing.
parents 6efda939 6dce3a7d
......@@ -67,6 +67,8 @@ target/
.project
.pydevproject
*.png
*.swp
# there might be ipython notebooks
.ipynb_checkpoints
......
......@@ -76,6 +76,10 @@ class FieldAnalyzer(object):
e.g. {'theta': [0, pi/2, pi, 3/2*pi]}
'''
name = {0: 'x', 1: 'y', 2: 'z', 90: 'r', 91: 'theta'}[helper.axesidentify[axis]]
Ntheta = kwargs.pop('Ntheta', None)
if Ntheta is not None and name == 'theta':
# linear spacing from fft
return np.linspace(0, 2*np.pi, Ntheta, endpoint=False)
grid = kwargs.pop(name, None)
if grid is not None:
# override with given grid
......
......@@ -32,6 +32,7 @@ from . import Simulationreader_ifc
import numpy as np
import re
from .. import helper
from ..helper_fft import fft
__all__ = ['OpenPMDreader', 'FileSeries',
'FbpicReader', 'FbpicFileSeries']
......@@ -199,7 +200,26 @@ class FbpicReader(OpenPMDreader):
super(FbpicReader, self).__init__(simidentifier, **kwargs)
@staticmethod
def modeexpansion(rawdata, theta=0):
def modeexpansion(rawdata, theta=None, Ntheta=None):
'''
rawdata has to be shaped (Nm, Nr, Nz).
Returns an array of shape (Nr, Ntheta, Nz), with
`Ntheta = (Nm+1)//2`. If Ntheta is given only larger
values are permitted.
The corresponding values for theta are given by
`np.linspace(0, 2*np.pi, Ntheta, endpoint=False)`
'''
rawdata = np.asarray(rawdata)
Nm, Nr, Nz = rawdata.shape
if Ntheta is not None or theta is None:
return FbpicReader._modeexpansion_fft(rawdata, Ntheta=Ntheta)
else:
return FbpicReader._modeexpansion_naiv(rawdata, theta=theta)
@staticmethod
def _modeexpansion_naiv_single(rawdata, theta=0):
'''
The mode representation will be expanded for a given theta.
rawdata has to have the shape (Nm, Nr, Nz).
......@@ -208,7 +228,7 @@ class FbpicReader(OpenPMDreader):
rawdata = np.float64(rawdata)
(Nm, Nr, Nz) = rawdata.shape
mult_above_axis = [1]
for mode in range(1, int(Nm / 2) + 1):
for mode in range(1, (Nm+1)//2):
cos = np.cos(mode * theta)
sin = np.sin(mode * theta)
mult_above_axis += [cos, sin]
......@@ -223,7 +243,7 @@ class FbpicReader(OpenPMDreader):
return F_total
@staticmethod
def _radialdata(rawdata, theta=0):
def _modeexpansion_naiv(rawdata, theta=0):
'''
converts to radial data using `modeexpansion`, possibly for multiple
theta at once.
......@@ -232,11 +252,32 @@ class FbpicReader(OpenPMDreader):
# single theta
theta = [theta]
# multiple theta
data = np.asarray([FbpicReader.modeexpansion(rawdata, theta=t) for t in theta])
data = np.asarray([FbpicReader._modeexpansion_naiv_single(rawdata, theta=t)
for t in theta])
# switch from (theta, r, z) to (r, theta, z)
data = data.swapaxes(0, 1)
return data
@staticmethod
def _modeexpansion_fft(rawdata, Ntheta=None):
'''
calculate the radialdata using an fft. This is by far the fastest
way to do the modeexpansion.
'''
Nm, Nr, Nz = rawdata.shape
Nth = (Nm+1)//2
if Ntheta is None or Ntheta < Nth:
Ntheta = Nth
fd = np.empty((Nr, Ntheta, Nz), dtype=np.complex128)
fd[:, 0, :].real = rawdata[0, :, :]
rawdatasw = np.swapaxes(rawdata, 0, 1)
fd[:, 1:Nth, :].real = rawdatasw[:, 1::2, :]
fd[:, 1:Nth, :].imag = rawdatasw[:, 2::2, :]
fd = fft.fft(fd, axis=1).real
return fd
# override inherited method to count points after mode expansion
def gridoffset(self, key, axis):
axid = helper.axesidentify[axis]
......@@ -265,10 +306,7 @@ class FbpicReader(OpenPMDreader):
# Ntheta does technically not exists because of the mode
# representation. To do a proper conversion from the modes to
# the grid, choose Ntheta based on the number of modes.
# Nyquist should be `Ntheta = Nm`, but be generous and
# make sure that theta and -theta are
# always included: 2, 4, 8, 16, 32, 64,...
Ntheta = 2 * 2**np.ceil(np.log2(Nm))
Ntheta = (Nm + 1) // 2
return (Nr, Ntheta, Nz)[axid]
# override
......@@ -276,22 +314,19 @@ class FbpicReader(OpenPMDreader):
return ('r', 'theta', 'z')
# override from OpenPMDreader
def data(self, key, theta=None):
def data(self, key, **kwargs):
raw = super(FbpicReader, self).data(key) # SI conversion
if key.startswith('particles'):
return raw
# for fields expand the modes into a spatial grid first:
if theta is None:
# guess a proper grid. Number of points defined in `self.gridpoints`
theta = self.grid(key, 'theta')
data = self._radialdata(raw, theta=theta) # modeexpansion
data = self.modeexpansion(raw, **kwargs) # modeexpansion
return data
def dataE(self, component, theta=None, **kwargs):
return self.data(self._keyE(component, **kwargs), theta=theta)
def dataE(self, component, theta=None, Ntheta=None, **kwargs):
return self.data(self._keyE(component, **kwargs), theta=theta, Ntheta=Ntheta)
def dataB(self, component, theta=None, **kwargs):
return self.data(self._keyB(component, **kwargs), theta=theta)
return self.data(self._keyB(component, **kwargs), theta=theta, Ntheta=Ntheta)
# override
def __str__(self):
......
......@@ -59,6 +59,10 @@ class _fft:
import pyfftw.interfaces.numpy_fft as fftw
fftw_cache.enable()
fft_module = fftw
# workaround for
# https://github.com/pyFFTW/pyFFTW/issues/135
# also, scaling is bad, so 1 process wont hurt too much
nproc = 1
fft_kwargs = dict(planner_effort='FFTW_ESTIMATE', threads=nproc)
except ImportError:
# pyFFTW is not available, just import numpys fft
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment