<jupyter_start><jupyter_text>Copyright 2022 Google LLCLicensed under the Apache License, Version 2.0 (the "License");you may not use this file except in compliance with the License.You may obtain a copy of the License at    https://www.apache.org/licenses/LICENSE-2.0Unless required by applicable law or agreed to in writing, softwaredistributed under the License is distributed on an "AS IS" BASIS,WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.See the License for the specific language governing permissions andlimitations under the License.  Colab notebook for demonstrating SCF calculations using the GAS22 functionalSCF calculations are performed using PySCF package (v1.7.6). Functional derivatives are evaluated using automatic differentiation with JAX.Link to this notebook: https://colab.research.google.com/github/google-research/google-research/blob/master/symbolic_functionals/colab/run_GAS22.ipynb<jupyter_code>!pip install pyscf==1.7.6
import numpy as np
import jax
import jax.numpy as jnp
from pyscf import dft
from pyscf import gto

jax.config.update('jax_enable_x64', True)

HYBRID_COEFF = dft.libxc.hybrid_coeff('wb97mv')
RSH_PARAMS = dft.libxc.rsh_coeff('wb97mv')<jupyter_output><empty_output><jupyter_text>Helper functions to compute LDA densities and auxiliary quantities<jupyter_code># _LDA_PW_PARAMETERS: dictionaries containing
# parameters a, alpha, beta1, beta2, beta3, beta4 and fpp_0, which are used to
# compute LDA correlation functional in the PW parametrization. The first six
# parameters are used to evaluate auxiliary G functions in the PW paper
# (10.1103/PhysRevB.45.13244). fpp_0 denotes the second derivative of
# f_zeta function in the PW paper at zeta = 0.
_LDA_PW_PARAMETERS = {
    'eps_c_unpolarized': {
        'a': 0.031091,
        'alpha1': 0.21370,
        'beta1': 7.5957,
        'beta2': 3.5876,
        'beta3': 1.6382,
        'beta4': 0.49294},
    'eps_c_polarized': {
        'a': 0.015545,
        'alpha1': 0.20548,
        'beta1': 14.1189,
        'beta2': 6.1977,
        'beta3': 3.3662,
        'beta4': 0.62517},
    'alpha_c': {
        'a': 0.016887,
        'alpha1': 0.11125,
        'beta1': 10.357,
        'beta2': 3.6231,
        'beta3': 0.88026,
        'beta4': 0.49671},
    'fpp_0': 1.709921
}

EPSILON = 1e-50

def e_x_lda_unpolarized(rho):
  """Evaluates LDA exchange energy density for spin unpolarized case.

  Parr & Yang Density-functional theory of atoms and molecules Eq. 6.1.20.

  Args:
    rho: Float numpy array with shape (num_grids,), the electron density.

  Returns:
    Float numpy array with shape (num_grids,), the LDA exchange energy density.
  """
  return -3 / 4 * (3 / jnp.pi) ** (1 / 3) * rho ** (4 / 3)


def e_x_lda_polarized(rhoa, rhob):
  """Evaluates LDA exchange energy density for spin polarized case.

  Parr & Yang Density-functional theory of atoms and molecules Eq. 8.2.18.

  Args:
    rhoa: Float numpy array with shape (num_grids,), the spin up
      electron density.
    rhob: Float numpy array with shape (num_grids,), the spin down
      electron density.

  Returns:
    Float numpy array with shape (num_grids,), the LDA exchange energy density.
  """
  rho = rhoa + rhob
  zeta = (rhoa - rhob) / (rho + EPSILON)  # spin polarization
  spin_scaling = 0.5 * ((1 + zeta) ** (4 / 3) + (1 - zeta) ** (4 / 3))
  return - 3 / 4 * (3 / jnp.pi) ** (1 / 3) * rho ** (4 / 3) * spin_scaling


def g_lda_pw(rs, a, alpha1, beta1, beta2, beta3, beta4):
  """Evaluates auxiliary function G in the PW parametrization of LDA functional.

  10.1103/PhysRevB.45.13244 Eq. 10.

  Args:
    rs: Float numpy array with shape (num_grids,), Wigner-Seitz radius.
    a: Float, parameter.
    alpha1: Float, parameter.
    beta1: Float, parameter.
    beta2: Float, parameter.
    beta3: Float, parameter.
    beta4: Float, parameter.
    use_jax: Boolean, if True, use jax.numpy for calculations, otherwise use
      numpy.

  Returns:
    Float numpy array with shape (num_grids,), auxiliary function G.
  """
  den = 2 * a * (
      beta1 * rs**(1 / 2) + beta2 * rs + beta3 * rs**(3 / 2) + beta4 * rs**2)
  return -2 * a * (1 + alpha1 * rs) * jnp.log(1 + 1 / den)


