#!/usr/bin/env python
# coding: utf-8

# # Carbonate-rich lakes modeling on the early Earth
# 
# <p class="acknowledgement">Written by Svetlana Kyas (ETH Zurich) on Mar 31th, 2022</p>
# 
# ```{attention}
# Always make sure you are using the [latest version of Reaktoro](https://anaconda.org/conda-forge/reaktoro). Otherwise, some new features documented on this website will not work on your machine and you may receive unintuitive errors. Follow these [update instructions](updating_reaktoro_via_conda) to get the latest version of Reaktoro!
# ```
# 
# Reaktoro can be used for various geobiological simulations. One of them is the modeling of carbonate-rich lakes, which were relatively common on the early Earth.
# 
# ```{note}
# This tutorial is one of two tutorials that follow the paper `Toner2020` and attempts to replicate the geobiological simulations performed in it (see also the second tutorial [**Phosphate accumulation in carbonate-rich brines**](geobiology-phreeqc-fixed-fugacity.ipynb)). This work was done in collaboration with Cara Magnabosco and Laura Murzakhmetov, ETH-Zurich.
# ```
# 
# ## Solubility of phosphate in the hydroxyl-, fluorapatite- and calcite-rich lakes
# 
# Carbonate-rich lakes can be explained by the strong chemical weathering of abundant, fresh volcanic rocks in the CO<sub>2</sub>-rich atmosphere of the early Earth. Weathering released phosphate from apatites (a group of phosphate minerals) and carbonate alkalinity from other minerals that accumulated in closed basins.
# 
# In Reaktoro, the Earth’s CO<sub>2</sub>-rich atmosphere can be modelled by fixing the fugacity of the simulated chemical states. In particular, a consequence of early Earth’s CO<sub>2</sub>-rich atmosphere (corresponding to the partial pressure of CO<sub>2</sub> from -2 to 0) is that it would have enhanced the weathering of hydroxyl- and fluorapatite in mafic rocks by lowering the pH of surface waters.
# 
# Below, we load python libraries important to carry out the simulations in this tutorial.

# In[1]:


from reaktoro import *
import numpy as np
import pandas as pd
import math as math


# The database is loaded from the extended version of the phreeqc database, including the properties of fluorapatite and hydroxylapatite minerals. The chemical system is composed of an aqueous phase and minerals calcite (CaCO<sub>3</sub>), fluorapatite (Ca<sub>5</sub>(F)(PO<sub>4</sub>)<sub>3</sub>), and hydroxyapatite (Ca<sub>5</sub>(PO<sub>4</sub>)<sub>3</sub>OH). The corresponding aqueous and chemical properties are also defined below.

# In[2]:


# Fluorapatite
#     Ca5(F)(PO4)3 = 5Ca+2 + F- + 3PO4-3
#     log_k     -59.6
#     -analytical_expression -1917.945184 0 87834.57783 631.9611081 0 0
# Hydroxylapatite
#     Ca5(OH)(PO4)3 = 5Ca+2 + OH- + 3PO4-3
#     log_k     -58.517
#     -analytical_expression -1.6657 -0.098215 -8219.41 0 0 0
db = PhreeqcDatabase.fromFile('phreeqc-extended.dat')

# Define the aqueous phase
solution = AqueousPhase(speciate(StringList("H O C Na Cl Ca P")))
solution.set(chain(
    ActivityModelHKF(),
    ActivityModelDrummond("CO2")
))

# Define minerals' phases
minerals = MineralPhases("Calcite Fluorapatite Hydroxylapatite")

# Define the chemical system
system = ChemicalSystem(db, solution, minerals)

# Define aqueous and chemical properties
aprops = AqueousProps(system)
props = ChemicalProps(system)


# To tell the solver that fugacity will be constrained in this chemical system, we need to define equilibrium specifications and define corresponding conditions. The first specifies what will be a constraint and the second by which value (defined below for the range of fugacities).

