From fa2edee85f3b72745c41d1e40479c1a4ddaac2dd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tam=C3=A1s=20K=2E=20Papp?= <tkpapp@gmail.com>
Date: Mon, 1 Sep 2025 15:00:30 +0200
Subject: [PATCH 01/10] start experimental model interface

---
 Project.toml              |  4 ++-
 src/SpectralKit.jl        |  1 +
 src/experimental.jl       | 54 +++++++++++++++++++++++++++++++
 test/Project.toml         |  1 +
 test/runtests.jl          |  1 +
 test/test_experimental.jl | 68 +++++++++++++++++++++++++++++++++++++++
 6 files changed, 128 insertions(+), 1 deletion(-)
 create mode 100644 src/experimental.jl
 create mode 100644 test/test_experimental.jl

diff --git a/Project.toml b/Project.toml
index 22ef9c1..12b2440 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,16 +1,18 @@
 name = "SpectralKit"
 uuid = "5c252ae7-b5b6-46ab-a016-b0e3d78320b7"
-authors = ["Tamas K. Papp <tkpapp@gmail.com>"]
 version = "0.16.1"
+authors = ["Tamas K. Papp <tkpapp@gmail.com>"]
 
 [deps]
 ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"
+Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 
 [compat]
 ArgCheck = "1, 2"
+Compat = "4.18.0"
 DocStringExtensions = "0.8, 0.9"
 OrderedCollections = "1"
 StaticArrays = "1"
diff --git a/src/SpectralKit.jl b/src/SpectralKit.jl
index e794d09..4757c2c 100644
--- a/src/SpectralKit.jl
+++ b/src/SpectralKit.jl
@@ -13,5 +13,6 @@ include("generic_api.jl")
 include("chebyshev.jl")
 include("smolyak_traversal.jl")
 include("smolyak_api.jl")
+include("experimental.jl") # experimental code is not part of the API, see its module docstring
 
 end # module
diff --git a/src/experimental.jl b/src/experimental.jl
new file mode 100644
index 0000000..506e2b1
--- /dev/null
+++ b/src/experimental.jl
@@ -0,0 +1,54 @@
+"""
+Experimental functionality. Officially not (yet) part of the API.
+
+A best effort is made to
+
+1. only change in breaking ways (including integration into the public API) with major or
+   minor releases.
+2. fix issues (users are encouraged to open them)
+"""
+module Experimental
+
+using Compat: @compat
+
+@compat public model_parameters_dimension, make_model_parameters, calculate_derived_quantities
+
+using DocStringExtensions: FUNCTIONNAME, SIGNATURES
+
+####
+#### endpoint mappings
+####
+
+# TODO
+
+####
+#### modelling API
+####
+
+"""
+$(FUNCTIONNAME)(model_family)
+
+Dimension of the model parameters as a flat vector.
+"""
+function model_parameters_dimension end
+
+"""
+$(FUNCTIONNAME)(model_family, x::AbstractVector)
+
+Convert a flat vector of dimension [`model_parameters_dimension`)(@ref) to the model
+parameters, which can be an arbitrary value (eg a `NamedTuple`), which is not used for
+dispatch.
+
+!!! NOTE
+    The method should accept all finite numbers in ``ℝ``, and transform them accordingly.
+"""
+function make_model_parameters end
+
+"""
+$(SIGNATURES)(model_family, model_parameters)
+
+Calculate derived quantities (for use in determining the bases and calculating the residuals).
+"""
+function calculate_derived_quantities end
+
+end
diff --git a/test/Project.toml b/test/Project.toml
index ddfa4a9..8f8be9c 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -4,6 +4,7 @@ BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
+LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Sobol = "ed01d8cd-4d21-5b2a-85b4-cc3bdc58bad4"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
diff --git a/test/runtests.jl b/test/runtests.jl
index 601fb1d..0bd0d81 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -11,6 +11,7 @@ include("test_chebyshev.jl")
 include("test_smolyak_traversal.jl")
 include("test_smolyak.jl")
 include("test_generic_api.jl")  # NOTE moved last as it used constructs from above
+include("test_experimental.jl") # NOTE experimental code is not public API
 
 using JET
 @testset "static analysis with JET.jl" begin
diff --git a/test/test_experimental.jl b/test/test_experimental.jl
new file mode 100644
index 0000000..67388f5
--- /dev/null
+++ b/test/test_experimental.jl
@@ -0,0 +1,68 @@
+import SpectralKit.Experimental as SK
+
+using LogExpFunctions: logistic
+
+"""
+A simple Ramsey model in discrete time.
+
+### Setup
+
+We document the model mainly to fix notation and make the code self-contained. Time is
+discrete, but below time periods are `0` (“this”) and `1` (“next”) without loss of
+generality.
+
+The resource constraint is
+```math
+c_0 + k_1 = f(k_0)
+```
+where ``f`` includes production and depreciation. The value function is
+```math
+V(k) = max_{0 < c ≤ f(k)} u(c) + β V(f(k) - c)
+```
+which is solved by the *policy function* ``c(k)``.
+
+The first order condition is
+```math
+u'(c(k)) = β V'(f(k) - c(k))
+```
+while the envelope condition is
+```math
+V'(k) = β V'(f(k) - c(k)) f'(k)
+```
+Combine to obtain
+```math
+V'(k) = u'(c(k)) f'(k)
+```
+and then
+```math
+u'(c(k)) f'(k) = β u'(c(f(k) - c(k))) f'(c(f(k) - c(k))) f'(k)
+```
+Then cancel ``f'(k)`` and rearrange as unitless
+```math
+1 = β u'(c(f(k) - c(k))) f'(c(f(k)) - c(k)) / u'(k)
+```
+From this it follows that the steady state ``k_s`` solves
+```
+f'(k_s) β = 1
+```
+and the resource constraint yields the steady state consumption as
+```
+c_s = f(k_s) - k_s
+```
+"""
+struct RamseyDiscrete end
+
+SK.model_parameters_dimension(::RamseyDiscrete) = 3
+
+function SK.make_model_parameters(::RamseyDiscrete, x)
+    pre_α, pre_β, pre_δ = x
+    (; a = logistic(pre_α), β = logistic(pre_β), δ = logistic(pre_δ))
+end
+
+function SK.calculate_derived_quantities(::RamseyDiscrete, model_parameters)
+    (; α, β, δ) = model_parameters
+    # f'(k) β = (α k^{α-1} - δ) β = 1
+    k_s = ((1 / β + δ) / α)^(α - 1)
+    c_s = k_s^α - δ_s * k_s - k_s
+    (; k_s, c_s)
+end

