def generate_mmodes(self):
"""Calculate the m-modes corresponding to the Timestream.
Perform an MPI transpose for efficiency.
"""
if os.path.exists(self.output_directory + "/mmodes/COMPLETED_M"):
if mpiutil.rank0:
print "******* m-files already generated ********"
return
tel = self.telescope
mmax = tel.mmax
nfreq = tel.nfreq
lfreq, sfreq, efreq = mpiutil.split_local(nfreq)
lm, sm, em = mpiutil.split_local(mmax + 1)
# Load in the local frequencies of the time stream
tstream = np.zeros((lfreq, tel.npairs, self.ntime),
dtype=np.complex128)
for lfi, fi in enumerate(range(sfreq, efreq)):
tstream[lfi] = self.timestream_f(fi)
# FFT to calculate the m-modes for the timestream
row_mmodes = np.fft.fft(tstream, axis=-1) / self.ntime
## Combine positive and negative m parts.
row_mpairs = np.zeros((lfreq, 2, tel.npairs, mmax + 1),
dtype=np.complex128)
row_mpairs[:, 0, ..., 0] = row_mmodes[..., 0]
for mi in range(1, mmax + 1):
row_mpairs[:, 0, ..., mi] = row_mmodes[..., mi]
row_mpairs[:, 1, ..., mi] = row_mmodes[..., -mi].conj()
# Transpose to get the entirety of an m-mode on each process (i.e. all frequencies)
col_mmodes = mpiutil.transpose_blocks(row_mpairs,
(nfreq, 2, tel.npairs, mmax + 1))
# Transpose the local section to make the m's first
col_mmodes = np.transpose(col_mmodes, (3, 0, 1, 2))
for lmi, mi in enumerate(range(sm, em)):
# Make directory for each m-mode
if not os.path.exists(self._mdir(mi)):
os.makedirs(self._mdir(mi))
# Create the m-file and save the result.
with h5py.File(self._mfile(mi), 'w') as f:
f.create_dataset('/mmode', data=col_mmodes[lmi])
f.attrs['m'] = mi
if mpiutil.rank0:
# Make file marker that the m's have been correctly generated:
open(self.output_directory + "/mmodes/COMPLETED_M", 'a').close()
mpiutil.barrier()
def generate_mmodes(self):
"""Calculate the m-modes corresponding to the Timestream.
Perform an MPI transpose for efficiency.
"""
if os.path.exists(self.output_directory + "/mmodes/COMPLETED_M"):
if mpiutil.rank0:
print "******* m-files already generated ********"
return
tel = self.telescope
mmax = tel.mmax
nfreq = tel.nfreq
lfreq, sfreq, efreq = mpiutil.split_local(nfreq)
lm, sm, em = mpiutil.split_local(mmax + 1)
# Load in the local frequencies of the time stream
tstream = np.zeros((lfreq, tel.npairs, self.ntime), dtype=np.complex128)
for lfi, fi in enumerate(range(sfreq, efreq)):
tstream[lfi] = self.timestream_f(fi)
# FFT to calculate the m-modes for the timestream
row_mmodes = np.fft.fft(tstream, axis=-1) / self.ntime
## Combine positive and negative m parts.
row_mpairs = np.zeros((lfreq, 2, tel.npairs, mmax+1), dtype=np.complex128)
row_mpairs[:, 0, ..., 0] = row_mmodes[..., 0]
for mi in range(1, mmax+1):
row_mpairs[:, 0, ..., mi] = row_mmodes[..., mi]
row_mpairs[:, 1, ..., mi] = row_mmodes[..., -mi].conj()
# Transpose to get the entirety of an m-mode on each process (i.e. all frequencies)
col_mmodes = mpiutil.transpose_blocks(row_mpairs, (nfreq, 2, tel.npairs, mmax + 1))
# Transpose the local section to make the m's first
col_mmodes = np.transpose(col_mmodes, (3, 0, 1, 2))
for lmi, mi in enumerate(range(sm, em)):
# Make directory for each m-mode
if not os.path.exists(self._mdir(mi)):
os.makedirs(self._mdir(mi))
# Create the m-file and save the result.
with h5py.File(self._mfile(mi), 'w') as f:
f.create_dataset('/mmode', data=col_mmodes[lmi])
f.attrs['m'] = mi
if mpiutil.rank0:
# Make file marker that the m's have been correctly generated:
open(self.output_directory + "/mmodes/COMPLETED_M", 'a').close()
mpiutil.barrier()
def simulate(beamtransfer,
outdir,
tsname,
maps=[],
ndays=None,
resolution=0,
add_noise=True,
seed=None,
**kwargs):
"""Create a simulated timestream and save it to disk.
Parameters
----------
m : ProductManager object
Products of telescope to simulate.
outdir : directoryname
Directory that we will save the timestream into.
maps : list
List of map filenames. The sum of these form the simulated sky.
ndays : int, optional
Number of days of observation. Setting `ndays = None` (default) uses
the default stored in the telescope object; `ndays = 0`, assumes the
observation time is infinite so that the noise is zero.
resolution : scalar, optional
Approximate time resolution in seconds. Setting `resolution = 0`
(default) calculates the value from the mmax.
Returns
-------
timestream : Timestream
"""
# Create timestream object
tstream = Timestream(outdir, tsname, beamtransfer)
completed_file = tstream._tsdir + '/COMPLETED_TIMESTREAM'
if os.path.exists(completed_file):
if mpiutil.rank0:
print "******* timestream-files already generated ********"
mpiutil.barrier()
return tstream
# Make directory if required
try:
os.makedirs(tstream._tsdir)
except OSError:
# directory exists
pass
if mpiutil.rank0:
# if not os.path.exists(tstream._tsdir):
# os.makedirs(tstream._tsdir)
tstream.save()
## Read in telescope system
bt = beamtransfer
tel = bt.telescope
lmax = tel.lmax
mmax = tel.mmax
nfreq = tel.nfreq
npol = tel.num_pol_sky
projmaps = (len(maps) > 0)
lfreq, sfreq, efreq = mpiutil.split_local(nfreq)
local_freq = range(sfreq, efreq)
lm, sm, em = mpiutil.split_local(mmax + 1)
# If ndays is not set use the default value.
if ndays is None:
ndays = tel.ndays
# Calculate the number of timesamples from the resolution
if resolution == 0:
# Set the minimum resolution required for the sky.
ntime = 2 * mmax + 1
else:
# Set the cl
ntime = int(np.round(24 * 3600.0 / resolution))
col_vis = np.zeros((tel.npairs, lfreq, ntime), dtype=np.complex128)
## If we want to add maps use the m-mode formalism to project a skymap
## into visibility space.
if projmaps:
# Load file to find out the map shapes.
with h5py.File(maps[0], 'r') as f:
mapshape = f['map'].shape
if lfreq > 0:
# Allocate array to store the local frequencies
row_map = np.zeros((lfreq, ) + mapshape[1:], dtype=np.float64)
# Read in and sum up the local frequencies of the supplied maps.
for mapfile in maps:
with h5py.File(mapfile, 'r') as f:
row_map += f['map'][sfreq:efreq]
# Calculate the alm's for the local sections
row_alm = hputil.sphtrans_sky(row_map, lmax=lmax).reshape(
(lfreq, npol * (lmax + 1), lmax + 1))
else:
row_alm = np.zeros((lfreq, npol * (lmax + 1), lmax + 1),
dtype=np.complex128)
# Perform the transposition to distribute different m's across processes. Neat
# tip, putting a shorter value for the number of columns, trims the array at
# the same time
col_alm = mpiutil.transpose_blocks(row_alm, (nfreq, npol *
(lmax + 1), mmax + 1))
# Transpose and reshape to shift m index first.
col_alm = np.transpose(col_alm,
(2, 0, 1)).reshape(lm, nfreq, npol, lmax + 1)
# Create storage for visibility data
vis_data = np.zeros((lm, nfreq, bt.ntel), dtype=np.complex128)
# Iterate over m's local to this process and generate the corresponding
# visibilities
for mp, mi in enumerate(range(sm, em)):
vis_data[mp] = bt.project_vector_sky_to_telescope(mi, col_alm[mp])
# Rearrange axes such that frequency is last (as we want to divide
# frequencies across processors)
row_vis = vis_data.transpose(
(0, 2, 1)) #.reshape((lm * bt.ntel, nfreq))
# Parallel transpose to get all m's back onto the same processor
col_vis_tmp = mpiutil.transpose_blocks(row_vis,
((mmax + 1), bt.ntel, nfreq))
col_vis_tmp = col_vis_tmp.reshape(mmax + 1, 2, tel.npairs, lfreq)
# Transpose the local section to make the m's the last axis and unwrap the
# positive and negative m at the same time.
col_vis[..., 0] = col_vis_tmp[0, 0]
for mi in range(1, mmax + 1):
col_vis[..., mi] = col_vis_tmp[mi, 0]
col_vis[..., -mi] = col_vis_tmp[
mi, 1].conj() # Conjugate only (not (-1)**m - see paper)
del col_vis_tmp
## If we're simulating noise, create a realisation and add it to col_vis
if ndays > 0:
# Fetch the noise powerspectrum
noise_ps = tel.noisepower(np.arange(tel.npairs)[:, np.newaxis],
np.array(local_freq)[np.newaxis, :],
ndays=ndays).reshape(tel.npairs,
lfreq)[:, :, np.newaxis]
# Seed random number generator to give consistent noise
if seed is not None:
# Must include rank such that we don't have massive power deficit from correlated noise
np.random.seed(seed + mpiutil.rank)
# Create and weight complex noise coefficients
noise_vis = (np.array([1.0, 1.0J]) *
np.random.standard_normal(col_vis.shape +
(2, ))).sum(axis=-1)
noise_vis *= (noise_ps / 2.0)**0.5
# Reset RNG
if seed is not None:
np.random.seed()
# Add into main noise sims
col_vis += noise_vis
del noise_vis
# Fourier transform m-modes back to get timestream.
vis_stream = np.fft.ifft(col_vis, axis=-1) * ntime
vis_stream = vis_stream.reshape(tel.npairs, lfreq, ntime)
# The time samples the visibility is calculated at
tphi = np.linspace(0, 2 * np.pi, ntime, endpoint=False)
# Create timestream object
tstream = Timestream(outdir, m)
## Iterate over the local frequencies and write them to disk.
for lfi, fi in enumerate(local_freq):
# Make directory if required
if not os.path.exists(tstream._fdir(fi)):
os.makedirs(tstream._fdir(fi))
# Write file contents
with h5py.File(tstream._ffile(fi), 'w') as f:
# Timestream data
f.create_dataset('/timestream', data=vis_stream[:, lfi])
f.create_dataset('/phi', data=tphi)
# Telescope layout data
f.create_dataset('/feedmap', data=tel.feedmap)
f.create_dataset('/feedconj', data=tel.feedconj)
f.create_dataset('/feedmask', data=tel.feedmask)
f.create_dataset('/uniquepairs', data=tel.uniquepairs)
f.create_dataset('/baselines', data=tel.baselines)
# Write metadata
f.attrs['beamtransfer_path'] = os.path.abspath(bt.directory)
f.attrs['ntime'] = ntime
mpiutil.barrier()
return tstream
def redistribute(self, axis):
"""Change the axis that the array is distributed over.
Parameters
----------
axis : integer
Axis to distribute over.
Returns
-------
array : MPIArray
A new copy of the array distributed over the specified axis. Note
that the local section will have changed.
"""
# Check to see if this is the current distributed axis
if self.axis == axis or self.comm is None:
return self
# Test to see if the datatype is one understood by MPI, this can
# probably be fixed up at somepoint by creating a datatype of the right
# number of bytes
try:
mpiutil.typemap(self.dtype)
except KeyError:
if self.comm is None or self.comm.rank == 0:
import warnings
warnings.warn('Cannot redistribute array of compound datatypes. Sorry!!')
return self
# Construct the list of the axes to swap around
axlist_f = list(range(len(self.shape)))
# Remove the axes we are going to swap around
axlist_f.remove(self.axis)
axlist_f.remove(axis)
# Move the current dist axis to the front, and the new to the end
axlist_f.insert(0, self.axis)
axlist_f.append(axis)
# Perform a local transpose on the array to get the axes in the correct order
trans_arr = self.view(np.ndarray).transpose(axlist_f).copy()
# Perform the global transpose
tmp_gshape = (self.global_shape[self.axis],) + trans_arr.shape[1:]
trans_arr = mpiutil.transpose_blocks(trans_arr, tmp_gshape, comm=self.comm)
axlist_b = list(range(len(self.shape)))
axlist_b.pop(0)
last = axlist_b.pop(-1)
if self.axis < axis: # This has to awkwardly depend on the order of the axes
axlist_b.insert(self.axis, 0)
axlist_b.insert(axis, last)
else:
axlist_b.insert(axis, last)
axlist_b.insert(self.axis, 0)
# Perform the local transpose to get the axes back in the correct order
trans_arr = trans_arr.transpose(axlist_b)
# Create a new MPIArray object out of the data
dist_arr = MPIArray(self.global_shape, axis=axis, dtype=self.dtype, comm=self.comm)
dist_arr[:] = trans_arr
return dist_arr
def simulate(m, outdir, maps=[], ndays=None, resolution=0, seed=None, **kwargs):
"""Create a simulated timestream and save it to disk.
Parameters
----------
m : ProductManager object
Products of telescope to simulate.
outdir : directoryname
Directory that we will save the timestream into.
maps : list
List of map filenames. The sum of these form the simulated sky.
ndays : int, optional
Number of days of observation. Setting `ndays = None` (default) uses
the default stored in the telescope object; `ndays = 0`, assumes the
observation time is infinite so that the noise is zero.
resolution : scalar, optional
Approximate time resolution in seconds. Setting `resolution = 0`
(default) calculates the value from the mmax.
Returns
-------
timestream : Timestream
"""
## Read in telescope system
bt = m.beamtransfer
tel = bt.telescope
lmax = tel.lmax
mmax = tel.mmax
nfreq = tel.nfreq
npol = tel.num_pol_sky
projmaps = (len(maps) > 0)
lfreq, sfreq, efreq = mpiutil.split_local(nfreq)
local_freq = range(sfreq, efreq)
lm, sm, em = mpiutil.split_local(mmax + 1)
# If ndays is not set use the default value.
if ndays is None:
ndays = tel.ndays
# Calculate the number of timesamples from the resolution
if resolution == 0:
# Set the minimum resolution required for the sky.
ntime = 2*mmax+1
else:
# Set the cl
ntime = int(np.round(24 * 3600.0 / resolution))
col_vis = np.zeros((tel.npairs, lfreq, ntime), dtype=np.complex128)
## If we want to add maps use the m-mode formalism to project a skymap
## into visibility space.
if projmaps:
# Load file to find out the map shapes.
with h5py.File(maps[0], 'r') as f:
mapshape = f['map'].shape
if lfreq > 0:
# Allocate array to store the local frequencies
row_map = np.zeros((lfreq,) + mapshape[1:], dtype=np.float64)
# Read in and sum up the local frequencies of the supplied maps.
for mapfile in maps:
with h5py.File(mapfile, 'r') as f:
row_map += f['map'][sfreq:efreq]
# Calculate the alm's for the local sections
row_alm = hputil.sphtrans_sky(row_map, lmax=lmax).reshape((lfreq, npol * (lmax+1), lmax+1))
else:
row_alm = np.zeros((lfreq, npol * (lmax+1), lmax+1), dtype=np.complex128)
# Perform the transposition to distribute different m's across processes. Neat
# tip, putting a shorter value for the number of columns, trims the array at
# the same time
col_alm = mpiutil.transpose_blocks(row_alm, (nfreq, npol * (lmax+1), mmax+1))
# Transpose and reshape to shift m index first.
col_alm = np.transpose(col_alm, (2, 0, 1)).reshape(lm, nfreq, npol, lmax+1)
# Create storage for visibility data
vis_data = np.zeros((lm, nfreq, bt.ntel), dtype=np.complex128)
# Iterate over m's local to this process and generate the corresponding
# visibilities
for mp, mi in enumerate(range(sm, em)):
vis_data[mp] = bt.project_vector_sky_to_telescope(mi, col_alm[mp])
# Rearrange axes such that frequency is last (as we want to divide
# frequencies across processors)
row_vis = vis_data.transpose((0, 2, 1))#.reshape((lm * bt.ntel, nfreq))
# Parallel transpose to get all m's back onto the same processor
col_vis_tmp = mpiutil.transpose_blocks(row_vis, ((mmax+1), bt.ntel, nfreq))
col_vis_tmp = col_vis_tmp.reshape(mmax + 1, 2, tel.npairs, lfreq)
# Transpose the local section to make the m's the last axis and unwrap the
# positive and negative m at the same time.
col_vis[..., 0] = col_vis_tmp[0, 0]
for mi in range(1, mmax+1):
col_vis[..., mi] = col_vis_tmp[mi, 0]
col_vis[..., -mi] = col_vis_tmp[mi, 1].conj() # Conjugate only (not (-1)**m - see paper)
del col_vis_tmp
## If we're simulating noise, create a realisation and add it to col_vis
if ndays > 0:
# Fetch the noise powerspectrum
noise_ps = tel.noisepower(np.arange(tel.npairs)[:, np.newaxis], np.array(local_freq)[np.newaxis, :], ndays=ndays).reshape(tel.npairs, lfreq)[:, :, np.newaxis]
# Seed random number generator to give consistent noise
if seed is not None:
# Must include rank such that we don't have massive power deficit from correlated noise
np.random.seed(seed + mpiutil.rank)
# Create and weight complex noise coefficients
noise_vis = (np.array([1.0, 1.0J]) * np.random.standard_normal(col_vis.shape + (2,))).sum(axis=-1)
noise_vis *= (noise_ps / 2.0)**0.5
# Reset RNG
if seed is not None:
np.random.seed()
# Add into main noise sims
col_vis += noise_vis
del noise_vis
# Fourier transform m-modes back to get timestream.
vis_stream = np.fft.ifft(col_vis, axis=-1) * ntime
vis_stream = vis_stream.reshape(tel.npairs, lfreq, ntime)
# The time samples the visibility is calculated at
tphi = np.linspace(0, 2*np.pi, ntime, endpoint=False)
# Create timestream object
tstream = Timestream(outdir, m)
## Iterate over the local frequencies and write them to disk.
for lfi, fi in enumerate(local_freq):
# Make directory if required
if not os.path.exists(tstream._fdir(fi)):
os.makedirs(tstream._fdir(fi))
# Write file contents
with h5py.File(tstream._ffile(fi), 'w') as f:
# Timestream data
f.create_dataset('/timestream', data=vis_stream[:, lfi])
f.create_dataset('/phi', data=tphi)
# Telescope layout data
f.create_dataset('/feedmap', data=tel.feedmap)
f.create_dataset('/feedconj', data=tel.feedconj)
f.create_dataset('/feedmask', data=tel.feedmask)
f.create_dataset('/uniquepairs', data=tel.uniquepairs)
f.create_dataset('/baselines', data=tel.baselines)
# Write metadata
f.attrs['beamtransfer_path'] = os.path.abspath(bt.directory)
f.attrs['ntime'] = ntime
tstream.save()
mpiutil.barrier()
return tstream
def redistribute(self, axis):
"""Change the axis that the array is distributed over.
Parameters
----------
axis : integer
Axis to distribute over.
Returns
-------
array : MPIArray
A new copy of the array distributed over the specified axis. Note
that the local section will have changed.
"""
# Check to see if this is the current distributed axis
if self.axis == axis or self.comm is None:
return self
# Test to see if the datatype is one understood by MPI, this can
# probably be fixed up at somepoint by creating a datatype of the right
# number of bytes
try:
mpiutil.typemap(self.dtype)
except KeyError:
if self.comm is None or self.comm.rank == 0:
import warnings
warnings.warn(
'Cannot redistribute array of compound datatypes. Sorry!!')
return self
# Construct the list of the axes to swap around
axlist_f = list(range(len(self.shape)))
# Remove the axes we are going to swap around
axlist_f.remove(self.axis)
axlist_f.remove(axis)
# Move the current dist axis to the front, and the new to the end
axlist_f.insert(0, self.axis)
axlist_f.append(axis)
# Perform a local transpose on the array to get the axes in the correct order
trans_arr = self.view(np.ndarray).transpose(axlist_f).copy()
# Perform the global transpose
tmp_gshape = (self.global_shape[self.axis], ) + trans_arr.shape[1:]
trans_arr = mpiutil.transpose_blocks(trans_arr,
tmp_gshape,
comm=self.comm)
axlist_b = list(range(len(self.shape)))
axlist_b.pop(0)
last = axlist_b.pop(-1)
if self.axis < axis: # This has to awkwardly depend on the order of the axes
axlist_b.insert(self.axis, 0)
axlist_b.insert(axis, last)
else:
axlist_b.insert(axis, last)
axlist_b.insert(self.axis, 0)
# Perform the local transpose to get the axes back in the correct order
trans_arr = trans_arr.transpose(axlist_b)
# Create a new MPIArray object out of the data
dist_arr = MPIArray(self.global_shape,
axis=axis,
dtype=self.dtype,
comm=self.comm)
dist_arr[:] = trans_arr
return dist_arr