Index

class Add(NamedLinop):
    """The sum of one or more linear operators.

    When ``threaded=True`` (default), each sub-linop is executed in parallel
    using a ThreadPoolExecutor, which is useful for I/O-bound operations or
    operations that release the GIL (e.g., PyTorch tensor operations).

    Shared linops (e.g., ``Add(A, A)``) are used directly without copying —
    each thread receives the same linop object but independent execution context.

    Attributes
    ----------
    linops : nn.ModuleList
        The list of linops being added together.
    threaded : bool
        Whether to run sub-linops in parallel. Default is True.
    num_workers : int | None
        Number of worker threads. If None, defaults to the number of sub-linops.
    """

    is_container = True

    def __init__(
        self,
        *linops,
        threaded: bool = True,
        num_workers: Optional[int] = None,
        **kwargs,
    ):
        """
        Parameters
        ----------
        *linops : tuple[NamedLinop]
            The linear operators to be added together.
        threaded : bool, optional
            Whether to run sub-linops in parallel. Default is True.
        num_workers : int | None, optional
            Number of worker threads. If None, defaults to the number of sub-linops.
        """
        if config.shape_inference:
            # Find the most specific shape across all linops
            max_ishape = max_shape([linop.ishape for linop in linops])
            max_oshape = max_shape([linop.oshape for linop in linops])

            for linop in linops:
                resolved_ishape = resolve_wildcards(linop.ishape, max_ishape)
                if resolved_ishape != linop.ishape:
                    linop.ishape = resolved_ishape

                resolved_oshape = resolve_wildcards(linop.oshape, max_oshape)
                if resolved_oshape != linop.oshape:
                    linop.oshape = resolved_oshape

        assert all(isequal(linop.ishape, linops[0].ishape)[0] for linop in linops), (
            f"Add: All linops must share same ishape. Found {linops}"
        )
        assert all(isequal(linop.oshape, linops[0].oshape)[0] for linop in linops), (
            f"Add: All linops must share same oshape. Linops: {linops}"
        )
        super().__init__(NS(linops[0].ishape, linops[0].oshape), **kwargs)
        self.threaded = threaded
        self.num_workers = num_workers
        self._linops = nn.ModuleList(linops)

    @property
    def linops(self):
        return self._linops

    @linops.setter
    def linops(self, new_linops):
        self._linops = new_linops

    def __setattr__(self, name, value):
        """Bypass PyTorch's setattr for linops."""
        if name == "linops":
            type(self).linops.fset(self, value)
        else:
            super().__setattr__(name, value)

    @staticmethod
    def fn(add, x: torch.Tensor, /, context=None):
        return parallel_execute(
            add.linops,
            [x] * len(add),
            context,
            reduce_fn=sum,
            threaded=add.threaded,
            num_workers=add.num_workers,
        )

    @staticmethod
    def adj_fn(add, x: torch.Tensor, /, context=None):
        return parallel_execute(
            [linop.H for linop in add.linops],
            [x] * len(add),
            context,
            reduce_fn=sum,
            threaded=add.threaded,
            num_workers=add.num_workers,
        )

    @staticmethod
    def split(add, tile):
        split = copy(add)
        linops = [type(linop).split(linop, tile) for linop in add.linops]
        split.linops = nn.ModuleList(linops)
        return split

    def adjoint(self):
        adj = copy(self)
        adj.linops = nn.ModuleList([linop.adjoint() for linop in self.linops])
        adj.shape = self.shape.adjoint()
        return adj

    def normal(self, inner=None):
        if inner is None:
            max_ishape = max_shape([linop.N.ishape for linop in self.linops])
            max_oshape = max_shape([linop.N.oshape for linop in self.linops])
            new_shape = NS(max_ishape, max_oshape)
            all_combinations = []
            for left_linop in self.linops:
                for right_linop in self.linops:
                    if left_linop == right_linop:
                        all_combinations.append(copy(left_linop.N))
                    else:
                        all_combinations.append(copy(left_linop.H) @ copy(right_linop))
            all_combinations = standardize_shapes(all_combinations, new_shape)
            normal = copy(self)
            normal.linops = nn.ModuleList(list(all_combinations))
            normal.shape = normal.linops[0].shape
            return normal
        return super().normal(inner)

    def size(self, dim):
        for linop in self.linops:
            out = linop.size(dim)
            if out is not None:
                return out
        return None

    @property
    def dims(self):
        return set().union(*[linop.dims for linop in self.linops])

    @property
    def H(self):
        try:
            if config.cache_adjoint_normal:
                if self._adjoint is None:
                    self._adjoint = [self.adjoint()]
                return self._adjoint[0]
            return self.adjoint()
        except AttributeError as e:
            raise RuntimeError(f"AttributeError in {type(self).__name__}.H: {e}") from e

    @property
    def N(self):
        try:
            if config.cache_adjoint_normal:
                if self._normal is None:
                    self._normal = [self.normal()]
                return self._normal[0]
            return self.normal()
        except AttributeError as e:
            raise RuntimeError(f"AttributeError in {type(self).__name__}.N: {e}") from e

    def flatten(self):
        return [self]

    def __getitem__(self, idx):
        linops = self.linops[idx]
        if isinstance(linops, NamedLinop):
            return linops
        new = copy(self)
        new.linops = nn.ModuleList(linops)
        return new

    def __len__(self):
        return len(self.linops)

    def __repr__(self):
        linop_chain = " + ".join(repr(linop) for linop in self.linops)
        return linop_chain

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))

    @staticmethod
    def split(linop, tile):
        return copy(linop)

    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

