def view_atoms_classifier(image, mg_label, mn_label, o_label, view=True):
x = image.reshape(-1, 3)
mg = x[2:10, :]
mn = x[10:18, :]
o = x[18:, :]
l = x[0, :] * 30
a = x[1, :] * 180
c = np.hstack((l, a))
atoms = Atoms('H')
atoms.set_cell(c)
cell = atoms.get_cell()
t = np.isnan(cell)
tt = np.sum(t)
isnan = False
if not tt == 0:
isnan = True
print(cell)
print(l)
print(a)
_, mg_index = torch.max(mg_label, dim=1)
_, mn_index = torch.max(mn_label, dim=1)
_, o_index = torch.max(o_label, dim=1)
mg_index = mg_index.reshape(8, ).detach().cpu().numpy()
mn_index = mn_index.reshape(8, ).detach().cpu().numpy()
o_index = o_index.reshape(12, ).detach().cpu().numpy()
mg_pos = mg[np.where(mg_index)]
mn_pos = mn[np.where(mn_index)]
o_pos = o[np.where(o_index)]
n_mg = len(mg_pos)
n_mn = len(mn_pos)
n_o = len(o_pos)
if n_mg == 0:
mg_pos = np.array([0.1667, 0.1667, 0.1667]).reshape(1, 3)
n_mg = 1
if n_mn == 0:
mn_pos = np.array([0.1667, 0.1667, 0.1667]).reshape(1, 3)
n_mn = 1
if n_o == 0:
o_pos = np.array([0.1667, 0.1667, 0.1667]).reshape(1, 3)
n_o = 1
pos = np.vstack((mg_pos, mn_pos, o_pos))
scaled_pos = back_to_10_cell(pos, n_mg, n_mn, n_o)
atoms = back_to_real_cell(scaled_pos, cell, n_mg, n_mn, n_o)
atoms.set_pbc([1, 1, 1])
if view:
atoms.edit()
return atoms, x, isnan
