diff --git a/Project.toml b/Project.toml
index 9b282cd..ffb61b8 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "NonlinearNormalForm"
 uuid = "05e19671-dec8-4f15-984f-54eaa6ca64be"
 authors = ["Matt Signorelli"]
-version = "0.3.3"
+version = "0.4.0"
 
 [deps]
 DelimitedFiles = "8bb1440f-4735-579b-a4ab-409b98df4dab"
@@ -26,7 +26,7 @@ Reexport = "1"
 ReferenceFrameRotations = "3"
 SkewLinearAlgebra = "1"
 StaticArrays = "1"
-TPSAInterface = "0.2.1"
+TPSAInterface = "0.2.2"
 julia = "1.9"
 
 [extras]
diff --git a/src/NonlinearNormalForm.jl b/src/NonlinearNormalForm.jl
index 0b37ba7..d799706 100644
--- a/src/NonlinearNormalForm.jl
+++ b/src/NonlinearNormalForm.jl
@@ -43,7 +43,6 @@ export        TaylorMap,
               VectorField,
               
               norm,
-              to_SO3,
       
               compose,
       
@@ -60,27 +59,18 @@ export        TaylorMap,
               normalize_eigenmode!,
               locate_modes!,
               moveback_unstable!,
-      
-              normal!,
+
               normal,
-              gofix!,
-              gofix,
-              linear_a!,
-              linear_a,
               c_map,
               ci_map,
               c_jacobian,
               ci_jacobian,
 
               log!,
-              pb,
-
-
 
-              inv_with_log,
               equilibrium_moments,
               factorize,
-              fast_canonize,
+              canonize,
               compute_de_moivre,
               compute_sagan_rubin
 
diff --git a/src/de_moivre.jl b/src/de_moivre.jl
index 9cfaf7c..162a46e 100644
--- a/src/de_moivre.jl
+++ b/src/de_moivre.jl
@@ -1,45 +1,45 @@
-function compute_de_moivre(a1::DAMap{V0}, ::Val{linear}=Val{false}()) where {V0<:StaticArray,linear}
+function compute_de_moivre(a1::DAMap{V0}, ::Val{linear}=Val{true}()) where {V0<:StaticArray,linear}
   # jp_mat[i] in FPP is J matrix restricted to i-th plane
   # ip_mat[i] in FPP is identity matrix restricted to i-th plane
   # jt_mat in FPP is symplectic s matrix
-  nv = isodd(length(V0)) ? length(V0)+1 : length(V0)
+  nhv = nhvars(a1)
   if linear
-    let a1_mat = jacobian(a1, VARS_CPARAM), a1i_mat = inv(a1_mat)
-      B = StaticArrays.sacollect(SVector{Int(nv/2),typeof(a1_mat)}, a1_mat*jp_mat(a1, i)*a1i_mat for i in 1:Int(nv/2))
-      K = StaticArrays.sacollect(SVector{Int(nv/2),typeof(a1_mat)}, -j_mat(a1)*B[i] for i in 1:Int(nv/2))
-      E = StaticArrays.sacollect(SVector{Int(nv/2),typeof(a1_mat)}, -B[i]*j_mat(a1) for i in 1:Int(nv/2))
-      H = StaticArrays.sacollect(SVector{Int(nv/2),typeof(a1_mat)}, a1_mat*ip_mat(a1, i)*a1i_mat for i in 1:Int(nv/2))
+    let a1_mat = jacobian(a1, HVARS), a1i_mat = inv(a1_mat)
+      B = StaticArrays.sacollect(SVector{Int(nhv/2),typeof(a1_mat)}, a1_mat*jp_mat(a1, i)*a1i_mat for i in 1:Int(nhv/2))
+      K = StaticArrays.sacollect(SVector{Int(nhv/2),typeof(a1_mat)}, -j_mat(a1)*B[i] for i in 1:Int(nhv/2))
+      E = StaticArrays.sacollect(SVector{Int(nhv/2),typeof(a1_mat)}, -B[i]*j_mat(a1) for i in 1:Int(nhv/2))
+      H = StaticArrays.sacollect(SVector{Int(nhv/2),typeof(a1_mat)}, a1_mat*ip_mat(a1, i)*a1i_mat for i in 1:Int(nhv/2))
       return (; H=H, B=B, E=E, K=K)
     end
   else
     let
       tmp = zero(a1)
       a1i = inv(a1)
