DiracSpinor.hpp

#pragma once
#include "Physics/PhysConst_constants.hpp" // for default alpha
#include <cstdint>
#include <memory>
#include <string>
#include <utility>
#include <vector>
class Grid;
 
/*!
  @brief Stores radial Dirac spinor: F_nk = (f, g)
 
  @details
  \f[
  \psi_{n\kappa m} = \frac{1}{r}
  \begin{pmatrix}
    f_{n\kappa}(r)\,\Omega_{\kappa m}\\
    g_{n\kappa}(r)\,\Omega_{-\kappa m}
  \end{pmatrix},
  \quad
  F_{n\kappa} =
  \begin{pmatrix}
    f_{n\kappa}(r)\\
    g_{n\kappa}(r)
  \end{pmatrix}
  \f]
 
  \par  Construction.
  Takes in n and k=kappa values + grid
    - A shared pointer to the Grid is stored (to avoid many copies of Grid)
 
  \par Operator Overloads
    - Intuitive operator overloads are provided (Fa, Fb are DiracSpinors):
    - v * Fa, where v is a double or vector works the obvious way
    - Fa * Fb = <Fa|Fb>
    - Fa == Fb returns true if {na,ka}=={nb,kb}
    - Fa > Fb : first compares n, and then kappa (via kappa_index)
    - Fa +/- Fb : Adds/subtracts the two spinors (and updates p0/pinf)
    - You can make copies: auto Fnew = Fa
    - And you can re-assign: Fb = Fa (copies all, including n and kappa)
    - n has no setter; kappa only via set_new_kappa(). Avoids angular errors.
    - non-const accessors (e.g., en(), f(), max_pt()) return mutable
      references to variables
 
  \par Continuum orbitals
  
  @note 
  Continuum (and other unlabelled) orbitals have n = 0. Since they are
  identified only by kappa, all n = 0 orbitals with the same kappa compare
  equal (operator==), and their nk_index() coincides with that of a bound
  state (see nk_index). They must not be mixed with each other, or with
  bound states, in anything that identifies orbitals by {n, kappa}.
  This might need to be updated later.
*/
class DiracSpinor {
 
public:
  /*!
    @brief Constructor: Requires n (PQN), kappa (Dirac QN), and grid.
 
    @details
    Grid is a shared pointer, as it's a shared resource. f and g are
    initialised to zero over the full grid (max_pt() = grid size).
 
    @param in_n             Principal quantum number
    @param in_kappa         Dirac quantum number, kappa
    @param in_rgrid         Radial grid (shared resource)
    @param negative_energy  Marks a negative-energy (Dirac sea) solution
  */
  DiracSpinor(int in_n, int in_kappa, std::shared_ptr<const Grid> in_rgrid,
              bool negative_energy = false);
  //! Integer type for the compressed (n,kappa) index
  using Index = uint16_t;
 
private:
  // Radial Grid; links F[i] to F(r)
  std::shared_ptr<const Grid> m_rgrid;
  // Principal quantum number
  int m_n;
  // Dirac quantum number, kappa
  int m_kappa;
  // Single-particle energy, not including rest energy
  double m_en = 0.0;
  // Upper (large) radial component
  std::vector<double> m_f;
  // Lower (small) radial component
  std::vector<double> m_g;
  // `practical zero': p0 is first non-zero point for f(r) [usually p0=0]
  std::size_t m_p0 = 0;
  // `practical infinity': pinf is last non-zero point for f(r)
  std::size_t m_pinf;
  // Number of iterations until energy convergence (for latest routine only)
  int m_its = -1;
  // Fractional energy convergence: eps = |(en'-en)/en|
  double m_eps = -1.0;
  // Occupation fraction. =1 for closed shells. =1/(2j+1) for valence
  double m_occ_frac = 0.0;
 
  // 2j, l, pi, kappa_index (for convenience): these are defined by n and kappa
  int m_twoj;
  int m_l;
  int m_parity;
  std::size_t m_kappa_index;
  Index m_nkappa_index;
 
