Simplify ODE construction

Author: ehermes

micki/model.py

@@ -563,6 +563,15 @@ def set_initial_conditions(self, U0):
                 newspecies.append(species)
         self._species = newspecies
 
+        # Also obtain a list of species that will be variables in the
+        # differential equations. This excludes fixed species and empty
+        # sites.
+        self._variable_species = []
+        for species in self._species:
+            if species.label not in self.fixed + [self.solvent]:
+                self._variable_species.append(species)
+        self.nvariables = len(self._variable_species)
+
         # Start with the incomplete user-provided initial conditions
         self.U0 = U0.copy()
 
@@ -614,11 +623,6 @@ def set_initial_conditions(self, U0):
             # Normalize the concentration of empty sites to match the
             # appropriate site ratio from the lattice.
             self.U0[name] = self.vactot[species] - occsites[species]
-#            if species not in U0:
-#                self.vactot[species] = 1.
-#                U0[species] = 1. - occsites[species]
-#            else:
-#                self.vactot[species] = U0[species] + occsites[species]
 
         # Populate dictionary of initial conditions for all species
         for name, species in self.species.items():
@@ -629,10 +633,6 @@ def set_initial_conditions(self, U0):
         # The number of variables that will be in our differential equations
         size = len(self._species)
 
-        # Initialize "mass matrix", which is a diagonal matrix with 1s for
-        # differential elements and 0s for algebraic elements (steady-state)
-        M = np.ones(size, dtype=int)
-
         # This creates a symbol for each species named modelparamX where X
         # is a three-digit numerical identifier that corresponds to its
         # position in the order of species
@@ -646,8 +646,8 @@ def set_initial_conditions(self, U0):
         for species in self._species:
             self.symbols_all.append(species.symbol)
             self.symbols_dict[species] = species.symbol
-            if species.label not in self.fixed and species.label != self.solvent:
-                self.symbols.append(species.symbol)
+
+        self.symbols = [species.symbol for species in self._variable_species]
 
         # subs converts a species symbol to either its initial value if
         # it is fixed or to a constraint (such as constraining the total
@@ -662,9 +662,6 @@ def set_initial_conditions(self, U0):
                 vacsymbols -= species.symbol
             subs[vacancy.symbol] = vacsymbols
 
-        # nsymbols is the size of the differential equations
-        self.nsymbols = len(self.symbols)
-
         # known_symbols keeps track of user-provided symbols that the
         # model has seen, so that symbols referring to species not in
         # the model can be later removed.
@@ -673,30 +670,43 @@ def set_initial_conditions(self, U0):
             known_symbols.add(species.symbol)
 
         # Create the final mass matrix of the proper dimensions
-        self.M = np.zeros((self.nsymbols, self.nsymbols), dtype=int)
+        self.M = np.eye(self.nvariables, dtype=int)
         # algvar tells the solver which variables are differential
         # and which are algebraic. It is the diagonal of the mass matrix.
-        algvar = np.zeros(self.nsymbols, dtype=float)
-        for i, symboli in enumerate(self.symbols_all):
-            for j, symbolj in enumerate(self.symbols):
-                if symboli == symbolj:
-                    self.M[j, j] = M[i]
-                    algvar[j] = M[i]
+        algvar = np.ones(self.nvariables, dtype=float)
+
+        # TODO: Make a way to set algebraic variables.
 
         # Initialize all rate expressions based on the above symbols
-        self._rate_init()
-        # f_sym is the SYMBOLIC master equation for all species
-        self.f_sym = []
+        nrxns = len(self._reactions)
+        # Array of symbolic rate expressions
+        self.rates = np.zeros(nrxns, dtype=object)
+        # Array of rate coefficients.
+        self.dfdr = np.zeros((self.nvariables, nrxns), dtype=int)
+
+        for j, rxn in enumerate(self._reactions):
+            rate_for = rxn.get_kfor(self.T, self.Asite, self.z)
+            rate_rev = rxn.get_krev(self.T, self.Asite, self.z)
+
+            for i, species in enumerate(self._variable_species):
+                rcount = rxn.reactants.species.count(species)
+                pcount = rxn.products.species.count(species)
+                self.dfdr[i, j] = -rcount + pcount
+
+            for species in self._species + self.vacancy:
+                rcount = rxn.reactants.species.count(species)
+                pcount = rxn.products.species.count(species)
+                if not isinstance(species, Electron):
+                    rate_for *= species.symbol**rcount
+                    rate_rev *= species.symbol**pcount
 
