Full-time postdoctoral researcher | Part-time design enthusiast | Aspiring cat herder.
visualization_legends = {
let margin = 10;
let plot_dict = [];
if (viz_inputs_a.visualization_checkboxes.includes("potential")) {
plot_dict = [...plot_dict, potential_dictionary];
}
if (viz_inputs_a.visualization_checkboxes.includes("probe")) {
plot_dict = [...plot_dict, probe_dictionary];
}
if (viz_inputs_a.visualization_checkboxes.includes("diffraction")) {
plot_dict = [...plot_dict, dp_dictionary];
}
return html`<div style="display:flex; flex-wrap: wrap;">
${plot_dict.map((x) =>
Plot.legend({
width: viz_inputs_b.width,
marginLeft: margin,
marginRight: margin,
label: x.label,
color: {
domain: x.domain,
scheme: x.scheme,
type: x.type,
exponent: x.exponent,
style: { background: "none" }
}
})
)}`;
}
viewof viz_inputs_a = Inputs.form({
visualization_checkboxes: Inputs.checkbox(
["potential", "probe", "diffraction"],
{
value: ["potential", "probe"],
label: "panels"
}
),
use_multislice: Inputs.toggle({
label: "use multislice?",
value: false
})
})
viewof viz_inputs_b = Inputs.form({
defocus: Inputs.range([-250, 250], {
value: 150,
step: 1,
label: "defocus, Å"
}),
semiangle: Inputs.range([5, 30], {
value: 25,
step: 0.5,
label: "semiangle, mrad"
}),
width: Inputs.range([100, 500], {
value: 265,
step: 1,
label: "width"
})
})
import {
nj,
ComplexNDArray,
fftfreq,
fftshift2D,
corner_crop,
meshgrid2D,
fourier_shift,
ComplexProbe,
electron_wavelength_angstroms,
electron_interaction_parameter
} from "@gvarnavi/ptychography-helper-functions"
potential_raw = {
let typed_array = new Float32Array(
await FileAttachment(
"FCC-slab-potential-7x244x242-float32.npy"
).arrayBuffer()
);
let array = Array.from(typed_array);
let potential_real = nj.zeros([7, 244, 242]);
potential_real.selection.data = array;
return potential_real;
}
potential = {
let scaling =
electron_interaction_parameter(energy_keV * 1e3) /
electron_interaction_parameter(80e3);
return potential_raw.multiply(scaling);
}
potential_dictionary = ({
label: "Projected Electrostatic Potential (rad)",
domain: [0, Math.PI],
scheme: "Purples",
width: potential.shape[2],
height: potential.shape[1],
values: projected_potential.flatten().tolist()
})
projected_potential = {
let [n_slices, nx, ny] = potential.shape;
let total_phase = nj.zeros([nx, ny]);
for (let i = 0; i < n_slices; i++) {
let phase = potential.pick(i, null, null);
total_phase = total_phase.add(phase);
}
return total_phase;
}
complex_potentials = {
let [n_slices, nx, ny] = potential.shape;
let complex_potentials = [...Array(n_slices + 1)].map((d, i) => {
if (i < n_slices) {
let complex_potential = new ComplexNDArray([nx, ny]);
let phase = potential.pick(i, null, null);
let cos_phase = nj.cos(phase);
let sin_phase = nj.sin(phase);
complex_potential.data = nj.stack([cos_phase, sin_phase], -1);
return complex_potential;
} else {
let complex_potential = new ComplexNDArray([nx, ny]);
let cos_phase = nj.cos(projected_potential);
let sin_phase = nj.sin(projected_potential);
complex_potential.data = nj.stack([cos_phase, sin_phase], -1);
return complex_potential;
}
});
return complex_potentials;
}
sampling = [0.1, 0.1]
gpts = [potential.shape[1], potential.shape[2]]
energy_keV = 80
real_space_probe = new ComplexProbe(
gpts,
sampling,
energy_keV * 1e3,
viz_inputs_b.semiangle,
viz_inputs_b.defocus
).build()
probe_dictionary = {
if (viz_inputs_a.visualization_checkboxes.includes("probe")) {
let probe_intensity = fourier_shift(real_space_probe._array, [
probe_xy[1],
-probe_xy[0]
]).abs_sqr();
return {
label: "Illuminating Probe Intensity",
scheme: "Greys",
domain: [0, real_space_probe._array.abs_sqr().max()],
width: probe_intensity.shape[1],
height: probe_intensity.shape[0],
values: probe_intensity.flatten().tolist()
};
}
}
function propagator_array(gpts, sampling, energy, dz) {
let prefactor = electron_wavelength_angstroms(energy) * Math.PI * dz;
let kx = fftfreq(gpts[0], sampling[0]);
let ky = fftfreq(gpts[1], sampling[1]);
let [KX, KY] = meshgrid2D(kx, ky);
let chi = nj.add(KX.multiply(KX), KY.multiply(KY)).multiply(prefactor);
let phase_re = nj.cos(chi);
let phase_im = nj.sin(chi);
let propagator = new ComplexNDArray(gpts);
propagator.data = nj.stack([phase_re, phase_im], -1);
return propagator;
}
function propagate_wavefunction(array, propagator_array) {
let array_fourier = new ComplexNDArray(array.shape);
array_fourier.data = nj.fft(array.data);
array_fourier.data = nj.ifft(array_fourier.multiply(propagator_array).data);
return array_fourier;
}
function multislice_propagation(potential, incoming_probe) {
let [n_slices, nx, ny] = potential.shape;
let wavefunction = incoming_probe.clone();
for (let s = 0; s < n_slices; s++) {
wavefunction = wavefunction.multiply(complex_potentials[s]);
if (s + 1 < n_slices) {
wavefunction = propagate_wavefunction(wavefunction, fixed_dz_propagator);
}
}
return wavefunction;
}
function singleslice_propagation(potential, incoming_probe) {
let wavefunction = incoming_probe.clone();
wavefunction = wavefunction.multiply(complex_potentials[potential.shape[0]]);
return wavefunction;
}
function diffraction_intensity(potential, incoming_probe, [sx, sy]) {
let propagator = viz_inputs_a.use_multislice
? multislice_propagation
: singleslice_propagation;
let exit_wave = propagator(potential, incoming_probe);
let exit_wave_fourier = new ComplexNDArray(exit_wave.shape);
exit_wave_fourier.data = nj.fft(exit_wave.data);
let exit_wave_cropped = corner_crop(exit_wave_fourier.data, [sx, sy]);
return exit_wave_cropped.abs_sqr();
}
fixed_dz_propagator = propagator_array(gpts, sampling, energy_keV * 1e3, 20 / 7)
dp_dictionary = {
if (viz_inputs_a.visualization_checkboxes.includes("diffraction")) {
let dp = diffraction_intensity(
potential,
fourier_shift(real_space_probe._array, [probe_xy[1], probe_xy[0]]),
gpts
);
let dp_intensity = fftshift2D(dp);
return {
label: "Diffraction Intensities",
scheme: "Magma",
domain: [0, 150],
type: "pow",
exponent: 0.5,
width: dp_intensity.shape[1],
height: dp_intensity.shape[0],
values: dp_intensity.flatten().tolist()
};
}
}
mutable probe_xy = [gpts[0] / 2, gpts[1] / 2]
raster_subplot = (dicts, a, viz_width, viz_height) =>
Plot.plot({
width: viz_width,
height: viz_height,
margin: 0,
x: { axis: null },
y: { axis: null },
color: {
label: dicts[a].label,
domain: dicts[a].domain,
scheme: dicts[a].scheme,
type: dicts[a].type,
exponent: dicts[a].exponent,
style: { background: "none" }
},
style: { background: "none" },
marks: [
Plot.raster(dicts[a].values, {
width: dicts[a].width,
height: dicts[a].height
}),
Plot.frame({ strokeWidth: 1 })
]
})
function left_panel_pixels_to_pos(
[px, py],
[x_margin, y_margin],
[x_bandwidth, y_bandwidth]
) {
let [fx, fy] = [(px - x_margin) / x_bandwidth, (py - y_margin) / y_bandwidth];
if (fx > 1) {
return null;
}
return [fx * gpts[0], fy * gpts[1]];
}
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.
