Box units conversion (#3)

* linting

* fix long line

* rename nb

* rename import

* add group catalog

* move imports out of functions

* add array_to_structured

* add references

* fix subsampling

* fix coord bug

* add 2M++ dists

* save nb

* fix comment

* add snapshot path

* add snapshot path

* add read_info

* Move transforms

* add import radec

* expand docs

* Move flipcols

* update nb

* add flip_cols

* create file

* add blank line

* Move units transfs

* add blank line

* add units import

* rm imports

* add import

* add box_units

* add comments

galomatch/__init__.py

@@ -13,4 +13,4 @@
 # with this program; if not, write to the Free Software Foundation, Inc.,
 # 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
 
-from galomatch import (io, match, utils)
+from galomatch import (io, match, utils, units)  # noqa

galomatch/fits/__init__.py

@@ -0,0 +1,14 @@
+# Copyright (C) 2022 Richard Stiskalek
+# This program is free software; you can redistribute it and/or modify it
+# under the terms of the GNU General Public License as published by the
+# Free Software Foundation; either version 3 of the License, or (at your
+# option) any later version.
+#
+# This program is distributed in the hope that it will be useful, but
+# WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General
+# Public License for more details.
+#
+# You should have received a copy of the GNU General Public License along
+# with this program; if not, write to the Free Software Foundation, Inc.,
+# 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.

galomatch/io/__init__.py

@@ -13,8 +13,8 @@
 # with this program; if not, write to the Free Software Foundation, Inc.,
 # 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
 
-from .readsim import (get_csiborg_ids, get_sim_path, open_particle,
+from .readsim import (get_csiborg_ids, get_sim_path, get_snapshot_path,  # noqa
-                      open_unbinding, read_particle, read_clumpid, read_clumps,
+                      read_info,  # noqa
-                      read_mmain,
+                      open_particle, open_unbinding, read_particle,  # noqa
-                      convert_mass_cols, convert_position_cols, flip_cols)
+                      read_clumpid, read_clumps, read_mmain)  # noqa
-from .readobs import (read_planck2015, read_2mpp)
+from .readobs import (read_planck2015, read_2mpp)  # noqa

galomatch/io/readobs.py

@@ -69,7 +69,7 @@ def read_planck2015(fpath, dist_cosmo, max_comdist=None):
     return out
 
 
-def read_2mpp(fpath):
+def read_2mpp(fpath, dist_cosmo):
     """
     Read in the 2M++ galaxy redshift catalogue [1], with the catalogue at [2].
     Removes fake galaxies used to fill the zone of avoidance.
@@ -83,16 +83,24 @@ def read_2mpp(fpath):
     -------
     out : structured array
         The catalogue.
+
+    References
+    ----------
+    [1] The 2M++ galaxy redshift catalogue; Lavaux, Guilhem, Hudson, Michael J.
+    [2] https://cdsarc.cds.unistra.fr/viz-bin/cat/J/MNRAS/416/2840#/article
     """
+    from scipy.constants import c
     # Read the catalogue and select non-fake galaxies
     cat = numpy.genfromtxt(fpath, delimiter="|", )
     cat = cat[cat[:, 12] == 0, :]
 
     F64 = numpy.float64
-    cols = [("RA", F64), ("DEC", F64), ("Ksmag", F64)]
+    cols = [("RA", F64), ("DEC", F64), ("Ksmag", F64), ("ZCMB", F64),
+            ("CDIST_CMB", F64)]
     out = cols_to_structured(cat.shape[0], cols)
-    out["RA"] = cat[:, 1] - 180
+    out["RA"] = cat[:, 1]
     out["DEC"] = cat[:, 2]
     out["Ksmag"] = cat[:, 5]
-
+    out["ZCMB"] = cat[:, 7] / (c * 1e-3)
+    out["CDIST_CMB"] = dist_cosmo.comoving_distance(out["ZCMB"]).value
     return out

galomatch/io/readsim.py

@@ -30,9 +30,6 @@ F32 = numpy.float32
 F64 = numpy.float64
 I32 = numpy.int32
 I64 = numpy.int64
-little_h = 0.705
-BOXSIZE = 677.7 / little_h  # Mpc. Otherwise positions in [0, 1].
-BOXMASS = 3.749e19  # Msun
 
 
 def get_csiborg_ids(srcdir):
@@ -82,11 +79,62 @@ def get_sim_path(n, fname="ramses_out_{}", srcdir="/mnt/extraspace/hdesmond"):
     Returns
     -------
     path : str
-        The complete path to the `n`th CSiBORG simulation.
+        Path to the `n`th CSiBORG simulation.
     """
     return join(srcdir, fname.format(n))
 
 
+def get_snapshot_path(Nsnap, simpath):
+    """
+    Get a path to a CSiBORG IC realisation snapshot.
+
+    Parameters
+    ----------
+    Nsnap : int
+        Snapshot index.
+    simpath : str
+        Path to the CSiBORG IC realisation.
+
+    Returns
+    -------
+    snappath : str
+        Path to the CSiBORG IC realisation snapshot.
+    """
+    return join(simpath, "output_{}".format(str(Nsnap).zfill(5)))
+
+
+def read_info(Nsnap, simpath):
+    """
+    Read CSiBORG simulation snapshot info.
+
+    Parameters
+    ----------
+    Nsnap : int
+        Snapshot index.
+    simpath : str
+        Path to the CSiBORG IC realisation.
+
+    Returns
+    -------
+    info : dict
+        Dictionary of info paramaters. Note that both keys and values are
+        strings.
+    """
+    # Open the info file
+    snappath = get_snapshot_path(Nsnap, simpath)
+    filename = join(snappath, "info_{}.txt".format(str(Nsnap).zfill(5)))
+    with open(filename, "r") as f:
+        info = f.read().split()
+    # Throw anything below ordering line out
+    info = numpy.asarray(info[:info.index("ordering")])
+    # Get indexes of lines with `=`. Indxs before/after be keys/vals
+    eqindxs = numpy.asarray([i for i in range(info.size) if info[i] == '='])
+
+    keys = info[eqindxs - 1]
+    vals = info[eqindxs + 1]
+    return {key: val for key, val in zip(keys, vals)}
+
+
 def open_particle(n, simpath, verbose=True):
     """
     Open particle files to a given CSiBORG simulation.
@@ -109,11 +157,9 @@ def open_particle(n, simpath, verbose=True):
     """
     # Zeros filled snapshot number and the snapshot path
     nout = str(n).zfill(5)
-    snappath = join(simpath, "output_{}".format(nout))
+    snappath = get_snapshot_path(n, simpath)
-    infopath = join(snappath, "info_{}.txt".format(nout))
+    ncpu = int(read_info(n, simpath)["ncpu"])
 
-    with open(infopath, "r") as f:
-        ncpu = int(f.readline().split()[-1])
     if verbose:
         print("Reading in output `{}` with ncpu = `{}`.".format(nout, ncpu))
 
@@ -136,6 +182,7 @@ def open_particle(n, simpath, verbose=True):
         # Read in this order
         ncpuloc = f.read_ints()
         if ncpuloc != ncpu:
+            infopath = join(snappath, "info_{}.txt".format(nout))
             raise ValueError("`ncpu = {}` of `{}` disagrees with `ncpu = {}` "
                              "of `{}`.".format(ncpu, infopath, ncpuloc, fpath))
         ndim = f.read_ints()
@@ -382,74 +429,3 @@ def read_mmain(n, srcdir, fname="Mmain_{}.npy"):
         out[name] = arr[:, i]
 
     return out
-
-
-def convert_mass_cols(arr, cols):
-    r"""
-    Convert mass columns from box units to :math:`M_{\odot}`. `arr` is passed
-    by reference and is not explicitly returned back.
-
-    Parameters
-    ----------
-    arr : structured array
-        The array whose columns are to be converted.
-    cols : str or list of str
-        The mass columns to be converted.
-
-    Returns
-    -------
-    None
-    """
-    cols = [cols] if isinstance(cols, str) else cols
-    for col in cols:
-        arr[col] *= BOXMASS
-
-
-def convert_position_cols(arr, cols, zero_centered=True):
-    r"""
-    Convert position columns from box units to :math:`\mathrm{Mpc}`. `arr` is
-    passed by reference and is not explicitly returned back.
-
-    Parameters
-    ----------
-    arr : structured array
-        The array whose columns are to be converted.
-    cols : str or list of str
-        The mass columns to be converted.
-    zero_centered : bool, optional
-        Whether to translate the well-resolved origin in the centre of the
-        simulation to the :math:`(0, 0 , 0)` point. By default `True`.
-
-    Returns
-    -------
-    None
-    """
-    cols = [cols] if isinstance(cols, str) else cols
-    for col in cols:
-        arr[col] *= BOXSIZE
-        if zero_centered:
-            arr[col] -= BOXSIZE / 2
-
-
-def flip_cols(arr, col1, col2):
-    """
-    Flip values in columns `col1` and `col2`. `arr` is passed by reference and
-    is not explicitly returned back.
-
-
-    Parameters
-    ----------
-    arr : structured array
-        The array whose columns are to be converted.
-    col1 : str
-        The first column name.
-    col2 : str
-        The second column name.
-
-    Returns
-    -------
-    nothing
-    """
-    dum = numpy.copy(arr[col1])
-    arr[col1] = arr[col2]
-    arr[col2] = dum

galomatch/match/__init__.py

@@ -13,5 +13,5 @@
 # with this program; if not, write to the Free Software Foundation, Inc.,
 # 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
 
-from .match import brute_spatial_separation
+from .match import brute_spatial_separation  # noqa
-from .correlation import (get_randoms_sphere, angular_tpcf)
+from .correlation import (get_randoms_sphere, sphere_angular_tpcf) # noqa

galomatch/match/correlation.py

@@ -72,7 +72,7 @@ def sphere_angular_tpcf(bins, RA1, DEC1, RA2=None, DEC2=None, nthreads=1,
     """
     Calculate the angular two-point correlation function. The coordinates must
     be provided in degrees. With the right ascension and degrees being
-    in range of :math:`[-180, 180]` and :math:`[-90, 90]` degrees, respectively.
+    in range of :math:`[-180, 180]` and :math:`[-90, 90]` degrees.
     If `RA2` and `DEC2` are provided cross-correlates the first data set with
     the second. Creates a uniformly sampled randoms on the surface of a sphere
     of size `Nmult` times the corresponding number of data points. Uses the

galomatch/units/__init__.py

@@ -0,0 +1,18 @@
+# Copyright (C) 2022 Richard Stiskalek
+# This program is free software; you can redistribute it and/or modify it
+# under the terms of the GNU General Public License as published by the
+# Free Software Foundation; either version 3 of the License, or (at your
+# option) any later version.
+#
+# This program is distributed in the hope that it will be useful, but
+# WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General
+# Public License for more details.
+#
+# You should have received a copy of the GNU General Public License along
+# with this program; if not, write to the Free Software Foundation, Inc.,
+# 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
+
+from .transforms import (cartesian_to_radec, convert_mass_cols,  # noqa
+                         convert_position_cols)  # noqa
+from .box_units import BoxUnits  # noqa

galomatch/units/box_units.py

@@ -0,0 +1,166 @@
+# Copyright (C) 2022 Richard Stiskalek, Deaglan Bartlett
+# This program is free software; you can redistribute it and/or modify it
+# under the terms of the GNU General Public License as published by the
+# Free Software Foundation; either version 3 of the License, or (at your
+# option) any later version.
+#
+# This program is distributed in the hope that it will be useful, but
+# WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General
+# Public License for more details.
+#
+# You should have received a copy of the GNU General Public License along
+# with this program; if not, write to the Free Software Foundation, Inc.,
+# 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
+"""
+Simulation box unit transformations.
+"""
+
+
+from astropy.cosmology import LambdaCDM
+from astropy import (constants, units)
+from ..io import read_info
+
+
+# Conversion factors
+MSUNCGS = constants.M_sun.cgs.value
+KPC_TO_CM = 3.08567758149137e21
+PI = 3.1415926535897932384626433
+
+
+class BoxUnits:
+    """
+    Box units class for converting between box and physical units.
+
+    Paramaters
+    ----------
+    Nsnap : int
+        Snapshot index.
+    simpath : str
+        Path to the simulation where its snapshot index folders are stored.
+    """
+
+    def __init__(self, Nsnap, simpath):
+        """
+        Read in the snapshot info file and set the units from it.
+        """
+        info = read_info(Nsnap, simpath)
+        pars = ["boxlen", "time", "aexp", "H0",
+                "omega_m", "omega_l", "omega_k", "omega_b",
+                "unit_l", "unit_d", "unit_t"]
+        for par in pars:
+            setattr(self, par, float(info[par]))
+
+        self.h = self.H0 / 100
+        self.cosmo = LambdaCDM(H0=self.H0, Om0=self.omega_m, Ode0=self.omega_l,
+                               Tcmb0=2.725 * units.K, Ob0=self.omega_b)
+        # Constants in box units
+        self.G = constants.G.cgs.value * (self.unit_d * self.unit_t ** 2)
+        self.H0 = self.H0 * 1e5 / (1e3 * KPC_TO_CM) * self.unit_t
+        self.c = constants.c.cgs.value * self.unit_t / self.unit_l
+        self.rho_crit = 3 * self.H0 ** 2 / (8 * PI * self.G)
+
+    def box2kpc(self, length):
+        r"""
+        Convert length from box units to :math:`\mathrm{kpc}`.
+
+        Parameters
+        ----------
+        length : float
+            Length in box units.
+
+        Returns
+        -------
+        length : foat
+            Length in :math:`\mathrm{kpc}`
+        """
+        return length * self.unit_l / KPC_TO_CM
+
+    def kpc2box(self, length):
+        r"""
+        Convert length from :math:`\mathrm{kpc}` to box units.
+
+        Parameters
+        ----------
+        length : float
+            Length in :math:`\mathrm{kpc}`
+
+        Returns
+        -------
+        length : foat
+            Length in box units.
+        """
+        return length / self.unit_l * KPC_TO_CM
+
+    def solarmass2box(self, mass):
+        r"""
+        Convert mass from :math:`M_\odot` to box units.
+
+        Parameters
+        ----------
+        mass : float
+            Mass in :math:`M_\odot`.
+
+        Returns
+        -------
+        mass : float
+            Mass in box units.
+        """
+        m = mass * MSUNCGS   # In cgs
+        unit_m = self.unit_d * self.unit_l ** 3
+        return m / unit_m
+
+    def box2solarmass(self, mass):
+        r"""
+        Convert mass from box units to :math:`M_\odot`.
+
+        TODO: check this.
+
+        Parameters
+        ----------
+        mass : float
+            Mass in box units.
+
+        Returns
+        -------
+        mass : float
+            Mass in :math:`M_\odot`.
+        """
+        unit_m = self.unit_d * self.unit_l**3
+        m = mass * unit_m  # In cgs
+        m = m / MSUNCGS
+        return m
+
+    def box2dens(self, density):
+        r"""
+        Convert density from box units to :math:`M_\odot / \mathrm{pc}^3`.
+
+        TODO: check this.
+
+        Parameters
+        ----------
+        density : float
+            Density in box units.
+        box : `BoxConstants`
+            Simulation box class with units.
+
+        Returns
+        -------
+        density : float
+            Density in :math:`M_\odot / \mathrm{pc}^3`.
+        """
+        rho = density * self.unit_d  # In cgs
+        rho = rho * (KPC_TO_CM * 1e-3)**3  # In g/pc^3
+        rho = rho / MSUNCGS
+        return rho
+
+    def dens2box(self, density):
+        r"""
+        Convert density from M_sun / pc^3
+
+        TODO: check this and write documentation.
+        """
+        rho = density * MSUNCGS
+        rho = rho / (KPC_TO_CM * 1e-3)**3  # In g/cm^3
+        rho = rho / self.unit_d
+        return rho

galomatch/units/transforms.py

@@ -0,0 +1,110 @@
+# Copyright (C) 2022 Richard Stiskalek
+# This program is free software; you can redistribute it and/or modify it
+# under the terms of the GNU General Public License as published by the
+# Free Software Foundation; either version 3 of the License, or (at your
+# option) any later version.
+#
+# This program is distributed in the hope that it will be useful, but
+# WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General
+# Public License for more details.
+#
+# You should have received a copy of the GNU General Public License along
+# with this program; if not, write to the Free Software Foundation, Inc.,
+# 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
+"""
+Various coordinate transformations.
+"""
+
+
+import numpy
+
+
+little_h = 0.705
+BOXSIZE = 677.7 / little_h  # Mpc. Otherwise positions in [0, 1].
+BOXMASS = 3.749e19  # Msun
+
+
+def cartesian_to_radec(arr, xpar="peak_x", ypar="peak_y", zpar="peak_z"):
+    r"""
+    Extract `x`, `y`, and `z` coordinates from a record array `arr` and
+    calculate the radial distance :math:`r` in coordinate units, right
+    ascension :math:`\mathrm{RA} \in [0, 360)` degrees and declination
+    :math:`\delta \in [-90, 90]` degrees.
+
+    Parameters
+    ----------
+    arr : record array
+        Record array with the Cartesian coordinates.
+    xpar : str, optional
+        Name of the x coordinate in the record array.
+    ypar : str, optional
+        Name of the y coordinate in the record array.
+    zpar : str, optional
+        Name of the z coordinate in the record array.
+
+    Returns
+    -------
+    dist : 1-dimensional array
+        Radial distance.
+    ra : 1-dimensional array
+        Right ascension.
+    dec : 1-dimensional array
+        Declination.
+    """
+    x, y, z = arr[xpar], arr[ypar], arr[zpar]
+
+    dist = numpy.sqrt(x**2 + y**2 + z**2)
+    dec = numpy.rad2deg(numpy.arcsin(z/dist))
+    ra = numpy.rad2deg(numpy.arctan2(y, x))
+    # Make sure RA in the correct range
+    ra[ra < 0] += 360
+
+    return dist, ra, dec
+
+
+def convert_mass_cols(arr, cols):
+    r"""
+    Convert mass columns from box units to :math:`M_{\odot}`. `arr` is passed
+    by reference and is not explicitly returned back.
+
+    Parameters
+    ----------
+    arr : structured array
+        The array whose columns are to be converted.
+    cols : str or list of str
+        The mass columns to be converted.
+
+    Returns
+    -------
+    None
+    """
+    cols = [cols] if isinstance(cols, str) else cols
+    for col in cols:
+        arr[col] *= BOXMASS
+
+
+def convert_position_cols(arr, cols, zero_centered=True):
+    r"""
+    Convert position columns from box units to :math:`\mathrm{Mpc}`. `arr` is
+    passed by reference and is not explicitly returned back.
+
+    Parameters
+    ----------
+    arr : structured array
+        The array whose columns are to be converted.
+    cols : str or list of str
+        The mass columns to be converted.
+    zero_centered : bool, optional
+        Whether to translate the well-resolved origin in the centre of the
+        simulation to the :math:`(0, 0 , 0)` point. By default `True`.
+
+    Returns
+    -------
+    None
+    """
+    cols = [cols] if isinstance(cols, str) else cols
+    for col in cols:
+        arr[col] *= BOXSIZE
+        if zero_centered:
+            arr[col] -= BOXSIZE / 2

galomatch/utils/__init__.py

@@ -13,6 +13,6 @@
 # with this program; if not, write to the Free Software Foundation, Inc.,
 # 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
 
-from .recarray_manip import (cols_to_structured, add_columns, rm_columns,
+from .recarray_manip import (cols_to_structured, add_columns, rm_columns,  # noqa
-                             list_to_ndarray)
+                             list_to_ndarray, array_to_structured,  # noqa
-from .transforms import cartesian_to_radec
+                             flip_cols)  # noqa

galomatch/utils/recarray_manip.py

@@ -12,11 +12,14 @@
 # You should have received a copy of the GNU General Public License along
 # with this program; if not, write to the Free Software Foundation, Inc.,
 # 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
+"""
+Utilility functions for manipulation structured arrays.
+"""
 
-"""Utilility functions for manipulation structured arrays."""
 
 import numpy
 
+
 def cols_to_structured(N, cols):
     """
     Allocate a structured array from `cols`.
@@ -84,6 +87,7 @@ def add_columns(arr, X, cols):
 
     return out
 
+
 def rm_columns(arr, cols):
     """
     Remove columns `cols` from a record array `arr`. Creates a new array.
@@ -143,7 +147,7 @@ def list_to_ndarray(arrs, cols):
         raise TypeError("`arrs` must be a list of structured arrays.")
     cols = [cols] if isinstance(cols, str) else cols
 
-    Narr  = len(arrs)
+    Narr = len(arrs)
     Nobj_max = max([arr.size for arr in arrs])
     Ncol = len(cols)
     # Preallocate the array and fill it
@@ -153,3 +157,55 @@ def list_to_ndarray(arrs, cols):
         for j in range(Ncol):
             out[i, :Nobj, j] = arrs[i][cols[j]]
     return out
+
+
+def array_to_structured(arr, cols):
+    """
+    Create a structured array from a 2-dimensional array.
+
+    Parameters
+    ----------
+    arr : 2-dimensional array
+        Original array of shape `(n_samples, n_cols)`.
+    cols : list of str
+        Columns of the structured array
+
+    Returns
+    -------
+    out : structured array
+        The output structured array.
+    """
+    cols = [cols] if isinstance(cols, str) else cols
+    if arr.ndim != 2 and arr.shape[1] != len(cols):
+        raise TypeError("`arr` must be a 2-dimensional array of "
+                        "shape `(n_samples, n_cols)`.")
+
+    dtype = {"names": cols, "formats": [arr.dtype] * len(cols)}
+    out = numpy.full(arr.shape[0], numpy.nan, dtype=dtype)
+    for i, col in enumerate(cols):
+        out[col] = arr[:, i]
+
+    return out
+
+
+def flip_cols(arr, col1, col2):
+    """
+    Flip values in columns `col1` and `col2`. `arr` is passed by reference and
+    is not explicitly returned back.
+
+    Parameters
+    ----------
+    arr : structured array
+        The array whose columns are to be converted.
+    col1 : str
+        The first column name.
+    col2 : str
+        The second column name.
+
+    Returns
+    -------
+    nothing
+    """
+    dum = numpy.copy(arr[col1])
+    arr[col1] = arr[col2]
+    arr[col2] = dum

galomatch/utils/transforms.py

@@ -1,56 +0,0 @@
-# Copyright (C) 2022 Richard Stiskalek
-# This program is free software; you can redistribute it and/or modify it
-# under the terms of the GNU General Public License as published by the
-# Free Software Foundation; either version 3 of the License, or (at your
-# option) any later version.
-#
-# This program is distributed in the hope that it will be useful, but
-# WITHOUT ANY WARRANTY; without even the implied warranty of
-# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General
-# Public License for more details.
-#
-# You should have received a copy of the GNU General Public License along
-# with this program; if not, write to the Free Software Foundation, Inc.,
-# 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
-
-import numpy
-
-
-def cartesian_to_radec(arr, xpar="peak_x", ypar="peak_y", zpar="peak_z", degrees=True):
-    """
-    Extract `x`, `y`, and `z` coordinates from a record array `arr` and
-    calculate their spherical coordinates representation.
-
-    Parameters
-    ----------
-    arr : record array
-        Record array with the Cartesian coordinates.
-    xpar : str, optional
-        Name of the x coordinate in the record array.
-    ypar : str, optional
-        Name of the y coordinate in the record array.
-    zpar : str, optional
-        Name of the z coordinate in the record array.
-    degrees : bool, optional
-        Whether to return angles in degrees. By default `True`.
-
-    Returns
-    -------
-    dist : 1-dimensional array
-        Radial distance.
-    ra : 1-dimensional array
-        Right ascension.
-    dec : 1-dimensional array
-        Declination.
-    """
-    x, y, z = arr[xpar], arr[ypar], arr[zpar]
-
-    dist = numpy.sqrt(x**2 + y**2 + z**2)
-    dec = numpy.arcsin(z / dist)
-    ra = numpy.arctan2(y, x)
-
-    if degrees:
-        dec = numpy.rad2deg(dec)
-        ra = numpy.rad2deg(ra)
-
-    return dist, ra, dec

scripts/utils.py

@@ -12,8 +12,14 @@
 # You should have received a copy of the GNU General Public License along
 # with this program; if not, write to the Free Software Foundation, Inc.,
 # 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
+"""
+Notebook utility functions.
+"""
 
-"""Notebook utility funnctions."""
+
+import numpy
+from tqdm import trange
+from astropy.cosmology import FlatLambdaCDM
 
 try:
     import galomatch
@@ -26,9 +32,9 @@ def load_mmain_convert(n):
     srcdir = "/users/hdesmond/Mmain"
     arr = galomatch.io.read_mmain(n, srcdir)
 
-    galomatch.io.convert_mass_cols(arr, "mass_cl")
+    galomatch.utils.convert_mass_cols(arr, "mass_cl")
-    galomatch.io.convert_position_cols(arr, ["peak_x", "peak_y", "peak_z"])
+    galomatch.utils.convert_position_cols(arr, ["peak_x", "peak_y", "peak_z"])
-    galomatch.io.flip_cols(arr, "peak_x", "peak_z")
+    galomatch.utils.flip_cols(arr, "peak_x", "peak_z")
 
     d, ra, dec = galomatch.utils.cartesian_to_radec(arr)
     arr = galomatch.utils.add_columns(arr, [d, ra, dec], ["dist", "ra", "dec"])
@@ -36,20 +42,25 @@ def load_mmain_convert(n):
 
 
 def load_mmains(N=None, verbose=True):
-    from tqdm import tqdm
     ids = galomatch.io.get_csiborg_ids("/mnt/extraspace/hdesmond")
     N = ids.size if N is None else N
     if N > ids.size:
         raise ValueError("`N` cannot be larger than 101.")
+    # If N less than num of CSiBORG, then radomly choose
+    if N == ids.size:
+        choices = numpy.arange(N)
+    else:
+        choices = numpy.random.choice(ids.size, N, replace=False)
+
     out = [None] * N
-    iters = tqdm(range(N)) if verbose else range(N)
+    iters = trange(N) if verbose else range(N)
     for i in iters:
-        out[i] = load_mmain_convert(ids[i])
+        j = choices[i]
+        out[i] = load_mmain_convert(ids[j])
     return out
 
 
 def load_planck2015(max_comdist=214):
-    from astropy.cosmology import FlatLambdaCDM
     cosmo = FlatLambdaCDM(H0=70.5, Om0=0.307, Tcmb0=2.728)
     fpath = ("/mnt/zfsusers/rstiskalek/galomatch/"
              + "data/HFI_PCCS_SZ-union_R2.08.fits")
@@ -57,4 +68,5 @@ def load_planck2015(max_comdist=214):
 
 
 def load_2mpp():
-    return galomatch.io.read_2mpp("../data/2M++_galaxy_catalog.dat")
+    cosmo = FlatLambdaCDM(H0=70.5, Om0=0.307, Tcmb0=2.728)
+    return galomatch.io.read_2mpp("../data/2M++_galaxy_catalog.dat", cosmo)