code: def trio_matrix(
      dataset, pedigree, complete_trios=False
  ) -> MatrixTable:
      """Builds and returns a matrix where columns correspond to trios and entries contain genotypes for the trio.
  
      .. include:: ../_templates/req_tstring.rst
  
      Examples
      --------
  
      Create a trio matrix:
  
      >>> pedigree = hl.Pedigree.read('data/case_control_study.fam')
      >>> trio_dataset = hl.trio_matrix(dataset, pedigree, complete_trios=True)
  
      Notes
      -----
  
      This method builds a new matrix table with one column per trio. If
      `complete_trios` is ``True``, then only trios that satisfy
      :meth:`.Trio.is_complete` are included. In this new dataset, the column
      identifiers are the sample IDs of the trio probands. The column fields and
      entries of the matrix are changed in the following ways:
  
      The new column fields consist of three structs (`proband`, `father`,
      `mother`), a Boolean field, and a string field:
  
      - **proband** (:class:`.tstruct`) - Column fields on the proband.
      - **father** (:class:`.tstruct`) - Column fields on the father.
      - **mother** (:class:`.tstruct`) - Column fields on the mother.
      - **id** (:py:data:`.tstr`) - Column key for the proband.
      - **is_female** (:py:data:`.tbool`) - Proband is female.
        ``True`` for female, ``False`` for male, missing if unknown.
      - **fam_id** (:py:data:`.tstr`) - Family ID.
  
      The new entry fields are:
  
      - **proband_entry** (:class:`.tstruct`) - Proband entry fields.
      - **father_entry** (:class:`.tstruct`) - Father entry fields.
      - **mother_entry** (:class:`.tstruct`) - Mother entry fields.
  
      Parameters
      ----------
      pedigree : :class:`.Pedigree`
  
      Returns
      -------
      :class:`.MatrixTable`
      """
      mt = dataset
      require_col_key_str(mt, "trio_matrix")
  
      k = mt.col_key.dtype.fields[0]
      samples = mt[k].collect()
  
      pedigree = pedigree.filter_to(samples)
      trios = (
          pedigree.complete_trios()
          if complete_trios
          else pedigree.trios
      )
      n_trios = len(trios)
  
      sample_idx = {}
      for i, s in enumerate(samples):
          sample_idx[s] = i
  
      trios = [
          hl.Struct(
              id=sample_idx[t.s],
              pat_id=None if t.pat_id is None else sample_idx[t.pat_id],
              mat_id=None if t.mat_id is None else sample_idx[t.mat_id],
              is_female=t.is_female,
              fam_id=t.fam_id,
          )
          for t in trios
      ]
      trios_type = hl.dtype(
          "array<struct{id:int,pat_id:int,mat_id:int,is_female:bool,fam_id:str}>"
      )
  
      trios_sym = Env.get_uid()
      entries_sym = Env.get_uid()
      cols_sym = Env.get_uid()
  
      mt = mt.annotate_globals(
          **{trios_sym: hl.literal(trios, trios_type)}
      )
      mt = mt._localize_entries(entries_sym, cols_sym)
      mt = mt.annotate_globals(
          **{
              cols_sym: hl.map(
                  lambda i: hl.bind(
                      lambda t: hl.struct(
                          id=mt[cols_sym][t.id][k],
                          proband=mt[cols_sym][t.id],
                          father=mt[cols_sym][t.pat_id],
                          mother=mt[cols_sym][t.mat_id],
                          is_female=t.is_female,
                          fam_id=t.fam_id,
                      ),
                      mt[trios_sym][i],
                  ),
                  hl.range(0, n_trios),
              )
          }
      )
      mt = mt.annotate(
          **{
              entries_sym: hl.map(
                  lambda i: hl.bind(
                      lambda t: hl.struct(
                          proband_entry=mt[entries_sym][t.id],
                          father_entry=mt[entries_sym][t.pat_id],
                          mother_entry=mt[entries_sym][t.mat_id],
                      ),
                      mt[trios_sym][i],
                  ),
                  hl.range(0, n_trios),
              )
          }
      )
      mt = mt.drop(trios_sym)
  
      return mt._unlocalize_entries(entries_sym, cols_sym, ["id"])
  
file_path: hail/hail/python/hail/methods/family_methods.py
imports: from typing import Tuple
  import hail
  import hail.expr.aggregators
  from hail.genetics.pedigree import Pedigree
  from hail.matrixtable import MatrixTable
  from hail.expr import expr_call, expr_float64
  from hail.table import Table
  from hail.typecheck import typecheck, numeric
  from hail.utils.java import Env
  from .misc import require_biallelic, require_col_key_str
label_desc: Use image node_id 1 for calling a global function with example usage: trio_matrix(dataset, pedigree, complete_trios) and returns: mt