generalize kwarg parsing in @onsite!/@hopping! macros

streamline macros complete test
pablosanjose · Apr 23, 2024 · 481963f · 481963f
1 parent 55f1659
commit 481963f
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 35 deletions.
diff --git a/src/models.jl b/src/models.jl
@@ -40,60 +40,58 @@ _hopping(t::ParametricHoppingTerm; kw...) =
 
 ## Parametric models ##
 
-# version with site selector kwargs
-macro onsite(kw, f)
-    f, N, params, spatial = parse_term(f, "Only @onsite(args -> body; kw...) syntax supported (or with -->). Mind the `;`.")
+macro onsite(x, ys...)
+    kw, f, N, params, spatial = parse_term("@onsite", x, ys...)
     return esc(:(Quantica.ParametricModel(Quantica.ParametricOnsiteTerm(
         Quantica.ParametricFunction{$N}($f, $(params)), Quantica.siteselector($kw), 1, $(spatial)))))
 end
 
-# version without site selector kwargs
-macro onsite(f)
-    f, N, params, spatial = parse_term(f, "Only @onsite(args -> body; kw...) syntax supported (or with -->). Mind the `;`.")
-    return esc(:(Quantica.ParametricModel(Quantica.ParametricOnsiteTerm(
-            Quantica.ParametricFunction{$N}($f, $(params)), Quantica.siteselector(), 1, $(spatial)))))
-end
-
-# version with hop selector kwargs
-macro hopping(kw, f)
-    f, N, params, spatial = parse_term(f, "Only @hopping(args -> body; kw...) syntax supported (or with -->). Mind the `;`.")
+macro hopping(x, ys...)
+    kw, f, N, params, spatial = parse_term("@hopping", x, ys...)
     return esc(:(Quantica.ParametricModel(Quantica.ParametricHoppingTerm(
         Quantica.ParametricFunction{$N}($f, $(params)), Quantica.hopselector($kw), 1, $(spatial)))))
 end
 
-# version without hop selector kwargs
-macro hopping(f)
-    f, N, params, spatial = parse_term(f, "Only @hopping(args -> body; kw...) syntax supported (or with -->). Mind the `;`.")
-    return esc(:(Quantica.ParametricModel(Quantica.ParametricHoppingTerm(
-        Quantica.ParametricFunction{$N}($f, $(params)), Quantica.hopselector(), 1, $(spatial)))))
-end
-
 ## Model modifiers ##
 
-macro onsite!(kw, f)
-    f, N, params, spatial = parse_term(f, "Only @onsite!(args -> body; kw...) syntax supported (or with -->). Mind the `;`.")
+macro onsite!(x, ys...)
+    kw, f, N, params, spatial = parse_term("@onsite!", x, ys...)
     return esc(:(Quantica.OnsiteModifier(Quantica.ParametricFunction{$N}($f, $(params)), Quantica.siteselector($kw), $(spatial))))
 end
 
-macro onsite!(f)
-    f, N, params, spatial = parse_term(f, "Only @onsite!(args -> body; kw...) syntax supported (or with -->).  Mind the `;`.")
-    return esc(:(Quantica.OnsiteModifier(Quantica.ParametricFunction{$N}($f, $(params)), Quantica.siteselector(), $(spatial))))
+macro hopping!(x, ys...)
+    kw, f, N, params, spatial = parse_term("@hopping!", x, ys...)
+    # Since the default hopping range is neighbors(1), we need change the default to Inf for @hopping!
+    return esc(:(Quantica.HoppingModifier(Quantica.ParametricFunction{$N}($f, $(params)), Quantica.hopselector_infrange($kw), $(spatial))))
 end
 
-# Since the default hopping range is neighbors(1), we need change the default to Inf for @hopping!
-macro hopping!(kw, f)
-    f, N, params, spatial = parse_term(f, "Only @hopping!(args -> body; kw...) syntax supported (or with -->). Mind the `;`.")
-    return esc(:(Quantica.HoppingModifier(Quantica.ParametricFunction{$N}($f, $(params)), Quantica.hopselector_infrange($kw), $(spatial))))
+function parse_term(macroname, x, ys...)
+    if x isa Expr && x.head == :parameters
+        kw = x
+        f, N, params, spatial = parse_term_body(macroname, only(ys))
+    else
+        kw = parse_term_parameters(ys...)
+        f, N, params, spatial = parse_term_body(macroname, x)
+    end
+    return kw, f, N, params, spatial
 end
 
-macro hopping!(f)
-    f, N, params, spatial = parse_term(f, "Only @hopping!(args -> body; kw...) syntax supported (or with -->). Mind the `;`.")
-    return esc(:(Quantica.HoppingModifier(Quantica.ParametricFunction{$N}($f, $(params)), Quantica.hopselector_infrange(), $(spatial))))
+# parse keywords after a comma as if they were keyword arguments
+function parse_term_parameters(exs...)
+    exs´ = maybe_kw.(exs)
+    paramex = Expr(:parameters, exs´...)
+    return paramex
 end
 
+maybe_kw(ex::Expr) = Expr(:kw, ex.args...)
+maybe_kw(ex::Symbol) = ex
+
 # Extracts normalized f, number of arguments and kwarg names from an anonymous function f
-function parse_term(f, msg)
-    (f isa Expr && (f.head == :-> || f.head == :-->)) || throw(ArgumentError(msg))
+function parse_term_body(macroname, f)
+    if !(f isa Expr && (f.head == :-> || f.head == :-->))
+        msg = "Only $(macroname)(args -> body; kw...) syntax supported (or with -->). Received $(macroname)($f, ...) instead."
+        throw(ArgumentError(msg))
+    end
     # change --> to -> and record change in spatial
     spatial = f.head == :->
     !spatial && (f.head = :->)

diff --git a/test/test_hamiltonian.jl b/test/test_hamiltonian.jl
@@ -316,7 +316,7 @@ end
     @test h((π/2, -π/2), λ=1) ≈ [4 1; 1 -4]
     # Non-numeric parameters
     @test h((π/2, -π/2); λ = 1, k = SA[1,0]) ≈ [-4 1; 1 4]
-    # # Issue 61, type stability
+    # # Issue #61, type stability
     # h = LatticePresets.honeycomb() |> hamiltonian(onsite(0))
     # @inferred hamiltonian(h, @onsite!((o;μ) -> o- μ))
     # @inferred hamiltonian(h, @onsite!(o->2o), @hopping!((t)->2t), @onsite!((o, r)->o+r[1]))
@@ -347,6 +347,13 @@ end
     @test iszero(h0())
     h0 = LP.square() |> hopping(0) |> supercell(3) |> @hopping!((t, r, dr) -> 1; dcells = SVector{2,Int}[])
     @test iszero(h0())
+    # Issue #118
+    sublats = :A;
+    h0 = LP.honeycomb() |>
+        @hopping(r->0, sublats, range = 2) + @onsite((r; p = 3) ->3p; sublats) + @onsite((r; q = 3) ->3q, sublats, region = RP.circle(2)) |>
+        @hopping!((t, r, dr; p = 1) -> p+r[2], dcells = SVector{2,Int}[]) |> @onsite!((o, r; q = 1) -> o + q, sublats, region = RP.circle(3))
+    @test h0 isa ParametricHamiltonian
+    @test Quantica.parameters(h0) == [:p, :q]
 end
 
 @testset "ExternalPresets.wannier90" begin