Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Removing the redundancy between A and H as inputs in IterativeModels #28

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 16 additions & 16 deletions Tutorials/iterative.ipynb

Large diffs are not rendered by default.

69 changes: 27 additions & 42 deletions src/deep_unfolding/iterative_solvers.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@
from abc import ABC, abstractmethod

import torch
from numpy.typing import NDArray
from numpy.typing import NDArray # possible to add some prevention on the type of H
from torch import Tensor

from .utils import _decompose_matrix, _device
Expand All @@ -27,7 +27,7 @@ class IterativeModel(ABC):
"""Random matrix $H$."""

_bs: int
"""Batch size."""
"""Number of samples."""

_y: Tensor
"""Solution tensor."""
Expand Down Expand Up @@ -65,7 +65,6 @@ def s_hat(self) -> Tensor:

def __init__(
self,
a: NDArray,
h: Tensor,
bs: int,
y: Tensor,
Expand All @@ -75,9 +74,8 @@ def __init__(

Args:

a: Input square matrix to decompose.
h: Random matrix $H$.
bs: Batch size.
h: Input random matrix $H$.
bs: Number of samples.
y: Solution tensor.
device: Device to run the model on ('cpu' or 'cuda').
"""
Expand All @@ -90,7 +88,7 @@ def __init__(
self._s_hats = []

self._A, self._D, self._L, self._U, self._Dinv, self._Minv = _decompose_matrix(
a, device
h, device
)

_logger.info(f"Code run on : {_device}")
Expand Down Expand Up @@ -200,7 +198,6 @@ class GaussSeidel(IterativeModel):

def __init__(
self,
a: NDArray,
h: Tensor,
bs: int,
y: Tensor,
Expand All @@ -209,13 +206,12 @@ def __init__(
"""Initialize the Gauss-Seidel solver.

Args:
a: Input square matrix to decompose.
h: Random matrix H.
bs: Batch size.
h: Input random matrix H.
bs: Number of samples.
y: Solution tensor.
device: Device to run the model on ('cpu' or 'cuda').
"""
super().__init__(a, h, bs, y, device)
super().__init__(h, bs, y, device)

def _iterate(self, num_itr: int, yMF: Tensor, s: Tensor) -> None:
for _ in range(num_itr):
Expand All @@ -232,7 +228,6 @@ class Richardson(IterativeModel):

def __init__(
self,
a: NDArray,
h: Tensor,
bs: int,
y: Tensor,
Expand All @@ -243,13 +238,13 @@ def __init__(

Args:
a: Input square matrix to decompose.
h: Random matrix $H$.
bs: Batch size.
h: Input random matrix $H$.
bs: Number of samples.
y: Solution tensor.
omega: TODO explain.
device: Device where to run the model.
"""
super().__init__(a, h, bs, y, device)
super().__init__(h, bs, y, device)

self.omega = torch.tensor(omega)

Expand All @@ -268,7 +263,6 @@ class Jacobi(IterativeModel):

def __init__(
self,
a: NDArray,
h: Tensor,
bs: int,
y: Tensor,
Expand All @@ -278,14 +272,13 @@ def __init__(
"""Initialize the Jacobi iteration solver.

Args:
a: Input square matrix to decompose.
h: Random matrix $H$.
bs: Batch size.
h: Input random matrix $H$.
bs: Number of samples.
y: Solution tensor.
omega: Relaxation parameter for Jacobi iterations.
device: Device where to run the model.
"""
super().__init__(a, h, bs, y, device)
super().__init__(h, bs, y, device)
self.omega = torch.tensor(omega)

def _iterate(self, num_itr: int, yMF: Tensor, s: Tensor) -> None:
Expand All @@ -304,7 +297,6 @@ class SOR(IterativeModel):

def __init__(
self,
a: NDArray,
h: Tensor,
bs: int,
y: Tensor,
Expand All @@ -314,14 +306,13 @@ def __init__(
"""Initialize the SOR solver.

Args:
A: Input square matrix to decompose.
H: Random matrix $H$.
bs: Batch size.
H: Input random matrix $H$.
bs: Number of samples.
y: Solution tensor.
omega: Relaxation parameter for SOR iterations.
device: Device where to run the model.
"""
super().__init__(a, h, bs, y, device)
super().__init__(h, bs, y, device)
self.omega = torch.tensor(omega)

def _iterate(self, num_itr: int, yMF: Tensor, s: Tensor) -> None:
Expand Down Expand Up @@ -351,7 +342,6 @@ class SORCheby(IterativeModel):

def __init__(
self,
a: NDArray,
h: Tensor,
bs: int,
y: Tensor,
Expand All @@ -363,16 +353,15 @@ def __init__(
"""Initialize the SOR-Chebyshev solver.

Args:
a: Input square matrix to decompose.
h: Random matrix $H$.
bs: Batch size.
h: Input random matrix $H$.
bs: Number of samples.
y: Solution tensor.
omega: Relaxation parameter for SOR iterations.
omegaa: Acceleration parameter for SOR-Chebyshev iterations.
gamma: Damping factor for SOR-Chebyshev iterations.
device: Device where to run the model.
"""
super().__init__(a, h, bs, y, device)
super().__init__(h, bs, y, device)
self.omega = torch.tensor(omega)
self.omegaa = torch.tensor(omegaa)
self.gamma = torch.tensor(gamma)
Expand Down Expand Up @@ -412,7 +401,6 @@ class AOR(IterativeModel):

def __init__(
self,
a: NDArray,
h: Tensor,
bs: int,
y: Tensor,
Expand All @@ -423,15 +411,14 @@ def __init__(
"""Initialize the AOR solver.

Args:
a: Input square matrix to decompose.
h: Random matrix $H$.
bs: Batch size.
h: Input random matrix $H$.
bs: Number of samples.
y: Solution tensor.
omega: Relaxation parameter for AOR iterations.
r: Relaxation parameter.
device: Device where to run the model.
"""
super().__init__(a, h, bs, y, device)
super().__init__(h, bs, y, device)
self.omega = torch.tensor(omega)
self.r = torch.tensor(r)

Expand Down Expand Up @@ -465,7 +452,6 @@ class AORCheby(IterativeModel):

def __init__(
self,
a: NDArray,
h: Tensor,
bs: int,
y: Tensor,
Expand All @@ -476,15 +462,14 @@ def __init__(
"""Initialize the AOR-Chebyshev solver.

Args:
a: Input square matrix to decompose.
h: Random matrix $H$.
bs: Batch size.
h: Input random matrix $H$.
bs: Number of samples.
y: Solution tensor.
omega: Relaxation parameter for AOR iterations.
r: Relaxation parameter.
device: Device where to run the model.
"""
super().__init__(a, h, bs, y, device)
super().__init__(h, bs, y, device)
self.omega = torch.tensor(omega)
self.r = torch.tensor(r)

Expand Down Expand Up @@ -524,4 +509,4 @@ def _iterate(self, num_itr: int, yMF: Tensor, s: Tensor) -> None:
mu0 = mu1
mu1 = mu

self._s_hats.append(y)
self._s_hats.append(y)
12 changes: 8 additions & 4 deletions src/deep_unfolding/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,6 @@
from __future__ import annotations

import logging
import math

import numpy as np
import torch
Expand Down Expand Up @@ -50,7 +49,7 @@ def gen_linear(
- Tensor $y$ resulting from `solution @ H` of shape (`bs`, `m`)
"""
np.random.seed(seed=seed)
h = np.random.normal(0, 1.0 / math.sqrt(n), (n, m))
h = np.random.normal(0, 1.0 / np.sqrt(n), (n, m))
a = np.dot(h, h.T)
eig = np.linalg.eigvals(a)

Expand All @@ -64,21 +63,22 @@ def gen_linear(
- Max eigenvalue of A: {np.max(eig)}"""
)

# this part should be changed
solution = torch.normal(torch.zeros(bs, n), 1.0).to(device).detach()
y = solution @ ht.detach()

return a, ht, wt, solution, y


def _decompose_matrix(
a: NDArray | Tensor,
h: NDArray | Tensor,
device: torch.device = _device,
) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]:
"""Decompose a given matrix into its diagonal, lower triangular, upper
triangular components and their inverses.

Args:
a: Input square matrix to decompose.
a: Input random matrix to decompose.
device: Device to run the computations on.

Returns:
Expand All @@ -90,6 +90,10 @@ def _decompose_matrix(
- $D^{-1}$: Inverse of the diagonal matrix $D$.
- $M^{-1}$: Inverse of the matrix $D + L$.
"""

# square matrix generated from the original matrix
a = np.dot(h, h.T)

# Decomposed matrix calculations
d = np.diag(np.diag(a)) # Diagonal matrix
l = np.tril(a, -1) # Lower triangular matrix # noqa: E741
Expand Down
Loading