added metallicity function

equilibrate/cython/ChemEquiAnalysis_pxd.pxd

@@ -20,6 +20,9 @@ cdef extern void chemequianalysis_solve_wrapper(ChemEquiAnalysis *ptr, double *P
                          cbool *molfracs_atoms_present, int *molfracs_atoms_dim, double *molfracs_atoms, 
                          cbool *molfracs_species_present, int *molfracs_species_dim, double *molfracs_species, 
                          cbool *converged, char *err)
+cdef extern void chemequianalysis_solve_metallicity_wrapper(ChemEquiAnalysis *ptr, double *P, double *T,
+                         double *metallicity, cbool *CtoO_present, double *CtoO, 
+                         cbool *converged, char *err)
 
 # Getters and setters
 
@@ -35,6 +38,10 @@ cdef extern void chemequianalysis_gas_names_get(ChemEquiAnalysis *ptr, int *dim1
 cdef extern void chemequianalysis_condensate_names_get_size(ChemEquiAnalysis *ptr, int *dim1)
 cdef extern void chemequianalysis_condensate_names_get(ChemEquiAnalysis *ptr, int *dim1, char* arr)
 
+cdef extern void chemequianalysis_molfracs_atoms_sun_get_size(ChemEquiAnalysis *ptr, int *dim1)
+cdef extern void chemequianalysis_molfracs_atoms_sun_get(ChemEquiAnalysis *ptr, int *dim1, double *arr)
+cdef extern void chemequianalysis_molfracs_atoms_sun_set(ChemEquiAnalysis *ptr, int *dim1, double *arr)
+
 cdef extern void chemequianalysis_molfracs_atoms_get_size(ChemEquiAnalysis *ptr, int *dim1)
 cdef extern void chemequianalysis_molfracs_atoms_get(ChemEquiAnalysis *ptr, int *dim1, double *arr)
 

equilibrate/cython/_equilibrate.pyx

@@ -137,6 +137,47 @@ cdef class ChemEquiAnalysis:
 
     return converged
 
+  def solve_metallicity(self, double P, double T, double metallicity, CtoO = None):
+    """Computes chemical equilibrium given an input metallicity and, optionally,
+    a C/O ratio. If successful, then the equilibrium composition will be stored in a number
+    of attributes (e.g., self%molfracs_species).
+
+    Parameters
+    ----------
+    P : double
+        Pressure in dynes/cm^2
+    T : double
+        Temperature in Kelvin
+    metallicity : double
+        Metallicity relative to the Sun
+    CtoO : double, optional
+        The C/O ratio relative to solar. CtoO = 1 would be the same composition as solar.
+
+    Results
+    -------
+    converged : bool
+        If true, then the calculation successfully achieved chemical equilibrium
+        to within the specified tolerances.
+    """
+
+    cdef double CtoO_ = 1.0
+    cdef cbool CtoO_present = False
+    if CtoO is not None:
+      CtoO_present = True
+      CtoO_ = CtoO
+
+    cdef cbool converged
+    cdef char err[ERR_LEN+1]
+
+    cea_pxd.chemequianalysis_solve_metallicity_wrapper(self._ptr, &P, &T,
+                         &metallicity, &CtoO_present, &CtoO_, 
+                         &converged, err)
+
+    if len(err.strip()) > 0:
+      raise EquilibrateException(err.decode("utf-8").strip())
+
+    return converged
+
   property atoms_names:
     "List. Names of atoms"
     def __get__(self):
@@ -173,6 +214,21 @@ cdef class ChemEquiAnalysis:
       cea_pxd.chemequianalysis_condensate_names_get(self._ptr, &dim1, <char *>arr_c.data)
       return c2stringarr(arr_c, S_STR_LEN, dim1)
 
+  property molfracs_atoms_sun:
+    "ndarray[double,ndim=1]. Assumed mole fractions of each atom in the Sun."
+    def __get__(self):
+      cdef int dim1
+      cea_pxd.chemequianalysis_molfracs_atoms_sun_get_size(self._ptr, &dim1)
+      cdef ndarray arr = np.empty(dim1, np.double)
+      cea_pxd.chemequianalysis_molfracs_atoms_sun_get(self._ptr, &dim1, <double *>arr.data)
+      return arr
+    def __set__(self,ndarray[double, ndim=1] arr):
+      cdef int dim1
+      cea_pxd.chemequianalysis_molfracs_atoms_sun_get_size(self._ptr, &dim1)
+      if arr.shape[0] != dim1:
+        raise EquilibrateException("Input array is the wrong size.")
+      cea_pxd.chemequianalysis_molfracs_atoms_sun_set(self._ptr, &dim1, <double *>arr.data)
+
   property molfracs_atoms:
     "ndarray[double,ndim=1]. Mole fractions of each atom."
     def __get__(self):

equilibrate/fortran/equilibrate_c_api.f90

@@ -130,6 +130,37 @@ subroutine chemequianalysis_solve_wrapper(ptr, P, T, &
 
   end subroutine
 
+  subroutine chemequianalysis_solve_metallicity_wrapper(ptr, P, T, &
+                                                        metallicity, CtoO_present, CtoO, &
+                                                        converged, err) bind(c)
+    use equilibrate, only: ChemEquiAnalysis
+    type(c_ptr), value, intent(in) :: ptr
+    real(c_double), intent(in) :: P
+    real(c_double), intent(in) :: T
+    real(c_double), intent(in) :: metallicity
+    logical(c_bool), intent(in) :: CtoO_present
+    real(c_double), intent(in) :: CtoO
+    logical(c_bool), intent(out) :: converged
+    character(kind=c_char), intent(out) :: err(err_len+1)
+
+    character(:), allocatable :: err_f
+    type(ChemEquiAnalysis), pointer :: cea
+
+    call c_f_pointer(ptr, cea)
+
+    if (CtoO_present) then
+      converged = cea%solve_metallicity(P, T, metallicity, CtoO=CtoO, err=err_f)
+    else
+      converged = cea%solve_metallicity(P, T, metallicity, err=err_f)
+    endif
+
+    err(1) = c_null_char
+    if (allocated(err_f)) then
+      call copy_string_ftoc(err_f, err)
+    endif
+
+  end subroutine
+
   !~~ Getters and setters ~~!
 
   subroutine chemequianalysis_atoms_names_get_size(ptr, dim1) bind(c)
@@ -248,6 +279,35 @@ subroutine chemequianalysis_condensate_names_get(ptr, dim1, arr) bind(c)
 
   end subroutine
 
+  subroutine chemequianalysis_molfracs_atoms_sun_get_size(ptr, dim1) bind(c)
+    use equilibrate, only: ChemEquiAnalysis
+    type(c_ptr), value, intent(in) :: ptr
+    integer(c_int), intent(out) :: dim1
+    type(ChemEquiAnalysis), pointer :: cea
+    call c_f_pointer(ptr, cea)
+    dim1 = size(cea%molfracs_atoms_sun)
+  end subroutine
+
+  subroutine chemequianalysis_molfracs_atoms_sun_get(ptr, dim1, arr) bind(c)
+    use equilibrate, only: ChemEquiAnalysis
+    type(c_ptr), value, intent(in) :: ptr
+    integer(c_int), intent(in) :: dim1
+    real(c_double), intent(out) :: arr(dim1)
+    type(ChemEquiAnalysis), pointer :: cea
+    call c_f_pointer(ptr, cea)
+    arr = cea%molfracs_atoms_sun
+  end subroutine
+
+  subroutine chemequianalysis_molfracs_atoms_sun_set(ptr, dim1, arr) bind(c)
+    use equilibrate, only: ChemEquiAnalysis
+    type(c_ptr), value, intent(in) :: ptr
+    integer(c_int), intent(in) :: dim1
+    real(c_double), intent(out) :: arr(dim1)
+    type(ChemEquiAnalysis), pointer :: cea
+    call c_f_pointer(ptr, cea)
+    cea%molfracs_atoms_sun = arr
+  end subroutine
+
   subroutine chemequianalysis_molfracs_atoms_get_size(ptr, dim1) bind(c)
     use equilibrate, only: ChemEquiAnalysis
     type(c_ptr), value, intent(in) :: ptr

src/equilibrate.f90

@@ -16,6 +16,9 @@ module equilibrate
     !> Describes the composition of each species (i.e., how many atoms)
     real(dp), allocatable, private :: species_composition(:,:)
 
+    !> Composition of the Sun.
+    real(dp), allocatable :: molfracs_atoms_sun(:)
+
     ! All below are results of an equilibrium solve.
 
     real(dp), allocatable :: molfracs_atoms(:) !! Mole fractions of each atom (size(atoms_names))
@@ -42,6 +45,7 @@ module equilibrate
 
   contains
     procedure :: solve
+    procedure :: solve_metallicity
   end type
   interface ChemEquiAnalysis
     module procedure :: create_ChemEquiAnalysis
@@ -155,12 +159,26 @@ function create_ChemEquiAnalysis(thermopath, atoms, species, err) result(cea)
       enddo
     enddo
 
+    block
+      use equilibrate_const, only: atoms_sun, molfracs_atoms_sun
+      integer :: ind
+      ! Default composition of Sun
+      allocate(cea%molfracs_atoms_sun(size(cea%atoms_names)))
+      cea%molfracs_atoms_sun = 1.0e-50_dp
+      do i = 1,size(cea%atoms_names)
+        ind = findloc(atoms_sun, cea%atoms_names(i), 1)
+        if (ind /= 0) then
+          cea%molfracs_atoms_sun(i) = molfracs_atoms_sun(ind)
+        endif
+      enddo
+      cea%molfracs_atoms_sun = cea%molfracs_atoms_sun/sum(cea%molfracs_atoms_sun)
+    endblock
+
   end function
 
   !> Computes chemical equilibrium given input atom or species mole fractions.
   !> If successful, then the equilibrium composition will be stored in a number
-  !> of attributes (e.g., self%molfracs_species). If unsuccessful, then the variable
-  !> `err` is allocated with an error message.
+  !> of attributes (e.g., self%molfracs_species).
   function solve(self, P, T, molfracs_atoms, molfracs_species, err) result(converged)
     class(ChemEquiAnalysis), intent(inout) :: self
     real(dp), intent(in) :: P !! Pressure in dynes/cm^2
@@ -273,4 +291,66 @@ function solve(self, P, T, molfracs_atoms, molfracs_species, err) result(converg
 
   end function
 
+  !> Computes chemical equilibrium given an input metallicity and, optionally,
+  !> a C/O ratio. If successful, then the equilibrium composition will be stored in a number
+  !> of attributes (e.g., self%molfracs_species).
+  function solve_metallicity(self, P, T, metallicity, CtoO, err) result(converged)
+    class(ChemEquiAnalysis), intent(inout) :: self
+    real(dp), intent(in) :: P !! Pressure in dynes/cm^2
+    real(dp), intent(in) :: T !! Temperature in Kelvin
+    real(dp), intent(in) :: metallicity !! Metallicity relative to the Sun
+    !> The C/O ratio relative to solar. CtoO = 1 would be the same
+    !> composition as solar.
+    real(dp), optional, intent(in) :: CtoO
+    character(:), allocatable, intent(out) :: err
+    logical :: converged
+
+    real(dp), allocatable :: molfracs_atoms(:)
+    integer :: i, indC, indO
+    real(dp) :: x, a
+
+    if (metallicity <= 0) then
+      err = '"metallicity" must be greater than 0.'
+      return
+    endif
+
+    ! Compute atoms based on the input metallicity.
+    molfracs_atoms = self%molfracs_atoms_sun
+    do i = 1,size(molfracs_atoms)
+      if (self%atoms_names(i) /= 'H' .or. self%atoms_names(i) /= 'He') then
+        molfracs_atoms(i) = self%molfracs_atoms_sun(i)*metallicity
+      endif
+    enddo
+    molfracs_atoms = molfracs_atoms/sum(molfracs_atoms)
+
+    ! Adjust C/O ratio, if specified
+    if (present(CtoO)) then
+      if (CtoO <= 0) then
+        err = '"CtoO" must be greater than 0.'
+        return
+      endif
+
+      indC = findloc(self%atoms_names, 'C', 1)
+      if (indC == 0) then
+        err = 'C must be an atom if CtoO is specified.'
+        return
+      endif
+
+      indO = findloc(self%atoms_names, 'O', 1)
+      if (indO == 0) then
+        err = 'O must be an atom if CtoO is specified.'
+        return
+      endif
+
+      x = CtoO*(molfracs_atoms(indC)/molfracs_atoms(indO))
+      a = (x*molfracs_atoms(indO) - molfracs_atoms(indC))/(1 + x)
+      molfracs_atoms(indC) = molfracs_atoms(indC) + a
+      molfracs_atoms(indO) = molfracs_atoms(indO) - a
+    endif
+
+    converged = self%solve(P, T, molfracs_atoms=molfracs_atoms, err=err)
+    if (allocated(err)) return
+
+  end function
+
 end module

src/equilibrate_const.f90

@@ -39,4 +39,41 @@ module equilibrate_const
   ! For thermo arrays
   integer, parameter      :: N_coeffs = 10, N_temps = 10
 
+  ! Composition of the Sun from Aspund et al. (2009)
+  ! https://doi.org/10.1146/annurev.astro.46.060407.145222
+  character(*), parameter :: atoms_sun(*) = &
+  [ &
+    'H ', 'He', 'Li', 'Be', 'B ', &
+    'C ', 'N ', 'O ', 'F ', 'Ne', &
+    'Na', 'Mg', 'Al', 'Si', 'P ', &
+    'S ', 'Cl', 'Ar', 'K ', 'Ca', &
+    'Sc', 'Ti', 'V ', 'Cr', 'Mn', &
+    'Fe', 'Co', 'Ni', 'Cu', 'Zn', &
+    'Ga', 'Ge', 'Kr', 'Rb', 'Sr', &
+    'Y ', 'Zr', 'Nb', 'Mo', 'Ru', &
+    'Rh', 'Pd', 'Ag', 'In', 'Sn', &
+    'Xe', 'Ba', 'La', 'Ce', 'Pr', &
+    'Nd', 'Sm', 'Eu', 'Gd', 'Tb', &
+    'Dy', 'Ho', 'Er', 'Tm', 'Yb', &
+    'Lu', 'Hf', 'W ', 'Os', 'Ir', &
+    'Au', 'Tl', 'Pb', 'Th' &
+  ]
+  real(dp), parameter :: molfracs_atoms_sun(*) = &
+  [ &
+    9.206789e-01_dp, 7.836248e-02_dp, 1.033019e-11_dp, 2.208555e-11_dp, 4.614325e-10_dp, &
+    2.478039e-04_dp, 6.224553e-05_dp, 4.509290e-04_dp, 3.342783e-08_dp, 7.836248e-05_dp, &
+    1.599957e-06_dp, 3.665289e-05_dp, 2.594826e-06_dp, 2.979258e-05_dp, 2.366509e-07_dp, &
+    1.213691e-05_dp, 2.911442e-07_dp, 2.312641e-06_dp, 9.865252e-08_dp, 2.014226e-06_dp, &
+    1.300493e-09_dp, 8.205559e-08_dp, 7.836248e-09_dp, 4.018909e-07_dp, 2.478039e-07_dp, &
+    2.911442e-05_dp, 8.997217e-08_dp, 1.527947e-06_dp, 1.425963e-08_dp, 3.342783e-08_dp, &
+    1.009504e-09_dp, 4.112522e-09_dp, 1.637224e-09_dp, 3.048654e-10_dp, 6.825087e-10_dp, &
+    1.493166e-10_dp, 3.500323e-10_dp, 2.655267e-11_dp, 6.984064e-11_dp, 5.177358e-11_dp, &
+    7.483559e-12_dp, 3.420646e-11_dp, 8.018778e-12_dp, 5.809091e-12_dp, 1.009504e-10_dp, &
+    1.599957e-10_dp, 1.393504e-10_dp, 1.159066e-11_dp, 3.500323e-11_dp, 4.831791e-12_dp, &
+    2.421632e-11_dp, 8.396691e-12_dp, 3.048654e-12_dp, 1.081703e-11_dp, 1.836996e-12_dp, &
+    1.159066e-11_dp, 2.780406e-12_dp, 7.657873e-12_dp, 1.159066e-12_dp, 6.369542e-12_dp, &
+    1.159066e-12_dp, 6.517907e-12_dp, 6.517907e-12_dp, 2.312641e-11_dp, 2.208555e-11_dp, &
+    7.657873e-12_dp, 7.313212e-12_dp, 5.177358e-11_dp, 9.640691e-13_dp &
+  ]
+
 end module

test/test_equilibrate.f90

@@ -309,6 +309,21 @@ subroutine test()
       endif
     enddo
 
+    converged = cea%solve_metallicity(1.0e6_dp, 1000.0_dp, 1.0_dp, err=err)
+    if (allocated(err)) then
+      print*,err
+      stop 1
+    endif
+    if (.not.converged) stop 1
+
+    do i = 1,size(cea%molfracs_species)
+      if (.not.is_close(cea%molfracs_species(i),correct_answer(i),tol=1.0e-4_dp) .and. cea%molfracs_species(i) > 1.0e-50_dp) then
+        print*,cea%molfracs_species(i),correct_answer(i)
+        print*,'ChemEquiAnalysis failed to compute the right equilibrium.'
+        stop 1
+      endif
+    enddo
+
     print*,'test passed.'
 
   end subroutine