[src] Implement ssspy.linalg.solve for backward compatibility

Author: tky823

ssspy/bss/_update_spatial_model.py

@@ -3,7 +3,9 @@
 
 import numpy as np
 
-from ..linalg import eigh2, inv2
+from ..linalg._solve import solve
+from ..linalg.eigh import eigh2
+from ..linalg.inv import inv2
 from ..linalg.lqpqm import lqpqm2
 from ..special.flooring import identity, max_flooring
 from ..special.psd import to_psd
@@ -64,7 +66,7 @@ def update_by_ip1(
         e_n = E[:, src_idx, :]  # (n_bins, n_n_channels)
 
         WU = W @ U_n
-        w_n = np.linalg.solve(WU, e_n)  # (n_bins, n_channels)
+        w_n = solve(WU, e_n)  # (n_bins, n_channels)
         wUw = w_n[:, np.newaxis, :].conj() @ U_n @ w_n[:, :, np.newaxis]
         wUw = np.real(wUw[..., 0])
         wUw = np.maximum(wUw, 0)
@@ -358,8 +360,8 @@ def update_by_ip2_one_pair(
     WU_m = W @ U_m
     WU_n = W @ U_n
 
-    P_m = np.linalg.solve(WU_m, E_mn)
-    P_n = np.linalg.solve(WU_n, E_mn)
+    P_m = solve(WU_m, E_mn)
+    P_n = solve(WU_n, E_mn)
 
     PUP_m = P_m.transpose(0, 2, 1).conj() @ U_m @ P_m
     PUP_n = P_n.transpose(0, 2, 1).conj() @ U_n @ P_n
@@ -457,7 +459,7 @@ def update_by_ipa(
         C_n = d_n @ E_n
         d_n = d_n[:, :, source_idx]
 
-        Cd_n = np.linalg.solve(C_n, d_n)
+        Cd_n = solve(C_n, d_n)
         dCd_n = np.sum(d_n.conj() * Cd_n, axis=-1)
         dCd_n = np.real(dCd_n)
         eUe_n = U_tilde_n_inverse[:, source_idx, source_idx]
@@ -490,7 +492,7 @@ def update_by_ipa(
         Eq_n = q_n.conj() @ E_n.transpose(1, 0)
         q_tilde_n = e_n.transpose(1, 0) - Eq_n
 
-        Uq_n = np.linalg.solve(U_tilde_n, q_tilde_n)
+        Uq_n = solve(U_tilde_n, q_tilde_n)
         qUq_n = np.sum(q_tilde_n.conj() * Uq_n, axis=-1, keepdims=True)
 
         qUq_n = np.real(qUq_n)
@@ -577,8 +579,8 @@ def _is_zero(x: np.ndarray) -> np.ndarray:
             gamma_in = np.sum(pad_mask_i[:, na] * RXY_in[..., 0], axis=1)
 
             WU_in = W[:, neighbor_idx, :, :] @ U_in
-            eta_in = np.linalg.solve(WU_in, e_n)
-            eta_hat_in = np.linalg.solve(U_in, gamma_in)
+            eta_in = solve(WU_in, e_n)
+            eta_hat_in = solve(U_in, gamma_in)
             eta_U_in = eta_in[:, na, :].conj() @ U_in
 
             xi_in = eta_U_in @ eta_in[:, :, na]

ssspy/bss/admmbss.py

@@ -4,6 +4,7 @@
 import numpy as np
 
 from ..linalg import prox
+from ..linalg._solve import solve
 from .proxbss import ProxBSSBase
 
 EPS = 1e-10
@@ -232,7 +233,7 @@ def update_once(self) -> None:
         VY_tilde = np.sum(V_tilde - Y_tilde, axis=0)
         XVY_tilde = X.transpose(1, 0, 2).conj() @ VY_tilde.transpose(1, 2, 0)
 
-        W = np.linalg.solve(n_penalties * XX + E, VY + XVY_tilde.transpose(0, 2, 1))
+        W = solve(n_penalties * XX + E, VY + XVY_tilde.transpose(0, 2, 1))
         XW = self.separate(X, demix_filter=W)
 
         U = alpha * W + (1 - alpha) * V
@@ -426,7 +427,7 @@ def update_once(self) -> None:
         VY_tilde = V_tilde - Y_tilde
         XVY_tilde = X.transpose(1, 0, 2).conj() @ VY_tilde.transpose(1, 2, 0)
 
-        W = np.linalg.solve(XX + E, VY + XVY_tilde.transpose(0, 2, 1))
+        W = solve(XX + E, VY + XVY_tilde.transpose(0, 2, 1))
         XW = self.separate(X, demix_filter=W)
 
         U = alpha * W + (1 - alpha) * V

ssspy/bss/mnmf.py

@@ -3,6 +3,7 @@
 
 import numpy as np
 
+from ..linalg._solve import solve
 from ..linalg.mean import gmeanmh
 from ..special.flooring import identity, max_flooring
 from ..special.psd import to_psd
@@ -753,7 +754,7 @@ def separate(self, input: np.ndarray) -> np.ndarray:
         R = np.sum(R_n, axis=0)
         R = to_psd(R, flooring_fn=self.flooring_fn)
         R = np.tile(R, reps=(n_sources, 1, 1, 1, 1))
-        W_Hermite = np.linalg.solve(R, R_n)
+        W_Hermite = solve(R, R_n)
         W = W_Hermite.transpose(0, 1, 2, 4, 3).conj()
         W_ref = W[:, :, :, reference_id, :]
         W_ref = W_ref.transpose(0, 3, 1, 2)
@@ -777,7 +778,7 @@ def compute_loss(self) -> float:
             R = self.reconstruct_mnmf(T, V, H)
 
         R = to_psd(R, flooring_fn=self.flooring_fn)
-        XXR_inv = np.linalg.solve(R, XX)  # Hermitian transpose of XX @ np.linalg.inv(R)
+        XXR_inv = solve(R, XX)  # Hermitian transpose of XX @ np.linalg.inv(R)
         trace = np.trace(XXR_inv, axis1=-2, axis2=-1)
         trace = np.real(trace)
         logdet = self.compute_logdet(R)
@@ -857,10 +858,10 @@ def update_basis(
         def _compute_traces(
             target: np.ndarray, reconstructed: np.ndarray, spatial: np.ndarray
         ) -> np.ndarray:
-            RXX = np.linalg.solve(reconstructed, target)
+            RXX = solve(reconstructed, target)
             R = np.tile(reconstructed, reps=(n_sources, 1, 1, 1, 1))
             H = np.tile(spatial[:, :, na, :, :], reps=(1, 1, n_frames, 1, 1))
-            RH = np.linalg.solve(R, H)
+            RH = solve(R, H)
 
             trace_RXXRH = np.trace(RXX @ RH, axis1=-2, axis2=-1)
             trace_RXXRH = np.real(trace_RXXRH)
@@ -924,10 +925,10 @@ def update_activation(
         def _compute_traces(
             target: np.ndarray, reconstructed: np.ndarray, spatial: np.ndarray
         ) -> np.ndarray:
-            RXX = np.linalg.solve(reconstructed, target)
+            RXX = solve(reconstructed, target)
             R = np.tile(reconstructed, reps=(n_sources, 1, 1, 1, 1))
             H = np.tile(spatial[:, :, na, :, :], reps=(1, 1, n_frames, 1, 1))
-            RH = np.linalg.solve(R, H)
+            RH = solve(R, H)
 
             trace_RXXRH = np.trace(RXX @ RH, axis1=-2, axis2=-1)
             trace_RXXRH = np.real(trace_RXXRH)
@@ -1039,10 +1040,10 @@ def update_latent(
         def _compute_traces(
             target: np.ndarray, reconstructed: np.ndarray, spatial: np.ndarray
         ) -> np.ndarray:
-            RXX = np.linalg.solve(reconstructed, target)
+            RXX = solve(reconstructed, target)
             R = np.tile(reconstructed, reps=(n_sources, 1, 1, 1, 1))
             H = np.tile(spatial[:, :, na, :, :], reps=(1, 1, n_frames, 1, 1))
-            RH = np.linalg.solve(R, H)
+            RH = solve(R, H)
 
             trace_RXXRH = np.trace(RXX @ RH, axis1=-2, axis2=-1)
             trace_RXXRH = np.real(trace_RXXRH)
@@ -1207,7 +1208,7 @@ def separate(self, input: np.ndarray) -> np.ndarray:
         R = np.sum(R_n, axis=0)
         R = to_psd(R, flooring_fn=self.flooring_fn)
         R = np.tile(R, reps=(n_sources, 1, 1, 1, 1))
-        W_Hermite = np.linalg.solve(R, R_n)
+        W_Hermite = solve(R, R_n)
         W = W_Hermite.transpose(0, 1, 2, 4, 3).conj()
         W_ref = W[:, :, :, reference_id, :]
         W_ref = W_ref.transpose(0, 3, 1, 2)

ssspy/linalg/__init__.py

@@ -1,3 +1,4 @@
+from ._solve import solve
 from .cubic import cbrt
 from .eigh import eigh, eigh2
 from .inv import inv2
@@ -18,4 +19,5 @@
     "gmeanmh",
     "solve_cubic",
     "lqpqm2",
+    "solve",
 ]

ssspy/linalg/_solve.py

@@ -0,0 +1,21 @@
+import numpy as np
+from packaging import version
+
+np_version = np.__version__
+
+IS_NUMPY_GE_2 = version.parse(np.__version__) >= version.parse("2")
+
+
+def solve(a: np.ndarray, b: np.ndarray) -> np.ndarray:
+    requires_new_axis = IS_NUMPY_GE_2 and a.ndim == b.ndim + 1
+
+    if requires_new_axis:
+        b = b[..., np.newaxis]
+
+    x = np.linalg.solve(a, b)
+
+    if requires_new_axis:
+        x = x[..., 0]
+        b = b[..., 0]
+
+    return x