def get_wigner_seitz_radius(rho):
  """Evaluates Wigner-Seitz radius of given density.

  Args:
    rho: Float numpy array with shape (num_grids,), the electron density.

  Returns:
    Float numpy array with shape (num_grids,), the Wigner-Seitz radius.
  """
  return (3 / (4 * jnp.pi)) ** (1/3) * (rho + EPSILON) ** (-1 / 3)


def e_c_lda_unpolarized(rho):
  """Evaluates LDA correlation energy density for spin unpolarized case.

  PW parametrization. 10.1103/PhysRevB.45.13244 Eq. 8-9.

  Args:
    rho: Float numpy array with shape (num_grids,), the electron density.
    use_pbe_params: Boolean, whether use PBE parameters.
    use_jax: Boolean, if True, use jax.numpy for calculations, otherwise use
      numpy.

  Returns:
    Float numpy array with shape (num_grids,), the LDA correlation energy
      density.
  """
  rs = get_wigner_seitz_radius(rho)
  return rho * g_lda_pw(rs=rs, **_LDA_PW_PARAMETERS['eps_c_unpolarized'])


def e_c_lda_polarized(rhoa, rhob):
  """Evaluates LDA correlation energy density for spin polarized case.

  PW parametrization. 10.1103/PhysRevB.45.13244 Eq. 8-9.

  Args:
    rhoa: Float numpy array with shape (num_grids,), the spin up
      electron density.
    rhob: Float numpy array with shape (num_grids,), the spin down
      electron density.
    use_pbe_params: Boolean, whether use PBE parameters.

  Returns:
    Float numpy array with shape (num_grids,), the LDA correlation energy
      density.
  """
  rho = rhoa + rhob
  rs = get_wigner_seitz_radius(rho)
  zeta = (rhoa - rhob) / (rho + EPSILON)

  # spin dependent interpolation coefficient
  f_zeta = 1 / (2 ** (4 / 3) - 2) * (
      (1 + zeta) ** (4 / 3) + (1 - zeta) ** (4 / 3) - 2)

  eps_c_unpolarized = g_lda_pw(
      rs=rs, **_LDA_PW_PARAMETERS['eps_c_unpolarized'])
  eps_c_polarized = g_lda_pw(
      rs=rs, **_LDA_PW_PARAMETERS['eps_c_polarized'])
  alpha_c = -g_lda_pw(rs=rs, **_LDA_PW_PARAMETERS['alpha_c'])

  return rho * (
      eps_c_unpolarized
      + (1 / _LDA_PW_PARAMETERS['fpp_0']) * alpha_c * f_zeta * (1 - zeta ** 4)
      + (eps_c_polarized - eps_c_unpolarized) * f_zeta * zeta ** 4
      )


def decomposed_e_c_lda_unpolarized(rho):
  """Evaluates LDA e_c decomposed into same-spin and opposite-spin components.

  This function returns the LDA correlation energy density partitioned into
  same-spin and opposite-spin components. 10.1063/1.475007 Eq. 7-8.

  Args:
    rho: Float numpy array with shape (num_grids,), the electron density.

  Returns:
    e_c_ss: Float numpy array with shape (num_grids,), the same-spin
      component of LDA correlation energy density.
    e_c_os: Float numpy array with shape (num_grids,), the opposite-spin
      component of LDA correlation energy density.
  """
  e_c = e_c_lda_unpolarized(rho)
  e_c_ss = 2 * e_c_lda_polarized(rho / 2, jnp.zeros_like(rho))
  e_c_os = e_c - e_c_ss

  return e_c_ss, e_c_os


