added NFMT module and tests

Author: passscoed

src/pyGroupedTransforms/GroupedTransform.py

@@ -35,9 +35,7 @@ def get_superposition_set(d, ds):
     return [tuple(item) for item in returnset]
 
 
-def get_setting(
-    system, N, U=None, d=None, ds=None
-):  # I have to add bases and NFMT stuff here!
+def get_setting(system, N, U=None, d=None, ds=None, basis_vect=None):
 
     if system not in systems:
         raise ValueError("System not found.")
@@ -52,18 +50,18 @@ def get_setting(
         bwl = [np.full(len(u), tmp[len(u)], "int32") for u in U]
 
         if systems[system] == NFMTtools:
-            if len(basis_vect) == 0:
-                ValueError(
+            if basis_vect is None or len(basis_vect) == 0:
+                raise ValueError(
                     "please call get_setting with basis_vect for a NFMT transform."
                 )
             if len(basis_vect) < d:
-                ValueError("basis_vect must have an entry for every dimension.")
+                raise ValueError("basis_vect must have an entry for every dimension.")
             return [
                 Setting(
                     u=U[idx],
                     mode=systems[system],
                     bandwidths=np.array(bwl[idx], "int32"),
-                    bases=basis_vect[U[idk]],
+                    bases=[basis_vect[j] for j in U[idx]],
                 )
                 for idx in range(len(U))
             ]
@@ -92,18 +90,18 @@ def get_setting(
                 bwl[i] = np.full(len(u), N[i])
 
         if systems[system] == NFMTtools:
-            if len(basis_vect) == 0:
-                ValueError(
+            if basis_vect is None or len(basis_vect) == 0:
+                raise ValueError(
                     "please call get_setting with basis_vect for a NFMT transform."
                 )
-            if len(basis_vect) < max(max(u) for u in U):
-                ValueError("basis_vect must have an entry for every dimension.")
+            if len(basis_vect) <= max(max(u) for u in U if len(u) > 0):
+                raise ValueError("basis_vect must have an entry for every dimension.")
             return [
                 Setting(
                     u=u,
                     mode=systems[system],
                     bandwidths=np.array(bwl[i], "int32"),
-                    bases=basis_vect[u],
+                    bases=[basis_vect[j] for j in u],
                 )
                 for i, u in enumerate(U)
             ]
@@ -131,18 +129,18 @@ def get_setting(
                 bwl[i] = N[i]
 
         if systems[system] == NFMTtools:
-            if len(basis_vect) == 0:
-                ValueError(
+            if basis_vect is None or len(basis_vect) == 0:
+                raise ValueError(
                     "please call get_setting with basis_vect for a NFMT transform."
                 )
-            if len(basis_vect) < max(max(u) for u in U):
-                ValueError("basis_vect must have an entry for every dimension.")
+            if len(basis_vect) <= max(max(u) for u in U if len(u) > 0):
+                raise ValueError("basis_vect must have an entry for every dimension.")
             return [
                 Setting(
                     u=U[idx],
                     mode=systems[system],
                     bandwidths=np.array(bwl[idx], "int32"),
-                    bases=basis_vect[U[idk]],
+                    bases=[basis_vect[j] for j in U[idx]],
                 )
                 for idx in range(len(U))
             ]
@@ -197,11 +195,11 @@ def __init__(
 
         if system == "mixed":
             if len(basis_vect) == 0:
-                ValueError(
+                raise ValueError(
                     "please call GroupedTransform with basis_vect for a NFMT transform."
                 )
             if len(basis_vect) != X.shape[1]:
-                ValueError("basis_vect must have an entry for every dimension.")
+                raise ValueError("basis_vect must have an entry for every dimension.")
 
         if system in {"exp", "chui1", "chui2", "chui3", "chui4"}:
             if np.min(X) < -0.5 or np.max(X) >= 0.5:
@@ -214,14 +212,14 @@ def __init__(
 
             cosine_mask = basis_vals > 0
             if np.sum(cosine_mask) > 0:
-                if (np.min(X[cosine_mask, :]) < 0) or (np.max(X[cosine_mask, :]) > 1):
+                if (np.min(X[:, cosine_mask]) < 0) or (np.max(X[:, cosine_mask]) > 1):
                     raise ValueError(
                         "Nodes must be between 0 and 1 for cosine or Chebyshev dimensions."
                     )
 
             exp_mask = ~cosine_mask
             if np.sum(exp_mask) > 0:
-                if (np.min(X[exp_mask, :]) < -0.5) or (np.max(X[exp_mask, :]) > 0.5):
+                if (np.min(X[:, exp_mask]) < -0.5) or (np.max(X[:, exp_mask]) > 0.5):
                     raise ValueError(
                         "Nodes must be between -0.5 and 0.5 for exponentional dimensions."
                     )
@@ -236,7 +234,9 @@ def __init__(
         self.parallel = parallel
 
         if len(settings) == 0:
-            self.settings = get_setting(system=system, N=N, U=U, d=d, ds=ds)
+            self.settings = get_setting(
+                system=system, N=N, U=U, d=d, ds=ds, basis_vect=basis_vect
+            )
         else:
             self.settings = settings
 
@@ -278,8 +278,9 @@ def __init__(
                 )
             elif system == "mixed":
                 matrix = np.array(
-                    s1.mode.get_matrix(bandwidths=s1.bandwidths, X=X[:, u1].T),
-                    bases=s1.bases,
+                    s1.mode.get_matrix(
+                        bandwidths=s1.bandwidths, X=X[:, u1].T, bases=s1.bases
+                    )
                 )
                 for s in self.settings[1:]:
                     if len(s.bandwidths) == 0:
@@ -290,8 +291,9 @@ def __init__(
                         [
                             matrix,
                             np.array(
-                                s.mode.get_matrix(s.bandwidths, X[:, u].T),
-                                bases=s.bases,
+                                s.mode.get_matrix(
+                                    s.bandwidths, X[:, u].T, bases=s.bases
+                                )
                             ),
                         ]
                     )
@@ -433,7 +435,7 @@ def get_matrix(self):
                 u1 = (0,)
             else:
                 u1 = s1.u
-            F_direct = s1.mode.get_matrix(s.bdanwidths, self.X[:, u1].T, bases=s.bases)
+            F_direct = s1.mode.get_matrix(s1.bandwidths, self.X[:, u1].T, bases=s1.bases)
             for idx, s in enumerate(self.settings):
                 if idx == 0:
                     continue
@@ -443,6 +445,7 @@ def get_matrix(self):
                     u = s.u
                 mat = s.mode.get_matrix(s.bandwidths, self.X[:, u].T, s.bases)
                 F_direct = np.hstack([F_direct, mat])
+            return F_direct
         else:
             s1 = self.settings[0]
             if len(s1.bandwidths) == 0:

src/pyGroupedTransforms/NFMTtools.py

@@ -1 +1,126 @@
-RuntimeError("NFMTtools is not implemented yet!")
+import numpy as np
+
+from pyGroupedTransforms import *
+from pyNFFT3.NFMT import BASES
+
+
+def datalength(bandwidths: np.ndarray) -> int:
+	"""Return the number of mixed-basis Fourier coefficients (zero-frequencies excluded).
+
+	For each dimension j the frequency range used is:
+	  exp  -> {-N_j/2, ..., -1, 1, ..., N_j/2 - 1}  (N_j - 1 values)
+	  cos/alg -> {1, 2, ..., N_j - 1}                 (N_j - 1 values)
+	Total: prod(N - 1).
+	"""
+	if bandwidths.ndim != 1 or bandwidths.dtype != "int32":
+		return "Please use an one-dimensional numpy.array with dtype 'int32' as input"
+	if len(bandwidths) == 0:
+		return 1
+	return int(np.prod(bandwidths - 1))
+
+
+def nfmt_index_set_without_zeros(bandwidths: np.ndarray, bases) -> np.ndarray:
+	"""Return the (d, prod(N-1)) integer frequency matrix.
+
+	Mirrors Julia's GroupedTransforms.NFMTtools.nfmt_index_set_without_zeros:
+	  exp  dimension j: {-N_j/2, ..., -1, 1, ..., N_j/2 - 1}
+	  cos/alg dimension j: {1, 2, ..., N_j - 1}
+	"""
+	d = len(bandwidths)
+	if d == 0:
+		return np.array([[0]], dtype=np.int64)
+
+	ranges = []
+	for n, basis in zip(bandwidths, bases):
+		if BASES[basis] > 0:
+			# cos or alg: {1, ..., N-1}
+			ranges.append(list(range(1, int(n))))
+		else:
+			# exp: {-N/2, ..., -1, 1, ..., N/2-1}
+			half = int(n) // 2
+			ranges.append(list(range(-half, 0)) + list(range(1, half)))
+
+	if d == 1:
+		return np.array(ranges[0], dtype=np.int64).reshape(1, -1)
+
+	# Cartesian product; last dimension changes fastest (row-major, matching Julia)
+	mesh = np.array(np.meshgrid(*ranges, indexing="ij"), dtype=np.int64)
+	return mesh.reshape(d, -1)
+
+
+def get_matrix(bandwidths, X, bases) -> np.ndarray:
+	"""Build the evaluation matrix F with F[m, j] = phi(x_m, k_j).
+
+	Mirrors Julia's GroupedTransforms.NFMTtools.get_matrix / get_phi:
+	  exp:     exp(-2pi i k x)
+	  cos:     sqrt(2) * cos(pi * k * x)   (k != 0 guaranteed by frequency set)
+	  alg:     sqrt(2) * cos(k * arccos(2x - 1))
+
+	X is expected in (d, M) format (columns are nodes), consistent with the
+	convention used in GroupedTransform.get_matrix.
+	"""
+	if X.ndim == 1 or X.shape[0] == 1:
+		X_eval = X.flatten().reshape(-1, 1)   # (M, 1)
+		d = 1
+		M = X_eval.shape[0]
+	else:
+		d, M = X.shape
+		X_eval = X.T                           # (M, d)
+
+	if len(bandwidths) == 0:
+		return np.ones((M, 1), dtype=np.complex128)
+
+	freq = nfmt_index_set_without_zeros(
+		np.asarray(bandwidths, dtype=np.int32), list(bases)
+	)                                          # (d, nf)
+
+	nf = freq.shape[1]
+	F = np.ones((M, nf), dtype=np.complex128)
+
+	for j in range(d):
+		n_j = freq[j]                          # (nf,)  all nonzero by construction
+		x_j = X_eval[:, j][:, None]           # (M, 1)
+		if BASES[bases[j]] == 1:
+			# cos basis: sqrt(2) * cos(pi * k * x)
+			F *= np.sqrt(2.0) * np.cos(np.pi * x_j * n_j)
+		elif BASES[bases[j]] == 2:
+			# alg (Chebyshev) basis: sqrt(2) * cos(k * arccos(2x - 1))
+			F *= np.sqrt(2.0) * np.cos(n_j * np.arccos(2.0 * x_j - 1.0))
+		else:
+			# exp basis: exp(-2pi i k x)
+			F *= np.exp(-2.0j * np.pi * x_j * n_j)
+
+	return F
+
+
+def get_transform(bandwidths: np.ndarray, X: np.ndarray, bases):
+	"""Return a DeferredLinearOperator for the mixed-basis grouped transform.
+
+	Uses the same matrix as get_matrix wrapped in a linear operator so that
+	trafo and adjoint are consistent by construction.
+	X is expected in (M, d) format (rows are nodes), consistent with the
+	convention used in GroupedTransform.transforms.
+	"""
+	if bandwidths.ndim > 1 or bandwidths.dtype != "int32":
+		return "Please use a zero or one-dimensional numpy.array with dtype 'int32' as input"
+
+	M = X.shape[0]
+
+	if len(bandwidths) == 0:
+		return DeferredLinearOperator(
+			dtype=np.complex128,
+			shape=(M, 1),
+			mfunc=lambda fhat: np.full(M, fhat[0], dtype=np.complex128),
+			rmfunc=lambda f: np.array([np.sum(f)], dtype=np.complex128),
+		)
+
+	# X is (M, d) here; get_matrix expects (d, M)
+	mat = get_matrix(bandwidths, X.T, list(bases))   # (M, nf)
+	N = int(np.prod(bandwidths - 1))
+
+	return DeferredLinearOperator(
+		dtype=np.complex128,
+		shape=(M, N),
+		mfunc=lambda fhat: mat @ np.asarray(fhat, dtype=np.complex128),
+		rmfunc=lambda f: mat.conj().T @ np.asarray(f, dtype=np.complex128),
+	)

tests/nfmt_U.py

@@ -0,0 +1,82 @@
+import os
+import sys
+
+src_aa = os.path.abspath(os.path.join(os.getcwd(), "src"))
+sys.path.insert(0, src_aa)
+
+import numpy as np
+
+from pyGroupedTransforms.GroupedTransforms import *
+
+rng = np.random.default_rng(42)
+
+d = 4
+ds = 3
+
+basis_vect = ["exp", "alg", "cos", "alg"]
+
+M = 1000
+# X in (M, d) format; exp dimension shifted to [-0.5, 0.5)
+X = rng.random((M, d))
+X[:, 0] -= 0.5
+
+U = [(), (1,), (1, 2)]
+
+# set up transform ###################################################
+
+F = GroupedTransform(
+    "mixed", X, U=U, N=[0, 64, 16], basis_vect=basis_vect
+)
+F_direct = F.get_matrix()
+
+# compute transform with NFMT ########################################
+
+fhat = GroupedCoefficients(F.settings)
+for i in range(len(F.settings)):
+    u = F.settings[i].u
+    fhat[u] = rng.random(len(fhat[u])) + 1.0j * rng.random(len(fhat[u]))
+
+# arithmetic tests ###################################################
+
+ghat = GroupedCoefficients(F.settings)
+for i in range(len(F.settings)):
+    u = F.settings[i].u
+    ghat[u] = rng.random(len(ghat[u])) + 1.0j * rng.random(len(ghat[u]))
+
+fhat[1]
+fhat[1] = 1.0 + 1.0j
+2 * fhat
+fhat + ghat
+fhat - ghat
+F[(1,)]
+fhat.set_data(ghat.data)
+
+###
+
+f = F * fhat
+
+# compute transform without NFMT #####################################
+
+f_direct = np.matmul(F_direct, fhat.vec())
+
+# compare results ####################################################
+
+error = np.linalg.norm(f - f_direct)
+assert error < 1e-5, f"trafo error {error} >= 1e-5"
+
+# generate random function values ####################################
+
+y = rng.random(M) + 1.0j * rng.random(M)
+
+# compute adjoint transform with NFMT ################################
+
+fhat = F * y
+
+# compute adjoint transform without NFMT #############################
+
+fhat_direct = np.matmul(np.conj(F_direct).T, y)
+
+# compare results ####################################################
+
+error = np.linalg.norm(fhat.vec() - fhat_direct)
+assert error < 1e-5, f"adjoint error {error} >= 1e-5"