{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "324469a0",
   "metadata": {},
   "source": [
    "# Force Fields for Materials"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de5d8c26",
   "metadata": {},
   "source": [
    "Machine learning force fields for studying materials has become increasingly popular. However, the large size of the physical systems associated with most materials requires some tricks to keep the amount of compute and required data sufficiently low. \n",
    "\n",
    "In this tutorial, we will describe some of those tricks and how they can be implemented in the SchNetPack pipeline, namely:\n",
    "\n",
    "- **periodic boundary conditions (PBC)**: PBC allow to effectively reduce the number of simulated particles to only a fraction of the actual system's size. This is achieved by considering a relatively small simulation box, which is periodically repeated at its boundaries. In most cases, the resulting simulated periodic structure is a good approximation of the system under consideration.\n",
    "- **cached neighbor lists**: For large systems, the computation of all neighbors is expensive. In the training procedure this problem can be circumvented by utilizing neighbor list caching. This way, the neighbors must only be computed in the first epoch. In the subsequent epochs the cached neighbor lists can be loaded, which reduces the training time tremendously.\n",
    "- **neighbor lists with skin buffer**: In the scope of molecular dynamics simulations or structure relaxations, caching neighbor lists is not possible since the neighborhoods change with each integration step. Hence, it is recommended to use a neighbor list that utilizes a so-called skin buffer. The latter represents a thin layer that extends the cutoff region. It allows to reuse the calculated neighbor lists for samples with similar atomic positions. Only when at least one atom penetrates the skin buffer, the neighbor list is recalculated.\n",
    "- **filtering out neighbors (neighbor list postprocessing)**: Also in the feed forward pass of the network, a large number of neighbors, and thus interactions, can result in slow inference and training. In some scenarios it is crucial to have as few operations as possible in the model to ensure fast inference. This can be achieved, e.g., by filtering out some neighbors from the neighbor list.\n",
    "- **prediction target emphasizing**: In some occasions it may be useful to exclude the properties of some atoms from the training procedure. For example you might want to focus the training on the forces of only some atoms and neglect the rest. An exemplary use case would be a model used for structure optimization where some atoms are fixed during the simulation (zero forces). Or when filtering out neighbors of certain atoms, it might be reasonable to exclude the corresponding atomic properties from the training loss. \n",
    "\n",
    "\n",
    "In the following tutorial, we will first describe how the dataset must be prepared to allow for utilizing the above-mentioned tricks. Subsequently, we explain how the configs in the SchNetPack framework must be adapted for training and inference, accordingly. The dataset preparation part is based on the tutorial script \"tutorial_01_preparing_data\". Please make sure you have understood the concepts introduced there, before continuing with this tutorial."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9374b076",
   "metadata": {},
   "source": [
    "## Preparing the Dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95cad555",
   "metadata": {},
   "source": [
    "First we will demonstrate how to prepare the dataset to enable the use of periodic boundary conditions (pbc). \"tutorial_01_preparing_data\" describes how to prepare your own data for SchNetPack. For the purpose of this tutorial we will create a new dataset and add the Si$_{16}$ diamond structure to the database. More structures can be added to the dataset by repeating the process with more datapoints in an iterative manner. In order to create the sample structure we need the atomic positions, the cell shape and the atom types. Furthermore, we want to add the properties \"energy\", \"forces\" and \"stress\" to our dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "666cc817",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "n_atoms = 16\n",
    "positions = np.array(\n",
    "    [\n",
    "        [5.88321935, 2.88297608, 6.11028356],\n",
    "        [3.55048572, 4.80964742, 2.77677731],\n",
    "        [0.40720652, 6.73142071, 3.88666154],\n",
    "        [2.73860477, 0.96144918, 0.55409947],\n",
    "        [0.40533082, 4.80977264, 0.55365277],\n",
    "        [2.73798189, 2.88306742, 3.88717458],\n",
    "        [5.88129464, 0.96133042, 2.77731562],\n",
    "        [3.54993842, 6.73126316, 6.10980443],\n",
    "        [0.4072795, 2.88356071, 3.88798019],\n",
    "        [2.7361945, 4.80839638, 0.55454406],\n",
    "        [5.8807282, 6.732434, 6.10842469],\n",
    "        [3.55049053, 0.96113024, 2.77692083],\n",
    "        [5.88118878, 4.80922032, 2.7759235],\n",
    "        [3.55233987, 2.8843493, 6.10940225],\n",
    "        [0.4078124, 0.96031906, 0.5555672],\n",
    "        [2.73802399, 6.73163254, 3.88698037],\n",
    "    ]\n",
    ")\n",
    "cell = np.array(\n",
    "    [\n",
    "        [6.287023489207423, -0.00034751886075738795, -0.0008093810364881463],\n",
    "        [0.00048186949720712026, 7.696440684406158, -0.001909478919115524],\n",
    "        [0.0010077843421425583, -0.0033467698530393886, 6.666654324468158],\n",
    "    ]\n",
    ")\n",
    "symbols = [\"Si\"] * n_atoms\n",
    "energy = np.array([-10169.33552017])\n",
    "forces = np.array(\n",
    "    [\n",
    "        [0.02808107, -0.02261815, -0.00868415],\n",
    "        [-0.03619687, -0.02530285, -0.00912962],\n",
    "        [-0.03512621, 0.02608594, 0.00913623],\n",
    "        [0.02955523, 0.02289934, 0.0089936],\n",
    "        [-0.02828359, 0.02255927, 0.00871455],\n",
    "        [0.03636321, 0.02545969, 0.00911801],\n",
    "        [0.0352177, -0.02613079, -0.00927739],\n",
    "        [-0.02963064, -0.0227443, -0.00894253],\n",
    "        [-0.03343582, 0.02324933, 0.00651742],\n",
    "        [0.03955335, 0.0259127, 0.00306112],\n",
    "        [0.03927719, -0.02677768, -0.00513233],\n",
    "        [-0.0332425, -0.02411682, -0.00464783],\n",
    "        [0.03358715, -0.02328505, -0.00626828],\n",
    "        [-0.03953832, -0.02600458, -0.00316128],\n",
    "        [-0.03932441, 0.02681881, 0.0048871],\n",
    "        [0.03314345, 0.0239951, 0.00481536],\n",
    "    ]\n",
    ")\n",
    "stress = np.array(\n",
    "    [\n",
    "        [\n",
    "            [-2.08967984e-02, 1.52890659e-06, 1.44133597e-06],\n",
    "            [1.52890659e-06, -6.45087059e-03, -7.26463797e-04],\n",
    "            [1.44133597e-06, -7.26463797e-04, -6.04950702e-03],\n",
    "        ]\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2b99067",
   "metadata": {},
   "source": [
    "For the purpose of filtering out certain neighbors from the neighbor list, one has to specify a list of atom indices. Between the corresponding atoms, all interactions are neglected. In our exemplary system we neglect all interactions between the atoms with index 4, 10, and 15."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af4bf0e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "filtered_out_neighbors = np.array([4, 10, 15])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "498bf755",
   "metadata": {},
   "source": [
    "To specify the atoms, which targets should be considered in the model optimization, one must define a list of booleans, indicating considered and neglected atoms. This boolean array should be stored in the database along with other sample properties such as, e.g., energy, forces, and the array of filtered out neighbors.\n",
    "\n",
    "For our exemplary system of 16 atoms, the array of considered atoms could be defined as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61424286",
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize array specifying considered atoms\n",
    "considered_atoms = np.ones(n_atoms, dtype=bool)\n",
    "\n",
    "# atom 4 and atom 5 should be neglected in the model optimization\n",
    "considered_atoms[[4, 10, 15]] = False"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e9b55c6",
   "metadata": {},
   "source": [
    "Before we can add our new data to a database for training, we will need to transform it to the correct format. The atomic structure of the new system is stored in an `ase.Atoms` object. In contrast to datasets without periodic boundary conditions we need to set the flag `pbc=True` when the `ase.Atoms` object is created. Both the chemical properties of the structure, as well as the settings arguments for filtering neighbors and defining the atoms to consider during training are stored in the data dictionary corresponding to the new structure. All properties of the data dictionary need to be stored as `np.ndarray`. Please note, that cell dependent properties need to be unsqueezed in the first dimension."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "130b34f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ase import Atoms\n",
    "\n",
    "atoms = Atoms(\n",
    "    symbols=symbols,\n",
    "    positions=positions,\n",
    "    cell=cell,\n",
    "    pbc=True,\n",
    ")\n",
    "data = dict(\n",
    "    energy=np.array(energy),\n",
    "    forces=np.array(forces),\n",
    "    stress=np.array(stress),\n",
    "    considered_atoms=considered_atoms,\n",
    "    filtered_out_neighbors=filtered_out_neighbors,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b628d4bf",
   "metadata": {},
   "source": [
    "Just as with molecular datasets, we can use the `create_dataset` function in order to build a new database. Since we added the `considered_atoms` and `filtered_out_neighbors` to our data dictionary, we also need to define these unitless properties in the `property_unit_dict`. Please note, that all structures in a dataset need to have an equal set of properties. In case some structures do not filter out any neighbors, the `filtered_out_neighbors` has to be set with an empty list."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fcb88102",
   "metadata": {},
   "outputs": [],
   "source": [
    "from schnetpack.data import create_dataset, AtomsDataFormat\n",
    "\n",
    "db_path = \"./si16.db\"\n",
    "new_dataset = create_dataset(\n",
    "    datapath=db_path,\n",
    "    format=AtomsDataFormat.ASE,\n",
    "    distance_unit=\"Ang\",\n",
    "    property_unit_dict=dict(\n",
    "        energy=\"eV\",\n",
    "        forces=\"eV/Ang\",\n",
    "        stress=\"eV/Ang/Ang/Ang\",\n",
    "        considered_atoms=\"\",\n",
    "        filtered_out_neighbors=\"\",\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cba91e55",
   "metadata": {},
   "source": [
    "We can now add our datapoint to the new dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11bedc4e",
   "metadata": {},
   "outputs": [],
   "source": [
    "new_dataset.add_systems(\n",
    "    property_list=[data],\n",
    "    atoms_list=[atoms],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08fc4404",
   "metadata": {},
   "source": [
    "It is important that the entries in the dataset have the appropriate dimensions. To this end we print out the shapes of the tensors in our first data point."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "098c16b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "for key, value in new_dataset[0].items():\n",
    "    print(key, value.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82172a3d",
   "metadata": {},
   "source": [
    "Let`s remove the database again, because it will no longer be needed during this tutorial:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b586f44",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.remove(db_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38b4456c",
   "metadata": {},
   "source": [
    "## Adapting the Configs in the SchNetPack Framework"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6f7ea4f",
   "metadata": {},
   "source": [
    "Now we will cover, how to adapt config files to enable the use of the above mentioned tricks in SchNetPack's training procedure and MD framework"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50391832",
   "metadata": {},
   "source": [
    "### SchNetPack Training"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fecc4a5f",
   "metadata": {},
   "source": [
    "Provided that appropriate neighbor list providers (ASE or MatScipy) are used, this is sufficient for using pbc in the SchNet/PaiNN framework.\n",
    "\n",
    "The neighbor list caching is implemented in the schnetpack transform ``schnetpack.transform.CachedNeighborList``. SchNetPack provides a transform for caching neighbor lists. It basically functions as a wrapper around a common neighbor lists. For further information regarding the CachedNeighborList please refer to the corresponding docstring in the SchNetPack code.\n",
    "\n",
    "Neighbors can be filtered out by using the neighbor list postprocessing transform ``schnetpack.transform.FilterNeighbors``.\n",
    "\n",
    "To ensure that only the specified atoms are considered for the training on a certain property, the respective ModelOutput object has to be adapted. This is achieved by using so-called constraints. Each ModelOutput object takes a list of constraints. For a precise explanation on how to use ``schnetpack.task.ModelOutput`` please refer to notebook \"tutorial_02_qm9\". To specify the selection of atoms for training we use the constraint transform ``schnetpack.task.ConsiderOnlySelectedAtoms``. It has the attribute selection_name, which should be a string linking to the array of specified atoms stored in the database.\n",
    "\n",
    "For adding the stress property to the learning task, we need to make some modifications to configs. First of all, the `Forces` output module needs a stress key and the flag `calc_stress=True` needs to be set. Then we also need to add the stress predictions to the `outputs`, so it can be included in the loss function. Since the absolute values of the stress tensors are generally significantly lower than for energy and forces we need to select a comparably high `loss_weight`. Finally, we need to add the `Strain` module to the input modules *before* pairwise distances.\n",
    "\n",
    "The following is an example of an experiment config file that utilizes the above-mentioned tricks:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "097913c7",
   "metadata": {},
   "source": [
    "```yaml\n",
    "# @package _global_\n",
    "\n",
    "defaults:\n",
    "  - override /model: nnp\n",
    "\n",
    "run.path: runs\n",
    "\n",
    "globals:\n",
    "  cutoff: 5.\n",
    "  lr: 5e-4\n",
    "  energy_key: energy\n",
    "  forces_key: forces\n",
    "  stress_key: stress\n",
    "data:\n",
    "  distance_unit: Ang\n",
    "  property_units:\n",
    "    energy: eV\n",
    "    forces: eV/Ang\n",
    "    stress: eV/Ang/Ang/Ang\n",
    "  transforms:\n",
    "    - _target_: schnetpack.transform.RemoveOffsets\n",
    "      property: energy\n",
    "      remove_mean: True\n",
    "    - _target_: schnetpack.transform.CachedNeighborList\n",
    "      cache_path: ${run.work_dir}/cache\n",
    "      keep_cache: False\n",
    "      neighbor_list:\n",
    "        _target_: schnetpack.transform.MatScipyNeighborList\n",
    "        cutoff: ${globals.cutoff}\n",
    "      nbh_transforms:\n",
    "        - _target_: schnetpack.transform.FilterNeighbors\n",
    "          selection_name: filtered_out_neighbors\n",
    "    - _target_: schnetpack.transform.CastTo32\n",
    "  test_transforms:\n",
    "    - _target_: schnetpack.transform.RemoveOffsets\n",
    "      property: energy\n",
    "      remove_mean: True\n",
    "    - _target_: schnetpack.transform.MatScipyNeighborList\n",
    "      cutoff: ${globals.cutoff}\n",
    "    - _target_: schnetpack.transform.FilterNeighbors\n",
    "      selection_name: filtered_out_neighbors\n",
    "    - _target_: schnetpack.transform.CastTo32\n",
    "\n",
    "model:\n",
    "  input_modules:\n",
    "    - _target_: schnetpack.atomistic.Strain\n",
    "    - _target_: schnetpack.atomistic.PairwiseDistances\n",
    "  output_modules:\n",
    "    - _target_: schnetpack.atomistic.Atomwise\n",
    "      output_key: ${globals.energy_key}\n",
    "      n_in: ${model.representation.n_atom_basis}\n",
    "      aggregation_mode: sum\n",
    "    - _target_: schnetpack.atomistic.Forces\n",
    "      energy_key: ${globals.energy_key}\n",
    "      force_key: ${globals.forces_key}\n",
    "      stress_key: ${globals.stress_key}\n",
    "      calc_stress: True\n",
    "  postprocessors:\n",
    "    - _target_: schnetpack.transform.CastTo64\n",
    "    - _target_: schnetpack.transform.AddOffsets\n",
    "      property: energy\n",
    "      add_mean: True\n",
    "\n",
    "task:\n",
    "  optimizer_cls: torch.optim.AdamW                                        \n",
    "  optimizer_args:                                                         \n",
    "    lr: ${globals.lr}                                                     \n",
    "    weight_decay: 0.01  \n",
    "  outputs:\n",
    "    - _target_: schnetpack.task.ModelOutput\n",
    "      name: ${globals.energy_key}\n",
    "      loss_fn:\n",
    "        _target_: torch.nn.MSELoss\n",
    "      metrics:\n",
    "        mae:\n",
    "          _target_: torchmetrics.regression.MeanAbsoluteError\n",
    "        mse:\n",
    "          _target_: torchmetrics.regression.MeanSquaredError\n",
    "      loss_weight: 0.0\n",
    "    - _target_: schnetpack.task.ModelOutput\n",
    "      name: ${globals.forces_key}\n",
    "      constraints: \n",
    "        - _target_: schnetpack.task.ConsiderOnlySelectedAtoms\n",
    "          selection_name: considered_atoms\n",
    "      loss_fn:\n",
    "        _target_: torch.nn.MSELoss\n",
    "      metrics:\n",
    "        mae:\n",
    "          _target_: torchmetrics.regression.MeanAbsoluteError\n",
    "        mse:\n",
    "          _target_: torchmetrics.regression.MeanSquaredError\n",
    "      loss_weight: 1.0\n",
    "    - _target_: schnetpack.task.ModelOutput\n",
    "      name: ${globals.stress}\n",
    "      loss_fn:\n",
    "        _target_: torch.nn.MSELoss\n",
    "      metrics:\n",
    "        mae:\n",
    "          _target_: torchmetrics.regression.MeanAbsoluteError\n",
    "        mse:\n",
    "          _target_: torchmetrics.regression.MeanSquaredError\n",
    "      loss_weight: 100.\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "938b08f2",
   "metadata": {},
   "source": [
    "### SchNetPack MD and Structure Relaxation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a4c49c0",
   "metadata": {},
   "source": [
    "In the framework of MD simulations and structure relaxations it is preferable to utilize neighbor lists with skin buffers. The corresponding class in SchNetPack is called ``schnetpack.transform.SkinNeighborList``. It takes as an argument a conventional neighbor list class such as, e.g., ASENeighborList, post-processing transforms for manipulating the neighbor lists and the cutoff skin which defines the size of the skin buffer around the actual cutoff region. Please choose a sufficiently large cutoff skin value to ensure that between two subsequent samples no atom can penetrate through the skin into the cutoff sphere of another atom if it is not in the neighbor list of that atom."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62137cf8-8f40-41d0-9d98-a2634394e5b0",
   "metadata": {},
   "source": [
    "```yaml\n",
    "_target_: schnetpack.md.calculators.SchNetPackCalculator\n",
    "required_properties:\n",
    "  - energy\n",
    "  - forces\n",
    "model: ???\n",
    "force_label: forces\n",
    "energy_units: kcal / mol\n",
    "position_units: Angstrom\n",
    "energy_label: energy\n",
    "stress_label: null\n",
    "script_model: false\n",
    "\n",
    "defaults:\n",
    "  - neighbor_list: \n",
    "    _target_: schnetpack.transform.SkinNeighborList\n",
    "    cutoff_skin: 2.0\n",
    "    neighbor_list:\n",
    "      _target_: schnetpack.transform.ASENeighborList\n",
    "      cutoff: 5.0\n",
    "    nbh_transforms:\n",
    "      - _target_: schnetpack.transform.FilteredNeighbors\n",
    "        selection_name: filtered_out_neighbors\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8b3bfdb",
   "metadata": {},
   "source": [
    "## Structure Optimization\n",
    "\n",
    "In order demonstrate the structure optimization with a trained model, we use our optimized `atoms` structure and add noise to the positions. Subsequently, we use a trained model and run the optimization on the noisy structure. This should return the local optimum again. Let`s first take a look our sample structure:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "705fe868",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ase.visualize.plot import plot_atoms\n",
    "\n",
    "plot_atoms(atoms, rotation=(\"90x,0y,0z\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc6bc8bf",
   "metadata": {},
   "source": [
    "Now we create the noisy structure by adding noise to the positions and the cell:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01fc07a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "noise = 0.7\n",
    "atms_noise = atoms.copy()\n",
    "atms_noise.positions += (np.random.random(atms_noise.positions.shape) - 0.5) * noise\n",
    "atms_noise.cell += (np.random.random(atms_noise.cell.shape) - 0.5) * noise\n",
    "plot_atoms(atms_noise, rotation=(\"90x,0y,0z\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6b675dc",
   "metadata": {},
   "source": [
    "As we can see, the new structure is visibly deformed. With the use of a trained model, the structure optimizer from `ase` and the `schnetpack.interfaces.ase_interface.SpkCalculator` we will now denoise our structure. The model has been trained for the purpose of this tutorial on a small dataset of 500 Si16 structures form different local minima. It uses a SO3Net representation with 64 features and 2 interaction layers. Since we also added noise to the cell, we will wrap the `atoms` object in an `ExpCellFilter`. Now let`s run the optimization:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d17e3ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "import schnetpack as spk\n",
    "from schnetpack.transform import ASENeighborList\n",
    "from schnetpack.interfaces.ase_interface import SpkCalculator\n",
    "from ase.optimize.lbfgs import LBFGS\n",
    "from ase.constraints import ExpCellFilter\n",
    "\n",
    "# Load model\n",
    "model_path = \"../../tests/testdata/si16.model\"\n",
    "\n",
    "# Get calculator and optimizer\n",
    "atms_noise.calc = SpkCalculator(\n",
    "    model=model_path,\n",
    "    stress_key=\"stress\",\n",
    "    neighbor_list=ASENeighborList(cutoff=7.0),\n",
    "    energy_unit=\"eV\",\n",
    ")\n",
    "optimizer = LBFGS(ExpCellFilter(atms_noise))\n",
    "\n",
    "# run optimization\n",
    "optimizer.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "caa86180",
   "metadata": {},
   "source": [
    "As we can see, the structure optimization has removed the noise, and we obtain our stable structure again (the stable structure may be rotated by some degree):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23cacd5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_atoms(atms_noise, rotation=(\"90x,0y,0z\"))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}