Calculate#

diatomic.calculate.electric_moment(mol, eigstates)#

Calculates the electric moments of each eigenstate.

Parameters:
  • mol – The molecule used to generate the eigenstates

  • eigstates (np.ndarray) – Eigenstates from diagonalisation with shape (basis, states) or (steps, basis, states).

Returns:

Electric moments with shape eigstates.shape[:-2] + (states,).

Return type:

ndarray

diatomic.calculate.label_states(mol, eigstates, labels, index_repeats=False, basis_idx=None, warn_mixed=True, min_weight=0.9)#

Labels the eigenstates of a molecule with quantum numbers corresponding to specified labels. The returned numbers will only be good if the state is well-represented in the desired basis.

This function processes a list of labels that specify which quantum numbers to calculate for the given eigenstates of a molecule. It uses the angular momentum operators for the molecule and the eigenstates to calculate these quantum numbers.

The function supports labels for the total angular momentum vectors (N, F, I, I1, I2) and for their z-components respectively (M<>). Labels are inferred from operator expectation values. If warn_mixed is True, the function warns when an assigned label sector carries less than min_weight of an eigenstate’s probability weight.

Parameters:
  • mol – An object representing the molecule, which contains the maximum angular momentum (Nmax), minimum angular momentum (Nmin), and intrinsic angular momenta (Ii).

  • eigstates (ndarray) – Eigenstates for which the labels are to be calculated. The expected shape is (basis, states). A single eigenvector with shape (basis,) is also accepted.

  • labels (list of str) – A list of strings indicating which quantum numbers to calculate. Valid labels are “MN”, “MF”, “MI”, “MI1”, “MI2” for M-components and “N”, “F”, “I”, “I1”, “I2” for squared total angular momentum vectors.

  • index_repeats (bool, optional) – If True, includes the label indices of repeated labels in the output. Defaults to False.

  • basis_idx (ndarray | None, optional) – Indices mapping a cropped eigstates basis back into the full uncoupled basis. Required when eigstates were diagonalised in a cropped basis.

  • warn_mixed (bool, optional) – If True, warns when labels are assigned to states with weak weight in the labelled sector. Defaults to True.

  • min_weight (float, optional) – Minimum sector weight before warning that a label may be unreliable. Defaults to 0.9.

Returns:

A two-dimensional array where each row corresponds to the desired

quantum numbers for each eigenstate, and each column corresponds to a label.

Return type:

ndarray

diatomic.calculate.magnetic_moment(mol, eigstates)#

Calculates the magnetic moments of each eigenstate.

Parameters:
  • mol – The molecule used to generate the eigenstates

  • eigstates (np.ndarray) – Eigenstates from diagonalisation with shape (basis, states) or (steps, basis, states).

Returns:

Magnetic moments with shape eigstates.shape[:-2] + (states,).

Return type:

ndarray

diatomic.calculate.solve_system(hamiltonians, num_diagonals=None, progress=False, chunk_size=None)#

Solves for the eigenenergies and eigenstates of given Hamiltonian(s) over one parameter.

This function diagonalizes the input Hamiltonian(s) to find the corresponding eigenenergies and eigenstates. It supports both single Hamiltonian and a set of Hamiltonians (e.g., varying with time or another parameter).

It takes care to make sure the eigenenergies and eigenstates vary smoothly with respect to the varying parameter, unlike np.eigh.

Parameters:
  • hamiltonians (np.ndarray) – A single Hamiltonian matrix or an array of Hamiltonian matrices.

  • num_diagonals (int | None, optional) – The number of local swaps above and below to consider when smoothing eigenfunctions. If None, compare all eigenstates to all other eigenstates when smoothing (safest).

  • progress (bool, optional) – If True, show a progress bar while diagonalising scan-shaped Hamiltonians. Uses tqdm if available. Defaults to False.

  • chunk_size (int | None, optional) – Number of Hamiltonians to diagonalise in each batch when using chunked diagonalisation. If None, the original fully batched diagonalisation is used unless progress=True, in which case chunk_size defaults to 1 so the progress bar can update after each step.

Returns:

(eigenenergies, eigenstates).

Return type:

tuple

Raises:

ValueError – If the input Hamiltonian has more than three dimensions.

diatomic.calculate.sort_by_labels(eigenlabels, eigenenergies, eigenstates)#

Sorts eigenstates and eigenenergies according to the provided labels lexicographically.

This function sorts the eigenstates and eigenenergies in ascending order based on the eigenlabels, which can be quantum numbers or any other identifiers for the states. They are sorted with the first label value taking preference, and then the second if the first drew, as with a language dictionary.

Parameters:
  • eigenlabels (np.ndarray) – Array of labels corresponding to eigenstates.

  • eigenenergies (np.ndarray) – Array of eigenenergies to be sorted along the last axis. Single-Hamiltonian and scan-shaped outputs are supported.

  • eigenstates (np.ndarray) – Array of eigenstates to be sorted along the last axis. Single-Hamiltonian and scan-shaped outputs are supported.

Returns:

The sorted eigenlabels, eigenenergies, and eigenstates.

Return type:

tuple

diatomic.calculate.sort_smooth(eigvals, eigvecs, num_diagonals=None)#

Smooths the eigenvalue trajectories as a function of an external parameter, rather than by energy order, as is the case with np.eigh.

This function sorts the eigenvalues and eigenvectors to maintain continuity as the external parameter changes, preventing sudden jumps in the eigenvalue ordering.

Parameters:
  • eigvals (np.ndarray) – Eigenvalues array with shape (steps, eigstate).

  • eigvecs (np.ndarray) – Eigenvectors array with shape (steps, basis, eigstate).

  • num_diagonals (int | None, optional) – The number of local swaps above and below to consider when smoothing eigenfunctions. Restricting this can reduce work if eigenstates only swap locally in energy order. If None, compare all eigenstates to all other eigenstates when smoothing (safest).

Returns:

The smoothed eigenvalues and eigenvectors.

Return type:

tuple

Raises:

ValueError – If eigvecs does not have three dimensions.

diatomic.calculate.transition_electric_moments(mol, eigenstates, h=0, from_states=slice(None, None, None), to_states=slice(None, None, None))#

Calculates electric dipole transition matrix elements between molecular eigenstates.

These are the absolute values of <from|d_h|to>, where d_h is the spherical component of the molecule-fixed electric dipole operator transformed into the lab basis and expanded in the eigenstate basis. Multiplying by an electric field amplitude gives the corresponding dipole coupling energy scale.

Parameters:
  • mol – An object representing the molecule.

  • eigenstates (ndarray) – An array containing the eigenstates of the molecule, with shape (…, basis, states).

  • h (int, optional) – The helicity of the transition, +1 for sigma+, 0 for pi, -1 for sigma-.

  • from_states (int, slice, or list, optional) – The indices of the initial states from which transitions are considered. If an integer is passed, it is internally converted to a list. By default, all states are considered (slice(None)). These indices refer to the last axis of the supplied eigenstates array, so they are relative to any cropped eigenstate array.

  • to_states (int, slice, or list, optional) – The indices of the final states to which transitions are considered. If an integer is passed, it is internally converted to a list. By default, all states are considered (slice(None)). These indices also refer to the last axis of the supplied eigenstates array, so they are relative to any cropped eigenstate array.

Returns:

Absolute transition electric moments with shape

eigenstates.shape[:-2] + (n_from, n_to). The second-to-last axis enumerates the selected initial/from-state columns, in from_states order. The last axis enumerates the selected final/to-state columns, in to_states order. Output indices are positions inside those supplied from_states and to_states selections, not necessarily the original eigenstate indices.

Return type:

ndarray