github.com/Schaudge/hts@v0.0.0-20240223063651-737b4d69d68c/bam/index.go (about)

     1  // Copyright ©2014 The bíogo Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package bam
     6  
     7  import (
     8  	"encoding/binary"
     9  	"io"
    10  
    11  	"github.com/Schaudge/grailbase/errors"
    12  	"github.com/Schaudge/hts/bgzf"
    13  	"github.com/Schaudge/hts/bgzf/index"
    14  	"github.com/Schaudge/hts/internal"
    15  	"github.com/Schaudge/hts/sam"
    16  )
    17  
    18  // Index is a BAI index.
    19  type Index struct {
    20  	idx internal.Index
    21  
    22  	// MergeStrategy is used to determine the
    23  	// the merge strategy used to prepare the
    24  	// slice of chunks returned by Chunks.
    25  	// If MergeStrategy is nil, index.MergeStrategy
    26  	// is used.
    27  	MergeStrategy index.MergeStrategy
    28  }
    29  
    30  // NumRefs returns the number of references in the index.
    31  func (i *Index) NumRefs() int {
    32  	return len(i.idx.Refs)
    33  }
    34  
    35  // ReferenceStats returns the index statistics for the given reference and true
    36  // if the statistics are valid.
    37  func (i *Index) ReferenceStats(id int) (stats index.ReferenceStats, ok bool) {
    38  	s := i.idx.Refs[id].Stats
    39  	if s == nil {
    40  		return index.ReferenceStats{}, false
    41  	}
    42  	return index.ReferenceStats(*s), true
    43  }
    44  
    45  // Unmapped returns the number of unmapped reads and true if the count is valid.
    46  func (i *Index) Unmapped() (n uint64, ok bool) {
    47  	if i.idx.Unmapped == nil {
    48  		return 0, false
    49  	}
    50  	return *i.idx.Unmapped, true
    51  }
    52  
    53  // Add records the SAM record as having being located at the given chunk.
    54  func (i *Index) Add(r *sam.Record, c bgzf.Chunk) error {
    55  	return i.idx.Add(r, uint32(r.Bin()), c, isPlaced(r), isMapped(r))
    56  }
    57  
    58  func isPlaced(r *sam.Record) bool {
    59  	return r.Ref != nil && r.Pos != -1
    60  }
    61  
    62  func isMapped(r *sam.Record) bool {
    63  	return r.Flags&sam.Unmapped == 0
    64  }
    65  
    66  // Chunks returns a []bgzf.Chunk that corresponds to the given genomic interval.
    67  func (i *Index) Chunks(r *sam.Reference, beg, end int) ([]bgzf.Chunk, error) {
    68  	chunks, err := i.idx.Chunks(r.ID(), beg, end)
    69  	if err != nil {
    70  		return nil, err
    71  	}
    72  	if i.MergeStrategy == nil {
    73  		return index.Adjacent(chunks), nil
    74  	}
    75  	return i.MergeStrategy(chunks), nil
    76  }
    77  
    78  // MergeChunks applies the given MergeStrategy to all bins in the Index.
    79  func (i *Index) MergeChunks(s index.MergeStrategy) {
    80  	i.idx.MergeChunks(s)
    81  }
    82  
    83  var baiMagic = [4]byte{'B', 'A', 'I', 0x1}
    84  
    85  // ReadIndex reads the BAI Index from the given io.Reader.
    86  func ReadIndex(r io.Reader) (*Index, error) {
    87  	var (
    88  		idx   Index
    89  		magic [4]byte
    90  		err   error
    91  	)
    92  	err = binary.Read(r, binary.LittleEndian, &magic)
    93  	if err != nil {
    94  		return nil, err
    95  	}
    96  	if magic != baiMagic {
    97  		return nil, errors.New("bam: magic number mismatch")
    98  	}
    99  
   100  	var n int32
   101  	err = binary.Read(r, binary.LittleEndian, &n)
   102  	if err != nil {
   103  		return nil, err
   104  	}
   105  	if n == 0 {
   106  		return nil, nil
   107  	}
   108  	idx.idx, err = internal.ReadIndex(r, n, "bam")
   109  	if err != nil {
   110  		return nil, err
   111  	}
   112  	return &idx, nil
   113  }
   114  
   115  // WriteIndex writes the Index to the given io.Writer.
   116  func WriteIndex(w io.Writer, idx *Index) error {
   117  	err := binary.Write(w, binary.LittleEndian, baiMagic)
   118  	if err != nil {
   119  		return err
   120  	}
   121  
   122  	err = binary.Write(w, binary.LittleEndian, int32(len(idx.idx.Refs)))
   123  	if err != nil {
   124  		return err
   125  	}
   126  	return internal.WriteIndex(w, &idx.idx, "bam")
   127  }