# In[3]:


# Define equilibrium specifications
specs = EquilibriumSpecs(system)
specs.temperature()
specs.pressure()
specs.fugacity("CO2")

# Define conditions to be satisfied at the chemical equilibrium state
conditions = EquilibriumConditions(specs)
conditions.pressure(1.0, "atm")

# Define the equilibrium solver
solver = EquilibriumSolver(specs)
opts = EquilibriumOptions()
opts.epsilon = 1e-13
solver.setOptions(opts)


# Below, we perform a series of experiments to determine how much phosphate can be accumulated by abiotic processes in carbonate-rich lakes. In the other words, we calculate the solubility of fluorapatite and hydroxyapatite in the presence of calcite buffer as a function of temperature and CO<sub>2</sub> partial pressure.
# 
# The block below defines the array of the partial CO<sub>2</sub> pressures and temperatures as well as the data blocks storing results for different temperatures. Finally, we run equilibrium calculations in the loop for different partial CO<sub>2</sub> pressure and temperatures and collect the following data:
# * pH of the equilibrated state,
# * phosphate level in the solution (the amount of element P), and
# * calcium level in the solution (the amount of element Ca).

# In[4]:


# Auxiliary arrays
num_temperatures = 3
num_log10pCO2s = 51
temperatures = np.array([0, 25, 50])
co2pressures = np.linspace(-4.0, 0.0, num=num_log10pCO2s)

# Output dataframe
data = pd.DataFrame(columns=["T", "ppCO2", "pH", "amount_P", "amount_Ca"])

for ppCO2 in co2pressures:
    for T in temperatures:

        conditions.temperature(T, "celsius")
        conditions.fugacity("CO2", 10 ** ppCO2, 'atm')

        state = ChemicalState(system)
        state.set("H2O"            ,   1.0, "kg")
        state.set("Calcite"        ,  10.0, "mol")
        state.set("Fluorapatite"   ,  10.0, "mol")
        state.set("Hydroxylapatite",  10.0, "mol")
        state.set("CO2"            , 100.0, "mol")

        # Equilibrate the solution with the given initial chemical state and desired conditions at the equilibrium
        solver.solve(state, conditions)

        aprops.update(state)
        props.update(state)

        # Collect the value to be added to the dataframe in the following order
        # "T", "ppCO2", "pH", "amount_P", "amount_Ca"
        data.loc[len(data)] = [T, ppCO2, float(aprops.pH()),
                               float(props.elementAmountInPhase("P", "AqueousPhase")),
                               float(props.elementAmountInPhase("Ca", "AqueousPhase"))]


# Apatites are more soluble at lower pH and weather more rapidly in CO<sub>2</sub>-acidified stream and rainwater, resulting in potentially high phosphate fluxes to carbonate-rich lakes on the early Earth. Below, we plot the dependency of the phosphates and carbonates solubility on the partial CO<sub>2</sub> pressures, which rises with growing log<sub>10</sub>(pCO<sub>2</sub>).

# In[6]:


from bokeh.plotting import figure, show, gridplot
from bokeh.models import HoverTool
from bokeh.io import output_notebook
from bokeh.models import ColumnDataSource
output_notebook()

# ----------------------------------- #
# Plot P amount
# ----------------------------------- #
hovertool1 = HoverTool()
hovertool1.tooltips = [("amount(P)", "@amount_P mol"), 
                       ("T", "@T °C"), 
                       ("ppCO2", "@ppCO2")]

p1 = figure(
    title="DEPENDENCE PHOSPHORUS AMOUNT ON TEMPERATURE",
    x_axis_label=r'PARTIAL PRESSURE CO2 [-]',
    y_axis_label='AMOUNT OF P IN SOLUTION [MOL]',
    sizing_mode="scale_width",
    height=300)

p1.add_tools(hovertool1)

