1024_real_beam.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8aabd29e-5b14-4f1d-9b26-c759c0234eb5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import healpy as hp\n",
    "from tqdm import tqdm\n",
    "import multiprocessing as mp\n",
    "import math\n",
    "from multiprocessing import Pool,Array,Lock\n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ee351e6-2e83-4296-bc0e-f3cd005ac99d",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "theta1 = 7.5 * np.pi / 180\n",
    "theta2 = 85 * np.pi / 180\n",
    "w1 = 2 * np.pi  # rad/min\n",
    "w2 = 2 * w1  # rad/min\n",
    "w3 = 0.000011954  # rad/min\n",
    "\n",
    "def get_vectors(t):\n",
    "    cos_w1t = np.cos(w1 * t)\n",
    "    sin_w1t = np.sin(w1 * t)\n",
    "\n",
    "    cos_w2t = np.cos(w2 * t)\n",
    "    sin_w2t = np.sin(w2 * t)\n",
    "\n",
    "    cos_theta1 = np.cos(theta1)\n",
    "    sin_theta1 = np.sin(theta1)\n",
    "\n",
    "    cos_theta2 = np.cos(theta2)\n",
    "    sin_theta2 = np.sin(theta2)\n",
    "\n",
    "    A = np.array([[cos_w1t, sin_w1t, 0],\n",
    "                  [-sin_w1t, cos_w1t, 0],\n",
    "                  [0, 0, 1]])\n",
    "\n",
    "    B = np.array([[1, 0, 0],\n",
    "                  [0, cos_w2t, sin_w2t],\n",
    "                  [0, -sin_w2t, cos_w2t]])\n",
    "\n",
    "    C = np.array([[cos_theta1, 0, sin_theta1],\n",
    "                  [0, 1, 0],\n",
    "                  [-sin_theta1, 0, cos_theta1]])\n",
    "\n",
    "    D_R = np.array([[cos_theta2],\n",
    "                    [sin_theta2 * np.cos(w3 * t)],\n",
    "                    [sin_theta2 * np.sin(w3 * t)]])\n",
    "\n",
    "    D_S = np.array([[1],\n",
    "                    [0],\n",
    "                    [0]])\n",
    "\n",
    "    result1 = np.dot(np.dot(A, B), C)\n",
    "    result_R = np.matmul(result1, D_R)\n",
    "    result_S = np.matmul(result1, D_S)\n",
    "\n",
    "    return result_R.T.flatten(), result_S.T.flatten()  # Return both flattened vectors\n",
    "\n",
    "\n",
    "#  Angle between two vector\n",
    "\n",
    "def angle_vec(A, B):\n",
    "    dot_product = np.dot(A, B)\n",
    "    mag_A = np.linalg.norm(A)\n",
    "    mag_B = np.linalg.norm(B)\n",
    "    if (mag_A * mag_B) == 0:\n",
    "        return 0 # To handle the case where one the vector becomes Zeros(R_i == Rc)\n",
    "    cos_theta = dot_product / (mag_A * mag_B)\n",
    "    angle = np.arccos(cos_theta)\n",
    "    return angle\n",
    "\n",
    "def anglev(vec1, vec2):\n",
    "    dot_product = np.dot(vec1, vec2)\n",
    "    clipped_dp = np.clip(dot_product, -1.0, 1.0) # Clip dot_product to the valid range for arccos to avoid NaNs\n",
    "    angle = np.arccos(clipped_dp)\n",
    "    return angle\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9a39c37-16bf-4e77-b8be-edddf80fdef6",
   "metadata": {},
   "outputs": [],
   "source": [
    "nside = 1024\n",
    "npix = 12*nside**2\n",
    "\n",
    "# time_step=scan_time\n",
    "scan_time = np.sqrt(4*np.pi/npix)/w1\n",
    "\n",
    "# temperature_map = hp.read_map(\"input_map.fits\")\n",
    "\n",
    "# Load the grid\n",
    "grid = np.loadtxt(\"grid.txt\")\n",
    "# grid_size = 2 (in arcsec)convert grid_size in radian\n",
    "grid_size = 2*math.pi / (180 * 3600)\n",
    "centre = (3001,3001)\n",
    "Radius = (50 / 60) * (math.pi / 180) #50 in arcmin"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e43954e7-e033-48ed-b3b7-f0f472b9b885",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "def process_time_step(time_step,R,S,pix_ring):\n",
    "\n",
    "    #3. Calculate Z, I and N (N = I for phi = 0)\n",
    "    Z_t = np.cross(R,S)\n",
    "    I_t = np.cross(R, Z_t)\n",
    "    N_t = I_t\n",
    "\n",
    "    # 4. Find neighboring pixels in RING format\n",
    "    Rc = hp.pix2vec(nside,pix_ring,nest=False)\n",
    "    neighbours = hp.query_disc(nside, Rc , radius=Radius)\n",
    "\n",
    "    # 5. angular separation between central pixel and neighbouring pixels\n",
    "    x = np.zeros_like(neighbours, dtype=float)\n",
    "    y = np.zeros(len(neighbours))\n",
    "    weight = np.zeros(len(neighbours))\n",
    "    # print(len(neighbours))\n",
    "    for i, neighbour_pix in enumerate(neighbours):\n",
    "\n",
    "        R_i = hp.pix2vec(nside,neighbour_pix,nest=False)\n",
    "        theta_i = anglev(Rc, R_i)\n",
    "\n",
    "        # 6. A_i = line joining central pixel and neighbour pixel\n",
    "        R_i = hp.pix2vec(nside,neighbour_pix,nest=False)\n",
    "        A_i = np.array(Rc)-np.array(R_i)\n",
    "        # print(\"A_i = \",A_i,\"\\nN_t = \",N_t)\n",
    "\n",
    "        # 7. angle between N & A_i\n",
    "        alpha_i = angle_vec(A_i, N_t)\n",
    "        # print(\"alpha_i=\",(alpha_i))\n",
    "        # 8. x_i and y_i\n",
    "        x[i] = theta_i * np.cos(alpha_i)\n",
    "        y[i] = theta_i * np.sin(alpha_i)\n",
    "        index_x = int(centre[0] + round(x[i]/grid_size,0))\n",
    "        index_y = int(centre[1] + round(y[i]/grid_size,0))\n",
    "        weight[i] = grid[index_x][index_y]\n",
    "        # print(centre[0] )\n",
    "    dictionary = {pix: weight[i] for i, pix in enumerate(neighbours)}\n",
    "    result = np.array(list(dictionary.items())).flatten()\n",
    "\n",
    "    return pix_ring, result\n",
    "    # return {pix: weight[i] for i, pix in enumerate(neighbours)}\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17a5e81d-a114-4ace-b5fd-e5bdcde6d170",
   "metadata": {},
   "outputs": [],
   "source": [
    "start = 0\n",
    "duration = 60*24 # in min (one month)\n",
    "steps = int(duration / scan_time)\n",
    "# steps = 10000\n",
    "length = 700 # no. of neigbours\n",
    "\n",
    "\n",
    "time_periods = np.linspace(start, start + duration, steps)\n",
    "# time_periods_iterator = tqdm(time_periods, desc=\"Processing\", total=len(time_periods))\n",
    "\n",
    "\n",
    "# To make an array having  occurance, centre_pix, neighbors&weights (size = 2 + no. of neighbour*2)\n",
    "\n",
    "format_string = '\\t'.join(['%d', '%.8e'] * length)\n",
    "centre_pix_format = \"%d\"\n",
    "occurance_format = \"%d\"\n",
    "\n",
    "\n",
    "fmt = [occurance_format] + [centre_pix_format] + format_string.split('\\t')  # Split format_string by tabs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e6826f9-c93d-40be-9b45-b7c74660800e",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def process_chunk(chunk_range, time_periods):\n",
    "  local_array_map = np.zeros((chunk_range[1] - chunk_range[0], (2 * length + 2)))\n",
    "  local_array_map[:, 1] = np.arange(chunk_range[0], chunk_range[1])\n",
    "  local_array_map[:, 0] = 0\n",
    "  local_array_map[:, 2] = 0\n",
    "  start_time_chunk = time.time()\n",
    "  for time_period in tqdm(time_periods, desc=f\"Processing chunk {chunk_range[0]} to {chunk_range[1]}\"):\n",
    "    R, S = get_vectors(time_period)\n",
    "    pixel = hp.vec2pix(nside, R[0], R[1], R[2], nest=False)\n",
    "    if chunk_range[0] <= pixel < chunk_range[1]:\n",
    "      idx = pixel - chunk_range[0]\n",
    "      pixel, weight = process_time_step(time_period, R, S, pixel)\n",
    "      result = np.pad(weight, (0, 2 * length - len(weight)), mode='constant', constant_values=0)\n",
    "      local_array_map[idx, 2:] += result\n",
    "      local_array_map[idx, 0] += 1  # occurrence = occurrence + 1\n",
    "\n",
    "  # Save local_array_map to a separate file with naming convention based on chunk range\n",
    "  end_time_chunk = time.time()\n",
    "  elapsed_time_chunk = end_time_chunk - start_time_chunk\n",
    "  filename = f\"1024_check_{chunk_range[0]}_{chunk_range[1]}.dat\"\n",
    "  np.savetxt(filename, local_array_map, fmt=fmt)\n",
    "  print(f\"Chunk {chunk_range[0]} to {chunk_range[1]} took {elapsed_time_chunk:.2f} seconds\")\n",
    "  return None  # Since we're saving data directly, no need to return local_array_map\n",
    "\n",
    "def main():\n",
    "  chunk_size = npix // 96  # This will be the base chunk size\n",
    "\n",
    "  chunk_ranges = []\n",
    "  current_start = 0\n",
    "\n",
    "  for i in range(48):\n",
    "    if i < 2 or i >= 46:  # First 2 and last 2 processors\n",
    "      chunks = 1\n",
    "    elif 22 <= i <= 25:  # Processors 23, 24, 25, 26 (0-indexed)\n",
    "      chunks = 3\n",
    "    else:  # All other processors\n",
    "      chunks = 2\n",
    "\n",
    "    for _ in range(chunks):\n",
    "      current_end = current_start + chunk_size\n",
    "      chunk_ranges.append((current_start, current_end))\n",
    "      current_start = current_end\n",
    "\n",
    "  start_time = time.time()\n",
    "\n",
    "  with mp.Pool(processes=48) as pool:\n",
    "    pool.starmap(process_chunk, [(chunk_range, time_periods) for chunk_range in chunk_ranges])\n",
    "\n",
    "  end_time = time.time()\n",
    "  elapsed_time = end_time - start_time\n",
    "\n",
    "  print(\"Elapsed time2:\", elapsed_time, \"seconds\")\n",
    "\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "  main()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "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.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}