def __setitem__(self, key, value):
key = (key, ) if not isinstance(key, tuple) else key
if all(isinstance(arg, basestring) for arg in key):
argidxs = EinsumTensor.split_indices(key)
#----------------------------------------#
if isinstance(value, EinsumContraction):
# Carry out the contraction and store the result...
# TODO Sanity checking for correct shape
value.contract(dest=EinsumTensor(argidxs, tensor=self))
elif isinstance(value, EinsumTensor):
idxs, subidxs = EinsumTensor.split_indices(argidxs,
include_sub=True)
eself = EinsumTensor(argidxs, tensor=self)
if idxs == value.indices and subidxs == value.sub_indices:
if value.coeff == 1.0:
super(Tensor, eself.sliced_tensor).__setitem__(
Ellipsis, value.sliced_tensor)
else:
super(Tensor, eself.sliced_tensor).__setitem__(
Ellipsis, value.coeff * value.sliced_tensor)
#np.ndarray.__setitem__(self, Ellipsis, value.tensor)
elif len(idxs) == len(value.indices):
# Just rearrange things...
rv = value.sort_to(dest=eself, multiplier=value.coeff)
else:
# We're doing an internal contraction. EinsumSum.sum_into handles this.
tmp = EinsumSum(value)
tmp.sum_into(eself)
elif isinstance(value, EinsumSum):
# TODO Sanity checking for correct shape
value.sum_into(dest=EinsumTensor(argidxs, tensor=self))
elif isinstance(value, Tensor):
if grendel.sanity_checking_enabled:
# Check the shape of the tensor to be assigned to the block
# TODO Sanity checking for correct shape
pass
#np.ndarray.__setitem__(dest_tens, tuple([Ellipsis]*self.ndim), value)
# The user probably made a mistake
raise ValueError
else:
super(Tensor, self).__setitem__(key, value)
def __getitem__(self, args):
"""
Returns the sub-``Tensor`` corresponding to the depth specified. If the resulting sub-``Tensor`` is just an item
(i.e. if the ``ndim`` attribute of self is equal to the number of arguments given), the ``numpy.ndarray``
behavior is used (which just returns the entry, which is a special numpy subclass of ``float`` or ``int``).
If the resulting sub-``Tensor`` has a ``ndim`` attribute of 1, a ``Vector`` object is returned.
If the resulting sub-``Tensor`` has a ``ndim`` attribute of 2, a ``Matrix`` object is returned.
:Examples:
TODO: Write example test cases
TODO: Move this to the class documentation since it doesn't show up in sphinx
"""
# if a single item is given, handle it as a length-1 tuple
args = (args, ) if not isinstance(args, tuple) else args
#----------------------------------------#
# Handle Einstein summation
if all(isinstance(arg, basestring) for arg in args):
argidxs = EinsumTensor.split_indices(args)
return EinsumTensor(argidxs, self)
#----------------------------------------#
else:
try:
ret_val = np.ndarray.__getitem__(self, args)
except:
raise
if np.isscalar(ret_val):
return ret_val
else:
shp = ret_val.shape
if len(shp) == 1:
return ret_val.view(Vector)
elif len(shp) == 2:
return ret_val.view(Matrix)
else:
return ret_val