colors = ['teal', 'darkred', 'indigo']
for T, color in zip(temperatures, colors):
    df_T = ColumnDataSource(data[data['T'] == T])
    p1.line("ppCO2", "amount_P", legend_label=f'{T} °C', line_width=3, line_cap="round", line_color=color, source=df_T)

p1.legend.location = 'top_left'

# ----------------------------------- #
# Plot Ca amount
# ----------------------------------- #
hovertool2 = HoverTool()
hovertool2.tooltips = [("amount(Ca)", "@amount_Ca mol"), 
                       ("T", "@T °C"),
                       ("ppCO2", "@ppCO2")]


p2 = figure(
    title="DEPENDENCE CALCIUM AMOUNT ON TEMPERATURE",
    x_axis_label=r'PARTIAL PRESSURE CO2 [-]',
    y_axis_label='AMOUNT OF CA IN SOLUTION [MOL]',
    sizing_mode="scale_width",
    height=300)

p2.add_tools(hovertool2)

colors = ['teal', 'darkred', 'indigo']
for T, color in zip(temperatures, colors):
    df_T = ColumnDataSource(data[data['T'] == T])
    p2.line("ppCO2", "amount_Ca", legend_label=f'{T} °C', line_width=3, line_cap="round", line_color=color, source=df_T)

p2.legend.location = 'top_left'

grid = gridplot([[p1], [p2]])

show(grid)


# Alternatively, we can plot the two-dimensional dependence of the phosphate solubility and pH on the range of partial CO<sub>2</sub> pressures and temperatures. The code below collects the data for such a plot:

# In[7]:


num_temperatures = 101
num_log10pCO2s = 101
temperatures =  np.linspace(0.0, 50.0, num=num_temperatures)
co2ppressures = np.linspace(-4.0, 0.0, num=num_log10pCO2s)

data_size = 2
data_pH = np.zeros((num_temperatures, num_log10pCO2s))
data_P = np.zeros((num_temperatures, num_log10pCO2s))

for i in range(0, num_temperatures):
    for j in range(0, num_log10pCO2s):

        conditions.temperature(temperatures[i], "celsius")
        conditions.pressure(1.0, "atm")
        conditions.fugacity("CO2", 10 ** co2ppressures[j], 'atm')

        state = ChemicalState(system)
        state.set("H2O"            ,   1.0, "kg")
        state.set("Calcite"        ,  10.0, "mol")
        state.set("Fluorapatite"   ,  10.0, "mol")
        state.set("Hydroxylapatite",  10.0, "mol")
        state.set("CO2"            , 100.0, "mol")

        solver.solve(state, conditions)

        aprops.update(state)
        props.update(state)

        data_pH[i, j] = float(aprops.pH())
        data_P[i, j] = float(math.log10(props.elementAmountInPhase("P", "AqueousPhase")))


# From the plots below, we see that relatively low CO<sub>2</sub> pressures log<sub>10</sub>(pCO<sub>2</sub>) = −3.5 (corresponding to the on present-day Earth) limit phosphate to ≤1 μM, which is consistent with phosphate concentrations found in a majority of present-day rivers and surface waters. However, in CO<sub>2</sub>-rich atmospheres relevant to the early Earth (corresponding to log<sub>10</sub>(pCO<sub>2</sub>) = 0.01 to 1 bar), resulting phosphate concentrations are one or two orders of magnitude higher. The latter implies a much higher phosphate weathering flux on the early Earth from streams into lakes.

# In[8]:


from bokeh.plotting import gridplot, figure, show
from bokeh.io import output_notebook
from bokeh.models import LinearColorMapper, BasicTicker, ColorBar
from bokeh.palettes import Spectral11 as palette
output_notebook()

x = co2ppressures
y = temperatures
xx, yy = np.meshgrid(co2ppressures, temperatures)

# Flip the colors in color palette
r_palette = palette[::-1]
color_mapper_1 = LinearColorMapper(palette=r_palette)

