tutorial_rename_tensors.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c041be2a",
   "metadata": {},
   "source": [
    "# Introduction"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a764e7e",
   "metadata": {},
   "source": [
    "This brief tutorial aims to explain how the Python module ``utils/rename_tensors.py`` can be used to rename the ``Labels`` names of ``TensorMap`` and ``TensorBlock`` objects.\n",
    "\n",
    "The purpose of these functions is allow for greater transferability and interoperability between equistore-based workflows, such as those produced by individuals/groups in different fields and with different naming conventions. This has already been raised as an [issue in the equistore repository](https://github.com/lab-cosmo/equistore/issues/61).\n",
    "\n",
    "It is worth pointing out that ``utils/rename_tensors.py`` only goes part of the way to solving the problem. It is only a Python-side implementation of the solution and is not rigorously tested. Furthermore, it works by creating **new** ``TensorBlock`` and ``TensorMap`` objects with ``Labels`` with the new names, instead of actuallly **renaming** the ``Labels`` of the existing object. Creating new objects altogether, instead of changing existing ones, has memory implications, as the data stored within the the ``TensorBlocks`` and ``TensorMaps`` being \"renamed\" is essentially copied.\n",
    "\n",
    "This is therefore a temporary solution to the problem. Ideally, the real solution will be implemented Rust-side, change the ``Labels`` names of the existing objects, be rigorously tested, and be more elegant both under-the-hood and to the user.\n",
    "\n",
    "This current Python-side implementation features a single \"rename\" method, where the user needs to specify explictly the whole name structure. The user-facing API would also have more specific methods such as ``TensorMap.rename_keys()``, ``TensorMap.rename_properties()``, ``TensorMap.rename_components``, ``TensorMap.rename_gradient_components()`` etc.\n",
    "\n",
    "Nevertheless, this script does a good enough job for now. In this tutorial we will \"rename\" some example ``TensorBlocks`` and ``TensorMaps``. Hopefully, while a more permanent solution is worked on, this can be useful enough to save others' time."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38427af4",
   "metadata": {},
   "source": [
    "# Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "68010046",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pprint import pprint as pprint\n",
    "\n",
    "from equistore import io, TensorMap, TensorBlock\n",
    "from utils.rename_tensors import *"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f381f93b",
   "metadata": {},
   "source": [
    "# Example Data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2bf634ff",
   "metadata": {},
   "source": [
    "We will load an example ``TensorMap`` that has been saved to file, ***data/tensormap_periodic.npz***, and use this to illustrate how renaming can be performed.\n",
    "\n",
    "For context, though not necessarily important for understanding of the tutorial, this example ``TensorMap`` is a descriptor for a small dataset of ~100 periodic crystal structures, generated using the Rascaline ``SphericalExpension`` calculator.  For the illustrative purposes of this tutorial, the descriptor was generated with calculation of both 'cell' and 'positions' gradients."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3965380c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TensorMap with 70 blocks\n",
       "keys: ['spherical_harmonics_l' 'species_center' 'species_neighbor']\n",
       "                  0                   1                1\n",
       "                  1                   1                1\n",
       "                  2                   1                1\n",
       "               ...\n",
       "                  2                   8                8\n",
       "                  3                   8                8\n",
       "                  4                   8                8"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tensormap = io.load('data/tensormap_periodic.npz')\n",
    "tensormap"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58fc400d",
   "metadata": {},
   "source": [
    "# Revealing the ``Labels`` Name Structure of a ``TensorMap``"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "475600a7",
   "metadata": {},
   "source": [
    "The first step to renaming a ``TensorMap`` is to find out its current ``Labels`` name structure. This can be done with the ``rename_tensors.get_tensormap_name_structure()`` function. This returns a nested dictionary specifying the **complete** name structure. We will 'pretty print' the ``dict`` output of this function and inspect the structure as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "307d027f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'components': [('spherical_harmonics_m',)],\n",
      " 'gradients': {'cell': {'components': [('direction_1',),\n",
      "                                       ('direction_2',),\n",
      "                                       ('spherical_harmonics_m',)],\n",
      "                        'properties': ('n',),\n",
      "                        'samples': ('sample',)},\n",
      "               'positions': {'components': [('direction',),\n",
      "                                            ('spherical_harmonics_m',)],\n",
      "                             'properties': ('n',),\n",
      "                             'samples': ('sample', 'structure', 'atom')}},\n",
      " 'keys': ('spherical_harmonics_l', 'species_center', 'species_neighbor'),\n",
      " 'properties': ('n',),\n",
      " 'samples': ('structure', 'center')}\n"
     ]
    }
   ],
   "source": [
    "pprint(get_tensormap_name_structure(tensormap))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "905cc09d",
   "metadata": {},
   "source": [
    "Upon inspection you can see the names of the various ``Labels`` objects that label the data in the ``TensorMap``. Remember that a ``TensorMap`` is a collection of one or more ``TensorBlock``s, indexed by a key. Here we can see that this descriptor contains some 70 ``TensorBlocks`` indicated by a key that specifies ``('spherical_harmonics_l', 'species_center', 'species_neighbor')``.\n",
    "\n",
    "Each ``TensorBlock`` is comprised of 'samples', 'components', and 'properties' dimensions. Here, also, each ``TensorBlock`` has associated ``Gradient TensorBlock``s for **both** 'cell' and 'positions' gradients, as these are what were calculated upon generation of the descriptor.\n",
    "\n",
    "Every ``Gradient TensorBlock``, like a normal ``TensorBlock``, has 'samples', 'components', and 'properties' dimensions. There are some important aspects to how these are named.\n",
    "\n",
    "* The 'properties' Labels names of the ``TensorBlock`` are/should be the **same** as those of the associated ``Gradient TensorBlock``s.\n",
    "    * i.e. above you can see that the properties for the radial channels are called ``'n'`` for the ``TensorBlock``s and associated ``Gradient TensorBlocks``.\n",
    "* Where there are $N$ components names for the ``TensorBlock``, **the final $N$ names** of the components for the associated ``Gradient TensorBlock``s are/should be be the same.\n",
    "    * i.e. the ``TensorBlock``s have $N=1$ name for the components, ``'spherical_harmonics_m'``, so the final 1 components names for the ``Gradient TensorBlock``s is also ``'spherical_harmonics_m'``. The preceding names are different: for the 'cell' gradients this correspond to pairs of distortions in each of the cartesian axes (i.e. xx, xy, xz, yx, yy, yz, zx, zy, or zz); and for the 'positions' gradients this just corresponds to displacement along a cartesian axis (i.e. x, y, or z).\n",
    "* Names of samples of the ``TensorBlock`` compared to its associated ``Gradient TensorBlocks`` are different.\n",
    "    * Here, each ``TensorBlock`` has an associated key for the combination of 'structure' (i.e. the periodic crystal structure) and 'center' (i.e. an atom in that structure). For the ``Gradient TensorBlock``s, the names of the samples are different. For 'cell' gradients, the samples just correspond to 'samples', i.e. how the values of a given ('structure, 'center') combination change upon cell distortion in each of the ('direction1', 'direction2') pair for each 'spherical_harmonic_m'. For 'positions' gradients the samples are given as ('sample', 'structure', 'atom'), which tracks the gradient for displacement of an 'atom' in the 'structure' changes for a given sample (i.e. 'structure' and 'center') and 'spherical_harmonics_m'."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a6a9e757",
   "metadata": {},
   "source": [
    "# Renaming a ``TensorMap``"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b226d49",
   "metadata": {},
   "source": [
    "The best way to use these functions is to copy the output from the ``get_tensormap_name_structure()`` function and edit the names as desired.\n",
    "\n",
    "The cell below features this copied-and-pasted ``dict`` from above, but with some changes, for no good reason at all except to illustrate the functionality:\n",
    "\n",
    "* ``TensorMap`` keys names changed:\n",
    "    * ``('spherical_harmonics_l', 'species_center', 'species_neighbor')`` to ``('l', 'species_center', 'species_neighbor')``\n",
    "* ``TensorBlock`` samples names changed:\n",
    "    * ``('structure', 'center')`` to ``('crystal', 'center')``"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "5da9a5e0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'components': [('spherical_harmonics_m',)],\n",
      " 'gradients': {'cell': {'components': [('direction_1',),\n",
      "                                       ('direction_2',),\n",
      "                                       ('spherical_harmonics_m',)],\n",
      "                        'properties': ('n',),\n",
      "                        'samples': ('sample',)},\n",
      "               'positions': {'components': [('direction',),\n",
      "                                            ('spherical_harmonics_m',)],\n",
      "                             'properties': ('n',),\n",
      "                             'samples': ('sample', 'structure', 'atom')}},\n",
      " 'keys': ('l', 'species_center', 'species_neighbor'),\n",
      " 'properties': ('n',),\n",
      " 'samples': ('crystal', 'center')}\n"
     ]
    }
   ],
   "source": [
    "# Change keys and samples names\n",
    "new_names = {\n",
    " 'components': [('spherical_harmonics_m',)],\n",
    " 'gradients': {'cell': {'components': [('direction_1',),\n",
    "                                       ('direction_2',),\n",
    "                                       ('spherical_harmonics_m',)],\n",
    "                        'properties': ('n',),\n",
    "                        'samples': ('sample',)},\n",
    "               'positions': {'components': [('direction',),\n",
    "                                            ('spherical_harmonics_m',)],\n",
    "                             'properties': ('n',),\n",
    "                             'samples': ('sample', 'structure', 'atom')}},\n",
    " 'keys': ('l', 'species_center', 'species_neighbor'),\n",
    " 'properties': ('n',),\n",
    " 'samples': ('crystal', 'center')\n",
    "}\n",
    "\n",
    "# \"Rename\" (i.e. create a new) tensormap with the desired names\n",
    "new_tensormap = rename_tensormap(tensormap, new_names)\n",
    "\n",
    "pprint(get_tensormap_name_structure(new_tensormap))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0d6088d",
   "metadata": {},
   "source": [
    "Now, knowing that the properties names should be consistent between ``TensorBlock``s and ``Gradient TensorBlock``s, let's try to changes one and not the other. An error is thrown."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "02458589",
   "metadata": {},
   "outputs": [
    {
     "ename": "AssertionError",
     "evalue": "The names of TensorBlock properties Labels should be the same as those for the Gradient TensorBlocks. You have passed ['radial_channel_n'] as the names of the properties of the TensorBlock, and [('n',), ('n',)] as the names of the properties for the Gradient TensorBlocks.",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mAssertionError\u001b[0m                            Traceback (most recent call last)",
      "Input \u001b[0;32mIn [5]\u001b[0m, in \u001b[0;36m<cell line: 3>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      1\u001b[0m new_names[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mproperties\u001b[39m\u001b[38;5;124m'\u001b[39m] \u001b[38;5;241m=\u001b[39m (\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mradial_channel_n\u001b[39m\u001b[38;5;124m'\u001b[39m,)\n\u001b[0;32m----> 3\u001b[0m new_tensormap \u001b[38;5;241m=\u001b[39m \u001b[43mrename_tensormap\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtensormap\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mnew_names\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/Documents/phd/code/equi/equistore-examples/utils/rename_tensors.py:265\u001b[0m, in \u001b[0;36mrename_tensormap\u001b[0;34m(tensormap, new_names)\u001b[0m\n\u001b[1;32m    222\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mrename_tensormap\u001b[39m(\n\u001b[1;32m    223\u001b[0m     tensormap: equistore\u001b[38;5;241m.\u001b[39mtensor\u001b[38;5;241m.\u001b[39mTensorMap, new_names: \u001b[38;5;28mdict\u001b[39m\n\u001b[1;32m    224\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m equistore\u001b[38;5;241m.\u001b[39mtensor\u001b[38;5;241m.\u001b[39mTensorMap:\n\u001b[1;32m    225\u001b[0m     \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    226\u001b[0m \u001b[38;5;124;03m    Takes a TensorMap and creates a new TensorMap with the same data\u001b[39;00m\n\u001b[1;32m    227\u001b[0m \u001b[38;5;124;03m    values, but different Labels names for the:\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    261\u001b[0m \u001b[38;5;124;03m        TensorMap, but will new Label names.\u001b[39;00m\n\u001b[1;32m    262\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[1;32m    263\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m TensorMap(\n\u001b[1;32m    264\u001b[0m         keys\u001b[38;5;241m=\u001b[39mLabels(names\u001b[38;5;241m=\u001b[39mnew_names[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mkeys\u001b[39m\u001b[38;5;124m\"\u001b[39m], values\u001b[38;5;241m=\u001b[39mtensormap\u001b[38;5;241m.\u001b[39mkeys\u001b[38;5;241m.\u001b[39masarray()),\n\u001b[0;32m--> 265\u001b[0m         blocks\u001b[38;5;241m=\u001b[39m[rename_tensorblock(block, new_names) \u001b[38;5;28;01mfor\u001b[39;00m block \u001b[38;5;129;01min\u001b[39;00m tensormap\u001b[38;5;241m.\u001b[39mblocks()],\n\u001b[1;32m    266\u001b[0m     )\n",
      "File \u001b[0;32m~/Documents/phd/code/equi/equistore-examples/utils/rename_tensors.py:265\u001b[0m, in \u001b[0;36m<listcomp>\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m    222\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mrename_tensormap\u001b[39m(\n\u001b[1;32m    223\u001b[0m     tensormap: equistore\u001b[38;5;241m.\u001b[39mtensor\u001b[38;5;241m.\u001b[39mTensorMap, new_names: \u001b[38;5;28mdict\u001b[39m\n\u001b[1;32m    224\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m equistore\u001b[38;5;241m.\u001b[39mtensor\u001b[38;5;241m.\u001b[39mTensorMap:\n\u001b[1;32m    225\u001b[0m     \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    226\u001b[0m \u001b[38;5;124;03m    Takes a TensorMap and creates a new TensorMap with the same data\u001b[39;00m\n\u001b[1;32m    227\u001b[0m \u001b[38;5;124;03m    values, but different Labels names for the:\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    261\u001b[0m \u001b[38;5;124;03m        TensorMap, but will new Label names.\u001b[39;00m\n\u001b[1;32m    262\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[1;32m    263\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m TensorMap(\n\u001b[1;32m    264\u001b[0m         keys\u001b[38;5;241m=\u001b[39mLabels(names\u001b[38;5;241m=\u001b[39mnew_names[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mkeys\u001b[39m\u001b[38;5;124m\"\u001b[39m], values\u001b[38;5;241m=\u001b[39mtensormap\u001b[38;5;241m.\u001b[39mkeys\u001b[38;5;241m.\u001b[39masarray()),\n\u001b[0;32m--> 265\u001b[0m         blocks\u001b[38;5;241m=\u001b[39m[\u001b[43mrename_tensorblock\u001b[49m\u001b[43m(\u001b[49m\u001b[43mblock\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mnew_names\u001b[49m\u001b[43m)\u001b[49m \u001b[38;5;28;01mfor\u001b[39;00m block \u001b[38;5;129;01min\u001b[39;00m tensormap\u001b[38;5;241m.\u001b[39mblocks()],\n\u001b[1;32m    266\u001b[0m     )\n",
      "File \u001b[0;32m~/Documents/phd/code/equi/equistore-examples/utils/rename_tensors.py:157\u001b[0m, in \u001b[0;36mrename_tensorblock\u001b[0;34m(tensorblock, new_names)\u001b[0m\n\u001b[1;32m    155\u001b[0m props \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mlist\u001b[39m(new_names[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mproperties\u001b[39m\u001b[38;5;124m\"\u001b[39m])\n\u001b[1;32m    156\u001b[0m grad_props_list \u001b[38;5;241m=\u001b[39m [grad[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mproperties\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;28;01mfor\u001b[39;00m grad \u001b[38;5;129;01min\u001b[39;00m new_names[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mgradients\u001b[39m\u001b[38;5;124m\"\u001b[39m]\u001b[38;5;241m.\u001b[39mvalues()]\n\u001b[0;32m--> 157\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m np\u001b[38;5;241m.\u001b[39mall(\n\u001b[1;32m    158\u001b[0m     [np\u001b[38;5;241m.\u001b[39marray_equal(props, grad_props) \u001b[38;5;28;01mfor\u001b[39;00m grad_props \u001b[38;5;129;01min\u001b[39;00m grad_props_list]\n\u001b[1;32m    159\u001b[0m ), (\n\u001b[1;32m    160\u001b[0m     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mThe names of TensorBlock properties Labels should be the same as those for the Gradient TensorBlocks. \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    161\u001b[0m     \u001b[38;5;241m+\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mYou have passed \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    162\u001b[0m     \u001b[38;5;241m+\u001b[39m \u001b[38;5;28mstr\u001b[39m(props)\n\u001b[1;32m    163\u001b[0m     \u001b[38;5;241m+\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m as the names of the properties of the TensorBlock, and \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    164\u001b[0m     \u001b[38;5;241m+\u001b[39m \u001b[38;5;28mstr\u001b[39m(grad_props_list)\n\u001b[1;32m    165\u001b[0m     \u001b[38;5;241m+\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m as the names of the properties for the Gradient TensorBlocks.\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    166\u001b[0m )\n\u001b[1;32m    168\u001b[0m \u001b[38;5;66;03m# Where there are N number of components names for a given TensorBlock,\u001b[39;00m\n\u001b[1;32m    169\u001b[0m \u001b[38;5;66;03m# the final N components names of the associated Gradient TensorBlocks\u001b[39;00m\n\u001b[1;32m    170\u001b[0m \u001b[38;5;66;03m# should be equivalent. Check that the user hasn't tried to define different\u001b[39;00m\n\u001b[1;32m    171\u001b[0m \u001b[38;5;66;03m# Gradient TensorBlock and parent TensorBlock components names.\u001b[39;00m\n\u001b[1;32m    172\u001b[0m comps \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mlist\u001b[39m(new_names[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mcomponents\u001b[39m\u001b[38;5;124m\"\u001b[39m])\n",
      "\u001b[0;31mAssertionError\u001b[0m: The names of TensorBlock properties Labels should be the same as those for the Gradient TensorBlocks. You have passed ['radial_channel_n'] as the names of the properties of the TensorBlock, and [('n',), ('n',)] as the names of the properties for the Gradient TensorBlocks."
     ]
    }
   ],
   "source": [
    "new_names['properties'] = ('radial_channel_n',)\n",
    "\n",
    "new_tensormap = rename_tensormap(tensormap, new_names)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50ee2716",
   "metadata": {},
   "source": [
    "We would need to change all the names of the properties consistently, as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "45c98972",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'components': [('spherical_harmonics_m',)],\n",
      " 'gradients': {'cell': {'components': [('direction_1',),\n",
      "                                       ('direction_2',),\n",
      "                                       ('spherical_harmonics_m',)],\n",
      "                        'properties': ('radial_channel_n',),\n",
      "                        'samples': ('sample',)},\n",
      "               'positions': {'components': [('direction',),\n",
      "                                            ('spherical_harmonics_m',)],\n",
      "                             'properties': ('radial_channel_n',),\n",
      "                             'samples': ('sample', 'structure', 'atom')}},\n",
      " 'keys': ('l', 'species_center', 'species_neighbor'),\n",
      " 'properties': ('radial_channel_n',),\n",
      " 'samples': ('crystal', 'center')}\n"
     ]
    }
   ],
   "source": [
    "# Change keys, samples, and properties names\n",
    "new_names = {\n",
    " 'components': [('spherical_harmonics_m',)],\n",
    " 'gradients': {'cell': {'components': [('direction_1',),\n",
    "                                       ('direction_2',),\n",
    "                                       ('spherical_harmonics_m',)],\n",
    "                        'properties': ('radial_channel_n',),\n",
    "                        'samples': ('sample',)},\n",
    "               'positions': {'components': [('direction',),\n",
    "                                            ('spherical_harmonics_m',)],\n",
    "                             'properties': ('radial_channel_n',),\n",
    "                             'samples': ('sample', 'structure', 'atom')}},\n",
    " 'keys': ('l', 'species_center', 'species_neighbor'),\n",
    " 'properties': ('radial_channel_n',),\n",
    " 'samples': ('crystal', 'center')\n",
    "}\n",
    "\n",
    "# \"Rename\" (i.e. create a new) tensormap with the desired names\n",
    "new_tensormap = rename_tensormap(tensormap, new_names)\n",
    "\n",
    "pprint(get_tensormap_name_structure(new_tensormap))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a59fde4",
   "metadata": {},
   "source": [
    "# Renaming a ``TensorBlock``"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "609f8761",
   "metadata": {},
   "source": [
    "This can be done in a very similar way to renaming a ``TensorMap``, except the ``rename_tensors.rename_tensorblock()`` function is used instead. In fact, under-the-hood, ``rename_tensormap()`` works by iterating over blocks and calling ``rename_tensorblock()``."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "c0ed1aa2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'components': [('spherical_harmonics_m',)],\n",
      " 'gradients': {'cell': {'components': [('direction_1',),\n",
      "                                       ('direction_2',),\n",
      "                                       ('spherical_harmonics_m',)],\n",
      "                        'properties': ('n',),\n",
      "                        'samples': ('sample',)},\n",
      "               'positions': {'components': [('direction',),\n",
      "                                            ('spherical_harmonics_m',)],\n",
      "                             'properties': ('n',),\n",
      "                             'samples': ('sample', 'structure', 'atom')}},\n",
      " 'properties': ('n',),\n",
      " 'samples': ('structure', 'center')}\n"
     ]
    }
   ],
   "source": [
    "tensorblock = tensormap.block(0)\n",
    "\n",
    "pprint(get_tensorblock_name_structure(tensorblock))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3c998f6",
   "metadata": {},
   "source": [
    "Here we will change only the 'properties' names, from ``'n'`` to ``'radial_channel_n'``:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "fa2b307a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'components': [('spherical_harmonics_m',)],\n",
      " 'gradients': {'cell': {'components': [('direction_1',),\n",
      "                                       ('direction_2',),\n",
      "                                       ('spherical_harmonics_m',)],\n",
      "                        'properties': ('radial_channel_n',),\n",
      "                        'samples': ('sample',)},\n",
      "               'positions': {'components': [('direction',),\n",
      "                                            ('spherical_harmonics_m',)],\n",
      "                             'properties': ('radial_channel_n',),\n",
      "                             'samples': ('sample', 'structure', 'atom')}},\n",
      " 'properties': ('radial_channel_n',),\n",
      " 'samples': ('structure', 'center')}\n"
     ]
    }
   ],
   "source": [
    "# Change properties names\n",
    "new_names_tensorblock = {'components': [('spherical_harmonics_m',)],\n",
    "                         'gradients': {'cell': {'components': [('direction_1',),\n",
    "                                                               ('direction_2',),\n",
    "                                                               ('spherical_harmonics_m',)],\n",
    "                                                'properties': ('radial_channel_n',),\n",
    "                                                'samples': ('sample',)},\n",
    "                                       'positions': {'components': [('direction',),\n",
    "                                                                    ('spherical_harmonics_m',)],\n",
    "                                                     'properties': ('radial_channel_n',),\n",
    "                                                     'samples': ('sample', 'structure', 'atom')}},\n",
    "                         'properties': ('radial_channel_n',),\n",
    "                         'samples': ('structure', 'center')}\n",
    "\n",
    "# \"Rename\" (i.e. create a new) tensormap with the desired names\n",
    "new_tensorblock = rename_tensorblock(tensorblock, new_names_tensorblock)\n",
    "\n",
    "pprint(get_tensorblock_name_structure(new_tensorblock))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0c1d240",
   "metadata": {},
   "source": [
    "# Summary"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0024a398",
   "metadata": {},
   "source": [
    "To re-emphasize, this is a clunky but functional way of renaming ``TensorMap`` and ``TensorBlock`` objects. The Rust-side and more-permanent implementation will edit the existing objects, not create new ones, and have specific methods to rename separate ``Labels``, i.e. for samples, components, properties, gradient components, etc."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "cosmo",
   "language": "python",
   "name": "cosmo"
  },
  "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.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}