  // flag for regular electron, or "exotic"
  bool m_exotic{false};
  // flag for negative-energy (Dirac sea) solution
  bool m_negative_energy{false};
 
public:
  //! Principal quantum number, n
  int n() const { return m_n; }
  //! Dirac quantum number, kappa
  int kappa() const { return m_kappa; }
  //! Orbital angular momentum Q number
  int l() const { return m_l; }
  //! j(j+1)
  double jjp1() const { return 0.25 * double(m_twoj * (m_twoj + 2)); }
  //! 2j (twice the total angular momentum)
  int twoj() const { return m_twoj; }
  //! 2j+1
  int twojp1() const { return m_twoj + 1; }
  //! (-1)^l, returns +/- 1
  int parity() const { return m_parity; }
  //! kappa index (see Angular::kappa_to_kindex)
  std::size_t k_index() const { return m_kappa_index; }
  /*!
    @brief (n,kappa) index
 
    @details
    See @ref Angular::nk_to_index.
 
    @warning Unique only for 1 <= n <= 256: larger n overflows Index (not
    checked here; see @ref SplineBasis::form_basis and @ref Angular::nk_to_index). 
    
    @note 
    For n = 0 (continuum orbitals), 
    it coincides with a bound-state index (that of n = 2: e.g.,
    continuum s and 2s, p_1/2 and 2p_1/2, d_3/2 and 3s), so continuum and
    bound states must not share a table keyed by this index.
  */
  Index nk_index() const { return m_nkappa_index; }
  //! Single-electron term symbol (e.g., 6s_1/2). Gnuplot=true => 6s_{1/2}
  std::string symbol(bool gnuplot = false) const;
  //! e.g., 6p_1/2 => 6p-, 6p_3/2 => 6p+
  std::string shortSymbol() const;
 
  //! Checks if spinor is for "exotic" lepton, or regular electron
  bool &exotic() { return m_exotic; }
  //! Checks if spinor is for "exotic" lepton, or regular electron
  bool exotic() const { return m_exotic; }
 
  /*!
    @brief Returns true if this is a negative-energy (Dirac sea) state.
 
    @details
    Set at construction (by the basis builder); independent of alpha and mass.
  */
  bool negativeEnergyStateQ() const { return m_negative_energy; }
 
  //! Changes 'kappa' angular quantum number. Use with caution!
  void set_new_kappa(int new_kappa);
 
  //! Returns a const reference to the radial grid
  const Grid &grid() const { return *m_rgrid; };
  /*!
    @brief Returns copy of shared_ptr to grid [shared resource].
 
    @details
    Used when we want to construct a new DiracSpinor that shares this grid.
  */
  std::shared_ptr<const Grid> grid_sptr() const { return m_rgrid; };
 
  //! Single-particle energy, not including rest energy
  double en() const { return m_en; }
  double &en() { return m_en; }
 
  //! Upper (large) radial component function, f
  const std::vector<double> &f() const { return m_f; }
  std::vector<double> &f() { return m_f; }
  //! Upper (large) radial component, f(r_i)
  double f(std::size_t i) const { return m_f.at(i); }
  double &f(std::size_t i) { return m_f.at(i); }
 
  //! Lower (small) radial component function, g
  const std::vector<double> &g() const { return m_g; }
  std::vector<double> &g() { return m_g; }
  //! Lower (small) radial component, g(r_i)
  double g(std::size_t i) const { return m_g.at(i); }
  double &g(std::size_t i) { return m_g.at(i); }
 
  /*!
    @brief First non-zero point (index for f[i])
 
    @deprecated Candidate for removal
  */
  auto min_pt() const { return m_p0; }
  auto &min_pt() { return m_p0; }
 
  //! Effective size(); index _after_ last non-zero point (index for f[i])
  auto max_pt() const { return m_pinf; }
  auto &max_pt() { return m_pinf; }
 
