Commit 5a79d4e4 authored by jansen31's avatar jansen31
Browse files

interpolation and zoom for incoherent iterate stacks

fixed random guess creation?
parent 513c9cd5
......@@ -66,7 +66,7 @@ class DegenerateOrbital(PlanarMolecule):
"""
self.u0 = np.array([self.data * np.exp(1j * 2 * np.pi * np.random.random_sample(self.data.shape))
for i in range(self.degeneracy)])
self.u0 = np.fft.fftn(self.u0, axes=(-1, -2))
self.u0 = shifted_fft(self.u0, axes=(-1, -2))
def setupProxOperators(self):
"""
......@@ -108,7 +108,7 @@ class DegenerateOrbital(PlanarMolecule):
ax[1][-1].set_title("Local density of states")
plt.show()
def show(self):
def show(self, **kwargs):
"""
Create basic result plots of the phase retrieval procedure
"""
......@@ -117,9 +117,14 @@ class DegenerateOrbital(PlanarMolecule):
self.output['u1'] = self.algorithm.prox1.eval(self.algorithm.u)
self.output['u2'] = self.algorithm.prox2.eval(self.algorithm.u)
figsize = kwargs.pop("figsize", (12,3))
for i, operator in enumerate(self.algorithm.proxOperators):
operator_name = self.proxOperators[i].__name__
f = self.plot_guess(operator.eval(self.algorithm.u), name="%s satisfied" % operator_name, show=False)
f = self.plot_guess(operator.eval(self.algorithm.u),
name="%s satisfied" % operator_name,
show=False,
interpolate_and_zoom=self.interpolate_and_zoom,
figsize=figsize)
self.output['plots'].append(f)
# f1 = self.plot_guess(self.output['u1'], name='Best approximation: physical constraint satisfied', show=False)
......@@ -136,18 +141,20 @@ class DegenerateOrbital(PlanarMolecule):
# def saveOutput(self, **kwargs):
# super(PlanarMolecule, self).saveOutput(**kwargs)
def plot_guess(self, u, name=None, show=True):
def plot_guess(self, u, name=None, show=True, interpolate_and_zoom=False, figsize=(12,3)):
""""Given a list of fields, plot the individual fields and the combined intensity"""
fig, ax = plt.subplots(1, len(u) + 2, figsize=(12, 3), num=name)
prop = self.propagator(self)
u_hat = prop.eval(u)
fourier_intensity = np.sqrt(np.sum(abs(u_hat) ** 2, axis=0))
if interpolate_and_zoom:
u_show = self.interp_zoom_field(u)
fig, ax = plt.subplots(1, len(u) + 2, figsize=figsize, num=name)
for ii in range(self.degeneracy):
im = ax[ii].imshow(complex_to_rgb(u[ii]))
im = ax[ii].imshow(complex_to_rgb(u_show[ii]))
ax[ii].set_title("Degenerate orbit %d" % ii)
im = ax[-2].imshow(np.sum(abs(u) ** 2, axis=0))
im = ax[-2].imshow(np.sum(abs(u_show) ** 2, axis=0))
ax[-2].set_title("Local density of states")
plt.colorbar(im, ax=ax[-2])
prop = self.propagator(self)
u_hat = prop.eval(u)
fourier_intensity = np.sqrt(np.sum(abs(u_hat)**2, axis=0))
im = ax[-1].imshow(fourier_intensity)
ax[-1].set_title("Fourier domain intensity")
plt.colorbar(im, ax=ax[-1])
......@@ -156,9 +163,15 @@ class DegenerateOrbital(PlanarMolecule):
plt.show()
return fig
def interpolate_and_zoom(self, u, interpolation=2, zoom=0.5):
def interp_zoom_field(self, u, interpolation=2, zoom=0.5):
"""
interpolate a field and zoom in to the center
"""
return u
nt, ny, nx = u.shape
zmy = int(ny * zoom) // 2
zmx = int(nx * zoom) // 2
zoomed = u[:, zmy:ny - zmy, zmx:nx - zmx]
interpolated = np.array([fourier_interpolate(u_i, factor=interpolation) for u_i in zoomed])
return interpolated
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment