Project: s41467-025-59232-0 | Function: entropy (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:** Computes the information entropy of atomistic datasets using kernel density estimation with Gaussian kernels. Implements batch processing for memory efficiency and handles both single and multiple bandwidth calculations. Core implementation of the Shannon entropy formula for atomistic environments.

**KEY TERMS:** {
  "kernel_sum": "kernel density estimation",
  "p_x": "probability density p(x)",
  "h": "bandwidth parameter",
  "batch_size": "memory optimization parameter"
}

**CONTEXT CODE:** import math
from typing import List, Tuple, Union

import numba as nb
import numpy as np

from .geometry import cutoff_fn
from .matrix import cdist, norm, sumexp, wsumexp

DEFAULT_BANDWIDTH = 0.015
DEFAULT_BATCH = 20000
DEFAULT_UQ_NBRS = 3
DEFAULT_GRAPH_NBRS = 10


def perfect_entropy(
    x: np.ndarray,
    h: Union[float, List[float]] = DEFAULT_BANDWIDTH,
    batch_size: int = DEFAULT_BATCH,
):
    """Deprecated. Please use `entropy`.

    Computes the perfect entropy of a dataset using a batch distance
        calculation. This is necessary because the full distance matrix
        often does not fit in the memory for a big dataset. This function
        can be SLOW, despite the optimization of the computation, as it
        does not approximate the results.

    Arguments:
        x (np.ndarray): an (N, d) matrix with the descriptors
        h (int or np.nadarray): bandwidth (value / vector) for the Gaussian kernel
        batch_size (int): maximum batch size to consider when
            performing a distance calculation.

    Returns:
        entropy (float): entropy of the dataset given by `x`.
            or (np.ndarray): if 'h' is a vector
    """
    return entropy(x, h, batch_size)


def entropy(
    x: np.ndarray,
    h: Union[float, List[float]] = DEFAULT_BANDWIDTH,
    batch_size: int = DEFAULT_BATCH,
):
    """Computes the perfect entropy of a dataset using a batch distance
        calculation. This is necessary because the full distance matrix
        often does not fit in the memory for a big dataset. This function
        can be SLOW, despite the optimization of the computation, as it
        does not approximate the results.

    Arguments:
        x (np.ndarray): an (N, d) matrix with the descriptors
        h (int or np.nadarray): bandwidth (value / vector) for the Gaussian kernel
        batch_size (int): maximum batch size to consider when
            performing a distance calculation.

    Returns:
        entropy (float): entropy of the dataset given by `x`.
            or (np.ndarray): if 'h' is a vector
    """
    # MASKED: Computes the information entropy of atomistic datasets using kernel density estimation with Gaussian kernels implementation
    raise NotImplementedError("This function is masked for testing purposes.")


def delta_entropy(
    y: np.ndarray,
    x: np.ndarray,
    h: float = DEFAULT_BANDWIDTH,
    batch_size: int = DEFAULT_BATCH,
):
    """Computes the differential entropy of a dataset `y` using the dataset
        `x` as reference. This function can be SLOW, despite the optimization
        of the computation, as it does not approximate the results.

    Arguments:
        y (np.ndarray): an (M, d) matrix with the descriptors of the test set
        x (np.ndarray): an (N, d) matrix with the descriptors of the reference
        h (int or np.nadarray): bandwidth (value / vector) for the Gaussian kernel
        batch_size (int): maximum batch size to consider when
            performing a distance calculation.

    Returns:
        dH (np.ndarray): an (M,) vector of differential entropy dH ( Y | X )
            or (np.ndarray): an (H,M) matrix if 'h' is a vector of length H
    """
    if type(h) is np.ndarray:
        p_y = kernel_sum_multi_bandwidth(y, x, h=h, batch_size=batch_size)
    else:
        p_y = kernel_sum(y, x, h=h, batch_size=batch_size)
    return -np.log(p_y)


def diversity(
    x: np.ndarray,
    h: Union[float, List[float]] = DEFAULT_BANDWIDTH,
    batch_size: int = DEFAULT_BATCH,
):
    """Computes the diversity of a dataset `x` by assuming a sum over the
        inverse p(x). This approximates the number of unique data points
        in the system, as Kij >= 1 for a kernel matrix of a dataset.

    Arguments:
        x (np.ndarray): an (N, d) matrix with the descriptors of the dataset
        h (int or np.nadarray): bandwidth (value / vector) for the Gaussian kernel
        batch_size (int): maximum batch size to consider when
            performing a distance calculation.

    Returns:
        entropy (float): entropy of the dataset given by `x`.
            or (np.ndarray): if 'h' is a vector
    """
    if type(h) is np.ndarray:
        p_x = kernel_sum_multi_bandwidth(x, x, h=h, batch_size=batch_size)
        return np.log(np.sum((1 / p_x), axis=0))
    else:
        p_x = kernel_sum(x, x, h=h, batch_size=batch_size)
        return np.log(np.sum(1 / p_x))


def get_all_metrics(
    x: np.ndarray,
    h: Union[float, List[float]] = DEFAULT_BANDWIDTH,
    batch_size: int = DEFAULT_BATCH,
):
    """Computes the entropy and diversity of a dataset `x`. This function is
        slightly faster than computing them separately, as one computes p(x)
        only once instead of twice.

    Arguments:
        x (np.ndarray): an (N, d) matrix with the descriptors of the dataset
        h (int or np.nadarray): bandwidth (value / vector) for the Gaussian kernel
        batch_size (int): maximum batch size to consider when
            performing a distance calculation.

    Returns:
        entropy (float): entropy of the dataset given by `x`.
        diversity (float): diversity of the dataset given by `x`.
        dH (np.ndarray): an (M,) vector of differential entropy dH ( X | X )
    """


Complete the function implementation replacing "raise NotImplementedError".