Mercurial > repos > public > sbplib_julia

comparison src/LazyTensors/lazy_tensor_operations.jl @ 402:1936e38fe51e

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
Merge feature/lazy_linear_map
author Vidar Stiernström <vidar.stiernstrom@it.uu.se>
date Mon, 05 Oct 2020 10:45:30 +0200
parents de4746d6d126 c640f37d1c74
children 4aa59af074ef d94891b8dfca
comparison
equal deleted inserted replaced
382:5c10cd0ed1fe 402:1936e38fe51e
100 # end 100 # end
101 101
102 # # Have i gone too crazy with the type parameters? Maybe they aren't all needed? 102 # # Have i gone too crazy with the type parameters? Maybe they aren't all needed?
103 103
104 # export → 104 # export →
105 """
106 LazyLinearMap{T,R,D,...}(A, range_indicies, domain_indicies)
107
108 TensorMapping defined by the AbstractArray A. `range_indicies` and `domain_indicies` define which indicies of A should
109 be considerd the range and domain of the TensorMapping. Each set of indices must be ordered in ascending order.
110
111 For instance, if A is a m x n matrix, and range_size = (1,), domain_size = (2,), then the LazyLinearMap performs the
112 standard matrix-vector product on vectors of size n.
113 """
114 struct LazyLinearMap{T,R,D, RD, AA<:AbstractArray{T,RD}} <: TensorMapping{T,R,D}
115 A::AA
116 range_indicies::NTuple{R,Int}
117 domain_indicies::NTuple{D,Int}
118
119 function LazyLinearMap(A::AA, range_indicies::NTuple{R,Int}, domain_indicies::NTuple{D,Int}) where {T,R,D, RD, AA<:AbstractArray{T,RD}}
120 if !issorted(range_indicies) || !issorted(domain_indicies)
121 throw(DomainError("range_indicies and domain_indicies must be sorted in ascending order"))
122 end
123
124 return new{T,R,D,RD,AA}(A,range_indicies,domain_indicies)
125 end
126 end
127 export LazyLinearMap
128
129 range_size(llm::LazyLinearMap) = size(llm.A)[[llm.range_indicies...]]
130 domain_size(llm::LazyLinearMap) = size(llm.A)[[llm.domain_indicies...]]
131
132 function apply(llm::LazyLinearMap{T,R,D}, v::AbstractArray{T,D}, I::Vararg{Index,R}) where {T,R,D}
133 view_index = ntuple(i->:,ndims(llm.A))
134 for i ∈ 1:R
135 view_index = Base.setindex(view_index, Int(I[i]), llm.range_indicies[i])
136 end
137 A_view = @view llm.A[view_index...]
138 return sum(A_view.*v)
139 end