p1 = figure(tooltips=[("pH", "@image"), ("T", "$y"),("ppCO2", "$x")],
           title="DEPENDENCE OF PH ON PARTIAL CO2 PRESSURE AND TEMPERATURE",
           x_axis_label=r'PARTIAL CO2 PRESSURE [-]',
           y_axis_label='TEMPERATURE [C]',
           sizing_mode="scale_width",
           x_range=(co2ppressures[0], co2ppressures[-1]),
           y_range=(temperatures[0], temperatures[-1]))

# must give a vector of image data for image parameter
p1.image(image=[data_pH], color_mapper=color_mapper_1,
         x=co2ppressures[0], y=temperatures[0],
         dw=co2ppressures[-1]-co2ppressures[0],
         dh=temperatures[-1]-temperatures[0],
         level="image")
p1.grid.grid_line_width = 0.5

color_bar = ColorBar(color_mapper=color_mapper_1, ticker= BasicTicker(), location=(0,0))
p1.add_layout(color_bar, 'right')

color_mapper_2 = LinearColorMapper(palette="Cividis256")

p2 = figure(tooltips=[("log10(P)", "@image"), ("T", "$y"), ("ppCO2", "$x")],
           title="DEPENDENCE OF LOG10(P) ON PARTIAL CO2 PRESSURE AND TEMPERATURE",
           x_axis_label=r'PARTIAL CO2 PRESSURE [-]',
           y_axis_label='TEMPERATURE [C]',
           sizing_mode="scale_width",
           x_range=(co2ppressures[0], co2ppressures[-1]),
           y_range=(temperatures[0], temperatures[-1]))

# must give a vector of image data for image parameter
p2.image(image=[data_P], color_mapper = color_mapper_2,
         x=co2ppressures[0], y=temperatures[0],
         dw=co2ppressures[-1]-co2ppressures[0],
         dh=temperatures[-1]-temperatures[0],
         level="image")
p2.grid.grid_line_width = 0.5

color_bar = ColorBar(color_mapper=color_mapper_2, ticker= BasicTicker(), location=(0,0))
p2.add_layout(color_bar, 'right')

grid = gridplot([[p1], [p2]])

show(grid)


# ## Solubility of phosphate in the vivianite- and siderite-rich lakes
# 
# A potential sink for soluble phosphate on the early Earth is reduced soluble iron (Fe<sup>+2</sup>), which precipitates with phosphate as the low-solubility mineral vivianite (Fe<sub>3</sub>(PO<sub>4</sub>)<sub>2</sub> · 8H<sub>2</sub>O). However, at the same time, Fe<sup>+2</sup> also precipitates as siderite (FeCO<sub>3</sub>) in carbonate-rich brines, which limits Fe<sup>+2</sup> to low levels and increases the solubility of phosphate from vivianite. Below, we model this behavior and compare the results to the system with apatities.

# In[9]:


# Define the aqueous phase (including Fe-containing species)
solution = AqueousPhase(speciate("H O C Na Cl Ca P Fe"))
solution.set(chain(
    ActivityModelHKF(),
    ActivityModelDrummond("CO2")
))

# Define mineral phases
minerals = MineralPhase("Siderite Vivianite")

# Create an updated chemical system
system = ChemicalSystem(db, solution, minerals)

aprops = AqueousProps(system)
props = ChemicalProps(system)

specs = EquilibriumSpecs(system)
specs.temperature()
specs.pressure()
specs.fugacity("CO2")

conditions = EquilibriumConditions(specs)

solver = EquilibriumSolver(specs)
opts = EquilibriumOptions()
opts.epsilon = 1e-13
solver.setOptions(opts)

# Auxiliary arrays
num_log10pCO2s = 51
temperatures = np.array([0, 25, 50])
co2pressures = np.linspace(-4.0, 0.0, num=num_log10pCO2s)

# Output dataframe
df = pd.DataFrame(columns=["T", "ppCO2", "pH", "amount_P", "amount_Fe"])

