def HaloConcentration(mass, cosmo, redshift, mdef='vir'):
"""
Return halo concentration from halo mass, based on the analytic fitting
formulas presented in
`Dutton and Maccio 2014 <https://arxiv.org/abs/1402.7073>`_.
.. note::
The units of the input mass are assumed to be :math:`M_{\odot}/h`
Parameters
----------
mass : array_like
either a numpy or dask array specifying the halo mass; units
assumed to be :math:`M_{\odot}/h`
cosmo : :class:`~nbodykit.cosmology.cosmology.Cosmology`
the cosmology instance used in the analytic formula
redshift : float
compute the c(M) relation at this redshift
mdef : str, optional
string specifying the halo mass definition to use; should be
'vir' or 'XXXc' or 'XXXm' where 'XXX' is an int specifying the
overdensity
Returns
-------
concen : :class:`dask.array.Array`
a dask array holding the analytic concentration values
References
----------
Dutton and Maccio, "Cold dark matter haloes in the Planck era: evolution
of structural parameters for Einasto and NFW profiles", 2014, arxiv:1402.7073
"""
from halotools.empirical_models import NFWProfile
mass, redshift = da.broadcast_arrays(mass, redshift)
kws = {
'cosmology': cosmo.to_astropy(),
'conc_mass_model': 'dutton_maccio14',
'mdef': mdef
}
def get_nfw_conc(mass, redshift):
kw1 = {}
kw1.update(kws)
kw1['redshift'] = redshift
model = NFWProfile(**kw1)
return model.conc_NFWmodel(prim_haloprop=mass)
return da.map_blocks(get_nfw_conc, mass, redshift, dtype=mass.dtype)
def SkyToCartesian(ra,
dec,
redshift,
cosmo,
observer=[0, 0, 0],
degrees=True,
frame='icrs'):
"""
Convert sky coordinates (``ra``, ``dec``, ``redshift``) to a
Cartesian ``Position`` column.
.. warning::
The returned Cartesian position is in units of Mpc/h.
Parameters
-----------
ra : :class:`dask.array.Array`; shape: (N,)
the right ascension angular coordinate
dec : :class:`dask.array.Array`; shape: (N,)
the declination angular coordinate
redshift : :class:`dask.array.Array`; shape: (N,)
the redshift coordinate
cosmo : :class:`~nbodykit.cosmology.cosmology.Cosmology`
the cosmology used to meausre the comoving distance from ``redshift``
degrees : bool, optional
specifies whether ``ra`` and ``dec`` are in degrees
frame : string ('icrs' or 'galactic')
speciefies which frame the Cartesian coordinates is.
Returns
-------
pos : :class:`dask.array.Array`; shape: (N,3)
the cartesian position coordinates, where columns represent
``x``, ``y``, and ``z`` in units of Mpc/h
Raises
------
TypeError
If the input columns are not dask arrays
"""
ra, dec, redshift = da.broadcast_arrays(ra, dec, redshift)
# pos on the unit sphere
pos = SkyToUnitSphere(ra, dec, degrees=degrees, frame=frame)
# multiply by the comoving distance in Mpc/h
r = redshift.map_blocks(lambda z: cosmo.comoving_distance(z),
dtype=redshift.dtype)
return r[:, None] * pos + observer
def HaloVelocityDispersion(mass, cosmo, redshift, mdef='vir'):
""" Compute the velocity dispersion of halo from Mass.
This is a simple model suggested by Martin White.
See http://adsabs.harvard.edu/abs/2008ApJ...672..122E
"""
mass, redshift = da.broadcast_arrays(mass, redshift)
def compute_vdisp(mass, redshift):
h = cosmo.efunc(redshift)
return 1100. * (h * mass / 1e15) ** 0.33333
return da.map_blocks(compute_vdisp, mass, redshift, dtype=mass.dtype)
def HaloVelocityDispersion(mass, cosmo, redshift, mdef='vir'):
""" Compute the velocity dispersion of halo from Mass.
This is a simple model suggested by Martin White.
See http://adsabs.harvard.edu/abs/2008ApJ...672..122E
"""
mass, redshift = da.broadcast_arrays(mass, redshift)
def compute_vdisp(mass, redshift):
h = cosmo.efunc(redshift)
return 1100. * (h * mass / 1e15) ** 0.33333
return da.map_blocks(compute_vdisp, mass, redshift, dtype=mass.dtype)
def HaloConcentration(mass, cosmo, redshift, mdef='vir'):
"""
Return halo concentration from halo mass, based on the analytic fitting
formulas presented in
`Dutton and Maccio 2014 <https://arxiv.org/abs/1402.7073>`_.
.. note::
The units of the input mass are assumed to be :math:`M_{\odot}/h`
Parameters
----------
mass : array_like
either a numpy or dask array specifying the halo mass; units
assumed to be :math:`M_{\odot}/h`
cosmo : :class:`~nbodykit.cosmology.cosmology.Cosmology`
the cosmology instance used in the analytic formula
redshift : float
compute the c(M) relation at this redshift
mdef : str, optional
string specifying the halo mass definition to use; should be
'vir' or 'XXXc' or 'XXXm' where 'XXX' is an int specifying the
overdensity
Returns
-------
concen : :class:`dask.array.Array`
a dask array holding the analytic concentration values
References
----------
Dutton and Maccio, "Cold dark matter haloes in the Planck era: evolution
of structural parameters for Einasto and NFW profiles", 2014, arxiv:1402.7073
"""
from halotools.empirical_models import NFWProfile
mass, redshift = da.broadcast_arrays(mass, redshift)
kws = {'cosmology':cosmo.to_astropy(), 'conc_mass_model':'dutton_maccio14', 'mdef':mdef}
def get_nfw_conc(mass, redshift):
kw1 = {}
kw1.update(kws)
kw1['redshift'] = redshift
model = NFWProfile(**kw1)
return model.conc_NFWmodel(prim_haloprop=mass)
return da.map_blocks(get_nfw_conc, mass, redshift, dtype=mass.dtype)
def HaloRadius(mass, cosmo, redshift, mdef='vir'):
r"""
Return proper halo radius from halo mass, based on the specified mass definition.
This is independent of halo profile, and simply returns
.. math::
R = \left [ 3 M /(4\pi\Delta) \right]^{1/3}
where :math:`\Delta` is the density threshold, which depends on cosmology,
redshift, and mass definition
.. note::
The units of the input mass are assumed to be :math:`M_{\odot}/h`
Parameters
----------
mass : array_like
either a numpy or dask array specifying the halo mass; units
assumed to be :math:`M_{\odot}/h`
cosmo : :class:`~nbodykit.cosmology.cosmology.Cosmology`
the cosmology instance
redshift : float
compute the density threshold which determines the R(M) relation
at this redshift
mdef : str, optional
string specifying the halo mass definition to use; should be
'vir' or 'XXXc' or 'XXXm' where 'XXX' is an int specifying the
overdensity
Returns
-------
radius : :class:`dask.array.Array`
a dask array holding the halo radius in 'physical Mpc/h [sic]'.
This is proper Mpc/h, to convert to comoving, divide this by scaling factor.
"""
from halotools.empirical_models import halo_mass_to_halo_radius
mass, redshift = da.broadcast_arrays(mass, redshift)
kws = {'cosmology':cosmo.to_astropy(), 'mdef':mdef}
def mass_to_radius(mass, redshift):
return halo_mass_to_halo_radius(mass=mass, redshift=redshift, **kws)
return da.map_blocks(mass_to_radius, mass, redshift, dtype=mass.dtype)
def HaloRadius(mass, cosmo, redshift, mdef='vir'):
r"""
Return proper halo radius from halo mass, based on the specified mass definition.
This is independent of halo profile, and simply returns
.. math::
R = \left [ 3 M /(4\pi\Delta) \right]^{1/3}
where :math:`\Delta` is the density threshold, which depends on cosmology,
redshift, and mass definition
.. note::
The units of the input mass are assumed to be :math:`M_{\odot}/h`
Parameters
----------
mass : array_like
either a numpy or dask array specifying the halo mass; units
assumed to be :math:`M_{\odot}/h`
cosmo : :class:`~nbodykit.cosmology.cosmology.Cosmology`
the cosmology instance
redshift : float
compute the density threshold which determines the R(M) relation
at this redshift
mdef : str, optional
string specifying the halo mass definition to use; should be
'vir' or 'XXXc' or 'XXXm' where 'XXX' is an int specifying the
overdensity
Returns
-------
radius : :class:`dask.array.Array`
a dask array holding the halo radius in 'physical Mpc/h [sic]'.
This is proper Mpc/h, to convert to comoving, divide this by scaling factor.
"""
from halotools.empirical_models import halo_mass_to_halo_radius
mass, redshift = da.broadcast_arrays(mass, redshift)
kws = {'cosmology':cosmo.to_astropy(), 'mdef':mdef}
def mass_to_radius(mass, redshift):
return halo_mass_to_halo_radius(mass=mass, redshift=redshift, **kws)
return da.map_blocks(mass_to_radius, mass, redshift, dtype=mass.dtype)
def test_dataframe_factory(test_nan_shapes):
nrow, nfreq = 100, 100
data1a = da.arange(nrow, chunks=(10, ))
# Generate nan chunk shapes in data1a if requested
if test_nan_shapes:
data1a = data1a[da.where(data1a > 4)]
data1b = da.random.random(size=(nfreq, ), chunks=(100, ))
df = dataframe_factory(("row", "chan"), data1a, ("row", ), data1b,
("chan", ))
assert isinstance(df, dd.DataFrame)
assert isinstance(df['x'], dd.Series)
assert isinstance(df['y'], dd.Series)
if test_nan_shapes:
# With unknown shapes, we broadcast our arrays in numpy
# to test
x, y = dask.compute(data1a[:, None], data1b[None, :])
x, y = np.broadcast_arrays(x, y)
x = x.ravel()
y = y.ravel()
# Unknown divisions
assert df.divisions == (None, ) * (df.npartitions + 1)
else:
# With know chunks we can broadcast our arrays in dask
x, y = da.broadcast_arrays(data1a[:, None], data1b[None, :])
x = x.ravel()
y = y.ravel()
# Known divisions
assert df.divisions == (0, 1000, 2000, 3000, 4000, 5000, 6000, 7000,
8000, 9000, 10000)
assert df['x'].npartitions == x.npartitions
assert df['y'].npartitions == y.npartitions
# Compare our lazy dataframe series vs (dask or numpy) arrays
assert_array_equal(df['x'], x)
assert_array_equal(df['y'], y)
assert_array_equal(df['x'].min(), data1a.min())
assert_array_equal(df['y'].min(), data1b.min())
assert_array_equal(df['x'].max(), data1a.max())
assert_array_equal(df['y'].max(), data1b.max())
def SkyToCartesian(ra, dec, redshift, cosmo, observer=[0, 0, 0], degrees=True, frame='icrs'):
"""
Convert sky coordinates (``ra``, ``dec``, ``redshift``) to a
Cartesian ``Position`` column.
.. warning::
The returned Cartesian position is in units of Mpc/h.
Parameters
-----------
ra : :class:`dask.array.Array`; shape: (N,)
the right ascension angular coordinate
dec : :class:`dask.array.Array`; shape: (N,)
the declination angular coordinate
redshift : :class:`dask.array.Array`; shape: (N,)
the redshift coordinate
cosmo : :class:`~nbodykit.cosmology.cosmology.Cosmology`
the cosmology used to meausre the comoving distance from ``redshift``
degrees : bool, optional
specifies whether ``ra`` and ``dec`` are in degrees
frame : string ('icrs' or 'galactic')
speciefies which frame the Cartesian coordinates is.
Returns
-------
pos : :class:`dask.array.Array`; shape: (N,3)
the cartesian position coordinates, where columns represent
``x``, ``y``, and ``z`` in units of Mpc/h
Raises
------
TypeError
If the input columns are not dask arrays
"""
ra, dec, redshift = da.broadcast_arrays(ra, dec, redshift)
# pos on the unit sphere
pos = SkyToUnitSphere(ra, dec, degrees=degrees, frame=frame)
# multiply by the comoving distance in Mpc/h
r = redshift.map_blocks(lambda z: cosmo.comoving_distance(z), dtype=redshift.dtype)
return r[:,None] * pos + observer
def StackColumns(*cols):
"""
Stack the input dask arrays vertically, column by column.
This uses :func:`dask.array.vstack`.
Parameters
----------
*cols : :class:`dask.array.Array`
the dask arrays to stack vertically together
Returns
-------
:class:`dask.array.Array` :
the dask array where columns correspond to the input arrays
Raises
------
TypeError
If the input columns are not dask arrays
"""
cols = da.broadcast_arrays(*cols)
return da.vstack(cols).T
def StackColumns(*cols):
"""
Stack the input dask arrays vertically, column by column.
This uses :func:`dask.array.vstack`.
Parameters
----------
*cols : :class:`dask.array.Array`
the dask arrays to stack vertically together
Returns
-------
:class:`dask.array.Array` :
the dask array where columns correspond to the input arrays
Raises
------
TypeError
If the input columns are not dask arrays
"""
cols = da.broadcast_arrays(*cols)
return da.vstack(cols).T
def test_dataframe_factory_multicol():
nrow, nfreq, ncorr = 100, 100, 4
data1a = da.random.random(size=nrow, chunks=(10, ))
data1b = da.random.random(size=(nfreq, ncorr), chunks=(100, 4))
data1c = da.random.random(size=(ncorr, ), chunks=(4, ))
df = dataframe_factory(("row", "chan", "corr"), data1a, ("row", ), data1b,
("chan", "corr"), data1c, ("corr", ))
assert isinstance(df, dd.DataFrame)
assert isinstance(df['x'], dd.Series)
assert isinstance(df['y'], dd.Series)
assert isinstance(df['c0'], dd.Series)
x, y, c0 = da.broadcast_arrays(data1a[:, None, None], data1b[None, :, :],
data1c[None, None, :])
assert_array_equal(df['x'], x.ravel())
assert_array_equal(df['y'], y.ravel())
assert_array_equal(df['c0'], c0.ravel())
assert_array_equal(df['x'].min(), data1a.min())
assert_array_equal(df['x'].max(), data1a.max())
assert_array_equal(df['y'].min(), data1b.min())
assert_array_equal(df['y'].max(), data1b.max())
assert_array_equal(df['c0'].min(), data1c.min())
assert_array_equal(df['c0'].max(), data1c.max())
df = df.append(df)
assert_array_equal(df['x'].min(), data1a.min())
assert_array_equal(df['x'].max(), data1a.max())
assert_array_equal(df['y'].min(), data1b.min())
assert_array_equal(df['y'].max(), data1b.max())
assert_array_equal(df['c0'].min(), data1c.min())
assert_array_equal(df['c0'].max(), data1c.max())
def get_plot_data(msinfo, group_cols, mytaql, chan_freqs,
chanslice, subset,
noflags, noconj,
iter_field, iter_spw, iter_scan,
join_corrs=False,
row_chunk_size=100000):
ms_cols = {'ANTENNA1', 'ANTENNA2'}
if not noflags:
ms_cols.update({'FLAG', 'FLAG_ROW'})
# get visibility columns
for axis in DataAxis.all_axes.values():
ms_cols.update(axis.columns)
# get MS data
msdata = daskms.xds_from_ms(msinfo.msname, columns=list(ms_cols), group_cols=group_cols, taql_where=mytaql,
chunks=dict(row=row_chunk_size))
log.info(f': Indexing MS and building dataframes (chunk size is {row_chunk_size})')
np = 0 # number of points to plot
# output dataframes, indexed by (field, spw, scan, antenna, correlation)
# If any of these axes is not being iterated over, then the index is None
output_dataframes = OrderedDict()
# # make prototype dataframe
# import pandas
#
#
# iterate over groups
for group in msdata:
ddid = group.DATA_DESC_ID # always present
fld = group.FIELD_ID # always present
if fld not in subset.field or ddid not in subset.spw:
log.debug(f"field {fld} ddid {ddid} not in selection, skipping")
continue
scan = getattr(group, 'SCAN_NUMBER', None) # will be present if iterating over scans
# TODO: antenna iteration. None forces no iteration, for now
antenna = None
# always read flags -- easier that way
flag = group.FLAG if not noflags else None
flag_row = group.FLAG_ROW if not noflags else None
baselines = group.ANTENNA1*len(msinfo.antenna) + group.ANTENNA2
freqs = chan_freqs[ddid]
chans = xarray.DataArray(range(len(freqs)), dims=("chan",))
wavel = freq_to_wavel(freqs)
extras = dict(chans=chans, freqs=freqs, wavel=wavel, rows=group.row, baselines=baselines)
nchan = len(group.chan)
if flag is not None:
flag = flag[dict(chan=chanslice)]
nchan = flag.shape[1]
shape = (len(group.row), nchan)
datums = OrderedDict()
for corr in subset.corr.numbers:
# make dictionary of extra values for DataMappers
extras['corr'] = corr
# loop over datums to be computed
for axis in DataAxis.all_axes.values():
value = datums[axis.label][-1] if axis.label in datums else None
# a datum was already computed?
if value is not None:
# if not joining correlations, then that's the only one we'll need, so continue
if not join_corrs:
continue
# joining correlations, and datum has a correlation dependence: compute another one
if axis.corr is None:
value = None
if value is None:
value = axis.get_value(group, corr, extras, flag=flag, flag_row=flag_row, chanslice=chanslice)
# reshape values of shape NTIME to (NTIME,1) and NFREQ to (1,NFREQ), and scalar to (NTIME,1)
if value.ndim == 1:
timefreq_axis = axis.mapper.axis or 0
assert value.shape[0] == shape[timefreq_axis], \
f"{axis.mapper.fullname}: size {value.shape[0]}, expected {shape[timefreq_axis]}"
shape1 = [1,1]
shape1[timefreq_axis] = value.shape[0]
value = value.reshape(shape1)
if timefreq_axis > 0:
value = da.broadcast_to(value, shape)
log.debug(f"axis {axis.mapper.fullname} has shape {value.shape}")
# else 2D value better match expected shape
else:
assert value.shape == shape, f"{axis.mapper.fullname}: shape {value.shape}, expected {shape}"
datums.setdefault(axis.label, []).append(value)
# if joining correlations, stick all elements together. Otherwise, we'd better have one per label
if join_corrs:
datums = OrderedDict({label: da.concatenate(arrs) for label, arrs in datums.items()})
else:
assert all([len(arrs) == 1 for arrs in datums.values()])
datums = OrderedDict({label: arrs[0] for label, arrs in datums.items()})
# broadcast to same shape, and unravel all datums
datums = OrderedDict({ key: arr.ravel() for key, arr in zip(datums.keys(),
da.broadcast_arrays(*datums.values()))})
# if any axis needs to be conjugated, double up all of them
if not noconj and any([axis.conjugate for axis in DataAxis.all_axes.values()]):
for axis in DataAxis.all_axes.values():
if axis.conjugate:
datums[axis.label] = da.concatenate([datums[axis.label], -datums[axis.label]])
else:
datums[axis.label] = da.concatenate([datums[axis.label], datums[axis.label]])
labels, values = list(datums.keys()), list(datums.values())
np += values[0].size
# now stack them all into a big dataframe
rectype = [(axis.label, numpy.int32 if axis.nlevels else numpy.float32) for axis in DataAxis.all_axes.values()]
recarr = da.empty_like(values[0], dtype=rectype)
ddf = dask_df.from_array(recarr)
for label, value in zip(labels, values):
ddf[label] = value
# now, are we iterating or concatenating? Make frame key accordingly
dataframe_key = (fld if iter_field else None,
ddid if iter_spw else None,
scan if iter_scan else None,
antenna)
# do we already have a frame for this key
ddf0 = output_dataframes.get(dataframe_key)
if ddf0 is None:
log.debug(f"first frame for {dataframe_key}")
output_dataframes[dataframe_key] = ddf
else:
log.debug(f"appending to frame for {dataframe_key}")
output_dataframes[dataframe_key] = ddf0.append(ddf)
# convert discrete axes into categoricals
if data_mappers.USE_COUNT_CAT:
categorical_axes = [axis.label for axis in DataAxis.all_axes.values() if axis.nlevels]
if categorical_axes:
log.info(": counting colours")
for key, ddf in list(output_dataframes.items()):
output_dataframes[key] = ddf.categorize(categorical_axes)
log.info(": complete")
return output_dataframes, np
def SkyToUnitSphere(ra, dec, degrees=True, frame='icrs'):
"""
Convert sky coordinates (``ra``, ``dec``) to Cartesian coordinates on
the unit sphere.
Parameters
----------
ra : :class:`dask.array.Array`; shape: (N,)
the right ascension angular coordinate
dec : :class:`dask.array.Array`; ; shape: (N,)
the declination angular coordinate
degrees : bool, optional
specifies whether ``ra`` and ``dec`` are in degrees or radians
frame : string ('icrs' or 'galactic')
speciefies which frame the Cartesian coordinates is. Useful if you know
the simulation (usually cartesian) is in galactic units but you want
to convert to the icrs (ra, dec) usually used in surveys.
Returns
-------
pos : :class:`dask.array.Array`; shape: (N,3)
the cartesian position coordinates, where columns represent
``x``, ``y``, and ``z``
Raises
------
TypeError
If the input columns are not dask arrays
"""
ra, dec = da.broadcast_arrays(ra, dec)
if frame == 'icrs':
# no frame transformation
# put into radians from degrees
if degrees:
ra = da.deg2rad(ra)
dec = da.deg2rad(dec)
# cartesian coordinates
x = da.cos( dec ) * da.cos( ra )
y = da.cos( dec ) * da.sin( ra )
z = da.sin( dec )
return da.vstack([x,y,z]).T
else:
from astropy.coordinates import SkyCoord
if degrees:
ra = da.deg2rad(ra)
dec = da.deg2rad(dec)
def eq_to_cart(ra, dec):
try:
sc = SkyCoord(ra, dec, unit='rad', representation_type='unitspherical', frame='icrs')
except:
sc = SkyCoord(ra, dec, unit='rad', representation='unitspherical', frame='icrs')
scg = sc.transform_to(frame=frame)
scg = scg.cartesian
x, y, z = scg.x.value, scg.y.value, scg.z.value
return numpy.stack([x, y, z], axis=1)
arr = da.apply_gufunc(eq_to_cart, '(),()->(p)', ra, dec, output_dtypes=[ra.dtype], output_sizes={'p': 3})
return arr
def SkyToUnitSphere(ra, dec, degrees=True, frame='icrs'):
"""
Convert sky coordinates (``ra``, ``dec``) to Cartesian coordinates on
the unit sphere.
Parameters
----------
ra : :class:`dask.array.Array`; shape: (N,)
the right ascension angular coordinate
dec : :class:`dask.array.Array`; ; shape: (N,)
the declination angular coordinate
degrees : bool, optional
specifies whether ``ra`` and ``dec`` are in degrees or radians
frame : string ('icrs' or 'galactic')
speciefies which frame the Cartesian coordinates is. Useful if you know
the simulation (usually cartesian) is in galactic units but you want
to convert to the icrs (ra, dec) usually used in surveys.
Returns
-------
pos : :class:`dask.array.Array`; shape: (N,3)
the cartesian position coordinates, where columns represent
``x``, ``y``, and ``z``
Raises
------
TypeError
If the input columns are not dask arrays
"""
ra, dec = da.broadcast_arrays(ra, dec)
if frame == 'icrs':
# no frame transformation
# put into radians from degrees
if degrees:
ra = da.deg2rad(ra)
dec = da.deg2rad(dec)
# cartesian coordinates
x = da.cos( dec ) * da.cos( ra )
y = da.cos( dec ) * da.sin( ra )
z = da.sin( dec )
return da.vstack([x,y,z]).T
else:
from astropy.coordinates import SkyCoord
if degrees:
ra = da.deg2rad(ra)
dec = da.deg2rad(dec)
def eq_to_cart(ra, dec):
try:
sc = SkyCoord(ra, dec, unit='rad', representation_type='unitspherical', frame='icrs')
except:
sc = SkyCoord(ra, dec, unit='rad', representation='unitspherical', frame='icrs')
scg = sc.transform_to(frame=frame)
scg = scg.cartesian
x, y, z = scg.x.value, scg.y.value, scg.z.value
return numpy.stack([x, y, z], axis=1)
arr = da.apply_gufunc(eq_to_cart, '(),()->(p)', ra, dec, output_dtypes=[ra.dtype], output_sizes={'p': 3})
return arr
