diff --git a/difi/metrics.py b/difi/metrics.py index 85c7fac..6649b1f 100644 --- a/difi/metrics.py +++ b/difi/metrics.py @@ -1028,8 +1028,9 @@ def run( by_object: bool = False, ignore_after_discovery: bool = False, num_jobs: Optional[int] = 1, + return_summary: bool = True, clear_on_failure: bool = True, - ) -> Tuple[pd.DataFrame, pd.DataFrame]: + ) -> Tuple[pd.DataFrame, Optional[pd.DataFrame]]: """ Run the findability metric on the observations. @@ -1070,6 +1071,9 @@ def run( num_jobs : int, optional The number of jobs to run in parallel. If 1, then run in serial. If None, then use the number of CPUs on the machine. + return_summary : bool, optional + If True, then return a summary of the number of observations, number of findable + objects and the start and end night of each window. clear_on_failure : bool, optional If a failure occurs and this is False, then the shared memory array will not be cleared. If True, then the shared memory array will be cleared. @@ -1116,7 +1120,10 @@ def run( clear_on_failure=clear_on_failure, ) - window_summary = self._create_window_summary(observations, windows, findable) + if return_summary: + window_summary = self._create_window_summary(observations, windows, findable) + else: + window_summary = None return findable, window_summary diff --git a/difi/tests/test_metrics.py b/difi/tests/test_metrics.py index bcecfc1..c9d97fc 100644 --- a/difi/tests/test_metrics.py +++ b/difi/tests/test_metrics.py @@ -1,6 +1,7 @@ import os import numpy as np +import pandas as pd import pytest from ..metrics import ( @@ -522,3 +523,14 @@ def test_FindabilityMetrics_shared_memory(test_observations): assert metric._shared_memory_name is None assert metric._num_observations == 0 assert metric._dtypes is None + + +def test_FindabilityMetrics_run_return_summary(test_observations): + # Check that the function returns the summary dataframe when desired + metric = MinObsMetric() + + findable, window_summary = metric.run(test_observations, return_summary=True) + assert isinstance(window_summary, pd.DataFrame) + + findable, window_summary = metric.run(test_observations, return_summary=False) + assert window_summary is None