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 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 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 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