Catalog.h

Go to the documentation of this file.

// -*- lsst-c++ -*-
#ifndef AFW_TABLE_Catalog_h_INCLUDED
#define AFW_TABLE_Catalog_h_INCLUDED

#include <type_traits>
#include <vector>

#include "boost/iterator/iterator_adaptor.hpp"
#include "boost/iterator/transform_iterator.hpp"

#include "lsst/base.h"
#include "lsst/pex/exceptions.h"
#include "lsst/afw/fitsDefaults.h"
#include "lsst/afw/table/fwd.h"
#include "lsst/afw/table/io/FitsWriter.h"
#include "lsst/afw/table/io/FitsReader.h"
#include "lsst/afw/table/SchemaMapper.h"

namespace lsst {
namespace afw {
namespace table {

template <typename BaseT>
class CatalogIterator : public boost::iterator_adaptor<CatalogIterator<BaseT>, BaseT,
                                                       typename BaseT::value_type::element_type> {
public:
    CatalogIterator() {}

    template <typename OtherBaseT>
    CatalogIterator(CatalogIterator<OtherBaseT> const& other)
            : CatalogIterator::iterator_adaptor_(other.base()) {}

    explicit CatalogIterator(BaseT const& base) : CatalogIterator::iterator_adaptor_(base) {}

    template <typename RecordT>
    operator std::shared_ptr<RecordT>() const {
        return *this->base();
    }

    template <typename RecordT>
    CatalogIterator& operator=(std::shared_ptr<RecordT> const& other) const {
        *this->base() = other;
        return *this;
    }

private:
    friend class boost::iterator_core_access;
    typename BaseT::value_type::element_type& dereference() const { return **this->base(); }
};

template <typename RecordT>
class CatalogT {
    typedef std::vector<std::shared_ptr<RecordT>> Internal;

public:
    typedef RecordT Record;
    typedef typename Record::Table Table;
    typedef typename Record::ColumnView ColumnView;

    typedef RecordT value_type;
    typedef RecordT& reference;
    typedef std::shared_ptr<RecordT> pointer;
    typedef typename Internal::size_type size_type;
    typedef typename Internal::difference_type difference_type;
    typedef CatalogIterator<typename Internal::iterator> iterator;
    typedef CatalogIterator<typename Internal::const_iterator> const_iterator;

    std::shared_ptr<Table> getTable() const { return _table; }

    Schema getSchema() const { return _table->getSchema(); }

    explicit CatalogT(std::shared_ptr<Table> const& table = std::shared_ptr<Table>())
            : _table(table), _internal() {}

    explicit CatalogT(Schema const& schema) : _table(Table::make(schema)), _internal() {}

    template <typename InputIterator>
    CatalogT(std::shared_ptr<Table> const& table, InputIterator first, InputIterator last, bool deep = false)
            : _table(table), _internal() {
        insert(end(), first, last, deep);
    }

    CatalogT(CatalogT const& other) : _table(other._table), _internal(other._internal) {}
    // Delegate to copy constructor for backward compatibility
    CatalogT(CatalogT&& other) : CatalogT(other) {}

    ~CatalogT() = default;

    template <typename OtherRecordT>
    CatalogT(CatalogT<OtherRecordT> const& other)
            : _table(other.getTable()), _internal(other.begin().base(), other.end().base()) {}

    CatalogT& operator=(CatalogT const& other) {
        if (&other != this) {
            _table = other._table;
            _internal = other._internal;
        }
        return *this;
    }
    // Delegate to copy assignment for backward compatibility
    CatalogT& operator=(CatalogT&& other) { return *this = other; }

    CatalogT<RecordT> subset(ndarray::Array<bool const, 1> const& mask) const {
        if (size_type(mask.size()) != size()) {
            throw LSST_EXCEPT(
                    pex::exceptions::LengthError,
                    (boost::format("Mask array with %d elements applied to catalog with %d elements") %
                     mask.size() % size())
                            .str());
        }
        CatalogT<RecordT> result(getTable());
        ndarray::Array<bool const, 1>::Iterator maskIter = mask.begin();
        const_iterator catIter = begin();
        for (; maskIter != mask.end(); ++maskIter, ++catIter) {
            if (*maskIter) result.push_back(catIter);
        }
        return result;
    }

