<commit_before>#!/usr/bin/python
# -*- coding: utf-8 -*-

################################################################################
#
#	RMG - Reaction Mechanism Generator
#
#	Copyright (c) 2002-2009 Prof. <NAME> (<EMAIL>) and the
#	RMG Team (<EMAIL>)
#
#	Permission is hereby granted, free of charge, to any person obtaining a
#	copy of this software and associated documentation files (the 'Software'),
#	to deal in the Software without restriction, including without limitation
#	the rights to use, copy, modify, merge, publish, distribute, sublicense,
#	and/or sell copies of the Software, and to permit persons to whom the
#	Software is furnished to do so, subject to the following conditions:
#
#	The above copyright notice and this permission notice shall be included in
#	all copies or substantial portions of the Software.
#
#	THE SOFTWARE IS PROVIDED 'AS IS', WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
#	IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
#	FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
#	AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
#	LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
#	FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
#	DEALINGS IN THE SOFTWARE.
#
################################################################################

"""
Contains classes for working with the reaction model generated by RMG.
"""

import log as logging
import numpy
import os

import constants
import settings
import reaction
import species
import unirxn.network

################################################################################

class ReactionModel:
	"""
	Represent a generic reaction model. A reaction model consists of `species`,
	a list of species, and `reactions`, a list of reactions.
	"""

	def __init__(self, species=None, reactions=None):
		self.species = species or []
		self.reactions = reactions or []

################################################################################