From 658c8cd0cc090625e7615cedb6594494d594572d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tam=C3=A1s=20K=2E=20Papp?= <tkpapp@gmail.com>
Date: Tue, 2 Sep 2025 14:28:04 +0200
Subject: [PATCH 02/10] start work on generic framework

---
 src/experimental.jl       | 72 +++++++++++++++++++++++++++++++++++----
 test/test_experimental.jl | 38 +++++++++++++++++----
 2 files changed, 97 insertions(+), 13 deletions(-)

diff --git a/src/experimental.jl b/src/experimental.jl
index 506e2b1..f35c30d 100644
--- a/src/experimental.jl
+++ b/src/experimental.jl
@@ -10,25 +10,44 @@ A best effort is made to
 module Experimental
 
 using Compat: @compat
+@compat public model_parameters_dimension, make_model_parameters,
+    calculate_derived_quantities, make_approximation_basis, describe_policy_transformations,
+    policy_coefficients_dimension, make_policy_functions
 
-@compat public model_parameters_dimension, make_model_parameters, calculate_derived_quantities
-
+using ArgCheck: @argcheck
 using DocStringExtensions: FUNCTIONNAME, SIGNATURES
+import ..SpectralKit
 
 ####
-#### endpoint mappings
+#### utilities
 ####
 
-# TODO
+"""
+$(SIGNATURES)
+
+FIXME docs
+"""
+function named_cumulative_ranges(lengths::NamedTuple{N}) where N
+    # NOTE this approach  relies on Base.afoldl, so it unrolled only up to 31 elements
+    ranges = accumulate(lengths, init = 0:0) do r, l
+        a = last(r)
+        (a + 1):(a + l)
+    end
+    NamedTuple{N}(ranges)
+end
 
 ####
 #### modelling API
 ####
 
+const USERNOTE = "**User should implement this method for the relevant `model_family`.**"
+
 """
 $(FUNCTIONNAME)(model_family)
 
 Dimension of the model parameters as a flat vector.
+
+$(USERNOTE)
 """
 function model_parameters_dimension end
 
@@ -36,8 +55,12 @@ function model_parameters_dimension end
 $(FUNCTIONNAME)(model_family, x::AbstractVector)
 
 Convert a flat vector of dimension [`model_parameters_dimension`)(@ref) to the model
-parameters, which can be an arbitrary value (eg a `NamedTuple`), which is not used for
-dispatch.
+parameters.
+
+$(USERNOTE)
+
+The returned value can be an arbitrary (eg a `NamedTuple`), as it does not need to be
+not used for dispatch.
 
 !!! NOTE
     The method should accept all finite numbers in ``ℝ``, and transform them accordingly.
@@ -48,7 +71,44 @@ function make_model_parameters end
 $(SIGNATURES)(model_family, model_parameters)
 
 Calculate derived quantities (for use in determining the bases and calculating the residuals).
+
+$(USERNOTE)
 """
 function calculate_derived_quantities end
 
+"""
+$(FUNCTIONNAME)(model_family, derived_quantities, approximation_scheme)
+
+Construct an approximation basis.
+
+$(USERNOTE)
+"""
+function make_approximation_basis end
+
+"""
+$(FUNCTIONNAME)(model_family)
+
+Return a `NamedTuple` is policy function approximation schemes. The values describe the
+transformations.
+
+$(USERNOTE)
+"""
+function describe_policy_transformations end
+
+function policy_coefficients_dimension(policy_transformations::NamedTuple, approximation_basis)
+    SpectralKit.dimension(approximation_basis) * length(policy_transformations)
+end
+
+function make_policy_functions(model_family, policy_transformations::NamedTuple,
+                               approximation_basis, coefficients)
+    d = SpectralKit.dimension(approximation_basis)
+    # QUESTION line below assumes all univariate, generalize?
+    ranges = named_cumulative_ranges(map(_ -> d, policy_transformations))
+    @argcheck firstindex(coefficients) == 1
+    @argcheck lastindex(coefficients) == last(last(ranges))
+    map(ranges, policy_transformations) do r, t
+        t ∘ SpectralKit.linear_combination(approximation_basis, @view coefficients[r])
+    end
+end
+
 end
diff --git a/test/test_experimental.jl b/test/test_experimental.jl
index 67388f5..7a76293 100644
--- a/test/test_experimental.jl
+++ b/test/test_experimental.jl
@@ -1,6 +1,7 @@
-import SpectralKit.Experimental as SK
-
+import SpectralKit.Experimental as SKX
 using LogExpFunctions: logistic
+using Test
+using SpectralKit
 
 """
 A simple Ramsey model in discrete time.
