Changing Newton's method matrix solver in fortran to python

Question

I've written a code in fortran77 that is essentially a simplified form of the BLAG model of atmospheric CO2 and ocean pH. I'm now trying to write this code in python, but I'm having trouble finding a matrix solver in python that can solve something with 4 functions and 4 initial guesses for variables. The fortran version uses LUDCMP and LUBKSB (matrix solvers) to get the solution. I've looked around and it seems like python ought to have a more streamlined way to solve this, but whatever that way is, I don't know it. I've attached the entire code below, annotated, for people to peruse -- I'm essentially just trying to convert this fortran program to python. I am relatively new to python, so I apologize in advance if this is a really obvious question. Any pointers would be greatly appreciated!

    module coeffs
    real::alpha = 3.74e-2    !! alpha value
    real::K1 = 8.73e-7       !! constant
    real::K2 = 5.58e-10      !! constant
    real::KSP = 4.7e-7       !! constant
    real::Ca = 1.03e-2       !! calcium ions
    real::VOC = 3.6e19       !! volume of ocean
    real::Mair = 1.7e20      !! mass of carbon in air
    real::Mtot = 1.5e17      !! total mass of carbon
    real::alk=2.6614e-3      !! constant alkalinity
    end module

    use coeffs

    real,dimension(4,4)::jac
    real,dimension(4)::b,x,f,dx,fp
    integer::n=4
    integer::np=4
    integer,dimension(4)::indx
    real::d,chg,pulse,pH,Mat,Moc,H2CO3,HCO3,CO3,H
    real::eps=1.e-3           !! constant
    real::pCO2                !! partial pressure of CO2 in atmosphere

    x(1) = 3.64205e-4       !!! initial guess for pCO2
    x(2) = 1.36212e-5       !!! initial guess for H2CO3
    x(3) = 2.20503e-3       !!! initial guess for HCO3
    x(4) = 2.28155e-4       !!! initial guess for CO3

    do m = 1,1                              !!! START M LOOP
    print *
    print *, 'Enter pulse of CO2...'
    read *, pulse           !! added pulse of CO2 into atmosphere (simulating something like a rapid release of carbon, or an anthropogenic input)... 0.0 is standard for no pulse
    if(pulse<0.0)exit

    do k = 1,10                             !!! START K LOOP
    call func(x,f,n,pulse)
    do j = 1,4                              !!! START J LOOP
    xsave = x(j)
    delxj = eps * x(j)
    x(j) = x(j) + delxj

    call func(x,fp,n,pulse) 
    do i = 1,4                              !!! START I LOOP
    jac(i,j) = (fp(i) - f(i))/delxj
    end do                                  !!! END I LOOP
    x(j) = xsave
    end do                                  !!! END J LOOP

    b = -f

    call ludcmp(jac,n,np,indx,d)
    call lubksb(jac,n,np,indx,b)

    dx = b
    x = x + dx

    errmax = 0.
    do i = 1,4                              !!! START I LOOP
    err = abs(dx(i)/x(i))
    errmax = amax1(errmax,err)
    end do                                  !!! END I LOOP
    if(errmax<1.e-5)exit

    pCO2 = x(1) * 1.e6
    pH = -log10((K1 * x(2))/x(3))

    print *
    print *, 'k = ', k, 'pH = ', pH
    print *, 'pCO2 = ', pCO2
    print *, 'H2CO3 = ', x(2), 'HCO3 = ', x(3), 'CO3 = ', x(4)

    end do                                  !!! END K LOOP
    end do                                  !!! END M LOOP
    stop
    end

!!! subroutine containing functions to be solved

    subroutine func(x,f,n,pulse)
    use coeffs
    real,dimension(n)::x,f
    f(1) = x(1) - (x(2)/alpha)                    !! Henry's law
    f(2) = x(2) - (K2*x(3)*x(3))/(K1*x(4))        !! combination of 1st and 2nd dissociations for H2CO3 (eliminating H as a variable)
    f(3) = x(3) - alk+(2.*x(4))                   !! constant alkalinity
    f(4) = 1.e-19 * (Mtot+pulse-(Mair*x(1))-(VOC*(x(2)+x(3)+x(4))))     !! conservation of total carbon in the atmosphere/ocean system
    return
    end

!!! below this point is fortran77 matrix solver subroutines

  SUBROUTINE LUDCMP(A,N,NP,INDX,D)
  PARAMETER (NMAX=100,TINY=1.0E-20)
  DIMENSION A(NP,NP),INDX(N),VV(NMAX)
  D=1.
  DO 12 I=1,N
    AAMAX=0.
    DO 11 J=1,N
      IF (ABS(A(I,J)).GT.AAMAX) AAMAX=ABS(A(I,J))
