<commit_before>#!/usr/bin/python

#=============================================================================================
# MODULE DOCSTRING
#=============================================================================================

"""
Alchemical factory for free energy calculations that operates directly on OpenMM System objects.

DESCRIPTION

This module contains enumerative factories for generating alchemically-modified System objects
usable for the calculation of free energy differences of hydration or ligand binding.

This version uses fused elecrostatic and steric alchemical modifications.

TODO

* Add functions for the automatic optimization of alchemical states?
* Can we store serialized form of Force objects so that we can save time in reconstituting
  Force objects when we make copies?  We can even manipulate the XML representation directly.
* Allow protocols to automatically be resized to arbitrary number of states, to 
  allow number of states to be enlarged to be an integral multiple of number of GPUs.
* Add GBVI support to AlchemicalFactory.
* Test AMOEBA support.
* Can alchemically-modified System objects share unmodified Force objects to avoid overhead
  of duplicating Forces that are not modified?

"""

#=============================================================================================
# GLOBAL IMPORTS
#=============================================================================================

import numpy as np
import copy
import time

import simtk.openmm as openmm
import simtk.unit as unit

import logging
logger = logging.getLogger(__name__)

#=============================================================================================
# PARAMETERS
#=============================================================================================

ONE_4PI_EPS0 = 138.935456 # OpenMM constant for Coulomb interactions (openmm/platforms/reference/include/SimTKOpenMMRealType.h) in OpenMM units

#=============================================================================================
# MODULE UTILITIES
#=============================================================================================

def _is_periodic(system):
    """
    Determine whether the specified system is periodic.

    Parameters
    ----------
    system : simtk.openmm.System
         The system to check.

    Returns
    -------
    is_periodic : bool
        If True, the system uses a nonbonded method recognized as being periodic.

    Examples
    --------

    >>> from openmmtools import testsystems

    Periodic water box.

    >>> waterbox = testsystems.WaterBox()
    >>> print _is_periodic(waterbox.system)
    True

    Non-periodic Lennard-Jones cluster.

    >>> cluster = testsystems.LennardJonesCluster()
    >>> print _is_periodic(cluster.system)
    False

    Notes
    -----
    This method simply checks to see if any nonbonded methods recognized to be periodic are in use.
    Addition of new nonbonded methods or force types may require adding to recognized_periodic_methods.

    """
    forces = { system.getForce(index).__class__.__name__ for index in range(system.getNumForces()) }

    # Only the following nonbonded methods are recognized as being periodic.
    recognized_periodic_methods = {
        'NonbondedForce' : [openmm.NonbondedForce.CutoffPeriodic, openmm.NonbondedForce.Ewald, openmm.NonbondedForce.PME],
        'CustomNonbondedForce' : [openmm.CustomNonbondedForce.CutoffPeriodic],
        'AmoebaVdwForce' : [openmm.AmoebaVdwForce.CutoffPeriodic],
        }

    for force_index in range(system.getNumForces()):
        force = system.getForce(force_index)
        force_name = force.__class__.__name__
        if force_name in recognized_periodic_methods:
            method = force.getNonbondedMethod()
            if method in recognized_periodic_methods[force_name]:
                return True

    # Nothing we recognize was found, so assume system was not periodic.
    return False

#=============================================================================================
# AlchemicalState
#=============================================================================================

class AlchemicalState(dict):
    """
    Alchemical state description.

    These parameters describe the parameters that affect computation of the energy.

    Attributes
    ----------
    lambda_restraints : float
        Scaling factor for remaining receptor-ligand relative restraint terms (to help keep ligand near protein).
    lambda_electrostatics : float
        Scaling factor for ligand charges, intrinsic Born radii, and surface area term.
    lambda_sterics : float
        Scaling factor for ligand sterics (Lennard-Jones and Halgren) interactions.
    labmda_torsions : float
        Scaling factor for ligand non-ring torsions.

    """
    def __init__(self, **kwargs):
        self['lambda_restraints'] = 0.0
        self['lambda_electrostatics'] = 1.0
        self['lambda_sterics'] = 1.0
        self['lambda_torsions'] = 1.0

        for key in kwargs.keys():
            self[key] = kwargs[key]

#=============================================================================================
# AbsoluteAlchemicalFactory
#=============================================================================================

