From 144afac50e574af8b579e0f6415f3c75e66632c6 Mon Sep 17 00:00:00 2001
From: jaimerzp <jaimerz011235813@gmail.com>
Date: Wed, 27 Nov 2024 16:39:24 +0000
Subject: [PATCH] working all

---
 heracles/fields.py   | 23 ++++++++++-------------
 tests/test_fields.py | 34 ++++++++++++++++++++++------------
 2 files changed, 32 insertions(+), 25 deletions(-)

diff --git a/heracles/fields.py b/heracles/fields.py
index 872ef97..ebf3b68 100644
--- a/heracles/fields.py
+++ b/heracles/fields.py
@@ -342,7 +342,7 @@ async def __call__(
 
         # total weighted variance from online algorithm
         ngal = 0
-        wmean, var = 0.0, 0.0
+        wmean, w2mean, var = 0.0, 0.0, 0.0
 
         # go through pages in catalogue and map values
         async for page in aiter_pages(catalog, progress):
@@ -358,6 +358,7 @@ async def __call__(
 
                 ngal += page.size
                 wmean += (w - wmean).sum() / ngal
+                w2mean += (w**2 - w2mean).sum() / ngal
                 var += (v**2 - var).sum() / ngal
 
                 del lon, lat, v, w
@@ -375,11 +376,9 @@ async def __call__(
         val /= wbar
 
         # compute bias from variance (per object)
-        musq = var / wmean**2
-        dens = ngal / (4 * np.pi * fsky)
-        # variance = var / w2mean
-        # deff = w2mean / wmean**2
-        # neff = ngal / (4 * np.pi * fsky) / deff
+        musq = var / w2mean
+        deff = w2mean / wmean**2
+        dens = ngal / (4 * np.pi * fsky) / deff
         bias = fsky * musq / dens
 
         # set metadata of array
@@ -420,7 +419,7 @@ async def __call__(
 
         # total weighted variance from online algorithm
         ngal = 0
-        wmean, var = 0.0, 0.0
+        wmean, w2mean, var = 0.0, 0.0, 0.0
 
         # go through pages in catalogue and get the shear values,
         async for page in aiter_pages(catalog, progress):
@@ -436,6 +435,7 @@ async def __call__(
 
                 ngal += page.size
                 wmean += (w - wmean).sum() / ngal
+                w2mean += (w**2 - w2mean).sum() / ngal
                 var += (re**2 + im**2 - var).sum() / ngal
 
                 del lon, lat, re, im, w
@@ -452,13 +452,10 @@ async def __call__(
         val /= wbar
 
         # bias from measured variance, for E/B decomposition
-        musq = var / wmean**2
-        dens = ngal / (4 * np.pi * fsky)
-        # variance = var / w2mean
-        # deff = w2mean / wmean**2
-        # neff = ngal / (2 * np.pi * fsky) / deff
+        musq = var / w2mean
+        deff = w2mean / wmean**2
+        dens = ngal / (4 * np.pi * fsky) / deff
         bias = (1 / 2) * fsky * musq / dens
-
         # set metadata of array
         update_metadata(
             val, catalog, wbar=wbar, musq=musq, dens=dens, fsky=fsky, bias=bias
diff --git a/tests/test_fields.py b/tests/test_fields.py
index db5a4e6..9a01035 100644
--- a/tests/test_fields.py
+++ b/tests/test_fields.py
@@ -299,14 +299,19 @@ def test_scalar_field(mapper, catalog):
 
     w = next(iter(catalog))["w"]
     v = next(iter(catalog))["g1"]
-    v2 = ((w * v) ** 2).sum()
+    v1 = w.sum()
+    v2 = (w ** 2).sum()
+    v3 = ((w * v) ** 2).sum()
     w = w.reshape(w.size // 4, 4).sum(axis=-1)
     wbar = w.mean()
-    v1 = w.sum()
-    bias = (4 * np.pi / npix / npix) * v2
     wmean = v1 / (4.0 * npix)
     w2mean = v2 / (4.0 * npix)
-    musq = w2mean / wmean**2
+    var = v3 / (4.0 * npix)
+    musq = var / w2mean
+    deff = w2mean / wmean**2
+    dens = npix / np.pi / deff
+    bias = (4 * np.pi / npix / npix) * v3
+    bias = bias / wbar**2
 
     assert m.shape == (npix,)
     assert m.dtype.metadata == {
@@ -319,9 +324,9 @@ def test_scalar_field(mapper, catalog):
         "lmax": mapper.lmax,
         "deconv": mapper.deconvolve,
         "musq": pytest.approx(musq),
-        "dens": pytest.approx(npix / np.pi),
+        "dens": pytest.approx(dens),
         "fsky": 1.0,
-        "bias": pytest.approx(bias / wbar**2),
+        "bias": pytest.approx(bias),
     }
     np.testing.assert_array_almost_equal(m, 0)
 
@@ -337,14 +342,19 @@ def test_complex_field(mapper, catalog):
     w = next(iter(catalog))["w"]
     re = next(iter(catalog))["g1"]
     im = next(iter(catalog))["g2"]
-    v2 = ((w * re) ** 2 + (w * im) ** 2).sum()
+    v1 = w.sum()
+    v2 = (w ** 2).sum()
+    v3 = ((w * re) ** 2 + (w * im) ** 2).sum()
     w = w.reshape(w.size // 4, 4).sum(axis=-1)
     wbar = w.mean()
-    v1 = w.sum()
-    bias = (4 * np.pi / npix / npix) * v2 / 2
     wmean = v1 / (4.0 * npix)
     w2mean = v2 / (4.0 * npix)
-    musq = w2mean / wmean**2
+    var = v3 / (4.0 * npix)
+    musq = var / w2mean
+    deff = w2mean / wmean**2
+    dens = npix / np.pi / deff
+    bias = (4 * np.pi / npix / npix) * v3 / 2
+    bias = bias / wbar**2
 
     assert m.shape == (2, npix)
     assert m.dtype.metadata == {
@@ -357,9 +367,9 @@ def test_complex_field(mapper, catalog):
         "lmax": mapper.lmax,
         "deconv": mapper.deconvolve,
         "fsky": 1.0,
-        "dens": npix / np.pi,
+        "dens": pytest.approx(dens),
         "musq": pytest.approx(musq),
-        "bias": pytest.approx(bias / wbar**2),
+        "bias": pytest.approx(bias),
     }
     np.testing.assert_array_almost_equal(m, 0)