class CoreEdgeReactionModel:
	"""
	Represent a reaction model constructed using a rate-based screening
	algorithm. The species and reactions in the model itself are called the
	*core*; the species and reactions identified as candidates for inclusion in
	the model are called the *edge*. The attributes are:

	=========================  ==============================================================
	Attribute                  Description
	=========================  ==============================================================
	`core`                     The species and reactions of the current model core
	`edge`                     The species and reactions of the current model edge
	`absoluteTolerance`        The absolute tolerance used in the ODE/DAE solver
	`relativeTolerance`        The relative tolerance used in the ODE/DAE solver
	`fluxToleranceKeepInEdge`  The relative species flux below which species are discarded from the edge
	`fluxToleranceMoveToCore`  The relative species flux above which species are moved from the edge to the core
	`fluxToleranceInterrupt`   The relative species flux above which the simulation will halt
	`maximumEdgeSpecies`       The maximum number of edge species allowed at any time
	`termination`              A list of termination targets (i.e :class:`TerminationTime` and :class:`TerminationConversion` objects)
	`unirxnNetworks`           A list of unimolecular reaction networks (:class:`unirxn.network.Network` objects)
	`networkCount`             A counter for the number of unirxn networks created
	=========================  ==============================================================


	"""

	def __init__(self, core=None, edge=None):
		if core is None:
			self.core = ReactionModel()
		else:
			self.core = core
		if edge is None:
			self.edge = ReactionModel()
		else:
			self.edge = edge
		# The default tolerances mimic the original RMG behavior; no edge
		# pruning takes place, and the simulation is interrupted as soon as
		# a species flux higher than the validity
		self.fluxToleranceKeepInEdge = 0.0
		self.fluxToleranceMoveToCore = 1.0
		self.fluxToleranceInterrupt = 1.0
		self.absoluteTolerance = 1.0e-8
		self.relativeTolerance = 1.0e-4
		self.maximumEdgeSpecies = 1000000
		self.termination = []
		self.unirxnNetworks = []
		self.networkCount = 0

	def initialize(self, coreSpecies):
		"""
		Initialize a reaction model with a list `coreSpecies` of species to
		start out with.
		"""

		logging.info('')

		# Add all species present in input file to model core
		for spec in coreSpecies:
			self.enlarge(spec)

	def enlarge(self, newObject):
		"""
		Enlarge a reaction model by processing `newObject`. If `newObject` is a
		:class:`rmg.species.Species` object, then the species is moved from
		the edge to the core and reactions generated for that species, reacting
		with itself and with all other species in the model core. If `newObject`
		is a :class:`rmg.unirxn.network.Network` object, then reactions are
		generated for the species in the network with the largest leak flux.
		"""

		rxnList = []

		if isinstance(newObject, species.Species):

			newSpecies = newObject
			# Find reactions involving the new species as unimolecular reactant
			# or product (e.g. A <---> products)
			rxnList.extend(reaction.kineticsDatabase.getReactions([newSpecies]))
			# Find reactions involving the new species as bimolecular reactants
			# or products with itself (e.g. A + A <---> products)
			rxnList.extend(reaction.kineticsDatabase.getReactions([newSpecies, newSpecies]))
			# Find reactions involving the new species as bimolecular reactants
			# or products with other core species (e.g. A + B <---> products)
			for coreSpecies in self.core.species:
				if coreSpecies.reactive:
					rxnList.extend(reaction.kineticsDatabase.getReactions([newSpecies, coreSpecies]))

			# Add new species
			self.addSpeciesToCore(newSpecies)

		elif isinstance(newObject, unirxn.network.Network) and settings.unimolecularReactionNetworks:

			network = newObject
			# Determine the species with the maximum leak flux
			maxSpecies, maxSpeciesFlux = network.getMaximumLeakSpecies()
			network.explored.append(maxSpecies)
			# Find reactions involving the found species as unimolecular
			# reactant or product (e.g. A <---> products)
			rxnList = reaction.kineticsDatabase.getReactions([maxSpecies])
			# Don't find reactions involving the new species as bimolecular
			# reactants or products with itself (e.g. A + A <---> products)
			# Don't find reactions involving the new species as bimolecular
			# reactants or products with other core species (e.g. A + B <---> products)

		else:
			raise TypeError('Unable to use object %s to enlarge reaction model; expecting an object of class rmg.species.Species or rmg.unirxn.network.Network.' % newObject)

		# Add new reactions generated in above
		for rxn in rxnList:
			allSpeciesInCore = True
			for spec in rxn.reactants:
				if spec not in self.core.species:
					allSpeciesInCore = False
					if spec not in self.edge.species:
						self.addSpeciesToEdge(spec)
			for spec in rxn.products:
				if spec not in self.core.species:
					allSpeciesInCore = False
					if spec not in self.edge.species:
						self.addSpeciesToEdge(spec)
			# If pressure dependence is on, we only add reactions that are not unimolecular;
			# unimolecular reactions will be added after processing the associated networks
			if not settings.unimolecularReactionNetworks or not (
				rxn.isIsomerization() or rxn.isDissociation() or rxn.isAssociation()):
				if allSpeciesInCore:
					self.addReactionToCore(rxn)
				else:
					self.addReactionToEdge(rxn)
			else:
				# Update unimolecular reaction networks
				self.addReactionToUnimolecularNetworks(rxn)

		# Output current model size information after enlargement
		logging.info('')
		logging.info('After network enlargement:')
		logging.info('\tThe model core has %s species and %s reactions' % (len(self.core.species), len(self.core.reactions)))
		logging.info('\tThe model edge has %s species and %s reactions' % (len(self.edge.species), len(self.edge.reactions)))
		logging.info('')

	def addSpeciesToCore(self, spec):
		"""
		Add a species `spec` to the reaction model core (and remove from edge if
		necessary). This function also moves any reactions in the edge that gain
		core status as a result of this change in status to the core.
		"""

		if spec in self.core.species: return

		# Add the species to the core
		self.core.species.append(spec)

		# Add it to the cantera list
		spec.toCantera()

		if spec in self.edge.species:

			# If species was in edge, remove it
			self.edge.species.remove(spec)

			# Search edge for reactions that now contain only core species;
			# these belong in the model core and will be moved there
			rxnList = []
			for rxn in self.edge.reactions:
				allCore = True
				for reactant in rxn.reactants:
					if reactant not in self.core.species: allCore = False
				for product in rxn.products:
					if product not in self.core.species: allCore = False
				if allCore: rxnList.append(rxn)

			# Move any identified reactions to the core
			for rxn in rxnList:
				self.addReactionToCore(rxn)

	def addSpeciesToEdge(self, spec):
		"""
		Add a species `spec` to the reaction model edge.
		"""
		self.edge.species.append(spec)

	def removeSpeciesFromEdge(self, spec):
		"""
		Remove species `spec` from the reaction model edge.
		"""
		# remove the species
		self.edge.species.remove(spec)
		# identify any reactions it's involved in
		rxnList = []
		for rxn in self.edge.reactions:
			if spec in rxn.reactants or spec in rxn.products:
				rxnList.append(rxn)
		# remove those reactions
		for rxn in rxnList:
			self.edge.reactions.remove(rxn)

		# Remove the species from any unirxn networks it is in
		if settings.unimolecularReactionNetworks:
			networksToDelete = []
			for network in self.unirxnNetworks:
				if spec in network.getSpeciesList():
					# Delete all path reactions involving the species
					rxnList = []
					for rxn in network.pathReactions:
						if spec in rxn.reactants or spec in rxn.products:
							rxnList.append(rxn)
					for rxn in rxnList:
						network.pathReactions.remove(rxn)
					# Delete all net reactions involving the species
					rxnList = []
					for rxn in network.netReactions:
						if spec in rxn.reactants or spec in rxn.products:
							rxnList.append(rxn)
					for rxn in rxnList:
						network.netReactions.remove(rxn)
					# Delete all isomers involving the species
					isomerList = []
					for isomer in network.isomers:
						if spec in isomer.species:
							isomerList.append(isomer)
					for isomer in isomerList:
						network.isomers.remove(isomer)
					# If no remaining reactions, delete the network (actually
					# add to list of networks to be deleted in a subsequent
					# step)
					if len(network.pathReactions) == 0 and len(network.netReactions) == 0:
						networksToDelete.append(network)

			# Complete deletion of empty networks
			for network in networksToDelete:
				logging.debug('Deleting empty unirxn network %s' % network.id)
				self.unirxnNetworks.remove(network)



		# remove from the global list of species, to free memory
		species.speciesList.remove(spec)

	def addReactionToCore(self, rxn):
		"""
		Add a reaction `rxn` to the reaction model core (and remove from edge if
		necessary). This function assumes `rxn` has already been checked to
		ensure it is supposed to be a core reaction (i.e. all of its reactants
		AND all of its products are in the list of core species).
		"""
		self.core.reactions.append(rxn)
		if rxn in self.edge.reactions:
			self.edge.reactions.remove(rxn)

		# add it to the Cantera list
		rxn.toCantera()

	def addReactionToEdge(self, rxn):
		"""
		Add a reaction `rxn` to the reaction model edge. This function assumes
		`rxn` has already been checked to ensure it is supposed to be an edge
		reaction (i.e. all of its reactants OR all of its products are in the
		list of core species, and the others are in either the core or the
		edge).
		"""
		self.edge.reactions.append(rxn)

	def getLists(self):
		"""
		Return lists of all of the species and reactions in the core and the
		edge.
		"""
		speciesList = []
		speciesList.extend(self.core.species)
		speciesList.extend(self.edge.species)
		reactionList = []
		reactionList.extend(self.core.reactions)
		reactionList.extend(self.edge.reactions)
		return speciesList, reactionList

	def getStoichiometryMatrix(self):
		"""
		Return the stoichiometry matrix for all generated species and reactions.
		The id of each species and reaction is the corresponding row and column,
		respectively, in the matrix.
		"""
		speciesList, reactionList = self.getLists()
		from scipy import sparse
		stoichiometry = sparse.dok_matrix((species.speciesCounter, reaction.reactionCounter), float)
		for rxn in reactionList:
			j = rxn.id - 1
			specList = rxn.reactants[:]; specList.extend(rxn.products)
			for spec in specList:
				i = spec.id - 1
				nu = rxn.getStoichiometricCoefficient(spec)
				if nu != 0: stoichiometry[i,j] = nu
		return stoichiometry.tocsr()

	def getReactionRates(self, T, P, Ci):
		"""
		Return an array of reaction rates for each reaction in the model core
		and edge. The id of the reaction is the index into the vector.
		"""
		speciesList, reactionList = self.getLists()
		rxnRate = numpy.zeros(reaction.reactionCounter, float)
		for rxn in reactionList:
			j = rxn.id - 1
			rxnRate[j] = rxn.getRate(T, P, Ci)
		return rxnRate

	def addReactionToUnimolecularNetworks(self, newReaction):
		"""
		Given a newly-created :class:`Reaction` object `newReaction`, update the
		corresponding unimolecular reaction network. If no network exists, a new
		one is created. If the new reaction is an isomerization that connects two
		existing networks, the two networks are merged. This function is called
		whenever a new high-pressure limit edge reaction is created. Returns the
		network containing the new reaction.
		"""

		def getNetworkContainingSpecies(spec, networks):
			for network in networks:
				if network.containsSpecies(spec): return network
			return None

		# If the reaction is not unimolecular, then there's nothing to do
		# Return None because we haven't added newReaction to a network
		if not newReaction.isIsomerization() and not newReaction.isDissociation() and not newReaction.isAssociation():
			return None

		# Don't add reactions that are dissociation of a diatomic or association
		# to form a diatomic
		if newReaction.isDissociation() and len(newReaction.reactants[0].structure[0].atoms()) == 2:
			return None
		elif newReaction.isAssociation() and len(newReaction.products[0].structure[0].atoms()) == 2:
			return None

		# Find networks containing either the reactant or the product as a
		# unimolecular isomer
		reactantNetwork = None; productNetwork = None
		if newReaction.isIsomerization() or newReaction.isDissociation():
			reactantNetwork = getNetworkContainingSpecies(newReaction.reactants[0], self.unirxnNetworks)
		if newReaction.isIsomerization() or newReaction.isAssociation():
			productNetwork = getNetworkContainingSpecies(newReaction.products[0], self.unirxnNetworks)

		# Action depends on results of above search
		network = None
		if reactantNetwork is not None and productNetwork is not None and reactantNetwork is productNetwork:
			# Found the same network twice, so we don't need to merge
			# This can happend when a net reaction is later generated by RMG as a
			# path reaction
			network = reactantNetwork
		elif reactantNetwork is not None and productNetwork is not None:
			# Found two different networks, so we need to merge
			network = reactantNetwork
			network.merge(productNetwork)
			self.unirxnNetworks.remove(productNetwork)
		elif reactantNetwork is not None and productNetwork is None:
			# Found one network, so we add to it
			network = reactantNetwork
		elif reactantNetwork is None and productNetwork is not None:
			# Found one network, so we add to it
			network = productNetwork
		else:
			# Didn't find any networks, so we need to make a new one
			import unirxn.network
			self.networkCount += 1
			network = unirxn.network.Network(id=self.networkCount)
			self.unirxnNetworks.append(network)

		# Add the new reaction to whatever network was found/created above
		network.addPathReaction(newReaction)

		# Return the network that the reaction was added to
		return network

	def updateUnimolecularReactionNetworks(self):
		"""
		Iterate through all of the currently-existing unimolecular reaction
		networks, updating those that have been marked as invalid. In each update,
		the phenomonological rate coefficients :math:`k(T,P)` are computed for
		each net reaction in the network, and the resulting reactions added or
		updated.
		"""

		from unirxn.network import Isomer, UnirxnNetworkException
		from reaction import PDepReaction, makeNewPDepReaction
		from kinetics import ChebyshevKinetics, PDepArrheniusKinetics

		count = sum([1 for network in self.unirxnNetworks if not network.valid])
		logging.info('Updating %i modified unimolecular reaction networks...' % count)

		for network in self.unirxnNetworks:
			if not network.valid:

				logging.verbose('Updating unimolecular reaction network %s' % network.id)

				# Other inputs
				method, Tmin, Tmax, Tlist, Pmin, Pmax, Plist, grainSize, numGrains, model = settings.unimolecularReactionNetworks

				network.bathGas = [spec for spec in self.core.species if not spec.reactive][0]
				network.bathGas.expDownParam = 4.86 * 4184

				# Generate isomers
				for rxn in network.pathReactions:

					# Create isomer for the reactant
					isomer = None
					for isom in network.isomers:
						if all([spec in isom.species for spec in rxn.reactants]):
							isomer = isom
					if isomer is None:
						isomer = Isomer(rxn.reactants)
						network.isomers.append(isomer)
					rxn.reactant = isomer

					# Create isomer for the product
					isomer = None
					for isom in network.isomers:
						if all([spec in isom.species for spec in rxn.products]):
							isomer = isom
					if isomer is None:
						isomer = Isomer(rxn.products)
						network.isomers.append(isomer)
					rxn.product = isomer

				# Update list of explored isomers to include all species in core
				for isom in network.isomers:
					if isom.isUnimolecular():
						spec = isom.species[0]
						if spec not in network.explored:
							if spec in self.core.species:
								network.explored.append(spec)

				# Remove any isomers that aren't found in any path reactions
				# Ideally this block of code wouldn't be needed, but it's here
				# just in case
				isomerList = []
				for isomer in network.isomers:
					found = False
					for rxn in network.pathReactions:
						if rxn.reactant is isomer or rxn.product is isomer:
							found = True
							break
					if not found:
						isomerList.append(isomer)
				if len(isomerList) > 0:
					logging.debug('Removed %i isomer(s) from network %i.' % (len(isomerList), network.id))
					for isomer in isomerList: network.isomers.remove(isomer)

				# Sort isomers so that all unimolecular isomers come first
				isomers = [isom for isom in network.isomers if isom.isUnimolecular()]
				isomers.extend([isom for isom in network.isomers if isom.isMultimolecular()])
				network.isomers = isomers

				# Get list of species in network
				speciesList = network.getSpeciesList()

				# Calculate ground-state energy of all species in network
				# For now we assume that this is equal to the enthalpy of formation
				# of the species
				for spec in speciesList:
					spec.E0 = spec.getEnthalpy(T=298)

				# Determine isomer ground-state energies
				for isomer in network.isomers:
					isomer.E0 = sum([spec.E0 for spec in isomer.species])
				# Determine transition state ground-state energies of the reactions
				for rxn in network.pathReactions:
					E0 = sum([spec.E0 for spec in rxn.reactants])
					rxn.E0 = E0 + rxn.kinetics[0].getActivationEnergy(rxn.getEnthalpyOfReaction(T=298))

				# Shift network such that lowest-energy isomer has a ground state of 0.0
				network.shiftToZeroEnergy()

				# Determine energy grains
				Elist = network.determineEnergyGrains(grainSize, numGrains, max(Tlist))

				# Calculate density of states for all isomers in network
				network.calculateDensitiesOfStates(Elist)

				# Determine phenomenological rate coefficients
				K = network.calculateRateCoefficients(Tlist, Plist, Elist, method)

				# Generate PDepReaction objects
				for i, product in enumerate(network.isomers):
					for j, reactant in enumerate(network.isomers[0:i]):
						if numpy.any(K[:,:,i,j]):
							if not numpy.all(K[:,:,i,j]):
								raise UnirxnNetworkException('Zero rate coefficient encountered while updating network %s.' % network)
							# Find the path reaction
							netReaction = None
							for r in network.netReactions:
								if r.hasTemplate(reactant.species, product.species):
									netReaction = r
							# If path reaction does not already exist, make a new one
							if netReaction is None:
								netReaction = makeNewPDepReaction(reactant.species, product.species, network, None)
								network.netReactions.append(netReaction)
								self.addReactionToEdge(netReaction)
							# Set its kinetics using interpolation model
							if model[0].lower() == 'chebyshev':
								modelType, degreeT, degreeP = model
								chebyshev = ChebyshevKinetics()
								chebyshev.fitToData(Tlist, Plist, K[:,:,i,j], degreeT, degreeP)
								netReaction.kinetics = chebyshev
							elif model.lower() == 'pdeparrhenius':
								pDepArrhenius = PDepArrheniusKinetics()
								pDepArrhenius.fitToData(Tlist, Plist, K[:,:,i,j])
								netReaction.kinetics = pDepArrhenius
							else:
								pass

							# Set temperature and pressure ranges explicitly
							netReaction.kinetics.Tmin = Tmin
							netReaction.kinetics.Tmax = Tmax
							netReaction.kinetics.Pmin = Pmin
							netReaction.kinetics.Pmax = Pmax

							# Update cantera if this is a core reaction
							if netReaction in self.core.reactions:
								netReaction.toCantera()

				for spec in speciesList:
					del spec.E0
				for rxn in network.pathReactions:
					del rxn.reactant
					del rxn.product
					del rxn.E0

				network.valid = True

	def loadSeedMechanism(self, path):
		"""
		Loads a seed mechanism from the folder indicated by `path` into the
		core-edge reaction model.
		"""

		import os.path
		import quantities as pq
		import data
		import species
		import kinetics
		import reaction
		
		# Load the species data from the file species.txt
		# This file has the format of a standard RMG dictionary
		d = data.Dictionary()
		d.load(os.path.join(path, 'species.txt'))
		d.toStructure(addH=True)
		
		# Load the thermo data from the file thermo.txt
		# This file has the format of a standard RMG thermo library
		thermoData = species.ThermoDatabase()
		thermoData.load(os.path.join(path, 'species.txt'), '', os.path.join(path, 'thermo.txt'))
		# Populate the main primary thermo library with this thermo data
		# This will overwrite keys (but not values), so the order that the
		# seed mechanisms are loaded matters!
		for key, value in d.iteritems():
			species.thermoDatabase.primaryDatabase.dictionary[key] = value
		for key, value in thermoData.library.iteritems():
			species.thermoDatabase.primaryDatabase.library[key] = value

		# Create new species based on items in species.txt
		seedSpeciesDict = {}; seedSpeciesList = []
		for label, struct in d.iteritems():
			spec, isNew = species.makeNewSpecies(struct, label, reactive=True)
			seedSpeciesDict[label] = spec
			seedSpeciesList.append(spec)
		
		# Load the reactions from the file reaction.txt
		seedReactionList = []
		f = open(os.path.join(path, 'reactions.txt'), 'r')
		for line in f:
			line = data.removeCommentFromLine(line)
			line.strip()
			if len(line) > 0:
				items = line.split()
				if len(items) > 0:
					rxn = items[0:-6]

					# Extract reactants and products
					if '<=>' in rxn: arrow = rxn.index('<=>')
					elif '=>' in rxn: arrow = rxn.index('=>')
					else: raise IOError('No arrow found in reaction equation from line %s' % line)
					reactants = rxn[0:arrow:2]
					products = rxn[arrow+1::2]

					# Remove third body 'M' if present
					thirdBody = False
					if 'M' in reactants and 'M' in products:
						thirdBody = True
						reactants.remove('M')
						products.remove('M')
					
					# Convert strings to species objects
					reactants = [seedSpeciesDict[r] for r in reactants]
					products = [seedSpeciesDict[r] for r in products]

					# Process Arrhenius parameters
					order = len(reactants)
					if (thirdBody): order += 1
					Aunits = 'cm^%i/(mol^%i*s)' % (3*(order-1), order-1)
					A = float(pq.Quantity(float(items[-6]), Aunits).simplified)
					n = float(items[-5])			# dimensionless
					Ea = float(pq.Quantity(float(items[-4]), 'cal/mol').simplified)
					kin = [kinetics.ArrheniusKinetics(A=A, n=n, Ea=Ea)]

					# Create reaction object and add to list
					rxn = reaction.Reaction(id=0, reactants=reactants, products=products, family='seed', kinetics=kin, thirdBody=thirdBody)
					reaction.processNewReaction(rxn)
					seedReactionList.append(rxn)
					
		f.close()
		
		# Add species to core
		for spec in seedSpeciesList:
			self.addSpeciesToCore(spec)
		# Add reactions to core
		for rxn in seedReactionList:
			self.addReactionToCore(rxn)