  /*!
    @brief r0 = r[min_pt] (in atomic units)
 
    @deprecated Candidate for removal
  */
  double r0() const;
  //! rinf = r[max_pt - 1], the last point of the spinor (au); 0 if max_pt = 0
  double rinf() const;
 
  //! Number of iterations until energy convergence (for latest routine only)
  auto its() const { return m_its; }
  auto &its() { return m_its; }
 
  //! Occupation fraction. =1 for closed shells. =1/(2j+1) for valence
  auto occ_frac() const { return m_occ_frac; }
  auto &occ_frac() { return m_occ_frac; }
 
  //! Fractional energy convergence: eps = |(en'-en)/en|
  auto eps() const { return m_eps; }
  auto &eps() { return m_eps; }
 
  //! Scales DiracSpinor (f and g) by constant
  const DiracSpinor &scale(const double factor);
  /*!
    @brief Scales DiracSpinor (f and g) by function of r
 
    @details
    Pointwise, over [min_pt, max_pt). Points beyond the end of @p v (up to
    max_pt) are set to zero.
  */
  const DiracSpinor &scale(const std::vector<double> &v);
  //! By default normalises to 1, but can normalise to other number.
  void normalise(double norm_to = 1.0);
  //! Forces f(r) and g(r) to be zero outside of [p0,pinf)
  void zero_boundaries();
 
  //! Returns the norm, defined: Norm = Sqrt[<a|a>]
  double norm() const;
  //! Returns the norm^2, defined: Norm2 = <a|a>
  double norm2() const;
 
  /*!
    @brief Number of nodes (sign changes) in f(r).
 
    @details
    Should be n-l-1 for bound states.
 
    Only counts out to the last point where |f| >= rel_cut*max|f|,
    so that small-amplitude sign changes in the tail (e.g., the
    exchange-dominated tail of Hartree-Fock orbitals), which are not physical
    nodes, are not counted. Likewise, sign changes before |f| first reaches
    1e-10*max|f| (near the nucleus, where high-l orbitals are negligible) are
    not counted.
 
    @param rel_cut  Relative amplitude cut-off for the tail (see above)
    @return Number of nodes; -1 if f is zero or not finite (failed solution).
  */
  int num_nodes(double rel_cut = 1.0e-3) const;
 
  //! Returns [f[p0]/f_max , f[pinf-1]/f_max] - for tests; {0,0} if f = 0
  std::pair<double, double> r0pinfratio() const;
 
  /*!
    @brief Radial density: rho(r) = (2j+1) * x_occ * |F^2|(r)
 
    @details
    \f$ \rho(r) = 4\pi r^2 \, x_{\rm occ} \sum_m |\psi_{n\kappa m}|^2
    = (2j+1) \, x_{\rm occ} \, (f^2 + g^2) \f$, with x_occ = occ_frac().
    Integrates (over r) to the number of electrons, for a normalised orbital.
  */
  std::vector<double> rho() const;
 
  //! Number of occupied electrons: (2j+1)*occ_frac, rounded to an integer
  int num_electrons() const;
 
public:
  // Operator overloads
 
  //! Returns radial integral (Fa,Fb) = Int(fa*fb + ga*gb)
  friend double operator*(const DiracSpinor &Fa, const DiracSpinor &Fb);
 
  /*!
    @brief Adds rhs to this, in place (should generally have same kappa, but forced)
 
    @details
    Kappa is not checked: the result keeps n and kappa of `this`. Range
    [min_pt, max_pt) is extended to include that of @p rhs.
  */
  DiracSpinor &operator+=(const DiracSpinor &rhs);
  //! Subtracts rhs from this, in place; should generally have same kappa (see +=)
  DiracSpinor &operator-=(const DiracSpinor &rhs);
  //! Adds x*rhs to this, in place (no temporary); should generally have same kappa
  DiracSpinor &add_scaled(double x, const DiracSpinor &rhs);
  //! Returns lhs + rhs; should generally have same kappa (see +=)
  friend DiracSpinor operator+(DiracSpinor lhs, const DiracSpinor &rhs);
  //! Returns lhs - rhs; should generally have same kappa (see +=)
  friend DiracSpinor operator-(DiracSpinor lhs, const DiracSpinor &rhs);
 