    CatalogT<RecordT> subset(std::ptrdiff_t startd, std::ptrdiff_t stopd, std::ptrdiff_t step) const {
        /* Python's slicing syntax is weird and wonderful.

         Both the "start" and "stop" indices can be negative, which means the
         abs() of the index less than the size; [-1] means the last item.
         Moreover, it's possible to have a negative index less than -len(); it
         will get clipped.  That is in fact one way to slice *backward* through
         the array *and* include element 0;

         >>> range(10)[5:-20:-1]
         [5, 4, 3, 2, 1, 0]

         The clipping tests in this function look more complicated than they
         need to be, but that is partly because there are some weird edge cases.

         Also, ptrdiff_t vs size_t introduces some annoying complexity.  Note
         that the args are "startd"/"stopd" (not "start"/"stop").

         There is a fairly complete set of tests in tests/ticket2026.py; if you
         try to simplify this function, be sure they continue to pass.
         */
        size_type S = size();
        size_type start, stop = 0;
        // Python doesn't allow step == 0
        if (step == 0) {
            throw LSST_EXCEPT(pex::exceptions::InvalidParameterError, "Step cannot be zero");
        }
        // Basic negative indexing rule: first add size
        if (startd < 0) {
            startd += S;
        }
        if (stopd < 0) {
            stopd += S;
        }
        // Start gets clipped to zero; stop does not (yet).
        if (startd < 0) {
            startd = 0;
        }
        // Now start is non-negative, so can cast to size_t.
        start = (size_type)startd;
        if (start > S) {
            start = S;
        }
        if (step > 0) {
            // When stepping forward, stop gets clipped at zero,
            // so is non-negative and can get cast to size_t.
            if (stopd < 0) {
                stopd = 0;
            }
            stop = (size_type)stopd;
            if (stop > S) {
                stop = S;
            }
        } else if (step < 0) {
            // When stepping backward, stop gets clipped at -1 so that slices
            // including 0 are possible.
            if (stopd < 0) {
                stopd = -1;
            }
        }

        if (((step > 0) && (start >= stop)) || ((step < 0) && ((std::ptrdiff_t)start <= stopd))) {
            // Empty range
            return CatalogT<RecordT>(getTable(), begin(), begin());
        }

        if (step == 1) {
            // Use the iterator-based constructor for this simple case
            assert(start >= 0);
            assert(stop > 0);
            assert(start < S);
            assert(stop <= S);
            return CatalogT<RecordT>(getTable(), begin() + start, begin() + stop);
        }

        // Build a new CatalogT and copy records into it.
        CatalogT<RecordT> cat(getTable());
        size_type N = 0;
        if (step >= 0) {
            N = (stop - start) / step + (((stop - start) % step) ? 1 : 0);
        } else {
            N = (size_t)((stopd - (std::ptrdiff_t)start) / step +
                         (((stopd - (std::ptrdiff_t)start) % step) ? 1 : 0));
        }
        cat.reserve(N);
        if (step >= 0) {
            for (size_type i = start; i < stop; i += step) {
                cat.push_back(get(i));
            }
        } else {
            for (std::ptrdiff_t i = (std::ptrdiff_t)start; i > stopd; i += step) {
                cat.push_back(get(i));
            }
        }
        return cat;
    }

    void writeFits(std::string const& filename, std::string const& mode = "w", int flags = 0) const {
        io::FitsWriter::apply(filename, mode, *this, flags);
    }

    void writeFits(fits::MemFileManager& manager, std::string const& mode = "w", int flags = 0) const {
        io::FitsWriter::apply(manager, mode, *this, flags);
    }

    void writeFits(fits::Fits& fitsfile, int flags = 0) const {
        io::FitsWriter::apply(fitsfile, *this, flags);
    }

    static CatalogT readFits(std::string const& filename, int hdu = fits::DEFAULT_HDU, int flags = 0) {
        return io::FitsReader::apply<CatalogT>(filename, hdu, flags);
    }

    static CatalogT readFits(fits::MemFileManager& manager, int hdu = fits::DEFAULT_HDU, int flags = 0) {
        return io::FitsReader::apply<CatalogT>(manager, hdu, flags);
    }

    static CatalogT readFits(fits::Fits& fitsfile, int flags = 0) {
        return io::FitsReader::apply<CatalogT>(fitsfile, flags);
    }

    ColumnView getColumnView() const {
        if (std::is_const<RecordT>::value) {
            throw LSST_EXCEPT(
                    pex::exceptions::LogicError,
                    "Cannot get a column view from a CatalogT<RecordT const> (as column views are always "
                    "non-const views).");
        }
        return ColumnView::make(_table, begin(), end());
    }

    bool isContiguous() const { return ColumnView::isRangeContiguous(_table, begin(), end()); }


