code: def get_cov_matrix(h2, rg, psd_rg=False):
      r"""Creates covariance matrix for simulating correlated SNP effects.
  
      Given a list of heritabilities and a list of genetic correlations, :func:`.get_cov_matrix`
      constructs the covariance matrix necessary to draw from a multivariate normal
      distribution to generate correlated SNP effects.
  
      Examples
      --------
      Suppose we have three traits enumerated as trait 1, trait 2, and trait 3.
      Each trait has a heritability: :math:`h^2_1`,:math:`h^2_2`,:math:`h^2_3`
      Traits have the following genetic correlations: :math:`r_{g, 12}`,:math:`r_{g, 13}`, :math:`r_{g, 23}`
      The ordering of indices in the subscript is arbitrary (e.g. :math:`r_{g, 12}` = :math:`r_{g, 21}`)
      as both values are the genetic correlation between trait 1 and trait 2.
      We can calculate :math:`\rho_{g,ab}`, the genetic covariance between two traits :math:`a` and :math:`b`,
      as :math:`\rho_{g,ab}=r_{g,ab}\sqrt{h^2_a\cdot h^2_b}`. The covariance matrix is thus:
  
      .. math::
  
          \begin{pmatrix}
          h^2_1                            & r_{g, 12}\sqrt{h^2_1\cdot h^2_2}  & r_{g, 13}\sqrt{h^2_1\cdot h^2_3} \\
          r_{g, 12}\sqrt{h^2_1\cdot h^2_2} & h^2_2                             & r_{g, 23}\sqrt{h^2_2\cdot h^2_3} \\
          r_{g, 13}\sqrt{h^2_1\cdot h^2_3} & r_{g, 23}*\sqrt{h^2_2\cdot h^2_3} & h^2_3
          \end{pmatrix}
  
      Now suppose we have four traits with the following heritabilities (:math:`h^2`): 0.1, 0.3, 0.2, 0.6.
      That is, trait 1 has an :math:`h^2` of 0.1, trait 2 has an :math:`h^2` of 0.3 and so on.
      Suppose the genetic correlations (:math:`r_g`) between traits are the following:
      trait 1 & trait 2 :math:`r_g` = 0.4
      trait 1 & trait 3 :math:`r_g` = 0.3
      trait 1 & trait 4 :math:`r_g` = 0.1
      trait 2 & trait 3 :math:`r_g` = 0.2
      trait 2 & trait 4 :math:`r_g` = 0.15
      trait 3 & trait 4 :math:`r_g` = 0.6
      To obtain the covariance matrix corresponding to this scenario :math:`h^2` values are
      ordered according to user specification and :math:`r_g` values are ordered by the
      order in which the corresponding genetic covariance terms will appear in the
      covariance matrix, reading lines in the upper triangular matrix from left to
      right, top to bottom (read first row left to right, read second row left to
      right, etc.), exluding the diagonal.
  
      >>> cov_matrix, rg = get_cov_matrix(h2=[0.1, 0.3, 0.2, 0.6], rg=[0.4, 0.3, 0.1, 0.2, 0.15, 0.6])
      >>> cov_matrix
      array([[0.1       , 0.06928203, 0.04242641, 0.0244949 ],
             [0.06928203, 0.3       , 0.04898979, 0.06363961],
             [0.04242641, 0.04898979, 0.2       , 0.2078461 ],
             [0.0244949 , 0.06363961, 0.2078461 , 0.6       ]])
  
      The diagonal corresponds directly to `h2`, the list of h2 values for all traits.
      In the upper triangular matrix, excluding the diagonal, the entry :math:`(a, b)`,
      where :math:`a` and :math:`b` are in :math:`{1,2,3,4}`, is the genetic covariance
      (:math:`\rho_g`) between traits :math:`a` and :math:`b`.
      Genetic covariance is calculated as :math:`\rho_g= r_g*\sqrt{h^2_a*h^2_b}`
      where :math:`r_g` is the genetic correlation between traits :math:`a` and
      :math:`b` and :math:`h^2_a` and :math:`h^2_b` are heritabilities corresponding
      to traits :math:`a` and :math:`b`.
  
      Notes
      -----
      Covariance matrix is not scaled by number of SNPs.
  
      If the h2 and rg parameters passed cause the resulting covariance matrix to
      not be positive semidefinite, this may cause the distribution of SNP effects
      generated by this covariance matrix to not have the properties specified by
      the h2 and rg parameters. To automatically adjust rg values so that the
      covariance matrix is positive semidefinite, set `psd_rg` = True.
  
      Parameters
      ----------
      h2 : :obj:`list` or :class:`numpy.ndarray`
          :math:`h^2` values for traits. :math:`h^2` values in list should be
          ordered by their order in the diagonal of the covariance array, reading
          from top left to bottom right.
      rg : :obj:`list` or :class:`numpy.ndarray`
          :math:`r_g` values for traits. :math:`r_g` values should be ordered in
          the order they appear in the upper triangle of the covariance matrix,
          from left to right, top to bottom.
      psd_rg :  :obj:`bool`
          Whether to automatically adjust rg values to get a positive semi-definite
          covariance matrix, which ensures that SNP effects simulated with that
          covariance matrix have the desired variance and correlation properties
          specified by the h2 and rg parameters.
  
      Returns
      -------
      cov_matrix : :class:`numpy.ndarray`
          Covariance matrix calculated using `h2` and (possibly altered) `rg` values.
      rg : :obj:`list`
          Genetic correlation between traits, possibly altered from input `rg` if
          covariance matrix was not positive semi-definite.
      """
      assert all(
          x >= 0 and x <= 1 for x in h2
      ), "h2 values must be between 0 and 1"
      assert all(
          x >= -1 and x <= 1 for x in rg
      ), "rg values must be between -1 and 1"
      rg = np.asarray(rg) if isinstance(rg, list) else rg
      n_rg = len(rg)
      n_h2 = len(h2)
      # expected number of rg values, given number of traits
      exp_n_rg = int((n_h2**2 - n_h2) / 2)
      assert (
          n_rg == exp_n_rg
      ), f"The number of rg values given is {n_rg}, expected {exp_n_rg}"
      cor = np.zeros(shape=(n_h2, n_h2))
      # set upper triangle of correlation matrix to be rg
      cor[np.triu_indices(n=n_h2, k=1)] = rg
      cor += cor.T
      cor[np.diag_indices(n=n_h2)] = 1
      if psd_rg:
          cor0 = cor
          cor = _nearpsd(cor)
          idx = np.triu_indices(n=n_h2, k=1)
          maxlines = 50
          msg = [
              "Adjusting rg values to make covariance matrix positive semidefinite"
          ]
          msg += (
              [
                  (
                      f"{cor0[idx[0][i],idx[1][i]]} -> {cor[idx[0][i],idx[1][i]]}"
                  )
                  for i in range(n_rg)
              ]
              if n_rg <= maxlines
              else [
                  (
                      f"{cor0[idx[0][i],idx[1][i]]} -> {cor[idx[0][i],idx[1][i]]}"
                  )
                  for i in range(maxlines)
              ]
              + [
                  f"[ printed first {maxlines} rg changes -- omitted {n_rg - maxlines} ]"
              ]
          )
          print("\n".join(msg))
          rg = np.ravel(cor[idx])
      S = np.diag(h2) ** (1 / 2)
      cov_matrix = S @ cor @ S  # covariance matrix decomposition
  
      # check positive semidefinite
      if not np.all(np.linalg.eigvals(cov_matrix) >= 0) and not psd_rg:
          msg = "WARNING: Covariance matrix is not positive semidefinite.\n"
          msg += "Multivariate Gaussian distributions generated with this \n"
          msg += (
              "covariance matrix may not have the desired h2 and rg.\n"
          )
          msg += (
              "To make the covariance matrix positive semidefinite,\n"
          )
          msg += "adjust h2 and rg values or set psd_rg=True."
          print(msg)
      rg = rg.tolist()
      return cov_matrix, rg
  
imports: import hail
  from hail.typecheck import typecheck, oneof, nullable
  from hail.expr.expressions import expr_float64, expr_int32, expr_array, expr_call
  from hail.matrixtable import MatrixTable
  from hail.table import Table
  from hail.utils.java import Env
  import numpy
  import pandas
  import scipy.stats
label_desc: Use image node_id 5 for calling a global function with example usage: get_cov_matrix(h2, rg, psd_rg) and returns: cov_matrix,  rg