Enquiry about fixing the probe during reconstruction

[tingyouli]

Dear Contributors,
Happy New Year!

I tried to fix a ground truth probe during reconstruction, this probe is what I used for simulating the diffraction patterns. To do this, I tried set
probe=np.load("data/probe.npy")
p.scans.scan_00.illumination.model = probe
p.engines.engine.probe_update_start = 10000

here, the probe_update_start is large than the total steps, so it won't update the probe until end.
Though the model seems to work well with the output, I carefully checked find that the probe that comes last after reconstruction is not the same as what I fixed in the illumination.model. I want the final reconstructed probe is the same as the probe loaded here.
e.g. list(P.probe.S.values())[0] is equal to np.load("data/probe.npy").

I don't know what is wrong with this, can you help me identify what is wrong with this code?
Thank you so much!!!

[kahntm]

Hi,

have you checked that the initial probe from the '..._None_0000.ptyr' dump file that goes into the iterations (where it is not touched anymore) matches the probe you put in.
It could happen that the probe you loaded got modified by ptypy. The "logic" of ptypy is to load or describe a probe and then being able to modify it further before using it in any engine. Some of these things are not optional and have default values. For example an aperture will always be applied... make sure to set it (way) larger than your loaded probe, so none of its pixels get set to zero. Then the probe can still be propagated or "diversity" (noise) applied to it.

Do you have images of the probe you load and the one that goes into the iterations?
And/or the full code of your reconstruction script?

[tingyouli]

Here is my code, yes, I found that the probe reconstructed from the code is added with a support, my original probe was in 64 64, but now what I have is a 6464 probe but with only central part having amplitude. besides, the max amplitude is 120 while what I gave to it is actually around 1200

yarr = f['points'][1]
xarr = f['points'][0]
if parameters["mode"]=="simulated":
    diff = f['diffamp'][:]**2
else:
    diff = np.fft.fftshift(f['diffamp'][:]**2, axes=[-2,-1])
z_m=f['z_m'][()]
ccd_pixel_um=f['ccd_pixel_um'][()]
path_to_data_xy=parameters["path_to_data_make"]


probe_model=probe.cpu().numpy()       # for known probe conditions

with h5py.File(path_to_data_xy,'w') as f1:
    posy_um = f1.create_dataset("posy_um", data=yarr)
    posx_um = f1.create_dataset("posx_um", data=xarr)
    diffinten = f1.create_dataset("diffinten", data=diff)
    z_m_2=f1.create_dataset("z_m", data=z_m)
    ccd_pixel_um_2=f1.create_dataset("ccd_pixel_um", data=ccd_pixel_um)
    


ptypy.load_ptyscan_module("hdf5_loader")

ptypy.load_gpu_engines(arch="cupy")

p = u.Param()

# Set verbose level to interactive
p.verbose_level = "interactive"

# Set io settings (no files saved)
p.io = u.Param()
p.io.autosave = u.Param(active=True)
p.io.interaction = u.Param(active=True)
# Path to final .ptyr output file 
# using variables p.run, engine name and total nr. of iterations
p.io.rfile =  "recons/%(run)s_%(engine)s_%(iterations)04d.ptyr"

# Use non-threaded live plotting
p.io.autoplot = u.Param()
p.io.autoplot.active=True
p.io.autoplot.threaded = False
p.io.autoplot.layout = "jupyter"
p.io.autoplot.interval = 1
p.io.autoplot.make_movie = True
# Save intermediate .ptyr files (dumps) every 10 iterations
p.io.autosave = u.Param()
p.io.autosave.active = False

# Live-plotting during the reconstruction
p.io.autoplot = u.Param()
p.io.autoplot.active=True
p.io.autoplot.threaded = False
p.io.autoplot.layout = "jupyter"
p.io.autoplot.interval = 1

# Define the scan model
p.scans = u.Param()
p.scans.scan_00 = u.Param()
p.scans.scan_00.name = 'Full'

# Initial illumination (based on simulated optics)
p.scans.scan_00.illumination = u.Param()
if parameters["probe_known"]==True:
    if parameters["mode"]=="simulated":
        probe=np.load("data/recon_203104_t1_mode_probe.npy")
        p.scans.scan_00.illumination.model = probe[1][24:-24,24:-24]
        print("no filling of code!")
    else:
        probe=np.load("data/recon_203104_t1_mode_probe.npy")
        p.scans.scan_00.illumination.model = probe[0]
else:
    p.scans.scan_00.illumination.aperture = u.Param()
    p.scans.scan_00.illumination.aperture.form = parameters["probe_shape"]

# Data loader
p.scans.scan_00.data = u.Param()
p.scans.scan_00.data.name = 'Hdf5Loader'

