Envoltura de scipy.optimize.fsolve para sistemas de ecuaciones mas "idiomáticas"

Ecuaciones.py

#!/usr/bin/env python3
#-*- coding: utf-8 -*-
#
# Copyright (C) 2013 Cristian Hernán Schmidt <[email protected]>
#
# Ecuaciones.py is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Ecuaciones.py is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Ecuaciones.py. If not, see <http://www.gnu.org/licenses/>.

import re
import builtins
import collections

from scipy import ones
from scipy.optimize import fsolve


class Sistema(object):
    """ Envoltura de scipy.optimize.fsolve
        para sistemas de ecuaciones mas "idiomáticas"
    """

    _normalizar = (  # Formato: (patrón, reemplazo)
        (re.compile('(\w?)\s*(\()\s*'), r'\1\2'),                 # (
        (re.compile('\s*(\))\s*'), r'\1'),                        # )
        (re.compile('^\s*(.+?)\s*=\s*(.+?)\s*$'), r'\1 - (\2)'),  # =
        (re.compile('\s*(\+|-|\*\*|\*|//|/|\^)\s*'), r' \1 '),    # operadores
        (re.compile('\s{2,}'), ' '),                              # espacios
    )

    def __init__(self, vinc=None, vini=None, nom_vinc='__var_'):
        """ Parámetros:
              vinc      Lista de nombres que se considerarán incognitas, útil
                        para definir como incognita a variables presentes en
                        el espacio de nombres. Si no se asigna (aquí o en
                        solve) se utilizarán las que no esten presente en el
                        espacio de nombre que no sean parte de algún módulo o
                        función.
              vini      Valor o lista de valores a tomar como valor inicial por
                        fsolve. Si pasa un valor o una lista de un elemento y
                        hay mas incognitas se repite ese valor para todas las
                        incognitas.
              nom_vinc  Nombre de la variable vectorizada (incognitas), por
                        defecto '__var_', no puede usarse variables con este
                        nombre ya que quedarían ocultadas al momento de evaluar
                        las soluciones.
        """
        self.vini = vini or 0.5
        self.vinc = set(vinc) if hasattr(vinc, '__iter__') else set()
        self.nom_vinc = nom_vinc
        self.ecuaciones = []

    def add(self, *ecuaciones):
        """ Agrega una o mas ecuaciones al sistema, las ecuaciones
            que no tengan el signo '=' se consideran iguales a cero.
            Ejemplo:
                s = Sistema()
                # Forma 1
                s.add('a = b', 'b = 2')
                # Forma 2 (equivalente a la Forma 1)
                s.add('a = b')
                s.add('b = 2')
        """
        self.ecuaciones.extend([self.normalizar(e) for e in ecuaciones])

    def normalizar(self, ecuacion):
        """ Comprueba el número de '=', '(' y ')', y normaliza los espacios
        """
        i = ecuacion.count('=')
        if i > 1:
            e = 'Las ecuaciones debe tener 0 o 1 signo =, no {}'
            raise ValueError(e.format(i))
        a = ecuacion.count('(')
        c = ecuacion.count(')')
        if a != c:
            e = 'Los paréntesis deben estar balanceados (apertura {} cierre)'
            raise ValueError(e.format('>' if a > c else '<'))
        for patron, reemplazo in self._normalizar:
            ecuacion = patron.sub(reemplazo, ecuacion)
        return ecuacion.strip()

    def solve(self, vini=None, vinc=None, **kwargs):
        """ Vectoriza las incognitas, ejecuta fsolve y devuelve
            los resultado en un collections.namedtuple
            Las incognitas: vinc o las que no estén en el espacio de nombre
        """
        codigo = '[{}]'.format(', '.join(self.ecuaciones))
        necu = len(self.ecuaciones)

        # Variables disponibles
        vdis = set(dir(builtins))
        vdis.update(set(locals()))
        vdis.update(set(globals()))

        # Variables necesarias
        vnec = set()
        for e in self.ecuaciones:
            vnec.update(eval('lambda: {}'.format(e)).__code__.co_names)

        # Variables incognitas
        _vinc = set(vinc) if hasattr(vinc, '__iter__') else self.vinc
        if not _vinc:
            _vinc.update(vnec.difference(vdis))
        vinc = []
        ninc = 0
        for var in sorted(_vinc):
            patron = re.compile(r'(?<!\.)(\b{}\b)(?!(\.|\[|\())'.format(var))
            _var = r'{}[{}]'.format(self.nom_vinc, ninc)
            codigo, reemplazos = patron.subn(_var, codigo)
            if reemplazos:
                vinc.append(var)
                ninc += 1
        if not ninc:
            raise TypeError('Debe haber al menos una variable como incognita')
        if necu < ninc:
            e = ('Debe suministrar al menos el mismo número de ecuaciones '
                 '({}) que incognitas ({}).\n Es posible que sea un problema '
                 'de scope; las variables consideradas incognitas son:\n  {}')
            raise TypeError(e.format(necu, ninc, ', '.join(vinc)))
        vini = vini or self.vini
        if not hasattr(vini, '__len__'):
            vini = (vini,)
        if len(vini) == 1 and ninc > 1:
            vini = ones(ninc) * vini[0]
        elif len(vini) != ninc:
            e = ('El número de incognitas debe ser igual al número de valores '
                 'iniciales (si es una lista de mas de un elemento)')
            raise TypeError(e)
        # Optimización
        funcion = eval('lambda {}: {}'.format(self.nom_vinc, codigo))
        raices = fsolve(funcion, vini, **kwargs)
        Resultados = collections.namedtuple('Resultados', vinc)
        return Resultados(*raices)


