github.com/qiaogw/arrgo@v0.0.8/index_opt.go (about)

     1  package arrgo
     2  
     3  type Range struct {
     4  	Start, Stop int
     5  }
     6  
     7  func (a *Arrf) Index(ranges ...Range) *Arrf {
     8  	var ndim = len(a.Shape)
     9  	totalRanges := make([]Range, ndim)
    10  	copy(totalRanges, ranges)
    11  	if len(ranges) < ndim {
    12  		for i := len(ranges); i < ndim; i++ {
    13  			totalRanges[i] = Range{Start: 0, Stop: a.Shape[i]}
    14  		}
    15  	}
    16  
    17  	Shape := make([]int, ndim)
    18  	for i := range Shape {
    19  		Shape[i] = totalRanges[i].Stop - totalRanges[i].Start
    20  	}
    21  
    22  	b := Zeros(Shape...)
    23  
    24  	totalCount := 1
    25  	for i := 0; i < ndim; i++ {
    26  		totalCount *= Shape[i]
    27  	}
    28  
    29  	counterSrc := make([]int, ndim)
    30  	counterDst := make([]int, ndim)
    31  	for i := range counterSrc {
    32  		counterSrc[i] = totalRanges[i].Start
    33  		counterDst[i] = counterSrc[i] - totalRanges[i].Start
    34  	}
    35  
    36  	for index := 0; index < totalCount; index++ {
    37  		var v = a.At(counterSrc...)
    38  		b.Set(v, counterDst...)
    39  		counterSrc[ndim-1]++
    40  		counterDst[ndim-1] = counterSrc[ndim-1] - totalRanges[ndim-1].Start
    41  		var j = ndim - 1
    42  		for {
    43  			if j > 0 && counterSrc[j] == totalRanges[j].Stop {
    44  				counterSrc[j-1]++
    45  				counterSrc[j] = totalRanges[j].Start
    46  				counterDst[j-1] = counterSrc[j-1] - totalRanges[j-1].Start
    47  				counterDst[j] = counterSrc[j] - totalRanges[j].Start
    48  				j--
    49  			} else {
    50  				break
    51  			}
    52  		}
    53  	}
    54  
    55  	return b
    56  }