class AbsoluteAlchemicalFactory(object):
    """
    Factory for generating OpenMM System objects that have been alchemically perturbed for absolute binding free energy calculation.

    Examples
    --------

    Create alchemical intermediates for default alchemical protocol for p-xylene in T4 lysozyme L99A in GBSA.

    >>> # Create a reference system.
    >>> from openmmtools import testsystems
    >>> complex = testsystems.LysozymeImplicit()
    >>> [reference_system, positions] = [complex.system, complex.positions]
    >>> # Create a factory to produce alchemical intermediates.
    >>> receptor_atoms = range(0,2603) # T4 lysozyme L99A
    >>> ligand_atoms = range(2603,2621) # p-xylene
    >>> factory = AbsoluteAlchemicalFactory(reference_system, ligand_atoms=ligand_atoms)
    >>> # Get the default protocol for 'denihilating' in complex in explicit solvent.
    >>> protocol = factory.defaultComplexProtocolImplicit()
    >>> # Create the perturbed systems using this protocol.
    >>> systems = factory.createPerturbedSystems(protocol)

    Create alchemical intermediates for default alchemical protocol for one water in a water box.

    >>> # Create a reference system.
    >>> from openmmtools import testsystems
    >>> waterbox = testsystems.WaterBox()
    >>> [reference_system, positions] = [waterbox.system, waterbox.positions]
    >>> # Create a factory to produce alchemical intermediates.
    >>> factory = AbsoluteAlchemicalFactory(reference_system, ligand_atoms=[0, 1, 2])
    >>> # Get the default protocol for 'denihilating' in solvent.
    >>> protocol = factory.defaultSolventProtocolExplicit()
    >>> # Create the perturbed systems using this protocol.
    >>> systems = factory.createPerturbedSystems(protocol)

    """

    # Factory initialization.
    def __init__(self, reference_system, ligand_atoms=list(), receptor_atoms=list(), annihilate_electrostatics=True, annihilate_sterics=False, test_positions=None, platform=None):
        """
        Initialize absolute alchemical intermediate factory with reference system.

        The reference system will not be modified when alchemical intermediates are generated.

        Parmeters
        ---------
        reference_system : simtk.openmm.System
            The reference system that is to be alchemically modified.
        ligand_atoms : list of int, optional, default = []
            List of atoms to be designated as 'ligand' for alchemical modification; everything else in system is considered the 'environment'.
        annihilateElectrostatics : bool
            If True, electrostatics should be annihilated, rather than decoupled.
        annihilateSterics : bool
            If True, sterics (Lennard-Jones or Halgren potential) will be annihilated, rather than decoupled.
        test_positions : simtk.unit.Quantity of dimension (natoms,3) with units compatible with nanometers, optional, default=None
            If provided, these coordinates will be used to test alchemically-modified system to ensure the potential energy is finite.
            If the potential energy is NaN, the energy for each force component will be computed for the Reference platform to aid in debugging.
        platform : simtk.openmm.Platform, optionl default=None
            If provided, this Platform will be used to check energies are finite.

        """

        # Store annihilation/decoupling information.
        self.annihilate_electrostatics = annihilate_electrostatics
        self.annihilate_sterics = annihilate_sterics

        # Store serialized form of reference system.
        self.reference_system = copy.deepcopy(reference_system)

        # Store copy of atom sets.
        self.ligand_atoms = copy.deepcopy(ligand_atoms)

        # Store atom sets
        self.ligand_atomset = set(self.ligand_atoms)

        # Create an alchemically-modified system to cache.
        self.alchemically_modified_system = self._createAlchemicallyModifiedSystem(self.reference_system)

        # Store information for use in aiding debugging of alchemical factory.
        self.test_positions = test_positions
        self.platform = platform
        if self.test_positions is not None:
            self._checkEnergyIsFinite(self.alchemically_modified_system, test_positions, platform=platform)

        # DEBUG: Write XML
        debug_write_xml = False
        if debug_write_xml:
            import os, os.path
            def write_file(filename, contents):
                with open(filename, 'w') as outfile:
                    outfile.write(contents)

            logger.info("Serializing to XML...")
            system_filename = os.path.join('setup', 'alchemical-system.xml')
            write_file(system_filename, openmm.XmlSerializer.serialize(self.alchemically_modified_system))

        return

    def _checkEnergyIsFinite(self, system, positions, platform=None):
        """
        Test that the potential energy is finite for the provided positions.
        If energy is NaN, compute energy from each force component on Reference platform to aid in debugging.

        Parameters
        ----------
        system : simtk.openmm.System
             The alchemical system to check.
        positions : simtk.unit.Quantity of dimension (natoms,3) with units compatible with nanometers
            Coordinates to use for energy test.
        platform : simtk.openmm.Platform, optional, default=None
            If specified, this platform will be used to compute test energy (but not by force component).

        """
        logger.debug("Checking alchemical system produces finite energies.")

        def compute_potential_energy(system, positions, platform=None, groups=-1):
            # Compute potential energy of reference system.
            integrator = openmm.VerletIntegrator(1.0 * unit.femtosecond)
            if platform is not None:
                # Compute potential energy with requested platform.
                context = openmm.Context(system, integrator, platform)
            else:
                # Use fastest available platform.
                context = openmm.Context(system, integrator)
            context.setPositions(positions)
            potential = context.getState(getEnergy=True,groups=groups).getPotentialEnergy()
            del context, integrator
            return potential

        # Compute potential energy error between reference and alchemical system.
        reference_potential = compute_potential_energy(self.reference_system, positions, platform)
        alchemical_potential = compute_potential_energy(system, positions, platform)
        energy_error = alchemical_potential - reference_potential

        # If alchemical potential energy is NaN, compute energy by component on Reference platform.
        if np.isnan(alchemical_potential / unit.kilocalories_per_mole):
            logger.debug("Energy for alchemically modified system is NaN.")
            logger.debug("Decomposing energy by force component on Reference platform:")

            # Assign all forces to separate components.
            system = self.alchemically_modified_system
            for force_index in range(system.getNumForces()):
                system.getForce(force_index).setForceGroup(force_index)

            # Compute potential energy for each force using Reference platform.
            reference_platform = openmm.Platform.getPlatformByName('Reference')
            for force_index in range(system.getNumForces()):
                groups = 1 << force_index # group index bitwise selector
                potential = compute_potential_energy(system, positions, reference_platform, groups)
                force_classname = system.getForce(force_index).__class__.__name__
                logger.debug("Force %5d / %5d [%24s] %12.3f kcal/mol" % (force_index, system.getNumForces(), force_classname, potential / unit.kilocalories_per_mole))

            # Clean up
            del context, integrator

            raise Exception("Energy for alchemically modified system is NaN.")

        # Return the energy error.
        logger.debug("Difference between alchemical and reference potential energy is %8.3f kcal/mol" % (energy_error / unit.kilocalories_per_mole))
        return energy_error

    @classmethod
    def defaultComplexProtocolImplicit(cls):
        """
        Return the default protocol for 'denihilating' a ligand in complex with a protein in implicit solvent.

        Returns
        -------
        alchemical_states : list of AlchemicalState
            The list of alchemical states defining the protocol.

        Notes
        -----
        The unrestrained, fully interacting system is always listed first.

        """

        alchemical_states = list()
        lambda_values = [1.0, 0.99, 0.98, 0.97, 0.96, 0.95, 0.94, 0.93, 0.92, 0.91, 0.90, 0.88, 0.86, 0.84, 0.81,
                         0.78, 0.74, 0.70, 0.60, 0.50, 0.40, 0.30, 0.20, 0.10, 0.05, 0.025, 0.01, 0.00]

        for lambda_value in lambda_values:
            alchemical_state = AlchemicalState()
            alchemical_state['lambda_electrostatics'] = lambda_value
            alchemical_state['lambda_sterics'] = lambda_value
            alchemical_states.append(alchemical_state)

        return alchemical_states

    @classmethod
    def defaultComplexProtocolExplicit(cls):
        """
        Return the default protocol for 'denihilating' a ligand in complex with a protein in explicit solvent.

        Returns<commit_msg>Fixed issue where restraints were not being imposed in complex in implicit solvent.<commit_after>#!/usr/bin/python

#=============================================================================================
# MODULE DOCSTRING
#=============================================================================================

"""
Alchemical factory for free energy calculations that operates directly on OpenMM System objects.

