MP_tests.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "# import sys\n",
    "# sys.path.insert(0,'/home/nigam/librascal_cs/librascal/build/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ase.io import read, write\n",
    "import ase\n",
    "import json\n",
    "from tqdm.notebook import tqdm\n",
    "class tqdm_reusable:\n",
    "    def __init__(self, *args, **kwargs):\n",
    "        self._args = args\n",
    "        self._kwargs = kwargs\n",
    "\n",
    "    def __iter__(self):\n",
    "        return tqdm(*self._args, **self._kwargs).__iter__()\n",
    "    \n",
    "from copy import deepcopy\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from rascal.representations import SphericalExpansion, SphericalInvariants\n",
    "from rascal.utils import (get_radial_basis_covariance, get_radial_basis_pca, \n",
    "                          get_radial_basis_projections, get_optimal_radial_basis_hypers )\n",
    "from rascal.utils import radial_basis\n",
    "from rascal.utils import WignerDReal, ClebschGordanReal, spherical_expansion_reshape, lm_slice, real2complex_matrix, compute_lambda_soap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# a library of utilities and wrapper to compute pair features and manipulate a Hamiltonian-like target\n",
    "from ncnice import * "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Check Bispectrum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ncnice.representations import *#compute_rho3i_lambda"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "spherical_expansion_hypers = {\n",
    "    \"interaction_cutoff\": 4,\n",
    "    \"max_radial\": 8,\n",
    "    \"max_angular\": 6,\n",
    "    \"gaussian_sigma_constant\": 0.3,\n",
    "    \"gaussian_sigma_type\": \"Constant\",\n",
    "    \"cutoff_smooth_width\": 0.5,\n",
    "    \"radial_basis\": \"GTO\",\n",
    "}\n",
    "\n",
    "spex = SphericalExpansion(**spherical_expansion_hypers)\n",
    "mycg = ClebschGordanReal(spherical_expansion_hypers[\"max_angular\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "frames = read('data/ethanol-structures.xyz', ':')\n",
    "# frames = read('water/water_randomized_1000.xyz', ':')\n",
    "# for f in frames:\n",
    "#     f.cell=[100,100,100]\n",
    "#     f.positions+=50\n",
    "#     # use same specie for all atoms so we get a single projection matrix\n",
    "#     f.numbers = f.numbers*0+1    \n",
    "\n",
    "# spherical_expansion_hypers = get_optimal_radial_basis_hypers(spherical_expansion_hypers, frames, expanded_max_radial=12)\n",
    "\n",
    "\n",
    "# # ... and we replicate it \n",
    "# pm = spherical_expansion_hypers['optimization']['RadialDimReduction']['projection_matrices'][1] \n",
    "# spherical_expansion_hypers['optimization']['RadialDimReduction']['projection_matrices'] = { i: pm for i in range(1000) }\n",
    "spex = SphericalExpansion(**spherical_expansion_hypers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "frames = read('data/ethanol-structures.xyz', ':')\n",
    "# frames = read('/water_randomized_1000.xyz', ':')\n",
    "for f in frames:\n",
    "    f.cell=[100,100,100]\n",
    "    f.positions+=50\n",
    "rhoi = compute_rhoi(frames[:1], spex, spherical_expansion_hypers)*1e0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "rho2i_l, prho2i_l = compute_all_rho2i_lambda(rhoi, mycg, rho2i_pca=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 0 0 0\n",
      "l1, l2, k, l3, prho2, p, expected 0 0 0 0 1 1 1\n",
      "0 1 1 0\n",
      "l1, l2, k, l3, prho2, p, expected 0 1 1 1 1 1 1\n",
      "0 2 2 0\n",
      "l1, l2, k, l3, prho2, p, expected 0 2 2 2 1 1 1\n",
      "0 3 3 0\n",
      "l1, l2, k, l3, prho2, p, expected 0 3 3 3 1 1 1\n",
      "0 4 4 0\n",
      "l1, l2, k, l3, prho2, p, expected 0 4 4 4 1 1 1\n",
      "0 5 5 0\n",
      "l1, l2, k, l3, prho2, p, expected 0 5 5 5 1 1 1\n",
      "0 6 6 0\n",
      "l1, l2, k, l3, prho2, p, expected 0 6 6 6 1 1 1\n",
      "1 1 1 1\n",
      "l1, l2, k, l3, prho2, p, expected 1 1 1 1 -1 -1 -1\n",
      "1 1 2 1\n",
      "l1, l2, k, l3, prho2, p, expected 1 1 2 2 1 1 1\n",
      "1 2 2 2\n",
      "l1, l2, k, l3, prho2, p, expected 1 2 2 2 -1 -1 -1\n",
      "1 2 3 1\n",
      "l1, l2, k, l3, prho2, p, expected 1 2 3 3 1 1 1\n",
      "1 3 3 2\n",
      "l1, l2, k, l3, prho2, p, expected 1 3 3 3 -1 -1 -1\n",
      "1 3 4 1\n",
      "l1, l2, k, l3, prho2, p, expected 1 3 4 4 1 1 1\n",
      "1 4 4 2\n",
      "l1, l2, k, l3, prho2, p, expected 1 4 4 4 -1 -1 -1\n",
      "1 4 5 1\n",
      "l1, l2, k, l3, prho2, p, expected 1 4 5 5 1 1 1\n",
      "1 5 5 2\n",
      "l1, l2, k, l3, prho2, p, expected 1 5 5 5 -1 -1 -1\n",
      "1 5 6 1\n",
      "l1, l2, k, l3, prho2, p, expected 1 5 6 6 1 1 1\n",
      "1 6 6 2\n",
      "l1, l2, k, l3, prho2, p, expected 1 6 6 6 -1 -1 -1\n",
      "2 2 2 4\n",
      "l1, l2, k, l3, prho2, p, expected 2 2 2 2 1 1 1\n",
      "2 2 3 4\n",
      "l1, l2, k, l3, prho2, p, expected 2 2 3 3 -1 -1 -1\n",
      "2 2 4 4\n",
      "l1, l2, k, l3, prho2, p, expected 2 2 4 4 1 1 1\n",
      "2 3 3 5\n",
      "l1, l2, k, l3, prho2, p, expected 2 3 3 3 1 1 1\n",
      "2 3 4 5\n",
      "l1, l2, k, l3, prho2, p, expected 2 3 4 4 -1 -1 -1\n",
      "2 3 5 4\n",
      "l1, l2, k, l3, prho2, p, expected 2 3 5 5 1 1 1\n",
      "2 4 4 6\n",
      "l1, l2, k, l3, prho2, p, expected 2 4 4 4 1 1 1\n",
      "2 4 5 5\n",
      "l1, l2, k, l3, prho2, p, expected 2 4 5 5 -1 -1 -1\n",
      "2 4 6 3\n",
      "l1, l2, k, l3, prho2, p, expected 2 4 6 6 1 1 1\n",
      "2 5 5 6\n",
      "l1, l2, k, l3, prho2, p, expected 2 5 5 5 1 1 1\n",
      "2 5 6 4\n",
      "l1, l2, k, l3, prho2, p, expected 2 5 6 6 -1 -1 -1\n",
      "2 6 6 5\n",
      "l1, l2, k, l3, prho2, p, expected 2 6 6 6 1 1 1\n",
      "3 3 3 8\n",
      "l1, l2, k, l3, prho2, p, expected 3 3 3 3 -1 -1 -1\n",
      "3 3 4 9\n",
      "l1, l2, k, l3, prho2, p, expected 3 3 4 4 1 1 1\n",
      "3 3 5 8\n",
      "l1, l2, k, l3, prho2, p, expected 3 3 5 5 -1 -1 -1\n",
      "3 3 6 6\n",
      "l1, l2, k, l3, prho2, p, expected 3 3 6 6 1 1 1\n",
      "3 4 4 10\n",
      "l1, l2, k, l3, prho2, p, expected 3 4 4 4 -1 -1 -1\n",
      "3 4 5 9\n",
      "l1, l2, k, l3, prho2, p, expected 3 4 5 5 1 1 1\n",
      "3 4 6 7\n",
      "l1, l2, k, l3, prho2, p, expected 3 4 6 6 -1 -1 -1\n",
      "3 5 5 10\n",
      "l1, l2, k, l3, prho2, p, expected 3 5 5 5 -1 -1 -1\n",
      "3 5 6 8\n",
      "l1, l2, k, l3, prho2, p, expected 3 5 6 6 1 1 1\n",
      "3 6 6 9\n",
      "l1, l2, k, l3, prho2, p, expected 3 6 6 6 -1 -1 -1\n",
      "4 4 4 13\n",
      "l1, l2, k, l3, prho2, p, expected 4 4 4 4 1 1 1\n",
      "4 4 5 12\n",
      "l1, l2, k, l3, prho2, p, expected 4 4 5 5 -1 -1 -1\n",
      "4 4 6 10\n",
      "l1, l2, k, l3, prho2, p, expected 4 4 6 6 1 1 1\n",
      "4 5 5 13\n",
      "l1, l2, k, l3, prho2, p, expected 4 5 5 5 1 1 1\n",
      "4 5 6 11\n",
      "l1, l2, k, l3, prho2, p, expected 4 5 6 6 -1 -1 -1\n",
      "4 6 6 12\n",
      "l1, l2, k, l3, prho2, p, expected 4 6 6 6 1 1 1\n",
      "5 5 5 15\n",
      "l1, l2, k, l3, prho2, p, expected 5 5 5 5 -1 -1 -1\n",
      "5 5 6 13\n",
      "l1, l2, k, l3, prho2, p, expected 5 5 6 6 1 1 1\n",
      "5 6 6 14\n",
      "l1, l2, k, l3, prho2, p, expected 5 6 6 6 -1 -1 -1\n",
      "6 6 6 15\n",
      "l1, l2, k, l3, prho2, p, expected 6 6 6 6 1 1 1\n"
     ]
    }
   ],
   "source": [
    "rho3ilambda, prho3 = compute_rho3i_lambda(rho2i_l,rhoi, 0, mycg, prho2i_l)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "#MANUAL computation of bispectrum from librascal/nice_demo.ipynb\n",
    "\n",
    "selframe = frames[:1];         # frame used for the test\n",
    "feat_scaling = 1e0         \n",
    "feats = spex.transform(selframe).get_features(spex)\n",
    "ref_feats = feat_scaling*spherical_expansion_reshape(feats, **spherical_expansion_hypers)\n",
    "\n",
    "nice2_full = mycg.combine_nice(ref_feats, ref_feats)\n",
    "bisp_nice = np.zeros(ref_feats.shape[:3] + (spherical_expansion_hypers[\"max_angular\"]+1,) +\n",
    "                         ref_feats.shape[1:3] + (spherical_expansion_hypers[\"max_angular\"]+1,) +\n",
    "                         ref_feats.shape[1:3] + (spherical_expansion_hypers[\"max_angular\"]+1,) )\n",
    "for l in range(spherical_expansion_hypers[\"max_angular\"]+1):\n",
    "    # while we are at it, we also reorder the indices in a bispectrum-like way\n",
    "    bisp_nice[...,l] = mycg.combine_einsum(nice2_full[...,lm_slice(l)], \n",
    "                                         ref_feats[...,lm_slice(l)], \n",
    "                                         L=0, \n",
    "                                         combination_string=\"ianANlL,ibp->ianlANLbp\" )[...,0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 2.07141323e-07,  0.00000000e+00,  0.00000000e+00,\n",
       "         0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "         0.00000000e+00],\n",
       "       [ 0.00000000e+00, -9.63960119e-08,  0.00000000e+00,\n",
       "         0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "         0.00000000e+00],\n",
       "       [ 0.00000000e+00,  0.00000000e+00,  1.01818857e-07,\n",
       "         0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "         0.00000000e+00],\n",
       "       [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "        -1.42010029e-07,  0.00000000e+00,  0.00000000e+00,\n",
       "         0.00000000e+00],\n",
       "       [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "         0.00000000e+00,  5.18225025e-08,  0.00000000e+00,\n",
       "         0.00000000e+00],\n",
       "       [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "         0.00000000e+00,  0.00000000e+00, -1.82675677e-08,\n",
       "         0.00000000e+00],\n",
       "       [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "         0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "         1.19413640e-08]])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bisp_nice[0,0,1,:,0,2,:,0,2,0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 2.07141323e-07, -9.63960119e-08,  1.01818857e-07, -1.42010029e-07,\n",
       "        5.18225025e-08, -1.82675677e-08,  1.19413640e-08,  8.76031266e-25,\n",
       "       -9.11113234e-10, -4.24433872e-09, -6.58195522e-08, -4.43080257e-09,\n",
       "        4.54856670e-08, -4.04601961e-10, -3.35661883e-09,  1.60444254e-09,\n",
       "        5.80968589e-09,  8.38271362e-10, -1.84908681e-08, -1.25559488e-24,\n",
       "        1.86479439e-08,  8.29472454e-08, -3.67235530e-09, -2.30273780e-08,\n",
       "       -1.57694580e-08, -2.16674422e-09,  9.06422688e-09,  2.39273987e-09,\n",
       "        6.94772614e-11, -6.75178533e-09, -9.29187773e-26, -5.29462405e-08,\n",
       "       -2.88929594e-25,  2.36051061e-08, -7.20717436e-10,  1.38368165e-08,\n",
       "       -1.25046418e-09, -5.66993837e-10, -7.33817985e-09, -1.00277126e-10,\n",
       "        1.41444235e-08, -1.94206029e-25, -7.20039030e-09, -2.61445379e-09,\n",
       "       -1.50203059e-11,  4.24586358e-09,  4.46966200e-26,  1.74992743e-09,\n",
       "        2.61066651e-10, -3.33186302e-09])"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rho3ilambda[0,0,1,0,2,0,2,:,0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "they are the **same** modulo noise. because we dont account for sorting l's in the manual computation and the storage format is inefficient, the shapes of the matrices are different"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# rho2iP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_rho1ijp_species(rho1ijp, frame, cg): \n",
    "    \"\"\"groups neighbors j by species\"\"\"\n",
    "    species = frame.numbers\n",
    "    nspecies = len(set(species))\n",
    "    shape = rho1ijp.shape[:2]+(nspecies,) +rho1ijp.shape[2:]\n",
    "    rho1ijp_spec = np.zeros(shape)\n",
    "    \n",
    "    for ispe, spe in enumerate(sorted(set(species))):\n",
    "        idx_spe = np.where(species==spe)[0]\n",
    "#         print(spe, idx_spe)\n",
    "        for j in range(rho1ijp.shape[1]):\n",
    "            if j in idx_spe:\n",
    "                rho1ijp_spec[:,j,ispe, ...]= rho1ijp[:,j,...]\n",
    "    return rho1ijp_spec"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "hypers_ij = deepcopy(spherical_expansion_hypers)\n",
    "hypers_ij[\"expansion_by_species_method\"] = \"structure wise\"\n",
    "spex_ij = SphericalExpansion(**hypers_ij)\n",
    "\n",
    "frame = frames[0]\n",
    "fgij = compute_gij(frame, spex_ij, hypers_ij)\n",
    "rhoi = compute_rhoi(frame, spex, spherical_expansion_hypers)\n",
    "rho1ijp, prhoijp = compute_rho1ijp_lambda(rhoi, fgij,0, mycg)\n",
    "rho1ij, prhoij = compute_rho1ij_lambda(rhoi, fgij,0, mycg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "rho1ijp_l, prho1ijp_l = compute_all_rho1ijp_lambda(rhoi, fgij, mycg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_rho1ip_lambda(rhoi, rho1ijp_l, fgij, frame, L, cg, prho1ijp_l):\n",
    "    \"\"\"compute combined rhoi and rhoi-MP\"\"\"\n",
    "    rho1ijp_spec = {}\n",
    "    prhoijp_spec = {}\n",
    "    rho1ip_lambda = {}\n",
    "    lmax = int(np.sqrt(rhoi.shape[-1])) -1\n",
    "    for l in range(lmax+1):\n",
    "        rho1ijp_spec[l] = compute_rho1ijp_species(rho1ijp_l[l], frame, cg)\n",
    "        rho1ip_lambda[l] = np.sum(rho1ijp_spec[l], axis=1)  # MP - atom density \n",
    "    \n",
    "    # do the usual gig of angular coupling \n",
    "    nl=0\n",
    "    for l1 in range(lmax+1):\n",
    "        for l2 in range(lmax+1):\n",
    "            if abs(l2 - l1) > L or l2 + l1 < L:\n",
    "                continue\n",
    "            nl += 1*rho1ip_lambda[l2].shape[4]\n",
    "# multiplication above to account for (l1,l2,l3) combinations. rho1ip_lambda[L].shape[4] accounts for l1,l2\n",
    "#  combinations of g(rij) and rhoj\n",
    "    shape = rhoi.shape[:3] + rho1ip_lambda[L].shape[1:4] +(nl, 2*L+1,)\n",
    "    rho_combined_lam = np.zeros(shape)\n",
    "    prho_combined_lam = np.ones(nl, dtype = int)*(1-2*(L%2))\n",
    "    \n",
    "    il=0\n",
    "    for l1 in range(lmax+1):\n",
    "        for l2 in range(lmax+1):\n",
    "            if abs(l2 - l1) > L or l2 + l1 < L:\n",
    "                continue\n",
    "            rho_combined_lam[...,il:il+rho1ip_lambda[l2].shape[4],:] = cg.combine_einsum(rhoi[...,lm_slice(l1)], rho1ip_lambda[l2],\n",
    "                                            L, combination_string=\"ian,ibNMl->ianbNMl\")\n",
    "\n",
    "            prho_combined_lam[il:il+rho1ip_lambda[l2].shape[4]] = (1-2*(l1%2)) * prho1ijp_l[l2]\n",
    "            il+=rho1ip_lambda[l2].shape[4]\n",
    "\n",
    "    return rho_combined_lam, prho_combined_lam"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "rho_combined_lam, prho_combined_lam = compute_rho1ip_lambda(rhoi, rho1ijp_l, fgij, frame, 1, mycg, prho1ijp_l)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(9, 3, 8, 3, 24, 8, 483, 3)"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rho_combined_lam.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_hamiltonian_representations(frames, orbs, hypers, lmax, nu, cg, scale=1,\n",
    "                     select_feats = None, half_hete = True, mp_feats = False,\n",
    "                     rhoi_pca = None, rho2i_pca = None,\n",
    "                     rhoij_rho2i_pca = None, rhoij_pca = None,\n",
    "                     verbose = False\n",
    "                     ):\n",
    "    \"\"\"\n",
    "        Computes the full set of features needed to learn matrix elements up to lmax.\n",
    "        Options are fluid, but here are some that need an explanation:\n",
    "\n",
    "        select_feats = dict(type=[\"diag\", \"offd_m\", \"offd_p\", \"hete\"], block = ('el1', ['el2',] L, pi) )\n",
    "        does the minimal amount of calculation to evaluate the selected block. other terms might be computed as well if they come for free.\n",
    "    \"\"\"\n",
    "\n",
    "    spex = SphericalExpansion(**hypers)\n",
    "    rhoi = compute_rhoi(frames, spex, hypers)\n",
    "\n",
    "    # compresses further the spherical expansion features across species\n",
    "    if rhoi_pca is not None:\n",
    "        rhoi = apply_rhoi_pca(rhoi, rhoi_pca)\n",
    "\n",
    "    # makes sure that the spex used for the pair terms uses adaptive species\n",
    "    hypers_ij = deepcopy(hypers)\n",
    "    hypers_ij[\"expansion_by_species_method\"] = \"structure wise\"\n",
    "    spex_ij = SphericalExpansion(**hypers_ij)\n",
    "\n",
    "    tnat = 0\n",
    "    els = list(orbs.keys())\n",
    "    nel = len(els)\n",
    "    # prepare storage\n",
    "    elL = list(itertools.product(els,range(lmax+1),[-1,1]))\n",
    "    hetL = [ (els[i1], els[i2], L, pi) for i1 in range(nel) for i2 in range((i1+1 if half_hete else 0), nel) for L in range(lmax+1) for pi in [-1,1] ]\n",
    "    feats = dict(diag = { L: [] for L in elL },\n",
    "                 offd_p = { L: [] for L in elL },\n",
    "                 offd_m = { L: [] for L in elL },\n",
    "                 hete =   { L: [] for L in hetL },)\n",
    "\n",
    "    if rhoij_rho2i_pca is None and rho2i_pca is not None:\n",
    "        rhoij_rho2i_pca = rho2i_pca\n",
    "\n",
    "    #before = tracemalloc.take_snapshot()\n",
    "    for f in frames:\n",
    "        fnat = len(f.numbers)\n",
    "        frhoi = rhoi[tnat:tnat+fnat]*scale\n",
    "        fgij = compute_gij(f, spex_ij, hypers_ij)*scale\n",
    "\n",
    "        if (select_feats is None or select_feats[\"type\"]!=\"diag\") and nu == 2:\n",
    "            rhonui, prhonui = compute_all_rho2i_lambda(frhoi, cg, rhoij_rho2i_pca)\n",
    "        else:\n",
    "            rhonui, prhonui = frhoi, None\n",
    "\n",
    "        for L in range(lmax+1):\n",
    "            if select_feats is not None and L>0 and select_feats[\"block\"][-2] != L:\n",
    "                continue\n",
    "\n",
    "            if nu==0:\n",
    "                lrhonui, lprhonui = np.ones((fnat, 1, 2*L+1)), np.ones((1))\n",
    "            elif nu==1:\n",
    "                lrhonui, lprhonui = compute_rho1i_lambda(frhoi, L, cg)\n",
    "            else:\n",
    "                if mp_feats:\n",
    "                    frho1ijp_l, fprho1ijp_l = compute_all_rho1ijp_lambda(frhoi, fgij, mycg)\n",
    "                    lrhonui, lprhonui = compute_rho1ip_lambda(frhoi, frho1ijp_l, fgij, frame, L, cg, fprho1ijp_l)\n",
    "#                     lrhonui, lprhonui = compute_rho1ip_lambda(frhoi, L, cg)\n",
    "                else:\n",
    "                    frho2i_l,fprho2i_l = compute_all_rho2i_lambda(frhoi, cg, rho2i_pca=None)\n",
    "                    lrhonui, lprhonui = compute_rho3i_lambda(frho2i_l, frhoi, L, cg,fprho2i_l )\n",
    "                    #lrhonui, lprhonui = compute_rho2i_lambda(frhoi, L, cg)\n",
    "                    #if rho2i_pca is not None:\n",
    "                    #     lrhonui, lprhonui = apply_rho2i_pca(lrhonui, lprhonui, rho2i_pca)\n",
    "\n",
    "            if select_feats is None or select_feats[\"type\"]!=\"diag\":\n",
    "                if nu==0:\n",
    "                    lrhoij, prhoij = compute_rho0ij_lambda(rhonui, fgij, L, cg, prhonui)\n",
    "                elif nu==1:\n",
    "                    if mp_feats:\n",
    "                        lrhoij, prhoij = compute_rho1ijp_lambda(rhonui, fgij, L, cg, prhonui)\n",
    "                    else:\n",
    "                        lrhoij, prhoij = compute_rho1ij_lambda(rhonui, fgij, L, cg, prhonui)\n",
    "                else:\n",
    "                    lrhoij, prhoij = compute_rho2ij_lambda(rhonui, fgij, L, cg, prhonui)\n",
    "                if rhoij_pca is not None:\n",
    "                    lrhoij, prhoij = apply_rhoij_pca(lrhoij, prhoij, rhoij_pca)\n",
    "\n",
    "            for i, el in enumerate(els):\n",
    "                iel = np.where(f.symbols==el)[0]\n",
    "                if len(iel) == 0:\n",
    "                    continue\n",
    "                if select_feats is not None and el != select_feats[\"block\"][0]:\n",
    "                    continue\n",
    "\n",
    "                for pi in [-1,1]:\n",
    "                    wherepi = np.where(lprhonui==pi)[0]\n",
    "                    if len(wherepi)==0:\n",
    "                        # add a vector of zeros\n",
    "                        feats['diag'][(el, L, pi)].append(np.zeros(shape=(len(iel), 1, 2*L+1)))\n",
    "                        continue\n",
    "                    feats['diag'][(el, L, pi)].append(lrhonui[...,wherepi,:][iel].reshape((len(iel), -1, 2*L+1) ) )\n",
    "\n",
    "                if select_feats is not None and select_feats[\"type\"]==\"diag\":\n",
    "                    continue\n",
    "\n",
    "                triu = np.triu_indices(len(iel), 1)\n",
    "                ij_up = (iel[triu[0]],iel[triu[1]]) # ij indices, i>j\n",
    "                ij_lw = (ij_up[1], ij_up[0]) # ij indices, i<j\n",
    "                lrhoij_p = (lrhoij[ij_up] + lrhoij[ij_lw])/np.sqrt(2)\n",
    "                lrhoij_m = (lrhoij[ij_up] - lrhoij[ij_lw])/np.sqrt(2)\n",
    "                for pi in [-1,1]:\n",
    "                    if len(ij_up[0])==0:\n",
    "                        continue\n",
    "                    wherepi = np.where(prhoij==pi)[0];\n",
    "                    if len(wherepi)==0:\n",
    "                        feats['offd_p'][(el, L, pi)].append( np.zeros((lrhoij_p.shape[0], 1, 2*L+1)) )\n",
    "                        feats['offd_m'][(el, L, pi)].append( np.zeros((lrhoij_p.shape[0], 1, 2*L+1)) )\n",
    "                        continue\n",
    "                    feats['offd_p'][(el, L, pi)].append(lrhoij_p[...,wherepi,:].reshape(lrhoij_p.shape[0], -1, 2*L+1))\n",
    "                    feats['offd_m'][(el, L, pi)].append(lrhoij_m[...,wherepi,:].reshape(lrhoij_m.shape[0], -1, 2*L+1))\n",
    "\n",
    "                if select_feats is not None and select_feats[\"type\"]!=\"hete\":\n",
    "                    continue\n",
    "                for elb in els[i+1:]:\n",
    "                    ielb = np.where(f.symbols==elb)[0]\n",
    "                    if len(ielb) == 0:\n",
    "                        continue\n",
    "                    if select_feats is not None and elb != select_feats[\"block\"][1]:\n",
    "                        continue\n",
    "\n",
    "                    # combines rho_ij and rho_ji\n",
    "                    lrhoij_het = lrhoij[iel][:,ielb]\n",
    "                    lrhoij_het_rev = np.swapaxes(lrhoij[ielb][:,iel],1,0)\n",
    "                    # make a copy and not a slice, so we keep better track\n",
    "                    for pi in [-1,1]:\n",
    "                        wherepi = np.where(prhoij==pi)[0];\n",
    "                        if len(wherepi)==0:\n",
    "                            feats['hete'][(el, elb, L, pi)].append(np.zeros((lrhoij_het.shape[0]*lrhoij_het.shape[1],1,2*L+1)))\n",
    "                            continue\n",
    "                        lrhoij_het_pi = lrhoij_het[...,wherepi,:]\n",
    "                        lrhoij_het_rev_pi = lrhoij_het_rev[...,wherepi,:]\n",
    "                        feats['hete'][(el, elb, L, pi)].append(\n",
    "                            np.concatenate([\n",
    "                            lrhoij_het_pi.reshape(\n",
    "                                (lrhoij_het.shape[0]*lrhoij_het.shape[1],-1,2*L+1) )\n",
    "        \n",
    "                            ,\n",
    "                            lrhoij_het_rev_pi.reshape(\n",
    "                                (lrhoij_het_rev.shape[0]*lrhoij_het_rev.shape[1],-1,2*L+1) )\n",
    "                            ], axis=-2)\n",
    "                        )\n",
    "                    #del(lrhoij_het)\n",
    "                #del(lrhoij_p, lrhoij_m)\n",
    "            #del(lrhoij, lrho2)\n",
    "        tnat+=fnat\n",
    "\n",
    "    # cleans up combining frames blocks into single vectors - splitting also odd and even blocks\n",
    "    for k in feats.keys():\n",
    "        for b in list(feats[k].keys()):\n",
    "            if len(feats[k][b]) == 0:\n",
    "                continue\n",
    "            block = np.vstack(feats[k][b])\n",
    "            feats[k].pop(b)\n",
    "            if len(block) == 0:\n",
    "                continue\n",
    "\n",
    "            feats[k][b] = block.reshape((block.shape[0], -1, 1+2*b[-2]))\n",
    "\n",
    "    return feats\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Regression Test "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[13  8 16  3 15 12  0 10  7 11]\n",
      "Calling all representation subroutines (no PCA)\n"
     ]
    }
   ],
   "source": [
    "nframes = 50\n",
    "frames = read('data/ethanol-structures.xyz',':')[:nframes]\n",
    "for f in frames:\n",
    "    f.cell=[100,100,100]\n",
    "    f.positions+=50\n",
    "    # use same specie for all atoms so we get a single projection matrix,\n",
    "    # which we can apply throughout. a bit less efficient but much more practical\n",
    "    f.numbers = f.numbers*0+1\n",
    "spherical_expansion_hypers = get_optimal_radial_basis_hypers(spherical_expansion_hypers, frames, expanded_max_radial=16)\n",
    "# ... and we replicate it\n",
    "pm = spherical_expansion_hypers['optimization']['RadialDimReduction']['projection_matrices'][1]\n",
    "spherical_expansion_hypers['optimization']['RadialDimReduction']['projection_matrices'] = { i: pm for i in range(99) }\n",
    "spex = SphericalExpansion(**spherical_expansion_hypers)\n",
    "hypers_ij = deepcopy(spherical_expansion_hypers)\n",
    "hypers_ij[\"expansion_by_species_method\"] = \"structure wise\"\n",
    "spex_gij = SphericalExpansion(**hypers_ij)\n",
    "\n",
    "\n",
    "orbs = json.load(open('data/ethanol-saph-orbs.json', \"r\"))\n",
    "frames = read('data/ethanol-structures.xyz', ':')[:nframes]\n",
    "ofocks = np.load('data/ethanol-saph-ofock.npy', allow_pickle=True)[:nframes]\n",
    "\n",
    "# hamiltonian to block coupling\n",
    "ofock_blocks, slices_idx = matrix_list_to_blocks(ofocks, frames, orbs, mycg)\n",
    "\n",
    "# training settings\n",
    "train_fraction = 0.5\n",
    "itrain = np.arange(len(frames))\n",
    "np.random.seed(12345)\n",
    "np.random.shuffle(itrain)\n",
    "ntrain = int(len(itrain)*train_fraction)\n",
    "itest = itrain[ntrain:]; itrain=itrain[:ntrain]\n",
    "train_slices = get_block_idx(itrain, slices_idx)\n",
    "print(itrain)\n",
    "\n",
    "FR = FockRegression(orbs, alpha = 1e-6, #alphas=np.geomspace(1e-8, 1e4, 7),\n",
    "                    fit_intercept=\"auto\")\n",
    "\n",
    "for f in frames:\n",
    "    f.cell=[100,100,100]\n",
    "    f.positions+=50\n",
    "\n",
    "print(\"Calling all representation subroutines (no PCA)\")\n",
    "rhoi = compute_rhoi(frames[0], spex, spherical_expansion_hypers)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "import itertools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "features:   0%|          | 0/20 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "features:   0%|          | 0/20 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e7a587c46b3c4820961db4743354bf4e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/7 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "diag (2, 0, 2, 0) ('O', 0, 1) 0 2.7753608473980735\n",
      "diag (2, 0, 2, 1) ('O', 1, 1) 1 0.4462237741637202\n",
      "diag (2, 1, 2, 1) ('O', 0, 1) 0 1.8107058087600478\n",
      "diag (2, 1, 2, 1) ('O', 2, 1) 2 0.7383770695421117\n",
      "diag (4, 0, 4, 0) ('C', 0, 1) 0 1.5047174907273053\n",
      "diag (4, 0, 4, 1) ('C', 1, 1) 1 0.17318842920340205\n",
      "diag (4, 1, 4, 1) ('C', 0, 1) 0 0.6245433127456875\n",
      "diag (4, 1, 4, 1) ('C', 2, 1) 2 0.1997065011012175\n",
      "diag (5, 0, 5, 0) ('H', 0, 1) 0 2.3634310716472076\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9803ca8bcde44e8ba2575a7b9ae97e55",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/4 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "offd_p (4, 0, 4, 0) ('C', 0, 1) 0 0.9407257978566762\n",
      "offd_p (4, 0, 4, 1) ('C', 1, 1) 1 0.06293572470407495\n",
      "offd_p (4, 1, 4, 1) ('C', 0, 1) 0 0.4455990241628668\n",
      "offd_p (4, 1, 4, 1) ('C', 2, 1) 2 1.4819606168029316\n",
      "offd_p (5, 0, 5, 0) ('H', 0, 1) 0 0.5849021565828035\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "741aa49a448745c8bf3a959934b769c1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/4 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "offd_m (4, 0, 4, 1) ('C', 1, 1) 1 1.1480843837102424\n",
      "offd_m (4, 1, 4, 1) ('C', 1, -1) 1 0.10865191868706488\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3de36f895f164b2abcd2431fe0a93e7f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/8 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "hete (2, 0, 4, 0) ('O', 'C', 0, 1) 0 0.772874429363242\n",
      "hete (2, 0, 4, 1) ('O', 'C', 1, 1) 1 1.1223810787455462\n",
      "hete (2, 0, 5, 0) ('O', 'H', 0, 1) 0 1.436981705589576\n",
      "hete (2, 1, 4, 0) ('O', 'C', 1, 1) 1 0.8719020699217738\n",
      "hete (2, 1, 4, 1) ('O', 'C', 0, 1) 0 0.2641701507043727\n",
      "hete (2, 1, 4, 1) ('O', 'C', 1, -1) 1 0.06980991794986062\n",
      "hete (2, 1, 4, 1) ('O', 'C', 2, 1) 2 1.255623878799006\n",
      "hete (2, 1, 5, 0) ('O', 'H', 1, 1) 1 1.2441809297317046\n",
      "hete (4, 0, 5, 0) ('C', 'H', 0, 1) 0 1.9900322823546825\n",
      "hete (4, 1, 5, 0) ('C', 'H', 1, 1) 1 2.4862517743722115\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5b16433953774acc8b5fcf9a426c9799",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/7 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ebc75e9867b642b89e0a97ecc20e7431",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/4 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5a40c5a2fc074ea49666e43dbfd0367f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/4 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "48507ebf51c74445a90c5fa981bd7488",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/8 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model size:  26\n",
      "Train RMSE:  0.012053177503949585\n",
      "Test RMSE:  0.05175830885752685\n"
     ]
    }
   ],
   "source": [
    "feats_nu1 = compute_hamiltonian_representations(tqdm_reusable(frames, desc=\"features\", leave=False),\n",
    "                        orbs, spherical_expansion_hypers, 2, nu=1, cg=mycg, scale=1e3, mp_feats = True)\n",
    "\n",
    "FR.fit(feats_nu1, ofock_blocks, train_slices, progress=tqdm)\n",
    "pred_blocks = FR.predict(feats_nu1, progress=tqdm)\n",
    "pred_ofocks = blocks_to_matrix_list(pred_blocks, frames, slices_idx, orbs, mycg)\n",
    "\n",
    "mse_train = 0\n",
    "for i in itrain:\n",
    "    mse_train += np.sum((pred_ofocks[i] - ofocks[i])**2)/len(ofocks[i])/len(itrain)\n",
    "\n",
    "mse_test = 0\n",
    "for i in itest:\n",
    "    mse_test += np.sum((pred_ofocks[i] - ofocks[i])**2)/len(ofocks[i])/len(itest)\n",
    "\n",
    "print(\"Model size: \", len(FR.cv_stats_))\n",
    "print(\"Train RMSE: \", np.sqrt(mse_train))\n",
    "print(\"Test RMSE: \", np.sqrt(mse_test))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "('O', 2, 1)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "features:   0%|          | 0/20 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "features:   0%|          | 0/20 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "('C', 0, 1)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "features:   0%|          | 0/20 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "features:   0%|          | 0/20 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "('H', 0, 1)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "features:   0%|          | 0/20 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "features:   0%|          | 0/20 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "('O', 1, 1)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "features:   0%|          | 0/20 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "features:   0%|          | 0/20 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "('O', 0, 1)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "features:   0%|          | 0/20 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "features:   0%|          | 0/20 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "('C', 2, 1)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "features:   0%|          | 0/20 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "features:   0%|          | 0/20 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "('C', 1, 1)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "features:   0%|          | 0/20 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "features:   0%|          | 0/20 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "blocks_c = ofock_blocks \n",
    "tblock = 'diag'\n",
    "sel_type  = list(blocks_c[tblock].keys())\n",
    "fblock=[]\n",
    "for i in range(len(sel_type)):\n",
    "    kblock = sel_type[i]\n",
    "    lblock = list(blocks_c['diag'][sel_type[i]].keys())\n",
    "    for l in lblock:\n",
    "        fblock.append(block_to_feat_index(tblock, kblock, l, orbs))\n",
    "\n",
    "fblock = list(set(fblock))\n",
    "for j in fblock:\n",
    "    print(j)\n",
    "    \n",
    "    feats_nu1['diag'][j] = compute_hamiltonian_representations(tqdm_reusable(frames, desc=\"features\", leave=False),\n",
    "                        orbs, spherical_expansion_hypers, 2, nu=2, cg=mycg, scale=1e3, mp_feats = False,  select_feats =dict(block=j, type=tblock))[tblock][j]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "090139f79b7d4ff2bbca87ae0fb632bc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/7 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "diag (2, 0, 2, 0) ('O', 0, 1) 0 2.7753608473980735\n",
      "diag (2, 0, 2, 1) ('O', 1, 1) 1 0.4462237741637202\n",
      "diag (2, 1, 2, 1) ('O', 0, 1) 0 1.8107058087600478\n",
      "diag (2, 1, 2, 1) ('O', 2, 1) 2 0.7383770695421117\n",
      "diag (4, 0, 4, 0) ('C', 0, 1) 0 1.5047174907273053\n",
      "diag (4, 0, 4, 1) ('C', 1, 1) 1 0.17318842920340205\n",
      "diag (4, 1, 4, 1) ('C', 0, 1) 0 0.6245433127456875\n",
      "diag (4, 1, 4, 1) ('C', 2, 1) 2 0.1997065011012175\n",
      "diag (5, 0, 5, 0) ('H', 0, 1) 0 2.3634310716472076\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9d6c93fbecc348e89d3ea1d839b0ced9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/4 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "offd_p (4, 0, 4, 0) ('C', 0, 1) 0 0.9407257978566762\n",
      "offd_p (4, 0, 4, 1) ('C', 1, 1) 1 0.06293572470407495\n",
      "offd_p (4, 1, 4, 1) ('C', 0, 1) 0 0.4455990241628668\n",
      "offd_p (4, 1, 4, 1) ('C', 2, 1) 2 1.4819606168029316\n",
      "offd_p (5, 0, 5, 0) ('H', 0, 1) 0 0.5849021565828035\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f87abeb0748346c5ab3f0b82e4471668",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/4 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "offd_m (4, 0, 4, 1) ('C', 1, 1) 1 1.1480843837102424\n",
      "offd_m (4, 1, 4, 1) ('C', 1, -1) 1 0.10865191868706488\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6c53db2706544d6b81ddc284f5c5845b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/8 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "hete (2, 0, 4, 0) ('O', 'C', 0, 1) 0 0.772874429363242\n",
      "hete (2, 0, 4, 1) ('O', 'C', 1, 1) 1 1.1223810787455462\n",
      "hete (2, 0, 5, 0) ('O', 'H', 0, 1) 0 1.436981705589576\n",
      "hete (2, 1, 4, 0) ('O', 'C', 1, 1) 1 0.8719020699217738\n",
      "hete (2, 1, 4, 1) ('O', 'C', 0, 1) 0 0.2641701507043727\n",
      "hete (2, 1, 4, 1) ('O', 'C', 1, -1) 1 0.06980991794986062\n",
      "hete (2, 1, 4, 1) ('O', 'C', 2, 1) 2 1.255623878799006\n",
      "hete (2, 1, 5, 0) ('O', 'H', 1, 1) 1 1.2441809297317046\n",
      "hete (4, 0, 5, 0) ('C', 'H', 0, 1) 0 1.9900322823546825\n",
      "hete (4, 1, 5, 0) ('C', 'H', 1, 1) 1 2.4862517743722115\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7f5e1ec22279463b94295a9bb1c7293a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/7 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "915bfed2012548b9aeae9892912d6e2c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/4 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "eb678ff1b6d6470cb0d7100e824fe5de",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/4 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3a93f63803564908b46569a3d266a379",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/8 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model size:  26\n",
      "Train RMSE:  0.005117145029683599\n",
      "Test RMSE:  0.09496571883986546\n"
     ]
    }
   ],
   "source": [
    "FR.fit(feats_nu1, ofock_blocks, train_slices, progress=tqdm)\n",
    "pred_blocks = FR.predict(feats_nu1, progress=tqdm)\n",
    "pred_ofocks = blocks_to_matrix_list(pred_blocks, frames, slices_idx, orbs, mycg)\n",
    "\n",
    "mse_train = 0\n",
    "for i in itrain:\n",
    "    mse_train += np.sum((pred_ofocks[i] - ofocks[i])**2)/len(ofocks[i])/len(itrain)\n",
    "\n",
    "mse_test = 0\n",
    "for i in itest:\n",
    "    mse_test += np.sum((pred_ofocks[i] - ofocks[i])**2)/len(ofocks[i])/len(itest)\n",
    "\n",
    "print(\"Model size: \", len(FR.cv_stats_))\n",
    "print(\"Train RMSE: \", np.sqrt(mse_train))\n",
    "print(\"Test RMSE: \", np.sqrt(mse_test))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "('O', 2, 1)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "features:   0%|          | 0/20 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "345353eda8dc43148c3b4ec5653ccc3f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "features:   0%|          | 0/20 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "ERROR:root:Internal Python error in the inspect module.\n",
      "Below is the traceback from this internal error.\n",
      "\n",
      "\n",
      "KeyboardInterrupt\n",
      "\n"
     ]
    }
   ],
   "source": [
    "blocks_c = ofock_blocks \n",
    "tblock = 'diag'\n",
    "sel_type  = list(blocks_c[tblock].keys())\n",
    "fblock=[]\n",
    "for i in range(len(sel_type)):\n",
    "    kblock = sel_type[i]\n",
    "    lblock = list(blocks_c['diag'][sel_type[i]].keys())\n",
    "    for l in lblock:\n",
    "        fblock.append(block_to_feat_index(tblock, kblock, l, orbs))\n",
    "\n",
    "fblock = list(set(fblock))\n",
    "for j in fblock:\n",
    "    print(j)\n",
    "    \n",
    "    feats_nu1['diag'][j] = compute_hamiltonian_representations(tqdm_reusable(frames, desc=\"features\", leave=False),\n",
    "                        orbs, spherical_expansion_hypers, 2, nu=2, cg=mycg, scale=1e3, mp_feats = True,  select_feats =dict(block=j, type=tblock))[tblock][j]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "FR.fit(feats_nu1, ofock_blocks, train_slices, progress=tqdm)\n",
    "pred_blocks = FR.predict(feats_nu1, progress=tqdm)\n",
    "pred_ofocks = blocks_to_matrix_list(pred_blocks, frames, slices_idx, orbs, mycg)\n",
    "\n",
    "mse_train = 0\n",
    "for i in itrain:\n",
    "    mse_train += np.sum((pred_ofocks[i] - ofocks[i])**2)/len(ofocks[i])/len(itrain)\n",
    "\n",
    "mse_test = 0\n",
    "for i in itest:\n",
    "    mse_test += np.sum((pred_ofocks[i] - ofocks[i])**2)/len(ofocks[i])/len(itest)\n",
    "\n",
    "print(\"Model size: \", len(FR.cv_stats_))\n",
    "print(\"Train RMSE: \", np.sqrt(mse_train))\n",
    "print(\"Test RMSE: \", np.sqrt(mse_test))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# Compute bispectrum as sum over rho2ij lambda - NOT RIGHT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "hypers_ij = deepcopy(spherical_expansion_hypers)\n",
    "# hypers_ij[\"expansion_by_species_method\"] = \"structure wise\"\n",
    "spex_ij = SphericalExpansion(**hypers_ij)\n",
    "scale=1e3\n",
    "for f in frames[:1]:\n",
    "    fgij = compute_gij(f, spex_ij, hypers_ij)\n",
    "    rhoi = compute_rhoi(f, spex, spherical_expansion_hypers)\n",
    "    rho2i_l, prho2i_l = compute_all_rho2i_lambda(rhoi, mycg, rho2i_pca=None)\n",
    "    rho2ij, prho2ij = compute_rho2ij_lambda(rho2i_l, fgij, 0, mycg, prho2i_l)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "shape = (rho2ij.shape[0],\n",
    "         rhoi.shape[1],rho2ij.shape[2], \n",
    "         rho2ij.shape[3], rho2ij.shape[4], \n",
    "         rho2ij.shape[5], rho2ij.shape[6],\n",
    "         rho2ij.shape[7], rho2ij.shape[8]\n",
    "        )\n",
    "bisp_gij = np.zeros(shape)\n",
    "# parity = np.ones(nl, dtype = int)*(1-2*(L%2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "rho2ij.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "for f in frames[:1]:\n",
    "    fgij = compute_gij(f, spex_ij, hypers_ij)\n",
    "    rhoi = compute_rhoi(f, spex, spherical_expansion_hypers)\n",
    "    rho2i_l, prho2i_l = compute_all_rho2i_lambda(rhoi, mycg, rho2i_pca=None)\n",
    "    rho2ij, prho2ij = compute_rho2ij_lambda(rho2i_l, fgij, 0, mycg, prho2i_l)\n",
    "    shape = (rho2ij.shape[0],\n",
    "         rhoi.shape[1],rho2ij.shape[2], \n",
    "         rho2ij.shape[3], rho2ij.shape[4], \n",
    "         rho2ij.shape[5], rho2ij.shape[6],\n",
    "         rho2ij.shape[7], rho2ij.shape[8]\n",
    "        )\n",
    "    bisp_gij = np.zeros(shape)\n",
    "#     parity = np.ones(nl, dtype = int)*(1-2*(L%2))\n",
    "    species = f.numbers\n",
    "    nspecies = rhoi.shape[1]\n",
    "    for ispe, spe in enumerate(set(species)):\n",
    "        idx_spe = np.where(species==spe)[0]\n",
    "        print(spe, idx_spe)\n",
    "        bisp_gij[:,ispe, ...]= np.sum(rho2ij[:,idx_spe,...],axis =1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "bisp_gij.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "bisp_gij[0,0,1,0,2,0,2,:,0][np.where(bisp_gij[0,0,1,0,2,0,2,:,0])]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "rho3ilambda[0,0,1,0,2,0,2,:,0][np.where(rho3ilambda[0,0,1,0,2,0,2,:,0])]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "rho3ilambda[0][np.where(rho3ilambda[0])]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "bisp_gij[0][np.where(bisp_gij[0])]*2.14dd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "bisp_gij.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def compute_mp(frames,hypers, lmax, nui, nuj, cg, rhoi_pca=None, rhoij_rho2i_pca=None,rho2i_pca=None, scale=1):\n",
    "    spex = SphericalExpansion(**hypers)\n",
    "    rhoi = compute_rhoi(frames, spex, hypers)\n",
    "\n",
    "    # compresses further the spherical expansion features across species\n",
    "    if rhoi_pca is not None:\n",
    "        rhoi = apply_rhoi_pca(rhoi, rhoi_pca)\n",
    "\n",
    "    # makes sure that the spex used for the pair terms uses adaptive species\n",
    "    hypers_ij = deepcopy(hypers)\n",
    "    hypers_ij[\"expansion_by_species_method\"] = \"structure wise\"\n",
    "    spex_ij = SphericalExpansion(**hypers_ij)\n",
    "\n",
    "    tnat = 0\n",
    "    els = list(orbs.keys())\n",
    "    nel = len(els)\n",
    "    # prepare storage\n",
    "#     elL = list(itertools.product(els,range(lmax+1),[-1,1]))\n",
    "#     hetL = [ (els[i1], els[i2], L, pi) for i1 in range(nel) for i2 in range((i1+1 if half_hete else 0), nel) for L in range(lmax+1) for pi in [-1,1] ]\n",
    "#     feats = dict(diag = { L: [] for L in elL },\n",
    "#                  offd_p = { L: [] for L in elL },\n",
    "#                  offd_m = { L: [] for L in elL },\n",
    "#                  hete =   { L: [] for L in hetL },)\n",
    "\n",
    "    if rhoij_rho2i_pca is None and rho2i_pca is not None:\n",
    "        rhoij_rho2i_pca = rho2i_pca\n",
    "\n",
    "    #before = tracemalloc.take_snapshot()\n",
    "    for f in frames:\n",
    "        fnat = len(f.numbers)\n",
    "        frhoi = rhoi[tnat:tnat+fnat]*scale\n",
    "        fgij = compute_gij(f, spex_ij, hypers_ij)*scale\n",
    "        for L in range(lmax+1):\n",
    "            if nu==0:\n",
    "                lrhonui, lprhonui = np.ones((fnat, 1, 2*L+1)), np.ones((1))\n",
    "            elif nui==1 or nuj==1:\n",
    "                lrhonui, lprhonui = compute_rho1i_lambda(frhoi, L, cg)\n",
    "        \n",
    "        rho_mp = \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "frames,hypers, lmax, nui, nuj, cg = frames, spherical_expansion_hypers, 2, 1, 1, mycg\n",
    "rhoi_pca=None\n",
    "rhoij_rho2i_pca=None\n",
    "rho2i_pca=None\n",
    "scale=1\n",
    "tnat = 0\n",
    "hypers_ij = deepcopy(hypers)\n",
    "hypers_ij[\"expansion_by_species_method\"] = \"structure wise\"\n",
    "spex_ij = SphericalExpansion(**hypers_ij)\n",
    "rhoi = compute_rhoi(frames[:1], spex, hypers)\n",
    "# els = list(orbs.keys())\n",
    "# nel = len(els)\n",
    "for f in frames[:1]:\n",
    "    fnat = len(f.numbers)\n",
    "    frhoi = rhoi[tnat:tnat+fnat]*scale\n",
    "    rhonui, prhonui = frhoi, None\n",
    "    fgij = compute_gij(f, spex_ij, hypers_ij)*scale\n",
    "    for L in range(lmax+1):\n",
    "        #lrho0i, lprho0i = np.ones((fnat, 1, 2*L+1)), np.ones((1))\n",
    "        lrho1i, lprho1i = compute_rho1i_lambda(frhoi, L, cg)\n",
    "        \n",
    "        lrho0ij, prho0ij = compute_rho0ij_lambda(rhonui, fgij, L, cg, prhonui)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "np.where(rhonui[:,0,:, lm_slice(2)]-lrho1i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "lrho0ij"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "?compute_rho0ij_lambda"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def compute_rho0ij_lambda(rhoi, gij, L, cg,  prfeats = None): # prfeats is (in analogy with rho2ijlambda) the parity, but is not really necessary)\n",
    "    \"\"\" computes |rho^0_{ij}; lm> \"\"\"\n",
    "    rho0ij = gij[..., lm_slice(L)].reshape((gij.shape[0], gij.shape[1], -1, 2*L+1))\n",
    "    return rho0ij, np.ones(rho0ij.shape[2])\n",
    "\n",
    "def compute_rho1ij_lambda(rhoj, gij, L, cg, prfeats = None): \n",
    "    \"\"\" computes |rho^1_{ij}; lm> \"\"\"\n",
    "\n",
    "    lmax = int(np.sqrt(gij.shape[-1])) -1\n",
    "    # can't work out analytically how many terms we have, so we precompute it here\n",
    "    nl = 0\n",
    "    for l1 in range(lmax + 1):\n",
    "        for l2 in range(lmax + 1):  # |rho_i> and |rho_ij; g> are not symmetric so we need all l2\n",
    "            if abs(l2 - l1) > L or l2 + l1 < L:\n",
    "                continue\n",
    "            nl += 1\n",
    "\n",
    "    rhoj = rhoj.reshape((rhoj.shape[0], -1, rhoj.shape[-1]))\n",
    "    # natom, natom, nel*nmax, nmax, lmax+1, lmax+1, M\n",
    "    shape = (rhoj.shape[0], rhoj.shape[0],\n",
    "             rhoj.shape[1] , gij.shape[2], nl, 2*L+1)\n",
    "    rho1jilambda = np.zeros(shape)\n",
    "    parity = np.ones(nl, dtype = int)*(1-2*(L%2))\n",
    "\n",
    "    il = 0\n",
    "    for l1 in range(lmax+1):\n",
    "        for l2 in range(lmax+1):\n",
    "            if abs(l2 - l1) > L or l2 + l1 < L:\n",
    "                continue\n",
    "            rho1jilambda[:,:,:,:,il] = cg.combine_einsum(rhoj[...,lm_slice(l1)], gij[...,lm_slice(l2)],\n",
    "                                                L, combination_string=\"in,ijN->ijnN\")\n",
    "            parity[il] *= (1-2*(l1%2)) * (1-2*(l2%2))\n",
    "            il+=1\n",
    "\n",
    "    return rho1jilambda, parity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "spex = SphericalExpansion(**hypers)\n",
    "rhoi = compute_rhoi(frames, spex, hypers)\n",
    "hypers_ij = deepcopy(hypers)\n",
    "hypers_ij[\"expansion_by_species_method\"] = \"structure wise\"\n",
    "spex_ij = SphericalExpansion(**hypers_ij)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# Test features and Hamiltonian transformation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "import re\n",
    "from ase.data import atomic_numbers\n",
    "import json\n",
    "class NpEncoder(json.JSONEncoder):\n",
    "    def default(self, obj):\n",
    "        if isinstance(obj, np.integer):\n",
    "            return int(obj)\n",
    "        elif isinstance(obj, np.floating):\n",
    "            return float(obj)\n",
    "        elif isinstance(obj, np.ndarray):\n",
    "            return obj.tolist()\n",
    "        else:\n",
    "            return super(NpEncoder, self).default(obj)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# frames = read('qm7/qm7b-chno.xyz', ':100')\n",
    "# fock = np.load('qm7/qm7-chno-fock.npy', allow_pickle=True)\n",
    "frames = read('water_random/water_randomized_1000.xyz',':')\n",
    "for f in frames:\n",
    "    f.pbc=False\n",
    "    \n",
    "focks = np.load('water/water-fock.npy', allow_pickle=True)\n",
    "orbs = json.loads(json.load(open('water/orbs.json', \"r\")))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "we have to fix the L=1 terms that are stored in a weird order. we do this as post-processing \n",
    "and one should undo the change if the matrix is to be read back into pyscf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "for i in range(len(focks)):\n",
    "    focks[i] = pyscf_fix_l1(focks[i], frames[i], orbs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "## Rotational behavior of features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "\n",
    "from rascal.representations import SphericalExpansion, SphericalInvariants\n",
    "from rascal.utils import (get_radial_basis_covariance, get_radial_basis_pca, \n",
    "                          get_radial_basis_projections, get_optimal_radial_basis_hypers )\n",
    "from rascal.utils import radial_basis\n",
    "from rascal.utils import (WignerDReal, ClebschGordanReal, \n",
    "                          spherical_expansion_reshape,\n",
    "                          lm_slice, real2complex_matrix, xyz_to_spherical, spherical_to_xyz)\n",
    "from rascal.utils.cg_utils import _r2c as r2c\n",
    "from rascal.utils.cg_utils import _c2r as c2r\n",
    "from rascal.utils.cg_utils import _cg as clebsch_gordan\n",
    "from rascal.utils.cg_utils import _rotation as rotation\n",
    "from rascal.utils.cg_utils import _wigner_d as wigner_d"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "### Rotation of features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "frame = frames[90]\n",
    "frame.cell = [100,100,100]\n",
    "frame.positions += 50\n",
    "rotframe = frame.copy()\n",
    "rotframe.positions = rotframe.positions @ mrot.T\n",
    "rotframe.cell = rotframe.cell @ mrot.T   # rotate also the cell\n",
    "\n",
    "feats = spex.transform(frame).get_features(spex)\n",
    "feats *= 1e6\n",
    "rfeats = spherical_expansion_reshape(feats, **spherical_expansion_hypers)\n",
    "\n",
    "g= get_gij_fast(frame, spex,spherical_expansion_hypers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "rotfeats = spex.transform(rotframe).get_features(spex)\n",
    "rotfeats *= 1e6\n",
    "rotfeats = spherical_expansion_reshape(rotfeats, **spherical_expansion_hypers)\n",
    "\n",
    "rotg = get_gij_fast(rotframe, spex,spherical_expansion_hypers)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Check rotation of $| \\overline{\\rho_{ij}^0; \\lambda \\mu} \\rangle \\equiv \\langle n | g; \\lambda \\mu\\rangle $"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "np.linalg.norm(WD.rotate(mk_rho0ijlambda_fast(g, 4, mycg)) - mk_rho0ijlambda_fast(rotg, 4, mycg))/np.linalg.norm(mk_rho0ijlambda_fast(g, 4, mycg))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Check rotation of $| \\overline{\\rho_{i}^1; \\lambda \\mu} \\rangle $"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "(np.linalg.norm(WD.rotate(rfeats[...,lm_slice(3)]) - rotfeats[...,lm_slice(3)]))/np.linalg.norm(rfeats[...,lm_slice(3)])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Check $| \\overline{\\rho_{i}^2; \\lambda \\mu} \\rangle $"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "lsoap=mk_rho2ilambda_fast(rfeats, 1, mycg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "rotlsoap =  mk_rho2ilambda_fast(rotfeats, 1, mycg)\n",
    "np.linalg.norm(WD.rotate(lsoap) -rotlsoap)/np.linalg.norm(lsoap)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Check rotation of $| \\overline{\\rho_{ij}^1; 00} \\rangle $"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "rhoij = mk_rho1ij_fast(rfeats, g, mycg)\n",
    "rhoij_rot = mk_rho1ij_fast(rotfeats, rotg, mycg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "np.linalg.norm(WD.rotate(rhoij[...,np.newaxis]) - rhoij_rot[...,np.newaxis])/np.linalg.norm(rhoij)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Check rotation of $| \\overline{\\rho_{ij}^1; \\lambda \\mu} \\rangle $"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "rhoijlm= mk_rho1ijlambda_fast(rfeats, g, 3, mycg)\n",
    "rhoijlm_rot = mk_rho1ijlambda_fast(rotfeats, rotg, 3, mycg)\n",
    "np.linalg.norm(WD.rotate(rhoijlm) - rhoijlm_rot)/np.linalg.norm(rhoijlm)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# Simple regression test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "iwater = 99\n",
    "fock = np.load('water/water-fock.npy', allow_pickle=True)[iwater]\n",
    "orbs = json.loads(json.load(open(\"water/orbs.json\", \"r\")))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "there has to be a model for each (n1,l1,n2,l2,L) entry in the coupled representation of the Fock matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "frame = read('water/water_coords_1000.xyz',':')[iwater]\n",
    "frame.cell = [100,100,100]\n",
    "frame.positions += 50\n",
    "frame.symbols"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "fock = pyscf_fix_l1(fock, frame, orbs)\n",
    "fock_blocks = pyscf_to_blocks(fock, frame, orbs)\n",
    "fock_blocks_c = to_coupled(fock_blocks, mycg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "ffeats = do_full_features([frame], orbs, spherical_expansion_hypers, 4, mycg, scale=1e3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "FR = FockRegression(orbs, alpha=1e-18, solver='svd')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "FR.fit(ffeats, fock_blocks_c)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "fpred = FR.predict(ffeats)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "fpred['diag'][(2,1,2,1)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "fock_blocks_c['diag'][(2,1,2,1)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "unfock = blocks_to_pyscf(to_decoupled(fpred, mycg), frame, orbs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "plt.matshow((unfock-fock).astype(float))\n",
    "np.linalg.norm(unfock-fock) # bingo!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "np.mean(np.abs(np.linalg.eigvalsh(fock.astype(float))-np.linalg.eigvalsh(unfock)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "plt.plot(np.linalg.eigvalsh(fock.astype(float)), 'b.')\n",
    "plt.plot(np.linalg.eigvalsh(unfock.astype(float)), 'r--')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "#### Rotation and permutation "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "iwater = 99\n",
    "frame = read('water/water_coords_1000.xyz',':')[iwater]\n",
    "frame.cell = [100,100,100]\n",
    "frame.positions += 50\n",
    "print(frame.symbols)\n",
    "fock = np.load('water/water-fock.npy', allow_pickle=True)[iwater]\n",
    "orbs = json.loads(json.load(open(\"water/orbs.json\", \"r\")))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "fock = pyscf_fix_l1(fock, frame, orbs)\n",
    "fock_blocks = pyscf_to_blocks(fock, frame, orbs)\n",
    "fock_blocks_c = to_coupled(fock_blocks, mycg)\n",
    "feats = do_full_features([frame], orbs, spherical_expansion_hypers, 4, mycg, scale=1e3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "FR = FockRegression(orbs, alpha=1e-18, solver='svd')\n",
    "FR.fit(feats, fock_blocks_c)\n",
    "fpred = FR.predict(feats)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "fock_original = blocks_to_pyscf(to_decoupled(fpred, mycg), frame, orbs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "frame_rotperm = frame.copy()\n",
    "iperm = np.arange(len(frame.numbers), dtype=int)\n",
    "np.random.shuffle(iperm)\n",
    "frame_rotperm.numbers = frame_rotperm.numbers[iperm]\n",
    "frame_rotperm.positions = frame_rotperm.positions[iperm]\n",
    "print(frame_rotperm.symbols)\n",
    "\n",
    "abc = np.random.uniform(size=(3))*np.pi\n",
    "WD = WignerDReal(spherical_expansion_hypers[\"max_angular\"], *abc)\n",
    "WD.rotate_frame(frame_rotperm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "feat_rotperm = do_full_features([frame_rotperm], orbs, spherical_expansion_hypers, 4, mycg, scale=1e3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "pred_rotperm = FR.predict(feat_rotperm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "fock_rotperm = blocks_to_pyscf(to_decoupled(pred_rotperm, mycg), frame_rotperm, orbs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "plt.matshow((fock_rotperm-fock_original).astype(float))\n",
    "np.linalg.norm(fock_rotperm-fock_original) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "plt.plot(np.linalg.eigvalsh(fock_original.astype(float)), 'b.')\n",
    "plt.plot(np.linalg.eigvalsh(fock_rotperm.astype(float)), 'r--') \n",
    "print(np.mean(np.abs(np.linalg.eigvalsh(fock_original)-np.linalg.eigvalsh(fock_rotperm))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.3"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {
    "height": "calc(100% - 180px)",
    "left": "10px",
    "top": "150px",
    "width": "288px"
   },
   "toc_section_display": true,
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}