diff --git a/CDTools/models/__init__.py b/CDTools/models/__init__.py
index f3b4823..5691059 100644
--- a/CDTools/models/__init__.py
+++ b/CDTools/models/__init__.py
@@ -29,9 +29,9 @@ __all__ = ['CDIModel', 'SimplePtycho', 'FancyPtycho', 'Bragg2DPtycho', 'SMatrixP
 from CDTools.models.base import CDIModel
 from CDTools.models.simple_ptycho import SimplePtycho
 from CDTools.models.fancy_ptycho import FancyPtycho
-from CDTools.models.pinhole_plane_ptycho import PinholePlanePtycho
-from CDTools.models.bragg_2d_ptycho import Bragg2DPtycho
-from CDTools.models.s_matrix_ptycho import SMatrixPtycho
-from CDTools.models.multislice_2d_ptycho import Multislice2DPtycho
-from CDTools.models.rpi import RPI
-from CDTools.models.unified_mode_ptycho import UnifiedModePtycho
+#from CDTools.models.pinhole_plane_ptycho import PinholePlanePtycho
+#from CDTools.models.bragg_2d_ptycho import Bragg2DPtycho
+#from CDTools.models.s_matrix_ptycho import SMatrixPtycho
+#from CDTools.models.multislice_2d_ptycho import Multislice2DPtycho
+#from CDTools.models.rpi import RPI
+#from CDTools.models.unified_mode_ptycho import UnifiedModePtycho
diff --git a/CDTools/models/base.py b/CDTools/models/base.py
index 9e28501..63d3aa7 100644
--- a/CDTools/models/base.py
+++ b/CDTools/models/base.py
@@ -40,6 +40,8 @@ import threading
 import queue
 import time
 #import pytorch_warmup
+from .complex_adam import MyAdam
+from .complex_lbfgs import MyLBFGS
 
 __all__ = ['CDIModel']
 
@@ -137,7 +139,6 @@ class CDIModel(t.nn.Module):
         for inputs, patterns in data_loader:
             normalization += t.sum(patterns).cpu().numpy()
             
-            
         def run_iteration(stop_event=None):
             loss = 0
             N = 0
@@ -169,13 +170,18 @@ class CDIModel(t.nn.Module):
 
                         loss.backward()
                         total_loss += loss.detach()    
-                            
+
+                    #print('probe grad')
+                    #print(t.mean(self.obj.grad))
+
                     if regularization_factor is not None \
                        and hasattr(self, 'regularizer'):
                         loss = self.regularizer(regularization_factor)
                         loss.backward()
                     return total_loss
 
+                                    
+
                 if warmup_scheduler is not None:
                     old_lrs = [group['lr'] for group in
                                optimizer.param_groups]
@@ -272,7 +278,8 @@ class CDIModel(t.nn.Module):
                                            shuffle=True)
 
         # Define the optimizer
-        optimizer = t.optim.Adam(self.parameters(), lr = lr, amsgrad=amsgrad)
+        #optimizer = t.optim.Adam(self.parameters(), lr = lr, amsgrad=amsgrad)
+        optimizer = MyAdam(self.parameters(), lr = lr, amsgrad=amsgrad)
 
 
         # Define the scheduler
@@ -339,8 +346,10 @@ class CDIModel(t.nn.Module):
 
 
         # Define the optimizer
