From 59e1a88e9727fe6839fdf66f43e9f8d2676bc798 Mon Sep 17 00:00:00 2001
From: Abe Levitan <alevitan@mit.edu>
Date: Mon, 1 Apr 2019 17:16:15 -0400
Subject: [PATCH] Add initialization for object size

---
 CDTools/tools/initializers.py    | 36 ++++++++++++++++++++++++++++++--
 tests/tools/test_initializers.py | 21 +++++++++++++++++++
 2 files changed, 55 insertions(+), 2 deletions(-)

diff --git a/CDTools/tools/initializers.py b/CDTools/tools/initializers.py
index eb4dfec..3479318 100644
--- a/CDTools/tools/initializers.py
+++ b/CDTools/tools/initializers.py
@@ -2,7 +2,7 @@ from __future__ import division, print_function, absolute_import
 import numpy as np
 import torch as t
 
-__all__ = ['exit_wave_geometry', 'gaussian']
+__all__ = ['exit_wave_geometry', 'calc_object_setup', 'gaussian']
 
 from CDTools.tools import cmath
 from scipy.fftpack import next_fast_len
@@ -68,7 +68,39 @@ def exit_wave_geometry(det_basis, det_shape, wavelength, distance, center=None,
     
     return real_space_basis, full_shape, det_slice
                       
-        
+
+def calc_object_setup(probe_shape, translations, padding=0):
+    """Returns an object shape and minimum pixel translation
+
+    Based on the given pixel-space translations, it will calculate the
+    required size for an object array and calculate the pixel translation
+    that corresponds to a shift by (0,0) of the probe. 
+    
+    Optionally a small extra border can be defined via the padding
+    attribute. If this is done, the calculated pixel translation will
+    correspond to (padding,padding)
+    
+    Args:
+        probe_shape (t.Size) : The size of the probe array
+        translations (t.Tensor) : Jx2 stack of pixel-valued (i,j) translations
+        padding (int) : Optional, the size of an extra border to include
+    """
+    # First we look at the translations to find the minimum translation
+    # and the range of translations
+    min_translation = t.min(translations, dim=0)[0]
+    translation_range = t.max(translations, dim=0)[0] - min_translation
+
+    # Calculate the required shape
+    translation_range = t.ceil(translation_range).numpy().astype(np.int32)
+    shape = translation_range + np.array(probe_shape) + 2 * padding
+    shape = t.Size(shape)
+
+    # And the minimum translation
+    min_translation = min_translation - padding
+
+    return shape, min_translation
+    
+    
 
 def gaussian(shape, amplitude, sigma, center = None):
     """Returns an array with a centered gaussian
diff --git a/tests/tools/test_initializers.py b/tests/tools/test_initializers.py
index 5cb6ba3..4107898 100644
--- a/tests/tools/test_initializers.py
+++ b/tests/tools/test_initializers.py
@@ -52,6 +52,27 @@ def test_exit_wave_geometry():
     assert t.ones(full_shape)[det_slice].shape ==  shape
 
 
+
+def test_calc_object_setup():
+    # First just try a simple case
+    probe_shape = t.Size([120,57])
+    translations = t.rand((30,2)) * 300
+    t_max = t.max(translations, dim=0)[0]
+    t_min = t.min(translations, dim=0)[0]
+    obj_shape, min_translation = initializers.calc_object_setup(probe_shape, translations)
+    exp_shape = t.ceil(t_max - t_min).to(t.int32) + t.Tensor(list(probe_shape)).to(t.int32)
+    assert t.allclose(min_translation, t_min)
+
+    assert obj_shape == t.Size(exp_shape)
+    
+    # Then add some padding
+    padding = 5
+    obj_shape, min_translation = initializers.calc_object_setup(probe_shape, translations, padding=padding)
+    assert t.allclose(min_translation, t_min - padding)
+    assert obj_shape == t.Size(exp_shape + 2 * padding)
+        
+
+    
 def test_gaussian():
     # Generate gaussian as a numpy array (square array)
     shape = [10, 10]