    iterator begin() { return iterator(_internal.begin()); }
    iterator end() { return iterator(_internal.end()); }
    const_iterator begin() const { return const_iterator(_internal.begin()); }
    const_iterator end() const { return const_iterator(_internal.end()); }
    const_iterator cbegin() const { return begin(); }
    const_iterator cend() const { return end(); }

    bool empty() const { return _internal.empty(); }

    size_type size() const { return _internal.size(); }

    size_type max_size() const { return _internal.max_size(); }

    size_type capacity() const { return _internal.size() + _table->getBufferSize(); }

    void reserve(size_type n) {
        if (n <= _internal.size()) return;
        _table->preallocate(n - _internal.size());
    }

    reference operator[](size_type i) const { return *_internal[i]; }

    reference at(size_type i) const { return *_internal.at(i); }

    reference front() const { return *_internal.front(); }

    reference back() const { return *_internal.back(); }

    std::shared_ptr<RecordT> const get(size_type i) const { return _internal[i]; }

    void set(size_type i, std::shared_ptr<RecordT> const& p) { _internal[i] = p; }

    template <typename InputIterator>
    void assign(InputIterator first, InputIterator last, bool deep = false) {
        clear();
        insert(end(), first, last, deep);
    }

    void push_back(Record const& r) {
        std::shared_ptr<RecordT> p = _table->copyRecord(r);
        _internal.push_back(p);
    }

    void push_back(std::shared_ptr<RecordT> const& p) { _internal.push_back(p); }

    std::shared_ptr<RecordT> addNew() {
        std::shared_ptr<RecordT> r = _table->makeRecord();
        _internal.push_back(r);
        return r;
    }

    void pop_back() { _internal.pop_back(); }

    CatalogT copy() const { return CatalogT(getTable()->clone(), begin(), end(), true); }

    template <typename InputIterator>
    void insert(iterator pos, InputIterator first, InputIterator last, bool deep = false) {
        _maybeReserve(pos, first, last, deep,
                      (typename std::iterator_traits<InputIterator>::iterator_category*)0);
        if (deep) {
            while (first != last) {
                pos = insert(pos, *first);
                ++pos;
                ++first;
            }
        } else {
            while (first != last) {
                pos = insert(pos, first);
                assert(pos != end());
                ++pos;
                ++first;
            }
        }
    }

    template <typename InputIterator>
    void insert(SchemaMapper const& mapper, iterator pos, InputIterator first, InputIterator last) {
        if (!_table->getSchema().contains(mapper.getOutputSchema())) {
            throw LSST_EXCEPT(pex::exceptions::InvalidParameterError,
                              "SchemaMapper's output schema does not match catalog's schema");
        }
        _maybeReserve(pos, first, last, true,
                      (typename std::iterator_traits<InputIterator>::iterator_category*)0);
        while (first != last) {
            pos = insert(pos, _table->copyRecord(*first, mapper));
            ++pos;
            ++first;
        }
    }

    iterator insert(iterator pos, Record const& r) {
        std::shared_ptr<RecordT> p = _table->copyRecord(r);
        return iterator(_internal.insert(pos.base(), p));
    }

    iterator insert(iterator pos, std::shared_ptr<RecordT> const& p) {
        return iterator(_internal.insert(pos.base(), p));
    }

    iterator erase(iterator pos) { return iterator(_internal.erase(pos.base())); }

    iterator erase(iterator first, iterator last) {
        return iterator(_internal.erase(first.base(), last.base()));
    }

    void swap(CatalogT& other) noexcept {
        _table.swap(other._table);
        _internal.swap(other._internal);
    }

    void clear() { _internal.clear(); }

    template <typename T>
    bool isSorted(Key<T> const& key) const;

    template <typename Compare>
    bool isSorted(Compare cmp) const;

    template <typename T>
    void sort(Key<T> const& key);

    template <typename Compare>
    void sort(Compare cmp);


    template <typename T>
    iterator find(typename Field<T>::Value const& value, Key<T> const& key);

    template <typename T>
    const_iterator find(typename Field<T>::Value const& value, Key<T> const& key) const;


    template <typename T>
    iterator lower_bound(typename Field<T>::Value const& value, Key<T> const& key);

    template <typename T>
    const_iterator lower_bound(typename Field<T>::Value const& value, Key<T> const& key) const;

    template <typename T>
    iterator upper_bound(typename Field<T>::Value const& value, Key<T> const& key);

    template <typename T>
    const_iterator upper_bound(typename Field<T>::Value const& value, Key<T> const& key) const;

    template <typename T>
    std::pair<iterator, iterator> equal_range(typename Field<T>::Value const& value, Key<T> const& key);

    template <typename T>
    std::pair<const_iterator, const_iterator> equal_range(typename Field<T>::Value const& value,
                                                          Key<T> const& key) const;


