Infrastructure improvements for power spectrum fitting

fitstack/mcmc.py

@@ -563,6 +563,7 @@ class RunMCMC(task.SingleTask):
     """
 
     max_iter = config.Property(proptype=int, default=1)
+    data_2d = config.Property(proptype=str, default=None)
 
     data = config.Property(proptype=str)
     mocks = config.Property(proptype=_list_or_glob)

fitstack/signal.py

@@ -12,7 +12,8 @@
 from draco.analysis.powerspec import get_1d_ps
 
 from . import utils
-
+from cora.util import cosmology
+from cora.util import units as u
 logger = logging.getLogger(__name__)
 
 
@@ -581,7 +582,7 @@ def __init__(
         derivs: Optional[Dict[str, Tuple[float, float]]] = None,
         factor: float = 1.0,
         aliases: Optional[Dict[str, str]] = None,
-        nbins: int = 10,
+        nbins: int = 7,
         logbins: bool = True,
     ):
 
@@ -648,13 +649,25 @@ def load_from_ps2Dfiles(
 
         # Find directories which match the right format
         for d in sorted(dirs):
-            mo = re.search(r"_compderiv-([^\/]+)", d)
-
-            if mo is None:
+            ###mo = re.search(r"_compderiv-([^\/]+)", d)
+            mo_comp = re.search(r"template_bias_([0-9\.]+)_Pk_([^_]+)_ps_HI", d)
+            # Check for shotnoise directory
+            mo_shotnoise = re.search(r"template_shotnoise", d)
+
+            if mo_comp:
+                bias = mo_comp.group(1)  # This will be "0", "0.5", or "1"
+                pk_type = mo_comp.group(2)  # This will be "base", "lin", "FoGh"
+                # Create a composite key that identifies the templates
+                key = f"{bias}-{pk_type}"
+
+            elif mo_shotnoise:
+                key = "shotnoise"
+            else:
                 print(f"Directory {d} does not match expected format, rejecting")
                 continue
 
-            key = mo.group(1)
+            print(f"Processing directory: {d}")
+
 
             if key in matching:
                 raise ValueError(
@@ -711,7 +724,7 @@ def _interpret_ps2Ds(
         # Generate the required templates from the 2d power spectra
 
         # Find all entries that have the linear component structure
-        compterms = [k.split("-")[1] for k in ps2Ds.keys() if k.startswith("0")]
+        compterms = [k.split("-")[1] for k in ps2Ds.keys() if k.startswith("0-")]
 
         ps2D_modes = {}
 
@@ -751,7 +764,7 @@ def _check_load_ps2D(key):
             logger.debug(f"Combining mode {term}")
 
             s0, v0 = _check_load_ps2D(f"0-{term}")
-            sh, vh = _check_load_ps2D(f"h-{term}")
+            sh, vh = _check_load_ps2D(f"0.5-{term}")
             s1, v1 = _check_load_ps2D(f"1-{term}")
 
             # Initialize arrays for b_HI = 0, 1/2, 1
@@ -957,9 +970,19 @@ def __init__(
         derivs: Optional[Dict[str, Tuple[float, float]]] = None,
         convolutions: Optional[Dict[str, Tuple[float, float]]] = None,
         kpar_range: Optional[Tuple[float, float]] = None,
+        nu21: float = 1420.40575177,  # MHz
+        z_eff: Optional[float] = None,  # effective redshift
+        cs: float = 299792.458,  # speed of light in km/s          
         *args,
         **kwargs,
     ):
+
+        # Set default z_eff if None
+        self.z_eff = z_eff if z_eff is not None else 1.0
+        cosmo_cora = cosmology.Cosmology()
+        self.H_z = cosmo_cora.H(self.z_eff) * u.mega_parsec / 1000.  # In km/s/Mpc
+        self.nu21 = nu21 
+        self.cs = cs           
 
         if derivs is None:
             derivs = {
@@ -1086,6 +1109,22 @@ def _interpret_ps2Ds(self, ps2Ds: Dict[str, Powerspec1D]):
             scale = self._solve_scale(base, ps2Ds[key], alpha)
             self._convolution_scale[name] = scale
 
+
+    def _get_factor(self) -> float:
+        """Calculate the conversion factor connecting tau and k_parallel.
+
+        C = -1/(2 pi nu21) * c/H(z) * (1+z)²
+
+        Returns
+        -------
+        C : float
+            Conversion factor
+        """
+
+        C = (-1.0 / (2 * np.pi * self.nu21)) * (self.cs / self.H_z) * (1 + self.z_eff)**2
+
+        return C            
+
     def multiply_pre_noncomp(self, signal: np.ndarray, **kwargs) -> np.ndarray:
         """Multiply the 2d power spectrum with the relative FoG kernel.
 
@@ -1101,7 +1140,9 @@ def multiply_pre_noncomp(self, signal: np.ndarray, **kwargs) -> np.ndarray:
         signal : np.ndarray[npol, nkpar, nkperp]
             The 2d power spectrum after multiplication with the relative FoG kernel.
         """
-
+        # Calculate the conversion factor
+        C = self._get_factor()
+
         # Loop over parameters corresponding to distinct kernels we'll need to
         # multiply into the signal
         for name, (_, x0) in self._convolutions.items():
@@ -1120,8 +1161,8 @@ def multiply_pre_noncomp(self, signal: np.ndarray, **kwargs) -> np.ndarray:
             scale = alpha * scale0
 
             # Multiply kernel into signal
-            signal *= (1.0 + (scale0 * self.kpar[np.newaxis, :, np.newaxis]) ** 2) / (
-                1.0 + (scale * self.kpar[np.newaxis, :, np.newaxis]) ** 2
+            signal *= (1.0 + (scale0 * C * self.kpar[np.newaxis, :, np.newaxis]) ** 2) / (
+                1.0 + (scale * C * self.kpar[np.newaxis, :, np.newaxis]) ** 2
             )
 
         return signal

fitstack/utils.py

@@ -424,7 +424,16 @@ def average_data(cnt, pol=None, combine=True, sort=True):
             distributed=False,
         )
         avg.ps2D[:] = np.mean(darr, axis=0)
-        avg.ps2D_weight[:] = tools.invert_no_zero(np.var(darr, axis=0))
+        #avg.ps2D_weight[:] = tools.invert_no_zero(np.var(darr, axis=0))
+
+        if darr.shape[0] == 1:
+            # Set weights to unity for single mock
+            avg.ps2D_weight[:] = np.ones_like(avg.ps2D[:])
+        else:
+            # variance calculation for multiple mocks
+            avg.ps2D_weight[:] = tools.invert_no_zero(np.var(darr, axis=0))
+
+
     else:
         avg = containers.Powerspec1D(
             pol=np.array(dpol), k=cnt.index_map["k"], attrs_from=cnt, distributed=False