  //! Scales this by scalar x, in place
  DiracSpinor &operator*=(const double x);
  //! Returns Fa scaled by scalar x
  friend DiracSpinor operator*(DiracSpinor Fa, const double x);
  //! Returns Fa scaled by scalar x
  friend DiracSpinor operator*(const double x, DiracSpinor Fa);
 
  //! Multiplies this by function v(r), pointwise, in place
  DiracSpinor &operator*=(const std::vector<double> &v);
  //! Returns Fa multiplied by function v(r), pointwise
  friend DiracSpinor operator*(const std::vector<double> &v, DiracSpinor Fa);
 
  /*!
    @brief Equality: true if {n,kappa} match
 
    @note Compares labels only: any two continuum orbitals (n = 0) with the
    same kappa compare equal, whatever their energy.
  */
  friend bool operator==(const DiracSpinor &lhs, const DiracSpinor &rhs);
  //! Inequality: true if {n,kappa} differ
  friend bool operator!=(const DiracSpinor &lhs, const DiracSpinor &rhs);
  //! Ordering: compares n, then kappa (via kappa_index)
  friend bool operator<(const DiracSpinor &lhs, const DiracSpinor &rhs);
  //! Ordering: compares n, then kappa (via kappa_index)
  friend bool operator>(const DiracSpinor &lhs, const DiracSpinor &rhs);
  //! Ordering: compares n, then kappa (via kappa_index)
  friend bool operator<=(const DiracSpinor &lhs, const DiracSpinor &rhs);
  //! Ordering: compares n, then kappa (via kappa_index)
  friend bool operator>=(const DiracSpinor &lhs, const DiracSpinor &rhs);
 
  //! Comparator (for sorting): by l
  static bool comp_l(const DiracSpinor &lhs, const DiracSpinor &rhs) {
    return lhs.m_l < rhs.m_l;
  }
  //! Comparator (for sorting): by 2j
  static bool comp_j(const DiracSpinor &lhs, const DiracSpinor &rhs) {
    return lhs.m_twoj < rhs.m_twoj;
  }
  //! Comparator (for sorting): by kappa_index
  static bool comp_ki(const DiracSpinor &lhs, const DiracSpinor &rhs) {
    return lhs.m_kappa_index < rhs.m_kappa_index;
  }
  //! Comparator (for sorting): by n
  static bool comp_n(const DiracSpinor &lhs, const DiracSpinor &rhs) {
    return lhs.m_n < rhs.m_n;
  }
  //! Comparator (for sorting): by energy
  static bool comp_en(const DiracSpinor &lhs, const DiracSpinor &rhs) {
    return lhs.en() < rhs.en();
  }
 
  //! Writes short symbol to ostream
  friend std::ostream &operator<<(std::ostream &ostr, const DiracSpinor &Fa) {
    return ostr << Fa.shortSymbol();
  }
 
  // Static (helper) functions:
 
  //! Returns worst |<a|b>| (or ||<a|b>|-1| for a=b) {val, state_names}
  static std::pair<double, std::string>
  check_ortho(const std::vector<DiracSpinor> &a,
              const std::vector<DiracSpinor> &b);
 
  /*!
    @brief (approximately) OrthoNormalises a set of any orbitals
 
    @details
    Symmetric: each orbital is orthogonalised against all others of the same
    kappa,
    \f$ \ket{a} \to \ket{a} - \frac{1}{2}\sum_{b \neq a}
    \ket{b}\langle b|a\rangle \f$ (overlaps from the start of the iteration),
    then normalised. Half of each overlap is removed from each of the pair,
    so the overlaps are removed to first order only. Repeated @p num_its
    times.
 
    @note Only updates orbs, not energies
  */
  static void orthonormaliseOrbitals(std::vector<DiracSpinor> &orbs,
                                     int num_its = 1);
 
