wehub-resource-sync
130 lines
3.6 KiB
Python
130 lines
3.6 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
"""strided_slice/set in python"""
|
|
|
|
|
|
def strided_slice_python(data, begin, end, strides, slice_mode="end", axes=None):
|
|
"""Python version of strided slice operator.
|
|
|
|
Parameters
|
|
----------
|
|
data : numpy.ndarray
|
|
Input data
|
|
|
|
begin : list
|
|
Beginning of the slices.
|
|
|
|
end : list
|
|
End of the slices.
|
|
|
|
strides : list
|
|
The stride of each slice.
|
|
|
|
slice_mode : str, optional
|
|
The slice mode [end, size].
|
|
|
|
- ``"end"``: The default slice mode, ending indices for the slice.
|
|
- ``"size"``: The input strides will be ignored, input end in this mode indicates
|
|
the size of a slice starting at the location specified by begin. If end[i] is -1,
|
|
all remaining elements in that dimension are included in the slice.
|
|
|
|
axes : list, optional
|
|
Axes along which slicing is applied
|
|
|
|
Returns
|
|
-------
|
|
result : numpy.ndarray
|
|
The sliced result.
|
|
"""
|
|
strides = [] if strides is None else strides
|
|
if axes is not None:
|
|
rank = len(data.shape)
|
|
new_begin = [0] * rank
|
|
new_end = [data.shape[i] for i in range(rank)]
|
|
new_strides = [1] * rank
|
|
|
|
for i, axis in enumerate(axes):
|
|
new_begin[axis] = begin[i]
|
|
new_end[axis] = end[i]
|
|
if len(strides) > i:
|
|
new_strides[axis] = strides[i]
|
|
|
|
begin = new_begin
|
|
end = new_end
|
|
strides = new_strides
|
|
|
|
slices = []
|
|
for i in range(len(data.shape)):
|
|
new_stride = None
|
|
if slice_mode == "end" and i < len(strides):
|
|
new_stride = strides[i]
|
|
|
|
new_begin = begin[i] if i < len(begin) else None
|
|
if i >= len(end):
|
|
new_end = None
|
|
elif slice_mode == "size":
|
|
if end[i] < 0:
|
|
new_end = None
|
|
else:
|
|
new_end = new_begin + end[i]
|
|
else:
|
|
new_end = end[i]
|
|
|
|
slices.append(slice(new_begin, new_end, new_stride))
|
|
|
|
return data[tuple(slices)]
|
|
|
|
|
|
def strided_set_python(data, v, begin, end, strides):
|
|
"""Python version of strided slice operator.
|
|
|
|
Parameters
|
|
----------
|
|
data : numpy.ndarray
|
|
Input data
|
|
|
|
v : numpy.ndarray
|
|
Value data
|
|
|
|
begin : list
|
|
Beginning of the slices.
|
|
|
|
end : list
|
|
End of the slices.
|
|
|
|
strides : list
|
|
The stride of each slice.
|
|
|
|
Returns
|
|
-------
|
|
result : numpy.ndarray
|
|
The updated result.
|
|
"""
|
|
strides = [] if strides is None else strides
|
|
slices = []
|
|
res = data.copy()
|
|
for i in range(len(data.shape)):
|
|
slices.append(
|
|
slice(
|
|
begin[i] if i < len(begin) else None,
|
|
end[i] if i < len(end) else None,
|
|
strides[i] if i < len(strides) else None,
|
|
)
|
|
)
|
|
res[tuple(slices)] = v
|
|
return res
|