#!/usr/bin/env python
# coding: utf-8

# ## Deformations
# 
# In PyCA, we have functionality for v-field (vector field) and h-field (deformation field) transformations. We set up a few dummy Image3Ds and Field3Ds to show the syntax.

# In[1]:


import PyCA.Core as ca
grid = ca.GridInfo(ca.Vec3Di(10, 10, 10), ca.Vec3Df(1.0, 1.0, 1.0), ca.Vec3Df(1.0, 1.0, 1.0))
mType = ca.MEM_HOST

# image/field to be deformed
im = ca.Image3D(grid, mType)
ca.SetMem(im, 1.0)
f = ca.Field3D(grid, mType)
ca.SetMem(f, ca.Vec3Df(1.0, 2.0, 3.0))

# Deformation field (v-field)
v = ca.Field3D(grid, mType)
ca.SetToZero(v)
# Deformation field (h-field)
h = ca.Field3D(grid, mType)
ca.SetToIdentity(h)


# Here, we just defined an image `im` and a field `f` that we will be deformed.
# The v-field `v` is set to 0, and the h-field `h` is set to x (identity).
# Each h-field has a corresponding v-field, and they are related by $h(x) = x + v(x)$.
# We can convert v-fields to h-fields and vice versa:

# In[2]:


ca.VtoH_I(v) # v(x) = v(x) + x
ca.HtoV_I(v) # v(x) = v(x) - x


# Applying a deformation to an image/field is done via a trilinear interpolation (or bilinear, for 2D images/fields).  Note that as in typical image registration, the deformation being applied is the inverse deformation, i.e. function composition: $I_{def}(x) = I(h(x)) = I(x+v(x))$.
# The PyCA functions are straightforward:

# In[3]:


alpha = 1.0 #scalar multiple for applying a v-field
im_def = ca.Image3D(grid, mType)
f_def = ca.Field3D(grid, mType)
bg = ca.BACKGROUND_STRATEGY_PARTIAL_ZERO

ca.ApplyV(im_def, im, v, alpha, bg) # im_def = im(x+alpha*v(x))
ca.ApplyVInv(im_def, im, v, alpha, bg) # im_def = im(x-alpha*v(x))
ca.ApplyH(im_def, im, h, bg) #im_def = im(h(x))
ca.ApplyV(f_def, f, v, bg) #f_def = f(x+v(x))
ca.ApplyVInv(f_def, f, v, bg) #f_def = f(x-v(x))
ca.ApplyH(f_def, f, h, bg) #f_def = f(h(x))


# Of course, these deformations don't actually do anything (since we have `v` and `h` set to zero and identity); this is just to demonstrate the syntax.
# 
# You will notice a few things here. First, ApplyV has a sometimes optional `alpha` as a constant multiple for the vector field. Second, we can choose a background strategy (what value to use for interpolations that lie outside the image boundary).  The options are:

# In[4]:


bg = ca.BACKGROUND_STRATEGY_CLAMP # set background to value at the image boundary
bg = ca.BACKGROUND_STRATEGY_PARTIAL_ZERO # background = 0
bg = ca.BACKGROUND_STRATEGY_PARTIAL_ID # background is identity (for fields only)
bg = ca.BACKGROUND_STRATEGY_WRAP # wrap image (interpolation on the torus)


# ### Deformation Composition
# You could do a composition of deformations using the previous functions, but since there is an intrinsic meaning to a v-field and h-field, we have composition functions that work specifically on these objects. This way, you don't need to set the appropriate background strategy.

# In[5]:


#initialize some variables
h_out = ca.Field3D(grid, mType)
h = ca.Field3D(grid, mType)
h0 = ca.Field3D(grid, mType)
h1 = ca.Field3D(grid, mType)
v = ca.Field3D(grid, mType)
ca.SetToIdentity(h)
ca.SetToIdentity(h0)
ca.SetToIdentity(h1)
ca.SetToZero(v)

#Apply Compositions
ca.ComposeHH(h_out, h0, h1) # h_out = h0(h1(x))
ca.ComposeHV(h_out, h, v, 1.0) # h_out = h(x+v(x))
ca.ComposeHVInv(h_out, h, v, 1.0) # h_out = h(x-v(x))
ca.ComposeVH(h_out, v, h, 1.0) # h_out = h(x)+v(h(x)) = (x+v(x))(h(x))
ca.ComposeVInvH(h_out, v, h, 1.0) # h_out = h(x)-v(h(x)) = (x-v(x))(h(x))


# If we have a list of deformations that we need composed, we can do this by going forward through the list or backward through the list; both methods are essentially equivalent. I will show this using `ComposeHH`, but doing it with a list of vector fields (or negative vector fields) is similar. 
# 
# For a list of 10 fields, our goal is to solve for $h(x) = h_0 \circ h_1 ... \circ h_9(x)$.
# 
# Going forward though the fields, we have

# In[6]:


# set up list of h fields
hlist = [ca.Field3D(grid, mType) for _ in xrange(10)]
scratchF = ca.Field3D(grid, mType)
for ht in hlist:
    ca.SetToIdentity(ht)
    
# Compose list        
ca.SetToIdentity(h)
for ht in hlist:
    ca.ComposeHH(scratchF, h, ht)   
    h.swap(scratchF)           # h = scratchF, scratchF = h  
                                                             
# h = x                         : initializiation        
# h = h(h0(x)) = h0(x)          : first iteration        
# h = h(h1(x)) = h0(h1(x))                               
# h = h(h2(x)) = h0(h1(h2(x)))                           
# ...                                                    
# h = h(h9(x)) = h0(h1(...(h8(h9))))


# Notice here that we need a scratch variable `scratchF` because `ComposeHH` must take three distinct Field3Ds. We then just call the swap method.
# 
# To go backward through the list, we have

# In[7]:


ca.SetToIdentity(h)
for ht in reversed(hlist):
    ca.ComposeHH(scratchF, ht, h)
    h.swap(scratchF)

# h = x                    : initialization
# h = h9(h(x)) = h9(x)     : first iteration
# h = h8(h(x)) = h8(h9(x)
# h = h7(h(x)) = h7(h8(h9(x)))
# ...
# h = h0(h(x)) = h0(h1(...(h8(h9))))