-        # TODO: Convert this to a matrix multiplication
-        for species in self._species:
-            f = 0
-            # DETAILED BALANCE NOTE: See discussion in _rate_init
-            if species.label not in self.fixed:
-                if species not in self.vacancy and species.label != self.solvent:
-                    for i, rate in enumerate(self.rates):
-                        f += self.rate_count[i][species] * rate
-                    self.f_sym.append(f)
+            # Overall reaction rate (flux)
+            self.rates[j] = rate_for
+            if rxn.reversible:
+                self.rates[j] -= rate_rev
+
+        # f_sym is the SYMBOLIC master equation for all species
+        self.f_sym = np.dot(self.dfdr, self.rates)
 
         # subs is a dictionary whose keys are internal species symbols and
         # whose values are the initial concentrations of that species if it
@@ -725,7 +735,7 @@ def set_initial_conditions(self, U0):
 
         # jac_sym is the SYMBOLIC Jacobian matrix, that is df/dc, where f is a
         # row of the master equation and c is a species.
-        self.jac_sym = np.zeros((self.nsymbols, self.nsymbols), dtype=object)
+        self.jac_sym = np.zeros((self.nvariables, self.nvariables), dtype=object)
         for i, f in enumerate(self.f_sym):
             for j, symbol in enumerate(self.symbols):
                 self.jac_sym[i, j] = sym.diff(f, symbol)
@@ -748,70 +758,10 @@ def set_initial_conditions(self, U0):
                     break
 
         # Pass initial values to the fortran module
-        self.finitialize(U0, 1e-10, [1e-10]*self.nsymbols, [], [], algvar)
+        self.finitialize(U0, 1e-10, [1e-10]*self.nvariables, [], [], algvar)
 
         self.initialized = True
 
-    def _rate_init(self):
-        # List of symbolic rate expressions
-        self.rates = []
-        # List of dicts. For a given reaction, the change in concentration
-        # of a species for "one unit" of reaction occuring in the forward
-        # direction
-        self.rate_count = []
-        # Keeps track of adsorption reactions
-        self.is_rate_ads = []
-
-        for reaction in self._reactions:
-            # Calculate kfor, krev, and keq
-            rate_for = reaction.get_kfor(self.T, self.Asite, self.z)
-            rate_rev = reaction.get_krev(self.T, self.Asite, self.z)
-
-            # IMPORTANT NOTE PERTAINING TO ABOVE:
-            # As far as the user can tell, all fluids (Gas, Liquid) are
-            # represented in units of concentration (i.e. mol/L) and all
-            # adsorbates are represented in coverage (i.e. N/M, where M is the
-            # total number of sites). INTERNALLY, this is not the case.
-            # Internally, ALL species are number fractions relative to the
-            # default number of catalytic sites (though this can be changed).
-            # This means that in order to achieve detailed balance, the
-            # reaction rates must be modulated by the volume. Further sections
-            # of the code that pertain to this behavior will be highlighted
-            # with "DETAILED BALANCE NOTE"
-
-            # Initialize dictionary for reaction stiochiometry to 0 for all
-            # species in the model
-            rate_count = {}
-            for species in self._species:
-                rate_count[species] = 0
-            for vacancy in self.vacancy:
-                rate_count[vacancy] = 0
-
-            # Reactants are consumed in the forward direction
-            for species in reaction.reactants:
-                # Multiply the rate by the species' symbol UNLESS it is an
-                # electron. All reactions are 0th order in electrons, even
-                # if they consume/produce electrons.
-                if not isinstance(species, Electron):
-                    rate_for *= species.symbol
-                rate_count[species] -= 1
-
-            # Products are created in the forward direction
-            for species in reaction.products:
-                if not isinstance(species, Electron):
-                    rate_rev *= species.symbol
-                rate_count[species] += 1
-
-            # Overall reaction rate (flux)
-            rate = rate_for
-            if reaction.reversible:
-                rate -= rate_rev
-
-            # Append all data to global rate lists
-            self.rates.append(rate)
-            self.rate_count.append(rate_count)
-            self.is_rate_ads.append(reaction.adsorption)
-
     def setup_execs(self):
         from micki.fortran import f90_template, pyf_template
         from numpy import f2py