@@ -52,17 +53,40 @@ c_s = f(k_s) - k_s
 """
 struct RamseyDiscrete end
 
-SK.model_parameters_dimension(::RamseyDiscrete) = 3
+SKX.model_parameters_dimension(::RamseyDiscrete) = 3
 
-function SK.make_model_parameters(::RamseyDiscrete, x)
+function SKX.make_model_parameters(::RamseyDiscrete, x)
     pre_α, pre_β, pre_δ = x
-    (; a = logistic(pre_α), β = logistic(pre_β), δ = logistic(pre_δ))
+    (; α = logistic(pre_α), β = logistic(pre_β), δ = logistic(pre_δ))
 end
 
-function SK.calculate_derived_quantities(::RamseyDiscrete, model_parameters)
+product(model_parameters, k) = k^model_parameters.α  - model_parameters.δ * k
+
+function SKX.calculate_derived_quantities(::RamseyDiscrete, model_parameters)
     (; α, β, δ) = model_parameters
     # f'(k) β = (α k^{α-1} - δ) β = 1
     k_s = ((1 / β + δ) / α)^(α - 1)
-    c_s = k_s^α - δ_s * k_s - k_s
+    c_s = product(model_parameters, k_s) - k_s
     (; k_s, c_s)
 end
+
+function SKX.make_approximation_basis(::RamseyDiscrete, derived_quantities, approximation_scheme)
+    (; policy_coefficients) = approximation_scheme
+    (; k_s) = derived_quantities
+    Chebyshev(InteriorGrid(), policy_coefficients) ∘ SemiInfRational(0.0, k_s)
+end
+
+SKX.describe_policy_transformations(::RamseyDiscrete) = (; c_share = logistic)
+
+model_family = RamseyDiscrete()
+model_parameters = SKX.make_model_parameters(model_family,
+                                            randn(SKX.model_parameters_dimension(model_family)))
+derived_quantities = SKX.calculate_derived_quantities(model_family, model_parameters)
+approximation_scheme = (; policy_coefficients = 10)
+approximation_basis = SKX.make_approximation_basis(model_family, derived_quantities,
+                                                   approximation_scheme)
+policy_transformations = SKX.describe_policy_transformations(model_family)
+policy_functions = SKX.make_policy_functions(model_family, policy_transformations,
+                                             approximation_basis,
+                                             zeros(SKX.policy_coefficients_dimension(policy_transformations,
+                                                                                     approximation_basis)))

From a34b4a16ddad51ff9bccb9c57dd5fd78c3ca7060 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tam=C3=A1s=20K=2E=20Papp?= <tkpapp@gmail.com>
Date: Tue, 2 Sep 2025 14:30:20 +0200
Subject: [PATCH 03/10] fix minor bug

---
 src/generic_api.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/generic_api.jl b/src/generic_api.jl
index 190e145..1b00d76 100644
--- a/src/generic_api.jl
+++ b/src/generic_api.jl
@@ -293,10 +293,10 @@ function basis_at(basis::TransformedBasis, x)
     basis_at(parent, transform_to(domain(parent), transformation, x))
 end
 
-function grid(basis::TransformedBasis)
+function grid(::Type{T}, basis::TransformedBasis) where T
     (; parent, transformation) = basis
     d = domain(parent)
-    Iterators.map(x -> transform_from(d, transformation, x), grid(parent))
+    Iterators.map(x -> transform_from(d, transformation, x), grid(T, parent))
 end
 
 function Base.:(∘)(linear_combination::LinearCombination, transformation)

From 2a3046085b7686f2b56dbf06e7a2edac8c4d29ab Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tam=C3=A1s=20K=2E=20Papp?= <tkpapp@gmail.com>
Date: Tue, 2 Sep 2025 14:58:46 +0200
Subject: [PATCH 04/10] implement constant coefficients

---
 docs/make.jl              |  2 +-
 docs/src/experimental.md  | 11 +++++++++++
 src/experimental.jl       | 33 +++++++++++++++++++++++++++++++++
 test/test_experimental.jl | 22 ++++++++++++++++++++++
 4 files changed, 67 insertions(+), 1 deletion(-)
 create mode 100644 docs/src/experimental.md

diff --git a/docs/make.jl b/docs/make.jl
index fb1921b..bef6cd9 100644
--- a/docs/make.jl
+++ b/docs/make.jl
@@ -7,7 +7,7 @@ makedocs(
     format = Documenter.HTML(; prettyurls = get(ENV, "CI", nothing) == "true"),
     authors = "Tamas K. Papp",
     sitename = "SpectralKit.jl",
-    pages = Any["index.md"],
+    pages = Any["index.md", "experimental.md"],
     clean = true,
     checkdocs = :exports,
 )
diff --git a/docs/src/experimental.md b/docs/src/experimental.md
new file mode 100644
index 0000000..1ebeb52
--- /dev/null
+++ b/docs/src/experimental.md
@@ -0,0 +1,11 @@
+# Experimental
+
+```@docs
+SpectralKit.Experimental
+```
+
+## Proposed additions to general API
+
+```@docs
+SpectralKit.Experimental.constant_coefficients
+```
diff --git a/src/experimental.jl b/src/experimental.jl
index f35c30d..8d75f67 100644
--- a/src/experimental.jl
+++ b/src/experimental.jl
@@ -9,6 +9,9 @@ A best effort is made to
 """
 module Experimental
 