# Read diffraction data
p.scans.scan_00.data.intensities = u.Param()
p.scans.scan_00.data.intensities.file = path_to_data_xy
p.scans.scan_00.data.intensities.key = "diffinten"
# p.scans.scan_00.data.intensities.file = path_to_data
# p.scans.scan_00.data.intensities.key = "diffamp"

# Read positions data
p.scans.scan_00.data.positions = u.Param()
p.scans.scan_00.data.positions.file = path_to_data_xy
p.scans.scan_00.data.positions.slow_key = "posx_um"        # changed
p.scans.scan_00.data.positions.slow_multiplier = 1e-6
p.scans.scan_00.data.positions.fast_key = "posy_um"
p.scans.scan_00.data.positions.fast_multiplier = 1e-6

lambda_nm=f['lambda_nm'][()]*1e-9
energy = C.h*C.c/lambda_nm/C.eV/1e3
p.scans.scan_00.data.energy = energy

# Read meta data: detector distance ???
p.scans.scan_00.data.recorded_distance = u.Param()
p.scans.scan_00.data.recorded_distance.file = path_to_data_xy
p.scans.scan_00.data.recorded_distance.key = "z_m"
p.scans.scan_00.data.recorded_distance.multiplier = 1

# Read meta data: detector pixelsize
p.scans.scan_00.data.recorded_psize = u.Param()
p.scans.scan_00.data.recorded_psize.file = path_to_data_xy
p.scans.scan_00.data.recorded_psize.key = "ccd_pixel_um"
p.scans.scan_00.data.recorded_psize.multiplier = 1e-6

# Define reconstruction engine (using DM)
p.engines = u.Param()
p.engines.engine = u.Param()




if method =="ePIE":
    p.engines.engine.name = "EPIE_cupy"
    p.engines.engine.numiter = parameters["total_steps"]
    p.engines.engine.numiter_contiguous = 1
    p.engines.engine.alpha =parameters["a"]  
    p.engines.engine.beta = parameters["b"]
    p.engines.engine.numiter_contiguous = 1
    p.engines.engine.object_norm_is_global = True
    p.engines.engine.probe_center_tol=parameters["probe_center_tol"]
    p.engines.engine.probe_update_start = parameters["probe_update_start"]
elif method == "RAAR":
    p.engines.engine.fourier_relax_factor=0
    p.engines.engine.subpix_start= parameters["total_steps"]+10
    p.engines.engine.name = "RAAR_cupy"
    p.engines.engine.numiter = parameters["total_steps"]  
    p.engines.engine.numiter_contiguous = 10
    p.engines.engine.beta = parameters["RAAR_beta"]
    p.engines.engine.probe_support = None
    p.engines.engine.probe_center_tol=parameters["probe_center_tol"]
    p.engines.engine.probe_update_start =parameters["probe_update_start"] #parameters["total_steps"]+10
    
    
elif method == "DM":
    p.engines.engine.fourier_relax_factor=parameters["fourier_relax_factor"]
    p.engines.engine.subpix_start= parameters["total_steps"]+10
    p.engines.engine.name = "DM_cupy"
    p.engines.engine.numiter = parameters["total_steps"]  
    p.engines.engine.numiter_contiguous = 1
    p.engines.engine.alpha =parameters["DM_alpha"]
    p.engines.engine.probe_support = None
    p.engines.engine.probe_center_tol=parameters["probe_center_tol"]
    p.engines.engine.probe_update_start =parameters["probe_update_start"]

start = time.time()
P = ptypy.core.Ptycho(p,level=5)
end = time.time()
print("running time is ",end - start,"s")
print(type(P.probe.S.values()))
S = list(P.obj.S.values())[0]
probe=list(P.probe.S.values())[0]
probe=probe.data[0].copy()
obj = S.data[0].copy()
np.save(parameters["save_path"]+parameters["tag"]+"_obj.npy",obj)
np.save(parameters["save_path"]+parameters["tag"]+"_probe.npy",probe)

[tingyouli]

I am checking the aperture code you mentioned, thank you for the guidance!!!

[kahntm]

Try adding the following lines when using a loaded probe that you want to keep unchanged:

p.scans.scan_00.illumination.aperture = u.Param()
p.scans.scan_00.illumination.aperture.form = 'rect'      # this aperture is not optional
p.scans.scan_00.illumination.aperture.size = 1              # in [meter] 
p.scans.scan_00.illumination.diversity = None

[jcesardasilva]

Hi everyone,

If I understood the problem well, one should notice from the PtyPy Tree Structure in the documentation that the array you pass in .illumination.model is considered an incoming wave that will be propagated with the aperture applied. Here is the link to the doc: https://ptycho.github.io/ptypy/rst/parameters.html#scan.Full.illumination.model

In case you have the .ptyr file from a previous reconstruction, you could pass it directly to .illumination = <filename>.ptyd, and then it should not be modified at the beginning.

