text: import logging
  import time
  from typing import Dict, List, Optional, Union
  
  import numpy as np
  from pymatgen.core.structure import Structure
  from pymatgen.electronic_structure.core import Spin
  from pymatgen.util.coord import pbc_diff
  
  from amset.constants import numeric_types
  from amset.log import log_time_taken
  from amset.util import (
      array_from_buffer,
      create_shared_array,
      create_shared_dict_array,
      dict_array_from_buffer,
      get_progress_bar,
      groupby,
  )
  
  __author__ = "Alex Ganose"
  __maintainer__ = "Alex Ganose"
  __email__ = "aganose@lbl.gov"
  
  """
       6-------7
      /|      /|
     / |     / |
    4-------5  |
    |  2----|--3
    | /     | /
    |/      |/
    0-------1
  
   i: vec        neighbours
   0: O          1, 2, 4
   1: a          0, 3, 5
   2: b          0, 3, 6
   3: a + b      1, 2, 7
   4: c          0, 5, 6
   5: c + a      1, 4, 7
   6: c + b      2, 4, 7
   7: c + a + b  3, 5, 6
  """
  
  _main_diagonals = (
      (1, 1, 1),  # 0-7
      (-1, 1, 1),  # 1-6
      (1, -1, 1),  # 2-5
      (1, 1, -1),  # 3-4
  )
  
  _tetrahedron_vectors = np.array(
      [
          [0, 0, 0],
          [1, 0, 0],
          [0, 1, 0],
          [1, 1, 0],
          [0, 0, 1],
          [1, 0, 1],
          [0, 1, 1],
          [1, 1, 1],
      ]
  )
  
  logger = logging.getLogger(__name__)
  
  
  def get_main_diagonal(reciprocal_lattice: np.ndarray) -> int:
      # want a list of tetrahedra as (k1, k2, k3, k4); as per the Bloechl paper,
      # the order of the k-points is irrelevant and therefore they should be ordered
      # in increasing number
      # All tetrahedrons for a specific submesh will share one common diagonal.
      # To minimise interpolaiton distances, we should choose the shortest
      # diagonal
      diagonal_lengths = []
      for diagonal in _main_diagonals:
          length = np.linalg.norm(reciprocal_lattice @ diagonal)
          diagonal_lengths.append(length)
  
      return int(np.argmin(diagonal_lengths))
  
  
  def get_relative_tetrahedron_vertices(reciprocal_lattice):
      shortest_index = get_main_diagonal(reciprocal_lattice)
  
      if shortest_index == 0:
          pairs = ((1, 3), (1, 5), (2, 3), (2, 6), (4, 5), (4, 6))
          main = (0, 7)
  
      elif shortest_index == 1:
          pairs = ((0, 2), (0, 4), (2, 3), (3, 7), (4, 5), (5, 7))
          main = (1, 6)
  
      elif shortest_index == 2:
          pairs = ((0, 1), (0, 4), (1, 3), (3, 7), (4, 6), (6, 7))
          main = (2, 5)
  
      elif shortest_index == 3:
          pairs = ((0, 1), (0, 2), (1, 5), (2, 6), (5, 7), (6, 7))
          main = (3, 4)
  
      else:
          assert False
  
      tetras = np.sort([main + x for x in pairs])
      return _tetrahedron_vectors[tetras]
  
  
  def get_tetrahedra(
      reciprocal_lattice: np.ndarray,
      grid_address: np.ndarray,
      mesh: np.ndarray,
      grid_address_mapping,
  ):
      tetrahedron_vertices = get_relative_tetrahedron_vertices(reciprocal_lattice)
  
      grid_order = [1, mesh[0], mesh[0] * mesh[1]]
  
      all_grid_points = np.repeat(grid_address, [24] * len(grid_address), axis=0)
      all_vertices = np.tile(tetrahedron_vertices, (len(grid_address), 1, 1))
      points = all_grid_points.reshape(all_vertices.shape) + all_vertices
  
      # fancy magic from phonopy to get neighboring indices given relative coordinates
      tetrahedra = np.dot(points % mesh, grid_order)
  
      ir_tetrahedra_vertices = grid_address_mapping[tetrahedra]
      _, ir_tetrahedra_idx, ir_tetrahedra_to_full_idx, ir_weights = np.unique(
          np.sort(ir_tetrahedra_vertices),
          axis=0,
          return_index=True,
          return_inverse=True,
          return_counts=True,
      )
  
      return tetrahedra, ir_tetrahedra_idx, ir_tetrahedra_to_full_idx, ir_weights
  
  
  class TetrahedralBandStructure:
      def __init__(
          self,
          energies: Dict[Spin, np.ndarray],
          kpoints: np.ndarray,
          ir_kpoints_idx: np.ndarray,
          ir_kpoint_mapping: np.ndarray,
          ir_kpoint_weights: np.ndarray,
          tetrahedra: Dict[Spin, np.ndarray],
          ir_tetrahedra: Dict[Spin, np.ndarray],
          ir_tetrahedra_energies: Dict[Spin, np.ndarray],
          ir_tetrahedra_idx: np.ndarray,
          ir_tetrahedra_to_full_idx: np.ndarray,
          ir_tetrahedra_weights: np.ndarray,
          e21: Dict[Spin, np.ndarray],
          e31: Dict[Spin, np.ndarray],
          e41: Dict[Spin, np.ndarray],
          e32: Dict[Spin, np.ndarray],
          e42: Dict[Spin, np.ndarray],
          e43: Dict[Spin, np.ndarray],
          max_tetrahedra_energies: Dict[Spin, np.ndarray],
          min_tetrahedra_energies: Dict[Spin, np.ndarray],
          cross_section_weights: Dict[Spin, np.ndarray],
          tetrahedron_volume: float,
          weights_cache: Optional[Dict[Spin, np.ndarray]] = None,
          weights_mask_cache: Optional[Dict[Spin, np.ndarray]] = None,
          energies_cache: Optional[Dict[Spin, np.ndarray]] = None,
      ):
          self.energies = energies
          self.kpoints = kpoints
          self.ir_kpoints_idx = ir_kpoints_idx
          self.ir_kpoint_mapping = ir_kpoint_mapping
          self.ir_kpoint_weights = ir_kpoint_weights
          self.tetrahedra = tetrahedra
          self.ir_tetrahedra = ir_tetrahedra
          self.ir_tetrahedra_energies = ir_tetrahedra_energies
          self.ir_tetrahedra_idx = ir_tetrahedra_idx
          self.ir_tetrahedra_to_full_idx = ir_tetrahedra_to_full_idx
          self.ir_tetrahedra_weights = ir_tetrahedra_weights
          self.e21 = e21
          self.e31 = e31
          self.e41 = e41
          self.e32 = e32
          self.e42 = e42
          self.e43 = e43
          self.max_tetrahedra_energies = max_tetrahedra_energies
          self.min_tetrahedra_energies = min_tetrahedra_energies
          self.cross_section_weights = cross_section_weights
          self._tetrahedron_volume = tetrahedron_volume
          self._weights_cache = {} if weights_cache is None else weights_cache
          self._weights_mask_cache = (
              {} if weights_mask_cache is None else weights_mask_cache
          )
          self._energies_cache = {} if energies_cache is None else energies_cache
  
          self.grouped_ir_to_full = groupby(
              np.arange(len(ir_tetrahedra_to_full_idx)), ir_tetrahedra_to_full_idx
          )
          self._ir_weights_shape = {
              s: (len(energies[s]), len(ir_kpoints_idx)) for s in energies
          }
  
      def to_reference(self):
          energies_buffer, self.energies = create_shared_dict_array(
              self.energies, return_shared_data=True
          )
          kpoints_buffer, self.kpoints = create_shared_array(
              self.kpoints, return_shared_data=True
          )
          ir_kpoints_idx_buffer, self.ir_kpoints_idx = create_shared_array(
              self.ir_kpoints_idx, return_shared_data=True
          )
          ir_kpoint_mapping_buffer, self.ir_kpoint_mapping = create_shared_array(
              self.ir_kpoint_mapping, return_shared_data=True
          )
          ir_kpoint_weights_buffer, self.ir_kpoint_weights = create_shared_array(
              self.ir_kpoint_weights, return_shared_data=True
          )
          tetrahedra_buffer, self.tetrahedra = create_shared_dict_array(
              self.tetrahedra, return_shared_data=True
          )
          ir_tetrahedra_buffer, self.ir_tetrahedra = create_shared_dict_array(
              self.ir_tetrahedra, return_shared_data=True
          )
          (
              ir_tetrahedra_energies_buffer,
              self.ir_tetrahedra_energies,
          ) = create_shared_dict_array(
              self.ir_tetrahedra_energies, return_shared_data=True
          )
          ir_tetrahedra_idx_buffer, self.ir_tetrahedra_idx = create_shared_array(
              self.ir_tetrahedra_idx, return_shared_data=True
          )
          (
              ir_tetrahedra_to_full_idx_buffer,
              self.ir_tetrahedra_to_full_idx,
          ) = create_shared_array(self.ir_tetrahedra_to_full_idx, return_shared_data=True)
          ir_tetrahedra_weights_buffer, self.ir_tetrahedra_weights = create_shared_array(
              self.ir_tetrahedra_weights, return_shared_data=True
          )
          e21_buffer, self.e21 = create_shared_dict_array(
              self.e21, return_shared_data=True
          )
          e31_buffer, self.e31 = create_shared_dict_array(
              self.e31, return_shared_data=True
          )
          e41_buffer, self.e41 = create_shared_dict_array(
              self.e41, return_shared_data=True
          )
          e32_buffer, self.e32 = create_shared_dict_array(
              self.e32, return_shared_data=True
          )
          e42_buffer, self.e42 = create_shared_dict_array(
              self.e42, return_shared_data=True
          )
          e43_buffer, self.e43 = create_shared_dict_array(
              self.e43, return_shared_data=True
          )
          (
              max_tetrahedra_energies_buffer,
              self.max_tetrahedra_energies,
          ) = create_shared_dict_array(
              self.max_tetrahedra_energies, return_shared_data=True
          )
          (
              min_tetrahedra_energies_buffer,
              self.min_tetrahedra_energies,
          ) = create_shared_dict_array(
              self.min_tetrahedra_energies, return_shared_data=True
          )
          (
              cross_section_weights_buffer,
              self.cross_section_weights,
          ) = create_shared_dict_array(
              self.cross_section_weights, return_shared_data=True
          )
          weights_cache_buffer, self._weights_cache = create_shared_dict_array(
              self._weights_cache, return_shared_data=True
          )
          weights_mask_cache_buffer, self._weights_mask_cache = create_shared_dict_array(
              self._weights_mask_cache, return_shared_data=True
          )
          energies_cache_buffer, self._energies_cache = create_shared_dict_array(
              self._energies_cache, return_shared_data=True
          )
  
          return (
              energies_buffer,
              kpoints_buffer,
              ir_kpoints_idx_buffer,
              ir_kpoint_mapping_buffer,
              ir_kpoint_weights_buffer,
              tetrahedra_buffer,
              ir_tetrahedra_buffer,
              ir_tetrahedra_energies_buffer,
              ir_tetrahedra_idx_buffer,
              ir_tetrahedra_to_full_idx_buffer,
              ir_tetrahedra_weights_buffer,
              e21_buffer,
              e31_buffer,
              e41_buffer,
              e32_buffer,
              e42_buffer,
              e43_buffer,
              max_tetrahedra_energies_buffer,
              min_tetrahedra_energies_buffer,
              cross_section_weights_buffer,
              self._tetrahedron_volume,
              weights_cache_buffer,
              weights_mask_cache_buffer,
              energies_cache_buffer,
          )
  
      @classmethod
      def from_reference(
          cls,
          energies_buffer,
          kpoints_buffer,
          ir_kpoints_idx_buffer,
          ir_kpoint_mapping_buffer,
          ir_kpoint_weights_buffer,
          tetrahedra_buffer,
          ir_tetrahedra_buffer,
          ir_tetrahedra_energies_buffer,
          ir_tetrahedra_idx_buffer,
          ir_tetrahedra_to_full_idx_buffer,
          ir_tetrahedra_weights_buffer,
          e21_buffer,
          e31_buffer,
          e41_buffer,
          e32_buffer,
          e42_buffer,
          e43_buffer,
          max_tetrahedra_energies_buffer,
          min_tetrahedra_energies_buffer,
          cross_section_weights_buffer,
          tetrahedron_volume,
          weights_cache_buffer,
          weights_mask_cache_buffer,
          energies_cache_buffer,
      ):
          return cls(
              dict_array_from_buffer(energies_buffer),
              array_from_buffer(kpoints_buffer),
              array_from_buffer(ir_kpoints_idx_buffer),
              array_from_buffer(ir_kpoint_mapping_buffer),
              array_from_buffer(ir_kpoint_weights_buffer),
              dict_array_from_buffer(tetrahedra_buffer),
              dict_array_from_buffer(ir_tetrahedra_buffer),
              dict_array_from_buffer(ir_tetrahedra_energies_buffer),
              array_from_buffer(ir_tetrahedra_idx_buffer),
              array_from_buffer(ir_tetrahedra_to_full_idx_buffer),
              array_from_buffer(ir_tetrahedra_weights_buffer),
              dict_array_from_buffer(e21_buffer),
              dict_array_from_buffer(e31_buffer),
              dict_array_from_buffer(e41_buffer),
              dict_array_from_buffer(e32_buffer),
              dict_array_from_buffer(e42_buffer),
              dict_array_from_buffer(e43_buffer),
              dict_array_from_buffer(max_tetrahedra_energies_buffer),
              dict_array_from_buffer(min_tetrahedra_energies_buffer),
              dict_array_from_buffer(cross_section_weights_buffer),
              tetrahedron_volume,
              dict_array_from_buffer(weights_cache_buffer),
              dict_array_from_buffer(weights_mask_cache_buffer),
              dict_array_from_buffer(energies_cache_buffer),
          )
  
      @classmethod
      def from_data(
          cls,
          energies: Dict[Spin, np.ndarray],
          kpoints: np.ndarray,
          tetrahedra: np.ndarray,
          structure: Structure,
          ir_kpoints_idx: np.ndarray,
          ir_kpoint_mapping: np.ndarray,
          ir_tetrahedra_idx: Optional[np.ndarray] = None,
          ir_tetrahedra_to_full_idx: Optional[np.ndarray] = None,
          ir_tetrahedra_weights: Optional[np.ndarray] = None,
      ):
          logger.info("Initializing tetrahedron band structure")
          t0 = time.perf_counter()
  
          tparams = (ir_tetrahedra_idx, ir_tetrahedra_to_full_idx, ir_tetrahedra_weights)
          if len({x is None for x in tparams}) != 1:
              raise ValueError(
                  "Either all or none of ir_tetrahedra_idx, ir_tetrahedra_to_full_idx and"
                  " ir_tetrahedra_weights should be set."
              )
  
          if ir_tetrahedra_idx is None:
              ir_tetrahedra_idx = np.arange(len(kpoints))
              ir_tetrahedra_to_full_idx = np.ones_like(ir_tetrahedra_idx)
              ir_tetrahedra_weights = np.ones_like(ir_tetrahedra_idx)
  
          ir_tetrahedra_to_full_idx = ir_tetrahedra_to_full_idx
          ir_kpoints_idx = ir_kpoints_idx
          ir_kpoint_mapping = ir_kpoint_mapping
  
          _, ir_kpoint_weights = np.unique(ir_kpoint_mapping, return_counts=True)
  
          # need to keep track of full tetrahedra to recover full k-point indices
          # when calculating scattering rates (i.e., k-k' is symmetry inequivalent).
          full_tetrahedra, _ = process_tetrahedra(tetrahedra, energies)
  
          # store irreducible tetrahedra and use energies to calculate diffs and min/maxes
          ir_tetrahedra, ir_tetrahedra_energies = process_tetrahedra(
              tetrahedra[ir_tetrahedra_idx], energies
          )
  
          # the remaining properties are given for each irreducible tetrahedra
          (e21, e31, e41, e32, e42, e43) = get_tetrahedra_energy_diffs(
              ir_tetrahedra_energies
          )
  
          (
              max_tetrahedra_energies,
              min_tetrahedra_energies,
          ) = get_max_min_tetrahedra_energies(ir_tetrahedra_energies)
  
          cross_section_weights = get_tetrahedra_cross_section_weights(
              structure.lattice.reciprocal_lattice.matrix,
              kpoints,
              ir_tetrahedra,
              e21,
              e31,
              e41,
          )
  
          tetrahedron_volume = 1 / len(tetrahedra)
  
          log_time_taken(t0)
          return cls(
              energies,
              kpoints,
              ir_kpoints_idx,
              ir_kpoint_mapping,
              ir_kpoint_weights,
              full_tetrahedra,
              ir_tetrahedra,
              ir_tetrahedra_energies,
              ir_tetrahedra_idx,
              ir_tetrahedra_to_full_idx,
              ir_tetrahedra_weights,
              e21,
              e31,
              e41,
              e32,
              e42,
              e43,
              max_tetrahedra_energies,
              min_tetrahedra_energies,
              cross_section_weights,
              tetrahedron_volume,
          )
  
      def get_connected_kpoints(self, kpoint_idx: Union[int, List[int], np.ndarray]):
          """Given one or more k-point indices, get a list of all k-points that are in
          the same tetrahedra
  
          Args:
              kpoint_idx: One or mode k-point indices.
  
          Returns:
              A list of k-point indices that are in the same tetrahedra.
          """
          if isinstance(kpoint_idx, numeric_types):
              kpoint_idx = [kpoint_idx]
  
          tetrahedra = self.tetrahedra[Spin.up][0]
          return np.unique(tetrahedra[np.isin(tetrahedra, kpoint_idx).any(axis=1)])
  
      def get_intersecting_tetrahedra(self, spin, energy, band_idx=None):
  
          max_energies = self.max_tetrahedra_energies[spin]
          min_energies = self.min_tetrahedra_energies[spin]
  
          if band_idx is not None:
              mask = np.full_like(max_energies, False, dtype=bool)
              mask[band_idx] = True
  
              return (min_energies < energy) & (max_energies > energy) & mask
  
          else:
              return (min_energies < energy) & (max_energies > energy)
  
      def get_tetrahedra_density_of_states(
          self,
          spin,
          energy,
          return_contributions=False,
          symmetry_reduce=True,
          band_idx=None,
      ):
          tetrahedra_mask = self.get_intersecting_tetrahedra(
              spin, energy, band_idx=band_idx
          )
  
          if not np.any(tetrahedra_mask):
              if return_contributions:
                  return [], [], [], []
              else:
                  return []
  
          energies = self.ir_tetrahedra_energies[spin][tetrahedra_mask]
          e21 = self.e21[spin][tetrahedra_mask]
          e31 = self.e31[spin][tetrahedra_mask]
          e41 = self.e41[spin][tetrahedra_mask]
          e32 = self.e32[spin][tetrahedra_mask]
          e42 = self.e42[spin][tetrahedra_mask]
          e43 = self.e43[spin][tetrahedra_mask]
          cs_weights = self.cross_section_weights[spin][tetrahedra_mask]
  
          cond_a_mask = (energies[:, 0] < energy) & (energy < energies[:, 1])
          cond_b_mask = (energies[:, 1] <= energy) & (energy < energies[:, 2])
          cond_c_mask = (energies[:, 2] <= energy) & (energy < energies[:, 3])
  
          ee1 = energy - energies[:, 0]
          ee2 = energy - energies[:, 1]
          e4e = energies[:, 3] - energy
  
          tetrahedra_dos = np.zeros(len(energies))
  
          tetrahedra_dos[cond_a_mask] = _get_density_of_states_a(
              ee1[cond_a_mask], e21[cond_a_mask], e31[cond_a_mask], e41[cond_a_mask]
          )
  
          tetrahedra_dos[cond_b_mask] = _get_density_of_states_b(
              ee2[cond_b_mask],
              e21[cond_b_mask],
              e31[cond_b_mask],
              e41[cond_b_mask],
              e32[cond_b_mask],
              e42[cond_b_mask],
          )
  
          tetrahedra_dos[cond_c_mask] = _get_density_of_states_c(
              e4e[cond_c_mask], e41[cond_c_mask], e42[cond_c_mask], e43[cond_c_mask]
          )
  
          tetrahedra_dos *= self._tetrahedron_volume
  
          band_idx, tetrahedra_idx = np.where(tetrahedra_mask)
          tetrahedra_weights = self.ir_tetrahedra_weights[tetrahedra_idx]
  
          if symmetry_reduce:
              tetrahedra_dos *= tetrahedra_weights
  
          else:
              # transform the mask to the full BZ
              band_idx = np.repeat(band_idx, tetrahedra_weights)
              tetrahedra_grouped = self.grouped_ir_to_full[tetrahedra_idx]
              tetrahedra_idx = np.concatenate(tetrahedra_grouped)
              tetrahedra_mask = (band_idx, tetrahedra_idx)
  
              # tetrahedra_mask = tetrahedra_mask[:, self.ir_tetrahedra_to_full_idx]
              # get a mask to expand the remaining properties with shape (ntetra,)
              # to the full BZ
              expand_tetrahedra = np.repeat(np.arange(len(ee1)), tetrahedra_weights)
              tetrahedra_dos = tetrahedra_dos[expand_tetrahedra]
              cond_a_mask = cond_a_mask[expand_tetrahedra]
              cond_b_mask = cond_b_mask[expand_tetrahedra]
              cond_c_mask = cond_c_mask[expand_tetrahedra]
              ee1 = ee1[expand_tetrahedra]
              ee2 = ee2[expand_tetrahedra]
              e4e = e4e[expand_tetrahedra]
              e21 = e21[expand_tetrahedra]
              e31 = e31[expand_tetrahedra]
              e41 = e41[expand_tetrahedra]
              e32 = e32[expand_tetrahedra]
              e42 = e42[expand_tetrahedra]
              e43 = e43[expand_tetrahedra]
              cs_weights = cs_weights[expand_tetrahedra]
  
          if return_contributions:
              frac_21 = ee1 / e21
              frac_31 = ee1 / e31
              frac_41 = ee1 / e41
              frac_32 = ee2 / e32
              frac_42 = ee2 / e42
              frac_c_41 = e4e / e41
              frac_c_42 = e4e / e42
              frac_c_43 = e4e / e43
  
              frac_21[(frac_21 == np.nan) | (frac_21 == np.inf)] = 1
              frac_31[(frac_31 == np.nan) | (frac_31 == np.inf)] = 1
              frac_41[(frac_41 == np.nan) | (frac_41 == np.inf)] = 1
              frac_32[(frac_32 == np.nan) | (frac_32 == np.inf)] = 1
              frac_42[(frac_42 == np.nan) | (frac_42 == np.inf)] = 1
              frac_c_41[(frac_c_41 == np.nan) | (frac_c_41 == np.inf)] = 1
              frac_c_42[(frac_c_42 == np.nan) | (frac_c_42 == np.inf)] = 1
              frac_c_43[(frac_c_43 == np.nan) | (frac_c_43 == np.inf)] = 1
  
              contributions = (
                  cond_a_mask,
                  cond_b_mask,
                  cond_c_mask,
                  frac_21,
                  frac_31,
                  frac_41,
                  frac_32,
                  frac_42,
                  frac_c_41,
                  frac_c_42,
                  frac_c_43,
              )
  
              return tetrahedra_dos, tetrahedra_mask, cs_weights, contributions
  
          else:
              return tetrahedra_dos
  
      def get_density_of_states(
          self,
          energies=None,
          integrand=None,
          sum_spins=False,
          band_idx=None,
          use_cached_weights=False,
          progress_bar=False,
      ):
          if energies is None:
              from amset.constants import defaults, ev_to_hartree
  
              min_e = np.min([np.min(e) for e in self.energies.values()])
              max_e = np.max([np.max(e) for e in self.energies.values()])
              energies = np.arange(min_e, max_e, defaults["dos_estep"] * ev_to_hartree)
  
          dos = {}
          for spin in self.energies.keys():
              if isinstance(integrand, dict):
                  # integrand given for each spin channel
                  spin_integrand = integrand[spin]
              else:
                  spin_integrand = integrand
  
              if isinstance(band_idx, dict):
                  # band indices given for each spin channel
                  spin_band_idx = band_idx[spin]
              else:
                  spin_band_idx = band_idx
  
              if spin_integrand is not None:
                  if spin_integrand.shape[:2] != self.energies[spin].shape:
                      raise ValueError(
                          "Unexpected integrand shape, should be (nbands, nkpoints, ...)"
                      )
  
                  nbands = len(spin_integrand)
                  integrand_shape = spin_integrand.shape[2:]
                  n_ir_kpoints = len(self.ir_kpoints_idx)
                  new_integrand = np.zeros((nbands, n_ir_kpoints) + integrand_shape)
  
                  flat_k = np.tile(self.ir_kpoint_mapping, nbands)
                  flat_b = np.repeat(np.arange(nbands), len(self.ir_kpoint_mapping))
                  flat_integrand = spin_integrand.reshape((-1,) + integrand_shape)
                  # flat_integrand = spin_integrand.reshape(-1, 3, 3)
  
                  # sum integrand at all symmetry equivalent points, new_integrand
                  # has shape (nbands, n_ir_kpoints)
                  np.add.at(new_integrand, (flat_b, flat_k), flat_integrand)
                  spin_integrand = new_integrand
  
              emesh, dos[spin] = self.get_spin_density_of_states(
                  spin,
                  energies,
                  integrand=spin_integrand,
                  band_idx=spin_band_idx,
                  use_cached_weights=use_cached_weights,
                  progress_bar=progress_bar,
              )
  
          if sum_spins:
              if Spin.down in dos:
                  dos = dos[Spin.up] + dos[Spin.down]
              else:
                  dos = dos[Spin.up]
  
          return energies, dos
  
      def get_spin_density_of_states(
          self,
          spin,
          energies,
          integrand=None,
          band_idx=None,
          use_cached_weights=False,
          progress_bar=False,
      ):
          # integrand should have the shape (nbands, n_ir_kpts, ...)
          # the integrand should have been summed at all equivalent k-points
          # TODO: add support for variable shaped integrands
          if integrand is None:
              dos = np.zeros_like(energies)
          else:
              integrand_shape = integrand.shape[2:]
              dos = np.zeros((len(energies),) + integrand_shape)
  
          if use_cached_weights:
              if self._weights_cache is None:
                  raise ValueError("No integrand have been cached")
  
              all_weights = self._weights_cache[spin]
              all_weights_mask = self._weights_mask_cache[spin]
              energies = self._energies_cache[spin]
          else:
              all_weights = []
              all_weights_mask = []
  
          nbands = len(self.energies[spin])
          kpoint_multiplicity = np.tile(self.ir_kpoint_weights, (nbands, 1))
  
          if band_idx is not None and integrand is not None:
              integrand = integrand[band_idx]
  
          if band_idx is not None and integrand is None:
              kpoint_multiplicity = kpoint_multiplicity[band_idx]
  
          energies_iter = list(enumerate(energies))
          if progress_bar:
              energies_iter = get_progress_bar(iterable=energies_iter, desc="DOS")
  
          for i, energy in energies_iter:
              if use_cached_weights:
                  weights = all_weights[i]
                  weights_mask = all_weights_mask[i]
              else:
                  weights = self.get_energy_dependent_integration_weights(spin, energy)
                  weights_mask = weights != 0
                  all_weights.append(weights)
                  all_weights_mask.append(weights_mask)
  
              if band_idx is not None:
                  weights = weights[band_idx]
                  weights_mask = weights_mask[band_idx]
  
              if integrand is None:
                  dos[i] = np.sum(
                      weights[weights_mask] * kpoint_multiplicity[weights_mask]
                  )
  
              else:
                  # expand weights to match the dimensions of the integrand
                  expand_axis = [1 + i for i in range(len(integrand.shape[2:]))]
                  expand_weights = np.expand_dims(weights[weights_mask], axis=expand_axis)
  
                  # don't need to include the k-point multiplicity as this is included by
                  # pre-summing the integrand at symmetry equivalent points
                  dos[i] = np.sum(expand_weights * integrand[weights_mask], axis=0)
  
          if not use_cached_weights:
              self._weights_cache[spin] = np.array(all_weights)
              self._weights_mask_cache[spin] = np.array(all_weights_mask)
              self._energies_cache[spin] = energies
  
          return energies, np.asarray(dos)
  
      def get_energy_dependent_integration_weights(self, spin, energy):
          integration_weights = np.zeros(self._ir_weights_shape[spin])
          tetrahedra_mask = self.get_intersecting_tetrahedra(spin, energy)
  
          if not np.any(tetrahedra_mask):
              return integration_weights
  
          energies = self.ir_tetrahedra_energies[spin][tetrahedra_mask]
          e21 = self.e21[spin][tetrahedra_mask]
          e31 = self.e31[spin][tetrahedra_mask]
          e41 = self.e41[spin][tetrahedra_mask]
          e32 = self.e32[spin][tetrahedra_mask]
          e42 = self.e42[spin][tetrahedra_mask]
          e43 = self.e43[spin][tetrahedra_mask]
  
          cond_a_mask = (energies[:, 0] < energy) & (energy < energies[:, 1])
          cond_b_mask = (energies[:, 1] <= energy) & (energy < energies[:, 2])
          cond_c_mask = (energies[:, 2] <= energy) & (energy < energies[:, 3])
  
          ee1 = energy - energies[:, 0]
          ee2 = energy - energies[:, 1]
          ee3 = energy - energies[:, 2]
          e2e = energies[:, 1] - energy
          e3e = energies[:, 2] - energy
          e4e = energies[:, 3] - energy
  
          kpoints_idx = self.ir_tetrahedra[spin][tetrahedra_mask]
          ir_kpoints_idx = self.ir_kpoint_mapping[kpoints_idx]
  
          # calculate the integrand for each vertices
          vert_weights = np.zeros_like(energies)
          vert_weights[cond_a_mask] = _get_energy_dependent_weight_a(
              ee1[cond_a_mask],
              e2e[cond_a_mask],
              e3e[cond_a_mask],
              e4e[cond_a_mask],
              e21[cond_a_mask],
              e31[cond_a_mask],
              e41[cond_a_mask],
          )
  
          vert_weights[cond_b_mask] = _get_energy_dependent_weight_b(
              ee1[cond_b_mask],
              ee2[cond_b_mask],
              e3e[cond_b_mask],
              e4e[cond_b_mask],
              e31[cond_b_mask],
              e41[cond_b_mask],
              e32[cond_b_mask],
              e42[cond_b_mask],
          )
  
          vert_weights[cond_c_mask] = _get_energy_dependent_weight_c(
              ee1[cond_c_mask],
              ee2[cond_c_mask],
              ee3[cond_c_mask],
              e4e[cond_c_mask],
              e41[cond_c_mask],
              e42[cond_c_mask],
              e43[cond_c_mask],
          )
  
          # finally, get the integrand for each ir_kpoint by summing over all
          # tetrahedra and multiplying by the tetrahedra multiplicity and
          # tetrahedra weight; Finally, divide by the k-point multiplicity
          # to get the final weight
          band_idx, tetrahedra_idx = np.where(tetrahedra_mask)
  
          # include tetrahedra multiplicity
          vert_weights *= self.ir_tetrahedra_weights[tetrahedra_idx][:, None]
  
          flat_ir_kpoints = np.ravel(ir_kpoints_idx)
          flat_ir_weights = np.ravel(vert_weights)
          flat_bands = np.repeat(band_idx, 4)
  
          # sum integrand, note this sums in place and is insanely fast
          np.add.at(integration_weights, (flat_bands, flat_ir_kpoints), flat_ir_weights)
          integration_weights *= (
              self._tetrahedron_volume / self.ir_kpoint_weights[None, :]
          )
  
          return integration_weights
  
      def get_masks(self, spin, tetrahedra_mask):
          # mask needs to be generated with symmetry_reduce=False
          band_idxs = tetrahedra_mask[0]
  
          # property_mask can be used to get the values of a band and k-dependent
          # property at all tetrahedra vertices for each tetrahedron selected by the
          # mask for example: energies[property_mask]
          property_mask = (
              np.repeat(band_idxs[:, None], 4, axis=1),
              self.tetrahedra[spin][tetrahedra_mask],
          )
  
          # band_kpoint_mask can be used to get the inequivalent band, k-point
          # combinations that are required to get the properties for all tetrahedra
          # vertices. Can be used to avoid calculation of band and k-point dependent
          # properties at all bands k-points and k-points (instead you only calculate
          # the properties at the bands and k-points involved in the tetrahedra)
          # and therefore reduce computational expense
          band_kpoint_mask = np.full(self.energies[spin].shape, False)
          band_kpoint_mask[property_mask] = True
  
          # kpoint_mask is the k-point contribution to band_kpoint_mask which allows
          # mixing of band and k-point dependent properties and just k-dependent
          # properties
          band_mask, kpoint_mask = np.where(band_kpoint_mask)
  
          return property_mask, band_kpoint_mask, band_mask, kpoint_mask
  
  
  def _get_density_of_states_a(ee1, e21, e31, e41):
      return 3 * ee1**2 / (e21 * e31 * e41)
  
  
  def _get_density_of_states_b(ee2, e21, e31, e41, e32, e42):
      return 3 * (e21 + 2 * ee2 - (e42 + e31) * ee2**2 / (e32 * e42)) / (e31 * e41)
  
  
  def _get_density_of_states_c(e4e, e41, e42, e43):
      return 3 * e4e**2 / (e41 * e42 * e43)
  
  
  def _get_energy_dependent_weight_a(ee1, e2e, e3e, e4e, e21, e31, e41):
      c = ee1**2 / (e21 * e31 * e41)
      i1 = c * (e2e / e21 + e3e / e31 + e4e / e41)
      i2 = c * (ee1 / e21)
      i3 = c * (ee1 / e31)
      i4 = c * (ee1 / e41)
      return np.stack([i1, i2, i3, i4], axis=1)
  
  
  def _get_energy_dependent_weight_b(ee1, ee2, e3e, e4e, e31, e41, e32, e42):
      c = (ee1 * e4e) / (e31 * e41 * e42)
      x = e3e / e31
      y = e4e / e42
      z = ee2 / (e32 * e42)
      zx = z * x
      k = ee1 / e31
      n = ee2 / e42
  
      i1 = c * (x + e4e / e41) + z * x**2
      i2 = c * y + zx * (e3e / e32 + y)
      i3 = c * k + zx * (k + ee2 / e32)
      i4 = c * (ee1 / e41 + n) + zx * n
      return np.stack([i1, i2, i3, i4], axis=1)
  
  
  def _get_energy_dependent_weight_c(ee1, ee2, ee3, e4e, e41, e42, e43):
      c = e4e**2 / (e41 * e42 * e43)
      i1 = c * e4e / e41
      i2 = c * e4e / e42
      i3 = c * e4e / e43
      i4 = c * (ee1 / e41 + ee2 / e42 + ee3 / e43)
      return np.stack([i1, i2, i3, i4], axis=1)
  
  
  def get_cross_section_values(
      property_values,
      cond_a_mask,
      cond_b_mask,
      cond_c_mask,
      frac_21,
      frac_31,
      frac_41,
      frac_32,
      frac_42,
      frac_c_41,
      frac_c_42,
      frac_c_43,
      average=True,
  ):
      # property_values is given for each tetrahedra
      # property can be scalar or vector
      p21 = property_values[:, 1] - property_values[:, 0]
      p31 = property_values[:, 2] - property_values[:, 0]
      p41 = property_values[:, 3] - property_values[:, 0]
      p32 = property_values[:, 2] - property_values[:, 1]
      p42 = property_values[:, 3] - property_values[:, 1]
      p43 = property_values[:, 3] - property_values[:, 2]
  
      prop_shape = property_values.shape[2:]
  
      if average:
          # remove the 4 vertices from shape
          value_shape = property_values.shape[:1] + prop_shape
      else:
          # shape is (ntetrahedra, 4, prop_shape)
          # the 4 is to account for the intersections of the cross section and tetrahedra
          # the number of intersections can be either 3 or 4;
          # for cross sections with 3 interesections, the final intersection will be
          # an array of 0's with the shape of prop_shape
          value_shape = property_values.shape[:1] + (4,) + prop_shape
  
      values = np.zeros(value_shape)
  
      if len(prop_shape) != 0:
          # property is multi-dimensional
          new_shape = (len(values),) + tuple([1] * len(prop_shape))
          frac_21 = frac_21.reshape(new_shape)
          frac_31 = frac_31.reshape(new_shape)
          frac_41 = frac_41.reshape(new_shape)
          frac_32 = frac_32.reshape(new_shape)
          frac_42 = frac_42.reshape(new_shape)
          frac_c_41 = frac_c_41.reshape(new_shape)
          frac_c_42 = frac_c_42.reshape(new_shape)
          frac_c_43 = frac_c_43.reshape(new_shape)
  
      values[cond_a_mask] = _get_cross_section_values_a(
          property_values[cond_a_mask],
          p21[cond_a_mask],
          p31[cond_a_mask],
          p41[cond_a_mask],
          frac_21[cond_a_mask],
          frac_31[cond_a_mask],
          frac_41[cond_a_mask],
          average=average,
      )
  
      values[cond_b_mask] = _get_cross_section_values_b(
          property_values[cond_b_mask],
          p31[cond_b_mask],
          p41[cond_b_mask],
          p32[cond_b_mask],
          p42[cond_b_mask],
          frac_31[cond_b_mask],
          frac_41[cond_b_mask],
          frac_32[cond_b_mask],
          frac_42[cond_b_mask],
          average=average,
      )
  
      values[cond_c_mask] = _get_cross_section_values_c(
          property_values[cond_c_mask],
          p41[cond_c_mask],
          p42[cond_c_mask],
          p43[cond_c_mask],
          frac_c_41[cond_c_mask],
          frac_c_42[cond_c_mask],
          frac_c_43[cond_c_mask],
          average=average,
      )
  
      return values
  
  
  def _get_cross_section_values_a(
      property_values, p21, p31, p41, frac_21, frac_31, frac_41, average=False
  ):
      # these are the property values at the intersection of the cross section section
      # with the tetrahedron edges
      interp_21 = p21 * frac_21 + property_values[:, 0]
      interp_31 = p31 * frac_31 + property_values[:, 0]
      interp_41 = p41 * frac_41 + property_values[:, 0]
  
      if average:
          return (interp_21 + interp_31 + interp_41) / 3
  
      else:
          return np.stack(
              [interp_21, interp_31, interp_41, np.zeros_like(interp_21)], axis=1
          )
  
  
  def _get_cross_section_values_b(
      property_values,
      p31,
      p41,
      p32,
      p42,
      frac_31,
      frac_41,
      frac_32,
      frac_42,
      average=False,
  ):
      # these are the property values at the intersection of the cross section section
      # with the tetrahedron edges
      interp_31 = p31 * frac_31 + property_values[:, 0]
      interp_41 = p41 * frac_41 + property_values[:, 0]
      interp_32 = p32 * frac_32 + property_values[:, 1]
      interp_42 = p42 * frac_42 + property_values[:, 1]
  
      if average:
          return (interp_31 + interp_41 + interp_32 + interp_42) / 4
  
      else:
          # the order of these matters for calculating the cross sectional
          # area. the first and last points are at opposite ends of the quadrilateral
          return np.stack([interp_31, interp_32, interp_41, interp_42], axis=1)
  
  
  def _get_cross_section_values_c(
      property_values, p41, p42, p43, frac_41, frac_42, frac_43, average=False
  ):
      # these are the property values at the intersection of the cross section
      # with the tetrahedron edges
      interp_41 = property_values[:, 3] - p41 * frac_41
      interp_42 = property_values[:, 3] - p42 * frac_42
      interp_43 = property_values[:, 3] - p43 * frac_43
  
      if average:
          return (interp_41 + interp_42 + interp_43) / 3
  
      else:
          return np.stack(
              [interp_41, interp_42, interp_43, np.zeros_like(interp_41)], axis=1
          )
  
  
  def get_projected_intersections(intersections):
      # project the intersections coordinates onto a two dimensional plane
  
      # take two of the triangle/quadrilateral sides and calculate the unit vectors
      # the first of these vectors will be a basis vector
      axis_a = intersections[:, 1] - intersections[:, 0]
      axis_a /= np.linalg.norm(axis_a, axis=1)[:, None]
  
      other_side = intersections[:, 2] - intersections[:, 0]
  
      # find the vector normal to the plane of the triangle/quadrilateral, this will be a
      # basis vector, however all points will have the same magnitude in this direction
      axis_c = np.cross(axis_a, other_side)
      axis_c /= np.linalg.norm(axis_c, axis=1)[:, None]
  
      # find the final basis vector by finding the normal to the a and c basis vectors
      # no need to normalise as it is the cross product of two orthogonal unit vectors
      axis_b = np.cross(axis_a, axis_c)
  
      # define a transformation matrix to transform the points to our new basis
      basis = np.stack([axis_a, axis_b, axis_c], axis=2)
      transform = np.linalg.inv(basis)
  
      # finally transform the intersection coordinates. This syntax is equivalent to
      # np.dot(transform[0], intersections[0][0]) for all intersections in all triangles/
      # quadrilaterals.
      return np.einsum("ikj,ilj->ilk", transform, intersections), basis
  
  
  def process_tetrahedra(tetrahedra, energies):
      all_tetrahedra = {}
      all_tetrahedra_energies = {}
  
      for spin, spin_energies in energies.items():
          data_shape = (len(spin_energies),) + tetrahedra.shape
          spin_tetrahedra = np.zeros(data_shape, dtype=int)
          spin_tetrahedra_energies = np.zeros(data_shape)
  
          for band_idx, band_energies in enumerate(spin_energies):
              band_tetrahedra_energies = band_energies[tetrahedra]
  
              sort_idx = np.argsort(band_tetrahedra_energies, axis=1)
              spin_tetrahedra_energies[band_idx, ...] = np.take_along_axis(
                  band_tetrahedra_energies, sort_idx, axis=1
              )
              spin_tetrahedra[band_idx, ...] = np.take_along_axis(
                  tetrahedra, sort_idx, axis=1
              )
  
          all_tetrahedra[spin] = spin_tetrahedra
          all_tetrahedra_energies[spin] = spin_tetrahedra_energies
  
      return all_tetrahedra, all_tetrahedra_energies
  
  
  def get_tetrahedra_energy_diffs(tetrahedra_energies):
      e21 = {}
      e31 = {}
      e41 = {}
      e32 = {}
      e42 = {}
      e43 = {}
  
      for spin, s_tetrahedra_energies in tetrahedra_energies.items():
          # each energy difference has the shape nbands, ntetrahedra
          e21[spin] = s_tetrahedra_energies[:, :, 1] - s_tetrahedra_energies[:, :, 0]
          e31[spin] = s_tetrahedra_energies[:, :, 2] - s_tetrahedra_energies[:, :, 0]
          e41[spin] = s_tetrahedra_energies[:, :, 3] - s_tetrahedra_energies[:, :, 0]
          e32[spin] = s_tetrahedra_energies[:, :, 2] - s_tetrahedra_energies[:, :, 1]
          e42[spin] = s_tetrahedra_energies[:, :, 3] - s_tetrahedra_energies[:, :, 1]
          e43[spin] = s_tetrahedra_energies[:, :, 3] - s_tetrahedra_energies[:, :, 2]
  
      return e21, e31, e41, e32, e42, e43
  
  
  def get_max_min_tetrahedra_energies(tetrahedra_energies):
      max_tetrahedra_energies = {}
      min_tetrahedra_energies = {}
  
      for spin, s_tetrahedra_energies in tetrahedra_energies.items():
          max_tetrahedra_energies[spin] = np.max(s_tetrahedra_energies, axis=2)
          min_tetrahedra_energies[spin] = np.min(s_tetrahedra_energies, axis=2)
  
      return max_tetrahedra_energies, min_tetrahedra_energies
  
  
  def get_tetrahedra_cross_section_weights(
      reciprocal_lattice, kpoints, tetrahedra, e21, e31, e41
  ):
      # weight (b) defined by equation 3.4 in https://doi.org/10.1002/pssb.2220540211
      # this weight is not the Bloechl integrand but a scaling needed to obtain the
      # DOS directly from the tetrahedron cross section
      cross_section_weights = {}
  
      # volume is 6 * the volume of one tetrahedron
      volume = np.linalg.det(reciprocal_lattice) / len(kpoints)
      for spin, s_tetrahedra in tetrahedra.items():
          tetrahedra_kpoints = kpoints[s_tetrahedra]
  
          k1 = pbc_diff(tetrahedra_kpoints[:, :, 1], tetrahedra_kpoints[:, :, 0])
          k2 = pbc_diff(tetrahedra_kpoints[:, :, 2], tetrahedra_kpoints[:, :, 0])
          k3 = pbc_diff(tetrahedra_kpoints[:, :, 3], tetrahedra_kpoints[:, :, 0])
  
          k1_cart = np.dot(k1, reciprocal_lattice)
          k2_cart = np.dot(k2, reciprocal_lattice)
          k3_cart = np.dot(k3, reciprocal_lattice)
  
          contragradient = np.stack(
              [
                  np.cross(k2_cart, k3_cart) / volume,
                  np.cross(k3_cart, k1_cart) / volume,
                  np.cross(k1_cart, k2_cart) / volume,
              ],
              axis=2,
          )
  
          energies = np.stack([e21[spin], e31[spin], e41[spin]], axis=2)
          b_vector = np.sum(contragradient * energies[..., None], axis=2)
  
          cross_section_weights[spin] = 1 / np.linalg.norm(b_vector, axis=2)
  
      return cross_section_weights
  