-        optimizer = t.optim.LBFGS(self.parameters(),
-                                  lr = lr, history_size=history_size)
+        #optimizer = t.optim.LBFGS(self.parameters(),
+        #                          lr = lr, history_size=history_size)
+        optimizer = MyLBFGS(self.parameters(),
+                            lr = lr, history_size=history_size)
         
         return self.AD_optimize(iterations, data_loader, optimizer,
                                 regularization_factor=regularization_factor,
@@ -515,7 +524,7 @@ class CDIModel(t.nn.Module):
         fig, axes = plt.subplots(1,3,figsize=(12,5.3))
         fig.tight_layout(rect=[0.02, 0.09, 0.98, 0.96])
         axslider = plt.axes([0.15,0.06,0.75,0.03])
-
+        
         
         def update_colorbar(im):
             # If the update brought the colorbar out of whack
@@ -532,7 +541,6 @@ class CDIModel(t.nn.Module):
             im.norecurse=True
             im.set_clim(vmin=np.min(im.get_array()),vmax=np.max(im.get_array()))
 
-        
         def update(idx):
             idx = int(idx) % len(dataset)
             fig.pattern_idx = idx
diff --git a/CDTools/models/complex_adam.py b/CDTools/models/complex_adam.py
new file mode 100644
index 0000000..1ded7d9
--- /dev/null
+++ b/CDTools/models/complex_adam.py
@@ -0,0 +1,163 @@
+import math
+import torch
+from torch import Tensor
+from torch.optim.optimizer import Optimizer
+from typing import List, Optional
+
+class MyAdam(Optimizer):
+    r"""Implements Adam algorithm.
+    It has been proposed in `Adam: A Method for Stochastic Optimization`_.
+    The implementation of the L2 penalty follows changes proposed in
+    `Decoupled Weight Decay Regularization`_.
+    Args:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups
+        lr (float, optional): learning rate (default: 1e-3)
+        betas (Tuple[float, float], optional): coefficients used for computing
+            running averages of gradient and its square (default: (0.9, 0.999))
+        eps (float, optional): term added to the denominator to improve
+            numerical stability (default: 1e-8)
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
+        amsgrad (boolean, optional): whether to use the AMSGrad variant of this
+            algorithm from the paper `On the Convergence of Adam and Beyond`_
+            (default: False)
+    .. _Adam\: A Method for Stochastic Optimization:
+        https://arxiv.org/abs/1412.6980
+    .. _Decoupled Weight Decay Regularization:
+        https://arxiv.org/abs/1711.05101
+    .. _On the Convergence of Adam and Beyond:
+        https://openreview.net/forum?id=ryQu7f-RZ
+    """
+
+    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,
+                 weight_decay=0, amsgrad=False):
+        if not 0.0 <= lr:
+            raise ValueError("Invalid learning rate: {}".format(lr))
+        if not 0.0 <= eps:
+            raise ValueError("Invalid epsilon value: {}".format(eps))
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
+        if not 0.0 <= weight_decay:
+            raise ValueError("Invalid weight_decay value: {}".format(weight_decay))
+        defaults = dict(lr=lr, betas=betas, eps=eps,
+                        weight_decay=weight_decay, amsgrad=amsgrad)
+        super(MyAdam, self).__init__(params, defaults)
+
+    def __setstate__(self, state):
+        super(MyAdam, self).__setstate__(state)
+        for group in self.param_groups:
+            group.setdefault('amsgrad', False)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        """Performs a single optimization step.
+        Args:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            params_with_grad = []
+            grads = []
+            exp_avgs = []
+            exp_avg_sqs = []
+            max_exp_avg_sqs = []
+            state_steps = []
+            beta1, beta2 = group['betas']
+
+            for p in group['params']:
+                if p.grad is not None:
+                    params_with_grad.append(p)
+                    if p.grad.is_sparse:
+                        raise RuntimeError('Adam does not support sparse gradients, please consider SparseAdam instead')
+                    grads.append(p.grad)
+
+                    state = self.state[p]
+                    # Lazy state initialization
+                    if len(state) == 0:
+                        state['step'] = 0
+                        # Exponential moving average of gradient values
+                        state['exp_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format)
+                        # Exponential moving average of squared gradient values
+                        state['exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)
+                        if group['amsgrad']:
+                            # Maintains max of all exp. moving avg. of sq. grad. values
+                            state['max_exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)
+
+                    exp_avgs.append(state['exp_avg'])
+                    exp_avg_sqs.append(state['exp_avg_sq'])
+
+                    if group['amsgrad']:
+                        max_exp_avg_sqs.append(state['max_exp_avg_sq'])
+
+                    # update the steps for each param group update
+                    state['step'] += 1
+                    # record the step after step update
+                    state_steps.append(state['step'])
+
+            adam(params_with_grad,
+                   grads,
+                   exp_avgs,
+                   exp_avg_sqs,
+                   max_exp_avg_sqs,
+                   state_steps,
+                   amsgrad=group['amsgrad'],
+                   beta1=beta1,
+                   beta2=beta2,
+                   lr=group['lr'],
+                   weight_decay=group['weight_decay'],
+                   eps=group['eps'])
+        return loss
+
+
+    
+def adam(params: List[Tensor],
+         grads: List[Tensor],
+         exp_avgs: List[Tensor],
+         exp_avg_sqs: List[Tensor],
+         max_exp_avg_sqs: List[Tensor],
+         state_steps: List[int],
+         *,
+         amsgrad: bool,
+         beta1: float,
+         beta2: float,
+         lr: float,
+         weight_decay: float,
+         eps: float):
+    r"""Functional API that performs Adam algorithm computation.
+    See :class:`~torch.optim.Adam` for details.
+    """
+
+    for i, param in enumerate(params):
+
+        grad = grads[i]
+        exp_avg = exp_avgs[i]
+        exp_avg_sq = exp_avg_sqs[i]
+        step = state_steps[i]
+
+        bias_correction1 = 1 - beta1 ** step
+        bias_correction2 = 1 - beta2 ** step
+
+        if weight_decay != 0:
+            grad = grad.add(param, alpha=weight_decay)
+
+        # Decay the first and second moment running average coefficient
+        exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
+        exp_avg_sq.mul_(beta2).addcmul_(grad, grad.conj(), value=1 - beta2)
+        if amsgrad:
+            # Maintains the maximum of all 2nd moment running avg. till now
+            torch.maximum(max_exp_avg_sqs[i], exp_avg_sq, out=max_exp_avg_sqs[i])
+            # Use the max. for normalizing running avg. of gradient
+            denom = (max_exp_avg_sqs[i].sqrt() / math.sqrt(bias_correction2)).add_(eps)
+        else:
+            denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(eps)
+
+        step_size = lr / bias_correction1
+
+        param.addcdiv_(exp_avg, denom, value=-step_size)
diff --git a/CDTools/models/complex_lbfgs.py b/CDTools/models/complex_lbfgs.py
new file mode 100644
index 0000000..2a03696
--- /dev/null
+++ b/CDTools/models/complex_lbfgs.py
@@ -0,0 +1,485 @@
+import torch
+from functools import reduce
+from torch.optim.optimizer import Optimizer
+
+
+def _cubic_interpolate(x1, f1, g1, x2, f2, g2, bounds=None):
+    # ported from https://github.com/torch/optim/blob/master/polyinterp.lua
+    # Compute bounds of interpolation area
+    if bounds is not None:
+        xmin_bound, xmax_bound = bounds
+    else:
+        xmin_bound, xmax_bound = (x1, x2) if x1 <= x2 else (x2, x1)
+
+    # Code for most common case: cubic interpolation of 2 points
+    #   w/ function and derivative values for both
+    # Solution in this case (where x2 is the farthest point):
+    #   d1 = g1 + g2 - 3*(f1-f2)/(x1-x2);
+    #   d2 = sqrt(d1^2 - g1*g2);
+    #   min_pos = x2 - (x2 - x1)*((g2 + d2 - d1)/(g2 - g1 + 2*d2));
+    #   t_new = min(max(min_pos,xmin_bound),xmax_bound);
+    d1 = g1 + g2 - 3 * (f1 - f2) / (x1 - x2)
+    d2_square = d1**2 - g1 * g2
+    if d2_square >= 0:
+        d2 = d2_square.sqrt()
+        if x1 <= x2:
+            min_pos = x2 - (x2 - x1) * ((g2 + d2 - d1) / (g2 - g1 + 2 * d2))
+        else:
+            min_pos = x1 - (x1 - x2) * ((g1 + d2 - d1) / (g1 - g2 + 2 * d2))
+        return min(max(min_pos, xmin_bound), xmax_bound)
+    else:
+        return (xmin_bound + xmax_bound) / 2.
+
+
+def _strong_wolfe(obj_func,
+                  x,
+                  t,
+                  d,
+                  f,
+                  g,
+                  gtd,
+                  c1=1e-4,
+                  c2=0.9,
+                  tolerance_change=1e-9,
+                  max_ls=25):
+    # ported from https://github.com/torch/optim/blob/master/lswolfe.lua
+    d_norm = d.abs().max()
+    g = g.clone(memory_format=torch.contiguous_format)
+    # evaluate objective and gradient using initial step
+    f_new, g_new = obj_func(x, t, d)
+    ls_func_evals = 1
+    gtd_new = g_new.dot(d)
+
+    # bracket an interval containing a point satisfying the Wolfe criteria
+    t_prev, f_prev, g_prev, gtd_prev = 0, f, g, gtd
+    done = False
+    ls_iter = 0
+    while ls_iter < max_ls:
+        # check conditions
+        if f_new > (f + c1 * t * gtd) or (ls_iter > 1 and f_new >= f_prev):
+            bracket = [t_prev, t]
+            bracket_f = [f_prev, f_new]
+            bracket_g = [g_prev, g_new.clone(memory_format=torch.contiguous_format)]
+            bracket_gtd = [gtd_prev, gtd_new]
+            break
+
+        if abs(gtd_new) <= -c2 * gtd:
+            bracket = [t]
+            bracket_f = [f_new]
+            bracket_g = [g_new]
+            done = True
+            break
+
+        if gtd_new >= 0:
+            bracket = [t_prev, t]
+            bracket_f = [f_prev, f_new]
+            bracket_g = [g_prev, g_new.clone(memory_format=torch.contiguous_format)]
+            bracket_gtd = [gtd_prev, gtd_new]
+            break
+
+        # interpolate
+        min_step = t + 0.01 * (t - t_prev)
+        max_step = t * 10
+        tmp = t
+        t = _cubic_interpolate(
+            t_prev,
+            f_prev,
+            gtd_prev,
+            t,
+            f_new,
+            gtd_new,
+            bounds=(min_step, max_step))
+
+        # next step
+        t_prev = tmp
+        f_prev = f_new
+        g_prev = g_new.clone(memory_format=torch.contiguous_format)
+        gtd_prev = gtd_new
+        f_new, g_new = obj_func(x, t, d)
+        ls_func_evals += 1
+        gtd_new = g_new.dot(d)
+        ls_iter += 1
+
+    # reached max number of iterations?
+    if ls_iter == max_ls:
+        bracket = [0, t]
+        bracket_f = [f, f_new]
+        bracket_g = [g, g_new]
+
+    # zoom phase: we now have a point satisfying the criteria, or
+    # a bracket around it. We refine the bracket until we find the
+    # exact point satisfying the criteria
+    insuf_progress = False
+    # find high and low points in bracket
+    low_pos, high_pos = (0, 1) if bracket_f[0] <= bracket_f[-1] else (1, 0)
+    while not done and ls_iter < max_ls:
+        # line-search bracket is so small
+        if abs(bracket[1] - bracket[0]) * d_norm < tolerance_change:
+            break
+
+        # compute new trial value
+        t = _cubic_interpolate(bracket[0], bracket_f[0], bracket_gtd[0],
+                               bracket[1], bracket_f[1], bracket_gtd[1])
+
+        # test that we are making sufficient progress:
+        # in case `t` is so close to boundary, we mark that we are making
+        # insufficient progress, and if
+        #   + we have made insufficient progress in the last step, or
+        #   + `t` is at one of the boundary,
+        # we will move `t` to a position which is `0.1 * len(bracket)`
+        # away from the nearest boundary point.
+        eps = 0.1 * (max(bracket) - min(bracket))
+        if min(max(bracket) - t, t - min(bracket)) < eps:
+            # interpolation close to boundary
+            if insuf_progress or t >= max(bracket) or t <= min(bracket):
+                # evaluate at 0.1 away from boundary
+                if abs(t - max(bracket)) < abs(t - min(bracket)):
+                    t = max(bracket) - eps
+                else:
+                    t = min(bracket) + eps
+                insuf_progress = False
+            else:
+                insuf_progress = True
+        else:
+            insuf_progress = False
+
+        # Evaluate new point
+        f_new, g_new = obj_func(x, t, d)
+        ls_func_evals += 1
+        gtd_new = g_new.dot(d)
+        ls_iter += 1
+
+        if f_new > (f + c1 * t * gtd) or f_new >= bracket_f[low_pos]:
+            # Armijo condition not satisfied or not lower than lowest point
+            bracket[high_pos] = t
+            bracket_f[high_pos] = f_new
+            bracket_g[high_pos] = g_new.clone(memory_format=torch.contiguous_format)
+            bracket_gtd[high_pos] = gtd_new
+            low_pos, high_pos = (0, 1) if bracket_f[0] <= bracket_f[1] else (1, 0)
+        else:
+            if abs(gtd_new) <= -c2 * gtd:
+                # Wolfe conditions satisfied
+                done = True
+            elif gtd_new * (bracket[high_pos] - bracket[low_pos]) >= 0:
+                # old high becomes new low
+                bracket[high_pos] = bracket[low_pos]
+                bracket_f[high_pos] = bracket_f[low_pos]
+                bracket_g[high_pos] = bracket_g[low_pos]
+                bracket_gtd[high_pos] = bracket_gtd[low_pos]
+
+            # new point becomes new low
+            bracket[low_pos] = t
+            bracket_f[low_pos] = f_new
+            bracket_g[low_pos] = g_new.clone(memory_format=torch.contiguous_format)
+            bracket_gtd[low_pos] = gtd_new
+
+    # return stuff
+    t = bracket[low_pos]
+    f_new = bracket_f[low_pos]
+    g_new = bracket_g[low_pos]
+    return f_new, g_new, t, ls_func_evals
+
+
+class MyLBFGS(Optimizer):
+    """Implements L-BFGS algorithm, heavily inspired by `minFunc
+    <https://www.cs.ubc.ca/~schmidtm/Software/minFunc.html>`_.
+    .. warning::
+        This optimizer doesn't support per-parameter options and parameter
+        groups (there can be only one).
+    .. warning::
+        Right now all parameters have to be on a single device. This will be
+        improved in the future.
+    .. note::
+        This is a very memory intensive optimizer (it requires additional
+        ``param_bytes * (history_size + 1)`` bytes). If it doesn't fit in memory
+        try reducing the history size, or use a different algorithm.
+    Args:
+        lr (float): learning rate (default: 1)
+        max_iter (int): maximal number of iterations per optimization step
+            (default: 20)
+        max_eval (int): maximal number of function evaluations per optimization
+            step (default: max_iter * 1.25).
+        tolerance_grad (float): termination tolerance on first order optimality
+            (default: 1e-5).
+        tolerance_change (float): termination tolerance on function
+            value/parameter changes (default: 1e-9).
+        history_size (int): update history size (default: 100).
+        line_search_fn (str): either 'strong_wolfe' or None (default: None).
+    """
+
+    def __init__(self,
+                 params,
+                 lr=1,
+                 max_iter=20,
+                 max_eval=None,
+                 tolerance_grad=1e-7,
+                 tolerance_change=1e-9,
+                 history_size=100,
+                 line_search_fn=None):
+        if max_eval is None:
+            max_eval = max_iter * 5 // 4
+        defaults = dict(
+            lr=lr,
+            max_iter=max_iter,
+            max_eval=max_eval,
+            tolerance_grad=tolerance_grad,
+            tolerance_change=tolerance_change,
+            history_size=history_size,
+            line_search_fn=line_search_fn)
+        super(MyLBFGS, self).__init__(params, defaults)
+
+        if len(self.param_groups) != 1:
+            raise ValueError("LBFGS doesn't support per-parameter options "
+                             "(parameter groups)")
+
+        self._params = self.param_groups[0]['params']
+        self._numel_cache = None
+
+    def _numel(self):
+        if self._numel_cache is None:
+            self._numel_cache = reduce(lambda total, p: total + p.numel(), self._params, 0)
+        return self._numel_cache
+
+    def _gather_flat_grad(self):
+        views = []
+        for p in self._params:
+            if p.grad is None:
+                view = p.new(p.numel()).zero_()
+            elif p.grad.is_sparse:
+                view = p.grad.to_dense().view(-1)
+            else:
+                view = p.grad.view(-1)
+            views.append(view)
+        return torch.cat(views, 0)
+
+    def _add_grad(self, step_size, update):
+        offset = 0
+        for p in self._params:
+            numel = p.numel()
+            if ((update.dtype==torch.complex64 or update.dtype==torch.complex128)
+                and (p.dtype==torch.float32 or p.dtype==torch.float64)):
+                p.add_(update[offset:offset + numel].real.view_as(p), alpha=step_size)
+            else:
+                p.add_(update[offset:offset + numel].view_as(p), alpha=step_size)                
+            # view as to avoid deprecated pointwise semantics
+            #try:
+            #    print(p.dtype)
+            #    print(update.dtype)
+            #    p.add_(update[offset:offset + numel].view_as(p), alpha=step_size)
+            #    print('Worked fine')
+            #    print(update[offset:offset + numel].view_as(p))
+            #except:
+            #    print('Failed')
+            #    print(update[offset:offset + numel].view_as(p))
+            #    exit()
+            offset += numel
+        assert offset == self._numel()
+
+    def _clone_param(self):
+        return [p.clone(memory_format=torch.contiguous_format) for p in self._params]
+
+    def _set_param(self, params_data):
+        for p, pdata in zip(self._params, params_data):
+            p.copy_(pdata)
+
+    def _directional_evaluate(self, closure, x, t, d):
+        self._add_grad(t, d)
+        loss = float(closure())
+        flat_grad = self._gather_flat_grad()
+        self._set_param(x)
+        return loss, flat_grad
+
+    @torch.no_grad()
+    def step(self, closure):
+        """Performs a single optimization step.
+        Args:
+            closure (callable): A closure that reevaluates the model
+                and returns the loss.
+        """
+        assert len(self.param_groups) == 1
+
+        # Make sure the closure is always called with grad enabled
+        closure = torch.enable_grad()(closure)
+
+        group = self.param_groups[0]
+        lr = group['lr']
+        max_iter = group['max_iter']
+        max_eval = group['max_eval']
+        tolerance_grad = group['tolerance_grad']
+        tolerance_change = group['tolerance_change']
+        line_search_fn = group['line_search_fn']
+        history_size = group['history_size']
+
+        # NOTE: LBFGS has only global state, but we register it as state for
+        # the first param, because this helps with casting in load_state_dict
+        state = self.state[self._params[0]]
+        state.setdefault('func_evals', 0)
+        state.setdefault('n_iter', 0)
+
+        # evaluate initial f(x) and df/dx
+        orig_loss = closure()
+        loss = float(orig_loss)
+        current_evals = 1
+        state['func_evals'] += 1
+
+        flat_grad = self._gather_flat_grad()
+        opt_cond = flat_grad.abs().max() <= tolerance_grad
+        
+        # optimal condition
+        if opt_cond:
+            return orig_loss
+
+        # tensors cached in state (for tracing)
+        d = state.get('d')
+        t = state.get('t')
+        old_dirs = state.get('old_dirs')
+        old_stps = state.get('old_stps')
+        ro = state.get('ro')
+        H_diag = state.get('H_diag')
+        prev_flat_grad = state.get('prev_flat_grad')
+        prev_loss = state.get('prev_loss')
+
+        n_iter = 0
+        # optimize for a max of max_iter iterations
+        while n_iter < max_iter:
+            # keep track of nb of iterations
+            n_iter += 1
+            state['n_iter'] += 1
+
+            ############################################################
+            # compute gradient descent direction
+            ############################################################
+            if state['n_iter'] == 1:
+                d = flat_grad.neg()
+                old_dirs = []
+                old_stps = []
+                ro = []
+                H_diag = 1
+            else:
+                # do lbfgs update (update memory)
+                y = flat_grad.sub(prev_flat_grad)
+                s = d.mul(t)
+                ys = y.dot(s)  # y*s
+                if ys.abs() > 1e-10:
+                    # updating memory
+                    if len(old_dirs) == history_size:
+                        # shift history by one (limited-memory)
+                        old_dirs.pop(0)
+                        old_stps.pop(0)
+                        ro.pop(0)
+
+                    # store new direction/step
+                    old_dirs.append(y)
+                    old_stps.append(s)
+                    ro.append(1. / ys)
+
+                    # update scale of initial Hessian approximation
+                    H_diag = ys / y.dot(y)  # (y*y)
+
+                # compute the approximate (L-BFGS) inverse Hessian
+                # multiplied by the gradient
+                num_old = len(old_dirs)
+
+                if 'al' not in state:
+                    state['al'] = [None] * history_size
+                al = state['al']
+
+                # iteration in L-BFGS loop collapsed to use just one buffer
+                q = flat_grad.neg()
+                for i in range(num_old - 1, -1, -1):
+                    al[i] = old_stps[i].dot(q) * ro[i]
+                    q.add_(old_dirs[i], alpha=-al[i])
+
+                # multiply by initial Hessian
+                # r/d is the final direction
+                d = r = torch.mul(q, H_diag)
+                for i in range(num_old):
+                    be_i = old_dirs[i].dot(r) * ro[i]
+                    r.add_(old_stps[i], alpha=al[i] - be_i)
+
+            if prev_flat_grad is None:
+                prev_flat_grad = flat_grad.clone(memory_format=torch.contiguous_format)
+            else:
+                prev_flat_grad.copy_(flat_grad)
+            prev_loss = loss
+
+            ############################################################
+            # compute step length
+            ############################################################
+            # reset initial guess for step size
+            if state['n_iter'] == 1:
+                t = min(1., 1. / flat_grad.abs().sum()) * lr
+            else:
+                t = lr
+
+            # directional derivative
+            gtd = flat_grad.dot(d)  # g * d
+            # Why did this used to be gtd > -tol_change?
+            
+            # directional derivative is below tolerance
+            if gtd.abs() < tolerance_change:
+                break
+
+            # optional line search: user function
+            ls_func_evals = 0
+            if line_search_fn is not None:
+                # perform line search, using user function
+                if line_search_fn != "strong_wolfe":
+                    raise RuntimeError("only 'strong_wolfe' is supported")
+                else:
+                    x_init = self._clone_param()
+
+                    def obj_func(x, t, d):
+                        return self._directional_evaluate(closure, x, t, d)
+
+                    loss, flat_grad, t, ls_func_evals = _strong_wolfe(
+                        obj_func, x_init, t, d, loss, flat_grad, gtd)
+                self._add_grad(t, d)
+                opt_cond = flat_grad.abs().max() <= tolerance_grad
+            else:
+                # no line search, simply move with fixed-step
+                self._add_grad(t, d)
+                if n_iter != max_iter:
+                    # re-evaluate function only if not in last iteration
+                    # the reason we do this: in a stochastic setting,
+                    # no use to re-evaluate that function here
+                    with torch.enable_grad():
+                        loss = float(closure())
+                    flat_grad = self._gather_flat_grad()
+                    opt_cond = flat_grad.abs().max() <= tolerance_grad
+                    ls_func_evals = 1
+
+            # update func eval
+            current_evals += ls_func_evals
+            state['func_evals'] += ls_func_evals
+
+            ############################################################
+            # check conditions
+            ############################################################
+            if n_iter == max_iter:
+                break
+
+            if current_evals >= max_eval:
+                break
+
+            # optimal condition
+            if opt_cond:
+                break
+
+            # lack of progress
+            if d.mul(t).abs().max() <= tolerance_change:
+                break
+
+            if abs(loss - prev_loss) < tolerance_change:
+                break
+
+        state['d'] = d
+        state['t'] = t
+        state['old_dirs'] = old_dirs
+        state['old_stps'] = old_stps
+        state['ro'] = ro
+        state['H_diag'] = H_diag
+        state['prev_flat_grad'] = prev_flat_grad
+        state['prev_loss'] = prev_loss
+
+        return orig_loss
diff --git a/CDTools/models/fancy_ptycho.py b/CDTools/models/fancy_ptycho.py
index 5d0ada9..a64e50d 100644
--- a/CDTools/models/fancy_ptycho.py
+++ b/CDTools/models/fancy_ptycho.py
@@ -4,7 +4,6 @@ import torch as t
 from CDTools.models import CDIModel
 from CDTools.datasets import Ptycho2DDataset
 from CDTools import tools
-from CDTools.tools import cmath
 from CDTools.tools import plotting as p
 from CDTools.tools import analysis
 from matplotlib import pyplot as plt
@@ -55,21 +54,21 @@ class FancyPtycho(CDIModel):
         
         # We rescale the probe here so it learns at the same rate as the
         # object
-        if probe_guess.dim() > 3:
-            self.probe_norm = 1 * t.max(tools.cmath.cabs(probe_guess[0].to(t.float32)))
+        if probe_guess.dim() > 2:
+            self.probe_norm = 1 * t.max(t.abs(probe_guess[0].to(t.complex64)))
         else:
-            self.probe_norm = 1 * t.max(tools.cmath.cabs(probe_guess.to(t.float32)))         
+            self.probe_norm = 1 * t.max(t.abs(probe_guess.to(t.complex64)))         
             
-        self.probe = t.nn.Parameter(probe_guess.to(t.float32)
+        self.probe = t.nn.Parameter(probe_guess.to(t.complex64)
                                     / self.probe_norm)
         
-        self.obj = t.nn.Parameter(obj_guess.to(t.float32))
+        self.obj = t.nn.Parameter(obj_guess.to(t.complex64))
         
         if background is None:
             if detector_slice is not None:
-                background = 1e-6 * t.ones(self.probe[0][self.detector_slice].shape[:-1])
+                background = 1e-6 * t.ones(self.probe[0][self.detector_slice])
             else:
-                background = 1e-6 * t.ones(self.probe[0].shape[:-1])
+                background = 1e-6 * t.ones(self.probe[0])
 
                 
         self.background = t.nn.Parameter(t.Tensor(background).to(t.float32))
@@ -85,9 +84,10 @@ class FancyPtycho(CDIModel):
             else:
                 # Now this is a matrix of weights, so we 
                 if type(weights) == type(t.zeros(1)):
-                    self.weights = t.nn.Parameter(weights.to(t.float32))        
+                    self.weights = t.nn.Parameter(weights.to(t.complex64))        
                 else:
-                    self.weights = t.nn.Parameter(cmath.complex_to_torch(weights).to(t.float32))
+                    # There is a good chance that this doesn't work
+                    self.weights = t.nn.Parameter(t.Tensor(weights).to(t.complex64))
 
                 
         if translation_offsets is None:
@@ -110,6 +110,7 @@ class FancyPtycho(CDIModel):
 
         self.oversampling = oversampling
 
+        
         # Here we set the appropriate loss function
         if loss.lower().strip() == 'amplitude mse'\
            or loss.lower().strip() == 'amplitude_mse':
@@ -192,12 +193,12 @@ class FancyPtycho(CDIModel):
 
 
         # Now we initialize all the subdominant probe modes
-        probe_max = t.max(cmath.cabs(probe))
+        probe_max = t.max(t.abs(probe))
         probe_stack = [0.01 * probe_max * t.rand(probe.shape,dtype=probe.dtype) for i in range(n_modes - 1)]
         probe = t.stack([probe,] + probe_stack)
         #probe = t.stack([tools.propagators.far_field(probe),] + probe_stack)
 
-        obj = tools.cmath.expi(randomize_ang * (t.rand(obj_size)-0.5))
+        obj = t.exp(1j*randomize_ang * (t.rand(obj_size)-0.5))
                               
         det_geo = dataset.detector_geometry
 
@@ -209,11 +210,12 @@ class FancyPtycho(CDIModel):
             elif dm_rank == -1:
                 # dm_rank == -1 is defined to mean full-rank
                 dm_rank = n_modes
-            Ws = t.zeros(len(dataset),dm_rank,n_modes,2)
+                
+            Ws = t.zeros(len(dataset),dm_rank,n_modes,dtype=t.complex64)
             # Start with as close to the identity matrix as possible,
             # cutting of when we hit the specified maximum rank
             for i in range(0,dm_rank):
-                Ws[:,i,i,0] = 1
+                Ws[:,i,i] = 1
         else:
             # dm_rank == None or dm_rank = 0 triggers a special case where
             # a standard incoherent multi-mode model is used. This is the
@@ -228,7 +230,7 @@ class FancyPtycho(CDIModel):
 
         if probe_support_radius is not None:
             probe_support = t.zeros_like(probe[0])
-            xs, ys = np.mgrid[:probe.shape[-3],:probe.shape[-2]]
+            xs, ys = np.mgrid[:probe.shape[-2],:probe.shape[-1]]
             xs = xs - np.mean(xs)
             ys = ys - np.mean(ys)
             Rs = np.sqrt(xs**2 + ys**2)
@@ -282,24 +284,25 @@ class FancyPtycho(CDIModel):
 
         # Now we construct the probes for each shot from the basis probes
         Ws = self.weights[index]
-
         if len(self.weights[0].shape) == 0:
             # If a purely stable coherent illumination is defined
-            # No cmult because Ws is real in this case
-            prs = Ws[...,None,None,None,None] * basis_prs
+            prs = Ws[...,None,None,None] * basis_prs
         else:
             # If a frame-by-frame weight matrix is defined
             # This takes the dot product of all the weight matrices with
             # the probes. The output has dimensions of translation, then
             # coherent mode index, then x,y, and then complex index
-            prs = t.sum(cmath.cmult(Ws[...,None,None,:], basis_prs),
-                    axis=-4)
+            # Maybe this can be done with a matmul now?
+            prs = t.sum(Ws[...,None,None] * basis_prs, axis=-3)
 
         # Now we actually do the interaction, using the sinc subpixel
         # translation model as per usual
         exit_waves = self.probe_norm * tools.interactions.ptycho_2D_sinc(
             prs, self.obj_support * self.obj,pix_trans,
             shift_probe=True, multiple_modes=True)
+        #exit_waves = self.probe_norm * tools.interactions.ptycho_2D_round(
+        #    prs, self.obj_support * self.obj,pix_trans,
+        #    multiple_modes=True)
         
         return exit_waves
     
@@ -387,7 +390,7 @@ class FancyPtycho(CDIModel):
 
     
     def corrected_translations(self,dataset):
-        translations = dataset.translations.to(dtype=self.probe.dtype,device=self.probe.device)
+        translations = dataset.translations.to(dtype=t.float32,device=self.probe.device)
         t_offset = tools.interactions.pixel_to_translations(self.probe_basis,self.translation_offsets*self.translation_scale,surface_normal=self.surface_normal)
         return translations + t_offset
 
@@ -395,7 +398,7 @@ class FancyPtycho(CDIModel):
     def get_rhos(self):
         # If this is the general unified mode model
         if self.weights.dim() >= 2:
-            Ws = cmath.torch_to_complex(self.weights.detach().cpu())
+            Ws = self.weights.detach().cpu().numpy()
             rhos_out = np.matmul(np.swapaxes(Ws,1,2), Ws.conj())
             return rhos_out
         # This is the purely incoherent case
@@ -418,10 +421,10 @@ class FancyPtycho(CDIModel):
         # to catch this case, but because it's so much simpler than the
         # unified mode case I think it's appropriate
         if self.weights.dim() == 1:
-            probe = cmath.torch_to_complex(self.probe.detach().cpu())
+            probe = self.probe.detach().cpu().numpy()
             ortho_probes = analysis.orthogonalize_probes(probe)
-            self.probe.data = cmath.complex_to_torch(ortho_probes).to(
-            device=self.probe.device,dtype=self.probe.dtype)
+            self.probe.data = t.as_tensor(ortho_probes,
+                    device=self.probe.device,dtype=self.probe.dtype)
             return
 
         # This is for the unified mode case
@@ -432,7 +435,7 @@ class FancyPtycho(CDIModel):
         
         rhos = self.get_rhos()
         overall_rho = np.mean(rhos,axis=0)
-        probe = cmath.torch_to_complex(self.probe.detach().cpu())
+        probe = self.probe.detach().cpu().numpy()
         ortho_probes, A = analysis.orthogonalize_probes(probe,
                                         density_matrix=overall_rho,
                                         keep_transform=True,
@@ -465,37 +468,35 @@ class FancyPtycho(CDIModel):
             
         new_Ws = np.array(new_Ws)
 
-        self.weights.data = cmath.complex_to_torch(new_Ws).to(
+        self.weights.data = t.as_tensor(new_Ws,
             dtype=self.weights.dtype,device=self.weights.device)
         
-        self.probe.data = cmath.complex_to_torch(ortho_probes).to(
+        self.probe.data = t.as_tensor(ortho_probes,
             device=self.probe.device,dtype=self.probe.dtype)
 
 
     def plot_wavefront_variation(self, dataset,fig=None,mode='amplitude',**kwargs):
         def get_probes(idx):
             basis_prs = self.probe * self.probe_support[...,:,:]
-            prs = t.sum(cmath.cmult(self.weights[idx,:,:,None,None,:],
-                                    basis_prs), axis=-4)
+            prs = t.sum(self.weights[idx,:,:,None,None] * basis_prs, axis=-4)
             ortho_probes = analysis.orthogonalize_probes(prs)
 
-            #return np.abs(cmath.torch_to_complex(prs.detach().cpu()))
             if mode.lower() == 'amplitude':
-                return np.abs(cmath.torch_to_complex(ortho_probes.detach().cpu()))
+                return np.abs(ortho_probes.detach().cpu().numpy())
             if mode.lower() == 'root_sum_intensity':
-                return np.sum(np.abs(cmath.torch_to_complex(ortho_probes.detach().cpu()))**2,axis=0)
+                return np.sum(np.abs(ortho_probes.detach().cpu().numpy())**2,axis=0)
             if mode.lower() == 'phase':
-                return np.angle(cmath.torch_to_complex(ortho_probes.detach().cpu()))
+                return np.angle(ortho_probes.detach().cpu().numpy())
             
         probe_matrix = np.zeros([self.probe.shape[0]]*2,
                                 dtype=np.complex64)
-        np_probes = cmath.torch_to_complex(self.probe.detach().cpu())
+        np_probes = self.probe.detach().cpu().numpy()
         for i in range(probe_matrix.shape[0]):
             for j in range(probe_matrix.shape[0]):
                 probe_matrix[i,j] = np.sum(np_probes[i]*np_probes[j].conj())
         
 
-        weights = cmath.torch_to_complex(self.weights.detach().cpu())
+        weights = self.weights.detach().cpu().numpy()
         
         probe_intensities = np.sum(np.tensordot(weights,probe_matrix,axes=1)*
                                    weights.conj(),axis=2)
@@ -546,14 +547,11 @@ class FancyPtycho(CDIModel):
     def save_results(self, dataset):
         basis = self.probe_basis.detach().cpu().numpy()
         translations = self.corrected_translations(dataset).detach().cpu().numpy()
-        probe = cmath.torch_to_complex(self.probe.detach().cpu())
+        probe = self.probe.detach().cpu().numpy()
         probe = probe * self.probe_norm.detach().cpu().numpy()
-        obj = cmath.torch_to_complex(self.obj.detach().cpu())
+        obj = self.obj.detach().cpu().numpy()
         background = self.background.detach().cpu().numpy()**2
-        if len(self.weights.shape) >=2:
-            weights = cmath.torch_to_complex(self.weights.detach().cpu())
-        else:
-            weights = self.weights.detach().cpu().numpy()
+        weights = self.weights.detach().cpu().numpy()
         
         return {'basis':basis, 'translation':translations,
                 'probe':probe,'obj':obj,
diff --git a/CDTools/models/simple_ptycho.py b/CDTools/models/simple_ptycho.py
index baa662c..a178d78 100644
--- a/CDTools/models/simple_ptycho.py
+++ b/CDTools/models/simple_ptycho.py
@@ -10,6 +10,7 @@ from torch.utils import data as torchdata
 from matplotlib import pyplot as plt
 from datetime import datetime
 import numpy as np
+from .complex_adam import MyAdam
 
 __all__ = ['SimplePtycho']
 
@@ -25,22 +26,22 @@ class SimplePtycho(CDIModel):
                  surface_normal=np.array([0.,0.,1.]), mask=None):
 
         super(SimplePtycho,self).__init__()
-        self.wavelength = t.Tensor([wavelength])
+        self.wavelength = t.tensor([wavelength])
         self.detector_geometry = copy(detector_geometry)
         det_geo = self.detector_geometry
         if hasattr(det_geo, 'distance'):
-            det_geo['distance'] = t.Tensor(det_geo['distance'])
+            det_geo['distance'] = t.tensor(det_geo['distance'])
         if hasattr(det_geo, 'basis'):
-            det_geo['basis'] = t.Tensor(det_geo['basis'])
+            det_geo['basis'] = t.tensor(det_geo['basis'])
         if hasattr(det_geo, 'corner'):
-            det_geo['corner'] = t.Tensor(det_geo['corner'])
+            det_geo['corner'] = t.tensor(det_geo['corner'])
 
-        self.min_translation = t.Tensor(min_translation)
+        self.min_translation = t.tensor(min_translation)
 
-        self.probe_basis = t.Tensor(probe_basis)
+        self.probe_basis = t.tensor(probe_basis)
         self.detector_slice = detector_slice
 
-        self.surface_normal = t.Tensor(surface_normal)
+        self.surface_normal = t.tensor(surface_normal)
 
         if mask is None:
             self.mask = None
@@ -49,11 +50,11 @@ class SimplePtycho(CDIModel):
 
         # We rescale the probe here so it learns at the same rate as the
         # object
-        self.probe_norm = t.max(tools.cmath.cabs(probe_guess.to(t.float32)))
+        self.probe_norm = t.max(t.abs(probe_guess.to(t.complex64)))
 
-        self.probe = t.nn.Parameter(probe_guess.to(t.float32)
+        self.probe = t.nn.Parameter(probe_guess.to(t.complex64)
                                     / self.probe_norm)
-        self.obj = t.nn.Parameter(obj_guess.to(t.float32))
+        self.obj = t.nn.Parameter(obj_guess.to(t.complex64))
 
 
 
@@ -94,8 +95,8 @@ class SimplePtycho(CDIModel):
 
         # Finally, initialize the probe and  object using this information
         probe = tools.initializers.SHARP_style_probe(dataset, probe_shape, det_slice)
-
-        obj = t.ones(obj_size+(2,))
+        
+        obj = t.ones(obj_size).to(dtype=t.complex64)
         det_geo = dataset.detector_geometry
 
 
@@ -112,6 +113,7 @@ class SimplePtycho(CDIModel):
                                                              translations,
                                             surface_normal=self.surface_normal)
         pix_trans -= self.min_translation
+
         return tools.interactions.ptycho_2D_round(self.probe_norm * self.probe,
                                                   self.obj,
                                                   pix_trans)
@@ -130,7 +132,7 @@ class SimplePtycho(CDIModel):
                                             detector_slice=self.detector_slice)
 
 
-    def loss(self, sim_data, real_data, mask=None):
+    def loss(self, real_data, sim_data, mask=None):
         return tools.losses.amplitude_mse(real_data, sim_data, mask=mask)
 
 
@@ -206,9 +208,9 @@ class SimplePtycho(CDIModel):
 
 
     def save_results(self):
-        probe = tools.cmath.torch_to_complex(self.probe.detach().cpu())
+        probe = self.probe.detach().cpu().numpy()
         probe = probe * self.probe_norm.detach().cpu().numpy()
-        obj = tools.cmath.torch_to_complex(self.obj.detach().cpu())
+        obj = self.obj.detach().cpu().numpy()
         return {'probe':probe,'obj':obj}
 
 
@@ -272,4 +274,4 @@ class SimplePtycho(CDIModel):
                     # Calculate loss
                     loss.append(self.loss(self.measurement(self.interaction(i, translations)), patterns))
 
-                yield t.mean(t.Tensor(loss)).cpu().numpy()
+                yield t.mean(t.tensor(loss)).cpu().numpy()
diff --git a/CDTools/tools/__init__.py b/CDTools/tools/__init__.py
index a94f74b..c65ac1a 100644
--- a/CDTools/tools/__init__.py
+++ b/CDTools/tools/__init__.py
@@ -17,7 +17,6 @@ having numpy and torch defined under CDTools.tools.cmath, you know?
 from __future__ import division, print_function, absolute_import
 
 
-from CDTools.tools import cmath
 from CDTools.tools import losses
 from CDTools.tools import data
 from CDTools.tools import image_processing
diff --git a/CDTools/tools/analysis/analysis.py b/CDTools/tools/analysis/analysis.py
index 854927b..49f2e50 100644
--- a/CDTools/tools/analysis/analysis.py
+++ b/CDTools/tools/analysis/analysis.py
@@ -9,7 +9,6 @@ from __future__ import division, print_function
 
 import torch as t
 import numpy as np
-from CDTools.tools import cmath
 from CDTools.tools import image_processing as ip
 from scipy import fftpack
 from scipy import linalg as sla
@@ -63,7 +62,7 @@ def orthogonalize_probes(probes, density_matrix=None, keep_transform=False, norm
     """
     
     try:
