Compute projected correlation with weights (#14)

* Add projected correlation test with weights

tests/corrgi/conftest.py

@@ -86,6 +86,26 @@ def acf_rans_weight_catalog(hipscat_catalogs_dir):
     return lsdb.read_hipscat(hipscat_catalogs_dir / "acf_rans_weight")
 
 
+@pytest.fixture
+def pcf_gals_weight_dir(hipscat_catalogs_dir):
+    return hipscat_catalogs_dir / "pcf_gals_weight"
+
+
+@pytest.fixture
+def pcf_gals_weight_catalog(pcf_gals_weight_dir):
+    return lsdb.read_hipscat(pcf_gals_weight_dir)
+
+
+@pytest.fixture
+def pcf_rans_weight_dir(hipscat_catalogs_dir):
+    return hipscat_catalogs_dir / "pcf_rans_weight"
+
+
+@pytest.fixture
+def pcf_rans_weight_catalog(pcf_rans_weight_dir):
+    return lsdb.read_hipscat(pcf_rans_weight_dir)
+
+
 @pytest.fixture
 def acf_bins_left_edges(acf_expected_results):
     return np.load(acf_expected_results / "l_binedges_acf.npy")
@@ -126,6 +146,16 @@ def pcf_rr_counts(pcf_expected_results):
     return np.load(pcf_expected_results / "rr_pcf.npy")
 
 
+@pytest.fixture
+def pcf_dd_counts_with_weights(pcf_expected_results):
+    return np.load(pcf_expected_results / "dd_pcf_weight.npy")
+
+
+@pytest.fixture
+def pcf_rr_counts_with_weights(pcf_expected_results):
+    return np.load(pcf_expected_results / "rr_pcf_weight.npy")
+
+
 @pytest.fixture
 def acf_nat_estimate(acf_expected_results):
     return np.load(acf_expected_results / "w_acf_nat.npy")

tests/corrgi/test_acf.py

@@ -22,12 +22,7 @@ def test_acf_bins_are_correct(acf_bins_left_edges, acf_bins_right_edges, acf_par
 
 
 def test_acf_counts_are_correct(
-    dask_client,
-    data_catalog,
-    rand_catalog,
-    acf_dd_counts,
-    acf_rr_counts,
-    acf_params,
+    dask_client, data_catalog, rand_catalog, acf_dd_counts, acf_rr_counts, acf_params
 ):
     ang_corr = AngularCorrelation(params=acf_params)
     counts_dd, counts_rr = compute_autocorrelation_counts(
@@ -38,11 +33,7 @@ def test_acf_counts_are_correct(
 
 
 def test_acf_natural_estimate_is_correct(
-    data_catalog_dir,
-    rand_catalog_dir,
-    acf_dd_counts,
-    acf_rr_counts,
-    acf_nat_estimate,
+    data_catalog_dir, rand_catalog_dir, acf_dd_counts, acf_rr_counts, acf_nat_estimate
 ):
     galaxy_hc_catalog = hipscat.read_from_hipscat(data_catalog_dir)
     random_hc_catalog = hipscat.read_from_hipscat(rand_catalog_dir)

tests/corrgi/test_pcf.py

@@ -17,6 +17,23 @@ def test_pcf_counts_are_correct(
     npt.assert_allclose(expected_rr, pcf_rr_counts, rtol=2e-3)
 
 
+def test_pcf_counts_with_weights_are_correct(
+    dask_client,
+    pcf_gals_weight_catalog,
+    pcf_rans_weight_catalog,
+    pcf_dd_counts_with_weights,
+    pcf_rr_counts_with_weights,
+    pcf_params,
+):
+    proj_corr = ProjectedCorrelation(params=pcf_params, use_weights=True)
+    counts_dd, counts_rr = compute_autocorrelation_counts(
+        pcf_gals_weight_catalog, pcf_rans_weight_catalog, proj_corr
+    )
+    expected_dd, expected_rr = counts_dd.transpose([1, 0]), counts_rr.transpose([1, 0])
+    npt.assert_allclose(expected_dd, pcf_dd_counts_with_weights, rtol=1e-3)
+    npt.assert_allclose(expected_rr, pcf_rr_counts_with_weights, rtol=2e-3)
+
+
 def test_pcf_catalog_has_no_redshift(data_catalog, rand_catalog, pcf_params):
     with pytest.raises(ValueError, match="ph_z not found"):
         compute_autocorrelation(

tests/data/generate_data.ipynb

@@ -52,7 +52,8 @@
     "generate_catalog(\"DR7-lrg\")\n",
     "\n",
     "# With weights\n",
-    "generate_catalog(\"acf_gals_weight\")"
+    "generate_catalog(\"acf_gals_weight\")\n",
+    "generate_catalog(\"pcf_gals_weight\")"
    ]
   },
   {
@@ -73,7 +74,8 @@
     "generate_catalog(\"DR7-lrg-rand\")\n",
     "\n",
     "# With weights\n",
-    "generate_catalog(\"acf_rans_weight\")"
+    "generate_catalog(\"acf_rans_weight\")\n",
+    "generate_catalog(\"pcf_rans_weight\")"
    ]
   }
  ],

tests/data/hipscat/pcf_gals_weight/catalog_info.json

@@ -0,0 +1,8 @@
+{
+    "catalog_name": "pcf_gals_weight",
+    "catalog_type": "object",
+    "total_rows": 84383,
+    "epoch": "J2000",
+    "ra_column": "ra",
+    "dec_column": "dec"
+}

tests/data/hipscat/pcf_gals_weight/partition_info.csv

@@ -0,0 +1,8 @@
+Norder,Dir,Npix,num_rows
+0,0,1,23743
+0,0,2,34965
+0,0,5,632
+0,0,6,21350
+0,0,7,3481
+0,0,9,3
+0,0,10,209

tests/data/hipscat/pcf_gals_weight/provenance_info.json

@@ -0,0 +1,25 @@
+{
+    "catalog_name": "pcf_gals_weight",
+    "catalog_type": "object",
+    "total_rows": 84383,
+    "epoch": "J2000",
+    "ra_column": "ra",
+    "dec_column": "dec",
+    "version": "0.3.4",
+    "generation_date": "2024.07.08",
+    "tool_args": {
+        "tool_name": "lsdb",
+        "version": "0.2.5",
+        "runtime_args": {
+            "catalog_name": "pcf_gals_weight",
+            "output_path": "hipscat/pcf_gals_weight",
+            "output_catalog_name": "pcf_gals_weight",
+            "catalog_path": "hipscat/pcf_gals_weight",
+            "catalog_type": "object",
+            "total_rows": 84383,
+            "epoch": "J2000",
+            "ra_column": "ra",
+            "dec_column": "dec"
+        }
+    }
+}

tests/data/hipscat/pcf_rans_weight/catalog_info.json

@@ -0,0 +1,8 @@
+{
+    "catalog_name": "pcf_rans_weight",
+    "catalog_type": "object",
+    "total_rows": 227336,
+    "epoch": "J2000",
+    "ra_column": "ra",
+    "dec_column": "dec"
+}

tests/data/hipscat/pcf_rans_weight/partition_info.csv

@@ -0,0 +1,8 @@
+Norder,Dir,Npix,num_rows
+0,0,1,73222
+0,0,2,82167
+0,0,5,1742
+0,0,6,60914
+0,0,7,8953
+0,0,9,17
+0,0,10,321

tests/data/hipscat/pcf_rans_weight/provenance_info.json

@@ -0,0 +1,25 @@
+{
+    "catalog_name": "pcf_rans_weight",
+    "catalog_type": "object",
+    "total_rows": 227336,
+    "epoch": "J2000",
+    "ra_column": "ra",
+    "dec_column": "dec",
+    "version": "0.3.4",
+    "generation_date": "2024.07.08",
+    "tool_args": {
+        "tool_name": "lsdb",
+        "version": "0.2.5",
+        "runtime_args": {
+            "catalog_name": "pcf_rans_weight",
+            "output_path": "hipscat/pcf_rans_weight",
+            "output_catalog_name": "pcf_rans_weight",
+            "catalog_path": "hipscat/pcf_rans_weight",
+            "catalog_type": "object",
+            "total_rows": 227336,
+            "epoch": "J2000",
+            "ra_column": "ra",
+            "dec_column": "dec"
+        }
+    }
+}