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 __new__(cls, global_shape, axis=0, comm=None, *args, **kwargs):
# if mpiutil.world is None:
# raise RuntimeError('There is no mpi4py installation. Aborting.')
if comm is None:
comm = mpiutil.world
# Determine local section of distributed axis
local_num, local_start, local_end = mpiutil.split_local(
global_shape[axis], comm=comm)
# Figure out the local shape and offset
lshape = list(global_shape)
lshape[axis] = local_num
loffset = [0] * len(global_shape)
loffset[axis] = local_start
# Create array
arr = np.ndarray.__new__(cls, lshape, *args, **kwargs)
# Set attributes of class
arr._global_shape = global_shape
arr._axis = axis
arr._local_shape = tuple(lshape)
arr._local_offset = tuple(loffset)
arr._comm = comm
return arr
def test_H5FileSelect_distributed(container_on_disk, fsel, isel):
"""Load H5 into parallel container while down-selecting axes."""
sel = {"dset1": (fsel, isel, slice(None)), "dset2": (fsel, slice(None))}
# Tests are designed to run for 1, 2 or 4 processes
assert 4 % comm.size == 0
m = MemGroup.from_hdf5(container_on_disk,
selections=sel,
distributed=True,
comm=comm)
d1 = dset1[(fsel, isel, slice(None))]
d2 = dset2[(fsel, slice(None))]
n, s, e = mpiutil.split_local(d1.shape[0], comm=comm)
dslice = slice(s, e)
# For debugging...
# Need to dereference datasets as this is collective
# md1 = m["dset1"][:]
# md2 = m["dset2"][:]
# for ri in range(comm.size):
# if ri == comm.rank:
# print(comm.rank)
# print(md1.shape, d1.shape, d1[dslice].shape)
# print(md1[0, :2, :2] if md1.size else "Empty")
# print(d1[dslice][0, :2, :2] if d1[dslice].size else "Empty")
# print()
# comm.Barrier()
assert np.all(m["dset1"][:] == d1[dslice])
assert np.all(m["dset2"][:] == d2[dslice])
def generate_mmodes(self, ts_data=None):
"""Calculate the m-modes corresponding to the Timestream.
Perform an MPI transpose for efficiency.
"""
completed_file = self._mdir + 'COMPLETED_M'
if os.path.exists(completed_file):
if mpiutil.rank0:
print "******* m-files already generated ********"
mpiutil.barrier()
return
# Make directory if required
# if mpiutil.rank0 and not os.path.exists(self._mdir):
# os.makedirs(self._mdir)
try:
os.makedirs(self._mdir)
except OSError:
# directory exists
pass
tel = self.telescope
mmax = tel.mmax
ntime = ts_data.shape[0] if ts_data is not None else self.ntime
nbl = tel.nbase
nfreq = tel.nfreq
indices = list(itertools.product(np.arange(nfreq), np.arange(nbl)))
lind, sind, eind = mpiutil.split_local(nfreq * nbl)
# load the local section of the time stream
tstream = np.zeros((ntime, lind), dtype=np.complex128)
for ind, (f_ind, bl_ind) in enumerate(indices[sind:eind]):
if ts_data is not None:
tstream[:, ind] = ts_data[:, f_ind, bl_ind]
else:
with h5py.File(self._tsfile, 'r') as f:
tstream[:, ind] = f['/timestream'][:, f_ind, bl_ind]
# FFT to get m-mode
mmodes = np.fft.fft(tstream, axis=0) / ntime # m = 0 is at left
mmodes = MPIArray.wrap(mmodes, axis=1)
# redistribute along different m
mmodes = mmodes.redistribute(axis=0)
# save m-modes to file
ms = np.concatenate([np.arange(0, mmax+1), np.arange(-mmax, 0)])
for ind, mi in enumerate(mpiutil.mpilist(ms, method='con')):
with h5py.File(self._mfile(mi), 'w') as f:
f.create_dataset('/mmode', data=mmodes[ind].view(np.ndarray).reshape(nfreq, nbl))
f.attrs['m'] = mi
mpiutil.barrier()
if mpiutil.rank0:
# Make file marker that the m's have been correctly generated:
open(completed_file, 'a').close()
def __new__(cls, global_shape, axis=0, comm=None, *args, **kwargs):
# if mpiutil.world is None:
# raise RuntimeError('There is no mpi4py installation. Aborting.')
if comm is None:
comm = mpiutil.world
# Determine local section of distributed axis
local_num, local_start, local_end = mpiutil.split_local(global_shape[axis], comm=comm)
# Figure out the local shape and offset
lshape = list(global_shape)
lshape[axis] = local_num
loffset = [0] * len(global_shape)
loffset[axis] = local_start
# Create array
arr = np.ndarray.__new__(cls, lshape, *args, **kwargs)
# Set attributes of class
arr._global_shape = global_shape
arr._axis = axis
arr._local_shape = tuple(lshape)
arr._local_offset = tuple(loffset)
arr._comm = comm
return arr
def make_clzz_array(self):
p_bands, s_bands, e_bands = mpiutil.split_all(self.nbands)
p, s, e = mpiutil.split_local(self.nbands)
self.clarray = np.zeros(
(
self.nbands,
self.telescope.lmax + 1,
self.telescope.nfreq,
self.telescope.nfreq,
),
dtype=np.float64,
)
for bi in range(s, e):
self.clarray[bi] = self.make_clzz(self.band_pk[bi])
bandsize = ((self.telescope.lmax + 1) * self.telescope.nfreq *
self.telescope.nfreq)
sizes = p_bands * bandsize
displ = s_bands * bandsize
MPI.COMM_WORLD.Allgatherv(MPI.IN_PLACE,
[self.clarray, sizes, displ, MPI.DOUBLE])
def test_io(self):
import h5py
# Cleanup directories
fname = 'testdset.hdf5'
if mpiutil.rank0 and os.path.exists(fname):
os.remove(fname)
mpiutil.barrier()
gshape = (19, 17)
ds = mpiarray.MPIArray(gshape, dtype=np.int64)
ga = np.arange(np.prod(gshape)).reshape(gshape)
l0, s0, e0 = mpiutil.split_local(gshape[0])
ds[:] = ga[s0:e0]
ds.redistribute(axis=1).to_hdf5(fname, 'testds', create=True)
if mpiutil.rank0:
with h5py.File(fname, 'r') as f:
h5ds = f['testds'][:]
assert (h5ds == ga).all()
ds2 = mpiarray.MPIArray.from_hdf5(fname, 'testds')
assert (ds2 == ds).all()
def wrap(cls, array, axis, comm=None):
"""Turn a set of numpy arrays into a distributed MPIArray object.
This is needed for functions such as `np.fft.fft` which always return
an `np.ndarray`.
Parameters
----------
array : np.ndarray
Array to wrap.
axis : integer
Axis over which the array is distributed. The lengths are checked
to try and ensure this is correct.
comm : MPI.Comm, optional
The communicator over which the array is distributed. If `None`
(default), use `MPI.COMM_WORLD`.
Returns
-------
dist_array : MPIArray
An MPIArray view of the input.
"""
# from mpi4py import MPI
if comm is None:
comm = mpiutil.world
# Get axis length, both locally, and globally
axlen = array.shape[axis]
totallen = mpiutil.allreduce(axlen, comm=comm)
# Figure out what the distributed layout should be
local_num, local_start, local_end = mpiutil.split_local(totallen, comm=comm)
# Check the local layout is consistent with what we expect, and send
# result to all ranks
layout_issue = mpiutil.allreduce(axlen != local_num, op=mpiutil.MAX, comm=comm)
if layout_issue:
raise Exception("Cannot wrap, distributed axis local length is incorrect.")
# Set shape and offset
lshape = array.shape
global_shape = list(lshape)
global_shape[axis] = totallen
loffset = [0] * len(lshape)
loffset[axis] = local_start
# Setup attributes of class
dist_arr = array.view(cls)
dist_arr._global_shape = tuple(global_shape)
dist_arr._axis = axis
dist_arr._local_shape = tuple(lshape)
dist_arr._local_offset = tuple(loffset)
dist_arr._comm = comm
return dist_arr
def test_redistribution(self):
gshape = (1, 11, 2, 14, 3, 4)
nelem = np.prod(gshape)
garr = np.arange(nelem).reshape(gshape)
l0, s0, e0 = mpiutil.split_local(11)
l1, s1, e1 = mpiutil.split_local(14)
l2, s2, e2 = mpiutil.split_local(4)
arr = mpiarray.MPIArray(gshape, axis=1, dtype=np.int64)
arr[:] = garr[:, s0:e0]
arr2 = arr.redistribute(axis=3)
assert (arr2 == garr[:, :, :, s1:e1]).view(np.ndarray).all()
arr3 = arr.redistribute(axis=5)
assert (arr3 == garr[:, :, :, :, :, s2:e2]).view(np.ndarray).all()
def test_redistribution(self):
gshape = (1, 11, 2, 14, 3, 4)
nelem = np.prod(gshape)
garr = np.arange(nelem).reshape(gshape)
l0, s0, e0 = mpiutil.split_local(11)
l1, s1, e1 = mpiutil.split_local(14)
l2, s2, e2 = mpiutil.split_local(4)
arr = mpiarray.MPIArray(gshape, axis=1, dtype=np.int64)
arr[:] = garr[:, s0:e0]
arr2 = arr.redistribute(axis=3)
assert (arr2 == garr[:, :, :, s1:e1]).view(np.ndarray).all()
arr3 = arr.redistribute(axis=5)
assert (arr3 == garr[:, :, :, :, :, s2:e2]).view(np.ndarray).all()
def test_construction(self):
arr = mpiarray.MPIArray((10, 11), axis=1)
l, s, e = mpiutil.split_local(11)
# Check that global shape is set correctly
assert arr.global_shape == (10, 11)
assert arr.shape == (10, l)
assert arr.local_offset == (0, s)
assert arr.local_shape == (10, l)
def test_construction(self):
arr = mpiarray.MPIArray((10, 11), axis=1)
l, s, e = mpiutil.split_local(11)
# Check that global shape is set correctly
assert arr.global_shape == (10, 11)
assert arr.shape == (10, l)
assert arr.local_offset == (0, s)
assert arr.local_shape == (10, l)
def make_clzz_array(self):
p_bands, s_bands, e_bands = mpiutil.split_all(self.nbands)
p, s, e = mpiutil.split_local(self.nbands)
self.clarray = np.zeros((self.nbands, self.telescope.lmax + 1,
self.telescope.nfreq, self.telescope.nfreq), dtype=np.float64)
for bi in range(s, e):
self.clarray[bi] = self.make_clzz(self.band_pk[bi])
bandsize = (self.telescope.lmax + 1) * self.telescope.nfreq * self.telescope.nfreq
sizes = p_bands * bandsize
displ = s_bands * bandsize
MPI.COMM_WORLD.Allgatherv(MPI.IN_PLACE, [self.clarray, sizes, displ, MPI.DOUBLE])
def test_LoadBasicCont_selection(ss_container, mpi_tmp_path):
fname = str(mpi_tmp_path / "ss.h5")
ss_container.save(fname)
freq_range = [1, 4]
ra_range = [3, 10]
task = io.LoadBasicCont()
task.files = [fname]
task.selections = {
"freq_range": freq_range,
"ra_range": ra_range,
}
task.setup()
ss_load = task.next()
ss_vis = ss_load.vis[:]
ss_weight = ss_load.weight[:]
# Check the datasets (only a maximum of three ranks will have anything)
nf = freq_range[1] - freq_range[0]
if ss_load.comm.rank < nf:
# Check the freq selection
n, s, e = mpiutil.split_local(nf, comm=ss_load.comm)
vis_real = np.arange(*freq_range)[s:e][:, np.newaxis, np.newaxis]
assert (ss_vis.real == vis_real).all()
# Check the ra selection
vis_imag = np.arange(*ra_range)[np.newaxis, np.newaxis, :]
assert (ss_vis.imag == vis_imag).all()
# Check that nothing funky happened on the stack axis
weight = np.arange(ss_container.vis.shape[1])[np.newaxis, :,
np.newaxis]
assert (ss_weight == weight).all()
# Check the attributes...
assert ss_load.attrs["test_attr1"] == "hello"
assert ss_load.vis.attrs["test_attr2"] == "hello2"
assert ss_load.weight.attrs["test_attr3"] == "hello3"
# As we only put one item into the queue, this should end the iterations
with pytest.raises(pipeline.PipelineStopIteration):
task.next()
def transfer_matrices(self, bl_indices, f_indices):
"""Calculate the spherical harmonic transfer matrices for baseline and
frequency combinations.
Parameters
----------
bl_indices : array_like
Indices of baselines to calculate.
f_indices : array_like
Indices of frequencies to calculate. Must be broadcastable against
`bl_indices`.
Returns
-------
transfer : np.ndarray, dtype=np.complex128
An array containing the transfer functions. The shape is somewhat
complicated, the first indices correspond to the broadcast size of
`bl_indices` and `f_indices`, then there may be some polarisation
indices, then finally the (l,m) indices, range (lside, 2*lside-1).
"""
## Check indices are all in range
if out_of_range(bl_indices, 0, self.npairs):
raise Exception("Baseline indices aren't valid")
if out_of_range(f_indices, 0, self.nfreq):
raise Exception("Frequency indices aren't valid")
# Generate the array for the Transfer functions
nfreq = f_indices.size
nbl = bl_indices.size
nprod = nfreq * nbl
indices = list(itertools.product(f_indices, bl_indices))
lind, sind, eind = mpiutil.split_local(nprod)
# local section of the array
tshape = (lind, self.num_pol_sky, self.theta_size, self.phi_size)
tarray = np.zeros(tshape, dtype=np.complex128)
for ind, (f_ind, bl_ind) in enumerate(indices[sind:eind]):
tarray[ind] = self._transfer_single(bl_ind, f_ind, self.lmax)
return tarray
def transfer_matrices(self, bl_indices, f_indices):
"""Calculate the spherical harmonic transfer matrices for baseline and
frequency combinations.
Parameters
----------
bl_indices : array_like
Indices of baselines to calculate.
f_indices : array_like
Indices of frequencies to calculate. Must be broadcastable against
`bl_indices`.
Returns
-------
transfer : np.ndarray, dtype=np.complex128
An array containing the transfer functions. The shape is somewhat
complicated, the first indices correspond to the broadcast size of
`bl_indices` and `f_indices`, then there may be some polarisation
indices, then finally the (l,m) indices, range (lside, 2*lside-1).
"""
## Check indices are all in range
if out_of_range(bl_indices, 0, self.npairs):
raise Exception("Baseline indices aren't valid")
if out_of_range(f_indices, 0, self.nfreq):
raise Exception("Frequency indices aren't valid")
# Generate the array for the Transfer functions
nfreq = f_indices.size
nbl = bl_indices.size
nprod = nfreq * nbl
indices = list(itertools.product(f_indices, bl_indices))
lind, sind, eind = mpiutil.split_local(nprod)
# local section of the array
tshape = (lind, self.num_pol_sky, self.theta_size, self.phi_size)
tarray = np.zeros(tshape, dtype=np.complex128)
for ind, (f_ind, bl_ind) in enumerate(indices[sind:eind]):
tarray[ind] = self._transfer_single(bl_ind, f_ind, self.lmax)
return tarray
def test_wrap(self):
ds = mpiarray.MPIArray((10, 17))
df = np.fft.rfft(ds, axis=1)
assert type(df) == np.ndarray
da = mpiarray.MPIArray.wrap(df, axis=0)
assert type(da) == mpiarray.MPIArray
assert da.global_shape == (10, 9)
l0, s0, e0 = mpiutil.split_local(10)
assert da.local_shape == (l0, 9)
if mpiutil.rank0:
df = df[:-1]
with self.assertRaises(Exception):
mpiarray.MPIArray.wrap(df, axis=0)
def test_io(self):
import h5py
# Cleanup directories
fname = 'testdset.hdf5'
if mpiutil.rank0 and os.path.exists(fname):
os.remove(fname)
mpiutil.barrier()
gshape = (19, 17)
ds = mpiarray.MPIArray(gshape, dtype=np.int64)
ga = np.arange(np.prod(gshape)).reshape(gshape)
l0, s0, e0 = mpiutil.split_local(gshape[0])
ds[:] = ga[s0:e0]
ds.redistribute(axis=1).to_hdf5(fname, 'testds', create=True)
if mpiutil.rank0:
with h5py.File(fname, 'r') as f:
h5ds = f['testds'][:]
assert (h5ds == ga).all()
ds2 = mpiarray.MPIArray.from_hdf5(fname, 'testds')
assert (ds2 == ds).all()
mpiutil.barrier()
if mpiutil.rank0 and os.path.exists(fname):
os.remove(fname)
def test_wrap(self):
ds = mpiarray.MPIArray((10, 17))
df = np.fft.rfft(ds, axis=1)
assert type(df) == np.ndarray
da = mpiarray.MPIArray.wrap(df, axis=0)
assert type(da) == mpiarray.MPIArray
assert da.global_shape == (10, 9)
l0, s0, e0 = mpiutil.split_local(10)
assert da.local_shape == (l0, 9)
if mpiutil.rank0:
df = df[:-1]
if mpiutil.size > 1:
with self.assertRaises(Exception):
mpiarray.MPIArray.wrap(df, axis=0)
def test_reshape(self):
gshape = (1, 11, 2, 14)
l0, s0, e0 = mpiutil.split_local(11)
arr = mpiarray.MPIArray(gshape, axis=1, dtype=np.int64)
arr2 = arr.reshape((None, 28))
# Check type
assert isinstance(arr2, mpiarray.MPIArray)
# Check global shape
assert arr2.global_shape == (11, 28)
# Check local shape
assert arr2.local_shape == (l0, 28)
# Check local offset
assert arr2.local_offset == (s0, 0)
# Check axis
assert arr2.axis == 0
def test_reshape(self):
gshape = (1, 11, 2, 14)
l0, s0, e0 = mpiutil.split_local(11)
arr = mpiarray.MPIArray(gshape, axis=1, dtype=np.int64)
arr2 = arr.reshape((None, 28))
# Check type
assert isinstance(arr2, mpiarray.MPIArray)
# Check global shape
assert arr2.global_shape == (11, 28)
# Check local shape
assert arr2.local_shape == (l0, 28)
# Check local offset
assert arr2.local_offset == (s0, 0)
# Check axis
assert arr2.axis == 0
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 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 process(self, map_):
"""Simulate a SiderealStream
Parameters
----------
map : :class:`containers.Map`
The sky map to process to into a sidereal stream. Frequencies in the
map, must match the Beam Transfer matrices.
Returns
-------
ss : SiderealStream
Stacked sidereal day.
feeds : list of CorrInput
Description of the feeds simulated.
"""
if self.done:
raise pipeline.PipelineStopIteration
# Read in telescope system
bt = self.beamtransfer
tel = self.telescope
lmax = tel.lmax
mmax = tel.mmax
nfreq = tel.nfreq
npol = tel.num_pol_sky
lfreq, sfreq, efreq = mpiutil.split_local(nfreq)
lm, sm, em = mpiutil.split_local(mmax + 1)
# Set the minimum resolution required for the sky.
ntime = 2 * mmax + 1
freqmap = map_.index_map["freq"][:]
row_map = map_.map[:]
if (tel.frequencies != freqmap["centre"]).any():
raise ValueError(
"Frequencies in map do not match those in Beam Transfers.")
# Calculate the alm's for the local sections
row_alm = hputil.sphtrans_sky(row_map, lmax=lmax).reshape(
(lfreq, npol * (lmax + 1), lmax + 1))
# Trim off excess m's and wrap into MPIArray
row_alm = row_alm[..., :(mmax + 1)]
row_alm = mpiarray.MPIArray.wrap(row_alm, axis=0)
# 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 = row_alm.redistribute(axis=2)
# Transpose and reshape to shift m index first.
col_alm = col_alm.transpose((2, 0, 1)).reshape(
(None, nfreq, npol, lmax + 1))
# Create storage for visibility data
vis_data = mpiarray.MPIArray((mmax + 1, nfreq, bt.ntel),
axis=0,
dtype=np.complex128)
vis_data[:] = 0.0
# Iterate over m's local to this process and generate the corresponding
# visibilities
for mp, mi in vis_data.enumerate(axis=0):
vis_data[mp] = bt.project_vector_sky_to_telescope(
mi, col_alm[mp].view(np.ndarray))
# Rearrange axes such that frequency is last (as we want to divide
# frequencies across processors)
row_vis = vis_data.transpose((0, 2, 1))
# Parallel transpose to get all m's back onto the same processor
col_vis_tmp = row_vis.redistribute(axis=2)
col_vis_tmp = col_vis_tmp.reshape((mmax + 1, 2, tel.npairs, None))
# 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 = mpiarray.MPIArray((tel.npairs, nfreq, ntime),
axis=1,
dtype=np.complex128)
col_vis[:] = 0.0
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
# Fourier transform m-modes back to get final timestream.
vis_stream = np.fft.ifft(col_vis, axis=-1) * ntime
vis_stream = vis_stream.reshape((tel.npairs, lfreq, ntime))
vis_stream = vis_stream.transpose((1, 0, 2)).copy()
# Try and fetch out the feed index and info from the telescope object.
try:
feed_index = tel.input_index
except AttributeError:
feed_index = tel.nfeed
# Construct a product map
prod_map = np.zeros(tel.uniquepairs.shape[0],
dtype=[("input_a", int), ("input_b", int)])
prod_map["input_a"] = tel.uniquepairs[:, 0]
prod_map["input_b"] = tel.uniquepairs[:, 1]
# Construct container and set visibility data
sstream = containers.SiderealStream(
freq=freqmap,
ra=ntime,
input=feed_index,
prod=prod_map,
distributed=True,
comm=map_.comm,
)
sstream.vis[:] = mpiarray.MPIArray.wrap(vis_stream, axis=0)
sstream.weight[:] = 1.0
self.done = True
return sstream
def test_io(self):
import h5py
# Cleanup directories
fname = 'testdset.hdf5'
if mpiutil.rank0 and os.path.exists(fname):
os.remove(fname)
mpiutil.barrier()
gshape = (19, 17)
ds = mpiarray.MPIArray(gshape, dtype=np.int64)
ga = np.arange(np.prod(gshape)).reshape(gshape)
l0, s0, e0 = mpiutil.split_local(gshape[0])
ds[:] = ga[s0:e0]
ds.redistribute(axis=1).to_hdf5(fname, 'testds', create=True)
if mpiutil.rank0:
with h5py.File(fname, 'r') as f:
h5ds = f['testds'][:]
assert (h5ds == ga).all()
ds2 = mpiarray.MPIArray.from_hdf5(fname, 'testds')
assert (ds2 == ds).all()
mpiutil.barrier()
# Check that reading over another distributed axis works
ds3 = mpiarray.MPIArray.from_hdf5(fname, 'testds', axis=1)
assert ds3.shape[0] == gshape[0]
assert ds3.shape[1] == mpiutil.split_local(gshape[1])[0]
ds3 = ds3.redistribute(axis=0)
assert (ds3 == ds).all()
mpiutil.barrier()
# Check a read with an arbitrary slice in there. This only checks the shape is correct.
ds4 = mpiarray.MPIArray.from_hdf5(fname, 'testds', axis=1, sel=(np.s_[3:10:2], np.s_[1:16:3]))
assert ds4.shape[0] == 4
assert ds4.shape[1] == mpiutil.split_local(5)[0]
mpiutil.barrier()
# Check the read with a slice along the axis being read
ds5 = mpiarray.MPIArray.from_hdf5(fname, 'testds', axis=1, sel=(np.s_[:], np.s_[3:15:2]))
assert ds5.shape[0] == gshape[0]
assert ds5.shape[1] == mpiutil.split_local(6)[0]
ds5 = ds5.redistribute(axis=0)
assert (ds5 == ds[:, 3:15:2]).all()
mpiutil.barrier()
# Check the read with a slice along the axis being read
ds6 = mpiarray.MPIArray.from_hdf5(fname, 'testds', axis=0, sel=(np.s_[:], np.s_[3:15:2]))
ds6 = ds6.redistribute(axis=0)
assert (ds6 == ds[:, 3:15:2]).all()
mpiutil.barrier()
if mpiutil.rank0 and os.path.exists(fname):
os.remove(fname)
def process(self, rt):
assert isinstance(rt, RawTimestream), '%s only works for RawTimestream object currently' % self.__class__.__name__
if not 'ns_on' in rt.iterkeys():
raise RuntimeError('No noise source info, can not do noise source calibration')
local_bl_size, local_bl_start, local_bl_end = mpiutil.split_local(len(rt['blorder']))
rt.redistribute('baseline')
num_mean = self.params['num_mean']
phs_only = self.params['phs_only']
save_gain = self.params['save_gain']
tag_output_iter = self.params['tag_output_iter']
gain_file = self.params['gain_file']
bl_incl = self.params['bl_incl']
bl_excl = self.params['bl_excl']
freq_incl = self.params['freq_incl']
freq_excl = self.params['freq_excl']
ns_stable = self.params['ns_stable']
last_transit2stable = self.params['last_transit2stable']
#===================================
normal_gain_file = self.params['normal_gain_file']
rt.gain_file = self.params['gain_file']
absolute_gain_filename = self.params['absolute_gain_filename']
use_center_data = self.params['use_center_data']
if use_center_data and rt['ns_on'].attrs['on_time'] <= 2:
warnings.warn('The period %d <= 2, cannot get rid of the beginning and ending points. Use the whole average automatically!')
use_center_data = False
#===================================
# apply abs gain to data if ps_first
if rt.ps_first:
if absolute_gain_filename is None:
raise NoPsGainFile('Absent parameter absolute_gain_file. In ps_first mode, absolute_gain_filename is required to process!')
if not os.path.exists(output_path(absolute_gain_filename)):
raise NoPsGainFile('No absolute gain file %s, do the ps calibration first!'%output_path(absolute_gain_filename))
with h5py.File(output_path(absolute_gain_filename, 'r')) as agfilein:
if not ns_stable:
# try: # if there is transit in this batch of data, build up transit amp and phs
# rt.normal_index = np.where(np.abs(rt['jul_date'] - agfilein.attrs['transit_time']) < 2.e-6)[0][0] # to avoid numeric error. 1.e-6 julian date is about 0.1s
if rt['sec1970'][0] <= agfilein.attrs['transit_time'] and rt['sec1970'][-1] >= agfilein.attrs['transit_time']:
rt.normal_index = np.int64(np.around((agfilein.attrs['transit_time'] - rt['sec1970'][0])/rt.attrs['inttime']))
if mpiutil.rank0:
print('Detected transit, time index %d, build up transit normalization gain file!'%rt.normal_index)
build_normal_file = True
rt.normal_phs = np.zeros((rt['vis'].shape[1], local_bl_size))
rt.normal_phs.fill(np.nan)
if not phs_only:
rt.normal_amp = np.zeros((rt['vis'].shape[1], local_bl_size))
rt.normal_amp.fill(np.nan)
# except IndexError: # no transit in this batch of data, load transit amp and phs from file
else:
rt.normal_index = None # no transit flag
build_normal_file = False
if not os.path.exists(output_path(normal_gain_file)):
time_info = (aipy.phs.juldate2ephem(agfilein.attrs['transit_jul'] + 8./24.), aipy.phs.juldate2ephem(rt['jul_date'][0] + 8./24.), aipy.phs.juldate2ephem(rt['jul_date'][-1] + 8./24.))
raise TransitGainNotRecorded('The transit %s is not in time range %s to %s and no transit normalization gain was recorded!'%time_info)
else:
if mpiutil.rank0:
print('No transit, use existing transit normalization gain file!')
with h5py.File(output_path(normal_gain_file), 'r') as filein:
rt.normal_phs = filein['phs'][:, local_bl_start:local_bl_end]
if not phs_only:
rt.normal_amp = filein['amp'][:, local_bl_start:local_bl_end]
else:
rt.normal_index = None # no transit flag
rt.normal_phs = np.zeros((rt['vis'].shape[1], local_bl_size))
rt.normal_phs.fill(np.nan)
if not phs_only:
rt.normal_amp = np.zeros((rt['vis'].shape[1], local_bl_size))
rt.normal_amp.fill(np.nan)
try:
stable_time = rt['pointingtime'][-1,-1]
except KeyError:
stable_time = rt['transitsource'][-1, 0]
if stable_time > 0: # the time when the antenna will be pointing at the target region was recorded, use it
stable_time += 300 # plus 5 min to ensure stable
else:
stable_time = rt['transitsource'][-2, 0] + last_transit2stable
stable_time_jul = datetime.utcfromtimestamp(stable_time)
stable_time_jul = ephem.julian_date(stable_time_jul) # convert to julian date, convenient to display
if not os.path.exists(output_path(normal_gain_file)):
if mpiutil.rank0:
print('Normalization gain file has not been built yet. Try to build it!')
print('Recorded transit Time: %s'%aipy.phs.juldate2ephem(agfilein.attrs['transit_jul'] + 8./24.))
print('Last transit Time: %s'%aipy.phs.juldate2ephem(ephem.julian_date(datetime.utcfromtimestamp(rt['transitsource'][-2,0])) + 8./24.))
print('First time point of this data: %s'%aipy.phs.juldate2ephem(rt['jul_date'][0] + 8./24.))
build_normal_file = True
# if rt['jul_date'][0] < stable_time:
if rt.attrs['sec1970'][0] < stable_time:
raise BeforeStableTime('The beginning time point is %s, but only after %s, will the noise source be stable. Abort the noise calibration!'%(aipy.phs.juldate2ephem(rt['jul_date'][0] + 8./24.), aipy.phs.juldate2ephem(stable_time_jul + 8./24.)))
else:
if mpiutil.rank0:
print('Use existing normalization gain file!')
build_normal_file = False
with h5py.File(output_path(normal_gain_file), 'r') as filein:
rt.normal_phs = filein['phs'][:, local_bl_start:local_bl_end]
if not phs_only:
rt.normal_amp = filein['amp'][:, local_bl_start:local_bl_end]
# apply absolute gain
absgain = agfilein['gain'][:]
polfeed, _ = bl2pol_feed_inds(rt.local_bl, agfilein['gain'].attrs['feed'][:], agfilein['gain'].attrs['pol'][:])
for ii, (ipol, ifeed, jpol, jfeed) in enumerate(polfeed):
# rt.local_vis[:,:,ii] /= (absgain[None,:,ipol,ifeed-1] * absgain[None,:,jpol,jfeed-1].conj())
rt.local_vis[:,:,ii] /= (absgain[None,:,ipol,ifeed] * absgain[None,:,jpol,jfeed].conj())
#===================================
nt = rt.local_vis.shape[0]
if num_mean <= 0:
raise RuntimeError('Invalid num_mean = %s' % num_mean)
ns_on = rt['ns_on'][:]
ns_on = np.where(ns_on, 1, 0)
diff_ns = np.diff(ns_on)
inds = np.where(diff_ns==1)[0] # NOTE: these are inds just 1 before the first ON
if not rt.FRB_cal: # for FRB there might be just one noise point, avoid waste
if inds[0]-1 < 0: # no off data in the beginning to use
inds = inds[1:]
if inds[-1]+2 > nt-1: # no on data in the end to use
inds = inds[:-1]
if save_gain:
num_inds = len(inds)
shp = (num_inds,)+rt.local_vis.shape[1:]
dtype = rt.local_vis.real.dtype
# create dataset to record ns_cal_time_inds
rt.create_time_ordered_dataset('ns_cal_time_inds', inds)
# create dataset to record ns_cal_phase
ns_cal_phase = np.empty(shp, dtype=dtype)
ns_cal_phase[:] = np.nan
ns_cal_phase = mpiarray.MPIArray.wrap(ns_cal_phase, axis=2, comm=rt.comm)
rt.create_freq_and_bl_ordered_dataset('ns_cal_phase', ns_cal_phase, axis_order=(None, 1, 2))
rt['ns_cal_phase'].attrs['unit'] = 'radians'
if not phs_only:
# create dataset to record ns_cal_amp
ns_cal_amp = np.empty(shp, dtype=dtype)
ns_cal_amp[:] = np.nan
ns_cal_amp = mpiarray.MPIArray.wrap(ns_cal_amp, axis=2, comm=rt.comm)
rt.create_freq_and_bl_ordered_dataset('ns_cal_amp', ns_cal_amp, axis_order=(None, 1, 2))
if bl_incl == 'all':
bls_plt = [ tuple(bl) for bl in rt.bl ]
else:
bls_plt = [ bl for bl in bl_incl if not bl in bl_excl ]
if freq_incl == 'all':
freq_plt = range(rt.freq.shape[0])
else:
freq_plt = [ fi for fi in freq_incl if not fi in freq_excl ]
show_progress = self.params['show_progress']
progress_step = self.params['progress_step']
if rt.ps_first:
rt.freq_and_bl_data_operate(self.cal, full_data=True, show_progress=show_progress, progress_step=progress_step, keep_dist_axis=False, num_mean=num_mean, inds=inds, bls_plt=bls_plt, freq_plt=freq_plt, build_normal_file = build_normal_file)
else:
rt.freq_and_bl_data_operate(self.cal, full_data=True, show_progress=show_progress, progress_step=progress_step, keep_dist_axis=False, num_mean=num_mean, inds=inds, bls_plt=bls_plt, freq_plt=freq_plt)
if save_gain:
interp_mask_ratio = mpiutil.allreduce(np.sum(rt.interp_mask_count))/1./mpiutil.allreduce(np.size(rt.interp_mask_count)) * 100.
if interp_mask_ratio > 50. and rt.normal_index is not None:
warnings.warn('%.1f%% of the data was masked due to shortage of noise points for interpolation(need at least 4 to perform cubic spline)! The pointsource calibration may not be done due to too many masked points!'%interp_mask_ratio, NotEnoughPointToInterpolateWarning)
if interp_mask_ratio > 80.:
rt.interp_all_masked = True
# gather bl_order to rank0
bl_order = mpiutil.gather_array(rt['blorder'].local_data, axis=0, root=0, comm=rt.comm)
# gather ns_cal_phase / ns_cal_amp to rank 0
ns_cal_phase = mpiutil.gather_array(rt['ns_cal_phase'].local_data, axis=2, root=0, comm=rt.comm)
phs_unit = rt['ns_cal_phase'].attrs['unit']
rt.delete_a_dataset('ns_cal_phase', reserve_hint=False)
if not phs_only:
ns_cal_amp = mpiutil.gather_array(rt['ns_cal_amp'].local_data, axis=2, root=0, comm=rt.comm)
rt.delete_a_dataset('ns_cal_amp', reserve_hint=False)
if tag_output_iter:
gain_file = output_path(gain_file, iteration=self.iteration)
else:
gain_file = output_path(gain_file)
if rt.ps_first:
phase_finite_count = mpiutil.allreduce(np.isfinite(rt.normal_phs).sum())
if not phs_only:
amp_finite_count = mpiutil.allreduce(np.isfinite(rt.normal_amp).sum())
if mpiutil.rank0:
if rt.ps_first and build_normal_file:
if phase_finite_count == 0:
raise AllMasked('All values are masked when calculating phase!')
if not phs_only:
if amp_finite_count == 0:
raise AllMasked('All values are masked when calculating amplitude!')
with h5py.File(output_path(normal_gain_file), 'w') as tapfilein:
tapfilein.create_dataset('amp',(rt['vis'].shape[1], rt['vis'].shape[2]))
tapfilein.create_dataset('phs',(rt['vis'].shape[1], rt['vis'].shape[2]))
with h5py.File(gain_file, 'w') as f:
# save time
f.create_dataset('time', data=rt['jul_date'][:])
f['time'].attrs['unit'] = 'Julian date'
# save freq
f.create_dataset('freq', data=rt['freq'][:])
f['freq'].attrs['unit'] = rt['freq'].attrs['unit']
# save bl
f.create_dataset('bl_order', data=bl_order)
# save ns_cal_time_inds
f.create_dataset('ns_cal_time_inds', data=rt['ns_cal_time_inds'][:])
# save ns_cal_phase
f.create_dataset('ns_cal_phase', data=ns_cal_phase)
f['ns_cal_phase'].attrs['unit'] = phs_unit
f['ns_cal_phase'].attrs['dim'] = '(time, freq, bl)'
if not phs_only:
# save ns_cal_amp
f.create_dataset('ns_cal_amp', data=ns_cal_amp)
# save channo
f.create_dataset('channo', data=rt['channo'][:])
f['channo'].attrs['dim'] = rt['channo'].attrs['dimname']
if rt.exclude_bad:
f['channo'].attrs['badchn'] = rt['channo'].attrs['badchn']
if not (absolute_gain_filename is None):
if not os.path.exists(output_path(absolute_gain_filename)):
raise NoPsGainFile('No absolute gain file %s, do the ps calibration first!'%output_path(absolute_gain_filename))
with h5py.File(output_path(absolute_gain_filename,'r')) as abs_gain:
new_gain = uni_gain(abs_gain, f, phs_only = phs_only)
f.create_dataset('uni_gain', data = new_gain)
f['uni_gain'].attrs['dim'] = '(time, freq, bl)'
if rt.ps_first and build_normal_file:
if mpiutil.rank0:
print('Start write normalization gain into %s'%output_path(normal_gain_file))
for i in range(10):
try:
for ii in range(mpiutil.size):
if ii == mpiutil.rank:
with h5py.File(output_path(normal_gain_file), 'r+') as fileout:
fileout['phs'][:,local_bl_start:local_bl_end] = rt.normal_phs[:,:]
if not phs_only:
fileout['amp'][:,local_bl_start:local_bl_end] = rt.normal_amp[:,:]
mpiutil.barrier()
break
except IOError:
time.sleep(0.5)
continue
rt.delete_a_dataset('ns_cal_time_inds', reserve_hint=False)
return super(NsCal, self).process(rt)
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
----------
beamtransfer : fmmode.core.beamtransfer.BeamTransfer
BeamTransfer object containing the analysis products.
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.
add_noise : bool, optional
Weather to add random noise to the simulated visibilities. Default True.
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
nbl = tel.nbase
npol = tel.num_pol_sky
# 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))
indices = list(itertools.product(np.arange(nfreq), np.arange(npol)))
lind, sind, eind = mpiutil.split_local(nfreq * npol)
# local section of the Tm array
theta_size = tel.theta_size
phi_size = tel.phi_size
Tm = np.zeros((lind, theta_size, phi_size), dtype=np.complex128)
for ind, (f_ind, p_ind) in enumerate(indices[sind:eind]):
hp_map = None
for idx, mapfile in enumerate(maps):
with h5py.File(mapfile, 'r') as f:
if idx == 0:
hp_map = f['map'][f_ind, p_ind, :]
else:
hp_map += f['map'][f_ind, p_ind, :]
if hp_map is not None:
cart_map = hpproj.cartesian_proj(hp_map, tel.cart_projector)
# Calculate the Tm's for the local sections
Tm[ind] = np.fft.ifft(cart_map, axis=1) # / phi_size # m = 0 is at left
Tm = MPIArray.wrap(Tm, axis=0)
# redistribute along different m
Tm = Tm.redistribute(axis=2)
Tm = Tm.reshape((nfreq, npol, theta_size, None))
Tm = Tm.reshape((nfreq, npol*theta_size, None))
ms = np.concatenate([np.arange(0, mmax+1), np.arange(-mmax, 0)])
lm, sm, em = mpiutil.split_local(phi_size)
# local section of mmode
# mmode = np.zeros((lm, nbl, nfreq), dtype=np.complex128)
mmode = np.zeros((lm, nfreq, nbl), dtype=np.complex128)
for ind, mi in enumerate(ms[sm:em]):
mmode[ind] = bt.project_vector_sky_to_telescope(mi, Tm[:, :, ind].view(np.ndarray))
mmode = MPIArray.wrap(mmode, axis=0)
mmode = mmode.redistribute(axis=2) # distribute along bl
# add noise if required
if add_noise:
lbl, sbl, ebl = mpiutil.split_local(nbl)
# Fetch the noise powerspectrum
noise_ps = tel.noisepower(np.arange(sbl, ebl)[:, np.newaxis], np.arange(nfreq)[np.newaxis, :], ndays=ndays).reshape(lbl, nfreq).T[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_mode = (np.array([1.0, 1.0J]) * np.random.standard_normal(mmode.shape + (2,))).sum(axis=-1)
noise_mode *= (noise_ps / 2.0)**0.5
mmode += noise_mode
del noise_mode
# Reset RNG
if seed is not None:
np.random.seed()
# The time samples the visibility is calculated at
tphi = np.linspace(0, 2*np.pi, ntime, endpoint=False)
# inverse FFT to get timestream
vis_stream = np.fft.ifft(mmode, axis=0) * ntime
vis_stream = MPIArray.wrap(vis_stream, axis=2)
# save vis_stream to file
vis_h5 = memh5.MemGroup(distributed=True)
vis_h5.create_dataset('/timestream', data=vis_stream)
vis_h5.create_dataset('/phi', data=tphi)
# Telescope layout data
vis_h5.create_dataset('/feedmap', data=tel.feedmap)
vis_h5.create_dataset('/feedconj', data=tel.feedconj)
vis_h5.create_dataset('/feedmask', data=tel.feedmask)
vis_h5.create_dataset('/uniquepairs', data=tel.uniquepairs)
vis_h5.create_dataset('/baselines', data=tel.baselines)
# Telescope frequencies
vis_h5.create_dataset('/frequencies', data=tel.frequencies)
# Write metadata
vis_h5.attrs['beamtransfer_path'] = os.path.abspath(bt.directory)
vis_h5.attrs['ntime'] = ntime
# save to file
vis_h5.to_hdf5(tstream._tsfile)
if mpiutil.rank0:
# Make file marker that the m's have been correctly generated:
open(completed_file, 'a').close()
mpiutil.barrier()
return tstream
def mapmake_full(self, nside, maptype):
mapfile = self._mapsdir + 'map_%s.hdf5' % maptype
Tmfile = self._Tmsdir + 'Tm_%s.hdf5' % maptype
if os.path.exists(mapfile):
if mpiutil.rank0:
print "File %s exists. Skipping..." % mapfile
mpiutil.barrier()
return
elif os.path.exists(Tmfile):
if mpiutil.rank0:
print "File %s exists. Read from it..." % Tmfile
Tm = MPIArray.from_hdf5(Tmfile, 'Tm')
else:
def _make_Tm(mi):
print "Making %i" % mi
mmode = self.mmode(mi)
return self.beamtransfer.project_vector_telescope_to_sky(mi, mmode)
# if mpiutil.rank0 and not os.path.exists(self._Tmsdir):
# # Make directory for Tms file
# os.makedirs(self._Tmsdir)
# Make directory for Tms file
try:
os.makedirs(self._Tmsdir)
except OSError:
# directory exists
pass
tel = self.telescope
mmax = tel.mmax
lm, sm, em = mpiutil.split_local(mmax+1)
nfreq = tel.nfreq
npol = tel.num_pol_sky
ntheta = tel.theta_size
# the local Tm array
Tm = np.zeros((nfreq, npol, ntheta, lm), dtype=np.complex128)
for ind, mi in enumerate(range(sm, em)):
Tm[..., ind] = _make_Tm(mi)
Tm = MPIArray.wrap(Tm, axis=3)
Tm = Tm.redistribute(axis=0) # redistribute along freq
# Save Tm
Tm.to_hdf5(Tmfile, 'Tm', create=True)
# if mpiutil.rank0 and not os.path.exists(self._mapsdir):
# # Make directory for maps file
# os.makedirs(self._mapsdir)
# Make directory for maps file
try:
os.makedirs(self._mapsdir)
except OSError:
# directory exists
pass
tel = self.telescope
npol = tel.num_pol_sky
ntime = self.ntime
# irfft to get map
# cart_map = np.fft.irfft(Tm, axis=3, n=ntime) * ntime # NOTE the normalization constant ntime here to be consistant with the simulation fft
cart_map = np.fft.hfft(Tm, axis=3, n=ntime)
lfreq = cart_map.shape[0]
hp_map = np.zeros((lfreq, npol, 12*nside**2), dtype=cart_map.dtype)
for fi in range(lfreq):
for pi in range(npol):
hp_map[fi, pi] = tel.cart_projector.inv_projmap(cart_map[fi, pi], nside)
mpiutil.barrier()
hp_map = MPIArray.wrap(hp_map, axis=0)
# save map
hp_map.to_hdf5(mapfile, 'map', create=True)
def generate_mmodes(self, ts_data=None):
"""Calculate the m-modes corresponding to the Timestream.
Perform an MPI transpose for efficiency.
"""
completed_file = self._mdir + 'COMPLETED_M'
if os.path.exists(completed_file):
if mpiutil.rank0:
print "******* m-files already generated ********"
mpiutil.barrier()
return
# Make directory if required
# if mpiutil.rank0 and not os.path.exists(self._mdir):
# os.makedirs(self._mdir)
try:
os.makedirs(self._mdir)
except OSError:
# directory exists
pass
tel = self.telescope
mmax = tel.mmax
ntime = ts_data.shape[0] if ts_data is not None else self.ntime
nbl = tel.nbase
nfreq = tel.nfreq
indices = list(itertools.product(np.arange(nfreq), np.arange(nbl)))
lind, sind, eind = mpiutil.split_local(nfreq * nbl)
# load the local section of the time stream
tstream = np.zeros((ntime, lind), dtype=np.complex128)
for ind, (f_ind, bl_ind) in enumerate(indices[sind:eind]):
if ts_data is not None:
tstream[:, ind] = ts_data[:, f_ind, bl_ind]
else:
with h5py.File(self._tsfile, 'r') as f:
tstream[:, ind] = f['/timestream'][:, f_ind, bl_ind]
# FFT to get m-mode
mmodes = np.fft.fft(tstream, axis=0) / ntime # m = 0 is at left
mmodes = MPIArray.wrap(mmodes, axis=1)
# redistribute along different m
mmodes = mmodes.redistribute(axis=0)
# save m-modes to file
ms = np.concatenate([np.arange(0, mmax + 1), np.arange(-mmax, 0)])
for ind, mi in enumerate(mpiutil.mpilist(ms, method='con')):
with h5py.File(self._mfile(mi), 'w') as f:
f.create_dataset('/mmode',
data=mmodes[ind].view(np.ndarray).reshape(
nfreq, nbl))
f.attrs['m'] = mi
mpiutil.barrier()
if mpiutil.rank0:
# Make file marker that the m's have been correctly generated:
open(completed_file, 'a').close()
def mapmake_full(self, nside, maptype):
mapfile = self._mapsdir + 'map_%s.hdf5' % maptype
Tmfile = self._Tmsdir + 'Tm_%s.hdf5' % maptype
if os.path.exists(mapfile):
if mpiutil.rank0:
print "File %s exists. Skipping..." % mapfile
mpiutil.barrier()
return
elif os.path.exists(Tmfile):
if mpiutil.rank0:
print "File %s exists. Read from it..." % Tmfile
Tm = MPIArray.from_hdf5(Tmfile, 'Tm')
else:
def _make_Tm(mi):
print "Making %i" % mi
mmode = self.mmode(mi)
return self.beamtransfer.project_vector_telescope_to_sky(
mi, mmode)
# if mpiutil.rank0 and not os.path.exists(self._Tmsdir):
# # Make directory for Tms file
# os.makedirs(self._Tmsdir)
# Make directory for Tms file
try:
os.makedirs(self._Tmsdir)
except OSError:
# directory exists
pass
tel = self.telescope
mmax = tel.mmax
lm, sm, em = mpiutil.split_local(mmax + 1)
nfreq = tel.nfreq
npol = tel.num_pol_sky
ntheta = tel.theta_size
# the local Tm array
Tm = np.zeros((nfreq, npol, ntheta, lm), dtype=np.complex128)
for ind, mi in enumerate(range(sm, em)):
Tm[..., ind] = _make_Tm(mi)
Tm = MPIArray.wrap(Tm, axis=3)
Tm = Tm.redistribute(axis=0) # redistribute along freq
# Save Tm
Tm.to_hdf5(Tmfile, 'Tm', create=True)
# if mpiutil.rank0 and not os.path.exists(self._mapsdir):
# # Make directory for maps file
# os.makedirs(self._mapsdir)
# Make directory for maps file
try:
os.makedirs(self._mapsdir)
except OSError:
# directory exists
pass
tel = self.telescope
npol = tel.num_pol_sky
ntime = self.ntime
# irfft to get map
# cart_map = np.fft.irfft(Tm, axis=3, n=ntime) * ntime # NOTE the normalization constant ntime here to be consistant with the simulation fft
cart_map = np.fft.hfft(Tm, axis=3, n=ntime)
lfreq = cart_map.shape[0]
hp_map = np.zeros((lfreq, npol, 12 * nside**2), dtype=cart_map.dtype)
for fi in range(lfreq):
for pi in range(npol):
hp_map[fi, pi] = tel.cart_projector.inv_projmap(
cart_map[fi, pi], nside)
mpiutil.barrier()
hp_map = MPIArray.wrap(hp_map, axis=0)
# save map
hp_map.to_hdf5(mapfile, 'map', create=True)
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
----------
beamtransfer : fmmode.core.beamtransfer.BeamTransfer
BeamTransfer object containing the analysis products.
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.
add_noise : bool, optional
Weather to add random noise to the simulated visibilities. Default True.
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
nbl = tel.nbase
npol = tel.num_pol_sky
# 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))
indices = list(itertools.product(np.arange(nfreq), np.arange(npol)))
lind, sind, eind = mpiutil.split_local(nfreq * npol)
# local section of the Tm array
theta_size = tel.theta_size
phi_size = tel.phi_size
Tm = np.zeros((lind, theta_size, phi_size), dtype=np.complex128)
for ind, (f_ind, p_ind) in enumerate(indices[sind:eind]):
hp_map = None
for idx, mapfile in enumerate(maps):
with h5py.File(mapfile, 'r') as f:
if idx == 0:
hp_map = f['map'][f_ind, p_ind, :]
else:
hp_map += f['map'][f_ind, p_ind, :]
if hp_map is not None:
cart_map = hpproj.cartesian_proj(hp_map, tel.cart_projector)
# Calculate the Tm's for the local sections
Tm[ind] = np.fft.ifft(cart_map,
axis=1) # / phi_size # m = 0 is at left
Tm = MPIArray.wrap(Tm, axis=0)
# redistribute along different m
Tm = Tm.redistribute(axis=2)
Tm = Tm.reshape((nfreq, npol, theta_size, None))
Tm = Tm.reshape((nfreq, npol * theta_size, None))
ms = np.concatenate([np.arange(0, mmax + 1), np.arange(-mmax, 0)])
lm, sm, em = mpiutil.split_local(phi_size)
# local section of mmode
# mmode = np.zeros((lm, nbl, nfreq), dtype=np.complex128)
mmode = np.zeros((lm, nfreq, nbl), dtype=np.complex128)
for ind, mi in enumerate(ms[sm:em]):
mmode[ind] = bt.project_vector_sky_to_telescope(
mi, Tm[:, :, ind].view(np.ndarray))
mmode = MPIArray.wrap(mmode, axis=0)
mmode = mmode.redistribute(axis=2) # distribute along bl
# add noise if required
if add_noise:
lbl, sbl, ebl = mpiutil.split_local(nbl)
# Fetch the noise powerspectrum
noise_ps = tel.noisepower(np.arange(sbl, ebl)[:, np.newaxis],
np.arange(nfreq)[np.newaxis, :],
ndays=ndays).reshape(
lbl, nfreq).T[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_mode = (np.array([1.0, 1.0J]) *
np.random.standard_normal(mmode.shape +
(2, ))).sum(axis=-1)
noise_mode *= (noise_ps / 2.0)**0.5
mmode += noise_mode
del noise_mode
# Reset RNG
if seed is not None:
np.random.seed()
# The time samples the visibility is calculated at
tphi = np.linspace(0, 2 * np.pi, ntime, endpoint=False)
# inverse FFT to get timestream
vis_stream = np.fft.ifft(mmode, axis=0) * ntime
vis_stream = MPIArray.wrap(vis_stream, axis=2)
# save vis_stream to file
vis_h5 = memh5.MemGroup(distributed=True)
vis_h5.create_dataset('/timestream', data=vis_stream)
vis_h5.create_dataset('/phi', data=tphi)
# Telescope layout data
vis_h5.create_dataset('/feedmap', data=tel.feedmap)
vis_h5.create_dataset('/feedconj', data=tel.feedconj)
vis_h5.create_dataset('/feedmask', data=tel.feedmask)
vis_h5.create_dataset('/uniquepairs', data=tel.uniquepairs)
vis_h5.create_dataset('/baselines', data=tel.baselines)
# Telescope frequencies
vis_h5.create_dataset('/frequencies', data=tel.frequencies)
# Write metadata
vis_h5.attrs['beamtransfer_path'] = os.path.abspath(bt.directory)
vis_h5.attrs['ntime'] = ntime
# save to file
vis_h5.to_hdf5(tstream._tsfile)
if mpiutil.rank0:
# Make file marker that the m's have been correctly generated:
open(completed_file, 'a').close()
mpiutil.barrier()
return tstream
def _generate_phase(self, time):
ntime = len(time)
freq = self.freq
nfreq = len(freq)
# Generate the correlation function
cf_delay = self._corr_func(self.corr_length_delay, self.sigma_delay)
# Check if we are simulating relative delays or common mode delays
if self.sim_type == "relative":
n_realisations = self.ninput_local
# Generate delay fluctuations
self.delay_error = gain.generate_fluctuations(
time, cf_delay, n_realisations, self._prev_time,
self._prev_delay)
gain_phase = (2.0 * np.pi * freq[:, np.newaxis, np.newaxis] * 1e6 *
self.delay_error[np.newaxis, :, :] /
np.sqrt(self.ndays))
if self.sim_type == "common_mode_cyl":
n_realisations = 1
ninput = self.ninput_global
# Generates as many random delay errors as there are cylinders
if self.comm.rank == 0:
if self.common_mode_type == "sinusoidal":
P1 = self.sinusoidal_period[0]
P2 = self.sinusoidal_period[1]
omega1 = 2 * np.pi / P1
omega2 = 2 * np.pi / P2
delay_error = (self.sigma_delay *
(np.sin(omega1 * time) -
np.sin(omega2 * time))[np.newaxis, :])
if self.common_mode_type == "random":
delay_error = gain.generate_fluctuations(
time,
cf_delay,
n_realisations,
self._prev_time,
self._prev_delay,
)
else:
delay_error = None
# Broadcast to other ranks
self.delay_error = self.comm.bcast(delay_error, root=0)
# Split frequencies to processes.
lfreq, sfreq, efreq = mpiutil.split_local(nfreq)
# Create an array to hold all inputs, which are common-mode within
# a cylinder
gain_phase = np.zeros((lfreq, ninput, ntime), dtype=complex)
# Since we have 2 cylinders populate half of them with a delay)
# TODO: generalize this for 3 or even 4 cylinders in the future.
gain_phase[:, ninput // self.ncyl:, :] = (
2.0 * np.pi * freq[sfreq:efreq, np.newaxis, np.newaxis] * 1e6 *
self.delay_error[np.newaxis, :, :] / np.sqrt(self.ndays))
gain_phase = mpiarray.MPIArray.wrap(gain_phase,
axis=0,
comm=self.comm)
# Redistribute over input to match rest of the code
gain_phase = gain_phase.redistribute(axis=1)
gain_phase = gain_phase.view(np.ndarray)
if self.sim_type == "common_mode_iceboard":
nchannel = self.nchannel
ninput = self.ninput_global
# Number of channels on a board
nboards = ninput // nchannel
# Generates as many random delay errors as there are iceboards
if self.comm.rank == 0:
delay_error = gain.generate_fluctuations(
time, cf_delay, nboards, self._prev_time, self._prev_delay)
else:
delay_error = None
# Broadcast to other ranks
self.delay_error = self.comm.bcast(delay_error, root=0)
# Calculate the corresponding phase by multiplying with frequencies
phase = (2.0 * np.pi * freq[:, np.newaxis, np.newaxis] * 1e6 *
self.delay_error[np.newaxis, :] / np.sqrt(self.ndays))
# Create an array to hold all inputs, which are common-mode within
# one iceboard
gain_phase = mpiarray.MPIArray((nfreq, ninput, ntime),
axis=1,
dtype=np.complex128,
comm=self.comm)
gain_phase[:] = 0.0
# Loop over inputs and and group common-mode phases on every board
for il, ig in gain_phase.enumerate(axis=1):
# Get the board number bi
bi = int(ig / nchannel)
gain_phase[:, il] = phase[:, bi]
gain_phase = gain_phase.view(np.ndarray)
self._prev_delay = self.delay_error
self._prev_time = time
return gain_phase
def wrap(cls, array, axis, comm=None):
"""Turn a set of numpy arrays into a distributed MPIArray object.
This is needed for functions such as `np.fft.fft` which always return
an `np.ndarray`.
Parameters
----------
array : np.ndarray
Array to wrap.
axis : integer
Axis over which the array is distributed. The lengths are checked
to try and ensure this is correct.
comm : MPI.Comm, optional
The communicator over which the array is distributed. If `None`
(default), use `MPI.COMM_WORLD`.
Returns
-------
dist_array : MPIArray
An MPIArray view of the input.
"""
# from mpi4py import MPI
if comm is None:
comm = mpiutil.world
# Get axis length, both locally, and globally
axlen = array.shape[axis]
totallen = mpiutil.allreduce(axlen, comm=comm)
# Figure out what the distributed layout should be
local_num, local_start, local_end = mpiutil.split_local(totallen,
comm=comm)
# Check the local layout is consistent with what we expect, and send
# result to all ranks
layout_issue = mpiutil.allreduce(axlen != local_num,
op=mpiutil.MAX,
comm=comm)
if layout_issue:
raise Exception(
"Cannot wrap, distributed axis local length is incorrect.")
# Set shape and offset
lshape = array.shape
global_shape = list(lshape)
global_shape[axis] = totallen
loffset = [0] * len(lshape)
loffset[axis] = local_start
# Setup attributes of class
dist_arr = array.view(cls)
dist_arr._global_shape = tuple(global_shape)
dist_arr._axis = axis
dist_arr._local_shape = tuple(lshape)
dist_arr._local_offset = tuple(loffset)
dist_arr._comm = comm
return dist_arr
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.rank == 0:
import warnings
warnings.warn(
"Cannot redistribute array of compound datatypes."
" Sorry!!")
return self
# Get a view of the array
arr = self.view(np.ndarray)
if self.comm.size == 1:
# only one process
if arr.shape[self.axis] == self.global_shape[self.axis]:
# We are working on a single node and being asked to do
# a trivial transpose.
trans_arr = arr.copy()
else:
raise ValueError(
"Global shape %s is incompatible with local arrays shape %s"
% (self.global_shape, self.shape))
else:
pc, sc, ec = mpiutil.split_local(arr.shape[axis], comm=self.comm)
par, sar, ear = mpiutil.split_all(self.global_shape[self.axis],
comm=self.comm)
pac, sac, eac = mpiutil.split_all(arr.shape[axis], comm=self.comm)
new_shape = np.asarray(self.global_shape)
new_shape[axis] = pc
requests_send = []
requests_recv = []
trans_arr = np.empty(new_shape, dtype=arr.dtype)
mpitype = mpiutil.typemap(arr.dtype)
buffers = list()
# Cut out the right blocks of the local array to send around
blocks = np.array_split(arr, np.insert(eac, 0, sac[0]), axis)[1:]
# Iterate over all processes row wise
for ir in range(self.comm.size):
# Iterate over all processes column wise
for ic in range(self.comm.size):
# Construct a unique tag
tag = ir * self.comm.size + ic
# Send and receive the messages as non-blocking passes
if self.comm.rank == ir:
# Send the message
request = self.comm.Isend(
[blocks[ic].flatten(), mpitype], dest=ic, tag=tag)
requests_send.append([ir, ic, request])
if self.comm.rank == ic:
buffer_shape = np.asarray(new_shape)
buffer_shape[axis] = eac[ic] - sac[ic]
buffer_shape[self.axis] = ear[ir] - sar[ir]
buffers.append(
np.ndarray(buffer_shape, dtype=arr.dtype))
request = self.comm.Irecv([buffers[ir], mpitype],
source=ir,
tag=tag)
requests_recv.append([ir, ic, request])
# Wait for all processes to have started their messages
self.comm.Barrier()
# For each node iterate over all sends and wait until completion
for ir, ic, request in requests_send:
stat = mpiutil.MPI.Status()
request.Wait(status=stat)
if stat.error != mpiutil.MPI.SUCCESS:
print(
"**** ERROR in MPI SEND (r: %i c: %i rank: %i) *****".
format(ir, ic, self.comm.rank))
self.comm.Barrier()
# For each frequency iterate over all receives and wait until
# completion
for ir, ic, request in requests_recv:
stat = mpiutil.MPI.Status()
request.Wait(status=stat)
if stat.error != mpiutil.MPI.SUCCESS:
print(
"**** ERROR in MPI RECV (r: %i c: %i rank: %i) *****".
format(ir, ir, self.comm.rank))
# Put together the blocks we received
np.concatenate(buffers, self.axis, trans_arr)
# 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 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.rank == 0:
import warnings
warnings.warn('Cannot redistribute array of compound datatypes.'
' Sorry!!')
return self
# Get a view of the array
arr = self.view(np.ndarray)
if self.comm.size == 1:
# only one process
if arr.shape[self.axis] == self.global_shape[self.axis]:
# We are working on a single node and being asked to do
# a trivial transpose.
trans_arr = arr.copy()
else:
raise ValueError(
'Global shape %s is incompatible with local arrays shape %s'
% (self.global_shape, self.shape))
else:
pc, sc, ec = mpiutil.split_local(arr.shape[axis], comm=self.comm)
par, sar, ear = mpiutil.split_all(self.global_shape[self.axis],
comm=self.comm)
pac, sac, eac = mpiutil.split_all(arr.shape[axis], comm=self.comm)
new_shape = np.asarray(self.global_shape)
new_shape[axis] = pc
requests_send = []
requests_recv = []
trans_arr = np.empty(new_shape, dtype=arr.dtype)
mpitype = mpiutil.typemap(arr.dtype)
buffers = list()
# Cut out the right blocks of the local array to send around
blocks = np.array_split(arr, np.insert(eac, 0, sac[0]), axis)[1:]
# Iterate over all processes row wise
for ir in range(self.comm.size):
# Iterate over all processes column wise
for ic in range(self.comm.size):
# Construct a unique tag
tag = ir * self.comm.size + ic
# Send and receive the messages as non-blocking passes
if self.comm.rank == ir:
# Send the message
request = self.comm.Isend([blocks[ic].flatten(),
mpitype], dest=ic, tag=tag)
requests_send.append([ir, ic, request])
if self.comm.rank == ic:
buffer_shape = np.asarray(new_shape)
buffer_shape[axis] = eac[ic] - sac[ic]
buffer_shape[self.axis] = ear[ir] - sar[ir]
buffers.append(np.ndarray(buffer_shape, dtype=arr.dtype))
request = self.comm.Irecv([buffers[ir], mpitype],
source=ir, tag=tag)
requests_recv.append([ir, ic, request])
# Wait for all processes to have started their messages
self.comm.Barrier()
# For each node iterate over all sends and wait until completion
for ir, ic, request in requests_send:
stat = mpiutil.MPI.Status()
request.Wait(status=stat)
if stat.error != mpiutil.MPI.SUCCESS:
print("**** ERROR in MPI SEND (r: %i c: %i rank: %i) *****"
.format(ir, ic, self.comm.rank))
self.comm.Barrier()
# For each frequency iterate over all receives and wait until
# completion
for ir, ic, request in requests_recv:
stat = mpiutil.MPI.Status()
request.Wait(status=stat)
if stat.error != mpiutil.MPI.SUCCESS:
print("**** ERROR in MPI RECV (r: %i c: %i rank: %i) *****"
.format(ir, ir, self.comm.rank))
# Put together the blocks we received
np.concatenate(buffers, self.axis, trans_arr)
# 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
import matplotlib
matplotlib.use('Agg')
if mpiutil.size != 20:
raise RuntimeError('Need 20 processes')
np.random.seed(0)
rands = np.random.rand(1000, 20)
N = 1000
m = 20 # number of augmentation by rotation
n = 256
dataset = np.zeros((N / 20, m, n, n, 1))
# for i in range(N):
ln, s, e = mpiutil.split_local(N)
for li, i in enumerate(range(s, e)):
if mpiutil.rank0:
print '%d of %d ...' % (li, ln)
# read in reconstructed map
with h5py.File('./output_sim_750MHz/map/ts/ts_%04d/map_full.hdf5' % i,
'r') as f:
rec_map = f['map'][0, 0]
for j in range(m):
if j == 0:
rot_map = rec_map.copy()
else:
# angle = 360 * np.random.rand()
angle = 360 * rands[i, j]
rot_map = rotate.np_rotate(rec_map.copy(), angle)
print 'rank %d has %s after gather_list' % (rank, lst)
# parallel_map
separator(sec, 'parallel_map')
glist = range(6)
result = mpiutil.parallel_map(lambda x: x * x, glist, root=0)
if rank == 0:
print 'result = %s' % result
# split_all
separator(sec, 'split_all')
print 'rank %d has: %s' % (rank, mpiutil.split_all(6))
# split_local
separator(sec, 'split_local')
print 'rank %d has: %s' % (rank, mpiutil.split_local(6))
# gather_array
separator(sec, 'gather_array')
if rank == 0:
local_ary = np.array([[0, 1], [6, 7]])
elif rank == 1:
local_ary = np.array([[2], [8]])
elif rank == 2:
local_ary = np.array([[3], [9]])
if rank == 3:
local_ary = np.array([[4, 5], [10, 11]])
global_ary = mpiutil.gather_array(local_ary, axis=1, root=0)
if rank == 0:
print 'global_ary = %s' % global_ary
def test_io(self):
import h5py
# Cleanup directories
fname = "testdset.hdf5"
if mpiutil.rank0 and os.path.exists(fname):
os.remove(fname)
mpiutil.barrier()
gshape = (19, 17)
ds = mpiarray.MPIArray(gshape, dtype=np.int64)
ga = np.arange(np.prod(gshape)).reshape(gshape)
l0, s0, e0 = mpiutil.split_local(gshape[0])
ds[:] = ga[s0:e0]
ds.redistribute(axis=1).to_hdf5(fname, "testds", create=True)
if mpiutil.rank0:
with h5py.File(fname, "r") as f:
h5ds = f["testds"][:]
assert (h5ds == ga).all()
ds2 = mpiarray.MPIArray.from_hdf5(fname, "testds")
assert (ds2 == ds).all()
mpiutil.barrier()
# Check that reading over another distributed axis works
ds3 = mpiarray.MPIArray.from_hdf5(fname, "testds", axis=1)
assert ds3.shape[0] == gshape[0]
assert ds3.shape[1] == mpiutil.split_local(gshape[1])[0]
ds3 = ds3.redistribute(axis=0)
assert (ds3 == ds).all()
mpiutil.barrier()
# Check a read with an arbitrary slice in there. This only checks the shape is correct.
ds4 = mpiarray.MPIArray.from_hdf5(fname,
"testds",
axis=1,
sel=(np.s_[3:10:2], np.s_[1:16:3]))
assert ds4.shape[0] == 4
assert ds4.shape[1] == mpiutil.split_local(5)[0]
mpiutil.barrier()
# Check the read with a slice along the axis being read
ds5 = mpiarray.MPIArray.from_hdf5(fname,
"testds",
axis=1,
sel=(np.s_[:], np.s_[3:15:2]))
assert ds5.shape[0] == gshape[0]
assert ds5.shape[1] == mpiutil.split_local(6)[0]
ds5 = ds5.redistribute(axis=0)
assert (ds5 == ds[:, 3:15:2]).all()
mpiutil.barrier()
# Check the read with a slice along the axis being read
ds6 = mpiarray.MPIArray.from_hdf5(fname,
"testds",
axis=0,
sel=(np.s_[:], np.s_[3:15:2]))
ds6 = ds6.redistribute(axis=0)
assert (ds6 == ds[:, 3:15:2]).all()
mpiutil.barrier()
if mpiutil.rank0 and os.path.exists(fname):
os.remove(fname)
