Halotools_weighted_tcf.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "acknowledged-enclosure",
   "metadata": {},
   "outputs": [],
   "source": [
    "import halotools\n",
    "import numpy as np\n",
    "from astropy.utils.misc import NumpyRNGContext\n",
    "from math import gamma\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "strategic-adrian",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "from halotools.mock_observables.two_point_clustering.clustering_helpers import process_optional_input_sample2\n",
    "\n",
    "from halotools.mock_observables.mock_observables_helpers import (enforce_sample_has_correct_shape,\n",
    "    get_separation_bins_array, get_period, get_num_threads)\n",
    "from halotools.mock_observables.pair_counters.mesh_helpers import _enforce_maximum_search_length\n",
    "\n",
    "from halotools.mock_observables.pair_counters import npairs_3d, marked_npairs_3d\n",
    "\n",
    "from halotools.custom_exceptions import HalotoolsError\n",
    "\n",
    "from halotools.mock_observables.two_point_clustering.tpcf_estimators import _TP_estimator\n",
    "from halotools.mock_observables.two_point_clustering.marked_tpcf import _marked_tpcf_process_args,marked_pair_counts\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "flying-temperature",
   "metadata": {},
   "outputs": [],
   "source": [
    "from halotools.mock_observables.two_point_clustering import tpcf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "wrapped-hughes",
   "metadata": {},
   "outputs": [],
   "source": [
    "def nball_volume(R, k=3):\n",
    "        \"\"\"\n",
    "        Calculate the volume of a n-shpere.\n",
    "        This is used for the analytical randoms.\n",
    "        \"\"\"\n",
    "        return (np.pi**(k/2.0)/gamma(k/2.0+1.0))*R**k #gamma should be math gamma?\n",
    "    \n",
    "def _random_counts(N1,N2,NR, rbins, period, num_threads,\n",
    "        _sample1_is_sample2,do_RR=True, do_DR=True, approx_cell1_size=None,\n",
    "        approx_cell2_size=None, approx_cellran_size=None,PBCs=None):\n",
    "    \"\"\"\n",
    "        modified version of https://halotools.readthedocs.io/en/latest/_modules/halotools/mock_observables/two_point_clustering/tpcf.html#tpcf\n",
    "        \n",
    "        Doing only the analytical version.\n",
    "    \"\"\"\n",
    "#     if randoms is None:\n",
    "#         NR = len(sample1)\n",
    "\n",
    "    # do volume calculations\n",
    "    v = nball_volume(rbins)  # volume of spheres\n",
    "    dv = np.diff(v)  # volume of shells\n",
    "    global_volume = period.prod()  # volume of simulation\n",
    "\n",
    "    # calculate randoms for sample1\n",
    "#         N1 = np.shape(sample1)[0]  # number of points in sample1\n",
    "    rho1 = N1/global_volume  # number density of points\n",
    "    D1R = (NR)*(dv*rho1)  # random counts are N**2*dv*rho\n",
    "\n",
    "    # calculate randoms for sample2\n",
    "#         N2 = np.shape(sample2)[0]  # number of points in sample2\n",
    "    rho2 = N2/global_volume  # number density of points\n",
    "    D2R = (NR)*(dv*rho2)  # random counts are N**2*dv*rho\n",
    "\n",
    "    # calculate the random-random pairs.\n",
    "    rhor = (NR**2)/global_volume\n",
    "    RR = (dv*rhor)\n",
    "\n",
    "    return D1R, D2R, RR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cooperative-monitoring",
   "metadata": {},
   "outputs": [],
   "source": [
    "def weighted_tpcf(sample1, rbins, sample2=None,\n",
    "        weights1=None, weights2=None, period=None, do_auto=True, do_cross=True,\n",
    "        num_threads=1, weight_func_id=1,\n",
    "        normalize_weights=True, seed=None,estimator='Natural'):\n",
    "    \"\"\"\n",
    "    \"\"\"\n",
    "    if normalize_weights:\n",
    "        if weights1 is not None:\n",
    "            weights1=weights1/weights1.mean()\n",
    "        else:\n",
    "            weights1=np.ones(len(sample1))\n",
    "        if weights2 is not None:\n",
    "            weights2=weights2/weights2.mean()\n",
    "        elif sample2 is not None:\n",
    "            weights2=np.ones(len(sample2))\n",
    "    \n",
    "    randomize_marks=False\n",
    "    normalize_by='random_marks'\n",
    "    iterations=1\n",
    "    function_args = (sample1, rbins, sample2, weights1,weights2,\n",
    "        period, do_auto, do_cross, num_threads,\n",
    "        weight_func_id, normalize_by, iterations, randomize_marks, seed)\n",
    "    \n",
    "    sample1, rbins, sample2, weights1,weights2, period, do_auto, do_cross, num_threads,\\\n",
    "        weight_func_id, normalize_by, _sample1_is_sample2, PBCs,\\\n",
    "        randomize_marks = _marked_tpcf_process_args(*function_args)      \n",
    "    \n",
    "    N1 = weights1.sum() #len(sample1)\n",
    "    N2=N1*1.\n",
    "    if sample2 is not None:\n",
    "        N2 = weights2.sum() #len(sample2)\n",
    "    NR=N1*1.\n",
    "\n",
    "    # calculate weighted pairs\n",
    "    W1W1, W1W2, W2W2 = marked_pair_counts(sample1, sample2, rbins, period,\n",
    "        num_threads, do_auto, do_cross, weights1, weights2, weight_func_id, _sample1_is_sample2)\n",
    "    randoms=None #use analytical calculations\n",
    "    \n",
    "    D1R, D2R, RR= _random_counts(N1=N1,N2=N2,NR=NR, rbins=rbins, period=period, num_threads= num_threads,\n",
    "        _sample1_is_sample2=_sample1_is_sample2,do_RR=True, do_DR=True,)\n",
    "    \n",
    "    print('weighted number of galaxies: ',N1,N2,NR)\n",
    "    \n",
    "    print('pair_counts: ',W1W1.sum()/1.e6,RR.sum()/1.e6,D1R.sum()/1.e6)\n",
    "    # run results through the estimator and return relavent/user specified results.\n",
    "    if _sample1_is_sample2:\n",
    "        xi_11 = _TP_estimator(W1W1, D1R, RR, N1, N1, NR, NR, estimator)\n",
    "        return xi_11\n",
    "    else:\n",
    "        if (do_auto is True) & (do_cross is True):\n",
    "            xi_11 = _TP_estimator(W1W1, D1R, RR, N1, N1, NR, NR, estimator)\n",
    "            xi_12 = _TP_estimator(W1W2, D1R, RR, N1, N2, NR, NR, estimator)\n",
    "            xi_22 = _TP_estimator(W2W2, D2R, RR, N2, N2, NR, NR, estimator)\n",
    "            return xi_11, xi_12, xi_22\n",
    "        elif (do_cross is True):\n",
    "            xi_12 = _TP_estimator(W1W2, D1R, RR, N1, N2, NR, NR, estimator)\n",
    "            return xi_12\n",
    "        elif (do_auto is True):\n",
    "            xi_11 = _TP_estimator(W1W1, D1R, RR, N1, N1, NR, NR, estimator)\n",
    "            xi_22 = _TP_estimator(W2W2, D2R, RR, N2, N2, NR, NR, estimator)\n",
    "            return xi_11, xi_22"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "antique-cooper",
   "metadata": {},
   "outputs": [],
   "source": [
    "N_gal1=10000\n",
    "N_gal2=20000\n",
    "\n",
    "period=40\n",
    "\n",
    "sample1=np.random.uniform(0,period,N_gal1*3).reshape(N_gal1,3)\n",
    "sample2=np.random.uniform(0,period,N_gal2*3).reshape(N_gal2,3)\n",
    "\n",
    "weights1=np.random.uniform(1,20,N_gal1)\n",
    "weights2=np.random.uniform(1,20,N_gal2)\n",
    "\n",
    "rbins=np.logspace(-1,1,20)\n",
    "\n",
    "# The default un-weighted calculation for testing\n",
    "xi0_11, xi0_12, xi0_22=tpcf(sample1, rbins, sample2=sample2,\n",
    "        period=period, do_auto=True, do_cross=True,\n",
    "        num_threads=1, seed=None,estimator='Natural') "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "demonstrated-collective",
   "metadata": {},
   "outputs": [],
   "source": [
    "#weighted calculation\n",
    "xi_11, xi_12, xi_22=weighted_tpcf(sample1=sample1, rbins=rbins, sample2=sample2,\n",
    "        weights1=weights1, weights2=weights2, period=period, do_auto=True, do_cross=True,\n",
    "        num_threads=1, weight_func_id=1,\n",
    "         normalize_weights=True, seed=None,estimator='Natural')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "comic-columbus",
   "metadata": {},
   "outputs": [],
   "source": [
    "rb=0.5*(rbins[1:]+rbins[:-1])\n",
    "plot(rb,xi_11)\n",
    "plot(rb,xi_12)\n",
    "plot(rb,xi_22)\n",
    "hlines(0,1,10,color='k')\n",
    "# loglog()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "sunset-string",
   "metadata": {},
   "outputs": [],
   "source": [
    "rb=0.5*(rbins[1:]+rbins[:-1])\n",
    "plot(rb,xi0_11)\n",
    "plot(rb,xi0_12)\n",
    "plot(rb,xi0_22)\n",
    "hlines(0,1,10,color='k')\n",
    "# loglog()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "oriented-contact",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cubic-compact",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda",
   "language": "python",
   "name": "conda"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}