Skip to content

ArrayToBlocks/BlocksToArray

torchlinops.linops.ArrayToBlocks

Bases: NamedLinop

Extract sliding windows from an array.

Adjoint of BlocksToArray.

Source code in src/torchlinops/linops/array_to_blocks.py 
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
class ArrayToBlocks(NamedLinop):
    """Extract sliding windows from an array.

    Adjoint of [BlocksToArray](#BlocksToArray).

    """

    def __init__(
        self,
        grid_size: tuple[int, ...],
        block_size: tuple[int, ...],
        stride: tuple[int, ...],
        mask: Optional[Tensor] = None,
        batch_shape: Optional[Shape] = None,
        array_shape: Optional[Shape] = None,
        blocks_shape: Optional[Shape] = None,
    ):
        """
        Parameters
        ----------
        grid_size : tuple[int, ...]
            Size of the input array spatial dimensions.
        block_size : tuple[int, ...]
            Size of each extracted block.
        stride : tuple[int, ...]
            Stride between consecutive blocks.
        mask : Tensor, optional
            Boolean mask selecting a subset of blocks.
        batch_shape : Shape, optional
            Named shape for batch dimensions.
        array_shape : Shape, optional
            Named shape for the input array dimensions.
        blocks_shape : Shape, optional
            Named shape for the output block dimensions.
        """
        self.grid_size = grid_size
        self.ndim = len(self.grid_size)
        self.block_size = block_size
        self.stride = stride

        self.batch_shape = default_to(("...",), batch_shape)
        self.array_shape = default_to(("...",), array_shape)
        self.blocks_shape = default_to(("...",), blocks_shape)
        shape = NS(self.batch_shape) + NS(self.array_shape, self.blocks_shape)
        super().__init__(shape)

        if mask is not None:
            self.mask = nn.Parameter(mask, requires_grad=False)
        else:
            self.mask = mask

    @staticmethod
    def fn(arraytoblocks, x, /):
        return F.array_to_blocks(
            x,
            arraytoblocks.block_size,
            arraytoblocks.stride,
            arraytoblocks.mask,
        )

    @staticmethod
    def adj_fn(arraytoblocks, x, /):
        return F.blocks_to_array(
            x,
            arraytoblocks.grid_size,
            arraytoblocks.block_size,
            arraytoblocks.stride,
            arraytoblocks.mask,
        )

    @staticmethod
    def normal_fn(arraytoblocks, x, /):
        return arraytoblocks.adj_fn(arraytoblocks, arraytoblocks.fn(arraytoblocks, x))

    def split_forward(self, ibatch, obatch):
        return copy(self)

    def adjoint(self):
        return BlocksToArray(
            self.grid_size,
            self.block_size,
            self.stride,
            self.mask,
            self.batch_shape,
            self.blocks_shape,
            self.array_shape,
        )

    def size(self, dim):
        ndim = len(self.grid_size)
        if dim in self.ishape[-ndim:]:
            i = self.ishape.index(dim) - len(self.ishape)
            return self.grid_size[i]
        return None

__init__ 

__init__(
    grid_size: tuple[int, ...],
    block_size: tuple[int, ...],
    stride: tuple[int, ...],
    mask: Optional[Tensor] = None,
    batch_shape: Optional[Shape] = None,
    array_shape: Optional[Shape] = None,
    blocks_shape: Optional[Shape] = None,
)
PARAMETER DESCRIPTION
grid_size

Size of the input array spatial dimensions.

TYPE: tuple[int, ...]

block_size

Size of each extracted block.

TYPE: tuple[int, ...]

stride

Stride between consecutive blocks.

TYPE: tuple[int, ...]

mask

Boolean mask selecting a subset of blocks.

TYPE: Tensor DEFAULT: None

batch_shape

Named shape for batch dimensions.

TYPE: Shape DEFAULT: None

array_shape

Named shape for the input array dimensions.

TYPE: Shape DEFAULT: None

blocks_shape

Named shape for the output block dimensions.

TYPE: Shape DEFAULT: None

Source code in src/torchlinops/linops/array_to_blocks.py 
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
def __init__(
    self,
    grid_size: tuple[int, ...],
    block_size: tuple[int, ...],
    stride: tuple[int, ...],
    mask: Optional[Tensor] = None,
    batch_shape: Optional[Shape] = None,
    array_shape: Optional[Shape] = None,
    blocks_shape: Optional[Shape] = None,
):
    """
    Parameters
    ----------
    grid_size : tuple[int, ...]
        Size of the input array spatial dimensions.
    block_size : tuple[int, ...]
        Size of each extracted block.
    stride : tuple[int, ...]
        Stride between consecutive blocks.
    mask : Tensor, optional
        Boolean mask selecting a subset of blocks.
    batch_shape : Shape, optional
        Named shape for batch dimensions.
    array_shape : Shape, optional
        Named shape for the input array dimensions.
    blocks_shape : Shape, optional
        Named shape for the output block dimensions.
    """
    self.grid_size = grid_size
    self.ndim = len(self.grid_size)
    self.block_size = block_size
    self.stride = stride

    self.batch_shape = default_to(("...",), batch_shape)
    self.array_shape = default_to(("...",), array_shape)
    self.blocks_shape = default_to(("...",), blocks_shape)
    shape = NS(self.batch_shape) + NS(self.array_shape, self.blocks_shape)
    super().__init__(shape)

    if mask is not None:
        self.mask = nn.Parameter(mask, requires_grad=False)
    else:
        self.mask = mask

