data2 = np.load('./data/PTCDA_Xe.npz')
X2 = data2['data']
afm_dim2 = (data2['lengthX'], data2['lengthY'])
assert afm_dim1 == afm_dim2
afm_dim = afm_dim1
X_exp = apply_preprocessing_exp([X1[None], X2[None]], afm_dim)
# Load model for sim
input_shape = afmulator.scan_dim[:2] + (10, )
model_sim = ESUNet(n_in=2,
n_out=2,
input_shape=input_shape,
last_relu=[False, True])
load_pretrained_weights(model_sim, tip_type='CO-Xe')
# Load model for exp (need two models because of different input sizes)
input_shape = X_exp[0].shape[1:]
model_exp = ESUNet(n_in=2,
n_out=2,
input_shape=input_shape,
last_relu=[False, True])
load_pretrained_weights(model_exp, tip_type='CO-Xe')
# Get predictions
pred_sim = model_sim.predict(X_sim)
pred_exp = model_exp.predict(X_exp)
# Create figure grid
width_ratios = [6, 4, 2]
X_bcb = apply_preprocessing_bcb([X_bcb[None]], afm_dim_bcb)
# Load PTCDA data and preprocess
data_ptcda = np.load('./data/PTCDA_CO.npz')
X_ptcda = data_ptcda['data']
afm_dim_ptcda = (data_ptcda['lengthX'], data_ptcda['lengthY'])
X_ptcda = apply_preprocessing_ptcda([X_ptcda[None]], afm_dim_ptcda)
# Load model for simulations
input_shape = afmulator.scan_dim[:2] + (afmulator.scan_dim[2] -
afmulator.df_steps, )
model_sim = ESUNet(n_in=1,
n_out=2,
input_shape=input_shape,
last_relu=[False, True])
load_pretrained_weights(model_sim, tip_type='CO')
# Load model for BCB
model_bcb = ESUNet(n_in=1,
n_out=2,
input_shape=X_bcb[0].shape[1:],
last_relu=[False, True])
load_pretrained_weights(model_bcb, tip_type='CO')
# Load model for PTCDA
model_ptcda = ESUNet(n_in=1,
n_out=2,
input_shape=X_ptcda[0].shape[1:],
last_relu=[False, True])
load_pretrained_weights(model_ptcda, tip_type='CO')
aux_maps,
molecules,
batch_size=1,
distAbove=5.2,
iZPPs=[54, 17], # Xe, Cl
Qs=[[30, -60, 30, 0], [-0.3, 0, 0, 0]],
QZs=[[0.1, 0, -0.1, 0], [0, 0, 0, 0]])
# Load model for Cl-CO
input_shape = afmulator.scan_dim[:2] + (afmulator.scan_dim[2] -
afmulator.df_steps, )
model_Cl_CO = ESUNet(n_in=2,
n_out=2,
input_shape=input_shape,
last_relu=[False, True])
load_pretrained_weights(model_Cl_CO, tip_type='Cl-CO')
# Load model for Xe-Cl
input_shape = afmulator.scan_dim[:2] + (afmulator.scan_dim[2] -
afmulator.df_steps, )
model_Xe_Cl = ESUNet(n_in=2,
n_out=2,
input_shape=input_shape,
last_relu=[False, False])
load_pretrained_weights(model_Xe_Cl, tip_type='Xe-Cl')
# Loop over molecules and plot
width_ratios = [6, 12]
fig = plt.figure(figsize=(sum(width_ratios), 6 * len(molecules)))
fig_grid = fig.add_gridspec(len(molecules), 1, wspace=0, hspace=0.03)
for ib, (batch_Cl_CO,
# Define AuxMaps
scan_window = [afmulator.scan_window[0][:2], afmulator.scan_window[1][:2]]
aux_maps = [aux.ESMap(scanner=afmulator.scanner, zmin=-2.0, iso=0.1)]
# Define generator
trainer = InverseAFMtrainer(afmulator, aux_maps, molecules, **generator_kwargs)
# Load model
input_shape = afmulator.scan_dim[:2] + (afmulator.scan_dim[2] -
afmulator.df_steps, )
model = ESUNet(n_in=2,
n_out=2,
input_shape=input_shape,
last_relu=[False, True])
load_pretrained_weights(model, tip_type='CO-Xe')
# Loop over molecules and plot
width_ratios = [3, 6, 8]
fig = plt.figure(figsize=(sum(width_ratios), 4.05 * len(molecules)))
fig_grid = fig.add_gridspec(len(molecules), 1, wspace=0, hspace=0.03)
for ib, batch in enumerate(trainer):
# Get batch and predict
X, Y, xyzs = apply_preprocessing(batch)
pred = model.predict(X)
# Create plot grid
sample_grid = fig_grid[ib, 0].subgridspec(1,
3,
wspace=0.01,