DM-13854: extend cache in Psf

include/lsst/afw/detection/Psf.h

@@ -29,6 +29,7 @@
 #include <memory>
 
 #include "lsst/daf/base.h"
+#include "lsst/utils/CacheFwd.h"
 #include "lsst/afw/geom/ellipses/Quadrupole.h"
 #include "lsst/afw/math/Kernel.h"
 #include "lsst/afw/image/Color.h"
@@ -37,6 +38,12 @@
 namespace lsst {
 namespace afw {
 namespace detection {
+namespace detail {
+
+/// Key for caching PSFs with lsst::utils::Cache
+struct PsfCacheKey;
+
+} // namespace detail
 
 class PsfFormatter;
 
@@ -86,12 +93,12 @@ class Psf : public daf::base::Citizen,
                       */
     };
 
-    Psf(Psf const&) = default;
-    Psf(Psf&&) = default;
+    Psf(Psf const&);
     Psf& operator=(Psf const&) = delete;
     Psf& operator=(Psf&&) = delete;
 
-    virtual ~Psf() = default;
+    Psf(Psf&&);
+    virtual ~Psf();
 
     /**
      *  Polymorphic deep-copy.
@@ -254,14 +261,27 @@ class Psf : public daf::base::Citizen,
                                                       std::string const& warpAlgorithm = "lanczos5",
                                                       unsigned int warpBuffer = 5);
 
+    /** Return the capacity of the caches
+     *
+     * Both the image and kernel image caches have the same capacity.
+     */
+    std::size_t getCacheCapacity() const;
+
+    /** Set the capacity of the caches
+     *
+     * Both the image and kernel image caches will be set to this capacity.
+     */
+    void setCacheCapacity(std::size_t capacity);
+
 protected:
     /**
      *  Main constructor for subclasses.
      *
      *  @param[in] isFixed  Should be true for Psf for which doComputeKernelImage always returns
      *                      the same image, regardless of color or position arguments.
+     *  @param[in] capacity  Capacity of the caches.
      */
-    explicit Psf(bool isFixed = false);
+    explicit Psf(bool isFixed = false, std::size_t capacity=100);
 
 private:
     //@{
@@ -284,12 +304,9 @@ class Psf : public daf::base::Citizen,
     //@}
 
     bool const _isFixed;
-    mutable std::shared_ptr<Image> _cachedImage;
-    mutable std::shared_ptr<Image> _cachedKernelImage;
-    mutable image::Color _cachedImageColor;
-    mutable image::Color _cachedKernelImageColor;
-    mutable geom::Point2D _cachedImagePosition;
-    mutable geom::Point2D _cachedKernelImagePosition;
+    using PsfCache = utils::Cache<detail::PsfCacheKey, std::shared_ptr<Image>>;
+    std::unique_ptr<PsfCache> _imageCache;
+    std::unique_ptr<PsfCache> _kernelImageCache;
 
     LSST_PERSIST_FORMATTER(PsfFormatter)
 };

include/lsst/afw/geom/Point.h

@@ -293,6 +293,10 @@ class Point<T, 3> : public PointBase<T, 3> {
     void swap(Point &other) { this->_swap(other); }
 };
 
+// Hash functions
+template <typename T, int N>
+std::size_t hash_value(Point<T, N> const& point);
+
 typedef Point<int, 2> PointI;
 typedef Point<int, 2> Point2I;
 typedef Point<int, 3> Point3I;

python/lsst/afw/detection/psf.cc

@@ -74,6 +74,8 @@ PYBIND11_PLUGIN(_psf) {
     cls.def("getAveragePosition", &Psf::getAveragePosition);
     cls.def_static("recenterKernelImage", &Psf::recenterKernelImage, "im"_a, "position"_a,
                    "warpAlgorithm"_a = "lanczos5", "warpBuffer"_a = 5);
+    cls.def("getCacheCapacity", &Psf::getCacheCapacity);
+    cls.def("setCacheCapacity", &Psf::setCacheCapacity);
 
     return mod.ptr();
 }

src/detection/Psf.cc

@@ -2,29 +2,82 @@
 #include <limits>
 #include <typeinfo>
 #include <cmath>
+#include <memory>
 
+#include "lsst/utils/Cache.h"
+#include "lsst/utils/hashCombine.h"
 #include "lsst/afw/detection/Psf.h"
 #include "lsst/afw/math/offsetImage.h"
 
+namespace lsst::afw::detection::detail {
+
+// Key for caching PSFs with lsst::utils::Cache
+//
+// We cache PSFs by their x,y position. Although there are placeholders
+// in the `Psf` class and here for `image::Color`, these are not used
+// in the cache because `image::Color` is not currently well-defined
+// or used.
+struct PsfCacheKey {
+    geom::Point2D const position;
+    image::Color const color;
+
+    PsfCacheKey(geom::Point2D const& position_, image::Color color_=image::Color())
+      : position(position_), color(color_) {}
+
+    bool operator==(PsfCacheKey const& other) const {
+        return position == other.position;  // Currently don't care about color
+    }
+
+    friend std::ostream & operator<<(std::ostream &os, PsfCacheKey const &key) {
+        return os << key.position;
+    }
+};
+
+} // namespace lsst::afw::detection::detail
+
+namespace std {
+
+// Template specialisation for hashing PsfCacheKey
+//
+// We currently ignore the color.
+template <>
+struct hash<lsst::afw::detection::detail::PsfCacheKey> {
+    std::size_t operator()(lsst::afw::detection::detail::PsfCacheKey const& key) const {
+        std::size_t seed = 0;
+        return lsst::utils::hashCombine(seed, key.position.getX(), key.position.getY());
+    }
+};
+
+} // namespace std
+
 namespace lsst {
 namespace afw {
 namespace detection {
 
 namespace {
 
-// Comparison function that determines when we used the cached image instead of recomputing it.
-// We'll probably want a tolerance for colors someday too, but they're just a placeholder right now
-// so it's not worth the effort.
-bool comparePsfEvalPoints(geom::Point2D const &a, geom::Point2D const &b) {
-    // n.b. desired tolerance is actually sqrt(eps), so tolerance squared is eps.
-    return (a - b).computeSquaredNorm() < std::numeric_limits<double>::epsilon();
-}
-
 bool isPointNull(geom::Point2D const &p) { return std::isnan(p.getX()) && std::isnan(p.getY()); }
 
 }  // anonymous
 
-Psf::Psf(bool isFixed) : daf::base::Citizen(typeid(this)), _isFixed(isFixed) {}
+Psf::Psf(bool isFixed, std::size_t capacity)
+  : daf::base::Citizen(typeid(this)),
+  _isFixed(isFixed)
+{
+  _imageCache = std::make_unique<PsfCache>(capacity);
+  _kernelImageCache = std::make_unique<PsfCache>(capacity);
+}
+
+Psf::~Psf() = default;
+
+Psf::Psf(Psf const& other) : Psf(other._isFixed, other.getCacheCapacity()) {}
+
+Psf::Psf(Psf && other)
+  : daf::base::Citizen(std::move(other)),
+    _isFixed(other._isFixed),
+    _imageCache(std::move(other._imageCache)),
+    _kernelImageCache(std::move(other._kernelImageCache))
+{}
 
 std::shared_ptr<image::Image<double>> Psf::recenterKernelImage(std::shared_ptr<Image> im,
                                                                geom::Point2D const &position,
@@ -46,15 +99,10 @@ std::shared_ptr<Psf::Image> Psf::computeImage(geom::Point2D position, image::Col
                                               ImageOwnerEnum owner) const {
     if (isPointNull(position)) position = getAveragePosition();
     if (color.isIndeterminate()) color = getAverageColor();
-    std::shared_ptr<Psf::Image> result;
-    if (_cachedImage && color == _cachedImageColor && comparePsfEvalPoints(position, _cachedImagePosition)) {
-        result = _cachedImage;
-    } else {
-        result = doComputeImage(position, color);
-        _cachedImage = result;
-        _cachedImageColor = color;
-        _cachedImagePosition = position;
-    }
+    std::shared_ptr<Psf::Image> result = (*_imageCache)(
+        detail::PsfCacheKey(position, color),
+        [this](detail::PsfCacheKey const& key) { return doComputeImage(key.position, key.color); }
+    );
     if (owner == COPY) {
         result = std::make_shared<Image>(*result, true);
     }
@@ -63,18 +111,12 @@ std::shared_ptr<Psf::Image> Psf::computeImage(geom::Point2D position, image::Col
 
 std::shared_ptr<Psf::Image> Psf::computeKernelImage(geom::Point2D position, image::Color color,
                                                     ImageOwnerEnum owner) const {
-    if (isPointNull(position)) position = getAveragePosition();
-    if (color.isIndeterminate()) color = getAverageColor();
-    std::shared_ptr<Psf::Image> result;
-    if (_cachedKernelImage && (_isFixed || (color == _cachedKernelImageColor &&
-                                            comparePsfEvalPoints(position, _cachedKernelImagePosition)))) {
-        result = _cachedKernelImage;
-    } else {
-        result = doComputeKernelImage(position, color);
-        _cachedKernelImage = result;
-        _cachedKernelImageColor = color;
-        _cachedKernelImagePosition = position;
-    }
+    if (_isFixed || isPointNull(position)) position = getAveragePosition();
+    if (_isFixed || color.isIndeterminate()) color = getAverageColor();
+    std::shared_ptr<Psf::Image> result = (*_kernelImageCache)(
+        detail::PsfCacheKey(position, color),
+        [this](detail::PsfCacheKey const& key) { return doComputeKernelImage(key.position, key.color); }
+    );
     if (owner == COPY) {
         result = std::make_shared<Image>(*result, true);
     }
@@ -121,6 +163,14 @@ std::shared_ptr<Psf::Image> Psf::doComputeImage(geom::Point2D const &position,
 }
 
 geom::Point2D Psf::getAveragePosition() const { return geom::Point2D(); }
+
+std::size_t Psf::getCacheCapacity() const { return _kernelImageCache->capacity(); }
+
+void Psf::setCacheCapacity(std::size_t capacity) {
+    _imageCache->reserve(capacity);
+    _kernelImageCache->reserve(capacity);
+}
+
 }
 }
 }  // namespace lsst::afw::detection

src/geom/Point.cc

@@ -22,6 +22,7 @@
 
 #include <cmath>
 
+#include "lsst/utils/hashCombine.h"
 #include "lsst/afw/geom/Point.h"
 
 namespace lsst {
@@ -115,19 +116,27 @@ CoordinateExpr<N> PointBase<T, N>::ge(Point<T, N> const &other) const {
     return r;
 }
 
+template <typename T, int N>
+std::size_t hash_value(Point<T, N> const& point) {
+    std::size_t result = 0;
+    for (int n = 0; n < N; ++n) result = utils::hashCombine(result, point[n]);
+    return result;
+}
+
+
 #ifndef DOXYGEN
-template class PointBase<int, 2>;
-template class PointBase<int, 3>;
-template class PointBase<double, 2>;
-template class PointBase<double, 3>;
-template class Point<int, 2>;
-template class Point<int, 3>;
-template class Point<double, 2>;
-template class Point<double, 3>;
-template Point<int, 2>::Point(Point<double, 2> const &);
-template Point<int, 3>::Point(Point<double, 3> const &);
-template Point<double, 2>::Point(Point<int, 2> const &);
-template Point<double, 3>::Point(Point<int, 3> const &);
+#define INSTANTIATE_TYPE_DIM(TYPE, DIM) \
+    template class PointBase<TYPE, DIM>; \
+    template class Point<TYPE, DIM>; \
+    template std::size_t hash_value(Point<TYPE, DIM> const&);
+#define INSTANTIATE_DIM(DIM) \
+    INSTANTIATE_TYPE_DIM(int, DIM); \
+    INSTANTIATE_TYPE_DIM(double, DIM); \
+    template Point<int, DIM>::Point(Point<double, DIM> const &); \
+    template Point<double, DIM>::Point(Point<int, DIM> const &);
+
+INSTANTIATE_DIM(2);
+INSTANTIATE_DIM(3);
 #endif
 }
 }