def decomposed_e_c_lda_polarized(rhoa, rhob):
  """Evaluates LDA e_c decomposed into same-spin and opposite-spin components.

  This function returns the LDA correlation energy density partitioned into
  same-spin and opposite-spin components. 10.1063/1.475007 Eq. 7-8.

  Args:
    rhoa: Float numpy array with shape (num_grids,), the spin up
      electron density.
    rhob: Float numpy array with shape (num_grids,), the spin down
      electron density.

  Returns:
    e_c_aa: Float numpy array with shape (num_grids,), the same-spin (aa)
      component of LDA correlation energy density.
    e_c_bb: Float numpy array with shape (num_grids,), the same-spin (bb)
      component of LDA correlation energy density.
    e_c_ab: Float numpy array with shape (num_grids,), the opposite-spin
      component of LDA correlation energy density.
  """
  zero = jnp.zeros_like(rhoa)

  e_c = e_c_lda_polarized(rhoa, rhob)
  e_c_aa = e_c_lda_polarized(rhoa, zero)
  e_c_bb = e_c_lda_polarized(zero, rhob)
  e_c_ab = e_c - e_c_aa - e_c_bb

  return e_c_aa, e_c_bb, e_c_ab


def get_reduced_density_gradient(rho, sigma):
  """Evaluates reduced density gradient.

  Args:
    rho: Float numpy array with shape (num_grids,), the electron density.
    sigma: Float numpy array with shape (num_grids,), the norm square of
      density gradient.

  Returns:
    Float numpy array with shape (num_grids,), the reduced density gradient.
  """
  # NOTE: EPSILON is inserted in such a way that the reduce density
  # gradient is zero if rho or sigma or both are zero.
  return jnp.sqrt(sigma + EPSILON ** 3) / (rho ** (4 / 3) + EPSILON)


def get_mgga_t(rho, tau, polarized):
  """Evaluates the auxiliary quantity t in meta-GGA functional forms.

  t = (tau_HEG / tau),  where tau_HEG is the kinetic energy density of
  homogeneous electron gass.

  Args:
    rho: Float numpy array with shape (num_grids,), the electron density.
    tau: Float numpy array with shape (num_grids,), the kinetic energy density.
    polarized: Boolean, whether the system is spin polarized.

  Returns:
    Float numpy array with shape (num_grids,), the auxiliary quantity w.
  """
  spin_factor = 1 if polarized else (1 / 2) ** (2 / 3)
  # 3 / 10 instead of 3 / 5 is used below because tau is defined with 1 / 2
  tau_heg = 3 / 10 * (6 * jnp.pi ** 2) ** (2 / 3) * rho ** (5 / 3)
  return spin_factor * tau_heg / (tau + EPSILON)


def f_rsh(rho, omega=RSH_PARAMS[0], polarized=False,):
  """Enchancement factor for evaluating short-range semilocal exchange.

  10.1063/1.4952647 Eq. 11.

  Args:
    rho: Float numpy array with shape (num_grids,), the electron density.
    omega: Float, the range seperation parameter.
    polarized: Boolean, whether the system is spin polarized.

  Returns:
    Float numpy array with shape (num_grids,), the RSH enhancement factor.
  """
  spin_factor = 1 if polarized else 2
  # Fermi wave vector
  kf = (6 * jnp.pi**2 * rho / spin_factor + EPSILON) ** (1 / 3)
  a = omega / kf + EPSILON  # variable a in Eq. 11
  return (1 - 2 / 3 * a * (2 * jnp.pi ** (1 / 2) * jax.scipy.special.erf(1 / a) - 3 * a
                           + a ** 3 + (2 * a - a ** 3) * jnp.exp(-1 / a ** 2)))<jupyter_output><empty_output><jupyter_text>GAS22 enhancement factorsThe following enhancement factors for the GAS22 functional corresponds to the Eq. 9-11 in the manuscript.<jupyter_code># NOTE: In the following functions, the square of reduced density gradient are
# used as function arguments (x2).

def get_u(x2, gamma):
  """Evaluates the auxiliary quantity u = gamma x^2 / (1 + gamma x^2)."""
  return gamma * x2 / (1 + gamma * x2)

def f_x_gas22(x2, w):
  """Evaluates the exchange enhancement factor for the GAS22 functional."""
  u = get_u(x2, gamma=0.003840616724010807)
  return 0.862139736374172 + 0.936993691972698 * u  + 0.317533683085033 * w

def f_css_gas22(x2, w):
  """Evaluates the same-spin correlation enhancement factor for the GAS22 functional."""
  u = get_u(x2, gamma=0.46914023462026644)
  return (u - 4.10753796482853 * w - 5.24218990333846 * w**2
          - 1.76643208454076 * u**6 + 7.5380689617542 * u**6 * w**4)

