From 3248a3b46e766a9dbc476d826135a372d3d6abfa Mon Sep 17 00:00:00 2001
From: Abe Levitan <alevitan@mit.edu>
Date: Mon, 1 Apr 2019 16:51:36 -0400
Subject: [PATCH] Add tool to convert from real space to pixel space
 translations

---
 CDTools/tools/interactions.py    | 44 ++++++++++++++++++++++++++++++++
 tests/tools/test_interactions.py | 30 ++++++++++++++++++++++
 2 files changed, 74 insertions(+)

diff --git a/CDTools/tools/interactions.py b/CDTools/tools/interactions.py
index 8593d6e..6a07974 100644
--- a/CDTools/tools/interactions.py
+++ b/CDTools/tools/interactions.py
@@ -9,6 +9,50 @@ import torch as t
 # area.
 #
 
+def translations_to_pixel(basis, translations, surface_normal=t.Tensor([0,0,1])):
+    """Takes real space translations and outputs them in pixel space
+    
+    This works for any 2D ptychography geometry. It takes in
+    A set of translations in (x,y) space and outputs the same translations
+    in internal pixel units perpendicular to the detector. 
+    
+    It uses information on the wavefield basis and, if defined, the
+    sample normal, to perform the conversion.
+    
+    The assumed geometry is incoming radiation with a wavevector parallel
+    to the +z axis, [0,0,1]. The default sample orientation has a surface
+    normal parallel to this direction
+
+    Args:
+        basis (torch.Tensor) : The real space basis the wavefields are defined in
+        translations (torch.Tensor) : A Jx3 stack of real-space translations
+        surface_normal (torch.Tensor) : Optional, the sample's surface normal
+    """
+    
+    projection_1 = t.Tensor([[1,0,0],
+                             [0,1,0],
+                             [0,0,0]])
+    projection_2 = t.inverse(t.Tensor([[1,0,0],
+                                       [0,1,0],
+                                       -surface_normal/
+                                       surface_normal[2]]))
+    basis_vectors_inv = t.pinverse(basis)
+    projection = t.mm(basis_vectors_inv,
+                      t.mm(projection_2,projection_1))
+    projection = projection.t()
+
+    single_translation = False
+    if len(translations.shape) == 1:
+        translations = translations[None,:]
+        single_translation = True
+
+    pixel_translations = t.mm(translations, projection)
+
+    if single_translation:
+        return pixel_translations[0]
+    else:
+        return pixel_translations
+    
 
 def ptycho_2D_round(probe, obj, translations):
     """Returns a stack of exit waves without accounting for subpixel shifts
diff --git a/tests/tools/test_interactions.py b/tests/tools/test_interactions.py
index 810d532..d787aa4 100644
--- a/tests/tools/test_interactions.py
+++ b/tests/tools/test_interactions.py
@@ -28,6 +28,36 @@ def single_pixel_probe(scope='module'):
     return probe
 
 
+def test_translations_to_pixel():
+    # First, try the case where everything is ones and simple
+    basis = t.Tensor([[0,-1,0],[-1,0,0]]).t()
+    translations = t.rand((10,3))
+    output = interactions.translations_to_pixel(basis, translations)
+    assert t.allclose(output, -translations[:,:2].flip(1))
+    
+    # Next, try a case with a single translation
+    translation = t.rand((3))
+    output = interactions.translations_to_pixel(basis, translation)
+    assert t.allclose(output, -translation[:2].flip(0))
+    
+    # Then, try a case with no surface normal but with a real conversion
+    basis = t.Tensor([[0,-2,0],[-1,0,0.1]]).t()
+    translations = t.rand((10,3))
+    output = interactions.translations_to_pixel(basis, translations)
+    basis_vectors_inv = t.pinverse(basis)
+    translations[:,2] = 0 # manually project off z component
+    assert t.allclose(output, t.mm(translations,basis_vectors_inv.t()))
+    
+    # Finally, try a case with a known surface normal (reflection)
+    basis = t.Tensor([[0,-1,0],[0,0,1]]).t()
+    surface_normal = t.Tensor([np.sqrt(2),0,-np.sqrt(2)])
+    translations = t.rand((10,3))
+    output = interactions.translations_to_pixel(basis, translations,
+                                                surface_normal=surface_normal)
+    exp_translations = t.stack((-translations[:,1],translations[:,0]),dim=1)
+    assert t.allclose(output, exp_translations)
+    
+
 def test_ptycho_2D_round(random_probe, random_obj):
     # Test a stack of images
     translations = np.random.rand(10,2) * 500