vincent-maillou / qttools

Quantum Transport Algorithms Toolbox
GNU General Public License v3.0
5 stars 1 forks source link

DSBSparse: Block-wise modification across a slice of the stack #30

Closed vetschn closed 2 months ago

vetschn commented 2 months ago

Particularly when doing batched things, we often want to modify only part of the local stack of a DSBSparse matrix, so something like

dsbsparse.blocks[0, 0][:3] = arr

Unfortunately this will not work, since the block indexer must return a new piece of memory. We should switch to some other way of modifying a specific stack slice. The easiest way of doing this would be to just include the slicing in the block indexer directly, so

dsbsparse.blocks[:3, ..., 0, 0] = arr

or maybe worse

dsbsparse.blocks[0, 0, :3] = arr

Other options we considered are things like

dsbsparse.blocks(slice(None, 3))[0, 0] = arr
dsbsparse[:3].blocks[0, 0] = arr

and personal favorite

x_ii.stack[:3].blocks[0, 0] = arr

Will try to implement this one asap...