@dataclass(frozen=True)
class BatchSpec:
    """Specification for splitting and distributing a linop across devices.

    Parameters
    ----------
    batch_sizes : dict[ND | str, int]
        Mapping from dimension names to chunk sizes for tiling.
    device_matrix : np.ndarray or list, optional
        Array of ``torch.device`` objects specifying target devices for each
        tile. Broadcast to match the tile grid shape.
    base_device : torch.device, optional
        The device where input/output data resides. Default is CPU.
    """

    batch_sizes: dict[ND | str, int]
    device_matrix: np.ndarray | None = None
    base_device: torch.device | None = None

    def resolve(self, linop, default_device: torch.device) -> ResolvedBatchSpec:
        """Return a new ResolvedBatchSpec with all fields filled in.

        Parameters
        ----------
        linop : NamedLinop
            The operator to resolve the device matrix against (for broadcasting).
        default_device : torch.device
            The default device to use for unspecified fields.

        Returns
        -------
        ResolvedBatchSpec
            A new frozen resolved specification with computed values.
        """
        base_device = default_to(default_device, self.base_device)
        device_matrix = default_to(np.array([base_device]), self.device_matrix)

        device_matrix = fuzzy_broadcast_to(
            ensure_ndarray(device_matrix), _tiled_shape(linop, self.batch_sizes)
        )
        return ResolvedBatchSpec(
            batch_sizes=self.batch_sizes,
            device_matrix=device_matrix,
            base_device=base_device,
        )

    def broadcast_device_matrix(self, linop):
        batch_dims = list(self.batch_sizes.keys())
        sizes = {dim: linop.size(dim) for dim in linop.dims}
        tiled_shape = tuple(
            ceil(sizes[dim] / self.batch_sizes[dim]) for dim in batch_dims
        )
        device_matrix = fuzzy_broadcast_to(
            ensure_ndarray(self.device_matrix), tiled_shape
        )
        return device_matrix

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,
        )

    @staticmethod
    def split(linop, tile):
        return copy(linop)

    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

class BreakpointLinop(NamedLinop):
    """Debugging identity operator that drops into ``pdb`` on forward/adjoint.

    Useful for inspecting intermediate tensor values inside a ``Chain``.
    """

    def __init__(self, ioshape: Optional[Shape] = None):
        super().__init__(NS(ioshape))

    @staticmethod
    def fn(linop, x, /):
        breakpoint()
        return x

    @staticmethod
    def adj_fn(linop, x, /):
        breakpoint()
        return x

    @staticmethod
    def split(linop, tile):
        return linop