In addition, concerning the issue you had about the amplitudes only being modified on the central region, this happens because the defaults for the parameter for engine.DM.probe_support or engine.EPIE.probe_support is 0.7, which defines a circular area centered on the probe frame, sized by 70% of the probe area, in which the probe is allowed to be nonzero. If you want to update the entire probe array, you must put None as value there, which will overwrite the default value of 0.7.

I hope this will help you.

[tingyouli]

@kahntm Thank you! I tried your advised code. Now the ground truth (up) and recovered probe (down) are looking like this, they are already very similar, except seems still remaining a little support at edges. And the max amplitude is still 120 while ground truth is around 1200. I think this should be the reason that I am providing a wave but not the probe directly. I will try to modify some parameters to make it better.

@jcesardasilva thank you for the guidance! Now I know that it is not directly providing the probe to the algorithm, instead it is a wave. However, in this case, if I am not running a .ptyr file before, instead, I am providing a numpy probe to it during the updates.
I wonder how can I do that. If I want the probe inside the reconstruction to be the same as the numpy file.

Because for most methods especially deep methods in my research on ptychography, they assume the probe to be known. So to be fair to the iterative methods, the probe should also be known and fixed. That is why I want to provide a probe numpy file to it directly and hope it could be fixed.

[kahntm]

Looking sharply at the features in your probe, it looks the one you end up with somehow is a pixel shifted compared to the "original one". And that in both directions.

(check the minima)

And if for some reason the first line and column are ignored ptypy will add zeros to create a probe of the correct size.
But here it looks like it was "rolled over" / shifted.

[kahntm]

I tried to recreate your problem... When creating the a numpy array as initial probe and setting all "processing" steps to None, I get the right probe phase in the dump file, just that the amplitudes have been scaled by a factor:

When changing the relative intensities in the probe I model, the overall integrated intensity in the probe that I get in the dump file stays constant. This sounds like ptypy is trying to set the initial probe intensity ... which is generally a good approach with modeled probes... For example to make the probe just as bright as the first / average diffraction pattern.

Looking at ptypy/ptypy/core/illumination.py, I can not quite see how / why that happens:

If an array is passed p.model is set to that array... aperture, propagation and diversity are set to None, the storage is initalized.

If p.model is now that array, model is set to point at p.model

If model has the wrong size, it is made to have the right size

finally aperture, propagation and diversity are applied

And those should all be None...

[jcesardasilva]

You might want to refer to issue #417, where this question has already been addressed.

[kahntm]

I was digging a bit... the given numpy array has the right overall intensity all the way until

... during which the intensity changes.
Even adding a line "p.photons = None" to the elif statement for the case where a numpy array and not a ptyr file is given does not change that. The probe is still scaled in amplitude.

correction... when I put it in the 2nd elif for np.arrays it does work

[daurer]

Thanks guys for looking into this. Like posted above I looked into a similar issue when loading from .ptyr files.

@kahntm maybe you could make a PR with the change suggested above which also works for ndarray probes?

[kahntm]

Hi,

yes... later today I will make a clean PR. right now I was just fiddling around in my messy version of ptypy.

[tingyouli]

Thank all of you so much for the discussions and suggestions! with @kahntm code, I think the only thing I need to see is the shift and scale problems, but it should be my own problems, I will solve myself.

And I look forward to the new version of ptypy you will update, it will create great convenience for those people who try to comapre their algorithms with known probe to those iterative methods.

[tingyouli]

1. I think I found the reason why the amplitude is still scaled, the line
   p.scans.scan_00.illumination.photons= None
   is not working. Somehow at least in my code that attached above, if I add this line after p.scans.scan_00.illumination.model = probe[1][24:-24,24:-24] . I don't know whether you had the same issue @kahntm. I just add one comman locally p.photons = None in line 476 in illumination file. Then it won't change the amplitude anymore.

2. And for the shift problem in the probe I get, I found that it is related to the the line p.engines.engine.probe_center_tol, if I set this value to be 0 (I want to center the probe at the middle), it will create errors. Right now I removed this line, it is working perfectly. The probe is fixed and the amplitude is also fixed.

[tingyouli]

To summarize, if I want to fix the probe, the below code is needed

p.scans.scan_00.illumination.model = np.load("probe.npy")
p.scans.scan_00.illumination.aperture = None
p.scans.scan_00.illumination.diversity = None
p.scans.scan_00.illumination.propagation= None

if I want to fix the probe amplitude as well, I need to change the inside code in illumination by setting p.photons = None in line 476 in illumination.

but I think the scale is not a problem, I can scale back after reconstruction by using distribution median.

[kahntm]

Yes, if you do not care about the scaling of the probe amplitude by a factor you can go ahead with those lines.

[kahntm]

Pushed the changes: commit

Created a pull request

Be aware... this fix makes it now impossible to scale the loaded probe (if it is coming as a numpy array) to a specific photon count.

[tingyouli]

Thank you so much for the updated version!!! I have no more questions on the issue.