+# migrate to generic API
+export constant_coefficients
+
 using Compat: @compat
 @compat public model_parameters_dimension, make_model_parameters,
     calculate_derived_quantities, make_approximation_basis, describe_policy_transformations,
@@ -36,6 +39,36 @@ function named_cumulative_ranges(lengths::NamedTuple{N}) where N
     NamedTuple{N}(ranges)
 end
 
+####
+#### generic API additions
+####
+
+"""
+$(FUNCTIONNAME)(basis, y)
+
+Approximate a constant value `y` on `basis`.
+
+Formally, return a set of coefficients `θ` such that `linear_combination(basis, θ, x) ≈
+y` for all `x` in the domain.
+"""
+function constant_coefficients end
+
+function constant_coefficients(basis::SpectralKit.Chebyshev, y)
+    θ = zeros(basis.N)
+    θ[1] = y
+    θ
+end
+
+function constant_coefficients(basis::SpectralKit.TransformedBasis, y)
+    constant_coefficients(parent(basis), y)
+end
+
+function constant_coefficients(basis::SpectralKit.SmolyakBasis, y)
+    θ = zeros(SpectralKit.dimension(basis))
+    θ[1] = y
+    θ
+end
+
 ####
 #### modelling API
 ####
diff --git a/test/test_experimental.jl b/test/test_experimental.jl
index 7a76293..417b81b 100644
--- a/test/test_experimental.jl
+++ b/test/test_experimental.jl
@@ -3,6 +3,28 @@ using LogExpFunctions: logistic
 using Test
 using SpectralKit
 
+####
+#### generic api
+####
+
+@testset "constant coefficients" begin
+    y = 3.8
+    b1 = Chebyshev(InteriorGrid(), 20)
+    t2 = InfRational(3.8, 1.3)
+    t3 = SemiInfRational(0.7, 1.9)
+    b2 = b1 ∘ t2
+    b3 = b1 ∘ t3
+    B1 = smolyak_basis(Chebyshev, InteriorGrid(), SmolyakParameters(3), Val(2))
+    B2 = B1 ∘ coordinate_transformations(t2, t3)
+    for basis in [b1, b2, b3, B1, B2]
+        θ = SKX.constant_coefficients(basis, y)
+        for _ in 1:100
+            x = rand_in_domain(basis)
+            @test linear_combination(basis, θ, x) ≈ y
+        end
+    end
+end
+
 """
 A simple Ramsey model in discrete time.
 

From 5f92e09298478d7f9623ac3f73caa4a08f0ef996 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tam=C3=A1s=20K=2E=20Papp?= <tkpapp@gmail.com>
Date: Wed, 3 Sep 2025 09:14:01 +0200
Subject: [PATCH 05/10] add constant solutions, initial guess

---
 Project.toml              |  2 ++
 src/experimental.jl       | 41 +++++++++++++++++++++++++++++++++++++--
 test/test_experimental.jl | 11 +++++++++--
 3 files changed, 50 insertions(+), 4 deletions(-)

diff --git a/Project.toml b/Project.toml
index 12b2440..d3e8a6d 100644
--- a/Project.toml
+++ b/Project.toml
@@ -7,6 +7,7 @@ authors = ["Tamas K. Papp <tkpapp@gmail.com>"]
 ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
+InverseFunctions = "3587e190-3f89-42d0-90ee-14403ec27112"
 OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 
@@ -14,6 +15,7 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 ArgCheck = "1, 2"
 Compat = "4.18.0"
 DocStringExtensions = "0.8, 0.9"
+InverseFunctions = "0.1"
 OrderedCollections = "1"
 StaticArrays = "1"
 julia = "1.10"
diff --git a/src/experimental.jl b/src/experimental.jl
index 8d75f67..6438ea3 100644
--- a/src/experimental.jl
+++ b/src/experimental.jl
@@ -15,11 +15,14 @@ export constant_coefficients
 using Compat: @compat
 @compat public model_parameters_dimension, make_model_parameters,
     calculate_derived_quantities, make_approximation_basis, describe_policy_transformations,
-    policy_coefficients_dimension, make_policy_functions
+    policy_coefficients_dimension, make_policy_functions, constant_initial_guess,
+    calculate_initial_guess
+
+import ..SpectralKit
 
 using ArgCheck: @argcheck
 using DocStringExtensions: FUNCTIONNAME, SIGNATURES
-import ..SpectralKit
+using InverseFunctions: inverse
 
 ####
 #### utilities
@@ -144,4 +147,38 @@ function make_policy_functions(model_family, policy_transformations::NamedTuple,
     end
 end
 
+"""
+$(FUNCTIONNAME)(model_family, derived_quantities)
+
+Return initial guesses in a type that supports `getproperty` (eg a `NamedTuple`).
+
+Should provide a scalar for each name in [`describe_policy_transformations`](@ref) (can
+be in any arbitrary order, and contain other names).
+"""
+function constant_initial_guess end
+
+"""
+$(SIGNATURES)
+
+Return an initial guess for the coefficients. Falls back to using
+[`constant_initial_guess`](@ref), or the user may define a method in case that is not
+sufficient.
+"""
+function calculate_initial_guess(model_family, derived_quantities,
+                                 policy_transformations::NamedTuple{N},
+                                 approximation_basis) where N
+    constant_guess = constant_initial_guess(model_family, derived_quantities)
+    d = SpectralKit.dimension(approximation_basis)
+    # QUESTION line below assumes all univariate, generalize?
+    ranges = named_cumulative_ranges(map(_ -> d, policy_transformations))
+    coefficients = zeros(last(last(ranges)))
+    for (name, transformation) in pairs(policy_transformations)
+        r = getproperty(ranges, name)
+        y = inverse(transformation)(getproperty(constant_guess, name))
+        # FIXME a constant_coefficient! API would have fewer allocations
+        coefficients[r] .= constant_coefficients(approximation_basis, y)
+    end
+    coefficients
+end
+
 end
diff --git a/test/test_experimental.jl b/test/test_experimental.jl
index 417b81b..69552f8 100644
--- a/test/test_experimental.jl
+++ b/test/test_experimental.jl
@@ -82,12 +82,12 @@ function SKX.make_model_parameters(::RamseyDiscrete, x)
     (; α = logistic(pre_α), β = logistic(pre_β), δ = logistic(pre_δ))
 end
 
-product(model_parameters, k) = k^model_parameters.α  - model_parameters.δ * k
+product(model_parameters, k) = k^model_parameters.α + (1 - model_parameters.δ) * k
 
 function SKX.calculate_derived_quantities(::RamseyDiscrete, model_parameters)
     (; α, β, δ) = model_parameters
     # f'(k) β = (α k^{α-1} - δ) β = 1
-    k_s = ((1 / β + δ) / α)^(α - 1)
+    k_s = ((1 / β + 1 - δ) / α)^(α - 1)
     c_s = product(model_parameters, k_s) - k_s
     (; k_s, c_s)
 end
@@ -100,6 +100,11 @@ end
 
 SKX.describe_policy_transformations(::RamseyDiscrete) = (; c_share = logistic)
 
+function SKX.constant_initial_guess(::RamseyDiscrete, derived_quantities)
+    (; k_s, c_s) = derived_quantities
+    (; c_share = c_s / k_s)
+end
+
 model_family = RamseyDiscrete()
 model_parameters = SKX.make_model_parameters(model_family,
                                             randn(SKX.model_parameters_dimension(model_family)))
@@ -112,3 +117,5 @@ policy_functions = SKX.make_policy_functions(model_family, policy_transformation
                                              approximation_basis,
                                              zeros(SKX.policy_coefficients_dimension(policy_transformations,
                                                                                      approximation_basis)))
+θ0 = SKX.calculate_initial_guess(model_family, derived_quantities, policy_transformations,
+                                 approximation_basis)

From fc051486aa78ad3aab7df1fef60472e15e36b7d6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tam=C3=A1s=20K=2E=20Papp?= <tkpapp@gmail.com>
Date: Wed, 3 Sep 2025 09:30:49 +0200
Subject: [PATCH 06/10] work around Heisenbug

---
 src/experimental.jl | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/experimental.jl b/src/experimental.jl
index 6438ea3..a1a868c 100644
--- a/src/experimental.jl
+++ b/src/experimental.jl
@@ -174,7 +174,8 @@ function calculate_initial_guess(model_family, derived_quantities,
     coefficients = zeros(last(last(ranges)))
     for (name, transformation) in pairs(policy_transformations)
         r = getproperty(ranges, name)
-        y = inverse(transformation)(getproperty(constant_guess, name))
+        transformed_y = getproperty(constant_guess, name)
+        y = inverse(transformation)(transformed_y)
         # FIXME a constant_coefficient! API would have fewer allocations
         coefficients[r] .= constant_coefficients(approximation_basis, y)
     end

From 86aa4c7f10eeb6e58bc3c55bab265effcd0f31a0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tam=C3=A1s=20K=2E=20Papp?= <tkpapp@gmail.com>
Date: Thu, 4 Sep 2025 15:11:32 +0200
Subject: [PATCH 07/10] algebra corrections, clarifications, residual
 calculations

---
 src/experimental.jl       | 15 ++++++++++---
 test/test_experimental.jl | 47 +++++++++++++++++++++++++++++----------
 2 files changed, 47 insertions(+), 15 deletions(-)

diff --git a/src/experimental.jl b/src/experimental.jl
index a1a868c..270cc9a 100644
--- a/src/experimental.jl
+++ b/src/experimental.jl
@@ -131,6 +131,15 @@ $(USERNOTE)
 """
 function describe_policy_transformations end
 
