Merge pull request #33 from bmad-sim/strict1

Strict1

fpp-ptc-sandbox/code/z_radiation_matt_fake_maps.f90

@@ -135,13 +135,13 @@ program example
 
   do i=1,c_%nd2
   m%v(i)=m%v(i)*decrement(i)
-  enddo
- m1=m + m
+ enddo
+ !m1=3.d0 + m
 
-   do i=1,6
-    write(6,format6) real(m1%e_ij(i,1:6) )
-  enddo
-  stop
+  ! do i=1,6
+  !  write(6,format6) real(m1%e_ij(i,1:6) )
+  !enddo
+  !stop
 
 ! do i=1,1000
 !  m1  = m1*m1

src/NonlinearNormalForm.jl

@@ -19,7 +19,9 @@ import Base: ∘,
              copy!,
              convert,
              show,
-             rand
+             rand,
+             promote_rule,
+             eltype
 
 import LinearAlgebra: norm,
                       dot,

src/map/compose.jl

@@ -28,8 +28,7 @@ See the documentation for `compose_it!` for information on `work_low` and `work_
 - `work_prom` -- Temporary vector of allocated `ComplexTPS`s when there is implicit promotion. See the `compose_it!` documentation for more details. Default is output from `prep_comp_work_prom(m, m2, m1)`
 """
 function compose!(m::DAMap, m2::DAMap, m1::DAMap; keep_scalar::Bool=true, work_ref::Union{Nothing,Vector{<:Union{Float64,ComplexF64}}}=nothing, dospin::Bool=true, work_low::Tuple{Vararg{Vector{<:Union{Ptr{RTPSA},Ptr{CTPSA}}}}}=prep_comp_work_low(m), work_prom::Union{Nothing,Tuple{Vararg{Vector{<:ComplexTPS}}}}=prep_comp_work_prom(m,m2,m1))
-  checkop(m, m2, m1)
-  checkpromotion(m, m2, m1)
+  checkinplace(m, m2, m1)
 
   # DAMap setup:
   desc = getdesc(m1)
@@ -92,8 +91,7 @@ See the documentation for `compose_it!` for information on `work_low` and `work_
 - `work_prom` -- Temporary vector of allocated `ComplexTPS`s when there is implicit promotion. See the `compose_it!` documentation for more details. Default is output from `prep_comp_work_prom(m, m2, m1)`
 """
 function compose!(m::TPSAMap, m2::TPSAMap, m1::TPSAMap; dospin::Bool=true, work_low::Tuple{Vararg{Vector{<:Union{Ptr{RTPSA},Ptr{CTPSA}}}}}=prep_comp_work_low(m), work_prom::Union{Nothing,Tuple{Vararg{Vector{<:ComplexTPS}}}}=prep_comp_work_prom(m,m2,m1))
-  checkop(m, m2, m1)
-  checkpromotion(m, m2, m1)
+  checkinplace(m, m2, m1)
 
   # TPSAMap setup:
   # For TPSA Map concatenation, we need to subtract w_0 (m2 x0) (Eq. 33)
@@ -134,41 +132,7 @@ $($t) composition, which calculates `m2 ∘ m1` $( $t == DAMap ? "ignoring the s
 """
 function compose(m2::$t,m1::$t)
   checkop(m2, m1)
-
-  desc = getdesc(m1)
-  nn = numnn(desc)
-  nv = numvars(desc)
-
-  outT = promote_type(eltype(m2.x),eltype(m1.x))
-
-  # set up outx0
-  outx0 = Vector{numtype(outT)}(undef, nv)
-
-  # Set up outx:
-  outx = Vector{outT}(undef, nn)
-  for i=1:nv  # no need to allocate immutable parameters taken care of inside compose_it!
-      @inbounds outx[i] = outT(use=desc)
-  end
-
-  # set up quaternion out:
-  if !isnothing(m1.Q)
-    outq = Vector{outT}(undef, 4)
-    for i=1:4
-      @inbounds outq[i] = outT(use=desc)
-    end
-    outQ = Quaternion(outq)
-  else
-    outQ = nothing
-  end
-
-  # set up stochastic out
-  if isnothing(m1.E) && isnothing(m2.E)
-    outE = nothing
-  else
-    outE = Matrix{numtype(outT)}(undef, nv, nv)
-  end
-
-  m = $t(outx0, outx, outQ, outE, m1.idpt)
+  m = zero_op(m1,m2)
   compose!(m, m2, m1)
 
   return m

src/map/compose_it.jl

@@ -42,9 +42,7 @@ Note that the `ComplexTPS`s in the vector(s) must be allocated and have the same
 If spin is included, not that the final quaternion concatenation step mul! will creat allocations
 """ 
 function compose_it!(m::$t, m2::$t, m1::$t; dospin::Bool=true, dostochastic::Bool=true, work_low::Tuple{Vararg{Vector{<:Union{Ptr{RTPSA},Ptr{CTPSA}}}}}=prep_comp_work_low(m), work_prom::Union{Nothing,Tuple{Vararg{Vector{<:ComplexTPS}}}}=prep_comp_work_prom(m,m2,m1))
-  checkop(m, m2, m1)
-  checkpromotion(m, m2, m1)
-
+  checkinplace(m, m2, m1)
   @assert !(m === m1) "Cannot compose_it!(m, m2, m1) with m === m1"
 
   desc = getdesc(m1)
@@ -153,16 +151,8 @@ function compose_it!(m::$t, m2::$t, m1::$t; dospin::Bool=true, dostochastic::Boo
   # Stochastic
   # MAKE THIS FASTER!
   if !isnothing(m.E) && dostochastic
-    if !isnothing(m1.E)
-      M2 = jacobian(m2)   
-      m.E .= M2*m1.E*transpose(M2) 
-    else 
-      m.E .= 0
-    end
-
-    if !isnothing(m2.E)
-      m.E .+= m2.E
-    end
+    M2 = jacobian(m2)   
+    m.E .= M2*m1.E*transpose(M2) + m2.E
   end
 
   return