def f_cos_gas22(x2, w):
  """Evaluates the opposite-spin correlation enhancement factor for the GAS22 functional."""
  return (0.805124374375355 + 7.98909430970845 * w**2
          - 1.76098915061634 * w**2 * jnp.cbrt(x2) - 7.54815900595292 * w**6
          + 2.00093961824784 * w**6 * jnp.cbrt(x2))<jupyter_output><empty_output><jupyter_text>Define GAS22 in the format of PySCF custom functions<jupyter_code>def gas22_unpolarized(rho, sigma, tau):
  """Evaluates XC energy density for spin unpolarized case.

  Args:
    rho: Float numpy array with shape (num_grids,), the electron density.
    sigma: Float numpy array with shape (num_grids,), the norm square of
      density gradient.
    tau: Float numpy array with shape (num_grids,), the kinetic energy
      density.

  Returns:
    Float numpy array with shape (num_grids,), the XC energy density.
  """
  rho_s = 0.5 * rho
  x_s = get_reduced_density_gradient(rho_s, 0.25 * sigma)
  x2_s = x_s ** 2
  t_s = get_mgga_t(rho_s, 0.5 * tau, polarized=True)
  w_s = (t_s - 1) / (t_s + 1)

  e_lda_x = e_x_lda_unpolarized(rho) * f_rsh(rho, polarized=False)
  e_lda_css, e_lda_cos = decomposed_e_c_lda_unpolarized(rho)
  
  f_x = f_x_gas22(x2_s, w_s)
  f_css = f_css_gas22(x2_s, w_s)
  f_cos = f_cos_gas22(x2_s, w_s)

  return e_lda_x * f_x + e_lda_css * f_css + e_lda_cos * f_cos

def gas22_polarized(rho_a, rho_b, sigma_aa, sigma_ab, sigma_bb, tau_a, tau_b):
  """Evaluates XC energy density for spin polarized case.

  Args:
    rho_a: Float numpy array with shape (num_grids,), the spin up electron
      density.
    rho_b: Float numpy array with shape (num_grids,), the spin down electron
      density.
    sigma_aa: Float numpy array with shape (num_grids,), the norm square of
      density gradient (aa component).
    sigma_ab: Float numpy array with shape (num_grids,), the norm square of
      density gradient (ab component).
    sigma_bb: Float numpy array with shape (num_grids,), the norm square of
      density gradient (bb component).
    tau_a: Float numpy array with shape (num_grids,), the spin up kinetic
      energy density.
    tau_b: Float numpy array with shape (num_grids,), the spin down kinetic
      energy density.

  Returns:
    Float numpy array with shape (num_grids,), the XC energy density.
  """
  del sigma_ab
  rho_ab = 0.5 * (rho_a + rho_b)

  x_a = get_reduced_density_gradient(rho_a, sigma_aa)
  x_b = get_reduced_density_gradient(rho_b, sigma_bb)
  x2_a = x_a ** 2
  x2_b = x_b ** 2
  x2_ab = 0.5 * (x2_a + x2_b)

  t_a = get_mgga_t(rho_a, tau_a, polarized=True)
  t_b = get_mgga_t(rho_b, tau_b, polarized=True)
  t_ab = 0.5 * (t_a + t_b)
  w_a = (t_a - 1) / (t_a + 1)
  w_b = (t_b - 1) / (t_b + 1)
  w_ab = (t_ab - 1) / (t_ab + 1)

  e_lda_x_a = 0.5 * e_x_lda_unpolarized(2 * rho_a) * f_rsh(
      rho_a, polarized=True)
  e_lda_x_b = 0.5 * e_x_lda_unpolarized(2 * rho_b) * f_rsh(
      rho_b, polarized=True)
  e_lda_css_a, e_lda_css_b, e_lda_cos = decomposed_e_c_lda_polarized(
      rho_a, rho_b)

  f_x_a = f_x_gas22(x2_a, w_a)
  f_x_b = f_x_gas22(x2_b, w_b)
  f_css_a = f_css_gas22(x2_a, w_a)
  f_css_b = f_css_gas22(x2_b, w_b)
  f_cos = f_cos_gas22(x2_ab, w_ab)

  return (e_lda_x_a * f_x_a + e_lda_x_b * f_x_b
          + e_lda_css_a * f_css_a + e_lda_css_b * f_css_b + e_lda_cos * f_cos)

