Project: s41467-025-56378-9 | Function: adaptive_coarsegrain_gpu (python)

system: You are an expert Python programmer implementing algorithms from scientific papers.
Complete the masked function based on the algorithm description.
Follow Python best practices and maintain consistent code style.
Return **only** the complete Python function code.

user: Complete this Python function:

**DESCRIPTION:** Adaptively coarsegrains Hi-C matrices based on local neighborhood pooling of counts. Implements a sophisticated algorithm that performs forward pass coarsegraining and reverse pass value replacement to produce smoothed contact matrices while preserving signal quality.

**KEY TERMS:** {
  "_coarsen": "2x2 pooling operation",
  "_expand": "upsampling for reverse pass",
  "cutoff": "minimum counts threshold",
  "max_levels": "coarsegraining levels",
  "armask": "valid pixel mask"
}

**CONTEXT CODE:**             assert end <= self.len_chrs[chrom] and start >= 0
            encoding = self.sequence_data[:, self.inds[chrom] + start : self.inds[chrom] + end]

        if strand == "-":
            encoding = encoding[::-1, ::-1]
        assert encoding.shape[1] == end - start
        return encoding.T

    @init
    def get_encoding_from_coords_check_unk(self, chrom, start, end, strand="+", pad=False):
        """Gets the one-hot encoding of the genomic sequence at the
        queried coordinates and check whether the sequence contains
        unknown base(s).

        Parameters
        ----------
        chrom : str
            The name of the chromosome or region, e.g. "chr1".
        start : int
            The 0-based start coordinate of the first position in the
            sequence.
        end : int
            One past the 0-based last position in the sequence.
        strand : {'+', '-', '.'}, optional
            Default is '+'. The strand the sequence is located on. '.' is
            treated as '+'.
        pad : bool, optional
            Default is `False`. Pad the output sequence with 'N' if `start`
            and/or `end` are out of bounds to return a sequence of length
            `end - start`.

        Returns
        -------
        tuple(numpy.ndarray, bool)
            * `tuple[0]` is the :math:`L \\times 4` encoding of the sequence, where
            :math:`L = end - start`.
            `L` = 0 for the NumPy array returned.
            * `tuple[1]` is the boolean value that indicates whether the
            sequence contains any unknown base(s) specified in self.UNK_BASE


        Raises
        ------
        AssertionError
            If it cannot retrieve encoding that matches the length `L = end - start`
            such as when end > chromosome length and pad=False
        """
        encoding = self.get_encoding_from_coords(chrom, start, end, strand=strand, pad=strand)
        return encoding, np.any(encoding[0, :] == 0.25)

def adaptive_coarsegrain_gpu(ar, countar, cutoff=5, max_levels=8, min_shape=8):
    """
    Adaptively coarsegrain a Hi-C matrix based on local neighborhood pooling
    of counts.

    Parameters
    ----------
    ar : torch.Tensor, shape (n, n)
        A square Hi-C matrix to coarsegrain. Usually this would be a balanced
        matrix.

    countar : torch.Tensor, shape (n, n)
        The raw count matrix for the same area. Has to be the same shape as the
        Hi-C matrix.

    cutoff : float, optional
        A minimum number of raw counts per pixel required to stop 2x2 pooling.
        Larger cutoff values would lead to a more coarse-grained, but smoother
        map. 3 is a good default value for display purposes, could be lowered
        to 1 or 2 to make the map less pixelated. Setting it to 1 will only
        ensure there are no zeros in the map.

    max_levels : int, optional
        How many levels of coarsening to perform. It is safe to keep this
        number large as very coarsened map will have large counts and no
        substitutions would be made at coarser levels.
    min_shape : int, optional
        Stop coarsegraining when coarsegrained array shape is less than that.

    Returns
    -------
    Smoothed array, shape (n, n)

    Notes
    -----
    The algorithm works as follows:

    First, it pads an array with NaNs to the nearest power of two. Second, it
    coarsens the array in powers of two until the size is less than minshape.

    Third, it starts with the most coarsened array, and goes one level up.
    It looks at all 4 pixels that make each pixel in the second-to-last
    coarsened array. If the raw counts for any valid (non-NaN) pixel are less
    than ``cutoff``, it replaces the values of the valid (4 or less) pixels
    with the NaN-aware average. It is then applied to the next
    (less coarsened) level until it reaches the original resolution.

    In the resulting matrix, there are guaranteed to be no zeros, unless very
    large zero-only areas were provided such that zeros were produced
    ``max_levels`` times when coarsening.

    Examples
    --------
    >>> c = cooler.Cooler("/path/to/some/cooler/at/about/2000bp/resolution")

    >>> # sample region of about 6000x6000
    >>> mat = c.matrix(balance=True).fetch("chr1:10000000-22000000")
    >>> mat_raw = c.matrix(balance=False).fetch("chr1:10000000-22000000")
    >>> mat_cg = adaptive_coarsegrain(mat, mat_raw)

    >>> plt.figure(figsize=(16,7))
    >>> ax = plt.subplot(121)
    >>> plt.imshow(np.log(mat), vmax=-3)
    >>> plt.colorbar()
    >>> plt.subplot(122, sharex=ax, sharey=ax)
    >>> plt.imshow(np.log(mat_cg), vmax=-3)
    >>> plt.colorbar()

    """
    # MASKED: Adaptively coarsegrains Hi-C matrices based on local neighborhood pooling of counts. Implements a sophisticated algorithm that performs forward pass coarsegraining and reverse pass value replacement to produce smoothed contact matrices while preserving signal quality.
    raise NotImplementedError("This function is masked for testing purposes.")


def _adaptive_coarsegrain(ar, countar, max_levels=12, cuda=False):
    """
    Wrapper for cooltools adaptive coarse-graining to add support 
    for non-square input for interchromosomal predictions.
    """
    global adaptive_coarsegrain_fn
    if cuda:
        adaptive_coarsegrain_fn = adaptive_coarsegrain_gpu
    else:
        adaptive_coarsegrain_fn = adaptive_coarsegrain


    assert np.all(ar.shape == countar.shape)
    if ar.shape[0] < 9 and ar.shape[1] < 9:
        ar_padded = np.empty((9, 9))
        ar_padded.fill(np.nan)
        ar_padded[: ar.shape[0], : ar.shape[1]] = ar

        countar_padded = np.empty((9, 9))
        countar_padded.fill(np.nan)
        countar_padded[: countar.shape[0], : countar.shape[1]] = countar
        return adaptive_coarsegrain_fn(ar_padded, countar_padded, max_levels=max_levels)[
            : ar.shape[0], : ar.shape[1]
        ]

    if ar.shape[0] == ar.shape[1]:
        return adaptive_coarsegrain_fn(ar, countar, max_levels=max_levels)


Complete the function implementation replacing "raise NotImplementedError".