class Chain(NamedLinop):
    """Composition (sequential application) of named linear operators.

    If ``Chain(A, B, C)`` is created, then the forward pass applies
    $A$ first, then $B$, then $C$: mathematically the operator is $C B A$.

    Attributes
    ----------
    linops : nn.ModuleList
        The constituent linops in **execution order** (inner to outer).
    """

    is_container = True

    def __init__(self, *linops, name: Optional[str] = None):
        """
        Parameters
        ----------
        *linops : NamedLinop
            Linops in order of execution. If ``linops = (A, B, C)``, the
            mathematical operator is $C B A$.
        name : str, optional
            Display name for this chain.
        """
        super().__init__(NS(linops[0].ishape, linops[-1].oshape), name=name)
        if len(linops) == 0:
            raise ValueError(f"Chain must contain at least one linop.")
        self.linops = nn.ModuleList(list(linops))

        if config.shape_inference:
            self._infer_shapes()

        self._check_inputs_outputs()

    @property
    def linops(self):
        return self._linops

    @linops.setter
    def linops(self, new_linops):
        self._linops = new_linops

    def __setattr__(self, name, value):
        """Bypasses pytorch's setattr, just for linops"""
        if name == "linops":
            # Force descriptor lookup for this name
            type(self).linops.fset(self, value)
        else:
            super().__setattr__(name, value)

    def _check_inputs_outputs(self):
        curr_shape = self.ishape
        for i, linop in enumerate(self.linops):
            if not isequal(linop.ishape, curr_shape)[0]:
                raise ValueError(
                    f"Mismatched shape: expected {linop.ishape}, got {curr_shape} at input to {linop}. Full stack: {self}, index {i}"
                )
            curr_shape = linop.oshape

    def _infer_shapes(self):
        """Propagate shapes forward through the chain, resolving wildcards."""
        for i in range(len(self.linops) - 1):
            curr = self.linops[i]
            next_linop = self.linops[i + 1]

            # Resolve wildcards in next_linop.ishape using curr.oshape
            resolved_ishape = resolve_wildcards(next_linop.ishape, curr.oshape)
            if resolved_ishape != next_linop.ishape:
                next_linop.ishape = resolved_ishape

    @staticmethod
    def fn(chain, x: torch.Tensor, context=None):
        x = chain[0](x, context)  # First linop inherits context
        for linop in chain.linops[1:]:
            x = linop(x)
        return x

    @staticmethod
    def adj_fn(chain, x: torch.Tensor, context=None):
        linops = list(reversed(chain.linops))
        x = linops[0].H(x, context)
        for linop in linops[1:]:
            x = linop.H(x)
        return x

    @staticmethod
    def normal_fn(chain, x: torch.Tensor, context=None):
        # fn does the reversing so it's unnecessary to do it here
        # If the normal hasn't been explicitly formed with`.N`, do things the naive way
        # Only the inner chain.fn inherits context
        return chain.adj_fn(chain, chain.fn(chain, x, context), context=None)

    @staticmethod
    def split(chain, tile: Mapping[ND | str, slice]):
        """Split a chain linop into sub-linops.

        Distributes the tile to each constituent linop based on dimension names.

        Parameters
        ----------
        chain : Chain
            The chain linop to split.
        tile : Mapping[ND | str, slice]
            Dictionary specifying how to slice the linop dimensions
        """
        split_linops = []
        for linop in chain.linops:
            sub_tile = {dim: tile.get(dim, slice(None)) for dim in linop.dims}
            split_linops.append(type(linop).split(linop, sub_tile))
        split = copy(chain)
        split.linops = nn.ModuleList(split_linops)
        return split

    def size(self, dim):
        out = None
        for linop in self.linops:
            tmp = linop.size(dim)
            if tmp is not None:
                if out is None:
                    out = tmp
                elif out != tmp:
                    raise ValueError(
                        f"Conflicting linop sizes found: {out} and {tmp} for dim {dim} in linop {linop} out of all linops {self.linops}"
                    )
        return out

    @property
    def dims(self):
        """Get the dims that appear anywhere in this linop chain."""
        return set().union(*[linop.dims for linop in self.linops])

    def adjoint(self):
        linops = list(linop.adjoint() for linop in reversed(self.linops))
        adj = copy(self)
        adj.linops = nn.ModuleList(linops)
        return adj

    def normal(self, inner=None):
        """Compute the normal operator by folding through the chain.

        For a chain $C B A$, the normal is computed as
        $A^H (B^H (C^H C (B (A \\cdot))))$ by iterating ``linop.normal(inner)``
        in reverse order. This enables Toeplitz embedding and other per-linop
        normal optimizations to compose correctly.

        Parameters
        ----------
        inner : NamedLinop, optional
            An inner operator seeded from an outer chain or ``None``.

        Returns
        -------
        NamedLinop
            The composed normal operator.
        """
        for linop in reversed(self.linops):
            inner = linop.normal(inner)
        return inner

    @property
    def shape(self):
        return NS(self.linops[0].ishape, self.linops[-1].oshape)

    @shape.setter
    def shape(self, val):
        self.ishape = val.ishape
        self.oshape = val.oshape

    @property
    def ishape(self):
        return self.linops[0].ishape

    @ishape.setter
    def ishape(self, val):
        self.linops[0].ishape = val

    @property
    def oshape(self):
        return self.linops[-1].oshape

    @oshape.setter
    def oshape(self, val):
        self.linops[-1].oshape = val

    def flatten(self):
        return list(self.linops)

    def __getitem__(self, idx):
        linops = self.linops[idx]
        if isinstance(linops, NamedLinop):
            return linops
        return Chain(*linops, name=self._name)

    def __len__(self):
        return len(self.linops)

    def __repr__(self):
        output = ""
        output += INDENT.indent(self.repr_name + "(\n")
        with INDENT:
            for linop in self.linops:
                output += repr(linop) + "\n"
        output += INDENT.indent(")")
        return output

    def __copy__(self):
        return super().__copy__()