def eval_xc_gas22(xc_code,
                  rho_and_derivs,
                  spin=0,
                  relativity=0,
                  deriv=1,
                  verbose=None):
  """Evaluates exchange-correlation energy densities and derivatives.

  Args:
    xc_code: A dummy argument, not used.
    rho_and_derivs: Float numpy array with shape (6, num_grids) for spin
      unpolarized case; 2-tuple of float numpy array with shape (6, num_grids)
      for spin polarized case. Electron density and its derivatives. For
      spin unpolarized case, the 6 subarrays represent (density, gradient_x,
      gradient_y, gradient_z, laplacian, tau); for spin polarized case, the
      spin up and spin down densities and derivatives are each represented
      with a (6, num_grids) array.
    spin: Integer, 0 for spin unpolarized and 1 for spin polarized
      calculations.
    relativity: A dummy argument, not used.
    deriv: Integer, the order of derivatives evaluated for XC energy density.
    verbose: A dummy argument, not used.

  Returns:
    eps_xc: Float numpy array with shape (num_grids,), the XC energy density
      per particle.
    v_xc: Tuple of float numpy arrays, the first derivatives of XC energy
      density per volume to various quantities. See pyscf/dft/libxc.py for
      more details.
    f_xc: A dummy return value, not used.
    k_xc: A dummy return value, not used.

  Raises:
    NotImplementedError: If derivative order higher than one is requested
      (deriv > 1).
  """
  del xc_code, relativity, verbose

  if deriv != 1:
    raise NotImplementedError('Only deriv = 1 is implemented.')

  if spin == 0:
    rho, grad_x, grad_y, grad_z, _, tau = rho_and_derivs
    sigma = grad_x**2 + grad_y**2 + grad_z**2

    e_xc, grads = jax.jit(jax.vmap(jax.value_and_grad(
        gas22_unpolarized, argnums=(0, 1, 2))))(rho, sigma, tau)
    vrho, vsigma, vtau = grads

  else:
    rhoa, grad_x_a, grad_y_a, grad_z_a, _, tau_a = rho_and_derivs[0]
    rhob, grad_x_b, grad_y_b, grad_z_b, _, tau_b = rho_and_derivs[1]
    rho = rhoa + rhob
    sigma_aa = grad_x_a**2 + grad_y_a**2 + grad_z_a**2
    sigma_ab = grad_x_a * grad_x_b + grad_y_a * grad_y_b + grad_z_a * grad_z_b
    sigma_bb = grad_x_b**2 + grad_y_b**2 + grad_z_b**2

    e_xc, grads = jax.jit(jax.vmap(jax.value_and_grad(
        gas22_polarized, argnums=(0, 1, 2, 3, 4, 5, 6))))(
            rhoa, rhob, sigma_aa, sigma_ab, sigma_bb, tau_a, tau_b)

    vrhoa, vrhob, vsigma_aa, vsigma_ab, vsigma_bb, vtau_a, vtau_b = grads
    vrho = np.stack((vrhoa, vrhob), axis=1)
    vsigma = np.stack((vsigma_aa, vsigma_ab, vsigma_bb), axis=1)
    vtau = np.stack((vtau_a, vtau_b), axis=1)

  eps_xc = e_xc / (rho + EPSILON)
  return (np.array(eps_xc),
          (np.array(vrho), np.array(vsigma),
            np.zeros_like(vtau), np.array(vtau)),
          None,
          None)<jupyter_output><empty_output><jupyter_text>SCF calculations with PySCF<jupyter_code>def run_scf_for_mol(atom,
                    basis,
                    charge,
                    spin,
                    conv_tol=1e-6,
                    use_sg1_prune_for_nlc=True):
  """Runs SCF calculations for molecule using GAS22 functional.
  
  Args:
    atom: String, the atomic structure.
    charge: Integer, the total charge of molecule.
    spin: Integer, the difference of spin up and spin down electron numbers.
    basis: String, the GTO basis.
    conv_tol: Float, the convergence threshold of total energy.
    use_sg1_prune_for_nlc: Boolean, whether use SG1 prune for NLC calculation.
  
  Returns:
    Float, the SCF total energy.
  """
  mol = gto.M(atom=atom, basis=basis, charge=charge, spin=spin, verbose=4)

  if spin == 0:
    ks = dft.RKS(mol)
  else:
    ks = dft.UKS(mol)

  ks.define_xc_(eval_xc_gas22, xctype='MGGA', hyb=HYBRID_COEFF, rsh=RSH_PARAMS)

  # NOTE: This is necessary because PySCF will use this name to decide the
  # parameters for VV10 NLC. GAS22 has identical nonlocal XC as wB97M-V, including
  # VV10 NLC. 
  ks.xc = 'wb97mv'
  ks.nlc = 'VV10'

  if use_sg1_prune_for_nlc:
    # NOTE: SG1 prune can be used to reduce the computational cost of
    # VV10 NLC. The use of SG1 prune has very little effect on the resulting
    # XC energy. SG1 prune is used in PySCF's example
    # pyscf/examples/dft/33-nlc_functionals.py and is also used in paper
    # 10.1080/00268976.2017.1333644. Note that SG1 prune in PySCF is not
    # available for some elements appeared in the MCGDB84 database.
    ks.nlcgrids.prune = dft.gen_grid.sg1_prune

  # NOTE: It is necessary to override ks._numint._xc_type method to
  # let PySCF correctly use a custom XC functional with NLC. Also, note that
  # ks.xc is used to determine NLC parameters.
  ks._numint._xc_type = lambda code: 'NLC' if 'VV10' in code else 'MGGA'

  ks.conv_tol = conv_tol

  ks.kernel()

  return ks.e_tot