def ejemplos():
    lista = (
        r'''
        s = Sistema()
        e = 'a - round(math.pi)'
        s.add(e)
        r = s.solve()

        print(('Prueba con modulos. La ecuación "{} = 0" '
            'da como resultado "a = {}"').format(e, r.a))
        ''',
        ####
        r'''
        # Libro: Química - Raymond Chang - 6ed - Pag 608
        global Co, ka, kw  # Necesario porque se ejecuta desde una función

        ka = 7.1e-4                         # HF <=> F⁻ + H⁺
        kw = 1e-14                          # H2O <=> H⁺ + OH⁻
        equilibrio = Sistema()
        equilibrio.add(
            'ka = H * F / HF',              # Cte. de acidez
            'kw = H * OH',                  # Cte. de autodisociación del agua
            'Co = F + HF',                  # Balance de materia
            'H = OH + F',                   # Balance de carga
        )
        for Co in (0.5, 0.05):              # Molar
            Ce = equilibrio.solve(vini=Co)
            pH = -math.log10(Ce.H)

            print(('Una disolución {} M de ácido fluorhídrico '
                'tiene un pH de {:.2f}').format(Co, pH))
        ''',
        ####
        r'''
        # Libro: Química - Raymond Chang - 6ed - Pag 618
        global Co, ka1, ka2, kw  # Necesario porque se ejecuta desde una función

        Co = 0.10                           # Molar
        ka1 = 6.5e-2                        # C2H2O4 <=> C2HO4⁻ + H⁺
        ka2 = 6.1e-5                        # C2HO4⁻ <=> C2O4⁼ + H⁺
        kw = 1e-14                          # H2O <=> H⁺ + OH⁻
        equilibrio = Sistema()
        equilibrio.add(
            'ka1 = H * C2HO4 / C2H2O4',     # Cte. de acidez 1
            'ka2 = H * C2O4 / C2HO4',       # Cte. de acidez 2
            'kw = H * OH',                  # Cte. de autodisociación del agua
            'Co = C2H2O4 + C2HO4 + C2O4',   # Balance de materia
            'H = OH + C2HO4 + C2O4',        # Balance de carga
        )
        Ce = equilibrio.solve(vini=Co, xtol=1e-10)
        pH = -math.log10(Ce.H)
        for c in ('C2H2O4', 'C2HO4', 'C2O4', 'H', 'OH'):
            print('{} = {:.3g}'.format(c, getattr(Ce, c)), end='  ')

        print(('\nUna disolución {} M de ácido oxálico '
            'tiene un pH de {:.2f}').format(Co, pH))
        ''',
    )
    line = '─' * 86
    from functools import partial
    cod1 = partial(re.compile('(\n) {4}').sub, r'\1')
    cod2 = partial(re.compile('(\n) {8}').sub, r'\1')
    for i, c in enumerate(lista, 1):
        print('{}\n\nEjemplo {}\nCódigo:{}\nResultado:'.format(
            line, i, cod1(c)))
        exec(cod2(c))
        print('')
    print(line)


if __name__ == "__main__":
    import math
    ejemplos()