+"""
+$(FUNCTIONNAME)(model_family, model_parameters, policy_functions, gridpoint)
+
+Calculate the residuals at the given gridpoint.
+
+$(USERNOTE)
+"""
+function calculate_residuals end
+
 function policy_coefficients_dimension(policy_transformations::NamedTuple, approximation_basis)
     SpectralKit.dimension(approximation_basis) * length(policy_transformations)
 end
@@ -148,7 +157,7 @@ function make_policy_functions(model_family, policy_transformations::NamedTuple,
 end
 
 """
-$(FUNCTIONNAME)(model_family, derived_quantities)
+$(FUNCTIONNAME)(model_family, model_parameters, derived_quantities)
 
 Return initial guesses in a type that supports `getproperty` (eg a `NamedTuple`).
 
@@ -164,10 +173,10 @@ Return an initial guess for the coefficients. Falls back to using
 [`constant_initial_guess`](@ref), or the user may define a method in case that is not
 sufficient.
 """
-function calculate_initial_guess(model_family, derived_quantities,
+function calculate_initial_guess(model_family, model_parameters, derived_quantities,
                                  policy_transformations::NamedTuple{N},
                                  approximation_basis) where N
-    constant_guess = constant_initial_guess(model_family, derived_quantities)
+    constant_guess = constant_initial_guess(model_family, model_parameters, derived_quantities)
     d = SpectralKit.dimension(approximation_basis)
     # QUESTION line below assumes all univariate, generalize?
     ranges = named_cumulative_ranges(map(_ -> d, policy_transformations))
diff --git a/test/test_experimental.jl b/test/test_experimental.jl
index 69552f8..cf3e801 100644
--- a/test/test_experimental.jl
+++ b/test/test_experimental.jl
@@ -58,11 +58,11 @@ V'(k) = u'(c(k)) f'(k)
 ```
 and then
 ```math
-u'(c(k)) f'(k) = β u'(c(f(k) - c(k))) f'(c(f(k) - c(k))) f'(k)
+u'(c(k)) f'(k) = β u'(c(f(k) - c(k))) f'(f(k) - c(k)) f'(k)
 ```
 Then cancel ``f'(k)`` and rearrange as unitless
 ```math
-1 = β u'(c(f(k) - c(k))) f'(c(f(k)) - c(k)) / u'(k)
+1 = β u'(c(f(k) - c(k))) f'(f(k) - c(k)) / u'(c(k))
 ```
 From this it follows that the steady state ``k_s`` solves
 ```
@@ -79,15 +79,21 @@ SKX.model_parameters_dimension(::RamseyDiscrete) = 3
 
 function SKX.make_model_parameters(::RamseyDiscrete, x)
     pre_α, pre_β, pre_δ = x
+    # NOTE: we are assuming log utility, doesn't need a parameter
     (; α = logistic(pre_α), β = logistic(pre_β), δ = logistic(pre_δ))
 end
 
 product(model_parameters, k) = k^model_parameters.α + (1 - model_parameters.δ) * k
 
+function marginal_product(model_parameters, k)
+    (; α, δ) =model_parameters
+    α * k^(α - 1) + (1 - model_parameters.δ)
+end
+
 function SKX.calculate_derived_quantities(::RamseyDiscrete, model_parameters)
     (; α, β, δ) = model_parameters
-    # f'(k) β = (α k^{α-1} - δ) β = 1
-    k_s = ((1 / β + 1 - δ) / α)^(α - 1)
+    # we are solving f'(k) β = (α k^{α-1} + 1 - δ) β = 1
+    k_s = ((1 / β - 1 + δ) / α)^(1/(α - 1))
     c_s = product(model_parameters, k_s) - k_s
     (; k_s, c_s)
 end
@@ -100,22 +106,39 @@ end
 
 SKX.describe_policy_transformations(::RamseyDiscrete) = (; c_share = logistic)
 
-function SKX.constant_initial_guess(::RamseyDiscrete, derived_quantities)
+function SKX.constant_initial_guess(::RamseyDiscrete, model_parameters, derived_quantities)
     (; k_s, c_s) = derived_quantities
-    (; c_share = c_s / k_s)
+    (; c_share = c_s / product(model_parameters, k_s))
+end
+
+function SKX.calculate_residuals(::RamseyDiscrete, model_parameters, policy_functions, k0)
+    (; β) = model_parameters
+    (; c_share) = policy_functions
+    y0 = product(model_parameters, k0)
+    c0 = c_share(k0) * y0
+    k1 = y0 - c0
+    c1 = c_share(k1) * product(model_parameters, k1)
+    # assuming log utility
+    (; Euler = β * c0/c1 * marginal_product(model_parameters, k1) - 1)
 end
 
 model_family = RamseyDiscrete()
 model_parameters = SKX.make_model_parameters(model_family,
-                                            randn(SKX.model_parameters_dimension(model_family)))
+                                             randn(SKX.model_parameters_dimension(model_family)))
 derived_quantities = SKX.calculate_derived_quantities(model_family, model_parameters)
+let
+    (; k_s, c_s) = derived_quantities
+    (; β, α, δ) = model_parameters
+    @test marginal_product(model_parameters, k_s) * model_parameters.β ≈ 1
+    @test product(model_parameters, k_s) - c_s ≈ k_s
+end
 approximation_scheme = (; policy_coefficients = 10)
 approximation_basis = SKX.make_approximation_basis(model_family, derived_quantities,
                                                    approximation_scheme)
 policy_transformations = SKX.describe_policy_transformations(model_family)
+θ0 = SKX.calculate_initial_guess(model_family, model_parameters, derived_quantities,
+                                 policy_transformations, approximation_basis)
 policy_functions = SKX.make_policy_functions(model_family, policy_transformations,
-                                             approximation_basis,
-                                             zeros(SKX.policy_coefficients_dimension(policy_transformations,
-                                                                                     approximation_basis)))
-θ0 = SKX.calculate_initial_guess(model_family, derived_quantities, policy_transformations,
-                                 approximation_basis)
+                                             approximation_basis, θ0)
+@test SKX.calculate_residuals(model_family, model_parameters, policy_functions,
+                              derived_quantities.k_s).Euler ≈ 0 atol = 1e-8