    Internal& getInternal() { return _internal; }
    Internal const& getInternal() const { return _internal; }

private:
    template <typename InputIterator>
    void _maybeReserve(iterator& pos, InputIterator first, InputIterator last, bool deep,
                       std::random_access_iterator_tag*) {
        if (deep) _table->preallocate(last - first);
    }

    template <typename InputIterator>
    void _maybeReserve(iterator pos, InputIterator first, InputIterator last, bool deep,
                       std::input_iterator_tag*) {}

    std::shared_ptr<Table> _table;
    Internal _internal;
};

namespace detail {

template <typename RecordT, typename T>
struct KeyComparisonFunctor {
    bool operator()(RecordT const& a, RecordT const& b) const { return a.get(key) < b.get(key); }

    Key<T> key;
};

template <typename RecordT, typename Adaptee>
struct ComparisonAdaptor {
    bool operator()(std::shared_ptr<RecordT> const& a, std::shared_ptr<RecordT> const& b) const {
        return adaptee(*a, *b);
    }

    Adaptee adaptee;
};

template <typename RecordT, typename T>
struct KeyExtractionFunctor {
    typedef typename Field<T>::Value result_type;

    result_type operator()(RecordT const& r) const { return r.get(key); }

    Key<T> key;
};

}  // namespace detail

template <typename RecordT>
template <typename Compare>
bool CatalogT<RecordT>::isSorted(Compare cmp) const {
    detail::ComparisonAdaptor<RecordT, Compare> f = {cmp};
    if (empty()) return true;
    const_iterator last = this->begin();
    const_iterator i = last;
    ++i;
    for (; i != this->end(); ++i) {
        if (f(i, last)) return false;
        last = i;
    }
    return true;
}

template <typename RecordT>
template <typename Compare>
void CatalogT<RecordT>::sort(Compare cmp) {
    detail::ComparisonAdaptor<RecordT, Compare> f = {cmp};
    std::stable_sort(_internal.begin(), _internal.end(), f);
}

template <typename RecordT>
template <typename T>
bool CatalogT<RecordT>::isSorted(Key<T> const& key) const {
    detail::KeyComparisonFunctor<RecordT, T> f = {key};
    return isSorted(f);
}

template <typename RecordT>
template <typename T>
void CatalogT<RecordT>::sort(Key<T> const& key) {
    detail::KeyComparisonFunctor<RecordT, T> f = {key};
    return sort(f);
}

template <typename RecordT>
template <typename T>
typename CatalogT<RecordT>::iterator CatalogT<RecordT>::find(typename Field<T>::Value const& value,
                                                             Key<T> const& key) {
    detail::KeyExtractionFunctor<RecordT, T> f = {key};
    // Iterator adaptor that makes a CatalogT iterator work like an iterator over field values.
    typedef boost::transform_iterator<detail::KeyExtractionFunctor<RecordT, T>, iterator> SearchIter;
    // Use binary search for log n search; requires sorted table.
    SearchIter i = std::lower_bound(SearchIter(begin(), f), SearchIter(end(), f), value);
    if (i.base() == end() || *i != value) return end();
    return i.base();
}

template <typename RecordT>
template <typename T>
typename CatalogT<RecordT>::const_iterator CatalogT<RecordT>::find(typename Field<T>::Value const& value,
                                                                   Key<T> const& key) const {
    detail::KeyExtractionFunctor<RecordT, T> f = {key};
    // Iterator adaptor that makes a CatalogT iterator work like an iterator over field values.
    typedef boost::transform_iterator<detail::KeyExtractionFunctor<RecordT, T>, const_iterator> SearchIter;
    // Use binary search for log n search; requires sorted table.
    SearchIter i = std::lower_bound(SearchIter(begin(), f), SearchIter(end(), f), value);
    if (i.base() == end() || *i != value) return end();
    return i.base();
}

template <typename RecordT>
template <typename T>
typename CatalogT<RecordT>::iterator CatalogT<RecordT>::lower_bound(typename Field<T>::Value const& value,
                                                                    Key<T> const& key) {
    detail::KeyExtractionFunctor<RecordT, T> f = {key};
    // Iterator adaptor that makes a CatalogT iterator work like an iterator over field values.
    typedef boost::transform_iterator<detail::KeyExtractionFunctor<RecordT, T>, iterator> SearchIter;
    SearchIter i = std::lower_bound(SearchIter(begin(), f), SearchIter(end(), f), value);
    return i.base();
}

template <typename RecordT>
template <typename T>
typename CatalogT<RecordT>::const_iterator CatalogT<RecordT>::lower_bound(
        typename Field<T>::Value const& value, Key<T> const& key) const {
    detail::KeyExtractionFunctor<RecordT, T> f = {key};
    // Iterator adaptor that makes a CatalogT iterator work like an iterator over field values.
    typedef boost::transform_iterator<detail::KeyExtractionFunctor<RecordT, T>, const_iterator> SearchIter;
    SearchIter i = std::lower_bound(SearchIter(begin(), f), SearchIter(end(), f), value);
    return i.base();
}

template <typename RecordT>
template <typename T>
typename CatalogT<RecordT>::iterator CatalogT<RecordT>::upper_bound(typename Field<T>::Value const& value,
                                                                    Key<T> const& key) {
    detail::KeyExtractionFunctor<RecordT, T> f = {key};
    // Iterator adaptor that makes a CatalogT iterator work like an iterator over field values.
    typedef boost::transform_iterator<detail::KeyExtractionFunctor<RecordT, T>, iterator> SearchIter;
    SearchIter i = std::upper_bound(SearchIter(begin(), f), SearchIter(end(), f), value);
    return i.base();
}

template <typename RecordT>
template <typename T>
typename CatalogT<RecordT>::const_iterator CatalogT<RecordT>::upper_bound(
        typename Field<T>::Value const& value, Key<T> const& key) const {
    detail::KeyExtractionFunctor<RecordT, T> f = {key};
    // Iterator adaptor that makes a CatalogT iterator work like an iterator over field values.
    typedef boost::transform_iterator<detail::KeyExtractionFunctor<RecordT, T>, const_iterator> SearchIter;
    SearchIter i = std::upper_bound(SearchIter(begin(), f), SearchIter(end(), f), value);
    return i.base();
}

template <typename RecordT>
template <typename T>
std::pair<typename CatalogT<RecordT>::iterator, typename CatalogT<RecordT>::iterator>
CatalogT<RecordT>::equal_range(typename Field<T>::Value const& value, Key<T> const& key) {
    detail::KeyExtractionFunctor<RecordT, T> f = {key};
    // Iterator adaptor that makes a CatalogT iterator work like an iterator over field values.
    typedef boost::transform_iterator<detail::KeyExtractionFunctor<RecordT, T>, iterator> SearchIter;
    std::pair<SearchIter, SearchIter> i =
            std::equal_range(SearchIter(begin(), f), SearchIter(end(), f), value);
    return std::make_pair(i.first.base(), i.second.base());
}

template <typename RecordT>
template <typename T>
std::pair<typename CatalogT<RecordT>::const_iterator, typename CatalogT<RecordT>::const_iterator>
CatalogT<RecordT>::equal_range(typename Field<T>::Value const& value, Key<T> const& key) const {
    detail::KeyExtractionFunctor<RecordT, T> f = {key};
    // Iterator adaptor that makes a CatalogT iterator work like an iterator over field values.
    typedef boost::transform_iterator<detail::KeyExtractionFunctor<RecordT, T>, const_iterator> SearchIter;
    std::pair<SearchIter, SearchIter> i =
            std::equal_range(SearchIter(begin(), f), SearchIter(end(), f), value);
    return std::make_pair(i.first.base(), i.second.base());
}


template <typename RecordT, typename Catalog, typename T>
std::shared_ptr<RecordT> _Catalog_find(Catalog const& catalog, T const& value, Key<T> const& key) {
    typename Catalog::const_iterator iter = catalog.find(value, key);
    if (iter == catalog.end()) {
        return std::shared_ptr<RecordT>();
    }
    return iter;  // n.b. CatalogIterator is explicitly convertible to shared_ptr
}

template <typename Catalog, typename T>
int _Catalog_lower_bound(Catalog const& catalog, T const& value, Key<T> const& key) {
    return catalog.lower_bound(value, key) - catalog.begin();
}

template <typename Catalog, typename T>
int _Catalog_upper_bound(Catalog const& catalog, T const& value, Key<T> const& key) {
    return catalog.upper_bound(value, key) - catalog.begin();
}

template <typename Catalog, typename T>
std::pair<int, int> _Catalog_equal_range(Catalog const& catalog, T const& value, Key<T> const& key) {
    std::pair<typename Catalog::const_iterator, typename Catalog::const_iterator> p =
            catalog.equal_range(value, key);
    return std::pair<int, int>(p.first - catalog.begin(), p.second - catalog.begin());
}

}  // namespace table
}  // namespace afw
}  // namespace lsst

#endif  // !AFW_TABLE_Catalog_h_INCLUDED