  /*!
    @brief Forces Fin to be orthogonal to orbs.
 
    @details
    If Fn (i.e., {n,k}) is in orbs, replaces Fin with that from orbs.
    Otherwise, subtracts the projection of @p Fin onto each orbital in
    @p orbs with the same kappa (exact only if @p orbs are orthonormal).
 
    @note Energy is not updated (unless Fin is replaced from @p orbs).
  */
  [[nodiscard]] static DiracSpinor
  orthogonaliseWrt(const DiracSpinor &Fin,
                   const std::vector<DiracSpinor> &orbs);
 
  //! As orthogonaliseWrt(), but normalises Fin afterwards
  [[nodiscard]] static DiracSpinor
  orthonormaliseWrt(const DiracSpinor &Fin,
                    const std::vector<DiracSpinor> &orbs);
 
  /*!
    @brief A more correct way to force orthogonality than normal
 
    @details
    |this> -> |this> - <this|rhs>|rhs>, only over the range where both
    spinors are defined: min_pt() and max_pt() of this are unchanged. Does
    nothing if kappa differs.
 
    @note Assumes @p rhs is normalised.
  */
  void orthog(const DiracSpinor &rhs);
 
  //! e.g., 6p_1/2 => 6p-, 6p_3/2 => 6p+
  static std::string shortSymbol(int n, int kappa);
 
  //! Which energy branch of a set of orbitals state_config() describes
  enum class Branch { electron, positron, both };
 
  /*!
    @brief Returns formatted states string (e.g., '7sp5d') given list of orbs
 
    @details
    For the electron states, the number is the maximum n for each l. For the
    negative-energy (positron) states, it is the number kept for each kappa
    (the larger of the two kappa, for each l).
    By default (both), the two are joined with '+' when negative-energy states
    are present (e.g., '30spd20f+9spdf'), so they are always visible.
  */
  static std::string state_config(const std::vector<DiracSpinor> &orbs,
                                  Branch branch = Branch::both);
 
  /*!
    @brief Formats a list of numbers into the compressed states-string notation
 
    @details
    One number per l, starting at l=0, e.g., {30,30,20,20} => '30sp20df'.
    Zero entries are skipped.
  */
  static std::string config_string(const std::vector<int> &n_per_l);
 
  /*!
    @brief Constructs an exact H-like (pointlike Coulomb) DiracSpinor
 
    @details
    Relativistic Dirac-Coulomb bound orbital for nuclear charge @p zeff (see
    DiracHydrogen).
 
    @note alpha <= 0 gives the non-relativistic (Schrodinger) solution instead:
    f = P_nl, g = 0.
  */
  static DiracSpinor exactHlike(int n, int k, std::shared_ptr<const Grid> rgrid,
                                double zeff, double alpha = PhysConst::alpha,
                                double mass = 1.0);
 
  /*!
    @brief Constructs an exact (Dirac-Coulomb) continuum DiracSpinor
 
    @details
    Energy-normalised relativistic continuum orbital (en > 0; n set to 0).
    Requires FLINT for the relativistic case (see DiracContinuum::available).
 
    @note alpha <= 0 gives the non-relativistic solution instead (f = P_el,
    g = 0), which does not require FLINT.
    @warning Relativistic case without FLINT: prints a warning and returns
    NaN (assert fails if NDEBUG is not defined).
  */
  static DiracSpinor exactHlike_cntm(double en, int k,
                                     std::shared_ptr<const Grid> rgrid,
                                     double zeff,
                                     double alpha = PhysConst::alpha,
                                     double mass = 1.0);
 