-        probes = cmath.torch_to_complex(probes.detach().cpu())
+        probes = probes.detach().cpu().numpy()
         send_to_torch = True
     except:
         send_to_torch = False
@@ -116,9 +115,8 @@ def orthogonalize_probes(probes, density_matrix=None, keep_transform=False, norm
         #A_dagger = np.dot(np.transpose(u).conj(),B_dagger)
     
     if send_to_torch:
-        ortho_probes = cmath.complex_to_torch(np.stack(ortho_probes))
-        A = cmath.complex_to_torch(A)
-        #A_dagger = cmath.complex_to_torch(A_dagger)
+        ortho_probes = t.as_tensor(np.stack(ortho_probes))
+        A = t.as_tensor(A)
 
     if keep_transform:
         return ortho_probes, A#_dagger
@@ -172,11 +170,11 @@ def standardize(probe, obj, obj_slice=None, correct_ramp=False):
     # First, we normalize the probe intensity to a fixed value.
     probe_np = False
     if isinstance(probe, np.ndarray):
-        probe = cmath.complex_to_torch(probe).to(t.float32)
+        probe = t.Tensor(probe).to(t.complex64)
         probe_np = True
     obj_np = False
     if isinstance(obj, np.ndarray):
-        obj = cmath.complex_to_torch(obj).to(t.float32)
+        obj = t.Tensor(obj).to(t.complex64)
         obj_np = True
 
     # If this is a single probe and not a stack of probes
@@ -186,7 +184,7 @@ def standardize(probe, obj, obj_slice=None, correct_ramp=False):
     else:
         single_probe = False
 
-    normalization = t.sqrt(t.sum(cmath.cabssq(probe[0])) / (len(probe[0].view(-1))/2))
+    normalization = t.sqrt(t.sum(t.abs(probe[0])**2) / (len(probe[0].view(-1))/2))
     probe = probe / normalization
     obj = obj * normalization
 
@@ -198,37 +196,38 @@ def standardize(probe, obj, obj_slice=None, correct_ramp=False):
 
     if correct_ramp:
         # Need to check if this is actually working and, if not, why not
-        center_freq = ip.centroid(cmath.cabssq(cmath.fftshift(t.fft(probe[0],2))))
+        center_freq = ip.centroid(t.abs(t.fft.fftshift(t.fft.fft2(probe[0]),
+                                                       dim=(-1,-2)))**2)
         center_freq -= (t.tensor(probe[0].shape[:-1]) // 2).to(t.float32)
         center_freq /= t.tensor(probe[0].shape[:-1]).to(t.float32)
 
         Is, Js = np.mgrid[:probe[0].shape[0],:probe[0].shape[1]]
-        probe_phase_ramp = cmath.expi(2 * np.pi *
-                                      (center_freq[0] * t.tensor(Is).to(t.float32) +
-                                       center_freq[1] * t.tensor(Js).to(t.float32)))
-        probe = cmath.cmult(probe, cmath.cconj(probe_phase_ramp))
+        probe_phase_ramp = t.exp(2j * np.pi *
+                                 (center_freq[0] * t.tensor(Is).to(t.float32) +
+                                  center_freq[1] * t.tensor(Js).to(t.float32)))
+        probe = probe *  t.conj(probe_phase_ramp)
         Is, Js = np.mgrid[:obj.shape[0],:obj.shape[1]]
-        obj_phase_ramp = cmath.expi(2*np.pi *
-                                    (center_freq[0] * t.tensor(Is).to(t.float32) +
-                                     center_freq[1] * t.tensor(Js).to(t.float32)))
-        obj = cmath.cmult(obj, obj_phase_ramp)
+        obj_phase_ramp = t.exp(2j*np.pi *
+                               (center_freq[0] * t.tensor(Is).to(t.float32) +
+                                center_freq[1] * t.tensor(Js).to(t.float32)))
+        obj = obj * obj_phase_ramp
 
     # Then, we set them to consistent absolute phases
 
-    obj_angle = cmath.cphase(t.sum(obj[obj_slice],dim=(0,1)))
-    obj = cmath.cmult(obj, cmath.expi(-obj_angle))
+    obj_angle = t.angle(t.sum(obj[obj_slice],dim=(0,1)))
+    obj = obj *  t.exp(-1j*obj_angle)
 
     for i in range(probe.shape[0]):
-        probe_angle = cmath.cphase(t.sum(probe[i],dim=(0,1)))
-        probe[i] = cmath.cmult(probe[i], cmath.expi(-probe_angle))
+        probe_angle = t.angle(t.sum(probe[i],dim=(0,1)))
+        probe[i] = probe[i] * t.exp(-1j*probe_angle)
 
     if single_probe:
         probe = probe[0]
 
     if probe_np:
-        probe = cmath.torch_to_complex(probe.detach().cpu())
+        probe = probe.detach().cpu().numpy()
     if obj_np:
-        obj = cmath.torch_to_complex(obj.detach().cpu())
+        obj = obj.detach().cpu().numpy()
 
     return probe, obj
 
@@ -268,11 +267,11 @@ def synthesize_reconstructions(probes, objects, use_probe=False, obj_slice=None,
 
     probe_np = False
     if isinstance(probes[0], np.ndarray):
-        probes = [cmath.complex_to_torch(probe).to(t.float32) for probe in probes]
+        probes = [t.Tensor(probe).to(t.complex64) for probe in probes]
         probe_np = True
     obj_np = False
     if isinstance(objects[0], np.ndarray):
-        objects = [cmath.complex_to_torch(obj).to(t.float32) for obj in objects]
+        objects = [t.Tensor(obj).to(t.complex64) for obj in objects]
         obj_np = True
 
     obj_shape = np.min(np.array([obj.shape[:-1] for obj in objects]),axis=0)
@@ -317,10 +316,10 @@ def synthesize_reconstructions(probes, objects, use_probe=False, obj_slice=None,
         i = -1
 
     if probe_np:
-        synth_probe = cmath.torch_to_complex(synth_probe)
+        synth_probe = synth_probe.numpy()
     if obj_np:
-        synth_obj = cmath.torch_to_complex(synth_obj)
-        obj_stack = [cmath.torch_to_complex(obj) for obj in obj_stack]
+        synth_obj = synth_obj.numpy()
+        obj_stack = [obj.numpy() for obj in obj_stack]
 
     return synth_probe/(i+2), synth_obj/(i+2), obj_stack
 
@@ -358,10 +357,10 @@ def calc_consistency_prtf(synth_obj, objects, basis, obj_slice=None,nbins=None):
 
     obj_np = False
     if isinstance(objects[0], np.ndarray):
-        objects = [cmath.complex_to_torch(obj).to(t.float32) for obj in objects]
+        objects = [t.Tensor(obj).to(t.complex64) for obj in objects]
         obj_np = True
     if isinstance(synth_obj, np.ndarray):
-        synth_obj = cmath.complex_to_torch(synth_obj).to(t.float32)
+        synth_obj = t.Tensor(synth_obj).to(t.complex64)
 
     if isinstance(basis, t.Tensor):
         basis = basis.detach().cpu().numpy()
@@ -373,7 +372,7 @@ def calc_consistency_prtf(synth_obj, objects, basis, obj_slice=None,nbins=None):
     if nbins is None:
         nbins = np.max(synth_obj[obj_slice].shape) // 4
 
-    synth_fft = cmath.cabssq(cmath.fftshift(t.fft(synth_obj[obj_slice],2))).numpy()
+    synth_fft = (t.abs(t.fft.fftshift(t.fft.fft2(synth_obj[obj_slice]), dim=(-1,-2)))**2).numpy()
 
 
     di = np.linalg.norm(basis[:,0])
@@ -391,7 +390,8 @@ def calc_consistency_prtf(synth_obj, objects, basis, obj_slice=None,nbins=None):
     prtfs = []
     for obj in objects:
         obj = obj[obj_slice]
-        single_fft = cmath.cabssq(cmath.fftshift(t.fft(obj,2))).numpy()
+        single_fft = (t.abs(t.fft.fftshift(t.fft.fft2(obj),
+                                           dim=(-1,-2)))**2).numpy()
         single_ints, bins = np.histogram(Rs,bins=nbins,weights=single_fft)
 
         prtfs.append(synth_ints/single_ints)
@@ -433,10 +433,10 @@ def calc_deconvolved_cross_correlation(im1, im2, im_slice=None):
 
     im_np = False
     if isinstance(im1, np.ndarray):
-        im1 = cmath.complex_to_torch(im1)
+        im1 = t.Tensor(im1)
         im_np = True
     if isinstance(im2, np.ndarray):
-        im2 = cmath.complex_to_torch(im2)
+        im2 = t.Tensor(im2)
         im_np = True
 
     # If last dimension is not 2, then convert to a complex tensor now
@@ -450,15 +450,15 @@ def calc_deconvolved_cross_correlation(im1, im2, im_slice=None):
                           (im1.shape[1]//8)*3:(im1.shape[1]//8)*5]
 
 
-    cor_fft = cmath.cmult(t.fft(im1[im_slice],2),
-                          cmath.cconj(t.fft(im2[im_slice],2)))
+    cor_fft = t.fft.fft2(im1[im_slice]) * \
+        t.conj(t.fft.fft2(im2[im_slice]))
 
     # Not sure if this is more or less stable than just the correlation
     # maximum - requires some testing
-    cor = t.ifft(cor_fft / cmath.cabs(cor_fft)[:,:,None],2)
+    cor = t.fft.ifft2(cor_fft / t.abs(cor_fft))
 
     if im_np:
-        cor = cmath.torch_to_complex(cor)
+        cor = cor.numpy()
 
     return cor
 
@@ -500,10 +500,10 @@ def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1.):
 
     im_np = False
     if isinstance(im1, np.ndarray):
-        im1 = cmath.complex_to_torch(im1)
+        im1 = t.Tensor(im1)
         im_np = True
     if isinstance(im2, np.ndarray):
-        im2 = cmath.complex_to_torch(im2)
+        im2 = t.Tensor(im2)
         im_np = True
 
     if isinstance(basis, np.ndarray):
@@ -524,17 +524,11 @@ def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1.):
         nbins = np.max(im1[im_slice].shape) // 4
 
 
-    cor_fft = cmath.cmult(cmath.fftshift(t.fft(im1[im_slice],2)),
-                          cmath.fftshift(cmath.cconj(t.fft(im2[im_slice],2))))
+    cor_fft = t.fft.fftshift(t.fft.fft2(im1[im_slice]),dim=(-1,-2)) * \
+        t.fft.fftshift(t.conj(t.fft.fft2(im2[im_slice])),dim=(-1,-2))
 
-    #from matplotlib import pyplot as plt
-    #plt.imshow(cmath.cphase(cor_fft))
-    #plt.figure()
-    #plt.imshow(np.log(cmath.cabs(cor_fft)))
-    #plt.show()
-    
-    F1 = cmath.cabs(cmath.fftshift(t.fft(im1[im_slice],2)))**2
-    F2 = cmath.cabs(cmath.fftshift(t.fft(im2[im_slice],2)))**2
+    F1 = t.abs(t.fft.fftshift(t.fft.fft2(im1[im_slice]),dim=(-1,-2)))**2
+    F2 = t.abs(t.fft.fftshift(t.fft.fft2(im2[im_slice]),dim=(-1,-2)))**2
 
 
     di = np.linalg.norm(basis[:,0])
@@ -548,7 +542,7 @@ def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1.):
 
 
 
-    numerator, bins = np.histogram(Rs,bins=nbins,weights=cmath.torch_to_complex(cor_fft))
+    numerator, bins = np.histogram(Rs,bins=nbins,weights=cor_fft.numpy())
     denominator_F1, bins = np.histogram(Rs,bins=nbins,weights=F1.detach().cpu().numpy())
     denominator_F2, bins = np.histogram(Rs,bins=nbins,weights=F2.detach().cpu().numpy())
     n_pix, bins = np.histogram(Rs,bins=nbins)
diff --git a/CDTools/tools/cmath/__init__.py b/CDTools/tools/cmath/__init__.py
deleted file mode 100644
index 828fd09..0000000
--- a/CDTools/tools/cmath/__init__.py
+++ /dev/null
@@ -1,3 +0,0 @@
-from __future__ import division, print_function, absolute_import
-
-from CDTools.tools.cmath.cmath import *
diff --git a/CDTools/tools/cmath/cmath.py b/CDTools/tools/cmath/cmath.py
deleted file mode 100644
index 282b840..0000000
--- a/CDTools/tools/cmath/cmath.py
+++ /dev/null
@@ -1,327 +0,0 @@
-"""Contains basic functions for dealing with complex numbers in pytorch.
-
-Since pytorch doesn't have built-in support for complex numbers, but the
-fast fourier transforms in pytorch assume a specific format for complex
-arrays, this module uses that format to store complex numbers. It exposes
-functions for converting between complex numpy arrays and torch tensors
-stored in that format, as well as basic complex math operations implemented
-on the torch tensors
-"""
-from __future__ import division, print_function, absolute_import
-import numpy as np
-import torch as t
-
-
-__all__ = ['complex_to_torch', 'torch_to_complex', 'cabssq', 'cabs', 'cconj',
-           'cmult', 'cdiv', 'cphase', 'fftshift', 'ifftshift', 'expi', 'cexpi']
-
-
-#
-# These define the conversions to and from this format
-#
-
-def complex_to_torch(x):
-    """Maps a complex numpy array to a torch tensor
-
-    Pytorch uses tensors with a final dimension of 2 to represent
-    complex numbers. This maps a complex type numpy array to a torch
-    tensor following this convention
-
-    Parameters
-    ----------
-    x : np.ndarray
-        A numpy array to convert
-
-    Returns
-    -------
-    torch.Tensor
-        A torch tensor representation of that array
-
-    """
-    return t.from_numpy(np.stack((np.real(x),np.imag(x)),axis=-1))
-
-
-def torch_to_complex(x):
-    """Maps a torch tensor to the a complex numpy array
-
-    Pytorch uses tensors with a final dimension of 2 to represent
-    complex numbers. This maps a torch tensor following that convention
-    to the appropriate numpy complex array. Note that, in order for this
-    function to work, the tensor must be detached from any parameters and
-    living on the CPU.
-
-    Parameters
-    ----------
-    x : torch.Tensor
-        A tensor to convert
-
-    Returns
-    -------
-    np.array
-        A complex typed numpy array corresponding to the input
-
-    """
-    x = np.array(x)
-    x = x[...,0] + x[...,1] * 1j
-    return x
-
-
-#
-# And these define the basic operations on these arrays. Note that
-# multiplication between a complex valued and real valued pytorch
-# tensor will proceed as expected because of torch's broadcasting
-# and thus doesn't need it's own function
-#
-
-def cabssq(x):
-    """Returns the square of the absolute value of a complex torch tensor
-
-    Pytorch uses tensors with a final dimension of 2 to represent
-    complex numbers. This calculates the elementwise absolute value
-    squared of any toch tensor following that standard.
-
-    Parameters
-    ----------
-    x : torch.Tensor
-        An input tensor
-
-    Returns
-    -------
-    torch.Tensor
-        A tensor storing the elementwise absolute value squared
-
-    """
-    return x[...,0]**2 + x[...,1]**2
-
-
-def cabs(x):
-    """Returns the absolute value of a complex torch tensor
-
-    Pytorch uses tensors with a final dimension of 2 to represent
-    complex numbers. This calculates the elementwise absolute value
-    of any torch tensor following that standard.
-
-    Parameters
-    ----------
-    x : torch.Tensor
-        An input tensor
-
-    Returns
-    -------
-    torch.Tensor
-        A tensor storing the elementwise absolute value
-
-    """
-    return t.sqrt(cabssq(x))
-
-
-def cphase(x):
-    """Returns the phase of a complex torch tensor
-
-    Pytorch uses tensors with a final dimension of 2 to represent
-    complex numbers. This calculates the elementwise complex phase
-    of any torch tensor following that standard.
-
-    Parameters
-    ----------
-    x : torch.Tensor
-        An input tensor
-
-    Returns
-    -------
-    torch.Tensor
-        A tensor storing the elementwise phase
-
-    """
-    return t.atan2(x[...,1],x[...,0])
-
-
-def cconj(x):
-    """Returns the complex conjugate of a complex torch tensor
-
-    Pytorch uses tensors with a final dimension of 2 to represent
-    complex numbers. This calculates the elementwise complex conjugate
-    of any torch tensor following that standard.
-
-    Parameters
-    ----------
-    x : torch.Tensor
-        An input tensor
-
-    Returns
-    -------
-    torch.Tensor
-        A tensor storing the elementwise complex conjugate
-
-    """
-    return t.stack((x[...,0],-x[...,1]),dim=-1)
-
-
-
-def cmult(a,b):
-    """Returns the complex product of two torch tensors
-
-    Pytorch uses tensors with a final dimension of 2 to represent
-    complex numbers. This calculates the elementwise product
-    of two torch tensors following that standard.
-
-    Parameters
-    ----------
-    a : torch.Tensor
-        An input tensor
-    b : torch.Tensor
-        A second input tensor
-
-    Returns
-    -------
-    torch.Tensor
-        A tensor storing the elementwise product
-
-    """
-
-    real = a[...,0] * b[...,0] - a[...,1] * b[...,1]
-    imag = a[...,0] * b[...,1] + a[...,1] * b[...,0]
-    return t.stack((real,imag),dim=-1)
-
-
-def cdiv(a,b):
-    """Returns the complex quotient of two torch tensors
-
-    Pytorch uses tensors with a final dimension of 2 to represent
-    complex numbers. This calculates the elementwise quotient
-    of two torch tensors following that standard.
-
-    Parameters
-    ----------
-    a : torch.Tensor
-        An input tensor
-    b : torch.Tensor
-        A second input tensor
-
-    Returns
-    -------
-    torch.Tensor
-        A tensor storing the elementwise complex quotient
-
-    """
-    return cmult(a, cconj(b)) / t.unsqueeze(cabssq(b),-1)
-
-
-
-#
-# Not entirely sure if these belong here, but heck with it.
-# We just need the ability to do fftshifts
-#
-
-
-def fftshift(array,dims=None):
-    """Drop-in torch replacement for scipy.fftpack.fftshift
-
-    This maps a tensor, assumed to be the output of a fast Fourier
-    transform, into a tensor whose zero-frequency element is at the
-    center of the tensor instead of the start. It will by default shift
-    every dimension in the tensor but the last (which is assumed to
-    represent the complex number and be of dimension 2), but can shift
-    any arbitrary set of dimensions.
-
-    Parameters
-    ----------
-    array : torch.Tensor
-        An array of data to be fftshifted
-    dims : iterable
-        A list of all dimensions to shift
-
-    Returns
-    -------
-    torch.Tensor
-        The fftshifted tensor
-
-    """
-
-    if dims is None:
-        dims=list(range(array.dim()))[:-1]
-    for dim in dims:
-        length = array.size()[dim]
-        cut_to = (length + 1) // 2
-        cut_len = length - cut_to
-        array = t.cat((array.narrow(dim,cut_to,cut_len),
-                       array.narrow(dim,0,cut_to)), dim)
-    return array
-
-
-
-def ifftshift(array,dims=None):
-    """Drop-in torch replacement for scipy.fftpack.iftshift
-
-    This maps a tensor, assumed to be the shifted output of a fast
-    Fourier transform, into a tensor whose zero-frequency element is
-    back at the start of the tensor instead of the center. It is the
-    inverse of the fftshift operator. It will by default shift
-    every dimension in the tensor but the last (which is assumed to
-    represent the complex number and be of dimension 2), but can shift
-    any arbitrary set of dimensions.
-
-    Parameters
-    ----------
-    array : torch.Tensor
-        An array of data to be ifftshifted
-    dims : list(int)
-        A list of all dimensions to shift
-
-    Returns
-    -------
-    torch.Tensor 
-        The ifftshifted tensor
-
-    """
-
-    if dims is None:
-        dims=list(range(array.dim()))[:-1]
-    for dim in dims:
-        length = array.size()[dim]
-        cut_to = length // 2
-        cut_len = length - cut_to
-
-        array = t.cat((array.narrow(dim,cut_to,cut_len),
-                       array.narrow(dim,0,cut_to)), dim)
-    return array
-
-
-def expi(x):
-    """Returns a complex-format tensor for exp(i* (x))
-    
-    Expects the input to be in the form of a real-valued tensor
-
-    Parameters
-    ----------
-    x : torch.Tensor
-        An array to be exponentiated
-    
-    Returns
-    -------
-    torch.Tensor
-        A complex-format tensor
-
-    """
-    return t.stack((t.cos(x),t.sin(x)),dim=-1)
-
-
-def cexpi(z):
-    """Returns a complex-format tensor for exp(i* (z))
-    
-    Expects the input to be in the form of a complex-valued tensor
-
-    Parameters
-    ----------
-    x : torch.Tensor
-        An array to be exponentiated
-    
-    Returns
-    -------
-    torch.Tensor
-        A complex-format tensor
-
-    """
-    real = t.cos(z[...,0]) * t.exp(-z[...,1])
-    imag = t.sin(z[...,0]) * t.exp(-z[...,1])
-    return t.stack((real, imag),dim=-1)
diff --git a/CDTools/tools/data/data.py b/CDTools/tools/data/data.py
index 65abcef..8aeafb7 100644
--- a/CDTools/tools/data/data.py
+++ b/CDTools/tools/data/data.py
@@ -277,7 +277,7 @@ def get_mask(cxi_file):
         mask = np.array(i1['detector_1/mask']).astype(np.uint32)
         mask_on = np.equal(mask,np.uint32(0))
         mask_has_signal = np.equal(mask,np.uint32(0x00001000))
-        return np.logical_or(mask_on,mask_has_signal).astype(np.bool)
+        return np.logical_or(mask_on,mask_has_signal).astype(bool)
     else:
         return None
 
diff --git a/CDTools/tools/image_processing/image_processing.py b/CDTools/tools/image_processing/image_processing.py
index 8681262..bd9ff13 100644
--- a/CDTools/tools/image_processing/image_processing.py
+++ b/CDTools/tools/image_processing/image_processing.py
@@ -10,7 +10,7 @@ a way that it is safe to include them in automatic differentiation models.
 from __future__ import division, print_function, absolute_import
 import numpy as np
 import torch as t
-from CDTools.tools import cmath, propagators
+from CDTools.tools import propagators
 
 __all__ = ['centroid', 'centroid_sq', 'sinc_subpixel_shift',
            'find_subpixel_shift', 'find_pixel_shift', 'find_shift',
@@ -75,7 +75,7 @@ def centroid_sq(im, dims=2, comp=False):
         An (i,j) index or stack of indices
     """
     if comp:
-        im_sq = cmath.cabssq(im)
+        im_sq = t.abs(im)**2
     else:
         im_sq = im**2
 
@@ -109,9 +109,9 @@ def sinc_subpixel_shift(im, shift):
     I = I.to(dtype=im.dtype,device=im.device)
     J = J.to(dtype=im.dtype,device=im.device)
 
-    fft_im = cmath.fftshift(t.fft(im, 2))
-    shifted_fft_im = cmath.cmult(fft_im, cmath.expi(-shift[0]*I - shift[1]*J))
-    return t.ifft(cmath.ifftshift(shifted_fft_im),2)
+    fft_im = t.fft.fftshift(t.fft.fft2(im),dim=(-2,-1))
+    shifted_fft_im = fft_im * t.exp(1j * (-shift[0]*I - shift[1]*J))
+    return t.fft.ifft2(t.fft.ifftshift(shifted_fft_im, dim=(-2,-1)))
 
 
 
@@ -155,11 +155,11 @@ def find_subpixel_shift(im1, im2, search_around=(0,0), resolution=10):
         im2 = t.stack((im2,t.zeros_like(im2)),dim=-1)
 
 
-    cor_fft = cmath.cmult(t.fft(im1,2),cmath.cconj(t.fft(im2,2)))
+    cor_fft = t.fft.fft2(im1) * t.conj(t.fft.fft2(im2))
 
     # Not sure if this is more or less stable than just the correlation
     # maximum - requires some testing
-    cor = t.ifft(cor_fft / cmath.cabs(cor_fft)[:,:,None],2)
+    cor = t.fft.ifft2(cor_fft / t.abs(cor_fft))
 
 
     # Now, I need to shift the array to pull out a contiguous window
@@ -174,13 +174,15 @@ def find_subpixel_shift(im1, im2, search_around=(0,0), resolution=10):
     cor_window = t.roll(cor, shift_zero, dims=(0,1))[:2*window_size,:2*window_size]
 
     # Now we upsample this window
-    cor_window_fft = cmath.fftshift(t.fft(cor_window,2))
+    cor_window_fft = t.fft.fftshift(t.fft.fft2(cor_window),dim=(-2,-1))
     upsampled = t.zeros(tuple(t.tensor(cor_window_fft.shape)[:-1] * resolution) + (2,),
                         dtype=cor.dtype,device=cor.device)
 
     upsampled[:2*window_size,:2*window_size] = cor_window_fft
     upsampled = t.roll(upsampled,(-window_size,-window_size),dims=(0,1))
-    upsampled = t.roll(cmath.cabssq(t.ifft(upsampled, 2)),(-window_size*resolution,-window_size*resolution), dims=(0,1))
+    upsampled = t.roll(t.abs(t.fft.ifft2(upsampled))**2,
+                       (-window_size*resolution,-window_size*resolution),
+                       dims=(0,1))
 
 
     # And we extract the shift from the window
@@ -220,11 +222,11 @@ def find_pixel_shift(im1, im2):
         im2 = t.stack((im2,t.zeros_like(im2)),dim=-1)
 
 
-    cor_fft = cmath.cmult(t.fft(im1,2),cmath.cconj(t.fft(im2,2)))
+    cor_fft = t.fft.fft2(im1) * t.conj(t.fft.fft2(im2))
 
     # Not sure if this is more or less stable than just the correlation
     # maximum - requires some testing
-    cor = cmath.cabs(t.ifft(cor_fft / cmath.cabs(cor_fft)[:,:,None],2))
+    cor = t.abs(t.fft.ifft2(cor_fft / t.abs(cor_fft)))
 
 
     sh = t.tensor(cor.shape).to(device=im1.device)
@@ -302,7 +304,7 @@ def convolve_1d(image, kernel, dim=0, fftshift_kernel=True):
         complex_things -= 1
 
     if fftshift_kernel:
-        kernel = cmath.ifftshift(kernel)
+        kernel = t.fft.ifftshift(kernel,dim=(-2,-1))
 
     # If the image wasn't originally complex, and the dimension
     # was passed with the nexative-indexing convention
@@ -314,9 +316,9 @@ def convolve_1d(image, kernel, dim=0, fftshift_kernel=True):
     trans_im = t.transpose(image, dim, -2)
         
     # Take a correlation
-    fft_im = t.fft(trans_im, 1)
-    fft_kernel = t.fft(kernel, 1)
-    trans_conv = t.ifft(cmath.cmult(fft_im,fft_kernel), 1)
+    fft_im = t.fft.fft(trans_im)
+    fft_kernel = t.fft.fft(kernel)
+    trans_conv = t.fft.ifft(fft_im * fft_kernel)
 
     conv_im = t.transpose(trans_conv, dim, -2)
 
diff --git a/CDTools/tools/initializers/initializers.py b/CDTools/tools/initializers/initializers.py
index aa7df9d..40ae992 100644
--- a/CDTools/tools/initializers/initializers.py
+++ b/CDTools/tools/initializers/initializers.py
@@ -12,7 +12,6 @@ __all__ = ['exit_wave_geometry', 'calc_object_setup', 'gaussian',
            'gaussian_probe', 'SHARP_style_probe', 'RPI_spectral_init',
            'generate_subdominant_modes'] 
 
-from CDTools.tools import cmath
 from CDTools.tools.propagators import *
 from CDTools.tools.analysis import orthogonalize_probes
 from scipy.fftpack import next_fast_len
@@ -61,14 +60,17 @@ def exit_wave_geometry(det_basis, det_shape, wavelength, distance, center=None,
         The slice corresponding to the physical detector
     """
     
-    det_shape = t.Tensor(tuple(det_shape)).to(t.int32)
-    det_basis = t.Tensor(det_basis)
+    det_shape = t.tensor(tuple(det_shape)).to(t.int32)
+    det_basis = t.tensor(det_basis)
     # First, set the center if it's not already specified
     # This definition matches the center pixel of an fftshifted array
     if center is None:
-        center = det_shape // 2
+        # this is center//2, but in pytorch 1.9.0 it throws a warning
+        # if you just do that
+
+        center = t.div(det_shape,2,rounding_mode='floor')# // 2
     else:
-        center = t.Tensor(center).to(t.int32)
+        center = t.tensor(center).to(t.int32)
         
     # Then, calculate the required detector size from the centering
     # This is a bit opaque but was worth doing accurately
@@ -81,11 +83,13 @@ def exit_wave_geometry(det_basis, det_shape, wavelength, distance, center=None,
 
     
     if opt_for_fft:
-        full_shape = t.Tensor([next_fast_len(dim) for dim in full_shape]).to(t.int32)
+        full_shape = t.tensor([next_fast_len(dim) for dim in full_shape]).to(t.int32)
     
     # Then, generate a slice that pops the actual detector from the full
     # detector shape
-    full_center = full_shape // 2
+    # this is full_center//2, but in pytorch 1.9.0 it throws a warning
+    # if you just do that
+    full_center = t.div(full_shape,2, rounding_mode='floor')# // 2
     det_slice = np.s_[int(full_center[0]-center[0]):
                       int(full_center[0]-center[0]+det_shape[0]),
                       int(full_center[1]-center[1]):
@@ -97,7 +101,7 @@ def exit_wave_geometry(det_basis, det_shape, wavelength, distance, center=None,
     # This method should work for a general parallelogram
     # shaped detector
     det_shape = det_basis * full_shape.to(t.float32)
-    pinv_basis = t.Tensor(np.linalg.pinv(det_shape).transpose()).to(t.float32)
+    pinv_basis = t.tensor(np.linalg.pinv(det_shape).transpose()).to(t.float32)
     real_space_basis = pinv_basis * wavelength * distance
 
     # This is definitely correct, but less simple. Included here
@@ -105,7 +109,7 @@ def exit_wave_geometry(det_basis, det_shape, wavelength, distance, center=None,
     #oop_dir = np.cross(det_basis[:,0],det_basis[:,1])
     #oop_dir /= np.linalg.norm(oop_dir)
     #full_basis = np.array([np.array(det_basis[:,0]),np.array(det_basis[:,1]),oop_dir]).transpose()
-    #inv_basis = t.Tensor(np.linalg.inv(full_basis)[:2,:].transpose()).to(t.float32)
+    #inv_basis = t.tensor(np.linalg.inv(full_basis)[:2,:].transpose()).to(t.float32)
     #real_space_basis = inv_basis*wavelength * distance / \
     #    full_shape.to(t.float32)
 
@@ -201,7 +205,7 @@ def gaussian(shape, sigma, amplitude=1, center = None, curvature=[0,0]):
     jsq = (j - center[1])**2
     result = np.exp((1j*curvature[0] / 2 - 1 / (2 * sigma[0]**2)) * isq + \
         (1j*curvature[1] / 2 - 1 / (2 * sigma[1]**2)) * jsq)
-    return cmath.complex_to_torch(amplitude*result)
+    return t.tensor(amplitude*result).to(t.complex64)
 
 
 
@@ -268,8 +272,8 @@ def gaussian_probe(dataset, basis, shape, sigma, propagation_distance=0):
     # Finally, we should calculate the average pattern intensity from the
     # dataset and normalize the gaussian probe. This should be done by
     avg_intensities = [t.sum(dataset[idx][1]) for idx in range(len(dataset))]
-    avg_intensity = t.mean(t.Tensor(avg_intensities))
-    probe_intensity = t.sum(cmath.cabssq(probe))
+    avg_intensity = t.mean(t.tensor(avg_intensities))
+    probe_intensity = t.sum(t.abs(probe)**2)
     return avg_intensity / probe_intensity * probe
     
 
@@ -324,14 +328,13 @@ def SHARP_style_probe(dataset, shape, det_slice, propagation_distance=None, over
     if hasattr(dataset, 'background') and dataset.background is not None:
         intensities[det_slice] = np.clip(intensities[det_slice] - dataset.background.cpu().numpy(), a_min=0,a_max=None)
     
-    probe_fft = cmath.complex_to_torch(np.sqrt(intensities))
-    
-    probe_guess = cmath.torch_to_complex(inverse_far_field(probe_fft))
-    
+    probe_fft = t.tensor(np.sqrt(intensities)).to(dtype=t.complex64)
+
+    probe_guess = inverse_far_field(probe_fft).numpy()
     # Now we remove the central pixel
     center = np.array(probe_guess.shape) // 2
     
-    # I'm always divided on whether to use this modification:
+    # I'm always unsure whether to use this modification:
     
     probe_guess[center[0], center[1]]=np.mean([
         probe_guess[center[0]-1, center[1]],
@@ -339,15 +342,15 @@ def SHARP_style_probe(dataset, shape, det_slice, propagation_distance=None, over
         probe_guess[center[0], center[1]-1],
         probe_guess[center[0], center[1]+1]])
 
-    probe_guess = cmath.complex_to_torch(probe_guess)
-    
+    probe_guess = t.tensor(probe_guess).to(dtype=t.complex64)
+
     if propagation_distance is not None:
         # First generate the propagation array
 
-        probe_shape = t.Tensor(tuple(probe_guess.shape))[:-1]
-
+        probe_shape = t.tensor(tuple(probe_guess.shape))
+        
         # Start by recalculating the probe basis from the given information
-        det_basis = t.Tensor(dataset.detector_geometry['basis'])
+        det_basis = t.tensor(dataset.detector_geometry['basis'])
         basis_dirs = det_basis / t.norm(det_basis, dim=0)
         distance = dataset.detector_geometry['distance']
         probe_basis = basis_dirs * dataset.wavelength * distance / \
@@ -362,14 +365,13 @@ def SHARP_style_probe(dataset, shape, det_slice, propagation_distance=None, over
         AS_prop = generate_angular_spectrum_propagator(probe_shape, probe_spacing, dataset.wavelength, propagation_distance)
 
         probe_guess = near_field(probe_guess,AS_prop)
-
     
     # Finally, place this probe in a full-sized array if there is oversampling
-    final_probe = t.zeros([dim for dim in shape] + [2])
+    final_probe = t.zeros(shape,dtype=t.complex64)
     left = shape[0]//2 - probe_guess.shape[0] // 2
     top = shape[1]//2 - probe_guess.shape[1] // 2 
     final_probe[left:left+probe_guess.shape[0],
-                top:top+probe_guess.shape[1],:] = probe_guess
+                top:top+probe_guess.shape[1]] = probe_guess
     
     return final_probe
 
@@ -388,20 +390,19 @@ def RPI_spectral_init(pattern, probe, obj_shape, n_modes=1, mask=None, backgroun
     pad1r = probe.shape[-2] - obj_shape[1] - pad1l
     
     def a_dagger(im):
-        im = cmath.complex_to_torch(im.reshape(obj_shape)).to(dtype=t.float32)
-        im = inverse_far_field(pad(far_field(im), (0,0,pad1l,pad1r,pad0l,pad0r)))
-        exit_wave = cmath.cmult(probe,im)
-        farfield = cmath.torch_to_complex(far_field(exit_wave))
-        return farfield.ravel()
+        im = t.tensor(im.reshape(obj_shape)).to(dtype=t.complex64)
+        im = inverse_far_field(pad(far_field(im), (pad1l,pad1r,pad0l,pad0r)))
+        exit_wave = probe * im
+        return far_field(exit_wave).numpy().ravel()
 
     def a(measured):
-        measured = cmath.complex_to_torch(measured.reshape(pattern.shape[0],pattern.shape[1])).to(dtype=t.float32)
+        measured = t.tensor(measured.reshape(pattern.shape[0],pattern.shape[1])).to(dtype=t.complex64)
         im = inverse_far_field(measured)
-        multiplied = cmath.cmult(cmath.cconj(probe), im)
+        multiplied = t.conj(probe) * im
         backplane = far_field(multiplied)
         clipped = backplane[pad0l:pad0l+obj_shape[0],
-                            pad1l:pad1l+obj_shape[1],:]
-        return cmath.torch_to_complex(inverse_far_field(clipped)).ravel()
+                            pad1l:pad1l+obj_shape[1]]
+        return (inverse_far_field(clipped)).numpy().ravel()
 
     patsize = pattern.shape[0]*pattern.shape[1]
     imsize = obj_shape[0]*obj_shape[1]
@@ -495,9 +496,9 @@ def generate_subdominant_modes(dominant_mode, n_modes, circular=True):
     center = ((shape[-3]-1)//2, (shape[-2]-1)//2)
         
     i, j = np.mgrid[:shape[-3], :shape[-2]]
-    i = t.Tensor(i - center[0]).to(dtype=dominant_fft.dtype,
+    i = t.tensor(i - center[0]).to(dtype=dominant_fft.dtype,
                                     device=dominant_fft.device)
-    j = t.Tensor(j - center[1]).to(dtype=dominant_fft.dtype,
+    j = t.tensor(j - center[1]).to(dtype=dominant_fft.dtype,
                                     device=dominant_fft.device)
 
     if circular:
diff --git a/CDTools/tools/interactions/interactions.py b/CDTools/tools/interactions/interactions.py
index 9d6a882..05ee0db 100644
--- a/CDTools/tools/interactions/interactions.py
+++ b/CDTools/tools/interactions/interactions.py
@@ -7,7 +7,6 @@ for ptychographic reconstruction.
 
 from __future__ import division, print_function, absolute_import
 
-from CDTools.tools.cmath import *
 import torch as t
 import numpy as np
 from CDTools.tools import propagators 
@@ -245,7 +244,7 @@ def ptycho_2D_round(probe, obj, translations, multiple_modes=False):
     Parameters
     ----------
     probe : torch.Tensor
-        A (P)xMxLx2 probe function to illuminate the object
+        A (P)xMxL probe function to illuminate the object
     object : torch.Tensor
         The object function to be probed
     translations : torch.Tensor
@@ -266,16 +265,16 @@ def ptycho_2D_round(probe, obj, translations, multiple_modes=False):
 
         
     integer_translations = t.round(translations).to(dtype=t.int32)
-    selections = t.stack([obj[tr[0]:tr[0]+probe.shape[-3],
-                              tr[1]:tr[1]+probe.shape[-2]]
+    selections = t.stack([obj[tr[0]:tr[0]+probe.shape[-2],
+                              tr[1]:tr[1]+probe.shape[-1]]
                           for tr in integer_translations])
 
     if multiple_modes:
         # if the probe dimension is 4, then this hasn't yet been broadcast
         # over the translation dimensions
-        output = cmult(probe,selections[:,None,:,:,:])
+        output = probe * selections[:,None,:,:]
     else:
-        output = cmult(probe,selections)
+        output = probe * selections
 
     if single_translation:
         return output[0]
@@ -345,7 +344,7 @@ def ptycho_2D_linear(probe, obj, translations, shift_probe=True):
             obj_slice = obj[tr[0]:tr[0]+probe.shape[0],
                             tr[1]:tr[1]+probe.shape[1]]
             
-            exit_waves.append(cmult(selection,obj_slice))
+            exit_waves.append(selection * obj_slice)
     else:
         for tr, sp in zip(integer_translations,
                           subpixel_translations):
@@ -371,7 +370,7 @@ def ptycho_2D_linear(probe, obj, translations, shift_probe=True):
                 sel10 * sp[0]*(1-sp[1]) + \
                 sel11 * sp[0]*sp[1]
 
-            exit_waves.append(cmult(probe,selection))
+            exit_waves.append(probe * selection)
 
     if single_translation:
         return exit_waves[0]
@@ -402,7 +401,7 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multi
     Parameters
     ----------
     probe : torch.Tensor
-        An (P)xMxLx2 probe function for the exit waves
+        An (P)xMxL probe function for the exit waves
     object : torch.Tensor
         The object function to be probed
     translations : torch.Tensor
@@ -415,7 +414,7 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multi
     Returns
     -------
     exit_waves : torch.Tensor
-        An (N)x(P)xMxLx2 tensor of the calculated exit waves
+        An (N)x(P)xMxL tensor of the calculated exit waves
     """
     single_translation = False
     if translations.dim() == 1:
@@ -428,38 +427,40 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multi
     subpixel_translations = translations - integer_translations
     integer_translations = integer_translations.to(dtype=t.int32)
 
-    selections = t.stack([obj[tr[0]:tr[0]+probe.shape[-3],
-                              tr[1]:tr[1]+probe.shape[-2]]
+    selections = t.stack([obj[tr[0]:tr[0]+probe.shape[-2],
+                              tr[1]:tr[1]+probe.shape[-1]]
                           for tr in integer_translations])
     
     exit_waves = []
     if shift_probe:
-        i = t.arange(probe.shape[-3],device=probe.device,dtype=probe.dtype) \
+        i = t.arange(probe.shape[-2],device=probe.device,dtype=t.float32) \
             - probe.shape[-3]//2
-        j = t.arange(probe.shape[-2],device=probe.device,dtype=probe.dtype) \
+        j = t.arange(probe.shape[-1],device=probe.device,dtype=t.float32) \
             - probe.shape[-2]//2
         I,J = t.meshgrid(i,j)
-        I = 2 * np.pi * I / probe.shape[-3]
-        J = 2 * np.pi * J / probe.shape[-2]
+        I = 2 * np.pi * I / probe.shape[-2]
+        J = 2 * np.pi * J / probe.shape[-1]
 
-        phase_masks = expi(-subpixel_translations[:,0,None,None]*I
-                           -subpixel_translations[:,1,None,None]*J)
-        fft_probe = fftshift(t.fft(probe, 2))
+        phase_masks = t.exp(1j*(-subpixel_translations[:,0,None,None]*I
+                                -subpixel_translations[:,1,None,None]*J))
+        
+        fft_probe = t.fft.fftshift(t.fft.fft2(probe),dim=(-1,-2))
         if multiple_modes:
             # if the probe dimension is 4, then this hasn't yet been broadcast
             # over the translation dimensions
-            shifted_fft_probe = cmult(fft_probe,phase_masks[:,None,:,:,:])
+            shifted_fft_probe = fft_probe * phase_masks[:,None,:,:]
         else:
-            shifted_fft_probe = cmult(fft_probe,phase_masks)
+            shifted_fft_probe = fft_probe * phase_masks
 
-        shifted_probe = t.ifft(ifftshift(shifted_fft_probe),2)
+        shifted_probe = t.fft.ifft2(t.fft.ifftshift(shifted_fft_probe,
+                                                    dim=(-1,-2)))
 
         if multiple_modes:
             # if the probe dimension is 4, then this hasn't yet been broadcast
             # over the translation dimensions
-            output = cmult(shifted_probe,selections[:,None,:,:,:])
+            output = shifted_probe * selections[:,None,:,:]
         else:
-            output = cmult(shifted_probe,selections)
+            output = shifted_probe * selections
         
     else:
         raise NotImplementedError('Object shift not yet implemented')
@@ -524,46 +525,37 @@ def ptycho_2D_sinc_s_matrix(probe, s_matrix, translations, shift_probe=True, pad
     B = s_matrix.shape[0]//2
     
     if shift_probe:
-        i = t.arange(probe.shape[0]) - probe.shape[0]//2
-        j = t.arange(probe.shape[1]) - probe.shape[1]//2
+        i = t.arange(probe.shape[-2]) - probe.shape[-2]//2
+        j = t.arange(probe.shape[-1]) - probe.shape[-1]//2
         I,J = t.meshgrid(i,j)
-        I = 2 * np.pi * I.to(t.float32) / probe.shape[0]
-        J = 2 * np.pi * J.to(t.float32) / probe.shape[1]
+        I = 2 * np.pi * I.to(t.float32) / probe.shape[-2]
+        J = 2 * np.pi * J.to(t.float32) / probe.shape[-1]
         I = I.to(dtype=probe.dtype,device=probe.device)
         J = J.to(dtype=probe.dtype,device=probe.device)
+        print('hi')
         
         for tr, sp in zip(integer_translations,
                           subpixel_translations):
-            fft_probe = fftshift(t.fft(probe, 2))
-            shifted_fft_probe = cmult(fft_probe, expi(-sp[0]*I - sp[1]*J))
-            shifted_probe = t.ifft(ifftshift(shifted_fft_probe),2)
+            fft_probe = t.fft.fftshift(t.fft.fft2(probe), dim=(-1,-2))
+            shifted_fft_probe = fft_probe * t.exp(1j*(-sp[0]*I - sp[1]*J))
+            shifted_probe = t.fft.ifft2(t.fft.ifftshift(shifted_fft_probe,
+                                                         dim=(-1,-2)))
             
-            s_matrix_slice = s_matrix[:,:,tr[0]:tr[0]+probe.shape[0]+2*B,
-                                      tr[1]:tr[1]+probe.shape[1]+2*B]
+            s_matrix_slice = s_matrix[:,:,tr[0]:tr[0]+probe.shape[-2]+2*B,
+                                      tr[1]:tr[1]+probe.shape[-1]+2*B]
 
 
-            output = t.zeros([probe.shape[0]+2*B,probe.shape[1]+2*B,2]).to(
+            output = t.zeros([probe.shape[-2]+2*B,probe.shape[-1]+2*B,2]).to(
                                   device=s_matrix_slice.device,
                                   dtype=s_matrix_slice.dtype)
 
             
             for i in range(s_matrix.shape[0]):
                 for j in range(s_matrix.shape[1]):
-                    output [i:i+probe.shape[0],j:j+probe.shape[1]] += \
-                        cmult(shifted_probe, s_matrix_slice[i,j,i:i+probe.shape[0],j:j+probe.shape[1],:])
-            
-            #output = t.zeros([s_matrix_slice.shape[2]+2*B,
-            #                  s_matrix_slice.shape[3]+2*B,2]).to(
-            #                      device=s_matrix_slice.device,
-            #                      dtype=s_matrix_slice.dtype)
-            
-            #for i in range(s_matrix.shape[0]):
-            #    for j in range(s_matrix.shape[1]):
-            #        output[i:i+probe.shape[0],j:j+probe.shape[1]] += \
-            #            cmult(shifted_probe, s_matrix_slice[i,j,:,:,:])
+                    output [i:i+probe.shape[-2],j:j+probe.shape[-1]] += \
+                        shifted_probe * s_matrix_slice[i,j,i:i+probe.shape[-2],j:j+probe.shape[-1]]
 
             exit_waves.append(output)
-            #exit_waves.append(cmult(shifted_probe, obj_slice))
         
     else:
         raise NotImplementedError('Object shift not yet implemented')
@@ -613,10 +605,10 @@ def RPI_interaction(probe, obj):
     pad1r = probe.shape[-2] - obj.shape[-2] - pad1l
         
     if obj.dim() == 3:
-        fftobj = t.nn.functional.pad(fftobj, (0, 0, pad1l, pad1r, pad0l, pad0r))
+        fftobj = t.nn.functional.pad(fftobj, (pad1l, pad1r, pad0l, pad0r))
     elif obj.dim() == 4:
         fftobj = t.nn.functional.pad(
-            fftobj, (0, 0, pad1l, pad1r, pad0l, pad0r, 0, 0))
+            fftobj, (pad1l, pad1r, pad0l, pad0r, 0,0))
     else:
         raise NotImplementedError('RPI interaction with obj of dimension higher than 4 (including complex dimension) is not supported.')
         
@@ -624,6 +616,6 @@ def RPI_interaction(probe, obj):
     upsampled_obj = propagators.inverse_far_field(fftobj)
 
     if obj.dim() == 4:
-        return cmult(probe[None,...], upsampled_obj)
+        return probe[None,...] *  upsampled_obj
     else:
-        return cmult(probe, upsampled_obj)
+        return probe * upsampled_obj
diff --git a/CDTools/tools/measurements/measurements.py b/CDTools/tools/measurements/measurements.py
index 9871341..1101814 100644
--- a/CDTools/tools/measurements/measurements.py
+++ b/CDTools/tools/measurements/measurements.py
@@ -7,7 +7,6 @@ thresholds, backgrounds, and more.
 
 from __future__ import division, print_function, absolute_import
 
-from CDTools.tools import cmath
 import torch as t
 import numpy as np
 from torch.nn.functional import avg_pool2d
@@ -44,18 +43,18 @@ def intensity(wavefield, detector_slice=None, epsilon=1e-7, saturation=None, ove
     sim_patterns : torch.Tensor
         A real MxN array storing the wavefield's intensities
     """
-    output = cmath.cabssq(wavefield)
-
+    output = t.abs(wavefield)**2
+    
     # Now we apply oversampling
     if oversampling != 1:
-        if wavefield.dim() == 3:
+        if wavefield.dim() == 2:
             output = avg_pool2d(output.unsqueeze(0), oversampling)[0]
         else:
             output = avg_pool2d(output, oversampling)
 
     # Then we grab the detector slice
     if detector_slice is not None:
-        if wavefield.dim() == 3:
+        if wavefield.dim() == 2:
             output = output[detector_slice]
         else:
             output = output[(np.s_[:],) + detector_slice]
@@ -126,7 +125,7 @@ def density_matrix(wavefields, density_matrix, detector_slice=None, epsilon=1e-7
                   for j in range(density_matrix.shape[-1])):
         if i == j: # diagonal
             output += density_matrix[...,i,j,None,None] \
-                * cmath.cabssq(wavefields[i])
+                * t.abs(wavefields[i])**2
         if i < j: # upper triangle, real part 
             output += 2 * density_matrix[...,i,j,None,None] \
                 * (wavefields[i,...,0] * wavefields[j,...,0]
@@ -191,7 +190,7 @@ def incoherent_sum(wavefields, detector_slice=None, epsilon=1e-7, saturation=Non
         A real LXMxN array storing the incoherently summed intensities
     """
 
-    output = t.sum(cmath.cabssq(wavefields),dim=-3) 
+    output = t.sum(t.abs(wavefields)**2,dim=-3) 
 
     # Now we apply oversampling
     if oversampling != 1:
diff --git a/CDTools/tools/plotting/plotting.py b/CDTools/tools/plotting/plotting.py
index c454f2e..08c98fa 100644
--- a/CDTools/tools/plotting/plotting.py
+++ b/CDTools/tools/plotting/plotting.py
@@ -8,7 +8,6 @@ images exist, as well as plotting scan patterns and nanomaps
 
 from __future__ import division, print_function, absolute_import
 
-from CDTools.tools import cmath
 import torch as t
 import numpy as np
 import matplotlib.pyplot as plt
@@ -129,7 +128,7 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m',
         # If final dimension is 2, assume it is a complex array. If not,
         # assume it represents a real array
         if im.shape[-1] == 2:
-            im = cmath.torch_to_complex(im.detach().cpu())
+            im = im.detach().cpu().numpy()
         else:
             im = im.detach().cpu().numpy()
 
diff --git a/CDTools/tools/projectors/projectors.py b/CDTools/tools/projectors/projectors.py
index eef5461..f853707 100644
--- a/CDTools/tools/projectors/projectors.py
+++ b/CDTools/tools/projectors/projectors.py
@@ -5,7 +5,6 @@ alongside the automatic differentiation ones, for comparison or in a
 situation where they might be needed.
 """
 from __future__ import division, print_function, absolute_import
-from CDTools.tools.cmath import *
 import torch as t
 
 __all__ = ['modulus', 'support']
@@ -41,7 +40,7 @@ def modulus(wavefront, intensities, mask = None):
     # Calculate amplitudes from intensities
     amplitudes = t.sqrt(intensities)
     # Normalize wavefront so the complex elements have modulus one
-    wavefront_mag = cabs(wavefront)
+    wavefront_mag = t.abs(wavefront)
     projected = wavefront * (amplitudes / wavefront_mag)[...,None]
     # Replace amplitude of wavefront with measured amplitude
     if mask is not None:
diff --git a/CDTools/tools/propagators/propagators.py b/CDTools/tools/propagators/propagators.py
index ed6d645..c81bcbd 100644
--- a/CDTools/tools/propagators/propagators.py
+++ b/CDTools/tools/propagators/propagators.py
@@ -5,7 +5,6 @@ ptychography model. Each function implements a different propagator.
 """
 from __future__ import division, print_function, absolute_import
 
-from CDTools.tools.cmath import *
 import torch as t
 from torch.nn.functional import grid_sample
 from scipy import fftpack
@@ -19,7 +18,6 @@ __all__ = ['far_field', 'near_field',
            'high_NA_far_field',
            'generate_generalized_angular_spectrum_propagator']
 
-
 def far_field(wavefront):
     """Implements a far-field propagator in torch
 
@@ -47,8 +45,10 @@ def far_field(wavefront):
     propagated : torch.Tensor
         The JxNxMx2 propagated wavefield
     """
-
-    return fftshift(t.fft(ifftshift(wavefront), 2, normalized=True))
+    
+    shifted = t.fft.ifftshift(wavefront, dim=(-1,-2))
+    propagated = t.fft.fft2(shifted, norm='ortho')
+    return t.fft.fftshift(propagated, dim=(-1,-2))
 
 
 def inverse_far_field(wavefront):
@@ -74,7 +74,9 @@ def inverse_far_field(wavefront):
     propagated : torch.Tensor
         The JxNxMx2 exit wavefield
     """
-    return fftshift(t.ifft(ifftshift(wavefront), 2, normalized=True))
+    shifted = t.fft.ifftshift(wavefront, dim=(-1,-2))
+    propagated = t.fft.ifft2(shifted, norm='ortho')
+    return t.fft.fftshift(propagated, dim=(-1,-2))
 
 
 def generate_high_NA_k_intensity_map(sample_basis, det_basis,det_shape,distance, wavelength, *args, lens=False, **kwargs):
@@ -355,8 +357,8 @@ def generate_angular_spectrum_propagator(shape, spacing, wavelength, z, *args, r
         A phase mask which accounts for the phase change that each plane wave will undergo.
     """
 
-    ki = 2 * np.pi * fftpack.fftfreq(shape[0],spacing[0])
-    kj = 2 * np.pi * fftpack.fftfreq(shape[1],spacing[1])
+    ki = 2 * np.pi * t.fft.fftfreq(shape[0],spacing[0]).numpy()
+    kj = 2 * np.pi * t.fftfreq(shape[1],spacing[1]).numpy()
     Kj, Ki = np.meshgrid(kj,ki)
 
     # Define this as complex so the square root properly gives
@@ -496,8 +498,8 @@ def generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, o
     inv_basis =  np.linalg.pinv(basis).transpose()
 
     # Then we calculate the frequencies in (i,j) space
-    ki = 2 * np.pi * fftpack.fftfreq(shape[0])
-    kj = 2 * np.pi * fftpack.fftfreq(shape[1])
+    ki = 2 * np.pi * t.fft.fftfreq(shape[0]).numpy()
+    kj = 2 * np.pi * t.fft.fftfreq(shape[1]).numpy()
     K_ij = np.stack(np.meshgrid(ki,kj, indexing='ij'))
 
     # Now we convert these to frequencies in reciprocal space
@@ -610,7 +612,7 @@ def near_field(wavefront, angular_spectrum_propagator):
     propagated : torch.Tensor
         The propagated wavefront 
     """
-    return t.ifft(cmult(angular_spectrum_propagator,t.fft(wavefront,2)), 2)
+    return t.fft.ifft2(angular_spectrum_propagator * t.fft.fft2(wavefront))
 
 
 
@@ -651,7 +653,8 @@ def inverse_near_field(wavefront, angular_spectrum_propagator):
     propagated : torch.Tensor
         The inverse propagated wavefront
     """
-    return t.ifft(cmult(t.fft(wavefront,2), cconj(angular_spectrum_propagator)), 2)
+    return t.fft.ifft2(t.fft.fft2(wavefront)
+                       * t.conj(angular_spectrum_propagator))
 
 
 
diff --git a/conda_requirements.txt b/conda_requirements.txt
index da4dcb2..186c07b 100644
--- a/conda_requirements.txt
+++ b/conda_requirements.txt
@@ -2,7 +2,7 @@ numpy>=1.0
 scipy>=1.0
 matplotlib>=2.0
 python-dateutil
-pytorch>=1.3.0
+pytorch>=1.8.0
 h5py>=2.1
 pytest
 sphinx
diff --git a/examples/gold_ball_ptycho.py b/examples/gold_ball_ptycho.py
index f822b94..fa4ccd7 100644
--- a/examples/gold_ball_ptycho.py
+++ b/examples/gold_ball_ptycho.py
@@ -11,23 +11,24 @@ dataset = CDTools.datasets.Ptycho2DDataset.from_cxi(filename)
 
 # Next, we create a ptychography model from the dataset
 # Note that we explicitly as for two incoherent probe modes
-model = CDTools.models.FancyPtycho.from_dataset(dataset, n_modes=2)
+model = CDTools.models.FancyPtycho.from_dataset(dataset, n_modes=2,dm_rank=0, probe_support_radius=50)
 
 # Let's do this reconstruction on the GPU, shall we? 
 model.to(device='cuda')
 dataset.get_as(device='cuda')
 
-for i, loss in enumerate(model.Adam_optimize(20, dataset, batch_size=50)):
-#for i, loss in enumerate(model.LBFGS_optimize(20, dataset, lr=1, history_size=5)):
+for loss in model.Adam_optimize(400, dataset, batch_size=50, schedule=True):
     # And we liveplot the updates to the model as they happen
-    print(i,loss)
+    print(model.report())
     model.inspect(dataset)
 
 # And we save the reconstruction out to a file
 #with open('example_reconstructions/gold_balls.pickle', 'wb') as f:
 #    pickle.dump(model.save_results(dataset),f)
 
+model.tidy_probes()
+
 # Finally, we plot the results
-#model.inspect(dataset)
+model.inspect(dataset)
 #model.compare(dataset)
-#plt.show()
+plt.show()
diff --git a/examples/simple_ptycho.py b/examples/simple_ptycho.py
index 866a7a0..768bbce 100644
--- a/examples/simple_ptycho.py
+++ b/examples/simple_ptycho.py
@@ -2,18 +2,19 @@ from __future__ import division, print_function, absolute_import
 
 import CDTools
 from matplotlib import pyplot as plt
-
+import time
 # First, we load an example dataset from a .cxi file
 filename = 'example_data/lab_ptycho_data.cxi'
 dataset = CDTools.datasets.Ptycho2DDataset.from_cxi(filename)
-plt.ion()
+
 # Next, we create a ptychography model from the dataset
 model = CDTools.models.SimplePtycho.from_dataset(dataset)
 
+t = time.time()
 # Now, we run a short reconstruction from the dataset!
-for i, loss in enumerate(model.Adam_optimize(10, dataset)):
+for i, loss in enumerate(model.Adam_optimize(40, dataset,lr=0.01,batch_size=25)):#,batch_size=10000)):#0.001)):
     print(i, loss)
-
+print(time.time() - t)
 # Finally, we plot the results
 model.inspect(dataset)
 model.compare(dataset)
diff --git a/setup.py b/setup.py
index 44ee5c6..2f4f5e8 100644
--- a/setup.py
+++ b/setup.py
@@ -17,7 +17,7 @@ setuptools.setup(
         "scipy>=1.0",
         "matplotlib>=2.0",
         "python-dateutil",
-        "torch>=1.3.0", #1.3.0 adds the align_corners option for grid_sample which is used for high NA far-field propagation. 1.2.0 introduced boolean tensors in a breaking way, we use the boolean tensors here for masking.
+        "torch>=1.9.0", #1.9.0 implements support for autograd on indexed complex tensors, key to allowing us to use complex tensors in the forward models
         "h5py>=2.1",
         "pathlib2 ; python_version<'3.4'"],
     extras_require={
diff --git a/tests/conftest.py b/tests/conftest.py
index 6b31585..bbf4ecb 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -1,4 +1,3 @@
-from __future__ import division, print_function, absolute_import
 import numpy as np
 import h5py
 import pytest
@@ -104,7 +103,7 @@ def ptycho_cxi_1():
     # Remember the format for the CXI file differs from the format used
     # internally
     mask = np.zeros((256,256)).astype(np.int32)
-    expected['mask'] = np.ones((256,256)).astype(np.bool)
+    expected['mask'] = np.ones((256,256)).astype(bool)
     d1f.create_dataset('mask',data=mask)
 
     # Create an initial background
@@ -272,7 +271,7 @@ def ptycho_cxi_3():
     # Remember the format for the CXI file differs from the format used
     # internally
     mask = np.ones((256,256)).astype(np.uint32) * 0x00001000
-    expected['mask'] = np.ones((256,256)).astype(np.bool)
+    expected['mask'] = np.ones((256,256)).astype(bool)
     d1f.create_dataset('mask',data=mask)
     expected['dark'] = None
     
diff --git a/tests/tools/test_analysis.py b/tests/tools/test_analysis.py
index 12e897f..ae06088 100644
--- a/tests/tools/test_analysis.py
+++ b/tests/tools/test_analysis.py
@@ -6,7 +6,7 @@ from scipy import fftpack as ffts
 import torch as t
 from itertools import combinations
 
-from CDTools.tools import analysis, cmath, initializers
+from CDTools.tools import analysis, initializers
 
 
 def test_orthogonalize_probes():
diff --git a/tests/tools/test_cmath.py b/tests/tools/test_cmath.py
deleted file mode 100644
index e39d821..0000000
--- a/tests/tools/test_cmath.py
+++ /dev/null
@@ -1,135 +0,0 @@
-from __future__ import division, print_function, absolute_import
-
-from CDTools.tools import cmath
-import numpy as np
-import torch as t
-from scipy.fftpack import fftshift, ifftshift
-
-def test_complex_to_torch():
-    arr = np.random.rand(100,4) + 1j * np.random.rand(100,4)
-    tensor = cmath.complex_to_torch(arr)
-    assert np.allclose(tensor[:,:,0].numpy(),np.real(arr))
-    assert np.allclose(tensor[:,:,1].numpy(),np.imag(arr))
-
-    
-def test_torch_to_complex():
-    tensor = t.rand(100,4,2)
-    arr = cmath.torch_to_complex(tensor)
-    assert np.allclose(tensor[:,:,0].numpy(),np.real(arr))
-    assert np.allclose(tensor[:,:,1].numpy(),np.imag(arr))
-
-
-def test_cabssq():
-    arr = np.random.rand(56,23,2) + 1j *  np.random.rand(56,23,2)
-    cabssq = cmath.cabssq(cmath.complex_to_torch(arr))
-    assert np.allclose(cabssq.numpy(),np.abs(arr)**2)
-
-
-def test_cabs():
-    arr = np.random.rand(2,3,4,5) + 1j * np.random.rand(2,3,4,5)
-    cabs = cmath.cabs(cmath.complex_to_torch(arr))
-    assert np.allclose(cabs.numpy(),np.abs(arr))
-
-    
-def test_cconj():
-    arr = np.random.rand(50) + 1j * np.random.rand(50)
-    cconj = cmath.cconj(cmath.complex_to_torch(arr))
-    assert np.allclose(cmath.torch_to_complex(cconj), np.conj(arr))
-
-    
-def test_cmult():
-    arr1 = np.random.rand(50) + 1j * np.random.rand(50)
-    arr2 = np.random.rand(50) + 1j * np.random.rand(50)
-    mult = cmath.cmult(cmath.complex_to_torch(arr1),
-                       cmath.complex_to_torch(arr2))
-    assert np.allclose(cmath.torch_to_complex(mult),arr1*arr2)
-
-    
-def test_cdiv():
-    arr1 = np.random.rand(50) + 1j * np.random.rand(50)
-    arr2 = np.random.rand(50) + 1j * np.random.rand(50)
-    div = cmath.cdiv(cmath.complex_to_torch(arr1),
-                     cmath.complex_to_torch(arr2))
-    assert np.allclose(cmath.torch_to_complex(div),arr1 / arr2)
-
-    
-def test_cphase():
-    arr = np.random.rand(50) + 1j * np.random.rand(50)
-    cphase = cmath.cphase(cmath.complex_to_torch(arr))
-    assert np.allclose(cphase.numpy(), np.angle(arr))
-
-
-def test_scalars():
-    arr1 = np.random.rand(1) + 1j * np.random.rand(1)
-    arr2 = np.random.rand(1) + 1j * np.random.rand(1)
-    div = cmath.cdiv(cmath.complex_to_torch(arr1),
-                     cmath.complex_to_torch(arr2))
-    assert np.allclose(cmath.torch_to_complex(div),arr1 / arr2)
-
-def test_fftshift():
-    #1D, even
-    arr = np.random.rand(300) + 1j * np.random.rand(300)
-    shifted = cmath.fftshift(cmath.complex_to_torch(arr))
-    assert np.allclose(fftshift(arr),
-                       cmath.torch_to_complex(shifted))
-    #1D, odd
-    arr = np.random.rand(301) + 1j * np.random.rand(301)
-    shifted = cmath.fftshift(cmath.complex_to_torch(arr))
-    assert np.allclose(fftshift(arr),
-                       cmath.torch_to_complex(shifted))
-
-    #2D
-    arr = np.random.rand(20,21) + 1j * np.random.rand(20,21)
-    shifted = cmath.fftshift(cmath.complex_to_torch(arr))
-    assert np.allclose(fftshift(arr),
-                       cmath.torch_to_complex(shifted))
-    #3D
-    arr = np.random.rand(15,16,17) + 1j * np.random.rand(15,16,17)
-    shifted = cmath.fftshift(cmath.complex_to_torch(arr))
-    assert np.allclose(fftshift(arr),
-                       cmath.torch_to_complex(shifted))
-
-    #3D, choosing specific axes
-    arr = np.random.rand(15,16,17) + 1j * np.random.rand(15,16,17)
-    shifted = cmath.fftshift(cmath.complex_to_torch(arr),dims=(0,1))
-    assert np.allclose(fftshift(arr,axes=(0,1)),
-                       cmath.torch_to_complex(shifted))
-
-
-def test_ifftshift():
-    #1D, even
-    arr = np.random.rand(300) + 1j * np.random.rand(300)
-    shifted = cmath.ifftshift(cmath.complex_to_torch(arr))
-    assert np.allclose(ifftshift(arr),
-                       cmath.torch_to_complex(shifted))
-    #1D, odd
-    arr = np.random.rand(301) + 1j * np.random.rand(301)
-    shifted = cmath.ifftshift(cmath.complex_to_torch(arr))
-    assert np.allclose(ifftshift(arr),
-                       cmath.torch_to_complex(shifted))
-
-    #2D
-    arr = np.random.rand(20,21) + 1j * np.random.rand(20,21)
-    shifted = cmath.ifftshift(cmath.complex_to_torch(arr))
-    assert np.allclose(ifftshift(arr),
-                       cmath.torch_to_complex(shifted))
-    #3D
-    arr = np.random.rand(15,16,17) + 1j * np.random.rand(15,16,17)
-    shifted = cmath.ifftshift(cmath.complex_to_torch(arr))
-    assert np.allclose(ifftshift(arr),
-                       cmath.torch_to_complex(shifted))
-
-    #3D, choosing specific axes
-    arr = np.random.rand(15,16,17) + 1j * np.random.rand(15,16,17)
-    shifted = cmath.ifftshift(cmath.complex_to_torch(arr),dims=(0,1))
-    assert np.allclose(ifftshift(arr,axes=(0,1)),
-                       cmath.torch_to_complex(shifted))
-
-
-def test_expi():
-    phases = np.random.rand(20,70) * 2 * np.pi
-    np_result = np.exp(1j * phases)
-
-    torch_result = cmath.expi(t.Tensor(phases))
-
-    assert np.allclose(np_result, cmath.torch_to_complex(torch_result))
diff --git a/tests/tools/test_data.py b/tests/tools/test_data.py
index 5e0b7d7..a39b15d 100644
--- a/tests/tools/test_data.py
+++ b/tests/tools/test_data.py
@@ -1,5 +1,3 @@
-from __future__ import division, print_function, absolute_import
-
 from CDTools.tools import data
 import numpy as np
 import torch as t
@@ -8,10 +6,7 @@ import pytest
 import os
 import datetime
 import numbers
-try:
-    import pathlib
-except ImportError:
-    import pathlib2 as pathlib
+import pathlib
 
 
 
diff --git a/tests/tools/test_image_processing.py b/tests/tools/test_image_processing.py
index f37581c..8f04f6f 100644
--- a/tests/tools/test_image_processing.py
+++ b/tests/tools/test_image_processing.py
@@ -5,7 +5,7 @@ import pytest
 import numpy as np
 import torch as t
 
-from CDTools.tools import image_processing, cmath, initializers, interactions
+from CDTools.tools import image_processing, initializers, interactions
 from scipy import ndimage
 from scipy.signal import fftconvolve
 
diff --git a/tests/tools/test_initializers.py b/tests/tools/test_initializers.py
index 4f6871c..92b69a4 100644
--- a/tests/tools/test_initializers.py
+++ b/tests/tools/test_initializers.py
@@ -1,7 +1,6 @@
 from __future__ import division, print_function, absolute_import
 
 from CDTools.tools import initializers
-from CDTools.tools import cmath
 from CDTools.datasets import Ptycho2DDataset
 import numpy as np
 import torch as t
diff --git a/tests/tools/test_interactions.py b/tests/tools/test_interactions.py
index 79e2bfe..fe0a95a 100644
--- a/tests/tools/test_interactions.py
+++ b/tests/tools/test_interactions.py
@@ -1,6 +1,5 @@
 from __future__ import division, print_function, absolute_import
 
-from CDTools.tools import cmath
 from CDTools.tools import interactions
 import numpy as np
 import torch as t
diff --git a/tests/tools/test_losses.py b/tests/tools/test_losses.py
index 5c372cf..6c7ec17 100644
--- a/tests/tools/test_losses.py
+++ b/tests/tools/test_losses.py
@@ -1,7 +1,6 @@
 from __future__ import division, print_function, absolute_import
 
 from CDTools.tools import losses
-from CDTools.tools import cmath
 import numpy as np
 import torch as t
 
diff --git a/tests/tools/test_measurements.py b/tests/tools/test_measurements.py
index 36375c4..0dd9019 100644
--- a/tests/tools/test_measurements.py
+++ b/tests/tools/test_measurements.py
@@ -1,7 +1,6 @@
 from __future__ import division, print_function, absolute_import
 
 from CDTools.tools import measurements
-from CDTools.tools import cmath
 import torch as t
 import numpy as np
 import pytest
diff --git a/tests/tools/test_plotting.py b/tests/tools/test_plotting.py
index a6dbe47..04cc811 100644
--- a/tests/tools/test_plotting.py
+++ b/tests/tools/test_plotting.py
@@ -1,6 +1,5 @@
 from __future__ import division, print_function, absolute_import
 
-from CDTools.tools import cmath
 from CDTools.tools import plotting
 from CDTools.tools import initializers
 import numpy as np
diff --git a/tests/tools/test_projectors.py b/tests/tools/test_projectors.py
index 060b70a..c814f61 100644
--- a/tests/tools/test_projectors.py
+++ b/tests/tools/test_projectors.py
@@ -1,6 +1,5 @@
 from __future__ import division, print_function, absolute_import
 
-from CDTools.tools import cmath
 from CDTools.tools import projectors
 import numpy as np
 import torch as t
diff --git a/tests/tools/test_propagators.py b/tests/tools/test_propagators.py
index e4ff41c..fdf13b4 100644
--- a/tests/tools/test_propagators.py
+++ b/tests/tools/test_propagators.py
@@ -1,6 +1,5 @@
 from __future__ import division, print_function, absolute_import
 
-from CDTools.tools import cmath
 from CDTools.tools import initializers
 from CDTools.tools import propagators