From 5c0bfe7c960e2ef87d8bd01438140e5628bf933d Mon Sep 17 00:00:00 2001
From: Sheehan Olver <solver@mac.com>
Date: Mon, 23 Oct 2023 08:52:52 +0100
Subject: [PATCH] WeightedFactorization -> WeightedPlan

---
 src/bases/bases.jl | 33 ++++++++++++++++++++++++---------
 1 file changed, 24 insertions(+), 9 deletions(-)

diff --git a/src/bases/bases.jl b/src/bases/bases.jl
index 4fbbb4c..606a48f 100644
--- a/src/bases/bases.jl
+++ b/src/bases/bases.jl
@@ -178,8 +178,8 @@ grid(P, n...) = grid_layout(MemoryLayout(P), P, n...)
 
 
 function plan_grid_transform(lay, L, szs::NTuple{N,Int}, dims=1:N) where N
-    p = grid(L)
-    p, InvPlan(factorize(L[p,:]), dims)
+    p = grid(L, szs...)
+    p, InvPlan(factorize(L[p,1:length(p)]), dims)
 end
 
 function plan_grid_transform(::MappedBasisLayout, L, szs::NTuple{N,Int}, dims=1:N) where N
@@ -187,6 +187,7 @@ function plan_grid_transform(::MappedBasisLayout, L, szs::NTuple{N,Int}, dims=1:
     invmap(parentindices(L)[1])[x], F
 end
 
+
 plan_grid_transform(L, szs::NTuple{N,Int}, dims=1:N) where N = plan_grid_transform(MemoryLayout(L), L, szs, dims)
 
 plan_grid_transform(L, arr::AbstractArray, dims...) = plan_grid_transform(L, size(arr), dims...)
@@ -316,19 +317,33 @@ end
 
 
 """
-    WeightedFactorization(w, F)
+    WeightedPlan(w, F)
 
-weights a factorization `F` by `w`.
+weights a plan `F` by `w`.
 """
-struct WeightedFactorization{T, WW, FAC<:Factorization{T}} <: Factorization{T}
+struct WeightedPlan{T, WW, FAC<:Plan{T}, Dims} <: Plan{T}
     w::WW
     F::FAC
+    dims::Dims
+end
+
+function *(F::WeightedPlan, b::AbstractVecOrMat)
+    ret = copy(b)
+    for d in F.dims
+        if d == 1
+            b .= b ./ F.w
+        else
+            @assert d == 2
+            b .= b ./ transpose(F.w)
+        end
+    end
+    F.F * b
 end
 
-_factorize(::WeightedBasisLayouts, wS, dims...; kws...) = WeightedFactorization(weight(wS), factorize(unweighted(wS), dims...; kws...))
-
-
-\(F::WeightedFactorization, b::AbstractQuasiVector) = F.F \ (b ./ F.w)
+function plan_grid_transform(::WeightedBasisLayout, L, szs::NTuple{N,Int}, dims=1:N) where N
+    x,F = plan_grid_transform(unweighted(L), szs, dims)
+    x, WeightedPlan(weight(L)[x], F, dims)
+end
 
 ##
 # Algebra