DESCRIPTION

This module contains enumerative factories for generating alchemically-modified System objects
usable for the calculation of free energy differences of hydration or ligand binding.

This version uses fused elecrostatic and steric alchemical modifications.

TODO

* Add functions for the automatic optimization of alchemical states?
* Can we store serialized form of Force objects so that we can save time in reconstituting
  Force objects when we make copies?  We can even manipulate the XML representation directly.
* Allow protocols to automatically be resized to arbitrary number of states, to
  allow number of states to be enlarged to be an integral multiple of number of GPUs.
* Add GBVI support to AlchemicalFactory.
* Test AMOEBA support.
* Can alchemically-modified System objects share unmodified Force objects to avoid overhead
  of duplicating Forces that are not modified?

"""

#=============================================================================================
# GLOBAL IMPORTS
#=============================================================================================

import numpy as np
import copy
import time

import simtk.openmm as openmm
import simtk.unit as unit

import logging
logger = logging.getLogger(__name__)

#=============================================================================================
# PARAMETERS
#=============================================================================================

ONE_4PI_EPS0 = 138.935456 # OpenMM constant for Coulomb interactions (openmm/platforms/reference/include/SimTKOpenMMRealType.h) in OpenMM units

#=============================================================================================
# MODULE UTILITIES
#=============================================================================================

def _is_periodic(system):
    """
    Determine whether the specified system is periodic.

    Parameters
    ----------
    system : simtk.openmm.System
         The system to check.

    Returns
    -------
    is_periodic : bool
        If True, the system uses a nonbonded method recognized as being periodic.

    Examples
    --------

    >>> from openmmtools import testsystems

    Periodic water box.

    >>> waterbox = testsystems.WaterBox()
    >>> print _is_periodic(waterbox.system)
    True

    Non-periodic Lennard-Jones cluster.

    >>> cluster = testsystems.LennardJonesCluster()
    >>> print _is_periodic(cluster.system)
    False

    Notes
    -----
    This method simply checks to see if any nonbonded methods recognized to be periodic are in use.
    Addition of new nonbonded methods or force types may require adding to recognized_periodic_methods.

    """
    forces = { system.getForce(index).__class__.__name__ for index in range(system.getNumForces()) }

    # Only the following nonbonded methods are recognized as being periodic.
    recognized_periodic_methods = {
        'NonbondedForce' : [openmm.NonbondedForce.CutoffPeriodic, openmm.NonbondedForce.Ewald, openmm.NonbondedForce.PME],
        'CustomNonbondedForce' : [openmm.CustomNonbondedForce.CutoffPeriodic],
        'AmoebaVdwForce' : [openmm.AmoebaVdwForce.CutoffPeriodic],
        }

    for force_index in range(system.getNumForces()):
        force = system.getForce(force_index)
        force_name = force.__class__.__name__
        if force_name in recognized_periodic_methods:
            method = force.getNonbondedMethod()
            if method in recognized_periodic_methods[force_name]:
                return True

    # Nothing we recognize was found, so assume system was not periodic.
    return False

#=============================================================================================
# AlchemicalState
#=============================================================================================

class AlchemicalState(dict):
    """
    Alchemical state description.

    These parameters describe the parameters that affect computation of the energy.

    Attributes
    ----------
    lambda_restraints : float
        Scaling factor for remaining receptor-ligand relative restraint terms (to help keep ligand near protein).
    lambda_electrostatics : float
        Scaling factor for ligand charges, intrinsic Born radii, and surface area term.
    lambda_sterics : float
        Scaling factor for ligand sterics (Lennard-Jones and Halgren) interactions.
    labmda_torsions : float
        Scaling factor for ligand non-ring torsions.

    """
    def __init__(self, **kwargs):
        self['lambda_restraints'] = 0.0
        self['lambda_electrostatics'] = 1.0
        self['lambda_sterics'] = 1.0
        self['lambda_torsions'] = 1.0

        for key in kwargs.keys():
            self[key] = kwargs[key]

#=============================================================================================
# AbsoluteAlchemicalFactory
#=============================================================================================

class AbsoluteAlchemicalFactory(object):
    """
    Factory for generating OpenMM System objects that have been alchemically perturbed for absolute binding free energy calculation.

    Examples
    --------

    Create alchemical intermediates for default alchemical protocol for p-xylene in T4 lysozyme L99A in GBSA.

    >>> # Create a reference system.
    >>> from openmmtools import testsystems
    >>> complex = testsystems.LysozymeImplicit()
    >>> [reference_system, positions] = [complex.system, complex.positions]
    >>> # Create a factory to produce alchemical intermediates.
    >>> receptor_atoms = range(0,2603) # T4 lysozyme L99A
    >>> ligand_atoms = range(2603,2621) # p-xylene
    >>> factory = AbsoluteAlchemicalFactory(reference_system, ligand_atoms=ligand_atoms)
    >>> # Get the default protocol for 'denihilating' in complex in explicit solvent.
    >>> protocol = factory.defaultComplexProtocolImplicit()
    >>> # Create the perturbed systems using this protocol.
    >>> systems = factory.createPerturbedSystems(protocol)

    Create alchemical intermediates for default alchemical protocol for one water in a water box.

    >>> # Create a reference system.
    >>> from openmmtools import testsystems
    >>> waterbox = testsystems.WaterBox()
    >>> [reference_system, positions] = [waterbox.system, waterbox.positions]
    >>> # Create a factory to produce alchemical intermediates.
    >>> factory = AbsoluteAlchemicalFactory(reference_system, ligand_atoms=[0, 1, 2])
    >>> # Get the default protocol for 'denihilating' in solvent.
    >>> protocol = factory.defaultSolventProtocolExplicit()
    >>> # Create the perturbed systems using this protocol.
    >>> systems = factory.createPerturbedSystems(protocol)

    """

    # Factory initialization.
    def __init__(self, reference_system, ligand_atoms=list(), receptor_atoms=list(), annihilate_electrostatics=True, annihilate_sterics=False, test_positions=None, platform=None):
        """
        Initialize absolute alchemical intermediate factory with reference system.

        The reference system will not be modified when alchemical intermediates are generated.

        Parmeters
        ---------
        reference_system : simtk.openmm.System
            The reference system that is to be alchemically modified.
        ligand_atoms : list of int, optional, default = []
            List of atoms to be designated as 'ligand' for alchemical modification; everything else in system is considered the 'environment'.
        annihilateElectrostatics : bool
            If True, electrostatics should be annihilated, rather than decoupled.
        annihilateSterics : bool
            If True, sterics (Lennard-Jones or Halgren potential) will be annihilated, rather than decoupled.
        test_positions : simtk.unit.Quantity of dimension (natoms,3) with units compatible with nanometers, optional, default=None
            If provided, these coordinates will be used to test alchemically-modified system to ensure the potential energy is finite.
            If the potential energy is NaN, the energy for each force component will be computed for the Reference platform to aid in debugging.
        platform : simtk.openmm.Platform, optionl default=None
            If provided, this Platform will be used to check energies are finite.

        """

        # Store annihilation/decoupling information.
        self.annihilate_electrostatics = annihilate_electrostatics
        self.annihilate_sterics = annihilate_sterics

        # Store serialized form of reference system.
        self.reference_system = copy.deepcopy(reference_system)

        # Store copy of atom sets.
        self.ligand_atoms = copy.deepcopy(ligand_atoms)

        # Store atom sets
        self.ligand_atomset = set(self.ligand_atoms)

        # Create an alchemically-modified system to cache.
        self.alchemically_modified_system = self._createAlchemicallyModifiedSystem(self.reference_system)

        # Store information for use in aiding debugging of alchemical factory.
        self.test_positions = test_positions
        self.platform = platform
        if self.test_positions is not None:
            self._checkEnergyIsFinite(self.alchemically_modified_system, test_positions, platform=platform)

        # DEBUG: Write XML
        debug_write_xml = False
        if debug_write_xml:
            import os, os.path
            def write_file(filename, contents):
                with open(filename, 'w') as outfile:
                    outfile.write(contents)

            logger.info("Serializing to XML...")
            system_filename = os.path.join('setup', 'alchemical-system.xml')
            write_file(system_filename, openmm.XmlSerializer.serialize(self.alchemically_modified_system))

        return

    def _checkEnergyIsFinite(self, system, positions, platform=None):
        """
        Test that the potential energy is finite for the provided positions.
        If energy is NaN, compute energy from each force component on Reference platform to aid in debugging.

        Parameters
        ----------
        system : simtk.openmm.System
             The alchemical system to check.
        positions : simtk.unit.Quantity of dimension (natoms,3) with units compatible with nanometers
            Coordinates to use for energy test.
        platform : simtk.openmm.Platform, optional, default=None
            If specified, this platform will be used to compute test energy (but not by force component).

        """
        logger.debug("Checking alchemical system produces finite energies.")

        def compute_potential_energy(system, positions, platform=None, groups=-1):
            # Compute potential energy of reference system.
            integrator = openmm.VerletIntegrator(1.0 * unit.femtosecond)
            if platform is not None:
                # Compute potential energy with requested platform.
                context = openmm.Context(system, integrator, platform)
            else:
                # Use fastest available platform.
                context = openmm.Context(system, integrator)
            context.setPositions(positions)
            potential = context.getState(getEnergy=True,groups=groups).getPotentialEnergy()
            del context, integrator
            return potential

        # Compute potential energy error between reference and alchemical system.
        reference_potential = compute_potential_energy(self.reference_system, positions, platform)
        alchemical_potential = compute_potential_energy(system, positions, platform)
        energy_error = alchemical_potential - reference_potential

        # If alchemical potential energy is NaN, compute energy by component on Reference platform.
        if np.isnan(alchemical_potential / unit.kilocalories_per_mole):
            logger.debug("Energy for alchemically modified system is NaN.")
            logger.debug("Decomposing energy by force component on Reference platform:")

            # Assign all forces to separate components.
            system = self.alchemically_modified_system
            for force_index in range(system.getNumForces()):
                system.getForce(force_index).setForceGroup(force_index)

            # Compute potential energy for each force using Reference platform.
            reference_platform = openmm.Platform.getPlatformByName('Reference')
            for force_index in range(system.getNumForces()):
                groups = 1 << force_index # group index bitwise selector
                potential = compute_potential_energy(system, positions, reference_platform, groups)
                force_classname = system.getForce(force_index).__class__.__name__
                logger.debug("Force %5d / %5d [%24s] %12.3f kcal/mol" % (force_index, system.getNumForces(), force_classname, potential / unit.kilocalories_per_mole))

            # Clean up
            del context, integrator

            raise Exception("Energy for alchemically modified system is NaN.")

        # Return the energy error.
        logger.debug("Difference between alchemical and reference potential energy is %8.3f kcal/mol" % (energy_error / unit.kilocalories_per_mole))
        return energy_error

    @classmethod
    def defaultComplexProtocolImplicit(cls):
        """
        Return the default protocol for 'denihilating' a ligand in complex with a protein in implicit solvent.

        Returns
        -------
        alchemical_states : list of AlchemicalState
            The list of alchemical states defining the protocol.

        Notes
        -----
        The unrestrained, fully interacting system is always listed first.

        """

        alchemical_states = list()
        lambda_values = [1.0, 0.99, 0.98, 0.97, 0.96, 0.95, 0.94, 0.93, 0.92, 0.91, 0.90, 0.88, 0.86, 0.84, 0.81,
                         0.78, 0.74, 0.70, 0.60, 0.50, 0.40, 0.30, 0.20, 0.10, 0.05, 0.025, 0.01, 0.00]

        for lambda_value in lambda_values:
            alchemical_state = AlchemicalState()
            alchemical_state['lambda_restraints'] = 1.0
            alchemical_state['lambda_electrostatics'] = lambda_value
            alchemical_state['lambda_sterics'] = lambda_value
            alchemical_states.append(alchemical_state)

        return alchemical_states

    @classmethod
    def defaultComplexProtocolExplicit(cls):
        """
        Return the default protocol for 'denihilating' a ligand in complex with a protein in explicit solvent.

        Returns
{'fraction_non_alphanumeric': 0.14289555951460553, 'fraction_numerical': 0.009479310228915258, 'mean_word_length': 3.2384955121812222, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 165, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 6}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '455500', 'n_tokens_mistral': 7764, 'n_tokens_neox': 6962, 'n_words': 2353}