torchlinops.linops.BlocksToArray

Bases: NamedLinop

Compose several equally-sized blocks into a larger array.

Adjoint of ArrayToBlocks.

Source code in src/torchlinops/linops/array_to_blocks.py 
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
class BlocksToArray(NamedLinop):
    """Compose several equally-sized blocks into a larger array.

    Adjoint of [ArrayToBlocks](#ArrayToBlocks).
    """

    def __init__(
        self,
        grid_size: tuple[int, ...],
        block_size: tuple[int, ...],
        stride: tuple[int, ...],
        mask: Optional[Tensor] = None,
        batch_shape: Optional = None,
        blocks_shape: Optional = None,
        array_shape: Optional = None,
    ):
        """
        Parameters
        ----------
        grid_size : tuple[int, ...]
            Size of the output array spatial dimensions.
        block_size : tuple[int, ...]
            Size of each block.
        stride : tuple[int, ...]
            Stride between consecutive blocks.
        mask : Tensor, optional
            Boolean mask selecting a subset of blocks.
        batch_shape : optional
            Named shape for batch dimensions.
        blocks_shape : optional
            Named shape for the input block dimensions.
        array_shape : optional
            Named shape for the output array dimensions.
        """
        self.grid_size = grid_size
        self.ndim = len(self.grid_size)
        self.block_size = block_size
        self.stride = stride

        self.batch_shape = default_to(("...",), batch_shape)
        self.blocks_shape = default_to(("...",), blocks_shape)
        self.array_shape = default_to(("...",), array_shape)
        shape = NS(self.batch_shape) + NS(self.blocks_shape, self.array_shape)
        super().__init__(shape)
        if mask is not None:
            self.mask = nn.Parameter(mask, requires_grad=False)
        else:
            self.mask = mask

    @staticmethod
    def fn(blockstoarray, x, /):
        return F.blocks_to_array(
            x,
            blockstoarray.grid_size,
            blockstoarray.block_size,
            blockstoarray.stride,
            blockstoarray.mask,
        )

    @staticmethod
    def adj_fn(blockstoarray, x, /):
        if x.shape[-blockstoarray.ndim :] != blockstoarray.grid_size:
            raise RuntimeError(
                f"BlocksToArray expected input with full size {blockstoarray.grid_size} but got {x.shape}"
            )
        return F.array_to_blocks(
            x,
            blockstoarray.block_size,
            blockstoarray.stride,
            blockstoarray.mask,
        )

    def split_forward(self, ibatch, obatch):
        return copy(self)

    def adjoint(self):
        return ArrayToBlocks(
            self.grid_size,
            self.block_size,
            self.stride,
            self.mask,
            self.batch_shape,
            self.array_shape,
            self.blocks_shape,
        )

    def size(self, dim):
        ndim = len(self.grid_size)
        if dim in self.oshape[-ndim:]:
            i = self.oshape.index(dim) - len(self.oshape)
            return self.grid_size[i]
        return None

__init__ 

__init__(
    grid_size: tuple[int, ...],
    block_size: tuple[int, ...],
    stride: tuple[int, ...],
    mask: Optional[Tensor] = None,
    batch_shape: Optional = None,
    blocks_shape: Optional = None,
    array_shape: Optional = None,
)
PARAMETER DESCRIPTION
grid_size

Size of the output array spatial dimensions.

TYPE: tuple[int, ...]

block_size

Size of each block.

TYPE: tuple[int, ...]

stride

Stride between consecutive blocks.

TYPE: tuple[int, ...]

mask

Boolean mask selecting a subset of blocks.

TYPE: Tensor DEFAULT: None

batch_shape

Named shape for batch dimensions.

TYPE: optional DEFAULT: None

blocks_shape

Named shape for the input block dimensions.

TYPE: optional DEFAULT: None

array_shape

Named shape for the output array dimensions.

TYPE: optional DEFAULT: None

Source code in src/torchlinops/linops/array_to_blocks.py 
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
def __init__(
    self,
    grid_size: tuple[int, ...],
    block_size: tuple[int, ...],
    stride: tuple[int, ...],
    mask: Optional[Tensor] = None,
    batch_shape: Optional = None,
    blocks_shape: Optional = None,
    array_shape: Optional = None,
):
    """
    Parameters
    ----------
    grid_size : tuple[int, ...]
        Size of the output array spatial dimensions.
    block_size : tuple[int, ...]
        Size of each block.
    stride : tuple[int, ...]
        Stride between consecutive blocks.
    mask : Tensor, optional
        Boolean mask selecting a subset of blocks.
    batch_shape : optional
        Named shape for batch dimensions.
    blocks_shape : optional
        Named shape for the input block dimensions.
    array_shape : optional
        Named shape for the output array dimensions.
    """
    self.grid_size = grid_size
    self.ndim = len(self.grid_size)
    self.block_size = block_size
    self.stride = stride

    self.batch_shape = default_to(("...",), batch_shape)
    self.blocks_shape = default_to(("...",), blocks_shape)
    self.array_shape = default_to(("...",), array_shape)
    shape = NS(self.batch_shape) + NS(self.blocks_shape, self.array_shape)
    super().__init__(shape)
    if mask is not None:
        self.mask = nn.Parameter(mask, requires_grad=False)
    else:
        self.mask = mask