Project: s41467-025-59232-0 | Function: descriptor_x1 (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 first-order structural descriptor for atomic environments based on distance-weighted neighbor contributions. Implements a cutoff function and distance-based weighting scheme to capture local atomic structure. Essential for generating feature vectors representing atomic environments.

**KEY TERMS:** {
  "dm": "distance matrix",
  "sorter": "sorted neighbor indices",
  "cutoff_fn": "cutoff weighting function",
  "x1": "first-order descriptor",
  "wij": "cutoff weights"
}

**CONTEXT CODE:** import math

import numba as nb
import numpy as np
from ase import Atoms
from numba import types
from numba.typed import Dict, List

from .geometry import cutoff_fn
from .matrix import argsort, cdist, inverse_3d, pdist, stack_xyz

IntList = types.ListType(types.int64)
FloatArrayList = types.Array(types.float64, 1, "C")

DEFAULT_CUTOFF: float = 5.0
DEFAULT_K: int = 32
EPS: float = 1e-15


@nb.njit(fastmath=True, cache=True)
def descriptor_x1(
    dm: np.ndarray,
    sorter: np.ndarray,
    k: int = 32,
    cutoff: float = 5.0,
    max_rows: int = -1,
    eps: float = 1e-16,
) -> np.ndarray:
    # MASKED: Computes the first-order structural descriptor for atomic environments based on distance-weighted neighbor contributions implementation
    raise NotImplementedError("This function is masked for testing purposes.")


@nb.njit(fastmath=True, cache=True)
def descriptor_x2(
    dm: np.ndarray,
    sorter: np.ndarray,
    k: int = 32,
    cutoff: float = 5.0,
    eps: float = 1e-15,
    max_rows: int = -1,
) -> np.ndarray:
    N = dm.shape[0]
    if max_rows <= 0:
        max_rows = N

    # Lazy initialization of the matrix
    x2 = np.full((max_rows, k - 1), fill_value=0.0)
    jmax = k if N > k else (N - 1)

    # Computes the second descriptor
    for i in range(max_rows):
        rjl = np.full((k, k), fill_value=0.0)

        # first compute the cross distances
        for j in range(jmax):
            atom_j = sorter[i, j + 1]
            rij = dm[i, atom_j]
            wij = cutoff_fn(rij, cutoff)

            for l in range(j + 1, jmax):
                atom_l = sorter[i, l + 1]
                ril = dm[i, atom_l]
                wil = cutoff_fn(ril, cutoff)

                x2_jl = math.sqrt(wij * wil) / (dm[atom_j, atom_l] + eps)
                rjl[j, l] = x2_jl
                rjl[l, j] = x2_jl

        r_sort = np.sort(rjl)

        # now compute the mean over rows
        # and sorts largest first in x2
        for l in range(1, k):
            _sum = 0.0
            for j in range(k):
                _sum += r_sort[j, l]

            # larger first
            x2[i, k - 1 - l] = _sum / k

    return x2


@nb.njit(fastmath=True, cache=True)
def descriptor_nopbc(
    xyz: np.ndarray,
    k: int = DEFAULT_K,
    cutoff: float = DEFAULT_CUTOFF,
    eps: float = EPS,
) -> np.ndarray:
    dm = pdist(xyz)
    sorter = argsort(dm)

    x1 = descriptor_x1(dm, sorter, k, cutoff)
    x2 = descriptor_x2(dm, sorter, k, cutoff)
    return x1, x2


@nb.njit(fastmath=True, cache=True)
def get_num_bins(cell: np.ndarray, cutoff: float):
    bx = np.cross(cell[1], cell[2])
    by = np.cross(cell[2], cell[0])
    bz = np.cross(cell[0], cell[1])

    bx_norm = np.sqrt(bx[0] * bx[0] + bx[1] * bx[1] + bx[2] * bx[2])
    by_norm = np.sqrt(by[0] * by[0] + by[1] * by[1] + by[2] * by[2])
    bz_norm = np.sqrt(bz[0] * bz[0] + bz[1] * bz[1] + bz[2] * bz[2])

    volume = np.dot(cell[0], bx)

    lx = volume / bx_norm
    ly = volume / by_norm
    lz = volume / bz_norm
    lengths = np.array([lx, ly, lz])

    nx = max(np.floor(lx / cutoff), 1)
    ny = max(np.floor(ly / cutoff), 1)
    nz = max(np.floor(lz / cutoff), 1)

    n_bins = np.array([nx, ny, nz], dtype=np.int64)


Complete the function implementation replacing "raise NotImplementedError".