@@ -820,13 +770,13 @@ def setup_execs(self):
         # concentrations provided by the differential equation solver inside
         # the Fortran code (that is, y_vec is an INPUT to the functions that
         # calculate the residual, Jacobian, and rate)
-        y_vec = sym.IndexedBase('yin', shape=(self.nsymbols,))
+        y_vec = sym.IndexedBase('yin', shape=(self.nvariables,))
         # Map y_vec elements (1-indexed, of course) onto 'modelparam' symbols
-        trans = {self.symbols[i]: y_vec[i + 1] for i in range(self.nsymbols)}
+        trans = {self.symbols[i]: y_vec[i + 1] for i in range(self.nvariables)}
         # Map string represntation of 'modelparam' symbols onto string
         # representation of y-vec elements
         str_trans = {}
-        for i in range(self.nsymbols):
+        for i in range(self.nvariables):
             str_trans[sym.fcode(self.symbols[i], source_format='free')] = \
                     sym.fcode(y_vec[i + 1], source_format='free')
 
@@ -841,7 +791,7 @@ def setup_execs(self):
         ratecode = []
 
         # Convert symbolic master equation into a valid Fortran string
-        for i in range(self.nsymbols):
+        for i in range(self.nvariables):
             fcode = sym.fcode(self.f_sym[i], source_format='free')
             # Replace modelparam symbols with their y_vec counterpart
             for key in str_list:
@@ -851,8 +801,8 @@ def setup_execs(self):
 
         # Effectively the same as above, except on the two-dimensional Jacobian
         # matrix.
-        for i in range(self.nsymbols):
-            for j in range(self.nsymbols):
+        for i in range(self.nvariables):
+            for j in range(self.nvariables):
                 expr = self.jac_sym[j, i]
                 # Unlike the residual, some elements of the Jacobian can be 0.
                 # We don't need to bother writing 'jac(x,y) = 0' a hundred
@@ -874,7 +824,7 @@ def setup_execs(self):
         # We insert all of the parameters of this differential equation into
         # the prewritten Fortran template, including the residual, Jacobian,
         # and rate expressions we just calculated.
-        program = f90_template.format(neq=self.nsymbols, nx=1,
+        program = f90_template.format(neq=self.nvariables, nx=1,
                                       nrates=len(self.rates),
                                       rescalc='\n'.join(rescode),
                                       jaccalc='\n'.join(jaccode),
@@ -893,7 +843,7 @@ def setup_execs(self):
 
         # Write the pertinent data into the temp directory
         with open(os.path.join(dname, pyfname), 'w') as f:
-            f.write(pyf_template.format(modname=modname, neq=self.nsymbols,
+            f.write(pyf_template.format(modname=modname, neq=self.nvariables,
                     nrates=len(self.rates)))
 
         # Compile the module with f2py
@@ -956,7 +906,7 @@ def _out_array_to_dict(self, U, dU, r):
         return Ui, dUi, ri
 
     def find_steady_state(self, dt=60, maxiter=2000, epsilon=1e-8):
-        t, U1, dU1, r1 = self.ffind_steady_state(self.nsymbols,
+        t, U1, dU1, r1 = self.ffind_steady_state(self.nvariables,
                                                  len(self.rates),
                                                  dt,
                                                  maxiter,
@@ -973,7 +923,7 @@ def find_steady_state(self, dt=60, maxiter=2000, epsilon=1e-8):
         return t, U, r
 
     def solve(self, t, ncp):
-        self.t, U1, dU1, r1 = self.fsolve(self.nsymbols,
+        self.t, U1, dU1, r1 = self.fsolve(self.nvariables,
                                           len(self.rates), ncp, t)
         self.U1 = U1.T
         self.dU1 = dU1.T
@@ -1017,7 +967,7 @@ def check_rates(self, U, epsilon=1e-6):
                                   RuntimeWarning, stacklevel=2)
 
     def copy(self, initialize=True):
-        newmodel = Model(self.T, self.Asite, self.z, self.nz, self.shape, self.lattice)
+        newmodel = Model(self.T, self.Asite, self.z, self.lattice)
         newmodel.add_reactions(self.reactions)
         newmodel.set_fixed(self.fixed)
         newmodel.set_solvent(self.solvent)