11      CONTINUE
    IF (AAMAX.EQ.0.) PAUSE 'Singular matrix.'
    VV(I)=1./AAMAX
12    CONTINUE
  DO 19 J=1,N
    IF (J.GT.1) THEN
      DO 14 I=1,J-1
        SUM=A(I,J)
        IF (I.GT.1)THEN
          DO 13 K=1,I-1
            SUM=SUM-A(I,K)*A(K,J)
13            CONTINUE
          A(I,J)=SUM
        ENDIF
14        CONTINUE
    ENDIF
    AAMAX=0.
    DO 16 I=J,N
      SUM=A(I,J)
      IF (J.GT.1)THEN
        DO 15 K=1,J-1
          SUM=SUM-A(I,K)*A(K,J)
15          CONTINUE
        A(I,J)=SUM
      ENDIF
      DUM=VV(I)*ABS(SUM)
      IF (DUM.GE.AAMAX) THEN
        IMAX=I
        AAMAX=DUM
      ENDIF
16      CONTINUE
    IF (J.NE.IMAX)THEN
      DO 17 K=1,N
        DUM=A(IMAX,K)
        A(IMAX,K)=A(J,K)
        A(J,K)=DUM
17        CONTINUE
      D=-D
      VV(IMAX)=VV(J)
    ENDIF
    INDX(J)=IMAX
    IF(J.NE.N)THEN
      IF(A(J,J).EQ.0.)A(J,J)=TINY
      DUM=1./A(J,J)
      DO 18 I=J+1,N
        A(I,J)=A(I,J)*DUM
18        CONTINUE
    ENDIF
19    CONTINUE
  IF(A(N,N).EQ.0.)A(N,N)=TINY
  RETURN
  END

  SUBROUTINE LUBKSB(A,N,NP,INDX,B)
  DIMENSION A(NP,NP),INDX(N),B(N)
  II=0
  DO 12 I=1,N
    LL=INDX(I)
    SUM=B(LL)
    B(LL)=B(I)
    IF (II.NE.0)THEN
      DO 11 J=II,I-1
        SUM=SUM-A(I,J)*B(J)
11        CONTINUE
    ELSE IF (SUM.NE.0.) THEN
      II=I
    ENDIF
    B(I)=SUM
12    CONTINUE
  DO 14 I=N,1,-1
    SUM=B(I)
    IF(I.LT.N)THEN
      DO 13 J=I+1,N
        SUM=SUM-A(I,J)*B(J)
13        CONTINUE
    ENDIF
    B(I)=SUM/A(I,I)
14    CONTINUE
  RETURN
  END

Answer 1

If I'm understanding your Fortran code correctly, you should be able to solve your system of equations using scipy.optimize.fsolve. fsolve is a wrapper around two robust Fortran quasi-newton solvers (hybrd and hybrdj) for systems of non-linear equations. To solve your system in python, you could do something like the following:

from scipy.optimize import fsolve
import numpy as np

alpha = 3.74e-2    # alpha value
K1 = 8.73e-7       # constant
K2 = 5.58e-10      # constant
KSP = 4.7e-7       # constant
Ca = 1.03e-2       # calcium ions
VOC = 3.6e19       # volume of ocean
Mair = 1.7e20      # mass of carbon in air
Mtot = 1.5e17      # total mass of carbon
alk=2.6614e-3      # constant alkalinity

x0 = np.array([
    3.64205e-4,      # initial guess for pCO2
    1.36212e-5,      # initial guess for H2CO3
    2.20503e-3,      # initial guess for HCO3
    2.28155e-4       # initial guess for CO3
])

def objective_func(x, pulse):
    return np.array([
        x[0] - x[1] / alpha,
        x[1] - K2 * x[2] * x[2] / (K1 * x[3]),
        x[2] - alk + 2 * x[3],
        1e-19 * (Mtot + pulse - Mair * x[0] - VOC * (x[1] + x[2] + x[3]))
    ])

pulse = 0
sol, info, conv_flag, conv_msg = fsolve(objective_func, x0, args=(pulse,), full_output=True)

print(conv_msg)
print('Solution is: ', sol)

OUTPUT:

The solution converged.
Solution is:  [  3.64195924e-04   1.36209276e-05   2.20506331e-03   2.28168347e-04]

If you are only interested in the solution, you can leave out full_output=True or set it to False.