From a21b9aa91809a9ec269a695334703c08b245664f Mon Sep 17 00:00:00 2001 From: Sandro Campos Date: Wed, 17 Jul 2024 09:06:19 -0400 Subject: [PATCH] Get correct catalog length (#26) --- src/corrgi/corrgi.py | 5 +++-- src/corrgi/utils.py | 13 +++++++++++++ 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/src/corrgi/corrgi.py b/src/corrgi/corrgi.py index 3e272ac..1c86a39 100644 --- a/src/corrgi/corrgi.py +++ b/src/corrgi/corrgi.py @@ -5,6 +5,7 @@ from corrgi.correlation.correlation import Correlation from corrgi.dask import compute_autocorrelation_counts from corrgi.estimators import calculate_natural_estimate +from corrgi.utils import compute_catalog_size def compute_autocorrelation( @@ -24,9 +25,9 @@ def compute_autocorrelation( """ correlation = corr_type(**kwargs) correlation.validate([catalog, random]) + num_galaxies = compute_catalog_size(catalog) + num_random = compute_catalog_size(random) counts_dd, counts_rr = compute_autocorrelation_counts(catalog, random, correlation) - num_galaxies = catalog.hc_structure.catalog_info.total_rows - num_random = random.hc_structure.catalog_info.total_rows return calculate_natural_estimate(counts_dd, counts_rr, num_galaxies, num_random) diff --git a/src/corrgi/utils.py b/src/corrgi/utils.py index 93e2d03..89428e0 100644 --- a/src/corrgi/utils.py +++ b/src/corrgi/utils.py @@ -1,5 +1,6 @@ import numpy as np from gundam import gundam +from lsdb import Catalog from numpy import deg2rad @@ -42,3 +43,15 @@ def join_count_histograms(partial_histograms: list[np.ndarray]) -> np.ndarray: The numpy array with the total counts for the partial histograms. """ return np.sum(np.stack(partial_histograms), axis=0) + + +def compute_catalog_size(catalog: Catalog) -> int: + """Compute the number of rows in a catalog. + + Args: + catalog (Catalog): An LSDB catalog. + + Returns: + The number of rows in the catalog. + """ + return catalog._ddf.shape[0].compute()