From da7dd6ab7a8e80d3d400caf3a51475467945f1ff Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tam=C3=A1s=20K=2E=20Papp?= <tkpapp@gmail.com>
Date: Tue, 9 Sep 2025 11:52:30 +0200
Subject: [PATCH 08/10] add grid and sum of squared residuals

---
 src/experimental.jl       | 41 ++++++++++++++++++++++++++++++++++++++-
 test/Project.toml         |  4 ++++
 test/test_experimental.jl |  5 +++++
 3 files changed, 49 insertions(+), 1 deletion(-)

diff --git a/src/experimental.jl b/src/experimental.jl
index 270cc9a..611fa3c 100644
--- a/src/experimental.jl
+++ b/src/experimental.jl
@@ -9,6 +9,13 @@ A best effort is made to
 """
 module Experimental
 
+# TODO
+# - [ ] clean up API, too many functions with positional arguments
+# - [ ] write up tutorial-like docs
+# - [ ] docs: an ascii diagram of various calculations, with https://asciiflow.com/
+# - [ ] solver functions
+# - [ ] actually implement a solution and AD it
+
 # migrate to generic API
 export constant_coefficients
 
@@ -16,7 +23,7 @@ using Compat: @compat
 @compat public model_parameters_dimension, make_model_parameters,
     calculate_derived_quantities, make_approximation_basis, describe_policy_transformations,
     policy_coefficients_dimension, make_policy_functions, constant_initial_guess,
-    calculate_initial_guess
+    calculate_initial_guess, sum_of_squared_residuals
 
 import ..SpectralKit
 
@@ -191,4 +198,36 @@ function calculate_initial_guess(model_family, model_parameters, derived_quantit
     coefficients
 end
 
+"""
+$(SIGNATURES)
+
+Return a grid for calculating the residuals. The default implementation just uses the
+grid that corresponds to the approximation basis.
+"""
+function make_approximation_grid(model_family, model_parameters, derived_quantities,
+                                 approximation_basis)
+    SpectralKit.grid(approximation_basis)
+end
+
+"""
+$(SIGNATURES)
+
+Sum of squares. Works for values returned by [`calculate_residuals`](@ref).
+"""
+sum_abs2(x::Real) = abs2(x)
+
+sum_abs2(nt::NamedTuple) = sum(abs2, values(nt))
+
+"""
+$(SIGNATURES)
+
+Calculate the sum of squared residuals.
+"""
+function sum_of_squared_residuals(model_family, model_parameters, policy_functions, grid)
+    mapreduce(+, grid) do gridpoint
+        sum_abs2(calculate_residuals(model_family, model_parameters, policy_functions,
+                                     gridpoint))
+    end
+end
+
 end
diff --git a/test/Project.toml b/test/Project.toml
index 8f8be9c..91e4688 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -7,5 +7,9 @@ JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Sobol = "ed01d8cd-4d21-5b2a-85b4-cc3bdc58bad4"
+SpectralKit = "5c252ae7-b5b6-46ab-a016-b0e3d78320b7"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+
+[sources]
+SpectralKit = {path = ".."}
diff --git a/test/test_experimental.jl b/test/test_experimental.jl
index cf3e801..a13ed64 100644
--- a/test/test_experimental.jl
+++ b/test/test_experimental.jl
@@ -142,3 +142,8 @@ policy_functions = SKX.make_policy_functions(model_family, policy_transformation
                                              approximation_basis, θ0)
 @test SKX.calculate_residuals(model_family, model_parameters, policy_functions,
                               derived_quantities.k_s).Euler ≈ 0 atol = 1e-8
+grid = SKX.make_approximation_grid(model_family, model_parameters, derived_quantities,
+                                   approximation_basis)
+
+@test @inferred SKX.sum_of_squared_residuals(model_family, model_parameters, policy_functions,
+                                             grid) isa Float64

From d87768a5b7ba01e7943c045a0acd0a9d8c24300a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tam=C3=A1s=20K=2E=20Papp?= <tkpapp@gmail.com>
Date: Tue, 9 Sep 2025 11:56:31 +0200
Subject: [PATCH 09/10] fix variable names

---
 test/test_experimental.jl | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/test/test_experimental.jl b/test/test_experimental.jl
index a13ed64..9afc74c 100644
--- a/test/test_experimental.jl
+++ b/test/test_experimental.jl
@@ -142,8 +142,8 @@ policy_functions = SKX.make_policy_functions(model_family, policy_transformation
                                              approximation_basis, θ0)
 @test SKX.calculate_residuals(model_family, model_parameters, policy_functions,
                               derived_quantities.k_s).Euler ≈ 0 atol = 1e-8
-grid = SKX.make_approximation_grid(model_family, model_parameters, derived_quantities,
-                                   approximation_basis)
+approximation_grid = SKX.make_approximation_grid(model_family, model_parameters,
+                                                 derived_quantities, approximation_basis)
 
 @test @inferred SKX.sum_of_squared_residuals(model_family, model_parameters, policy_functions,
-                                             grid) isa Float64
+                                             approximation_grid) isa Float64

From 0043309f1d53c45656abb25673390f8857a33426 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tam=C3=A1s=20K=2E=20Papp?= <tkpapp@gmail.com>
Date: Thu, 15 Jan 2026 10:50:19 +0100
Subject: [PATCH 10/10] slight renaming, fix issues after JET

---
 src/derivatives.jl        |  4 ++--
 src/experimental.jl       |  6 +++---
 src/transformations.jl    |  4 ++--
 test/runtests.jl          | 27 +++++++++++++++------------
 test/test_experimental.jl |  9 +++++----
 5 files changed, 27 insertions(+), 23 deletions(-)

diff --git a/src/derivatives.jl b/src/derivatives.jl
index 58962e8..b4e6b26 100644
--- a/src/derivatives.jl
+++ b/src/derivatives.jl
@@ -510,9 +510,9 @@ end
     end
 end
 
-function _product_type(::Type{D}, source_eltypes) where {Ps,D<:∂Derivatives{Ps}}
+function _product_type(::Type{D}, source_eltypes::NTuple{M}) where {Ps,D<:∂Derivatives{Ps},M}
     T = _product_type(Nothing, map(eltype, source_eltypes))
-    N = length(_partials_canonical_expansion(Val(N), fieldtypes(Ps)))
+    N = length(_partials_canonical_expansion(Val(M), fieldtypes(Ps)))
     ∂Expansion{D,N,T}
 end
 
diff --git a/src/experimental.jl b/src/experimental.jl
index 611fa3c..d39e168 100644
--- a/src/experimental.jl
+++ b/src/experimental.jl
@@ -120,7 +120,7 @@ $(USERNOTE)
 function calculate_derived_quantities end
 
 """
