Add property to DimensionGroup to make it easy to get best time/regio…

…n dim for query
lsst · Sep 6, 2024 · b186361 · b186361
1 parent 767a556
commit b186361
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 0 deletions.
diff --git a/python/lsst/daf/butler/dimensions/_group.py b/python/lsst/daf/butler/dimensions/_group.py
@@ -449,6 +449,31 @@ def add_to_order(element: DimensionElement) -> None:
         order.extend(element for element in self.elements if element not in done)
         return tuple(order)
 
+    def _choose_dimension(self, families: NamedValueAbstractSet[TopologicalFamily]) -> str | None:
+        if len(families) != 1:
+            return None
+        return list(families)[0].choose(self.elements, self.universe).name
+
+    @property
+    def region_dimension(self) -> str | None:
+        """Return the most appropriate spatial dimension to use when looking
+        up a region.
+
+        Returns `None` if there are no appropriate dimensions or more than one
+        spatial family.
+        """
+        return self._choose_dimension(self.spatial)
+
+    @property
+    def timespan_dimension(self) -> str | None:
+        """Return the most appropriate temporal dimension to use when looking
+        up a time span.
+
+        Returns `None` if there are no appropriate dimensions or more than one
+        temporal family.
+        """
+        return self._choose_dimension(self.temporal)
+
     @property
     def spatial(self) -> NamedValueAbstractSet[TopologicalFamily]:
         """Families represented by the spatial elements in this graph."""

diff --git a/tests/test_dimensions.py b/tests/test_dimensions.py
@@ -285,6 +285,8 @@ def testCalibrationDimensions(self):
         self.assertCountEqual(group.implied, ("band",))
         self.assertCountEqual(group.elements, group.names)
         self.assertCountEqual(group.governors, {"instrument"})
+        self.assertIsNone(group.region_dimension)
+        self.assertIsNone(group.timespan_dimension)
 
     def testObservationDimensions(self):
         group = self.universe.conform(["exposure", "detector", "visit"])
@@ -298,6 +300,8 @@ def testObservationDimensions(self):
         self.assertCountEqual(group.spatial.names, ("observation_regions",))
         self.assertCountEqual(group.temporal.names, ("observation_timespans",))
         self.assertCountEqual(group.governors, {"instrument"})
+        self.assertEqual(group.region_dimension, "visit_detector_region")
+        self.assertEqual(group.timespan_dimension, "exposure")
         self.assertEqual(group.spatial.names, {"observation_regions"})
         self.assertEqual(group.temporal.names, {"observation_timespans"})
         self.assertEqual(next(iter(group.spatial)).governor, self.universe["instrument"])
@@ -334,6 +338,9 @@ def testSubsetCalculation(self):
         group = self.universe.conform(["visit", "detector", "tract", "patch", "htm7", "exposure"])
         self.assertCountEqual(group.spatial.names, ("observation_regions", "skymap_regions", "htm"))
         self.assertCountEqual(group.temporal.names, ("observation_timespans",))
+        self.assertEqual(group.timespan_dimension, "exposure")
+        # Can not choose between visit_detector_region or htm7 or tract/patch.
+        self.assertIsNone(group.region_dimension)
 
     def testSchemaGeneration(self):
         tableSpecs: NamedKeyDict[DimensionElement, ddl.TableSpec] = NamedKeyDict({})