Merge pull request #44 from astro-turing/Use_solve_ivp_without_py-pde…

…_wrapper

Use solve_ivp without the py-pde wrapper, but retain fields and grids from py-pde, for more readable code

Pipfile

@@ -15,8 +15,6 @@ py-pde = "*"
 tqdm = "*"
 h5py = "*"
 pint = "*"
-ipython = "*"
-pandas = "*"
 
 [dev-packages]
 pylint = "*"

README.md

@@ -1,17 +1,17 @@
 # Integrating diagenetic equations using Python
 
-This repo was created as an attempt to reproduce the plots shown at the kickoff of the AstroTOM ("Turing or Milankovitch") project by Niklas Hohmann, from his Matlab scripts (available at [github.com/MindTheGap-ERC/LMA-Matlab](https://github.com/MindTheGap-ERC/LMA-Matlab)). 
+This repo was created in an attempt to reproduce the plots shown at the kickoff of the AstroTOM ("Turing or Milankovitch") project by Niklas Hohmann, from his Matlab scripts (available at [github.com/MindTheGap-ERC/LMA-Matlab](https://github.com/MindTheGap-ERC/LMA-Matlab)). 
 
 AstroTOM is an OpenSSI 2021b project from the Netherlands eScience Center and Utrecht University (UU).
 
 Dr. Emilia Jarochowska (UU) is the lead applicant of this project.
 
 After replacing central differencing for the gradients in the five diagenetic equations 40-43 from [L'Heureux (2018)](https://www.hindawi.com/journals/geofluids/2018/4968315/) by forward and backward differencing depending on the sign of U and W as a first step and a Fiadeiro-Veronis spatial difference scheme as a second step, it turns out that these equations can be integrated for more than 13.190 years (the full T*) with an implicit or explicit (in time) solver, but not with a simple Eulerian scheme. A Runge-Kutta solver, with an adaptive timestep will, however, suffice.
-After correcting the value of b (5-->5e-4) it turned out that a stable integration is also possible without a Fiadeiro-Veronis scheme. The `main` branch makes use of a constant porosity diffusion coefficient.
+After correcting the value of b (5-->5e-4) it turned out that a stable integration is also possible without a Fiadeiro-Veronis scheme. We currently make use of a constant porosity diffusion coefficient.
 
-Implicit (in time) solvers with use of Jacobians (in functional forms, so without numerical approximations) are available in the `Use_solve_ivp_without_py-pde_wrapper` branch.
+The implicit (in time) solvers that `solve_ivp` offers can be deployed with its numerically approximated Jacobians and a Jacobian sparsity matrix.
 
-Wide use is made of the [py-pde](https://py-pde.readthedocs.io/en/latest/) package, especially in the `main` branch.
+Wide use is made of the [py-pde](https://py-pde.readthedocs.io/en/latest/) package, especially `CartesianGrid` and `ScalarField`.
 
 ## Installing and using
 To run this code, you need `git` and `conda` or `pip` to install .

marlpde/Evolve_scenario.py

@@ -1,16 +1,20 @@
 #!/usr/bin/env python
 
-from datetime import datetime
-import os
 from dataclasses import asdict
 import inspect
-import matplotlib.pyplot as plt
+import os
+import time
+from datetime import datetime
 import h5py
-from pde import CartesianGrid, ScalarField, FileStorage, LivePlotTracker
-from pde import DataTracker
-from pde.grids.operators.cartesian import _make_derivative
-from parameters import Map_Scenario, Solver, Tracker
 from LHeureux_model import LMAHeureuxPorosityDiff
+from parameters import Map_Scenario, Solver, Tracker
+from pde import CartesianGrid, ScalarField
+from pde.grids.operators.cartesian import _make_derivative
+from scipy.integrate import solve_ivp
+from tqdm import tqdm
+import matplotlib
+import matplotlib.pyplot as plt
+matplotlib.use("AGG")
 
 def integrate_equations(solver_parms, tracker_parms, pde_parms):
     '''
@@ -41,12 +45,6 @@ def integrate_equations(solver_parms, tracker_parms, pde_parms):
     depths.register_operator("grad_forw", \
         lambda grid: _make_derivative(grid, method="forward"))
 
-    AragoniteSurface = ScalarField(depths, CAIni)
-    CalciteSurface = ScalarField(depths, CCIni)
-    CaSurface = ScalarField(depths, cCaIni)
-    CO3Surface = ScalarField(depths, cCO3Ini)
-    PorSurface = ScalarField(depths, PhiIni)
-
     # I need those two fields for computing coA, which is rather involved.
     # There may be a simpler way of selecting these depths, but I haven't
     # found one yet. For now these two Heaviside step functions.
@@ -55,77 +53,152 @@ def integrate_equations(solver_parms, tracker_parms, pde_parms):
     not_too_deep = ScalarField.from_expression(depths,
                    f"heaviside({DeepLimit/Xstar}-x, 0)")    
 
-    # Not all keys from kwargs are LMAHeureuxPorosityDiff arguments.
+    # Not all keys from pde_parms are LMAHeureuxPorosityDiff arguments.
     # Taken from https://stackoverflow.com/questions/334655/passing-a-\
     # dictionary-to-a-function-as-keyword-parameters
-    filtered_kwargs = {k: v for k, v in pde_parms.items() if k in [p.name for p
-                       in inspect.signature(LMAHeureuxPorosityDiff).parameters.\
-                        values()]}
+    filtered_pde_parms = {k: v for k, v in pde_parms.items() if k in [p.name for
+                          p in
+                          inspect.signature(LMAHeureuxPorosityDiff).parameters.\
+                          values()]}
 
-    eq = LMAHeureuxPorosityDiff(AragoniteSurface, CalciteSurface, CaSurface, 
-                                CO3Surface, PorSurface, not_too_shallow, 
-                                not_too_deep, **filtered_kwargs)             
+    slices_all_fields = [slice(i * Number_of_depths, (i+1) * Number_of_depths) \
+                             for i in range(5)]            
+
+    eq = LMAHeureuxPorosityDiff(depths, slices_all_fields, not_too_shallow, 
+                                not_too_deep, **filtered_pde_parms)     
+
+    # Setting initial values for all five fields in three steps.
+    # This is perhaps more involved than necessary, since the
+    # end result is a 1D Numpy array with a length of five times
+    # the number of depths. The number five comes from the five
+    # fields we are tracking, by integrating five pdes.
+    # First step.
+    AragoniteSurface = ScalarField(depths, CAIni)
+    CalciteSurface = ScalarField(depths, CCIni)
+    CaSurface = ScalarField(depths, cCaIni)
+    CO3Surface = ScalarField(depths, cCO3Ini)
+    PorSurface = ScalarField(depths, PhiIni)
 
+    # Second step.
     state = eq.get_state(AragoniteSurface, CalciteSurface, CaSurface, 
                          CO3Surface, PorSurface)
+    # Third step.
+    y0 = state.data.ravel()   
 
+    no_progress_updates = tracker_parms["no_progress_updates"]
+
+    start_computing = time.time()
+    with tqdm(total=no_progress_updates) as pbar:
+        t0 = solver_parms["t_span"][0]
+        end_time = solver_parms["t_span"][1]
+        progress_bar_args = [pbar, (end_time - t0) / no_progress_updates, t0]
+
+        # The backend parameter determines which function should be called to 
+        # determine the right-hand sides of the five pdes. It can be either
+        # a Numpy-based function or a Numba-based function. The latter is 
+        # faster.
+        backend = solver_parms["backend"]
+        # "backend" is not an argument for solve_ivp, so remove it now.
+        del solver_parms["backend"]
+
+        sol = solve_ivp(eq.fun if backend=="numpy" else eq.fun_numba, 
+                        y0=y0, **solver_parms,
+                        t_eval= tracker_parms["t_eval"], 
+                        events = [eq.zeros, eq.zeros_CA, eq.zeros_CC, \
+                        eq.ones_CA_plus_CC, eq.ones_Phi, eq.zeros_U, \
+                        eq.zeros_W], args=progress_bar_args)
+    end_computing = time.time()
+
+    print()
+    print(f"Number of rhs evaluations = {sol.nfev} \n")
+    print(f"Number of Jacobian evaluations = {sol.njev} \n")
+    print(f"Number of LU decompositions = {sol.nlu} \n")
+    print(f"Status = {sol.status} \n")
+    print(f"Success = {sol.success} \n")
+    f = sol.t_events[0]
+    print(("Times, in years, at which any field at any depth crossed zero: "\
+      +', '.join(['%.2f']*len(f))+"") % tuple([Tstar * time for time in f]))
+    print()
+    g = sol.t_events[1]
+    print(("Times, in years, at which CA at any depth crossed zero: "\
+      +', '.join(['%.2f']*len(g))+"") % tuple([Tstar * time for time in g]))
+    print()
+    h = sol.t_events[2]
+    print(("Times, in years, at which CC at any depth crossed zero: "\
+      +', '.join(['%.2f']*len(h))+"") % tuple([Tstar * time for time in h]))
+    print()
+    k = sol.t_events[3]
+    print(("Times, in years, at which CA + CC at any depth crossed one: "\
+      +', '.join(['%.2f']*len(k))+"") % tuple([Tstar * time for time in k]))
+    print()
+    l = sol.t_events[4]
+    print(("Times, in years, at which the porosity at any depth crossed one: "\
+      +', '.join(['%.2f']*len(l))+"") % tuple([Tstar * time for time in l]))
+    print()
+    m = sol.t_events[5]
+    print(("Times, in years, at which U at any depth crossed zero: "\
+      +', '.join(['%.2f']*len(m))+"") % tuple([Tstar * time for time in m]))
+    print()
+    n = sol.t_events[6]
+    print(("Times, in years, at which W at any depth crossed zero: "\
+      +', '.join(['%.2f']*len(n))+"") % tuple([Tstar * time for time in n]))
+    print()
+
+    print(f"Message from solve_ivp = {sol.message} \n")
+    print(("Time taken for solve_ivp is "
+          f"{end_computing - start_computing:.2e}s. \n"))
+
+    if sol.status == 0:
+        covered_time = Tstar * end_time
+    else:
+        covered_time = pbar.n * Tstar * end_time / no_progress_updates 
+
     # Store your results somewhere in a subdirectory of a parent directory.
     store_folder = "../Results/" + \
                    datetime.now().strftime("%d_%m_%Y_%H_%M_%S" + "/")
     os.makedirs(store_folder)
     stored_results = store_folder + "LMAHeureuxPorosityDiff.hdf5"
     # Keep a record of all parameters.
-    storage_parms = solver_parms | tracker_parms | pde_parms
-    storage = FileStorage(stored_results, info=storage_parms)
-
-    if tracker_parms["live_plotting"]:
-        live_plots = LivePlotTracker(interval=\
-                                     tracker_parms["plotting_interval"], \
-                                     title="Integration results",
-                                     show=True, max_fps=1, \
-                                     plot_args ={"ax_style": {"ylim": (0, 1.5)}})
-    else:
-        live_plots = None
+    stored_parms = solver_parms | tracker_parms | pde_parms
+    # Remove items which will raise a problem when storing as metadata in an
+    # hdf5 file
+    if "jac_sparsity" in stored_parms:
+        del stored_parms["jac_sparsity"]
 
-    if tracker_parms["track_U_at_bottom"]:
-        data_tracker = DataTracker(eq.track_U_at_bottom, \
-                               interval = tracker_parms["data_tracker_interval"])
-    else:
-        data_tracker = None
+    # Reshape the solutions such that you get one row per field.
+    # The third axis, i.e.the time axis, can remain unchanged.
+    field_solutions = sol.y.reshape(5, Number_of_depths, sol.y.shape[-1])
 
-    sol, info = eq.solve(state, **solver_parms, tracker=["progress", \
-                         storage.tracker(\
-                             tracker_parms["progress_tracker_interval"]),\
-                         live_plots, data_tracker])
+    with h5py.File(stored_results, "w") as stored:
+        stored.create_dataset("solutions", data=field_solutions)
+        stored.create_dataset("times", data=sol.t)
+        for event_index, _ in enumerate(sol.t_events):
+            stored.create_dataset("event_" + str(event_index), 
+                                  data=sol.t_events[event_index])
+        stored.attrs.update(stored_parms)
 
-    print()
-    print(f"Meta-information about the solution : {info} \n")        
-
-    covered_time_span = Tstar * info["controller"]["t_final"]
-
-    if tracker_parms["track_U_at_bottom"]:
-        with h5py.File(stored_results, 'a') as hf:
-            U_grp = hf.create_group("U")
-            U_grp.create_dataset("U_at_bottom", \
-                              data=data_tracker.dataframe.to_numpy())
+    # We will be plotting only the distributions corresponding to the last time.
+    # Thus no point in returning all data. Moreover, all data have been saved.
 
-    return sol, covered_time_span, depths, Xstar, store_folder
+    return field_solutions[:, :, -1], covered_time, depths, Xstar, store_folder
 
-def Plot_results(sol, covered_time, depths, Xstar, store_folder):
+def Plot_results(last_field_sol, covered_time, depths, Xstar, store_folder):
     '''
     Plot the five fields at the end of the integration interval as a function
     of depth.
     '''
     fig, ax = plt.subplots()
-    fig.suptitle(f"Distributions after {covered_time:.2f} years")
+    fig.suptitle(f"Distributions after {covered_time:.2e} years")
     # Marker size
     ms = 5
     plotting_depths = ScalarField.from_expression(depths, "x").data * Xstar
-    ax.plot(plotting_depths, sol.data[0], "v", ms = ms, label = "CA")
-    ax.plot(plotting_depths, sol.data[1], "^", ms = ms, label = "CC")
-    ax.plot(plotting_depths, sol.data[2], ">", ms = ms, label = "cCa")
-    ax.plot(plotting_depths, sol.data[3], "<", ms = ms, label = "cCO3")
-    ax.plot(plotting_depths, sol.data[4], "o", ms = ms, label = "Phi")
+
+    ax.plot(plotting_depths, last_field_sol[0], "v", ms = ms, label = "CA")
+    ax.plot(plotting_depths, last_field_sol[1], "^", ms = ms, label = "CC")
+    ax.plot(plotting_depths, last_field_sol[2], ">", ms = ms, label = "cCa")
+    ax.plot(plotting_depths, last_field_sol[3], "<", ms = ms, label = "cCO3")
+    ax.plot(plotting_depths, last_field_sol[4], "o", ms = ms, label = "Phi")
+
     ax.set_xlabel("Depth (cm)")
     ax.set_ylabel("Compositions and concentrations (dimensionless)")
     ax.legend(loc='upper right')
@@ -135,6 +208,4 @@ def Plot_results(sol, covered_time, depths, Xstar, store_folder):
     pde_parms = asdict(Map_Scenario()) 
     solver_parms = asdict(Solver()) 
     tracker_parms = asdict(Tracker())
-    solution, covered_time, depths, Xstar, store_folder = \
-        integrate_equations(solver_parms, tracker_parms, pde_parms)
-    Plot_results(solution, covered_time, depths, Xstar, store_folder)
+    Plot_results(*integrate_equations(solver_parms, tracker_parms, pde_parms))