diff --git a/CDTools/models/multislice_2d_ptycho.py b/CDTools/models/multislice_2d_ptycho.py index 0032044..63019ab 100644 --- a/CDTools/models/multislice_2d_ptycho.py +++ b/CDTools/models/multislice_2d_ptycho.py @@ -111,7 +111,6 @@ class Multislice2DPtycho(CDIModel): self.bandlimit = bandlimit - # Big question: Should there be a minus sign before self.dz, or not? self.as_prop = tools.propagators.generate_angular_spectrum_propagator(shape, spacing, self.wavelength, self.dz, bandlimit=self.bandlimit) @@ -273,12 +272,13 @@ class Multislice2DPtycho(CDIModel): else: strip_first_index = False - all_exit_waves = [] + + # For a Fourier-space probe + prs = tools.propagators.inverse_far_field(self.probe*self.probe_fourier_support[None,:,:]) + # Here is where the mixing would happen, if it happened + for i in range(self.probe.shape[0]): - # For a Fourier-space probe - pr = tools.propagators.inverse_far_field(self.probe[i] * self.probe_fourier_support) - # For a real-space probe - #pr = self.probe[i] * self.probe_support + pr = prs[i] #exit_waves = pr #print(self.probe_norm) @@ -341,11 +341,8 @@ class Multislice2DPtycho(CDIModel): def loss(self, sim_data, real_data, mask=None): - regularizer = t.mean(t.abs(self.obj)) - loss = tools.losses.amplitude_mse(real_data, sim_data, mask=mask) - lambd = 1 - return loss + lambd * regularizer - #return tools.losses.poisson_nll(real_data, sim_data, mask=mask) + return tools.losses.amplitude_mse(real_data, sim_data, mask=mask) + #return tools.losses.poisson_nll(real_data, sim_data, mask=mask,eps=0.5) def to(self, *args, **kwargs): diff --git a/CDTools/tools/losses/losses.py b/CDTools/tools/losses/losses.py index 2ec2816..ad05faa 100644 --- a/CDTools/tools/losses/losses.py +++ b/CDTools/tools/losses/losses.py @@ -139,7 +139,7 @@ def poisson_nll(intensities, sim_intensities, mask=None, eps=1e-6): """ if mask is None: return t.sum(sim_intensities+eps - - intensities * t.log(sim_intensities+eps)) \ + (intensities+eps) * t.log(sim_intensities+eps)) \ / intensities.view(-1).shape[0] else: