python(denoising): Copied files over

Signed-off-by: prescientmoon <git@moonythm.dev>

python/denoising/.gitignore

@@ -0,0 +1,4 @@
+.direnv
+.ipynb_checkpoints
+.envrc
+.pyc

python/denoising/flake.lock

@@ -0,0 +1,43 @@
+{
+  "nodes": {
+    "flake-utils": {
+      "locked": {
+        "lastModified": 1676283394,
+        "narHash": "sha256-XX2f9c3iySLCw54rJ/CZs+ZK6IQy7GXNY4nSOyu2QG4=",
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "rev": "3db36a8b464d0c4532ba1c7dda728f4576d6d073",
+        "type": "github"
+      },
+      "original": {
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "type": "github"
+      }
+    },
+    "nixpkgs": {
+      "locked": {
+        "lastModified": 1676428076,
+        "narHash": "sha256-8caqXsVSUvZqp4/pOLjbWWabzLEV/ZfDUiEv32WumKw=",
+        "owner": "nixos",
+        "repo": "nixpkgs",
+        "rev": "e0a054002198445ef987b55c34cde892cecc0de7",
+        "type": "github"
+      },
+      "original": {
+        "owner": "nixos",
+        "ref": "release-22.11",
+        "repo": "nixpkgs",
+        "type": "github"
+      }
+    },
+    "root": {
+      "inputs": {
+        "flake-utils": "flake-utils",
+        "nixpkgs": "nixpkgs"
+      }
+    }
+  },
+  "root": "root",
+  "version": 7
+}

python/denoising/flake.nix

@@ -0,0 +1,23 @@
+{
+  inputs = {
+    nixpkgs.url = "github:nixos/nixpkgs/release-22.11";
+    flake-utils.url = "github:numtide/flake-utils";
+  };
+
+  outputs = { self, nixpkgs, flake-utils }:
+    flake-utils.lib.eachDefaultSystem
+      (system:
+        let
+          pkgs = nixpkgs.legacyPackages.${system};
+          pyDeps = p: with p; [ numpy scipy matplotlib ];
+        in
+        rec {
+          devShell = pkgs.mkShell {
+            buildInputs = with pkgs; [
+              (python3.withPackages pyDeps)
+              jupyter
+              pandoc
+            ];
+          };
+        });
+}

python/denoising/modules/decompose.py

@@ -0,0 +1,44 @@
+import numpy as np
+
+def decompose(M):
+  n = len(M)
+
+  U = M.copy()
+  L = np.identity(n)
+
+  for i in range(0, n - 1):
+    for j in range(i + 1, n):
+      L[j, i] = U[j, i] / U[i, i]
+      U[j, i:n] -= L[j, i] * U[i, i:n]
+
+  assert np.allclose(M, L@U)
+
+  return L, U
+
+def solve_lower(m, v):
+  x = np.zeros(len(v))
+
+  for i in range(0, len(v)):
+    x[i] = (v[i] - m[i, :i] @ x[:i]) / m[i, i]
+
+  assert np.allclose(m@x, v)
+
+  return x
+
+def solve_upper(m, v):
+  x = np.zeros(len(v))
+
+  for i in reversed(range(0, len(m))):
+    x[i] = (v[i] - m[i, i+1:] @ x[i+1:]) / m[i, i]
+
+  assert np.allclose(m@x, v)
+
+  return x
+
+def solve(m, v):
+  L, U = decompose(m)
+  x = solve_upper(U, solve_lower(L, v))
+
+  assert np.allclose(m@x, v)
+
+  return x

python/denoising/modules/tridiagonal.py

@@ -0,0 +1,132 @@
+import numpy as np
+import scipy as sp
+
+def decompose(a, c, e):
+  """
+  Computes the LU decomposition of a tridiagonal matrix.
+  """
+  α = np.zeros(len(c))
+  β = a.copy()
+
+  for i in range(len(c)):
+    α[i] = e[i] / β[i]
+    β[i + 1] -= c[i] * α[i]
+
+  # Sanity check
+  if len(a) <= 10:
+    assert np.allclose(
+      to_array(a, c, e),
+      to_array(*from_lower(α)) @ to_array(*from_upper(β, c))
+    )
+
+  return (α, β)
+
+def to_csr(a, c, e):
+  """
+  Converts a tridiagonal matrix into a scipy csr sparse matrix.
+  """
+  n = len(c)
+
+  values = np.zeros(n * 3 + 1)
+  values[::3] = a
+  values[1::3] = c
+  values[2::3] = e
+  
+  col_indices = np.zeros_like(values)
+  col_indices[1::3] = np.arange(1, n + 1)
+  col_indices[2::3] = np.arange(0, n)
+  col_indices[3::3] = np.arange(1, n + 1)
+
+  index_ptr = np.zeros(n + 2)
+  index_ptr[1:n+1] = np.arange(2, n * 3 + 2, 3)
+  index_ptr[n+1] = n * 3 + 1
+
+  return sp.sparse.csr_array((values, col_indices, index_ptr))
+
+def to_array(a, c, e):
+  """
+  Converts a tridiagonal matrix into a numpy matrix.
+  """
+  return to_csr(a, c, e).toarray()
+
+def from_lower(α):
+  """
+  Turns the lower vector of a decomposition into a tridiagonal matrix.
+
+  Example ussage:
+  ```py 
+  α, β = decompose(m)
+  print(from_lower(α))
+  ```
+  """
+  return (np.ones(len(α) + 1), np.zeros(len(α)), α) 
+
+def from_upper(β, c):
+  """
+  Turns the upper vectors of a decomposition into a tridiagonal matrix.
+
+  Example ussage:
+  ```py 
+  α, β = decompose((a, c, e))
+  print(from_upper(β, c))
+  ```
+  """
+  return (β, c, np.zeros(len(c)))
+
+def solve_lower(α, rhs):
+  """
+  Solve a linear system of equations Mx = v
+  where M is a lower triangular matrix constructed
+  by LU decomposing a tridiagonal matrix.
+  """
+  x = np.zeros_like(rhs)
+
+  x[0] = rhs[0]
+
+  for i in range(1, len(rhs)):
+    x[i] = rhs[i] - α[i - 1] * x[i - 1]
+
+  if len(α) <= 10:
+    assert np.allclose(to_array(*from_lower(α)) @ x, rhs)
+
+  return x
+
+def solve_upper(β, c, rhs):
+  """
+  Solve a linear system of equations Mx = v
+  where M is an upper triangular matrix constructed
+  by LU decomposing a tridiagonal matrix.
+  """
+  x = np.zeros_like(rhs)
+
+  x[-1] = rhs[-1] / β[-1]
+
+  for i in reversed(range(len(rhs) - 1)):
+    x[i] = (rhs[i] - c[i] * x[i+1]) / β[i]
+
+  if len(β) <= 10:
+    assert np.allclose(to_array(*from_upper(β, c)) @ x, rhs)
+
+  return x
+
+def solve(a, c, e, rhs):
+  α, β = decompose(a, c, e)
+  x = solve_upper(β, c, solve_lower(α, rhs))
+
+  if len(α) <= 10:
+    assert np.allclose(to_array(a, c, e)@x, rhs)
+
+  return x
+
+# Small sanity check for the above code
+def main():
+  a, c, e = (np.ones(4), 2*np.ones(3), 3*np.ones(3))
+
+  rhs = np.ones(4)
+  result = solve(a, c, e, rhs)
+  print(f"m={to_array(a, c, e)}")
+  print(f"{rhs=}")
+  print(f"{result=}")
+  print(to_array(a, c, e) @ result)
+
+main()