-$(FUNCTIONNAME)(model_family, derived_quantities, approximation_scheme)
+$(FUNCTIONNAME)(model_family, derived_quantities, approximation_parameters)
 
 Construct an approximation basis.
 
@@ -204,8 +204,8 @@ $(SIGNATURES)
 Return a grid for calculating the residuals. The default implementation just uses the
 grid that corresponds to the approximation basis.
 """
-function make_approximation_grid(model_family, model_parameters, derived_quantities,
-                                 approximation_basis)
+function make_approximation_grid(model_family, model_parameters, approximation_parameters,
+                                 derived_quantities, approximation_basis)
     SpectralKit.grid(approximation_basis)
 end
 
diff --git a/src/transformations.jl b/src/transformations.jl
index 2a338aa..e535d97 100644
--- a/src/transformations.jl
+++ b/src/transformations.jl
@@ -123,7 +123,7 @@ end
 
 function transform_to(domain::CoordinateDomains, ct::CoordinateTransformations,
                       Dx::∂CoordinateExpansion)
-    ∂CoordinateExpansion(x.∂D, transform_to(domain, ct, x.x))
+    ∂CoordinateExpansion(Dx.∂D, transform_to(domain, ct, Dx.x))
 end
 
 function transform_from(domain::CoordinateDomains, ct::CoordinateTransformations, x::Tuple)
@@ -323,7 +323,7 @@ function transform_to(domain::PM1, t::InfRational, y::𝑑Expansion{Dp1}) where
     (; A, L) = t
     (; coefficients) = y
     x0 = transform_to(domain, t, coefficients[1])
-    Dp1 == 1 && return Coefficients(SVector(x0))
+    Dp1 == 1 && return SVector(x0)
     # based on Boyd (2001), Table E.5
     Q = 1 - abs2(x0)
     sQ = √Q
diff --git a/test/runtests.jl b/test/runtests.jl
index 0bd0d81..d119afe 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1,5 +1,18 @@
-using SpectralKit
-using Test, DocStringExtensions, StaticArrays, BenchmarkTools, FiniteDifferences
+using SpectralKit, Test
+
+using JET
+@testset "static analysis with JET.jl" begin
+    @test isempty(JET.get_reports(report_package(SpectralKit,
+                                                 target_modules=(SpectralKit,),
+                                                 ignored_modules = (SpectralKit.Experimental,))))
+end
+
+@testset "QA with Aqua" begin
+    import Aqua
+    Aqua.test_all(SpectralKit)
+end
+
+using DocStringExtensions, StaticArrays, BenchmarkTools, FiniteDifferences
 
 include("utilities.jl")
 
@@ -12,13 +25,3 @@ include("test_smolyak_traversal.jl")
 include("test_smolyak.jl")
 include("test_generic_api.jl")  # NOTE moved last as it used constructs from above
 include("test_experimental.jl") # NOTE experimental code is not public API
-
-using JET
-@testset "static analysis with JET.jl" begin
-    @test isempty(JET.get_reports(report_package(SpectralKit, target_modules=(SpectralKit,))))
-end
-
-@testset "QA with Aqua" begin
-    import Aqua
-    Aqua.test_all(SpectralKit)
-end
diff --git a/test/test_experimental.jl b/test/test_experimental.jl
index 9afc74c..a86da41 100644
--- a/test/test_experimental.jl
+++ b/test/test_experimental.jl
@@ -98,8 +98,8 @@ function SKX.calculate_derived_quantities(::RamseyDiscrete, model_parameters)
     (; k_s, c_s)
 end
 
-function SKX.make_approximation_basis(::RamseyDiscrete, derived_quantities, approximation_scheme)
-    (; policy_coefficients) = approximation_scheme
+function SKX.make_approximation_basis(::RamseyDiscrete, derived_quantities, approximation_parameters)
+    (; policy_coefficients) = approximation_parameters
     (; k_s) = derived_quantities
     Chebyshev(InteriorGrid(), policy_coefficients) ∘ SemiInfRational(0.0, k_s)
 end
@@ -132,9 +132,9 @@ let
     @test marginal_product(model_parameters, k_s) * model_parameters.β ≈ 1
     @test product(model_parameters, k_s) - c_s ≈ k_s
 end
-approximation_scheme = (; policy_coefficients = 10)
+approximation_parameters = (; policy_coefficients = 10)
 approximation_basis = SKX.make_approximation_basis(model_family, derived_quantities,
-                                                   approximation_scheme)
+                                                   approximation_parameters)
 policy_transformations = SKX.describe_policy_transformations(model_family)
 θ0 = SKX.calculate_initial_guess(model_family, model_parameters, derived_quantities,
                                  policy_transformations, approximation_basis)
@@ -143,6 +143,7 @@ policy_functions = SKX.make_policy_functions(model_family, policy_transformation
 @test SKX.calculate_residuals(model_family, model_parameters, policy_functions,
                               derived_quantities.k_s).Euler ≈ 0 atol = 1e-8
 approximation_grid = SKX.make_approximation_grid(model_family, model_parameters,
+                                                 approximation_parameters,
                                                  derived_quantities, approximation_basis)
 
 @test @inferred SKX.sum_of_squared_residuals(model_family, model_parameters, policy_functions,