docs: def query(self, coords, return_sigma=False):
          """
          Returns 2MASS Ks-band extinction at the given coordinates.
  
          Args:
              coords (:obj:`astropy.coordinates.SkyCoord`): The coordinates to query.
                  Must contain distances.
              return_sigma (Optional[:obj:`bool`]): If ``True``, returns the uncertainty in
                  extinction as well. Defaults to ``False``.
  
          Returns:
              Extinction at the specified coordinates, in mags of 2MASS Ks-band
              extinction. If ``return_sigma`` is ``True``, then the uncertainty
              in reddening is also returned, so that the output is
              ``(A, sigma_A)``, where both ``A`` and ``sigma_A`` have the same
              shape as the input coordinates.
          """
  
          # Ensure that distance has been requested
          has_dist = hasattr(coords.distance, 'kpc')
          if not has_dist:
              raise ValueError('Input `coords` must specify distance.')
  
          # Convert coordinates to pixel indices
          j, k, mask_idx = self._gal2idx(coords)
  
          # Which distances to extract
          d = coords.distance.kpc
          dist_idx_ceil = np.sum(d[:,None] > self._dist[j, k, :], axis=1)
  
          # Initialize return arrays
          A_ret = np.full(coords.shape, np.nan, dtype='f4')
          if return_sigma:
              sigma_ret = np.full(coords.shape, np.nan, dtype='f4')
  
          # d < d(nearest distance slice)
          idx_near = (dist_idx_ceil == 0) & mask_idx
  
          if np.any(idx_near):
              a = d[idx_near] / self._dist[j[idx_near], k[idx_near], 0]
              A_ret[idx_near] = a * self._A[j[idx_near], k[idx_near], 0]
  
              if return_sigma:
                  sigma_ret[idx_near] = (
                      self._sigma_A[
                          j[idx_near],
                          k[idx_near],
                          0
                      ])
  
          # d > d(farthest distance slice)
          idx_far = (dist_idx_ceil == self._n_dists[j,k]) & mask_idx
  
          if np.any(idx_far):
              A_ret[idx_far] = (
                  self._A[
                      j[idx_far],
                      k[idx_far],
                      self._n_dists[j[idx_far],k[idx_far]]-1
                  ])
  
              if return_sigma:
                  sigma_ret[idx_far] = (
                      self._sigma_A[
                          j[idx_far],
                          k[idx_far],
                          self._n_dists[j[idx_far],k[idx_far]]-1
                      ])
  
          # d(nearest distance slice) < d < d(farthest distance slice)
          idx_btw = (~idx_near & ~idx_far) & mask_idx
  
          if np.any(idx_btw):
              d_ceil = (
                  self._dist[
                      j[idx_btw],
                      k[idx_btw],
                      dist_idx_ceil[idx_btw]
                  ])
              d_floor = (
                  self._dist[
                      j[idx_btw],
                      k[idx_btw],
                      dist_idx_ceil[idx_btw]-1
                  ])
              a = (d_ceil - d[idx_btw]) / (d_ceil - d_floor)
  
              A_ret[idx_btw] = (
                  (1.-a) * self._A[j[idx_btw], k[idx_btw], dist_idx_ceil[idx_btw]]
                  + a * self._A[j[idx_btw], k[idx_btw], dist_idx_ceil[idx_btw]-1]
              )
  
              if return_sigma:
                  w0 = (1.-a)**2
                  w1 = a**2
                  norm = 1. / (w0 + w1)
                  w0 *= norm
                  w1 *= norm
                  sigma_ret[idx_btw] = np.sqrt(
                      w0 * self._sigma_A[j[idx_btw], k[idx_btw], dist_idx_ceil[idx_btw]]**2
                      + w1 * self._sigma_A[j[idx_btw], k[idx_btw], dist_idx_ceil[idx_btw]-1]**2
                  )
  
          if return_sigma:
              return A_ret, sigma_ret
  
          return A_ret