  /*!
    @brief Constructs a basis of H-like (pointlike) DiracSpinors
 
    @param max_l  maximum l (both kappas calculated)
    @param num_ns number of principal quantum numbers per l (not max n)
    @param mass   lepton mass in units of m_e (default 1 = electron)
    @note alpha <= 0 gives the non-relativistic solutions (f = P_nl, g = 0)
  */
  static std::vector<DiracSpinor>
  HlikeBasis(int max_l, int num_ns, std::shared_ptr<const Grid> rgrid,
             double zeff, double alpha = PhysConst::alpha, double mass = 1.0);
 
  /*!
    @brief Constructs a basis of H-like (pointlike) DiracSpinors
 
    @param basis_string e.g., "6sp5d": max n for each kappa; all n down to
                        n_min = l+1 are included
    @param mass         lepton mass in units of m_e (default 1 = electron)
    @note alpha <= 0 gives the non-relativistic solutions (f = P_nl, g = 0)
  */
  static std::vector<DiracSpinor>
  HlikeBasis(const std::string &basis_string, std::shared_ptr<const Grid> rgrid,
             double zeff, double alpha = PhysConst::alpha, double mass = 1.0);
 
  //! Searches for {n,k} in list of orbitals, returns pointer (may be null)
  static const DiracSpinor *find(int n, int k,
                                 const std::vector<DiracSpinor> &orbs);
 
  //! Returns maximum (2j) found in {orbs}
  static int max_tj(const std::vector<DiracSpinor> &orbs);
  //! Returns maximum l found in {orbs}
  static int max_l(const std::vector<DiracSpinor> &orbs);
  //! Returns maximum n found in {orbs}
  static int max_n(const std::vector<DiracSpinor> &orbs);
  //! Returns maximum n found in {orbs}, for given kappa
  static int max_n(const std::vector<DiracSpinor> &orbs, int kappa);
  //! Returns maximum kappa_index found in {orbs}
  static std::size_t max_kindex(const std::vector<DiracSpinor> &orbs);
 
  //! Returns maximum Energy found in {orbs}
  static double max_En(const std::vector<DiracSpinor> &orbs);
  //! Returns minimum Energy found in {orbs}
  static double min_En(const std::vector<DiracSpinor> &orbs);
 
  /*!
    @brief Splits orbitals into two groups (i.e., core, excited) by energy
 
    @details
    - The first group contains orbitals with energy <= @p Fermi_energy
    - The second group contains orbitals with energy > @p Fermi_energy
    - The first group is also limited to have n >= @p n_min_core (i.e.,
      exclude deep core states)
    - The second group is limited to have n <= @p n_max_excited
    - Negative-energy states (if present) all go to the second group
      (@p n_max_excited is not applied to them). Set
      @p include_negative_energy = false to drop them entirely instead
  */
  static std::pair<std::vector<DiracSpinor>, std::vector<DiracSpinor>>
  split_by_energy(const std::vector<DiracSpinor> &orbitals, double Fermi_energy,
                  int n_min_core = 1, int n_max_excited = 9999,
                  bool include_negative_energy = true);
 
  /*!
    @brief Splits orbitals into two groups (i.e., core, excited)
 
    @details
    - The first group contains orbitals with {n,kappa} in the given @p core
    - The second group contains all of the rest
    - The first group is also limited to have n >= @p n_min_core (i.e.,
      exclude deep core states)
  */
  static std::pair<std::vector<DiracSpinor>, std::vector<DiracSpinor>>
  split_by_core(const std::vector<DiracSpinor> &orbitals,
                const std::vector<DiracSpinor> &core, int n_min_core = 1);
 
  /*!
    @brief Takes a subset of a basis (by copy), according to subset_string
 
    @details
    - Includes only states matching the @p subset_string (e.g., "6sp5d": max
      n for each kappa)
    - Negative-energy states are excluded, unless @p exclude_negative_energy
      is false: then all those with kappa in @p subset_string are included
      (the n limit is not applied to them)
  */
  static std::vector<DiracSpinor> subset(const std::vector<DiracSpinor> &basis,
                                         const std::string &subset_string,
                                         bool exclude_negative_energy = true);
};