################################################################################

class TemperatureModel:
	"""
	Represent a temperature profile. Currently the only implemented model is
	isothermal (constant temperature).
	"""

	def __init__(self):
		self.type = ''
		self.temperatures = []

	def isIsothermal(self):
		return self.type == 'isothermal'

	def setIsothermal(self, temperature):<commit_msg>Seed reactions now have reverse, and reactants and products are sorted.<commit_after>#!/usr/bin/python
# -*- coding: utf-8 -*-

################################################################################
#
#	RMG - Reaction Mechanism Generator
#
#	Copyright (c) 2002-2009 Prof. <NAME> (<EMAIL>) and the
#	RMG Team (<EMAIL>)
#
#	Permission is hereby granted, free of charge, to any person obtaining a
#	copy of this software and associated documentation files (the 'Software'),
#	to deal in the Software without restriction, including without limitation
#	the rights to use, copy, modify, merge, publish, distribute, sublicense,
#	and/or sell copies of the Software, and to permit persons to whom the
#	Software is furnished to do so, subject to the following conditions:
#
#	The above copyright notice and this permission notice shall be included in
#	all copies or substantial portions of the Software.
#
#	THE SOFTWARE IS PROVIDED 'AS IS', WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
#	IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
#	FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
#	AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
#	LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
#	FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
#	DEALINGS IN THE SOFTWARE.
#
################################################################################

