Merge pull request #1307 from SamuelSun2000/main

Record the min and max to recover original U

Author: gaow

code/reference_data/rss_ld_sketch.ipynb

@@ -623,12 +623,6 @@
     "        fh.write(f\"n_all_na={n_all_na}\\nn_high_msng={n_high_msng}\\n\")\n",
     "        fh.write(f\"n_low_maf={n_low_maf}\\nn_low_mac={n_low_mac}\\n\")\n",
     "\n",
-    "    # Write .afreq\n",
-    "    with open(afreq_path, \"w\") as fh:\n",
-    "        fh.write(\"#CHROM\\tID\\tREF\\tALT\\tALT_FREQS\\tOBS_CT\\n\")\n",
-    "        for v in var_info:\n",
-    "            fh.write(f\"{v['chr']}\\t{v['id']}\\t{v['ref']}\\t{v['alt']}\\t\"\n",
-    "                     f\"{v['af']:.6f}\\t{v['obs_ct']}\\n\")\n",
     "\n",
     "    # Build G from collected dosages, compute U = $W^T G$, write dosage.gz\n",
     "    # Dosage format=1: ID ALT REF val_S1 ... val_SB\n",
@@ -656,6 +650,16 @@
     "    U = 2.0 * (U - col_min) / denom\n",
     "    U = np.round(U, 4)\n",
     "\n",
+    "    # Record the col min and max for U\n",
+    "\n",
+    "    with open(afreq_path, \"w\") as fh:\n",
+    "        # Add column headers\n",
+    "        fh.write(\"#CHROM\\tID\\tREF\\tALT\\tALT_FREQS\\tOBS_CT\\tU_MIN\\tU_MAX\\n\")\n",
+    "        for j, v in enumerate(var_info):\n",
+    "            fh.write(f\"{v['chr']}\\t{v['id']}\\t{v['ref']}\\t{v['alt']}\\t\"\n",
+    "                     f\"{v['af']:.6f}\\t{v['obs_ct']}\\t\"\n",
+    "                     f\"{col_min[j]:.6f}\\t{col_max[j]:.6f}\\n\")\n",
+    "\n",
     "    with gzip.open(dosage_path, \"wt\", compresslevel=4) as gz:\n",
     "        for j, v in enumerate(var_info):\n",
     "            vals = \" \".join(f\"{x:.4f}\" for x in U[:, j])\n",
@@ -779,6 +783,20 @@
     "        fi\n",
     "    done\n",
     "\n",
+    "    # Step 4: Add U_MIN and U_MAX to the final .pvar file\n",
+    "    # Extract ID, MIN, and MAX from the concatenated .afreq (columns 2, 7, 8)\n",
+    "    awk 'NR > 1 {print $2, $7, $8}' \"${final_prefix}.afreq\" > \"${final_prefix}.recovery_params\"\n",
+    "\n",
+    "    # Decompress the .pvar.zst \n",
+    "    $[plink2_bin] --zst-decompress \"${final_prefix}.pvar.zst\" > \"${final_prefix}.pvar.tmp\"\n",
+    "\n",
+    "    awk 'NR==FNR {min[$1]=$2; max[$1]=$3; next} \n",
+    "         /^#/ {print $0 \"\\tU_MIN\\tU_MAX\"; next} \n",
+    "         {if ($3 in min) print $0 \"\\t\" min[$3] \"\\t\" max[$3]; else print $0 \"\\tNA\\tNA\"}' \\\n",
+    "         \"${final_prefix}.recovery_params\" \"${final_prefix}.pvar.tmp\" > \"${final_prefix}.pvar\"\n",
+    "\n",
+    "    # Replace the compressed file with the updated text file\n",
+    "    rm \"${final_prefix}.pvar.tmp\" \"${final_prefix}.pvar.zst\" \"${final_prefix}.recovery_params\"\n",
     "\n",
     "\n",
     "R: expand = \"$[ ]\"\n",