-      B = StaticArrays.sacollect(SVector{Int(nv/2)}, begin
+      B = StaticArrays.sacollect(SVector{Int(nhv/2)}, begin
         setray!(tmp.v, v_matrix=jp_mat(a1, i))
         a1∘tmp∘a1i
-      end for i in 1:Int(nv/2)
+      end for i in 1:Int(nhv/2)
       )
       setray!(tmp.v, v_matrix=-j_mat(a1))
-      K = StaticArrays.sacollect(SVector{Int(nv/2)}, begin
+      K = StaticArrays.sacollect(SVector{Int(nhv/2)}, begin
         tmp∘B[i]
-      end for i in 1:Int(nv/2)
+      end for i in 1:Int(nhv/2)
       )
-      E = StaticArrays.sacollect(SVector{Int(nv/2)}, begin
+      E = StaticArrays.sacollect(SVector{Int(nhv/2)}, begin
         B[i]∘tmp
-      end for i in 1:Int(nv/2)
+      end for i in 1:Int(nhv/2)
       )
-      H = StaticArrays.sacollect(SVector{Int(nv/2)}, begin
+      H = StaticArrays.sacollect(SVector{Int(nhv/2)}, begin
         setray!(tmp.v, v_matrix=ip_mat(a1, i))
         a1∘tmp∘a1i
-      end for i in 1:Int(nv/2)
+      end for i in 1:Int(nhv/2)
       )
       return (; H=H, B=B, E=E, K=K)
     end
   end
 end
 
-function compute_de_moivre(a1::DAMap, ::Val{linear}=Val{false}()) where {linear}
+function compute_de_moivre(a1::DAMap, ::Val{linear}=Val{true}()) where {linear}
   # jp_mat[i] in FPP is J matrix restricted to i-th plane
   # ip_mat[i] in FPP is identity matrix restricted to i-th plane
   # jt_mat in FPP is symplectic s matrix
@@ -83,8 +83,8 @@ end
 # =================================================================================== #
 # Construct special matrices for lattice functions
 function ip_mat(m::DAMap{V0}, i) where {V0<:StaticArray}
-  nv = isodd(length(V0)) ? length(V0)+1 : length(V0)
-  return StaticArrays.sacollect(SMatrix{nv,nv,real(eltype(V0))}, 
+  nhv = nhvars(m) #isodd(length(V0)) ? length(V0)+1 : length(V0)
+  return StaticArrays.sacollect(SMatrix{nhv,nhv,real(eltype(V0))}, 
   begin
     if col == 2*i-1 && row == 2*i-1
       1
@@ -93,14 +93,11 @@ function ip_mat(m::DAMap{V0}, i) where {V0<:StaticArray}
     else
       0
     end
-  end for col in 1:nv for row in 1:nv)
+  end for col in 1:nhv for row in 1:nhv)
 end
 
 function ip_mat(m::DAMap, i)
-  nv = nvars(m)
-  if isodd(nv) 
-    nv += 1
-  end
+  nhv = nhvars(m)
   ip = zeros(real(eltype(m.v0)), nhv, nhv)
   ip[2*i-1, 2*i-1] = 1
   ip[2*i, 2*i] = 1
@@ -108,35 +105,30 @@ function ip_mat(m::DAMap, i)
 end
 
 function jp_mat(m::DAMap{V0}, i) where {V0<:StaticArray}
-  nv = isodd(length(V0)) ? length(V0)+1 : length(V0)
-  coast = isodd(length(V0))
-  coast_plane = coast && i == Int(nv/2)
-  return StaticArrays.sacollect(SMatrix{nv,nv,real(eltype(V0))}, 
+  nhv = nhvars(m) 
+  return StaticArrays.sacollect(SMatrix{nhv,nhv,real(eltype(V0))}, 
   begin
     if col == 2*i && row == 2*i-1
       1
-    elseif col == 2*i-1 && row == 2*i #&& !coast_plane
+    elseif col == 2*i-1 && row == 2*i
       -1
     else
       0
     end
-  end for col in 1:nv for row in 1:nv)
+  end for col in 1:nhv for row in 1:nhv)
 end
 
 function jp_mat(m::DAMap, i)
-  nv = nvars(m)
-  if isodd(nv) 
-    nv += 1
-  end
-  jp = zeros(real(eltype(m.v0)), nv, nv)
+  nhv = nhvars(m)
+  jp = zeros(real(eltype(m.v0)), nhv, nhv)
   jp[2*i-1, 2*i] = 1
   jp[2*i, 2*i-1] = -1
   return jp
 end
 
 function j_mat(m::DAMap{V0}) where {V0<:StaticArray}
-  nv = isodd(length(V0)) ? length(V0)+1 : length(V0)
-  return StaticArrays.sacollect(SMatrix{nv,nv,real(eltype(V0))}, 
+  nhv = nhvars(m) 
+  return StaticArrays.sacollect(SMatrix{nhv,nhv,real(eltype(V0))}, 
   begin
     if fld1(col,2) != fld1(row,2) 
       0
@@ -149,16 +141,13 @@ function j_mat(m::DAMap{V0}) where {V0<:StaticArray}
         0
       end
     end
-  end for col in 1:nv for row in 1:nv)
+  end for col in 1:nhv for row in 1:nhv)
 end
 
 function j_mat(m::DAMap)
-  nv = nvars(m)
-  if isodd(nv) 
-    nv += 1
-  end
-  j = zeros(real(eltype(m.v0)), nv, nv)
-  for i in 1:Int(nv/2)
+  nhv = nhvars(m)
+  j = zeros(real(eltype(m.v0)), nhv, nhv)
+  for i in 1:Int(nhv/2)
     j[2*i-1, 2*i] = 1
     j[2*i, 2*i-1] = -1
   end
diff --git a/src/map.jl b/src/map.jl
index c5fc13c..b3871b3 100644
--- a/src/map.jl
+++ b/src/map.jl
@@ -265,10 +265,11 @@ end
 
 # zero but new type potentially
 function zero(::Type{$t{V0,V,Q,S}}, m::TaylorMap) where {V0,V,Q,S}
-  return _zero($t{V0,V,Q,S}, getinit(eltype(V)), nvars(m), m)
+  return _zero($t{V0,V,Q,S}, getinit(m), nvars(m), m)
 end
 
-zero(::Type{$t}, m::TaylorMap{V0,V,Q,S}) where {V0,V,Q,S} = zero($t{V0,V,Q,S}, m)
+# Keep same initializer if not specified
+zero(::Type{$t}, m::TaylorMap{V0,V,Q,S}) where {V0,V,Q,S} = _zero($t{V0,V,Q,S}, getinit(m), nvars(m), m)
 
 # Copy ctor including optional TPSA init change
 function $t(m::TaylorMap; init::AbstractTPSAInit=getinit(m))
@@ -347,11 +348,6 @@ function $t(;
     np = ndiffs(init) - nv
   end
 
-
-  #println(nv)
-  #println(np)
-  
-  #println(np)
   nn = nv+np
   # Check if nv+np agrees with ndiffs in init
   nn == ndiffs(init) || error("Number of variables + parameters does not agree with the number of differentials in the TPSA")
diff --git a/src/normal.jl b/src/normal.jl
index db0d102..c2bdd05 100644
--- a/src/normal.jl
+++ b/src/normal.jl
@@ -323,7 +323,12 @@ function factorize(a)
     nt = nv
     ndpt = nv+1
     id = one(a)
-    vf = VectorField(v=view(a0.v, 1:nv), q=a0.q)
+    vf = zero(VectorField, a)
+    setray!(vf.v; v=view(a0.v, 1:nv))
+    if !isnothing(a0.q)
+      setquat!(vf.q; q=a0.q)
+    end
+    #vf = VectorField(v=view(a0.v, 1:nv), q=a0.q)
     TI.clear!(vf.v[nt])
     t1 = zero(vf.v[nt])
     t2 = zero(t1)
@@ -592,82 +597,188 @@ end
 
 # Returns the rotation map to put a in Courant Snyder form 
 # The phase advance can be acquired from this map by atan(r11,r22), etc etc
-function fast_canonize(a::DAMap, damping::Bool=!isnothing(a.s); phase=nothing)
+function canonize(
+  a::DAMap,
+  ::Val{linear}=Val{true}();
+   phase=nothing, 
+  damping::Bool=!isnothing(a.s)
+) where {linear}
   nv = nvars(a)
   nhv = nhvars(a)
   coast = iscoasting(a)
 
-  canonizer = zero(a)
-
-  a_matrix = real.(jacobian(a, VARS))
-  
 
-  #phase = zeros(nvars(a))
+  # tries to be fast for linear calc
+  if linear
+    a_matrix = real.(jacobian(a, VARS))
+    canonizer = zero(a)
+    for i in 1:Int(nhv/2) # for each harmonic oscillator
+      t = sqrt(a_matrix[2*i-1,2*i-1]^2 + a_matrix[2*i-1,2*i]^2)
+      cphi = a_matrix[2*i-1,2*i-1]/t
+      sphi = a_matrix[2*i-1,2*i]/t
+      if sphi*a_matrix[2*i-1,2*i] + cphi*a_matrix[2*i-1,2*i-1] < 0
+        cphi = -cphi
+        sphi = -sphi
+      end
+      TI.seti!(canonizer.v[2*i-1],  cphi, 2*i-1)
+      TI.seti!(canonizer.v[2*i],    cphi, 2*i)
+      TI.seti!(canonizer.v[2*i-1], -sphi, 2*i)
+      TI.seti!(canonizer.v[2*i],    sphi, 2*i-1)
+      
+      if !isnothing(phase)
+        phase[i] += atan(sphi,cphi)/(2*pi)
+      end
+    end
 
-  for i in 1:Int(nhv/2) # for each harmonic oscillator
-    t = sqrt(a_matrix[2*i-1,2*i-1]^2 + a_matrix[2*i-1,2*i]^2)
-    cphi = a_matrix[2*i-1,2*i-1]/t
-    sphi = a_matrix[2*i-1,2*i]/t
-    if sphi*a_matrix[2*i-1,2*i] + cphi*a_matrix[2*i-1,2*i-1] < 0
-      cphi = -cphi
-      sphi = -sphi
+    if coast
+      nt = nv
+      ndpt = nv + 1
+      TI.seti!(canonizer.v[nt], 1, nt)
+      TI.seti!(canonizer.v[ndpt], 1, ndpt)
+      TI.seti!(canonizer.v[nt], -TI.geti(a.v[nt], ndpt), ndpt)
+      if !isnothing(phase)
+        phase[end] += -TI.geti(a.v[nt], ndpt)
+      end
+    end
+    #a_rot = a_matrix*ri
+    #return a_rot
+
+    # Now we have rotated a so that a_12, a_34, a_56, etc are 0 (Courant Snyder)
+    # But if we have damping, we also have
+    # A*S*transpose(A) != S
+    # We can multiply the normalizing map A by some dilation to make it so that, 
+    # even though we don't have exactly A*S*transpose(A) == S, that 
+    # we atleast have (A*S*transpose(A))[1,2] == 1, (A*S*transpose(A))[2,1] == -1, etc 
+
+    # note that with damping we have M as
+    # A*Λ*R*A^-1  where R is the amplitude dependent rotation (diagonal matrix with 
+    # complex values on unit circle) and Λ is a diagonal matrix with real values 
+    # which correspond to the damping (same in each plane, Diagonal(lambda1, lambda1, lambda2, lambda2, etc)
+
+    if damping
+      error("Canonization including damping is not implemented yet")
+      damp = zeros(Int(nvars(a)/2))
+      tmp = zeros(Int(nvars(a)/2), Int(nvars(a)/2))
+      for i in 1:Int(nvars(a)/2)
+        tmp[i,i] = a_rot[2*i-1,2*i-1]*a_rot[2*i,2*i]-a_rot[2*i-1,2*i]*a_rot[2*i,2*i-1]
+        for j in 1:Int(nvars(a)/2)
+          if i != j
+            tmp[i,j] = a_rot[2*i-1,2*j-1]*a_rot[2*i,2*j]-a_rot[2*i-1,2*j]*a_rot[2*i,2*j-1]
+          end
+        end
+      end
+      tmp = inv(tmp)
+      damp .= sqrt.(sum(tmp, dims=2))
+      a_rot = a_rot*Diagonal(repeat(damp,inner=2))
+      damp .= log.(damp)
     end
+    return canonizer
+  else
+    mo = maxord(a)
+    # For nonlinear case we have to exponentiate to ensure symplecticity
+    id = one(a)
+    lin_canonizer = zero(a) # linear part separately to speed up exponent
+    canonizerf = zero(VectorField, a)
+    for i in 1:Int(nhv/2) # for each harmonic oscillator
+      # Need to include parameter dependence if nonlinear, so "var-par"
+      t1 = fast_var_slice(a.v[2*i-1], 2*i-1, nv; par=true)
+      t2 = fast_var_slice(a.v[2*i-1], 2*i,   nv; par=true)
+      t = TI.cutord(sqrt(TI.cutord(t1^2, mo)+TI.cutord(t2^2, mo)), mo)
+      cphi = TI.cutord(t1/t, mo)
+      sphi = TI.cutord(t2/t, mo)
+      if TI.scalar(sphi)*TI.scalar(t2) + TI.scalar(cphi)*TI.scalar(t1) < 0
+        cphi = -cphi
+        sphi = -sphi
+      end
+      phi = TI.cutord(atan(sphi,cphi), mo)
 
-    TI.seti!(canonizer.v[2*i-1],  cphi, 2*i-1)
-    TI.seti!(canonizer.v[2*i],    cphi, 2*i)
-    TI.seti!(canonizer.v[2*i-1], -sphi, 2*i)
-    TI.seti!(canonizer.v[2*i],    sphi, 2*i-1)
+      TI.seti!(lin_canonizer.v[2*i-1], TI.scalar( cphi), 2*i-1)
+      TI.seti!(lin_canonizer.v[2*i],   TI.scalar( cphi), 2*i)
+      TI.seti!(lin_canonizer.v[2*i-1], TI.scalar(-sphi), 2*i)
+      TI.seti!(lin_canonizer.v[2*i],   TI.scalar( sphi), 2*i-1)
 
-    if !isnothing(phase)
-      phase[i] += atan(sphi,cphi)/(2*pi)
+      if !isnothing(phase)
+        TI.add!(phase[i], phase[i], TI.cutord(phi/(2*pi), mo))
+      end
+
+      TI.seti!(phi, 0, 0) # set scalar part to zero bc linear canonization separate
+      factor_in!(canonizerf.v[2*i-1], -phi, 2*i)
+      factor_in!(canonizerf.v[2*i],  phi, 2*i-1)
     end
-  end
 
-  if coast
-    nt = nv
-    ndpt = nv + 1
-    TI.seti!(canonizer.v[nt], 1, nt)
-    TI.seti!(canonizer.v[ndpt], 1, ndpt)
-    TI.seti!(canonizer.v[nt], -TI.geti(a.v[nt], ndpt), ndpt)
-    if !isnothing(phase)
-      phase[end] += -TI.geti(a.v[nt], ndpt)
+    if coast
+      nt = nv
+      ndpt = nv + 1
+      
+      TI.clear!(canonizerf.v[nt])
+      tmp1 = zero(canonizerf.v[nt])
+      tmp2 = zero(canonizerf.v[nt])
+      tmp3 = zero(canonizerf.v[nt])
+      # set the timelike variable so poisson bracket does not change
+      # have :mu_x(delta)(x^2+px^2) + mu_y(delta)(y^2+py^2):
+      # so for time like variable we have 
+      # -\partial_\delta  (mu_x(delta)(x^2+px^2) + mu_y(delta)(y^2+py^2))
+      # = -dmu_x/ddelta*(x^2+px^2) - dmu_y/ddelta(y^2+py^2)
+      for i in 1:Int(nhv/2)
+        #TI.add!(canonizerf.v[nt], canonizerf.v[nt], tm)
+        TI.clear!(tmp3)
+        TI.deriv!(tmp1, canonizerf.v[2*i-1], ndpt)
+        TI.seti!(tmp3, 1, 2*i)
+        TI.mul!(tmp2, tmp1, tmp3)
+        TI.div!(tmp3, tmp2, 2)
+        TI.add!(canonizerf.v[nt], canonizerf.v[nt], tmp3)
+
+        TI.deriv!(tmp1, canonizerf.v[2*i], ndpt)
+        TI.clear!(tmp3)
+        TI.seti!(tmp3, -1, 2*i-1)
+        TI.mul!(tmp2, tmp1, tmp3)
+        TI.div!(tmp3, tmp2, 2)
+        TI.add!(canonizerf.v[nt], canonizerf.v[nt], tmp3)
+      end
+              
+      # get delta dependent part only
+      # this makes sure doesn't blow up, have to remove
+      slip = fast_var_slice(a.v[nt], ndpt, nv; all_ords=true)
+      if !isnothing(phase)
+        TI.sub!(phase[end], phase[end], slip)
+      end
+      TI.seti!(lin_canonizer.v[nt], 1, nt)
+      TI.seti!(lin_canonizer.v[ndpt], 1, ndpt)
+      TI.seti!(lin_canonizer.v[nt], -TI.geti(slip, ndpt), ndpt)
+      
+      # set linear part to zero for nonlinear part
+      TI.seti!(slip, 0, ndpt)
+      TI.add!(canonizerf.v[nt], canonizerf.v[nt], -slip)
+
+      #=
+        TI.seti!(canonizer.v[nt], 1, nt)
+        TI.seti!(canonizer.v[ndpt], 1, ndpt)
+        TI.seti!(canonizer.v[nt], -TI.geti(a.v[nt], ndpt), ndpt)
+        if !isnothing(phase)
+          phase[end] += -TI.geti(a.v[nt], ndpt)
+        end
+      =#
     end
-  end
-  #a_rot = a_matrix*ri
-  #return a_rot
-
-  # Now we have rotated a so that a_12, a_34, a_56, etc are 0 (Courant Snyder)
-  # But if we have damping, we also have
-  # A*S*transpose(A) != S
-  # We can multiply the normalizing map A by some dilation to make it so that, 
-  # even though we don't have exactly A*S*transpose(A) == S, that 
-  # we atleast have (A*S*transpose(A))[1,2] == 1, (A*S*transpose(A))[2,1] == -1, etc 
-
-  # note that with damping we have M as
-  # A*Λ*R*A^-1  where R is the amplitude dependent rotation (diagonal matrix with 
-  # complex values on unit circle) and Λ is a diagonal matrix with real values 
-  # which correspond to the damping (same in each plane, Diagonal(lambda1, lambda1, lambda2, lambda2, etc)
-
-  if damping
-    error("need to finish")
-    damp = zeros(Int(nvars(a)/2))
-    tmp = zeros(Int(nvars(a)/2), Int(nvars(a)/2))
-    for i in 1:Int(nvars(a)/2)
-      tmp[i,i] = a_rot[2*i-1,2*i-1]*a_rot[2*i,2*i]-a_rot[2*i-1,2*i]*a_rot[2*i,2*i-1]
-      for j in 1:Int(nvars(a)/2)
-        if i != j
-          tmp[i,j] = a_rot[2*i-1,2*j-1]*a_rot[2*i,2*j]-a_rot[2*i-1,2*j]*a_rot[2*i,2*j-1]
+
+    if damping
+      error("Canonization including damping is not implemented yet")
+      damp = zeros(Int(nvars(a)/2))
+      tmp = zeros(Int(nvars(a)/2), Int(nvars(a)/2))
+      for i in 1:Int(nvars(a)/2)
+        tmp[i,i] = a_rot[2*i-1,2*i-1]*a_rot[2*i,2*i]-a_rot[2*i-1,2*i]*a_rot[2*i,2*i-1]
+        for j in 1:Int(nvars(a)/2)
+          if i != j
+            tmp[i,j] = a_rot[2*i-1,2*j-1]*a_rot[2*i,2*j]-a_rot[2*i-1,2*j]*a_rot[2*i,2*j-1]
+          end
         end
       end
+      tmp = inv(tmp)
+      damp .= sqrt.(sum(tmp, dims=2))
+      a_rot = a_rot*Diagonal(repeat(damp,inner=2))
+      damp .= log.(damp)
     end
-    tmp = inv(tmp)
-    damp .= sqrt.(sum(tmp, dims=2))
-    a_rot = a_rot*Diagonal(repeat(damp,inner=2))
-    damp .= log.(damp)
+    return lin_canonizer ∘ exp(canonizerf, id)
   end
-
-
-  return canonizer #a_rot
 end
 
 # This can help give you the fixed point
diff --git a/src/sagan_rubin.jl b/src/sagan_rubin.jl
index fc11a86..465a58a 100644
--- a/src/sagan_rubin.jl
+++ b/src/sagan_rubin.jl
@@ -3,26 +3,104 @@ function compute_sagan_rubin(a1::DAMap{V0}, ::Val{linear}=Val{false}()) where {V
   # ip_mat[i] in FPP is identity matrix restricted to i-th plane
   # jt_mat in FPP is symplectic s matrix
 
-  nv = isodd(length(V0)) ? length(V0)+1 : length(V0)
-  nhv = isodd(length(V0)) ? length(V0)-1 : length(V0)
+  # If we are NOT coasting, then we will return the H[3][:,5] and H[3][:,6]
+  # column vectors (eta and zeta in Etienne's blue book). These are all that 
+  # is needed for expressions correct to leading order in the coupling. If 
+  # that is not enough, then the de Moivre lattice functions should be used.
+  # Note that in such an approximation the longitudinal tune will be the 
+  # phase slip
+
+  # If we are coasting, then we will not return any dispersion stuff, since 
+  # that is all well-defined in the parameter-dependent fixed point.
+  nv = nvars(a1)
+  nhv = nhvars(a1)
+  coast = iscoasting(a1)
   if linear
-    let a1_mat = jacobian(a1, VARS_CPARAM), a1i_mat = inv(a1_mat), a1_mat_hv = jacobian(a1, HVARS)
-      H = StaticArrays.sacollect(SVector{Int(nv/2),typeof(a1_mat)}, a1_mat*ip_mat(a1, i)*a1i_mat for i in 1:Int(nv/2))
-      gamma_c = sqrt(H[1][1,1])
-      C = nv > 2 ? StaticArrays.sacollect(SMatrix{2,2,eltype(H[1])}, H[1][i,j] for i in 3:4 for j in 3:4) : 0
-      Ct = nv > 2 ? SA[C[2,2] -C[1,2]; -C[2,1] C[1,1]] : 0
-      eta_full = StaticArrays.sacollect(SVector{nhv,eltype(H[1])}, H[end][i,nv] for i in 1:nhv)
-      eta = StaticArrays.sacollect(SVector{Int(nhv/2),eltype(eta_full)}, eta_full[i] for i in 1:2:nhv) 
-      etap = StaticArrays.sacollect(SVector{Int(nhv/2),eltype(eta_full)}, eta_full[i] for i in 2:2:nhv)
-      Vi = gamma_c*I + vcat(hcat(zero(C), -C), hcat(Ct, zero(Ct))) #transpose(SDiagonal(Ct,-C)) #[gamma_c*I -C; Ct gamma_c*I]
-      N = Vi*a1_mat_hv
-      beta = StaticArrays.sacollect(SVector{Int(nhv/2),eltype(N)}, N[i,i]^2 for i in 1:2:Int(nv/2))
-      alpha = StaticArrays.sacollect(SVector{Int(nhv/2),eltype(N)}, -N[i+1,i]*N[i,i] for i in 1:2:Int(nv/2))
-      return (; beta=beta, alpha=alpha, eta=eta, etap=etap, gamma_c=gamma_c, C=C)
+    if !coast
+      let a1_mat = jacobian(a1, HVARS), a1i_mat = inv(a1_mat), nvm = nv-2
+        H = StaticArrays.sacollect(SVector{Int(nv/2),typeof(a1_mat)}, a1_mat*ip_mat(a1, i)*a1i_mat for i in 1:Int(nv/2))
+        gamma_c = sqrt(H[1][1,1])
+        C = nv > 2 ? StaticArrays.sacollect(SMatrix{2,2,eltype(H[1])}, H[1][i,j] for i in 3:4 for j in 3:4) : 0
+        Ct = nv > 2 ? SA[C[2,2] -C[1,2]; -C[2,1] C[1,1]] : 0
+        Vi = gamma_c*I + vcat(hcat(zero(C), -C), hcat(Ct, zero(Ct)))
+        N = Vi*StaticArrays.sacollect(SMatrix{nvm,nvm,eltype(a1_mat)}, a1_mat[i,j] for i in 1:nvm for j in 1:nvm)
+        beta = StaticArrays.sacollect(SVector{Int(nvm/2),eltype(N)}, N[i,i]^2 for i in 1:2:Int(nvm))
+        alpha = StaticArrays.sacollect(SVector{Int(nvm/2),eltype(N)}, -N[i,i+1]*N[i,i] for i in 1:2:Int(nvm))
+        # Note that eta and zeta here are in x,px,y,py, not in a,pa,b,pb
+        eta = StaticArrays.sacollect(SVector{nvm,eltype(H[1])}, H[end][i,nv] for i in 1:(nvm))
+        zeta = StaticArrays.sacollect(SVector{nvm,eltype(H[1])}, H[end][i,nv-1] for i in 1:(nvm))
+        # we also return B[3][5,6], which can then be multiplied by sin(mu_3) to approximate the
+        # phase slip
+        approx_slip = (a1_mat*jp_mat(a1, Int(nv/2))*a1i_mat)[5,6]
+        # I don't really understand the utility of this, but for consistency with Sagan-Rubin
+        # multiply by Vi to put in "normalized" coordinates
+        return (; beta=beta, alpha=alpha, gamma_c=gamma_c, C=C, eta=Vi*eta, zeta=Vi*zeta, approx_slip=approx_slip)
+      end
+    else
+      let a1_mat = jacobian(a1, HVARS), a1i_mat = inv(a1_mat)
+        H = StaticArrays.sacollect(SVector{Int(nhv/2),typeof(a1_mat)}, a1_mat*ip_mat(a1, i)*a1i_mat for i in 1:Int(nhv/2))
+        gamma_c = sqrt(H[1][1,1])
+        C = nv > 2 ? StaticArrays.sacollect(SMatrix{2,2,eltype(H[1])}, H[1][i,j] for i in 3:4 for j in 3:4) : 0
+        Ct = nv > 2 ? SA[C[2,2] -C[1,2]; -C[2,1] C[1,1]] : 0
+        Vi = gamma_c*I + vcat(hcat(zero(C), -C), hcat(Ct, zero(Ct)))
+        N = Vi*StaticArrays.sacollect(SMatrix{nhv,nhv,eltype(a1_mat)}, a1_mat[i,j] for i in 1:nhv for j in 1:nhv)
+        beta = StaticArrays.sacollect(SVector{Int(nhv/2),eltype(N)}, N[i,i]^2 for i in 1:2:nhv)
+        alpha = StaticArrays.sacollect(SVector{Int(nhv/2),eltype(N)}, -N[i,i+1]*N[i,i] for i in 1:2:nhv)
+        return (; beta=beta, alpha=alpha, gamma_c=gamma_c, C=C)
+      end
     end
   else
-    let
-      error("Not implemented yet")
+    if !coast
+      let nvm = nv-2
+        tmp = zero(a1)
+        a1i = inv(a1)
+        H = StaticArrays.sacollect(SVector{Int(nhv/2)}, begin
+          setray!(tmp.v, v_matrix=ip_mat(a1, i))
+          a1∘tmp∘a1i
+        end for i in 1:Int(nhv/2)
+        )
+        mo = maxord(a1)
+        gamma_c = TI.cutord(sqrt(factor_out(H[1].v[1], 1)), mo)
+        C = nv > 2 ? StaticArrays.sacollect(SMatrix{2,2,eltype(H[1].v)}, factor_out(H[1].v[i], j) for i in 3:4 for j in 3:4) : 0
+        Ct = nv > 2 ? SA[C[2,2] -C[1,2]; -C[2,1] C[1,1]] : 0
+        Vi = gamma_c*I + vcat(hcat(zero(C), -C), hcat(Ct, zero(Ct)))
+        N = TI.cutord.(Vi*StaticArrays.sacollect(SMatrix{nvm,nvm,eltype(a1.v)}, factor_out(a1.v[i], j) for i in 1:nvm for j in 1:nvm), mo)
+        beta = StaticArrays.sacollect(SVector{Int(nvm/2),eltype(N)}, TI.cutord(N[i,i]^2, mo) for i in 1:2:Int(nvm))
+        alpha = StaticArrays.sacollect(SVector{Int(nvm/2),eltype(N)}, TI.cutord(-N[i,i+1]*N[i,i], mo) for i in 1:2:Int(nvm))
+        # Note that eta and zeta here are in x,px,y,py, not in a,pa,b,pb
+        eta = StaticArrays.sacollect(SVector{nvm,eltype(H[1].v)}, factor_out(H[end].v[i], nv) for i in 1:(nvm))
+        zeta = StaticArrays.sacollect(SVector{nvm,eltype(H[1].v)}, factor_out(H[end].v[i], nv-1) for i in 1:(nvm))
+        # I don't really understand the utility of this, but for consistency with Sagan-Rubin
+        # multiply dispersions by Vi to put in "normalized" coordinates
+
+        # we also return B[3][5,6], which can then be multiplied by sin(mu_3) to approximate the
+        # phase slip
+        setray!(tmp.v, v_matrix=jp_mat(a1, 3))
+        B = a1∘tmp∘a1i
+        approx_slip = factor_out(B.v[nv-1], nv)
+        # In the non-coasting case, we return quantities which are coasting-like (now 
+        # contained in a1 and not a0). Of course in coasting case everything is in a0
+        return (; beta=beta, alpha=alpha, gamma_c=gamma_c, C=C,  eta=Vi*eta, zeta=Vi*zeta, approx_slip=approx_slip)
+      end
+    else
+      let
+        tmp = zero(a1)
+        a1i = inv(a1)
+        H = StaticArrays.sacollect(SVector{Int(nhv/2)}, begin
+          setray!(tmp.v, v_matrix=ip_mat(a1, i))
+          a1∘tmp∘a1i
+        end for i in 1:Int(nhv/2)
+        )
+        mo = maxord(a1)
+        gamma_c = TI.cutord(sqrt(factor_out(H[1].v[1], 1)), mo)
+        C = nv > 2 ? StaticArrays.sacollect(SMatrix{2,2,eltype(H[1].v)}, factor_out(H[1].v[i], j) for i in 3:4 for j in 3:4) : 0
+        Ct = nv > 2 ? SA[C[2,2] -C[1,2]; -C[2,1] C[1,1]] : 0
+        Vi = gamma_c*I + vcat(hcat(zero(C), -C), hcat(Ct, zero(Ct)))
+        N = TI.cutord.(Vi*StaticArrays.sacollect(SMatrix{nhv,nhv,eltype(a1.v)}, factor_out(a1.v[i], j) for i in 1:nhv for j in 1:nhv), mo)
+        beta = StaticArrays.sacollect(SVector{Int(nhv/2),eltype(N)}, TI.cutord(N[i,i]^2, mo) for i in 1:2:Int(nhv))
+        alpha = StaticArrays.sacollect(SVector{Int(nhv/2),eltype(N)}, TI.cutord(-N[i,i+1]*N[i,i], mo) for i in 1:2:Int(nhv))
+        return (; beta=beta, alpha=alpha, gamma_c=gamma_c, C=C)
+      end
     end
   end
 end
diff --git a/src/utils.jl b/src/utils.jl
index cc77966..eee6900 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -9,17 +9,96 @@ field array type promotion, checking if the last plane of a map is coasting.
 # Helper functions
 getinit(m::Union{TaylorMap,VectorField}) = TI.getinit(first(m.v))
 ndiffs(m::TaylorMap) = length(m.v)
+ndiffs(m::TaylorMap{<:Any,V}) where {V} = length(V)
 ndiffs(F::VectorField) = TI.ndiffs(first(F.v))
 nmonos(m::Union{TaylorMap,VectorField}) = TI.nmonos(first(m.v))
 maxord(m::Union{TaylorMap,VectorField}) = TI.maxord(first(m.v))
 
 # NNF-specific helpers:
 nvars(m::TaylorMap) = length(m.v0)
+nvars(m::TaylorMap{V0}) where {V0<:StaticArray} = length(V0)
 nvars(F::VectorField) = length(F.v)
+nvars(F::VectorField{V0}) where {V0<:StaticArray} = length(V0)
 nparams(m::Union{TaylorMap,VectorField}) = ndiffs(m) - nvars(m)
 nhvars(m::Union{TaylorMap,VectorField}) = iseven(nvars(m)) ? nvars(m) : nvars(m)-1 # number of "harmonic" variables
 
 iscoasting(m::Union{TaylorMap,VectorField}) = !iseven(nvars(m))
+# =================================================================================== #
+# Tool to "factor" out a first order dependence on a particular variable
+# e.g. for a TPSA with monomial [1,2,3,4], for 1 will return same TPSA but 
+# that monomial is now [0,2,3,4]
+# This is useful for lattice functions (first order) with nonlinear parameter dependence
+
+factor_out(t, var::Int) = (out = zero(t); factor_out!(out, t, var))
+
+function factor_out!(out, t, var::Int)
+  TI.is_tps_type(typeof(t)) isa TI.IsTPSType || error("Function only accepts TPS types")
+  TI.is_tps_type(typeof(out)) isa TI.IsTPSType || error("Function only accepts TPS types")
+  nn = ndiffs(t)
+  v = Ref{TI.numtype(t)}()
+  tmpmono = zeros(UInt8, nn) 
+
+  idx = TI.cycle!(t, 0, mono=tmpmono, val=v)
+  while idx > 0
+    if tmpmono[var] != 0
+      tmpmono[var] -= 1
+      TI.setm!(out, v[], tmpmono)
+    end
+    idx = TI.cycle!(t, idx, mono=tmpmono, val=v)
+  end
+  return out
+end
+
+factor_in(t, var::Int, n::Int=1) = (out = zero(t); factor_in!(out, t, var, n))
+
+function factor_in!(out, t, var::Int, n::Int=1)
+  TI.is_tps_type(typeof(t)) isa TI.IsTPSType || error("Function only accepts TPS types")
+  TI.is_tps_type(typeof(out)) isa TI.IsTPSType || error("Function only accepts TPS types")
+  nn = ndiffs(t)
+  v = Ref{TI.numtype(t)}()
+  tmpmono = zeros(UInt8, nn) 
+
+  # First do scalar part
+  tmpmono[var] += n
+  TI.setm!(out, TI.geti(t, 0), tmpmono)
+
+  idx = TI.cycle!(t, 0, mono=tmpmono, val=v)
+  while idx > 0
+    tmpmono[var] += n
+    TI.setm!(out, v[], tmpmono)
+    idx = TI.cycle!(t, idx, mono=tmpmono, val=v)
+  end
+  return out
+end
+
+fast_var_slice(t, var::Int, nv::Int; par::Bool=false, all_ords::Bool=false) = (out = zero(t); fast_var_slice!(out, t, var, nv; par=par, all_ords=all_ords); return out)
+
+# This will par out the polynomial with only 
+# 1st order dependence on variable, nonlinear 
+# parameter dependence
+function fast_var_slice!(out, t, var::Int, nv::Int; par::Bool=false, all_ords::Bool=false)
+  TI.is_tps_type(typeof(t)) isa TI.IsTPSType || error("Function only accepts TPS types")
+  TI.is_tps_type(typeof(out)) isa TI.IsTPSType || error("Function only accepts TPS types")
+  nn = ndiffs(t)
+  np = nn-nv
+  v = Ref{TI.numtype(t)}()
+  tmpmono = zeros(UInt8, nn) 
+
+  idx = TI.cycle!(t, 0, mono=tmpmono, val=v)
+  while idx > 0
+    if all_ords || tmpmono[var] == 1
+      if all(t->t==0, view(tmpmono, 1:(var-1))) && all(t->t==0, view(tmpmono, (var+1):nv))
+        if par
+          tmpmono[var] -= 1
+        end
+        TI.setm!(out, v[], tmpmono)
+      end
+    end
+    idx = TI.cycle!(t, idx, mono=tmpmono, val=v)
+  end
+  return out
+end
+
 # =================================================================================== #
 # Jacobian/jacobiant
 # Useful options should be 1) harmonic variables, 2) variables, and 3) variables + parameters
@@ -411,4 +490,24 @@ function checksymp(M)
   res = transpose(M)*S*M-S
   return res
 end
+
+function checksymp(m::TaylorMap)
+  nv = nvars(m)
+  if iscoasting(m)
+    nv += 1
+  end
+  Sij = S(nv)
+  mo = maxord(m)
+  out = zeros(eltype(m.v), nv, nv)
+  for i in 1:nv
+    for j in 1:nv
+      grad1 = [TI.deriv(m.v[j], ell) for ell in 1:nv]
+      grad2 = [TI.deriv(m.v[i], k) for k in 1:nv]
+      # per thesis, for given k (index) we sum
+      s1 = [sum([TI.cutord(Sij[k,ell]*grad1[ell],mo) for ell in 1:nv]) for k in 1:nv]
+      out[i,j] = sum([TI.cutord(grad2[k]*s1[k],mo) for k in 1:nv])
+    end
+  end
+  return out-S
+end
 # =================================================================================== #
\ No newline at end of file
diff --git a/src/vectorfield.jl b/src/vectorfield.jl
index 4968d0f..c98cdf2 100644
--- a/src/vectorfield.jl
+++ b/src/vectorfield.jl
@@ -118,12 +118,12 @@ end
 
 # zero but new type potentially
 function zero(::Type{VectorField{V,Q}}, F::Union{VectorField,TaylorMap}) where {V,Q}
-  return _zero(VectorField{V,Q}, getinit(eltype(V)), nvars(F))
+  return _zero(VectorField{V,Q}, getinit(F), nvars(F))
 end
 
-zero(::Type{VectorField}, F::VectorField{V,Q}) where {V,Q} = zero(VectorField{V,Q}, F)
-
-zero(::Type{VectorField}, m::TaylorMap{V0,V,Q,S}) where {V0,V,Q,S} = zero(VectorField{similar_eltype(V0,eltype(V)),Q}, m)
+# Keep same initializer if not specified
+zero(::Type{VectorField}, F::VectorField{V,Q}) where {V,Q} = _zero(VectorField{V,Q}, getinit(F), nvars(F))
+zero(::Type{VectorField}, m::TaylorMap{V0,V,Q,S}) where {V0,V,Q,S} = _zero(VectorField{similar_eltype(V0,eltype(V)),Q}, getinit(m), nvars(m))
 
 # Copy ctor including optional TPSA init change
 function VectorField(F::VectorField; init::AbstractTPSAInit=getinit(F))
@@ -483,27 +483,19 @@ function *(F::VectorField, m1::Union{DAMap,UniformScaling})
   return m
 end
 
-function exp(F::VectorField, m1::Union{UniformScaling,DAMap}=I)
+function exp(F::VectorField, m1::DAMap)
   Fprom = F
-  if m1 isa UniformScaling
-    error("Not currently supported please provide a map")
-    # TO-DO: How/should I store nparams in VectorField?
-    # It is just equal to nvars - ndiffs, though this will be slow...
-    m1prom = one(DAMap{typeof(F.v),typeof(F.q)})
-    Fprom = F
-  else
-    if TI.numtype(eltype(m1.v)) != TI.numtype(eltype(F.v))
-      if promote_type(typeof(F), TI.numtype(eltype(m1.v))) != typeof(F)
-        Fprom = zero(promote_type(typeof(F), TI.numtype(eltype(m1.v))), F)
-        copy!(Fprom, F)
-      else
-        m1prom = zero(promote_type(typeof(m1), TI.numtype(eltype(F.v))), m1)
-        copy!(m1prom, m1)
-      end
+  if TI.numtype(eltype(m1.v)) != TI.numtype(eltype(F.v))
+    if promote_type(typeof(F), TI.numtype(eltype(m1.v))) != typeof(F)
+      Fprom = zero(promote_type(typeof(F), TI.numtype(eltype(m1.v))), F)
+      copy!(Fprom, F)
     else
-      Fprom = F
-      m1prom = m1
+      m1prom = zero(promote_type(typeof(m1), TI.numtype(eltype(F.v))), m1)
+      copy!(m1prom, m1)
     end
+  else
+    Fprom = F
+    m1prom = m1
   end
   m = zero(m1prom)
   exp!(m, Fprom, m1prom)