# Example: water molecule (spin unpolarized)
# Atomic structure is taken from 01a_water_monA_3B-69.xyz in MGCDB84
# Expected Etot: 76.430060 au
run_scf_for_mol(
    atom = """\
    O -0.0848890000 0.0568040000 0.0552050000
    H 0.7103750000 0.5177290000 0.4139710000
    H -0.8909580000 0.4985880000 0.4142740000
    """,
    basis = 'def2qzvppd',
    charge = 0,
    spin = 0,
)
# Example: P atom (spin polarized)
# Atomic structure is taken from 25_P_AE18.xyz in MGCDB84
# Expected Etot: -341.253901 au
run_scf_for_mol(
    atom = "P 0.0000000000 0.0000000000 0.0000000000",
    basis = 'def2qzvppd',
    charge = 0,
    spin = 3,
)<jupyter_output>System: uname_result(system='Linux', node='adbbb022342a', release='5.4.144+', version='#1 SMP Tue Dec 7 09:58:10 PST 2021', machine='x86_64', processor='x86_64')  Threads 2
Python 3.7.12 (default, Jan 15 2022, 18:48:18) 
[GCC 7.5.0]
numpy 1.19.5  scipy 1.4.1
Date: Tue Feb  8 18:44:05 2022
PySCF version 1.7.6
PySCF path  /usr/local/lib/python3.7/dist-packages/pyscf

[CONFIG] conf_file None
[INPUT] verbose = 4
[INPUT] num. atoms = 1
[INPUT] num. electrons = 15
[INPUT] charge = 0
[INPUT] spin (= nelec alpha-beta = 2S) = 3
[INPUT] symmetry False subgroup None
[INPUT] Mole.unit = angstrom
[INPUT]  1 P      0.000000000000   0.000000000000   0.000000000000 AA    0.000000000000   0.000000000000   0.000000000000 Bohr

nuclear repulsion = 0
number of shells = 24
number of NR pGTOs = 111
number of NR cGTOs = 76
basis = def2qzvppd
ecp = {}
CPU time:       103.62


******** <class 'pyscf.dft.uks.UKS'> ********
method = UKS-UHF
initial guess = minao
damping factor = 0
level_shift factor = 0
DIIS = <[...]
{'fraction_non_alphanumeric': 0.10339249179235857, 'fraction_numerical': 0.055621639625065424, 'mean_word_length': 3.243488794669897, 'pattern_counts': {'":': 0, '<': 22, '<?xml version=': 0, '>': 22, 'https://': 2, 'lorem ipsum': 0, 'www.': 1, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 9, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '514386', 'license': 'Apache-2.0', 'n_tokens_mistral': 8730, 'n_tokens_neox': 7608, 'n_words': 2342, 'path': 'symbolic_functionals/colab/run_GAS22.ipynb', 'repo_name': 'm-niemeyer/google-research', 'size': '21017'}