Atoms

import torch
import nfflr

cell = 4.1 * torch.eye(3)
scaled_positions = torch.tensor([[0,0,0], [0.5, 0.5, 0.5]])
numbers = torch.tensor([24, 22])
atoms = nfflr.Atoms(cell, scaled_positions @ cell, numbers)
atoms

/opt/hostedtoolcache/Python/3.10.13/x64/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
  from .autonotebook import tqdm as notebook_tqdm

2024-03-08 20:08:17,125	INFO util.py:154 -- Missing packages: ['ipywidgets']. Run `pip install -U ipywidgets`, then restart the notebook server for rich notebook output.

2024-03-08 20:08:17,279	INFO util.py:154 -- Missing packages: ['ipywidgets']. Run `pip install -U ipywidgets`, then restart the notebook server for rich notebook output.

Atoms(cell=tensor([[4.1000, 0.0000, 0.0000],
        [0.0000, 4.1000, 0.0000],
        [0.0000, 0.0000, 4.1000]]), positions=tensor([[0.0000, 0.0000, 0.0000],
        [2.0500, 2.0500, 2.0500]]), numbers=tensor([24, 22]), _batch_num_atoms=None)

import ase
from ase.data.colors import jmol_colors
from ase.visualize.plot import plot_atoms

import matplotlib.pyplot as plt
%matplotlib inline

fig, ax = plt.subplots(figsize=(3,3))
ase_atoms = nfflr.to_ase(atoms)
plot_atoms(ase_atoms, ax, radii=0.3, rotation=("10x,20y,0z"), show_unit_cell=2)
ax.axis("off");