occupation

occupation#

Occupation functions.

jrystal.occupation.gamma(num_k: int, num_electrons: int, spin: int = 0, num_bands: int | None = None, spin_restricted: bool = True) → Float[Array, 'spin kpt band'][source]#

Calculate occupation numbers for Gamma point sampling.

This function returns an occupation array where electrons are placed only at the Gamma point (first \(k\)-point). For the Gamma point, occupation is 1 for occupied bands and 0 for unoccupied bands. All other \(k\)-points have zero occupation.

The occupation array has the following structure for each spin channel:

\[\begin{split}\begin{pmatrix} 1 & 1 & 1 & 0 & \cdots & 0 \\ 0 & 0 & 0 & 0 & \cdots & 0 \\ \vdots & \vdots & \vdots & \vdots & \ddots & \vdots \\ 0 & 0 & 0 & 0 & \cdots & 0 \end{pmatrix}\end{split}\]
Parameters:

  • num_k (int) – The number of k-points.

  • num_electrons (int) – The number of electrons.

  • spin (int, optional) – The number of unpaired electrons. Defaults to 0.

  • num_bands (int | None, optional) – The number of bands (orbitals). If num_bands is not provided, then num_bands is the same as num_electrons. Defaults to None.

  • spin_restricted (bool, optional) – Indicate spin channel. If True, the first axis of output is 1. If False, the first axis of output is 2. Defaults to True.

Returns:

An occupation array with shape [num_spin, num_k, num_bands], where num_spin=1 if spin_restricted is True, else num_spin=2. The sum of occupation array equals to num_electrons.

Return type:

Float[Array, ‘spin kpt band’]

jrystal.occupation.idempotent(params: dict, num_electrons: int, num_kpts: int, spin: int = 0, spin_restricted: bool = True) → Float[Array, 'spin kpt band'][source]#

Calculate the occupation number using the idempotent method.

This function is an implementation of the method proposed in:

Li, Tianbo, et al. “Diagonalization without Diagonalization: A Direct Optimization Approach for Solid-State Density Functional Theory.” arXiv preprint arXiv:2411.05033 (2024).

Please also refer to the tutorial Occupation for more details.

Examples:

from jrystal import occupation
key = jax.random.PRNGKey(0)
num_kpts = 1
num_electrons = 10
# number of bands must be greater than or equal to
# (num_electrons + 1) // 2 for spin-restricted calculation.
num_bands = 10
params = occupation.idempotent_param_init(key, num_bands, num_kpts)
occ = occupation.idempotent(params, num_electrons, num_kpts, spin=0)
Parameters:

  • params (dict) – The parameters for the idempotent occupation. It can be generated by idempotent_param_init().

  • num_electrons (int) – The number of electrons.

  • num_kpts (int) – The number of \(k\)-points.

  • spin (int, optional) – The number of unpaired electrons. Defaults to 0.

  • spin_restricted (bool, optional) – Whether to use spin-restricted calculation. Defaults to True.

Returns:

The occupation number.

Return type:

Float[Array, ‘spin kpt band’]

jrystal.occupation.idempotent_param_init(key: Array, num_bands: int, num_kpts: int) → dict[source]#

Initialize the parameters for the idempotent occupation.

This function is an implementation of the method proposed in:

Li, Tianbo, et al. “Diagonalization without Diagonalization: A Direct Optimization Approach for Solid-State Density Functional Theory.” arXiv preprint arXiv:2411.05033 (2024).

Please also refer to the tutorial Occupation for more details.

Examples:

from jrystal import occupation
key = jax.random.PRNGKey(0)
num_kpts = 1
num_electrons = 10
# number of bands must be greater than or equal to
# (num_electrons + 1) // 2 for spin-restricted calculation.
num_bands = 10
params = occupation.idempotent_param_init(key, num_bands, num_kpts)
occ = occupation.idempotent(params, num_electrons, num_kpts, spin=0)
Parameters:

  • key (Array) – The random key for initialization.

  • num_bands (int) – The number of bands.

  • num_kpts (int) – The number of \(k\)-points.

Returns:

A dictionary containing the initialized parameters.

Return type:

dict

jrystal.occupation.uniform(num_k: int, num_electrons: int, spin: int = 0, num_bands: int | None = None, spin_restricted: bool = True) → Float[Array, 'spin kpt band'][source]#

Calculate uniform occupation numbers across k-points.

This function returns an occupation array where electrons are distributed uniformly across \(k\)-points. For each \(k\)-point, the occupation is \(1/k\) for occupied bands and 0 for unoccupied bands.

The occupation array has the following structure for each spin channel:

\[\begin{split}\begin{pmatrix} 1/k & 1/k & 1/k & 0 & \cdots & 0 \\ 1/k & 1/k & 1/k & 0 & \cdots & 0 \\ \vdots & \vdots & \vdots & \vdots & \ddots & \vdots \\ 1/k & 1/k & 1/k & 0 & \cdots & 0 \end{pmatrix}\end{split}\]

where \(k\) is the number of k-points.

Parameters:

  • num_k (int) – The number of \(k\)-points.

  • num_electrons (int) – The number of electrons.

  • spin (int, optional) – The number of unpaired electrons. Defaults to 0.

  • num_bands (int | None, optional) – The number of bands (orbitals). If num_bands is not provided, then num_bands is the same as num_electrons.Defaults to None.

  • spin_restricted (bool, optional) – Indicate spin channel. If True, the first axis of output is 1. If False, the first axis of output is 2. Defaults to True.

Returns:

An occupation array with shape [num_spin, num_k, num_bands], where num_spin=1 if spin_restricted is True, else num_spin=2. The sum of occupation array equals to num_electrons.

Return type:

Float[Array, ‘spin kpt band’]

Contents