"""
Contains classes for working with the reaction model generated by RMG.
"""

import log as logging
import numpy
import os

import constants
import settings
import reaction
import species
import unirxn.network

################################################################################

class ReactionModel:
	"""
	Represent a generic reaction model. A reaction model consists of `species`,
	a list of species, and `reactions`, a list of reactions.
	"""

	def __init__(self, species=None, reactions=None):
		self.species = species or []
		self.reactions = reactions or []

################################################################################

class CoreEdgeReactionModel:
	"""
	Represent a reaction model constructed using a rate-based screening
	algorithm. The species and reactions in the model itself are called the
	*core*; the species and reactions identified as candidates for inclusion in
	the model are called the *edge*. The attributes are:

	=========================  ==============================================================
	Attribute                  Description
	=========================  ==============================================================
	`core`                     The species and reactions of the current model core
	`edge`                     The species and reactions of the current model edge
	`absoluteTolerance`        The absolute tolerance used in the ODE/DAE solver
	`relativeTolerance`        The relative tolerance used in the ODE/DAE solver
	`fluxToleranceKeepInEdge`  The relative species flux below which species are discarded from the edge
	`fluxToleranceMoveToCore`  The relative species flux above which species are moved from the edge to the core
	`fluxToleranceInterrupt`   The relative species flux above which the simulation will halt
	`maximumEdgeSpecies`       The maximum number of edge species allowed at any time
	`termination`              A list of termination targets (i.e :class:`TerminationTime` and :class:`TerminationConversion` objects)
	`unirxnNetworks`           A list of unimolecular reaction networks (:class:`unirxn.network.Network` objects)
	`networkCount`             A counter for the number of unirxn networks created
	=========================  ==============================================================


	"""

	def __init__(self, core=None, edge=None):
		if core is None:
			self.core = ReactionModel()
		else:
			self.core = core
		if edge is None:
			self.edge = ReactionModel()
		else:
			self.edge = edge
		# The default tolerances mimic the original RMG behavior; no edge
		# pruning takes place, and the simulation is interrupted as soon as
		# a species flux higher than the validity
		self.fluxToleranceKeepInEdge = 0.0
		self.fluxToleranceMoveToCore = 1.0
		self.fluxToleranceInterrupt = 1.0
		self.absoluteTolerance = 1.0e-8
		self.relativeTolerance = 1.0e-4
		self.maximumEdgeSpecies = 1000000
		self.termination = []
		self.unirxnNetworks = []
		self.networkCount = 0

	def initialize(self, coreSpecies):
		"""
		Initialize a reaction model with a list `coreSpecies` of species to
		start out with.
		"""

		logging.info('')

		# Add all species present in input file to model core
		for spec in coreSpecies:
			self.enlarge(spec)

	def enlarge(self, newObject):
		"""
		Enlarge a reaction model by processing `newObject`. If `newObject` is a
		:class:`rmg.species.Species` object, then the species is moved from
		the edge to the core and reactions generated for that species, reacting
		with itself and with all other species in the model core. If `newObject`
		is a :class:`rmg.unirxn.network.Network` object, then reactions are
		generated for the species in the network with the largest leak flux.
		"""

		rxnList = []

		if isinstance(newObject, species.Species):

			newSpecies = newObject
			# Find reactions involving the new species as unimolecular reactant
			# or product (e.g. A <---> products)
			rxnList.extend(reaction.kineticsDatabase.getReactions([newSpecies]))
			# Find reactions involving the new species as bimolecular reactants
			# or products with itself (e.g. A + A <---> products)
			rxnList.extend(reaction.kineticsDatabase.getReactions([newSpecies, newSpecies]))
			# Find reactions involving the new species as bimolecular reactants
			# or products with other core species (e.g. A + B <---> products)
			for coreSpecies in self.core.species:
				if coreSpecies.reactive:
					rxnList.extend(reaction.kineticsDatabase.getReactions([newSpecies, coreSpecies]))

			# Add new species
			self.addSpeciesToCore(newSpecies)

		elif isinstance(newObject, unirxn.network.Network) and settings.unimolecularReactionNetworks:

			network = newObject
			# Determine the species with the maximum leak flux
			maxSpecies, maxSpeciesFlux = network.getMaximumLeakSpecies()
			network.explored.append(maxSpecies)
			# Find reactions involving the found species as unimolecular
			# reactant or product (e.g. A <---> products)
			rxnList = reaction.kineticsDatabase.getReactions([maxSpecies])
			# Don't find reactions involving the new species as bimolecular
			# reactants or products with itself (e.g. A + A <---> products)
			# Don't find reactions involving the new species as bimolecular
			# reactants or products with other core species (e.g. A + B <---> products)

		else:
			raise TypeError('Unable to use object %s to enlarge reaction model; expecting an object of class rmg.species.Species or rmg.unirxn.network.Network.' % newObject)

		# Add new reactions generated in above
		for rxn in rxnList:
			allSpeciesInCore = True
			for spec in rxn.reactants:
				if spec not in self.core.species:
					allSpeciesInCore = False
					if spec not in self.edge.species:
						self.addSpeciesToEdge(spec)
			for spec in rxn.products:
				if spec not in self.core.species:
					allSpeciesInCore = False
					if spec not in self.edge.species:
						self.addSpeciesToEdge(spec)
			# If pressure dependence is on, we only add reactions that are not unimolecular;
			# unimolecular reactions will be added after processing the associated networks
			if not settings.unimolecularReactionNetworks or not (
				rxn.isIsomerization() or rxn.isDissociation() or rxn.isAssociation()):
				if allSpeciesInCore:
					self.addReactionToCore(rxn)
				else:
					self.addReactionToEdge(rxn)
			else:
				# Update unimolecular reaction networks
				self.addReactionToUnimolecularNetworks(rxn)

		# Output current model size information after enlargement
		logging.info('')
		logging.info('After network enlargement:')
		logging.info('\tThe model core has %s species and %s reactions' % (len(self.core.species), len(self.core.reactions)))
		logging.info('\tThe model edge has %s species and %s reactions' % (len(self.edge.species), len(self.edge.reactions)))
		logging.info('')

	def addSpeciesToCore(self, spec):
		"""
		Add a species `spec` to the reaction model core (and remove from edge if
		necessary). This function also moves any reactions in the edge that gain
		core status as a result of this change in status to the core.
		"""

		if spec in self.core.species: return

		# Add the species to the core
		self.core.species.append(spec)

		# Add it to the cantera list
		spec.toCantera()

		if spec in self.edge.species:

			# If species was in edge, remove it
			self.edge.species.remove(spec)

			# Search edge for reactions that now contain only core species;
			# these belong in the model core and will be moved there
			rxnList = []
			for rxn in self.edge.reactions:
				allCore = True
				for reactant in rxn.reactants:
					if reactant not in self.core.species: allCore = False
				for product in rxn.products:
					if product not in self.core.species: allCore = False
				if allCore: rxnList.append(rxn)

			# Move any identified reactions to the core
			for rxn in rxnList:
				self.addReactionToCore(rxn)

	def addSpeciesToEdge(self, spec):
		"""
		Add a species `spec` to the reaction model edge.
		"""
		self.edge.species.append(spec)

	def removeSpeciesFromEdge(self, spec):
		"""
		Remove species `spec` from the reaction model edge.
		"""
		# remove the species
		self.edge.species.remove(spec)
		# identify any reactions it's involved in
		rxnList = []
		for rxn in self.edge.reactions:
			if spec in rxn.reactants or spec in rxn.products:
				rxnList.append(rxn)
		# remove those reactions
		for rxn in rxnList:
			self.edge.reactions.remove(rxn)

		# Remove the species from any unirxn networks it is in
		if settings.unimolecularReactionNetworks:
			networksToDelete = []
			for network in self.unirxnNetworks:
				if spec in network.getSpeciesList():
					# Delete all path reactions involving the species
					rxnList = []
					for rxn in network.pathReactions:
						if spec in rxn.reactants or spec in rxn.products:
							rxnList.append(rxn)
					for rxn in rxnList:
						network.pathReactions.remove(rxn)
					# Delete all net reactions involving the species
					rxnList = []
					for rxn in network.netReactions:
						if spec in rxn.reactants or spec in rxn.products:
							rxnList.append(rxn)
					for rxn in rxnList:
						network.netReactions.remove(rxn)
					# Delete all isomers involving the species
					isomerList = []
					for isomer in network.isomers:
						if spec in isomer.species:
							isomerList.append(isomer)
					for isomer in isomerList:
						network.isomers.remove(isomer)
					# If no remaining reactions, delete the network (actually
					# add to list of networks to be deleted in a subsequent
					# step)
					if len(network.pathReactions) == 0 and len(network.netReactions) == 0:
						networksToDelete.append(network)

			# Complete deletion of empty networks
			for network in networksToDelete:
				logging.debug('Deleting empty unirxn network %s' % network.id)
				self.unirxnNetworks.remove(network)



		# remove from the global list of species, to free memory
		species.speciesList.remove(spec)

	def addReactionToCore(self, rxn):
		"""
		Add a reaction `rxn` to the reaction model core (and remove from edge if
		necessary). This function assumes `rxn` has already been checked to
		ensure it is supposed to be a core reaction (i.e. all of its reactants
		AND all of its products are in the list of core species).
		"""
		self.core.reactions.append(rxn)
		if rxn in self.edge.reactions:
			self.edge.reactions.remove(rxn)

		# add it to the Cantera list
		rxn.toCantera()

	def addReactionToEdge(self, rxn):
		"""
		Add a reaction `rxn` to the reaction model edge. This function assumes
		`rxn` has already been checked to ensure it is supposed to be an edge
		reaction (i.e. all of its reactants OR all of its products are in the
		list of core species, and the others are in either the core or the
		edge).
		"""
		self.edge.reactions.append(rxn)

	def getLists(self):
		"""
		Return lists of all of the species and reactions in the core and the
		edge.
		"""
		speciesList = []
		speciesList.extend(self.core.species)
		speciesList.extend(self.edge.species)
		reactionList = []
		reactionList.extend(self.core.reactions)
		reactionList.extend(self.edge.reactions)
		return speciesList, reactionList

	def getStoichiometryMatrix(self):
		"""
		Return the stoichiometry matrix for all generated species and reactions.
		The id of each species and reaction is the corresponding row and column,
		respectively, in the matrix.
		"""
		speciesList, reactionList = self.getLists()
		from scipy import sparse
		stoichiometry = sparse.dok_matrix((species.speciesCounter, reaction.reactionCounter), float)
		for rxn in reactionList:
			j = rxn.id - 1
			specList = rxn.reactants[:]; specList.extend(rxn.products)
			for spec in specList:
				i = spec.id - 1
				nu = rxn.getStoichiometricCoefficient(spec)
				if nu != 0: stoichiometry[i,j] = nu
		return stoichiometry.tocsr()

	def getReactionRates(self, T, P, Ci):
		"""
		Return an array of reaction rates for each reaction in the model core
		and edge. The id of the reaction is the index into the vector.
		"""
		speciesList, reactionList = self.getLists()
		rxnRate = numpy.zeros(reaction.reactionCounter, float)
		for rxn in reactionList:
			j = rxn.id - 1
			rxnRate[j] = rxn.getRate(T, P, Ci)
		return rxnRate

	def addReactionToUnimolecularNetworks(self, newReaction):
		"""
		Given a newly-created :class:`Reaction` object `newReaction`, update the
		corresponding unimolecular reaction network. If no network exists, a new
		one is created. If the new reaction is an isomerization that connects two
		existing networks, the two networks are merged. This function is called
		whenever a new high-pressure limit edge reaction is created. Returns the
		network containing the new reaction.
		"""

		def getNetworkContainingSpecies(spec, networks):
			for network in networks:
				if network.containsSpecies(spec): return network
			return None

		# If the reaction is not unimolecular, then there's nothing to do
		# Return None because we haven't added newReaction to a network
		if not newReaction.isIsomerization() and not newReaction.isDissociation() and not newReaction.isAssociation():
			return None

		# Don't add reactions that are dissociation of a diatomic or association
		# to form a diatomic
		if newReaction.isDissociation() and len(newReaction.reactants[0].structure[0].atoms()) == 2:
			return None
		elif newReaction.isAssociation() and len(newReaction.products[0].structure[0].atoms()) == 2:
			return None

		# Find networks containing either the reactant or the product as a
		# unimolecular isomer
		reactantNetwork = None; productNetwork = None
		if newReaction.isIsomerization() or newReaction.isDissociation():
			reactantNetwork = getNetworkContainingSpecies(newReaction.reactants[0], self.unirxnNetworks)
		if newReaction.isIsomerization() or newReaction.isAssociation():
			productNetwork = getNetworkContainingSpecies(newReaction.products[0], self.unirxnNetworks)

		# Action depends on results of above search
		network = None
		if reactantNetwork is not None and productNetwork is not None and reactantNetwork is productNetwork:
			# Found the same network twice, so we don't need to merge
			# This can happend when a net reaction is later generated by RMG as a
			# path reaction
			network = reactantNetwork
		elif reactantNetwork is not None and productNetwork is not None:
			# Found two different networks, so we need to merge
			network = reactantNetwork
			network.merge(productNetwork)
			self.unirxnNetworks.remove(productNetwork)
		elif reactantNetwork is not None and productNetwork is None:
			# Found one network, so we add to it
			network = reactantNetwork
		elif reactantNetwork is None and productNetwork is not None:
			# Found one network, so we add to it
			network = productNetwork
		else:
			# Didn't find any networks, so we need to make a new one
			import unirxn.network
			self.networkCount += 1
			network = unirxn.network.Network(id=self.networkCount)
			self.unirxnNetworks.append(network)

		# Add the new reaction to whatever network was found/created above
		network.addPathReaction(newReaction)

		# Return the network that the reaction was added to
		return network

	def updateUnimolecularReactionNetworks(self):
		"""
		Iterate through all of the currently-existing unimolecular reaction
		networks, updating those that have been marked as invalid. In each update,
		the phenomonological rate coefficients :math:`k(T,P)` are computed for
		each net reaction in the network, and the resulting reactions added or
		updated.
		"""

		from unirxn.network import Isomer, UnirxnNetworkException
		from reaction import PDepReaction, makeNewPDepReaction
		from kinetics import ChebyshevKinetics, PDepArrheniusKinetics

		count = sum([1 for network in self.unirxnNetworks if not network.valid])
		logging.info('Updating %i modified unimolecular reaction networks...' % count)

		for network in self.unirxnNetworks:
			if not network.valid:

				logging.verbose('Updating unimolecular reaction network %s' % network.id)

				# Other inputs
				method, Tmin, Tmax, Tlist, Pmin, Pmax, Plist, grainSize, numGrains, model = settings.unimolecularReactionNetworks

				network.bathGas = [spec for spec in self.core.species if not spec.reactive][0]
				network.bathGas.expDownParam = 4.86 * 4184

				# Generate isomers
				for rxn in network.pathReactions:

					# Create isomer for the reactant
					isomer = None
					for isom in network.isomers:
						if all([spec in isom.species for spec in rxn.reactants]):
							isomer = isom
					if isomer is None:
						isomer = Isomer(rxn.reactants)
						network.isomers.append(isomer)
					rxn.reactant = isomer

					# Create isomer for the product
					isomer = None
					for isom in network.isomers:
						if all([spec in isom.species for spec in rxn.products]):
							isomer = isom
					if isomer is None:
						isomer = Isomer(rxn.products)
						network.isomers.append(isomer)
					rxn.product = isomer

				# Update list of explored isomers to include all species in core
				for isom in network.isomers:
					if isom.isUnimolecular():
						spec = isom.species[0]
						if spec not in network.explored:
							if spec in self.core.species:
								network.explored.append(spec)

				# Remove any isomers that aren't found in any path reactions
				# Ideally this block of code wouldn't be needed, but it's here
				# just in case
				isomerList = []
				for isomer in network.isomers:
					found = False
					for rxn in network.pathReactions:
						if rxn.reactant is isomer or rxn.product is isomer:
							found = True
							break
					if not found:
						isomerList.append(isomer)
				if len(isomerList) > 0:
					logging.debug('Removed %i isomer(s) from network %i.' % (len(isomerList), network.id))
					for isomer in isomerList: network.isomers.remove(isomer)

				# Sort isomers so that all unimolecular isomers come first
				isomers = [isom for isom in network.isomers if isom.isUnimolecular()]
				isomers.extend([isom for isom in network.isomers if isom.isMultimolecular()])
				network.isomers = isomers

				# Get list of species in network
				speciesList = network.getSpeciesList()

				# Calculate ground-state energy of all species in network
				# For now we assume that this is equal to the enthalpy of formation
				# of the species
				for spec in speciesList:
					spec.E0 = spec.getEnthalpy(T=298)

				# Determine isomer ground-state energies
				for isomer in network.isomers:
					isomer.E0 = sum([spec.E0 for spec in isomer.species])
				# Determine transition state ground-state energies of the reactions
				for rxn in network.pathReactions:
					E0 = sum([spec.E0 for spec in rxn.reactants])
					rxn.E0 = E0 + rxn.kinetics[0].getActivationEnergy(rxn.getEnthalpyOfReaction(T=298))

				# Shift network such that lowest-energy isomer has a ground state of 0.0
				network.shiftToZeroEnergy()

				# Determine energy grains
				Elist = network.determineEnergyGrains(grainSize, numGrains, max(Tlist))

				# Calculate density of states for all isomers in network
				network.calculateDensitiesOfStates(Elist)

				# Determine phenomenological rate coefficients
				K = network.calculateRateCoefficients(Tlist, Plist, Elist, method)

				# Generate PDepReaction objects
				for i, product in enumerate(network.isomers):
					for j, reactant in enumerate(network.isomers[0:i]):
						if numpy.any(K[:,:,i,j]):
							if not numpy.all(K[:,:,i,j]):
								raise UnirxnNetworkException('Zero rate coefficient encountered while updating network %s.' % network)
							# Find the path reaction
							netReaction = None
							for r in network.netReactions:
								if r.hasTemplate(reactant.species, product.species):
									netReaction = r
							# If path reaction does not already exist, make a new one
							if netReaction is None:
								netReaction = makeNewPDepReaction(reactant.species, product.species, network, None)
								network.netReactions.append(netReaction)
								self.addReactionToEdge(netReaction)
							# Set its kinetics using interpolation model
							if model[0].lower() == 'chebyshev':
								modelType, degreeT, degreeP = model
								chebyshev = ChebyshevKinetics()
								chebyshev.fitToData(Tlist, Plist, K[:,:,i,j], degreeT, degreeP)
								netReaction.kinetics = chebyshev
							elif model.lower() == 'pdeparrhenius':
								pDepArrhenius = PDepArrheniusKinetics()
								pDepArrhenius.fitToData(Tlist, Plist, K[:,:,i,j])
								netReaction.kinetics = pDepArrhenius
							else:
								pass

							# Set temperature and pressure ranges explicitly
							netReaction.kinetics.Tmin = Tmin
							netReaction.kinetics.Tmax = Tmax
							netReaction.kinetics.Pmin = Pmin
							netReaction.kinetics.Pmax = Pmax

							# Update cantera if this is a core reaction
							if netReaction in self.core.reactions:
								netReaction.toCantera()

				for spec in speciesList:
					del spec.E0
				for rxn in network.pathReactions:
					del rxn.reactant
					del rxn.product
					del rxn.E0

				network.valid = True

	def loadSeedMechanism(self, path):
		"""
		Loads a seed mechanism from the folder indicated by `path` into the
		core-edge reaction model.
		"""

		import os.path
		import quantities as pq
		import data
		import species
		import kinetics
		import reaction
		
		# Load the species data from the file species.txt
		# This file has the format of a standard RMG dictionary
		d = data.Dictionary()
		d.load(os.path.join(path, 'species.txt'))
		d.toStructure(addH=True)
		
		# Load the thermo data from the file thermo.txt
		# This file has the format of a standard RMG thermo library
		thermoData = species.ThermoDatabase()
		thermoData.load(os.path.join(path, 'species.txt'), '', os.path.join(path, 'thermo.txt'))
		# Populate the main primary thermo library with this thermo data
		# This will overwrite keys (but not values), so the order that the
		# seed mechanisms are loaded matters!
		for key, value in d.iteritems():
			species.thermoDatabase.primaryDatabase.dictionary[key] = value
		for key, value in thermoData.library.iteritems():
			species.thermoDatabase.primaryDatabase.library[key] = value

		# Create new species based on items in species.txt
		seedSpeciesDict = {}; seedSpeciesList = []
		for label, struct in d.iteritems():
			spec, isNew = species.makeNewSpecies(struct, label, reactive=True)
			seedSpeciesDict[label] = spec
			seedSpeciesList.append(spec)
		
		# Load the reactions from the file reaction.txt
		seedReactionList = []
		f = open(os.path.join(path, 'reactions.txt'), 'r')
		for line in f:
			line = data.removeCommentFromLine(line)
			line.strip()
			if len(line) > 0:
				items = line.split()
				if len(items) > 0:
					rxn = items[0:-6]

					# Extract reactants and products
					if '<=>' in rxn: arrow = rxn.index('<=>')
					elif '=>' in rxn: arrow = rxn.index('=>')
					else: raise IOError('No arrow found in reaction equation from line %s' % line)
					reactants = rxn[0:arrow:2]
					products = rxn[arrow+1::2]

					# Remove third body 'M' if present
					thirdBody = False
					if 'M' in reactants and 'M' in products:
						thirdBody = True
						reactants.remove('M')
						products.remove('M')
					
					# Convert strings to species objects
					reactants = [seedSpeciesDict[r] for r in reactants]
					products = [seedSpeciesDict[r] for r in products]
					reactants.sort()
					products.sort()

					# Process Arrhenius parameters
					order = len(reactants)
					if (thirdBody): order += 1
					Aunits = 'cm^%i/(mol^%i*s)' % (3*(order-1), order-1)
					A = float(pq.Quantity(float(items[-6]), Aunits).simplified)
					n = float(items[-5])			# dimensionless
					Ea = float(pq.Quantity(float(items[-4]), 'cal/mol').simplified)
					kin = [kinetics.ArrheniusKinetics(A=A, n=n, Ea=Ea)]

					# Create reaction object and add to list
					rxn = reaction.Reaction(id=0, reactants=reactants, products=products, family='seed', kinetics=kin, thirdBody=thirdBody)
					rxn.reverse = reaction.Reaction(id=0, reactants=products, products=reactants, family='seed', kinetics=None, thirdBody=thirdBody)
					rxn.reverse.reverse = rxn
					reaction.processNewReaction(rxn)
					seedReactionList.append(rxn)
					
		f.close()
		
		# Add species to core
		for spec in seedSpeciesList:
			self.addSpeciesToCore(spec)
		# Add reactions to core
		for rxn in seedReactionList:
			self.addReactionToCore(rxn)

################################################################################

class TemperatureModel:
	"""
	Represent a temperature profile. Currently the only implemented model is
	isothermal (constant temperature).
	"""

	def __init__(self):
		self.type = ''
		self.temperatures = []

	def isIsothermal(self):
		return self.type == 'isothermal'

	def setIsothermal(self, temperature):
{'fraction_non_alphanumeric': 0.10405369024268298, 'fraction_numerical': 0.003519095419406946, 'mean_word_length': 3.6427104722792607, 'pattern_counts': {'":': 0, '<': 25, '<?xml version=': 0, '>': 35, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, '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': '6469139', 'n_tokens_mistral': 18384, 'n_tokens_neox': 14915, 'n_words': 5364}