for ppCO2 in co2pressures:
    for T in temperatures:

        conditions.temperature(T, "celsius")
        conditions.pressure(1.0, "atm")
        conditions.fugacity("CO2", 10 ** (ppCO2), 'atm')

        state = ChemicalState(system)
        state.set("H2O"      ,   1.0, "kg")
        state.set("Siderite" ,  10.0, "mol")
        state.set("Vivianite",  10.0, "mol")
        state.set("CO2"      , 100.0, "mol")

        solver.solve(state, conditions)

        aprops.update(state)
        props.update(state)
        
        # Collect the value to be added to the dataframe in the following order
        # "T", "ppCO2", "pH", "amount_P", "amount_Fe"
        df.loc[len(df)] = [T, ppCO2, float(aprops.pH()),
                               float(props.elementAmountInPhase("P", "AqueousPhase")),
                               float(props.elementAmountInPhase("Fe", "AqueousPhase"))]


# Apatites are more soluble at lower pH and weather faster in CO<sub>2</sub>-acidified stream and rainwater, resulting in potentially high phosphate fluxes to carbonate-rich lakes on the early Earth.

# In[10]:


# ----------------------------------- #
# Plot P amount
# ----------------------------------- #
hovertool1 = HoverTool()
hovertool1.tooltips = [("T", "@T °C"),
                      ("ppCO2", "@ppCO2"),
                      ("amount(P)", "@amount_P mol")]

p1 = figure(
    title="DEPENDENCE PHOSPHORUS AMOUNT ON TEMPERATURE \n (IN THE SYSTEM WITH SIDERITE AND VIVIANITE)",
    x_axis_label=r'PARTIAL PRESSURE CO2 [-]',
    y_axis_label='AMOUNT OF P IN SOLUTION [MOL]',
    sizing_mode="scale_width",
    height=300)

p1.add_tools(hovertool1)

colors = ['rosybrown', 'steelblue', 'seagreen', 'palevioletred']
for T, color in zip(temperatures, colors):
    df_T = ColumnDataSource(df[df['T'] == T])
    p1.line("ppCO2", "amount_P", legend_label=f'{T} °C', line_width=3, line_cap="round", line_color=color, source=df_T)

p1.legend.location = 'top_left'

# ----------------------------------- #
# Plot Ca amount
# ----------------------------------- #
hovertool2 = HoverTool()
hovertool2.tooltips = [("T", "@T °C"),
                      ("ppCO2", "@ppCO2"),
                      ("amount(Fe)", "@amount_Fe mol")]


p2 = figure(
    title="DEPENDENCE IRON AMOUNT ON TEMPERATURE \n (IN THE SYSTEM WITH SIDERITE AND VIVIANITE)",
    x_axis_label=r'PARTIAL PRESSURE CO2 [-]',
    y_axis_label='AMOUNT OF FE IN SOLUTION [MOL]',
    sizing_mode="scale_width",
    height=300)

p2.add_tools(hovertool2)

colors = ['teal', 'darkred', 'indigo']
for T, color in zip(temperatures, colors):
    df_T = ColumnDataSource(df[df['T'] == T])
    p2.line("ppCO2", "amount_Fe", legend_label=f'{T} °C', line_width=3, line_cap="round", line_color=color, source=df_T)

p2.legend.location = 'top_left'

grid = gridplot([[p1], [p2]])

show(grid)


# Due to the lower solubility of siderite (in comparison to calcite) as well as the fact that its nucleation and growth are not inhibited by phosphate, Fe<sup>+2</sup> concentrations in anoxic, phosphate- and carbonate-rich brines is expected to be lower than Ca<sup>+2</sup>. We see it from the plot of the solubility of iron. This suggests that Ca<sup>+2</sup>, not Fe<sup>+2</sup>, would have controlled phosphate concentration in carbonate-rich lakes on the early Earth.

#