Full-time postdoctoral researcher | Part-time design enthusiast | Aspiring cat herder.
4DSTEM Multislice SimulationProjection Set AlgorithmsPtychography Helper FunctionsComplex Electron Probe
Extended Ptychographic Iterative Engine
viewof viz_inputs = Inputs.form({
visualization_checkboxes: Inputs.checkbox(
["probe", "potential", "model", "diffraction", "gradient"],
{
value: ["probe", "potential", "model"],
label: "panels"
}
),
options: Inputs.checkbox(
["reconstruct", "random order", "object positivity"],
{
value: ["random order", "object positivity"],
label: "ePIE"
}
),
width: Inputs.range([100, 500], {
value: 157.5,
step: 0.5,
label: "width"
}),
reset_potential_and_probe_position: Inputs.button(
"reset potential and probe position"
)
})
diffraction_amplitudes_shifted = {
let typed_array = new Float32Array(
await FileAttachment(
"FCC-slab-dps-25x25x96x96-float32-shifted.npy"
).arrayBuffer()
);
let array = Array.from(typed_array);
let dps = nj.zeros([25, 25, 96, 96]);
dps.selection.data = array;
return nj.sqrt(dps);
}
diffraction_amplitudes = {
let typed_array = new Float32Array(
await FileAttachment("FCC-slab-dps-25x25x96x96-float32.npy").arrayBuffer()
);
let array = Array.from(typed_array);
let dps = nj.zeros([25, 25, 96, 96]);
dps.selection.data = array;
return nj.sqrt(dps);
}
dp_dictionary = {
if (viz_inputs.visualization_checkboxes.includes("diffraction")) {
let dp_intensity = diffraction_amplitudes_shifted
.pick(probe_xy[1] / 4, probe_xy[0] / 4, null, null)
.pow(2);
return {
label: "Diffraction Intensity",
scheme: "Magma",
domain: [0, 0.00295],
type: "pow",
exponent: 0.375,
width: dp_intensity.shape[1],
height: dp_intensity.shape[0],
values: dp_intensity.flatten().tolist()
};
}
}
gpts = diffraction_amplitudes.shape.slice(2)
real_space_probe = new ComplexProbe(
gpts,
[0.255, 0.255],
80 * 1e3,
25,
150
).build()
probe_dictionary = {
let probe_intensity = fourier_shift(real_space_probe._array, [
-probe_xy[1],
-probe_xy[0]
]).abs_sqr();
return {
label: "Probe Intensity",
scheme: "Greys",
domain: [0, probe_normalization],
width: probe_intensity.shape[1],
height: probe_intensity.shape[0],
values: probe_intensity.flatten().tolist()
};
}
probe_normalization = real_space_probe._array.abs_sqr().max()
probe_xy
mouseover_pointer = {
let arr = [];
for (let i = 0; i < gpts[0]; i += 4) {
for (let j = 0; j < gpts[1]; j += 4) {
arr = [...arr, { x: i, y: j }];
}
}
return Plot.rect(
arr,
Plot.pointer({
x1: "x",
x2: (d) => d.x + 4,
y1: "y",
y2: (d) => d.y + 4,
stroke: "none",
fill: "none"
})
);
}
function overlap_projection(potential, incoming_probe) {
let exit_wave = incoming_probe.clone();
let complex_potential = new ComplexNDArray(potential.shape);
let cos_phase = nj.cos(potential);
let sin_phase = nj.sin(potential);
complex_potential.data = nj.stack([cos_phase, sin_phase], -1);
exit_wave = exit_wave.multiply(complex_potential);
return [exit_wave, complex_potential];
}
function fourier_projection(amplitudes, exit_wave) {
let exit_wave_fourier = new ComplexNDArray(exit_wave.shape);
exit_wave_fourier.data = nj.fft(exit_wave.data);
mutable exit_wave = exit_wave_fourier;
let exit_wave_fourier_angle = exit_wave_fourier.angle();
let modified_exit_wave_fourier = new ComplexNDArray(exit_wave.shape);
let cos_phase = nj.cos(exit_wave_fourier_angle);
let sin_phase = nj.sin(exit_wave_fourier_angle);
modified_exit_wave_fourier.data = nj.stack([cos_phase, sin_phase], -1);
modified_exit_wave_fourier =
modified_exit_wave_fourier.real_multiply(amplitudes);
let gradient_fourier = modified_exit_wave_fourier.subtract(exit_wave_fourier);
let gradient = new ComplexNDArray(exit_wave.shape);
gradient.data = nj.ifft(gradient_fourier.data);
return gradient;
}
mutable potential = nj.zeros(gpts)
mutable potential_static = nj.zeros(gpts)
potential_dictionary = {
let zero_point = viz_inputs.options.includes("object positivity")
? 0.0
: potential.min() * 0.75 - 1e-4;
let cmap = viz_inputs.options.includes("object positivity")
? "Magma"
: "PiYG";
return {
label: "Potential",
scheme: cmap,
domain: [zero_point, 0.75 * potential.max() + 1e-4],
width: potential.shape[1],
height: potential.shape[0],
values: potential.flatten().tolist()
};
}
function forward_operator(potential, probe, amplitudes) {
let [exit_wave, complex_potential] = overlap_projection(potential, probe);
let gradient = fourier_projection(amplitudes, exit_wave);
return [exit_wave, complex_potential, gradient];
}
function gradient_descent(potential, [probe_x, probe_y], step_size) {
let j = probe_x / 4;
let i = ((gpts[1] - probe_y) % gpts[1]) / 4;
let amp = diffraction_amplitudes.pick(i, j, null, null);
let probe = fourier_shift(real_space_probe._array, [-probe_y, -probe_x]);
let [exit_wave, complex_potential, gradient] = forward_operator(
potential,
probe,
amp
);
let numerator = gradient
.multiply(complex_potential.conjugate())
.multiply(probe.conjugate())
.scalar_multiply(0.0, -1 / probe_normalization);
let scaled_gradient = potential.add(numerator.re().multiply(step_size));
return scaled_gradient;
}
mutable exit_wave = real_space_probe._array
mutable gradient_step = gradient_descent(potential_static, probe_xy, 1.0)
gradient_dictionary = {
if (viz_inputs.visualization_checkboxes.includes("gradient")) {
let gradient_update = gradient_step.subtract(potential_static);
return {
label: "Gradient Step",
scheme: "PuOr",
domain: [gradient_update.min(), gradient_update.max()],
width: gradient_update.shape[1],
height: gradient_update.shape[0],
values: gradient_update.flatten().tolist()
};
}
}
function positivity_constraint(potential) {
return nj.clip(potential, 0.0);
}
mutable reconstruct_xy = {
if (viz_inputs.options.includes("random order")) {
let random_int = d3.randomInt(25);
let probe_x = random_int() * 4;
let probe_y = random_int() * 4;
return [probe_x, probe_y];
} else {
return [0, 0];
}
}
mutable mouse_pos = null
mutable probe_xy = [gpts[0] / 2, gpts[1] / 2]
probe_moving = {
let probe_moving = false;
if (mouse_pos !== null) {
if (mouse_pos[0] != probe_xy[0] || mouse_pos[1] != gpts[1] - probe_xy[1]) {
probe_moving = true;
mutable probe_xy = [mouse_pos[0], gpts[1] - mouse_pos[1]];
}
}
return probe_moving;
}
unshuffled_order = d3.range(25 * 25)
mutable shuffled_order = unshuffled_order.slice()
mutable iteration_counter = 620
reconstruct_mutable = {
if (viz_inputs.options.includes("reconstruct")) {
if (viz_inputs.options.includes("random order")) {
if (iteration_counter % (25 * 25) == 0) {
mutable shuffled_order = d3.shuffle(shuffled_order);
}
let index = shuffled_order[iteration_counter % 625];
let probe_x = (index % 25) * 4;
let probe_y = ((index / 25) | 0 % 25) * 4;
mutable reconstruct_xy = [probe_x, probe_y];
} else {
let index = unshuffled_order[iteration_counter % 625];
let probe_x = (index % 25) * 4;
let probe_y = ((index / 25) | 0 % 25) * 4;
mutable reconstruct_xy = [probe_x, probe_y];
}
if (viz_inputs.visualization_checkboxes.includes("probe")) {
mutable probe_xy = reconstruct_xy;
}
let gr_step = gradient_descent(potential, reconstruct_xy, 1.0);
if (viz_inputs.options.includes("object positivity")) {
mutable potential = positivity_constraint(gr_step);
} else {
mutable potential = gr_step;
}
mutable iteration_counter += 1;
} else if (!probe_moving) {
mutable potential_static = potential.clone();
}
}
moving_mutable = {
if (probe_moving) {
if (viz_inputs.visualization_checkboxes.includes("gradient")) {
let gr_step = gradient_descent(potential_static, probe_xy, 1.0);
mutable gradient_step = gr_step;
if (viz_inputs.visualization_checkboxes.includes("potential")) {
if (viz_inputs.options.includes("object positivity")) {
mutable potential_static = positivity_constraint(gradient_step);
} else {
mutable potential_static = gradient_step;
}
mutable potential = potential_static.clone();
}
}
}
}
mutable previous_potential_reset_value = 0
reset_potential_mutable = {
if (
viz_inputs.reset_potential_and_probe_position !=
previous_potential_reset_value
) {
mutable mouse_pos = null;
mutable potential = nj.zeros(gpts);
mutable potential_static = nj.zeros(gpts);
mutable probe_xy = [gpts[0] / 2, gpts[1] / 2];
mutable reconstruct_xy = [0, 0];
mutable previous_potential_reset_value =
viz_inputs.reset_potential_and_probe_position;
mutable iteration_counter = 0;
}
}
Purpose-built for displays of data
Observable is your go-to platform for exploring data and creating expressive data visualizations. Use reactive JavaScript notebooks for prototyping and a collaborative canvas for visual data exploration and dashboard creation.
